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

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

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

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

Known §9 leftover (deliberate, for the M3 sweep): the BYOK default base
URL string. Gates: workspace check/clippy 0/0, fmt, deny ok; suites
green (workspace 1042, shell 4918, tui 6634, tools 2608, tool-runtime
47, mcp 154).
This commit is contained in:
ZacharyZhang-NY
2026-07-17 22:34:10 -04:00
parent 5919526e91
commit fa75eb139a
90 changed files with 452 additions and 9702 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/codegen/kigi-workspace/src/config.rs
-863
View File
@@ -1,6 +1,5 @@
//! Workspace and session configuration types.
use crate::capability::CapabilityMode;
use crate::hub::HubConfig;
use kigi_tools::registry::types::{SessionContext, ToolRegistryBuilder, ToolServerConfig};
use std::collections::HashMap;
use std::path::PathBuf;
@@ -97,605 +96,6 @@ pub trait SessionContextFactory: Send + Sync {
#[derive(Debug, Clone, Default)]
#[non_exhaustive]
pub struct MemoryConfig {}
/// Per-session toolset/capability selection from the `session.bind`
/// metadata. Absent fields fall back to the workspace default and `CapabilityMode::All`.
#[derive(Debug, Clone, Default)]
#[non_exhaustive]
pub struct WorkspaceBindConfig {
/// Named toolset preset from the wire. **Never resolved** (see
/// [`Self::resolve`]); parsed only so it can be logged.
pub preset: Option<String>,
/// Capability mode applied to the session's toolset.
pub capability_mode: Option<CapabilityMode>,
/// Fully-specified toolset in the runtime serde shape. Takes precedence
/// over `tools`.
pub tool_config: Option<ToolServerConfig>,
/// Per-user feature-flag bag. `None` on legacy payloads → tools
/// fall back to their safe defaults.
pub viewer_ctx: Option<kigi_tool_runtime::WorkspaceViewerContext>,
/// Initial auto-approve (YOLO) state. `None` on legacy payloads →
/// fail-closed (false).
pub yolo_mode: Option<bool>,
/// Plane-configured toolset in the gRPC wire shape. An empty list is
/// treated as unset (proto3 repeated default).
pub tools: Option<Vec<kigi_tools_api::ToolConfigEntry>>,
pub manifest_version: Option<String>,
pub manifest_hash: Option<String>,
/// Opt-in: forward `BackgroundTaskCompleted` system notifications for this session.
pub system_notifications: bool,
pub rpc_only: bool,
}
/// Outcome of resolving a [`WorkspaceBindConfig`]; lets callers fail closed
/// instead of widening to the default toolset. Deliberately has **no preset
/// arm** (see [`WorkspaceBindConfig::resolve`]).
#[derive(Debug)]
pub enum ResolvedToolset {
/// An explicit toolset (`tool_config` or `tools`).
Toolset(ResolvedTools),
/// No explicit toolset was specified and the workspace allows falling
/// back to its default catalog (local/CLI embedders only).
UseDefault,
/// No explicit toolset was specified and the workspace requires one
/// (sandbox-launched standalone servers) — fail closed.
MissingToolConfig,
/// `tools` entries were specified but at least one failed to convert.
InvalidToolConfig(kigi_tools::registry::proto_convert::ToolConfigEntryError),
}
/// A resolved toolset plus the pinned entries this binary could not serve.
#[derive(Debug)]
pub struct ResolvedTools {
pub toolset: ToolServerConfig,
/// Pinned `tools` ids unknown to this binary's registry, sorted. Always
/// empty for `tool_config` resolutions.
pub unserved_tool_ids: Vec<String>,
}
impl ResolvedTools {
/// A fully-served toolset (no divergence).
fn full(toolset: ToolServerConfig) -> Self {
Self {
toolset,
unserved_tool_ids: Vec::new(),
}
}
}
impl WorkspaceBindConfig {
/// Parse hub `session.bind` metadata. The envelope is the shared
/// [`kigi_tool_runtime::WorkspaceBindMetadata`] (same type the emitter
/// serializes); `tool_config` is a consumer-only raw escape hatch read
/// separately.
pub fn from_metadata(metadata: &serde_json::Value) -> Self {
let wire: kigi_tool_runtime::WorkspaceBindMetadata =
serde_json::from_value(metadata.clone()).unwrap_or_default();
Self {
preset: wire.preset,
capability_mode: wire
.capability_mode
.and_then(|s| serde_json::from_value(serde_json::Value::String(s)).ok()),
tool_config: metadata
.as_object()
.and_then(|obj| obj.get("tool_config"))
.and_then(|v| parse_field("tool_config", v)),
viewer_ctx: wire.viewer_ctx,
yolo_mode: wire.yolo_mode,
tools: Some(wire.tools).filter(|tools| !tools.is_empty()),
manifest_version: wire.manifest_version,
manifest_hash: wire.manifest_hash,
system_notifications: wire.system_notifications.unwrap_or(false),
rpc_only: wire.rpc_only,
}
}
/// Resolve the selected toolset.
///
/// Precedence: `tool_config` > `tools` (wire entries) > default/fail-closed.
/// Pinned `tools` are served per entry: ids `known_id` rejects are dropped
/// and reported in [`ResolvedTools::unserved_tool_ids`] instead of
/// silently falling back to a different toolset.
///
/// **Presets are never resolved** — a `preset` on the wire is logged and
/// ignored; only explicit `tools`/`tool_config` may select a toolset.
///
/// With `require_explicit_toolset` (sandbox standalone servers) a bind
/// without an explicit toolset fails closed instead of widening to the
/// binary's default catalog.
pub fn resolve(
&self,
known_id: &dyn Fn(&str) -> bool,
require_explicit_toolset: bool,
) -> ResolvedToolset {
if let Some(cfg) = &self.tool_config {
for (idx, tool) in cfg.tools.iter().enumerate() {
if let Err(err) = kigi_tools_api::config_validation::validate_name_override(
idx,
&tool.id,
tool.name_override.as_deref(),
) {
return ResolvedToolset::InvalidToolConfig(err);
}
}
return ResolvedToolset::Toolset(ResolvedTools::full(cfg.clone()));
}
if let Some(tools) = &self.tools {
let mut unserved_tool_ids: Vec<String> = Vec::new();
let mut served = Vec::with_capacity(tools.len());
for (idx, entry) in tools.iter().enumerate() {
if !known_id(&entry.id) {
unserved_tool_ids.push(entry.id.clone());
continue;
}
match kigi_tools::registry::proto_convert::tool_config_from_entry(
idx,
entry.clone(),
) {
Ok(tc) => served.push(tc),
Err(err) => return ResolvedToolset::InvalidToolConfig(err),
}
}
unserved_tool_ids.sort_unstable();
if !unserved_tool_ids.is_empty() {
tracing::warn!(
unserved = ? unserved_tool_ids, config_manifest_version = ? self
.manifest_version, running_version = kigi_version::VERSION,
"session.bind: serving known subset of pinned tools"
);
}
return ResolvedToolset::Toolset(ResolvedTools {
toolset: ToolServerConfig {
tools: served,
behavior_preset: None,
},
unserved_tool_ids,
});
}
if let Some(preset) = self.preset.as_deref() {
tracing::warn!(
preset,
"session.bind: toolset presets are not resolved by the workspace \
server; pass an explicit `tools` config"
);
}
if require_explicit_toolset {
ResolvedToolset::MissingToolConfig
} else {
ResolvedToolset::UseDefault
}
}
}
/// Parse a single bind-metadata field, ignoring (and logging) a malformed value.
fn parse_field<T: serde::de::DeserializeOwned>(name: &str, value: &serde_json::Value) -> Option<T> {
match serde_json::from_value(value.clone()) {
Ok(parsed) => Some(parsed),
Err(e) => {
tracing::warn!(
field = name, error = % e,
"session.bind metadata: ignoring malformed field"
);
None
}
}
}
#[cfg(test)]
mod bind_config_tests {
use super::*;
/// Predicate for tests where every pinned id is known to the binary.
fn all_known(_: &str) -> bool {
true
}
/// Predicate for tests simulating a binary that knows none of the ids.
fn none_known(_: &str) -> bool {
false
}
#[test]
fn parses_preset_and_capability() {
let v = serde_json::json!(
{ "preset" : "explore", "capability_mode" : "read_only" }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert_eq!(cfg.preset.as_deref(), Some("explore"));
assert_eq!(
cfg.capability_mode,
Some(crate::capability::CapabilityMode::ReadOnly)
);
}
#[test]
fn defaults_on_empty_or_mismatched_metadata() {
let empty = WorkspaceBindConfig::from_metadata(&serde_json::json!({}));
assert!(empty.preset.is_none());
assert!(empty.capability_mode.is_none());
assert!(matches!(
empty.resolve(&all_known, false),
ResolvedToolset::UseDefault
));
let weird = WorkspaceBindConfig::from_metadata(&serde_json::json!("hello"));
assert!(matches!(
weird.resolve(&all_known, false),
ResolvedToolset::UseDefault
));
}
/// Presets are banned: any preset (known or not) is ignored — never
/// resolved to a toolset, and never widened to the default in strict mode.
#[test]
fn presets_are_never_resolved() {
for preset in ["explore", "grok-computer", "bogus"] {
let cfg = WorkspaceBindConfig::from_metadata(&serde_json::json!({ "preset" : preset }));
assert!(
matches!(cfg.resolve(&all_known, false), ResolvedToolset::UseDefault),
"lax mode must fall through to the default, preset={preset}"
);
assert!(
matches!(
cfg.resolve(&all_known, true),
ResolvedToolset::MissingToolConfig
),
"strict mode must fail closed, preset={preset}"
);
}
}
/// Strict mode (sandbox standalone server): no explicit toolset on the
/// bind ⇒ fail closed instead of widening to the default catalog.
#[test]
fn strict_mode_requires_explicit_toolset() {
let empty = WorkspaceBindConfig::from_metadata(&serde_json::json!({}));
assert!(matches!(
empty.resolve(&all_known, true),
ResolvedToolset::MissingToolConfig
));
}
#[test]
fn malformed_field_does_not_discard_valid_siblings() {
let v = serde_json::json!(
{ "preset" : "explore", "capability_mode" : "raed_only" }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert_eq!(cfg.preset.as_deref(), Some("explore"));
assert!(cfg.capability_mode.is_none());
}
#[test]
fn workspace_bind_config_from_metadata_extracts_viewer_ctx() {
let v = serde_json::json!(
{ "preset" : "explore", "viewer_ctx" : { "stream_tool_progress" : true }, }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert_eq!(cfg.preset.as_deref(), Some("explore"));
let viewer = cfg.viewer_ctx.expect("viewer_ctx parsed");
assert!(viewer.stream_tool_progress);
}
/// Legacy payload without `viewer_ctx` still parses (mixed-version
/// proxy/workspace deploys).
#[test]
fn workspace_bind_config_from_metadata_legacy_omitted_viewer_ctx() {
let v = serde_json::json!({ "preset" : "explore" });
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert!(cfg.viewer_ctx.is_none());
}
#[test]
fn workspace_bind_config_from_metadata_extracts_yolo_mode() {
let v = serde_json::json!({ "preset" : "explore", "yolo_mode" : true });
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert_eq!(cfg.yolo_mode, Some(true));
}
#[test]
fn workspace_bind_config_yolo_mode_omitted_or_malformed_fails_closed() {
let omitted =
WorkspaceBindConfig::from_metadata(&serde_json::json!({ "preset" : "explore" }));
assert!(omitted.yolo_mode.is_none());
let malformed = WorkspaceBindConfig::from_metadata(
&serde_json::json!({ "preset" : "explore", "yolo_mode" : "yes" }),
);
assert!(malformed.yolo_mode.is_none());
assert_eq!(malformed.preset.as_deref(), Some("explore"));
}
#[test]
fn workspace_bind_config_extracts_system_notifications_flag() {
let on = WorkspaceBindConfig::from_metadata(
&serde_json::json!({ "system_notifications" : true }),
);
assert!(on.system_notifications);
let off = WorkspaceBindConfig::from_metadata(&serde_json::json!({ "preset" : "explore" }));
assert!(!off.system_notifications);
let explicit_off = WorkspaceBindConfig::from_metadata(
&serde_json::json!({ "system_notifications" : false }),
);
assert!(!explicit_off.system_notifications);
}
#[test]
fn workspace_bind_config_extracts_rpc_only_flag() {
let on = WorkspaceBindConfig::from_metadata(&serde_json::json!({ "rpc_only" : true }));
assert!(on.rpc_only);
let off = WorkspaceBindConfig::from_metadata(&serde_json::json!({ "preset" : "explore" }));
assert!(!off.rpc_only);
let explicit_off =
WorkspaceBindConfig::from_metadata(&serde_json::json!({ "rpc_only" : false }));
assert!(!explicit_off.rpc_only);
}
#[test]
fn workspace_bind_config_from_metadata_extracts_manifest_fields() {
let v = serde_json::json!(
{ "preset" : "explore", "manifest_version" : "v1", "manifest_hash" :
"abc123", }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert_eq!(cfg.manifest_version.as_deref(), Some("v1"));
assert_eq!(cfg.manifest_hash.as_deref(), Some("abc123"));
}
#[test]
fn workspace_bind_config_manifest_fields_default_to_none_when_absent() {
let cfg = WorkspaceBindConfig::from_metadata(&serde_json::json!({ "preset" : "explore" }));
assert!(cfg.manifest_version.is_none());
assert!(cfg.manifest_hash.is_none());
}
/// Consumer-side parity test for the bind-metadata `tools` contract;
/// pairs with the producer-side pin test in agentic-sampler's
/// `configs::plane` tests.
#[test]
fn tools_entries_resolve_to_tool_server_config() {
let v = serde_json::json!(
{ "preset" : "explore", "tools" : [{ "id" : "GrokBuild:grep", "params_json" :
"{\"max_results\":50}", "name_override" : "search", "params_name_overrides" :
{ "pattern" : "query" }, "behavior_version" : "legacy-0.4.10",
"description_override" : "Search the codebase", }, { "id" :
"GrokBuild:read_file" },], }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&all_known, false) else {
panic!("tools entries must resolve to an explicit toolset");
};
assert!(resolved.unserved_tool_ids.is_empty());
let toolset = resolved.toolset;
assert_eq!(
toolset.behavior_preset, None,
"always the 'current' default"
);
assert_eq!(toolset.tools.len(), 2);
let grep = &toolset.tools[0];
assert_eq!(grep.id, "GrokBuild:grep");
assert_eq!(
grep.params,
serde_json::json!({ "max_results" : 50 })
.as_object()
.cloned()
);
assert_eq!(grep.name_override.as_deref(), Some("search"));
assert_eq!(
grep.params_name_overrides.as_ref().unwrap()["pattern"],
"query"
);
assert_eq!(grep.behavior_version.as_deref(), Some("legacy-0.4.10"));
assert_eq!(
grep.description_override.as_deref(),
Some("Search the codebase")
);
assert_eq!(grep.kind, None);
assert_eq!(toolset.tools[1].id, "GrokBuild:read_file");
}
#[test]
fn explicit_tool_config_wins_over_tools_entries() {
let v = serde_json::json!(
{ "tool_config" : { "tools" : [{ "id" : "raw:tool" }] }, "tools" : [{ "id" :
"wire:tool" }], }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&all_known, false) else {
panic!("must resolve to a toolset");
};
assert_eq!(resolved.toolset.tools.len(), 1);
assert_eq!(resolved.toolset.tools[0].id, "raw:tool");
}
#[test]
fn tools_entries_win_even_with_preset_present() {
let v = serde_json::json!(
{ "preset" : "explore", "tools" : [{ "id" : "wire:tool" }], }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&all_known, false) else {
panic!("must resolve to a toolset");
};
assert_eq!(resolved.toolset.tools.len(), 1);
assert_eq!(resolved.toolset.tools[0].id, "wire:tool");
}
#[test]
fn empty_tools_array_is_treated_as_unset() {
let v = serde_json::json!({ "preset" : "explore", "tools" : [] });
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert!(cfg.tools.is_none());
assert!(matches!(
cfg.resolve(&all_known, false),
ResolvedToolset::UseDefault
));
assert!(matches!(
cfg.resolve(&all_known, true),
ResolvedToolset::MissingToolConfig
));
let no_preset = serde_json::json!({ "tools" : [] });
let cfg = WorkspaceBindConfig::from_metadata(&no_preset);
assert!(matches!(
cfg.resolve(&all_known, false),
ResolvedToolset::UseDefault
));
}
#[test]
fn invalid_tools_entry_fails_closed() {
let v = serde_json::json!(
{ "preset" : "explore", "tools" : [{ "id" : "bad:tool", "params_json" :
"{not json" }], }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
match cfg.resolve(&all_known, false) {
ResolvedToolset::InvalidToolConfig(err) => {
assert_eq!(err.tool_id, "bad:tool");
assert_eq!(err.index, 0);
}
other => panic!("expected InvalidToolConfig, got {other:?}"),
}
}
#[test]
fn invalid_name_override_fails_closed() {
let v = serde_json::json!(
{ "tools" : [{ "id" : "wire:ok", "name_override" : "fine_name" }, { "id" :
"wire:bad", "name_override" : "not a tool id!" },], }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
match cfg.resolve(&all_known, true) {
ResolvedToolset::InvalidToolConfig(err) => {
assert_eq!(err.tool_id, "wire:bad");
assert_eq!(err.field_path(), "tools[1].name_override");
}
other => panic!("expected InvalidToolConfig, got {other:?}"),
}
}
#[test]
fn tool_config_escape_hatch_invalid_name_override_fails_closed() {
let v = serde_json::json!(
{ "tool_config" : { "tools" : [{ "id" : "raw:ok", "name_override" :
"fine_name" }, { "id" : "raw:bad", "name_override" : "not a tool id!" },] },
}
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
match cfg.resolve(&all_known, true) {
ResolvedToolset::InvalidToolConfig(err) => {
assert_eq!(err.tool_id, "raw:bad");
assert_eq!(err.field_path(), "tools[1].name_override");
}
other => panic!("expected InvalidToolConfig, got {other:?}"),
}
let v = serde_json::json!(
{ "tool_config" : { "tools" : [{ "id" : "raw:ok", "name_override" :
"fine_name" }] }, }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&all_known, true) else {
panic!("valid escape-hatch config must resolve");
};
assert_eq!(resolved.toolset.tools.len(), 1);
}
#[test]
fn invalid_entry_error_reports_wire_index_after_unknown_drop() {
let v = serde_json::json!(
{ "tools" : [{ "id" : "wire:unknown" }, { "id" : "wire:bad", "params_json" :
"{not json" },], }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let known = |id: &str| id != "wire:unknown";
match cfg.resolve(&known, false) {
ResolvedToolset::InvalidToolConfig(err) => {
assert_eq!(err.tool_id, "wire:bad");
assert_eq!(
err.index, 1,
"index must be the wire position, not the known-subset position"
);
}
other => panic!("expected InvalidToolConfig, got {other:?}"),
}
}
#[test]
fn valid_name_overrides_resolve_intact() {
let v = serde_json::json!(
{ "tools" : [{ "id" : "wire:a", "name_override" : "renamed_a" }, { "id" :
"wire:b" },], }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&all_known, true) else {
panic!("well-formed overrides must resolve to a toolset");
};
assert_eq!(resolved.toolset.tools.len(), 2);
assert_eq!(
resolved.toolset.tools[0].name_override.as_deref(),
Some("renamed_a")
);
assert_eq!(resolved.toolset.tools[1].name_override, None);
}
#[test]
fn pinned_tools_all_known_serves_full_expansion() {
let v = serde_json::json!(
{ "preset" : "explore", "tools" : [{ "id" : "wire:tool" }],
"manifest_version" : "9.9.9-any", }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&all_known, false) else {
panic!("known pinned tools must use the tools expansion");
};
assert!(resolved.unserved_tool_ids.is_empty());
assert_eq!(resolved.toolset.tools.len(), 1);
assert_eq!(resolved.toolset.tools[0].id, "wire:tool");
}
/// Unknown ids must be partitioned and reported, never silently replaced
/// by live preset resolution.
#[test]
fn pinned_tools_unknown_ids_are_partitioned_and_reported() {
let v = serde_json::json!(
{ "preset" : "explore", "tools" : [{ "id" : "wire:known" }, { "id" :
"wire:zz_unknown" }, { "id" : "wire:aa_unknown" },], "manifest_version" :
"0.0.0-stale", }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let known = |id: &str| id == "wire:known";
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&known, false) else {
panic!("partial coverage must still resolve to the known subset");
};
assert_eq!(resolved.toolset.tools.len(), 1);
assert_eq!(resolved.toolset.tools[0].id, "wire:known");
assert_eq!(
resolved.unserved_tool_ids,
vec!["wire:aa_unknown".to_owned(), "wire:zz_unknown".to_owned()],
"unserved ids are reported sorted"
);
}
/// A fully-unknown expansion serves empty and reports every id — it never
/// widens to preset/default.
#[test]
fn pinned_tools_all_unknown_serves_empty_and_reports_all() {
let v = serde_json::json!(
{ "preset" : "explore", "tools" : [{ "id" : "wire:tool" }],
"manifest_version" : "0.0.0-stale", }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&none_known, false) else {
panic!("all-unknown expansion must resolve (empty), not fall back");
};
assert!(resolved.toolset.tools.is_empty());
assert_eq!(resolved.unserved_tool_ids, vec!["wire:tool".to_owned()]);
}
#[test]
fn legacy_tools_without_manifest_version_are_not_gated() {
let v = serde_json::json!(
{ "preset" : "explore", "tools" : [{ "id" : "wire:tool" }], }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert!(cfg.manifest_version.is_none());
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&all_known, false) else {
panic!("legacy unpinned tools must resolve without gating");
};
assert_eq!(resolved.toolset.tools.len(), 1);
assert_eq!(resolved.toolset.tools[0].id, "wire:tool");
}
#[test]
fn tool_config_wins_regardless_of_stale_manifest_version() {
let v = serde_json::json!(
{ "tool_config" : { "tools" : [{ "id" : "raw:tool" }] }, "tools" : [{ "id" :
"wire:tool" }], "manifest_version" : "0.0.0-stale", }
);
let cfg = WorkspaceBindConfig::from_metadata(&v);
let ResolvedToolset::Toolset(resolved) = cfg.resolve(&none_known, false) else {
panic!("tool_config must always win");
};
assert!(resolved.unserved_tool_ids.is_empty());
assert_eq!(resolved.toolset.tools.len(), 1);
assert_eq!(resolved.toolset.tools[0].id, "raw:tool");
}
#[test]
fn malformed_tools_field_is_dropped_keeping_siblings() {
let v = serde_json::json!({ "preset" : "explore", "tools" : "not-a-list" });
let cfg = WorkspaceBindConfig::from_metadata(&v);
assert!(cfg.tools.is_none());
assert!(matches!(
cfg.resolve(&all_known, false),
ResolvedToolset::UseDefault
));
assert!(matches!(
cfg.resolve(&all_known, true),
ResolvedToolset::MissingToolConfig
));
}
}
/// Top-level config required to construct a [`crate::handle::WorkspaceHandle`].
///
/// `#[non_exhaustive]` so future fields are non-breaking.
@@ -726,18 +126,6 @@ pub struct WorkspaceConfig {
/// and disabled/enabled lists. Stored on `WorkspaceShared` for
/// `discover_plugins` calls. Defaults to empty.
pub plugin_discovery_config: crate::discovery::PluginDiscoveryConfig,
/// Optional server configuration. When `Some`, the workspace
/// can connect to the server after construction via
/// [`WorkspaceHandle::connect_hub`](crate::handle::WorkspaceHandle::connect_hub).
pub hub_config: Option<HubConfig>,
/// Auth provider for xAI service calls made from workspace-scoped code.
/// `None` for workspaces that do not configure service auth.
pub auth_provider: Option<kigi_computer_hub_sdk::SharedAuthProvider>,
/// Metadata attached to the tool server registration.
/// Propagated through the server to `ServerInfo.metadata` in
/// `servers.list` responses so harness clients can identify the
/// sandbox that started the tool server.
pub server_metadata: Option<serde_json::Value>,
/// Runtime-tunable timing/threshold config for the tool server.
pub status_config: crate::status_config::StatusConfig,
/// Folder-trust verdict for repo-local (project-scoped) LSP servers from
@@ -745,108 +133,12 @@ pub struct WorkspaceConfig {
/// shell caller resolves the verdict and threads it in; callers without a
/// folder-trust decision pass `true`.
pub project_lsp_trusted: bool,
/// Fail `session.bind`s without an explicit toolset closed instead of
/// widening to `default_tool_config`. Set by sandbox-launched standalone
/// servers; local/CLI embedders keep the default-catalog fallback.
pub require_explicit_toolset: bool,
/// Confine `x.ai/fs/*` / `workspace.fs_*` resolution to the workspace root
/// (reject `..`, absolute-outside-root, symlink escapes). Default `false`
/// (unconfined) — set to `true` only by the workspace server on a remote
/// sandbox, where the root is a real tenant boundary.
pub confine_fs_to_workspace_root: bool,
}
/// Metadata a tool server announces so hub consumers can identify and route
/// to it. Every field is optional and independently sourced; a local process
/// announces none.
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
pub struct WorkspaceServerMetadata {
/// Sandbox that provisioned this server. Absent for local servers.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub sandbox_id: Option<String>,
/// Logical sandbox-service session UUID, from the `KIGI_SESSION_ID` env
/// var. Present whenever that var is set (every sandbox container, start
/// and restore), absent otherwise.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub session_id: Option<String>,
/// Provider that provisioned this server. Populated on the start path
/// only (no container-side source on restore); absent for local servers.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_id: Option<String>,
/// Per-spawn launch nonce minted by the sandbox orchestrator and echoed
/// verbatim on the diagnostics `/ready` endpoint. Absent for local/legacy
/// launches.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub launch_id: Option<String>,
}
impl WorkspaceServerMetadata {
/// Merge an env-sourced logical session id into caller-supplied
/// tool-server metadata (`None` on the restore/local path).
///
/// `env_session_id` is the raw `KIGI_SESSION_ID`; empty is normalized to
/// absent. An explicit `session_id` already in `metadata` is never
/// clobbered. A non-object `metadata` value is returned unchanged (a
/// defensive no-op — the sole caller always sends an object).
pub fn merge_session_metadata(
metadata: Option<serde_json::Value>,
env_session_id: Option<String>,
) -> Option<serde_json::Value> {
let env_session_id = env_session_id.filter(|s| !s.is_empty());
match metadata {
Some(mut value) => {
if let Some(session_id) = env_session_id
&& let Some(obj) = value.as_object_mut()
&& !obj.contains_key("session_id")
{
obj.insert(
"session_id".to_owned(),
serde_json::Value::String(session_id),
);
}
Some(value)
}
None => serde_json::to_value(WorkspaceServerMetadata {
sandbox_id: None,
session_id: env_session_id,
provider_id: None,
launch_id: None,
})
.ok(),
}
}
}
impl WorkspaceConfig {
/// Construct a minimal config suitable for proxy-mode workspaces
/// where the workspace is used primarily as a ToolServer host.
pub fn new_for_proxy(
root_cwd: PathBuf,
session_factory: Arc<dyn SessionContextFactory>,
hub_config: HubConfig,
auth_provider: kigi_computer_hub_sdk::SharedAuthProvider,
server_metadata: Option<serde_json::Value>,
status_config: crate::status_config::StatusConfig,
tool_config: ToolServerConfig,
) -> Self {
Self {
root_cwd,
default_tool_config: tool_config,
respect_gitignore: false,
memory_config: None,
event_buffer_capacity: crate::config::DEFAULT_EVENT_BUFFER_CAPACITY,
session_factory,
hook_global_sources: vec![],
hook_project_sources: vec![],
skills_config: Default::default(),
plugin_discovery_config: Default::default(),
auth_provider: Some(auth_provider),
hub_config: Some(hub_config),
server_metadata,
project_lsp_trusted: true,
require_explicit_toolset: false,
confine_fs_to_workspace_root: false,
status_config,
}
}
}
/// Configuration for spawning a subagent session within a workspace.
#[derive(Clone)]
#[non_exhaustive]
@@ -928,158 +220,3 @@ pub enum IsolationMode {
/// Run the subagent inside a sandbox/container.
Sandbox,
}
#[cfg(test)]
mod tests {
use super::WorkspaceServerMetadata;
#[test]
fn workspace_server_metadata_serializes_all_present_fields() {
let meta = WorkspaceServerMetadata {
sandbox_id: Some("sb-123".to_owned()),
session_id: Some("11111111-1111-1111-1111-111111111111".to_owned()),
provider_id: Some("test-provider".to_owned()),
launch_id: Some("33333333-3333-3333-3333-333333333333".to_owned()),
};
let value = serde_json::to_value(&meta).unwrap();
assert_eq!(
value,
serde_json::json!({ "sandbox_id" : "sb-123", "session_id" :
"11111111-1111-1111-1111-111111111111", "provider_id" : "test-provider",
"launch_id" : "33333333-3333-3333-3333-333333333333", })
);
}
#[test]
fn workspace_server_metadata_omits_none_fields() {
let meta = WorkspaceServerMetadata {
sandbox_id: Some("sb-123".to_owned()),
session_id: None,
provider_id: None,
launch_id: None,
};
let value = serde_json::to_value(&meta).unwrap();
assert_eq!(value, serde_json::json!({ "sandbox_id" : "sb-123" }));
let empty = serde_json::to_value(WorkspaceServerMetadata::default()).unwrap();
assert_eq!(empty, serde_json::json!({}));
}
#[test]
fn workspace_server_metadata_deserializes_legacy_payload_without_new_fields() {
let legacy = serde_json::json!(
{ "sandbox_id" : "sb-legacy", "cwd" : "/workspace", "mode" : "remote", }
);
let meta: WorkspaceServerMetadata = serde_json::from_value(legacy).unwrap();
assert_eq!(meta.sandbox_id.as_deref(), Some("sb-legacy"));
assert_eq!(meta.session_id, None);
assert_eq!(meta.provider_id, None);
}
#[test]
fn workspace_server_metadata_round_trips_with_new_fields() {
let meta = WorkspaceServerMetadata {
sandbox_id: Some("sb-123".to_owned()),
session_id: Some("22222222-2222-2222-2222-222222222222".to_owned()),
provider_id: Some("test-provider".to_owned()),
launch_id: None,
};
let json = serde_json::to_string(&meta).unwrap();
let back: WorkspaceServerMetadata = serde_json::from_str(&json).unwrap();
assert_eq!(back.sandbox_id, meta.sandbox_id);
assert_eq!(back.session_id, meta.session_id);
assert_eq!(back.provider_id, meta.provider_id);
}
#[test]
fn workspace_server_metadata_deserializes_partial_new_fields() {
let only_session = serde_json::json!(
{ "sandbox_id" : "sb-1", "session_id" :
"33333333-3333-3333-3333-333333333333", }
);
let meta: WorkspaceServerMetadata = serde_json::from_value(only_session).unwrap();
assert_eq!(
meta.session_id.as_deref(),
Some("33333333-3333-3333-3333-333333333333")
);
assert_eq!(meta.provider_id, None);
let only_provider = serde_json::json!(
{ "sandbox_id" : "sb-1", "provider_id" : "test-provider", }
);
let meta: WorkspaceServerMetadata = serde_json::from_value(only_provider).unwrap();
assert_eq!(meta.provider_id.as_deref(), Some("test-provider"));
assert_eq!(meta.session_id, None);
}
#[test]
fn workspace_server_metadata_reads_start_path_shaped_payload() {
let start_path = serde_json::json!(
{ "cwd" : "/workspace", "mode" : "remote", "sandbox_id" : "sb-start",
"session_id" : "44444444-4444-4444-4444-444444444444", "provider_id" :
"test-provider", }
);
let meta: WorkspaceServerMetadata = serde_json::from_value(start_path).unwrap();
assert_eq!(meta.sandbox_id.as_deref(), Some("sb-start"));
assert_eq!(
meta.session_id.as_deref(),
Some("44444444-4444-4444-4444-444444444444")
);
assert_eq!(meta.provider_id.as_deref(), Some("test-provider"));
}
#[test]
fn merge_session_metadata_builds_struct_from_env_on_none_branch() {
let merged =
WorkspaceServerMetadata::merge_session_metadata(None, Some("sess-1".to_owned()))
.unwrap();
assert_eq!(merged, serde_json::json!({ "session_id" : "sess-1" }));
let empty = WorkspaceServerMetadata::merge_session_metadata(None, None).unwrap();
assert_eq!(empty, serde_json::json!({}));
}
#[test]
fn merge_session_metadata_overlays_into_object_without_clobbering() {
let base = serde_json::json!({ "sandbox_id" : "sb-9", "mode" : "remote" });
let merged =
WorkspaceServerMetadata::merge_session_metadata(Some(base), Some("env-id".to_owned()))
.unwrap();
assert_eq!(
merged,
serde_json::json!({ "sandbox_id" : "sb-9", "mode" : "remote",
"session_id" : "env-id", })
);
let explicit = serde_json::json!({ "session_id" : "explicit" });
let merged = WorkspaceServerMetadata::merge_session_metadata(
Some(explicit),
Some("env-id".to_owned()),
)
.unwrap();
assert_eq!(merged, serde_json::json!({ "session_id" : "explicit" }));
}
#[test]
fn merge_session_metadata_leaves_object_untouched_when_no_env_id() {
let base = serde_json::json!({ "sandbox_id" : "sb-9" });
let merged =
WorkspaceServerMetadata::merge_session_metadata(Some(base.clone()), None).unwrap();
assert_eq!(merged, base);
}
#[test]
fn merge_session_metadata_non_object_is_returned_unchanged() {
let scalar = serde_json::json!("just-a-string");
let merged = WorkspaceServerMetadata::merge_session_metadata(
Some(scalar.clone()),
Some("env-id".to_owned()),
)
.unwrap();
assert_eq!(merged, scalar);
}
#[test]
fn merge_session_metadata_treats_empty_env_id_as_absent() {
let none_branch =
WorkspaceServerMetadata::merge_session_metadata(None, Some(String::new())).unwrap();
assert_eq!(none_branch, serde_json::json!({}));
let base = serde_json::json!({ "sandbox_id" : "sb-9" });
let overlay = WorkspaceServerMetadata::merge_session_metadata(
Some(base.clone()),
Some(String::new()),
)
.unwrap();
assert_eq!(overlay, base);
}
#[test]
fn workspace_server_metadata_rejects_wrong_typed_field() {
let bad = serde_json::json!({ "sandbox_id" : "sb-1", "session_id" : 42 });
let result: Result<WorkspaceServerMetadata, _> = serde_json::from_value(bad);
assert!(result.is_err());
}
}
crates/codegen/kigi-workspace/src/error.rs
+3 -3
View File
@@ -53,9 +53,9 @@ pub enum WorkspaceError {
#[error("hunk action failed: {0}")]
HunkActionFailed(String),
/// An error from the server connection or tool server.
#[error("hub error: {0}")]
HubError(String),
/// An internal workspace error.
#[error("workspace error: {0}")]
Internal(String),
/// Deploy-service error tagged with its gRPC status class; see
/// [`DeployError`] for how the class crosses the workspace RPC boundary.
crates/codegen/kigi-workspace/src/file_system/ext_fs.rs
+8 -8
View File
@@ -62,7 +62,7 @@ impl WorkspaceOp for FsListReq {
let req = self.clone();
tokio::task::spawn_blocking(move || list(&abs, &req, confine_root))
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?
.map_err(|e| WorkspaceError::Internal(e.to_string()))?
}
}
@@ -98,7 +98,7 @@ impl WorkspaceOp for FsReadFileReq {
if !ranged {
let bytes = tokio::fs::read(&abs)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
return Ok(build_file_entry(&bytes));
}
@@ -106,9 +106,9 @@ impl WorkspaceOp for FsReadFileReq {
// chunk is `[offset, offset + min(length, max_bytes, cap))`.
let md = tokio::fs::metadata(&abs)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
if md.is_dir() {
return Err(WorkspaceError::HubError(format!(
return Err(WorkspaceError::Internal(format!(
"not a file: {}",
self.path
)));
@@ -120,7 +120,7 @@ impl WorkspaceOp for FsReadFileReq {
let length = super::walk::clamp_read_length(self.length, self.max_bytes);
let chunk = super::walk::read_range(&abs, offset, length)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
Ok(build_ranged_entry(chunk, size, self.encoding))
}
}
@@ -143,8 +143,8 @@ impl WorkspaceOp for FsWriteFileReq {
std::fs::write(&abs, content.as_bytes())
})
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
.map_err(|e| WorkspaceError::Internal(e.to_string()))?
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
Ok(())
}
}
@@ -160,7 +160,7 @@ impl WorkspaceOp for FsDeleteFileReq {
let (abs, _) = ws.confine_to_workspace_root(&abs_unconfined).await?;
tokio::fs::remove_file(&abs)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
Ok(())
}
}
crates/codegen/kigi-workspace/src/file_system/mod.rs
+1 -6
View File
@@ -7,13 +7,8 @@ pub use ext_fs::{
FsReadFileReq, FsWriteFileReq,
};
// Client-facing read-only fs ops (`workspace.client_fs_*`). Not re-exported:
// its wire types live in `kigi_workspace_types::rpc::fs` (the `ClientFs*`
// types) and would collide with the shell-facing `ext_fs` names above.
pub(crate) mod client_fs;
// Shared filesystem core: paginated listing + binary-safe ranged reads,
// used by `client_fs`, `ext_fs`, and the shell-local `session::file_system`.
// used by `ext_fs` and the shell-local `session::file_system`.
mod walk;
pub use walk::{
ChunkPayload, ListOptions, ListPage, ListedEntry, MAX_LIST_COLLECT, MAX_READ_BYTES,
crates/codegen/kigi-workspace/src/handle.rs
+45 -3930
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-workspace/src/lib.rs
+7 -29
View File
@@ -10,8 +10,6 @@ pub mod activity;
pub mod capability;
pub mod channel;
pub mod config;
pub mod daemonize;
pub mod diag_server;
pub mod discovery;
pub mod envrc;
pub mod error;
@@ -20,17 +18,8 @@ pub mod folder_trust;
pub mod foreign_sessions;
pub mod fs_notify;
pub mod handle;
pub mod hub;
pub mod hub_auth;
pub mod hub_channel;
pub mod hub_ids;
pub mod hub_server;
pub mod mcp;
pub mod permission;
pub mod preview_supervisor;
pub mod project_config;
pub mod recovery;
pub mod rpc_envelope;
pub mod session;
pub mod status_config;
pub use status_config::StatusConfig;
@@ -46,25 +35,18 @@ pub use config::{
};
pub use error::{WorkspaceError, WorkspaceResult};
pub use file_system::*;
pub use handle::{
DrainOutcome, DrainReason, WorkspaceHandle, connect_local_workspace, resolve_workspace_home,
termination_grace_from_env,
};
pub use hub::HubConfig;
pub use handle::WorkspaceHandle;
pub use kigi_hunk_tracker::HunkTrackerHandle;
pub use kigi_workspace_client::WorkspaceClient;
pub use kigi_workspace_types::WorkspaceEvent;
pub use permission::*;
pub use session::{WorkspaceSession, WorkspaceShared};
pub use session::{file_state, git, jj};
pub use workspace_ops::{WorkspaceOp, WorkspaceOps};
/// Zero-init every workspace metric family so idle panels render a `0` baseline
/// instead of "No data". Idempotent; call once at workspace-server startup.
/// instead of "No data". Idempotent; call once at startup.
pub fn init_metrics() {
handle::init_metrics();
session::swap_policy::init_metrics();
permission::init_metrics();
hub_server::init_metrics();
}
/// Crate-wide lock serializing every test that mutates the process-global
/// environment (`KIGI_SHARE_DIR`, `HOME`, …). nextest isolates each test in its own
@@ -167,22 +149,18 @@ mod init_metrics_tests {
})
})
};
assert!(has(
"grok_workspace_rpc_requests_total",
&[("method", "unknown"), ("result", "error")]
));
assert!(has(
"grok_workspace_drain_started_total",
&[("reason", "sigterm")]
));
assert!(has(
"grok_workspace_toolset_swap_rejected_total",
&[("reason", "turn_active"), ("trigger", "update_tool_config")]
));
assert!(has(
"grok_workspace_rewind_checkpoint_capture_total",
&[("domain", "fs"), ("outcome", "completed")]
));
assert!(
families
.iter()
.any(|mf| mf.name() == "grok_workspace_permission_timeout_total")
.any(|mf| mf.name() == "grok_workspace_terminal_backend_orphaned_total")
);
}
}
crates/codegen/kigi-workspace/src/permission/manager.rs
+4 -228
View File
@@ -820,6 +820,10 @@ fn session_grant_pre_decision(
/// Spawns the permission manager actor, returning a handle and the telemetry
/// event receiver.
///
/// `remember_tool_approvals` — resolved gate: shows the per-tool always-allow
/// options and lets an explicit grant satisfy an `ask` rule (ask once, remember).
#[allow(clippy::too_many_arguments)]
pub fn spawn_permission_manager(
session_id: acp::SessionId,
gateway: GatewaySender,
@@ -833,45 +837,7 @@ pub fn spawn_permission_manager(
web_fetch_allowed_domains: Vec<String>,
initial_yolo: bool,
client_identifier: Option<String>,
) -> (PermissionHandle, mpsc::UnboundedReceiver<PermissionEvent>) {
spawn_permission_manager_with_hub(
session_id,
gateway,
cwd,
client_type,
permission_config,
deny_read_globs,
web_fetch_allowed_domains,
initial_yolo,
client_identifier,
// Legacy/test entry point: preserve the full option set. Production uses
// `spawn_permission_manager_with_hub` with the resolved gate.
true,
None,
)
}
/// Like [`spawn_permission_manager`] but routes the permission prompt to chat
/// over the server (the HITL live path) when `hub_permission` is `Some`. The
/// caller builds the transport only when [`hitl_permission_live_enabled`] and a
/// server is connected; `None` keeps the local ACP prompt.
///
/// [`hitl_permission_live_enabled`]: crate::permission::hitl_permission_live_enabled
#[allow(clippy::too_many_arguments)]
pub fn spawn_permission_manager_with_hub(
session_id: acp::SessionId,
gateway: GatewaySender,
cwd: AbsPathBuf,
client_type: ClientType,
permission_config: Option<crate::permission::types::PermissionConfig>,
deny_read_globs: Vec<String>,
web_fetch_allowed_domains: Vec<String>,
initial_yolo: bool,
client_identifier: Option<String>,
// Resolved `remember_tool_approvals` gate: shows the per-tool always-allow
// options and lets an explicit grant satisfy an `ask` rule (ask once, remember).
remember_tool_approvals: bool,
hub_permission: Option<Arc<dyn crate::permission::PermissionHookTransport>>,
) -> (PermissionHandle, mpsc::UnboundedReceiver<PermissionEvent>) {
// Read the pin ONCE (file I/O) and cache it; never re-read per tool-call.
// Every yolo ingestion path funnels through construction or SetYoloMode.
@@ -887,7 +853,6 @@ pub fn spawn_permission_manager_with_hub(
client_identifier,
remember_tool_approvals,
crate::permission::resolution::yolo_disabled_by_policy(),
hub_permission,
)
}
@@ -905,7 +870,6 @@ fn spawn_permission_manager_with_pin(
client_identifier: Option<String>,
remember_tool_approvals: bool,
yolo_pin: Option<&'static str>,
hub_permission: Option<Arc<dyn crate::permission::PermissionHookTransport>>,
) -> (PermissionHandle, mpsc::UnboundedReceiver<PermissionEvent>) {
let (tx, mut rx) = mpsc::unbounded_channel::<PermissionCommand>();
let (event_tx, event_rx) = mpsc::unbounded_channel::<PermissionEvent>();
@@ -969,7 +933,6 @@ fn spawn_permission_manager_with_pin(
}
let prompter = AcpPrompter::new(session_id.clone(), gateway.clone(), client_type)
.with_hub_permission(hub_permission)
.with_remember_tool_approvals(remember_tool_approvals);
let mut yolo_mode = initial_yolo;
let mut auto_mode = seed_auto;
@@ -1849,7 +1812,6 @@ mod tests {
None,
true,
yolo_pin,
None,
)
}
@@ -1871,7 +1833,6 @@ mod tests {
None,
true,
None,
None,
)
}
@@ -1939,187 +1900,6 @@ mod tests {
.await;
}
/// Like [`test_manager`] but routes prompts through a hub permission transport.
fn test_manager_with_hub(
cwd: &AbsPathBuf,
hub_permission: Arc<dyn crate::permission::PermissionHookTransport>,
) -> (PermissionHandle, mpsc::UnboundedReceiver<PermissionEvent>) {
let (tx, _rx) = mpsc::unbounded_channel();
spawn_permission_manager_with_pin(
acp::SessionId::new(Arc::from("test-session")),
GatewaySender::new(tx),
cwd.clone(),
ClientType::Generic,
None,
vec![],
vec![],
false,
None,
true,
None,
Some(hub_permission),
)
}
/// Records every emitted payload and replies with a canned decision, so the
/// hub permission prompt path is exercised without a live hub.
struct FakeHubTransport {
reply: serde_json::Value,
seen: std::sync::Mutex<Vec<serde_json::Value>>,
}
#[async_trait::async_trait]
impl crate::permission::PermissionHookTransport for FakeHubTransport {
async fn request_permission(
&self,
payload: serde_json::Value,
) -> Result<serde_json::Value, String> {
self.seen.lock().unwrap().push(payload);
Ok(self.reply.clone())
}
}
fn fake_hub(reply: serde_json::Value) -> Arc<FakeHubTransport> {
Arc::new(FakeHubTransport {
reply,
seen: std::sync::Mutex::new(Vec::new()),
})
}
#[tokio::test]
async fn hub_permission_approve_allows_and_emits_payload() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let tmp = tempfile::tempdir().unwrap();
let cwd = AbsPathBuf::new(tmp.path().to_path_buf()).unwrap();
let transport = fake_hub(serde_json::json!({ "outcome": "approve" }));
let (mgr, _e) = test_manager_with_hub(&cwd, transport.clone());
let d = mgr
.request(
AccessKind::Edit("src/main.rs".into()),
tool_call(),
None,
None,
None,
)
.await;
assert_eq!(d, Decision::Allow);
let seen = transport.seen.lock().unwrap();
assert_eq!(seen.len(), 1, "exactly one permission hook emitted");
assert_eq!(seen[0]["tool_call_id"], "tc");
assert_eq!(seen[0]["tool_name"], "search_replace");
assert_eq!(seen[0]["description"], "Edit src/main.rs");
assert_eq!(seen[0]["scope"], "write");
assert_eq!(
seen[0]["edit_file_paths"],
serde_json::json!(["src/main.rs"])
);
})
.await;
}
#[tokio::test]
async fn hub_permission_reject_aborts() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let tmp = tempfile::tempdir().unwrap();
let cwd = AbsPathBuf::new(tmp.path().to_path_buf()).unwrap();
let (mgr, _e) = test_manager_with_hub(
&cwd,
fake_hub(serde_json::json!({ "outcome": "reject" })),
);
let d = mgr
.request(
AccessKind::Edit("a.rs".into()),
tool_call(),
None,
None,
None,
)
.await;
assert!(
matches!(d, Decision::Reject(_)),
"reject must abort, got {d:?}"
);
})
.await;
}
/// `cancelled` reply (turn-end drain) → abort, distinct from a user reject.
#[tokio::test]
async fn hub_permission_cancelled_aborts_distinctly() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let tmp = tempfile::tempdir().unwrap();
let cwd = AbsPathBuf::new(tmp.path().to_path_buf()).unwrap();
let (mgr, _e) = test_manager_with_hub(
&cwd,
fake_hub(serde_json::json!({ "outcome": "cancelled" })),
);
let d = mgr
.request(
AccessKind::Edit("a.rs".into()),
tool_call(),
None,
None,
None,
)
.await;
assert_eq!(d, Decision::Cancelled);
})
.await;
}
#[tokio::test]
async fn hub_permission_always_approve_persists_scope() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let tmp = tempfile::tempdir().unwrap();
let cwd = AbsPathBuf::new(tmp.path().to_path_buf()).unwrap();
let transport = fake_hub(serde_json::json!({
"outcome": "always_approve",
"scope": { "kind": "server_prefix", "value": "linear" },
}));
let (mgr, _e) = test_manager_with_hub(&cwd, transport.clone());
let first = mgr
.request(
AccessKind::MCPTool {
name: "linear__list".into(),
input: serde_json::Value::Null,
},
tool_call(),
None,
None,
None,
)
.await;
assert_eq!(first, Decision::Allow);
let second = mgr
.request(
AccessKind::MCPTool {
name: "linear__create".into(),
input: serde_json::Value::Null,
},
tool_call(),
None,
None,
None,
)
.await;
assert_eq!(second, Decision::Allow);
assert_eq!(
transport.seen.lock().unwrap().len(),
1,
"always_approve must persist so the second call needs no hook"
);
})
.await;
}
/// A managed `Ask` rule on a direct `Read`/`Grep` must reach the prompt, not
/// the unconditional auto-allow. With no responder wired, that surfaces as a
/// non-`Allow` decision; a non-ask read still auto-allows.
@@ -2484,7 +2264,6 @@ mod tests {
None,
true,
None,
None,
);
assert_eq!(
handle.deny_read_globs(),
@@ -2654,7 +2433,6 @@ mod tests {
None,
remember_tool_approvals,
None,
None,
)
}
@@ -3386,7 +3164,6 @@ mod tests {
None,
true,
None,
None,
);
let PermissionHandle::Actor { ref cmd_tx, .. } = mgr else {
panic!("manager must be actor-backed");
@@ -3577,7 +3354,6 @@ mod tests {
None,
true,
None,
None,
);
// Request A parks in the gated prompt; B then arrives and overlaps it.
crates/codegen/kigi-workspace/src/permission/mod.rs
+1 -14
View File
@@ -1,6 +1,5 @@
pub mod auto_mode;
pub mod claude_settings;
mod hub_permission;
mod manager;
mod policy;
mod prompter;
@@ -19,19 +18,7 @@ pub use auto_mode::{
classifier_output_json_schema, default_auto_mode_classifier, is_auto_mode_allowlisted_access,
is_auto_mode_allowlisted_tool_name, parse_classifier_model_text, permission_decision_args,
};
pub use hub_permission::{
PermissionHookTransport, ToolServerPermissionTransport, access_kind_for_hub_tool,
hitl_permission_live_enabled, prompt_outcome_allows, request_permission_via_hub,
};
/// Zero-init this module's metric families. See [`crate::init_metrics`].
pub(crate) fn init_metrics() {
hub_permission::init_metrics();
}
pub use manager::{
PermissionHandle, default_always_allow_scope, spawn_permission_manager,
spawn_permission_manager_with_hub,
};
pub use manager::{PermissionHandle, default_always_allow_scope, spawn_permission_manager};
pub use policy::CompiledPolicy;
pub use prompter::{
ALLOW_EDITS_SESSION_OPTION_ID, AcpPrompter, BashCommandPermission, BashCommandSelectedTerms,
crates/codegen/kigi-workspace/src/permission/prompter.rs
+23 -49
View File
@@ -321,9 +321,6 @@ pub struct AcpPrompter {
/// at decision-time through it. `EventWriter::noop()` when events recording
/// is disabled (the default for the permission scaffolding's own tests).
event_writer: EventWriter,
/// Server permission transport: when set, [`request`](Self::request) asks chat for the
/// decision over the server; `None` keeps the local prompt.
hub_permission: Option<Arc<dyn crate::permission::PermissionHookTransport>>,
/// When `false` (default, fail-safe), the per-tool "Always allow …" options
/// are stripped (see [`REMEMBER_TOOL_APPROVALS_GATED_IDS`]).
remember_tool_approvals: bool,
@@ -496,7 +493,6 @@ impl AcpPrompter {
// must NOT double-emit. A workspace-server-side caller that owns the
// per-session `events.jsonl` opts in via [`with_event_writer`].
event_writer: EventWriter::noop(),
hub_permission: None,
// Fail-safe default; opt in via `with_remember_tool_approvals`.
remember_tool_approvals: false,
}
@@ -509,16 +505,6 @@ impl AcpPrompter {
self
}
/// Route the permission prompt to chat over the server when `Some`;
/// `None` keeps the local prompt.
pub fn with_hub_permission(
mut self,
hub_permission: Option<Arc<dyn crate::permission::PermissionHookTransport>>,
) -> Self {
self.hub_permission = hub_permission;
self
}
/// Attach a per-session `events.jsonl` writer so [`request`](Self::request)
/// records `PermissionRequested` / `PermissionResolved`. Used by the
/// workspace-server permission path (which owns the session log); the shell
@@ -737,41 +723,29 @@ impl AcpPrompter {
prompt_start,
};
let outcome = match &self.hub_permission {
// Route the prompt to chat over the server (see
// `ToolServerPermissionTransport` for the await/release contract).
Some(transport) => {
crate::permission::hub_permission::request_permission_via_hub(
transport.as_ref(),
access,
tool_call_update.tool_call_id.0.as_ref(),
)
.await
}
None => {
let permission_options = self.build_options(access);
let req = acp::RequestPermissionRequest::new(
self.session_id.clone(),
tool_call_update.clone(),
permission_options.values().cloned().collect(),
)
.meta(self.bash_selection_meta(access));
match self.gateway.request_permission(req).await {
Ok(resp) => match resp.outcome {
acp::RequestPermissionOutcome::Cancelled => PromptOutcome::Cancelled,
acp::RequestPermissionOutcome::Selected(selected) => map_selected_outcome(
&permission_options,
&selected.option_id,
resp.meta.as_ref(),
access,
),
// TODO(acp-0.10): `RequestPermissionOutcome` is #[non_exhaustive].
_ => PromptOutcome::Error("unknown permission outcome".to_owned()),
},
Err(e) => {
tracing::error!(?e, "failed to request permission");
PromptOutcome::Error("failed to request permission".to_owned())
}
let outcome = {
let permission_options = self.build_options(access);
let req = acp::RequestPermissionRequest::new(
self.session_id.clone(),
tool_call_update.clone(),
permission_options.values().cloned().collect(),
)
.meta(self.bash_selection_meta(access));
match self.gateway.request_permission(req).await {
Ok(resp) => match resp.outcome {
acp::RequestPermissionOutcome::Cancelled => PromptOutcome::Cancelled,
acp::RequestPermissionOutcome::Selected(selected) => map_selected_outcome(
&permission_options,
&selected.option_id,
resp.meta.as_ref(),
access,
),
// TODO(acp-0.10): `RequestPermissionOutcome` is #[non_exhaustive].
_ => PromptOutcome::Error("unknown permission outcome".to_owned()),
},
Err(e) => {
tracing::error!(?e, "failed to request permission");
PromptOutcome::Error("failed to request permission".to_owned())
}
}
};
crates/codegen/kigi-workspace/src/session/mod.rs
+3 -122
View File
@@ -8,12 +8,8 @@ pub mod tool_config;
use crate::capability::CapabilityMode;
use crate::config::{MemoryConfig, SessionContextFactory};
use crate::file_system::{AsyncFsWrapper, LocalFs};
use crate::hub::{HubConfig, HubHandle};
use crate::session::file_state::FileStateTracker;
use kigi_computer_hub_mcp_adapter::McpBridgeHandle;
use kigi_hunk_tracker::HunkTrackerHandle;
use kigi_mcp::servers::McpState;
use kigi_tool_protocol::ToolId;
use kigi_tool_runtime::WorkspaceViewerContext;
use kigi_tools::notification::types::{ToolNotification, ToolNotificationHandle};
use kigi_tools::registry::types::{FinalizedToolset, ToolConfig, ToolServerConfig};
@@ -76,12 +72,6 @@ pub struct WorkspaceSession {
inner: RwLock<WorkspaceSessionInner>,
/// Per-session lock that serialises `update_tool_config` calls.
pub(crate) update_lock: tokio::sync::Mutex<()>,
/// Per-session MCP state (owned clients, etc.).
pub(crate) mcp_state: Arc<tokio::sync::Mutex<McpState>>,
/// MCP bridges kept alive for the session lifetime.
pub(crate) mcp_bridges: tokio::sync::Mutex<Vec<McpBridgeHandle>>,
/// Qualified tool IDs registered on the server for this session's MCP tools.
pub(crate) mcp_tool_ids: tokio::sync::Mutex<Vec<ToolId>>,
/// Per-user feature-flag bag resolved at session-bind time, frozen for
/// the session lifetime. `None` → tools use their safe defaults.
pub(crate) viewer_ctx: Option<WorkspaceViewerContext>,
@@ -107,8 +97,7 @@ pub struct WorkspaceSession {
/// created (or last rebound) with. `None` when the session was resolved
/// from the workspace default (no explicit toolset in the bind metadata).
/// Lets a rebind detect a config change and re-resolve instead of silently
/// reusing a stale toolset (e.g. a session created by a metadata-less
/// hub revive bind that a config-carrying client rebind must correct).
/// reusing a stale toolset.
bind_tool_config_fingerprint: std::sync::Mutex<Option<serde_json::Value>>,
/// The last snapshot-driven rebuild failed and kept a stale toolset;
/// cleared by any successful install. While set, an identical-config
@@ -190,9 +179,6 @@ impl WorkspaceSession {
update_lock: tokio::sync::Mutex::new(()),
bind_tool_config_fingerprint: std::sync::Mutex::new(None),
stale_resolve: std::sync::atomic::AtomicBool::new(false),
mcp_state: Arc::new(tokio::sync::Mutex::new(McpState::new(vec![]))),
mcp_bridges: tokio::sync::Mutex::new(Vec::new()),
mcp_tool_ids: tokio::sync::Mutex::new(Vec::new()),
viewer_ctx,
yolo_mode: std::sync::atomic::AtomicBool::new(false),
system_notifications,
@@ -440,9 +426,6 @@ pub type ClientExtSink = std::sync::Arc<dyn Fn(String, serde_json::Value) + Send
/// Workspace-wide shared state.
pub struct WorkspaceShared {
pub(crate) default_tool_config: ToolServerConfig,
/// Require an explicit toolset on every `session.bind`; see
/// [`crate::config::WorkspaceConfig::require_explicit_toolset`].
pub(crate) require_explicit_toolset: bool,
/// See [`crate::config::WorkspaceConfig::confine_fs_to_workspace_root`].
/// Default `false`; enabled only for remote-sandbox workspace servers.
pub(crate) confine_fs_to_workspace_root: bool,
@@ -464,38 +447,16 @@ pub struct WorkspaceShared {
/// disabled/enabled lists). Used by `discover_plugins` via the
/// `discovery` module.
pub(crate) plugin_discovery_config: crate::discovery::PluginDiscoveryConfig,
/// Live server connection handle. `None` until
/// [`WorkspaceHandle::connect_hub`](crate::handle::WorkspaceHandle::connect_hub)
/// is called (or if no [`HubConfig`] was provided).
///
/// Uses `tokio::sync::Mutex` so the guard can be held across the
/// async `HubHandle::connect()` call, preventing TOCTOU races.
pub(crate) hub_handle: tokio::sync::Mutex<Option<HubHandle>>,
/// Remote-origin tool configs (consumer direction), updated by the
/// notification listener.
pub(crate) hub_tools_snapshot: arc_swap::ArcSwap<Vec<ToolConfig>>,
/// Server config stashed at construction time for deferred connect.
pub(crate) hub_config: Option<HubConfig>,
/// Auth provider for xAI service calls.
pub(crate) auth_provider: Option<kigi_computer_hub_sdk::SharedAuthProvider>,
/// Connection-level sink feeding the `ActivityTracker` (drained by
/// `run_activity_feed`); not a network egress. `None` until `connect_hub()` sets it.
pub(crate) activity_notify_handle:
arc_swap::ArcSwap<Option<kigi_tools::notification::types::ToolNotificationHandle>>,
/// Sink for workspace-originated ext-notifications to the client (e.g.
/// `x.ai/search/fuzzy/status`). Mode-agnostic: the shell wires it to the
/// agent gateway in local mode, and to the server in proxy mode. `None` until
/// set via [`WorkspaceHandle::set_client_ext_sink`](crate::handle::WorkspaceHandle::set_client_ext_sink).
pub(crate) client_ext_sink: arc_swap::ArcSwap<Option<ClientExtSink>>,
pub(crate) local_registry: kigi_computer_hub_sdk::LocalRegistry,
pub(crate) local_registry: kigi_tool_runtime::LocalRegistry,
pub(crate) activity_tracker: std::sync::Arc<crate::activity::ActivityTracker>,
/// Runtime-tunable timing/threshold config for the tool server.
/// Read by the status publisher task and at shutdown.
pub(crate) status_config: crate::status_config::StatusConfig,
/// Opaque metadata for the tool server registration, forwarded verbatim to
/// the server; structured access goes through
/// [`WorkspaceShared::server_metadata_typed`].
pub(crate) server_metadata: Option<serde_json::Value>,
/// Workspace-level fuzzy search manager. Separate from the shell's
/// own `FuzzySearchManager` — this instance serves remote (hub/RPC)
/// clients.
@@ -533,7 +494,6 @@ pub struct WorkspaceShared {
/// turn start inside the check→install window deterministically.
#[cfg(test)]
pub(crate) post_resolve_test_hook: parking_lot::Mutex<Option<Box<dyn Fn() + Send + Sync>>>,
pub(crate) client_fs_hash_memo: crate::file_system::client_fs::FileHashMemo,
}
impl WorkspaceShared {
/// Workspace root directory.
@@ -578,40 +538,6 @@ impl WorkspaceShared {
.get(session_id)
.map(|w| w.value().clone())
}
/// Stable hub server id (`--server-id`), if a hub config is present.
pub(crate) fn server_id(&self) -> Option<String> {
self.hub_config.as_ref().and_then(|c| c.server_id.clone())
}
/// Auth provider used for xAI service calls.
pub fn auth_provider(&self) -> Option<&kigi_computer_hub_sdk::SharedAuthProvider> {
self.auth_provider.as_ref()
}
/// Parse the opaque [`server_metadata`](Self::server_metadata) blob into
/// the typed subset the workspace needs (currently `sandbox_id`);
/// unknown/missing fields default cleanly. A present-but-malformed blob is
/// logged and salvaged field-by-field (a bad sibling field must not
/// silently drop `sandbox_id` from every environment artifact).
pub(crate) fn server_metadata_typed(&self) -> crate::config::WorkspaceServerMetadata {
let Some(v) = self.server_metadata.as_ref() else {
return Default::default();
};
match serde_json::from_value(v.clone()) {
Ok(typed) => typed,
Err(e) => {
tracing::warn!(
error = % e,
"workspace: malformed server_metadata; salvaging sandbox_id field-wise"
);
crate::config::WorkspaceServerMetadata {
sandbox_id: v
.get("sandbox_id")
.and_then(serde_json::Value::as_str)
.map(str::to_owned),
..Default::default()
}
}
}
}
pub fn default_tool_config(&self) -> &ToolServerConfig {
&self.default_tool_config
}
@@ -624,49 +550,6 @@ impl WorkspaceShared {
pub fn mcp_tools_snapshot(&self) -> Arc<Vec<ToolConfig>> {
self.mcp_tools_snapshot.load_full()
}
/// The tool server, if a server connection is active.
///
/// Returns a clone of the [`ToolServer`](kigi_computer_hub_sdk::ToolServer)
/// which is cheap (`Arc` bump). Uses `try_lock` to avoid blocking
/// on the async mutex from synchronous contexts. Returns `None` if
/// the lock is held (i.e. a `connect_hub` call is in progress).
pub fn hub_server(&self) -> Option<kigi_computer_hub_sdk::ToolServer> {
self.hub_handle
.try_lock()
.ok()
.and_then(|guard| guard.as_ref().map(|h| h.server.clone()))
}
/// Like [`Self::hub_server`] but awaits the `hub_handle` lock instead of
/// returning `None` on contention. Use from async contexts that must not
/// confuse a transient `connect_hub` lock-hold with "no hub connected";
/// `None` means no hub is connected.
pub async fn hub_server_blocking(&self) -> Option<kigi_computer_hub_sdk::ToolServer> {
self.hub_handle
.lock()
.await
.as_ref()
.map(|h| h.server.clone())
}
/// Current snapshot of hub-provided tool configs (consumer direction).
pub fn hub_tools_snapshot(&self) -> Arc<Vec<ToolConfig>> {
self.hub_tools_snapshot.load_full()
}
/// Compose a session's tool `ctx.notification_handle` as a fan-out of the
/// connection-level activity feed (internal tracker accounting) and the
/// opt-in per-session `system.notify` sender. Only the `system.notify` leg
/// reaches a client, so the fan-out can't double-wake. `None` → factory default.
pub(crate) fn compose_session_notification_handle(
&self,
system_notify_handle: Option<ToolNotificationHandle>,
) -> Option<ToolNotificationHandle> {
let activity = self.activity_notify_handle.load_full().as_ref().clone();
match (activity, system_notify_handle) {
(None, None) => None,
(Some(a), None) => Some(a),
(None, Some(s)) => Some(s),
(Some(a), Some(s)) => Some(ToolNotificationHandle::tee(vec![a, s])),
}
}
pub fn activity_tracker(&self) -> &std::sync::Arc<crate::activity::ActivityTracker> {
&self.activity_tracker
}
@@ -717,7 +600,6 @@ impl WorkspaceShared {
};
let trigger = SwapTrigger::from_rebuild_source(source);
let mcp_snap = self.mcp_tools_snapshot.load_full();
let hub_snap = self.hub_tools_snapshot.load_full();
let sessions: Vec<(String, Arc<WorkspaceSession>)> = {
let guard = self.sessions.read();
guard
@@ -779,7 +661,6 @@ impl WorkspaceShared {
baseline,
session.capability_mode(),
&mcp_snap,
&hub_snap,
session.cwd().to_path_buf(),
session.session_env().clone(),
&sid,
@@ -787,7 +668,7 @@ impl WorkspaceShared {
Some(self.local_registry.clone()),
self.lsp.clone(),
session.viewer_ctx().cloned(),
self.compose_session_notification_handle(session.system_notify_handle()),
session.system_notify_handle(),
session.terminal_backend().clone(),
) {
Ok((effective, toolset)) => {
crates/codegen/kigi-workspace/src/session/tool_config.rs
+43 -389
View File
@@ -1,11 +1,10 @@
//! Tool config resolution pipeline.
//!
//! Five-step resolution:
//! Four-step resolution:
//! 1. `effective_tool_config = config.tool_config.unwrap_or_else(|| parent.effective_tool_config.clone())`
//! 2. `merged = merge_mcp_tools(effective_tool_config, shared.mcp_servers.snapshot())`
//! 3. `merged = merge_hub_tools(merged, shared.hub_tools_snapshot())`
//! 4. `filtered = config.capability_mode.filter(merged)`
//! 5. `toolset = build_finalized_toolset(filtered, &session.cwd, &session.session_env, ...)`
//! 3. `filtered = config.capability_mode.filter(merged)`
//! 4. `toolset = build_finalized_toolset(filtered, &session.cwd, &session.session_env, ...)`
use crate::capability::{CapabilityMode, kind_allowed};
use crate::config::SessionContextFactory;
use crate::error::{WorkspaceError, WorkspaceResult};
@@ -19,19 +18,16 @@ use std::sync::Arc;
/// Create-shaped entry of the resolution pipeline: run
/// [`resolve_session_toolset_rebuild`] around a FRESH factory-built
/// session-lifetime terminal backend, and return that backend so the caller
/// can store it on the session it is creating. Session-less resolves (the
/// `__template__` catalog resolve in `connect_hub`) also use this entry and
/// simply drop the returned backend with the toolset.
/// can store it on the session it is creating.
pub(crate) fn resolve_session_toolset(
effective_tool_config: ToolServerConfig,
capability_mode: CapabilityMode,
mcp_snapshot: &[ToolConfig],
hub_snapshot: &[ToolConfig],
cwd: PathBuf,
session_env: Arc<HashMap<String, String>>,
session_id: &str,
factory: &dyn SessionContextFactory,
local_registry: Option<kigi_computer_hub_sdk::LocalRegistry>,
local_registry: Option<kigi_tool_runtime::LocalRegistry>,
lsp: Option<std::sync::Arc<dyn kigi_tools::implementations::lsp::LspBackend>>,
viewer_ctx: Option<kigi_tool_runtime::WorkspaceViewerContext>,
notification_handle: Option<kigi_tools::notification::types::ToolNotificationHandle>,
@@ -45,7 +41,6 @@ pub(crate) fn resolve_session_toolset(
effective_tool_config,
capability_mode,
mcp_snapshot,
hub_snapshot,
cwd,
session_env,
session_id,
@@ -66,24 +61,23 @@ pub(crate) fn resolve_session_toolset(
///
/// Returns the *unmodified* `effective_tool_config` (step-1 baseline) so
/// the caller can store it on the session. The FinalizedToolset reflects
/// MCP + hub merging and capability filtering on top of that baseline.
/// MCP merging and capability filtering on top of that baseline.
///
/// **MCP-origin and hub-origin `kind: None` tools are dropped under
/// every non-`All` mode.** Baseline `kind: None` tools are always kept —
/// but before filtering, kind-less baseline entries whose id the binary's
/// registry knows get their [`ToolKind`] backfilled (see
/// [`backfill_tool_kinds`]), so the capability filter applies to pinned
/// server-bind toolsets whose wire entries cannot carry a kind.
/// **MCP-origin `kind: None` tools are dropped under every non-`All`
/// mode.** Baseline `kind: None` tools are always kept — but before
/// filtering, kind-less baseline entries whose id the binary's registry
/// knows get their [`ToolKind`] backfilled (see [`backfill_tool_kinds`]),
/// so the capability filter applies to pinned toolsets whose wire entries
/// cannot carry a kind.
pub(crate) fn resolve_session_toolset_rebuild(
effective_tool_config: ToolServerConfig,
capability_mode: CapabilityMode,
mcp_snapshot: &[ToolConfig],
hub_snapshot: &[ToolConfig],
cwd: PathBuf,
session_env: Arc<HashMap<String, String>>,
session_id: &str,
factory: &dyn SessionContextFactory,
local_registry: Option<kigi_computer_hub_sdk::LocalRegistry>,
local_registry: Option<kigi_tool_runtime::LocalRegistry>,
lsp: Option<std::sync::Arc<dyn kigi_tools::implementations::lsp::LspBackend>>,
viewer_ctx: Option<kigi_tool_runtime::WorkspaceViewerContext>,
notification_handle: Option<kigi_tools::notification::types::ToolNotificationHandle>,
@@ -94,24 +88,7 @@ pub(crate) fn resolve_session_toolset_rebuild(
builder = builder.with_local_registry(lr);
}
let baseline = backfill_tool_kinds(&effective_tool_config, &builder.known_tool_kinds());
let filtered = merge_and_filter(
&baseline,
mcp_snapshot,
hub_snapshot,
capability_mode,
session_id,
);
let hub_ids: std::collections::HashSet<&str> =
hub_snapshot.iter().map(|t| t.id.as_str()).collect();
let finalize_config = ToolServerConfig {
tools: filtered
.tools
.iter()
.filter(|t| !hub_ids.contains(t.id.as_str()))
.cloned()
.collect(),
behavior_preset: filtered.behavior_preset.clone(),
};
let finalize_config = merge_and_filter(&baseline, mcp_snapshot, capability_mode, session_id);
let mut ctx = factory.build_session_context(session_id, cwd, session_env, terminal_backend);
if let Some(lsp_handle) = lsp {
ctx.lsp = Some(lsp_handle);
@@ -157,22 +134,20 @@ fn backfill_tool_kinds(
behavior_preset: config.behavior_preset.clone(),
}
}
/// Steps 2-4 of the resolution pipeline, without step 5 (`finalize`):
/// Steps 2-3 of the resolution pipeline, without the `finalize` step:
///
/// - **Step 2** -- MCP merge: append MCP-origin tools, skipping ID/name collisions with baseline.
/// - **Step 3** -- Hub merge: append hub-origin tools, skipping ID/name collisions with baseline or MCP.
/// - **Step 4** -- Capability filter: drop tools whose `kind` is not allowed by the mode.
/// External (MCP/hub) `kind: None` tools are only kept under `CapabilityMode::All`.
/// - **Step 3** -- Capability filter: drop tools whose `kind` is not allowed by the mode.
/// External (MCP) `kind: None` tools are only kept under `CapabilityMode::All`.
///
/// Priority on ID/name collision: baseline wins > MCP wins > hub is skipped.
/// Priority on ID/name collision: baseline wins over MCP.
pub(crate) fn merge_and_filter(
baseline: &ToolServerConfig,
mcp_snapshot: &[ToolConfig],
hub_snapshot: &[ToolConfig],
mode: CapabilityMode,
session_id: &str,
) -> ToolServerConfig {
if mcp_snapshot.is_empty() && hub_snapshot.is_empty() {
if mcp_snapshot.is_empty() {
return mode.filter(baseline);
}
let baseline_ids: std::collections::HashSet<&str> =
@@ -187,7 +162,6 @@ pub(crate) fn merge_and_filter(
.collect();
let mut tagged: Vec<(ToolConfig, bool)> =
baseline.tools.iter().cloned().map(|t| (t, false)).collect();
let mut mcp_tool_ids: std::collections::HashSet<&str> = std::collections::HashSet::new();
for mcp_tool in mcp_snapshot {
if baseline_ids.contains(mcp_tool.id.as_str()) {
tracing::warn!(
@@ -205,35 +179,8 @@ pub(crate) fn merge_and_filter(
);
continue;
}
mcp_tool_ids.insert(mcp_tool.id.as_str());
tagged.push((mcp_tool.clone(), true));
}
for hub_tool in hub_snapshot {
if baseline_ids.contains(hub_tool.id.as_str()) {
tracing::debug!(
hub_id = % hub_tool.id, session = % session_id,
"skipping remote tool: id collides with baseline"
);
continue;
}
if mcp_tool_ids.contains(hub_tool.id.as_str()) {
tracing::debug!(
hub_id = % hub_tool.id, session = % session_id,
"skipping remote tool: id collides with MCP tool"
);
continue;
}
let client_name = hub_tool.resolve_client_name(&hub_tool.id);
if !taken_names.insert(client_name.clone()) {
tracing::debug!(
hub_id = % hub_tool.id, client_name = % client_name, session = %
session_id,
"skipping remote tool: resolved client name collides with another tool"
);
continue;
}
tagged.push((hub_tool.clone(), true));
}
let kept: Vec<ToolConfig> = tagged
.into_iter()
.filter(|(tool, is_external)| match tool.kind {
@@ -250,19 +197,13 @@ pub(crate) fn merge_and_filter(
}
/// Alias for backward compatibility.
pub type NoopSessionContextFactory = WorkspaceSessionContextFactory;
/// Whether per-session `tool_state.json` persistence + per-turn upload is
/// enabled (`KIGI_WORKSPACE_TOOL_STATE_ENABLED=true`; any other value keeps
/// legacy behavior).
pub fn tool_state_enabled() -> bool {
std::env::var("KIGI_WORKSPACE_TOOL_STATE_ENABLED").as_deref() == Ok("true")
}
/// Sanitize a `session_id` into a single safe filesystem path segment: chars
/// outside `[A-Za-z0-9_-]` become `_`, empty becomes `anon`. When any
/// replacement happened, an 8-hex digest of the ORIGINAL id is appended so the
/// mapping stays injective — plain substitution would collide distinct ids
/// (`sess/1` and `sess_1`) into one directory, cross-contaminating
/// persistence, rehydration, and [`crate::recovery::cleanup_stale_sessions`].
/// Already-safe ids (the common UUID case) map to themselves.
/// persistence and rehydration. Already-safe ids (the common UUID case)
/// map to themselves.
fn sanitize_session_id(session_id: &str) -> String {
let mut safe = String::with_capacity(session_id.len());
let mut modified = false;
@@ -297,16 +238,7 @@ fn ensure_session_dir(root: &std::path::Path, session_id: &str) -> (PathBuf, std
/// hazard is the global `environ` array, not the variable's value).
#[cfg(test)]
pub(crate) use crate::ENV_TEST_LOCK as TOOL_STATE_ENV_LOCK;
/// [`SessionContextFactory`] for workspace server sessions.
///
/// When constructed with an [`AuthProvider`] and API base URL, gen tools
/// (image_gen, video_gen) are enabled using the provider's current
/// OAuth token. Without auth, gen tools default to `Disabled`.
///
/// When [`with_tool_state_home`](Self::with_tool_state_home) is set, each
/// session's [`SessionContext::state_path`] is rooted at
/// `<home>/sessions/<session_id>/`; left unset, `state_path` stays empty
/// (legacy behavior).
/// [`SessionContextFactory`] for workspace sessions.
///
/// [`SessionContext::session_folder`] is `/tmp/sessions/<sanitized_id>/`
/// (terminal logs and other tool artifacts — not the project `cwd`).
@@ -318,64 +250,11 @@ pub(crate) use crate::ENV_TEST_LOCK as TOOL_STATE_ENV_LOCK;
/// [`build_terminal_backend`]: crate::config::SessionContextFactory::build_terminal_backend
/// [`build_session_context`]: crate::config::SessionContextFactory::build_session_context
/// [`LocalTerminalBackend`]: kigi_tools::computer::local::LocalTerminalBackend
pub struct WorkspaceSessionContextFactory {
auth: Option<kigi_computer_hub_sdk::SharedAuthProvider>,
api_base_url: Option<String>,
/// Resolved `$KIGI_WORKSPACE_HOME` when tool-state persistence is enabled;
/// `None` disables it. Resolved once by the caller so the factory performs
/// no per-build env reads.
tool_state_home: Option<PathBuf>,
}
impl Default for WorkspaceSessionContextFactory {
fn default() -> Self {
Self::new()
}
}
#[derive(Default)]
pub struct WorkspaceSessionContextFactory;
impl WorkspaceSessionContextFactory {
pub fn new() -> Self {
Self {
auth: None,
api_base_url: None,
tool_state_home: None,
}
}
/// Factory with auth — gen tools use the provider's live token.
pub fn with_auth(
auth: kigi_computer_hub_sdk::SharedAuthProvider,
api_base_url: String,
) -> Self {
Self {
auth: Some(auth),
api_base_url: Some(api_base_url),
tool_state_home: None,
}
}
/// Enable session-keyed tool-state persistence rooted at `home`
/// (`$KIGI_WORKSPACE_HOME`). Callers should only invoke this when
/// [`tool_state_enabled`] is `true`.
pub fn with_tool_state_home(mut self, home: PathBuf) -> Self {
self.tool_state_home = Some(home);
self
}
/// `<tool_state_home>/sessions/<sanitized_id>/tool_state.json`, or empty
/// when persistence is disabled / dir creation fails.
fn resolve_state_path(&self, session_id: &str) -> PathBuf {
let Some(home) = self.tool_state_home.as_ref() else {
return PathBuf::new();
};
let (dir, created) = ensure_session_dir(home, session_id);
if let Err(e) = created {
tracing::warn!(
session = % session_id, dir = % dir.display(), error = % e,
"tool_state: failed to create session dir; persistence disabled for session"
);
return PathBuf::new();
}
tracing::debug!(
session = % session_id, dir = % dir.display(),
"tool_state: persistence bound to session-keyed dir"
);
dir.join("tool_state.json")
Self
}
/// `/tmp/sessions/<sanitized_id>/` for terminal logs and other tool artifacts.
fn resolve_session_folder(session_id: &str) -> PathBuf {
@@ -397,59 +276,9 @@ impl SessionContextFactory for WorkspaceSessionContextFactory {
session_env: Arc<HashMap<String, String>>,
backend: Arc<dyn kigi_tools::computer::types::TerminalBackend>,
) -> kigi_tools::registry::types::SessionContext {
use kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig;
use kigi_tools::implementations::grok_build::image_gen::ImageGenConfig;
use kigi_tools::implementations::grok_build::video_gen::VideoGenConfig;
use kigi_tools::implementations::web_search::WebSearchConfig;
let fs = Arc::new(kigi_tools::computer::local::LocalFs)
as Arc<dyn kigi_tools::computer::types::AsyncFileSystem>;
let notification_handle = kigi_tools::notification::ToolNotificationHandle::noop();
let (image_gen_config, video_gen_config, web_search_config, app_builder_deployer_config) =
if let (Some(auth), Some(url)) = (&self.auth, &self.api_base_url) {
let cred = auth.current();
match cred {
kigi_computer_hub_sdk::AuthCredential::Bearer { token, .. } => {
let headers = build_proxy_headers(url);
(
ImageGenConfig::Enabled {
api_key: token.clone(),
base_url: url.clone(),
extra_headers: headers.clone(),
image_gen_enabled: true,
image_edit_enabled: true,
model_override: None,
tier_restricted: false,
},
VideoGenConfig::Enabled {
api_key: token.clone(),
base_url: url.clone(),
extra_headers: headers.clone(),
zdr_video_output_s3: None,
tier_restricted: false,
},
WebSearchConfig::Enabled {
search_url: format!("{}/search", url.trim_end_matches('/')),
api_key: token,
extra_headers: headers,
},
AppBuilderDeployerConfig::default(),
)
}
_ => (
ImageGenConfig::default(),
VideoGenConfig::default(),
WebSearchConfig::default(),
AppBuilderDeployerConfig::default(),
),
}
} else {
(
ImageGenConfig::default(),
VideoGenConfig::default(),
WebSearchConfig::default(),
AppBuilderDeployerConfig::default(),
)
};
kigi_tools::registry::types::SessionContext {
backend,
fs,
@@ -460,16 +289,18 @@ impl SessionContextFactory for WorkspaceSessionContextFactory {
owner_session_id: None,
parent_scheduler_handle: None,
skills: vec![],
state_path: self.resolve_state_path(session_id),
state_path: PathBuf::new(),
memory_backend: None,
web_search_config,
web_search_config: kigi_tools::implementations::web_search::WebSearchConfig::default(),
web_fetch_config: build_web_fetch_config(),
lsp: None,
image_gen_config,
video_gen_config,
app_builder_deployer_config,
image_gen_config:
kigi_tools::implementations::grok_build::image_gen::ImageGenConfig::default(),
video_gen_config:
kigi_tools::implementations::grok_build::video_gen::VideoGenConfig::default(),
app_builder_deployer_config:
kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig::default(),
api_key_provider: None,
auth_provider: self.auth.clone(),
attribution_callback: None,
system_reminder_tag: kigi_tools::reminders::DEFAULT_REMINDER_TAG,
}
@@ -488,18 +319,6 @@ impl SessionContextFactory for WorkspaceSessionContextFactory {
IDS.clone()
}
}
/// Build extra headers for API calls routed through the chat proxy.
/// Mirrors the shell's `inject_proxy_headers` logic.
fn build_proxy_headers(base_url: &str) -> indexmap::IndexMap<String, String> {
let mut headers = indexmap::IndexMap::new();
let version = kigi_version::VERSION;
headers.insert(
"user-agent".to_string(),
format!("kigi-workspace/{version}"),
);
headers.insert("x-grok-client-version".to_string(), version.to_string());
headers
}
/// Build web fetch config. Enabled with default params unless
/// `KIGI_DISABLE_WEB_FETCH=1` is set.
fn build_web_fetch_config() -> kigi_tools::implementations::grok_build::web_fetch::WebFetchConfig {
@@ -574,7 +393,6 @@ pub mod test_support {
video_gen_config: Default::default(),
app_builder_deployer_config: Default::default(),
api_key_provider: None,
auth_provider: None,
attribution_callback: None,
system_reminder_tag: kigi_tools::reminders::DEFAULT_REMINDER_TAG,
}
@@ -635,7 +453,6 @@ mod tests {
baseline,
CapabilityMode::ReadWrite,
&[],
&[],
cwd,
empty_env(),
"main",
@@ -672,7 +489,6 @@ mod tests {
baseline,
CapabilityMode::ReadWrite,
&snapshot,
&[],
PathBuf::from("/tmp"),
empty_env(),
"main",
@@ -752,7 +568,6 @@ mod tests {
baseline,
CapabilityMode::ReadOnly,
&[],
&[],
PathBuf::from("/tmp"),
empty_env(),
"main",
@@ -792,13 +607,7 @@ mod tests {
behavior_preset: None,
};
let mcp_edit = test_support::tc("mcp.editor", Some(ToolKind::Edit));
let filtered = merge_and_filter(
&baseline,
&[mcp_edit],
&[],
CapabilityMode::ReadOnly,
"test",
);
let filtered = merge_and_filter(&baseline, &[mcp_edit], CapabilityMode::ReadOnly, "test");
assert!(!filtered.tools.iter().any(|t| t.id == "mcp.editor"));
}
#[tokio::test]
@@ -812,13 +621,7 @@ mod tests {
behavior_preset: None,
};
let mcp = vec![test_support::tc("mcp.opaque", None)];
let filtered = merge_and_filter(
&baseline,
&mcp,
&[],
CapabilityMode::ReadOnly,
"test_session",
);
let filtered = merge_and_filter(&baseline, &mcp, CapabilityMode::ReadOnly, "test_session");
let kept_ids: Vec<&str> = filtered.tools.iter().map(|t| t.id.as_str()).collect();
assert!(
kept_ids.contains(&"baseline.opaque"),
@@ -841,7 +644,7 @@ mod tests {
behavior_preset: None,
};
let mcp = vec![test_support::tc("mcp.opaque", None)];
let filtered = merge_and_filter(&baseline, &mcp, &[], CapabilityMode::All, "test_session");
let filtered = merge_and_filter(&baseline, &mcp, CapabilityMode::All, "test_session");
let kept_ids: Vec<&str> = filtered.tools.iter().map(|t| t.id.as_str()).collect();
assert!(
kept_ids.contains(&"mcp.opaque"),
@@ -859,13 +662,7 @@ mod tests {
let mut mcp_b = test_support::tc("mcp.tool_b", Some(ToolKind::Read));
mcp_b.name_override = Some("shared_name".into());
let mcp = vec![mcp_a, mcp_b];
let filtered = merge_and_filter(
&baseline,
&mcp,
&[],
CapabilityMode::ReadOnly,
"test_session",
);
let filtered = merge_and_filter(&baseline, &mcp, CapabilityMode::ReadOnly, "test_session");
let ids: Vec<&str> = filtered.tools.iter().map(|t| t.id.as_str()).collect();
assert!(ids.contains(&"mcp.tool_a"), "first wins: {ids:?}");
assert!(
@@ -874,144 +671,6 @@ mod tests {
);
}
#[test]
fn hub_tool_merged_into_empty_baseline() {
let baseline = ToolServerConfig {
tools: vec![],
behavior_preset: None,
};
let hub = vec![test_support::tc("hub:remote_exec", None)];
let filtered = merge_and_filter(&baseline, &[], &hub, CapabilityMode::All, "test");
let ids: Vec<&str> = filtered.tools.iter().map(|t| t.id.as_str()).collect();
assert!(
ids.contains(&"hub:remote_exec"),
"remote tool should appear under All mode: {ids:?}"
);
}
#[test]
fn hub_tool_dropped_under_readonly_because_kind_none() {
let baseline = ToolServerConfig {
tools: vec![test_support::tc(
"GrokBuild:read_file",
Some(ToolKind::Read),
)],
behavior_preset: None,
};
let hub = vec![test_support::tc("hub:remote_exec", None)];
let filtered = merge_and_filter(&baseline, &[], &hub, CapabilityMode::ReadOnly, "test");
let ids: Vec<&str> = filtered.tools.iter().map(|t| t.id.as_str()).collect();
assert!(
!ids.contains(&"hub:remote_exec"),
"hub kind: None MUST be dropped under ReadOnly: {ids:?}"
);
}
#[test]
fn hub_tool_dedup_baseline_wins() {
let baseline = ToolServerConfig {
tools: vec![test_support::tc("hub:read_file", Some(ToolKind::Read))],
behavior_preset: None,
};
let hub = vec![test_support::tc("hub:read_file", None)];
let filtered = merge_and_filter(&baseline, &[], &hub, CapabilityMode::All, "test");
let count = filtered
.tools
.iter()
.filter(|t| t.id == "hub:read_file")
.count();
assert_eq!(count, 1, "duplicate should be deduped");
}
#[test]
fn hub_tool_dedup_mcp_wins_over_hub() {
let baseline = ToolServerConfig {
tools: vec![],
behavior_preset: None,
};
let mcp = vec![test_support::tc("hub:shared_tool", Some(ToolKind::Read))];
let hub = vec![test_support::tc("hub:shared_tool", None)];
let filtered = merge_and_filter(&baseline, &mcp, &hub, CapabilityMode::All, "test");
let count = filtered
.tools
.iter()
.filter(|t| t.id == "hub:shared_tool")
.count();
assert_eq!(count, 1, "MCP wins; hub duplicate skipped");
let tool = filtered
.tools
.iter()
.find(|t| t.id == "hub:shared_tool")
.unwrap();
assert_eq!(tool.kind, Some(ToolKind::Read));
}
#[test]
fn hub_tool_name_collision_with_baseline_skipped() {
let baseline = ToolServerConfig {
tools: vec![test_support::tc(
"GrokBuild:read_file",
Some(ToolKind::Read),
)],
behavior_preset: None,
};
let mut hub_tool = test_support::tc("hub:read_file_v2", None);
hub_tool.name_override = Some("read_file".into());
let hub = vec![hub_tool];
let filtered = merge_and_filter(&baseline, &[], &hub, CapabilityMode::All, "test");
let ids: Vec<&str> = filtered.tools.iter().map(|t| t.id.as_str()).collect();
assert!(
!ids.contains(&"hub:read_file_v2"),
"remote tool with colliding client name must be skipped: {ids:?}"
);
}
#[test]
fn empty_hub_snapshot_is_noop() {
let baseline = test_support::baseline_config();
let baseline_ids: Vec<String> = baseline.tools.iter().map(|t| t.id.clone()).collect();
let filtered = merge_and_filter(&baseline, &[], &[], CapabilityMode::ReadWrite, "test");
let filtered_ids: Vec<String> = filtered.tools.iter().map(|t| t.id.clone()).collect();
assert_eq!(filtered_ids, baseline_ids);
}
/// Only the literal `"true"` enables tool-state persistence.
#[test]
fn tool_state_enabled_only_true_enables() {
let _guard = super::TOOL_STATE_ENV_LOCK
.lock()
.unwrap_or_else(|e| e.into_inner());
let var = "KIGI_WORKSPACE_TOOL_STATE_ENABLED";
unsafe { std::env::remove_var(var) };
assert!(!tool_state_enabled(), "unset → disabled");
unsafe { std::env::set_var(var, "false") };
assert!(!tool_state_enabled(), "false → disabled");
unsafe { std::env::set_var(var, "1") };
assert!(!tool_state_enabled(), "1 → disabled (only \"true\")");
unsafe { std::env::set_var(var, "true") };
assert!(tool_state_enabled(), "true → enabled");
unsafe { std::env::remove_var(var) };
}
/// With a tool-state home set, state is rooted at
/// `<home>/sessions/<session_id>/tool_state.json` and the dir is created.
#[test]
fn factory_resolves_session_keyed_state_path_when_home_set() {
let home = tempfile::TempDir::new().unwrap();
let factory =
WorkspaceSessionContextFactory::new().with_tool_state_home(home.path().to_path_buf());
let p = factory.resolve_state_path("sess-1");
assert_eq!(
p,
home.path()
.join("sessions")
.join("sess-1")
.join("tool_state.json")
);
assert!(
home.path().join("sessions").join("sess-1").is_dir(),
"the session dir must be created so the persistence writer can rename into it"
);
}
/// Without a tool-state home, `state_path` stays empty (legacy behavior).
#[test]
fn factory_state_path_empty_when_home_unset() {
let factory = WorkspaceSessionContextFactory::new();
assert_eq!(factory.resolve_state_path("sess-1"), PathBuf::new());
}
#[test]
fn factory_session_folder_is_tmp_sessions_not_project_cwd() {
let cwd = PathBuf::from("/workspace");
let folder = WorkspaceSessionContextFactory::resolve_session_folder("sess-1");
@@ -1063,18 +722,16 @@ mod tests {
/// A hostile `session_id` (`../../etc`) is sanitized to a single safe
/// segment and cannot traverse outside `<home>/sessions/`.
#[test]
fn factory_sanitizes_malicious_session_id_no_traversal() {
fn ensure_session_dir_sanitizes_malicious_session_id_no_traversal() {
let home = tempfile::TempDir::new().unwrap();
let factory =
WorkspaceSessionContextFactory::new().with_tool_state_home(home.path().to_path_buf());
let sessions = home.path().join("sessions");
let p = factory.resolve_state_path("../../etc");
let (session_dir, created) = ensure_session_dir(home.path(), "../../etc");
assert!(created.is_ok());
assert!(
p.starts_with(&sessions),
"state path escaped sessions/: {}",
p.display()
session_dir.starts_with(&sessions),
"session dir escaped sessions/: {}",
session_dir.display()
);
let session_dir = p.parent().expect("state path has a parent dir");
assert_eq!(
session_dir.parent(),
Some(sessions.as_path()),
@@ -1123,7 +780,6 @@ mod tests {
test_support::baseline_config(),
CapabilityMode::ReadWrite,
&[],
&[],
cwd.clone(),
empty_env(),
"sess-A",
@@ -1145,7 +801,6 @@ mod tests {
test_support::baseline_config(),
CapabilityMode::ReadWrite,
&[],
&[],
cwd.clone(),
empty_env(),
"sess-A",
@@ -1170,7 +825,6 @@ mod tests {
test_support::baseline_config(),
CapabilityMode::ReadWrite,
&[],
&[],
cwd,
empty_env(),
"sess-B",
crates/codegen/kigi-workspace/src/workspace_ops.rs
+91 -323
View File
@@ -1,30 +1,22 @@
//! [`WorkspaceOps`] — dual-mode workspace operations handle.
//! [`WorkspaceOps`] — workspace operations handle.
//!
//! Two modes:
//!
//! - **`Local`** — extensions dispatch through [`WorkspaceHandle`]; tool
//! calls dispatch through the workspace session's [`FinalizedToolset`].
//! The toolset is installed via [`WorkspaceOps::bind_local_session`]
//! after the agent is built.
//!
//! - **`Proxy`** — everything routes through hub WebSocket to a remote
//! workspace server.
//! Extensions dispatch through [`WorkspaceHandle`]; tool calls dispatch
//! through the workspace session's [`FinalizedToolset`]. The toolset is
//! installed via [`WorkspaceOps::bind_local_session`] after the agent is
//! built.
//!
//! ## Type safety
//!
//! Each RPC method has a corresponding request struct that implements
//! Each operation has a corresponding request struct that implements
//! [`WorkspaceRpc`]. The struct carries a `METHOD` constant and derives
//! `Serialize + Deserialize`. Both the proxy client (`WorkspaceOps`) and
//! the server (`WorkspaceRpcHandler::dispatch`) use the same struct —
//! add/rename a field and the compiler catches both sides.
//! `Serialize + Deserialize`, so op identity and payload shape live in
//! one place.
use crate::error::{WorkspaceError, WorkspaceResult};
use crate::file_system::ContentSearchRequest;
use crate::handle::WorkspaceHandle;
use crate::worktree::{ApplyWorktreeRequest, CreateWorktreeRequest, RemoveWorktreeRequest};
use async_trait::async_trait;
use kigi_computer_hub_sdk::ToolHarness;
use kigi_tools::types::output::ToolRunResult;
use kigi_workspace_client::{WorkspaceClient, is_transport_fatal};
pub use kigi_workspace_types::rpc::WorkspaceRpc;
pub use kigi_workspace_types::rpc::agents_md::DiscoverAgentsMdReq;
pub use kigi_workspace_types::rpc::code_nav::{
@@ -32,9 +24,8 @@ pub use kigi_workspace_types::rpc::code_nav::{
CodeIndexStats, CodeIndexStatusReq, CodeIndexStatusResponse, CodeNavLocation, CodeNavResponse,
};
pub use kigi_workspace_types::rpc::fs::{
ClientFsListNode, ClientFsListReq, ClientFsListRes, ClientFsReadFileReq, ClientFsReadFileRes,
ClientFsStatReq, ClientFsStatRes, GetFileEntry, GetFileResult, GetFilesReq, GetFilesRes,
PutFileEntry, PutFileResult, PutFilesReq, PutFilesRes,
GetFileEntry, GetFileResult, GetFilesReq, GetFilesRes, PutFileEntry, PutFileResult,
PutFilesReq, PutFilesRes,
};
pub use kigi_workspace_types::rpc::git::{
BinaryFileInfoData, CheckoutCommitResponse, CommitWithPatchData, DetectVcsKindReq,
@@ -69,7 +60,6 @@ use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
/// Implements [`WorkspaceRpc`] for request types whose responses
/// reference crate-internal types and so cannot live in the types crate.
macro_rules! workspace_rpc {
@@ -214,7 +204,7 @@ fn session_tracker(
session_id: Option<&str>,
) -> WorkspaceResult<kigi_hunk_tracker::HunkTrackerHandle> {
let sid = session_id
.ok_or_else(|| WorkspaceError::HubError("per-session hunk op requires a session".into()))?;
.ok_or_else(|| WorkspaceError::Internal("per-session hunk op requires a session".into()))?;
let session = ws
.session(sid)
.ok_or_else(|| WorkspaceError::SessionNotFound(sid.to_owned()))?;
@@ -251,13 +241,13 @@ impl WorkspaceOp for GitStatusExtReq {
self.include_patches,
)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
Ok(GitStatusExtResponse::structured(data))
}
GitStatusFormat::Prompt => {
let result = crate::file_system::git_status(cwd)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
Ok(GitStatusExtResponse::prompt(result))
}
}
@@ -273,7 +263,7 @@ impl WorkspaceOp for GitFilesReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::read_files(&cwd, &self.paths, &self.version)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -294,7 +284,7 @@ impl WorkspaceOp for GitDiffReq {
self.merge_base,
)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -307,7 +297,7 @@ impl WorkspaceOp for GitStageReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::stage(&cwd, self.paths.clone())
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -320,7 +310,7 @@ impl WorkspaceOp for GitStageContentReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::stage_content(&cwd, &self.path, &self.content)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -333,7 +323,7 @@ impl WorkspaceOp for GitUnstageReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::unstage(&cwd, self.paths.clone())
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -346,7 +336,7 @@ impl WorkspaceOp for GitDiscardReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::discard(&cwd, self.paths.clone(), self.scope, self.include_untracked)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -366,7 +356,7 @@ impl WorkspaceOp for GitCommitReq {
self.sync,
)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -379,7 +369,7 @@ impl WorkspaceOp for GitCheckoutReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::checkout_branch(&cwd, &self.branch, self.create)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -392,7 +382,7 @@ impl WorkspaceOp for GitStashReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::stash(&cwd, self.include_untracked)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -405,7 +395,7 @@ impl WorkspaceOp for GitInfoReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::git_info(&cwd)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -418,7 +408,7 @@ impl WorkspaceOp for GitBranchesReq {
let cwd = git_op_cwd(ws, &self.git_root)?;
crate::session::git::list_branches(&cwd)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -429,7 +419,7 @@ impl WorkspaceOp for GitCollectChangesReq {
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
{
return Err(WorkspaceError::HubError(
return Err(WorkspaceError::Internal(
"git collect changes is unavailable in this build".to_string(),
));
}
@@ -550,7 +540,7 @@ impl WorkspaceOp for PrepareWorktreeFromWorktreeReq {
spawn_task: result.spawn_task,
response: Some(
serde_json::to_value(&resp)
.map_err(|e| WorkspaceError::HubError(e.to_string()))?,
.map_err(|e| WorkspaceError::Internal(e.to_string()))?,
),
error: None,
}),
@@ -572,7 +562,7 @@ impl WorkspaceOp for CreateWorktreeFromWorktreeSyncReq {
let req = crate::worktree::CreateWorktreeFromWorktreeRequest::from(self.inner.clone());
crate::worktree::create_worktree_from_worktree_sync(&req)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -583,8 +573,8 @@ impl WorkspaceOp for WorktreeDbRebuildReq {
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
let report = crate::worktree::worktree_db_rebuild()
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::to_value(report).map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::to_value(report).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -595,7 +585,7 @@ impl WorkspaceOp for WorktreeDbPathReq {
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
let path = crate::worktree::worktree_db_path()
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
Ok(WorktreeDbPathResponse {
path: Some(path.display().to_string()),
})
@@ -864,8 +854,8 @@ fn hook_registry_to_wire(
registry: &kigi_hooks::discovery::HookRegistry,
) -> WorkspaceResult<HookRegistryWire> {
let value =
serde_json::to_value(registry).map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::from_value(value).map_err(|e| WorkspaceError::HubError(e.to_string()))
serde_json::to_value(registry).map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::from_value(value).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
/// Inverse of [`hook_registry_to_wire`]. The compiled `matcher` is absent from
/// the wire (and from this result); callers recompile it via
@@ -873,8 +863,8 @@ fn hook_registry_to_wire(
fn wire_to_hook_registry(
wire: &HookRegistryWire,
) -> WorkspaceResult<kigi_hooks::discovery::HookRegistry> {
let value = serde_json::to_value(wire).map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::from_value(value).map_err(|e| WorkspaceError::HubError(e.to_string()))
let value = serde_json::to_value(wire).map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::from_value(value).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
#[async_trait]
impl WorkspaceOp for HookRegistryReq {
@@ -886,56 +876,6 @@ impl WorkspaceOp for HookRegistryReq {
hook_registry_to_wire(&ws.hook_registry())
}
}
#[async_trait]
impl WorkspaceOp for PutFilesReq {
async fn execute(
&self,
ws: &WorkspaceHandle,
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
ws.put_files(self.files.clone()).await
}
}
#[async_trait]
impl WorkspaceOp for GetFilesReq {
async fn execute(
&self,
ws: &WorkspaceHandle,
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
ws.get_files(self.files.clone()).await
}
}
#[async_trait]
impl WorkspaceOp for ClientFsListReq {
async fn execute(
&self,
ws: &WorkspaceHandle,
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
crate::file_system::client_fs::list(ws, self).await
}
}
#[async_trait]
impl WorkspaceOp for ClientFsStatReq {
async fn execute(
&self,
ws: &WorkspaceHandle,
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
crate::file_system::client_fs::stat(ws, self).await
}
}
#[async_trait]
impl WorkspaceOp for ClientFsReadFileReq {
async fn execute(
&self,
ws: &WorkspaceHandle,
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
crate::file_system::client_fs::read_file(ws, self).await
}
}
/// Resolve the index root for a code-nav op. Prefers the explicit per-session
/// `root` (the cwd the client sends per window), else the workspace root.
/// Without this, code nav in a non-primary window would query the launch
@@ -972,7 +912,7 @@ impl WorkspaceOp for CodeGotoDefinitionReq {
let result = handle
.goto_definition(std::path::PathBuf::from(&self.file), self.line, self.col)
.await
.map_err(|e| WorkspaceError::HubError(format!("index channel closed: {e}")))?;
.map_err(|e| WorkspaceError::Internal(format!("index channel closed: {e}")))?;
Ok(query_result_to_response(result))
}
}
@@ -992,7 +932,7 @@ impl WorkspaceOp for CodeGotoReferencesReq {
self.include_definition,
)
.await
.map_err(|e| WorkspaceError::HubError(format!("index channel closed: {e}")))?;
.map_err(|e| WorkspaceError::Internal(format!("index channel closed: {e}")))?;
Ok(query_result_to_response(result))
}
}
@@ -1010,7 +950,7 @@ impl WorkspaceOp for CodeFindDefinitionsReq {
self.context_file.as_ref().map(std::path::PathBuf::from),
)
.await
.map_err(|e| WorkspaceError::HubError(format!("index channel closed: {e}")))?;
.map_err(|e| WorkspaceError::Internal(format!("index channel closed: {e}")))?;
Ok(symbol_locations_to_response(result))
}
}
@@ -1028,7 +968,7 @@ impl WorkspaceOp for CodeFindReferencesReq {
self.context_file.as_ref().map(std::path::PathBuf::from),
)
.await
.map_err(|e| WorkspaceError::HubError(format!("index channel closed: {e}")))?;
.map_err(|e| WorkspaceError::Internal(format!("index channel closed: {e}")))?;
Ok(symbol_locations_to_response(result))
}
}
@@ -1105,9 +1045,9 @@ impl WorkspaceOp for CreateWorktreeRequest {
let result = crate::worktree::prepare_worktree_creation(self).await;
match result.response {
Ok(resp) => {
serde_json::to_value(resp).map_err(|e| WorkspaceError::HubError(e.to_string()))
serde_json::to_value(resp).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
Err(e) => Err(WorkspaceError::HubError(e.to_string())),
Err(e) => Err(WorkspaceError::Internal(e.to_string())),
}
}
}
@@ -1121,8 +1061,8 @@ impl WorkspaceOp for RemoveWorktreeRequest {
let copy_ctx = crate::worktree::BackgroundCopyContext::new();
let result = crate::worktree::remove_worktree(self, &copy_ctx)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::to_value(result).map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::to_value(result).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -1134,8 +1074,8 @@ impl WorkspaceOp for ApplyWorktreeRequest {
) -> WorkspaceResult<Self::Response> {
let result = crate::worktree::apply_worktree(self)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::to_value(result).map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::to_value(result).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -1147,8 +1087,8 @@ impl WorkspaceOp for WorktreeListReq {
) -> WorkspaceResult<Self::Response> {
let records =
crate::worktree::list_worktrees(self.repo.as_deref(), &self.types, self.include_all)
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::to_value(records).map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::to_value(records).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -1159,8 +1099,8 @@ impl WorkspaceOp for WorktreeShowReq {
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
let record = crate::worktree::show_worktree(&self.id_or_path)
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::to_value(record).map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::to_value(record).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -1175,9 +1115,9 @@ impl WorkspaceOp for WorktreeGcReq {
crate::worktree::gc_worktrees_mgmt(dry_run, max_age_secs, force)
})
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))?
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::to_value(report).map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))?
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::to_value(report).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
#[async_trait]
@@ -1188,27 +1128,22 @@ impl WorkspaceOp for WorktreeDbStatsReq {
_session_id: Option<&str>,
) -> WorkspaceResult<Self::Response> {
let stats = crate::worktree::worktree_db_stats()
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
serde_json::to_value(stats).map_err(|e| WorkspaceError::HubError(e.to_string()))
.map_err(|e| WorkspaceError::Internal(e.to_string()))?;
serde_json::to_value(stats).map_err(|e| WorkspaceError::Internal(e.to_string()))
}
}
/// Dual-mode workspace operations handle.
///
/// - **`Local`** — wraps a [`WorkspaceHandle`]. Extensions dispatch
/// through the handle; tool calls dispatch through the workspace
/// session's [`FinalizedToolset`](kigi_tools::registry::types::FinalizedToolset).
/// Call [`bind_local_session`](Self::bind_local_session) after building
/// the agent to install the toolset on the workspace session.
///
/// - **`Proxy`** — wraps a [`WorkspaceClient`] connected to a remote hub.
/// Everything routes through hub WebSocket to a remote workspace server.
/// Wraps a [`WorkspaceHandle`]. Extensions dispatch through the handle;
/// tool calls dispatch through the workspace session's
/// [`FinalizedToolset`](kigi_tools::registry::types::FinalizedToolset).
/// Call [`bind_local_session`](Self::bind_local_session) after building
/// the agent to install the toolset on the workspace session.
#[derive(Clone)]
pub enum WorkspaceOps {
/// Local in-process mode — extensions through the handle, tool calls
/// through the workspace session's toolset.
Local { handle: WorkspaceHandle },
/// Proxy mode — routes through hub RPC.
Proxy { client: WorkspaceClient },
}
impl WorkspaceOps {
/// Construct a local-mode ops handle.
@@ -1219,36 +1154,10 @@ impl WorkspaceOps {
pub fn local(handle: WorkspaceHandle) -> Self {
Self::Local { handle }
}
/// Construct a proxy-mode ops handle.
pub fn proxy(harness: Arc<ToolHarness>) -> Self {
Self::Proxy {
client: WorkspaceClient::new((*harness).clone()),
}
}
/// Construct a proxy-mode ops handle sharing a pre-created connected
/// flag. The same `Arc<AtomicBool>` should be wired into the harness
/// builder's `on_reconnect` callback so reconnects reset the flag.
pub fn proxy_with_connected(harness: Arc<ToolHarness>, connected: Arc<AtomicBool>) -> Self {
Self::Proxy {
client: WorkspaceClient::with_connected_flag((*harness).clone(), connected),
}
}
/// Whether this handle routes through the server (proxy mode).
pub fn is_proxy(&self) -> bool {
matches!(self, Self::Proxy { .. })
}
/// Access the underlying workspace RPC client (proxy mode only).
pub fn client(&self) -> Option<&WorkspaceClient> {
match self {
Self::Proxy { client } => Some(client),
Self::Local { .. } => None,
}
}
/// Access the underlying workspace handle (local mode only).
/// Access the underlying workspace handle.
pub fn workspace_handle(&self) -> Option<&WorkspaceHandle> {
match self {
Self::Local { handle } => Some(handle),
Self::Proxy { .. } => None,
}
}
/// Create the workspace session and bind the agent's toolset for local mode.
@@ -1274,9 +1183,7 @@ impl WorkspaceOps {
toolset: Arc<kigi_tools::registry::types::FinalizedToolset>,
viewer_ctx: Option<kigi_tool_runtime::WorkspaceViewerContext>,
) -> WorkspaceResult<()> {
let Self::Local { handle } = self else {
return Ok(());
};
let Self::Local { handle } = self;
if handle.session(session_id).is_none() {
handle.create_session_with_tracker_and_viewer_ctx(
session_id,
@@ -1296,9 +1203,7 @@ impl WorkspaceOps {
}
/// Release the workspace session. No-op in proxy mode.
pub fn end_local_session(&self, session_id: &str) {
let Self::Local { handle } = self else {
return;
};
let Self::Local { handle } = self;
handle.on_session_ended(session_id);
if let Err(e) = handle.drop_session(session_id, session_id) {
tracing::debug!(
@@ -1312,82 +1217,26 @@ impl WorkspaceOps {
session_id: &str,
payload: &kigi_tool_protocol::turn_hook::BeforeTurnPayload,
) {
match self {
Self::Local { handle } => {
handle.on_before_turn(session_id, payload).await;
}
Self::Proxy { .. } => {
tracing::debug!("on_before_turn called on Proxy WorkspaceOps (no-op)");
}
}
let Self::Local { handle } = self;
handle.on_before_turn(session_id, payload).await;
}
pub async fn on_after_turn(
&self,
session_id: &str,
payload: &kigi_tool_protocol::turn_hook::AfterTurnPayload,
) {
match self {
Self::Local { handle } => {
handle.on_after_turn(session_id, payload).await;
}
Self::Proxy { .. } => {
tracing::debug!("on_after_turn called on Proxy WorkspaceOps (no-op)");
}
}
let Self::Local { handle } = self;
handle.on_after_turn(session_id, payload).await;
}
pub async fn rpc_raw(&self, method: &str, params: Value) -> WorkspaceResult<Value> {
let client = match self {
Self::Proxy { client } => client,
Self::Local { .. } => {
return Err(WorkspaceError::HubError(
"rpc not available in local mode".into(),
));
}
};
client
.rpc_raw(method, params)
.await
.map_err(|e| WorkspaceError::HubError(e.to_string()))
}
async fn rpc<R: WorkspaceRpc>(&self, req: &R) -> WorkspaceResult<R::Response> {
let params = serde_json::to_value(req)
.map_err(|e| WorkspaceError::HubError(format!("serialize failed: {e}")))?;
let terminal = self.rpc_raw(R::METHOD, params).await?;
let envelope: crate::rpc_envelope::RpcEnvelope<R::Response> =
serde_json::from_value(terminal)
.map_err(|e| WorkspaceError::HubError(format!("envelope parse failed: {e}")))?;
envelope
.into_result()
.map_err(crate::rpc_envelope::rpc_error_to_workspace)
}
/// Dispatch a typed operation in either local or proxy mode.
///
/// - **Local mode**: calls `op.execute(handle, session_id)` directly.
/// - **Proxy mode**: serializes the op and routes through the server RPC.
/// The server handler owns session context, so `session_id` is only
/// needed for local `execute()`.
/// Dispatch a typed operation: calls `op.execute(handle, session_id)`.
pub async fn dispatch<Op: WorkspaceOp>(
&self,
op: &Op,
session_id: Option<&str>,
) -> WorkspaceResult<Op::Response> {
let mode = match self {
Self::Local { .. } => "local",
Self::Proxy { .. } => "proxy",
};
tracing::debug!(method = Op::METHOD, mode, "WorkspaceOps::dispatch");
match self {
Self::Local { handle } => op.execute(handle, session_id).await,
Self::Proxy { .. } => self.rpc(op).await,
}
}
pub async fn workspace_info(&self) -> WorkspaceResult<Value> {
self.rpc(&WorkspaceInfoReq {}).await
}
/// **DEPRECATED**: Use [`Self::git_status_ext`] with `format: GitStatusFormat::Prompt`
/// instead. This method will be removed in a future release.
pub async fn git_status(&self) -> WorkspaceResult<Value> {
self.rpc(&GitStatusReq {}).await
tracing::debug!(method = Op::METHOD, "WorkspaceOps::dispatch");
let Self::Local { handle } = self;
op.execute(handle, session_id).await
}
/// Get git status with configurable output format.
///
@@ -1408,52 +1257,9 @@ impl WorkspaceOps {
let wire = self.dispatch(&HookRegistryReq {}, None).await?;
wire_to_hook_registry(&wire)
}
pub async fn begin_prompt(&self, session_id: &str, prompt_index: usize) -> WorkspaceResult<()> {
self.rpc(&BeginPromptReq {
session_id: session_id.to_owned(),
prompt_index,
})
.await
}
pub async fn end_prompt(&self, session_id: &str, prompt_index: usize) -> WorkspaceResult<()> {
self.rpc(&EndPromptReq {
session_id: session_id.to_owned(),
prompt_index,
})
.await
}
pub async fn get_rewind_points(
&self,
session_id: &str,
) -> WorkspaceResult<Vec<crate::session::file_state::RewindPoint>> {
self.rpc(&GetRewindPointsReq {
session_id: session_id.to_owned(),
})
.await
}
pub async fn rewind_to(
&self,
session_id: &str,
target_prompt_index: usize,
) -> WorkspaceResult<crate::session::file_state::FileRewindResponse> {
self.rpc(&RewindToReq {
session_id: session_id.to_owned(),
target_prompt_index,
})
.await
}
pub async fn put_files(&self, req: PutFilesReq) -> WorkspaceResult<PutFilesRes> {
self.dispatch(&req, None).await
}
pub async fn get_files(&self, req: GetFilesReq) -> WorkspaceResult<GetFilesRes> {
self.dispatch(&req, None).await
}
/// Dispatch a tool call through the workspace.
///
/// - **Local**: dispatches through the workspace session's
/// [`FinalizedToolset`](kigi_tools::registry::types::FinalizedToolset)
/// (in-process). Requires `session_id` to look up the session.
/// - **Proxy**: routes through the server `ToolHarness` (remote).
/// Dispatch a tool call through the workspace session's
/// [`FinalizedToolset`](kigi_tools::registry::types::FinalizedToolset)
/// (in-process). Requires `session_id` to look up the session.
pub async fn call_tool(
&self,
name: &str,
@@ -1461,57 +1267,23 @@ impl WorkspaceOps {
call_id: &str,
session_id: Option<&str>,
) -> Result<ToolRunResult, kigi_tool_runtime::ToolError> {
match self {
Self::Local { handle } => {
let session_id = session_id.ok_or_else(|| {
kigi_tool_runtime::ToolError::custom(
"missing_session",
"session_id required for local tool dispatch",
)
})?;
let session = handle.session(session_id).ok_or_else(|| {
kigi_tool_runtime::ToolError::custom(
"session_not_found",
format!(
"workspace session not found: {session_id} \
— call bind_local_session() first"
),
)
})?;
session.toolset().call(name, args, call_id, None).await
}
Self::Proxy { client } => {
if !client.is_connected() {
return Err(kigi_tool_runtime::ToolError::network_error(
"The workspace server connection was lost. \
Please restart your session to reconnect.",
));
}
let tool_id = kigi_tool_protocol::ToolId::new(name).map_err(|e| {
kigi_tool_runtime::ToolError::custom(
"hub_proxy_error",
format!("invalid tool name: {e}"),
)
})?;
let mut ctx = kigi_tool_runtime::ToolCallContext::default();
ctx.call_id =
kigi_tool_protocol::ToolCallId::new(call_id.to_owned()).unwrap_or(ctx.call_id);
let mut stream = client.harness().call(tool_id, args, ctx).await;
let typed = crate::hub_channel::consume_stream_terminal(&mut stream)
.await
.inspect_err(|e| {
if is_transport_fatal(e) {
client.mark_disconnected();
}
})?;
serde_json::from_value::<ToolRunResult>(typed.value).map_err(|e| {
kigi_tool_runtime::ToolError::custom(
"tool_result_deserialize",
format!("tool result deserialization failed: {e}"),
)
})
}
}
let Self::Local { handle } = self;
let session_id = session_id.ok_or_else(|| {
kigi_tool_runtime::ToolError::custom(
"missing_session",
"session_id required for local tool dispatch",
)
})?;
let session = handle.session(session_id).ok_or_else(|| {
kigi_tool_runtime::ToolError::custom(
"session_not_found",
format!(
"workspace session not found: {session_id} \
— call bind_local_session() first"
),
)
})?;
session.toolset().call(name, args, call_id, None).await
}
}
#[cfg(any(test, feature = "test-support"))]
@@ -1569,9 +1341,7 @@ mod tests {
#[test]
fn git_op_cwd_uses_explicit_git_root_per_window() {
let ops = WorkspaceOps::for_test();
let WorkspaceOps::Local { handle } = &ops else {
unreachable!("for_test builds a local handle");
};
let WorkspaceOps::Local { handle } = &ops;
let workspace_root = handle.root_cwd().unwrap();
let window_a = std::path::PathBuf::from("/repos/xai-main");
let window_b = std::path::PathBuf::from("/repos/xai-main-2");
@@ -1594,9 +1364,7 @@ mod tests {
#[tokio::test]
async fn end_local_session_drops_bound_toolset() {
let ops = WorkspaceOps::for_test();
let WorkspaceOps::Local { handle } = &ops else {
unreachable!("for_test builds a local handle");
};
let WorkspaceOps::Local { handle } = &ops;
let sid = "sess-teardown";
let toolset =
std::sync::Arc::new(kigi_tools::registry::types::FinalizedToolset::empty_for_test());