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" - ); - } -}