ZacharyZhang-NY/Kigi-CLI
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

M2 audit: excise the Computer Hub stack — Kigi's last remote-cloud surface

Browse Source 
Removed root-and-branch for the zero-egress guarantee (the hub was xAI's
remote-workspace/cloud-sandbox service):

- Crates deleted: kigi-computer-hub-core, kigi-computer-hub-sdk,
  kigi-computer-hub-mcp-adapter, kigi-workspace-client (hub-proxied
  workspace RPC client), and kigi-tracing (its sole network path was the
  OTLP gRPC exporter; zero consumers remained). kigi-tracing-macros
  (purely local) stays.
- kigi-workspace: every hub surface deleted — hub server/channel/auth,
  HITL-over-hub permissions, donation/metrics pumps, file upload RPCs,
  hub tool-snapshot merge (resolve pipeline is MCP-only now),
  WorkspaceOps::Proxy. Local worktrees, sessions, leader IPC, MCP, and
  the ACP permission prompt path are untouched; LocalRegistry re-homed
  into kigi-tool-runtime on the existing ToolDyn types so in-process
  tool dispatch is unchanged.
- kigi-shell: leader workspace-exposure control surface (incl. the
  wss://computer-hub... URL), [hub] config, ObservabilityBridge, hub
  WebSocket proxy, dead OTLP config knobs. ClientMode::Headless (never
  constructed) removed.
- kigi-tui/bin: hidden `kigi workspace` command removed (`kigi
  worktree` stays).
- Renames: --xai-api-base-url → --api-base-url / KIGI_API_BASE_URL /
  [endpoints] api_base_url (serde alias keeps old configs working; the
  flag feeds BYOK/custom-endpoint routing, not main inference);
  grok_version → kigi_version in inspect/models-cache/trace metadata
  (old caches self-heal via version-mismatch refetch).
- Dependency tree: dropped fastrace*, opentelemetry-otlp/http/proto,
  tokio-tungstenite from the workspace; fixed the 4 real useless_format
  violations the fastrace lint allowance was masking and removed the
  allowance.
- marketplaceAllowlist kept: it gates the LOCAL plugin-marketplace
  feature, not an xAI service.

Known §9 leftover (deliberate, for the M3 sweep): the BYOK default base
URL string. Gates: workspace check/clippy 0/0, fmt, deny ok; suites
green (workspace 1042, shell 4918, tui 6634, tools 2608, tool-runtime
47, mcp 154).
This commit is contained in:
ZacharyZhang-NY
2026-07-17 22:34:10 -04:00
parent 5919526e91
commit fa75eb139a
90 changed files with 452 additions and 9702 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/codegen/kigi-agent/src/builder.rs
-1
View File
@@ -1028,7 +1028,6 @@ impl AgentBuilder {
video_gen_config: self.video_gen_config,
app_builder_deployer_config: self.app_builder_deployer_config,
api_key_provider: self.api_key_provider,
auth_provider: None,
attribution_callback: self.attribution_callback,
system_reminder_tag: self.system_reminder_tag,
},
crates/codegen/kigi-bin/src/main.rs
+2 -212
View File
@@ -37,7 +37,6 @@ use kigi_tui::app::{
AgentCmd, Command, LeaderMgmtArgs, LeaderMgmtCommand, LeaderTargetArgs, PagerArgs,
resolve_use_leader,
};
use kigi_tui::app::{WorkspaceMgmtArgs, WorkspaceMgmtCommand, WorkspaceStartArgs};
use kigi_tui::client_identity::PAGER_CLIENT_VERSION;
use kigi_update::{UpdateConfig, auto_update, enforce_minimum_version_or_exit};
use std::env;
@@ -48,8 +47,8 @@ fn apply_agent_endpoint_args(agent_args: &kigi_tui::app::AgentArgs, config: &mut
if let Some(v) = &agent_args.coding_api_base_url {
config.endpoints.coding_api_base_url = Some(v.clone());
}
if let Some(v) = &agent_args.xai_api_base_url {
config.endpoints.xai_api_base_url = v.clone();
if let Some(v) = &agent_args.api_base_url {
config.endpoints.api_base_url = v.clone();
}
}
/// Resolve --agent-profile path: canonicalize and verify the file exists.
@@ -324,194 +323,6 @@ fn ensure_control_caps(reg: &LeaderRegistration) -> Result<&LeaderCapabilities>
.as_ref()
.ok_or_else(|| anyhow::anyhow!("Leader does not advertise capabilities (legacy version)"))
}
/// Env override for the `kigi workspace` gate: any truthy value enables the
/// command locally, a falsy one disables it. This is the only gate now that
/// the server-side feature flag (xAI remote settings) is gone.
const WORKSPACE_COMMAND_ENV: &str = "KIGI_WORKSPACE_COMMAND";
/// The `KIGI_WORKSPACE_COMMAND` override, if set (`Some(true)`/`Some(false)`);
/// `None` means unset (the command stays disabled by default).
fn workspace_command_env_override() -> Option<bool> {
std::env::var(WORKSPACE_COMMAND_ENV)
.ok()
.map(|v| env_flag_enabled(&v))
}
/// Resolve the gate: enabled exactly when the env override says so.
fn workspace_command_gate(env_override: Option<bool>) -> bool {
env_override.unwrap_or(false)
}
/// Truthy parse for grok on/off env vars: everything enables except the common
/// falsy spellings (`0`, `false`, `off`, `no`, empty).
fn env_flag_enabled(value: &str) -> bool {
!matches!(
value.trim().to_ascii_lowercase().as_str(),
"" | "0" | "false" | "off" | "no"
)
}
async fn run_workspace_mgmt(args: WorkspaceMgmtArgs) -> Result<()> {
if !workspace_command_gate(workspace_command_env_override()) {
anyhow::bail!(
"`kigi workspace` is experimental and disabled by default. \
Set {WORKSPACE_COMMAND_ENV}=1 to enable it."
)
}
match args.command {
WorkspaceMgmtCommand::Start(a) => workspace_start(a, false).await,
WorkspaceMgmtCommand::Restart(a) => workspace_start(a, true).await,
WorkspaceMgmtCommand::Pause { target, json } => {
workspace_control(&target, json, ControlCommand::WorkspacePause).await
}
WorkspaceMgmtCommand::Resume { target, json } => {
workspace_control(&target, json, ControlCommand::WorkspaceResume).await
}
WorkspaceMgmtCommand::Stop { target, json } => {
workspace_control(&target, json, ControlCommand::WorkspaceStop).await
}
WorkspaceMgmtCommand::Status { target, json } => {
workspace_control(&target, json, ControlCommand::WorkspaceStatus).await
}
}
}
fn ensure_workspace_caps(reg: &LeaderRegistration) -> Result<()> {
let caps = ensure_control_caps(reg)?;
if !caps.workspace_exposure {
anyhow::bail!(
"the running leader does not support workspace exposure — stop the \
leader process and re-run to pick up the new version"
);
}
Ok(())
}
async fn connect_workspace_control(
_agent_config: &AgentConfig,
target: &LeaderTargetArgs,
) -> Result<LeaderClient> {
if target.pid.is_some() {
let (_descriptor, client) = connect_to_leader(target).await?;
return Ok(client);
}
let socket = default_socket_path();
LeaderClient::connect(
socket,
"grok-workspace-cli",
ClientMode::Stdio,
ClientCapabilities::default(),
)
.await
.map_err(|e| {
anyhow::anyhow!(
"no running leader ({e}). \
Start a grok session, or run `grok workspace start`."
)
})
}
async fn workspace_control(
target: &LeaderTargetArgs,
json: bool,
command: ControlCommand,
) -> Result<()> {
let raw_config = kigi_shell::config::load_effective_config_disk_only()
.map_err(|e| anyhow::anyhow!("Failed to load config: {e}"))?;
let agent_config = AgentConfig::new_from_toml_cfg(&raw_config)
.map_err(|e| anyhow::anyhow!("Failed to create agent config: {e}"))?;
let client = connect_workspace_control(&agent_config, target).await?;
ensure_workspace_caps(client.registration())?;
let payload = client.send_control(command).await??;
render_workspace_payload(&payload, json);
client.cancel();
Ok(())
}
async fn workspace_start(args: WorkspaceStartArgs, restart: bool) -> Result<()> {
use kigi_shell::auth::ensure_authenticated;
let raw_config = kigi_shell::config::load_effective_config()
.map_err(|e| anyhow::anyhow!("Failed to load config: {e}"))?;
let agent_config = AgentConfig::new_from_toml_cfg(&raw_config)
.map_err(|e| anyhow::anyhow!("Failed to create agent config: {e}"))?;
let (use_leader, _) = resolve_use_leader(args.leader, args.no_leader, &raw_config, true);
if !use_leader {
anyhow::bail!(
"`grok workspace` requires leader mode (the workspace is shared via the leader).\n\
Enable it with `[cli] use_leader = true` in ~/.kigi/config.toml, or pass --leader."
);
}
ensure_authenticated(
&agent_config.kimi_code_config,
false,
Some("No cached credentials found. Run `kigi login` first."),
)
.await?;
let capabilities = ClientCapabilities {
client_version: Some(PAGER_CLIENT_VERSION.to_string()),
..Default::default()
};
let conn = connect_or_spawn("grok-workspace-cli", ClientMode::Stdio, capabilities)
.await
.map_err(|e| anyhow::anyhow!("failed to start or connect to leader: {e}"))?;
drop(conn);
let target = LeaderTargetArgs::default();
let client = connect_workspace_control(&agent_config, &target).await?;
ensure_workspace_caps(client.registration())?;
if restart {
let _ = client.send_control(ControlCommand::WorkspaceStop).await;
}
let cwd = match args.cwd {
Some(p) => p,
None => std::env::current_dir()
.map_err(|e| anyhow::anyhow!("cannot determine current directory: {e}"))?,
};
let cwd = std::path::absolute(&cwd).unwrap_or(cwd);
let payload = client
.send_control(ControlCommand::WorkspaceStart {
hub_url: args.hub_url.clone(),
cwd: cwd.display().to_string(),
})
.await??;
render_workspace_payload(&payload, args.json);
client.cancel();
Ok(())
}
fn render_workspace_payload(payload: &ControlPayload, json: bool) {
let ControlPayload::WorkspaceStatus {
state,
hub_url,
cwd,
uptime_ms,
active_tool_calls,
sessions,
pid,
} = payload
else {
eprintln!("unexpected control response: {payload:?}");
return;
};
if json {
let value = serde_json::json!(
{ "state" : state, "hubUrl" : hub_url, "cwd" : cwd, "uptimeMs" : uptime_ms,
"activeToolCalls" : active_tool_calls, "sessions" : sessions, "pid" : pid, }
);
println!("{}", serde_json::to_string(&value).unwrap_or_default());
return;
}
if state == "none" {
println!("Workspace exposure: not running (leader PID {pid})");
return;
}
println!("Workspace exposure: {state}");
if let Some(url) = hub_url {
println!(" hub: {url}");
}
if let Some(dir) = cwd {
println!(" cwd: {dir}");
}
println!(" uptime: {}s", uptime_ms / 1000);
println!(" active: {active_tool_calls} tool call(s)");
let session_list = if sessions.is_empty() {
"-".to_string()
} else {
sessions.join(", ")
};
println!(" sessions: {} ({session_list})", sessions.len());
println!(" leader: PID {pid}");
}
/// How to rebuild one session's `session/load` after a leader reconnect.
#[derive(Default, Clone)]
struct CachedSession {
@@ -1604,10 +1415,6 @@ async fn async_main() -> Result<()> {
.map_err(|e| anyhow::anyhow!("Failed to create agent config: {e}"))?;
return kigi_tui::worktree_cmd::run(worktree_args, &agent_config).await;
}
Command::Workspace(workspace_args) => {
init_tracing_simple("cli");
return run_workspace_mgmt(workspace_args).await;
}
Command::Sessions(sessions_args) => {
init_tracing_simple("cli");
return kigi_tui::sessions_cmd::run(sessions_args).await;
@@ -2094,23 +1901,6 @@ mod tests {
"failure path must not flag the startup hook",
);
}
#[test]
fn workspace_command_gate_resolution() {
assert!(workspace_command_gate(Some(true)));
assert!(!workspace_command_gate(Some(false)));
assert!(!workspace_command_gate(None), "unset env defaults to off");
}
#[serial_test::serial(KIGI_WORKSPACE_COMMAND)]
#[test]
fn workspace_command_env_override_parsing() {
unsafe { std::env::remove_var("KIGI_WORKSPACE_COMMAND") };
assert_eq!(workspace_command_env_override(), None);
unsafe { std::env::set_var("KIGI_WORKSPACE_COMMAND", "1") };
assert_eq!(workspace_command_env_override(), Some(true));
unsafe { std::env::set_var("KIGI_WORKSPACE_COMMAND", "off") };
assert_eq!(workspace_command_env_override(), Some(false));
unsafe { std::env::remove_var("KIGI_WORKSPACE_COMMAND") };
}
fn make_state() -> std::sync::Mutex<StdioReplayState> {
std::sync::Mutex::new(StdioReplayState::default())
}
crates/codegen/kigi-config-types/src/lib.rs
-23
View File
@@ -694,11 +694,6 @@ pub struct RemoteSettings {
/// See `Config::resolve_image_edit`.
#[serde(default)]
pub imagine_tools_disabled: Option<Vec<String>>,
/// remote settings gate for the `grok workspace` CLI command (Computer Hub
/// workspace exposure), from `grok_build_settings.workspace_command_enabled`.
/// `Some(true)` enables it; `None`/`Some(false)` (the default) keep it off.
#[serde(default)]
pub workspace_command_enabled: Option<bool>,
/// Master switch for jemalloc heap sampling + threshold dumps.
/// `Some(true)` enables, `Some(false)` kill-switch, `None` = client default off.
#[serde(default)]
@@ -1308,24 +1303,6 @@ mod tests {
assert_eq!(s.folder_trust_enabled, None);
}
#[test]
fn remote_settings_workspace_command_enabled_present() {
let json = r#"{"workspace_command_enabled": true}"#;
let s: RemoteSettings = serde_json::from_str(json).unwrap();
assert_eq!(s.workspace_command_enabled, Some(true));
}
#[test]
fn remote_settings_workspace_command_enabled_false() {
let json = r#"{"workspace_command_enabled": false}"#;
let s: RemoteSettings = serde_json::from_str(json).unwrap();
assert_eq!(s.workspace_command_enabled, Some(false));
}
#[test]
fn remote_settings_workspace_command_enabled_absent() {
let json = r#"{}"#;
let s: RemoteSettings = serde_json::from_str(json).unwrap();
assert_eq!(s.workspace_command_enabled, None);
}
#[test]
fn remote_settings_permission_mode_deserializes() {
let s: RemoteSettings = serde_json::from_str(r#"{"permission_mode": "auto"}"#).unwrap();
assert_eq!(s.permission_mode.as_deref(), Some("auto"));
crates/codegen/kigi-mcp/Cargo.toml
-1
View File
@@ -67,7 +67,6 @@ libc = { workspace = true }
[dev-dependencies]
# Used in servers.rs test `test_same_raw_name_different_servers_no_local_registry_collision`
# to verify MCP tools register into a `LocalRegistry` without collision.
kigi-computer-hub-sdk = { workspace = true }
# `test-util` enables `tokio::time::pause` / `advance`, required by
# the deterministic timing tests in `liveness.rs`.
tokio = { workspace = true, features = ["test-util"] }
crates/codegen/kigi-mcp/src/servers.rs
+1 -1
View File
@@ -5301,7 +5301,7 @@ mod tests {
/// into a `LocalRegistry` must preserve both entries (no silent overwrite).
#[test]
fn test_same_raw_name_different_servers_no_local_registry_collision() {
use kigi_computer_hub_sdk::LocalRegistry;
use kigi_tool_runtime::LocalRegistry;
use kigi_tool_runtime::Tool;
let mcp_state = Arc::new(Mutex::new(McpState::new(vec![])));
crates/codegen/kigi-memory/src/schema.rs
+2 -3
View File
@@ -21,8 +21,7 @@ pub const SCHEMA_VERSION: u32 = 1;
/// path (`kigi_sqlite_journal::JournalMode::open`) — the journal mode depends
/// on the database's filesystem.
pub fn schema_sql(dimensions: usize, vec_available: bool) -> String {
let mut sql = format!(
r#"
let mut sql = r#"
CREATE TABLE IF NOT EXISTS meta (
key TEXT PRIMARY KEY,
value TEXT NOT NULL
@@ -50,7 +49,7 @@ CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(text, content='');
INSERT OR IGNORE INTO meta(key, value) VALUES ('reindex_claim', '');
"#
);
.to_string();
if vec_available {
sql.push_str(&format!(
crates/codegen/kigi-pager-pty-harness/src/content.rs
+2 -2
View File
@@ -91,7 +91,7 @@ impl ContentController {
// KIGI_SHARE_DIR is set in the test runner's env).
("KIGI_SHARE_DIR".into(), kigi_home),
("KIGI_CODE_BASE_URL".into(), self.url()),
("KIGI_XAI_API_BASE_URL".into(), self.url()),
("KIGI_API_BASE_URL".into(), self.url()),
("XAI_API_KEY".into(), "test-key-for-ci".into()),
("KIGI_TELEMETRY_ENABLED".into(), "false".into()),
("KIGI_FEEDBACK_ENABLED".into(), "false".into()),
@@ -279,7 +279,7 @@ mod tests {
content.home().join(".kigi").to_str()
);
assert_eq!(get("KIGI_CODE_BASE_URL"), Some(content.url()));
assert_eq!(get("KIGI_XAI_API_BASE_URL"), Some(content.url()));
assert_eq!(get("KIGI_API_BASE_URL"), Some(content.url()));
assert_eq!(get("XAI_API_KEY").as_deref(), Some("test-key-for-ci"));
assert_eq!(get("KIGI_TELEMETRY_ENABLED").as_deref(), Some("false"));
assert_eq!(get("KIGI_FEEDBACK_ENABLED").as_deref(), Some("false"));
crates/codegen/kigi-shell/Cargo.toml
+1 -3
View File
@@ -79,7 +79,6 @@ kigi-acp-lib = { workspace = true }
axum = { workspace = true, features = ["ws", "multipart"] }
backon = { workspace = true }
webbrowser = { workspace = true }
tokio-tungstenite = { workspace = true, features = ["rustls-tls-native-roots"] }
tokio-rustls = { version = "0.26", default-features = false, features = [
"ring",
"logging",
@@ -161,8 +160,7 @@ parking_lot.workspace = true
dashmap.workspace = true
# Used by acp_session.rs and mcp_servers.rs for the unified
# kigi_tool_runtime::Tool dispatch model and the computer-hub MCP adapter.
kigi-computer-hub-sdk = { workspace = true }
# kigi_tool_runtime::Tool dispatch model.
kigi-tool-runtime = { workspace = true }
kigi-tool-protocol = { workspace = true }
kigi-tool-types = { workspace = true }
crates/codegen/kigi-shell/src/agent/app.rs
+1 -7
View File
@@ -518,11 +518,7 @@ pub async fn run_leader(
lock_path: lock.lock_path().clone(),
socket_suffix: socket_suffix_from_paths(lock.lock_path(), &socket_path).unwrap_or_default(),
leader_binary_version: kigi_version::VERSION.to_string(),
})
.with_default_hub_url(agent_config.hub.url.clone());
// Cloned before control_state moves into the IPC server; auth wired below.
let workspace_control = control_state.workspace.clone();
});
// ── Phase 3: Bind socket and start IPC server (BEFORE auth/prefetch) ──────
//
@@ -654,8 +650,6 @@ pub async fn run_leader(
// process so a refresh can't straddle a suspend.
shared_auth_manager.start_system_power_listener();
// Same manager as the leader, so the exposure never writes auth.json itself.
workspace_control.set_auth_manager(shared_auth_manager.clone());
let auth_manager_for_agent = shared_auth_manager.clone();
let auth_manager_for_config = shared_auth_manager;
crates/codegen/kigi-shell/src/agent/config.rs
+19 -588
View File
@@ -43,7 +43,7 @@ pub fn default_agent_type() -> String {
DEFAULT_AGENT_TYPE.to_owned()
}
/// Default base URL for the public xAI API.
pub const XAI_API_BASE_URL_DEFAULT: &str = "https://api.x.ai/v1";
pub const API_BASE_URL_DEFAULT: &str = "https://api.x.ai/v1";
/// One or more environment variable names that may hold a model API key.
///
/// Serde `untagged`: accepts a string or an array in TOML/JSON.
@@ -142,8 +142,10 @@ pub struct EndpointsConfig {
/// default value) lets an org pin the proxy to the default on purpose.
#[serde(skip_serializing_if = "Option::is_none")]
pub coding_api_base_url: Option<String>,
/// Base URL for the public xAI API.
pub xai_api_base_url: String,
/// Base URL for direct (BYOK / external-API-key) API calls.
/// Accepts the legacy `xai_api_base_url` config key.
#[serde(alias = "xai_api_base_url")]
pub api_base_url: String,
/// Optional extra access-header value (applied only with the optional
/// non-production feature, and only for matching first-party hosts).
#[serde(skip_serializing_if = "Option::is_none")]
@@ -166,54 +168,6 @@ pub struct EndpointsConfig {
/// Defaults to `{proxy_url()}/deployment/config`.
#[serde(skip_serializing_if = "Option::is_none")]
pub managed_config_url: Option<String>,
/// Env: `OTEL_EXPORTER_OTLP_ENDPOINT`. OTLP collector base; `/v1/traces` is
/// appended. Legacy repoint of the INTERNAL trace pipeline — deprecated in
/// favor of `KIGI_INTERNAL_OTLP_TRACES_ENDPOINT`, and ignored by the internal
/// pipeline when `KIGI_EXTERNAL_OTEL` is set (the standard `OTEL_*` vars then
/// route the external stream only).
#[serde(skip_serializing_if = "Option::is_none")]
pub otel_exporter_otlp_endpoint: Option<String>,
/// Env: `OTEL_EXPORTER_OTLP_TRACES_ENDPOINT`. Full traces endpoint, used
/// verbatim; overrides `otel_exporter_otlp_endpoint`. Same legacy/deprecation
/// semantics as `otel_exporter_otlp_endpoint`.
#[serde(skip_serializing_if = "Option::is_none")]
pub otel_exporter_otlp_traces_endpoint: Option<String>,
/// Env: `OTEL_EXPORTER_OTLP_HEADERS`. `k=v,k2=v2`; merged onto export headers.
/// Same legacy/deprecation semantics as `otel_exporter_otlp_endpoint`.
#[serde(skip_serializing_if = "Option::is_none")]
pub otel_exporter_otlp_headers: Option<String>,
/// Env: `KIGI_INTERNAL_OTLP_TRACES_ENDPOINT`. Full INTERNAL traces endpoint,
/// used verbatim. Dev/debug repoint of the internal span firehose (replaces
/// the legacy `OTEL_EXPORTER_OTLP_TRACES_ENDPOINT` behavior; used by
/// local-ic-testing / internal dev flows). Wins over the legacy `OTEL_*` vars.
#[serde(skip_serializing_if = "Option::is_none")]
pub grok_internal_otlp_traces_endpoint: Option<String>,
/// Env: `KIGI_INTERNAL_OTLP_HEADERS`. `k=v,k2=v2` extra headers for the
/// internal export (debug). Wins over the legacy `OTEL_EXPORTER_OTLP_HEADERS`.
#[serde(skip_serializing_if = "Option::is_none")]
pub grok_internal_otlp_headers: Option<String>,
/// External-OTEL master switch, captured at construction via
/// [`external_otel_master_switch_resolved`] — the same layered resolution
/// (requirement pin > `KIGI_EXTERNAL_OTEL` env > `[telemetry].otel_enabled`
/// config, managed layers included) that activates the external stream.
/// When set, the standard `OTEL_EXPORTER_OTLP_*` vars are reserved for the
/// external OTEL stream and the internal trace pipeline ignores them
/// entirely — an admin who opts in (by *any* layer, including an org
/// enable distributed via managed config with no env var) never receives
/// the internally-authed firehose. Held as a field (not re-read in the
/// resolvers) so the resolvers stay pure and testable without env races.
#[serde(skip)]
pub external_otel_master_switch: bool,
/// Env: `OTEL_TRACES_EXPORTER`. `otlp` (default) or `none` to disable spans.
#[serde(skip_serializing_if = "Option::is_none")]
pub otel_traces_exporter: Option<String>,
/// Env: `OTEL_BSP_SCHEDULE_DELAY` (OTel) or `OTEL_TRACES_EXPORT_INTERVAL`
/// (Claude alias). Batch flush interval (ms).
#[serde(skip_serializing_if = "Option::is_none")]
pub otel_traces_export_interval: Option<u64>,
/// Env: `OTEL_EXPORTER_OTLP_TIMEOUT`. Export HTTP timeout (ms).
#[serde(skip_serializing_if = "Option::is_none")]
pub otel_exporter_otlp_timeout: Option<u64>,
/// Read by `load_management_api_key_sync()`. Declared for `serde_ignored`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub management_api_key: Option<String>,
@@ -228,18 +182,6 @@ fn blank_as_unset(opt: &Option<String>) -> Option<String> {
.filter(|s| !s.trim().is_empty())
.map(str::to_owned)
}
/// Parse a `k=v,k2=v2` OTLP header list (the `OTEL_EXPORTER_OTLP_HEADERS`
/// format, shared with `KIGI_INTERNAL_OTLP_HEADERS`): split on `,`,
/// `split_once('=')`, trim key/value, skip blank keys, keep empty values.
fn parse_otlp_header_list(raw: &str) -> Vec<(String, String)> {
raw.split(',')
.filter_map(|kv| {
let (k, v) = kv.split_once('=')?;
let k = k.trim();
(!k.is_empty()).then(|| (k.to_string(), v.trim().to_string()))
})
.collect()
}
impl EndpointsConfig {
pub fn has_custom_endpoint(&self) -> bool {
self.models_base_url.is_some() || self.models_list_url.is_some()
@@ -257,22 +199,18 @@ impl EndpointsConfig {
/// Layer the `[endpoints]` table from `config` over the env/default base.
/// No field is derived from another — defaulting is done by the resolvers.
pub(crate) fn from_config_value(config: &toml::Value) -> Self {
let default = Self::default();
let external_otel_master_switch = default.external_otel_master_switch;
let mut base = match toml::Value::try_from(default) {
let mut base = match toml::Value::try_from(Self::default()) {
Ok(v) => v,
Err(_) => return Self::default(),
};
if let Some(endpoints) = config.get("endpoints") {
crate::config::deep_merge_toml(&mut base, endpoints);
}
let mut resolved: Self = base.try_into().unwrap_or_default();
resolved.external_otel_master_switch = external_otel_master_switch;
resolved
base.try_into().unwrap_or_default()
}
/// The subscription proxy base URL through which all auxiliary services (and
/// OAuth/session inference) resolve: explicit `coding_api_base_url`, else
/// [`kigi_env::coding_api_base_url`]. NEVER falls back to `xai_api_base_url` —
/// [`kigi_env::coding_api_base_url`]. NEVER falls back to `api_base_url` —
/// that is the inference endpoint (API-key auth) only.
pub fn proxy_url(&self) -> String {
blank_as_unset(&self.coding_api_base_url).unwrap_or_else(kigi_env::coding_api_base_url)
@@ -283,12 +221,12 @@ impl EndpointsConfig {
.unwrap_or_else(|| self.proxy_url())
}
/// Feedback endpoint — an auxiliary service, so it defaults to the
/// cli-chat-proxy, never `xai_api_base_url`.
/// cli-chat-proxy, never `api_base_url`.
pub fn resolve_feedback_base_url(&self) -> String {
blank_as_unset(&self.feedback_base_url).unwrap_or_else(|| self.proxy_url())
}
/// Managed deployment-config URL (`grok setup`): explicit `managed_config_url`,
/// else `proxy_url` + `/deployment/config`. Never `xai_api_base_url`, so the
/// else `proxy_url` + `/deployment/config`. Never `api_base_url`, so the
/// deployment key reaches the proxy, not the inference host.
pub fn resolve_managed_config_url(&self) -> String {
blank_as_unset(&self.managed_config_url).unwrap_or_else(|| {
@@ -298,99 +236,6 @@ impl EndpointsConfig {
)
})
}
/// INTERNAL OTLP traces endpoint. Precedence:
/// 1. `grok_internal_otlp_traces_endpoint` (verbatim)
/// 2. legacy `otel_exporter_otlp_traces_endpoint` (verbatim) >
/// `otel_exporter_otlp_endpoint` + `/v1/traces` — ONLY when the
/// external-OTEL master switch is unset (back-compat; deprecated)
/// 3. `proxy_url` + `/traces`.
/// Uses the proxy default (not the `xai_api_base_url` fallback) so
/// telemetry reports to xAI even when inference is overridden. When the
/// master switch IS set, the standard `OTEL_EXPORTER_OTLP_*` values are
/// completely ignored here so the internally-authed firehose never lands
/// at an external collector.
pub fn resolve_otlp_traces_endpoint(&self) -> String {
if let Some(full) = blank_as_unset(&self.grok_internal_otlp_traces_endpoint) {
return full.trim_end_matches('/').to_string();
}
if !self.external_otel_master_switch
&& let Some(legacy) = self.legacy_internal_otlp_traces_endpoint()
{
tracing::warn!(
"Repointing the internal trace pipeline via OTEL_EXPORTER_OTLP_ENDPOINT / \
OTEL_EXPORTER_OTLP_TRACES_ENDPOINT is deprecated; use \
KIGI_INTERNAL_OTLP_TRACES_ENDPOINT instead — the standard OTEL_* vars will \
route the external OTEL stream only in a future release"
);
return legacy;
}
format!("{}/traces", self.proxy_url().trim_end_matches('/'))
}
/// Legacy (standard-OTEL-var) internal traces endpoint, if any:
/// `otel_exporter_otlp_traces_endpoint` verbatim, else
/// `otel_exporter_otlp_endpoint` + `/v1/traces`. Ignores the master switch.
fn legacy_internal_otlp_traces_endpoint(&self) -> Option<String> {
if let Some(full) = blank_as_unset(&self.otel_exporter_otlp_traces_endpoint) {
return Some(full.trim_end_matches('/').to_string());
}
blank_as_unset(&self.otel_exporter_otlp_endpoint)
.map(|base| format!("{}/v1/traces", base.trim_end_matches('/')))
}
/// Extra headers for the INTERNAL export: `grok_internal_otlp_headers`
/// first; legacy fallback to `otel_exporter_otlp_headers` ONLY when the
/// external-OTEL master switch is unset (back-compat for existing users).
pub fn resolve_otlp_headers(&self) -> Vec<(String, String)> {
if let Some(headers) = blank_as_unset(&self.grok_internal_otlp_headers) {
return parse_otlp_header_list(&headers);
}
if !self.external_otel_master_switch {
return parse_otlp_header_list(
self.otel_exporter_otlp_headers.as_deref().unwrap_or(""),
);
}
Vec::new()
}
/// Whether the legacy fallback actually supplied the internal endpoint OR
/// internal headers from the standard `OTEL_EXPORTER_OTLP_*` vars — i.e.
/// the master switch is unset AND (`otel_exporter_otlp_traces_endpoint` /
/// `otel_exporter_otlp_endpoint` is non-blank for the endpoint, or
/// `otel_exporter_otlp_headers` is non-blank for headers) AND no
/// `grok_internal_otlp_*` override shadowed that half.
///
/// CONTRACT: this flag is passed to the external OTEL stream's init, which
/// MUST refuse to activate when it is true — the same standard vars cannot
/// feed both pipelines (no-double-send invariant, enforced in code).
pub fn internal_otlp_consumed_standard_vars(&self) -> bool {
if self.external_otel_master_switch {
return false;
}
let endpoint_consumed = blank_as_unset(&self.grok_internal_otlp_traces_endpoint).is_none()
&& self.legacy_internal_otlp_traces_endpoint().is_some();
let headers_consumed = blank_as_unset(&self.grok_internal_otlp_headers).is_none()
&& blank_as_unset(&self.otel_exporter_otlp_headers).is_some();
endpoint_consumed || headers_consumed
}
/// Trace export enabled unless `OTEL_TRACES_EXPORTER=none`. Deliberately
/// still honored by the internal pipeline even with `KIGI_EXTERNAL_OTEL`
/// set: disabling internal span export is the safe direction.
pub fn resolve_traces_export_enabled(&self) -> bool {
!matches!(
self.otel_traces_exporter.as_deref().map(str::trim),
Some("none")
)
}
/// `OTEL_BSP_SCHEDULE_DELAY` / `OTEL_TRACES_EXPORT_INTERVAL` — tuning-only,
/// deliberately shared between the internal and external pipelines.
pub fn resolve_otlp_export_interval(&self) -> Option<std::time::Duration> {
self.otel_traces_export_interval
.map(std::time::Duration::from_millis)
}
/// `OTEL_EXPORTER_OTLP_TIMEOUT` — tuning-only, deliberately shared between
/// the internal and external pipelines.
pub fn resolve_otlp_timeout(&self) -> Option<std::time::Duration> {
self.otel_exporter_otlp_timeout
.map(std::time::Duration::from_millis)
}
/// `models_list_url` > `{models_base_url}/models` > `{proxy_base_url}/models`.
pub fn resolve_models_list_url(&self) -> String {
if let Some(ref url) = self.models_list_url {
@@ -407,26 +252,14 @@ impl Default for EndpointsConfig {
fn default() -> Self {
Self {
coding_api_base_url: std::env::var("KIGI_CODE_BASE_URL").ok(),
xai_api_base_url: std::env::var("KIGI_XAI_API_BASE_URL")
.unwrap_or_else(|_| XAI_API_BASE_URL_DEFAULT.to_owned()),
api_base_url: std::env::var("KIGI_API_BASE_URL")
.unwrap_or_else(|_| API_BASE_URL_DEFAULT.to_owned()),
alpha_test_key: None,
models_base_url: env_string("KIGI_MODELS_BASE_URL"),
models_list_url: env_string("KIGI_MODELS_LIST_URL"),
feedback_base_url: env_string("KIGI_FEEDBACK_BASE_URL"),
deployment_key: env_string("KIGI_DEPLOYMENT_KEY"),
managed_config_url: env_string("KIGI_MANAGED_CONFIG_URL"),
otel_exporter_otlp_endpoint: env_string("OTEL_EXPORTER_OTLP_ENDPOINT"),
otel_exporter_otlp_traces_endpoint: env_string("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"),
otel_exporter_otlp_headers: env_string("OTEL_EXPORTER_OTLP_HEADERS"),
grok_internal_otlp_traces_endpoint: env_string("KIGI_INTERNAL_OTLP_TRACES_ENDPOINT"),
grok_internal_otlp_headers: env_string("KIGI_INTERNAL_OTLP_HEADERS"),
external_otel_master_switch: external_otel_master_switch_resolved(),
otel_traces_exporter: env_string("OTEL_TRACES_EXPORTER"),
otel_traces_export_interval: env_string("OTEL_BSP_SCHEDULE_DELAY")
.or_else(|| env_string("OTEL_TRACES_EXPORT_INTERVAL"))
.and_then(|s| s.parse().ok()),
otel_exporter_otlp_timeout: env_string("OTEL_EXPORTER_OTLP_TIMEOUT")
.and_then(|s| s.parse().ok()),
management_api_key: None,
gcs_service_account_key: None,
}
@@ -1005,30 +838,6 @@ pub struct RemoteConfig {
#[serde(skip_serializing_if = "Option::is_none")]
pub secret: Option<String>,
}
/// `[hub]` section from config.toml.
///
/// Optional default Computer Hub URL for **workspace provider** exposure
/// (`grok workspace` / leader `with_default_hub_url`). Does **not** enable
/// agent-side harness/client connections or alter local session behavior.
///
/// ```toml
/// [hub]
/// url = "wss://hub.x.ai/ws"
/// ```
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(default)]
pub struct HubConfig {
/// Hub WebSocket URL (`ws://` or `wss://`) used as the leader default for
/// `grok workspace start` when the CLI does not pass `--hub-url`.
#[serde(skip_serializing_if = "Option::is_none")]
pub url: Option<String>,
}
impl HubConfig {
/// Whether a non-empty hub URL is configured (workspace default only).
pub fn is_enabled(&self) -> bool {
self.url.as_ref().is_some_and(|u| !u.trim().is_empty())
}
}
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(default)]
pub struct WorktreePoolConfig {
@@ -1250,9 +1059,6 @@ pub struct Config {
pub harness: HarnessConfig,
#[serde(default, skip_serializing)]
pub remote: RemoteConfig,
/// Computer Hub configuration (`[hub]` in config.toml).
#[serde(default, skip_serializing)]
pub hub: HubConfig,
#[serde(default, skip_serializing)]
pub worktree_pool: WorktreePoolConfig,
#[serde(default, skip_serializing)]
@@ -1632,7 +1438,6 @@ impl Default for Config {
platforms: PlatformsConfig::default(),
harness: HarnessConfig::default(),
remote: RemoteConfig::default(),
hub: HubConfig::default(),
worktree_pool: WorktreePoolConfig::default(),
sandbox: SandboxSettingsConfig::default(),
mcp_servers: std::collections::HashMap::new(),
@@ -2703,41 +2508,6 @@ pub(crate) fn read_requirements_toml() -> Option<toml::Value> {
pub fn deployment_id_from_key(key: &str) -> String {
uuid::Uuid::new_v5(&uuid::Uuid::NAMESPACE_OID, key.as_bytes()).to_string()
}
/// Resolve the external-OTEL master switch exactly the way the external
/// stream's activation does: **requirement pin > `KIGI_EXTERNAL_OTEL` env >
/// `[telemetry].otel_enabled` config layer (managed config included) > off**.
///
/// The internal trace pipeline keys its "ignore `OTEL_EXPORTER_OTLP_*`"
/// behavior off this value ([`EndpointsConfig::external_otel_master_switch`]),
/// so an org enable distributed via managed config / requirements (no env
/// var) flips **both** sides together. A desync here would leave the
/// internally-authed firehose honoring legacy `OTEL_*` repointing while
/// `internal_pipeline_consumed_otel_vars` simultaneously blocks the external
/// stream — exactly the split this design forbids.
pub(crate) fn external_otel_master_switch_resolved() -> bool {
external_otel_master_switch_from(
kigi_config::load_merged_requirements().as_ref(),
env_bool("KIGI_EXTERNAL_OTEL"),
crate::config::load_effective_config().ok().as_ref(),
)
}
/// Testable core of [`external_otel_master_switch_resolved`].
pub(crate) fn external_otel_master_switch_from(
requirements: Option<&toml::Value>,
env_switch: Option<bool>,
effective_config: Option<&toml::Value>,
) -> bool {
let table_enabled = |v: Option<&toml::Value>| -> Option<bool> {
v?.get("telemetry")?.get("otel_enabled")?.as_bool()
};
if let Some(pinned) = table_enabled(requirements) {
return pinned;
}
if let Some(env) = env_switch {
return env;
}
table_enabled(effective_config).unwrap_or(false)
}
/// Seed free-function remote caches after writing `Config.remote_settings`.
pub fn apply_remote_settings_side_effects(settings: Option<&crate::util::config::RemoteSettings>) {
crate::util::config::cache_remote_mcp_startup_timeout_secs(
@@ -4983,7 +4753,7 @@ reasoning_effort = "low"
auth_scheme: AuthScheme::Bearer,
};
assert_eq!(
api_key_creds.base_url, endpoints.xai_api_base_url,
api_key_creds.base_url, endpoints.api_base_url,
"{model_id}: ExternalApiKey must route to api.x.ai"
);
}
@@ -6890,24 +6660,18 @@ reasoning_effort = "low"
for k in [
"KIGI_CODE_BASE_URL",
kigi_env::CODE_BASE_URL_ENV,
"KIGI_XAI_API_BASE_URL",
"KIGI_API_BASE_URL",
"KIGI_FEEDBACK_BASE_URL",
"KIGI_TRACE_UPLOAD_URL",
"KIGI_MANAGED_CONFIG_URL",
"KIGI_MODELS_BASE_URL",
"KIGI_MODELS_LIST_URL",
"OTEL_EXPORTER_OTLP_ENDPOINT",
"OTEL_EXPORTER_OTLP_TRACES_ENDPOINT",
"OTEL_EXPORTER_OTLP_HEADERS",
"KIGI_INTERNAL_OTLP_TRACES_ENDPOINT",
"KIGI_INTERNAL_OTLP_HEADERS",
"KIGI_EXTERNAL_OTEL",
] {
unsafe { std::env::remove_var(k) };
}
}
/// INVARIANT: auxiliary-service resolvers resolve to the cli-chat-proxy, never
/// `xai_api_base_url` — overriding ONLY inference keeps every aux endpoint on
/// `api_base_url` — overriding ONLY inference keeps every aux endpoint on
/// the proxy; explicit per-service overrides win verbatim.
#[test]
#[serial]
@@ -6915,7 +6679,7 @@ reasoning_effort = "low"
unset_endpoint_env_vars();
let inference = "https://inference.acme-corp.example/xai/v1";
let cfg = EndpointsConfig {
xai_api_base_url: inference.to_string(),
api_base_url: inference.to_string(),
coding_api_base_url: None,
..Default::default()
};
@@ -6928,11 +6692,7 @@ reasoning_effort = "low"
format!("{proxy}/deployment/config")
);
assert_eq!(cfg.resolve_feedback_base_url(), proxy);
assert_eq!(
cfg.resolve_otlp_traces_endpoint(),
format!("{proxy}/traces")
);
assert_eq!(cfg.xai_api_base_url, inference);
assert_eq!(cfg.api_base_url, inference);
let overridden = EndpointsConfig {
coding_api_base_url: Some("https://proxy.enterprise.example/v1".to_string()),
managed_config_url: Some(
@@ -6945,10 +6705,6 @@ reasoning_effort = "low"
overridden.proxy_url(),
"https://proxy.enterprise.example/v1"
);
assert_eq!(
overridden.resolve_otlp_traces_endpoint(),
"https://proxy.enterprise.example/v1/traces"
);
assert_eq!(
overridden.resolve_managed_config_url(),
"https://control.enterprise.example/deployment/config"
@@ -6958,7 +6714,7 @@ reasoning_effort = "low"
"https://feedback.enterprise.example"
);
}
/// REGRESSION: the managed-config URL never follows `xai_api_base_url`
/// REGRESSION: the managed-config URL never follows `api_base_url`
/// through the full loader `Config::new_from_toml_cfg` — a distinct construction
/// path from `from_config_value`, so the deployment key never reaches the
/// inference host on either.
@@ -6969,7 +6725,7 @@ reasoning_effort = "low"
let cfg = Config::new_from_toml_cfg(
&toml::from_str(
r#"[endpoints]
xai_api_base_url = "https://inference.acme-corp.example/xai/v1""#,
api_base_url = "https://inference.acme-corp.example/xai/v1""#,
)
.unwrap(),
)
@@ -8467,311 +8223,6 @@ agent_type = "cursor"
"exactly the typo'd key must be flagged"
);
}
#[test]
fn otlp_traces_endpoint_precedence() {
let proxy = "https://inference.acme.com/v1".to_string();
let derived = EndpointsConfig {
coding_api_base_url: Some(proxy.clone()),
..Default::default()
};
assert_eq!(
derived.resolve_otlp_traces_endpoint(),
"https://inference.acme.com/v1/traces"
);
let base = EndpointsConfig {
coding_api_base_url: Some(proxy.clone()),
otel_exporter_otlp_endpoint: Some("https://otel.acme.com".to_string()),
..Default::default()
};
assert_eq!(
base.resolve_otlp_traces_endpoint(),
"https://otel.acme.com/v1/traces"
);
let full = EndpointsConfig {
coding_api_base_url: Some(proxy),
otel_exporter_otlp_endpoint: Some("https://ignored.example".to_string()),
otel_exporter_otlp_traces_endpoint: Some("https://otel.acme.com/v1/traces".to_string()),
..Default::default()
};
assert_eq!(
full.resolve_otlp_traces_endpoint(),
"https://otel.acme.com/v1/traces"
);
}
#[test]
fn otlp_headers_parse() {
let cfg = EndpointsConfig {
otel_exporter_otlp_headers: Some("a=1, b = 2 ,=skip,c=".to_string()),
..Default::default()
};
assert_eq!(
cfg.resolve_otlp_headers(),
vec![
("a".to_string(), "1".to_string()),
("b".to_string(), "2".to_string()),
("c".to_string(), String::new()),
]
);
}
/// Base config for the internal-OTLP tests: pinned proxy, every OTLP knob
/// explicitly unset so ambient env (via `Default`) can't leak in.
fn internal_otlp_test_config() -> EndpointsConfig {
EndpointsConfig {
coding_api_base_url: Some("https://proxy.example/v1".to_string()),
otel_exporter_otlp_endpoint: None,
otel_exporter_otlp_traces_endpoint: None,
otel_exporter_otlp_headers: None,
grok_internal_otlp_traces_endpoint: None,
grok_internal_otlp_headers: None,
external_otel_master_switch: false,
..Default::default()
}
}
/// `grok_internal_otlp_traces_endpoint` wins over the legacy `OTEL_*`
/// fields regardless of the master switch.
#[test]
fn internal_otlp_endpoint_grok_internal_wins_regardless_of_switch() {
for switch in [false, true] {
let cfg = EndpointsConfig {
grok_internal_otlp_traces_endpoint: Some(
"https://internal.example/traces/".to_string(),
),
otel_exporter_otlp_traces_endpoint: Some(
"https://legacy.example/v1/traces".to_string(),
),
otel_exporter_otlp_endpoint: Some("https://legacy-base.example".to_string()),
external_otel_master_switch: switch,
..internal_otlp_test_config()
};
assert_eq!(
cfg.resolve_otlp_traces_endpoint(),
"https://internal.example/traces",
"switch={switch}: KIGI_INTERNAL_OTLP_TRACES_ENDPOINT must win verbatim (trailing / trimmed)"
);
}
}
/// Master switch unset → legacy fallback preserved (back-compat).
#[test]
fn internal_otlp_endpoint_legacy_fallback_when_switch_unset() {
let traces = EndpointsConfig {
otel_exporter_otlp_traces_endpoint: Some(
"https://legacy.example/v1/traces".to_string(),
),
..internal_otlp_test_config()
};
assert_eq!(
traces.resolve_otlp_traces_endpoint(),
"https://legacy.example/v1/traces"
);
let base = EndpointsConfig {
otel_exporter_otlp_endpoint: Some("https://legacy-base.example/".to_string()),
..internal_otlp_test_config()
};
assert_eq!(
base.resolve_otlp_traces_endpoint(),
"https://legacy-base.example/v1/traces"
);
}
/// Master switch SET → legacy `OTEL_*` endpoint/headers are completely
/// ignored by the internal pipeline (the external stream owns them); the
/// internal pipeline falls back to the proxy default and
/// `internal_otlp_consumed_standard_vars()` is false.
#[test]
fn internal_otlp_ignores_legacy_vars_when_switch_set() {
let cfg = EndpointsConfig {
otel_exporter_otlp_traces_endpoint: Some(
"https://admin-collector.example/v1/traces".to_string(),
),
otel_exporter_otlp_endpoint: Some("https://admin-collector.example".to_string()),
otel_exporter_otlp_headers: Some("authorization=Bearer admin".to_string()),
external_otel_master_switch: true,
..internal_otlp_test_config()
};
assert_eq!(
cfg.resolve_otlp_traces_endpoint(),
"https://proxy.example/v1/traces",
"internal firehose must never follow OTEL_* to the external collector"
);
assert_eq!(cfg.resolve_otlp_headers(), Vec::<(String, String)>::new());
assert!(!cfg.internal_otlp_consumed_standard_vars());
}
/// `internal_otlp_consumed_standard_vars()` truth table.
#[test]
fn internal_otlp_consumed_standard_vars_cases() {
struct Case {
switch: bool,
legacy_traces_ep: bool,
legacy_base_ep: bool,
legacy_headers: bool,
internal_ep: bool,
internal_headers: bool,
expected: bool,
why: &'static str,
}
let unset = Case {
switch: false,
legacy_traces_ep: false,
legacy_base_ep: false,
legacy_headers: false,
internal_ep: false,
internal_headers: false,
expected: false,
why: "nothing set",
};
let cases = [
Case { ..unset },
Case {
legacy_traces_ep: true,
expected: true,
why: "legacy traces endpoint consumed",
..unset
},
Case {
legacy_base_ep: true,
expected: true,
why: "legacy base endpoint consumed",
..unset
},
Case {
legacy_headers: true,
expected: true,
why: "legacy headers consumed",
..unset
},
Case {
legacy_traces_ep: true,
internal_ep: true,
expected: false,
why: "internal endpoint shadows legacy",
..unset
},
Case {
legacy_headers: true,
internal_headers: true,
expected: false,
why: "internal headers shadow legacy",
..unset
},
Case {
legacy_traces_ep: true,
legacy_headers: true,
internal_ep: true,
expected: true,
why: "endpoint shadowed but legacy headers still consumed (headers half)",
..unset
},
Case {
switch: true,
legacy_traces_ep: true,
legacy_base_ep: true,
legacy_headers: true,
expected: false,
why: "switch set: legacy vars ignored",
..unset
},
];
for case in cases {
let cfg = EndpointsConfig {
external_otel_master_switch: case.switch,
otel_exporter_otlp_traces_endpoint: case
.legacy_traces_ep
.then(|| "https://legacy.example/v1/traces".to_string()),
otel_exporter_otlp_endpoint: case
.legacy_base_ep
.then(|| "https://legacy-base.example".to_string()),
otel_exporter_otlp_headers: case.legacy_headers.then(|| "k=v".to_string()),
grok_internal_otlp_traces_endpoint: case
.internal_ep
.then(|| "https://internal.example/traces".to_string()),
grok_internal_otlp_headers: case.internal_headers.then(|| "ik=iv".to_string()),
..internal_otlp_test_config()
};
assert_eq!(
cfg.internal_otlp_consumed_standard_vars(),
case.expected,
"case: {}",
case.why
);
}
}
/// Headers precedence: `grok_internal_otlp_headers` wins; legacy
/// `otel_exporter_otlp_headers` only when the master switch is unset.
#[test]
fn internal_otlp_headers_precedence() {
for switch in [false, true] {
let cfg = EndpointsConfig {
grok_internal_otlp_headers: Some("x-debug=1".to_string()),
otel_exporter_otlp_headers: Some("legacy=1".to_string()),
external_otel_master_switch: switch,
..internal_otlp_test_config()
};
assert_eq!(
cfg.resolve_otlp_headers(),
vec![("x-debug".to_string(), "1".to_string())],
"switch={switch}"
);
}
let legacy = EndpointsConfig {
otel_exporter_otlp_headers: Some("legacy=1".to_string()),
..internal_otlp_test_config()
};
assert_eq!(
legacy.resolve_otlp_headers(),
vec![("legacy".to_string(), "1".to_string())]
);
}
/// Regression: an org enable via `[telemetry].otel_enabled`
/// (managed config / requirements — no `KIGI_EXTERNAL_OTEL` env var) must
/// flip the master switch the *internal* pipeline keys off, so legacy
/// `OTEL_EXPORTER_OTLP_*` repointing shuts off in lockstep with the
/// external stream activating. A desync would point the internally-authed
/// firehose at the customer collector while
/// `internal_pipeline_consumed_otel_vars` blocks the external stream.
#[test]
fn external_otel_master_switch_resolves_from_all_layers() {
let enabled_table: toml::Value =
toml::from_str("[telemetry]\notel_enabled = true").unwrap();
let disabled_table: toml::Value =
toml::from_str("[telemetry]\notel_enabled = false").unwrap();
assert!(external_otel_master_switch_from(
None,
None,
Some(&enabled_table)
));
assert!(!external_otel_master_switch_from(None, None, None));
assert!(!external_otel_master_switch_from(
None,
Some(false),
Some(&enabled_table)
));
assert!(external_otel_master_switch_from(
None,
Some(true),
Some(&disabled_table)
));
assert!(!external_otel_master_switch_from(
Some(&disabled_table),
Some(true),
Some(&enabled_table)
));
assert!(external_otel_master_switch_from(
Some(&enabled_table),
Some(false),
None
));
let cfg = EndpointsConfig {
otel_exporter_otlp_traces_endpoint: Some(
"https://collector.corp:4318/v1/traces".into(),
),
external_otel_master_switch: true,
..internal_otlp_test_config()
};
assert!(!cfg.internal_otlp_consumed_standard_vars());
assert!(
!cfg.resolve_otlp_traces_endpoint()
.contains("collector.corp")
);
}
fn empty_config() -> toml::Value {
toml::Value::Table(toml::map::Map::new())
}
@@ -9820,26 +9271,6 @@ default = "grok-4.5"
}
}
#[test]
fn hub_config_default_has_no_url() {
assert!(HubConfig::default().url.is_none());
assert!(!HubConfig::default().is_enabled());
}
#[test]
fn hub_config_is_enabled_only_for_nonempty_url() {
assert!(
HubConfig {
url: Some("wss://hub.example/ws".into()),
}
.is_enabled()
);
assert!(
!HubConfig {
url: Some(" ".into()),
}
.is_enabled()
);
}
#[test]
fn resolve_model_list_prunes_bundled_entries_not_in_prefetch() {
let cfg = Config::default();
let mut defs = default_model_entries(&EndpointsConfig::default());
crates/codegen/kigi-shell/src/agent/mod.rs
-1
View File
@@ -12,7 +12,6 @@ pub mod init;
pub mod models;
pub(crate) mod models_fetch;
pub mod mvp_agent;
pub(crate) mod proxy;
pub(crate) mod restore_code;
pub mod roster;
pub mod server;
crates/codegen/kigi-shell/src/agent/models.rs
+7 -7
View File
@@ -1323,7 +1323,7 @@ const CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(300);
struct ModelsCache {
fetched_at: DateTime<Utc>,
#[serde(default, skip_serializing_if = "Option::is_none")]
grok_version: Option<String>,
kigi_version: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
auth_method: Option<CacheAuthMethod>,
/// Models-list URL this catalog was fetched from
@@ -1418,7 +1418,7 @@ impl ModelsCacheManager {
) -> Option<ModelsCache> {
let data = std::fs::read(&self.path).ok()?;
let cache: ModelsCache = serde_json::from_slice(&data).ok()?;
if cache.grok_version.as_deref() != Some(kigi_version::VERSION) {
if cache.kigi_version.as_deref() != Some(kigi_version::VERSION) {
tracing::debug!("models cache version mismatch");
return None;
}
@@ -1447,7 +1447,7 @@ impl ModelsCacheManager {
) {
let cache = ModelsCache {
fetched_at: Utc::now(),
grok_version: Some(kigi_version::VERSION.to_string()),
kigi_version: Some(kigi_version::VERSION.to_string()),
auth_method: Some(auth_method),
origin: Some(origin.to_string()),
etag: etag.map(|s| s.to_string()),
@@ -3064,7 +3064,7 @@ mod tests {
let auth_method = mgr.inner.fetch_auth.read().cache_auth_method();
let stale = ModelsCache {
fetched_at: Utc::now() - ChronoDuration::seconds(3600),
grok_version: Some(kigi_version::VERSION.to_string()),
kigi_version: Some(kigi_version::VERSION.to_string()),
auth_method: Some(auth_method),
origin: Some(mgr.cache_origin()),
etag: Some("etag-stale".into()),
@@ -3140,7 +3140,7 @@ mod tests {
let auth_method = mgr.inner.fetch_auth.read().cache_auth_method();
let legacy = ModelsCache {
fetched_at: Utc::now(),
grok_version: Some(kigi_version::VERSION.to_string()),
kigi_version: Some(kigi_version::VERSION.to_string()),
auth_method: Some(auth_method),
origin: None,
etag: Some("etag-legacy".into()),
@@ -4100,7 +4100,7 @@ mod tests {
let cache = ModelsCacheManager::new();
let stale = ModelsCache {
fetched_at: Utc::now() - ChronoDuration::seconds(86_400),
grok_version: Some(kigi_version::VERSION.to_string()),
kigi_version: Some(kigi_version::VERSION.to_string()),
auth_method: Some(CacheAuthMethod::Platforms),
origin: Some(origin),
etag: None,
@@ -4187,7 +4187,7 @@ mod tests {
let cache = ModelsCacheManager::new();
cache.atomic_write(&ModelsCache {
fetched_at: Utc::now() - ChronoDuration::seconds(86_400),
grok_version: Some(kigi_version::VERSION.to_string()),
kigi_version: Some(kigi_version::VERSION.to_string()),
auth_method: Some(CacheAuthMethod::Platforms),
origin: Some(with_key_origin),
etag: None,
crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs
+3 -3
View File
@@ -762,7 +762,7 @@ impl MvpAgent {
}
/// Build image generation config.
///
/// Both BYOK and session (OAuth) users go direct to `xai_api_base_url`.
/// Both BYOK and session (OAuth) users go direct to `api_base_url`.
/// `sampling_config.api_key` carries the OAuth bearer for session users (the
/// `api_key_provider` refreshes it per request), so IC authenticates and
/// meters Imagine usage per-user.
@@ -775,7 +775,7 @@ impl MvpAgent {
return ImageGenConfig::Disabled;
};
let cfg = self.cfg.borrow();
let base_url = cfg.endpoints.xai_api_base_url.clone();
let base_url = cfg.endpoints.api_base_url.clone();
let version = cfg
.client_version
.clone()
@@ -818,7 +818,7 @@ impl MvpAgent {
tracing::info!("video_gen disabled by tools.disable_zdr_incompatible_tools");
return VideoGenConfig::Disabled;
}
let base_url = cfg.endpoints.xai_api_base_url.clone();
let base_url = cfg.endpoints.api_base_url.clone();
let version = cfg
.client_version
.clone()
crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs
-2
View File
@@ -664,8 +664,6 @@ pub struct MvpAgent {
plugin_registry_initialized: std::cell::Cell<bool>,
persona_io_summaries: Vec<String>,
/// Local workspace ops, built lazily via [`Self::ensure_local_workspace_ops`].
/// The agent never opens Computer Hub as a harness/client; remote cloud
/// sandboxes are gateway-owned (`gateway_bridge` / `computer_sessions`).
workspace_ops: RefCell<Option<kigi_workspace::WorkspaceOps>>,
/// Sessions opened with `require_gateway` / chat light-frontend (K13).
/// Prompt-time guard consults this when the bridge map entry is missing,
crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs
-34
View File
@@ -652,7 +652,6 @@ async fn file_toolset_override_e2e_to_finalized_toolset() {
video_gen_config: kigi_tools::implementations::grok_build::video_gen::VideoGenConfig::default(),
app_builder_deployer_config: kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig::default(),
api_key_provider: None,
auth_provider: None,
attribution_callback: None,
system_reminder_tag: kigi_tools::reminders::DEFAULT_REMINDER_TAG,
};
@@ -3183,7 +3182,6 @@ fn interactive_trust_prompt_reprompts_after_untrust() {
}
mod direct_hub_cloud_removed {
use super::super::{DIRECT_HUB_CLOUD_REMOVED_MSG, reject_direct_hub_cloud_meta};
use crate::agent::config::HubConfig;
fn assert_direct_hub_error(err: agent_client_protocol::Error) {
assert_eq!(
err.data.as_ref(),
@@ -3238,38 +3236,6 @@ mod direct_hub_cloud_removed {
.is_ok()
);
}
#[test]
fn hub_url_gating_matrix() {
let with_url = HubConfig {
url: Some("wss://hub.example/ws".into()),
};
let without_url = HubConfig { url: None };
let blank = HubConfig {
url: Some(" ".into()),
};
assert!(with_url.is_enabled());
assert!(!without_url.is_enabled());
assert!(!blank.is_enabled());
}
#[test]
fn hub_config_is_url_only_workspace_default() {
let json = serde_json::to_value(HubConfig {
url: Some("wss://hub.example/ws".into()),
})
.expect("serialize");
let obj = json.as_object().expect("object");
assert_eq!(
obj.keys().collect::<Vec<_>>(),
vec!["url"],
"HubConfig must only serialize url (no proxy-mode fields)"
);
let from_legacy: HubConfig = serde_json::from_value(serde_json::json!(
{ "url" : "wss://hub.example/ws", "workspace_mode" : "remote",
"send_turn_hooks" : false, }
))
.expect("ignore unknown fields");
assert_eq!(from_legacy.url.as_deref(), Some("wss://hub.example/ws"));
}
}
mod soft_default_settings_emit {
use super::*;
crates/codegen/kigi-shell/src/agent/mvp_agent/tests/subagent_spawn_context_tests.rs
+1
View File
@@ -39,6 +39,7 @@ async fn subagent_spawn_context_inherits_parent_permission_handle() {
Vec::new(),
false,
None,
true,
);
let mut handle = make_test_handle("test-model", false, None);
crates/codegen/kigi-shell/src/agent/proxy.rs
-619
View File
@@ -1,619 +0,0 @@
//! HTTP CONNECT proxy support for WebSocket connections.
//!
//! When running behind a corporate egress proxy,
//! `tokio-tungstenite`'s `connect_async` cannot reach external
//! hosts directly because it does not read the standard `HTTPS_PROXY` /
//! `HTTP_PROXY` environment variables.
//!
//! This module provides:
//! - [`resolve_proxy_for_host`]: reads proxy env vars and `NO_PROXY`, returning
//! the proxy URL to use for a given target host (or `None` for direct).
//! - [`connect_via_proxy`]: opens a TCP connection to the proxy, sends an HTTP
//! CONNECT request to create a tunnel, wraps the result in TLS, and returns a
//! stream suitable for `tokio_tungstenite::client_async`.
use std::sync::{Arc, OnceLock};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::TcpStream;
use tokio_tungstenite::MaybeTlsStream;
use tracing::debug;
// ---------------------------------------------------------------------------
// Environment-variable resolution
// ---------------------------------------------------------------------------
/// Read proxy configuration from the environment and decide whether `target_host`
/// should be connected through a proxy.
///
/// Resolution order (matches `curl` / `reqwest` behaviour):
/// 1. If `NO_PROXY` contains `target_host` (or a matching domain suffix / CIDR),
/// return `None`.
/// 2. If `HTTPS_PROXY` (or `https_proxy`) is set, return its value.
/// 3. If `HTTP_PROXY` (or `http_proxy`) is set, return its value.
/// 4. Otherwise return `None`.
pub fn resolve_proxy_for_host(target_host: &str) -> Option<String> {
resolve_proxy_for_host_with(target_host, |key| std::env::var(key))
}
/// Testable inner implementation that accepts a custom env-var reader.
fn resolve_proxy_for_host_with<F>(target_host: &str, env: F) -> Option<String>
where
F: for<'a> Fn(&'a str) -> Result<String, std::env::VarError>,
{
// Check NO_PROXY / no_proxy.
let no_proxy = env("NO_PROXY")
.or_else(|_| env("no_proxy"))
.unwrap_or_default();
if is_host_bypassed(target_host, &no_proxy) {
return None;
}
// HTTPS_PROXY takes precedence (our target is always wss://).
if let Ok(url) = env("HTTPS_PROXY").or_else(|_| env("https_proxy")) {
let url = url.trim().to_string();
if !url.is_empty() {
return Some(url);
}
}
// Fall back to HTTP_PROXY.
if let Ok(url) = env("HTTP_PROXY").or_else(|_| env("http_proxy")) {
let url = url.trim().to_string();
if !url.is_empty() {
return Some(url);
}
}
None
}
/// Check whether `host` is in the `no_proxy` list.
///
/// The `no_proxy` value is a comma-separated list of hostnames, domain
/// suffixes (with or without a leading dot), IP addresses, or CIDR ranges.
/// The special value `*` matches everything.
fn is_host_bypassed(host: &str, no_proxy: &str) -> bool {
let host_lower = host.to_ascii_lowercase();
for entry in no_proxy.split(',') {
let entry = entry.trim().to_ascii_lowercase();
if entry.is_empty() {
continue;
}
// Wildcard — bypass all hosts.
if entry == "*" {
return true;
}
// Exact match.
if host_lower == entry {
return true;
}
// Domain suffix match: ".example.com" matches "foo.example.com".
// Also handle the common convention of omitting the leading dot:
// "example.com" in NO_PROXY should match "sub.example.com".
let matches_suffix = if entry.starts_with('.') {
host_lower.ends_with(entry.as_str())
} else {
host_lower.len() > entry.len()
&& host_lower.ends_with(entry.as_str())
&& host_lower.as_bytes()[host_lower.len() - entry.len() - 1] == b'.'
};
if matches_suffix {
return true;
}
// CIDR / IP matching is intentionally omitted here — our target host
// is always a DNS name, not an IP literal. Keeping this simple avoids
// pulling in a CIDR parsing dependency.
}
false
}
// ---------------------------------------------------------------------------
// HTTP CONNECT tunnel
// ---------------------------------------------------------------------------
/// Establish a TLS-wrapped TCP stream through an HTTP CONNECT proxy.
///
/// Steps:
/// 1. Parse the proxy URL to get host + port.
/// 2. Open a TCP connection to the proxy and perform the CONNECT handshake.
/// 3. Wrap the tunnel in TLS (using rustls with native root certificates).
/// 4. Return the stream as `MaybeTlsStream<TcpStream>` so it is compatible
/// with `tokio_tungstenite::client_async`.
pub async fn connect_via_proxy(
proxy_url: &str,
target_host: &str,
target_port: u16,
) -> anyhow::Result<MaybeTlsStream<TcpStream>> {
let stream = open_connect_tunnel(proxy_url, target_host, target_port).await?;
let tls_stream = tls_wrap(stream, target_host).await?;
Ok(MaybeTlsStream::Rustls(tls_stream))
}
/// Open a raw TCP tunnel through an HTTP CONNECT proxy (no TLS).
///
/// 1. Parse the proxy URL to get host + port.
/// 2. Open a plain TCP connection to the proxy.
/// 3. Send `CONNECT target_host:target_port HTTP/1.1\r\n\r\n`.
/// 4. Read the proxy's response; expect `HTTP/1.x 200 …`.
/// 5. Return the raw `TcpStream` positioned after the CONNECT response.
async fn open_connect_tunnel(
proxy_url: &str,
target_host: &str,
target_port: u16,
) -> anyhow::Result<TcpStream> {
// 1. Parse proxy URL.
let (proxy_host, proxy_port) = parse_proxy_url(proxy_url)?;
// 2. TCP connect to proxy.
let proxy_addr = format!("{proxy_host}:{proxy_port}");
debug!(proxy_addr = %proxy_addr, "Opening TCP to proxy");
let stream = TcpStream::connect(&proxy_addr)
.await
.map_err(|e| anyhow::anyhow!("Failed to connect to proxy at {proxy_addr}: {e}"))?;
// 3. Send HTTP CONNECT.
let connect_req = format!(
"CONNECT {target_host}:{target_port} HTTP/1.1\r\n\
Host: {target_host}:{target_port}\r\n\
\r\n"
);
let (reader_half, mut writer_half) = stream.into_split();
writer_half.write_all(connect_req.as_bytes()).await?;
writer_half.flush().await?;
// 4. Read the status line from the proxy.
let mut reader = BufReader::new(reader_half);
let mut status_line = String::new();
reader.read_line(&mut status_line).await?;
debug!(status_line = %status_line.trim(), "Proxy CONNECT response");
if !status_line.starts_with("HTTP/1.1 200") && !status_line.starts_with("HTTP/1.0 200") {
anyhow::bail!("Proxy CONNECT failed: {}", status_line.trim());
}
// Consume remaining response headers (until empty line).
loop {
let mut line = String::new();
reader.read_line(&mut line).await?;
if line.trim().is_empty() {
break;
}
}
// 5. Assert the BufReader's internal buffer is empty before reuniting.
// BufReader::read_line may have read ahead into its buffer. If extra
// bytes were consumed beyond the HTTP headers (e.g., from a proxy that
// eagerly forwards data or coalesced TCP segments), dropping them would
// corrupt the subsequent TLS handshake.
let remaining = reader.buffer();
if !remaining.is_empty() {
anyhow::bail!(
"Proxy sent {} unexpected byte(s) after CONNECT response headers",
remaining.len()
);
}
// 6. Reunite the split halves back into a TcpStream.
let stream = reader.into_inner().reunite(writer_half)?;
Ok(stream)
}
/// Lazily-initialized TLS client configuration.
///
/// Loading native root certificates involves syscalls (reading `/etc/ssl/certs/`
/// or the macOS Keychain) and the cert store never changes at runtime. We build
/// the `ClientConfig` once and reuse it across all proxy connections / reconnects.
///
/// Stores `Ok(config)` on success or `Err(message)` if cert loading fails.
static TLS_CONFIG: OnceLock<Result<Arc<rustls::ClientConfig>, String>> = OnceLock::new();
/// Build (or return the cached) TLS client configuration.
fn get_tls_config() -> anyhow::Result<Arc<rustls::ClientConfig>> {
let result = TLS_CONFIG.get_or_init(|| {
let mut root_store = rustls::RootCertStore::empty();
let cert_result = rustls_native_certs::load_native_certs();
if cert_result.certs.is_empty() {
let errors: Vec<_> = cert_result.errors.iter().map(|e| e.to_string()).collect();
return Err(format!(
"No native root certificates found. Errors: {}",
if errors.is_empty() {
"(none)".to_string()
} else {
errors.join("; ")
}
));
}
for cert in cert_result.certs {
if let Err(e) = root_store.add(cert) {
tracing::warn!(error = %e, "Skipping unparseable native root certificate");
}
}
let config = rustls::ClientConfig::builder()
.with_root_certificates(root_store)
.with_no_client_auth();
Ok(Arc::new(config))
});
match result {
Ok(config) => Ok(config.clone()),
Err(msg) => anyhow::bail!("{msg}"),
}
}
/// Perform a TLS handshake over an existing TCP stream using rustls with
/// native root certificates (cached via [`TLS_CONFIG`]).
async fn tls_wrap(
stream: TcpStream,
server_name: &str,
) -> anyhow::Result<tokio_rustls::client::TlsStream<TcpStream>> {
let tls_config = get_tls_config()?;
let connector = tokio_rustls::TlsConnector::from(tls_config);
let dns_name = rustls::pki_types::ServerName::try_from(server_name.to_string())
.map_err(|e| anyhow::anyhow!("Invalid TLS server name '{server_name}': {e}"))?;
let tls_stream = connector
.connect(dns_name, stream)
.await
.map_err(|e| anyhow::anyhow!("TLS handshake through proxy failed: {e}"))?;
Ok(tls_stream)
}
/// Parse a proxy URL into (host, port).
///
/// Accepted formats:
/// - `http://host:port`
/// - `http://host` (defaults to port 80)
/// - `host:port`
fn parse_proxy_url(url: &str) -> anyhow::Result<(String, u16)> {
// Strip scheme if present.
let without_scheme = url
.strip_prefix("http://")
.or_else(|| url.strip_prefix("https://"))
.unwrap_or(url);
// Strip trailing path/slash.
let authority = without_scheme.split('/').next().unwrap_or(without_scheme);
if let Some((host, port_str)) = authority.rsplit_once(':') {
let port: u16 = port_str
.parse()
.map_err(|_| anyhow::anyhow!("Invalid proxy port in '{url}'"))?;
Ok((host.to_string(), port))
} else {
// No port — default to 80 for HTTP proxies.
Ok((authority.to_string(), 80))
}
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
// ===== parse_proxy_url =====
#[test]
fn test_parse_proxy_url_with_scheme_and_port() {
let (host, port) = parse_proxy_url("http://proxy.example.com:3140").unwrap();
assert_eq!(host, "proxy.example.com");
assert_eq!(port, 3140);
}
#[test]
fn test_parse_proxy_url_without_scheme() {
let (host, port) = parse_proxy_url("proxy.example.com:8080").unwrap();
assert_eq!(host, "proxy.example.com");
assert_eq!(port, 8080);
}
#[test]
fn test_parse_proxy_url_without_port() {
let (host, port) = parse_proxy_url("http://proxy.example.com").unwrap();
assert_eq!(host, "proxy.example.com");
assert_eq!(port, 80);
}
#[test]
fn test_parse_proxy_url_with_trailing_slash() {
let (host, port) = parse_proxy_url("http://proxy.example.com:3140/").unwrap();
assert_eq!(host, "proxy.example.com");
assert_eq!(port, 3140);
}
#[test]
fn test_parse_proxy_url_https_scheme() {
let (host, port) = parse_proxy_url("https://secure-proxy:443").unwrap();
assert_eq!(host, "secure-proxy");
assert_eq!(port, 443);
}
#[test]
fn test_parse_proxy_url_multi_label_host() {
let (host, port) =
parse_proxy_url("http://http-proxy.services.internal.example:3128").unwrap();
assert_eq!(host, "http-proxy.services.internal.example");
assert_eq!(port, 3128);
}
#[test]
fn test_parse_proxy_url_invalid_port() {
assert!(parse_proxy_url("http://proxy:notaport").is_err());
}
// ===== is_host_bypassed =====
#[test]
fn test_bypass_exact_match() {
assert!(is_host_bypassed("localhost", "localhost,127.0.0.1"));
}
#[test]
fn test_bypass_domain_suffix_with_dot() {
assert!(is_host_bypassed(
"api.corp.example",
"localhost,.corp.example"
));
}
#[test]
fn test_bypass_domain_suffix_without_dot() {
// Common convention: "example.com" in NO_PROXY matches "api.example.com".
assert!(is_host_bypassed("api.example.com", "localhost,example.com"));
}
#[test]
fn test_bypass_wildcard() {
assert!(is_host_bypassed("anything.example.com", "*"));
}
#[test]
fn test_no_bypass_when_not_listed() {
assert!(!is_host_bypassed(
"api.external.example",
"localhost,127.0.0.1,.corp.example,.internal.example"
));
}
#[test]
fn test_bypass_case_insensitive() {
assert!(is_host_bypassed("API.Corp.EXAMPLE", ".corp.example"));
}
#[test]
fn test_bypass_empty_no_proxy() {
assert!(!is_host_bypassed("api.external.example", ""));
}
#[test]
fn test_bypass_spaces_in_entries() {
assert!(is_host_bypassed(
"foo.example.com",
" localhost , .example.com , .other.com "
));
}
#[test]
fn test_bypass_cidr_not_matched_for_dns_names() {
// CIDR entries like 10.0.0.0/8 should not match DNS names.
assert!(!is_host_bypassed("api.external.example", "10.0.0.0/8"));
}
#[test]
fn test_bypass_combined_no_proxy_list() {
// A typical corporate NO_PROXY mixes loopback, private CIDRs, and domain suffixes.
let no_proxy = "localhost,127.0.0.1,10.0.0.0/8,.internal.example,.corp.example";
assert!(!is_host_bypassed("api.external.example", no_proxy));
assert!(is_host_bypassed("db.internal.example", no_proxy));
assert!(is_host_bypassed("git.corp.example", no_proxy));
assert!(is_host_bypassed("localhost", no_proxy));
}
// ===== resolve_proxy_for_host_with =====
#[test]
fn test_resolve_no_proxy_vars_set() {
let result = resolve_proxy_for_host_with("api.external.example", |_| {
Err(std::env::VarError::NotPresent)
});
assert_eq!(result, None);
}
#[test]
fn test_resolve_https_proxy_used() {
let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
"HTTPS_PROXY" => Ok("http://proxy.example.com:3128".to_string()),
"NO_PROXY" => Err(std::env::VarError::NotPresent),
_ => Err(std::env::VarError::NotPresent),
});
assert_eq!(result, Some("http://proxy.example.com:3128".to_string()));
}
#[test]
fn test_resolve_http_proxy_fallback() {
let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
"HTTP_PROXY" => Ok("http://proxy.example.com:8080".to_string()),
"NO_PROXY" => Err(std::env::VarError::NotPresent),
_ => Err(std::env::VarError::NotPresent),
});
assert_eq!(result, Some("http://proxy.example.com:8080".to_string()));
}
#[test]
fn test_resolve_no_proxy_bypasses() {
let result = resolve_proxy_for_host_with("api.corp.example", |key| match key {
"HTTPS_PROXY" => Ok("http://proxy.example.com:3128".to_string()),
"NO_PROXY" => Ok("localhost,.corp.example".to_string()),
_ => Err(std::env::VarError::NotPresent),
});
assert_eq!(result, None);
}
#[test]
fn test_resolve_https_proxy_takes_precedence() {
let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
"HTTPS_PROXY" => Ok("http://https-proxy.example.com:443".to_string()),
"HTTP_PROXY" => Ok("http://http-proxy.example.com:80".to_string()),
"NO_PROXY" => Err(std::env::VarError::NotPresent),
_ => Err(std::env::VarError::NotPresent),
});
assert_eq!(
result,
Some("http://https-proxy.example.com:443".to_string())
);
}
#[test]
fn test_resolve_lowercase_env_vars() {
let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
"https_proxy" => Ok("http://proxy.example.com:3128".to_string()),
"no_proxy" => Err(std::env::VarError::NotPresent),
_ => Err(std::env::VarError::NotPresent),
});
assert_eq!(result, Some("http://proxy.example.com:3128".to_string()));
}
#[test]
fn test_resolve_empty_proxy_ignored() {
let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
"HTTPS_PROXY" => Ok(" ".to_string()),
"HTTP_PROXY" => Ok("http://proxy.example.com:8080".to_string()),
_ => Err(std::env::VarError::NotPresent),
});
assert_eq!(result, Some("http://proxy.example.com:8080".to_string()));
}
#[test]
fn test_resolve_respects_no_proxy_when_proxy_set() {
let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
"HTTPS_PROXY" | "HTTP_PROXY" => Ok("http://proxy.example.com:3128".to_string()),
"NO_PROXY" => {
Ok("localhost,127.0.0.1,10.0.0.0/8,.internal.example,.corp.example".to_string())
}
_ => Err(std::env::VarError::NotPresent),
});
assert_eq!(result, Some("http://proxy.example.com:3128".to_string()));
}
// ===== HTTP CONNECT tunnel (integration-style) =====
/// Helper: spawn a mock HTTP CONNECT proxy that accepts one connection.
///
/// On receiving a CONNECT request, it validates the request format,
/// replies with `status_line`, and then echoes data (simulating a tunnel).
/// Returns the proxy's listen address.
async fn spawn_mock_proxy(status_line: &'static str) -> std::net::SocketAddr {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
// Read CONNECT request (read until \r\n\r\n).
let mut buf = vec![0u8; 4096];
let mut total = 0;
loop {
let n = stream.read(&mut buf[total..]).await.unwrap();
if n == 0 {
return;
}
total += n;
let so_far = std::str::from_utf8(&buf[..total]).unwrap_or("");
if so_far.contains("\r\n\r\n") {
break;
}
}
let request = std::str::from_utf8(&buf[..total]).unwrap().to_string();
assert!(
request.contains("CONNECT ") && request.contains(" HTTP/1.1"),
"Expected CONNECT request, got: {request}"
);
// Reply with the provided status line.
stream.write_all(status_line.as_bytes()).await.unwrap();
// Echo loop (simulates the transparent tunnel).
let mut echo_buf = [0u8; 1024];
loop {
let n = match stream.read(&mut echo_buf).await {
Ok(0) | Err(_) => break,
Ok(n) => n,
};
if stream.write_all(&echo_buf[..n]).await.is_err() {
break;
}
}
});
addr
}
/// Tests that `open_connect_tunnel` sends a correct CONNECT request,
/// parses the proxy's 200 response, and returns a usable tunnel stream.
#[tokio::test]
async fn test_open_connect_tunnel_success() {
let addr =
spawn_mock_proxy("HTTP/1.1 200 Connection Established\r\nServer: mock\r\n\r\n").await;
let proxy_url = format!("http://{addr}");
// Call the real function under test.
let mut stream = open_connect_tunnel(&proxy_url, "example.com", 443)
.await
.expect("tunnel should succeed");
// Verify the tunnel works by echoing data through it.
stream.write_all(b"hello tunnel").await.unwrap();
stream.flush().await.unwrap();
let mut response = vec![0u8; 12];
stream.read_exact(&mut response).await.unwrap();
assert_eq!(&response, b"hello tunnel");
}
/// Tests that `open_connect_tunnel` with a non-default port sends the
/// correct CONNECT target.
#[tokio::test]
async fn test_open_connect_tunnel_custom_port() {
let addr = spawn_mock_proxy("HTTP/1.1 200 OK\r\n\r\n").await;
let proxy_url = format!("http://{addr}");
let stream = open_connect_tunnel(&proxy_url, "internal.example.com", 8443).await;
assert!(stream.is_ok(), "tunnel should succeed for custom port");
}
/// Tests that `open_connect_tunnel` returns an error when the proxy
/// rejects the CONNECT request with a non-200 status.
#[tokio::test]
async fn test_open_connect_tunnel_proxy_rejects() {
let addr = spawn_mock_proxy("HTTP/1.1 403 Forbidden\r\n\r\n").await;
let proxy_url = format!("http://{addr}");
let result = open_connect_tunnel(&proxy_url, "blocked.example.com", 443).await;
assert!(result.is_err());
let err_msg = result.unwrap_err().to_string();
assert!(
err_msg.contains("403"),
"Error should mention 403: {err_msg}"
);
}
/// Tests that `open_connect_tunnel` returns an error when connecting
/// to a proxy that isn't listening.
#[tokio::test]
async fn test_open_connect_tunnel_proxy_unreachable() {
let result = open_connect_tunnel("http://127.0.0.1:1", "example.com", 443).await;
assert!(result.is_err());
let err_msg = result.unwrap_err().to_string();
assert!(
err_msg.contains("Failed to connect to proxy"),
"Error should mention proxy connection failure: {err_msg}"
);
}
}
crates/codegen/kigi-shell/src/config/mod.rs
+4 -4
View File
@@ -1059,11 +1059,11 @@ fn apply_requirements_inner(
enforce_str!("models", "default", config.models.default);
enforce_str!("cli", "channel", config.cli.channel);
enforce_str!("cli", "minimum_version", config.cli.minimum_version);
if let Some(val) = req_str(req, "endpoints", "xai_api_base_url")
&& config.endpoints.xai_api_base_url != val
if let Some(val) = req_str(req, "endpoints", "api_base_url")
&& config.endpoints.api_base_url != val
{
config.endpoints.xai_api_base_url = val.to_owned();
push("endpoints.xai_api_base_url", val.to_owned());
config.endpoints.api_base_url = val.to_owned();
push("endpoints.api_base_url", val.to_owned());
}
if let Some(val) = req_str(req, "endpoints", "coding_api_base_url")
&& config.endpoints.coding_api_base_url.as_deref() != Some(val)
crates/codegen/kigi-shell/src/config/tests.rs
+5 -5
View File
@@ -2490,7 +2490,7 @@ fn enterprise_two_file_merge_routes_deployment_key_to_proxy() {
let managed = toml::from_str(
r#"
[endpoints]
xai_api_base_url = "https://inference.acme-corp.example/xai/v1"
api_base_url = "https://inference.acme-corp.example/xai/v1"
coding_api_base_url = "https://cli-chat-proxy.kigi.com/v1"
[model.kigi-build]
@@ -2511,7 +2511,7 @@ telemetry = false
[endpoints]
deployment_key = "xai-token-ENTERPRISE"
xai_api_base_url = "https://inference.acme-corp.example/xai/v1"
api_base_url = "https://inference.acme-corp.example/xai/v1"
"#,
)
.unwrap();
@@ -2586,13 +2586,13 @@ fn config_layers_system_managed_lowest_priority() {
#[test]
fn apply_requirements_value_overrides_user_settings() {
let raw_config: toml::Value = toml::from_str(
"[cli]\nauto_update = true\nchannel = \"beta\"\n\n[features]\nfeedback = true\nlsp_tools = true\nweb_fetch = true\nwrite_file = true\n\n[ui]\nyolo = true\n\n[models]\ndefault = \"user-model\"\n\n[endpoints]\ncoding_api_base_url = \"https://user-proxy.example/v1\"\nxai_api_base_url = \"https://user-api.example/v1\"\nmodels_base_url = \"https://user-models.example/v1\"\nmodels_list_url = \"https://user-models.example/v1/models\"\n",
"[cli]\nauto_update = true\nchannel = \"beta\"\n\n[features]\nfeedback = true\nlsp_tools = true\nweb_fetch = true\nwrite_file = true\n\n[ui]\nyolo = true\n\n[models]\ndefault = \"user-model\"\n\n[endpoints]\ncoding_api_base_url = \"https://user-proxy.example/v1\"\napi_base_url = \"https://user-api.example/v1\"\nmodels_base_url = \"https://user-models.example/v1\"\nmodels_list_url = \"https://user-models.example/v1/models\"\n",
)
.unwrap();
let mut cfg = crate::agent::config::Config::new_from_toml_cfg(&raw_config).unwrap();
cfg.default_yolo_mode = true;
let requirements: toml::Value = toml::from_str(
"[cli]\nauto_update = false\nchannel = \"stable\"\n\n[features]\nfeedback = false\nlsp_tools = false\nweb_fetch = false\nwrite_file = false\nremote_fetch = false\n\n[ui]\nyolo = false\n\n[models]\ndefault = \"managed-model\"\n\n[endpoints]\ncoding_api_base_url = \"https://managed-proxy.example/v1\"\nxai_api_base_url = \"https://managed-api.example/v1\"\nmodels_base_url = \"https://managed-models.example/v1\"\nmodels_list_url = \"https://managed-models.example/v1/models\"\ndeployment_key = \"enterprise-deploy-key-should-not-log\"\n",
"[cli]\nauto_update = false\nchannel = \"stable\"\n\n[features]\nfeedback = false\nlsp_tools = false\nweb_fetch = false\nwrite_file = false\nremote_fetch = false\n\n[ui]\nyolo = false\n\n[models]\ndefault = \"managed-model\"\n\n[endpoints]\ncoding_api_base_url = \"https://managed-proxy.example/v1\"\napi_base_url = \"https://managed-api.example/v1\"\nmodels_base_url = \"https://managed-models.example/v1\"\nmodels_list_url = \"https://managed-models.example/v1/models\"\ndeployment_key = \"enterprise-deploy-key-should-not-log\"\n",
)
.unwrap();
let source = RequirementSource::Requirements {
@@ -2617,7 +2617,7 @@ fn apply_requirements_value_overrides_user_settings() {
Some("https://managed-proxy.example/v1"), cfg.endpoints.coding_api_base_url
.as_deref()
);
assert_eq!("https://managed-api.example/v1", cfg.endpoints.xai_api_base_url);
assert_eq!("https://managed-api.example/v1", cfg.endpoints.api_base_url);
assert_eq!(
Some("https://managed-models.example/v1"), cfg.endpoints.models_base_url
.as_deref()
crates/codegen/kigi-shell/src/inspect/mod.rs
+3 -3
View File
@@ -53,7 +53,7 @@ impl std::fmt::Display for Scope {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct InspectReport {
pub grok_version: String,
pub kigi_version: String,
pub channel: String,
pub cwd: String,
pub project_root: Option<String>,
@@ -370,7 +370,7 @@ async fn build_report(cwd: &Path) -> InspectReport {
.unwrap_or_default();
InspectReport {
grok_version: kigi_version::VERSION.to_string(),
kigi_version: kigi_version::VERSION.to_string(),
channel: crate::util::config::channel_name_from_cache()
.unwrap_or("unknown")
.to_string(),
@@ -1224,7 +1224,7 @@ fn render_harness_compatibility(report: &ExternalCompatReport) -> String {
fn print_human(r: &InspectReport) {
println!();
println!(" Environment");
println!(" {TREE} Version: {} [{}]", r.grok_version, r.channel);
println!(" {TREE} Version: {} [{}]", r.kigi_version, r.channel);
println!(" {TREE} CWD: {}", r.cwd);
if let Some(ref root) = r.project_root {
println!(" {TREE} Git root: {}", root);
crates/codegen/kigi-shell/src/leader/mod.rs
+1 -1
View File
@@ -1113,7 +1113,7 @@ async fn evict_leader(conn: LeaderConnection, lock: &LeaderLock) {
/// # Arguments
///
/// * `client_type` - Identifier for the client type (e.g., "grok-tui", "vscode")
/// * `mode` - Communication mode (Stdio or Headless)
/// * `mode` - Communication mode (Stdio)
/// * `capabilities` - Client capabilities (e.g., yolo_mode) to register with the leader
pub async fn connect_or_spawn(
client_type: &str,
crates/codegen/kigi-shell/src/leader/protocol.rs
-93
View File
@@ -108,9 +108,6 @@ impl Default for ClientId {
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ClientMode {
/// Headless mode (grok agent, grok agent headless) - uses websocket relay.
/// Leader connects to websocket relay once and forwards messages.
Headless,
/// Stdio mode (grok agent stdio, grok -p) - uses local IPC.
/// Client sends/receives ACP messages directly via IPC.
Stdio,
@@ -178,8 +175,6 @@ pub struct LeaderCapabilities {
pub runtime_cpu_profile: bool,
#[serde(default)]
pub profile_formats: Vec<ProfileArtifactFormat>,
#[serde(default)]
pub workspace_exposure: bool,
/// Whether the leader supports [`ControlCommand::RelaunchForUpdate`] — a
/// disruptive, bounded-grace relaunch onto a freshly-installed binary
/// (driven by `grok update`). Old leaders default to `false`, so a new
@@ -200,15 +195,6 @@ pub enum ControlCommand {
frequency_hz: Option<i32>,
},
StopCpuProfile,
WorkspaceStart {
#[serde(default)]
hub_url: Option<String>,
cwd: String,
},
WorkspacePause,
WorkspaceResume,
WorkspaceStop,
WorkspaceStatus,
/// Ask the leader to relaunch onto a freshly-installed binary (driven by
/// `grok update`). The leader stops admitting new turns, waits a bounded
/// grace period for in-flight turns to finish, flushes session state, then
@@ -260,18 +246,6 @@ pub enum ControlPayload {
started_at: String,
stopped_at: String,
},
WorkspaceStatus {
state: String,
#[serde(default)]
hub_url: Option<String>,
#[serde(default)]
cwd: Option<String>,
uptime_ms: u64,
active_tool_calls: u32,
#[serde(default)]
sessions: Vec<String>,
pid: u32,
},
/// Ack for [`ControlCommand::RelaunchForUpdate`]: the leader accepted the
/// request and will exit after a bounded grace period of `grace_ms`.
Relaunching {
@@ -531,7 +505,6 @@ mod tests {
control_v1: true,
runtime_cpu_profile: true,
profile_formats: vec![ProfileArtifactFormat::Svg],
workspace_exposure: true,
relaunch_v1: true,
}),
};
@@ -549,7 +522,6 @@ mod tests {
control_v1: true,
runtime_cpu_profile: true,
profile_formats,
workspace_exposure: true,
relaunch_v1: true,
}),
} if profile_formats == vec![ProfileArtifactFormat::Svg]
@@ -622,71 +594,6 @@ mod tests {
));
}
#[tokio::test]
async fn workspace_control_command_roundtrip() {
let (mut client, mut server) = duplex(1024);
let msg = ClientMessage::Control {
request_id: "ws-1".into(),
command: ControlCommand::WorkspaceStart {
hub_url: Some("wss://hub.example/v1/tools".into()),
cwd: "/home/u/proj".into(),
},
};
write_message(&mut client, &msg).await.unwrap();
let received: ClientMessage = read_message(&mut server).await.unwrap();
assert!(matches!(
received,
ClientMessage::Control {
request_id,
command: ControlCommand::WorkspaceStart { hub_url: Some(url), cwd },
} if request_id == "ws-1"
&& url == "wss://hub.example/v1/tools"
&& cwd == "/home/u/proj"
));
}
#[test]
fn workspace_status_payload_roundtrip() {
let payload = ControlPayload::WorkspaceStatus {
state: "running".into(),
hub_url: Some("wss://hub.example/v1/tools".into()),
cwd: Some("/home/u/proj".into()),
uptime_ms: 4200,
active_tool_calls: 2,
sessions: vec!["grok-a".into(), "grok-b".into()],
pid: 4242,
};
let json = serde_json::to_string(&payload).unwrap();
let decoded: ControlPayload = serde_json::from_str(&json).unwrap();
assert_eq!(decoded, payload);
assert!(json.contains("\"type\":\"workspace_status\""));
}
#[test]
fn workspace_status_payload_defaults_optional_fields() {
let json = r#"{"type":"workspace_status","state":"none","uptime_ms":0,"active_tool_calls":0,"pid":1}"#;
let decoded: ControlPayload = serde_json::from_str(json).unwrap();
assert!(matches!(
decoded,
ControlPayload::WorkspaceStatus {
state,
hub_url: None,
cwd: None,
sessions,
..
} if state == "none" && sessions.is_empty()
));
}
#[test]
fn workspace_exposure_capability_defaults_false() {
let json = r#"{"control_v1":true,"runtime_cpu_profile":false,"profile_formats":[]}"#;
let caps: LeaderCapabilities = serde_json::from_str(json).unwrap();
assert!(!caps.workspace_exposure);
}
#[test]
fn client_id_is_unique() {
let ids: Vec<_> = (0..100).map(|_| ClientId::new()).collect();
crates/codegen/kigi-shell/src/leader/server.rs
-367
View File
@@ -27,8 +27,6 @@ use crate::cpu_profile::{
};
use agent_client_protocol::AGENT_METHOD_NAMES;
use kanal::{AsyncReceiver, AsyncSender};
use kigi_computer_hub_sdk::{AuthCredential, AuthIdentity, AuthProvider};
use kigi_workspace::WorkspaceHandle;
use parking_lot::Mutex;
use tokio::sync::{mpsc, watch};
use tokio_util::sync::CancellationToken;
@@ -129,105 +127,24 @@ pub struct LeaderServerMetadata {
pub struct LeaderServerControlState {
pub metadata: LeaderServerMetadata,
pub cpu_profile: Arc<Mutex<CpuProfileManager>>,
pub workspace: Arc<WorkspaceControl>,
}
impl LeaderServerControlState {
pub fn new(metadata: LeaderServerMetadata) -> Self {
Self {
metadata,
cpu_profile: Arc::new(Mutex::new(CpuProfileManager::new())),
workspace: Arc::new(WorkspaceControl::new(None)),
}
}
pub fn with_default_hub_url(mut self, default_hub_url: Option<String>) -> Self {
self.workspace = Arc::new(WorkspaceControl::new(default_hub_url));
self
}
fn leader_capabilities(&self) -> LeaderCapabilities {
let manager = self.cpu_profile.lock();
LeaderCapabilities {
control_v1: true,
runtime_cpu_profile: manager.runtime_cpu_profile(),
profile_formats: manager.profile_formats().to_vec(),
workspace_exposure: true,
relaunch_v1: true,
}
}
}
pub struct WorkspaceControl {
default_hub_url: Option<String>,
/// Hub credential, wired to the leader's `AuthManager` once auth is ready.
/// A `watch` so a starting leader (socket up, auth pending) can be awaited
/// instead of failing the command.
auth: tokio::sync::watch::Sender<Option<Arc<dyn AuthProvider>>>,
/// Serializes mutating commands (start/pause/resume/stop) so their long
/// awaits (drain, reconnect) never interleave.
lock: tokio::sync::Mutex<()>,
/// Current exposure, published for lock-free reads so `status` never
/// blocks behind an in-flight drain/reconnect.
exposure: arc_swap::ArcSwapOption<WorkspaceExposure>,
}
impl WorkspaceControl {
fn new(default_hub_url: Option<String>) -> Self {
Self {
default_hub_url,
auth: tokio::sync::watch::channel(None).0,
lock: tokio::sync::Mutex::new(()),
exposure: arc_swap::ArcSwapOption::empty(),
}
}
/// Wire the hub credential to the leader's shared `AuthManager` (sole
/// owner of refresh + persistence).
pub fn set_auth_manager(&self, auth_manager: Arc<AuthManager>) {
self.auth
.send_replace(Some(Arc::new(LeaderAuthProvider { auth_manager })));
}
}
impl std::fmt::Debug for WorkspaceControl {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("WorkspaceControl")
.field("default_hub_url", &self.default_hub_url)
.finish_non_exhaustive()
}
}
/// Hub [`AuthProvider`] backed by the leader's `AuthManager`: returns the
/// current token at each connect/reconnect; never writes auth.json.
struct LeaderAuthProvider {
auth_manager: Arc<AuthManager>,
}
impl std::fmt::Debug for LeaderAuthProvider {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LeaderAuthProvider").finish_non_exhaustive()
}
}
impl AuthProvider for LeaderAuthProvider {
fn current(&self) -> AuthCredential {
let token = self
.auth_manager
.current_or_expired()
.map(|a| a.key)
.unwrap_or_default();
AuthCredential::bearer(token)
}
/// Owner identity from the leader's `AuthManager`, surfaced on the auth
/// provider instead of a separate auth.json read. The Kimi credential
/// carries no principal metadata; only the (possibly empty) user id.
fn identity(&self) -> Option<AuthIdentity> {
let a = self.auth_manager.current_or_expired()?;
Some(AuthIdentity {
user_id: a.user_id,
principal_type: None,
principal_id: None,
})
}
}
struct WorkspaceExposure {
handle: WorkspaceHandle,
hub_url: String,
cwd: PathBuf,
started_at: Instant,
paused: std::sync::atomic::AtomicBool,
}
/// Rewrite JSON-RPC request ID **in place** by prefixing with client ID to
/// avoid collisions.
///
@@ -934,233 +851,6 @@ fn leader_info_payload(control_state: &LeaderServerControlState) -> ControlPaylo
profile_formats: manager.profile_formats().to_vec(),
}
}
const PROD_COMPUTER_HUB_URL: &str = "wss://computer-hub.kigi.com/v1/tools";
const WORKSPACE_DRAIN_TIMEOUT: Duration = Duration::from_secs(10);
fn workspace_err(message: impl Into<String>) -> ControlError {
ControlError {
code: ControlErrorCode::InternalError,
message: message.into(),
details: None,
}
}
/// Resolve the hub credential, waiting if the leader is still wiring auth
/// (the IPC socket comes up first). Resolves the instant auth is wired or the
/// leader cancels — event-driven, no timeout.
async fn wait_for_leader_auth(
ws: &WorkspaceControl,
cancel: &CancellationToken,
) -> Result<Arc<dyn AuthProvider>, ControlError> {
let mut rx = ws.auth.subscribe();
tokio::select! {
result = rx.wait_for(| v | v.is_some()) => match result { Ok(guard) => Ok(guard
.clone().expect("waited for Some")), Err(_) =>
Err(workspace_err("leader is shutting down; cannot expose workspace to the hub",)),
}, _ = cancel.cancelled() =>
Err(workspace_err("leader is shutting down; cannot expose workspace to the hub",)),
}
}
fn workspace_server_id() -> String {
let raw = gethostname::gethostname()
.to_string_lossy()
.to_ascii_lowercase();
let sanitized: String = raw
.chars()
.map(|c| {
if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
c
} else {
'-'
}
})
.collect();
let name = sanitized.trim_matches('-');
if name.is_empty() {
"grok-workspace".to_string()
} else {
name.to_string()
}
}
async fn drain_and_disconnect(handle: &WorkspaceHandle) {
let tracker = handle.activity_tracker().clone();
tracker.set_draining();
if tokio::time::timeout(WORKSPACE_DRAIN_TIMEOUT, tracker.wait_until_drained())
.await
.is_err()
{
warn!(
active = tracker.total_active(),
"workspace drain timed out; disconnecting hub anyway"
);
}
handle.shutdown_hub().await;
}
fn build_workspace_status(
metadata: &LeaderServerMetadata,
exposure: Option<&WorkspaceExposure>,
) -> ControlPayload {
match exposure {
None => ControlPayload::WorkspaceStatus {
state: "none".to_string(),
hub_url: None,
cwd: None,
uptime_ms: 0,
active_tool_calls: 0,
sessions: Vec::new(),
pid: metadata.pid,
},
Some(exp) => {
let snapshot = exp.handle.activity_tracker().snapshot();
let mut sessions = exp.handle.session_ids();
sessions.sort();
ControlPayload::WorkspaceStatus {
state: if exp.paused.load(std::sync::atomic::Ordering::Relaxed) {
"paused"
} else {
"running"
}
.to_string(),
hub_url: Some(exp.hub_url.clone()),
cwd: Some(exp.cwd.display().to_string()),
uptime_ms: exp.started_at.elapsed().as_millis() as u64,
active_tool_calls: snapshot.active_tool_calls,
sessions,
pid: metadata.pid,
}
}
}
}
async fn handle_workspace_start(
control_state: LeaderServerControlState,
hub_url: Option<String>,
cwd: String,
cancel: CancellationToken,
) -> Result<ControlPayload, ControlError> {
let ws = &control_state.workspace;
let url_str = hub_url
.filter(|u| !u.trim().is_empty())
.or_else(|| ws.default_hub_url.clone())
.unwrap_or_else(|| PROD_COMPUTER_HUB_URL.to_string());
let url = url::Url::parse(&url_str)
.map_err(|e| workspace_err(format!("invalid hub url {url_str}: {e}")))?;
let cwd_path = PathBuf::from(&cwd);
let _serialize = ws.lock.lock().await;
if let Some(existing) = ws.exposure.load_full()
&& !existing.paused.load(Ordering::Relaxed)
&& existing.cwd == cwd_path
&& existing.hub_url == url_str
{
return Ok(build_workspace_status(
&control_state.metadata,
Some(existing.as_ref()),
));
}
let allow_insecure_ws =
url.scheme() == "ws" && matches!(url.host_str(), Some("localhost" | "127.0.0.1" | "::1"));
let status_config = kigi_workspace::StatusConfig::from_env();
let alpha_test_key = None;
let auth = wait_for_leader_auth(ws, &cancel).await?;
let server_id = workspace_server_id();
let metadata = serde_json::json!(
{ "source" : "grok-workspace", "hostname" : gethostname::gethostname()
.to_string_lossy(), "cwd" : cwd_path.display().to_string(), }
);
crate::agent::folder_trust::resolve_and_record(&cwd_path, None, false);
let project_lsp_trusted = crate::agent::folder_trust::project_scope_allowed(&cwd_path);
let handle = kigi_workspace::connect_local_workspace(
cwd_path.clone(),
url,
auth,
Some(metadata),
Some(server_id),
alpha_test_key,
allow_insecure_ws,
status_config,
project_lsp_trusted,
None,
None,
false,
false,
)
.await
.map_err(|e| workspace_err(format!("failed to connect workspace to hub: {e}")))?;
let exposure = Arc::new(WorkspaceExposure {
handle,
hub_url: url_str,
cwd: cwd_path,
started_at: Instant::now(),
paused: AtomicBool::new(false),
});
let payload = build_workspace_status(&control_state.metadata, Some(exposure.as_ref()));
if let Some(old) = ws.exposure.swap(Some(exposure)) {
drain_and_disconnect(&old.handle).await;
}
Ok(payload)
}
async fn handle_workspace_pause(
control_state: LeaderServerControlState,
) -> Result<ControlPayload, ControlError> {
let ws = &control_state.workspace;
let _serialize = ws.lock.lock().await;
let Some(exp) = ws.exposure.load_full() else {
return Err(workspace_err("no workspace exposure is running"));
};
if !exp.paused.load(Ordering::Relaxed) {
drain_and_disconnect(&exp.handle).await;
exp.paused.store(true, Ordering::Relaxed);
}
Ok(build_workspace_status(
&control_state.metadata,
Some(exp.as_ref()),
))
}
async fn handle_workspace_resume(
control_state: LeaderServerControlState,
) -> Result<ControlPayload, ControlError> {
let ws = &control_state.workspace;
let _serialize = ws.lock.lock().await;
let Some(exp) = ws.exposure.load_full() else {
return Err(workspace_err("no workspace exposure is running"));
};
if exp.paused.load(Ordering::Relaxed) {
exp.handle.activity_tracker().set_active();
if let Err(e) = exp.handle.connect_hub().await {
exp.handle.activity_tracker().set_draining();
return Err(workspace_err(format!("failed to reconnect to hub: {e}")));
}
exp.paused.store(false, Ordering::Relaxed);
}
Ok(build_workspace_status(
&control_state.metadata,
Some(exp.as_ref()),
))
}
async fn handle_workspace_stop(
control_state: LeaderServerControlState,
) -> Result<ControlPayload, ControlError> {
let ws = &control_state.workspace;
let _serialize = ws.lock.lock().await;
if let Some(exp) = ws.exposure.swap(None) {
drain_and_disconnect(&exp.handle).await;
}
Ok(build_workspace_status(&control_state.metadata, None))
}
async fn handle_workspace_status(
control_state: LeaderServerControlState,
) -> Result<ControlPayload, ControlError> {
let exposure = control_state.workspace.exposure.load_full();
Ok(build_workspace_status(
&control_state.metadata,
exposure.as_deref(),
))
}
async fn finalize_workspace_on_shutdown(control_state: LeaderServerControlState) {
let ws = &control_state.workspace;
let _serialize = ws.lock.lock().await;
if let Some(exp) = ws.exposure.swap(None) {
info!("Draining workspace exposure on leader shutdown");
drain_and_disconnect(&exp.handle).await;
}
}
fn handle_control_command(
control_state: &LeaderServerControlState,
command: ControlCommand,
@@ -1208,13 +898,6 @@ fn handle_control_command(
ControlCommand::StopCpuProfile => {
unreachable!("StopCpuProfile must be handled asynchronously")
}
ControlCommand::WorkspaceStart { .. }
| ControlCommand::WorkspacePause
| ControlCommand::WorkspaceResume
| ControlCommand::WorkspaceStop
| ControlCommand::WorkspaceStatus => {
unreachable!("workspace control commands are handled asynchronously")
}
ControlCommand::RelaunchForUpdate { .. } => {
unreachable!("RelaunchForUpdate must be handled asynchronously")
}
@@ -1564,13 +1247,6 @@ pub async fn run_leader_server(
agent_activity = agent_activity.clone(); let relaunching = relaunching
.clone(); tokio::spawn(async move { let result = match command {
ControlCommand::StopCpuProfile => { handle_stop_cpu_profile(control_state).
await } ControlCommand::WorkspaceStart { hub_url, cwd } => {
handle_workspace_start(control_state, hub_url, cwd, cancel.clone(),). await }
ControlCommand::WorkspacePause => { handle_workspace_pause(control_state).
await } ControlCommand::WorkspaceResume => {
handle_workspace_resume(control_state). await } ControlCommand::WorkspaceStop
=> { handle_workspace_stop(control_state). await }
ControlCommand::WorkspaceStatus => { handle_workspace_status(control_state).
await } ControlCommand::RelaunchForUpdate { to_version } => {
decide_relaunch_for_update(& control_state, to_version, & relaunching,) }
other => handle_control_command(& control_state, other), }; let arm_relaunch
@@ -1806,7 +1482,6 @@ pub async fn run_leader_server(
debug!("No client available for notification routing, message dropped"); } }
}
}
finalize_workspace_on_shutdown(control_state.clone()).await;
finalize_cpu_profile_on_shutdown(control_state).await;
let _ = std::fs::remove_file(&socket_path);
Ok(())
@@ -2214,48 +1889,6 @@ mod tests {
Ok(ControlPayload::RelaunchDeclined { .. })
));
}
#[derive(Debug)]
struct TestAuth;
impl AuthProvider for TestAuth {
fn current(&self) -> AuthCredential {
AuthCredential::bearer("test-token")
}
}
#[tokio::test]
async fn wait_for_leader_auth_returns_when_already_wired() {
let ws = WorkspaceControl::new(None);
ws.auth.send_replace(Some(Arc::new(TestAuth)));
let cancel = CancellationToken::new();
let auth = wait_for_leader_auth(&ws, &cancel).await.expect("wired");
assert!(matches!(auth.current(), AuthCredential::Bearer { .. }));
}
#[tokio::test]
async fn wait_for_leader_auth_resolves_when_wired_late() {
let ws = Arc::new(WorkspaceControl::new(None));
let cancel = CancellationToken::new();
let waiter = {
let ws = ws.clone();
let cancel = cancel.clone();
tokio::spawn(async move { wait_for_leader_auth(&ws, &cancel).await.is_ok() })
};
tokio::task::yield_now().await;
ws.auth.send_replace(Some(Arc::new(TestAuth)));
assert!(waiter.await.unwrap(), "auth wired late should resolve Ok");
}
#[tokio::test]
async fn workspace_start_errors_when_cancelled_before_auth() {
let state = default_test_control_state(Path::new("/tmp/grok-ws-auth-test.sock"));
let cancel = CancellationToken::new();
cancel.cancel();
let err = handle_workspace_start(state, None, "/tmp".to_string(), cancel)
.await
.unwrap_err();
assert!(
err.message.contains("shutting down"),
"unexpected error: {}",
err.message
);
}
async fn setup_test_server(
temp: &TempDir,
) -> (PathBuf, CancellationToken, mpsc::UnboundedReceiver<String>) {
crates/codegen/kigi-shell/src/leader/test_support.rs
-1
View File
@@ -35,7 +35,6 @@ pub(crate) fn fake_caps(control_v1: bool, relaunch_v1: bool) -> LeaderCapabiliti
control_v1,
runtime_cpu_profile: false,
profile_formats: Vec::new(),
workspace_exposure: false,
relaunch_v1,
}
}
crates/codegen/kigi-shell/src/session/acp_session.rs
-5
View File
@@ -928,9 +928,6 @@ pub(crate) struct SessionActor {
/// Centralized event tracking: event log, turn-end guard, active tool,
/// doom loop terminate flag. All event-related state lives here.
pub(crate) events: crate::session::events::EventTracker,
/// Optional hub-side session event emitter (always constructed without a
/// harness client in the agent; methods no-op with `None` transport).
pub(crate) observability_bridge: kigi_computer_hub_sdk::ObservabilityBridge,
/// Turn number captured at the start of each turn (before prompt index
/// increment). Used by `ToolCallStarted` bridge emissions so they
/// report the same turn number as `TurnStarted` / `TurnEnded`.
@@ -1522,8 +1519,6 @@ mod fs_injection_regression_tests;
#[path = "acp_session_tests/interjection_actor_tests.rs"]
mod interjection_actor_tests;
#[cfg(test)]
#[path = "acp_session_tests/observability_bridge_mapping_tests.rs"]
mod observability_bridge_mapping_tests;
#[cfg(test)]
#[path = "acp_session_tests/permission_auto_mode_tests.rs"]
mod permission_auto_mode_tests;
crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs
+1 -34
View File
@@ -251,34 +251,8 @@ pub(crate) async fn spawn_session_actor(
.as_ref()
.map(kigi_workspace::permission::resolution::deny_read_globs_from_config)
.unwrap_or_default();
let hub_permission = if kigi_workspace::permission::hitl_permission_live_enabled() {
let server = match workspace_ops.workspace_handle() {
Some(handle) => handle.hub_server_blocking().await,
None => None,
};
let transport = server
.and_then(|server| {
kigi_workspace::permission::ToolServerPermissionTransport::from_session_id(
server,
session_info.id.0.as_ref(),
)
})
.map(|t| {
std::sync::Arc::new(t)
as std::sync::Arc<dyn kigi_workspace::permission::PermissionHookTransport>
});
if transport.is_none() {
tracing::debug!(
session_id = % session_info.id.0,
"hitl permission live enabled but no remote transport available; using local prompt"
);
}
transport
} else {
None
};
let (permissions, _permission_events_rx) =
kigi_workspace::permission::spawn_permission_manager_with_hub(
kigi_workspace::permission::spawn_permission_manager(
session_info.id.clone(),
gateway.clone(),
tool_context.cwd.clone(),
@@ -289,7 +263,6 @@ pub(crate) async fn spawn_session_actor(
session_yolo_mode,
session_client_identifier.clone(),
crate::util::config::remember_tool_approvals_from_disk(),
hub_permission,
);
if crate::util::config::auto_mode_session_active(
crate::util::config::auto_permission_mode_enabled_from_disk(),
@@ -975,11 +948,6 @@ pub(crate) async fn spawn_session_actor(
let (goal_update_tx, goal_update_rx) = tokio::sync::mpsc::unbounded_channel::<
kigi_tools::implementations::grok_build::update_goal::UpdateGoalEnvelope,
>();
let obs_bridge = {
let sid = kigi_tool_protocol::SessionId::new(&*session_info.id.0)
.unwrap_or_else(|_| kigi_tool_protocol::SessionId::new("unknown").expect("valid"));
kigi_computer_hub_sdk::ObservabilityBridge::new(None, sid)
};
let mut effective_config = crate::config::load_effective_config()
.ok()
.and_then(|raw| crate::agent::config::Config::new_from_toml_cfg(&raw).ok())
@@ -1199,7 +1167,6 @@ pub(crate) async fn spawn_session_actor(
events: crate::session::events::EventTracker::new(
&crate::session::persistence::session_dir(&session_info),
),
observability_bridge: obs_bridge,
current_turn_number: std::cell::Cell::new(0),
last_recap_main_turn: std::cell::Cell::new(0),
recap_in_flight: std::cell::Cell::new(false),
crates/codegen/kigi-shell/src/session/acp_session_impl/tool_calls.rs
-19
View File
@@ -323,15 +323,6 @@ impl SessionActor {
self.emit_event(crate::session::events::Event::ToolStarted {
tool_name: call.function.name.clone(),
});
self.observability_bridge
.emit(
kigi_tool_protocol::session_event::SessionEvent::ToolCallStarted {
tool_call_id: call.id.clone(),
tool_name: call.function.name.clone(),
turn_number: self.current_turn_number.get(),
},
)
.await;
let call_name = call.function.name.clone();
match self
.prepare_tool_call(call, &mut deferred_followups)
@@ -653,16 +644,6 @@ impl SessionActor {
duration_ms,
outcome: tool_outcome,
});
self.observability_bridge
.emit(
kigi_tool_protocol::session_event::SessionEvent::ToolCallCompleted {
tool_call_id: prepared.call_id.clone(),
tool_name: prepared.tool_name.clone(),
duration_ms,
outcome: map_tool_outcome(tool_outcome),
},
)
.await;
tracing::info_span!(
"tool.execution", tool_name = % prepared.tool_name, tool_use_id = %
prepared.call_id, tool_input_size_bytes = prepared.raw_arguments.len() as
crates/codegen/kigi-shell/src/session/acp_session_impl/turn.rs
-32
View File
@@ -394,15 +394,6 @@ impl SessionActor {
schema_version: crate::session::events::EVENT_SCHEMA_VERSION.into(),
redirect_kind,
});
self.observability_bridge
.emit(
kigi_tool_protocol::session_event::SessionEvent::TurnStarted {
turn_number,
model_id: model_id.clone(),
yolo_mode,
},
)
.await;
self.send_before_turn_event(kigi_tool_protocol::turn_hook::BeforeTurnPayload {
turn_number: self.chat_state_handle.get_prompt_index().await as u64,
model_id: model_id.clone(),
@@ -733,15 +724,6 @@ impl SessionActor {
);
let turn_tool_count = self.events.tool_count_this_turn();
let bridge_outcome = turn_result_to_hook_outcome(&result);
self.observability_bridge
.emit(kigi_tool_protocol::session_event::SessionEvent::TurnEnded {
turn_number: current_prompt_index as u64,
outcome: bridge_outcome,
duration_ms: turn_duration_ms,
tool_call_count: turn_tool_count,
model_id: turn_model_id.clone(),
})
.await;
match &result {
Ok(TurnOutcome::Completed { .. }) => {
self.emit_turn_ended(
@@ -1678,13 +1660,6 @@ impl SessionActor {
self.emit_event(crate::session::events::Event::PhaseChanged {
phase: crate::session::events::Phase::WaitingForModel,
});
self.observability_bridge
.emit(
kigi_tool_protocol::session_event::SessionEvent::PhaseChanged {
phase: kigi_tool_protocol::session_event::SessionPhase::Sampling,
},
)
.await;
kigi_log::unified_log::info(
"shell.turn.inference_start",
Some(self.session_info.id.0.as_ref()),
@@ -1989,13 +1964,6 @@ impl SessionActor {
self.emit_event(crate::session::events::Event::PhaseChanged {
phase: crate::session::events::Phase::ToolExecution,
});
self.observability_bridge
.emit(
kigi_tool_protocol::session_event::SessionEvent::PhaseChanged {
phase: kigi_tool_protocol::session_event::SessionPhase::ToolExecution,
},
)
.await;
let execute_tool_calls_result = self.execute_tool_calls(tool_call_responses).await;
match execute_tool_calls_result {
Ok(ToolLoop::PermissionReject { tool_name, reason }) => {
crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs
-4
View File
@@ -256,7 +256,6 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
plugin_registry: std::cell::RefCell::new(None),
plugin_registry_handle: None,
events: crate::session::events::EventTracker::new(std::path::Path::new("/tmp")),
observability_bridge: noop_observability_bridge(),
current_turn_number: std::cell::Cell::new(0),
last_recap_main_turn: std::cell::Cell::new(0),
recap_in_flight: std::cell::Cell::new(false),
@@ -695,7 +694,6 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
plugin_registry: std::cell::RefCell::new(None),
plugin_registry_handle: None,
events: crate::session::events::EventTracker::new(std::path::Path::new("/tmp")),
observability_bridge: noop_observability_bridge(),
current_turn_number: std::cell::Cell::new(0),
last_recap_main_turn: std::cell::Cell::new(0),
recap_in_flight: std::cell::Cell::new(false),
@@ -943,7 +941,6 @@ async fn cancel_running_task_teardown_clears_running_and_pending_work() {
plugin_registry: std::cell::RefCell::new(None),
plugin_registry_handle: None,
events: crate::session::events::EventTracker::new(std::path::Path::new("/tmp")),
observability_bridge: noop_observability_bridge(),
current_turn_number: std::cell::Cell::new(0),
last_recap_main_turn: std::cell::Cell::new(0),
recap_in_flight: std::cell::Cell::new(false),
@@ -1924,7 +1921,6 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() {
plugin_registry: std::cell::RefCell::new(None),
plugin_registry_handle: None,
events: crate::session::events::EventTracker::new(std::path::Path::new("/tmp")),
observability_bridge: noop_observability_bridge(),
current_turn_number: std::cell::Cell::new(0),
last_recap_main_turn: std::cell::Cell::new(0),
recap_in_flight: std::cell::Cell::new(false),
crates/codegen/kigi-shell/src/session/acp_session_tests/fs_injection_regression_tests.rs
-1
View File
@@ -62,7 +62,6 @@ async fn tool_bridge_routes_writes_through_injected_fs() {
video_gen_config: Default::default(),
app_builder_deployer_config: Default::default(),
api_key_provider: None,
auth_provider: None,
attribution_callback: None,
system_reminder_tag: kigi_tools::reminders::DEFAULT_REMINDER_TAG,
};
crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs
-1
View File
@@ -285,7 +285,6 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
plugin_registry: std::cell::RefCell::new(None),
plugin_registry_handle: None,
events: crate::session::events::EventTracker::new(std::path::Path::new("/tmp")),
observability_bridge: noop_observability_bridge(),
current_turn_number: std::cell::Cell::new(0),
last_recap_main_turn: std::cell::Cell::new(0),
recap_in_flight: std::cell::Cell::new(false),
crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs
-3
View File
@@ -214,7 +214,6 @@ async fn create_test_actor(
plugin_registry: std::cell::RefCell::new(None),
plugin_registry_handle: None,
events: crate::session::events::EventTracker::new(std::path::Path::new("/tmp")),
observability_bridge: noop_observability_bridge(),
current_turn_number: std::cell::Cell::new(0),
last_recap_main_turn: std::cell::Cell::new(0),
recap_in_flight: std::cell::Cell::new(false),
@@ -655,7 +654,6 @@ async fn create_test_actor_with_memory(
plugin_registry: std::cell::RefCell::new(None),
plugin_registry_handle: None,
events: crate::session::events::EventTracker::new(std::path::Path::new("/tmp")),
observability_bridge: noop_observability_bridge(),
current_turn_number: std::cell::Cell::new(0),
last_recap_main_turn: std::cell::Cell::new(0),
recap_in_flight: std::cell::Cell::new(false),
@@ -1410,7 +1408,6 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
plugin_registry: std::cell::RefCell::new(None),
plugin_registry_handle: None,
events: crate::session::events::EventTracker::new(std::path::Path::new("/tmp")),
observability_bridge: noop_observability_bridge(),
current_turn_number: std::cell::Cell::new(0),
last_recap_main_turn: std::cell::Cell::new(0),
recap_in_flight: std::cell::Cell::new(false),
crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs
-1
View File
@@ -276,7 +276,6 @@ async fn create_test_actor_with_memory(
plugin_registry: std::cell::RefCell::new(None),
plugin_registry_handle: None,
events: crate::session::events::EventTracker::new(std::path::Path::new("/tmp")),
observability_bridge: noop_observability_bridge(),
current_turn_number: std::cell::Cell::new(0),
last_recap_main_turn: std::cell::Cell::new(0),
recap_in_flight: std::cell::Cell::new(false),
crates/codegen/kigi-shell/src/session/acp_session_tests/permission_auto_mode_tests.rs
+1
View File
@@ -33,6 +33,7 @@ fn install_real_permissions(actor: &mut SessionActor) {
vec![],
false,
None,
true,
);
actor.permissions = handle;
}
crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs
-1
View File
@@ -222,7 +222,6 @@ pub(super) async fn make_replay_send_update_fixture() -> ReplaySendUpdateFixture
plugin_registry: std::cell::RefCell::new(None),
plugin_registry_handle: None,
events: crate::session::events::EventTracker::new(std::path::Path::new("/tmp")),
observability_bridge: noop_observability_bridge(),
current_turn_number: std::cell::Cell::new(0),
last_recap_main_turn: std::cell::Cell::new(0),
recap_in_flight: std::cell::Cell::new(false),
crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs
-8
View File
@@ -12,12 +12,6 @@ pub(crate) const HARNESS_VERIFIES_SENTENCE: &str =
pub(crate) const PLAN_SEED_TODOS_PHRASE: &str =
"Seed todos from the plan's acceptance criteria via";
#[cfg(test)]
pub(crate) fn noop_observability_bridge() -> kigi_computer_hub_sdk::ObservabilityBridge {
kigi_computer_hub_sdk::ObservabilityBridge::new(
None,
kigi_tool_protocol::SessionId::new("test").expect("valid"),
)
}
#[cfg(test)]
pub(crate) async fn test_agent_default() -> kigi_agent::Agent {
test_agent_with_tools(vec![]).await
@@ -99,7 +93,6 @@ async fn test_agent_from_config(
video_gen_config: Default::default(),
app_builder_deployer_config: Default::default(),
api_key_provider: None,
auth_provider: None,
attribution_callback: None,
system_reminder_tag: kigi_tools::reminders::DEFAULT_REMINDER_TAG,
};
@@ -336,7 +329,6 @@ pub(crate) async fn create_test_actor_ex(
plugin_registry: std::cell::RefCell::new(None),
plugin_registry_handle: None,
events: crate::session::events::EventTracker::new(std::path::Path::new("/tmp")),
observability_bridge: noop_observability_bridge(),
current_turn_number: std::cell::Cell::new(0),
last_recap_main_turn: std::cell::Cell::new(0),
recap_in_flight: std::cell::Cell::new(false),
crates/codegen/kigi-shell/src/session/acp_session_tests/web_search_e2e_tests.rs
-1
View File
@@ -41,7 +41,6 @@ async fn web_search_errors_when_disabled() {
video_gen_config: Default::default(),
app_builder_deployer_config: Default::default(),
api_key_provider: None,
auth_provider: None,
attribution_callback: None,
system_reminder_tag: kigi_tools::reminders::DEFAULT_REMINDER_TAG,
};
crates/codegen/kigi-shell/src/session/compaction.rs
-1
View File
@@ -2312,7 +2312,6 @@ mod inline_auto_compact_flow_tests {
plugin_registry: std::cell::RefCell::new(None),
plugin_registry_handle: None,
events: crate::session::events::EventTracker::new(std::path::Path::new("/tmp")),
observability_bridge: noop_observability_bridge(),
current_turn_number: std::cell::Cell::new(0),
last_recap_main_turn: std::cell::Cell::new(0),
recap_in_flight: std::cell::Cell::new(false),
crates/codegen/kigi-shell/src/session/storage/jsonl/tests.rs
+1 -1
View File
@@ -1052,7 +1052,7 @@ async fn test_load_prompts_only_large_session() {
info.id.clone(),
acp::SessionUpdate::UserMessageChunk(
acp::ContentChunk::new(
acp::ContentBlock::Text(acp::TextContent::new(format!("part2"))),
acp::ContentBlock::Text(acp::TextContent::new("part2".to_string())),
),
),
);
crates/codegen/kigi-shell/tests/git_contention_e2e.rs
+1 -1
View File
@@ -432,7 +432,7 @@ fn git_rebase_refresh_storm_e2e() {
unsafe {
std::env::set_var("KIGI_SHARE_DIR", kigi_home.path());
std::env::set_var("KIGI_CODE_BASE_URL", server.url());
std::env::set_var("KIGI_XAI_API_BASE_URL", server.url());
std::env::set_var("KIGI_API_BASE_URL", server.url());
std::env::set_var("XAI_API_KEY", "test-key-for-ci");
std::env::set_var("KIGI_TELEMETRY_ENABLED", "false");
std::env::set_var("KIGI_FEEDBACK_ENABLED", "false");
crates/codegen/kigi-shell/tests/session_load_perf.rs
+1 -1
View File
@@ -626,7 +626,7 @@ async fn full_session_load_e2e() {
std::env::set_var("KIGI_INSTRUMENTATION", "log");
std::env::set_var("KIGI_INSTRUMENTATION_LOG", &instr_log);
std::env::set_var("KIGI_CODE_BASE_URL", server.url());
std::env::set_var("KIGI_XAI_API_BASE_URL", server.url());
std::env::set_var("KIGI_API_BASE_URL", server.url());
std::env::set_var("XAI_API_KEY", "test-key-for-ci");
std::env::set_var("KIGI_TELEMETRY_ENABLED", "false");
std::env::set_var("KIGI_FEEDBACK_ENABLED", "false");
crates/codegen/kigi-shell/tests/test_built_binary_e2e.rs
+1 -1
View File
@@ -1330,7 +1330,7 @@ async fn test_headless_managed_config_byok_sends_authorized_requests() {
r#"
[endpoints]
deployment_key = "test-deployment-key"
xai_api_base_url = "{url}"
api_base_url = "{url}"
[model.kigi-build]
api_backend = "responses"
crates/codegen/kigi-shell/tests/test_leader_soak.rs
+1 -1
View File
@@ -137,7 +137,7 @@ async fn leader_soak_churning_clients_no_leaks_no_zombies() {
unsafe {
std::env::set_var("KIGI_SHARE_DIR", kigi_home.path());
std::env::set_var("KIGI_CODE_BASE_URL", server.url());
std::env::set_var("KIGI_XAI_API_BASE_URL", server.url());
std::env::set_var("KIGI_API_BASE_URL", server.url());
std::env::set_var("XAI_API_KEY", "test-key-for-ci");
std::env::set_var("KIGI_TELEMETRY_ENABLED", "false");
std::env::set_var("KIGI_FEEDBACK_ENABLED", "false");
crates/codegen/kigi-shell/tests/test_mcp_permission_persistence.rs
+5 -5
View File
@@ -16,7 +16,7 @@ use kigi_workspace::permission::types::{
};
use kigi_workspace::permission::{
AccessKind, ClientType, Decision, PermissionCommand, PermissionHandle, PermissionState,
spawn_permission_manager, spawn_permission_manager_with_hub,
spawn_permission_manager,
};
use serial_test::serial;
use tokio::sync::{mpsc, oneshot};
@@ -221,6 +221,7 @@ async fn run_actor_test_full<F, Fut>(
vec![],
initial_yolo,
None,
true,
);
body(handle, gw, cwd).await;
})
@@ -357,7 +358,7 @@ async fn policy_ask_suppresses_mcp_tool_allowlist() {
let (gw, _gw_task) = fake_gateway();
// Gate OFF so the `ask` rule stays a hard floor over the grant.
let (handle, _events) = spawn_permission_manager_with_hub(
let (handle, _events) = spawn_permission_manager(
make_session_id(),
gw.sender.clone(),
cwd.clone(),
@@ -368,7 +369,6 @@ async fn policy_ask_suppresses_mcp_tool_allowlist() {
false,
None,
false, // remember_tool_approvals
None,
);
// Script an outright reject so we can confirm the prompt fires.
@@ -408,7 +408,7 @@ async fn policy_ask_suppresses_mcp_server_allowlist() {
let (gw, _gw_task) = fake_gateway();
// Gate OFF so the `ask` rule stays a hard floor over the grant.
let (handle, _events) = spawn_permission_manager_with_hub(
let (handle, _events) = spawn_permission_manager(
make_session_id(),
gw.sender.clone(),
cwd.clone(),
@@ -419,7 +419,6 @@ async fn policy_ask_suppresses_mcp_server_allowlist() {
false,
None,
false, // remember_tool_approvals
None,
);
gw.expected.send(("reject-once".to_string(), None)).unwrap();
@@ -463,6 +462,7 @@ async fn policy_deny_takes_precedence_over_mcp_allowlist() {
vec![],
false,
None,
true,
);
// Do NOT script a response: a policy Deny must short-circuit
crates/codegen/kigi-test-support/Cargo.toml
-1
View File
@@ -17,7 +17,6 @@ serde = { workspace = true }
serde_json = { workspace = true }
tempfile = { workspace = true }
tokio = { workspace = true }
tokio-tungstenite = { workspace = true }
tokio-util = { workspace = true, features = ["compat"] }
tracing = { workspace = true }
tracing-subscriber = { workspace = true, features = ["fmt"] }
crates/codegen/kigi-test-support/src/env.rs
+1 -1
View File
@@ -165,7 +165,7 @@ pub fn test_env_cmd_tokio(
// Mirrors `leader.rs` and the pty-harness `env_for_pager`.
.env("KIGI_SHARE_DIR", home.join(".kigi"))
.env("KIGI_CODE_BASE_URL", mock_url)
.env("KIGI_XAI_API_BASE_URL", mock_url)
.env("KIGI_API_BASE_URL", mock_url)
.env("XAI_API_KEY", "test-key-for-ci")
.env("KIGI_TELEMETRY_ENABLED", "false")
.env("KIGI_FEEDBACK_ENABLED", "false")
crates/codegen/kigi-test-support/src/leader.rs
+1 -1
View File
@@ -142,7 +142,7 @@ impl LeaderStdioClient {
// (re-)elected leader binds the same sandboxed path.
.env("KIGI_LEADER_SOCKET", home.join(".kigi").join("leader.sock"))
.env("KIGI_CODE_BASE_URL", server.url())
.env("KIGI_XAI_API_BASE_URL", server.url())
.env("KIGI_API_BASE_URL", server.url())
.env("XAI_API_KEY", "test-key-for-ci")
.env("KIGI_TELEMETRY_ENABLED", "false")
.env("KIGI_FEEDBACK_ENABLED", "false")
crates/codegen/kigi-tools-api/build.rs
+2 -3
View File
@@ -4,9 +4,8 @@ fn main() {
".", // match every message & enum
"#[derive(serde::Serialize, serde::Deserialize)]",
)
// ToolConfigEntry is embedded in external JSON contracts (Computer Hub
// `session.bind` metadata and agent-config JSON) where sparse payloads
// must deserialize. Defaults are applied per optional field (not
// ToolConfigEntry is embedded in external JSON contracts
// (agent-config JSON) where sparse payloads must deserialize. Defaults are applied per optional field (not
// type-level) so the required `id` field still fails deserialization
// when missing instead of silently becoming "". See tests/wire_shape.rs.
.field_attribute(
crates/codegen/kigi-tools/Cargo.toml
-3
View File
@@ -52,9 +52,6 @@ serde_path_to_error = { workspace = true }
kigi-tool-runtime = { workspace = true }
kigi-tool-types = { workspace = true }
kigi-tool-protocol = { workspace = true }
kigi-computer-hub-core = { workspace = true }
kigi-computer-hub-sdk = { workspace = true }
futures = { workspace = true }
similar = { workspace = true }
strum = { version = "0.26", features = ["derive"] }
crates/codegen/kigi-tools/src/implementations/grok_build_hashline/edit/apply.rs
+2 -2
View File
@@ -2093,11 +2093,11 @@ mod tests {
assert!(validate_anchor(&padded, &lines, &scheme).is_err());
// Stale anchor with arrow reports the stripped anchor in error metadata.
let stale = format!("2:zzz:zzz\u{2192}content");
let stale = "2:zzz:zzz\u{2192}content".to_string();
let err = validate_anchor(&stale, &lines, &scheme).unwrap_err();
assert_eq!(err.requested_anchor.as_deref(), Some("2:zzz:zzz"));
let stale_ascii = format!("2:zzz:zzz->content");
let stale_ascii = "2:zzz:zzz->content".to_string();
let err = validate_anchor(&stale_ascii, &lines, &scheme).unwrap_err();
assert_eq!(err.requested_anchor.as_deref(), Some("2:zzz:zzz"));
}
crates/codegen/kigi-tools/src/implementations/search_tool/mod.rs
+1 -1
View File
@@ -90,7 +90,7 @@ pub fn build_server_reminder(
return None;
}
let mut text = format!("Connected MCP servers:\n",);
let mut text = "Connected MCP servers:\n".to_string();
for server in servers {
text.push_str(&format_server_line(server));
}
crates/codegen/kigi-tools/src/registry/types.rs
+7 -20
View File
@@ -283,12 +283,6 @@ pub struct SessionContext {
/// instead of using the key baked into their config at construction time.
/// Prevents 401 failures when a session outlives the initial token lifetime.
pub api_key_provider: Option<crate::types::SharedApiKeyProvider>,
/// Auth provider which returns a kigi_computer_hub_sdk::AuthCredential. Can be used by
/// tools that need to authenticate with services.
///
/// Not to be confused with the api_key_provider, which is a legacy
/// provider used by the shell's auth manager.
pub auth_provider: Option<kigi_computer_hub_sdk::SharedAuthProvider>,
/// Optional 401-attribution callback for tool HTTP clients. When
/// set, a 401 from `image_gen` / `video_gen` / `web_search`
/// emits an `auth_401_attribution` event via this hook. Hosts can
@@ -351,7 +345,7 @@ type OutputConverter =
/// `.await`.
struct DispatchParts {
/// Resolved `LocalRegistry` handle to dispatch through.
lr_handle: Arc<dyn kigi_computer_hub_core::ToolHandle>,
lr_handle: kigi_tool_runtime::ArcTool,
/// Runtime context built for the call (resources, renderer, cwd,
/// behavior version, inner-dispatch).
ctx: kigi_tool_runtime::ToolCallContext,
@@ -394,7 +388,7 @@ struct ToolEntry {
>,
/// Registers this tool into a `LocalRegistry` using the concrete type.
/// Captured at `register::<T>()` time when T is known.
register_in_local: Box<dyn Fn(&kigi_computer_hub_sdk::LocalRegistry) + Send + Sync>,
register_in_local: Box<dyn Fn(&kigi_tool_runtime::LocalRegistry) + Send + Sync>,
}
/// Per-reminder metadata stored in the builder.
struct ReminderEntry {
@@ -450,7 +444,7 @@ pub struct FinalizedToolset {
scheduler_cancel: Option<tokio_util::sync::CancellationToken>,
/// Shared local registry for in-process dispatch.
/// Contains only config-enabled tools. Can be shared with ToolHarness.
local_registry: kigi_computer_hub_sdk::LocalRegistry,
local_registry: kigi_tool_runtime::LocalRegistry,
/// Lock-free access to the template renderer for tool name/param resolution.
/// Cloned into `ToolCallContext::extensions` on each `call()` so tools
/// can resolve names without acquiring the `resources` mutex.
@@ -519,7 +513,7 @@ impl RequirementError {
pub struct ToolRegistryBuilder {
tools: HashMap<String, ToolEntry>,
reminders: Vec<ReminderEntry>,
shared_local_registry: Option<kigi_computer_hub_sdk::LocalRegistry>,
shared_local_registry: Option<kigi_tool_runtime::LocalRegistry>,
}
impl Default for ToolRegistryBuilder {
fn default() -> Self {
@@ -613,7 +607,7 @@ impl ToolRegistryBuilder {
let typed = serde_json::from_value::<T::Args>(json)?;
Ok(typed.into())
}),
register_in_local: Box::new(|lr: &kigi_computer_hub_sdk::LocalRegistry| {
register_in_local: Box::new(|lr: &kigi_tool_runtime::LocalRegistry| {
lr.register(T::default());
}),
},
@@ -744,7 +738,7 @@ impl ToolRegistryBuilder {
}
b
}
pub fn with_local_registry(mut self, registry: kigi_computer_hub_sdk::LocalRegistry) -> Self {
pub fn with_local_registry(mut self, registry: kigi_tool_runtime::LocalRegistry) -> Self {
self.shared_local_registry = Some(registry);
self
}
@@ -990,9 +984,6 @@ impl ToolRegistryBuilder {
if let Some(memory_backend) = ctx.memory_backend {
resources.insert(memory_backend);
}
if let Some(auth_provider) = ctx.auth_provider.clone() {
resources.insert(auth_provider);
}
if let Ok(client) = crate::implementations::web_search::client::WebSearchClient::new(
&ctx.web_search_config,
ctx.api_key_provider.clone(),
@@ -1265,7 +1256,7 @@ impl FinalizedToolset {
)),
resources_persistence: Arc::new(ResourcesPersistence::noop()),
scheduler_cancel: None,
local_registry: kigi_computer_hub_sdk::LocalRegistry::new(),
local_registry: kigi_tool_runtime::LocalRegistry::new(),
renderer: Arc::new(TemplateRenderer::new(
std::collections::HashMap::new(),
std::collections::HashMap::new(),
@@ -1274,9 +1265,6 @@ impl FinalizedToolset {
workspace_viewer_ctx: None,
}
}
pub fn local_registry(&self) -> &kigi_computer_hub_sdk::LocalRegistry {
&self.local_registry
}
/// Get all tool definitions to send to the client.
pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
self.tools
@@ -2018,7 +2006,6 @@ mod tests {
app_builder_deployer_config:
crate::implementations::grok_build::deploy_app::AppBuilderDeployerConfig::default(),
api_key_provider: None,
auth_provider: None,
attribution_callback: None,
system_reminder_tag: crate::reminders::DEFAULT_REMINDER_TAG,
}
crates/codegen/kigi-tui/src/acp/mod.rs
-3
View File
@@ -175,9 +175,6 @@ pub async fn connect(cancel: &CancellationToken, flags: ConnectFlags) -> Result<
if let Some(effort) = flags.reasoning_effort_override {
agent_config.reasoning_effort_override = Some(effort);
}
// Agent connect intentionally leaves hub URL unset; provider hub is
// WorkspaceStartArgs only.
if !flags.permission_rules.is_empty() {
agent_config.cli_agent_overrides.permission_rules = flags.permission_rules.clone();
}
crates/codegen/kigi-tui/src/app/cli.rs
+4 -73
View File
@@ -110,12 +110,6 @@ See ~/.kigi/README.md for more information.
},
/// Manage git worktrees
Worktree(crate::worktree_cmd::WorktreeArgs),
/// Expose this workspace to the Computer Hub (via the leader).
///
/// Disabled by default and enabled server-side per account; set
/// `KIGI_WORKSPACE_COMMAND=1` to enable it locally for testing.
#[command(hide = true)]
Workspace(WorkspaceMgmtArgs),
/// Open the Agent Dashboard view at startup.
///
/// Centralised, agent-native overview of every session (top-level and
@@ -138,7 +132,7 @@ pub struct WrapArgs {
)]
pub command: Vec<String>,
}
/// Targets a running leader process by PID (used by `kigi leader` / `kigi workspace`).
/// Targets a running leader process by PID (used by `kigi leader`).
#[derive(Debug, clap::Args, Clone, Default)]
pub struct LeaderTargetArgs {
/// Leader process ID from `kigi leader list`.
@@ -169,69 +163,6 @@ pub enum LeaderMgmtCommand {
/// Stop all running leader processes
Kill,
}
#[derive(Debug, clap::Args, Clone)]
pub struct WorkspaceMgmtArgs {
#[command(subcommand)]
pub command: WorkspaceMgmtCommand,
}
#[derive(Debug, Subcommand, Clone)]
pub enum WorkspaceMgmtCommand {
/// Start (or update) the workspace→hub exposure.
Start(WorkspaceStartArgs),
/// Drain and disconnect from the hub, keeping the exposure warm.
Pause {
#[command(flatten)]
target: LeaderTargetArgs,
/// Emit machine-readable JSON output.
#[arg(long)]
json: bool,
},
/// Reconnect a paused exposure to the hub.
Resume {
#[command(flatten)]
target: LeaderTargetArgs,
/// Emit machine-readable JSON output.
#[arg(long)]
json: bool,
},
/// Stop exposing the workspace (the leader keeps running).
Stop {
#[command(flatten)]
target: LeaderTargetArgs,
/// Emit machine-readable JSON output.
#[arg(long)]
json: bool,
},
/// Restart the exposure (stop, then start with the given options).
Restart(WorkspaceStartArgs),
/// Show the current workspace-exposure status.
#[command(visible_alias = "list")]
Status {
#[command(flatten)]
target: LeaderTargetArgs,
/// Emit machine-readable JSON output.
#[arg(long)]
json: bool,
},
}
#[derive(Debug, clap::Args, Clone)]
pub struct WorkspaceStartArgs {
/// Computer Hub WebSocket URL (default: `[hub].url`, then the prod hub).
#[arg(long, value_name = "URL")]
pub hub_url: Option<String>,
/// Workspace root directory to expose. Defaults to the current directory.
#[arg(long, value_name = "DIR", value_hint = ValueHint::DirPath)]
pub cwd: Option<PathBuf>,
/// Force leader mode for this command, overriding config.
#[arg(long, conflicts_with = "no_leader")]
pub leader: bool,
/// Refuse to start even when config enables leader mode.
#[arg(long, conflicts_with = "leader")]
pub no_leader: bool,
/// Emit machine-readable JSON output.
#[arg(long)]
pub json: bool,
}
/// Arguments for the `agent` subcommand.
#[derive(Debug, clap::Args, Clone)]
pub struct AgentArgs {
@@ -276,9 +207,9 @@ pub struct AgentArgs {
/// Override the CLI chat proxy base URL.
#[arg(long = "coding-api-base-url")]
pub coding_api_base_url: Option<String>,
/// Override the public xAI API base URL.
#[arg(long = "xai-api-base-url")]
pub xai_api_base_url: Option<String>,
/// Override the direct (BYOK / external-API-key) API base URL.
#[arg(long = "api-base-url")]
pub api_base_url: Option<String>,
/// Agent runtime mode. Optional: bare `kigi agent` (and the `kigi acp`
/// alias) default to stdio.
#[command(subcommand)]
crates/codegen/kigi-tui/src/app/effects/tests.rs
+1 -1
View File
@@ -1619,7 +1619,7 @@ fn sanitize_user_error_strips_auth_prefixes() {
#[test]
fn sanitize_user_error_collapses_disk_full() {
assert_eq!(
sanitize_user_error("couldn't create worktree: Internal error: \"hub error: Worktree creation failed: not enough free disk space\""),
sanitize_user_error("couldn't create worktree: Internal error: \"workspace error: Worktree creation failed: not enough free disk space\""),
"Out of disk space."
);
assert_eq!(
crates/codegen/kigi-tui/src/app/leader_cluster/mod.rs
+1 -1
View File
@@ -318,7 +318,7 @@ impl PagerLeaderCluster {
let env = vec![
crate::test_util::EnvVarGuard::set("KIGI_SHARE_DIR", kigi_home.path()),
crate::test_util::EnvVarGuard::set("KIGI_CODE_BASE_URL", server.url()),
crate::test_util::EnvVarGuard::set("KIGI_XAI_API_BASE_URL", server.url()),
crate::test_util::EnvVarGuard::set("KIGI_API_BASE_URL", server.url()),
crate::test_util::EnvVarGuard::set("XAI_API_KEY", "test-key-for-ci"),
crate::test_util::EnvVarGuard::set("KIGI_TELEMETRY_ENABLED", "false"),
crate::test_util::EnvVarGuard::set("KIGI_FEEDBACK_ENABLED", "false"),
crates/codegen/kigi-tui/src/app/mod.rs
-1
View File
@@ -48,7 +48,6 @@ pub use cli::{
AgentArgs, AgentCmd, Command, LeaderArgs, LeaderMgmtArgs, LeaderMgmtCommand, LeaderTargetArgs,
OutputFormat, PagerArgs, ServeArgs, WrapArgs,
};
pub use cli::{WorkspaceMgmtArgs, WorkspaceMgmtCommand, WorkspaceStartArgs};
use crossterm::cursor::{self, SetCursorStyle};
use crossterm::event;
use crossterm::execute;
crates/codegen/kigi-tui/src/headless.rs
-3
View File
@@ -888,9 +888,6 @@ pub async fn run_single_turn(
None,
);
// No agent-level hub client URL (gateway-only cloud; workspace provider
// hub_url lives on `grok workspace` / WorkspaceStartArgs only).
apply_agent_flag(&options.agent, &mut agent_config);
if let Some(ref json) = options.agents_json {
crates/codegen/kigi-tui/src/trace_cmd.rs
+2 -2
View File
@@ -54,7 +54,7 @@ pub fn build_session_tar(session_dir: &Path, session_id: &str) -> Result<Vec<u8>
let metadata = ExportMetadata {
session_id: session_id.to_owned(),
grok_version: env!("VERSION_WITH_COMMIT").to_owned(),
kigi_version: env!("VERSION_WITH_COMMIT").to_owned(),
os: std::env::consts::OS.to_owned(),
arch: std::env::consts::ARCH.to_owned(),
exported_at: chrono::Utc::now().to_rfc3339(),
@@ -86,7 +86,7 @@ pub fn build_session_tar(session_dir: &Path, session_id: &str) -> Result<Vec<u8>
#[derive(serde::Serialize)]
struct ExportMetadata {
session_id: String,
grok_version: String,
kigi_version: String,
os: String,
arch: String,
exported_at: String,
crates/codegen/kigi-workspace/Cargo.toml
+1 -36
View File
@@ -8,9 +8,6 @@ description = "Core host-local workspace library (FS, VCS, execution, discovery)
[dependencies]
anyhow = { workspace = true }
arc-swap = { workspace = true }
# Diagnostics HTTP server (in-guest readiness/status endpoint); pinned to the
# workspace version already used by the sibling preview proxy.
axum = { workspace = true }
dunce = { workspace = true }
kigi-version = { workspace = true }
async-stream = { workspace = true }
@@ -41,7 +38,6 @@ uuid = { workspace = true, features = ["v4", "v5", "v7"] }
kigi-agent = { path = "../kigi-agent" }
kigi-tools = { path = "../kigi-tools" }
kigi-tools-api = { path = "../kigi-tools-api" }
kigi-workspace-client = { path = "../kigi-workspace-client" }
kigi-workspace-types = { path = "../kigi-workspace-types" }
kigi-config = { workspace = true }
# Leaf config value types (RemoteSettings, BoolFlag) for the folder-trust decision.
@@ -73,53 +69,28 @@ glob = "0.3"
kigi-sandbox = { path = "../kigi-sandbox", default-features = false }
kigi-hooks = { path = "../kigi-hooks" }
kigi-hunk-tracker = { path = "../kigi-hunk-tracker" }
# `metrics` enables the SDK's metric-donation client (periodic Prometheus
# registry gather → OTLP → hub donation pump); see metric_donation_reporter.
kigi-computer-hub-sdk = { workspace = true, features = ["metrics"] }
kigi-computer-hub-mcp-adapter = { path = "../../common/kigi-computer-hub-mcp-adapter" }
kigi-mcp = { path = "../kigi-mcp" }
kigi-file-utils = { path = "../kigi-file-utils" }
kigi-auth = { path = "../kigi-auth" }
kigi-log = { workspace = true }
kigi-tty-utils = { workspace = true }
kigi-sqlite-journal = { workspace = true }
reqwest = { workspace = true }
kigi-tool-protocol = { workspace = true }
kigi-tool-runtime = { workspace = true }
kigi-tool-types = { workspace = true }
tokio-util = { workspace = true }
urlencoding = "2"
kigi-fast-worktree = { path = "../kigi-fast-worktree", features = ["metadata"] }
# tonic stays for `tonic::Status`/`Code` mapping in workspace_ops deploy errors.
tonic = { workspace = true }
kigi-fsnotify = { path = "../kigi-fsnotify" }
clap = { workspace = true }
tracing-subscriber = { workspace = true }
# "enable" required for the SDK's spans to record at all.
fastrace = { workspace = true, features = ["enable"] }
kigi-tracing = { workspace = true }
rustls = { version = "0.23", default-features = false, features = ["ring", "logging", "std", "tls12"] }
tokio-tungstenite = { workspace = true, features = ["rustls-tls-webpki-roots"] }
tempfile = { workspace = true }
zstd = { workspace = true }
# Only referenced by the Unix self-daemonize path (fork/setsid/dup2/chdir).
# Only referenced by the Unix foreign-session capability probe (O_DIRECTORY etc.).
[target.'cfg(unix)'.dependencies]
libc = { workspace = true }
# Only referenced by the Windows self-daemonize path (SetStdHandle).
[target.'cfg(windows)'.dependencies]
windows = { workspace = true }
[[bin]]
name = "kigi-workspace-server"
path = "src/bin/workspace_server.rs"
[[bin]]
name = "workspace-server-probe"
path = "src/bin/workspace_server_probe.rs"
[features]
default = ["sandbox-enforce"]
compression = []
@@ -137,12 +108,6 @@ filetime = { workspace = true }
tempfile = { workspace = true }
tokio = { workspace = true, features = ["test-util"] }
kigi-test-utils = { path = "../../common/kigi-test-utils" }
# Decode tar.gz archives produced by workspace upload helpers in tests.
flate2 = { workspace = true }
tar = { workspace = true }
[package.metadata.cargo-shear]
ignored = ["tokio-tungstenite"]
[lints]
workspace = true
crates/codegen/kigi-workspace/src/config.rs
-863
View File
@@ -1,6 +1,5 @@
//! Workspace and session configuration types.
use crate::capability::CapabilityMode;
use crate::hub::HubConfig;
use kigi_tools::registry::types::{SessionContext, ToolRegistryBuilder, ToolServerConfig};
use std::collections::HashMap;
use std::path::PathBuf;
@@ -97,605 +96,6 @@ pub trait SessionContextFactory: Send + Sync {
#[derive(Debug, Clone, Default)]
#[non_exhaustive]
pub struct MemoryConfig {}
/// Per-session toolset/capability selection from the `session.bind`
/// metadata. Absent fields fall back to the workspace default and `CapabilityMode::All`.
#[derive(Debug, Clone, Default)]
#[non_exhaustive]
pub struct WorkspaceBindConfig {
/// Named toolset preset from the wire. **Never resolved** (see
/// [`Self::resolve`]); parsed only so it can be logged.
pub preset: Option<String>,
/// Capability mode applied to the session's toolset.
pub capability_mode: Option<CapabilityMode>,
/// Fully-specified toolset in the runtime serde shape. Takes precedence
/// over `tools`.
pub tool_config: Option<ToolServerConfig>,
/// Per-user feature-flag bag. `None` on legacy payloads → tools
/// fall back to their safe defaults.
pub viewer_ctx: Option<kigi_tool_runtime::WorkspaceViewerContext>,
/// Initial auto-approve (YOLO) state. `None` on legacy payloads →
/// fail-closed (false).
pub yolo_mode: Option<bool>,
/// Plane-configured toolset in the gRPC wire shape. An empty list is
/// treated as unset (proto3 repeated default).
pub tools: Option<Vec<kigi_tools_api::ToolConfigEntry>>,
pub manifest_version: Option<String>,
pub manifest_hash: Option<String>,
/// Opt-in: forward `BackgroundTaskCompleted` system notifications for this session.
pub system_notifications: bool,
pub rpc_only: bool,
}
/// Outcome of resolving a [`WorkspaceBindConfig`]; lets callers fail closed
/// instead of widening to the default toolset. Deliberately has **no preset
/// arm** (see [`WorkspaceBindConfig::resolve`]).
#[derive(Debug)]
pub enum ResolvedToolset {
/// An explicit toolset (`tool_config` or `tools`).
Toolset(ResolvedTools),
/// No explicit toolset was specified and the workspace allows falling
/// back to its default catalog (local/CLI embedders only).
UseDefault,
/// No explicit toolset was specified and the workspace requires one
/// (sandbox-launched standalone servers) — fail closed.
MissingToolConfig,
/// `tools` entries were specified but at least one failed to convert.
InvalidToolConfig(kigi_tools::registry::proto_convert::ToolConfigEntryError),
}
/// A resolved toolset plus the pinned entries this binary could not serve.
#[derive(Debug)]
pub struct ResolvedTools {
pub toolset: ToolServerConfig,
/// Pinned `tools` ids unknown to this binary's registry, sorted. Always
/// empty for `tool_config` resolutions.
pub unserved_tool_ids: Vec<String>,
}
impl ResolvedTools {
/// A fully-served toolset (no divergence).
fn full(toolset: ToolServerConfig) -> Self {
Self {
toolset,
unserved_tool_ids: Vec::new(),
}
}
}
impl WorkspaceBindConfig {
/// Parse hub `session.bind` metadata. The envelope is the shared
/// [`kigi_tool_runtime::WorkspaceBindMetadata`] (same type the emitter
/// serializes); `tool_config` is a consumer-only raw escape hatch read
/// separately.
pub fn from_metadata(metadata: &serde_json::Value) -> Self {
let wire: kigi_tool_runtime::WorkspaceBindMetadata =
serde_json::from_value(metadata.clone()).unwrap_or_default();
Self {
preset: wire.preset,
capability_mode: wire
.capability_mode
.and_then(|s| serde_json::from_value(serde_json::Value::String(s)).ok()),
tool_config: metadata
.as_object()
.and_then(|obj| obj.get("tool_config"))
.and_then(|v| parse_field("tool_config", v)),
viewer_ctx: wire.viewer_ctx,
yolo_mode: wire.yolo_mode,
tools: Some(wire.tools).filter(|tools| !tools.is_empty()),
manifest_version: wire.manifest_version,
manifest_hash: wire.manifest_hash,
system_notifications: wire.system_notifications.unwrap_or(false),
rpc_only: wire.rpc_only,
}
}
/// Resolve the selected toolset.
///
/// Precedence: `tool_config` > `tools` (wire entries) > default/fail-closed.
/// Pinned `tools` are served per entry: ids `known_id` rejects are dropped
/// and reported in [`ResolvedTools::unserved_tool_ids`] instead of
/// silently falling back to a different toolset.
///
/// **Presets are never resolved** — a `preset` on the wire is logged and
/// ignored; only explicit `tools`/`tool_config` may select a toolset.
///
/// With `require_explicit_toolset` (sandbox standalone servers) a bind
/// without an explicit toolset fails closed instead of widening to the
/// binary's default catalog.
pub fn resolve(
&self,
known_id: &dyn Fn(&str) -> bool,
require_explicit_toolset: bool,
) -> ResolvedToolset {
if let Some(cfg) = &self.tool_config {
for (idx, tool) in cfg.tools.iter().enumerate() {
if let Err(err) = kigi_tools_api::config_validation::validate_name_override(
idx,
&tool.id,
tool.name_override.as_deref(),
) {
return ResolvedToolset::InvalidToolConfig(err);
}
}
return ResolvedToolset::Toolset(ResolvedTools::full(cfg.clone()));
}
if let Some(tools) = &self.tools {
let mut unserved_tool_ids: Vec<String> = Vec::new();
let mut served = Vec::with_capacity(tools.len());
for (idx, entry) in tools.iter().enumerate() {
if !known_id(&entry.id) {
unserved_tool_ids.push(entry.id.clone());
continue;
}
match kigi_tools::registry::proto_convert::tool_config_from_entry(
idx,
entry.clone(),
) {
Ok(tc) => served.push(tc),
Err(err) => return ResolvedToolset::InvalidToolConfig(err),
}
}
unserved_tool_ids.sort_unstable();
if !unserved_tool_ids.is_empty() {
tracing::warn!(
unserved = ? unserved_tool_ids, config_manifest_version = ? self
.manifest_version, running_version = kigi_version::VERSION,
"session.bind: serving known subset of pinned tools"
);
}
return ResolvedToolset::Toolset(ResolvedTools {
toolset: ToolServerConfig {
tools: served,
behavior_preset: None,
},
unserved_tool_ids,
});
}
if let Some(preset) = self.preset.as_deref() {
tracing::warn!(
preset,
"session.bind: toolset presets are not resolved by the workspace \
server; pass an explicit `tools` config"
);
}
if require_explicit_toolset {
ResolvedToolset::MissingToolConfig
} else {
ResolvedToolset::UseDefault
}
}
}
/// Parse a single bind-metadata field, ignoring (and logging) a malformed value.
fn parse_field<T: serde::de::DeserializeOwned>(name: &str, value: &serde_json::Value) -> Option<T> {
match serde_json::from_value(value.clone()) {
Ok(parsed) => Some(parsed),
Err(e) => {
tracing::warn!(
field = name, error = % e,
"session.bind metadata: ignoring malformed field"
);
None
}
}
}
#[cfg(test)]
mod bind_config_tests {
use super::*;
/// Predicate for tests where every pinned id is known to the binary.
fn all_known(_: &str) -> bool {
true
}
/// Predicate for tests simulating a binary that knows none of the ids.
fn none_known(_: &str) -> bool {
false
}
#[test]
fn parses_preset_and_capability() {
let v = serde_json::json!(
{ "preset" : "explore", "capability_mode" : "read_only" }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert_eq!(cfg.preset.as_deref(), Some("explore"));
assert_eq!(
cfg.capability_mode,
Some(crate::capability::CapabilityMode::ReadOnly)
);
}
#[test]
fn defaults_on_empty_or_mismatched_metadata() {
let empty = WorkspaceBindConfig::from_metadata(&serde_json::json!({}));
assert!(empty.preset.is_none());
assert!(empty.capability_mode.is_none());
assert!(matches!(
empty.resolve(&all_known, false),
ResolvedToolset::UseDefault
));
let weird = WorkspaceBindConfig::from_metadata(&serde_json::json!("hello"));
assert!(matches!(
weird.resolve(&all_known, false),
ResolvedToolset::UseDefault
));
}
/// Presets are banned: any preset (known or not) is ignored — never
/// resolved to a toolset, and never widened to the default in strict mode.
#[test]
fn presets_are_never_resolved() {
for preset in ["explore", "grok-computer", "bogus"] {
let cfg = WorkspaceBindConfig::from_metadata(&serde_json::json!({ "preset" : preset }));
assert!(
matches!(cfg.resolve(&all_known, false), ResolvedToolset::UseDefault),
"lax mode must fall through to the default, preset={preset}"
);
assert!(
matches!(
cfg.resolve(&all_known, true),
ResolvedToolset::MissingToolConfig
),
"strict mode must fail closed, preset={preset}"
);
}
}
/// Strict mode (sandbox standalone server): no explicit toolset on the
/// bind ⇒ fail closed instead of widening to the default catalog.
#[test]
fn strict_mode_requires_explicit_toolset() {
let empty = WorkspaceBindConfig::from_metadata(&serde_json::json!({}));
assert!(matches!(
empty.resolve(&all_known, true),
ResolvedToolset::MissingToolConfig
));
}
#[test]
fn malformed_field_does_not_discard_valid_siblings() {
let v = serde_json::json!(
{ "preset" : "explore", "capability_mode" : "raed_only" }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert_eq!(cfg.preset.as_deref(), Some("explore"));
assert!(cfg.capability_mode.is_none());
}
#[test]
fn workspace_bind_config_from_metadata_extracts_viewer_ctx() {
let v = serde_json::json!(
{ "preset" : "explore", "viewer_ctx" : { "stream_tool_progress" : true }, }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert_eq!(cfg.preset.as_deref(), Some("explore"));
let viewer = cfg.viewer_ctx.expect("viewer_ctx parsed");
assert!(viewer.stream_tool_progress);
}
/// Legacy payload without `viewer_ctx` still parses (mixed-version
/// proxy/workspace deploys).
#[test]
fn workspace_bind_config_from_metadata_legacy_omitted_viewer_ctx() {
let v = serde_json::json!({ "preset" : "explore" });
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert!(cfg.viewer_ctx.is_none());
}
#[test]
fn workspace_bind_config_from_metadata_extracts_yolo_mode() {
let v = serde_json::json!({ "preset" : "explore", "yolo_mode" : true });
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert_eq!(cfg.yolo_mode, Some(true));
}
#[test]
fn workspace_bind_config_yolo_mode_omitted_or_malformed_fails_closed() {
let omitted =
WorkspaceBindConfig::from_metadata(&serde_json::json!({ "preset" : "explore" }));
assert!(omitted.yolo_mode.is_none());
let malformed = WorkspaceBindConfig::from_metadata(
&serde_json::json!({ "preset" : "explore", "yolo_mode" : "yes" }),
);
assert!(malformed.yolo_mode.is_none());
assert_eq!(malformed.preset.as_deref(), Some("explore"));
}
#[test]
fn workspace_bind_config_extracts_system_notifications_flag() {
let on = WorkspaceBindConfig::from_metadata(
&serde_json::json!({ "system_notifications" : true }),
);
assert!(on.system_notifications);
let off = WorkspaceBindConfig::from_metadata(&serde_json::json!({ "preset" : "explore" }));
assert!(!off.system_notifications);
let explicit_off = WorkspaceBindConfig::from_metadata(
&serde_json::json!({ "system_notifications" : false }),
);
assert!(!explicit_off.system_notifications);
}
#[test]
fn workspace_bind_config_extracts_rpc_only_flag() {
let on = WorkspaceBindConfig::from_metadata(&serde_json::json!({ "rpc_only" : true }));
assert!(on.rpc_only);
let off = WorkspaceBindConfig::from_metadata(&serde_json::json!({ "preset" : "explore" }));
assert!(!off.rpc_only);
let explicit_off =
WorkspaceBindConfig::from_metadata(&serde_json::json!({ "rpc_only" : false }));
assert!(!explicit_off.rpc_only);
}
#[test]
fn workspace_bind_config_from_metadata_extracts_manifest_fields() {
let v = serde_json::json!(
{ "preset" : "explore", "manifest_version" : "v1", "manifest_hash" :
"abc123", }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert_eq!(cfg.manifest_version.as_deref(), Some("v1"));
assert_eq!(cfg.manifest_hash.as_deref(), Some("abc123"));
}
#[test]
fn workspace_bind_config_manifest_fields_default_to_none_when_absent() {
let cfg = WorkspaceBindConfig::from_metadata(&serde_json::json!({ "preset" : "explore" }));
assert!(cfg.manifest_version.is_none());
assert!(cfg.manifest_hash.is_none());
}
/// Consumer-side parity test for the bind-metadata `tools` contract;
/// pairs with the producer-side pin test in agentic-sampler's
/// `configs::plane` tests.
#[test]
fn tools_entries_resolve_to_tool_server_config() {
let v = serde_json::json!(
{ "preset" : "explore", "tools" : [{ "id" : "GrokBuild:grep", "params_json" :
"{\"max_results\":50}", "name_override" : "search", "params_name_overrides" :
{ "pattern" : "query" }, "behavior_version" : "legacy-0.4.10",
"description_override" : "Search the codebase", }, { "id" :
"GrokBuild:read_file" },], }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&all_known, false) else {
panic!("tools entries must resolve to an explicit toolset");
};
assert!(resolved.unserved_tool_ids.is_empty());
let toolset = resolved.toolset;
assert_eq!(
toolset.behavior_preset, None,
"always the 'current' default"
);
assert_eq!(toolset.tools.len(), 2);
let grep = &toolset.tools[0];
assert_eq!(grep.id, "GrokBuild:grep");
assert_eq!(
grep.params,
serde_json::json!({ "max_results" : 50 })
.as_object()
.cloned()
);
assert_eq!(grep.name_override.as_deref(), Some("search"));
assert_eq!(
grep.params_name_overrides.as_ref().unwrap()["pattern"],
"query"
);
assert_eq!(grep.behavior_version.as_deref(), Some("legacy-0.4.10"));
assert_eq!(
grep.description_override.as_deref(),
Some("Search the codebase")
);
assert_eq!(grep.kind, None);
assert_eq!(toolset.tools[1].id, "GrokBuild:read_file");
}
#[test]
fn explicit_tool_config_wins_over_tools_entries() {
let v = serde_json::json!(
{ "tool_config" : { "tools" : [{ "id" : "raw:tool" }] }, "tools" : [{ "id" :
"wire:tool" }], }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&all_known, false) else {
panic!("must resolve to a toolset");
};
assert_eq!(resolved.toolset.tools.len(), 1);
assert_eq!(resolved.toolset.tools[0].id, "raw:tool");
}
#[test]
fn tools_entries_win_even_with_preset_present() {
let v = serde_json::json!(
{ "preset" : "explore", "tools" : [{ "id" : "wire:tool" }], }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&all_known, false) else {
panic!("must resolve to a toolset");
};
assert_eq!(resolved.toolset.tools.len(), 1);
assert_eq!(resolved.toolset.tools[0].id, "wire:tool");
}
#[test]
fn empty_tools_array_is_treated_as_unset() {
let v = serde_json::json!({ "preset" : "explore", "tools" : [] });
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert!(cfg.tools.is_none());
assert!(matches!(
cfg.resolve(&all_known, false),
ResolvedToolset::UseDefault
));
assert!(matches!(
cfg.resolve(&all_known, true),
ResolvedToolset::MissingToolConfig
));
let no_preset = serde_json::json!({ "tools" : [] });
let cfg = WorkspaceBindConfig::from_metadata(&no_preset);
assert!(matches!(
cfg.resolve(&all_known, false),
ResolvedToolset::UseDefault
));
}
#[test]
fn invalid_tools_entry_fails_closed() {
let v = serde_json::json!(
{ "preset" : "explore", "tools" : [{ "id" : "bad:tool", "params_json" :
"{not json" }], }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
match cfg.resolve(&all_known, false) {
ResolvedToolset::InvalidToolConfig(err) => {
assert_eq!(err.tool_id, "bad:tool");
assert_eq!(err.index, 0);
}
other => panic!("expected InvalidToolConfig, got {other:?}"),
}
}
#[test]
fn invalid_name_override_fails_closed() {
let v = serde_json::json!(
{ "tools" : [{ "id" : "wire:ok", "name_override" : "fine_name" }, { "id" :
"wire:bad", "name_override" : "not a tool id!" },], }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
match cfg.resolve(&all_known, true) {
ResolvedToolset::InvalidToolConfig(err) => {
assert_eq!(err.tool_id, "wire:bad");
assert_eq!(err.field_path(), "tools[1].name_override");
}
other => panic!("expected InvalidToolConfig, got {other:?}"),
}
}
#[test]
fn tool_config_escape_hatch_invalid_name_override_fails_closed() {
let v = serde_json::json!(
{ "tool_config" : { "tools" : [{ "id" : "raw:ok", "name_override" :
"fine_name" }, { "id" : "raw:bad", "name_override" : "not a tool id!" },] },
}
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
match cfg.resolve(&all_known, true) {
ResolvedToolset::InvalidToolConfig(err) => {
assert_eq!(err.tool_id, "raw:bad");
assert_eq!(err.field_path(), "tools[1].name_override");
}
other => panic!("expected InvalidToolConfig, got {other:?}"),
}
let v = serde_json::json!(
{ "tool_config" : { "tools" : [{ "id" : "raw:ok", "name_override" :
"fine_name" }] }, }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&all_known, true) else {
panic!("valid escape-hatch config must resolve");
};
assert_eq!(resolved.toolset.tools.len(), 1);
}
#[test]
fn invalid_entry_error_reports_wire_index_after_unknown_drop() {
let v = serde_json::json!(
{ "tools" : [{ "id" : "wire:unknown" }, { "id" : "wire:bad", "params_json" :
"{not json" },], }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let known = |id: &str| id != "wire:unknown";
match cfg.resolve(&known, false) {
ResolvedToolset::InvalidToolConfig(err) => {
assert_eq!(err.tool_id, "wire:bad");
assert_eq!(
err.index, 1,
"index must be the wire position, not the known-subset position"
);
}
other => panic!("expected InvalidToolConfig, got {other:?}"),
}
}
#[test]
fn valid_name_overrides_resolve_intact() {
let v = serde_json::json!(
{ "tools" : [{ "id" : "wire:a", "name_override" : "renamed_a" }, { "id" :
"wire:b" },], }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&all_known, true) else {
panic!("well-formed overrides must resolve to a toolset");
};
assert_eq!(resolved.toolset.tools.len(), 2);
assert_eq!(
resolved.toolset.tools[0].name_override.as_deref(),
Some("renamed_a")
);
assert_eq!(resolved.toolset.tools[1].name_override, None);
}
#[test]
fn pinned_tools_all_known_serves_full_expansion() {
let v = serde_json::json!(
{ "preset" : "explore", "tools" : [{ "id" : "wire:tool" }],
"manifest_version" : "9.9.9-any", }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&all_known, false) else {
panic!("known pinned tools must use the tools expansion");
};
assert!(resolved.unserved_tool_ids.is_empty());
assert_eq!(resolved.toolset.tools.len(), 1);
assert_eq!(resolved.toolset.tools[0].id, "wire:tool");
}
/// Unknown ids must be partitioned and reported, never silently replaced
/// by live preset resolution.
#[test]
fn pinned_tools_unknown_ids_are_partitioned_and_reported() {
let v = serde_json::json!(
{ "preset" : "explore", "tools" : [{ "id" : "wire:known" }, { "id" :
"wire:zz_unknown" }, { "id" : "wire:aa_unknown" },], "manifest_version" :
"0.0.0-stale", }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let known = |id: &str| id == "wire:known";
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&known, false) else {
panic!("partial coverage must still resolve to the known subset");
};
assert_eq!(resolved.toolset.tools.len(), 1);
assert_eq!(resolved.toolset.tools[0].id, "wire:known");
assert_eq!(
resolved.unserved_tool_ids,
vec!["wire:aa_unknown".to_owned(), "wire:zz_unknown".to_owned()],
"unserved ids are reported sorted"
);
}
/// A fully-unknown expansion serves empty and reports every id — it never
/// widens to preset/default.
#[test]
fn pinned_tools_all_unknown_serves_empty_and_reports_all() {
let v = serde_json::json!(
{ "preset" : "explore", "tools" : [{ "id" : "wire:tool" }],
"manifest_version" : "0.0.0-stale", }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&none_known, false) else {
panic!("all-unknown expansion must resolve (empty), not fall back");
};
assert!(resolved.toolset.tools.is_empty());
assert_eq!(resolved.unserved_tool_ids, vec!["wire:tool".to_owned()]);
}
#[test]
fn legacy_tools_without_manifest_version_are_not_gated() {
let v = serde_json::json!(
{ "preset" : "explore", "tools" : [{ "id" : "wire:tool" }], }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert!(cfg.manifest_version.is_none());
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&all_known, false) else {
panic!("legacy unpinned tools must resolve without gating");
};
assert_eq!(resolved.toolset.tools.len(), 1);
assert_eq!(resolved.toolset.tools[0].id, "wire:tool");
}
#[test]
fn tool_config_wins_regardless_of_stale_manifest_version() {
let v = serde_json::json!(
{ "tool_config" : { "tools" : [{ "id" : "raw:tool" }] }, "tools" : [{ "id" :
"wire:tool" }], "manifest_version" : "0.0.0-stale", }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&none_known, false) else {
panic!("tool_config must always win");
};
assert!(resolved.unserved_tool_ids.is_empty());
assert_eq!(resolved.toolset.tools.len(), 1);
assert_eq!(resolved.toolset.tools[0].id, "raw:tool");
}
#[test]
fn malformed_tools_field_is_dropped_keeping_siblings() {
let v = serde_json::json!({ "preset" : "explore", "tools" : "not-a-list" });
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert!(cfg.tools.is_none());
assert!(matches!(
cfg.resolve(&all_known, false),
ResolvedToolset::UseDefault
));
assert!(matches!(
cfg.resolve(&all_known, true),
ResolvedToolset::MissingToolConfig
));
}
}
/// Top-level config required to construct a [`crate::handle::WorkspaceHandle`].
///
/// `#[non_exhaustive]` so future fields are non-breaking.
@@ -726,18 +126,6 @@ pub struct WorkspaceConfig {
/// and disabled/enabled lists. Stored on `WorkspaceShared` for
/// `discover_plugins` calls. Defaults to empty.
pub plugin_discovery_config: crate::discovery::PluginDiscoveryConfig,
/// Optional server configuration. When `Some`, the workspace
/// can connect to the server after construction via
/// [`WorkspaceHandle::connect_hub`](crate::handle::WorkspaceHandle::connect_hub).
pub hub_config: Option<HubConfig>,
/// Auth provider for xAI service calls made from workspace-scoped code.
/// `None` for workspaces that do not configure service auth.
pub auth_provider: Option<kigi_computer_hub_sdk::SharedAuthProvider>,
/// Metadata attached to the tool server registration.
/// Propagated through the server to `ServerInfo.metadata` in
/// `servers.list` responses so harness clients can identify the
/// sandbox that started the tool server.
pub server_metadata: Option<serde_json::Value>,
/// Runtime-tunable timing/threshold config for the tool server.
pub status_config: crate::status_config::StatusConfig,
/// Folder-trust verdict for repo-local (project-scoped) LSP servers from
@@ -745,108 +133,12 @@ pub struct WorkspaceConfig {
/// shell caller resolves the verdict and threads it in; callers without a
/// folder-trust decision pass `true`.
pub project_lsp_trusted: bool,
/// Fail `session.bind`s without an explicit toolset closed instead of
/// widening to `default_tool_config`. Set by sandbox-launched standalone
/// servers; local/CLI embedders keep the default-catalog fallback.
pub require_explicit_toolset: bool,
/// Confine `x.ai/fs/*` / `workspace.fs_*` resolution to the workspace root
/// (reject `..`, absolute-outside-root, symlink escapes). Default `false`
/// (unconfined) — set to `true` only by the workspace server on a remote
/// sandbox, where the root is a real tenant boundary.
pub confine_fs_to_workspace_root: bool,
}
/// Metadata a tool server announces so hub consumers can identify and route
/// to it. Every field is optional and independently sourced; a local process
/// announces none.
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
pub struct WorkspaceServerMetadata {
/// Sandbox that provisioned this server. Absent for local servers.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub sandbox_id: Option<String>,
/// Logical sandbox-service session UUID, from the `KIGI_SESSION_ID` env
/// var. Present whenever that var is set (every sandbox container, start
/// and restore), absent otherwise.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub session_id: Option<String>,
/// Provider that provisioned this server. Populated on the start path
/// only (no container-side source on restore); absent for local servers.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_id: Option<String>,
/// Per-spawn launch nonce minted by the sandbox orchestrator and echoed
/// verbatim on the diagnostics `/ready` endpoint. Absent for local/legacy
/// launches.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub launch_id: Option<String>,
}
impl WorkspaceServerMetadata {
/// Merge an env-sourced logical session id into caller-supplied
/// tool-server metadata (`None` on the restore/local path).
///
/// `env_session_id` is the raw `KIGI_SESSION_ID`; empty is normalized to
/// absent. An explicit `session_id` already in `metadata` is never
/// clobbered. A non-object `metadata` value is returned unchanged (a
/// defensive no-op — the sole caller always sends an object).
pub fn merge_session_metadata(
metadata: Option<serde_json::Value>,
env_session_id: Option<String>,
) -> Option<serde_json::Value> {
let env_session_id = env_session_id.filter(|s| !s.is_empty());
match metadata {
Some(mut value) => {
if let Some(session_id) = env_session_id
&& let Some(obj) = value.as_object_mut()
&& !obj.contains_key("session_id")
{
obj.insert(
"session_id".to_owned(),
serde_json::Value::String(session_id),
);
}
Some(value)
}
None => serde_json::to_value(WorkspaceServerMetadata {
sandbox_id: None,
session_id: env_session_id,
provider_id: None,
launch_id: None,
})
.ok(),
}
}
}
impl WorkspaceConfig {
/// Construct a minimal config suitable for proxy-mode workspaces
/// where the workspace is used primarily as a ToolServer host.
pub fn new_for_proxy(
root_cwd: PathBuf,
session_factory: Arc<dyn SessionContextFactory>,
hub_config: HubConfig,
auth_provider: kigi_computer_hub_sdk::SharedAuthProvider,
server_metadata: Option<serde_json::Value>,
status_config: crate::status_config::StatusConfig,
tool_config: ToolServerConfig,
) -> Self {
Self {
root_cwd,
default_tool_config: tool_config,
respect_gitignore: false,
memory_config: None,
event_buffer_capacity: crate::config::DEFAULT_EVENT_BUFFER_CAPACITY,
session_factory,
hook_global_sources: vec![],
hook_project_sources: vec![],
skills_config: Default::default(),
plugin_discovery_config: Default::default(),
auth_provider: Some(auth_provider),
hub_config: Some(hub_config),
server_metadata,
project_lsp_trusted: true,
require_explicit_toolset: false,
confine_fs_to_workspace_root: false,
status_config,
}
}
}
/// Configuration for spawning a subagent session within a workspace.
#[derive(Clone)]
#[non_exhaustive]
@@ -928,158 +220,3 @@ pub enum IsolationMode {
/// Run the subagent inside a sandbox/container.
Sandbox,
}
#[cfg(test)]
mod tests {
use super::WorkspaceServerMetadata;
#[test]
fn workspace_server_metadata_serializes_all_present_fields() {
let meta = WorkspaceServerMetadata {
sandbox_id: Some("sb-123".to_owned()),
session_id: Some("11111111-1111-1111-1111-111111111111".to_owned()),
provider_id: Some("test-provider".to_owned()),
launch_id: Some("33333333-3333-3333-3333-333333333333".to_owned()),
};
let value = serde_json::to_value(&meta).unwrap();
assert_eq!(
value,
serde_json::json!({ "sandbox_id" : "sb-123", "session_id" :
"11111111-1111-1111-1111-111111111111", "provider_id" : "test-provider",
"launch_id" : "33333333-3333-3333-3333-333333333333", })
);
}
#[test]
fn workspace_server_metadata_omits_none_fields() {
let meta = WorkspaceServerMetadata {
sandbox_id: Some("sb-123".to_owned()),
session_id: None,
provider_id: None,
launch_id: None,
};
let value = serde_json::to_value(&meta).unwrap();
assert_eq!(value, serde_json::json!({ "sandbox_id" : "sb-123" }));
let empty = serde_json::to_value(WorkspaceServerMetadata::default()).unwrap();
assert_eq!(empty, serde_json::json!({}));
}
#[test]
fn workspace_server_metadata_deserializes_legacy_payload_without_new_fields() {
let legacy = serde_json::json!(
{ "sandbox_id" : "sb-legacy", "cwd" : "/workspace", "mode" : "remote", }
);
let meta: WorkspaceServerMetadata = serde_json::from_value(legacy).unwrap();
assert_eq!(meta.sandbox_id.as_deref(), Some("sb-legacy"));
assert_eq!(meta.session_id, None);
assert_eq!(meta.provider_id, None);
}
#[test]
fn workspace_server_metadata_round_trips_with_new_fields() {
let meta = WorkspaceServerMetadata {
sandbox_id: Some("sb-123".to_owned()),
session_id: Some("22222222-2222-2222-2222-222222222222".to_owned()),
provider_id: Some("test-provider".to_owned()),
launch_id: None,
};
let json = serde_json::to_string(&meta).unwrap();
let back: WorkspaceServerMetadata = serde_json::from_str(&json).unwrap();
assert_eq!(back.sandbox_id, meta.sandbox_id);
assert_eq!(back.session_id, meta.session_id);
assert_eq!(back.provider_id, meta.provider_id);
}
#[test]
fn workspace_server_metadata_deserializes_partial_new_fields() {
let only_session = serde_json::json!(
{ "sandbox_id" : "sb-1", "session_id" :
"33333333-3333-3333-3333-333333333333", }
);
let meta: WorkspaceServerMetadata = serde_json::from_value(only_session).unwrap();
assert_eq!(
meta.session_id.as_deref(),
Some("33333333-3333-3333-3333-333333333333")
);
assert_eq!(meta.provider_id, None);
let only_provider = serde_json::json!(
{ "sandbox_id" : "sb-1", "provider_id" : "test-provider", }
);
let meta: WorkspaceServerMetadata = serde_json::from_value(only_provider).unwrap();
assert_eq!(meta.provider_id.as_deref(), Some("test-provider"));
assert_eq!(meta.session_id, None);
}
#[test]
fn workspace_server_metadata_reads_start_path_shaped_payload() {
let start_path = serde_json::json!(
{ "cwd" : "/workspace", "mode" : "remote", "sandbox_id" : "sb-start",
"session_id" : "44444444-4444-4444-4444-444444444444", "provider_id" :
"test-provider", }
);
let meta: WorkspaceServerMetadata = serde_json::from_value(start_path).unwrap();
assert_eq!(meta.sandbox_id.as_deref(), Some("sb-start"));
assert_eq!(
meta.session_id.as_deref(),
Some("44444444-4444-4444-4444-444444444444")
);
assert_eq!(meta.provider_id.as_deref(), Some("test-provider"));
}
#[test]
fn merge_session_metadata_builds_struct_from_env_on_none_branch() {
let merged =
WorkspaceServerMetadata::merge_session_metadata(None, Some("sess-1".to_owned()))
.unwrap();
assert_eq!(merged, serde_json::json!({ "session_id" : "sess-1" }));
let empty = WorkspaceServerMetadata::merge_session_metadata(None, None).unwrap();
assert_eq!(empty, serde_json::json!({}));
}
#[test]
fn merge_session_metadata_overlays_into_object_without_clobbering() {
let base = serde_json::json!({ "sandbox_id" : "sb-9", "mode" : "remote" });
let merged =
WorkspaceServerMetadata::merge_session_metadata(Some(base), Some("env-id".to_owned()))
.unwrap();
assert_eq!(
merged,
serde_json::json!({ "sandbox_id" : "sb-9", "mode" : "remote",
"session_id" : "env-id", })
);
let explicit = serde_json::json!({ "session_id" : "explicit" });
let merged = WorkspaceServerMetadata::merge_session_metadata(
Some(explicit),
Some("env-id".to_owned()),
)
.unwrap();
assert_eq!(merged, serde_json::json!({ "session_id" : "explicit" }));
}
#[test]
fn merge_session_metadata_leaves_object_untouched_when_no_env_id() {
let base = serde_json::json!({ "sandbox_id" : "sb-9" });
let merged =
WorkspaceServerMetadata::merge_session_metadata(Some(base.clone()), None).unwrap();
assert_eq!(merged, base);
}
#[test]
fn merge_session_metadata_non_object_is_returned_unchanged() {
let scalar = serde_json::json!("just-a-string");
let merged = WorkspaceServerMetadata::merge_session_metadata(
Some(scalar.clone()),
Some("env-id".to_owned()),
)
.unwrap();
assert_eq!(merged, scalar);
}
#[test]
fn merge_session_metadata_treats_empty_env_id_as_absent() {
let none_branch =
WorkspaceServerMetadata::merge_session_metadata(None, Some(String::new())).unwrap();
assert_eq!(none_branch, serde_json::json!({}));
let base = serde_json::json!({ "sandbox_id" : "sb-9" });
let overlay = WorkspaceServerMetadata::merge_session_metadata(
Some(base.clone()),
Some(String::new()),
)
.unwrap();
assert_eq!(overlay, base);
}
#[test]
fn workspace_server_metadata_rejects_wrong_typed_field() {
let bad = serde_json::json!({ "sandbox_id" : "sb-1", "session_id" : 42 });
let result: Result<WorkspaceServerMetadata, _> = serde_json::from_value(bad);
assert!(result.is_err());
}
}
crates/codegen/kigi-workspace/src/error.rs
+3 -3
View File
@@ -53,9 +53,9 @@ pub enum WorkspaceError {
#[error("hunk action failed: {0}")]
HunkActionFailed(String),
/// An error from the server connection or tool server.
#[error("hub error: {0}")]
HubError(String),
/// An internal workspace error.
#[error("workspace error: {0}")]
Internal(String),
/// Deploy-service error tagged with its gRPC status class; see
/// [`DeployError`] for how the class crosses the workspace RPC boundary.
crates/codegen/kigi-workspace/src/file_system/ext_fs.rs
+8 -8
View File
@@ -62,7 +62,7 @@ impl WorkspaceOp for FsListReq {
let req = self.clone();
tokio::task::spawn_blocking(move || list(&abs, &req, confine_root))
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?
.map_err(|e| WorkspaceError::Internal(e.to_string()))?
}
}
@@ -98,7 +98,7 @@ impl WorkspaceOp for FsReadFileReq {
if !ranged {
let bytes = tokio::fs::read(&abs)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
return Ok(build_file_entry(&bytes));
}
@@ -106,9 +106,9 @@ impl WorkspaceOp for FsReadFileReq {
// chunk is `[offset, offset + min(length, max_bytes, cap))`.
let md = tokio::fs::metadata(&abs)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
if md.is_dir() {
return Err(WorkspaceError::HubError(format!(
return Err(WorkspaceError::Internal(format!(
"not a file: {}",
self.path
)));
@@ -120,7 +120,7 @@ impl WorkspaceOp for FsReadFileReq {
let length = super::walk::clamp_read_length(self.length, self.max_bytes);
let chunk = super::walk::read_range(&abs, offset, length)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
Ok(build_ranged_entry(chunk, size, self.encoding))
}
}
@@ -143,8 +143,8 @@ impl WorkspaceOp for FsWriteFileReq {
std::fs::write(&abs, content.as_bytes())
})
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
.map_err(|e| WorkspaceError::Internal(e.to_string()))?
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
Ok(())
}
}
@@ -160,7 +160,7 @@ impl WorkspaceOp for FsDeleteFileReq {
let (abs, _) = ws.confine_to_workspace_root(&abs_unconfined).await?;
tokio::fs::remove_file(&abs)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
Ok(())
}
}
crates/codegen/kigi-workspace/src/file_system/mod.rs
+1 -6
View File
@@ -7,13 +7,8 @@ pub use ext_fs::{
FsReadFileReq, FsWriteFileReq,
};
// Client-facing read-only fs ops (`workspace.client_fs_*`). Not re-exported:
// its wire types live in `kigi_workspace_types::rpc::fs` (the `ClientFs*`
// types) and would collide with the shell-facing `ext_fs` names above.
pub(crate) mod client_fs;
// Shared filesystem core: paginated listing + binary-safe ranged reads,
// used by `client_fs`, `ext_fs`, and the shell-local `session::file_system`.
// used by `ext_fs` and the shell-local `session::file_system`.
mod walk;
pub use walk::{
ChunkPayload, ListOptions, ListPage, ListedEntry, MAX_LIST_COLLECT, MAX_READ_BYTES,
crates/codegen/kigi-workspace/src/handle.rs
+45 -3930
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-workspace/src/lib.rs
+7 -29
View File
@@ -10,8 +10,6 @@ pub mod activity;
pub mod capability;
pub mod channel;
pub mod config;
pub mod daemonize;
pub mod diag_server;
pub mod discovery;
pub mod envrc;
pub mod error;
@@ -20,17 +18,8 @@ pub mod folder_trust;
pub mod foreign_sessions;
pub mod fs_notify;
pub mod handle;
pub mod hub;
pub mod hub_auth;
pub mod hub_channel;
pub mod hub_ids;
pub mod hub_server;
pub mod mcp;
pub mod permission;
pub mod preview_supervisor;
pub mod project_config;
pub mod recovery;
pub mod rpc_envelope;
pub mod session;
pub mod status_config;
pub use status_config::StatusConfig;
@@ -46,25 +35,18 @@ pub use config::{
};
pub use error::{WorkspaceError, WorkspaceResult};
pub use file_system::*;
pub use handle::{
DrainOutcome, DrainReason, WorkspaceHandle, connect_local_workspace, resolve_workspace_home,
termination_grace_from_env,
};
pub use hub::HubConfig;
pub use handle::WorkspaceHandle;
pub use kigi_hunk_tracker::HunkTrackerHandle;
pub use kigi_workspace_client::WorkspaceClient;
pub use kigi_workspace_types::WorkspaceEvent;
pub use permission::*;
pub use session::{WorkspaceSession, WorkspaceShared};
pub use session::{file_state, git, jj};
pub use workspace_ops::{WorkspaceOp, WorkspaceOps};
/// Zero-init every workspace metric family so idle panels render a `0` baseline
/// instead of "No data". Idempotent; call once at workspace-server startup.
/// instead of "No data". Idempotent; call once at startup.
pub fn init_metrics() {
handle::init_metrics();
session::swap_policy::init_metrics();
permission::init_metrics();
hub_server::init_metrics();
}
/// Crate-wide lock serializing every test that mutates the process-global
/// environment (`KIGI_SHARE_DIR`, `HOME`, …). nextest isolates each test in its own
@@ -167,22 +149,18 @@ mod init_metrics_tests {
})
})
};
assert!(has(
"grok_workspace_rpc_requests_total",
&[("method", "unknown"), ("result", "error")]
));
assert!(has(
"grok_workspace_drain_started_total",
&[("reason", "sigterm")]
));
assert!(has(
"grok_workspace_toolset_swap_rejected_total",
&[("reason", "turn_active"), ("trigger", "update_tool_config")]
));
assert!(has(
"grok_workspace_rewind_checkpoint_capture_total",
&[("domain", "fs"), ("outcome", "completed")]
));
assert!(
families
.iter()
.any(|mf| mf.name() == "grok_workspace_permission_timeout_total")
.any(|mf| mf.name() == "grok_workspace_terminal_backend_orphaned_total")
);
}
}
crates/codegen/kigi-workspace/src/permission/manager.rs
+4 -228
View File
@@ -820,6 +820,10 @@ fn session_grant_pre_decision(
/// Spawns the permission manager actor, returning a handle and the telemetry
/// event receiver.
///
/// `remember_tool_approvals` — resolved gate: shows the per-tool always-allow
/// options and lets an explicit grant satisfy an `ask` rule (ask once, remember).
#[allow(clippy::too_many_arguments)]
pub fn spawn_permission_manager(
session_id: acp::SessionId,
gateway: GatewaySender,
@@ -833,45 +837,7 @@ pub fn spawn_permission_manager(
web_fetch_allowed_domains: Vec<String>,
initial_yolo: bool,
client_identifier: Option<String>,
) -> (PermissionHandle, mpsc::UnboundedReceiver<PermissionEvent>) {
spawn_permission_manager_with_hub(
session_id,
gateway,
cwd,
client_type,
permission_config,
deny_read_globs,
web_fetch_allowed_domains,
initial_yolo,
client_identifier,
// Legacy/test entry point: preserve the full option set. Production uses
// `spawn_permission_manager_with_hub` with the resolved gate.
true,
None,
)
}
/// Like [`spawn_permission_manager`] but routes the permission prompt to chat
/// over the server (the HITL live path) when `hub_permission` is `Some`. The
/// caller builds the transport only when [`hitl_permission_live_enabled`] and a
/// server is connected; `None` keeps the local ACP prompt.
///
/// [`hitl_permission_live_enabled`]: crate::permission::hitl_permission_live_enabled
#[allow(clippy::too_many_arguments)]
pub fn spawn_permission_manager_with_hub(
session_id: acp::SessionId,
gateway: GatewaySender,
cwd: AbsPathBuf,
client_type: ClientType,
permission_config: Option<crate::permission::types::PermissionConfig>,
deny_read_globs: Vec<String>,
web_fetch_allowed_domains: Vec<String>,
initial_yolo: bool,
client_identifier: Option<String>,
// Resolved `remember_tool_approvals` gate: shows the per-tool always-allow
// options and lets an explicit grant satisfy an `ask` rule (ask once, remember).
remember_tool_approvals: bool,
hub_permission: Option<Arc<dyn crate::permission::PermissionHookTransport>>,
) -> (PermissionHandle, mpsc::UnboundedReceiver<PermissionEvent>) {
// Read the pin ONCE (file I/O) and cache it; never re-read per tool-call.
// Every yolo ingestion path funnels through construction or SetYoloMode.
@@ -887,7 +853,6 @@ pub fn spawn_permission_manager_with_hub(
client_identifier,
remember_tool_approvals,
crate::permission::resolution::yolo_disabled_by_policy(),
hub_permission,
)
}
@@ -905,7 +870,6 @@ fn spawn_permission_manager_with_pin(
client_identifier: Option<String>,
remember_tool_approvals: bool,
yolo_pin: Option<&'static str>,
hub_permission: Option<Arc<dyn crate::permission::PermissionHookTransport>>,
) -> (PermissionHandle, mpsc::UnboundedReceiver<PermissionEvent>) {
let (tx, mut rx) = mpsc::unbounded_channel::<PermissionCommand>();
let (event_tx, event_rx) = mpsc::unbounded_channel::<PermissionEvent>();
@@ -969,7 +933,6 @@ fn spawn_permission_manager_with_pin(
}
let prompter = AcpPrompter::new(session_id.clone(), gateway.clone(), client_type)
.with_hub_permission(hub_permission)
.with_remember_tool_approvals(remember_tool_approvals);
let mut yolo_mode = initial_yolo;
let mut auto_mode = seed_auto;
@@ -1849,7 +1812,6 @@ mod tests {
None,
true,
yolo_pin,
None,
)
}
@@ -1871,7 +1833,6 @@ mod tests {
None,
true,
None,
None,
)
}
@@ -1939,187 +1900,6 @@ mod tests {
.await;
}
/// Like [`test_manager`] but routes prompts through a hub permission transport.
fn test_manager_with_hub(
cwd: &AbsPathBuf,
hub_permission: Arc<dyn crate::permission::PermissionHookTransport>,
) -> (PermissionHandle, mpsc::UnboundedReceiver<PermissionEvent>) {
let (tx, _rx) = mpsc::unbounded_channel();
spawn_permission_manager_with_pin(
acp::SessionId::new(Arc::from("test-session")),
GatewaySender::new(tx),
cwd.clone(),
ClientType::Generic,
None,
vec![],
vec![],
false,
None,
true,
None,
Some(hub_permission),
)
}
/// Records every emitted payload and replies with a canned decision, so the
/// hub permission prompt path is exercised without a live hub.
struct FakeHubTransport {
reply: serde_json::Value,
seen: std::sync::Mutex<Vec<serde_json::Value>>,
}
#[async_trait::async_trait]
impl crate::permission::PermissionHookTransport for FakeHubTransport {
async fn request_permission(
&self,
payload: serde_json::Value,
) -> Result<serde_json::Value, String> {
self.seen.lock().unwrap().push(payload);
Ok(self.reply.clone())
}
}
fn fake_hub(reply: serde_json::Value) -> Arc<FakeHubTransport> {
Arc::new(FakeHubTransport {
reply,
seen: std::sync::Mutex::new(Vec::new()),
})
}
#[tokio::test]
async fn hub_permission_approve_allows_and_emits_payload() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let tmp = tempfile::tempdir().unwrap();
let cwd = AbsPathBuf::new(tmp.path().to_path_buf()).unwrap();
let transport = fake_hub(serde_json::json!({ "outcome": "approve" }));
let (mgr, _e) = test_manager_with_hub(&cwd, transport.clone());
let d = mgr
.request(
AccessKind::Edit("src/main.rs".into()),
tool_call(),
None,
None,
None,
)
.await;
assert_eq!(d, Decision::Allow);
let seen = transport.seen.lock().unwrap();
assert_eq!(seen.len(), 1, "exactly one permission hook emitted");
assert_eq!(seen[0]["tool_call_id"], "tc");
assert_eq!(seen[0]["tool_name"], "search_replace");
assert_eq!(seen[0]["description"], "Edit src/main.rs");
assert_eq!(seen[0]["scope"], "write");
assert_eq!(
seen[0]["edit_file_paths"],
serde_json::json!(["src/main.rs"])
);
})
.await;
}
#[tokio::test]
async fn hub_permission_reject_aborts() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let tmp = tempfile::tempdir().unwrap();
let cwd = AbsPathBuf::new(tmp.path().to_path_buf()).unwrap();
let (mgr, _e) = test_manager_with_hub(
&cwd,
fake_hub(serde_json::json!({ "outcome": "reject" })),
);
let d = mgr
.request(
AccessKind::Edit("a.rs".into()),
tool_call(),
None,
None,
None,
)
.await;
assert!(
matches!(d, Decision::Reject(_)),
"reject must abort, got {d:?}"
);
})
.await;
}
/// `cancelled` reply (turn-end drain) → abort, distinct from a user reject.
#[tokio::test]
async fn hub_permission_cancelled_aborts_distinctly() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let tmp = tempfile::tempdir().unwrap();
let cwd = AbsPathBuf::new(tmp.path().to_path_buf()).unwrap();
let (mgr, _e) = test_manager_with_hub(
&cwd,
fake_hub(serde_json::json!({ "outcome": "cancelled" })),
);
let d = mgr
.request(
AccessKind::Edit("a.rs".into()),
tool_call(),
None,
None,
None,
)
.await;
assert_eq!(d, Decision::Cancelled);
})
.await;
}
#[tokio::test]
async fn hub_permission_always_approve_persists_scope() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let tmp = tempfile::tempdir().unwrap();
let cwd = AbsPathBuf::new(tmp.path().to_path_buf()).unwrap();
let transport = fake_hub(serde_json::json!({
"outcome": "always_approve",
"scope": { "kind": "server_prefix", "value": "linear" },
}));
let (mgr, _e) = test_manager_with_hub(&cwd, transport.clone());
let first = mgr
.request(
AccessKind::MCPTool {
name: "linear__list".into(),
input: serde_json::Value::Null,
},
tool_call(),
None,
None,
None,
)
.await;
assert_eq!(first, Decision::Allow);
let second = mgr
.request(
AccessKind::MCPTool {
name: "linear__create".into(),
input: serde_json::Value::Null,
},
tool_call(),
None,
None,
None,
)
.await;
assert_eq!(second, Decision::Allow);
assert_eq!(
transport.seen.lock().unwrap().len(),
1,
"always_approve must persist so the second call needs no hook"
);
})
.await;
}
/// A managed `Ask` rule on a direct `Read`/`Grep` must reach the prompt, not
/// the unconditional auto-allow. With no responder wired, that surfaces as a
/// non-`Allow` decision; a non-ask read still auto-allows.
@@ -2484,7 +2264,6 @@ mod tests {
None,
true,
None,
None,
);
assert_eq!(
handle.deny_read_globs(),
@@ -2654,7 +2433,6 @@ mod tests {
None,
remember_tool_approvals,
None,
None,
)
}
@@ -3386,7 +3164,6 @@ mod tests {
None,
true,
None,
None,
);
let PermissionHandle::Actor { ref cmd_tx, .. } = mgr else {
panic!("manager must be actor-backed");
@@ -3577,7 +3354,6 @@ mod tests {
None,
true,
None,
None,
);
// Request A parks in the gated prompt; B then arrives and overlaps it.
crates/codegen/kigi-workspace/src/permission/mod.rs
+1 -14
View File
@@ -1,6 +1,5 @@
pub mod auto_mode;
pub mod claude_settings;
mod hub_permission;
mod manager;
mod policy;
mod prompter;
@@ -19,19 +18,7 @@ pub use auto_mode::{
classifier_output_json_schema, default_auto_mode_classifier, is_auto_mode_allowlisted_access,
is_auto_mode_allowlisted_tool_name, parse_classifier_model_text, permission_decision_args,
};
pub use hub_permission::{
PermissionHookTransport, ToolServerPermissionTransport, access_kind_for_hub_tool,
hitl_permission_live_enabled, prompt_outcome_allows, request_permission_via_hub,
};
/// Zero-init this module's metric families. See [`crate::init_metrics`].
pub(crate) fn init_metrics() {
hub_permission::init_metrics();
}
pub use manager::{
PermissionHandle, default_always_allow_scope, spawn_permission_manager,
spawn_permission_manager_with_hub,
};
pub use manager::{PermissionHandle, default_always_allow_scope, spawn_permission_manager};
pub use policy::CompiledPolicy;
pub use prompter::{
ALLOW_EDITS_SESSION_OPTION_ID, AcpPrompter, BashCommandPermission, BashCommandSelectedTerms,
crates/codegen/kigi-workspace/src/permission/prompter.rs
+23 -49
View File
@@ -321,9 +321,6 @@ pub struct AcpPrompter {
/// at decision-time through it. `EventWriter::noop()` when events recording
/// is disabled (the default for the permission scaffolding's own tests).
event_writer: EventWriter,
/// Server permission transport: when set, [`request`](Self::request) asks chat for the
/// decision over the server; `None` keeps the local prompt.
hub_permission: Option<Arc<dyn crate::permission::PermissionHookTransport>>,
/// When `false` (default, fail-safe), the per-tool "Always allow …" options
/// are stripped (see [`REMEMBER_TOOL_APPROVALS_GATED_IDS`]).
remember_tool_approvals: bool,
@@ -496,7 +493,6 @@ impl AcpPrompter {
// must NOT double-emit. A workspace-server-side caller that owns the
// per-session `events.jsonl` opts in via [`with_event_writer`].
event_writer: EventWriter::noop(),
hub_permission: None,
// Fail-safe default; opt in via `with_remember_tool_approvals`.
remember_tool_approvals: false,
}
@@ -509,16 +505,6 @@ impl AcpPrompter {
self
}
/// Route the permission prompt to chat over the server when `Some`;
/// `None` keeps the local prompt.
pub fn with_hub_permission(
mut self,
hub_permission: Option<Arc<dyn crate::permission::PermissionHookTransport>>,
) -> Self {
self.hub_permission = hub_permission;
self
}
/// Attach a per-session `events.jsonl` writer so [`request`](Self::request)
/// records `PermissionRequested` / `PermissionResolved`. Used by the
/// workspace-server permission path (which owns the session log); the shell
@@ -737,41 +723,29 @@ impl AcpPrompter {
prompt_start,
};
let outcome = match &self.hub_permission {
// Route the prompt to chat over the server (see
// `ToolServerPermissionTransport` for the await/release contract).
Some(transport) => {
crate::permission::hub_permission::request_permission_via_hub(
transport.as_ref(),
access,
tool_call_update.tool_call_id.0.as_ref(),
)
.await
}
None => {
let permission_options = self.build_options(access);
let req = acp::RequestPermissionRequest::new(
self.session_id.clone(),
tool_call_update.clone(),
permission_options.values().cloned().collect(),
)
.meta(self.bash_selection_meta(access));
match self.gateway.request_permission(req).await {
Ok(resp) => match resp.outcome {
acp::RequestPermissionOutcome::Cancelled => PromptOutcome::Cancelled,
acp::RequestPermissionOutcome::Selected(selected) => map_selected_outcome(
&permission_options,
&selected.option_id,
resp.meta.as_ref(),
access,
),
// TODO(acp-0.10): `RequestPermissionOutcome` is #[non_exhaustive].
_ => PromptOutcome::Error("unknown permission outcome".to_owned()),
},
Err(e) => {
tracing::error!(?e, "failed to request permission");
PromptOutcome::Error("failed to request permission".to_owned())
}
let outcome = {
let permission_options = self.build_options(access);
let req = acp::RequestPermissionRequest::new(
self.session_id.clone(),
tool_call_update.clone(),
permission_options.values().cloned().collect(),
)
.meta(self.bash_selection_meta(access));
match self.gateway.request_permission(req).await {
Ok(resp) => match resp.outcome {
acp::RequestPermissionOutcome::Cancelled => PromptOutcome::Cancelled,
acp::RequestPermissionOutcome::Selected(selected) => map_selected_outcome(
&permission_options,
&selected.option_id,
resp.meta.as_ref(),
access,
),
// TODO(acp-0.10): `RequestPermissionOutcome` is #[non_exhaustive].
_ => PromptOutcome::Error("unknown permission outcome".to_owned()),
},
Err(e) => {
tracing::error!(?e, "failed to request permission");
PromptOutcome::Error("failed to request permission".to_owned())
}
}
};
crates/codegen/kigi-workspace/src/session/mod.rs
+3 -122
View File
@@ -8,12 +8,8 @@ pub mod tool_config;
use crate::capability::CapabilityMode;
use crate::config::{MemoryConfig, SessionContextFactory};
use crate::file_system::{AsyncFsWrapper, LocalFs};
use crate::hub::{HubConfig, HubHandle};
use crate::session::file_state::FileStateTracker;
use kigi_computer_hub_mcp_adapter::McpBridgeHandle;
use kigi_hunk_tracker::HunkTrackerHandle;
use kigi_mcp::servers::McpState;
use kigi_tool_protocol::ToolId;
use kigi_tool_runtime::WorkspaceViewerContext;
use kigi_tools::notification::types::{ToolNotification, ToolNotificationHandle};
use kigi_tools::registry::types::{FinalizedToolset, ToolConfig, ToolServerConfig};
@@ -76,12 +72,6 @@ pub struct WorkspaceSession {
inner: RwLock<WorkspaceSessionInner>,
/// Per-session lock that serialises `update_tool_config` calls.
pub(crate) update_lock: tokio::sync::Mutex<()>,
/// Per-session MCP state (owned clients, etc.).
pub(crate) mcp_state: Arc<tokio::sync::Mutex<McpState>>,
/// MCP bridges kept alive for the session lifetime.
pub(crate) mcp_bridges: tokio::sync::Mutex<Vec<McpBridgeHandle>>,
/// Qualified tool IDs registered on the server for this session's MCP tools.
pub(crate) mcp_tool_ids: tokio::sync::Mutex<Vec<ToolId>>,
/// Per-user feature-flag bag resolved at session-bind time, frozen for
/// the session lifetime. `None` → tools use their safe defaults.
pub(crate) viewer_ctx: Option<WorkspaceViewerContext>,
@@ -107,8 +97,7 @@ pub struct WorkspaceSession {
/// created (or last rebound) with. `None` when the session was resolved
/// from the workspace default (no explicit toolset in the bind metadata).
/// Lets a rebind detect a config change and re-resolve instead of silently
/// reusing a stale toolset (e.g. a session created by a metadata-less
/// hub revive bind that a config-carrying client rebind must correct).
/// reusing a stale toolset.
bind_tool_config_fingerprint: std::sync::Mutex<Option<serde_json::Value>>,
/// The last snapshot-driven rebuild failed and kept a stale toolset;
/// cleared by any successful install. While set, an identical-config
@@ -190,9 +179,6 @@ impl WorkspaceSession {
update_lock: tokio::sync::Mutex::new(()),
bind_tool_config_fingerprint: std::sync::Mutex::new(None),
stale_resolve: std::sync::atomic::AtomicBool::new(false),
mcp_state: Arc::new(tokio::sync::Mutex::new(McpState::new(vec![]))),
mcp_bridges: tokio::sync::Mutex::new(Vec::new()),
mcp_tool_ids: tokio::sync::Mutex::new(Vec::new()),
viewer_ctx,
yolo_mode: std::sync::atomic::AtomicBool::new(false),
system_notifications,
@@ -440,9 +426,6 @@ pub type ClientExtSink = std::sync::Arc<dyn Fn(String, serde_json::Value) + Send
/// Workspace-wide shared state.
pub struct WorkspaceShared {
pub(crate) default_tool_config: ToolServerConfig,
/// Require an explicit toolset on every `session.bind`; see
/// [`crate::config::WorkspaceConfig::require_explicit_toolset`].
pub(crate) require_explicit_toolset: bool,
/// See [`crate::config::WorkspaceConfig::confine_fs_to_workspace_root`].
/// Default `false`; enabled only for remote-sandbox workspace servers.
pub(crate) confine_fs_to_workspace_root: bool,
@@ -464,38 +447,16 @@ pub struct WorkspaceShared {
/// disabled/enabled lists). Used by `discover_plugins` via the
/// `discovery` module.
pub(crate) plugin_discovery_config: crate::discovery::PluginDiscoveryConfig,
/// Live server connection handle. `None` until
/// [`WorkspaceHandle::connect_hub`](crate::handle::WorkspaceHandle::connect_hub)
/// is called (or if no [`HubConfig`] was provided).
///
/// Uses `tokio::sync::Mutex` so the guard can be held across the
/// async `HubHandle::connect()` call, preventing TOCTOU races.
pub(crate) hub_handle: tokio::sync::Mutex<Option<HubHandle>>,
/// Remote-origin tool configs (consumer direction), updated by the
/// notification listener.
pub(crate) hub_tools_snapshot: arc_swap::ArcSwap<Vec<ToolConfig>>,
/// Server config stashed at construction time for deferred connect.
pub(crate) hub_config: Option<HubConfig>,
/// Auth provider for xAI service calls.
pub(crate) auth_provider: Option<kigi_computer_hub_sdk::SharedAuthProvider>,
/// Connection-level sink feeding the `ActivityTracker` (drained by
/// `run_activity_feed`); not a network egress. `None` until `connect_hub()` sets it.
pub(crate) activity_notify_handle:
arc_swap::ArcSwap<Option<kigi_tools::notification::types::ToolNotificationHandle>>,
/// Sink for workspace-originated ext-notifications to the client (e.g.
/// `x.ai/search/fuzzy/status`). Mode-agnostic: the shell wires it to the
/// agent gateway in local mode, and to the server in proxy mode. `None` until
/// set via [`WorkspaceHandle::set_client_ext_sink`](crate::handle::WorkspaceHandle::set_client_ext_sink).
pub(crate) client_ext_sink: arc_swap::ArcSwap<Option<ClientExtSink>>,
pub(crate) local_registry: kigi_computer_hub_sdk::LocalRegistry,
pub(crate) local_registry: kigi_tool_runtime::LocalRegistry,
pub(crate) activity_tracker: std::sync::Arc<crate::activity::ActivityTracker>,
/// Runtime-tunable timing/threshold config for the tool server.
/// Read by the status publisher task and at shutdown.
pub(crate) status_config: crate::status_config::StatusConfig,
/// Opaque metadata for the tool server registration, forwarded verbatim to
/// the server; structured access goes through
/// [`WorkspaceShared::server_metadata_typed`].
pub(crate) server_metadata: Option<serde_json::Value>,
/// Workspace-level fuzzy search manager. Separate from the shell's
/// own `FuzzySearchManager` — this instance serves remote (hub/RPC)
/// clients.
@@ -533,7 +494,6 @@ pub struct WorkspaceShared {
/// turn start inside the check→install window deterministically.
#[cfg(test)]
pub(crate) post_resolve_test_hook: parking_lot::Mutex<Option<Box<dyn Fn() + Send + Sync>>>,
pub(crate) client_fs_hash_memo: crate::file_system::client_fs::FileHashMemo,
}
impl WorkspaceShared {
/// Workspace root directory.
@@ -578,40 +538,6 @@ impl WorkspaceShared {
.get(session_id)
.map(|w| w.value().clone())
}
/// Stable hub server id (`--server-id`), if a hub config is present.
pub(crate) fn server_id(&self) -> Option<String> {
self.hub_config.as_ref().and_then(|c| c.server_id.clone())
}
/// Auth provider used for xAI service calls.
pub fn auth_provider(&self) -> Option<&kigi_computer_hub_sdk::SharedAuthProvider> {
self.auth_provider.as_ref()
}
/// Parse the opaque [`server_metadata`](Self::server_metadata) blob into
/// the typed subset the workspace needs (currently `sandbox_id`);
/// unknown/missing fields default cleanly. A present-but-malformed blob is
/// logged and salvaged field-by-field (a bad sibling field must not
/// silently drop `sandbox_id` from every environment artifact).
pub(crate) fn server_metadata_typed(&self) -> crate::config::WorkspaceServerMetadata {
let Some(v) = self.server_metadata.as_ref() else {
return Default::default();
};
match serde_json::from_value(v.clone()) {
Ok(typed) => typed,
Err(e) => {
tracing::warn!(
error = % e,
"workspace: malformed server_metadata; salvaging sandbox_id field-wise"
);
crate::config::WorkspaceServerMetadata {
sandbox_id: v
.get("sandbox_id")
.and_then(serde_json::Value::as_str)
.map(str::to_owned),
..Default::default()
}
}
}
}
pub fn default_tool_config(&self) -> &ToolServerConfig {
&self.default_tool_config
}
@@ -624,49 +550,6 @@ impl WorkspaceShared {
pub fn mcp_tools_snapshot(&self) -> Arc<Vec<ToolConfig>> {
self.mcp_tools_snapshot.load_full()
}
/// The tool server, if a server connection is active.
///
/// Returns a clone of the [`ToolServer`](kigi_computer_hub_sdk::ToolServer)
/// which is cheap (`Arc` bump). Uses `try_lock` to avoid blocking
/// on the async mutex from synchronous contexts. Returns `None` if
/// the lock is held (i.e. a `connect_hub` call is in progress).
pub fn hub_server(&self) -> Option<kigi_computer_hub_sdk::ToolServer> {
self.hub_handle
.try_lock()
.ok()
.and_then(|guard| guard.as_ref().map(|h| h.server.clone()))
}
/// Like [`Self::hub_server`] but awaits the `hub_handle` lock instead of
/// returning `None` on contention. Use from async contexts that must not
/// confuse a transient `connect_hub` lock-hold with "no hub connected";
/// `None` means no hub is connected.
pub async fn hub_server_blocking(&self) -> Option<kigi_computer_hub_sdk::ToolServer> {
self.hub_handle
.lock()
.await
.as_ref()
.map(|h| h.server.clone())
}
/// Current snapshot of hub-provided tool configs (consumer direction).
pub fn hub_tools_snapshot(&self) -> Arc<Vec<ToolConfig>> {
self.hub_tools_snapshot.load_full()
}
/// Compose a session's tool `ctx.notification_handle` as a fan-out of the
/// connection-level activity feed (internal tracker accounting) and the
/// opt-in per-session `system.notify` sender. Only the `system.notify` leg
/// reaches a client, so the fan-out can't double-wake. `None` → factory default.
pub(crate) fn compose_session_notification_handle(
&self,
system_notify_handle: Option<ToolNotificationHandle>,
) -> Option<ToolNotificationHandle> {
let activity = self.activity_notify_handle.load_full().as_ref().clone();
match (activity, system_notify_handle) {
(None, None) => None,
(Some(a), None) => Some(a),
(None, Some(s)) => Some(s),
(Some(a), Some(s)) => Some(ToolNotificationHandle::tee(vec![a, s])),
}
}
pub fn activity_tracker(&self) -> &std::sync::Arc<crate::activity::ActivityTracker> {
&self.activity_tracker
}
@@ -717,7 +600,6 @@ impl WorkspaceShared {
};
let trigger = SwapTrigger::from_rebuild_source(source);
let mcp_snap = self.mcp_tools_snapshot.load_full();
let hub_snap = self.hub_tools_snapshot.load_full();
let sessions: Vec<(String, Arc<WorkspaceSession>)> = {
let guard = self.sessions.read();
guard
@@ -779,7 +661,6 @@ impl WorkspaceShared {
baseline,
session.capability_mode(),
&mcp_snap,
&hub_snap,
session.cwd().to_path_buf(),
session.session_env().clone(),
&sid,
@@ -787,7 +668,7 @@ impl WorkspaceShared {
Some(self.local_registry.clone()),
self.lsp.clone(),
session.viewer_ctx().cloned(),
self.compose_session_notification_handle(session.system_notify_handle()),
session.system_notify_handle(),
session.terminal_backend().clone(),
) {
Ok((effective, toolset)) => {
crates/codegen/kigi-workspace/src/session/tool_config.rs
+43 -389
View File
@@ -1,11 +1,10 @@
//! Tool config resolution pipeline.
//!
//! Five-step resolution:
//! Four-step resolution:
//! 1. `effective_tool_config = config.tool_config.unwrap_or_else(|| parent.effective_tool_config.clone())`
//! 2. `merged = merge_mcp_tools(effective_tool_config, shared.mcp_servers.snapshot())`
//! 3. `merged = merge_hub_tools(merged, shared.hub_tools_snapshot())`
//! 4. `filtered = config.capability_mode.filter(merged)`
//! 5. `toolset = build_finalized_toolset(filtered, &session.cwd, &session.session_env, ...)`
//! 3. `filtered = config.capability_mode.filter(merged)`
//! 4. `toolset = build_finalized_toolset(filtered, &session.cwd, &session.session_env, ...)`
use crate::capability::{CapabilityMode, kind_allowed};
use crate::config::SessionContextFactory;
use crate::error::{WorkspaceError, WorkspaceResult};
@@ -19,19 +18,16 @@ use std::sync::Arc;
/// Create-shaped entry of the resolution pipeline: run
/// [`resolve_session_toolset_rebuild`] around a FRESH factory-built
/// session-lifetime terminal backend, and return that backend so the caller
/// can store it on the session it is creating. Session-less resolves (the
/// `__template__` catalog resolve in `connect_hub`) also use this entry and
/// simply drop the returned backend with the toolset.
/// can store it on the session it is creating.
pub(crate) fn resolve_session_toolset(
effective_tool_config: ToolServerConfig,
capability_mode: CapabilityMode,
mcp_snapshot: &[ToolConfig],
hub_snapshot: &[ToolConfig],
cwd: PathBuf,
session_env: Arc<HashMap<String, String>>,
session_id: &str,
factory: &dyn SessionContextFactory,
local_registry: Option<kigi_computer_hub_sdk::LocalRegistry>,
local_registry: Option<kigi_tool_runtime::LocalRegistry>,
lsp: Option<std::sync::Arc<dyn kigi_tools::implementations::lsp::LspBackend>>,
viewer_ctx: Option<kigi_tool_runtime::WorkspaceViewerContext>,
notification_handle: Option<kigi_tools::notification::types::ToolNotificationHandle>,
@@ -45,7 +41,6 @@ pub(crate) fn resolve_session_toolset(
effective_tool_config,
capability_mode,
mcp_snapshot,
hub_snapshot,
cwd,
session_env,
session_id,
@@ -66,24 +61,23 @@ pub(crate) fn resolve_session_toolset(
///
/// Returns the *unmodified* `effective_tool_config` (step-1 baseline) so
/// the caller can store it on the session. The FinalizedToolset reflects
/// MCP + hub merging and capability filtering on top of that baseline.
/// MCP merging and capability filtering on top of that baseline.
///
/// **MCP-origin and hub-origin `kind: None` tools are dropped under
/// every non-`All` mode.** Baseline `kind: None` tools are always kept —
/// but before filtering, kind-less baseline entries whose id the binary's
/// registry knows get their [`ToolKind`] backfilled (see
/// [`backfill_tool_kinds`]), so the capability filter applies to pinned
/// server-bind toolsets whose wire entries cannot carry a kind.
/// **MCP-origin `kind: None` tools are dropped under every non-`All`
/// mode.** Baseline `kind: None` tools are always kept — but before
/// filtering, kind-less baseline entries whose id the binary's registry
/// knows get their [`ToolKind`] backfilled (see [`backfill_tool_kinds`]),
/// so the capability filter applies to pinned toolsets whose wire entries
/// cannot carry a kind.
pub(crate) fn resolve_session_toolset_rebuild(
effective_tool_config: ToolServerConfig,
capability_mode: CapabilityMode,
mcp_snapshot: &[ToolConfig],
hub_snapshot: &[ToolConfig],
cwd: PathBuf,
session_env: Arc<HashMap<String, String>>,
session_id: &str,
factory: &dyn SessionContextFactory,
local_registry: Option<kigi_computer_hub_sdk::LocalRegistry>,
local_registry: Option<kigi_tool_runtime::LocalRegistry>,
lsp: Option<std::sync::Arc<dyn kigi_tools::implementations::lsp::LspBackend>>,
viewer_ctx: Option<kigi_tool_runtime::WorkspaceViewerContext>,
notification_handle: Option<kigi_tools::notification::types::ToolNotificationHandle>,
@@ -94,24 +88,7 @@ pub(crate) fn resolve_session_toolset_rebuild(
builder = builder.with_local_registry(lr);
}
let baseline = backfill_tool_kinds(&effective_tool_config, &builder.known_tool_kinds());
let filtered = merge_and_filter(
&baseline,
mcp_snapshot,
hub_snapshot,
capability_mode,
session_id,
);
let hub_ids: std::collections::HashSet<&str> =
hub_snapshot.iter().map(|t| t.id.as_str()).collect();
let finalize_config = ToolServerConfig {
tools: filtered
.tools
.iter()
.filter(|t| !hub_ids.contains(t.id.as_str()))
.cloned()
.collect(),
behavior_preset: filtered.behavior_preset.clone(),
};
let finalize_config = merge_and_filter(&baseline, mcp_snapshot, capability_mode, session_id);
let mut ctx = factory.build_session_context(session_id, cwd, session_env, terminal_backend);
if let Some(lsp_handle) = lsp {
ctx.lsp = Some(lsp_handle);
@@ -157,22 +134,20 @@ fn backfill_tool_kinds(
behavior_preset: config.behavior_preset.clone(),
}
}
/// Steps 2-4 of the resolution pipeline, without step 5 (`finalize`):
/// Steps 2-3 of the resolution pipeline, without the `finalize` step:
///
/// - **Step 2** -- MCP merge: append MCP-origin tools, skipping ID/name collisions with baseline.
/// - **Step 3** -- Hub merge: append hub-origin tools, skipping ID/name collisions with baseline or MCP.
/// - **Step 4** -- Capability filter: drop tools whose `kind` is not allowed by the mode.
/// External (MCP/hub) `kind: None` tools are only kept under `CapabilityMode::All`.
/// - **Step 3** -- Capability filter: drop tools whose `kind` is not allowed by the mode.
/// External (MCP) `kind: None` tools are only kept under `CapabilityMode::All`.
///
/// Priority on ID/name collision: baseline wins > MCP wins > hub is skipped.
/// Priority on ID/name collision: baseline wins over MCP.
pub(crate) fn merge_and_filter(
baseline: &ToolServerConfig,
mcp_snapshot: &[ToolConfig],
hub_snapshot: &[ToolConfig],
mode: CapabilityMode,
session_id: &str,
) -> ToolServerConfig {
if mcp_snapshot.is_empty() && hub_snapshot.is_empty() {
if mcp_snapshot.is_empty() {
return mode.filter(baseline);
}
let baseline_ids: std::collections::HashSet<&str> =
@@ -187,7 +162,6 @@ pub(crate) fn merge_and_filter(
.collect();
let mut tagged: Vec<(ToolConfig, bool)> =
baseline.tools.iter().cloned().map(|t| (t, false)).collect();
let mut mcp_tool_ids: std::collections::HashSet<&str> = std::collections::HashSet::new();
for mcp_tool in mcp_snapshot {
if baseline_ids.contains(mcp_tool.id.as_str()) {
tracing::warn!(
@@ -205,35 +179,8 @@ pub(crate) fn merge_and_filter(
);
continue;
}
mcp_tool_ids.insert(mcp_tool.id.as_str());
tagged.push((mcp_tool.clone(), true));
}
for hub_tool in hub_snapshot {
if baseline_ids.contains(hub_tool.id.as_str()) {
tracing::debug!(
hub_id = % hub_tool.id, session = % session_id,
"skipping remote tool: id collides with baseline"
);
continue;
}
if mcp_tool_ids.contains(hub_tool.id.as_str()) {
tracing::debug!(
hub_id = % hub_tool.id, session = % session_id,
"skipping remote tool: id collides with MCP tool"
);
continue;
}
let client_name = hub_tool.resolve_client_name(&hub_tool.id);
if !taken_names.insert(client_name.clone()) {
tracing::debug!(
hub_id = % hub_tool.id, client_name = % client_name, session = %
session_id,
"skipping remote tool: resolved client name collides with another tool"
);
continue;
}
tagged.push((hub_tool.clone(), true));
}
let kept: Vec<ToolConfig> = tagged
.into_iter()
.filter(|(tool, is_external)| match tool.kind {
@@ -250,19 +197,13 @@ pub(crate) fn merge_and_filter(
}
/// Alias for backward compatibility.
pub type NoopSessionContextFactory = WorkspaceSessionContextFactory;
/// Whether per-session `tool_state.json` persistence + per-turn upload is
/// enabled (`KIGI_WORKSPACE_TOOL_STATE_ENABLED=true`; any other value keeps
/// legacy behavior).
pub fn tool_state_enabled() -> bool {
std::env::var("KIGI_WORKSPACE_TOOL_STATE_ENABLED").as_deref() == Ok("true")
}
/// Sanitize a `session_id` into a single safe filesystem path segment: chars
/// outside `[A-Za-z0-9_-]` become `_`, empty becomes `anon`. When any
/// replacement happened, an 8-hex digest of the ORIGINAL id is appended so the
/// mapping stays injective — plain substitution would collide distinct ids
/// (`sess/1` and `sess_1`) into one directory, cross-contaminating
/// persistence, rehydration, and [`crate::recovery::cleanup_stale_sessions`].
/// Already-safe ids (the common UUID case) map to themselves.
/// persistence and rehydration. Already-safe ids (the common UUID case)
/// map to themselves.
fn sanitize_session_id(session_id: &str) -> String {
let mut safe = String::with_capacity(session_id.len());
let mut modified = false;
@@ -297,16 +238,7 @@ fn ensure_session_dir(root: &std::path::Path, session_id: &str) -> (PathBuf, std
/// hazard is the global `environ` array, not the variable's value).
#[cfg(test)]
pub(crate) use crate::ENV_TEST_LOCK as TOOL_STATE_ENV_LOCK;
/// [`SessionContextFactory`] for workspace server sessions.
///
/// When constructed with an [`AuthProvider`] and API base URL, gen tools
/// (image_gen, video_gen) are enabled using the provider's current
/// OAuth token. Without auth, gen tools default to `Disabled`.
///
/// When [`with_tool_state_home`](Self::with_tool_state_home) is set, each
/// session's [`SessionContext::state_path`] is rooted at
/// `<home>/sessions/<session_id>/`; left unset, `state_path` stays empty
/// (legacy behavior).
/// [`SessionContextFactory`] for workspace sessions.
///
/// [`SessionContext::session_folder`] is `/tmp/sessions/<sanitized_id>/`
/// (terminal logs and other tool artifacts — not the project `cwd`).
@@ -318,64 +250,11 @@ pub(crate) use crate::ENV_TEST_LOCK as TOOL_STATE_ENV_LOCK;
/// [`build_terminal_backend`]: crate::config::SessionContextFactory::build_terminal_backend
/// [`build_session_context`]: crate::config::SessionContextFactory::build_session_context
/// [`LocalTerminalBackend`]: kigi_tools::computer::local::LocalTerminalBackend
pub struct WorkspaceSessionContextFactory {
auth: Option<kigi_computer_hub_sdk::SharedAuthProvider>,
api_base_url: Option<String>,
/// Resolved `$KIGI_WORKSPACE_HOME` when tool-state persistence is enabled;
/// `None` disables it. Resolved once by the caller so the factory performs
/// no per-build env reads.
tool_state_home: Option<PathBuf>,
}
impl Default for WorkspaceSessionContextFactory {
fn default() -> Self {
Self::new()
}
}
#[derive(Default)]
pub struct WorkspaceSessionContextFactory;
impl WorkspaceSessionContextFactory {
pub fn new() -> Self {
Self {
auth: None,
api_base_url: None,
tool_state_home: None,
}
}
/// Factory with auth — gen tools use the provider's live token.
pub fn with_auth(
auth: kigi_computer_hub_sdk::SharedAuthProvider,
api_base_url: String,
) -> Self {
Self {
auth: Some(auth),
api_base_url: Some(api_base_url),
tool_state_home: None,
}
}
/// Enable session-keyed tool-state persistence rooted at `home`
/// (`$KIGI_WORKSPACE_HOME`). Callers should only invoke this when
/// [`tool_state_enabled`] is `true`.
pub fn with_tool_state_home(mut self, home: PathBuf) -> Self {
self.tool_state_home = Some(home);
self
}
/// `<tool_state_home>/sessions/<sanitized_id>/tool_state.json`, or empty
/// when persistence is disabled / dir creation fails.
fn resolve_state_path(&self, session_id: &str) -> PathBuf {
let Some(home) = self.tool_state_home.as_ref() else {
return PathBuf::new();
};
let (dir, created) = ensure_session_dir(home, session_id);
if let Err(e) = created {
tracing::warn!(
session = % session_id, dir = % dir.display(), error = % e,
"tool_state: failed to create session dir; persistence disabled for session"
);
return PathBuf::new();
}
tracing::debug!(
session = % session_id, dir = % dir.display(),
"tool_state: persistence bound to session-keyed dir"
);
dir.join("tool_state.json")
Self
}
/// `/tmp/sessions/<sanitized_id>/` for terminal logs and other tool artifacts.
fn resolve_session_folder(session_id: &str) -> PathBuf {
@@ -397,59 +276,9 @@ impl SessionContextFactory for WorkspaceSessionContextFactory {
session_env: Arc<HashMap<String, String>>,
backend: Arc<dyn kigi_tools::computer::types::TerminalBackend>,
) -> kigi_tools::registry::types::SessionContext {
use kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig;
use kigi_tools::implementations::grok_build::image_gen::ImageGenConfig;
use kigi_tools::implementations::grok_build::video_gen::VideoGenConfig;
use kigi_tools::implementations::web_search::WebSearchConfig;
let fs = Arc::new(kigi_tools::computer::local::LocalFs)
as Arc<dyn kigi_tools::computer::types::AsyncFileSystem>;
let notification_handle = kigi_tools::notification::ToolNotificationHandle::noop();
let (image_gen_config, video_gen_config, web_search_config, app_builder_deployer_config) =
if let (Some(auth), Some(url)) = (&self.auth, &self.api_base_url) {
let cred = auth.current();
match cred {
kigi_computer_hub_sdk::AuthCredential::Bearer { token, .. } => {
let headers = build_proxy_headers(url);
(
ImageGenConfig::Enabled {
api_key: token.clone(),
base_url: url.clone(),
extra_headers: headers.clone(),
image_gen_enabled: true,
image_edit_enabled: true,
model_override: None,
tier_restricted: false,
},
VideoGenConfig::Enabled {
api_key: token.clone(),
base_url: url.clone(),
extra_headers: headers.clone(),
zdr_video_output_s3: None,
tier_restricted: false,
},
WebSearchConfig::Enabled {
search_url: format!("{}/search", url.trim_end_matches('/')),
api_key: token,
extra_headers: headers,
},
AppBuilderDeployerConfig::default(),
)
}
_ => (
ImageGenConfig::default(),
VideoGenConfig::default(),
WebSearchConfig::default(),
AppBuilderDeployerConfig::default(),
),
}
} else {
(
ImageGenConfig::default(),
VideoGenConfig::default(),
WebSearchConfig::default(),
AppBuilderDeployerConfig::default(),
)
};
kigi_tools::registry::types::SessionContext {
backend,
fs,
@@ -460,16 +289,18 @@ impl SessionContextFactory for WorkspaceSessionContextFactory {
owner_session_id: None,
parent_scheduler_handle: None,
skills: vec![],
state_path: self.resolve_state_path(session_id),
state_path: PathBuf::new(),
memory_backend: None,
web_search_config,
web_search_config: kigi_tools::implementations::web_search::WebSearchConfig::default(),
web_fetch_config: build_web_fetch_config(),
lsp: None,
image_gen_config,
video_gen_config,
app_builder_deployer_config,
image_gen_config:
kigi_tools::implementations::grok_build::image_gen::ImageGenConfig::default(),
video_gen_config:
kigi_tools::implementations::grok_build::video_gen::VideoGenConfig::default(),
app_builder_deployer_config:
kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig::default(),
api_key_provider: None,
auth_provider: self.auth.clone(),
attribution_callback: None,
system_reminder_tag: kigi_tools::reminders::DEFAULT_REMINDER_TAG,
}
@@ -488,18 +319,6 @@ impl SessionContextFactory for WorkspaceSessionContextFactory {
IDS.clone()
}
}
/// Build extra headers for API calls routed through the chat proxy.
/// Mirrors the shell's `inject_proxy_headers` logic.
fn build_proxy_headers(base_url: &str) -> indexmap::IndexMap<String, String> {
let mut headers = indexmap::IndexMap::new();
let version = kigi_version::VERSION;
headers.insert(
"user-agent".to_string(),
format!("kigi-workspace/{version}"),
);
headers.insert("x-grok-client-version".to_string(), version.to_string());
headers
}
/// Build web fetch config. Enabled with default params unless
/// `KIGI_DISABLE_WEB_FETCH=1` is set.
fn build_web_fetch_config() -> kigi_tools::implementations::grok_build::web_fetch::WebFetchConfig {
@@ -574,7 +393,6 @@ pub mod test_support {
video_gen_config: Default::default(),
app_builder_deployer_config: Default::default(),
api_key_provider: None,
auth_provider: None,
attribution_callback: None,
system_reminder_tag: kigi_tools::reminders::DEFAULT_REMINDER_TAG,
}
@@ -635,7 +453,6 @@ mod tests {
baseline,
CapabilityMode::ReadWrite,
&[],
&[],
cwd,
empty_env(),
"main",
@@ -672,7 +489,6 @@ mod tests {
baseline,
CapabilityMode::ReadWrite,
&snapshot,
&[],
PathBuf::from("/tmp"),
empty_env(),
"main",
@@ -752,7 +568,6 @@ mod tests {
baseline,
CapabilityMode::ReadOnly,
&[],
&[],
PathBuf::from("/tmp"),
empty_env(),
"main",
@@ -792,13 +607,7 @@ mod tests {
behavior_preset: None,
};
let mcp_edit = test_support::tc("mcp.editor", Some(ToolKind::Edit));
let filtered = merge_and_filter(
&baseline,
&[mcp_edit],
&[],
CapabilityMode::ReadOnly,
"test",
);
let filtered = merge_and_filter(&baseline, &[mcp_edit], CapabilityMode::ReadOnly, "test");
assert!(!filtered.tools.iter().any(|t| t.id == "mcp.editor"));
}
#[tokio::test]
@@ -812,13 +621,7 @@ mod tests {
behavior_preset: None,
};
let mcp = vec![test_support::tc("mcp.opaque", None)];
let filtered = merge_and_filter(
&baseline,
&mcp,
&[],
CapabilityMode::ReadOnly,
"test_session",
);
let filtered = merge_and_filter(&baseline, &mcp, CapabilityMode::ReadOnly, "test_session");
let kept_ids: Vec<&str> = filtered.tools.iter().map(|t| t.id.as_str()).collect();
assert!(
kept_ids.contains(&"baseline.opaque"),
@@ -841,7 +644,7 @@ mod tests {
behavior_preset: None,
};
let mcp = vec![test_support::tc("mcp.opaque", None)];
let filtered = merge_and_filter(&baseline, &mcp, &[], CapabilityMode::All, "test_session");
let filtered = merge_and_filter(&baseline, &mcp, CapabilityMode::All, "test_session");
let kept_ids: Vec<&str> = filtered.tools.iter().map(|t| t.id.as_str()).collect();
assert!(
kept_ids.contains(&"mcp.opaque"),
@@ -859,13 +662,7 @@ mod tests {
let mut mcp_b = test_support::tc("mcp.tool_b", Some(ToolKind::Read));
mcp_b.name_override = Some("shared_name".into());
let mcp = vec![mcp_a, mcp_b];
let filtered = merge_and_filter(
&baseline,
&mcp,
&[],
CapabilityMode::ReadOnly,
"test_session",
);
let filtered = merge_and_filter(&baseline, &mcp, CapabilityMode::ReadOnly, "test_session");
let ids: Vec<&str> = filtered.tools.iter().map(|t| t.id.as_str()).collect();
assert!(ids.contains(&"mcp.tool_a"), "first wins: {ids:?}");
assert!(
@@ -874,144 +671,6 @@ mod tests {
);
}
#[test]
fn hub_tool_merged_into_empty_baseline() {
let baseline = ToolServerConfig {
tools: vec![],
behavior_preset: None,
};
let hub = vec![test_support::tc("hub:remote_exec", None)];
let filtered = merge_and_filter(&baseline, &[], &hub, CapabilityMode::All, "test");
let ids: Vec<&str> = filtered.tools.iter().map(|t| t.id.as_str()).collect();
assert!(
ids.contains(&"hub:remote_exec"),
"remote tool should appear under All mode: {ids:?}"
);
}
#[test]
fn hub_tool_dropped_under_readonly_because_kind_none() {
let baseline = ToolServerConfig {
tools: vec![test_support::tc(
"GrokBuild:read_file",
Some(ToolKind::Read),
)],
behavior_preset: None,
};
let hub = vec![test_support::tc("hub:remote_exec", None)];
let filtered = merge_and_filter(&baseline, &[], &hub, CapabilityMode::ReadOnly, "test");
let ids: Vec<&str> = filtered.tools.iter().map(|t| t.id.as_str()).collect();
assert!(
!ids.contains(&"hub:remote_exec"),
"hub kind: None MUST be dropped under ReadOnly: {ids:?}"
);
}
#[test]
fn hub_tool_dedup_baseline_wins() {
let baseline = ToolServerConfig {
tools: vec![test_support::tc("hub:read_file", Some(ToolKind::Read))],
behavior_preset: None,
};
let hub = vec![test_support::tc("hub:read_file", None)];
let filtered = merge_and_filter(&baseline, &[], &hub, CapabilityMode::All, "test");
let count = filtered
.tools
.iter()
.filter(|t| t.id == "hub:read_file")
.count();
assert_eq!(count, 1, "duplicate should be deduped");
}
#[test]
fn hub_tool_dedup_mcp_wins_over_hub() {
let baseline = ToolServerConfig {
tools: vec![],
behavior_preset: None,
};
let mcp = vec![test_support::tc("hub:shared_tool", Some(ToolKind::Read))];
let hub = vec![test_support::tc("hub:shared_tool", None)];
let filtered = merge_and_filter(&baseline, &mcp, &hub, CapabilityMode::All, "test");
let count = filtered
.tools
.iter()
.filter(|t| t.id == "hub:shared_tool")
.count();
assert_eq!(count, 1, "MCP wins; hub duplicate skipped");
let tool = filtered
.tools
.iter()
.find(|t| t.id == "hub:shared_tool")
.unwrap();
assert_eq!(tool.kind, Some(ToolKind::Read));
}
#[test]
fn hub_tool_name_collision_with_baseline_skipped() {
let baseline = ToolServerConfig {
tools: vec![test_support::tc(
"GrokBuild:read_file",
Some(ToolKind::Read),
)],
behavior_preset: None,
};
let mut hub_tool = test_support::tc("hub:read_file_v2", None);
hub_tool.name_override = Some("read_file".into());
let hub = vec![hub_tool];
let filtered = merge_and_filter(&baseline, &[], &hub, CapabilityMode::All, "test");
let ids: Vec<&str> = filtered.tools.iter().map(|t| t.id.as_str()).collect();
assert!(
!ids.contains(&"hub:read_file_v2"),
"remote tool with colliding client name must be skipped: {ids:?}"
);
}
#[test]
fn empty_hub_snapshot_is_noop() {
let baseline = test_support::baseline_config();
let baseline_ids: Vec<String> = baseline.tools.iter().map(|t| t.id.clone()).collect();
let filtered = merge_and_filter(&baseline, &[], &[], CapabilityMode::ReadWrite, "test");
let filtered_ids: Vec<String> = filtered.tools.iter().map(|t| t.id.clone()).collect();
assert_eq!(filtered_ids, baseline_ids);
}
/// Only the literal `"true"` enables tool-state persistence.
#[test]
fn tool_state_enabled_only_true_enables() {
let _guard = super::TOOL_STATE_ENV_LOCK
.lock()
.unwrap_or_else(|e| e.into_inner());
let var = "KIGI_WORKSPACE_TOOL_STATE_ENABLED";
unsafe { std::env::remove_var(var) };
assert!(!tool_state_enabled(), "unset → disabled");
unsafe { std::env::set_var(var, "false") };
assert!(!tool_state_enabled(), "false → disabled");
unsafe { std::env::set_var(var, "1") };
assert!(!tool_state_enabled(), "1 → disabled (only \"true\")");
unsafe { std::env::set_var(var, "true") };
assert!(tool_state_enabled(), "true → enabled");
unsafe { std::env::remove_var(var) };
}
/// With a tool-state home set, state is rooted at
/// `<home>/sessions/<session_id>/tool_state.json` and the dir is created.
#[test]
fn factory_resolves_session_keyed_state_path_when_home_set() {
let home = tempfile::TempDir::new().unwrap();
let factory =
WorkspaceSessionContextFactory::new().with_tool_state_home(home.path().to_path_buf());
let p = factory.resolve_state_path("sess-1");
assert_eq!(
p,
home.path()
.join("sessions")
.join("sess-1")
.join("tool_state.json")
);
assert!(
home.path().join("sessions").join("sess-1").is_dir(),
"the session dir must be created so the persistence writer can rename into it"
);
}
/// Without a tool-state home, `state_path` stays empty (legacy behavior).
#[test]
fn factory_state_path_empty_when_home_unset() {
let factory = WorkspaceSessionContextFactory::new();
assert_eq!(factory.resolve_state_path("sess-1"), PathBuf::new());
}
#[test]
fn factory_session_folder_is_tmp_sessions_not_project_cwd() {
let cwd = PathBuf::from("/workspace");
let folder = WorkspaceSessionContextFactory::resolve_session_folder("sess-1");
@@ -1063,18 +722,16 @@ mod tests {
/// A hostile `session_id` (`../../etc`) is sanitized to a single safe
/// segment and cannot traverse outside `<home>/sessions/`.
#[test]
fn factory_sanitizes_malicious_session_id_no_traversal() {
fn ensure_session_dir_sanitizes_malicious_session_id_no_traversal() {
let home = tempfile::TempDir::new().unwrap();
let factory =
WorkspaceSessionContextFactory::new().with_tool_state_home(home.path().to_path_buf());
let sessions = home.path().join("sessions");
let p = factory.resolve_state_path("../../etc");
let (session_dir, created) = ensure_session_dir(home.path(), "../../etc");
assert!(created.is_ok());
assert!(
p.starts_with(&sessions),
"state path escaped sessions/: {}",
p.display()
session_dir.starts_with(&sessions),
"session dir escaped sessions/: {}",
session_dir.display()
);
let session_dir = p.parent().expect("state path has a parent dir");
assert_eq!(
session_dir.parent(),
Some(sessions.as_path()),
@@ -1123,7 +780,6 @@ mod tests {
test_support::baseline_config(),
CapabilityMode::ReadWrite,
&[],
&[],
cwd.clone(),
empty_env(),
"sess-A",
@@ -1145,7 +801,6 @@ mod tests {
test_support::baseline_config(),
CapabilityMode::ReadWrite,
&[],
&[],
cwd.clone(),
empty_env(),
"sess-A",
@@ -1170,7 +825,6 @@ mod tests {
test_support::baseline_config(),
CapabilityMode::ReadWrite,
&[],
&[],
cwd,
empty_env(),
"sess-B",
crates/codegen/kigi-workspace/src/workspace_ops.rs
+91 -323
View File
@@ -1,30 +1,22 @@
//! [`WorkspaceOps`] — dual-mode workspace operations handle.
//! [`WorkspaceOps`] — workspace operations handle.
//!
//! Two modes:
//!
//! - **`Local`** — extensions dispatch through [`WorkspaceHandle`]; tool
//! calls dispatch through the workspace session's [`FinalizedToolset`].
//! The toolset is installed via [`WorkspaceOps::bind_local_session`]
//! after the agent is built.
//!
//! - **`Proxy`** — everything routes through hub WebSocket to a remote
//! workspace server.
//! Extensions dispatch through [`WorkspaceHandle`]; tool calls dispatch
//! through the workspace session's [`FinalizedToolset`]. The toolset is
//! installed via [`WorkspaceOps::bind_local_session`] after the agent is
//! built.
//!
//! ## Type safety
//!
//! Each RPC method has a corresponding request struct that implements
//! Each operation has a corresponding request struct that implements
//! [`WorkspaceRpc`]. The struct carries a `METHOD` constant and derives
//! `Serialize + Deserialize`. Both the proxy client (`WorkspaceOps`) and
//! the server (`WorkspaceRpcHandler::dispatch`) use the same struct —
//! add/rename a field and the compiler catches both sides.
//! `Serialize + Deserialize`, so op identity and payload shape live in
//! one place.
use crate::error::{WorkspaceError, WorkspaceResult};
use crate::file_system::ContentSearchRequest;
use crate::handle::WorkspaceHandle;
use crate::worktree::{ApplyWorktreeRequest, CreateWorktreeRequest, RemoveWorktreeRequest};
use async_trait::async_trait;
use kigi_computer_hub_sdk::ToolHarness;
use kigi_tools::types::output::ToolRunResult;
use kigi_workspace_client::{WorkspaceClient, is_transport_fatal};
pub use kigi_workspace_types::rpc::WorkspaceRpc;
pub use kigi_workspace_types::rpc::agents_md::DiscoverAgentsMdReq;
pub use kigi_workspace_types::rpc::code_nav::{
@@ -32,9 +24,8 @@ pub use kigi_workspace_types::rpc::code_nav::{
CodeIndexStats, CodeIndexStatusReq, CodeIndexStatusResponse, CodeNavLocation, CodeNavResponse,
};
pub use kigi_workspace_types::rpc::fs::{
ClientFsListNode, ClientFsListReq, ClientFsListRes, ClientFsReadFileReq, ClientFsReadFileRes,
ClientFsStatReq, ClientFsStatRes, GetFileEntry, GetFileResult, GetFilesReq, GetFilesRes,
PutFileEntry, PutFileResult, PutFilesReq, PutFilesRes,
GetFileEntry, GetFileResult, GetFilesReq, GetFilesRes, PutFileEntry, PutFileResult,
PutFilesReq, PutFilesRes,
};
pub use kigi_workspace_types::rpc::git::{
BinaryFileInfoData, CheckoutCommitResponse, CommitWithPatchData, DetectVcsKindReq,
@@ -69,7 +60,6 @@ use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
/// Implements [`WorkspaceRpc`] for request types whose responses
/// reference crate-internal types and so cannot live in the types crate.
macro_rules! workspace_rpc {
@@ -214,7 +204,7 @@ fn session_tracker(
session_id: Option<&str>,
) -> WorkspaceResult<kigi_hunk_tracker::HunkTrackerHandle> {
let sid = session_id
.ok_or_else(|| WorkspaceError::HubError("per-session hunk op requires a session".into()))?;
.ok_or_else(|| WorkspaceError::Internal("per-session hunk op requires a session".into()))?;
let session = ws
.session(sid)
.ok_or_else(|| WorkspaceError::SessionNotFound(sid.to_owned()))?;
@@ -251,13 +241,13 @@ impl WorkspaceOp for GitStatusExtReq {
self.include_patches,
)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
Ok(GitStatusExtResponse::structured(data))
}
GitStatusFormat::Prompt => {
let result = crate::file_system::git_status(cwd)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
Ok(GitStatusExtResponse::prompt(result))
}
}
@@ -273,7 +263,7 @@ impl WorkspaceOp for GitFilesReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::read_files(&cwd, &self.paths, &self.version)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -294,7 +284,7 @@ impl WorkspaceOp for GitDiffReq {
self.merge_base,
)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -307,7 +297,7 @@ impl WorkspaceOp for GitStageReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::stage(&cwd, self.paths.clone())
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -320,7 +310,7 @@ impl WorkspaceOp for GitStageContentReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::stage_content(&cwd, &self.path, &self.content)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -333,7 +323,7 @@ impl WorkspaceOp for GitUnstageReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::unstage(&cwd, self.paths.clone())
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -346,7 +336,7 @@ impl WorkspaceOp for GitDiscardReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::discard(&cwd, self.paths.clone(), self.scope, self.include_untracked)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -366,7 +356,7 @@ impl WorkspaceOp for GitCommitReq {
self.sync,
)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -379,7 +369,7 @@ impl WorkspaceOp for GitCheckoutReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::checkout_branch(&cwd, &self.branch, self.create)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -392,7 +382,7 @@ impl WorkspaceOp for GitStashReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::stash(&cwd, self.include_untracked)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -405,7 +395,7 @@ impl WorkspaceOp for GitInfoReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::git_info(&cwd)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -418,7 +408,7 @@ impl WorkspaceOp for GitBranchesReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::list_branches(&cwd)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -429,7 +419,7 @@ impl WorkspaceOp for GitCollectChangesReq {
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
{
return Err(WorkspaceError::HubError(
return Err(WorkspaceError::Internal(
"git collect changes is unavailable in this build".to_string(),
));
}
@@ -550,7 +540,7 @@ impl WorkspaceOp for PrepareWorktreeFromWorktreeReq {
spawn_task: result.spawn_task,
response: Some(
serde_json::to_value(&resp)
.map_err(|e| WorkspaceError::HubError(e.to_string()))?,
.map_err(|e| WorkspaceError::Internal(e.to_string()))?,
),
error: None,
}),
@@ -572,7 +562,7 @@ impl WorkspaceOp for CreateWorktreeFromWorktreeSyncReq {
let req = crate::worktree::CreateWorktreeFromWorktreeRequest::from(self.inner.clone());
crate::worktree::create_worktree_from_worktree_sync(&req)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -583,8 +573,8 @@ impl WorkspaceOp for WorktreeDbRebuildReq {
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
let report = crate::worktree::worktree_db_rebuild()
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::to_value(report).map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::to_value(report).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -595,7 +585,7 @@ impl WorkspaceOp for WorktreeDbPathReq {
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
let path = crate::worktree::worktree_db_path()
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
Ok(WorktreeDbPathResponse {
path: Some(path.display().to_string()),
})
@@ -864,8 +854,8 @@ fn hook_registry_to_wire(
registry: &kigi_hooks::discovery::HookRegistry,
) -> WorkspaceResult<HookRegistryWire> {
let value =
serde_json::to_value(registry).map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::from_value(value).map_err(|e| WorkspaceError::HubError(e.to_string()))
serde_json::to_value(registry).map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::from_value(value).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
/// Inverse of [`hook_registry_to_wire`]. The compiled `matcher` is absent from
/// the wire (and from this result); callers recompile it via
@@ -873,8 +863,8 @@ fn hook_registry_to_wire(
fn wire_to_hook_registry(
wire: &HookRegistryWire,
) -> WorkspaceResult<kigi_hooks::discovery::HookRegistry> {
let value = serde_json::to_value(wire).map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::from_value(value).map_err(|e| WorkspaceError::HubError(e.to_string()))
let value = serde_json::to_value(wire).map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::from_value(value).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
#[async_trait]
impl WorkspaceOp for HookRegistryReq {
@@ -886,56 +876,6 @@ impl WorkspaceOp for HookRegistryReq {
hook_registry_to_wire(&ws.hook_registry())
}
}
#[async_trait]
impl WorkspaceOp for PutFilesReq {
async fn execute(
&self,
ws: &WorkspaceHandle,
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
ws.put_files(self.files.clone()).await
}
}
#[async_trait]
impl WorkspaceOp for GetFilesReq {
async fn execute(
&self,
ws: &WorkspaceHandle,
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
ws.get_files(self.files.clone()).await
}
}
#[async_trait]
impl WorkspaceOp for ClientFsListReq {
async fn execute(
&self,
ws: &WorkspaceHandle,
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
crate::file_system::client_fs::list(ws, self).await
}
}
#[async_trait]
impl WorkspaceOp for ClientFsStatReq {
async fn execute(
&self,
ws: &WorkspaceHandle,
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
crate::file_system::client_fs::stat(ws, self).await
}
}
#[async_trait]
impl WorkspaceOp for ClientFsReadFileReq {
async fn execute(
&self,
ws: &WorkspaceHandle,
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
crate::file_system::client_fs::read_file(ws, self).await
}
}
/// Resolve the index root for a code-nav op. Prefers the explicit per-session
/// `root` (the cwd the client sends per window), else the workspace root.
/// Without this, code nav in a non-primary window would query the launch
@@ -972,7 +912,7 @@ impl WorkspaceOp for CodeGotoDefinitionReq {
let result = handle
.goto_definition(std::path::PathBuf::from(&self.file), self.line, self.col)
.await
.map_err(|e| WorkspaceError::HubError(format!("index channel closed: {e}")))?;
.map_err(|e| WorkspaceError::Internal(format!("index channel closed: {e}")))?;
Ok(query_result_to_response(result))
}
}
@@ -992,7 +932,7 @@ impl WorkspaceOp for CodeGotoReferencesReq {
self.include_definition,
)
.await
.map_err(|e| WorkspaceError::HubError(format!("index channel closed: {e}")))?;
.map_err(|e| WorkspaceError::Internal(format!("index channel closed: {e}")))?;
Ok(query_result_to_response(result))
}
}
@@ -1010,7 +950,7 @@ impl WorkspaceOp for CodeFindDefinitionsReq {
self.context_file.as_ref().map(std::path::PathBuf::from),
)
.await
.map_err(|e| WorkspaceError::HubError(format!("index channel closed: {e}")))?;
.map_err(|e| WorkspaceError::Internal(format!("index channel closed: {e}")))?;
Ok(symbol_locations_to_response(result))
}
}
@@ -1028,7 +968,7 @@ impl WorkspaceOp for CodeFindReferencesReq {
self.context_file.as_ref().map(std::path::PathBuf::from),
)
.await
.map_err(|e| WorkspaceError::HubError(format!("index channel closed: {e}")))?;
.map_err(|e| WorkspaceError::Internal(format!("index channel closed: {e}")))?;
Ok(symbol_locations_to_response(result))
}
}
@@ -1105,9 +1045,9 @@ impl WorkspaceOp for CreateWorktreeRequest {
let result = crate::worktree::prepare_worktree_creation(self).await;
match result.response {
Ok(resp) => {
serde_json::to_value(resp).map_err(|e| WorkspaceError::HubError(e.to_string()))
serde_json::to_value(resp).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
Err(e) => Err(WorkspaceError::HubError(e.to_string())),
Err(e) => Err(WorkspaceError::Internal(e.to_string())),
}
}
}
@@ -1121,8 +1061,8 @@ impl WorkspaceOp for RemoveWorktreeRequest {
let copy_ctx = crate::worktree::BackgroundCopyContext::new();
let result = crate::worktree::remove_worktree(self, &copy_ctx)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::to_value(result).map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::to_value(result).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -1134,8 +1074,8 @@ impl WorkspaceOp for ApplyWorktreeRequest {
) -> WorkspaceResult<Self::Response> {
let result = crate::worktree::apply_worktree(self)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::to_value(result).map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::to_value(result).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -1147,8 +1087,8 @@ impl WorkspaceOp for WorktreeListReq {
) -> WorkspaceResult<Self::Response> {
let records =
crate::worktree::list_worktrees(self.repo.as_deref(), &self.types, self.include_all)
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::to_value(records).map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::to_value(records).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -1159,8 +1099,8 @@ impl WorkspaceOp for WorktreeShowReq {
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
let record = crate::worktree::show_worktree(&self.id_or_path)
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::to_value(record).map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::to_value(record).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -1175,9 +1115,9 @@ impl WorkspaceOp for WorktreeGcReq {
crate::worktree::gc_worktrees_mgmt(dry_run, max_age_secs, force)
})
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::to_value(report).map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))?
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::to_value(report).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -1188,27 +1128,22 @@ impl WorkspaceOp for WorktreeDbStatsReq {
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
let stats = crate::worktree::worktree_db_stats()
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::to_value(stats).map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::to_value(stats).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
/// Dual-mode workspace operations handle.
///
/// - **`Local`** — wraps a [`WorkspaceHandle`]. Extensions dispatch
/// through the handle; tool calls dispatch through the workspace
/// session's [`FinalizedToolset`](kigi_tools::registry::types::FinalizedToolset).
/// Call [`bind_local_session`](Self::bind_local_session) after building
/// the agent to install the toolset on the workspace session.
///
/// - **`Proxy`** — wraps a [`WorkspaceClient`] connected to a remote hub.
/// Everything routes through hub WebSocket to a remote workspace server.
/// Wraps a [`WorkspaceHandle`]. Extensions dispatch through the handle;
/// tool calls dispatch through the workspace session's
/// [`FinalizedToolset`](kigi_tools::registry::types::FinalizedToolset).
/// Call [`bind_local_session`](Self::bind_local_session) after building
/// the agent to install the toolset on the workspace session.
#[derive(Clone)]
pub enum WorkspaceOps {
/// Local in-process mode — extensions through the handle, tool calls
/// through the workspace session's toolset.
Local { handle: WorkspaceHandle },
/// Proxy mode — routes through hub RPC.
Proxy { client: WorkspaceClient },
}
impl WorkspaceOps {
/// Construct a local-mode ops handle.
@@ -1219,36 +1154,10 @@ impl WorkspaceOps {
pub fn local(handle: WorkspaceHandle) -> Self {
Self::Local { handle }
}
/// Construct a proxy-mode ops handle.
pub fn proxy(harness: Arc<ToolHarness>) -> Self {
Self::Proxy {
client: WorkspaceClient::new((*harness).clone()),
}
}
/// Construct a proxy-mode ops handle sharing a pre-created connected
/// flag. The same `Arc<AtomicBool>` should be wired into the harness
/// builder's `on_reconnect` callback so reconnects reset the flag.
pub fn proxy_with_connected(harness: Arc<ToolHarness>, connected: Arc<AtomicBool>) -> Self {
Self::Proxy {
client: WorkspaceClient::with_connected_flag((*harness).clone(), connected),
}
}
/// Whether this handle routes through the server (proxy mode).
pub fn is_proxy(&self) -> bool {
matches!(self, Self::Proxy { .. })
}
/// Access the underlying workspace RPC client (proxy mode only).
pub fn client(&self) -> Option<&WorkspaceClient> {
match self {
Self::Proxy { client } => Some(client),
Self::Local { .. } => None,
}
}
/// Access the underlying workspace handle (local mode only).
/// Access the underlying workspace handle.
pub fn workspace_handle(&self) -> Option<&WorkspaceHandle> {
match self {
Self::Local { handle } => Some(handle),
Self::Proxy { .. } => None,
}
}
/// Create the workspace session and bind the agent's toolset for local mode.
@@ -1274,9 +1183,7 @@ impl WorkspaceOps {
toolset: Arc<kigi_tools::registry::types::FinalizedToolset>,
viewer_ctx: Option<kigi_tool_runtime::WorkspaceViewerContext>,
) -> WorkspaceResult<()> {
let Self::Local { handle } = self else {
return Ok(());
};
let Self::Local { handle } = self;
if handle.session(session_id).is_none() {
handle.create_session_with_tracker_and_viewer_ctx(
session_id,
@@ -1296,9 +1203,7 @@ impl WorkspaceOps {
}
/// Release the workspace session. No-op in proxy mode.
pub fn end_local_session(&self, session_id: &str) {
let Self::Local { handle } = self else {
return;
};
let Self::Local { handle } = self;
handle.on_session_ended(session_id);
if let Err(e) = handle.drop_session(session_id, session_id) {
tracing::debug!(
@@ -1312,82 +1217,26 @@ impl WorkspaceOps {
session_id: &str,
payload: &kigi_tool_protocol::turn_hook::BeforeTurnPayload,
) {
match self {
Self::Local { handle } => {
handle.on_before_turn(session_id, payload).await;
}
Self::Proxy { .. } => {
tracing::debug!("on_before_turn called on Proxy WorkspaceOps (no-op)");
}
}
let Self::Local { handle } = self;
handle.on_before_turn(session_id, payload).await;
}
pub async fn on_after_turn(
&self,
session_id: &str,
payload: &kigi_tool_protocol::turn_hook::AfterTurnPayload,
) {
match self {
Self::Local { handle } => {
handle.on_after_turn(session_id, payload).await;
}
Self::Proxy { .. } => {
tracing::debug!("on_after_turn called on Proxy WorkspaceOps (no-op)");
}
}
let Self::Local { handle } = self;
handle.on_after_turn(session_id, payload).await;
}
pub async fn rpc_raw(&self, method: &str, params: Value) -> WorkspaceResult<Value> {
let client = match self {
Self::Proxy { client } => client,
Self::Local { .. } => {
return Err(WorkspaceError::HubError(
"rpc not available in local mode".into(),
));
}
};
client
.rpc_raw(method, params)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
}
async fn rpc<R: WorkspaceRpc>(&self, req: &R) -> WorkspaceResult<R::Response> {
let params = serde_json::to_value(req)
.map_err(|e| WorkspaceError::HubError(format!("serialize failed: {e}")))?;
let terminal = self.rpc_raw(R::METHOD, params).await?;
let envelope: crate::rpc_envelope::RpcEnvelope<R::Response> =
serde_json::from_value(terminal)
.map_err(|e| WorkspaceError::HubError(format!("envelope parse failed: {e}")))?;
envelope
.into_result()
.map_err(crate::rpc_envelope::rpc_error_to_workspace)
}
/// Dispatch a typed operation in either local or proxy mode.
///
/// - **Local mode**: calls `op.execute(handle, session_id)` directly.
/// - **Proxy mode**: serializes the op and routes through the server RPC.
/// The server handler owns session context, so `session_id` is only
/// needed for local `execute()`.
/// Dispatch a typed operation: calls `op.execute(handle, session_id)`.
pub async fn dispatch<Op: WorkspaceOp>(
&self,
op: &Op,
session_id: Option<&str>,
) -> WorkspaceResult<Op::Response> {
let mode = match self {
Self::Local { .. } => "local",
Self::Proxy { .. } => "proxy",
};
tracing::debug!(method = Op::METHOD, mode, "WorkspaceOps::dispatch");
match self {
Self::Local { handle } => op.execute(handle, session_id).await,
Self::Proxy { .. } => self.rpc(op).await,
}
}
pub async fn workspace_info(&self) -> WorkspaceResult<Value> {
self.rpc(&WorkspaceInfoReq {}).await
}
/// **DEPRECATED**: Use [`Self::git_status_ext`] with `format: GitStatusFormat::Prompt`
/// instead. This method will be removed in a future release.
pub async fn git_status(&self) -> WorkspaceResult<Value> {
self.rpc(&GitStatusReq {}).await
tracing::debug!(method = Op::METHOD, "WorkspaceOps::dispatch");
let Self::Local { handle } = self;
op.execute(handle, session_id).await
}
/// Get git status with configurable output format.
///
@@ -1408,52 +1257,9 @@ impl WorkspaceOps {
let wire = self.dispatch(&HookRegistryReq {}, None).await?;
wire_to_hook_registry(&wire)
}
pub async fn begin_prompt(&self, session_id: &str, prompt_index: usize) -> WorkspaceResult<()> {
self.rpc(&BeginPromptReq {
session_id: session_id.to_owned(),
prompt_index,
})
.await
}
pub async fn end_prompt(&self, session_id: &str, prompt_index: usize) -> WorkspaceResult<()> {
self.rpc(&EndPromptReq {
session_id: session_id.to_owned(),
prompt_index,
})
.await
}
pub async fn get_rewind_points(
&self,
session_id: &str,
) -> WorkspaceResult<Vec<crate::session::file_state::RewindPoint>> {
self.rpc(&GetRewindPointsReq {
session_id: session_id.to_owned(),
})
.await
}
pub async fn rewind_to(
&self,
session_id: &str,
target_prompt_index: usize,
) -> WorkspaceResult<crate::session::file_state::FileRewindResponse> {
self.rpc(&RewindToReq {
session_id: session_id.to_owned(),
target_prompt_index,
})
.await
}
pub async fn put_files(&self, req: PutFilesReq) -> WorkspaceResult<PutFilesRes> {
self.dispatch(&req, None).await
}
pub async fn get_files(&self, req: GetFilesReq) -> WorkspaceResult<GetFilesRes> {
self.dispatch(&req, None).await
}
/// Dispatch a tool call through the workspace.
///
/// - **Local**: dispatches through the workspace session's
/// [`FinalizedToolset`](kigi_tools::registry::types::FinalizedToolset)
/// (in-process). Requires `session_id` to look up the session.
/// - **Proxy**: routes through the server `ToolHarness` (remote).
/// Dispatch a tool call through the workspace session's
/// [`FinalizedToolset`](kigi_tools::registry::types::FinalizedToolset)
/// (in-process). Requires `session_id` to look up the session.
pub async fn call_tool(
&self,
name: &str,
@@ -1461,57 +1267,23 @@ impl WorkspaceOps {
call_id: &str,
session_id: Option<&str>,
) -> Result<ToolRunResult, kigi_tool_runtime::ToolError> {
match self {
Self::Local { handle } => {
let session_id = session_id.ok_or_else(|| {
kigi_tool_runtime::ToolError::custom(
"missing_session",
"session_id required for local tool dispatch",
)
})?;
let session = handle.session(session_id).ok_or_else(|| {
kigi_tool_runtime::ToolError::custom(
"session_not_found",
format!(
"workspace session not found: {session_id} \
— call bind_local_session() first"
),
)
})?;
session.toolset().call(name, args, call_id, None).await
}
Self::Proxy { client } => {
if !client.is_connected() {
return Err(kigi_tool_runtime::ToolError::network_error(
"The workspace server connection was lost. \
Please restart your session to reconnect.",
));
}
let tool_id = kigi_tool_protocol::ToolId::new(name).map_err(|e| {
kigi_tool_runtime::ToolError::custom(
"hub_proxy_error",
format!("invalid tool name: {e}"),
)
})?;
let mut ctx = kigi_tool_runtime::ToolCallContext::default();
ctx.call_id =
kigi_tool_protocol::ToolCallId::new(call_id.to_owned()).unwrap_or(ctx.call_id);
let mut stream = client.harness().call(tool_id, args, ctx).await;
let typed = crate::hub_channel::consume_stream_terminal(&mut stream)
.await
.inspect_err(|e| {
if is_transport_fatal(e) {
client.mark_disconnected();
}
})?;
serde_json::from_value::<ToolRunResult>(typed.value).map_err(|e| {
kigi_tool_runtime::ToolError::custom(
"tool_result_deserialize",
format!("tool result deserialization failed: {e}"),
)
})
}
}
let Self::Local { handle } = self;
let session_id = session_id.ok_or_else(|| {
kigi_tool_runtime::ToolError::custom(
"missing_session",
"session_id required for local tool dispatch",
)
})?;
let session = handle.session(session_id).ok_or_else(|| {
kigi_tool_runtime::ToolError::custom(
"session_not_found",
format!(
"workspace session not found: {session_id} \
— call bind_local_session() first"
),
)
})?;
session.toolset().call(name, args, call_id, None).await
}
}
#[cfg(any(test, feature = "test-support"))]
@@ -1569,9 +1341,7 @@ mod tests {
#[test]
fn git_op_cwd_uses_explicit_git_root_per_window() {
let ops = WorkspaceOps::for_test();
let WorkspaceOps::Local { handle } = &ops else {
unreachable!("for_test builds a local handle");
};
let WorkspaceOps::Local { handle } = &ops;
let workspace_root = handle.root_cwd().unwrap();
let window_a = std::path::PathBuf::from("/repos/xai-main");
let window_b = std::path::PathBuf::from("/repos/xai-main-2");
@@ -1594,9 +1364,7 @@ mod tests {
#[tokio::test]
async fn end_local_session_drops_bound_toolset() {
let ops = WorkspaceOps::for_test();
let WorkspaceOps::Local { handle } = &ops else {
unreachable!("for_test builds a local handle");
};
let WorkspaceOps::Local { handle } = &ops;
let sid = "sess-teardown";
let toolset =
std::sync::Arc::new(kigi_tools::registry::types::FinalizedToolset::empty_for_test());
crates/common/kigi-tool-protocol/Cargo.toml
+1 -1
View File
@@ -3,7 +3,7 @@ license = "Apache-2.0"
name = "kigi-tool-protocol"
version.workspace = true
edition.workspace = true
description = "Wire-protocol types for the xAI Computer Hub"
description = "Tool wire-protocol types"
[dependencies]
kigi-tool-types = { workspace = true }
crates/common/kigi-tool-protocol/src/lib.rs
+1 -1
View File
@@ -1,4 +1,4 @@
//! xAI Computer Hub — wire-protocol types.
//! Tool wire-protocol types.
//!
//! Identifier newtypes, registration payloads, capabilities, hook events,
//! handshake messages, the JSON-RPC 2.0 envelope and method catalog, the
crates/common/kigi-tool-runtime/Cargo.toml
+3 -2
View File
@@ -3,17 +3,18 @@ license = "Apache-2.0"
name = "kigi-tool-runtime"
version.workspace = true
edition.workspace = true
description = "Unified Tool trait, dispatch trait, error taxonomy, notifications, and search index for the xAI Computer Hub"
description = "Unified Tool trait, dispatch trait, error taxonomy, notifications, local tool registry, and search index"
[dependencies]
anyhow = { workspace = true }
async-trait = { workspace = true }
futures = { workspace = true }
indexmap = { workspace = true }
parking_lot = { workspace = true }
schemars = { workspace = true }
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
tokio-util = { workspace = true }
kigi-tools-api = { workspace = true }
kigi-tool-protocol = { workspace = true }
kigi-tool-types = { workspace = true }
crates/common/kigi-tool-runtime/src/context.rs
-182
View File
@@ -152,185 +152,3 @@ pub struct WorkspaceViewerContext {
#[serde(default)]
pub stream_tool_progress: bool,
}
/// Wire shape of the Computer Hub `session.bind` metadata — one definition
/// shared by the emitter (serializes) and the workspace consumer
/// (deserializes), so the two can't drift on field names/types.
///
/// Excludes anything not meant for the workspace (cached tool definitions,
/// and terminal-provisioning inputs like image/fuse/isolation) so they can
/// never reach the wire. Every field tolerates a missing/malformed value
/// (drops to default) to keep valid siblings and mixed-version compatibility.
#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
pub struct WorkspaceBindMetadata {
#[serde(
default,
deserialize_with = "ok_or_default",
skip_serializing_if = "Option::is_none"
)]
pub preset: Option<String>,
/// Raw string; the workspace maps it to its own capability enum.
#[serde(
default,
deserialize_with = "ok_or_default",
skip_serializing_if = "Option::is_none"
)]
pub capability_mode: Option<String>,
/// Explicit toolset in the grok-tools gRPC wire shape. Empty = unset.
#[serde(
default,
deserialize_with = "ok_or_default",
skip_serializing_if = "Vec::is_empty"
)]
pub tools: Vec<kigi_tools_api::ToolConfigEntry>,
#[serde(
default,
deserialize_with = "ok_or_default",
skip_serializing_if = "Option::is_none"
)]
pub viewer_ctx: Option<WorkspaceViewerContext>,
/// Initial auto-approve (YOLO) state for the bound session. Omitted when
/// unset (legacy emitters / wire compat with older workspace servers);
/// consumers fail closed on `None`.
#[serde(
default,
deserialize_with = "ok_or_default",
skip_serializing_if = "Option::is_none"
)]
pub yolo_mode: Option<bool>,
/// Optional/additive: omitted by emitters that don't yet write it.
#[serde(
default,
deserialize_with = "ok_or_default",
skip_serializing_if = "Option::is_none"
)]
pub manifest_version: Option<String>,
#[serde(
default,
deserialize_with = "ok_or_default",
skip_serializing_if = "Option::is_none"
)]
pub manifest_hash: Option<String>,
/// Opt-in: forward SystemNotifications produced in this session to the gateway.
#[serde(
default,
deserialize_with = "ok_or_default",
skip_serializing_if = "Option::is_none"
)]
pub system_notifications: Option<bool>,
#[serde(
default,
deserialize_with = "ok_or_default",
skip_serializing_if = "std::ops::Not::not"
)]
pub rpc_only: bool,
}
/// Deserialize a field, falling back to its default on a malformed value
/// instead of failing the whole struct.
fn ok_or_default<'de, D, T>(deserializer: D) -> Result<T, D::Error>
where
D: serde::Deserializer<'de>,
T: serde::de::DeserializeOwned + Default,
{
let value = <serde_json::Value as serde::Deserialize>::deserialize(deserializer)?;
Ok(serde_json::from_value(value).unwrap_or_default())
}
#[cfg(test)]
mod bind_metadata_tests {
use super::WorkspaceBindMetadata;
#[test]
fn serialize_omits_empty_fields() {
let md = WorkspaceBindMetadata::default();
assert_eq!(serde_json::to_value(&md).unwrap(), serde_json::json!({}));
}
#[test]
fn round_trips_populated() {
let md = WorkspaceBindMetadata {
preset: Some("explore".to_owned()),
capability_mode: Some("read_only".to_owned()),
tools: vec![kigi_tools_api::ToolConfigEntry {
id: "GrokBuild:grep".to_owned(),
..Default::default()
}],
viewer_ctx: Some(super::WorkspaceViewerContext {
stream_tool_progress: true,
}),
yolo_mode: Some(true),
manifest_version: Some("v1".to_owned()),
manifest_hash: Some("abc123".to_owned()),
system_notifications: Some(true),
rpc_only: true,
};
let value = serde_json::to_value(&md).unwrap();
let back: WorkspaceBindMetadata = serde_json::from_value(value).unwrap();
assert_eq!(back.preset.as_deref(), Some("explore"));
assert_eq!(back.capability_mode.as_deref(), Some("read_only"));
assert_eq!(back.tools.len(), 1);
assert!(back.viewer_ctx.unwrap().stream_tool_progress);
assert_eq!(back.yolo_mode, Some(true));
assert_eq!(back.manifest_version.as_deref(), Some("v1"));
assert_eq!(back.manifest_hash.as_deref(), Some("abc123"));
assert_eq!(back.system_notifications, Some(true));
assert!(back.rpc_only);
}
#[test]
fn rpc_only_omitted_when_false_wire_compatible() {
let md = WorkspaceBindMetadata::default();
let value = serde_json::to_value(&md).unwrap();
assert!(value.get("rpc_only").is_none());
let md: WorkspaceBindMetadata =
serde_json::from_value(serde_json::json!({"preset": "explore"})).unwrap();
assert!(!md.rpc_only);
let md: WorkspaceBindMetadata =
serde_json::from_value(serde_json::json!({"rpc_only": true})).unwrap();
assert!(md.rpc_only);
}
#[test]
fn system_notifications_is_wire_compatible() {
let md = WorkspaceBindMetadata::default();
let value = serde_json::to_value(&md).unwrap();
assert!(value.get("system_notifications").is_none());
let md = WorkspaceBindMetadata {
system_notifications: Some(true),
..Default::default()
};
let value = serde_json::to_value(&md).unwrap();
let back: WorkspaceBindMetadata = serde_json::from_value(value).unwrap();
assert_eq!(back.system_notifications, Some(true));
let md: WorkspaceBindMetadata =
serde_json::from_value(serde_json::json!({"preset": "explore"})).unwrap();
assert!(md.system_notifications.is_none());
}
#[test]
fn malformed_field_falls_back_to_default_keeping_siblings() {
// `tools` is the wrong type and `capability_mode` is fine: the bad
// field drops to default, the good sibling survives.
let value = serde_json::json!({
"preset": "explore",
"capability_mode": "read_only",
"tools": "not-a-list",
});
let md: WorkspaceBindMetadata = serde_json::from_value(value).unwrap();
assert_eq!(md.preset.as_deref(), Some("explore"));
assert_eq!(md.capability_mode.as_deref(), Some("read_only"));
assert!(md.tools.is_empty());
}
#[test]
fn legacy_payload_without_viewer_ctx_parses() {
let md: WorkspaceBindMetadata =
serde_json::from_value(serde_json::json!({"preset": "explore"})).unwrap();
assert!(md.viewer_ctx.is_none());
}
}
crates/common/kigi-tool-runtime/src/lib.rs
+7 -5
View File
@@ -1,16 +1,17 @@
//! xAI Computer Hub — unified runtime contract.
//! Unified tool runtime contract.
//!
//! Single home for the `Tool` trait, `ToolDispatch`, `ToolError`,
//! `ToolNotification`, `ToolSearchIndex`, `ToolCallContext`, `ToolStream`,
//! and the helper constructors that build well-formed streams. Adapters
//! for individual tool sources re-export from here so every tool author
//! sees the same surface.
//! the in-process `LocalRegistry`, and the helper constructors that build
//! well-formed streams. Adapters for individual tool sources re-export
//! from here so every tool author sees the same surface.
#![forbid(unsafe_code)]
pub mod context;
pub mod dispatch;
pub mod error;
pub mod local_registry;
pub mod notification;
pub mod render;
pub mod search;
@@ -19,10 +20,11 @@ pub mod tool;
pub use context::{
BehaviorVersion, Cancellation, Cwd, ListToolsContext, SessionContext, ToolCallContext,
TraceContext, TypedExtensions, WorkspaceBindMetadata, WorkspaceViewerContext,
TraceContext, TypedExtensions, WorkspaceViewerContext,
};
pub use dispatch::ToolDispatch;
pub use error::{ToolError, ToolErrorKind};
pub use local_registry::LocalRegistry;
pub use notification::{
BashExecutionBackgrounded, BashExecutionComplete, BashExecutionFailed, BashExecutionTimeout,
BashNotificationBase, BashOutputChunk, FileRead, FileWritten, LspServerCrashed,
crates/common/kigi-tool-runtime/src/local_registry.rs
+117
View File
@@ -0,0 +1,117 @@
//! In-process tool registry for local dispatch.
//!
//! [`LocalRegistry`] maps [`ToolId`]s to type-erased
//! [`ToolDyn`](crate::tool::ToolDyn) handles.
//! Toolset finalization registers every config-enabled tool here and
//! dispatch resolves handles via [`LocalRegistry::find`], so a call
//! executes in-process without any wire round-trip.
use std::sync::Arc;
use indexmap::IndexMap;
use parking_lot::RwLock;
use crate::context::ListToolsContext;
use crate::tool::{ArcTool, Tool};
use kigi_tool_protocol::ToolId;
use kigi_tool_types::ToolDescription;
/// In-process registry of tool handles.
///
/// Mutations are concurrency-safe (`RwLock` on the entry map), so
/// callers MAY hot-add or hot-remove tools while dispatch is in use.
///
/// Entries use `RwLock<IndexMap>` to preserve insertion order so that
/// [`list_tools`](Self::list_tools) returns descriptions in the same
/// order tools were registered (matching the config-defined order).
#[derive(Clone, Default)]
pub struct LocalRegistry {
entries: Arc<RwLock<IndexMap<ToolId, ArcTool>>>,
}
impl std::fmt::Debug for LocalRegistry {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LocalRegistry")
.field("entries", &self.entries.read().len())
.finish()
}
}
impl LocalRegistry {
/// Construct an empty registry.
pub fn new() -> Self {
Self::default()
}
/// Register a typed [`Tool`] implementation by value. Subsequent
/// registrations of the same id replace the previous handle and
/// return the displaced handle for inspection / drop ordering.
pub fn register<T>(&self, tool: T) -> Option<ArcTool>
where
T: Tool + 'static,
{
self.register_arc(Arc::new(tool))
}
/// Register a typed [`Tool`] implementation already wrapped in `Arc`.
pub fn register_arc<T>(&self, tool: Arc<T>) -> Option<ArcTool>
where
T: Tool + 'static,
{
let id = tool.id();
self.entries.write().insert(id, tool as ArcTool)
}
/// Register a type-erased [`ToolDyn`](crate::tool::ToolDyn) directly.
///
/// Use this for inherently dynamic tools (e.g. MCP tools retrieved
/// from a registry as `Arc<dyn ToolDyn>`) where the concrete type
/// is not available. For native tools with a concrete type, prefer
/// [`register`](Self::register).
pub fn register_dyn(&self, tool: ArcTool) -> Option<ArcTool> {
let id = tool.id();
self.entries.write().insert(id, tool)
}
/// Resolve `tool_id` to its in-process handle, if registered.
/// Returns a clone of the handle so the caller can dispatch without
/// holding the lock across an await point.
pub fn find(&self, tool_id: &ToolId) -> Option<ArcTool> {
self.entries.read().get(tool_id).cloned()
}
/// Drop the handle bound to `tool_id`. Returns `true` iff a
/// matching entry was removed.
pub fn unregister(&self, tool_id: &ToolId) -> bool {
self.entries.write().shift_remove(tool_id).is_some()
}
/// Number of tools currently registered.
pub fn len(&self) -> usize {
self.entries.read().len()
}
/// `true` iff no tools are registered.
pub fn is_empty(&self) -> bool {
self.entries.read().is_empty()
}
/// `true` iff `tool_id` is currently registered.
pub fn contains(&self, tool_id: &ToolId) -> bool {
self.entries.read().contains_key(tool_id)
}
/// Descriptions of registered tools filtered by `should_list`.
///
/// Returns descriptions in **insertion order** — the order tools
/// were registered — so the caller sees the same ordering as the
/// config-defined tool list.
pub fn list_tools(&self, ctx: &ListToolsContext) -> Vec<ToolDescription> {
self.entries
.read()
.values()
.filter(|handle| handle.should_list(ctx))
.map(|handle| handle.description(ctx))
.collect()
}
}
crates/common/kigi-tool-types/Cargo.toml
+1 -1
View File
@@ -3,7 +3,7 @@ license = "Apache-2.0"
name = "kigi-tool-types"
version.workspace = true
edition.workspace = true
description = "Canonical tool-description types for the xAI platform"
description = "Canonical tool-description types"
[features]
# Enables `BuiltinSubagent::render_prompt` (MiniJinja rendering of the
crates/common/kigi-tracing/Cargo.toml
-37
View File
@@ -1,37 +0,0 @@
[package]
license = "Apache-2.0"
edition.workspace = true
name = "kigi-tracing"
version.workspace = true
[dependencies]
async-trait = { workspace = true }
http = { workspace = true }
log = { workspace = true }
fastrace = { workspace = true }
fastrace-tonic = { workspace = true }
fastrace-opentelemetry = { workspace = true }
fastrace-reqwest = { workspace = true }
opentelemetry = { workspace = true }
opentelemetry-http = { workspace = true }
opentelemetry-otlp = { workspace = true }
opentelemetry_sdk = { workspace = true, features = ["testing"] }
reqwest = { workspace = true }
reqwest-middleware = { workspace = true }
tokio = { workspace = true }
tonic = { workspace = true }
tower = { workspace = true, features = ["full"] }
tower-http = { workspace = true, features = ["trace"] }
tracing = { workspace = true }
tracing-opentelemetry = { workspace = true }
uuid = { workspace = true }
[dev-dependencies]
bytes = { workspace = true }
http-body-util = { workspace = true }
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
tracing-subscriber = { workspace = true }
wiremock = { workspace = true }
[lints]
workspace = true
crates/common/kigi-tracing/src/dispatch.rs
-45
View File
@@ -1,45 +0,0 @@
use tracing::subscriber::NoSubscriber;
/// Returns `true` when a `tracing` dispatcher (subscriber) is active in the
/// current context — either the thread-scoped default
/// (`tracing::subscriber::with_default` / `set_default`) or the global one.
///
/// When this returns `false`, spans have no consumer. Worse than useless: when
/// `tracing` is compiled with its `log` compatibility feature, every span
/// creation and every later `Span::record(...)` is downgraded to a `log`
/// record at the span's level. In processes that only configure a `log`
/// logger — e.g. integration tests or fastrace-only binaries — that prints
/// noise like:
///
/// ```text
/// I grpc; otel.name="POST_/pkg.Service/Method" ...
/// I grpc; status_code=200 OK
/// I grpc; trace_id=00000000000000000000000000000000
/// ```
///
/// Request-span factories (gRPC/HTTP server and client middleware) call this
/// and return `Span::none()` when no dispatcher is active, so the span is
/// neither built nor downgraded to log spam.
pub fn dispatcher_active() -> bool {
tracing::dispatcher::get_default(|dispatch| !dispatch.is::<NoSubscriber>())
}
#[cfg(test)]
mod tests {
use super::*;
// NOTE: relies on no test in this binary installing a *global* subscriber
// (`OtelTestEnv` and friends use thread-scoped `set_default` guards).
#[test]
fn without_dispatcher_inactive() {
assert!(!dispatcher_active());
}
#[test]
fn scoped_dispatcher_active() {
tracing::subscriber::with_default(tracing_subscriber::registry(), || {
assert!(dispatcher_active());
});
assert!(!dispatcher_active());
}
}
crates/common/kigi-tracing/src/fastrace.rs
-83
View File
@@ -1,83 +0,0 @@
use fastrace::prelude::*;
use fastrace_opentelemetry::OpenTelemetryReporter;
use opentelemetry::InstrumentationScope;
use opentelemetry::KeyValue;
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_otlp::{ExporterBuildError, SpanExporter};
use opentelemetry_sdk::Resource;
use std::borrow::Cow;
use std::iter;
// Fastrace initialization
pub fn init_fastrace(
endpoint: String,
name: String,
resource_attributes: impl IntoIterator<Item = (String, String)>,
) -> Result<(), ExporterBuildError> {
let exporter = SpanExporter::builder()
.with_tonic()
.with_endpoint(endpoint)
.with_protocol(opentelemetry_otlp::Protocol::Grpc)
.with_timeout(opentelemetry_otlp::OTEL_EXPORTER_OTLP_TIMEOUT_DEFAULT)
.build()?;
let attributes = resource_attributes
.into_iter()
.chain(iter::once(("service.name".into(), name.clone())))
.map(|(k, v)| KeyValue::new(k, v));
let reporter = OpenTelemetryReporter::new(
exporter,
Cow::Owned(Resource::builder().with_attributes(attributes).build()),
InstrumentationScope::builder(name)
.with_version(env!("CARGO_PKG_VERSION"))
.build(),
);
fastrace::set_reporter(reporter, fastrace::collector::Config::default());
Ok(())
}
pub fn current_trace_id() -> Option<String> {
SpanContext::current_local_parent().map(|current| current.encode_w3c_traceparent())
}
pub fn local_or_random_span_ctx() -> SpanContext {
SpanContext::current_local_parent().unwrap_or_else(SpanContext::random)
}
pub fn enter_span_with_traceparent(name: impl Into<Cow<'static, str>>, traceparent: &str) -> Span {
if let Some(span_ctx) = SpanContext::decode_w3c_traceparent(traceparent) {
Span::root(name, span_ctx)
} else {
Span::enter_with_local_parent(name)
}
}
// Tonic channel (TODO: Move into grpc_client when deprecated tracing)
#[allow(dead_code)]
pub type FastraceChannel = fastrace_tonic::FastraceClientService<tonic::transport::Channel>;
pub fn fastrace_channel(
channel: tonic::transport::Channel,
) -> fastrace_tonic::FastraceClientService<tonic::transport::Channel> {
tower::ServiceBuilder::new()
.layer(fastrace_tonic::FastraceClientLayer)
.service(channel)
}
// Request middleware (TODO: Move into http_client when deprecated tracing)
#[derive(Clone)]
#[allow(dead_code)]
pub struct TraceparentMiddleware;
#[async_trait::async_trait]
impl reqwest_middleware::Middleware for TraceparentMiddleware {
async fn handle(
&self,
mut req: reqwest::Request,
extensions: &mut http::Extensions,
next: reqwest_middleware::Next<'_>,
) -> reqwest_middleware::Result<reqwest::Response> {
req.headers_mut()
.extend(fastrace_reqwest::traceparent_headers());
next.run(req, extensions).await
}
}
crates/common/kigi-tracing/src/grpc_client.rs
-388
View File
@@ -1,388 +0,0 @@
use http::{HeaderMap, Request};
use opentelemetry::{global, propagation::Extractor, propagation::Injector};
use std::task::{Context, Poll};
use tonic::transport::Channel;
use tonic::{
Status,
metadata::{MetadataKey, MetadataMap, MetadataValue},
};
use tower::{Layer, Service, ServiceBuilder};
use tower_http::classify::{GrpcErrorsAsFailures, SharedClassifier};
use tower_http::trace::{MakeSpan, Trace, TraceLayer};
use tracing::{Span, warn};
use tracing_opentelemetry::OpenTelemetrySpanExt;
pub type TracedChannel = Trace<
InjectTraceContextService<Channel>,
SharedClassifier<GrpcErrorsAsFailures>,
MakeClientSpan,
>;
/// Wraps the input channel with a tracing layer. This function can be used to create a traced gRPC
/// client as follows:
///
/// ```rust
/// use tonic::transport::Endpoint;
/// use std::str::FromStr;
/// use kigi_tracing::traced_channel;
///
/// let channel = Endpoint::from_str("http://foo").unwrap();
/// //let client = SomeClient::new(traced_channel(channel));
///```
pub fn traced_channel(channel: Channel) -> TracedChannel {
ServiceBuilder::new()
.layer(TraceLayer::new_for_grpc().make_span_with(MakeClientSpan))
.layer(InjectTraceContextLayer)
.service(channel)
}
/// Implements the [`MakeSpan`] trait, to trace outgoing gRPC requests.
#[derive(Debug, Clone, Copy)]
pub struct MakeClientSpan;
impl<B> MakeSpan<B> for MakeClientSpan {
fn make_span(&mut self, request: &Request<B>) -> Span {
// No active dispatcher → the span has no consumer and would only be
// downgraded to `log` spam. See `crate::dispatcher_active`.
if !crate::dispatcher_active() {
return Span::none();
}
tracing::info_span!(
"grpc_request",
otel.kind = "client",
method = %request.method(),
uri = %request.uri(),
version = ?request.version(),
)
}
}
#[derive(Clone, Copy, Debug, Default)]
pub struct InjectTraceContextLayer;
impl<S> Layer<S> for InjectTraceContextLayer {
type Service = InjectTraceContextService<S>;
fn layer(&self, inner: S) -> Self::Service {
InjectTraceContextService { inner }
}
}
#[derive(Clone, Debug)]
pub struct InjectTraceContextService<S> {
inner: S,
}
impl<S, B> Service<Request<B>> for InjectTraceContextService<S>
where
S: Service<Request<B>>,
{
type Response = S::Response;
type Error = S::Error;
type Future = S::Future;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, mut req: Request<B>) -> Self::Future {
crate::http_client::attach_trace_to_http_request(req.headers_mut());
self.inner.call(req)
}
}
/// Inject W3C `traceparent` / `tracestate` from the active span into gRPC
/// metadata. Mutates in place so callers never lose the request body.
pub fn attach_trace_to_grpc_request_mut(metadata: &mut MetadataMap) {
global::get_text_map_propagator(|propagator| {
let context = Span::current().context();
propagator.inject_context(&context, &mut MetadataInjector(metadata));
});
}
/// Trace context propagation: send the trace context by injecting it into the metadata of the given
/// request.
pub fn attach_trace_to_grpc_request<T>(
mut request: tonic::Request<T>,
) -> Result<tonic::Request<T>, Status> {
attach_trace_to_grpc_request_mut(request.metadata_mut());
Ok(request)
}
// Need a custom Injector to inject OTel headers
pub struct MetadataInjector<'a>(&'a mut MetadataMap);
impl Injector for MetadataInjector<'_> {
fn set(&mut self, key: &str, value: String) {
match MetadataKey::from_bytes(key.as_bytes()) {
Ok(key) => match MetadataValue::try_from(&value) {
Ok(value) => {
self.0.insert(key, value);
}
Err(error) => warn!(value, error = format!("{error:#}"), "parse metadata value"),
},
Err(error) => warn!(key, error = format!("{error:#}"), "parse metadata key"),
}
}
}
pub struct HeaderExtractor<'a>(pub &'a HeaderMap);
impl Extractor for HeaderExtractor<'_> {
fn get(&self, key: &str) -> Option<&str> {
self.0.get(key).and_then(|v| {
let s = v.to_str();
if let Err(ref error) = s {
warn!(%error, ?v, "cannot convert header value to ASCII")
};
s.ok()
})
}
fn keys(&self) -> Vec<&str> {
self.0.keys().map(|k| k.as_str()).collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::testing::{OtelTestEnv, otel_span_id_hex, otel_trace_id_hex, parse_traceparent};
use http_body_util::Empty;
use std::convert::Infallible;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tower_http::classify::GrpcFailureClass;
use tracing::Instrument;
type EmptyBody = Empty<bytes::Bytes>;
#[derive(Clone)]
struct CaptureService {
seen: Arc<Mutex<Option<HeaderMap>>>,
response_grpc_status: Option<&'static str>,
}
impl CaptureService {
fn new() -> Self {
Self {
seen: Arc::new(Mutex::new(None)),
response_grpc_status: None,
}
}
fn with_grpc_status(status: &'static str) -> Self {
Self {
seen: Arc::new(Mutex::new(None)),
response_grpc_status: Some(status),
}
}
}
impl<B> Service<Request<B>> for CaptureService {
type Response = http::Response<EmptyBody>;
type Error = Infallible;
type Future = std::future::Ready<Result<Self::Response, Self::Error>>;
fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, req: Request<B>) -> Self::Future {
*self.seen.lock().unwrap() = Some(req.headers().clone());
let mut builder = http::Response::builder().status(200);
if let Some(status) = self.response_grpc_status {
builder = builder.header("grpc-status", status);
}
std::future::ready(Ok(builder.body(Empty::new()).unwrap()))
}
}
fn post_req() -> Request<EmptyBody> {
Request::builder()
.method("POST")
.uri("http://svc/package.Service/Method")
.body(Empty::new())
.unwrap()
}
// With no dispatcher active, the client span must not be created —
// `tracing`'s `log` compat would downgrade it into `grpc_request; ...`
// log spam in processes that only configure a `log` logger.
// See `crate::dispatcher_active`.
#[test]
fn make_client_span_without_dispatcher_is_none() {
assert!(MakeClientSpan.make_span(&post_req()).is_none());
}
#[test]
fn make_client_span_with_scoped_dispatcher_is_enabled() {
let _env = OtelTestEnv::install();
assert!(!MakeClientSpan.make_span(&post_req()).is_disabled());
}
#[tokio::test]
async fn inject_under_trace_layer_uses_client_span_not_parent() {
let _env = OtelTestEnv::install();
let capture = CaptureService::new();
let seen = Arc::clone(&capture.seen);
let mut svc = ServiceBuilder::new()
.layer(TraceLayer::new_for_grpc().make_span_with(MakeClientSpan))
.layer(InjectTraceContextLayer)
.service(capture);
let parent = tracing::info_span!("parent_handler");
let parent_span_id = otel_span_id_hex(&parent);
let parent_trace_id = otel_trace_id_hex(&parent);
assert_ne!(parent_span_id, "0000000000000000");
async {
let fut = Service::call(&mut svc, post_req());
fut.await.unwrap();
}
.instrument(parent)
.await;
let headers = seen.lock().unwrap().clone().expect("headers");
let tp = headers
.get("traceparent")
.expect("traceparent")
.to_str()
.unwrap();
let (_ver, injected_trace_id, injected_span_id) = parse_traceparent(tp);
assert_eq!(injected_trace_id, parent_trace_id);
assert_ne!(injected_span_id, parent_span_id);
}
#[tokio::test]
async fn inject_without_trace_layer_uses_parent_span() {
let _env = OtelTestEnv::install();
let capture = CaptureService::new();
let seen = Arc::clone(&capture.seen);
let mut svc = ServiceBuilder::new()
.layer(InjectTraceContextLayer)
.service(capture);
let parent = tracing::info_span!("parent_handler");
let parent_span_id = otel_span_id_hex(&parent);
let parent_trace_id = otel_trace_id_hex(&parent);
async {
let fut = Service::call(&mut svc, post_req());
fut.await.unwrap();
}
.instrument(parent)
.await;
let headers = seen.lock().unwrap().clone().expect("headers captured");
let tp = headers.get("traceparent").unwrap().to_str().unwrap();
let (_ver, injected_trace_id, injected_span_id) = parse_traceparent(tp);
assert_eq!(injected_trace_id, parent_trace_id);
assert_eq!(injected_span_id, parent_span_id);
}
#[tokio::test]
async fn grpc_status_non_ok_invokes_on_failure_classifier() {
let _env = OtelTestEnv::install();
let failures: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
let failures_cb = Arc::clone(&failures);
let capture = CaptureService::with_grpc_status("13");
let mut svc = ServiceBuilder::new()
.layer(
TraceLayer::new_for_grpc()
.make_span_with(MakeClientSpan)
.on_failure(
move |class: GrpcFailureClass,
_latency: Duration,
_span: &tracing::Span| {
failures_cb.lock().unwrap().push(class.to_string());
},
),
)
.layer(InjectTraceContextLayer)
.service(capture);
let fut = Service::call(&mut svc, post_req());
let _ = fut.await.unwrap();
let recorded = failures.lock().unwrap().clone();
assert_eq!(recorded.len(), 1, "{recorded:?}");
assert!(
recorded[0].contains("13") || recorded[0].to_lowercase().contains("code"),
"{}",
recorded[0]
);
}
#[tokio::test]
async fn grpc_status_ok_does_not_invoke_on_failure() {
let _env = OtelTestEnv::install();
let failures: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
let failures_cb = Arc::clone(&failures);
let capture = CaptureService::with_grpc_status("0");
let mut svc = ServiceBuilder::new()
.layer(
TraceLayer::new_for_grpc()
.make_span_with(MakeClientSpan)
.on_failure(
move |class: GrpcFailureClass,
_latency: Duration,
_span: &tracing::Span| {
failures_cb.lock().unwrap().push(class.to_string());
},
),
)
.layer(InjectTraceContextLayer)
.service(capture);
let fut = Service::call(&mut svc, post_req());
let _ = fut.await.unwrap();
assert!(failures.lock().unwrap().is_empty());
}
#[tokio::test]
async fn attach_trace_to_grpc_request_sets_traceparent_metadata() {
let _env = OtelTestEnv::install();
let span = tracing::info_span!("handler");
let span_id = otel_span_id_hex(&span);
let _enter = span.enter();
let req = attach_trace_to_grpc_request(tonic::Request::new(())).unwrap();
let tp = req
.metadata()
.get("traceparent")
.expect("traceparent in metadata")
.to_str()
.unwrap();
let (_ver, _tid, injected_span_id) = parse_traceparent(tp);
assert_eq!(injected_span_id, span_id);
}
#[tokio::test]
async fn make_client_span_records_otel_kind_client() {
let env = OtelTestEnv::install();
{
let req = post_req();
let mut make = MakeClientSpan;
let span = make.make_span(&req);
let _e = span.enter();
}
let spans = env.finished_spans();
let grpc = spans
.iter()
.find(|s| s.name == "grpc_request")
.expect("grpc_request");
assert_eq!(grpc.span_kind, opentelemetry::trace::SpanKind::Client);
}
}
crates/common/kigi-tracing/src/http_client.rs
-155
View File
@@ -1,155 +0,0 @@
use async_trait::async_trait;
use opentelemetry::global;
use opentelemetry_http::HeaderInjector;
use reqwest::header::HeaderMap;
use reqwest_middleware::{ClientBuilder, ClientWithMiddleware, Middleware};
use tracing::{Instrument, Span, field};
use tracing_opentelemetry::OpenTelemetrySpanExt;
pub fn attach_trace_to_http_request(headers: &mut HeaderMap) {
global::get_text_map_propagator(|propagator| {
let context = Span::current().context();
propagator.inject_context(&context, &mut HeaderInjector(headers));
});
}
pub type TracedHttpClient = ClientWithMiddleware;
pub fn traced_client(client: reqwest::Client) -> TracedHttpClient {
ClientBuilder::new(client).with(TracingMiddleware).build()
}
pub fn traced_client_new() -> TracedHttpClient {
traced_client(reqwest::Client::new())
}
pub fn traced_client_from_builder(
builder: reqwest::ClientBuilder,
) -> Result<TracedHttpClient, reqwest::Error> {
Ok(traced_client(builder.build()?))
}
#[derive(Clone, Debug, Default)]
struct TracingMiddleware;
#[async_trait]
impl Middleware for TracingMiddleware {
async fn handle(
&self,
mut req: reqwest::Request,
extensions: &mut http::Extensions,
next: reqwest_middleware::Next<'_>,
) -> reqwest_middleware::Result<reqwest::Response> {
let method = req.method().as_str().to_owned();
let url = req.url().clone();
// No active dispatcher → the span has no consumer and would only be
// downgraded to `log` spam. See `crate::dispatcher_active`.
let span = if crate::dispatcher_active() {
tracing::info_span!(
"http_request",
otel.kind = "client",
"http.request.method" = %method,
"url.full" = %url,
"http.response.status_code" = field::Empty,
)
} else {
Span::none()
};
let result = async move {
attach_trace_to_http_request(req.headers_mut());
next.run(req, extensions).await
}
.instrument(span.clone())
.await;
if let Ok(ref response) = result {
span.record("http.response.status_code", response.status().as_u16());
}
result
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::testing::{OtelTestEnv, otel_span_id_hex, otel_trace_id_hex, parse_traceparent};
use tracing::Instrument;
use wiremock::matchers::method;
use wiremock::{Mock, MockServer, ResponseTemplate};
#[tokio::test]
async fn attach_trace_to_http_request_writes_traceparent() {
let _env = OtelTestEnv::install();
let span = tracing::info_span!("http_request", otel.kind = "client");
let span_id = otel_span_id_hex(&span);
let mut headers = HeaderMap::new();
let _enter = span.enter();
attach_trace_to_http_request(&mut headers);
let tp = headers.get("traceparent").unwrap().to_str().unwrap();
let (_ver, _tid, injected_span_id) = parse_traceparent(tp);
assert_eq!(injected_span_id, span_id);
}
#[tokio::test]
async fn traced_client_injects_client_span_not_parent_on_wire() {
let _env = OtelTestEnv::install();
let server = MockServer::start().await;
Mock::given(method("GET"))
.respond_with(ResponseTemplate::new(200))
.mount(&server)
.await;
let client = traced_client(reqwest::Client::new());
let parent = tracing::info_span!("parent_handler");
let parent_span_id = otel_span_id_hex(&parent);
let parent_trace_id = otel_trace_id_hex(&parent);
async {
client
.get(format!("{}/health", server.uri()))
.send()
.await
.unwrap();
}
.instrument(parent)
.await;
let received = server.received_requests().await.unwrap();
assert_eq!(received.len(), 1);
let tp = received[0]
.headers
.get("traceparent")
.expect("traceparent on wire")
.to_str()
.unwrap();
let (_ver, injected_trace_id, injected_span_id) = parse_traceparent(tp);
assert_eq!(injected_trace_id, parent_trace_id);
assert_ne!(injected_span_id, parent_span_id);
}
#[tokio::test]
async fn traced_client_returns_response_status() {
let _env = OtelTestEnv::install();
let server = MockServer::start().await;
Mock::given(method("GET"))
.respond_with(ResponseTemplate::new(404))
.mount(&server)
.await;
let client = traced_client_new();
let resp = client
.get(format!("{}/missing", server.uri()))
.send()
.await
.unwrap();
assert_eq!(resp.status().as_u16(), 404);
assert!(
server.received_requests().await.unwrap()[0]
.headers
.get("traceparent")
.is_some()
);
}
}
crates/common/kigi-tracing/src/lib.rs
-19
View File
@@ -1,19 +0,0 @@
mod dispatch;
mod grpc_client;
mod timer;
pub mod fastrace;
pub mod http_client;
pub mod tokio;
#[cfg(test)]
mod testing;
pub use dispatch::*;
pub use fastrace::*;
pub use grpc_client::*;
pub use http_client::{
TracedHttpClient, attach_trace_to_http_request, traced_client, traced_client_from_builder,
traced_client_new,
};
pub use timer::*;
crates/common/kigi-tracing/src/testing.rs
-65
View File
@@ -1,65 +0,0 @@
use opentelemetry::global;
use opentelemetry::trace::{SpanContext, TraceContextExt, TracerProvider as _};
use opentelemetry_sdk::propagation::TraceContextPropagator;
use opentelemetry_sdk::trace::{
InMemorySpanExporter, InMemorySpanExporterBuilder, SdkTracerProvider, SimpleSpanProcessor,
};
use tracing::Span;
use tracing_opentelemetry::OpenTelemetrySpanExt;
use tracing_subscriber::prelude::*;
pub struct OtelTestEnv {
_guard: tracing::subscriber::DefaultGuard,
provider: SdkTracerProvider,
exporter: InMemorySpanExporter,
}
impl OtelTestEnv {
pub fn install() -> Self {
global::set_text_map_propagator(TraceContextPropagator::new());
let exporter = InMemorySpanExporterBuilder::new().build();
let provider = SdkTracerProvider::builder()
.with_span_processor(SimpleSpanProcessor::new(exporter.clone()))
.build();
let tracer = provider.tracer("kigi-tracing-test");
let otel_layer = tracing_opentelemetry::layer()
.with_tracer(tracer)
.with_context_activation(false)
.with_filter(tracing_subscriber::filter::LevelFilter::INFO);
let guard = tracing_subscriber::registry()
.with(otel_layer)
.set_default();
Self {
_guard: guard,
provider,
exporter,
}
}
pub fn finished_spans(&self) -> Vec<opentelemetry_sdk::trace::SpanData> {
let _ = self.provider.force_flush();
self.exporter.get_finished_spans().unwrap_or_default()
}
}
pub fn parse_traceparent(value: &str) -> (&str, &str, &str) {
let mut parts = value.split('-');
let version = parts.next().expect("version");
let trace_id = parts.next().expect("trace_id");
let span_id = parts.next().expect("span_id");
(version, trace_id, span_id)
}
pub fn otel_span_id_hex(span: &Span) -> String {
let cx = span.context();
let span_ref = cx.span();
let sc: &SpanContext = span_ref.span_context();
format!("{:016x}", u64::from_be_bytes(sc.span_id().to_bytes()))
}
pub fn otel_trace_id_hex(span: &Span) -> String {
let cx = span.context();
let span_ref = cx.span();
let sc: &SpanContext = span_ref.span_context();
format!("{:032x}", u128::from_be_bytes(sc.trace_id().to_bytes()))
}
crates/common/kigi-tracing/src/timer.rs
-62
View File
@@ -1,62 +0,0 @@
use log::error;
use log::info;
use tokio::time::Instant;
/// A simple timer that logs the runtime of an operation.
pub struct Timer {
/// Time the operation started.
start: Instant,
/// An ID shown in the logs to associate the log messages from the timer with each other-
id: uuid::Uuid,
/// A string that is being logged.
message: String,
/// True if the timer has been stopped already.
stopped: bool,
}
impl Timer {
/// Creates a new Timer instance and starts the timer.
pub fn new<S: AsRef<str>>(message: S) -> Self {
let id = uuid::Uuid::new_v4();
info!("[{}] START: {}", id, message.as_ref());
Self {
start: Instant::now(),
id,
message: message.as_ref().to_string(),
stopped: false,
}
}
/// Stops the timer and logs the result.
pub fn stop<T>(&mut self, result: T) -> T {
if !self.stopped {
let runtime = self.start.elapsed().as_secs_f32();
info!(
"[{}] FINISHED in {:.3}s: {}",
self.id, runtime, self.message
);
self.stopped = true;
}
result
}
/// Stop the timer prematurely and logs an error.
pub fn force_stop(&mut self) {
if !self.stopped {
let runtime = self.start.elapsed().as_secs_f32();
error!(
"[{}] FAILED after {:.3}s: {}",
self.id, runtime, self.message
);
self.stopped = true;
}
}
}
/// Automatically report the runtime when the object is dropped.
impl Drop for Timer {
fn drop(&mut self) {
self.force_stop();
}
}
crates/common/kigi-tracing/src/tokio.rs
-22
View File
@@ -1,22 +0,0 @@
use std::future::Future;
use tokio::task::JoinHandle;
use tracing::{Instrument, Span};
/// Utility macro for propagating the current tracing context to a newly spawned task.
///
/// Note: The spawned task will be associated with the currently active span. To create a *new*
/// span for the spawned task, manually instrument the future using [tracing::Instrument] instead of
/// using this macro. For example:
///
/// use tracing::{info_span, Instrument};
///
/// let fut = tokio::spawn(async move {
/// print!("do stuff")
/// }.instrument(info_span!("spawned task")));
pub fn spawn_traced<F>(future: F) -> JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
tokio::spawn(future.instrument(Span::current()))
}