From fa75eb139a5f16c5e2550453227c30c0aaf09bb7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Fri, 17 Jul 2026 22:34:10 -0400 Subject: [PATCH] =?UTF-8?q?M2=20audit:=20excise=20the=20Computer=20Hub=20s?= =?UTF-8?q?tack=20=E2=80=94=20Kigi's=20last=20remote-cloud=20surface?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- Cargo.lock | 402 +- Cargo.toml | 20 - crates/codegen/kigi-agent/src/builder.rs | 1 - crates/codegen/kigi-bin/src/main.rs | 214 +- crates/codegen/kigi-config-types/src/lib.rs | 23 - crates/codegen/kigi-mcp/Cargo.toml | 1 - crates/codegen/kigi-mcp/src/servers.rs | 2 +- crates/codegen/kigi-memory/src/schema.rs | 5 +- .../kigi-pager-pty-harness/src/content.rs | 4 +- crates/codegen/kigi-shell/Cargo.toml | 4 +- crates/codegen/kigi-shell/src/agent/app.rs | 8 +- crates/codegen/kigi-shell/src/agent/config.rs | 607 +-- crates/codegen/kigi-shell/src/agent/mod.rs | 1 - crates/codegen/kigi-shell/src/agent/models.rs | 14 +- .../src/agent/mvp_agent/agent_ops.rs | 6 +- .../kigi-shell/src/agent/mvp_agent/mod.rs | 2 - .../kigi-shell/src/agent/mvp_agent/tests.rs | 34 - .../tests/subagent_spawn_context_tests.rs | 1 + crates/codegen/kigi-shell/src/agent/proxy.rs | 619 --- crates/codegen/kigi-shell/src/config/mod.rs | 8 +- crates/codegen/kigi-shell/src/config/tests.rs | 10 +- crates/codegen/kigi-shell/src/inspect/mod.rs | 6 +- crates/codegen/kigi-shell/src/leader/mod.rs | 2 +- .../codegen/kigi-shell/src/leader/protocol.rs | 93 - .../codegen/kigi-shell/src/leader/server.rs | 367 -- .../kigi-shell/src/leader/test_support.rs | 1 - .../kigi-shell/src/session/acp_session.rs | 5 - .../src/session/acp_session_impl/spawn.rs | 35 +- .../session/acp_session_impl/tool_calls.rs | 19 - .../src/session/acp_session_impl/turn.rs | 32 - .../cancel_running_task_tests.rs | 4 - .../fs_injection_regression_tests.rs | 1 - .../acp_session_tests/idle_resume_tests.rs | 1 - .../inline_auto_compact_flow_tests.rs | 3 - .../acp_session_tests/memory_config_tests.rs | 1 - .../permission_auto_mode_tests.rs | 1 + .../replay_buffer_send_update_tests.rs | 1 - .../src/session/acp_session_tests/support.rs | 8 - .../acp_session_tests/web_search_e2e_tests.rs | 1 - .../kigi-shell/src/session/compaction.rs | 1 - .../src/session/storage/jsonl/tests.rs | 2 +- .../kigi-shell/tests/git_contention_e2e.rs | 2 +- .../kigi-shell/tests/session_load_perf.rs | 2 +- .../kigi-shell/tests/test_built_binary_e2e.rs | 2 +- .../kigi-shell/tests/test_leader_soak.rs | 2 +- .../tests/test_mcp_permission_persistence.rs | 10 +- crates/codegen/kigi-test-support/Cargo.toml | 1 - crates/codegen/kigi-test-support/src/env.rs | 2 +- .../codegen/kigi-test-support/src/leader.rs | 2 +- crates/codegen/kigi-tools-api/build.rs | 5 +- crates/codegen/kigi-tools/Cargo.toml | 3 - .../grok_build_hashline/edit/apply.rs | 4 +- .../src/implementations/search_tool/mod.rs | 2 +- .../codegen/kigi-tools/src/registry/types.rs | 27 +- crates/codegen/kigi-tui/src/acp/mod.rs | 3 - crates/codegen/kigi-tui/src/app/cli.rs | 77 +- .../codegen/kigi-tui/src/app/effects/tests.rs | 2 +- .../kigi-tui/src/app/leader_cluster/mod.rs | 2 +- crates/codegen/kigi-tui/src/app/mod.rs | 1 - crates/codegen/kigi-tui/src/headless.rs | 3 - crates/codegen/kigi-tui/src/trace_cmd.rs | 4 +- crates/codegen/kigi-workspace/Cargo.toml | 37 +- crates/codegen/kigi-workspace/src/config.rs | 863 ---- crates/codegen/kigi-workspace/src/error.rs | 6 +- .../kigi-workspace/src/file_system/ext_fs.rs | 16 +- .../kigi-workspace/src/file_system/mod.rs | 7 +- crates/codegen/kigi-workspace/src/handle.rs | 3975 +---------------- crates/codegen/kigi-workspace/src/lib.rs | 36 +- .../kigi-workspace/src/permission/manager.rs | 232 +- .../kigi-workspace/src/permission/mod.rs | 15 +- .../kigi-workspace/src/permission/prompter.rs | 72 +- .../codegen/kigi-workspace/src/session/mod.rs | 125 +- .../kigi-workspace/src/session/tool_config.rs | 432 +- .../kigi-workspace/src/workspace_ops.rs | 414 +- crates/common/kigi-tool-protocol/Cargo.toml | 2 +- crates/common/kigi-tool-protocol/src/lib.rs | 2 +- crates/common/kigi-tool-runtime/Cargo.toml | 5 +- .../common/kigi-tool-runtime/src/context.rs | 182 - crates/common/kigi-tool-runtime/src/lib.rs | 12 +- .../kigi-tool-runtime/src/local_registry.rs | 117 + crates/common/kigi-tool-types/Cargo.toml | 2 +- crates/common/kigi-tracing/Cargo.toml | 37 - crates/common/kigi-tracing/src/dispatch.rs | 45 - crates/common/kigi-tracing/src/fastrace.rs | 83 - crates/common/kigi-tracing/src/grpc_client.rs | 388 -- crates/common/kigi-tracing/src/http_client.rs | 155 - crates/common/kigi-tracing/src/lib.rs | 19 - crates/common/kigi-tracing/src/testing.rs | 65 - crates/common/kigi-tracing/src/timer.rs | 62 - crates/common/kigi-tracing/src/tokio.rs | 22 - 90 files changed, 452 insertions(+), 9702 deletions(-) delete mode 100644 crates/codegen/kigi-shell/src/agent/proxy.rs create mode 100644 crates/common/kigi-tool-runtime/src/local_registry.rs delete mode 100644 crates/common/kigi-tracing/Cargo.toml delete mode 100644 crates/common/kigi-tracing/src/dispatch.rs delete mode 100644 crates/common/kigi-tracing/src/fastrace.rs delete mode 100644 crates/common/kigi-tracing/src/grpc_client.rs delete mode 100644 crates/common/kigi-tracing/src/http_client.rs delete mode 100644 crates/common/kigi-tracing/src/lib.rs delete mode 100644 crates/common/kigi-tracing/src/testing.rs delete mode 100644 crates/common/kigi-tracing/src/timer.rs delete mode 100644 crates/common/kigi-tracing/src/tokio.rs diff --git a/Cargo.lock b/Cargo.lock index 13f450e..9a6710f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1177,7 +1177,7 @@ dependencies = [ "sha1 0.10.7", "sync_wrapper", "tokio", - "tokio-tungstenite 0.29.0", + "tokio-tungstenite", "tower", "tower-layer", "tower-service", @@ -1912,18 +1912,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "const-hex" -version = "1.19.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "33e2a781ebdf4467d1428dc4593067825fb646f6871475098d8577421af73558" -dependencies = [ - "cfg-if", - "cpufeatures 0.2.17", - "proptest", - "serde_core", -] - [[package]] name = "const-oid" version = "0.9.6" @@ -3110,16 +3098,6 @@ dependencies = [ "thiserror 2.0.18", ] -[[package]] -name = "fastant" -version = "0.1.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e825441bfb2d831c47c97d05821552db8832479f44c571b97fededbf0099c07" -dependencies = [ - "small_ctor", - "web-time", -] - [[package]] name = "faster-hex" version = "0.10.0" @@ -3130,67 +3108,6 @@ dependencies = [ "serde", ] -[[package]] -name = "fastrace" -version = "0.7.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74539685f6baa3c33251571272eac7939fc3b93dff97c30381d7a6267394e528" -dependencies = [ - "fastant", - "fastrace-macro", - "parking_lot", - "pin-project", - "rand 0.10.2", - "rtrb", - "serde", -] - -[[package]] -name = "fastrace-macro" -version = "0.7.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "390d4f754b751905fa9294716c4ead717ffb5fa733a0c2417358b7779fceaeed" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.119", -] - -[[package]] -name = "fastrace-opentelemetry" -version = "0.18.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67d87b05a82ea59d6d2a11066f4e0cc1e169ce50bc8c276279249ba374936d91" -dependencies = [ - "fastrace", - "log", - "opentelemetry", - "opentelemetry_sdk", - "pollster", -] - -[[package]] -name = "fastrace-reqwest" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4705b6c98308ef10f178dadc516a460ca9191c0e1c6e07c73104c9679e82d4e" -dependencies = [ - "fastrace", - "reqwest 0.12.28", -] - -[[package]] -name = "fastrace-tonic" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9bfe91596cd82e4b4fbd390f0e789fe3cf327b635fa713347f1509db97aa482c" -dependencies = [ - "fastrace", - "http 1.4.2", - "tower-layer", - "tower-service", -] - [[package]] name = "fastrand" version = "2.4.1" @@ -4686,16 +4603,6 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bab0c77d09c0d9b144d429d0bf0fcd4c63c01d5f6c33f9b8ed501283e0f1ef76" -[[package]] -name = "hdrhistogram" -version = "7.5.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "765c9198f173dd59ce26ff9f95ef0aafd0a0fe01fb9d72841bc5066a4c06511d" -dependencies = [ - "byteorder", - "num-traits", -] - [[package]] name = "heapless" version = "0.8.0" @@ -4927,7 +4834,7 @@ dependencies = [ "tokio", "tokio-rustls", "tower-service", - "webpki-roots 1.0.8", + "webpki-roots", ] [[package]] @@ -5700,81 +5607,6 @@ dependencies = [ "tracing", ] -[[package]] -name = "kigi-computer-hub-core" -version = "0.1.0" -dependencies = [ - "async-trait", - "chrono", - "dashmap", - "futures", - "kigi-tool-protocol", - "kigi-tool-runtime", - "kigi-tool-types", - "schemars 1.2.1", - "serde", - "serde_json", - "tokio", - "tracing", -] - -[[package]] -name = "kigi-computer-hub-mcp-adapter" -version = "0.1.0" -dependencies = [ - "async-trait", - "futures", - "kigi-computer-hub-sdk", - "kigi-tool-protocol", - "kigi-tool-runtime", - "kigi-tool-types", - "prometheus", - "serde", - "serde_json", - "thiserror 2.0.18", - "tokio", - "tracing", -] - -[[package]] -name = "kigi-computer-hub-sdk" -version = "0.1.0" -dependencies = [ - "arc-swap", - "async-trait", - "axum", - "base64", - "chrono", - "dashmap", - "fastrace", - "fastrace-opentelemetry", - "futures", - "http 1.4.2", - "indexmap", - "kigi-computer-hub-core", - "kigi-tool-protocol", - "kigi-tool-runtime", - "kigi-tool-types", - "kigi-tracing", - "opentelemetry", - "opentelemetry-proto", - "opentelemetry_sdk", - "parking_lot", - "prometheus", - "prost", - "reqwest 0.12.28", - "schemars 1.2.1", - "serde", - "serde_json", - "thiserror 2.0.18", - "tokio", - "tokio-tungstenite 0.27.0", - "tokio-util", - "tracing", - "tracing-subscriber", - "url", -] - [[package]] name = "kigi-config" version = "0.1.0" @@ -6046,7 +5878,6 @@ dependencies = [ "axum", "futures", "http 1.4.2", - "kigi-computer-hub-sdk", "kigi-config", "kigi-file-utils", "kigi-log", @@ -6440,7 +6271,6 @@ dependencies = [ "kigi-chat-state", "kigi-codebase-graph", "kigi-compaction", - "kigi-computer-hub-sdk", "kigi-config", "kigi-config-types", "kigi-env", @@ -6516,7 +6346,6 @@ dependencies = [ "tokio-retry", "tokio-rustls", "tokio-stream", - "tokio-tungstenite 0.27.0", "tokio-util", "toml", "toml_edit 0.22.27", @@ -6632,7 +6461,6 @@ dependencies = [ "serde_json", "tempfile", "tokio", - "tokio-tungstenite 0.27.0", "tokio-util", "tracing", "tracing-subscriber", @@ -6669,9 +6497,10 @@ dependencies = [ "anyhow", "async-trait", "futures", + "indexmap", "kigi-tool-protocol", "kigi-tool-types", - "kigi-tools-api", + "parking_lot", "schemars 1.2.1", "serde", "serde_json", @@ -6717,8 +6546,6 @@ dependencies = [ "image", "indexmap", "infer 0.19.0", - "kigi-computer-hub-core", - "kigi-computer-hub-sdk", "kigi-config", "kigi-env", "kigi-file-utils", @@ -6781,36 +6608,6 @@ dependencies = [ "tonic-prost", ] -[[package]] -name = "kigi-tracing" -version = "0.1.0" -dependencies = [ - "async-trait", - "bytes", - "fastrace", - "fastrace-opentelemetry", - "fastrace-reqwest", - "fastrace-tonic", - "http 1.4.2", - "http-body-util", - "log", - "opentelemetry", - "opentelemetry-http", - "opentelemetry-otlp", - "opentelemetry_sdk", - "reqwest 0.12.28", - "reqwest-middleware", - "tokio", - "tonic", - "tower", - "tower-http", - "tracing", - "tracing-opentelemetry", - "tracing-subscriber", - "uuid", - "wiremock", -] - [[package]] name = "kigi-tracing-macros" version = "0.1.0" @@ -6953,17 +6750,13 @@ dependencies = [ "arc-swap", "async-stream", "async-trait", - "axum", "base64", "bstr", "chrono", - "clap", "dashmap", "dirs 6.0.0", "dunce", - "fastrace", "filetime", - "flate2", "fs2", "futures", "git2", @@ -6977,8 +6770,6 @@ dependencies = [ "kigi-agent", "kigi-auth", "kigi-codebase-graph", - "kigi-computer-hub-mcp-adapter", - "kigi-computer-hub-sdk", "kigi-config", "kigi-config-types", "kigi-env", @@ -6998,10 +6789,8 @@ dependencies = [ "kigi-tool-types", "kigi-tools", "kigi-tools-api", - "kigi-tracing", "kigi-tty-utils", "kigi-version", - "kigi-workspace-client", "kigi-workspace-types", "libc", "md5", @@ -7010,22 +6799,17 @@ dependencies = [ "parking_lot", "prometheus", "regex", - "reqwest 0.12.28", "rusqlite", "rustc-hash", - "rustls", "serde", "serde_json", "sha2 0.10.9", "smallvec", - "tar", "tempfile", "thiserror 2.0.18", "tokio", - "tokio-tungstenite 0.27.0", "tokio-util", "toml", - "tonic", "tracing", "tracing-subscriber", "tree-sitter", @@ -7033,27 +6817,9 @@ dependencies = [ "url", "urlencoding", "uuid", - "windows 0.61.3", "zstd", ] -[[package]] -name = "kigi-workspace-client" -version = "0.1.0" -dependencies = [ - "kigi-computer-hub-sdk", - "kigi-tool-protocol", - "kigi-tool-runtime", - "kigi-tool-types", - "kigi-workspace-types", - "schemars 1.2.1", - "serde", - "serde_json", - "thiserror 2.0.18", - "tokio", - "tracing", -] - [[package]] name = "kigi-workspace-types" version = "0.1.0" @@ -8184,54 +7950,6 @@ dependencies = [ "tracing", ] -[[package]] -name = "opentelemetry-http" -version = "0.32.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5683015d09e2df236ef005b17f6f196f0d5f6313c4fa43a7b6a53b52776e4331" -dependencies = [ - "async-trait", - "bytes", - "http 1.4.2", - "opentelemetry", - "reqwest 0.13.4", -] - -[[package]] -name = "opentelemetry-otlp" -version = "0.32.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9966929966d17620d7c316c643ba62631826e10021409357772d5eea84f62c35" -dependencies = [ - "http 1.4.2", - "opentelemetry", - "opentelemetry-http", - "opentelemetry-proto", - "opentelemetry_sdk", - "prost", - "reqwest 0.13.4", - "thiserror 2.0.18", - "tokio", - "tonic", - "tonic-types", -] - -[[package]] -name = "opentelemetry-proto" -version = "0.32.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "56d658ba1faf63f7b9c492cfbe6e0ec365440a16132d3270c1065f7b33f1b638" -dependencies = [ - "base64", - "const-hex", - "opentelemetry", - "opentelemetry_sdk", - "prost", - "serde", - "tonic", - "tonic-prost", -] - [[package]] name = "opentelemetry_sdk" version = "0.32.1" @@ -8246,7 +7964,6 @@ dependencies = [ "portable-atomic", "rand 0.9.5", "thiserror 2.0.18", - "tokio", ] [[package]] @@ -8859,12 +8576,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "pollster" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f3a9f18d041e6d0e102a0a46750538147e5e8992d3b4873aaafee2520b00ce3" - [[package]] name = "polycool" version = "0.4.0" @@ -9070,21 +8781,6 @@ dependencies = [ "thiserror 2.0.18", ] -[[package]] -name = "proptest" -version = "1.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b45fcc2344c680f5025fe57779faef368840d0bd1f42f216291f0dc4ace4744" -dependencies = [ - "bitflags 2.13.1", - "num-traits", - "rand 0.9.5", - "rand_chacha 0.9.0", - "rand_xorshift", - "regex-syntax", - "unarray", -] - [[package]] name = "prost" version = "0.14.4" @@ -9421,15 +9117,6 @@ dependencies = [ "rand_core 0.10.1", ] -[[package]] -name = "rand_xorshift" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "513962919efc330f829edb2535844d1b912b0fbe2ca165d613e4e8788bb05a5a" -dependencies = [ - "rand_core 0.9.5", -] - [[package]] name = "rapidhash" version = "4.5.1" @@ -9673,7 +9360,7 @@ dependencies = [ "wasm-bindgen-futures", "wasm-streams 0.4.2", "web-sys", - "webpki-roots 1.0.8", + "webpki-roots", ] [[package]] @@ -9872,12 +9559,6 @@ dependencies = [ "zeroize", ] -[[package]] -name = "rtrb" -version = "0.3.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ade083ccbb4bf536df69d1f6432cc23deb7acccff86b183f3923a6fd56a1153" - [[package]] name = "runfiles" version = "0.1.0" @@ -10852,12 +10533,6 @@ dependencies = [ "version_check", ] -[[package]] -name = "small_ctor" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88414a5ca1f85d82cc34471e975f0f74f6aa54c40f062efa42c0080e7f763f81" - [[package]] name = "smallvec" version = "1.15.2" @@ -11650,23 +11325,6 @@ dependencies = [ "tokio", ] -[[package]] -name = "tokio-tungstenite" -version = "0.27.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "489a59b6730eda1b0171fcfda8b121f4bee2b35cba8645ca35c5f7ba3eb736c1" -dependencies = [ - "futures-util", - "log", - "rustls", - "rustls-native-certs", - "rustls-pki-types", - "tokio", - "tokio-rustls", - "tungstenite 0.27.0", - "webpki-roots 0.26.11", -] - [[package]] name = "tokio-tungstenite" version = "0.29.0" @@ -11676,7 +11334,7 @@ dependencies = [ "futures-util", "log", "tokio", - "tungstenite 0.29.0", + "tungstenite", ] [[package]] @@ -11803,7 +11461,6 @@ dependencies = [ "hyper-util", "percent-encoding", "pin-project", - "rustls-native-certs", "socket2", "sync_wrapper", "tokio", @@ -11855,17 +11512,6 @@ dependencies = [ "tonic-build", ] -[[package]] -name = "tonic-types" -version = "0.14.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73ab1b02061f83d519bba3caa167f88f261ef05720ab8ebc954ade70de3348e8" -dependencies = [ - "prost", - "prost-types", - "tonic", -] - [[package]] name = "tower" version = "0.5.3" @@ -11874,7 +11520,6 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", - "hdrhistogram", "indexmap", "pin-project-lite", "slab", @@ -11906,7 +11551,6 @@ dependencies = [ "tower", "tower-layer", "tower-service", - "tracing", "url", ] @@ -12154,25 +11798,6 @@ dependencies = [ "ratatui-core", ] -[[package]] -name = "tungstenite" -version = "0.27.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eadc29d668c91fcc564941132e17b28a7ceb2f3ebf0b9dae3e03fd7a6748eb0d" -dependencies = [ - "bytes", - "data-encoding", - "http 1.4.2", - "httparse", - "log", - "rand 0.9.5", - "rustls", - "rustls-pki-types", - "sha1 0.10.7", - "thiserror 2.0.18", - "utf-8", -] - [[package]] name = "tungstenite" version = "0.29.0" @@ -12276,12 +11901,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "unarray" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94" - [[package]] name = "unicase" version = "2.9.0" @@ -12834,15 +12453,6 @@ dependencies = [ "rustls-pki-types", ] -[[package]] -name = "webpki-roots" -version = "0.26.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" -dependencies = [ - "webpki-roots 1.0.8", -] - [[package]] name = "webpki-roots" version = "1.0.8" diff --git a/Cargo.toml b/Cargo.toml index cf335a1..0d44d4c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -52,7 +52,6 @@ members = [ "crates/codegen/kigi-update", "crates/codegen/kigi-version", "crates/codegen/kigi-workspace", - "crates/codegen/kigi-workspace-client", "crates/codegen/kigi-workspace-types", "crates/codegen/kigi-hooks-plugins-types", "crates/codegen/kigi-hunk-tracker", @@ -65,16 +64,12 @@ members = [ "crates/codegen/kigi-token-estimation", "crates/codegen/kigi-tracing-macros", "crates/codegen/kigi-tty-utils", - "crates/common/kigi-computer-hub-core", - "crates/common/kigi-computer-hub-mcp-adapter", - "crates/common/kigi-computer-hub-sdk", "crates/common/kigi-compaction", "crates/common/kigi-interjection-core", "crates/common/kigi-test-utils", "crates/common/kigi-tool-protocol", "crates/common/kigi-tool-runtime", "crates/common/kigi-tool-types", - "crates/common/kigi-tracing", "third_party/dagre_rust", "third_party/graphlib_rust", "third_party/mermaid-to-svg", @@ -130,10 +125,6 @@ educe = "0.6.0" enum_delegate = "0.2" env_logger = "0.11" eventsource-stream = "0.2" -fastrace = { version = "0.7" } -fastrace-opentelemetry = { version = "0.18" } -fastrace-reqwest = { version = "0.2" } -fastrace-tonic = { version = "0.1" } fastrand = "2" filetime = "0.2.25" flate2 = { version = "1", default-features = false, features = ["zlib-rs"] } @@ -178,9 +169,6 @@ oauth2 = "5" obfstr = "0.4" once_cell = "1" opentelemetry = "0.32" -opentelemetry-http = { version = "0.32", features = ["reqwest-blocking"] } -opentelemetry-otlp = { version = "0.32", features = ["grpc-tonic", "reqwest-blocking-client", "tls-roots"] } -opentelemetry-proto = { version = "0.32", features = ["gen-tonic"] } opentelemetry_sdk = { version = "0.32.1", features = ["spec_unstable_metrics_views"] } parking_lot = "0.12.4" pbjson-build = "0.9" @@ -239,7 +227,6 @@ tiny-skia = "0.12" tokio = { version = "1", features = ["full"] } tokio-retry = "0.3" tokio-stream = "0.1" -tokio-tungstenite = "0.27" tokio-util = { version = "0.7", features = ["rt"] } toml = "0.9" toml_edit = "0.22" @@ -269,8 +256,6 @@ wiremock = "0.6" wl-clipboard-rs = "0.9" kigi-acp-lib = { path = "crates/codegen/kigi-acp-lib" } kigi-agent-lifecycle = { path = "crates/codegen/kigi-agent-lifecycle" } -kigi-computer-hub-core = { path = "crates/common/kigi-computer-hub-core" } -kigi-computer-hub-sdk = { path = "crates/common/kigi-computer-hub-sdk" } kigi-gix-status = { path = "crates/codegen/kigi-gix-status" } kigi-agent = { path = "crates/codegen/kigi-agent" } kigi-auth = { path = "crates/codegen/kigi-auth" } @@ -310,7 +295,6 @@ kigi-token-estimation = { path = "crates/codegen/kigi-token-estimation" } kigi-tool-protocol = { path = "crates/common/kigi-tool-protocol" } kigi-tool-runtime = { path = "crates/common/kigi-tool-runtime" } kigi-tool-types = { path = "crates/common/kigi-tool-types" } -kigi-tracing = { path = "crates/common/kigi-tracing" } kigi-tty-utils = { path = "crates/codegen/kigi-tty-utils" } zbus = { version = "5" } zstd = "0.13" @@ -388,7 +372,3 @@ too_many_arguments = "allow" # TODO: -> "deny" once/if merge queue enabled uninlined_format_args = "allow" -# The `fastrace::trace(properties = { … })` proc-macro expands `"{param}"` -# into `format!("{}", param)`, which clippy flags as useless_format. -# Suppressed until the upstream crate fixes its codegen (fixed in 0.7.16+). -useless_format = "allow" diff --git a/crates/codegen/kigi-agent/src/builder.rs b/crates/codegen/kigi-agent/src/builder.rs index 4e642a5..29195b5 100644 --- a/crates/codegen/kigi-agent/src/builder.rs +++ b/crates/codegen/kigi-agent/src/builder.rs @@ -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, }, diff --git a/crates/codegen/kigi-bin/src/main.rs b/crates/codegen/kigi-bin/src/main.rs index df3ba92..e9ffde7 100644 --- a/crates/codegen/kigi-bin/src/main.rs +++ b/crates/codegen/kigi-bin/src/main.rs @@ -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 { - 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 { - 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 { - 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 { std::sync::Mutex::new(StdioReplayState::default()) } diff --git a/crates/codegen/kigi-config-types/src/lib.rs b/crates/codegen/kigi-config-types/src/lib.rs index 36eba7c..fd5f922 100644 --- a/crates/codegen/kigi-config-types/src/lib.rs +++ b/crates/codegen/kigi-config-types/src/lib.rs @@ -694,11 +694,6 @@ pub struct RemoteSettings { /// See `Config::resolve_image_edit`. #[serde(default)] pub imagine_tools_disabled: Option>, - /// 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, /// 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")); diff --git a/crates/codegen/kigi-mcp/Cargo.toml b/crates/codegen/kigi-mcp/Cargo.toml index d74f61a..f3e5de9 100644 --- a/crates/codegen/kigi-mcp/Cargo.toml +++ b/crates/codegen/kigi-mcp/Cargo.toml @@ -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"] } diff --git a/crates/codegen/kigi-mcp/src/servers.rs b/crates/codegen/kigi-mcp/src/servers.rs index 614050f..8b4c6a7 100644 --- a/crates/codegen/kigi-mcp/src/servers.rs +++ b/crates/codegen/kigi-mcp/src/servers.rs @@ -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![]))); diff --git a/crates/codegen/kigi-memory/src/schema.rs b/crates/codegen/kigi-memory/src/schema.rs index b32eab0..e4baa4b 100644 --- a/crates/codegen/kigi-memory/src/schema.rs +++ b/crates/codegen/kigi-memory/src/schema.rs @@ -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!( diff --git a/crates/codegen/kigi-pager-pty-harness/src/content.rs b/crates/codegen/kigi-pager-pty-harness/src/content.rs index d1daf4f..ffdf840 100644 --- a/crates/codegen/kigi-pager-pty-harness/src/content.rs +++ b/crates/codegen/kigi-pager-pty-harness/src/content.rs @@ -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")); diff --git a/crates/codegen/kigi-shell/Cargo.toml b/crates/codegen/kigi-shell/Cargo.toml index 1069724..d7446b9 100644 --- a/crates/codegen/kigi-shell/Cargo.toml +++ b/crates/codegen/kigi-shell/Cargo.toml @@ -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 } diff --git a/crates/codegen/kigi-shell/src/agent/app.rs b/crates/codegen/kigi-shell/src/agent/app.rs index e9e9197..bc9a49e 100644 --- a/crates/codegen/kigi-shell/src/agent/app.rs +++ b/crates/codegen/kigi-shell/src/agent/app.rs @@ -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; diff --git a/crates/codegen/kigi-shell/src/agent/config.rs b/crates/codegen/kigi-shell/src/agent/config.rs index 3cfa1ca..f3ff708 100644 --- a/crates/codegen/kigi-shell/src/agent/config.rs +++ b/crates/codegen/kigi-shell/src/agent/config.rs @@ -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, - /// 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, - /// 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, - /// 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, - /// 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, - /// 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, - /// 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, - /// 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, - /// 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, - /// Env: `OTEL_EXPORTER_OTLP_TIMEOUT`. Export HTTP timeout (ms). - #[serde(skip_serializing_if = "Option::is_none")] - pub otel_exporter_otlp_timeout: Option, /// Read by `load_management_api_key_sync()`. Declared for `serde_ignored`. #[serde(default, skip_serializing_if = "Option::is_none")] pub management_api_key: Option, @@ -228,18 +182,6 @@ fn blank_as_unset(opt: &Option) -> Option { .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 { - 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 { - 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 { - 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, } -/// `[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, -} -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 { 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, - effective_config: Option<&toml::Value>, -) -> bool { - let table_enabled = |v: Option<&toml::Value>| -> Option { - 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()); diff --git a/crates/codegen/kigi-shell/src/agent/mod.rs b/crates/codegen/kigi-shell/src/agent/mod.rs index 3cd6a96..5d068b6 100644 --- a/crates/codegen/kigi-shell/src/agent/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/mod.rs @@ -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; diff --git a/crates/codegen/kigi-shell/src/agent/models.rs b/crates/codegen/kigi-shell/src/agent/models.rs index 343a7e5..a31d960 100644 --- a/crates/codegen/kigi-shell/src/agent/models.rs +++ b/crates/codegen/kigi-shell/src/agent/models.rs @@ -1323,7 +1323,7 @@ const CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(300); struct ModelsCache { fetched_at: DateTime, #[serde(default, skip_serializing_if = "Option::is_none")] - grok_version: Option, + kigi_version: Option, #[serde(default, skip_serializing_if = "Option::is_none")] auth_method: Option, /// Models-list URL this catalog was fetched from @@ -1418,7 +1418,7 @@ impl ModelsCacheManager { ) -> Option { 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, diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs index 365474d..b6e97ca 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs @@ -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() diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs index c71098f..17d5352 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs @@ -664,8 +664,6 @@ pub struct MvpAgent { plugin_registry_initialized: std::cell::Cell, persona_io_summaries: Vec, /// 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>, /// Sessions opened with `require_gateway` / chat light-frontend (K13). /// Prompt-time guard consults this when the bridge map entry is missing, diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs index c71f413..7994be5 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs @@ -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!["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::*; diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests/subagent_spawn_context_tests.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests/subagent_spawn_context_tests.rs index f4b437e..0eb3b9a 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests/subagent_spawn_context_tests.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests/subagent_spawn_context_tests.rs @@ -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); diff --git a/crates/codegen/kigi-shell/src/agent/proxy.rs b/crates/codegen/kigi-shell/src/agent/proxy.rs deleted file mode 100644 index eec7a56..0000000 --- a/crates/codegen/kigi-shell/src/agent/proxy.rs +++ /dev/null @@ -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 { - 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(target_host: &str, env: F) -> Option -where - F: for<'a> Fn(&'a str) -> Result, -{ - // 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` 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> { - 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 { - // 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, String>> = OnceLock::new(); - -/// Build (or return the cached) TLS client configuration. -fn get_tls_config() -> anyhow::Result> { - 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> { - 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}" - ); - } -} diff --git a/crates/codegen/kigi-shell/src/config/mod.rs b/crates/codegen/kigi-shell/src/config/mod.rs index 38b0014..44db347 100644 --- a/crates/codegen/kigi-shell/src/config/mod.rs +++ b/crates/codegen/kigi-shell/src/config/mod.rs @@ -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) diff --git a/crates/codegen/kigi-shell/src/config/tests.rs b/crates/codegen/kigi-shell/src/config/tests.rs index 5ba499a..cdf8fb3 100644 --- a/crates/codegen/kigi-shell/src/config/tests.rs +++ b/crates/codegen/kigi-shell/src/config/tests.rs @@ -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() diff --git a/crates/codegen/kigi-shell/src/inspect/mod.rs b/crates/codegen/kigi-shell/src/inspect/mod.rs index 42a60a4..19cfca2 100644 --- a/crates/codegen/kigi-shell/src/inspect/mod.rs +++ b/crates/codegen/kigi-shell/src/inspect/mod.rs @@ -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, @@ -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); diff --git a/crates/codegen/kigi-shell/src/leader/mod.rs b/crates/codegen/kigi-shell/src/leader/mod.rs index 35a887b..e4637a5 100644 --- a/crates/codegen/kigi-shell/src/leader/mod.rs +++ b/crates/codegen/kigi-shell/src/leader/mod.rs @@ -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, diff --git a/crates/codegen/kigi-shell/src/leader/protocol.rs b/crates/codegen/kigi-shell/src/leader/protocol.rs index 20f33e5..3f5682b 100644 --- a/crates/codegen/kigi-shell/src/leader/protocol.rs +++ b/crates/codegen/kigi-shell/src/leader/protocol.rs @@ -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, - #[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, }, StopCpuProfile, - WorkspaceStart { - #[serde(default)] - hub_url: Option, - 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, - #[serde(default)] - cwd: Option, - uptime_ms: u64, - active_tool_calls: u32, - #[serde(default)] - sessions: Vec, - 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(); diff --git a/crates/codegen/kigi-shell/src/leader/server.rs b/crates/codegen/kigi-shell/src/leader/server.rs index 33d4dd4..0aa1735 100644 --- a/crates/codegen/kigi-shell/src/leader/server.rs +++ b/crates/codegen/kigi-shell/src/leader/server.rs @@ -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>, - pub workspace: Arc, } 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) -> 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, - /// 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>>, - /// 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, -} -impl WorkspaceControl { - fn new(default_hub_url: Option) -> 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) { - 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, -} -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 { - 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) -> 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, 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, - cwd: String, - cancel: CancellationToken, -) -> Result { - 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 { - 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 { - 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 { - 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 { - 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) { diff --git a/crates/codegen/kigi-shell/src/leader/test_support.rs b/crates/codegen/kigi-shell/src/leader/test_support.rs index 6c41713..ddfc2f2 100644 --- a/crates/codegen/kigi-shell/src/leader/test_support.rs +++ b/crates/codegen/kigi-shell/src/leader/test_support.rs @@ -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, } } diff --git a/crates/codegen/kigi-shell/src/session/acp_session.rs b/crates/codegen/kigi-shell/src/session/acp_session.rs index e883834..b878ded 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session.rs @@ -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; diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs index 67505e8..a005e8a 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs @@ -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 - }); - 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), diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/tool_calls.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/tool_calls.rs index 2e07757..da80b97 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/tool_calls.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/tool_calls.rs @@ -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 diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/turn.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/turn.rs index 03f604e..3a8664e 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/turn.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/turn.rs @@ -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 }) => { diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs index 70c0af1..72afcf8 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs @@ -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), diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/fs_injection_regression_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/fs_injection_regression_tests.rs index 670ae04..0cfabb9 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/fs_injection_regression_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/fs_injection_regression_tests.rs @@ -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, }; diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs index ca2144e..ec126e6 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs @@ -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), diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs index 5a42214..f466d3f 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs @@ -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), diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs index 7689481..057ae00 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs @@ -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), diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/permission_auto_mode_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/permission_auto_mode_tests.rs index 1b1a4d0..b19c31c 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/permission_auto_mode_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/permission_auto_mode_tests.rs @@ -33,6 +33,7 @@ fn install_real_permissions(actor: &mut SessionActor) { vec![], false, None, + true, ); actor.permissions = handle; } diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs index ed64fde..2c310a2 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs @@ -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), diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs index 6b63542..f841c97 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs @@ -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), diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/web_search_e2e_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/web_search_e2e_tests.rs index 4fce030..67e55cd 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/web_search_e2e_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/web_search_e2e_tests.rs @@ -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, }; diff --git a/crates/codegen/kigi-shell/src/session/compaction.rs b/crates/codegen/kigi-shell/src/session/compaction.rs index 31f7795..2ecec33 100644 --- a/crates/codegen/kigi-shell/src/session/compaction.rs +++ b/crates/codegen/kigi-shell/src/session/compaction.rs @@ -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), diff --git a/crates/codegen/kigi-shell/src/session/storage/jsonl/tests.rs b/crates/codegen/kigi-shell/src/session/storage/jsonl/tests.rs index a69723f..a52bf13 100644 --- a/crates/codegen/kigi-shell/src/session/storage/jsonl/tests.rs +++ b/crates/codegen/kigi-shell/src/session/storage/jsonl/tests.rs @@ -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())), ), ), ); diff --git a/crates/codegen/kigi-shell/tests/git_contention_e2e.rs b/crates/codegen/kigi-shell/tests/git_contention_e2e.rs index 010a30e..e115f79 100644 --- a/crates/codegen/kigi-shell/tests/git_contention_e2e.rs +++ b/crates/codegen/kigi-shell/tests/git_contention_e2e.rs @@ -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"); diff --git a/crates/codegen/kigi-shell/tests/session_load_perf.rs b/crates/codegen/kigi-shell/tests/session_load_perf.rs index afba299..7a7ca44 100644 --- a/crates/codegen/kigi-shell/tests/session_load_perf.rs +++ b/crates/codegen/kigi-shell/tests/session_load_perf.rs @@ -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"); diff --git a/crates/codegen/kigi-shell/tests/test_built_binary_e2e.rs b/crates/codegen/kigi-shell/tests/test_built_binary_e2e.rs index 91ed5a6..e662983 100644 --- a/crates/codegen/kigi-shell/tests/test_built_binary_e2e.rs +++ b/crates/codegen/kigi-shell/tests/test_built_binary_e2e.rs @@ -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" diff --git a/crates/codegen/kigi-shell/tests/test_leader_soak.rs b/crates/codegen/kigi-shell/tests/test_leader_soak.rs index c733071..4ab0c76 100644 --- a/crates/codegen/kigi-shell/tests/test_leader_soak.rs +++ b/crates/codegen/kigi-shell/tests/test_leader_soak.rs @@ -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"); diff --git a/crates/codegen/kigi-shell/tests/test_mcp_permission_persistence.rs b/crates/codegen/kigi-shell/tests/test_mcp_permission_persistence.rs index eb7dd10..2d443be 100644 --- a/crates/codegen/kigi-shell/tests/test_mcp_permission_persistence.rs +++ b/crates/codegen/kigi-shell/tests/test_mcp_permission_persistence.rs @@ -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( 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 diff --git a/crates/codegen/kigi-test-support/Cargo.toml b/crates/codegen/kigi-test-support/Cargo.toml index d9fcf72..cda5abb 100644 --- a/crates/codegen/kigi-test-support/Cargo.toml +++ b/crates/codegen/kigi-test-support/Cargo.toml @@ -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"] } diff --git a/crates/codegen/kigi-test-support/src/env.rs b/crates/codegen/kigi-test-support/src/env.rs index dba42e9..228182c 100644 --- a/crates/codegen/kigi-test-support/src/env.rs +++ b/crates/codegen/kigi-test-support/src/env.rs @@ -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") diff --git a/crates/codegen/kigi-test-support/src/leader.rs b/crates/codegen/kigi-test-support/src/leader.rs index ba0193e..c0e8c93 100644 --- a/crates/codegen/kigi-test-support/src/leader.rs +++ b/crates/codegen/kigi-test-support/src/leader.rs @@ -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") diff --git a/crates/codegen/kigi-tools-api/build.rs b/crates/codegen/kigi-tools-api/build.rs index d76784a..17f619c 100644 --- a/crates/codegen/kigi-tools-api/build.rs +++ b/crates/codegen/kigi-tools-api/build.rs @@ -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( diff --git a/crates/codegen/kigi-tools/Cargo.toml b/crates/codegen/kigi-tools/Cargo.toml index 421acc0..5f42a01 100644 --- a/crates/codegen/kigi-tools/Cargo.toml +++ b/crates/codegen/kigi-tools/Cargo.toml @@ -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"] } diff --git a/crates/codegen/kigi-tools/src/implementations/grok_build_hashline/edit/apply.rs b/crates/codegen/kigi-tools/src/implementations/grok_build_hashline/edit/apply.rs index 5db3e94..34fe9c6 100644 --- a/crates/codegen/kigi-tools/src/implementations/grok_build_hashline/edit/apply.rs +++ b/crates/codegen/kigi-tools/src/implementations/grok_build_hashline/edit/apply.rs @@ -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")); } diff --git a/crates/codegen/kigi-tools/src/implementations/search_tool/mod.rs b/crates/codegen/kigi-tools/src/implementations/search_tool/mod.rs index 958358e..b56d991 100644 --- a/crates/codegen/kigi-tools/src/implementations/search_tool/mod.rs +++ b/crates/codegen/kigi-tools/src/implementations/search_tool/mod.rs @@ -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)); } diff --git a/crates/codegen/kigi-tools/src/registry/types.rs b/crates/codegen/kigi-tools/src/registry/types.rs index 8029dc3..d9c6971 100644 --- a/crates/codegen/kigi-tools/src/registry/types.rs +++ b/crates/codegen/kigi-tools/src/registry/types.rs @@ -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, - /// 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, /// 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, + 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::()` time when T is known. - register_in_local: Box, + register_in_local: Box, } /// Per-reminder metadata stored in the builder. struct ReminderEntry { @@ -450,7 +444,7 @@ pub struct FinalizedToolset { scheduler_cancel: Option, /// 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, reminders: Vec, - shared_local_registry: Option, + shared_local_registry: Option, } impl Default for ToolRegistryBuilder { fn default() -> Self { @@ -613,7 +607,7 @@ impl ToolRegistryBuilder { let typed = serde_json::from_value::(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 { 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, } diff --git a/crates/codegen/kigi-tui/src/acp/mod.rs b/crates/codegen/kigi-tui/src/acp/mod.rs index 7fe361c..4d560d6 100644 --- a/crates/codegen/kigi-tui/src/acp/mod.rs +++ b/crates/codegen/kigi-tui/src/acp/mod.rs @@ -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(); } diff --git a/crates/codegen/kigi-tui/src/app/cli.rs b/crates/codegen/kigi-tui/src/app/cli.rs index bd65702..ffb7ec5 100644 --- a/crates/codegen/kigi-tui/src/app/cli.rs +++ b/crates/codegen/kigi-tui/src/app/cli.rs @@ -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, } -/// 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, - /// Workspace root directory to expose. Defaults to the current directory. - #[arg(long, value_name = "DIR", value_hint = ValueHint::DirPath)] - pub cwd: Option, - /// 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, - /// Override the public xAI API base URL. - #[arg(long = "xai-api-base-url")] - pub xai_api_base_url: Option, + /// Override the direct (BYOK / external-API-key) API base URL. + #[arg(long = "api-base-url")] + pub api_base_url: Option, /// Agent runtime mode. Optional: bare `kigi agent` (and the `kigi acp` /// alias) default to stdio. #[command(subcommand)] diff --git a/crates/codegen/kigi-tui/src/app/effects/tests.rs b/crates/codegen/kigi-tui/src/app/effects/tests.rs index 393ef94..36753dd 100644 --- a/crates/codegen/kigi-tui/src/app/effects/tests.rs +++ b/crates/codegen/kigi-tui/src/app/effects/tests.rs @@ -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!( diff --git a/crates/codegen/kigi-tui/src/app/leader_cluster/mod.rs b/crates/codegen/kigi-tui/src/app/leader_cluster/mod.rs index a1a9066..9e36be4 100644 --- a/crates/codegen/kigi-tui/src/app/leader_cluster/mod.rs +++ b/crates/codegen/kigi-tui/src/app/leader_cluster/mod.rs @@ -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"), diff --git a/crates/codegen/kigi-tui/src/app/mod.rs b/crates/codegen/kigi-tui/src/app/mod.rs index c72cf95..f678826 100644 --- a/crates/codegen/kigi-tui/src/app/mod.rs +++ b/crates/codegen/kigi-tui/src/app/mod.rs @@ -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; diff --git a/crates/codegen/kigi-tui/src/headless.rs b/crates/codegen/kigi-tui/src/headless.rs index 75437cc..cb9b3c0 100644 --- a/crates/codegen/kigi-tui/src/headless.rs +++ b/crates/codegen/kigi-tui/src/headless.rs @@ -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 { diff --git a/crates/codegen/kigi-tui/src/trace_cmd.rs b/crates/codegen/kigi-tui/src/trace_cmd.rs index 5e1581c..5e577c6 100644 --- a/crates/codegen/kigi-tui/src/trace_cmd.rs +++ b/crates/codegen/kigi-tui/src/trace_cmd.rs @@ -54,7 +54,7 @@ pub fn build_session_tar(session_dir: &Path, session_id: &str) -> Result 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 #[derive(serde::Serialize)] struct ExportMetadata { session_id: String, - grok_version: String, + kigi_version: String, os: String, arch: String, exported_at: String, diff --git a/crates/codegen/kigi-workspace/Cargo.toml b/crates/codegen/kigi-workspace/Cargo.toml index 8dc7eb9..0ab9b28 100644 --- a/crates/codegen/kigi-workspace/Cargo.toml +++ b/crates/codegen/kigi-workspace/Cargo.toml @@ -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 diff --git a/crates/codegen/kigi-workspace/src/config.rs b/crates/codegen/kigi-workspace/src/config.rs index bc06651..2e1180b 100644 --- a/crates/codegen/kigi-workspace/src/config.rs +++ b/crates/codegen/kigi-workspace/src/config.rs @@ -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, - /// Capability mode applied to the session's toolset. - pub capability_mode: Option, - /// Fully-specified toolset in the runtime serde shape. Takes precedence - /// over `tools`. - pub tool_config: Option, - /// Per-user feature-flag bag. `None` on legacy payloads → tools - /// fall back to their safe defaults. - pub viewer_ctx: Option, - /// Initial auto-approve (YOLO) state. `None` on legacy payloads → - /// fail-closed (false). - pub yolo_mode: Option, - /// Plane-configured toolset in the gRPC wire shape. An empty list is - /// treated as unset (proto3 repeated default). - pub tools: Option>, - pub manifest_version: Option, - pub manifest_hash: Option, - /// 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, -} -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 = 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(name: &str, value: &serde_json::Value) -> Option { - 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, - /// Auth provider for xAI service calls made from workspace-scoped code. - /// `None` for workspaces that do not configure service auth. - pub auth_provider: Option, - /// 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, /// 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, - /// 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, - /// 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, - /// 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, -} -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, - env_session_id: Option, - ) -> Option { - 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, - hub_config: HubConfig, - auth_provider: kigi_computer_hub_sdk::SharedAuthProvider, - server_metadata: Option, - 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 = serde_json::from_value(bad); - assert!(result.is_err()); - } -} diff --git a/crates/codegen/kigi-workspace/src/error.rs b/crates/codegen/kigi-workspace/src/error.rs index 9030761..066dfe8 100644 --- a/crates/codegen/kigi-workspace/src/error.rs +++ b/crates/codegen/kigi-workspace/src/error.rs @@ -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. diff --git a/crates/codegen/kigi-workspace/src/file_system/ext_fs.rs b/crates/codegen/kigi-workspace/src/file_system/ext_fs.rs index 711044c..9528afb 100644 --- a/crates/codegen/kigi-workspace/src/file_system/ext_fs.rs +++ b/crates/codegen/kigi-workspace/src/file_system/ext_fs.rs @@ -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(()) } } diff --git a/crates/codegen/kigi-workspace/src/file_system/mod.rs b/crates/codegen/kigi-workspace/src/file_system/mod.rs index 366a4cb..1a4ad8a 100644 --- a/crates/codegen/kigi-workspace/src/file_system/mod.rs +++ b/crates/codegen/kigi-workspace/src/file_system/mod.rs @@ -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, diff --git a/crates/codegen/kigi-workspace/src/handle.rs b/crates/codegen/kigi-workspace/src/handle.rs index 7293053..dfdfbd5 100644 --- a/crates/codegen/kigi-workspace/src/handle.rs +++ b/crates/codegen/kigi-workspace/src/handle.rs @@ -1,6 +1,4 @@ //! [`WorkspaceHandle`] -- public handle to a workspace instance. -use fastrace::future::FutureExt as _; -use fastrace::local::LocalSpan; use kigi_hunk_tracker::{HunkTrackerActor, HunkTrackerHandle, TrackingMode}; use kigi_tool_protocol::ToolServerStatusPayload; use kigi_tool_protocol::turn_hook::TurnHookOutcome; @@ -10,81 +8,6 @@ use prometheus::{ }; use std::path::PathBuf; use std::sync::Arc; -/// Default SIGTERM drain budget (ms); override via -/// `KIGI_WORKSPACE_TERMINATION_GRACE_MS`. 45s fits under the K8s grace period. -const DEFAULT_TERMINATION_GRACE_MS: u64 = 45_000; -/// preStop-hook drain marker; override via `KIGI_WORKSPACE_DRAINING_FILE`. -const DEFAULT_DRAINING_FILE: &str = "/tmp/workspace-server.draining"; -static DRAIN_STARTED_TOTAL: std::sync::LazyLock = std::sync::LazyLock::new(|| { - register_int_counter_vec!( - "grok_workspace_drain_started_total", - "Graceful drains started, by trigger reason", - &["reason"] - ) - .unwrap() -}); -static DRAIN_COMPLETED_TOTAL: std::sync::LazyLock = std::sync::LazyLock::new(|| { - register_int_counter_vec!( - "grok_workspace_drain_completed_total", - "Graceful drains completed, by outcome", - &["outcome"] - ) - .unwrap() -}); -static DRAIN_DURATION: std::sync::LazyLock = std::sync::LazyLock::new(|| { - register_histogram!( - "grok_workspace_drain_duration_seconds", - "Wall-clock duration of a graceful two-phase drain", - vec![0.5, 1.0, 2.0, 5.0, 10.0, 30.0, 60.0, 120.0] - ) - .unwrap() -}); -/// `session.bind` resolutions advertising zero model-facing tools, by reason. -/// At most one reason is counted per zero-tool bind. -static WORKSPACE_BIND_ZERO_TOOLS_TOTAL: std::sync::LazyLock = - std::sync::LazyLock::new(|| { - register_int_counter_vec!( - "grok_workspace_bind_zero_tools_total", - "session.bind resolutions advertising zero model-facing tools, by reason", - &["reason"] - ) - .unwrap() - }); -/// `session.bind` resolutions that FAILED the bind (the server reports -/// bind-unavailable and the harness re-provisions), by reason. Distinct from -/// [`WORKSPACE_BIND_ZERO_TOOLS_TOTAL`], which counts binds that *completed* -/// while advertising zero model-facing tools. -static WORKSPACE_BIND_FAILED_TOTAL: std::sync::LazyLock = - std::sync::LazyLock::new(|| { - register_int_counter_vec!( - "grok_workspace_bind_failed_total", - "session.bind resolutions that failed the bind, by reason", - &["reason"] - ) - .unwrap() - }); -/// Pinned tool ids this binary could not serve at `session.bind`. -static WORKSPACE_BIND_UNSERVED_TOOLS_TOTAL: std::sync::LazyLock = - std::sync::LazyLock::new(|| { - register_int_counter!( - "grok_workspace_bind_unserved_tools_total", - "Pinned tool ids unknown to this binary at session.bind (reported, not served)" - ) - .unwrap() - }); -/// Model-facing tools advertised per successful `session.bind` (the RPC infra -/// handler is not counted). Catches silent shrinkage of a session's toolset. -static WORKSPACE_BIND_ADVERTISED_TOOLS: std::sync::LazyLock = - std::sync::LazyLock::new(|| { - register_histogram!( - "grok_workspace_bind_advertised_tools", - "Model-facing tools advertised per successful session.bind", - vec![ - 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 8.0, 10.0, 12.0, 16.0, 20.0, 30.0 - ] - ) - .unwrap() - }); /// Tripwire, expected 0 in production. `path="swap"`: a toolset swap found /// the outgoing toolset's `Terminal` resource pointing at a backend other /// than the session-owned one — a resolve path bypassed the session-owned @@ -106,7 +29,6 @@ use crate::capability::CapabilityMode; use crate::config::{ AgentSessionConfig, DEFAULT_EVENT_BUFFER_CAPACITY, HookSourceConfig, WorkspaceConfig, }; -use crate::diag_server::DiagHandle; use crate::error::{WorkspaceError, WorkspaceResult}; use crate::session::swap_policy::{ DeferReason, SessionSnapshot, SwapAction, SwapDecision, SwapPolicy, SwapTrigger, @@ -114,9 +36,6 @@ use crate::session::swap_policy::{ }; use crate::session::tool_config::resolve_session_toolset; use crate::session::{WorkspaceSession, WorkspaceShared}; -use crate::workspace_ops::{ - GetFileEntry, GetFileResult, GetFilesRes, PutFileEntry, PutFileResult, PutFilesRes, -}; use kigi_file_utils::events::types::CancellationCategory; use kigi_file_utils::events::{Event, SessionRelationship, TurnOutcomeLabel}; use kigi_tool_protocol::turn_hook::{AfterTurnAckPayload, AfterTurnAckStatus}; @@ -240,33 +159,6 @@ pub(crate) fn record_non_completed_finalize_canary(outcome: TurnHookOutcome) { } /// Zero-init this module's metric families. See [`crate::init_metrics`]. pub(crate) fn init_metrics() { - for reason in [DrainReason::Sigterm, DrainReason::Evict] { - DRAIN_STARTED_TOTAL - .with_label_values(&[reason.as_str()]) - .inc_by(0); - } - for outcome in [DrainOutcome::Full, DrainOutcome::Partial] { - DRAIN_COMPLETED_TOTAL - .with_label_values(&[outcome.as_str()]) - .inc_by(0); - } - WORKSPACE_BIND_UNSERVED_TOOLS_TOTAL.inc_by(0); - std::sync::LazyLock::force(&DRAIN_DURATION); - std::sync::LazyLock::force(&WORKSPACE_BIND_ADVERTISED_TOOLS); - for reason in [ - "workspace_shutdown", - "session_lookup_failed", - "session_error", - ] { - WORKSPACE_BIND_FAILED_TOTAL - .with_label_values(&[reason]) - .inc_by(0); - } - for reason in ["empty_after_filter", "missing_tool_config"] { - WORKSPACE_BIND_ZERO_TOOLS_TOTAL - .with_label_values(&[reason]) - .inc_by(0); - } WORKSPACE_TERMINAL_BACKEND_ORPHANED_TOTAL .with_label_values(&["swap"]) .inc_by(0); @@ -292,42 +184,6 @@ pub(crate) fn init_metrics() { .inc_by(0); } } -/// Outcome of a hub `session.bind` against an already-existing session -/// (see [`WorkspaceHandle::rebind_existing_hub_session`]). -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum RebindOutcome { - /// Same (or no) explicit toolset — session reused untouched. - Reused, - /// Changed explicit toolset — re-resolved and swapped in. - Reresolved, - /// Changed explicit toolset, but the re-resolve failed; existing kept. - ReresolveFailed, - /// Changed explicit toolset, but the session's toolset is externally - /// owned (local-bind shape) — nothing was resolved or swapped; the - /// existing toolset (and fingerprint) kept. Reused-semantics for the - /// bind reply: advertise the KEPT toolset, drop any unserved set from - /// the unapplied resolve. - KeptExternallyOwned, - /// Changed explicit toolset while the session had tool calls in flight - /// (`explicit → different-explicit` transition only) — existing kept; - /// a later rebind with no calls in flight applies the correction. - ReresolveDeferredInFlight, -} -/// What [`WorkspaceHandle::resolve_and_swap_session_toolset`] actually did — -/// so no caller can mistake a deliberate skip for an installed swap (the -/// skip leaves toolset AND fingerprint untouched). -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -#[must_use = "a skip means the config was NOT applied; callers must not report success"] -pub(crate) enum SwapOutcome { - /// Toolset re-resolved and installed; fingerprint updated. - Swapped, - /// Identical fingerprint ([`SwapDecision::Reuse`]): the live toolset - /// already reflects the config, nothing resolved or changed. - Reused, - /// Externally-owned (local-bind) toolset: rebuild skipped, nothing - /// changed. See `toolset_terminal_is_session_owned`. - SkippedExternallyOwned, -} /// Public handle to a workspace instance. Owns shared state (sessions, /// MCP snapshot, tool config, event bus) and session lifecycle. #[derive(Clone)] @@ -335,62 +191,6 @@ pub struct WorkspaceHandle { pub(crate) shared: Arc, } impl WorkspaceHandle { - /// `None` when not connected. Never hands out an owned - /// `ToolServer` — a clone-drop begins server teardown. - pub async fn trace_donation_reporter( - &self, - service_name: &str, - ) -> Option<( - kigi_computer_hub_sdk::HubDonatingReporter, - kigi_computer_hub_sdk::TraceDonationPump, - )> { - self.shared - .hub_handle - .lock() - .await - .as_ref() - .map(|hub| hub.server.trace_donation_reporter(service_name)) - } - /// Post-connect entry point for the log export layer, the analogue of - /// [`Self::trace_donation_reporter`]. Returns `None` when not connected - /// (the layer stays inert). On - /// `Some`, yields a [`LogDonationSender`] to swap into the - /// already-installed inert `DonatingLogLayer` plus a drain handle. - /// Never hands out an owned `ToolServer` — a clone-drop begins server - /// teardown. - /// - /// [`LogDonationSender`]: kigi_computer_hub_sdk::LogDonationSender - pub async fn log_donation_layer( - &self, - service_name: &str, - ) -> Option<( - kigi_computer_hub_sdk::LogDonationSender, - kigi_computer_hub_sdk::LogDonationPump, - )> { - self.shared - .hub_handle - .lock() - .await - .as_ref() - .map(|hub| hub.server.log_donation_layer(service_name)) - } - /// Post-connect entry point for metric export, the analogue of - /// [`Self::trace_donation_reporter`]. Returns `None` when not connected - /// (no reporter is spawned). On - /// `Some`, spawns the periodic Prometheus-registry gather → OTLP → - /// export pump and yields a drain handle. Never hands out an owned - /// `ToolServer` — a clone-drop begins server teardown. - pub async fn metric_donation_reporter( - &self, - service_name: &str, - ) -> Option { - self.shared - .hub_handle - .lock() - .await - .as_ref() - .map(|hub| hub.server.metric_donation_reporter(service_name)) - } /// Construct a handle with zero sessions. /// /// Sessions are created explicitly via [`Self::create_session`] or @@ -408,24 +208,6 @@ impl WorkspaceHandle { rewind_all_outcomes_from_env(), ) } - /// Construct a handle with an explicit `$KIGI_WORKSPACE_HOME`. - /// - /// [`connect_local_workspace`] calls this; [`Self::new`] takes the - /// ephemeral-home path for tests and local mode. - /// - /// # Panics - /// Requires a Tokio runtime to be entered (for broadcast channel). - pub(crate) fn new_with_home( - config: WorkspaceConfig, - workspace_home: std::path::PathBuf, - ) -> WorkspaceResult { - Self::build( - config, - workspace_home, - events_enabled(), - rewind_all_outcomes_from_env(), - ) - } fn build( config: WorkspaceConfig, workspace_home: std::path::PathBuf, @@ -433,7 +215,7 @@ impl WorkspaceHandle { workspace_rewind_all_outcomes: bool, ) -> WorkspaceResult { let sessions = std::collections::HashMap::new(); - let local_registry = kigi_computer_hub_sdk::LocalRegistry::new(); + let local_registry = kigi_tool_runtime::LocalRegistry::new(); let capacity = if config.event_buffer_capacity == 0 { DEFAULT_EVENT_BUFFER_CAPACITY } else { @@ -513,7 +295,6 @@ impl WorkspaceHandle { activity_tracker.set_event_writers(session_event_writers.clone()); let shared = WorkspaceShared { default_tool_config: config.default_tool_config, - require_explicit_toolset: config.require_explicit_toolset, confine_fs_to_workspace_root: config.confine_fs_to_workspace_root, root_cwd: config.root_cwd.clone(), sessions: parking_lot::RwLock::new(sessions), @@ -526,16 +307,10 @@ impl WorkspaceHandle { hook_load_errors, skills_config: config.skills_config, plugin_discovery_config: config.plugin_discovery_config, - hub_handle: tokio::sync::Mutex::new(None), - hub_tools_snapshot: arc_swap::ArcSwap::new(Arc::new(vec![])), - hub_config: config.hub_config, - auth_provider: config.auth_provider, - activity_notify_handle: arc_swap::ArcSwap::new(Arc::new(None)), client_ext_sink: arc_swap::ArcSwap::new(Arc::new(None)), local_registry, activity_tracker, status_config: config.status_config, - server_metadata: config.server_metadata, fuzzy_searches: Arc::new(tokio::sync::Mutex::new( crate::file_system::FuzzySearchManager::new(std::time::Duration::from_secs(300)), )), @@ -549,7 +324,6 @@ impl WorkspaceHandle { session_event_writers, #[cfg(test)] post_resolve_test_hook: parking_lot::Mutex::new(None), - client_fs_hash_memo: Default::default(), }; Ok(Self { shared: Arc::new(shared), @@ -562,21 +336,6 @@ impl WorkspaceHandle { pub fn activity_tracker(&self) -> &std::sync::Arc { &self.shared.activity_tracker } - /// The [`ToolServer`](kigi_computer_hub_sdk::ToolServer) for this - /// workspace, if a server connection is active. - /// - /// Non-blocking: returns `None` both when no server is connected and when the - /// handle is momentarily locked (e.g. a concurrent connect), so callers - /// must treat `None` as "no server available right now" and degrade gracefully. - pub fn hub_server(&self) -> Option { - self.shared.hub_server() - } - /// Like [`Self::hub_server`] but awaits the connection lock instead of returning - /// `None` on contention, so a transient `connect_hub` lock is not mistaken - /// for "no server". `None` means no server is connected. Use from async callers. - pub async fn hub_server_blocking(&self) -> Option { - self.shared.hub_server_blocking().await - } /// Get the workspace root directory. pub(crate) fn root_cwd(&self) -> crate::error::WorkspaceResult { Ok(self.shared.root_cwd.clone()) @@ -689,29 +448,22 @@ impl WorkspaceHandle { let session_env = Arc::new(std::collections::HashMap::new()); let config = tool_config.unwrap_or_else(|| self.shared.default_tool_config.clone()); let mcp_snapshot = self.shared.mcp_tools_snapshot.load_full(); - let hub_snapshot = self.shared.hub_tools_snapshot.load_full(); let system_notify_channel = system_notifications .then(kigi_tools::notification::types::ToolNotificationHandle::channel); let system_notify_handle = system_notify_channel.as_ref().map(|(h, _)| h.clone()); - let (effective, toolset, terminal_backend) = { - let _span = LocalSpan::enter_with_local_parent("tool_server.toolset_resolve") - .with_property(|| ("session_id", session_id.clone())); - resolve_session_toolset( - config, - capability, - &mcp_snapshot, - &hub_snapshot, - cwd.clone(), - session_env.clone(), - &session_id, - self.shared.session_factory.as_ref(), - Some(self.shared.local_registry.clone()), - self.shared.lsp.clone(), - viewer_ctx.clone(), - self.shared - .compose_session_notification_handle(system_notify_handle), - ) - }?; + let (effective, toolset, terminal_backend) = resolve_session_toolset( + config, + capability, + &mcp_snapshot, + cwd.clone(), + session_env.clone(), + &session_id, + self.shared.session_factory.as_ref(), + Some(self.shared.local_registry.clone()), + self.shared.lsp.clone(), + viewer_ctx.clone(), + system_notify_handle, + )?; let session = Arc::new(WorkspaceSession::new( session_id.clone(), cwd, @@ -736,324 +488,6 @@ impl WorkspaceHandle { ); Ok(session) } - /// Update a session's tool config with auth and serialization; the RPC - /// handler derives `caller_session_id` from the server-bound envelope. - /// Swap gating (retryable `TurnActive`, stale heal): [`SwapPolicy::evaluate`]. - pub(crate) async fn update_tool_config( - &self, - caller_session_id: &str, - session_id: &str, - new_config: kigi_tools::registry::types::ToolServerConfig, - ) -> crate::error::WorkspaceResult<()> { - let session = self - .session(session_id) - .ok_or_else(|| crate::error::WorkspaceError::SessionNotFound(session_id.to_owned()))?; - if caller_session_id != session_id { - return Err(crate::error::WorkspaceError::Unauthorized { - caller: caller_session_id.to_owned(), - target: session_id.to_owned(), - }); - } - match self - .resolve_and_swap_session_toolset(&session, new_config, SwapTrigger::UpdateRpc) - .await? - { - SwapOutcome::Swapped | SwapOutcome::Reused => Ok(()), - SwapOutcome::SkippedExternallyOwned => Err( - crate::error::WorkspaceError::ToolsetExternallyOwned(session_id.to_owned()), - ), - } - } - /// Re-resolve `new_config` against the session's frozen bind-time inputs - /// and atomically swap its toolset (`ToolsChanged`). Update-RPC entry: - /// gated by [`SwapPolicy::evaluate`], twice (entry + post-resolve). - pub(crate) async fn resolve_and_swap_session_toolset( - &self, - session: &Arc, - new_config: kigi_tools::registry::types::ToolServerConfig, - trigger: SwapTrigger, - ) -> crate::error::WorkspaceResult { - let _update_guard = session.update_lock.lock().await; - let session_id = session.session_id(); - let new_fingerprint = serde_json::to_value(&new_config).ok(); - let snapshot = SessionSnapshot::capture( - session, - &self.shared.activity_tracker, - new_fingerprint.as_ref(), - ) - .await; - match SwapPolicy::evaluate(&snapshot, trigger) { - SwapDecision::Reuse => { - tracing::debug!( - session_id = % session_id, trigger = trigger.metric_label(), - "toolset config identical to the stored bind fingerprint — \ - reused untouched" - ); - Ok(SwapOutcome::Reused) - } - SwapDecision::Skip(reason) => { - record_swap_decision( - &self.shared.activity_tracker, - trigger, - session_id, - SwapAction::Skipped(reason), - ); - tracing::warn!( - session_id = % session_id, trigger = trigger.metric_label(), - "toolset swap skipped: toolset terminal backend is externally \ - owned (local bind)" - ); - Ok(SwapOutcome::SkippedExternallyOwned) - } - SwapDecision::Defer(reason) => { - record_swap_decision( - &self.shared.activity_tracker, - trigger, - session_id, - SwapAction::Deferred(reason), - ); - tracing::info!( - session_id = % session_id, trigger = trigger.metric_label(), - "toolset mutation rejected: turn active — retry at the turn boundary" - ); - Err(crate::error::WorkspaceError::TurnActive( - session_id.to_owned(), - )) - } - SwapDecision::Apply => { - self.resolve_and_swap_session_toolset_locked( - session, - new_config, - new_fingerprint, - trigger, - ) - .await - } - } - } - /// The [`SwapDecision::Apply`] arm: resolve `new_config` (whose - /// fingerprint `new_fingerprint` must be) and install it. Callers hold - /// `update_lock` and evaluated [`SwapPolicy`] to `Apply` under that hold. - async fn resolve_and_swap_session_toolset_locked( - &self, - session: &Arc, - new_config: kigi_tools::registry::types::ToolServerConfig, - new_fingerprint: Option, - trigger: SwapTrigger, - ) -> crate::error::WorkspaceResult { - let session_id = session.session_id().to_owned(); - let mcp_snapshot = self.shared.mcp_tools_snapshot.load_full(); - let hub_snapshot = self.shared.hub_tools_snapshot.load_full(); - let cwd = session.cwd().to_path_buf(); - let session_env = session.session_env().clone(); - let cap = session.capability_mode(); - let factory = self.shared.session_factory.clone(); - let lr = self.shared.local_registry.clone(); - let lsp = self.shared.lsp.clone(); - let sid = session_id.to_owned(); - let viewer_ctx = session.viewer_ctx().cloned(); - let notification_handle = self - .shared - .compose_session_notification_handle(session.system_notify_handle()); - let terminal_backend = session.terminal_backend().clone(); - let resolve_result = tokio::task::spawn_blocking(move || { - crate::session::tool_config::resolve_session_toolset_rebuild( - new_config, - cap, - &mcp_snapshot, - &hub_snapshot, - cwd, - session_env, - &sid, - factory.as_ref(), - Some(lr), - lsp, - viewer_ctx, - notification_handle, - terminal_backend, - ) - }) - .await - .map_err(|e| crate::error::WorkspaceError::JoinError(e.to_string()))?; - let (effective, new_toolset) = resolve_result?; - #[cfg(test)] - if let Some(hook) = self.shared.post_resolve_test_hook.lock().as_ref() { - hook(); - } - if trigger.rechecks_after_resolve() { - let snapshot = SessionSnapshot::capture( - session, - &self.shared.activity_tracker, - new_fingerprint.as_ref(), - ) - .await; - match SwapPolicy::evaluate(&snapshot, trigger) { - SwapDecision::Apply => {} - SwapDecision::Reuse => { - tracing::debug!( - session_id = % session_id, trigger = trigger.metric_label(), - "resolved toolset discarded post-resolve: a concurrent \ - bind installed the identical fingerprint during the \ - re-resolve" - ); - return Ok(SwapOutcome::Reused); - } - SwapDecision::Skip(reason) => { - record_swap_decision( - &self.shared.activity_tracker, - trigger, - &session_id, - SwapAction::Skipped(reason), - ); - tracing::warn!( - session_id = % session_id, trigger = trigger.metric_label(), - "toolset swap skipped: toolset terminal backend is externally \ - owned (local bind)" - ); - return Ok(SwapOutcome::SkippedExternallyOwned); - } - SwapDecision::Defer(reason) => { - let reason = match reason { - DeferReason::TurnActive => DeferReason::TurnActiveLate, - other => other, - }; - record_swap_decision( - &self.shared.activity_tracker, - trigger, - &session_id, - SwapAction::Deferred(reason), - ); - tracing::info!( - session_id = % session_id, trigger = trigger.metric_label(), - "toolset mutation rejected post-resolve: a turn started during \ - the re-resolve — resolved toolset discarded; retry at the \ - turn boundary" - ); - return Err(crate::error::WorkspaceError::TurnActive(session_id)); - } - } - } - session - .replace_carrying_browser_service(Arc::new(effective), new_toolset) - .await; - session.set_bind_tool_config_fingerprint(new_fingerprint); - session.clear_stale_resolve(); - record_swap_decision( - &self.shared.activity_tracker, - trigger, - &session_id, - SwapAction::Applied, - ); - let _ = self - .shared - .events - .send(kigi_workspace_types::WorkspaceEvent::ToolsChanged { - session_id: session_id.to_owned(), - }); - Ok(SwapOutcome::Swapped) - } - /// Hub `session.bind` against an existing session: reuse, or re-resolve - /// and swap per the owner-rebind policy rows (incl. the identical stale - /// heal). `explicit_cfg=None` never overwrites; `None` = session vanished. - pub(crate) async fn rebind_existing_hub_session( - &self, - session_id: &str, - explicit_cfg: Option, - bind_fingerprint: Option, - ) -> Option<(Arc, RebindOutcome)> { - let session = self.session(session_id)?; - let Some(cfg) = explicit_cfg else { - return Some((session, RebindOutcome::Reused)); - }; - let outcome = { - let _update_guard = session.update_lock.lock().await; - let snapshot = SessionSnapshot::capture( - &session, - &self.shared.activity_tracker, - bind_fingerprint.as_ref(), - ) - .await; - match SwapPolicy::evaluate(&snapshot, SwapTrigger::OwnerRebind) { - SwapDecision::Reuse => RebindOutcome::Reused, - SwapDecision::Defer(reason) => { - record_swap_decision( - &self.shared.activity_tracker, - SwapTrigger::OwnerRebind, - session_id, - SwapAction::Deferred(reason), - ); - tracing::warn!( - session_id = % session_id, in_flight = snapshot - .in_flight_calls(), - "session.bind: rebind swap (changed explicit toolset or stale-heal \ - re-apply) deferred: tool calls in flight — keeping the existing \ - toolset" - ); - RebindOutcome::ReresolveDeferredInFlight - } - SwapDecision::Skip(reason) => { - record_swap_decision( - &self.shared.activity_tracker, - SwapTrigger::OwnerRebind, - session_id, - SwapAction::Skipped(reason), - ); - tracing::warn!( - session_id = % session_id, - "session.bind: rebind carried a changed toolset config, but the \ - session's toolset is externally owned (local bind) — keeping the \ - existing toolset; the new config did NOT take effect" - ); - RebindOutcome::KeptExternallyOwned - } - SwapDecision::Apply => { - match self - .resolve_and_swap_session_toolset_locked( - &session, - cfg, - bind_fingerprint, - SwapTrigger::OwnerRebind, - ) - .await - { - Ok(SwapOutcome::Swapped) => { - tracing::info!( - session_id = % session_id, - "session.bind: rebind carried a changed toolset config — re-resolved \ - and swapped" - ); - RebindOutcome::Reresolved - } - Ok(SwapOutcome::Reused) => RebindOutcome::Reused, - Ok(SwapOutcome::SkippedExternallyOwned) => { - tracing::warn!( - session_id = % session_id, - "session.bind: rebind carried a changed toolset config, but the \ - session's toolset is externally owned (local bind) — keeping the \ - existing toolset; the new config did NOT take effect" - ); - RebindOutcome::KeptExternallyOwned - } - Err(e) => { - record_swap_decision( - &self.shared.activity_tracker, - SwapTrigger::OwnerRebind, - session_id, - SwapAction::ApplyFailed, - ); - tracing::warn!( - session_id = % session_id, error = % e, - "session.bind: rebind toolset re-resolve failed — keeping the \ - existing toolset" - ); - RebindOutcome::ReresolveFailed - } - } - } - } - }; - Some((session, outcome)) - } pub async fn on_before_turn( &self, session_id: &str, @@ -1177,72 +611,6 @@ impl WorkspaceHandle { self.on_yolo_toggled(session_id, yolo_mode); } } - /// Preemption-aware graceful drain: waits for in-flight tool calls and - /// background tasks within `grace_budget`. Shared by the SIGTERM and - /// server-evict triggers so they can't diverge. - /// - /// The preStop drain marker is (re)written at the start and end with the - /// live count of outstanding work (active tool calls + background tasks), - /// so a preStop hook never reads `0` while a tool call is still running. - /// - /// Returns that same outstanding count after the deadline, so `0` means a - /// fully clean drain — consistent with the final marker and - /// [`DrainOutcome::Full`]; a wedged tool call keeps it non-zero. - pub async fn two_phase_drain( - &self, - grace_budget: std::time::Duration, - reason: DrainReason, - ) -> usize { - let tracker = self.shared.activity_tracker.clone(); - let start = std::time::Instant::now(); - tracker.set_draining(); - tracker.poke(); - DRAIN_STARTED_TOTAL - .with_label_values(&[reason.as_str()]) - .inc(); - let active_at_start = tracker.total_active() as usize; - let drain_file = draining_file_path(); - write_draining_marker(&drain_file, active_at_start); - tracing::info!( - drain_reason = reason.as_str(), - grace_ms = grace_budget.as_millis() as u64, - active_at_start, - "workspace: graceful drain commencing" - ); - let tools_idle = tokio::time::timeout(grace_budget, tracker.wait_until_tools_idle()) - .await - .is_ok(); - if !tools_idle { - tracing::warn!( - active = tracker.total_active(), - "drain deadline exceeded — tool calls still in flight" - ); - } - let total_unfinished = self.shared.activity_tracker.total_active() as usize; - let outcome = classify_drain_outcome(tools_idle); - DRAIN_COMPLETED_TOTAL - .with_label_values(&[outcome.as_str()]) - .inc(); - DRAIN_DURATION.observe(start.elapsed().as_secs_f64()); - write_draining_marker(&drain_file, total_unfinished); - if total_unfinished > 0 { - tracing::warn!( - reason = reason.as_str(), - outcome = outcome.as_str(), - total_unfinished, - duration_ms = start.elapsed().as_millis() as u64, - "workspace: graceful drain finished with work still outstanding" - ); - } else { - tracing::info!( - reason = reason.as_str(), - outcome = outcome.as_str(), - duration_ms = start.elapsed().as_millis() as u64, - "workspace: graceful drain complete" - ); - } - total_unfinished - } /// Bookkeeping for a cancelled in-flight tool call: marks it as /// completed in the activity tracker. Does **not** abort execution /// of the tool — that requires `CancellationToken` plumbing (future work). @@ -1328,7 +696,7 @@ impl WorkspaceHandle { async fn canonicalize_root_dir(root: &std::path::Path) -> WorkspaceResult { #[allow(clippy::disallowed_methods)] let canonical = tokio::fs::canonicalize(root).await.map_err(|e| { - WorkspaceError::HubError(format!("failed to canonicalize workspace root: {e}")) + WorkspaceError::Internal(format!("failed to canonicalize workspace root: {e}")) })?; Ok(dunce::simplified(&canonical).to_path_buf()) } @@ -1359,7 +727,7 @@ impl WorkspaceHandle { ) -> WorkspaceResult { use std::path::{Component, Path}; if req_path.is_empty() { - return Err(WorkspaceError::HubError("empty path not allowed".into())); + return Err(WorkspaceError::Internal("empty path not allowed".into())); } let path = Path::new(req_path); let joined = if path.is_absolute() { @@ -1383,7 +751,7 @@ impl WorkspaceHandle { } let normalized: PathBuf = components.into_iter().collect(); if !normalized.starts_with(root) && !normalized.starts_with(canonical_root) { - return Err(WorkspaceError::HubError(format!( + return Err(WorkspaceError::Internal(format!( "path escapes workspace root: {req_path}" ))); } @@ -1396,7 +764,7 @@ impl WorkspaceHandle { Ok(canonical) => { let canonical = dunce::simplified(&canonical).to_path_buf(); if !canonical.starts_with(canonical_root) { - return Err(WorkspaceError::HubError(format!( + return Err(WorkspaceError::Internal(format!( "path resolves outside workspace root (symlink escape): {req_path}" ))); } @@ -1410,12 +778,12 @@ impl WorkspaceHandle { && md.file_type().is_symlink() { if symlink_hops >= MAX_SYMLINK_HOPS { - return Err(WorkspaceError::HubError(format!( + return Err(WorkspaceError::Internal(format!( "path resolves outside workspace root (unresolved symlink chain): {req_path}" ))); } let Ok(target) = tokio::fs::read_link(&check_path).await else { - return Err(WorkspaceError::HubError(format!( + return Err(WorkspaceError::Internal(format!( "failed to resolve symlink for containment: {req_path}" ))); }; @@ -1443,7 +811,7 @@ impl WorkspaceHandle { } } Err(e) => { - return Err(WorkspaceError::HubError(format!( + return Err(WorkspaceError::Internal(format!( "failed to verify path containment: {e}" ))); } @@ -1465,7 +833,7 @@ impl WorkspaceHandle { return Ok((path.to_path_buf(), None)); } let path_str = path.to_str().ok_or_else(|| { - WorkspaceError::HubError(format!("non-UTF-8 path: {}", path.display())) + WorkspaceError::Internal(format!("non-UTF-8 path: {}", path.display())) })?; let canonical_root = self.canonical_root().await?; let confined = self.resolve_service_path(path_str, &canonical_root).await?; @@ -1482,253 +850,12 @@ impl WorkspaceHandle { return Ok((path.to_path_buf(), None)); } let path_str = path.to_str().ok_or_else(|| { - WorkspaceError::HubError(format!("non-UTF-8 path: {}", path.display())) + WorkspaceError::Internal(format!("non-UTF-8 path: {}", path.display())) })?; let canonical_root = Self::canonicalize_root_dir(root).await?; let confined = Self::resolve_path_within_root(path_str, root, &canonical_root).await?; Ok((confined, Some(canonical_root))) } - /// Write files to the workspace filesystem (service-level, no hunk tracking). - /// - /// Files are written sequentially. If file N fails, files 1..N-1 are - /// already on disk and will NOT be rolled back. Callers must inspect - /// per-file results in the response to detect partial failures. - pub async fn put_files(&self, files: Vec) -> WorkspaceResult { - let canonical_root = self.canonical_root().await?; - let mut results = Vec::with_capacity(files.len()); - for entry in files { - let result = self.put_single_file(&entry, &canonical_root).await; - results.push(result); - } - Ok(PutFilesRes { results }) - } - async fn put_single_file( - &self, - entry: &PutFileEntry, - canonical_root: &std::path::Path, - ) -> PutFileResult { - let resolved = match self.resolve_service_path(&entry.path, canonical_root).await { - Ok(p) => p, - Err(e) => { - return PutFileResult { - path: entry.path.clone(), - ok: false, - error: Some(e.to_string()), - hash: None, - }; - } - }; - if entry.create_dirs - && let Some(parent) = resolved.parent() - && let Err(e) = tokio::fs::create_dir_all(parent).await - { - return PutFileResult { - path: entry.path.clone(), - ok: false, - error: Some(format!("failed to create directories: {e}")), - hash: None, - }; - } - let write_result = if entry.append { - use tokio::io::AsyncWriteExt; - async { - let mut f = tokio::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&resolved) - .await?; - f.write_all(entry.content.as_bytes()).await?; - f.flush().await - } - .await - } else { - tokio::fs::write(&resolved, entry.content.as_bytes()).await - }; - match write_result { - Ok(()) => { - let hash = sha256_hex(entry.content.as_bytes()); - PutFileResult { - path: entry.path.clone(), - ok: true, - error: None, - hash: Some(hash), - } - } - Err(e) => PutFileResult { - path: entry.path.clone(), - ok: false, - error: Some(e.to_string()), - hash: None, - }, - } - } - /// Read files from the workspace filesystem with optional cache - /// validation and byte-range support. - /// - /// Files are read sequentially. Each result includes: - /// - `exists`: whether the file exists on disk. - /// - `content`: file content (full or requested byte range as UTF-8). - /// - `hash`: SHA-256 hex digest of the **full** file content. - /// - `matched`: true if `if_none_match` matched the current hash. - /// - `size`: total file size in bytes. - pub async fn get_files(&self, files: Vec) -> WorkspaceResult { - let canonical_root = self.canonical_root().await?; - let mut results = Vec::with_capacity(files.len()); - for entry in files { - let result = self.get_single_file(&entry, &canonical_root).await; - results.push(result); - } - Ok(GetFilesRes { results }) - } - async fn get_single_file( - &self, - entry: &GetFileEntry, - canonical_root: &std::path::Path, - ) -> GetFileResult { - let resolved = match self.resolve_service_path(&entry.path, canonical_root).await { - Ok(p) => p, - Err(e) => { - return GetFileResult { - path: entry.path.clone(), - exists: false, - content: None, - hash: None, - matched: false, - size: None, - error: Some(e.to_string()), - }; - } - }; - let is_chunked = entry.offset.is_some() || entry.length.is_some(); - let metadata = match tokio::fs::metadata(&resolved).await { - Ok(m) => m, - Err(e) - if e.kind() == std::io::ErrorKind::NotFound - || e.kind() == std::io::ErrorKind::NotADirectory => - { - return GetFileResult { - path: entry.path.clone(), - exists: false, - content: None, - hash: None, - matched: false, - size: None, - error: None, - }; - } - Err(e) => { - return GetFileResult { - path: entry.path.clone(), - exists: true, - content: None, - hash: None, - matched: false, - size: None, - error: Some(e.to_string()), - }; - } - }; - let file_size = metadata.len(); - if is_chunked { - let req_offset = entry.offset.unwrap_or(0); - let req_length = entry.length.unwrap_or(file_size.saturating_sub(req_offset)); - let read_result = stream_hash_and_range(&resolved, req_offset, req_length).await; - match read_result { - Ok((hash, chunk_bytes, _streamed)) => { - if let Some(ref etag) = entry.if_none_match - && *etag == hash - { - return GetFileResult { - path: entry.path.clone(), - exists: true, - content: None, - hash: Some(hash), - matched: true, - size: Some(file_size), - error: None, - }; - } - match String::from_utf8(chunk_bytes) { - Ok(content) => GetFileResult { - path: entry.path.clone(), - exists: true, - content: Some(content), - hash: Some(hash), - matched: false, - size: Some(file_size), - error: None, - }, - Err(e) => GetFileResult { - path: entry.path.clone(), - exists: true, - content: None, - hash: Some(hash), - matched: false, - size: Some(file_size), - error: Some(format!("not valid UTF-8 in range: {e}")), - }, - } - } - Err(e) => GetFileResult { - path: entry.path.clone(), - exists: true, - content: None, - hash: None, - matched: false, - size: Some(file_size), - error: Some(e.to_string()), - }, - } - } else { - match tokio::fs::read(&resolved).await { - Ok(bytes) => { - let hash = sha256_hex(&bytes); - if let Some(ref etag) = entry.if_none_match - && *etag == hash - { - return GetFileResult { - path: entry.path.clone(), - exists: true, - content: None, - hash: Some(hash), - matched: true, - size: Some(file_size), - error: None, - }; - } - match String::from_utf8(bytes) { - Ok(content) => GetFileResult { - path: entry.path.clone(), - exists: true, - content: Some(content), - hash: Some(hash), - matched: false, - size: Some(file_size), - error: None, - }, - Err(e) => GetFileResult { - path: entry.path.clone(), - exists: true, - content: None, - hash: Some(hash), - matched: false, - size: Some(file_size), - error: Some(format!("not valid UTF-8: {e}")), - }, - } - } - Err(e) => GetFileResult { - path: entry.path.clone(), - exists: true, - content: None, - hash: None, - matched: false, - size: Some(file_size), - error: Some(e.to_string()), - }, - } - } - } /// Open a fuzzy file search index rooted at the workspace cwd. pub async fn fuzzy_open( &self, @@ -1918,7 +1045,7 @@ impl WorkspaceHandle { handle.emit_client_ext("x.ai/search/content/status".to_string(), params); }) .await - .map_err(|e| WorkspaceError::HubError(e.to_string())) + .map_err(|e| WorkspaceError::Internal(e.to_string())) } pub fn get_or_create_codebase_index( &self, @@ -1973,193 +1100,6 @@ impl WorkspaceHandle { tracing::debug!("codebase index event forwarder exited"); }) } - /// Start MCP servers for a session and bridge them to the server. - pub async fn start_session_mcp_servers( - &self, - session_id: &str, - configs: Vec, - ) -> crate::error::WorkspaceResult { - use crate::mcp::{ - McpClientTransportAdapter, McpStartFailure, McpStartResult, QualifiedMcpToolHandler, - make_bridge_config, server_name_from_mcp_error, - }; - use kigi_computer_hub_mcp_adapter::McpBridge; - use kigi_computer_hub_sdk::ToolServerHandler as _; - use kigi_mcp::servers::MCP_TOOL_NAME_DELIMITER; - use kigi_tool_protocol::SessionId; - let tool_server = { - let hub_guard = self.shared.hub_handle.lock().await; - let hub = hub_guard - .as_ref() - .ok_or_else(|| WorkspaceError::HubError("no hub connection".into()))?; - hub.server.clone() - }; - let session = self - .session(session_id) - .ok_or_else(|| WorkspaceError::SessionNotFound(session_id.to_owned()))?; - let sid = SessionId::new(session_id) - .map_err(|e| WorkspaceError::HubError(format!("invalid session_id: {e}")))?; - { - let mut tool_ids = session.mcp_tool_ids.lock().await; - for tid in tool_ids.drain(..) { - let _ = tool_server.unregister_tool_dynamic(&tid, &sid).await; - } - let mut existing_bridges = session.mcp_bridges.lock().await; - existing_bridges.clear(); - let mut state = session.mcp_state.lock().await; - state.owned_clients.clear(); - } - let session_id_owned = session_id.to_owned(); - let event_writer = self.shared.session_event_writer(session_id); - let rt_handle = tokio::runtime::Handle::current(); - let mcp_results: Vec> = - tokio::task::spawn_blocking(move || { - use kigi_mcp::oauth_config::McpOAuthConfigMap; - use kigi_mcp::servers::{McpClientTimeoutOverrides, McpMetaConfigMap}; - use std::collections::HashMap; - let overrides_map: HashMap = HashMap::new(); - let meta_config_map = McpMetaConfigMap::new(); - let oauth_config_map = McpOAuthConfigMap::new(); - rt_handle.block_on(kigi_mcp::servers::start_mcp_servers( - configs, - Some(&session_id_owned), - &overrides_map, - &meta_config_map, - &oauth_config_map, - &event_writer, - kigi_mcp::servers::OauthInteractivity::Interactive, - )) - }) - .await - .map_err(|e| WorkspaceError::JoinError(e.to_string()))?; - let mcp_state = session.mcp_state.clone(); - let mut started = Vec::new(); - let mut failed = Vec::new(); - let mut bridges = Vec::new(); - let mut registered_tool_ids = Vec::new(); - for result in mcp_results { - match result { - Ok(client) => { - let server_name = client.server_name().to_owned(); - let client = Arc::new(client); - { - let mut state = mcp_state.lock().await; - state - .owned_clients - .insert(server_name.clone(), Arc::clone(&client)); - } - let transport: Arc = - Arc::new(McpClientTransportAdapter::new(Arc::clone(&client))); - let bridge_config = make_bridge_config(sid.clone(), &server_name); - match McpBridge::connect(transport, &bridge_config).await { - Ok(handle) => { - for handler in handle.bridge.handlers() { - let qualified_name = format!( - "{}{}{}", - server_name, - MCP_TOOL_NAME_DELIMITER, - handler.tool_id() - ); - let qualified = match QualifiedMcpToolHandler::try_new( - qualified_name.clone(), - handler.clone(), - ) { - Some(h) => Arc::new(h), - None => continue, - }; - if let Err(e) = tool_server - .register_tool_dynamic(qualified, vec![sid.clone()]) - .await - { - tracing::warn!( - server = % server_name, tool = % qualified_name, error = % - e, "failed to register MCP tool on hub" - ); - } else if let Ok(tid) = - kigi_tool_protocol::ToolId::new(&qualified_name) - { - registered_tool_ids.push(tid); - } - } - bridges.push(handle); - started.push(server_name); - } - Err(e) => { - { - let mut state = mcp_state.lock().await; - state.owned_clients.remove(&server_name); - } - tracing::warn!( - server = % server_name, error = % e, - "McpBridge::connect failed" - ); - failed.push(McpStartFailure { - name: server_name, - error: e.to_string(), - }); - } - } - } - Err(e) => { - let name = server_name_from_mcp_error(&e).to_owned(); - tracing::warn!( - server = % name, error = % e, "MCP server start failed" - ); - failed.push(McpStartFailure { - name, - error: e.to_string(), - }); - } - } - } - { - let mut session_bridges = session.mcp_bridges.lock().await; - session_bridges.extend(bridges); - } - { - let mut ids = session.mcp_tool_ids.lock().await; - ids.extend(registered_tool_ids); - } - tracing::info!( - session_id = % session_id, started = ? started, failed_count = failed.len(), - "session MCP servers initialized" - ); - if !started.is_empty() { - let _ = self - .shared - .events - .send(kigi_workspace_types::WorkspaceEvent::ToolsChanged { - session_id: session_id.to_owned(), - }); - } - Ok(McpStartResult { started, failed }) - } - /// Unregister all MCP tools for a session from the server. - pub async fn teardown_session_mcp(&self, session_id: &str) { - let tool_server = { - let hub_guard = self.shared.hub_handle.lock().await; - match hub_guard.as_ref() { - Some(hub) => hub.server.clone(), - None => return, - } - }; - let session = match self.session(session_id) { - Some(s) => s, - None => return, - }; - let sid = match kigi_tool_protocol::SessionId::new(session_id) { - Ok(s) => s, - Err(_) => return, - }; - let mut tool_ids = session.mcp_tool_ids.lock().await; - for tid in tool_ids.drain(..) { - let _ = tool_server.unregister_tool_dynamic(&tid, &sid).await; - } - let mut bridges = session.mcp_bridges.lock().await; - bridges.clear(); - let mut state = session.mcp_state.lock().await; - state.owned_clients.clear(); - } /// Look up an existing session. pub fn session(&self, session_id: &str) -> Option> { self.shared.sessions.read().get(session_id).cloned() @@ -2216,13 +1156,11 @@ impl WorkspaceHandle { env.extend(config.extra_env.clone()); let session_env = Arc::new(env); let mcp_snapshot = self.shared.mcp_tools_snapshot.load_full(); - let hub_snapshot = self.shared.hub_tools_snapshot.load_full(); let inherited_viewer_ctx = parent.viewer_ctx().cloned(); let (effective, toolset, terminal_backend) = resolve_session_toolset( baseline, config.capability_mode, &mcp_snapshot, - &hub_snapshot, cwd.clone(), session_env.clone(), &config.agent_id, @@ -2230,7 +1168,7 @@ impl WorkspaceHandle { Some(self.shared.local_registry.clone()), self.shared.lsp.clone(), inherited_viewer_ctx.clone(), - self.shared.compose_session_notification_handle(None), + None, )?; let (hunk_event_tx, _hunk_event_rx) = tokio::sync::mpsc::unbounded_channel(); let hunk_cancel = tokio_util::sync::CancellationToken::new(); @@ -2300,986 +1238,6 @@ impl WorkspaceHandle { ) }) } - /// Bulk-replace hub tool configs and re-resolve every session. - pub fn on_hub_tools_changed( - &self, - new_hub_tools: Vec, - ) -> usize { - self.shared - .hub_tools_snapshot - .store(Arc::new(new_hub_tools)); - tokio::task::block_in_place(|| { - tokio::runtime::Handle::current().block_on( - self.shared - .re_resolve_all_sessions("hub_tools_changed", true), - ) - }) - } - /// Per-`session.bind` handler resolver: resolves the bind metadata into a - /// session toolset (fail-closed in strict mode) and returns the handlers - /// plus the bind-report fields. Extracted from `connect_hub` so tests can - /// drive the full bind path without a hub connection. - pub(crate) fn session_bind_resolver( - &self, - catalog: Arc>>, - rpc_tool_id: kigi_tool_protocol::ToolId, - ) -> kigi_computer_hub_sdk::SessionHandlerResolver { - let weak_shared = Arc::downgrade(&self.shared); - Arc::new( - move |sid: kigi_tool_protocol::SessionId, params: Option| { - let catalog = catalog.clone(); - let rpc_tool_id = rpc_tool_id.clone(); - let weak_shared = weak_shared.clone(); - let bind_parent = params - .as_ref() - .and_then(|p| p.pointer("/trace_context")) - .and_then(serde_json::Value::as_str) - .and_then(fastrace::collector::SpanContext::decode_w3c_traceparent) - .unwrap_or_else(kigi_tracing::local_or_random_span_ctx); - let bind_span = fastrace::Span::root("tool_server.session_bind", bind_parent) - .with_properties(|| { - [ - ("session_id", sid.to_string()), - ("force_tracing", "true".to_owned()), - ] - }); - Box::pin( - async move { - let Some(shared) = weak_shared.upgrade() else { - WORKSPACE_BIND_FAILED_TOTAL - .with_label_values(&["workspace_shutdown"]) - .inc(); - return Err( - kigi_tool_runtime::ToolError::service_unavailable( - "workspace is shutting down; cannot bind session", - ), - ); - }; - let ws = WorkspaceHandle { shared }; - let sid_str = sid.to_string(); - let params = params.unwrap_or(serde_json::Value::Null); - let bind_cwd = params - .pointer("/cwd") - .and_then(serde_json::Value::as_str) - .map(std::path::PathBuf::from); - let bind_config = params - .pointer("/metadata") - .map(crate::config::WorkspaceBindConfig::from_metadata) - .unwrap_or_default(); - let empty_toolset = || kigi_tools::registry::types::ToolServerConfig { - tools: vec![], - behavior_preset: None, - }; - let mut resolve_zero_reason: Option<&'static str> = None; - let mut resolve_error: Option = None; - let mut unserved_tool_ids: Vec = Vec::new(); - let known_ids = ws.shared.session_factory.known_tool_ids(); - let known_id = |id: &str| known_ids.contains(id); - let require_explicit = ws.shared.require_explicit_toolset; - let tool_config = match bind_config - .resolve(&known_id, require_explicit) - { - crate::config::ResolvedToolset::Toolset(resolved) => { - unserved_tool_ids = resolved.unserved_tool_ids; - Some(resolved.toolset) - } - crate::config::ResolvedToolset::UseDefault => None, - crate::config::ResolvedToolset::MissingToolConfig => { - if bind_config.rpc_only { - tracing::info!( - session_id = % sid_str, - "session.bind: rpc_only bind with no toolset — \ - failing closed with an empty toolset" - ); - } else { - tracing::warn!( - session_id = % sid_str, - "session.bind: no explicit tool configuration passed and this \ - workspace requires one — failing closed with an empty toolset" - ); - } - resolve_zero_reason = Some("missing_tool_config"); - resolve_error = Some( - format!( - "missing_tool_config: no usable explicit tool configuration \ - on session.bind (absent, or dropped as malformed — see \ - server logs) and this workspace requires one (presets are \ - not supported; server version {})", - kigi_version::VERSION - ), - ); - Some(empty_toolset()) - } - crate::config::ResolvedToolset::InvalidToolConfig(err) => { - tracing::warn!( - session_id = % sid_str, error = % err, - "session.bind: invalid tool config entry — failing closed with an empty toolset" - ); - resolve_zero_reason = Some("invalid_tool_config"); - resolve_error = Some( - format!( - "invalid_tool_config: {err} (server version {})", - kigi_version::VERSION - ), - ); - Some(empty_toolset()) - } - }; - let (explicit_cfg, bind_fingerprint) = match ( - &tool_config, - resolve_zero_reason, - ) { - (Some(cfg), None) if !cfg.tools.is_empty() => { - (Some(cfg.clone()), serde_json::to_value(cfg).ok()) - } - _ => (None, None), - }; - let capability = bind_config - .capability_mode - .unwrap_or(crate::capability::CapabilityMode::All); - let yolo_mode = bind_config.yolo_mode.unwrap_or(false); - tracing::info!( - session_id = % sid_str, cwd = ? bind_cwd, preset = ? bind_config - .preset, capability = ? capability, yolo_mode, - "session.bind: resolving workspace session toolset" - ); - let created = { - let _span = LocalSpan::enter_with_local_parent( - "tool_server.session_bind.create_session", - ) - .with_property(|| ("session_id", sid_str.clone())); - ws.create_session_with_config( - sid_str.clone(), - bind_cwd, - tool_config, - capability, - bind_config.viewer_ctx.clone(), - bind_config.system_notifications, - ) - }; - let session = match created { - Ok(session) => { - session.set_yolo_mode(yolo_mode); - session - .set_bind_tool_config_fingerprint_if_unset( - bind_fingerprint.clone(), - ); - tracing::info!( - session_id = % sid_str, - "workspace session created for hub bind" - ); - session - } - Err(crate::error::WorkspaceError::SessionAlreadyExists(_)) => { - match ws - .rebind_existing_hub_session( - &sid_str, - explicit_cfg, - bind_fingerprint, - ) - .await - { - Some((session, RebindOutcome::Reresolved)) => session, - Some((session, _)) => { - unserved_tool_ids.clear(); - if resolve_zero_reason != Some("invalid_tool_config") - && !session.effective_tool_config().tools.is_empty() - { - resolve_error = None; - resolve_zero_reason = None; - } - session - } - None => { - WORKSPACE_BIND_FAILED_TOTAL - .with_label_values(&["session_lookup_failed"]) - .inc(); - return Err( - kigi_tool_runtime::ToolError::service_unavailable( - format!( - "session rebind raced teardown for `{sid_str}`; retry" - ), - ), - ); - } - } - } - Err(e) => { - tracing::error!( - session_id = % sid_str, error = % e, - "failed to create workspace session for hub bind" - ); - WORKSPACE_BIND_FAILED_TOTAL - .with_label_values(&["session_error"]) - .inc(); - return Err( - kigi_tool_runtime::ToolError::service_unavailable( - format!("failed to create workspace session: {e}"), - ), - ); - } - }; - let mut handlers = { - let _span = LocalSpan::enter_with_local_parent( - "tool_server.session_bind.handlers", - ) - .with_property(|| ("session_id", sid_str.clone())); - build_session_routed_handlers(&session.toolset(), &ws) - }; - let advertised: Vec = handlers - .iter() - .map(|h| h.tool_id().as_str().to_owned()) - .collect(); - WORKSPACE_BIND_ADVERTISED_TOOLS.observe(advertised.len() as f64); - if advertised.is_empty() { - let reason = resolve_zero_reason.unwrap_or("empty_after_filter"); - let skip_zero_metric = bind_config.rpc_only - && reason == "missing_tool_config"; - if skip_zero_metric { - tracing::info!( - session_id = % sid_str, reason, - "session.bind: advertising zero model-facing tools (rpc_only)" - ); - } else { - tracing::warn!( - session_id = % sid_str, - "session.bind: advertising zero model-facing tools (RPC handler only)" - ); - WORKSPACE_BIND_ZERO_TOOLS_TOTAL - .with_label_values(&[reason]) - .inc(); - } - } - handlers - .extend( - catalog - .iter() - .filter(|h| h.tool_id() == rpc_tool_id) - .cloned(), - ); - if !unserved_tool_ids.is_empty() { - WORKSPACE_BIND_UNSERVED_TOOLS_TOTAL - .inc_by(unserved_tool_ids.len() as u64); - tracing::warn!( - session_id = % sid_str, unserved = ? unserved_tool_ids, - "session.bind: serving partial pinned toolset" - ); - } - tracing::info!( - session_id = % sid_str, advertised = advertised.len(), tools = ? - advertised, unserved = ? unserved_tool_ids, - "session.bind: advertising finalized session toolset" - ); - Ok(kigi_computer_hub_sdk::ResolvedSessionHandlers { - handlers, - unserved_tool_ids, - resolve_error, - }) - } - .in_span(bind_span), - ) - }, - ) - } - /// Connect to the server, start the tool server (provider - /// direction) and notification listener (consumer direction). - /// - /// No-op if no `hub_config` was provided or already connected. - /// - /// The tool server exposes the workspace's main session tools so - /// the server can dispatch `tool_call_request` frames to them. The - /// notification listener updates `hub_tools_snapshot` and - /// re-resolves every session's toolset whenever the server announces - /// tool changes. - pub async fn connect_hub(&self) -> WorkspaceResult<()> { - use crate::hub::{HubHandle, apply_tools_changed, hub_result}; - tracing::info!("WorkspaceHandle::connect_hub — starting"); - let hub_config = match &self.shared.hub_config { - Some(c) => { - let mut cfg = c.clone(); - cfg.activity_tracker = Some(self.shared.activity_tracker.clone()); - cfg - } - None => { - tracing::info!("WorkspaceHandle::connect_hub — no hub config, skipping"); - return Ok(()); - } - }; - let mut hub_guard = self.shared.hub_handle.lock().await; - if hub_guard.is_some() { - return Ok(()); - } - tracing::info!( - url = % hub_config.url, "WorkspaceHandle::connect_hub — connecting to hub" - ); - let (template_handlers, rpc_tool_id) = { - let session_env = Arc::new(std::collections::HashMap::new()); - let mcp_snapshot = self.shared.mcp_tools_snapshot.load_full(); - let hub_snapshot = self.shared.hub_tools_snapshot.load_full(); - let (_, template_toolset, _template_backend) = resolve_session_toolset( - self.shared.default_tool_config.clone(), - crate::capability::CapabilityMode::All, - &mcp_snapshot, - &hub_snapshot, - self.shared.root_cwd.clone(), - session_env, - "__template__", - self.shared.session_factory.as_ref(), - Some(self.shared.local_registry.clone()), - self.shared.lsp.clone(), - None, - None, - )?; - let mut handlers = build_session_routed_handlers(&template_toolset, self); - let tool_names: Vec = handlers - .iter() - .map(|h| h.tool_id().as_str().to_owned()) - .collect(); - let rpc_handler: Arc = - Arc::new(crate::hub_server::WorkspaceRpcHandler::new(self.clone())); - let rpc_tool_id = rpc_handler.tool_id(); - handlers.push(rpc_handler); - tracing::info!( - tool_count = handlers.len(), tools = ? tool_names, - "Registering server tool catalog on hub" - ); - (handlers, rpc_tool_id) - }; - let catalog: Arc>> = - Arc::new(template_handlers.clone()); - let resolver = self.session_bind_resolver(catalog, rpc_tool_id); - let mut handle = hub_result( - HubHandle::connect( - &hub_config, - self.shared.status_config.ws_ping, - self.shared.status_config.ws_reconnect_backoff.clone(), - template_handlers, - self.shared.server_metadata.clone(), - Some(resolver), - ) - .await, - )?; - tracing::info!("WorkspaceHandle::connect_hub — connected, starting server + listeners"); - let (activity_notify_handle, activity_notify_rx) = - kigi_tools::notification::types::ToolNotificationHandle::channel(); - let activity_feed_task = tokio::spawn(run_activity_feed( - self.shared.activity_tracker.clone(), - activity_notify_rx, - )); - handle.set_activity_feed_task(activity_feed_task); - self.shared - .activity_notify_handle - .store(Arc::new(Some(activity_notify_handle))); - let server = handle.server.clone(); - let server_task = tokio::spawn(async move { - if let Err(e) = server.run().await { - tracing::warn!( - error = % e, "hub tool server run loop exited with error" - ); - } - }); - handle.set_server_task(server_task); - let mut notification_rx = handle.server.subscribe_notifications(); - let shared = self.shared.clone(); - let listener_task = tokio::spawn(async move { - while let Some(notification) = notification_rx.recv().await { - match notification { - kigi_computer_hub_sdk::HubNotification::ToolsChanged { - added, - removed, - updated, - .. - } => { - let current = shared.hub_tools_snapshot.load_full(); - let new_tools = apply_tools_changed(¤t, &added, &removed, &updated); - shared.hub_tools_snapshot.store(Arc::new(new_tools)); - shared - .re_resolve_all_sessions("hub_notification", true) - .await; - } - other => { - tracing::debug!(?other, "hub notification (unhandled type)"); - } - } - } - tracing::debug!("hub notification listener exited"); - }); - handle.set_notification_task(listener_task); - let hub_warn_threshold = self.shared.status_config.hub_warn_threshold; - let hub_backoff_base = self.shared.status_config.hub_backoff_base; - /// Compute exponential backoff: `base` * 2^min(n, 7). - fn hub_backoff(base: std::time::Duration, consecutive_errors: u32) -> std::time::Duration { - base.saturating_mul(2u32.pow(consecutive_errors.min(7))) - } - let events_rx = self.shared.events.subscribe(); - let server_for_events = handle.server.clone(); - let event_publisher_task = tokio::spawn(async move { - let mut rx = events_rx; - let mut consecutive_errors: u32 = 0; - loop { - match rx.recv().await { - Ok(event) => { - let payload = - serde_json::to_value(&event).unwrap_or(serde_json::Value::Null); - let frame = kigi_tool_protocol::ToolNotificationFrame::custom( - kigi_tool_protocol::ToolId::new( - crate::hub_ids::WORKSPACE_EVENTS_TOOL_ID, - ) - .expect("constant tool id"), - "workspace_event", - payload, - ); - if let Err(e) = server_for_events.send_notification(frame).await { - consecutive_errors += 1; - if consecutive_errors <= hub_warn_threshold { - tracing::warn!( - error = % e, "failed to send workspace event to hub" - ); - } else { - tracing::debug!( - error = % e, consecutive = consecutive_errors, - "workspace event send failed (backoff)" - ); - } - tokio::time::sleep(hub_backoff(hub_backoff_base, consecutive_errors)) - .await; - } else { - consecutive_errors = 0; - } - } - Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => { - tracing::warn!(skipped = n, "workspace event publisher lagged"); - } - Err(tokio::sync::broadcast::error::RecvError::Closed) => break, - } - } - tracing::debug!("workspace event publisher exited"); - }); - handle.set_event_publisher_task(event_publisher_task); - let tracker_for_status = self.shared.activity_tracker.clone(); - let server_conn = handle.server.connection().clone(); - let heartbeat = self.shared.status_config.heartbeat; - let keepalive = self.shared.status_config.keepalive; - let status_publisher_task = tokio::spawn(async move { - /// Attempt to send a status frame. - /// - /// Returns `Some(true)` on success, `Some(false)` on transport - /// failure (hub unreachable), and `None` when the send was - /// skipped due to a local error (serialization, id allocation) - /// that does not indicate a dead connection. - async fn send_status( - conn: &kigi_computer_hub_sdk::HubConnection, - payload: ToolServerStatusPayload, - ) -> Option { - let params = match serde_json::to_value(&payload) { - Ok(v) => v, - Err(e) => { - tracing::warn!( - error = % e, "failed to serialize tool server status" - ); - return None; - } - }; - let request_id = match conn.try_alloc_request_id() { - Ok(id) => id, - Err(e) => { - tracing::warn!( - error = % e, "failed to alloc request id for status" - ); - return None; - } - }; - let req = kigi_tool_protocol::JsonRpcRequest { - jsonrpc: kigi_tool_protocol::JsonRpcVersion, - id: kigi_tool_protocol::JsonRpcId::from_request_id(&request_id), - session_id: None, - method: kigi_tool_protocol::Method::ToolServerStatus - .as_wire_str() - .to_owned(), - params, - }; - if let Err(e) = conn.call_request(request_id, &req).await { - tracing::debug!(error = % e, "tool_server.status send failed"); - return Some(false); - } - Some(true) - } - fn dedup_key(p: &ToolServerStatusPayload) -> ToolServerStatusPayload { - let mut k = p.clone(); - k.uptime_ms = 0; - k - } - let mut last_sent: std::collections::HashMap, ToolServerStatusPayload> = - std::collections::HashMap::new(); - let mut consecutive_errors: u32 = 0; - let mut last_successful_send = std::time::Instant::now(); - { - let payload = tracker_for_status.snapshot(); - if send_status(&server_conn, payload.clone()).await == Some(true) { - last_sent.insert(None, payload); - last_successful_send = std::time::Instant::now(); - } - } - const MIN_REPUBLISH_INTERVAL: std::time::Duration = - std::time::Duration::from_millis(250); - let mut last_cycle = tokio::time::Instant::now() - MIN_REPUBLISH_INTERVAL; - loop { - tracker_for_status.wait_for_change(heartbeat).await; - let since_last = last_cycle.elapsed(); - if since_last < MIN_REPUBLISH_INTERVAL { - tokio::time::sleep(MIN_REPUBLISH_INTERVAL - since_last).await; - } - last_cycle = tokio::time::Instant::now(); - let mut any_attempt = false; - let mut any_success = false; - let session_ids = tracker_for_status.known_sessions(); - for sid in &session_ids { - let payload = tracker_for_status.snapshot_session(sid); - let key = Some(sid.clone()); - if last_sent.get(&key).map(dedup_key) == Some(dedup_key(&payload)) { - continue; - } - if let Some(ok) = send_status(&server_conn, payload.clone()).await { - any_attempt = true; - if ok { - any_success = true; - last_sent.insert(key, payload); - last_successful_send = std::time::Instant::now(); - } - } - } - last_sent.retain(|k, _| match k { - None => true, - Some(sid) => session_ids.iter().any(|s| s == sid), - }); - let payload = tracker_for_status.snapshot(); - let needs_send = last_sent.get(&None).map(dedup_key) != Some(dedup_key(&payload)); - let force_keepalive = - !needs_send && !any_success && last_successful_send.elapsed() >= keepalive; - if (needs_send || force_keepalive) - && let Some(ok) = send_status(&server_conn, payload.clone()).await - { - any_attempt = true; - if ok { - any_success = true; - last_sent.insert(None, payload); - last_successful_send = std::time::Instant::now(); - } - } - if any_attempt && !any_success { - consecutive_errors += 1; - if consecutive_errors <= hub_warn_threshold { - tracing::warn!( - "status publisher: hub unreachable ({} consecutive failed cycles)", - consecutive_errors, - ); - } else { - tracing::debug!( - consecutive = consecutive_errors, - "status publish failed (backoff)" - ); - } - tokio::time::sleep(hub_backoff(hub_backoff_base, consecutive_errors)).await; - } else if any_success { - consecutive_errors = 0; - } - } - }); - handle.set_status_publisher_task(status_publisher_task); - { - let (ext_tx, mut ext_rx) = - tokio::sync::mpsc::unbounded_channel::<(String, serde_json::Value)>(); - let server_for_ext = handle.server.clone(); - let ext_task = tokio::spawn(async move { - while let Some((method, params)) = ext_rx.recv().await { - let frame = kigi_tool_protocol::ToolNotificationFrame::custom( - kigi_tool_protocol::ToolId::new( - crate::hub_ids::WORKSPACE_CLIENT_EXT_NOTIFICATIONS_TOOL_ID, - ) - .expect("constant tool id"), - "client_ext_notification", - serde_json::json!({ "method" : method, "params" : params }), - ); - let _ = server_for_ext.send_notification(frame).await; - } - }); - handle.set_client_ext_forwarder_task(ext_task); - self.set_client_ext_sink(Arc::new(move |method, params| { - let _ = ext_tx.send((method, params)); - })); - } - handle.set_codebase_index_forwarder_task(self.spawn_codebase_index_event_forwarder()); - *hub_guard = Some(handle); - Ok(()) - } - /// Shutdown the server connection, if active. - pub async fn shutdown_hub(&self) { - let handle = self.shared.hub_handle.lock().await.take(); - if let Some(h) = handle { - h.shutdown().await; - } - } -} -/// Build one [`SessionRoutedToolHandler`](crate::hub::SessionRoutedToolHandler) -/// per tool in `toolset`, keyed by client (function) name. Shared by the -/// connect-time catalog and the per-`session.bind` resolver so the two -/// construction paths cannot drift. -/// -/// `finalize` already rejects duplicate client names, so the `seen` set is -/// defense-in-depth: it guards a regression from ever emitting two handlers -/// with the same `tool_id` (which would duplicate the bind response and -/// silently first-win at dispatch). -fn build_session_routed_handlers( - toolset: &kigi_tools::registry::types::FinalizedToolset, - ws: &WorkspaceHandle, -) -> Vec> { - let tool_kinds = toolset.tool_kinds(); - let mut seen = std::collections::HashSet::new(); - let mut handlers = Vec::new(); - for def in toolset.tool_definitions() { - if !seen.insert(def.function.name.clone()) { - tracing::warn!( - tool = % def.function.name, - "duplicate client name in finalized toolset; skipping" - ); - continue; - } - let mut desc = kigi_tool_types::ToolDescription::new( - def.function.name.clone(), - def.function.description.clone().unwrap_or_default(), - ); - desc.arguments_schema = Some(def.function.parameters.clone()); - desc.kind = tool_kinds.get(&def.function.name).cloned(); - match crate::hub::SessionRoutedToolHandler::new( - def.function.name.clone(), - desc, - Some(def.function.parameters.clone()), - ws.clone(), - ) { - Ok(handler) => handlers - .push(Arc::new(handler) as Arc), - Err(e) => { - tracing::warn!( - tool = % def.function.name, error = % e, - "client name is not a valid ToolId; skipping hub registration" - ); - } - } - } - handlers -} -/// Apply a tool notification to the ActivityTracker background-task count. -/// `started` must precede `completed`, else the unknown `completed` no-ops and -/// strands the count. -pub(crate) fn apply_background_task_notification( - tracker: &crate::activity::ActivityTracker, - notification: &kigi_tools::notification::types::ToolNotification, -) { - use kigi_tools::notification::types::ToolNotification; - match notification { - ToolNotification::BashExecutionBackgrounded(bg) => { - tracker.background_task_started(&bg.task_id); - } - ToolNotification::TaskCompleted(snap) => { - tracker.background_task_completed(&snap.task_id); - } - _ => {} - } -} -/// Tracker-only drain of the session tool-notification stream — not a network -/// send, so the hibernation decrement isn't delayed by send backoff and -/// notifications aren't misattributed across sessions. -pub(crate) async fn run_activity_feed( - tracker: Arc, - mut rx: tokio::sync::mpsc::UnboundedReceiver, -) { - while let Some(notification) = rx.recv().await { - apply_background_task_notification(&tracker, ¬ification); - } -} -/// Compute SHA-256 hex digest. -fn sha256_hex(data: &[u8]) -> String { - use sha2::Digest; - format!("{:x}", sha2::Sha256::digest(data)) -} -/// What triggered a [`WorkspaceHandle::two_phase_drain`] — the metric label. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum DrainReason { - /// Process received SIGTERM / Ctrl-C (standalone `workspace_server`). - Sigterm, - /// Hub sent `tool_server.evict`. - Evict, -} -impl DrainReason { - /// Stable `reason` label for `grok_workspace_drain_started_total`. - pub fn as_str(self) -> &'static str { - match self { - DrainReason::Sigterm => "sigterm", - DrainReason::Evict => "evict", - } - } -} -/// Terminal classification of a graceful drain — the metric label. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum DrainOutcome { - /// All in-flight tool calls and background tasks finished within budget. - Full, - /// Tool calls still in flight at the deadline. - Partial, -} -impl DrainOutcome { - /// Stable `outcome` label for `grok_workspace_drain_completed_total`. - pub fn as_str(self) -> &'static str { - match self { - DrainOutcome::Full => "full", - DrainOutcome::Partial => "partial", - } - } -} -/// Classify a drain: `Partial` when the deadline expired with tool calls still -/// in flight, `Full` otherwise. -fn classify_drain_outcome(tools_idle: bool) -> DrainOutcome { - if !tools_idle { - DrainOutcome::Partial - } else { - DrainOutcome::Full - } -} -/// The SIGTERM drain budget from `KIGI_WORKSPACE_TERMINATION_GRACE_MS` -/// (default [`DEFAULT_TERMINATION_GRACE_MS`]). The hub-evict path uses the -/// hub-provided `grace_period_ms` instead. -pub fn termination_grace_from_env() -> std::time::Duration { - grace_budget_from_raw(std::env::var("KIGI_WORKSPACE_TERMINATION_GRACE_MS").ok()) -} -/// Pure parse of the termination-grace env value: a positive integer ms wins, -/// anything else (absent, unparseable, zero) falls back to the default. -fn grace_budget_from_raw(raw: Option) -> std::time::Duration { - let ms = raw - .and_then(|s| s.trim().parse::().ok()) - .filter(|&ms| ms > 0) - .unwrap_or(DEFAULT_TERMINATION_GRACE_MS); - std::time::Duration::from_millis(ms) -} -/// Path of the preStop drain marker (`KIGI_WORKSPACE_DRAINING_FILE` or -/// [`DEFAULT_DRAINING_FILE`]). -fn draining_file_path() -> std::path::PathBuf { - std::env::var("KIGI_WORKSPACE_DRAINING_FILE") - .map(std::path::PathBuf::from) - .unwrap_or_else(|_| std::path::PathBuf::from(DEFAULT_DRAINING_FILE)) -} -/// Atomically write `outstanding` (total durability work still pending: upload -/// queue depth + in-flight artifact producers) to the drain marker (temp + -/// fsync + rename) so the preStop hook never reads a torn value and never sees -/// `0` while a producer could still enqueue. Best-effort. The temp name is -/// unique (pid + counter) so concurrent evict drains don't race on a fixed -/// `.tmp`. -fn write_draining_marker(path: &std::path::Path, outstanding: usize) { - use std::io::Write; - use std::sync::atomic::{AtomicU64, Ordering}; - static TMP_SEQ: AtomicU64 = AtomicU64::new(0); - let mut tmp = path.as_os_str().to_owned(); - tmp.push(format!( - ".{}.{}.draining.tmp", - std::process::id(), - TMP_SEQ.fetch_add(1, Ordering::Relaxed) - )); - let tmp = std::path::PathBuf::from(tmp); - let result = std::fs::File::create(&tmp) - .and_then(|mut f| { - f.write_all(outstanding.to_string().as_bytes())?; - f.sync_all() - }) - .and_then(|()| std::fs::rename(&tmp, path)); - if let Err(e) = result { - tracing::warn!( - path = % path.display(), error = % e, "failed to write drain marker" - ); - let _ = std::fs::remove_file(&tmp); - } -} -/// Stream a file once: SHA-256 over every byte while capturing the -/// `[offset, offset + length)` overlap. Returns -/// `(hash_hex, range_bytes, total_streamed_bytes)`. -/// -/// Shared by [`WorkspaceHandle::get_files`]' chunked reads and the -/// `file_system::client_fs` ops so the overlap arithmetic lives in one -/// place. -pub(crate) async fn stream_hash_and_range( - path: &std::path::Path, - offset: u64, - length: u64, -) -> std::io::Result<(String, Vec, u64)> { - use sha2::{Digest, Sha256}; - use tokio::io::AsyncReadExt; - let req_end = offset.saturating_add(length); - let mut f = tokio::fs::File::open(path).await?; - let mut hasher = Sha256::new(); - let mut chunk = Vec::new(); - let mut pos: u64 = 0; - let mut buf = [0u8; 64 * 1024]; - loop { - let n = f.read(&mut buf).await?; - if n == 0 { - break; - } - hasher.update(&buf[..n]); - let start = pos.max(offset); - let end = (pos + n as u64).min(req_end); - if start < end { - let local_start = (start - pos) as usize; - let local_end = (end - pos) as usize; - chunk.extend_from_slice(&buf[local_start..local_end]); - } - pos += n as u64; - } - Ok((format!("{:x}", hasher.finalize()), chunk, pos)) -} -/// Create a [`WorkspaceHandle`] and connect it to the hub. -/// -/// This is the shared setup used by both the standalone `workspace_server` -/// binary and the TUI's in-process local workspace server. The workspace -/// registers its tools on the server so external clients can reach them. -/// Sessions are bound dynamically by clients calling `bind_server`. -/// -/// `confine_fs_to_workspace_root` confines `x.ai/fs/*` resolution to the root. -/// The standalone workspace server defaults it on (it always backs a remote -/// sandbox; override via `KIGI_WORKSPACE_CONFINE_FS_TO_ROOT`); the CLI leader -/// passes `false`. -/// -/// Returns the connected handle (caller should keep it alive for the -/// lifetime of the server connection). -pub async fn connect_local_workspace( - cwd: std::path::PathBuf, - hub_url: url::Url, - auth: kigi_computer_hub_sdk::SharedAuthProvider, - metadata: Option, - server_id: Option, - alpha_test_key: Option, - allow_insecure_ws: bool, - status_config: crate::status_config::StatusConfig, - project_lsp_trusted: bool, - ready_file: Option, - diag: Option, - require_explicit_toolset: bool, - confine_fs_to_workspace_root: bool, -) -> WorkspaceResult { - use crate::session::tool_config::WorkspaceSessionContextFactory; - let workspace_home = resolve_workspace_home(); - std::fs::create_dir_all(&workspace_home).map_err(|e| { - WorkspaceError::HubError(format!( - "failed to create workspace home {}: {e}", - workspace_home.display() - )) - })?; - let api_base_url = kigi_env::coding_api_base_url(); - let mut factory = WorkspaceSessionContextFactory::with_auth(auth.clone(), api_base_url.clone()); - if crate::session::tool_config::tool_state_enabled() { - factory = factory.with_tool_state_home(workspace_home.clone()); - } - let hub_cfg = crate::hub::HubConfig { - url: hub_url, - auth: auth.clone(), - activity_tracker: None, - server_id, - alpha_test_key, - allow_insecure_ws, - ready_file, - diag, - }; - let tool_config = kigi_agent::workspace_grok_build_toolset(); - let mut ws_config = WorkspaceConfig::new_for_proxy( - cwd, - Arc::new(factory), - hub_cfg, - auth.clone(), - metadata, - status_config, - tool_config, - ); - ws_config.project_lsp_trusted = project_lsp_trusted; - ws_config.require_explicit_toolset = require_explicit_toolset; - ws_config.confine_fs_to_workspace_root = confine_fs_to_workspace_root; - if let Ok(dir) = std::env::var("KIGI_WORKSPACE_SERVER_SKILLS_DIR") - && !dir.is_empty() - { - ws_config.skills_config.server_skill_dirs = vec![dir]; - } - if let Ok(dir) = std::env::var("KIGI_WORKSPACE_BUNDLED_SKILLS_DIR") - && !dir.is_empty() - { - let allowlist = std::env::var("KIGI_WORKSPACE_BUNDLED_SKILLS_ALLOWLIST").ok(); - ws_config - .skills_config - .ignore - .extend(bundled_allowlist_ignore_dirs(&dir, allowlist.as_deref())); - ws_config.skills_config.bundled_skill_dirs = vec![dir]; - } - if crate::session::tool_config::tool_state_enabled() { - let home = workspace_home.clone(); - tokio::spawn(async move { - crate::recovery::cleanup_stale_sessions( - &home, - crate::recovery::DEFAULT_SESSION_MAX_AGE, - ) - .await; - }); - } - let ws_handle = WorkspaceHandle::new_with_home(ws_config, workspace_home) - .map_err(|e| WorkspaceError::HubError(format!("failed to create workspace: {e}")))?; - ws_handle.connect_hub().await?; - Ok(ws_handle) -} -/// Resolve `$KIGI_WORKSPACE_HOME` — the workspace-owned on-disk state root. -/// -/// Precedence: -/// 1. `$KIGI_WORKSPACE_HOME` (operator override). -/// 2. `/workspace`, where `` honours `$KIGI_SHARE_DIR` and -/// otherwise falls back to `~/.kigi` (see [`kigi_config::kigi_home`]). -pub fn resolve_workspace_home() -> std::path::PathBuf { - if let Ok(p) = std::env::var("KIGI_WORKSPACE_HOME") - && !p.trim().is_empty() - { - return std::path::PathBuf::from(p); - } - kigi_config::kigi_home().join("workspace") -} -/// Skill `ignore` entries for the allow-list: subdirs of `dir` not in the -/// comma-separated list (`bundled__` prefix optional). Blank list → none; -/// unreadable `dir` → ignore `dir` itself (fail closed). -fn bundled_allowlist_ignore_dirs(dir: &str, allowlist: Option<&str>) -> Vec { - let allowed: std::collections::HashSet<&str> = allowlist - .unwrap_or_default() - .split(',') - .map(|s| s.trim()) - .map(|s| s.strip_prefix("bundled__").unwrap_or(s)) - .filter(|s| !s.is_empty()) - .collect(); - if allowed.is_empty() { - return vec![]; - } - let entries = match std::fs::read_dir(dir) { - Ok(entries) => entries, - Err(err) => { - tracing::warn!( - dir, % err, - "bundled skills dir unreadable; allow-list ignores the whole dir" - ); - return vec![dir.to_string()]; - } - }; - let mut dirs: Vec = entries - .filter_map(|entry| entry.ok()) - .filter(|entry| entry.path().is_dir()) - .filter(|entry| { - let name = entry.file_name(); - let name = name.to_string_lossy(); - let stripped = name.strip_prefix("bundled__").unwrap_or(&name); - !allowed.contains(stripped) - }) - .map(|entry| entry.path().to_string_lossy().into_owned()) - .collect(); - dirs.sort(); - dirs } /// Whether per-session `events.jsonl` recording is enabled /// (`KIGI_WORKSPACE_EVENTS_ENABLED=true`). Any other value — including unset — @@ -3329,141 +1287,6 @@ fn ephemeral_workspace_home() -> std::path::PathBuf { fn rewind_all_outcomes_from_env() -> bool { kigi_config::env_bool("KIGI_WORKSPACE_REWIND_ALL_OUTCOMES").unwrap_or(false) } -/// `ToolHandle` adapter that delegates to a workspace session's -/// [`FinalizedToolset`]. Used by [`WorkspaceHandle::create_local_harness`] -/// to populate a [`LocalRegistry`] for in-process tool dispatch. -/// -/// This is the same dispatch pattern as [`SessionRoutedToolHandler`] in -/// `hub.rs`, but implements `ToolHandle` (for `LocalRegistry`) instead -/// of `ToolServerHandler` (for `ToolServer`). -struct SessionToolHandle { - tool_id: kigi_tool_protocol::ToolId, - desc: kigi_tool_types::ToolDescription, - workspace: WorkspaceHandle, - session_id: String, -} -impl SessionToolHandle { - fn new( - tool_name: String, - desc: kigi_tool_types::ToolDescription, - workspace: WorkspaceHandle, - session_id: String, - ) -> Result { - Ok(Self { - tool_id: kigi_tool_protocol::ToolId::new(tool_name)?, - desc, - workspace, - session_id, - }) - } - fn name(&self) -> &str { - self.tool_id.as_str() - } -} -impl std::fmt::Debug for SessionToolHandle { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("SessionToolHandle") - .field("tool_name", &self.name()) - .field("session_id", &self.session_id) - .finish() - } -} -#[async_trait::async_trait] -impl kigi_tool_runtime::ToolDyn for SessionToolHandle { - fn id(&self) -> kigi_tool_protocol::ToolId { - self.tool_id.clone() - } - fn description( - &self, - _ctx: &::kigi_tool_runtime::ListToolsContext, - ) -> kigi_tool_types::ToolDescription { - self.desc.clone() - } - async fn execute( - &self, - ctx: kigi_tool_runtime::ToolCallContext, - args: serde_json::Value, - ) -> kigi_tool_runtime::ToolStream { - use kigi_tool_runtime::{ToolError, ToolErrorKind, ToolStreamItem, terminal_only}; - let session = match self.workspace.session(&self.session_id) { - Some(s) => s, - None => { - return terminal_only(Err(ToolError::new( - ToolErrorKind::InvalidArguments, - format!("session not bound: {}", self.session_id), - ))); - } - }; - let toolset = session.toolset(); - let call_id = ctx.call_id.to_string(); - let tool_id = self.id(); - let tool_name = self.name().to_owned(); - let session_label = self.session_id.clone(); - tracing::debug!( - tool = % self.name(), call_id = % call_id, session = % self.session_id, - "local harness: dispatching tool call" - ); - let inner = toolset.call_streaming(self.name(), args, &call_id, None); - Box::pin(async_stream::stream! { - use futures::StreamExt; let mut inner = inner; while let Some(item) = - inner.next(). await { match item { ToolStreamItem::Progress(p) => { yield - ToolStreamItem::Progress(p); } ToolStreamItem::Terminal(Ok(run_result)) - => { yield ToolStreamItem::Terminal(Ok(run_result - .into_typed_tool_output(tool_id),)); return; } - ToolStreamItem::Terminal(Err(e)) => { tracing::error!(tool = % tool_name, - session = % session_label, error = % e, - "local harness tool call failed"); yield - ToolStreamItem::Terminal(Err(ToolError::new(ToolErrorKind::TerminalError, - e.to_string(),))); return; } } } yield - ToolStreamItem::Terminal(Err(ToolError::new(ToolErrorKind::TerminalError, - "tool stream ended without a terminal",))); - }) - } -} -impl WorkspaceHandle { - /// Create a local-only [`ToolHarness`] backed by this workspace's - /// session toolset. - /// - /// Tools are dispatched in-process via a [`LocalRegistry`] — no hub - /// connection needed. Each tool is resolved dynamically from the - /// session's live [`FinalizedToolset`] at call time, so tool config - /// hot-reloads (via `update_tool_config()`) take effect automatically. - pub fn create_local_harness( - &self, - session_id: &str, - ) -> WorkspaceResult { - let session = self - .session(session_id) - .ok_or_else(|| WorkspaceError::SessionNotFound(session_id.to_string()))?; - let toolset = session.toolset(); - let registry = kigi_computer_hub_sdk::LocalRegistry::new(); - for def in toolset.tool_definitions() { - let tool_name = def.function.name.clone(); - let desc = kigi_tool_types::ToolDescription::new( - tool_name.clone(), - def.function.description.clone().unwrap_or_default(), - ); - match SessionToolHandle::new(tool_name, desc, self.clone(), session_id.to_string()) { - Ok(tool) => { - registry.register_dyn(Arc::new(tool) as Arc); - } - Err(e) => { - tracing::warn!( - tool = % def.function.name, error = % e, - "client name is not a valid ToolId; skipping local-harness registration" - ); - } - } - } - let session_id = kigi_tool_protocol::SessionId::new(session_id.to_string()) - .map_err(|e| WorkspaceError::HubError(format!("invalid session id: {e}")))?; - Ok(kigi_computer_hub_sdk::ToolHarness::local_only_with( - registry, - session_id, - kigi_tool_runtime::TypedExtensions::default(), - )) - } -} impl WorkspaceHandle { /// Minimal handle for local mode (no hub). Requires Tokio runtime. pub fn new_minimal( @@ -3485,12 +1308,8 @@ impl WorkspaceHandle { hook_project_sources: vec![], skills_config: Default::default(), plugin_discovery_config: Default::default(), - hub_config: None, - auth_provider: None, - server_metadata: None, status_config: Default::default(), project_lsp_trusted, - require_explicit_toolset: false, confine_fs_to_workspace_root: false, }; Self::build( @@ -3522,12 +1341,8 @@ impl WorkspaceHandle { hook_project_sources: vec![], skills_config: Default::default(), plugin_discovery_config: Default::default(), - hub_config: None, - auth_provider: None, - server_metadata: None, status_config: Default::default(), project_lsp_trusted: true, - require_explicit_toolset: false, confine_fs_to_workspace_root: false, } } @@ -3565,32 +1380,16 @@ pub(crate) mod tests { pub(crate) fn make_handle() -> WorkspaceHandle { make_handle_with_rewind_all_outcomes(false) } - /// [`make_handle`] with `require_explicit_toolset` (strict sandbox mode). - pub(crate) fn make_strict_handle() -> WorkspaceHandle { - make_handle_with_options(false, true) - } - /// [`make_handle`] with fs confinement on (mirrors a remote-sandbox server). + /// [`make_handle`] with fs confinement on. pub(crate) fn make_confining_handle() -> WorkspaceHandle { - make_handle_inner(false, false, Default::default(), true) + make_handle_inner(false, Default::default(), true) } /// [`make_handle`] with an explicit `workspace_rewind_all_outcomes` value. pub(crate) fn make_handle_with_rewind_all_outcomes(enabled: bool) -> WorkspaceHandle { - make_handle_inner(enabled, false, Default::default(), false) - } - pub(crate) fn make_handle_with_options( - rewind_all_outcomes: bool, - require_explicit_toolset: bool, - ) -> WorkspaceHandle { - make_handle_inner( - rewind_all_outcomes, - require_explicit_toolset, - Default::default(), - false, - ) + make_handle_inner(enabled, Default::default(), false) } fn make_handle_inner( rewind_all_outcomes: bool, - require_explicit_toolset: bool, status_config: crate::StatusConfig, confine_fs_to_workspace_root: bool, ) -> WorkspaceHandle { @@ -3607,14 +1406,8 @@ pub(crate) mod tests { hook_project_sources: vec![], skills_config: Default::default(), plugin_discovery_config: Default::default(), - hub_config: None, - auth_provider: Some(Arc::new(kigi_computer_hub_sdk::AuthCredential::bearer( - "test-token", - ))), - server_metadata: None, status_config, project_lsp_trusted: true, - require_explicit_toolset, confine_fs_to_workspace_root, }; let handle = WorkspaceHandle::build( @@ -3629,146 +1422,6 @@ pub(crate) mod tests { .expect("create main session should succeed"); handle } - pub(crate) const BASH_CCO_STUB_NAME: &str = "bash_cco_stub"; - pub(crate) const BASH_CCO_STUB_STDOUT: &str = "cco-stdout"; - #[derive(Debug)] - pub(crate) struct BashCcoStub; - impl kigi_tools::types::tool_metadata::ToolMetadata for BashCcoStub { - fn kind(&self) -> ToolKind { - ToolKind::Execute - } - fn tool_namespace(&self) -> kigi_tools::types::tool::ToolNamespace { - kigi_tools::types::tool::ToolNamespace::MCP - } - fn description_template(&self) -> &str { - "bash cco stub" - } - } - impl kigi_tool_runtime::Tool for BashCcoStub { - type Args = serde_json::Value; - type Output = kigi_tools::types::output::ToolOutput; - fn id(&self) -> kigi_tool_protocol::ToolId { - kigi_tool_protocol::ToolId::new(BASH_CCO_STUB_NAME).expect("valid tool id") - } - fn description( - &self, - _ctx: &::kigi_tool_runtime::ListToolsContext, - ) -> kigi_tool_types::ToolDescription { - kigi_tool_types::ToolDescription::new(BASH_CCO_STUB_NAME, "bash cco stub") - } - async fn run( - &self, - _ctx: kigi_tool_runtime::ToolCallContext, - _input: serde_json::Value, - ) -> Result { - let output = BASH_CCO_STUB_STDOUT.as_bytes(); - Ok(kigi_tools::types::output::ToolOutput::Bash( - kigi_tools::types::output::BashOutput { - output: output.to_vec(), - output_for_prompt: - kigi_tools::types::output::BashOutput::make_output_for_prompt( - BASH_CCO_STUB_STDOUT, - ), - exit_code: 0, - command: format!("echo {BASH_CCO_STUB_STDOUT}"), - truncated: false, - signal: None, - timed_out: false, - description: None, - current_dir: "/tmp".into(), - output_file: String::new(), - total_bytes: output.len(), - output_delta: None, - was_bare_echo: false, - }, - )) - } - } - pub(crate) fn register_bash_cco_stub(handle: &WorkspaceHandle) { - let session = handle.session("main").expect("main session present"); - session - .toolset() - .register_tool( - BASH_CCO_STUB_NAME.to_owned(), - BashCcoStub, - Some(serde_json::json!({ "type" : "object", "properties" : {} })), - ) - .expect("register bash_cco_stub"); - } - pub(crate) fn assert_bash_cco_terminal(typed: &kigi_tool_runtime::TypedToolOutput) { - use kigi_tool_runtime::ToolOutput as _; - let resp = typed - .chat_completion_output() - .expect("bash chat_completion_output must be preserved"); - let cer = resp - .result - .as_ref() - .and_then(|r| r.code_execution_result.as_ref()) - .expect("code_execution_result"); - assert_eq!(cer.stdout, BASH_CCO_STUB_STDOUT); - assert_eq!(cer.exit_code, 0); - assert!(!cer.command_timed_out); - } - pub(crate) async fn drain_terminal_ok( - mut stream: impl futures::Stream< - Item = kigi_tool_runtime::ToolStreamItem, - > + Unpin, - ) -> kigi_tool_runtime::TypedToolOutput { - use futures::StreamExt; - use kigi_tool_runtime::ToolStreamItem; - while let Some(item) = stream.next().await { - match item { - ToolStreamItem::Terminal(Ok(t)) => return t, - ToolStreamItem::Progress(_) => {} - ToolStreamItem::Terminal(Err(e)) => { - panic!("expected Terminal(Ok), got Err: {e}") - } - } - } - panic!("stream ended without terminal") - } - #[tokio::test] - async fn local_harness_preserves_bash_chat_completion_output() { - use kigi_tool_runtime::ToolCallContext; - let handle = make_handle(); - register_bash_cco_stub(&handle); - let harness = handle.create_local_harness("main").expect("local harness"); - let tool_id = kigi_tool_protocol::ToolId::new(BASH_CCO_STUB_NAME).expect("valid tool id"); - let stream = harness - .call(tool_id, serde_json::json!({}), ToolCallContext::default()) - .await; - let typed = drain_terminal_ok(stream).await; - assert_bash_cco_terminal(&typed); - } - /// No connection ⇒ every export entry point returns `None`, so the - /// binary leaves the `DonatingLogLayer` inert and spawns no metric reporter. - /// This is the flag-free "activate only on connection" contract that log - /// and metric export share with the pre-existing `trace_donation_reporter`. - #[tokio::test] - async fn donation_entry_points_are_inert_without_a_hub() { - let handle = make_handle(); - assert!( - handle - .trace_donation_reporter("prod_grok_workspace") - .await - .is_none(), - "trace export must stay inert without a connection" - ); - assert!( - handle - .log_donation_layer("prod_grok_workspace") - .await - .is_none(), - "log export must stay inert without a connection" - ); - assert!( - handle - .metric_donation_reporter("prod_grok_workspace") - .await - .is_none(), - "metric export must stay inert without a connection" - ); - } #[test] fn rewind_outcome_label_maps_each_variant() { assert_eq!( @@ -3789,144 +1442,6 @@ pub(crate) mod tests { assert_eq!(rewind_result_label(true), "success"); assert_eq!(rewind_result_label(false), "failure"); } - /// The per-bind handler builder maps the session's finalized toolset 1:1 — - /// one handler per `tool_definitions()` entry, keyed by client name, with no - /// extra handlers and no RPC handler (that is appended by the resolver / - /// `connect_hub`, not here). The resolver-level "no intersection, no silent - /// drop" guarantee is covered by - /// [`resolver_advertises_tool_absent_from_connect_catalog`]. - #[tokio::test] - async fn build_session_routed_handlers_covers_finalized_toolset() { - let handle = make_handle(); - let session = handle.session("main").expect("main session exists"); - let toolset = session.toolset(); - let expected: std::collections::HashSet = toolset - .tool_definitions() - .iter() - .map(|d| d.function.name.clone()) - .collect(); - assert!( - expected.contains("read_file"), - "baseline toolset should expose read_file" - ); - let handlers = build_session_routed_handlers(&toolset, &handle); - let got: std::collections::HashSet = handlers - .iter() - .map(|h| h.tool_id().as_str().to_owned()) - .collect(); - assert_eq!(handlers.len(), expected.len(), "one handler per tool def"); - assert_eq!( - got, expected, - "advertised handlers must equal the finalized toolset (no intersection)" - ); - } - #[tokio::test] - async fn build_session_routed_handlers_skips_invalid_client_name_without_panic() { - let handle = make_handle(); - let mut renamed = tc("GrokBuild:read_file", Some(ToolKind::Read)); - renamed.name_override = Some("bad name!".to_owned()); - let session = handle - .create_session_with_config( - "sess-invalid-name", - None, - Some(ToolServerConfig { - tools: vec![renamed, tc("GrokBuild:grep", Some(ToolKind::Read))], - behavior_preset: None, - }), - CapabilityMode::All, - None, - false, - ) - .expect("create session with invalidly renamed tool"); - let handlers = build_session_routed_handlers(&session.toolset(), &handle); - let names: Vec = handlers - .iter() - .map(|h| h.tool_id().as_str().to_owned()) - .collect(); - assert!( - !names.iter().any(|n| n == "bad name!"), - "the invalid client name must be skipped: {names:?}" - ); - assert!( - names.iter().any(|n| n == "grep"), - "valid tools must still get handlers: {names:?}" - ); - } - /// Regression for the deleted catalog intersection. Reproduces the - /// `session.bind` resolver tail's composition — `build_session_routed_handlers` - /// for the session toolset, plus the single RPC handler filtered from the - /// connect-time catalog — and proves a session tool whose client name is - /// ABSENT from that (grok-build) catalog is still advertised. The old - /// `catalog ∩ session-names` filter silently dropped exactly such tools - /// (grok-build renames → 6/11). - #[tokio::test] - async fn resolver_advertises_tool_absent_from_connect_catalog() { - let handle = make_handle(); - let catalog_toolset = handle - .session("main") - .expect("main session exists") - .toolset(); - let mut catalog = build_session_routed_handlers(&catalog_toolset, &handle); - let rpc_handler: Arc = - Arc::new(crate::hub_server::WorkspaceRpcHandler::new(handle.clone())); - let rpc_tool_id = rpc_handler.tool_id(); - catalog.push(rpc_handler); - let catalog_names: std::collections::HashSet = catalog - .iter() - .map(|h| h.tool_id().as_str().to_owned()) - .collect(); - let mut renamed = tc("GrokBuild:read_file", Some(ToolKind::Read)); - renamed.name_override = Some("non_catalog_tool".to_owned()); - let session = handle - .create_session_with_config( - "sess-non-catalog", - None, - Some(ToolServerConfig { - tools: vec![renamed], - behavior_preset: None, - }), - CapabilityMode::All, - None, - false, - ) - .expect("create session with renamed tool"); - assert!( - !catalog_names.contains("non_catalog_tool"), - "precondition: the renamed tool must be absent from the catalog" - ); - let toolset = session.toolset(); - let mut handlers = build_session_routed_handlers(&toolset, &handle); - handlers.extend( - catalog - .iter() - .filter(|h| h.tool_id() == rpc_tool_id) - .cloned(), - ); - let advertised: std::collections::HashSet = handlers - .iter() - .map(|h| h.tool_id().as_str().to_owned()) - .collect(); - assert!( - advertised.contains("non_catalog_tool"), - "a session tool absent from the catalog must still be advertised" - ); - assert_eq!( - handlers - .iter() - .filter(|h| h.tool_id() == rpc_tool_id) - .count(), - 1, - "exactly one RPC handler appended" - ); - let mut expected: std::collections::HashSet = toolset - .tool_definitions() - .iter() - .map(|d| d.function.name.clone()) - .collect(); - expected.insert(rpc_tool_id.as_str().to_owned()); - assert_eq!(advertised, expected); - } - /// Client names advertised by a session's current toolset. fn session_tool_names(session: &Arc) -> Vec { session .toolset() @@ -3935,463 +1450,6 @@ pub(crate) mod tests { .map(|d| d.function.name.clone()) .collect() } - /// The sandbox-resume regression (`workspace_tool_coverage_incomplete`): a - /// session created by a metadata-less bind resolves the workspace default; - /// a later rebind that carries the client's explicit toolset must - /// re-resolve and swap it in — not silently reuse the default — so the - /// bind response advertises the configured (renamed) tools. A repeat - /// rebind with the identical config is a no-op reuse. - #[tokio::test] - async fn rebind_with_changed_explicit_toolset_reresolves_and_swaps() { - let handle = make_handle(); - let session = handle - .create_session_with_config("resumed", None, None, CapabilityMode::All, None, false) - .expect("create default-resolved session"); - session.set_bind_tool_config_fingerprint(None); - assert!( - session_tool_names(&session) - .iter() - .all(|n| n != "renamed_read"), - "precondition: the default toolset must not carry the override name" - ); - let mut renamed = tc("GrokBuild:read_file", Some(ToolKind::Read)); - renamed.name_override = Some("renamed_read".to_owned()); - let cfg = ToolServerConfig { - tools: vec![renamed], - behavior_preset: None, - }; - let fingerprint = serde_json::to_value(&cfg).ok(); - let (rebound, outcome) = handle - .rebind_existing_hub_session("resumed", Some(cfg.clone()), fingerprint.clone()) - .await - .expect("session exists"); - assert_eq!(outcome, RebindOutcome::Reresolved); - assert_eq!( - session_tool_names(&rebound), - vec!["renamed_read".to_owned()], - "the rebind must swap in the explicit toolset's resolution" - ); - let (_, outcome) = handle - .rebind_existing_hub_session("resumed", Some(cfg), fingerprint) - .await - .expect("session exists"); - assert_eq!(outcome, RebindOutcome::Reused); - } - /// A rebind without an explicit toolset (default resolution, or the - /// fail-closed placeholders which the caller maps to `None`) must never - /// downgrade an explicitly-configured session to the default toolset. - #[tokio::test] - async fn rebind_without_explicit_toolset_reuses_existing() { - let handle = make_handle(); - let mut renamed = tc("GrokBuild:read_file", Some(ToolKind::Read)); - renamed.name_override = Some("renamed_read".to_owned()); - let cfg = ToolServerConfig { - tools: vec![renamed], - behavior_preset: None, - }; - let session = handle - .create_session_with_config( - "configured", - None, - Some(cfg.clone()), - CapabilityMode::All, - None, - false, - ) - .expect("create configured session"); - session.set_bind_tool_config_fingerprint(serde_json::to_value(&cfg).ok()); - let (rebound, outcome) = handle - .rebind_existing_hub_session("configured", None, None) - .await - .expect("session exists"); - assert_eq!(outcome, RebindOutcome::Reused); - assert_eq!( - session_tool_names(&rebound), - vec!["renamed_read".to_owned()], - "a metadata-less rebind must not clobber the configured toolset" - ); - } - /// The create arm's fingerprint write is set-if-unset: a concurrent - /// rebind that already swapped in its toolset (and recorded its - /// fingerprint under `update_lock`) must not be clobbered by the create - /// task's deferred write, or a later identical rebind would `Reused`-skip - /// against a fingerprint that no longer describes the live toolset. - #[tokio::test] - async fn create_fingerprint_write_does_not_clobber_concurrent_rebind() { - let handle = make_handle(); - let session = handle - .create_session_with_config("racy", None, None, CapabilityMode::All, None, false) - .expect("create session"); - let mut renamed = tc("GrokBuild:read_file", Some(ToolKind::Read)); - renamed.name_override = Some("renamed_read".to_owned()); - let cfg_b = ToolServerConfig { - tools: vec![renamed], - behavior_preset: None, - }; - let fp_b = serde_json::to_value(&cfg_b).ok(); - let (_, outcome) = handle - .rebind_existing_hub_session("racy", Some(cfg_b.clone()), fp_b.clone()) - .await - .expect("session exists"); - assert_eq!(outcome, RebindOutcome::Reresolved); - let fp_a = serde_json::to_value(&ToolServerConfig { - tools: vec![tc("GrokBuild:list_dir", Some(ToolKind::ListDir))], - behavior_preset: None, - }) - .ok(); - session.set_bind_tool_config_fingerprint_if_unset(fp_a); - let (rebound, outcome) = handle - .rebind_existing_hub_session("racy", Some(cfg_b), fp_b) - .await - .expect("session exists"); - assert_eq!(outcome, RebindOutcome::Reused); - assert_eq!( - session_tool_names(&rebound), - vec!["renamed_read".to_owned()] - ); - } - /// A vanished session yields `None` (the caller falls back to RPC-only). - #[tokio::test] - async fn rebind_missing_session_returns_none() { - let handle = make_handle(); - assert!( - handle - .rebind_existing_hub_session("no-such-session", None, None) - .await - .is_none() - ); - } - fn swap_rejected_count(reason: &str, trigger: &str) -> u64 { - crate::session::swap_policy::WORKSPACE_TOOLSET_SWAP_REJECTED_TOTAL - .with_label_values(&[reason, trigger]) - .get() - } - /// The lazy-bind / resume-correction regression lock: a - /// default-resolved session (stored fingerprint `None`) must accept the - /// owner's explicit-config rebind even mid-turn with a call in flight — - /// the owner bind is designed to land mid-turn, and deferring it would - /// serve a toolset that contradicts the config-built prompt. - #[tokio::test] - async fn rebind_none_to_explicit_swaps_mid_turn() { - let handle = make_handle(); - let session = handle - .create_session_with_config("lazy", None, None, CapabilityMode::All, None, false) - .expect("create default-resolved session"); - session.set_bind_tool_config_fingerprint(None); - let tracker = handle.activity_tracker().clone(); - tracker.turn_started("lazy", 1); - tracker.tool_call_started("lazy-c1", "read_file", Some("lazy")); - let cfg = explicit_cfg("renamed_read"); - let fingerprint = serde_json::to_value(&cfg).ok(); - let (rebound, outcome) = handle - .rebind_existing_hub_session("lazy", Some(cfg), fingerprint) - .await - .expect("session exists"); - assert_eq!( - outcome, - RebindOutcome::Reresolved, - "a None → explicit correction must swap even mid-turn with calls in flight" - ); - assert_eq!( - session_tool_names(&rebound), - vec!["renamed_read".to_owned()] - ); - } - /// `explicit → different-explicit` under dispatch: the rebind keeps the - /// existing toolset (`ReresolveDeferredInFlight`, counted); once the - /// call completes, a later rebind applies the correction. - #[tokio::test] - async fn rebind_explicit_to_explicit_with_in_flight_call_defers_then_corrects() { - use kigi_file_utils::events::ToolOutcome; - let rejected_before = swap_rejected_count("in_flight", "owner_rebind"); - let handle = make_handle(); - let cfg_a = explicit_cfg("read_a"); - let session = handle - .create_session_with_config( - "busy", - None, - Some(cfg_a.clone()), - CapabilityMode::All, - None, - false, - ) - .expect("create session with cfg A"); - session.set_bind_tool_config_fingerprint(serde_json::to_value(&cfg_a).ok()); - let tracker = handle.activity_tracker().clone(); - tracker.tool_call_started("busy-c1", "read_a", Some("busy")); - let cfg_b = explicit_cfg("read_b"); - let fp_b = serde_json::to_value(&cfg_b).ok(); - let (kept, outcome) = handle - .rebind_existing_hub_session("busy", Some(cfg_b.clone()), fp_b.clone()) - .await - .expect("session exists"); - assert_eq!(outcome, RebindOutcome::ReresolveDeferredInFlight); - assert_eq!( - session_tool_names(&kept), - vec!["read_a".to_owned()], - "the existing toolset must be kept while a call is in flight" - ); - assert!( - swap_rejected_count("in_flight", "owner_rebind") > rejected_before, - "the deferred swap must be counted" - ); - tracker.tool_call_completed("busy-c1", Some("busy"), ToolOutcome::Success); - let (rebound, outcome) = handle - .rebind_existing_hub_session("busy", Some(cfg_b), fp_b) - .await - .expect("session exists"); - assert_eq!( - outcome, - RebindOutcome::Reresolved, - "the correction must apply once no calls are in flight" - ); - assert_eq!(session_tool_names(&rebound), vec!["read_b".to_owned()]); - } - /// A reconnect's identical `session.bind` heals a stale session: reuse - /// without the marker, defer in-flight, rebuild + clear once idle. - #[tokio::test] - async fn rebind_identical_reapply_repairs_stale_resolve() { - use kigi_file_utils::events::ToolOutcome; - let handle = make_handle(); - let cfg = explicit_cfg("renamed_read"); - let fingerprint = serde_json::to_value(&cfg).ok(); - let session = handle - .create_session_with_config( - "stale-rebind", - None, - Some(cfg.clone()), - CapabilityMode::All, - None, - false, - ) - .expect("create session"); - session.set_bind_tool_config_fingerprint(fingerprint.clone()); - let toolset_before = session.toolset(); - let (_, outcome) = handle - .rebind_existing_hub_session("stale-rebind", Some(cfg.clone()), fingerprint.clone()) - .await - .expect("session exists"); - assert_eq!(outcome, RebindOutcome::Reused); - assert!( - Arc::ptr_eq(&session.toolset(), &toolset_before), - "without the stale marker the identical rebind must not rebuild" - ); - session.mark_stale_resolve(); - let tracker = handle.activity_tracker().clone(); - tracker.tool_call_started("stale-c1", "read_file", Some("stale-rebind")); - let rejected_before = swap_rejected_count("in_flight", "owner_rebind"); - let (kept, outcome) = handle - .rebind_existing_hub_session("stale-rebind", Some(cfg.clone()), fingerprint.clone()) - .await - .expect("session exists"); - assert_eq!( - outcome, - RebindOutcome::ReresolveDeferredInFlight, - "the heal must defer while a call is in flight" - ); - assert!( - Arc::ptr_eq(&kept.toolset(), &toolset_before), - "the deferred heal must keep the existing toolset" - ); - assert!(kept.stale_resolve(), "the deferred heal keeps the marker"); - assert!( - swap_rejected_count("in_flight", "owner_rebind") > rejected_before, - "the deferred heal must be counted" - ); - tracker.tool_call_completed("stale-c1", Some("stale-rebind"), ToolOutcome::Success); - let (healed, outcome) = handle - .rebind_existing_hub_session("stale-rebind", Some(cfg), fingerprint) - .await - .expect("session exists"); - assert_eq!( - outcome, - RebindOutcome::Reresolved, - "the idle reconnect must repair the stale toolset" - ); - assert!( - !Arc::ptr_eq(&healed.toolset(), &toolset_before), - "the heal must install a freshly resolved toolset" - ); - assert!( - !healed.stale_resolve(), - "a successful install must clear the stale marker" - ); - } - /// The RPC path rejects a mid-turn config change with the retryable - /// `TurnActive` error (counted); the retry at the turn boundary succeeds. - #[tokio::test] - async fn update_tool_config_rejects_mid_turn_then_succeeds_at_boundary() { - let rejected_before = swap_rejected_count("turn_active", "update_tool_config"); - let handle = make_handle(); - handle.activity_tracker().turn_started("main", 1); - let cfg = explicit_cfg("renamed_read"); - let err = handle - .update_tool_config("main", "main", cfg.clone()) - .await - .expect_err("a mid-turn config change must be rejected"); - assert!( - matches!(err, WorkspaceError::TurnActive(ref s) if s == "main"), - "got {err:?}" - ); - assert!( - swap_rejected_count("turn_active", "update_tool_config") > rejected_before, - "the rejection must be counted" - ); - let session = handle.session("main").expect("main session exists"); - assert!( - session_tool_names(&session) - .iter() - .all(|n| n != "renamed_read"), - "the rejected config must not take effect" - ); - handle.activity_tracker().turn_completed("main", 1, 0); - handle - .update_tool_config("main", "main", cfg) - .await - .expect("the retry at the turn boundary must succeed"); - let session = handle.session("main").expect("main session exists"); - assert_eq!( - session_tool_names(&session), - vec!["renamed_read".to_owned()] - ); - } - /// TOCTOU lock: a turn that starts DURING the re-resolve (after the - /// entry check passed) must still abort the install — the resolved - /// toolset is discarded, the fingerprint stays unchanged, and the - /// rejection is counted under `reason="turn_active_late"`. The retry at - /// the turn boundary then succeeds. - #[tokio::test] - async fn update_tool_config_rejects_turn_started_during_resolve() { - let late_rejected_before = swap_rejected_count("turn_active_late", "update_tool_config"); - let handle = make_handle(); - let session = handle.session("main").expect("main session exists"); - let toolset_before = session.toolset(); - let hook_handle = handle.clone(); - *handle.shared.post_resolve_test_hook.lock() = Some(Box::new(move || { - hook_handle.activity_tracker().turn_started("main", 7); - })); - let cfg = explicit_cfg("late_read"); - let err = handle - .update_tool_config("main", "main", cfg.clone()) - .await - .expect_err("a turn starting mid-resolve must abort the install"); - assert!( - matches!(err, WorkspaceError::TurnActive(ref s) if s == "main"), - "got {err:?}" - ); - assert!( - swap_rejected_count("turn_active_late", "update_tool_config") > late_rejected_before, - "the post-resolve rejection must be counted distinctly" - ); - let session = handle.session("main").expect("main session exists"); - assert!( - Arc::ptr_eq(&session.toolset(), &toolset_before), - "the resolved toolset must be discarded, not installed" - ); - assert!( - session.bind_tool_config_matches(None), - "the unapplied config's fingerprint must NOT be recorded" - ); - *handle.shared.post_resolve_test_hook.lock() = None; - handle.activity_tracker().turn_completed("main", 7, 0); - handle - .update_tool_config("main", "main", cfg) - .await - .expect("the retry at the turn boundary must succeed"); - let session = handle.session("main").expect("main session exists"); - assert_eq!(session_tool_names(&session), vec!["late_read".to_owned()]); - } - /// Re-applying the session's current config mid-turn stays allowed - /// (matching fingerprint), so hot-reload re-applies keep working - /// during turns. - #[tokio::test] - async fn update_tool_config_reapply_of_current_config_allowed_mid_turn() { - let handle = make_handle(); - let cfg = explicit_cfg("renamed_read"); - let session = handle - .create_session_with_config( - "hot", - None, - Some(cfg.clone()), - CapabilityMode::All, - None, - false, - ) - .expect("create session"); - session.set_bind_tool_config_fingerprint(serde_json::to_value(&cfg).ok()); - handle.activity_tracker().turn_started("hot", 1); - handle - .update_tool_config("hot", "hot", cfg) - .await - .expect("an identical-config re-apply must not be turn_active-rejected"); - } - #[tokio::test] - async fn update_tool_config_identical_reapply_repairs_stale_resolve() { - let handle = make_handle(); - let cfg = explicit_cfg("renamed_read"); - let session = handle - .create_session_with_config( - "stale", - None, - Some(cfg.clone()), - CapabilityMode::All, - None, - false, - ) - .expect("create session"); - session.set_bind_tool_config_fingerprint(serde_json::to_value(&cfg).ok()); - let toolset_before = session.toolset(); - handle - .update_tool_config("stale", "stale", cfg.clone()) - .await - .expect("an identical re-apply must succeed"); - assert!( - Arc::ptr_eq(&session.toolset(), &toolset_before), - "without the stale marker the identical re-apply must not rebuild" - ); - session.mark_stale_resolve(); - let rejected_before = swap_rejected_count("turn_active", "update_tool_config"); - handle.activity_tracker().turn_started("stale", 1); - let err = handle - .update_tool_config("stale", "stale", cfg.clone()) - .await - .expect_err("a mid-turn recovery re-apply must be rejected"); - assert!( - matches!(err, WorkspaceError::TurnActive(ref s) if s == "stale"), - "got {err:?}" - ); - assert!( - swap_rejected_count("turn_active", "update_tool_config") > rejected_before, - "the rejected recovery must be counted" - ); - assert!( - session.stale_resolve(), - "the rejected recovery must keep the stale marker" - ); - assert!( - Arc::ptr_eq(&session.toolset(), &toolset_before), - "the rejected recovery must not install" - ); - handle.activity_tracker().turn_completed("stale", 1, 0); - handle - .update_tool_config("stale", "stale", cfg.clone()) - .await - .expect("the boundary retry must repair the stale toolset"); - let session = handle.session("stale").expect("session exists"); - assert!( - !Arc::ptr_eq(&session.toolset(), &toolset_before), - "the recovery re-apply must install a freshly resolved toolset" - ); - assert!( - !session.stale_resolve(), - "a successful install must clear the stale marker" - ); - assert!( - session.bind_tool_config_matches(serde_json::to_value(&cfg).ok().as_ref()), - "the stored fingerprint must be unchanged by the identical recovery" - ); - } - /// The `Terminal` resource of a session's current toolset. async fn toolset_terminal( toolset: &Arc, ) -> Arc { @@ -4468,57 +1526,9 @@ pub(crate) mod tests { .await .expect("start background task") } - /// A rebind that swaps in a different explicit toolset must rebuild the - /// toolset AROUND the session-owned terminal backend, not a fresh one — - /// that identity is what keeps background tasks alive across the swap. - #[tokio::test] - async fn rebind_swap_preserves_session_terminal_backend() { - let orphaned_before = orphaned_swap_count(); - let handle = make_handle(); - let cfg_a = explicit_cfg("read_a"); - let session = handle - .create_session_with_config( - "owned", - None, - Some(cfg_a.clone()), - CapabilityMode::All, - None, - false, - ) - .expect("create session with cfg A"); - session.set_bind_tool_config_fingerprint(serde_json::to_value(&cfg_a).ok()); - let backend = session.terminal_backend().clone(); - assert!( - Arc::ptr_eq(&backend, &toolset_terminal(&session.toolset()).await), - "create must wire the session-owned backend into the toolset" - ); - let cfg_b = explicit_cfg("read_b"); - let fingerprint_b = serde_json::to_value(&cfg_b).ok(); - let (rebound, outcome) = handle - .rebind_existing_hub_session("owned", Some(cfg_b), fingerprint_b) - .await - .expect("session exists"); - assert_eq!(outcome, RebindOutcome::Reresolved); - assert_eq!(session_tool_names(&rebound), vec!["read_b".to_owned()]); - assert!( - Arc::ptr_eq(&backend, rebound.terminal_backend()), - "the session-owned backend must not be replaced by a swap" - ); - assert!( - Arc::ptr_eq(&backend, &toolset_terminal(&rebound.toolset()).await), - "the swapped-in toolset must reference the session-owned backend" - ); - assert_eq!( - orphaned_swap_count(), - orphaned_before, - "the orphaned-backend tripwire must stay 0" - ); - } - /// A snapshot-driven `re_resolve_all_sessions` rebuild (MCP snapshot - /// change) must also rebuild around the session-owned backend — with a - /// LIVE background task riding through the rebuild. This is the - /// regression lock for snapshot-triggered swaps killing background - /// tasks by minting a fresh backend per session. + /// A snapshot-driven rebuild must rebuild the toolset AROUND the + /// session-owned terminal backend, not a fresh one — that identity is + /// what keeps background tasks alive across the swap. #[tokio::test] async fn re_resolve_all_sessions_preserves_session_terminal_backend() { let orphaned_before = orphaned_swap_count(); @@ -4622,112 +1632,6 @@ pub(crate) mod tests { orphaned_before, "the skip must not fire the orphaned-backend tripwire" ); - let outcome = handle - .resolve_and_swap_session_toolset( - &local, - explicit_cfg("read_new"), - SwapTrigger::UpdateRpc, - ) - .await - .expect("the skip is not an internal error at the choke point"); - assert_eq!(outcome, SwapOutcome::SkippedExternallyOwned); - assert!( - Arc::ptr_eq(&local.toolset(), &external_toolset), - "the choke point must not swap an externally-owned toolset" - ); - assert_eq!(orphaned_swap_count(), orphaned_before); - let err = handle - .update_tool_config("local", "local", explicit_cfg("read_new")) - .await - .expect_err("update_tool_config must refuse an externally-owned toolset"); - assert!( - matches!(err, crate ::error::WorkspaceError::ToolsetExternallyOwned(ref s) if - s == "local"), - "expected ToolsetExternallyOwned, got: {err:?}" - ); - assert!( - Arc::ptr_eq(&local.toolset(), &external_toolset), - "the refused update must leave the toolset untouched" - ); - let fp_local = serde_json::to_value(explicit_cfg("read_local")).ok(); - local.set_bind_tool_config_fingerprint(fp_local.clone()); - let cfg_new = explicit_cfg("read_new2"); - let fp_new = serde_json::to_value(&cfg_new).ok(); - let (rebound, outcome) = handle - .rebind_existing_hub_session("local", Some(cfg_new), fp_new.clone()) - .await - .expect("session exists"); - assert_eq!(outcome, RebindOutcome::KeptExternallyOwned); - assert!( - Arc::ptr_eq(&rebound.toolset(), &external_toolset), - "the rebind must keep the externally-owned toolset" - ); - assert!( - rebound.bind_tool_config_matches(fp_local.as_ref()), - "the stored fingerprint must be unchanged by the skipped swap" - ); - assert!( - !rebound.bind_tool_config_matches(fp_new.as_ref()), - "the unapplied config's fingerprint must NOT be recorded" - ); - assert_eq!(orphaned_swap_count(), orphaned_before); - handle - .update_tool_config("local", "local", explicit_cfg("read_local")) - .await - .expect("an identical config on an externally-owned toolset is a no-op success"); - assert!( - Arc::ptr_eq(&local.toolset(), &external_toolset), - "the identical no-op must leave the externally-owned toolset untouched" - ); - assert!( - local.bind_tool_config_matches(fp_local.as_ref()), - "the identical no-op must leave the stored fingerprint untouched" - ); - assert_eq!(orphaned_swap_count(), orphaned_before); - } - /// A background task started before a toolset swap must still be - /// queryable through the NEW toolset's `Terminal` resource — the - /// swap ⇒ empty task table + SIGKILL incident class. - #[tokio::test] - async fn background_task_survives_toolset_swap() { - let orphaned_before = orphaned_swap_count(); - let handle = make_handle(); - let cfg_a = explicit_cfg("read_a"); - let session = handle - .create_session_with_config( - "bg", - None, - Some(cfg_a.clone()), - CapabilityMode::All, - None, - false, - ) - .expect("create session"); - session.set_bind_tool_config_fingerprint(serde_json::to_value(&cfg_a).ok()); - let out_dir = tempfile::tempdir().expect("temp dir"); - let bg = start_background_sleep(&session, out_dir.path(), "bg-task").await; - let cfg_b = explicit_cfg("read_b"); - let fingerprint_b = serde_json::to_value(&cfg_b).ok(); - let (rebound, outcome) = handle - .rebind_existing_hub_session("bg", Some(cfg_b), fingerprint_b) - .await - .expect("session exists"); - assert_eq!(outcome, RebindOutcome::Reresolved); - let new_terminal = toolset_terminal(&rebound.toolset()).await; - let task = new_terminal - .get_task(&bg.task_id) - .await - .expect("the task table must survive the toolset swap"); - assert!( - !task.completed, - "the task's process must still be running after the swap" - ); - assert_eq!( - orphaned_swap_count(), - orphaned_before, - "the orphaned-backend tripwire must stay 0" - ); - new_terminal.kill_task(&bg.task_id).await; } /// Test factory whose sessions own a PERSISTENT-shell backend (the /// production factory shape). The plain [`TestSessionContextFactory`] @@ -4774,12 +1678,8 @@ pub(crate) mod tests { hook_project_sources: vec![], skills_config: Default::default(), plugin_discovery_config: Default::default(), - hub_config: None, - auth_provider: None, - server_metadata: None, status_config: Default::default(), project_lsp_trusted: true, - require_explicit_toolset: false, confine_fs_to_workspace_root: false, }; WorkspaceHandle::build(config, ephemeral_workspace_home(), false, false) @@ -4830,16 +1730,16 @@ pub(crate) mod tests { "the shell must have entered the subdir: {}", cwd_before.display() ); - let cfg_b = explicit_cfg("read_b"); - let (rebound, outcome) = handle - .rebind_existing_hub_session( - "shell-swap", - Some(cfg_b.clone()), - serde_json::to_value(&cfg_b).ok(), - ) - .await - .expect("session exists"); - assert_eq!(outcome, RebindOutcome::Reresolved); + handle.shared.mcp_tools_snapshot.store(Arc::new(vec![tc( + "GrokBuild:read_file", + Some(ToolKind::Read), + )])); + let rebuilt = handle + .shared + .re_resolve_all_sessions("mcp_snapshot_changed", true) + .await; + assert!(rebuilt >= 1, "the session must be rebuilt"); + let rebound = handle.session("shell-swap").expect("session still exists"); let cwd_after = toolset_terminal(&rebound.toolset()) .await .get_shell_cwd() @@ -5119,12 +2019,8 @@ pub(crate) mod tests { hook_project_sources: vec![], skills_config: Default::default(), plugin_discovery_config: Default::default(), - hub_config: None, - auth_provider: None, - server_metadata: None, status_config: Default::default(), project_lsp_trusted: true, - require_explicit_toolset: false, confine_fs_to_workspace_root: false, }; let home = tempfile::tempdir().unwrap(); @@ -5468,36 +2364,10 @@ pub(crate) mod tests { "default workspace_home must live under the temp dir, got {}", home.display() ); - assert_ne!( - home, - resolve_workspace_home(), - "default construction must NOT use the real $KIGI_WORKSPACE_HOME" - ); - } - /// `server_metadata_typed` defaults cleanly when no metadata is configured. - #[tokio::test] - async fn server_metadata_typed_defaults_without_metadata() { - let handle = make_handle(); - assert_eq!(handle.shared().server_metadata_typed().sandbox_id, None); - } - /// `server_metadata_typed` parses `sandbox_id` from the registration blob. - #[tokio::test] - async fn server_metadata_typed_parses_sandbox_id() { - use crate::session::tool_config::test_support::TestSessionContextFactory; - let factory = Arc::new(TestSessionContextFactory::new()); - let cwd = factory.temp.path().to_path_buf(); - let mut config = WorkspaceHandle::test_config(cwd, factory); - config.server_metadata = - Some(serde_json::json!({ "sandbox_id" : "sb_test123", "mode" : "remote", })); - let handle = WorkspaceHandle::build(config, ephemeral_workspace_home(), false, false) - .expect("handle construction should succeed"); - assert_eq!( - handle - .shared() - .server_metadata_typed() - .sandbox_id - .as_deref(), - Some("sb_test123") + assert!( + home.starts_with(std::env::temp_dir()), + "default construction must use an ephemeral temp home, got {}", + home.display() ); } #[tokio::test] @@ -5524,12 +2394,10 @@ pub(crate) mod tests { }; let factory = handle.shared.session_factory.clone(); let mcp_snapshot = handle.shared.mcp_tools_snapshot.load_full(); - let hub_snapshot = handle.shared.hub_tools_snapshot.load_full(); let (eff, ts, _backend) = resolve_session_toolset( new_parent_baseline, parent.capability_mode(), &mcp_snapshot, - &hub_snapshot, parent.cwd().to_path_buf(), parent.session_env().clone(), "main", @@ -5590,12 +2458,10 @@ pub(crate) mod tests { let main = handle.session("main").expect("main present"); let factory = handle.shared.session_factory.clone(); let mcp_snapshot = handle.shared.mcp_tools_snapshot.load_full(); - let hub_snapshot = handle.shared.hub_tools_snapshot.load_full(); let (eff, ts, _backend) = resolve_session_toolset( marker_config, main.capability_mode(), &mcp_snapshot, - &hub_snapshot, main.cwd().to_path_buf(), main.session_env().clone(), "main", @@ -6010,12 +2876,8 @@ pub(crate) mod tests { hook_project_sources: vec![], skills_config: Default::default(), plugin_discovery_config: Default::default(), - hub_config: None, - auth_provider: None, - server_metadata: None, status_config: Default::default(), project_lsp_trusted: true, - require_explicit_toolset: false, confine_fs_to_workspace_root: false, }; let handle = WorkspaceHandle::new(config).expect("ok"); @@ -6045,12 +2907,8 @@ pub(crate) mod tests { hook_project_sources: vec![HookSourceConfig::Directory(hooks_dir)], skills_config: Default::default(), plugin_discovery_config: Default::default(), - hub_config: None, - auth_provider: None, - server_metadata: None, status_config: Default::default(), project_lsp_trusted: true, - require_explicit_toolset: false, confine_fs_to_workspace_root: false, }; let handle = WorkspaceHandle::new(config).expect("ok"); @@ -6101,12 +2959,8 @@ pub(crate) mod tests { hook_project_sources: vec![], skills_config: Default::default(), plugin_discovery_config: Default::default(), - hub_config: None, - auth_provider: None, - server_metadata: None, status_config: Default::default(), project_lsp_trusted: true, - require_explicit_toolset: false, confine_fs_to_workspace_root: false, }; let handle = WorkspaceHandle::new(config).expect("construction must still succeed"); @@ -6142,12 +2996,8 @@ pub(crate) mod tests { hook_project_sources: vec![HookSourceConfig::SettingsFile(project_settings)], skills_config: Default::default(), plugin_discovery_config: Default::default(), - hub_config: None, - auth_provider: None, - server_metadata: None, status_config: Default::default(), project_lsp_trusted: true, - require_explicit_toolset: false, confine_fs_to_workspace_root: false, }; let handle = WorkspaceHandle::new(config).expect("ok"); @@ -6170,12 +3020,8 @@ pub(crate) mod tests { hook_project_sources: vec![], skills_config: Default::default(), plugin_discovery_config: Default::default(), - hub_config: None, - auth_provider: None, - server_metadata: None, status_config: Default::default(), project_lsp_trusted: true, - require_explicit_toolset: false, confine_fs_to_workspace_root: false, }; let handle = WorkspaceHandle::new(config).expect("must not panic on missing source"); @@ -6202,12 +3048,8 @@ pub(crate) mod tests { hook_project_sources: vec![HookSourceConfig::Directory(empty_dir)], skills_config: Default::default(), plugin_discovery_config: Default::default(), - hub_config: None, - auth_provider: None, - server_metadata: None, status_config: Default::default(), project_lsp_trusted: true, - require_explicit_toolset: false, confine_fs_to_workspace_root: false, }; let handle = WorkspaceHandle::new(config).expect("ok"); @@ -6215,108 +3057,6 @@ pub(crate) mod tests { assert!(handle.hook_load_errors().is_empty()); } #[tokio::test] - async fn hub_tools_snapshot_starts_empty() { - let handle = make_handle(); - assert!(handle.shared().hub_tools_snapshot().is_empty()); - assert!(handle.shared().hub_server().is_none()); - } - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn on_hub_tools_changed_emits_per_session_events() { - let handle = make_handle(); - handle - .fork_session(fork_cfg_with( - "hubA", - CapabilityMode::ReadWrite, - None, - Some("main"), - )) - .await - .expect("hubA ok"); - let mut rx = handle.shared.events.subscribe(); - let hub_tool = tc("hub:remote_exec", None); - let rebuilt = handle.on_hub_tools_changed(vec![hub_tool]); - assert_eq!(rebuilt, 2, "main + 1 subagent"); - let mut got: std::collections::BTreeSet = std::collections::BTreeSet::new(); - for _ in 0..2 { - let ev = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv()) - .await - .expect("event arrives") - .expect("not closed"); - match ev { - WorkspaceEvent::ToolsChanged { session_id } => { - got.insert(session_id); - } - other => panic!("unexpected event: {other:?}"), - } - } - assert_eq!( - got, - ["main".to_string(), "hubA".to_string()] - .into_iter() - .collect::>() - ); - } - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn on_hub_tools_changed_updates_snapshot() { - let handle = make_handle(); - assert!(handle.shared().hub_tools_snapshot().is_empty()); - let hub_tool = tc("hub:remote_exec", None); - handle.on_hub_tools_changed(vec![hub_tool]); - let snapshot = handle.shared().hub_tools_snapshot(); - assert_eq!(snapshot.len(), 1); - assert_eq!(snapshot[0].id, "hub:remote_exec"); - } - #[tokio::test] - async fn connect_hub_noop_when_no_config() { - let handle = make_handle(); - let result = handle.connect_hub().await; - assert!(result.is_ok()); - assert!(handle.shared().hub_server().is_none()); - } - #[test] - fn workspace_shared_auth_provider_uses_workspace_config() { - let temp = tempfile::tempdir().unwrap(); - let service_auth: kigi_computer_hub_sdk::SharedAuthProvider = Arc::new( - kigi_computer_hub_sdk::auth::AuthCredential::bearer("xai-service-token"), - ); - let hub_auth: kigi_computer_hub_sdk::SharedAuthProvider = Arc::new( - kigi_computer_hub_sdk::auth::AuthCredential::bearer("hub-token"), - ); - let hub_cfg = crate::hub::HubConfig { - url: url::Url::parse("ws://127.0.0.1:9/ws").unwrap(), - auth: hub_auth.clone(), - activity_tracker: None, - server_id: Some("server-1".to_string()), - alpha_test_key: None, - allow_insecure_ws: true, - ready_file: None, - diag: None, - }; - let config = WorkspaceConfig::new_for_proxy( - temp.path().to_path_buf(), - Arc::new(TestSessionContextFactory::new()), - hub_cfg, - service_auth.clone(), - None, - Default::default(), - baseline_config(), - ); - let handle = WorkspaceHandle::build(config, ephemeral_workspace_home(), false, false) - .expect("handle construction should succeed"); - let shared_auth = handle - .shared() - .auth_provider() - .expect("WorkspaceConfig auth provider must populate WorkspaceShared"); - assert_eq!(shared_auth.current(), service_auth.current()); - assert_ne!(shared_auth.current(), hub_auth.current()); - } - #[tokio::test] - async fn shutdown_hub_noop_when_not_connected() { - let handle = make_handle(); - handle.shutdown_hub().await; - assert!(handle.shared().hub_server().is_none()); - } - #[tokio::test] async fn codebase_index_forwarder_abort_releases_shared() { let handle = make_handle(); tokio::task::yield_now().await; @@ -6708,481 +3448,6 @@ pub(crate) mod tests { "child must inherit the parent's stream_tool_progress flag" ); } - /// Build the resolver exactly the way `connect_hub` does: session catalog - /// handlers + the workspace RPC handler. - fn bind_resolver_fixture( - handle: &WorkspaceHandle, - ) -> kigi_computer_hub_sdk::SessionHandlerResolver { - let catalog_toolset = handle.session("main").expect("main session").toolset(); - let mut catalog = build_session_routed_handlers(&catalog_toolset, handle); - let rpc_handler: Arc = - Arc::new(crate::hub_server::WorkspaceRpcHandler::new(handle.clone())); - let rpc_tool_id = rpc_handler.tool_id(); - catalog.push(rpc_handler); - handle.session_bind_resolver(Arc::new(catalog), rpc_tool_id) - } - fn handler_names(resolved: &kigi_computer_hub_sdk::ResolvedSessionHandlers) -> Vec { - resolved - .handlers - .iter() - .map(|h| h.tool_id().as_str().to_owned()) - .collect() - } - /// Strict mode, preset-only bind: the full resolver path fails closed — - /// RPC-only advertise + a `missing_tool_config` reason in the bind report. - #[tokio::test] - async fn strict_bind_without_explicit_toolset_fails_closed_end_to_end() { - let handle = make_strict_handle(); - let resolver = bind_resolver_fixture(&handle); - let resolved = resolver( - kigi_tool_protocol::SessionId::new("bind-e2e-strict").unwrap(), - Some(serde_json::json!( - { "metadata" : { "preset" : "grok-computer", "capability_mode" : - "all" }, } - )), - ) - .await - .expect("bind must succeed"); - assert_eq!( - handler_names(&resolved), - vec![crate::hub_ids::WORKSPACE_RPC_TOOL_ID.to_owned()], - "must advertise the RPC handler only" - ); - let reason = resolved.resolve_error.expect("resolve_error must be set"); - assert!( - reason.starts_with("missing_tool_config:"), - "reason must name the fail-closed cause: {reason}" - ); - assert!( - reason.contains(kigi_version::VERSION), - "reason must carry the server version: {reason}" - ); - } - #[tokio::test] - async fn strict_rpc_only_bind_fails_closed_with_resolve_error_end_to_end() { - let handle = make_strict_handle(); - let resolver = bind_resolver_fixture(&handle); - let resolved = resolver( - kigi_tool_protocol::SessionId::new("bind-e2e-rpc-only").unwrap(), - Some(serde_json::json!( - { "metadata" : { "capability_mode" : "read_write", "rpc_only" : - true, "system_notifications" : true, }, } - )), - ) - .await - .expect("bind must succeed"); - assert_eq!( - handler_names(&resolved), - vec![crate::hub_ids::WORKSPACE_RPC_TOOL_ID.to_owned()], - ); - let reason = resolved.resolve_error.expect("resolve_error must be set"); - assert!(reason.starts_with("missing_tool_config:"), "{reason}"); - } - /// Strict mode, explicit `tools`: resolves and advertises the configured - /// tool with no resolve_error. - #[tokio::test] - async fn strict_bind_with_explicit_toolset_serves_it_end_to_end() { - let handle = make_strict_handle(); - let resolver = bind_resolver_fixture(&handle); - let resolved = resolver( - kigi_tool_protocol::SessionId::new("bind-e2e-tools").unwrap(), - Some(serde_json::json!( - { "metadata" : { "tools" : [{ "id" : "GrokBuild:read_file" }] }, - } - )), - ) - .await - .expect("bind must succeed"); - let names = handler_names(&resolved); - assert!( - names.iter().any(|n| n == "read_file"), - "configured tool must be advertised: {names:?}" - ); - assert_eq!(resolved.resolve_error, None); - assert!(resolved.unserved_tool_ids.is_empty()); - } - /// Lax mode (CLI/local embedders), metadata-less bind: falls back to the - /// default catalog with no resolve_error. - #[tokio::test] - async fn lax_bind_without_metadata_uses_default_catalog_end_to_end() { - let handle = make_handle(); - let resolver = bind_resolver_fixture(&handle); - let resolved = resolver( - kigi_tool_protocol::SessionId::new("bind-e2e-lax").unwrap(), - None, - ) - .await - .expect("bind must succeed"); - let names = handler_names(&resolved); - assert!( - names.iter().any(|n| n == "read_file") && names.iter().any(|n| n == "grep"), - "default catalog must be advertised: {names:?}" - ); - assert_eq!(resolved.resolve_error, None); - } - /// A rebind whose explicit config is REJECTED (invalid entry) keeps the - /// fail-closed reason even though the healthy session's previous toolset - /// is reused — the client must learn its new config did not take effect. - #[tokio::test] - async fn rejected_rebind_config_keeps_resolve_error_end_to_end() { - let handle = make_strict_handle(); - let resolver = bind_resolver_fixture(&handle); - let sid = kigi_tool_protocol::SessionId::new("bind-e2e-rejected").unwrap(); - let first = resolver( - sid.clone(), - Some(serde_json::json!( - { "metadata" : { "tools" : [{ "id" : "GrokBuild:read_file" }] }, - } - )), - ) - .await - .expect("healthy bind"); - assert_eq!(first.resolve_error, None); - let second = resolver( - sid, - Some(serde_json::json!( - { "metadata" : { "tools" : [{ "id" : "GrokBuild:read_file", - "params_json" : "{not json" }] }, } - )), - ) - .await - .expect("rejected rebind still advertises the previous toolset"); - assert!( - handler_names(&second).iter().any(|n| n == "read_file"), - "previous toolset must still be served" - ); - let reason = second - .resolve_error - .expect("rejected config must keep the fail-closed reason"); - assert!(reason.starts_with("invalid_tool_config:"), "{reason}"); - } - /// An explicit EMPTY toolset (RPC-only clients, e.g. deploy binds) must - /// reuse an existing session unchanged — never swap its tools away. - #[tokio::test] - async fn explicit_empty_toolset_rebind_never_swaps_session_tools() { - let handle = make_strict_handle(); - let resolver = bind_resolver_fixture(&handle); - let sid = kigi_tool_protocol::SessionId::new("bind-e2e-rpc-only").unwrap(); - let first = resolver( - sid.clone(), - Some(serde_json::json!( - { "metadata" : { "tools" : [{ "id" : "GrokBuild:read_file" }] }, - } - )), - ) - .await - .expect("agent bind"); - assert!(handler_names(&first).iter().any(|n| n == "read_file")); - let rpc_bind = resolver( - sid, - Some(serde_json::json!( - { "metadata" : { "tool_config" : { "tools" : [] } }, } - )), - ) - .await - .expect("rpc-only rebind"); - assert!( - handler_names(&rpc_bind).iter().any(|n| n == "read_file"), - "agent session tools must survive an RPC-only rebind" - ); - assert_eq!(rpc_bind.resolve_error, None); - } - /// Rebind heal end-to-end: a strict fail-closed bind leaves the session - /// empty; a corrected rebind with explicit tools rebuilds and advertises - /// them with the report cleared. - #[tokio::test] - async fn strict_rebind_with_corrected_toolset_heals_end_to_end() { - let handle = make_strict_handle(); - let resolver = bind_resolver_fixture(&handle); - let sid = kigi_tool_protocol::SessionId::new("bind-e2e-heal").unwrap(); - let first = resolver( - sid.clone(), - Some(serde_json::json!({ "metadata" : { "preset" : "grok-computer" } })), - ) - .await - .expect("fail-closed bind still succeeds with an RPC-only advertise"); - assert!(first.resolve_error.is_some(), "first bind must fail closed"); - let second = resolver( - sid, - Some(serde_json::json!( - { "metadata" : { "tools" : [{ "id" : "GrokBuild:read_file" }] }, - } - )), - ) - .await - .expect("bind must succeed"); - let names = handler_names(&second); - assert!( - names.iter().any(|n| n == "read_file"), - "corrected rebind must advertise the explicit toolset: {names:?}" - ); - assert_eq!( - second.resolve_error, None, - "healed rebind must not carry the stale fail-closed reason" - ); - } - /// Owner bind: capability `all` + explicit toolset (strict servers fail - /// closed otherwise). - fn owner_full_bind_metadata() -> serde_json::Value { - serde_json::json!( - { "metadata" : { "capability_mode" : "all", "tools" : [{ "id" : - "GrokBuild:read_file" }, { "id" : "GrokBuild:search_replace" }, { "id" : - "GrokBuild:grep" }, { "id" : "GrokBuild:list_dir" },], }, } - ) - } - const OWNER_TOOLS: [&str; 4] = ["read_file", "search_replace", "grep", "list_dir"]; - #[track_caller] - fn assert_advertises_owner_tools(names: &[String], context: &str) { - for tool in OWNER_TOOLS { - assert!( - names.iter().any(|n| n == tool), - "{context}: owner tool `{tool}` missing from advertised set {names:?}" - ); - } - } - /// Consumer-shaped rebinds against a live owner session must `Reuse` it - /// unchanged — never shrink its toolset or narrow its frozen capability. - #[tokio::test] - async fn owner_toolset_survives_concurrent_consumer_shaped_rebinds() { - let handle = make_strict_handle(); - let resolver = bind_resolver_fixture(&handle); - let sid = kigi_tool_protocol::SessionId::new("bind-e2e-consumer-storm").unwrap(); - let owner = resolver(sid.clone(), Some(owner_full_bind_metadata())) - .await - .expect("owner bind"); - assert_advertises_owner_tools(&handler_names(&owner), "owner bind"); - assert_eq!(owner.resolve_error, None); - let consumer_shapes: Vec> = vec![ - Some( - serde_json::json!({ "metadata" : { "capability_mode" : "read_only" } - }), - ), - Some( - serde_json::json!({ "metadata" : { "capability_mode" : "read_write" - } }), - ), - None, - Some(serde_json::json!({ "metadata" : { "tool_config" : { - "tools" : [] } } })), - ]; - let storm = futures::future::join_all( - consumer_shapes - .iter() - .cycle() - .take(12) - .cloned() - .map(|metadata| resolver(sid.clone(), metadata)), - ) - .await; - for (i, result) in storm.into_iter().enumerate() { - let resolved = result.expect("consumer-shaped rebind must not error"); - assert_advertises_owner_tools( - &handler_names(&resolved), - &format!("consumer-shaped rebind #{i}"), - ); - assert_eq!( - resolved.resolve_error, None, - "reuse against a healthy owner session must not surface a resolve error" - ); - } - let session = handle - .session("bind-e2e-consumer-storm") - .expect("owner session survives the storm"); - assert_eq!( - session.capability_mode(), - CapabilityMode::All, - "consumer-shaped rebinds must never narrow the owner's frozen capability" - ); - assert_advertises_owner_tools( - &session - .toolset() - .tool_definitions() - .into_iter() - .map(|d| d.function.name) - .collect::>(), - "post-storm session toolset", - ); - } - /// On a fresh workspace-server the FIRST bind freezes `capability_mode`: consumer-shaped - /// first binds strand the session narrow (why consumers never bind); owner-first is whole. - #[tokio::test] - async fn restored_server_first_bind_ordering_decides_capability_and_toolset() { - let handle = make_strict_handle(); - let resolver = bind_resolver_fixture(&handle); - let sid = kigi_tool_protocol::SessionId::new("bind-e2e-restore-read-first").unwrap(); - let read_first = resolver( - sid.clone(), - Some(serde_json::json!( - { "metadata" : { "capability_mode" : "read_only" } } - )), - ) - .await - .expect("consumer-shaped bind resolves"); - assert_eq!( - handler_names(&read_first), - vec![crate::hub_ids::WORKSPACE_RPC_TOOL_ID.to_owned()], - "strict fail-closed create advertises the RPC handler only" - ); - let agent = resolver(sid, Some(owner_full_bind_metadata())) - .await - .expect("agent bind resolves"); - let names = handler_names(&agent); - assert!( - names.iter().any(|n| n == "read_file"), - "agent bind heals the read-class toolset: {names:?}" - ); - assert!( - !names.iter().any(|n| n == "search_replace"), - "frozen read_only capability keeps filtering Edit-class tools — \ - the incident's shrunken toolset: {names:?}" - ); - let session = handle - .session("bind-e2e-restore-read-first") - .expect("session exists"); - assert_eq!( - session.capability_mode(), - CapabilityMode::ReadOnly, - "the consumer-shaped first bind froze the capability for good" - ); - let handle = make_strict_handle(); - let resolver = bind_resolver_fixture(&handle); - let sid = kigi_tool_protocol::SessionId::new("bind-e2e-restore-write-first").unwrap(); - resolver( - sid.clone(), - Some(serde_json::json!( - { "metadata" : { "capability_mode" : "read_write" } } - )), - ) - .await - .expect("consumer-shaped bind resolves"); - resolver(sid, Some(owner_full_bind_metadata())) - .await - .expect("agent bind resolves"); - let session = handle - .session("bind-e2e-restore-write-first") - .expect("session exists"); - assert_eq!( - session.capability_mode(), - CapabilityMode::ReadWrite, - "the agent's `all` must not take on a session a deploy/write-shaped \ - bind created first — this narrower freeze is why deploy and fs \ - writes are consumers now" - ); - let handle = make_strict_handle(); - let resolver = bind_resolver_fixture(&handle); - let sid = kigi_tool_protocol::SessionId::new("bind-e2e-restore-owner-first").unwrap(); - let owner = resolver(sid, Some(owner_full_bind_metadata())) - .await - .expect("owner bind resolves"); - assert_advertises_owner_tools(&handler_names(&owner), "owner-first bind"); - assert_eq!(owner.resolve_error, None); - let session = handle - .session("bind-e2e-restore-owner-first") - .expect("session exists"); - assert_eq!( - session.capability_mode(), - CapabilityMode::All, - "owner-first ordering yields the full capability the agent declared" - ); - } - /// Isolation matrix #1–#3 through the REAL `session.bind` resolver (the - /// closure `connect_hub` installs — the exact path both a soft rebind and - /// an SDK dead-loop FULL rebind re-run): with a live background task, - /// an identical rebind (`Reused`) and a changed-explicit-toolset rebind - /// (`Reresolved`, driven with no in-flight tool calls) both keep the - /// session-owned backend (`Arc::ptr_eq`) and the running task, while the - /// changed rebind swaps the advertised handler set. - /// - /// The remaining matrix-#3 sub-asserts live beside the swap tests above: - /// persistent-shell cwd preservation - /// (`reresolved_swap_preserves_persistent_shell_cwd`) and the - /// snapshot-driven rebuild with a live task - /// (`re_resolve_all_sessions_preserves_session_terminal_backend`). - #[tokio::test] - async fn bind_flow_rebinds_keep_backend_and_task_alive_end_to_end() { - let orphaned_before = orphaned_swap_count(); - let handle = make_handle(); - let resolver = bind_resolver_fixture(&handle); - let sid = kigi_tool_protocol::SessionId::new("bind-e2e-bg").unwrap(); - let bg_metadata = serde_json::json!( - { "metadata" : { "tools" : [{ "id" : "GrokBuild:read_file" }, { "id" : - "GrokBuild:run_terminal_cmd" }, { "id" : "GrokBuild:get_task_output" }, { - "id" : "GrokBuild:kill_task" },] }, } - ); - let first = resolver(sid.clone(), Some(bg_metadata.clone())) - .await - .expect("owner bind"); - assert!( - handler_names(&first) - .iter() - .any(|n| n == "run_terminal_cmd"), - "owner bind must serve the execute tool" - ); - let session = handle.session("bind-e2e-bg").expect("session created"); - let backend = session.terminal_backend().clone(); - let out_dir = tempfile::tempdir().expect("temp dir"); - let bg = start_background_sleep(&session, out_dir.path(), "bind-e2e-bg-task").await; - let reused = resolver(sid.clone(), Some(bg_metadata)) - .await - .expect("identical rebind"); - assert!( - handler_names(&reused) - .iter() - .any(|n| n == "run_terminal_cmd"), - "a reused rebind keeps advertising the existing toolset" - ); - let session = handle.session("bind-e2e-bg").expect("session kept"); - assert!( - Arc::ptr_eq(&backend, session.terminal_backend()), - "an identical rebind must keep the session-owned backend" - ); - assert!( - !backend - .get_task(&bg.task_id) - .await - .expect("task listed across the reused rebind") - .completed, - "the task must still be running after the reused rebind" - ); - let swapped = resolver( - sid, - Some(serde_json::json!( - { "metadata" : { "tools" : [{ "id" : "GrokBuild:read_file" }] }, - } - )), - ) - .await - .expect("changed-toolset rebind"); - let names = handler_names(&swapped); - assert!( - names.iter().any(|n| n == "read_file") - && !names.iter().any(|n| n == "run_terminal_cmd"), - "the changed rebind must advertise the NEW toolset only: {names:?}" - ); - let session = handle.session("bind-e2e-bg").expect("session kept"); - assert!( - Arc::ptr_eq(&backend, session.terminal_backend()), - "a toolset-swapping rebind must keep the session-owned backend" - ); - assert!( - Arc::ptr_eq(&backend, &toolset_terminal(&session.toolset()).await), - "the swapped-in toolset must reference the session-owned backend" - ); - assert!( - !backend - .get_task(&bg.task_id) - .await - .expect("task table must survive the toolset swap") - .completed, - "the task's process must still be running after the swap" - ); - assert_eq!( - orphaned_swap_count(), - orphaned_before, - "the orphaned-backend tripwire must stay 0" - ); - backend.kill_task(&bg.task_id).await; - } /// Dropping and rebinding a session with the same ID surfaces the /// new `viewer_ctx` (kill-switch for mid-session staleness). #[tokio::test] @@ -7364,154 +3629,4 @@ pub(crate) mod tests { serde_json::json!({ "recovery" : false }) ); } - fn bundled_dir_fixture(subdirs: &[&str]) -> tempfile::TempDir { - let tmp = tempfile::tempdir().expect("tempdir"); - for name in subdirs { - std::fs::create_dir(tmp.path().join(name)).expect("create subdir"); - } - std::fs::write(tmp.path().join("BUILD.bazel"), b"").expect("create file"); - tmp - } - #[test] - fn bundled_allowlist_blank_ignores_nothing() { - let tmp = bundled_dir_fixture(&["bundled__pdf", "bundled__xlsx"]); - let dir = tmp.path().to_string_lossy().into_owned(); - for allowlist in [None, Some(""), Some(" "), Some(" , ,")] { - assert_eq!( - bundled_allowlist_ignore_dirs(&dir, allowlist), - Vec::::new(), - "allowlist {allowlist:?} must produce no ignore entries" - ); - } - } - #[test] - fn bundled_allowlist_ignores_complement() { - let tmp = bundled_dir_fixture(&["bundled__pdf", "bundled__xlsx", "bundled__docx"]); - let dir = tmp.path().to_string_lossy().into_owned(); - let got = bundled_allowlist_ignore_dirs(&dir, Some("xlsx, pdf")); - let want = vec![ - tmp.path() - .join("bundled__docx") - .to_string_lossy() - .into_owned(), - ]; - assert_eq!(got, want); - } - #[test] - fn bundled_allowlist_strips_bundled_prefix() { - let tmp = bundled_dir_fixture(&["bundled__pdf", "xlsx", "bundled__skip"]); - let dir = tmp.path().to_string_lossy().into_owned(); - let got = bundled_allowlist_ignore_dirs(&dir, Some("bundled__pdf,bundled__xlsx")); - let want = vec![ - tmp.path() - .join("bundled__skip") - .to_string_lossy() - .into_owned(), - ]; - assert_eq!(got, want); - } - #[test] - fn bundled_allowlist_unreadable_dir_fails_closed() { - let got = bundled_allowlist_ignore_dirs("/nonexistent/bundled-skills", Some("pdf")); - assert_eq!(got, vec!["/nonexistent/bundled-skills".to_string()]); - } - /// Unique skill names: discovery also reads the dev machine's `~/.kigi`. - #[tokio::test] - async fn bundled_allowlist_filters_discovery() { - let tmp = tempfile::tempdir().expect("tempdir"); - for name in ["allowlist-e2e-kept", "allowlist-e2e-blocked"] { - let skill_dir = tmp.path().join(format!("bundled__{name}")); - std::fs::create_dir(&skill_dir).expect("create subdir"); - std::fs::write( - skill_dir.join("SKILL.md"), - format!("---\nname: {name}\ndescription: test\n---\nbody"), - ) - .expect("write SKILL.md"); - } - let dir = tmp.path().to_string_lossy().into_owned(); - let cwd = tempfile::tempdir().expect("tempdir"); - let mut config = crate::discovery::SkillsConfig { - bundled_skill_dirs: vec![dir.clone()], - ..Default::default() - }; - config.ignore.extend(bundled_allowlist_ignore_dirs( - &dir, - Some("allowlist-e2e-kept"), - )); - let skills = crate::discovery::discover_skills(cwd.path(), &config).await; - let names: Vec<&str> = skills - .iter() - .filter_map(|s| s["name"].as_str()) - .filter(|n| n.starts_with("allowlist-e2e-")) - .collect(); - assert_eq!( - names, - vec!["allowlist-e2e-kept"], - "only the allowlisted skill survives" - ); - } - #[test] - fn classify_drain_outcome_covers_all_arms() { - assert_eq!(classify_drain_outcome(false), DrainOutcome::Partial); - assert_eq!(classify_drain_outcome(true), DrainOutcome::Full); - } - #[test] - fn drain_reason_and_outcome_labels_are_stable() { - assert_eq!(DrainReason::Sigterm.as_str(), "sigterm"); - assert_eq!(DrainReason::Evict.as_str(), "evict"); - assert_eq!(DrainOutcome::Full.as_str(), "full"); - assert_eq!(DrainOutcome::Partial.as_str(), "partial"); - } - #[test] - fn grace_budget_from_raw_parses_and_falls_back() { - let d = |ms| std::time::Duration::from_millis(ms); - assert_eq!(grace_budget_from_raw(None), d(DEFAULT_TERMINATION_GRACE_MS)); - assert_eq!(grace_budget_from_raw(Some("120000".into())), d(120_000)); - assert_eq!(grace_budget_from_raw(Some(" 90000 ".into())), d(90_000)); - assert_eq!( - grace_budget_from_raw(Some("0".into())), - d(DEFAULT_TERMINATION_GRACE_MS) - ); - assert_eq!( - grace_budget_from_raw(Some("nonsense".into())), - d(DEFAULT_TERMINATION_GRACE_MS) - ); - } - #[test] - fn write_draining_marker_writes_count_atomically() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("workspace-server.draining"); - write_draining_marker(&path, 5); - assert_eq!(std::fs::read_to_string(&path).unwrap(), "5"); - let leftover_tmp = std::fs::read_dir(dir.path()) - .unwrap() - .filter_map(|e| e.ok()) - .any(|e| e.file_name().to_string_lossy().ends_with(".draining.tmp")); - assert!(!leftover_tmp, "temp file must be renamed away"); - write_draining_marker(&path, 0); - assert_eq!(std::fs::read_to_string(&path).unwrap(), "0"); - } - #[tokio::test] - async fn two_phase_drain_no_queue_marks_draining_and_returns_zero() { - let handle = make_handle(); - let tracker = handle.activity_tracker().clone(); - assert!(!tracker.is_draining()); - let unfinished = handle - .two_phase_drain(std::time::Duration::from_millis(300), DrainReason::Sigterm) - .await; - assert_eq!(unfinished, 0, "no queue → nothing pending to lose"); - assert!( - tracker.is_draining(), - "drain must mark the tracker draining" - ); - let snap = tracker.snapshot(); - assert_eq!( - snap.status, - kigi_tool_protocol::ToolServerLifecycleStatus::Draining - ); - assert!( - snap.drain_started_ms.is_some(), - "drain_started_ms must be stamped at drain start" - ); - } } diff --git a/crates/codegen/kigi-workspace/src/lib.rs b/crates/codegen/kigi-workspace/src/lib.rs index 0a56b03..d28856e 100644 --- a/crates/codegen/kigi-workspace/src/lib.rs +++ b/crates/codegen/kigi-workspace/src/lib.rs @@ -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") ); } } diff --git a/crates/codegen/kigi-workspace/src/permission/manager.rs b/crates/codegen/kigi-workspace/src/permission/manager.rs index 9ebd1cf..89fbc13 100644 --- a/crates/codegen/kigi-workspace/src/permission/manager.rs +++ b/crates/codegen/kigi-workspace/src/permission/manager.rs @@ -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, initial_yolo: bool, client_identifier: Option, -) -> (PermissionHandle, mpsc::UnboundedReceiver) { - 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, - deny_read_globs: Vec, - web_fetch_allowed_domains: Vec, - initial_yolo: bool, - client_identifier: Option, - // 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>, ) -> (PermissionHandle, mpsc::UnboundedReceiver) { // 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, remember_tool_approvals: bool, yolo_pin: Option<&'static str>, - hub_permission: Option>, ) -> (PermissionHandle, mpsc::UnboundedReceiver) { let (tx, mut rx) = mpsc::unbounded_channel::(); let (event_tx, event_rx) = mpsc::unbounded_channel::(); @@ -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, - ) -> (PermissionHandle, mpsc::UnboundedReceiver) { - 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>, - } - - #[async_trait::async_trait] - impl crate::permission::PermissionHookTransport for FakeHubTransport { - async fn request_permission( - &self, - payload: serde_json::Value, - ) -> Result { - self.seen.lock().unwrap().push(payload); - Ok(self.reply.clone()) - } - } - - fn fake_hub(reply: serde_json::Value) -> Arc { - 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. diff --git a/crates/codegen/kigi-workspace/src/permission/mod.rs b/crates/codegen/kigi-workspace/src/permission/mod.rs index 8f7ae75..55b6b83 100644 --- a/crates/codegen/kigi-workspace/src/permission/mod.rs +++ b/crates/codegen/kigi-workspace/src/permission/mod.rs @@ -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, diff --git a/crates/codegen/kigi-workspace/src/permission/prompter.rs b/crates/codegen/kigi-workspace/src/permission/prompter.rs index 55ebc74..89bdc46 100644 --- a/crates/codegen/kigi-workspace/src/permission/prompter.rs +++ b/crates/codegen/kigi-workspace/src/permission/prompter.rs @@ -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>, /// 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>, - ) -> 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()) } } }; diff --git a/crates/codegen/kigi-workspace/src/session/mod.rs b/crates/codegen/kigi-workspace/src/session/mod.rs index 9c847f4..5626977 100644 --- a/crates/codegen/kigi-workspace/src/session/mod.rs +++ b/crates/codegen/kigi-workspace/src/session/mod.rs @@ -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, /// 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>, - /// MCP bridges kept alive for the session lifetime. - pub(crate) mcp_bridges: tokio::sync::Mutex>, - /// Qualified tool IDs registered on the server for this session's MCP tools. - pub(crate) mcp_tool_ids: tokio::sync::Mutex>, /// 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, @@ -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>, /// 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>, - /// Remote-origin tool configs (consumer direction), updated by the - /// notification listener. - pub(crate) hub_tools_snapshot: arc_swap::ArcSwap>, - /// Server config stashed at construction time for deferred connect. - pub(crate) hub_config: Option, - /// Auth provider for xAI service calls. - pub(crate) auth_provider: Option, - /// 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>, /// 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>, - pub(crate) local_registry: kigi_computer_hub_sdk::LocalRegistry, + pub(crate) local_registry: kigi_tool_runtime::LocalRegistry, pub(crate) activity_tracker: std::sync::Arc, /// 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, /// 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>>, - 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 { - 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> { 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 { - 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 { - 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> { - 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, - ) -> Option { - 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 { &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)> = { 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)) => { diff --git a/crates/codegen/kigi-workspace/src/session/tool_config.rs b/crates/codegen/kigi-workspace/src/session/tool_config.rs index f645c04..60bc039 100644 --- a/crates/codegen/kigi-workspace/src/session/tool_config.rs +++ b/crates/codegen/kigi-workspace/src/session/tool_config.rs @@ -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>, session_id: &str, factory: &dyn SessionContextFactory, - local_registry: Option, + local_registry: Option, lsp: Option>, viewer_ctx: Option, notification_handle: Option, @@ -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>, session_id: &str, factory: &dyn SessionContextFactory, - local_registry: Option, + local_registry: Option, lsp: Option>, viewer_ctx: Option, notification_handle: Option, @@ -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 = 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 -/// `/sessions//`; left unset, `state_path` stays empty -/// (legacy behavior). +/// [`SessionContextFactory`] for workspace sessions. /// /// [`SessionContext::session_folder`] is `/tmp/sessions//` /// (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, - api_base_url: Option, - /// 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, -} -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 - } - /// `/sessions//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//` 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>, backend: Arc, ) -> 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; 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 { - 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 = baseline.tools.iter().map(|t| t.id.clone()).collect(); - let filtered = merge_and_filter(&baseline, &[], &[], CapabilityMode::ReadWrite, "test"); - let filtered_ids: Vec = 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 - /// `/sessions//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 `/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", diff --git a/crates/codegen/kigi-workspace/src/workspace_ops.rs b/crates/codegen/kigi-workspace/src/workspace_ops.rs index 1595cf0..c9ba29c 100644 --- a/crates/codegen/kigi-workspace/src/workspace_ops.rs +++ b/crates/codegen/kigi-workspace/src/workspace_ops.rs @@ -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 { 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 { { - 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 { 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 { 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 { 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 { - 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 { - ws.put_files(self.files.clone()).await - } -} -#[async_trait] -impl WorkspaceOp for GetFilesReq { - async fn execute( - &self, - ws: &WorkspaceHandle, - _session_id: Option<&str>, - ) -> WorkspaceResult { - ws.get_files(self.files.clone()).await - } -} -#[async_trait] -impl WorkspaceOp for ClientFsListReq { - async fn execute( - &self, - ws: &WorkspaceHandle, - _session_id: Option<&str>, - ) -> WorkspaceResult { - 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 { - 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 { - 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, ©_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 { 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 { 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 { 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 { 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) -> Self { - Self::Proxy { - client: WorkspaceClient::new((*harness).clone()), - } - } - /// Construct a proxy-mode ops handle sharing a pre-created connected - /// flag. The same `Arc` should be wired into the harness - /// builder's `on_reconnect` callback so reconnects reset the flag. - pub fn proxy_with_connected(harness: Arc, connected: Arc) -> 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, viewer_ctx: Option, ) -> 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 { - 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(&self, req: &R) -> WorkspaceResult { - 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 = - 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( &self, op: &Op, session_id: Option<&str>, ) -> WorkspaceResult { - 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 { - 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 { - 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> { - self.rpc(&GetRewindPointsReq { - session_id: session_id.to_owned(), - }) - .await - } - pub async fn rewind_to( - &self, - session_id: &str, - target_prompt_index: usize, - ) -> WorkspaceResult { - self.rpc(&RewindToReq { - session_id: session_id.to_owned(), - target_prompt_index, - }) - .await - } - pub async fn put_files(&self, req: PutFilesReq) -> WorkspaceResult { - self.dispatch(&req, None).await - } - pub async fn get_files(&self, req: GetFilesReq) -> WorkspaceResult { - 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 { - 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::(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()); diff --git a/crates/common/kigi-tool-protocol/Cargo.toml b/crates/common/kigi-tool-protocol/Cargo.toml index 5ce4697..a33302b 100644 --- a/crates/common/kigi-tool-protocol/Cargo.toml +++ b/crates/common/kigi-tool-protocol/Cargo.toml @@ -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 } diff --git a/crates/common/kigi-tool-protocol/src/lib.rs b/crates/common/kigi-tool-protocol/src/lib.rs index 3061b1a..f8d22e1 100644 --- a/crates/common/kigi-tool-protocol/src/lib.rs +++ b/crates/common/kigi-tool-protocol/src/lib.rs @@ -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 diff --git a/crates/common/kigi-tool-runtime/Cargo.toml b/crates/common/kigi-tool-runtime/Cargo.toml index ddbab26..f4557f9 100644 --- a/crates/common/kigi-tool-runtime/Cargo.toml +++ b/crates/common/kigi-tool-runtime/Cargo.toml @@ -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 } diff --git a/crates/common/kigi-tool-runtime/src/context.rs b/crates/common/kigi-tool-runtime/src/context.rs index 8eb196a..3040851 100644 --- a/crates/common/kigi-tool-runtime/src/context.rs +++ b/crates/common/kigi-tool-runtime/src/context.rs @@ -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, - /// 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, - /// 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, - #[serde( - default, - deserialize_with = "ok_or_default", - skip_serializing_if = "Option::is_none" - )] - pub viewer_ctx: Option, - /// 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, - /// 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, - #[serde( - default, - deserialize_with = "ok_or_default", - skip_serializing_if = "Option::is_none" - )] - pub manifest_hash: Option, - /// 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, - #[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 -where - D: serde::Deserializer<'de>, - T: serde::de::DeserializeOwned + Default, -{ - let value = ::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()); - } -} diff --git a/crates/common/kigi-tool-runtime/src/lib.rs b/crates/common/kigi-tool-runtime/src/lib.rs index 49690d8..3158eb4 100644 --- a/crates/common/kigi-tool-runtime/src/lib.rs +++ b/crates/common/kigi-tool-runtime/src/lib.rs @@ -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, diff --git a/crates/common/kigi-tool-runtime/src/local_registry.rs b/crates/common/kigi-tool-runtime/src/local_registry.rs new file mode 100644 index 0000000..55c85a9 --- /dev/null +++ b/crates/common/kigi-tool-runtime/src/local_registry.rs @@ -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` 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>>, +} + +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(&self, tool: T) -> Option + where + T: Tool + 'static, + { + self.register_arc(Arc::new(tool)) + } + + /// Register a typed [`Tool`] implementation already wrapped in `Arc`. + pub fn register_arc(&self, tool: Arc) -> Option + 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`) 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 { + 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 { + 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 { + self.entries + .read() + .values() + .filter(|handle| handle.should_list(ctx)) + .map(|handle| handle.description(ctx)) + .collect() + } +} diff --git a/crates/common/kigi-tool-types/Cargo.toml b/crates/common/kigi-tool-types/Cargo.toml index 3512689..dee0535 100644 --- a/crates/common/kigi-tool-types/Cargo.toml +++ b/crates/common/kigi-tool-types/Cargo.toml @@ -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 diff --git a/crates/common/kigi-tracing/Cargo.toml b/crates/common/kigi-tracing/Cargo.toml deleted file mode 100644 index a952455..0000000 --- a/crates/common/kigi-tracing/Cargo.toml +++ /dev/null @@ -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 diff --git a/crates/common/kigi-tracing/src/dispatch.rs b/crates/common/kigi-tracing/src/dispatch.rs deleted file mode 100644 index b9fd838..0000000 --- a/crates/common/kigi-tracing/src/dispatch.rs +++ /dev/null @@ -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::()) -} - -#[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()); - } -} diff --git a/crates/common/kigi-tracing/src/fastrace.rs b/crates/common/kigi-tracing/src/fastrace.rs deleted file mode 100644 index 0092cdc..0000000 --- a/crates/common/kigi-tracing/src/fastrace.rs +++ /dev/null @@ -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, -) -> 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 { - 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>, 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; - -pub fn fastrace_channel( - channel: tonic::transport::Channel, -) -> fastrace_tonic::FastraceClientService { - 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 { - req.headers_mut() - .extend(fastrace_reqwest::traceparent_headers()); - next.run(req, extensions).await - } -} diff --git a/crates/common/kigi-tracing/src/grpc_client.rs b/crates/common/kigi-tracing/src/grpc_client.rs deleted file mode 100644 index ef8052a..0000000 --- a/crates/common/kigi-tracing/src/grpc_client.rs +++ /dev/null @@ -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, - SharedClassifier, - 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 MakeSpan for MakeClientSpan { - fn make_span(&mut self, request: &Request) -> 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 Layer for InjectTraceContextLayer { - type Service = InjectTraceContextService; - - fn layer(&self, inner: S) -> Self::Service { - InjectTraceContextService { inner } - } -} - -#[derive(Clone, Debug)] -pub struct InjectTraceContextService { - inner: S, -} - -impl Service> for InjectTraceContextService -where - S: Service>, -{ - type Response = S::Response; - type Error = S::Error; - type Future = S::Future; - - fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { - self.inner.poll_ready(cx) - } - - fn call(&mut self, mut req: Request) -> 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( - mut request: tonic::Request, -) -> Result, 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; - - #[derive(Clone)] - struct CaptureService { - seen: Arc>>, - 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 Service> for CaptureService { - type Response = http::Response; - type Error = Infallible; - type Future = std::future::Ready>; - - fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } - - fn call(&mut self, req: Request) -> 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 { - 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>> = 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>> = 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); - } -} diff --git a/crates/common/kigi-tracing/src/http_client.rs b/crates/common/kigi-tracing/src/http_client.rs deleted file mode 100644 index ca66e1d..0000000 --- a/crates/common/kigi-tracing/src/http_client.rs +++ /dev/null @@ -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 { - 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 { - 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() - ); - } -} diff --git a/crates/common/kigi-tracing/src/lib.rs b/crates/common/kigi-tracing/src/lib.rs deleted file mode 100644 index 9f7a1d0..0000000 --- a/crates/common/kigi-tracing/src/lib.rs +++ /dev/null @@ -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::*; diff --git a/crates/common/kigi-tracing/src/testing.rs b/crates/common/kigi-tracing/src/testing.rs deleted file mode 100644 index bf327da..0000000 --- a/crates/common/kigi-tracing/src/testing.rs +++ /dev/null @@ -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 { - 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())) -} diff --git a/crates/common/kigi-tracing/src/timer.rs b/crates/common/kigi-tracing/src/timer.rs deleted file mode 100644 index 21990a6..0000000 --- a/crates/common/kigi-tracing/src/timer.rs +++ /dev/null @@ -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>(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(&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(); - } -} diff --git a/crates/common/kigi-tracing/src/tokio.rs b/crates/common/kigi-tracing/src/tokio.rs deleted file mode 100644 index 2b1638e..0000000 --- a/crates/common/kigi-tracing/src/tokio.rs +++ /dev/null @@ -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(future: F) -> JoinHandle -where - F: Future + Send + 'static, - F::Output: Send + 'static, -{ - tokio::spawn(future.instrument(Span::current())) -}