From 5919526e91bf31366a17fb827212d00f5ba1252a 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 21:44:43 -0400 Subject: [PATCH] Show reasoning effort in the model's own vocabulary (K3: max, not xhigh) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The welcome and prompt model labels (and /effort's 'current' hint) rendered the canonical internal level name, so K3 at its default effort showed 'K3 (xhigh)' even though the server's vocabulary for that level is 'max' (live /models think_efforts: low/high/max). New ModelState::reasoning_effort_display() resolves the current effort through the model's own effort menu (option id whose value matches), falling back to the canonical name only when the model has no entry for the level. All three display sites route through it; test pins the K3 mapping (Xhigh → 'max', Low → 'low', no-menu-entry → canonical). --- .../observability_bridge_mapping_tests.rs | 97 - .../codegen/kigi-tui/src/acp/model_state.rs | 40 + .../kigi-tui/src/app/agent_view/render.rs | 2 +- crates/codegen/kigi-tui/src/app/app_view.rs | 2 +- .../kigi-tui/src/slash/commands/effort.rs | 2 +- .../codegen/kigi-workspace-client/Cargo.toml | 28 - .../codegen/kigi-workspace-client/src/lib.rs | 810 ----- .../src/bin/workspace_server.rs | 583 ---- .../src/bin/workspace_server_probe.rs | 214 -- .../codegen/kigi-workspace/src/daemonize.rs | 1080 ------ .../codegen/kigi-workspace/src/diag_server.rs | 561 --- .../src/file_system/client_fs.rs | 775 ----- crates/codegen/kigi-workspace/src/hub.rs | 1429 -------- crates/codegen/kigi-workspace/src/hub_auth.rs | 490 --- .../codegen/kigi-workspace/src/hub_channel.rs | 127 - crates/codegen/kigi-workspace/src/hub_ids.rs | 10 - .../codegen/kigi-workspace/src/hub_server.rs | 2999 ----------------- crates/codegen/kigi-workspace/src/mcp.rs | 220 -- .../src/permission/hub_permission.rs | 498 --- .../kigi-workspace/src/preview_supervisor.rs | 1326 -------- crates/codegen/kigi-workspace/src/recovery.rs | 158 - .../kigi-workspace/src/rpc_envelope.rs | 227 -- .../common/kigi-computer-hub-core/Cargo.toml | 25 - .../kigi-computer-hub-core/src/inner.rs | 73 - .../common/kigi-computer-hub-core/src/lib.rs | 29 - .../kigi-computer-hub-core/src/local.rs | 77 - .../kigi-computer-hub-core/src/registry.rs | 305 -- .../kigi-computer-hub-core/src/remote.rs | 541 --- .../kigi-computer-hub-core/src/resolver.rs | 280 -- .../kigi-computer-hub-core/src/transport.rs | 116 - .../tests/compound_resolver.rs | 335 -- .../tests/inner_dispatch.rs | 317 -- .../tests/internal_error_detail.rs | 65 - .../tests/local_transport.rs | 369 -- .../tests/remote_proxy.rs | 727 ---- .../tests/tool_registry.rs | 607 ---- .../tests/transport_trait.rs | 137 - .../tests/workspace_unavailable.rs | 203 -- .../kigi-computer-hub-mcp-adapter/Cargo.toml | 30 - .../src/bridge.rs | 781 ----- .../kigi-computer-hub-mcp-adapter/src/lib.rs | 52 - .../src/metrics.rs | 61 - .../src/transport.rs | 38 - .../src/types.rs | 107 - .../common/kigi-computer-hub-sdk/Cargo.toml | 57 - .../kigi-computer-hub-sdk/src/admission.rs | 333 -- .../common/kigi-computer-hub-sdk/src/auth.rs | 238 -- .../kigi-computer-hub-sdk/src/cancel.rs | 301 -- .../kigi-computer-hub-sdk/src/connection.rs | 2694 --------------- .../src/connection_borrow.rs | 224 -- .../common/kigi-computer-hub-sdk/src/demux.rs | 973 ------ .../kigi-computer-hub-sdk/src/donate_pump.rs | 206 -- .../common/kigi-computer-hub-sdk/src/error.rs | 346 -- .../kigi-computer-hub-sdk/src/handshake.rs | 93 - .../kigi-computer-hub-sdk/src/harness.rs | 2940 ---------------- .../common/kigi-computer-hub-sdk/src/lib.rs | 71 - .../kigi-computer-hub-sdk/src/log_donate.rs | 592 ---- .../src/metric_donate.rs | 396 --- .../kigi-computer-hub-sdk/src/metrics.rs | 552 --- .../kigi-computer-hub-sdk/src/notification.rs | 362 -- .../src/observability.rs | 275 -- .../src/oidc_provider.rs | 336 -- .../common/kigi-computer-hub-sdk/src/pool.rs | 344 -- .../kigi-computer-hub-sdk/src/refcount.rs | 123 - .../kigi-computer-hub-sdk/src/server.rs | 2649 --------------- .../kigi-computer-hub-sdk/src/trace_donate.rs | 206 -- 66 files changed, 43 insertions(+), 31221 deletions(-) delete mode 100644 crates/codegen/kigi-shell/src/session/acp_session_tests/observability_bridge_mapping_tests.rs delete mode 100644 crates/codegen/kigi-workspace-client/Cargo.toml delete mode 100644 crates/codegen/kigi-workspace-client/src/lib.rs delete mode 100644 crates/codegen/kigi-workspace/src/bin/workspace_server.rs delete mode 100644 crates/codegen/kigi-workspace/src/bin/workspace_server_probe.rs delete mode 100644 crates/codegen/kigi-workspace/src/daemonize.rs delete mode 100644 crates/codegen/kigi-workspace/src/diag_server.rs delete mode 100644 crates/codegen/kigi-workspace/src/file_system/client_fs.rs delete mode 100644 crates/codegen/kigi-workspace/src/hub.rs delete mode 100644 crates/codegen/kigi-workspace/src/hub_auth.rs delete mode 100644 crates/codegen/kigi-workspace/src/hub_channel.rs delete mode 100644 crates/codegen/kigi-workspace/src/hub_ids.rs delete mode 100644 crates/codegen/kigi-workspace/src/hub_server.rs delete mode 100644 crates/codegen/kigi-workspace/src/mcp.rs delete mode 100644 crates/codegen/kigi-workspace/src/permission/hub_permission.rs delete mode 100644 crates/codegen/kigi-workspace/src/preview_supervisor.rs delete mode 100644 crates/codegen/kigi-workspace/src/recovery.rs delete mode 100644 crates/codegen/kigi-workspace/src/rpc_envelope.rs delete mode 100644 crates/common/kigi-computer-hub-core/Cargo.toml delete mode 100644 crates/common/kigi-computer-hub-core/src/inner.rs delete mode 100644 crates/common/kigi-computer-hub-core/src/lib.rs delete mode 100644 crates/common/kigi-computer-hub-core/src/local.rs delete mode 100644 crates/common/kigi-computer-hub-core/src/registry.rs delete mode 100644 crates/common/kigi-computer-hub-core/src/remote.rs delete mode 100644 crates/common/kigi-computer-hub-core/src/resolver.rs delete mode 100644 crates/common/kigi-computer-hub-core/src/transport.rs delete mode 100644 crates/common/kigi-computer-hub-core/tests/compound_resolver.rs delete mode 100644 crates/common/kigi-computer-hub-core/tests/inner_dispatch.rs delete mode 100644 crates/common/kigi-computer-hub-core/tests/internal_error_detail.rs delete mode 100644 crates/common/kigi-computer-hub-core/tests/local_transport.rs delete mode 100644 crates/common/kigi-computer-hub-core/tests/remote_proxy.rs delete mode 100644 crates/common/kigi-computer-hub-core/tests/tool_registry.rs delete mode 100644 crates/common/kigi-computer-hub-core/tests/transport_trait.rs delete mode 100644 crates/common/kigi-computer-hub-core/tests/workspace_unavailable.rs delete mode 100644 crates/common/kigi-computer-hub-mcp-adapter/Cargo.toml delete mode 100644 crates/common/kigi-computer-hub-mcp-adapter/src/bridge.rs delete mode 100644 crates/common/kigi-computer-hub-mcp-adapter/src/lib.rs delete mode 100644 crates/common/kigi-computer-hub-mcp-adapter/src/metrics.rs delete mode 100644 crates/common/kigi-computer-hub-mcp-adapter/src/transport.rs delete mode 100644 crates/common/kigi-computer-hub-mcp-adapter/src/types.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/Cargo.toml delete mode 100644 crates/common/kigi-computer-hub-sdk/src/admission.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/auth.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/cancel.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/connection.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/connection_borrow.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/demux.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/donate_pump.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/error.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/handshake.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/harness.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/lib.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/log_donate.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/metric_donate.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/metrics.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/notification.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/observability.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/oidc_provider.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/pool.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/refcount.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/server.rs delete mode 100644 crates/common/kigi-computer-hub-sdk/src/trace_donate.rs diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/observability_bridge_mapping_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/observability_bridge_mapping_tests.rs deleted file mode 100644 index 156f38d..0000000 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/observability_bridge_mapping_tests.rs +++ /dev/null @@ -1,97 +0,0 @@ -use super::*; -use crate::session::events::ToolOutcome; -use kigi_tool_protocol::session_event::ToolCallOutcome; -use kigi_tool_protocol::turn_hook::TurnHookOutcome; -#[test] -fn map_tool_outcome_success() { - assert_eq!( - map_tool_outcome(ToolOutcome::Success), - ToolCallOutcome::Success - ); -} -#[test] -fn map_tool_outcome_errors() { - assert_eq!(map_tool_outcome(ToolOutcome::Error), ToolCallOutcome::Error); - assert_eq!( - map_tool_outcome(ToolOutcome::InvalidTool), - ToolCallOutcome::Error - ); -} -#[test] -fn map_tool_outcome_cancellations() { - for variant in [ - ToolOutcome::PermissionRejected, - ToolOutcome::PermissionCancelled, - ToolOutcome::Followup, - ToolOutcome::HookDenied, - ToolOutcome::Cancelled, - ] { - assert_eq!( - map_tool_outcome(variant), - ToolCallOutcome::Cancelled, - "expected Cancelled for {variant:?}", - ); - } -} -#[test] -fn turn_result_completed() { - let result: Result = Ok(TurnOutcome::Completed { - snapshot: Box::new(None), - tools_called: vec![], - structured_output: None, - refusal: false, - }); - assert_eq!( - turn_result_to_hook_outcome(&result), - TurnHookOutcome::Completed - ); -} -#[test] -fn turn_result_cancelled() { - let result: Result = Ok(TurnOutcome::Cancelled { - category: None, - context: None, - }); - assert_eq!( - turn_result_to_hook_outcome(&result), - TurnHookOutcome::Cancelled - ); -} -#[test] -fn turn_result_error() { - let result: Result = Err(acp::Error::internal_error()); - assert_eq!(turn_result_to_hook_outcome(&result), TurnHookOutcome::Error); -} -#[test] -fn is_remote_image_url_classifies_schemes() { - assert!(is_remote_image_url("https://example.com/x.png")); - assert!(is_remote_image_url("http://example.com/x.png")); - assert!(!is_remote_image_url("file:///Users/me/x.png")); - assert!(!is_remote_image_url("data:image/png;base64,AAAA")); - assert!(!is_remote_image_url("")); - assert!(!is_remote_image_url("FILE:///Users/me/x.png")); -} -#[test] -fn pick_image_url_prefers_base64_over_file_uri() { - let img = agent_client_protocol::ImageContent::new("AAAA", "image/png") - .uri(Some("file:///Users/me/Downloads/screenshot.png".into())); - assert_eq!(pick_user_image_url(&img), "data:image/png;base64,AAAA"); -} -#[test] -fn pick_image_url_prefers_base64_when_https_uri_also_present() { - let img = agent_client_protocol::ImageContent::new("BBBB", "image/jpeg") - .uri(Some("https://example.com/x.jpg".into())); - assert_eq!(pick_user_image_url(&img), "data:image/jpeg;base64,BBBB"); -} -#[test] -fn pick_image_url_falls_back_to_https_uri_when_data_empty() { - let img = agent_client_protocol::ImageContent::new(String::new(), "image/png") - .uri(Some("https://example.com/x.png".into())); - assert_eq!(pick_user_image_url(&img), "https://example.com/x.png"); -} -#[test] -fn pick_image_url_ignores_file_uri_when_data_empty() { - let img = agent_client_protocol::ImageContent::new(String::new(), "image/png") - .uri(Some("file:///Users/me/missing.png".into())); - assert_eq!(pick_user_image_url(&img), "data:image/png;base64,"); -} diff --git a/crates/codegen/kigi-tui/src/acp/model_state.rs b/crates/codegen/kigi-tui/src/acp/model_state.rs index 017e41f..89db5da 100644 --- a/crates/codegen/kigi-tui/src/acp/model_state.rs +++ b/crates/codegen/kigi-tui/src/acp/model_state.rs @@ -204,6 +204,22 @@ impl ModelState { parse_reasoning_efforts_meta(info.meta.as_ref()).unwrap_or_else(legacy_effort_options) } + /// Display token for the current reasoning effort, in the MODEL's own + /// vocabulary: the menu option id whose value matches (e.g. K3 spells + /// `Xhigh` as `max`), falling back to the canonical name only when the + /// model has no menu entry for it. Displays must never leak the internal + /// canonical spelling for a level the server names differently. + pub fn reasoning_effort_display(&self) -> Option { + let effort = self.reasoning_effort?; + let display = self + .reasoning_effort_options() + .into_iter() + .find(|opt| opt.value == effort) + .map(|opt| opt.id) + .unwrap_or_else(|| effort.as_str().to_owned()); + Some(display) + } + /// Map a typed/selected effort token to its canonical value for the current /// model. Accepts a menu option id (case-insensitive) or a canonical level /// that appears as a **value** in that model's menu. Levels the model does @@ -479,6 +495,30 @@ mod tests { assert_eq!(opts[1].description.as_deref(), Some("Max")); } + /// K3 spells `Xhigh` as `max` (live /models `think_efforts`): every + /// display surface must show the model's own vocabulary, never the + /// internal canonical name. + #[test] + fn reasoning_effort_display_uses_model_vocabulary() { + let mut state = state_with_meta(Some(serde_json::json!({ + "supportsReasoningEffort": true, + "reasoningEfforts": [ + { "id": "low", "value": "low", "label": "Low" }, + { "id": "high", "value": "high", "label": "High" }, + { "id": "max", "value": "xhigh", "label": "Max" }, + ], + }))); + state.reasoning_effort = Some(ReasoningEffort::Xhigh); + assert_eq!(state.reasoning_effort_display().as_deref(), Some("max")); + state.reasoning_effort = Some(ReasoningEffort::Low); + assert_eq!(state.reasoning_effort_display().as_deref(), Some("low")); + // No menu entry for the level → canonical fallback (never hide it). + state.reasoning_effort = Some(ReasoningEffort::Medium); + assert_eq!(state.reasoning_effort_display().as_deref(), Some("medium")); + state.reasoning_effort = None; + assert_eq!(state.reasoning_effort_display(), None); + } + #[test] fn reasoning_effort_options_gate_first_empty_when_unsupported() { // No current model → empty. diff --git a/crates/codegen/kigi-tui/src/app/agent_view/render.rs b/crates/codegen/kigi-tui/src/app/agent_view/render.rs index ec03ff1..15fe1f7 100644 --- a/crates/codegen/kigi-tui/src/app/agent_view/render.rs +++ b/crates/codegen/kigi-tui/src/app/agent_view/render.rs @@ -2038,7 +2038,7 @@ impl AgentView { } let mode_flags: &[PromptFlag] = &mode_flags_vec; let multiline = self.multiline_mode; - let model_label = match self.session.models.reasoning_effort { + let model_label = match self.session.models.reasoning_effort_display() { Some(eff) => format!("{model_id} ({eff})"), None => model_id, }; diff --git a/crates/codegen/kigi-tui/src/app/app_view.rs b/crates/codegen/kigi-tui/src/app/app_view.rs index 7fef5b4..da529e4 100644 --- a/crates/codegen/kigi-tui/src/app/app_view.rs +++ b/crates/codegen/kigi-tui/src/app/app_view.rs @@ -3060,7 +3060,7 @@ impl AppView { self.tip.as_deref() }; let model_name_base = self.models.current_model_name().unwrap_or_default(); - let model_name = match self.models.reasoning_effort { + let model_name = match self.models.reasoning_effort_display() { Some(eff) => format!("{model_name_base} ({eff})"), None => model_name_base, }; diff --git a/crates/codegen/kigi-tui/src/slash/commands/effort.rs b/crates/codegen/kigi-tui/src/slash/commands/effort.rs index 859e945..cda1aa4 100644 --- a/crates/codegen/kigi-tui/src/slash/commands/effort.rs +++ b/crates/codegen/kigi-tui/src/slash/commands/effort.rs @@ -69,7 +69,7 @@ impl SlashCommand for EffortCommand { .collect(); let current = ctx .models - .reasoning_effort + .reasoning_effort_display() .map(|e| format!(" (current: {e})")) .unwrap_or_default(); let levels = if offered.is_empty() { diff --git a/crates/codegen/kigi-workspace-client/Cargo.toml b/crates/codegen/kigi-workspace-client/Cargo.toml deleted file mode 100644 index e46f6c7..0000000 --- a/crates/codegen/kigi-workspace-client/Cargo.toml +++ /dev/null @@ -1,28 +0,0 @@ -[package] -license = "Apache-2.0" -name = "kigi-workspace-client" -version.workspace = true -edition.workspace = true -description = "Lightweight typed client for hub-proxied workspace.* RPCs (shared by kigi-shell proxy mode and other consumers)" - -[dependencies] -serde = { workspace = true, features = ["derive"] } -serde_json = { workspace = true } -thiserror = { workspace = true } -tokio = { workspace = true, features = ["time"] } -tracing = { workspace = true } -kigi-computer-hub-sdk = { workspace = true } -kigi-workspace-types = { workspace = true } -kigi-tool-protocol = { workspace = true } -kigi-tool-runtime = { workspace = true } - -[features] -default = [] -default-bazel = [] -[dev-dependencies] -schemars = { workspace = true } -tokio = { workspace = true, features = ["macros", "rt", "test-util"] } -kigi-tool-types = { workspace = true } - -[lints] -workspace = true diff --git a/crates/codegen/kigi-workspace-client/src/lib.rs b/crates/codegen/kigi-workspace-client/src/lib.rs deleted file mode 100644 index f894c95..0000000 --- a/crates/codegen/kigi-workspace-client/src/lib.rs +++ /dev/null @@ -1,810 +0,0 @@ -#![allow( - unused_imports, - unused_variables, - unused_mut, - unreachable_code, - dead_code -)] -//! Typed client for hub-proxied `workspace.*` RPC methods — the single -//! transport for the `workspace_rpc` channel, shared by `WorkspaceOps` -//! proxy mode and by consumers that cannot depend on -//! `kigi-workspace`. Wire types live in -//! `kigi_workspace_types::rpc`; this crate adds the connected-state -//! latch, the generic [`WorkspaceClient::rpc`] core, and error mapping. -//! -//! No deadline is imposed by default ([`WorkspaceClient::with_deadline`] -//! opts in), preserving `WorkspaceOps::rpc_raw` semantics where callers -//! own their timeouts. -use kigi_computer_hub_sdk::harness::ToolHarness; -use kigi_tool_runtime::{ToolCallContext, ToolStreamItem, TypedToolOutput}; -use kigi_workspace_types::rpc::agents_md::{AgentConfigFile, DiscoverAgentsMdReq}; -use kigi_workspace_types::rpc::code_nav::{ - CodeFindDefinitionsReq, CodeFindReferencesReq, CodeGotoDefinitionReq, CodeGotoReferencesReq, - CodeIndexStatusReq, CodeIndexStatusResponse, CodeNavResponse, -}; -use kigi_workspace_types::rpc::fs::{ - FsDeleteFileReq, FsExistsData, FsExistsReq, FsListData, FsListReq, FsReadFileData, - FsReadFileReq, FsWriteFileReq, GetFilesReq, GetFilesRes, PutFilesReq, PutFilesRes, -}; -use kigi_workspace_types::rpc::git::{ - CheckoutCommitResponse, CommitResult, DetectVcsKindReq, GitBranchInfoReq, GitBranchListData, - GitBranchesReq, GitCheckoutCommitReq, GitCheckoutReq, GitCollectChangesReq, - GitCollectChangesResponse, GitCommitReq, GitCurrentCommitReq, GitDiffReq, GitDiffsData, - GitDiscardReq, GitFilesReq, GitInfoData, GitInfoReq, GitMetadataReq, GitReadFilesData, - GitResolveRootReq, GitStageContentReq, GitStageReq, GitStashReq, GitStatusExtReq, - GitStatusExtResponse, GitStatusReq, GitUnstageReq, StageData, VcsKind, -}; -use kigi_workspace_types::rpc::hunks::{ - BulkHunkActionResponse, FileSummary, HunkActionResponse, HunkAllActionReq, HunkFileActionReq, - HunkGetFileSummariesReq, HunkGetStagedFilesReq, HunkSingleActionReq, HunkTurnActionReq, -}; -use kigi_workspace_types::rpc::search::{ - ContentSearchData, ContentSearchRequest, FuzzyChangeReq, FuzzyCloseReq, FuzzyOpenReq, - FuzzyStatusReq, -}; -use kigi_workspace_types::rpc::session::{ - BeginPromptReq, EndPromptReq, FileRewindResponse, RewindToReq, -}; -use kigi_workspace_types::rpc::skills::{DiscoverPluginsReq, DiscoverSkillsReq, SkillInfo}; -use kigi_workspace_types::rpc::workspace::{ - ConfigureMcpReq, DropSessionReq, InstallPluginReq, LoadEnvrcReq, LoadPermissionsReq, - LoadProjectConfigReq, RefreshPluginsReq, ResolveFileReferencesReq, ToolDefinitionsReq, - UpdateToolConfigReq, WorkspaceInfo, WorkspaceInfoReq, -}; -use kigi_workspace_types::rpc::worktree::{ - ApplyWorktreeRequest, CreateWorktreeRequest, RemoveWorktreeRequest, WorktreeCreateSyncReq, - WorktreeDbPathReq, WorktreeDbPathResponse, WorktreeDbRebuildReq, WorktreeDbStatsReq, - WorktreeGcReq, WorktreeListReq, WorktreeShowReq, -}; -use kigi_workspace_types::rpc::{RpcEnvelope, RpcError, WORKSPACE_RPC_TOOL_ID, WorkspaceRpc}; -use serde_json::Value; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::time::Duration; -#[derive(Debug, thiserror::Error)] -pub enum WorkspaceClientError { - /// A previous call observed a fatal transport error and no - /// reconnect has been signalled since. - #[error("hub connection lost (previously disconnected)")] - NotConnected, - #[error("rpc failed: {0}")] - Transport(String), - #[error("{method} timed out after {after:?}")] - Timeout { method: String, after: Duration }, - #[error("{method}: response decode: {source}")] - Decode { - method: String, - #[source] - source: serde_json::Error, - }, - /// The server returned an error envelope. - #[error("workspace rpc error: {0}")] - Rpc(RpcError), -} -/// Consume a `ToolStream` to its terminal item, -/// discarding progress frames. -/// -/// Returns the terminal result, or a `ToolError::NetworkError` if the -/// stream ended without producing a terminal item. -pub async fn consume_stream_terminal( - stream: &mut kigi_tool_runtime::ToolStream, -) -> Result { - loop { - let item = std::future::poll_fn(|cx| stream.as_mut().poll_next(cx)).await; - match item { - Some(ToolStreamItem::Progress(_)) => {} - Some(ToolStreamItem::Terminal(result)) => return result, - None => { - return Err(kigi_tool_runtime::ToolError::network_error( - "stream ended without terminal item", - )); - } - } - } -} -/// Check whether a [`ToolError`](kigi_tool_runtime::ToolError) indicates -/// a fatal transport failure that should mark the hub as disconnected. -/// -/// Returns `true` for: -/// - `NetworkError` — direct transport failure (socket dropped, stream -/// ended without terminal item, etc.) -/// - `Custom` with `details.code == "protocol_error"` — half-closed -/// WebSocket producing malformed frames -pub fn is_transport_fatal(err: &kigi_tool_runtime::ToolError) -> bool { - match err.kind { - kigi_tool_runtime::ToolErrorKind::NetworkError => true, - kigi_tool_runtime::ToolErrorKind::Custom => err - .details - .as_ref() - .and_then(|d| d.get("code")) - .and_then(|c| c.as_str()) - .is_some_and(|c| c == "protocol_error"), - _ => false, - } -} -/// Typed client over a bound [`ToolHarness`] for `workspace.*` RPCs. -/// -/// Clones share the harness and the connected latch, which fast-fails -/// calls after a fatal transport error until -/// [`mark_connected`](Self::mark_connected) resets it (e.g. from an SDK -/// `on_reconnect` callback sharing the flag via -/// [`with_connected_flag`](Self::with_connected_flag)). -#[derive(Clone)] -pub struct WorkspaceClient { - harness: ToolHarness, - connected: Arc, - deadline: Option, -} -impl std::fmt::Debug for WorkspaceClient { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("WorkspaceClient") - .field("connected", &self.is_connected()) - .field("deadline", &self.deadline) - .finish_non_exhaustive() - } -} -impl WorkspaceClient { - pub fn new(harness: ToolHarness) -> Self { - Self { - harness, - connected: Arc::new(AtomicBool::new(true)), - deadline: None, - } - } - /// Shares a pre-created connected flag, so an SDK `on_reconnect` - /// callback holding the same `Arc` can reset it. - pub fn with_connected_flag(harness: ToolHarness, connected: Arc) -> Self { - Self { - harness, - connected, - deadline: None, - } - } - /// Set a per-call deadline covering dispatch and stream consumption. - pub fn with_deadline(mut self, deadline: Duration) -> Self { - self.deadline = Some(deadline); - self - } - pub fn harness(&self) -> &ToolHarness { - &self.harness - } - /// Whether the hub connection is believed to be alive. - pub fn is_connected(&self) -> bool { - self.connected.load(Ordering::Relaxed) - } - pub fn mark_disconnected(&self) { - self.connected.store(false, Ordering::Relaxed); - } - /// Reset after an SDK reconnect. - pub fn mark_connected(&self) { - self.connected.store(true, Ordering::Relaxed); - } - /// Untyped RPC call: `{"method": .., "params": ..}` through the - /// `workspace_rpc` hub tool, returning the raw envelope value. - pub async fn rpc_raw( - &self, - method: &str, - params: Value, - ) -> Result { - if !self.is_connected() { - return Err(WorkspaceClientError::NotConnected); - } - let tool_id = kigi_tool_protocol::ToolId::new(WORKSPACE_RPC_TOOL_ID) - .expect("constant tool id is valid"); - let args = serde_json::json!({ "method" : method, "params" : params }); - tracing::debug!(method, "WorkspaceClient::rpc"); - let fut = async { - let mut stream = self - .harness - .call(tool_id, args, ToolCallContext::default()) - .await; - consume_stream_terminal(&mut stream).await - }; - let result = match self.deadline { - Some(deadline) => match tokio::time::timeout(deadline, fut).await { - Ok(r) => r, - Err(_) => { - return Err(WorkspaceClientError::Timeout { - method: method.to_owned(), - after: deadline, - }); - } - }, - None => fut.await, - }; - let typed = result.map_err(|e| { - if is_transport_fatal(&e) { - self.mark_disconnected(); - } - WorkspaceClientError::Transport(e.to_string()) - })?; - Ok(typed.value) - } - /// Typed RPC call: derives the method and response type from the - /// request type's [`WorkspaceRpc`] impl and decodes the envelope. - pub async fn rpc(&self, req: &R) -> Result { - let params = serde_json::to_value(req).map_err(|e| WorkspaceClientError::Decode { - method: R::METHOD.to_owned(), - source: e, - })?; - let raw = self.rpc_raw(R::METHOD, params).await?; - let envelope: RpcEnvelope = - serde_json::from_value(raw).map_err(|e| WorkspaceClientError::Decode { - method: R::METHOD.to_owned(), - source: e, - })?; - envelope.into_result().map_err(WorkspaceClientError::Rpc) - } - /// `workspace.info`, decoded into the typed shape - /// (`WorkspaceInfoReq::Response` is the raw `Value` for - /// `WorkspaceOps` compat). - pub async fn info(&self) -> Result { - let raw = self.rpc(&WorkspaceInfoReq {}).await?; - serde_json::from_value(raw).map_err(|e| WorkspaceClientError::Decode { - method: WorkspaceInfoReq::METHOD.to_owned(), - source: e, - }) - } - /// `workspace.git_status` (JSON string value, ~1 KB server-side cap). - pub async fn git_status(&self) -> Result { - self.rpc(&GitStatusReq {}).await - } - pub async fn discover_skills(&self) -> Result, WorkspaceClientError> { - self.rpc(&DiscoverSkillsReq {}).await - } - pub async fn discover_agents_md(&self) -> Result, WorkspaceClientError> { - self.rpc(&DiscoverAgentsMdReq {}).await - } - pub async fn git_status_ext( - &self, - req: &GitStatusExtReq, - ) -> Result { - self.rpc(req).await - } - pub async fn git_files( - &self, - req: &GitFilesReq, - ) -> Result { - self.rpc(req).await - } - pub async fn git_diff(&self, req: &GitDiffReq) -> Result { - self.rpc(req).await - } - pub async fn git_stage(&self, req: &GitStageReq) -> Result { - self.rpc(req).await - } - pub async fn git_stage_content( - &self, - req: &GitStageContentReq, - ) -> Result<(), WorkspaceClientError> { - self.rpc(req).await - } - pub async fn git_unstage(&self, req: &GitUnstageReq) -> Result<(), WorkspaceClientError> { - self.rpc(req).await - } - pub async fn git_discard(&self, req: &GitDiscardReq) -> Result<(), WorkspaceClientError> { - self.rpc(req).await - } - pub async fn git_commit( - &self, - req: &GitCommitReq, - ) -> Result { - self.rpc(req).await - } - pub async fn git_checkout(&self, req: &GitCheckoutReq) -> Result<(), WorkspaceClientError> { - self.rpc(req).await - } - pub async fn git_stash(&self, req: &GitStashReq) -> Result<(), WorkspaceClientError> { - self.rpc(req).await - } - pub async fn git_info(&self, req: &GitInfoReq) -> Result { - self.rpc(req).await - } - pub async fn git_branches( - &self, - req: &GitBranchesReq, - ) -> Result { - self.rpc(req).await - } - pub async fn git_resolve_root( - &self, - req: &GitResolveRootReq, - ) -> Result, WorkspaceClientError> { - self.rpc(req).await - } - pub async fn git_current_commit( - &self, - req: &GitCurrentCommitReq, - ) -> Result, WorkspaceClientError> { - self.rpc(req).await - } - pub async fn detect_vcs_kind( - &self, - req: &DetectVcsKindReq, - ) -> Result { - self.rpc(req).await - } - pub async fn git_checkout_commit( - &self, - req: &GitCheckoutCommitReq, - ) -> Result { - self.rpc(req).await - } - pub async fn git_branch_info(&self) -> Result, WorkspaceClientError> { - self.rpc(&GitBranchInfoReq {}).await - } - pub async fn git_metadata(&self) -> Result { - self.rpc(&GitMetadataReq {}).await - } - /// `workspace.git_collect_changes` — collect repository changes for serialization. - pub async fn git_collect_changes( - &self, - req: &GitCollectChangesReq, - ) -> Result { - self.rpc(req).await - } - pub async fn put_files(&self, req: &PutFilesReq) -> Result { - self.rpc(req).await - } - pub async fn get_files(&self, req: &GetFilesReq) -> Result { - self.rpc(req).await - } - pub async fn fs_list(&self, req: &FsListReq) -> Result { - self.rpc(req).await - } - pub async fn fs_exists(&self, req: &FsExistsReq) -> Result { - self.rpc(req).await - } - pub async fn fs_read_file( - &self, - req: &FsReadFileReq, - ) -> Result { - self.rpc(req).await - } - pub async fn fs_write_file(&self, req: &FsWriteFileReq) -> Result<(), WorkspaceClientError> { - self.rpc(req).await - } - pub async fn fs_delete_file(&self, req: &FsDeleteFileReq) -> Result<(), WorkspaceClientError> { - self.rpc(req).await - } - pub async fn hunk_action( - &self, - req: &HunkSingleActionReq, - ) -> Result { - self.rpc(req).await - } - pub async fn hunk_file_action( - &self, - req: &HunkFileActionReq, - ) -> Result { - self.rpc(req).await - } - pub async fn hunk_turn_action( - &self, - req: &HunkTurnActionReq, - ) -> Result { - self.rpc(req).await - } - pub async fn hunk_all_action( - &self, - req: &HunkAllActionReq, - ) -> Result { - self.rpc(req).await - } - pub async fn hunk_get_staged_files(&self) -> Result, WorkspaceClientError> { - self.rpc(&HunkGetStagedFilesReq {}).await - } - pub async fn hunk_get_file_summaries(&self) -> Result, WorkspaceClientError> { - self.rpc(&HunkGetFileSummariesReq {}).await - } - pub async fn code_goto_definition( - &self, - req: &CodeGotoDefinitionReq, - ) -> Result { - self.rpc(req).await - } - pub async fn code_goto_references( - &self, - req: &CodeGotoReferencesReq, - ) -> Result { - self.rpc(req).await - } - pub async fn code_find_definitions( - &self, - req: &CodeFindDefinitionsReq, - ) -> Result { - self.rpc(req).await - } - pub async fn code_find_references( - &self, - req: &CodeFindReferencesReq, - ) -> Result { - self.rpc(req).await - } - pub async fn code_index_status( - &self, - req: &CodeIndexStatusReq, - ) -> Result { - self.rpc(req).await - } - pub async fn ripgrep( - &self, - req: &ContentSearchRequest, - ) -> Result { - self.rpc(req).await - } - pub async fn fuzzy_open(&self, req: &FuzzyOpenReq) -> Result { - self.rpc(req).await - } - pub async fn fuzzy_change(&self, req: &FuzzyChangeReq) -> Result { - self.rpc(req).await - } - pub async fn fuzzy_close(&self, req: &FuzzyCloseReq) -> Result { - self.rpc(req).await - } - pub async fn fuzzy_status(&self, req: &FuzzyStatusReq) -> Result { - self.rpc(req).await - } - pub async fn create_worktree( - &self, - req: &CreateWorktreeRequest, - ) -> Result { - self.rpc(req).await - } - pub async fn worktree_create_sync( - &self, - req: &WorktreeCreateSyncReq, - ) -> Result { - self.rpc(req).await - } - pub async fn remove_worktree( - &self, - req: &RemoveWorktreeRequest, - ) -> Result { - self.rpc(req).await - } - pub async fn apply_worktree( - &self, - req: &ApplyWorktreeRequest, - ) -> Result { - self.rpc(req).await - } - pub async fn worktree_list( - &self, - req: &WorktreeListReq, - ) -> Result { - self.rpc(req).await - } - pub async fn worktree_show( - &self, - req: &WorktreeShowReq, - ) -> Result { - self.rpc(req).await - } - pub async fn worktree_gc(&self, req: &WorktreeGcReq) -> Result { - self.rpc(req).await - } - pub async fn worktree_db_rebuild(&self) -> Result { - self.rpc(&WorktreeDbRebuildReq {}).await - } - pub async fn worktree_db_path(&self) -> Result { - self.rpc(&WorktreeDbPathReq {}).await - } - pub async fn worktree_db_stats(&self) -> Result { - self.rpc(&WorktreeDbStatsReq {}).await - } - pub async fn begin_prompt(&self, req: &BeginPromptReq) -> Result<(), WorkspaceClientError> { - self.rpc(req).await - } - pub async fn end_prompt(&self, req: &EndPromptReq) -> Result<(), WorkspaceClientError> { - self.rpc(req).await - } - pub async fn rewind_to( - &self, - req: &RewindToReq, - ) -> Result { - self.rpc(req).await - } - pub async fn load_project_config(&self) -> Result { - self.rpc(&LoadProjectConfigReq {}).await - } - pub async fn load_permissions(&self) -> Result { - self.rpc(&LoadPermissionsReq {}).await - } - pub async fn load_envrc(&self) -> Result { - self.rpc(&LoadEnvrcReq {}).await - } - pub async fn tool_definitions( - &self, - req: &ToolDefinitionsReq, - ) -> Result { - self.rpc(req).await - } - pub async fn resolve_file_references( - &self, - req: &ResolveFileReferencesReq, - ) -> Result { - self.rpc(req).await - } - pub async fn update_tool_config( - &self, - req: &UpdateToolConfigReq, - ) -> Result { - self.rpc(req).await - } - pub async fn drop_session(&self, req: &DropSessionReq) -> Result { - self.rpc(req).await - } - pub async fn configure_mcp( - &self, - req: &ConfigureMcpReq, - ) -> Result { - self.rpc(req).await - } - pub async fn install_plugin(&self) -> Result { - self.rpc(&InstallPluginReq {}).await - } - pub async fn refresh_plugins(&self) -> Result { - self.rpc(&RefreshPluginsReq {}).await - } - pub async fn discover_plugins(&self) -> Result, WorkspaceClientError> { - self.rpc(&DiscoverPluginsReq {}).await - } -} -#[cfg(test)] -mod tests { - use super::*; - use kigi_computer_hub_sdk::harness::LocalRegistry; - use kigi_tool_protocol::{SessionId, ToolId}; - use kigi_tool_runtime::{Tool, ToolError}; - use kigi_tool_types::ToolDescription; - use kigi_workspace_types::rpc::skills::SkillScope; - use schemars::JsonSchema; - use serde::Deserialize; - #[derive(Debug, Deserialize, JsonSchema)] - struct RpcArgs { - method: String, - params: serde_json::Value, - } - #[derive(Debug, serde::Serialize)] - #[serde(transparent)] - struct RawOut(serde_json::Value); - impl kigi_tool_runtime::ToolOutput for RawOut {} - #[derive(Debug)] - struct FakeWorkspaceRpc; - impl Tool for FakeWorkspaceRpc { - type Args = RpcArgs; - type Output = RawOut; - fn id(&self) -> ToolId { - ToolId::new(WORKSPACE_RPC_TOOL_ID).unwrap() - } - fn description(&self, _ctx: &::kigi_tool_runtime::ListToolsContext) -> ToolDescription { - ToolDescription::new(WORKSPACE_RPC_TOOL_ID, "fake workspace rpc") - } - async fn run(&self, _ctx: ToolCallContext, args: Self::Args) -> Result { - let ok = |v: serde_json::Value| Ok(RawOut(serde_json::json!({ "ok" : v }))); - match args.method.as_str() { - "workspace.info" => ok(serde_json::json!( - { "os" : "linux", "shell" : "bash", "cwd" : "/workspace", } - )), - "workspace.git_status" => ok(serde_json::json!("On branch main")), - "workspace.discover_skills" => ok(serde_json::json!( - [{ "name" : "my-skill", "description" : "A test skill", - "path" : "/workspace/.kigi/skills/my-skill/SKILL.md", "scope" - : "local", }] - )), - "workspace.discover_agents_md" => ok(serde_json::json!( - [{ "file_name" : "AGENTS.md", "file_path" : - "/workspace/AGENTS.md", "content" : "# Project instructions", - }] - )), - "workspace.echo_params" => ok(args.params), - "workspace.err" => Ok(RawOut(serde_json::json!( - { "err" : { "code" : "session_not_found", "message" : - "ghost" }, } - ))), - "workspace.malformed" => Ok(RawOut(serde_json::json!({ "neither" : true }))), - "workspace.slow" => { - tokio::time::sleep(Duration::from_secs(60)).await; - ok(serde_json::Value::Null) - } - "workspace.netfail" => Err(ToolError::network_error("socket dropped")), - "workspace.toolfail" => Err(ToolError::custom("some_code", "boom")), - other => panic!("unexpected method {other}"), - } - } - } - fn client() -> WorkspaceClient { - let registry = LocalRegistry::new(); - registry.register(FakeWorkspaceRpc); - let harness = ToolHarness::local_only_with( - registry, - SessionId::new("test").unwrap(), - Default::default(), - ); - WorkspaceClient::new(harness) - } - #[tokio::test] - async fn info_decodes_typed_response() { - let info = client().info().await.unwrap(); - assert_eq!( - info, - WorkspaceInfo { - os: "linux".into(), - shell: "bash".into(), - cwd: "/workspace".into(), - } - ); - } - #[tokio::test] - async fn git_status_returns_raw_value() { - let v = client().git_status().await.unwrap(); - assert_eq!(v, serde_json::json!("On branch main")); - } - #[tokio::test] - async fn discover_skills_decodes_typed_list() { - let skills = client().discover_skills().await.unwrap(); - assert_eq!(skills.len(), 1); - assert_eq!(skills[0].name, "my-skill"); - assert_eq!(skills[0].scope, SkillScope::Local); - } - #[tokio::test] - async fn discover_agents_md_decodes_typed_list() { - let files = client().discover_agents_md().await.unwrap(); - assert_eq!(files.len(), 1); - assert_eq!(files[0].file_name, "AGENTS.md"); - assert_eq!(files[0].file_path, "/workspace/AGENTS.md"); - } - #[tokio::test] - async fn typed_request_params_round_trip() { - #[derive(Debug, serde::Serialize)] - struct EchoReq { - flag: bool, - n: u32, - } - impl WorkspaceRpc for EchoReq { - const METHOD: &'static str = "workspace.echo_params"; - type Response = Value; - } - let echoed = client().rpc(&EchoReq { flag: true, n: 7 }).await.unwrap(); - assert_eq!(echoed, serde_json::json!({ "flag" : true, "n" : 7 })); - } - #[tokio::test] - async fn err_envelope_maps_to_rpc_error() { - let c = client(); - let raw = c.rpc_raw("workspace.err", Value::Null).await.unwrap(); - assert!(raw.get("err").is_some()); - #[derive(Debug, serde::Serialize)] - struct ErrReq; - impl WorkspaceRpc for ErrReq { - const METHOD: &'static str = "workspace.err"; - type Response = Value; - } - let err = c.rpc(&ErrReq).await.unwrap_err(); - match err { - WorkspaceClientError::Rpc(e) => { - assert_eq!(e.code, "session_not_found"); - assert_eq!(e.message, "ghost"); - } - other => panic!("expected Rpc, got {other:?}"), - } - assert!(c.is_connected(), "envelope errors must not trip the latch"); - } - #[tokio::test] - async fn malformed_envelope_maps_to_decode() { - let c = client(); - #[derive(Debug, serde::Serialize)] - struct MalformedReq; - impl WorkspaceRpc for MalformedReq { - const METHOD: &'static str = "workspace.malformed"; - type Response = Value; - } - let err = c.rpc(&MalformedReq).await.unwrap_err(); - assert!( - matches!(err, WorkspaceClientError::Decode { .. }), - "{err:?}" - ); - assert!(c.is_connected()); - } - #[tokio::test] - async fn network_error_trips_latch_and_fast_fails() { - let c = client(); - let err = c - .rpc_raw("workspace.netfail", Value::Null) - .await - .unwrap_err(); - assert!(matches!(err, WorkspaceClientError::Transport(_)), "{err:?}"); - assert!(!c.is_connected(), "network error must trip the latch"); - let err = c.rpc_raw("workspace.info", Value::Null).await.unwrap_err(); - assert!(matches!(err, WorkspaceClientError::NotConnected)); - c.mark_connected(); - assert!(c.rpc_raw("workspace.info", Value::Null).await.is_ok()); - } - #[tokio::test] - async fn non_fatal_tool_error_leaves_latch_up() { - let c = client(); - let err = c - .rpc_raw("workspace.toolfail", Value::Null) - .await - .unwrap_err(); - assert!(matches!(err, WorkspaceClientError::Transport(_)), "{err:?}"); - assert!( - c.is_connected(), - "non-fatal tool errors must not trip the latch" - ); - } - #[tokio::test(start_paused = true)] - async fn deadline_times_out_slow_calls() { - let c = client().with_deadline(Duration::from_secs(3)); - let err = c.rpc_raw("workspace.slow", Value::Null).await.unwrap_err(); - match err { - WorkspaceClientError::Timeout { method, after } => { - assert_eq!(method, "workspace.slow"); - assert_eq!(after, Duration::from_secs(3)); - } - other => panic!("expected Timeout, got {other:?}"), - } - assert!(c.is_connected(), "timeouts must not trip the latch"); - } - #[tokio::test(start_paused = true)] - async fn no_deadline_by_default_waits_out_slow_calls() { - let v = client() - .rpc_raw("workspace.slow", Value::Null) - .await - .unwrap(); - assert!(v.get("ok").is_some()); - } - #[tokio::test] - async fn shared_connected_flag_is_externally_controllable() { - let registry = LocalRegistry::new(); - registry.register(FakeWorkspaceRpc); - let harness = ToolHarness::local_only_with( - registry, - SessionId::new("test").unwrap(), - Default::default(), - ); - let flag = Arc::new(AtomicBool::new(true)); - let c = WorkspaceClient::with_connected_flag(harness, flag.clone()); - flag.store(false, Ordering::Relaxed); - let err = c.rpc_raw("workspace.info", Value::Null).await.unwrap_err(); - assert!(matches!(err, WorkspaceClientError::NotConnected)); - flag.store(true, Ordering::Relaxed); - assert!(c.rpc_raw("workspace.info", Value::Null).await.is_ok()); - } - #[tokio::test] - async fn clones_share_the_latch() { - let a = client(); - let b = a.clone(); - a.mark_disconnected(); - assert!(!b.is_connected()); - } - #[tokio::test] - async fn consume_stream_terminal_returns_ok() { - let value = serde_json::json!({ "result" : "hello" }); - let typed = TypedToolOutput::from_value(ToolId::new("t").unwrap(), value.clone()); - let mut stream = kigi_tool_runtime::terminal_only(Ok(typed)); - assert_eq!( - consume_stream_terminal(&mut stream).await.unwrap().value, - value - ); - } - #[tokio::test] - async fn consume_stream_terminal_returns_err() { - let mut stream: kigi_tool_runtime::ToolStream = - kigi_tool_runtime::terminal_only(Err(ToolError::network_error("oops"))); - let err = consume_stream_terminal(&mut stream).await.unwrap_err(); - assert!(err.to_string().contains("oops")); - } - #[tokio::test] - async fn consume_stream_terminal_exhausted_stream_is_network_error() { - let typed = TypedToolOutput::from_value(ToolId::new("t").unwrap(), Value::Null); - let mut stream = kigi_tool_runtime::terminal_only::(Ok(typed)); - let _ = consume_stream_terminal(&mut stream).await; - let err = consume_stream_terminal(&mut stream).await.unwrap_err(); - assert!( - err.to_string() - .contains("stream ended without terminal item") - ); - assert!(is_transport_fatal(&err)); - } -} diff --git a/crates/codegen/kigi-workspace/src/bin/workspace_server.rs b/crates/codegen/kigi-workspace/src/bin/workspace_server.rs deleted file mode 100644 index fe13c40..0000000 --- a/crates/codegen/kigi-workspace/src/bin/workspace_server.rs +++ /dev/null @@ -1,583 +0,0 @@ -//! Standalone workspace ToolServer for remote sandboxes. -//! -//! Reads OIDC credentials from `~/.kigi/auth.json`, connects to a -//! server, exposes workspace tools, and refreshes tokens -//! automatically. -use clap::Parser; -use kigi_workspace::config::WorkspaceServerMetadata; -use kigi_workspace::daemonize; -use kigi_workspace::diag_server; -use kigi_workspace::preview_supervisor::{self, PreviewArgs, PreviewVisibility}; -use std::path::PathBuf; -use url::Url; -/// OTLP `service.name` for this binary's exported traces/logs/metrics and -/// direct-OTLP fastrace export. Single source so the call sites can't drift. -const SERVICE_NAME: &str = "prod_grok_workspace"; -const EXIT_SERVER_ID_INVALID: i32 = 3; -const INVALID_SERVER_ID_MARKER: &str = "workspace-server: invalid --server-id"; -fn server_id_startup_error(id: &str) -> Option { - id.parse::() - .err() - .map(|e| format!("{INVALID_SERVER_ID_MARKER} {id:?}: {e}")) -} -#[derive(Parser)] -#[command(name = "kigi-workspace-server")] -#[command(about = "Standalone workspace ToolServer for the server connection")] -struct Args { - /// Print the capability manifest as JSON to stdout and exit 0. Legacy - /// binaries reject the unknown flag via clap (non-zero exit), giving the - /// launcher a definitive feature probe. - #[arg(long)] - capabilities: bool, - #[arg(long, default_value = "wss://computer-hub.kigi.com/v1/tools")] - hub_url: String, - #[arg(long)] - auth_config: Option, - #[arg(long)] - cwd: Option, - /// Stable server identity for hub registration. Used as the - /// `server_id` in `servers.list` and `server.bind` so clients - /// can address this specific workspace server. - /// When omitted, the SDK default ("workspace-server") is used. - #[arg(long)] - server_id: Option, - /// JSON metadata attached to the tool server registration. - /// Propagated to `ServerInfo.metadata` in `servers.list` responses. - #[arg(long)] - metadata: Option, - /// Path to write a PID file once the server connection is established. - /// The sandbox service polls this file to determine readiness. - #[arg(long, default_value = daemonize::DEFAULT_READY_PATH)] - ready_file: PathBuf, - /// Unix-socket path for the in-guest diagnostics HTTP server - /// (`/ready`, `/statusz`). - #[cfg(unix)] - #[arg(long, default_value = diag_server::DEFAULT_DIAG_SOCKET_PATH)] - diag_socket: PathBuf, - /// Loopback TCP port for the diagnostics HTTP server (Windows guests, - /// which lack a reliable Unix-socket HTTP client). - #[cfg(windows)] - #[arg(long, default_value_t = diag_server::DEFAULT_DIAG_PORT)] - diag_port: u16, - /// Permit a plaintext `ws://` hub on a non-loopback host. Only for a - /// mesh-secured transport; the bearer crosses the network otherwise. - #[arg(long)] - allow_insecure_ws: bool, - /// Fail `session.bind`s without an explicit toolset closed (RPC-only) - /// instead of widening to the built-in default catalog. Passed by the - /// sandbox service; doubles as a version tripwire (a stale revived binary - /// rejects the argv and never reports ready). - #[arg(long)] - require_explicit_toolset: bool, - /// Confine `x.ai/fs/*` resolution to the workspace root (reject `..`, - /// absolute-outside-root, symlink escapes). On by default: the standalone - /// server always backs a remote-sandbox workspace, a real tenant boundary. - /// Override with `KIGI_WORKSPACE_CONFINE_FS_TO_ROOT=false` (e.g. local dev). - #[arg( - long, - env = "KIGI_WORKSPACE_CONFINE_FS_TO_ROOT", - default_value_t = true, - action = clap::ArgAction::Set, - )] - confine_fs_to_workspace_root: bool, - /// Self-daemonize at startup: double-fork + `setsid()` into a new session - /// and process group (escaping the launcher's process-group reap), - /// redirect stdio to a log file, and hold a single-instance pidfile lock. - /// - /// Off by default — opt-in, passed by the launcher in the supervised - /// deployment mode. With the flag absent, startup is unchanged. - #[arg(long)] - daemonize: bool, - /// Where `--daemonize` redirects stdout+stderr. Ignored without - /// `--daemonize`. - #[arg(long, default_value = daemonize::DEFAULT_LOG_PATH)] - log_file: PathBuf, - /// Single-instance pidfile lock path used with `--daemonize`. Ignored - /// without `--daemonize`. - #[arg(long, default_value = daemonize::DEFAULT_PIDFILE_PATH)] - pid_file: PathBuf, - #[command(flatten)] - preview: PreviewCliArgs, -} -/// Preview-proxy supervision flags. Forwarded 1:1 to the -/// `/usr/local/bin/xai-grok-preview-proxy` child (see `cli.rs` for the proxy's -/// flag names). Off by default — when `--preview-enabled` is absent the -/// supervisor is never started and startup is byte-for-byte the non-preview -/// path. -#[derive(clap::Args, Debug)] -struct PreviewCliArgs { - /// Spawn and supervise the in-sandbox preview-proxy. The launcher passes - /// this only when the proxy binary was mounted into this container. - #[arg(long)] - preview_enabled: bool, - /// Proxy `--preview-port` (externally exposed listener). Absent ⇒ proxy default. - #[arg(long)] - preview_port: Option, - /// Proxy `--control-port` (loopback control). Absent ⇒ proxy default. - #[arg(long)] - preview_control_port: Option, - /// Proxy `--visibility` (`owner` | `public`). Absent ⇒ proxy default. - /// Validated here so a bad value fails fast instead of crash-looping the proxy. - #[arg(long, value_enum)] - preview_visibility: Option, - /// Proxy `--instance-suffix` for inbound Host validation. - #[arg(long)] - preview_instance_suffix: Option, - /// Proxy `--auth-redirect`: URL the unauthenticated handshake redirects to. - /// Required for the default `owner` gate to redirect rather than deny. - #[arg(long)] - preview_auth_redirect: Option, - /// Proxy `--allow-public` org public-policy gate (forwarded only when set). - #[arg(long)] - preview_allow_public: bool, - /// Proxy `--workspace-server-port` (added to the discovery denylist). - #[arg(long)] - preview_workspace_server_port: Option, -} -impl PreviewCliArgs { - fn into_preview_args(self, workspace_dir: PathBuf) -> PreviewArgs { - PreviewArgs { - enabled: self.preview_enabled, - port: self.preview_port, - control_port: self.preview_control_port, - visibility: self.preview_visibility, - instance_suffix: self.preview_instance_suffix, - auth_redirect: self.preview_auth_redirect, - allow_public: self.preview_allow_public, - workspace_server_port: self.preview_workspace_server_port, - workspace_dir, - } - } -} -/// Capability manifest printed by `--capabilities`, consumed by the sandbox -/// launcher to pick a launch protocol. Additions are backward-compatible. -#[derive(Debug, serde::Serialize)] -struct Capabilities { - /// The in-guest diagnostics HTTP server (`/ready`, `/statusz`, `/logs`). - diag: bool, -} -const CAPABILITIES: Capabilities = Capabilities { diag: true }; -fn main() -> anyhow::Result<()> { - let mut args = Args::parse(); - if args.capabilities { - println!("{}", serde_json::to_string(&CAPABILITIES)?); - return Ok(()); - } - if let Some(msg) = args.server_id.as_deref().and_then(server_id_startup_error) { - eprintln!("{msg}"); - std::process::exit(EXIT_SERVER_ID_INVALID); - } - let cwd = match args.cwd { - Some(ref p) => dunce::canonicalize(p)?, - None => std::env::current_dir()?, - }; - let _pidfile_guard = if args.daemonize { - let anchor = |p: PathBuf| if p.is_absolute() { p } else { cwd.join(p) }; - args.log_file = anchor(std::mem::take(&mut args.log_file)); - args.pid_file = anchor(std::mem::take(&mut args.pid_file)); - args.ready_file = anchor(std::mem::take(&mut args.ready_file)); - #[cfg(unix)] - { - args.diag_socket = anchor(std::mem::take(&mut args.diag_socket)); - } - args.auth_config = args.auth_config.take().map(anchor); - daemonize::daemonize(&args.log_file)?; - match daemonize::PidFile::acquire_or_take_over(&args.pid_file, daemonize::TAKEOVER_GRACE)? { - Some(guard) => Some(guard), - None => return Ok(()), - } - } else { - None - }; - tokio::runtime::Builder::new_multi_thread() - .enable_all() - .build()? - .block_on(run(args, cwd)) -} -async fn run(args: Args, cwd: PathBuf) -> anyhow::Result<()> { - let _ = rustls::crypto::ring::default_provider().install_default(); - use tracing_subscriber::layer::SubscriberExt as _; - use tracing_subscriber::util::SubscriberInitExt as _; - let env_filter = tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")); - let donating = kigi_computer_hub_sdk::DonatingLogLayer::new_inert(); - tracing_subscriber::registry() - .with(env_filter) - .with(tracing_subscriber::fmt::layer()) - .with(donating.clone()) - .init(); - let direct_otlp = match std::env::var("KIGI_WORKSPACE_OTLP_ENDPOINT") { - Ok(endpoint) if !endpoint.is_empty() => { - match kigi_tracing::init_fastrace(endpoint.clone(), SERVICE_NAME.to_owned(), None) { - Ok(()) => { - tracing::info!(% endpoint, "trace export enabled (direct OTLP)"); - true - } - Err(e) => { - tracing::warn!(error = % e, "direct OTLP trace export init failed"); - false - } - } - } - _ => false, - }; - let url = Url::parse(&args.hub_url).map_err(|e| anyhow::anyhow!("invalid --hub-url: {e}"))?; - { - use kigi_sandbox::{ProfileName, SandboxManager}; - let profile = match std::env::var("KIGI_SANDBOX_PROFILE").ok() { - Some(val) => { - let parsed = val - .parse::() - .expect("ProfileName::from_str is infallible"); - if matches!(parsed, ProfileName::Custom(_)) { - tracing::warn!( - value = % val, - "Unrecognized KIGI_SANDBOX_PROFILE, defaulting to workspace" - ); - ProfileName::Workspace - } else { - parsed - } - } - None if kigi_sandbox::trust_bwrap_marker_for_devbox() => ProfileName::Devbox, - None => ProfileName::Workspace, - }; - let profile_name = profile.to_string(); - if profile == ProfileName::Off { - tracing::info!( - profile = % profile_name, - "Sandbox explicitly disabled via KIGI_SANDBOX_PROFILE=off" - ); - } else { - let mut sandbox = SandboxManager::new(profile, &cwd); - if let Err(e) = sandbox.apply(&cwd) { - tracing::warn!( - error = % e, "Sandbox apply returned error, continuing unsandboxed" - ); - } else if !sandbox.is_applied() { - tracing::warn!("Sandbox could not be applied (unsupported platform)"); - } - sandbox.install(); - let active = kigi_sandbox::is_active(); - let status_msg = if active { - "Workspace server sandbox active" - } else { - "Workspace server sandbox NOT active" - }; - tracing::info!( - profile = % profile_name, active, restrict_network = - kigi_sandbox::should_restrict_child_network(), "{status_msg}" - ); - } - } - let auth_provider = kigi_workspace::hub_auth::provider(&url, args.auth_config.as_deref())?; - tracing::info!(hub_url = % url, cwd = % cwd.display(), "Starting workspace server"); - let cwd_display = cwd.display().to_string(); - let session_id = std::env::var("KIGI_SESSION_ID").ok(); - let parsed_metadata = match args.metadata { - Some(json_str) => Some( - serde_json::from_str(&json_str) - .map_err(|e| anyhow::anyhow!("invalid --metadata JSON: {e}"))?, - ), - None => None, - }; - let metadata = WorkspaceServerMetadata::merge_session_metadata(parsed_metadata, session_id); - let launch_id = metadata - .as_ref() - .and_then(|v| v.get("launch_id")) - .and_then(serde_json::Value::as_str) - .map(str::to_owned); - let diag_handle = diag_server::DiagHandle::new(launch_id); - #[cfg(unix)] - let diag_listener = diag_server::DiagListener::Unix(args.diag_socket); - #[cfg(windows)] - let diag_listener = diag_server::DiagListener::Tcp(args.diag_port); - let diag_log_file = args.daemonize.then_some(args.log_file); - let _diag_server = - match diag_server::serve(diag_listener, diag_handle.clone(), diag_log_file).await { - Ok(bound) => { - tracing::info!(addr = % bound.addr, "diagnostics server listening"); - Some(bound) - } - Err(e) => { - if args.daemonize { - tracing::error!(error = % e, "{}", diag_server::DIAG_BIND_FAILED_MARKER); - std::process::exit(diag_server::EXIT_DIAG_BIND_FAILED); - } - tracing::warn!( - error = % e, "{} (continuing without)", - diag_server::DIAG_BIND_FAILED_MARKER - ); - None - } - }; - tracing::info!( - cwd = % cwd_display, - "Workspace server starting — sessions created dynamically via server bind" - ); - let server_id = args.server_id.clone(); - let status_config = kigi_workspace::StatusConfig::from_env(); - let preview_shutdown = if args.preview.preview_enabled { - let control_port = args.preview.preview_control_port; - let cfg = args.preview.into_preview_args(cwd.clone()); - let (tx, rx) = tokio::sync::watch::channel(false); - tokio::spawn(preview_supervisor::supervise_preview(cfg, rx)); - Some((tx, control_port)) - } else { - None - }; - let project_lsp_trusted = true; - let preview_scrape_interval = status_config.preview_activity_scrape_interval; - kigi_workspace::init_metrics(); - let ws_handle = kigi_workspace::handle::connect_local_workspace( - cwd, - url, - auth_provider, - metadata, - server_id.clone(), - None, - args.allow_insecure_ws, - status_config, - project_lsp_trusted, - Some(args.ready_file.clone()), - Some(diag_handle.clone()), - args.require_explicit_toolset, - args.confine_fs_to_workspace_root, - ) - .await - .map_err(|e| anyhow::anyhow!("failed to connect workspace to hub: {e}"))?; - if let Some((tx, control_port)) = &preview_shutdown { - tokio::spawn(preview_supervisor::supervise_preview_activity( - *control_port, - ws_handle.activity_tracker().clone(), - preview_scrape_interval, - tx.subscribe(), - )); - } - let mut donation_pump = None; - if !direct_otlp { - match ws_handle.trace_donation_reporter(SERVICE_NAME).await { - Some((reporter, pump)) => { - fastrace::set_reporter(reporter, fastrace::collector::Config::default()); - donation_pump = Some(pump); - tracing::info!("trace export enabled"); - } - None => tracing::info!("trace export disabled (not connected)"), - } - } - let mut log_donation_pump = None; - match ws_handle.log_donation_layer(SERVICE_NAME).await { - Some((sender, pump)) => { - donating.activate(sender); - log_donation_pump = Some(pump); - tracing::info!("log export enabled"); - } - None => tracing::info!("log export disabled (not connected)"), - } - let mut metric_donation_pump = None; - match ws_handle.metric_donation_reporter(SERVICE_NAME).await { - Some(pump) => { - metric_donation_pump = Some(pump); - tracing::info!("metric export enabled"); - } - None => tracing::info!("metric export disabled (not connected)"), - } - tracing::info!( - server_id = ? server_id, "Workspace server connected to hub. Serving tools." - ); - #[cfg(unix)] - { - use tokio::signal::unix::{SignalKind, signal}; - let mut sigterm = signal(SignalKind::terminate())?; - tokio::select! { - _ = tokio::signal::ctrl_c() => {} _ = sigterm.recv() => {} - } - } - #[cfg(not(unix))] - tokio::signal::ctrl_c().await?; - if let Some((tx, _)) = &preview_shutdown { - let _ = tx.send(true); - } - let _ = std::fs::remove_file(&args.ready_file); - diag_handle.set_shutting_down(); - tracing::info!("Received shutdown signal, draining..."); - let tracker = ws_handle.activity_tracker().clone(); - let grace_budget = kigi_workspace::handle::termination_grace_from_env(); - ws_handle - .two_phase_drain(grace_budget, kigi_workspace::handle::DrainReason::Sigterm) - .await; - tracker.set_shutting_down(); - tracing::info!("Shutting down..."); - fastrace::flush(); - if let Some(pump) = &donation_pump { - pump.drain().await; - } - kigi_computer_hub_sdk::flush_log_layer(); - if let Some(pump) = &log_donation_pump { - pump.drain().await; - } - if let Some(pump) = &metric_donation_pump { - pump.drain().await; - } - ws_handle.shutdown_hub().await; - kigi_sandbox::flush(); - Ok(()) -} -#[cfg(test)] -mod tests { - use super::*; - #[test] - fn capabilities_flag_parses_and_defaults_off() { - let args = Args::try_parse_from(["kigi-workspace-server"]).unwrap(); - assert!(!args.capabilities); - let args = Args::try_parse_from(["kigi-workspace-server", "--capabilities"]).unwrap(); - assert!(args.capabilities); - } - #[test] - fn capabilities_manifest_shape() { - let value = serde_json::to_value(CAPABILITIES).unwrap(); - assert_eq!(value, serde_json::json!({ "diag" : true })); - } - #[test] - fn capabilities_probe_of_legacy_binary_exits_nonzero() { - /// A stand-in for a pre-`--capabilities` Args surface. - #[derive(Debug, Parser)] - struct LegacyArgs { - #[arg(long)] - daemonize: bool, - } - let err = LegacyArgs::try_parse_from(["kigi-workspace-server", "--capabilities"]) - .expect_err("a legacy binary must reject the flag"); - assert_eq!(err.kind(), clap::error::ErrorKind::UnknownArgument); - assert_ne!( - err.exit_code(), - 0, - "the probe relies on a non-zero exit distinguishing legacy binaries" - ); - } - #[test] - fn daemonize_defaults_are_inert() { - let args = Args::try_parse_from(["kigi-workspace-server"]).unwrap(); - assert!(!args.daemonize); - assert_eq!(args.log_file, PathBuf::from(daemonize::DEFAULT_LOG_PATH)); - assert_eq!( - args.pid_file, - PathBuf::from(daemonize::DEFAULT_PIDFILE_PATH) - ); - assert_eq!( - args.ready_file, - PathBuf::from(daemonize::DEFAULT_READY_PATH) - ); - } - #[test] - fn invalid_server_id_produces_the_marker_line() { - for bad in ["auto:tool:x", ""] { - let msg = server_id_startup_error(bad) - .unwrap_or_else(|| panic!("server id {bad:?} must be rejected")); - assert!( - msg.starts_with(INVALID_SERVER_ID_MARKER), - "startup error must START with the greppable marker prefix: {msg}" - ); - } - assert_eq!(server_id_startup_error("session-abc-123"), None); - } - #[test] - fn argv_rejection_exit_code_is_distinct_from_server_id_exit_code() { - let err = Args::try_parse_from(["kigi-workspace-server", "--flag-from-the-future"]) - .err() - .expect("unknown argv must be rejected"); - assert_eq!(err.exit_code(), 2, "clap argv rejection exits 2"); - assert_ne!(err.exit_code(), EXIT_SERVER_ID_INVALID); - assert_ne!(EXIT_SERVER_ID_INVALID, 0); - assert_ne!(EXIT_SERVER_ID_INVALID, 1); - } - #[test] - fn preview_defaults_are_inert() { - let args = Args::try_parse_from(["kigi-workspace-server"]).unwrap(); - assert!(!args.preview.preview_enabled); - let cfg = args.preview.into_preview_args(PathBuf::from("/workspace")); - assert!(!cfg.enabled); - assert!( - cfg.to_argv().is_empty(), - "an inert config forwards no proxy args" - ); - } - #[test] - fn preview_flags_parse_and_lower_to_supervisor_config() { - let args = Args::try_parse_from([ - "kigi-workspace-server", - "--preview-enabled", - "--preview-port", - "6014", - "--preview-control-port", - "6015", - "--preview-visibility", - "public", - "--preview-instance-suffix", - ".inst.example", - "--preview-auth-redirect", - "https://grok.com/preview-auth", - "--preview-allow-public", - "--preview-workspace-server-port", - "8470", - ]) - .unwrap(); - assert!(args.preview.preview_enabled); - let cfg = args.preview.into_preview_args(PathBuf::from("/workspace")); - assert!(cfg.enabled); - assert_eq!(cfg.port, Some(6014)); - assert_eq!(cfg.control_port, Some(6015)); - assert_eq!(cfg.visibility, Some(PreviewVisibility::Public)); - assert_eq!(cfg.instance_suffix.as_deref(), Some(".inst.example")); - assert_eq!( - cfg.auth_redirect.as_deref(), - Some("https://grok.com/preview-auth") - ); - assert!(cfg.allow_public); - assert_eq!(cfg.workspace_server_port, Some(8470)); - assert_eq!(cfg.workspace_dir, PathBuf::from("/workspace")); - assert_eq!( - cfg.to_argv(), - vec![ - "--preview-port", - "6014", - "--control-port", - "6015", - "--visibility", - "public", - "--instance-suffix", - ".inst.example", - "--auth-redirect", - "https://grok.com/preview-auth", - "--allow-public", - "--workspace-server-port", - "8470", - ], - ); - } - #[test] - fn preview_visibility_rejects_invalid_value() { - let err = Args::try_parse_from([ - "kigi-workspace-server", - "--preview-enabled", - "--preview-visibility", - "nobody", - ]) - .err() - .expect("an invalid --preview-visibility must be rejected"); - assert_eq!(err.kind(), clap::error::ErrorKind::InvalidValue); - } - #[test] - fn preview_visibility_owner_parses_and_lowers() { - let args = Args::try_parse_from([ - "kigi-workspace-server", - "--preview-enabled", - "--preview-visibility", - "owner", - ]) - .unwrap(); - let cfg = args.preview.into_preview_args(PathBuf::from("/workspace")); - assert_eq!(cfg.visibility, Some(PreviewVisibility::Owner)); - assert_eq!(cfg.to_argv(), vec!["--visibility", "owner"]); - } -} diff --git a/crates/codegen/kigi-workspace/src/bin/workspace_server_probe.rs b/crates/codegen/kigi-workspace/src/bin/workspace_server_probe.rs deleted file mode 100644 index 2663cd7..0000000 --- a/crates/codegen/kigi-workspace/src/bin/workspace_server_probe.rs +++ /dev/null @@ -1,214 +0,0 @@ -//! Probe that verifies a running workspace-server actually serves tools -//! over the server connection (not just that it reached READY). -//! -//! Connects to the server as a client harness, binds the workspace-server's -//! session (its `server_id` equals the sandbox `session_id`), then -//! invokes real tools and asserts the results: -//! - `run_terminal_command` echoes a nonce that must come back, -//! - `read_file` reads a file the first command wrote. -//! -//! Exits 0 and prints `PROBE_OK` on success; non-zero with a diagnostic -//! on failure. Intended for sandbox end-to-end tests. -//! -//! The client connects to the *local* server (the same one the -//! workspace-server reaches back to, e.g. `ws://localhost:10030/v1/tools`) -//! using a bearer token. `servers.list` is scoped per-user on the server, so -//! the bearer must resolve to the same user that owns the session — the -//! access token from `~/.kigi/auth.json` does (same identity). - -use base64::Engine; -use clap::Parser; -use kigi_computer_hub_sdk::pool::HubConnectionPool; -use kigi_computer_hub_sdk::{AuthCredential, ToolHarnessBuilder}; -use kigi_tool_protocol::{SessionId, ToolId}; -use kigi_tool_runtime::{ToolCallContext, ToolStreamItem, TypedToolOutput}; -use serde_json::{Value, json}; -use url::Url; -use uuid::Uuid; - -#[derive(Parser)] -#[command(name = "workspace-server-probe")] -#[command(about = "Invoke tools on a running workspace-server via the server connection")] -struct Args { - /// Server WebSocket URL the client connects to (the local server). - #[arg(long, default_value = "ws://localhost:10030/v1/tools")] - hub_url: String, - - /// The workspace-server's server_id (equals the sandbox session_id). - #[arg(long)] - session_id: String, - - /// Hub user to authenticate as. Must match the user the - /// workspace-server registered under (`local-dev` in local-auth-dev). - #[arg(long, default_value = "local-dev")] - user_id: String, - - /// Explicit bearer token (overrides --user-id). For real tokens. - #[arg(long)] - bearer: Option, - - /// Workspace directory to bind on the server. - #[arg(long, default_value = "/workspace")] - cwd: String, -} - -/// Mint an unsigned JWT carrying `{"sub": user_id}`. The local-auth-dev -/// hub derives the principal's user from the bearer's JWT `sub` and does -/// NOT verify the signature, so this is sufficient to authenticate as a -/// specific local-dev user. Not usable against a real (verifying) hub. -fn dev_bearer(user_id: &str) -> String { - let b64 = base64::engine::general_purpose::URL_SAFE_NO_PAD; - let header = b64.encode(br#"{"alg":"none","typ":"JWT"}"#); - let payload = b64.encode(json!({ "sub": user_id }).to_string()); - format!("{header}.{payload}.sig") -} - -fn bearer(args: &Args) -> String { - args.bearer - .clone() - .unwrap_or_else(|| dev_bearer(&args.user_id)) -} - -/// Drive a tool call to its terminal item, discarding progress. -async fn call_tool( - harness: &kigi_computer_hub_sdk::ToolHarness, - name: &str, - args: Value, -) -> anyhow::Result { - let tool_id = ToolId::new(name).map_err(|e| anyhow::anyhow!("invalid tool id {name}: {e}"))?; - let mut stream = harness - .call(tool_id, args, ToolCallContext::default()) - .await; - loop { - let item = std::future::poll_fn(|cx| stream.as_mut().poll_next(cx)).await; - match item { - Some(ToolStreamItem::Progress(_)) => {} - Some(ToolStreamItem::Terminal(Ok(typed))) => { - let typed: TypedToolOutput = typed; - return Ok(typed.value); - } - Some(ToolStreamItem::Terminal(Err(e))) => { - anyhow::bail!("tool `{name}` returned error: {e}") - } - None => anyhow::bail!("tool `{name}` stream ended without a terminal item"), - } - } -} - -#[tokio::main] -async fn main() -> anyhow::Result<()> { - let _ = rustls::crypto::ring::default_provider().install_default(); - tracing_subscriber::fmt() - .with_env_filter( - tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn")), - ) - .init(); - let args = Args::parse(); - - let credential = AuthCredential::bearer(bearer(&args)); - - // The server connection can drop once right after connect (the SDK - // reconnects with backoff). build()'s eager session_open doesn't - // survive that first blip, so retry the connect + bind a few times. - let mut last_err = None; - for attempt in 1..=8u32 { - match connect_and_bind(&args, &credential).await { - Ok(harness) => return run_checks(&harness, &args).await, - Err(e) => { - eprintln!("[probe] attempt {attempt} failed: {e}"); - last_err = Some(e); - tokio::time::sleep(std::time::Duration::from_millis(1500)).await; - } - } - } - Err(last_err.unwrap_or_else(|| anyhow::anyhow!("probe failed"))) -} - -/// Connect to the server, open a session, and bind the workspace-server. -async fn connect_and_bind( - args: &Args, - credential: &AuthCredential, -) -> anyhow::Result { - let harness_session = SessionId::new(format!("probe-{}", Uuid::new_v4())) - .map_err(|e| anyhow::anyhow!("invalid harness session id: {e}"))?; - let url = Url::parse(&format!("{}?role=harness", args.hub_url)) - .map_err(|e| anyhow::anyhow!("invalid --hub-url: {e}"))?; - - let harness = ToolHarnessBuilder::default() - .pool(HubConnectionPool::new()) - .url(url) - .auth(credential.clone()) - .session(harness_session) - .build() - .await - .map_err(|e| anyhow::anyhow!("failed to build server harness: {e}"))?; - - let servers = harness - .list_servers() - .await - .map_err(|e| anyhow::anyhow!("servers.list failed: {e}"))?; - let server_ids: Vec = servers - .iter() - .map(|s| s.server_id.as_str().to_owned()) - .collect(); - eprintln!("[probe] servers visible to this user: {server_ids:?}"); - - let server_id = args.session_id.as_str(); - // Strict servers (`--require-explicit-toolset`) fail metadata-less binds - // closed: bind with exactly the tools the checks below invoke. - let metadata = json!({ - "tools": [ - {"id": "GrokBuild:run_terminal_cmd", "name_override": "run_terminal_command"}, - {"id": "GrokBuild:read_file"}, - ], - }); - let tools = harness - .session_bind(server_id, Some(&args.cwd), Some(metadata)) - .await - .map_err(|e| { - anyhow::anyhow!( - "session_bind({server_id}) failed: {e} (is the workspace-server registered \ - for this user? visible servers: {server_ids:?})" - ) - })?; - eprintln!( - "[probe] bound workspace-server {server_id}: {} tools available", - tools.len() - ); - Ok(harness) -} - -/// Invoke real tools on the bound workspace-server and assert results. -async fn run_checks( - harness: &kigi_computer_hub_sdk::ToolHarness, - _args: &Args, -) -> anyhow::Result<()> { - // 1) run_terminal_command must echo our nonce back. - let nonce = format!("probe-nonce-{}", Uuid::new_v4()); - let marker = format!("/tmp/{nonce}.txt"); - let out = call_tool( - harness, - "run_terminal_command", - json!({ "command": format!("echo {nonce} | tee {marker}") }), - ) - .await?; - let out_text = serde_json::to_string(&out).unwrap_or_default(); - anyhow::ensure!( - out_text.contains(&nonce), - "run_terminal_command output did not contain the nonce; got: {out_text}" - ); - eprintln!("[probe] run_terminal_command OK (nonce echoed)"); - - // 2) read_file must read the file the command wrote. - let read_out = call_tool(harness, "read_file", json!({ "target_file": marker })).await?; - let read_text = serde_json::to_string(&read_out).unwrap_or_default(); - anyhow::ensure!( - read_text.contains(&nonce), - "read_file did not return the written nonce; got: {read_text}" - ); - eprintln!("[probe] read_file OK (read back written nonce)"); - - println!("PROBE_OK"); - Ok(()) -} diff --git a/crates/codegen/kigi-workspace/src/daemonize.rs b/crates/codegen/kigi-workspace/src/daemonize.rs deleted file mode 100644 index 4d217eb..0000000 --- a/crates/codegen/kigi-workspace/src/daemonize.rs +++ /dev/null @@ -1,1080 +0,0 @@ -//! Self-daemonization and single-instance locking for the workspace-server. -//! -//! The server is launched fire-and-forget by the sandbox orchestrator, which -//! only ever holds a handle to the originally-spawned PID / process group. -//! After a double-fork + `setsid()` the surviving daemon lives in a new -//! session and process group, so a later process-group kill on the original -//! pgid cannot reach it. -//! -//! The double-fork MUST run before the tokio runtime — or `tracing_subscriber` -//! / the rustls provider — start any threads: forking a multi-threaded process -//! leaves every lock held by a non-forking thread permanently locked in the -//! child, which can deadlock it. - -use std::fs::{self, File, OpenOptions}; -use std::io::{self, Write}; -#[cfg(target_os = "linux")] -use std::os::fd::{FromRawFd as _, OwnedFd}; -use std::path::Path; -use std::time::{Duration, Instant}; -use std::{process, thread}; -#[cfg(windows)] -use windows::Win32::Foundation::HANDLE; - -use fs2::FileExt; - -use crate::util::is_lock_contended; - -#[cfg(unix)] -use std::os::unix::io::{AsRawFd, RawFd}; - -/// stdout + stderr redirect target when no `--log-file` is given. -#[cfg(unix)] -pub const DEFAULT_LOG_PATH: &str = "/tmp/workspace-server.log"; -#[cfg(windows)] -pub const DEFAULT_LOG_PATH: &str = "C:\\Windows\\Temp\\workspace-server.log"; - -/// Single-instance lock file used when no `--pid-file` is given. -#[cfg(unix)] -pub const DEFAULT_PIDFILE_PATH: &str = "/tmp/workspace-server.pid"; -#[cfg(windows)] -pub const DEFAULT_PIDFILE_PATH: &str = "C:\\Windows\\Temp\\workspace-server.pid"; - -/// Readiness marker written once the server connection is established; the control -/// plane polls it and may override the path with `--ready-file`. -#[cfg(unix)] -pub const DEFAULT_READY_PATH: &str = "/tmp/workspace-server.ready"; -#[cfg(windows)] -pub const DEFAULT_READY_PATH: &str = "C:\\Windows\\Temp\\workspace-server.ready"; - -/// How long a takeover waits for the gracefully-terminated predecessor to -/// release the pidfile lock before escalating to a forceful kill. -/// -/// Intentionally far below the server's own SIGTERM drain budget -/// (`KIGI_WORKSPACE_TERMINATION_GRACE_MS`, default 45s): a takeover only -/// happens when the orchestrator has already declared the incumbent stale, -/// so a bounded ready time for the replacement outranks completing the -/// predecessor's drain. -pub const TAKEOVER_GRACE: Duration = Duration::from_secs(2); - -/// How long a takeover waits for the lock after the forceful kill (process -/// death releases the flock) before declining. -const TAKEOVER_KILL_GRACE: Duration = Duration::from_secs(1); - -/// Poll interval while waiting for the predecessor to release the lock. -const TAKEOVER_POLL: Duration = Duration::from_millis(50); - -/// Invocation fragment identifying a pidfile holder as a workspace-server. -const WORKSPACE_SERVER_NAME_FRAGMENT: &str = "workspace-server"; - -/// Double-fork + `setsid()` into a new session, `chdir("/")`, and redirect -/// stdio (stdin ← `/dev/null`, stdout+stderr appended to `log_path`). -/// -/// Must be called before any runtime/tracing/TLS threads start (see module docs). -#[cfg(unix)] -pub fn daemonize(log_path: &Path) -> io::Result<()> { - // First fork: the launcher-tracked parent exits, orphaning the child. - fork_and_exit_parent()?; - - // New session/process group, detaching the controlling terminal. Must - // follow a fork — a process-group leader cannot call setsid(). - // SAFETY: `setsid()` takes no pointers; it only changes session membership. - if unsafe { libc::setsid() } == -1 { - return Err(io::Error::last_os_error()); - } - - // Second fork: a non-session-leader can never reacquire a controlling tty. - fork_and_exit_parent()?; - - // Detach from the launch directory (callers capture cwd beforehand). - // SAFETY: `c"/"` is a 'static, NUL-terminated string valid for the call. - if unsafe { libc::chdir(c"/".as_ptr()) } == -1 { - return Err(io::Error::last_os_error()); - } - - redirect_stdio(log_path) -} - -/// Windows daemonization: no fork/setsid (the launcher already backgrounds the -/// server) — only redirect stdout+stderr to the log file. Must run before any -/// stdout/stderr use (Rust caches the std handles on first access). The -/// single-instance lock is taken separately via [`PidFile`]. -#[cfg(windows)] -pub fn daemonize(log_path: &Path) -> io::Result<()> { - use std::os::windows::io::AsRawHandle; - - use windows::Win32::Foundation::HANDLE; - use windows::Win32::System::Console::{STD_ERROR_HANDLE, STD_OUTPUT_HANDLE, SetStdHandle}; - - if let Some(parent) = log_path.parent() { - let _ = fs::create_dir_all(parent); - } - let log = daemon_file_options() - .create(true) - .append(true) - .open(log_path)?; - - let handle = HANDLE(log.as_raw_handle()); - // SAFETY: `handle` is a live file handle owned by `log`; SetStdHandle only - // records it as the process stdout/stderr. `forget(log)` keeps it open for - // the process lifetime (the std streams reference it now). - unsafe { - SetStdHandle(STD_OUTPUT_HANDLE, handle).map_err(io::Error::other)?; - SetStdHandle(STD_ERROR_HANDLE, handle).map_err(io::Error::other)?; - } - std::mem::forget(log); - Ok(()) -} - -#[cfg(not(any(unix, windows)))] -pub fn daemonize(_log_path: &Path) -> io::Result<()> { - Err(io::Error::new( - io::ErrorKind::Unsupported, - "daemonize is only supported on Unix and Windows", - )) -} - -/// `fork()`; the parent exits 0, the child returns `Ok(())` to continue. -#[cfg(unix)] -fn fork_and_exit_parent() -> io::Result<()> { - // SAFETY: only called pre-runtime while single-threaded, so the fork - // cannot strand another thread's lock in the child. - match unsafe { libc::fork() } { - -1 => Err(io::Error::last_os_error()), - 0 => Ok(()), - _ => process::exit(0), - } -} - -/// `OpenOptions` for a daemon-owned file (log or pidfile). On Unix it adds -/// `O_NOFOLLOW` + mode `0600` as symlink/permission defense-in-depth; the -/// per-tenant sandbox namespace is the primary control. Shared with the -/// preview-proxy log (`preview_supervisor`) so both daemon-owned files get the -/// same posture. -#[cfg(unix)] -pub(crate) fn daemon_file_options() -> OpenOptions { - use std::os::unix::fs::OpenOptionsExt; - let mut opts = OpenOptions::new(); - opts.custom_flags(libc::O_NOFOLLOW).mode(0o600); - opts -} - -#[cfg(not(unix))] -pub(crate) fn daemon_file_options() -> OpenOptions { - OpenOptions::new() -} - -/// Open `/dev/null` (read) for stdin and `log_path` (created, append) for -/// stdout + stderr. -#[cfg(unix)] -fn open_stdio_targets(log_path: &Path) -> io::Result<(File, File)> { - if let Some(parent) = log_path.parent() { - let _ = fs::create_dir_all(parent); - } - let stdin_src = OpenOptions::new().read(true).open("/dev/null")?; - let log = daemon_file_options() - .create(true) - .append(true) - .open(log_path)?; - Ok((stdin_src, log)) -} - -/// `dup2(source, target)`, mapping the `-1` sentinel to an `io::Error`. -#[cfg(unix)] -fn redirect_fd(target: RawFd, source: &File) -> io::Result<()> { - // SAFETY: `source` is an open File and `target` a standard descriptor — - // both valid for `dup2`. - if unsafe { libc::dup2(source.as_raw_fd(), target) } == -1 { - return Err(io::Error::last_os_error()); - } - Ok(()) -} - -#[cfg(unix)] -fn redirect_stdio(log_path: &Path) -> io::Result<()> { - let (stdin_src, log) = open_stdio_targets(log_path)?; - redirect_fd(libc::STDIN_FILENO, &stdin_src)?; - redirect_fd(libc::STDOUT_FILENO, &log)?; - redirect_fd(libc::STDERR_FILENO, &log)?; - // `stdin_src` / `log` close here; fds 0/1/2 keep their dup'd copies. - Ok(()) -} - -/// Single-instance lock backed by an advisory `flock` on a pidfile, held for -/// the daemon's lifetime. Dropping it closes the file, releasing the lock; the -/// pidfile itself is left on disk for diagnostics. -#[derive(Debug)] -pub struct PidFile { - _file: File, -} - -impl PidFile { - /// Take the exclusive lock and record the current PID. - /// - /// - `Ok(Some(_))` — lock acquired; hold the returned guard. - /// - `Ok(None)` — another live process holds the lock (caller should - /// no-op and exit cleanly). - /// - `Err(_)` — an I/O error opening or locking the file. - pub fn acquire(path: &Path) -> io::Result> { - if let Some(parent) = path.parent() { - let _ = fs::create_dir_all(parent); - } - let mut file = daemon_file_options() - .read(true) - .write(true) - .create(true) - .truncate(false) - .open(path)?; - - match file.try_lock_exclusive() { - Ok(()) => {} - Err(e) if is_lock_contended(&e) => return Ok(None), - Err(e) => return Err(e), - } - - // PID contents are advisory diagnostics; the flock provides exclusion. - // `set_len(0)` clears any stale (possibly longer) value first. - file.set_len(0)?; - file.write_all(process::id().to_string().as_bytes())?; - file.flush()?; - - Ok(Some(Self { _file: file })) - } - - /// Acquire the lock, taking over from a live predecessor workspace-server - /// if one holds it: graceful termination (its normal drain runs), `grace` - /// to release the lock, then a forceful kill (process death releases the - /// flock). The lock is never bypassed — a guard is returned only with the - /// flock held. - /// - /// `Ok(None)` means the caller should exit quietly: the holder is not an - /// identifiable workspace-server, or the lock could not be won after the - /// escalation (e.g. a concurrent newer spawn took it). - pub fn acquire_or_take_over(path: &Path, grace: Duration) -> io::Result> { - Self::acquire_or_take_over_matching(path, grace, WORKSPACE_SERVER_NAME_FRAGMENT) - } - - /// [`Self::acquire_or_take_over`] with an injectable name fragment so - /// tests can match their own predecessor processes. - fn acquire_or_take_over_matching( - path: &Path, - grace: Duration, - name_fragment: &str, - ) -> io::Result> { - if let Some(guard) = Self::acquire(path)? { - return Ok(Some(guard)); - } - - let Some(pid) = read_pidfile_pid(path) else { - return Ok(None); - }; - if pid == process::id() { - return Ok(None); - } - let Some(predecessor) = PredecessorTarget::open(pid, name_fragment) else { - return Ok(None); - }; - - // tracing is not initialized this early; in daemonized mode stderr is - // already redirected to the log file, so eprintln! is the log channel. - eprintln!("taking over from predecessor workspace-server (pid {pid})"); - if let Err(e) = predecessor.signal(false) { - eprintln!("failed to signal predecessor (pid {pid}): {e}"); - } - if let Some(guard) = Self::poll_acquire(path, grace)? { - return Ok(Some(guard)); - } - - eprintln!("predecessor (pid {pid}) did not release the pidfile lock in time; killing it"); - if let Err(e) = predecessor.signal(true) { - eprintln!("failed to kill predecessor (pid {pid}): {e}"); - } - if let Some(guard) = Self::poll_acquire(path, TAKEOVER_KILL_GRACE)? { - return Ok(Some(guard)); - } - - // The holder we signaled is dead yet the lock is still owned — a - // concurrent newer spawn won it. Decline rather than double-run. - eprintln!("pidfile lock is still held after killing pid {pid}; exiting"); - Ok(None) - } - - /// Retry [`Self::acquire`] until it succeeds or `budget` elapses. - fn poll_acquire(path: &Path, budget: Duration) -> io::Result> { - let deadline = Instant::now() + budget; - loop { - if let Some(guard) = Self::acquire(path)? { - return Ok(Some(guard)); - } - if Instant::now() >= deadline { - return Ok(None); - } - thread::sleep(TAKEOVER_POLL); - } - } -} - -/// Advisory pid recorded in the pidfile by its holder; `None` if unreadable -/// or not a positive integer. -fn read_pidfile_pid(path: &Path) -> Option { - fs::read_to_string(path) - .ok()? - .trim() - .parse::() - .ok() - .filter(|&pid| pid > 0) -} - -/// True if the basename of `name` (path separators `/` and `\` both count) -/// contains `fragment`. Matching the basename rather than the whole path -/// keeps a directory component like `/home/workspace-server-data/foo` from -/// satisfying the kill gate. -#[cfg(any(test, target_os = "linux", windows))] -fn basename_contains(name: &str, fragment: &str) -> bool { - name.rsplit(['/', '\\']).next().is_some_and(|base| { - base.to_ascii_lowercase() - .contains(&fragment.to_ascii_lowercase()) - }) -} - -/// True if `pid`'s argv0 basename (from `/proc//cmdline`) matches -/// `fragment`. -#[cfg(target_os = "linux")] -fn process_name_matches(pid: u32, fragment: &str) -> bool { - match fs::read(format!("/proc/{pid}/cmdline")) { - Ok(cmdline) => cmdline - .split(|&b| b == 0) - .next() - .is_some_and(|argv0| basename_contains(&String::from_utf8_lossy(argv0), fragment)), - Err(_) => false, - } -} - -/// A pinned, verified handle to the predecessor process: `pidfd_open(2)` on -/// Linux, an `OpenProcess` handle on Windows. -/// -/// Pinning happens **before** verification and every signal is delivered -/// through the pin, closing the check-then-kill pid-reuse race by -/// construction: a recycled pid is unreachable — at worst a signal lands on -/// the already-dead pinned instance and is a no-op. -#[cfg(target_os = "linux")] -struct PredecessorTarget { - pid: u32, - /// `None` = pidfd unsupported on this kernel; plain-`kill` fallback mode - /// (retains only the historical residual race). - pidfd: Option, -} - -#[cfg(target_os = "linux")] -impl PredecessorTarget { - /// Pin `pid` and verify its executable basename matches `fragment`. - /// `None` if the process is gone, inaccessible, or not a match. - fn open(pid: u32, fragment: &str) -> Option { - // SAFETY: pidfd_open takes value arguments only; the returned fd is - // fresh and exclusively owned here. - let ret = unsafe { libc::syscall(libc::SYS_pidfd_open, pid as libc::pid_t, 0u32) }; - let pidfd = if ret >= 0 { - // SAFETY: `ret` is a freshly returned, unowned fd. - Some(unsafe { OwnedFd::from_raw_fd(ret as RawFd) }) - } else { - let e = io::Error::last_os_error(); - if e.raw_os_error() == Some(libc::ESRCH) { - return None; - } - // ENOSYS or seccomp-filtered: degrade to unpinned kill(). - None - }; - - // Verify after pinning: a pid recycled before the pin fails the name - // match; recycled after, the pin targets the dead predecessor. - if !process_name_matches(pid, fragment) { - return None; - } - Some(Self { pid, pidfd }) - } - - /// Deliver graceful (SIGTERM) or forceful (SIGKILL) termination to the - /// pinned instance. Already-dead is `Ok`. - fn signal(&self, forceful: bool) -> io::Result<()> { - let signal = if forceful { - libc::SIGKILL - } else { - libc::SIGTERM - }; - let ret = match &self.pidfd { - // SAFETY: the pidfd is owned and open; the siginfo pointer is - // documented-null (kernel builds a default), flags are zero. - Some(fd) => unsafe { - libc::syscall( - libc::SYS_pidfd_send_signal, - fd.as_raw_fd(), - signal, - std::ptr::null::(), - 0u32, - ) - }, - // SAFETY: kill() takes no pointers. - None => unsafe { libc::kill(self.pid as libc::pid_t, signal) }.into(), - }; - if ret == 0 { - return Ok(()); - } - match io::Error::last_os_error() { - e if e.raw_os_error() == Some(libc::ESRCH) => Ok(()), - e => Err(e), - } - } -} - -#[cfg(windows)] -struct PredecessorTarget { - handle: HANDLE, -} - -// SAFETY: the HANDLE is an owned kernel object reference; it is not tied to -// the creating thread and is only used behind &self. -#[cfg(windows)] -unsafe impl Send for PredecessorTarget {} - -#[cfg(windows)] -impl PredecessorTarget { - /// Pin `pid` with query + terminate rights and verify the image basename - /// matches `fragment` on the pinned handle. - fn open(pid: u32, fragment: &str) -> Option { - use windows::Win32::System::Threading::{ - OpenProcess, PROCESS_NAME_WIN32, PROCESS_QUERY_LIMITED_INFORMATION, PROCESS_TERMINATE, - QueryFullProcessImageNameW, - }; - use windows::core::PWSTR; - - // SAFETY: OpenProcess is FFI with value args; windows-rs returns Err - // on absence/permission failure. - let handle = unsafe { - OpenProcess( - PROCESS_QUERY_LIMITED_INFORMATION | PROCESS_TERMINATE, - false, - pid, - ) - } - .ok()?; - let target = Self { handle }; - - // QueryFullProcessImageNameW writes a NUL-terminated UTF-16 path into - // the buffer; `size` is updated to the chars written (excluding NUL). - let mut buf: Vec = vec![0; 1024]; - let mut size: u32 = buf.len() as u32; - // SAFETY: the handle is pinned by `target`; buf outlives the call; - // size is in/out. - let result = unsafe { - QueryFullProcessImageNameW( - target.handle, - PROCESS_NAME_WIN32, - PWSTR(buf.as_mut_ptr()), - &mut size, - ) - }; - if result.is_err() { - return None; - } - basename_contains(&String::from_utf16_lossy(&buf[..size as usize]), fragment) - .then_some(target) - } - - /// `TerminateProcess` on the pinned handle (Windows has no graceful - /// signal for a detached process). Already-dead is `Ok`. - fn signal(&self, _forceful: bool) -> io::Result<()> { - use windows::Win32::System::Threading::TerminateProcess; - - // SAFETY: the handle is the pinned kernel object owned by self. - match unsafe { TerminateProcess(self.handle, 0) } { - Ok(()) => Ok(()), - // Already exited: terminating a dead (but pinned) process fails - // with access-style errors; the takeover treats that as done. - Err(e) => Err(io::Error::other(format!("TerminateProcess: {e}"))), - } - } -} - -#[cfg(windows)] -impl Drop for PredecessorTarget { - fn drop(&mut self) { - use windows::Win32::Foundation::CloseHandle; - // SAFETY: the handle is owned by self and closed exactly once. - let _ = unsafe { CloseHandle(self.handle) }; - } -} - -#[cfg(not(any(target_os = "linux", windows)))] -struct PredecessorTarget; - -#[cfg(not(any(target_os = "linux", windows)))] -impl PredecessorTarget { - /// Unsupported platform: never identify a predecessor (takeover declines - /// rather than kill blind). - fn open(_pid: u32, _fragment: &str) -> Option { - None - } - - fn signal(&self, _forceful: bool) -> io::Result<()> { - Err(io::Error::new( - io::ErrorKind::Unsupported, - "process termination is only supported on Linux and Windows", - )) - } -} - -#[cfg(test)] -mod tests { - // Used only by the linux-gated predecessor-takeover tests below. - #[cfg(target_os = "linux")] - use std::process::{Child, Command, Stdio}; - - use tempfile::TempDir; - - use super::*; - - #[test] - fn pidfile_acquire_is_exclusive() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("ws.pid"); - - let first = PidFile::acquire(&path).unwrap(); - assert!(first.is_some(), "first acquire should win the lock"); - - // A second open of the same path conflicts on the advisory flock, - // even within the same process (flock is per open file description). - let second = PidFile::acquire(&path).unwrap(); - assert!(second.is_none(), "contended acquire must report None"); - - drop(first); - - // Dropping the guard closes the fd and releases the flock. Retry briefly: - // under the parallel test runner a concurrent `fork`/`Command::spawn` can - // transiently duplicate this flock'd fd, holding the lock until the child - // `execve`s (the fd is `O_CLOEXEC`). That window is microseconds, so a - // short bounded retry makes the release deterministic without weakening - // the held-exclusion assertion above. - let deadline = Instant::now() + Duration::from_secs(2); - let third = loop { - match PidFile::acquire(&path).unwrap() { - Some(guard) => break Some(guard), - None if Instant::now() < deadline => { - thread::sleep(Duration::from_millis(5)); - } - none => break none, - } - }; - assert!(third.is_some(), "acquire should succeed after release"); - } - - #[test] - fn pidfile_records_current_pid() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("ws.pid"); - - let guard = PidFile::acquire(&path).unwrap().unwrap(); - let contents = fs::read_to_string(&path).unwrap(); - assert_eq!(contents.trim().parse::().unwrap(), process::id()); - drop(guard); - } - - #[test] - fn pidfile_persists_on_disk_after_drop() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("ws.pid"); - - let guard = PidFile::acquire(&path).unwrap().unwrap(); - assert!(path.exists()); - drop(guard); - // The file is intentionally left behind for diagnostics; only the - // lock is released (re-acquirable, covered by the exclusivity test). - assert!(path.exists(), "pidfile should remain on disk after drop"); - } - - #[test] - fn pidfile_acquire_creates_parent_dir() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("nested/sub/ws.pid"); - - let guard = PidFile::acquire(&path).unwrap(); - assert!(guard.is_some()); - assert!(path.exists()); - } - - #[test] - fn pidfile_acquire_truncates_stale_longer_content() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("ws.pid"); - // A leftover value longer than our PID would leave trailing bytes if - // `set_len(0)` were missing. - fs::write(&path, "999999999999 stale junk\n").unwrap(); - - let guard = PidFile::acquire(&path).unwrap().unwrap(); - let contents = fs::read_to_string(&path).unwrap(); - assert_eq!( - contents, - process::id().to_string(), - "stale content must be fully truncated, no trailing bytes" - ); - drop(guard); - } - - #[test] - fn contended_acquire_does_not_modify_pidfile() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("ws.pid"); - - let holder = PidFile::acquire(&path).unwrap().unwrap(); - let before = fs::read_to_string(&path).unwrap(); - - let contended = PidFile::acquire(&path).unwrap(); - assert!(contended.is_none()); - - let after = fs::read_to_string(&path).unwrap(); - assert_eq!(before, after, "contended acquire must not rewrite the file"); - drop(holder); - } - - #[test] - fn pidfile_acquire_errors_on_directory() { - let dir = TempDir::new().unwrap(); - let as_dir = dir.path().join("a_dir"); - fs::create_dir(&as_dir).unwrap(); - - // Opening a directory for writing yields EISDIR — a real error that - // must surface as `Err`, never be swallowed into `Ok(None)`. - assert!( - PidFile::acquire(&as_dir).is_err(), - "acquiring a directory path must error, not report Ok(None)" - ); - } - - #[cfg(unix)] - #[test] - fn open_stdio_targets_opens_devnull_and_log() { - use std::io::{Read, Write}; - - let dir = TempDir::new().unwrap(); - let log_path = dir.path().join("logs/ws.log"); - - let (mut stdin_src, mut log) = open_stdio_targets(&log_path).unwrap(); - assert!(log_path.exists(), "log file should be created"); - - log.write_all(b"hello").unwrap(); - log.flush().unwrap(); - assert_eq!(fs::read_to_string(&log_path).unwrap(), "hello"); - - // The stdin source is /dev/null: reads yield EOF immediately. - let mut buf = [0u8; 4]; - assert_eq!(stdin_src.read(&mut buf).unwrap(), 0); - } - - #[cfg(unix)] - #[test] - fn open_stdio_targets_appends_to_existing_log() { - use std::io::Write; - - let dir = TempDir::new().unwrap(); - let log_path = dir.path().join("ws.log"); - fs::write(&log_path, "prior\n").unwrap(); - - let (_stdin_src, mut log) = open_stdio_targets(&log_path).unwrap(); - log.write_all(b"more\n").unwrap(); - log.flush().unwrap(); - - assert_eq!(fs::read_to_string(&log_path).unwrap(), "prior\nmore\n"); - } - - #[cfg(unix)] - #[test] - fn open_stdio_targets_errors_when_parent_is_a_file() { - let dir = TempDir::new().unwrap(); - let parent_file = dir.path().join("not_a_dir"); - fs::write(&parent_file, "x").unwrap(); - // `not_a_dir` is a regular file, so a log path under it is ENOTDIR. - let log_path = parent_file.join("ws.log"); - - assert!( - open_stdio_targets(&log_path).is_err(), - "a log path whose parent is a file must error" - ); - } - - // O_NOFOLLOW makes a symlinked final component fail with ELOOP rather than - // being followed — deterministic and uid-independent (no chmod, root-safe). - #[cfg(unix)] - #[test] - fn open_stdio_targets_rejects_symlinked_log() { - let dir = TempDir::new().unwrap(); - let target = dir.path().join("real.log"); - fs::write(&target, "").unwrap(); - let link = dir.path().join("link.log"); - std::os::unix::fs::symlink(&target, &link).unwrap(); - - let err = open_stdio_targets(&link).unwrap_err(); - assert_eq!(err.raw_os_error(), Some(libc::ELOOP)); - } - - #[cfg(unix)] - #[test] - fn pidfile_acquire_rejects_symlinked_path() { - let dir = TempDir::new().unwrap(); - let target = dir.path().join("real.pid"); - let link = dir.path().join("link.pid"); - std::os::unix::fs::symlink(&target, &link).unwrap(); - - let err = PidFile::acquire(&link).unwrap_err(); - assert_eq!(err.raw_os_error(), Some(libc::ELOOP)); - // The truncate-through-symlink primitive is blocked: O_CREAT did not - // follow the link to create (and truncate) its target. - assert!(!target.exists()); - } - - #[cfg(unix)] - #[test] - fn pidfile_created_mode_is_owner_only() { - use std::os::unix::fs::PermissionsExt; - let dir = TempDir::new().unwrap(); - let path = dir.path().join("ws.pid"); - - let _guard = PidFile::acquire(&path).unwrap().unwrap(); - let mode = fs::metadata(&path).unwrap().permissions().mode(); - // No group/other bits, regardless of umask (0600 & ~umask keeps them 0). - assert_eq!( - mode & 0o077, - 0, - "pidfile must not be group/other-accessible" - ); - } - - #[cfg(not(any(unix, windows)))] - #[test] - fn daemonize_unsupported_off_unix_and_windows() { - let err = daemonize(Path::new("ignored")).unwrap_err(); - assert_eq!(err.kind(), io::ErrorKind::Unsupported); - } - - #[test] - fn take_over_uncontended_acquires_normally() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("ws.pid"); - - let guard = PidFile::acquire_or_take_over(&path, Duration::from_millis(100)).unwrap(); - assert!(guard.is_some()); - assert_eq!( - fs::read_to_string(&path).unwrap(), - process::id().to_string() - ); - } - - #[test] - fn take_over_declines_unreadable_pidfile() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("ws.pid"); - - let _holder = PidFile::acquire(&path).unwrap().unwrap(); - fs::write(&path, "not a pid").unwrap(); - - let taken = - PidFile::acquire_or_take_over_matching(&path, Duration::from_millis(100), "sleep") - .unwrap(); - assert!(taken.is_none(), "an unidentifiable holder must be declined"); - } - - #[test] - fn take_over_declines_own_pid() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("ws.pid"); - - // The in-process holder wrote our own pid; a takeover must not - // signal ourselves. - let _holder = PidFile::acquire(&path).unwrap().unwrap(); - let taken = - PidFile::acquire_or_take_over_matching(&path, Duration::from_millis(100), "").unwrap(); - assert!(taken.is_none()); - } - - /// Spawn a long-sleeping child to stand in for a predecessor process. - #[cfg(target_os = "linux")] - fn spawn_predecessor() -> Child { - Command::new("sleep") - .arg("300") - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn() - .expect("spawn sleep") - } - - /// Wait (bounded) for a child to exit; returns true if it did. - #[cfg(target_os = "linux")] - fn wait_for_exit(child: &mut Child, budget: Duration) -> bool { - let deadline = Instant::now() + budget; - while Instant::now() < deadline { - if child.try_wait().expect("try_wait").is_some() { - return true; - } - thread::sleep(Duration::from_millis(10)); - } - false - } - - #[cfg(target_os = "linux")] - #[test] - fn take_over_declines_non_matching_holder() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("ws.pid"); - - let _holder = PidFile::acquire(&path).unwrap().unwrap(); - let mut child = spawn_predecessor(); - fs::write(&path, child.id().to_string()).unwrap(); - - let taken = PidFile::acquire_or_take_over_matching( - &path, - Duration::from_millis(100), - "definitely-not-this-process", - ) - .unwrap(); - assert!(taken.is_none(), "a foreign holder must not be taken over"); - assert!( - child.try_wait().expect("try_wait").is_none(), - "a foreign holder must not be killed" - ); - - child.kill().expect("cleanup kill"); - let _ = child.wait(); - } - - #[cfg(target_os = "linux")] - #[test] - fn take_over_declines_when_lock_is_never_released() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("ws.pid"); - - // The flock is held in-process for the whole test — after the child - // named in the pidfile is dead, the lock is still owned by "someone - // else" (a concurrent-spawn stand-in), so the takeover must decline - // rather than run without single-instance protection. - let _holder = PidFile::acquire(&path).unwrap().unwrap(); - let mut child = spawn_predecessor(); - let child_pid = child.id(); - fs::write(&path, child_pid.to_string()).unwrap(); - - let taken = - PidFile::acquire_or_take_over_matching(&path, Duration::from_millis(300), "sleep") - .unwrap(); - - assert!( - wait_for_exit(&mut child, Duration::from_secs(2)), - "the predecessor must be terminated" - ); - assert!( - taken.is_none(), - "a takeover that cannot win the lock must decline, never proceed lockless" - ); - assert_eq!( - fs::read_to_string(&path).unwrap(), - child_pid.to_string(), - "a declined takeover must not rewrite the pidfile" - ); - } - - #[cfg(target_os = "linux")] - #[test] - fn take_over_escalates_to_sigkill_for_stuck_predecessor() { - use std::os::unix::process::ExitStatusExt as _; - - let dir = TempDir::new().unwrap(); - let path = dir.path().join("ws.pid"); - - // A predecessor that ignores the graceful signal: only the SIGKILL - // escalation can end it. It touches a marker once the trap is - // installed so the test cannot signal it during bash startup. - let trap_ready = dir.path().join("trap-ready"); - let mut child = Command::new("bash") - .arg("-c") - .arg(format!( - "trap '' TERM; touch {}; while true; do sleep 1; done", - trap_ready.display() - )) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn() - .expect("spawn stubborn child"); - let trap_deadline = Instant::now() + Duration::from_secs(5); - while !trap_ready.exists() { - assert!(Instant::now() < trap_deadline, "child never set its trap"); - thread::sleep(Duration::from_millis(10)); - } - - // Stand in for the stuck predecessor's flock: released only after the - // graceful grace has expired, inside the post-kill window. - let holder = PidFile::acquire(&path).unwrap().unwrap(); - fs::write(&path, child.id().to_string()).unwrap(); - let release = thread::spawn(move || { - thread::sleep(Duration::from_millis(600)); - drop(holder); - }); - - let taken = - PidFile::acquire_or_take_over_matching(&path, Duration::from_millis(300), "bash") - .unwrap(); - release.join().expect("release thread"); - - let status = child.wait().expect("child wait"); - assert_eq!( - status.signal(), - Some(libc::SIGKILL), - "a SIGTERM-immune predecessor must be ended by the SIGKILL escalation" - ); - assert!(taken.is_some(), "the lock freed within the kill window"); - assert_eq!( - fs::read_to_string(&path).unwrap(), - process::id().to_string() - ); - } - - #[cfg(target_os = "linux")] - #[test] - fn take_over_acquires_cleanly_when_predecessor_releases() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("ws.pid"); - - let holder = PidFile::acquire(&path).unwrap().unwrap(); - let mut child = spawn_predecessor(); - fs::write(&path, child.id().to_string()).unwrap(); - - // Release the lock shortly after the takeover starts waiting, - // simulating the predecessor finishing its drain within grace. - let release = thread::spawn(move || { - thread::sleep(Duration::from_millis(100)); - drop(holder); - }); - - let taken = - PidFile::acquire_or_take_over_matching(&path, Duration::from_secs(5), "sleep").unwrap(); - release.join().expect("release thread"); - - assert!( - wait_for_exit(&mut child, Duration::from_secs(2)), - "the predecessor must be terminated" - ); - assert!(taken.is_some()); - assert_eq!( - fs::read_to_string(&path).unwrap(), - process::id().to_string() - ); - } - - #[cfg(target_os = "linux")] - #[test] - fn process_name_matches_own_argv0() { - let pid = process::id(); - // Derive the fragment from this process's real argv0 basename rather - // than hardcoding a name: different test runners name the binary - // differently (e.g. Cargo uses `kigi_workspace-`), so a - // hardcoded fragment matches under one runner but not another. - let cmdline = fs::read(format!("/proc/{pid}/cmdline")).expect("read own cmdline"); - let argv0 = cmdline.split(|&b| b == 0).next().expect("argv0 present"); - let basename = String::from_utf8_lossy(argv0) - .rsplit(['/', '\\']) - .next() - .expect("basename") - .to_owned(); - assert!(!basename.is_empty(), "argv0 basename must not be empty"); - assert!(process_name_matches(pid, &basename)); - assert!(!process_name_matches(pid, "definitely-not-this-process")); - } - - #[cfg(target_os = "linux")] - #[test] - fn predecessor_target_pins_verifies_and_signals() { - let mut child = spawn_predecessor(); - - assert!( - PredecessorTarget::open(child.id(), "not-a-match").is_none(), - "a non-matching name must not produce a target" - ); - - // /proc//cmdline can lag briefly after spawn under remote CI - // executors; derive the fragment from the live cmdline (handles - // busybox-as-sleep) and retry pin open instead of a one-shot expect. - let fragment = { - let deadline = Instant::now() + Duration::from_secs(2); - loop { - if let Ok(cmdline) = fs::read(format!("/proc/{}/cmdline", child.id())) { - let argv0 = cmdline.split(|&b| b == 0).next().unwrap_or_default(); - let basename = String::from_utf8_lossy(argv0) - .rsplit(['/', '\\']) - .next() - .unwrap_or("") - .to_owned(); - if !basename.is_empty() { - break basename; - } - } - if Instant::now() >= deadline { - panic!("child cmdline never became readable"); - } - thread::sleep(Duration::from_millis(10)); - } - }; - - let target = { - let deadline = Instant::now() + Duration::from_secs(2); - loop { - if let Some(t) = PredecessorTarget::open(child.id(), &fragment) { - break t; - } - if Instant::now() >= deadline { - panic!("pin child (fragment={fragment:?})"); - } - thread::sleep(Duration::from_millis(10)); - } - }; - target.signal(false).expect("graceful signal"); - assert!( - wait_for_exit(&mut child, Duration::from_secs(2)), - "the pinned child must receive the signal" - ); - - // Signalling the dead pinned instance is a no-op, never a stray kill. - target - .signal(true) - .expect("signal to dead pinned instance is Ok"); - } - - #[test] - fn basename_contains_ignores_directory_components() { - assert!(basename_contains( - "/usr/local/bin/kigi-workspace-server", - "workspace-server" - )); - assert!(basename_contains( - "C:\\Program Files\\XAI-Workspace-Server.exe", - "workspace-server" - )); - assert!( - !basename_contains("/home/workspace-server-data/unrelated", "workspace-server"), - "a matching directory component must not satisfy the kill gate" - ); - } - - #[test] - fn read_pidfile_pid_parses_and_rejects() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("ws.pid"); - - fs::write(&path, "1234\n").unwrap(); - assert_eq!(read_pidfile_pid(&path), Some(1234)); - - fs::write(&path, "0").unwrap(); - assert_eq!(read_pidfile_pid(&path), None, "pid 0 is not a process"); - - fs::write(&path, "garbage").unwrap(); - assert_eq!(read_pidfile_pid(&path), None); - - assert_eq!(read_pidfile_pid(&dir.path().join("missing")), None); - } -} diff --git a/crates/codegen/kigi-workspace/src/diag_server.rs b/crates/codegen/kigi-workspace/src/diag_server.rs deleted file mode 100644 index c4ac5d6..0000000 --- a/crates/codegen/kigi-workspace/src/diag_server.rs +++ /dev/null @@ -1,561 +0,0 @@ -//! In-guest diagnostics HTTP server (`/ready`, `/statusz`, `/logs`) for the -//! standalone workspace-server. -//! -//! The surface is reachable by any process inside the user's own sandbox -//! (loopback-only TCP, or a 0600 Unix socket) and is never exposed through -//! the sandbox port mapping. `/logs` returns the raw daemon log: treat its -//! output as sensitive and keep the log stream free of secrets. - -use std::io::{self, Read as _, Seek as _, SeekFrom}; -use std::net::Ipv4Addr; -use std::path::{Path, PathBuf}; -use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; -use std::time::{SystemTime, UNIX_EPOCH}; -use std::{env, fs, process}; - -use anyhow::anyhow; -use axum::Router; -use axum::extract::{Query, State}; -use axum::http::{StatusCode, header}; -use axum::response::{IntoResponse, Response}; -use axum::routing::get; -use serde::{Deserialize, Serialize}; -use tokio::net::TcpListener; -#[cfg(unix)] -use tokio::net::UnixListener; -use tokio::task::JoinHandle; - -/// Default Unix socket path (resolves host-visibly under the production -/// launcher's `/tmp` bind mount, next to the ready/pid files). -#[cfg(unix)] -pub const DEFAULT_DIAG_SOCKET_PATH: &str = "/tmp/workspace-server.sock"; - -/// Default loopback TCP port for Windows guests. -pub const DEFAULT_DIAG_PORT: u16 = 6016; - -/// Grep-able daemon-log marker for a diagnostics bind failure. -pub const DIAG_BIND_FAILED_MARKER: &str = "diagnostics server bind failed"; - -/// Process exit code for a fatal diagnostics bind failure in `--daemonize` mode. -pub const EXIT_DIAG_BIND_FAILED: i32 = 5; - -/// Default `/logs` tail size when `tail_bytes` is not given. -pub const DEFAULT_LOG_TAIL_BYTES: u64 = 64 * 1024; - -/// Hard cap on a `/logs` response; larger `tail_bytes` values are clamped. -pub const MAX_LOG_TAIL_BYTES: u64 = 256 * 1024; - -/// Hub connection state as reported on `/ready`. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] -#[serde(rename_all = "lowercase")] -pub enum DiagState { - Starting, - Connected, - Disconnected, -} - -/// Response body for `/ready`. The field set is a frozen contract with the -/// sandbox readiness gate: never rename or remove fields; additions are -/// backward-compatible. -#[derive(Debug, Serialize)] -struct ReadyBody { - /// Serialized as an explicit `null` (never omitted) for nonce-less - /// launches. - launch_id: Option, - state: DiagState, - pid: u32, - connected_at: Option, - state_changed_at: u64, - version: &'static str, -} - -/// Response body for `/statusz`: the `/ready` fields plus debug extras. -#[derive(Debug, Serialize)] -struct StatuszBody { - #[serde(flatten)] - ready: ReadyBody, - os: &'static str, -} - -#[derive(Debug)] -struct Inner { - state: DiagState, - connected_at: Option, - state_changed_at: u64, - shutting_down: bool, -} - -/// Cloneable handle publishing hub lifecycle transitions to the server. -#[derive(Debug, Clone)] -pub struct DiagHandle { - launch_id: Option, - inner: Arc>, -} - -impl DiagHandle { - /// `launch_id` is the caller-minted per-spawn nonce, echoed verbatim on - /// `/ready` (`null` for nonce-less local launches). - pub fn new(launch_id: Option) -> Self { - Self { - launch_id, - inner: Arc::new(Mutex::new(Inner { - state: DiagState::Starting, - connected_at: None, - state_changed_at: now_ms(), - shutting_down: false, - })), - } - } - - /// Initial hello completed, or a reconnect's serve replay settled. - /// Ignored after [`Self::set_shutting_down`]: a reconnect that settles - /// during the shutdown drain must not republish `connected`. - pub fn set_connected(&self) { - let mut inner = self.lock(); - if inner.shutting_down { - return; - } - inner.state = DiagState::Connected; - let now = now_ms(); - inner.connected_at.get_or_insert(now); - inner.state_changed_at = now; - } - - /// Server socket dropped. - pub fn set_disconnected(&self) { - let mut inner = self.lock(); - inner.state = DiagState::Disconnected; - inner.state_changed_at = now_ms(); - } - - /// Latch `disconnected` for process shutdown: reported as `disconnected` - /// on `/ready`, and later `set_connected` calls become no-ops. - pub fn set_shutting_down(&self) { - let mut inner = self.lock(); - inner.shutting_down = true; - inner.state = DiagState::Disconnected; - inner.state_changed_at = now_ms(); - } - - /// True after [`Self::set_shutting_down`]. - pub fn is_shutting_down(&self) -> bool { - self.lock().shutting_down - } - - fn lock(&self) -> MutexGuard<'_, Inner> { - self.inner.lock().unwrap_or_else(PoisonError::into_inner) - } - - fn ready_body(&self) -> ReadyBody { - let inner = self.lock(); - ReadyBody { - launch_id: self.launch_id.clone(), - state: inner.state, - pid: process::id(), - connected_at: inner.connected_at, - state_changed_at: inner.state_changed_at, - version: kigi_version::VERSION, - } - } - - fn statusz_body(&self) -> StatuszBody { - StatuszBody { - ready: self.ready_body(), - os: env::consts::OS, - } - } -} - -fn now_ms() -> u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0) -} - -/// Where the diagnostics server listens: a Unix socket on Linux, loopback TCP -/// on Windows. Both variants compile everywhere so the TCP path is testable -/// on Linux. -#[derive(Debug, Clone)] -pub enum DiagListener { - #[cfg(unix)] - Unix(PathBuf), - Tcp(u16), -} - -/// Shared request state: the lifecycle handle plus the daemon log path -/// (`None` when logs go to a terminal instead of a file — `/logs` is 404). -#[derive(Debug, Clone)] -struct DiagContext { - handle: DiagHandle, - log_file: Option>, -} - -#[derive(Debug, Deserialize)] -struct LogsQuery { - tail_bytes: Option, -} - -async fn logs(State(ctx): State, Query(query): Query) -> Response { - let Some(path) = ctx.log_file else { - return StatusCode::NOT_FOUND.into_response(); - }; - let tail = query - .tail_bytes - .unwrap_or(DEFAULT_LOG_TAIL_BYTES) - .min(MAX_LOG_TAIL_BYTES); - match tokio::task::spawn_blocking(move || tail_file(&path, tail)).await { - Ok(Ok(bytes)) => ( - [(header::CONTENT_TYPE, "text/plain; charset=utf-8")], - String::from_utf8_lossy(&bytes).into_owned(), - ) - .into_response(), - Ok(Err(e)) if e.kind() == io::ErrorKind::NotFound => StatusCode::NOT_FOUND.into_response(), - // Generic body: an io::Error would echo the log-file path to clients. - Ok(Err(e)) => { - tracing::warn!(error = %e, "failed to read log tail"); - (StatusCode::INTERNAL_SERVER_ERROR, "failed to read log").into_response() - } - Err(e) => { - tracing::warn!(error = %e, "log tail task failed"); - (StatusCode::INTERNAL_SERVER_ERROR, "failed to read log").into_response() - } - } -} - -/// Read at most the last `max` bytes of `path`. -fn tail_file(path: &Path, max: u64) -> io::Result> { - let mut file = fs::File::open(path)?; - let len = file.metadata()?.len(); - file.seek(SeekFrom::Start(len.saturating_sub(max)))?; - let mut buf = Vec::new(); - // `take` bounds the read even if the file grows underneath us. - file.take(max).read_to_end(&mut buf)?; - Ok(buf) -} - -fn router(ctx: DiagContext) -> Router { - Router::new() - .route( - "/ready", - get(|State(ctx): State| async move { - let body = ctx.handle.ready_body(); - // Non-2xx for "not ready" so naive HTTP probes agree with - // consumers that parse `state`. The body is served either way. - let status = if body.state == DiagState::Connected { - StatusCode::OK - } else { - StatusCode::SERVICE_UNAVAILABLE - }; - (status, axum::Json(body)) - }), - ) - .route( - "/statusz", - get(|State(ctx): State| async move { axum::Json(ctx.handle.statusz_body()) }), - ) - .route("/logs", get(logs)) - .with_state(ctx) -} - -/// Bind the listener and spawn the server task. Binding happens before this -/// returns, so a bind failure surfaces synchronously. `log_file` is the -/// daemon log served by `/logs` (`None` ⇒ `/logs` is 404). -pub async fn serve( - listener: DiagListener, - handle: DiagHandle, - log_file: Option, -) -> anyhow::Result { - let ctx = DiagContext { - handle, - log_file: log_file.map(Arc::new), - }; - match listener { - #[cfg(unix)] - DiagListener::Unix(path) => { - let _ = fs::remove_file(&path); - let listener = - UnixListener::bind(&path).map_err(|e| anyhow!("bind {}: {e}", path.display()))?; - use std::os::unix::fs::PermissionsExt as _; - if let Err(e) = fs::set_permissions(&path, fs::Permissions::from_mode(0o600)) { - tracing::warn!( - path = %path.display(), - error = %e, - "failed to restrict diagnostics socket permissions" - ); - } - let task = tokio::spawn(async move { - if let Err(e) = axum::serve(listener, router(ctx)).await { - tracing::warn!(error = %e, "diagnostics server exited"); - } - }); - Ok(BoundDiag { - addr: format!("unix:{}", path.display()), - port: None, - task, - }) - } - DiagListener::Tcp(port) => { - let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, port)) - .await - .map_err(|e| anyhow!("bind 127.0.0.1:{port}: {e}"))?; - let local = listener.local_addr()?; - let task = tokio::spawn(async move { - if let Err(e) = axum::serve(listener, router(ctx)).await { - tracing::warn!(error = %e, "diagnostics server exited"); - } - }); - Ok(BoundDiag { - addr: format!("http://{local}"), - port: Some(local.port()), - task, - }) - } - } -} - -/// A successfully bound diagnostics server. -#[derive(Debug)] -pub struct BoundDiag { - /// Human-readable bound address for the startup log line. - pub addr: String, - /// Bound TCP port (`None` for Unix sockets). - pub port: Option, - /// The serve task; held by the production launcher for the process lifetime. - pub task: JoinHandle<()>, -} - -#[cfg(test)] -mod tests { - use serde_json::Value; - - use super::*; - - async fn get_json(port: u16, path: &str) -> (u16, Value) { - let response = reqwest::get(format!("http://127.0.0.1:{port}{path}")) - .await - .expect("request"); - let status = response.status().as_u16(); - let body = response.text().await.expect("body"); - (status, serde_json::from_str(&body).expect("json body")) - } - - #[tokio::test] - async fn ready_response_contract_is_frozen() { - let handle = DiagHandle::new(Some("nonce-1".to_owned())); - let bound = serve(DiagListener::Tcp(0), handle, None) - .await - .expect("bind"); - let (status, body) = get_json(bound.port.expect("tcp port"), "/ready").await; - - assert_eq!(status, 503, "not yet connected must not probe as ready"); - let obj = body.as_object().expect("object"); - for key in [ - "launch_id", - "state", - "pid", - "connected_at", - "state_changed_at", - "version", - ] { - assert!(obj.contains_key(key), "missing frozen key {key}"); - } - assert_eq!(body["launch_id"], "nonce-1"); - assert_eq!(body["state"], "starting"); - assert_eq!(body["connected_at"], Value::Null); - } - - #[tokio::test] - async fn state_follows_hub_lifecycle_and_freezes_connected_at() { - let handle = DiagHandle::new(None); - let bound = serve(DiagListener::Tcp(0), handle.clone(), None) - .await - .expect("bind"); - let port = bound.port.expect("tcp port"); - - handle.set_connected(); - let (status, connected) = get_json(port, "/ready").await; - assert_eq!(status, 200); - assert_eq!(connected["state"], "connected"); - assert!(connected["connected_at"].is_u64()); - assert_eq!(connected["launch_id"], Value::Null); - - handle.set_disconnected(); - let (status, disconnected) = get_json(port, "/ready").await; - assert_eq!(status, 503); - assert_eq!(disconnected["state"], "disconnected"); - assert_eq!( - disconnected["connected_at"], connected["connected_at"], - "connected_at is frozen at first connect and echoed on disconnect" - ); - - handle.set_connected(); - let (status, reconnected) = get_json(port, "/ready").await; - assert_eq!(status, 200); - assert_eq!(reconnected["state"], "connected"); - assert_eq!( - reconnected["connected_at"], connected["connected_at"], - "reconnect must not re-mint connected_at" - ); - } - - #[tokio::test] - async fn shutting_down_latches_disconnected_across_reconnects() { - let handle = DiagHandle::new(None); - let bound = serve(DiagListener::Tcp(0), handle.clone(), None) - .await - .expect("bind"); - let port = bound.port.expect("tcp port"); - - handle.set_connected(); - handle.set_shutting_down(); - // A reconnect settling during the shutdown drain must not republish - // `connected`. - handle.set_connected(); - - let (status, body) = get_json(port, "/ready").await; - assert_eq!(status, 503); - assert_eq!(body["state"], "disconnected"); - assert!(handle.is_shutting_down()); - } - - #[cfg(unix)] - #[tokio::test] - async fn unix_socket_serves_ready_and_rebinds_over_stale_socket() { - let dir = tempfile::tempdir().expect("tempdir"); - let sock = dir.path().join("ws.sock"); - - drop(std::os::unix::net::UnixListener::bind(&sock).expect("stale bind")); - - let handle = DiagHandle::new(Some("nonce-uds".to_owned())); - let _bound = serve(DiagListener::Unix(sock.clone()), handle.clone(), None) - .await - .expect("bind over stale socket"); - handle.set_connected(); - - use std::os::unix::fs::PermissionsExt as _; - let mode = fs::metadata(&sock) - .expect("socket meta") - .permissions() - .mode(); - assert_eq!(mode & 0o777, 0o600, "socket must be owner-only"); - - use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; - let mut stream = tokio::net::UnixStream::connect(&sock) - .await - .expect("connect"); - stream - .write_all(b"GET /ready HTTP/1.1\r\nHost: ws\r\nConnection: close\r\n\r\n") - .await - .expect("write request"); - let mut response = Vec::new(); - stream - .read_to_end(&mut response) - .await - .expect("read response"); - let response = String::from_utf8_lossy(&response); - assert!(response.starts_with("HTTP/1.1 200"), "got: {response}"); - let body = response.split("\r\n\r\n").nth(1).expect("body"); - let json_start = body.find('{').expect("json start"); - let json_end = body.rfind('}').expect("json end"); - let parsed: Value = serde_json::from_str(&body[json_start..=json_end]).expect("json"); - assert_eq!(parsed["launch_id"], "nonce-uds"); - assert_eq!(parsed["state"], "connected"); - } - - #[tokio::test] - async fn tcp_bind_conflict_surfaces_as_error() { - let first = serve(DiagListener::Tcp(0), DiagHandle::new(None), None) - .await - .expect("first bind"); - let port = first.port.expect("tcp port"); - let err = serve(DiagListener::Tcp(port), DiagHandle::new(None), None).await; - assert!(err.is_err(), "second bind on the same port must fail"); - } - - async fn get_text(port: u16, path: &str) -> (u16, Option, String) { - let response = reqwest::get(format!("http://127.0.0.1:{port}{path}")) - .await - .expect("request"); - let status = response.status().as_u16(); - let content_type = response - .headers() - .get(reqwest::header::CONTENT_TYPE) - .map(|v| v.to_str().expect("content-type").to_owned()); - (status, content_type, response.text().await.expect("body")) - } - - async fn serve_with_log(log_file: Option) -> u16 { - let bound = serve(DiagListener::Tcp(0), DiagHandle::new(None), log_file) - .await - .expect("bind"); - bound.port.expect("tcp port") - } - - #[tokio::test] - async fn logs_tails_requested_bytes_as_plain_text() { - let dir = tempfile::tempdir().expect("tempdir"); - let log = dir.path().join("ws.log"); - fs::write(&log, "0123456789").expect("write log"); - let port = serve_with_log(Some(log)).await; - - let (status, content_type, body) = get_text(port, "/logs?tail_bytes=4").await; - assert_eq!(status, 200); - assert_eq!(content_type.as_deref(), Some("text/plain; charset=utf-8")); - assert_eq!(body, "6789"); - - // A tail larger than the file returns the whole file. - let (status, _, body) = get_text(port, "/logs?tail_bytes=1000").await; - assert_eq!(status, 200); - assert_eq!(body, "0123456789"); - } - - #[tokio::test] - async fn logs_default_and_hard_cap_bound_the_response() { - let dir = tempfile::tempdir().expect("tempdir"); - let log = dir.path().join("ws.log"); - // Larger than the hard cap; ends with a marker to prove we got the tail. - let mut content = vec![b'a'; (MAX_LOG_TAIL_BYTES + 4096) as usize]; - content.extend_from_slice(b"END-MARKER"); - fs::write(&log, &content).expect("write log"); - let port = serve_with_log(Some(log)).await; - - let (status, _, body) = get_text(port, "/logs").await; - assert_eq!(status, 200); - assert_eq!(body.len() as u64, DEFAULT_LOG_TAIL_BYTES); - assert!(body.ends_with("END-MARKER")); - - let (status, _, body) = get_text(port, "/logs?tail_bytes=999999999").await; - assert_eq!(status, 200); - assert_eq!(body.len() as u64, MAX_LOG_TAIL_BYTES, "hard cap applies"); - assert!(body.ends_with("END-MARKER")); - } - - #[tokio::test] - async fn logs_tail_is_lossy_utf8() { - let dir = tempfile::tempdir().expect("tempdir"); - let log = dir.path().join("ws.log"); - // 'é' is 0xC3 0xA9; a 1-byte tail cuts the sequence mid-char. - fs::write(&log, "aé").expect("write log"); - let port = serve_with_log(Some(log)).await; - - let (status, _, body) = get_text(port, "/logs?tail_bytes=1").await; - assert_eq!(status, 200); - assert_eq!(body, "\u{FFFD}", "a torn UTF-8 boundary must be lossy"); - } - - #[tokio::test] - async fn logs_without_log_file_is_404() { - let port = serve_with_log(None).await; - let (status, _, _) = get_text(port, "/logs").await; - assert_eq!(status, 404); - } - - #[tokio::test] - async fn logs_missing_log_file_is_404() { - let dir = tempfile::tempdir().expect("tempdir"); - let port = serve_with_log(Some(dir.path().join("never-created.log"))).await; - let (status, _, _) = get_text(port, "/logs").await; - assert_eq!(status, 404); - } -} diff --git a/crates/codegen/kigi-workspace/src/file_system/client_fs.rs b/crates/codegen/kigi-workspace/src/file_system/client_fs.rs deleted file mode 100644 index 82cb6df..0000000 --- a/crates/codegen/kigi-workspace/src/file_system/client_fs.rs +++ /dev/null @@ -1,775 +0,0 @@ -//! Read-only filesystem helpers backing the client-facing -//! `workspace.client_fs_*` RPCs (the grok.com conversation-files UI, -//! tunneled through the server). -//! -//! Deliberately separate from the shell-facing ext ops in -//! [`ext_fs`](super::ext_fs): every path here is workspace-root-relative -//! and resolves through the root-confinement helper -//! (`WorkspaceHandle::resolve_service_path`), the list walk excludes -//! symlinks that resolve outside the root (and never descends into -//! them), listings paginate with stable post-sort slices, and reads are -//! binary-safe (base64 chunks). -//! -//! Wire types live in `kigi_workspace_types::rpc::fs` (the -//! `ClientFs*` types), shared with the backend caller. - -use std::collections::HashMap; -use std::path::{Path, PathBuf}; - -use kigi_workspace_types::rpc::fs::{ - ClientFsListNode as FsListNode, ClientFsListReq as FsListReq, ClientFsListRes as FsListRes, - ClientFsReadFileReq as FsReadFileReq, ClientFsReadFileRes as FsReadFileRes, - ClientFsStatReq as FsStatReq, ClientFsStatRes as FsStatRes, FsContentType, FsNodeType, -}; - -use crate::error::{WorkspaceError, WorkspaceResult}; -use crate::handle::WorkspaceHandle; - -/// Hard cap on entries collected per list call before sorting (shared -/// across all fs surfaces; see [`super::walk::MAX_LIST_COLLECT`]). -const MAX_LIST_COLLECT: usize = super::walk::MAX_LIST_COLLECT; - -/// Server-side cap on `FsListReq::limit`. -const MAX_LIST_LIMIT: u32 = 1000; - -/// Server-side cap on a single read's effective byte budget (shared -/// across all fs surfaces; see [`super::walk::MAX_READ_BYTES`]). Only -/// referenced by tests now that the clamp lives in `walk::clamp_read_length`. -#[cfg(test)] -const MAX_READ_BYTES: u64 = super::walk::MAX_READ_BYTES; - -/// Bound on memoized hashes; the memo is cleared (not LRU-evicted) when -/// full — entries simply re-hash on next use. -const HASH_MEMO_CAPACITY: usize = 4096; - -// ========================================================================= -// (path, size, mtime_ms) → hash memo -// ========================================================================= - -#[derive(Debug, Clone)] -struct MemoEntry { - size: u64, - mtime_ms: i64, - hash: String, -} - -/// Memo of full-content SHA-256 digests keyed by absolute path and -/// validated against `(size, mtime_ms)`, so unchanged files hash once -/// instead of on every `client_fs_stat`. The memo only avoids redundant -/// hashing — it never substitutes mtime for content addressing: a -/// `(size, mtime_ms)` mismatch is a miss and the caller re-hashes. -#[derive(Debug, Default)] -pub(crate) struct FileHashMemo { - entries: parking_lot::Mutex>, -} - -impl FileHashMemo { - /// Return the memoized hash when `(size, mtime_ms)` still match. - pub(crate) fn lookup(&self, path: &Path, size: u64, mtime_ms: i64) -> Option { - let entries = self.entries.lock(); - let entry = entries.get(path)?; - (entry.size == size && entry.mtime_ms == mtime_ms).then(|| entry.hash.clone()) - } - - /// Record a freshly computed hash, replacing any stale entry for the - /// same path. Clears the whole memo when inserting a new path would - /// exceed [`HASH_MEMO_CAPACITY`]. - pub(crate) fn store(&self, path: &Path, size: u64, mtime_ms: i64, hash: String) { - let mut entries = self.entries.lock(); - if !entries.contains_key(path) && entries.len() >= HASH_MEMO_CAPACITY { - entries.clear(); - } - entries.insert( - path.to_path_buf(), - MemoEntry { - size, - mtime_ms, - hash, - }, - ); - } -} - -// ========================================================================= -// Path resolution -// ========================================================================= - -/// Resolve a root-relative request path through the workspace's -/// root-confinement helper, returning the resolved path together with the -/// canonical root it was checked against. `""` and `"."` mean the -/// workspace root; absolute paths, `..` escapes, and symlink escapes are -/// rejected there. -async fn resolve_with_root( - ws: &WorkspaceHandle, - path: &str, -) -> WorkspaceResult<(PathBuf, PathBuf)> { - let rel = if path.is_empty() { "." } else { path }; - let canonical_root = ws.canonical_root().await?; - let abs = ws.resolve_service_path(rel, &canonical_root).await?; - Ok((abs, canonical_root)) -} - -/// [`resolve_with_root`] for callers that don't need the canonical root. -async fn resolve(ws: &WorkspaceHandle, path: &str) -> WorkspaceResult { - resolve_with_root(ws, path).await.map(|(abs, _)| abs) -} - -fn system_time_ms(st: std::time::SystemTime) -> i64 { - match st.duration_since(std::time::UNIX_EPOCH) { - Ok(d) => i64::try_from(d.as_millis()).unwrap_or(i64::MAX), - Err(e) => -i64::try_from(e.duration().as_millis()).unwrap_or(i64::MAX), - } -} - -// ========================================================================= -// list -// ========================================================================= - -/// List `req.path` (workspace-root-relative) with stable pagination: -/// collect the full walk (bounded by [`MAX_LIST_COLLECT`]), sort -/// directories-first / case-insensitive by name, then slice -/// `[offset, offset + limit)`. Symlinks resolving outside the workspace -/// root are excluded from the walk (and never descended into). -pub(crate) async fn list(ws: &WorkspaceHandle, req: &FsListReq) -> WorkspaceResult { - let (abs, canonical_root) = resolve_with_root(ws, &req.path).await?; - let root = ws.root_cwd()?; - let req = req.clone(); - // The walk does synchronous traversal + metadata syscalls; run it off - // the async executor (matching the ext_fs ops). - tokio::task::spawn_blocking(move || { - list_blocking(&abs, &root, &canonical_root, &req, MAX_LIST_COLLECT) - }) - .await - .map_err(|e| WorkspaceError::JoinError(e.to_string()))? -} - -fn list_blocking( - abs_dir: &Path, - root: &Path, - canonical_root: &Path, - req: &FsListReq, - max_collect: usize, -) -> WorkspaceResult { - // Root confinement also holds mid-walk: a symlink inside the root - // pointing outside must not enumerate outside metadata. - let page = super::walk::list_directory_paged( - abs_dir, - super::walk::ListOptions { - depth: req.depth as usize, - follow_symlinks: req.follow_symlinks, - respect_git_ignore: req.respect_git_ignore, - include_hidden: req.include_hidden, - include_globs: &req.include_globs, - exclude_globs: &req.exclude_globs, - offset: req.offset, - limit: req.limit.min(MAX_LIST_LIMIT) as usize, - confine_to_canonical_root: Some(canonical_root.to_path_buf()), - }, - max_collect, - ); - - let nodes: Vec = page - .entries - .into_iter() - .map(|e| FsListNode { - node_type: if e.is_dir { - FsNodeType::Directory - } else { - FsNodeType::File - }, - size: e.size, - mtime_ms: e.modified.map(system_time_ms), - is_symlink: e.is_symlink.then_some(true), - // Root-relative path (divergent from the shell's absolute path). - // A walk under a symlinked root yields canonical-root-spelled - // entries, so strip either spelling. - path: e - .abs_path - .strip_prefix(root) - .or_else(|_| e.abs_path.strip_prefix(canonical_root)) - .unwrap_or(&e.abs_path) - .to_string_lossy() - .into_owned(), - name: e.name, - }) - .collect(); - - Ok(FsListRes { - nodes, - truncated: page.truncated, - }) -} - -// ========================================================================= -// stat -// ========================================================================= - -/// Stat `req.path`: existence, kind, size, mtime, and — for files — a -/// full-content SHA-256 served through the workspace hash memo. -pub(crate) async fn stat(ws: &WorkspaceHandle, req: &FsStatReq) -> WorkspaceResult { - let abs = resolve(ws, &req.path).await?; - let md = match tokio::fs::metadata(&abs).await { - Ok(md) => md, - // NotADirectory: a *file* sits mid-path (e.g. `a.txt/sub`) — for an - // existence probe that is a miss, not an RPC error. - Err(e) - if e.kind() == std::io::ErrorKind::NotFound - || e.kind() == std::io::ErrorKind::NotADirectory => - { - return Ok(FsStatRes { - exists: false, - node_type: None, - size: None, - mtime_ms: None, - hash: None, - }); - } - Err(e) => { - return Err(WorkspaceError::HubError(format!( - "stat failed for {}: {e}", - req.path - ))); - } - }; - let mtime_ms = md.modified().ok().map(system_time_ms); - if md.is_dir() { - return Ok(FsStatRes { - exists: true, - node_type: Some(FsNodeType::Directory), - size: None, - mtime_ms, - hash: None, - }); - } - let size = md.len(); - let memo = &ws.shared.client_fs_hash_memo; - let hash = match mtime_ms.and_then(|m| memo.lookup(&abs, size, m)) { - Some(hash) => hash, - None => { - let (hash, _, _) = crate::handle::stream_hash_and_range(&abs, 0, 0) - .await - .map_err(|e| { - WorkspaceError::HubError(format!("hash failed for {}: {e}", req.path)) - })?; - if let Some(m) = mtime_ms { - memo.store(&abs, size, m, hash.clone()); - } - hash - } - }; - Ok(FsStatRes { - exists: true, - node_type: Some(FsNodeType::File), - size: Some(size), - mtime_ms, - hash: Some(hash), - }) -} - -// ========================================================================= -// read_file -// ========================================================================= - -/// Read a byte range of `req.path` (binary-safe, capped at -/// `min(req.max_bytes, MAX_READ_BYTES)`) together with the full-file -/// SHA-256. When the hash is memoized for the current `(size, mtime)` -/// only the requested range is read; otherwise the whole file streams -/// once (via the shared [`crate::handle::stream_hash_and_range`]) to -/// hash it. -pub(crate) async fn read_file( - ws: &WorkspaceHandle, - req: &FsReadFileReq, -) -> WorkspaceResult { - let abs = resolve(ws, &req.path).await?; - let read_err = - |e: std::io::Error| WorkspaceError::HubError(format!("read failed for {}: {e}", req.path)); - let md = tokio::fs::metadata(&abs).await.map_err(read_err)?; - if md.is_dir() { - return Err(WorkspaceError::HubError(format!( - "not a file: {}", - req.path - ))); - } - let size = md.len(); - let mtime_ms = md.modified().ok().map(system_time_ms); - let offset = req.offset.unwrap_or(0); - // Server-side clamp: a hostile/buggy caller cannot lift the per-chunk - // budget past MAX_READ_BYTES regardless of `maxBytes`. - let length = super::walk::clamp_read_length(req.length, req.max_bytes); - - let memo = &ws.shared.client_fs_hash_memo; - let (hash, chunk, size) = match mtime_ms.and_then(|m| memo.lookup(&abs, size, m)) { - Some(hash) => { - let chunk = super::walk::read_range(&abs, offset, length) - .await - .map_err(read_err)?; - (hash, chunk, size) - } - None => { - let (hash, chunk, streamed) = - crate::handle::stream_hash_and_range(&abs, offset, length) - .await - .map_err(read_err)?; - if let Some(m) = mtime_ms { - memo.store(&abs, streamed, m, hash.clone()); - } - (hash, chunk, streamed) - } - }; - - // Shared encoder keeps the paired wire fields coherent with one UTF-8 - // validation pass; `type` is text iff the bytes were valid UTF-8. - let (payload, is_text) = super::walk::encode_chunk(chunk, req.encoding); - let (content, content_base64) = match payload { - super::walk::ChunkPayload::Text(t) => (Some(t), None), - super::walk::ChunkPayload::Base64(b) => (None, Some(b)), - }; - let content_type = if is_text { - FsContentType::Text - } else { - FsContentType::Binary - }; - Ok(FsReadFileRes { - content, - content_base64, - size, - hash, - content_type, - }) -} - -#[cfg(test)] -mod tests { - use base64::Engine; - use kigi_workspace_types::rpc::fs::FsReadEncoding; - - use super::*; - use crate::handle::tests::make_handle; - - fn list_req(path: &str) -> FsListReq { - FsListReq { - path: path.to_owned(), - depth: 1, - include_hidden: true, - limit: 1000, - offset: 0, - follow_symlinks: true, - respect_git_ignore: false, - include_globs: vec![], - exclude_globs: vec![], - } - } - - /// Fixture: root with files `b.txt`, `A.txt`, `c.txt` and dirs - /// `Zeta`, `alpha`. Expected order: dirs first case-insensitive - /// (`alpha`, `Zeta`), then files (`A.txt`, `b.txt`, `c.txt`). - fn populate(root: &Path) { - std::fs::write(root.join("b.txt"), b"bb").unwrap(); - std::fs::write(root.join("A.txt"), b"a").unwrap(); - std::fs::write(root.join("c.txt"), b"ccc").unwrap(); - std::fs::create_dir(root.join("Zeta")).unwrap(); - std::fs::create_dir(root.join("alpha")).unwrap(); - } - - /// `list_blocking` against `dir` as both walk root and workspace root - /// (canonicalized for the confinement check, like production). - fn list_dir(dir: &Path, req: &FsListReq, max_collect: usize) -> FsListRes { - let canonical = dunce::canonicalize(dir).unwrap(); - list_blocking(dir, dir, &canonical, req, max_collect).unwrap() - } - - #[test] - fn list_sorts_dirs_first_case_insensitive() { - let dir = tempfile::tempdir().unwrap(); - populate(dir.path()); - let res = list_dir(dir.path(), &list_req(""), MAX_LIST_COLLECT); - let names: Vec<&str> = res.nodes.iter().map(|n| n.name.as_str()).collect(); - assert_eq!(names, ["alpha", "Zeta", "A.txt", "b.txt", "c.txt"]); - assert!(!res.truncated); - assert_eq!(res.nodes[0].node_type, FsNodeType::Directory); - assert_eq!(res.nodes[2].node_type, FsNodeType::File); - assert_eq!(res.nodes[2].size, Some(1)); - assert!(res.nodes[2].mtime_ms.is_some()); - // Paths are workspace-root-relative. - assert_eq!(res.nodes[2].path, "A.txt"); - } - - /// Pagination slices the *sorted* listing, so consecutive pages have - /// stable boundaries and concatenate to the full listing. - #[test] - fn list_paginates_post_sort_with_stable_boundaries() { - let dir = tempfile::tempdir().unwrap(); - populate(dir.path()); - let full = list_dir(dir.path(), &list_req(""), MAX_LIST_COLLECT); - - let mut paged = Vec::new(); - for page_start in [0u64, 2, 4] { - let req = FsListReq { - limit: 2, - offset: page_start, - ..list_req("") - }; - let page = list_dir(dir.path(), &req, MAX_LIST_COLLECT); - // truncated while more entries remain past this page. - assert_eq!(page.truncated, page_start + 2 < full.nodes.len() as u64); - paged.extend(page.nodes); - } - assert_eq!(paged, full.nodes); - - // Offset past the end yields an empty, non-truncated page. - let req = FsListReq { - offset: 100, - ..list_req("") - }; - let page = list_dir(dir.path(), &req, MAX_LIST_COLLECT); - assert!(page.nodes.is_empty()); - assert!(!page.truncated); - } - - /// The collection cap marks the result truncated even when the page - /// itself is not full. - #[test] - fn list_collection_cap_truncates() { - let dir = tempfile::tempdir().unwrap(); - populate(dir.path()); - let res = list_dir(dir.path(), &list_req(""), 2); - assert_eq!(res.nodes.len(), 2); - assert!(res.truncated); - } - - #[test] - fn list_caps_limit_at_server_max() { - let dir = tempfile::tempdir().unwrap(); - populate(dir.path()); - let req = FsListReq { - limit: u32::MAX, - ..list_req("") - }; - // Must not panic / overflow; the page is everything (< 1000). - let res = list_dir(dir.path(), &req, MAX_LIST_COLLECT); - assert_eq!(res.nodes.len(), 5); - } - - /// Regression: the list walk must not traverse — or - /// even surface — in-root symlinks that resolve outside the workspace - /// root, while symlinks staying inside the root keep working. - #[test] - #[cfg(unix)] - fn list_excludes_symlink_escapes_mid_walk() { - let dir = tempfile::tempdir().unwrap(); - let root = dir.path(); - let outside = tempfile::tempdir().unwrap(); - std::fs::write(outside.path().join("secret.txt"), b"secret").unwrap(); - // Escaping symlink: root/escape_link -> . - std::os::unix::fs::symlink(outside.path(), root.join("escape_link")).unwrap(); - // In-root symlink: root/good_link -> root/real_dir. - std::fs::create_dir(root.join("real_dir")).unwrap(); - std::fs::write(root.join("real_dir/inner.txt"), b"inner").unwrap(); - std::os::unix::fs::symlink(root.join("real_dir"), root.join("good_link")).unwrap(); - - let req = FsListReq { - depth: 2, - follow_symlinks: true, - ..list_req("") - }; - let res = list_dir(root, &req, MAX_LIST_COLLECT); - let paths: Vec<&str> = res.nodes.iter().map(|n| n.path.as_str()).collect(); - assert!( - !paths.iter().any(|p| p.contains("escape_link")), - "escaping symlink (and its subtree) must be excluded: {paths:?}" - ); - assert!( - !paths.iter().any(|p| p.contains("secret.txt")), - "outside entries must not be enumerated: {paths:?}" - ); - // Confinement must not over-filter: in-root symlinks survive, - // including descent through them. - assert!(paths.contains(&"good_link"), "{paths:?}"); - assert!(paths.contains(&"good_link/inner.txt"), "{paths:?}"); - assert!(paths.contains(&"real_dir/inner.txt"), "{paths:?}"); - let good = res.nodes.iter().find(|n| n.path == "good_link").unwrap(); - assert_eq!(good.is_symlink, Some(true)); - } - - #[test] - fn memo_lookup_hits_and_invalidates_on_mismatch() { - let memo = FileHashMemo::default(); - let path = Path::new("/ws/a.txt"); - memo.store(path, 10, 1000, "h1".into()); - assert_eq!(memo.lookup(path, 10, 1000).as_deref(), Some("h1")); - // Size change ⇒ miss. - assert_eq!(memo.lookup(path, 11, 1000), None); - // Mtime change ⇒ miss. - assert_eq!(memo.lookup(path, 10, 2000), None); - // Re-store replaces the stale entry. - memo.store(path, 11, 2000, "h2".into()); - assert_eq!(memo.lookup(path, 11, 2000).as_deref(), Some("h2")); - assert_eq!(memo.lookup(path, 10, 1000), None); - } - - /// `stat` consults the memo (no re-hash for an unchanged file) and - /// recomputes when `(size, mtime)` no longer match. - #[tokio::test] - async fn stat_uses_memo_until_file_changes() { - let ws = make_handle(); - let root = ws.root_cwd().unwrap(); - std::fs::write(root.join("data.txt"), b"hello world").unwrap(); - - let req = FsStatReq { - path: "data.txt".into(), - }; - let first = stat(&ws, &req).await.unwrap(); - assert!(first.exists); - assert_eq!(first.node_type, Some(FsNodeType::File)); - assert_eq!(first.size, Some(11)); - let real_hash = first.hash.clone().expect("hash for files"); - - // Plant a sentinel hash for the file's current (size, mtime). A - // second stat must return the sentinel — proof it did not re-hash. - let abs = root.join("data.txt"); - let md = std::fs::metadata(&abs).unwrap(); - let mtime = system_time_ms(md.modified().unwrap()); - ws.shared - .client_fs_hash_memo - .store(&abs, md.len(), mtime, "sentinel".into()); - let memoized = stat(&ws, &req).await.unwrap(); - assert_eq!(memoized.hash.as_deref(), Some("sentinel")); - - // A size change invalidates the memo entry and re-hashes. - std::fs::write(&abs, b"hello brave new world").unwrap(); - let rehashed = stat(&ws, &req).await.unwrap(); - let new_hash = rehashed.hash.expect("hash for files"); - assert_ne!(new_hash, "sentinel"); - assert_ne!(new_hash, real_hash); - } - - /// A path with a *file* as an intermediate component (`ENOTDIR`) is an - /// existence miss, not an RPC error. - #[tokio::test] - async fn stat_enotdir_intermediate_reports_not_exists() { - let ws = make_handle(); - let root = ws.root_cwd().unwrap(); - std::fs::write(root.join("file.txt"), b"x").unwrap(); - let res = stat( - &ws, - &FsStatReq { - path: "file.txt/nested".into(), - }, - ) - .await - .unwrap(); - assert!(!res.exists); - assert_eq!(res.node_type, None); - assert_eq!(res.hash, None); - } - - #[tokio::test] - async fn stat_missing_path_reports_not_exists() { - let ws = make_handle(); - let res = stat( - &ws, - &FsStatReq { - path: "nope.txt".into(), - }, - ) - .await - .unwrap(); - assert!(!res.exists); - assert_eq!(res.node_type, None); - assert_eq!(res.hash, None); - } - - #[tokio::test] - async fn read_file_chunks_are_binary_safe_and_capped() { - let ws = make_handle(); - let root = ws.root_cwd().unwrap(); - // Non-UTF-8 payload: every byte value once. - let payload: Vec = (0u8..=255).collect(); - std::fs::write(root.join("blob.bin"), &payload).unwrap(); - - let req = FsReadFileReq { - path: "blob.bin".into(), - // Bytes 200..210 are bare continuation bytes — never valid UTF-8. - offset: Some(200), - length: Some(50), - max_bytes: 10, // cap below the requested length - encoding: FsReadEncoding::Base64, - }; - let res = read_file(&ws, &req).await.unwrap(); - assert_eq!(res.size, 256); - assert_eq!(res.content, None); - assert_eq!(res.content_type, FsContentType::Binary); - let bytes = base64::engine::general_purpose::STANDARD - .decode(res.content_base64.unwrap()) - .unwrap(); - assert_eq!(bytes, payload[200..210], "maxBytes caps the chunk"); - - // Full-file hash regardless of the requested range. - use sha2::{Digest, Sha256}; - assert_eq!(res.hash, format!("{:x}", Sha256::digest(&payload))); - - // Memoized second read (range-only fast path) returns the - // identical chunk + hash. - let again = read_file(&ws, &req).await.unwrap(); - assert_eq!(again.hash, res.hash); - let again_bytes = base64::engine::general_purpose::STANDARD - .decode(again.content_base64.unwrap()) - .unwrap(); - assert_eq!(again_bytes, payload[200..210]); - } - - /// `maxBytes` is server-capped at [`MAX_READ_BYTES`]: - /// a caller-supplied huge budget cannot make the workspace buffer the - /// whole file. - #[tokio::test] - async fn read_file_server_caps_max_bytes() { - let ws = make_handle(); - let root = ws.root_cwd().unwrap(); - let payload = vec![0u8; (MAX_READ_BYTES + 100) as usize]; - std::fs::write(root.join("big.bin"), &payload).unwrap(); - - let res = read_file( - &ws, - &FsReadFileReq { - path: "big.bin".into(), - offset: None, - length: None, - max_bytes: u64::MAX, - encoding: FsReadEncoding::Base64, - }, - ) - .await - .unwrap(); - assert_eq!(res.size, payload.len() as u64); - let bytes = base64::engine::general_purpose::STANDARD - .decode(res.content_base64.unwrap()) - .unwrap(); - assert_eq!( - bytes.len() as u64, - MAX_READ_BYTES, - "clamped to the server cap" - ); - } - - #[tokio::test] - async fn read_file_utf8_default_and_binary_fallback() { - let ws = make_handle(); - let root = ws.root_cwd().unwrap(); - std::fs::write(root.join("text.txt"), "héllo").unwrap(); - std::fs::write(root.join("bin.dat"), [0xff, 0xfe, 0x00]).unwrap(); - - let text = read_file( - &ws, - &FsReadFileReq { - path: "text.txt".into(), - offset: None, - length: None, - max_bytes: 1_048_576, - encoding: FsReadEncoding::Utf8, - }, - ) - .await - .unwrap(); - assert_eq!(text.content.as_deref(), Some("héllo")); - assert_eq!(text.content_base64, None); - assert_eq!(text.content_type, FsContentType::Text); - - // Invalid UTF-8 under the utf8 default degrades to base64. - let bin = read_file( - &ws, - &FsReadFileReq { - path: "bin.dat".into(), - offset: None, - length: None, - max_bytes: 1_048_576, - encoding: FsReadEncoding::Utf8, - }, - ) - .await - .unwrap(); - assert_eq!(bin.content, None); - assert_eq!(bin.content_type, FsContentType::Binary); - let bytes = base64::engine::general_purpose::STANDARD - .decode(bin.content_base64.unwrap()) - .unwrap(); - assert_eq!(bytes, [0xff, 0xfe, 0x00]); - } - - #[tokio::test] - async fn resolve_rejects_escapes() { - let ws = make_handle(); - for path in ["/etc/passwd", "../escape.txt"] { - let err = stat( - &ws, - &FsStatReq { - path: path.to_owned(), - }, - ) - .await - .expect_err("escape must be rejected"); - assert!(matches!(err, WorkspaceError::HubError(_)), "{err:?}"); - } - } - - /// An absolute path *inside* the workspace root is accepted and stats - /// the same file as its root-relative form. - #[tokio::test] - async fn resolve_accepts_absolute_within_root() { - let ws = make_handle(); - let root = ws.root_cwd().unwrap(); - std::fs::write(root.join("data.txt"), b"hello").unwrap(); - - let rel = stat( - &ws, - &FsStatReq { - path: "data.txt".into(), - }, - ) - .await - .unwrap(); - assert!(rel.exists); - - let abs_path = root.join("data.txt").to_string_lossy().into_owned(); - let abs = stat(&ws, &FsStatReq { path: abs_path }).await.unwrap(); - assert!(abs.exists); - assert_eq!(abs.node_type, rel.node_type); - assert_eq!(abs.size, rel.size); - assert_eq!(abs.hash, rel.hash); - } - - #[tokio::test] - #[cfg(unix)] - async fn resolve_rejects_symlink_escape() { - let ws = make_handle(); - let root = ws.root_cwd().unwrap(); - let outside = tempfile::tempdir().unwrap(); - std::fs::write(outside.path().join("secret.txt"), b"secret").unwrap(); - std::os::unix::fs::symlink(outside.path(), root.join("escape_link")).unwrap(); - - let err = read_file( - &ws, - &FsReadFileReq { - path: "escape_link/secret.txt".into(), - offset: None, - length: None, - max_bytes: 1_048_576, - encoding: FsReadEncoding::Base64, - }, - ) - .await - .expect_err("symlink escape must be rejected"); - assert!( - err.to_string().contains("symlink escape"), - "unexpected error: {err}" - ); - } - - #[tokio::test] - async fn list_empty_path_lists_root() { - let ws = make_handle(); - let root = ws.root_cwd().unwrap(); - std::fs::write(root.join("rooted.txt"), b"x").unwrap(); - let res = list(&ws, &list_req("")).await.unwrap(); - assert!(res.nodes.iter().any(|n| n.name == "rooted.txt")); - } -} diff --git a/crates/codegen/kigi-workspace/src/hub.rs b/crates/codegen/kigi-workspace/src/hub.rs deleted file mode 100644 index 8214343..0000000 --- a/crates/codegen/kigi-workspace/src/hub.rs +++ /dev/null @@ -1,1429 +0,0 @@ -//! Server integration for the workspace. -//! -//! Provides server integration via a single [`ToolServer`] connection: -//! -//! **Provider direction:** The workspace exposes its session tools to -//! the server via [`ToolServer`] + [`WorkspaceToolHandler`]. When the server -//! receives a `tool_call_request` it routes to the workspace's handler, -//! which dispatches to the workspace session matching the -//! `session_id`. Sessions are created on demand via -//! `session.bind` — there is no privileged "main" session. -//! -//! **Session multiplexing:** Multiple sessions can be bound to the -//! same workspace server concurrently. Each gets its own workspace -//! session (isolated CWD, shell state, toolset). Sessions are created -//! when the server sends a `session.bind` notification, and -//! cleaned up on disconnect or explicit unbind. -//! -//! **Notifications:** The same `ToolServer` connection is used for -//! subscribing to notifications (tool changes) and sending -//! workspace events / tool notifications back to the server. -//! -//! The [`HubConnectionPool`] and auth credential are shared, so -//! everything multiplexes over one WebSocket per `(url, principal)`. -//! -//! # Security considerations -//! -//! - **Provider direction** returns full `result.prompt_text` to the -//! remote server. This may contain sensitive workspace data (file -//! contents, env vars). Callers must ensure the server endpoint is -//! trusted. -//! - **Consumer direction** remote tools are merged with `kind: None` and -//! are only visible under `CapabilityMode::All`. They are dropped in -//! subagent sessions with restricted capability modes. -use crate::diag_server::DiagHandle; -use crate::error::{WorkspaceError, WorkspaceResult}; -use crate::handle::WorkspaceHandle; -use async_trait::async_trait; -use kigi_computer_hub_sdk::{ - AuthProvider, ClientError, HubConnectionPool, ToolServer, ToolServerBuilder, ToolServerHandler, -}; -use kigi_tool_protocol::ToolId; -use kigi_tool_runtime::{ - ToolCallContext, ToolError, ToolErrorKind, ToolStream, ToolStreamItem, TypedToolOutput, - terminal_only, -}; -use kigi_tool_types::ToolDescription; -use kigi_tools::registry::types::ToolConfig; -use serde_json::Value; -use std::sync::Arc; -use tokio::task::JoinHandle; -use url::Url; -/// Configuration for connecting to a server instance. -/// -/// Passed via [`WorkspaceConfig::hub_config`](crate::config::WorkspaceConfig::hub_config). -/// When `Some`, the workspace can connect to the server after construction -/// via [`WorkspaceHandle::connect_hub`](crate::handle::WorkspaceHandle::connect_hub). -#[derive(Clone)] -pub struct HubConfig { - /// Server WebSocket URL (`ws://` or `wss://`). - pub url: Url, - pub auth: Arc, - /// Activity tracker to poke on reconnect so the status publisher - /// sends an immediate heartbeat (prevents status reverting to null). - pub activity_tracker: Option>, - /// Stable server ID for `register_server` / `servers.list` / - /// `server.bind`. When `None`, the SDK default (`"workspace-server"`) - /// is used. Set to the sandbox `session_id` in production so each - /// workspace server has a unique, predictable identity. - pub server_id: Option, - /// Optional extra access key attached on the server connection when the - /// non-production feature set is enabled. `None` on prod / local-dev. - pub alpha_test_key: Option, - /// Permit a plaintext `ws://` server on a non-loopback host (mesh-secured). - pub allow_insecure_ws: bool, - /// When set, the ready file tracks hub readiness for the sandbox reconnect - /// gate: written on initial connect (hello / server registered), removed on - /// disconnect, rewritten only after reconnect **serve replay settles** - /// (not merely socket-up). Presence means the tool-server is registered and - /// prior sessions have been re-served when applicable. `None` = unmanaged. - pub ready_file: Option, - /// Diagnostics-server state handle, driven from the same lifecycle points - /// as the ready file. `None` = no diagnostics server (embedded/local use). - pub diag: Option, -} -/// Write the workspace-server ready file (pid as contents). Presence means the -/// tool-server is hub-ready for the sandbox gate (initial hello, or reconnect -/// after serve replay settled); failures are logged, not fatal. -fn write_ready_file(path: &std::path::Path) { - if let Err(e) = std::fs::write(path, std::process::id().to_string()) { - tracing::warn!( - path = % path.display(), error = % e, - "failed to write workspace-server ready file" - ); - } -} -/// Publishes hub readiness to the ready file and the diagnostics server from -/// the same lifecycle transitions, so the two protocols cannot disagree while -/// a diagnostics handle is configured (its shutdown latch gates both). -struct ReadyPublisher { - ready_file: Option, - diag: Option, -} -impl ReadyPublisher { - fn connected(&self) { - if self.diag.as_ref().is_some_and(DiagHandle::is_shutting_down) { - return; - } - if let Some(path) = &self.ready_file { - write_ready_file(path); - } - if let Some(diag) = &self.diag { - diag.set_connected(); - } - } - fn disconnected(&self) { - if let Some(path) = &self.ready_file { - let _ = std::fs::remove_file(path); - } - if let Some(diag) = &self.diag { - diag.set_disconnected(); - } - } -} -impl std::fmt::Debug for HubConfig { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("HubConfig") - .field("url", &self.url.as_str()) - .field("auth", &"") - .field("server_id", &self.server_id) - .finish() - } -} -/// Live handle to a server connection, tool server, and notification -/// listener. -/// -/// Stored on [`WorkspaceShared`](crate::session::WorkspaceShared) as -/// `Option`. Created by -/// [`WorkspaceHandle::connect_hub`](crate::handle::WorkspaceHandle::connect_hub). -pub(crate) struct HubHandle { - /// The tool server exposing workspace tools to the server (provider direction). - /// Also used for subscribing to and sending notifications. - pub(crate) server: ToolServer, - /// Kept alive so the underlying WebSocket connection is not dropped. - /// Dropping this last reference tears down the connection. - /// Not directly accessed — its lifetime keeps connections alive. - #[allow(dead_code)] - pub(crate) pool: Arc, - /// Background tool server run loop task handle. - server_task: Option>, - /// Background notification listener task handle. - notification_task: Option>, - /// Background task that forwards `WorkspaceEvent`s as `tool.notify` - /// custom frames through the server. - event_publisher_task: Option>, - /// Background drain feeding the `ActivityTracker` from the session - /// tool-notification stream (see `run_activity_feed`). - activity_feed_task: Option>, - /// Background task that publishes `tool_server.status` to the server. - status_publisher_task: Option>, - /// Background task that listens for `session.bind` and - /// creates workspace sessions. - session_bind_task: Option>, - /// Background codebase-index event forwarder. Tracked so shutdown aborts it - /// (it holds the `events` sender and cannot self-terminate). - codebase_index_forwarder_task: Option>, - /// Background client ext-notification forwarder. - client_ext_forwarder_task: Option>, -} -impl std::fmt::Debug for HubHandle { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("HubHandle") - .field("server", &self.server) - .field( - "server_task", - if self.server_task.is_some() { - &"Some()" - } else { - &"None" - }, - ) - .field( - "notification_task", - if self.notification_task.is_some() { - &"Some()" - } else { - &"None" - }, - ) - .field( - "event_publisher_task", - if self.event_publisher_task.is_some() { - &"Some()" - } else { - &"None" - }, - ) - .field( - "activity_feed_task", - if self.activity_feed_task.is_some() { - &"Some()" - } else { - &"None" - }, - ) - .field( - "status_publisher_task", - if self.status_publisher_task.is_some() { - &"Some()" - } else { - &"None" - }, - ) - .field( - "session_bind_task", - if self.session_bind_task.is_some() { - &"Some()" - } else { - &"None" - }, - ) - .field( - "codebase_index_forwarder_task", - if self.codebase_index_forwarder_task.is_some() { - &"Some()" - } else { - &"None" - }, - ) - .field( - "client_ext_forwarder_task", - if self.client_ext_forwarder_task.is_some() { - &"Some()" - } else { - &"None" - }, - ) - .finish() - } -} -impl HubHandle { - /// Build server connection pool, tool server, and return a handle. - /// - /// The tool server starts with zero sessions — all sessions are - /// bound dynamically via `session.bind` at runtime. - /// The tool server run loop and notification listener are NOT - /// started here — call [`Self::set_server_task`] and - /// [`Self::set_notification_task`] after spawning. - pub(crate) async fn connect( - config: &HubConfig, - ws_ping: std::time::Duration, - ws_reconnect_backoff: Option>, - tool_handlers: Vec>, - server_metadata: Option, - session_handler_resolver: Option, - ) -> Result { - let pool = HubConnectionPool::new(); - let server_url = config.url.clone(); - let mut server_builder = ToolServerBuilder::default() - .pool(pool.clone()) - .url(server_url) - .auth_provider(config.auth.clone()) - .allow_insecure_ws(config.allow_insecure_ws) - .binary_version(kigi_version::VERSION) - .with_ws_ping_interval(ws_ping); - if let Some(schedule) = ws_reconnect_backoff { - server_builder = server_builder.with_reconnect_backoff(schedule); - } - let activity_tracker = config.activity_tracker.clone(); - if activity_tracker.is_some() { - server_builder = server_builder.on_reconnect(move |_| { - if let Some(tracker) = &activity_tracker { - tracker.poke(); - } - }); - } - if config.ready_file.is_some() || config.diag.is_some() { - let publisher = Arc::new(ReadyPublisher { - ready_file: config.ready_file.clone(), - diag: config.diag.clone(), - }); - let on_connect = Arc::clone(&publisher); - let on_disconnect = Arc::clone(&publisher); - server_builder = server_builder - .on_connect(move || on_connect.connected()) - .on_disconnect(move || on_disconnect.disconnected()) - .on_reconnect_settled(move || publisher.connected()); - } - if let Some(ref id) = config.server_id { - server_builder = server_builder.server_id(parse_server_id(id)?); - } - for handler in tool_handlers { - server_builder = server_builder.tool_dyn(handler); - } - if let Some(meta) = server_metadata { - server_builder = server_builder.metadata(meta); - } - if let Some(resolver) = session_handler_resolver { - server_builder = server_builder.session_handler_resolver(resolver); - } - let server = server_builder.build().await?; - Ok(Self { - server, - pool, - server_task: None, - notification_task: None, - event_publisher_task: None, - activity_feed_task: None, - status_publisher_task: None, - session_bind_task: None, - codebase_index_forwarder_task: None, - client_ext_forwarder_task: None, - }) - } - /// Attach the background tool server run loop task. - pub(crate) fn set_server_task(&mut self, task: JoinHandle<()>) { - self.server_task = Some(task); - } - /// Attach the background notification listener task. - pub(crate) fn set_notification_task(&mut self, task: JoinHandle<()>) { - self.notification_task = Some(task); - } - /// Attach the background workspace event publisher task. - pub(crate) fn set_event_publisher_task(&mut self, task: JoinHandle<()>) { - self.event_publisher_task = Some(task); - } - /// Attach the background activity-feed drain task. - pub(crate) fn set_activity_feed_task(&mut self, task: JoinHandle<()>) { - self.activity_feed_task = Some(task); - } - /// Attach the background status publisher task. - pub(crate) fn set_status_publisher_task(&mut self, task: JoinHandle<()>) { - self.status_publisher_task = Some(task); - } - /// Attach the background codebase-index event forwarder task. - pub(crate) fn set_codebase_index_forwarder_task(&mut self, task: JoinHandle<()>) { - self.codebase_index_forwarder_task = Some(task); - } - /// Attach the background client ext-notification forwarder task. - pub(crate) fn set_client_ext_forwarder_task(&mut self, task: JoinHandle<()>) { - self.client_ext_forwarder_task = Some(task); - } - /// Cooperative shutdown with timeout. - /// - /// 1. Shuts down the tool server (unregisters tools + sessions). - /// 2. Aborts background tasks. - /// - /// The shutdown call is guarded by a 5-second timeout to prevent - /// blocking indefinitely if the server is unreachable. - pub(crate) async fn shutdown(self) { - const SHUTDOWN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); - match tokio::time::timeout(SHUTDOWN_TIMEOUT, self.server.shutdown()).await { - Ok(Ok(())) => {} - Ok(Err(e)) => tracing::warn!(error = % e, "tool server shutdown error"), - Err(_) => tracing::warn!("tool server shutdown timed out"), - } - if let Some(task) = self.server_task { - task.abort(); - let _ = task.await; - } - if let Some(task) = self.notification_task { - task.abort(); - let _ = task.await; - } - if let Some(task) = self.event_publisher_task { - task.abort(); - let _ = task.await; - } - if let Some(task) = self.activity_feed_task { - task.abort(); - let _ = task.await; - } - if let Some(task) = self.status_publisher_task { - task.abort(); - let _ = task.await; - } - if let Some(task) = self.session_bind_task { - task.abort(); - let _ = task.await; - } - if let Some(task) = self.codebase_index_forwarder_task { - task.abort(); - let _ = task.await; - } - if let Some(task) = self.client_ext_forwarder_task { - task.abort(); - let _ = task.await; - } - } -} -/// [`ToolServerHandler`] for an individual tool, dispatched to the -/// workspace session matching the `session_id`. -/// -/// One instance is created per tool discovered from the workspace's -/// `default_tool_config`. The server sees individual tools (bash, -/// read_file, etc.) and routes `tool_call_request` frames directly by -/// `tool_id`. No meta-wrapper, no envelope — the server has full per-tool -/// visibility for routing, listing, and per-session binding. -/// -/// Sessions must be bound via `session.bind` before tool calls -/// are accepted. There is no implicit default session. -pub(crate) struct SessionRoutedToolHandler { - tool_id: ToolId, - desc: ToolDescription, - schema: Option, - workspace: WorkspaceHandle, -} -impl SessionRoutedToolHandler { - pub(crate) fn new( - name: String, - desc: ToolDescription, - schema: Option, - workspace: WorkspaceHandle, - ) -> Result { - Ok(Self { - tool_id: ToolId::new(name)?, - desc, - schema, - workspace, - }) - } - fn name(&self) -> &str { - self.tool_id.as_str() - } -} -/// RAII guard that brackets a tool call's activity-tracker accounting. -/// -/// [`SessionRoutedToolHandler::handle_call`] calls -/// [`ActivityTracker::tool_call_started`](crate::activity::ActivityTracker::tool_call_started) -/// at stream construction and moves this guard into the returned stream. Its -/// [`Drop`] calls -/// [`tool_call_completed`](crate::activity::ActivityTracker::tool_call_completed), -/// so completion bookkeeping fires whether the stream reaches its terminal -/// item *or* the consumer drops the stream early (e.g. harness disconnect). -struct CallCompletedGuard { - tracker: Arc, - call_id: String, - session_id: Option, - outcome: kigi_file_utils::events::ToolOutcome, -} -impl CallCompletedGuard { - fn new( - tracker: Arc, - call_id: String, - session_id: Option, - ) -> Self { - Self { - tracker, - call_id, - session_id, - outcome: kigi_file_utils::events::ToolOutcome::Cancelled, - } - } - fn set_outcome(&mut self, outcome: kigi_file_utils::events::ToolOutcome) { - self.outcome = outcome; - } -} -impl Drop for CallCompletedGuard { - fn drop(&mut self) { - self.tracker - .tool_call_completed(&self.call_id, self.session_id.as_deref(), self.outcome); - } -} -#[async_trait] -impl ToolServerHandler for SessionRoutedToolHandler { - fn tool_id(&self) -> ToolId { - self.tool_id.clone() - } - fn description(&self) -> ToolDescription { - self.desc.clone() - } - fn input_schema(&self) -> Option { - self.schema.clone() - } - async fn handle_call(&self, ctx: ToolCallContext, args: Value) -> ToolStream { - let tool_id = self.tool_id(); - let hub_session = ctx - .extensions - .get::() - .map(|s| s.0.clone()); - let tracker = &self.workspace.shared.activity_tracker; - if tracker.is_draining() { - return terminal_only(Err(ToolError::new( - ToolErrorKind::TerminalError, - "workspace is draining — no new tool calls accepted", - ))); - } - let session_id = match &hub_session { - Some(sid) => sid.as_str(), - None => { - return terminal_only(Err(ToolError::new( - ToolErrorKind::InvalidArguments, - "tool_call_request missing session_id", - ))); - } - }; - let session = match self.workspace.session(session_id) { - Some(s) => s, - None => { - return terminal_only(Err(ToolError::new( - ToolErrorKind::InvalidArguments, - format!("session not bound: {session_id}"), - ))); - } - }; - let call_id = ctx.call_id.to_string(); - if crate::permission::hitl_permission_live_enabled() - && !session.yolo_mode() - && let Some(access) = crate::permission::access_kind_for_hub_tool(self.name(), &args) - { - let transport = self - .workspace - .hub_server_blocking() - .await - .and_then(|server| { - crate::permission::ToolServerPermissionTransport::from_session_id( - server, session_id, - ) - }); - match transport { - Some(transport) => { - let outcome = crate::permission::request_permission_via_hub( - &transport, &access, &call_id, - ) - .await; - if !crate::permission::prompt_outcome_allows(&outcome) { - use crate::permission::PromptOutcome; - let deny_msg = match &outcome { - PromptOutcome::FollowupMessage(msg) => { - format!("tool permission redirected: {msg}") - } - _ => format!("tool permission denied for {}", self.name()), - }; - tracing::info!( - tool = % self.name(), session = % session_id, call_id = % - call_id, ? outcome, - "tool-permission denied via hub; rejecting tool call" - ); - return terminal_only(Err(ToolError::new( - ToolErrorKind::PermissionDenied, - deny_msg, - ))); - } - } - None => { - tracing::warn!( - tool = % self.name(), session = % session_id, - "KIGI_HITL_PERMISSION_LIVE set but no hub ToolServer; rejecting guarded tool" - ); - return terminal_only(Err(ToolError::new( - ToolErrorKind::PermissionDenied, - "tool permission unavailable (no hub transport)", - ))); - } - } - } - let toolset = session.toolset(); - tracing::debug!( - tool = % self.name(), call_id = % call_id, session = % session_id, - "dispatching tool call" - ); - tracker.tool_call_started(&call_id, self.name(), hub_session.as_deref()); - let inner = toolset.call_streaming(self.name(), args, &call_id, None); - let tracker = self.workspace.shared.activity_tracker.clone(); - let name = self.name().to_owned(); - let session_label = session_id.to_owned(); - let guard = CallCompletedGuard::new(tracker, call_id, Some(session_label.clone())); - Box::pin(async_stream::stream! { - use futures::StreamExt; let mut _guard = guard; 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)) => { _guard - .set_outcome(kigi_file_utils::events::ToolOutcome::Success); yield - ToolStreamItem::Terminal(Ok(run_result - .into_typed_tool_output(tool_id),)); return; } - ToolStreamItem::Terminal(Err(e)) => { tracing::error!(tool = % name, - session = % session_label, error = % e, kind = % e.variant_name(), - "tool call failed"); _guard - .set_outcome(kigi_file_utils::events::ToolOutcome::Error); yield - ToolStreamItem::Terminal(Err(e)); return; } } } yield - ToolStreamItem::Terminal(Err(ToolError::new(ToolErrorKind::TerminalError, - "tool stream ended without a terminal",))); - }) - } -} -/// Convert a set of remote [`ToolId`]s into workspace [`ToolConfig`]s. -/// -/// Each remote tool gets a `ToolConfig` with: -/// - `id` prefixed with `hub:` to avoid collisions with baseline/MCP tools -/// - `kind: None` (remote tools have unknown capability kind) -/// - `name_override` set to the bare tool name -/// -/// # Capability mode filtering -/// -/// Remote-origin `kind: None` tools are dropped under non-`All` capability -/// modes (e.g. `ReadWrite`, `ReadOnly` in subagent sessions), matching -/// MCP-origin tool behavior. They are only visible in the main session -/// which uses `CapabilityMode::All`. -pub(crate) fn hub_tool_ids_to_tool_configs(tool_ids: &[ToolId]) -> Vec { - if !tool_ids.is_empty() { - tracing::info!( - count = tool_ids.len(), tools = ? tool_ids.iter().map(| id | id.as_str()) - .collect::< Vec < _ >> (), "Registering remote tools" - ); - } - tool_ids - .iter() - .map(|id| { - let name = id.as_str().to_owned(); - let mut tc = ToolConfig::from_id(format!("hub:{name}")); - tc.name_override = Some(name); - tc - }) - .collect() -} -/// Apply a `ToolsChanged` notification to the current remote tools snapshot. -/// -/// Returns the new snapshot. Extracted as a named function for -/// testability. -pub(crate) fn apply_tools_changed( - current: &[ToolConfig], - added: &[ToolId], - removed: &[ToolId], - updated: &[ToolId], -) -> Vec { - let evicted: std::collections::HashSet = removed - .iter() - .chain(updated.iter()) - .map(|id| format!("hub:{}", id.as_str())) - .collect(); - let mut new_tools: Vec = current - .iter() - .filter(|t| !evicted.contains(&t.id)) - .cloned() - .collect(); - let mut to_add = Vec::with_capacity(added.len() + updated.len()); - to_add.extend_from_slice(added); - to_add.extend_from_slice(updated); - let added_configs = hub_tool_ids_to_tool_configs(&to_add); - let existing_ids: std::collections::HashSet = - new_tools.iter().map(|t| t.id.clone()).collect(); - for tc in added_configs { - if !existing_ids.contains(&tc.id) { - new_tools.push(tc); - } - } - new_tools -} -fn parse_server_id(id: &str) -> Result { - kigi_tool_protocol::ServerId::new(id) - .map_err(|e| ClientError::InvalidConfig(format!("invalid server_id {id:?}: {e}"))) -} -/// Map a [`ClientError`] into a [`WorkspaceError::HubError`]. -pub(crate) fn client_error_to_workspace(err: ClientError) -> WorkspaceError { - WorkspaceError::HubError(err.to_string()) -} -/// Map a server connection failure into a [`WorkspaceResult`]. -pub(crate) fn hub_result(result: Result) -> WorkspaceResult { - result.map_err(client_error_to_workspace) -} -#[cfg(test)] -mod tests { - use super::*; - use kigi_tools::types::tool::ToolKind; - #[test] - fn hub_tool_ids_to_tool_configs_basic() { - let ids = vec![ - ToolId::new("read_file").unwrap(), - ToolId::new("web_search").unwrap(), - ]; - let configs = hub_tool_ids_to_tool_configs(&ids); - assert_eq!(configs.len(), 2); - assert_eq!(configs[0].id, "hub:read_file"); - assert_eq!(configs[0].name_override.as_deref(), Some("read_file")); - assert_eq!(configs[0].kind, None::); - assert_eq!(configs[1].id, "hub:web_search"); - assert_eq!(configs[1].name_override.as_deref(), Some("web_search")); - } - #[test] - fn hub_tool_ids_to_tool_configs_empty() { - let configs = hub_tool_ids_to_tool_configs(&[]); - assert!(configs.is_empty()); - } - #[test] - fn apply_tools_changed_adds_and_removes() { - let initial = hub_tool_ids_to_tool_configs(&[ - ToolId::new("tool_a").unwrap(), - ToolId::new("tool_b").unwrap(), - ]); - let result = apply_tools_changed( - &initial, - &[ToolId::new("tool_c").unwrap()], - &[ToolId::new("tool_a").unwrap()], - &[], - ); - let ids: Vec<&str> = result.iter().map(|t| t.id.as_str()).collect(); - assert!(ids.contains(&"hub:tool_b")); - assert!(ids.contains(&"hub:tool_c")); - assert!(!ids.contains(&"hub:tool_a")); - } - #[test] - fn apply_tools_changed_updates_replace() { - let initial = hub_tool_ids_to_tool_configs(&[ToolId::new("tool_a").unwrap()]); - let result = apply_tools_changed(&initial, &[], &[], &[ToolId::new("tool_a").unwrap()]); - assert_eq!(result.len(), 1); - assert_eq!(result[0].id, "hub:tool_a"); - } - use futures::StreamExt; - use kigi_tool_runtime::{SessionContext, ToolCallId}; - fn make_handler(workspace: &WorkspaceHandle, tool_name: &str) -> SessionRoutedToolHandler { - SessionRoutedToolHandler::new( - tool_name.to_owned(), - ToolDescription::new(tool_name.to_owned(), String::new()), - None, - workspace.clone(), - ) - .expect("test tool name is a valid ToolId") - } - #[tokio::test] - async fn handler_construction_rejects_invalid_tool_name_without_panic() { - let handle = crate::handle::tests::make_handle(); - let err = SessionRoutedToolHandler::new( - "not a tool id!".to_owned(), - ToolDescription::new("not a tool id!".to_owned(), String::new()), - None, - handle.clone(), - ); - assert!( - err.is_err(), - "invalid name must be rejected at construction" - ); - } - #[tokio::test] - async fn handler_tool_id_round_trips_the_validated_name() { - let handle = crate::handle::tests::make_handle(); - let handler = make_handler(&handle, "read_file"); - assert_eq!(handler.tool_id().as_str(), "read_file"); - } - #[test] - fn parse_server_id_maps_invalid_id_to_invalid_config_error() { - for bad in ["auto:tool:x", ""] { - assert!( - matches!(parse_server_id(bad), Err(ClientError::InvalidConfig(_))), - "server_id {bad:?} must map to InvalidConfig" - ); - } - assert!(parse_server_id("sess-abc123").is_ok()); - } - fn make_ctx(session_id: &str) -> (ToolCallContext, String) { - let mut ctx = ToolCallContext::new(ToolCallId::new_v7()); - ctx.insert(SessionContext(session_id.to_owned())); - let call_id = ctx.call_id.to_string(); - (ctx, call_id) - } - #[tokio::test] - async fn handle_call_is_passthrough_zero_progress_one_terminal() { - let handle = crate::handle::tests::make_handle(); - let handler = make_handler(&handle, "read_file"); - let (ctx, _call_id) = make_ctx("main"); - let stream = handler - .handle_call( - ctx, - serde_json::json!({ "target_file" : "does-not-exist.txt" }), - ) - .await; - let items: Vec<_> = stream.collect().await; - let progress = items - .iter() - .filter(|i| matches!(i, ToolStreamItem::Progress(_))) - .count(); - let terminal = items - .iter() - .filter(|i| matches!(i, ToolStreamItem::Terminal(_))) - .count(); - assert_eq!(progress, 0, "gate-off pass-through must emit zero Progress"); - assert_eq!(terminal, 1, "must emit exactly one Terminal"); - assert!(matches!(items.last(), Some(ToolStreamItem::Terminal(_)))); - } - #[tokio::test] - async fn handle_call_terminal_matches_non_streaming_call() { - let handle = crate::handle::tests::make_handle(); - let session = handle.session("main").expect("main session present"); - let toolset = session.toolset(); - let args = serde_json::json!({ "target_file" : "missing-file.txt" }); - let reference = toolset - .call("read_file", args.clone(), "ref-call", None) - .await; - let handler = make_handler(&handle, "read_file"); - let (ctx, _call_id) = make_ctx("main"); - let stream = handler.handle_call(ctx, args).await; - let items: Vec<_> = stream.collect().await; - let terminal = items - .into_iter() - .find_map(|i| match i { - ToolStreamItem::Terminal(t) => Some(t), - ToolStreamItem::Progress(_) => None, - }) - .expect("a terminal item"); - match reference { - Ok(run_result) => { - let reference_value = serde_json::to_value(&run_result).unwrap(); - let typed = terminal.expect("streaming terminal must be Ok when call() is Ok"); - assert_eq!( - typed.value, reference_value, - "streaming terminal must serialize identically to the non-streaming path" - ); - assert_eq!( - typed.model_output, - vec![kigi_tool_runtime::ContentBlock::Text { - text: run_result.prompt_text.clone(), - }], - "model_output must be the prompt_text, not a JSON dump" - ); - } - Err(_) => { - assert!( - terminal.is_err(), - "streaming terminal must be Err when call() is Err" - ); - } - } - } - #[tokio::test] - async fn handle_call_draining_returns_single_terminal_error() { - let handle = crate::handle::tests::make_handle(); - let tracker = handle.activity_tracker().clone(); - tracker.set_draining(); - let handler = make_handler(&handle, "read_file"); - let (ctx, _call_id) = make_ctx("main"); - let stream = handler - .handle_call(ctx, serde_json::json!({ "target_file" : "x.txt" })) - .await; - let items: Vec<_> = stream.collect().await; - assert_eq!(items.len(), 1, "draining yields exactly one item"); - match &items[0] { - ToolStreamItem::Terminal(Err(e)) => { - assert!(e.to_string().contains("draining"), "got: {e}"); - } - ToolStreamItem::Terminal(Ok(_)) => { - panic!("expected Terminal(Err), got Terminal(Ok)") - } - ToolStreamItem::Progress(_) => panic!("expected Terminal(Err), got Progress"), - } - assert_eq!( - tracker.snapshot().active_tool_calls, - 0, - "draining must not start a tool call" - ); - } - #[tokio::test] - async fn handle_call_guard_completes_on_early_drop() { - let handle = crate::handle::tests::make_handle(); - let tracker = handle.activity_tracker().clone(); - let handler = make_handler(&handle, "read_file"); - let (ctx, _call_id) = make_ctx("main"); - let stream = handler - .handle_call(ctx, serde_json::json!({ "target_file" : "x.txt" })) - .await; - assert_eq!( - tracker.snapshot().active_tool_calls, - 1, - "tool_call_started fires at stream construction" - ); - drop(stream); - assert_eq!( - tracker.snapshot().active_tool_calls, - 0, - "dropping the stream must run the RAII completion guard" - ); - } - use kigi_tools::types::tool_metadata::ToolMetadata as XaiToolMetadata; - #[derive(Debug)] - struct GateStreamingStub; - impl XaiToolMetadata for GateStreamingStub { - fn kind(&self) -> ToolKind { - ToolKind::Other - } - fn tool_namespace(&self) -> kigi_tools::types::tool::ToolNamespace { - kigi_tools::types::tool::ToolNamespace::MCP - } - fn description_template(&self) -> &str { - "gate streaming stub" - } - } - impl kigi_tool_runtime::Tool for GateStreamingStub { - type Args = serde_json::Value; - type Output = String; - fn id(&self) -> kigi_tool_protocol::ToolId { - kigi_tool_protocol::ToolId::new("gate_streaming_stub").expect("valid tool id") - } - fn description( - &self, - _ctx: &::kigi_tool_runtime::ListToolsContext, - ) -> kigi_tool_types::ToolDescription { - kigi_tool_types::ToolDescription::new("gate_streaming_stub", "gate streaming stub") - } - async fn run( - &self, - _ctx: kigi_tool_runtime::ToolCallContext, - _input: serde_json::Value, - ) -> Result { - Ok("terminal-value".into()) - } - async fn execute( - &self, - _ctx: kigi_tool_runtime::ToolCallContext, - _input: serde_json::Value, - ) -> kigi_tool_runtime::ToolStream { - use kigi_tool_runtime::{ToolProgress, ToolStreamItem}; - Box::pin(futures::stream::iter(vec![ - ToolStreamItem::Progress(ToolProgress::Text { - text: "stub-progress-1".into(), - }), - ToolStreamItem::Progress(ToolProgress::Text { - text: "stub-progress-2".into(), - }), - ToolStreamItem::Terminal(Ok("stub-terminal".into())), - ])) - } - } - fn register_gate_stub(handle: &WorkspaceHandle, tool_name: &str) { - let session = handle.session("main").expect("main session present"); - let toolset = session.toolset(); - toolset - .register_tool( - tool_name.to_owned(), - GateStreamingStub, - Some(serde_json::json!({ "type" : "object", "properties" : {} })), - ) - .expect("register_tool must succeed"); - } - async fn drain_counts(mut stream: kigi_tool_runtime::ToolStream) -> (usize, usize, bool) { - let mut progress = 0; - let mut terminal = 0; - let mut last_is_terminal = false; - while let Some(item) = stream.next().await { - match item { - ToolStreamItem::Progress(_) => { - progress += 1; - last_is_terminal = false; - } - ToolStreamItem::Terminal(_) => { - terminal += 1; - last_is_terminal = true; - } - } - } - (progress, terminal, last_is_terminal) - } - #[tokio::test] - async fn handle_call_forwards_inner_streaming_tool_progress() { - let handle = crate::handle::tests::make_handle(); - register_gate_stub(&handle, "gate_streamer_forward"); - let handler = make_handler(&handle, "gate_streamer_forward"); - let (ctx, _call_id) = make_ctx("main"); - let stream = handler.handle_call(ctx, serde_json::json!({})).await; - let (progress, terminal, last_is_terminal) = drain_counts(stream).await; - assert_eq!( - progress, 2, - "workspace must forward both stub Progress items end-to-end" - ); - assert_eq!(terminal, 1, "exactly one Terminal"); - assert!(last_is_terminal, "Terminal must be the final item"); - } - #[tokio::test] - async fn handle_call_preserves_bash_chat_completion_output() { - let handle = crate::handle::tests::make_handle(); - crate::handle::tests::register_bash_cco_stub(&handle); - let handler = make_handler(&handle, crate::handle::tests::BASH_CCO_STUB_NAME); - let (ctx, _call_id) = make_ctx("main"); - let stream = handler.handle_call(ctx, serde_json::json!({})).await; - let typed = crate::handle::tests::drain_terminal_ok(stream).await; - crate::handle::tests::assert_bash_cco_terminal(&typed); - } - use crate::capability::CapabilityMode; - use crate::session::tool_config::test_support::tc; - use kigi_tools::notification::types::{ToolNotification, ToolNotificationHandle}; - use kigi_tools::registry::types::ToolServerConfig; - use std::time::Duration; - fn bg_config() -> ToolServerConfig { - ToolServerConfig { - tools: vec![ - tc("GrokBuild:run_terminal_cmd", Some(ToolKind::Execute)), - tc( - "GrokBuild:get_task_output", - Some(ToolKind::BackgroundTaskAction), - ), - tc("GrokBuild:kill_task", Some(ToolKind::KillTaskAction)), - tc("GrokBuild:monitor", Some(ToolKind::Monitor)), - ], - behavior_preset: None, - } - } - fn install_activity_feed(handle: &WorkspaceHandle) { - let (sink, rx) = ToolNotificationHandle::channel(); - handle - .shared() - .activity_notify_handle - .store(Arc::new(Some(sink))); - tokio::spawn(crate::handle::run_activity_feed( - handle.activity_tracker().clone(), - rx, - )); - } - fn make_bg_handle_with_config(cfg: ToolServerConfig) -> WorkspaceHandle { - let handle = WorkspaceHandle::for_test(); - install_activity_feed(&handle); - handle - .create_session_with_config("main", None, Some(cfg), CapabilityMode::All, None, false) - .expect("create main session with background tools"); - handle - } - fn make_bg_tracking_handle() -> WorkspaceHandle { - make_bg_handle_with_config(bg_config()) - } - async fn wait_until( - tracker: &crate::activity::ActivityTracker, - pred: impl Fn(&kigi_tool_protocol::ToolServerStatusPayload) -> bool, - timeout: Duration, - ) -> kigi_tool_protocol::ToolServerStatusPayload { - let deadline = tokio::time::Instant::now() + timeout; - loop { - let snap = tracker.snapshot(); - if pred(&snap) || tokio::time::Instant::now() >= deadline { - return snap; - } - tokio::time::sleep(Duration::from_millis(20)).await; - } - } - async fn run_tool_in_session(handle: &WorkspaceHandle, session: &str, tool: &str, args: Value) { - let handler = make_handler(handle, tool); - let (ctx, _call_id) = make_ctx(session); - let mut stream = handler.handle_call(ctx, args).await; - while let Some(item) = stream.next().await { - if matches!(item, ToolStreamItem::Terminal(_)) { - break; - } - } - } - fn bg_started_notif(task_id: &str) -> ToolNotification { - use kigi_tools::notification::types::{BashExecutionBackgrounded, BashNotificationBase}; - ToolNotification::BashExecutionBackgrounded(BashExecutionBackgrounded { - base: BashNotificationBase { - tool_call_id: task_id.to_owned(), - command: "sleep 1".to_owned(), - output: Vec::new(), - total_bytes: 0, - truncated: false, - cwd: std::path::PathBuf::from("/tmp"), - }, - output_file: std::path::PathBuf::from("/tmp/x.log"), - task_id: task_id.to_owned(), - monitor_description: None, - description: None, - }) - } - fn task_completed_notif(task_id: &str) -> ToolNotification { - use kigi_tools::computer::types::{TaskKind, TaskSnapshot}; - ToolNotification::TaskCompleted(TaskSnapshot { - task_id: task_id.to_owned(), - command: "sleep 1".to_owned(), - display_command: None, - cwd: "/tmp".to_owned(), - start_time: std::time::SystemTime::UNIX_EPOCH, - end_time: Some(std::time::SystemTime::UNIX_EPOCH), - output: String::new(), - output_file: std::path::PathBuf::from("/tmp/x.log"), - truncated: false, - exit_code: Some(0), - signal: None, - completed: true, - kind: TaskKind::Bash, - block_waited: false, - explicitly_killed: false, - owner_session_id: None, - }) - } - fn started_id(n: &ToolNotification) -> &str { - match n { - ToolNotification::BashExecutionBackgrounded(b) => &b.task_id, - _ => "", - } - } - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn backgrounded_bash_increments_then_decrements_through_real_wiring() { - let handle = make_bg_tracking_handle(); - let tracker = handle.activity_tracker().clone(); - run_tool_in_session( - &handle, - "main", - "run_terminal_cmd", - serde_json::json!( - { "command" : "sleep 2", "description" : "test", "is_background" : - true } - ), - ) - .await; - let busy = wait_until( - &tracker, - |s| s.background_tasks == 1 && s.idle_since_ms.is_none(), - Duration::from_secs(5), - ) - .await; - assert_eq!(busy.background_tasks, 1, "running bg bash must increment"); - assert!(busy.idle_since_ms.is_none(), "idle withheld while bg runs"); - let idle = wait_until( - &tracker, - |s| s.background_tasks == 0 && s.idle_since_ms.is_some(), - Duration::from_secs(15), - ) - .await; - assert_eq!( - idle.background_tasks, 0, - "TaskCompleted must decrement to zero" - ); - assert!( - idle.idle_since_ms.is_some(), - "idle restored when nothing runs" - ); - } - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn auto_background_on_timeout_increments_then_decrements_through_real_wiring() { - let mut cfg = bg_config(); - cfg.tools[0].params = serde_json::json!( - { "enabled_background" : true, "auto_background_on_timeout" : true, } - ) - .as_object() - .cloned(); - let handle = make_bg_handle_with_config(cfg); - let tracker = handle.activity_tracker().clone(); - run_tool_in_session( - &handle, - "main", - "run_terminal_cmd", - serde_json::json!( - { "command" : "sleep 2", "description" : "test", "timeout" : 300 } - ), - ) - .await; - let busy = wait_until( - &tracker, - |s| s.background_tasks == 1 && s.idle_since_ms.is_none(), - Duration::from_secs(5), - ) - .await; - assert_eq!( - busy.background_tasks, 1, - "auto-backgrounded task must increment" - ); - assert!(busy.idle_since_ms.is_none()); - let idle = wait_until( - &tracker, - |s| s.background_tasks == 0 && s.idle_since_ms.is_some(), - Duration::from_secs(15), - ) - .await; - assert_eq!( - idle.background_tasks, 0, - "auto-bg completion must decrement (matching task_id)" - ); - assert!(idle.idle_since_ms.is_some()); - } - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn monitor_increments_then_decrements_through_real_wiring() { - let handle = make_bg_tracking_handle(); - let tracker = handle.activity_tracker().clone(); - run_tool_in_session( - &handle, - "main", - "monitor", - serde_json::json!( - { "command" : "sleep 2", "description" : "test monitor" } - ), - ) - .await; - let busy = wait_until( - &tracker, - |s| s.background_tasks == 1, - Duration::from_secs(5), - ) - .await; - assert_eq!(busy.background_tasks, 1, "a started monitor must increment"); - let idle = wait_until( - &tracker, - |s| s.background_tasks == 0 && s.idle_since_ms.is_some(), - Duration::from_secs(15), - ) - .await; - assert_eq!( - idle.background_tasks, 0, - "monitor completion must decrement (the previously-lost decrement)" - ); - assert!(idle.idle_since_ms.is_some()); - } - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn concurrent_background_tasks_track_independently() { - let handle = make_bg_tracking_handle(); - let tracker = handle.activity_tracker().clone(); - run_tool_in_session( - &handle, - "main", - "run_terminal_cmd", - serde_json::json!( - { "command" : "sleep 2", "description" : "test", "is_background" : - true } - ), - ) - .await; - run_tool_in_session( - &handle, - "main", - "run_terminal_cmd", - serde_json::json!( - { "command" : "sleep 5", "description" : "test", "is_background" : - true } - ), - ) - .await; - let two = wait_until( - &tracker, - |s| s.background_tasks == 2, - Duration::from_secs(5), - ) - .await; - assert_eq!( - two.background_tasks, 2, - "both concurrent bg tasks must count" - ); - assert!(two.idle_since_ms.is_none(), "not idle with two bg tasks"); - let one = wait_until( - &tracker, - |s| s.background_tasks == 1 && s.idle_since_ms.is_none(), - Duration::from_secs(10), - ) - .await; - assert_eq!( - one.background_tasks, 1, - "one finishing must not zero the counter" - ); - assert!( - one.idle_since_ms.is_none(), - "still not idle while one remains" - ); - let zero = wait_until( - &tracker, - |s| s.background_tasks == 0 && s.idle_since_ms.is_some(), - Duration::from_secs(15), - ) - .await; - assert_eq!(zero.background_tasks, 0); - assert!( - zero.idle_since_ms.is_some(), - "idle restored only after the last ends" - ); - } - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn forked_child_background_task_feeds_tracker() { - let handle = make_bg_tracking_handle(); - let tracker = handle.activity_tracker().clone(); - let mut cfg = crate::config::AgentSessionConfig::new("child"); - cfg.parent_session_id = Some("main".to_owned()); - cfg.capability_mode = CapabilityMode::All; - cfg.tool_config = Some(bg_config()); - handle.fork_session(cfg).await.expect("fork child session"); - run_tool_in_session( - &handle, - "child", - "run_terminal_cmd", - serde_json::json!( - { "command" : "sleep 2", "description" : "test", "is_background" : - true } - ), - ) - .await; - let busy = wait_until( - &tracker, - |s| s.background_tasks == 1, - Duration::from_secs(5), - ) - .await; - assert_eq!( - busy.background_tasks, 1, - "a forked subagent's bg task must feed the connection-level tracker" - ); - let zero = wait_until( - &tracker, - |s| s.background_tasks == 0, - Duration::from_secs(15), - ) - .await; - assert_eq!( - zero.background_tasks, 0, - "the fork's bg task must decrement on completion" - ); - } - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn compose_session_notification_handle_covers_all_branches() { - let handle = WorkspaceHandle::for_test(); - let shared = handle.shared(); - assert!(shared.compose_session_notification_handle(None).is_none()); - let (sys, mut sys_rx) = ToolNotificationHandle::channel(); - shared - .compose_session_notification_handle(Some(sys)) - .expect("system-only sink") - .send(bg_started_notif("sys-only")); - assert!(matches!(sys_rx.try_recv(), Ok(n) if started_id(& n) == "sys-only")); - let (activity, mut activity_rx) = ToolNotificationHandle::channel(); - shared - .activity_notify_handle - .store(Arc::new(Some(activity))); - shared - .compose_session_notification_handle(None) - .expect("activity-only sink") - .send(bg_started_notif("act-only")); - assert!(matches!(activity_rx.try_recv(), Ok(n) if started_id(& n) == "act-only")); - let (sys2, mut sys2_rx) = ToolNotificationHandle::channel(); - shared - .compose_session_notification_handle(Some(sys2)) - .expect("tee sink") - .send(bg_started_notif("both")); - assert!( - matches!(activity_rx.try_recv(), Ok(n) if started_id(& n) == "both"), - "tee must deliver to the activity (tracker) leg" - ); - assert!( - matches!(sys2_rx.try_recv(), Ok(n) if started_id(& n) == "both"), - "tee must deliver to the system.notify leg" - ); - } - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn activity_feed_drains_and_dedups_by_task_id() { - let tracker = Arc::new(crate::activity::ActivityTracker::new()); - let (sink, rx) = ToolNotificationHandle::channel(); - let feed = tokio::spawn(crate::handle::run_activity_feed(tracker.clone(), rx)); - sink.send(bg_started_notif("dup")); - sink.send(bg_started_notif("dup")); - let one = wait_until( - &tracker, - |s| s.background_tasks == 1, - Duration::from_secs(5), - ) - .await; - assert_eq!( - one.background_tasks, 1, - "duplicate started must not double-count" - ); - sink.send(bg_started_notif("other")); - let two = wait_until( - &tracker, - |s| s.background_tasks == 2, - Duration::from_secs(5), - ) - .await; - assert_eq!(two.background_tasks, 2); - sink.send(task_completed_notif("dup")); - sink.send(task_completed_notif("other")); - let zero = wait_until( - &tracker, - |s| s.background_tasks == 0, - Duration::from_secs(5), - ) - .await; - assert_eq!( - zero.background_tasks, 0, - "completions drain via run_activity_feed" - ); - drop(sink); - feed.await - .expect("activity feed must exit cleanly once senders drop"); - } - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn update_tool_config_preserves_tracker_feed() { - let handle = make_bg_tracking_handle(); - let tracker = handle.activity_tracker().clone(); - handle - .update_tool_config("main", "main", bg_config()) - .await - .expect("update_tool_config rebuilds the toolset"); - run_tool_in_session( - &handle, - "main", - "run_terminal_cmd", - serde_json::json!( - { "command" : "sleep 2", "description" : "test", "is_background" : - true } - ), - ) - .await; - let busy = wait_until( - &tracker, - |s| s.background_tasks == 1, - Duration::from_secs(5), - ) - .await; - assert_eq!( - busy.background_tasks, 1, - "a bg task after update_tool_config must still feed the tracker" - ); - } - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn re_resolve_all_sessions_preserves_tracker_feed() { - let handle = make_bg_tracking_handle(); - let tracker = handle.activity_tracker().clone(); - let rebuilt = handle - .shared() - .re_resolve_all_sessions("test_preserves_feed", true) - .await; - assert!(rebuilt >= 1, "the main session must be re-resolved"); - run_tool_in_session( - &handle, - "main", - "run_terminal_cmd", - serde_json::json!( - { "command" : "sleep 2", "description" : "test", "is_background" : - true } - ), - ) - .await; - let busy = wait_until( - &tracker, - |s| s.background_tasks == 1, - Duration::from_secs(5), - ) - .await; - assert_eq!( - busy.background_tasks, 1, - "a bg task after re_resolve_all_sessions must still feed the tracker" - ); - } -} diff --git a/crates/codegen/kigi-workspace/src/hub_auth.rs b/crates/codegen/kigi-workspace/src/hub_auth.rs deleted file mode 100644 index f2bf81e..0000000 --- a/crates/codegen/kigi-workspace/src/hub_auth.rs +++ /dev/null @@ -1,490 +0,0 @@ -//! Hub [`AuthProvider`] from `~/.kigi/auth.json` for the standalone -//! `workspace_server` binary: loopback `ws://` uses a plain bearer, otherwise -//! an auto-refreshing OIDC provider that persists rotated tokens to disk. -//! -//! The in-leader `grok workspace` exposure does NOT use this path — it sources -//! an in-memory provider from the leader's `AuthManager` (see -//! `LeaderAuthProvider`) to avoid racing the leader's own auth.json writer. - -use std::collections::BTreeMap; -use std::path::{Path, PathBuf}; -use std::sync::Arc; - -use kigi_computer_hub_sdk::{ - AuthCredential, AuthIdentity, AuthProvider, OidcAuthProviderBuilder, RefreshEvent, -}; -use url::Url; - -/// Plain bearer provider that also carries the owner identity parsed from the -/// same auth.json entry. Used for the loopback / local-dev path (no OIDC -/// refresh) so consumers can still read the owner identity from the auth -/// provider — without a second auth.json read. -struct BearerWithIdentity { - token: String, - identity: AuthIdentity, -} - -impl std::fmt::Debug for BearerWithIdentity { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - // Never log the bearer token; surface only the (non-secret) identity. - f.debug_struct("BearerWithIdentity") - .field("identity", &self.identity) - .finish_non_exhaustive() - } -} - -impl AuthProvider for BearerWithIdentity { - fn current(&self) -> AuthCredential { - AuthCredential::bearer(self.token.clone()) - } - - fn identity(&self) -> Option { - Some(self.identity.clone()) - } -} - -/// Owner identity parsed from an auth.json entry, for the [`AuthProvider`]s -/// built here to surface via [`AuthProvider::identity`]. -fn identity_from_entry(entry: &AuthEntry) -> AuthIdentity { - AuthIdentity { - user_id: entry.user_id.clone(), - principal_type: entry.principal_type.clone(), - principal_id: entry.principal_id.clone(), - } -} - -#[derive(Debug, serde::Deserialize)] -struct AuthEntry { - key: String, - #[serde(default)] - user_id: String, - #[serde(default)] - refresh_token: Option, - #[serde(default)] - oidc_issuer: Option, - #[serde(default)] - oidc_client_id: Option, - #[serde(default)] - principal_type: Option, - #[serde(default)] - principal_id: Option, - #[serde(default)] - expires_at: Option>, -} - -fn default_auth_path() -> anyhow::Result { - let grok = kigi_config::user_kigi_home() - .ok_or_else(|| anyhow::anyhow!("no user grok home (set $KIGI_SHARE_DIR or $HOME)"))?; - Ok(grok.join("auth.json")) -} - -/// Read the active OIDC entry and its scope key. The key is threaded to the -/// refresh write so rotation updates exactly the entry that was read. -fn read_auth_entry(path: &Path) -> anyhow::Result<(String, AuthEntry)> { - if !path.exists() { - anyhow::bail!( - "No auth credentials found at {}. Run `kigi login` first.", - path.display() - ); - } - - let content = std::fs::read_to_string(path) - .map_err(|e| anyhow::anyhow!("failed to read {}: {e}", path.display()))?; - let entries: BTreeMap = serde_json::from_str(&content) - .map_err(|e| anyhow::anyhow!("failed to parse {}: {e}", path.display()))?; - - entries - .into_iter() - .find(|(_, e)| e.refresh_token.is_some() && e.oidc_issuer.is_some()) - .ok_or_else(|| { - anyhow::anyhow!( - "no OIDC auth entry found in {}. Run `kigi login` first.", - path.display() - ) - }) -} - -fn build_oidc_provider( - scope_key: String, - entry: &AuthEntry, - auth_path: PathBuf, -) -> anyhow::Result> { - let refresh_token = entry.refresh_token.as_ref().ok_or_else(|| { - anyhow::anyhow!("auth entry has no refresh_token — cannot refresh expired tokens") - })?; - let issuer = entry.oidc_issuer.as_ref().ok_or_else(|| { - anyhow::anyhow!("auth entry has no oidc_issuer — cannot refresh expired tokens") - })?; - let client_id = entry.oidc_client_id.as_ref().ok_or_else(|| { - anyhow::anyhow!("auth entry has no oidc_client_id — cannot refresh expired tokens") - })?; - - let mut builder = OidcAuthProviderBuilder::new(&entry.key, refresh_token, issuer, client_id); - - // Owner identity is surfaced via `AuthProvider::identity()` so consumers - // read it from this provider — no separate auth.json read. - builder = builder.user_id(&entry.user_id); - if let Some(ref pt) = entry.principal_type { - builder = builder.principal_type(pt); - } - if let Some(ref pid) = entry.principal_id { - builder = builder.principal_id(pid); - } - if let Some(exp) = entry.expires_at { - builder = builder.expires_at(exp); - } - - builder = builder.on_refresh(Arc::new(move |event: &RefreshEvent| { - if let Err(e) = write_refreshed_token(&auth_path, &scope_key, event) { - tracing::warn!(error = %e, "failed to persist refreshed token to auth.json"); - } - })); - - Ok(Arc::new(builder.build())) -} - -fn write_refreshed_token(path: &Path, scope_key: &str, event: &RefreshEvent) -> anyhow::Result<()> { - let content = std::fs::read_to_string(path)?; - let mut raw: serde_json::Value = serde_json::from_str(&content)?; - - let Some(obj) = raw.get_mut(scope_key).and_then(|e| e.as_object_mut()) else { - anyhow::bail!("auth entry '{scope_key}' not found while persisting refreshed token"); - }; - obj.insert( - "key".to_owned(), - serde_json::Value::String(event.access_token.clone()), - ); - if let Some(ref rt) = event.new_refresh_token { - obj.insert( - "refresh_token".to_owned(), - serde_json::Value::String(rt.clone()), - ); - } - if let Some(exp) = event.expires_at { - obj.insert( - "expires_at".to_owned(), - serde_json::Value::String(exp.to_rfc3339()), - ); - } - - write_json_atomic(path, &raw)?; - tracing::info!(path = %path.display(), "persisted refreshed token to auth.json"); - Ok(()) -} - -/// Atomically replace `path`: temp file (0600 on Unix) + fsync + rename. Avoids -/// the truncate-in-place corruption window when the long-lived binary rewrites -/// auth.json. -fn write_json_atomic(path: &Path, value: &serde_json::Value) -> anyhow::Result<()> { - use std::io::Write; - - let json = serde_json::to_string_pretty(value)?; - let tmp = path.with_extension(format!("json.{}.tmp", std::process::id())); - - let mut opts = std::fs::OpenOptions::new(); - opts.write(true).create(true).truncate(true); - #[cfg(unix)] - { - use std::os::unix::fs::OpenOptionsExt; - opts.mode(0o600); - } - - let mut file = opts - .open(&tmp) - .map_err(|e| anyhow::anyhow!("failed to open {}: {e}", tmp.display()))?; - file.write_all(json.as_bytes())?; - file.sync_all()?; - drop(file); - - #[cfg(windows)] - let _ = std::fs::remove_file(path); - - if let Err(e) = std::fs::rename(&tmp, path) { - let _ = std::fs::remove_file(&tmp); - return Err(anyhow::anyhow!("failed to replace {}: {e}", path.display())); - } - Ok(()) -} - -/// Build a hub auth provider for `hub_url`. `auth_config` overrides -/// the default credential path (`~/.kigi/auth.json`). -pub fn provider( - hub_url: &Url, - auth_config: Option<&Path>, -) -> anyhow::Result> { - let auth_path = match auth_config { - Some(p) => p.to_path_buf(), - None => default_auth_path()?, - }; - let (scope_key, entry) = read_auth_entry(&auth_path)?; - - let is_loopback = hub_url.scheme() == "ws" - && matches!(hub_url.host_str(), Some("localhost" | "127.0.0.1" | "::1")); - - if is_loopback { - tracing::info!("Using local-dev auth (loopback hub)"); - Ok(Arc::new(BearerWithIdentity { - identity: identity_from_entry(&entry), - token: entry.key.clone(), - })) - } else { - build_oidc_provider(scope_key, &entry, auth_path) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::io::Write; - - fn write_auth_json(dir: &std::path::Path, json: &str) -> PathBuf { - let path = dir.join("auth.json"); - let mut f = std::fs::File::create(&path).unwrap(); - f.write_all(json.as_bytes()).unwrap(); - path - } - - #[test] - fn read_auth_entry_picks_oidc_entry() { - let dir = tempfile::tempdir().unwrap(); - let path = write_auth_json( - dir.path(), - r#"{ - "legacy": { "key": "xai-plainkey", "user_id": "u1" }, - "oidc": { - "key": "eyJhbGciOiJFUzI1NiJ9.test", - "user_id": "u2", - "refresh_token": "rt", - "oidc_issuer": "https://auth.example.com", - "oidc_client_id": "client1" - } - }"#, - ); - - let (key, entry) = read_auth_entry(&path).unwrap(); - assert_eq!(key, "oidc"); - assert_eq!(entry.refresh_token.as_deref(), Some("rt")); - assert_eq!( - entry.oidc_issuer.as_deref(), - Some("https://auth.example.com") - ); - } - - #[test] - fn read_auth_entry_rejects_non_oidc() { - let dir = tempfile::tempdir().unwrap(); - let path = write_auth_json( - dir.path(), - r#"{ - "api_key": { "key": "xai-plainkey", "user_id": "u1" } - }"#, - ); - - let err = read_auth_entry(&path).unwrap_err(); - assert!(err.to_string().contains("no OIDC auth entry")); - } - - #[test] - fn read_auth_entry_missing_file() { - let path = PathBuf::from("/nonexistent/auth.json"); - let err = read_auth_entry(&path).unwrap_err(); - assert!(err.to_string().contains("No auth credentials")); - } - - #[test] - fn read_auth_entry_tolerates_extra_fields() { - let dir = tempfile::tempdir().unwrap(); - let path = write_auth_json( - dir.path(), - r#"{ - "scope": { - "key": "eyJhbGciOiJFUzI1NiJ9.tok", - "user_id": "u1", - "auth_mode": "oauth", - "create_time": "2026-01-01T00:00:00Z", - "email": "test@x.ai", - "first_name": "Test", - "refresh_token": "rt1", - "oidc_issuer": "https://auth.x.ai", - "oidc_client_id": "c1", - "some_future_field": true - } - }"#, - ); - - let (_key, entry) = read_auth_entry(&path).unwrap(); - assert_eq!(entry.refresh_token.as_deref(), Some("rt1")); - } - - #[test] - fn build_oidc_provider_requires_refresh_token() { - let entry = AuthEntry { - key: "eyJ.tok".into(), - user_id: "u1".into(), - refresh_token: None, - oidc_issuer: Some("https://auth.x.ai".into()), - oidc_client_id: Some("c1".into()), - principal_type: None, - principal_id: None, - expires_at: None, - }; - let err = build_oidc_provider("oidc".into(), &entry, PathBuf::from("/tmp/x")).unwrap_err(); - assert!(err.to_string().contains("refresh_token")); - } - - #[test] - fn build_oidc_provider_requires_issuer() { - let entry = AuthEntry { - key: "eyJ.tok".into(), - user_id: "u1".into(), - refresh_token: Some("rt".into()), - oidc_issuer: None, - oidc_client_id: Some("c1".into()), - principal_type: None, - principal_id: None, - expires_at: None, - }; - let err = build_oidc_provider("oidc".into(), &entry, PathBuf::from("/tmp/x")).unwrap_err(); - assert!(err.to_string().contains("oidc_issuer")); - } - - #[test] - fn build_oidc_provider_requires_client_id() { - let entry = AuthEntry { - key: "eyJ.tok".into(), - user_id: "u1".into(), - refresh_token: Some("rt".into()), - oidc_issuer: Some("https://auth.x.ai".into()), - oidc_client_id: None, - principal_type: None, - principal_id: None, - expires_at: None, - }; - let err = build_oidc_provider("oidc".into(), &entry, PathBuf::from("/tmp/x")).unwrap_err(); - assert!(err.to_string().contains("oidc_client_id")); - } - - #[test] - fn build_oidc_provider_succeeds_with_all_fields() { - let entry = AuthEntry { - key: "eyJ.tok".into(), - user_id: "u1".into(), - refresh_token: Some("rt".into()), - oidc_issuer: Some("https://auth.x.ai".into()), - oidc_client_id: Some("c1".into()), - principal_type: Some("Team".into()), - principal_id: Some("t1".into()), - expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), - }; - let provider = build_oidc_provider("oidc".into(), &entry, PathBuf::from("/tmp/x")).unwrap(); - let cred = provider.current(); - match cred { - kigi_computer_hub_sdk::AuthCredential::Bearer { token } => { - assert_eq!(token, "eyJ.tok"); - } - _ => panic!("expected Bearer"), - } - // Identity is surfaced from the parsed entry (no second auth.json read). - let id = provider.identity().expect("identity present"); - assert_eq!(id.user_id, "u1"); - assert_eq!(id.principal_type.as_deref(), Some("Team")); - assert_eq!(id.principal_id.as_deref(), Some("t1")); - } - - #[test] - fn write_refreshed_token_updates_jwt_entry() { - let dir = tempfile::tempdir().unwrap(); - let path = write_auth_json( - dir.path(), - r#"{ - "legacy": { "key": "xai-old", "user_id": "u1" }, - "oidc": { "key": "eyJ.old", "user_id": "u2", "refresh_token": "rt-old", "oidc_issuer": "https://auth.x.ai" } - }"#, - ); - - let event = RefreshEvent { - access_token: "eyJ.new".into(), - new_refresh_token: Some("rt-new".into()), - expires_at: None, - }; - write_refreshed_token(&path, "oidc", &event).unwrap(); - - let updated: serde_json::Value = - serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap(); - assert_eq!(updated["oidc"]["key"], "eyJ.new"); - assert_eq!(updated["oidc"]["refresh_token"], "rt-new"); - assert_eq!(updated["legacy"]["key"], "xai-old"); // untouched - } - - #[test] - fn write_refreshed_token_targets_exact_scope_key() { - // Non-sorted order: refresh must update the read-selected key ("aaa"), - // not the first in file order ("zzz"). - let dir = tempfile::tempdir().unwrap(); - let path = write_auth_json( - dir.path(), - r#"{ - "zzz": { "key": "eyJ.z", "refresh_token": "rt-z", "oidc_issuer": "https://auth.x.ai" }, - "aaa": { "key": "eyJ.a", "refresh_token": "rt-a", "oidc_issuer": "https://auth.x.ai" } - }"#, - ); - - let (key, _entry) = read_auth_entry(&path).unwrap(); - assert_eq!(key, "aaa"); - - let event = RefreshEvent { - access_token: "eyJ.a-new".into(), - new_refresh_token: Some("rt-a-new".into()), - expires_at: None, - }; - write_refreshed_token(&path, &key, &event).unwrap(); - - let updated: serde_json::Value = - serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap(); - assert_eq!(updated["aaa"]["key"], "eyJ.a-new"); - assert_eq!(updated["aaa"]["refresh_token"], "rt-a-new"); - assert_eq!(updated["zzz"]["key"], "eyJ.z"); - assert_eq!(updated["zzz"]["refresh_token"], "rt-z"); - } - - #[test] - fn write_refreshed_token_preserves_existing_rt_when_not_rotated() { - let dir = tempfile::tempdir().unwrap(); - let path = write_auth_json( - dir.path(), - r#"{ - "oidc": { "key": "eyJ.old", "user_id": "u1", "refresh_token": "rt-keep", "oidc_issuer": "https://auth.x.ai" } - }"#, - ); - - let event = RefreshEvent { - access_token: "eyJ.new".into(), - new_refresh_token: None, - expires_at: None, - }; - write_refreshed_token(&path, "oidc", &event).unwrap(); - - let updated: serde_json::Value = - serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap(); - assert_eq!(updated["oidc"]["key"], "eyJ.new"); - assert_eq!(updated["oidc"]["refresh_token"], "rt-keep"); - } - - #[test] - fn provider_loopback_uses_bearer() { - let dir = tempfile::tempdir().unwrap(); - let path = write_auth_json( - dir.path(), - r#"{ "oidc": { "key": "eyJ.tok", "user_id": "u1", "refresh_token": "rt", "oidc_issuer": "https://auth.x.ai", "oidc_client_id": "c1" } }"#, - ); - let url = Url::parse("ws://localhost:9988/v1/tools").unwrap(); - let auth = provider(&url, Some(&path)).unwrap(); - match auth.current() { - AuthCredential::Bearer { token } => assert_eq!(token, "eyJ.tok"), - _ => panic!("expected Bearer"), - } - // Loopback still surfaces identity from the same entry. - let id = auth.identity().expect("loopback identity present"); - assert_eq!(id.user_id, "u1"); - } -} diff --git a/crates/codegen/kigi-workspace/src/hub_channel.rs b/crates/codegen/kigi-workspace/src/hub_channel.rs deleted file mode 100644 index e8f5b75..0000000 --- a/crates/codegen/kigi-workspace/src/hub_channel.rs +++ /dev/null @@ -1,127 +0,0 @@ -//! Server-proxied workspace utilities. -//! -//! Helper functions for consuming server tool streams and extracting typed -//! notifications from server notification frames. These are used by both -//! the workspace crate (hub_server) and the shell crate (proxy-mode -//! session actors) to interact with the server. -//! -//! The `HubWorkspaceChannel` struct that previously lived here has been -//! removed. Sessions now call the server harness directly via `ToolContext`. - -use kigi_tools::notification::types::ToolNotification; -use kigi_workspace_types::WorkspaceEvent; - -pub use crate::hub_ids::WORKSPACE_RPC_TOOL_ID; -// Canonical in the client crate; re-exported for existing importers. -pub use kigi_workspace_client::consume_stream_terminal; - -/// Extract a `WorkspaceEvent` from a custom `ToolNotificationFrame`. -pub fn extract_workspace_event( - frame: &kigi_tool_protocol::ToolNotificationFrame, -) -> Option { - use kigi_tool_protocol::WireToolNotification; - match &frame.notification { - WireToolNotification::Custom(c) => serde_json::from_value(c.payload.clone()).ok(), - _ => None, - } -} - -/// Extract a `ToolNotification` from a custom `ToolNotificationFrame`. -/// -/// Currently has no producer; kept for a future per-session `tool.notify` path. -pub fn extract_tool_notification( - frame: &kigi_tool_protocol::ToolNotificationFrame, -) -> Option { - use kigi_tool_protocol::WireToolNotification; - match &frame.notification { - WireToolNotification::Custom(c) => { - serde_json::from_value::(c.payload.clone()).ok() - } - _ => None, - } -} - -// --------------------------------------------------------------------------- -// Tests -// --------------------------------------------------------------------------- - -#[cfg(test)] -mod tests { - use super::*; - use kigi_tool_protocol::{ToolId, ToolNotificationFrame, WireToolNotification}; - - #[test] - fn extract_workspace_event_custom_valid() { - let event = WorkspaceEvent::ToolsChanged { - session_id: "main".into(), - }; - let payload = serde_json::to_value(&event).unwrap(); - let frame = ToolNotificationFrame::custom( - ToolId::new("workspace_events").unwrap(), - "workspace_event", - payload, - ); - let result = extract_workspace_event(&frame); - assert!(result.is_some(), "should parse valid WorkspaceEvent"); - match result.unwrap() { - WorkspaceEvent::ToolsChanged { session_id } => { - assert_eq!(session_id, "main"); - } - other => panic!("expected ToolsChanged, got {other:?}"), - } - } - - #[test] - fn extract_workspace_event_invalid_payload() { - let frame = ToolNotificationFrame::custom( - ToolId::new("workspace_events").unwrap(), - "workspace_event", - serde_json::json!({"not_a_valid_event": true}), - ); - assert!( - extract_workspace_event(&frame).is_none(), - "invalid payload should return None" - ); - } - - #[test] - fn extract_workspace_event_known_variant_returns_none() { - let frame = ToolNotificationFrame { - tool_call_id: None, - tool_id: Some(ToolId::new("workspace_events").unwrap()), - notification: WireToolNotification::Known(serde_json::json!({})), - }; - assert!( - extract_workspace_event(&frame).is_none(), - "Known variant should return None" - ); - } - - #[test] - fn extract_tool_notification_invalid_payload() { - let frame = ToolNotificationFrame::custom( - ToolId::new("workspace_tool_notifications").unwrap(), - "tool_notification", - serde_json::json!({"not_a_notification": true}), - ); - assert!( - extract_tool_notification(&frame).is_none(), - "invalid payload should return None" - ); - } - - #[test] - fn extract_tool_notification_known_variant_returns_none() { - let frame = ToolNotificationFrame { - tool_call_id: None, - tool_id: Some(ToolId::new("workspace_tool_notifications").unwrap()), - notification: WireToolNotification::Known(serde_json::json!({})), - }; - assert!( - extract_tool_notification(&frame).is_none(), - "Known variant should return None" - ); - } - - // consume_stream_terminal tests live in kigi-workspace-client. -} diff --git a/crates/codegen/kigi-workspace/src/hub_ids.rs b/crates/codegen/kigi-workspace/src/hub_ids.rs deleted file mode 100644 index f427b17..0000000 --- a/crates/codegen/kigi-workspace/src/hub_ids.rs +++ /dev/null @@ -1,10 +0,0 @@ -//! Hub tool ID constants, canonical in `kigi_workspace_types::rpc` -//! and re-exported here for existing importers. -//! -//! `WORKSPACE_TOOL_NOTIFICATIONS_TOOL_ID` is intentionally producer-less today; -//! see [`crate::hub_channel::extract_tool_notification`]. - -pub use kigi_workspace_types::rpc::{ - WORKSPACE_CLIENT_EXT_NOTIFICATIONS_TOOL_ID, WORKSPACE_EVENTS_TOOL_ID, WORKSPACE_RPC_TOOL_ID, - WORKSPACE_TOOL_NOTIFICATIONS_TOOL_ID, -}; diff --git a/crates/codegen/kigi-workspace/src/hub_server.rs b/crates/codegen/kigi-workspace/src/hub_server.rs deleted file mode 100644 index d47d528..0000000 --- a/crates/codegen/kigi-workspace/src/hub_server.rs +++ /dev/null @@ -1,2999 +0,0 @@ -//! Workspace-side RPC handler for server-proxied workspace method calls. -//! -//! [`WorkspaceRpcHandler`] implements [`ToolServerHandler`] and dispatches -//! `workspace.*` JSON-RPC methods to [`WorkspaceHandle`]. Registered on -//! the `ToolServer` with tool_id `workspace_rpc`. -use crate::error::{WorkspaceError, WorkspaceResult}; -use crate::handle::WorkspaceHandle; -use crate::hub_ids::WORKSPACE_RPC_TOOL_ID; -use crate::rpc_envelope::{RpcEnvelope, envelope_err}; -use crate::workspace_ops::WorkspaceOp; -use async_trait::async_trait; -use chrono::{DateTime, Utc}; -use kigi_computer_hub_sdk::ToolServerHandler; -use kigi_tool_protocol::{HookEvent, HookFrame, SessionId, ToolId, ToolServerEvictParams}; -use kigi_tool_runtime::{ - ToolCallContext, ToolError, ToolErrorKind, ToolStream, TypedToolOutput, terminal_only, -}; -use kigi_tool_types::ToolDescription; -use kigi_tools::computer::types::TaskKind; -use kigi_tools::implementations::grok_build::scheduler::interval::interval_to_human; -use kigi_tools::implementations::grok_build::scheduler::types::{ - SchedulerCommand, SchedulerHandle, -}; -use kigi_tools::registry::types::FinalizedToolset; -use kigi_tools::types::resources::Terminal; -use kigi_workspace_types::rpc::workspace::{ - BackgroundTaskSnapshotWire, ScheduledTaskSnapshotWire, TasksSnapshotResponse, -}; -use prometheus::{HistogramVec, IntCounterVec, register_histogram_vec, register_int_counter_vec}; -use serde_json::Value; -/// Deprecation monitor for the self-attested `caller_session_id` param: -/// `kind="param_mismatch"` — the param disagreed with the server-bound envelope -/// session (envelope trusted); `kind="envelope_absent"` — no envelope -/// session, the param was used as a compat fallback. Enforcement -/// (envelope-only identity) waits for this to be flat zero. -static WORKSPACE_RPC_CALLER_MISMATCH_TOTAL: std::sync::LazyLock = - std::sync::LazyLock::new(|| { - register_int_counter_vec!( - "grok_workspace_rpc_caller_mismatch_total", - "Mutation RPCs whose caller_session_id param was not backed by a matching \ - server-bound envelope session, by method and kind", - &["method", "kind"] - ) - .unwrap() - }); -/// Audit trail for the deliberate-mutation RPC surface -/// (`update_tool_config` / `drop_session` / `configure_mcp`). -static WORKSPACE_RPC_MUTATION_TOTAL: std::sync::LazyLock = - std::sync::LazyLock::new(|| { - register_int_counter_vec!( - "grok_workspace_rpc_mutation_total", - "Session-mutating workspace RPC calls, by method and outcome", - &["method", "outcome"] - ) - .unwrap() - }); -/// Every dispatched `workspace.*` RPC, by method and result. Unrecognized -/// methods collapse to the `unknown` label to keep cardinality bounded. -static WORKSPACE_RPC_REQUESTS_TOTAL: std::sync::LazyLock = - std::sync::LazyLock::new(|| { - register_int_counter_vec!( - "grok_workspace_rpc_requests_total", - "Workspace RPC dispatches, by method and result", - &["method", "result"] - ) - .unwrap() - }); -/// Per-method wall-clock duration of a `workspace.*` RPC dispatch. -static WORKSPACE_RPC_DURATION_SECONDS: std::sync::LazyLock = - std::sync::LazyLock::new(|| { - register_histogram_vec!( - "grok_workspace_rpc_duration_seconds", - "Workspace RPC dispatch duration", - &["method"], - vec![ - 0.001, 0.005, 0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 5.0, 10.0 - ] - ) - .unwrap() - }); -const UNKNOWN_METHOD_LABEL: &str = "unknown"; -/// Prefix of the [`WorkspaceError::HubError`] for an unrecognized method. Shared -/// by the dispatch default arm and the metric classifier so the "collapse to -/// `unknown`" decision cannot drift from the error it keys on. -const UNKNOWN_METHOD_ERR_PREFIX: &str = "unknown workspace method:"; -/// Zero-init this module's metric families. See [`crate::init_metrics`]. -pub(crate) fn init_metrics() { - WORKSPACE_RPC_REQUESTS_TOTAL - .with_label_values(&[UNKNOWN_METHOD_LABEL, "error"]) - .inc_by(0); - let _ = WORKSPACE_RPC_DURATION_SECONDS.with_label_values(&[UNKNOWN_METHOD_LABEL]); -} -/// Resolve the caller identity for a mutation RPC: the server-bound envelope -/// session is authoritative; the deprecated `caller_session_id` param is -/// only used when no envelope session exists (old call paths). Both -/// divergences are counted on [`WORKSPACE_RPC_CALLER_MISMATCH_TOTAL`]. -fn resolve_mutation_caller<'a>( - method: &'static str, - bound_session: Option<&'a str>, - param_caller: Option<&'a str>, -) -> WorkspaceResult<&'a str> { - match (bound_session, param_caller) { - (Some(envelope), Some(param)) => { - if envelope != param { - WORKSPACE_RPC_CALLER_MISMATCH_TOTAL - .with_label_values(&[method, "param_mismatch"]) - .inc(); - tracing::warn!( - method, envelope_session = % envelope, param_caller = % param, - "caller_session_id param disagrees with the server-bound envelope session; \ - trusting the envelope" - ); - } - Ok(envelope) - } - (Some(envelope), None) => Ok(envelope), - (None, Some(param)) => { - WORKSPACE_RPC_CALLER_MISMATCH_TOTAL - .with_label_values(&[method, "envelope_absent"]) - .inc(); - Ok(param) - } - (None, None) => Err(WorkspaceError::HubError(format!( - "{method}: missing caller identity (no bound session and no caller_session_id)" - ))), - } -} -/// Audit-log and count a mutation RPC on [`WORKSPACE_RPC_MUTATION_TOTAL`]. -/// Failures log at WARN because that arm carries rejected cross-session -/// forgeries (`Unauthorized`), the audit trail's most interesting event. -fn record_mutation_rpc( - method: &'static str, - caller: &str, - target: &str, - result: &WorkspaceResult, -) { - let outcome = match result { - Ok(_) => "ok", - Err(_) => "error", - }; - WORKSPACE_RPC_MUTATION_TOTAL - .with_label_values(&[method, outcome]) - .inc(); - match result { - Ok(_) => tracing::info!(method, caller, target, "workspace mutation rpc"), - Err(e) => { - tracing::warn!( - method, caller, target, error = % e, "workspace mutation rpc failed" - ); - } - } -} -/// No-op notifier for RPC-driven worktree creation. -struct NoOpNotifier; -#[async_trait] -impl crate::worktree::WorktreeNotificationSender for NoOpNotifier { - async fn send_worktree_status(&self, _progress: crate::worktree::WorktreeStatus) {} -} -/// Env escape hatch for the client-facing `workspace.client_fs_*` ops. -/// -/// Default **on**; setting `WORKSPACE_CLIENT_FS_QUERIES=0` (or `false`) -/// disables the ops with a graceful `HubError` that the remote caller -/// maps to a fallback. Read per call — flipping the variable needs no -/// process restart and tests can toggle it under a lock. -fn client_fs_queries_enabled() -> bool { - !matches!( - std::env::var("WORKSPACE_CLIENT_FS_QUERIES").as_deref(), - Ok("0") | Ok("false") - ) -} -/// Reject `workspace.client_fs_*` dispatch when the escape hatch is off. -fn ensure_client_fs_queries_enabled() -> WorkspaceResult<()> { - if client_fs_queries_enabled() { - Ok(()) - } else { - Err(WorkspaceError::HubError( - "client fs queries disabled on this workspace".into(), - )) - } -} -/// Generic dispatch helper: deserialize params, execute, serialize result. -async fn dispatch_op( - params: Value, - ws: &WorkspaceHandle, - session_id: Option<&str>, -) -> WorkspaceResult { - let req: Op = serde_json::from_value(params) - .map_err(|e| WorkspaceError::HubError(format!("invalid params for {}: {e}", Op::METHOD)))?; - let result = req.execute(ws, session_id).await?; - serde_json::to_value(result) - .map_err(|e| WorkspaceError::HubError(format!("{}: {e}", Op::METHOD))) -} -/// List a session's outstanding (not-completed) background terminal tasks from -/// the session toolset's `TerminalBackend` resource, mapped to the slim wire -/// DTO. Empty when the session has no terminal backend. Source of truth for the -/// `workspace.list_background_tasks` RPC (post-compaction system-reminder state). -async fn list_outstanding_background_tasks( - toolset: &kigi_tools::registry::types::FinalizedToolset, -) -> Vec { - use kigi_tools::computer::types::TaskKind; - use kigi_tools::types::resources::Terminal; - use kigi_tools::types::tool::ToolKind; - use kigi_workspace_types::rpc::workspace::BackgroundTaskSummaryWire; - let terminal = { - let res = toolset.resources.lock().await; - res.get::().map(|t| t.0.clone()) - }; - let Some(terminal) = terminal else { - return Vec::new(); - }; - let execute_name = toolset.tool_name_for_kind(ToolKind::Execute); - let monitor_name = toolset.tool_name_for_kind(ToolKind::Monitor); - terminal - .list_tasks() - .await - .into_iter() - .filter(|t| !t.completed) - .map(|t| { - let command = t - .display_command - .clone() - .filter(|s| !s.is_empty()) - .unwrap_or(t.command); - let tool_name = match t.kind { - TaskKind::Monitor => monitor_name.clone(), - TaskKind::Bash => execute_name.clone(), - }; - BackgroundTaskSummaryWire { - task_id: t.task_id, - command, - tool_name, - } - }) - .collect() -} -/// Point-in-time snapshot of the session's outstanding background terminal -/// tasks and live scheduled tasks. -async fn tasks_snapshot(toolset: &FinalizedToolset) -> TasksSnapshotResponse { - let (terminal, scheduler) = { - let res = toolset.resources.lock().await; - ( - res.get::().map(|t| t.0.clone()), - res.get::().cloned(), - ) - }; - let background_tasks = match terminal { - Some(terminal) => terminal - .list_tasks() - .await - .into_iter() - .filter(|t| !t.completed) - .map(|t| { - let command = t - .display_command - .clone() - .filter(|s| !s.is_empty()) - .unwrap_or(t.command); - BackgroundTaskSnapshotWire { - task_id: t.task_id, - command, - kind: match t.kind { - TaskKind::Bash => "bash".to_owned(), - TaskKind::Monitor => "monitor".to_owned(), - }, - started_at: DateTime::::from(t.start_time).to_rfc3339(), - } - }) - .collect(), - None => Vec::new(), - }; - let scheduled_tasks = match scheduler { - Some(handle) => { - let (reply_tx, reply_rx) = tokio::sync::oneshot::channel(); - let _ = handle.0.send(SchedulerCommand::List { reply: reply_tx }); - reply_rx - .await - .unwrap_or_default() - .into_iter() - .map(|t| ScheduledTaskSnapshotWire { - task_id: t.id.clone(), - prompt: t.prompt.clone(), - human_schedule: interval_to_human(t.interval_secs), - next_fire_at: t.next_fire_at().to_rfc3339(), - recurring: t.recurring, - created_at: t.created_at.to_rfc3339(), - }) - .collect() - } - None => Vec::new(), - }; - TasksSnapshotResponse { - background_tasks, - scheduled_tasks, - } -} -/// List the session's TODO items (via `todo_write`) from the session toolset's -/// `State` resource, mapped to the slim wire DTO. Empty when the -/// session has no todo state. Source of truth for the `workspace.list_todos` -/// RPC (post-compaction system-reminder state). -async fn list_session_todos( - toolset: &kigi_tools::registry::types::FinalizedToolset, -) -> Vec { - use kigi_tools::implementations::grok_build::todo::{TodoState, TodoStatus}; - use kigi_tools::types::resources::State; - use kigi_workspace_types::rpc::workspace::TodoSummaryWire; - let res = toolset.resources.lock().await; - let Some(state) = res.get::>() else { - return Vec::new(); - }; - state - .0 - .todo_items_with_ids() - .map(|(id, item)| { - let status = match item.status { - TodoStatus::Pending => "pending", - TodoStatus::InProgress => "in_progress", - TodoStatus::Completed => "completed", - TodoStatus::Cancelled => "cancelled", - }; - TodoSummaryWire { - id: id.to_string(), - content: item.content.clone(), - status: status.to_string(), - } - }) - .collect() -} -/// Routes JSON-RPC `workspace.*` method calls to [`WorkspaceHandle`]. -pub(crate) struct WorkspaceRpcHandler { - workspace: WorkspaceHandle, -} -impl WorkspaceRpcHandler { - pub(crate) fn new(workspace: WorkspaceHandle) -> Self { - Self { workspace } - } - /// Route a `workspace.*` method; `bound_session` is the caller's server-bound session. - async fn dispatch( - &self, - method: &str, - params: Value, - bound_session: Option<&str>, - ) -> WorkspaceResult { - use crate::file_system::ContentSearchRequest; - use crate::file_system::{ - FsDeleteFileReq, FsExistsReq, FsListReq, FsReadFileReq, FsWriteFileReq, - }; - use crate::session::checkpoint::TurnBoundary; - use crate::workspace_ops::*; - use crate::worktree::{ApplyWorktreeRequest, CreateWorktreeRequest, RemoveWorktreeRequest}; - use kigi_workspace_types::rpc::git::{GitBranchInfoReq, GitMetadataReq}; - use kigi_workspace_types::rpc::search::FuzzyStatusReq; - use kigi_workspace_types::rpc::skills::DiscoverPluginsReq; - use kigi_workspace_types::rpc::workspace::{ - ConfigureMcpReq, DropSessionReq, InstallPluginReq, ListBackgroundTasksReq, - ListBackgroundTasksResponse, ListTodosReq, ListTodosResponse, LoadEnvrcReq, - LoadPermissionsReq, LoadProjectConfigReq, RefreshPluginsReq, ResolveFileReferencesReq, - TasksSnapshotReq, ToolDefinitionsReq, UpdateToolConfigReq, - }; - use kigi_workspace_types::rpc::worktree::WorktreeCreateSyncReq; - tracing::debug!(method, "workspace rpc dispatch"); - let params = if params.is_null() { - serde_json::json!({}) - } else { - params - }; - match method { - ::METHOD => { - let cwd = self.workspace.root_cwd()?; - let cwd_str = cwd.to_string_lossy().to_string(); - let os = std::env::consts::OS; - let shell = std::env::var("SHELL") - .ok() - .and_then(|s| { - std::path::Path::new(&s) - .file_name() - .map(|n| n.to_string_lossy().to_string()) - }) - .unwrap_or_else(|| "sh".to_string()); - Ok(serde_json::json!({ "os" : os, "shell" : shell, "cwd" : cwd_str, })) - } - ::METHOD => { - static DEPRECATION_WARNING: std::sync::Once = std::sync::Once::new(); - DEPRECATION_WARNING.call_once(|| { - tracing::warn!( - "workspace.git_status is deprecated and will be removed in a future \ - release. Use workspace.git_status_ext with format: \"prompt\" instead." - ); - }); - let cwd = self.workspace.root_cwd()?; - let result = crate::file_system::git_status(cwd) - .await - .map_err(|e| WorkspaceError::HubError(e.to_string()))?; - Ok(Value::String(result)) - } - ::METHOD => { - let cwd = self.workspace.root_cwd()?; - match crate::session::git::git_info(&cwd).await { - Ok(info) => serde_json::to_value(info) - .map_err(|e| WorkspaceError::HubError(e.to_string())), - Err(_) => Ok(Value::Null), - } - } - ::METHOD => { - let session_id = params - .get("session_id") - .and_then(Value::as_str) - .ok_or_else(|| WorkspaceError::HubError("missing session_id".into()))?; - let session = self - .workspace - .session(session_id) - .ok_or_else(|| WorkspaceError::SessionNotFound(session_id.into()))?; - let defs = session.toolset().tool_definitions(); - serde_json::to_value(defs).map_err(|e| WorkspaceError::HubError(e.to_string())) - } - ::METHOD => { - let session_id = params - .get("session_id") - .and_then(Value::as_str) - .ok_or_else(|| WorkspaceError::HubError("missing session_id".into()))?; - let session = self - .workspace - .session(session_id) - .ok_or_else(|| WorkspaceError::SessionNotFound(session_id.into()))?; - let toolset = session.toolset(); - let tasks = list_outstanding_background_tasks(toolset.as_ref()).await; - serde_json::to_value(ListBackgroundTasksResponse { tasks }) - .map_err(|e| WorkspaceError::HubError(e.to_string())) - } - ::METHOD => { - let session_id = params - .get("session_id") - .and_then(Value::as_str) - .ok_or_else(|| WorkspaceError::HubError("missing session_id".into()))?; - let session = self - .workspace - .session(session_id) - .ok_or_else(|| WorkspaceError::SessionNotFound(session_id.into()))?; - let toolset = session.toolset(); - let snapshot = tasks_snapshot(toolset.as_ref()).await; - serde_json::to_value(snapshot).map_err(|e| WorkspaceError::HubError(e.to_string())) - } - ::METHOD => { - let session_id = params - .get("session_id") - .and_then(Value::as_str) - .ok_or_else(|| WorkspaceError::HubError("missing session_id".into()))?; - let session = self - .workspace - .session(session_id) - .ok_or_else(|| WorkspaceError::SessionNotFound(session_id.into()))?; - let toolset = session.toolset(); - let todos = list_session_todos(toolset.as_ref()).await; - serde_json::to_value(ListTodosResponse { todos }) - .map_err(|e| WorkspaceError::HubError(e.to_string())) - } - ::METHOD => { - let caller = resolve_mutation_caller( - "update_tool_config", - bound_session, - params - .get("caller_session_id") - .and_then(Value::as_str) - .filter(|s| !s.is_empty()), - )?; - let session_id = params - .get("session_id") - .and_then(Value::as_str) - .ok_or_else(|| WorkspaceError::HubError("missing session_id".into()))? - .to_owned(); - let new_config = serde_json::from_value( - params - .get("new_config") - .cloned() - .ok_or_else(|| WorkspaceError::HubError("missing new_config".into()))?, - ) - .map_err(|e| WorkspaceError::HubError(format!("invalid new_config: {e}")))?; - let result = self - .workspace - .update_tool_config(caller, &session_id, new_config) - .await; - record_mutation_rpc("update_tool_config", caller, &session_id, &result); - result.map(|()| Value::Null) - } - ::METHOD => { - let caller = resolve_mutation_caller( - "drop_session", - bound_session, - params - .get("caller_session_id") - .and_then(Value::as_str) - .filter(|s| !s.is_empty()), - )?; - let target = params - .get("session_id") - .and_then(Value::as_str) - .ok_or_else(|| WorkspaceError::HubError("missing session_id".into()))?; - let result = self.workspace.drop_session(caller, target); - record_mutation_rpc("drop_session", caller, target, &result); - result.map(|()| Value::Null) - } - ::METHOD => { - let refs: Vec = params - .get("refs") - .and_then(|v| serde_json::from_value(v.clone()).ok()) - .unwrap_or_default(); - let cwd = self.workspace.root_cwd()?; - let mut results = Vec::new(); - for ref_path in &refs { - let full_path = if std::path::Path::new(ref_path).is_absolute() { - std::path::PathBuf::from(ref_path) - } else { - cwd.join(ref_path) - }; - let exists = full_path.exists(); - let content = if exists { - tokio::fs::read_to_string(&full_path).await.ok() - } else { - None - }; - results.push(serde_json::json!( - { "path" : full_path.to_string_lossy(), "ref" : ref_path, - "exists" : exists, "content" : content, } - )); - } - Ok(Value::Array(results)) - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - ensure_client_fs_queries_enabled()?; - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - ensure_client_fs_queries_enabled()?; - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - ensure_client_fs_queries_enabled()?; - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - let cwd = self.workspace.root_cwd()?; - let skills = - crate::discovery::discover_skills(&cwd, self.workspace.shared.skills_config()) - .await; - Ok(Value::Array(skills)) - } - ::METHOD => { - let cwd = self.workspace.root_cwd()?; - let files = crate::discovery::discover_agents_md(&cwd).await; - Ok(Value::Array(files)) - } - ::METHOD => { - let cwd = self.workspace.root_cwd()?; - let plugins = crate::discovery::discover_plugins( - &cwd, - self.workspace.shared.plugin_discovery_config(), - &crate::discovery::PluginTrustStore::load(), - true, - ); - Ok(Value::Array(plugins)) - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - let cwd = self.workspace.root_cwd()?; - Ok(crate::discovery::load_project_config(&cwd)) - } - ::METHOD => { - let cwd = self.workspace.root_cwd()?; - Ok(crate::discovery::load_permissions(&cwd).await) - } - ::METHOD => { - let cwd = self.workspace.root_cwd()?; - let env = crate::envrc::load_envrc_or_empty(&cwd); - serde_json::to_value(env).map_err(|e| WorkspaceError::HubError(e.to_string())) - } - ::METHOD => { - let _ = params; - Ok(Value::Null) - } - ::METHOD => { - let cwd = self.workspace.root_cwd()?; - let plugins = crate::discovery::discover_plugins( - &cwd, - self.workspace.shared.plugin_discovery_config(), - &crate::discovery::PluginTrustStore::load(), - true, - ); - Ok(Value::Array(plugins)) - } - ::METHOD => { - let session_id = bound_session.ok_or_else(|| { - WorkspaceError::HubError("configure_mcp requires a bound session".into()) - })?; - let configs: Vec = serde_json::from_value( - params - .get("mcp_servers") - .cloned() - .unwrap_or(Value::Array(vec![])), - ) - .map_err(|e| WorkspaceError::HubError(format!("invalid mcp_servers: {e}")))?; - let result = async { - if self.workspace.session(session_id).is_none() { - tracing::info!( - session_id, - "workspace.configure_mcp: session not found, creating on demand" - ); - match self - .workspace - .create_session_with_config( - session_id, - None, - None, - crate::capability::CapabilityMode::All, - None, - true, - ) - { - Ok(_session) => {} - Err(WorkspaceError::SessionAlreadyExists(_)) => { - tracing::debug!( - session_id, - "workspace.configure_mcp: session created concurrently, using existing" - ); - } - Err(e) => return Err(e), - } - } - self.workspace.start_session_mcp_servers(session_id, configs).await - } - .await; - record_mutation_rpc("configure_mcp", "self", session_id, &result); - serde_json::to_value(&result?) - .map_err(|e| WorkspaceError::HubError(format!("serialize McpStartResult: {e}"))) - } - ::METHOD => { - let cwd = self.workspace.root_cwd()?; - let metadata = - crate::session::git::resolve_persisted_session_git_metadata_sync(&cwd); - Ok(serde_json::to_value(metadata).unwrap_or(Value::Null)) - } - ::METHOD => { - let search_id = params - .get("search_id") - .and_then(Value::as_str) - .ok_or_else(|| WorkspaceError::HubError("missing search_id".into()))?; - let results = self.workspace.fuzzy_get_results(search_id).await; - match results { - Some(data) => serde_json::to_value(data) - .map_err(|e| WorkspaceError::HubError(e.to_string())), - None => Ok(Value::Null), - } - } - "workspace.worktree_create_from_worktree" - | ::METHOD => { - dispatch_op::(params, &self.workspace, None) - .await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - let req: crate::worktree::CreateWorktreeRequest = serde_json::from_value(params) - .map_err(|e| { - WorkspaceError::HubError(format!("invalid create_sync params: {e}")) - })?; - let result = crate::worktree::create_worktree_streaming(&req, &NoOpNotifier).await; - serde_json::to_value(result).map_err(|e| WorkspaceError::HubError(e.to_string())) - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, bound_session).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, bound_session).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, bound_session).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, bound_session).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, bound_session) - .await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, bound_session) - .await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, bound_session).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, bound_session).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, bound_session).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, bound_session).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - dispatch_op::(params, &self.workspace, None).await - } - ::METHOD => { - let req: BeginPromptReq = serde_json::from_value(params).map_err(|e| { - WorkspaceError::HubError(format!("invalid params for begin_prompt: {e}")) - })?; - self.workspace - .session(&req.session_id) - .ok_or_else(|| WorkspaceError::SessionNotFound(req.session_id.clone()))?; - self.workspace - .on_turn_boundary( - &req.session_id, - TurnBoundary::rewind_begin(req.prompt_index), - ) - .await; - Ok(Value::Null) - } - ::METHOD => { - let req: EndPromptReq = serde_json::from_value(params).map_err(|e| { - WorkspaceError::HubError(format!("invalid params for end_prompt: {e}")) - })?; - self.workspace - .session(&req.session_id) - .ok_or_else(|| WorkspaceError::SessionNotFound(req.session_id.clone()))?; - self.workspace - .on_turn_boundary( - &req.session_id, - TurnBoundary::rewind_finalize(req.prompt_index), - ) - .await; - Ok(Value::Null) - } - ::METHOD => { - let req: GetRewindPointsReq = serde_json::from_value(params).map_err(|e| { - WorkspaceError::HubError(format!("invalid params for get_rewind_points: {e}")) - })?; - let session = self - .workspace - .session(&req.session_id) - .ok_or_else(|| WorkspaceError::SessionNotFound(req.session_id.clone()))?; - let points = session - .file_state_tracker() - .get_rewind_points_normalized(session.cwd()) - .await; - serde_json::to_value(points).map_err(|e| WorkspaceError::HubError(e.to_string())) - } - ::METHOD => { - let req: RewindToReq = serde_json::from_value(params).map_err(|e| { - WorkspaceError::HubError(format!("invalid params for rewind_to: {e}")) - })?; - self.workspace - .session(&req.session_id) - .ok_or_else(|| WorkspaceError::SessionNotFound(req.session_id.clone()))?; - let response = self - .workspace - .rewind_to(&req.session_id, req.target_prompt_index) - .await; - serde_json::to_value(response).map_err(|e| WorkspaceError::HubError(e.to_string())) - } - _ => { - tracing::warn!(method, "unknown workspace rpc method"); - Err(WorkspaceError::HubError(format!( - "{UNKNOWN_METHOD_ERR_PREFIX} {method}" - ))) - } - } - } -} -#[async_trait] -impl ToolServerHandler for WorkspaceRpcHandler { - fn tool_id(&self) -> ToolId { - ToolId::new(WORKSPACE_RPC_TOOL_ID).expect("constant is a valid ToolId") - } - fn description(&self) -> ToolDescription { - ToolDescription::new( - WORKSPACE_RPC_TOOL_ID, - "Routes workspace RPC calls to the local workspace handle.", - ) - } - fn input_schema(&self) -> Option { - Some(serde_json::json!( - { "type" : "object", "properties" : { "method" : { "type" : "string", - "description" : "The workspace.* method to invoke" }, "params" : { "type" - : "object", "description" : "Method parameters" } }, "required" : - ["method"] } - )) - } - async fn handle_call(&self, ctx: ToolCallContext, args: Value) -> ToolStream { - let tool_id = self.tool_id(); - let method = match args.get("method").and_then(Value::as_str) { - Some(m) => m, - None => { - return terminal_only(Err(ToolError::new( - ToolErrorKind::InvalidArguments, - "missing required field: method", - ))); - } - }; - tracing::debug!("workspace rpc call from server"); - let params = args - .get("params") - .cloned() - .unwrap_or(Value::Object(Default::default())); - let bound_session = ctx.extensions.get::(); - let start = std::time::Instant::now(); - let result = self - .dispatch( - method, - params, - bound_session.as_deref().map(|s| s.0.as_str()), - ) - .await; - let is_unknown_method = matches!( - & result, Err(WorkspaceError::HubError(msg)) if msg - .starts_with(UNKNOWN_METHOD_ERR_PREFIX) - ); - let method_label = if is_unknown_method { - UNKNOWN_METHOD_LABEL - } else { - method - }; - WORKSPACE_RPC_REQUESTS_TOTAL - .with_label_values(&[method_label, if result.is_ok() { "ok" } else { "error" }]) - .inc(); - WORKSPACE_RPC_DURATION_SECONDS - .with_label_values(&[method_label]) - .observe(start.elapsed().as_secs_f64()); - let envelope = match result { - Ok(value) => RpcEnvelope::ok(value), - Err(ref e) => envelope_err(e), - }; - let envelope = - serde_json::to_value(envelope).expect("RpcEnvelope serialization is infallible"); - terminal_only(Ok(TypedToolOutput::from_value(tool_id, envelope))) - } - async fn handle_hook(&self, session_id: SessionId, frame: HookFrame) { - match frame.event { - HookEvent::Cancel => { - if let Some(call_id) = &frame.call_id { - tracing::info!(% session_id, % call_id, "cancel hook received"); - self.workspace - .cancel_tool_call(session_id.as_str(), call_id.as_str()); - } else { - tracing::info!(% session_id, "cancel hook received (session-wide)"); - self.workspace.cancel_all_tool_calls(session_id.as_str()); - } - } - HookEvent::SessionEnded => { - tracing::info!(% session_id, "session_ended hook received"); - self.workspace - .teardown_session_mcp(session_id.as_str()) - .await; - self.workspace.on_session_ended(session_id.as_str()); - } - HookEvent::Custom { kind, payload } => { - use kigi_tool_protocol::turn_hook::{ - AFTER_TURN_KIND, AfterTurnPayload, BEFORE_TURN_KIND, BeforeTurnPayload, - }; - match kind.as_str() { - BEFORE_TURN_KIND => { - match serde_json::from_value::(payload) { - Ok(p) => { - tracing::info!( - session = % session_id, turn = p.turn_number, model = % p - .model_id, "before_turn hook received" - ); - self.workspace.on_before_turn(session_id.as_str(), &p).await; - } - Err(e) => { - tracing::warn!( - error = % e, "before_turn payload deserialization failed" - ); - } - } - } - AFTER_TURN_KIND => match serde_json::from_value::(payload) { - Ok(p) => { - tracing::info!( - session = % session_id, turn = p.turn_number, outcome = ? p - .outcome, duration_ms = p.duration_ms, - "after_turn hook received" - ); - self.workspace.on_after_turn(session_id.as_str(), &p).await; - } - Err(e) => { - tracing::warn!( - error = % e, "after_turn payload deserialization failed" - ); - } - }, - _ => { - tracing::debug!( - kind = % kind, session = % session_id, - "unrecognized custom hook kind" - ); - } - } - } - HookEvent::Pause | HookEvent::Resume => { - tracing::debug!( - % session_id, event = ? frame.event, "hook not yet implemented" - ); - } - } - } - async fn handle_hook_request(&self, session_id: SessionId, frame: HookFrame) -> Option { - use kigi_tool_protocol::turn_hook::{self, TurnHookRequest}; - let HookEvent::Custom { kind, payload } = frame.event else { - return None; - }; - if kind != turn_hook::TURN_HOOK_KIND { - return None; - } - let no_op = || serde_json::to_value(turn_hook::HookReply::default()).ok(); - if self.workspace.shared.activity_tracker.is_draining() - || self.workspace.session(session_id.as_str()).is_none() - { - return no_op(); - } - let request: TurnHookRequest = match serde_json::from_value(payload) { - Ok(r) => r, - Err(e) => { - tracing::warn!(error = % e, % session_id, "invalid turn hook request"); - return no_op(); - } - }; - let reply = self - .workspace - .compute_turn_injections(session_id.as_str(), &request) - .await; - Some(serde_json::to_value(&reply).unwrap_or(Value::Null)) - } - /// Hub-issued `tool_server.evict`. Always tears the evicted session down - /// (MCP bridges + activity/writer state, like the `SessionEnded` hook), then - /// runs the global two-phase drain **only** when no other session survives — - /// a global drain shuts down the *shared* upload queue, which must not happen - /// while another session is live. Idempotent across fan-out and safe for an - /// already-gone session id. - /// - /// Contract: the server-supplied `grace_period_ms` budgets the drain and is - /// therefore honored only when evicting the **last** live session. For a - /// multi-session workspace the evicted session is dropped immediately - /// (no per-session drain) because the shared upload queue cannot be flushed - /// or closed without affecting the survivors. - async fn handle_evict(&self, params: ToolServerEvictParams) { - let sid = params.session_id.as_str(); - self.workspace.teardown_session_mcp(sid).await; - self.workspace.on_session_ended(sid); - let (became_empty, start_drain) = { - let mut sessions = self.workspace.shared.sessions.write(); - if let Some(session) = sessions.remove(sid) { - session.abort_system_notify_forwarder(); - session.shutdown_terminal_backend(); - } - let empty = sessions.is_empty(); - let already_winding_down = self.workspace.activity_tracker().is_draining(); - let start = empty && !already_winding_down; - if start { - self.workspace.activity_tracker().set_draining(); - } - (empty, start) - }; - if !start_drain { - if became_empty { - tracing::info!( - session = % params.session_id, reason = % params.reason, - "workspace: hub evict — already draining/shutting down; dropped session only" - ); - } else { - tracing::info!( - session = % params.session_id, reason = % params.reason, - "workspace: hub evict — other sessions live; dropped session only" - ); - } - return; - } - let grace = std::time::Duration::from_millis(params.grace_period_ms); - tracing::info!( - session = % params.session_id, reason = % params.reason, grace_period_ms = - params.grace_period_ms, - "workspace: hub evict — last session; commencing two-phase drain" - ); - let unfinished = self - .workspace - .two_phase_drain(grace, crate::handle::DrainReason::Evict) - .await; - if unfinished > 0 { - tracing::warn!( - session = % params.session_id, unfinished, - "workspace: hub evict drain left items pending" - ); - } - self.workspace.activity_tracker().set_shutting_down(); - } -} -#[cfg(test)] -mod tests { - use super::*; - use crate::capability::CapabilityMode; - use crate::handle::tests::{background_capable_cfg, make_handle, start_background_sleep}; - use kigi_tool_protocol::turn_hook; - use kigi_tools::implementations::grok_build::scheduler::types::ScheduledTask; - /// Helper: consume the first item from a ToolStream. - async fn next_item( - stream: &mut ToolStream, - ) -> Option> { - use std::task::Context; - std::future::poll_fn(|cx: &mut Context<'_>| stream.as_mut().poll_next(cx)).await - } - fn turn_hook_frame(session: &str, req: &turn_hook::TurnHookRequest) -> HookFrame { - HookFrame::custom_request( - SessionId::new(session).unwrap(), - "hk-test".to_owned(), - turn_hook::TURN_HOOK_KIND.to_owned(), - serde_json::to_value(req).unwrap(), - ) - } - #[tokio::test] - async fn handle_hook_request_turn_hook_returns_reply() { - let handler = WorkspaceRpcHandler::new(make_handle()); - let req = turn_hook::TurnHookRequest::Before(turn_hook::BeforeTurnPayload { - turn_number: 1, - model_id: "grok-3".to_owned(), - yolo_mode: false, - conversation_message_count: 0, - session_relationship: "primary".to_owned(), - schema_version: "1.0".to_owned(), - }); - let value = handler - .handle_hook_request( - SessionId::new("main").unwrap(), - turn_hook_frame("main", &req), - ) - .await - .expect("turn hook claimed"); - let reply: turn_hook::HookReply = serde_json::from_value(value).unwrap(); - assert_eq!(reply, turn_hook::HookReply::default()); - } - #[tokio::test] - async fn handle_hook_request_ignores_non_turn_hook_kind() { - let handler = WorkspaceRpcHandler::new(make_handle()); - let frame = HookFrame::custom_request( - SessionId::new("main").unwrap(), - "hk-x".to_owned(), - "some_other_kind".to_owned(), - serde_json::json!({}), - ); - assert!( - handler - .handle_hook_request(SessionId::new("main").unwrap(), frame) - .await - .is_none() - ); - } - #[tokio::test] - async fn handle_hook_request_unbound_session_is_noop() { - let handler = WorkspaceRpcHandler::new(make_handle()); - let req = turn_hook::TurnHookRequest::Before(turn_hook::BeforeTurnPayload { - turn_number: 1, - model_id: "grok-3".to_owned(), - yolo_mode: false, - conversation_message_count: 0, - session_relationship: "primary".to_owned(), - schema_version: "1.0".to_owned(), - }); - let value = handler - .handle_hook_request( - SessionId::new("never-bound").unwrap(), - turn_hook_frame("never-bound", &req), - ) - .await - .expect("fail-open no-op reply"); - let reply: turn_hook::HookReply = serde_json::from_value(value).unwrap(); - assert_eq!(reply, turn_hook::HookReply::default()); - } - #[tokio::test] - async fn dispatch_unknown_method_returns_hub_error() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let result = handler - .dispatch("workspace.nonexistent", Value::Null, None) - .await; - assert!(matches!(result, Err(WorkspaceError::HubError(msg)) if msg - .contains("unknown workspace method"))); - } - /// A hub evict runs the two-phase drain then settles into terminal - /// ShuttingDown (not a lingering Draining) for an evicted workspace. - #[tokio::test] - async fn handle_evict_triggers_two_phase_drain() { - use kigi_tool_protocol::ToolServerLifecycleStatus; - let handle = make_handle(); - let tracker = handle.activity_tracker().clone(); - let handler = WorkspaceRpcHandler::new(handle); - assert!(!tracker.is_draining(), "not draining before evict"); - handler - .handle_evict(ToolServerEvictParams { - session_id: SessionId::new("main").expect("valid session id"), - reason: "preemption".into(), - grace_period_ms: 200, - }) - .await; - let snap = tracker.snapshot(); - assert_eq!( - snap.status, - ToolServerLifecycleStatus::ShuttingDown, - "an evicted workspace must end in terminal ShuttingDown, not lingering Draining" - ); - assert!( - snap.drain_started_ms.is_some(), - "evict drain must stamp drain_started_ms" - ); - } - /// A hub evict shuts the evicted session's terminal backend down - /// explicitly: the actor stops even while other `Arc`s to the backend are - /// still alive (mirrors `drop_session_shuts_down_terminal_backend_explicitly`). - #[tokio::test] - async fn handle_evict_shuts_down_terminal_backend_explicitly() { - let handle = make_handle(); - let session = handle.session("main").expect("main session exists"); - let retained_backend = session.terminal_backend().clone(); - let retained_toolset = session.toolset(); - drop(session); - let handler = WorkspaceRpcHandler::new(handle); - handler - .handle_evict(ToolServerEvictParams { - session_id: SessionId::new("main").expect("valid session id"), - reason: "preemption".into(), - grace_period_ms: 100, - }) - .await; - crate::handle::tests::assert_backend_stops(&retained_backend).await; - drop(retained_toolset); - } - /// Isolation matrix #1/#3 at the RPC surface: `workspace.list_background_tasks` - /// (the post-compaction reminder source of truth) stays truthful across - /// both rebind shapes. The task stays listed through a `Reused` rebind - /// AND a `Reresolved` toolset swap — reading it through each rebind's - /// CURRENT toolset — and leaves the list only when explicitly killed. - #[tokio::test] - async fn list_background_tasks_rpc_stays_truthful_across_rebinds() { - use crate::capability::CapabilityMode; - use crate::handle::RebindOutcome; - use crate::handle::tests::{background_capable_cfg, start_background_sleep}; - use crate::session::tool_config::test_support::tc; - use kigi_tools::registry::types::ToolServerConfig; - use kigi_workspace_types::rpc::workspace::ListBackgroundTasksResponse; - let handle = make_handle(); - let cfg = background_capable_cfg(); - let session = handle - .create_session_with_config( - "bg-rpc", - None, - Some(cfg.clone()), - CapabilityMode::All, - None, - false, - ) - .expect("create background-capable session"); - session.set_bind_tool_config_fingerprint(serde_json::to_value(&cfg).ok()); - let out_dir = tempfile::tempdir().expect("temp dir"); - let bg = start_background_sleep(&session, out_dir.path(), "bg-rpc-task").await; - let handler = WorkspaceRpcHandler::new(handle.clone()); - async fn list_tasks( - handler: &WorkspaceRpcHandler, - ) -> Vec { - let value = handler - .dispatch( - "workspace.list_background_tasks", - serde_json::json!({ "session_id" : "bg-rpc" }), - Some("bg-rpc"), - ) - .await - .expect("list_background_tasks rpc"); - serde_json::from_value::(value) - .expect("decode response") - .tasks - } - let tasks = list_tasks(&handler).await; - assert_eq!(tasks.len(), 1, "the running task must be listed"); - assert_eq!(tasks[0].task_id, bg.task_id); - assert_eq!( - tasks[0].tool_name.as_deref(), - Some("run_terminal_cmd"), - "the creator tool is named from the live toolset" - ); - let (_, outcome) = handle - .rebind_existing_hub_session( - "bg-rpc", - Some(cfg.clone()), - serde_json::to_value(&cfg).ok(), - ) - .await - .expect("session exists"); - assert_eq!(outcome, RebindOutcome::Reused); - let tasks = list_tasks(&handler).await; - assert_eq!(tasks.len(), 1, "the task must survive a reused rebind"); - let read_only = ToolServerConfig { - tools: vec![tc( - "GrokBuild:read_file", - Some(kigi_tools::types::tool::ToolKind::Read), - )], - behavior_preset: None, - }; - let (_, outcome) = handle - .rebind_existing_hub_session( - "bg-rpc", - Some(read_only.clone()), - serde_json::to_value(&read_only).ok(), - ) - .await - .expect("session exists"); - assert_eq!(outcome, RebindOutcome::Reresolved); - let tasks = list_tasks(&handler).await; - assert_eq!(tasks.len(), 1, "the task must survive the toolset swap"); - assert_eq!(tasks[0].task_id, bg.task_id); - assert_eq!( - tasks[0].tool_name, None, - "the swapped-in toolset has no execute tool to name" - ); - session.terminal_backend().kill_task(&bg.task_id).await; - let tasks = list_tasks(&handler).await; - assert!( - tasks.is_empty(), - "a killed task must leave the outstanding list: {tasks:?}" - ); - } - /// `workspace.tasks_snapshot` (GC-614 part 3): returns the outstanding - /// background task with kind/started_at, plus scheduled tasks (empty when - /// no scheduler resource exists), and drops the task once killed. - #[tokio::test] - async fn tasks_snapshot_rpc_lists_outstanding_background_tasks() { - let handle = make_handle(); - let cfg = background_capable_cfg(); - let session = handle - .create_session_with_config( - "snap-rpc", - None, - Some(cfg.clone()), - CapabilityMode::All, - None, - false, - ) - .expect("create background-capable session"); - session.set_bind_tool_config_fingerprint(serde_json::to_value(&cfg).ok()); - let out_dir = tempfile::tempdir().expect("temp dir"); - let bg = start_background_sleep(&session, out_dir.path(), "snap-rpc-task").await; - let handler = WorkspaceRpcHandler::new(handle.clone()); - async fn snapshot(handler: &WorkspaceRpcHandler) -> TasksSnapshotResponse { - let value = handler - .dispatch( - "workspace.tasks_snapshot", - serde_json::json!({ "session_id" : "snap-rpc" }), - Some("snap-rpc"), - ) - .await - .expect("tasks_snapshot rpc"); - serde_json::from_value(value).expect("decode response") - } - let snap = snapshot(&handler).await; - assert_eq!( - snap.background_tasks.len(), - 1, - "the running task must be listed" - ); - let task = &snap.background_tasks[0]; - assert_eq!(task.task_id, bg.task_id); - assert_eq!(task.kind, "bash"); - assert!( - DateTime::parse_from_rfc3339(&task.started_at).is_ok(), - "started_at must be RFC3339: {}", - task.started_at - ); - assert!( - snap.scheduled_tasks.is_empty(), - "no scheduler resource in this toolset: {:?}", - snap.scheduled_tasks - ); - session.terminal_backend().kill_task(&bg.task_id).await; - let snap = snapshot(&handler).await; - assert!( - snap.background_tasks.is_empty(), - "a killed task must leave the snapshot: {:?}", - snap.background_tasks - ); - { - let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); - tokio::spawn(async move { - while let Some(cmd) = rx.recv().await { - if let SchedulerCommand::List { reply } = cmd { - let mut task = ScheduledTask::new(300, "check CI".into(), true, false); - task.id = "loop-1".into(); - let _ = reply.send(vec![task]); - } - } - }); - let toolset = session.toolset(); - toolset.resources.lock().await.insert(SchedulerHandle(tx)); - } - let snap = snapshot(&handler).await; - assert_eq!(snap.scheduled_tasks.len(), 1); - let loop_task = &snap.scheduled_tasks[0]; - assert_eq!(loop_task.task_id, "loop-1"); - assert_eq!(loop_task.prompt, "check CI"); - assert_eq!(loop_task.human_schedule, "every 5 minutes"); - assert!(loop_task.recurring); - assert!( - DateTime::parse_from_rfc3339(&loop_task.next_fire_at).is_ok(), - "next_fire_at must be RFC3339: {}", - loop_task.next_fire_at - ); - } - /// Evicting one session while another is live must NOT global-drain (which - /// would close the shared queue for the survivor) — even when the evicted - /// id is no longer in the session map. - #[tokio::test] - async fn handle_evict_keeps_queue_when_other_sessions_live() { - let handle = make_handle(); - handle - .create_session("other") - .expect("create second session"); - let tracker = handle.activity_tracker().clone(); - let handler = WorkspaceRpcHandler::new(handle); - handler - .handle_evict(ToolServerEvictParams { - session_id: SessionId::new("ghost").expect("valid session id"), - reason: "idle_timeout".into(), - grace_period_ms: 200, - }) - .await; - assert!( - !tracker.is_draining(), - "evict of an absent id with live sessions must not global-drain" - ); - } - /// Evicting one of several live sessions removes *that* session (full - /// teardown), keeps the survivors, and does not global-drain the shared - /// queue. The drain decision is made on the post-removal map. - #[tokio::test] - async fn handle_evict_nonlast_removes_session_and_preserves_survivors() { - let handle = make_handle(); - handle - .create_session("other") - .expect("create second session"); - let tracker = handle.activity_tracker().clone(); - let handler = WorkspaceRpcHandler::new(handle.clone()); - handler - .handle_evict(ToolServerEvictParams { - session_id: SessionId::new("other").expect("valid session id"), - reason: "idle_timeout".into(), - grace_period_ms: 200, - }) - .await; - assert!( - handle.session("other").is_none(), - "the evicted session must be removed from the map" - ); - assert!( - handle.session("main").is_some(), - "a surviving session must be kept" - ); - assert!( - !tracker.is_draining(), - "evicting a non-last session must not global-drain the shared queue" - ); - } - /// Once a terminal evict drain has started, a racing `bind`/create must be - /// rejected so the shared upload queue is never torn down under a fresh - /// session (race #3). - #[tokio::test] - async fn bind_rejected_after_evict_drain() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle.clone()); - handler - .handle_evict(ToolServerEvictParams { - session_id: SessionId::new("main").expect("valid session id"), - reason: "preemption".into(), - grace_period_ms: 100, - }) - .await; - assert!(matches!( - handle.create_session("late"), - Err(WorkspaceError::ShuttingDown) - )); - } - /// A duplicate / retried evict of the last session must not re-run the - /// drain or downgrade terminal `ShuttingDown` back to `Draining`. - #[tokio::test] - async fn repeat_evict_does_not_redrain() { - use kigi_tool_protocol::ToolServerLifecycleStatus; - let handle = make_handle(); - let tracker = handle.activity_tracker().clone(); - let handler = WorkspaceRpcHandler::new(handle); - let params = || ToolServerEvictParams { - session_id: SessionId::new("main").expect("valid session id"), - reason: "preemption".into(), - grace_period_ms: 100, - }; - handler.handle_evict(params()).await; - assert_eq!( - tracker.snapshot().status, - ToolServerLifecycleStatus::ShuttingDown - ); - handler.handle_evict(params()).await; - assert_eq!( - tracker.snapshot().status, - ToolServerLifecycleStatus::ShuttingDown, - "a repeat evict must not downgrade terminal ShuttingDown to Draining" - ); - } - #[tokio::test] - async fn dispatch_tool_definitions_returns_known_tools() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let params = serde_json::json!({ "session_id" : "main" }); - let result = handler - .dispatch("workspace.tool_definitions", params, None) - .await; - let value = result.expect("should succeed"); - let arr = value.as_array().expect("should be array"); - assert!(!arr.is_empty(), "main session should have tools"); - let names: Vec<&str> = arr - .iter() - .filter_map(|d| { - d.get("function") - .and_then(|f| f.get("name")) - .and_then(|n| n.as_str()) - }) - .collect(); - assert!( - names.contains(&"read_file"), - "should contain read_file: {names:?}" - ); - } - #[tokio::test] - async fn dispatch_tool_definitions_unknown_session() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let params = serde_json::json!({ "session_id" : "ghost" }); - let result = handler - .dispatch("workspace.tool_definitions", params, None) - .await; - assert!(matches!(result, Err(WorkspaceError::SessionNotFound(_)))); - } - #[tokio::test] - async fn dispatch_get_all_hunks_returns_array() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let result = handler - .dispatch("workspace.get_all_hunks", Value::Null, Some("main")) - .await; - assert!(result.is_ok()); - assert!(result.unwrap().is_array()); - } - #[tokio::test] - async fn dispatch_get_session_summary_returns_object_or_null() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let result = handler - .dispatch("workspace.get_session_summary", Value::Null, Some("main")) - .await; - let value = result.expect("should succeed"); - assert!( - value.is_object() || value.is_null(), - "expected object or null, got {value}" - ); - } - #[tokio::test] - async fn dispatch_discover_skills_returns_array() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let result = handler - .dispatch("workspace.discover_skills", Value::Null, None) - .await; - assert!(result.is_ok()); - assert!(result.unwrap().is_array()); - } - #[tokio::test] - async fn dispatch_load_envrc_returns_object() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let result = handler - .dispatch("workspace.load_envrc", Value::Null, None) - .await; - assert!(result.is_ok()); - assert!(result.unwrap().is_object()); - } - #[tokio::test] - async fn dispatch_drop_session_self_succeeds() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle.clone()); - let params = serde_json::json!( - { "caller_session_id" : "main", "session_id" : "main" } - ); - let result = handler - .dispatch("workspace.drop_session", params, None) - .await; - assert!(result.is_ok(), "dropping own session should succeed"); - assert!(handle.session("main").is_none(), "session should be gone"); - } - #[tokio::test] - async fn dispatch_update_tool_config_missing_params() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let result = handler - .dispatch("workspace.update_tool_config", serde_json::json!({}), None) - .await; - assert!( - matches!(result, Err(WorkspaceError::HubError(ref msg)) if msg - .contains("missing")), - "got {result:?}" - ); - } - fn caller_mismatch_count(method: &str, kind: &str) -> u64 { - WORKSPACE_RPC_CALLER_MISMATCH_TOTAL - .with_label_values(&[method, kind]) - .get() - } - fn baseline_config_value() -> Value { - serde_json::to_value(crate::session::tool_config::test_support::baseline_config()) - .expect("baseline config serializes") - } - /// With both an envelope session and a (spoofed) param, the envelope - /// wins: the call is authorized as the envelope session and the - /// mismatch is counted. - #[tokio::test] - async fn dispatch_update_tool_config_envelope_overrides_param() { - let mismatch_before = caller_mismatch_count("update_tool_config", "param_mismatch"); - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let params = serde_json::json!( - { "caller_session_id" : "spoofed", "session_id" : "main", "new_config" : - baseline_config_value(), } - ); - let result = handler - .dispatch("workspace.update_tool_config", params, Some("main")) - .await; - assert!( - result.is_ok(), - "envelope caller == target must authorize: {result:?}" - ); - assert!( - caller_mismatch_count("update_tool_config", "param_mismatch") > mismatch_before, - "the param/envelope disagreement must be counted" - ); - } - /// A forged `caller_session_id` param cannot authorize a cross-session - /// mutation: the envelope session is the caller and differs from the - /// target, so the target's caller-equals-target check rejects it. - #[tokio::test] - async fn dispatch_update_tool_config_envelope_cross_session_unauthorized() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle.clone()); - let params = serde_json::json!( - { "caller_session_id" : "main", "session_id" : "main", "new_config" : - baseline_config_value(), } - ); - let result = handler - .dispatch("workspace.update_tool_config", params, Some("other")) - .await; - assert!( - matches!(result, Err(WorkspaceError::Unauthorized { .. })), - "got {result:?}" - ); - assert!( - handle.session("main").is_some(), - "the target session must be untouched" - ); - } - /// Compat: without an envelope session (old call paths) the param is - /// still honored, and the fallback is counted for the deprecation - /// monitor. - #[tokio::test] - async fn dispatch_update_tool_config_param_fallback_without_envelope() { - let absent_before = caller_mismatch_count("update_tool_config", "envelope_absent"); - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let params = serde_json::json!( - { "caller_session_id" : "main", "session_id" : "main", "new_config" : - baseline_config_value(), } - ); - let result = handler - .dispatch("workspace.update_tool_config", params, None) - .await; - assert!(result.is_ok(), "param fallback must authorize: {result:?}"); - assert!( - caller_mismatch_count("update_tool_config", "envelope_absent") > absent_before, - "the envelope-absent fallback must be counted" - ); - } - /// The intended steady state once clients drop the deprecated param: - /// envelope-only identity (no `caller_session_id` in params) authorizes. - /// Counter non-advance is asserted by - /// [`resolve_mutation_caller_clean_arms_count_nothing`], which uses a - /// test-unique method label — the real label is shared with concurrently - /// running dispatch tests, so an equality assert here would flake. - #[tokio::test] - async fn dispatch_update_tool_config_envelope_only_without_param() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let params = serde_json::json!( - { "session_id" : "main", "new_config" : baseline_config_value(), } - ); - let result = handler - .dispatch("workspace.update_tool_config", params, Some("main")) - .await; - assert!( - result.is_ok(), - "envelope-only identity must authorize: {result:?}" - ); - } - /// The two clean `resolve_mutation_caller` arms — envelope-only and - /// envelope+matching-param — resolve to the envelope without ticking - /// either deprecation-monitor kind. - #[test] - fn resolve_mutation_caller_clean_arms_count_nothing() { - const METHOD: &str = "test_clean_arms"; - let mismatch_before = caller_mismatch_count(METHOD, "param_mismatch"); - let absent_before = caller_mismatch_count(METHOD, "envelope_absent"); - let caller = resolve_mutation_caller(METHOD, Some("sess"), None) - .expect("envelope-only must resolve"); - assert_eq!(caller, "sess"); - let caller = resolve_mutation_caller(METHOD, Some("sess"), Some("sess")) - .expect("matching param must resolve"); - assert_eq!(caller, "sess"); - assert_eq!( - caller_mismatch_count(METHOD, "param_mismatch"), - mismatch_before, - "clean arms must not count a mismatch" - ); - assert_eq!( - caller_mismatch_count(METHOD, "envelope_absent"), - absent_before, - "clean arms must not count an envelope-absent fallback" - ); - } - /// `drop_session` gets the same envelope-derived identity: a spoofed - /// param is ignored when the envelope authorizes the drop, and the - /// mutation audit counter advances. - #[tokio::test] - async fn dispatch_drop_session_envelope_overrides_param() { - let mutation_before = WORKSPACE_RPC_MUTATION_TOTAL - .with_label_values(&["drop_session", "ok"]) - .get(); - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle.clone()); - let params = serde_json::json!( - { "caller_session_id" : "spoofed", "session_id" : "main" } - ); - let result = handler - .dispatch("workspace.drop_session", params, Some("main")) - .await; - assert!(result.is_ok(), "{result:?}"); - assert!(handle.session("main").is_none(), "session should be gone"); - assert!( - WORKSPACE_RPC_MUTATION_TOTAL - .with_label_values(&["drop_session", "ok"]) - .get() - > mutation_before, - "the mutation audit counter must advance" - ); - } - /// A cross-session drop forged via the param is rejected off the - /// envelope identity and the target survives. - #[tokio::test] - async fn dispatch_drop_session_envelope_cross_session_unauthorized() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle.clone()); - let params = serde_json::json!( - { "caller_session_id" : "main", "session_id" : "main" } - ); - let result = handler - .dispatch("workspace.drop_session", params, Some("observer-ish")) - .await; - assert!( - matches!(result, Err(WorkspaceError::Unauthorized { .. })), - "got {result:?}" - ); - assert!( - handle.session("main").is_some(), - "the target session must survive" - ); - } - /// `configure_mcp`'s on-demand session create opts into system - /// notifications, like every other sandbox-path creator. - #[tokio::test] - async fn dispatch_configure_mcp_on_demand_create_enables_system_notifications() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle.clone()); - let _ = handler - .dispatch( - "workspace.configure_mcp", - serde_json::json!({ "mcp_servers" : [] }), - Some("mcp-fresh"), - ) - .await; - let session = handle - .session("mcp-fresh") - .expect("session created on demand"); - assert!( - session.system_notifications(), - "the on-demand created session must forward system notifications" - ); - } - #[tokio::test] - async fn dispatch_hunk_action_unknown_action() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let params = serde_json::json!( - { "action" : { "hunk_id" : "test-id", "action" : "dance" } } - ); - let result = handler - .dispatch("workspace.hunk_action", params, None) - .await; - assert!( - matches!(result, Err(WorkspaceError::HubError(_))), - "expected HubError for invalid action enum, got {result:?}" - ); - } - #[tokio::test] - async fn dispatch_hunk_action_malformed_json() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let params = serde_json::json!({ "action" : "not-an-object" }); - let result = handler - .dispatch("workspace.hunk_action", params, None) - .await; - assert!( - matches!(result, Err(WorkspaceError::HubError(_))), - "expected HubError for malformed action, got {result:?}" - ); - } - #[tokio::test] - async fn dispatch_hunk_action_missing_action_field() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let params = serde_json::json!({}); - let result = handler - .dispatch("workspace.hunk_action", params, None) - .await; - assert!( - matches!(result, Err(WorkspaceError::HubError(ref msg)) if msg - .contains("missing field")), - "got {result:?}" - ); - } - #[tokio::test] - async fn dispatch_hunk_file_action_missing_path() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let params = serde_json::json!({ "action" : "accept" }); - let result = handler - .dispatch("workspace.hunk_file_action", params, None) - .await; - assert!( - matches!(result, Err(WorkspaceError::HubError(ref msg)) if msg - .contains("missing field")), - "got {result:?}" - ); - } - #[tokio::test] - async fn dispatch_hunk_turn_action_missing_prompt_index() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let params = serde_json::json!({ "action" : "accept" }); - let result = handler - .dispatch("workspace.hunk_turn_action", params, None) - .await; - assert!( - matches!(result, Err(WorkspaceError::HubError(ref msg)) if msg - .contains("missing field")), - "got {result:?}" - ); - } - #[tokio::test] - async fn dispatch_hunk_all_action_invalid_action() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let params = serde_json::json!({ "action" : "explode" }); - let result = handler - .dispatch("workspace.hunk_all_action", params, None) - .await; - assert!( - matches!(result, Err(WorkspaceError::HubError(_))), - "expected HubError for invalid action enum, got {result:?}" - ); - } - #[tokio::test] - async fn dispatch_hunk_get_all_file_contents_returns_array() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let result = handler - .dispatch( - "workspace.hunk_get_all_file_contents", - Value::Null, - Some("main"), - ) - .await; - assert!(result.is_ok()); - assert!(result.unwrap().is_array()); - } - #[tokio::test] - async fn dispatch_hunk_get_staged_files_returns_array() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let result = handler - .dispatch("workspace.hunk_get_staged_files", Value::Null, Some("main")) - .await; - assert!(result.is_ok()); - assert!(result.unwrap().is_array()); - } - #[tokio::test] - async fn dispatch_fuzzy_open_returns_search_id() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let result = handler - .dispatch( - "workspace.fuzzy_open", - serde_json::json!({ "hidden" : false }), - None, - ) - .await; - let value = result.expect("should succeed"); - assert!( - value.as_str().is_some_and(|s| !s.is_empty()), - "response should be a non-empty search_id string: {value}" - ); - } - #[tokio::test] - async fn dispatch_fuzzy_close_unknown_id() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let result = handler - .dispatch( - "workspace.fuzzy_close", - serde_json::json!({ "search_id" : "nonexistent" }), - None, - ) - .await; - let value = result.expect("should succeed"); - assert!(!value.as_bool().expect("response should be a bool")); - } - #[tokio::test] - async fn dispatch_fuzzy_change_missing_search_id() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let result = handler - .dispatch( - "workspace.fuzzy_change", - serde_json::json!({ "query" : "test" }), - None, - ) - .await; - assert!( - matches!(result, Err(WorkspaceError::HubError(ref msg)) if msg - .contains("missing field")), - "got {result:?}" - ); - } - #[tokio::test] - async fn dispatch_fuzzy_search_missing_search_id() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let result = handler - .dispatch("workspace.fuzzy_search", serde_json::json!({}), None) - .await; - assert!( - matches!(result, Err(WorkspaceError::HubError(ref msg)) if msg - .contains("missing search_id")), - "got {result:?}" - ); - } - #[tokio::test] - async fn dispatch_fuzzy_open_then_close_roundtrip() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let open_result = handler - .dispatch( - "workspace.fuzzy_open", - serde_json::json!({ "hidden" : false }), - None, - ) - .await - .expect("open should succeed"); - let search_id = open_result - .as_str() - .expect("open response should be a search_id string") - .to_owned(); - let close_result = handler - .dispatch( - "workspace.fuzzy_close", - serde_json::json!({ "search_id" : search_id }), - None, - ) - .await - .expect("close should succeed"); - assert!( - close_result - .as_bool() - .expect("close response should be a bool") - ); - let close_again = handler - .dispatch( - "workspace.fuzzy_close", - serde_json::json!({ "search_id" : search_id }), - None, - ) - .await - .expect("close again should succeed"); - assert!( - !close_again - .as_bool() - .expect("close-again response should be a bool") - ); - } - #[tokio::test] - async fn handle_call_wraps_in_envelope_with_value() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let mut ctx = ToolCallContext::default(); - ctx.extensions - .insert(kigi_tool_runtime::SessionContext("main".to_owned())); - let args = serde_json::json!( - { "method" : "workspace.get_session_summary", "params" : {} } - ); - let mut stream = handler.handle_call(ctx, args).await; - let item = next_item(&mut stream).await.expect("should have terminal"); - match item { - kigi_tool_runtime::ToolStreamItem::Terminal(Ok(typed)) => { - let ok_val = typed - .value - .get("ok") - .expect("envelope should have 'ok' key"); - assert!( - ok_val.is_object() || ok_val.is_null(), - "ok value should be object or null, got {ok_val}" - ); - } - other => panic!("expected Terminal(Ok), got {other:?}"), - } - } - #[tokio::test] - async fn handle_call_error_envelope() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let ctx = ToolCallContext::default(); - let args = serde_json::json!( - { "method" : "workspace.nonexistent", "params" : {} } - ); - let mut stream = handler.handle_call(ctx, args).await; - let item = next_item(&mut stream).await.expect("should have terminal"); - match item { - kigi_tool_runtime::ToolStreamItem::Terminal(Ok(typed)) => { - assert!( - typed.value.get("err").is_some(), - "envelope should have 'err' key: {}", - typed.value - ); - let err = typed.value.get("err").unwrap(); - assert!(err.get("code").is_some()); - assert!(err.get("message").is_some()); - } - other => panic!("expected Terminal(Ok(envelope)), got {other:?}"), - } - } - /// `handle_call` records the RPC metrics: a known method increments its - /// per-method `ok` series, and an unrecognized method collapses to - /// `method="unknown",result="error"` — never creating a per-bad-method - /// series (the cardinality-bounding guarantee). - #[tokio::test] - async fn handle_call_records_rpc_metrics_and_collapses_unknown_method() { - let handler = WorkspaceRpcHandler::new(make_handle()); - let ok_before = WORKSPACE_RPC_REQUESTS_TOTAL - .with_label_values(&["workspace.get_session_summary", "ok"]) - .get(); - let dur_samples_before = WORKSPACE_RPC_DURATION_SECONDS - .with_label_values(&["workspace.get_session_summary"]) - .get_sample_count(); - let mut ctx = ToolCallContext::default(); - ctx.extensions - .insert(kigi_tool_runtime::SessionContext("main".to_owned())); - let mut stream = handler - .handle_call( - ctx, - serde_json::json!( - { "method" : "workspace.get_session_summary", "params" : {} } - ), - ) - .await; - let _ = next_item(&mut stream).await; - assert!( - WORKSPACE_RPC_REQUESTS_TOTAL - .with_label_values(&["workspace.get_session_summary", "ok"]) - .get() - > ok_before, - "a known ok RPC must increment its per-method ok counter" - ); - assert!( - WORKSPACE_RPC_DURATION_SECONDS - .with_label_values(&["workspace.get_session_summary"]) - .get_sample_count() - > dur_samples_before, - "the dispatch must observe the per-method duration histogram" - ); - const BOGUS: &str = "workspace.__test_bogus_method_zzz"; - let unknown_before = WORKSPACE_RPC_REQUESTS_TOTAL - .with_label_values(&[UNKNOWN_METHOD_LABEL, "error"]) - .get(); - let mut stream = handler - .handle_call( - ToolCallContext::default(), - serde_json::json!({ "method" : BOGUS, "params" : {} }), - ) - .await; - let _ = next_item(&mut stream).await; - assert!( - WORKSPACE_RPC_REQUESTS_TOTAL - .with_label_values(&[UNKNOWN_METHOD_LABEL, "error"]) - .get() - > unknown_before, - "an unrecognized method must increment the collapsed unknown/error counter" - ); - let has_bogus_series = prometheus::gather() - .iter() - .filter(|mf| mf.name() == "grok_workspace_rpc_requests_total") - .flat_map(|mf| mf.get_metric()) - .any(|m| { - m.get_label() - .iter() - .any(|l| l.name() == "method" && l.value() == BOGUS) - }); - assert!( - !has_bogus_series, - "the raw bad method must collapse to `unknown`, never its own series" - ); - } - #[tokio::test] - async fn dispatch_git_stage_non_git_dir_returns_error() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let result = handler - .dispatch("workspace.git_stage", serde_json::json!({}), None) - .await; - assert!(result.is_err(), "non-git dir should error"); - } - #[tokio::test] - async fn dispatch_git_commit_missing_message_returns_error() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let result = handler - .dispatch("workspace.git_commit", serde_json::json!({}), None) - .await; - assert!( - matches!(result, Err(WorkspaceError::HubError(ref msg)) if msg - .contains("missing field")) - ); - } - #[tokio::test] - async fn dispatch_git_checkout_missing_branch_returns_error() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let result = handler - .dispatch("workspace.git_checkout", serde_json::json!({}), None) - .await; - assert!( - matches!(result, Err(WorkspaceError::HubError(ref msg)) if msg - .contains("missing field")) - ); - } - #[tokio::test] - async fn dispatch_git_stage_content_missing_fields() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let result = handler - .dispatch("workspace.git_stage_content", serde_json::json!({}), None) - .await; - assert!( - matches!(result, Err(WorkspaceError::HubError(ref msg)) if msg - .contains("missing")) - ); - } - #[tokio::test] - async fn handle_hook_before_turn_sets_turn_state() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle.clone()); - let payload = turn_hook::BeforeTurnPayload { - turn_number: 1, - model_id: "grok-3".to_string(), - yolo_mode: false, - conversation_message_count: 0, - session_relationship: "primary".to_string(), - schema_version: "1.0".to_string(), - }; - let frame = HookFrame { - session_id: SessionId::new("main").unwrap(), - tool_id: None, - call_id: None, - hook_id: None, - event: HookEvent::Custom { - kind: turn_hook::BEFORE_TURN_KIND.to_string(), - payload: serde_json::to_value(&payload).unwrap(), - }, - trace_context: None, - }; - handler - .handle_hook(SessionId::new("main").unwrap(), frame) - .await; - let tracker = handle.activity_tracker(); - assert!( - tracker.known_sessions().contains(&"main".to_string()), - "before_turn hook should create a session entry in the activity tracker" - ); - } - #[tokio::test] - async fn handle_hook_after_turn_does_not_panic() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle.clone()); - handle.activity_tracker().turn_started("main", 1); - let payload = turn_hook::AfterTurnPayload { - turn_number: 1, - outcome: turn_hook::TurnHookOutcome::Completed, - duration_ms: 500, - tool_call_count: 3, - model_id: "grok-3".to_string(), - written_repo_paths: Vec::new(), - cancellation_category: None, - cancellation_context: None, - }; - let frame = HookFrame { - session_id: SessionId::new("main").unwrap(), - tool_id: None, - call_id: None, - hook_id: None, - event: HookEvent::Custom { - kind: turn_hook::AFTER_TURN_KIND.to_string(), - payload: serde_json::to_value(&payload).unwrap(), - }, - trace_context: None, - }; - handler - .handle_hook(SessionId::new("main").unwrap(), frame) - .await; - } - #[tokio::test] - async fn handle_hook_malformed_payload_does_not_panic() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let frame = HookFrame { - session_id: SessionId::new("main").unwrap(), - tool_id: None, - call_id: None, - hook_id: None, - event: HookEvent::Custom { - kind: turn_hook::BEFORE_TURN_KIND.to_string(), - payload: serde_json::json!({ "garbage" : true }), - }, - trace_context: None, - }; - handler - .handle_hook(SessionId::new("main").unwrap(), frame) - .await; - } - #[tokio::test] - async fn handle_hook_unrecognized_custom_kind_does_not_panic() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let frame = HookFrame { - session_id: SessionId::new("main").unwrap(), - tool_id: None, - call_id: None, - hook_id: None, - event: HookEvent::Custom { - kind: "unknown_kind".to_string(), - payload: serde_json::json!({}), - }, - trace_context: None, - }; - handler - .handle_hook(SessionId::new("main").unwrap(), frame) - .await; - } - #[tokio::test] - async fn handle_hook_cancel_marks_call_completed() { - use kigi_tool_protocol::ToolCallId; - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle.clone()); - let tracker = handle.activity_tracker(); - tracker.tool_call_started("call-42", "read_file", Some("main")); - assert_eq!(tracker.snapshot().active_tool_calls, 1); - let frame = HookFrame::cancel( - SessionId::new("main").unwrap(), - ToolId::new("read_file").unwrap(), - ToolCallId::new("call-42").unwrap(), - ); - handler - .handle_hook(SessionId::new("main").unwrap(), frame) - .await; - assert_eq!( - tracker.snapshot().active_tool_calls, - 0, - "cancel hook should mark the call as completed" - ); - } - #[tokio::test] - async fn handle_hook_cancel_without_call_id_cancels_all_session_calls() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle.clone()); - let tracker = handle.activity_tracker(); - tracker.tool_call_started("call-a", "grep", Some("main")); - tracker.tool_call_started("call-b", "read_file", Some("main")); - tracker.tool_call_started("call-c", "write", Some("other")); - assert_eq!(tracker.snapshot().active_tool_calls, 3); - let frame = HookFrame { - session_id: SessionId::new("main").unwrap(), - tool_id: None, - call_id: None, - hook_id: None, - event: HookEvent::Cancel, - trace_context: None, - }; - handler - .handle_hook(SessionId::new("main").unwrap(), frame) - .await; - assert_eq!( - tracker.snapshot_session("main").active_tool_calls, - 0, - "session-wide cancel should complete all calls for the session" - ); - assert_eq!( - tracker.snapshot_session("other").active_tool_calls, - 1, - "cancel must not affect calls in other sessions" - ); - assert_eq!(tracker.snapshot().active_tool_calls, 1); - } - #[tokio::test] - async fn handle_hook_session_ended_clears_turn_active() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle.clone()); - let tracker = handle.activity_tracker(); - tracker.turn_started("main", 1); - assert!(tracker.is_turn_active("main")); - let frame = HookFrame::session_ended(SessionId::new("main").unwrap()); - handler - .handle_hook(SessionId::new("main").unwrap(), frame) - .await; - assert!( - !tracker.is_turn_active("main"), - "session_ended hook should clear turn_active" - ); - } - use crate::workspace_ops::{GetFilesRes, PutFilesRes}; - /// Helper: compute SHA-256 hex digest for test assertions. - fn test_sha256(data: &[u8]) -> String { - use sha2::{Digest, Sha256}; - format!("{:x}", Sha256::digest(data)) - } - #[tokio::test] - async fn dispatch_put_files_writes_and_returns_hash() { - let handle = make_handle(); - let root = handle.root_cwd().unwrap(); - let handler = WorkspaceRpcHandler::new(handle); - let params = serde_json::json!( - { "files" : [{ "path" : "test_file.txt", "content" : "hello world" }] } - ); - let result = handler - .dispatch("workspace.put_files", params, None) - .await - .expect("dispatch should succeed"); - let res: PutFilesRes = serde_json::from_value(result).unwrap(); - assert_eq!(res.results.len(), 1); - assert!(res.results[0].ok, "write should succeed"); - let expected_hash = test_sha256(b"hello world"); - assert_eq!( - res.results[0].hash.as_deref(), - Some(expected_hash.as_str()), - "hash should be SHA-256 of written content" - ); - assert!(res.results[0].error.is_none(), "no error expected"); - let on_disk = std::fs::read_to_string(root.join("test_file.txt")).unwrap(); - assert_eq!(on_disk, "hello world"); - } - #[tokio::test] - async fn dispatch_put_files_rejects_path_traversal() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let params = serde_json::json!( - { "files" : [{ "path" : "../escape.txt", "content" : "evil" }] } - ); - let result = handler - .dispatch("workspace.put_files", params, None) - .await - .expect("dispatch itself should succeed"); - let res: PutFilesRes = serde_json::from_value(result).unwrap(); - assert_eq!(res.results.len(), 1); - assert!(!res.results[0].ok, "path traversal should be rejected"); - assert!( - res.results[0] - .error - .as_ref() - .unwrap() - .contains("escapes workspace root"), - "error should mention escape: {:?}", - res.results[0].error - ); - } - #[tokio::test] - async fn handle_hook_pause_resume_are_noops() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - for event in [HookEvent::Pause, HookEvent::Resume] { - let frame = HookFrame { - session_id: SessionId::new("main").unwrap(), - tool_id: None, - call_id: None, - hook_id: None, - event, - trace_context: None, - }; - handler - .handle_hook(SessionId::new("main").unwrap(), frame) - .await; - } - } - #[tokio::test] - async fn dispatch_put_files_rejects_absolute_outside_root() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let params = serde_json::json!( - { "files" : [{ "path" : "/etc/passwd", "content" : "evil" }] } - ); - let result = handler - .dispatch("workspace.put_files", params, None) - .await - .expect("dispatch itself should succeed"); - let res: PutFilesRes = serde_json::from_value(result).unwrap(); - assert_eq!(res.results.len(), 1); - assert!( - !res.results[0].ok, - "absolute path outside root should be rejected" - ); - assert!( - res.results[0] - .error - .as_ref() - .unwrap() - .contains("escapes workspace root"), - "error should mention escape: {:?}", - res.results[0].error - ); - } - #[tokio::test] - async fn dispatch_put_files_accepts_absolute_within_root() { - let handle = make_handle(); - let root = handle.root_cwd().unwrap(); - let handler = WorkspaceRpcHandler::new(handle); - let abs = root.join("sub/abs.txt"); - let params = serde_json::json!( - { "files" : [{ "path" : abs.to_str().expect("utf-8 path"), "content" : - "hello" }] } - ); - let result = handler - .dispatch("workspace.put_files", params, None) - .await - .expect("dispatch itself should succeed"); - let res: PutFilesRes = serde_json::from_value(result).unwrap(); - assert_eq!(res.results.len(), 1); - assert!( - res.results[0].ok, - "absolute path within root should be accepted: {:?}", - res.results[0].error - ); - assert_eq!( - std::fs::read_to_string(root.join("sub/abs.txt")).unwrap(), - "hello" - ); - } - #[tokio::test] - #[cfg(unix)] - async fn dispatch_put_files_rejects_symlink_escape() { - let handle = make_handle(); - let root = handle.root_cwd().unwrap(); - let handler = WorkspaceRpcHandler::new(handle); - let outside = tempfile::tempdir().expect("create outside dir"); - std::os::unix::fs::symlink(outside.path(), root.join("escape_link")) - .expect("create symlink"); - let params = serde_json::json!( - { "files" : [{ "path" : "escape_link/evil.txt", "content" : "pwned" }] } - ); - let result = handler - .dispatch("workspace.put_files", params, None) - .await - .expect("dispatch itself should succeed"); - let res: PutFilesRes = serde_json::from_value(result).unwrap(); - assert_eq!(res.results.len(), 1); - assert!(!res.results[0].ok, "symlink escape should be rejected"); - assert!( - res.results[0] - .error - .as_ref() - .unwrap() - .contains("symlink escape"), - "error should mention symlink: {:?}", - res.results[0].error - ); - assert!( - !outside.path().join("evil.txt").exists(), - "file must not be created outside workspace" - ); - } - #[tokio::test] - async fn dispatch_put_files_partial_failure() { - let handle = make_handle(); - let root = handle.root_cwd().unwrap(); - let handler = WorkspaceRpcHandler::new(handle); - let params = serde_json::json!( - { "files" : [{ "path" : "good.txt", "content" : "valid content" }, { "path" : - "../bad.txt", "content" : "should fail" },] } - ); - let result = handler - .dispatch("workspace.put_files", params, None) - .await - .expect("dispatch should succeed"); - let res: PutFilesRes = serde_json::from_value(result).unwrap(); - assert_eq!(res.results.len(), 2); - assert!(res.results[0].ok, "first file should succeed"); - assert!(res.results[0].hash.is_some(), "first file should have hash"); - assert!(!res.results[1].ok, "second file should fail"); - assert!( - res.results[1].error.is_some(), - "second file should have error" - ); - assert!( - res.results[1].hash.is_none(), - "failed file should have no hash" - ); - let on_disk = std::fs::read_to_string(root.join("good.txt")).unwrap(); - assert_eq!(on_disk, "valid content"); - } - #[tokio::test] - async fn dispatch_get_files_reads_existing_file() { - let handle = make_handle(); - let root = handle.root_cwd().unwrap(); - let handler = WorkspaceRpcHandler::new(handle); - let content = "read me back"; - std::fs::write(root.join("readable.txt"), content).unwrap(); - let params = serde_json::json!({ "files" : [{ "path" : "readable.txt" }] }); - let result = handler - .dispatch("workspace.get_files", params, None) - .await - .expect("dispatch should succeed"); - let res: GetFilesRes = serde_json::from_value(result).unwrap(); - assert_eq!(res.results.len(), 1); - assert!(res.results[0].exists, "file should exist"); - assert_eq!( - res.results[0].content.as_deref(), - Some(content), - "content should match what was written" - ); - let expected_hash = test_sha256(content.as_bytes()); - assert_eq!( - res.results[0].hash.as_deref(), - Some(expected_hash.as_str()), - "hash should be SHA-256 of file content" - ); - assert!(!res.results[0].matched); - assert_eq!( - res.results[0].size, - Some(content.len() as u64), - "size should match content length" - ); - assert!(res.results[0].error.is_none()); - } - #[tokio::test] - async fn dispatch_get_files_nonexistent_returns_not_exists() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let params = serde_json::json!( - { "files" : [{ "path" : "does_not_exist.txt" }] } - ); - let result = handler - .dispatch("workspace.get_files", params, None) - .await - .expect("dispatch should succeed"); - let res: GetFilesRes = serde_json::from_value(result).unwrap(); - assert_eq!(res.results.len(), 1); - assert!(!res.results[0].exists, "file should not exist"); - assert!(res.results[0].content.is_none()); - assert!(res.results[0].hash.is_none()); - assert!(!res.results[0].matched); - assert!( - res.results[0].error.is_none(), - "missing file is not an error" - ); - } - #[tokio::test] - async fn dispatch_get_files_io_error_returns_exists_true() { - let handle = make_handle(); - let root = handle.root_cwd().unwrap(); - let handler = WorkspaceRpcHandler::new(handle); - std::fs::create_dir_all(root.join("a_directory")).unwrap(); - let params = serde_json::json!({ "files" : [{ "path" : "a_directory" }] }); - let result = handler - .dispatch("workspace.get_files", params, None) - .await - .expect("dispatch should succeed"); - let res: GetFilesRes = serde_json::from_value(result).unwrap(); - assert_eq!(res.results.len(), 1); - assert!(res.results[0].exists, "directory exists on disk"); - assert!( - res.results[0].error.is_some(), - "reading a directory as file should fail: {:?}", - res.results[0] - ); - assert!(res.results[0].content.is_none(), "no content on error"); - } - #[tokio::test] - async fn dispatch_get_files_non_utf8_returns_error_with_hash() { - let handle = make_handle(); - let root = handle.root_cwd().unwrap(); - let handler = WorkspaceRpcHandler::new(handle); - let binary_content: &[u8] = b"\xff\xfe\x00\x01"; - std::fs::write(root.join("binary.bin"), binary_content).unwrap(); - let params = serde_json::json!({ "files" : [{ "path" : "binary.bin" }] }); - let result = handler - .dispatch("workspace.get_files", params, None) - .await - .expect("dispatch should succeed"); - let res: GetFilesRes = serde_json::from_value(result).unwrap(); - assert_eq!(res.results.len(), 1); - assert!(res.results[0].exists, "file should exist"); - assert!( - res.results[0].content.is_none(), - "non-UTF-8 content should be None" - ); - let expected_hash = test_sha256(binary_content); - assert_eq!( - res.results[0].hash.as_deref(), - Some(expected_hash.as_str()), - "hash should be SHA-256 of file content even for non-UTF-8 files" - ); - assert!( - res.results[0] - .error - .as_ref() - .unwrap() - .contains("not valid UTF-8"), - "error should mention UTF-8: {:?}", - res.results[0].error - ); - assert_eq!( - res.results[0].size, - Some(4), - "size should still be reported" - ); - } - #[tokio::test] - async fn dispatch_get_files_cache_hit() { - let handle = make_handle(); - let root = handle.root_cwd().unwrap(); - let handler = WorkspaceRpcHandler::new(handle); - let content = "cacheable content"; - std::fs::write(root.join("cached.txt"), content).unwrap(); - let expected_hash = test_sha256(content.as_bytes()); - let params = serde_json::json!( - { "files" : [{ "path" : "cached.txt", "if_none_match" : expected_hash }] } - ); - let result = handler - .dispatch("workspace.get_files", params, None) - .await - .expect("dispatch should succeed"); - let res: GetFilesRes = serde_json::from_value(result).unwrap(); - assert_eq!(res.results.len(), 1); - assert!(res.results[0].exists); - assert!(res.results[0].matched, "should be a cache hit"); - assert!( - res.results[0].content.is_none(), - "content should be omitted on cache hit" - ); - assert_eq!( - res.results[0].hash.as_deref(), - Some(expected_hash.as_str()), - "hash should still be returned" - ); - assert!(res.results[0].error.is_none()); - } - #[tokio::test] - async fn dispatch_get_files_cache_miss() { - let handle = make_handle(); - let root = handle.root_cwd().unwrap(); - let handler = WorkspaceRpcHandler::new(handle); - let content = "fresh content"; - std::fs::write(root.join("stale.txt"), content).unwrap(); - let params = serde_json::json!( - { "files" : [{ "path" : "stale.txt", "if_none_match" : - "0000000000000000000000000000000000000000000000000000000000000000" }] } - ); - let result = handler - .dispatch("workspace.get_files", params, None) - .await - .expect("dispatch should succeed"); - let res: GetFilesRes = serde_json::from_value(result).unwrap(); - assert_eq!(res.results.len(), 1); - assert!(res.results[0].exists); - assert!(!res.results[0].matched, "should be a cache miss"); - assert_eq!( - res.results[0].content.as_deref(), - Some(content), - "content should be returned on miss" - ); - let expected_hash = test_sha256(content.as_bytes()); - assert_eq!( - res.results[0].hash.as_deref(), - Some(expected_hash.as_str()), - "current hash should be returned" - ); - } - #[tokio::test] - async fn dispatch_put_then_get_round_trip() { - let handle = make_handle(); - let handler = WorkspaceRpcHandler::new(handle); - let content = "round trip content"; - let put_params = serde_json::json!( - { "files" : [{ "path" : "round_trip.txt", "content" : content }] } - ); - let put_result = handler - .dispatch("workspace.put_files", put_params, None) - .await - .expect("put should succeed"); - let put_res: PutFilesRes = serde_json::from_value(put_result).unwrap(); - assert!(put_res.results[0].ok); - let put_hash = put_res.results[0].hash.clone().unwrap(); - let get_params = serde_json::json!( - { "files" : [{ "path" : "round_trip.txt" }] } - ); - let get_result = handler - .dispatch("workspace.get_files", get_params, None) - .await - .expect("get should succeed"); - let get_res: GetFilesRes = serde_json::from_value(get_result).unwrap(); - assert!(get_res.results[0].exists); - assert_eq!( - get_res.results[0].content.as_deref(), - Some(content), - "content should match what was written" - ); - assert_eq!( - get_res.results[0].hash.as_deref(), - Some(put_hash.as_str()), - "get hash should match put hash" - ); - } - #[tokio::test] - async fn dispatch_put_files_append_mode() { - let handle = make_handle(); - let root = handle.root_cwd().unwrap(); - let handler = WorkspaceRpcHandler::new(handle); - let params1 = serde_json::json!( - { "files" : [{ "path" : "chunked.txt", "content" : "hello", "append" : false - }] } - ); - let res1 = handler - .dispatch("workspace.put_files", params1, None) - .await - .expect("first chunk should succeed"); - let put1: PutFilesRes = serde_json::from_value(res1).unwrap(); - assert!(put1.results[0].ok); - let chunk1_hash = put1.results[0].hash.clone().unwrap(); - assert_eq!( - chunk1_hash, - test_sha256(b"hello"), - "hash should be of the appended chunk, not full file" - ); - let params2 = serde_json::json!( - { "files" : [{ "path" : "chunked.txt", "content" : " world", "append" : true - }] } - ); - let res2 = handler - .dispatch("workspace.put_files", params2, None) - .await - .expect("second chunk should succeed"); - let put2: PutFilesRes = serde_json::from_value(res2).unwrap(); - assert!(put2.results[0].ok); - let chunk2_hash = put2.results[0].hash.clone().unwrap(); - assert_eq!( - chunk2_hash, - test_sha256(b" world"), - "hash should be of the appended chunk only" - ); - let on_disk = std::fs::read_to_string(root.join("chunked.txt")).unwrap(); - assert_eq!(on_disk, "hello world"); - } - #[tokio::test] - async fn dispatch_get_files_byte_range() { - let handle = make_handle(); - let root = handle.root_cwd().unwrap(); - let handler = WorkspaceRpcHandler::new(handle); - let content = "0123456789"; - std::fs::write(root.join("range.txt"), content).unwrap(); - let params = serde_json::json!( - { "files" : [{ "path" : "range.txt", "offset" : 3, "length" : 4 }] } - ); - let result = handler - .dispatch("workspace.get_files", params, None) - .await - .expect("dispatch should succeed"); - let res: GetFilesRes = serde_json::from_value(result).unwrap(); - assert_eq!(res.results.len(), 1); - assert!(res.results[0].exists); - assert_eq!( - res.results[0].content.as_deref(), - Some("3456"), - "should return only the requested byte range" - ); - let full_hash = test_sha256(content.as_bytes()); - assert_eq!( - res.results[0].hash.as_deref(), - Some(full_hash.as_str()), - "hash should be of the full file, not the chunk" - ); - assert!(!res.results[0].matched); - assert_eq!( - res.results[0].size, - Some(content.len() as u64), - "size should be full file size" - ); - } - #[tokio::test] - async fn dispatch_get_files_byte_range_cache_hit() { - let handle = make_handle(); - let root = handle.root_cwd().unwrap(); - let handler = WorkspaceRpcHandler::new(handle); - let content = "abcdefghij"; - std::fs::write(root.join("range_cache.txt"), content).unwrap(); - let full_hash = test_sha256(content.as_bytes()); - let params = serde_json::json!( - { "files" : [{ "path" : "range_cache.txt", "offset" : 2, "length" : 3, - "if_none_match" : full_hash, }] } - ); - let result = handler - .dispatch("workspace.get_files", params, None) - .await - .expect("dispatch should succeed"); - let res: GetFilesRes = serde_json::from_value(result).unwrap(); - assert_eq!(res.results.len(), 1); - assert!(res.results[0].exists); - assert!(res.results[0].matched, "should be a cache hit"); - assert!( - res.results[0].content.is_none(), - "content should be omitted on cache hit" - ); - assert_eq!( - res.results[0].hash.as_deref(), - Some(full_hash.as_str()), - "hash should still be returned" - ); - assert_eq!(res.results[0].size, Some(10)); - } - /// Every type with a `WorkspaceRpc` impl must be routed by `dispatch()`. - /// - /// Each entry is compiler-checked via `::METHOD`. - /// Dispatching `{}` may fail with any per-method error (invalid params, - /// session not found, not a git repo) — only an "unknown workspace - /// method" error fails the test. - #[tokio::test] - async fn dispatch_knows_every_typed_method() { - use crate::file_system::{ - ContentSearchRequest, FsDeleteFileReq, FsExistsReq, FsListReq, FsReadFileReq, - FsWriteFileReq, - }; - use crate::workspace_ops::*; - use crate::worktree::{ApplyWorktreeRequest, CreateWorktreeRequest, RemoveWorktreeRequest}; - use kigi_workspace_types::rpc::git::{GitBranchInfoReq, GitMetadataReq}; - use kigi_workspace_types::rpc::search::FuzzyStatusReq; - use kigi_workspace_types::rpc::skills::DiscoverPluginsReq; - use kigi_workspace_types::rpc::workspace::{ - ConfigureMcpReq, DropSessionReq, InstallPluginReq, LoadEnvrcReq, LoadPermissionsReq, - LoadProjectConfigReq, RefreshPluginsReq, ResolveFileReferencesReq, ToolDefinitionsReq, - UpdateToolConfigReq, - }; - use kigi_workspace_types::rpc::worktree::WorktreeCreateSyncReq; - let handler = WorkspaceRpcHandler::new(make_handle()); - let methods = [ - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ::METHOD, - ]; - let skipped_global_db_mutators = [ - ::METHOD, - ::METHOD, - ]; - assert_eq!(skipped_global_db_mutators.len(), 2); - for method in methods { - let result = handler.dispatch(method, serde_json::json!({}), None).await; - if let Err(e) = &result { - assert!( - !e.to_string().contains("unknown workspace method"), - "dispatch does not know {method}: {e}" - ); - } - } - } -} diff --git a/crates/codegen/kigi-workspace/src/mcp.rs b/crates/codegen/kigi-workspace/src/mcp.rs deleted file mode 100644 index 5c15182..0000000 --- a/crates/codegen/kigi-workspace/src/mcp.rs +++ /dev/null @@ -1,220 +0,0 @@ -//! MCP integration for the workspace server. -//! -//! Bridges [`McpClient`] to the server's [`McpTransport`] trait and wraps -//! tool handlers with qualified `server__tool` names. - -use std::sync::Arc; - -use async_trait::async_trait; -use kigi_computer_hub_mcp_adapter::{ - McpBridgeConfig, McpCallResult, McpContent, McpServerInfo, McpToolDefinition, McpToolHandler, - McpTransport, -}; -use kigi_computer_hub_sdk::ToolServerHandler; -use kigi_mcp::rmcp; -use kigi_mcp::servers::McpClient; -use kigi_tool_protocol::ToolId; -use kigi_tool_runtime::{ToolCallContext, ToolStream, TypedToolOutput}; -use kigi_tool_types::ToolDescription; -use serde_json::Value; - -/// Adapts [`McpClient`] to the [`McpTransport`] trait for [`McpBridge`]. -pub(crate) struct McpClientTransportAdapter { - client: Arc, -} - -impl McpClientTransportAdapter { - pub fn new(client: Arc) -> Self { - Self { client } - } -} - -#[async_trait] -impl McpTransport for McpClientTransportAdapter { - async fn initialize(&self) -> Result { - let service = self - .client - .ensure_initialized() - .await - .map_err(|e| kigi_computer_hub_mcp_adapter::McpError::Transport(e.to_string()))?; - let info = service.peer_info().ok_or_else(|| { - kigi_computer_hub_mcp_adapter::McpError::Transport("no peer info after init".into()) - })?; - Ok(McpServerInfo { - name: info.server_info.name.clone(), - version: info.server_info.version.clone(), - capabilities: serde_json::to_value(&info.capabilities).unwrap_or_default(), - }) - } - - async fn list_tools( - &self, - ) -> Result, kigi_computer_hub_mcp_adapter::McpError> { - let service = self - .client - .ensure_initialized() - .await - .map_err(|e| kigi_computer_hub_mcp_adapter::McpError::Transport(e.to_string()))?; - - let mut all_tools = Vec::new(); - let mut cursor: Option = None; - loop { - let result = service - .list_tools(Some( - rmcp::model::PaginatedRequestParams::default().with_cursor(cursor.clone()), - )) - .await - .map_err(|e| kigi_computer_hub_mcp_adapter::McpError::Transport(e.to_string()))?; - - all_tools.extend(result.tools.into_iter().map(|t| McpToolDefinition { - name: t.name.to_string(), - description: t.description.map(|d| d.to_string()), - input_schema: serde_json::to_value(&t.input_schema).ok(), - })); - - match result.next_cursor { - Some(next) => cursor = Some(next), - None => break, - } - } - - Ok(all_tools) - } - - async fn call_tool( - &self, - name: &str, - arguments: Value, - ) -> Result { - let service = self - .client - .ensure_initialized() - .await - .map_err(|e| kigi_computer_hub_mcp_adapter::McpError::Transport(e.to_string()))?; - // MCP spec requires arguments to be an object; coerce if needed. - let args_object = match arguments { - Value::Object(map) => Some(map), - Value::Null => None, - other => { - let mut wrapper = serde_json::Map::new(); - wrapper.insert("value".to_string(), other); - Some(wrapper) - } - }; - let result = service - .call_tool({ - let mut params = rmcp::model::CallToolRequestParams::new(name.to_string()); - params.arguments = args_object; - params - }) - .await - .map_err(|e| kigi_computer_hub_mcp_adapter::McpError::Transport(e.to_string()))?; - - Ok(McpCallResult { - content: result - .content - .into_iter() - .map(|c| match c { - rmcp::model::ContentBlock::Text(t) => McpContent::Text { text: t.text }, - rmcp::model::ContentBlock::Image(img) => McpContent::Image { - mime_type: img.mime_type, - data: img.data, - }, - _ => McpContent::Text { - text: "[unsupported content type]".to_string(), - }, - }) - .collect(), - is_error: result.is_error.unwrap_or(false), - }) - } - - async fn close(&self) -> Result<(), kigi_computer_hub_mcp_adapter::McpError> { - // No-op: cleanup happens when McpClient is dropped. - Ok(()) - } -} - -/// Wraps an [`McpToolHandler`] to qualify tool names as `server__tool`. -pub(crate) struct QualifiedMcpToolHandler { - qualified_id: ToolId, - qualified_name: String, - inner: Arc, -} - -impl QualifiedMcpToolHandler { - /// Returns `None` if the qualified name is not a valid `ToolId`. - pub fn try_new(qualified_name: String, inner: Arc) -> Option { - let qualified_id = match ToolId::new(&qualified_name) { - Ok(id) => id, - Err(err) => { - tracing::warn!( - qualified_name = %qualified_name, - error = %err, - "skipping MCP tool: qualified name is not a valid ToolId" - ); - return None; - } - }; - Some(Self { - qualified_id, - qualified_name, - inner, - }) - } -} - -#[async_trait] -impl ToolServerHandler for QualifiedMcpToolHandler { - fn tool_id(&self) -> ToolId { - self.qualified_id.clone() - } - - fn description(&self) -> ToolDescription { - let inner_desc = self.inner.description(); - ToolDescription::new(self.qualified_name.clone(), inner_desc.description) - } - - fn input_schema(&self) -> Option { - self.inner.input_schema() - } - - async fn handle_call(&self, ctx: ToolCallContext, args: Value) -> ToolStream { - self.inner.handle_call(ctx, args).await - } -} - -/// Result of a `workspace.configure_mcp` RPC call. -#[derive(Debug, Clone, serde::Serialize)] -pub struct McpStartResult { - /// Server names that started successfully. - pub started: Vec, - /// Servers that failed to start. - pub failed: Vec, -} - -/// A single MCP server startup failure. -#[derive(Debug, Clone, serde::Serialize)] -pub struct McpStartFailure { - /// Server name. - pub name: String, - /// Human-readable error description. - pub error: String, -} - -/// Extract a server name from an [`McpError`](kigi_mcp::servers::McpError), -/// falling back to `"unknown"`. -pub(crate) fn server_name_from_mcp_error(e: &kigi_mcp::servers::McpError) -> &str { - e.server_name().unwrap_or("unknown") -} - -/// Bridge config factory for MCP bridge connections. -pub(crate) fn make_bridge_config( - session_id: kigi_tool_protocol::SessionId, - server_name: &str, -) -> McpBridgeConfig { - McpBridgeConfig { - session_id, - namespace: Some(server_name.to_owned()), - } -} diff --git a/crates/codegen/kigi-workspace/src/permission/hub_permission.rs b/crates/codegen/kigi-workspace/src/permission/hub_permission.rs deleted file mode 100644 index fa40aa6..0000000 --- a/crates/codegen/kigi-workspace/src/permission/hub_permission.rs +++ /dev/null @@ -1,498 +0,0 @@ -//! Tool-permission emit: when the rules engine returns "ask" for a guarded -//! tool, request the decision from chat over the server instead of prompting a -//! local ACP client, then map chat's reply back onto a [`PromptOutcome`] so the -//! manager's existing decision + `ALWAYS_*` persistence applies unchanged. -use crate::permission::prompter::{PromptOutcome, tool_name_for_access}; -use crate::permission::types::AccessKind; -use async_trait::async_trait; -use kigi_computer_hub_sdk::harness::PERMISSION_REQUEST_KIND; -use kigi_computer_hub_sdk::{ToolServer, WeakToolServer}; -use kigi_tool_protocol::SessionId; -use prometheus::{HistogramVec, IntCounter, register_histogram_vec, register_int_counter}; -use serde_json::Value; -use std::sync::LazyLock; -/// Wall-clock time the workspace awaits chat's decision on a `permission_request` -/// hook. `outcome` is `ok` (chat replied) or `error` (transport failure / -/// backstop deadline). -static PERMISSION_REPLY_DURATION: LazyLock = LazyLock::new(|| { - register_histogram_vec!( - "grok_workspace_permission_reply_seconds", - "Wall-clock time awaiting chat's reply to a permission_request hook", - &["outcome"], - vec![0.5, 1.0, 2.0, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0] - ) - .expect("grok_workspace_permission_reply_seconds must register once") -}); -/// Permission requests whose reply timed out (the server backstop deadline fired). -/// A subset of the histogram's `error` outcome, promoted to its own counter so a -/// stuck/lost reply is distinguishable from other transport failures. -static PERMISSION_TIMEOUT_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "grok_workspace_permission_timeout_total", - "permission_request hooks whose reply timed out (backstop deadline fired)" - ) - .expect("grok_workspace_permission_timeout_total must register once") -}); -/// Zero-init this module's metric families. See [`crate::init_metrics`]. -pub(crate) fn init_metrics() { - for outcome in ["ok", "error"] { - let _ = PERMISSION_REPLY_DURATION.with_label_values(&[outcome]); - } - PERMISSION_TIMEOUT_TOTAL.inc_by(0); -} -/// Identifies the reply backstop-deadline timeout by its rendered message; the -/// server SDK exposes no typed timeout variant to match on. If that message text -/// changes, such a reply is recorded under the histogram's `error` outcome but -/// not counted in `permission_timeout_total`. -fn is_timeout_err(msg: &str) -> bool { - msg.contains("timed out") -} -/// Env var that enables the HITL-live **tool-permission** emit (workspace → -/// chat over the server) for local e2e and gradual rollout. Prefer server capability -/// negotiation long-term; this is the interim gate so tool-permission can be -/// exercised without waiting on that wire format. -pub const HITL_PERMISSION_LIVE_ENV: &str = "KIGI_HITL_PERMISSION_LIVE"; -/// Whether the HITL-live permission path is enabled. -/// -/// Intended long-term gate: the chat flag `grok_chat_enable_hitl_live_path`, -/// propagated by the server at session-bind (capability negotiation). Until that -/// lands, honor [`HITL_PERMISSION_LIVE_ENV`] (`1` / `true` / `yes`) so local -/// stacks and e2e can turn the emit on explicitly. Default remains **off** -/// (fail closed to the local ACP prompt). -pub fn hitl_permission_live_enabled() -> bool { - match std::env::var(HITL_PERMISSION_LIVE_ENV) { - Ok(v) => { - matches!( - v.trim().to_ascii_lowercase().as_str(), - "1" | "true" | "yes" | "on" - ) - } - Err(_) => false, - } -} -/// Sends a `permission_request` hook to chat and awaits the decision reply. -#[async_trait] -pub trait PermissionHookTransport: Send + Sync { - /// Emit the permission-request `payload` and return chat's decision reply. - async fn request_permission(&self, payload: Value) -> Result; -} -/// Hub-backed permission transport (weak server handle; upgrades per request). -pub struct ToolServerPermissionTransport { - server: WeakToolServer, - session_id: SessionId, -} -impl ToolServerPermissionTransport { - pub fn new(server: ToolServer, session_id: SessionId) -> Self { - Self { - server: server.downgrade(), - session_id, - } - } - /// Build from a session id held as a string; `None` if it is not a valid - /// [`SessionId`]. - pub fn from_session_id(server: ToolServer, session_id: &str) -> Option { - SessionId::new(session_id) - .ok() - .map(|sid| Self::new(server, sid)) - } -} -#[async_trait] -impl PermissionHookTransport for ToolServerPermissionTransport { - async fn request_permission(&self, payload: Value) -> Result { - let start = std::time::Instant::now(); - let Some(server) = self.server.upgrade() else { - PERMISSION_REPLY_DURATION - .with_label_values(&["error"]) - .observe(start.elapsed().as_secs_f64()); - return Err("tool server gone (weak upgrade failed)".to_owned()); - }; - let raw = server - .request_hook( - self.session_id.clone(), - PERMISSION_REQUEST_KIND.to_owned(), - payload, - ) - .await; - let outcome = match &raw { - Ok(_) => "ok", - Err(e) => { - if is_timeout_err(&e.to_string()) { - PERMISSION_TIMEOUT_TOTAL.inc(); - } - "error" - } - }; - PERMISSION_REPLY_DURATION - .with_label_values(&[outcome]) - .observe(start.elapsed().as_secs_f64()); - raw.map_err(|e| e.to_string()) - } -} -fn scope_for_access(access: &AccessKind) -> &'static str { - match access { - AccessKind::Bash(_) | AccessKind::Edit(_) | AccessKind::MCPTool { .. } => "write", - AccessKind::Read(_) - | AccessKind::Grep { .. } - | AccessKind::WebFetch(_) - | AccessKind::WebSearch(_) => "read", - } -} -fn describe_access(access: &AccessKind) -> String { - match access { - AccessKind::Bash(_) => "Run a terminal command".to_owned(), - AccessKind::Edit(path) => format!("Edit {path}"), - AccessKind::MCPTool { name, .. } => format!("Run MCP tool {name}"), - AccessKind::WebFetch(url) => format!("Fetch {url}"), - AccessKind::WebSearch(query) => format!("Search the web for {query}"), - AccessKind::Read(_) => "Read a file".to_owned(), - AccessKind::Grep { .. } => "Search file contents".to_owned(), - } -} -/// Build the server → chat `permission_request` payload. The field set matches -/// chat's `PermissionRequestPayload` parser: `tool_call_id`, `tool_name`, -/// `description`, `scope`, and the bash/edit context. -pub(crate) fn build_permission_payload(access: &AccessKind, tool_call_id: &str) -> Value { - let mut payload = serde_json::json!( - { "tool_call_id" : tool_call_id, "tool_name" : tool_name_for_access(access), - "description" : describe_access(access), "scope" : scope_for_access(access), } - ); - if let Some(map) = payload.as_object_mut() { - match access { - AccessKind::Bash(command) => { - map.insert("bash_command".to_owned(), Value::from(command.clone())); - } - AccessKind::Edit(path) => { - map.insert( - "edit_file_paths".to_owned(), - Value::from(vec![path.clone()]), - ); - } - _ => {} - } - } - payload -} -/// Decode chat's decision reply onto a [`PromptOutcome`]. The reply is chat's -/// `permission_answer_to_json` output: `{ "outcome", "scope"?, "followup_message"? }`. -/// An unknown / `unspecified` outcome fails closed (reject). -pub(crate) fn reply_to_outcome(reply: &Value) -> PromptOutcome { - let outcome = match reply.get("outcome") { - Some(Value::String(s)) => s.as_str(), - Some(Value::Number(n)) => match n.as_i64() { - Some(1) => "approve", - Some(2) => "reject", - Some(3) => "always_approve", - Some(4) => "always_reject", - _ => "", - }, - _ => "", - }; - let followup = reply - .get("followup_message") - .and_then(Value::as_str) - .filter(|s| !s.is_empty()); - match outcome { - "approve" => PromptOutcome::AllowOnce, - "always_approve" => match scope_kind_value(reply) { - Some(("bash_command", Some(value))) => PromptOutcome::AllowAlwaysBashCommand(value), - Some(("server_prefix", Some(value))) => PromptOutcome::AllowAlwaysMcpServer(value), - Some(("domain", Some(value))) => PromptOutcome::AllowAlwaysDomain(value), - _ => PromptOutcome::AllowAlways, - }, - "reject" => match followup { - Some(message) => PromptOutcome::FollowupMessage(message.to_owned()), - None => PromptOutcome::RejectOnce, - }, - "always_reject" => match scope_kind_value(reply) { - Some(("bash_command", Some(value))) => PromptOutcome::RejectAlwaysBashCommand(value), - _ => PromptOutcome::RejectOnce, - }, - "cancelled" => PromptOutcome::Cancelled, - _ => PromptOutcome::RejectOnce, - } -} -fn scope_kind_value(reply: &Value) -> Option<(&str, Option)> { - let scope = reply.get("scope")?; - let kind = scope.get("kind").and_then(Value::as_str)?; - let value = scope - .get("value") - .and_then(Value::as_str) - .map(str::to_owned); - Some((kind, value)) -} -/// Map a hub-served tool name + JSON args onto an [`AccessKind`] for the -/// permission gate in [`crate::hub::SessionRoutedToolHandler`]. Returns `None` -/// for tools that never need a user prompt (reads / todos / dynamic). -pub fn access_kind_for_hub_tool(tool_name: &str, args: &Value) -> Option { - let name = tool_name.rsplit(':').next().unwrap_or(tool_name); - let name = name.strip_prefix("GrokBuild:").unwrap_or(name); - match name { - "run_terminal_command" | "run_terminal_cmd" | "bash" | "shell" => { - let cmd = args - .get("command") - .or_else(|| args.get("full_command")) - .and_then(Value::as_str) - .unwrap_or("") - .to_owned(); - Some(AccessKind::Bash(cmd)) - } - "search_replace" | "hashline_edit" => { - let path = args - .get("file_path") - .or_else(|| args.get("path")) - .and_then(Value::as_str) - .unwrap_or("unknown") - .to_owned(); - Some(AccessKind::Edit(path)) - } - "write" | "write_file" => { - let path = args - .get("file_path") - .or_else(|| args.get("path")) - .and_then(Value::as_str) - .unwrap_or("unknown") - .to_owned(); - Some(AccessKind::Edit(path)) - } - "apply_patch" => Some(AccessKind::Edit("apply_patch".to_owned())), - "web_fetch" => { - let url = args - .get("url") - .and_then(Value::as_str) - .unwrap_or("") - .to_owned(); - Some(AccessKind::WebFetch(url)) - } - n if n.contains("__") || n.starts_with("mcp") => Some(AccessKind::MCPTool { - name: tool_name.to_owned(), - input: args.clone(), - }), - _ => None, - } -} -/// Whether a [`PromptOutcome`] allows the tool call to proceed. -pub fn prompt_outcome_allows(outcome: &PromptOutcome) -> bool { - matches!( - outcome, - PromptOutcome::AllowOnce - | PromptOutcome::AllowAlways - | PromptOutcome::AllowEditsForSession - | PromptOutcome::AllowAlwaysBashCommand(_) - | PromptOutcome::AllowAlwaysDomain(_) - | PromptOutcome::AllowAlwaysMcpTool(_) - | PromptOutcome::AllowAlwaysMcpServer(_) - ) -} -/// Request a permission decision from chat over `transport` and map the reply -/// to a [`PromptOutcome`]. A transport error fails closed (the manager turns an -/// `Error` outcome into a reject) so a lost server connection never silently runs -/// a guarded tool. -pub async fn request_permission_via_hub( - transport: &dyn PermissionHookTransport, - access: &AccessKind, - tool_call_id: &str, -) -> PromptOutcome { - let payload = build_permission_payload(access, tool_call_id); - match transport.request_permission(payload).await { - Ok(reply) => match reply_to_outcome(&reply) { - PromptOutcome::AllowAlways if matches!(access, AccessKind::Edit(_)) => { - PromptOutcome::AllowEditsForSession - } - other => other, - }, - Err(e) => { - tracing::error!(error = % e, "hub permission request failed; rejecting"); - PromptOutcome::Error(format!("hub permission request failed: {e}")) - } - } -} -#[cfg(test)] -mod tests { - use super::*; - use std::sync::Mutex; - /// Pins the current SDK timeout wording the classifier matches on. - #[test] - fn is_timeout_err_matches_backstop_wording_only() { - assert!(is_timeout_err("request timed out after 600s")); - assert!(is_timeout_err("request timed out after 600.0s")); - assert!(!is_timeout_err("connection lost")); - assert!(!is_timeout_err("tool server gone (weak upgrade failed)")); - } - #[test] - fn payload_for_bash_carries_command_and_write_scope() { - let payload = build_permission_payload(&AccessKind::Bash("rm -rf /tmp/x".into()), "tc-1"); - assert_eq!(payload["tool_call_id"], "tc-1"); - assert_eq!(payload["tool_name"], "run_terminal_command"); - assert_eq!(payload["description"], "Run a terminal command"); - assert_eq!(payload["scope"], "write"); - assert_eq!(payload["bash_command"], "rm -rf /tmp/x"); - assert!(payload.get("edit_file_paths").is_none()); - } - #[test] - fn payload_for_edit_carries_file_paths() { - let payload = build_permission_payload(&AccessKind::Edit("src/main.rs".into()), "tc-2"); - assert_eq!(payload["tool_name"], "search_replace"); - assert_eq!(payload["description"], "Edit src/main.rs"); - assert_eq!(payload["scope"], "write"); - assert_eq!( - payload["edit_file_paths"], - serde_json::json!(["src/main.rs"]) - ); - assert!(payload.get("bash_command").is_none()); - assert!(payload.get("edit_kind").is_none()); - } - #[test] - fn payload_for_mcp_has_no_tool_context() { - let payload = build_permission_payload( - &AccessKind::MCPTool { - name: "linear__list".into(), - input: serde_json::Value::Null, - }, - "tc-3", - ); - assert_eq!(payload["tool_name"], "mcp:linear__list"); - assert_eq!(payload["description"], "Run MCP tool linear__list"); - assert_eq!(payload["scope"], "write"); - assert!(payload.get("bash_command").is_none()); - assert!(payload.get("edit_file_paths").is_none()); - } - #[test] - fn reply_outcomes_map_to_prompt_outcomes() { - assert!(matches!( - reply_to_outcome(&serde_json::json!({ "outcome" : "approve" })), - PromptOutcome::AllowOnce - )); - assert!(matches!( - reply_to_outcome(&serde_json::json!({ "outcome" : "reject" })), - PromptOutcome::RejectOnce - )); - assert!(matches!( - reply_to_outcome(&serde_json::json!({ "outcome" : "cancelled" })), - PromptOutcome::Cancelled - )); - assert!(matches!( - reply_to_outcome(&serde_json::json!({ "outcome" : "unspecified" - })), - PromptOutcome::RejectOnce - )); - assert!(matches!( - reply_to_outcome(&serde_json::json!({})), - PromptOutcome::RejectOnce - )); - } - #[test] - fn reject_with_followup_routes_message_to_model() { - let reply = serde_json::json!( - { "outcome" : "reject", "followup_message" : "use cargo instead" } - ); - match reply_to_outcome(&reply) { - PromptOutcome::FollowupMessage(m) => assert_eq!(m, "use cargo instead"), - other => panic!("expected FollowupMessage, got {other:?}"), - } - } - #[test] - fn always_approve_maps_scope_to_persistent_outcome() { - let bash = serde_json::json!( - { "outcome" : "always_approve", "scope" : { "kind" : "bash_command", "value" - : "cargo build" }, } - ); - match reply_to_outcome(&bash) { - PromptOutcome::AllowAlwaysBashCommand(v) => assert_eq!(v, "cargo build"), - other => panic!("expected AllowAlwaysBashCommand, got {other:?}"), - } - let server = serde_json::json!( - { "outcome" : "always_approve", "scope" : { "kind" : "server_prefix", "value" - : "linear" }, } - ); - match reply_to_outcome(&server) { - PromptOutcome::AllowAlwaysMcpServer(v) => assert_eq!(v, "linear"), - other => panic!("expected AllowAlwaysMcpServer, got {other:?}"), - } - assert!(matches!( - reply_to_outcome(&serde_json::json!({ "outcome" : "always_approve" - })), - PromptOutcome::AllowAlways - )); - } - #[test] - fn always_reject_with_bash_scope_persists_the_denied_prefix() { - let reply = serde_json::json!( - { "outcome" : "always_reject", "scope" : { "kind" : "bash_command", "value" : - "curl" }, } - ); - match reply_to_outcome(&reply) { - PromptOutcome::RejectAlwaysBashCommand(v) => assert_eq!(v, "curl"), - other => panic!("expected RejectAlwaysBashCommand, got {other:?}"), - } - } - struct StubTransport { - reply: Result, - seen: Mutex>, - } - #[async_trait] - impl PermissionHookTransport for StubTransport { - async fn request_permission(&self, payload: Value) -> Result { - *self.seen.lock().unwrap() = Some(payload); - self.reply.clone() - } - } - #[tokio::test] - async fn request_sends_payload_and_decodes_reply() { - let transport = StubTransport { - reply: Ok(serde_json::json!({ "outcome" : "approve" })), - seen: Mutex::new(None), - }; - let outcome = - request_permission_via_hub(&transport, &AccessKind::Bash("ls -la".into()), "tc-7") - .await; - assert!(matches!(outcome, PromptOutcome::AllowOnce)); - let seen = transport - .seen - .lock() - .unwrap() - .clone() - .expect("payload sent"); - assert_eq!(seen["tool_call_id"], "tc-7"); - assert_eq!(seen["bash_command"], "ls -la"); - } - #[tokio::test] - async fn transport_error_fails_closed() { - let transport = StubTransport { - reply: Err("connection lost".to_owned()), - seen: Mutex::new(None), - }; - let outcome = - request_permission_via_hub(&transport, &AccessKind::Edit("a.rs".into()), "tc-8").await; - assert!(matches!(outcome, PromptOutcome::Error(_))); - } - #[tokio::test] - async fn edit_always_approve_maps_to_session_scope() { - let transport = StubTransport { - reply: Ok(serde_json::json!({ "outcome" : "always_approve" })), - seen: Mutex::new(None), - }; - let outcome = - request_permission_via_hub(&transport, &AccessKind::Edit("a.rs".into()), "tc-9").await; - assert!(matches!(outcome, PromptOutcome::AllowEditsForSession)); - let transport = StubTransport { - reply: Ok(serde_json::json!({ "outcome" : "always_approve" })), - seen: Mutex::new(None), - }; - let outcome = request_permission_via_hub( - &transport, - &AccessKind::MCPTool { - name: "x".into(), - input: serde_json::Value::Null, - }, - "tc-10", - ) - .await; - assert!(matches!(outcome, PromptOutcome::AllowAlways)); - } - #[test] - fn hitl_permission_live_defaults_off_without_env() { - if std::env::var(HITL_PERMISSION_LIVE_ENV).is_err() { - assert!(!hitl_permission_live_enabled()); - } - } -} diff --git a/crates/codegen/kigi-workspace/src/preview_supervisor.rs b/crates/codegen/kigi-workspace/src/preview_supervisor.rs deleted file mode 100644 index 307eb90..0000000 --- a/crates/codegen/kigi-workspace/src/preview_supervisor.rs +++ /dev/null @@ -1,1326 +0,0 @@ -//! One-child supervisor for the in-sandbox preview-proxy. -//! -//! After the workspace-server self-daemonizes (see [`crate::daemonize`]) it -//! spawns the unchanged `/usr/local/bin/xai-grok-preview-proxy` binary as a -//! child process and supervises exactly that one child: fork/exec → `wait` → -//! restart-on-exit with capped backoff that resets after a healthy run. -//! -//! Two properties depend on *where* this runs: -//! - The child is spawned only from [`supervise_preview`], which the bin invokes -//! **after** daemonize. The child therefore inherits the daemon's new -//! session/pgid and escapes the launcher's process-group reap — one daemonize -//! protects both processes. -//! - `PR_SET_PDEATHSIG(SIGKILL)` binds the child's lifetime to the -//! workspace-server so a WS crash cannot orphan the proxy holding the ports. -//! PDEATHSIG keys off the *spawning thread*, so the fork→exec race (WS dies -//! before the child arms `prctl`) is closed by re-checking `getppid()` in -//! `pre_exec`, and the supervise task is spawned on tokio's long-lived -//! `multi_thread` workers to avoid a spurious worker-thread-death kill. - -use std::fs::{self, File}; -use std::io; -use std::net::Ipv4Addr; -use std::path::{Path, PathBuf}; -use std::sync::{Arc, LazyLock}; -use std::time::{Duration, Instant}; - -use prometheus::{IntCounterVec, register_int_counter_vec}; -use tokio::sync::watch; - -use crate::activity::ActivityTracker; - -/// Absolute path of the preview-proxy binary the supervisor execs. -pub const PREVIEW_PROXY_BIN_PATH: &str = "/usr/local/bin/xai-grok-preview-proxy"; - -/// WS-owned, per-restart-truncated log capturing the proxy's stdout+stderr. A -/// sibling of `WORKSPACE_SERVER_LOG_PATH` on the snapshot-excluded `/var/tmp` -/// overlay (NOT `/tmp`, which is the in-namespace tmpfs rebind), so it persists -/// and is retrievable via the sandbox session log retrieval path. -pub const PREVIEW_PROXY_LOG_PATH: &str = "/var/tmp/workspace-server/tmp/preview-proxy.log"; - -/// A child that ran at least this long is treated as a healthy run, resetting -/// the restart backoff. -pub const PREVIEW_PROXY_HEALTHY_RUN_SECS: u64 = 30; - -/// First/base restart delay; doubled on each consecutive unhealthy restart. -pub const PREVIEW_PROXY_RESTART_BACKOFF_BASE_SECS: u64 = 1; - -/// Ceiling on the restart backoff so a crash-loop pins at this interval rather -/// than growing unbounded. -pub const PREVIEW_PROXY_RESTART_BACKOFF_CAP_SECS: u64 = 30; - -/// `grok_workspace_preview_proxy_restart_total{reason}` — (re)start events the -/// in-sandbox supervisor emits, by reason; tracks preview-proxy restart pressure. -static PREVIEW_PROXY_RESTART_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter_vec!( - "grok_workspace_preview_proxy_restart_total", - "preview-proxy (re)start events emitted by the in-sandbox supervisor, by reason", - &["reason"] - ) - .unwrap() -}); - -/// Reason label for the restart metric. -enum RestartReason { - Exit, - SpawnError, -} - -impl RestartReason { - fn as_str(&self) -> &'static str { - match self { - RestartReason::Exit => "exit", - RestartReason::SpawnError => "spawn_error", - } - } -} - -fn record_restart(reason: RestartReason) { - PREVIEW_PROXY_RESTART_TOTAL - .with_label_values(&[reason.as_str()]) - .inc(); -} - -/// Access policy forwarded to the proxy's `--visibility`. Mirrors the proxy's -/// own enum values (`owner` | `public`) without depending on its crate, and -/// constrains the workspace-server CLI so a bad value fails fast at startup -/// rather than crash-looping the proxy. -#[derive(Clone, Copy, Debug, PartialEq, Eq, clap::ValueEnum)] -pub enum PreviewVisibility { - Owner, - Public, -} - -impl PreviewVisibility { - fn as_str(self) -> &'static str { - match self { - PreviewVisibility::Owner => "owner", - PreviewVisibility::Public => "public", - } - } -} - -/// Supervisor config forwarded to the proxy child. `Option` fields are omitted -/// from argv when absent (the proxy applies its own defaults); per-session -/// secrets stay in the inherited env, never argv. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct PreviewArgs { - /// Gate: the supervisor is started only when this is true. Not forwarded — - /// the proxy has no such flag. - pub enabled: bool, - /// → proxy `--preview-port`. - pub port: Option, - /// → proxy `--control-port`. - pub control_port: Option, - /// → proxy `--visibility` (`owner` | `public`). - pub visibility: Option, - /// → proxy `--instance-suffix`. - pub instance_suffix: Option, - /// → proxy `--auth-redirect` (URL the unauthenticated handshake redirects - /// to). Without it the owner gate denies instead of redirecting. - pub auth_redirect: Option, - /// → proxy `--allow-public` (a bare flag, emitted only when true). - pub allow_public: bool, - /// → proxy `--workspace-server-port`. - pub workspace_server_port: Option, - /// `current_dir` for the spawned child. Not forwarded as an arg. - pub workspace_dir: PathBuf, -} - -impl PreviewArgs { - /// Map the forwarded fields to the proxy's exact CLI flag names (see - /// `xai-grok-preview-proxy/src/cli.rs`). Absent options and a false - /// `allow_public` contribute nothing; the `enabled` gate is never emitted. - pub fn to_argv(&self) -> Vec { - let mut argv = Vec::new(); - if let Some(port) = self.port { - argv.push("--preview-port".to_owned()); - argv.push(port.to_string()); - } - if let Some(port) = self.control_port { - argv.push("--control-port".to_owned()); - argv.push(port.to_string()); - } - if let Some(visibility) = self.visibility { - argv.push("--visibility".to_owned()); - argv.push(visibility.as_str().to_owned()); - } - if let Some(suffix) = &self.instance_suffix { - argv.push("--instance-suffix".to_owned()); - argv.push(suffix.clone()); - } - if let Some(redirect) = &self.auth_redirect { - argv.push("--auth-redirect".to_owned()); - argv.push(redirect.clone()); - } - if self.allow_public { - argv.push("--allow-public".to_owned()); - } - if let Some(port) = self.workspace_server_port { - argv.push("--workspace-server-port".to_owned()); - argv.push(port.to_string()); - } - argv - } -} - -/// Exponential restart backoff with a hard ceiling. The step counter is the -/// number of consecutive unhealthy restarts; a healthy run resets it. -#[derive(Clone, Copy, Debug)] -struct BackoffPolicy { - base: Duration, - cap: Duration, -} - -impl BackoffPolicy { - fn new(base: Duration, cap: Duration) -> Self { - Self { base, cap } - } - - /// `base * 2^step`, saturating to `cap` (overflow ⇒ cap). - fn delay(self, step: u32) -> Duration { - let factor = 2u32.saturating_pow(step); - self.base - .checked_mul(factor) - .unwrap_or(self.cap) - .min(self.cap) - } -} - -/// Delay before the next spawn and the next step counter: a healthy run resets to -/// the base delay; an unhealthy run (or spawn failure) advances the backoff. -fn next_step(policy: BackoffPolicy, healthy: bool, step: u32) -> (Duration, u32) { - if healthy { - (policy.delay(0), 0) - } else { - (policy.delay(step), step.saturating_add(1)) - } -} - -/// Healthy once the run reached `healthy_run` (inclusive); resets the backoff. -fn is_healthy(elapsed: Duration, healthy_run: Duration) -> bool { - elapsed >= healthy_run -} - -/// Open the WS-owned proxy log, truncating it on every (re)start so a crash-loop -/// pinned at the backoff cap cannot grow it unbounded. Reuses the daemon file -/// options (`O_NOFOLLOW` + mode `0600` on Unix) for the same symlink/permission -/// defense as the workspace-server log. -fn open_truncated_log(path: &Path) -> io::Result { - if let Some(parent) = path.parent() { - let _ = fs::create_dir_all(parent); - } - crate::daemonize::daemon_file_options() - .create(true) - .write(true) - .truncate(true) - .open(path) -} - -/// Build the unspawned proxy command. Secrets (`KIGI_SERVER_KEY` / -/// `KIGI_SESSION_ID`) reach the proxy by env inheritance — never argv. -fn build_preview_command(cfg: &PreviewArgs) -> io::Result { - use std::process::Stdio; - - let log = open_truncated_log(Path::new(PREVIEW_PROXY_LOG_PATH))?; - let log_err = log.try_clone()?; - - let mut cmd = std::process::Command::new(PREVIEW_PROXY_BIN_PATH); - cmd.args(cfg.to_argv()) - .current_dir(&cfg.workspace_dir) - .stdin(Stdio::null()) - .stdout(Stdio::from(log)) - .stderr(Stdio::from(log_err)); - - // Linux-only: PDEATHSIG (prctl) does not exist on macOS/other unixes, so this - // parent-death binding is gated to Linux. The proxy simply runs without the - // binding elsewhere. - #[cfg(target_os = "linux")] - { - use std::os::unix::process::CommandExt; - - // Raw pre_exec, NOT kigi_tty_utils::detach_command: the proxy must stay in - // the workspace-server's session/pgid to share its reap-escape, so the - // setsid that detach_command performs would be actively wrong here (and - // the daemonized server owns no controlling TTY, so the detach rationale - // does not apply). - // - // PDEATHSIG keys off the spawning thread, so capture our PID to also - // close the fork→exec race in the child. - let parent_pid = std::process::id(); - // SAFETY: the closure runs in the forked child between fork and exec, so - // it calls only async-signal-safe libc functions (`prctl`, `getppid`, - // `_exit`) and touches no allocation/locks/Rust runtime state. Its error - // path returns `io::Error::last_os_error()`, which only wraps the raw - // errno (no allocation) — never use `io::Error::new`/`other` here, as - // they allocate. - unsafe { - cmd.pre_exec(move || { - // Bind the proxy's lifetime to the workspace-server: a WS crash - // makes the kernel SIGKILL the proxy so it can't orphan and hold - // the preview/control ports. This binding survives the proxy's - // own execve only because that binary is non-setuid and carries - // no file capabilities (the kernel clears PDEATHSIG across a - // privileged exec). - if libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGKILL as libc::c_ulong) == -1 { - return Err(io::Error::last_os_error()); - } - // If the WS already exited (PDEATHSIG won't fire), bail before - // exec rather than orphan. - if libc::getppid() as u32 != parent_pid { - libc::_exit(0); - } - Ok(()) - }); - } - } - - Ok(tokio::process::Command::from(cmd)) -} - -/// Supervise the preview-proxy child until `shutdown` flips. Spawn this **after** -/// daemonize (see module docs) so the child inherits the new session/pgid. -pub async fn supervise_preview(cfg: PreviewArgs, shutdown: watch::Receiver) { - tracing::info!( - bin = PREVIEW_PROXY_BIN_PATH, - log = PREVIEW_PROXY_LOG_PATH, - workspace_dir = %cfg.workspace_dir.display(), - "starting preview-proxy supervisor", - ); - let policy = BackoffPolicy::new( - Duration::from_secs(PREVIEW_PROXY_RESTART_BACKOFF_BASE_SECS), - Duration::from_secs(PREVIEW_PROXY_RESTART_BACKOFF_CAP_SECS), - ); - let healthy_run = Duration::from_secs(PREVIEW_PROXY_HEALTHY_RUN_SECS); - supervise_loop( - move || build_preview_command(&cfg), - policy, - healthy_run, - shutdown, - ) - .await; -} - -/// Core supervise loop, generic over the command factory so tests can drive a -/// fake child without the real proxy binary. -async fn supervise_loop( - mut make_command: F, - policy: BackoffPolicy, - healthy_run: Duration, - mut shutdown: watch::Receiver, -) where - F: FnMut() -> io::Result, -{ - let mut step = 0u32; - while !*shutdown.borrow() { - let started = Instant::now(); - let spawned = make_command().and_then(|mut cmd| cmd.kill_on_drop(true).spawn()); - let mut child = match spawned { - Ok(child) => child, - Err(e) => { - // Spawn/log-open failure must back off and retry — never drop the task. - tracing::error!(error = %e, bin = PREVIEW_PROXY_BIN_PATH, "preview-proxy spawn failed; backing off"); - record_restart(RestartReason::SpawnError); - let (delay, next) = next_step(policy, false, step); - step = next; - if sleep_or_shutdown(delay, &mut shutdown).await { - return; - } - continue; - } - }; - tracing::info!(pid = child.id(), "preview-proxy started"); - - tokio::select! { - status = child.wait() => { - // A shutdown racing the exit must not be counted as a restart. - if *shutdown.borrow() { - return; - } - let ran = started.elapsed(); - let healthy = is_healthy(ran, healthy_run); - record_restart(RestartReason::Exit); - let (delay, next) = next_step(policy, healthy, step); - step = next; - tracing::warn!( - ?status, - healthy, - ran_secs = ran.as_secs(), - backoff_secs = delay.as_secs(), - "preview-proxy exited; restarting", - ); - if sleep_or_shutdown(delay, &mut shutdown).await { - return; - } - } - // SIGKILL on teardown: preview is best-effort and the container is - // going away. - _ = shutdown.changed() => { - let _ = child.kill().await; - tracing::info!("supervisor received shutdown; killed preview-proxy"); - return; - } - } - } -} - -/// Sleep for `delay`, returning early with `true` if `shutdown` flips (or all -/// senders drop) during the wait — the caller then tears down and returns. -async fn sleep_or_shutdown(delay: Duration, shutdown: &mut watch::Receiver) -> bool { - tokio::select! { - _ = tokio::time::sleep(delay) => false, - _ = shutdown.changed() => true, - } -} - -// ── Preview-activity scraper ─────────────────────────────────────────────── -// -// Polls the proxy's loopback `/__control/activity` and feeds the workspace -// `ActivityTracker` so in-sandbox preview traffic withholds idle. - -/// Proxy control path exposing the last-activity stamp (mirrors -/// `xai-grok-preview-proxy`'s `/__control/activity` route). -const PREVIEW_ACTIVITY_PATH: &str = "/__control/activity"; - -/// Effective proxy `--control-port` when the supervisor didn't pin one (mirrors -/// the default in `xai-grok-preview-proxy/src/cli.rs`). -const DEFAULT_PREVIEW_CONTROL_PORT: u16 = 6015; - -/// Per-scrape budget. The endpoint is loopback and trivial, so a small bound is -/// ample and a wedged proxy can't stall the scraper. -const PREVIEW_ACTIVITY_SCRAPE_TIMEOUT: Duration = Duration::from_secs(2); - -/// The proxy's loopback activity URL for `control_port`. Shared by the scrape -/// loop and its tests so a URL-shape change can't drift between them. -fn activity_url(control_port: u16) -> String { - format!( - "http://{}:{control_port}{PREVIEW_ACTIVITY_PATH}", - Ipv4Addr::LOCALHOST - ) -} - -/// Classified result of one scrape, so a missing proxy (quiet no-op) is never -/// confused with a genuine error response or with real activity. -#[derive(Debug, PartialEq, Eq)] -enum ScrapeOutcome { - /// The proxy answered with a parseable activity stamp (epoch-ms). - Stamp(u64), - /// The proxy isn't reachable (connection refused / not up yet): quiet no-op. - Absent, - /// The proxy answered but the response was unusable (error status / bad body). - BadResponse, -} - -/// Parse `{ "last_activity_ms": }`; `None` for a malformed body, a missing -/// field, or a non-integer value. -fn parse_activity_body(body: &str) -> Option { - serde_json::from_str::(body) - .ok()? - .get("last_activity_ms")? - .as_u64() -} - -/// Classify a completed response by status + body (transport failures are -/// classified by the caller). Only a 2xx carrying a parseable stamp is data. -fn classify_activity_response(status: u16, body: &str) -> ScrapeOutcome { - if !(200..300).contains(&status) { - return ScrapeOutcome::BadResponse; - } - match parse_activity_body(body) { - Some(ms) => ScrapeOutcome::Stamp(ms), - None => ScrapeOutcome::BadResponse, - } -} - -/// Whether a scraped stamp is strictly newer than the last seen — i.e. the proxy -/// recorded preview traffic since. A non-increasing value (incl. a proxy -/// restart-to-zero) is not an advance. -fn preview_activity_advanced(last_seen: u64, current: u64) -> bool { - current > last_seen -} - -/// One scrape, classified fail-open: a connect/timeout failure (proxy absent or -/// still starting) is `Absent` (quiet), kept distinct from a genuine error -/// response so neither is ever read as activity. -async fn scrape_activity(client: &reqwest::Client, url: &str) -> ScrapeOutcome { - match client.get(url).send().await { - Ok(resp) => { - let status = resp.status().as_u16(); - match resp.text().await { - Ok(body) => classify_activity_response(status, &body), - Err(_) => ScrapeOutcome::BadResponse, - } - } - Err(e) if e.is_connect() || e.is_timeout() => ScrapeOutcome::Absent, - Err(_) => ScrapeOutcome::BadResponse, - } -} - -/// Poll the proxy's loopback activity endpoint until `shutdown` flips, feeding -/// the tracker on each advance. Spawn after the `ActivityTracker` exists (post -/// hub-connect), gated on preview being enabled. `control_port` is the proxy's -/// loopback control port; `None` falls back to [`DEFAULT_PREVIEW_CONTROL_PORT`]. -/// `scrape_interval` comes from `StatusConfig` (kept strictly below the withhold -/// window by `StatusConfig::validate`). -pub async fn supervise_preview_activity( - control_port: Option, - tracker: Arc, - scrape_interval: Duration, - shutdown: watch::Receiver, -) { - scrape_activity_loop( - control_port.unwrap_or(DEFAULT_PREVIEW_CONTROL_PORT), - tracker, - scrape_interval, - shutdown, - ) - .await; -} - -/// Core scrape loop, parameterized by the scrape interval so tests drive it fast. -async fn scrape_activity_loop( - control_port: u16, - tracker: Arc, - interval: Duration, - mut shutdown: watch::Receiver, -) { - if *shutdown.borrow() { - return; - } - let url = activity_url(control_port); - // A fixed loopback control endpoint never redirects, so a 3xx is anomalous — - // don't follow it; it classifies as `BadResponse`. - let client = match reqwest::Client::builder() - .timeout(PREVIEW_ACTIVITY_SCRAPE_TIMEOUT) - .redirect(reqwest::redirect::Policy::none()) - .build() - { - Ok(client) => client, - Err(e) => { - tracing::warn!(error = %e, "preview-activity scraper: HTTP client build failed; disabled"); - return; - } - }; - tracing::info!(%url, "starting preview-activity scraper"); - - // `None` until the first successful scrape establishes a baseline. Baselining - // (rather than starting at 0) avoids a spurious withhold when a workspace-server - // restart meets a proxy whose stamp is already non-zero but stale. - let mut last_seen: Option = None; - loop { - if sleep_or_shutdown(interval, &mut shutdown).await { - return; - } - match scrape_activity(&client, &url).await { - ScrapeOutcome::Stamp(current) => { - if last_seen.is_some_and(|prev| preview_activity_advanced(prev, current)) { - tracker.note_preview_activity(); - } - last_seen = Some(current); - } - // Proxy absent (preview disabled / starting / restarting): no-op. - ScrapeOutcome::Absent => {} - ScrapeOutcome::BadResponse => { - tracing::debug!(%url, "preview-activity scrape returned an unusable response"); - } - } - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - use std::sync::atomic::{AtomicUsize, Ordering::SeqCst}; - - use super::*; - - fn sample_cfg() -> PreviewArgs { - PreviewArgs { - enabled: true, - port: Some(6014), - control_port: Some(6015), - visibility: Some(PreviewVisibility::Public), - instance_suffix: Some(".inst.example".to_owned()), - auth_redirect: Some("https://grok.com/preview-auth".to_owned()), - allow_public: true, - workspace_server_port: Some(8470), - workspace_dir: PathBuf::from("/workspace"), - } - } - - #[test] - fn to_argv_maps_every_flag_to_the_proxy_cli_names() { - // Flag names must match xai-grok-preview-proxy/src/cli.rs exactly. - assert_eq!( - sample_cfg().to_argv(), - vec![ - "--preview-port", - "6014", - "--control-port", - "6015", - "--visibility", - "public", - "--instance-suffix", - ".inst.example", - "--auth-redirect", - "https://grok.com/preview-auth", - "--allow-public", - "--workspace-server-port", - "8470", - ], - ); - } - - #[test] - fn to_argv_omits_absent_options_and_false_allow_public() { - let cfg = PreviewArgs { - enabled: true, - port: None, - control_port: None, - visibility: None, - instance_suffix: None, - auth_redirect: None, - allow_public: false, - workspace_server_port: None, - workspace_dir: PathBuf::from("/workspace"), - }; - assert!( - cfg.to_argv().is_empty(), - "absent options + false allow_public ⇒ the proxy uses its own defaults" - ); - } - - #[test] - fn to_argv_never_emits_the_enabled_gate() { - // `enabled` gates whether the supervisor runs; it is not a proxy flag and - // must never leak into argv, regardless of its value. - let mut cfg = sample_cfg(); - cfg.enabled = true; - let enabled_argv = cfg.to_argv(); - cfg.enabled = false; - let disabled_argv = cfg.to_argv(); - assert_eq!(enabled_argv, disabled_argv, "the gate must not affect argv"); - assert!(!enabled_argv.iter().any(|a| a.contains("enabled"))); - } - - #[test] - fn to_argv_lowers_owner_visibility() { - // The common (default) case: `Owner` lowers to the proxy's `owner` value. - let mut cfg = sample_cfg(); - cfg.visibility = Some(PreviewVisibility::Owner); - let argv = cfg.to_argv(); - let i = argv - .iter() - .position(|a| a == "--visibility") - .expect("--visibility present"); - assert_eq!(argv[i + 1], "owner"); - } - - #[test] - fn backoff_doubles_caps_at_30s_and_resets_after_a_healthy_run() { - let policy = BackoffPolicy::new(Duration::from_secs(1), Duration::from_secs(30)); - - // Consecutive unhealthy restarts: 1, 2, 4, 8, 16, then pinned at the 30s - // cap (32 → 30, 64 → 30). - let mut step = 0u32; - for want in [1u64, 2, 4, 8, 16, 30, 30] { - let (delay, next) = next_step(policy, false, step); - assert_eq!(delay, Duration::from_secs(want), "step {step}"); - step = next; - } - - // A healthy run resets to the base delay and zeroes the step… - let (delay, next) = next_step(policy, true, step); - assert_eq!(delay, Duration::from_secs(1)); - assert_eq!(next, 0); - - // …and the exponential progression starts over from the base. - let (delay, next) = next_step(policy, false, next); - assert_eq!(delay, Duration::from_secs(1)); - assert_eq!(next, 1); - } - - #[test] - fn backoff_does_not_overflow_at_extreme_steps() { - let policy = BackoffPolicy::new(Duration::from_secs(1), Duration::from_secs(30)); - // Huge step must saturate to the cap, never panic/overflow. - assert_eq!(policy.delay(u32::MAX), Duration::from_secs(30)); - let (delay, next) = next_step(policy, false, u32::MAX); - assert_eq!(delay, Duration::from_secs(30)); - assert_eq!(next, u32::MAX, "step saturates instead of wrapping"); - } - - #[test] - fn is_healthy_boundary_is_inclusive() { - let threshold = Duration::from_secs(PREVIEW_PROXY_HEALTHY_RUN_SECS); - assert!(!is_healthy(threshold - Duration::from_millis(1), threshold)); - assert!(is_healthy(threshold, threshold), "boundary is healthy"); - assert!(is_healthy(threshold + Duration::from_millis(1), threshold)); - } - - #[test] - fn open_truncated_log_truncates_prior_content() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("preview-proxy.log"); - std::fs::write(&path, "stale output from a prior crash-looping run\n").unwrap(); - - let _file = open_truncated_log(&path).unwrap(); - assert_eq!( - std::fs::metadata(&path).unwrap().len(), - 0, - "the log must be truncated to 0 on each (re)start" - ); - } - - #[test] - fn record_restart_increments_the_labeled_counter() { - // The reason → label wiring must be exact (no cross-wiring). - assert_eq!(RestartReason::Exit.as_str(), "exit"); - assert_eq!(RestartReason::SpawnError.as_str(), "spawn_error"); - - // Concurrency-safe: assert a strict increase (other tests may also bump - // these labels), matching the crate's metric-test convention. - let exit_before = PREVIEW_PROXY_RESTART_TOTAL - .with_label_values(&["exit"]) - .get(); - let spawn_before = PREVIEW_PROXY_RESTART_TOTAL - .with_label_values(&["spawn_error"]) - .get(); - - record_restart(RestartReason::Exit); - record_restart(RestartReason::SpawnError); - - assert!( - PREVIEW_PROXY_RESTART_TOTAL - .with_label_values(&["exit"]) - .get() - > exit_before - ); - assert!( - PREVIEW_PROXY_RESTART_TOTAL - .with_label_values(&["spawn_error"]) - .get() - > spawn_before - ); - } - - #[test] - fn default_preview_control_port_is_6015() { - assert_eq!( - DEFAULT_PREVIEW_CONTROL_PORT, 6015, - "guards only this local fallback; keep in step with the proxy's --control-port default" - ); - } - - #[test] - fn parse_activity_body_reads_stamp_and_rejects_bad_shapes() { - assert_eq!( - parse_activity_body(r#"{"last_activity_ms":1234}"#), - Some(1234) - ); - assert_eq!( - parse_activity_body(r#"{"last_activity_ms":0,"extra":true}"#), - Some(0) - ); - assert_eq!(parse_activity_body(r#"{"other":1}"#), None); - assert_eq!(parse_activity_body(r#"{"last_activity_ms":"7"}"#), None); - assert_eq!(parse_activity_body(r#"{"last_activity_ms":1.5}"#), None); - assert_eq!(parse_activity_body("not json"), None); - assert_eq!(parse_activity_body(""), None); - } - - #[test] - fn classify_activity_response_distinguishes_stamp_from_bad() { - assert_eq!( - classify_activity_response(200, r#"{"last_activity_ms":42}"#), - ScrapeOutcome::Stamp(42) - ); - assert_eq!( - classify_activity_response(200, "garbage"), - ScrapeOutcome::BadResponse - ); - assert_eq!( - classify_activity_response(204, ""), - ScrapeOutcome::BadResponse - ); - for status in [301u16, 302, 400, 404, 500, 503] { - assert_eq!( - classify_activity_response(status, r#"{"last_activity_ms":42}"#), - ScrapeOutcome::BadResponse, - "a non-2xx ({status}) must not be read as a stamp" - ); - } - } - - #[test] - fn preview_activity_advanced_detects_only_strictly_newer() { - assert!(!preview_activity_advanced(0, 0)); - assert!(preview_activity_advanced(0, 1)); - assert!(preview_activity_advanced(5, 6)); - assert!(!preview_activity_advanced(5, 5)); - assert!(!preview_activity_advanced(5, 4)); - } - - fn scrape_client() -> reqwest::Client { - reqwest::Client::builder() - .timeout(Duration::from_secs(2)) - .redirect(reqwest::redirect::Policy::none()) - .build() - .expect("build client") - } - - async fn serve_canned(status_line: &'static str, body: &'static str, repeat: bool) -> u16 { - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - let listener = tokio::net::TcpListener::bind((Ipv4Addr::LOCALHOST, 0)) - .await - .expect("bind ephemeral loopback"); - let port = listener.local_addr().expect("addr").port(); - tokio::spawn(async move { - loop { - let Ok((mut sock, _)) = listener.accept().await else { - return; - }; - let mut buf = [0u8; 1024]; - let _ = sock.read(&mut buf).await; - let resp = format!( - "{status_line}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", - body.len() - ); - let _ = sock.write_all(resp.as_bytes()).await; - let _ = sock.flush().await; - if !repeat { - return; - } - } - }); - port - } - - async fn serve_incrementing_stamp() -> u16 { - use std::sync::atomic::{AtomicU64, Ordering}; - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - let listener = tokio::net::TcpListener::bind((Ipv4Addr::LOCALHOST, 0)) - .await - .expect("bind ephemeral loopback"); - let port = listener.local_addr().expect("addr").port(); - tokio::spawn(async move { - let counter = AtomicU64::new(1); - while let Ok((mut sock, _)) = listener.accept().await { - let mut buf = [0u8; 1024]; - let _ = sock.read(&mut buf).await; - let stamp = counter.fetch_add(1, Ordering::Relaxed); - let body = format!(r#"{{"last_activity_ms":{stamp}}}"#); - let resp = format!( - "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", - body.len() - ); - let _ = sock.write_all(resp.as_bytes()).await; - let _ = sock.flush().await; - } - }); - port - } - - async fn serve_accept_then_hang() -> u16 { - let listener = tokio::net::TcpListener::bind((Ipv4Addr::LOCALHOST, 0)) - .await - .expect("bind ephemeral loopback"); - let port = listener.local_addr().expect("addr").port(); - tokio::spawn(async move { - if let Ok((_sock, _)) = listener.accept().await { - tokio::time::sleep(Duration::from_secs(3600)).await; - } - }); - port - } - - #[tokio::test] - async fn scrape_activity_returns_stamp_from_a_live_endpoint() { - let port = serve_canned("HTTP/1.1 200 OK", r#"{"last_activity_ms":9876}"#, false).await; - assert_eq!( - scrape_activity(&scrape_client(), &activity_url(port)).await, - ScrapeOutcome::Stamp(9876) - ); - } - - #[tokio::test] - async fn scrape_activity_classifies_error_status_as_bad_response() { - let port = serve_canned("HTTP/1.1 500 Internal Server Error", "boom", false).await; - assert_eq!( - scrape_activity(&scrape_client(), &activity_url(port)).await, - ScrapeOutcome::BadResponse - ); - } - - #[tokio::test] - async fn scrape_activity_treats_a_closed_port_as_absent_not_error() { - let probe = tokio::net::TcpListener::bind((Ipv4Addr::LOCALHOST, 0)) - .await - .expect("reserve"); - let port = probe.local_addr().expect("addr").port(); - drop(probe); - assert_eq!( - scrape_activity(&scrape_client(), &activity_url(port)).await, - ScrapeOutcome::Absent, - "a refused connection (proxy absent) must be a quiet no-op, not BadResponse" - ); - } - - #[tokio::test] - async fn scrape_activity_treats_a_hung_endpoint_as_absent() { - let port = serve_accept_then_hang().await; - let client = reqwest::Client::builder() - .timeout(Duration::from_millis(150)) - .redirect(reqwest::redirect::Policy::none()) - .build() - .expect("build client"); - assert_eq!( - scrape_activity(&client, &activity_url(port)).await, - ScrapeOutcome::Absent, - "a connection that never responds must time out to Absent, not BadResponse" - ); - } - - #[tokio::test] - async fn scrape_loop_withholds_idle_on_advance_then_stops_on_shutdown() { - let port = serve_incrementing_stamp().await; - let tracker = Arc::new(ActivityTracker::new()); - assert!(tracker.snapshot().idle_since_ms.is_some()); - - let (tx, rx) = watch::channel(false); - let handle = tokio::spawn(scrape_activity_loop( - port, - tracker.clone(), - Duration::from_millis(5), - rx, - )); - - let deadline = Instant::now() + Duration::from_secs(5); - while tracker.snapshot().idle_since_ms.is_some() { - assert!( - Instant::now() < deadline, - "the scrape loop must withhold idle once it observes an advance" - ); - tokio::time::sleep(Duration::from_millis(5)).await; - } - - tx.send(true).expect("receiver alive"); - tokio::time::timeout(Duration::from_secs(5), handle) - .await - .expect("scrape loop must stop promptly on shutdown") - .expect("task should not panic"); - } - - #[tokio::test] - async fn scrape_loop_keeps_idle_reported_on_error_responses() { - let port = serve_canned("HTTP/1.1 500 Internal Server Error", "boom", true).await; - let tracker = Arc::new(ActivityTracker::new()); - let (tx, rx) = watch::channel(false); - let handle = tokio::spawn(scrape_activity_loop( - port, - tracker.clone(), - Duration::from_millis(5), - rx, - )); - - tokio::time::sleep(Duration::from_millis(80)).await; - assert!( - tracker.snapshot().idle_since_ms.is_some(), - "an error response must never be read as activity, so idle stays reported" - ); - - tx.send(true).expect("receiver alive"); - tokio::time::timeout(Duration::from_secs(5), handle) - .await - .expect("scrape loop must stop promptly on shutdown") - .expect("task should not panic"); - } - - #[tokio::test] - async fn scrape_loop_baselines_a_stale_stamp_without_withholding() { - let port = serve_canned("HTTP/1.1 200 OK", r#"{"last_activity_ms":777}"#, true).await; - let tracker = Arc::new(ActivityTracker::new()); - let (tx, rx) = watch::channel(false); - let handle = tokio::spawn(scrape_activity_loop( - port, - tracker.clone(), - Duration::from_millis(5), - rx, - )); - - tokio::time::sleep(Duration::from_millis(80)).await; - assert!( - tracker.snapshot().idle_since_ms.is_some(), - "a constant (stale) stamp must only baseline, never count as an advance" - ); - - tx.send(true).expect("receiver alive"); - tokio::time::timeout(Duration::from_secs(5), handle) - .await - .expect("scrape loop must stop promptly on shutdown") - .expect("task should not panic"); - } - - #[tokio::test] - async fn scrape_loop_returns_immediately_when_already_shut_down() { - let tracker = Arc::new(ActivityTracker::new()); - let (_tx, rx) = watch::channel(true); - tokio::time::timeout( - Duration::from_secs(5), - scrape_activity_loop(6015, tracker, Duration::from_millis(5), rx), - ) - .await - .expect("a pre-flipped shutdown must return without scraping"); - } - - /// A `tokio::process::Command` that exits immediately with `code`. - fn fake_exit_command(code: i32) -> tokio::process::Command { - use std::process::Stdio; - let mut cmd = tokio::process::Command::new("/bin/sh"); - cmd.arg("-c") - .arg(format!("exit {code}")) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()); - cmd - } - - /// A long-lived child that records its own (post-`exec`) PID to `pid_path` - /// then blocks, so the test can assert the process is actually killed. - #[cfg(unix)] - fn fake_pid_recording_command(pid_path: &Path) -> tokio::process::Command { - use std::process::Stdio; - let mut cmd = tokio::process::Command::new("/bin/sh"); - // `$0` is the script's first positional arg; `exec sleep` keeps the PID. - cmd.arg("-c") - .arg(r#"echo $$ > "$0"; exec sleep 3600"#) - .arg(pid_path) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()); - cmd - } - - /// Poll `pid_path` until the child has written a valid PID. - #[cfg(unix)] - async fn read_recorded_pid(pid_path: &Path) -> i32 { - let deadline = Instant::now() + Duration::from_secs(5); - loop { - if let Ok(text) = std::fs::read_to_string(pid_path) - && let Ok(pid) = text.trim().parse::() - && pid > 0 - { - return pid; - } - assert!(Instant::now() < deadline, "child never recorded its pid"); - tokio::time::sleep(Duration::from_millis(5)).await; - } - } - - /// `kill(pid, 0)` existence probe (sends no signal). - #[cfg(unix)] - fn process_alive(pid: i32) -> bool { - // SAFETY: signal 0 only checks for the process's existence/permission. - unsafe { libc::kill(pid, 0) == 0 } - } - - /// Fast policy so backoff sleeps don't slow the real-time integration tests. - fn fast_policy() -> BackoffPolicy { - BackoffPolicy::new(Duration::from_millis(1), Duration::from_millis(2)) - } - - async fn wait_until(counter: &Arc, at_least: usize) { - loop { - if counter.load(SeqCst) >= at_least { - return; - } - tokio::time::sleep(Duration::from_millis(2)).await; - } - } - - #[tokio::test] - async fn supervisor_restarts_on_exit_then_stops_on_shutdown() { - let (tx, rx) = watch::channel(false); - let spawns = Arc::new(AtomicUsize::new(0)); - let factory = { - let spawns = spawns.clone(); - move || { - spawns.fetch_add(1, SeqCst); - Ok(fake_exit_command(1)) - } - }; - // healthy_run far above the child's lifetime ⇒ every exit is unhealthy. - let handle = tokio::spawn(supervise_loop( - factory, - fast_policy(), - Duration::from_secs(3600), - rx, - )); - - // The child is re-spawned after it exits (≥2 spawns ⇒ at least one - // restart). Kept low to minimize fork churn in the parallel test runner; - // the backoff progression itself is covered by the pure `next_step` test. - tokio::time::timeout(Duration::from_secs(5), wait_until(&spawns, 2)) - .await - .expect("supervisor should restart the child after it exits"); - - // Cooperative shutdown: flipping the watch returns the task. - tx.send(true).expect("receiver alive"); - tokio::time::timeout(Duration::from_secs(5), handle) - .await - .expect("supervisor should stop after shutdown") - .expect("task should not panic"); - - assert!(spawns.load(SeqCst) >= 2); - } - - #[tokio::test] - async fn supervisor_returns_cleanly_when_shutdown_races_a_child_exit() { - // Best-effort guard for the post-`wait()` shutdown re-check: with a - // fast-exiting child and shutdown flipped right after a spawn, the exit - // and the shutdown land close together. `select!` is random, so we can't - // deterministically force the re-check branch over the `changed()` arm, - // but both must yield a prompt, panic-free return with no hang or leak. - // Repeated to cover the timing window. - for _ in 0..8 { - let (tx, rx) = watch::channel(false); - let spawns = Arc::new(AtomicUsize::new(0)); - let factory = { - let spawns = spawns.clone(); - move || { - spawns.fetch_add(1, SeqCst); - Ok(fake_exit_command(0)) - } - }; - let handle = tokio::spawn(supervise_loop( - factory, - fast_policy(), - Duration::from_secs(3600), - rx, - )); - wait_until(&spawns, 1).await; - tx.send(true).expect("receiver alive"); - tokio::time::timeout(Duration::from_secs(5), handle) - .await - .expect("supervisor must return on a shutdown/exit race") - .expect("task should not panic"); - } - } - - #[cfg(unix)] - #[tokio::test] - async fn supervisor_shutdown_kills_running_child_without_restart() { - let dir = tempfile::tempdir().unwrap(); - let pid_path = dir.path().join("child.pid"); - let (tx, rx) = watch::channel(false); - let spawns = Arc::new(AtomicUsize::new(0)); - let factory = { - let spawns = spawns.clone(); - let pid_path = pid_path.clone(); - move || { - spawns.fetch_add(1, SeqCst); - Ok(fake_pid_recording_command(&pid_path)) - } - }; - let handle = tokio::spawn(supervise_loop( - factory, - fast_policy(), - Duration::from_secs(1), - rx, - )); - - // The child is up and alive before we ask for shutdown. - let pid = read_recorded_pid(&pid_path).await; - assert!( - process_alive(pid), - "child should be running before shutdown" - ); - - tx.send(true).expect("receiver alive"); - tokio::time::timeout(Duration::from_secs(5), handle) - .await - .expect("supervisor should stop after shutdown") - .expect("task should not panic"); - - assert_eq!( - spawns.load(SeqCst), - 1, - "a live child is killed on shutdown, never restarted" - ); - - // Positively assert the child process is gone (SIGKILL + reap on the - // shutdown path), not merely that the supervise task returned. - let deadline = Instant::now() + Duration::from_secs(5); - while process_alive(pid) { - assert!(Instant::now() < deadline, "child {pid} survived shutdown"); - tokio::time::sleep(Duration::from_millis(10)).await; - } - } - - #[tokio::test] - async fn supervisor_survives_persistent_spawn_failure() { - let (tx, rx) = watch::channel(false); - let attempts = Arc::new(AtomicUsize::new(0)); - let factory = { - let attempts = attempts.clone(); - move || -> io::Result { - attempts.fetch_add(1, SeqCst); - Err(io::Error::new(io::ErrorKind::NotFound, "no such binary")) - } - }; - let handle = tokio::spawn(supervise_loop( - factory, - fast_policy(), - Duration::from_secs(3600), - rx, - )); - - // Repeated spawn failures must back off and retry, never drop the task. - tokio::time::timeout(Duration::from_secs(5), wait_until(&attempts, 5)) - .await - .expect("spawn failures should be retried"); - assert!( - !handle.is_finished(), - "the supervise task must not terminate on spawn failure" - ); - - tx.send(true).expect("receiver alive"); - tokio::time::timeout(Duration::from_secs(5), handle) - .await - .expect("supervisor should still honor shutdown") - .expect("task should not panic"); - } - - /// Helper-process env switch and the success exit code for the PDEATHSIG - /// test below. A distinct (non-zero) success code so a filter that matched - /// no test (libtest would exit 0) can't masquerade as a pass. - #[cfg(target_os = "linux")] - const PDEATHSIG_HELPER_ENV: &str = "KIGI_PDEATHSIG_HELPER"; - #[cfg(target_os = "linux")] - const PDEATHSIG_HELPER_OK: i32 = 42; - - /// End-to-end PDEATHSIG guard: a grandchild armed exactly like the proxy - /// child (`PR_SET_PDEATHSIG(SIGKILL)` + `getppid` race-close) must not - /// outlive its parent. Validates the actual kernel mechanism, not just - /// wiring. (The "parent lives ⇒ child lives" direction is covered by - /// `supervisor_shutdown_kills_running_child_without_restart`, where a child - /// runs until the still-alive parent flips shutdown.) - /// - /// The fork(P)/fork(G) scenario runs in a freshly **re-exec'd** helper - /// process, not the test process: the test binary's descriptors are - /// `O_CLOEXEC`, so the helper's exec closes them all (including other - /// concurrent tests' `flock`'d pidfiles). Forking only that isolated process - /// therefore can't pin a lock another test holds — the hazard of forking the - /// multi-threaded test runner directly. - #[cfg(target_os = "linux")] - #[test] - fn pdeathsig_does_not_let_a_child_outlive_its_parent() { - // Helper mode: run the scenario in this isolated process and exit with - // the verdict. - if std::env::var_os(PDEATHSIG_HELPER_ENV).is_some() { - std::process::exit(run_pdeathsig_scenario()); - } - - // Driver mode: launch the helper and assert its verdict. stdio is - // silenced so the nested libtest banner doesn't pollute this run's output. - let exe = std::env::current_exe().expect("current_exe"); - let status = std::process::Command::new(exe) - .arg("pdeathsig_does_not_let_a_child_outlive_its_parent") // unique substring filter - .arg("--nocapture") - .env(PDEATHSIG_HELPER_ENV, "1") - .stdout(std::process::Stdio::null()) - .stderr(std::process::Stdio::null()) - .status() - .expect("spawn pdeathsig helper process"); - assert_eq!( - status.code(), - Some(PDEATHSIG_HELPER_OK), - "armed grandchild must not outlive its parent (helper verdict {status:?})", - ); - } - - /// fork P; P forks the armed grandchild G then dies; assert G terminates. - /// Returns [`PDEATHSIG_HELPER_OK`] on success, `1` otherwise. Runs only in - /// the isolated helper process (no concurrent tests, no inherited locks). - #[cfg(target_os = "linux")] - fn run_pdeathsig_scenario() -> i32 { - const FAIL: i32 = 1; - // SAFETY: after each fork the child/grandchild call only async-signal-safe - // libc functions (`prctl`, `getppid`, `pause`, `close`, `write`, - // `_exit`); no allocation and no locks. - unsafe { - // Subreaper so G reparents to us when P dies, letting us reap it. - libc::prctl(libc::PR_SET_CHILD_SUBREAPER, 1 as libc::c_ulong); - - let mut fds = [0i32; 2]; - if libc::pipe(fds.as_mut_ptr()) != 0 { - return FAIL; - } - let (rfd, wfd) = (fds[0], fds[1]); - - let intermediate = libc::fork(); - if intermediate < 0 { - return FAIL; - } - if intermediate == 0 { - // Intermediate parent P — stand-in for the workspace-server. - libc::close(rfd); - let p_pid = libc::getpid(); - let grandchild = libc::fork(); - if grandchild == 0 { - // Grandchild G — stand-in for the proxy child. - libc::close(wfd); - libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGKILL as libc::c_ulong); - if libc::getppid() != p_pid { - libc::_exit(0); // race-close: P already gone. - } - loop { - libc::pause(); // wait for the PDEATHSIG SIGKILL. - } - } - // P: report G's pid, then die so PDEATHSIG fires on G. - let _ = libc::write( - wfd, - &grandchild as *const i32 as *const libc::c_void, - std::mem::size_of::(), - ); - libc::close(wfd); - libc::_exit(0); - } - - libc::close(wfd); - let mut grandchild: i32 = -1; - let n = libc::read( - rfd, - &mut grandchild as *mut i32 as *mut libc::c_void, - std::mem::size_of::(), - ); - libc::close(rfd); - if n as usize != std::mem::size_of::() || grandchild <= 0 { - return FAIL; - } - - let mut status = 0i32; - libc::waitpid(intermediate, &mut status, 0); // reap P - - // Once P dies G must terminate (SIGKILL via PDEATHSIG, or exit(0) - // via the race-close). We reap it as the subreaper, or — if it - // reparented to init — `kill(_, 0)` reports `ESRCH`. - let deadline = Instant::now() + Duration::from_secs(5); - loop { - if libc::waitpid(grandchild, &mut status, libc::WNOHANG) == grandchild { - return PDEATHSIG_HELPER_OK; - } - if libc::kill(grandchild, 0) == -1 { - return PDEATHSIG_HELPER_OK; - } - if Instant::now() >= deadline { - libc::kill(grandchild, libc::SIGKILL); - let _ = libc::waitpid(grandchild, &mut status, 0); - return FAIL; - } - std::thread::sleep(Duration::from_millis(10)); - } - } - } -} diff --git a/crates/codegen/kigi-workspace/src/recovery.rs b/crates/codegen/kigi-workspace/src/recovery.rs deleted file mode 100644 index 1bb4370..0000000 --- a/crates/codegen/kigi-workspace/src/recovery.rs +++ /dev/null @@ -1,158 +0,0 @@ -//! Startup janitor for workspace-owned on-disk session state. - -use std::path::Path; -use std::time::Duration; - -/// Default maximum age for a per-session state directory before the -/// [`cleanup_stale_sessions`] janitor reclaims it (7 days). -pub const DEFAULT_SESSION_MAX_AGE: Duration = Duration::from_secs(7 * 24 * 60 * 60); - -/// Remove per-session state directories under `/sessions/` -/// whose mtime is older than `max_age`, bounding the unbounded growth a -/// long-lived workspace (or reused sandbox) would otherwise accumulate. -/// -/// A directory's mtime advances on every atomic-rename persistence write (the -/// rename mutates the directory entry), so it tracks last activity closely -/// enough for a best-effort reclaim. All errors are swallowed — a startup -/// janitor must never fail boot. Only directories with a resolvable, expired -/// mtime are removed; stray files and future-mtime entries are left untouched. -pub async fn cleanup_stale_sessions(workspace_home: &Path, max_age: Duration) { - let sessions_dir = workspace_home.join("sessions"); - let Ok(mut entries) = tokio::fs::read_dir(&sessions_dir).await else { - return; // No sessions dir yet (first boot). - }; - - let mut removed = 0u32; - while let Ok(Some(entry)) = entries.next_entry().await { - let Ok(ft) = entry.file_type().await else { - continue; - }; - if !ft.is_dir() { - continue; - } - let path = entry.path(); - if let Ok(metadata) = tokio::fs::metadata(&path).await - && let Ok(modified) = metadata.modified() - && let Ok(age) = modified.elapsed() - && age > max_age - { - match tokio::fs::remove_dir_all(&path).await { - Ok(()) => { - removed += 1; - tracing::info!( - path = %path.display(), - age_secs = age.as_secs(), - "cleanup_stale_sessions: removed stale session dir" - ); - } - Err(e) => { - tracing::warn!( - path = %path.display(), - error = %e, - "cleanup_stale_sessions: failed to remove stale session dir" - ); - } - } - } - } - - if removed > 0 { - tracing::info!( - removed, - "cleanup_stale_sessions: stale session-dir sweep complete" - ); - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::time::Duration; - - /// A session dir older than `max_age` is removed. - #[tokio::test] - async fn cleanup_removes_stale_session_dir() { - let home = tempfile::TempDir::new().unwrap(); - let stale = home.path().join("sessions").join("sess-old"); - std::fs::create_dir_all(&stale).unwrap(); - std::fs::write(stale.join("tool_state.json"), b"{}").unwrap(); - // Guarantee `age > Duration::ZERO` regardless of mtime resolution. - tokio::time::sleep(Duration::from_millis(15)).await; - - cleanup_stale_sessions(home.path(), Duration::ZERO).await; - - assert!( - !stale.exists(), - "a session dir older than max_age must be removed" - ); - } - - /// A session dir younger than `max_age` is kept. - #[tokio::test] - async fn cleanup_keeps_fresh_session_dir() { - let home = tempfile::TempDir::new().unwrap(); - let fresh = home.path().join("sessions").join("sess-new"); - std::fs::create_dir_all(&fresh).unwrap(); - std::fs::write(fresh.join("tool_state.json"), b"{}").unwrap(); - - cleanup_stale_sessions(home.path(), Duration::from_secs(3600)).await; - - assert!( - fresh.exists(), - "a session dir younger than max_age must be kept" - ); - assert!( - fresh.join("tool_state.json").exists(), - "a kept session dir must retain its contents" - ); - } - - /// No `sessions/` directory (first boot) is a silent no-op, not a panic. - #[tokio::test] - async fn cleanup_missing_sessions_dir_is_noop() { - let home = tempfile::TempDir::new().unwrap(); - cleanup_stale_sessions(home.path(), Duration::ZERO).await; - assert!(home.path().exists(), "home is left untouched"); - } - - /// Stray files under `sessions/` are never removed — directories only. - #[tokio::test] - async fn cleanup_ignores_non_dir_entries() { - let home = tempfile::TempDir::new().unwrap(); - let sessions = home.path().join("sessions"); - std::fs::create_dir_all(&sessions).unwrap(); - let stray = sessions.join("stray.txt"); - std::fs::write(&stray, b"not a session dir").unwrap(); - tokio::time::sleep(Duration::from_millis(15)).await; - - cleanup_stale_sessions(home.path(), Duration::ZERO).await; - - assert!( - stray.exists(), - "stray files under sessions/ must never be removed" - ); - } - - /// Mixed sweep: the `max_age` comparison is per-entry, not all-or-nothing. - #[tokio::test] - async fn cleanup_removes_only_expired_dirs() { - let home = tempfile::TempDir::new().unwrap(); - let sessions = home.path().join("sessions"); - let old = sessions.join("sess-old"); - std::fs::create_dir_all(&old).unwrap(); - // Wide gap (100ms) vs a 50ms threshold so neither side is sensitive to - // scheduler jitter. - tokio::time::sleep(Duration::from_millis(100)).await; - let max_age = Duration::from_millis(50); - let fresh = sessions.join("sess-fresh"); - std::fs::create_dir_all(&fresh).unwrap(); - - cleanup_stale_sessions(home.path(), max_age).await; - - assert!(!old.exists(), "the >max_age dir must be removed"); - assert!( - fresh.exists(), - "the wire-code mapping for the workspace RPC envelope -//! (the envelope types are canonical in `kigi_workspace_types::rpc`). -//! -//! `error_code` uses a non-wildcard match so the compiler enforces -//! coverage of new `WorkspaceError` variants. - -pub use kigi_workspace_types::rpc::{RpcEnvelope, RpcError}; - -use crate::error::WorkspaceError; - -/// Build an error envelope from a `WorkspaceError`. -pub fn envelope_err(error: &WorkspaceError) -> RpcEnvelope { - RpcEnvelope::err_parts(error_code(error), error.to_string()) -} - -/// Map a `WorkspaceError` to its wire code string. -/// -/// Uses an exhaustive match with no wildcard -- the compiler will -/// error when new variants are added to `WorkspaceError`, forcing -/// the implementer to assign a wire code. -pub fn error_code(err: &WorkspaceError) -> &'static str { - match err { - WorkspaceError::ParentSessionNotFound(_) => "parent_session_not_found", - WorkspaceError::SessionNotFound(_) => "session_not_found", - WorkspaceError::SessionAlreadyExists(_) => "session_already_exists", - WorkspaceError::EmptyAgentId => "empty_agent_id", - WorkspaceError::CannotDropMainSession => "cannot_drop_main", - WorkspaceError::Finalize(_) => "finalize", - WorkspaceError::CapabilityWidening { .. } => "capability_widening", - WorkspaceError::Unauthorized { .. } => "unauthorized", - WorkspaceError::TurnActive(_) => kigi_workspace_types::rpc::envelope::TURN_ACTIVE, - WorkspaceError::MaxDepthExceeded { .. } => "max_depth_exceeded", - WorkspaceError::JoinError(_) => "join_error", - WorkspaceError::InvalidHunkAction(_) => "invalid_hunk_action", - WorkspaceError::HunkActionFailed(_) => "hunk_action_failed", - WorkspaceError::HubError(_) => "hub_error", - WorkspaceError::DeployError { kind, .. } => kind.wire_code(), - WorkspaceError::ShuttingDown => "shutting_down", - WorkspaceError::ToolsetExternallyOwned(_) => "toolset_externally_owned", - } -} - -/// Map a wire [`RpcError`] back to a [`WorkspaceError`]. -/// -/// Known codes are mapped to their specific variants. Unknown codes -/// degrade gracefully to `WorkspaceError::HubError`, ensuring -/// forward compatibility when a newer workspace sends codes an older -/// shell does not recognise. -/// -/// # Intentional degradation -/// -/// The structured variants `CapabilityWidening`, `Unauthorized`, and -/// `MaxDepthExceeded` carry multiple fields that are not preserved in -/// the wire `message` string. These are mapped to `HubError` on the -/// deserializing side because reconstructing the original struct fields -/// from a flattened `Display` string would be fragile and error-prone. -/// Callers that need to distinguish these errors can match on the -/// `HubError` message which contains the original error code as a -/// prefix (e.g. `"capability_widening: ..."`). -pub fn rpc_error_to_workspace(err: RpcError) -> WorkspaceError { - if let Some(kind) = kigi_workspace_types::rpc::deploy::DeployError::from_wire_code(&err.code) { - return WorkspaceError::DeployError { - kind, - message: err.message, - }; - } - match err.code.as_str() { - "parent_session_not_found" => WorkspaceError::ParentSessionNotFound(err.message), - "session_not_found" => WorkspaceError::SessionNotFound(err.message), - "session_already_exists" => WorkspaceError::SessionAlreadyExists(err.message), - "empty_agent_id" => WorkspaceError::EmptyAgentId, - "cannot_drop_main" => WorkspaceError::CannotDropMainSession, - "finalize" => WorkspaceError::Finalize(err.message), - "capability_widening" => { - WorkspaceError::HubError(format!("capability_widening: {}", err.message)) - } - "unauthorized" => WorkspaceError::HubError(format!("unauthorized: {}", err.message)), - "turn_active" => WorkspaceError::TurnActive(err.message), - "max_depth_exceeded" => { - WorkspaceError::HubError(format!("max_depth_exceeded: {}", err.message)) - } - "join_error" => WorkspaceError::JoinError(err.message), - "invalid_hunk_action" => WorkspaceError::InvalidHunkAction(err.message), - "hunk_action_failed" => WorkspaceError::HunkActionFailed(err.message), - "hub_error" => WorkspaceError::HubError(err.message), - "shutting_down" => WorkspaceError::ShuttingDown, - "toolset_externally_owned" => WorkspaceError::ToolsetExternallyOwned(err.message), - unknown => { - WorkspaceError::HubError(format!("unknown error code: {unknown}: {}", err.message)) - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::capability::CapabilityMode; - - /// Verify round-trip fidelity for every `WorkspaceError` variant. - #[test] - fn error_code_round_trip_all_variants() { - let mut variants: Vec = vec![ - WorkspaceError::ParentSessionNotFound("p".into()), - WorkspaceError::SessionNotFound("s".into()), - WorkspaceError::SessionAlreadyExists("s".into()), - WorkspaceError::EmptyAgentId, - WorkspaceError::CannotDropMainSession, - WorkspaceError::Finalize("f".into()), - WorkspaceError::CapabilityWidening { - parent: CapabilityMode::ReadOnly, - child: CapabilityMode::All, - }, - WorkspaceError::Unauthorized { - caller: "a".into(), - target: "b".into(), - }, - WorkspaceError::TurnActive("s".into()), - WorkspaceError::MaxDepthExceeded { parent: "p".into() }, - WorkspaceError::JoinError("j".into()), - WorkspaceError::InvalidHunkAction("h".into()), - WorkspaceError::HunkActionFailed("h".into()), - WorkspaceError::HubError("hub".into()), - WorkspaceError::ShuttingDown, - WorkspaceError::ToolsetExternallyOwned("s".into()), - ]; - variants.extend( - kigi_workspace_types::rpc::deploy::DeployError::ALL - .into_iter() - .map(|kind| WorkspaceError::DeployError { - kind, - message: "deploy".into(), - }), - ); - - for err in &variants { - let code = error_code(err); - assert!(!code.is_empty(), "code must not be empty for {err:?}"); - - // Round-trip through RpcError - let rpc_err = RpcError { - code: code.to_owned(), - message: err.to_string(), - }; - let recovered = rpc_error_to_workspace(rpc_err); - // The recovered error's code should match the original code - let recovered_code = error_code(&recovered); - - // Structured variants (CapabilityWidening, Unauthorized, - // MaxDepthExceeded) lose their fields on the wire and - // degrade to HubError, which is the expected behavior. - // Their error messages are preserved in the HubError string. - match err { - WorkspaceError::CapabilityWidening { .. } => { - assert_eq!(recovered_code, "hub_error"); - let msg = recovered.to_string(); - assert!( - msg.contains("capability_widening"), - "degraded error should contain original code: {msg}" - ); - } - WorkspaceError::Unauthorized { .. } => { - assert_eq!(recovered_code, "hub_error"); - let msg = recovered.to_string(); - assert!( - msg.contains("unauthorized"), - "degraded error should contain original code: {msg}" - ); - } - WorkspaceError::MaxDepthExceeded { .. } => { - assert_eq!(recovered_code, "hub_error"); - let msg = recovered.to_string(); - assert!( - msg.contains("max_depth_exceeded"), - "degraded error should contain original code: {msg}" - ); - } - _ => { - assert_eq!( - recovered_code, code, - "round-trip mismatch for {err:?}: got code {recovered_code}" - ); - } - } - } - } - - /// Verify unknown codes degrade to HubError. - #[test] - fn unknown_code_degrades_to_hub_error() { - let rpc_err = RpcError { - code: "future_new_variant".into(), - message: "something new".into(), - }; - let recovered = rpc_error_to_workspace(rpc_err); - assert!(matches!(recovered, WorkspaceError::HubError(_))); - let msg = recovered.to_string(); - assert!(msg.contains("future_new_variant")); - } - - /// Verify serde round-trip of RpcEnvelope. - #[test] - fn envelope_serde_round_trip_ok() { - let env: RpcEnvelope = RpcEnvelope::ok("hello".into()); - let json = serde_json::to_value(&env).unwrap(); - let recovered: RpcEnvelope = serde_json::from_value(json).unwrap(); - match recovered.into_result() { - Ok(v) => assert_eq!(v, "hello"), - Err(e) => panic!("expected Ok, got {e:?}"), - } - } - - /// Verify serde round-trip of RpcEnvelope error, through the - /// `WorkspaceError` mapping in both directions. - #[test] - fn envelope_serde_round_trip_err() { - let err = WorkspaceError::SessionNotFound("ghost".into()); - let env: RpcEnvelope = envelope_err(&err); - let json = serde_json::to_value(&env).unwrap(); - let recovered: RpcEnvelope = serde_json::from_value(json).unwrap(); - match recovered.into_result().map_err(rpc_error_to_workspace) { - Ok(_) => panic!("expected Err"), - Err(e) => { - assert_eq!(error_code(&e), "session_not_found"); - } - } - } -} diff --git a/crates/common/kigi-computer-hub-core/Cargo.toml b/crates/common/kigi-computer-hub-core/Cargo.toml deleted file mode 100644 index 55ddb95..0000000 --- a/crates/common/kigi-computer-hub-core/Cargo.toml +++ /dev/null @@ -1,25 +0,0 @@ -[package] -license = "Apache-2.0" -name = "kigi-computer-hub-core" -version.workspace = true -edition.workspace = true -description = "Transport, ToolRegistry, and resolver abstractions for the xAI Computer Hub" - -[dependencies] -async-trait = { workspace = true } -chrono = { workspace = true } -futures = { workspace = true } -serde_json = { workspace = true } -tracing = { workspace = true } -kigi-tool-protocol = { workspace = true } -kigi-tool-runtime = { workspace = true } -kigi-tool-types = { workspace = true } - -[dev-dependencies] -dashmap = { workspace = true } -schemars = { workspace = true } -tokio = { workspace = true, features = ["rt", "rt-multi-thread", "macros", "test-util", "sync"] } -serde = { workspace = true, features = ["derive"] } - -[lints] -workspace = true diff --git a/crates/common/kigi-computer-hub-core/src/inner.rs b/crates/common/kigi-computer-hub-core/src/inner.rs deleted file mode 100644 index 4985f2e..0000000 --- a/crates/common/kigi-computer-hub-core/src/inner.rs +++ /dev/null @@ -1,73 +0,0 @@ -//! `InnerDispatchForResolver` — an object-safe `ToolDispatch` that routes -//! through a `Weak` bound to a single session. -//! -//! Tools that need to call other tools (the inner-dispatch pattern) ask -//! the runtime for an `Arc`. This adapter answers that -//! question with a resolver-backed implementation. Holding the resolver -//! by [`Weak`] lets the router own the resolver while inner-dispatch -//! handles created from the same resolver release naturally when the -//! router is torn down. - -use std::sync::Weak; - -use async_trait::async_trait; -use serde_json::Value; - -use kigi_tool_protocol::{SessionId, ToolId}; -use kigi_tool_runtime::{ - ToolCallContext, ToolDispatch, ToolError, ToolStream, TypedToolOutput, terminal_only, -}; - -use crate::resolver::CompoundResolver; - -/// Resolver-backed `ToolDispatch` implementation. -/// -/// The resolver is held by [`Weak`] so the inner-dispatch handle never -/// keeps the router alive past its natural lifetime — when the owning -/// router drops the resolver, in-flight inner calls fail cleanly with -/// [`ToolError::Custom`] keyed `computer_hub_dropped`. -/// -/// Bound to a single [`SessionId`] at construction (rather than reading a -/// session from [`ToolCallContext`]) so the inner-dispatch path mirrors -/// the per-session lifetime of the outer router. -#[derive(Debug)] -pub struct InnerDispatchForResolver { - resolver: Weak, - session_id: SessionId, -} - -impl InnerDispatchForResolver { - /// Build an inner-dispatch handle bound to `session_id`, resolving - /// through `resolver`. - pub fn new(resolver: Weak, session_id: SessionId) -> Self { - Self { - resolver, - session_id, - } - } - - /// Borrow the bound session identifier. - pub fn session_id(&self) -> &SessionId { - &self.session_id - } -} - -#[async_trait] -impl ToolDispatch for InnerDispatchForResolver { - async fn call( - &self, - tool_id: ToolId, - args: Value, - ctx: ToolCallContext, - ) -> ToolStream { - let Some(resolver) = self.resolver.upgrade() else { - return terminal_only(Err(ToolError::custom( - "computer_hub_dropped", - "computer hub dropped before inner call could execute", - ))); - }; - resolver - .resolve_and_dispatch(&self.session_id, tool_id, args, ctx) - .await - } -} diff --git a/crates/common/kigi-computer-hub-core/src/lib.rs b/crates/common/kigi-computer-hub-core/src/lib.rs deleted file mode 100644 index fab24fd..0000000 --- a/crates/common/kigi-computer-hub-core/src/lib.rs +++ /dev/null @@ -1,29 +0,0 @@ -//! xAI Computer Hub — transport + registry + resolver core. -//! -//! Object-safe abstractions used by every router build: a [`Transport`] -//! that authorises and dispatches calls, a [`ToolRegistry`] trait shared -//! by both storage planes, a [`CompoundResolver`] that applies the -//! local-shadows-remote rule, and the local + remote transports plus -//! inner-dispatch glue that sit on top. - -#![forbid(unsafe_code)] - -pub mod inner; -pub mod local; -pub mod registry; -pub mod remote; -pub mod resolver; -pub mod transport; - -pub use inner::InnerDispatchForResolver; -pub use local::{LOCAL_INVOKE_SCOPE, LocalTransport}; -pub use registry::{ - ConnectionCleanupReport, SessionCleanupReport, ToolRegistry, ToolSessionBindOutcome, - ToolSessionUnbindOutcome, -}; -pub use remote::{ - ConnectionClient, RemoteToolProxy, RemoteTransport, decode_call_result, error_from_envelope, - is_workspace_unavailable, output_to_value, progress_from_frame, tool_error_from_wire, -}; -pub use resolver::{CompoundResolver, ErasedTool, ResolvedTool, ToolHandle}; -pub use transport::{Principal, Transport, TransportKind}; diff --git a/crates/common/kigi-computer-hub-core/src/local.rs b/crates/common/kigi-computer-hub-core/src/local.rs deleted file mode 100644 index 52e12b8..0000000 --- a/crates/common/kigi-computer-hub-core/src/local.rs +++ /dev/null @@ -1,77 +0,0 @@ -//! In-process transport that resolves through a [`CompoundResolver`]. -//! -//! `LocalTransport` is bound to a single `(user_id, session_id)` at -//! construction. Authorisation returns a principal pre-populated with the -//! bound session and the `tool.invoke` scope; per-call dispatch resolves -//! against the bound session's view of the resolver. - -use std::sync::Arc; - -use async_trait::async_trait; -use serde_json::Value; - -use kigi_tool_protocol::{SessionId, ToolId, UserId}; -use kigi_tool_runtime::{ToolCallContext, ToolError, ToolStream, TypedToolOutput}; - -use crate::resolver::CompoundResolver; -use crate::transport::{Principal, Transport, TransportKind}; - -/// The scope `LocalTransport::authorize` grants to its principal. -/// -/// Hoisted so adapters that authorise principals through other paths -/// can match the local convention without restating the literal. -pub const LOCAL_INVOKE_SCOPE: &str = "tool.invoke"; - -/// Transport that dispatches against an in-process resolver. -#[derive(Debug)] -pub struct LocalTransport { - resolver: Arc, - user_id: UserId, - session_id: SessionId, -} - -impl LocalTransport { - /// Build a transport bound to `(user_id, session_id)` and resolving - /// through `resolver`. - pub fn new(resolver: Arc, user_id: UserId, session_id: SessionId) -> Self { - Self { - resolver, - user_id, - session_id, - } - } - - /// Bound user identity for this transport. - pub fn user_id(&self) -> &UserId { - &self.user_id - } - - /// Bound session for this transport. - pub fn session_id(&self) -> &SessionId { - &self.session_id - } -} - -#[async_trait] -impl Transport for LocalTransport { - fn kind(&self) -> TransportKind { - TransportKind::Local - } - - async fn authorize(&self) -> Result { - Ok(Principal::new(self.user_id.clone()) - .with_session(self.session_id.clone()) - .with_scope(LOCAL_INVOKE_SCOPE)) - } - - async fn call( - &self, - tool_id: ToolId, - args: Value, - ctx: ToolCallContext, - ) -> ToolStream { - self.resolver - .resolve_and_dispatch(&self.session_id, tool_id, args, ctx) - .await - } -} diff --git a/crates/common/kigi-computer-hub-core/src/registry.rs b/crates/common/kigi-computer-hub-core/src/registry.rs deleted file mode 100644 index 0f368f2..0000000 --- a/crates/common/kigi-computer-hub-core/src/registry.rs +++ /dev/null @@ -1,305 +0,0 @@ -//! Object-safe `ToolRegistry` trait shared by every storage plane. -//! -//! Two registry implementations are expected: one in-memory plane for -//! statically-registered local tools, and one connection-keyed plane fed by -//! incoming remote registrations. Both expose the same trait so the -//! router can compose them through [`crate::CompoundResolver`] without -//! caring which is which. -//! -//! Mutations are connection-scoped: each registered tool belongs to the -//! [`ConnectionId`] that introduced it. Per-tool session bindings live -//! alongside the tool's record and are mutated independently via -//! [`ToolRegistry::bind_tool_session`] / [`ToolRegistry::unbind_tool_session`]. -//! Reads (`find_tool`, `list_tools`, `search`) remain session-scoped — the -//! router resolves a tool by `(session_id, tool_id)`, never by -//! connection id. -//! -//! The concrete in-memory implementation is intentionally **out of scope** -//! for this crate — it requires a concurrency story (sharded maps, an -//! actor, etc.) that belongs alongside the registry's collision matrix and -//! generation handling. Tests exercise the trait via per-test mock impls. - -use std::collections::HashSet; -use std::sync::atomic::{AtomicU64, Ordering}; - -use async_trait::async_trait; - -use kigi_tool_protocol::{ - ConnectionId, RegistrationOutcome, ServerId, SessionId, ToolDefinitionMode, ToolId, - ToolRegistration, ToolServerRegistration, UserId, -}; -use kigi_tool_runtime::{SearchSnapshot, ServerSummary}; -use kigi_tool_types::ToolDescription; - -use crate::resolver::ResolvedTool; - -/// Outcome of a single [`ToolRegistry::bind_tool_session`] call. -/// -/// This enum is the source of truth for storage outcomes; the wire enum -/// [`kigi_tool_protocol::ToolSessionBindOutcome`] is a strict subset with -/// one extra wire-only variant. The two layers diverge deliberately: -/// -/// - `Conflict` (cross-connection race on the `(session_id, tool_id)` -/// reverse-index slot) is registry-internal: the router lifts it to a -/// top-level `ServerError::ToolBindingConflict` (-32600) instead of -/// mirroring it to the wire ack, so the contended caller gets a -/// dedicated error code rather than overloading `UnknownTool`. -/// - The wire enum's `SessionNotBound` is router-injected by the -/// per-frame envelope pre-check (the connection's bound-session set -/// lives in router state, not the registry) and is never produced by -/// any registry call — so it has no counterpart here. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ToolSessionBindOutcome { - /// Added to the tool's session set. - Bound, - /// Session id was already in the tool's session set; no-op. - AlreadyBound, - /// No tool with the given id is registered against this connection. - UnknownTool, - /// Cross-connection conflict: another connection already holds the - /// `(session_id, tool_id)` reverse-index slot. The router lifts this - /// into a top-level `ToolBindingConflict` server error so the wire - /// reply uses the dedicated -32600 code instead of the structurally - /// dishonest `UnknownTool`. No registry state was mutated. - Conflict, -} - -/// Outcome of a single [`ToolRegistry::unbind_tool_session`] call. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ToolSessionUnbindOutcome { - /// Removed from the tool's session set. - Unbound, - /// Session id was not in the tool's session set; no-op. - NotBound, - /// No tool with the given id is registered against this connection. - UnknownTool, -} - -/// Aggregated summary of a connection-scoped cleanup pass. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct ConnectionCleanupReport { - /// Number of distinct `(connection, tool_id)` records dropped. - pub tools_dropped: usize, - /// Number of reverse-index `(session_id, tool_id)` rows cleaned up - /// across every session the dropped tools were bound to. - pub session_bindings_cleared: usize, -} - -/// Aggregated summary of a session-scoped cleanup pass. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct SessionCleanupReport { - /// Number of tools whose session set lost the unregistered session id. - pub tools_touched: usize, - /// Number of tools whose session set became empty after the - /// unregistration. The tool record itself is NOT removed — the owning - /// connection still owns it and may rebind via - /// [`ToolRegistry::bind_tool_session`] later. - pub tools_left_orphaned: usize, -} - -/// Backend-agnostic registry of tools available within a router. -/// -/// Methods are split into mutating (`async fn` — registration changes may -/// touch shared state and require coordination) and read-only views -/// (synchronous — implementations should answer from a consistent snapshot -/// without awaiting). The split mirrors how callers use the registry: the -/// hot path is the `find_tool` / `list_tools` view; mutations happen on the -/// rarer registration boundary. -#[async_trait] -pub trait ToolRegistry: Send + Sync + std::fmt::Debug { - /// Register a single tool against `connection_id`. - /// - /// The outcome reports whether the registration created a new entry, - /// updated an existing one, was shadowed by a higher-priority - /// registration, or was rejected. `reg.sessions` may be empty — the - /// tool is registered but unreachable until - /// [`Self::bind_tool_session`] adds at least one session binding. - /// Implementations must enforce per-`(connection_id, tool_id)` - /// uniqueness within their plane. - async fn register_tool( - &self, - connection_id: ConnectionId, - reg: ToolRegistration, - ) -> RegistrationOutcome; - - /// Register a multi-tool batch from a single tool server against - /// `connection_id`. - /// - /// Returns one [`RegistrationOutcome`] per tool in input order. Batch - /// semantics are best-effort: per-tool failures do not abort the rest - /// of the batch. The whole batch shares `reg.sessions` (which may be - /// empty). - async fn register_server( - &self, - connection_id: ConnectionId, - reg: ToolServerRegistration, - ) -> Vec; - - /// Drop the tool registered under `(connection_id, tool_id)`. Returns - /// `true` if a matching entry was removed, `false` if no such entry - /// existed. The tool is removed from every session it was bound to in - /// one shot — use [`Self::unbind_tool_session`] for per-session removal. - async fn unregister_tool(&self, connection_id: &ConnectionId, tool: &ToolId) -> bool; - - /// Drop every tool registered by `connection_id` under `server_id`. - /// Returns the number of entries removed. - async fn unregister_server(&self, connection_id: &ConnectionId, server: &ServerId) -> usize; - - /// Add `session_id` to the per-tool session set of - /// `(connection_id, tool_id)`. The caller (typically the WebSocket - /// router) is responsible for verifying that `session_id` is in the - /// connection's bound-session set before calling this method. - async fn bind_tool_session( - &self, - connection_id: &ConnectionId, - tool: &ToolId, - session_id: &SessionId, - ) -> ToolSessionBindOutcome; - - /// Remove `session_id` from the per-tool session set of - /// `(connection_id, tool_id)`. Does not unregister the tool itself. - async fn unbind_tool_session( - &self, - connection_id: &ConnectionId, - tool: &ToolId, - session_id: &SessionId, - ) -> ToolSessionUnbindOutcome; - - /// Drop every tool registered by `connection_id`. Used by the WebSocket - /// transport on disconnect cleanup. Returns counters describing how - /// much state was released. - async fn drop_connection(&self, connection_id: &ConnectionId) -> ConnectionCleanupReport; - - /// Look up the active resolution for `(session, tool)`. - /// - /// Returns `None` when no entry exists or when an entry exists but is - /// shadowed. A shadowed entry is never returned — the caller sees only - /// the active resolution. - fn find_tool(&self, session: &SessionId, tool: &ToolId) -> Option; - - /// Enumerate every active tool description for `session`, filtered by - /// the requested presentation `mode`. Implementations decide how to - /// honour the mode (e.g. omit non-meta tools when `Concise` is set). - fn list_tools(&self, session: &SessionId, mode: &ToolDefinitionMode) -> Vec; - - /// Enumerate active server summaries for `session`. Useful for - /// rendering connected-integrations system reminders. - fn list_servers(&self, session: &SessionId) -> Vec; - - /// Run a search query against the registry's index for `session`. - /// `limit` caps the result count; the snapshot reports how many - /// matches were hidden by the cap. - fn search(&self, session: &SessionId, query: &str, limit: usize) -> SearchSnapshot; - - /// Drop the binding to `session` from every tool that has it. The - /// affected tool records are NOT removed — their owning connection - /// retains them and may rebind via [`Self::bind_tool_session`]. Called - /// by the WebSocket transport when a session ends globally (no peer - /// connection still holds the binding) and by the connection actor - /// during per-disconnect cleanup. - async fn unregister_session(&self, session: &SessionId) -> SessionCleanupReport; - - /// Helper: set of session ids currently bound to `(connection_id, tool_id)`. - /// Returns an empty set when the tool is not registered. Mainly used - /// by tests to assert per-tool session set invariants without leaning - /// on the reverse index. - fn tool_sessions(&self, connection_id: &ConnectionId, tool: &ToolId) -> HashSet; - - /// All servers registered by this user across all connections. - fn list_servers_for_user(&self, user_id: &UserId) -> Vec; - - /// Look up a server by its connection ID. - fn get_server_record(&self, connection_id: &ConnectionId) -> Option; - - /// Look up only a server's id by its connection ID. Lighter than - /// [`Self::get_server_record`] for callers that need nothing else: - /// implementations should override the default to avoid deep-cloning - /// the whole record (notably its `metadata` JSON). - fn get_server_id(&self, connection_id: &ConnectionId) -> Option { - self.get_server_record(connection_id) - .map(|record| record.server_id) - } -} - -/// Server identity captured at `register_server` time. -#[derive(Debug, Clone)] -pub struct ServerRecord { - pub connection_id: ConnectionId, - pub user_id: UserId, - pub server_id: ServerId, - pub description: String, - pub metadata: serde_json::Value, - pub registered_at: chrono::DateTime, - /// Monotonic registration stamp ([`next_registration_seq`]) — the - /// stale-vs-revived discriminator for newest-wins (`registered_at` is display-only). - pub registration_seq: u64, -} - -/// Process-global hybrid logical clock: per-process strictly-increasing (no ties, -/// immune to NTP step-back) and epoch-seeded so stamps also roughly order across -/// replicas — only while inter-replica clock skew stays within the revive window -/// (`tool_route_ttl_ms`); past that, TTL eviction, not seq order, is the backstop. -/// The recency key for bind newest-wins and strictly-older eviction. -static REGISTRATION_CLOCK: AtomicU64 = AtomicU64::new(0); - -/// Issue the next monotonic registration stamp. See [`REGISTRATION_CLOCK`]. -pub fn next_registration_seq() -> u64 { - let now_ms = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0); - let candidate = now_ms << 10; - let mut prev = REGISTRATION_CLOCK.load(Ordering::Relaxed); - loop { - let next = candidate.max(prev + 1); - match REGISTRATION_CLOCK.compare_exchange_weak( - prev, - next, - Ordering::Relaxed, - Ordering::Relaxed, - ) { - Ok(_) => return next, - Err(actual) => prev = actual, - } - } -} - -#[cfg(test)] -mod seq_tests { - use super::next_registration_seq; - - fn now_ms() -> u64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0) - } - - #[test] - fn next_registration_seq_is_monotonic_and_epoch_seeded_under_burst() { - let before_ms = now_ms(); - let first = next_registration_seq(); - let mut prev = first; - const N: u64 = 50_000; - for _ in 0..N { - let s = next_registration_seq(); - assert!(s > prev, "must be strictly increasing: {prev} -> {s}"); - prev = s; - } - let after_ms = now_ms(); - - assert!( - prev - first >= N, - "burst must advance by at least one per call: {first} -> {prev}", - ); - let high = prev >> 10; - assert!( - high >= before_ms, - "high bits ({high}) must be epoch-seeded (>= {before_ms})", - ); - assert!( - high <= after_ms + 1_000, - "high bits ({high}) must track wall clock (<= {after_ms} + slack)", - ); - } -} diff --git a/crates/common/kigi-computer-hub-core/src/remote.rs b/crates/common/kigi-computer-hub-core/src/remote.rs deleted file mode 100644 index 67d351a..0000000 --- a/crates/common/kigi-computer-hub-core/src/remote.rs +++ /dev/null @@ -1,541 +0,0 @@ -//! `ConnectionClient` abstraction, `RemoteToolProxy`, and -//! `RemoteTransport`. -//! -//! `ConnectionClient` is the thin contract a downstream WebSocket SDK (or -//! an in-test channel-backed mock) implements; this crate stays free of -//! tokio-runtime / tokio-tungstenite deps so callers can pick their own. -//! -//! `RemoteToolProxy` wraps a remote tool registration so it implements -//! [`ToolHandle`] — the router routes through the same handle -//! type for local and remote registrations. `RemoteTransport` is the -//! transport-side equivalent: it forwards arbitrary `(tool_id, args)` -//! pairs over a [`ConnectionClient`] without needing a per-tool handle. - -use std::pin::Pin; -use std::sync::Arc; -use std::task::{Context, Poll}; - -use async_trait::async_trait; -use futures::Stream; -use futures::future::BoxFuture; -use futures::stream::BoxStream; -use serde_json::Value; -use tracing::warn; - -use kigi_tool_protocol::{ - JsonRpcId, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, JsonRpcVersion, Method, - ResponseOutcome, SessionId, ToolCallId, ToolCallParams, ToolCallProgressFrame, ToolCallResult, - ToolCapabilities, ToolErrorWire, ToolId, ToolOutputWire, UserId, WORKSPACE_UNAVAILABLE_SUBCODE, -}; -use kigi_tool_runtime::{ - BehaviorVersion, ContentBlock, Cwd, ListToolsContext, ToolCallContext, - ToolChatCompletionResponse, ToolError, ToolErrorKind, ToolProgress, ToolStream, ToolStreamItem, - TypedToolOutput, terminal_only, -}; -use kigi_tool_types::ToolDescription; - -use crate::resolver::ToolHandle; -use crate::transport::{Principal, Transport, TransportKind}; - -/// Object-safe contract for a connected remote endpoint. -/// -/// Concrete implementations supply the wire transport — the Rust SDK uses -/// `tokio_tungstenite`; tests use channel-backed mocks. Implementations -/// are expected to: -/// -/// - correlate request/response pairs by [`JsonRpcId`]; -/// - deliver progress notifications matching `tool_call_id` to whichever -/// subscriber registered for them; -/// - surface transport-level disconnects as [`ToolError::NetworkError`]. -#[async_trait] -pub trait ConnectionClient: Send + Sync + std::fmt::Debug { - /// Send a JSON-RPC request and await the matching response. Errors - /// signal a transport-level failure (write failed, connection closed - /// before the response arrived); a successful return carries the - /// response envelope verbatim, including method-level error outcomes. - async fn request(&self, request: JsonRpcRequest) -> Result; - - /// Subscribe to progress notifications for `tool_call_id`. - /// - /// The returned stream closes when the call's terminal frame arrives, - /// when the connection drops, or when the caller drops the receiver. - /// Subscribers MUST be registered before the corresponding request is - /// sent — otherwise progress frames that arrive before subscription - /// is complete are lost. - async fn subscribe_progress( - &self, - tool_call_id: ToolCallId, - ) -> BoxStream<'static, ToolCallProgressFrame>; - - /// Send a one-way notification (no response expected). Useful for - /// hook frames such as cancel. - async fn notify(&self, notification: JsonRpcNotification) -> Result<(), ToolError>; -} - -/// Wraps a remote registration so it dispatches through a connection. -/// -/// Identity, description, and capabilities come from the registration -/// snapshot held on the proxy; execution forwards a `tool_call_request` -/// over the connection and merges progress + terminal frames into a -/// single [`ToolStream`]. -#[derive(Debug, Clone)] -pub struct RemoteToolProxy { - tool_id: ToolId, - session_id: SessionId, - description: ToolDescription, - capabilities: ToolCapabilities, - connection: Arc, -} - -impl RemoteToolProxy { - /// Build a proxy bound to a single remote registration. - pub fn new( - tool_id: ToolId, - session_id: SessionId, - description: ToolDescription, - capabilities: ToolCapabilities, - connection: Arc, - ) -> Self { - Self { - tool_id, - session_id, - description, - capabilities, - connection, - } - } - - /// Bound session identifier. - pub fn session_id(&self) -> &SessionId { - &self.session_id - } -} - -#[async_trait] -impl ToolHandle for RemoteToolProxy { - fn id(&self) -> ToolId { - self.tool_id.clone() - } - - fn description(&self, _ctx: &ListToolsContext) -> ToolDescription { - self.description.clone() - } - - fn capabilities(&self) -> ToolCapabilities { - self.capabilities.clone() - } - - async fn execute(&self, ctx: ToolCallContext, args: Value) -> ToolStream { - dispatch_via_connection( - Arc::clone(&self.connection), - self.tool_id.clone(), - self.session_id.clone(), - args, - ctx, - ) - .await - } -} - -/// Transport that forwards calls over a [`ConnectionClient`]. -/// -/// The transport is bound to a single `(user_id, session_id)` at -/// construction. Calls do not require a pre-built proxy — the transport -/// builds the request frame from the `tool_id` it is asked to dispatch. -#[derive(Debug)] -pub struct RemoteTransport { - connection: Arc, - session_id: SessionId, - user_id: UserId, -} - -impl RemoteTransport { - /// Build a transport over `connection`, bound to `(user_id, - /// session_id)`. - pub fn new( - connection: Arc, - session_id: SessionId, - user_id: UserId, - ) -> Self { - Self { - connection, - session_id, - user_id, - } - } - - /// Bound session identifier. - pub fn session_id(&self) -> &SessionId { - &self.session_id - } - - /// Bound user identifier. - pub fn user_id(&self) -> &UserId { - &self.user_id - } -} - -#[async_trait] -impl Transport for RemoteTransport { - fn kind(&self) -> TransportKind { - TransportKind::Remote - } - - async fn authorize(&self) -> Result { - Ok(Principal::new(self.user_id.clone()).with_session(self.session_id.clone())) - } - - async fn call( - &self, - tool_id: ToolId, - args: Value, - ctx: ToolCallContext, - ) -> ToolStream { - dispatch_via_connection( - Arc::clone(&self.connection), - tool_id, - self.session_id.clone(), - args, - ctx, - ) - .await - } -} - -/// Subscribe to progress for `ctx.call_id`, send the `tool_call_request`, -/// and return a stream interleaving progress frames with the eventual -/// terminal item. -/// -/// Subscribing **before** sending is the contract that -/// [`ConnectionClient::subscribe_progress`] requires; doing so here keeps -/// individual transports / proxies from re-implementing the dance. -async fn dispatch_via_connection( - connection: Arc, - tool_id: ToolId, - session_id: SessionId, - arguments: Value, - ctx: ToolCallContext, -) -> ToolStream { - let cwd = ctx - .extensions - .get::() - .map(|c| c.0.to_string_lossy().into_owned()); - let behavior_version = ctx.extensions.get::().map(|v| v.0.clone()); - let call_id = ctx.call_id; - - // Subscribe BEFORE sending. The single remaining `call_id.clone()` - // is unavoidable: subscription needs an owned id and the same id has - // to land in the request params below. - let progress = connection.subscribe_progress(call_id.clone()).await; - - let params = ToolCallParams { - tool_call_id: call_id, - tool_id, - arguments, - deadline_ms: None, - behavior_version, - cwd, - // The ctx `TraceContext` extension is receive-side state. - trace_context: None, - }; - - let request = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::new_uuid_v7(), - session_id: Some(session_id), - method: Method::ToolCallRequest.as_wire_str().to_string(), - params: match serde_json::to_value(¶ms) { - Ok(v) => v, - Err(e) => { - return terminal_only(Err(ToolError::custom("request_encoding", e.to_string()))); - } - }, - }; - - // Build the response future without awaiting it here so progress and - // terminal can be polled concurrently from the returned stream. - let request_fut = Box::pin(async move { connection.request(request).await }); - - Box::pin(RequestStream { - tool_id: Some(params.tool_id), - progress, - request: Some(request_fut), - done: false, - }) -} - -/// Owned response future with `'static` lifetime so the stream can hold -/// it across polls. -type ResponseFuture = BoxFuture<'static, Result>; - -/// Stream that interleaves wire-side progress frames with the eventual -/// JSON-RPC response, ending with exactly one terminal item. -struct RequestStream { - /// Consumed exactly once when the terminal is built. - tool_id: Option, - progress: BoxStream<'static, ToolCallProgressFrame>, - request: Option, - done: bool, -} - -impl Stream for RequestStream { - type Item = ToolStreamItem; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - if self.done { - return Poll::Ready(None); - } - - // Poll the response first so the terminal short-circuits the - // moment it lands. Any progress frames that arrived alongside - // the response are dropped — once `Terminal` is emitted, `done` - // is set and the next poll returns `None` immediately without - // re-polling the progress stream. The router invariant is - // "`Progress* Terminal`, exactly one terminal"; dropping any - // post-terminal progress is what makes that invariant hold here. - if let Some(req_fut) = self.request.as_mut() { - match req_fut.as_mut().poll(cx) { - Poll::Ready(result) => { - self.done = true; - self.request = None; - let Some(tool_id) = self.tool_id.take() else { - return Poll::Ready(None); - }; - let terminal = match result { - Ok(resp) => terminal_from_response(tool_id, resp), - Err(err) => Err(err), - }; - return Poll::Ready(Some(ToolStreamItem::Terminal(terminal))); - } - Poll::Pending => {} - } - } else { - self.done = true; - return Poll::Ready(None); - } - - // Poll the progress stream while the request is pending. Closing - // the progress stream is fine — the response future is still - // registered for wake-up. - match Pin::new(&mut self.progress).poll_next(cx) { - Poll::Ready(Some(frame)) => { - Poll::Ready(Some(ToolStreamItem::Progress(progress_from_frame(frame)))) - } - Poll::Ready(None) | Poll::Pending => Poll::Pending, - } - } -} - -/// Map a wire-side [`ToolCallProgressFrame`] into a runtime -/// [`ToolProgress`]. `kind` becomes the `Custom` subkind so callers can -/// dispatch on the producer-defined identifier without losing the body. -pub fn progress_from_frame(frame: ToolCallProgressFrame) -> ToolProgress { - ToolProgress::Custom { - subkind: frame.kind, - payload: frame.body, - } -} - -/// Decode the response envelope into the terminal -/// `Result` the runtime expects. -fn terminal_from_response( - tool_id: ToolId, - resp: JsonRpcResponse, -) -> Result { - match resp.outcome { - ResponseOutcome::Result(value) => decode_call_result(tool_id, value), - ResponseOutcome::Error(err) => Err(error_from_envelope(err)), - } -} - -/// Decode a `tool_call_result` success body into the terminal -/// [`TypedToolOutput`]. Shared by the core remote proxy and the SDK -/// harness so both wire decoders reconstruct `chat_completion_output` -/// identically. -/// -/// A body with a `tool_call_id` is decoded strictly (`response_decoding` on -/// failure), reconstructing `chat_completion_output` (an unparseable cco -/// degrades to `None`). A bare body — e.g. a hub-local tool's raw output — -/// passes through unchanged. -pub fn decode_call_result(tool_id: ToolId, value: Value) -> Result { - if value.get("tool_call_id").is_none() { - return Ok(TypedToolOutput::from_value(tool_id, value)); - } - let result: ToolCallResult = serde_json::from_value(value) - .map_err(|e| ToolError::custom("response_decoding", e.to_string()))?; - let chat_completion_output = result.chat_completion_output.and_then(|cco| { - serde_json::from_value::(cco) - .inspect_err(|e| { - warn!(tool_id = %tool_id, error = %e, "dropping unparseable chat_completion_output"); - }) - .ok() - }); - let value = output_to_value(result.output); - Ok(TypedToolOutput::from_value(tool_id, value) - .with_chat_completion_output(chat_completion_output)) -} - -/// Project a wire [`ToolOutputWire`] into a JSON [`Value`]. -/// -/// Three shapes collapse to one runtime type: -/// - `Text` becomes a JSON string; -/// - `Json` is forwarded verbatim; -/// - `Mcp { blocks }` is re-serialised as `{ "blocks": [ContentBlock, ...] }` -/// so the same downstream decoder used for in-process content blocks -/// works without case-by-case adaptation. -pub fn output_to_value(output: ToolOutputWire) -> Value { - match output { - ToolOutputWire::Text(s) => Value::String(s), - ToolOutputWire::Json(v) => v, - ToolOutputWire::Mcp { blocks } => { - let runtime_blocks: Vec = blocks.into_iter().map(map_block).collect(); - // `ContentBlock`'s derived `Serialize` impl never fails for any - // valid in-memory variant, but `to_value` is fallible at the - // type level; collapse a hypothetical failure to `Value::Null` - // before wrapping so this function stays total without an - // `unwrap`. The outer `json!` only sees a `Value` expression - // (which `to_value` round-trips infallibly), so the macro's - // hidden `to_value` call cannot panic here. - let blocks_value = serde_json::to_value(&runtime_blocks).unwrap_or(Value::Null); - serde_json::json!({ "blocks": blocks_value }) - } - } -} - -fn map_block(block: kigi_tool_protocol::McpBlock) -> ContentBlock { - use kigi_tool_protocol::McpBlock; - match block { - McpBlock::Text { text } => ContentBlock::Text { text }, - McpBlock::Image { mime_type, data } => ContentBlock::Image { - mime_type, - data, - media_id: None, - filename: None, - path: None, - metadata: Default::default(), - }, - McpBlock::Resource { - uri, - mime_type, - text, - } => ContentBlock::Resource { - uri, - mime_type, - text, - }, - } -} - -/// Decode a JSON-RPC error envelope into a [`ToolError`]. The envelope's -/// `data` field is expected to carry a serialised [`ToolErrorWire`] when -/// available; falls back to a [`ToolError::Custom`] keyed by the numeric -/// envelope code when the data shape is unknown. -pub fn error_from_envelope(err: kigi_tool_protocol::JsonRpcError) -> ToolError { - if let Some(data) = err.data.clone() - && let Ok(wire) = serde_json::from_value::(data) - { - return tool_error_from_wire(wire); - } - let mut e = ToolError::custom(format!("jsonrpc_{}", err.code), err.message); - if let Some(data) = err.data { - e = e.with_details(data); - } - e -} - -/// Recognize the hub's `workspace_unavailable` error on an already-decoded -/// [`ToolError`]. Keys on `details["code"]` — the field that survives -/// `ToolError::custom` + `with_details` — not the numeric code or the wire -/// `Custom.subcode`. -pub fn is_workspace_unavailable(err: &ToolError) -> bool { - err.kind == ToolErrorKind::Custom - && err - .details - .as_ref() - .and_then(|d| d.get("code")) - .and_then(|v| v.as_str()) - == Some(WORKSPACE_UNAVAILABLE_SUBCODE) -} - -/// Map [`ToolErrorWire`] back into the runtime [`ToolError`]. The runtime -/// error variants are the source-of-truth taxonomy; the wire form is a -/// lossy projection onto stable codes for serialisation, so a few wire -/// variants land on [`ToolError::Custom`] keyed by their wire code -/// rather than a dedicated runtime variant. -pub fn tool_error_from_wire(wire: ToolErrorWire) -> ToolError { - match wire { - ToolErrorWire::InvalidArguments { message, details } => { - let e = ToolError::invalid_arguments(message); - match details { - Some(d) => e.with_details(d), - None => e, - } - } - ToolErrorWire::ToolNotFound { tool_id } => { - let detail = format!("tool not found: {tool_id}"); - ToolError::not_found(tool_id, detail) - } - ToolErrorWire::PermissionDenied { reason } => ToolError::permission_denied(reason), - ToolErrorWire::Timeout { - tool_id, - elapsed_ms, - } => ToolError::new( - ToolErrorKind::Timeout, - format!("timed out after {elapsed_ms}ms"), - ) - .with_details(serde_json::json!({"tool_id": tool_id.as_str(), "elapsed_ms": elapsed_ms})), - ToolErrorWire::Cancelled { tool_id } => ToolError::cancelled(tool_id, "cancelled"), - ToolErrorWire::Execution { tool_id, message } => ToolError::execution(tool_id, message), - ToolErrorWire::BehaviorVersionUnsupported { tool_id, requested } => ToolError::new( - ToolErrorKind::BehaviorVersionUnsupported, - format!("behavior version {requested} not supported"), - ) - .with_details(serde_json::json!({"tool_id": tool_id.as_str(), "requested": requested})), - ToolErrorWire::RenderLimited { - tool_id, - card_id, - reason, - } => ToolError::new(ToolErrorKind::RenderLimited, reason) - .with_details(serde_json::json!({"tool_id": tool_id.as_str(), "card_id": card_id})), - ToolErrorWire::TerminalError { tool_id, message } => { - ToolError::terminal_error(tool_id, message) - } - ToolErrorWire::Custom { - subcode, - message, - details, - } => { - let e = ToolError::custom(subcode, message); - match details { - Some(d) => e.with_details(d), - None => e, - } - } - ToolErrorWire::SessionMismatch => ToolError::custom("session_mismatch", "session mismatch"), - ToolErrorWire::TransportClosed { tool_id } => { - ToolError::network_error(format!("transport closed for {tool_id}")) - } - ToolErrorWire::UnsupportedProtocolVersion { supported } => ToolError::custom( - "unsupported_protocol_version", - format!("supported versions: {supported:?}"), - ), - ToolErrorWire::PayloadTooLarge { bytes, limit } => ToolError::custom( - "payload_too_large", - format!("payload {bytes} bytes exceeds limit {limit}"), - ), - ToolErrorWire::Internal { request_id, detail } => { - let e = ToolError::custom( - "internal_error", - detail.unwrap_or_else(|| "internal router error".to_owned()), - ); - match request_id { - // Keep `code` alongside `request_id`: `with_details` replaces - // the `{"code": …}` object `ToolError::custom` installed. - Some(id) => e.with_details( - serde_json::json!({ "code": "internal_error", "request_id": id.as_str() }), - ), - None => e, - } - } - } -} diff --git a/crates/common/kigi-computer-hub-core/src/resolver.rs b/crates/common/kigi-computer-hub-core/src/resolver.rs deleted file mode 100644 index 24ded90..0000000 --- a/crates/common/kigi-computer-hub-core/src/resolver.rs +++ /dev/null @@ -1,280 +0,0 @@ -//! `CompoundResolver` plus the `ResolvedTool` and `ToolHandle` -//! types it returns. -//! -//! `Tool` carries associated `Args` / `Output` types and is therefore not -//! object-safe. [`ToolHandle`] is the dyn-compatible projection used -//! by every router build: typed tools are wrapped via -//! [`ErasedTool::new`]; remote registrations expose -//! [`crate::RemoteToolProxy`] which implements [`ToolHandle`] -//! directly without an intermediate typed `Tool` impl. - -use std::sync::Arc; - -use async_trait::async_trait; -use futures::StreamExt; -use serde_json::Value; - -use kigi_tool_protocol::{SessionId, ToolCapabilities, ToolId, ToolRegistration}; -use kigi_tool_runtime::{ - ListToolsContext, Tool, ToolCallContext, ToolError, ToolOutput, ToolStream, ToolStreamItem, - TypedToolOutput, terminal_only, -}; -use kigi_tool_types::ToolDescription; - -use crate::registry::ToolRegistry; - -/// Active resolution returned by [`CompoundResolver::resolve`]. -/// -/// Variants share the same `tool` handle and `registration` shape; the -/// discriminant only tells callers whether the executing handle dispatches -/// in-process or forwards over a connection. Differentiating the variants -/// is useful for metrics, log tags, and the local-shadows-remote rule -/// applied when both planes register the same `tool_id`. -#[derive(Debug, Clone)] -pub enum ResolvedTool { - /// In-process tool resolved from the local registry. - Local { - /// Object-safe handle to the tool's `execute` entry point. - tool: Arc, - /// Wire-shape registration record. Carries `tool_id`, the - /// schema-bearing description, capabilities, and ownership data. - registration: ToolRegistration, - }, - /// Remote registration resolved through a connection-backed proxy. - Remote { - /// Object-safe handle whose `execute` forwards over the - /// owning connection. - proxy: Arc, - /// Wire-shape registration record (same shape as the local - /// variant — both store the active registration so callers do not - /// have to round-trip the registry for description / capabilities). - registration: ToolRegistration, - }, -} - -impl ResolvedTool { - /// Borrow the registration record regardless of variant. - pub fn registration(&self) -> &ToolRegistration { - match self { - Self::Local { registration, .. } | Self::Remote { registration, .. } => registration, - } - } - - /// Borrow the executing handle regardless of variant. - pub fn handle(&self) -> &Arc { - match self { - Self::Local { tool, .. } => tool, - Self::Remote { proxy, .. } => proxy, - } - } -} - -/// Object-safe projection of a registered tool. -/// -/// The router only needs identity, description, capabilities, and a -/// JSON-typed `execute` entry point — exactly what this trait exposes. -/// Adapters that wrap a typed `Tool` impl get [`ErasedTool`] for free; -/// non-`Tool` handles (notably remote proxies) implement this trait -/// directly. -#[async_trait] -pub trait ToolHandle: Send + Sync + std::fmt::Debug { - /// Stable identity used by the router to route calls. - fn id(&self) -> ToolId; - - /// Model-facing description of the tool's argument schema. - /// - /// Receives the per-turn [`ListToolsContext`] so handles backed by a - /// typed [`Tool`] can produce context-aware descriptions at listing - /// time. Callers outside a listing turn pass - /// [`ListToolsContext::default`]. - fn description(&self, ctx: &ListToolsContext) -> ToolDescription; - - /// Per-tool capability flags. - fn capabilities(&self) -> ToolCapabilities; - - /// Per-turn listing predicate. - fn should_list(&self, _ctx: &ListToolsContext) -> bool { - true - } - - /// Streaming execution entry point. - /// - /// Implementations encode the tool's typed `Output` to - /// [`serde_json::Value`] and surface argument-decoding failures as - /// [`ToolError::InvalidArguments`] within the terminal item. - async fn execute(&self, ctx: ToolCallContext, args: Value) -> ToolStream; -} - -/// Type-erasing wrapper for any [`Tool`] implementation. -/// -/// Decodes `args` into `T::Args`, drives `T::execute`, and re-encodes each -/// `T::Output` (terminal and progress items pass through unchanged -/// otherwise). The wrapper holds the inner tool by `Arc` so the same -/// underlying instance can back multiple registrations cheaply. -pub struct ErasedTool { - inner: Arc, -} - -impl ErasedTool { - /// Wrap an `Arc` for use as an [`ToolHandle`]. - pub fn from_arc(inner: Arc) -> Self { - Self { inner } - } - - /// Wrap an owned tool, taking the `Arc` allocation internally. - pub fn new(inner: T) -> Self { - Self::from_arc(Arc::new(inner)) - } -} - -impl std::fmt::Debug for ErasedTool { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ErasedTool") - .field("inner", &self.inner) - .finish() - } -} - -impl Clone for ErasedTool { - fn clone(&self) -> Self { - Self { - inner: Arc::clone(&self.inner), - } - } -} - -#[async_trait] -impl ToolHandle for ErasedTool -where - T: Tool + std::fmt::Debug + 'static, - T::Output: ToolOutput, -{ - fn id(&self) -> ToolId { - self.inner.id() - } - - fn description(&self, ctx: &ListToolsContext) -> ToolDescription { - self.inner.description(ctx) - } - - fn capabilities(&self) -> ToolCapabilities { - self.inner.capabilities() - } - - fn should_list(&self, ctx: &ListToolsContext) -> bool { - self.inner.should_list(ctx) - } - - async fn execute(&self, ctx: ToolCallContext, args: Value) -> ToolStream { - let typed_args: T::Args = match serde_json::from_value(args) { - Ok(a) => a, - Err(e) => { - return terminal_only(Err(ToolError::invalid_arguments(e.to_string()))); - } - }; - let tool_id = self.inner.id(); - let stream = self.inner.execute(ctx, typed_args).await; - let mapped = stream.map(move |item| match item { - ToolStreamItem::Progress(p) => ToolStreamItem::Progress(p), - ToolStreamItem::Terminal(Ok(out)) => match serde_json::to_value(&out) { - Ok(value) => { - let custom = out.model_output(); - let model_output = if custom.is_empty() { - kigi_tool_runtime::extract_content_blocks(&value) - } else { - custom - }; - let chat_completion_output = out.chat_completion_output(); - ToolStreamItem::Terminal(Ok(TypedToolOutput { - tool_id: tool_id.clone(), - value, - model_output, - chat_completion_output, - })) - } - Err(e) => ToolStreamItem::Terminal(Err(ToolError::custom( - "output_encoding", - e.to_string(), - ))), - }, - ToolStreamItem::Terminal(Err(err)) => ToolStreamItem::Terminal(Err(err)), - }); - Box::pin(mapped) - } -} - -/// Compose a local-first lookup over one (`local_only`) or two -/// (`compound`) registries. -/// -/// The lookup contract: `find_tool` is called on the local registry first; -/// only if it returns `None` is the remote registry consulted. Any local -/// registration shadows a same-id remote registration. Cross-session -/// lookups return `None` — the caller may surface this as a -/// [`ToolError::NotFound`] to keep ownership invisible to the requester. -#[derive(Debug)] -pub struct CompoundResolver { - local: Arc, - remote: Option>, -} - -impl CompoundResolver { - /// Compose a resolver that consults a single local registry. - pub fn local_only(local: Arc) -> Self { - Self { - local, - remote: None, - } - } - - /// Compose a resolver with both planes; `local` is consulted first. - pub fn compound(local: Arc, remote: Arc) -> Self { - Self { - local, - remote: Some(remote), - } - } - - /// Borrow the local plane. - pub fn local(&self) -> &Arc { - &self.local - } - - /// Borrow the optional remote plane. - pub fn remote(&self) -> Option<&Arc> { - self.remote.as_ref() - } - - /// Resolve `(session, tool_id)` honouring the local-first rule. - pub fn resolve(&self, session: &SessionId, tool_id: &ToolId) -> Option { - if let Some(hit) = self.local.find_tool(session, tool_id) { - return Some(hit); - } - self.remote - .as_ref() - .and_then(|r| r.find_tool(session, tool_id)) - } - - /// Resolve `(session, tool_id)` and dispatch through the active - /// handle, returning the tool's stream verbatim. Misses produce a - /// single-item terminal stream carrying [`ToolError::NotFound`]. - /// - /// Centralises the resolve-then-dispatch sequence so both the - /// transport-side `LocalTransport::call` and the inner-dispatch path - /// share one implementation: a future change to the miss-shape (or - /// to the dispatch contract) lands once. - pub async fn resolve_and_dispatch( - &self, - session: &SessionId, - tool_id: ToolId, - args: Value, - ctx: ToolCallContext, - ) -> ToolStream { - match self.resolve(session, &tool_id) { - Some(resolved) => resolved.handle().execute(ctx, args).await, - None => terminal_only(Err(ToolError::not_found( - tool_id.clone(), - format!("tool not found: {tool_id}"), - ))), - } - } -} diff --git a/crates/common/kigi-computer-hub-core/src/transport.rs b/crates/common/kigi-computer-hub-core/src/transport.rs deleted file mode 100644 index b27b079..0000000 --- a/crates/common/kigi-computer-hub-core/src/transport.rs +++ /dev/null @@ -1,116 +0,0 @@ -//! Object-safe `Transport` trait plus the `Principal` value carried across -//! authorize/call boundaries. -//! -//! [`TransportKind`] is re-exported from [`kigi_tool_protocol`] so the wire -//! and dispatch layers share one canonical enum and there is no duplicate -//! `Local` / `Remote` definition to keep in sync. - -use async_trait::async_trait; -use serde_json::Value; - -use kigi_tool_protocol::{SessionId, ToolId, UserId}; -use kigi_tool_runtime::{ToolCallContext, ToolError, ToolStream, TypedToolOutput}; - -pub use kigi_tool_protocol::TransportKind; - -/// Authenticated identity bound to a transport at handshake time. -/// -/// The transport authorises **once** at connect; subsequent dispatch calls -/// carry no extra credentials. `session_ids` is plural because a JWT may -/// authorise more than one session (multi-tenant tooling sessions sharing -/// a single user identity); the router narrows by [`SessionId`] at the -/// per-call boundary. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct Principal { - /// Authenticated user identity. - pub user_id: UserId, - - /// Sessions this principal is authorised to act on. Empty when the - /// transport authorises a user but has not yet bound a session - /// (e.g. a fresh harness connection that has not opened a session). - pub session_ids: Vec, - - /// OAuth-style scopes granted to this principal, e.g. `"tool.invoke"`. - pub scopes: Vec, - - /// Token audiences claimed by the credential, e.g. the router's - /// expected `aud` values. Used by callers that need defence-in-depth - /// audience checks beyond what the transport already validated. - pub audiences: Vec, -} - -impl Principal { - /// Build a principal for `user_id` with no sessions, scopes, or - /// audiences. Use the `with_*` builders to populate the rest. - pub fn new(user_id: UserId) -> Self { - Self { - user_id, - session_ids: Vec::new(), - scopes: Vec::new(), - audiences: Vec::new(), - } - } - - /// Append `session_id` to the authorised set. - pub fn with_session(mut self, session_id: SessionId) -> Self { - self.session_ids.push(session_id); - self - } - - /// Append `scope` to the granted scopes. - pub fn with_scope(mut self, scope: impl Into) -> Self { - self.scopes.push(scope.into()); - self - } - - /// Append `aud` to the token's audience list. - pub fn with_audience(mut self, aud: impl Into) -> Self { - self.audiences.push(aud.into()); - self - } - - /// Whether `scope` is present in the granted scopes. - pub fn has_scope(&self, scope: &str) -> bool { - self.scopes.iter().any(|s| s == scope) - } - - /// Whether `session_id` is in the principal's authorised session set. - pub fn authorizes_session(&self, session_id: &SessionId) -> bool { - self.session_ids.iter().any(|s| s == session_id) - } -} - -/// Object-safe transport for dispatching tool calls. -/// -/// Implementations come in two flavours: [`TransportKind::Local`] resolves -/// against an in-process registry, while [`TransportKind::Remote`] forwards -/// a `tool_call_request` over a [`crate::ConnectionClient`]. -#[async_trait] -pub trait Transport: Send + Sync + std::fmt::Debug { - /// Whether the underlying transport is local (in-process) or remote - /// (forwarded over a connection). - fn kind(&self) -> TransportKind; - - /// One-time authorisation handshake. - /// - /// Local transports return a principal derived from the bound OS user - /// (or whatever ambient identity the host process provides). Remote - /// transports return the principal extracted from a validated - /// credential. Subsequent [`Self::call`] invocations reuse this - /// principal — the router never re-authorises per call. - async fn authorize(&self) -> Result; - - /// Dispatch a tool call. - /// - /// The returned [`ToolStream`] follows the runtime invariant: zero or - /// more `Progress` items followed by exactly one `Terminal`. A - /// not-found result is reported as a single-item terminal stream - /// carrying [`ToolError::NotFound`]; transport-level disconnects - /// surface as [`ToolError::NetworkError`]. - async fn call( - &self, - tool_id: ToolId, - args: Value, - ctx: ToolCallContext, - ) -> ToolStream; -} diff --git a/crates/common/kigi-computer-hub-core/tests/compound_resolver.rs b/crates/common/kigi-computer-hub-core/tests/compound_resolver.rs deleted file mode 100644 index 6e35792..0000000 --- a/crates/common/kigi-computer-hub-core/tests/compound_resolver.rs +++ /dev/null @@ -1,335 +0,0 @@ -//! `CompoundResolver` and `ResolvedTool` coverage. Exercises local-only, -//! local-shadows-remote, remote-fallback, and cross-session scenarios. - -use std::sync::Arc; - -use dashmap::DashMap; - -use async_trait::async_trait; - -use futures::StreamExt; -use kigi_computer_hub_core::{ - CompoundResolver, ConnectionCleanupReport, ErasedTool, ResolvedTool, SessionCleanupReport, - ToolHandle, ToolRegistry, ToolSessionBindOutcome, ToolSessionUnbindOutcome, -}; -use kigi_tool_protocol::{ - ConnectionId, RegistrationOutcome, ServerId, SessionId, ToolDefinitionMode, ToolId, - ToolRegistration, ToolServerRegistration, TransportKind, UserId, -}; -use kigi_tool_runtime::{ - SearchSnapshot, ServerSummary, Tool, ToolCallContext, ToolError, ToolStreamItem, -}; -use kigi_tool_types::ToolDescription; -use serde::{Deserialize, Serialize}; - -#[derive(Debug, Default, Clone, Serialize, Deserialize, schemars::JsonSchema)] -struct EmptyArgs {} - -#[derive(Debug)] -struct StubTool { - id: ToolId, -} - -impl Tool for StubTool { - type Args = EmptyArgs; - type Output = serde_json::Value; - - fn id(&self) -> ToolId { - self.id.clone() - } - - fn description(&self, _ctx: &::kigi_tool_runtime::ListToolsContext) -> ToolDescription { - ToolDescription::new(self.id.as_str(), format!("stub for {}", self.id)) - } - - async fn run( - &self, - _ctx: ToolCallContext, - _args: Self::Args, - ) -> Result { - Ok(serde_json::json!({"id": self.id.as_str()})) - } -} - -#[derive(Debug)] -struct PlaneRegistry { - // Set once at construction; `TransportKind` is `Copy` so a direct - // field is the obvious choice — no interior mutability required. - transport_kind: TransportKind, - entries: DashMap<(SessionId, ToolId), ToolRegistration>, - handles: DashMap>, -} - -impl PlaneRegistry { - fn new(kind: TransportKind) -> Self { - Self { - transport_kind: kind, - entries: DashMap::new(), - handles: DashMap::new(), - } - } - - fn install(&self, session: &SessionId, id: &ToolId) { - let reg = ToolRegistration { - tool_id: id.clone(), - sessions: Some(vec![session.clone()]), - user_id: UserId::new("alice").expect("user id"), - server_id: None, - description: ToolDescription::new(id.as_str(), format!("stub for {id}")), - input_schema: None, - capabilities: None, - notification_schemas: None, - transport_kind: self.transport_kind, - if_match_generation: None, - metadata: None, - }; - self.entries.insert((session.clone(), id.clone()), reg); - self.handles.insert( - id.clone(), - Arc::new(ErasedTool::new(StubTool { id: id.clone() })), - ); - } -} - -#[async_trait] -impl ToolRegistry for PlaneRegistry { - async fn register_tool( - &self, - _connection_id: ConnectionId, - _reg: ToolRegistration, - ) -> RegistrationOutcome { - unreachable!("resolver tests pre-populate via install()") - } - async fn register_server( - &self, - _connection_id: ConnectionId, - _reg: ToolServerRegistration, - ) -> Vec { - unreachable!() - } - async fn unregister_tool(&self, _connection_id: &ConnectionId, _tool: &ToolId) -> bool { - unreachable!() - } - async fn unregister_server(&self, _connection_id: &ConnectionId, _server: &ServerId) -> usize { - unreachable!() - } - async fn bind_tool_session( - &self, - _connection_id: &ConnectionId, - _tool: &ToolId, - _session_id: &SessionId, - ) -> ToolSessionBindOutcome { - unreachable!() - } - async fn unbind_tool_session( - &self, - _connection_id: &ConnectionId, - _tool: &ToolId, - _session_id: &SessionId, - ) -> ToolSessionUnbindOutcome { - unreachable!() - } - async fn drop_connection(&self, _connection_id: &ConnectionId) -> ConnectionCleanupReport { - ConnectionCleanupReport::default() - } - - fn find_tool(&self, session: &SessionId, tool: &ToolId) -> Option { - let registration = self - .entries - .get(&(session.clone(), tool.clone()))? - .value() - .clone(); - let handle = self.handles.get(tool)?.value().clone(); - match registration.transport_kind { - TransportKind::Local => Some(ResolvedTool::Local { - tool: handle, - registration, - }), - TransportKind::Remote => Some(ResolvedTool::Remote { - proxy: handle, - registration, - }), - } - } - - fn list_tools(&self, _session: &SessionId, _mode: &ToolDefinitionMode) -> Vec { - vec![] - } - fn list_servers(&self, _session: &SessionId) -> Vec { - vec![] - } - fn search(&self, _session: &SessionId, _query: &str, _limit: usize) -> SearchSnapshot { - SearchSnapshot { - results: vec![], - total_hidden_tools: 0, - is_ready: true, - } - } - async fn unregister_session(&self, _session: &SessionId) -> SessionCleanupReport { - SessionCleanupReport::default() - } - fn tool_sessions( - &self, - _connection_id: &ConnectionId, - _tool: &ToolId, - ) -> std::collections::HashSet { - std::collections::HashSet::new() - } - - fn list_servers_for_user( - &self, - _user_id: &kigi_tool_protocol::UserId, - ) -> Vec { - Vec::new() - } - - fn get_server_record( - &self, - _connection_id: &ConnectionId, - ) -> Option { - None - } -} - -fn sid(s: &str) -> SessionId { - SessionId::new(s).expect("session id") -} - -fn tid(s: &str) -> ToolId { - ToolId::new(s).expect("tool id") -} - -#[tokio::test] -async fn local_only_resolves_local_hits() { - let local = Arc::new(PlaneRegistry::new(TransportKind::Local)); - local.install(&sid("sess-1"), &tid("foo")); - let resolver = CompoundResolver::local_only(local as Arc); - match resolver.resolve(&sid("sess-1"), &tid("foo")) { - Some(ResolvedTool::Local { registration, .. }) => { - assert_eq!(registration.tool_id, tid("foo")); - } - other => panic!("expected Local, got {other:?}"), - } -} - -#[tokio::test] -async fn local_only_returns_none_for_unknown() { - let local = Arc::new(PlaneRegistry::new(TransportKind::Local)); - let resolver = CompoundResolver::local_only(local as Arc); - assert!(resolver.resolve(&sid("sess-1"), &tid("missing")).is_none()); -} - -#[tokio::test] -async fn compound_falls_through_to_remote_when_local_misses() { - let local = Arc::new(PlaneRegistry::new(TransportKind::Local)); - let remote = Arc::new(PlaneRegistry::new(TransportKind::Remote)); - remote.install(&sid("sess-1"), &tid("foo")); - let resolver = CompoundResolver::compound( - local as Arc, - remote as Arc, - ); - match resolver.resolve(&sid("sess-1"), &tid("foo")) { - Some(ResolvedTool::Remote { registration, .. }) => { - assert_eq!(registration.tool_id, tid("foo")); - assert_eq!(registration.transport_kind, TransportKind::Remote); - } - other => panic!("expected Remote, got {other:?}"), - } -} - -#[tokio::test] -async fn local_shadows_same_id_remote() { - let local = Arc::new(PlaneRegistry::new(TransportKind::Local)); - local.install(&sid("sess-1"), &tid("foo")); - let remote = Arc::new(PlaneRegistry::new(TransportKind::Remote)); - remote.install(&sid("sess-1"), &tid("foo")); - let resolver = CompoundResolver::compound( - local as Arc, - remote as Arc, - ); - match resolver.resolve(&sid("sess-1"), &tid("foo")) { - Some(ResolvedTool::Local { registration, .. }) => { - assert_eq!(registration.transport_kind, TransportKind::Local); - } - other => panic!("expected local resolution to shadow remote, got {other:?}"), - } -} - -#[tokio::test] -async fn cross_session_lookup_returns_none() { - let local = Arc::new(PlaneRegistry::new(TransportKind::Local)); - local.install(&sid("sess-1"), &tid("foo")); - let resolver = CompoundResolver::local_only(local as Arc); - assert!(resolver.resolve(&sid("sess-other"), &tid("foo")).is_none()); -} - -#[tokio::test] -async fn compound_returns_none_when_neither_plane_holds_id() { - let local = Arc::new(PlaneRegistry::new(TransportKind::Local)); - let remote = Arc::new(PlaneRegistry::new(TransportKind::Remote)); - let resolver = CompoundResolver::compound( - local as Arc, - remote as Arc, - ); - assert!(resolver.resolve(&sid("sess-1"), &tid("foo")).is_none()); -} - -#[tokio::test] -async fn resolved_tool_helpers_borrow_active_handle_and_registration() { - let local = Arc::new(PlaneRegistry::new(TransportKind::Local)); - local.install(&sid("sess-1"), &tid("foo")); - let resolver = CompoundResolver::local_only(local as Arc); - let resolved = resolver.resolve(&sid("sess-1"), &tid("foo")).expect("hit"); - assert_eq!(resolved.registration().tool_id, tid("foo")); - assert_eq!(resolved.handle().id(), tid("foo")); -} - -#[tokio::test] -async fn resolve_and_dispatch_drives_the_resolved_handle() { - let local = Arc::new(PlaneRegistry::new(TransportKind::Local)); - local.install(&sid("sess-1"), &tid("foo")); - let resolver = CompoundResolver::local_only(local as Arc); - let mut stream = resolver - .resolve_and_dispatch( - &sid("sess-1"), - tid("foo"), - serde_json::json!({}), - ToolCallContext::default(), - ) - .await; - let item = stream.next().await.expect("terminal"); - match item { - ToolStreamItem::Terminal(Ok(typed)) => { - assert_eq!(typed.value, serde_json::json!({"id": "foo"})); - } - other => panic!("expected Terminal(Ok), got {other:?}"), - } - assert!(stream.next().await.is_none()); -} - -#[tokio::test] -async fn resolve_and_dispatch_misses_yield_terminal_not_found() { - let local = Arc::new(PlaneRegistry::new(TransportKind::Local)); - let resolver = CompoundResolver::local_only(local as Arc); - let mut stream = resolver - .resolve_and_dispatch( - &sid("sess-1"), - tid("missing"), - serde_json::json!(null), - ToolCallContext::default(), - ) - .await; - let item = stream.next().await.expect("terminal"); - match item { - ToolStreamItem::Terminal(Err(ref e)) - if e.kind == kigi_tool_runtime::ToolErrorKind::NotFound => - { - assert!( - e.detail.contains("missing"), - "detail should mention tool id: {}", - e.detail - ); - } - other => panic!("expected Terminal(Err(NotFound)), got {other:?}"), - } -} diff --git a/crates/common/kigi-computer-hub-core/tests/inner_dispatch.rs b/crates/common/kigi-computer-hub-core/tests/inner_dispatch.rs deleted file mode 100644 index 10cda3b..0000000 --- a/crates/common/kigi-computer-hub-core/tests/inner_dispatch.rs +++ /dev/null @@ -1,317 +0,0 @@ -//! `InnerDispatchForResolver` coverage. Verifies the cycle-safe `Weak` -//! resolver semantics and the session-bound resolution path. - -use std::sync::Arc; - -use dashmap::DashMap; - -use async_trait::async_trait; -use futures::StreamExt; -use serde::{Deserialize, Serialize}; - -use kigi_computer_hub_core::{ - CompoundResolver, ConnectionCleanupReport, ErasedTool, InnerDispatchForResolver, ResolvedTool, - SessionCleanupReport, ToolHandle, ToolRegistry, ToolSessionBindOutcome, - ToolSessionUnbindOutcome, -}; -use kigi_tool_protocol::{ - ConnectionId, RegistrationOutcome, ServerId, SessionId, ToolDefinitionMode, ToolId, - ToolRegistration, ToolServerRegistration, TransportKind, UserId, -}; -use kigi_tool_runtime::{ - SearchSnapshot, ServerSummary, Tool, ToolCallContext, ToolDispatch, ToolError, ToolStreamItem, -}; -use kigi_tool_types::ToolDescription; - -#[derive(Debug, Default, Clone, Serialize, Deserialize, schemars::JsonSchema)] -struct EchoArgs { - payload: String, -} - -#[derive(Debug)] -struct EchoTool; - -impl Tool for EchoTool { - type Args = EchoArgs; - type Output = serde_json::Value; - - fn id(&self) -> ToolId { - ToolId::new("echo").expect("tool id") - } - - fn description(&self, _ctx: &::kigi_tool_runtime::ListToolsContext) -> ToolDescription { - ToolDescription::new("echo", "Echoes its input.") - } - - async fn run( - &self, - _ctx: ToolCallContext, - args: Self::Args, - ) -> Result { - Ok(serde_json::json!({"echoed": args.payload})) - } -} - -type RegistryEntry = (ToolRegistration, Arc); - -#[derive(Debug, Default)] -struct InMemRegistry { - entries: DashMap<(SessionId, ToolId), RegistryEntry>, -} - -impl InMemRegistry { - fn install(&self, session: SessionId, tool: ToolId, handle: Arc) { - let registration = ToolRegistration { - tool_id: tool.clone(), - sessions: Some(vec![session.clone()]), - user_id: UserId::new("alice").expect("user id"), - server_id: None, - description: handle.description(&kigi_tool_runtime::ListToolsContext::default()), - input_schema: None, - capabilities: Some(handle.capabilities()), - notification_schemas: None, - transport_kind: TransportKind::Local, - if_match_generation: None, - metadata: None, - }; - self.entries.insert((session, tool), (registration, handle)); - } -} - -#[async_trait] -impl ToolRegistry for InMemRegistry { - async fn register_tool( - &self, - _connection_id: ConnectionId, - _reg: ToolRegistration, - ) -> RegistrationOutcome { - unreachable!() - } - async fn register_server( - &self, - _connection_id: ConnectionId, - _reg: ToolServerRegistration, - ) -> Vec { - unreachable!() - } - async fn unregister_tool(&self, _connection_id: &ConnectionId, _tool: &ToolId) -> bool { - unreachable!() - } - async fn unregister_server(&self, _connection_id: &ConnectionId, _server: &ServerId) -> usize { - unreachable!() - } - async fn bind_tool_session( - &self, - _connection_id: &ConnectionId, - _tool: &ToolId, - _session_id: &SessionId, - ) -> ToolSessionBindOutcome { - unreachable!() - } - async fn unbind_tool_session( - &self, - _connection_id: &ConnectionId, - _tool: &ToolId, - _session_id: &SessionId, - ) -> ToolSessionUnbindOutcome { - unreachable!() - } - async fn drop_connection(&self, _connection_id: &ConnectionId) -> ConnectionCleanupReport { - ConnectionCleanupReport::default() - } - fn find_tool(&self, session: &SessionId, tool: &ToolId) -> Option { - let (registration, handle) = self - .entries - .get(&(session.clone(), tool.clone()))? - .value() - .clone(); - Some(ResolvedTool::Local { - tool: handle, - registration, - }) - } - fn list_tools(&self, _session: &SessionId, _mode: &ToolDefinitionMode) -> Vec { - vec![] - } - fn list_servers(&self, _session: &SessionId) -> Vec { - vec![] - } - fn search(&self, _session: &SessionId, _query: &str, _limit: usize) -> SearchSnapshot { - SearchSnapshot { - results: vec![], - total_hidden_tools: 0, - is_ready: true, - } - } - async fn unregister_session(&self, _session: &SessionId) -> SessionCleanupReport { - SessionCleanupReport::default() - } - fn tool_sessions( - &self, - _connection_id: &ConnectionId, - _tool: &ToolId, - ) -> std::collections::HashSet { - std::collections::HashSet::new() - } - - fn list_servers_for_user( - &self, - _user_id: &kigi_tool_protocol::UserId, - ) -> Vec { - Vec::new() - } - - fn get_server_record( - &self, - _connection_id: &ConnectionId, - ) -> Option { - None - } -} - -fn sid(s: &str) -> SessionId { - SessionId::new(s).expect("session id") -} - -fn tid(s: &str) -> ToolId { - ToolId::new(s).expect("tool id") -} - -#[tokio::test] -async fn inner_dispatch_resolves_through_bound_session() { - let registry = Arc::new(InMemRegistry::default()); - registry.install( - sid("sess-1"), - tid("echo"), - Arc::new(ErasedTool::new(EchoTool)), - ); - let resolver = Arc::new(CompoundResolver::local_only( - registry as Arc, - )); - let inner = InnerDispatchForResolver::new(Arc::downgrade(&resolver), sid("sess-1")); - assert_eq!(inner.session_id(), &sid("sess-1")); - - let result = inner - .call_terminal( - tid("echo"), - serde_json::json!({"payload": "x"}), - ToolCallContext::default(), - ) - .await - .expect("terminal ok"); - assert_eq!(result.value, serde_json::json!({"echoed": "x"})); -} - -#[tokio::test] -async fn inner_dispatch_returns_not_found_when_tool_absent() { - let registry = Arc::new(InMemRegistry::default()); - let resolver = Arc::new(CompoundResolver::local_only( - registry as Arc, - )); - let inner = InnerDispatchForResolver::new(Arc::downgrade(&resolver), sid("sess-1")); - let mut stream = inner - .call( - tid("ghost"), - serde_json::json!(null), - ToolCallContext::default(), - ) - .await; - let item = stream.next().await.expect("terminal"); - match item { - ToolStreamItem::Terminal(Err(ref e)) - if e.kind == kigi_tool_runtime::ToolErrorKind::NotFound => - { - assert!( - e.detail.contains("ghost"), - "detail should mention tool id: {}", - e.detail - ); - } - other => panic!("expected Terminal(NotFound), got {other:?}"), - } -} - -#[tokio::test] -async fn inner_dispatch_uses_bound_session_not_context_session() { - // Even if the context were to carry a different session, the inner - // dispatch handle resolves against its construction-time session. - let registry = Arc::new(InMemRegistry::default()); - registry.install( - sid("sess-A"), - tid("echo"), - Arc::new(ErasedTool::new(EchoTool)), - ); - let resolver = Arc::new(CompoundResolver::local_only( - registry as Arc, - )); - let inner = InnerDispatchForResolver::new(Arc::downgrade(&resolver), sid("sess-B")); - let mut stream = inner - .call( - tid("echo"), - serde_json::json!({"payload": "x"}), - ToolCallContext::default(), - ) - .await; - let item = stream.next().await.expect("terminal"); - match item { - ToolStreamItem::Terminal(Err(ref e)) - if e.kind == kigi_tool_runtime::ToolErrorKind::NotFound => {} - other => panic!("session-A registration must not be visible from session-B, got {other:?}"), - } -} - -#[tokio::test] -async fn inner_dispatch_after_resolver_drop_returns_computer_hub_dropped() { - let registry = Arc::new(InMemRegistry::default()); - let resolver = Arc::new(CompoundResolver::local_only( - registry as Arc, - )); - let weak = Arc::downgrade(&resolver); - let inner = InnerDispatchForResolver::new(weak, sid("sess-1")); - drop(resolver); - let mut stream = inner - .call( - tid("echo"), - serde_json::json!(null), - ToolCallContext::default(), - ) - .await; - let item = stream.next().await.expect("terminal"); - match item { - ToolStreamItem::Terminal(Err(ref e)) - if e.kind == kigi_tool_runtime::ToolErrorKind::Custom => - { - assert!( - e.detail.contains("computer_hub_dropped") - || e.details - .as_ref() - .and_then(|d| d.get("code")) - .and_then(|c| c.as_str()) - == Some("computer_hub_dropped"), - "expected computer_hub_dropped code, got: {:?}", - e - ); - } - other => panic!("expected Terminal(Custom(computer_hub_dropped)), got {other:?}"), - } -} - -#[tokio::test] -async fn inner_dispatch_implements_object_safe_tool_dispatch() { - let registry = Arc::new(InMemRegistry::default()); - let resolver = Arc::new(CompoundResolver::local_only( - registry as Arc, - )); - let inner: Arc = Arc::new(InnerDispatchForResolver::new( - Arc::downgrade(&resolver), - sid("sess-1"), - )); - let result = inner - .call_terminal( - tid("ghost"), - serde_json::json!(null), - ToolCallContext::default(), - ) - .await; - assert!(matches!(result, Err(ref e) if e.kind == kigi_tool_runtime::ToolErrorKind::NotFound)); -} diff --git a/crates/common/kigi-computer-hub-core/tests/internal_error_detail.rs b/crates/common/kigi-computer-hub-core/tests/internal_error_detail.rs deleted file mode 100644 index e7e66f1..0000000 --- a/crates/common/kigi-computer-hub-core/tests/internal_error_detail.rs +++ /dev/null @@ -1,65 +0,0 @@ -//! Decode-side coverage for `ToolErrorWire::Internal`'s optional `detail`: -//! a populated detail must become the reconstructed `ToolError`'s message, -//! and its absence (frames from older peers) must fall back to the historic -//! constant. - -use kigi_computer_hub_core::{error_from_envelope, tool_error_from_wire}; -use kigi_tool_protocol::{JsonRpcError, RequestId, ToolErrorWire}; -use kigi_tool_runtime::ToolErrorKind; -use serde_json::json; - -#[test] -fn internal_with_detail_reconstructs_the_wire_detail() { - let err = tool_error_from_wire(ToolErrorWire::Internal { - request_id: None, - detail: Some("cross-instance tool.call timed out".to_owned()), - }); - assert_eq!(err.kind, ToolErrorKind::Custom); - assert_eq!(err.detail, "cross-instance tool.call timed out"); - // The `internal_error` code survives so callers can still classify it. - assert_eq!( - err.details - .as_ref() - .and_then(|d| d.get("code")) - .and_then(|v| v.as_str()), - Some("internal_error"), - ); -} - -#[test] -fn internal_without_detail_falls_back_to_the_historic_constant() { - let err = tool_error_from_wire(ToolErrorWire::Internal { - request_id: None, - detail: None, - }); - assert_eq!(err.kind, ToolErrorKind::Custom); - assert_eq!(err.detail, "internal router error"); -} - -#[test] -fn internal_with_request_id_keeps_both_code_and_request_id() { - let err = tool_error_from_wire(ToolErrorWire::Internal { - request_id: Some(RequestId::new("req-7").unwrap()), - detail: Some("relay publish failed".to_owned()), - }); - assert_eq!(err.detail, "relay publish failed"); - let details = err.details.expect("details present"); - assert_eq!(details["code"], json!("internal_error")); - assert_eq!(details["request_id"], json!("req-7")); -} - -#[test] -fn envelope_with_internal_data_prefers_data_detail_over_message() { - // The hub's `-32000 "internal error"` envelope keeps its constant message; - // the harness must read the cause from `error.data`, not the message. - let err = error_from_envelope(JsonRpcError { - code: -32000, - message: "internal error".to_owned(), - data: Some(json!({ - "code": "internal_error", - "detail": "cross-instance call cancelled", - })), - }); - assert_eq!(err.kind, ToolErrorKind::Custom); - assert_eq!(err.detail, "cross-instance call cancelled"); -} diff --git a/crates/common/kigi-computer-hub-core/tests/local_transport.rs b/crates/common/kigi-computer-hub-core/tests/local_transport.rs deleted file mode 100644 index 6f640d7..0000000 --- a/crates/common/kigi-computer-hub-core/tests/local_transport.rs +++ /dev/null @@ -1,369 +0,0 @@ -//! `LocalTransport` end-to-end coverage. Verifies that the transport -//! resolves through the bound resolver, drives both blocking and -//! streaming tools, and surfaces missing tools as `Terminal(NotFound)`. - -use std::sync::Arc; - -use dashmap::DashMap; - -use async_trait::async_trait; -use futures::StreamExt; -use serde::{Deserialize, Serialize}; - -use kigi_computer_hub_core::{ - CompoundResolver, ConnectionCleanupReport, ErasedTool, LocalTransport, ResolvedTool, - SessionCleanupReport, ToolHandle, ToolRegistry, ToolSessionBindOutcome, - ToolSessionUnbindOutcome, Transport, TransportKind, -}; -use kigi_tool_protocol::{ - ConnectionId, RegistrationOutcome, ServerId, SessionId, ToolDefinitionMode, ToolId, - ToolRegistration, ToolServerRegistration, TransportKind as WireTransportKind, UserId, -}; -use kigi_tool_runtime::{ - SearchSnapshot, ServerSummary, Tool, ToolCallContext, ToolError, ToolProgress, ToolStream, - ToolStreamItem, terminal_only, with_progress, -}; -use kigi_tool_types::ToolDescription; - -#[derive(Debug, Default, Clone, Serialize, Deserialize, schemars::JsonSchema)] -struct EchoArgs { - payload: String, -} - -#[derive(Debug)] -struct EchoTool; - -impl Tool for EchoTool { - type Args = EchoArgs; - type Output = serde_json::Value; - - fn id(&self) -> ToolId { - ToolId::new("echo").expect("tool id") - } - - fn description(&self, _ctx: &::kigi_tool_runtime::ListToolsContext) -> ToolDescription { - ToolDescription::new("echo", "Echoes its input.") - } - - async fn run( - &self, - _ctx: ToolCallContext, - args: Self::Args, - ) -> Result { - Ok(serde_json::json!({ "echoed": args.payload })) - } -} - -#[derive(Debug)] -struct StreamerTool; - -impl Tool for StreamerTool { - type Args = EchoArgs; - type Output = serde_json::Value; - - fn id(&self) -> ToolId { - ToolId::new("streamer").expect("tool id") - } - - fn description(&self, _ctx: &::kigi_tool_runtime::ListToolsContext) -> ToolDescription { - ToolDescription::new("streamer", "Emits three progress chunks.") - } - - async fn execute(&self, _ctx: ToolCallContext, args: Self::Args) -> ToolStream { - let chunks = futures::stream::iter(vec![ - ToolProgress::Text { - text: "tick".to_string(), - }, - ToolProgress::Text { - text: "tock".to_string(), - }, - ToolProgress::Text { - text: "boom".to_string(), - }, - ]); - with_progress(chunks, async move { - Ok(serde_json::json!({ "echoed": args.payload })) - }) - } -} - -type RegistryEntry = (ToolRegistration, Arc); - -#[derive(Debug, Default)] -struct InMemRegistry { - entries: DashMap<(SessionId, ToolId), RegistryEntry>, -} - -impl InMemRegistry { - fn install(&self, session: SessionId, tool_id: ToolId, handle: Arc) { - let registration = ToolRegistration { - tool_id: tool_id.clone(), - sessions: Some(vec![session.clone()]), - user_id: UserId::new("alice").expect("user id"), - server_id: None, - description: handle.description(&kigi_tool_runtime::ListToolsContext::default()), - input_schema: None, - capabilities: Some(handle.capabilities()), - notification_schemas: None, - transport_kind: WireTransportKind::Local, - if_match_generation: None, - metadata: None, - }; - self.entries - .insert((session, tool_id), (registration, handle)); - } -} - -#[async_trait] -impl ToolRegistry for InMemRegistry { - async fn register_tool( - &self, - _connection_id: ConnectionId, - _reg: ToolRegistration, - ) -> RegistrationOutcome { - unreachable!("transport tests pre-populate via install()") - } - async fn register_server( - &self, - _connection_id: ConnectionId, - _reg: ToolServerRegistration, - ) -> Vec { - unreachable!() - } - async fn unregister_tool(&self, _connection_id: &ConnectionId, _tool: &ToolId) -> bool { - unreachable!() - } - async fn unregister_server(&self, _connection_id: &ConnectionId, _server: &ServerId) -> usize { - unreachable!() - } - async fn bind_tool_session( - &self, - _connection_id: &ConnectionId, - _tool: &ToolId, - _session_id: &SessionId, - ) -> ToolSessionBindOutcome { - unreachable!() - } - async fn unbind_tool_session( - &self, - _connection_id: &ConnectionId, - _tool: &ToolId, - _session_id: &SessionId, - ) -> ToolSessionUnbindOutcome { - unreachable!() - } - async fn drop_connection(&self, _connection_id: &ConnectionId) -> ConnectionCleanupReport { - ConnectionCleanupReport::default() - } - fn find_tool(&self, session: &SessionId, tool: &ToolId) -> Option { - let (registration, handle) = self - .entries - .get(&(session.clone(), tool.clone()))? - .value() - .clone(); - Some(ResolvedTool::Local { - tool: handle, - registration, - }) - } - fn list_tools(&self, _session: &SessionId, _mode: &ToolDefinitionMode) -> Vec { - vec![] - } - fn list_servers(&self, _session: &SessionId) -> Vec { - vec![] - } - fn search(&self, _session: &SessionId, _query: &str, _limit: usize) -> SearchSnapshot { - SearchSnapshot { - results: vec![], - total_hidden_tools: 0, - is_ready: true, - } - } - async fn unregister_session(&self, _session: &SessionId) -> SessionCleanupReport { - SessionCleanupReport::default() - } - fn tool_sessions( - &self, - _connection_id: &ConnectionId, - _tool: &ToolId, - ) -> std::collections::HashSet { - std::collections::HashSet::new() - } - - fn list_servers_for_user( - &self, - _user_id: &kigi_tool_protocol::UserId, - ) -> Vec { - Vec::new() - } - - fn get_server_record( - &self, - _connection_id: &ConnectionId, - ) -> Option { - None - } -} - -fn sid(s: &str) -> SessionId { - SessionId::new(s).expect("session id") -} - -fn tid(s: &str) -> ToolId { - ToolId::new(s).expect("tool id") -} - -fn uid(s: &str) -> UserId { - UserId::new(s).expect("user id") -} - -async fn collect( - stream: &mut ToolStream, -) -> Vec> { - let mut items = Vec::new(); - while let Some(item) = stream.next().await { - items.push(item); - } - items -} - -#[tokio::test] -async fn dispatches_blocking_tool_to_terminal_value() { - let registry = Arc::new(InMemRegistry::default()); - registry.install( - sid("sess-1"), - tid("echo"), - Arc::new(ErasedTool::new(EchoTool)), - ); - let resolver = Arc::new(CompoundResolver::local_only( - registry as Arc, - )); - let transport = LocalTransport::new(resolver, uid("alice"), sid("sess-1")); - - let mut stream = transport - .call( - tid("echo"), - serde_json::json!({"payload": "hi"}), - ToolCallContext::default(), - ) - .await; - let items = collect(&mut stream).await; - assert_eq!(items.len(), 1); - match &items[0] { - ToolStreamItem::Terminal(Ok(typed)) => { - assert_eq!(typed.value, serde_json::json!({"echoed": "hi"})); - } - other => panic!("expected Terminal(Ok), got {other:?}"), - } -} - -#[tokio::test] -async fn dispatches_streaming_tool_with_three_progress_then_terminal() { - let registry = Arc::new(InMemRegistry::default()); - registry.install( - sid("sess-1"), - tid("streamer"), - Arc::new(ErasedTool::new(StreamerTool)), - ); - let resolver = Arc::new(CompoundResolver::local_only( - registry as Arc, - )); - let transport = LocalTransport::new(resolver, uid("alice"), sid("sess-1")); - - let mut stream = transport - .call( - tid("streamer"), - serde_json::json!({"payload": "hi"}), - ToolCallContext::default(), - ) - .await; - let items = collect(&mut stream).await; - assert_eq!(items.len(), 4); - for item in &items[..3] { - assert!(matches!(item, ToolStreamItem::Progress(_))); - } - assert!(matches!(items[3], ToolStreamItem::Terminal(Ok(_)))); -} - -#[tokio::test] -async fn missing_tool_resolves_as_terminal_not_found() { - let registry = Arc::new(InMemRegistry::default()); - let resolver = Arc::new(CompoundResolver::local_only( - registry as Arc, - )); - let transport = LocalTransport::new(resolver, uid("alice"), sid("sess-1")); - - let mut stream = transport - .call( - tid("ghost"), - serde_json::json!(null), - ToolCallContext::default(), - ) - .await; - let items = collect(&mut stream).await; - assert_eq!(items.len(), 1); - match &items[0] { - ToolStreamItem::Terminal(Err(e)) - if e.kind == kigi_tool_runtime::ToolErrorKind::NotFound => - { - assert!( - e.detail.contains("ghost"), - "detail should mention tool id: {}", - e.detail - ); - } - other => panic!("expected Terminal(Err(NotFound)), got {other:?}"), - } -} - -#[tokio::test] -async fn invalid_arguments_surface_as_terminal_error() { - let registry = Arc::new(InMemRegistry::default()); - registry.install( - sid("sess-1"), - tid("echo"), - Arc::new(ErasedTool::new(EchoTool)), - ); - let resolver = Arc::new(CompoundResolver::local_only( - registry as Arc, - )); - let transport = LocalTransport::new(resolver, uid("alice"), sid("sess-1")); - - let mut stream = transport - .call( - tid("echo"), - // Missing required `payload` field. - serde_json::json!({}), - ToolCallContext::default(), - ) - .await; - let items = collect(&mut stream).await; - assert_eq!(items.len(), 1); - match &items[0] { - ToolStreamItem::Terminal(Err(e)) - if e.kind == kigi_tool_runtime::ToolErrorKind::InvalidArguments => {} - other => panic!("expected Terminal(Err(InvalidArguments)), got {other:?}"), - } -} - -#[tokio::test] -async fn authorize_returns_bound_principal_with_invoke_scope() { - let registry = Arc::new(InMemRegistry::default()); - let resolver = Arc::new(CompoundResolver::local_only( - registry as Arc, - )); - let transport = LocalTransport::new(resolver, uid("alice"), sid("sess-1")); - let principal = transport.authorize().await.expect("authorize"); - assert_eq!(principal.user_id, uid("alice")); - assert!(principal.authorizes_session(&sid("sess-1"))); - assert!(principal.has_scope(kigi_computer_hub_core::LOCAL_INVOKE_SCOPE)); - assert_eq!(transport.kind(), TransportKind::Local); -} - -#[test] -fn unused_helpers_silenced() { - // `terminal_only` is re-exported for adapter authors; touch it here so - // a future refactor that drops the import does not silently break the - // re-export surface. - let _: ToolStream = terminal_only(Ok(serde_json::Value::Null)); -} diff --git a/crates/common/kigi-computer-hub-core/tests/remote_proxy.rs b/crates/common/kigi-computer-hub-core/tests/remote_proxy.rs deleted file mode 100644 index 28adbc8..0000000 --- a/crates/common/kigi-computer-hub-core/tests/remote_proxy.rs +++ /dev/null @@ -1,727 +0,0 @@ -//! `RemoteToolProxy` and `RemoteTransport` coverage. A channel-backed -//! mock `ConnectionClient` lets the test inspect outgoing frames and -//! drive synthetic responses + progress without any tokio I/O. - -use std::sync::{Arc, Mutex}; - -use dashmap::DashMap; - -use async_trait::async_trait; -use futures::StreamExt; -use futures::channel::{mpsc, oneshot}; -use futures::stream::BoxStream; - -use kigi_computer_hub_core::{ - ConnectionClient, RemoteToolProxy, RemoteTransport, ToolHandle, Transport, TransportKind, -}; -use kigi_tool_protocol::{ - JsonRpcError, JsonRpcId, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, JsonRpcVersion, - Method, ResponseOutcome, SessionId, ToolCallId, ToolCallParams, ToolCallProgressFrame, - ToolCallResult, ToolCapabilities, ToolErrorWire, ToolId, ToolOutputWire, UserId, -}; -use kigi_tool_runtime::{ - ContentBlock, ToolCallContext, ToolError, ToolOutput, ToolProgress, ToolStreamItem, -}; -use kigi_tool_types::ToolDescription; - -/// Programmable `ConnectionClient`. Each request gets a pre-staged -/// response; progress frames are pushed through per-call senders. -#[derive(Debug, Default)] -struct MockConnection { - /// Senders keyed by `tool_call_id`. Pulled out of the inner state so - /// per-call subscription touches a lock-free DashMap rather than the - /// shared Mutex that guards the rest of the queue + capture state. - progress_senders: DashMap>, - /// Three-Vec state guarded by one Mutex. The lock provides atomic - /// pop-from-`responses` + push-to-`captured_requests` semantics that - /// some tests rely on. - inner: Mutex, -} - -#[derive(Default)] -struct MockState { - /// FIFO queue of responses to return for each `request` call. - responses: Vec, - /// Captured outgoing requests so tests can assert on them. - captured_requests: Vec, - /// Captured one-way notifications. - captured_notifications: Vec, -} - -impl std::fmt::Debug for MockState { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("MockState") - .field("responses_len", &self.responses.len()) - .field("captured_reqs", &self.captured_requests.len()) - .field("captured_notifs", &self.captured_notifications.len()) - .finish() - } -} - -enum MockResponse { - Ok(serde_json::Value), - Err(JsonRpcError), - /// Resolves a oneshot when the request arrives so the test can - /// release progress before allowing the response. - Gated { - gate: oneshot::Receiver<()>, - body: serde_json::Value, - }, - /// Fail at the transport layer (e.g. socket dropped). - Network(String), -} - -impl MockConnection { - fn enqueue_ok(&self, body: serde_json::Value) { - self.inner - .lock() - .expect("mutex") - .responses - .push(MockResponse::Ok(body)); - } - - fn enqueue_err(&self, code: i32, message: impl Into, data: Option) { - self.inner - .lock() - .expect("mutex") - .responses - .push(MockResponse::Err(JsonRpcError { - code, - message: message.into(), - data, - })); - } - - fn enqueue_gated(&self, gate: oneshot::Receiver<()>, body: serde_json::Value) { - self.inner - .lock() - .expect("mutex") - .responses - .push(MockResponse::Gated { gate, body }); - } - - fn enqueue_network_failure(&self, message: impl Into) { - self.inner - .lock() - .expect("mutex") - .responses - .push(MockResponse::Network(message.into())); - } - - fn last_request(&self) -> Option { - self.inner - .lock() - .expect("mutex") - .captured_requests - .last() - .cloned() - } - - fn captured_request_count(&self) -> usize { - self.inner.lock().expect("mutex").captured_requests.len() - } - - fn push_progress(&self, tool_call_id: &ToolCallId, frame: ToolCallProgressFrame) { - if let Some(tx) = self.progress_senders.get(tool_call_id) { - let _ = tx.value().unbounded_send(frame); - } - } -} - -#[async_trait] -impl ConnectionClient for MockConnection { - async fn request(&self, request: JsonRpcRequest) -> Result { - let response = { - let mut guard = self.inner.lock().expect("mutex"); - guard.captured_requests.push(request.clone()); - if guard.responses.is_empty() { - return Err(ToolError::custom( - "mock_response_missing", - "no response staged", - )); - } - guard.responses.remove(0) - }; - match response { - MockResponse::Ok(body) => Ok(JsonRpcResponse { - jsonrpc: JsonRpcVersion, - id: request.id, - session_id: request.session_id, - outcome: ResponseOutcome::Result(body), - }), - MockResponse::Err(err) => Ok(JsonRpcResponse { - jsonrpc: JsonRpcVersion, - id: request.id, - session_id: request.session_id, - outcome: ResponseOutcome::Error(err), - }), - MockResponse::Gated { gate, body } => { - let _ = gate.await; - Ok(JsonRpcResponse { - jsonrpc: JsonRpcVersion, - id: request.id, - session_id: request.session_id, - outcome: ResponseOutcome::Result(body), - }) - } - MockResponse::Network(msg) => Err(ToolError::network_error(msg)), - } - } - - async fn subscribe_progress( - &self, - tool_call_id: ToolCallId, - ) -> BoxStream<'static, ToolCallProgressFrame> { - let (tx, rx) = mpsc::unbounded(); - self.progress_senders.insert(tool_call_id, tx); - rx.boxed() - } - - async fn notify(&self, notification: JsonRpcNotification) -> Result<(), ToolError> { - self.inner - .lock() - .expect("mutex") - .captured_notifications - .push(notification); - Ok(()) - } -} - -fn sid(s: &str) -> SessionId { - SessionId::new(s).expect("session id") -} - -fn tid(s: &str) -> ToolId { - ToolId::new(s).expect("tool id") -} - -fn uid(s: &str) -> UserId { - UserId::new(s).expect("user id") -} - -fn description_for(name: &str) -> ToolDescription { - ToolDescription::new(name, format!("desc for {name}")) -} - -fn ok_call_result(call_id: &ToolCallId, output: ToolOutputWire) -> serde_json::Value { - serde_json::to_value(ToolCallResult { - tool_call_id: call_id.clone(), - output, - follow_ups: vec![], - reminders: vec![], - chat_completion_output: None, - }) - .expect("serialise call result") -} - -fn ok_call_result_with_cco( - call_id: &ToolCallId, - output: ToolOutputWire, - chat_completion_output: serde_json::Value, -) -> serde_json::Value { - serde_json::to_value(ToolCallResult { - tool_call_id: call_id.clone(), - output, - follow_ups: vec![], - reminders: vec![], - chat_completion_output: Some(chat_completion_output), - }) - .expect("serialise call result") -} - -#[tokio::test] -async fn proxy_sends_well_formed_tool_call_request() { - let conn = Arc::new(MockConnection::default()); - let proxy = RemoteToolProxy::new( - tid("foo"), - sid("sess-1"), - description_for("foo"), - ToolCapabilities::default(), - conn.clone(), - ); - let call_id = ToolCallId::new_v7(); - let ctx = ToolCallContext::new(call_id.clone()); - conn.enqueue_ok(ok_call_result( - &call_id, - ToolOutputWire::Text("hello".to_string()), - )); - let mut stream = proxy.execute(ctx, serde_json::json!({"k": "v"})).await; - while stream.next().await.is_some() {} - let req = conn.last_request().expect("captured request"); - assert_eq!(req.method, Method::ToolCallRequest.as_wire_str()); - let params: ToolCallParams = serde_json::from_value(req.params).expect("decode params"); - assert_eq!(params.tool_id, tid("foo")); - assert_eq!(params.tool_call_id, call_id); - assert_eq!(params.arguments, serde_json::json!({"k": "v"})); -} - -#[tokio::test] -async fn progress_then_terminal_orders_correctly() { - let conn = Arc::new(MockConnection::default()); - let proxy = RemoteToolProxy::new( - tid("foo"), - sid("sess-1"), - description_for("foo"), - ToolCapabilities::default(), - conn.clone(), - ); - let call_id = ToolCallId::new_v7(); - let ctx = ToolCallContext::new(call_id.clone()); - let (gate_tx, gate_rx) = oneshot::channel(); - conn.enqueue_gated( - gate_rx, - ok_call_result(&call_id, ToolOutputWire::Text("done".to_string())), - ); - let mut stream = proxy.execute(ctx, serde_json::json!(null)).await; - - // Push two progress frames before the terminal is unblocked. - conn.push_progress( - &call_id, - ToolCallProgressFrame { - tool_call_id: call_id.clone(), - kind: "log".to_string(), - body: serde_json::json!({"text": "tick"}), - dropped_count: None, - }, - ); - conn.push_progress( - &call_id, - ToolCallProgressFrame { - tool_call_id: call_id.clone(), - kind: "log".to_string(), - body: serde_json::json!({"text": "tock"}), - dropped_count: None, - }, - ); - - let first = stream.next().await.expect("first item"); - let second = stream.next().await.expect("second item"); - match (&first, &second) { - (ToolStreamItem::Progress(p1), ToolStreamItem::Progress(p2)) => { - match p1 { - ToolProgress::Custom { subkind, payload } => { - assert_eq!(subkind, "log"); - assert_eq!(payload, &serde_json::json!({"text": "tick"})); - } - other => panic!("expected Custom progress, got {other:?}"), - } - match p2 { - ToolProgress::Custom { subkind, .. } => assert_eq!(subkind, "log"), - other => panic!("expected Custom progress, got {other:?}"), - } - } - other => panic!("expected two Progress items, got {other:?}"), - } - - // Release the response and consume the terminal. - let _ = gate_tx.send(()); - let terminal = stream.next().await.expect("terminal"); - match terminal { - ToolStreamItem::Terminal(Ok(typed)) => { - assert_eq!(typed.value, serde_json::json!("done")); - } - other => panic!("expected Terminal(Ok), got {other:?}"), - } - assert!(stream.next().await.is_none()); -} - -#[tokio::test] -async fn json_rpc_error_response_decodes_into_tool_error() { - let conn = Arc::new(MockConnection::default()); - let proxy = RemoteToolProxy::new( - tid("foo"), - sid("sess-1"), - description_for("foo"), - ToolCapabilities::default(), - conn.clone(), - ); - let wire = ToolErrorWire::ToolNotFound { - tool_id: tid("foo"), - }; - conn.enqueue_err( - -32011, - "tool not found", - Some(serde_json::to_value(&wire).unwrap()), - ); - let mut stream = proxy - .execute(ToolCallContext::default(), serde_json::json!(null)) - .await; - let item = stream.next().await.expect("terminal"); - match item { - ToolStreamItem::Terminal(Err(ref e)) - if e.kind == kigi_tool_runtime::ToolErrorKind::NotFound => - { - assert!( - e.detail.contains("foo"), - "detail should mention tool id: {}", - e.detail - ); - } - other => panic!("expected Terminal(NotFound), got {other:?}"), - } -} - -#[tokio::test] -async fn network_failure_surfaces_as_terminal_network_error() { - let conn = Arc::new(MockConnection::default()); - let proxy = RemoteToolProxy::new( - tid("foo"), - sid("sess-1"), - description_for("foo"), - ToolCapabilities::default(), - conn.clone(), - ); - conn.enqueue_network_failure("socket closed"); - let mut stream = proxy - .execute(ToolCallContext::default(), serde_json::json!(null)) - .await; - let item = stream.next().await.expect("terminal"); - match item { - ToolStreamItem::Terminal(Err(ref e)) - if e.kind == kigi_tool_runtime::ToolErrorKind::NetworkError => - { - assert!( - e.detail.contains("socket closed"), - "detail should mention cause: {}", - e.detail - ); - } - other => panic!("expected Terminal(NetworkError), got {other:?}"), - } -} - -#[tokio::test] -async fn mcp_output_re_serialises_into_blocks_value() { - let conn = Arc::new(MockConnection::default()); - let proxy = RemoteToolProxy::new( - tid("foo"), - sid("sess-1"), - description_for("foo"), - ToolCapabilities::default(), - conn.clone(), - ); - let call_id = ToolCallId::new_v7(); - let ctx = ToolCallContext::new(call_id.clone()); - let blocks = vec![kigi_tool_protocol::McpBlock::Text { - text: "hello".to_string(), - }]; - conn.enqueue_ok(ok_call_result(&call_id, ToolOutputWire::Mcp { blocks })); - let mut stream = proxy.execute(ctx, serde_json::json!(null)).await; - let item = stream.next().await.expect("terminal"); - match item { - ToolStreamItem::Terminal(Ok(typed)) => { - // The wire blocks round-trip through ContentBlock; assert the - // text value survives the transformation. - let blocks_value = typed - .value - .get("blocks") - .and_then(|v| v.as_array()) - .cloned() - .expect("blocks array"); - assert_eq!(blocks_value.len(), 1); - let block: ContentBlock = - serde_json::from_value(blocks_value[0].clone()).expect("decode runtime block"); - match block { - ContentBlock::Text { text } => assert_eq!(text, "hello"), - other => panic!("expected Text block, got {other:?}"), - } - } - other => panic!("expected Terminal(Ok), got {other:?}"), - } -} - -#[tokio::test] -async fn terminal_carries_chat_completion_output_from_wire() { - let conn = Arc::new(MockConnection::default()); - let proxy = RemoteToolProxy::new( - tid("bash"), - sid("sess-1"), - description_for("bash"), - ToolCapabilities::default(), - conn.clone(), - ); - let call_id = ToolCallId::new_v7(); - let ctx = ToolCallContext::new(call_id.clone()); - let cco = serde_json::json!({ - "result": { - "sender": "assistant", - "message": "", - "code_execution_result": { - "stdout": "hi\n", - "stderr": "", - "exit_code": 0, - "command_timed_out": false - } - } - }); - conn.enqueue_ok(ok_call_result_with_cco( - &call_id, - ToolOutputWire::Json(serde_json::json!({"stdout": "hi\n"})), - cco, - )); - let mut stream = proxy.execute(ctx, serde_json::json!(null)).await; - let item = stream.next().await.expect("terminal"); - match item { - ToolStreamItem::Terminal(Ok(typed)) => { - let response = typed - .chat_completion_output() - .expect("chat completion output survives the wire"); - let completion = response.result.expect("completion result present"); - let exec = completion - .code_execution_result - .expect("code execution result present"); - assert_eq!(exec.stdout, "hi\n"); - assert_eq!(exec.exit_code, 0); - assert!(!exec.command_timed_out); - } - other => panic!("expected Terminal(Ok), got {other:?}"), - } -} - -#[tokio::test] -async fn terminal_without_chat_completion_output_is_none() { - let conn = Arc::new(MockConnection::default()); - let proxy = RemoteToolProxy::new( - tid("bash"), - sid("sess-1"), - description_for("bash"), - ToolCapabilities::default(), - conn.clone(), - ); - let call_id = ToolCallId::new_v7(); - let ctx = ToolCallContext::new(call_id.clone()); - conn.enqueue_ok(ok_call_result( - &call_id, - ToolOutputWire::Json(serde_json::json!({"stdout": "hi\n"})), - )); - let mut stream = proxy.execute(ctx, serde_json::json!(null)).await; - let item = stream.next().await.expect("terminal"); - match item { - ToolStreamItem::Terminal(Ok(typed)) => { - assert!(typed.chat_completion_output().is_none()); - } - other => panic!("expected Terminal(Ok), got {other:?}"), - } -} - -#[tokio::test] -async fn malformed_inner_chat_completion_output_degrades_to_none() { - let conn = Arc::new(MockConnection::default()); - let proxy = RemoteToolProxy::new( - tid("bash"), - sid("sess-1"), - description_for("bash"), - ToolCapabilities::default(), - conn.clone(), - ); - let call_id = ToolCallId::new_v7(); - let ctx = ToolCallContext::new(call_id.clone()); - conn.enqueue_ok(ok_call_result_with_cco( - &call_id, - ToolOutputWire::Json(serde_json::json!({"stdout": "hi\n"})), - serde_json::json!({"result": "not-a-completion-object"}), - )); - let mut stream = proxy.execute(ctx, serde_json::json!(null)).await; - let item = stream.next().await.expect("terminal"); - match item { - ToolStreamItem::Terminal(Ok(typed)) => { - assert_eq!(typed.value, serde_json::json!({"stdout": "hi\n"})); - assert!(typed.chat_completion_output().is_none()); - } - other => panic!("expected Terminal(Ok) with degraded cco, got {other:?}"), - } -} - -#[tokio::test] -async fn bare_non_enveloped_success_body_passes_through() { - let conn = Arc::new(MockConnection::default()); - let proxy = RemoteToolProxy::new( - tid("bash"), - sid("sess-1"), - description_for("bash"), - ToolCapabilities::default(), - conn.clone(), - ); - let ctx = ToolCallContext::new(ToolCallId::new_v7()); - conn.enqueue_ok(serde_json::json!({"stdout": "hi\n"})); - let mut stream = proxy.execute(ctx, serde_json::json!(null)).await; - let item = stream.next().await.expect("terminal"); - match item { - ToolStreamItem::Terminal(Ok(typed)) => { - assert_eq!(typed.value, serde_json::json!({"stdout": "hi\n"})); - assert!(typed.chat_completion_output().is_none()); - } - other => panic!("expected Terminal(Ok) passthrough, got {other:?}"), - } -} - -#[tokio::test] -async fn malformed_envelope_with_tool_call_id_surfaces_decode_error() { - let conn = Arc::new(MockConnection::default()); - let proxy = RemoteToolProxy::new( - tid("bash"), - sid("sess-1"), - description_for("bash"), - ToolCapabilities::default(), - conn.clone(), - ); - let ctx = ToolCallContext::new(ToolCallId::new_v7()); - conn.enqueue_ok(serde_json::json!({"tool_call_id": "call_x", "output": 123})); - let mut stream = proxy.execute(ctx, serde_json::json!(null)).await; - let item = stream.next().await.expect("terminal"); - match item { - ToolStreamItem::Terminal(Err(ref e)) - if e.kind == kigi_tool_runtime::ToolErrorKind::Custom => - { - let code = e - .details - .as_ref() - .and_then(|d| d.get("code")) - .and_then(|c| c.as_str()); - assert_eq!(code, Some("response_decoding"), "error: {e:?}"); - } - other => panic!("expected Terminal(Err) decode failure, got {other:?}"), - } -} - -#[tokio::test] -async fn remote_transport_call_dispatches_via_connection() { - let conn = Arc::new(MockConnection::default()); - let transport = RemoteTransport::new(conn.clone(), sid("sess-1"), uid("alice")); - let call_id = ToolCallId::new_v7(); - let ctx = ToolCallContext::new(call_id.clone()); - conn.enqueue_ok(ok_call_result(&call_id, ToolOutputWire::Text("hi".into()))); - let mut stream = transport - .call(tid("foo"), serde_json::json!({"k": "v"}), ctx) - .await; - let _ = stream.next().await; - assert_eq!(conn.captured_request_count(), 1); - assert_eq!(transport.kind(), TransportKind::Remote); -} - -#[tokio::test] -async fn remote_transport_authorize_returns_bound_principal() { - let conn = Arc::new(MockConnection::default()); - let transport = RemoteTransport::new(conn, sid("sess-1"), uid("alice")); - let principal = transport.authorize().await.expect("authorize"); - assert_eq!(principal.user_id, uid("alice")); - assert!(principal.authorizes_session(&sid("sess-1"))); -} - -#[tokio::test] -async fn proxy_subscribe_happens_before_request_send() { - // Locks in BOTH halves of the subscribe-before-send contract: - // 1. the subscription IS active by the time `execute` returns; - // 2. the request HAS NOT been sent yet at that point. - // A future refactor that eagerly sent the request inside - // `execute` would still satisfy (1) but would break (2). - let conn = Arc::new(MockConnection::default()); - let proxy = RemoteToolProxy::new( - tid("foo"), - sid("sess-1"), - description_for("foo"), - ToolCapabilities::default(), - conn.clone(), - ); - let call_id = ToolCallId::new_v7(); - let ctx = ToolCallContext::new(call_id.clone()); - conn.enqueue_ok(ok_call_result(&call_id, ToolOutputWire::Text("ok".into()))); - let mut stream = proxy.execute(ctx, serde_json::json!(null)).await; - { - // The DashMap subscription read and the captured-requests check - // are individually atomic. Single-threaded `#[tokio::test]` - // execution means no other task can mutate either between the - // two checks, so the pair is observationally simultaneous. - assert!( - conn.progress_senders.contains_key(&call_id), - "subscription must be active before request send" - ); - let guard = conn.inner.lock().expect("mutex"); - assert!( - guard.captured_requests.is_empty(), - "request must not be sent before stream is polled" - ); - } - // Polling the stream is what actually drives the request future, - // so the captured-requests vec only fills in once we start consuming. - while stream.next().await.is_some() {} - { - let guard = conn.inner.lock().expect("mutex"); - assert_eq!( - guard.captured_requests.len(), - 1, - "request must have been sent during stream polling" - ); - } -} - -#[tokio::test] -async fn notify_round_trips_through_connection_client() { - let conn = Arc::new(MockConnection::default()); - let notification = JsonRpcNotification { - jsonrpc: JsonRpcVersion, - session_id: Some(sid("sess-1")), - seq: None, - method: Method::Hook.as_wire_str().to_string(), - params: serde_json::json!({ - "session_id": "sess-1", - "tool_id": "foo", - "call_id": "call-1", - "event": { "type": "Cancel" } - }), - }; - let trait_handle: &dyn ConnectionClient = conn.as_ref(); - trait_handle - .notify(notification.clone()) - .await - .expect("notify succeeds"); - let guard = conn.inner.lock().expect("mutex"); - assert_eq!(guard.captured_notifications.len(), 1); - let captured = &guard.captured_notifications[0]; - assert_eq!(captured.method, Method::Hook.as_wire_str()); - assert_eq!(captured.session_id, Some(sid("sess-1"))); - assert_eq!( - captured.params.get("event").and_then(|v| v.get("type")), - Some(&serde_json::Value::String("Cancel".to_string())) - ); - assert_eq!(captured, ¬ification); -} - -#[tokio::test] -async fn json_rpc_id_is_unique_per_call() { - let conn = Arc::new(MockConnection::default()); - let transport = RemoteTransport::new(conn.clone(), sid("sess-1"), uid("alice")); - let call_a = ToolCallId::new_v7(); - let call_b = ToolCallId::new_v7(); - conn.enqueue_ok(ok_call_result(&call_a, ToolOutputWire::Text("a".into()))); - conn.enqueue_ok(ok_call_result(&call_b, ToolOutputWire::Text("b".into()))); - - let mut s1 = transport - .call( - tid("foo"), - serde_json::json!(null), - ToolCallContext::new(call_a.clone()), - ) - .await; - while s1.next().await.is_some() {} - let mut s2 = transport - .call( - tid("foo"), - serde_json::json!(null), - ToolCallContext::new(call_b.clone()), - ) - .await; - while s2.next().await.is_some() {} - - let guard = conn.inner.lock().expect("mutex"); - assert_eq!(guard.captured_requests.len(), 2); - let id_a = match &guard.captured_requests[0].id { - JsonRpcId::String(s) => s.clone(), - JsonRpcId::Number(n) => n.to_string(), - }; - let id_b = match &guard.captured_requests[1].id { - JsonRpcId::String(s) => s.clone(), - JsonRpcId::Number(n) => n.to_string(), - }; - assert_ne!(id_a, id_b, "envelope ids must differ across calls"); -} diff --git a/crates/common/kigi-computer-hub-core/tests/tool_registry.rs b/crates/common/kigi-computer-hub-core/tests/tool_registry.rs deleted file mode 100644 index 16410b2..0000000 --- a/crates/common/kigi-computer-hub-core/tests/tool_registry.rs +++ /dev/null @@ -1,607 +0,0 @@ -//! `ToolRegistry` trait coverage via a per-test mock backed by `DashMap` -//! — lock-free per-key concurrent access mirrors the production -//! direction even at the test layer. The mock implements the -//! connection-scoped `ToolRegistry` trait surface. - -use std::collections::{HashMap, HashSet}; -use std::sync::Arc; - -use async_trait::async_trait; -use dashmap::DashMap; - -use kigi_computer_hub_core::{ - ConnectionCleanupReport, ErasedTool, ResolvedTool, SessionCleanupReport, ToolHandle, - ToolRegistry, ToolSessionBindOutcome, ToolSessionUnbindOutcome, resolver::CompoundResolver, -}; -use kigi_tool_protocol::{ - ConnectionId, RegistrationOutcome, ServerId, SessionId, ToolDefinitionMode, ToolId, - ToolRegistration, ToolServerRegistration, TransportKind, UserId, -}; -use kigi_tool_runtime::{SearchSnapshot, ServerSummary, Tool, ToolCallContext, ToolError}; -use kigi_tool_types::ToolDescription; -use serde::{Deserialize, Serialize}; - -#[derive(Debug, Default, Clone, Serialize, Deserialize, schemars::JsonSchema)] -struct EmptyArgs {} - -#[derive(Debug)] -struct StubTool { - id: ToolId, -} - -impl Tool for StubTool { - type Args = EmptyArgs; - type Output = serde_json::Value; - - fn id(&self) -> ToolId { - self.id.clone() - } - - fn description(&self, _ctx: &::kigi_tool_runtime::ListToolsContext) -> ToolDescription { - ToolDescription::new(self.id.as_str(), format!("stub for {}", self.id)) - } - - async fn run( - &self, - _ctx: ToolCallContext, - _args: Self::Args, - ) -> Result { - unreachable!("registry tests do not exercise execution") - } -} - -#[derive(Debug, Clone)] -struct MockEntry { - registration: ToolRegistration, - sessions: HashSet, -} - -/// Mock registry. Last-write-wins on duplicate registrations within a -/// `(connection, tool_id)` slot — pinned here so the trait contract has -/// a clear test fixture. -#[derive(Debug, Default)] -struct MockRegistry { - entries: DashMap<(ConnectionId, ToolId), MockEntry>, - by_session: DashMap<(SessionId, ToolId), ConnectionId>, - handles: DashMap>, -} - -impl MockRegistry { - fn install_handle(&self, tool: Arc) { - self.handles.insert(tool.id(), tool); - } -} - -fn build_registration(tool: &ToolId, sessions: &[SessionId]) -> ToolRegistration { - ToolRegistration { - tool_id: tool.clone(), - sessions: Some(sessions.to_vec()), - user_id: UserId::new("alice").expect("valid user id"), - server_id: None, - description: ToolDescription::new(tool.as_str(), format!("desc for {tool}")), - input_schema: None, - capabilities: None, - notification_schemas: None, - transport_kind: TransportKind::Local, - if_match_generation: None, - metadata: None, - } -} - -#[async_trait] -impl ToolRegistry for MockRegistry { - async fn register_tool( - &self, - connection_id: ConnectionId, - reg: ToolRegistration, - ) -> RegistrationOutcome { - let key = (connection_id.clone(), reg.tool_id.clone()); - let sessions: HashSet = reg - .sessions - .as_ref() - .map(|v| v.iter().cloned().collect()) - .unwrap_or_default(); - let updated = self - .entries - .insert( - key, - MockEntry { - registration: reg.clone(), - sessions: sessions.clone(), - }, - ) - .is_some(); - for session in &sessions { - self.by_session.insert( - (session.clone(), reg.tool_id.clone()), - connection_id.clone(), - ); - } - if updated { - RegistrationOutcome::Updated { - tool_id: reg.tool_id, - generation: 1, - } - } else { - RegistrationOutcome::Registered { - tool_id: reg.tool_id, - generation: 0, - } - } - } - - async fn register_server( - &self, - connection_id: ConnectionId, - reg: ToolServerRegistration, - ) -> Vec { - let mut outcomes = Vec::with_capacity(reg.tools.len()); - for tool in reg.tools { - let tool_id = tool - .derive_tool_id() - .expect("test descriptions have valid tool ids"); - let registration = ToolRegistration { - tool_id: tool_id.clone(), - sessions: reg.sessions.clone(), - user_id: reg.user_id.clone(), - server_id: Some(reg.server_id.clone()), - description: tool.description, - input_schema: tool.input_schema, - capabilities: tool.capabilities, - notification_schemas: tool.notification_schemas, - transport_kind: TransportKind::Remote, - if_match_generation: None, - metadata: None, - }; - outcomes.push( - self.register_tool(connection_id.clone(), registration) - .await, - ); - } - outcomes - } - - async fn unregister_tool(&self, connection_id: &ConnectionId, tool: &ToolId) -> bool { - let Some((_, removed)) = self.entries.remove(&(connection_id.clone(), tool.clone())) else { - return false; - }; - for session in &removed.sessions { - self.by_session - .remove_if(&(session.clone(), tool.clone()), |_, owner| { - owner == connection_id - }); - } - true - } - - async fn unregister_server(&self, connection_id: &ConnectionId, server: &ServerId) -> usize { - let to_remove: Vec = self - .entries - .iter() - .filter(|r| { - r.key().0 == *connection_id - && r.value().registration.server_id.as_ref() == Some(server) - }) - .map(|r| r.key().1.clone()) - .collect(); - let mut removed = 0usize; - for tool in to_remove { - if self.unregister_tool(connection_id, &tool).await { - removed += 1; - } - } - removed - } - - async fn bind_tool_session( - &self, - connection_id: &ConnectionId, - tool: &ToolId, - session_id: &SessionId, - ) -> ToolSessionBindOutcome { - let key = (connection_id.clone(), tool.clone()); - let Some(mut entry) = self.entries.get_mut(&key) else { - return ToolSessionBindOutcome::UnknownTool; - }; - if !entry.value_mut().sessions.insert(session_id.clone()) { - return ToolSessionBindOutcome::AlreadyBound; - } - self.by_session - .insert((session_id.clone(), tool.clone()), connection_id.clone()); - ToolSessionBindOutcome::Bound - } - - async fn unbind_tool_session( - &self, - connection_id: &ConnectionId, - tool: &ToolId, - session_id: &SessionId, - ) -> ToolSessionUnbindOutcome { - let key = (connection_id.clone(), tool.clone()); - let Some(mut entry) = self.entries.get_mut(&key) else { - return ToolSessionUnbindOutcome::UnknownTool; - }; - if !entry.value_mut().sessions.remove(session_id) { - return ToolSessionUnbindOutcome::NotBound; - } - self.by_session - .remove_if(&(session_id.clone(), tool.clone()), |_, owner| { - owner == connection_id - }); - ToolSessionUnbindOutcome::Unbound - } - - async fn drop_connection(&self, connection_id: &ConnectionId) -> ConnectionCleanupReport { - let to_remove: Vec = self - .entries - .iter() - .filter(|r| r.key().0 == *connection_id) - .map(|r| r.key().1.clone()) - .collect(); - let mut report = ConnectionCleanupReport::default(); - for tool in to_remove { - if let Some((_, removed)) = self.entries.remove(&(connection_id.clone(), tool.clone())) - { - report.tools_dropped += 1; - for session in removed.sessions { - if self - .by_session - .remove_if(&(session, tool.clone()), |_, owner| owner == connection_id) - .is_some() - { - report.session_bindings_cleared += 1; - } - } - } - } - report - } - - fn find_tool(&self, session: &SessionId, tool: &ToolId) -> Option { - let owner = self - .by_session - .get(&(session.clone(), tool.clone()))? - .value() - .clone(); - let entry = self.entries.get(&(owner, tool.clone()))?; - let registration = entry.value().registration.clone(); - let handle = self.handles.get(tool)?.value().clone(); - match registration.transport_kind { - TransportKind::Local => Some(ResolvedTool::Local { - tool: handle, - registration, - }), - TransportKind::Remote => Some(ResolvedTool::Remote { - proxy: handle, - registration, - }), - } - } - - fn list_tools(&self, session: &SessionId, _mode: &ToolDefinitionMode) -> Vec { - self.by_session - .iter() - .filter(|r| r.key().0 == *session) - .filter_map(|r| { - let owner = r.value().clone(); - let tool_id = r.key().1.clone(); - self.entries - .get(&(owner, tool_id)) - .map(|e| e.value().registration.description.clone()) - }) - .collect() - } - - fn list_servers(&self, session: &SessionId) -> Vec { - let mut by_server: HashMap> = HashMap::new(); - for r in self.by_session.iter().filter(|r| r.key().0 == *session) { - let owner = r.value().clone(); - let tool_id = r.key().1.clone(); - if let Some(entry) = self.entries.get(&(owner, tool_id)) { - let reg = &entry.value().registration; - let server = reg - .server_id - .clone() - .unwrap_or_else(|| ServerId::synthesize_for_tool(r.value(), ®.tool_id)); - by_server - .entry(server) - .or_default() - .push(reg.tool_id.as_str().to_string()); - } - } - by_server - .into_iter() - .map(|(server, mut names)| { - names.sort(); - ServerSummary { - name: server.into_inner(), - description: None, - tool_names: names, - } - }) - .collect() - } - - fn search(&self, session: &SessionId, query: &str, limit: usize) -> SearchSnapshot { - let matches: Vec<_> = self - .by_session - .iter() - .filter(|r| r.key().0 == *session) - .filter_map(|r| { - let owner = r.value().clone(); - let tool_id = r.key().1.clone(); - let entry = self.entries.get(&(owner, tool_id))?; - let reg = &entry.value().registration; - if reg.tool_id.as_str().contains(query) { - Some(kigi_tool_runtime::ToolSearchResult { - tool_name: reg.tool_id.as_str().to_string(), - server_name: reg - .server_id - .as_ref() - .map(|s| s.as_str().to_string()) - .unwrap_or_default(), - description: reg.description.description.clone(), - score: 1.0, - parameters: vec![], - input_schema: serde_json::Value::Null, - }) - } else { - None - } - }) - .take(limit) - .collect(); - SearchSnapshot { - results: matches, - total_hidden_tools: 0, - is_ready: true, - } - } - - async fn unregister_session(&self, session: &SessionId) -> SessionCleanupReport { - let pairs: Vec<(ToolId, ConnectionId)> = self - .by_session - .iter() - .filter(|r| r.key().0 == *session) - .map(|r| (r.key().1.clone(), r.value().clone())) - .collect(); - let mut report = SessionCleanupReport::default(); - for (tool_id, owner) in pairs { - self.by_session - .remove_if(&(session.clone(), tool_id.clone()), |_, value| { - value == &owner - }); - if let Some(mut entry) = self.entries.get_mut(&(owner, tool_id)) { - entry.value_mut().sessions.remove(session); - report.tools_touched += 1; - if entry.value().sessions.is_empty() { - report.tools_left_orphaned += 1; - } - } - } - report - } - - fn tool_sessions(&self, connection_id: &ConnectionId, tool: &ToolId) -> HashSet { - self.entries - .get(&(connection_id.clone(), tool.clone())) - .map(|r| r.value().sessions.clone()) - .unwrap_or_default() - } - - fn list_servers_for_user( - &self, - _user_id: &kigi_tool_protocol::UserId, - ) -> Vec { - Vec::new() - } - - fn get_server_record( - &self, - _connection_id: &ConnectionId, - ) -> Option { - None - } -} - -fn sid(s: &str) -> SessionId { - SessionId::new(s).expect("valid session id") -} - -fn tid(s: &str) -> ToolId { - ToolId::new(s).expect("valid tool id") -} - -fn cid(s: &str) -> ConnectionId { - ConnectionId::new(s).expect("valid connection id") -} - -#[tokio::test] -async fn register_then_find_returns_local_resolution() { - let reg = MockRegistry::default(); - reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") }))); - let outcome = reg - .register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")])) - .await; - assert!(matches!(outcome, RegistrationOutcome::Registered { .. })); - let resolved = reg - .find_tool(&sid("sess-1"), &tid("foo")) - .expect("registration found"); - match resolved { - ResolvedTool::Local { registration, .. } => { - assert_eq!(registration.tool_id, tid("foo")); - assert!( - registration - .sessions - .as_ref() - .is_some_and(|s| s.contains(&sid("sess-1"))) - ); - } - other => panic!("expected Local, got {other:?}"), - } -} - -#[tokio::test] -async fn find_in_other_session_returns_none() { - let reg = MockRegistry::default(); - reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") }))); - reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")])) - .await; - assert!(reg.find_tool(&sid("sess-2"), &tid("foo")).is_none()); -} - -#[tokio::test] -async fn duplicate_registration_yields_updated_outcome() { - let reg = MockRegistry::default(); - reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") }))); - let first = reg - .register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")])) - .await; - let second = reg - .register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")])) - .await; - assert!(matches!(first, RegistrationOutcome::Registered { .. })); - assert!(matches!(second, RegistrationOutcome::Updated { .. })); -} - -#[tokio::test] -async fn unregister_tool_removes_only_that_entry() { - let reg = MockRegistry::default(); - reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") }))); - reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("bar") }))); - reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")])) - .await; - reg.register_tool(cid("c1"), build_registration(&tid("bar"), &[sid("sess-1")])) - .await; - assert!(reg.unregister_tool(&cid("c1"), &tid("foo")).await); - assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_none()); - assert!(reg.find_tool(&sid("sess-1"), &tid("bar")).is_some()); -} - -#[tokio::test] -async fn unregister_session_drops_session_binding_and_leaves_orphan_count() { - let reg = MockRegistry::default(); - reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") }))); - reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("bar") }))); - reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")])) - .await; - reg.register_tool( - cid("c1"), - build_registration(&tid("bar"), &[sid("sess-1"), sid("sess-2")]), - ) - .await; - let report = reg.unregister_session(&sid("sess-1")).await; - assert_eq!(report.tools_touched, 2); - // `foo` had only sess-1 → orphaned. `bar` had sess-2 left → not orphaned. - assert_eq!(report.tools_left_orphaned, 1); - assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_none()); - assert!(reg.find_tool(&sid("sess-1"), &tid("bar")).is_none()); - assert!(reg.find_tool(&sid("sess-2"), &tid("bar")).is_some()); -} - -#[tokio::test] -async fn list_tools_filters_by_session() { - let reg = MockRegistry::default(); - reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") }))); - reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("bar") }))); - reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")])) - .await; - reg.register_tool(cid("c1"), build_registration(&tid("bar"), &[sid("sess-2")])) - .await; - let s1 = reg.list_tools(&sid("sess-1"), &ToolDefinitionMode::Full); - let s2 = reg.list_tools(&sid("sess-2"), &ToolDefinitionMode::Full); - assert_eq!(s1.len(), 1); - assert_eq!(s1[0].name, "foo"); - assert_eq!(s2.len(), 1); - assert_eq!(s2[0].name, "bar"); -} - -#[tokio::test] -async fn list_servers_groups_by_owning_server() { - let reg = MockRegistry::default(); - reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") }))); - reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")])) - .await; - let summaries = reg.list_servers(&sid("sess-1")); - assert_eq!(summaries.len(), 1); - assert_eq!(summaries[0].tool_count(), 1); - assert_eq!(summaries[0].tool_names[0], "foo"); -} - -#[tokio::test] -async fn search_returns_substring_matches() { - let reg = MockRegistry::default(); - reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") }))); - reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foobar") }))); - reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")])) - .await; - reg.register_tool( - cid("c1"), - build_registration(&tid("foobar"), &[sid("sess-1")]), - ) - .await; - let snap = reg.search(&sid("sess-1"), "foo", 10); - assert_eq!(snap.results.len(), 2); - assert!(snap.is_ready); - assert_eq!(snap.total_hidden_tools, 0); -} - -#[tokio::test] -async fn registry_drives_compound_resolver() { - let registry = Arc::new(MockRegistry::default()); - registry.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") }))); - registry - .register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")])) - .await; - let resolver = CompoundResolver::local_only(registry as Arc); - assert!(resolver.resolve(&sid("sess-1"), &tid("foo")).is_some()); - assert!(resolver.resolve(&sid("sess-1"), &tid("missing")).is_none()); -} - -#[tokio::test] -async fn bind_and_unbind_tool_session_round_trips_visibility() { - let reg = MockRegistry::default(); - reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") }))); - reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[])) - .await; - // Empty sessions: tool is registered but unreachable. - assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_none()); - let outcome = reg - .bind_tool_session(&cid("c1"), &tid("foo"), &sid("sess-1")) - .await; - assert_eq!(outcome, ToolSessionBindOutcome::Bound); - assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_some()); - let again = reg - .bind_tool_session(&cid("c1"), &tid("foo"), &sid("sess-1")) - .await; - assert_eq!(again, ToolSessionBindOutcome::AlreadyBound); - let unbind = reg - .unbind_tool_session(&cid("c1"), &tid("foo"), &sid("sess-1")) - .await; - assert_eq!(unbind, ToolSessionUnbindOutcome::Unbound); - assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_none()); - let unknown = reg - .bind_tool_session(&cid("c1"), &tid("missing"), &sid("sess-1")) - .await; - assert_eq!(unknown, ToolSessionBindOutcome::UnknownTool); -} - -#[tokio::test] -async fn drop_connection_releases_every_owned_tool() { - let reg = MockRegistry::default(); - reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") }))); - reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("bar") }))); - reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")])) - .await; - reg.register_tool( - cid("c1"), - build_registration(&tid("bar"), &[sid("sess-1"), sid("sess-2")]), - ) - .await; - let report = reg.drop_connection(&cid("c1")).await; - assert_eq!(report.tools_dropped, 2); - assert_eq!(report.session_bindings_cleared, 3); - assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_none()); - assert!(reg.find_tool(&sid("sess-2"), &tid("bar")).is_none()); - assert!(reg.tool_sessions(&cid("c1"), &tid("foo")).is_empty()); -} diff --git a/crates/common/kigi-computer-hub-core/tests/transport_trait.rs b/crates/common/kigi-computer-hub-core/tests/transport_trait.rs deleted file mode 100644 index 0370f16..0000000 --- a/crates/common/kigi-computer-hub-core/tests/transport_trait.rs +++ /dev/null @@ -1,137 +0,0 @@ -//! Behavioural coverage for the `Transport` trait, `Principal` builder, -//! and `TransportKind` re-export. - -use async_trait::async_trait; -use serde_json::{Value, json}; - -use kigi_computer_hub_core::{Principal, Transport, TransportKind}; -use kigi_tool_protocol::{SessionId, ToolId, UserId}; -use kigi_tool_runtime::{ - ToolCallContext, ToolError, ToolStream, ToolStreamItem, TypedToolOutput, terminal_only, -}; - -fn uid(s: &str) -> UserId { - UserId::new(s).expect("test user id") -} - -fn sid(s: &str) -> SessionId { - SessionId::new(s).expect("test session id") -} - -fn tid(s: &str) -> ToolId { - ToolId::new(s).expect("test tool id") -} - -#[derive(Debug)] -struct EchoTransport { - kind: TransportKind, - user: UserId, - session: SessionId, -} - -#[async_trait] -impl Transport for EchoTransport { - fn kind(&self) -> TransportKind { - self.kind - } - - async fn authorize(&self) -> Result { - Ok(Principal::new(self.user.clone()) - .with_session(self.session.clone()) - .with_scope("tool.invoke")) - } - - async fn call( - &self, - tool_id: ToolId, - args: Value, - _ctx: ToolCallContext, - ) -> ToolStream { - terminal_only(Ok(TypedToolOutput::from_value(tool_id, args))) - } -} - -#[tokio::test] -async fn boxed_transport_compiles_and_dispatches() { - let boxed: Box = Box::new(EchoTransport { - kind: TransportKind::Local, - user: uid("alice"), - session: sid("sess-1"), - }); - let mut stream = boxed - .call(tid("echo"), json!({"k": "v"}), ToolCallContext::default()) - .await; - let item = futures::StreamExt::next(&mut stream) - .await - .expect("at least one item"); - match item { - ToolStreamItem::Terminal(Ok(typed)) => assert_eq!(typed.value, json!({"k": "v"})), - other => panic!("expected Terminal(Ok), got {other:?}"), - } -} - -#[tokio::test] -async fn kind_distinguishes_local_and_remote() { - let local = EchoTransport { - kind: TransportKind::Local, - user: uid("alice"), - session: sid("sess-1"), - }; - let remote = EchoTransport { - kind: TransportKind::Remote, - user: uid("alice"), - session: sid("sess-1"), - }; - assert_eq!(local.kind(), TransportKind::Local); - assert_eq!(remote.kind(), TransportKind::Remote); - assert_ne!(local.kind(), remote.kind()); -} - -#[tokio::test] -async fn authorize_returns_bound_principal() { - let t = EchoTransport { - kind: TransportKind::Local, - user: uid("alice"), - session: sid("sess-1"), - }; - let principal = t.authorize().await.expect("authorize succeeds"); - assert_eq!(principal.user_id, uid("alice")); - assert!(principal.authorizes_session(&sid("sess-1"))); - assert!(!principal.authorizes_session(&sid("sess-other"))); - assert!(principal.has_scope("tool.invoke")); - assert!(!principal.has_scope("admin")); -} - -#[test] -fn principal_builder_chains_in_order() { - let principal = Principal::new(uid("alice")) - .with_session(sid("sess-a")) - .with_session(sid("sess-b")) - .with_scope("tool.invoke") - .with_scope("tool.search") - .with_audience("dispatcher.example"); - assert_eq!(principal.session_ids, vec![sid("sess-a"), sid("sess-b")]); - assert_eq!(principal.scopes, vec!["tool.invoke", "tool.search"]); - assert_eq!(principal.audiences, vec!["dispatcher.example"]); -} - -#[test] -fn principal_supports_multi_session_tokens() { - let p = Principal::new(uid("alice")) - .with_session(sid("sess-1")) - .with_session(sid("sess-2")); - assert!(p.authorizes_session(&sid("sess-1"))); - assert!(p.authorizes_session(&sid("sess-2"))); - assert!(!p.authorizes_session(&sid("sess-3"))); - assert_eq!(p.session_ids.len(), 2); -} - -#[test] -fn principal_default_state_is_empty() { - let p = Principal::new(uid("alice")); - assert!(p.session_ids.is_empty()); - assert!(p.scopes.is_empty()); - assert!(p.audiences.is_empty()); - assert!(!p.has_scope("anything")); - assert!(!p.authorizes_session(&sid("sess"))); -} diff --git a/crates/common/kigi-computer-hub-core/tests/workspace_unavailable.rs b/crates/common/kigi-computer-hub-core/tests/workspace_unavailable.rs deleted file mode 100644 index 1d1f7fb..0000000 --- a/crates/common/kigi-computer-hub-core/tests/workspace_unavailable.rs +++ /dev/null @@ -1,203 +0,0 @@ -//! `is_workspace_unavailable` recognizer coverage, pinned against the real -//! wire decode path (`error_from_envelope` / `tool_error_from_wire`). - -use kigi_computer_hub_core::{error_from_envelope, is_workspace_unavailable, tool_error_from_wire}; -use kigi_tool_protocol::{ - JsonRpcError, ToolErrorWire, WORKSPACE_UNAVAILABLE_SUBCODE, WorkspaceGonePhase, - WorkspaceGoneReason, WorkspaceUnavailableDetails, workspace_unavailable_wire, -}; -use kigi_tool_runtime::{ToolError, ToolErrorKind}; -use serde_json::json; - -const REASONS: [WorkspaceGoneReason; 5] = [ - WorkspaceGoneReason::IdleTimeout, - WorkspaceGoneReason::Disconnect, - WorkspaceGoneReason::Shutdown, - WorkspaceGoneReason::NotBound, - WorkspaceGoneReason::InstanceGone, -]; -const PHASES: [WorkspaceGonePhase; 2] = [ - WorkspaceGonePhase::InFlightCancelled, - WorkspaceGonePhase::RouteMissing, -]; - -fn envelope_for(wire: &ToolErrorWire) -> JsonRpcError { - JsonRpcError { - // -32005 is the best-effort numeric companion (`tool_server_gone`); - // recognition keys on `data.details.code`, not the numeric. - code: -32005, - message: "workspace server gone".to_owned(), - data: Some(serde_json::to_value(wire).unwrap()), - } -} - -#[test] -fn round_trip_through_envelope_is_recognized_for_every_reason_and_phase() { - for reason in REASONS { - for phase in PHASES { - let wire = workspace_unavailable_wire(reason, phase); - let err = error_from_envelope(envelope_for(&wire)); - - assert!( - is_workspace_unavailable(&err), - "should recognize {reason:?}/{phase:?}", - ); - assert_eq!(err.kind, ToolErrorKind::Custom); - - // The full structured payload survives into `ToolError::details`, - // so a caller can branch on code/reason/phase/retryable. - let details: WorkspaceUnavailableDetails = - serde_json::from_value(err.details.expect("details survive")).unwrap(); - assert_eq!( - details, - WorkspaceUnavailableDetails { - code: WORKSPACE_UNAVAILABLE_SUBCODE.to_owned(), - reason, - phase, - retryable: true, - }, - ); - } - } -} - -#[test] -fn tool_error_from_wire_directly_is_recognized() { - let wire = workspace_unavailable_wire( - WorkspaceGoneReason::Disconnect, - WorkspaceGonePhase::RouteMissing, - ); - let err = tool_error_from_wire(wire); - assert!(is_workspace_unavailable(&err)); - let details = err.details.expect("details survive"); - assert_eq!(details["code"], json!(WORKSPACE_UNAVAILABLE_SUBCODE)); - assert_eq!(details["reason"], json!("disconnect")); - assert_eq!(details["phase"], json!("route_missing")); - assert_eq!(details["retryable"], json!(true)); -} - -#[test] -fn wire_to_tool_error_to_wire_preserves_outer_subcode() { - // Keying the identity on details.code lets From for ToolErrorWire - // rebuild the outer subcode on re-serialization. - let original = workspace_unavailable_wire( - WorkspaceGoneReason::IdleTimeout, - WorkspaceGonePhase::InFlightCancelled, - ); - let tool_error = tool_error_from_wire(original); - let back: ToolErrorWire = tool_error.into(); - let ToolErrorWire::Custom { subcode, .. } = back else { - panic!("expected Custom variant"); - }; - assert_eq!(subcode, WORKSPACE_UNAVAILABLE_SUBCODE); -} - -#[test] -fn recognized_with_unknown_reason_and_phase() { - // Recognition is decoupled from the typed reason/phase enums: a newer hub - // emitting unknown values is still recognized (it keys only on `code`). - let wire = ToolErrorWire::Custom { - subcode: WORKSPACE_UNAVAILABLE_SUBCODE.to_owned(), - message: "from a newer hub".to_owned(), - details: Some(json!({ - "code": WORKSPACE_UNAVAILABLE_SUBCODE, - "reason": "brand_new_reason", - "phase": "brand_new_phase", - "retryable": true, - })), - }; - let err = error_from_envelope(envelope_for(&wire)); - assert!(is_workspace_unavailable(&err)); - // End-to-end decode → typed-parse → `Unknown`, the path consumers read by. - let details: WorkspaceUnavailableDetails = - serde_json::from_value(err.details.expect("details survive")).unwrap(); - assert_eq!(details.reason, WorkspaceGoneReason::Unknown); - assert_eq!(details.phase, WorkspaceGonePhase::Unknown); -} - -#[test] -fn decoded_custom_with_none_details_is_recognized_via_canonical_code() { - // Wire `details: None` decodes through `ToolError::custom`, which repopulates - // `details = {"code": subcode}`, so it IS recognized — contrast the hand-built - // no-details case in `custom_error_without_any_details_is_not_recognized`. - let wire = ToolErrorWire::Custom { - subcode: WORKSPACE_UNAVAILABLE_SUBCODE.to_owned(), - message: "no structured details".to_owned(), - details: None, - }; - let err = error_from_envelope(envelope_for(&wire)); - assert_eq!(err.kind, ToolErrorKind::Custom); - assert!(is_workspace_unavailable(&err)); -} - -#[test] -fn decoded_custom_without_code_key_is_not_recognized() { - // The central correctness property: recognition keys on the surviving - // `details.code`, NOT the outer `Custom.subcode`. Here the outer subcode - // matches, but `with_details` overwrote the auto-populated `code`, so the - // decoded error must NOT be recognized. - let wire = ToolErrorWire::Custom { - subcode: WORKSPACE_UNAVAILABLE_SUBCODE.to_owned(), - message: "details lack code".to_owned(), - details: Some(json!({ "reason": "disconnect" })), - }; - let err = error_from_envelope(envelope_for(&wire)); - assert_eq!(err.kind, ToolErrorKind::Custom); - assert!(!is_workspace_unavailable(&err)); -} - -#[test] -fn different_custom_code_is_not_recognized() { - let wire = ToolErrorWire::Custom { - subcode: "some_other_error".to_owned(), - message: "nope".to_owned(), - details: Some(json!({ "code": "some_other_error" })), - }; - let err = error_from_envelope(envelope_for(&wire)); - assert_eq!(err.kind, ToolErrorKind::Custom); - assert!(!is_workspace_unavailable(&err)); -} - -#[test] -fn numeric_only_tool_server_gone_without_data_is_not_recognized() { - // Recognition is by the data payload, never the numeric code: a bare -32005 - // with no `data` decodes to a `jsonrpc_-32005` custom error, not recognized. - let err = error_from_envelope(JsonRpcError { - code: -32005, - message: "tool server gone".to_owned(), - data: None, - }); - assert!(!is_workspace_unavailable(&err)); -} - -#[test] -fn custom_error_without_any_details_is_not_recognized() { - // Hand-built Custom with no `details` (no `code`) — unlike a wire `details: - // None`, nothing repopulates `code` here, so it is not recognized. - let err = ToolError::new(ToolErrorKind::Custom, "no details at all"); - assert!(!is_workspace_unavailable(&err)); -} - -#[test] -fn non_custom_error_with_matching_code_is_not_recognized() { - // The kind guard matters: a non-Custom error carrying a matching - // `details.code` must still be rejected. - let err = ToolError::new(ToolErrorKind::NetworkError, "socket closed") - .with_details(json!({ "code": WORKSPACE_UNAVAILABLE_SUBCODE })); - assert_ne!(err.kind, ToolErrorKind::Custom); - assert!(!is_workspace_unavailable(&err)); -} - -#[test] -fn non_custom_decoded_error_is_not_recognized() { - let wire = ToolErrorWire::ToolNotFound { - tool_id: kigi_tool_protocol::ToolId::new("ns:tool").unwrap(), - }; - let err = error_from_envelope(envelope_for(&wire)); - assert_ne!(err.kind, ToolErrorKind::Custom); - assert!(!is_workspace_unavailable(&err)); - - assert!(!is_workspace_unavailable(&ToolError::network_error( - "socket closed" - ))); -} diff --git a/crates/common/kigi-computer-hub-mcp-adapter/Cargo.toml b/crates/common/kigi-computer-hub-mcp-adapter/Cargo.toml deleted file mode 100644 index 3a38f9c..0000000 --- a/crates/common/kigi-computer-hub-mcp-adapter/Cargo.toml +++ /dev/null @@ -1,30 +0,0 @@ -[package] -license = "Apache-2.0" -name = "kigi-computer-hub-mcp-adapter" -version.workspace = true -edition.workspace = true -description = "Bridge between MCP servers and the xAI Computer Hub, registering MCP-discovered tools as native hub tools." - -[features] -metrics = ["dep:prometheus"] - -[dependencies] -async-trait = { workspace = true } -serde = { workspace = true, features = ["derive"] } -serde_json = { workspace = true } -thiserror = { workspace = true } -tokio = { workspace = true, features = ["sync"] } -tracing = { workspace = true } -prometheus = { workspace = true, optional = true } - -kigi-tool-protocol = { workspace = true } -kigi-tool-runtime = { workspace = true } -kigi-tool-types = { workspace = true } -kigi-computer-hub-sdk = { workspace = true } - -[dev-dependencies] -tokio = { workspace = true, features = ["full", "test-util"] } -futures = { workspace = true } - -[lints] -workspace = true diff --git a/crates/common/kigi-computer-hub-mcp-adapter/src/bridge.rs b/crates/common/kigi-computer-hub-mcp-adapter/src/bridge.rs deleted file mode 100644 index df7115d..0000000 --- a/crates/common/kigi-computer-hub-mcp-adapter/src/bridge.rs +++ /dev/null @@ -1,781 +0,0 @@ -//! Core bridge that connects an MCP server to the computer hub. -//! -//! [`McpBridge`] discovers tools from an [`McpTransport`] and registers -//! them with a hub `ToolServer` via one `ToolServerHandler` per -//! tool. Incoming hub calls are translated to MCP `tools/call` and the -//! response is mapped back to [`ToolOutputWire`]. - -use std::sync::Arc; - -use async_trait::async_trait; -use kigi_tool_protocol::{McpBlock, SessionId, ToolId, ToolOutputWire}; -use kigi_tool_runtime::{ToolCallContext, ToolError, ToolStream, TypedToolOutput, terminal_only}; -use kigi_tool_types::ToolDescription; -use serde_json::Value; -use tracing::{debug, info, warn}; - -use crate::transport::McpTransport; -use crate::types::{McpCallResult, McpContent, McpError, McpServerInfo, McpToolDefinition}; - -/// Configuration for an [`McpBridge`] instance. -#[derive(Debug, Clone)] -pub struct McpBridgeConfig { - /// Hub session to bind tools to. - pub session_id: SessionId, - /// Optional namespace prefix for tool descriptions. - pub namespace: Option, -} - -/// Result of a successful [`McpBridge::connect`] call. -/// -/// Contains the bridge handle and the server info returned during the -/// MCP initialize handshake. -pub struct McpBridgeHandle { - /// The bridge managing the MCP-to-hub tool registrations. - pub bridge: McpBridge, - /// Server metadata from the MCP `initialize` response. - pub server_info: McpServerInfo, -} - -impl std::fmt::Debug for McpBridgeHandle { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("McpBridgeHandle") - .field("server_info", &self.server_info.name) - .field("tool_count", &self.bridge.tool_count()) - .finish_non_exhaustive() - } -} - -/// Bridges an MCP server's tools into the computer hub. -/// -/// On construction the bridge performs the MCP `initialize` handshake, -/// discovers tools via `tools/list`, and builds a handler -/// for each one. Callers wire these handlers into a -/// [`kigi_computer_hub_sdk::ToolServerBuilder`] to register them -/// with the hub. -/// -/// Callers **must** call [`McpBridge::shutdown`] before dropping to -/// close the underlying MCP transport cleanly. If the bridge is dropped -/// without an explicit shutdown, a best-effort `close()` is spawned on -/// the tokio runtime (mirroring `ToolServer`'s drop behavior). -pub struct McpBridge { - transport: Arc, - handlers: Vec>, - server_info: McpServerInfo, -} - -impl std::fmt::Debug for McpBridge { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("McpBridge") - .field("server", &self.server_info.name) - .field("tool_count", &self.handlers.len()) - .finish_non_exhaustive() - } -} - -impl McpBridge { - /// Initialize the MCP server, discover its tools, and build handlers. - /// - /// Returns `Err` if the MCP handshake or tool discovery fails. - pub async fn connect( - transport: Arc, - config: &McpBridgeConfig, - ) -> Result { - let server_info = match transport.initialize().await { - Ok(info) => info, - Err(e) => { - crate::metrics::mcp_error(); - return Err(e); - } - }; - info!( - server_name = %server_info.name, - version = %server_info.version, - "MCP server initialized" - ); - - let tools = match transport.list_tools().await { - Ok(t) => t, - Err(e) => { - // Close the transport so the initialized connection is not leaked - // when list_tools fails after a successful initialize. - if let Err(close_err) = transport.close().await { - warn!( - ?close_err, - "failed to close transport after list_tools error" - ); - } - crate::metrics::mcp_error(); - return Err(e); - } - }; - debug!( - server_name = %server_info.name, - tool_count = tools.len(), - "discovered MCP tools" - ); - - let handlers: Vec> = tools - .into_iter() - .filter_map(|def| { - let tool_id = match ToolId::new(&def.name) { - Ok(id) => id, - Err(err) => { - warn!( - tool_name = %def.name, - %err, - "skipping MCP tool with invalid name" - ); - return None; - } - }; - Some(Arc::new(McpToolHandler { - tool_id, - definition: def, - transport: Arc::clone(&transport), - namespace: config.namespace.clone(), - })) - }) - .collect(); - - crate::metrics::mcp_tools_bridged_set(handlers.len() as i64); - - let bridge = McpBridge { - transport, - handlers, - server_info: server_info.clone(), - }; - - Ok(McpBridgeHandle { - bridge, - server_info, - }) - } - - /// Handlers to register with a [`kigi_computer_hub_sdk::ToolServerBuilder`]. - /// - /// Each handler implements `ToolServerHandler` for one MCP tool. - pub fn handlers(&self) -> &[Arc] { - &self.handlers - } - - /// Server metadata from the MCP `initialize` response. - pub fn server_info(&self) -> &McpServerInfo { - &self.server_info - } - - /// Number of tools discovered and registered. - pub fn tool_count(&self) -> usize { - self.handlers.len() - } - - /// Close the underlying MCP transport. - pub async fn shutdown(&self) -> Result<(), McpError> { - crate::metrics::mcp_tools_bridged_set(0); - self.transport.close().await - } -} - -impl Drop for McpBridge { - fn drop(&mut self) { - crate::metrics::mcp_tools_bridged_set(0); - let transport = Arc::clone(&self.transport); - if tokio::runtime::Handle::try_current().is_ok() { - tokio::spawn(async move { - if let Err(err) = transport.close().await { - warn!(?err, "best-effort transport close on drop failed"); - } - }); - } - } -} - -/// Hub-facing handler for a single MCP tool. -/// -/// Translates hub `tool_call_request` frames into MCP `tools/call` -/// invocations and maps the result back to [`ToolOutputWire`]. -pub struct McpToolHandler { - tool_id: ToolId, - definition: McpToolDefinition, - transport: Arc, - namespace: Option, -} - -impl std::fmt::Debug for McpToolHandler { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("McpToolHandler") - .field("tool_id", &self.tool_id) - .finish_non_exhaustive() - } -} - -#[async_trait] -impl kigi_computer_hub_sdk::ToolServerHandler for McpToolHandler { - fn tool_id(&self) -> ToolId { - self.tool_id.clone() - } - - fn description(&self) -> ToolDescription { - let desc = ToolDescription::new( - self.definition.name.clone(), - self.definition.description.clone().unwrap_or_default(), - ); - match self.namespace { - Some(ref ns) => desc.with_namespace(ns.clone()), - None => desc, - } - } - - fn input_schema(&self) -> Option { - self.definition.input_schema.clone() - } - - async fn handle_call(&self, _ctx: ToolCallContext, args: Value) -> ToolStream { - let _start = std::time::Instant::now(); - let tool_id = self.tool_id.clone(); - let result = self - .transport - .call_tool(self.definition.name.as_str(), args) - .await; - crate::metrics::mcp_call_duration_observe(_start.elapsed().as_secs_f64()); - - let terminal = match result { - Ok(call_result) => { - let output = translate_mcp_result(&call_result); - serde_json::to_value(output) - .map(|value| TypedToolOutput::from_value(tool_id, value)) - .map_err(|e| { - crate::metrics::mcp_error(); - ToolError::execution(self.tool_id.clone(), e.to_string()).with_source(e) - }) - } - Err(mcp_err) => { - crate::metrics::mcp_error(); - Err(ToolError::execution( - self.tool_id.clone(), - format!("{mcp_err}"), - )) - } - }; - - terminal_only(terminal) - } -} - -/// Convert an [`McpCallResult`] into the wire output format. -/// -/// - **Error responses** (`is_error: true`): concatenates text-only blocks -/// into a single [`ToolOutputWire::Text`], discarding non-text content -/// (with a warning when content is dropped). -/// - **Empty content**: returns `ToolOutputWire::Text("")` regardless of -/// `is_error` — matches side-effect-only MCP tools. -/// - **Single text block**: returns [`ToolOutputWire::Text`] directly. -/// - **Multi-block / non-text**: returns [`ToolOutputWire::Mcp`] with -/// structured blocks. -fn translate_mcp_result(result: &McpCallResult) -> ToolOutputWire { - if result.content.is_empty() { - return ToolOutputWire::Text(String::new()); - } - - if result.is_error { - let error_text = result - .content - .iter() - .filter_map(|c| match c { - McpContent::Text { text } => Some(text.as_str()), - _ => None, - }) - .collect::>() - .join("\n"); - if error_text.is_empty() { - warn!( - content_count = result.content.len(), - "MCP error response contained only non-text blocks; content dropped" - ); - } - return ToolOutputWire::Text(error_text); - } - - // Single text block → flat text output. - if result.content.len() == 1 - && let Some(McpContent::Text { text }) = result.content.first() - { - return ToolOutputWire::Text(text.clone()); - } - - let blocks: Vec = result - .content - .iter() - .map(|c| match c { - McpContent::Text { text } => McpBlock::Text { text: text.clone() }, - McpContent::Image { mime_type, data } => McpBlock::Image { - mime_type: mime_type.clone(), - data: data.clone(), - }, - McpContent::Resource { - uri, - mime_type, - text, - } => McpBlock::Resource { - uri: uri.clone(), - mime_type: mime_type.clone(), - text: text.clone(), - }, - }) - .collect(); - - ToolOutputWire::Mcp { blocks } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::types::{McpCallResult, McpContent, McpServerInfo, McpToolDefinition}; - use std::sync::atomic::{AtomicBool, Ordering}; - use tokio::sync::Mutex; - - struct MockTransport { - server_info: McpServerInfo, - tools: Vec, - call_response: Mutex>, - call_error: Mutex>, - closed: AtomicBool, - last_call: Mutex>, - } - - impl MockTransport { - fn new(server_info: McpServerInfo, tools: Vec) -> Self { - Self { - server_info, - tools, - call_response: Mutex::new(None), - call_error: Mutex::new(None), - closed: AtomicBool::new(false), - last_call: Mutex::new(None), - } - } - - fn with_call_response(self, response: McpCallResult) -> Self { - Self { - call_response: Mutex::new(Some(response)), - ..self - } - } - - fn with_call_error(self, error: McpError) -> Self { - Self { - call_error: Mutex::new(Some(error)), - ..self - } - } - } - - #[async_trait] - impl McpTransport for MockTransport { - async fn initialize(&self) -> Result { - Ok(self.server_info.clone()) - } - - async fn list_tools(&self) -> Result, McpError> { - Ok(self.tools.clone()) - } - - async fn call_tool(&self, name: &str, arguments: Value) -> Result { - *self.last_call.lock().await = Some((name.to_string(), arguments)); - - if let Some(err) = self.call_error.lock().await.take() { - return Err(err); - } - self.call_response - .lock() - .await - .clone() - .ok_or_else(|| McpError::Transport("no canned response".into())) - } - - async fn close(&self) -> Result<(), McpError> { - self.closed.store(true, Ordering::SeqCst); - Ok(()) - } - } - - fn sample_server_info() -> McpServerInfo { - McpServerInfo { - name: "test-server".into(), - version: "1.0.0".into(), - capabilities: Value::Null, - } - } - - fn sample_tools() -> Vec { - vec![ - McpToolDefinition { - name: "search".into(), - description: Some("Search for items".into()), - input_schema: Some(serde_json::json!({ - "type": "object", - "properties": { "query": { "type": "string" } } - })), - }, - McpToolDefinition { - name: "create".into(), - description: Some("Create an item".into()), - input_schema: None, - }, - ] - } - - fn make_transport(mock: MockTransport) -> Arc { - Arc::new(mock) as Arc - } - - #[tokio::test] - async fn bridge_discovers_and_builds_handlers() { - let transport = make_transport(MockTransport::new(sample_server_info(), sample_tools())); - let config = McpBridgeConfig { - session_id: SessionId::new("test-session").unwrap(), - namespace: None, - }; - - let handle = McpBridge::connect(transport, &config).await.unwrap(); - assert_eq!(handle.server_info.name, "test-server"); - assert_eq!(handle.bridge.tool_count(), 2); - - let ids: Vec = handle - .bridge - .handlers() - .iter() - .map(|h| h.tool_id.as_str().to_string()) - .collect(); - assert!(ids.contains(&"search".to_string())); - assert!(ids.contains(&"create".to_string())); - } - - #[tokio::test] - async fn bridge_handler_descriptions() { - let transport = make_transport(MockTransport::new(sample_server_info(), sample_tools())); - let config = McpBridgeConfig { - session_id: SessionId::new("test-session").unwrap(), - namespace: Some("mcp".into()), - }; - - let handle = McpBridge::connect(transport, &config).await.unwrap(); - let handler = handle - .bridge - .handlers() - .iter() - .find(|h| h.tool_id.as_str() == "search") - .unwrap(); - - use kigi_computer_hub_sdk::ToolServerHandler; - let desc = handler.description(); - assert_eq!(desc.name, "search"); - assert_eq!(desc.description, "Search for items"); - assert_eq!(desc.namespace.as_deref(), Some("mcp")); - assert!(handler.input_schema().is_some()); - } - - #[tokio::test] - async fn bridge_forwards_call_text_response() { - let call_result = McpCallResult { - content: vec![McpContent::Text { - text: "found 3 results".into(), - }], - is_error: false, - }; - let transport = make_transport( - MockTransport::new(sample_server_info(), sample_tools()) - .with_call_response(call_result), - ); - let config = McpBridgeConfig { - session_id: SessionId::new("test-session").unwrap(), - namespace: None, - }; - - let handle = McpBridge::connect(Arc::clone(&transport), &config) - .await - .unwrap(); - let handler = handle - .bridge - .handlers() - .iter() - .find(|h| h.tool_id.as_str() == "search") - .unwrap(); - - use futures::StreamExt; - use kigi_computer_hub_sdk::ToolServerHandler; - - let ctx = ToolCallContext::default(); - let args = serde_json::json!({"query": "test"}); - let mut stream = handler.handle_call(ctx, args).await; - - let item = stream.next().await.unwrap(); - match item { - kigi_tool_runtime::ToolStreamItem::Terminal(Ok(typed)) => { - let output: ToolOutputWire = serde_json::from_value(typed.value).unwrap(); - assert_eq!(output, ToolOutputWire::Text("found 3 results".into())); - } - other => panic!("expected Terminal(Ok(_)), got {other:?}"), - } - } - - #[tokio::test] - async fn bridge_forwards_call_mcp_blocks_response() { - let call_result = McpCallResult { - content: vec![ - McpContent::Text { - text: "result text".into(), - }, - McpContent::Image { - mime_type: "image/png".into(), - data: "base64data".into(), - }, - ], - is_error: false, - }; - let transport = make_transport( - MockTransport::new(sample_server_info(), sample_tools()) - .with_call_response(call_result), - ); - let config = McpBridgeConfig { - session_id: SessionId::new("test-session").unwrap(), - namespace: None, - }; - - let handle = McpBridge::connect(transport, &config).await.unwrap(); - let handler = &handle.bridge.handlers()[0]; - - use futures::StreamExt; - use kigi_computer_hub_sdk::ToolServerHandler; - - let ctx = ToolCallContext::default(); - let mut stream = handler - .handle_call(ctx, Value::Object(Default::default())) - .await; - - let item = stream.next().await.unwrap(); - match item { - kigi_tool_runtime::ToolStreamItem::Terminal(Ok(typed)) => { - let output: ToolOutputWire = serde_json::from_value(typed.value).unwrap(); - match output { - ToolOutputWire::Mcp { blocks } => { - assert_eq!(blocks.len(), 2); - assert!( - matches!(&blocks[0], McpBlock::Text { text } if text == "result text") - ); - assert!( - matches!(&blocks[1], McpBlock::Image { mime_type, .. } if mime_type == "image/png") - ); - } - other => panic!("expected Mcp blocks, got {other:?}"), - } - } - other => panic!("expected Terminal(Ok(_)), got {other:?}"), - } - } - - #[tokio::test] - async fn bridge_handles_mcp_error_response() { - let call_result = McpCallResult { - content: vec![McpContent::Text { - text: "permission denied".into(), - }], - is_error: true, - }; - let transport = make_transport( - MockTransport::new(sample_server_info(), sample_tools()) - .with_call_response(call_result), - ); - let config = McpBridgeConfig { - session_id: SessionId::new("test-session").unwrap(), - namespace: None, - }; - - let handle = McpBridge::connect(transport, &config).await.unwrap(); - let handler = &handle.bridge.handlers()[0]; - - use futures::StreamExt; - use kigi_computer_hub_sdk::ToolServerHandler; - - let ctx = ToolCallContext::default(); - let mut stream = handler.handle_call(ctx, Value::Null).await; - - let item = stream.next().await.unwrap(); - match item { - kigi_tool_runtime::ToolStreamItem::Terminal(Ok(typed)) => { - let output: ToolOutputWire = serde_json::from_value(typed.value).unwrap(); - assert_eq!(output, ToolOutputWire::Text("permission denied".into())); - } - other => panic!("expected Terminal(Ok(_)), got {other:?}"), - } - } - - #[tokio::test] - async fn bridge_handles_transport_error() { - let transport = make_transport( - MockTransport::new(sample_server_info(), sample_tools()) - .with_call_error(McpError::Transport("connection reset".into())), - ); - let config = McpBridgeConfig { - session_id: SessionId::new("test-session").unwrap(), - namespace: None, - }; - - let handle = McpBridge::connect(transport, &config).await.unwrap(); - let handler = &handle.bridge.handlers()[0]; - - use futures::StreamExt; - use kigi_computer_hub_sdk::ToolServerHandler; - - let ctx = ToolCallContext::default(); - let mut stream = handler.handle_call(ctx, Value::Null).await; - - let item = stream.next().await.unwrap(); - match item { - kigi_tool_runtime::ToolStreamItem::Terminal(Err(ref e)) - if e.kind == kigi_tool_runtime::ToolErrorKind::Execution => - { - assert!( - e.detail.contains("connection reset"), - "expected 'connection reset' in: {}", - e.detail - ); - } - other => panic!("expected Terminal(Err(Execution)), got {other:?}"), - } - } - - #[tokio::test] - async fn bridge_shutdown_closes_transport() { - let mock = Arc::new(MockTransport::new(sample_server_info(), sample_tools())); - let transport: Arc = Arc::clone(&mock) as Arc; - let config = McpBridgeConfig { - session_id: SessionId::new("test-session").unwrap(), - namespace: None, - }; - - let handle = McpBridge::connect(transport, &config).await.unwrap(); - assert!(!mock.closed.load(Ordering::SeqCst)); - - handle.bridge.shutdown().await.unwrap(); - assert!(mock.closed.load(Ordering::SeqCst)); - } - - #[tokio::test] - async fn bridge_skips_tools_with_invalid_names() { - let tools = vec![ - McpToolDefinition { - name: "valid_tool".into(), - description: Some("a valid tool".into()), - input_schema: None, - }, - McpToolDefinition { - name: "".into(), - description: Some("empty name".into()), - input_schema: None, - }, - ]; - let transport = make_transport(MockTransport::new(sample_server_info(), tools)); - let config = McpBridgeConfig { - session_id: SessionId::new("test-session").unwrap(), - namespace: None, - }; - - let handle = McpBridge::connect(transport, &config).await.unwrap(); - assert_eq!(handle.bridge.tool_count(), 1); - assert_eq!(handle.bridge.handlers()[0].tool_id.as_str(), "valid_tool"); - } - - #[test] - fn translate_mcp_result_single_text() { - let result = McpCallResult { - content: vec![McpContent::Text { - text: "hello".into(), - }], - is_error: false, - }; - assert_eq!( - translate_mcp_result(&result), - ToolOutputWire::Text("hello".into()) - ); - } - - #[test] - fn translate_mcp_result_error_concatenates_text() { - let result = McpCallResult { - content: vec![ - McpContent::Text { - text: "line 1".into(), - }, - McpContent::Text { - text: "line 2".into(), - }, - ], - is_error: true, - }; - assert_eq!( - translate_mcp_result(&result), - ToolOutputWire::Text("line 1\nline 2".into()) - ); - } - - #[test] - fn translate_mcp_result_mixed_content_uses_blocks() { - let result = McpCallResult { - content: vec![ - McpContent::Text { - text: "hello".into(), - }, - McpContent::Resource { - uri: "file:///test".into(), - mime_type: Some("text/plain".into()), - text: Some("content".into()), - }, - ], - is_error: false, - }; - match translate_mcp_result(&result) { - ToolOutputWire::Mcp { blocks } => assert_eq!(blocks.len(), 2), - other => panic!("expected Mcp, got {other:?}"), - } - } - - #[test] - fn translate_mcp_result_empty_content_returns_empty_text() { - let result = McpCallResult { - content: vec![], - is_error: false, - }; - assert_eq!( - translate_mcp_result(&result), - ToolOutputWire::Text(String::new()) - ); - } - - #[test] - fn translate_mcp_result_empty_error_content_returns_empty_text() { - let result = McpCallResult { - content: vec![], - is_error: true, - }; - assert_eq!( - translate_mcp_result(&result), - ToolOutputWire::Text(String::new()) - ); - } - - #[test] - fn translate_mcp_result_error_with_only_image_drops_content() { - let result = McpCallResult { - content: vec![McpContent::Image { - mime_type: "image/png".into(), - data: "base64data".into(), - }], - is_error: true, - }; - assert_eq!( - translate_mcp_result(&result), - ToolOutputWire::Text(String::new()) - ); - } -} diff --git a/crates/common/kigi-computer-hub-mcp-adapter/src/lib.rs b/crates/common/kigi-computer-hub-mcp-adapter/src/lib.rs deleted file mode 100644 index 23bec75..0000000 --- a/crates/common/kigi-computer-hub-mcp-adapter/src/lib.rs +++ /dev/null @@ -1,52 +0,0 @@ -//! Unified MCP adapter for the xAI Computer Hub. -//! -//! This crate bridges MCP (Model Context Protocol) servers into the -//! computer hub's tool routing infrastructure. An [`McpBridge`] connects -//! to an MCP server via an [`McpTransport`], discovers the server's -//! tools, and produces [`ToolServerHandler`](kigi_computer_hub_sdk::ToolServerHandler) -//! implementations that can be registered with a hub -//! [`ToolServerBuilder`](kigi_computer_hub_sdk::ToolServerBuilder). -//! -//! # Architecture -//! -//! ```text -//! MCP Server <──McpTransport──> McpBridge ──handlers──> ToolServerBuilder -//! (stdio/SSE) (discover+forward) (register with hub) -//! ``` -//! -//! The [`McpTransport`] trait abstracts the wire protocol so the bridge -//! is testable with in-memory mocks. Concrete transport implementations -//! (stdio, HTTP+SSE) are provided by downstream crates. -//! -//! # Usage -//! -//! ```rust,ignore -//! let transport: Arc = /* ... */; -//! let config = McpBridgeConfig { -//! session_id: SessionId::new("session-1").unwrap(), -//! namespace: Some("my-mcp-server".into()), -//! }; -//! let handle = McpBridge::connect(transport, &config).await?; -//! -//! let mut builder = ToolServerBuilder::default() -//! .pool(pool) -//! .url(hub_url) -//! .auth(auth); -//! -//! for handler in handle.bridge.handlers() { -//! builder = builder.tool(handler.clone()); -//! } -//! -//! let server = builder.build().await?; -//! ``` - -#![forbid(unsafe_code)] - -pub mod bridge; -pub(crate) mod metrics; -pub mod transport; -pub mod types; - -pub use bridge::{McpBridge, McpBridgeConfig, McpBridgeHandle, McpToolHandler}; -pub use transport::McpTransport; -pub use types::{McpCallResult, McpContent, McpError, McpServerInfo, McpToolDefinition}; diff --git a/crates/common/kigi-computer-hub-mcp-adapter/src/metrics.rs b/crates/common/kigi-computer-hub-mcp-adapter/src/metrics.rs deleted file mode 100644 index 1b7b110..0000000 --- a/crates/common/kigi-computer-hub-mcp-adapter/src/metrics.rs +++ /dev/null @@ -1,61 +0,0 @@ -//! Feature-gated Prometheus metrics for the MCP adapter bridge. -//! -//! When the `metrics` cargo feature is enabled, each helper records to a -//! lazily-registered Prometheus counter / gauge / histogram. When -//! disabled (the default), every helper compiles to an empty function -//! body so the crate carries zero prometheus dependency. - -#[cfg(feature = "metrics")] -mod inner { - use prometheus::{ - Histogram, IntCounter, IntGauge, exponential_buckets, register_histogram, - register_int_counter, register_int_gauge, - }; - use std::sync::LazyLock; - - static MCP_CALL_DURATION_SECONDS: LazyLock = LazyLock::new(|| { - register_histogram!( - "computer_hub_mcp_call_duration_seconds", - "MCP server response latency for tool calls.", - exponential_buckets(0.01, 2.0, 14).expect("valid bucket params") - ) - .expect("computer_hub_mcp_call_duration_seconds must register once") - }); - - static MCP_ERRORS_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "computer_hub_mcp_errors_total", - "Errors in the MCP adapter pipeline (transport, protocol, or serialization)." - ) - .expect("computer_hub_mcp_errors_total must register once") - }); - - static MCP_TOOLS_BRIDGED: LazyLock = LazyLock::new(|| { - register_int_gauge!( - "computer_hub_mcp_tools_bridged", - "MCP tools currently bridged into the computer hub." - ) - .expect("computer_hub_mcp_tools_bridged must register once") - }); - - pub(crate) fn mcp_call_duration_observe(secs: f64) { - MCP_CALL_DURATION_SECONDS.observe(secs); - } - - pub(crate) fn mcp_error() { - MCP_ERRORS_TOTAL.inc(); - } - - pub(crate) fn mcp_tools_bridged_set(count: i64) { - MCP_TOOLS_BRIDGED.set(count); - } -} - -#[cfg(not(feature = "metrics"))] -mod inner { - pub(crate) fn mcp_call_duration_observe(_secs: f64) {} - pub(crate) fn mcp_error() {} - pub(crate) fn mcp_tools_bridged_set(_count: i64) {} -} - -pub(crate) use inner::*; diff --git a/crates/common/kigi-computer-hub-mcp-adapter/src/transport.rs b/crates/common/kigi-computer-hub-mcp-adapter/src/transport.rs deleted file mode 100644 index 241417f..0000000 --- a/crates/common/kigi-computer-hub-mcp-adapter/src/transport.rs +++ /dev/null @@ -1,38 +0,0 @@ -//! MCP transport abstraction. -//! -//! [`McpTransport`] defines the async interface consumed by [`crate::McpBridge`]. -//! Concrete implementations (stdio, HTTP+SSE) live outside this crate; -//! the trait boundary keeps the bridge testable with in-memory mocks. - -use async_trait::async_trait; -use serde_json::Value; - -use crate::types::{McpCallResult, McpError, McpServerInfo, McpToolDefinition}; - -/// Async interface to a single MCP server connection. -/// -/// Implementations manage the underlying JSON-RPC framing (stdio pipe, -/// HTTP+SSE stream, etc.) and expose the four lifecycle operations the -/// bridge needs. -#[async_trait] -pub trait McpTransport: Send + Sync { - /// Perform the MCP `initialize` handshake with the server. - /// - /// Must be called exactly once before any other method. Returns - /// the server's advertised name, version, and capabilities. - async fn initialize(&self) -> Result; - - /// Discover available tools via MCP `tools/list`. - async fn list_tools(&self) -> Result, McpError>; - - /// Invoke a tool via MCP `tools/call`. - /// - /// `arguments` is the JSON object the model produced for the tool's - /// input schema. - async fn call_tool(&self, name: &str, arguments: Value) -> Result; - - /// Gracefully shut down the transport (close pipes, drop connections). - /// Implementations must be idempotent — a second call after a - /// successful close must return `Ok(())` without error. - async fn close(&self) -> Result<(), McpError>; -} diff --git a/crates/common/kigi-computer-hub-mcp-adapter/src/types.rs b/crates/common/kigi-computer-hub-mcp-adapter/src/types.rs deleted file mode 100644 index 18ac956..0000000 --- a/crates/common/kigi-computer-hub-mcp-adapter/src/types.rs +++ /dev/null @@ -1,107 +0,0 @@ -//! MCP protocol types used by the adapter. -//! -//! These mirror the MCP specification's JSON-RPC shapes for server -//! metadata, tool definitions, and call results. They are intentionally -//! decoupled from any specific transport implementation so the bridge -//! stays testable with in-memory mocks. - -use serde::{Deserialize, Serialize}; - -/// Metadata returned by a successful MCP `initialize` handshake. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct McpServerInfo { - /// Human-readable server name (e.g. `"linear"`, `"github"`). - pub name: String, - /// Semver-ish version reported by the server. - pub version: String, - /// Free-form capability flags advertised during init. - #[serde(default)] - pub capabilities: serde_json::Value, -} - -/// A single tool definition from MCP `tools/list`. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct McpToolDefinition { - /// Unqualified tool name (e.g. `"create_issue"`). - pub name: String, - /// Model-facing description of the tool. - #[serde(default)] - pub description: Option, - /// JSON Schema for the tool's input arguments. - #[serde(default)] - pub input_schema: Option, -} - -/// Result of an MCP `tools/call` invocation. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct McpCallResult { - /// Content blocks returned by the tool. - #[serde(default)] - pub content: Vec, - /// When `true`, the tool signalled an application-level error. - #[serde(default)] - pub is_error: bool, -} - -/// A single content block inside an [`McpCallResult`]. -/// -/// Covers the three content types defined by the MCP specification: -/// text, image (base64-encoded), and embedded resource. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(tag = "type", rename_all = "camelCase")] -pub enum McpContent { - /// Plain text content. - #[serde(rename = "text")] - Text { - /// The text payload. - text: String, - }, - /// Base64-encoded image content. - #[serde(rename = "image")] - Image { - /// MIME type (e.g. `"image/png"`). - #[serde(rename = "mimeType")] - mime_type: String, - /// Base64-encoded image bytes. - data: String, - }, - /// Embedded resource content. - #[serde(rename = "resource")] - Resource { - /// Resource URI. - uri: String, - /// Optional MIME type. - #[serde(default, rename = "mimeType")] - mime_type: Option, - /// Optional text body. - #[serde(default)] - text: Option, - }, -} - -/// Errors originating from MCP transport or protocol handling. -#[derive(Debug, Clone, thiserror::Error)] -pub enum McpError { - /// The underlying transport failed (connection refused, pipe broken, etc.). - #[error("transport error: {0}")] - Transport(String), - - /// The server returned a JSON-RPC error response. - #[error("protocol error (code {code}): {message}")] - Protocol { - /// JSON-RPC error code. - code: i64, - /// Human-readable error message. - message: String, - }, - - /// Timeout waiting for MCP server response. - #[error("timeout: {0}")] - Timeout(String), - - /// The response could not be decoded. - #[error("decode error: {0}")] - Decode(String), -} diff --git a/crates/common/kigi-computer-hub-sdk/Cargo.toml b/crates/common/kigi-computer-hub-sdk/Cargo.toml deleted file mode 100644 index 13f712b..0000000 --- a/crates/common/kigi-computer-hub-sdk/Cargo.toml +++ /dev/null @@ -1,57 +0,0 @@ -[package] -license = "Apache-2.0" -name = "kigi-computer-hub-sdk" -version.workspace = true -edition.workspace = true -description = "SDK for the xAI Computer Hub: connection pool, transparent reconnect, tool harness, and tool-server runtime." - -[features] -metrics = ["dep:prometheus"] - -[dependencies] -tokio = { workspace = true, features = ["rt", "sync", "time", "macros"] } -tokio-tungstenite = { workspace = true, features = ["rustls-tls-native-roots"] } -tokio-util = { workspace = true } -futures = { workspace = true } -serde = { workspace = true, features = ["derive"] } -serde_json = { workspace = true } -dashmap = { workspace = true } -indexmap = { workspace = true } -arc-swap = { workspace = true } -async-trait = { workspace = true } -thiserror = { workspace = true } -tracing = { workspace = true } -tracing-subscriber = { workspace = true } -url = { workspace = true } -http = { workspace = true } -prometheus = { workspace = true, optional = true } -reqwest = { workspace = true } -chrono = { workspace = true } -parking_lot = { workspace = true } - -fastrace = { workspace = true } -# Trace donation: spans convert via the stock fastrace -> OTel reporter -# and ship as standard OTLP payloads. -fastrace-opentelemetry = { workspace = true } -opentelemetry = { workspace = true } -opentelemetry_sdk = { workspace = true } -opentelemetry-proto = { workspace = true } -prost = { workspace = true } -base64 = { workspace = true } - -kigi-tool-protocol = { workspace = true } -kigi-tool-runtime = { workspace = true } -kigi-tool-types = { workspace = true } -kigi-computer-hub-core = { workspace = true } -kigi-tracing = { workspace = true } - -# Integration tests that need heavier backend deps live in a separate sibling crate to keep this dev-dep set minimal. -[dev-dependencies] -tokio = { workspace = true, features = ["full", "test-util"] } -axum = { workspace = true, features = ["ws", "macros"] } -chrono = { workspace = true } -base64 = { workspace = true } -schemars = { workspace = true } - -[lints] -workspace = true diff --git a/crates/common/kigi-computer-hub-sdk/src/admission.rs b/crates/common/kigi-computer-hub-sdk/src/admission.rs deleted file mode 100644 index c4e4e62..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/admission.rs +++ /dev/null @@ -1,333 +0,0 @@ -//! Three-tier semaphore admission + bounded-wait backpressure. -//! -//! Concurrent *running* calls are bounded at three scopes, acquired in a -//! fixed **session → connection → global** order. A consistent -//! most-local-first order is deadlock-free and never holds a scarce -//! global permit while blocking on a local one. A single shared deadline -//! spans all three acquisitions, so total admission latency is bounded by -//! `wait_timeout`, not `3 × wait_timeout`. -//! -//! Under moderate pressure `admit` waits; under very high pressure the -//! deadline elapses and the caller emits the shared overloaded JSON-RPC -//! error (`-32016` "tool_busy") instead of silently dropping the request. - -use std::sync::{Arc, OnceLock}; -use std::time::Duration; - -use dashmap::DashMap; -use kigi_tool_protocol::{ - JsonRpcError, JsonRpcId, JsonRpcResponse, JsonRpcVersion, ResponseOutcome, SessionId, -}; -use serde_json::Value; -use tokio::sync::{OwnedSemaphorePermit, Semaphore}; -use tokio::time::Instant; - -/// Numeric JSON-RPC code for overload rejection (`kigi-tool-protocol` -/// `error_codes.rs`: `-32016` "tool_busy"). -pub(crate) const TOOL_BUSY_CODE: i32 = -32016; - -const TOOL_BUSY_MESSAGE: &str = "tool server busy; tool call rejected"; - -/// Default ceiling for the process-wide concurrency guard. -pub(crate) const DEFAULT_GLOBAL_MAX_INFLIGHT: usize = 1024; -/// Default per-session concurrent running calls. -pub(crate) const DEFAULT_SESSION_MAX_INFLIGHT: usize = 16; -/// Default per-connection concurrent running calls. -pub(crate) const DEFAULT_CONN_MAX_INFLIGHT: usize = 256; -/// Default bounded wait before an overloaded rejection. -pub(crate) const DEFAULT_ADMISSION_WAIT_TIMEOUT: Duration = Duration::from_secs(3); - -/// Ops-tunable override for the process-wide global cap (Helm `env:`). -const GLOBAL_MAX_INFLIGHT_ENV: &str = "XAI_TOOL_SERVER_GLOBAL_MAX_INFLIGHT"; - -/// Inflight-gauge scope labels, in acquisition order. A held [`AdmitGuard`] -/// counts against all three. -const SCOPES: [&str; 3] = ["session", "conn", "global"]; - -/// Build the shared overloaded (`-32016` "tool_busy") JSON-RPC error -/// response. This is the single source of the overload wire shape, reused -/// by BOTH the admission-timeout path (`server::execute_call`) and the -/// demux inbox-full path (`demux::route_session`) so the two never drift. -pub(crate) fn overloaded_response(id: JsonRpcId, session_id: SessionId) -> JsonRpcResponse { - JsonRpcResponse { - jsonrpc: JsonRpcVersion, - id, - session_id: Some(session_id), - outcome: ResponseOutcome::Error(JsonRpcError { - code: TOOL_BUSY_CODE, - message: TOOL_BUSY_MESSAGE.to_owned(), - data: Some(serde_json::json!({ "code": "tool_busy", "retryable": true })), - }), - } -} - -/// Process-wide global admission semaphore, shared by every connection. -/// -/// Initialized once at first use: the value comes from -/// `XAI_TOOL_SERVER_GLOBAL_MAX_INFLIGHT` when present and parseable as a -/// positive integer, otherwise `default_cap` (the builder knob, default -/// [`DEFAULT_GLOBAL_MAX_INFLIGHT`]). Because the cell initializes exactly -/// once, the first caller's `default_cap` and the env var at that instant -/// fix the process-wide capacity. -pub(crate) fn global_semaphore(default_cap: usize) -> Arc { - static SEM: OnceLock> = OnceLock::new(); - SEM.get_or_init(|| { - let raw = std::env::var(GLOBAL_MAX_INFLIGHT_ENV).ok(); - Arc::new(Semaphore::new(resolve_global_cap( - raw.as_deref(), - default_cap, - ))) - }) - .clone() -} - -/// Resolve the process-wide global cap from the raw env value, falling -/// back to `default_cap`. Pure (no global state) so the -/// fall-back-never-panic guarantee is unit-tested: a non-numeric, -/// negative, empty, or zero value all yield `default_cap`. -fn resolve_global_cap(raw: Option<&str>, default_cap: usize) -> usize { - raw.and_then(|v| v.parse::().ok()) - .filter(|n| *n > 0) - .unwrap_or(default_cap) -} - -/// Why admission was refused. -#[derive(Debug, PartialEq, Eq)] -pub(crate) enum Overloaded { - /// The bounded admission deadline elapsed under very high pressure. - Timeout, - /// A semaphore was closed — the server is shutting down. - Shutdown, -} - -/// RAII guard holding all three permits for the call's lifetime. -/// -/// Fields drop in declaration order, so permits are released in reverse -/// of acquisition: global → connection → session. -#[derive(Debug)] -pub(crate) struct AdmitGuard { - _global: OwnedSemaphorePermit, - _conn: OwnedSemaphorePermit, - _session: OwnedSemaphorePermit, -} - -impl Drop for AdmitGuard { - fn drop(&mut self) { - for scope in SCOPES { - crate::metrics::tool_call_inflight_dec(scope); - } - } -} - -/// Three-tier admission controller. One per connection (`conn_sem`); the -/// per-session map is created/destroyed alongside each session loop. -#[derive(Debug)] -pub(crate) struct Admission { - session_sems: DashMap>, - session_max: usize, - conn_sem: Arc, - global_sem: Arc, - wait_timeout: Duration, -} - -impl Admission { - pub(crate) fn new( - session_max: usize, - conn_max: usize, - global_sem: Arc, - wait_timeout: Duration, - ) -> Self { - Self { - session_sems: DashMap::new(), - session_max, - conn_sem: Arc::new(Semaphore::new(conn_max)), - global_sem, - wait_timeout, - } - } - - /// Create the per-session semaphore entry. Called from - /// `bind_session_local` so the entry's lifetime is tied to the - /// session-loop task, not lazily minted in [`Self::admit`]. - pub(crate) fn ensure_session(&self, session_id: &SessionId) { - self.session_sems - .entry(session_id.clone()) - .or_insert_with(|| Arc::new(Semaphore::new(self.session_max))); - } - - /// Remove the per-session semaphore entry on unbind / loop exit. - pub(crate) fn remove_session(&self, session_id: &SessionId) { - self.session_sems.remove(session_id); - } - - /// Acquire one permit at each scope (session → connection → global) - /// against a single shared deadline. - pub(crate) async fn admit(&self, session_id: &SessionId) -> Result { - let start = Instant::now(); - let deadline = start + self.wait_timeout; - // The entry is created in `bind_session_local`; a straggler call - // admitted just after unbind cleanup falls back to a private, - // un-tracked semaphore rather than recreating a leaked entry. - let session_sem = self - .session_sems - .get(session_id) - .map(|s| s.clone()) - .unwrap_or_else(|| Arc::new(Semaphore::new(self.session_max))); - - let session = acquire_until(&session_sem, deadline).await?; - let conn = acquire_until(&self.conn_sem, deadline).await?; - let global = acquire_until(&self.global_sem, deadline).await?; - - crate::metrics::admission_wait_observe(start.elapsed().as_secs_f64()); - for scope in SCOPES { - crate::metrics::tool_call_inflight_inc(scope); - } - Ok(AdmitGuard { - _global: global, - _conn: conn, - _session: session, - }) - } -} - -/// Acquire one owned permit before `deadline`, mapping closed/elapsed to -/// the matching [`Overloaded`] variant. -async fn acquire_until( - sem: &Arc, - deadline: Instant, -) -> Result { - match tokio::time::timeout_at(deadline, sem.clone().acquire_owned()).await { - Ok(Ok(permit)) => Ok(permit), - Ok(Err(_closed)) => Err(Overloaded::Shutdown), - Err(_elapsed) => Err(Overloaded::Timeout), - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn sid(s: &str) -> SessionId { - SessionId::new(s).expect("valid session id") - } - - fn test_admission(session_max: usize, conn_max: usize, global_max: usize) -> Admission { - Admission::new( - session_max, - conn_max, - Arc::new(Semaphore::new(global_max)), - Duration::from_millis(150), - ) - } - - #[test] - fn resolve_global_cap_falls_back_on_bad_input_and_honors_valid() { - // Absent / non-numeric / negative / empty / zero → default (never panic). - assert_eq!(resolve_global_cap(None, 1024), 1024); - assert_eq!(resolve_global_cap(Some("abc"), 1024), 1024); - assert_eq!(resolve_global_cap(Some("-5"), 1024), 1024); - assert_eq!(resolve_global_cap(Some(""), 1024), 1024); - assert_eq!(resolve_global_cap(Some("0"), 1024), 1024); - assert_eq!(resolve_global_cap(Some(" 7"), 1024), 1024); // leading space → parse fails - // A valid positive integer overrides the default. - assert_eq!(resolve_global_cap(Some("2048"), 1024), 2048); - assert_eq!(resolve_global_cap(Some("1"), 1024), 1); - } - - #[test] - fn overloaded_response_carries_minus_32016_and_data_marker() { - let id: JsonRpcId = serde_json::from_value(serde_json::json!("call-1")).expect("id"); - let resp = overloaded_response(id, sid("s1")); - let wire: Value = serde_json::from_str(&serde_json::to_string(&resp).expect("ser")) - .expect("round-trips to json"); - assert_eq!(wire["error"]["code"], TOOL_BUSY_CODE); - assert_eq!(wire["error"]["code"], -32016); - assert_eq!(wire["error"]["data"]["code"], "tool_busy"); - assert_eq!(wire["error"]["data"]["retryable"], true); - assert_eq!(wire["session_id"], "s1"); - assert_eq!(wire["id"], "call-1"); - assert!( - wire.get("result").is_none(), - "overload is an error, never a result" - ); - } - - #[tokio::test(start_paused = true)] - async fn admit_times_out_when_session_saturated() { - let admission = test_admission(2, 16, 64); - let session = sid("sat"); - admission.ensure_session(&session); - - // Hold both session permits. - let g1 = admission.admit(&session).await.expect("first admit"); - let _g2 = admission.admit(&session).await.expect("second admit"); - - // Third admit must elapse the deadline → Timeout (not a hang). - let result = admission.admit(&session).await; - assert_eq!(result.unwrap_err(), Overloaded::Timeout); - - // Releasing one permit frees a slot for the next admit. - drop(g1); - admission - .admit(&session) - .await - .expect("permit released → admit succeeds"); - } - - #[tokio::test(start_paused = true)] - async fn admit_blocks_on_connection_scope_when_conn_saturated() { - // conn_max = 1 is the binding constraint even though session has - // room; a second admit on a *different* session still times out. - let admission = test_admission(8, 1, 64); - let a = sid("a"); - let b = sid("b"); - admission.ensure_session(&a); - admission.ensure_session(&b); - - let _held = admission.admit(&a).await.expect("first admit"); - let result = admission.admit(&b).await; - assert_eq!( - result.unwrap_err(), - Overloaded::Timeout, - "connection cap binds across sessions" - ); - } - - #[tokio::test] - async fn admit_succeeds_repeatedly_under_capacity() { - let admission = test_admission(4, 16, 64); - let session = sid("ok"); - admission.ensure_session(&session); - let mut guards = Vec::new(); - for _ in 0..4 { - guards.push(admission.admit(&session).await.expect("within capacity")); - } - assert_eq!(guards.len(), 4); - } - - #[tokio::test] - async fn closed_semaphore_maps_to_shutdown() { - let global = Arc::new(Semaphore::new(0)); - global.close(); - let admission = Admission::new(4, 16, global, Duration::from_secs(5)); - let session = sid("closed"); - admission.ensure_session(&session); - let result = admission.admit(&session).await; - assert_eq!(result.unwrap_err(), Overloaded::Shutdown); - } - - #[tokio::test(start_paused = true)] - async fn straggler_admit_after_remove_uses_private_permit() { - let admission = test_admission(1, 16, 64); - let session = sid("gone"); - // No ensure_session: simulate a straggler after unbind removed it. - admission.remove_session(&session); - // Falls back to a private semaphore and still admits (no panic, - // no leaked tracked entry). - let _g = admission.admit(&session).await.expect("private fallback"); - assert!( - admission.session_sems.get(&session).is_none(), - "straggler must not recreate a tracked entry" - ); - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/auth.rs b/crates/common/kigi-computer-hub-sdk/src/auth.rs deleted file mode 100644 index 592d8cb..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/auth.rs +++ /dev/null @@ -1,238 +0,0 @@ -//! Auth credentials and pool-dedup principal keys. -//! -//! [`AuthCredential`] models the credential the client attaches at -//! handshake time. Two variants are supported: -//! -//! - [`AuthCredential::Bearer`] for the simple "Authorization: Bearer -//! …" path (e.g. JWT-against-OAuth2 deployments). -//! - [`AuthCredential::Headers`] for callers that already hold a -//! pre-built header bundle (e.g. signed identity headers generated -//! by an upstream proxy or test harness). -//! -//! [`PrincipalKey`] is the stable hashable projection of an -//! `AuthCredential`; the pool keys connections by -//! `(url, principal_key)` so two [`crate::ToolServer`] builds with the -//! same credential reuse one socket while distinct credentials open -//! distinct sockets. The server derives `user_id` from the credential at -//! upgrade time and returns it in the hello ack — the SDK never needs -//! to carry `user_id` alongside the credential. -//! -//! ## Pool dedup and credential refresh -//! -//! Both variants include the secret material in the `PrincipalKey` -//! fingerprint. This is deliberate: distinct secrets imply distinct -//! credentials, so two callers with different tokens open distinct -//! sockets. The trade-off is that a caller that rotates its bearer JWT -//! every N minutes will open a new socket on each rotation. -//! Long-running tool servers should reuse the SAME [`AuthCredential`] -//! instance across builds and refresh the credential out-of-band rather -//! than hand a fresh JWT to every build. - -use std::collections::BTreeMap; -use std::fmt; - -use http::HeaderName; -use http::header::AUTHORIZATION; - -use crate::error::ClientError; - -/// Credential carried into the WebSocket upgrade. -/// -/// Clones are cheap (the secret material is at most a small number of -/// owned strings). The server derives `user_id` from the credential at -/// upgrade time and returns it in the [`kigi_tool_protocol::HelloAckMsg`]. -#[derive(Clone, PartialEq, Eq, Hash)] -pub enum AuthCredential { - /// Bearer token attached as the `Authorization: Bearer …` header. - Bearer { token: String }, - /// Pre-built header bundle. Used when the auth flow lives outside - /// the SDK (e.g. an upstream proxy that already produced signed - /// identity headers). Header order is canonicalised for stable - /// hashing via [`BTreeMap`]; names are lowercased and validated - /// as `HeaderName` at construction time so an invalid name - /// surfaces as [`ClientError::InvalidConfig`] instead of being - /// silently dropped at upgrade time. - Headers { headers: BTreeMap }, -} - -impl AuthCredential { - /// Convenience constructor for the bearer-token shape. - pub fn bearer(token: impl Into) -> Self { - Self::Bearer { - token: token.into(), - } - } - - /// Convenience constructor for the raw-header bundle shape. - /// - /// Names are canonicalised to lowercase and validated as - /// [`HeaderName`] at construction so an invalid header (e.g. one - /// containing a newline injection attempt) returns - /// [`ClientError::InvalidConfig`] rather than being silently - /// filtered out at upgrade time. - pub fn headers(headers: I) -> Result - where - I: IntoIterator, - K: AsRef, - V: Into, - { - let mut map: BTreeMap = BTreeMap::new(); - for (raw_name, raw_value) in headers { - let name = raw_name.as_ref().to_ascii_lowercase(); - HeaderName::from_bytes(name.as_bytes()).map_err(|err| { - ClientError::InvalidConfig(format!("invalid header name {name:?}: {err}")) - })?; - map.insert(name, raw_value.into()); - } - Ok(Self::Headers { headers: map }) - } - - /// Stable hashable projection used as the pool dedup key. - /// - /// Distinct credentials hash equal iff they carry the same secret - /// material. See the module-level "Pool dedup and credential - /// refresh" section for the implications when bearer tokens are - /// rotated. - pub fn principal_key(&self) -> PrincipalKey { - match self { - Self::Bearer { token } => PrincipalKey { - fingerprint: format!("bearer:{token}"), - }, - Self::Headers { headers } => { - // Concatenate canonicalised name=value pairs so the - // fingerprint is order-independent. - let mut joined = String::with_capacity(headers.len() * 32); - for (name, value) in headers { - joined.push_str(name); - joined.push('='); - joined.push_str(value); - joined.push('\n'); - } - PrincipalKey { - fingerprint: format!("headers:{joined}"), - } - } - } - } - - /// Headers to attach to the WebSocket upgrade request. - /// - /// `Headers` variant entries are infallible at this point — names - /// were validated by [`Self::headers`]. - pub fn upgrade_headers(&self) -> Vec<(HeaderName, String)> { - match self { - Self::Bearer { token, .. } => { - vec![(AUTHORIZATION, format!("Bearer {token}"))] - } - Self::Headers { headers, .. } => headers - .iter() - .filter_map(|(name, value)| { - HeaderName::from_bytes(name.as_bytes()) - .ok() - .map(|n| (n, value.clone())) - }) - .collect(), - } - } -} - -impl fmt::Debug for AuthCredential { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - // Never log the secret; surface only the variant. - match self { - Self::Bearer { .. } => f - .debug_struct("AuthCredential::Bearer") - .finish_non_exhaustive(), - Self::Headers { headers } => f - .debug_struct("AuthCredential::Headers") - .field("header_count", &headers.len()) - .finish_non_exhaustive(), - } - } -} - -/// Stable hashable projection of an [`AuthCredential`] used as the -/// pool dedup key alongside the connect URL. Two connections with the -/// same token fingerprint will get the same server-assigned `user_id`. -#[derive(Clone, PartialEq, Eq, Hash)] -pub struct PrincipalKey { - fingerprint: String, -} - -impl fmt::Debug for PrincipalKey { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("PrincipalKey").finish_non_exhaustive() - } -} - -/// Owner identity surfaced by an [`AuthProvider`] alongside its credential. -/// -/// Mirrors the OAuth principal fields the provider parsed from its auth source. -/// It is kept separate from [`AuthCredential`] on purpose: identity must NOT -/// participate in pool-dedup hashing (that keys only on the secret), and the -/// credential's `Eq`/`Hash` derives must stay token-only. Consumers (e.g. the -/// workspace) map this onto their own identity record. -#[derive(Clone, Debug, Default, PartialEq, Eq)] -pub struct AuthIdentity { - /// Stable user identifier (owner of the bearer token). - pub user_id: String, - /// OAuth `principal_type` wire string (`"User"` / `"Team"`), when known. - pub principal_type: Option, - /// Team id when `principal_type == "Team"`; otherwise `None`. - pub principal_id: Option, -} - -/// Credential provider called on every connect/reconnect. -pub trait AuthProvider: Send + Sync + std::fmt::Debug { - fn current(&self) -> AuthCredential; - - /// Stable pool-dedup key, decoupled from the per-connect credential. - /// - /// Defaults to the current credential's key (existing behavior). A provider - /// that re-mints a rotating secret on every [`Self::current`] call (e.g. a - /// refresh-before-use bearer) MUST override this to key only on stable - /// identity, otherwise each rotation fragments the connection pool. - fn principal_key(&self) -> PrincipalKey { - self.current().principal_key() - } - - /// Owner identity behind the credential, when the provider can surface it. - /// - /// Defaults to `None` for providers that only carry a bearer token (e.g. a - /// bare [`AuthCredential`]). Providers that parse OAuth principal fields - /// (e.g. OIDC) override this so downstream consumers can attribute - /// requests without a second auth-source read. - fn identity(&self) -> Option { - None - } -} - -pub type SharedAuthProvider = std::sync::Arc; - -impl AuthProvider for AuthCredential { - fn current(&self) -> AuthCredential { - self.clone() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn invalid_header_name_rejected_at_construction() { - let cred = AuthCredential::headers([("authorization\nx-injected", "value")]); - match cred { - Err(ClientError::InvalidConfig(msg)) => { - assert!(msg.contains("invalid header name"), "got {msg}") - } - other => panic!("expected InvalidConfig; got {other:?}"), - } - } - - #[test] - fn valid_headers_accepted() { - let cred = AuthCredential::headers([("authorization", "Bearer token")]).expect("valid"); - assert_eq!(cred.upgrade_headers().len(), 1); - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/cancel.rs b/crates/common/kigi-computer-hub-sdk/src/cancel.rs deleted file mode 100644 index 0e1d93b..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/cancel.rs +++ /dev/null @@ -1,301 +0,0 @@ -//! Per-session strict-cancellation registry. -//! -//! Maps each in-flight `tool_call_id` to its [`CancellationToken`] so a -//! `Cancel` hook (or session teardown) can hard-cancel the running call -//! by dropping its future. A small `pending` tombstone set covers the -//! race where a `Cancel` arrives *before* the dispatcher registered the -//! token (the symmetric window to pre-spawn registration): the id is -//! tombstoned and the dispatcher cancels it at registration time. -//! -//! One registry per session, tied to the session-loop lifetime alongside -//! the inbox and the per-session admission semaphore. - -use std::sync::atomic::{AtomicBool, Ordering}; - -use dashmap::{DashMap, DashSet}; -use kigi_tool_protocol::ToolCallId; -use tokio_util::sync::CancellationToken; - -/// Upper bound on outstanding pre-registration tombstones. Tombstones -/// cover the microscopic window between a `Cancel` hook and the matching -/// `register`, so in steady state the set holds a handful of entries. A -/// `Cancel` whose call never registers (e.g. one racing call completion, -/// after `deregister` already removed the live token) leaves a tombstone -/// that no `register` ever consumes; this cap reclaims such stragglers so -/// a single long-lived session cannot grow `pending` without bound. -const MAX_PENDING_TOMBSTONES: usize = 8192; - -/// Per-session `tool_call_id -> CancellationToken` map plus a pending -/// tombstone set for cancels that land before registration. -#[derive(Default, Debug)] -pub(crate) struct CancelRegistry { - map: DashMap, - pending: DashSet, - /// Set once by [`Self::cancel_all`] (teardown). After this, every new - /// `register` starts cancelled so a request dispatched in the teardown - /// window cannot escape as an orphaned, uncancellable task. - closed: AtomicBool, -} - -impl CancelRegistry { - /// Register `token` for `call_id` before the call is spawned. If a - /// `Cancel` already tombstoned this id, the token is cancelled - /// immediately so the call starts cancelled. Returns whether the - /// token was pre-cancelled (by a tombstone or because the registry was - /// torn down). - pub(crate) fn register(&self, call_id: ToolCallId, token: &CancellationToken) -> bool { - if self.closed.load(Ordering::Acquire) { - token.cancel(); - return true; - } - let pre_cancelled = self.pending.remove(&call_id).is_some(); - if pre_cancelled { - token.cancel(); - } - self.map.insert(call_id.clone(), token.clone()); - // Re-check after the insert: if `cancel_all` drained the map - // between our closed-check and the insert, our entry would be - // missed. The DashMap shard lock orders the insert against the - // drain, so observing `closed` here guarantees we cancel + drop - // any entry the drain could not reach (closes the teardown race). - if self.closed.load(Ordering::Acquire) { - if let Some((_, missed)) = self.map.remove(&call_id) { - missed.cancel(); - } - return true; - } - pre_cancelled - } - - /// Cancel a live call, else tombstone the id so the dispatcher - /// cancels it at registration time. Returns true when a live token - /// was found and cancelled. - pub(crate) fn cancel(&self, call_id: &ToolCallId) -> bool { - if let Some((_, token)) = self.map.remove(call_id) { - token.cancel(); - true - } else { - if self.pending.len() >= MAX_PENDING_TOMBSTONES { - // Evict one straggler tombstone (a cancel whose call never - // registered) before inserting so the set stays bounded. - // Collect the key first, then remove, so we never hold a - // shard iterator across the removal. - let stale = self.pending.iter().next().map(|e| e.key().clone()); - if let Some(stale) = stale { - self.pending.remove(&stale); - } - } - self.pending.insert(call_id.clone()); - false - } - } - - /// Deregister a call's token on completion or cancel. Idempotent. - pub(crate) fn deregister(&self, call_id: &ToolCallId) { - self.map.remove(call_id); - } - - /// Whether [`Self::cancel_all`] has closed this registry. A closed - /// registry marks a session whose loop is (or is about to be) torn - /// down — used by the soft-rebind liveness gate. - pub(crate) fn is_closed(&self) -> bool { - self.closed.load(Ordering::Acquire) - } - - /// Drain-and-cancel every live token and close the registry. Used on - /// session teardown (`unbind_session` / `shutdown` / full rebind of a - /// dead loop — a soft rebind of a live session keeps its registry) so - /// detached `execute_call` tasks wind down promptly AND any call - /// dispatched in - /// the teardown window starts cancelled (see [`Self::register`]). - /// Returns the number of tokens cancelled. - pub(crate) fn cancel_all(&self) -> usize { - // Mark closed BEFORE draining so a concurrent `register` either - // observes the close (and self-cancels) or has its entry drained - // here — never both-miss. - self.closed.store(true, Ordering::Release); - let mut cancelled = 0; - self.map.retain(|_, token| { - token.cancel(); - cancelled += 1; - false - }); - // Drop tombstones too: teardown closes the registry, so no future - // `register` will consume them. Leaving them would let a stale - // straggler set survive to the end of the (already-done) session. - self.pending.clear(); - cancelled - } - - #[cfg(test)] - pub(crate) fn live_count(&self) -> usize { - self.map.len() - } - - #[cfg(test)] - pub(crate) fn pending_count(&self) -> usize { - self.pending.len() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn cid() -> ToolCallId { - ToolCallId::new_v7() - } - - #[test] - fn cancel_live_token_fires_and_removes_entry() { - let reg = CancelRegistry::default(); - let id = cid(); - let token = CancellationToken::new(); - assert!( - !reg.register(id.clone(), &token), - "fresh register, no tombstone" - ); - assert_eq!(reg.live_count(), 1); - - assert!(reg.cancel(&id), "live token must report a hit"); - assert!( - token.is_cancelled(), - "the registered token must be cancelled" - ); - assert_eq!(reg.live_count(), 0, "cancel removes the live entry"); - assert_eq!(reg.pending_count(), 0, "a live hit leaves no tombstone"); - } - - #[test] - fn cancel_before_registration_tombstones_then_register_pre_cancels() { - let reg = CancelRegistry::default(); - let id = cid(); - - // Cancel arrives first: no live token, so it tombstones. - assert!(!reg.cancel(&id), "no live token yet → miss"); - assert_eq!(reg.pending_count(), 1); - assert_eq!(reg.live_count(), 0); - - // Registration consumes the tombstone and starts cancelled. - let token = CancellationToken::new(); - assert!( - reg.register(id.clone(), &token), - "register must report the pre-cancel" - ); - assert!(token.is_cancelled(), "tombstone must pre-cancel the token"); - assert_eq!(reg.pending_count(), 0, "tombstone consumed at registration"); - assert_eq!(reg.live_count(), 1); - } - - #[test] - fn deregister_clears_live_entry_without_cancel() { - let reg = CancelRegistry::default(); - let id = cid(); - let token = CancellationToken::new(); - reg.register(id.clone(), &token); - - reg.deregister(&id); - assert_eq!(reg.live_count(), 0); - assert!( - !token.is_cancelled(), - "deregister on normal completion must NOT cancel the token" - ); - // A later cancel for a completed call only tombstones (harmless). - assert!(!reg.cancel(&id)); - assert_eq!(reg.pending_count(), 1); - } - - #[test] - fn cancel_all_drains_and_cancels_every_live_token() { - let reg = CancelRegistry::default(); - let ids: Vec = (0..5).map(|_| cid()).collect(); - let tokens: Vec = ids - .iter() - .map(|id| { - let t = CancellationToken::new(); - reg.register(id.clone(), &t); - t - }) - .collect(); - assert_eq!(reg.live_count(), 5); - - assert_eq!( - reg.cancel_all(), - 5, - "cancel_all reports every drained token" - ); - assert_eq!(reg.live_count(), 0, "registry is empty after teardown"); - for token in &tokens { - assert!(token.is_cancelled(), "every live token must be cancelled"); - } - // Idempotent: a second teardown cancels nothing. - assert_eq!(reg.cancel_all(), 0); - } - - #[test] - fn register_after_cancel_all_starts_cancelled() { - // Teardown race regression: once `cancel_all` has closed the - // registry, a call dispatched in the teardown window must start - // cancelled and must NOT linger as a live, uncancellable entry. - let reg = CancelRegistry::default(); - assert_eq!(reg.cancel_all(), 0, "empty teardown cancels nothing"); - - let id = cid(); - let token = CancellationToken::new(); - assert!( - reg.register(id.clone(), &token), - "register on a closed registry must report pre-cancel" - ); - assert!( - token.is_cancelled(), - "a call dispatched after teardown must start cancelled" - ); - assert_eq!( - reg.live_count(), - 0, - "a closed-registry register must not leave a live (orphan) entry" - ); - } - - #[test] - fn register_without_tombstone_does_not_cancel() { - let reg = CancelRegistry::default(); - let id = cid(); - let token = CancellationToken::new(); - assert!(!reg.register(id, &token)); - assert!( - !token.is_cancelled(), - "a clean registration must leave the token live" - ); - } - - #[test] - fn cancel_all_clears_pending_tombstones() { - let reg = CancelRegistry::default(); - reg.cancel(&cid()); - reg.cancel(&cid()); - assert_eq!(reg.pending_count(), 2); - reg.cancel_all(); - assert_eq!( - reg.pending_count(), - 0, - "teardown must drop pending tombstones" - ); - } - - #[test] - fn pending_tombstones_stay_bounded_under_spurious_cancels() { - // A long-lived session that keeps receiving cancels for call_ids - // that never register (e.g. cancels racing call completion) must - // not grow `pending` without bound. - let reg = CancelRegistry::default(); - for _ in 0..(MAX_PENDING_TOMBSTONES + 256) { - assert!(!reg.cancel(&cid()), "never-registered id is a miss"); - } - assert!( - reg.pending_count() <= MAX_PENDING_TOMBSTONES, - "tombstone set must stay within its cap, got {}", - reg.pending_count() - ); - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/connection.rs b/crates/common/kigi-computer-hub-sdk/src/connection.rs deleted file mode 100644 index 834ed4a..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/connection.rs +++ /dev/null @@ -1,2694 +0,0 @@ -//! Per-`(url, principal)` WebSocket connection actor. -//! -//! # Why this exists -//! -//! Multiple [`crate::ToolServer`] instances in the same process MAY -//! attach to the same server URL with the same credential. Opening one -//! socket per server would multiply server-side connection cost, fan-out -//! the per-tool ack chatter, and make per-frame envelope checks -//! ambiguous (the server can't tell which of N sockets owns a session -//! binding). The pool collapses every `(url, principal)` to one -//! [`HubConnection`]; refcounted session bindings make the collapse -//! safe. -//! -//! # The reconnect / replay state machine -//! -//! When the underlying socket drops, in-flight `tool_call_request` -//! responses CANNOT be recovered (the server holds no replay log). The -//! connection actor therefore: -//! -//! 1. Drains every parked response waiter with -//! [`crate::ClientError::NetworkError`] so callers can fast-fail -//! instead of deadlocking. -//! 2. Reconnects with exponential backoff (capped). -//! 3. Re-runs the `hello` handshake. -//! 4. The ToolServer replays `serve{session_id, tools}` per active -//! session via the on_reconnect callback. The server auto-registers -//! sessions from `serve` so no separate wire call is needed. -//! 5. Drains any outbound frames that buffered during step 1-4. -use crate::auth::{AuthCredential, AuthProvider, PrincipalKey}; -use crate::demux::Demux; -use crate::error::ClientError; -use crate::handshake::send_hello; -use crate::refcount::RefCountedSet; -use futures::stream::SplitSink; -use futures::stream::SplitStream; -use futures::{SinkExt, Stream, StreamExt}; -use http::HeaderName; -use http::header::HeaderValue; -use kigi_tool_protocol::{ - ConnectionId, ConnectionKind, JsonRpcId, JsonRpcRequest, JsonRpcResponse, JsonRpcVersion, - Method, PongFrame, ResponseOutcome, SessionId, -}; -use serde_json::Value; -use std::sync::{Arc, Weak}; -use std::time::{Duration, Instant, SystemTime}; -use tokio::net::TcpStream; -use tokio::sync::{Mutex, broadcast, mpsc, oneshot}; -use tokio::time::sleep; -use tokio_tungstenite::tungstenite::Message; -use tokio_tungstenite::tungstenite::client::IntoClientRequest; -use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; -use tokio_util::sync::CancellationToken; -use tracing::{info, warn}; -use url::Url; -/// Outbound mpsc bound. Picked to match the server's per-actor outbound -/// buffer so a single-process roundtrip never dead-blocks on sender -/// capacity. -const OUTBOUND_BUFFER: usize = 256; -/// Backoff schedule (in ms) for reconnect attempts. The last value is -/// reused for any further attempts so the cap is `10s`. -const RECONNECT_BACKOFF_MS: &[u64] = &[100, 200, 500, 1_000, 2_000, 5_000, 10_000]; -/// Floor for the per-attempt reconnect budget: a small liveness override -/// must not shrink it below what a WAN handshake + session replay needs, -/// or the retry loop would livelock aborting every attempt at the bound. -const RECONNECT_ATTEMPT_MIN_BUDGET: Duration = Duration::from_secs(30); -/// Per-attempt reconnect budget: the liveness deadline, floored so liveness -/// tuning bounds detection, not connection establishment. -fn reconnect_attempt_budget(liveness_deadline: Duration) -> Duration { - liveness_deadline.max(RECONNECT_ATTEMPT_MIN_BUDGET) -} -/// Default WebSocket keepalive ping cadence when a connection does not -/// override [`ConnectionTuning::ws_ping_interval`]. -const DEFAULT_WS_PING_INTERVAL: Duration = Duration::from_secs(30); -const SERVE_ATTEMPT_TIMEOUT: Duration = Duration::from_secs(30); -const SERVE_MAX_ATTEMPTS: u32 = 3; -const CLOCK_PROBE_INTERVAL: Duration = Duration::from_secs(5); -const CLOCK_JUMP_ACCUM_MIN_MS: u64 = 100; -const CLOCK_JUMP_REPORT_MIN_MS: u64 = 2_000; -type WriteErrorSlot = Arc>>; -struct HealthState { - last_inbound: Instant, - mono_ref: Instant, - wall_ref: SystemTime, - clock_jump_accum_ms: u64, -} -struct HealthSnapshot { - last_inbound: Instant, - /// Monotonic time elapsed since the last probe window rolled (the most - /// recent inbound frame or 5s clock probe) — NOT since connection start. - /// Healthy traffic keeps this small (<= ~5s); the meaningful freeze - /// signal in this snapshot is `clock_jump_ms`. - since_last_probe_monotonic_ms: u64, - /// Wall-clock time elapsed over the same probe window as - /// `since_last_probe_monotonic_ms`. - since_last_probe_wall_ms: u64, - clock_jump_ms: u64, -} -struct ConnHealth { - state: parking_lot::Mutex, -} -impl ConnHealth { - fn new() -> Self { - Self { - state: parking_lot::Mutex::new(Self::fresh_state()), - } - } - fn fresh_state() -> HealthState { - HealthState { - last_inbound: Instant::now(), - mono_ref: Instant::now(), - wall_ref: SystemTime::now(), - clock_jump_accum_ms: 0, - } - } - fn deltas(state: &HealthState) -> (u64, u64) { - let mono_ms = state.mono_ref.elapsed().as_millis() as u64; - let wall_ms = SystemTime::now() - .duration_since(state.wall_ref) - .map(|d| d.as_millis() as u64) - .unwrap_or(0); - (mono_ms, wall_ms) - } - fn roll(state: &mut HealthState) { - let (mono_ms, wall_ms) = Self::deltas(state); - let excess = wall_ms.saturating_sub(mono_ms); - if excess >= CLOCK_JUMP_ACCUM_MIN_MS { - state.clock_jump_accum_ms = state.clock_jump_accum_ms.saturating_add(excess); - } - state.mono_ref = Instant::now(); - state.wall_ref = SystemTime::now(); - } - fn record_inbound(&self) { - let mut state = self.state.lock(); - Self::roll(&mut state); - state.last_inbound = Instant::now(); - } - fn refresh_clock(&self) { - let mut state = self.state.lock(); - Self::roll(&mut state); - } - fn snapshot(&self) -> HealthSnapshot { - let state = self.state.lock(); - let (mono_ms, wall_ms) = Self::deltas(&state); - let excess = wall_ms.saturating_sub(mono_ms); - let total = - state - .clock_jump_accum_ms - .saturating_add(if excess >= CLOCK_JUMP_ACCUM_MIN_MS { - excess - } else { - 0 - }); - HealthSnapshot { - last_inbound: state.last_inbound, - since_last_probe_monotonic_ms: mono_ms, - since_last_probe_wall_ms: wall_ms, - clock_jump_ms: if total >= CLOCK_JUMP_REPORT_MIN_MS { - total - } else { - 0 - }, - } - } - fn reset(&self) { - *self.state.lock() = Self::fresh_state(); - } -} -enum DisconnectCause { - CloseFrame(Option), - Eof, - ReadError(String), - WriteError(String), - Forced, - /// No inbound frame arrived within the inbound-liveness deadline, so the - /// transport is silently dead (snapshot-restored VM, NAT/LB flow expiry). - LivenessDeadline, -} -impl DisconnectCause { - fn label(&self) -> &'static str { - match self { - Self::CloseFrame(_) => "close_frame", - Self::Eof => "eof", - Self::ReadError(_) => "transport_read_error", - Self::WriteError(_) => "transport_write_error", - Self::Forced => "forced", - Self::LivenessDeadline => "liveness_deadline", - } - } - fn close_code(&self) -> Option { - match self { - Self::CloseFrame(code) => *code, - _ => None, - } - } - fn detail(&self) -> Option<&str> { - match self { - Self::ReadError(detail) | Self::WriteError(detail) => Some(detail), - _ => None, - } - } -} -struct OutageInfo { - cause: DisconnectCause, - prev_connection_id: Option, - prev_connection_duration_ms: u64, - last_inbound: Instant, - detect_ms: u64, - since_last_probe_monotonic_ms: u64, - since_last_probe_wall_ms: u64, - clock_jump_ms: u64, -} -enum DeadlineCallError { - TimedOut(Duration), - Other(ClientError), -} -impl From for ClientError { - fn from(err: DeadlineCallError) -> Self { - match err { - DeadlineCallError::TimedOut(timeout) => { - ClientError::NetworkError(format!("request timed out after {timeout:?}")) - } - DeadlineCallError::Other(e) => e, - } - } -} -struct WaiterGuard<'a> { - demux: &'a Demux, - request_id: &'a kigi_tool_protocol::RequestId, -} -impl Drop for WaiterGuard<'_> { - fn drop(&mut self) { - let _ = self.demux.take_response_waiter(self.request_id); - } -} -/// Process-wide default reconnect schedule, materialised once from -/// [`RECONNECT_BACKOFF_MS`]. Connections that do not override -/// [`ConnectionTuning::reconnect_backoff`] share this `Arc` (cheap clone, -/// no per-connect allocation). -fn default_reconnect_backoff() -> Arc<[Duration]> { - static DEFAULT: std::sync::OnceLock> = std::sync::OnceLock::new(); - DEFAULT - .get_or_init(|| { - RECONNECT_BACKOFF_MS - .iter() - .map(|&ms| Duration::from_millis(ms)) - .collect() - }) - .clone() -} -/// Resolve a configured backoff schedule, falling back to the built-in -/// table when unset (or empty, which would be degenerate). -fn resolve_reconnect_backoff(configured: Option>) -> Arc<[Duration]> { - match configured { - Some(schedule) if !schedule.is_empty() => schedule, - _ => default_reconnect_backoff(), - } -} -/// Resolve the keepalive ping cadence, clamping an unset *or zero* value to -/// [`DEFAULT_WS_PING_INTERVAL`]. `tokio::time::interval` panics on a zero -/// period, so a configured `Duration::ZERO` (e.g. via -/// `with_ws_ping_interval(0)` or a `StatusConfig.ws_ping` of 0) must never -/// reach the writer task. -fn resolve_ws_ping_interval(configured: Option) -> Duration { - match configured { - Some(interval) if !interval.is_zero() => interval, - _ => DEFAULT_WS_PING_INTERVAL, - } -} -/// Resolve the inbound-liveness deadline, clamping an unset *or zero* value -/// to 2.5× the (already-resolved) keepalive ping cadence — 75s at the -/// default 30s ping. -/// -/// The default multiple is chosen for fleet-wide false-positive safety: a -/// healthy connection delivers at least one inbound frame per ping period -/// (the server must answer each WS `Ping` with a `Pong`, and any data frame -/// also counts), so 2.5× tolerates a fully lost/coalesced pong plus -/// scheduling jitter before declaring death. It still detects a silently -/// dead transport (snapshot-restored VM, NAT/LB flow expiry) within ~1–2 -/// keepalive cycles instead of TCP-retransmission timescales (15+ min). -/// Explicit overrides are honored verbatim; keep them comfortably above -/// the ping interval or a healthy-but-idle connection will be churned. -fn resolve_ws_liveness_deadline(configured: Option, ping_interval: Duration) -> Duration { - match configured { - Some(deadline) if !deadline.is_zero() => deadline, - _ => ping_interval.saturating_mul(5) / 2, - } -} -/// Optional, default-preserving connection-tuning knobs carried from the -/// pool/builder into [`ConnectionConfig`]. `Default` leaves every value -/// `None`, reproducing the historical hardcoded behaviour — and lets -/// config constructors write `tuning: ConnectionTuning::default()` so new -/// knobs never churn every [`ConnectionConfig`] literal. -#[derive(Clone, Default)] -pub struct ConnectionTuning { - /// Override for the keepalive ping cadence. `None` (or zero) ⇒ - /// [`DEFAULT_WS_PING_INTERVAL`]. - pub ws_ping_interval: Option, - /// Override for the inbound-liveness deadline: with no inbound frame of - /// any kind for this long, the reader declares the socket dead and - /// reconnects. `None` (or zero) ⇒ 2.5× the effective ping cadence (see - /// [`resolve_ws_liveness_deadline`]). - pub ws_liveness_deadline: Option, - /// Override for the reconnect backoff schedule. `None` (or empty) ⇒ - /// the built-in [`RECONNECT_BACKOFF_MS`] table. Stored as - /// `Arc<[Duration]>` so it is shared (not deep-copied) per reconnect. - pub reconnect_backoff: Option>, -} -/// Pool dedup key. Two connections are pooled together iff their -/// `(url, principal)` match. -#[derive(Clone, PartialEq, Eq, Hash)] -pub struct ConnKey { - /// Normalised URL (parsed by [`Url::parse`]). - pub url: String, - /// Principal projection of the [`AuthCredential`]. - pub principal: PrincipalKey, -} -impl std::fmt::Debug for ConnKey { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ConnKey") - .field("url", &self.url) - .field("principal", &self.principal) - .finish() - } -} -/// Reconnect-callback payload. Dispatched once per successful reconnect -/// so consumers can record metrics or surface UI hints. -#[derive(Debug, Clone)] -pub struct ReconnectEvent { - /// Server-issued connection id of the FRESH connection (different - /// from the dropped one). - pub connection_id: ConnectionId, - /// Number of session bindings replayed. - pub sessions_replayed: usize, - /// Reconnect attempt index (1 for the first reconnect). - pub attempt: u32, -} -/// Boxed reconnect callback. -pub type ReconnectCallback = Box; -/// Boxed disconnect callback, fired when the live socket drops (before a -/// reconnect attempt) and on a terminal close. -pub type DisconnectCallback = Box; -/// Boxed connect callback, fired once on the initial successful connect, -/// before the reader actor task spawns. It therefore strictly happens-before -/// any disconnect/reconnect callback, so a connect/disconnect pair can never -/// be observed out of order (e.g. a readiness marker resurrected after the -/// socket has already dropped). -pub type ConnectCallback = Box; -/// A live (or reconnecting) connection to the server. -/// -/// Cheap to clone via the `Arc` returned from -/// [`crate::HubConnectionPool::get_or_connect`]. Methods on the inner -/// `HubConnection` are `&self` so multiple consumers can share the -/// same instance without external locking. -/// -/// Dropping the last `Arc` runs [`Drop`], which sends -/// a stop signal to the connection actor; the actor drains every -/// in-flight response waiter with [`ClientError::NetworkError`] and -/// exits asynchronously. [`Self::request_shutdown`] triggers the -/// same stop-and-drain sequence without giving up the `Arc`. -pub struct HubConnection { - inner: Arc, -} -impl std::fmt::Debug for HubConnection { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("HubConnection") - .field("key", &self.inner.key) - .field("kind", &self.inner.kind) - .finish_non_exhaustive() - } -} -/// Configuration for a [`HubConnection`]. -/// -/// Consumed by [`HubConnection::connect`]; not `Clone` because the -/// only path that wants a copy is the pool, and the pool builds a -/// fresh config per attempt rather than cloning. -pub struct ConnectionConfig { - /// `ws://` or `wss://` URL of the server. - pub url: Url, - pub credential: Arc, - /// Connection role announced in the hello frame. - pub kind: ConnectionKind, - /// Optional reconnect-event callback. - pub on_reconnect: Option>, - /// Optional disconnect callback, fired when the live socket drops or the - /// server sends a terminal close. - pub on_disconnect: Option>, - /// Optional connect callback, fired once on the initial successful connect - /// before the actor starts (so it happens-before any disconnect/reconnect). - pub on_connect: Option>, - /// Stable server identity sent in the hello frame. Only meaningful - /// for [`ConnectionKind::ToolServer`] connections. - pub server_id: Option, - /// One-line server description for `servers.list`. - pub server_description: Option, - /// Opaque metadata surfaced in `ServerInfo.metadata`. - pub server_metadata: Option, - /// Optional override for the outbound mpsc bound. `None` uses the - /// crate default (matched to the server's per-actor outbound - /// buffer). Tests use this to exercise the - /// bounded-wait fast-fail path without flooding production-sized - /// buffers. - pub outbound_buffer: Option, - /// Optional tuning knobs (ping cadence, liveness deadline, reconnect - /// backoff). `ConnectionTuning::default()` keeps every historical - /// default. - pub tuning: ConnectionTuning, - /// When set, attached as an extra access header on every - /// (re)connect, unconditionally. Harmless when the peer ignores it. - pub alpha_test_key: Option, - /// Allow a plaintext `ws://` connection to a non-loopback host. - /// Only enable when the transport is otherwise secured (e.g. a - /// private network or TLS-terminating proxy) — otherwise the bearer - /// credential crosses the network in cleartext. - pub allow_insecure_ws: bool, - /// Optional weak handle to the owning pool, set by - /// [`crate::HubConnectionPool::get_or_connect`]. On a fatal - /// handshake-auth failure the reconnect driver evicts its own pool - /// entry through this so the next caller opens a fresh socket. - /// `None` for the unpooled [`HubConnection::connect`] path (tests / - /// one-shot) — nothing to evict. Weak so the pool↔connection edge - /// is not an ownership cycle. - pub on_fatal: Option>, -} -impl std::fmt::Debug for ConnectionConfig { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ConnectionConfig") - .field("url", &self.url.as_str()) - .field("credential", &self.credential) - .field("kind", &self.kind) - .field("allow_insecure_ws", &self.allow_insecure_ws) - .field("on_reconnect", &self.on_reconnect.is_some()) - .finish() - } -} -struct HubConnectionInner { - key: ConnKey, - kind: ConnectionKind, - credential: Arc, - on_reconnect: Option>, - on_disconnect: Option>, - server_id: Option, - server_description: Option, - server_metadata: Option, - /// Attached as an extra access header on every (re)connect when set. - alpha_test_key: Option, - /// Permit plaintext `ws://` to a non-loopback host (transport otherwise secured). - allow_insecure_ws: bool, - /// See [`ConnectionConfig::on_fatal`]. - on_fatal: Option>, - /// Resolved reconnect backoff schedule (configured override or the - /// built-in table). Resolved once at connect; shared per reconnect. - reconnect_backoff: Arc<[Duration]>, - /// Outbound frames waiting to be written. Filled by `send_*` - /// helpers; drained by the writer half of the actor. - outbound_tx: mpsc::Sender, - /// Inbound demux state (response waiters + session inboxes). - demux: Arc, - /// Refcounted bound-session set. Used by the reconnect path to - /// re-issue `register_session` for every still-live session. - bound_sessions: Arc>, - /// Cached server-issued `connection_id`. Updated on every (re)connect. - connection_id: Arc>>, - /// Optional capabilities the server advertised in the most recent - /// `hello_ack` (wire method strings). Refreshed on every (re)connect - /// handshake. Empty when the ack carried none — on the wire that is - /// indistinguishable from a server predating the field, so - /// [`HubConnection::supports`] reports unknown in that case. - hello_capabilities: parking_lot::RwLock>, - /// Monotonically-increasing JSON-RPC request id counter. - next_request_id: std::sync::atomic::AtomicU64, - /// Cancelled by the actor task once it has fully exited so - /// `await_shutdown` resolves promptly. `CancellationToken` has - /// persistent semantics so a wait that arrives AFTER the actor - /// has already cancelled the token still wakes immediately. - shutdown: CancellationToken, - /// Stops the actor on `Drop`. - stop_tx: mpsc::Sender<()>, - reconnect_tx: mpsc::Sender<()>, - early_notif_rx: parking_lot::Mutex>>, - health: ConnHealth, - writer_error: WriteErrorSlot, -} -type WsStream = WebSocketStream>; -impl HubConnection { - /// Open a brand-new [`HubConnection`] and spawn its actor task. - /// - /// The pool is the canonical caller; outside callers MAY use this - /// for tests or one-shot programs but lose pool dedup. - pub async fn connect(config: ConnectionConfig) -> Result, ClientError> { - let initial_cred = config.credential.current(); - let key = ConnKey { - url: config.url.as_str().to_owned(), - principal: config.credential.principal_key(), - }; - let ws_ping_interval = resolve_ws_ping_interval(config.tuning.ws_ping_interval); - let ws_liveness_deadline = - resolve_ws_liveness_deadline(config.tuning.ws_liveness_deadline, ws_ping_interval); - if ws_liveness_deadline <= ws_ping_interval { - warn!( - ?ws_liveness_deadline, - ?ws_ping_interval, - "ws liveness deadline is not greater than the keepalive ping interval; healthy idle connections will be killed and reconnected every window" - ); - } - let reconnect_backoff = resolve_reconnect_backoff(config.tuning.reconnect_backoff); - let buffer = config.outbound_buffer.unwrap_or(OUTBOUND_BUFFER); - let (outbound_tx, outbound_rx) = mpsc::channel::(buffer); - let (stop_tx, stop_rx) = mpsc::channel::<()>(1); - let (reconnect_tx, reconnect_rx) = mpsc::channel::<()>(1); - let demux = Arc::new(Demux::with_outbound(outbound_tx.clone())); - let bound_sessions = Arc::new(RefCountedSet::::new()); - let connection_id = Arc::new(Mutex::new(None)); - let shutdown = CancellationToken::new(); - let ws = open_socket( - &config.url, - &initial_cred, - config.kind, - config.alpha_test_key.as_deref(), - config.allow_insecure_ws, - ) - .await?; - let (sink, stream) = ws.split(); - let (sink, stream, ack) = run_handshake( - sink, - stream, - config.kind, - config.server_id.clone(), - config.server_description.clone(), - config.server_metadata.clone(), - ) - .await?; - *connection_id.lock().await = Some(ack.connection_id.clone()); - info!( - url = % config.url, connection_id = % ack.connection_id, - "server connection established" - ); - if let Some(cb) = &config.on_connect { - cb(); - } - let early_notif_rx = parking_lot::Mutex::new(match config.kind { - ConnectionKind::ToolServer => Some(demux.subscribe_notifications()), - _ => None, - }); - let writer_error: WriteErrorSlot = Arc::new(parking_lot::Mutex::new(None)); - let inner = Arc::new(HubConnectionInner { - key, - kind: config.kind, - credential: config.credential, - on_reconnect: config.on_reconnect.clone(), - on_disconnect: config.on_disconnect.clone(), - server_id: config.server_id, - server_description: config.server_description, - server_metadata: config.server_metadata, - alpha_test_key: config.alpha_test_key, - allow_insecure_ws: config.allow_insecure_ws, - on_fatal: config.on_fatal, - reconnect_backoff, - outbound_tx, - demux: demux.clone(), - bound_sessions: bound_sessions.clone(), - connection_id, - hello_capabilities: parking_lot::RwLock::new(ack.capabilities), - next_request_id: std::sync::atomic::AtomicU64::new(1), - shutdown, - stop_tx, - reconnect_tx, - early_notif_rx, - health: ConnHealth::new(), - writer_error: writer_error.clone(), - }); - let (writer_ctl_tx, writer_ctl_rx) = - mpsc::channel::>>(2); - let (writer_stop_tx, writer_stop_rx) = mpsc::channel::<()>(1); - let writer_handle = tokio::spawn(run_writer( - sink, - outbound_rx, - writer_ctl_rx, - writer_stop_rx, - ws_ping_interval, - writer_error, - )); - let reader_inner = inner.clone(); - tokio::spawn(run_reader_actor( - reader_inner, - stream, - stop_rx, - reconnect_rx, - writer_ctl_tx, - writer_stop_tx, - writer_handle, - config.url, - ws_liveness_deadline, - )); - Ok(Arc::new(Self { inner })) - } - /// Pool dedup key for this connection. - pub fn key(&self) -> &ConnKey { - &self.inner.key - } - /// Connection role. - pub fn kind(&self) -> ConnectionKind { - self.inner.kind - } - /// Stable identity of this connection's actor state. Lets the pool - /// evict by identity so a connection only ever forgets its own slot. - pub(crate) fn actor_id(&self) -> usize { - Arc::as_ptr(&self.inner) as *const () as usize - } - /// Server-issued connection id of the most recently established - /// (post-handshake) socket. During a reconnect gap this still names the - /// dropped connection until the next handshake + replay completes. - pub async fn connection_id(&self) -> Option { - self.inner.connection_id.lock().await.clone() - } - /// Whether the server advertised `capability` (a wire method string, - /// e.g. `"session_attach_server"`) in the CURRENT connection's - /// `hello_ack`. - /// - /// - `Some(true)`: advertised. - /// - `Some(false)`: the ack carried a non-empty capability list that - /// does not include it. - /// - `None`: the ack advertised nothing — servers predating the - /// `capabilities` field are indistinguishable from an empty list, so - /// support is unknown and callers should probe per call. - pub fn supports(&self, capability: &str) -> Option { - let caps = self.inner.hello_capabilities.read(); - if caps.is_empty() { - return None; - } - Some(caps.iter().any(|c| c == capability)) - } - /// Demux (used by the server-side run loop to register session - /// inboxes). Cheap to clone (Arc bump). - pub fn demux(&self) -> Arc { - self.inner.demux.clone() - } - pub(crate) fn take_early_notifications(&self) -> Option> { - self.inner.early_notif_rx.lock().take() - } - pub(crate) fn force_reconnect(&self) { - let _ = self.inner.reconnect_tx.try_send(()); - } - /// Future that resolves once the connection actor has shut down. - pub async fn await_shutdown(&self) { - self.inner.shutdown.cancelled().await; - } - /// Signal the connection actor to begin shutdown. The actor - /// drains its in-flight waiters with `NetworkError` and exits; - /// the outbound channel closes shortly after, so subsequent - /// [`Self::send_outbound`] calls return - /// [`ClientError::NetworkError`]. [`Self::await_shutdown`] - /// resolves once the actor task has terminated. - /// - /// Idempotent: redundant calls are no-ops. Equivalent to - /// dropping the last `Arc`, but lets a holder - /// trigger shutdown without giving up its reference. - pub fn request_shutdown(&self) { - let _ = self.inner.stop_tx.try_send(()); - } - /// Increment the refcount on `session_id`. The session is tracked - /// locally for reconnect-replay; the server learns about it via - /// `serve` (auto-registration on the server side). - pub fn track_session(&self, session_id: SessionId) { - self.inner.bound_sessions.increment(session_id); - } - /// Decrement the refcount on `session_id`. Removes tracking when - /// the last borrower drops. - pub fn untrack_session(&self, session_id: &SessionId) { - self.inner.bound_sessions.decrement(session_id); - } - /// Send a JSON-RPC request and await the response. - /// - /// The waiter is registered before the frame is sent so a fast response - /// can never arrive before its waiter exists, and is reclaimed on send - /// failure (via [`WaiterGuard`]) so a call that never reached the wire - /// cannot leak a parked waiter across a reconnect episode. - pub async fn call_request

( - &self, - request_id: kigi_tool_protocol::RequestId, - request: &JsonRpcRequest

, - ) -> Result - where - P: serde::Serialize, - { - let text = serde_json::to_string(request)?; - let (tx, rx) = oneshot::channel(); - self.inner - .demux - .register_response_waiter(request_id.clone(), tx); - let _guard = WaiterGuard { - demux: &self.inner.demux, - request_id: &request_id, - }; - self.send_outbound(text).await?; - rx.await? - } - /// Send a JSON-RPC request and await the response, bounded by `timeout`. - pub async fn call_request_with_timeout

( - &self, - request_id: kigi_tool_protocol::RequestId, - request: &JsonRpcRequest

, - timeout: Duration, - ) -> Result - where - P: serde::Serialize, - { - self.call_request_with_deadline(request_id, request, timeout) - .await - .map_err(ClientError::from) - } - async fn call_request_with_deadline

( - &self, - request_id: kigi_tool_protocol::RequestId, - request: &JsonRpcRequest

, - timeout: Duration, - ) -> Result - where - P: serde::Serialize, - { - let text = - serde_json::to_string(request).map_err(|e| DeadlineCallError::Other(e.into()))?; - let (tx, rx) = oneshot::channel(); - self.inner - .demux - .register_response_waiter(request_id.clone(), tx); - let _guard = WaiterGuard { - demux: &self.inner.demux, - request_id: &request_id, - }; - self.send_outbound(text) - .await - .map_err(DeadlineCallError::Other)?; - match tokio::time::timeout(timeout, rx).await { - Ok(Ok(result)) => result.map_err(DeadlineCallError::Other), - Ok(Err(recv_err)) => Err(DeadlineCallError::Other(recv_err.into())), - Err(_elapsed) => Err(DeadlineCallError::TimedOut(timeout)), - } - } - /// Send a fully-formed JSON text frame onto the outbound channel. - /// Used by the server-side handler when replying to a - /// `tool_call_request` (the response flows out without going - /// through a waiter). - pub async fn send_outbound(&self, text: String) -> Result<(), ClientError> { - match self.inner.outbound_tx.try_send(text) { - Ok(()) => Ok(()), - Err(mpsc::error::TrySendError::Full(text)) => { - match tokio::time::timeout( - Duration::from_millis(250), - self.inner.outbound_tx.send(text), - ) - .await - { - Ok(Ok(())) => Ok(()), - Ok(Err(_)) => Err(ClientError::NetworkError( - "outbound channel closed".to_owned(), - )), - Err(_) => Err(ClientError::BackpressureError( - "outbound mpsc full beyond bounded wait".to_owned(), - )), - } - } - Err(mpsc::error::TrySendError::Closed(_)) => Err(ClientError::NetworkError( - "outbound channel closed".to_owned(), - )), - } - } - /// Non-blocking enqueue for synchronous drop paths that cannot - /// `.await` (e.g. `RemoteCallStream::Drop` cancel-on-drop). A full - /// or closed channel returns `Err` and the caller abandons the - /// frame — best-effort, mirroring the heartbeat-pong drop discipline. - pub(crate) fn try_send_outbound(&self, text: String) -> Result<(), ClientError> { - match self.inner.outbound_tx.try_send(text) { - Ok(()) => Ok(()), - Err(mpsc::error::TrySendError::Full(_)) => Err(ClientError::BackpressureError( - "outbound mpsc full".to_owned(), - )), - Err(mpsc::error::TrySendError::Closed(_)) => Err(ClientError::NetworkError( - "outbound channel closed".to_owned(), - )), - } - } - /// Allocate a fresh request id. Monotonic per-connection. - /// - /// Returns `Err` only if a future-added `RequestId` invariant - /// rejects the formatted `c{value}` string (today the only - /// failure path is the empty-string check, which `format!` cannot - /// produce). Callers in non-fallible contexts should propagate - /// the error rather than panic. - pub fn try_alloc_request_id(&self) -> Result { - let value = self - .inner - .next_request_id - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - kigi_tool_protocol::RequestId::new(format!("c{value}")).map_err(ClientError::from) - } - /// Number of sessions currently bound to this connection. - /// Stable observable for monitoring and tests; not on the hot path. - pub fn bound_session_count(&self) -> usize { - self.inner.bound_sessions.len() - } - /// Send a `serve` frame: full tool snapshot for a session. - /// - /// Idempotent: re-sending replaces the tool set. The server diffs - /// against the previous snapshot and emits `tools_changed` to - /// subscribed harnesses. - pub async fn serve( - &self, - session_id: SessionId, - params: kigi_tool_protocol::ServeParams, - ) -> Result { - let mut last_err: Option = None; - for attempt in 1..=SERVE_MAX_ATTEMPTS { - let request_id = self.try_alloc_request_id()?; - let req = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: Some(session_id.clone()), - method: Method::Serve.as_wire_str().to_owned(), - params: ¶ms, - }; - match self - .call_request_with_deadline(request_id, &req, SERVE_ATTEMPT_TIMEOUT) - .await - { - Ok(resp) => { - return match resp.outcome { - ResponseOutcome::Result(value) => serde_json::from_value(value) - .map_err(|e| ClientError::Serde(e.to_string())), - ResponseOutcome::Error(err) => Err(ClientError::from_jsonrpc_error(err)), - }; - } - Err(DeadlineCallError::TimedOut(timeout)) => { - crate::metrics::serve_replay_timeout(); - warn!( - % session_id, attempt, ? timeout, - "serve attempt timed out; will retry" - ); - last_err = Some(DeadlineCallError::TimedOut(timeout).into()); - } - Err(DeadlineCallError::Other(e)) => return Err(e), - } - } - warn!( - % session_id, attempts = SERVE_MAX_ATTEMPTS, - "serve timed out every bounded attempt; forcing reconnect to restart replay" - ); - self.force_reconnect(); - Err(last_err.unwrap_or_else(|| { - ClientError::NetworkError("serve failed after bounded retries".to_owned()) - })) - } -} -impl Drop for HubConnection { - fn drop(&mut self) { - let _ = self.inner.stop_tx.try_send(()); - } -} -/// True iff `url`'s host is one of the canonical loopback names. Case -/// insensitive on the hostname; IP literals match the standard loopback -/// addresses for IPv4 and IPv6. -pub(crate) fn host_is_loopback(url: &Url) -> bool { - use std::net::{Ipv4Addr, Ipv6Addr}; - match url.host() { - Some(url::Host::Ipv4(ip)) => ip == Ipv4Addr::LOCALHOST, - Some(url::Host::Ipv6(ip)) => ip == Ipv6Addr::LOCALHOST, - Some(url::Host::Domain(host)) => host.eq_ignore_ascii_case("localhost"), - None => false, - } -} -/// Open a fresh `ws://` / `wss://` socket. No handshake yet. -/// -/// Refuses to send the credential over `ws://` to any non-loopback host -/// so the bearer token never crosses the network in plaintext. Local -/// loopback (`127.0.0.1`, `::1`, `localhost`) is the explicit exception -/// for development and local-proxy use; every other host must be -/// reached over `wss://`. -async fn open_socket( - url: &Url, - credential: &AuthCredential, - kind: ConnectionKind, - alpha_test_key: Option<&str>, - allow_insecure_ws: bool, -) -> Result { - let is_plaintext_remote = url.scheme() != "wss" && !host_is_loopback(url); - if is_plaintext_remote && !allow_insecure_ws { - return Err(ClientError::InsecureScheme { url: url.clone() }); - } - if is_plaintext_remote { - warn!( - host = % url.host_str().unwrap_or(""), - "opening server connection over plaintext ws:// (allow_insecure_ws=true); bearer crosses the network in cleartext" - ); - } - let mut connect_url = url.clone(); - let expected_role = match kind { - ConnectionKind::Harness => "harness", - ConnectionKind::ToolServer => "tool_server", - }; - if let Some(existing) = connect_url - .query_pairs() - .find(|(k, _)| k == "role") - .map(|(_, v)| v.to_string()) - { - if existing != expected_role { - return Err(ClientError::InvalidConfig(format!( - "URL query parameter role={existing} conflicts with ConnectionKind::{kind:?} (expected role={expected_role})" - ))); - } - } else { - connect_url - .query_pairs_mut() - .append_pair("role", expected_role); - } - let mut request = connect_url - .as_str() - .into_client_request() - .map_err(|e| ClientError::InvalidConfig(format!("invalid ws request: {e}")))?; - let headers = request.headers_mut(); - for (name, value) in credential.upgrade_headers() { - let header_name: HeaderName = name; - let header_value: HeaderValue = HeaderValue::from_str(&value) - .map_err(|e| ClientError::InvalidConfig(format!("invalid auth header value: {e}")))?; - headers.insert(header_name, header_value); - } - let _ = alpha_test_key; - kigi_tracing::http_client::attach_trace_to_http_request(headers); - let (ws, _resp) = connect_async(request) - .await - .map_err(ClientError::from_handshake_error)?; - Ok(ws) -} -/// Drive the hello / hello_ack exchange and hand back the (sink, -/// stream) pair for steady-state use. -async fn run_handshake( - mut sink: SplitSink, - mut stream: SplitStream, - kind: ConnectionKind, - server_id: Option, - server_description: Option, - server_metadata: Option, -) -> Result< - ( - SplitSink, - SplitStream, - kigi_tool_protocol::HelloAckMsg, - ), - ClientError, -> { - let ack = send_hello( - &mut sink, - &mut stream, - kind, - server_id, - server_description, - server_metadata, - ) - .await?; - Ok((sink, stream, ack)) -} -/// Outcome of the connected-phase loop. -enum ConnectedExit { - /// Stop signal — actor terminates. - Stop, - /// Socket closed / errored — actor enters reconnect. - SocketClosed(DisconnectCause), - /// Server sent a close frame with a code that means "do not reconnect" - /// (e.g. force eviction, session expired, admin disconnect). - TerminalClose(u16), -} -/// Current Unix time in milliseconds (saturating to 0 before the epoch). -fn now_unix_millis() -> u64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_millis() as u64 -} -/// Decode an inbound text frame. Returns the serialized [`PongFrame`] -/// to send back when the frame is an app-level server `ping`; otherwise -/// routes the frame through the demux and returns `None`. -fn route_or_pong(inner: &HubConnectionInner, text: &str) -> Option { - match serde_json::from_str::(text) { - Ok(value) => { - if value.get("method").and_then(Value::as_str) == Some(Method::Ping.as_wire_str()) { - serde_json::to_string(&PongFrame::new(now_unix_millis())).ok() - } else { - let _ = inner.demux.route(value); - None - } - } - Err(e) => { - warn!(?e, "discarding unparseable inbound text frame"); - None - } - } -} -/// Map a websocket close frame's code to the connected-phase exit. Close -/// codes 4100-4199 are terminal (the server intentionally ended the -/// connection: eviction, session expiry, admin disconnect, rate limit). -/// The range is deliberately wide so new terminal codes added server-side -/// are recognised without a client update. -fn exit_for_close_code(code: Option) -> ConnectedExit { - match code { - Some(code) if (4100..4200).contains(&code) => ConnectedExit::TerminalClose(code), - _ => ConnectedExit::SocketClosed(DisconnectCause::CloseFrame(code)), - } -} -/// Classify why the inbound stream ended, preferring a write error the -/// writer task recorded over what the reader observed. -/// -/// Best-effort: the writer task populates `writer_error` asynchronously -/// after its send fails, so the reader can observe the resulting stream -/// EOF/error and classify it here *before* the slot is set. In that -/// (telemetry-only) race a genuine write-side failure is reported as -/// `eof` / `transport_read_error` instead of `transport_write_error`. -fn classify_stream_end(inner: &HubConnectionInner, read_error: Option) -> DisconnectCause { - if let Some(detail) = inner.writer_error.lock().take() { - return DisconnectCause::WriteError(detail); - } - match read_error { - Some(detail) => DisconnectCause::ReadError(detail), - None => DisconnectCause::Eof, - } -} -/// Control messages handed to the dedicated writer task. -/// -/// The reader is the sole reconnect driver; it `Pause`s the writer the -/// instant the socket is known dead so no buffered frame is dequeued -/// onto the corpse, then `Resume`s it with the fresh sink once the -/// handshake completes. Carried on a cap-2 channel so a `Pause` is never -/// dropped. -enum WriterControl { - /// Socket is dead; stop draining `outbound_rx` (frames stay buffered). - Pause, - /// Reconnected; install the fresh sink and resume draining. - Resume(S), -} -/// Dedicated writer task: owns the sink, drains `outbound_rx`, and fires -/// the keepalive ping (`ping_period`) — but only while `live`. Between a `Pause` and -/// the matching `Resume` it parks on the control/stop channels only, so -/// frames enqueued during the reconnect gap stay buffered in -/// `outbound_rx` and flush after `Resume` (no multi-frame loss; the -/// single in-flight frame whose `send` fails is the only loss, matching -/// the pre-split worst case). -/// -/// Generic over the sink so it can be unit-tested with an in-memory sink -/// without a live socket. -async fn run_writer( - mut sink: S, - mut outbound_rx: mpsc::Receiver, - mut writer_ctl_rx: mpsc::Receiver>, - mut writer_stop_rx: mpsc::Receiver<()>, - ping_period: Duration, - write_error: WriteErrorSlot, -) where - S: futures::Sink + Unpin, - S::Error: std::fmt::Display, -{ - let mut ping_interval = tokio::time::interval(ping_period); - ping_interval.tick().await; - let mut live = true; - loop { - tokio::select! { - biased; _ = writer_stop_rx.recv() => break, ctl = writer_ctl_rx.recv() => - match ctl { Some(WriterControl::Pause) => live = false, - Some(WriterControl::Resume(new_sink)) => { sink = new_sink; live = true; - write_error.lock().take(); ping_interval = - tokio::time::interval(ping_period); ping_interval.tick(). await; } None => - break, }, _ = ping_interval.tick(), if live => { if let Err(e) = sink - .send(Message::Ping(Vec::new().into())). await { * write_error.lock() = - Some(format!("ping send failed: {e}")); crate - ::metrics::writer_sink_send_error(); live = false; } } outbound = outbound_rx - .recv(), if live => match outbound { Some(text) => { if let Err(e) = sink - .send(Message::Text(text.into())). await { * write_error.lock() = - Some(format!("frame send failed: {e}")); crate - ::metrics::writer_sink_send_error(); live = false; } } None => break, }, - } - } -} -/// Invoke the optional disconnect callback (best-effort, sync). -fn fire_on_disconnect(inner: &HubConnectionInner) { - if let Some(cb) = &inner.on_disconnect { - cb(); - } -} -/// Reader half of the split actor: owns the stream, routes inbound -/// frames, and drives reconnect. Never touches the sink — it asks the -/// writer task to `Pause`/`Resume` instead. -async fn run_reader_actor( - inner: Arc, - mut stream: SplitStream, - mut stop_rx: mpsc::Receiver<()>, - mut reconnect_rx: mpsc::Receiver<()>, - writer_ctl_tx: mpsc::Sender>>, - writer_stop_tx: mpsc::Sender<()>, - writer_handle: tokio::task::JoinHandle<()>, - url: Url, - liveness_deadline: Duration, -) { - let mut attempt: u32 = 0; - let mut connected_at = Instant::now(); - 'actor: loop { - match run_reader_phase( - inner.as_ref(), - &mut stream, - &mut stop_rx, - &mut reconnect_rx, - liveness_deadline, - ) - .await - { - ConnectedExit::Stop => break, - ConnectedExit::TerminalClose(code) => { - info!(code, url = % url, "server sent terminal close; not reconnecting"); - fire_on_disconnect(inner.as_ref()); - inner.demux.drain_waiters_with(|| { - ClientError::Closed(format!("server terminal close (code {code})")) - }); - inner.demux.drain_progress(); - break; - } - ConnectedExit::SocketClosed(cause) => { - let detected_at = Instant::now(); - let health = inner.health.snapshot(); - let prev_connection_id = inner.connection_id.lock().await.clone(); - let outage = OutageInfo { - prev_connection_id, - prev_connection_duration_ms: detected_at - .duration_since(connected_at) - .as_millis() as u64, - last_inbound: health.last_inbound, - detect_ms: detected_at.duration_since(health.last_inbound).as_millis() as u64, - since_last_probe_monotonic_ms: health.since_last_probe_monotonic_ms, - since_last_probe_wall_ms: health.since_last_probe_wall_ms, - clock_jump_ms: health.clock_jump_ms, - cause, - }; - warn!( - url = % url, cause = outage.cause.label(), close_code = ? outage - .cause.close_code(), error_detail = ? outage.cause.detail(), - connection_id = ? outage.prev_connection_id, - prev_connection_duration_ms = outage.prev_connection_duration_ms, - detect_ms = outage.detect_ms, since_last_probe_monotonic_ms = outage - .since_last_probe_monotonic_ms, since_last_probe_wall_ms = outage - .since_last_probe_wall_ms, clock_jump_ms = outage.clock_jump_ms, - "server connection lost; scheduling reconnect" - ); - fire_on_disconnect(inner.as_ref()); - if writer_ctl_tx.send(WriterControl::Pause).await.is_err() { - break; - } - inner.demux.drain_waiters_with(|| { - ClientError::NetworkError("socket dropped during in-flight call".to_owned()) - }); - inner.demux.drain_progress(); - let mut backoff_total = Duration::ZERO; - loop { - attempt = attempt.saturating_add(1); - let backoff = backoff_for(attempt, &inner.reconnect_backoff); - info!( - ? backoff, attempt, url = % url, "reconnecting server connection" - ); - tokio::select! { - _ = stop_rx.recv() => break 'actor, _ = sleep(backoff) => {} - } - backoff_total += backoff; - let reconnect_start = std::time::Instant::now(); - let attempt_budget = reconnect_attempt_budget(liveness_deadline); - let outcome = tokio::select! { - _ = stop_rx.recv() => break 'actor, outcome = - tokio::time::timeout(attempt_budget, reconnect_and_replay(inner - .as_ref(), & url, attempt, & outage, backoff_total,),) => outcome - .unwrap_or_else(| _elapsed | { - Err(ClientError::NetworkError(format!("reconnect attempt timed out after {attempt_budget:?}"))) - }), - }; - match outcome { - Ok((new_sink, new_stream)) => { - let elapsed = reconnect_start.elapsed().as_secs_f64(); - crate::metrics::reconnect_succeeded(); - crate::metrics::reconnect_duration_observe(elapsed); - inner.health.reset(); - inner.writer_error.lock().take(); - connected_at = Instant::now(); - drain_reconnect_signals(&mut reconnect_rx); - stream = new_stream; - if writer_ctl_tx - .send(WriterControl::Resume(new_sink)) - .await - .is_err() - { - break 'actor; - } - crate::metrics::reconnect_writer_resume(); - break; - } - Err(ClientError::HandshakeAuthFailed { status }) => { - warn!( - status, - attempt, - "reconnect rejected with handshake auth failure; evicting pool entry and stopping" - ); - crate::metrics::reconnect_failed("handshake_auth"); - inner.demux.drain_waiters_with(|| { - ClientError::AuthError(format!( - "server rejected reconnect handshake (HTTP {status})" - )) - }); - inner.demux.drain_progress(); - if let Some(pool) = inner.on_fatal.as_ref().and_then(Weak::upgrade) { - let own_id = Arc::as_ptr(&inner) as *const () as usize; - pool.forget_if(&inner.key, move |conn| conn.actor_id() == own_id); - } - break 'actor; - } - Err(err) => { - crate::metrics::reconnect_failed("transport"); - warn!( - ?err, - attempt, - cause = outage.cause.label(), - backoff_total_ms = backoff_total.as_millis() as u64, - "reconnect attempt failed; will retry" - ); - } - } - } - } - } - } - inner - .demux - .drain_waiters_with(|| ClientError::NetworkError("connection actor exited".to_owned())); - inner.demux.drain_progress(); - let _ = writer_stop_tx.send(()).await; - drop(writer_ctl_tx); - drop(stop_rx); - drop(stream); - if let Err(e) = writer_handle.await { - warn!(?e, "writer task panicked during shutdown"); - } - inner.shutdown.cancel(); -} -fn drain_reconnect_signals(reconnect_rx: &mut mpsc::Receiver<()>) { - while reconnect_rx.try_recv().is_ok() {} -} -/// Reader-only steady-state loop for the split actor: drives the inbound -/// half but never writes (app-level pongs route through `outbound_tx`; WS -/// pings are auto-answered by tungstenite on poll). -/// -/// Enforces the inbound-liveness deadline: no inbound frame of any kind for -/// the deadline window (default 2.5× the ping cadence, see -/// [`resolve_ws_liveness_deadline`]) means the transport is silently dead -/// (snapshot-restored VM, NAT/LB flow expiry), so exit via -/// [`ConnectedExit::SocketClosed`] onto the normal reconnect path. The -/// deadline runs only in this phase and re-arms on every (re)entry. -/// -/// Generic over the stream for in-memory unit tests, mirroring -/// [`run_writer`]. -async fn run_reader_phase( - inner: &HubConnectionInner, - stream: &mut S, - stop_rx: &mut mpsc::Receiver<()>, - reconnect_rx: &mut mpsc::Receiver<()>, - liveness_deadline: Duration, -) -> ConnectedExit -where - S: Stream> + Unpin, -{ - let mut clock_probe = tokio::time::interval(CLOCK_PROBE_INTERVAL); - clock_probe.tick().await; - let deadline = sleep(liveness_deadline); - tokio::pin!(deadline); - loop { - tokio::select! { - biased; _ = stop_rx.recv() => return ConnectedExit::Stop, _ = reconnect_rx - .recv() => { info!("forced reconnect requested; dropping current socket"); - return ConnectedExit::SocketClosed(DisconnectCause::Forced); } msg = stream - .next() => { if matches!(msg, Some(Ok(ref m)) if ! matches!(m, - Message::Close(_))) { inner.health.record_inbound(); } match msg { - Some(Ok(msg)) => { let now = tokio::time::Instant::now(); let rearm = now - .checked_add(liveness_deadline).unwrap_or_else(|| now + - Duration::from_secs(86400 * 365 * 30)); deadline.as_mut().reset(rearm); match - msg { Message::Text(text) => { if let Some(pong_text) = route_or_pong(inner, - text.as_ref()) && inner.outbound_tx.try_send(pong_text).is_err() { crate - ::metrics::heartbeat_pong_dropped(); } } Message::Ping(_) | Message::Pong(_) - | Message::Frame(_) => {} Message::Binary(_) => { - warn!("server sent binary frame; ignoring"); } Message::Close(frame) => { - return exit_for_close_code(frame.map(| f | f.code.into())); } } } - Some(Err(e)) => { return - ConnectedExit::SocketClosed(classify_stream_end(inner, Some(e - .to_string()),)); } None => { return - ConnectedExit::SocketClosed(classify_stream_end(inner, None)); } } } _ = - clock_probe.tick() => inner.health.refresh_clock(), _ = & mut deadline => { - crate ::metrics::liveness_deadline_expired(); warn!(? liveness_deadline, - "no inbound frame within the liveness deadline; declaring the socket dead and reconnecting"); - return ConnectedExit::SocketClosed(DisconnectCause::LivenessDeadline); } - } - } -} -/// Reconnect once and replay every session binding + tool registration. -async fn reconnect_and_replay( - inner: &HubConnectionInner, - url: &Url, - attempt: u32, - outage: &OutageInfo, - backoff_total: Duration, -) -> Result<(SplitSink, SplitStream), ClientError> { - let fresh_cred = inner.credential.current(); - let ws = open_socket( - url, - &fresh_cred, - inner.kind, - inner.alpha_test_key.as_deref(), - inner.allow_insecure_ws, - ) - .await?; - let (sink, stream) = ws.split(); - let (mut sink, mut stream, mut ack) = run_handshake( - sink, - stream, - inner.kind, - inner.server_id.clone(), - inner.server_description.clone(), - inner.server_metadata.clone(), - ) - .await?; - let sessions = inner.bound_sessions.snapshot_keys(); - if inner.kind == ConnectionKind::Harness { - for sid in &sessions { - let req = kigi_tool_protocol::JsonRpcRequest { - jsonrpc: kigi_tool_protocol::JsonRpcVersion, - id: kigi_tool_protocol::JsonRpcId::new_uuid_v7(), - session_id: Some(sid.clone()), - method: Method::SessionOpen.as_wire_str().to_owned(), - params: kigi_tool_protocol::SessionOpenParams { - resume: false, - last_seq: None, - }, - }; - if let Ok(text) = serde_json::to_string(&req) { - let _ = SinkExt::send(&mut sink, Message::Text(text.into())).await; - let _ = tokio::time::timeout(Duration::from_secs(5), StreamExt::next(&mut stream)) - .await; - } - } - } - let sessions_replayed = sessions.len(); - let silent_gap_ms = outage.last_inbound.elapsed().as_millis() as u64; - info!( - attempt, sessions_replayed, cause = outage.cause.label(), close_code = ? outage - .cause.close_code(), error_detail = ? outage.cause.detail(), prev_connection_id = - ? outage.prev_connection_id, connection_id = % ack.connection_id, - prev_connection_duration_ms = outage.prev_connection_duration_ms, silent_gap_ms, - detect_ms = outage.detect_ms, backoff_total_ms = backoff_total.as_millis() as - u64, since_last_probe_monotonic_ms = outage.since_last_probe_monotonic_ms, - since_last_probe_wall_ms = outage.since_last_probe_wall_ms, clock_jump_ms = - outage.clock_jump_ms, "server reconnect succeeded" - ); - crate::metrics::reconnect_cause(outage.cause.label()); - crate::metrics::reconnect_gap_observe(silent_gap_ms as f64 / 1_000.0); - *inner.connection_id.lock().await = Some(ack.connection_id.clone()); - *inner.hello_capabilities.write() = std::mem::take(&mut ack.capabilities); - if let Some(cb) = &inner.on_reconnect { - cb(ReconnectEvent { - connection_id: ack.connection_id, - sessions_replayed, - attempt, - }); - } - Ok((sink, stream)) -} -/// Look up the backoff for `attempt` in `schedule`, clamping past the end -/// to the final (cap) slot. Call sites pass a non-empty schedule (resolved -/// via [`resolve_reconnect_backoff`]); the lookup is nonetheless -/// self-contained — an empty slice yields `Duration::ZERO` rather than -/// panicking. -fn backoff_for(attempt: u32, schedule: &[Duration]) -> Duration { - let idx = (attempt as usize) - .saturating_sub(1) - .min(schedule.len().saturating_sub(1)); - schedule.get(idx).copied().unwrap_or_default() -} -#[cfg(test)] -mod tests { - use super::*; - #[test] - fn backoff_for_follows_exponential_schedule() { - let schedule = default_reconnect_backoff(); - assert_eq!(backoff_for(1, &schedule), Duration::from_millis(100)); - assert_eq!(backoff_for(2, &schedule), Duration::from_millis(200)); - assert_eq!(backoff_for(3, &schedule), Duration::from_millis(500)); - assert_eq!(backoff_for(4, &schedule), Duration::from_millis(1_000)); - assert_eq!(backoff_for(5, &schedule), Duration::from_millis(2_000)); - assert_eq!(backoff_for(6, &schedule), Duration::from_millis(5_000)); - assert_eq!(backoff_for(7, &schedule), Duration::from_millis(10_000)); - } - #[test] - fn backoff_for_caps_at_last_slot() { - let schedule = default_reconnect_backoff(); - let cap = Duration::from_millis(10_000); - assert_eq!(backoff_for(8, &schedule), cap); - assert_eq!(backoff_for(50, &schedule), cap); - assert_eq!(backoff_for(u32::MAX, &schedule), cap); - } - #[test] - fn backoff_for_zero_attempt_uses_first_slot() { - let schedule = default_reconnect_backoff(); - assert_eq!(backoff_for(0, &schedule), Duration::from_millis(100)); - } - #[test] - fn backoff_for_honors_configured_schedule() { - let schedule = resolve_reconnect_backoff(Some(Arc::from([ - Duration::from_millis(5), - Duration::from_millis(15), - ]))); - assert_eq!(backoff_for(1, &schedule), Duration::from_millis(5)); - assert_eq!(backoff_for(2, &schedule), Duration::from_millis(15)); - assert_eq!(backoff_for(3, &schedule), Duration::from_millis(15)); - assert_eq!(backoff_for(99, &schedule), Duration::from_millis(15)); - } - #[test] - fn backoff_for_empty_schedule_is_zero_not_panic() { - assert_eq!(backoff_for(1, &[]), Duration::ZERO); - assert_eq!(backoff_for(0, &[]), Duration::ZERO); - assert_eq!(backoff_for(u32::MAX, &[]), Duration::ZERO); - } - #[test] - fn resolve_reconnect_backoff_falls_back_when_unset_or_empty() { - let from_none = resolve_reconnect_backoff(None); - let from_empty = resolve_reconnect_backoff(Some(Arc::from([]))); - for schedule in [from_none, from_empty] { - assert_eq!(backoff_for(1, &schedule), Duration::from_millis(100)); - assert_eq!(backoff_for(7, &schedule), Duration::from_millis(10_000)); - assert_eq!(backoff_for(99, &schedule), Duration::from_millis(10_000)); - } - } - /// A zero or unset ping interval must resolve to the default. A zero - /// period would otherwise reach `tokio::time::interval`, which panics on - /// `Duration::ZERO`; a positive override is honored verbatim. - #[test] - fn resolve_ws_ping_interval_clamps_zero_and_unset_to_default() { - assert_eq!(resolve_ws_ping_interval(None), DEFAULT_WS_PING_INTERVAL); - assert_eq!( - resolve_ws_ping_interval(Some(Duration::ZERO)), - DEFAULT_WS_PING_INTERVAL - ); - let custom = Duration::from_secs(7); - assert_eq!(resolve_ws_ping_interval(Some(custom)), custom); - } - /// Resolving a zero ping interval to a non-zero default means - /// `tokio::time::interval` can be constructed without panicking. - #[tokio::test] - async fn resolved_zero_ping_interval_builds_interval_without_panic() { - let resolved = resolve_ws_ping_interval(Some(Duration::ZERO)); - assert!(!resolved.is_zero()); - let _interval = tokio::time::interval(resolved); - } - fn bearer_credential() -> AuthCredential { - AuthCredential::bearer("test-token") - } - #[tokio::test] - async fn open_socket_refuses_plaintext_ws_to_remote_host() { - let url = Url::parse("ws://hub.example.com:8080/v1/tools").expect("valid url"); - let credential = bearer_credential(); - match open_socket(&url, &credential, ConnectionKind::Harness, None, false).await { - Err(ClientError::InsecureScheme { url: rejected }) => { - assert_eq!(rejected, url); - } - other => panic!("expected InsecureScheme; got {other:?}"), - } - } - #[tokio::test] - async fn open_socket_allows_plaintext_ws_to_loopback() { - let url = Url::parse("ws://127.0.0.1:1/").expect("valid url"); - let credential = bearer_credential(); - if let Err(ClientError::InsecureScheme { .. }) = - open_socket(&url, &credential, ConnectionKind::Harness, None, false).await - { - panic!("loopback ws:// must not be rejected by the scheme guard") - } - } - #[tokio::test] - async fn open_socket_allows_wss_to_remote_host() { - let url = Url::parse("wss://hub.example.com/").expect("valid url"); - let credential = bearer_credential(); - if let Err(ClientError::InsecureScheme { .. }) = - open_socket(&url, &credential, ConnectionKind::Harness, None, false).await - { - panic!("wss:// must not be rejected by the scheme guard") - } - } - #[tokio::test] - async fn open_socket_allows_plaintext_ws_when_insecure_opt_in() { - let url = Url::parse("ws://hub.example.com:1/").expect("valid url"); - let credential = bearer_credential(); - if let Err(ClientError::InsecureScheme { .. }) = - open_socket(&url, &credential, ConnectionKind::Harness, None, true).await - { - panic!("allow_insecure_ws must bypass the scheme guard") - } - } - #[tokio::test] - async fn open_socket_rejects_role_mismatch() { - let url = Url::parse("ws://127.0.0.1:1/?role=harness").expect("valid url"); - let credential = bearer_credential(); - match open_socket(&url, &credential, ConnectionKind::ToolServer, None, false).await { - Err(ClientError::InvalidConfig(msg)) => { - assert!( - msg.contains("conflicts with"), - "message should mention conflict; got: {msg}" - ); - } - other => panic!("expected InvalidConfig; got {other:?}"), - } - } - #[test] - fn host_is_loopback_recognises_canonical_names() { - for raw in [ - "ws://127.0.0.1/", - "ws://[::1]/", - "ws://localhost/", - "ws://LOCALHOST/", - ] { - let url = Url::parse(raw).expect("valid url"); - assert!(host_is_loopback(&url), "{raw} must be treated as loopback"); - } - for raw in ["ws://hub.example.com/", "ws://10.0.0.1/", "ws://127.0.0.2/"] { - let url = Url::parse(raw).expect("valid url"); - assert!( - !host_is_loopback(&url), - "{raw} must NOT be treated as loopback", - ); - } - } - #[test] - fn exit_for_close_code_classifies_terminal_range() { - assert!(matches!( - exit_for_close_code(Some(4100)), - ConnectedExit::TerminalClose(4100) - )); - assert!(matches!( - exit_for_close_code(Some(4199)), - ConnectedExit::TerminalClose(4199) - )); - assert!(matches!( - exit_for_close_code(Some(4099)), - ConnectedExit::SocketClosed(DisconnectCause::CloseFrame(Some(4099))) - )); - assert!(matches!( - exit_for_close_code(Some(4200)), - ConnectedExit::SocketClosed(DisconnectCause::CloseFrame(Some(4200))) - )); - assert!(matches!( - exit_for_close_code(Some(1000)), - ConnectedExit::SocketClosed(DisconnectCause::CloseFrame(Some(1000))) - )); - assert!(matches!( - exit_for_close_code(None), - ConnectedExit::SocketClosed(DisconnectCause::CloseFrame(None)) - )); - } - #[test] - fn disconnect_cause_labels_and_fields() { - assert_eq!( - DisconnectCause::CloseFrame(Some(1006)).label(), - "close_frame" - ); - assert_eq!( - DisconnectCause::CloseFrame(Some(1006)).close_code(), - Some(1006) - ); - assert_eq!(DisconnectCause::Eof.label(), "eof"); - assert_eq!(DisconnectCause::Eof.close_code(), None); - assert_eq!(DisconnectCause::Eof.detail(), None); - let read = DisconnectCause::ReadError("reset".to_owned()); - assert_eq!(read.label(), "transport_read_error"); - assert_eq!(read.detail(), Some("reset")); - let write = DisconnectCause::WriteError("pipe".to_owned()); - assert_eq!(write.label(), "transport_write_error"); - assert_eq!(write.detail(), Some("pipe")); - assert_eq!(DisconnectCause::Forced.label(), "forced"); - } - #[test] - fn conn_health_snapshot_without_clock_skew_reports_zero_jump() { - let health = ConnHealth::new(); - health.record_inbound(); - health.refresh_clock(); - let snap = health.snapshot(); - assert_eq!(snap.clock_jump_ms, 0); - assert!(snap.since_last_probe_monotonic_ms < 2_000); - } - #[test] - fn conn_health_snapshot_reports_wall_clock_jump() { - let health = ConnHealth::new(); - { - let mut state = health.state.lock(); - state.wall_ref = SystemTime::now() - Duration::from_secs(10); - } - let snap = health.snapshot(); - assert!(snap.since_last_probe_wall_ms >= 9_000); - assert!(snap.since_last_probe_monotonic_ms < 2_000); - assert!(snap.clock_jump_ms >= 8_000); - health.reset(); - assert_eq!(health.snapshot().clock_jump_ms, 0); - } - #[test] - fn conn_health_accumulates_jump_across_refreshes() { - let health = ConnHealth::new(); - { - let mut state = health.state.lock(); - state.wall_ref = SystemTime::now() - Duration::from_secs(5); - } - health.refresh_clock(); - { - let mut state = health.state.lock(); - state.wall_ref = SystemTime::now() - Duration::from_secs(4); - } - let snap = health.snapshot(); - assert!(snap.clock_jump_ms >= 8_000); - } - use std::pin::Pin; - use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; - use std::task::{Context, Poll}; - /// Default ping period for writer tests that don't exercise the - /// keepalive: long enough that no ping fires during the test. - const TEST_PING_NEVER: Duration = Duration::from_secs(3_600); - /// In-memory [`futures::Sink`] for `run_writer` tests. Records the - /// text payload of every `Message::Text` sent and counts every - /// `Message::Ping` (keepalive). When the `fail` flag is set, `send` - /// errors at `poll_ready`, modelling a dead socket. - #[derive(Clone)] - struct RecordingSink { - recorded: Arc>>, - pings: Arc, - fail: Arc, - } - impl RecordingSink { - fn new() -> Self { - Self { - recorded: Arc::new(std::sync::Mutex::new(Vec::new())), - pings: Arc::new(AtomicUsize::new(0)), - fail: Arc::new(AtomicBool::new(false)), - } - } - fn recorded(&self) -> Arc>> { - self.recorded.clone() - } - fn pings(&self) -> Arc { - self.pings.clone() - } - fn fail_flag(&self) -> Arc { - self.fail.clone() - } - } - impl futures::Sink for RecordingSink { - type Error = std::io::Error; - fn poll_ready( - self: Pin<&mut Self>, - _cx: &mut Context<'_>, - ) -> Poll> { - if self.fail.load(Ordering::SeqCst) { - Poll::Ready(Err(std::io::Error::new( - std::io::ErrorKind::BrokenPipe, - "sink dead", - ))) - } else { - Poll::Ready(Ok(())) - } - } - fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { - match item { - Message::Text(text) => { - self.recorded - .lock() - .expect("recorded lock") - .push(text.as_str().to_owned()); - } - Message::Ping(_) => { - self.pings.fetch_add(1, Ordering::SeqCst); - } - _ => {} - } - Ok(()) - } - fn poll_flush( - self: Pin<&mut Self>, - _cx: &mut Context<'_>, - ) -> Poll> { - Poll::Ready(Ok(())) - } - fn poll_close( - self: Pin<&mut Self>, - _cx: &mut Context<'_>, - ) -> Poll> { - Poll::Ready(Ok(())) - } - } - type TestCtl = WriterControl; - fn idle_write_error_slot() -> WriteErrorSlot { - Arc::new(parking_lot::Mutex::new(None)) - } - /// Poll `predicate` every 5ms up to ~2s. Keeps the writer-task tests - /// off arbitrary fixed sleeps for the positive assertions. - async fn wait_until bool>(predicate: F, label: &str) { - for _ in 0..400 { - if predicate() { - return; - } - tokio::time::sleep(Duration::from_millis(5)).await; - } - panic!("timed out waiting for: {label}"); - } - #[tokio::test] - async fn writer_drains_outbound_while_live() { - let sink = RecordingSink::new(); - let recorded = sink.recorded(); - let (out_tx, out_rx) = mpsc::channel::(8); - let (_ctl_tx, ctl_rx) = mpsc::channel::(2); - let (stop_tx, stop_rx) = mpsc::channel::<()>(1); - let writer = tokio::spawn(run_writer( - sink, - out_rx, - ctl_rx, - stop_rx, - TEST_PING_NEVER, - idle_write_error_slot(), - )); - out_tx.send("a".to_owned()).await.expect("send a"); - out_tx.send("b".to_owned()).await.expect("send b"); - wait_until( - || recorded.lock().expect("lock").len() == 2, - "two frames drained", - ) - .await; - assert_eq!( - *recorded.lock().expect("lock"), - vec!["a".to_owned(), "b".to_owned()], - "frames must be written to the live sink in order" - ); - stop_tx.send(()).await.expect("stop"); - writer.await.expect("writer task joins"); - } - #[tokio::test] - async fn writer_honors_custom_ping_interval() { - let sink = RecordingSink::new(); - let pings = sink.pings(); - let (_out_tx, out_rx) = mpsc::channel::(4); - let (_ctl_tx, ctl_rx) = mpsc::channel::(2); - let (stop_tx, stop_rx) = mpsc::channel::<()>(1); - let writer = tokio::spawn(run_writer( - sink, - out_rx, - ctl_rx, - stop_rx, - Duration::from_millis(20), - idle_write_error_slot(), - )); - wait_until( - || pings.load(Ordering::SeqCst) >= 3, - "three keepalive pings at the configured cadence", - ) - .await; - stop_tx.send(()).await.expect("stop"); - writer.await.expect("writer task joins"); - } - #[tokio::test] - async fn writer_re_arms_custom_ping_interval_after_resume() { - let dead = RecordingSink::new(); - let (_out_tx, out_rx) = mpsc::channel::(4); - let (ctl_tx, ctl_rx) = mpsc::channel::(2); - let (stop_tx, stop_rx) = mpsc::channel::<()>(1); - let writer = tokio::spawn(run_writer( - dead, - out_rx, - ctl_rx, - stop_rx, - Duration::from_millis(20), - idle_write_error_slot(), - )); - ctl_tx.send(WriterControl::Pause).await.expect("pause"); - let fresh = RecordingSink::new(); - let fresh_pings = fresh.pings(); - ctl_tx - .send(WriterControl::Resume(fresh)) - .await - .expect("resume"); - wait_until( - || fresh_pings.load(Ordering::SeqCst) >= 3, - "keepalive pings resume on the configured cadence after Resume", - ) - .await; - stop_tx.send(()).await.expect("stop"); - writer.await.expect("writer task joins"); - } - #[tokio::test] - async fn writer_buffers_during_pause_and_flushes_on_resume() { - let dead = RecordingSink::new(); - let dead_log = dead.recorded(); - let (out_tx, out_rx) = mpsc::channel::(16); - let (ctl_tx, ctl_rx) = mpsc::channel::(2); - let (stop_tx, stop_rx) = mpsc::channel::<()>(1); - let writer = tokio::spawn(run_writer( - dead, - out_rx, - ctl_rx, - stop_rx, - TEST_PING_NEVER, - idle_write_error_slot(), - )); - ctl_tx.send(WriterControl::Pause).await.expect("pause"); - tokio::time::sleep(Duration::from_millis(20)).await; - for frame in ["g1", "g2", "g3"] { - out_tx - .send(frame.to_owned()) - .await - .expect("enqueue during gap"); - } - tokio::time::sleep(Duration::from_millis(50)).await; - assert!( - dead_log.lock().expect("lock").is_empty(), - "paused writer must not drain onto the dead sink; got {:?}", - dead_log.lock().expect("lock") - ); - let fresh = RecordingSink::new(); - let fresh_log = fresh.recorded(); - ctl_tx - .send(WriterControl::Resume(fresh)) - .await - .expect("resume"); - wait_until( - || fresh_log.lock().expect("lock").len() == 3, - "buffered frames flush after resume", - ) - .await; - assert_eq!( - *fresh_log.lock().expect("lock"), - vec!["g1".to_owned(), "g2".to_owned(), "g3".to_owned()], - "all gap frames flush, in order, to the fresh sink" - ); - assert!( - dead_log.lock().expect("lock").is_empty(), - "no frame must ever reach the dead sink" - ); - stop_tx.send(()).await.expect("stop"); - writer.await.expect("writer task joins"); - } - #[tokio::test] - async fn writer_send_error_pauses_until_resume_without_multi_frame_loss() { - let failing = RecordingSink::new(); - let failing_log = failing.recorded(); - let fail_flag = failing.fail_flag(); - let (out_tx, out_rx) = mpsc::channel::(16); - let (ctl_tx, ctl_rx) = mpsc::channel::(2); - let (stop_tx, stop_rx) = mpsc::channel::<()>(1); - let write_error = idle_write_error_slot(); - let writer = tokio::spawn(run_writer( - failing, - out_rx, - ctl_rx, - stop_rx, - TEST_PING_NEVER, - write_error.clone(), - )); - out_tx.send("ok".to_owned()).await.expect("send ok"); - wait_until( - || failing_log.lock().expect("lock").len() == 1, - "first frame drained before failure", - ) - .await; - fail_flag.store(true, Ordering::SeqCst); - out_tx.send("lost".to_owned()).await.expect("enqueue lost"); - out_tx - .send("kept1".to_owned()) - .await - .expect("enqueue kept1"); - out_tx - .send("kept2".to_owned()) - .await - .expect("enqueue kept2"); - tokio::time::sleep(Duration::from_millis(50)).await; - assert_eq!( - *failing_log.lock().expect("lock"), - vec!["ok".to_owned()], - "only the pre-failure frame should have been recorded on the dead sink" - ); - assert!( - write_error - .lock() - .as_deref() - .is_some_and(|detail| detail.contains("sink dead")), - "failed send must record the write-error detail for disconnect classification" - ); - let fresh = RecordingSink::new(); - let fresh_log = fresh.recorded(); - ctl_tx - .send(WriterControl::Resume(fresh)) - .await - .expect("resume"); - wait_until( - || fresh_log.lock().expect("lock").len() == 2, - "buffered post-failure frames flush after resume", - ) - .await; - assert_eq!( - *fresh_log.lock().expect("lock"), - vec!["kept1".to_owned(), "kept2".to_owned()], - "post-failure frames survive; only the in-flight 'lost' frame is gone" - ); - stop_tx.send(()).await.expect("stop"); - writer.await.expect("writer task joins"); - } - #[tokio::test] - async fn writer_resume_discards_stale_write_error() { - let sink = RecordingSink::new(); - let (_out_tx, out_rx) = mpsc::channel::(4); - let (ctl_tx, ctl_rx) = mpsc::channel::(2); - let (stop_tx, stop_rx) = mpsc::channel::<()>(1); - let write_error = idle_write_error_slot(); - let writer = tokio::spawn(run_writer( - sink, - out_rx, - ctl_rx, - stop_rx, - TEST_PING_NEVER, - write_error.clone(), - )); - ctl_tx.send(WriterControl::Pause).await.expect("pause"); - *write_error.lock() = Some("frame send failed: stale broken pipe".to_owned()); - ctl_tx - .send(WriterControl::Resume(RecordingSink::new())) - .await - .expect("resume"); - wait_until( - || write_error.lock().is_none(), - "Resume must clear a stale write-error left by a late old-sink send", - ) - .await; - stop_tx.send(()).await.expect("stop"); - writer.await.expect("writer task joins"); - } - #[tokio::test] - async fn writer_exits_on_stop_signal() { - let sink = RecordingSink::new(); - let (_out_tx, out_rx) = mpsc::channel::(4); - let (_ctl_tx, ctl_rx) = mpsc::channel::(2); - let (stop_tx, stop_rx) = mpsc::channel::<()>(1); - let writer = tokio::spawn(run_writer( - sink, - out_rx, - ctl_rx, - stop_rx, - TEST_PING_NEVER, - idle_write_error_slot(), - )); - stop_tx.send(()).await.expect("stop"); - tokio::time::timeout(Duration::from_secs(2), writer) - .await - .expect("writer must exit on the stop signal") - .expect("writer task joins"); - } - #[tokio::test] - async fn writer_exits_when_outbound_channel_closes() { - let sink = RecordingSink::new(); - let (out_tx, out_rx) = mpsc::channel::(4); - let (_ctl_tx, ctl_rx) = mpsc::channel::(2); - let (_stop_tx, stop_rx) = mpsc::channel::<()>(1); - let writer = tokio::spawn(run_writer( - sink, - out_rx, - ctl_rx, - stop_rx, - TEST_PING_NEVER, - idle_write_error_slot(), - )); - drop(out_tx); - tokio::time::timeout(Duration::from_secs(2), writer) - .await - .expect("writer must exit when outbound closes") - .expect("writer task joins"); - } - #[tokio::test] - async fn writer_exits_when_control_channel_closes() { - let sink = RecordingSink::new(); - let (_out_tx, out_rx) = mpsc::channel::(4); - let (ctl_tx, ctl_rx) = mpsc::channel::(2); - let (_stop_tx, stop_rx) = mpsc::channel::<()>(1); - let writer = tokio::spawn(run_writer( - sink, - out_rx, - ctl_rx, - stop_rx, - TEST_PING_NEVER, - idle_write_error_slot(), - )); - drop(ctl_tx); - tokio::time::timeout(Duration::from_secs(2), writer) - .await - .expect("writer must exit when the control channel closes") - .expect("writer task joins"); - } - /// Socket-less `HubConnection` for tests: observe the sent frame and - /// resolve the response waiter without a live server or actor task. - fn test_connection() -> (Arc, Arc, mpsc::Receiver) { - let (outbound_tx, outbound_rx) = mpsc::channel::(8); - let demux = Arc::new(Demux::with_outbound(outbound_tx.clone())); - let credential: Arc = Arc::new(AuthCredential::bearer("test-token")); - let (stop_tx, _stop_rx) = mpsc::channel::<()>(1); - let (reconnect_tx, _reconnect_rx) = mpsc::channel::<()>(1); - let inner = Arc::new(HubConnectionInner { - key: ConnKey { - url: "ws://test/v1/tools".to_owned(), - principal: credential.principal_key(), - }, - kind: ConnectionKind::ToolServer, - credential, - on_reconnect: None, - on_disconnect: None, - server_id: None, - server_description: None, - server_metadata: None, - alpha_test_key: None, - allow_insecure_ws: false, - on_fatal: None, - reconnect_backoff: resolve_reconnect_backoff(None), - outbound_tx, - demux: demux.clone(), - bound_sessions: Arc::new(RefCountedSet::new()), - connection_id: Arc::new(Mutex::new(None)), - hello_capabilities: parking_lot::RwLock::new(Vec::new()), - next_request_id: std::sync::atomic::AtomicU64::new(1), - shutdown: CancellationToken::new(), - stop_tx, - reconnect_tx, - early_notif_rx: parking_lot::Mutex::new(Some(demux.subscribe_notifications())), - health: ConnHealth::new(), - writer_error: Arc::new(parking_lot::Mutex::new(None)), - }); - (Arc::new(HubConnection { inner }), demux, outbound_rx) - } - #[test] - fn classify_stream_end_prefers_recorded_write_error() { - let (conn, _demux, _outbound_rx) = test_connection(); - let inner = conn.inner.as_ref(); - assert!(matches!( - classify_stream_end(inner, None), - DisconnectCause::Eof - )); - assert!( - matches!(classify_stream_end(inner, Some("reset by peer".to_owned())), - DisconnectCause::ReadError(detail) if detail == "reset by peer") - ); - *inner.writer_error.lock() = Some("ping send failed: broken pipe".to_owned()); - assert!(matches!(classify_stream_end(inner, None), - DisconnectCause::WriteError(detail) if detail == - "ping send failed: broken pipe")); - assert!( - inner.writer_error.lock().is_none(), - "classification must consume the recorded write error" - ); - *inner.writer_error.lock() = Some("frame send failed: broken pipe".to_owned()); - assert!(matches!( - classify_stream_end(inner, Some("reset".to_owned())), - DisconnectCause::WriteError(_) - )); - } - #[test] - fn supports_is_unknown_until_capabilities_advertised() { - let (conn, _demux, _outbound_rx) = test_connection(); - assert_eq!(conn.supports("session_attach_server"), None); - *conn.inner.hello_capabilities.write() = vec!["session_attach_server".to_owned()]; - assert_eq!(conn.supports("session_attach_server"), Some(true)); - assert_eq!(conn.supports("some_other_method"), Some(false)); - } - #[tokio::test] - async fn call_request_with_timeout_round_trips_via_demux() { - let (conn, demux, mut outbound_rx) = test_connection(); - let session = SessionId::new("rt_session").expect("valid"); - let request_id = conn.try_alloc_request_id().expect("request id"); - let id_str = request_id.to_string(); - let req = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: Some(session.clone()), - method: Method::Hook.as_wire_str().to_owned(), - params: serde_json::json!({ "k" : "v" }), - }; - let call = tokio::spawn(async move { - conn.call_request_with_timeout(request_id, &req, Duration::from_secs(5)) - .await - }); - let sent = tokio::time::timeout(Duration::from_secs(1), outbound_rx.recv()) - .await - .expect("frame sent before deadline") - .expect("outbound frame present"); - let sent_value: Value = serde_json::from_str(&sent).expect("sent frame is valid json"); - assert_eq!(sent_value["id"].as_str(), Some(id_str.as_str())); - assert_eq!( - sent_value["method"].as_str(), - Some(Method::Hook.as_wire_str()) - ); - let outcome = demux.route(serde_json::json!( - { "jsonrpc" : "2.0", "id" : id_str, "session_id" : session.as_str(), - "result" : { "ok" : true }, } - )); - assert_eq!(outcome, crate::demux::RouteOutcome::Response); - let resp = call - .await - .expect("call task joins") - .expect("call resolves with a response"); - let ResponseOutcome::Result(value) = resp.outcome else { - panic!("expected a result outcome"); - }; - assert_eq!(value, serde_json::json!({ "ok" : true })); - } - #[tokio::test] - async fn call_request_reclaims_waiter_on_send_failure() { - let (conn, demux, outbound_rx) = test_connection(); - drop(outbound_rx); - let request_id = conn.try_alloc_request_id().expect("request id"); - let probe_id = request_id.clone(); - let req = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: None, - method: Method::Hook.as_wire_str().to_owned(), - params: serde_json::json!({}), - }; - let result = conn.call_request(request_id, &req).await; - assert!(matches!(result, Err(ClientError::NetworkError(_)))); - assert!( - demux.take_response_waiter(&probe_id).is_none(), - "the failed-send waiter is reclaimed so it cannot leak" - ); - } - #[tokio::test] - async fn serve_send_failure_fails_fast_without_retry() { - let (conn, demux, outbound_rx) = test_connection(); - drop(outbound_rx); - let session = SessionId::new("serve_session").expect("valid"); - let result = tokio::time::timeout( - Duration::from_secs(5), - conn.serve(session, kigi_tool_protocol::ServeParams { tools: vec![] }), - ) - .await - .expect("serve must fail bounded, not park"); - assert!(matches!(result, Err(ClientError::NetworkError(_)))); - let request_id = kigi_tool_protocol::RequestId::new("c1").expect("valid"); - assert!( - demux.take_response_waiter(&request_id).is_none(), - "the failed attempt must not leak a waiter" - ); - assert_eq!( - conn.try_alloc_request_id().expect("request id").to_string(), - "c2", - "a non-timeout failure must consume a single attempt, not retry" - ); - } - #[tokio::test(start_paused = true)] - async fn serve_times_out_bounded_and_reclaims_every_attempt_waiter() { - let (conn, demux, mut outbound_rx) = test_connection(); - let session = SessionId::new("serve_timeout").expect("valid"); - let result = conn - .serve(session, kigi_tool_protocol::ServeParams { tools: vec![] }) - .await; - assert!(matches!(result, Err(ClientError::NetworkError(_)))); - for id in ["c1", "c2", "c3"] { - let sent = outbound_rx.try_recv().expect("attempt frame sent"); - let value: Value = serde_json::from_str(&sent).expect("valid json"); - assert_eq!(value["id"].as_str(), Some(id)); - let request_id = kigi_tool_protocol::RequestId::new(id).expect("valid"); - assert!( - demux.take_response_waiter(&request_id).is_none(), - "attempt {id} must not leak a waiter" - ); - } - assert!( - outbound_rx.try_recv().is_err(), - "exactly SERVE_MAX_ATTEMPTS frames are sent" - ); - } - #[tokio::test] - async fn call_request_reclaims_waiter_on_caller_cancellation() { - let (conn, demux, mut outbound_rx) = test_connection(); - let request_id = conn.try_alloc_request_id().expect("request id"); - let probe_id = request_id.clone(); - let conn_for_call = conn.clone(); - let call = tokio::spawn(async move { - let req = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: None, - method: Method::Hook.as_wire_str().to_owned(), - params: serde_json::json!({}), - }; - conn_for_call.call_request(request_id, &req).await - }); - tokio::time::timeout(Duration::from_secs(1), outbound_rx.recv()) - .await - .expect("frame sent") - .expect("outbound frame present"); - call.abort(); - let _ = call.await; - assert!( - demux.take_response_waiter(&probe_id).is_none(), - "the cancelled caller's waiter is reclaimed so it cannot leak" - ); - } - #[tokio::test] - async fn reader_phase_exits_socket_closed_on_forced_reconnect_signal() { - let (conn, _demux, _outbound_rx) = test_connection(); - let (reconnect_tx, mut reconnect_rx) = mpsc::channel::<()>(1); - let (_stop_tx, mut stop_rx) = mpsc::channel::<()>(1); - reconnect_tx.try_send(()).expect("queue forced reconnect"); - let mut stream = - futures::stream::pending::>(); - let exit = tokio::time::timeout( - Duration::from_secs(1), - run_reader_phase( - conn.inner.as_ref(), - &mut stream, - &mut stop_rx, - &mut reconnect_rx, - Duration::from_secs(75), - ), - ) - .await - .expect("forced reconnect must break the reader phase"); - assert!( - matches!(exit, ConnectedExit::SocketClosed(DisconnectCause::Forced)), - "a forced reconnect exits as SocketClosed (drives the reconnect path)" - ); - } - #[tokio::test] - async fn stop_signal_outranks_forced_reconnect() { - let (conn, _demux, _outbound_rx) = test_connection(); - let (reconnect_tx, mut reconnect_rx) = mpsc::channel::<()>(1); - let (stop_tx, mut stop_rx) = mpsc::channel::<()>(1); - reconnect_tx.try_send(()).expect("queue forced reconnect"); - stop_tx.try_send(()).expect("queue stop"); - let mut stream = - futures::stream::pending::>(); - let exit = tokio::time::timeout( - Duration::from_secs(1), - run_reader_phase( - conn.inner.as_ref(), - &mut stream, - &mut stop_rx, - &mut reconnect_rx, - Duration::from_secs(75), - ), - ) - .await - .expect("stop must break the reader phase"); - assert!(matches!(exit, ConnectedExit::Stop)); - } - #[tokio::test] - async fn drain_reconnect_signals_clears_stale_signal_only() { - let (reconnect_tx, mut reconnect_rx) = mpsc::channel::<()>(1); - reconnect_tx.try_send(()).expect("queue stale signal"); - drain_reconnect_signals(&mut reconnect_rx); - assert!( - reconnect_rx.try_recv().is_err(), - "a stale pre-reconnect signal is consumed by the drain" - ); - reconnect_tx.try_send(()).expect("queue fresh signal"); - assert!( - reconnect_rx.try_recv().is_ok(), - "the drain must not disable the channel for future signals" - ); - } - #[tokio::test] - async fn early_subscribed_receiver_buffers_pre_run_connection_notifications() { - let (conn, demux, _outbound_rx) = test_connection(); - let outcome = demux.route(serde_json::json!( - { "jsonrpc" : "2.0", "id" : "b1", "method" : "session.bind", "params" - : { "session_id" : "s1" }, } - )); - assert_eq!(outcome, crate::demux::RouteOutcome::Notification); - let mut rx = conn - .take_early_notifications() - .expect("receiver retained until taken"); - let frame = rx.try_recv().expect("pre-run frame buffered"); - assert_eq!(frame["method"], "session.bind"); - assert!( - conn.take_early_notifications().is_none(), - "the early receiver is handed off exactly once" - ); - } - #[tokio::test] - async fn call_request_with_timeout_reclaims_waiter_on_deadline() { - let (conn, demux, mut outbound_rx) = test_connection(); - let session = SessionId::new("to_session").expect("valid"); - let request_id = conn.try_alloc_request_id().expect("request id"); - let probe_id = request_id.clone(); - let req = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: Some(session), - method: Method::Hook.as_wire_str().to_owned(), - params: serde_json::json!({}), - }; - let result = conn - .call_request_with_timeout(request_id, &req, Duration::from_millis(50)) - .await; - assert!(matches!(result, Err(ClientError::NetworkError(_)))); - assert!( - outbound_rx.try_recv().is_ok(), - "the request frame is sent before the deadline fires" - ); - assert!( - demux.take_response_waiter(&probe_id).is_none(), - "the timed-out waiter is reclaimed so it cannot leak" - ); - } - /// Regression: a *forced* reconnect abandons a still-healthy socket. If - /// the first reconnect attempt then fails, the actor must keep retrying - /// off the abandoned stream — falling back into the reader phase would - /// park in `stream.next()` on the live old connection forever (the - /// reconnect signal was already consumed), stalling the retry loop. - /// - /// Mock: conn #0 (initial) completes the handshake and stays healthy; - /// conn #1 (first reconnect) is dropped before the ack (transport - /// failure); conn #2 must then be attempted and complete. With the bug, - /// upgrade #2 never happens and the test times out. - #[tokio::test] - async fn forced_reconnect_retries_past_failed_attempt_without_repolling_old_stream() { - use futures::{SinkExt as _, StreamExt as _}; - use std::sync::atomic::{AtomicUsize, Ordering}; - let listener = tokio::net::TcpListener::bind("127.0.0.1:0") - .await - .expect("bind mock hub"); - let addr = listener.local_addr().expect("mock addr"); - let upgrades = Arc::new(AtomicUsize::new(0)); - let upgrades_srv = upgrades.clone(); - tokio::spawn(async move { - loop { - let Ok((tcp, _)) = listener.accept().await else { - return; - }; - let n = upgrades_srv.fetch_add(1, Ordering::SeqCst); - tokio::spawn(async move { - let Ok(mut ws) = tokio_tungstenite::accept_async(tcp).await else { - return; - }; - if n == 1 { - return; - } - let _ = ws.next().await; - let ack = serde_json::json!( - { "connection_id" : format!("mock-conn-{n}"), "user_id" : "test", - "computer_hub_version" : "test", "supported_protocol_versions" : - ["1.0.0"], } - ); - if ws - .send(tokio_tungstenite::tungstenite::Message::Text( - ack.to_string().into(), - )) - .await - .is_err() - { - return; - } - while let Some(msg) = ws.next().await { - if msg.is_err() { - return; - } - } - }); - } - }); - let credential: Arc = Arc::new(AuthCredential::bearer("test-token")); - let conn = HubConnection::connect(ConnectionConfig { - url: url::Url::parse(&format!("ws://{addr}/v1/tools")).expect("mock url"), - credential, - kind: ConnectionKind::ToolServer, - on_reconnect: None, - on_disconnect: None, - on_connect: None, - server_id: None, - server_description: None, - server_metadata: None, - outbound_buffer: None, - tuning: ConnectionTuning { - reconnect_backoff: Some(Arc::from([Duration::from_millis(10)])), - ..Default::default() - }, - alpha_test_key: None, - allow_insecure_ws: false, - on_fatal: None, - }) - .await - .expect("initial connect"); - conn.force_reconnect(); - let deadline = tokio::time::Instant::now() + Duration::from_secs(5); - while upgrades.load(Ordering::SeqCst) < 3 { - assert!( - tokio::time::Instant::now() < deadline, - "retry stalled after a failed forced-reconnect attempt: \ - {} upgrades observed (expected 3: initial + failed + successful)", - upgrades.load(Ordering::SeqCst) - ); - tokio::time::sleep(Duration::from_millis(20)).await; - } - conn.request_shutdown(); - conn.await_shutdown().await; - } - #[tokio::test] - async fn call_request_serialization_failure_registers_no_waiter() { - struct FailingParams; - impl serde::Serialize for FailingParams { - fn serialize(&self, _serializer: S) -> Result { - Err(serde::ser::Error::custom("intentionally unserializable")) - } - } - let (conn, demux, _outbound_rx) = test_connection(); - let session = SessionId::new("serde_fail_session").expect("valid"); - let request_id = conn.try_alloc_request_id().expect("request id"); - let probe_id = request_id.clone(); - let req = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: Some(session), - method: Method::Hook.as_wire_str().to_owned(), - params: FailingParams, - }; - let result = conn.call_request(request_id, &req).await; - assert!(result.is_err(), "serialization failure must surface"); - assert!( - demux.take_response_waiter(&probe_id).is_none(), - "no waiter may be registered when serialization fails" - ); - } - type WsError = tokio_tungstenite::tungstenite::Error; - type InboundTx = futures::channel::mpsc::UnboundedSender>; - type InboundRx = futures::channel::mpsc::UnboundedReceiver>; - /// In-memory inbound frame source for `run_reader_phase` tests - /// (mirrors `RecordingSink` for the writer half). - fn test_inbound() -> (InboundTx, InboundRx) { - futures::channel::mpsc::unbounded() - } - /// A zero or unset liveness deadline resolves to 2.5× the effective - /// ping cadence; a positive override is honored verbatim. Mirrors the - /// `resolve_ws_ping_interval` clamp semantics. - #[test] - fn resolve_ws_liveness_deadline_clamps_zero_and_unset_to_default() { - let ping = Duration::from_secs(30); - assert_eq!( - resolve_ws_liveness_deadline(None, ping), - Duration::from_secs(75) - ); - assert_eq!( - resolve_ws_liveness_deadline(Some(Duration::ZERO), ping), - Duration::from_secs(75) - ); - let custom = Duration::from_secs(120); - assert_eq!(resolve_ws_liveness_deadline(Some(custom), ping), custom); - } - /// The per-attempt reconnect budget tracks the liveness deadline above - /// the floor and is clamped to the floor below it, so a small liveness - /// override can never starve connection establishment. - #[test] - fn reconnect_attempt_budget_floors_small_deadlines() { - assert_eq!( - reconnect_attempt_budget(Duration::from_millis(2_500)), - RECONNECT_ATTEMPT_MIN_BUDGET - ); - assert_eq!( - reconnect_attempt_budget(RECONNECT_ATTEMPT_MIN_BUDGET), - RECONNECT_ATTEMPT_MIN_BUDGET - ); - let large = Duration::from_secs(300); - assert_eq!(reconnect_attempt_budget(large), large); - } - #[test] - fn resolve_ws_liveness_deadline_scales_with_ping_override() { - assert_eq!( - resolve_ws_liveness_deadline(None, Duration::from_secs(10)), - Duration::from_secs(25) - ); - } - #[tokio::test(start_paused = true)] - async fn reader_deadline_kills_silently_dead_connection() { - let (conn, _demux, _outbound_rx) = test_connection(); - let (inbound_tx, mut inbound_rx) = test_inbound(); - let (_stop_tx, mut stop_rx) = mpsc::channel::<()>(1); - let (_reconnect_tx, mut reconnect_rx) = mpsc::channel::<()>(1); - let liveness = Duration::from_secs(75); - let start = tokio::time::Instant::now(); - let exit = run_reader_phase( - &conn.inner, - &mut inbound_rx, - &mut stop_rx, - &mut reconnect_rx, - liveness, - ) - .await; - assert!(matches!( - exit, - ConnectedExit::SocketClosed(DisconnectCause::LivenessDeadline) - )); - assert_eq!( - start.elapsed(), - liveness, - "expiry exactly one liveness window after (re)entry" - ); - drop(inbound_tx); - } - #[tokio::test(start_paused = true)] - async fn reader_deadline_rearms_on_any_inbound_frame() { - let (conn, _demux, _outbound_rx) = test_connection(); - let (inbound_tx, mut inbound_rx) = test_inbound(); - let (_stop_tx, mut stop_rx) = mpsc::channel::<()>(1); - let (_reconnect_tx, mut reconnect_rx) = mpsc::channel::<()>(1); - let liveness = Duration::from_secs(75); - let phase = run_reader_phase( - &conn.inner, - &mut inbound_rx, - &mut stop_rx, - &mut reconnect_rx, - liveness, - ); - tokio::pin!(phase); - let frames = [ - Message::Pong(Vec::new().into()), - Message::Ping(Vec::new().into()), - Message::Text(r#"{"jsonrpc":"2.0","method":"noop","params":{}}"#.into()), - Message::Pong(Vec::new().into()), - ]; - for frame in frames { - tokio::time::advance(liveness * 3 / 4).await; - inbound_tx.unbounded_send(Ok(frame)).expect("send frame"); - assert!( - futures::poll!(phase.as_mut()).is_pending(), - "phase must stay live while frames keep arriving" - ); - } - tokio::time::advance(liveness - Duration::from_millis(1)).await; - assert!( - futures::poll!(phase.as_mut()).is_pending(), - "still inside the window re-armed by the last frame" - ); - tokio::time::advance(Duration::from_millis(1)).await; - match futures::poll!(phase.as_mut()) { - std::task::Poll::Ready(exit) => { - assert!(matches!( - exit, - ConnectedExit::SocketClosed(DisconnectCause::LivenessDeadline) - )); - } - std::task::Poll::Pending => { - panic!("deadline must fire one window after the last frame") - } - } - } - #[tokio::test(start_paused = true)] - async fn reader_deadline_huge_override_saturates_instead_of_panicking() { - let (conn, _demux, _outbound_rx) = test_connection(); - let (inbound_tx, mut inbound_rx) = test_inbound(); - let (_stop_tx, mut stop_rx) = mpsc::channel::<()>(1); - let (_reconnect_tx, mut reconnect_rx) = mpsc::channel::<()>(1); - let phase = run_reader_phase( - &conn.inner, - &mut inbound_rx, - &mut stop_rx, - &mut reconnect_rx, - Duration::MAX, - ); - tokio::pin!(phase); - inbound_tx - .unbounded_send(Ok(Message::Pong(Vec::new().into()))) - .expect("send frame"); - assert!( - futures::poll!(phase.as_mut()).is_pending(), - "saturating re-arm must neither panic nor fire" - ); - } - /// Sink for writer↔reader composition tests: echoes every keepalive - /// `Ping` back as a `Pong` on the reader's inbound channel, emulating a - /// healthy server whose only traffic is the keepalive exchange. - struct PongEchoSink { - inbound: InboundTx, - } - impl futures::Sink for PongEchoSink { - type Error = std::io::Error; - fn poll_ready( - self: Pin<&mut Self>, - _cx: &mut Context<'_>, - ) -> Poll> { - Poll::Ready(Ok(())) - } - fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { - if let Message::Ping(payload) = item { - let _ = self.inbound.unbounded_send(Ok(Message::Pong(payload))); - } - Ok(()) - } - fn poll_flush( - self: Pin<&mut Self>, - _cx: &mut Context<'_>, - ) -> Poll> { - Poll::Ready(Ok(())) - } - fn poll_close( - self: Pin<&mut Self>, - _cx: &mut Context<'_>, - ) -> Poll> { - Poll::Ready(Ok(())) - } - } - #[tokio::test(start_paused = true)] - async fn default_ping_pong_composition_keeps_idle_connection_alive() { - let ping = resolve_ws_ping_interval(None); - let deadline = resolve_ws_liveness_deadline(None, ping); - let (conn, _demux, _outbound_rx) = test_connection(); - let (inbound_tx, mut inbound_rx) = test_inbound(); - let (_out_tx, out_rx) = mpsc::channel::(4); - let (ctl_tx, ctl_rx) = mpsc::channel::>(2); - let (writer_stop_tx, writer_stop_rx) = mpsc::channel::<()>(1); - let writer = tokio::spawn(run_writer( - PongEchoSink { - inbound: inbound_tx.clone(), - }, - out_rx, - ctl_rx, - writer_stop_rx, - ping, - idle_write_error_slot(), - )); - let (_stop_tx, mut stop_rx) = mpsc::channel::<()>(1); - let (_reconnect_tx, mut reconnect_rx) = mpsc::channel::<()>(1); - { - let phase = run_reader_phase( - &conn.inner, - &mut inbound_rx, - &mut stop_rx, - &mut reconnect_rx, - deadline, - ); - tokio::pin!(phase); - tokio::select! { - _ = phase.as_mut() => - panic!("idle-but-healthy connection tripped the deadline"), _ = - tokio::time::sleep(deadline * 4) => {} - } - } - ctl_tx.send(WriterControl::Pause).await.expect("pause"); - let (fresh_tx, mut fresh_rx) = test_inbound(); - ctl_tx - .send(WriterControl::Resume(PongEchoSink { inbound: fresh_tx })) - .await - .expect("resume"); - { - let phase = run_reader_phase( - &conn.inner, - &mut fresh_rx, - &mut stop_rx, - &mut reconnect_rx, - deadline, - ); - tokio::pin!(phase); - tokio::select! { - _ = phase.as_mut() => { - panic!("idle connection tripped the deadline after Pause→Resume") } _ = - tokio::time::sleep(deadline * 4) => {} - } - } - writer_stop_tx.send(()).await.expect("stop"); - writer.await.expect("writer task joins"); - } - #[tokio::test] - async fn reader_phase_close_frame_classification_unchanged() { - use tokio_tungstenite::tungstenite::protocol::CloseFrame; - use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; - let (conn, _demux, _outbound_rx) = test_connection(); - let (inbound_tx, mut inbound_rx) = test_inbound(); - let (_stop_tx, mut stop_rx) = mpsc::channel::<()>(1); - let (_reconnect_tx, mut reconnect_rx) = mpsc::channel::<()>(1); - inbound_tx - .unbounded_send(Ok(Message::Close(Some(CloseFrame { - code: CloseCode::from(4100), - reason: "evicted".into(), - })))) - .expect("send close"); - let exit = run_reader_phase( - &conn.inner, - &mut inbound_rx, - &mut stop_rx, - &mut reconnect_rx, - Duration::from_secs(75), - ) - .await; - assert!(matches!(exit, ConnectedExit::TerminalClose(4100))); - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/connection_borrow.rs b/crates/common/kigi-computer-hub-sdk/src/connection_borrow.rs deleted file mode 100644 index fe4839a..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/connection_borrow.rs +++ /dev/null @@ -1,224 +0,0 @@ -//! Shared connection-borrow lifecycle for `ToolServer` and `ToolHarness`. -//! -//! Wraps a pooled [`HubConnection`] with a [`CancellationToken`] for -//! shutdown coordination and an at-most-once `torn_down` guard. - -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; - -use kigi_tool_protocol::ConnectionKind; -use tokio_util::sync::CancellationToken; -use url::Url; - -use crate::auth::AuthProvider; -use crate::connection::{ - ConnectCallback, ConnectionTuning, DisconnectCallback, HubConnection, ReconnectCallback, -}; -use crate::error::ClientError; -use crate::pool::HubConnectionPool; - -/// Borrowed slice of a pooled [`HubConnection`] plus the refcount of -/// session bindings the borrower owns. Drop guard lives here so the -/// teardown sequence is at-most-once across explicit `shutdown` and -/// the `Drop` fallback. -pub(crate) struct ConnectionBorrow { - connection: Arc, - shutdown: CancellationToken, - /// At-most-once guard coordinated via `compare_exchange`. - torn_down: AtomicBool, -} - -impl std::fmt::Debug for ConnectionBorrow { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ConnectionBorrow") - .field( - "torn_down", - &self.torn_down.load(std::sync::atomic::Ordering::Relaxed), - ) - .finish_non_exhaustive() - } -} - -impl ConnectionBorrow { - /// Resolve a pool entry, refcount-bind every requested session, - /// and return a borrow. On any per-session bind failure the - /// already-bound sessions are unregistered before returning the - /// error so partial state never leaks. - pub(crate) async fn acquire( - pool: Arc, - url: Url, - auth: Arc, - kind: ConnectionKind, - on_reconnect: Option>, - on_disconnect: Option>, - on_connect: Option>, - server_id: Option, - server_description: Option, - server_metadata: Option, - alpha_test_key: Option, - allow_insecure_ws: bool, - tuning: ConnectionTuning, - ) -> Result { - let connection = pool - .get_or_connect_tuned( - url, - auth, - kind, - on_reconnect, - on_disconnect, - on_connect, - server_id, - server_description, - server_metadata, - alpha_test_key, - allow_insecure_ws, - tuning, - ) - .await?; - Ok(Self { - connection, - shutdown: CancellationToken::new(), - torn_down: AtomicBool::new(false), - }) - } - - pub(crate) fn connection(&self) -> &Arc { - &self.connection - } - - pub(crate) fn shutdown_token(&self) -> &CancellationToken { - &self.shutdown - } - - /// Returns `true` if this caller won the at-most-once teardown. - pub(crate) fn begin_teardown(&self) -> bool { - self.torn_down - .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) - .is_ok() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - use std::net::SocketAddr; - use std::sync::Arc; - - use crate::auth::AuthCredential; - use axum::Router; - use axum::extract::WebSocketUpgrade; - use axum::extract::ws::{Message, WebSocket}; - use axum::response::IntoResponse; - use axum::routing::get; - use serde_json::json; - use tokio::net::TcpListener; - - /// Spawn an in-process mock server that completes the WebSocket - /// handshake and ignores everything else. Returned address is - /// bound on `127.0.0.1`. - async fn spawn_borrow_mock_hub() -> SocketAddr { - let app = Router::new().route("/v1/tools", get(ws_upgrade)); - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("bind ephemeral"); - let addr = listener.local_addr().expect("local addr"); - tokio::spawn(async move { - let _ = axum::serve(listener, app.into_make_service()).await; - }); - tokio::task::yield_now().await; - addr - } - - async fn ws_upgrade(ws: WebSocketUpgrade) -> impl IntoResponse { - ws.on_upgrade(handle_socket) - } - - async fn handle_socket(mut socket: WebSocket) { - let _ = socket.recv().await; - let ack = json!({ - "connection_id": "borrow-mock", - "user_id": "test", - "computer_hub_version": "test", - "supported_protocol_versions": ["1.0.0"], - }); - let _ = socket.send(Message::Text(ack.to_string().into())).await; - // Keep the WebSocket alive until the client disconnects. - // These tests only exercise borrow lifecycle (teardown - // atomicity), not protocol frames. - while let Some(Ok(_msg)) = socket.recv().await {} - } - - async fn acquire_borrow() -> ConnectionBorrow { - let addr = spawn_borrow_mock_hub().await; - let url = Url::parse(&format!("ws://{addr}/v1/tools")).expect("valid url"); - let cred: Arc = Arc::new(AuthCredential::bearer("ignored")); - let pool = HubConnectionPool::new(); - ConnectionBorrow::acquire( - pool, - url, - cred, - ConnectionKind::Harness, - None, // on_reconnect - None, // on_disconnect - None, // on_connect - None, // server_id - None, // server_description - None, // server_metadata - None, // alpha_test_key - false, - ConnectionTuning::default(), - ) - .await - .expect("acquire borrow") - } - - #[tokio::test] - async fn begin_teardown_returns_true_once_and_false_after() { - let borrow = acquire_borrow().await; - assert!( - borrow.begin_teardown(), - "first call wins the at-most-once transition" - ); - assert!( - !borrow.begin_teardown(), - "subsequent calls observe the already-torn-down state" - ); - assert!( - !borrow.begin_teardown(), - "the at-most-once transition is sticky" - ); - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 4)] - async fn begin_teardown_is_atomic_under_concurrent_callers() { - let borrow = Arc::new(acquire_borrow().await); - let n_callers = 64; - let barrier = Arc::new(tokio::sync::Barrier::new(n_callers)); - let mut handles = Vec::with_capacity(n_callers); - for _ in 0..n_callers { - let borrow = borrow.clone(); - let barrier = barrier.clone(); - handles.push(tokio::spawn(async move { - barrier.wait().await; - borrow.begin_teardown() - })); - } - let mut wins = 0usize; - for h in handles { - if h.await.expect("join") { - wins += 1; - } - } - assert_eq!( - wins, 1, - "exactly one of {n_callers} concurrent callers must win the at-most-once transition" - ); - } - - #[tokio::test] - async fn acquire_returns_zero_bound_sessions() { - let borrow = acquire_borrow().await; - assert_eq!(borrow.connection().bound_session_count(), 0); - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/demux.rs b/crates/common/kigi-computer-hub-sdk/src/demux.rs deleted file mode 100644 index 9d08a58..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/demux.rs +++ /dev/null @@ -1,973 +0,0 @@ -//! Inbound frame demultiplexer. -//! -//! Frames inbound from the WebSocket fall into four buckets: -//! -//! 1. JSON-RPC **responses** correlated to a previously-issued request -//! by `id`. Routed through the crate-internal response-waiter map. -//! 2. **`tool_call_progress` notifications** correlated to a per-call -//! `tool_call_id` carried in `params`. Routed through the -//! crate-internal progress-waiter map registered via -//! `Demux::try_register_progress_waiter` (crate-internal). -//! 3. JSON-RPC **requests / notifications** carrying a `session_id` — -//! routed to the per-session inbox registered via -//! [`Demux::register_session_inbox`]. -//! 4. Connection-level frames (handshake, ping/pong) that the -//! connection actor handles directly without going through the demux. -//! -//! The demux owns the session inbox map, the in-flight response -//! waiters, and the per-call progress waiters; the connection actor -//! parses each text frame, classifies it, -//! and pushes it through this module. -//! -//! Routing inbound frames is non-blocking: a full session inbox or a -//! dropped receiver returns a typed [`RouteOutcome`] variant rather -//! than awaiting the inbox. Blocking on a slow consumer would back up -//! the entire connection actor and starve every other session sharing -//! the socket. - -use dashmap::DashMap; -use serde_json::Value; -use tokio::sync::oneshot; -use tracing::warn; - -use kigi_tool_protocol::{ - JsonRpcId, JsonRpcResponse, RequestId, SessionId, ToolCallId, ToolCallProgressFrame, -}; - -use crate::error::ClientError; - -/// Frame routed to a session inbox. -#[derive(Debug, Clone)] -pub enum InboundFrame { - /// A request the inbox owner must answer (any session frame carrying - /// an `id`): a `tool_call_request`, or a reverse-direction `hook` - /// answered via `ToolHarness::send_hook_reply`. Carries raw JSON. - Request(Value), - /// Server-issued notification (e.g. `tool.notification`) — fire-and- - /// forget, no reply expected. - Notification(Value), -} - -/// Outcome of [`Demux::route`]. -#[derive(Debug, PartialEq, Eq)] -pub enum RouteOutcome { - /// Matched a response waiter; the oneshot was fulfilled. - Response, - /// Forwarded to a session inbox. - Session, - /// Matched a progress waiter; the progress frame was forwarded to - /// the per-call progress channel. - Progress, - /// No inbox is bound for the targeted session. - UnknownSession, - /// No progress waiter is parked for the targeted `tool_call_id`. The - /// caller's stream is no longer subscribed (typical post-terminal), - /// so the frame is dropped. - UnknownProgress, - /// Connection-level notification broadcast to subscribers. - Notification, - /// No waiter is parked for the targeted request id, OR the frame - /// was unaddressable. - Unrouted, - /// The session inbox sender was full; the frame was dropped to - /// avoid blocking the connection actor. - InboxFull, - /// The session inbox receiver was dropped (e.g. the consumer's - /// run loop exited); the binding is now stale and the frame was - /// dropped. - SessionDropped, - /// The progress channel was full; the frame was dropped to avoid - /// blocking the connection actor. The caller's stream consumer - /// fell behind on draining progress. - ProgressFull, - /// The progress receiver was dropped (e.g. the caller's stream was - /// dropped); the waiter binding is now stale and the frame was - /// dropped. - ProgressDropped, -} - -/// Demux state. Cheap to construct; uses [`DashMap`] internally so -/// concurrent registers and routes never block each other. -#[derive(Debug)] -pub struct Demux { - sessions: DashMap>, - waiters: DashMap>>, - /// Session index for `tool.call` response waiters only. Lets the SDK - /// in-flight short-circuit fail every parked call for a session on a - /// workspace Disconnected notification without waiting for the server. - /// Turn-hook / session-RPC waiters are NOT indexed here, so the - /// short-circuit never touches them. - call_sessions: DashMap, - progress: DashMap>, - /// Broadcast channel for connection-level notifications (no session_id). - notifications: tokio::sync::broadcast::Sender, - /// Clone of the connection's outbound sender. Used to synthesize the - /// overloaded (-32016) response when a session inbox is full so a - /// Request is rejected with an error rather than silently dropped. - /// `None` in unit tests that construct a bare demux. - outbound: Option>, -} - -impl Default for Demux { - fn default() -> Self { - let (notifications, _) = tokio::sync::broadcast::channel(64); - Self { - sessions: DashMap::new(), - waiters: DashMap::new(), - call_sessions: DashMap::new(), - progress: DashMap::new(), - notifications, - outbound: None, - } - } -} - -impl Demux { - pub fn new() -> Self { - Self::default() - } - - /// Construct a demux wired to the connection's outbound sender so the - /// inbox-full Request path can ship an overloaded (-32016) response. - pub fn with_outbound(outbound: tokio::sync::mpsc::Sender) -> Self { - Self { - outbound: Some(outbound), - ..Self::default() - } - } - - /// Subscribe to connection-level notifications (no session_id). - pub fn subscribe_notifications(&self) -> tokio::sync::broadcast::Receiver { - self.notifications.subscribe() - } - - /// Bind `session_id` to `inbox`; replaces any existing binding. - /// Returns the previous sender if one existed; the caller may - /// drop or drain it as appropriate. - pub fn register_session_inbox( - &self, - session_id: SessionId, - inbox: tokio::sync::mpsc::Sender, - ) -> Option> { - self.sessions.insert(session_id, inbox) - } - - /// Remove the inbox bound to `session_id`. The returned sender (if - /// present) is dropped by the caller, signalling EOF to its - /// receiver task. - pub fn unregister_session_inbox( - &self, - session_id: &SessionId, - ) -> Option> { - self.sessions.remove(session_id).map(|(_, sender)| sender) - } - - /// Park a oneshot waiter for `request_id`. Crate-internal: only - /// the connection actor allocates request ids. - pub(crate) fn register_response_waiter( - &self, - request_id: RequestId, - waiter: oneshot::Sender>, - ) { - self.waiters.insert(request_id, waiter); - } - - /// Park a `tool.call` response waiter and record its `session_id` so the - /// SDK in-flight short-circuit ([`Self::fail_calls_for_session`]) can - /// resolve it on a workspace Disconnected notification. Crate-internal. - pub(crate) fn register_call_response_waiter( - &self, - request_id: RequestId, - session_id: SessionId, - waiter: oneshot::Sender>, - ) { - self.call_sessions.insert(request_id.clone(), session_id); - self.waiters.insert(request_id, waiter); - } - - /// Pop the waiter for `request_id`, if any. Also drops the session index - /// entry so the two maps stay consistent. Crate-internal. - pub(crate) fn take_response_waiter( - &self, - request_id: &RequestId, - ) -> Option>> { - self.call_sessions.remove(request_id); - self.waiters.remove(request_id).map(|(_, waiter)| waiter) - } - - /// Fail every in-flight `tool.call` waiter bound to `session_id`, - /// completing each with `result_factory`. Returns the number resolved. - /// - /// Drives the SDK in-flight short-circuit: on a workspace - /// `ToolServerStatusChanged(Disconnected)` notification the harness fails - /// its parked calls for that session promptly instead of parking until - /// `rpc_ttl_ms`. Idempotent with the server-side cancel — each waiter is - /// taken at most once, so a call already resolved by the server is skipped. - pub(crate) fn fail_calls_for_session( - &self, - session_id: &SessionId, - result_factory: F, - ) -> usize - where - F: Fn() -> ClientError, - { - // Snapshot the matching request ids first so we never hold a DashMap - // shard lock across the oneshot send. - let request_ids: Vec = self - .call_sessions - .iter() - .filter(|kv| kv.value() == session_id) - .map(|kv| kv.key().clone()) - .collect(); - let mut resolved = 0; - for request_id in request_ids { - if let Some(waiter) = self.take_response_waiter(&request_id) - && waiter.send(Err(result_factory())).is_ok() - { - resolved += 1; - } - } - resolved - } - - /// Park a per-call progress sender keyed by `tool_call_id`. - /// Returns `Err(progress)` (handing the not-yet-inserted sender - /// back) when another in-flight call already owns the id, leaving - /// the prior waiter intact. The caller drops the matching - /// receiver to terminate the subscription — subsequent inbound - /// progress for the same id is silently dropped via - /// [`RouteOutcome::ProgressDropped`]. - /// - /// Atomic check-then-insert under a single shard lock so a - /// concurrent caller cannot observe a transient empty slot. - pub(crate) fn try_register_progress_waiter( - &self, - tool_call_id: ToolCallId, - progress: tokio::sync::mpsc::Sender, - ) -> Result<(), tokio::sync::mpsc::Sender> { - use dashmap::mapref::entry::Entry; - match self.progress.entry(tool_call_id) { - Entry::Occupied(_) => Err(progress), - Entry::Vacant(slot) => { - slot.insert(progress); - Ok(()) - } - } - } - - /// Remove the progress sender bound to `tool_call_id`. Crate-internal; - /// called by the harness once the terminal frame for `tool_call_id` - /// has been observed. - pub(crate) fn unregister_progress_waiter( - &self, - tool_call_id: &ToolCallId, - ) -> Option> { - self.progress.remove(tool_call_id).map(|(_, tx)| tx) - } - - /// Drain every parked waiter, completing each with `result_factory`. - /// Used by the reconnect path to fast-fail in-flight calls with - /// [`ClientError::NetworkError`]. Crate-internal. - pub(crate) fn drain_waiters_with(&self, result_factory: F) - where - F: Fn() -> ClientError, - { - // Snapshot keys, then remove individually so we never hold - // a DashMap shard lock across the oneshot send. - let keys: Vec = self.waiters.iter().map(|kv| kv.key().clone()).collect(); - for key in keys { - if let Some((_, waiter)) = self.waiters.remove(&key) { - self.call_sessions.remove(&key); - let _ = waiter.send(Err(result_factory())); - } - } - } - - /// Drop every parked progress sender. Used by the reconnect path - /// after [`Self::drain_waiters_with`]: the response waiter resolves - /// with `NetworkError` and the matching progress channel closes, - /// so any in-flight harness call stream terminates promptly - /// instead of stalling on a half-empty progress channel. - pub(crate) fn drain_progress(&self) { - let keys: Vec = self.progress.iter().map(|kv| kv.key().clone()).collect(); - for key in keys { - self.progress.remove(&key); - } - } - - /// Route a parsed JSON value. Classification rules: - /// - /// - presence of `result`/`error` → response, routed to waiter; - /// - method == `tool_call_progress` notification → progress waiter - /// keyed by `params.tool_call_id`; - /// - presence of `session_id` → session inbox, request vs. - /// notification distinguished by the presence of `id`; - /// - otherwise → [`RouteOutcome::Unrouted`]. - /// - /// Routing to a session inbox or progress channel uses non-blocking - /// `try_send`. A full inbox or progress channel returns the matching - /// `*Full` variant; a dropped receiver returns the matching - /// `*Dropped` variant. Either way the frame is dropped without - /// awaiting the consumer, so a slow handler never starves other - /// sessions or calls multiplexed onto the same connection. - pub fn route(&self, frame: Value) -> RouteOutcome { - crate::metrics::demux_inbox_depth_set(self.sessions.len() as i64); - if frame.get("result").is_some() || frame.get("error").is_some() { - return self.route_response(frame); - } - if frame.get("method").and_then(Value::as_str) == Some("tool_call_progress") { - return self.route_progress(frame); - } - if frame.get("session_id").is_some() { - return self.route_session(frame); - } - // Connection-level notification (e.g. session.bind, session.unbind). - if frame.get("method").is_some() { - let _ = self.notifications.send(frame); - return RouteOutcome::Notification; - } - RouteOutcome::Unrouted - } - - fn route_progress(&self, frame: Value) -> RouteOutcome { - let Some(params) = frame.get("params") else { - return RouteOutcome::Unrouted; - }; - let Some(call_id_str) = params.get("tool_call_id").and_then(Value::as_str) else { - return RouteOutcome::Unrouted; - }; - let Ok(tool_call_id) = ToolCallId::new(call_id_str) else { - return RouteOutcome::Unrouted; - }; - let Some(sender) = self.progress.get(&tool_call_id) else { - return RouteOutcome::UnknownProgress; - }; - let tx = sender.value().clone(); - drop(sender); - let progress_frame: ToolCallProgressFrame = match serde_json::from_value(params.clone()) { - Ok(p) => p, - Err(err) => { - warn!(%tool_call_id, ?err, "failed to decode tool_call_progress params"); - return RouteOutcome::Unrouted; - } - }; - match tx.try_send(progress_frame) { - Ok(()) => RouteOutcome::Progress, - Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { - warn!(%tool_call_id, "progress channel full; dropping inbound progress frame"); - RouteOutcome::ProgressFull - } - Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => { - self.progress.remove(&tool_call_id); - RouteOutcome::ProgressDropped - } - } - } - - fn route_response(&self, frame: Value) -> RouteOutcome { - let Some(id_value) = frame.get("id") else { - return RouteOutcome::Unrouted; - }; - let request_id = match id_value { - Value::String(s) => RequestId::new(s.as_str()).ok(), - Value::Number(n) => RequestId::new(n.to_string()).ok(), - _ => None, - }; - let Some(request_id) = request_id else { - return RouteOutcome::Unrouted; - }; - let Some(waiter) = self.take_response_waiter(&request_id) else { - return RouteOutcome::Unrouted; - }; - let parsed: Result = - serde_json::from_value::(frame).map_err(ClientError::from); - let _ = waiter.send(parsed); - RouteOutcome::Response - } - - fn route_session(&self, frame: Value) -> RouteOutcome { - let Some(sid_str) = frame.get("session_id").and_then(Value::as_str) else { - return RouteOutcome::Unrouted; - }; - let Ok(session_id) = SessionId::new(sid_str) else { - return RouteOutcome::Unrouted; - }; - let Some(sender) = self.sessions.get(&session_id) else { - return RouteOutcome::UnknownSession; - }; - let inbox = sender.value().clone(); - drop(sender); - let kind = if frame.get("id").is_some() { - InboundFrame::Request(frame) - } else { - InboundFrame::Notification(frame) - }; - match inbox.try_send(kind) { - Ok(()) => RouteOutcome::Session, - Err(tokio::sync::mpsc::error::TrySendError::Full(frame)) => { - self.reject_inbox_full(&session_id, frame); - RouteOutcome::InboxFull - } - Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => { - warn!(%session_id, "session inbox dropped; binding stale"); - self.sessions.remove(&session_id); - RouteOutcome::SessionDropped - } - } - } - - /// Handle a full session inbox without blocking the reader. - /// - /// A Request (has an `id`) is rejected with the shared overloaded - /// (-32016 "tool_busy") response on a best-effort `try_send`; if the - /// outbound is *also* full the rejection itself is dropped and metered - /// (`inbox_full_reject_send_failed`). A Notification (no `id`) stays - /// fire-and-forget and is metered (`inbox_full_notification_dropped`). - fn reject_inbox_full(&self, session_id: &SessionId, frame: InboundFrame) { - let InboundFrame::Request(value) = frame else { - crate::metrics::inbox_full_notification_dropped(); - return; - }; - crate::metrics::inbox_full_request_rejected(); - warn!(%session_id, "session inbox full; rejecting request with tool_busy"); - let Some(out) = &self.outbound else { - return; - }; - // A `Request` always carries an `id` (that is how `route_session` - // classifies it). A well-formed id deserializes into a `JsonRpcId`; - // a malformed id (object/array/bool/null) cannot, but the request - // must STILL get an overloaded response rather than be silently - // dropped, so we fall back to echoing the raw id JSON as a string. - let raw_id = value.get("id"); - let id = raw_id - .and_then(|v| serde_json::from_value::(v.clone()).ok()) - .unwrap_or_else(|| { - JsonRpcId::new_string(raw_id.map(ToString::to_string).unwrap_or_default()) - }); - let response = crate::admission::overloaded_response(id, session_id.clone()); - let Ok(text) = serde_json::to_string(&response) else { - return; - }; - if out.try_send(text).is_err() { - crate::metrics::inbox_full_reject_send_failed(); - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - use tokio::sync::mpsc; - - #[tokio::test] - async fn response_route_matches_waiter() { - let demux = Demux::new(); - let request_id = RequestId::new("r1").expect("valid"); - let (tx, rx) = oneshot::channel(); - demux.register_response_waiter(request_id.clone(), tx); - let outcome = demux.route(json!({ - "jsonrpc": "2.0", - "id": "r1", - "result": {"outcome": "bound"}, - })); - assert_eq!(outcome, RouteOutcome::Response); - let resp = rx.await.expect("waiter").expect("ok"); - assert_eq!(resp.id.to_string(), "r1"); - } - - #[tokio::test] - async fn fail_calls_for_session_resolves_only_matching_call_waiters() { - // Fails exactly the session's `tool.call` waiters; other sessions' - // calls and non-call (turn-hook) waiters stay parked. - let demux = Demux::new(); - let s1 = SessionId::new("s1").expect("valid"); - let s2 = SessionId::new("s2").expect("valid"); - - let (tx_a, rx_a) = oneshot::channel(); - let (tx_b, rx_b) = oneshot::channel(); - let (tx_other, rx_other) = oneshot::channel(); - // Two calls on s1, one on s2. - demux.register_call_response_waiter(RequestId::new("a").unwrap(), s1.clone(), tx_a); - demux.register_call_response_waiter(RequestId::new("b").unwrap(), s1.clone(), tx_b); - demux.register_call_response_waiter(RequestId::new("c").unwrap(), s2.clone(), tx_other); - // A non-call waiter (e.g. a turn hook) on s1 — NOT session-indexed. - let (tx_hook, rx_hook) = oneshot::channel(); - demux.register_response_waiter(RequestId::new("hook").unwrap(), tx_hook); - - let n = demux.fail_calls_for_session(&s1, || ClientError::NetworkError("gone".to_owned())); - assert_eq!(n, 2, "only the two s1 call waiters are failed"); - - assert!(matches!(rx_a.await, Ok(Err(ClientError::NetworkError(_))))); - assert!(matches!(rx_b.await, Ok(Err(ClientError::NetworkError(_))))); - // s2's call and the turn-hook waiter are untouched (still parked). - assert!( - demux - .take_response_waiter(&RequestId::new("c").unwrap()) - .is_some(), - "the s2 call must remain parked" - ); - assert!( - demux - .take_response_waiter(&RequestId::new("hook").unwrap()) - .is_some(), - "the turn-hook waiter must remain parked" - ); - // Keep the receivers alive until the asserts above ran. - drop((rx_other, rx_hook)); - } - - #[tokio::test] - async fn fail_calls_for_session_is_idempotent_after_resolution() { - // A call already resolved (waiter taken) must not be double-counted by - // the short-circuit. - let demux = Demux::new(); - let s1 = SessionId::new("s1").expect("valid"); - let (tx_a, rx_a) = oneshot::channel(); - demux.register_call_response_waiter(RequestId::new("a").unwrap(), s1.clone(), tx_a); - // Simulate the server-side resolution taking the waiter first. - let waiter = demux - .take_response_waiter(&RequestId::new("a").unwrap()) - .expect("waiter present"); - drop(waiter); - drop(rx_a); - let n = demux.fail_calls_for_session(&s1, || ClientError::NetworkError("gone".to_owned())); - assert_eq!(n, 0, "already-resolved call is not re-failed"); - } - - #[tokio::test] - async fn short_circuit_then_late_response_is_unrouted() { - // A short-circuit that resolves first leaves no waiter, so a late server - // response for the same id is dropped (no double-resolve). - let demux = Demux::new(); - let s1 = SessionId::new("s1").expect("valid"); - let (tx_a, rx_a) = oneshot::channel(); - demux.register_call_response_waiter(RequestId::new("a").unwrap(), s1.clone(), tx_a); - - let n = demux.fail_calls_for_session(&s1, || ClientError::NetworkError("gone".to_owned())); - assert_eq!(n, 1); - assert!(matches!(rx_a.await, Ok(Err(ClientError::NetworkError(_))))); - - let outcome = demux.route(json!({ "jsonrpc": "2.0", "id": "a", "result": {} })); - assert_eq!( - outcome, - RouteOutcome::Unrouted, - "the late normal response must not double-resolve the call" - ); - } - - #[tokio::test] - async fn session_route_pushes_to_inbox() { - let demux = Demux::new(); - let session = SessionId::new("s1").expect("valid"); - let (tx, mut rx) = mpsc::channel(4); - demux.register_session_inbox(session.clone(), tx); - let frame = json!({ - "jsonrpc": "2.0", - "id": "x", - "session_id": "s1", - "method": "tool_call_request", - "params": {}, - }); - let outcome = demux.route(frame.clone()); - assert_eq!(outcome, RouteOutcome::Session); - match rx.recv().await { - Some(InboundFrame::Request(value)) => assert_eq!(value, frame), - other => panic!("expected request inbound; got {other:?}"), - } - } - - #[tokio::test] - async fn reverse_hook_request_routes_to_inbox_as_request() { - // A reverse hook request carries an `id`, so it must route to the - // inbox as `Request` (not `Notification`) for the harness to answer. - let demux = Demux::new(); - let session = SessionId::new("s1").expect("valid"); - let (tx, mut rx) = mpsc::channel(4); - demux.register_session_inbox(session.clone(), tx); - let hook = kigi_tool_protocol::HookFrame::custom_request( - session.clone(), - "hook-7".to_owned(), - crate::harness::PERMISSION_REQUEST_KIND.to_owned(), - json!({}), - ); - let frame = json!({ - "jsonrpc": "2.0", - "id": "h1", - "session_id": "s1", - "method": kigi_tool_protocol::Method::Hook.as_wire_str(), - "params": serde_json::to_value(&hook).expect("serialize hook"), - }); - assert_eq!(demux.route(frame.clone()), RouteOutcome::Session); - match rx.recv().await { - Some(InboundFrame::Request(value)) => assert_eq!(value, frame), - other => panic!("expected request inbound; got {other:?}"), - } - } - - #[tokio::test] - async fn notification_classified_without_id() { - let demux = Demux::new(); - let session = SessionId::new("s1").expect("valid"); - let (tx, mut rx) = mpsc::channel(4); - demux.register_session_inbox(session.clone(), tx); - let frame = json!({ - "jsonrpc": "2.0", - "session_id": "s1", - "method": "tool.notification", - "params": {}, - }); - let outcome = demux.route(frame); - assert_eq!(outcome, RouteOutcome::Session); - match rx.recv().await { - Some(InboundFrame::Notification(_)) => {} - other => panic!("expected notification; got {other:?}"), - } - } - - #[tokio::test] - async fn unknown_session_returns_unknown_session() { - let demux = Demux::new(); - let outcome = demux.route( - json!({"jsonrpc":"2.0","id":"x","session_id":"missing","method":"x","params":{}}), - ); - assert_eq!(outcome, RouteOutcome::UnknownSession); - } - - #[tokio::test] - async fn unknown_request_id_returns_unrouted() { - let demux = Demux::new(); - let outcome = demux.route(json!({ - "jsonrpc": "2.0", - "id": "missing", - "result": {}, - })); - assert_eq!(outcome, RouteOutcome::Unrouted); - } - - #[tokio::test] - async fn full_inbox_returns_inbox_full_without_blocking() { - let demux = Demux::new(); - let session = SessionId::new("backed_up").expect("valid"); - let (tx, _rx) = mpsc::channel(1); - demux.register_session_inbox(session.clone(), tx); - let frame = || { - json!({ - "jsonrpc": "2.0", - "id": "x", - "session_id": "backed_up", - "method": "tool_call_request", - "params": {}, - }) - }; - // First send fills capacity. - assert_eq!(demux.route(frame()), RouteOutcome::Session); - // Second send must NOT block; it returns InboxFull. - assert_eq!(demux.route(frame()), RouteOutcome::InboxFull); - } - - #[tokio::test] - async fn dropped_receiver_returns_session_dropped() { - let demux = Demux::new(); - let session = SessionId::new("gone").expect("valid"); - let (tx, rx) = mpsc::channel(1); - demux.register_session_inbox(session.clone(), tx); - drop(rx); - let frame = json!({ - "jsonrpc": "2.0", - "id": "x", - "session_id": "gone", - "method": "tool_call_request", - "params": {}, - }); - assert_eq!(demux.route(frame), RouteOutcome::SessionDropped); - // Stale binding should have been removed. - assert!(demux.sessions.get(&session).is_none()); - } - - #[tokio::test] - async fn inbox_full_request_synthesizes_overloaded_response_onto_outbound() { - // A full session inbox for a Request must produce the shared - // -32016 "tool_busy" response on outbound, not a silent drop. - let (out_tx, mut out_rx) = mpsc::channel::(4); - let demux = Demux::with_outbound(out_tx); - let session = SessionId::new("busy").expect("valid"); - let (tx, _rx) = mpsc::channel(1); - demux.register_session_inbox(session.clone(), tx); - let frame = |id: &str| { - json!({ - "jsonrpc": "2.0", - "id": id, - "session_id": "busy", - "method": "tool_call_request", - "params": {}, - }) - }; - // First fills capacity (cap 1); second overflows → InboxFull. - assert_eq!(demux.route(frame("a")), RouteOutcome::Session); - assert_eq!(demux.route(frame("b")), RouteOutcome::InboxFull); - - let text = out_rx.try_recv().expect("overloaded response enqueued"); - let wire: Value = serde_json::from_str(&text).expect("valid json"); - assert_eq!(wire["id"], "b"); - assert_eq!(wire["session_id"], "busy"); - assert_eq!(wire["error"]["code"], -32016); - assert_eq!(wire["error"]["data"]["code"], "tool_busy"); - assert_eq!(wire["error"]["data"]["retryable"], true); - assert!( - out_rx.try_recv().is_err(), - "exactly one rejection emitted for one overflow" - ); - } - - #[tokio::test] - async fn inbox_full_request_with_malformed_id_still_emits_overloaded_response() { - // A Request whose `id` is present but not a valid JsonRpcId - // (object/array/null) must NOT be silently dropped on a full - // inbox: it still gets the shared -32016 response, with the raw - // id echoed back as a string. - let (out_tx, mut out_rx) = mpsc::channel::(4); - let demux = Demux::with_outbound(out_tx); - let session = SessionId::new("bad_id").expect("valid"); - let (tx, _rx) = mpsc::channel(1); - demux.register_session_inbox(session.clone(), tx); - let frame = |id: Value| { - json!({ - "jsonrpc": "2.0", - "id": id, - "session_id": "bad_id", - "method": "tool_call_request", - "params": {}, - }) - }; - // First fills capacity (cap 1); the malformed-id second overflows. - assert_eq!(demux.route(frame(json!("a"))), RouteOutcome::Session); - assert_eq!( - demux.route(frame(json!({ "nested": 1 }))), - RouteOutcome::InboxFull - ); - - let text = out_rx.try_recv().expect("overloaded response enqueued"); - let wire: Value = serde_json::from_str(&text).expect("valid json"); - assert_eq!( - wire["id"], "{\"nested\":1}", - "malformed id is echoed back as its raw JSON text" - ); - assert_eq!(wire["error"]["code"], -32016); - assert_eq!(wire["error"]["data"]["code"], "tool_busy"); - } - - #[tokio::test] - async fn inbox_full_notification_is_dropped_without_outbound_response() { - // A Notification (no id) on a full inbox stays fire-and-forget: - // no synthesized response is emitted. - let (out_tx, mut out_rx) = mpsc::channel::(4); - let demux = Demux::with_outbound(out_tx); - let session = SessionId::new("notif_busy").expect("valid"); - let (tx, _rx) = mpsc::channel(1); - demux.register_session_inbox(session.clone(), tx); - let notif = || { - json!({ - "jsonrpc": "2.0", - "session_id": "notif_busy", - "method": "tool.notification", - "params": {}, - }) - }; - // First notification fills capacity; second overflows. - assert_eq!(demux.route(notif()), RouteOutcome::Session); - assert_eq!(demux.route(notif()), RouteOutcome::InboxFull); - assert!( - out_rx.try_recv().is_err(), - "notifications must not synthesize an outbound response" - ); - } - - #[tokio::test] - async fn inbox_full_request_without_outbound_does_not_panic() { - // A bare demux (no outbound, e.g. unit context) must still report - // InboxFull cleanly when it cannot synthesize a rejection. - let demux = Demux::new(); - let session = SessionId::new("no_out").expect("valid"); - let (tx, _rx) = mpsc::channel(1); - demux.register_session_inbox(session.clone(), tx); - let frame = || { - json!({ - "jsonrpc": "2.0", - "id": "x", - "session_id": "no_out", - "method": "tool_call_request", - "params": {}, - }) - }; - assert_eq!(demux.route(frame()), RouteOutcome::Session); - assert_eq!(demux.route(frame()), RouteOutcome::InboxFull); - } - - #[tokio::test] - async fn progress_route_pushes_to_progress_waiter() { - let demux = Demux::new(); - let call_id = ToolCallId::new_v7(); - let (tx, mut rx) = mpsc::channel(4); - demux - .try_register_progress_waiter(call_id.clone(), tx) - .expect("first registration"); - let frame = json!({ - "jsonrpc": "2.0", - "session_id": "any", - "method": "tool_call_progress", - "params": { - "tool_call_id": call_id.as_str(), - "kind": "log_chunk", - "body": {"text": "hello"}, - }, - }); - let outcome = demux.route(frame); - assert_eq!(outcome, RouteOutcome::Progress); - let progress = rx.recv().await.expect("progress frame"); - assert_eq!(progress.tool_call_id, call_id); - assert_eq!(progress.kind, "log_chunk"); - assert_eq!(progress.body, json!({"text": "hello"})); - } - - #[tokio::test] - async fn progress_with_no_waiter_returns_unknown_progress() { - let demux = Demux::new(); - let call_id = ToolCallId::new_v7(); - let frame = json!({ - "jsonrpc": "2.0", - "session_id": "any", - "method": "tool_call_progress", - "params": { - "tool_call_id": call_id.as_str(), - "kind": "log_chunk", - "body": {}, - }, - }); - assert_eq!(demux.route(frame), RouteOutcome::UnknownProgress); - } - - #[tokio::test] - async fn dropped_progress_receiver_returns_progress_dropped() { - let demux = Demux::new(); - let call_id = ToolCallId::new_v7(); - let (tx, rx) = mpsc::channel(1); - demux - .try_register_progress_waiter(call_id.clone(), tx) - .expect("first registration"); - drop(rx); - let frame = json!({ - "jsonrpc": "2.0", - "session_id": "any", - "method": "tool_call_progress", - "params": { - "tool_call_id": call_id.as_str(), - "kind": "x", - "body": {}, - }, - }); - assert_eq!(demux.route(frame), RouteOutcome::ProgressDropped); - assert!(demux.progress.get(&call_id).is_none()); - } - - #[tokio::test] - async fn unregister_progress_waiter_returns_sender_when_present() { - let demux = Demux::new(); - let call_id = ToolCallId::new_v7(); - let (tx, _rx) = mpsc::channel::(1); - demux - .try_register_progress_waiter(call_id.clone(), tx) - .expect("first registration"); - assert!(demux.unregister_progress_waiter(&call_id).is_some()); - assert!(demux.unregister_progress_waiter(&call_id).is_none()); - } - - #[tokio::test] - async fn try_register_progress_waiter_rejects_collision_and_preserves_existing() { - let demux = Demux::new(); - let call_id = ToolCallId::new_v7(); - let (tx_first, mut rx_first) = mpsc::channel::(1); - demux - .try_register_progress_waiter(call_id.clone(), tx_first) - .expect("first registration"); - let (tx_second, _rx_second) = mpsc::channel::(1); - let returned = demux - .try_register_progress_waiter(call_id.clone(), tx_second) - .expect_err("collision returns the rejected sender"); - // Returned sender is independent of the live one: dropping - // it must not close the original receiver. - drop(returned); - let frame = json!({ - "jsonrpc": "2.0", - "session_id": "any", - "method": "tool_call_progress", - "params": { - "tool_call_id": call_id.as_str(), - "kind": "log_chunk", - "body": {"text": "first"}, - }, - }); - assert_eq!(demux.route(frame), RouteOutcome::Progress); - let progress = rx_first.recv().await.expect("original receiver still live"); - assert_eq!(progress.body, json!({"text": "first"})); - } - - #[tokio::test] - async fn full_progress_channel_returns_progress_full_without_blocking() { - let demux = Demux::new(); - let call_id = ToolCallId::new_v7(); - let (tx, _rx) = mpsc::channel::(1); - demux - .try_register_progress_waiter(call_id.clone(), tx) - .expect("first registration"); - let frame = || { - json!({ - "jsonrpc": "2.0", - "session_id": "any", - "method": "tool_call_progress", - "params": { - "tool_call_id": call_id.as_str(), - "kind": "x", - "body": {}, - }, - }) - }; - // First send fills capacity (mpsc(1)). - assert_eq!(demux.route(frame()), RouteOutcome::Progress); - // Second send must NOT block; it returns ProgressFull. - assert_eq!(demux.route(frame()), RouteOutcome::ProgressFull); - } - - #[tokio::test] - async fn drain_progress_removes_all_waiters_and_drops_senders() { - let demux = Demux::new(); - let call_a = ToolCallId::new_v7(); - let call_b = ToolCallId::new_v7(); - let (tx_a, mut rx_a) = mpsc::channel::(1); - let (tx_b, mut rx_b) = mpsc::channel::(1); - demux - .try_register_progress_waiter(call_a.clone(), tx_a) - .expect("first registration"); - demux - .try_register_progress_waiter(call_b.clone(), tx_b) - .expect("first registration"); - assert_eq!(demux.progress.len(), 2); - - demux.drain_progress(); - - // Post-drain: every entry removed. - assert_eq!(demux.progress.len(), 0); - assert!(demux.progress.get(&call_a).is_none()); - assert!(demux.progress.get(&call_b).is_none()); - // The senders held by the demux were dropped, so each - // receiver sees `None` (channel closed). - assert!( - rx_a.recv().await.is_none(), - "sender dropped → receiver closes" - ); - assert!( - rx_b.recv().await.is_none(), - "sender dropped → receiver closes" - ); - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/donate_pump.rs b/crates/common/kigi-computer-hub-sdk/src/donate_pump.rs deleted file mode 100644 index 3307408..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/donate_pump.rs +++ /dev/null @@ -1,206 +0,0 @@ -//! Shared donation transport: a bounded retry buffer + in-order drain -//! barrier, parameterized over a `donate` closure. Traces, logs, and -//! metrics all pump through this; failed sends are retained briefly, -//! overflow drops payloads — telemetry, never correctness. - -use std::collections::VecDeque; -use std::time::{SystemTime, UNIX_EPOCH}; - -use opentelemetry_proto::tonic::common::v1::{AnyValue, KeyValue, any_value}; -use opentelemetry_proto::tonic::resource::v1::Resource; -use tokio::sync::{mpsc, oneshot}; - -/// Bound on payloads queued before the pump drains them. -pub(crate) const PENDING_FLUSHES: usize = 8; -/// Payloads retained across failed sends (disconnect/reconnect window). -pub(crate) const RETRY_CAP: usize = 8; - -// --------------------------------------------------------------------------- -// Shared OTLP encoding helpers -// -// Reused by the log and metric donation clients so the AnyValue/KeyValue/ -// Resource construction lives in one place instead of being copy-pasted per -// client. (`trace_donate` builds its payload via `opentelemetry_sdk`'s own -// conversion and does not use these.) -// --------------------------------------------------------------------------- - -/// Current wall-clock time as Unix-epoch nanoseconds (OTLP `time_unix_nano`). -pub(crate) fn now_unix_nanos() -> u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_nanos() as u64) - .unwrap_or(0) -} - -/// OTLP string `AnyValue`. -pub(crate) fn string_value(s: String) -> AnyValue { - AnyValue { - value: Some(any_value::Value::StringValue(s)), - } -} - -/// OTLP string-valued `KeyValue`. -pub(crate) fn string_kv(key: &str, value: String) -> KeyValue { - KeyValue { - key: key.to_owned(), - value: Some(string_value(value)), - ..Default::default() - } -} - -/// OTLP `Resource` carrying just `service.name`. -pub(crate) fn make_resource(service_name: String) -> Resource { - Resource { - attributes: vec![string_kv("service.name", service_name)], - ..Default::default() - } -} - -pub(crate) enum PumpMsg { - /// Base64 OTLP request, ready for the wire. - Payload(String), - /// In-order drain fence — a barrier, not a timeout. - Barrier(oneshot::Sender<()>), -} - -/// Resolves once every payload queued before this call has had a send -/// attempt. Call after the producer's flush (e.g. `fastrace::flush()`). -pub(crate) async fn drain_via(tx: &mpsc::Sender) { - let (ack_tx, ack_rx) = oneshot::channel(); - if tx.send(PumpMsg::Barrier(ack_tx)).await.is_ok() { - let _ = ack_rx.await; - } -} - -/// `donate` hands the payload back so a failed send retains it -/// without cloning. -pub(crate) async fn run_pump(mut rx: mpsc::Receiver, donate: D) -where - D: Fn(String) -> F, - F: std::future::Future, -{ - let mut retry: VecDeque = VecDeque::new(); - while let Some(msg) = rx.recv().await { - match msg { - PumpMsg::Payload(payload) => { - if retry.len() == RETRY_CAP { - retry.pop_front(); - tracing::debug!("donation retry buffer full; dropping oldest payload"); - } - retry.push_back(payload); - } - PumpMsg::Barrier(ack) => { - attempt_sends(&mut retry, &donate).await; - let _ = ack.send(()); - continue; - } - } - attempt_sends(&mut retry, &donate).await; - } -} - -/// Send in order, stopping at the first failure; the remainder stays -/// queued for the next wake. -async fn attempt_sends(retry: &mut VecDeque, donate: &D) -where - D: Fn(String) -> F, - F: std::future::Future, -{ - while let Some(payload) = retry.pop_front() { - let (ok, payload) = donate(payload).await; - if !ok { - tracing::debug!("donation send failed; retaining payload for retry"); - retry.push_front(payload); - break; - } - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - use std::sync::atomic::{AtomicBool, Ordering}; - - use parking_lot::Mutex; - - use super::*; - - fn payload(tag: u64) -> PumpMsg { - PumpMsg::Payload(format!("payload-{tag}")) - } - - /// The drain barrier acks even while the link is down. - #[tokio::test] - async fn pump_retries_failed_payloads_across_reconnect() { - let healthy = Arc::new(AtomicBool::new(false)); - let sent: Arc>> = Arc::new(Mutex::new(Vec::new())); - let (tx, rx) = mpsc::channel::(PENDING_FLUSHES); - let pump = { - let healthy = Arc::clone(&healthy); - let sent = Arc::clone(&sent); - tokio::spawn(run_pump(rx, move |p: String| { - let healthy = Arc::clone(&healthy); - let sent = Arc::clone(&sent); - async move { - if healthy.load(Ordering::SeqCst) { - sent.lock().push(p.clone()); - (true, p) - } else { - (false, p) - } - } - })) - }; - - tx.send(payload(1)).await.unwrap(); - tx.send(payload(2)).await.unwrap(); - drain_via(&tx).await; - assert!(sent.lock().is_empty(), "nothing sent while link is down"); - - healthy.store(true, Ordering::SeqCst); - drain_via(&tx).await; - assert_eq!(*sent.lock(), vec!["payload-1", "payload-2"]); - - drop(tx); - pump.await.expect("pump must exit cleanly"); - } - - #[tokio::test] - async fn pump_retry_buffer_drops_oldest_beyond_cap() { - let sent: Arc>> = Arc::new(Mutex::new(Vec::new())); - let healthy = Arc::new(AtomicBool::new(false)); - let (tx, rx) = mpsc::channel::(RETRY_CAP + 2); - let pump = { - let healthy = Arc::clone(&healthy); - let sent = Arc::clone(&sent); - tokio::spawn(run_pump(rx, move |p: String| { - let healthy = Arc::clone(&healthy); - let sent = Arc::clone(&sent); - async move { - if healthy.load(Ordering::SeqCst) { - sent.lock().push(p.clone()); - (true, p) - } else { - (false, p) - } - } - })) - }; - - for i in 0..=(RETRY_CAP as u64) { - tx.send(payload(i + 1)).await.unwrap(); - } - drain_via(&tx).await; - - healthy.store(true, Ordering::SeqCst); - drain_via(&tx).await; - { - let sent = sent.lock(); - assert_eq!(sent.len(), RETRY_CAP, "buffer bounded at RETRY_CAP"); - assert_eq!(sent[0], "payload-2", "oldest payload evicted first"); - } - - drop(tx); - pump.await.expect("pump must exit cleanly"); - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/error.rs b/crates/common/kigi-computer-hub-sdk/src/error.rs deleted file mode 100644 index f83dec5..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/error.rs +++ /dev/null @@ -1,346 +0,0 @@ -//! Client-side error taxonomy. -//! -//! Wire-level [`kigi_tool_protocol::ToolErrorWire`] variants and JSON-RPC -//! error envelopes are mapped into the smaller [`ClientError`] vocabulary -//! at the SDK boundary so consumers can match on a single enum without -//! re-deriving the numeric/string code mapping. - -use kigi_tool_protocol::{IdError, JsonRpcError, ToolCallId, ToolErrorWire}; -use thiserror::Error; -use url::Url; - -/// Errors surfaced by the client SDK. -#[derive(Debug, Error)] -pub enum ClientError { - /// WebSocket transport failure: failed to connect, dropped socket, - /// or in-flight request interrupted by a reconnect cycle. - #[error("network error: {0}")] - NetworkError(String), - - /// Wire-protocol violation: malformed JSON, unexpected method, - /// hello/hello_ack mismatch, or unsupported `protocol_version`. - #[error("protocol error: {0}")] - ProtocolError(String), - - /// Authentication or authorisation rejected by the server. - #[error("auth error: {0}")] - AuthError(String), - - /// Server rejected the WebSocket upgrade with an HTTP auth status - /// (401/403). Non-retryable: replaying the same credential is - /// rejected identically, so the reconnect loop classifies this as - /// fatal instead of retrying forever. - #[error("handshake auth failed: HTTP {status}")] - HandshakeAuthFailed { status: u16 }, - - /// `register_tool` / `register_session` ack reported a conflict - /// (cross-connection contention or an already-bound entry the - /// caller did not expect). - #[error("registration conflict: {0}")] - RegistrationConflict(String), - - /// Outbound mpsc full or call-site bounded wait elapsed before the - /// frame could be enqueued. Distinct from [`Self::NetworkError`]: - /// the socket may still be healthy. - #[error("backpressure: {0}")] - BackpressureError(String), - - /// JSON serialise / deserialise failure inside the SDK. - #[error("serde error: {0}")] - Serde(String), - - /// Builder consistency error: missing URL, missing auth, etc. - #[error("invalid configuration: {0}")] - InvalidConfig(String), - - /// Wrapped wire-format tool error; surfaces the upstream - /// [`ToolErrorWire`] variant verbatim for callers that need to - /// switch on the stable string code. - #[error(transparent)] - Wire(ToolErrorWire), - - /// Server-side close / shutdown signal received during steady state. - #[error("server closed connection: {0}")] - Closed(String), - - /// Refused to send credentials over an insecure `ws://` scheme to a - /// non-loopback host. Local-loopback (`127.0.0.1`, `::1`, - /// `localhost`) is the only exception; every other host MUST be - /// reached over `wss://` so the bearer token never crosses the - /// network in plaintext. - #[error( - "insecure scheme: refusing to send credentials over plaintext ws:// to non-loopback host {url}" - )] - InsecureScheme { url: Url }, - - /// Caller passed a `ToolCallId` that already keys an in-flight - /// dispatch on the same connection. The prior call's progress - /// waiter and response correlation are left intact; this error - /// surfaces synchronously so the second caller can retry with a - /// fresh id. Mint a fresh [`ToolCallId::new_v7`] (or use - /// [`kigi_tool_runtime::ToolCallContext::default`], which does so) - /// per call. This is client misuse, not a transport or server - /// failure. - #[error("call_id {call_id} already in flight on this connection")] - CallIdInUse { call_id: ToolCallId }, -} - -impl ClientError { - /// Map a JSON-RPC envelope error into a [`ClientError`]. The - /// envelope's `data` payload (when present) carries the stable - /// [`ToolErrorWire`] discriminator; the numeric `code` is used as a - /// coarse fallback when `data` is absent or undecodable. - pub fn from_jsonrpc_error(err: JsonRpcError) -> Self { - if let Some(data) = err.data - && let Ok(wire) = serde_json::from_value::(data) - { - return Self::from_wire(wire); - } - match err.code { - -32002 | -32003 => Self::AuthError(err.message), - -32004 => Self::NetworkError(err.message), - -32600..=-32500 => Self::ProtocolError(err.message), - _ => Self::Wire(ToolErrorWire::Custom { - subcode: format!("jsonrpc_{}", err.code), - message: err.message, - details: None, - }), - } - } - - /// `true` when a `data`-less envelope collapsed to the given `jsonrpc_` - /// subcode (see [`Self::from_jsonrpc_error`]); shared by the bind recognizers. - fn has_collapsed_jsonrpc_subcode(&self, subcode: &str) -> bool { - matches!( - self, - Self::Wire(ToolErrorWire::Custom { subcode: s, .. }) if s == subcode - ) - } - - /// `true` for the server's "server not found" bind rejection (JSON-RPC `-32601`): - /// no workspace-server is registered for this user. - pub fn is_server_not_found(&self) -> bool { - self.has_collapsed_jsonrpc_subcode("jsonrpc_-32601") - } - - /// `true` for the server's `-32013` "server found but bind did not complete" error - /// (the `ServerBindOutcome::Unavailable` cases). Recognized so the harness - /// re-provisions this recoverable case, distinct from [`Self::is_server_not_found`]. - pub fn is_tool_unavailable(&self) -> bool { - self.has_collapsed_jsonrpc_subcode("jsonrpc_-32013") - } - - /// Map a [`ToolErrorWire`] variant into the SDK error taxonomy. - pub fn from_wire(wire: ToolErrorWire) -> Self { - match wire { - ToolErrorWire::PermissionDenied { reason } => Self::AuthError(reason), - ToolErrorWire::TransportClosed { tool_id } => { - Self::NetworkError(format!("transport closed for {tool_id}")) - } - ToolErrorWire::UnsupportedProtocolVersion { supported } => { - Self::ProtocolError(format!("unsupported protocol; supported: {supported:?}")) - } - other => Self::Wire(other), - } - } -} - -impl From for ClientError { - fn from(err: serde_json::Error) -> Self { - Self::Serde(err.to_string()) - } -} - -impl From for ClientError { - fn from(err: IdError) -> Self { - Self::ProtocolError(err.to_string()) - } -} - -impl From for ClientError { - fn from(err: url::ParseError) -> Self { - Self::InvalidConfig(format!("invalid url: {err}")) - } -} - -impl From for ClientError { - fn from(err: tokio_tungstenite::tungstenite::Error) -> Self { - Self::NetworkError(err.to_string()) - } -} - -impl ClientError { - /// Classify a failed WebSocket upgrade. A `401`/`403` on the HTTP - /// upgrade is a non-retryable auth rejection - /// ([`Self::HandshakeAuthFailed`]); every other failure stays a - /// transport [`Self::NetworkError`] via the blanket `From` impl. The - /// distinction must be made here, before `From` collapses the typed - /// `Http` response status into an opaque string. - pub(crate) fn from_handshake_error(err: tokio_tungstenite::tungstenite::Error) -> Self { - if let tokio_tungstenite::tungstenite::Error::Http(resp) = &err { - let status = resp.status().as_u16(); - if status == 401 || status == 403 { - return Self::HandshakeAuthFailed { status }; - } - } - Self::from(err) - } -} - -impl From for ClientError { - fn from(_: tokio::sync::oneshot::error::RecvError) -> Self { - Self::NetworkError("response waiter dropped (connection closed)".to_owned()) - } -} - -#[cfg(test)] -mod tests { - use kigi_tool_protocol::{ - WORKSPACE_UNAVAILABLE_SUBCODE, WorkspaceGonePhase, WorkspaceGoneReason, - workspace_unavailable_wire, - }; - use serde_json::json; - - use super::*; - - fn workspace_gone_envelope() -> JsonRpcError { - let wire = workspace_unavailable_wire( - WorkspaceGoneReason::Disconnect, - WorkspaceGonePhase::RouteMissing, - ); - JsonRpcError { - code: -32005, - message: "workspace server gone".to_owned(), - data: Some(serde_json::to_value(&wire).unwrap()), - } - } - - fn http_upgrade_error(status: u16) -> tokio_tungstenite::tungstenite::Error { - let resp = tokio_tungstenite::tungstenite::http::Response::builder() - .status(status) - .body(None::>) - .expect("response builds"); - tokio_tungstenite::tungstenite::Error::Http(resp) - } - - #[test] - fn handshake_401_and_403_map_to_handshake_auth_failed() { - for status in [401u16, 403] { - match ClientError::from_handshake_error(http_upgrade_error(status)) { - ClientError::HandshakeAuthFailed { status: got } => assert_eq!(got, status), - other => panic!("expected HandshakeAuthFailed for {status}; got {other:?}"), - } - } - } - - #[test] - fn handshake_non_auth_status_stays_network_error() { - for status in [500u16, 502, 429] { - match ClientError::from_handshake_error(http_upgrade_error(status)) { - ClientError::NetworkError(_) => {} - other => panic!("expected NetworkError for {status}; got {other:?}"), - } - } - } - - #[test] - fn from_jsonrpc_error_preserves_workspace_subcode_and_details() { - // The `data` payload decodes as `ToolErrorWire` first, so the stable - // subcode and structured details reach the SDK consumer intact rather - // than collapsing to the numeric code. - match ClientError::from_jsonrpc_error(workspace_gone_envelope()) { - ClientError::Wire(ToolErrorWire::Custom { - subcode, details, .. - }) => { - assert_eq!(subcode, WORKSPACE_UNAVAILABLE_SUBCODE); - let details = details.expect("details present"); - assert_eq!(details["code"], json!(WORKSPACE_UNAVAILABLE_SUBCODE)); - assert_eq!(details["reason"], json!("disconnect")); - assert_eq!(details["phase"], json!("route_missing")); - assert_eq!(details["retryable"], json!(true)); - } - other => panic!("expected Wire(Custom), got {other:?}"), - } - } - - #[test] - fn is_server_not_found_recognizes_bare_minus_32601() { - // data-less -32601 -> custom subcode. - let err = ClientError::from_jsonrpc_error(JsonRpcError { - code: -32601, - message: "server abc not found for user".to_owned(), - data: None, - }); - assert!(err.is_server_not_found()); - } - - #[test] - fn is_tool_unavailable_recognizes_bare_minus_32013() { - let err = ClientError::from_jsonrpc_error(JsonRpcError { - code: -32013, - message: "server abc did not complete the bind".to_owned(), - data: None, - }); - assert!(err.is_tool_unavailable()); - } - - #[test] - fn is_server_not_found_rejects_other_errors() { - let auth = ClientError::from_jsonrpc_error(JsonRpcError { - code: -32002, - message: "nope".to_owned(), - data: None, - }); - assert!(!auth.is_server_not_found()); - // workspace-gone is the tool-call re-provision path, not bind ServerNotFound. - assert!(!ClientError::from_jsonrpc_error(workspace_gone_envelope()).is_server_not_found()); - } - - #[test] - fn bind_recognizers_are_mutually_exclusive() { - let not_found = ClientError::from_jsonrpc_error(JsonRpcError { - code: -32601, - message: "not found".to_owned(), - data: None, - }); - let unavailable = ClientError::from_jsonrpc_error(JsonRpcError { - code: -32013, - message: "unavailable".to_owned(), - data: None, - }); - assert!(not_found.is_server_not_found()); - assert!( - !not_found.is_tool_unavailable(), - "-32601 must not be recognized as tool_unavailable" - ); - assert!(unavailable.is_tool_unavailable()); - assert!( - !unavailable.is_server_not_found(), - "-32013 must not be recognized as server_not_found" - ); - } - - #[test] - fn sdk_reexported_recognizer_matches_decoded_error() { - // SDK-only consumers reach the recognizer through the SDK re-export and - // the core decode path. - let err = kigi_computer_hub_core::error_from_envelope(workspace_gone_envelope()); - assert!(crate::is_workspace_unavailable(&err)); - } - - #[test] - fn sdk_reexported_recognizer_rejects_unrelated_custom_error() { - let wire = ToolErrorWire::Custom { - subcode: "unrelated".to_owned(), - message: "nope".to_owned(), - details: Some(json!({ "code": "unrelated" })), - }; - let env = JsonRpcError { - code: -32000, - message: "nope".to_owned(), - data: Some(serde_json::to_value(&wire).unwrap()), - }; - let err = kigi_computer_hub_core::error_from_envelope(env); - assert!(!crate::is_workspace_unavailable(&err)); - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/handshake.rs b/crates/common/kigi-computer-hub-sdk/src/handshake.rs deleted file mode 100644 index e50f7f7..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/handshake.rs +++ /dev/null @@ -1,93 +0,0 @@ -//! Hello handshake helpers used by the connection actor and the -//! reconnect-replay path. -//! -//! Splitting these into a dedicated module keeps the connection state -//! machine readable: send the frame, parse the ack, surface a typed -//! [`crate::ClientError`]. - -use futures::{SinkExt, StreamExt}; -use kigi_tool_protocol::{ConnectionKind, HelloAckMsg, HelloMsg}; -use tokio_tungstenite::tungstenite::Message; - -use crate::error::ClientError; - -/// Wire-protocol version both ends speak. Re-exported from the -/// protocol crate so the SDK and the IC service share one source of -/// truth. -pub use kigi_tool_protocol::PROTOCOL_VERSION; - -/// Send the [`HelloMsg`] and wait for the matching [`HelloAckMsg`]. -/// -/// `kind` should be [`ConnectionKind::ToolServer`] for tool-server -/// builds (the only consumer today). The function returns the parsed -/// ack so callers can observe the server-issued `connection_id` and -/// the server-derived `user_id`. -/// -/// When `server_id` is `Some`, it is included in the hello frame so the -/// server can identify itself without a separate `register_server` call. -pub async fn send_hello( - sink: &mut Si, - stream: &mut St, - kind: ConnectionKind, - server_id: Option, - description: Option, - metadata: Option, -) -> Result -where - Si: SinkExt + Unpin, - Si::Error: std::fmt::Display, - St: StreamExt> + Unpin, -{ - let hello = HelloMsg { - protocol_version: PROTOCOL_VERSION.to_owned(), - kind, - server_id, - description, - metadata, - }; - let text = serde_json::to_string(&hello)?; - sink.send(Message::Text(text.into())) - .await - .map_err(|e| ClientError::NetworkError(format!("hello send failed: {e}")))?; - - while let Some(msg) = stream.next().await { - let msg = msg?; - match msg { - Message::Text(text) => { - let ack: HelloAckMsg = serde_json::from_str(text.as_ref()) - .map_err(|e| ClientError::ProtocolError(format!("malformed hello_ack: {e}")))?; - if !ack - .supported_protocol_versions - .iter() - .any(|v| v == PROTOCOL_VERSION) - { - return Err(ClientError::ProtocolError(format!( - "server does not support {PROTOCOL_VERSION}; supported: {:?}", - ack.supported_protocol_versions - ))); - } - return Ok(ack); - } - Message::Ping(payload) => { - sink.send(Message::Pong(payload)) - .await - .map_err(|e| ClientError::NetworkError(format!("pong send failed: {e}")))?; - } - Message::Close(frame) => { - let reason = frame.map(|f| f.reason.to_string()).unwrap_or_default(); - return Err(ClientError::Closed(format!( - "server closed during handshake: {reason}" - ))); - } - Message::Pong(_) | Message::Frame(_) => continue, - Message::Binary(_) => { - return Err(ClientError::ProtocolError( - "server sent binary frame during handshake".to_owned(), - )); - } - } - } - Err(ClientError::NetworkError( - "server closed before hello_ack".to_owned(), - )) -} diff --git a/crates/common/kigi-computer-hub-sdk/src/harness.rs b/crates/common/kigi-computer-hub-sdk/src/harness.rs deleted file mode 100644 index dfb2442..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/harness.rs +++ /dev/null @@ -1,2940 +0,0 @@ -//! Harness-side dispatch surface. -//! -//! A [`ToolHarness`] is the SDK-side counterpart to a server-side -//! `Harness` connection. Build it via [`ToolHarnessBuilder`], seed -//! the in-process [`LocalRegistry`] with `Tool` implementations that -//! should resolve without a wire round-trip, and call -//! [`ToolHarness::call`] to dispatch a tool. -//! -//! Local-first dispatch: every call queries the bound -//! [`LocalRegistry`] first; on a hit the tool's `execute` runs -//! in-process and the returned [`ToolStream`] is forwarded verbatim. -//! Misses fall through to a remote `tool.call` JSON-RPC request over -//! the shared [`HubConnection`]; the demux routes the matching -//! response and any intermediate `tool_call_progress` notifications -//! back into the call's stream. -//! -//! Connection lifecycle mirrors [`crate::ToolServer`]: the harness -//! attaches to a pooled connection under a `(url, principal)` key, -//! refcount-binds its session, and runs cooperative shutdown -//! through a `ConnectionBorrow` (crate-internal) that -//! both ends share. - -use std::pin::Pin; -use std::sync::Arc; -use std::task::{Context, Poll}; - -use dashmap::DashMap; -use futures::FutureExt; -use futures::Stream; -use futures::future::{BoxFuture, Shared}; -use indexmap::IndexMap; -use kigi_computer_hub_core::{ - ErasedTool, ToolHandle, decode_call_result, error_from_envelope, progress_from_frame, - tool_error_from_wire, -}; -use kigi_tool_protocol::notification_wire::{WireCustomNotification, WireToolNotification}; -use kigi_tool_protocol::session_event::{SessionEvent, ToolCallOutcome}; -use kigi_tool_protocol::{ - ConnectionKind, JsonRpcId, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, - JsonRpcVersion, Method, RequestId, ResponseOutcome, SessionId, ToolCallId, ToolCallParams, - ToolCallProgressFrame, ToolId, ToolNotificationFrame, ToolServerLifecycleStatus, - WorkspaceGonePhase, WorkspaceGoneReason, workspace_unavailable_wire, -}; -use kigi_tool_runtime::{ - BehaviorVersion, Cwd, ListToolsContext, Tool, ToolCallContext, ToolError, ToolStream, - ToolStreamItem, TypedToolOutput, terminal_only, -}; -use kigi_tool_types::ToolDescription; -use parking_lot::RwLock; -use serde_json::Value; -use tokio::sync::{mpsc, oneshot}; -use url::Url; - -use crate::auth::{AuthCredential, AuthProvider}; -use crate::connection::{HubConnection, ReconnectCallback, ReconnectEvent}; -use crate::connection_borrow::ConnectionBorrow; -use crate::error::ClientError; -use crate::pool::HubConnectionPool; - -/// Host-supplied source of the current W3C `traceparent`. -pub type TraceContextProvider = Arc Option + Send + Sync>; - -/// Host-registered sink for inbound reverse-direction hook requests -/// (server → harness); invoked by the inbox loop with the decoded -/// [`HookFrame`](kigi_tool_protocol::HookFrame), answered via -/// [`ToolHarness::send_hook_reply`]. -type HookRequestHandler = Arc; - -/// Well-known [`HookEvent::Custom`](kigi_tool_protocol::HookEvent::Custom) kind -/// for a server → harness permission request. Sibling of -/// [`kigi_tool_protocol::turn_hook::TURN_HOOK_KIND`]. -pub const PERMISSION_REQUEST_KIND: &str = "permission_request"; - -/// Buffer size for the per-call progress channel. Picked to absorb a -/// brief consumer pause without blocking the connection actor's -/// inbound dispatch loop. A slow stream consumer surfaces as -/// `RouteOutcome::ProgressFull` and the dropped frame is logged. -const PROGRESS_BUFFER: usize = 64; - -/// Opt-in per-call flag (a [`ToolCallContext`] extension): when set to -/// `true`, the remote call's [`ToolStream`] emits one best-effort -/// call-scoped cancel hook on `Drop` so the workspace hard-cancels the -/// in-flight call. Absent or `false` (the default) → `Drop` emits -/// nothing. No effect on local-dispatch calls. -#[derive(Clone, Copy, Debug)] -pub struct CancelOnDrop(pub bool); - -pub type ModelOutputExtractor = - Arc Option> + Send + Sync>; - -pub fn extractor_for() -> ModelOutputExtractor -where - T: kigi_tool_runtime::ToolOutput + serde::de::DeserializeOwned + 'static, -{ - Arc::new(|value: &Value| { - serde_json::from_value::(value.clone()) - .ok() - .map(|output| output.model_output().to_vec()) - }) -} - -/// In-process registry of tool handles owned by a [`ToolHarness`]. -/// -/// Tools registered here resolve in-process — `ToolHarness::call` -/// short-circuits the wire dispatch and invokes the handle directly. -/// Mutations are concurrency-safe (`RwLock` on `entries`, `DashMap` -/// on `extractors`), so callers MAY hot-add or hot-remove tools -/// while the harness is in use. -/// -/// `entries` uses `RwLock` to preserve insertion order so -/// that `list_tools` returns descriptions in the same order tools -/// were registered (matching the config-defined order). -/// -/// Optionally stores a per-tool [`ModelOutputExtractor`] for client-side -/// model output extraction. Use [`register_with_model_output`](Self::register_with_model_output) -/// to capture the extractor at registration time, or [`register_extractor`](Self::register_extractor) -/// to add one separately. -#[derive(Default)] -struct LocalRegistryInner { - entries: RwLock>>, - extractors: DashMap, -} - -#[derive(Clone, Default)] -pub struct LocalRegistry { - inner: 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.inner.entries.read().len()) - .field("extractors", &self.inner.extractors.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 + std::fmt::Debug + '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 + std::fmt::Debug + 'static, - { - let id = tool.id(); - let handle: Arc = Arc::new(ErasedTool::from_arc(tool)); - self.inner.entries.write().insert(id, handle) - } - - /// Resolve `tool_id` to its in-process handle, if registered. - /// Returns a clone of the `Arc` so the - /// caller can read without holding the lock across an await point. - pub fn find(&self, tool_id: &ToolId) -> Option> { - self.inner.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.inner.entries.write().shift_remove(tool_id).is_some() - } - - /// Number of tools currently registered. - pub fn len(&self) -> usize { - self.inner.entries.read().len() - } - - /// `true` iff no tools are registered. - pub fn is_empty(&self) -> bool { - self.inner.entries.read().is_empty() - } - - /// `true` iff `tool_id` is currently registered. - pub fn contains(&self, tool_id: &ToolId) -> bool { - self.inner.entries.read().contains_key(tool_id) - } - - /// Register `alias_id` as an alias pointing to the same handle as - /// `target_id`. Returns `true` if the alias was created (i.e. the - /// target exists), `false` otherwise. - /// - /// Used for MCP prefix-fallback: the model may emit the bare remote - /// name (`search_channels`) instead of the full prefixed name - /// (`slack___search_channels`). Registering the bare name as an - /// alias lets `find` resolve it without prefix-scanning logic. - pub fn register_alias(&self, alias_id: ToolId, target_id: &ToolId) -> bool { - if let Some(handle) = self.find(target_id) { - if let Some(extractor) = self.inner.extractors.get(target_id) { - self.inner - .extractors - .insert(alias_id.clone(), extractor.clone()); - } - self.inner.entries.write().insert(alias_id, handle); - true - } else { - false - } - } - - /// Register a type-erased [`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: Arc, - ) -> Option> { - let id = tool.id(); - // ToolDyn already implements ToolHandle via the blanket impl - // in kigi-computer-hub-core (ErasedTool). We wrap it in a thin - // adapter that delegates execute → ToolDyn::execute. - let handle: Arc = Arc::new(DynToolAdapter(tool)); - self.inner.entries.write().insert(id, handle) - } - - /// Register a tool and capture a type-safe model output extractor. - /// - /// Equivalent to calling [`register`](Self::register) followed by - /// [`register_extractor`](Self::register_extractor) with an extractor - /// built from `T::Output`. - pub fn register_with_model_output(&self, tool: T) -> Option> - where - T: Tool + std::fmt::Debug + 'static, - T::Output: kigi_tool_runtime::ToolOutput + serde::de::DeserializeOwned + 'static, - { - let id = tool.id(); - self.inner - .extractors - .insert(id, extractor_for::()); - self.register(tool) - } - - /// Attach a model output extractor for `tool_id`. Replaces any previous. - pub fn register_extractor(&self, tool_id: ToolId, extractor: ModelOutputExtractor) { - self.inner.extractors.insert(tool_id, extractor); - } - - /// Extract model-facing content blocks from a tool's output. - /// - /// Returns `None` if no extractor is registered for `tool_id` or if - /// the value fails to deserialize into the expected output type. - pub fn model_output( - &self, - tool_id: &ToolId, - output: &Value, - ) -> Option> { - self.inner - .extractors - .get(tool_id) - .and_then(|e| e.value()(output)) - } - - /// 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.inner - .entries - .read() - .values() - .filter(|handle| handle.should_list(ctx)) - .map(|handle| handle.description(ctx)) - .collect() - } -} - -/// Thin adapter from `Arc` to `ToolHandle`. -/// -/// `ToolDyn::execute` returns `ToolStream` which matches -/// `ToolHandle::execute`, so the adapter is a trivial delegation. -struct DynToolAdapter(Arc); - -impl std::fmt::Debug for DynToolAdapter { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("DynToolAdapter") - .field("id", &self.0.id()) - .finish() - } -} - -#[async_trait::async_trait] -impl ToolHandle for DynToolAdapter { - fn id(&self) -> ToolId { - self.0.id() - } - fn description(&self, ctx: &ListToolsContext) -> ToolDescription { - self.0.description(ctx) - } - fn capabilities(&self) -> kigi_tool_protocol::ToolCapabilities { - self.0.capabilities() - } - fn should_list(&self, ctx: &ListToolsContext) -> bool { - self.0.should_list(ctx) - } - async fn execute( - &self, - ctx: ToolCallContext, - args: Value, - ) -> ToolStream { - self.0.execute(ctx, args).await - } -} - -/// Builder for [`ToolHarness`]. See module docs for end-to-end usage. -#[derive(Default)] -pub struct ToolHarnessBuilder { - pool: Option>, - url: Option, - auth: Option>, - session: Option, - local_registry: LocalRegistry, - default_extensions: Option, - trace_context_provider: Option, - on_reconnect: Option>, - /// Sampler label for `hub_harness_connect_total` metric - /// (`"chat"` or `"shell"`). Defaults to `"unknown"`. - sampler: Option, - alpha_test_key: Option, - allow_insecure_ws: bool, - /// Resume the build-time `session.open`; default `false`. Does not affect - /// the transport auto-reconnect loop, which always uses `resume: false`. - resume: bool, - last_seq: Option, -} - -impl ToolHarnessBuilder { - /// Attach an extra access header on every (re)connect. - pub fn alpha_test_key(mut self, key: impl Into) -> Self { - self.alpha_test_key = Some(key.into()); - self - } - - /// Permit plaintext `ws://` to a non-loopback host. Only enable - /// when the transport is otherwise secured (e.g. a private network - /// or TLS-terminating proxy) — the bearer would otherwise cross the - /// wire in cleartext. - pub fn allow_insecure_ws(mut self, allow: bool) -> Self { - self.allow_insecure_ws = allow; - self - } - - /// Connection pool to attach to. Required. - pub fn pool(mut self, pool: Arc) -> Self { - self.pool = Some(pool); - self - } - - /// Server URL (`ws://` / `wss://`). Required. - pub fn url(mut self, url: Url) -> Self { - self.url = Some(url); - self - } - - pub fn auth(mut self, cred: AuthCredential) -> Self { - self.auth = Some(Arc::new(cred)); - self - } - - pub fn auth_provider(mut self, provider: Arc) -> Self { - self.auth = Some(provider); - self - } - - /// Bind `session_id` on the underlying connection and use it as - /// the envelope `session_id` for outgoing `tool.call` requests. - /// Calling repeatedly replaces the previous binding. - pub fn session(mut self, session_id: SessionId) -> Self { - self.session = Some(session_id); - self - } - - /// Register an in-process tool. Additive: subsequent calls add - /// more tools to the same [`LocalRegistry`]. - pub fn local_tool(self, tool: T) -> Self - where - T: Tool + std::fmt::Debug + 'static, - { - self.local_registry.register(tool); - self - } - - pub fn local_registry(mut self, registry: LocalRegistry) -> Self { - self.local_registry = registry; - self - } - - /// Default extensions merged into every `ToolCallContext` before dispatch. - pub fn default_extensions(mut self, extensions: kigi_tool_runtime::TypedExtensions) -> Self { - self.default_extensions = Some(extensions); - self - } - - /// Sampled per outgoing tool call / hook on the caller's task. - /// Keeps the host's tracing stack out of the SDK. - pub fn trace_context_provider(mut self, provider: F) -> Self - where - F: Fn() -> Option + Send + Sync + 'static, - { - self.trace_context_provider = Some(Arc::new(provider)); - self - } - - /// Optional callback fired once per successful reconnect cycle. - pub fn on_reconnect(mut self, cb: F) -> Self - where - F: Fn(ReconnectEvent) + Send + Sync + 'static, - { - self.on_reconnect = Some(Arc::new(Box::new(cb) as ReconnectCallback)); - self - } - - /// Sampler label for the `hub_harness_connect_total` metric. - /// Pass `"chat"` or `"shell"` to identify the calling sampler. - pub fn sampler(mut self, sampler: impl Into) -> Self { - self.sampler = Some(sampler.into()); - self - } - - /// Resume the build-time `session.open` (e.g. a cloud reconnect re-attaching - /// to its existing server session). Default `false`; auto-reconnect is unaffected. - pub fn resume(mut self, resume: bool) -> Self { - self.resume = resume; - self - } - - /// Last-seen `(connection_id, seq)` paired with [`Self::resume`] for replay dedup. - pub fn last_seq(mut self, last_seq: kigi_tool_protocol::LastSeq) -> Self { - self.last_seq = Some(last_seq); - self - } - - /// Resolve the pool entry, refcount-bind the session, and return - /// a [`ToolHarness`] ready to dispatch tool calls. On a session- - /// bind failure the previously-bound sessions are rolled back - /// before propagating the original error; the pooled connection - /// stays in the pool for future borrowers (the failed build's - /// local `Arc` is dropped, but the pool keeps its - /// own clone). - pub async fn build(self) -> Result { - let pool = self - .pool - .ok_or_else(|| ClientError::InvalidConfig("missing pool".to_owned()))?; - let url = self - .url - .ok_or_else(|| ClientError::InvalidConfig("missing url".to_owned()))?; - let auth = self - .auth - .ok_or_else(|| ClientError::InvalidConfig("missing auth".to_owned()))?; - let session = self - .session - .ok_or_else(|| ClientError::InvalidConfig("missing session".to_owned()))?; - let sampler = self.sampler.as_deref().unwrap_or("unknown"); - let borrow = match ConnectionBorrow::acquire( - pool, - url, - auth, - ConnectionKind::Harness, - self.on_reconnect.clone(), - None, // on_disconnect (unused for harness connections) - None, // on_connect (unused for harness connections) - None, - None, - None, - self.alpha_test_key.clone(), - self.allow_insecure_ws, - crate::connection::ConnectionTuning::default(), - ) - .await - { - Ok(b) => { - crate::metrics::harness_connect("ok", sampler); - b - } - Err(err) => { - crate::metrics::harness_connect("error", sampler); - return Err(err); - } - }; - // Track the harness session locally for reconnect replay. - borrow.connection().track_session(session.clone()); - - // Register the session on the server so session-scoped RPCs - // (session_bind_server, etc.) are accepted. The reconnect path - // already replays session_open per tracked session (connection.rs), - // but the initial connect must do it explicitly. - { - let connection = borrow.connection(); - let request_id = connection.try_alloc_request_id()?; - let req = kigi_tool_protocol::JsonRpcRequest { - jsonrpc: kigi_tool_protocol::JsonRpcVersion, - id: kigi_tool_protocol::JsonRpcId::from_request_id(&request_id), - session_id: Some(session.clone()), - method: kigi_tool_protocol::Method::SessionOpen - .as_wire_str() - .to_owned(), - params: kigi_tool_protocol::SessionOpenParams { - resume: self.resume, - last_seq: self.last_seq, - }, - }; - connection - .call_request(request_id, &req) - .await - .map_err(|e| { - tracing::warn!(error = %e, "session_open failed during harness build"); - e - })?; - } - - let inner = Arc::new(ToolHarnessInner { - borrow: Some(borrow), - local_registry: self.local_registry, - session, - default_extensions: self.default_extensions.unwrap_or_default(), - trace_context_provider: self.trace_context_provider, - remote_tools: arc_swap::ArcSwap::from_pointee(Vec::new()), - last_bind_report: arc_swap::ArcSwapOption::empty(), - discovery_handle: parking_lot::Mutex::new(None), - pending_bind: None, - hook_request_handler: Arc::new(parking_lot::Mutex::new(None)), - }); - Ok(ToolHarness { inner }) - } -} - -/// Typed bind-contract report from a `session.bind` response. -#[derive(Debug, Clone, Default)] -pub struct SessionBindReport { - /// Version of the tool-server binary that served the bind. - pub binary_version: Option, - /// Configured tool ids the server could not serve. - pub unserved_tool_ids: Vec, - /// Server-stated reason the toolset resolution failed closed (the bind - /// advertises no model-facing tools by design when set). - pub resolve_error: Option, -} - -/// Harness attached to a pooled [`HubConnection`]. -/// -/// `ToolHarness` is `Clone`-cheap (`Arc` bump). Cooperative teardown -/// via [`Self::shutdown`] is preferred; the `Drop` impl schedules a -/// best-effort asynchronous cleanup as a fallback when no explicit -/// shutdown ran. Cleanup fires at most once across all clones — the -/// first drop (or shutdown) to flip the underlying `torn_down` flag -/// wins; subsequent drops no-op. -pub struct ToolHarness { - inner: Arc, -} - -/// An owned, type-erased server-bind future, resolving to the server-connected -/// [`ToolHarness`] (or a stringified bind error). -type BindFuture = BoxFuture<'static, Result>>; - -/// Cloneable handle to the deferred server bind; every clone observes the same -/// single bind, resolving to the server-connected [`ToolHarness`]. -type PendingBind = Shared; - -/// Spawn `bind` on the runtime as a cloneable [`PendingBind`], projecting a -/// task-join panic into the bind's `Err`. Shared by the eager constructor and -/// `LazyBind::start` so the two spawn paths can't drift. -fn spawn_pending_bind(bind: F) -> PendingBind -where - F: std::future::Future>> + Send + 'static, -{ - let task = tokio::spawn(bind); - async move { - match task.await { - Ok(result) => result, - Err(join_err) => Err(Arc::::from( - format!("server bind task panicked: {join_err}").as_str(), - )), - } - } - .boxed() - .shared() -} - -/// A bind future kept unspawned until the first [`ToolHarness::await_bound`], so -/// the server connection — and the sandbox provisioning it performs — is deferred -/// to the first remote tool dispatch. See [`ToolHarness::local_with_lazy_bind`]. -struct LazyBind { - fut: parking_lot::Mutex>, - started: std::sync::OnceLock, -} - -impl LazyBind { - /// Spawn the bind exactly once and return the cloneable shared handle. - fn start(&self) -> PendingBind { - self.started - .get_or_init(|| { - let fut = self - .fut - .lock() - .take() - .expect("LazyBind future taken more than once"); - spawn_pending_bind(fut) - }) - .clone() - } -} - -/// Deferred server bind: `Eager` is spawned at construction (races sampling); -/// `Lazy` spawns on the first `await_bound` (provisioning deferred to first call). -enum DeferredBind { - Eager(PendingBind), - Lazy(LazyBind), -} - -struct ToolHarnessInner { - /// `None` for local-only harnesses (no server connection). - borrow: Option, - local_registry: LocalRegistry, - session: SessionId, - /// Default extensions merged into every `ToolCallContext` before dispatch. - default_extensions: kigi_tool_runtime::TypedExtensions, - /// See [`ToolHarnessBuilder::trace_context_provider`]. - trace_context_provider: Option, - remote_tools: arc_swap::ArcSwap>, - last_bind_report: arc_swap::ArcSwapOption, - discovery_handle: parking_lot::Mutex>>, - /// Deferred server bind (prompt-before-bind): set when this local-only harness - /// resolves to a server-connected one once the bind completes. Eager variant - /// races sampling; lazy variant defers provisioning to the first remote - /// tool dispatch. - pending_bind: Option, - /// Optional sink for inbound reverse-direction hook requests. Held in - /// its own `Arc` so the inbox loop can clone this slot — not the whole - /// `inner` — keeping the `Drop` strong-count teardown gate intact. - hook_request_handler: Arc>>, -} - -impl ToolHarnessInner { - /// When the SDK sees a workspace `Disconnected` notification on its own - /// socket, fail this session's parked `call` futures with the recognizable - /// `workspace_unavailable` error (`phase: InFlightCancelled`). Client-side - /// guarantee that closes the window where the server-side cancel is missed - /// (no subscription / lost broadcast). - fn fail_inflight_calls_on_disconnect(&self, session_id: &SessionId) { - let Some(borrow) = self.borrow.as_ref() else { - return; // local-only harness has no wire calls to fail - }; - let resolved = borrow - .connection() - .demux() - .fail_calls_for_session(session_id, || { - ClientError::Wire(workspace_unavailable_wire( - WorkspaceGoneReason::Disconnect, - WorkspaceGonePhase::InFlightCancelled, - )) - }); - if resolved > 0 { - tracing::info!( - %session_id, - resolved, - "SDK short-circuit: failed in-flight calls on workspace disconnect" - ); - } - } - - async fn refresh_remote_tools(&self) -> Result, ClientError> { - let borrow = self.borrow.as_ref().ok_or_else(|| { - ClientError::InvalidConfig("local-only harness has no server connection".to_owned()) - })?; - let connection = borrow.connection(); - let request_id = connection.try_alloc_request_id()?; - let params = kigi_tool_protocol::ToolsListParams { - session_id: self.session.clone(), - mode: kigi_tool_protocol::ToolDefinitionMode::Full, - }; - let req = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: Some(self.session.clone()), - method: Method::ToolsList.as_wire_str().to_owned(), - params, - }; - let resp = connection.call_request(request_id, &req).await?; - match resp.outcome { - ResponseOutcome::Result(value) => { - let result: kigi_tool_protocol::ToolsListResult = - serde_json::from_value(value).map_err(|e| ClientError::Serde(e.to_string()))?; - self.remote_tools.store(Arc::new(result.tools.clone())); - Ok(result.tools) - } - ResponseOutcome::Error(err) => Err(ClientError::from_jsonrpc_error(err)), - } - } -} - -impl Clone for ToolHarness { - fn clone(&self) -> Self { - Self { - inner: self.inner.clone(), - } - } -} - -impl std::fmt::Debug for ToolHarness { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ToolHarness") - .field("session", &self.inner.session) - .field("local_tool_count", &self.inner.local_registry.len()) - .finish_non_exhaustive() - } -} - -impl ToolHarness { - /// Construct a local-only harness (no server connection). - /// - /// Tools are resolved exclusively from the `LocalRegistry`. Remote - /// dispatch returns `ToolError::NotFound`. `default_extensions` are - /// merged into every `ToolCallContext` before dispatch — use this - /// to inject `GrokSharedState`, `GrokAgentState`, etc. - pub fn local_only_with( - registry: LocalRegistry, - session: SessionId, - default_extensions: kigi_tool_runtime::TypedExtensions, - ) -> Self { - let inner = Arc::new(ToolHarnessInner { - borrow: None, - local_registry: registry, - session, - default_extensions, - trace_context_provider: None, - remote_tools: arc_swap::ArcSwap::from_pointee(Vec::new()), - last_bind_report: arc_swap::ArcSwapOption::empty(), - discovery_handle: parking_lot::Mutex::new(None), - pending_bind: None, - hook_request_handler: Arc::new(parking_lot::Mutex::new(None)), - }); - Self { inner } - } - - /// Construct a local-only harness whose server bind is deferred to a background - /// task, spawned eagerly so the connection races with sampling instead of - /// blocking it. Local tools dispatch immediately; remote work awaits - /// [`Self::await_bound`] (or probes [`Self::try_bound`]). - pub fn local_with_pending_bind( - registry: LocalRegistry, - session: SessionId, - default_extensions: kigi_tool_runtime::TypedExtensions, - bind: F, - ) -> Self - where - F: std::future::Future>> + Send + 'static, - { - let pending = spawn_pending_bind(bind); - let inner = Arc::new(ToolHarnessInner { - borrow: None, - local_registry: registry, - session, - default_extensions, - trace_context_provider: None, - remote_tools: arc_swap::ArcSwap::from_pointee(Vec::new()), - last_bind_report: arc_swap::ArcSwapOption::empty(), - discovery_handle: parking_lot::Mutex::new(None), - pending_bind: Some(DeferredBind::Eager(pending)), - hook_request_handler: Arc::new(parking_lot::Mutex::new(None)), - }); - Self { inner } - } - - /// Construct a local-only harness whose server bind is deferred *and not - /// started* until the first [`Self::await_bound`] call. Unlike - /// [`Self::local_with_pending_bind`], the bind future is stored unspawned, - /// so the server connection — and the sandbox provisioning it triggers — only - /// begins on the first remote tool dispatch. Local tools dispatch - /// immediately and never start the bind; [`Self::try_bound`] returns `None` - /// until the bind has been started (it does not itself kick it off). - pub fn local_with_lazy_bind( - registry: LocalRegistry, - session: SessionId, - default_extensions: kigi_tool_runtime::TypedExtensions, - bind: F, - ) -> Self - where - F: std::future::Future>> + Send + 'static, - { - let lazy = LazyBind { - fut: parking_lot::Mutex::new(Some(bind.boxed())), - started: std::sync::OnceLock::new(), - }; - let inner = Arc::new(ToolHarnessInner { - borrow: None, - local_registry: registry, - session, - default_extensions, - trace_context_provider: None, - remote_tools: arc_swap::ArcSwap::from_pointee(Vec::new()), - last_bind_report: arc_swap::ArcSwapOption::empty(), - discovery_handle: parking_lot::Mutex::new(None), - pending_bind: Some(DeferredBind::Lazy(lazy)), - hook_request_handler: Arc::new(parking_lot::Mutex::new(None)), - }); - Self { inner } - } - - pub fn has_pending_bind(&self) -> bool { - self.inner.pending_bind.is_some() - } - - /// Await the deferred server bind and return the server-connected harness, or a - /// clone of `self` when there is no pending bind (so callers dispatch - /// through the result uniformly). - /// - /// For a lazy bind (see [`Self::local_with_lazy_bind`]), this is what - /// actually starts the bind — the first call spawns it. - pub async fn await_bound(&self) -> Result> { - match &self.inner.pending_bind { - Some(DeferredBind::Eager(pending)) => pending.clone().await, - Some(DeferredBind::Lazy(lazy)) => lazy.start().await, - None => Ok(self.clone()), - } - } - - /// Non-blocking probe of the deferred bind: `None` while in flight (or no - /// pending bind), `Some(Ok)`/`Some(Err)` once resolved. Uses `now_or_never` - /// (not `peek`): the bind runs in a spawned task, so `Shared` only observes - /// completion once polled. - /// - /// For a lazy bind, this never *starts* the bind — it returns `None` until - /// a prior [`Self::await_bound`] call has spawned it, then probes that. - pub fn try_bound(&self) -> Option>> { - match self.inner.pending_bind.as_ref()? { - DeferredBind::Eager(pending) => pending.clone().now_or_never(), - DeferredBind::Lazy(lazy) => lazy.started.get()?.clone().now_or_never(), - } - } - - /// Underlying connection. Useful for tests that need to assert - /// pool dedup. - pub fn connection(&self) -> Result<&Arc, ClientError> { - self.require_connection() - } - - /// Bound session. - pub fn session(&self) -> &SessionId { - &self.inner.session - } - - pub fn local_registry(&self) -> LocalRegistry { - self.inner.local_registry.clone() - } - - /// Extract model-facing content blocks from a tool's output. - /// - /// Returns `None` if no extractor is registered for `tool_id` - /// or if deserialization fails. - pub fn model_output( - &self, - tool_id: &ToolId, - output: &Value, - ) -> Option> { - self.inner.local_registry.model_output(tool_id, output) - } - - fn require_connection(&self) -> Result<&Arc, ClientError> { - self.inner - .borrow - .as_ref() - .map(|b| b.connection()) - .ok_or_else(|| { - ClientError::InvalidConfig( - "operation requires a server connection (local-only harness)".to_owned(), - ) - }) - } - - /// Discover available tool servers for the current user. - pub async fn list_servers(&self) -> Result, ClientError> { - let connection = self.require_connection()?; - let request_id = connection.try_alloc_request_id()?; - let req = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: Some(self.inner.session.clone()), - method: kigi_tool_protocol::Method::ServersList - .as_wire_str() - .to_owned(), - params: kigi_tool_protocol::ServersListParams {}, - }; - let resp = connection.call_request(request_id, &req).await?; - match resp.outcome { - ResponseOutcome::Result(value) => { - let result: kigi_tool_protocol::ServersListResult = - serde_json::from_value(value).map_err(|e| ClientError::Serde(e.to_string()))?; - Ok(result.servers) - } - ResponseOutcome::Error(err) => Err(ClientError::from_jsonrpc_error(err)), - } - } - - /// Open a session on the server. Does not bind any server. - /// - /// Registers the session on the server connection and claims - /// ownership. Server binding is a separate step via - /// [`Self::session_bind`]. - pub async fn session_open(&self) -> Result<(), ClientError> { - self.session_open_with(false, None).await - } - - /// [`Self::session_open`] with an explicit `resume` flag and optional `last_seq`. - pub async fn session_open_with( - &self, - resume: bool, - last_seq: Option, - ) -> Result<(), ClientError> { - let start = std::time::Instant::now(); - let result: Result<(), ClientError> = async { - let connection = self.require_connection()?; - let request_id = connection.try_alloc_request_id()?; - let params = kigi_tool_protocol::SessionOpenParams { resume, last_seq }; - let req = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: Some(self.inner.session.clone()), - method: Method::SessionOpen.as_wire_str().to_owned(), - params, - }; - let resp = connection.call_request(request_id, &req).await?; - match resp.outcome { - ResponseOutcome::Result(_) => Ok(()), - ResponseOutcome::Error(err) => Err(ClientError::from_jsonrpc_error(err)), - } - } - .await; - crate::metrics::session_op_observe( - "open", - if result.is_ok() { "ok" } else { "error" }, - start.elapsed().as_secs_f64(), - ); - result - } - - /// Bind a server's tools to the current session. - /// - /// Returns the tools available after the server is bound. Updates - /// the in-memory remote tools snapshot. Optional `cwd` and - /// `metadata` are forwarded to the tool server's session creation. - pub async fn session_bind( - &self, - server_id: &str, - cwd: Option<&str>, - metadata: Option, - ) -> Result, ClientError> { - self.session_bind_with_report(server_id, cwd, metadata) - .await - .map(|r| r.tools) - } - - /// [`Self::session_bind`], returning the full bind result including the - /// server's bind report (`binary_version`, `unserved_tool_ids`). - pub async fn session_bind_with_report( - &self, - server_id: &str, - cwd: Option<&str>, - metadata: Option, - ) -> Result { - let start = std::time::Instant::now(); - let result: Result = async { - let connection = self.require_connection()?; - let request_id = connection.try_alloc_request_id()?; - let parsed_server_id = kigi_tool_protocol::ServerId::new(server_id) - .map_err(|e| ClientError::InvalidConfig(format!("invalid server_id: {e}")))?; - let params = kigi_tool_protocol::SessionBindServerParams { - server_id: parsed_server_id, - cwd: cwd.map(String::from), - metadata, - }; - let req = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: Some(self.inner.session.clone()), - method: Method::SessionBindServer.as_wire_str().to_owned(), - params, - }; - let resp = connection.call_request(request_id, &req).await?; - match resp.outcome { - ResponseOutcome::Result(value) => { - let bind_result: kigi_tool_protocol::SessionBindServerResult = - serde_json::from_value(value) - .map_err(|e| ClientError::Serde(e.to_string()))?; - let arc = Arc::new(bind_result.tools.clone()); - self.inner.remote_tools.store(arc); - self.inner - .last_bind_report - .store(Some(Arc::new(SessionBindReport { - binary_version: bind_result.binary_version.clone(), - unserved_tool_ids: bind_result.unserved_tool_ids.clone(), - resolve_error: bind_result.resolve_error.clone(), - }))); - Ok(bind_result) - } - ResponseOutcome::Error(err) => Err(ClientError::from_jsonrpc_error(err)), - } - } - .await; - crate::metrics::session_op_observe( - "bind", - if result.is_ok() { "ok" } else { "error" }, - start.elapsed().as_secs_f64(), - ); - result - } - - /// Attach to the current session as an observer - /// (`session_attach_server`): a server-local check that a tool-server is - /// routed for the envelope session, replying the tool snapshot and the - /// route it was found on. Never binds, never creates a workspace - /// session, never touches toolsets or handlers. Updates the in-memory - /// remote tools snapshot like [`Self::session_bind`]. - /// - /// `server_id` is an optional diagnostics cross-check (the envelope - /// session is the authoritative key); `caller` is a free-form label - /// surfaced in server metrics/logs. An attach miss is the retryable - /// `workspace_unavailable` family (`reason: not_bound`). Server support - /// is answered before calling, not from the reply: the server advertises - /// `session_attach_server` in `hello_ack` (capability and method ship - /// in the same server commit) — gate on `HubConnection::supports`; a server - /// that predates the method advertises nothing and replies `-32601`. - pub async fn session_attach( - &self, - server_id: Option<&str>, - caller: &str, - ) -> Result { - let start = std::time::Instant::now(); - let result: Result = async { - let connection = self.require_connection()?; - let request_id = connection.try_alloc_request_id()?; - let parsed_server_id = server_id - .map(|s| { - kigi_tool_protocol::ServerId::new(s) - .map_err(|e| ClientError::InvalidConfig(format!("invalid server_id: {e}"))) - }) - .transpose()?; - let params = kigi_tool_protocol::SessionAttachServerParams { - server_id: parsed_server_id, - caller: Some(caller.to_owned()), - }; - let req = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: Some(self.inner.session.clone()), - method: Method::SessionAttachServer.as_wire_str().to_owned(), - params, - }; - let resp = connection.call_request(request_id, &req).await?; - match resp.outcome { - ResponseOutcome::Result(value) => { - let attach_result: kigi_tool_protocol::SessionAttachServerResult = - serde_json::from_value(value) - .map_err(|e| ClientError::Serde(e.to_string()))?; - self.inner - .remote_tools - .store(Arc::new(attach_result.tools.clone())); - Ok(attach_result) - } - ResponseOutcome::Error(err) => Err(ClientError::from_jsonrpc_error(err)), - } - } - .await; - crate::metrics::session_op_observe( - "attach", - if result.is_ok() { "ok" } else { "error" }, - start.elapsed().as_secs_f64(), - ); - result - } - - /// Unbind a server from the current session. - pub async fn session_unbind(&self, server_id: &str) -> Result<(), ClientError> { - let connection = self.require_connection()?; - let request_id = connection.try_alloc_request_id()?; - let parsed_server_id = kigi_tool_protocol::ServerId::new(server_id) - .map_err(|e| ClientError::InvalidConfig(format!("invalid server_id: {e}")))?; - let params = kigi_tool_protocol::SessionUnbindServerParams { - server_id: parsed_server_id, - }; - let req = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: Some(self.inner.session.clone()), - method: Method::SessionUnbindServer.as_wire_str().to_owned(), - params, - }; - let resp = connection.call_request(request_id, &req).await?; - match resp.outcome { - ResponseOutcome::Result(_) => { - // Clear cached remote tools — the server's tools are no - // longer available after unbind. - self.inner.remote_tools.store(Arc::new(Vec::new())); - Ok(()) - } - ResponseOutcome::Error(err) => Err(ClientError::from_jsonrpc_error(err)), - } - } - - /// Close a session, unbinding all servers. - /// - /// On the server side, this drops all tool bindings for the session - /// and sends `session.unbind` to any bound tool servers. - pub async fn session_close(&self) -> Result<(), ClientError> { - let connection = self.require_connection()?; - let request_id = connection.try_alloc_request_id()?; - let params = kigi_tool_protocol::SessionCloseParams { reason: None }; - let req = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: Some(self.inner.session.clone()), - method: Method::SessionClose.as_wire_str().to_owned(), - params, - }; - let resp = connection.call_request(request_id, &req).await?; - match resp.outcome { - ResponseOutcome::Result(_) => Ok(()), - ResponseOutcome::Error(err) => Err(ClientError::from_jsonrpc_error(err)), - } - } - - /// All tool descriptions: local registry plus cached remote tools. - pub fn list_tools(&self, ctx: &ListToolsContext) -> Vec { - let mut tools = self.list_local_tools(ctx); - tools.extend(self.list_remote_tools()); - tools - } - - /// Whether `name` is a known remote tool, per the cached - /// `session_bind` / `tools.list` result. Empty before the first - /// bind/discovery and after unbind; always `false` for local-only - /// harnesses. Remote tools are not mirrored into the local registry. - pub fn has_remote_tool(&self, name: &str) -> bool { - self.inner - .remote_tools - .load() - .iter() - .any(|t| t.name == name) - } - - /// Test-only: seed the remote-tools cache without a server round-trip. - #[doc(hidden)] - pub fn seed_remote_tools_for_tests(&self, tools: Vec) { - self.inner.remote_tools.store(Arc::new(tools)); - } - - /// Tool descriptions from the local registry only. - pub fn list_local_tools(&self, ctx: &ListToolsContext) -> Vec { - self.inner.local_registry.list_tools(ctx) - } - - /// Cached tool descriptions advertised by remote tool servers. - pub fn list_remote_tools(&self) -> Vec { - self.inner.remote_tools.load().iter().cloned().collect() - } - - /// Shared snapshot of the cached remote tool descriptions, cloning only the - /// backing `Arc` rather than the `Vec`. - pub fn remote_tools_snapshot(&self) -> Arc> { - self.inner.remote_tools.load_full() - } - - /// The bind-contract report from the most recent successful `session.bind`, if any. - pub fn last_bind_report(&self) -> Option> { - self.inner.last_bind_report.load_full() - } - - /// Test-only: seed the bind report. - pub fn seed_bind_report_for_tests(&self, report: SessionBindReport) { - self.inner.last_bind_report.store(Some(Arc::new(report))); - } - - /// Dispatch a tool call. - /// - /// Local-first: if `tool_id` is registered in the - /// [`LocalRegistry`] the tool's `execute` runs in-process and the - /// returned [`ToolStream`] is forwarded verbatim — no wire - /// round-trip. The hot path on a hit is a single `DashMap` lookup - /// plus an `Arc` clone of the handle. - /// - /// Otherwise the harness sends a `tool.call` JSON-RPC request - /// over the shared connection. The returned stream interleaves - /// any intermediate `tool_call_progress` notifications (matched - /// by `tool_call_id`) with the eventual JSON-RPC response, which - /// becomes the terminal item. - /// - /// Each call MUST use a fresh `ToolCallId`. The default - /// [`ToolCallContext::default`] mints a UUIDv7 via - /// [`ToolCallId::new_v7`]; callers that build a context manually - /// MUST do the same. Reusing a `ToolCallId` already in flight on - /// this connection is detected synchronously: the second call's - /// stream resolves with a single `Terminal(Err(_))` carrying - /// `ToolError::Custom { code: "call_id_in_use", .. }` and the - /// FIRST call's progress and response correlation are left intact - /// (no stomp). The `tool_call_id` keys the per-call progress - /// channel; concurrent ids would otherwise observe each other's - /// progress frames. - pub async fn call( - &self, - tool_id: ToolId, - args: Value, - mut ctx: ToolCallContext, - ) -> ToolStream { - // Merge harness-level default extensions (SharedState, AgentState, etc.) - // into the per-call context. Per-call values take priority. - ctx.extensions - .merge_defaults(&self.inner.default_extensions); - - // Sampled on the caller's task while the dispatching span is - // active; never rides ctx extensions. - let trace_context = self - .inner - .trace_context_provider - .as_ref() - .and_then(|provider| provider()); - - // Capture identifiers before `ctx` moves into the dispatch path — - // they feed both the local-only / remote branches AND the - // `ObservedToolStream` wrapper. - let observed_call_id = ctx.call_id.clone(); - - let raw_stream: ToolStream = - if let Some(handle) = self.inner.local_registry.find(&tool_id) { - handle.execute(ctx, args).await - } else if let Some(ref borrow) = self.inner.borrow { - let start = std::time::Instant::now(); - let stream = dispatch_remote( - borrow.connection(), - &self.inner.session, - tool_id.clone(), - args, - ctx, - trace_context, - ) - .await; - crate::metrics::call_dispatch_observe(start.elapsed().as_secs_f64()); - stream - } else { - let message = format!("tool not found (local-only harness): {tool_id}"); - return terminal_only(Err(ToolError::not_found(tool_id, message))); - }; - - // Server observability: only wrap when the harness has a server - // connection — otherwise emission is a no-op and the extra - // Started/Completed bookkeeping is pure waste. Local-only - // harnesses (and the no-`borrow` not-found branch above) skip - // wrapping entirely. - if self.inner.borrow.is_none() { - return raw_stream; - } - self.emit_session_event(SessionEvent::ToolCallStarted { - tool_call_id: observed_call_id.as_str().to_owned(), - tool_name: tool_id.as_str().to_owned(), - turn_number: 0, - }) - .await; - Box::pin(ObservedToolStream::new( - raw_stream, - self.clone(), - observed_call_id, - tool_id, - )) - } - - /// Emit a session-level event to the server as a `session_event` - /// custom notification. - /// - /// No-op on a local-only harness — without a server connection there is - /// nothing to deliver to and the wire frame would be discarded. Server - /// errors are silently ignored: emission is fire-and-forget and must - /// never affect dispatch. - /// - /// `ToolHarness::call` invokes this automatically for - /// `ToolCallStarted` / `ToolCallCompleted`; higher-level events - /// (turn lifecycle, phase changes) are emitted by callers that - /// already hold the harness. - pub async fn emit_session_event(&self, event: SessionEvent) { - if self.inner.borrow.is_none() { - return; - } - let _ = self - .send_notification(build_session_event_frame(&event)) - .await; - } - - /// Send a `tool.notify` frame to the server. - /// - /// The server fans the notification out to every connection that has an - /// active `subscribe_notifications` subscription for this session. - /// The frame is fire-and-forget: this method returns `Ok` once the - /// outbound message is queued, without waiting for a server ack. - /// Server-side errors (e.g. unknown tool, invalid session) are not - /// surfaced to the caller. - pub async fn send_notification( - &self, - notification: kigi_tool_protocol::ToolNotificationFrame, - ) -> Result<(), ClientError> { - self.send_fire_and_forget(Method::ToolNotify, notification) - .await - } - - /// Send a `hook` frame to the server. - /// - /// The server routes the hook to the tool server that owns the targeted - /// tool (or broadcasts to all servers bound to the session for - /// session-wide hooks like `Pause` / `Resume` / `SessionEnded`). - /// The frame is fire-and-forget: this method returns `Ok` once the - /// outbound message is queued, without waiting for a server ack. - /// Server-side errors (e.g. unknown tool, invalid session) are not - /// surfaced to the caller. - pub async fn send_hook( - &self, - mut hook: kigi_tool_protocol::HookFrame, - ) -> Result<(), ClientError> { - if hook.trace_context.is_none() - && let Some(provider) = &self.inner.trace_context_provider - { - hook.trace_context = provider(); - } - let hook_type = match &hook.event { - kigi_tool_protocol::HookEvent::Cancel => "cancel", - kigi_tool_protocol::HookEvent::Pause => "pause", - kigi_tool_protocol::HookEvent::Resume => "resume", - kigi_tool_protocol::HookEvent::SessionEnded => "session_ended", - kigi_tool_protocol::HookEvent::Custom { .. } => "custom", - }; - crate::metrics::hook_send(hook_type); - self.send_fire_and_forget(Method::Hook, hook).await - } - - /// Cancel an in-flight remote call. - /// - /// Sends the call-scoped `Cancel` [`HookFrame`](kigi_tool_protocol::HookFrame) - /// over the fire-and-forget [`Self::send_hook`] path. Idempotent: the - /// server routes it to the owning tool server, which hard-cancels a live - /// call or tombstones an unknown / already-completed `call_id`. `Ok` - /// means the frame was queued; a late or repeated id is never an error - /// here. - pub async fn cancel_call( - &self, - tool_id: &ToolId, - call_id: &ToolCallId, - ) -> Result<(), ClientError> { - let hook = kigi_tool_protocol::HookFrame::cancel( - self.inner.session.clone(), - tool_id.clone(), - call_id.clone(), - ); - self.send_hook(hook).await - } - - /// Send a `before_turn` hook to all tool servers bound to this session. - /// - /// Fire-and-forget: returns `Ok` once the frame is queued on the - /// outbound channel. The workspace (or any other tool server) receives - /// the hook via `ToolServerHandler::handle_hook`. - pub async fn send_before_turn_hook( - &self, - payload: kigi_tool_protocol::turn_hook::BeforeTurnPayload, - ) -> Result<(), ClientError> { - let hook = kigi_tool_protocol::HookFrame::custom( - self.session().clone(), - kigi_tool_protocol::turn_hook::BEFORE_TURN_KIND.to_owned(), - serde_json::to_value(&payload).map_err(|e| ClientError::Serde(e.to_string()))?, - ); - self.send_hook(hook).await - } - - /// Send an `after_turn` hook to all tool servers bound to this session. - /// - /// Fire-and-forget: returns `Ok` once the frame is queued on the - /// outbound channel. The workspace (or any other tool server) receives - /// the hook via `ToolServerHandler::handle_hook`. - pub async fn send_after_turn_hook( - &self, - payload: kigi_tool_protocol::turn_hook::AfterTurnPayload, - ) -> Result<(), ClientError> { - let hook = kigi_tool_protocol::HookFrame::custom( - self.session().clone(), - kigi_tool_protocol::turn_hook::AFTER_TURN_KIND.to_owned(), - serde_json::to_value(&payload).map_err(|e| ClientError::Serde(e.to_string()))?, - ); - self.send_hook(hook).await - } - - /// Max time to await a turn-hook reply before giving up. - pub const TURN_HOOK_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); - - /// Request turn-boundary injections + a loop-control decision from the bound workspace server, - /// bounded by [`TURN_HOOK_TIMEOUT`](Self::TURN_HOOK_TIMEOUT). Sent as a request/response hook - /// (not a tool); any error is treated as a no-op by the caller. - pub async fn request_turn_hook( - &self, - request: &kigi_tool_protocol::turn_hook::TurnHookRequest, - ) -> Result { - self.request_turn_hook_with_timeout(request, Self::TURN_HOOK_TIMEOUT) - .await - } - - /// [`Self::request_turn_hook`] with a caller-supplied reply deadline - /// (the responder's own watchdog must stay below it). - pub async fn request_turn_hook_with_timeout( - &self, - request: &kigi_tool_protocol::turn_hook::TurnHookRequest, - timeout: std::time::Duration, - ) -> Result { - let connection = self.require_connection()?; - let payload = serde_json::to_value(request).map_err(ClientError::from)?; - // hook_id keys the server's parked-request table — must be globally unique. - let hook_id = kigi_tool_protocol::ToolCallId::new_v7().to_string(); - let hook = kigi_tool_protocol::HookFrame::custom_request( - self.inner.session.clone(), - hook_id, - kigi_tool_protocol::turn_hook::TURN_HOOK_KIND.to_owned(), - payload, - ) - .with_trace_context( - self.inner - .trace_context_provider - .as_ref() - .and_then(|provider| provider()), - ); - let request_id = connection.try_alloc_request_id()?; - let req = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: Some(self.inner.session.clone()), - method: Method::Hook.as_wire_str().to_owned(), - params: hook, - }; - - // Tag transport errors with turn-hook context; other variants pass through. - let resp = connection - .call_request_with_timeout(request_id, &req, timeout) - .await - .map_err(|e| match e { - ClientError::NetworkError(msg) => { - ClientError::NetworkError(format!("turn hook: {msg}")) - } - other => other, - })?; - - match resp.outcome { - ResponseOutcome::Result(value) => { - serde_json::from_value(value).map_err(|e| ClientError::Serde(e.to_string())) - } - ResponseOutcome::Error(err) => Err(ClientError::from_jsonrpc_error(err)), - } - } - - /// Answer a reverse-direction request/response [`HookFrame`](kigi_tool_protocol::HookFrame) - /// (one whose `hook_id` is set). - /// - /// Sent as a fire-and-forget `hook_reply` notification (not a JSON-RPC - /// response): the server correlates it to the parked request by `hook_id`. - /// `Ok` once the frame is queued on the outbound channel. - pub async fn send_hook_reply( - &self, - reply: kigi_tool_protocol::HookReplyFrame, - ) -> Result<(), ClientError> { - let connection = self.require_connection()?; - let notif = build_hook_reply_notification(&self.inner.session, reply); - let text = serde_json::to_string(¬if).map_err(ClientError::from)?; - connection.send_outbound(text).await - } - - /// Best-effort, non-blocking twin of [`Self::send_hook_reply`] for - /// synchronous `Drop`/teardown paths that cannot `.await`. - /// - /// Enqueues via [`HubConnection::try_send_outbound`]; a full or closed - /// outbound channel returns `Err` and the frame is abandoned — the server's - /// parked-request backstop then releases the await. Mirrors - /// `RemoteCallStream`'s cancel-on-drop discipline. The async variant's only - /// edge over this is a brief bounded wait when the channel is momentarily - /// full, which best-effort teardown does not need. - pub fn try_send_hook_reply( - &self, - reply: kigi_tool_protocol::HookReplyFrame, - ) -> Result<(), ClientError> { - let connection = self.require_connection()?; - let notif = build_hook_reply_notification(&self.inner.session, reply); - let text = serde_json::to_string(¬if).map_err(ClientError::from)?; - connection.try_send_outbound(text) - } - - /// Register the sink for inbound reverse-direction hook requests - /// (server → harness). Replaces any prior handler; the inbox loop loads - /// it per frame, so registering before or after - /// [`Self::subscribe_notifications`] both work. - /// - /// `handler` runs **inline** on the shared inbox loop that also delivers - /// [`HubNotification`](crate::notification::HubNotification)s, so it MUST - /// NOT block: only enqueue / hand the frame off (e.g. a non-blocking - /// channel `try_send`) and return promptly. Blocking here stalls - /// notification delivery for the whole session. A panic is caught (the - /// frame is dropped and the loop continues), but should still be avoided. - pub fn set_hook_request_handler(&self, handler: F) - where - F: Fn(kigi_tool_protocol::HookFrame) + Send + Sync + 'static, - { - *self.inner.hook_request_handler.lock() = Some(Arc::new(handler)); - } - - /// Shared implementation for fire-and-forget JSON-RPC requests. - /// - /// Allocates a request id, constructs a [`JsonRpcRequest`] with the - /// given method and params, serializes it, and queues it on the - /// outbound channel. Used by [`Self::send_notification`] and - /// [`Self::send_hook`] to avoid duplicating the boilerplate. - async fn send_fire_and_forget( - &self, - method: Method, - params: P, - ) -> Result<(), ClientError> { - let connection = self.require_connection()?; - let (_request_id, text) = - build_request_frame(connection, &self.inner.session, method, params)?; - connection.send_outbound(text).await - } - - /// Subscribe to server notifications for the bound session. - /// - /// The server auto-subscribes on `register_session`, so no wire - /// request is needed. Registers a session inbox on the demux and - /// spawns a task that parses notifications onto the returned channel - /// and routes reverse-direction hook requests to the handler set via - /// [`Self::set_hook_request_handler`]. - pub async fn subscribe_notifications( - &self, - ) -> Result, ClientError> { - let connection = self.require_connection()?; - - let (inbox_tx, mut inbox_rx) = mpsc::channel::(64); - connection - .demux() - .register_session_inbox(self.inner.session.clone(), inbox_tx); - - let (event_tx, event_rx) = mpsc::channel::(64); - // Clone only the handler slot (a standalone `Arc`), never `inner`: - // an `inner` clone here would keep the strong count above 1 and - // suppress the `Drop` teardown gate. - let hook_request_handler = self.inner.hook_request_handler.clone(); - tokio::spawn(async move { - while let Some(frame) = inbox_rx.recv().await { - match frame { - crate::demux::InboundFrame::Notification(value) => { - if let Some(event) = crate::notification::HubNotification::parse(&value) - && event_tx.send(event).await.is_err() - { - break; - } - } - crate::demux::InboundFrame::Request(value) => { - dispatch_inbound_hook_request(&value, &hook_request_handler); - } - } - } - }); - - Ok(event_rx) - } - - /// Query the server for remote tool descriptions via `tools.list` RPC - /// and store the result in the in-memory cache. Returns the - /// discovered tools. - pub async fn query_remote_tools(&self) -> Result, ClientError> { - self.inner.refresh_remote_tools().await - } - - /// Start background tool discovery: populate the cache, then - /// re-query on every `ToolsChanged` notification. - pub async fn start_tool_discovery(&self) { - let Ok(mut rx) = self.subscribe_notifications().await else { - tracing::warn!("tool discovery: failed to subscribe to notifications"); - return; - }; - - if let Err(e) = self.query_remote_tools().await { - tracing::warn!(error = %e, "tool discovery: initial query failed"); - } - - // Clone only the inner Arc, not a full ToolHarness — dropping - // a ToolHarness triggers begin_teardown which unregisters sessions. - let inner = self.inner.clone(); - let handle = tokio::spawn(async move { - while let Some(notification) = rx.recv().await { - match notification { - crate::notification::HubNotification::ToolsChanged { .. } => { - if let Err(e) = inner.refresh_remote_tools().await { - tracing::warn!(error = %e, "tool discovery: refresh after ToolsChanged failed"); - } - } - crate::notification::HubNotification::ToolServerStatusChanged { - session_id, - status, - } if status.status == ToolServerLifecycleStatus::Disconnected => { - inner.fail_inflight_calls_on_disconnect(&session_id); - } - _ => {} - } - } - }); - *self.inner.discovery_handle.lock() = Some(handle); - } - - /// Cooperatively release the harness's session refcount. - /// - /// Marks the harness as torn down (atomic `compare_exchange` on the - /// shared `torn_down` flag) and refcount-decrements the bound - /// session through the underlying [`HubConnection`]. The wire-level - /// `unregister_session` only fires when this is the LAST borrower - /// of the session id; otherwise the binding stays live for the - /// remaining peers. Idempotent across all clones — the first - /// caller wins the `compare_exchange`; later callers return - /// `Ok(())` without sending any frames. - /// - /// **In-flight `call(...)` futures are NOT cancelled** by - /// `shutdown`. The harness owns no run-loop — the underlying - /// connection actor keeps reading inbound frames and each - /// per-call [`ToolStream`] resolves naturally on its terminal - /// frame. To force-drain in-flight calls, call - /// `connection().request_shutdown()` (which closes the connection - /// and surfaces every parked waiter as `NetworkError`) or drop - /// the per-call stream. - pub async fn shutdown(&self) -> Result<(), ClientError> { - let Some(ref borrow) = self.inner.borrow else { - return Ok(()); // local-only: nothing to tear down - }; - if !borrow.begin_teardown() { - return Ok(()); - } - if let Some(h) = self.inner.discovery_handle.lock().take() { - h.abort(); - } - borrow.shutdown_token().cancel(); - borrow.connection().untrack_session(&self.inner.session); - Ok(()) - } -} - -/// Build the `ToolNotificationFrame` carrying a session-level event. -/// -/// Both `tool_id` and `tool_call_id` are intentionally `None`: session -/// events are not associated with any single tool dispatch. If -/// `serde_json::to_value` were to fail (it cannot — every `SessionEvent` -/// field is a primitive), an empty `null` payload is sent rather than -/// panicking. -fn build_session_event_frame(event: &SessionEvent) -> ToolNotificationFrame { - ToolNotificationFrame { - tool_call_id: None, - tool_id: None, - notification: WireToolNotification::Custom(WireCustomNotification { - kind: "session_event".to_owned(), - payload: serde_json::to_value(event).unwrap_or_default(), - }), - } -} - -/// Classify an inbound `Request` frame as a reverse-direction permission-request -/// hook, returning the decoded [`HookFrame`](kigi_tool_protocol::HookFrame) or `None`. -fn parse_permission_request_hook(value: &Value) -> Option { - let method = value.get("method").and_then(Value::as_str)?; - if method != Method::Hook.as_wire_str() { - return None; - } - let hook = - ::deserialize(value.get("params")?) - .ok()?; - // Only request/response hooks (those with a reply leg) qualify. - hook.hook_id.as_ref()?; - match &hook.event { - kigi_tool_protocol::HookEvent::Custom { kind, .. } if kind == PERMISSION_REQUEST_KIND => { - Some(hook) - } - _ => None, - } -} - -/// Dispatch one inbound `Request`: hand a permission-request hook to `handler`; -/// drop anything else (and permission requests with no handler registered). -fn dispatch_inbound_hook_request( - value: &Value, - handler: &parking_lot::Mutex>, -) { - let Some(hook) = parse_permission_request_hook(value) else { - tracing::debug!("inbound request frame is not a permission-request hook; dropping"); - return; - }; - // Clone out of the lock so the handler never runs while it is held. - let handler = handler.lock().clone(); - match handler { - // Isolate the host handler: an un-caught panic here would unwind the - // shared inbox task and silently stop ALL notification delivery for the - // session. Catch it, log, and keep the loop alive. - Some(handler) => { - if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| handler(hook))).is_err() { - tracing::error!( - "inbound hook request handler panicked; dropping frame and continuing" - ); - } - } - None => { - tracing::debug!("inbound hook request received but no handler is registered; dropping") - } - } -} - -/// Build the `hook_reply` notification answering a reverse-direction hook, -/// correlated to the request by -/// [`HookReplyFrame::hook_id`](kigi_tool_protocol::HookReplyFrame::hook_id). -fn build_hook_reply_notification( - session_id: &SessionId, - mut reply: kigi_tool_protocol::HookReplyFrame, -) -> JsonRpcNotification { - // Pin the frame's session id to the harness's bound session so params can - // never disagree with the envelope; callers supply only hook_id + result. - reply.session_id = session_id.clone(); - JsonRpcNotification { - jsonrpc: JsonRpcVersion, - session_id: Some(session_id.clone()), - seq: None, - method: Method::HookReply.as_wire_str().to_owned(), - params: reply, - } -} - -/// Owned per-call state needed to build and dispatch the matching -/// `ToolCallCompleted` event. Lives inside an `Option` on -/// [`ObservedToolStream`] so the IDs can be moved into the event by -/// value (no string clones) on emission, and so a present `Some` is the -/// only "still owes a Completed" signal we need. -struct EmissionState { - harness: ToolHarness, - tool_call_id: ToolCallId, - tool_id: ToolId, - start: std::time::Instant, -} - -/// Observability wrapper around a [`ToolStream`] returned by -/// [`ToolHarness::call`]. -/// -/// Emits exactly one [`SessionEvent::ToolCallCompleted`] for the call: -/// -/// - `Success` — terminal `Ok` flowed through. -/// - `Error` — terminal `Err` flowed through. -/// - `Cancelled` — stream was dropped before any terminal item, i.e. -/// the consumer (typically a `tokio::select!` on a cancel token) gave -/// up mid-dispatch. -/// -/// Emission is fire-and-forget via `tokio::spawn` because both -/// `poll_next` and `Drop` are synchronous. The matching -/// `ToolCallStarted` was emitted by `call` before the stream was built. -struct ObservedToolStream { - inner: ToolStream, - /// `Some` until the `ToolCallCompleted` event is scheduled, then - /// `None` so neither a subsequent `poll_next` nor `Drop` double-emits. - emission: Option, -} - -impl ObservedToolStream { - fn new( - inner: ToolStream, - harness: ToolHarness, - tool_call_id: ToolCallId, - tool_id: ToolId, - ) -> Self { - Self { - inner, - emission: Some(EmissionState { - harness, - tool_call_id, - tool_id, - start: std::time::Instant::now(), - }), - } - } - - /// Spawn a fire-and-forget task that emits `ToolCallCompleted` with - /// the given outcome, consuming the stashed [`EmissionState`]. - /// Best-effort no-op when no Tokio runtime is current — this only - /// happens during process teardown, when missing observability - /// events are an acceptable trade-off for a clean shutdown. - fn spawn_completed(&mut self, outcome: ToolCallOutcome) { - let Some(state) = self.emission.take() else { - return; - }; - let Ok(handle) = tokio::runtime::Handle::try_current() else { - return; // `state` drops here — emission lost on shutdown. - }; - let event = SessionEvent::ToolCallCompleted { - tool_call_id: state.tool_call_id.into_inner(), - tool_name: state.tool_id.into_inner(), - duration_ms: state.start.elapsed().as_millis() as u64, - outcome, - }; - let harness = state.harness; - handle.spawn(async move { - harness.emit_session_event(event).await; - }); - } -} - -impl Stream for ObservedToolStream { - type Item = ToolStreamItem; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let poll = self.inner.as_mut().poll_next(cx); - if let Poll::Ready(Some(ToolStreamItem::Terminal(result))) = &poll - && self.emission.is_some() - { - let outcome = if result.is_ok() { - ToolCallOutcome::Success - } else { - ToolCallOutcome::Error - }; - self.spawn_completed(outcome); - } - poll - } -} - -impl Drop for ObservedToolStream { - fn drop(&mut self) { - if self.emission.is_some() { - self.spawn_completed(ToolCallOutcome::Cancelled); - } - } -} - -impl Drop for ToolHarness { - fn drop(&mut self) { - let Some(ref borrow) = self.inner.borrow else { - return; // local-only: nothing to tear down - }; - // Skip teardown when other ToolHarness clones still exist. The - // harness is cloned into ObservedToolStream by `call()`; that - // internal clone's Drop must NOT race the user-held harness - // into begin_teardown (which is at-most-once and would cause a - // premature unregister_session while the user is still calling). - // strong_count == 1 means this is the last Arc reference. - if Arc::strong_count(&self.inner) > 1 { - return; - } - if !borrow.begin_teardown() { - return; - } - // Abort the discovery loop (matches shutdown() behavior). - // Lock is safe: only held briefly for take(); no async work under lock. - if let Some(h) = self.inner.discovery_handle.lock().take() { - h.abort(); - } - let inner = self.inner.clone(); - if tokio::runtime::Handle::try_current().is_ok() { - tokio::spawn(async move { - // Best-effort cleanup; a closed server WebSocket will - // surface as an error here — that is expected and - // must not panic. - if let Some(ref borrow) = inner.borrow { - borrow.shutdown_token().cancel(); - borrow.connection().untrack_session(&inner.session); - } - }); - } - } -} - -/// Send a `tool.call` JSON-RPC request, register a progress waiter, -/// and return a stream that interleaves progress notifications with -/// the eventual response terminal. -/// -/// The progress waiter is registered BEFORE the request is sent so -/// any progress notification the server forwards while the request is -/// in flight lands in the per-call channel rather than being -/// dropped as `RouteOutcome::UnknownProgress`. -async fn dispatch_remote( - connection: &Arc, - session_id: &SessionId, - tool_id: ToolId, - args: Value, - ctx: ToolCallContext, - trace_context: Option, -) -> ToolStream { - let call_id = ctx.call_id.clone(); - let cwd = ctx - .extensions - .get::() - .map(|c| c.0.to_string_lossy().into_owned()); - let behavior_version = ctx.extensions.get::().map(|v| v.0.clone()); - // Clone the cancel-on-drop identifiers ONLY when opted in — the - // default path pays no extra clone. `tool_id` itself moves into the - // request params below. - let cancel_on_drop = ctx - .extensions - .get::() - .is_some_and(|c| c.0) - .then(|| (session_id.clone(), tool_id.clone())); - - let (progress_tx, progress_rx) = mpsc::channel::(PROGRESS_BUFFER); - let demux = connection.demux(); - // Reject a concurrent caller passing the same call_id synchronously - // instead of silently overwriting the live waiter, which would strand - // the prior call's progress stream. - if demux - .try_register_progress_waiter(call_id.clone(), progress_tx) - .is_err() - { - crate::metrics::call_id_collision(); - return terminal_only(Err(client_error_to_tool_error(ClientError::CallIdInUse { - call_id, - }))); - } - - let params = ToolCallParams { - tool_call_id: call_id.clone(), - tool_id, - arguments: args, - deadline_ms: None, - behavior_version, - cwd, - trace_context, - }; - // Serialize params to a Value first so the wire shape and THIS step's - // `request_encoding` subcode stay unchanged. The envelope `to_string` - // in `build_request_frame` can't fail for a valid `Value`, so its - // differing error subcode is unreachable. - let params_value = match serde_json::to_value(¶ms) { - Ok(v) => v, - Err(err) => { - demux.unregister_progress_waiter(&call_id); - return terminal_only(Err(ToolError::custom("request_encoding", err.to_string()))); - } - }; - let (request_id, request_text) = - match build_request_frame(connection, session_id, Method::ToolCall, params_value) { - Ok(frame) => frame, - Err(err) => { - demux.unregister_progress_waiter(&call_id); - return terminal_only(Err(client_error_to_tool_error(err))); - } - }; - - let (response_tx, response_rx) = oneshot::channel(); - // Register with the session index so the in-flight short-circuit can fail - // this call on a workspace Disconnected notification for the session. - demux.register_call_response_waiter(request_id.clone(), session_id.clone(), response_tx); - - if let Err(err) = connection.send_outbound(request_text).await { - // The demux still holds the parked waiters; pull them out so - // the response oneshot doesn't sit forever on a request that - // never reached the wire. - demux.unregister_progress_waiter(&call_id); - let _ = demux.take_response_waiter(&request_id); - return terminal_only(Err(client_error_to_tool_error(err))); - } - - Box::pin(RemoteCallStream::new( - connection.clone(), - params.tool_id, - call_id, - request_id, - progress_rx, - Box::pin(response_rx), - cancel_on_drop, - )) -} - -/// Assemble a JSON-RPC request frame: allocate a request id, wrap -/// `params` under `method`, and serialize to text. The id is returned -/// for callers that park a response waiter; fire-and-forget callers -/// discard it. The enqueue (async [`HubConnection::send_outbound`] vs -/// sync [`HubConnection::try_send_outbound`]) stays with the caller. -fn build_request_frame( - connection: &HubConnection, - session_id: &SessionId, - method: Method, - params: P, -) -> Result<(RequestId, String), ClientError> { - let request_id = connection.try_alloc_request_id()?; - let req = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: Some(session_id.clone()), - method: method.as_wire_str().to_owned(), - params, - }; - let text = serde_json::to_string(&req).map_err(ClientError::from)?; - Ok((request_id, text)) -} - -/// Unified stream that interleaves per-call progress with the -/// eventual JSON-RPC response and ends after exactly one terminal. -/// -/// The progress receiver is polled while the response is pending; -/// once the response resolves the stream emits the matching -/// terminal item and returns `None` thereafter. The `Drop` impl -/// unregisters BOTH the per-call progress waiter (keyed by -/// `tool_call_id`) AND the response waiter (keyed by -/// `request_id`) from the demux — a stream that is dropped -/// before the response lands MUST NOT leak either map entry. -/// -/// `cancel_on_drop` is `Some((session_id, tool_id))` only when the caller -/// opted in (the default path stores no extra clone); on `Drop` before a -/// terminal is polled it emits one best-effort call-scoped `Cancel` hook -/// so the workspace hard-cancels the in-flight call. -struct RemoteCallStream { - connection: Arc, - /// Consumed exactly once when the terminal is built. - tool_id: Option, - call_id: ToolCallId, - request_id: RequestId, - progress_rx: Option>, - response_rx: Option< - BoxFuture<'static, Result, oneshot::error::RecvError>>, - >, - done: bool, - cancel_on_drop: Option<(SessionId, ToolId)>, -} - -impl RemoteCallStream { - fn new( - connection: Arc, - tool_id: ToolId, - call_id: ToolCallId, - request_id: RequestId, - progress_rx: mpsc::Receiver, - response_rx: BoxFuture< - 'static, - Result, oneshot::error::RecvError>, - >, - cancel_on_drop: Option<(SessionId, ToolId)>, - ) -> Self { - Self { - connection, - tool_id: Some(tool_id), - call_id, - request_id, - progress_rx: Some(progress_rx), - response_rx: Some(response_rx), - done: false, - cancel_on_drop, - } - } - - /// Best-effort, non-blocking call-scoped `Cancel` hook for the - /// cancel-on-drop path. `Drop` cannot `.await`, so the frame is - /// try-enqueued onto the outbound channel; a full or closed channel - /// drops it (the connection is already winding down / abandoning the - /// call), matching the heartbeat-pong drop discipline. - fn try_emit_cancel_on_drop(&self, session_id: &SessionId, tool_id: &ToolId) { - let hook = kigi_tool_protocol::HookFrame::cancel( - session_id.clone(), - tool_id.clone(), - self.call_id.clone(), - ); - // Counts the attempt (including a frame later dropped on a full / - // closed channel), matching `send_hook`'s count-before-send. - crate::metrics::hook_send("cancel"); - let Ok((_request_id, text)) = - build_request_frame(&self.connection, session_id, Method::Hook, hook) - else { - return; - }; - if self.connection.try_send_outbound(text).is_err() { - tracing::debug!( - call_id = %self.call_id, - "cancel-on-drop hook dropped (outbound channel full or closed)" - ); - } - } -} - -impl Drop for RemoteCallStream { - fn drop(&mut self) { - let demux = self.connection.demux(); - // Pull both waiters out of the demux. Either may already be - // gone (the response waiter is consumed by `route_response` - // when the terminal frame lands; the progress waiter is - // already removed by `unregister_progress_waiter` if that - // ran). The `Option`-returning APIs make this idempotent. - demux.unregister_progress_waiter(&self.call_id); - let _ = demux.take_response_waiter(&self.request_id); - - // `cancel_on_drop` is `Some` only when opted in. `done` flips when - // the consumer POLLS a terminal, so a drop after an unpolled - // terminal still emits one cancel — benign, the server tombstones - // the finished call (mirrors `ObservedToolStream`'s window). - if !self.done - && let Some((session_id, tool_id)) = self.cancel_on_drop.as_ref() - { - self.try_emit_cancel_on_drop(session_id, tool_id); - } - } -} - -impl Stream for RemoteCallStream { - type Item = ToolStreamItem; - - fn poll_next( - mut self: std::pin::Pin<&mut Self>, - cx: &mut std::task::Context<'_>, - ) -> std::task::Poll> { - use std::task::Poll; - if self.done { - return Poll::Ready(None); - } - - // Poll progress first so any frames that have already been - // routed by the demux drain in arrival order. The runtime's - // stream invariant is `Progress* Terminal`; checking the - // progress channel before the response future is what makes - // that invariant hold even when the server has already shipped - // the terminal frame by the time the consumer first polls. - if let Some(rx) = self.progress_rx.as_mut() { - match rx.poll_recv(cx) { - Poll::Ready(Some(frame)) => { - return Poll::Ready(Some(ToolStreamItem::Progress(progress_from_frame(frame)))); - } - Poll::Ready(None) | Poll::Pending => {} - } - } - - // Progress is empty (or closed); now check the response - // future. A `Ready` outcome here is the terminal item. - if let Some(fut) = self.response_rx.as_mut() { - match fut.as_mut().poll(cx) { - Poll::Ready(Ok(Ok(resp))) => { - self.done = true; - self.response_rx = None; - self.progress_rx = None; - let terminal = match resp.outcome { - ResponseOutcome::Result(value) => match self.tool_id.take() { - Some(tool_id) => decode_call_result(tool_id, value), - None => return Poll::Ready(None), - }, - ResponseOutcome::Error(err) => Err(error_from_envelope(err)), - }; - return Poll::Ready(Some(ToolStreamItem::Terminal(terminal))); - } - Poll::Ready(Ok(Err(err))) => { - self.done = true; - self.response_rx = None; - self.progress_rx = None; - return Poll::Ready(Some(ToolStreamItem::Terminal(Err( - client_error_to_tool_error(err), - )))); - } - Poll::Ready(Err(_)) => { - self.done = true; - self.response_rx = None; - self.progress_rx = None; - return Poll::Ready(Some(ToolStreamItem::Terminal(Err( - ToolError::network_error("response waiter dropped (connection closed)"), - )))); - } - Poll::Pending => {} - } - } else { - self.done = true; - return Poll::Ready(None); - } - - Poll::Pending - } -} - -// The `tool_call_result` success-body decode (`decode_call_result`), -// the progress-frame mapping (`progress_from_frame`), and the JSON-RPC -// envelope / `ToolErrorWire` error projections are all canonical in -// `kigi_computer_hub_core::remote`. Routing the harness through the same -// functions as the core remote proxy keeps both wire-decoding paths -// identical, so a future variant addition lands in one place. - -/// Project a [`ClientError`] into a runtime [`ToolError`]. SDK -/// transport / protocol failures collapse to -/// [`ToolError::NetworkError`]; structurally-typed wire errors -/// pass through to [`tool_error_from_wire`]; the rest land on -/// [`ToolError::Custom`] keyed by a stable subcode. -fn client_error_to_tool_error(err: ClientError) -> ToolError { - match err { - ClientError::NetworkError(message) => ToolError::network_error(message), - ClientError::ProtocolError(message) => ToolError::custom("protocol_error", message), - ClientError::AuthError(message) => ToolError::permission_denied(message), - ClientError::HandshakeAuthFailed { status } => { - ToolError::permission_denied(format!("handshake auth failed (HTTP {status})")) - } - ClientError::RegistrationConflict(message) => { - ToolError::custom("registration_conflict", message) - } - ClientError::BackpressureError(message) => ToolError::custom("backpressure", message), - ClientError::Serde(message) => ToolError::custom("serde", message), - ClientError::InvalidConfig(message) => ToolError::custom("invalid_config", message), - ClientError::Wire(wire) => tool_error_from_wire(wire), - ClientError::Closed(message) => ToolError::network_error(message), - ClientError::InsecureScheme { url } => ToolError::custom( - "insecure_scheme", - format!("refusing plaintext ws:// to non-loopback host {url}"), - ), - err @ ClientError::CallIdInUse { .. } => { - ToolError::custom("call_id_in_use", err.to_string()) - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use kigi_tool_types::ToolDescription; - use schemars::JsonSchema; - use serde::{Deserialize, Serialize}; - - #[derive(Debug)] - struct EchoTool { - id: ToolId, - } - - #[derive(Debug, Deserialize, JsonSchema)] - struct EchoArgs { - msg: String, - } - - #[derive(Debug, Serialize)] - struct EchoOut { - echoed: String, - } - impl kigi_tool_runtime::ToolOutput for EchoOut {} - - impl Tool for EchoTool { - type Args = EchoArgs; - type Output = EchoOut; - fn id(&self) -> ToolId { - self.id.clone() - } - fn description(&self, _ctx: &ListToolsContext) -> ToolDescription { - ToolDescription::new(self.id.as_str(), "echo") - } - async fn run( - &self, - _ctx: ToolCallContext, - args: Self::Args, - ) -> Result { - Ok(EchoOut { echoed: args.msg }) - } - } - - #[tokio::test] - async fn request_turn_hook_errors_without_hub_connection() { - use kigi_tool_protocol::turn_hook::{AfterTurnPayload, TurnHookOutcome, TurnHookRequest}; - - let harness = ToolHarness::local_only_with( - LocalRegistry::new(), - SessionId::new("test-session").expect("valid session"), - kigi_tool_runtime::TypedExtensions::default(), - ); - let req = TurnHookRequest::After(AfterTurnPayload { - turn_number: 1, - outcome: TurnHookOutcome::Completed, - duration_ms: 1, - tool_call_count: 0, - model_id: "grok-3".to_string(), - written_repo_paths: Vec::new(), - cancellation_category: None, - cancellation_context: None, - }); - assert!(harness.request_turn_hook(&req).await.is_err()); - } - - fn pending_bind_test_harness() -> ToolHarness { - ToolHarness::local_only_with( - LocalRegistry::new(), - SessionId::new("pending-bind-test").expect("valid session"), - kigi_tool_runtime::TypedExtensions::default(), - ) - } - - #[tokio::test] - async fn no_pending_bind_awaits_to_self_and_probes_none() { - let harness = pending_bind_test_harness(); - assert!(!harness.has_pending_bind()); - assert!(harness.try_bound().is_none()); - assert!(harness.await_bound().await.is_ok()); - } - - #[tokio::test] - async fn pending_bind_resolves_to_connected_harness() { - let harness = ToolHarness::local_with_pending_bind( - LocalRegistry::new(), - SessionId::new("twin").expect("valid session"), - kigi_tool_runtime::TypedExtensions::default(), - async move { Ok(pending_bind_test_harness()) }, - ); - assert!(harness.has_pending_bind()); - let bound = harness.await_bound().await.expect("bind ok"); - assert!(!bound.has_pending_bind()); - } - - #[tokio::test] - async fn pending_bind_failure_surfaces_error() { - let harness = ToolHarness::local_with_pending_bind( - LocalRegistry::new(), - SessionId::new("twin").expect("valid session"), - kigi_tool_runtime::TypedExtensions::default(), - async { Err(Arc::::from("boom")) }, - ); - let err = harness.await_bound().await.expect_err("bind failed"); - assert!(err.contains("boom")); - } - - #[tokio::test] - async fn try_bound_is_none_until_resolved_then_some() { - let (tx, rx) = tokio::sync::oneshot::channel::<()>(); - let harness = ToolHarness::local_with_pending_bind( - LocalRegistry::new(), - SessionId::new("twin").expect("valid session"), - kigi_tool_runtime::TypedExtensions::default(), - async move { - let _ = rx.await; - Ok(pending_bind_test_harness()) - }, - ); - assert!(harness.try_bound().is_none(), "still binding"); - tx.send(()).unwrap(); - harness.await_bound().await.expect("bind ok"); - assert!(harness.try_bound().is_some(), "resolved"); - } - - #[tokio::test] - async fn try_bound_flips_to_some_without_await() { - let harness = ToolHarness::local_with_pending_bind( - LocalRegistry::new(), - SessionId::new("twin").expect("valid session"), - kigi_tool_runtime::TypedExtensions::default(), - async move { Ok(pending_bind_test_harness()) }, - ); - let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); - loop { - match harness.try_bound() { - Some(result) => { - result.expect("bind ok"); - break; - } - None if std::time::Instant::now() < deadline => { - tokio::task::yield_now().await; - } - None => panic!("try_bound never observed the completed background bind"), - } - } - } - - #[tokio::test] - async fn lazy_bind_does_not_start_until_await_bound() { - let started = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let started_in_fut = Arc::clone(&started); - let harness = ToolHarness::local_with_lazy_bind( - LocalRegistry::new(), - SessionId::new("lazy").expect("valid session"), - kigi_tool_runtime::TypedExtensions::default(), - async move { - started_in_fut.store(true, std::sync::atomic::Ordering::SeqCst); - Ok(pending_bind_test_harness()) - }, - ); - assert!(harness.has_pending_bind()); - - // Probing must NOT kick off the bind, even after giving the runtime - // several chances to make progress. - for _ in 0..16 { - assert!(harness.try_bound().is_none(), "lazy bind not started yet"); - tokio::task::yield_now().await; - } - assert!( - !started.load(std::sync::atomic::Ordering::SeqCst), - "lazy bind future must not run until await_bound" - ); - - // First await_bound starts and resolves it. - let bound = harness.await_bound().await.expect("bind ok"); - assert!(!bound.has_pending_bind()); - assert!( - started.load(std::sync::atomic::Ordering::SeqCst), - "await_bound must start the lazy bind" - ); - assert!(harness.try_bound().is_some(), "resolved after await_bound"); - } - - #[tokio::test] - async fn lazy_bind_failure_surfaces_error() { - let harness = ToolHarness::local_with_lazy_bind( - LocalRegistry::new(), - SessionId::new("lazy").expect("valid session"), - kigi_tool_runtime::TypedExtensions::default(), - async { Err(Arc::::from("boom")) }, - ); - let err = harness.await_bound().await.expect_err("bind failed"); - assert!(err.contains("boom")); - } - - #[test] - fn local_registry_starts_empty() { - let registry = LocalRegistry::new(); - assert_eq!(registry.len(), 0); - assert!(registry.is_empty()); - let id = ToolId::new("missing").expect("valid"); - assert!(!registry.contains(&id)); - assert!(registry.find(&id).is_none()); - assert!(!registry.unregister(&id)); - } - - #[test] - fn has_remote_tool_consults_only_the_remote_cache() { - let registry = LocalRegistry::new(); - registry.register(EchoTool { - id: ToolId::new("local_echo").expect("valid"), - }); - let harness = ToolHarness::local_only_with( - registry, - SessionId::new("remote-cache-session").expect("valid session"), - kigi_tool_runtime::TypedExtensions::default(), - ); - - assert!(!harness.has_remote_tool("bash")); - assert!(!harness.has_remote_tool("local_echo")); - - harness.seed_remote_tools_for_tests(vec![ - ToolDescription::new("bash", "run a shell command"), - ToolDescription::new("write_file", "write a file"), - ]); - assert!(harness.has_remote_tool("bash")); - assert!(harness.has_remote_tool("write_file")); - assert!(!harness.has_remote_tool("read_file")); - assert!(!harness.has_remote_tool("local_echo")); - - // Cleared on unbind. - harness.seed_remote_tools_for_tests(Vec::new()); - assert!(!harness.has_remote_tool("bash")); - } - - #[test] - fn local_registry_register_then_find_returns_handle() { - let registry = LocalRegistry::new(); - let id = ToolId::new("echo").expect("valid"); - let prev = registry.register(EchoTool { id: id.clone() }); - assert!(prev.is_none(), "first register has nothing to displace"); - assert_eq!(registry.len(), 1); - assert!(registry.contains(&id)); - let handle = registry.find(&id).expect("registered tool resolves"); - assert_eq!(handle.id(), id); - } - - #[test] - fn local_registry_register_returns_displaced_handle_on_duplicate_id() { - let registry = LocalRegistry::new(); - let id = ToolId::new("echo").expect("valid"); - let first = registry.register(EchoTool { id: id.clone() }); - assert!(first.is_none()); - let displaced = registry.register(EchoTool { id: id.clone() }); - assert!( - displaced.is_some(), - "second register on same id must return the displaced handle" - ); - assert_eq!(registry.len(), 1, "id remains unique"); - } - - #[test] - fn local_registry_unregister_returns_true_then_false() { - let registry = LocalRegistry::new(); - let id = ToolId::new("echo").expect("valid"); - registry.register(EchoTool { id: id.clone() }); - assert!( - registry.unregister(&id), - "first unregister removes the entry" - ); - assert!( - !registry.unregister(&id), - "second unregister sees the entry already gone" - ); - assert!(registry.is_empty()); - } - - #[test] - fn local_registry_register_arc_uses_shared_allocation() { - let registry = LocalRegistry::new(); - let id = ToolId::new("echo_arc").expect("valid"); - let tool = Arc::new(EchoTool { id: id.clone() }); - let prev = registry.register_arc(tool.clone()); - assert!(prev.is_none()); - // The original Arc is still held by the test (refcount ≥ 2): - // the registry stores its own clone via `ErasedTool::from_arc`, - // so dropping the test's clone does not invalidate the registry. - drop(tool); - let handle = registry.find(&id).expect("handle still present"); - assert_eq!(handle.id(), id); - } - - #[test] - fn local_registry_register_alias_resolves_both_ids() { - let registry = LocalRegistry::new(); - let full_id = ToolId::new("slack___search_channels").expect("valid"); - let bare_id = ToolId::new("search_channels").expect("valid"); - - registry.register(EchoTool { - id: full_id.clone(), - }); - assert!(registry.contains(&full_id)); - assert!(!registry.contains(&bare_id)); - - assert!(registry.register_alias(bare_id.clone(), &full_id)); - assert!(registry.contains(&bare_id)); - - let h1 = registry.find(&full_id).expect("full id resolves"); - let h2 = registry.find(&bare_id).expect("bare alias resolves"); - assert_eq!(h1.id(), h2.id(), "both resolve to the same tool"); - } - - #[test] - fn local_registry_register_alias_returns_false_for_missing_target() { - let registry = LocalRegistry::new(); - let alias = ToolId::new("alias").expect("valid"); - let missing = ToolId::new("missing").expect("valid"); - assert!(!registry.register_alias(alias.clone(), &missing)); - assert!(!registry.contains(&alias)); - } - - #[test] - fn local_registry_list_tools_preserves_insertion_order() { - let registry = LocalRegistry::new(); - let names = [ - "web_search", - "code_execution", - "generate_image", - "browse_page", - ]; - for name in &names { - registry.register(EchoTool { - id: ToolId::new(*name).expect("valid"), - }); - } - let ctx = ListToolsContext::new(); - let listed: Vec = registry - .list_tools(&ctx) - .into_iter() - .map(|d| d.name) - .collect(); - assert_eq!( - listed, - names.iter().map(|s| s.to_string()).collect::>(), - "list_tools must return tools in registration (insertion) order" - ); - } - - #[tokio::test] - async fn builder_missing_pool_errors() { - let url = Url::parse("ws://127.0.0.1:0/v1/tools").expect("valid url"); - let cred = AuthCredential::bearer("ignored"); - let session = SessionId::new("s").expect("valid"); - let err = ToolHarnessBuilder::default() - .url(url) - .auth(cred) - .session(session) - .build() - .await - .expect_err("missing pool must fail"); - assert!(matches!(err, ClientError::InvalidConfig(msg) if msg.contains("missing pool"))); - } - - #[tokio::test] - async fn builder_missing_url_errors() { - let pool = HubConnectionPool::new(); - let cred = AuthCredential::bearer("ignored"); - let session = SessionId::new("s").expect("valid"); - let err = ToolHarnessBuilder::default() - .pool(pool) - .auth(cred) - .session(session) - .build() - .await - .expect_err("missing url must fail"); - assert!(matches!(err, ClientError::InvalidConfig(msg) if msg.contains("missing url"))); - } - - #[tokio::test] - async fn builder_missing_auth_errors() { - let pool = HubConnectionPool::new(); - let url = Url::parse("ws://127.0.0.1:0/v1/tools").expect("valid url"); - let session = SessionId::new("s").expect("valid"); - let err = ToolHarnessBuilder::default() - .pool(pool) - .url(url) - .session(session) - .build() - .await - .expect_err("missing auth must fail"); - assert!(matches!(err, ClientError::InvalidConfig(msg) if msg.contains("missing auth"))); - } - - #[tokio::test] - async fn builder_missing_session_errors() { - let pool = HubConnectionPool::new(); - let url = Url::parse("ws://127.0.0.1:0/v1/tools").expect("valid url"); - let cred = AuthCredential::bearer("ignored"); - let err = ToolHarnessBuilder::default() - .pool(pool) - .url(url) - .auth(cred) - .build() - .await - .expect_err("missing session must fail"); - assert!(matches!(err, ClientError::InvalidConfig(msg) if msg.contains("missing session"))); - } - - #[test] - fn local_registry_clone_shares_backing_store() { - let a = LocalRegistry::new(); - let id = ToolId::new("shared").expect("valid"); - a.register(EchoTool { id: id.clone() }); - - let b = a.clone(); - assert_eq!(b.len(), 1); - assert!(b.find(&id).is_some()); - - // Register through b, visible through a - let id2 = ToolId::new("shared2").expect("valid"); - b.register(EchoTool { id: id2.clone() }); - assert_eq!(a.len(), 2); - assert!(a.find(&id2).is_some()); - } - - #[test] - fn local_registry_unregister_visible_to_clones() { - let a = LocalRegistry::new(); - let id = ToolId::new("ephemeral").expect("valid"); - a.register(EchoTool { id: id.clone() }); - - let b = a.clone(); - assert!(b.unregister(&id)); - assert!(a.is_empty()); - } - - // ── Session event frame construction ──────────────────────────── - - #[test] - fn session_event_frame_has_no_tool_correlation_ids() { - let event = SessionEvent::TurnStarted { - turn_number: 1, - model_id: "grok-3".into(), - yolo_mode: false, - }; - let frame = build_session_event_frame(&event); - assert!(frame.tool_id.is_none()); - assert!(frame.tool_call_id.is_none()); - match &frame.notification { - WireToolNotification::Custom(c) => { - assert_eq!(c.kind, "session_event"); - assert_eq!(c.payload["event_type"], "turn_started"); - assert_eq!(c.payload["turn_number"], 1); - } - other => panic!("expected Custom, got {other:?}"), - } - } - - #[test] - fn session_event_frame_round_trips_through_serde() { - let event = SessionEvent::ToolCallStarted { - tool_call_id: "call-42".into(), - tool_name: "read_file".into(), - turn_number: 3, - }; - let frame = build_session_event_frame(&event); - let json = serde_json::to_value(&frame).unwrap(); - let back: ToolNotificationFrame = serde_json::from_value(json).unwrap(); - assert_eq!(back, frame); - } - - // ── Local-only call() must not wrap with observability ───────── - - fn build_local_only_harness_with(tool: EchoTool) -> ToolHarness { - let session = SessionId::new("test-observed").expect("valid"); - let registry = LocalRegistry::new(); - registry.register(tool); - ToolHarness::local_only_with(registry, session, Default::default()) - } - - #[tokio::test] - async fn local_only_call_returns_raw_stream_no_observability_wrap() { - // Sanity: a local-only harness has `borrow == None`, so `call()` - // must skip the `ObservedToolStream` wrap and emit nothing — the - // existing local-tool test surface stays byte-for-byte identical. - let tool_id = ToolId::new("echo").expect("valid"); - let harness = build_local_only_harness_with(EchoTool { - id: tool_id.clone(), - }); - let mut stream = harness - .call( - tool_id, - serde_json::json!({ "msg": "hi" }), - ToolCallContext::default(), - ) - .await; - - let mut saw_terminal_ok = false; - while let Some(item) = futures::StreamExt::next(&mut stream).await { - if let ToolStreamItem::Terminal(Ok(_)) = item { - saw_terminal_ok = true; - } - } - assert!(saw_terminal_ok, "local echo must yield Terminal(Ok)"); - } - - #[tokio::test] - async fn local_only_emit_session_event_is_noop() { - // No server borrow → `emit_session_event` must short-circuit before - // building a frame. The test passes if the call returns without - // touching `send_notification` (no panic from a missing - // connection actor). - let harness = ToolHarness::local_only_with( - LocalRegistry::new(), - SessionId::new("test-no-hub").expect("valid"), - Default::default(), - ); - harness - .emit_session_event(SessionEvent::PhaseChanged { - phase: kigi_tool_protocol::session_event::SessionPhase::Idle, - }) - .await; - } - - /// Wrap a [`HookFrame`](kigi_tool_protocol::HookFrame) in the JSON-RPC `Request` envelope. - fn inbound_hook_request_frame(hook: &kigi_tool_protocol::HookFrame) -> Value { - serde_json::json!({ - "jsonrpc": "2.0", - "id": "h1", - "session_id": hook.session_id.as_str(), - "method": Method::Hook.as_wire_str(), - "params": serde_json::to_value(hook).expect("serialize hook"), - }) - } - - #[test] - fn parse_permission_request_hook_matches_request_response_hook() { - let hook = kigi_tool_protocol::HookFrame::custom_request( - SessionId::new("s1").expect("valid"), - "hook-7".to_owned(), - PERMISSION_REQUEST_KIND.to_owned(), - serde_json::json!({ "tool_call_id": "call-1" }), - ); - let parsed = parse_permission_request_hook(&inbound_hook_request_frame(&hook)) - .expect("permission request matches"); - assert_eq!(parsed.hook_id.as_deref(), Some("hook-7")); - match parsed.event { - kigi_tool_protocol::HookEvent::Custom { kind, payload } => { - assert_eq!(kind, PERMISSION_REQUEST_KIND); - assert_eq!(payload, serde_json::json!({ "tool_call_id": "call-1" })); - } - other => panic!("expected Custom event, got {other:?}"), - } - } - - #[test] - fn parse_permission_request_hook_rejects_other_custom_kind() { - let hook = kigi_tool_protocol::HookFrame::custom_request( - SessionId::new("s1").expect("valid"), - "hook-7".to_owned(), - kigi_tool_protocol::turn_hook::TURN_HOOK_KIND.to_owned(), - serde_json::json!({}), - ); - assert!(parse_permission_request_hook(&inbound_hook_request_frame(&hook)).is_none()); - } - - #[test] - fn parse_permission_request_hook_rejects_missing_hook_id() { - let hook = kigi_tool_protocol::HookFrame::custom( - SessionId::new("s1").expect("valid"), - PERMISSION_REQUEST_KIND.to_owned(), - serde_json::json!({}), - ); - assert!(hook.hook_id.is_none()); - assert!(parse_permission_request_hook(&inbound_hook_request_frame(&hook)).is_none()); - } - - #[test] - fn parse_permission_request_hook_rejects_non_custom_event() { - let hook = kigi_tool_protocol::HookFrame { - session_id: SessionId::new("s1").expect("valid"), - tool_id: None, - call_id: None, - hook_id: Some("hook-7".to_owned()), - event: kigi_tool_protocol::HookEvent::Pause, - trace_context: None, - }; - assert!(parse_permission_request_hook(&inbound_hook_request_frame(&hook)).is_none()); - } - - #[test] - fn parse_permission_request_hook_rejects_non_hook_method() { - let hook = kigi_tool_protocol::HookFrame::custom_request( - SessionId::new("s1").expect("valid"), - "hook-7".to_owned(), - PERMISSION_REQUEST_KIND.to_owned(), - serde_json::json!({}), - ); - let mut frame = inbound_hook_request_frame(&hook); - frame["method"] = serde_json::json!("tool_call_request"); - assert!(parse_permission_request_hook(&frame).is_none()); - } - - #[tokio::test] - async fn registered_handler_receives_matching_request_only() { - let harness = ToolHarness::local_only_with( - LocalRegistry::new(), - SessionId::new("s1").expect("valid session"), - kigi_tool_runtime::TypedExtensions::default(), - ); - let (tx, mut rx) = mpsc::channel::(4); - harness.set_hook_request_handler(move |hook| { - tx.try_send(hook).expect("handler channel has capacity"); - }); - let slot = &harness.inner.hook_request_handler; - - let other = kigi_tool_protocol::HookFrame::custom_request( - SessionId::new("s1").expect("valid"), - "hook-0".to_owned(), - kigi_tool_protocol::turn_hook::TURN_HOOK_KIND.to_owned(), - serde_json::json!({}), - ); - dispatch_inbound_hook_request(&inbound_hook_request_frame(&other), slot); - assert!( - rx.try_recv().is_err(), - "non-permission request must be dropped" - ); - - let perm = kigi_tool_protocol::HookFrame::custom_request( - SessionId::new("s1").expect("valid"), - "hook-7".to_owned(), - PERMISSION_REQUEST_KIND.to_owned(), - serde_json::json!({ "tool_call_id": "call-1" }), - ); - dispatch_inbound_hook_request(&inbound_hook_request_frame(&perm), slot); - let received = rx.recv().await.expect("handler invoked"); - assert_eq!(received.hook_id.as_deref(), Some("hook-7")); - } - - #[test] - fn dispatch_inbound_hook_request_without_handler_invokes_nothing() { - use std::sync::atomic::{AtomicUsize, Ordering}; - - let calls = Arc::new(AtomicUsize::new(0)); - let counter = calls.clone(); - let slot: parking_lot::Mutex> = - parking_lot::Mutex::new(Some(Arc::new(move |_hook| { - counter.fetch_add(1, Ordering::SeqCst); - }))); - let perm = kigi_tool_protocol::HookFrame::custom_request( - SessionId::new("s1").expect("valid"), - "hook-7".to_owned(), - PERMISSION_REQUEST_KIND.to_owned(), - serde_json::json!({}), - ); - let frame = inbound_hook_request_frame(&perm); - - dispatch_inbound_hook_request(&frame, &slot); - assert_eq!(calls.load(Ordering::SeqCst), 1); - - *slot.lock() = None; - dispatch_inbound_hook_request(&frame, &slot); - assert_eq!( - calls.load(Ordering::SeqCst), - 1, - "no-handler path must invoke nothing" - ); - } - - #[test] - fn hook_reply_notification_has_correct_wire_shape() { - let session = SessionId::new("s1").expect("valid"); - let reply = kigi_tool_protocol::HookReplyFrame { - session_id: session.clone(), - hook_id: "hook-7".to_owned(), - result: serde_json::json!({ "outcome": "approve" }), - }; - let notif = build_hook_reply_notification(&session, reply); - assert_eq!(notif.method, Method::HookReply.as_wire_str()); - let json: Value = serde_json::from_str(&serde_json::to_string(¬if).unwrap()).unwrap(); - assert_eq!(json["jsonrpc"], "2.0"); - assert_eq!(json["method"], "hook_reply"); - assert_eq!(json["session_id"], "s1"); - assert!( - json.get("id").is_none(), - "hook_reply is a notification — must carry no id" - ); - assert!(json.get("seq").is_none(), "None seq must be omitted"); - assert_eq!(json["params"]["hook_id"], "hook-7"); - assert_eq!(json["params"]["session_id"], "s1"); - assert_eq!(json["params"]["result"]["outcome"], "approve"); - } - - #[tokio::test] - async fn send_hook_reply_errors_without_hub_connection() { - let harness = ToolHarness::local_only_with( - LocalRegistry::new(), - SessionId::new("test-session").expect("valid session"), - kigi_tool_runtime::TypedExtensions::default(), - ); - let reply = kigi_tool_protocol::HookReplyFrame { - session_id: harness.session().clone(), - hook_id: "hook-7".to_owned(), - result: Value::Null, - }; - assert!(harness.send_hook_reply(reply).await.is_err()); - } - - #[test] - fn try_send_hook_reply_errors_without_hub_connection() { - let harness = ToolHarness::local_only_with( - LocalRegistry::new(), - SessionId::new("test-session").expect("valid session"), - kigi_tool_runtime::TypedExtensions::default(), - ); - let reply = kigi_tool_protocol::HookReplyFrame { - session_id: harness.session().clone(), - hook_id: "hook-7".to_owned(), - result: Value::Null, - }; - assert!(harness.try_send_hook_reply(reply).is_err()); - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/lib.rs b/crates/common/kigi-computer-hub-sdk/src/lib.rs deleted file mode 100644 index f6853f0..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/lib.rs +++ /dev/null @@ -1,71 +0,0 @@ -//! Tool-server and harness SDK. -//! -//! Single crate hosting both the tool-server runtime and the -//! harness-side dispatch surface. The shared substrate — -//! [`HubConnectionPool`], [`HubConnection`], the inbound demux, the -//! refcount-managed bound-session set, and the transparent reconnect / -//! replay state machine — lives here so both ends speak through one -//! frame multiplex on top of one WebSocket per `(url, principal)`. -//! -//! The server entry point is [`ToolServer`]: build it via -//! [`ToolServerBuilder`], wire one or more [`ToolServerHandler`] -//! implementations, and call [`ToolServer::run`] to drive the inbound -//! loop. The harness entry point is [`ToolHarness`]: build it via -//! [`ToolHarnessBuilder`], optionally seed it with in-process -//! [`kigi_tool_runtime::Tool`] implementations, and call -//! [`ToolHarness::call`] to dispatch a tool call. Authorisation -//! credentials (`AuthCredential`) plus the target URL determine -//! which pool entry the consumer attaches to; multiple -//! [`ToolServer`] / [`ToolHarness`] instances against the same -//! `(url, principal)` share a single connection and refcount their -//! session bindings. - -#![forbid(unsafe_code)] - -pub(crate) mod admission; -pub mod auth; -pub(crate) mod cancel; -pub mod connection; -pub(crate) mod connection_borrow; -pub mod demux; -pub(crate) mod donate_pump; -pub mod error; -pub mod handshake; -pub mod harness; -pub mod log_donate; -#[cfg(feature = "metrics")] -pub mod metric_donate; -pub mod metrics; -pub mod notification; -pub mod observability; -pub mod pool; -pub mod refcount; -pub mod server; -pub mod trace_donate; - -pub mod oidc_provider; - -pub use auth::{AuthCredential, AuthIdentity, AuthProvider, PrincipalKey, SharedAuthProvider}; -pub use connection::{ConnKey, HubConnection, ReconnectEvent}; -pub use error::ClientError; -pub use harness::{ - CancelOnDrop, LocalRegistry, ModelOutputExtractor, SessionBindReport, ToolHarness, - ToolHarnessBuilder, extractor_for, -}; -pub use log_donate::{DonatingLogLayer, LogDonationPump, LogDonationSender, flush_log_layer}; -#[cfg(feature = "metrics")] -pub use metric_donate::MetricDonationPump; -pub use notification::HubNotification; -pub use observability::ObservabilityBridge; -pub use oidc_provider::{ - OidcAuthProvider, OidcAuthProviderBuilder, OnRefreshCallback, RefreshEvent, -}; -pub use pool::HubConnectionPool; -pub use server::{ - ResolvedSessionHandlers, SessionHandlerResolver, SystemNotifyAck, ToolServer, - ToolServerBuilder, ToolServerHandler, WeakToolServer, -}; -pub use trace_donate::{HubDonatingReporter, TraceDonationPump}; -// Re-exported so consumers that depend only on the SDK can recognize the -// server's `workspace_unavailable` error without also pulling in the core crate. -pub use kigi_computer_hub_core::is_workspace_unavailable; diff --git a/crates/common/kigi-computer-hub-sdk/src/log_donate.rs b/crates/common/kigi-computer-hub-sdk/src/log_donate.rs deleted file mode 100644 index e881549..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/log_donate.rs +++ /dev/null @@ -1,592 +0,0 @@ -//! Forward curated `tracing` events to the connected server over the -//! WebSocket transport (`logs.donate`). -//! -//! [`DonatingLogLayer`] is installed **inert** at startup and activated -//! post-connect by swapping in a [`LogDonationSender`] (a global-subscriber -//! constraint); while inert, selected events are dropped before enqueueing. -//! -//! Only events on the [`TELEMETRY_TARGET`] target at `>= INFO` are -//! forwarded, and only fields in [`ALLOWED_FIELDS`] are included; other -//! fields such as `reason`/`error` are omitted. - -use std::sync::{Arc, LazyLock}; -use std::time::{Duration, Instant}; - -use arc_swap::ArcSwapOption; -use base64::Engine as _; -use fastrace::collector::SpanContext; -use kigi_tool_protocol::{MAX_DONATION_BYTES, MAX_LOG_RECORDS_PER_DONATION}; -use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; -use opentelemetry_proto::tonic::common::v1::{AnyValue, InstrumentationScope, KeyValue, any_value}; -use opentelemetry_proto::tonic::logs::v1::{LogRecord, ResourceLogs, ScopeLogs}; -use opentelemetry_proto::tonic::resource::v1::Resource; -use prost::Message as _; -use tokio::sync::mpsc; -use tracing::Level; -use tracing::field::{Field, Visit}; -use tracing_subscriber::Layer; -use tracing_subscriber::layer::Context; - -use crate::donate_pump::{ - PENDING_FLUSHES, PumpMsg, drain_via, make_resource, now_unix_nanos, run_pump, string_kv, - string_value, -}; -use crate::server::ToolServer; - -/// Stable target the workspace routes selected events through. -/// The layer selects exactly this target, ignoring global -/// `RUST_LOG`. The server re-stamps it as the OTLP scope name. -pub const TELEMETRY_TARGET: &str = "workspace::telemetry"; - -/// Set of forwardable field names — guaranteed-literal or numeric. -/// Only the listed fields are included; other fields such as -/// `reason`/`error`/`object_path`/`gcs_path` are omitted. -const ALLOWED_FIELDS: &[&str] = &[ - "session_id", - "turn_number", - "phase", - "bytes", - "file_count", - "pending", - "pending_bytes", - "sample_period_secs", - "error_category", - "outcome", - "skip_reason", - "drain_reason", - "grace_ms", - "active_at_start", - "pending_at_start", - "producers_at_start", -]; - -/// Flush a buffered batch once it reaches this many records. -const LOG_BATCH_FLUSH_RECORDS: usize = 32; -/// Flush a partial batch once its oldest record is at least this old -/// (checked on the next event; the tail is fenced by teardown). -const LOG_BATCH_MAX_AGE: Duration = Duration::from_secs(2); - -fn is_allowed(name: &str) -> bool { - ALLOWED_FIELDS.contains(&name) -} - -/// `tracing::Level` → OTLP (`SeverityText`, `SeverityNumber`). -fn severity(level: &Level) -> (&'static str, i32) { - match *level { - Level::ERROR => ("ERROR", 17), - Level::WARN => ("WARN", 13), - Level::INFO => ("INFO", 9), - Level::DEBUG => ("DEBUG", 5), - Level::TRACE => ("TRACE", 1), - } -} - -/// `>= INFO` in severity terms (INFO/WARN/ERROR). Note tracing orders -/// `ERROR < WARN < INFO < DEBUG < TRACE`, so this is `level <= INFO`. -fn at_least_info(level: &Level) -> bool { - *level <= Level::INFO -} - -/// Big-endian byte encoding of the local parent's ids into the OTLP -/// 16-byte / 8-byte fields; empty when no fastrace local parent is -/// active (the common case for detached producer tasks). -fn current_ids() -> (Vec, Vec) { - match SpanContext::current_local_parent() { - Some(ctx) => encode_ids(&ctx), - None => (Vec::new(), Vec::new()), - } -} - -fn encode_ids(ctx: &SpanContext) -> (Vec, Vec) { - ( - ctx.trace_id.0.to_be_bytes().to_vec(), - ctx.span_id.0.to_be_bytes().to_vec(), - ) -} - -/// Field visitor: keeps the message as the OTLP `Body` and only -/// allowlisted fields as attributes; everything else is dropped. -#[derive(Default)] -struct AllowlistVisitor { - body: Option, - attributes: Vec, -} - -impl Visit for AllowlistVisitor { - fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) { - let name = field.name(); - if name == "message" { - self.body = Some(format!("{value:?}")); - } else if is_allowed(name) { - self.attributes.push(string_kv(name, format!("{value:?}"))); - } - } - - fn record_str(&mut self, field: &Field, value: &str) { - let name = field.name(); - if name == "message" { - self.body = Some(value.to_owned()); - } else if is_allowed(name) { - self.attributes.push(string_kv(name, value.to_owned())); - } - } - - fn record_i64(&mut self, field: &Field, value: i64) { - if is_allowed(field.name()) { - self.push_int(field.name(), value); - } - } - - fn record_u64(&mut self, field: &Field, value: u64) { - if is_allowed(field.name()) { - self.push_int(field.name(), value as i64); - } - } - - fn record_bool(&mut self, field: &Field, value: bool) { - if is_allowed(field.name()) { - self.attributes.push(KeyValue { - key: field.name().to_owned(), - value: Some(AnyValue { - value: Some(any_value::Value::BoolValue(value)), - }), - ..Default::default() - }); - } - } - - fn record_f64(&mut self, field: &Field, value: f64) { - if is_allowed(field.name()) { - self.attributes.push(KeyValue { - key: field.name().to_owned(), - value: Some(AnyValue { - value: Some(any_value::Value::DoubleValue(value)), - }), - ..Default::default() - }); - } - } -} - -impl AllowlistVisitor { - fn push_int(&mut self, name: &str, value: i64) { - self.attributes.push(KeyValue { - key: name.to_owned(), - value: Some(AnyValue { - value: Some(any_value::Value::IntValue(value)), - }), - ..Default::default() - }); - } -} - -fn build_log_record(level: &Level, visitor: AllowlistVisitor) -> LogRecord { - let now_nanos = now_unix_nanos(); - let (text, number) = severity(level); - let (trace_id, span_id) = current_ids(); - LogRecord { - time_unix_nano: now_nanos, - observed_time_unix_nano: now_nanos, - severity_number: number, - severity_text: text.to_owned(), - body: visitor.body.map(string_value), - attributes: visitor.attributes, - trace_id, - span_id, - ..Default::default() - } -} - -/// Encodes batches of OTLP `LogRecord`s onto the pump channel. Chunks at -/// [`MAX_LOG_RECORDS_PER_DONATION`], drops payloads over -/// [`MAX_DONATION_BYTES`], and never blocks. -#[derive(Clone)] -struct PumpLogExporter { - tx: mpsc::Sender, - resource: Resource, -} - -impl PumpLogExporter { - fn export(&self, mut records: Vec) { - while !records.is_empty() { - let chunk = if records.len() > MAX_LOG_RECORDS_PER_DONATION { - let rest = records.split_off(MAX_LOG_RECORDS_PER_DONATION); - std::mem::replace(&mut records, rest) - } else { - std::mem::take(&mut records) - }; - let request = ExportLogsServiceRequest { - resource_logs: vec![ResourceLogs { - resource: Some(self.resource.clone()), - scope_logs: vec![ScopeLogs { - scope: Some(InstrumentationScope { - name: TELEMETRY_TARGET.to_owned(), - ..Default::default() - }), - log_records: chunk, - schema_url: String::new(), - }], - schema_url: String::new(), - }], - }; - let bytes = request.encode_to_vec(); - if bytes.len() > MAX_DONATION_BYTES { - tracing::debug!(len = bytes.len(), "dropping oversized log donation payload"); - continue; - } - let payload = base64::engine::general_purpose::STANDARD.encode(bytes); - if self.tx.try_send(PumpMsg::Payload(payload)).is_err() { - tracing::debug!("log donation queue full; dropping log batch"); - } - } - } -} - -/// Activation handle swapped into an inert [`DonatingLogLayer`]. Wraps -/// the pump sender plus the resource (`service.name`) the layer needs to -/// encode batches. -pub struct LogDonationSender { - exporter: PumpLogExporter, -} - -impl LogDonationSender { - fn export(&self, records: Vec) { - self.exporter.export(records); - } -} - -#[derive(Default)] -struct LogBatch { - records: Vec, - oldest: Option, -} - -struct LogLayerShared { - sender: ArcSwapOption, - batch: parking_lot::Mutex, -} - -impl LogLayerShared { - /// Buffer a record; return any records due for flush (count/age). - fn push(&self, record: LogRecord) -> Vec { - let mut batch = self.batch.lock(); - if batch.records.is_empty() { - batch.oldest = Some(Instant::now()); - } - batch.records.push(record); - let due = batch.records.len() >= LOG_BATCH_FLUSH_RECORDS - || batch - .oldest - .is_some_and(|t| t.elapsed() >= LOG_BATCH_MAX_AGE); - if due { - batch.oldest = None; - std::mem::take(&mut batch.records) - } else { - Vec::new() - } - } - - /// Force the buffered batch onto the pump (teardown analogue of - /// `fastrace::flush()`); no-op while inert or empty. - fn flush(&self) { - let records = { - let mut batch = self.batch.lock(); - batch.oldest = None; - std::mem::take(&mut batch.records) - }; - if records.is_empty() { - return; - } - if let Some(sender) = self.sender.load_full() { - sender.export(records); - } - } -} - -/// Process-global handle to the active layer's shared state so -/// [`flush_log_layer`] can drive a teardown flush without a reference. -static ACTIVE_LOG_LAYER: LazyLock> = - LazyLock::new(ArcSwapOption::empty); - -/// A composable [`tracing_subscriber::Layer`] that converts selected -/// events into OTLP log records and batches them onto the pump. -/// Installed inert; activated by [`Self::activate`]. -#[derive(Clone)] -pub struct DonatingLogLayer { - shared: Arc, -} - -impl DonatingLogLayer { - /// Install inert (no sender): selected events are dropped until - /// [`Self::activate`] swaps a sender in. Registers itself as the - /// process-global flush target. - pub fn new_inert() -> Self { - let shared = Arc::new(LogLayerShared { - sender: ArcSwapOption::empty(), - batch: parking_lot::Mutex::new(LogBatch::default()), - }); - ACTIVE_LOG_LAYER.store(Some(shared.clone())); - Self { shared } - } - - /// Swap in the donation sender, activating donation. - pub fn activate(&self, sender: LogDonationSender) { - self.shared.sender.store(Some(Arc::new(sender))); - } -} - -impl Layer for DonatingLogLayer { - fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) { - let Some(sender) = self.shared.sender.load_full() else { - return; - }; - let meta = event.metadata(); - if meta.target() != TELEMETRY_TARGET || !at_least_info(meta.level()) { - return; - } - let mut visitor = AllowlistVisitor::default(); - event.record(&mut visitor); - let record = build_log_record(meta.level(), visitor); - let due = self.shared.push(record); - if !due.is_empty() { - sender.export(due); - } - } -} - -/// Flush the active [`DonatingLogLayer`]'s in-memory batch onto the -/// pump. Called from `ToolServer` teardown before the pump drain so a -/// crash-y shutdown does not abandon a partial batch. -pub fn flush_log_layer() { - if let Some(shared) = ACTIVE_LOG_LAYER.load_full() { - shared.flush(); - } -} - -/// Shutdown fence: drains queued log donations before the connection -/// closes. Call after [`flush_log_layer`]. -pub struct LogDonationPump { - tx: mpsc::Sender, -} - -impl LogDonationPump { - /// Resolves once every payload queued before this call has had a - /// send attempt. - pub async fn drain(&self) { - drain_via(&self.tx).await; - } -} - -impl ToolServer { - /// Post-connect entry point: spawn the log donation pump (wiring - /// [`ToolServer::donate_logs`]) and return a sender to swap into the - /// already-installed inert [`DonatingLogLayer`] plus a drain handle. - /// Does **not** return a `Layer` — a layer cannot be added to an - /// already-set global subscriber. `service_name` must be - /// server-allowlisted. - pub fn log_donation_layer( - &self, - service_name: impl Into, - ) -> (LogDonationSender, LogDonationPump) { - let (tx, rx) = mpsc::channel::(PENDING_FLUSHES); - let server = self.downgrade(); - tokio::spawn(run_pump(rx, move |payload: String| { - let server = server.clone(); - async move { - let Some(server) = server.upgrade() else { - return (false, payload); - }; - let ok = server.donate_logs(&payload).await.is_ok(); - (ok, payload) - } - })); - self.set_log_donation_pump(tx.clone()); - - let sender = LogDonationSender { - exporter: PumpLogExporter { - tx: tx.clone(), - resource: make_resource(service_name.into()), - }, - }; - (sender, LogDonationPump { tx }) - } -} - -#[cfg(test)] -mod tests { - use fastrace::collector::{SpanId, TraceId}; - use tracing_subscriber::layer::SubscriberExt; - - use super::*; - - fn test_sender(tx: mpsc::Sender) -> LogDonationSender { - LogDonationSender { - exporter: PumpLogExporter { - tx, - resource: make_resource("test-service".to_owned()), - }, - } - } - - fn decode(payload: String) -> ExportLogsServiceRequest { - let bytes = base64::engine::general_purpose::STANDARD - .decode(payload) - .expect("payload must be base64"); - ExportLogsServiceRequest::decode(bytes.as_slice()).expect("payload must be OTLP") - } - - #[test] - fn severity_maps_levels_to_otlp_numbers() { - assert_eq!(severity(&Level::ERROR), ("ERROR", 17)); - assert_eq!(severity(&Level::WARN), ("WARN", 13)); - assert_eq!(severity(&Level::INFO), ("INFO", 9)); - assert_eq!(severity(&Level::DEBUG), ("DEBUG", 5)); - assert_eq!(severity(&Level::TRACE), ("TRACE", 1)); - } - - #[test] - fn donation_filter_selects_info_and_above() { - assert!(at_least_info(&Level::ERROR)); - assert!(at_least_info(&Level::WARN)); - assert!(at_least_info(&Level::INFO)); - assert!(!at_least_info(&Level::DEBUG)); - assert!(!at_least_info(&Level::TRACE)); - } - - #[test] - fn encode_ids_is_big_endian_16_and_8_bytes() { - let ctx = SpanContext::new( - TraceId(0x0af7651916cd43dd8448eb211c80319c), - SpanId(0xb7ad6b7169203331), - ); - let (trace_id, span_id) = encode_ids(&ctx); - assert_eq!(trace_id.len(), 16); - assert_eq!(span_id.len(), 8); - assert_eq!( - format!("{:032x}", u128::from_be_bytes(trace_id.try_into().unwrap())), - "0af7651916cd43dd8448eb211c80319c" - ); - assert_eq!( - format!("{:016x}", u64::from_be_bytes(span_id.try_into().unwrap())), - "b7ad6b7169203331" - ); - } - - #[test] - fn current_ids_empty_without_local_parent() { - let (trace_id, span_id) = current_ids(); - assert!(trace_id.is_empty()); - assert!(span_id.is_empty()); - } - - #[test] - fn layer_converts_event_and_redacts_free_form_fields() { - let layer = DonatingLogLayer::new_inert(); - let (tx, mut rx) = mpsc::channel::(8); - layer.activate(test_sender(tx)); - let flusher = layer.clone(); - - let subscriber = tracing_subscriber::registry().with(layer); - tracing::subscriber::with_default(subscriber, || { - tracing::warn!( - target: "workspace::telemetry", - session_id = "s1", - turn_number = 3u64, - phase = "tool_state", - error_category = "archive_failed", - error = "secret git stderr with /home/user/path", - "archive build failed (queued path)" - ); - // Off-target event must never be forwarded. - tracing::warn!(session_id = "s2", "unrelated chatter"); - // DEBUG on-target is below the threshold. - tracing::debug!(target: "workspace::telemetry", session_id = "s3", "verbose"); - }); - flusher.shared.flush(); - - let PumpMsg::Payload(payload) = rx.try_recv().expect("one batch must be queued") else { - panic!("expected a payload"); - }; - let request = decode(payload); - let scope_logs = &request.resource_logs[0].scope_logs[0]; - assert_eq!( - scope_logs.scope.as_ref().unwrap().name, - "workspace::telemetry" - ); - assert_eq!( - scope_logs.log_records.len(), - 1, - "only the WARN on-target row" - ); - - let record = &scope_logs.log_records[0]; - assert_eq!(record.severity_text, "WARN"); - assert_eq!(record.severity_number, 13); - assert_eq!( - record.body.as_ref().unwrap().value, - Some(any_value::Value::StringValue( - "archive build failed (queued path)".to_owned() - )) - ); - let keys: Vec<&str> = record.attributes.iter().map(|kv| kv.key.as_str()).collect(); - assert!(keys.contains(&"session_id")); - assert!(keys.contains(&"turn_number")); - assert!(keys.contains(&"phase")); - assert!(keys.contains(&"error_category")); - assert!( - !keys.contains(&"error"), - "free-form `error` must be dropped, got {keys:?}" - ); - - // Resource carries the donor service.name. - let service_name = request.resource_logs[0] - .resource - .as_ref() - .unwrap() - .attributes - .iter() - .find(|kv| kv.key == "service.name") - .and_then(|kv| kv.value.as_ref()) - .and_then(|v| v.value.clone()); - assert_eq!( - service_name, - Some(any_value::Value::StringValue("test-service".to_owned())) - ); - - assert!(rx.try_recv().is_err(), "no further payloads"); - } - - #[test] - fn inert_layer_drops_selected_events() { - let layer = DonatingLogLayer::new_inert(); - let flusher = layer.clone(); - // No sender activated. - let subscriber = tracing_subscriber::registry().with(layer); - tracing::subscriber::with_default(subscriber, || { - tracing::warn!(target: "workspace::telemetry", session_id = "s1", "dropped"); - }); - flusher.shared.flush(); - // Nothing to assert beyond not panicking: with no sender the - // batch never fills and flush is a no-op. - } - - #[test] - fn exporter_chunks_at_max_records_per_donation() { - let (tx, mut rx) = mpsc::channel::(8); - let exporter = PumpLogExporter { - tx, - resource: make_resource("test-service".to_owned()), - }; - let records = vec![LogRecord::default(); MAX_LOG_RECORDS_PER_DONATION + 1]; - exporter.export(records); - - let mut total = 0; - let mut payloads = 0; - while let Ok(PumpMsg::Payload(p)) = rx.try_recv() { - payloads += 1; - total += decode(p).resource_logs[0].scope_logs[0].log_records.len(); - } - assert_eq!(payloads, 2, "one full chunk + remainder"); - assert_eq!(total, MAX_LOG_RECORDS_PER_DONATION + 1); - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/metric_donate.rs b/crates/common/kigi-computer-hub-sdk/src/metric_donate.rs deleted file mode 100644 index 2d5bcd5..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/metric_donate.rs +++ /dev/null @@ -1,396 +0,0 @@ -//! Forward the process's Prometheus metrics to the connected server over -//! the WebSocket transport (`metrics.donate`). -//! -//! A [`MetricDonationReporter`] periodically snapshots the default -//! Prometheus registry via [`prometheus::gather`], converts the -//! `MetricFamily` set to native OTLP metrics (Counter→Sum, Gauge→Gauge, -//! Histogram→Histogram, labels preserved, cumulative temporality), and -//! pumps the batch over the shared [`crate::donate_pump`]. Because it -//! gathers the whole registry, every current and future metric is -//! exported with zero per-metric wiring. Metrics are **process-aggregate** -//! — [`ToolServer::donate_metrics`] requires no bound session. - -use std::sync::{Arc, LazyLock}; -use std::time::Duration; - -use arc_swap::ArcSwapOption; -use base64::Engine as _; -use kigi_tool_protocol::{MAX_DONATION_BYTES, MAX_METRICS_PER_DONATION}; -use opentelemetry_proto::tonic::collector::metrics::v1::ExportMetricsServiceRequest; -use opentelemetry_proto::tonic::common::v1::KeyValue; -use opentelemetry_proto::tonic::metrics::v1::{ - AggregationTemporality, Gauge, Histogram, HistogramDataPoint, Metric, NumberDataPoint, - ResourceMetrics, ScopeMetrics, Sum, metric, number_data_point, -}; -use opentelemetry_proto::tonic::resource::v1::Resource; -use prometheus::proto::{MetricFamily, MetricType}; -use prost::Message as _; -use tokio::sync::mpsc; -use tokio_util::sync::CancellationToken; - -use crate::donate_pump::{ - PENDING_FLUSHES, PumpMsg, make_resource, now_unix_nanos, run_pump, string_kv, -}; -use crate::server::ToolServer; - -/// How often the reporter snapshots the registry. The server re-stamps -/// attribution; cumulative temporality means missed ticks only delay -/// freshness, never lose monotonic state. -const DEFAULT_GATHER_INTERVAL: Duration = Duration::from_secs(60); - -fn labels_to_kv(labels: &[prometheus::proto::LabelPair]) -> Vec { - labels - .iter() - .map(|l| string_kv(l.name(), l.value().to_owned())) - .collect() -} - -fn number_point(metric: &prometheus::proto::Metric, value: f64, now: u64) -> NumberDataPoint { - NumberDataPoint { - attributes: labels_to_kv(metric.get_label()), - time_unix_nano: now, - value: Some(number_data_point::Value::AsDouble(value)), - ..Default::default() - } -} - -/// Prometheus histogram buckets are **cumulative** (`le` counts); OTLP -/// wants per-bucket counts plus an implicit `+Inf` bucket, so the -/// cumulative counts are differenced here. -fn histogram_point(metric: &prometheus::proto::Metric, now: u64) -> HistogramDataPoint { - let hist = metric.get_histogram(); - let mut bucket_counts = Vec::new(); - let mut explicit_bounds = Vec::new(); - let mut prev = 0u64; - for bucket in hist.get_bucket() { - let cumulative = bucket.cumulative_count(); - bucket_counts.push(cumulative.saturating_sub(prev)); - explicit_bounds.push(bucket.upper_bound()); - prev = cumulative; - } - let total = hist.get_sample_count(); - bucket_counts.push(total.saturating_sub(prev)); - HistogramDataPoint { - attributes: labels_to_kv(metric.get_label()), - time_unix_nano: now, - count: total, - sum: Some(hist.get_sample_sum()), - bucket_counts, - explicit_bounds, - ..Default::default() - } -} - -/// Convert a gathered `MetricFamily` set to OTLP metrics. Summaries and -/// untyped families are skipped defensively (none registered today). -fn convert_families(families: &[MetricFamily]) -> Vec { - let now = now_unix_nanos(); - let cumulative = AggregationTemporality::Cumulative as i32; - let mut out = Vec::new(); - for family in families { - let name = family.name().to_owned(); - let data = match family.get_field_type() { - MetricType::COUNTER => metric::Data::Sum(Sum { - data_points: family - .get_metric() - .iter() - .map(|m| number_point(m, m.get_counter().value(), now)) - .collect(), - aggregation_temporality: cumulative, - is_monotonic: true, - }), - MetricType::GAUGE => metric::Data::Gauge(Gauge { - data_points: family - .get_metric() - .iter() - .map(|m| number_point(m, m.get_gauge().value(), now)) - .collect(), - }), - MetricType::HISTOGRAM => metric::Data::Histogram(Histogram { - data_points: family - .get_metric() - .iter() - .map(|m| histogram_point(m, now)) - .collect(), - aggregation_temporality: cumulative, - }), - MetricType::SUMMARY | MetricType::UNTYPED => continue, - }; - out.push(Metric { - name, - data: Some(data), - ..Default::default() - }); - } - out -} - -/// Encodes batches of OTLP metrics onto the pump channel. Chunks at -/// [`MAX_METRICS_PER_DONATION`], drops payloads over -/// [`MAX_DONATION_BYTES`], and never blocks. -#[derive(Clone)] -struct MetricExporter { - tx: mpsc::Sender, - resource: Resource, -} - -impl MetricExporter { - fn export(&self, mut metrics: Vec) { - while !metrics.is_empty() { - let chunk = if metrics.len() > MAX_METRICS_PER_DONATION { - let rest = metrics.split_off(MAX_METRICS_PER_DONATION); - std::mem::replace(&mut metrics, rest) - } else { - std::mem::take(&mut metrics) - }; - let request = ExportMetricsServiceRequest { - resource_metrics: vec![ResourceMetrics { - resource: Some(self.resource.clone()), - scope_metrics: vec![ScopeMetrics { - metrics: chunk, - ..Default::default() - }], - schema_url: String::new(), - }], - }; - let bytes = request.encode_to_vec(); - if bytes.len() > MAX_DONATION_BYTES { - tracing::debug!( - len = bytes.len(), - "dropping oversized metric donation payload" - ); - continue; - } - let payload = base64::engine::general_purpose::STANDARD.encode(bytes); - if self.tx.try_send(PumpMsg::Payload(payload)).is_err() { - tracing::debug!("metric donation queue full; dropping metric batch"); - } - } - } - - fn gather_and_send(&self) { - let metrics = convert_families(&prometheus::gather()); - if metrics.is_empty() { - return; - } - self.export(metrics); - } -} - -/// Process-global handle to the active exporter so [`gather_and_send`] -/// can drive a final teardown gather without a reference. -static ACTIVE_METRIC_EXPORTER: LazyLock> = - LazyLock::new(ArcSwapOption::empty); - -/// Final registry gather onto the active metric pump. Called from -/// `ToolServer` teardown before the pump drain so a crash-y shutdown -/// captures the latest values. -pub(crate) fn gather_and_send() { - if let Some(exporter) = ACTIVE_METRIC_EXPORTER.load_full() { - exporter.gather_and_send(); - } -} - -/// Drop the process-global exporter on teardown so its pump `Sender` is released -/// and the metric pump can wind down. Called from `flush_donations_inner` after -/// the final [`gather_and_send`] (and alongside clearing the stored pump -/// senders), so a dropped `ToolServer` doesn't leak the pump task. -pub(crate) fn clear_active_exporter() { - ACTIVE_METRIC_EXPORTER.store(None); -} - -/// Periodic registry gatherer spawned by -/// [`ToolServer::metric_donation_reporter`]. Internal: constructed and run -/// only by `metric_donation_reporter`; not part of the crate's public API. -pub(crate) struct MetricDonationReporter { - exporter: MetricExporter, - interval: Duration, - shutdown: CancellationToken, -} - -impl MetricDonationReporter { - async fn run(self) { - let mut ticker = tokio::time::interval(self.interval); - loop { - tokio::select! { - _ = ticker.tick() => self.exporter.gather_and_send(), - // Stop on teardown so this task (and the pump `tx` clone it - // holds) doesn't outlive `ToolServer::shutdown` and keep - // gathering/sending forever. - _ = self.shutdown.cancelled() => break, - } - } - } -} - -/// Shutdown fence: drains queued metric donations before the connection -/// closes. -pub struct MetricDonationPump { - tx: mpsc::Sender, -} - -impl MetricDonationPump { - /// Resolves once every payload queued before this call has had a - /// send attempt. - pub async fn drain(&self) { - crate::donate_pump::drain_via(&self.tx).await; - } -} - -impl ToolServer { - /// Post-connect entry point: spawn the metric donation pump (wiring - /// [`ToolServer::donate_metrics`]) plus the periodic registry - /// gatherer, and return a drain handle. Activates on server presence - /// (no env flag). `service_name` must be server-allowlisted. - pub fn metric_donation_reporter(&self, service_name: impl Into) -> MetricDonationPump { - let (tx, rx) = mpsc::channel::(PENDING_FLUSHES); - let server = self.downgrade(); - tokio::spawn(run_pump(rx, move |payload: String| { - let server = server.clone(); - async move { - let Some(server) = server.upgrade() else { - return (false, payload); - }; - let ok = server.donate_metrics(&payload).await.is_ok(); - (ok, payload) - } - })); - self.set_metric_donation_pump(tx.clone()); - - let exporter = MetricExporter { - tx: tx.clone(), - resource: make_resource(service_name.into()), - }; - ACTIVE_METRIC_EXPORTER.store(Some(Arc::new(exporter.clone()))); - tokio::spawn( - MetricDonationReporter { - exporter, - interval: DEFAULT_GATHER_INTERVAL, - shutdown: self.shutdown_token(), - } - .run(), - ); - - MetricDonationPump { tx } - } -} - -#[cfg(test)] -mod tests { - use opentelemetry_proto::tonic::common::v1::any_value; - use prometheus::{Histogram, HistogramOpts, IntCounterVec, IntGauge, Opts, Registry}; - - use super::*; - - fn decode(payload: String) -> ExportMetricsServiceRequest { - let bytes = base64::engine::general_purpose::STANDARD - .decode(payload) - .expect("payload must be base64"); - ExportMetricsServiceRequest::decode(bytes.as_slice()).expect("payload must be OTLP") - } - - fn label_map(attrs: &[KeyValue]) -> std::collections::HashMap { - attrs - .iter() - .filter_map(|kv| match kv.value.as_ref().and_then(|v| v.value.clone()) { - Some(any_value::Value::StringValue(s)) => Some((kv.key.clone(), s)), - _ => None, - }) - .collect() - } - - #[test] - fn converts_counter_gauge_histogram_with_labels() { - let registry = Registry::new(); - - let counter = - IntCounterVec::new(Opts::new("grok_test_total", "help"), &["reason"]).unwrap(); - registry.register(Box::new(counter.clone())).unwrap(); - counter.with_label_values(&["zdr"]).inc_by(5); - - let gauge = IntGauge::new("grok_test_pending", "help").unwrap(); - registry.register(Box::new(gauge.clone())).unwrap(); - gauge.set(7); - - let hist = Histogram::with_opts( - HistogramOpts::new("grok_test_seconds", "help").buckets(vec![0.5, 1.0]), - ) - .unwrap(); - registry.register(Box::new(hist.clone())).unwrap(); - hist.observe(0.25); - hist.observe(0.75); - hist.observe(5.0); - - let metrics = convert_families(®istry.gather()); - let by_name: std::collections::HashMap<_, _> = - metrics.iter().map(|m| (m.name.clone(), m)).collect(); - - // Counter -> Sum (monotonic, cumulative), label preserved. - let metric::Data::Sum(sum) = by_name["grok_test_total"].data.as_ref().unwrap() else { - panic!("counter must convert to Sum"); - }; - assert!(sum.is_monotonic); - assert_eq!( - sum.aggregation_temporality, - AggregationTemporality::Cumulative as i32 - ); - let dp = &sum.data_points[0]; - assert_eq!(dp.value, Some(number_data_point::Value::AsDouble(5.0))); - assert_eq!( - label_map(&dp.attributes).get("reason").map(String::as_str), - Some("zdr") - ); - - // Gauge -> Gauge. - let metric::Data::Gauge(g) = by_name["grok_test_pending"].data.as_ref().unwrap() else { - panic!("gauge must convert to Gauge"); - }; - assert_eq!( - g.data_points[0].value, - Some(number_data_point::Value::AsDouble(7.0)) - ); - - // Histogram -> Histogram with cumulative buckets differenced and - // a +Inf bucket appended. - let metric::Data::Histogram(h) = by_name["grok_test_seconds"].data.as_ref().unwrap() else { - panic!("histogram must convert to Histogram"); - }; - assert_eq!( - h.aggregation_temporality, - AggregationTemporality::Cumulative as i32 - ); - let hdp = &h.data_points[0]; - assert_eq!(hdp.count, 3); - assert_eq!(hdp.sum, Some(6.0)); - assert_eq!(hdp.explicit_bounds, vec![0.5, 1.0]); - // (<=0.5): 0.25 -> 1 ; (0.5,1.0]: 0.75 -> 1 ; (+Inf): 5.0 -> 1 - assert_eq!(hdp.bucket_counts, vec![1, 1, 1]); - } - - #[test] - fn exporter_chunks_at_max_metrics_per_donation() { - let (tx, mut rx) = mpsc::channel::(8); - let exporter = MetricExporter { - tx, - resource: make_resource("test-service".to_owned()), - }; - let metrics = vec![Metric::default(); MAX_METRICS_PER_DONATION + 1]; - exporter.export(metrics); - - let mut payloads = 0; - let mut total = 0; - while let Ok(PumpMsg::Payload(p)) = rx.try_recv() { - payloads += 1; - total += decode(p).resource_metrics[0].scope_metrics[0].metrics.len(); - } - assert_eq!(payloads, 2, "one full chunk + remainder"); - assert_eq!(total, MAX_METRICS_PER_DONATION + 1); - } - - #[test] - fn summary_and_untyped_families_are_skipped() { - // An empty registry gathers nothing; convert yields nothing. - let registry = Registry::new(); - assert!(convert_families(®istry.gather()).is_empty()); - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/metrics.rs b/crates/common/kigi-computer-hub-sdk/src/metrics.rs deleted file mode 100644 index 125eada..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/metrics.rs +++ /dev/null @@ -1,552 +0,0 @@ -//! Feature-gated Prometheus metrics for the SDK. -//! -//! When the `metrics` cargo feature is enabled, each helper records to a -//! lazily-registered Prometheus counter / gauge / histogram. When -//! disabled (the default), every helper compiles to an empty function -//! body so the SDK carries zero prometheus dependency. - -#[cfg(feature = "metrics")] -mod inner { - use prometheus::{ - Histogram, HistogramVec, IntCounter, IntCounterVec, IntGauge, IntGaugeVec, - exponential_buckets, register_histogram, register_histogram_vec, register_int_counter, - register_int_counter_vec, register_int_gauge, register_int_gauge_vec, - }; - use std::sync::LazyLock; - - static POOL_CONNECTIONS: LazyLock = LazyLock::new(|| { - register_int_gauge!( - "computer_hub_client_pool_connections", - "Active pooled connections in the SDK connection pool." - ) - .expect("computer_hub_client_pool_connections must register once") - }); - - static POOL_EVICTIONS_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "computer_hub_client_pool_evictions_total", - "Pooled connections closed by the idle reaper (unused past the idle TTL)." - ) - .expect("computer_hub_client_pool_evictions_total must register once") - }); - - static RECONNECTS_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "computer_hub_client_reconnects_total", - "Cumulative reconnect attempts that succeeded." - ) - .expect("computer_hub_client_reconnects_total must register once") - }); - - static RECONNECT_FAILED_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter_vec!( - "computer_hub_client_reconnect_failed_total", - "Cumulative reconnect attempts that failed, by reason \ - (handshake_auth = fatal 401/403, transport = retryable).", - &["reason"] - ) - .expect("computer_hub_client_reconnect_failed_total must register once") - }); - - static RECONNECT_DURATION_SECONDS: LazyLock = LazyLock::new(|| { - register_histogram!( - "computer_hub_client_reconnect_duration_seconds", - "Time to complete a reconnect cycle (handshake + session/tool replay).", - exponential_buckets(0.01, 2.0, 14).expect("valid bucket params") - ) - .expect("computer_hub_client_reconnect_duration_seconds must register once") - }); - - static RECONNECTS_BY_CAUSE_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter_vec!( - "computer_hub_client_reconnects_by_cause_total", - "Successful reconnects by disconnect cause of the previous connection \ - (close_frame, eof, transport_read_error, transport_write_error, forced). \ - Cause-labeled companion to computer_hub_client_reconnects_total.", - &["cause"] - ) - .expect("computer_hub_client_reconnects_by_cause_total must register once") - }); - - static RECONNECT_GAP_SECONDS: LazyLock = LazyLock::new(|| { - register_histogram!( - "computer_hub_client_reconnect_gap_seconds", - "Time from the last inbound frame on the dead connection to a successful reconnect.", - exponential_buckets(0.1, 2.0, 14).expect("valid bucket params") - ) - .expect("computer_hub_client_reconnect_gap_seconds must register once") - }); - - static CALL_DISPATCH_SECONDS: LazyLock = LazyLock::new(|| { - register_histogram!( - "computer_hub_client_call_dispatch_seconds", - "Time to set up and queue the outbound remote dispatch.", - exponential_buckets(0.0001, 2.0, 14).expect("valid bucket params") - ) - .expect("computer_hub_client_call_dispatch_seconds must register once") - }); - - static DEMUX_INBOX_DEPTH: LazyLock = LazyLock::new(|| { - register_int_gauge!( - "computer_hub_client_demux_inbox_depth", - "Number of session inboxes registered in the inbound demux." - ) - .expect("computer_hub_client_demux_inbox_depth must register once") - }); - - static CALL_ID_COLLISIONS_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "computer_hub_client_call_id_collisions_total", - "Call-id collisions detected in the harness dispatch path." - ) - .expect("computer_hub_client_call_id_collisions_total must register once") - }); - - // ── Server integration metrics ────────────────────────────────── - - static HARNESS_CONNECT_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter_vec!( - "hub_harness_connect_total", - "Hub connection attempts by outcome and sampler.", - &["status", "sampler"] - ) - .expect("hub_harness_connect_total must register once") - }); - - static SESSION_EVENT_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter_vec!( - "hub_session_event_total", - "SessionEvent emissions by event type.", - &["event_type"] - ) - .expect("hub_session_event_total must register once") - }); - - static SESSION_OP_DURATION_SECONDS: LazyLock = LazyLock::new(|| { - register_histogram_vec!( - "hub_session_op_duration_seconds", - "Latency of hub session lifecycle operations (open/bind) by op and outcome.", - &["op", "status"], - exponential_buckets(0.001, 2.0, 14).expect("valid bucket params") - ) - .expect("hub_session_op_duration_seconds must register once") - }); - - static SESSION_SOFT_REBIND_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "hub_session_soft_rebind_total", - "Redundant session.bind frames for a session with a live dispatch loop, \ - handled as a non-destructive soft rebind (serve state refreshed, \ - in-flight calls preserved)." - ) - .expect("hub_session_soft_rebind_total must register once") - }); - - static NO_HANDLER_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "hub_sdk_no_handler_total", - "tool_call_request frames rejected with -32011 because the session's \ - current handler set has no handler for the requested tool_id." - ) - .expect("hub_sdk_no_handler_total must register once") - }); - - static HOOK_SEND_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter_vec!( - "hub_hook_send_total", - "Hook sends by hook type.", - &["hook_type"] - ) - .expect("hub_hook_send_total must register once") - }); - - static PROGRESS_FRAMES_FORWARDED_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "hub_progress_frames_forwarded_total", - "Progress frames forwarded by ToolServer." - ) - .expect("hub_progress_frames_forwarded_total must register once") - }); - - static CANCEL_HOOK_RECEIVED_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "hub_cancel_hook_received_total", - "Cancel hooks received by workspace tool server." - ) - .expect("hub_cancel_hook_received_total must register once") - }); - - static WRITER_SINK_SEND_ERRORS_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "computer_hub_client_writer_sink_send_errors_total", - "Writer-task sink send failures; each signals the reader to reconnect." - ) - .expect("computer_hub_client_writer_sink_send_errors_total must register once") - }); - - static RECONNECT_WRITER_RESUME_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "computer_hub_client_reconnect_writer_resume_total", - "Fresh-sink Resume handoffs delivered to the writer task after reconnect." - ) - .expect("computer_hub_client_reconnect_writer_resume_total must register once") - }); - - static LIVENESS_DEADLINE_EXPIRED_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "computer_hub_client_liveness_deadline_expired_total", - "Liveness-deadline expiries in the reader (no inbound WebSocket \ - frame within the deadline); each declares the connection dead and \ - drives the normal reconnect path." - ) - .expect("computer_hub_client_liveness_deadline_expired_total must register once") - }); - - static HEARTBEAT_PONG_DROPPED_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "computer_hub_client_heartbeat_pong_dropped_total", - "App-level heartbeat pongs dropped because outbound_tx was saturated (split reader)." - ) - .expect("computer_hub_client_heartbeat_pong_dropped_total must register once") - }); - - static CANCEL_APPLIED_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "hub_cancel_applied_total", - "Cancel hooks that hit a live in-flight call and cancelled it." - ) - .expect("hub_cancel_applied_total must register once") - }); - - static CANCEL_PENDING_TOMBSTONED_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "hub_cancel_pending_tombstoned_total", - "Cancel hooks recorded as a pending tombstone (call not yet registered or already done)." - ) - .expect("hub_cancel_pending_tombstoned_total must register once") - }); - - static CANCEL_NO_TARGET_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "hub_cancel_no_target_total", - "Cancel hooks with no call_id (session-wide, no specific call to cancel)." - ) - .expect("hub_cancel_no_target_total must register once") - }); - - static TOOL_CALL_REJECTED_OVERLOADED_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "hub_tool_call_rejected_overloaded_total", - "Tool calls rejected by admission timeout (-32016 tool_busy)." - ) - .expect("hub_tool_call_rejected_overloaded_total must register once") - }); - - static INBOX_FULL_REQUEST_REJECTED_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "hub_inbox_full_request_rejected_total", - "Requests rejected with an overloaded response on a full session inbox." - ) - .expect("hub_inbox_full_request_rejected_total must register once") - }); - - static INBOX_FULL_REJECT_SEND_FAILED_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "hub_inbox_full_reject_send_failed_total", - "Overloaded rejections dropped because outbound was also full (residual silent loss)." - ) - .expect("hub_inbox_full_reject_send_failed_total must register once") - }); - - static INBOX_FULL_NOTIFICATION_DROPPED_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "hub_inbox_full_notification_dropped_total", - "Notifications (no id) dropped on a full session inbox." - ) - .expect("hub_inbox_full_notification_dropped_total must register once") - }); - - static SERVE_REPLAY_TIMEOUT_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "computer_hub_client_serve_replay_timeout_total", - "serve attempts that hit the per-attempt reply deadline, from any \ - serve call site (reconnect replay, run(), bind, tool updates)." - ) - .expect("computer_hub_client_serve_replay_timeout_total must register once") - }); - - static NOTIF_LAGGED_RECOVERED_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "hub_notif_lagged_recovered_total", - "Connection-notification broadcast Lagged events recovered by \ - continuing the loop instead of exiting." - ) - .expect("hub_notif_lagged_recovered_total must register once") - }); - - static EARLY_NOTIF_BUFFERED_TOTAL: LazyLock = LazyLock::new(|| { - register_int_counter!( - "hub_early_notif_buffered_total", - "Connection-level notification frames (binds, unbinds, evicts, \ - ...) buffered between connect and ToolServer::run() and replayed \ - instead of dropped." - ) - .expect("hub_early_notif_buffered_total must register once") - }); - - static TOOL_CALL_INFLIGHT: LazyLock = LazyLock::new(|| { - register_int_gauge_vec!( - "hub_tool_call_inflight", - "Concurrent running tool calls holding an admission permit, by scope.", - &["scope"] - ) - .expect("hub_tool_call_inflight must register once") - }); - - static ADMISSION_WAIT_SECONDS: LazyLock = LazyLock::new(|| { - register_histogram!( - "hub_tool_call_admission_wait_seconds", - "Time blocked acquiring the three admission permits (one shared deadline).", - exponential_buckets(0.0001, 2.0, 16).expect("valid bucket params") - ) - .expect("hub_tool_call_admission_wait_seconds must register once") - }); - - pub(crate) fn pool_connections_inc() { - POOL_CONNECTIONS.inc(); - } - - pub(crate) fn pool_connections_dec() { - POOL_CONNECTIONS.dec(); - } - - pub(crate) fn pool_evictions_inc() { - POOL_EVICTIONS_TOTAL.inc(); - } - - pub(crate) fn reconnect_succeeded() { - RECONNECTS_TOTAL.inc(); - } - - pub(crate) fn reconnect_failed(reason: &str) { - RECONNECT_FAILED_TOTAL.with_label_values(&[reason]).inc(); - } - - pub(crate) fn reconnect_duration_observe(secs: f64) { - RECONNECT_DURATION_SECONDS.observe(secs); - } - - pub(crate) fn reconnect_cause(cause: &str) { - RECONNECTS_BY_CAUSE_TOTAL.with_label_values(&[cause]).inc(); - } - - pub(crate) fn reconnect_gap_observe(secs: f64) { - RECONNECT_GAP_SECONDS.observe(secs); - } - - pub(crate) fn call_dispatch_observe(secs: f64) { - CALL_DISPATCH_SECONDS.observe(secs); - } - - pub(crate) fn demux_inbox_depth_set(depth: i64) { - DEMUX_INBOX_DEPTH.set(depth); - } - - pub(crate) fn call_id_collision() { - CALL_ID_COLLISIONS_TOTAL.inc(); - } - - /// Record a harness connection attempt. - /// - /// `sampler` identifies the caller (`"chat"` or `"shell"`). - /// `status` is `"ok"`, `"error"`, or `"fallback"` (fallback is - /// emitted by the caller in `AgentBuilder::build_harness()`, not - /// by the SDK). - pub fn harness_connect(status: &str, sampler: &str) { - HARNESS_CONNECT_TOTAL - .with_label_values(&[status, sampler]) - .inc(); - } - - pub(crate) fn session_event(event_type: &str) { - SESSION_EVENT_TOTAL.with_label_values(&[event_type]).inc(); - } - - /// Observe the latency of a session lifecycle operation. - /// `op` is `"open"` or `"bind"`; `status` is `"ok"` or `"error"`. - pub(crate) fn session_op_observe(op: &str, status: &str, secs: f64) { - SESSION_OP_DURATION_SECONDS - .with_label_values(&[op, status]) - .observe(secs); - } - - pub(crate) fn session_soft_rebind() { - SESSION_SOFT_REBIND_TOTAL.inc(); - } - - pub(crate) fn no_handler() { - NO_HANDLER_TOTAL.inc(); - } - - pub(crate) fn hook_send(hook_type: &str) { - HOOK_SEND_TOTAL.with_label_values(&[hook_type]).inc(); - } - - pub(crate) fn progress_frame_forwarded() { - PROGRESS_FRAMES_FORWARDED_TOTAL.inc(); - } - - pub(crate) fn cancel_hook_received() { - CANCEL_HOOK_RECEIVED_TOTAL.inc(); - } - - pub(crate) fn writer_sink_send_error() { - WRITER_SINK_SEND_ERRORS_TOTAL.inc(); - } - - pub(crate) fn reconnect_writer_resume() { - RECONNECT_WRITER_RESUME_TOTAL.inc(); - } - - pub(crate) fn liveness_deadline_expired() { - LIVENESS_DEADLINE_EXPIRED_TOTAL.inc(); - } - - pub(crate) fn heartbeat_pong_dropped() { - HEARTBEAT_PONG_DROPPED_TOTAL.inc(); - } - - pub(crate) fn cancel_applied() { - CANCEL_APPLIED_TOTAL.inc(); - } - - pub(crate) fn cancel_pending_tombstoned() { - CANCEL_PENDING_TOMBSTONED_TOTAL.inc(); - } - - pub(crate) fn cancel_no_target() { - CANCEL_NO_TARGET_TOTAL.inc(); - } - - pub(crate) fn tool_call_rejected_overloaded() { - TOOL_CALL_REJECTED_OVERLOADED_TOTAL.inc(); - } - - pub(crate) fn inbox_full_request_rejected() { - INBOX_FULL_REQUEST_REJECTED_TOTAL.inc(); - } - - pub(crate) fn inbox_full_reject_send_failed() { - INBOX_FULL_REJECT_SEND_FAILED_TOTAL.inc(); - } - - pub(crate) fn inbox_full_notification_dropped() { - INBOX_FULL_NOTIFICATION_DROPPED_TOTAL.inc(); - } - - pub(crate) fn serve_replay_timeout() { - SERVE_REPLAY_TIMEOUT_TOTAL.inc(); - } - - pub(crate) fn notif_lagged_recovered() { - NOTIF_LAGGED_RECOVERED_TOTAL.inc(); - } - - pub(crate) fn early_notif_buffered(frames: u64) { - EARLY_NOTIF_BUFFERED_TOTAL.inc_by(frames); - } - - pub(crate) fn tool_call_inflight_inc(scope: &str) { - TOOL_CALL_INFLIGHT.with_label_values(&[scope]).inc(); - } - - pub(crate) fn tool_call_inflight_dec(scope: &str) { - TOOL_CALL_INFLIGHT.with_label_values(&[scope]).dec(); - } - - pub(crate) fn admission_wait_observe(secs: f64) { - ADMISSION_WAIT_SECONDS.observe(secs); - } -} - -#[cfg(not(feature = "metrics"))] -mod inner { - pub(crate) fn pool_connections_inc() {} - pub(crate) fn pool_connections_dec() {} - pub(crate) fn pool_evictions_inc() {} - pub(crate) fn reconnect_succeeded() {} - pub(crate) fn reconnect_failed(_reason: &str) {} - pub(crate) fn reconnect_duration_observe(_secs: f64) {} - pub(crate) fn reconnect_cause(_cause: &str) {} - pub(crate) fn reconnect_gap_observe(_secs: f64) {} - pub(crate) fn call_dispatch_observe(_secs: f64) {} - pub(crate) fn demux_inbox_depth_set(_depth: i64) {} - pub(crate) fn call_id_collision() {} - pub fn harness_connect(_status: &str, _sampler: &str) {} - pub(crate) fn session_event(_event_type: &str) {} - pub(crate) fn session_op_observe(_op: &str, _status: &str, _secs: f64) {} - pub(crate) fn session_soft_rebind() {} - pub(crate) fn no_handler() {} - pub(crate) fn hook_send(_hook_type: &str) {} - pub(crate) fn progress_frame_forwarded() {} - pub(crate) fn cancel_hook_received() {} - pub(crate) fn writer_sink_send_error() {} - pub(crate) fn reconnect_writer_resume() {} - pub(crate) fn liveness_deadline_expired() {} - pub(crate) fn heartbeat_pong_dropped() {} - pub(crate) fn cancel_applied() {} - pub(crate) fn cancel_pending_tombstoned() {} - pub(crate) fn cancel_no_target() {} - pub(crate) fn tool_call_rejected_overloaded() {} - pub(crate) fn inbox_full_request_rejected() {} - pub(crate) fn inbox_full_reject_send_failed() {} - pub(crate) fn inbox_full_notification_dropped() {} - pub(crate) fn serve_replay_timeout() {} - pub(crate) fn notif_lagged_recovered() {} - pub(crate) fn early_notif_buffered(_frames: u64) {} - pub(crate) fn tool_call_inflight_inc(_scope: &str) {} - pub(crate) fn tool_call_inflight_dec(_scope: &str) {} - pub(crate) fn admission_wait_observe(_secs: f64) {} -} - -pub(crate) use inner::admission_wait_observe; -pub(crate) use inner::call_dispatch_observe; -pub(crate) use inner::call_id_collision; -pub(crate) use inner::cancel_applied; -pub(crate) use inner::cancel_hook_received; -pub(crate) use inner::cancel_no_target; -pub(crate) use inner::cancel_pending_tombstoned; -pub(crate) use inner::demux_inbox_depth_set; -pub(crate) use inner::early_notif_buffered; -pub(crate) use inner::heartbeat_pong_dropped; -pub(crate) use inner::hook_send; -pub(crate) use inner::inbox_full_notification_dropped; -pub(crate) use inner::inbox_full_reject_send_failed; -pub(crate) use inner::inbox_full_request_rejected; -pub(crate) use inner::liveness_deadline_expired; -pub(crate) use inner::no_handler; -pub(crate) use inner::notif_lagged_recovered; -pub(crate) use inner::pool_connections_dec; -pub(crate) use inner::pool_connections_inc; -pub(crate) use inner::pool_evictions_inc; -pub(crate) use inner::progress_frame_forwarded; -pub(crate) use inner::reconnect_cause; -pub(crate) use inner::reconnect_duration_observe; -pub(crate) use inner::reconnect_failed; -pub(crate) use inner::reconnect_gap_observe; -pub(crate) use inner::reconnect_succeeded; -pub(crate) use inner::reconnect_writer_resume; -pub(crate) use inner::serve_replay_timeout; -pub(crate) use inner::session_event; -pub(crate) use inner::session_op_observe; -pub(crate) use inner::session_soft_rebind; -pub(crate) use inner::tool_call_inflight_dec; -pub(crate) use inner::tool_call_inflight_inc; -pub(crate) use inner::tool_call_rejected_overloaded; -pub(crate) use inner::writer_sink_send_error; - -/// Record a harness connection attempt. Public so callers outside -/// the SDK (e.g. `AgentBuilder::build_harness()` in the agentic sampler) -/// can emit `status="fallback"` when the server connection fails and the -/// builder falls back to a local-only harness. -pub use inner::harness_connect; diff --git a/crates/common/kigi-computer-hub-sdk/src/notification.rs b/crates/common/kigi-computer-hub-sdk/src/notification.rs deleted file mode 100644 index 7f6f5ad..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/notification.rs +++ /dev/null @@ -1,362 +0,0 @@ -//! Parsed server notification events. -//! -//! [`HubNotification`] is the typed representation of server-pushed -//! notification frames that arrive on a session inbox. The -//! [`HubNotification::parse`] constructor classifies a raw JSON value -//! by its `method` field and deserializes the known shapes; anything -//! unrecognised lands in [`HubNotification::Unknown`] so callers never -//! lose data. - -use kigi_tool_protocol::{ - SessionId, ToolId, ToolNotificationFrame, ToolServerStatusPayload, ToolsChanged, -}; -use serde_json::Value; -use tracing::warn; - -/// A typed server notification event parsed from a raw JSON-RPC notification frame. -#[derive(Debug, Clone, PartialEq)] -pub enum HubNotification { - /// The active tool set for a session changed (tools added, removed, or updated). - ToolsChanged { - session_id: SessionId, - added: Vec, - removed: Vec, - updated: Vec, - }, - /// A tool notification forwarded by the server to all subscribers. - ToolNotification { - session_id: SessionId, - frame: ToolNotificationFrame, - }, - /// Tool server lifecycle status change, extracted from - /// `__tool_server_status` / `status_changed` notification frames. - ToolServerStatusChanged { - session_id: SessionId, - status: ToolServerStatusPayload, - }, - /// A notification whose `method` is not recognised by this SDK version. - Unknown { method: String, params: Value }, -} - -impl HubNotification { - /// Parse a raw JSON-RPC notification into a typed [`HubNotification`]. - /// - /// Returns `None` when the value lacks a `method` field (i.e. it is - /// not a notification at all). - pub fn parse(value: &Value) -> Option { - let method = value.get("method")?.as_str()?; - let params = value - .get("params") - .cloned() - .unwrap_or(Value::Object(Default::default())); - - match method { - // `ToolsChanged` carries `session_id` inside `params`. - "tools_changed" => match serde_json::from_value::(params.clone()) { - Ok(tc) => Some(HubNotification::ToolsChanged { - session_id: tc.session_id, - added: tc.added, - removed: tc.removed, - updated: tc.updated, - }), - Err(err) => { - warn!(%err, "tools_changed params failed to deserialize; falling back to Unknown"); - Some(HubNotification::Unknown { - method: method.to_owned(), - params, - }) - } - }, - // `ToolNotificationFrame` has no `session_id`; use the envelope field. - "tool.notification" => { - let frame_result = serde_json::from_value::(params.clone()); - let session_id = value - .get("session_id") - .and_then(Value::as_str) - .and_then(|s| SessionId::new(s).ok()); - match (frame_result, session_id) { - (Ok(frame), Some(session_id)) => { - if frame - .tool_id - .as_ref() - .is_some_and(|id| id.as_str() == "__tool_server_status") - && let kigi_tool_protocol::notification_wire::WireToolNotification::Custom(ref c) = frame.notification - && c.kind == "status_changed" - { - match serde_json::from_value::( - c.payload.clone(), - ) { - Ok(status) => { - return Some(HubNotification::ToolServerStatusChanged { - session_id, - status, - }); - } - Err(err) => { - warn!(%err, "tool_server status payload failed to deserialize"); - } - } - } - Some(HubNotification::ToolNotification { session_id, frame }) - } - (Err(err), _) => { - warn!(%err, "tool.notification params failed to deserialize; falling back to Unknown"); - Some(HubNotification::Unknown { - method: method.to_owned(), - params, - }) - } - (_, None) => { - warn!( - "tool.notification missing or invalid session_id; falling back to Unknown" - ); - Some(HubNotification::Unknown { - method: method.to_owned(), - params, - }) - } - } - } - _ => Some(HubNotification::Unknown { - method: method.to_owned(), - params, - }), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - - #[test] - fn parse_tools_changed() { - let value = json!({ - "jsonrpc": "2.0", - "session_id": "s1", - "method": "tools_changed", - "params": { - "session_id": "s1", - "added": ["echo", "add"], - "removed": [], - } - }); - let notif = HubNotification::parse(&value).expect("should parse"); - match notif { - HubNotification::ToolsChanged { - session_id, - added, - removed, - updated, - } => { - assert_eq!(session_id.as_str(), "s1"); - assert_eq!(added.len(), 2); - assert!(removed.is_empty()); - assert!(updated.is_empty()); - } - other => panic!("expected ToolsChanged, got {other:?}"), - } - } - - #[test] - fn parse_tools_changed_with_updated() { - let value = json!({ - "jsonrpc": "2.0", - "session_id": "s1", - "method": "tools_changed", - "params": { - "session_id": "s1", - "added": ["new_tool"], - "removed": ["old_tool"], - "updated": ["echo", "add"], - } - }); - let notif = HubNotification::parse(&value).expect("should parse"); - match notif { - HubNotification::ToolsChanged { - session_id, - added, - removed, - updated, - } => { - assert_eq!(session_id.as_str(), "s1"); - assert_eq!(added.len(), 1); - assert_eq!(removed.len(), 1); - assert_eq!(updated.len(), 2); - assert_eq!(updated[0].as_str(), "echo"); - assert_eq!(updated[1].as_str(), "add"); - } - other => panic!("expected ToolsChanged, got {other:?}"), - } - } - - #[test] - fn parse_tool_notification_custom() { - let value = json!({ - "jsonrpc": "2.0", - "session_id": "s1", - "method": "tool.notification", - "params": { - "tool_id": "echo", - "notification": { - "shape": "custom", - "value": { - "kind": "echo.status", - "payload": { "status": "idle" } - } - } - } - }); - let notif = HubNotification::parse(&value).expect("should parse"); - match notif { - HubNotification::ToolNotification { session_id, frame } => { - assert_eq!(session_id.as_str(), "s1"); - assert_eq!(frame.tool_id.as_ref().unwrap().as_str(), "echo"); - } - other => panic!("expected ToolNotification, got {other:?}"), - } - } - - #[test] - fn parse_tool_notification_missing_session_id_falls_back_to_unknown() { - let value = json!({ - "jsonrpc": "2.0", - "method": "tool.notification", - "params": { - "tool_id": "echo", - "notification": { - "shape": "custom", - "value": { "kind": "test", "payload": {} } - } - } - }); - let notif = HubNotification::parse(&value).expect("should parse as Unknown, not None"); - assert!( - matches!(notif, HubNotification::Unknown { ref method, .. } if method == "tool.notification"), - "tool.notification without envelope session_id should fall back to Unknown, got {notif:?}" - ); - } - - #[test] - fn parse_unknown_method() { - let value = json!({ - "jsonrpc": "2.0", - "session_id": "s1", - "method": "future.method", - "params": { "key": "value" } - }); - let notif = HubNotification::parse(&value).expect("should parse"); - match notif { - HubNotification::Unknown { method, params } => { - assert_eq!(method, "future.method"); - assert_eq!(params["key"], "value"); - } - other => panic!("expected Unknown, got {other:?}"), - } - } - - #[test] - fn parse_missing_method_returns_none() { - let value = json!({ "jsonrpc": "2.0", "id": "123", "result": {} }); - assert!(HubNotification::parse(&value).is_none()); - } - - #[test] - fn parse_tools_changed_bad_params_falls_back_to_unknown() { - // `params` has wrong shape (missing required fields) — should fall - // back to Unknown instead of returning None and dropping the event. - let value = json!({ - "jsonrpc": "2.0", - "method": "tools_changed", - "params": { "unexpected_field": true } - }); - let notif = HubNotification::parse(&value).expect("should parse as Unknown, not None"); - assert!( - matches!(notif, HubNotification::Unknown { ref method, .. } if method == "tools_changed"), - "malformed tools_changed should fall back to Unknown, got {notif:?}" - ); - } - - #[test] - fn parse_tool_server_status_changed() { - let value = json!({ - "jsonrpc": "2.0", - "session_id": "s1", - "method": "tool.notification", - "params": { - "tool_id": "__tool_server_status", - "notification": { - "shape": "custom", - "value": { - "kind": "status_changed", - "payload": { - "status": "busy", - "active_tool_calls": 2, - "active_tool_names": ["read_file", "grep"], - "background_tasks": 0, - "pending_tool_calls": 0, - "last_tool_call_started_ms": 100, - "last_tool_call_completed_ms": 0, - "uptime_ms": 5000, - } - } - } - } - }); - let notif = HubNotification::parse(&value).expect("should parse"); - match notif { - HubNotification::ToolServerStatusChanged { session_id, status } => { - assert_eq!(session_id.as_str(), "s1"); - assert_eq!( - status.status, - kigi_tool_protocol::ToolServerLifecycleStatus::Busy - ); - assert_eq!(status.active_tool_calls, 2); - } - other => panic!("expected ToolServerStatusChanged, got {other:?}"), - } - } - - #[test] - fn parse_tool_server_status_non_status_tool_id_stays_generic() { - // A tool.notification with a different tool_id should remain - // as ToolNotification, not be intercepted. - let value = json!({ - "jsonrpc": "2.0", - "session_id": "s1", - "method": "tool.notification", - "params": { - "tool_id": "some_other_tool", - "notification": { - "shape": "custom", - "value": { - "kind": "status_changed", - "payload": { "status": "ready" } - } - } - } - }); - let notif = HubNotification::parse(&value).expect("should parse"); - assert!( - matches!(notif, HubNotification::ToolNotification { .. }), - "non-__tool_server_status tool_id should stay as ToolNotification, got {notif:?}" - ); - } - - #[test] - fn parse_tool_notification_bad_params_falls_back_to_unknown() { - // `params` has wrong shape — should fall back to Unknown. - let value = json!({ - "jsonrpc": "2.0", - "session_id": "s1", - "method": "tool.notification", - "params": { "not_a_valid_frame": true } - }); - let notif = HubNotification::parse(&value).expect("should parse as Unknown, not None"); - assert!( - matches!(notif, HubNotification::Unknown { ref method, .. } if method == "tool.notification"), - "malformed tool.notification should fall back to Unknown, got {notif:?}" - ); - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/observability.rs b/crates/common/kigi-computer-hub-sdk/src/observability.rs deleted file mode 100644 index ec7a96b..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/observability.rs +++ /dev/null @@ -1,275 +0,0 @@ -//! Server-side session event emitter. -//! -//! [`ObservabilityBridge`] is a thin facade for emitting session-level -//! events (turn lifecycle, phase changes) to the connected server. Tool-call events -//! (`ToolCallStarted` / `ToolCallCompleted`) are emitted automatically -//! by [`crate::harness::ToolHarness::call`] and do not need the bridge. -//! -//! The caller is responsible for also emitting to the local sink -//! (`EventTracker` in the shell, `EventProcPublisher` in the -//! chat service) — the bridge handles only the server leg. -//! -//! This separation is deliberate: each sampler's local sink has a -//! different type and API surface. Forcing a trait/callback into the -//! bridge would add abstraction overhead without benefit, since the -//! call sites already have the local sink in scope. - -use std::sync::Arc; - -use kigi_tool_protocol::{SessionId, session_event::SessionEvent}; - -use crate::harness::ToolHarness; - -/// Emits [`SessionEvent`]s to the connected server as `ToolNotificationFrame` custom -/// notifications with `kind = "session_event"`. -/// -/// No-ops gracefully when no harness is present (i.e. `harness` is -/// `None`). Server notification failures are silently ignored — the bridge -/// is fire-and-forget so server issues never affect the sampler's main loop. -/// -/// Callers MUST also emit to their local sink separately: -/// - Shell: `self.events.emit(Event::...)` -/// - Chat service: `publisher.publish_agent_event(...)` -pub struct ObservabilityBridge { - harness: Option>, - /// Retained for future payload enrichment and logging. - session_id: SessionId, -} - -impl ObservabilityBridge { - pub fn new(harness: Option>, session_id: SessionId) -> Self { - Self { - harness, - session_id, - } - } - - /// The session id this bridge was created for. - pub fn session_id(&self) -> &SessionId { - &self.session_id - } - - /// Whether a harness is present (i.e. server emission is active). - pub fn has_harness(&self) -> bool { - self.harness.is_some() - } - - /// Emit a session event to the connected server. No-ops if no harness is present. - /// - /// Delegates frame construction + wire dispatch to - /// [`ToolHarness::emit_session_event`] so the SDK keeps a single - /// canonical encoding path. - /// - /// Callers MUST also emit to their local sink separately: - /// - Shell: `self.events.emit(Event::...)` - /// - Chat service: `publisher.publish_agent_event(...)` - pub async fn emit(&self, event: SessionEvent) { - let event_type = match &event { - SessionEvent::TurnStarted { .. } => "turn_started", - SessionEvent::TurnEnded { .. } => "turn_ended", - SessionEvent::ToolCallStarted { .. } => "tool_call_started", - SessionEvent::ToolCallCompleted { .. } => "tool_call_completed", - SessionEvent::PhaseChanged { .. } => "phase_changed", - SessionEvent::Unknown => "unknown", - }; - crate::metrics::session_event(event_type); - if let Some(harness) = &self.harness { - harness.emit_session_event(event).await; - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use kigi_tool_protocol::session_event::{SessionEvent, SessionPhase, ToolCallOutcome}; - use kigi_tool_protocol::turn_hook::TurnHookOutcome; - - fn test_session_id() -> SessionId { - SessionId::new("test-obs-session").expect("valid") - } - - // ── No-harness path ───────────────────────────────────────────── - - #[tokio::test] - async fn emit_without_harness_is_noop() { - let bridge = ObservabilityBridge::new(None, test_session_id()); - // Must not panic and should return immediately. - bridge - .emit(SessionEvent::TurnStarted { - turn_number: 1, - model_id: "grok-3".into(), - yolo_mode: false, - }) - .await; - } - - #[test] - fn has_harness_returns_false_when_none() { - let bridge = ObservabilityBridge::new(None, test_session_id()); - assert!(!bridge.has_harness()); - } - - // ── Constructor field storage ─────────────────────────────────── - - #[test] - fn new_stores_session_id() { - let sid = test_session_id(); - let bridge = ObservabilityBridge::new(None, sid.clone()); - assert_eq!(bridge.session_id(), &sid); - } - - #[test] - fn has_harness_returns_true_when_present() { - let harness = ToolHarness::local_only_with( - crate::harness::LocalRegistry::new(), - test_session_id(), - Default::default(), - ); - let bridge = ObservabilityBridge::new(Some(Arc::new(harness)), test_session_id()); - assert!(bridge.has_harness()); - } - - // ── Serialization correctness ─────────────────────────────────── - - #[test] - fn session_event_serializes_to_expected_json() { - let event = SessionEvent::TurnStarted { - turn_number: 1, - model_id: "grok-3".into(), - yolo_mode: true, - }; - let value = serde_json::to_value(&event).unwrap(); - assert_eq!(value["event_type"], "turn_started"); - assert_eq!(value["turn_number"], 1); - assert_eq!(value["model_id"], "grok-3"); - assert_eq!(value["yolo_mode"], true); - } - - #[test] - fn session_event_turn_ended_serializes_correctly() { - let event = SessionEvent::TurnEnded { - turn_number: 5, - outcome: TurnHookOutcome::Completed, - duration_ms: 3200, - tool_call_count: 12, - model_id: "grok-3".into(), - }; - let value = serde_json::to_value(&event).unwrap(); - assert_eq!(value["event_type"], "turn_ended"); - assert_eq!(value["outcome"], "completed"); - assert_eq!(value["tool_call_count"], 12); - } - - #[test] - fn session_event_tool_call_completed_serializes_correctly() { - let event = SessionEvent::ToolCallCompleted { - tool_call_id: "call-1".into(), - tool_name: "bash".into(), - duration_ms: 500, - outcome: ToolCallOutcome::Success, - }; - let value = serde_json::to_value(&event).unwrap(); - assert_eq!(value["event_type"], "tool_call_completed"); - assert_eq!(value["outcome"], "success"); - } - - #[test] - fn session_event_phase_changed_serializes_correctly() { - let event = SessionEvent::PhaseChanged { - phase: SessionPhase::Sampling, - }; - let value = serde_json::to_value(&event).unwrap(); - assert_eq!(value["event_type"], "phase_changed"); - assert_eq!(value["phase"], "sampling"); - } - - // Frame-construction invariants moved to `crate::harness` where the - // builder now lives (`ToolHarness::emit_session_event`). - - // ── End-to-end with local-only harness ─────────────────────────── - - #[tokio::test] - async fn emit_with_local_only_harness_does_not_panic() { - // A local-only harness has no server connection, so - // `send_notification` returns `Err` — but the bridge ignores - // errors, so this must succeed silently. - let harness = ToolHarness::local_only_with( - crate::harness::LocalRegistry::new(), - test_session_id(), - Default::default(), - ); - let bridge = ObservabilityBridge::new(Some(Arc::new(harness)), test_session_id()); - bridge - .emit(SessionEvent::PhaseChanged { - phase: SessionPhase::ToolExecution, - }) - .await; - } - - #[tokio::test] - async fn emit_all_event_variants_does_not_panic() { - let harness = ToolHarness::local_only_with( - crate::harness::LocalRegistry::new(), - test_session_id(), - Default::default(), - ); - let bridge = ObservabilityBridge::new(Some(Arc::new(harness)), test_session_id()); - - // Smoke test: every variant emits without panic through a - // local-only harness. Includes error/cancelled outcomes to - // cover non-happy-path enum values. - let events = vec![ - SessionEvent::TurnStarted { - turn_number: 1, - model_id: "grok-3".into(), - yolo_mode: false, - }, - SessionEvent::ToolCallStarted { - tool_call_id: "c1".into(), - tool_name: "bash".into(), - turn_number: 1, - }, - SessionEvent::ToolCallCompleted { - tool_call_id: "c1".into(), - tool_name: "bash".into(), - duration_ms: 100, - outcome: ToolCallOutcome::Success, - }, - SessionEvent::ToolCallCompleted { - tool_call_id: "c2".into(), - tool_name: "read_file".into(), - duration_ms: 50, - outcome: ToolCallOutcome::Error, - }, - SessionEvent::ToolCallCompleted { - tool_call_id: "c3".into(), - tool_name: "grep".into(), - duration_ms: 10, - outcome: ToolCallOutcome::Cancelled, - }, - SessionEvent::PhaseChanged { - phase: SessionPhase::Idle, - }, - SessionEvent::TurnEnded { - turn_number: 1, - outcome: TurnHookOutcome::Completed, - duration_ms: 500, - tool_call_count: 3, - model_id: "grok-3".into(), - }, - SessionEvent::TurnEnded { - turn_number: 2, - outcome: TurnHookOutcome::Error, - duration_ms: 100, - tool_call_count: 0, - model_id: "grok-3".into(), - }, - SessionEvent::Unknown, - ]; - - for event in events { - bridge.emit(event).await; - } - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/oidc_provider.rs b/crates/common/kigi-computer-hub-sdk/src/oidc_provider.rs deleted file mode 100644 index abe3e20..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/oidc_provider.rs +++ /dev/null @@ -1,336 +0,0 @@ -//! [`AuthProvider`] that refreshes OIDC tokens before they expire. -//! -//! `current()` checks token expiry and, if needed, performs OIDC -//! discovery + token exchange before returning the credential. - -use std::sync::Arc; -use std::time::Duration; - -use chrono::{DateTime, Utc}; -use parking_lot::Mutex; - -use crate::auth::{AuthCredential, AuthIdentity, AuthProvider}; - -pub type OnRefreshCallback = Arc; - -#[derive(Debug, Clone)] -pub struct RefreshEvent { - pub access_token: String, - pub new_refresh_token: Option, - pub expires_at: Option>, -} - -struct TokenState { - access_token: String, - refresh_token: String, - expires_at: Option>, -} - -pub struct OidcAuthProvider { - state: Mutex, - issuer: String, - client_id: String, - user_id: Option, - principal_type: Option, - principal_id: Option, - on_refresh: Option, -} - -const REFRESH_MARGIN: Duration = Duration::from_secs(60); - -impl std::fmt::Debug for OidcAuthProvider { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("OidcAuthProvider") - .field("issuer", &self.issuer) - .field("client_id", &self.client_id) - .finish_non_exhaustive() - } -} - -pub struct OidcAuthProviderBuilder { - access_token: String, - refresh_token: String, - issuer: String, - client_id: String, - expires_at: Option>, - user_id: Option, - principal_type: Option, - principal_id: Option, - on_refresh: Option, -} - -impl OidcAuthProviderBuilder { - pub fn new( - access_token: impl Into, - refresh_token: impl Into, - issuer: impl Into, - client_id: impl Into, - ) -> Self { - Self { - access_token: access_token.into(), - refresh_token: refresh_token.into(), - issuer: issuer.into(), - client_id: client_id.into(), - expires_at: None, - user_id: None, - principal_type: None, - principal_id: None, - on_refresh: None, - } - } - - pub fn expires_at(mut self, expires_at: DateTime) -> Self { - self.expires_at = Some(expires_at); - self - } - - /// Owner user id parsed from the auth source, surfaced via - /// [`AuthProvider::identity`]. - pub fn user_id(mut self, user_id: impl Into) -> Self { - self.user_id = Some(user_id.into()); - self - } - - pub fn principal_type(mut self, pt: impl Into) -> Self { - self.principal_type = Some(pt.into()); - self - } - - pub fn principal_id(mut self, pid: impl Into) -> Self { - self.principal_id = Some(pid.into()); - self - } - - pub fn on_refresh(mut self, cb: OnRefreshCallback) -> Self { - self.on_refresh = Some(cb); - self - } - - pub fn build(self) -> OidcAuthProvider { - OidcAuthProvider { - state: Mutex::new(TokenState { - access_token: self.access_token, - refresh_token: self.refresh_token, - expires_at: self.expires_at, - }), - issuer: self.issuer, - client_id: self.client_id, - user_id: self.user_id, - principal_type: self.principal_type, - principal_id: self.principal_id, - on_refresh: self.on_refresh, - } - } -} - -impl AuthProvider for OidcAuthProvider { - fn current(&self) -> AuthCredential { - let expired = { - let s = self.state.lock(); - s.expires_at.is_some_and(|exp| { - Utc::now() + chrono::Duration::from_std(REFRESH_MARGIN).unwrap() >= exp - }) - }; - if expired && let Err(e) = self.try_refresh() { - tracing::warn!(error = %e, "OIDC refresh failed, using stale token"); - } - let s = self.state.lock(); - AuthCredential::bearer(&s.access_token) - } - - /// Surface the principal fields parsed from the auth source. `None` only - /// when no `user_id` was supplied (nothing to attribute). - fn identity(&self) -> Option { - let user_id = self.user_id.clone()?; - Some(AuthIdentity { - user_id, - principal_type: self.principal_type.clone(), - principal_id: self.principal_id.clone(), - }) - } -} - -impl OidcAuthProvider { - fn try_refresh(&self) -> Result<(), Box> { - tracing::info!(issuer = %self.issuer, "refreshing OIDC token"); - if let Ok(handle) = tokio::runtime::Handle::try_current() { - tokio::task::block_in_place(|| handle.block_on(self.do_refresh())) - } else { - tokio::runtime::Builder::new_current_thread() - .enable_all() - .build()? - .block_on(self.do_refresh()) - } - } - - async fn do_refresh(&self) -> Result<(), Box> { - let refresh_token = self.state.lock().refresh_token.clone(); - let issuer = self.issuer.trim_end_matches('/'); - let client = reqwest::Client::new(); - - #[derive(serde::Deserialize)] - struct Discovery { - token_endpoint: String, - } - - let disc: Discovery = client - .get(format!("{issuer}/.well-known/openid-configuration")) - .timeout(Duration::from_secs(10)) - .send() - .await? - .error_for_status()? - .json() - .await?; - - let mut params = vec![ - ("grant_type", "refresh_token"), - ("refresh_token", refresh_token.as_str()), - ("client_id", self.client_id.as_str()), - ]; - let pt = self.principal_type.clone(); - let pid = self.principal_id.clone(); - if let Some(ref v) = pt { - params.push(("principal_type", v)); - } - if let Some(ref v) = pid { - params.push(("principal_id", v)); - } - - #[derive(serde::Deserialize)] - struct Tokens { - access_token: String, - #[serde(default)] - refresh_token: Option, - #[serde(default)] - expires_in: Option, - } - - let tokens: Tokens = client - .post(&disc.token_endpoint) - .form(¶ms) - .timeout(Duration::from_secs(15)) - .send() - .await? - .error_for_status()? - .json() - .await?; - - let expires_at = tokens - .expires_in - .map(|s| Utc::now() + chrono::Duration::seconds(s as i64)); - - tracing::info!(expires_at = ?expires_at, "OIDC token refreshed"); - - if let Some(ref cb) = self.on_refresh { - cb(&RefreshEvent { - access_token: tokens.access_token.clone(), - new_refresh_token: tokens.refresh_token.clone(), - expires_at, - }); - } - - let mut s = self.state.lock(); - s.access_token = tokens.access_token; - if let Some(rt) = tokens.refresh_token { - s.refresh_token = rt; - } - s.expires_at = expires_at; - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn current_returns_token_when_not_expired() { - let provider = OidcAuthProviderBuilder::new( - "access-tok", - "refresh-tok", - "https://auth.example.com", - "client1", - ) - .expires_at(Utc::now() + chrono::Duration::hours(1)) - .build(); - - let cred = provider.current(); - match cred { - AuthCredential::Bearer { token } => { - assert_eq!(token, "access-tok"); - } - _ => panic!("expected Bearer"), - } - } - - #[test] - fn current_returns_token_when_no_expiry() { - let provider = OidcAuthProviderBuilder::new( - "no-expiry-tok", - "refresh-tok", - "https://auth.example.com", - "client1", - ) - .build(); - - let cred = provider.current(); - match cred { - AuthCredential::Bearer { token } => assert_eq!(token, "no-expiry-tok"), - _ => panic!("expected Bearer"), - } - } - - #[test] - fn current_returns_stale_token_when_refresh_fails() { - // Expired token, but issuer is unreachable — should return stale - let provider = OidcAuthProviderBuilder::new( - "stale-tok", - "refresh-tok", - "https://localhost:1", // unreachable - "client1", - ) - .expires_at(Utc::now() - chrono::Duration::hours(1)) - .build(); - - let cred = provider.current(); - match cred { - AuthCredential::Bearer { token } => assert_eq!(token, "stale-tok"), - _ => panic!("expected Bearer"), - } - } - - #[test] - fn identity_surfaces_principal_fields() { - let provider = OidcAuthProviderBuilder::new("tok", "rt", "https://auth.example.com", "c1") - .user_id("user-1") - .principal_type("Team") - .principal_id("team-9") - .build(); - let id = provider.identity().expect("identity present"); - assert_eq!(id.user_id, "user-1"); - assert_eq!(id.principal_type.as_deref(), Some("Team")); - assert_eq!(id.principal_id.as_deref(), Some("team-9")); - } - - #[test] - fn identity_none_without_user_id() { - let provider = - OidcAuthProviderBuilder::new("tok", "rt", "https://auth.example.com", "c1").build(); - assert!(provider.identity().is_none()); - } - - #[test] - fn debug_does_not_leak_tokens() { - let provider = OidcAuthProviderBuilder::new( - "secret-access-token", - "secret-refresh-token", - "https://auth.example.com", - "client1", - ) - .build(); - - let debug = format!("{provider:?}"); - assert!(!debug.contains("secret-access-token")); - assert!(!debug.contains("secret-refresh-token")); - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/pool.rs b/crates/common/kigi-computer-hub-sdk/src/pool.rs deleted file mode 100644 index 816bc31..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/pool.rs +++ /dev/null @@ -1,344 +0,0 @@ -//! Process-wide connection pool keyed by `(url, principal)`. -//! -//! Two [`crate::ToolServer`] builds with the same `(url, credential)` -//! observe the same `Arc`; distinct credentials open -//! distinct sockets. The pool is the canonical entry point — direct -//! [`crate::HubConnection::connect`] calls are reserved for tests and -//! one-shot programs that explicitly want unpooled behaviour. - -use std::sync::Arc; -use std::time::{Duration, Instant}; - -use dashmap::DashMap; -use kigi_tool_protocol::ConnectionKind; -use tokio::sync::OnceCell; -use tokio::task::JoinHandle; -use url::Url; - -use crate::auth::AuthProvider; -use crate::connection::{ - ConnKey, ConnectCallback, ConnectionConfig, ConnectionTuning, DisconnectCallback, - HubConnection, ReconnectCallback, -}; -use crate::error::ClientError; - -/// Idle window for the reaper: a pooled connection is evictable once it is -/// unused (`Arc::strong_count == 1`, i.e. only the pool holds it) **and** -/// `now - last_handout >= DEFAULT_POOL_IDLE_TTL`. -/// -/// Note the clock is `last_handout` (the last time the pool returned the -/// connection), not the moment the last consumer `Arc` was dropped: a -/// connection held longer than the TTL and then released is eligible on the -/// very next sweep, with no extra post-drop grace period. The only hard -/// guarantee is that an in-use connection (`strong_count > 1`) is never -/// reaped. Tuned well above the server's own 90s dead-peer idle timeout so a -/// short borrow between turns of an active conversation isn't churned. -pub const DEFAULT_POOL_IDLE_TTL: Duration = Duration::from_secs(300); - -/// How often the shared pool's idle reaper scans for evictable entries. -pub const DEFAULT_POOL_SWEEP_INTERVAL: Duration = Duration::from_secs(60); - -/// A pooled connection plus the last time it was handed out to a caller. -/// -/// `last_handout` is refreshed on every [`HubConnectionPool::get_or_connect`] -/// hit (and on the initial insert), so a connection that is repeatedly -/// re-fetched never looks idle even if its [`Arc`] strong count briefly -/// returns to 1 between fetches. Eviction additionally requires -/// `Arc::strong_count == 1` (only the pool holds it), so a connection a -/// consumer still holds is never reaped regardless of `last_handout`. -struct Pooled { - conn: Arc, - last_handout: Instant, -} - -/// The process-global pool used by [`HubConnectionPool::shared`]. -/// -/// `tokio::sync::OnceCell` is preferred over `std::sync::OnceLock` / -/// `LazyLock` here because the pool is only ever observed from -/// async contexts (the connection actor lives on a tokio runtime -/// already), so the async-aware `get_or_init` semantics avoid the -/// blocking-init footgun of the sync alternatives without taking a -/// hard dependency on additional sync primitives. -/// -/// Tests MUST use [`HubConnectionPool::new`] to avoid cross-test -/// pollution: cargo runs all integration tests in the same binary -/// unless otherwise configured, so any test that touches -/// `HubConnectionPool::shared()` leaves the pool populated for -/// subsequent tests. -static SHARED: OnceCell> = OnceCell::const_new(); - -/// Pool of live server connections. -pub struct HubConnectionPool { - connections: DashMap, -} - -impl std::fmt::Debug for HubConnectionPool { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("HubConnectionPool") - .field("connection_count", &self.connections.len()) - .finish() - } -} - -impl HubConnectionPool { - /// Build a fresh, unshared pool. Tests typically use this so each - /// test sees an isolated registry. - pub fn new() -> Arc { - Arc::new(Self { - connections: DashMap::new(), - }) - } - - /// Return the process-wide shared pool, lazily initialising it on - /// the first call. Subsequent callers in the same process observe - /// the same `Arc`. - /// - /// The shared pool spawns an idle reaper (see [`Self::spawn_idle_reaper`]) - /// exactly once, so a connection that is unused (`strong_count == 1`) and - /// has not been handed out for [`DEFAULT_POOL_IDLE_TTL`] is closed instead - /// of living for the whole process lifetime. (Unpooled / test pools built - /// via [`Self::new`] do not - /// get a reaper; they can call [`Self::sweep_idle`] directly.) - pub async fn shared() -> Arc { - SHARED - .get_or_init(|| async { - let pool = Self::new(); - pool.spawn_idle_reaper(DEFAULT_POOL_IDLE_TTL, DEFAULT_POOL_SWEEP_INTERVAL); - pool - }) - .await - .clone() - } - - /// Look up an existing pooled connection for `(url, credential)`, - /// or open a fresh one if no pooled entry exists. - /// - /// `kind` is the connection role announced in the hello frame. The - /// pool is keyed by `(url, principal)` only; mixing - /// [`ConnectionKind`] values for the same `(url, principal)` is a - /// caller error and surfaces as a [`ClientError::InvalidConfig`]. - /// - /// The optional extra access key is not part of the pool key, so the first - /// caller's key is the one carried on a shared connection's handshake (in - /// practice it is a per-deployment constant). The plaintext-scheme guard is - /// re-checked on every call below so it can't be bypassed by a cached - /// insecure entry. - pub async fn get_or_connect( - self: &Arc, - url: Url, - credential: Arc, - kind: ConnectionKind, - on_reconnect: Option>, - on_disconnect: Option>, - server_id: Option, - alpha_test_key: Option, - allow_insecure_ws: bool, - ) -> Result, ClientError> { - self.get_or_connect_tuned( - url, - credential, - kind, - on_reconnect, - on_disconnect, - None, // on_connect (unused by the simple wrapper) - server_id, - None, - None, - alpha_test_key, - allow_insecure_ws, - ConnectionTuning::default(), - ) - .await - } - - /// Like [`Self::get_or_connect`] but carries optional connection-tuning - /// knobs ([`ConnectionTuning`]) onto a freshly-opened connection. A - /// `Default` tuning is behaviourally identical to `get_or_connect`, so - /// existing callers are unaffected. - /// - /// Tuning binds to the socket at open time: it takes effect only when - /// THIS call opens the connection. Because the pool dedups by - /// `(url, principal)`, a hit on an existing entry returns that - /// connection as-is and the `tuning` argument is ignored — the first - /// opener's ping/backoff settings win for the lifetime of the pooled - /// connection. Callers that need distinct tuning must use a distinct - /// `(url, principal)` or an unpooled [`HubConnection::connect`]. - pub(crate) async fn get_or_connect_tuned( - self: &Arc, - url: Url, - credential: Arc, - kind: ConnectionKind, - on_reconnect: Option>, - on_disconnect: Option>, - on_connect: Option>, - server_id: Option, - server_description: Option, - server_metadata: Option, - alpha_test_key: Option, - allow_insecure_ws: bool, - tuning: ConnectionTuning, - ) -> Result, ClientError> { - if url.scheme() != "wss" && !crate::connection::host_is_loopback(&url) && !allow_insecure_ws - { - return Err(ClientError::InsecureScheme { url }); - } - let key = ConnKey { - url: url.as_str().to_owned(), - principal: credential.principal_key(), - }; - if let Some(mut existing) = self.connections.get_mut(&key) { - existing.last_handout = Instant::now(); - let conn = existing.conn.clone(); - drop(existing); - if conn.kind() != kind { - return Err(ClientError::InvalidConfig(format!( - "pool entry for {} bound to {:?}; rebuild requested {:?}", - key.url, - conn.kind(), - kind - ))); - } - return Ok(conn); - } - let config = ConnectionConfig { - url, - credential, - kind, - on_reconnect, - on_disconnect, - on_connect, - server_id, - server_description, - server_metadata, - outbound_buffer: None, - tuning, - alpha_test_key, - allow_insecure_ws, - on_fatal: Some(Arc::downgrade(self)), - }; - let conn = HubConnection::connect(config).await?; - // Race window: another caller may have inserted between our - // `get` and `connect`. Resolve via `entry().or_insert_with` - // semantics — if we lose the race we drop our fresh - // connection and adopt the winning one. - match self.connections.entry(key.clone()) { - dashmap::Entry::Occupied(mut existing) => { - existing.get_mut().last_handout = Instant::now(); - let winner = existing.get().conn.clone(); - drop(conn); - if winner.kind() != kind { - return Err(ClientError::InvalidConfig(format!( - "pool entry for {} bound to {:?}; rebuild requested {:?}", - key.url, - winner.kind(), - kind - ))); - } - Ok(winner) - } - dashmap::Entry::Vacant(slot) => { - crate::metrics::pool_connections_inc(); - slot.insert(Pooled { - conn: conn.clone(), - last_handout: Instant::now(), - }); - Ok(conn) - } - } - } - - /// Number of pooled connections. Intended for tests and metrics. - pub fn len(&self) -> usize { - self.connections.len() - } - - /// `true` when no connection is pooled. - pub fn is_empty(&self) -> bool { - self.connections.is_empty() - } - - /// Forget the pooled connection for `key`. The actual underlying - /// `Arc` is dropped only when no other holder - /// keeps a reference; the next [`Self::get_or_connect`] for the - /// same key opens a fresh socket. - pub fn forget(&self, key: &ConnKey) { - if self.connections.remove(key).is_some() { - crate::metrics::pool_connections_dec(); - } - } - - /// Close and remove every pooled connection that is BOTH unused (no live - /// consumer holds an `Arc` — only the pool does, so `strong_count == 1`) - /// AND idle longer than `idle_ttl` (no hand-out within the window). - /// Removing the entry drops the pool's last `Arc`, whose - /// `Drop` closes the socket. - /// - /// The strong-count check runs inside the map's per-shard lock (via - /// [`DashMap::retain`]), serialised against `get_or_connect`, so a - /// connection handed out concurrently is never evicted out from under a - /// caller. Returns the number of connections evicted. - pub fn sweep_idle(&self, idle_ttl: Duration) -> usize { - let now = Instant::now(); - let mut evicted = 0usize; - self.connections.retain(|_key, pooled| { - let idle_for = now.saturating_duration_since(pooled.last_handout); - // `strong_count == 1` ⇒ only this pool entry references the - // connection, so no consumer can still be using it. - let unused = Arc::strong_count(&pooled.conn) == 1; - let evict = unused && idle_for >= idle_ttl; - if evict { - evicted += 1; - } - !evict - }); - for _ in 0..evicted { - crate::metrics::pool_connections_dec(); - crate::metrics::pool_evictions_inc(); - } - evicted - } - - /// Spawn a background task that calls [`Self::sweep_idle`] every - /// `sweep_interval`, closing connections idle longer than `idle_ttl`. - /// - /// The task holds a [`std::sync::Weak`] to the pool, so it exits on its - /// own once the last strong `Arc` is dropped (it never - /// keeps the pool alive). The first interval tick is skipped so a - /// freshly-handed-out connection is never swept on the immediate tick. - pub fn spawn_idle_reaper( - self: &Arc, - idle_ttl: Duration, - sweep_interval: Duration, - ) -> JoinHandle<()> { - let weak = Arc::downgrade(self); - tokio::spawn(async move { - let mut ticker = tokio::time::interval(sweep_interval); - ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - // `interval`'s first tick resolves immediately; skip it. - ticker.tick().await; - loop { - ticker.tick().await; - let Some(pool) = weak.upgrade() else { break }; - pool.sweep_idle(idle_ttl); - } - }) - } - - /// Like [`Self::forget`] but identity-checked: only removes the slot - /// when `predicate` accepts the currently-stored connection. The - /// self-evicting actor passes an `Arc::ptr_eq` check so a race-loser - /// can never drop the winner's fresh entry (ABA-safe). - pub(crate) fn forget_if( - &self, - key: &ConnKey, - predicate: impl FnOnce(&Arc) -> bool, - ) { - if self - .connections - .remove_if(key, |_, pooled| predicate(&pooled.conn)) - .is_some() - { - crate::metrics::pool_connections_dec(); - } - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/refcount.rs b/crates/common/kigi-computer-hub-sdk/src/refcount.rs deleted file mode 100644 index d9e2a28..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/refcount.rs +++ /dev/null @@ -1,123 +0,0 @@ -//! Generic refcounted-binding helper used by the connection's -//! bound-session set. -//! -//! Multiple [`crate::ToolServer`] instances can share one -//! [`crate::HubConnection`] when they target the same `(url, principal)`. -//! Each instance independently asks for a session binding; the substrate -//! must `register_session` once per session (not once per consumer) and -//! `unregister_session` only when the LAST consumer drops its borrow. -//! [`RefCountedSet`] tracks the per-key borrow count behind a -//! [`dashmap::DashMap`] so increments and decrements never serialise on -//! a single mutex. - -use std::hash::Hash; - -use dashmap::DashMap; - -/// Refcounted set keyed by `K`. Each [`Self::increment`] returns the -/// new count; the corresponding [`Self::decrement`] returns the count -/// AFTER the decrement (so callers fire teardown when the result is -/// `Some(0)`). -#[derive(Debug, Default)] -pub struct RefCountedSet { - counts: DashMap, -} - -impl RefCountedSet { - /// Empty set. - pub fn new() -> Self { - Self { - counts: DashMap::new(), - } - } - - /// Increment `key`'s refcount. Returns `(prev_count, new_count)` - /// so callers can detect the 0→1 edge (when the protocol-level - /// register call must fire). - pub fn increment(&self, key: K) -> (u64, u64) - where - K: Clone, - { - let mut entry = self.counts.entry(key).or_insert(0); - let prev = *entry; - *entry = prev.saturating_add(1); - (prev, *entry) - } - - /// Decrement `key`'s refcount. Returns the post-decrement count; - /// `Some(0)` means the entry was removed and callers should fire - /// the protocol-level unregister. `None` means the key was not - /// present (idempotent drop). - pub fn decrement(&self, key: &K) -> Option { - let mut current = None; - self.counts.remove_if_mut(key, |_, value| { - *value = value.saturating_sub(1); - current = Some(*value); - *value == 0 - }); - current - } - - /// Snapshot the live keys. Allocates a fresh `Vec` — only used by - /// the reconnect-replay path which fires once per disconnect. - pub fn snapshot_keys(&self) -> Vec - where - K: Clone, - { - self.counts.iter().map(|kv| kv.key().clone()).collect() - } - - /// `true` when no key has a non-zero refcount. - pub fn is_empty(&self) -> bool { - self.counts.is_empty() - } - - /// Number of distinct live keys. - pub fn len(&self) -> usize { - self.counts.len() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn increment_returns_new_count() { - let set = RefCountedSet::<&'static str>::new(); - assert_eq!(set.increment("a"), (0, 1)); - assert_eq!(set.increment("a"), (1, 2)); - assert_eq!(set.increment("b"), (0, 1)); - assert_eq!(set.len(), 2); - } - - #[test] - fn decrement_removes_at_zero() { - let set = RefCountedSet::<&'static str>::new(); - set.increment("a"); - set.increment("a"); - assert_eq!(set.decrement(&"a"), Some(1)); - assert!(!set.is_empty()); - assert_eq!(set.decrement(&"a"), Some(0)); - assert!(set.is_empty()); - } - - #[test] - fn decrement_unknown_returns_none() { - let set = RefCountedSet::<&'static str>::new(); - assert!(set.decrement(&"missing").is_none()); - } - - #[test] - fn increment_saturates_at_u64_max() { - let set = RefCountedSet::<&'static str>::new(); - // Pre-load the entry to MAX-1 via direct DashMap access. The - // public API only ever reaches this region via overflow, - // which is impossible in practice; this test pins the - // saturating_add defensive line so it can't silently regress - // to wrapping_add. - set.counts.insert("max", u64::MAX - 1); - assert_eq!(set.increment("max"), (u64::MAX - 1, u64::MAX)); - assert_eq!(set.increment("max"), (u64::MAX, u64::MAX)); - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/server.rs b/crates/common/kigi-computer-hub-sdk/src/server.rs deleted file mode 100644 index 14bf917..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/server.rs +++ /dev/null @@ -1,2649 +0,0 @@ -//! Tool-server runtime: builder, handler trait, and inbound dispatch loop. -//! -//! A [`ToolServer`] is the SDK-side counterpart to a server-side -//! `ToolServer` connection. The builder collects: -//! -//! - the [`crate::HubConnectionPool`] to attach to, -//! - the server URL, -//! - an [`crate::AuthCredential`], -//! - one or more [`ToolServerHandler`] implementations, -//! - zero or more sessions (bound during [`ToolServer::run`]). -//! -//! On [`ToolServerBuilder::build`] the server registers its identity -//! and tools with the server. [`ToolServer::run`] drives the inbound loop: -//! every server-issued `tool_call_request` is decoded, dispatched to -//! the matching handler, and the response is shipped back over the -//! shared connection. [`ToolServer::shutdown`] cooperatively stops -//! `run` and unregisters everything; [`Drop`] is a best-effort -//! fallback that schedules the same cleanup on a background task. - -use async_trait::async_trait; -use dashmap::DashMap; -use futures::StreamExt; -use kigi_tool_protocol::{ - ConnectionKind, HookEvent, HookFrame, HookReplyFrame, JsonRpcError, JsonRpcId, - JsonRpcNotification, JsonRpcResponse, JsonRpcVersion, Method, ResponseOutcome, SessionId, - ToolCallId, ToolCallParams, ToolCallProgressFrame, ToolCallResult, ToolErrorWire, ToolId, - ToolOutputWire, ToolServerEvictParams, error_codes, -}; -use kigi_tool_runtime::{ - BehaviorVersion, Cancellation, Cwd, ToolCallContext, ToolError, ToolProgress, ToolStream, - ToolStreamItem, TraceContext, TypedToolOutput, -}; -use kigi_tool_types::ToolDescription; -use serde_json::Value; -use std::collections::HashMap; -use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; -use tokio::sync::broadcast::error::RecvError; -use tokio::sync::mpsc; -use tokio_util::sync::CancellationToken; -use tracing::{debug, warn}; -use url::Url; - -use crate::auth::{AuthCredential, AuthProvider}; -use crate::cancel::CancelRegistry; -use crate::connection::{ - ConnectCallback, ConnectionTuning, DisconnectCallback, HubConnection, ReconnectCallback, - ReconnectEvent, -}; -use crate::connection_borrow::ConnectionBorrow; -use crate::demux::InboundFrame; -use crate::error::ClientError; -use crate::pool::HubConnectionPool; - -/// Fired after reconnect `serve` replay completes (async settle). -pub type ReconnectSettledCallback = Box; - -/// Outcome of a `system.notify` request. `Accepted` is the server's ack to forward; -/// `ForwardingUnsupported` is an older server that lacks the method (`-32601`). -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum SystemNotifyAck { - Accepted, - ForwardingUnsupported, -} - -/// Serialized JSON byte length without allocating the string. -fn json_serialized_len(value: &Value) -> Result { - struct Counter(usize); - impl std::io::Write for Counter { - fn write(&mut self, buf: &[u8]) -> std::io::Result { - self.0 += buf.len(); - Ok(buf.len()) - } - fn flush(&mut self) -> std::io::Result<()> { - Ok(()) - } - } - let mut counter = Counter(0); - serde_json::to_writer(&mut counter, value).map_err(|e| ClientError::Serde(e.to_string()))?; - Ok(counter.0) -} - -fn system_notify_ack_from_outcome( - outcome: ResponseOutcome, -) -> Result { - match outcome { - ResponseOutcome::Result(_) => Ok(SystemNotifyAck::Accepted), - // Plain `-32601` (no `data` discriminator) means the server lacks the method; - // require `data` absent so a richer error still flows through the normal taxonomy. - ResponseOutcome::Error(err) - if err.data.is_none() - && kigi_tool_protocol::error_codes::string_for(err.code) - == Some("method_not_found") => - { - Ok(SystemNotifyAck::ForwardingUnsupported) - } - ResponseOutcome::Error(err) => Err(ClientError::from_jsonrpc_error(err)), - } -} - -/// Per-session inbound queue depth. The spawn-per-request dispatcher -/// dequeues immediately, so the inbox barely lags and the -/// inbox-full path is now a rare relief valve rather than the steady -/// state; 64 is kept as a comfortable burst buffer ahead of the -/// admission-deadline backpressure. -const SESSION_INBOX_BUFFER: usize = 64; - -type SessionHandlerMap = - Arc>>>>; - -/// A resolved session binding: handlers plus the bind-report fields echoed -/// in the `session.bind` response. -#[derive(Default)] -pub struct ResolvedSessionHandlers { - pub handlers: Vec>, - /// Tool ids the resolver declined to serve; forwarded as - /// [`kigi_tool_protocol::SessionBindResult::unserved_tool_ids`]. - pub unserved_tool_ids: Vec, - /// Human-readable reason the resolver failed the toolset closed; - /// forwarded as [`kigi_tool_protocol::SessionBindResult::resolve_error`]. - pub resolve_error: Option, -} - -impl ResolvedSessionHandlers { - /// A fully-served binding with no divergence to report. - pub fn full(handlers: Vec>) -> Self { - Self { - handlers, - unserved_tool_ids: Vec::new(), - resolve_error: None, - } - } -} - -/// Resolves a session's served handler set from the raw `session.bind` -/// params. When unset, sessions bind with a clone of `initial_handlers`. -/// -/// Returning `Err` **fails the bind**: the server receives an error response -/// (and classifies it as bind-unavailable so the harness can re-provision) -/// instead of a "successful" bind that advertises zero model-facing tools — -/// which would make every subsequent tool call fail as route-missing with no -/// hint of the real cause. -pub type SessionHandlerResolver = Arc< - dyn Fn( - SessionId, - Option, - ) - -> futures::future::BoxFuture<'static, Result> - + Send - + Sync, ->; - -/// User-facing tool implementation. -/// -/// The server speaks JSON, so handlers receive `serde_json::Value` -/// arguments and return a `ToolStream`. Implementations that -/// already use [`kigi_tool_runtime::Tool`] can adapt by calling the -/// underlying tool's `execute` and serialising the typed output. -#[async_trait] -pub trait ToolServerHandler: Send + Sync + 'static { - /// Stable identity used by the server to route to this handler. - fn tool_id(&self) -> ToolId; - - /// Model-facing description and argument schema. - fn description(&self) -> ToolDescription; - - /// Argument schema. Default `None` — many handlers don't expose a - /// schema independent of the description. - fn input_schema(&self) -> Option { - None - } - - /// Execute one tool call. - /// - /// Implementations MUST honour the [`ToolStream`] invariant: zero - /// or more `Progress` items followed by exactly one `Terminal`. - /// - /// `Progress` items are forwarded as `tool_call_progress` - /// notifications; the `Terminal` item is shipped as the response. - async fn handle_call(&self, ctx: ToolCallContext, args: Value) -> ToolStream; - - /// Receive a harness-issued hook for `session_id`. - /// - /// `frame.tool_id` was set when the harness routed the hook to a - /// specific tool (e.g. [`HookEvent::Cancel`] with the matching - /// `tool_id`); `None` when the hook is session-wide (broadcast). - /// Implementations route by `frame.event` shape and may - /// abort in-flight calls correlated by `frame.call_id`. - /// - /// Default is a no-op so existing handlers do not need to opt in. - /// Override to receive cancel / pause / resume / session-ended / - /// custom hooks. - #[allow(unused_variables)] - async fn handle_hook(&self, session_id: SessionId, frame: HookFrame) {} - - /// Answer a request/response hook (one whose `hook_id` is set); first `Some` wins, `None` declines. - #[allow(unused_variables)] - async fn handle_hook_request(&self, session_id: SessionId, frame: HookFrame) -> Option { - None - } - - /// Handle a server-issued `tool_server.evict` (graceful-shutdown request), - /// fanned out to the evicted session's handlers. Implementations should - /// drain in-flight work within `params.grace_period_ms`, after which the - /// server force-closes the connection. Default is a no-op. - #[allow(unused_variables)] - async fn handle_evict(&self, params: ToolServerEvictParams) {} -} - -/// Builder for [`ToolServer`]. See module docs for end-to-end usage. -#[derive(Default)] -pub struct ToolServerBuilder { - pool: Option>, - url: Option, - auth: Option>, - sessions: Vec, - handlers: Vec>, - on_reconnect: Option>, - /// Fired after reconnect serve replay finishes (async settle, not the sync - /// socket-up `on_reconnect`). Use for readiness markers that must not - /// precede server session re-serve. - on_reconnect_settled: Option>, - on_disconnect: Option>, - on_connect: Option>, - metadata: Option, - server_id: Option, - server_description: Option, - alpha_test_key: Option, - allow_insecure_ws: bool, - session_max_inflight: Option, - conn_max_inflight: Option, - global_max_inflight: Option, - admission_wait_timeout: Option, - ws_ping_interval: Option, - ws_liveness_deadline: Option, - reconnect_backoff: Option>, - session_handler_resolver: Option, - binary_version: Option, -} - -impl ToolServerBuilder { - /// Attach an extra access header on every (re)connect. - pub fn alpha_test_key(mut self, key: impl Into) -> Self { - self.alpha_test_key = Some(key.into()); - self - } - - /// Permit plaintext `ws://` to a non-loopback host. Only enable - /// when the transport is otherwise secured (e.g. a private network - /// or TLS-terminating proxy) — the bearer would otherwise cross the - /// wire in cleartext. - pub fn allow_insecure_ws(mut self, allow: bool) -> Self { - self.allow_insecure_ws = allow; - self - } - - /// Max concurrent *running* calls per session (default 16), - /// enforced by the spawned per-request dispatcher. - pub fn session_max_inflight(mut self, max: usize) -> Self { - self.session_max_inflight = Some(max); - self - } - - /// Max concurrent *running* calls across all sessions on this - /// connection (default 256). - pub fn conn_max_inflight(mut self, max: usize) -> Self { - self.conn_max_inflight = Some(max); - self - } - - /// Process-wide concurrent running-call ceiling (default 1024). Shared - /// by every connection via a once-initialized semaphore; the - /// `XAI_TOOL_SERVER_GLOBAL_MAX_INFLIGHT` env var overrides this at - /// startup. Because the global cell initializes once, the first server - /// built in the process fixes the process-wide value. - pub fn global_max_inflight(mut self, max: usize) -> Self { - self.global_max_inflight = Some(max); - self - } - - /// Bounded wait before an admission attempt is rejected with the - /// overloaded (-32016 "tool_busy") error (default 3s). A single - /// deadline spans all three semaphore acquisitions. - pub fn admission_wait_timeout(mut self, timeout: std::time::Duration) -> Self { - self.admission_wait_timeout = Some(timeout); - self - } - - /// Override the WebSocket keepalive ping cadence on a freshly-opened - /// connection (default 30s). Not setting it preserves the default. - pub fn with_ws_ping_interval(mut self, interval: std::time::Duration) -> Self { - self.ws_ping_interval = Some(interval); - self - } - - /// Override the inbound-liveness deadline on a freshly-opened - /// connection: if no inbound WebSocket frame of any kind arrives within - /// this window, the connection is declared dead and reconnected. This - /// catches silently dead transports (e.g. a VM snapshot restore or - /// NAT/LB flow expiry) that a send-only keepalive never notices. - /// - /// Default (also used for a zero value): 2.5× the effective ping - /// interval — 75s at the default 30s ping — which guarantees at least - /// two keepalive pings fit in every window, so a healthy-but-idle - /// connection (one pong per ping) can never trip it. Explicit values - /// are honored verbatim; keep them comfortably above the ping interval - /// for the same reason (a value at or below the ping interval churns - /// healthy idle connections and is logged as a warning at connect). - pub fn with_ws_liveness_deadline(mut self, deadline: std::time::Duration) -> Self { - self.ws_liveness_deadline = Some(deadline); - self - } - - /// Override the reconnect backoff schedule on a freshly-opened - /// connection (default: the built-in exponential table capped at 10s). - /// Each attempt uses the next slot, clamping at the last; an empty - /// schedule falls back to the default. Not setting it preserves the - /// default. - pub fn with_reconnect_backoff(mut self, schedule: Vec) -> Self { - self.reconnect_backoff = Some(schedule.into()); - self - } - - /// Connection pool to attach to. Required. - pub fn pool(mut self, pool: Arc) -> Self { - self.pool = Some(pool); - self - } - - /// Server URL (`ws://` / `wss://`). Required. - pub fn url(mut self, url: Url) -> Self { - self.url = Some(url); - self - } - - pub fn auth(mut self, cred: AuthCredential) -> Self { - self.auth = Some(Arc::new(cred)); - self - } - - pub fn auth_provider(mut self, provider: Arc) -> Self { - self.auth = Some(provider); - self - } - - /// Bind `session_id` on the underlying connection. May be called - /// repeatedly; each call adds one session that `run()` will bind - /// via `bind_session_local`. - pub fn session(mut self, session_id: SessionId) -> Self { - self.sessions.push(session_id); - self - } - - /// Register a tool. May be called repeatedly. - pub fn tool(mut self, handler: H) -> Self { - self.handlers.push(Arc::new(handler)); - self - } - - /// Register a dynamically-typed tool handler. - pub fn tool_dyn(mut self, handler: Arc) -> Self { - self.handlers.push(handler); - self - } - - /// Optional callback fired once per successful reconnect cycle. - pub fn on_reconnect(mut self, cb: F) -> Self - where - F: Fn(ReconnectEvent) + Send + Sync + 'static, - { - self.on_reconnect = Some(Arc::new(Box::new(cb) as ReconnectCallback)); - self - } - - /// Optional callback fired after reconnect `serve` replay completes for all - /// active sessions (runs on the reconnect task, after the sync - /// [`Self::on_reconnect`] / hello). Prefer this for readiness markers so - /// "server-ready" means registered **and** session tools re-served. - /// - /// It only fires when **every** session re-served successfully **and** no - /// disconnect raced the (async) replay; otherwise it is skipped and the - /// next reconnect's replay settles instead. This keeps a readiness marker - /// from being resurrected while the socket is already down again. - pub fn on_reconnect_settled(mut self, cb: F) -> Self - where - F: Fn() + Send + Sync + 'static, - { - self.on_reconnect_settled = Some(Arc::new(Box::new(cb) as ReconnectSettledCallback)); - self - } - - /// Optional callback fired when the live server socket drops or closes. - pub fn on_disconnect(mut self, cb: F) -> Self - where - F: Fn() + Send + Sync + 'static, - { - self.on_disconnect = Some(Arc::new(Box::new(cb) as DisconnectCallback)); - self - } - - /// Optional callback fired once on the initial successful connect, before - /// the actor starts (so it happens-before any disconnect/reconnect). - pub fn on_connect(mut self, cb: F) -> Self - where - F: Fn() + Send + Sync + 'static, - { - self.on_connect = Some(Arc::new(Box::new(cb) as ConnectCallback)); - self - } - - /// Stable server identity for `servers.list` discovery and - /// `session.open` addressing. Sent in the hello frame. - pub fn server_id(mut self, id: kigi_tool_protocol::ServerId) -> Self { - self.server_id = Some(id); - self - } - - /// Server description for `servers.list` discovery. - pub fn server_description(mut self, desc: impl Into) -> Self { - self.server_description = Some(desc.into()); - self - } - - /// Attach opaque metadata to every tool registered by this server. - /// Propagated to `ServerInfo.metadata` in `servers.list` responses. - pub fn metadata(mut self, metadata: serde_json::Value) -> Self { - self.metadata = Some(metadata); - self - } - - /// Install a per-session handler resolver (the binding path). - pub fn session_handler_resolver(mut self, resolver: SessionHandlerResolver) -> Self { - self.session_handler_resolver = Some(resolver); - self - } - - /// Version of the embedding binary, echoed as - /// [`kigi_tool_protocol::SessionBindResult::binary_version`]. - pub fn binary_version(mut self, version: impl Into) -> Self { - self.binary_version = Some(version.into()); - self - } - - /// Resolve the pool entry, bind sessions, register tools. - /// - /// Returns a [`ToolServer`] ready to be driven via - /// [`ToolServer::run`]. On any mid-loop failure (a tool's - /// `register_tool` returning `Err`, a session bind failing) the - /// builder rolls back every successfully-registered tool and - /// every successfully-bound session before returning the original - /// error, so a failed `build()` does not leak server-side state. - pub async fn build(self) -> Result { - let pool = self - .pool - .ok_or_else(|| ClientError::InvalidConfig("missing pool".to_owned()))?; - let url = self - .url - .ok_or_else(|| ClientError::InvalidConfig("missing url".to_owned()))?; - let auth = self - .auth - .ok_or_else(|| ClientError::InvalidConfig("missing auth".to_owned()))?; - if self.handlers.is_empty() { - return Err(ClientError::InvalidConfig( - "ToolServer must register at least one tool".to_owned(), - )); - } - - // Pre-create shared state so the on_reconnect callback can - // signal `serve` replay after a reconnect. - let active_sessions = parking_lot::Mutex::new(Vec::::new()); - let reconnect_notify = Arc::new(tokio::sync::Notify::new()); - - // Compose the internal reconnect handler (signal serve replay) - // with the user's optional callback. - let user_on_reconnect = self.on_reconnect.clone(); - let notify_clone = Arc::clone(&reconnect_notify); - let combined_reconnect: Arc = - Arc::new(Box::new(move |event: ReconnectEvent| { - notify_clone.notify_one(); - if let Some(ref cb) = user_on_reconnect { - cb(event); - } - })); - - // Compose the internal disconnect handler (bump the epoch so a - // disconnect racing an in-flight serve replay is observed by the - // reconnect task) with the user's optional callback. - let disconnect_epoch = Arc::new(AtomicU64::new(0)); - let user_on_disconnect = self.on_disconnect.clone(); - let epoch_for_disconnect = Arc::clone(&disconnect_epoch); - let combined_disconnect: Arc = Arc::new(Box::new(move || { - epoch_for_disconnect.fetch_add(1, Ordering::Release); - if let Some(ref cb) = user_on_disconnect { - cb(); - } - })); - - let tuning = ConnectionTuning { - ws_ping_interval: self.ws_ping_interval, - ws_liveness_deadline: self.ws_liveness_deadline, - reconnect_backoff: self.reconnect_backoff, - }; - let borrow = ConnectionBorrow::acquire( - pool, - url, - auth, - ConnectionKind::ToolServer, - Some(combined_reconnect), - Some(combined_disconnect), - self.on_connect, - self.server_id, - self.server_description, - self.metadata, - self.alpha_test_key, - self.allow_insecure_ws, - tuning, - ) - .await?; - - let global_sem = crate::admission::global_semaphore( - self.global_max_inflight - .unwrap_or(crate::admission::DEFAULT_GLOBAL_MAX_INFLIGHT), - ); - let admission = Arc::new(crate::admission::Admission::new( - self.session_max_inflight - .unwrap_or(crate::admission::DEFAULT_SESSION_MAX_INFLIGHT), - self.conn_max_inflight - .unwrap_or(crate::admission::DEFAULT_CONN_MAX_INFLIGHT), - global_sem, - self.admission_wait_timeout - .unwrap_or(crate::admission::DEFAULT_ADMISSION_WAIT_TIMEOUT), - )); - - let inner = Arc::new(ToolServerInner { - borrow, - initial_handlers: self.handlers, - initial_sessions: self.sessions, - session_handler_resolver: self.session_handler_resolver, - session_handlers: Arc::new(parking_lot::RwLock::new(HashMap::new())), - session_unserved: parking_lot::RwLock::new(HashMap::new()), - session_resolve_errors: parking_lot::RwLock::new(HashMap::new()), - binary_version: self.binary_version, - notification_fwd: Arc::new(parking_lot::Mutex::new(None)), - parsed_notif_tx: Arc::new(parking_lot::Mutex::new(None)), - session_handles: Arc::new(parking_lot::Mutex::new(HashMap::new())), - active_sessions, - dynamic_tool_mu: tokio::sync::Mutex::new(()), - session_bind_mu: tokio::sync::Mutex::new(()), - reconnect_notify, - on_reconnect_settled: self.on_reconnect_settled, - disconnect_epoch, - admission, - cancels: Arc::new(DashMap::new()), - donation_pumps: parking_lot::Mutex::new(DonationPumps::default()), - }); - Ok(ToolServer { inner: Some(inner) }) - } -} - -/// Running tool-server attached to a pooled [`HubConnection`]. -/// -/// `Clone` is an `Arc` bump. Prefer [`Self::shutdown`]; [`Drop`] tears down -/// only for the last strong owner. Use [`Self::downgrade`] for observers. -pub struct ToolServer { - /// `Option` so [`Drop`] can take the `Arc` for [`Arc::into_inner`]. - inner: Option>, -} - -/// Non-owning handle to a [`ToolServer`]; [`Self::upgrade`] per use. -#[derive(Clone, Default)] -pub struct WeakToolServer { - inner: std::sync::Weak, -} - -impl WeakToolServer { - pub fn upgrade(&self) -> Option { - self.inner - .upgrade() - .map(|inner| ToolServer { inner: Some(inner) }) - } -} - -struct ToolServerInner { - borrow: ConnectionBorrow, - /// Handlers passed to the builder — cloned into each new session. - initial_handlers: Vec>, - /// Sessions passed to the builder. `run()` binds these via - /// `bind_session_local` so tools are registered and inboxes are - /// created before the dispatch loop starts. - initial_sessions: Vec, - session_handler_resolver: Option, - /// Per-session handler maps. Each session owns its own handler vec. - session_handlers: SessionHandlerMap, - /// Per-session unserved tool ids from the last resolver run; same - /// lifetime as the `session_handlers` entry. - session_unserved: parking_lot::RwLock>>, - /// Per-session fail-closed resolve reason from the last resolver run; - /// same lifetime as the `session_handlers` entry. - session_resolve_errors: parking_lot::RwLock>, - binary_version: Option, - - /// Raw notification forwarding channel. Session loops write here; - /// the parsing task (spawned in `run()`) reads and parses into - /// `HubNotification` events sent to `parsed_notif_tx`. - notification_fwd: Arc>>>, - /// Parsed notification sender. Set by `subscribe_notifications`; - /// the parsing task in `run()` bridges `notification_fwd` → this. - parsed_notif_tx: - Arc>>>, - /// Session loops spawned by `bind_session_local`, keyed by session ID. - session_handles: Arc>>>, - /// All sessions currently active on this server. - active_sessions: parking_lot::Mutex>, - /// Serializes `register_tool_dynamic` / `unregister_tool_dynamic`. - /// `tokio::sync::Mutex` because the critical section spans `.await`. - dynamic_tool_mu: tokio::sync::Mutex<()>, - /// Serializes session binds against each other and against - /// `unbind_session`, so the soft-rebind liveness decision and the - /// destructive full-rebind setup are atomic (no check-then-act race - /// between two concurrent binds, and no unbind sneaking between a - /// bind's liveness check and its return). Binds/unbinds are rare - /// lifecycle events; a global mutex is contention-free in practice. - session_bind_mu: tokio::sync::Mutex<()>, - /// Signalled by the on_reconnect callback so `run()` can replay - /// `serve` for every active session after a reconnect. - reconnect_notify: Arc, - /// Optional readiness / settle hook after serve replay (see - /// [`ToolServerBuilder::on_reconnect_settled`]). - on_reconnect_settled: Option>, - /// Bumped on every disconnect. The reconnect task snapshots this before - /// `serve` replay and only fires `on_reconnect_settled` if it is unchanged - /// afterward — so a disconnect racing the (async) replay cannot resurrect a - /// stale ready marker while the socket is already down. - disconnect_epoch: Arc, - /// Three-tier (session/connection/global) admission controller. Drives - /// the bounded-wait-then-overloaded backpressure on the spawned path. - admission: Arc, - /// Per-session strict-cancellation registries. Created in - /// `bind_session_local` alongside the inbox + admission semaphore, and - /// drained-and-cancelled on `unbind_session` / `shutdown` so detached - /// `execute_call` tasks wind down promptly. - cancels: Arc>>, - /// Trace/log/metric donation pump senders, fenced by - /// [`ToolServer::flush_donations`] on unbind/shutdown. - donation_pumps: parking_lot::Mutex, -} - -/// The three symmetric donation pump senders. Each is fenced -/// independently by `flush_donations_inner` so a teardown never abandons -/// a queued batch. -#[derive(Default)] -struct DonationPumps { - traces: Option>, - logs: Option>, - metrics: Option>, -} - -impl Clone for ToolServer { - fn clone(&self) -> Self { - Self { - inner: self.inner.as_ref().map(Arc::clone), - } - } -} - -impl std::fmt::Debug for ToolServer { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let inner = self.inner(); - f.debug_struct("ToolServer") - .field("active_sessions", &*inner.active_sessions.lock()) - .field("handler_count", &inner.initial_handlers.len()) - .finish_non_exhaustive() - } -} - -impl ToolServer { - fn inner(&self) -> &Arc { - self.inner - .as_ref() - .expect("ToolServer used after Drop took the Arc") - } - - pub fn downgrade(&self) -> WeakToolServer { - WeakToolServer { - inner: Arc::downgrade(self.inner()), - } - } - - /// Underlying connection. Useful for tests that need to assert - /// pool dedup. - pub fn connection(&self) -> &Arc { - self.inner().borrow.connection() - } - - /// Snapshot of all currently active sessions. - pub fn active_sessions(&self) -> Vec { - self.inner().active_sessions.lock().clone() - } - - /// Handlers for a specific session. - pub fn handlers_for_session(&self, session_id: &SessionId) -> Vec> { - self.inner() - .session_handlers - .read() - .get(session_id) - .cloned() - .unwrap_or_default() - } - - /// Unserved tool ids reported by the resolver for a session's last bind. - pub fn unserved_for_session(&self, session_id: &SessionId) -> Vec { - self.inner() - .session_unserved - .read() - .get(session_id) - .cloned() - .unwrap_or_default() - } - - /// Fail-closed resolve reason reported by the resolver for a session's - /// last bind, if any. - pub fn resolve_error_for_session(&self, session_id: &SessionId) -> Option { - self.inner() - .session_resolve_errors - .read() - .get(session_id) - .cloned() - } - - /// Bind a new session in two steps: - /// - /// 1. **Local setup** (`bind_session_local`): register the session - /// on the connection, create the demux inbox, and spawn the - /// per-session dispatch loop. - /// 2. **Publish** (`serve`): send a `serve` frame to the server with - /// the full tool snapshot so harnesses see the tools immediately. - pub async fn bind_session(&self, session_id: SessionId) -> Result<(), ClientError> { - self.bind_session_local(session_id.clone()).await?; - self.serve(session_id).await - } - - /// Register a session on the connection, create the demux inbox, - /// and spawn the per-session dispatch loop. - /// - /// After binding, the caller should send a `serve` frame to - /// publish the tool snapshot. - pub async fn bind_session_local(&self, session_id: SessionId) -> Result<(), ClientError> { - self.bind_session_local_with_metadata(session_id, None) - .await - } - - /// [`bind_session_local`] with the raw `session.bind` params for an - /// installed [`SessionHandlerResolver`]. - /// - /// Non-destructive for a live session (**soft rebind**): when the - /// session's dispatch loop is still running, a repeated bind only - /// refreshes serve state (resolver re-run → `session_handlers` / - /// `session_unserved`) and leaves the inbox, cancel registry, and - /// dispatch loop untouched, so in-flight tool calls survive. The - /// full destructive setup (new inbox + registry + loop, cancelling - /// anything stale) only runs when the previous loop is dead or the - /// session was never bound. - /// - /// Handler-set caveat: a soft rebind re-runs the resolver, so the - /// resolver must tolerate re-execution while handler instances from a - /// previous bind may still be mid-call; subsequent hook frames - /// (pause/resume/cancel fan-out) target the refreshed handler set, - /// while in-flight calls keep their pre-rebind handler clones and the - /// shared cancel registry. - pub async fn bind_session_local_with_metadata( - &self, - session_id: SessionId, - bind_params: Option, - ) -> Result<(), ClientError> { - // Serialized against other binds and `unbind_session` so the - // liveness decision below cannot race a concurrent bind/teardown. - let _bind_guard = self.inner().session_bind_mu.lock().await; - let connection = self.inner().borrow.connection(); - let sid = session_id; - - // Resolve handlers BEFORE mutating any session state so a resolver - // failure fails the bind cleanly (nothing tracked, no inbox, no - // session loop) and a retry re-runs from scratch. - let resolved = match &self.inner().session_handler_resolver { - // Re-run on every bind — including a soft rebind of a live - // session — so a retry after a failed bind recreates the - // session (resolver-owned) and refreshes advertised tools. - Some(resolver) => Some( - resolver(sid.clone(), bind_params) - .await - .map_err(|err| ClientError::Wire(ToolErrorWire::from(err)))?, - ), - None => None, - }; - - connection.track_session(sid.clone()); - { - let mut sessions = self.inner().active_sessions.lock(); - if !sessions.contains(&sid) { - sessions.push(sid.clone()); - } - } - - match resolved { - Some(resolved) => { - self.inner() - .session_handlers - .write() - .insert(sid.clone(), resolved.handlers); - self.inner() - .session_unserved - .write() - .insert(sid.clone(), resolved.unserved_tool_ids); - { - let mut errors = self.inner().session_resolve_errors.write(); - match resolved.resolve_error { - Some(reason) => { - errors.insert(sid.clone(), reason); - } - None => { - errors.remove(&sid); - } - } - } - } - None => { - self.inner() - .session_handlers - .write() - .entry(sid.clone()) - .or_insert_with(|| self.inner().initial_handlers.clone()); - } - } - - // Soft rebind: a live dispatch loop means this is a redundant - // `session.bind` for a healthy session — refresh serve state only - // (done above; the server's bind response reads `handlers_for_session`). - // Replacing the inbox/registry/loop would cancel every in-flight - // tool call. `session_bind_mu` keeps a concurrent bind/unbind from - // invalidating this check before we return; the registry check - // additionally rejects a loop in its post-teardown exit tail - // (`cancel_all` closes the registry before the JoinHandle finishes), - // making the gate best-effort-safe against non-serialized teardown. - // A dead loop falls through to the full (destructive) rebind, which - // preserves stale-token cleanup. - { - let loop_alive = { - let handles = self.inner().session_handles.lock(); - handles.get(&sid).is_some_and(|h| !h.is_finished()) - }; - let registry_open = self - .inner() - .cancels - .get(&sid) - .is_some_and(|r| !r.is_closed()); - if loop_alive && registry_open { - crate::metrics::session_soft_rebind(); - tracing::debug!( - %sid, - "bind: session loop alive — soft rebind (serve state refreshed, \ - in-flight calls preserved)" - ); - return Ok(()); - } - } - - let (tx, rx) = mpsc::channel(SESSION_INBOX_BUFFER); - connection.demux().register_session_inbox(sid.clone(), tx); - // Tie the per-session admission semaphore to the session-loop - // lifetime (created here, removed on unbind / loop exit) so a - // straggler call cannot recreate a leaked entry after teardown. - self.inner().admission.ensure_session(&sid); - // Per-session cancellation registry, same lifetime as the loop. - // A full rebind (previous loop dead — the live-loop case returned - // above) drains-and-cancels any stale registry before installing - // a fresh one so tokens from a previous loop never linger. - let cancels = Arc::new(CancelRegistry::default()); - if let Some(old) = self.inner().cancels.insert(sid.clone(), cancels.clone()) { - old.cancel_all(); - } - - let sh = self.inner().session_handlers.clone(); - let notification_fwd = self.inner().notification_fwd.clone(); - let conn = connection.clone(); - let loop_sid = sid.clone(); - let admission = self.inner().admission.clone(); - let cancels_owner = self.inner().cancels.clone(); - let handle = tokio::spawn(async move { - run_session_loop( - loop_sid, - rx, - conn, - sh, - notification_fwd, - admission, - cancels, - cancels_owner, - ) - .await; - }); - // Replace the previous (dead) loop's handle, if any. Abort is a - // no-op for a finished handle; with binds serialized under - // `session_bind_mu` no concurrent full rebind can have installed a - // live handle in between, so this never kills live work. - if let Some(old_handle) = self.inner().session_handles.lock().insert(sid, handle) { - old_handle.abort(); - } - - Ok(()) - } - - /// Send a `serve` frame for a session, publishing the full tool - /// snapshot. Idempotent — the server diffs and emits `tools_changed`. - /// - /// On reconnect, call `serve()` per active session to replay state. - pub async fn serve(&self, session_id: SessionId) -> Result<(), ClientError> { - // Build tool descriptions while holding the read lock so the - // handler list cannot mutate between read and serialization. - let tools: Vec = { - let map = self.inner().session_handlers.read(); - let handlers = map.get(&session_id); - handlers - .map(|h| { - h.iter() - .map(|h| kigi_tool_protocol::ToolDescriptionWithSchema { - description: h.description(), - input_schema: h.input_schema(), - capabilities: None, - notification_schemas: None, - }) - .collect() - }) - .unwrap_or_default() - }; - - let params = kigi_tool_protocol::ServeParams { tools }; - - let connection = self.inner().borrow.connection(); - connection.serve(session_id, params).await?; - Ok(()) - } - - /// Register a tool handler at runtime for the given sessions. - /// - /// Adds the handler to the local session handler map and replays - /// `serve` for each affected session so the server sees the updated - /// tool set. - pub async fn register_tool_dynamic( - &self, - handler: Arc, - sessions: Vec, - ) -> Result<(), ClientError> { - let _guard = self.inner().dynamic_tool_mu.lock().await; - - let tool_id = handler.tool_id(); - - // Reject duplicates within the target sessions. - { - let map = self.inner().session_handlers.read(); - for sid in &sessions { - if let Some(handlers) = map.get(sid) - && handlers.iter().any(|h| h.tool_id() == tool_id) - { - return Err(ClientError::InvalidConfig(format!( - "tool_id {tool_id} is already registered for session {sid}" - ))); - } - } - } - - // Insert into each target session's handler list. - { - let mut map = self.inner().session_handlers.write(); - for sid in &sessions { - map.entry(sid.clone()).or_default().push(handler.clone()); - } - } - - // Replay `serve` for each affected session so the server sees - // the updated tool set. - for sid in sessions { - self.serve(sid).await?; - } - - Ok(()) - } - - /// Remove a dynamically registered tool from a specific session. - /// Returns `Ok(false)` if not found. - pub async fn unregister_tool_dynamic( - &self, - tool_id: &ToolId, - session_id: &SessionId, - ) -> Result { - let _guard = self.inner().dynamic_tool_mu.lock().await; - - // Check existence in the target session. - { - let map = self.inner().session_handlers.read(); - let found = map - .get(session_id) - .is_some_and(|h| h.iter().any(|h| h.tool_id() == *tool_id)); - if !found { - return Ok(false); - } - } - - // Remove from the session's handler list. - { - let mut map = self.inner().session_handlers.write(); - if let Some(handlers) = map.get_mut(session_id) { - handlers.retain(|h| h.tool_id() != *tool_id); - } - } - - // Replay `serve` for the affected session so the server sees the - // tool removal. - self.serve(session_id.clone()).await?; - - Ok(true) - } - - /// Unbind a session: tear down the session loop, remove handlers, - /// and unregister the session. - pub async fn unbind_session(&self, session_id: &SessionId) -> Result<(), ClientError> { - // Serialized against binds (see `session_bind_mu`): an unbind must - // not interleave with a bind's liveness check / setup. - let _bind_guard = self.inner().session_bind_mu.lock().await; - let connection = self.inner().borrow.connection(); - - self.inner() - .active_sessions - .lock() - .retain(|s| s != session_id); - - self.inner().session_handlers.write().remove(session_id); - self.inner().session_unserved.write().remove(session_id); - self.inner() - .session_resolve_errors - .write() - .remove(session_id); - - // Teardown ordering: drain-and-cancel every in-flight - // call's token FIRST so detached `execute_call` tasks wind down via - // their `select!`, THEN abort the dispatcher and remove the inbox / - // admission / registry entries. - if let Some((_, registry)) = self.inner().cancels.remove(session_id) { - registry.cancel_all(); - } - - if let Some(handle) = self.inner().session_handles.lock().remove(session_id) { - handle.abort(); - } - - connection.demux().unregister_session_inbox(session_id); - self.inner().admission.remove_session(session_id); - - connection.untrack_session(session_id); - - Ok(()) - } - - /// Subscribe to server notifications across all bound sessions. - /// - /// The server auto-subscribes harness sessions on `register_session`, - /// so no wire request is needed. Call before `run()` — the - /// parsing task is spawned inside `run()`. - pub fn subscribe_notifications(&self) -> mpsc::Receiver { - let (event_tx, event_rx) = mpsc::channel::(64); - *self.inner().parsed_notif_tx.lock() = Some(event_tx); - event_rx - } - - /// Send a `tool.notify` frame to the server. - /// - /// Mirrors [`ToolHarness::send_notification`] but over a - /// `tool_server` connection. The server allows `tool.notify` for both - /// `Harness` and `ToolServer` connection kinds. - /// - /// The frame is fire-and-forget: this method returns `Ok` once the - /// outbound message is queued, without waiting for a server ack. - pub async fn send_notification( - &self, - notification: kigi_tool_protocol::ToolNotificationFrame, - ) -> Result<(), ClientError> { - let session = self - .inner() - .active_sessions - .lock() - .first() - .cloned() - .ok_or_else(|| { - ClientError::InvalidConfig( - "send_notification requires at least one bound session".to_owned(), - ) - })?; - let connection = self.inner().borrow.connection(); - let request_id = connection.try_alloc_request_id()?; - let req = kigi_tool_protocol::JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: Some(session), - method: Method::ToolNotify.as_wire_str().to_owned(), - params: notification, - }; - let text = serde_json::to_string(&req).map_err(ClientError::from)?; - connection.send_outbound(text).await - } - - /// Send a `system.notify` frame scoped to an explicit session and await the - /// server's ack (a JSON-RPC request, unlike the fire-and-forget `send_notification`). - pub async fn send_system_notification( - &self, - session_id: SessionId, - params: kigi_tool_protocol::SystemNotifyParams, - ) -> Result { - // Fail fast on an oversized payload instead of round-tripping to the server. - let payload_len = json_serialized_len(¶ms.payload)?; - if payload_len > kigi_tool_protocol::MAX_SYSTEM_NOTIFY_PAYLOAD_BYTES { - return Err(ClientError::ProtocolError(format!( - "system.notify payload {payload_len} bytes exceeds {} byte cap", - kigi_tool_protocol::MAX_SYSTEM_NOTIFY_PAYLOAD_BYTES - ))); - } - let connection = self.inner().borrow.connection(); - let request_id = connection.try_alloc_request_id()?; - let req = kigi_tool_protocol::JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: Some(session_id), - method: Method::SystemNotify.as_wire_str().to_owned(), - params, - }; - let resp = connection.call_request(request_id, &req).await?; - system_notify_ack_from_outcome(resp.outcome) - } - - /// Backstop deadline for [`Self::request_hook`]. - /// - /// Deliberately long: the hook is normally released by the real reply - /// or requester teardown; this only bounds a request whose reply is - /// lost (e.g. a dead connection), so it sits above any turn deadline. - pub const HOOK_REQUEST_BACKSTOP_TIMEOUT: std::time::Duration = - std::time::Duration::from_secs(600); - - /// Send a request/response `Custom` hook on the shared connection and await its reply. - /// - /// The tool-server counterpart to the harness requester: the server - /// originates the hook and the bound harness answers it. Bounded by - /// [`HOOK_REQUEST_BACKSTOP_TIMEOUT`](Self::HOOK_REQUEST_BACKSTOP_TIMEOUT); - /// use [`Self::request_hook_with_timeout`] for a different deadline. - /// - /// Only `permission_request` is answered by the bound harness today; any - /// other `kind` is dropped by the responder and the call resolves only when - /// the backstop timeout fires. Callers must pass a supported `kind`. - pub async fn request_hook( - &self, - session_id: SessionId, - kind: String, - payload: Value, - ) -> Result { - self.request_hook_with_timeout( - session_id, - kind, - payload, - Self::HOOK_REQUEST_BACKSTOP_TIMEOUT, - ) - .await - } - - /// [`Self::request_hook`] with a caller-supplied backstop deadline. - pub async fn request_hook_with_timeout( - &self, - session_id: SessionId, - kind: String, - payload: Value, - timeout: std::time::Duration, - ) -> Result { - let connection = self.inner().borrow.connection(); - // hook_id keys the server's parked-request table — must be globally unique. - let hook_id = ToolCallId::new_v7().to_string(); - let hook = HookFrame::custom_request(session_id.clone(), hook_id, kind, payload); - let request_id = connection.try_alloc_request_id()?; - let req = kigi_tool_protocol::JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: Some(session_id), - method: Method::Hook.as_wire_str().to_owned(), - params: hook, - }; - let resp = connection - .call_request_with_timeout(request_id, &req, timeout) - .await?; - match resp.outcome { - ResponseOutcome::Result(value) => Ok(value), - ResponseOutcome::Error(err) => Err(ClientError::from_jsonrpc_error(err)), - } - } - - /// Fire-and-forget `traces.donate`: `Ok` = queued; rejects surface - /// only in server metrics. `otlp_request_b64` is a base64 - /// `ExportTraceServiceRequest` with a server-allowlisted `service.name`. - pub async fn donate_traces(&self, otlp_request_b64: &str) -> Result<(), ClientError> { - /// Borrowed wire shape so retried payloads are never cloned. - #[derive(serde::Serialize)] - struct ParamsRef<'a> { - otlp_request: &'a str, - } - let session = self - .inner() - .active_sessions - .lock() - .first() - .cloned() - .ok_or_else(|| { - ClientError::InvalidConfig( - "donate_traces requires at least one bound session".to_owned(), - ) - })?; - let connection = self.inner().borrow.connection(); - let notification = kigi_tool_protocol::JsonRpcNotification { - jsonrpc: JsonRpcVersion, - session_id: Some(session), - seq: None, - method: Method::TracesDonate.as_wire_str().to_owned(), - params: ParamsRef { - otlp_request: otlp_request_b64, - }, - }; - let text = serde_json::to_string(¬ification).map_err(ClientError::from)?; - connection.send_outbound(text).await - } - - /// Fire-and-forget `logs.donate`: `Ok` = queued; rejects surface - /// only in server metrics. `otlp_request_b64` is a base64 - /// `ExportLogsServiceRequest` with a server-allowlisted `service.name`. - /// Requires a bound session (mirrors [`Self::donate_traces`]). - pub async fn donate_logs(&self, otlp_request_b64: &str) -> Result<(), ClientError> { - /// Borrowed wire shape so retried payloads are never cloned. - #[derive(serde::Serialize)] - struct ParamsRef<'a> { - otlp_request: &'a str, - } - let session = self - .inner() - .active_sessions - .lock() - .first() - .cloned() - .ok_or_else(|| { - ClientError::InvalidConfig( - "donate_logs requires at least one bound session".to_owned(), - ) - })?; - let connection = self.inner().borrow.connection(); - let notification = kigi_tool_protocol::JsonRpcNotification { - jsonrpc: JsonRpcVersion, - session_id: Some(session), - seq: None, - method: Method::LogsDonate.as_wire_str().to_owned(), - params: ParamsRef { - otlp_request: otlp_request_b64, - }, - }; - let text = serde_json::to_string(¬ification).map_err(ClientError::from)?; - connection.send_outbound(text).await - } - - /// Fire-and-forget `metrics.donate`: `Ok` = queued; rejects surface - /// only in server metrics. `otlp_request_b64` is a base64 - /// `ExportMetricsServiceRequest` with a server-allowlisted `service.name`. - /// Unlike [`Self::donate_logs`], metrics are process-aggregate, so - /// this does **not** require a bound session. - pub async fn donate_metrics(&self, otlp_request_b64: &str) -> Result<(), ClientError> { - /// Borrowed wire shape so retried payloads are never cloned. - #[derive(serde::Serialize)] - struct ParamsRef<'a> { - otlp_request: &'a str, - } - let connection = self.inner().borrow.connection(); - let notification = kigi_tool_protocol::JsonRpcNotification { - jsonrpc: JsonRpcVersion, - session_id: None, - seq: None, - method: Method::MetricsDonate.as_wire_str().to_owned(), - params: ParamsRef { - otlp_request: otlp_request_b64, - }, - }; - let text = serde_json::to_string(¬ification).map_err(ClientError::from)?; - connection.send_outbound(text).await - } - - /// Drive the inbound loop until either the connection actor - /// signals shutdown OR [`Self::shutdown`] is called. - /// - /// Each bound session gets its own per-session task that pulls - /// frames from the connection's demux and dispatches them to the - /// matching handler by `tool_id`. - pub async fn run(&self) -> Result<(), ClientError> { - let connection = self.inner().borrow.connection().clone(); - let demux = connection.demux(); - - for sid in &self.inner().initial_sessions { - if let Err(e) = self.bind_session_local(sid.clone()).await { - warn!(%sid, error = %e, "run: bind_session_local failed for builder session"); - continue; - } - // Publish the tool snapshot so the server registers the tools - // for this session. bind_session_local only does local setup. - if let Err(e) = self.serve(sid.clone()).await { - warn!(%sid, error = %e, "run: serve failed for builder session"); - } - } - - // If subscribe_notifications() was called before run(), - // install the forwarding channel and spawn the parsing task. - if let Some(event_tx) = self.inner().parsed_notif_tx.lock().take() { - let (fwd_tx, mut fwd_rx) = mpsc::channel::(64); - *self.inner().notification_fwd.lock() = Some(fwd_tx); - tokio::spawn(async move { - while let Some(value) = fwd_rx.recv().await { - if let Some(event) = crate::notification::HubNotification::parse(&value) - && event_tx.send(event).await.is_err() - { - break; - } - } - }); - } - - let mut notif_rx = connection - .take_early_notifications() - .unwrap_or_else(|| demux.subscribe_notifications()); - let buffered = notif_rx.len(); - if buffered > 0 { - crate::metrics::early_notif_buffered(buffered as u64); - tracing::info!( - buffered, - "replaying connection-level notifications buffered before run()" - ); - } - // Wrap in Arc so spawned per-bind tasks hold Arc::clone - // (refcount bump) instead of ToolServer::clone(). A - // ToolServer::clone() going out of scope triggers - // Drop::begin_teardown() on the shared AtomicBool, which - // tears down *all* sessions as soon as the first spawned - // bind task completes. - let server_for_notif = Arc::new(self.clone()); - let connection_for_notif = connection.clone(); - let notif_handle = tokio::spawn(async move { - loop { - let frame = match notif_rx.recv().await { - Ok(frame) => frame, - Err(RecvError::Lagged(skipped)) => { - crate::metrics::notif_lagged_recovered(); - tracing::warn!( - skipped, - "connection notification stream lagged; continuing" - ); - continue; - } - Err(RecvError::Closed) => break, - }; - let method = frame - .get("method") - .and_then(serde_json::Value::as_str) - .unwrap_or(""); - let request_id = frame.get("id").filter(|v| !v.is_null()).cloned(); - - match method { - "session.bind" => { - let Some(sid_str) = frame - .pointer("/params/session_id") - .and_then(serde_json::Value::as_str) - else { - continue; - }; - let Ok(sid) = kigi_tool_protocol::SessionId::new(sid_str) else { - continue; - }; - tracing::info!(%sid, "session.bind: binding new session"); - - let bind_params = frame.get("params").cloned(); - let server = Arc::clone(&server_for_notif); - let conn = connection_for_notif.clone(); - tokio::spawn(async move { - let result = server - .bind_session_local_with_metadata(sid.clone(), bind_params) - .await; - - // Respond with tools on success, error on failure. - // The server registers tools from this response directly - // (v2 protocol — no separate serve RPC needed). - if let Some(id) = request_id { - let response = match result { - Ok(()) => { - let tools: Vec = server - .handlers_for_session(&sid) - .iter() - .map(|h| h.description()) - .collect(); - let result = kigi_tool_protocol::SessionBindResult { - tools, - binary_version: server.inner().binary_version.clone(), - unserved_tool_ids: server.unserved_for_session(&sid), - resolve_error: server.resolve_error_for_session(&sid), - }; - serde_json::json!({ - "jsonrpc": "2.0", - "id": id, - "result": result - }) - } - Err(ref e) => { - tracing::warn!( - error = %e, session = %sid, - "session.bind: bind_session_local failed" - ); - // Resolver failures carry a decodable - // ToolErrorWire; forward its numeric - // code + payload so the cause survives - // past the server instead of collapsing - // to a bare -32603. - let (code, data) = match e { - ClientError::Wire(wire) => ( - error_codes::from_tool_error_wire(wire), - serde_json::to_value(wire).ok(), - ), - _ => (-32603, None), - }; - serde_json::json!({ - "jsonrpc": "2.0", - "id": id, - "error": { - "code": code, - "message": format!("bind failed: {e}"), - "data": data - } - }) - } - }; - if let Ok(text) = serde_json::to_string(&response) - && let Err(e) = conn.send_outbound(text).await - { - tracing::warn!( - error = %e, - "failed to send session.bind response" - ); - } - } - }); - } - "session.unbind" => { - let Some(sid_str) = frame - .pointer("/params/session_id") - .and_then(serde_json::Value::as_str) - else { - continue; - }; - let Ok(sid) = kigi_tool_protocol::SessionId::new(sid_str) else { - continue; - }; - tracing::info!(%sid, "session.unbind: unbinding session"); - - let server = server_for_notif.clone(); - let conn = connection_for_notif.clone(); - tokio::spawn(async move { - // Unbind precedes hibernate: flush while up. - server.flush_donations().await; - let _ = server.unbind_session(&sid).await; - - // Respond so session.close returns synchronously. - if let Some(id) = request_id { - let response = serde_json::json!({ - "jsonrpc": "2.0", - "id": id, - "result": {} - }); - if let Ok(text) = serde_json::to_string(&response) - && let Err(e) = conn.send_outbound(text).await - { - tracing::warn!( - error = %e, - "failed to send session.unbind response" - ); - } - } - }); - } - // Server-issued graceful-shutdown request, fanned out to the - // evicted session's handlers (mirroring the hook fan-out). - "tool_server.evict" => { - let Some(params) = frame.get("params") else { - continue; - }; - // Deserialize straight from the borrowed `&Value` - // (`&Value: Deserializer`) — no need to clone the params. - let evict: ToolServerEvictParams = - match ::deserialize(params) - { - Ok(p) => p, - Err(e) => { - tracing::warn!( - error = %e, - "tool_server.evict: failed to decode params" - ); - continue; - } - }; - tracing::info!( - session = %evict.session_id, - reason = %evict.reason, - grace_period_ms = evict.grace_period_ms, - "tool_server.evict received; draining" - ); - let server = Arc::clone(&server_for_notif); - let conn = connection_for_notif.clone(); - tokio::spawn(async move { - // Ack first (best-effort) so the server's request - // resolves promptly; the drain runs in the background. - if let Some(id) = request_id { - let response = serde_json::json!({ - "jsonrpc": "2.0", - "id": id, - "result": {} - }); - if let Ok(text) = serde_json::to_string(&response) - && let Err(e) = conn.send_outbound(text).await - { - tracing::warn!( - error = %e, - "failed to send tool_server.evict ack" - ); - } - } - for handler in server.handlers_for_session(&evict.session_id) { - handler.handle_evict(evict.clone()).await; - } - }); - } - _ => {} - } - } - }); - - // Spawn a task that replays `serve` for every active session - // after a reconnect. The on_reconnect callback (sync) signals - // via Notify; this async task picks up the event and does - // the actual serve calls, then fires on_reconnect_settled so - // readiness markers can wait until tools are re-served. - let server_for_reconnect = Arc::new(self.clone()); - let reconnect_handle = { - let server = server_for_reconnect; - let notify = Arc::clone(&self.inner().reconnect_notify); - let settled = self.inner().on_reconnect_settled.clone(); - let epoch = Arc::clone(&self.inner().disconnect_epoch); - tokio::spawn(async move { - loop { - notify.notified().await; - // Snapshot the disconnect epoch for the connection we are - // now replaying onto; if it advances during replay a fresh - // disconnect raced us and `settled` must not fire (it would - // resurrect a stale ready marker over a downed socket). - let epoch_at_start = epoch.load(Ordering::Acquire); - let sessions: Vec = server.active_sessions(); - tracing::info!( - sessions = sessions.len(), - "reconnect: replaying serve for active sessions" - ); - let mut all_served = true; - for sid in sessions { - if let Err(e) = server.serve(sid.clone()).await { - all_served = false; - tracing::warn!( - error = %e, - session = %sid, - "reconnect: serve replay failed" - ); - } - } - // Only settle when every session re-served AND no disconnect - // raced this replay — otherwise the next reconnect's notify - // re-runs this loop and settles then. - let raced = epoch.load(Ordering::Acquire) != epoch_at_start; - if !all_served || raced { - tracing::info!( - all_served, - raced, - "reconnect: not settling ready (replay incomplete or disconnect raced)" - ); - continue; - } - if let Some(ref cb) = settled { - cb(); - } - } - }) - }; - - let shutdown = self.inner().borrow.shutdown_token().clone(); - tokio::select! { - biased; - _ = shutdown.cancelled() => {} - _ = connection.await_shutdown() => {} - } - notif_handle.abort(); - reconnect_handle.abort(); - for (_, handle) in self.inner().session_handles.lock().drain() { - handle.abort(); - } - Ok(()) - } - - /// Cooperatively shut the server down: signal `run` to return, - /// unregister each session binding (refcount-aware), and unbind - /// each registered tool. - /// - /// Errors during teardown are aggregated rather than short- - /// circuiting so a partial cleanup still releases everything it - /// can. The first error (if any) is returned. - pub async fn shutdown(&self) -> Result<(), ClientError> { - // Mark torn_down BEFORE the cleanup so the Drop fallback - // doesn't double-schedule. - if !self.inner().borrow.begin_teardown() { - return Ok(()); - } - // Real teardown: flush AND clear the pumps to break the reference cycle. - flush_donations_inner(self.inner().as_ref(), true).await; - teardown_sessions(self.inner().as_ref()).await; - Ok(()) - } - - pub(crate) fn set_donation_pump(&self, tx: mpsc::Sender) { - self.inner().donation_pumps.lock().traces = Some(tx); - } - - pub(crate) fn set_log_donation_pump(&self, tx: mpsc::Sender) { - self.inner().donation_pumps.lock().logs = Some(tx); - } - - #[cfg(feature = "metrics")] - pub(crate) fn set_metric_donation_pump(&self, tx: mpsc::Sender) { - self.inner().donation_pumps.lock().metrics = Some(tx); - } - - /// Clone of the connection's shutdown token so the periodic metric - /// reporter (the only perpetually-running donation task) can stop on - /// teardown instead of gathering and sending forever. - #[cfg(feature = "metrics")] - pub(crate) fn shutdown_token(&self) -> CancellationToken { - self.inner().borrow.shutdown_token().clone() - } - - /// Flush each producer and fence its donation pump; no-op without a - /// pump. Drives all three signals so a teardown never abandons a batch. - pub async fn flush_donations(&self) { - flush_donations_inner(self.inner().as_ref(), false).await; - } -} - -/// Fence each donation pump. `clear_pumps` only from `shutdown` / `Drop`. -async fn flush_donations_inner(inner: &ToolServerInner, clear_pumps: bool) { - let (traces, logs, metrics) = { - let pumps = inner.donation_pumps.lock(); - ( - pumps.traces.clone(), - pumps.logs.clone(), - pumps.metrics.clone(), - ) - }; - if let Some(tx) = traces { - fastrace::flush(); - crate::donate_pump::drain_via(&tx).await; - } - if let Some(tx) = logs { - crate::log_donate::flush_log_layer(); - crate::donate_pump::drain_via(&tx).await; - } - if let Some(tx) = metrics { - #[cfg(feature = "metrics")] - crate::metric_donate::gather_and_send(); - crate::donate_pump::drain_via(&tx).await; - } - - // Teardown only: drop pump senders so pump tasks exit. Keep them on - // while-running unbind flushes (`clear_pumps = false`). - if clear_pumps { - { - let mut pumps = inner.donation_pumps.lock(); - pumps.traces = None; - pumps.logs = None; - pumps.metrics = None; - } - #[cfg(feature = "metrics")] - crate::metric_donate::clear_active_exporter(); - } -} - -impl Drop for ToolServer { - fn drop(&mut self) { - // `into_inner` (not `try_unwrap`): exactly one concurrent dropper wins. - let Some(inner) = self.inner.take() else { - return; - }; - let Some(owned) = Arc::into_inner(inner) else { - return; - }; - if !owned.borrow.begin_teardown() { - return; - } - if tokio::runtime::Handle::try_current().is_ok() { - tokio::spawn(async move { - flush_donations_inner(&owned, true).await; - teardown_sessions(&owned).await; - }); - } - } -} - -/// Shared teardown for `shutdown` and `Drop`. Drain-and-cancel every -/// session's in-flight tokens BEFORE aborting the dispatchers, so detached -/// `execute_call` tasks wind down via their `select!` rather than being -/// orphaned, then unregister inboxes, drop admission entries, -/// and untrack sessions. Callers own the `begin_teardown` guard. -async fn teardown_sessions(inner: &ToolServerInner) { - let connection = inner.borrow.connection(); - let all_sessions: Vec = inner.active_sessions.lock().clone(); - - push_disconnect_status(connection, &all_sessions).await; - inner.borrow.shutdown_token().cancel(); - - for entry in inner.cancels.iter() { - entry.value().cancel_all(); - } - inner.cancels.clear(); - for (_, handle) in inner.session_handles.lock().drain() { - handle.abort(); - } - let demux = connection.demux(); - for sid in &all_sessions { - let _ = demux.unregister_session_inbox(sid); - inner.admission.remove_session(sid); - connection.untrack_session(sid); - } -} - -/// Push `tool_server.status(Disconnected)` for each bound session. -/// Must be called before `unregister_session` (the server needs the -/// session bindings to route the notification). -async fn push_disconnect_status(connection: &HubConnection, sessions: &[SessionId]) { - use kigi_tool_protocol::{JsonRpcRequest, ToolServerLifecycleStatus, ToolServerStatusPayload}; - - for sid in sessions { - let mut payload = - ToolServerStatusPayload::terminal(ToolServerLifecycleStatus::Disconnected); - payload.session_id = Some(sid.clone()); - let params = match serde_json::to_value(&payload) { - Ok(v) => v, - Err(_) => continue, - }; - let Ok(request_id) = connection.try_alloc_request_id() else { - continue; - }; - let req = JsonRpcRequest { - jsonrpc: JsonRpcVersion, - id: JsonRpcId::from_request_id(&request_id), - session_id: None, - method: Method::ToolServerStatus.as_wire_str().to_owned(), - params, - }; - let text = match serde_json::to_string(&req) { - Ok(t) => t, - Err(_) => continue, - }; - if let Err(e) = connection.send_outbound(text).await { - debug!(error = %e, session = %sid, "push_disconnect_status failed"); - } - } -} - -/// Per-session inbound dispatcher. -/// -/// Dequeues one frame at a time. Notifications (including `Cancel` -/// hooks) are handled **inline** because they are cheap and must never -/// queue behind a running call. A `tool_call_request` is dispatched to a -/// spawned task that performs three-tier admission before running, and -/// the loop immediately returns to `rx.recv()`, so calls within a session -/// run concurrently. -async fn run_session_loop( - session_id: SessionId, - mut rx: mpsc::Receiver, - connection: Arc, - session_handlers: SessionHandlerMap, - notification_fwd: Arc>>>, - admission: Arc, - cancels: Arc, - cancels_owner: Arc>>, -) { - while let Some(frame) = rx.recv().await { - let handlers = session_handlers - .read() - .get(&session_id) - .cloned() - .unwrap_or_default(); - match frame { - // Cheap, inline — keeps Cancel ahead of running calls. - InboundFrame::Notification(value) => { - handle_notification( - &session_id, - value, - &handlers, - ¬ification_fwd, - &cancels, - &connection, - ) - .await; - } - // Hot path: never await execution in the loop. Admission runs - // *inside* the spawned task (acquiring before spawn would - // head-of-line block the loop and Cancel hooks). - InboundFrame::Request(value) => { - // Register the cancellation token BEFORE spawn - // so an inline `Cancel` dequeued immediately after can never - // race ahead of registration and silently no-op. A pending - // tombstone (cancel-before-registration) pre-cancels here. - let token = CancellationToken::new(); - let call_id = parse_tool_call_id(&value); - if let Some(id) = &call_id { - cancels.register(id.clone(), &token); - } - let sid = session_id.clone(); - let conn = connection.clone(); - let adm = admission.clone(); - let cancels = cancels.clone(); - tokio::spawn(async move { - execute_call(&sid, value, &conn, &handlers, &adm, token).await; - // Always deregister on completion/cancel, - // regardless of which path `execute_call` returned by. - if let Some(id) = call_id { - cancels.deregister(&id); - } - }); - } - } - } - // Loop exited (inbox closed): the session is being torn down, so wind - // down this loop's own in-flight detached calls. Doing it here makes - // exit self-sufficient — it never depends on the rebind path's - // `insert -> old.cancel_all()` running (which can be skipped if this - // exit removes the entry first). `cancel_all` on an empty/already-closed - // registry is a safe no-op. - cancels.cancel_all(); - // Symmetric cleanup of BOTH per-session entries. `remove_if` guards - // against evicting a fresh registry that a concurrent rebind just - // installed (only this loop's own `Arc` is removed). Unbind/shutdown - // also remove these when they abort the loop. - admission.remove_session(&session_id); - cancels_owner.remove_if(&session_id, |_, registry| Arc::ptr_eq(registry, &cancels)); -} - -/// Extract the `params.tool_call_id` from a raw `tool_call_request` -/// frame so the dispatcher can register a cancellation token under it -/// before spawning the call. -fn parse_tool_call_id(value: &Value) -> Option { - value - .pointer("/params/tool_call_id") - .cloned() - .and_then(|v| serde_json::from_value(v).ok()) -} - -/// Dispatch an inbound notification frame. -/// -/// Hook frames (`method == "hook"`) are forwarded to every registered -/// handler whose `tool_id` matches the frame's target, or to every -/// handler when the hook is session-wide (no `tool_id`). -/// -/// All other notifications (e.g. `ToolsChanged`, `tool.notification`) -/// are forwarded to the `notification_fwd` channel if one has been -/// installed by [`ToolServer::subscribe_notifications`]. This lets -/// the session loop (which owns the demux inbox) coexist with the -/// notification subscriber without a registration conflict. -async fn handle_notification( - session_id: &SessionId, - value: Value, - handlers: &[Arc], - notification_fwd: &parking_lot::Mutex>>, - cancels: &CancelRegistry, - connection: &Arc, -) { - let method = value.get("method").and_then(Value::as_str).unwrap_or(""); - if method == Method::Hook.as_wire_str() { - let Some(params) = value.get("params").cloned() else { - warn!(%session_id, "hook notification missing params; ignoring"); - return; - }; - let frame: HookFrame = match serde_json::from_value(params) { - Ok(f) => f, - Err(err) => { - warn!(?err, %session_id, "hook notification failed to decode"); - return; - } - }; - let hook_span = frame - .trace_context - .as_deref() - .and_then(fastrace::collector::SpanContext::decode_w3c_traceparent) - .map(|parent| { - // Lets the server scope this span to its owning session. - fastrace::Span::root("tool_server.hook", parent) - .with_property(|| ("session_id", session_id.as_str().to_owned())) - }) - .unwrap_or_else(fastrace::Span::noop); - if let Some(hook_id) = frame.hook_id.clone() { - use fastrace::future::FutureExt as _; - // Spawned: a hook-request handler may legitimately take seconds - // (e.g. a handler that enqueues follow-up work for its `After` - // ack); inline it would head-of-line block this loop and Cancel - // hooks. Correlation rides `hook_id`, so ordering is not - // load-bearing. - let session_id = session_id.clone(); - let handlers = handlers.to_vec(); - let connection = connection.clone(); - tokio::spawn( - async move { - let mut result = None; - for handler in &handlers { - if let Some(r) = handler - .handle_hook_request(session_id.clone(), frame.clone()) - .await - { - result = Some(r); - break; - } - } - let notif = JsonRpcNotification { - jsonrpc: JsonRpcVersion, - session_id: Some(session_id.clone()), - seq: None, - method: Method::HookReply.as_wire_str().to_owned(), - params: HookReplyFrame { - session_id: session_id.clone(), - hook_id, - result: result.unwrap_or(Value::Null), - }, - }; - match serde_json::to_string(¬if) { - Ok(text) => { - let _ = connection.send_outbound(text).await; - } - Err(err) => warn!(?err, %session_id, "failed to encode hook_reply"), - } - } - .in_span(hook_span), - ); - return; - } - // Apply `Cancel` BEFORE the per-handler `handle_hook` fan-out - // — a slow user `handle_hook` must never delay the - // cancel. Compiler-enforced exhaustiveness keeps new HookEvent - // variants from silently skipping this dispatch. - match &frame.event { - HookEvent::Cancel => { - crate::metrics::cancel_hook_received(); - if let Some(call_id) = &frame.call_id { - if cancels.cancel(call_id) { - crate::metrics::cancel_applied(); - } else { - crate::metrics::cancel_pending_tombstoned(); - } - } else { - crate::metrics::cancel_no_target(); - } - } - HookEvent::Pause - | HookEvent::Resume - | HookEvent::SessionEnded - | HookEvent::Custom { .. } => {} - } - let frame_tool_id = frame.tool_id.clone(); - { - use fastrace::future::FutureExt as _; - async { - for handler in handlers { - if let Some(target) = &frame_tool_id - && handler.tool_id() != *target - { - continue; - } - handler.handle_hook(session_id.clone(), frame.clone()).await; - } - } - .in_span(hook_span) - .await; - } - return; - } - - // Forward to notification subscriber if one is registered. - if let Some(fwd) = notification_fwd.lock().as_ref() { - if fwd.try_send(value).is_err() { - warn!(%session_id, "notification forwarding channel full or closed"); - } - return; - } - debug!(?value, %session_id, method = %method, "tool-server received notification (no subscriber)"); -} - -/// Execute one `tool_call_request`: parse id/params, admit via the -/// [`Admission`](crate::admission::Admission) controller, locate the -/// handler, build the call context, drain the handler stream -/// (forwarding progress), build the JSON-RPC response, and ship it -/// back. Invoked from a spawned task by the dispatcher. -/// -/// Admission happens AFTER parsing id/params (so an overload can be -/// addressed to the request) but BEFORE invoking the handler. On -/// overload it emits the shared `-32016` "tool_busy" error — never a -/// silent drop — and the `AdmitGuard` holds all three permits for the -/// handler's lifetime, releasing them on return. -/// -/// `token` is the per-call cancellation handle registered by the -/// dispatcher before spawn. It is exposed to the tool via -/// the [`Cancellation`] extension and the handler-stream drain is wrapped -/// in a biased `select!` on it, so a `Cancel` hook hard-cancels by -/// dropping the call future and yields a `ToolError::Cancelled` response. -async fn execute_call( - session_id: &SessionId, - value: Value, - connection: &Arc, - handlers: &[Arc], - admission: &crate::admission::Admission, - token: CancellationToken, -) { - let id_value = match value.get("id").cloned() { - Some(v) => v, - None => { - warn!("tool_call_request missing id; ignoring"); - return; - } - }; - let json_id: JsonRpcId = match serde_json::from_value(id_value) { - Ok(v) => v, - Err(err) => { - warn!(?err, "tool_call_request id failed to decode"); - return; - } - }; - let params: ToolCallParams = match value - .get("params") - .cloned() - .and_then(|p| serde_json::from_value(p).ok()) - { - Some(p) => p, - None => { - send_error( - connection, - json_id, - session_id.clone(), - -32602, - "invalid params", - ) - .await; - return; - } - }; - let Some(handler) = handlers.iter().find(|h| h.tool_id() == params.tool_id) else { - crate::metrics::no_handler(); - send_error( - connection, - json_id, - session_id.clone(), - -32011, - &format!("no handler for tool_id {}", params.tool_id), - ) - .await; - return; - }; - - // Admission held for the handler's lifetime; overload -> -32016 reply. - let _guard = match admission.admit(session_id).await { - Ok(guard) => guard, - Err(crate::admission::Overloaded::Timeout) => { - crate::metrics::tool_call_rejected_overloaded(); - send_overloaded(connection, json_id, session_id.clone()).await; - return; - } - Err(crate::admission::Overloaded::Shutdown) => return, - }; - - let call_span = params - .trace_context - .as_deref() - .and_then(fastrace::collector::SpanContext::decode_w3c_traceparent) - .map(|parent| { - fastrace::Span::root("tool_server.tool_call", parent).with_properties(|| { - [ - ("tool_id", params.tool_id.as_str().to_owned()), - ("tool_call_id", params.tool_call_id.as_str().to_owned()), - // Lets the server scope this span to its owning session. - ("session_id", session_id.as_str().to_owned()), - ] - }) - }) - .unwrap_or_else(fastrace::Span::noop); - - let mut ctx = ToolCallContext::new(params.tool_call_id.clone()); - ctx.extensions.insert(kigi_tool_runtime::SessionContext( - session_id.as_str().to_owned(), - )); - if let Some(cwd) = params.cwd { - ctx.extensions.insert(Cwd(std::path::PathBuf::from(cwd))); - } - if let Some(version) = params.behavior_version { - ctx.extensions.insert(BehaviorVersion(version)); - } - if let Some(trace) = params.trace_context { - ctx.extensions.insert(TraceContext(trace)); - } - // Expose the cancellation handle so cooperative tools can poll/await - // it; the dispatcher still hard-cancels via the `select!` below. - ctx.extensions.insert(Cancellation(token.clone())); - - // Move `tool_id` (unused afterward) into the cancellation arm; clone - // only `tool_call_id`, which the response build still needs. - let tool_id = params.tool_id; - let tool_call_id = params.tool_call_id.clone(); - let arguments = params.arguments; - - // Drain the handler stream, forwarding progress and capturing the - // terminal. Wrapped in a biased `select!` on the cancellation token - // so a `Cancel` DROPS this future (hard cancel of the in-flight - // await) and yields a `ToolError::Cancelled` response. - let drain = async { - let mut stream = handler.handle_call(ctx, arguments).await; - let mut terminal: Option> = None; - while let Some(item) = stream.next().await { - match item { - ToolStreamItem::Progress(progress) => { - let frame = progress_to_frame(progress, tool_call_id.clone()); - let notif = JsonRpcNotification { - jsonrpc: JsonRpcVersion, - session_id: Some(session_id.clone()), - seq: None, - method: Method::ToolCallProgress.as_wire_str().to_owned(), - params: frame, - }; - if let Ok(text) = serde_json::to_string(¬if) { - crate::metrics::progress_frame_forwarded(); - let _ = connection.send_outbound(text).await; - } - } - ToolStreamItem::Terminal(result) => { - terminal = Some(result); - break; - } - } - } - terminal - }; - // Span duration = actual tool execution (the per-hop leaf). - let drain = { - use fastrace::future::FutureExt as _; - drain.in_span(call_span) - }; - let terminal: Option> = tokio::select! { - biased; - _ = token.cancelled() => Some(Err(ToolError::cancelled( - tool_id, - "tool call cancelled by server", - ))), - t = drain => t, - }; - - let response = match terminal { - Some(Ok(typed)) => { - // A cco encode failure degrades to `None` (mirroring the decode - // side) rather than failing an otherwise-successful call. - let chat_completion_output = typed.chat_completion_output.as_ref().and_then(|cco| { - serde_json::to_value(cco) - .inspect_err(|err| warn!(%err, "dropping unencodable chat_completion_output")) - .ok() - }); - let encoded = serde_json::to_value(ToolCallResult { - tool_call_id: params.tool_call_id.clone(), - output: ToolOutputWire::Json(typed.value), - follow_ups: Vec::new(), - reminders: Vec::new(), - chat_completion_output, - }); - match encoded { - Ok(payload) => JsonRpcResponse { - jsonrpc: JsonRpcVersion, - id: json_id, - session_id: Some(session_id.clone()), - outcome: ResponseOutcome::Result(payload), - }, - Err(err) => { - let message = format!("failed to encode tool_call_result: {err}"); - JsonRpcResponse { - jsonrpc: JsonRpcVersion, - id: json_id, - session_id: Some(session_id.clone()), - outcome: ResponseOutcome::Error(JsonRpcError { - code: -32603, - data: serde_json::to_value(ToolErrorWire::Internal { - request_id: None, - detail: Some(message.clone()), - }) - .ok(), - message, - }), - } - } - } - } - Some(Err(err)) => build_error_response(json_id, session_id.clone(), err), - None => { - warn!("handler stream ended without terminal"); - JsonRpcResponse { - jsonrpc: JsonRpcVersion, - id: json_id, - session_id: Some(session_id.clone()), - outcome: ResponseOutcome::Error(JsonRpcError { - code: -32603, - message: "handler produced no terminal".to_owned(), - data: serde_json::to_value(ToolErrorWire::Internal { - request_id: None, - detail: Some("handler produced no terminal".to_owned()), - }) - .ok(), - }), - } - } - }; - - let text = match serde_json::to_string(&response) { - Ok(s) => s, - Err(err) => { - warn!(?err, "failed to serialise tool_call_result; dropping"); - return; - } - }; - if let Err(err) = connection.send_outbound(text).await { - warn!(?err, "failed to send tool_call_result"); - } -} - -/// Convert a [`ToolProgress`] into a wire [`ToolCallProgressFrame`]. -fn progress_to_frame(progress: ToolProgress, tool_call_id: ToolCallId) -> ToolCallProgressFrame { - let (kind, body) = match progress { - ToolProgress::Text { text } => ("text".to_owned(), serde_json::json!({ "text": text })), - ToolProgress::Content { blocks } => ( - "content".to_owned(), - serde_json::to_value(blocks).unwrap_or_else(|err| { - warn!( - ?err, - "failed to serialize Content blocks for progress frame" - ); - Value::default() - }), - ), - ToolProgress::Custom { subkind, payload } => (subkind, payload), - }; - ToolCallProgressFrame { - tool_call_id, - kind, - body, - dropped_count: None, - } -} - -/// Build the error response for a failed tool call, preserving the full -/// [`ToolError`] as a decodable [`ToolErrorWire`] in `error.data` (and the -/// matching numeric code) so the harness recovers kind + detail + structured -/// details instead of collapsing everything to a bare `-32603` string. -fn build_error_response(id: JsonRpcId, session_id: SessionId, err: ToolError) -> JsonRpcResponse { - let message = err.to_string(); - let wire = ToolErrorWire::from(err); - JsonRpcResponse { - jsonrpc: JsonRpcVersion, - id, - session_id: Some(session_id), - outcome: ResponseOutcome::Error(JsonRpcError { - code: error_codes::from_tool_error_wire(&wire), - message, - data: serde_json::to_value(&wire).ok(), - }), - } -} - -async fn send_error( - connection: &Arc, - id: JsonRpcId, - session_id: SessionId, - code: i32, - message: &str, -) { - send_error_with_data(connection, id, session_id, code, message, None).await; -} - -/// Like [`send_error`] but attaches a machine-readable `data` payload so -/// receivers can switch on `data.code` (see `error_codes.rs`). -async fn send_error_with_data( - connection: &Arc, - id: JsonRpcId, - session_id: SessionId, - code: i32, - message: &str, - data: Option, -) { - let response = JsonRpcResponse:: { - jsonrpc: JsonRpcVersion, - id, - session_id: Some(session_id), - outcome: ResponseOutcome::Error(JsonRpcError { - code, - message: message.to_owned(), - data, - }), - }; - if let Ok(text) = serde_json::to_string(&response) { - let _ = connection.send_outbound(text).await; - } -} - -/// Ship the shared overloaded (-32016 "tool_busy") response built by -/// [`crate::admission::overloaded_response`] — the single source of the -/// overload wire shape, reused by the demux inbox-full path too. -async fn send_overloaded(connection: &Arc, id: JsonRpcId, session_id: SessionId) { - let response = crate::admission::overloaded_response(id, session_id); - if let Ok(text) = serde_json::to_string(&response) { - let _ = connection.send_outbound(text).await; - } -} - -#[cfg(test)] -mod tests { - use super::*; - use kigi_tool_runtime::ContentBlock; - - fn call_id() -> ToolCallId { - ToolCallId::new_v7() - } - - // ── build_error_response: wire fidelity ───────────────────────── - - #[test] - fn build_error_response_preserves_kind_and_details_in_data() { - let sid = SessionId::new("sess").expect("valid"); - let tool = ToolId::new("run_terminal_command").expect("valid"); - let err = ToolError::execution(tool, "command exited 127: bash: not found"); - let resp = build_error_response(JsonRpcId::Number(1), sid, err); - match resp.outcome { - ResponseOutcome::Error(rpc) => { - assert_eq!(rpc.code, -32603, "Execution maps to internal_error numeric"); - assert!(rpc.message.contains("command exited 127")); - let data = rpc.data.expect("ToolErrorWire payload must ride in data"); - let wire: ToolErrorWire = - serde_json::from_value(data).expect("data decodes as ToolErrorWire"); - match wire { - ToolErrorWire::Execution { tool_id, message } => { - assert_eq!(tool_id.as_str(), "run_terminal_command"); - assert_eq!(message, "command exited 127: bash: not found"); - } - other => panic!("expected Execution, got {other:?}"), - } - } - ResponseOutcome::Result(_) => panic!("expected error outcome"), - } - } - - #[test] - fn build_error_response_maps_invalid_arguments_numeric() { - let sid = SessionId::new("sess").expect("valid"); - let err = ToolError::invalid_arguments("missing `path`"); - let resp = build_error_response(JsonRpcId::Number(2), sid, err); - match resp.outcome { - ResponseOutcome::Error(rpc) => { - assert_eq!(rpc.code, -32602, "InvalidArguments maps to invalid_params"); - assert_eq!( - rpc.data - .expect("data present") - .get("code") - .and_then(|v| v.as_str()), - Some("invalid_params"), - ); - } - ResponseOutcome::Result(_) => panic!("expected error outcome"), - } - } - - // ── progress_to_frame: basic variant tests ────────────────────── - - #[test] - fn progress_text_produces_text_kind_frame() { - let id = call_id(); - let progress = ToolProgress::Text { - text: "hello world".to_owned(), - }; - let frame = progress_to_frame(progress, id.clone()); - assert_eq!(frame.tool_call_id, id); - assert_eq!(frame.kind, "text"); - assert_eq!(frame.body, serde_json::json!({ "text": "hello world" })); - assert_eq!(frame.dropped_count, None); - } - - #[test] - fn progress_content_serializes_blocks_into_body() { - let id = call_id(); - let blocks = vec![ - ContentBlock::Text { - text: "line one".to_owned(), - }, - ContentBlock::Text { - text: "line two".to_owned(), - }, - ]; - let expected_body = serde_json::to_value(&blocks).unwrap(); - let progress = ToolProgress::Content { blocks }; - let frame = progress_to_frame(progress, id.clone()); - assert_eq!(frame.tool_call_id, id); - assert_eq!(frame.kind, "content"); - assert_eq!(frame.body, expected_body); - assert_eq!(frame.dropped_count, None); - } - - #[test] - fn progress_custom_preserves_subkind_and_payload() { - let id = call_id(); - let payload = serde_json::json!({ "cursor": 42, "partial": true }); - let progress = ToolProgress::Custom { - subkind: "my_tool.cursor".to_owned(), - payload: payload.clone(), - }; - let frame = progress_to_frame(progress, id.clone()); - assert_eq!(frame.tool_call_id, id); - assert_eq!(frame.kind, "my_tool.cursor"); - assert_eq!(frame.body, payload); - assert_eq!(frame.dropped_count, None); - } - - // ── edge cases ────────────────────────────────────────────────── - - #[test] - fn progress_text_empty_string() { - let id = call_id(); - let progress = ToolProgress::Text { - text: String::new(), - }; - let frame = progress_to_frame(progress, id); - assert_eq!(frame.kind, "text"); - assert_eq!(frame.body, serde_json::json!({ "text": "" })); - } - - #[test] - fn progress_content_empty_blocks() { - let id = call_id(); - let progress = ToolProgress::Content { blocks: vec![] }; - let frame = progress_to_frame(progress, id); - assert_eq!(frame.kind, "content"); - assert_eq!(frame.body, serde_json::json!([])); - } - - #[test] - fn progress_custom_null_payload_and_empty_subkind() { - let id = call_id(); - let progress = ToolProgress::Custom { - subkind: String::new(), - payload: Value::Null, - }; - let frame = progress_to_frame(progress, id); - assert_eq!(frame.kind, ""); - assert_eq!(frame.body, Value::Null); - } - - // ── content block coverage ────────────────────────────────────── - - #[test] - fn progress_content_with_image_block() { - let id = call_id(); - let blocks = vec![ContentBlock::Image { - mime_type: "image/png".to_owned(), - data: "base64data".to_owned(), - media_id: Some("img-1".to_owned()), - filename: None, - path: None, - metadata: Default::default(), - }]; - let expected_body = serde_json::to_value(&blocks).unwrap(); - let progress = ToolProgress::Content { blocks }; - let frame = progress_to_frame(progress, id); - assert_eq!(frame.kind, "content"); - assert_eq!(frame.body, expected_body); - } - - #[test] - fn progress_content_with_resource_block() { - let id = call_id(); - let blocks = vec![ContentBlock::Resource { - uri: "file:///tmp/data.csv".to_owned(), - mime_type: Some("text/csv".to_owned()), - text: Some("a,b\n1,2".to_owned()), - }]; - let expected_body = serde_json::to_value(&blocks).unwrap(); - let progress = ToolProgress::Content { blocks }; - let frame = progress_to_frame(progress, id); - assert_eq!(frame.kind, "content"); - assert_eq!(frame.body, expected_body); - assert_eq!(frame.body[0]["type"], "resource"); - assert_eq!(frame.body[0]["uri"], "file:///tmp/data.csv"); - assert_eq!(frame.body[0]["mime_type"], "text/csv"); - assert_eq!(frame.body[0]["text"], "a,b\n1,2"); - } - - // ── serde round-trips ─────────────────────────────────────────── - - #[test] - fn progress_frame_text_round_trips() { - let progress = ToolProgress::Text { - text: "round-trip".to_owned(), - }; - let frame = progress_to_frame(progress, call_id()); - let json = serde_json::to_value(&frame).expect("serialize"); - let back: ToolCallProgressFrame = serde_json::from_value(json).expect("deserialize"); - assert_eq!(back, frame); - } - - #[test] - fn progress_frame_content_round_trips() { - let blocks = vec![ - ContentBlock::Text { - text: "hello".to_owned(), - }, - ContentBlock::Image { - mime_type: "image/png".to_owned(), - data: "abc".to_owned(), - media_id: None, - filename: None, - path: None, - metadata: Default::default(), - }, - ]; - let progress = ToolProgress::Content { blocks }; - let frame = progress_to_frame(progress, call_id()); - let json = serde_json::to_value(&frame).expect("serialize"); - let back: ToolCallProgressFrame = serde_json::from_value(json).expect("deserialize"); - assert_eq!(back, frame); - } - - #[test] - fn progress_frame_custom_round_trips() { - let progress = ToolProgress::Custom { - subkind: "streaming.chunk".to_owned(), - payload: serde_json::json!({ "offset": 1024, "data": [1, 2, 3] }), - }; - let frame = progress_to_frame(progress, call_id()); - let json = serde_json::to_value(&frame).expect("serialize"); - let back: ToolCallProgressFrame = serde_json::from_value(json).expect("deserialize"); - assert_eq!(back, frame); - } - - // ── wire notification shape ───────────────────────────────────── - - #[test] - fn progress_notification_has_correct_wire_shape() { - let id = call_id(); - let sid = SessionId::new("test-sess").expect("valid"); - let progress = ToolProgress::Text { - text: "wire test".to_owned(), - }; - let frame = progress_to_frame(progress, id.clone()); - // Build with typed params — same generic instantiation as production. - let notif = JsonRpcNotification { - jsonrpc: JsonRpcVersion, - session_id: Some(sid.clone()), - seq: None, - method: Method::ToolCallProgress.as_wire_str().to_owned(), - params: frame, - }; - let json: Value = serde_json::from_str(&serde_json::to_string(¬if).unwrap()).unwrap(); - assert_eq!(json["jsonrpc"], "2.0"); - assert_eq!(json["method"], "tool_call_progress"); - assert_eq!(json["session_id"], sid.as_str()); - assert!(json.get("id").is_none(), "notifications must not have id"); - assert!(json.get("seq").is_none(), "None seq must be omitted"); - // params.kind - assert_eq!(json["params"]["kind"], "text"); - // params.tool_call_id present as string - assert_eq!(json["params"]["tool_call_id"], id.as_str()); - // params.body.text matches - assert_eq!(json["params"]["body"]["text"], "wire test"); - // dropped_count absent when None - assert!( - json["params"].get("dropped_count").is_none(), - "None dropped_count must be omitted" - ); - } - - #[test] - fn json_serialized_len_matches_to_string() { - let v = serde_json::json!({"a": 1, "b": ["x", "y"], "c": {"d": true}}); - assert_eq!( - json_serialized_len(&v).unwrap(), - serde_json::to_string(&v).unwrap().len() - ); - } - - #[test] - fn system_notify_data_error_propagates_not_unsupported() { - let outcome = ResponseOutcome::Error(JsonRpcError { - code: -32601, - message: "server not found".to_owned(), - data: Some(serde_json::json!({"code": "tool_server_not_found"})), - }); - assert!(system_notify_ack_from_outcome(outcome).is_err()); - } - - #[test] - fn system_notify_ok_reply_maps_to_accepted() { - let outcome = ResponseOutcome::Result(serde_json::json!({})); - assert_eq!( - system_notify_ack_from_outcome(outcome).unwrap(), - SystemNotifyAck::Accepted - ); - } - - #[test] - fn system_notify_method_not_found_maps_to_forwarding_unsupported() { - let outcome = ResponseOutcome::Error(JsonRpcError { - code: -32601, - message: "method not found".to_owned(), - data: None, - }); - assert_eq!( - system_notify_ack_from_outcome(outcome).unwrap(), - SystemNotifyAck::ForwardingUnsupported - ); - } - - #[test] - fn system_notify_other_error_propagates() { - let outcome = ResponseOutcome::Error(JsonRpcError { - code: -32602, - message: "invalid params".to_owned(), - data: None, - }); - assert!(system_notify_ack_from_outcome(outcome).is_err()); - } -} diff --git a/crates/common/kigi-computer-hub-sdk/src/trace_donate.rs b/crates/common/kigi-computer-hub-sdk/src/trace_donate.rs deleted file mode 100644 index db9e927..0000000 --- a/crates/common/kigi-computer-hub-sdk/src/trace_donate.rs +++ /dev/null @@ -1,206 +0,0 @@ -//! Forward selected spans to the connected server over the WebSocket -//! transport (`traces.donate`). The bounded retry buffer + drain barrier -//! live in [`crate::donate_pump`]; overflow drops spans — telemetry, -//! never correctness. - -use std::borrow::Cow; - -use base64::Engine as _; -use fastrace::collector::{Reporter, SpanRecord}; -use fastrace_opentelemetry::OpenTelemetryReporter; -use kigi_tool_protocol::{MAX_DONATION_BYTES, MAX_SPANS_PER_DONATION}; -use opentelemetry::InstrumentationScope; -use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; -use opentelemetry_proto::transform::common::tonic::ResourceAttributesWithSchema; -use opentelemetry_proto::transform::trace::tonic::group_spans_by_resource_and_scope; -use opentelemetry_sdk::Resource; -use opentelemetry_sdk::error::OTelSdkResult; -use opentelemetry_sdk::trace::{SpanData, SpanExporter}; -use prost::Message as _; -use tokio::sync::mpsc; - -use crate::donate_pump::{PENDING_FLUSHES, PumpMsg, drain_via, run_pump}; -use crate::server::ToolServer; - -/// fastrace [`Reporter`] feeding the donation pump. -pub struct HubDonatingReporter(OpenTelemetryReporter); - -impl Reporter for HubDonatingReporter { - fn report(&mut self, spans: Vec) { - if spans.is_empty() { - return; - } - self.0.report(spans); - } -} - -/// [`SpanExporter`] that encodes OTLP requests onto the pump channel. -/// Runs on fastrace's collector thread; must never block. -#[derive(Debug)] -struct PumpSpanExporter { - tx: mpsc::Sender, - resource: ResourceAttributesWithSchema, -} - -impl SpanExporter for PumpSpanExporter { - fn export( - &self, - batch: Vec, - ) -> impl std::future::Future + Send { - let mut remaining = batch; - while !remaining.is_empty() { - let chunk = if remaining.len() > MAX_SPANS_PER_DONATION { - let rest = remaining.split_off(MAX_SPANS_PER_DONATION); - std::mem::replace(&mut remaining, rest) - } else { - std::mem::take(&mut remaining) - }; - let request = ExportTraceServiceRequest { - resource_spans: group_spans_by_resource_and_scope(chunk, &self.resource), - }; - let bytes = request.encode_to_vec(); - if bytes.len() > MAX_DONATION_BYTES { - tracing::debug!(len = bytes.len(), "dropping oversized donation payload"); - continue; - } - let payload = base64::engine::general_purpose::STANDARD.encode(bytes); - if self.tx.try_send(PumpMsg::Payload(payload)).is_err() { - tracing::debug!("trace donation queue full; dropping span batch"); - } - } - std::future::ready(Ok(())) - } - - fn set_resource(&mut self, resource: &Resource) { - self.resource = resource.into(); - } -} - -/// Shutdown fence: drains queued donations before the connection closes. -pub struct TraceDonationPump { - tx: mpsc::Sender, -} - -impl TraceDonationPump { - /// Resolves once every payload queued before this call has had a - /// send attempt. Call after `fastrace::flush()`. - pub async fn drain(&self) { - drain_via(&self.tx).await; - } -} - -impl ToolServer { - /// Spawn the donation pump and return its reporter + drain handle. - /// `service_name` must be server-allowlisted. - pub fn trace_donation_reporter( - &self, - service_name: impl Into, - ) -> (HubDonatingReporter, TraceDonationPump) { - let (tx, rx) = mpsc::channel::(PENDING_FLUSHES); - let server = self.downgrade(); - tokio::spawn(run_pump(rx, move |payload: String| { - let server = server.clone(); - async move { - let Some(server) = server.upgrade() else { - return (false, payload); - }; - let ok = server.donate_traces(&payload).await.is_ok(); - (ok, payload) - } - })); - self.set_donation_pump(tx.clone()); - - let resource = Resource::builder() - .with_service_name(service_name.into()) - .build(); - let exporter = PumpSpanExporter { - tx: tx.clone(), - resource: (&resource).into(), - }; - let reporter = OpenTelemetryReporter::new( - exporter, - Cow::Owned(resource), - InstrumentationScope::default(), - ); - (HubDonatingReporter(reporter), TraceDonationPump { tx }) - } -} - -#[cfg(test)] -mod tests { - use std::time::SystemTime; - - use opentelemetry::trace::{SpanContext, SpanKind, Status, TraceFlags, TraceState}; - - use super::*; - - #[tokio::test] - async fn exporter_encodes_standard_otlp_with_resource() { - let resource = Resource::builder() - .with_service_name("test-service") - .build(); - let (tx, mut rx) = mpsc::channel::(4); - let exporter = PumpSpanExporter { - tx, - resource: (&resource).into(), - }; - - let span = SpanData { - span_context: SpanContext::new( - 0x0af7651916cd43dd8448eb211c80319c_u128.into(), - 0xb7ad6b7169203331_u64.into(), - TraceFlags::SAMPLED, - false, - TraceState::default(), - ), - parent_span_id: 0_u64.into(), - parent_span_is_remote: false, - span_kind: SpanKind::Internal, - name: "tool_server.tool_call".into(), - start_time: SystemTime::UNIX_EPOCH, - end_time: SystemTime::UNIX_EPOCH, - attributes: vec![opentelemetry::KeyValue::new("tool_id", "bash")], - dropped_attributes_count: 0, - events: opentelemetry_sdk::trace::SpanEvents::default(), - links: opentelemetry_sdk::trace::SpanLinks::default(), - status: Status::Unset, - instrumentation_scope: InstrumentationScope::default(), - }; - exporter - .export(vec![span]) - .await - .expect("export must succeed"); - - let Some(PumpMsg::Payload(payload)) = rx.try_recv().ok() else { - panic!("exporter must enqueue one payload"); - }; - let bytes = base64::engine::general_purpose::STANDARD - .decode(payload) - .expect("payload must be base64"); - let request = - ExportTraceServiceRequest::decode(bytes.as_slice()).expect("payload must be OTLP"); - let resource_spans = &request.resource_spans[0]; - let service_name = resource_spans - .resource - .as_ref() - .unwrap() - .attributes - .iter() - .find(|kv| kv.key == "service.name") - .and_then(|kv| kv.value.as_ref()) - .map(|v| format!("{v:?}")); - assert!( - service_name.unwrap_or_default().contains("test-service"), - "resource must carry the donor service.name" - ); - let span = &resource_spans.scope_spans[0].spans[0]; - assert_eq!(span.name, "tool_server.tool_call"); - assert_eq!( - format!( - "{:032x}", - u128::from_be_bytes(span.trace_id.as_slice().try_into().unwrap()) - ), - "0af7651916cd43dd8448eb211c80319c" - ); - } -}