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

Show reasoning effort in the model's own vocabulary (K3: max, not xhigh)

Browse Source 
The welcome and prompt model labels (and /effort's 'current' hint)
rendered the canonical internal level name, so K3 at its default effort
showed 'K3 (xhigh)' even though the server's vocabulary for that level is
'max' (live /models think_efforts: low/high/max).

New ModelState::reasoning_effort_display() resolves the current effort
through the model's own effort menu (option id whose value matches),
falling back to the canonical name only when the model has no entry for
the level. All three display sites route through it; test pins the K3
mapping (Xhigh → 'max', Low → 'low', no-menu-entry → canonical).
This commit is contained in:
ZacharyZhang-NY
2026-07-17 21:44:43 -04:00
parent 28050e8e75
commit 5919526e91
66 changed files with 43 additions and 31221 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/codegen/kigi-shell/src/session/acp_session_tests/observability_bridge_mapping_tests.rs
-97
View File
@@ -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<TurnOutcome, acp::Error> = 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<TurnOutcome, acp::Error> = 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<TurnOutcome, acp::Error> = 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,");
}
crates/codegen/kigi-tui/src/acp/model_state.rs
+40
View File
@@ -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<String> {
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.
crates/codegen/kigi-tui/src/app/agent_view/render.rs
+1 -1
View File
@@ -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,
};
crates/codegen/kigi-tui/src/app/app_view.rs
+1 -1
View File
@@ -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,
};
crates/codegen/kigi-tui/src/slash/commands/effort.rs
+1 -1
View File
@@ -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() {
crates/codegen/kigi-workspace-client/Cargo.toml
-28
View File
@@ -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
crates/codegen/kigi-workspace-client/src/lib.rs
-810
View File
@@ -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<TypedToolOutput>` 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<TypedToolOutput>,
) -> Result<TypedToolOutput, kigi_tool_runtime::ToolError> {
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<AtomicBool>,
deadline: Option<Duration>,
}
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<AtomicBool>) -> 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<Value, WorkspaceClientError> {
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<R: WorkspaceRpc>(&self, req: &R) -> Result<R::Response, WorkspaceClientError> {
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<R::Response> =
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<WorkspaceInfo, WorkspaceClientError> {
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<Value, WorkspaceClientError> {
self.rpc(&GitStatusReq {}).await
}
pub async fn discover_skills(&self) -> Result<Vec<SkillInfo>, WorkspaceClientError> {
self.rpc(&DiscoverSkillsReq {}).await
}
pub async fn discover_agents_md(&self) -> Result<Vec<AgentConfigFile>, WorkspaceClientError> {
self.rpc(&DiscoverAgentsMdReq {}).await
}
pub async fn git_status_ext(
&self,
req: &GitStatusExtReq,
) -> Result<GitStatusExtResponse, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn git_files(
&self,
req: &GitFilesReq,
) -> Result<GitReadFilesData, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn git_diff(&self, req: &GitDiffReq) -> Result<GitDiffsData, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn git_stage(&self, req: &GitStageReq) -> Result<StageData, WorkspaceClientError> {
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<CommitResult, WorkspaceClientError> {
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<GitInfoData, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn git_branches(
&self,
req: &GitBranchesReq,
) -> Result<GitBranchListData, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn git_resolve_root(
&self,
req: &GitResolveRootReq,
) -> Result<Option<std::path::PathBuf>, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn git_current_commit(
&self,
req: &GitCurrentCommitReq,
) -> Result<Option<String>, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn detect_vcs_kind(
&self,
req: &DetectVcsKindReq,
) -> Result<VcsKind, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn git_checkout_commit(
&self,
req: &GitCheckoutCommitReq,
) -> Result<CheckoutCommitResponse, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn git_branch_info(&self) -> Result<Option<GitInfoData>, WorkspaceClientError> {
self.rpc(&GitBranchInfoReq {}).await
}
pub async fn git_metadata(&self) -> Result<Value, WorkspaceClientError> {
self.rpc(&GitMetadataReq {}).await
}
/// `workspace.git_collect_changes` — collect repository changes for serialization.
pub async fn git_collect_changes(
&self,
req: &GitCollectChangesReq,
) -> Result<GitCollectChangesResponse, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn put_files(&self, req: &PutFilesReq) -> Result<PutFilesRes, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn get_files(&self, req: &GetFilesReq) -> Result<GetFilesRes, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn fs_list(&self, req: &FsListReq) -> Result<FsListData, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn fs_exists(&self, req: &FsExistsReq) -> Result<FsExistsData, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn fs_read_file(
&self,
req: &FsReadFileReq,
) -> Result<FsReadFileData, WorkspaceClientError> {
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<HunkActionResponse, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn hunk_file_action(
&self,
req: &HunkFileActionReq,
) -> Result<BulkHunkActionResponse, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn hunk_turn_action(
&self,
req: &HunkTurnActionReq,
) -> Result<BulkHunkActionResponse, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn hunk_all_action(
&self,
req: &HunkAllActionReq,
) -> Result<BulkHunkActionResponse, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn hunk_get_staged_files(&self) -> Result<Vec<String>, WorkspaceClientError> {
self.rpc(&HunkGetStagedFilesReq {}).await
}
pub async fn hunk_get_file_summaries(&self) -> Result<Vec<FileSummary>, WorkspaceClientError> {
self.rpc(&HunkGetFileSummariesReq {}).await
}
pub async fn code_goto_definition(
&self,
req: &CodeGotoDefinitionReq,
) -> Result<CodeNavResponse, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn code_goto_references(
&self,
req: &CodeGotoReferencesReq,
) -> Result<CodeNavResponse, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn code_find_definitions(
&self,
req: &CodeFindDefinitionsReq,
) -> Result<CodeNavResponse, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn code_find_references(
&self,
req: &CodeFindReferencesReq,
) -> Result<CodeNavResponse, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn code_index_status(
&self,
req: &CodeIndexStatusReq,
) -> Result<CodeIndexStatusResponse, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn ripgrep(
&self,
req: &ContentSearchRequest,
) -> Result<ContentSearchData, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn fuzzy_open(&self, req: &FuzzyOpenReq) -> Result<String, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn fuzzy_change(&self, req: &FuzzyChangeReq) -> Result<bool, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn fuzzy_close(&self, req: &FuzzyCloseReq) -> Result<bool, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn fuzzy_status(&self, req: &FuzzyStatusReq) -> Result<Value, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn create_worktree(
&self,
req: &CreateWorktreeRequest,
) -> Result<Value, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn worktree_create_sync(
&self,
req: &WorktreeCreateSyncReq,
) -> Result<Value, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn remove_worktree(
&self,
req: &RemoveWorktreeRequest,
) -> Result<Value, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn apply_worktree(
&self,
req: &ApplyWorktreeRequest,
) -> Result<Value, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn worktree_list(
&self,
req: &WorktreeListReq,
) -> Result<Value, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn worktree_show(
&self,
req: &WorktreeShowReq,
) -> Result<Value, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn worktree_gc(&self, req: &WorktreeGcReq) -> Result<Value, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn worktree_db_rebuild(&self) -> Result<Value, WorkspaceClientError> {
self.rpc(&WorktreeDbRebuildReq {}).await
}
pub async fn worktree_db_path(&self) -> Result<WorktreeDbPathResponse, WorkspaceClientError> {
self.rpc(&WorktreeDbPathReq {}).await
}
pub async fn worktree_db_stats(&self) -> Result<Value, WorkspaceClientError> {
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<FileRewindResponse, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn load_project_config(&self) -> Result<Value, WorkspaceClientError> {
self.rpc(&LoadProjectConfigReq {}).await
}
pub async fn load_permissions(&self) -> Result<Value, WorkspaceClientError> {
self.rpc(&LoadPermissionsReq {}).await
}
pub async fn load_envrc(&self) -> Result<Value, WorkspaceClientError> {
self.rpc(&LoadEnvrcReq {}).await
}
pub async fn tool_definitions(
&self,
req: &ToolDefinitionsReq,
) -> Result<Value, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn resolve_file_references(
&self,
req: &ResolveFileReferencesReq,
) -> Result<Value, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn update_tool_config(
&self,
req: &UpdateToolConfigReq,
) -> Result<Value, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn drop_session(&self, req: &DropSessionReq) -> Result<Value, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn configure_mcp(
&self,
req: &ConfigureMcpReq,
) -> Result<Value, WorkspaceClientError> {
self.rpc(req).await
}
pub async fn install_plugin(&self) -> Result<Value, WorkspaceClientError> {
self.rpc(&InstallPluginReq {}).await
}
pub async fn refresh_plugins(&self) -> Result<Value, WorkspaceClientError> {
self.rpc(&RefreshPluginsReq {}).await
}
pub async fn discover_plugins(&self) -> Result<Vec<Value>, 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<RawOut, ToolError> {
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<TypedToolOutput> =
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::<TypedToolOutput>(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));
}
}
crates/codegen/kigi-workspace/src/bin/workspace_server.rs
-583
View File
@@ -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<String> {
id.parse::<kigi_tool_protocol::ServerId>()
.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<PathBuf>,
#[arg(long)]
cwd: Option<PathBuf>,
/// 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<String>,
/// JSON metadata attached to the tool server registration.
/// Propagated to `ServerInfo.metadata` in `servers.list` responses.
#[arg(long)]
metadata: Option<String>,
/// 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<u16>,
/// Proxy `--control-port` (loopback control). Absent ⇒ proxy default.
#[arg(long)]
preview_control_port: Option<u16>,
/// 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<PreviewVisibility>,
/// Proxy `--instance-suffix` for inbound Host validation.
#[arg(long)]
preview_instance_suffix: Option<String>,
/// 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<String>,
/// 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<u16>,
}
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::<ProfileName>()
.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"]);
}
}
crates/codegen/kigi-workspace/src/bin/workspace_server_probe.rs
-214
View File
@@ -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<String>,
/// 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<Value> {
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<kigi_computer_hub_sdk::ToolHarness> {
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<String> = 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(())
}
crates/codegen/kigi-workspace/src/daemonize.rs
-1080
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-workspace/src/diag_server.rs
-561
View File
@@ -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<String>,
state: DiagState,
pid: u32,
connected_at: Option<u64>,
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<u64>,
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<String>,
inner: Arc<Mutex<Inner>>,
}
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<String>) -> 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<Arc<PathBuf>>,
}
#[derive(Debug, Deserialize)]
struct LogsQuery {
tail_bytes: Option<u64>,
}
async fn logs(State(ctx): State<DiagContext>, Query(query): Query<LogsQuery>) -> 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<Vec<u8>> {
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<DiagContext>| 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<DiagContext>| 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<PathBuf>,
) -> anyhow::Result<BoundDiag> {
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<u16>,
/// 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>, 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<PathBuf>) -> 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);
}
}
crates/codegen/kigi-workspace/src/file_system/client_fs.rs
-775
View File
@@ -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<HashMap<PathBuf, MemoEntry>>,
}
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<String> {
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<PathBuf> {
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<FsListRes> {
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<FsListRes> {
// 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<FsListNode> = 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<FsStatRes> {
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<FsReadFileRes> {
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 -> <outside>.
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<u8> = (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"));
}
}
crates/codegen/kigi-workspace/src/hub.rs
-1429
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-workspace/src/hub_auth.rs
-490
View File
@@ -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<AuthIdentity> {
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<String>,
#[serde(default)]
oidc_issuer: Option<String>,
#[serde(default)]
oidc_client_id: Option<String>,
#[serde(default)]
principal_type: Option<String>,
#[serde(default)]
principal_id: Option<String>,
#[serde(default)]
expires_at: Option<chrono::DateTime<chrono::Utc>>,
}
fn default_auth_path() -> anyhow::Result<PathBuf> {
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<String, AuthEntry> = 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<Arc<dyn AuthProvider>> {
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<Arc<dyn AuthProvider>> {
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");
}
}
crates/codegen/kigi-workspace/src/hub_channel.rs
-127
View File
@@ -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<WorkspaceEvent> {
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<ToolNotification> {
use kigi_tool_protocol::WireToolNotification;
match &frame.notification {
WireToolNotification::Custom(c) => {
serde_json::from_value::<ToolNotification>(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.
}
crates/codegen/kigi-workspace/src/hub_ids.rs
-10
View File
@@ -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,
};
crates/codegen/kigi-workspace/src/hub_server.rs
-2999
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-workspace/src/mcp.rs
-220
View File
@@ -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<McpClient>,
}
impl McpClientTransportAdapter {
pub fn new(client: Arc<McpClient>) -> Self {
Self { client }
}
}
#[async_trait]
impl McpTransport for McpClientTransportAdapter {
async fn initialize(&self) -> Result<McpServerInfo, 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 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<Vec<McpToolDefinition>, 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<String> = 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<McpCallResult, 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()))?;
// 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<McpToolHandler>,
}
impl QualifiedMcpToolHandler {
/// Returns `None` if the qualified name is not a valid `ToolId`.
pub fn try_new(qualified_name: String, inner: Arc<McpToolHandler>) -> Option<Self> {
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<Value> {
self.inner.input_schema()
}
async fn handle_call(&self, ctx: ToolCallContext, args: Value) -> ToolStream<TypedToolOutput> {
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<String>,
/// Servers that failed to start.
pub failed: Vec<McpStartFailure>,
}
/// 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()),
}
}
crates/codegen/kigi-workspace/src/permission/hub_permission.rs
-498
View File
@@ -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<HistogramVec> = 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<IntCounter> = 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<Value, String>;
}
/// 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<Self> {
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<Value, String> {
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<String>)> {
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<AccessKind> {
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<Value, String>,
seen: Mutex<Option<Value>>,
}
#[async_trait]
impl PermissionHookTransport for StubTransport {
async fn request_permission(&self, payload: Value) -> Result<Value, String> {
*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());
}
}
}
crates/codegen/kigi-workspace/src/preview_supervisor.rs
-1326
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-workspace/src/recovery.rs
-158
View File
@@ -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 `<workspace_home>/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 <max_age dir must survive the same sweep"
);
}
}
crates/codegen/kigi-workspace/src/rpc_envelope.rs
-227
View File
@@ -1,227 +0,0 @@
//! `WorkspaceError` <-> 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<T>(error: &WorkspaceError) -> RpcEnvelope<T> {
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<WorkspaceError> = 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<String> = RpcEnvelope::ok("hello".into());
let json = serde_json::to_value(&env).unwrap();
let recovered: RpcEnvelope<String> = 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<String> = envelope_err(&err);
let json = serde_json::to_value(&env).unwrap();
let recovered: RpcEnvelope<String> = 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");
}
}
}
}
crates/common/kigi-computer-hub-core/Cargo.toml
-25
View File
@@ -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
crates/common/kigi-computer-hub-core/src/inner.rs
-73
View File
@@ -1,73 +0,0 @@
//! `InnerDispatchForResolver` — an object-safe `ToolDispatch` that routes
//! through a `Weak<CompoundResolver>` bound to a single session.
//!
//! Tools that need to call other tools (the inner-dispatch pattern) ask
//! the runtime for an `Arc<dyn ToolDispatch>`. 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<CompoundResolver>,
session_id: SessionId,
}
impl InnerDispatchForResolver {
/// Build an inner-dispatch handle bound to `session_id`, resolving
/// through `resolver`.
pub fn new(resolver: Weak<CompoundResolver>, 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<TypedToolOutput> {
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
}
}
crates/common/kigi-computer-hub-core/src/lib.rs
-29
View File
@@ -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};
crates/common/kigi-computer-hub-core/src/local.rs
-77
View File
@@ -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<CompoundResolver>,
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<CompoundResolver>, 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<Principal, ToolError> {
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<TypedToolOutput> {
self.resolver
.resolve_and_dispatch(&self.session_id, tool_id, args, ctx)
.await
}
}
crates/common/kigi-computer-hub-core/src/registry.rs
-305
View File
@@ -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<RegistrationOutcome>;
/// 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<ResolvedTool>;
/// 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<ToolDescription>;
/// Enumerate active server summaries for `session`. Useful for
/// rendering connected-integrations system reminders.
fn list_servers(&self, session: &SessionId) -> Vec<ServerSummary>;
/// 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<SessionId>;
/// All servers registered by this user across all connections.
fn list_servers_for_user(&self, user_id: &UserId) -> Vec<ServerRecord>;
/// Look up a server by its connection ID.
fn get_server_record(&self, connection_id: &ConnectionId) -> Option<ServerRecord>;
/// 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<ServerId> {
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<chrono::Utc>,
/// 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)",
);
}
}
crates/common/kigi-computer-hub-core/src/remote.rs
-541
View File
@@ -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<JsonRpcResponse, ToolError>;
/// 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<dyn ConnectionClient>,
}
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<dyn ConnectionClient>,
) -> 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<TypedToolOutput> {
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<dyn ConnectionClient>,
session_id: SessionId,
user_id: UserId,
}
impl RemoteTransport {
/// Build a transport over `connection`, bound to `(user_id,
/// session_id)`.
pub fn new(
connection: Arc<dyn ConnectionClient>,
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<Principal, ToolError> {
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<TypedToolOutput> {
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<dyn ConnectionClient>,
tool_id: ToolId,
session_id: SessionId,
arguments: Value,
ctx: ToolCallContext,
) -> ToolStream<TypedToolOutput> {
let cwd = ctx
.extensions
.get::<Cwd>()
.map(|c| c.0.to_string_lossy().into_owned());
let behavior_version = ctx.extensions.get::<BehaviorVersion>().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(&params) {
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<JsonRpcResponse, ToolError>>;
/// 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<ToolId>,
progress: BoxStream<'static, ToolCallProgressFrame>,
request: Option<ResponseFuture>,
done: bool,
}
impl Stream for RequestStream {
type Item = ToolStreamItem<TypedToolOutput>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
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<TypedToolOutput, _>` the runtime expects.
fn terminal_from_response(
tool_id: ToolId,
resp: JsonRpcResponse,
) -> Result<TypedToolOutput, ToolError> {
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<TypedToolOutput, ToolError> {
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::<ToolChatCompletionResponse>(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<ContentBlock> = 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::<ToolErrorWire>(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,
}
}
}
}
crates/common/kigi-computer-hub-core/src/resolver.rs
-280
View File
@@ -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<dyn ToolHandle>,
/// 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<dyn ToolHandle>,
/// 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<dyn ToolHandle> {
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<TypedToolOutput>;
}
/// 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<T> {
inner: Arc<T>,
}
impl<T> ErasedTool<T> {
/// Wrap an `Arc<T>` for use as an [`ToolHandle`].
pub fn from_arc(inner: Arc<T>) -> 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<T: std::fmt::Debug> std::fmt::Debug for ErasedTool<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ErasedTool")
.field("inner", &self.inner)
.finish()
}
}
impl<T> Clone for ErasedTool<T> {
fn clone(&self) -> Self {
Self {
inner: Arc::clone(&self.inner),
}
}
}
#[async_trait]
impl<T> ToolHandle for ErasedTool<T>
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<TypedToolOutput> {
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<dyn ToolRegistry>,
remote: Option<Arc<dyn ToolRegistry>>,
}
impl CompoundResolver {
/// Compose a resolver that consults a single local registry.
pub fn local_only(local: Arc<dyn ToolRegistry>) -> Self {
Self {
local,
remote: None,
}
}
/// Compose a resolver with both planes; `local` is consulted first.
pub fn compound(local: Arc<dyn ToolRegistry>, remote: Arc<dyn ToolRegistry>) -> Self {
Self {
local,
remote: Some(remote),
}
}
/// Borrow the local plane.
pub fn local(&self) -> &Arc<dyn ToolRegistry> {
&self.local
}
/// Borrow the optional remote plane.
pub fn remote(&self) -> Option<&Arc<dyn ToolRegistry>> {
self.remote.as_ref()
}
/// Resolve `(session, tool_id)` honouring the local-first rule.
pub fn resolve(&self, session: &SessionId, tool_id: &ToolId) -> Option<ResolvedTool> {
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<TypedToolOutput> {
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}"),
))),
}
}
}
crates/common/kigi-computer-hub-core/src/transport.rs
-116
View File
@@ -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<SessionId>,
/// OAuth-style scopes granted to this principal, e.g. `"tool.invoke"`.
pub scopes: Vec<String>,
/// 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<String>,
}
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<String>) -> Self {
self.scopes.push(scope.into());
self
}
/// Append `aud` to the token's audience list.
pub fn with_audience(mut self, aud: impl Into<String>) -> 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<Principal, ToolError>;
/// 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<TypedToolOutput>;
}
crates/common/kigi-computer-hub-core/tests/compound_resolver.rs
-335
View File
@@ -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<Self::Output, ToolError> {
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<ToolId, Arc<dyn ToolHandle>>,
}
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<RegistrationOutcome> {
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<ResolvedTool> {
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<ToolDescription> {
vec![]
}
fn list_servers(&self, _session: &SessionId) -> Vec<ServerSummary> {
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<SessionId> {
std::collections::HashSet::new()
}
fn list_servers_for_user(
&self,
_user_id: &kigi_tool_protocol::UserId,
) -> Vec<kigi_computer_hub_core::registry::ServerRecord> {
Vec::new()
}
fn get_server_record(
&self,
_connection_id: &ConnectionId,
) -> Option<kigi_computer_hub_core::registry::ServerRecord> {
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<dyn ToolRegistry>);
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<dyn ToolRegistry>);
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<dyn ToolRegistry>,
remote as Arc<dyn ToolRegistry>,
);
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<dyn ToolRegistry>,
remote as Arc<dyn ToolRegistry>,
);
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<dyn ToolRegistry>);
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<dyn ToolRegistry>,
remote as Arc<dyn ToolRegistry>,
);
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<dyn ToolRegistry>);
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<dyn ToolRegistry>);
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<dyn ToolRegistry>);
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:?}"),
}
}
crates/common/kigi-computer-hub-core/tests/inner_dispatch.rs
-317
View File
@@ -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<Self::Output, ToolError> {
Ok(serde_json::json!({"echoed": args.payload}))
}
}
type RegistryEntry = (ToolRegistration, Arc<dyn ToolHandle>);
#[derive(Debug, Default)]
struct InMemRegistry {
entries: DashMap<(SessionId, ToolId), RegistryEntry>,
}
impl InMemRegistry {
fn install(&self, session: SessionId, tool: ToolId, handle: Arc<dyn ToolHandle>) {
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<RegistrationOutcome> {
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<ResolvedTool> {
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<ToolDescription> {
vec![]
}
fn list_servers(&self, _session: &SessionId) -> Vec<ServerSummary> {
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<SessionId> {
std::collections::HashSet::new()
}
fn list_servers_for_user(
&self,
_user_id: &kigi_tool_protocol::UserId,
) -> Vec<kigi_computer_hub_core::registry::ServerRecord> {
Vec::new()
}
fn get_server_record(
&self,
_connection_id: &ConnectionId,
) -> Option<kigi_computer_hub_core::registry::ServerRecord> {
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<dyn ToolRegistry>,
));
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<dyn ToolRegistry>,
));
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<dyn ToolRegistry>,
));
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<dyn ToolRegistry>,
));
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<dyn ToolRegistry>,
));
let inner: Arc<dyn ToolDispatch> = 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));
}
crates/common/kigi-computer-hub-core/tests/internal_error_detail.rs
-65
View File
@@ -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");
}
crates/common/kigi-computer-hub-core/tests/local_transport.rs
-369
View File
@@ -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<Self::Output, ToolError> {
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<Self::Output> {
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<dyn ToolHandle>);
#[derive(Debug, Default)]
struct InMemRegistry {
entries: DashMap<(SessionId, ToolId), RegistryEntry>,
}
impl InMemRegistry {
fn install(&self, session: SessionId, tool_id: ToolId, handle: Arc<dyn ToolHandle>) {
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<RegistrationOutcome> {
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<ResolvedTool> {
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<ToolDescription> {
vec![]
}
fn list_servers(&self, _session: &SessionId) -> Vec<ServerSummary> {
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<SessionId> {
std::collections::HashSet::new()
}
fn list_servers_for_user(
&self,
_user_id: &kigi_tool_protocol::UserId,
) -> Vec<kigi_computer_hub_core::registry::ServerRecord> {
Vec::new()
}
fn get_server_record(
&self,
_connection_id: &ConnectionId,
) -> Option<kigi_computer_hub_core::registry::ServerRecord> {
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<kigi_tool_runtime::TypedToolOutput>,
) -> Vec<ToolStreamItem<kigi_tool_runtime::TypedToolOutput>> {
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<dyn ToolRegistry>,
));
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<dyn ToolRegistry>,
));
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<dyn ToolRegistry>,
));
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<dyn ToolRegistry>,
));
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<dyn ToolRegistry>,
));
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<serde_json::Value> = terminal_only(Ok(serde_json::Value::Null));
}
crates/common/kigi-computer-hub-core/tests/remote_proxy.rs
-727
View File
@@ -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<ToolCallId, mpsc::UnboundedSender<ToolCallProgressFrame>>,
/// 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<MockState>,
}
#[derive(Default)]
struct MockState {
/// FIFO queue of responses to return for each `request` call.
responses: Vec<MockResponse>,
/// Captured outgoing requests so tests can assert on them.
captured_requests: Vec<JsonRpcRequest>,
/// Captured one-way notifications.
captured_notifications: Vec<JsonRpcNotification>,
}
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<String>, data: Option<serde_json::Value>) {
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<String>) {
self.inner
.lock()
.expect("mutex")
.responses
.push(MockResponse::Network(message.into()));
}
fn last_request(&self) -> Option<JsonRpcRequest> {
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<JsonRpcResponse, ToolError> {
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, &notification);
}
#[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");
}
crates/common/kigi-computer-hub-core/tests/tool_registry.rs
-607
View File
@@ -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<Self::Output, ToolError> {
unreachable!("registry tests do not exercise execution")
}
}
#[derive(Debug, Clone)]
struct MockEntry {
registration: ToolRegistration,
sessions: HashSet<SessionId>,
}
/// 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<ToolId, Arc<dyn ToolHandle>>,
}
impl MockRegistry {
fn install_handle(&self, tool: Arc<dyn ToolHandle>) {
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<SessionId> = 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<RegistrationOutcome> {
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<ToolId> = 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<ToolId> = 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<ResolvedTool> {
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<ToolDescription> {
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<ServerSummary> {
let mut by_server: HashMap<ServerId, Vec<String>> = 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(), &reg.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<SessionId> {
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<kigi_computer_hub_core::registry::ServerRecord> {
Vec::new()
}
fn get_server_record(
&self,
_connection_id: &ConnectionId,
) -> Option<kigi_computer_hub_core::registry::ServerRecord> {
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<dyn ToolRegistry>);
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());
}
crates/common/kigi-computer-hub-core/tests/transport_trait.rs
-137
View File
@@ -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<Principal, ToolError> {
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<TypedToolOutput> {
terminal_only(Ok(TypedToolOutput::from_value(tool_id, args)))
}
}
#[tokio::test]
async fn boxed_transport_compiles_and_dispatches() {
let boxed: Box<dyn Transport> = 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")));
}
crates/common/kigi-computer-hub-core/tests/workspace_unavailable.rs
-203
View File
@@ -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<ToolError> 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"
)));
}
crates/common/kigi-computer-hub-mcp-adapter/Cargo.toml
-30
View File
@@ -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
crates/common/kigi-computer-hub-mcp-adapter/src/bridge.rs
-781
View File
@@ -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<String>,
}
/// 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<dyn McpTransport>,
handlers: Vec<Arc<McpToolHandler>>,
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<dyn McpTransport>,
config: &McpBridgeConfig,
) -> Result<McpBridgeHandle, McpError> {
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<Arc<McpToolHandler>> = 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<McpToolHandler>] {
&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<dyn McpTransport>,
namespace: Option<String>,
}
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<Value> {
self.definition.input_schema.clone()
}
async fn handle_call(&self, _ctx: ToolCallContext, args: Value) -> ToolStream<TypedToolOutput> {
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::<Vec<_>>()
.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<McpBlock> = 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<McpToolDefinition>,
call_response: Mutex<Option<McpCallResult>>,
call_error: Mutex<Option<McpError>>,
closed: AtomicBool,
last_call: Mutex<Option<(String, Value)>>,
}
impl MockTransport {
fn new(server_info: McpServerInfo, tools: Vec<McpToolDefinition>) -> 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<McpServerInfo, McpError> {
Ok(self.server_info.clone())
}
async fn list_tools(&self) -> Result<Vec<McpToolDefinition>, McpError> {
Ok(self.tools.clone())
}
async fn call_tool(&self, name: &str, arguments: Value) -> Result<McpCallResult, McpError> {
*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<McpToolDefinition> {
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<dyn McpTransport> {
Arc::new(mock) as Arc<dyn McpTransport>
}
#[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<String> = 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<dyn McpTransport> = Arc::clone(&mock) as Arc<dyn McpTransport>;
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())
);
}
}
crates/common/kigi-computer-hub-mcp-adapter/src/lib.rs
-52
View File
@@ -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<dyn McpTransport> = /* ... */;
//! 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};
crates/common/kigi-computer-hub-mcp-adapter/src/metrics.rs
-61
View File
@@ -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<Histogram> = 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<IntCounter> = 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<IntGauge> = 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::*;
crates/common/kigi-computer-hub-mcp-adapter/src/transport.rs
-38
View File
@@ -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<McpServerInfo, McpError>;
/// Discover available tools via MCP `tools/list`.
async fn list_tools(&self) -> Result<Vec<McpToolDefinition>, 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<McpCallResult, McpError>;
/// 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>;
}
crates/common/kigi-computer-hub-mcp-adapter/src/types.rs
-107
View File
@@ -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<String>,
/// JSON Schema for the tool's input arguments.
#[serde(default)]
pub input_schema: Option<serde_json::Value>,
}
/// 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<McpContent>,
/// 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<String>,
/// Optional text body.
#[serde(default)]
text: Option<String>,
},
}
/// 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),
}
crates/common/kigi-computer-hub-sdk/Cargo.toml
-57
View File
@@ -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
crates/common/kigi-computer-hub-sdk/src/admission.rs
-333
View File
@@ -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<Value> {
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<Semaphore> {
static SEM: OnceLock<Arc<Semaphore>> = 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::<usize>().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<SessionId, Arc<Semaphore>>,
session_max: usize,
conn_sem: Arc<Semaphore>,
global_sem: Arc<Semaphore>,
wait_timeout: Duration,
}
impl Admission {
pub(crate) fn new(
session_max: usize,
conn_max: usize,
global_sem: Arc<Semaphore>,
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<AdmitGuard, Overloaded> {
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<Semaphore>,
deadline: Instant,
) -> Result<OwnedSemaphorePermit, Overloaded> {
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"
);
}
}
crates/common/kigi-computer-hub-sdk/src/auth.rs
-238
View File
@@ -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<String, String> },
}
impl AuthCredential {
/// Convenience constructor for the bearer-token shape.
pub fn bearer(token: impl Into<String>) -> 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<I, K, V>(headers: I) -> Result<Self, ClientError>
where
I: IntoIterator<Item = (K, V)>,
K: AsRef<str>,
V: Into<String>,
{
let mut map: BTreeMap<String, String> = 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<String>,
/// Team id when `principal_type == "Team"`; otherwise `None`.
pub principal_id: Option<String>,
}
/// 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<AuthIdentity> {
None
}
}
pub type SharedAuthProvider = std::sync::Arc<dyn AuthProvider>;
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);
}
}
crates/common/kigi-computer-hub-sdk/src/cancel.rs
-301
View File
@@ -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<ToolCallId, CancellationToken>,
pending: DashSet<ToolCallId>,
/// 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<ToolCallId> = (0..5).map(|_| cid()).collect();
let tokens: Vec<CancellationToken> = 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()
);
}
}
crates/common/kigi-computer-hub-sdk/src/connection.rs
-2694
View File
File diff suppressed because it is too large Load Diff
crates/common/kigi-computer-hub-sdk/src/connection_borrow.rs
-224
View File
@@ -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<HubConnection>,
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<HubConnectionPool>,
url: Url,
auth: Arc<dyn AuthProvider>,
kind: ConnectionKind,
on_reconnect: Option<Arc<ReconnectCallback>>,
on_disconnect: Option<Arc<DisconnectCallback>>,
on_connect: Option<Arc<ConnectCallback>>,
server_id: Option<kigi_tool_protocol::ServerId>,
server_description: Option<String>,
server_metadata: Option<serde_json::Value>,
alpha_test_key: Option<String>,
allow_insecure_ws: bool,
tuning: ConnectionTuning,
) -> Result<Self, ClientError> {
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<HubConnection> {
&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<dyn AuthProvider> = 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);
}
}
crates/common/kigi-computer-hub-sdk/src/demux.rs
-973
View File
@@ -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<SessionId, tokio::sync::mpsc::Sender<InboundFrame>>,
waiters: DashMap<RequestId, oneshot::Sender<Result<JsonRpcResponse, ClientError>>>,
/// 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<RequestId, SessionId>,
progress: DashMap<ToolCallId, tokio::sync::mpsc::Sender<ToolCallProgressFrame>>,
/// Broadcast channel for connection-level notifications (no session_id).
notifications: tokio::sync::broadcast::Sender<Value>,
/// 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<tokio::sync::mpsc::Sender<String>>,
}
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<String>) -> Self {
Self {
outbound: Some(outbound),
..Self::default()
}
}
/// Subscribe to connection-level notifications (no session_id).
pub fn subscribe_notifications(&self) -> tokio::sync::broadcast::Receiver<Value> {
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<InboundFrame>,
) -> Option<tokio::sync::mpsc::Sender<InboundFrame>> {
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<tokio::sync::mpsc::Sender<InboundFrame>> {
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<Result<JsonRpcResponse, ClientError>>,
) {
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<Result<JsonRpcResponse, ClientError>>,
) {
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<oneshot::Sender<Result<JsonRpcResponse, ClientError>>> {
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<F>(
&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<RequestId> = 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<ToolCallProgressFrame>,
) -> Result<(), tokio::sync::mpsc::Sender<ToolCallProgressFrame>> {
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<tokio::sync::mpsc::Sender<ToolCallProgressFrame>> {
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<F>(&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<RequestId> = 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<ToolCallId> = 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<JsonRpcResponse, ClientError> =
serde_json::from_value::<JsonRpcResponse>(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::<JsonRpcId>(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::<String>(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::<String>(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::<String>(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::<ToolCallProgressFrame>(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::<ToolCallProgressFrame>(1);
demux
.try_register_progress_waiter(call_id.clone(), tx_first)
.expect("first registration");
let (tx_second, _rx_second) = mpsc::channel::<ToolCallProgressFrame>(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::<ToolCallProgressFrame>(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::<ToolCallProgressFrame>(1);
let (tx_b, mut rx_b) = mpsc::channel::<ToolCallProgressFrame>(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"
);
}
}
crates/common/kigi-computer-hub-sdk/src/donate_pump.rs
-206
View File
@@ -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<PumpMsg>) {
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<D, F>(mut rx: mpsc::Receiver<PumpMsg>, donate: D)
where
D: Fn(String) -> F,
F: std::future::Future<Output = (bool, String)>,
{
let mut retry: VecDeque<String> = 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<D, F>(retry: &mut VecDeque<String>, donate: &D)
where
D: Fn(String) -> F,
F: std::future::Future<Output = (bool, String)>,
{
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<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
let (tx, rx) = mpsc::channel::<PumpMsg>(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<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
let healthy = Arc::new(AtomicBool::new(false));
let (tx, rx) = mpsc::channel::<PumpMsg>(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");
}
}
crates/common/kigi-computer-hub-sdk/src/error.rs
-346
View File
@@ -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::<ToolErrorWire>(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_<code>`
/// 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<serde_json::Error> for ClientError {
fn from(err: serde_json::Error) -> Self {
Self::Serde(err.to_string())
}
}
impl From<IdError> for ClientError {
fn from(err: IdError) -> Self {
Self::ProtocolError(err.to_string())
}
}
impl From<url::ParseError> for ClientError {
fn from(err: url::ParseError) -> Self {
Self::InvalidConfig(format!("invalid url: {err}"))
}
}
impl From<tokio_tungstenite::tungstenite::Error> 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<tokio::sync::oneshot::error::RecvError> 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::<Vec<u8>>)
.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));
}
}
crates/common/kigi-computer-hub-sdk/src/handshake.rs
-93
View File
@@ -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<Si, St>(
sink: &mut Si,
stream: &mut St,
kind: ConnectionKind,
server_id: Option<kigi_tool_protocol::ServerId>,
description: Option<String>,
metadata: Option<serde_json::Value>,
) -> Result<HelloAckMsg, ClientError>
where
Si: SinkExt<Message> + Unpin,
Si::Error: std::fmt::Display,
St: StreamExt<Item = Result<Message, tokio_tungstenite::tungstenite::Error>> + 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(),
))
}
crates/common/kigi-computer-hub-sdk/src/harness.rs
-2940
View File
File diff suppressed because it is too large Load Diff
crates/common/kigi-computer-hub-sdk/src/lib.rs
-71
View File
@@ -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;
crates/common/kigi-computer-hub-sdk/src/log_donate.rs
-592
View File
@@ -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<u8>, Vec<u8>) {
match SpanContext::current_local_parent() {
Some(ctx) => encode_ids(&ctx),
None => (Vec::new(), Vec::new()),
}
}
fn encode_ids(ctx: &SpanContext) -> (Vec<u8>, Vec<u8>) {
(
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<String>,
attributes: Vec<KeyValue>,
}
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<PumpMsg>,
resource: Resource,
}
impl PumpLogExporter {
fn export(&self, mut records: Vec<LogRecord>) {
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<LogRecord>) {
self.exporter.export(records);
}
}
#[derive(Default)]
struct LogBatch {
records: Vec<LogRecord>,
oldest: Option<Instant>,
}
struct LogLayerShared {
sender: ArcSwapOption<LogDonationSender>,
batch: parking_lot::Mutex<LogBatch>,
}
impl LogLayerShared {
/// Buffer a record; return any records due for flush (count/age).
fn push(&self, record: LogRecord) -> Vec<LogRecord> {
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<ArcSwapOption<LogLayerShared>> =
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<LogLayerShared>,
}
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<S: tracing::Subscriber> Layer<S> 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<PumpMsg>,
}
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<String>,
) -> (LogDonationSender, LogDonationPump) {
let (tx, rx) = mpsc::channel::<PumpMsg>(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<PumpMsg>) -> 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::<PumpMsg>(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::<PumpMsg>(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);
}
}
crates/common/kigi-computer-hub-sdk/src/metric_donate.rs
-396
View File
@@ -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<KeyValue> {
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<Metric> {
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<PumpMsg>,
resource: Resource,
}
impl MetricExporter {
fn export(&self, mut metrics: Vec<Metric>) {
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<ArcSwapOption<MetricExporter>> =
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<PumpMsg>,
}
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<String>) -> MetricDonationPump {
let (tx, rx) = mpsc::channel::<PumpMsg>(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<String, String> {
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(&registry.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::<PumpMsg>(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(&registry.gather()).is_empty());
}
}
crates/common/kigi-computer-hub-sdk/src/metrics.rs
-552
View File
@@ -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<IntGauge> = 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<IntCounter> = 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<IntCounter> = 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<IntCounterVec> = 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<Histogram> = 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<IntCounterVec> = 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<Histogram> = 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<Histogram> = 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<IntGauge> = 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<IntCounter> = 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<IntCounterVec> = 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<IntCounterVec> = 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<HistogramVec> = 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<IntCounter> = 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<IntCounter> = 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<IntCounterVec> = 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<IntCounter> = 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<IntCounter> = 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<IntCounter> = 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<IntCounter> = 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<IntCounter> = 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<IntCounter> = 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<IntCounter> = 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<IntCounter> = 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<IntCounter> = 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<IntCounter> = 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<IntCounter> = 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<IntCounter> = 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<IntCounter> = 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<IntCounter> = 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<IntCounter> = 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<IntCounter> = 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<IntGaugeVec> = 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<Histogram> = 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;
crates/common/kigi-computer-hub-sdk/src/notification.rs
-362
View File
@@ -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<ToolId>,
removed: Vec<ToolId>,
updated: Vec<ToolId>,
},
/// 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<Self> {
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::<ToolsChanged>(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::<ToolNotificationFrame>(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::<ToolServerStatusPayload>(
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:?}"
);
}
}
crates/common/kigi-computer-hub-sdk/src/observability.rs
-275
View File
@@ -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<Arc<ToolHarness>>,
/// Retained for future payload enrichment and logging.
session_id: SessionId,
}
impl ObservabilityBridge {
pub fn new(harness: Option<Arc<ToolHarness>>, 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;
}
}
}
crates/common/kigi-computer-hub-sdk/src/oidc_provider.rs
-336
View File
@@ -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<dyn Fn(&RefreshEvent) + Send + Sync>;
#[derive(Debug, Clone)]
pub struct RefreshEvent {
pub access_token: String,
pub new_refresh_token: Option<String>,
pub expires_at: Option<DateTime<Utc>>,
}
struct TokenState {
access_token: String,
refresh_token: String,
expires_at: Option<DateTime<Utc>>,
}
pub struct OidcAuthProvider {
state: Mutex<TokenState>,
issuer: String,
client_id: String,
user_id: Option<String>,
principal_type: Option<String>,
principal_id: Option<String>,
on_refresh: Option<OnRefreshCallback>,
}
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<DateTime<Utc>>,
user_id: Option<String>,
principal_type: Option<String>,
principal_id: Option<String>,
on_refresh: Option<OnRefreshCallback>,
}
impl OidcAuthProviderBuilder {
pub fn new(
access_token: impl Into<String>,
refresh_token: impl Into<String>,
issuer: impl Into<String>,
client_id: impl Into<String>,
) -> 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<Utc>) -> 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<String>) -> Self {
self.user_id = Some(user_id.into());
self
}
pub fn principal_type(mut self, pt: impl Into<String>) -> Self {
self.principal_type = Some(pt.into());
self
}
pub fn principal_id(mut self, pid: impl Into<String>) -> 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<AuthIdentity> {
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<dyn std::error::Error>> {
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<dyn std::error::Error>> {
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<String>,
#[serde(default)]
expires_in: Option<u64>,
}
let tokens: Tokens = client
.post(&disc.token_endpoint)
.form(&params)
.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"));
}
}
crates/common/kigi-computer-hub-sdk/src/pool.rs
-344
View File
@@ -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<HubConnection>`; 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<HubConnection>,
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<Arc<HubConnectionPool>> = OnceCell::const_new();
/// Pool of live server connections.
pub struct HubConnectionPool {
connections: DashMap<ConnKey, Pooled>,
}
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<Self> {
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<Self> {
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<Self>,
url: Url,
credential: Arc<dyn AuthProvider>,
kind: ConnectionKind,
on_reconnect: Option<Arc<ReconnectCallback>>,
on_disconnect: Option<Arc<DisconnectCallback>>,
server_id: Option<kigi_tool_protocol::ServerId>,
alpha_test_key: Option<String>,
allow_insecure_ws: bool,
) -> Result<Arc<HubConnection>, 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<Self>,
url: Url,
credential: Arc<dyn AuthProvider>,
kind: ConnectionKind,
on_reconnect: Option<Arc<ReconnectCallback>>,
on_disconnect: Option<Arc<DisconnectCallback>>,
on_connect: Option<Arc<ConnectCallback>>,
server_id: Option<kigi_tool_protocol::ServerId>,
server_description: Option<String>,
server_metadata: Option<serde_json::Value>,
alpha_test_key: Option<String>,
allow_insecure_ws: bool,
tuning: ConnectionTuning,
) -> Result<Arc<HubConnection>, 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<HubConnection>` 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<HubConnection>`, 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<HubConnectionPool>` 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<Self>,
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<HubConnection>) -> bool,
) {
if self
.connections
.remove_if(key, |_, pooled| predicate(&pooled.conn))
.is_some()
{
crate::metrics::pool_connections_dec();
}
}
}
crates/common/kigi-computer-hub-sdk/src/refcount.rs
-123
View File
@@ -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<K: Eq + Hash> {
counts: DashMap<K, u64>,
}
impl<K: Eq + Hash> RefCountedSet<K> {
/// 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<u64> {
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<K>
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));
}
}
crates/common/kigi-computer-hub-sdk/src/server.rs
-2649
View File
File diff suppressed because it is too large Load Diff
crates/common/kigi-computer-hub-sdk/src/trace_donate.rs
-206
View File
@@ -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<SpanRecord>) {
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<PumpMsg>,
resource: ResourceAttributesWithSchema,
}
impl SpanExporter for PumpSpanExporter {
fn export(
&self,
batch: Vec<SpanData>,
) -> impl std::future::Future<Output = OTelSdkResult> + 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<PumpMsg>,
}
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<String>,
) -> (HubDonatingReporter, TraceDonationPump) {
let (tx, rx) = mpsc::channel::<PumpMsg>(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::<PumpMsg>(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"
);
}
}