M2 audit: excise managed connectors and xAI media-gen tools

Managed connectors (grok.com MCP admin) removed root-and-branch:
- The managed-MCP fetch/injection pipeline is gone, including the whole
  kigi-shell-session-support crate (managed-config fetch client, gateway
  tool catalog + dispatch, header injection, refresh task), reactive
  managed re-auth, mcp_doctor's grok.com-source discovery, and the
  [managed_mcps] config surface.
- TUI: the 'Managed by grok.com' section, connectors URL/deep-link,
  Action::OpenManagedConnectors, and session_team_id are gone. Local MCP
  management (list/toggle/add/remove/auth/tools) is fully intact.
- Kept as LOCAL policy: managed-settings.json MCP allow/deny enforcement,
  the multi-source local MCP merge, folder-trust gating. PluginOrigin
  Project/User labels kept (they tag locally discovered plugin dirs).

imagine/media-gen tools (xAI image/video generation) removed:
- image_gen, image_edit, video_gen, image_to_video, reference_to_video
  implementations, registrations, ToolKind/ToolInput/Output variants
  (serde-safe), config plumbing end to end, ZDR video machinery,
  /imagine + /imagine-video commands and guidance text, the bundled
  imagine skill (added to legacy cleanup so user installs delete it),
  and the media-gen render path.
- Kept: image INPUT (paste/attach, [Image #N] meta, pdf/image fetch,
  clipboard wrap), generic media-ref rendering, and the generic tool
  401-retry machinery (tests renamed, assertions unweakened).
- deploy_app stays: it is a permanently-disabled local stub deploying
  nowhere.

121 files changed, 8 deleted. Gates: workspace check/clippy 0/0, fmt,
deny ok; suites green (tools 2554, shell 4862, tui 6608, workspace
1042). Remaining grok.com strings live only in the auth-method ids and
changelog archives (§9/M3 sweep).
This commit is contained in:
2026-07-17 23:45:05 -04:00
parent fa75eb139a
commit 5e4e24db99
120 changed files with 301 additions and 11327 deletions
Generated
-22
View File
@@ -6293,7 +6293,6 @@ dependencies = [
"kigi-sandbox", "kigi-sandbox",
"kigi-shared", "kigi-shared",
"kigi-shell-base", "kigi-shell-base",
"kigi-shell-session-support",
"kigi-sqlite-journal", "kigi-sqlite-journal",
"kigi-subagent-resolution", "kigi-subagent-resolution",
"kigi-system-power", "kigi-system-power",
@@ -6389,27 +6388,6 @@ dependencies = [
"windows 0.61.3", "windows 0.61.3",
] ]
[[package]]
name = "kigi-shell-session-support"
version = "0.1.0"
dependencies = [
"agent-client-protocol",
"axum",
"chrono",
"kigi-file-utils",
"kigi-http",
"kigi-shell-base",
"kigi-version",
"kigi-workspace",
"reqwest 0.12.28",
"serde",
"serde_json",
"thiserror 2.0.18",
"tokio",
"tokio-util",
"tracing",
]
[[package]] [[package]]
name = "kigi-sqlite-journal" name = "kigi-sqlite-journal"
version = "0.1.0" version = "0.1.0"
-2
View File
@@ -44,7 +44,6 @@ members = [
"crates/codegen/kigi-shared", "crates/codegen/kigi-shared",
"crates/codegen/kigi-shell", "crates/codegen/kigi-shell",
"crates/codegen/kigi-shell-base", "crates/codegen/kigi-shell-base",
"crates/codegen/kigi-shell-session-support",
"crates/codegen/kigi-subagent-resolution", "crates/codegen/kigi-subagent-resolution",
"crates/codegen/kigi-test-support", "crates/codegen/kigi-test-support",
"crates/codegen/kigi-tools", "crates/codegen/kigi-tools",
@@ -274,7 +273,6 @@ kigi-secrets = { path = "crates/codegen/kigi-secrets" }
kigi-shared = { path = "crates/codegen/kigi-shared" } kigi-shared = { path = "crates/codegen/kigi-shared" }
kigi-shell = { path = "crates/codegen/kigi-shell" } kigi-shell = { path = "crates/codegen/kigi-shell" }
kigi-shell-base = { path = "crates/codegen/kigi-shell-base" } kigi-shell-base = { path = "crates/codegen/kigi-shell-base" }
kigi-shell-session-support = { path = "crates/codegen/kigi-shell-session-support" }
kigi-log = { path = "crates/codegen/kigi-log" } kigi-log = { path = "crates/codegen/kigi-log" }
kigi-test-support = { path = "crates/codegen/kigi-test-support" } kigi-test-support = { path = "crates/codegen/kigi-test-support" }
kigi-tools = { path = "crates/codegen/kigi-tools" } kigi-tools = { path = "crates/codegen/kigi-tools" }
+3 -55
View File
@@ -93,8 +93,6 @@ pub struct AgentBuilder {
backend_search: bool, backend_search: bool,
web_fetch_config: kigi_tools::implementations::grok_build::web_fetch::WebFetchConfig, web_fetch_config: kigi_tools::implementations::grok_build::web_fetch::WebFetchConfig,
lsp: Option<std::sync::Arc<dyn kigi_tools::implementations::lsp::LspBackend>>, lsp: Option<std::sync::Arc<dyn kigi_tools::implementations::lsp::LspBackend>>,
image_gen_config: kigi_tools::implementations::grok_build::image_gen::ImageGenConfig,
video_gen_config: kigi_tools::implementations::grok_build::video_gen::VideoGenConfig,
app_builder_deployer_config: app_builder_deployer_config:
kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig, kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig,
write_file_enabled: bool, write_file_enabled: bool,
@@ -218,8 +216,6 @@ impl AgentBuilder {
backend_search: false, backend_search: false,
web_fetch_config: Default::default(), web_fetch_config: Default::default(),
lsp: None, lsp: None,
image_gen_config: Default::default(),
video_gen_config: Default::default(),
app_builder_deployer_config: Default::default(), app_builder_deployer_config: Default::default(),
write_file_enabled: true, write_file_enabled: true,
subagents_enabled: false, subagents_enabled: false,
@@ -451,33 +447,6 @@ impl AgentBuilder {
self.lsp = Some(handle); self.lsp = Some(handle);
self self
} }
/// Set the image generation configuration.
///
/// When `Enabled`, an `ImageGenClient` is created and injected into
/// the ToolBridge's resources and the `image_gen` tool is registered,
/// allowing image generation via the xAI Imagine API with session
/// credentials. When `Disabled` (default), the tool is not registered.
pub fn with_image_gen_config(
mut self,
config: kigi_tools::implementations::grok_build::image_gen::ImageGenConfig,
) -> Self {
self.image_gen_config = config;
self
}
/// Set the video generation configuration.
///
/// When `Enabled`, a `VideoGenClient` is created and injected into
/// the ToolBridge's resources and the `video_gen` tool is registered,
/// allowing video generation via the xAI Video Generation API with
/// session credentials. When `Disabled` (default), the tool is not
/// registered.
pub fn with_video_gen_config(
mut self,
config: kigi_tools::implementations::grok_build::video_gen::VideoGenConfig,
) -> Self {
self.video_gen_config = config;
self
}
/// Set the deploy service configuration. /// Set the deploy service configuration.
pub fn with_app_builder_deployer_config( pub fn with_app_builder_deployer_config(
mut self, mut self,
@@ -495,10 +464,9 @@ impl AgentBuilder {
self self
} }
/// Set the 401-attribution callback for tool HTTP clients /// Set the 401-attribution callback for tool HTTP clients
/// (`image_gen`, `video_gen`, `web_search`). When set, a 401 /// (`web_search`). When set, a 401 from such a tool emits an
/// from any of those tools emits an `auth_401_attribution` /// `auth_401_attribution` event with `consumer` of
/// event with `consumer` of `"ImageGen"` / `"VideoGen.start"` / /// `"WebSearch"`. Callers should pass the
/// `"VideoGen.poll"` / `"WebSearch"`. Callers should pass the
/// same `ShellAttribution` instance they wire into /// same `ShellAttribution` instance they wire into
/// `kigi_sampler::SamplerConfig::attribution_callback` so /// `kigi_sampler::SamplerConfig::attribution_callback` so
/// all 401s share the same `AuthManager` reference and land in /// all 401s share the same `AuthManager` reference and land in
@@ -712,24 +680,6 @@ impl AgentBuilder {
.tools .tools
.push((&kigi_tools::implementations::grok_build::LspTool).into()); .push((&kigi_tools::implementations::grok_build::LspTool).into());
} }
if self.image_gen_config.image_gen_enabled() {
tool_config
.tools
.push((&kigi_tools::implementations::grok_build::ImageGenTool).into());
}
if self.image_gen_config.image_edit_enabled() {
tool_config
.tools
.push((&kigi_tools::implementations::grok_build::ImageEditTool).into());
}
if self.video_gen_config.is_enabled() {
tool_config
.tools
.push((&kigi_tools::implementations::grok_build::ImageToVideoTool).into());
tool_config
.tools
.push((&kigi_tools::implementations::grok_build::ReferenceToVideoTool).into());
}
let has_write_tool = tool_config let has_write_tool = tool_config
.tools .tools
.iter() .iter()
@@ -1024,8 +974,6 @@ impl AgentBuilder {
web_search_config: self.web_search_config, web_search_config: self.web_search_config,
web_fetch_config: self.web_fetch_config, web_fetch_config: self.web_fetch_config,
lsp: self.lsp, lsp: self.lsp,
image_gen_config: self.image_gen_config,
video_gen_config: self.video_gen_config,
app_builder_deployer_config: self.app_builder_deployer_config, app_builder_deployer_config: self.app_builder_deployer_config,
api_key_provider: self.api_key_provider, api_key_provider: self.api_key_provider,
attribution_callback: self.attribution_callback, attribution_callback: self.attribution_callback,
+1 -7
View File
@@ -180,9 +180,6 @@ pub fn workspace_grok_build_toolset() -> ToolServerConfig {
tools.push((&grok_build::ExitPlanModeTool).into()); tools.push((&grok_build::ExitPlanModeTool).into());
tools.push((&grok_build::AskUserQuestionTool).into()); tools.push((&grok_build::AskUserQuestionTool).into());
tools.push((&grok_build::WebSearchTool).into()); tools.push((&grok_build::WebSearchTool).into());
tools.push((&grok_build::ImageGenTool).into());
tools.push((&grok_build::ImageToVideoTool).into());
tools.push((&grok_build::ReferenceToVideoTool).into());
tools.push((&grok_build::WebFetchTool).into()); tools.push((&grok_build::WebFetchTool).into());
tools.push((&memory::search_tool::MemorySearchImpl).into()); tools.push((&memory::search_tool::MemorySearchImpl).into());
tools.push((&memory::get_tool::MemoryGetImpl).into()); tools.push((&memory::get_tool::MemoryGetImpl).into());
@@ -447,9 +444,6 @@ fn orchestrator_toolset() -> ToolServerConfig {
(&grok_build::MonitorTool).into(), (&grok_build::MonitorTool).into(),
(&grok_build::WebSearchTool).into(), (&grok_build::WebSearchTool).into(),
(&grok_build::WebFetchTool).into(), (&grok_build::WebFetchTool).into(),
(&grok_build::ImageGenTool).into(),
(&grok_build::ImageToVideoTool).into(),
(&grok_build::ReferenceToVideoTool).into(),
(&memory::MemorySearchImpl).into(), (&memory::MemorySearchImpl).into(),
(&memory::MemoryGetImpl).into(), (&memory::MemoryGetImpl).into(),
], ],
@@ -744,7 +738,7 @@ pub struct AgentDefinition {
pub agents_md: bool, pub agents_md: bool,
/// When true (the default), the AgentBuilder layers session-level optional /// When true (the default), the AgentBuilder layers session-level optional
/// tools on top of the agent's declared `tool_config`: memory_search/get, /// tools on top of the agent's declared `tool_config`: memory_search/get,
/// web_search, web_fetch, lsp, image_gen, video_gen, OpenCode write /// web_search, web_fetch, lsp, OpenCode write
/// fallback, and the plan-mode tools. /// fallback, and the plan-mode tools.
/// ///
/// Set this to `false` for harnesses that need an exact, minimal toolset /// Set this to `false` for harnesses that need an exact, minimal toolset
@@ -171,7 +171,7 @@ fn process_hooks_content(
// Substitute plugin env vars in command paths at config-load time so // Substitute plugin env vars in command paths at config-load time so
// that hooks like `${CLAUDE_PLUGIN_ROOT}/hooks/foo.sh` resolve to the // that hooks like `${CLAUDE_PLUGIN_ROOT}/hooks/foo.sh` resolve to the
// real plugin directory regardless of which spawn branch the runner // real plugin directory regardless of which spawn branch the runner
// takes (mirrors what managed_mcp does for MCP server commands). // takes (mirrors what the MCP server-command loader does).
if let Some(cmd) = &spec.command { if let Some(cmd) = &spec.command {
let cmd_str = cmd.to_string_lossy(); let cmd_str = cmd.to_string_lossy();
// Mirror what `managed_mcp::load_plugin_mcp_servers_from_config` // Mirror what `managed_mcp::load_plugin_mcp_servers_from_config`
@@ -445,11 +445,6 @@ pub struct RemoteSettings {
skip_serializing_if = "Vec::is_empty" skip_serializing_if = "Vec::is_empty"
)] )]
pub goal_skeptic_models: Vec<GoalRoleModel>, pub goal_skeptic_models: Vec<GoalRoleModel>,
/// Remote fallback for managed MCP connector fetching.
#[serde(default)]
pub managed_mcps_enabled: Option<bool>,
#[serde(default)]
pub managed_mcp_gateway_tools_enabled: Option<bool>,
/// Enable user-facing feedback (heuristic popups, `/feedback` command). /// Enable user-facing feedback (heuristic popups, `/feedback` command).
#[serde(default)] #[serde(default)]
pub feedback_enabled: Option<bool>, pub feedback_enabled: Option<bool>,
@@ -517,20 +512,6 @@ pub struct RemoteSettings {
/// is set in config.toml. Absent → default (**disabled** — ships dark). /// is set in config.toml. Absent → default (**disabled** — ships dark).
#[serde(default)] #[serde(default)]
pub subagent_worktree_snapshot_enabled: Option<bool>, pub subagent_worktree_snapshot_enabled: Option<bool>,
/// When `Some(true)`, enable the `image_gen` tool for session-based auth users.
/// When `Some(false)` or absent, the tool is hidden regardless of credentials.
#[serde(default)]
pub image_gen_enabled: Option<bool>,
/// remote settings flag: optional Imagine model override for `image_gen`.
/// When present and non-empty, `image_gen` uses this model slug
/// (e.g. `grok-imagine-image`) instead of the default quality model
/// (`grok-imagine-image-quality`). Absent/empty → default model.
#[serde(default)]
pub image_gen_model_override: Option<String>,
/// When `Some(true)`, enable the `video_gen` tool for session-based auth users.
/// When `Some(false)` or absent, the tool is hidden regardless of credentials.
#[serde(default)]
pub video_gen_enabled: Option<bool>,
/// When `Some(true)`, enable the process-wide image normalize cache that /// When `Some(true)`, enable the process-wide image normalize cache that
/// amortises decode + integrity-check + re-encode work across SessionActors. /// amortises decode + integrity-check + re-encode work across SessionActors.
/// Default: disabled. See `session::normalize_cache`. /// Default: disabled. See `session::normalize_cache`.
@@ -687,13 +668,6 @@ pub struct RemoteSettings {
/// remote settings verbatim-input flag; env (`KIGI_COMPACTION_VERBATIM_INPUT`) and config override it. `None` = default (true). /// remote settings verbatim-input flag; env (`KIGI_COMPACTION_VERBATIM_INPUT`) and config override it. `None` = default (true).
#[serde(default)] #[serde(default)]
pub compaction_verbatim_input: Option<bool>, pub compaction_verbatim_input: Option<bool>,
/// remote settings denylist of optional imagine tools to disable
/// (e.g. `["image_edit"]`). When a tool is listed it is authoritatively
/// removed from the toolset and local env/config can't re-enable it.
/// Absent or not listed → each tool keeps its own default.
/// See `Config::resolve_image_edit`.
#[serde(default)]
pub imagine_tools_disabled: Option<Vec<String>>,
/// Master switch for jemalloc heap sampling + threshold dumps. /// Master switch for jemalloc heap sampling + threshold dumps.
/// `Some(true)` enables, `Some(false)` kill-switch, `None` = client default off. /// `Some(true)` enables, `Some(false)` kill-switch, `None` = client default off.
#[serde(default)] #[serde(default)]
@@ -706,16 +680,6 @@ pub struct RemoteSettings {
#[serde(default)] #[serde(default)]
pub jemalloc_heap_profile_poll_interval_secs: Option<u64>, pub jemalloc_heap_profile_poll_interval_secs: Option<u64>,
} }
impl RemoteSettings {
/// Denylist check for an optional imagine tool. Returns `true` when the
/// server sent `imagine_tools_disabled` and it contains `tool` (force-off);
/// otherwise `false` (defer to the tool's own default).
pub fn imagine_tool_disabled(&self, tool: &str) -> bool {
self.imagine_tools_disabled
.as_ref()
.is_some_and(|list| list.iter().any(|t| t == tool))
}
}
/// Remote enable tier for the per-tip contextual hints (mirrors the client's /// Remote enable tier for the per-tip contextual hints (mirrors the client's
/// `[ui.contextual_hints]` shape). Each field is a soft default for one tip; /// `[ui.contextual_hints]` shape). Each field is a soft default for one tip;
/// `None` defers to the client default (on). All fields `#[serde(default)]` so /// `None` defers to the client default (on). All fields `#[serde(default)]` so
+3 -17
View File
@@ -46,7 +46,7 @@ pub use kigi_workspace_types::MCP_TOOL_NAME_DELIMITER;
/// Normalize an MCP server URL for comparison: strip trailing slashes. /// Normalize an MCP server URL for comparison: strip trailing slashes.
/// Must match the normalization the host's managed-config layer uses /// Must match the normalization the host's managed-config layer uses
/// (e.g. shell's `session::managed_mcp::normalize_url`) so refresh /// (e.g. shell's `session::managed_mcp::normalize_url`) so URL
/// lookup keys agree. /// lookup keys agree.
fn normalize_url(url: &str) -> String { fn normalize_url(url: &str) -> String {
url.trim_end_matches('/').to_string() url.trim_end_matches('/').to_string()
@@ -112,7 +112,7 @@ pub struct McpConfigDiff {
pub retained: Vec<McpServerName>, pub retained: Vec<McpServerName>,
} }
/// MCP server name used as the key in client/tool maps (e.g. `"github"`, `"grok_com_linear"`). /// MCP server name used as the key in client/tool maps (e.g. `"github"`).
pub type McpServerName = String; pub type McpServerName = String;
/// Unqualified MCP tool name (e.g. `"create_issue"`, without the `server__` prefix). /// Unqualified MCP tool name (e.g. `"create_issue"`, without the `server__` prefix).
@@ -4012,16 +4012,6 @@ fn is_figma_mcp(server_name: &str, url: &str) -> bool {
if server_name.eq_ignore_ascii_case("figma") { if server_name.eq_ignore_ascii_case("figma") {
return true; return true;
} }
// Legacy direct managed name (`grok_com_figma`); newer clients use gateway tools (`managed_mcp_gateway_tools_enabled`).
const MANAGED_PREFIX: &str = "grok_com_";
if let (Some(prefix), Some(rest)) = (
server_name.get(..MANAGED_PREFIX.len()),
server_name.get(MANAGED_PREFIX.len()..),
) && prefix.eq_ignore_ascii_case(MANAGED_PREFIX)
&& rest.eq_ignore_ascii_case("figma")
{
return true;
}
reqwest::Url::parse(url) reqwest::Url::parse(url)
.ok() .ok()
.and_then(|u| u.host_str().map(|h| h.to_ascii_lowercase())) .and_then(|u| u.host_str().map(|h| h.to_ascii_lowercase()))
@@ -4510,14 +4500,10 @@ mod tests {
fn is_figma_mcp_matches_name_and_host() { fn is_figma_mcp_matches_name_and_host() {
assert!(is_figma_mcp("figma", "https://example.com/mcp")); assert!(is_figma_mcp("figma", "https://example.com/mcp"));
assert!(is_figma_mcp("Figma", "https://example.com/mcp")); assert!(is_figma_mcp("Figma", "https://example.com/mcp"));
assert!(is_figma_mcp("grok_com_figma", "https://example.com/mcp"));
assert!(is_figma_mcp("GROK_COM_FIGMA", "https://example.com/mcp"));
assert!(is_figma_mcp("grok_com_FIGMA", "https://example.com/mcp"));
assert!(is_figma_mcp("other", "https://mcp.figma.com/mcp")); assert!(is_figma_mcp("other", "https://mcp.figma.com/mcp"));
assert!(is_figma_mcp("other", "https://figma.com/mcp")); assert!(is_figma_mcp("other", "https://figma.com/mcp"));
assert!(!is_figma_mcp("linear", "https://mcp.linear.app/mcp")); assert!(!is_figma_mcp("linear", "https://mcp.linear.app/mcp"));
assert!(!is_figma_mcp("figma_extra", "https://example.com/mcp")); assert!(!is_figma_mcp("figma_extra", "https://example.com/mcp"));
assert!(!is_figma_mcp("grok_com_linear", "https://example.com/mcp"));
assert!(!is_figma_mcp("linear", "not-a-url")); assert!(!is_figma_mcp("linear", "not-a-url"));
assert!(!is_figma_mcp("linear", "https://notfigma.com/mcp")); assert!(!is_figma_mcp("linear", "https://notfigma.com/mcp"));
assert!(!is_figma_mcp("linear", "https://figma.com.evil/mcp")); assert!(!is_figma_mcp("linear", "https://figma.com.evil/mcp"));
@@ -6762,7 +6748,7 @@ mod tests {
fn is_auth_rejection_message_matches_auth_signals() { fn is_auth_rejection_message_matches_auth_signals() {
// The verbatim string captured in production for a managed handshake. // The verbatim string captured in production for a managed handshake.
assert!(is_auth_rejection_message( assert!(is_auth_rejection_message(
"MCP server 'grok_com_notion' handshake failed: Auth required, when send initialize request" "MCP server 'notion' handshake failed: Auth required, when send initialize request"
)); ));
assert!(is_auth_rejection_message("401 Unauthorized")); assert!(is_auth_rejection_message("401 Unauthorized"));
assert!(is_auth_rejection_message("unauthorized")); assert!(is_auth_rejection_message("unauthorized"));
@@ -461,7 +461,6 @@ fn render_mcps(
badge: badges[i].as_str(), badge: badges[i].as_str(),
badge_color: badge_colors[i], badge_color: badge_colors[i],
collapsible: collapsibles[i], collapsible: collapsibles[i],
underline_last_desc: false,
}) })
}) })
.collect(); .collect();
@@ -516,7 +515,7 @@ fn measure_entries(entries: &[PickerEntry<'_>]) -> u16 {
mod tests { mod tests {
use super::*; use super::*;
use kigi_tui::views::extensions_modal::ExtensionsModalState; use kigi_tui::views::extensions_modal::ExtensionsModalState;
use kigi_tui::views::mcps_modal::{McpServerDisplayStatus, McpServerInfo, McpWireSource}; use kigi_tui::views::mcps_modal::{McpServerDisplayStatus, McpServerInfo};
use ratatui::layout::Rect; use ratatui::layout::Rect;
fn agent() -> AgentView { fn agent() -> AgentView {
@@ -533,9 +532,7 @@ mod tests {
tools: Vec::new(), tools: Vec::new(),
enabled: true, enabled: true,
source: "local".to_string(), source: "local".to_string(),
wire_source: McpWireSource::Local,
plugin_name: None, plugin_name: None,
is_managed_gateway: false,
} }
} }
@@ -2,7 +2,7 @@
//! the terminal via the kitty graphics protocol. //! the terminal via the kitty graphics protocol.
//! //!
//! Typing `/gboom` (and nothing else) opens a modal overlay — the same //! Typing `/gboom` (and nothing else) opens a modal overlay — the same
//! surface the imagine-video player uses — and streams PNG frames via //! surface the inline video player uses — and streams PNG frames via
//! per-frame kitty `a=T` retransmission at the ~30 fps animation tick. The //! per-frame kitty `a=T` retransmission at the ~30 fps animation tick. The
//! simulation steps with wall-clock `dt`, so gameplay speed is independent //! simulation steps with wall-clock `dt`, so gameplay speed is independent
//! of the achieved frame rate. //! of the achieved frame rate.
@@ -92,7 +92,7 @@ fn build_open_path_command(path: &std::path::Path) -> std::process::Command {
/// - **Windows**: `explorer.exe /select,<path>` reveals + highlights the file /// - **Windows**: `explorer.exe /select,<path>` reveals + highlights the file
/// in Explorer. We deliberately avoid `cmd /c start`, whose `%VAR%` /// in Explorer. We deliberately avoid `cmd /c start`, whose `%VAR%`
/// expansion corrupts the percent-encoded session-directory segment in /// expansion corrupts the percent-encoded session-directory segment in
/// imagine media paths (e.g. `…\C%3A%5CUsers…`). /// media paths (e.g. `…\C%3A%5CUsers…`).
/// - **macOS / Linux**: `open` / `xdg-open` open the file in its default app. /// - **macOS / Linux**: `open` / `xdg-open` open the file in its default app.
pub fn open_path(path: &std::path::Path) -> bool { pub fn open_path(path: &std::path::Path) -> bool {
// Never launch a real GUI app in tests. // Never launch a real GUI app in tests.
@@ -210,28 +210,6 @@ pub fn open_url_if_safe(url: &str, filter: SchemeFilter) -> bool {
} }
} }
/// Ensure `url` carries the given query parameter, returning the rewritten URL.
///
/// If the URL already contains a parameter with that name, its value is left
/// untouched (the caller upstream may have intentionally set one). On parse
/// failure, the original string is returned unchanged so this is safe to apply
/// to opener input from untrusted sources.
///
/// Used by the SuperGrok upsell flow to attribute clicks to `referrer=grok-build`,
/// matching the OAuth consent screen and x.ai/cli marketing links regardless of
/// what the remote settings `gate_url` value happens to be.
pub fn ensure_query_param(url: &str, key: &str, value: &str) -> String {
let Ok(mut parsed) = url::Url::parse(url) else {
return url.to_string();
};
let already_present = parsed.query_pairs().any(|(k, _)| k == key);
if already_present {
return parsed.to_string();
}
parsed.query_pairs_mut().append_pair(key, value);
parsed.to_string()
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
@@ -383,55 +361,6 @@ mod tests {
)); ));
} }
#[test]
fn ensure_query_param_appends_when_missing() {
let out = ensure_query_param("https://grok.com/supergrok", "referrer", "grok-build");
assert_eq!(out, "https://grok.com/supergrok?referrer=grok-build");
}
#[test]
fn ensure_query_param_preserves_existing_value() {
let out = ensure_query_param(
"https://grok.com/supergrok?referrer=other",
"referrer",
"grok-build",
);
assert_eq!(out, "https://grok.com/supergrok?referrer=other");
}
#[test]
fn ensure_query_param_keeps_other_query_pairs() {
let out = ensure_query_param(
"https://grok.com/supergrok?heavy=1",
"referrer",
"grok-build",
);
assert_eq!(
out,
"https://grok.com/supergrok?heavy=1&referrer=grok-build"
);
}
#[test]
fn ensure_query_param_preserves_fragment() {
// The current remote settings value uses a hash fragment for client-side
// routing (`grok.com/#supergrok`); we still want the referrer attached.
let out = ensure_query_param("https://grok.com/#supergrok", "referrer", "grok-build");
assert_eq!(out, "https://grok.com/?referrer=grok-build#supergrok");
}
#[test]
fn ensure_query_param_returns_unchanged_on_parse_failure() {
let out = ensure_query_param("not a url", "referrer", "grok-build");
assert_eq!(out, "not a url");
}
#[test]
fn ensure_query_param_url_encodes_value() {
let out = ensure_query_param("https://grok.com/supergrok", "referrer", "grok build");
assert_eq!(out, "https://grok.com/supergrok?referrer=grok+build");
}
#[test] #[test]
fn fallback_scheme_case_insensitive() { fn fallback_scheme_case_insensitive() {
// Uppercase scheme that url::Url::parse rejects triggers fallback path; // Uppercase scheme that url::Url::parse rejects triggers fallback path;
@@ -1879,7 +1879,7 @@ impl ScrollbackVideoRef {
pub fn extract_video_refs(text: &str) -> Vec<ScrollbackVideoRef> { pub fn extract_video_refs(text: &str) -> Vec<ScrollbackVideoRef> {
use std::sync::LazyLock; use std::sync::LazyLock;
// Reuse the markdown image ref pattern — video_gen uses ![prompt](path.mp4). // Reuse the markdown image ref pattern — video refs use ![prompt](path.mp4).
static MD_RE: LazyLock<regex::Regex> = static MD_RE: LazyLock<regex::Regex> =
LazyLock::new(|| regex::Regex::new(MARKDOWN_IMAGE_REF_PATTERN).unwrap()); LazyLock::new(|| regex::Regex::new(MARKDOWN_IMAGE_REF_PATTERN).unwrap());
@@ -254,7 +254,7 @@ struct RowSegment {
/// joiner to the *previous* row (see `BlockLine::joiner`): `None` = hard /// joiner to the *previous* row (see `BlockLine::joiner`): `None` = hard
/// break, `Some("")` = mid-word wrap, `Some(" ")` = word wrap. Consecutive /// break, `Some("")` = mid-word wrap, `Some(" ")` = word wrap. Consecutive
/// rows connected by `Some(..)` joiners are re-joined into one logical line /// rows connected by `Some(..)` joiners are re-joined into one logical line
/// before matching, so a long path or URL soft-wrapped across rows (imagine /// before matching, so a long path or URL soft-wrapped across rows (media
/// media lives at `~/.kigi/sessions/%2F…/images/1.jpg`, which wraps in /// media lives at `~/.kigi/sessions/%2F…/images/1.jpg`, which wraps in
/// narrow panes) is detected whole and each row's fragment gets its own /// narrow panes) is detected whole and each row's fragment gets its own
/// clickable overlay region. Spans within a row are likewise concatenated so /// clickable overlay region. Spans within a row are likewise concatenated so
@@ -858,7 +858,7 @@ mod tests {
#[test] #[test]
fn scan_detects_grok_session_media_path() { fn scan_detects_grok_session_media_path() {
// Dot-directory (`.kigi`), percent-encoded session segment, and a // Dot-directory (`.kigi`), percent-encoded session segment, and a
// trailing sentence period — the shape of `image_gen` output prose. // trailing sentence period — the shape of media-tool output prose.
let line = make_line("Saved to /Users/alice/.kigi/sessions/%2Fabc/00000000/images/1.jpg."); let line = make_line("Saved to /Users/alice/.kigi/sessions/%2Fabc/00000000/images/1.jpg.");
let mut overlay = LinkOverlay::new(); let mut overlay = LinkOverlay::new();
scan_unjoined(std::iter::once((0, &line)), 0, &[], &mut overlay); scan_unjoined(std::iter::once((0, &line)), 0, &[], &mut overlay);
@@ -873,7 +873,7 @@ mod tests {
#[test] #[test]
fn scan_detects_media_path_soft_wrapped_across_rows() { fn scan_detects_media_path_soft_wrapped_across_rows() {
// Regression: `image_gen` output prose wraps the long session path // Regression: media-tool output prose wraps the long session path
// across visual rows (`joiner: Some("")` mid-word break). Previously // across visual rows (`joiner: Some("")` mid-word break). Previously
// each row was scanned in isolation, so only the `/Users/alice` // each row was scanned in isolation, so only the `/Users/alice`
// fragment on the first row matched and became clickable. // fragment on the first row matched and became clickable.
@@ -79,7 +79,8 @@ pub const MAX_PLACEHOLDER_AGGREGATE_BYTES: usize = 200 * 1024 * 1024;
/// `_meta` key under which an attached image's `[Image #N]` display number /// `_meta` key under which an attached image's `[Image #N]` display number
/// is recorded on its ACP image block, so the server can resolve /// is recorded on its ACP image block, so the server can resolve
/// `[Image #N]` tokens to the right attachment by number rather than list /// `[Image #N]` tokens to the right attachment by number rather than list
/// position (the two diverge — see `AttachedImages` in `kigi-tools`). /// position (numbers are not renumbered when a chip is removed
/// mid-compose, so the two diverge).
pub const IMAGE_DISPLAY_NUMBER_META_KEY: &str = "xai.dev/imageDisplayNumber"; pub const IMAGE_DISPLAY_NUMBER_META_KEY: &str = "xai.dev/imageDisplayNumber";
/// Build an ACP image-block `_meta` value carrying `display_number` under /// Build an ACP image-block `_meta` value carrying `display_number` under
@@ -102,34 +103,6 @@ pub fn display_number_from_meta(meta: Option<&agent_client_protocol::Meta>) -> O
.and_then(|n| usize::try_from(n).ok()) .and_then(|n| usize::try_from(n).ok())
} }
/// Build the per-turn `[Image #N]` → reference registry (see
/// [`AttachedImages`](kigi_tools::types::resources::AttachedImages))
/// from the user's inline attached images.
///
/// The display number comes from each block's `_meta` (set by the TUI),
/// falling back to 1-based position for callers that don't record it. The
/// reference is one `image_edit`'s resolver can read directly: the bare
/// durable path (from the `file://` URI) when present, else a
/// `data:<mime>;base64,<data>` URL.
pub fn attached_image_references(
images: &[agent_client_protocol::ImageContent],
) -> Vec<(usize, String)> {
images
.iter()
.enumerate()
.map(|(idx, image)| {
let display_number = display_number_from_meta(image.meta.as_ref()).unwrap_or(idx + 1);
let reference =
if let Some(path) = image.uri.as_deref().and_then(|u| u.strip_prefix("file://")) {
path.to_owned()
} else {
format!("data:{};base64,{}", image.mime_type, image.data)
};
(display_number, reference)
})
.collect()
}
/// File extensions accepted by the placeholder loader. /// File extensions accepted by the placeholder loader.
/// ///
/// SVG is intentionally **not** in this list: SVG is XML text with no /// SVG is intentionally **not** in this list: SVG is XML text with no
@@ -1460,67 +1433,4 @@ mod tests {
assert_eq!(loaded.mime_type, "image/png"); assert_eq!(loaded.mime_type, "image/png");
assert_eq!(loaded.data, PNG_BYTES); assert_eq!(loaded.data, PNG_BYTES);
} }
#[test]
fn attached_image_references_prefers_file_path_over_data() {
// `[Image #N]` resolution should hand `image_edit` a bare on-disk
// path (from the durable `file://` URI) so it reads the session
// copy rather than re-decoding a large base64 blob.
let img = agent_client_protocol::ImageContent::new("AAAA", "image/png")
.uri(Some(
"file:///Users/me/.kigi/sessions/s/images/image-1.png".into(),
))
.meta(display_number_meta(1));
let refs = attached_image_references(std::slice::from_ref(&img));
assert_eq!(
refs,
vec![(
1,
"/Users/me/.kigi/sessions/s/images/image-1.png".to_string()
)]
);
}
#[test]
fn attached_image_references_falls_back_to_data_url() {
// No durable URI (e.g. persistence failed): keep the inline bytes
// as a data URL so the token still resolves to the right image.
let img = agent_client_protocol::ImageContent::new("BBBB", "image/jpeg")
.meta(display_number_meta(2));
let refs = attached_image_references(std::slice::from_ref(&img));
assert_eq!(refs, vec![(2, "data:image/jpeg;base64,BBBB".to_string())]);
}
#[test]
fn attached_image_references_keys_by_meta_number_not_position() {
// Non-contiguous numbers (`#1`, `#3`) survive a mid-compose chip
// removal; the registry must key on the recorded number, not the
// list position.
let mk = |data: &str, n: usize| {
agent_client_protocol::ImageContent::new(data, "image/png").meta(display_number_meta(n))
};
let refs = attached_image_references(&[mk("first", 1), mk("third", 3)]);
assert_eq!(
refs,
vec![
(1, "data:image/png;base64,first".to_string()),
(3, "data:image/png;base64,third".to_string()),
]
);
}
#[test]
fn attached_image_references_falls_back_to_position_without_meta() {
// Older client / non-TUI caller with no recorded number: fall back
// to 1-based position so the common contiguous case still resolves.
let mk = |data: &str| agent_client_protocol::ImageContent::new(data, "image/png");
let refs = attached_image_references(&[mk("first"), mk("second")]);
assert_eq!(
refs,
vec![
(1, "data:image/png;base64,first".to_string()),
(2, "data:image/png;base64,second".to_string()),
]
);
}
} }
@@ -1,32 +0,0 @@
[package]
license = "Apache-2.0"
name = "kigi-shell-session-support"
version.workspace = true
edition.workspace = true
description = "Session-support modules for the grok shell crate family: managed MCP credential/catalog caching and file-access tracking."
[dependencies]
agent-client-protocol = { workspace = true }
chrono = { workspace = true }
reqwest = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
thiserror = { workspace = true }
tokio = { workspace = true, features = ["sync", "rt", "time", "macros"] }
tokio-util = { workspace = true }
tracing = { workspace = true }
kigi-file-utils = { path = "../kigi-file-utils", optional = true }
kigi-http = { workspace = true }
kigi-shell-base = { workspace = true }
kigi-version = { workspace = true }
kigi-workspace = { workspace = true }
[features]
default-bazel = []
[dev-dependencies]
axum = { workspace = true }
tokio = { workspace = true, features = ["net", "rt-multi-thread"] }
[lints]
workspace = true
@@ -1,11 +0,0 @@
#![allow(
unused_imports,
unused_variables,
unused_mut,
unreachable_code,
dead_code
)]
//! Session-support modules extracted from `kigi-shell`'s `session/` tree
//! (which re-exports them at their original paths) so they build in parallel
//! and stop rebuilding on shell edits.
pub mod managed_mcp;
File diff suppressed because it is too large Load Diff
-1
View File
@@ -115,7 +115,6 @@ prost = { workspace = true, optional = true }
kigi-paths = { path = "../kigi-paths" } kigi-paths = { path = "../kigi-paths" }
kigi-shell-base = { workspace = true } kigi-shell-base = { workspace = true }
kigi-env = { workspace = true } kigi-env = { workspace = true }
kigi-shell-session-support = { workspace = true }
kigi-shared = { workspace = true } kigi-shared = { workspace = true }
kigi-tracing-macros = { path = "../kigi-tracing-macros" } kigi-tracing-macros = { path = "../kigi-tracing-macros" }
kigi-fsnotify = { path = "../kigi-fsnotify" } kigi-fsnotify = { path = "../kigi-fsnotify" }
@@ -1,131 +0,0 @@
---
name: imagine
description: >
How to use the image_gen and image_edit tool calls in Grok Build: when to
build a visual with code instead of generating it, prompt-craft,
reference-first handling of real people, factual grounding, and
asset-consistency. Load this whenever generating or editing an image is on the
table, i.e. when an image_gen or image_edit call is being considered or about
to be made. Tool-usage-driven, not triggered by a user merely mentioning
images.
metadata:
short-description: "Prompting and workflow guidance for Imagine image tools"
---
# Imagine
Guidance for the two image tool calls in Grok Build:
- `image_gen` - generate a **new** image from a text prompt.
- `image_edit` - modify an **existing** image using a text prompt and source image.
Apply this whenever you're considering or about to call either tool.
## Build accurate visuals with code, not the image tools
1. **Image models are unreliable at exact text, numbers, and structure.** They can handle short text or a simple layout, but they often garble words, invent numbers, draw chart bars that match no data, or point diagram arrows nowhere, and the more that has to be exact, the worse they do. A detailed prompt doesn't make it dependable, and an `image_edit` pass usually won't fix it. So when a result needs specific text, data, or structure to be correct (charts from real numbers, labeled or technical diagrams, math explainers, tables, screens with real copy), construct the asset with code, where you control the exact content. Prefer HTML and CSS, which give much better layout, typography, and polish than Python plotting. When only the look matters (photos, illustrations, characters, scenes, decorative art), the image tools are the right choice. Which one fits depends on what the output needs to get right, not on how the request is worded.
## Verifying discrete accuracy (loop)
When the output must get specific text, numbers, data, or structure right, don't trust the first result - verify it in a loop:
1. Produce the result (generate, or per *Build accurate visuals with code*, construct it in code).
2. Inspect the actual output - use image understanding to read a generated image back (or check the rendered code) - and confirm every word, number, label, and structural detail matches the requirement, and that nothing overlaps, clips, or runs off-canvas.
3. If anything is wrong, fix and re-verify:
- Garbled text, invented numbers, or broken layout from an image model? Don't just re-prompt - it will likely garble it again. Rebuild it with code.
- Overlapping or clipped elements in code-built output? Re-lay-out with auto-layout (HTML/CSS) rather than nudging coordinates by hand.
- Otherwise make one targeted edit.
4. Only finish when the discrete content is exactly correct. If it can't be made accurate, tell the user instead of shipping something wrong.
## Core Principles
1. **You own the prompt.** If the user gives a detailed prompt or asks you to use theirs, use it verbatim. Otherwise craft the final prompt: front-load the subject, give strong high-level direction for mood, composition, lighting, and style without over-specifying every detail, write natural prose rather than keyword tags, and describe positively instead of using negative prompts. For edits, describe only what changes. Target 2-5 sentences.
2. **Reference-first for real people.** Never use pure `image_gen` for a named real person or group, including face swaps, posters, cartoons, and cinematic or editorial depictions. Use `image_edit` with a real reference instead, and never produce non-consensual, sexualized, or minor-involving likenesses. See Real People and References for the procedure.
3. **Ground facts with search first.** If any part of the request depends on a real-world fact, identity, brand or product, place, event, or top/latest/current result, search the web before generating and put the actual verified details into the prompt. Don't rely on memory, and don't write vague placeholders like "the current president"; write the verified name.
4. **Reuse a base image for consistency.** When the same character, object, or setting must appear across multiple images, generate one base image first, then use it as the input to `image_edit` for every variation. Don't re-run `image_gen` from scratch for a recurring subject.
5. **Handle failures gracefully.** On a moderation or safety block, stop; don't retry and don't paraphrase the prompt to evade the filter. Tell the user it was blocked and offer a different direction. If a reference is weak or a result looks off-target, say so and ask for an upload or redirect rather than silently iterating.
6. **Plan multi-step workflows.** Sequence the steps; only parallelize generations that belong to the same step.
7. **Review at the end.** Confirm the generations you intended actually executed and match what was asked.
8. **Don't assume tool behavior.** Don't invent tool parameters, return values, or environment capabilities that aren't actually provided; verify rather than guess.
## Choosing the Tool
| Situation | Tool |
|-----------|------|
| New image, no source image | `image_gen` |
| Edit, restyle, recolor, add, remove, or extend an existing image | `image_edit` |
| Iterate on a previous result while keeping composition | `image_edit` |
| Named real person or group | `image_edit` with a real reference after a web search |
| Generic, invented, or non-factual subject from scratch | `image_gen` |
Rule of thumb: **no source image -> `image_gen`; source image -> `image_edit`.**
## `image_gen`
Generates a new image from a text prompt.
Inputs:
- `prompt` (required) - full description of the desired image.
- `aspect_ratio` - e.g. `1:1`, `16:9`, `9:16`, `4:3`, `3:4`, or `auto`.
Use for generic or invented subjects, or to create a base image you'll edit later. Not for named real people; see Reference-first for real people.
To produce multiple variations, make multiple `image_gen` calls with distinct prompts. The tool does not expose `n` or `count` parameters.
## `image_edit`
Transforms an existing image according to a prompt.
Inputs:
- `prompt` (required) - describe the desired transformation, and note what should stay the same.
- `image` (required) - one or more source/reference images as filesystem paths or `data:image/...;base64,...` URLs. Prefer a single clean reference for reliable results.
- `aspect_ratio` - optional; used for multi-image edits. Single-image edits preserve the input image aspect ratio.
Use to restyle, recolor, add or remove elements, preserve likeness, transfer style, remix, or iterate on a generated result.
To produce multiple variations, make multiple `image_edit` calls. The tool does not expose `n` or `count` parameters.
## Writing Strong Prompts
Describe, roughly in this order: **subject -> action/pose -> setting -> style -> composition -> lighting/mood -> key details.**
- Be specific and concrete; lead with the most important elements.
- State what to include rather than what to exclude.
- Use one coherent scene per prompt.
- Match `aspect_ratio` to the use case when using `image_gen`: `9:16` for phone/story, `16:9` for banner/video frame, `1:1` for avatar/icon.
## Real People and References
1. Search the web first to confirm identity, role, relationship, or event, even when it seems obvious.
2. Use a single strong reference with `image_edit`. A user-uploaded photo is best; otherwise use a high-quality found reference and cite the source. `image_edit` can take more than one reference, but one clean reference is more reliable.
3. If no suitable reference exists, ask the user to upload one rather than generating from a weak base.
## Video
> The video tools below may not exist - verify they're available before calling them; if they're not, the user cannot do video gen with Imagine.
Video starts from an image - there is no text-to-video tool. Default to `image_to_video`.
**Think in shots.** Build video as a planned sequence of short shots, not one long take:
1. **Plan the story as shots** - break the idea into distinct shots, one beat each.
2. **Favor frequent, short shots** - prefer more 6s shots over fewer long ones; more cuts keep it dynamic and interesting.
3. **Create each shot's source image** with `image_gen` (or a multi-image `image_edit` when a shot must combine references), keeping characters and settings consistent (Core Principle 4).
4. **Animate each shot with `image_to_video`** - the source becomes frame 1.
Use `reference_to_video` only if the user asks for it or a shot genuinely needs multiple references - and even then, prefer composing those references with a multi-image `image_edit` and animating the result with `image_to_video`.
Key behaviors:
- **Prompt-craft:** one short, vivid moment in present tense with a clear camera movement, in 1-2 sentences.
- **Minimal but interesting:** keep each shot to one clear subject and a single, simple motion or camera move. Avoid complex or multi-action animation (models handle it poorly); make the shot interesting through composition, lighting, and a strong moment, not busy motion.
- **Complex source image?** An intricate frame (busy geometry, fine detail, heavy reflections) warps when animated. If you must use it, keep the subject fixed and move only the camera (slow push-in, orbit, or parallax), or break it into tighter, simpler shots. For new shots, generate a simpler, animation-friendly base image up front instead of animating a busy one.
- **`image_to_video` animates from frame 1**, so stage the intended first frame with `image_gen`/`image_edit` first.
- **Aspect ratio:** set it on the source image (`image_gen` `aspect_ratio`); don't re-crop an existing video.
- **Duration:** 6s or 10s only (prefer 6s shots); round to the nearest.
- **Real people:** reference-first - drive the video from a verified reference image; never animate a named person without one.
- Don't loop the same clip unless asked.
**Assemble shots with FFmpeg** using stream copy so there's no quality loss: `ffmpeg -f concat ... -c copy` - never re-encode. Keep every shot at the same resolution and frame rate so the copy works.
@@ -917,15 +917,6 @@ pub async fn run_leader(
} }
ConfigUpdate::AuthCleared => { ConfigUpdate::AuthCleared => {
auth_manager_for_config.clear_in_memory(); auth_manager_for_config.clear_in_memory();
let line = internal_reload_request_line(
"config-auth-cleared",
"x.ai/internal/auth_cleared",
serde_json::json!({}),
);
let mut tx = acp_tx_for_config.lock().await;
if let Err(e) = tx.write_all(line.as_bytes()).await {
warn!(error = %e, "failed to inject auth-cleared cleanup into ACP stream");
}
models_manager_for_config.on_auth_changed().await; models_manager_for_config.on_auth_changed().await;
kigi_log::unified_log::warn( kigi_log::unified_log::warn(
"auth cleared from disk", "auth cleared from disk",
@@ -1144,14 +1135,6 @@ mod tests {
); );
let msg: serde_json::Value = serde_json::from_str(line.trim_end()).unwrap(); let msg: serde_json::Value = serde_json::from_str(line.trim_end()).unwrap();
assert_eq!(msg["params"]["cwd"], "/repo/x"); assert_eq!(msg["params"]["cwd"], "/repo/x");
let line = internal_reload_request_line(
"config-auth-cleared",
"x.ai/internal/auth_cleared",
serde_json::json!({}),
);
let msg: serde_json::Value = serde_json::from_str(line.trim_end()).unwrap();
assert_eq!(msg["method"], "_x.ai/internal/auth_cleared");
} }
#[tokio::test] #[tokio::test]
@@ -301,9 +301,6 @@ pub struct Requirements {
pub tool_search: Constrained<bool>, pub tool_search: Constrained<bool>,
pub web_fetch: Constrained<bool>, pub web_fetch: Constrained<bool>,
pub ask_user_question: Constrained<bool>, pub ask_user_question: Constrained<bool>,
pub image_gen: Constrained<bool>,
pub image_edit: Constrained<bool>,
pub video_gen: Constrained<bool>,
pub write_file: Constrained<bool>, pub write_file: Constrained<bool>,
/// Voice dictation (STT). Pin via requirements/managed `[features] voice_mode`. /// Voice dictation (STT). Pin via requirements/managed `[features] voice_mode`.
pub voice_mode: Constrained<bool>, pub voice_mode: Constrained<bool>,
@@ -1075,8 +1072,6 @@ pub struct Config {
pub memory: crate::config::MemoryConfig, pub memory: crate::config::MemoryConfig,
#[serde(default, skip_serializing)] #[serde(default, skip_serializing)]
pub compaction: CompactionConfig, pub compaction: CompactionConfig,
#[serde(default, skip_serializing)]
pub managed_mcps: crate::config::ManagedMcpsConfig,
/// `[auth]` alias — consumed by `expand_auth_alias` before serde. /// `[auth]` alias — consumed by `expand_auth_alias` before serde.
/// Typed as `KimiCodeConfig` (same schema) so sub-field typos are caught. /// Typed as `KimiCodeConfig` (same schema) so sub-field typos are caught.
#[serde(default, skip_serializing)] #[serde(default, skip_serializing)]
@@ -1206,18 +1201,6 @@ pub struct Config {
/// Resolved by [`crate::config::ToolsConfig::resolve`]. /// Resolved by [`crate::config::ToolsConfig::resolve`].
#[serde(skip)] #[serde(skip)]
pub respect_gitignore: bool, pub respect_gitignore: bool,
/// When `true`, `MvpAgent::prepare_video_gen_config` returns
/// `VideoGenConfig::Disabled`, dropping `video_gen` (and any
/// future ZDR-incompatible tools) from the model's tool set.
/// Resolved by [`crate::config::ToolsConfig::resolve`].
#[serde(skip)]
pub disable_zdr_incompatible_tools: bool,
/// S3 config for ZDR video output (presigned upload to team bucket).
/// Only used when `disable_zdr_incompatible_tools` is `true` and the
/// config is valid. Resolved by [`crate::config::ToolsConfig::resolve`].
#[serde(skip)]
pub zdr_video_output_s3:
Option<kigi_tools::implementations::grok_build::video_gen::ZdrVideoOutputS3Config>,
/// Whether to enrich path-not-found errors with CWD reminders, /// Whether to enrich path-not-found errors with CWD reminders,
/// "dropped repo folder" correction, and similar-name suggestions. /// "dropped repo folder" correction, and similar-name suggestions.
/// Default `false`. Enabled via remote settings. /// Default `false`. Enabled via remote settings.
@@ -1225,13 +1208,6 @@ pub struct Config {
/// which sessions had path-not-found hints active. /// which sessions had path-not-found hints active.
#[serde(default)] #[serde(default)]
pub path_not_found_hints: bool, pub path_not_found_hints: bool,
/// Whether to fetch managed MCP configs from the managed connectors service at startup.
/// Resolved by [`crate::config::ManagedMcpsConfig::resolve`]: env var >
/// config.toml > remote settings > default (off in headless, on in interactive).
#[serde(skip)]
pub managed_mcps_enabled: bool,
#[serde(skip)]
pub managed_mcp_gateway_tools_enabled: bool,
/// Whether auto-wake is enabled: when a background task or subagent /// Whether auto-wake is enabled: when a background task or subagent
/// completes, immediately inject a synthetic prompt instead of waiting /// completes, immediately inject a synthetic prompt instead of waiting
/// for the idle-gated notification drain. /// for the idle-gated notification drain.
@@ -1446,7 +1422,6 @@ impl Default for Config {
subagents: crate::config::SubagentsConfig::default(), subagents: crate::config::SubagentsConfig::default(),
memory: crate::config::MemoryConfig::default(), memory: crate::config::MemoryConfig::default(),
compaction: CompactionConfig::default(), compaction: CompactionConfig::default(),
managed_mcps: crate::config::ManagedMcpsConfig::default(),
auth: None, auth: None,
desktop: None, desktop: None,
tips: None, tips: None,
@@ -1476,15 +1451,11 @@ impl Default for Config {
todo_gate: false, todo_gate: false,
laziness_debug_log: None, laziness_debug_log: None,
respect_gitignore: false, respect_gitignore: false,
disable_zdr_incompatible_tools: false,
zdr_video_output_s3: None,
path_not_found_hints: false, path_not_found_hints: false,
cli_experimental_memory: false, cli_experimental_memory: false,
cli_no_memory: false, cli_no_memory: false,
cli_subagents: None, cli_subagents: None,
memory_config: None, memory_config: None,
managed_mcps_enabled: true,
managed_mcp_gateway_tools_enabled: false,
auto_wake_enabled: true, auto_wake_enabled: true,
compat_resolved: CompatConfig::default(), compat_resolved: CompatConfig::default(),
requirements: Requirements::default(), requirements: Requirements::default(),
@@ -1603,8 +1574,6 @@ impl Config {
/// Call immediately after `new_from_toml_cfg()`. Fields resolved: /// Call immediately after `new_from_toml_cfg()`. Fields resolved:
/// - subagents (6 fields) via `SubagentsConfig::resolve` /// - subagents (6 fields) via `SubagentsConfig::resolve`
/// - respect_gitignore via `ToolsConfig::resolve` /// - respect_gitignore via `ToolsConfig::resolve`
/// - disable_zdr_incompatible_tools via `ToolsConfig::resolve`
/// - managed_mcps_enabled via `ManagedMcpsConfig::resolve`
/// - session_summary_model / image_description_model / /// - session_summary_model / image_description_model /
/// prompt_suggest_model_pin via `ModelOverrideConfig::resolve` /// prompt_suggest_model_pin via `ModelOverrideConfig::resolve`
/// - memory_config via `MemoryConfig::resolve` /// - memory_config via `MemoryConfig::resolve`
@@ -1624,15 +1593,6 @@ impl Config {
Some(pinned) => pinned, Some(pinned) => pinned,
None => tools.respect_gitignore, None => tools.respect_gitignore,
}; };
self.disable_zdr_incompatible_tools = tools.disable_zdr_incompatible_tools;
self.zdr_video_output_s3 = tools.zdr_video_output_s3;
let mcps = crate::config::ManagedMcpsConfig::resolve(
ctx.raw_config,
ctx.remote_settings,
ctx.is_headless,
);
self.managed_mcps_enabled = mcps.enabled;
self.managed_mcp_gateway_tools_enabled = mcps.gateway_tools_enabled;
let models = crate::config::ModelOverrideConfig::resolve( let models = crate::config::ModelOverrideConfig::resolve(
ctx.cli_session_summary_model, ctx.cli_session_summary_model,
ctx.raw_config, ctx.raw_config,
@@ -1867,51 +1827,6 @@ impl Config {
.default(true) .default(true)
.resolve() .resolve()
} }
/// `image_gen` tool gate. Default on; gated only by the `KIGI_IMAGE_GEN`
/// env var and managed-config requirement pin.
pub(crate) fn resolve_image_gen(&self) -> Resolved<bool> {
BoolFlag::env("KIGI_IMAGE_GEN")
.requirement(self.requirements.image_gen.pinned())
.default(true)
.resolve()
}
/// `image_edit` tool gate.
///
/// The remote settings `imagine_tools_disabled` denylist is authoritative:
/// when it lists `image_edit`, the tool is force-removed and local
/// env/config can't re-enable it. A managed requirement pin still outranks
/// it; otherwise the tool defaults on and is overridable via
/// `KIGI_IMAGE_EDIT`.
pub(crate) fn resolve_image_edit(&self) -> Resolved<bool> {
use kigi_tools::implementations::grok_build::IMAGE_EDIT_TOOL_NAME;
if let Some(pinned) = self.requirements.image_edit.pinned() {
return Resolved::new(pinned, ConfigSource::Requirement);
}
if self
.remote_settings
.as_ref()
.is_some_and(|s| s.imagine_tool_disabled(IMAGE_EDIT_TOOL_NAME))
{
return Resolved::new(false, ConfigSource::Remote);
}
BoolFlag::env("KIGI_IMAGE_EDIT").default(true).resolve()
}
/// Optional Imagine model override for `image_gen`. When set (non-empty),
/// `image_gen` calls this model slug instead of the default quality model.
/// Precedence: env `KIGI_IMAGE_GEN_MODEL_OVERRIDE` > `[features]
/// image_gen_model_override` config > remote settings `image_gen_model_override`.
/// `None` → default model (`grok-imagine-image-quality`).
pub(crate) fn resolve_image_gen_model_override(&self) -> Option<String> {
resolve_string_flag(
None,
"KIGI_IMAGE_GEN_MODEL_OVERRIDE",
self.features.image_gen_model_override.as_deref(),
self.remote_settings
.as_ref()
.and_then(|s| s.image_gen_model_override.as_deref()),
)
.map(|r| r.value)
}
/// Goal mode (`/goal`) master switch. Default ON: deployments that can't /// Goal mode (`/goal`) master switch. Default ON: deployments that can't
/// reach cli-chat-proxy `/v1/settings` (custom `models_base_url`, external /// reach cli-chat-proxy `/v1/settings` (custom `models_base_url`, external
/// `auth_provider_command`, air-gapped proxies) never receive the /// `auth_provider_command`, air-gapped proxies) never receive the
@@ -3646,13 +3561,6 @@ pub struct Features {
/// compaction. `None` = defer to remote settings / env / default (`false`). /// compaction. `None` = defer to remote settings / env / default (`false`).
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub two_pass_compaction: Option<bool>, pub two_pass_compaction: Option<bool>,
/// Video generation tool. `None` = defer to remote settings / env / default (false).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub video_gen: Option<bool>,
/// `image_gen` Imagine model override. `None`/empty = defer to remote settings
/// (`image_gen_model_override`) / env / default (`grok-imagine-image-quality`).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub image_gen_model_override: Option<String>,
/// Write file tool. `None` = defer to remote settings / env / default (true). /// Write file tool. `None` = defer to remote settings / env / default (true).
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub write_file: Option<bool>, pub write_file: Option<bool>,
@@ -4472,34 +4380,6 @@ reasoning_effort = "low"
assert_eq!(cfg.toolset.bash.timeout_secs, Some(30.5)); assert_eq!(cfg.toolset.bash.timeout_secs, Some(30.5));
} }
#[test] #[test]
fn resolve_runtime_fields_propagates_disable_zdr_incompatible_tools() {
fn ctx(raw: &toml::Value) -> RuntimeResolutionContext<'_> {
RuntimeResolutionContext {
raw_config: raw,
remote_settings: None,
cwd: None,
is_headless: false,
cli_subagents: None,
cli_session_summary_model: None,
cli_experimental_memory: false,
cli_no_memory: false,
disable_web_search: false,
todo_gate: false,
laziness_debug_log: None,
storage_mode: None,
}
}
let empty: toml::Value = toml::Value::Table(toml::map::Map::new());
let mut cfg = Config::new_from_toml_cfg(&empty).unwrap();
cfg.resolve_runtime_fields(&ctx(&empty));
assert!(!cfg.disable_zdr_incompatible_tools);
let zdr: toml::Value =
toml::from_str("[tools]\ndisable_zdr_incompatible_tools = true").unwrap();
let mut cfg = Config::new_from_toml_cfg(&zdr).unwrap();
cfg.resolve_runtime_fields(&ctx(&zdr));
assert!(cfg.disable_zdr_incompatible_tools);
}
#[test]
fn resolve_runtime_fields_propagates_disable_web_search() { fn resolve_runtime_fields_propagates_disable_web_search() {
fn ctx(raw: &toml::Value, disable_web_search: bool) -> RuntimeResolutionContext<'_> { fn ctx(raw: &toml::Value, disable_web_search: bool) -> RuntimeResolutionContext<'_> {
RuntimeResolutionContext { RuntimeResolutionContext {
@@ -7215,51 +7095,6 @@ reasoning_effort = "low"
assert_eq!(r.source, ConfigSource::Remote); assert_eq!(r.source, ConfigSource::Remote);
assert!(!r.value); assert!(!r.value);
} }
#[test]
#[serial]
fn resolve_image_gen_model_override_remote_settings_or_config() {
unsafe { std::env::remove_var("KIGI_IMAGE_GEN_MODEL_OVERRIDE") };
let with = |config: Option<&str>, gb: Option<&str>| Config {
features: Features {
image_gen_model_override: config.map(String::from),
..Default::default()
},
remote_settings: Some(crate::util::config::RemoteSettings {
image_gen_model_override: gb.map(String::from),
..Default::default()
}),
..Default::default()
};
assert_eq!(Config::default().resolve_image_gen_model_override(), None);
assert_eq!(
with(None, Some("grok-imagine-image")).resolve_image_gen_model_override(),
Some("grok-imagine-image".to_owned())
);
assert_eq!(
with(Some("grok-imagine-image-pro"), Some("grok-imagine-image"))
.resolve_image_gen_model_override(),
Some("grok-imagine-image-pro".to_owned())
);
}
#[test]
#[serial]
fn imagine_tools_disabled_gates_image_edit() {
unsafe { std::env::remove_var("KIGI_IMAGE_EDIT") };
let with_list = |tools: Vec<&str>| Config {
remote_settings: Some(crate::util::config::RemoteSettings {
imagine_tools_disabled: Some(tools.into_iter().map(String::from).collect()),
..Default::default()
}),
..Default::default()
};
unsafe { std::env::set_var("KIGI_IMAGE_EDIT", "1") };
let off = with_list(vec!["image_edit"]).resolve_image_edit();
assert!(!off.value);
assert_eq!(off.source, ConfigSource::Remote);
unsafe { std::env::remove_var("KIGI_IMAGE_EDIT") };
assert!(with_list(vec!["image_to_video"]).resolve_image_edit().value);
assert!(Config::default().resolve_image_edit().value);
}
/// Clear every env var the goal/companion resolvers read so tests /// Clear every env var the goal/companion resolvers read so tests
/// start from a known baseline regardless of run order. /// start from a known baseline regardless of run order.
fn clear_goal_envs() { fn clear_goal_envs() {
@@ -8103,8 +7938,6 @@ agent_type = "cursor"
secret = "value" secret = "value"
[worktree_pool] [worktree_pool]
pool_size = 4 pool_size = 4
[managed_mcps]
enabled = true
[mcp_servers.test] [mcp_servers.test]
url = "https://mcp.test.com" url = "https://mcp.test.com"
[toolset.bash] [toolset.bash]
@@ -8239,12 +8072,6 @@ agent_type = "cursor"
std::env::remove_var("KIGI_CLAUDE_AGENTS_ENABLED"); std::env::remove_var("KIGI_CLAUDE_AGENTS_ENABLED");
} }
} }
fn clear_managed_mcp_env_vars() {
unsafe {
std::env::remove_var("KIGI_MANAGED_MCPS_ENABLED");
std::env::remove_var("KIGI_MANAGED_MCP_GATEWAY_TOOLS_ENABLED");
}
}
fn isolate_compat_env() -> Vec<EnvGuard> { fn isolate_compat_env() -> Vec<EnvGuard> {
COMPAT_CELLS COMPAT_CELLS
.into_iter() .into_iter()
@@ -8531,7 +8358,6 @@ hooks = true
#[serial] #[serial]
fn resolve_runtime_fields_interactive_defaults() { fn resolve_runtime_fields_interactive_defaults() {
clear_runtime_env_vars(); clear_runtime_env_vars();
clear_managed_mcp_env_vars();
let raw = empty_config(); let raw = empty_config();
let mut cfg = Config::new_from_toml_cfg(&raw).unwrap(); let mut cfg = Config::new_from_toml_cfg(&raw).unwrap();
cfg.resolve_runtime_fields(&RuntimeResolutionContext { cfg.resolve_runtime_fields(&RuntimeResolutionContext {
@@ -8550,8 +8376,6 @@ hooks = true
}); });
assert!(cfg.subagents_enabled); assert!(cfg.subagents_enabled);
assert!(!cfg.respect_gitignore); assert!(!cfg.respect_gitignore);
assert!(cfg.managed_mcps_enabled);
assert!(!cfg.managed_mcp_gateway_tools_enabled);
assert_eq!( assert_eq!(
cfg.session_summary_model, cfg.session_summary_model,
Some(crate::models::default_session_summary_model().to_owned()) Some(crate::models::default_session_summary_model().to_owned())
@@ -8560,60 +8384,6 @@ hooks = true
} }
#[test] #[test]
#[serial] #[serial]
fn resolve_runtime_fields_headless_defaults() {
clear_runtime_env_vars();
clear_managed_mcp_env_vars();
let raw = empty_config();
let mut cfg = Config::new_from_toml_cfg(&raw).unwrap();
cfg.resolve_runtime_fields(&RuntimeResolutionContext {
raw_config: &raw,
remote_settings: None,
cwd: None,
is_headless: true,
cli_subagents: None,
cli_session_summary_model: None,
cli_experimental_memory: false,
cli_no_memory: false,
disable_web_search: false,
todo_gate: false,
laziness_debug_log: None,
storage_mode: None,
});
assert!(
!cfg.managed_mcps_enabled,
"headless should default managed_mcps to false"
);
assert!(!cfg.managed_mcp_gateway_tools_enabled);
}
#[test]
#[serial]
fn resolve_runtime_fields_managed_gateway_tools_from_remote() {
clear_runtime_env_vars();
clear_managed_mcp_env_vars();
let raw = empty_config();
let remote = crate::util::config::RemoteSettings {
managed_mcp_gateway_tools_enabled: Some(true),
..Default::default()
};
let mut cfg = Config::new_from_toml_cfg(&raw).unwrap();
cfg.resolve_runtime_fields(&RuntimeResolutionContext {
raw_config: &raw,
remote_settings: Some(&remote),
cwd: None,
is_headless: false,
cli_subagents: None,
cli_session_summary_model: None,
cli_experimental_memory: false,
cli_no_memory: false,
disable_web_search: false,
todo_gate: false,
laziness_debug_log: None,
storage_mode: None,
});
assert!(cfg.managed_mcp_gateway_tools_enabled);
}
#[test]
#[serial]
fn resolve_runtime_fields_subagents_from_config() { fn resolve_runtime_fields_subagents_from_config() {
clear_runtime_env_vars(); clear_runtime_env_vars();
let raw: toml::Value = toml::from_str("[subagents]\nenabled = true").unwrap(); let raw: toml::Value = toml::from_str("[subagents]\nenabled = true").unwrap();
@@ -8751,12 +8521,10 @@ hooks = true
cfg.resolve_runtime_fields(&ctx); cfg.resolve_runtime_fields(&ctx);
let first_subagents = cfg.subagents_enabled; let first_subagents = cfg.subagents_enabled;
let first_gitignore = cfg.respect_gitignore; let first_gitignore = cfg.respect_gitignore;
let first_mcps = cfg.managed_mcps_enabled;
let first_ss = cfg.session_summary_model.clone(); let first_ss = cfg.session_summary_model.clone();
cfg.resolve_runtime_fields(&ctx); cfg.resolve_runtime_fields(&ctx);
assert_eq!(cfg.subagents_enabled, first_subagents); assert_eq!(cfg.subagents_enabled, first_subagents);
assert_eq!(cfg.respect_gitignore, first_gitignore); assert_eq!(cfg.respect_gitignore, first_gitignore);
assert_eq!(cfg.managed_mcps_enabled, first_mcps);
assert_eq!(cfg.session_summary_model, first_ss); assert_eq!(cfg.session_summary_model, first_ss);
} }
#[test] #[test]
@@ -270,13 +270,7 @@ impl acp::Agent for MvpAgent {
let current_working_directory = self.launch_cwd.clone(); let current_working_directory = self.launch_cwd.clone();
let hostname = gethostname::gethostname(); let hostname = gethostname::gethostname();
let mcp_servers: Vec<crate::extensions::mcp::McpServerEntry> = Vec::new(); let mcp_servers: Vec<crate::extensions::mcp::McpServerEntry> = Vec::new();
let fetch_managed_mcps = self.cfg.borrow().managed_mcps_enabled self.spawn_initialize_launch_mcp_setup();
&& self.can_fetch_managed_mcps();
if self.cfg.borrow().managed_mcps_enabled && !fetch_managed_mcps {
tracing::info!("Managed MCP fetch: DISABLED");
}
self.spawn_initialize_launch_mcp_setup(fetch_managed_mcps);
self.spawn_managed_gateway_tool_catalog_fetch();
let init_model_state = self.model_state(None); let init_model_state = self.model_state(None);
Ok( Ok(
acp::InitializeResponse::new(acp::ProtocolVersion::V1) acp::InitializeResponse::new(acp::ProtocolVersion::V1)
@@ -545,9 +539,7 @@ impl acp::Agent for MvpAgent {
let remote_settings = self.cfg.borrow().remote_settings.clone(); let remote_settings = self.cfg.borrow().remote_settings.clone();
folder_trust::resolve_and_record(cwd.as_path(), remote_settings.as_ref(), false); folder_trust::resolve_and_record(cwd.as_path(), remote_settings.as_ref(), false);
let initial_client_mcp_servers = arguments.mcp_servers.clone(); let initial_client_mcp_servers = arguments.mcp_servers.clone();
let (mcp_servers, managed_mcp_expires_at) = self let mcp_servers = self.resolve_mcp_servers(arguments.mcp_servers, cwd.as_path());
.resolve_mcp_servers(arguments.mcp_servers, cwd.as_path())
.await;
let mcp_meta_config_map = parse_mcp_meta_config(arguments.meta.as_ref()); let mcp_meta_config_map = parse_mcp_meta_config(arguments.meta.as_ref());
let client_session_id = arguments let client_session_id = arguments
.meta .meta
@@ -756,7 +748,6 @@ impl acp::Agent for MvpAgent {
persisted_goal_mode: None, persisted_goal_mode: None,
persisted_announcement_state: None, persisted_announcement_state: None,
session_meta: arguments.meta.as_ref(), session_meta: arguments.meta.as_ref(),
managed_mcp_expires_at,
model_agent_type: model_agent_type.as_deref(), model_agent_type: model_agent_type.as_deref(),
session_model_id, session_model_id,
session_yolo_mode, session_yolo_mode,
@@ -902,9 +893,7 @@ impl acp::Agent for MvpAgent {
let remote_settings = self.cfg.borrow().remote_settings.clone(); let remote_settings = self.cfg.borrow().remote_settings.clone();
folder_trust::resolve_and_record(cwd.as_path(), remote_settings.as_ref(), false); folder_trust::resolve_and_record(cwd.as_path(), remote_settings.as_ref(), false);
let initial_client_mcp_servers = client_mcp_servers.clone(); let initial_client_mcp_servers = client_mcp_servers.clone();
let (mcp_servers, managed_mcp_expires_at) = self let mcp_servers = self.resolve_mcp_servers(client_mcp_servers, cwd.as_path());
.resolve_mcp_servers(client_mcp_servers, cwd.as_path())
.await;
let mcp_meta_config_map = parse_mcp_meta_config(request_meta.as_ref()); let mcp_meta_config_map = parse_mcp_meta_config(request_meta.as_ref());
let mut load_timer = crate::instrumentation_timer!("session.load_session"); let mut load_timer = crate::instrumentation_timer!("session.load_session");
load_timer.with_field("session_id", session_id.0.as_ref()); load_timer.with_field("session_id", session_id.0.as_ref());
@@ -1212,7 +1201,6 @@ impl acp::Agent for MvpAgent {
persisted_goal_mode: _persisted_goal_mode, persisted_goal_mode: _persisted_goal_mode,
persisted_announcement_state, persisted_announcement_state,
session_meta: request_meta.as_ref(), session_meta: request_meta.as_ref(),
managed_mcp_expires_at,
model_agent_type: persisted_agent_name.as_deref(), model_agent_type: persisted_agent_name.as_deref(),
session_model_id: summary.current_model_id.clone(), session_model_id: summary.current_model_id.clone(),
session_yolo_mode, session_yolo_mode,
@@ -2083,7 +2071,7 @@ impl acp::Agent for MvpAgent {
| "x.ai/internal/reload_all_mcp_servers" | "x.ai/internal/reload_all_mcp_servers"
| "x.ai/internal/reload_project_mcp_servers" | "x.ai/internal/reload_skills" | "x.ai/internal/reload_project_mcp_servers" | "x.ai/internal/reload_skills"
| "x.ai/internal/reload_models" | "x.ai/internal/reload_models_cache" | "x.ai/internal/reload_models" | "x.ai/internal/reload_models_cache"
| "x.ai/internal/auth_cleared" | "x.ai/plugins/reload" | "x.ai/plugins/reload"
| "x.ai/commands/list" => { | "x.ai/commands/list" => {
crate::extensions::session_admin::handle(self, &args).await crate::extensions::session_admin::handle(self, &args).await
} }
@@ -84,123 +84,6 @@ impl MvpAgent {
} }
}) })
} }
fn has_managed_mcp_auth(&self) -> bool {
self.auth_manager
.current_or_expired()
.is_some_and(|a| a.is_session_auth())
}
/// Requires feature flag AND xAI authentication (OIDC or legacy WebLogin).
pub(super) fn can_fetch_managed_mcps(&self) -> bool {
let cfg = self.cfg.borrow();
cfg.managed_mcps_enabled && !cfg.managed_mcp_gateway_tools_enabled
&& self.has_managed_mcp_auth()
}
fn can_fetch_managed_mcp_gateway_tools(&self) -> bool {
self.cfg.borrow().managed_mcp_gateway_tools_enabled
&& self.has_managed_mcp_auth()
}
pub async fn get_managed_mcp_configs(
&self,
) -> Vec<crate::session::managed_mcp::ManagedMcpConfig> {
if !self.can_fetch_managed_mcps() {
return vec![];
}
let proxy_url = self.cfg.borrow().endpoints.proxy_url();
crate::session::managed_mcp::fetch_managed_mcp_configs(
&self.managed_mcp_cache,
&proxy_url,
&self.auth_manager,
)
.await
}
pub async fn get_managed_mcp_gateway_tool_catalog(
&self,
) -> Option<crate::session::managed_mcp::GatewayToolCatalog> {
if !self.can_fetch_managed_mcp_gateway_tools() {
self.managed_mcp_cache.lock().await.disable_gateway_tools();
return None;
}
self.managed_mcp_cache.lock().await.enable_gateway_tools();
let proxy_url = self.cfg.borrow().endpoints.proxy_url();
let auth_key = self
.auth_manager
.get_valid_token()
.await
.ok()
.or_else(|| self.auth_manager.current_or_expired().map(|a| a.key));
crate::session::managed_mcp::get_or_fetch_gateway_tool_catalog(
&self.managed_mcp_cache,
&proxy_url,
auth_key.as_deref(),
)
.await
}
pub fn managed_mcp_cache(
&self,
) -> &crate::session::managed_mcp::ManagedMcpStateHandle {
&self.managed_mcp_cache
}
pub(crate) fn disable_managed_gateway_tools_and_refresh_sessions(&self) {
self.disable_managed_gateway_tools_and_refresh_sessions_with_txs(
self.sessions.borrow().values().map(|handle| handle.cmd_tx.clone()).collect(),
);
}
fn disable_managed_gateway_tools_and_refresh_sessions_with_txs(
&self,
session_txs: Vec<tokio::sync::mpsc::UnboundedSender<SessionCommand>>,
) {
let cache = self.managed_mcp_cache.clone();
tokio::task::spawn_local(async move {
cache.lock().await.disable_gateway_tools();
for tx in session_txs {
let _ = tx.send(SessionCommand::RefreshMcpSearchIndex);
}
});
}
pub(crate) fn spawn_managed_gateway_tool_catalog_fetch(&self) {
let session_txs: Vec<_> = self
.sessions
.borrow()
.values()
.map(|handle| handle.cmd_tx.clone())
.collect();
if !self.can_fetch_managed_mcp_gateway_tools() {
self.disable_managed_gateway_tools_and_refresh_sessions_with_txs(
session_txs,
);
return;
}
let cache = self.managed_mcp_cache.clone();
let proxy_url = self.cfg.borrow().endpoints.proxy_url();
let auth_manager = self.auth_manager.clone();
tokio::task::spawn_local(async move {
let auth_key = auth_manager
.get_valid_token()
.await
.ok()
.or_else(|| auth_manager.current_or_expired().map(|a| a.key));
if !auth_manager
.current_or_expired()
.is_some_and(|a| a.is_session_auth())
{
cache.lock().await.disable_gateway_tools();
for tx in session_txs {
let _ = tx.send(SessionCommand::RefreshMcpSearchIndex);
}
return;
}
cache.lock().await.enable_gateway_tools();
crate::session::managed_mcp::get_or_fetch_gateway_tool_catalog(
&cache,
&proxy_url,
auth_key.as_deref(),
)
.await;
for tx in session_txs {
let _ = tx.send(SessionCommand::RefreshMcpSearchIndex);
}
});
}
/// Resolve the launch dir's project-scope trust verdict ONCE and return it /// Resolve the launch dir's project-scope trust verdict ONCE and return it
/// with its path. /// with its path.
/// ///
@@ -226,15 +109,12 @@ impl MvpAgent {
/// Resolve folder trust and load launch-dir MCP configs after `initialize` /// Resolve folder trust and load launch-dir MCP configs after `initialize`
/// returns. The walks are synchronous and expensive in large monorepos; they /// returns. The walks are synchronous and expensive in large monorepos; they
/// must not block the ACP response (grok-desktop sends `initialize` immediately). /// must not block the ACP response (grok-desktop sends `initialize` immediately).
pub(super) fn spawn_initialize_launch_mcp_setup(&self, fetch_managed_mcps: bool) { pub(super) fn spawn_initialize_launch_mcp_setup(&self) {
let cwd = self.launch_cwd.clone(); let cwd = self.launch_cwd.clone();
let compat = self.cfg.borrow().compat_resolved; let compat = self.cfg.borrow().compat_resolved;
let remote_settings = self.cfg.borrow().remote_settings.clone(); let remote_settings = self.cfg.borrow().remote_settings.clone();
let gateway = self.gateway.clone(); let gateway = self.gateway.clone();
let agent_mcp_state = self.agent_mcp_state.clone(); let agent_mcp_state = self.agent_mcp_state.clone();
let managed_mcp_cache = self.managed_mcp_cache.clone();
let proxy_url = self.cfg.borrow().endpoints.proxy_url();
let auth_manager = self.auth_manager.clone();
tokio::task::spawn_local(async move { tokio::task::spawn_local(async move {
let local_mcp_servers = match tokio::task::spawn_blocking(move || { let local_mcp_servers = match tokio::task::spawn_blocking(move || {
let local = crate::util::config::load_mcp_servers(&cwd, &compat); let local = crate::util::config::load_mcp_servers(&cwd, &compat);
@@ -258,27 +138,9 @@ impl MvpAgent {
} }
crate::extensions::mcp::notify_servers_updated( crate::extensions::mcp::notify_servers_updated(
&gateway, &gateway,
&[],
&local_mcp_servers, &local_mcp_servers,
) )
.await; .await;
if !fetch_managed_mcps {
return;
}
let managed = crate::session::managed_mcp::fetch_managed_mcp_configs(
&managed_mcp_cache,
&proxy_url,
&auth_manager,
)
.await;
if !managed.is_empty() {
crate::extensions::mcp::notify_servers_updated(
&gateway,
&managed,
&local_mcp_servers,
)
.await;
}
}); });
} }
pub fn agent_mcp_state( pub fn agent_mcp_state(
@@ -311,23 +173,19 @@ impl MvpAgent {
plugin_count = count, "lazily populated plugin registry snapshot" plugin_count = count, "lazily populated plugin registry snapshot"
); );
} }
/// Fetch managed configs, merge with client servers, return merged list + earliest expiry. /// Merge on-disk/plugin MCP servers with client servers.
pub(super) async fn resolve_mcp_servers( pub(super) fn resolve_mcp_servers(
&self, &self,
client_servers: Vec<acp::McpServer>, client_servers: Vec<acp::McpServer>,
cwd: &std::path::Path, cwd: &std::path::Path,
) -> (Vec<acp::McpServer>, Option<chrono::DateTime<chrono::Utc>>) { ) -> Vec<acp::McpServer> {
self.ensure_plugin_registry(); self.ensure_plugin_registry();
let managed = self.get_managed_mcp_configs().await; crate::session::managed_mcp::merge_managed_mcp_servers(
let expires_at = managed.iter().filter_map(|c| c.token_expires_at).min();
let merged = crate::session::managed_mcp::merge_managed_mcp_servers(
client_servers, client_servers,
cwd, cwd,
&managed,
self.plugin_registry_handle.snapshot().as_deref(), self.plugin_registry_handle.snapshot().as_deref(),
&self.cfg.borrow().compat_resolved, &self.cfg.borrow().compat_resolved,
); )
(merged, expires_at)
} }
/// Set the memory configuration (called from TUI after config resolution). /// Set the memory configuration (called from TUI after config resolution).
pub fn set_memory_config(&mut self, config: crate::config::MemoryConfig) { pub fn set_memory_config(&mut self, config: crate::config::MemoryConfig) {
@@ -760,39 +618,6 @@ impl MvpAgent {
); );
(id.clone(), new_config) (id.clone(), new_config)
} }
/// Build image generation config.
///
/// Both BYOK and session (OAuth) users go direct to `api_base_url`.
/// `sampling_config.api_key` carries the OAuth bearer for session users (the
/// `api_key_provider` refreshes it per request), so IC authenticates and
/// meters Imagine usage per-user.
pub(super) fn prepare_image_gen_config(
&self,
) -> kigi_tools::implementations::grok_build::image_gen::ImageGenConfig {
use kigi_tools::implementations::grok_build::image_gen::ImageGenConfig;
let sampling_config = self.sampling_config.borrow();
let Some(ref api_key) = sampling_config.api_key else {
return ImageGenConfig::Disabled;
};
let cfg = self.cfg.borrow();
let base_url = cfg.endpoints.api_base_url.clone();
let version = cfg
.client_version
.clone()
.unwrap_or_else(|| kigi_version::VERSION.to_string());
let alpha_test_key = cfg.endpoints.alpha_test_key.clone();
let mut headers = indexmap::IndexMap::new();
headers.insert("user-agent".to_string(), format!("kigi/{version}"));
ImageGenConfig::Enabled {
api_key: api_key.clone(),
base_url,
extra_headers: headers,
image_gen_enabled: cfg.resolve_image_gen().value,
image_edit_enabled: cfg.resolve_image_edit().value,
model_override: cfg.resolve_image_gen_model_override(),
tier_restricted: false,
}
}
/// Build deploy-service config. The tool talks directly to the deployer service. /// Build deploy-service config. The tool talks directly to the deployer service.
pub(super) fn prepare_app_builder_deployer_config( pub(super) fn prepare_app_builder_deployer_config(
&self, &self,
@@ -800,40 +625,6 @@ impl MvpAgent {
use kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig; use kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig;
AppBuilderDeployerConfig::Disabled AppBuilderDeployerConfig::Disabled
} }
/// Build video generation config. Video tools call the xAI API directly.
pub(super) fn prepare_video_gen_config(
&self,
) -> kigi_tools::implementations::grok_build::video_gen::VideoGenConfig {
use kigi_tools::implementations::grok_build::video_gen::VideoGenConfig;
let Some(api_key) = self.sampling_config.borrow().api_key.clone() else {
return VideoGenConfig::Disabled;
};
let cfg = self.cfg.borrow();
let zdr_video_output_s3 = cfg
.disable_zdr_incompatible_tools
.then(|| cfg.zdr_video_output_s3.clone())
.flatten()
.filter(|s3| s3.is_valid());
if cfg.disable_zdr_incompatible_tools && zdr_video_output_s3.is_none() {
tracing::info!("video_gen disabled by tools.disable_zdr_incompatible_tools");
return VideoGenConfig::Disabled;
}
let base_url = cfg.endpoints.api_base_url.clone();
let version = cfg
.client_version
.clone()
.unwrap_or_else(|| kigi_version::VERSION.to_string());
let alpha_test_key = cfg.endpoints.alpha_test_key.clone();
let mut headers = indexmap::IndexMap::new();
headers.insert("user-agent".to_string(), format!("kigi/{version}"));
VideoGenConfig::Enabled {
api_key,
base_url,
extra_headers: headers,
zdr_video_output_s3: zdr_video_output_s3.map(Box::new),
tier_restricted: false,
}
}
/// Web search config (PRD F5). The Kimi search service exists only on /// Web search config (PRD F5). The Kimi search service exists only on
/// the Kimi Code subscription channel (`POST {coding_base}/search`, /// the Kimi Code subscription channel (`POST {coding_base}/search`,
/// kimi-cli `auth/platforms.py`), so this is `Enabled` only for OAuth /// kimi-cli `auth/platforms.py`), so this is `Enabled` only for OAuth
@@ -1016,7 +807,6 @@ impl MvpAgent {
worktree_type, worktree_type,
restore_code, restore_code,
session_registry_local, session_registry_local,
managed_mcp_cache: Default::default(),
agent_mcp_state: std::sync::Arc::new( agent_mcp_state: std::sync::Arc::new(
tokio::sync::Mutex::new( tokio::sync::Mutex::new(
crate::session::mcp_servers::McpState::new(vec![]), crate::session::mcp_servers::McpState::new(vec![]),
@@ -1829,7 +1619,6 @@ impl MvpAgent {
persisted_goal_mode, persisted_goal_mode,
persisted_announcement_state, persisted_announcement_state,
session_meta, session_meta,
managed_mcp_expires_at,
model_agent_type, model_agent_type,
session_model_id, session_model_id,
session_yolo_mode, session_yolo_mode,
@@ -2258,8 +2047,6 @@ impl MvpAgent {
.and_then(|entry| entry.info.max_retries); .and_then(|entry| entry.info.max_retries);
let origin_client = self.origin_client_info_from_meta(init.meta.as_ref()); let origin_client = self.origin_client_info_from_meta(init.meta.as_ref());
let web_search_config = self.prepare_web_search_config(); let web_search_config = self.prepare_web_search_config();
let image_gen_config = self.prepare_image_gen_config();
let video_gen_config = self.prepare_video_gen_config();
let app_builder_deployer_config = self.prepare_app_builder_deployer_config(); let app_builder_deployer_config = self.prepare_app_builder_deployer_config();
let web_fetch_config = self.prepare_web_fetch_config(); let web_fetch_config = self.prepare_web_fetch_config();
let write_file_enabled = self.cfg.borrow().resolve_write_file().value; let write_file_enabled = self.cfg.borrow().resolve_write_file().value;
@@ -2309,7 +2096,6 @@ impl MvpAgent {
let cfg = self.cfg.borrow(); let cfg = self.cfg.borrow();
cfg.resolve_backend_tools().value cfg.resolve_backend_tools().value
}; };
let managed_mcp_proxy_url = self.cfg.borrow().endpoints.proxy_url();
let init_meta = self let init_meta = self
.initialize_request .initialize_request
.get() .get()
@@ -2471,9 +2257,6 @@ impl MvpAgent {
persisted_announcement_state, persisted_announcement_state,
self.memory_config.clone(), self.memory_config.clone(),
feedback_flags, feedback_flags,
self.managed_mcp_cache.clone(),
managed_mcp_expires_at,
managed_mcp_proxy_url,
session_model_id, session_model_id,
session_yolo_mode, session_yolo_mode,
session_auto_mode, session_auto_mode,
@@ -2482,8 +2265,6 @@ impl MvpAgent {
model_max_retries, model_max_retries,
web_search_config, web_search_config,
web_fetch_config, web_fetch_config,
image_gen_config,
video_gen_config,
app_builder_deployer_config, app_builder_deployer_config,
write_file_enabled, write_file_enabled,
goal_enabled, goal_enabled,
@@ -2617,7 +2398,6 @@ impl MvpAgent {
self.notify_session_cwd_for_watch(std::path::Path::new(&session_info.cwd)); self.notify_session_cwd_for_watch(std::path::Path::new(&session_info.cwd));
self.activity.register_session(&session_info.id.0, &handle); self.activity.register_session(&session_info.id.0, &handle);
self.sessions.borrow_mut().insert(session_info.id.clone(), handle); self.sessions.borrow_mut().insert(session_info.id.clone(), handle);
self.spawn_managed_gateway_tool_catalog_fetch();
let cwd_for_maintenance = session_info.cwd.clone(); let cwd_for_maintenance = session_info.cwd.clone();
tokio::spawn(async move { tokio::spawn(async move {
crate::session::prompt_history::truncate_if_needed_async(cwd_for_maintenance) crate::session::prompt_history::truncate_if_needed_async(cwd_for_maintenance)
@@ -150,10 +150,6 @@ impl MvpAgent {
let gateway = self.gateway.clone(); let gateway = self.gateway.clone();
let plugin_handle = self.plugin_registry_handle.clone(); let plugin_handle = self.plugin_registry_handle.clone();
let managed_mcp_cache = self.managed_mcp_cache.clone();
let auth_manager = self.auth_manager.clone();
let can_fetch_managed = self.can_fetch_managed_mcps();
let proxy_url = self.cfg.borrow().endpoints.proxy_url();
let compat = self.cfg.borrow().compat_resolved; let compat = self.cfg.borrow().compat_resolved;
let remote = remote.cloned(); let remote = remote.cloned();
let cwd = cwd.to_path_buf(); let cwd = cwd.to_path_buf();
@@ -258,10 +254,6 @@ impl MvpAgent {
gateway: &gateway, gateway: &gateway,
targets, targets,
plugin_handle: &plugin_handle, plugin_handle: &plugin_handle,
managed_mcp_cache: &managed_mcp_cache,
auth_manager: &auth_manager,
can_fetch_managed,
proxy_url: &proxy_url,
compat: &compat, compat: &compat,
prompt_cwd: &cwd, prompt_cwd: &cwd,
}) })
@@ -291,10 +283,6 @@ struct ReloadAfterGrant<'a> {
/// Every session sharing the granted workspace, each with its own cwd. /// Every session sharing the granted workspace, each with its own cwd.
targets: Vec<ReloadTarget>, targets: Vec<ReloadTarget>,
plugin_handle: &'a kigi_agent::plugins::SharedPluginRegistryHandle, plugin_handle: &'a kigi_agent::plugins::SharedPluginRegistryHandle,
managed_mcp_cache: &'a crate::session::managed_mcp::ManagedMcpStateHandle,
auth_manager: &'a std::sync::Arc<AuthManager>,
can_fetch_managed: bool,
proxy_url: &'a str,
compat: &'a kigi_tools::types::CompatConfig, compat: &'a kigi_tools::types::CompatConfig,
/// The prompting session's cwd — used only for the client catalog push. /// The prompting session's cwd — used only for the client catalog push.
prompt_cwd: &'a std::path::Path, prompt_cwd: &'a std::path::Path,
@@ -303,27 +291,15 @@ struct ReloadAfterGrant<'a> {
/// Reload each granted-workspace session's now-trusted project servers in place /// Reload each granted-workspace session's now-trusted project servers in place
/// (no restart), driving the canonical primitives the normal spawn/reload paths /// (no restart), driving the canonical primitives the normal spawn/reload paths
/// use — PER SESSION CWD, like `handle_reload_project_mcp_servers` / /// use — PER SESSION CWD, like `handle_reload_project_mcp_servers` /
/// `broadcast_plugin_registry_to_sessions`: `fetch_managed_mcp_configs` + /// `broadcast_plugin_registry_to_sessions`:
/// `merge_managed_mcp_servers` (`SessionCommand::UpdateMcpServers`), `build_for_cwd` /// `merge_managed_mcp_servers` (`SessionCommand::UpdateMcpServers`), `build_for_cwd`
/// (`SessionCommand::ReloadPlugins`), and `reload_hooks_impl` /// (`SessionCommand::ReloadPlugins`), and `reload_hooks_impl`
/// (`SessionCommand::ReloadHooks`), then push the refreshed MCP catalog. LSP is /// (`SessionCommand::ReloadHooks`), then push the refreshed MCP catalog. LSP is
/// spawn-baked and applies on the next session open (see module docs). Caller /// spawn-baked and applies on the next session open (see module docs). Caller
/// must have granted + recorded trust first. /// must have granted + recorded trust first.
async fn reload_project_servers_after_grant(ctx: ReloadAfterGrant<'_>) { async fn reload_project_servers_after_grant(ctx: ReloadAfterGrant<'_>) {
// Managed (gateway/Toolbox) servers must survive the re-merge; fetch them once // The plugin MCP snapshot is global, so it is fine to reuse across cwds
// (cwd-independent) via the shared helper (single-sources the auth-key dance // for the merge.
// with `MvpAgent::get_managed_mcp_configs`). The plugin MCP snapshot is also
// global, so it is fine to reuse across cwds for the merge.
let managed = if ctx.can_fetch_managed {
crate::session::managed_mcp::fetch_managed_mcp_configs(
ctx.managed_mcp_cache,
ctx.proxy_url,
ctx.auth_manager,
)
.await
} else {
vec![]
};
let plugin_snapshot = ctx.plugin_handle.snapshot(); let plugin_snapshot = ctx.plugin_handle.snapshot();
for target in ctx.targets { for target in ctx.targets {
@@ -336,7 +312,6 @@ async fn reload_project_servers_after_grant(ctx: ReloadAfterGrant<'_>) {
let merged = crate::session::managed_mcp::merge_managed_mcp_servers( let merged = crate::session::managed_mcp::merge_managed_mcp_servers(
target.initial_client_mcp_servers, target.initial_client_mcp_servers,
session_cwd, session_cwd,
&managed,
plugin_snapshot.as_deref(), plugin_snapshot.as_deref(),
ctx.compat, ctx.compat,
); );
@@ -375,7 +350,7 @@ async fn reload_project_servers_after_grant(ctx: ReloadAfterGrant<'_>) {
ctx.prompt_cwd, ctx.prompt_cwd,
crate::util::config::load_mcp_servers(ctx.prompt_cwd, ctx.compat), crate::util::config::load_mcp_servers(ctx.prompt_cwd, ctx.compat),
); );
crate::extensions::mcp::notify_servers_updated(ctx.gateway, &managed, &local).await; crate::extensions::mcp::notify_servers_updated(ctx.gateway, &local).await;
} }
#[cfg(test)] #[cfg(test)]
@@ -124,7 +124,6 @@ pub(crate) struct SessionSpawnOptions<'a> {
crate::session::announcement_state::AnnouncementState, crate::session::announcement_state::AnnouncementState,
>, >,
pub session_meta: Option<&'a acp::Meta>, pub session_meta: Option<&'a acp::Meta>,
pub managed_mcp_expires_at: Option<chrono::DateTime<chrono::Utc>>,
pub model_agent_type: Option<&'a str>, pub model_agent_type: Option<&'a str>,
pub session_model_id: acp::ModelId, pub session_model_id: acp::ModelId,
pub session_yolo_mode: bool, pub session_yolo_mode: bool,
@@ -260,7 +259,6 @@ pub(crate) fn chat_session_spawn_options<'a>(
persisted_goal_mode: None, persisted_goal_mode: None,
persisted_announcement_state: None, persisted_announcement_state: None,
session_meta, session_meta,
managed_mcp_expires_at: None,
model_agent_type, model_agent_type,
session_model_id, session_model_id,
session_yolo_mode, session_yolo_mode,
@@ -591,8 +589,6 @@ pub struct MvpAgent {
/// Local config.toml override for session registry (`[cli] session_registry`). /// Local config.toml override for session registry (`[cli] session_registry`).
/// `Some(true)` enables, `Some(false)` disables, `None` defers to remote settings. /// `Some(true)` enables, `Some(false)` disables, `None` defers to remote settings.
session_registry_local: Option<bool>, session_registry_local: Option<bool>,
/// Managed MCP configs and gateway tool catalog; lazily fetched.
managed_mcp_cache: crate::session::managed_mcp::ManagedMcpStateHandle,
/// Agent-level MCP server state. LEADER-SAFE(shared): MCP servers are /// Agent-level MCP server state. LEADER-SAFE(shared): MCP servers are
/// agent-scoped, not per-client. /// agent-scoped, not per-client.
agent_mcp_state: std::sync::Arc< agent_mcp_state: std::sync::Arc<
@@ -304,7 +304,6 @@ impl MvpAgent {
session_env, session_env,
parent_attribution_callback, parent_attribution_callback,
parent_agent_name, parent_agent_name,
parent_managed_mcp_proxy_base_url,
) = { ) = {
let sessions = self.sessions.borrow(); let sessions = self.sessions.borrow();
let ps = sessions.get(&parent_sid); let ps = sessions.get(&parent_sid);
@@ -339,7 +338,6 @@ impl MvpAgent {
.unwrap_or_else(|| std::sync::Arc::new(std::collections::HashMap::new())), .unwrap_or_else(|| std::sync::Arc::new(std::collections::HashMap::new())),
ps.and_then(|h| h.attribution_callback.clone()), ps.and_then(|h| h.attribution_callback.clone()),
ps.map(|h| h.agent_name.clone()), ps.map(|h| h.agent_name.clone()),
ps.map(|h| h.managed_mcp_proxy_base_url.clone()),
) )
}; };
let ( let (
@@ -401,8 +399,6 @@ impl MvpAgent {
gateway: self.gateway.clone(), gateway: self.gateway.clone(),
client_hooks: Default::default(), client_hooks: Default::default(),
sampling_config: self.sampling_config.borrow().clone(), sampling_config: self.sampling_config.borrow().clone(),
managed_mcp_proxy_base_url: parent_managed_mcp_proxy_base_url
.unwrap_or_else(|| self.coding_api_base_url()),
alpha_test_key: self.alpha_test_key(), alpha_test_key: self.alpha_test_key(),
auth_method_id: self auth_method_id: self
.auth_method_id .auth_method_id
@@ -429,8 +425,6 @@ impl MvpAgent {
memory_config: self.memory_config.clone(), memory_config: self.memory_config.clone(),
web_search_config: self.prepare_web_search_config(), web_search_config: self.prepare_web_search_config(),
web_fetch_config: self.prepare_web_fetch_config(), web_fetch_config: self.prepare_web_fetch_config(),
image_gen_config: self.prepare_image_gen_config(),
video_gen_config: self.prepare_video_gen_config(),
app_builder_deployer_config: self.prepare_app_builder_deployer_config(), app_builder_deployer_config: self.prepare_app_builder_deployer_config(),
write_file_enabled: self.cfg.borrow().resolve_write_file().value, write_file_enabled: self.cfg.borrow().resolve_write_file().value,
goal_enabled: self.cfg.borrow().resolve_goal().value, goal_enabled: self.cfg.borrow().resolve_goal().value,
@@ -505,7 +499,6 @@ impl MvpAgent {
.map(|h| h.mcp_servers.clone()) .map(|h| h.mcp_servers.clone())
.unwrap_or_default() .unwrap_or_default()
}, },
managed_mcp_state: self.managed_mcp_cache.clone(),
parent_mcp_pool: None, parent_mcp_pool: None,
parent_tool_snapshot: None, parent_tool_snapshot: None,
parent_skills: None, parent_skills: None,
@@ -648,8 +648,6 @@ async fn file_toolset_override_e2e_to_finalized_toolset() {
web_search_config: kigi_tools::implementations::web_search::WebSearchConfig::default(), web_search_config: kigi_tools::implementations::web_search::WebSearchConfig::default(),
web_fetch_config: Default::default(), web_fetch_config: Default::default(),
lsp: None, lsp: None,
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(), app_builder_deployer_config: kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig::default(),
api_key_provider: None, api_key_provider: None,
attribution_callback: None, attribution_callback: None,
@@ -744,7 +742,6 @@ fn make_test_handle(
permission_handle: kigi_workspace::permission::PermissionHandle::allow_all(), permission_handle: kigi_workspace::permission::PermissionHandle::allow_all(),
attribution_callback: None, attribution_callback: None,
agent_name: "grok-build".to_string(), agent_name: "grok-build".to_string(),
managed_mcp_proxy_base_url: String::new(),
session_default_agent_profile: None, session_default_agent_profile: None,
allowed_subagent_types: None, allowed_subagent_types: None,
hook_registry: None, hook_registry: None,
@@ -1136,38 +1133,6 @@ fn test_sessionless_request_requires_session_id() {
"cwd-only requests with no sessionId must return SessionRequired" "cwd-only requests with no sessionId must return SessionRequired"
); );
} }
#[tokio::test(flavor = "current_thread")]
async fn ext_method_routes_auth_cleared_and_refreshes_resident_sessions() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let agent = build_agent_with_auth(crate::auth::KimiAuth {
key: "eligible".into(),
auth_mode: crate::auth::AuthMode::OAuth,
..crate::auth::KimiAuth::test_default()
});
use acp::Agent as _;
agent.managed_mcp_cache.lock().await.enable_gateway_tools();
let sid = acp::SessionId::new("sess-auth-cleared");
let (handle, _tx, mut cmd_rx) = make_live_session_handle(&sid, None);
agent.sessions.borrow_mut().insert(sid, handle);
let params = serde_json::json!({});
agent
.ext_method(acp::ExtRequest::new(
"x.ai/internal/auth_cleared",
std::sync::Arc::from(serde_json::value::to_raw_value(&params).unwrap()),
))
.await
.expect("auth_cleared must route through session-admin");
let cmd = tokio::time::timeout(std::time::Duration::from_secs(1), cmd_rx.recv())
.await
.expect("refresh command should be sent")
.expect("channel should stay open until command is received");
assert!(matches!(cmd, SessionCommand::RefreshMcpSearchIndex));
assert!(!agent.managed_mcp_cache.lock().await.gateway_tools_active);
})
.await;
}
/// Build a minimal MvpAgent suitable for testing extension methods. /// Build a minimal MvpAgent suitable for testing extension methods.
fn build_minimal_agent_for_tests() -> MvpAgent { fn build_minimal_agent_for_tests() -> MvpAgent {
use crate::agent::config::Config as AgentConfig; use crate::agent::config::Config as AgentConfig;
@@ -1860,7 +1825,7 @@ async fn auth_type_session_based_no_current_returns_session_token() {
/// BYOK guard. Users with `xai.api_key` must continue to report `ApiKey` /// BYOK guard. Users with `xai.api_key` must continue to report `ApiKey`
/// regardless of live-token state -- BYOK sessions have nothing to refresh, /// regardless of live-token state -- BYOK sessions have nothing to refresh,
/// and reporting `SessionToken` would route through cli-chat-proxy paths /// and reporting `SessionToken` would route through cli-chat-proxy paths
/// (image_gen / video_gen base_url) that don't apply to BYOK keys. /// that don't apply to BYOK keys.
#[tokio::test(flavor = "current_thread")] #[tokio::test(flavor = "current_thread")]
async fn auth_type_xai_api_key_no_current_returns_api_key() { async fn auth_type_xai_api_key_no_current_returns_api_key() {
let agent = build_minimal_agent_for_tests(); let agent = build_minimal_agent_for_tests();
@@ -1952,26 +1917,6 @@ async fn cached_token_fallthrough_falls_to_grok_com_without_credentials() {
"no API-key creds and no kill switch -> interactive grok.com login", "no API-key creds and no kill switch -> interactive grok.com login",
); );
} }
/// The imagine tier gate fails **open**: with no resolved auth we can't confirm
/// a restricted personal tier, so the tools stay advertised and un-flagged (the
/// server 429 remains the authoritative backstop). Guards against accidentally
/// disabling a paid feature when tier info hasn't loaded.
#[tokio::test(flavor = "current_thread")]
async fn prepare_image_gen_config_fails_open_without_auth() {
use kigi_tools::implementations::grok_build::image_gen::ImageGenConfig;
let agent = build_minimal_agent_for_tests();
agent.sampling_config.borrow_mut().api_key = Some("test-key".to_string());
let ImageGenConfig::Enabled {
tier_restricted, ..
} = agent.prepare_image_gen_config()
else {
panic!("expected Enabled");
};
assert!(
!tier_restricted,
"no resolved auth ⇒ fail open (tools not tier-restricted)"
);
}
/// `parse_session_kind` routes `session/load` to the gateway Chat path vs. the /// `parse_session_kind` routes `session/load` to the gateway Chat path vs. the
/// disk-backed Build path. Anything but an explicit `kind: "chat"` is Build. /// disk-backed Build path. Anything but an explicit `kind: "chat"` is Build.
#[test] #[test]
@@ -2107,7 +2052,6 @@ fn chat_session_spawn_options_matches_thin_profile() {
assert!(!opts.client_fs_read); assert!(!opts.client_fs_read);
assert!(!opts.client_fs_write); assert!(!opts.client_fs_write);
assert!(opts.chat_history.is_empty()); assert!(opts.chat_history.is_empty());
assert!(opts.managed_mcp_expires_at.is_none());
assert!(!opts.session_auto_mode); assert!(!opts.session_auto_mode);
assert!( assert!(
opts.persistence.is_noop(), opts.persistence.is_noop(),
@@ -1054,9 +1054,6 @@ pub(crate) async fn handle_subagent_request(
ctx.memory_config.clone() ctx.memory_config.clone()
}, },
Default::default(), Default::default(),
ctx.managed_mcp_state.clone(),
None,
ctx.managed_mcp_proxy_base_url.clone(),
effective_model_id, effective_model_id,
ctx.yolo_mode ctx.yolo_mode
|| matches!( || matches!(
@@ -1069,8 +1066,6 @@ pub(crate) async fn handle_subagent_request(
None, None,
ctx.web_search_config.clone(), ctx.web_search_config.clone(),
ctx.web_fetch_config.clone(), ctx.web_fetch_config.clone(),
ctx.image_gen_config.clone(),
ctx.video_gen_config.clone(),
ctx.app_builder_deployer_config.clone(), ctx.app_builder_deployer_config.clone(),
ctx.write_file_enabled, ctx.write_file_enabled,
ctx.goal_enabled, ctx.goal_enabled,
@@ -144,7 +144,6 @@ pub(crate) struct SubagentSpawnContext {
/// context is built (an async snapshot from the parent session actor). /// context is built (an async snapshot from the parent session actor).
pub client_hooks: crate::extensions::hooks::ClientHooks, pub client_hooks: crate::extensions::hooks::ClientHooks,
pub sampling_config: kigi_sampler::SamplerConfig, pub sampling_config: kigi_sampler::SamplerConfig,
pub managed_mcp_proxy_base_url: String,
/// The staging auth header value propagated from the parent. Used /// The staging auth header value propagated from the parent. Used
/// when materialising subagent `SamplerConfig`s for auth-flow tracking /// when materialising subagent `SamplerConfig`s for auth-flow tracking
/// and for `inject_url_derived_headers` in the construction helpers. /// and for `inject_url_derived_headers` in the construction helpers.
@@ -209,10 +208,6 @@ pub(crate) struct SubagentSpawnContext {
pub web_search_config: kigi_tools::implementations::WebSearchConfig, pub web_search_config: kigi_tools::implementations::WebSearchConfig,
/// Resolved config for web fetch. /// Resolved config for web fetch.
pub web_fetch_config: kigi_tools::implementations::grok_build::web_fetch::WebFetchConfig, pub web_fetch_config: kigi_tools::implementations::grok_build::web_fetch::WebFetchConfig,
/// Image generation config (parent-inherited).
pub image_gen_config: kigi_tools::implementations::grok_build::image_gen::ImageGenConfig,
/// Resolved config for video generation.
pub video_gen_config: kigi_tools::implementations::grok_build::video_gen::VideoGenConfig,
/// Resolved config for the deploy service. /// Resolved config for the deploy service.
pub app_builder_deployer_config: pub app_builder_deployer_config:
kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig, kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig,
@@ -320,8 +315,6 @@ pub(crate) struct SubagentSpawnContext {
/// time). Servers added later via `UpdateMcpServers` (managed MCPs, plugin reload) /// time). Servers added later via `UpdateMcpServers` (managed MCPs, plugin reload)
/// will not appear here. Named references only resolve against the initial config. /// will not appear here. Named references only resolve against the initial config.
pub parent_mcp_configs: Vec<agent_client_protocol::McpServer>, pub parent_mcp_configs: Vec<agent_client_protocol::McpServer>,
/// Parent's managed MCP state handle (Arc-shared, no re-fetch).
pub managed_mcp_state: crate::session::managed_mcp::ManagedMcpStateHandle,
/// Snapshot of the parent session's MCP client pool at spawn time. /// Snapshot of the parent session's MCP client pool at spawn time.
pub parent_mcp_pool: Option<crate::session::mcp_servers::SharedMcpPool>, pub parent_mcp_pool: Option<crate::session::mcp_servers::SharedMcpPool>,
/// Snapshot of the parent session's resolved tool schema at spawn time. /// Snapshot of the parent session's resolved tool schema at spawn time.
@@ -143,16 +143,6 @@ fn subagent_inherits_parent_lsp_via_context() {
); );
} }
#[test] #[test]
fn subagent_inherits_managed_mcp_state_via_context() {
let handle = crate::session::managed_mcp::ManagedMcpStateHandle::default();
let mut ctx = ctx_with_toggle(HashMap::new());
ctx.managed_mcp_state = handle.clone();
assert!(
Arc::ptr_eq(& handle, & ctx.managed_mcp_state),
"child should share parent's managed MCP state (Arc identity)"
);
}
#[test]
fn no_parent_lsp_means_child_gets_none() { fn no_parent_lsp_means_child_gets_none() {
let ctx = ctx_with_toggle(HashMap::new()); let ctx = ctx_with_toggle(HashMap::new());
assert!(ctx.lsp.is_none()); assert!(ctx.lsp.is_none());
@@ -1110,7 +1100,6 @@ fn dummy_tracker(
permission_handle: kigi_workspace::permission::PermissionHandle::allow_all(), permission_handle: kigi_workspace::permission::PermissionHandle::allow_all(),
attribution_callback: None, attribution_callback: None,
agent_name: "grok-build".to_string(), agent_name: "grok-build".to_string(),
managed_mcp_proxy_base_url: String::new(),
session_default_agent_profile: None, session_default_agent_profile: None,
allowed_subagent_types: None, allowed_subagent_types: None,
hook_registry: None, hook_registry: None,
@@ -123,7 +123,7 @@ impl ShellAttribution {
/// Tool-side counterpart of [`Self::new`]: returns /// Tool-side counterpart of [`Self::new`]: returns
/// `Arc<dyn kigi_tools::Auth401AttributionCallback>` for the /// `Arc<dyn kigi_tools::Auth401AttributionCallback>` for the
/// `with_attribution_callback(...)` builder on each tool HTTP /// `with_attribution_callback(...)` builder on each tool HTTP
/// client (`ImageGenClient`, `VideoGenClient`, `WebSearchClient`). /// client (`WebSearchClient`).
/// The two callbacks share the same underlying impl and emit the /// The two callbacks share the same underlying impl and emit the
/// same `auth_401_attribution` event format -- only the trait /// same `auth_401_attribution` event format -- only the trait
/// signature differs (`SamplingConsumer` vs. `ToolConsumer`). /// signature differs (`SamplingConsumer` vs. `ToolConsumer`).
@@ -161,20 +161,13 @@ impl Auth401AttributionCallback for ShellAttribution {
} }
} }
/// Tool-side hook: each tool client (image_gen, video_gen, web_search) /// Tool-side hook: each tool client (web_search) in `kigi-tools`
/// in `kigi-tools` emits a 401 attribution event through this /// emits a 401 attribution event through this trait when its HTTP
/// trait when its HTTP request returns UNAUTHORIZED. Same shape as /// request returns UNAUTHORIZED. Same shape as the sampler-side impl
/// the sampler-side impl above; routes to the same pair of sinks. /// above; routes to the same pair of sinks.
///
/// `ToolConsumer::VideoGenStart` and `VideoGenPoll` collapse to the
/// same [`ConsumerKind::VideoGen`] with different op strings so the
/// gate query can break down video-gen 401s by phase.
impl ToolAuth401AttributionCallback for ShellAttribution { impl ToolAuth401AttributionCallback for ShellAttribution {
fn record_401(&self, consumer: ToolConsumer, sent_bearer_prefix: Option<&str>) { fn record_401(&self, consumer: ToolConsumer, sent_bearer_prefix: Option<&str>) {
let (kind, op) = match consumer { let (kind, op) = match consumer {
ToolConsumer::ImageGen => (ConsumerKind::ImageGen, ""),
ToolConsumer::VideoGenStart => (ConsumerKind::VideoGen, "start"),
ToolConsumer::VideoGenPoll => (ConsumerKind::VideoGen, "poll"),
ToolConsumer::WebSearch => (ConsumerKind::WebSearch, ""), ToolConsumer::WebSearch => (ConsumerKind::WebSearch, ""),
}; };
record_consumer_401( record_consumer_401(
@@ -207,15 +200,6 @@ pub(crate) enum ConsumerKind {
/// No per-op discriminator -- the consumer string is just /// No per-op discriminator -- the consumer string is just
/// `"IdleResumeModelRefresh"`. /// `"IdleResumeModelRefresh"`.
IdleResumeModelRefresh, IdleResumeModelRefresh,
/// `kigi_tools::ToolConsumer::ImageGen` -- Imagine API
/// (`POST /images/generations`). No per-op discriminator;
/// consumer string is just `"ImageGen"`.
ImageGen,
/// `kigi_tools::ToolConsumer::VideoGenStart` and
/// `VideoGenPoll` -- Video Generation API. The op string is
/// `"start"` (`POST /videos/generations`) or `"poll"`
/// (`GET /videos/{request_id}`).
VideoGen,
/// `kigi_tools::ToolConsumer::WebSearch` -- web search via /// `kigi_tools::ToolConsumer::WebSearch` -- web search via
/// `POST /responses` with a `WebSearch` tool. No per-op /// `POST /responses` with a `WebSearch` tool. No per-op
/// discriminator; consumer string is just `"WebSearch"`. /// discriminator; consumer string is just `"WebSearch"`.
@@ -230,8 +214,6 @@ impl ConsumerKind {
Self::FeedbackClient => "FeedbackClient", Self::FeedbackClient => "FeedbackClient",
Self::SessionRegistryClient => "SessionRegistryClient", Self::SessionRegistryClient => "SessionRegistryClient",
Self::IdleResumeModelRefresh => "IdleResumeModelRefresh", Self::IdleResumeModelRefresh => "IdleResumeModelRefresh",
Self::ImageGen => "ImageGen",
Self::VideoGen => "VideoGen",
Self::WebSearch => "WebSearch", Self::WebSearch => "WebSearch",
} }
} }
@@ -239,13 +221,10 @@ impl ConsumerKind {
/// `true` for variants that take a per-operation discriminator /// `true` for variants that take a per-operation discriminator
/// appended as `<prefix>.<op>`. `false` for variants whose /// appended as `<prefix>.<op>`. `false` for variants whose
/// `consumer` string is just the prefix /// `consumer` string is just the prefix
/// (`IdleResumeModelRefresh`, `ImageGen`, `WebSearch` -- each is /// (`IdleResumeModelRefresh`, `WebSearch` -- each is
/// a single endpoint with no sub-operation). /// a single endpoint with no sub-operation).
fn takes_op(self) -> bool { fn takes_op(self) -> bool {
!matches!( !matches!(self, Self::IdleResumeModelRefresh | Self::WebSearch)
self,
Self::IdleResumeModelRefresh | Self::ImageGen | Self::WebSearch
)
} }
} }
@@ -261,7 +240,7 @@ fn format_consumer(kind: ConsumerKind, op: &str) -> String {
/// Emit a single `auth 401 attribution` event for a per-consumer 401. /// Emit a single `auth 401 attribution` event for a per-consumer 401.
/// ///
/// Wraps [`record_auth_401`] with the design-doc `consumer` formatting /// Wraps [`record_auth_401`] with the design-doc `consumer` formatting
/// (e.g., `"FeedbackClient.submit"`, `"VideoGen.start"`). /// (e.g., `"FeedbackClient.submit"`, `"WebSearch"`).
/// All 401 emit sites in `kigi-shell` go through this helper -- the /// All 401 emit sites in `kigi-shell` go through this helper -- the
/// per-client `record_401_attribution` wrappers in /// per-client `record_401_attribution` wrappers in
/// `agent/feedback_client.rs` and `agent/session_registry_client.rs` each /// `agent/feedback_client.rs` and `agent/session_registry_client.rs` each
@@ -606,10 +585,6 @@ mod tests {
"ignored", "ignored",
"IdleResumeModelRefresh", "IdleResumeModelRefresh",
), ),
(ConsumerKind::ImageGen, "", "ImageGen"),
(ConsumerKind::ImageGen, "ignored", "ImageGen"),
(ConsumerKind::VideoGen, "start", "VideoGen.start"),
(ConsumerKind::VideoGen, "poll", "VideoGen.poll"),
(ConsumerKind::WebSearch, "", "WebSearch"), (ConsumerKind::WebSearch, "", "WebSearch"),
(ConsumerKind::WebSearch, "ignored", "WebSearch"), (ConsumerKind::WebSearch, "ignored", "WebSearch"),
]; ];
@@ -647,12 +622,7 @@ mod tests {
let cb: Arc<dyn ToolAuth401AttributionCallback> = let cb: Arc<dyn ToolAuth401AttributionCallback> =
ShellAttribution::new_tool_callback(am_arc.clone(), Some("sid-tool".into())); ShellAttribution::new_tool_callback(am_arc.clone(), Some("sid-tool".into()));
let cases = [ let cases = [(ToolConsumer::WebSearch, "WebSearch")];
(ToolConsumer::ImageGen, "ImageGen"),
(ToolConsumer::VideoGenStart, "VideoGen.start"),
(ToolConsumer::VideoGenPoll, "VideoGen.poll"),
(ToolConsumer::WebSearch, "WebSearch"),
];
for (consumer, expected_consumer_str) in cases { for (consumer, expected_consumer_str) in cases {
cb.record_401(consumer, Some("bearer-1234567890")); cb.record_401(consumer, Some("bearer-1234567890"));
@@ -1877,7 +1877,7 @@ impl AuthManager {
} }
/// Bridges `Arc<AuthManager>` into the `ApiKeyProvider` trait used by /// Bridges `Arc<AuthManager>` into the `ApiKeyProvider` trait used by
/// tool clients (image_gen, video_gen, web_search, embedding). Sync callers /// tool clients (web_search, embedding). Sync callers
/// get the buffered snapshot; async callers drive the refresh chain. /// get the buffered snapshot; async callers drive the refresh chain.
pub(crate) struct SharedAuthKeyProvider(pub Arc<AuthManager>); pub(crate) struct SharedAuthKeyProvider(pub Arc<AuthManager>);
+6 -7
View File
@@ -5,7 +5,6 @@ const BUNDLED_FILES: &[(&str, &str)] = &[("README.md", include_str!("../README.m
const HELP_SKILL_MD: &str = include_str!("../skills/help/SKILL.md"); const HELP_SKILL_MD: &str = include_str!("../skills/help/SKILL.md");
const CREATE_SKILL_MD: &str = include_str!("../skills/create-skill/SKILL.md"); const CREATE_SKILL_MD: &str = include_str!("../skills/create-skill/SKILL.md");
const CODE_REVIEW_SKILL_MD: &str = include_str!("../skills/code-review/SKILL.md"); const CODE_REVIEW_SKILL_MD: &str = include_str!("../skills/code-review/SKILL.md");
const IMAGINE_SKILL_MD: &str = include_str!("../skills/imagine/SKILL.md");
/// Compiled-in SKILL.md content for `/check-work` (available to headless mode). /// Compiled-in SKILL.md content for `/check-work` (available to headless mode).
pub const CHECK_SKILL_MD: &str = include_str!("../skills/check-work/SKILL.md"); pub const CHECK_SKILL_MD: &str = include_str!("../skills/check-work/SKILL.md");
/// Compiled-in SKILL.md content for headless `--best-of-n` (not extracted as /// Compiled-in SKILL.md content for headless `--best-of-n` (not extracted as
@@ -40,7 +39,8 @@ pub const BEST_OF_N_SKILL_MD: &str = include_str!("../skills/best-of-n/SKILL.md"
/// them for many releases. /// them for many releases.
/// - After the rename has had time to propagate, you **may** clean old /// - After the rename has had time to propagate, you **may** clean old
/// strings out of this list for hygiene. /// strings out of this list for hygiene.
const LEGACY_BUNDLED_SKILL_NAMES: &[&str] = &["check", "best-of-n", "docx", "pptx", "xlsx"]; const LEGACY_BUNDLED_SKILL_NAMES: &[&str] =
&["check", "best-of-n", "docx", "pptx", "xlsx", "imagine"];
/// All bundled skill SKILL.md files. Single source of truth used by both /// All bundled skill SKILL.md files. Single source of truth used by both
/// the full extraction path (version bump) and the missing-file fast path /// the full extraction path (version bump) and the missing-file fast path
@@ -56,7 +56,6 @@ const BUNDLED_SKILLS: &[(&str, &str)] = &[
("help", HELP_SKILL_MD), ("help", HELP_SKILL_MD),
("create-skill", CREATE_SKILL_MD), ("create-skill", CREATE_SKILL_MD),
("code-review", CODE_REVIEW_SKILL_MD), ("code-review", CODE_REVIEW_SKILL_MD),
("imagine", IMAGINE_SKILL_MD),
("check-work", CHECK_SKILL_MD), ("check-work", CHECK_SKILL_MD),
]; ];
@@ -210,13 +209,13 @@ mod tests {
std::fs::write(home.join(filename), "old").unwrap(); std::fs::write(home.join(filename), "old").unwrap();
} }
std::fs::write(home.join("skills/help/SKILL.md"), "old").unwrap(); std::fs::write(home.join("skills/help/SKILL.md"), "old").unwrap();
for name in ["check-work", "imagine", "code-review"] { for name in ["check-work", "code-review"] {
std::fs::write(home.join(format!("skills/{name}/SKILL.md")), "old").unwrap(); std::fs::write(home.join(format!("skills/{name}/SKILL.md")), "old").unwrap();
} }
std::fs::write(home.join(".metadata_version"), "0.0.0-stale").unwrap(); std::fs::write(home.join(".metadata_version"), "0.0.0-stale").unwrap();
// Simulate legacy skills that should be cleaned up. // Simulate legacy skills that should be cleaned up.
for name in ["check", "best-of-n", "docx", "pptx", "xlsx"] { for name in ["check", "best-of-n", "docx", "pptx", "xlsx", "imagine"] {
std::fs::create_dir_all(home.join(format!("skills/{name}"))).unwrap(); std::fs::create_dir_all(home.join(format!("skills/{name}"))).unwrap();
std::fs::write( std::fs::write(
home.join(format!("skills/{name}/SKILL.md")), home.join(format!("skills/{name}/SKILL.md")),
@@ -238,7 +237,7 @@ mod tests {
std::fs::read_to_string(home.join("skills/help/SKILL.md")).unwrap(), std::fs::read_to_string(home.join("skills/help/SKILL.md")).unwrap(),
"old" "old"
); );
for name in ["check-work", "imagine", "code-review"] { for name in ["check-work", "code-review"] {
assert_ne!( assert_ne!(
std::fs::read_to_string(home.join(format!("skills/{name}/SKILL.md"))).unwrap(), std::fs::read_to_string(home.join(format!("skills/{name}/SKILL.md"))).unwrap(),
"old", "old",
@@ -248,7 +247,7 @@ mod tests {
// Legacy skill directories must have been removed (the key part of // Legacy skill directories must have been removed (the key part of
// supporting renames like check → check-work without leaving orphans). // supporting renames like check → check-work without leaving orphans).
for name in ["check", "best-of-n", "docx", "pptx", "xlsx"] { for name in ["check", "best-of-n", "docx", "pptx", "xlsx", "imagine"] {
assert!( assert!(
!home.join(format!("skills/{name}")).exists(), !home.join(format!("skills/{name}")).exists(),
"legacy '{name}' skill directory should have been deleted during version bump" "legacy '{name}' skill directory should have been deleted during version bump"
+3 -117
View File
@@ -474,62 +474,6 @@ impl SubagentsConfig {
result result
} }
} }
/// Managed MCP connector fetching config (`[managed_mcps]` in config.toml).
///
/// See [`Self::resolve`] for full priority chain.
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
#[serde(default)]
pub struct ManagedMcpsConfig {
pub enabled: bool,
pub gateway_tools_enabled: bool,
}
impl Default for ManagedMcpsConfig {
fn default() -> Self {
Self {
enabled: true,
gateway_tools_enabled: false,
}
}
}
impl ManagedMcpsConfig {
/// Priority: env var > TOML > remote > default (enabled interactive, disabled headless).
pub fn resolve(
config: &toml::Value,
remote: Option<&crate::util::config::RemoteSettings>,
is_headless: bool,
) -> Self {
let mut result: Self = config
.get("managed_mcps")
.and_then(|v| v.clone().try_into().ok())
.unwrap_or(Self {
enabled: !is_headless,
gateway_tools_enabled: false,
});
let managed_mcps_table = config.get("managed_mcps").and_then(|v| v.as_table());
let has_local_enabled = managed_mcps_table.is_some_and(|t| t.contains_key("enabled"));
let resolved = crate::agent::config::resolve_enabled(
None,
"KIGI_MANAGED_MCPS_ENABLED",
result.enabled,
has_local_enabled,
remote.and_then(|r| r.managed_mcps_enabled),
!is_headless,
);
result.enabled = resolved.value;
let has_local_gateway_tools =
managed_mcps_table.is_some_and(|t| t.contains_key("gateway_tools_enabled"));
let gateway_resolved = crate::agent::config::resolve_enabled(
None,
"KIGI_MANAGED_MCP_GATEWAY_TOOLS_ENABLED",
result.gateway_tools_enabled,
has_local_gateway_tools,
remote.and_then(|r| r.managed_mcp_gateway_tools_enabled),
false,
);
result.gateway_tools_enabled = result.enabled && gateway_resolved.value;
result
}
}
/// Auxiliary model overrides under `[models]`. /// Auxiliary model overrides under `[models]`.
#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)] #[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
#[serde(default)] #[serde(default)]
@@ -656,8 +600,7 @@ impl ModelOverrideConfig {
/// ///
/// ```toml /// ```toml
/// [tools] /// [tools]
/// disable_zdr_incompatible_tools = true /// respect_gitignore = true
/// # [tools.zdr_video_output_s3] — see ZdrVideoOutputS3Config
/// ``` /// ```
#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)] #[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
#[serde(default)] #[serde(default)]
@@ -665,29 +608,12 @@ pub struct ToolsConfig {
/// When `true`, all tools (including `read_file`) filter gitignored /// When `true`, all tools (including `read_file`) filter gitignored
/// files. When `false` (default), each tool picks its own default. /// files. When `false` (default), each tool picks its own default.
pub respect_gitignore: bool, pub respect_gitignore: bool,
/// Drop tools whose xAI API requires server-side artifact storage
/// (currently just `video_gen`). Intended for ZDR-bound teams via
/// `~/.kigi/managed_config.toml`. Defaults to `false`.
pub disable_zdr_incompatible_tools: bool,
/// Optional S3 bucket config for ZDR video output. When present (and
/// valid), video tools presign an upload URL and pass it to the API so
/// the generated video lands in a team-owned bucket instead of being
/// downloaded locally. Only effective when `disable_zdr_incompatible_tools`
/// is `true`. Populated from `[tools.zdr_video_output_s3]` in config.
pub zdr_video_output_s3:
Option<kigi_tools::implementations::grok_build::video_gen::ZdrVideoOutputS3Config>,
} }
impl ToolsConfig { impl ToolsConfig {
/// Resolve the final tools config, in priority order: /// Resolve the final tools config, in priority order:
/// 1. Env vars `KIGI_RESPECT_GITIGNORE` and /// 1. Env var `KIGI_RESPECT_GITIGNORE` (`0`/`false` off, `1`/`true` on).
/// `KIGI_DISABLE_ZDR_INCOMPATIBLE_TOOLS` (`0`/`false` off,
/// `1`/`true` on).
/// 2. `[tools]` block from the merged effective config. /// 2. `[tools]` block from the merged effective config.
/// 3. Defaults (both `false`). /// 3. Default (`false`).
///
/// Fields are read individually so a malformed
/// `[tools.zdr_video_output_s3]` cannot wipe `disable_zdr_incompatible_tools`
/// (or any other tools flag) via whole-table deserialize failure.
pub fn resolve(config: &toml::Value) -> Self { pub fn resolve(config: &toml::Value) -> Self {
let tools = config.get("tools"); let tools = config.get("tools");
let mut result = Self { let mut result = Self {
@@ -695,33 +621,6 @@ impl ToolsConfig {
.and_then(|t| t.get("respect_gitignore")) .and_then(|t| t.get("respect_gitignore"))
.and_then(|v| v.as_bool()) .and_then(|v| v.as_bool())
.unwrap_or(false), .unwrap_or(false),
disable_zdr_incompatible_tools: tools
.and_then(|t| t.get("disable_zdr_incompatible_tools"))
.and_then(|v| v.as_bool())
.unwrap_or(false),
zdr_video_output_s3: tools
.and_then(|t| t.get("zdr_video_output_s3"))
.and_then(|s3_val| match s3_val
.clone()
.try_into::<
kigi_tools::implementations::grok_build::video_gen::ZdrVideoOutputS3Config,
>()
{
Ok(cfg) if cfg.is_valid() => Some(cfg),
Ok(_) => {
tracing::warn!(
"tools.zdr_video_output_s3 is present but incomplete; ignoring ZDR video output config"
);
None
}
Err(e) => {
tracing::warn!(
error = % e,
"tools.zdr_video_output_s3 failed to parse; ignoring ZDR video output config"
);
None
}
}),
}; };
match std::env::var("KIGI_RESPECT_GITIGNORE").as_deref() { match std::env::var("KIGI_RESPECT_GITIGNORE").as_deref() {
Ok("0") | Ok("false") => { Ok("0") | Ok("false") => {
@@ -732,15 +631,6 @@ impl ToolsConfig {
} }
_ => {} _ => {}
} }
match std::env::var("KIGI_DISABLE_ZDR_INCOMPATIBLE_TOOLS").as_deref() {
Ok("0") | Ok("false") => {
result.disable_zdr_incompatible_tools = false;
}
Ok("1") | Ok("true") => {
result.disable_zdr_incompatible_tools = true;
}
_ => {}
}
result result
} }
} }
@@ -1009,9 +899,6 @@ fn apply_requirements_inner(
pin_feature!(tool_search); pin_feature!(tool_search);
pin_feature!(web_fetch); pin_feature!(web_fetch);
pin_feature!(ask_user_question); pin_feature!(ask_user_question);
pin_requirement_only!(image_gen);
pin_requirement_only!(image_edit);
pin_feature!(video_gen);
pin_feature!(write_file); pin_feature!(write_file);
pin_feature!(voice_mode); pin_feature!(voice_mode);
pin_requirement_only!(remote_fetch); pin_requirement_only!(remote_fetch);
@@ -1020,7 +907,6 @@ fn apply_requirements_inner(
enforce_opt!("cli", "show_tips", config.cli.show_tips); enforce_opt!("cli", "show_tips", config.cli.show_tips);
enforce_val!("memory", "enabled", config.memory.enabled); enforce_val!("memory", "enabled", config.memory.enabled);
enforce_val!("subagents", "enabled", config.subagents.enabled); enforce_val!("subagents", "enabled", config.subagents.enabled);
enforce_val!("managed_mcps", "enabled", config.managed_mcps.enabled);
if let Some(val) = req_bool(req, "tools", "respect_gitignore") { if let Some(val) = req_bool(req, "tools", "respect_gitignore") {
config config
.requirements .requirements
+6 -247
View File
@@ -1201,154 +1201,6 @@ fn subagents_config_is_subagent_enabled_false_when_toggled_off() {
); );
assert!(sa.is_subagent_enabled("explore"), "explore = true should return enabled"); assert!(sa.is_subagent_enabled("explore"), "explore = true should return enabled");
} }
fn with_managed_mcp_env<T>(
managed_mcps: Option<&str>,
gateway_tools: Option<&str>,
f: impl FnOnce() -> T,
) -> T {
static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
let _guard = LOCK.lock().unwrap_or_else(|e| e.into_inner());
with_env_var_opt(
"KIGI_MANAGED_MCPS_ENABLED",
managed_mcps,
|| with_env_var_opt("KIGI_MANAGED_MCP_GATEWAY_TOOLS_ENABLED", gateway_tools, f),
)
}
#[test]
#[serial_test::serial]
fn managed_mcps_interactive_default_enabled() {
with_managed_mcp_env(
None,
None,
|| {
let empty = toml::Value::Table(toml::map::Map::new());
let cfg = ManagedMcpsConfig::resolve(&empty, None, false);
assert!(cfg.enabled);
},
);
}
#[test]
#[serial_test::serial]
fn managed_mcps_headless_default_disabled() {
with_managed_mcp_env(
None,
None,
|| {
let empty = toml::Value::Table(toml::map::Map::new());
let cfg = ManagedMcpsConfig::resolve(&empty, None, true);
assert!(! cfg.enabled);
},
);
}
#[test]
#[serial_test::serial]
fn managed_mcp_gateway_tools_default_disabled() {
with_managed_mcp_env(
None,
None,
|| {
let empty = toml::Value::Table(toml::map::Map::new());
let cfg = ManagedMcpsConfig::resolve(&empty, None, false);
assert!(! cfg.gateway_tools_enabled);
},
);
}
#[test]
#[serial_test::serial]
fn managed_mcp_gateway_tools_require_managed_master() {
with_managed_mcp_env(
None,
None,
|| {
let config: toml::Value = toml::from_str(
r#"
[managed_mcps]
gateway_tools_enabled = true
"#,
)
.unwrap();
let remote = crate::util::config::RemoteSettings {
managed_mcps_enabled: Some(false),
..Default::default()
};
let cfg = ManagedMcpsConfig::resolve(&config, Some(&remote), true);
assert!(! cfg.enabled);
assert!(! cfg.gateway_tools_enabled);
},
);
}
#[test]
#[serial_test::serial]
fn managed_mcp_gateway_tools_remote_enabled() {
with_managed_mcp_env(
None,
None,
|| {
let empty = toml::Value::Table(toml::map::Map::new());
let remote = crate::util::config::RemoteSettings {
managed_mcp_gateway_tools_enabled: Some(true),
..Default::default()
};
let cfg = ManagedMcpsConfig::resolve(&empty, Some(&remote), false);
assert!(cfg.gateway_tools_enabled);
},
);
}
#[test]
#[serial_test::serial]
fn managed_mcp_gateway_tools_env_overrides_remote() {
with_managed_mcp_env(
None,
Some("0"),
|| {
let empty = toml::Value::Table(toml::map::Map::new());
let remote = crate::util::config::RemoteSettings {
managed_mcp_gateway_tools_enabled: Some(true),
..Default::default()
};
let cfg = ManagedMcpsConfig::resolve(&empty, Some(&remote), false);
assert!(! cfg.gateway_tools_enabled);
},
);
}
#[test]
#[serial_test::serial]
fn managed_mcp_gateway_tools_env_on_overrides_remote_off() {
with_managed_mcp_env(
None,
Some("1"),
|| {
let empty = toml::Value::Table(toml::map::Map::new());
let remote = crate::util::config::RemoteSettings {
managed_mcp_gateway_tools_enabled: Some(false),
..Default::default()
};
let cfg = ManagedMcpsConfig::resolve(&empty, Some(&remote), false);
assert!(cfg.gateway_tools_enabled);
},
);
}
#[test]
#[serial_test::serial]
fn managed_mcp_gateway_tools_enabled_with_managed_master() {
with_managed_mcp_env(
None,
None,
|| {
let config: toml::Value = toml::from_str(
r#"
[managed_mcps]
enabled = true
gateway_tools_enabled = true
"#,
)
.unwrap();
let cfg = ManagedMcpsConfig::resolve(&config, None, false);
assert!(cfg.enabled);
assert!(cfg.gateway_tools_enabled);
},
);
}
fn with_model_overrides_env_full<T>( fn with_model_overrides_env_full<T>(
ss: Option<&str>, ss: Option<&str>,
id: Option<&str>, id: Option<&str>,
@@ -1777,28 +1629,16 @@ fn model_overrides_prompt_suggestion_blank_values_are_unset() {
}, },
); );
} }
/// Lock shared by every test that touches the env vars read by /// Lock shared by every test that touches the env var read by
/// `ToolsConfig::resolve`, so tests across both fields can't race. /// `ToolsConfig::resolve`, so tests can't race.
static TOOLS_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); static TOOLS_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
/// Set both `ToolsConfig` env vars for the duration of `f`, then
/// restore. `None` clears the var.
fn with_tools_env<T>(
respect_gitignore: Option<&str>,
disable_zdr: Option<&str>,
f: impl FnOnce() -> T,
) -> T {
let _guard = TOOLS_ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
with_env_var_opt(
"KIGI_RESPECT_GITIGNORE",
respect_gitignore,
|| with_env_var_opt("KIGI_DISABLE_ZDR_INCOMPATIBLE_TOOLS", disable_zdr, f),
)
}
fn without_grok_respect_gitignore<T>(f: impl FnOnce() -> T) -> T { fn without_grok_respect_gitignore<T>(f: impl FnOnce() -> T) -> T {
with_tools_env(None, None, f) let _guard = TOOLS_ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
with_env_var_opt("KIGI_RESPECT_GITIGNORE", None, f)
} }
fn with_grok_respect_gitignore<T>(value: &str, f: impl FnOnce() -> T) -> T { fn with_grok_respect_gitignore<T>(value: &str, f: impl FnOnce() -> T) -> T {
with_tools_env(Some(value), None, f) let _guard = TOOLS_ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
with_env_var_opt("KIGI_RESPECT_GITIGNORE", Some(value), f)
} }
#[test] #[test]
fn tools_config_default_disabled() { fn tools_config_default_disabled() {
@@ -1858,87 +1698,6 @@ fn tools_config_env_false_overrides_toml_true() {
); );
} }
#[test] #[test]
fn zdr_incompatible_tools_env_overrides_toml_false() {
with_tools_env(
None,
Some("true"),
|| {
let config: toml::Value = toml::from_str(
"[tools]\ndisable_zdr_incompatible_tools = false",
)
.unwrap();
let tc = ToolsConfig::resolve(&config);
assert!(tc.disable_zdr_incompatible_tools, "env must override TOML");
},
);
}
#[test]
fn zdr_video_output_s3_deserializes_from_tools_block() {
let config: toml::Value = toml::from_str(
r#"
[tools]
disable_zdr_incompatible_tools = true
[tools.zdr_video_output_s3]
bucket = "team-videos"
endpoint = "https://s3.example.com"
region = "us-east-1"
[tools.zdr_video_output_s3.read_write]
access_key_id = "AKIA..."
secret_access_key = "secret"
"#,
)
.unwrap();
let tc = ToolsConfig::resolve(&config);
let s3 = tc.zdr_video_output_s3.expect("zdr_video_output_s3 should deserialize");
assert_eq!(s3.bucket, "team-videos");
assert!(s3.is_valid());
}
#[test]
fn incomplete_zdr_video_output_s3_is_ignored() {
without_grok_respect_gitignore(|| {
let config: toml::Value = toml::from_str(
r#"
[tools]
disable_zdr_incompatible_tools = true
[tools.zdr_video_output_s3]
bucket = "team-videos"
"#,
)
.unwrap();
let tc = ToolsConfig::resolve(&config);
assert!(tc.zdr_video_output_s3.is_none());
assert!(
tc.disable_zdr_incompatible_tools,
"incomplete zdr_video_output_s3 must not drop disable_zdr_incompatible_tools"
);
});
}
#[test]
fn malformed_zdr_video_output_s3_preserves_zdr_flag() {
without_grok_respect_gitignore(|| {
let config: toml::Value = toml::from_str(
r#"
[tools]
disable_zdr_incompatible_tools = true
respect_gitignore = true
[tools.zdr_video_output_s3]
bucket = "team-videos"
endpoint = "https://s3.example.com"
region = "us-east-1"
"#,
)
.unwrap();
let tc = ToolsConfig::resolve(&config);
assert!(tc.zdr_video_output_s3.is_none());
assert!(tc.disable_zdr_incompatible_tools);
assert!(tc.respect_gitignore);
});
}
#[test]
fn roles_parse_from_toml() { fn roles_parse_from_toml() {
let toml_str = r#" let toml_str = r#"
[roles.researcher] [roles.researcher]
+46 -701
View File
@@ -2,7 +2,7 @@
//! //!
//! - `x.ai/mcp/list` — list available MCP servers (agent-scoped or session-annotated) //! - `x.ai/mcp/list` — list available MCP servers (agent-scoped or session-annotated)
//! - `x.ai/mcp/call` — invoke an MCP tool directly, outside the LLM loop //! - `x.ai/mcp/call` — invoke an MCP tool directly, outside the LLM loop
//! - `x.ai/mcp/servers_updated` — notification pushed when managed configs resolve //! - `x.ai/mcp/servers_updated` — notification pushed when the server list changes
//! - `x.ai/mcp/server_status` — per-server delta pushed by the //! - `x.ai/mcp/server_status` — per-server delta pushed by the
//! `StatusDispatcher` (transport-closed pollers, handshake failures, //! `StatusDispatcher` (transport-closed pollers, handshake failures,
//! config diffs, server-pushed list-changed notifications). See //! config diffs, server-pushed list-changed notifications). See
@@ -48,7 +48,6 @@ pub mod mcp_methods {
pub const INIT_PROGRESS: &str = "x.ai/mcp/init_progress"; pub const INIT_PROGRESS: &str = "x.ai/mcp/init_progress";
} }
use crate::agent::MvpAgent; use crate::agent::MvpAgent;
use crate::session::managed_mcp::MANAGED_MCP_PREFIX;
use crate::session::mcp_servers::{MCP_TOOL_NAME_DELIMITER, McpClient, McpServerName, McpState}; use crate::session::mcp_servers::{MCP_TOOL_NAME_DELIMITER, McpClient, McpServerName, McpState};
// ── Wire types: mcp/list ──────────────────────────────────────────── // ── Wire types: mcp/list ────────────────────────────────────────────
@@ -58,8 +57,9 @@ use crate::session::mcp_servers::{MCP_TOOL_NAME_DELIMITER, McpClient, McpServerN
pub struct McpListRequest { pub struct McpListRequest {
#[serde(default)] #[serde(default)]
pub session_id: Option<String>, pub session_id: Option<String>,
/// When false, bypasses the managed MCP config cache and fetches fresh /// When false, forces an OAuth retry pass on servers stuck in
/// from cli-chat-proxy. Set this after OAuth enrollment or disconnect. /// auth-required state before building the list. Set this after OAuth
/// enrollment or disconnect.
#[serde(default = "default_true")] #[serde(default = "default_true")]
pub cache: bool, pub cache: bool,
} }
@@ -92,21 +92,13 @@ pub struct McpServerEntry {
/// MCP server config for the `mcp/list` catalog response. /// MCP server config for the `mcp/list` catalog response.
/// ///
/// Distinct from `acp::McpServer` (session/new input) because: /// Distinct from `acp::McpServer` (session/new input) because:
/// - HTTP: exposes `scope`/`scope_id`/`scope_name` for connector selection, NOT headers (auth tokens stay private) /// - HTTP: exposes the URL, NOT headers (auth tokens stay private)
/// - Stdio: same structure but optimized for JSON wire format /// - Stdio: same structure but optimized for JSON wire format
#[derive(Debug, Clone, Serialize)] #[derive(Debug, Clone, Serialize)]
#[serde(tag = "type", rename_all = "camelCase")] #[serde(tag = "type", rename_all = "camelCase")]
pub enum McpServerConfig { pub enum McpServerConfig {
#[serde(rename = "http")] #[serde(rename = "http")]
Http { Http { url: String },
url: String,
#[serde(skip_serializing_if = "Option::is_none")]
scope: Option<String>,
#[serde(rename = "scopeId", skip_serializing_if = "Option::is_none")]
scope_id: Option<String>,
#[serde(rename = "scopeName", skip_serializing_if = "Option::is_none")]
scope_name: Option<String>,
},
#[serde(rename = "stdio")] #[serde(rename = "stdio")]
Stdio { Stdio {
command: std::path::PathBuf, command: std::path::PathBuf,
@@ -115,8 +107,6 @@ pub enum McpServerConfig {
#[serde(default, skip_serializing_if = "Vec::is_empty")] #[serde(default, skip_serializing_if = "Vec::is_empty")]
env: Vec<McpEnvVar>, env: Vec<McpEnvVar>,
}, },
#[serde(rename = "managedGateway")]
ManagedGateway,
} }
#[derive(Debug, Clone, Serialize)] #[derive(Debug, Clone, Serialize)]
@@ -128,7 +118,6 @@ pub struct McpEnvVar {
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")] #[serde(rename_all = "lowercase")]
pub enum McpServerSource { pub enum McpServerSource {
Managed,
Local, Local,
} }
@@ -306,14 +295,13 @@ pub struct McpReadResourceContent {
pub meta: Option<serde_json::Value>, pub meta: Option<serde_json::Value>,
} }
/// Push the full MCP catalog to the client. Called in the background after /// Push the full MCP catalog to the client. Called in the background so
/// managed configs resolve so `initialize()` isn't blocked by the network fetch. /// `initialize()` isn't blocked by catalog construction.
pub async fn notify_servers_updated( pub async fn notify_servers_updated(
gateway: &kigi_acp_lib::AcpAgentGatewaySender, gateway: &kigi_acp_lib::AcpAgentGatewaySender,
managed_configs: &[crate::session::managed_mcp::ManagedMcpConfig],
local_servers: &[acp::McpServer], local_servers: &[acp::McpServer],
) { ) {
let catalog = build_mcp_catalog(managed_configs, local_servers); let catalog = build_mcp_catalog(local_servers);
let payload = McpServersUpdated { let payload = McpServersUpdated {
mcp_servers: catalog, mcp_servers: catalog,
}; };
@@ -389,115 +377,22 @@ fn mcp_server_url(server: &acp::McpServer) -> Option<&str> {
} }
} }
/// Build MCP server catalog: managed + local servers, deduplicated by name. /// Build MCP server catalog from local servers, deduplicated by name.
/// Pure function — no I/O. Used by `mcp/list`, `InitializeResponse._meta`, /// Pure function — no I/O. Used by `mcp/list`, `InitializeResponse._meta`,
/// and `mcp/servers_updated`. /// and `mcp/servers_updated`.
pub fn build_mcp_catalog( pub fn build_mcp_catalog(local_servers: &[acp::McpServer]) -> Vec<McpServerEntry> {
managed_configs: &[crate::session::managed_mcp::ManagedMcpConfig],
local_servers: &[acp::McpServer],
) -> Vec<McpServerEntry> {
build_mcp_catalog_with_gateway_tools(managed_configs, local_servers, None, &Default::default())
}
pub fn build_mcp_catalog_with_gateway_tools(
managed_configs: &[crate::session::managed_mcp::ManagedMcpConfig],
local_servers: &[acp::McpServer],
gateway_catalog: Option<&crate::session::managed_mcp::GatewayToolCatalog>,
disabled_tools: &HashMap<String, HashSet<String>>,
) -> Vec<McpServerEntry> {
let mut servers: Vec<McpServerEntry> = Vec::new(); let mut servers: Vec<McpServerEntry> = Vec::new();
let mut seen = std::collections::HashSet::new(); let mut seen = std::collections::HashSet::new();
// Managed servers (always HTTP)
for config in managed_configs {
let name = crate::session::managed_mcp::to_managed_name(&config.name);
if seen.insert(name.clone()) {
servers.push(McpServerEntry {
name,
display_name: None,
source: McpServerSource::Managed,
config: McpServerConfig::Http {
url: config.endpoint.clone(),
scope: config.scope.clone(),
scope_id: config.scope_id.clone(),
scope_name: config.scope_name.clone(),
},
source_label: None,
session: None,
});
}
}
if let Some(catalog) = gateway_catalog {
let reauth: HashSet<&str> = catalog
.connectors_needing_reauth
.iter()
.map(String::as_str)
.collect();
let mut by_connector: BTreeMap<&str, Vec<&crate::session::managed_mcp::GatewayTool>> =
BTreeMap::new();
for tool in &catalog.tools {
by_connector
.entry(tool.connector_id.as_str())
.or_default()
.push(tool);
}
for (connector_id, tools) in by_connector {
let connector_name = tools
.first()
.map(|tool| tool.connector_name.as_str())
.unwrap_or(connector_id);
let disabled = disabled_tools.get(connector_id);
let server_disabled = disabled_tools
.get(crate::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY)
.is_some_and(|set| set.contains(connector_id));
let auth_required = reauth.contains(connector_id) || reauth.contains(connector_name);
servers.push(McpServerEntry {
name: managed_gateway_entry_name(connector_id),
display_name: Some(connector_name.to_owned()),
source: McpServerSource::Managed,
config: McpServerConfig::ManagedGateway,
source_label: None,
session: Some(McpServerSessionState {
enabled: !server_disabled,
status: (!auth_required && !server_disabled).then_some(McpSessionStatus::Ready),
tools: tools
.into_iter()
.map(|tool| {
let qualified_name = tool.qualified_name();
McpToolEntry {
name: qualified_name.clone(),
display_name: Some(tool.tool_name.clone()),
description: Some(tool.description.clone()),
meta: None,
enabled: disabled.is_none_or(|set| !set.contains(&qualified_name)),
}
})
.collect(),
auth_required,
}),
});
}
}
// Local servers (HTTP or Stdio) // Local servers (HTTP or Stdio)
for server in local_servers { for server in local_servers {
let name = crate::session::mcp_servers::mcp_server_name(server).to_string(); let name = crate::session::mcp_servers::mcp_server_name(server).to_string();
if seen.insert(name.clone()) { if seen.insert(name.clone()) {
let source = if name.starts_with(MANAGED_MCP_PREFIX) {
McpServerSource::Managed
} else {
McpServerSource::Local
};
let config = match server { let config = match server {
acp::McpServer::Http(acp::McpServerHttp { url, .. }) acp::McpServer::Http(acp::McpServerHttp { url, .. })
| acp::McpServer::Sse(acp::McpServerSse { url, .. }) => McpServerConfig::Http { | acp::McpServer::Sse(acp::McpServerSse { url, .. }) => {
url: url.clone(), McpServerConfig::Http { url: url.clone() }
scope: None, }
scope_id: None,
scope_name: None,
},
acp::McpServer::Stdio(acp::McpServerStdio { acp::McpServer::Stdio(acp::McpServerStdio {
command, args, env, .. command, args, env, ..
}) => McpServerConfig::Stdio { }) => McpServerConfig::Stdio {
@@ -517,7 +412,7 @@ pub fn build_mcp_catalog_with_gateway_tools(
servers.push(McpServerEntry { servers.push(McpServerEntry {
name, name,
display_name: None, display_name: None,
source, source: McpServerSource::Local,
config, config,
source_label: None, source_label: None,
session: None, session: None,
@@ -528,51 +423,17 @@ pub fn build_mcp_catalog_with_gateway_tools(
servers servers
} }
pub const MANAGED_GATEWAY_ENTRY_PREFIX: &str = "managed_gateway:";
fn managed_gateway_entry_name(connector_id: &str) -> String {
format!("{MANAGED_GATEWAY_ENTRY_PREFIX}{connector_id}")
}
fn managed_gateway_connector_id(entry_name: &str) -> Option<&str> {
entry_name.strip_prefix(MANAGED_GATEWAY_ENTRY_PREFIX)
}
fn should_append_disabled_mcp_placeholder(
name: &str,
catalog_names: &std::collections::HashSet<String>,
gateway_tools_enabled: bool,
) -> bool {
if catalog_names.contains(name) {
return false;
}
!gateway_tools_enabled
}
fn disabled_server_placeholder_entry(name: &str) -> McpServerEntry { fn disabled_server_placeholder_entry(name: &str) -> McpServerEntry {
let is_managed_gateway = name.starts_with(MANAGED_GATEWAY_ENTRY_PREFIX); McpServerEntry {
let source = if is_managed_gateway || name.starts_with(MANAGED_MCP_PREFIX) { name: name.to_owned(),
McpServerSource::Managed display_name: None,
} else { source: McpServerSource::Local,
McpServerSource::Local source_label: None,
}; config: McpServerConfig::Stdio {
let config = if is_managed_gateway {
McpServerConfig::ManagedGateway
} else {
McpServerConfig::Stdio {
command: std::path::PathBuf::new(), command: std::path::PathBuf::new(),
args: Vec::new(), args: Vec::new(),
env: Vec::new(), env: Vec::new(),
} },
};
McpServerEntry {
name: name.to_owned(),
display_name: name
.strip_prefix(MANAGED_GATEWAY_ENTRY_PREFIX)
.map(str::to_owned),
source,
source_label: None,
config,
session: Some(McpServerSessionState { session: Some(McpServerSessionState {
enabled: false, enabled: false,
status: None, status: None,
@@ -876,13 +737,6 @@ pub async fn call_mcp_tool(
// ── mcp/list handler ──────────────────────────────────────────────── // ── mcp/list handler ────────────────────────────────────────────────
async fn handle_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { async fn handle_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
// Latency layout: the two costly awaits — the managed-MCP
// proxy fetch (~1-2s) and the session-state branch (conditional
// `retry_auth_required_servers` followed by `build_mcp_status`, the
// latter cheap since is_healthy is a state-mutex inspection) —
// are independent and now run concurrently via tokio::join!. OAuth
// retries only fire on explicit refresh (cache=false); cached opens
// skip them so the warm path stays fast.
let req = parse_params::<McpListRequest>(args)?; let req = parse_params::<McpListRequest>(args)?;
let cwd = req let cwd = req
@@ -891,19 +745,6 @@ async fn handle_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
.and_then(|sid| agent.get_session_cwd(&acp::SessionId::new(sid.clone()))) .and_then(|sid| agent.get_session_cwd(&acp::SessionId::new(sid.clone())))
.unwrap_or_else(|| std::env::current_dir().unwrap_or_default()); .unwrap_or_else(|| std::env::current_dir().unwrap_or_default());
// NOTE: `invalidate_cache` must remain O(μs) (in-memory `Mutex` clear)
// so this serial pre-step does not eat into the latency budget. If
// it ever grows IO (fsync, contended lock, network), fold it into the
// managed-fetch arm of the `tokio::join!` below instead of keeping it
// here — otherwise the cache=false path silently re-introduces the
// sequential ~500ms+ gap the concurrent layout removed.
if !req.cache {
crate::session::managed_mcp::invalidate_cache(agent.managed_mcp_cache()).await;
crate::session::managed_mcp::invalidate_gateway_tool_cache(agent.managed_mcp_cache()).await;
}
// Resolve the session handle synchronously up front so the session-state
// future can be polled alongside the managed-MCP proxy fetch.
let session_handle = req.session_id.as_ref().and_then(|sid| { let session_handle = req.session_id.as_ref().and_then(|sid| {
let acp_id = acp::SessionId::new(sid.clone()); let acp_id = acp::SessionId::new(sid.clone());
agent.get_session_handle(&acp_id) agent.get_session_handle(&acp_id)
@@ -915,39 +756,22 @@ async fn handle_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
); );
} }
let cache = req.cache; let session_snapshot = match session_handle.as_ref() {
let session_state_fut = async { Some(handle) => {
let handle = session_handle.as_ref()?; // Auth retries belong on explicit refresh: skipping them on
// Auth retries belong on explicit refresh: skipping them on cached // cached opens saves ~500ms when multiple OAuth servers are
// opens saves ~500ms when multiple OAuth servers are configured. // configured.
if !cache { if !req.cache {
handle.retry_auth_required_servers().await; handle.retry_auth_required_servers().await;
} }
Some(handle.get_mcp_status().await) Some(handle.get_mcp_status().await)
};
let gateway_tools_enabled = agent.cfg.borrow().managed_mcp_gateway_tools_enabled;
let (managed_configs, gateway_catalog, session_snapshot) = tokio::join!(
agent.get_managed_mcp_configs(),
async {
if gateway_tools_enabled {
agent.get_managed_mcp_gateway_tool_catalog().await
} else {
None
} }
}, None => None,
session_state_fut };
);
let local_servers = let local_servers =
crate::util::config::load_mcp_servers(&cwd, &agent.cfg.borrow().compat_resolved); crate::util::config::load_mcp_servers(&cwd, &agent.cfg.borrow().compat_resolved);
let disabled_tools = crate::util::config::get_all_mcp_disabled_tools(&cwd); let mut servers = build_mcp_catalog(&local_servers);
let mut servers = build_mcp_catalog_with_gateway_tools(
&managed_configs,
&local_servers,
gateway_catalog.as_ref(),
&disabled_tools,
);
// Include disabled servers from config so they appear in the list // Include disabled servers from config so they appear in the list
// with enabled=false and can be re-enabled by the user. // with enabled=false and can be re-enabled by the user.
@@ -955,42 +779,18 @@ async fn handle_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
let catalog_names: std::collections::HashSet<String> = let catalog_names: std::collections::HashSet<String> =
servers.iter().map(|s| s.name.clone()).collect(); servers.iter().map(|s| s.name.clone()).collect();
for name in &disabled_names { for name in &disabled_names {
if should_append_disabled_mcp_placeholder(name, &catalog_names, gateway_tools_enabled) { if !catalog_names.contains(name) {
servers.push(disabled_server_placeholder_entry(name)); servers.push(disabled_server_placeholder_entry(name));
} }
} }
if let Some(snapshot) = session_snapshot { if let Some(snapshot) = session_snapshot {
if gateway_catalog.is_some()
&& let Some(disabled) = match session_handle.as_ref() {
Some(h) => Some(h.managed_gateway_disabled_tool_names().await),
None => None,
}
{
for entry in &mut servers {
if entry.source == McpServerSource::Managed
&& let Some(session) = entry.session.as_mut()
{
let connector_id =
managed_gateway_connector_id(&entry.name).unwrap_or(&entry.name);
if let Some(tools) = disabled.get(connector_id) {
for tool in &mut session.tools {
if tools.contains(&tool.name) {
tool.enabled = false;
}
}
}
}
}
}
// `session_snapshot` is `Some` only when `session_handle` resolved, // `session_snapshot` is `Some` only when `session_handle` resolved,
// which requires `req.session_id` to have been `Some`. Rather than // which requires `req.session_id` to have been `Some`. Rather than
// assert that non-local invariant with `expect` (which a future // assert that non-local invariant with `expect` (which a future
// refactor of `session_state_fut` could silently turn into a panic // refactor could silently turn into a panic in a request handler),
// in a request handler), use a local `if let` guard around the only // use a local `if let` guard around the only consumer — the debug
// consumer — the debug log. We emit `%sid` (Display) to match the // log.
// sibling "session not found" log; `?req.session_id` would wrap the
// bare string as `Some("...")` and diverge from the earlier format.
if let Some(sid) = req.session_id.as_ref() { if let Some(sid) = req.session_id.as_ref() {
tracing::debug!(session_id = %sid, "Annotating mcp/list with session state"); tracing::debug!(session_id = %sid, "Annotating mcp/list with session state");
} }
@@ -999,19 +799,6 @@ async fn handle_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
// Annotate catalog entries with session state. // Annotate catalog entries with session state.
for entry in &mut servers { for entry in &mut servers {
let managed_gateway_session = entry.source == McpServerSource::Managed
&& matches!(&entry.config, McpServerConfig::ManagedGateway);
if managed_gateway_session {
if let Some(session) = entry.session.as_mut() {
let connector_id =
managed_gateway_connector_id(&entry.name).unwrap_or(&entry.name);
let managed_disabled = disabled_tools
.get(crate::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY)
.is_some_and(|set| set.contains(connector_id));
session.enabled = !disabled_names.contains(&entry.name) && !managed_disabled;
}
continue;
}
let enabled = snapshot let enabled = snapshot
.configs .configs
.iter() .iter()
@@ -1444,38 +1231,13 @@ async fn handle_toggle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
.get_session_handle(&acp_id) .get_session_handle(&acp_id)
.ok_or_else(|| acp::Error::invalid_params().data("session not found"))?; .ok_or_else(|| acp::Error::invalid_params().data("session not found"))?;
let gateway_connector_id = managed_gateway_connector_id(&req.server_name); // Build the server config outside the session actor. The actual config
// mutation happens atomically inside the session actor via
// Build the server config outside the session actor (may need async I/O // ToggleMcpServer.
// for managed config fetch). The actual config mutation happens atomically
// inside the session actor via ToggleMcpServer.
let server_config = if req.enabled { let server_config = if req.enabled {
let cwd = agent let cwd = agent
.get_session_cwd(&acp_id) .get_session_cwd(&acp_id)
.unwrap_or_else(|| std::env::current_dir().unwrap_or_default()); .unwrap_or_else(|| std::env::current_dir().unwrap_or_default());
if let Some(connector_id) = gateway_connector_id {
if let Err(e) =
crate::util::config::save_mcp_server_enabled(&req.server_name, true).await
{
tracing::warn!(
server = req.server_name.as_str(),
error = %e,
"Failed to clear disabled MCP server entry for managed gateway connector"
);
}
handle
.toggle_managed_gateway_tool(connector_id.to_string(), String::new(), true)
.await
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
return to_ext_response(Ok(McpToggleResponse { ok: true }));
}
if req
.server_name
.starts_with(crate::session::managed_mcp::MANAGED_MCP_PREFIX)
{
crate::session::managed_mcp::invalidate_cache(agent.managed_mcp_cache()).await;
}
let managed_configs = agent.get_managed_mcp_configs().await;
if let Err(e) = crate::util::config::save_mcp_server_enabled(&req.server_name, true).await { if let Err(e) = crate::util::config::save_mcp_server_enabled(&req.server_name, true).await {
tracing::warn!( tracing::warn!(
server = req.server_name.as_str(), server = req.server_name.as_str(),
@@ -1488,7 +1250,6 @@ async fn handle_toggle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
crate::session::managed_mcp::merge_managed_mcp_servers_with_policy( crate::session::managed_mcp::merge_managed_mcp_servers_with_policy(
vec![], vec![],
&cwd, &cwd,
&managed_configs,
agent.plugin_registry_snapshot().as_deref(), agent.plugin_registry_snapshot().as_deref(),
&agent.cfg.borrow().compat_resolved, &agent.cfg.borrow().compat_resolved,
); );
@@ -1497,15 +1258,11 @@ async fn handle_toggle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
.find(|s| crate::session::mcp_servers::mcp_server_name(&s.server) == req.server_name); .find(|s| crate::session::mcp_servers::mcp_server_name(&s.server) == req.server_name);
match found { match found {
Some(s) if s.disabled_reason.is_some() => { Some(s) if s.disabled_reason.is_some() => {
let display = req
.server_name
.strip_prefix(crate::session::managed_mcp::MANAGED_MCP_PREFIX)
.unwrap_or(&req.server_name);
// Capitalize first letter for display. // Capitalize first letter for display.
let mut chars = display.chars(); let mut chars = req.server_name.chars();
let capitalized: String = match chars.next() { let capitalized: String = match chars.next() {
Some(c) => c.to_uppercase().chain(chars).collect(), Some(c) => c.to_uppercase().chain(chars).collect(),
None => display.to_string(), None => req.server_name.clone(),
}; };
let path = match &s.disabled_reason { let path = match &s.disabled_reason {
Some( Some(
@@ -1525,12 +1282,6 @@ async fn handle_toggle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
_ => {} _ => {}
} }
found.map(|s| s.server) found.map(|s| s.server)
} else if let Some(connector_id) = gateway_connector_id {
handle
.toggle_managed_gateway_tool(connector_id.to_string(), String::new(), false)
.await
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
return to_ext_response(Ok(McpToggleResponse { ok: true }));
} else { } else {
None None
}; };
@@ -1561,24 +1312,9 @@ async fn handle_toggle_tool(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResu
.get_session_handle(&acp_id) .get_session_handle(&acp_id)
.ok_or_else(|| acp::Error::invalid_params().data("session not found"))?; .ok_or_else(|| acp::Error::invalid_params().data("session not found"))?;
// `managed_gateway:` is reserved, so route by prefix alone — never consult
// the catalog, or a stale tool toggle would fall back to the local path.
let gateway_connector_id = managed_gateway_connector_id(&req.server_name);
let is_managed_gateway = gateway_connector_id.is_some();
if is_managed_gateway {
handle
.toggle_managed_gateway_tool(
gateway_connector_id.unwrap_or(&req.server_name).to_string(),
req.tool_name,
req.enabled,
)
.await
} else {
handle handle
.toggle_mcp_tool(req.server_name, req.tool_name, req.enabled) .toggle_mcp_tool(req.server_name, req.tool_name, req.enabled)
.await .await
}
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?; .map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
to_ext_response(Ok(McpToggleResponse { ok: true })) to_ext_response(Ok(McpToggleResponse { ok: true }))
@@ -1687,149 +1423,16 @@ mod tests {
assert_eq!(route_mcp_method(wire::MCP_CALL), Some(McpRoute::Call)); assert_eq!(route_mcp_method(wire::MCP_CALL), Some(McpRoute::Call));
} }
fn gateway_tool(
connector_id: &str,
connector_name: &str,
tool_id: &str,
tool_name: &str,
call_id: &str,
description: &str,
) -> crate::session::managed_mcp::GatewayTool {
crate::session::managed_mcp::GatewayTool {
connector_id: connector_id.into(),
connector_name: connector_name.into(),
tool_id: tool_id.into(),
tool_name: tool_name.into(),
call_id: call_id.into(),
description: description.into(),
json_schema: serde_json::json!({"type": "object"}),
}
}
/// **Pattern-regression test, not an end-to-end `handle_list` test.**
///
/// `handle_list` takes an `&MvpAgent`, which has no lightweight test
/// constructor; spinning up a fake agent here would be a much larger
/// refactor than this test warrants. Instead this test mirrors the exact
/// production structure (resolve session handle synchronously, then
/// `tokio::join!` a managed-fetch arm with a session-state arm whose
/// inner future conditionally awaits `retry_auth_required_servers` then
/// `build_mcp_status`) using stand-in futures, and asserts the two
/// latency invariants `handle_list` guarantees:
///
/// 1. The two `tokio::join!` arms — `get_managed_mcp_configs` on one
/// side, and the session-state branch (`retry_auth_required_servers?`
/// + `build_mcp_status`) on the other — are polled concurrently, so
/// total wall-time ≈ max(t_managed, t_session) rather than the sum.
/// 2. `retry_auth_required_servers` is gated on `cache=false`. On cached
/// opens it is skipped entirely, removing ~500ms of OAuth retry
/// overhead when multiple OAuth servers are configured.
///
/// If a future refactor of `handle_list` changes the structure (e.g.
/// awaits the arms sequentially, or runs auth retry on cache=true),
/// this test will *not* fail — it only guards the pattern. The real
/// behavioural guard is reading the diff against the structure
/// documented here.
#[tokio::test(start_paused = true)]
async fn handle_list_parallel_join_pattern_regression() {
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use tokio::time::{Duration, Instant};
async fn run(cache: bool) -> (Duration, bool, usize) {
let auth_retried = Arc::new(AtomicBool::new(false));
let max_concurrent = Arc::new(AtomicUsize::new(0));
let in_flight = Arc::new(AtomicUsize::new(0));
let bump = {
let max_concurrent = Arc::clone(&max_concurrent);
let in_flight = Arc::clone(&in_flight);
move || {
let n = in_flight.fetch_add(1, Ordering::SeqCst) + 1;
max_concurrent.fetch_max(n, Ordering::SeqCst);
}
};
let drop_ = {
let in_flight = Arc::clone(&in_flight);
move || {
in_flight.fetch_sub(1, Ordering::SeqCst);
}
};
// Stand-in for `agent.get_managed_mcp_configs()` (~1-2s proxy fetch).
let managed_fut = {
let bump = bump.clone();
let drop_ = drop_.clone();
async move {
bump();
tokio::time::sleep(Duration::from_millis(1500)).await;
drop_();
}
};
// Stand-in for the session-state branch: conditional auth retry
// followed by `build_mcp_status`. Mirrors the closure in
// `handle_list`.
let session_fut = {
let auth_retried = Arc::clone(&auth_retried);
let bump = bump.clone();
let drop_ = drop_.clone();
async move {
bump();
if !cache {
auth_retried.store(true, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(500)).await;
}
// build_mcp_status is cheap (state-mutex inspect).
tokio::time::sleep(Duration::from_millis(50)).await;
drop_();
}
};
let start = Instant::now();
tokio::join!(managed_fut, session_fut);
(
start.elapsed(),
auth_retried.load(Ordering::SeqCst),
max_concurrent.load(Ordering::SeqCst),
)
}
// cache=true: no auth retry, total ≈ managed fetch alone.
let (cached_elapsed, cached_auth, cached_overlap) = run(true).await;
assert!(!cached_auth, "auth retry must be skipped on cache=true");
assert_eq!(cached_overlap, 2, "futures must run concurrently");
assert!(
cached_elapsed < Duration::from_millis(1600),
"cached handle_list should finish in ~1.5s, got {:?}",
cached_elapsed
);
// cache=false: auth retry runs, but still concurrent with managed
// fetch — total ≈ max(1500, 500+50) ≈ 1500ms, not 2050ms.
let (refresh_elapsed, refresh_auth, refresh_overlap) = run(false).await;
assert!(refresh_auth, "auth retry must run on cache=false");
assert_eq!(refresh_overlap, 2, "futures must run concurrently");
assert!(
refresh_elapsed < Duration::from_millis(1600),
"refresh handle_list should still finish in ~1.5s (parallel), got {:?}",
refresh_elapsed
);
}
#[test] #[test]
fn test_mcp_list_response_serialization() { fn test_mcp_list_response_serialization() {
let resp = McpListResponse { let resp = McpListResponse {
servers: vec![ servers: vec![
McpServerEntry { McpServerEntry {
name: "grok_com_linear".to_string(), name: "linear".to_string(),
display_name: None, display_name: None,
source: McpServerSource::Managed, source: McpServerSource::Local,
config: McpServerConfig::Http { config: McpServerConfig::Http {
url: "https://mcp.linear.app".to_string(), url: "https://mcp.linear.app".to_string(),
scope: Some("team".to_string()),
scope_id: Some("team-uuid-123".to_string()),
scope_name: Some("Grok CLI".to_string()),
}, },
source_label: None, source_label: None,
session: None, session: None,
@@ -1860,34 +1463,11 @@ mod tests {
], ],
}; };
let json = serde_json::to_value(&resp).unwrap(); let json = serde_json::to_value(&resp).unwrap();
// [0] managed HTTP // [0] local HTTP
assert_eq!(json["servers"][0]["source"], "managed"); assert_eq!(json["servers"][0]["source"], "local");
assert_eq!(json["servers"][0]["type"], "http"); assert_eq!(json["servers"][0]["type"], "http");
assert_eq!(json["servers"][0]["url"], "https://mcp.linear.app"); assert_eq!(json["servers"][0]["url"], "https://mcp.linear.app");
assert_eq!(json["servers"][0]["scope"], "team");
assert_eq!(json["servers"][0]["scopeId"], "team-uuid-123");
assert_eq!(json["servers"][0]["scopeName"], "Grok CLI");
assert!(json["servers"][0].get("session").is_none()); assert!(json["servers"][0].get("session").is_none());
// Managed gateway connectors are not serialized as local transports.
let gateway = serde_json::to_value(McpServerEntry {
name: managed_gateway_entry_name("linear"),
display_name: Some("linear".to_string()),
source: McpServerSource::Managed,
source_label: None,
config: McpServerConfig::ManagedGateway,
session: Some(McpServerSessionState {
enabled: true,
status: Some(McpSessionStatus::Ready),
tools: vec![],
auth_required: false,
}),
})
.unwrap();
assert_eq!(gateway["name"], "managed_gateway:linear");
assert_eq!(gateway["displayName"], "linear");
assert_eq!(gateway["type"], "managedGateway");
assert!(gateway.get("command").is_none());
assert!(gateway.get("url").is_none());
// [1] local Stdio // [1] local Stdio
assert_eq!(json["servers"][1]["source"], "local"); assert_eq!(json["servers"][1]["source"], "local");
assert_eq!(json["servers"][1]["type"], "stdio"); assert_eq!(json["servers"][1]["type"], "stdio");
@@ -1905,236 +1485,6 @@ mod tests {
); );
} }
#[test]
fn gateway_catalog_groups_by_connector_name_and_exact_tool_names() {
let catalog = crate::session::managed_mcp::GatewayToolCatalog {
tools: vec![
gateway_tool(
"linear",
"Linear",
"list_issues",
"List issues",
"linear.list_issues",
"List Linear issues",
),
gateway_tool(
"linear",
"Linear",
"create_issue",
"Create issue",
"linear.create_issue",
"Create a Linear issue",
),
gateway_tool(
"slack",
"Slack",
"search",
"Search",
"slack.search",
"Search Slack",
),
],
total_tools: 3,
connectors_needing_reauth: vec!["slack".into()],
};
let servers =
build_mcp_catalog_with_gateway_tools(&[], &[], Some(&catalog), &Default::default());
assert_eq!(servers.len(), 2);
assert_eq!(servers[0].name, "managed_gateway:linear");
assert_eq!(servers[0].display_name.as_deref(), Some("Linear"));
assert_eq!(servers[0].source, McpServerSource::Managed);
assert!(matches!(servers[0].config, McpServerConfig::ManagedGateway));
let linear_session = servers[0].session.as_ref().unwrap();
assert_eq!(linear_session.status, Some(McpSessionStatus::Ready));
assert!(!linear_session.auth_required);
let linear_names: Vec<&str> = linear_session
.tools
.iter()
.map(|tool| tool.name.as_str())
.collect();
assert_eq!(
linear_names,
vec!["linear__list_issues", "linear__create_issue"]
);
assert_eq!(servers[1].name, "managed_gateway:slack");
assert_eq!(servers[1].display_name.as_deref(), Some("Slack"));
let slack_session = servers[1].session.as_ref().unwrap();
assert!(slack_session.auth_required);
assert!(slack_session.status.is_none());
assert_eq!(slack_session.tools[0].name, "slack__search");
assert_eq!(
slack_session.tools[0].display_name.as_deref(),
Some("Search")
);
}
#[test]
fn gateway_catalog_preserves_local_name_collision() {
let catalog = crate::session::managed_mcp::GatewayToolCatalog {
tools: vec![gateway_tool(
"linear",
"Linear",
"list_issues",
"List issues",
"linear.list_issues",
"List Linear issues",
)],
total_tools: 1,
connectors_needing_reauth: vec![],
};
let local = acp::McpServer::Stdio(
acp::McpServerStdio::new("linear", "/usr/bin/local-linear")
.args(vec![])
.env(vec![]),
);
let servers = build_mcp_catalog_with_gateway_tools(
&[],
&[local],
Some(&catalog),
&Default::default(),
);
assert_eq!(servers.len(), 2);
assert_eq!(servers[0].name, "managed_gateway:linear");
assert_eq!(servers[0].display_name.as_deref(), Some("Linear"));
assert_eq!(servers[0].source, McpServerSource::Managed);
assert_eq!(servers[1].name, "linear");
assert_eq!(servers[1].display_name, None);
assert_eq!(servers[1].source, McpServerSource::Local);
assert!(matches!(servers[1].config, McpServerConfig::Stdio { .. }));
}
#[test]
fn gateway_toggle_classification_requires_managed_gateway_entry_id() {
assert_eq!(
managed_gateway_connector_id("managed_gateway:linear"),
Some("linear")
);
assert_eq!(managed_gateway_connector_id("linear"), None);
}
#[test]
fn disabled_managed_http_rows_keep_non_gateway_placeholder_config() {
let entry = disabled_server_placeholder_entry("grok_com_slack");
assert_eq!(entry.source, McpServerSource::Managed);
assert!(matches!(entry.config, McpServerConfig::Stdio { .. }));
}
/// Mirrors `handle_list` set construction: catalog names from
/// `build_mcp_catalog_with_gateway_tools` (same inputs as production),
/// then disabled placeholders via `should_append_disabled_mcp_placeholder`.
fn append_disabled_like_handle_list(
servers: &mut Vec<McpServerEntry>,
disabled_names: &[&str],
gateway_tools_enabled: bool,
) {
let catalog_names: std::collections::HashSet<String> =
servers.iter().map(|s| s.name.clone()).collect();
for name in disabled_names {
if should_append_disabled_mcp_placeholder(name, &catalog_names, gateway_tools_enabled) {
servers.push(disabled_server_placeholder_entry(name));
}
}
}
#[test]
fn disabled_placeholders_match_handle_list_catalog_relationships() {
// Empty loads (gateway on, nothing in catalog) + orphan legacy disables
// only in disabled_mcp_servers — production ghost-stub regression.
let mut servers = build_mcp_catalog_with_gateway_tools(&[], &[], None, &Default::default());
append_disabled_like_handle_list(
&mut servers,
&["grok_com_slack", "grok_mcp_linear"],
true,
);
assert!(
servers.is_empty(),
"gateway on + no catalog rows → no stubs for orphan disables"
);
// Same orphans with gateway off → still placeholders (legacy UX).
let mut servers = build_mcp_catalog_with_gateway_tools(&[], &[], None, &Default::default());
append_disabled_like_handle_list(&mut servers, &["grok_com_slack"], false);
assert_eq!(servers.len(), 1);
assert_eq!(servers[0].name, "grok_com_slack");
assert!(!servers[0].session.as_ref().unwrap().enabled);
// Name already in catalog (gateway row) → never double-append.
let gateway = crate::session::managed_mcp::GatewayToolCatalog {
tools: vec![gateway_tool(
"linear",
"Linear",
"list_issues",
"List issues",
"linear.list_issues",
"List Linear issues",
)],
total_tools: 1,
connectors_needing_reauth: vec![],
};
let mut servers =
build_mcp_catalog_with_gateway_tools(&[], &[], Some(&gateway), &Default::default());
let gateway_entry = managed_gateway_entry_name("linear");
assert!(servers.iter().any(|s| s.name == gateway_entry));
let before = servers.len();
append_disabled_like_handle_list(&mut servers, &[gateway_entry.as_str()], true);
append_disabled_like_handle_list(&mut servers, &[gateway_entry.as_str()], false);
assert_eq!(
servers.len(),
before,
"disabled name already in catalog must not add a second row"
);
assert_eq!(
servers.iter().filter(|s| s.name == gateway_entry).count(),
1
);
}
#[test]
fn gateway_catalog_honors_disabled_connectors_and_tools() {
let catalog = crate::session::managed_mcp::GatewayToolCatalog {
tools: vec![
gateway_tool(
"linear",
"Linear",
"list_issues",
"List issues",
"linear.list_issues",
"List Linear issues",
),
gateway_tool(
"linear",
"Linear",
"create_issue",
"Create issue",
"linear.create_issue",
"Create a Linear issue",
),
],
total_tools: 2,
connectors_needing_reauth: vec![],
};
let disabled: HashMap<String, HashSet<String>> = HashMap::from([
(
crate::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY.to_string(),
HashSet::from(["linear".to_string()]),
),
(
"linear".to_string(),
HashSet::from(["linear__create_issue".to_string()]),
),
]);
let servers = build_mcp_catalog_with_gateway_tools(&[], &[], Some(&catalog), &disabled);
let session = servers[0].session.as_ref().unwrap();
assert!(!session.enabled);
assert!(session.status.is_none());
assert!(session.tools[0].enabled);
assert!(!session.tools[1].enabled);
}
#[test] #[test]
fn test_mcp_call_response_serialization() { fn test_mcp_call_response_serialization() {
let resp = McpCallResponse { let resp = McpCallResponse {
@@ -2192,15 +1542,12 @@ mod tests {
#[test] #[test]
fn test_disabled_session_state_serialization() { fn test_disabled_session_state_serialization() {
let entry = McpServerEntry { let entry = McpServerEntry {
name: "grok_com_slack".to_string(), name: "slack".to_string(),
display_name: None, display_name: None,
source: McpServerSource::Managed, source: McpServerSource::Local,
source_label: None, source_label: None,
config: McpServerConfig::Http { config: McpServerConfig::Http {
url: "https://mcp.slack.com".to_string(), url: "https://mcp.slack.com".to_string(),
scope: Some("user".to_string()),
scope_id: Some("user-uuid-456".to_string()),
scope_name: None,
}, },
session: Some(McpServerSessionState { session: Some(McpServerSessionState {
enabled: false, enabled: false,
@@ -2211,8 +1558,6 @@ mod tests {
}; };
let json = serde_json::to_value(&entry).unwrap(); let json = serde_json::to_value(&entry).unwrap();
assert_eq!(json["type"], "http"); assert_eq!(json["type"], "http");
assert_eq!(json["scope"], "user");
assert_eq!(json["scopeId"], "user-uuid-456");
assert_eq!(json["session"]["enabled"], false); assert_eq!(json["session"]["enabled"], false);
assert!(json["session"].get("status").is_none()); assert!(json["session"].get("status").is_none());
assert!(json["session"].get("tools").is_none()); assert!(json["session"].get("tools").is_none());
@@ -13,7 +13,6 @@
//! - `x.ai/internal/reload_skills` skills file watcher fan-out //! - `x.ai/internal/reload_skills` skills file watcher fan-out
//! - `x.ai/internal/reload_models` model list hot-reload from config.toml //! - `x.ai/internal/reload_models` model list hot-reload from config.toml
//! - `x.ai/internal/reload_models_cache` model catalog hot-reload from disk cache //! - `x.ai/internal/reload_models_cache` model catalog hot-reload from disk cache
//! - `x.ai/internal/auth_cleared` auth hot-clear cleanup
//! - `x.ai/plugins/reload` rebuild shared plugin registry //! - `x.ai/plugins/reload` rebuild shared plugin registry
//! - `x.ai/commands/list` list slash commands //! - `x.ai/commands/list` list slash commands
@@ -46,7 +45,6 @@ pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
"x.ai/internal/reload_skills" => handle_reload_skills(agent), "x.ai/internal/reload_skills" => handle_reload_skills(agent),
"x.ai/internal/reload_models" => handle_reload_models(agent), "x.ai/internal/reload_models" => handle_reload_models(agent),
"x.ai/internal/reload_models_cache" => handle_reload_models_cache(agent), "x.ai/internal/reload_models_cache" => handle_reload_models_cache(agent),
"x.ai/internal/auth_cleared" => handle_auth_cleared(agent),
"x.ai/plugins/reload" => handle_plugins_reload(agent).await, "x.ai/plugins/reload" => handle_plugins_reload(agent).await,
"x.ai/commands/list" => handle_commands_list(agent, args).await, "x.ai/commands/list" => handle_commands_list(agent, args).await,
_ => Err(acp::Error::method_not_found()), _ => Err(acp::Error::method_not_found()),
@@ -219,11 +217,9 @@ async fn handle_update_mcp_servers(agent: &MvpAgent, args: &acp::ExtRequest) ->
(h, cwd) (h, cwd)
}; };
let managed = agent.get_managed_mcp_configs().await;
let merged = crate::session::managed_mcp::merge_managed_mcp_servers( let merged = crate::session::managed_mcp::merge_managed_mcp_servers(
params.mcp_servers.clone(), params.mcp_servers.clone(),
&cwd, &cwd,
&managed,
agent.plugin_registry_handle().snapshot().as_deref(), agent.plugin_registry_handle().snapshot().as_deref(),
&agent.cfg.borrow().compat_resolved, &agent.cfg.borrow().compat_resolved,
); );
@@ -285,7 +281,6 @@ async fn handle_reload_all_mcp_servers(agent: &MvpAgent) -> ExtResult {
.map_err(|e| acp::Error::internal_error().data(e.to_string())); .map_err(|e| acp::Error::internal_error().data(e.to_string()));
} }
let managed = agent.get_managed_mcp_configs().await;
let mut updated = 0u32; let mut updated = 0u32;
for session_id in &session_ids { for session_id in &session_ids {
let Some(handle) = agent.sessions.borrow().get(session_id).cloned() else { let Some(handle) = agent.sessions.borrow().get(session_id).cloned() else {
@@ -294,8 +289,8 @@ async fn handle_reload_all_mcp_servers(agent: &MvpAgent) -> ExtResult {
let cwd = std::path::PathBuf::from(&handle.info.cwd); let cwd = std::path::PathBuf::from(&handle.info.cwd);
let compat = agent.cfg.borrow().compat_resolved; let compat = agent.cfg.borrow().compat_resolved;
// Re-seed the merge with the session's original client-provided MCP // Re-seed the merge with the session's original client-provided MCP
// servers (e.g. a managed connector injected at `session/new` by a // servers (e.g. a server injected at `session/new` by a client
// client session binding). `merge_managed_mcp_servers` already // session binding). `merge_managed_mcp_servers` already
// re-reads every disk source (config.toml, plugins, ~/.claude.json, // re-reads every disk source (config.toml, plugins, ~/.claude.json,
// ~/.cursor/mcp.json, .mcp.json) internally, so passing // ~/.cursor/mcp.json, .mcp.json) internally, so passing
// `load_mcp_servers()` output here was redundant — and silently // `load_mcp_servers()` output here was redundant — and silently
@@ -304,7 +299,6 @@ async fn handle_reload_all_mcp_servers(agent: &MvpAgent) -> ExtResult {
let merged = crate::session::managed_mcp::merge_managed_mcp_servers( let merged = crate::session::managed_mcp::merge_managed_mcp_servers(
handle.initial_client_mcp_servers.clone(), handle.initial_client_mcp_servers.clone(),
&cwd, &cwd,
&managed,
agent.plugin_registry_handle().snapshot().as_deref(), agent.plugin_registry_handle().snapshot().as_deref(),
&compat, &compat,
); );
@@ -369,7 +363,6 @@ async fn handle_reload_project_mcp_servers(agent: &MvpAgent, args: &acp::ExtRequ
.map_err(|e| acp::Error::internal_error().data(e.to_string())); .map_err(|e| acp::Error::internal_error().data(e.to_string()));
} }
let managed = agent.get_managed_mcp_configs().await;
let mut updated = 0u32; let mut updated = 0u32;
for (session_id, cwd) in &session_ids { for (session_id, cwd) in &session_ids {
let Some(handle) = agent.sessions.borrow().get(session_id).cloned() else { let Some(handle) = agent.sessions.borrow().get(session_id).cloned() else {
@@ -382,7 +375,6 @@ async fn handle_reload_project_mcp_servers(agent: &MvpAgent, args: &acp::ExtRequ
let merged = crate::session::managed_mcp::merge_managed_mcp_servers( let merged = crate::session::managed_mcp::merge_managed_mcp_servers(
handle.initial_client_mcp_servers.clone(), handle.initial_client_mcp_servers.clone(),
cwd, cwd,
&managed,
agent.plugin_registry_handle().snapshot().as_deref(), agent.plugin_registry_handle().snapshot().as_deref(),
&agent.cfg.borrow().compat_resolved, &agent.cfg.borrow().compat_resolved,
); );
@@ -497,13 +489,6 @@ fn handle_reload_models_cache(agent: &MvpAgent) -> ExtResult {
.map_err(|e| acp::Error::internal_error().data(e.to_string())) .map_err(|e| acp::Error::internal_error().data(e.to_string()))
} }
fn handle_auth_cleared(agent: &MvpAgent) -> ExtResult {
agent.disable_managed_gateway_tools_and_refresh_sessions();
ExtMethodResult::success(serde_json::json!({ "ok": true }))
.to_ext_response()
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
}
// plugins/reload // plugins/reload
async fn handle_plugins_reload(agent: &MvpAgent) -> ExtResult { async fn handle_plugins_reload(agent: &MvpAgent) -> ExtResult {
@@ -5297,20 +5297,16 @@ mod tests {
let (mut reader_a, _writer_a) = connect_and_register(&sock_path, "client-a").await; let (mut reader_a, _writer_a) = connect_and_register(&sock_path, "client-a").await;
let (mut reader_b, _writer_b) = connect_and_register(&sock_path, "client-b").await; let (mut reader_b, _writer_b) = connect_and_register(&sock_path, "client-b").await;
tokio::time::sleep(Duration::from_millis(20)).await; tokio::time::sleep(Duration::from_millis(20)).await;
let update = r#"{"jsonrpc":"2.0","method":"_x.ai/mcp/servers_updated","params":{"method":"x.ai/mcp/servers_updated","params":{"mcpServers":[{"name":"grok_com_slack","source":"managed"}]}}}"#; let update = r#"{"jsonrpc":"2.0","method":"_x.ai/mcp/servers_updated","params":{"method":"x.ai/mcp/servers_updated","params":{"mcpServers":[{"name":"team_slack","source":"local"}]}}}"#;
response_tx.send(update.to_string()).unwrap(); response_tx.send(update.to_string()).unwrap();
let got_a = next_acp_payload(&mut reader_a).await; let got_a = next_acp_payload(&mut reader_a).await;
let got_b = next_acp_payload(&mut reader_b).await; let got_b = next_acp_payload(&mut reader_b).await;
assert!( assert!(
got_a got_a.as_deref().is_some_and(|p| p.contains("team_slack")),
.as_deref()
.is_some_and(|p| p.contains("grok_com_slack")),
"client A must receive the MCP catalog broadcast, got {got_a:?}" "client A must receive the MCP catalog broadcast, got {got_a:?}"
); );
assert!( assert!(
got_b got_b.as_deref().is_some_and(|p| p.contains("team_slack")),
.as_deref()
.is_some_and(|p| p.contains("grok_com_slack")),
"client B must receive the MCP catalog broadcast, got {got_b:?}" "client B must receive the MCP catalog broadcast, got {got_b:?}"
); );
cancel.cancel(); cancel.cancel();
+1 -86
View File
@@ -7,8 +7,6 @@ use std::sync::Arc;
use kigi_tools::types::config_source::ConfigSource; use kigi_tools::types::config_source::ConfigSource;
use serde::Serialize; use serde::Serialize;
use crate::auth::KimiCodeConfig;
use crate::session::managed_mcp;
use crate::session::mcp_servers; use crate::session::mcp_servers;
// ── Report types ──────────────────────────────────────────────── // ── Report types ────────────────────────────────────────────────
@@ -238,85 +236,6 @@ fn discover_servers(cwd: &Path) -> (Vec<ConfigSourceStatus>, Vec<DiscoveredServe
(sources, servers) (sources, servers)
} }
// ── Managed (grok.com) server discovery ─────────────────────────
const MANAGED_SOURCE_LABEL: &str = "grok.com";
fn managed_skipped(reason: impl Into<String>) -> (ConfigSourceStatus, Vec<DiscoveredServer>) {
(
ConfigSourceStatus {
path: MANAGED_SOURCE_LABEL.to_string(),
status: ConfigSourceState::Skipped {
reason: reason.into(),
},
},
vec![],
)
}
fn managed_found(
count: usize,
servers: Vec<DiscoveredServer>,
) -> (ConfigSourceStatus, Vec<DiscoveredServer>) {
(
ConfigSourceStatus {
path: MANAGED_SOURCE_LABEL.to_string(),
status: ConfigSourceState::Found {
server_count: count,
},
},
servers,
)
}
/// Discover managed `grok_com_*` servers if the user has xAI auth on disk.
async fn try_discover_managed_servers() -> (ConfigSourceStatus, Vec<DiscoveredServer>) {
let kigi_home = kigi_tools::util::kigi_home::kigi_home();
let kimi_code_config = KimiCodeConfig::default();
let auth_manager = Arc::new(crate::auth::AuthManager::new(&kigi_home, kimi_code_config));
let Some(snapshot) = auth_manager.current_or_expired() else {
return managed_skipped("not logged in");
};
if !snapshot.is_session_auth() {
return managed_skipped(format!("{:?} auth (not xAI OIDC)", snapshot.auth_mode));
}
let token = match auth_manager.get_valid_token().await {
Ok(key) => key,
Err(_) => return managed_skipped("auth expired — run `kigi login`"),
};
let proxy_url = crate::agent::config::EndpointsConfig::from_effective_config().proxy_url();
let configs = match managed_mcp::fetch_managed_configs(&proxy_url, &token).await {
Ok(configs) => configs,
Err(e) => return managed_skipped(format!("fetch failed: {e}")),
};
if configs.is_empty() {
return managed_found(0, vec![]);
}
let mut servers: Vec<agent_client_protocol::McpServer> = vec![];
managed_mcp::auto_inject_managed_servers_with_disabled(
&mut servers,
&configs,
&Default::default(),
);
managed_mcp::inject_managed_headers(&mut servers, &configs);
let source = ConfigSource::Managed { path: None };
let discovered: Vec<DiscoveredServer> = servers
.into_iter()
.map(|server| DiscoveredServer {
server,
source: source.clone(),
})
.collect();
managed_found(discovered.len(), discovered)
}
// ── Check functions ───────────────────────────────────────────── // ── Check functions ─────────────────────────────────────────────
fn resolve_command(command: &str) -> Option<String> { fn resolve_command(command: &str) -> Option<String> {
@@ -558,11 +477,7 @@ pub async fn run_auth(cwd: &Path, name: &str) -> Result<usize, String> {
} }
pub async fn run_doctor(cwd: &Path, name_filter: Option<&str>) -> DoctorReport { pub async fn run_doctor(cwd: &Path, name_filter: Option<&str>) -> DoctorReport {
let (mut sources, mut discovered) = discover_servers(cwd); let (mut sources, discovered) = discover_servers(cwd);
let (managed_source, managed_servers) = try_discover_managed_servers().await;
sources.push(managed_source);
discovered.extend(managed_servers);
let allowlist = &kigi_workspace::permission::resolution::managed_settings().mcp_allowlist; let allowlist = &kigi_workspace::permission::resolution::managed_settings().mcp_allowlist;
if allowlist.is_restricted() { if allowlist.is_restricted() {
@@ -533,19 +533,6 @@ pub fn acp_tool_update(
)])) )]))
.raw_output(raw_output_json(output, rewriter)), .raw_output(raw_output_json(output, rewriter)),
)), )),
// Dual channel: prose for non-pager clients, typed `raw_output` for the pager.
ToolOutput::ImageGen(_)
| ToolOutput::ImageToVideo(_)
| ToolOutput::ReferenceToVideo(_)
| ToolOutput::ImageEdit(_) => Some(acp::ToolCallUpdate::new(
acp::ToolCallId::new(Arc::from(tool_call_id)),
acp::ToolCallUpdateFields::new()
.status(Some(acp::ToolCallStatus::Completed))
.content(Some(vec![acp::ToolCallContent::from(
acp::ContentBlock::Text(acp::TextContent::new(output.to_prompt_format())),
)]))
.raw_output(raw_output_json(output, rewriter)),
)),
ToolOutput::SubagentCompleted(sub) => { ToolOutput::SubagentCompleted(sub) => {
// Text includes resume handle for discoverability + meta for TUI. // Text includes resume handle for discoverability + meta for TUI.
// Shared with the chat-bidi server via `to_model_text` so both // Shared with the chat-bidi server via `to_model_text` so both
@@ -1035,34 +1022,6 @@ mod tests {
} }
} }
#[test]
fn test_media_gen_acp_update_emits_prose_and_raw_output() {
// Dual channel: prompt-format JSON in content, typed variant in raw_output.
let output = ToolOutput::ImageToVideo(MediaGenOutput::new(PathBuf::from(
"/tmp/session/videos/3.mp4",
)));
let update = acp_tool_update(&output, "tc-1", None, None).expect("update");
let content = update.fields.content.expect("content");
let text = match &content[0] {
acp::ToolCallContent::Content(acp::Content {
content: acp::ContentBlock::Text(t),
..
}) => t.text.clone(),
other => panic!("expected text content, got {other:?}"),
};
let prompt_json: serde_json::Value = serde_json::from_str(&text).expect("prompt json");
assert_eq!(prompt_json["path"], "/tmp/session/videos/3.mp4");
assert_eq!(prompt_json["filename"], "3.mp4");
assert_eq!(prompt_json["session_folder"], "videos");
assert_eq!(
prompt_json["message"],
"Video generated and saved to /tmp/session/videos/3.mp4. Do not read or re-display it, and do not describe how it appears to the user."
);
let raw = update.fields.raw_output.expect("raw_output");
assert_eq!(raw["type"], "ImageToVideo");
assert_eq!(raw["path"], "/tmp/session/videos/3.mp4");
}
#[test] #[test]
fn test_path_rewriter_rewrites_list_dir_raw_output() { fn test_path_rewriter_rewrites_list_dir_raw_output() {
let rw = PathRewriter::new( let rw = PathRewriter::new(
@@ -353,166 +353,6 @@ pub(crate) fn state_is_busy(state: &State) -> bool {
state.running_task.is_some() || !state.pending_inputs.is_empty() state.running_task.is_some() || !state.pending_inputs.is_empty()
} }
use crate::auth::AuthManager; use crate::auth::AuthManager;
#[derive(Clone)]
struct ShellManagedGatewayToolClient {
proxy_base_url: String,
auth_manager: Arc<AuthManager>,
}
#[async_trait::async_trait]
impl kigi_tools::types::resources::ManagedGatewayToolCaller for ShellManagedGatewayToolClient {
async fn call_tool(
&self,
call_id: &str,
arguments: serde_json::Value,
caller: &str,
) -> Result<
kigi_tools::types::resources::ManagedGatewayToolCallResponse,
kigi_tool_runtime::ToolError,
> {
let auth_key = self
.auth_manager
.get_valid_token()
.await
.ok()
.or_else(|| self.auth_manager.current_or_expired().map(|a| a.key))
.ok_or_else(|| kigi_tool_runtime::ToolError::unauthorized("no auth token available"))?;
let response = crate::session::managed_mcp::call_gateway_tool(
&self.proxy_base_url,
&auth_key,
call_id,
arguments,
)
.await
.map_err(|error| managed_gateway_error_to_tool_error(error, caller))?;
Ok(
kigi_tools::types::resources::ManagedGatewayToolCallResponse {
result: response.result,
connectors_needing_reauth: response.connectors_needing_reauth,
},
)
}
}
fn managed_gateway_error_to_tool_error(
error: crate::session::managed_mcp::ManagedMcpFetchError,
caller: &str,
) -> kigi_tool_runtime::ToolError {
match error {
crate::session::managed_mcp::ManagedMcpFetchError::Status { status, message } => {
let detail = format!("Managed MCP gateway tool call failed: {message}");
let mut err = if status == reqwest::StatusCode::UNAUTHORIZED {
kigi_tool_runtime::ToolError::unauthorized(detail)
} else if status == reqwest::StatusCode::FORBIDDEN {
kigi_tool_runtime::ToolError::permission_denied(detail)
} else {
let tool_id = kigi_tool_protocol::ToolId::new(caller).unwrap_or_else(|_| {
kigi_tool_protocol::ToolId::new("use_tool").expect("valid")
});
kigi_tool_runtime::ToolError::execution(tool_id, detail)
};
match err.details.as_mut() {
Some(serde_json::Value::Object(map)) => {
map.insert(
HTTP_STATUS_DETAILS_KEY.to_string(),
serde_json::json!(status.as_u16()),
);
}
_ => {
err.details =
Some(serde_json::json!({ HTTP_STATUS_DETAILS_KEY : status.as_u16(), }));
}
}
err
}
crate::session::managed_mcp::ManagedMcpFetchError::Transport(e) => {
kigi_tool_runtime::ToolError::network_error(format!(
"Managed MCP gateway tool call failed: {}",
e.without_url()
))
}
crate::session::managed_mcp::ManagedMcpFetchError::NoAuth => {
kigi_tool_runtime::ToolError::unauthorized("no auth token available")
}
}
}
#[cfg(test)]
mod managed_gateway_error_tests {
use super::*;
fn status_error(code: u16, message: &str) -> crate::session::managed_mcp::ManagedMcpFetchError {
crate::session::managed_mcp::ManagedMcpFetchError::Status {
status: reqwest::StatusCode::from_u16(code).unwrap(),
message: message.to_string(),
}
}
#[test]
fn unauthorized_status_maps_to_unauthorized_and_carries_status() {
let err = managed_gateway_error_to_tool_error(status_error(401, "expired"), "use_tool");
assert_eq!(err.kind, kigi_tool_runtime::ToolErrorKind::Unauthorized);
assert!(err.detail.contains("expired"));
let details = err.details.as_ref().unwrap();
assert_eq!(
details.get(HTTP_STATUS_DETAILS_KEY),
Some(&serde_json::json!(401))
);
}
#[test]
fn forbidden_status_maps_to_permission_denied_and_carries_status() {
let err = managed_gateway_error_to_tool_error(status_error(403, "denied"), "use_tool");
assert_eq!(err.kind, kigi_tool_runtime::ToolErrorKind::PermissionDenied);
let details = err.details.as_ref().unwrap();
assert_eq!(
details.get(HTTP_STATUS_DETAILS_KEY),
Some(&serde_json::json!(403))
);
}
#[test]
fn general_status_maps_to_execution_with_caller_tool_id() {
let err = managed_gateway_error_to_tool_error(status_error(500, "boom"), "CallMcpTool");
assert_eq!(err.kind, kigi_tool_runtime::ToolErrorKind::Execution);
let details = err.details.as_ref().unwrap();
assert_eq!(
details.get(HTTP_STATUS_DETAILS_KEY),
Some(&serde_json::json!(500))
);
assert_eq!(
details.get("tool_id"),
Some(&serde_json::json!("CallMcpTool"))
);
}
#[test]
fn general_status_falls_back_to_use_tool_for_unknown_caller() {
let err = managed_gateway_error_to_tool_error(status_error(500, "boom"), "not a tool id");
assert_eq!(err.kind, kigi_tool_runtime::ToolErrorKind::Execution);
let details = err.details.as_ref().unwrap();
assert_eq!(details.get("tool_id"), Some(&serde_json::json!("use_tool")));
}
#[test]
fn no_auth_maps_to_unauthorized() {
let err = managed_gateway_error_to_tool_error(
crate::session::managed_mcp::ManagedMcpFetchError::NoAuth,
"use_tool",
);
assert_eq!(err.kind, kigi_tool_runtime::ToolErrorKind::Unauthorized);
}
#[tokio::test]
async fn transport_error_maps_to_network_error_without_url() {
let transport = reqwest::Client::new()
.post("http://127.0.0.1:1/mcp/tools/call")
.send()
.await
.expect_err("connection to a dead port should fail");
let err = managed_gateway_error_to_tool_error(
crate::session::managed_mcp::ManagedMcpFetchError::Transport(transport),
"use_tool",
);
assert_eq!(err.kind, kigi_tool_runtime::ToolErrorKind::NetworkError);
assert!(err.detail.contains("Managed MCP gateway tool call failed"));
assert!(
!err.detail.contains("http://"),
"transport detail must not leak the proxy URL: {}",
err.detail
);
}
}
/// Data carried from prepare_tool_call → dispatch_tool → finalize. /// Data carried from prepare_tool_call → dispatch_tool → finalize.
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct PreparedToolCall { pub(crate) struct PreparedToolCall {
@@ -858,10 +698,6 @@ pub(crate) struct SessionActor {
/// flag is set short-circuits through /// flag is set short-circuits through
/// [`Self::account_not_achieved_without_sampler`]. /// [`Self::account_not_achieved_without_sampler`].
pub(crate) goal_classifier_in_flight: std::sync::atomic::AtomicBool, pub(crate) goal_classifier_in_flight: std::sync::atomic::AtomicBool,
/// Agent-level managed MCP config cache (refreshed in background).
pub(crate) managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle,
/// Earliest managed MCP token expiry; checked before tool dispatch.
pub(crate) managed_mcp_expires_at: std::sync::Mutex<Option<chrono::DateTime<chrono::Utc>>>,
/// Original client-provided MCP servers from session creation. /// Original client-provided MCP servers from session creation.
/// Retained for re-merge during plugin reload. /// Retained for re-merge during plugin reload.
pub(crate) initial_client_mcp_servers: Vec<acp::McpServer>, pub(crate) initial_client_mcp_servers: Vec<acp::McpServer>,
@@ -1315,200 +1151,6 @@ fn load_prompt_context_from_dir(
#[cfg(test)] #[cfg(test)]
#[path = "acp_session_tests/client_hooks_tests.rs"] #[path = "acp_session_tests/client_hooks_tests.rs"]
mod client_hooks_tests; mod client_hooks_tests;
#[cfg(test)]
#[path = "acp_session_tests/replace_system_prompt_tests.rs"]
mod replace_system_prompt_tests;
#[cfg(test)]
#[path = "acp_session_tests/support.rs"]
mod support;
#[cfg(test)]
#[path = "acp_session_tests/usage_categories_tests.rs"]
mod usage_categories_tests;
#[cfg(test)]
mod managed_gateway_descriptor_tests {
use super::*;
use kigi_tools::types::output::{MCPOutput, ToolOutput};
use kigi_tools::types::tool::{ToolKind, ToolNamespace};
#[derive(Debug, Default)]
struct FixtureMcpTool;
impl kigi_tools::types::tool_metadata::ToolMetadata for FixtureMcpTool {
fn kind(&self) -> ToolKind {
ToolKind::Other
}
fn tool_namespace(&self) -> ToolNamespace {
ToolNamespace::MCP
}
fn description_template(&self) -> &str {
"fixture"
}
}
impl kigi_tool_runtime::Tool for FixtureMcpTool {
type Args = serde_json::Value;
type Output = ToolOutput;
fn id(&self) -> kigi_tool_protocol::ToolId {
kigi_tool_protocol::ToolId::new("server__tool").expect("valid")
}
fn description(
&self,
_ctx: &::kigi_tool_runtime::ListToolsContext,
) -> kigi_tool_types::ToolDescription {
kigi_tool_types::ToolDescription::new("server__tool", "fixture")
}
async fn run(
&self,
_ctx: kigi_tool_runtime::ToolCallContext,
_args: serde_json::Value,
) -> Result<ToolOutput, kigi_tool_runtime::ToolError> {
Ok(ToolOutput::MCP(MCPOutput::okay_output(
"server__tool".to_string(),
"server".to_string(),
"ok".to_string(),
)))
}
}
#[tokio::test]
async fn refresh_snapshot_indexes_only_admitted_gateway_tools() {
let bridge = Arc::new(crate::tools::bridge::ToolBridge::for_test());
bridge
.register_mcp_tools(
"server__tool".to_string(),
FixtureMcpTool,
Some(serde_json::json!({ "type" : "object" })),
)
.await
.expect("local fixture registration succeeds");
let mcp_state = Arc::new(TokioMutex::new(McpState::new(vec![])));
let managed = crate::session::managed_mcp::ManagedMcpStateHandle::default();
{
let mut state = managed.lock().await;
state.enable_gateway_tools();
let epoch = state.start_gateway_tool_fetch().unwrap();
assert!(state.complete_gateway_tool_fetch(
epoch,
crate::session::managed_mcp::GatewayToolCatalog {
tools: vec![
crate::session::managed_mcp::GatewayTool {
connector_id: "server".to_string(),
connector_name: "Gateway Collision".to_string(),
tool_id: "tool".to_string(),
tool_name: "Collision".to_string(),
call_id: "gateway.collision".to_string(),
description: "Gateway collision".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
crate::session::managed_mcp::GatewayTool {
connector_id: "gateway".to_string(),
connector_name: "Gateway".to_string(),
tool_id: "search".to_string(),
tool_name: "Search".to_string(),
call_id: "gateway.search".to_string(),
description: "Gateway search".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
],
total_tools: 2,
connectors_needing_reauth: vec![],
}
));
}
let snapshot = Arc::new(std::sync::Mutex::new(
crate::session::tool_index::ToolMetadataSnapshot::default(),
));
refresh_mcp_snapshot_for_test(bridge, mcp_state, managed, snapshot.clone()).await;
let snapshot = snapshot.lock().unwrap();
let names: std::collections::HashSet<&str> = snapshot
.tools
.iter()
.map(|tool| tool.qualified_name.as_str())
.collect();
assert!(names.contains("gateway__search"));
let server_tool = snapshot
.tools
.iter()
.find(|tool| tool.qualified_name == "server__tool")
.expect("local MCP tool remains indexed");
assert_eq!(server_tool.server_name, "server");
assert_eq!(server_tool.description, "fixture");
}
#[tokio::test]
async fn refresh_snapshot_excludes_disabled_gateway_tools_and_connectors() {
let bridge = Arc::new(crate::tools::bridge::ToolBridge::for_test());
let mcp_state = Arc::new(TokioMutex::new(McpState::new(vec![])));
let managed = crate::session::managed_mcp::ManagedMcpStateHandle::default();
{
let mut state = managed.lock().await;
state.enable_gateway_tools();
let epoch = state.start_gateway_tool_fetch().unwrap();
assert!(state.complete_gateway_tool_fetch(
epoch,
crate::session::managed_mcp::GatewayToolCatalog {
tools: vec![
crate::session::managed_mcp::GatewayTool {
connector_id: "linear".to_string(),
connector_name: "Linear".to_string(),
tool_id: "list_issues".to_string(),
tool_name: "List".to_string(),
call_id: "linear.list_issues".to_string(),
description: "List issues".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
crate::session::managed_mcp::GatewayTool {
connector_id: "linear".to_string(),
connector_name: "Linear".to_string(),
tool_id: "create_issue".to_string(),
tool_name: "Create".to_string(),
call_id: "linear.create_issue".to_string(),
description: "Create issue".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
crate::session::managed_mcp::GatewayTool {
connector_id: "slack".to_string(),
connector_name: "Slack".to_string(),
tool_id: "search".to_string(),
tool_name: "Search".to_string(),
call_id: "slack.search".to_string(),
description: "Search Slack".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
],
total_tools: 3,
connectors_needing_reauth: vec![],
}
));
}
let snapshot = Arc::new(std::sync::Mutex::new(
crate::session::tool_index::ToolMetadataSnapshot::default(),
));
let disabled: std::collections::HashMap<String, std::collections::HashSet<String>> =
std::collections::HashMap::from([
(
"linear".to_string(),
std::collections::HashSet::from(["linear__create_issue".to_string()]),
),
(
crate::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY.to_string(),
std::collections::HashSet::from(["slack".to_string()]),
),
]);
refresh_mcp_snapshot_for_test_with_disabled(
bridge,
mcp_state,
managed,
snapshot.clone(),
&disabled,
)
.await;
let snapshot = snapshot.lock().unwrap();
let names: std::collections::HashSet<&str> = snapshot
.tools
.iter()
.map(|tool| tool.qualified_name.as_str())
.collect();
assert!(names.contains("linear__list_issues"));
assert!(!names.contains("linear__create_issue"));
assert!(!names.contains("slack__search"));
}
}
/// ToolBridge must route file operations through the injected FileSystem, /// ToolBridge must route file operations through the injected FileSystem,
/// not direct disk I/O. When `.with_fs()` is dropped from the builder, /// not direct disk I/O. When `.with_fs()` is dropped from the builder,
/// tools fall back to LocalFs and ACP client-side enforcement stops working. /// tools fall back to LocalFs and ACP client-side enforcement stops working.
@@ -1519,7 +1161,6 @@ mod fs_injection_regression_tests;
#[path = "acp_session_tests/interjection_actor_tests.rs"] #[path = "acp_session_tests/interjection_actor_tests.rs"]
mod interjection_actor_tests; mod interjection_actor_tests;
#[cfg(test)] #[cfg(test)]
#[cfg(test)]
#[path = "acp_session_tests/permission_auto_mode_tests.rs"] #[path = "acp_session_tests/permission_auto_mode_tests.rs"]
mod permission_auto_mode_tests; mod permission_auto_mode_tests;
/// Resume re-park of the parked `exit_plan_mode` approval. /// Resume re-park of the parked `exit_plan_mode` approval.
@@ -1559,6 +1200,9 @@ mod prompt_queue_actor_tests;
#[path = "acp_session_tests/record_response_token_usage_tests.rs"] #[path = "acp_session_tests/record_response_token_usage_tests.rs"]
mod record_response_token_usage_tests; mod record_response_token_usage_tests;
#[cfg(test)] #[cfg(test)]
#[path = "acp_session_tests/replace_system_prompt_tests.rs"]
mod replace_system_prompt_tests;
#[cfg(test)]
#[path = "acp_session_tests/replay_buffer_send_update_tests.rs"] #[path = "acp_session_tests/replay_buffer_send_update_tests.rs"]
mod replay_buffer_send_update_tests; mod replay_buffer_send_update_tests;
#[cfg(test)] #[cfg(test)]
@@ -1578,9 +1222,15 @@ mod rewrite_zero_turn_prefix_tests;
#[path = "acp_session_tests/subagent_usage_fold_tests.rs"] #[path = "acp_session_tests/subagent_usage_fold_tests.rs"]
mod subagent_usage_fold_tests; mod subagent_usage_fold_tests;
#[cfg(test)] #[cfg(test)]
#[path = "acp_session_tests/support.rs"]
mod support;
#[cfg(test)]
#[path = "acp_session_tests/turn_completion_emit_tests.rs"] #[path = "acp_session_tests/turn_completion_emit_tests.rs"]
mod turn_completion_emit_tests; mod turn_completion_emit_tests;
#[cfg(test)] #[cfg(test)]
#[path = "acp_session_tests/usage_categories_tests.rs"]
mod usage_categories_tests;
#[cfg(test)]
mod tool_meta_stamp_tests { mod tool_meta_stamp_tests {
//! Pin the `x.ai/tool` stamps on the harness emission paths: the early //! Pin the `x.ai/tool` stamps on the harness emission paths: the early
//! ToolCall registered by `prepare_tool_call` and the permission-request //! ToolCall registered by `prepare_tool_call` and the permission-request
@@ -1759,15 +1409,33 @@ mod cancel_running_task_tests;
#[path = "acp_session_tests/feedback_turn_lookup_tests.rs"] #[path = "acp_session_tests/feedback_turn_lookup_tests.rs"]
mod feedback_turn_lookup_tests; mod feedback_turn_lookup_tests;
#[cfg(test)] #[cfg(test)]
#[path = "acp_session_tests/goal/goal_backoff_tests.rs"]
mod goal_backoff_tests;
#[cfg(test)]
#[path = "acp_session_tests/goal/goal_classifier_e2e_tests.rs"]
mod goal_classifier_e2e_tests;
#[cfg(test)]
#[path = "acp_session_tests/goal/goal_planner_e2e_tests.rs"]
mod goal_planner_e2e_tests;
#[cfg(test)]
#[path = "acp_session_tests/goal/goal_reminder_subagent_rules_tests.rs"] #[path = "acp_session_tests/goal/goal_reminder_subagent_rules_tests.rs"]
mod goal_reminder_subagent_rules_tests; mod goal_reminder_subagent_rules_tests;
#[cfg(test)] #[cfg(test)]
#[path = "acp_session_tests/goal/goal_strategist_e2e_tests.rs"]
mod goal_strategist_e2e_tests;
#[cfg(test)]
#[path = "acp_session_tests/goal/goal_summarizer_e2e_tests.rs"]
mod goal_summarizer_e2e_tests;
#[cfg(test)]
#[path = "acp_session_tests/idle_resume_tests.rs"] #[path = "acp_session_tests/idle_resume_tests.rs"]
mod idle_resume_tests; mod idle_resume_tests;
#[cfg(test)] #[cfg(test)]
#[path = "acp_session_tests/inline_auto_compact_flow_tests.rs"] #[path = "acp_session_tests/inline_auto_compact_flow_tests.rs"]
mod inline_auto_compact_flow_tests; mod inline_auto_compact_flow_tests;
#[cfg(test)] #[cfg(test)]
#[path = "acp_session_tests/interjection_tests.rs"]
mod interjection_tests;
#[cfg(test)]
#[path = "acp_session_tests/laziness/laziness_debug_tests.rs"] #[path = "acp_session_tests/laziness/laziness_debug_tests.rs"]
mod laziness_debug_tests; mod laziness_debug_tests;
#[cfg(test)] #[cfg(test)]
@@ -1780,9 +1448,6 @@ mod laziness_integration_tests;
#[path = "acp_session_tests/load_user_prompts_tests.rs"] #[path = "acp_session_tests/load_user_prompts_tests.rs"]
mod load_user_prompts_tests; mod load_user_prompts_tests;
#[cfg(test)] #[cfg(test)]
#[path = "acp_session_tests/media_gen_auth_retry_tests.rs"]
mod media_gen_auth_retry_tests;
#[cfg(test)]
#[path = "acp_session_tests/memory_config_tests.rs"] #[path = "acp_session_tests/memory_config_tests.rs"]
mod memory_config_tests; mod memory_config_tests;
#[cfg(test)] #[cfg(test)]
@@ -1792,15 +1457,18 @@ mod parallel_dispatch_tests;
#[path = "acp_session_tests/prompt_context_persistence_tests.rs"] #[path = "acp_session_tests/prompt_context_persistence_tests.rs"]
mod prompt_context_persistence_tests; mod prompt_context_persistence_tests;
#[cfg(test)] #[cfg(test)]
#[path = "acp_session_tests/reactive_managed_reauth_e2e_tests.rs"] #[path = "acp_session_tests/recap_display_only_tests.rs"]
mod reactive_managed_reauth_e2e_tests; mod recap_display_only_tests;
#[cfg(test)] #[cfg(test)]
#[path = "acp_session_tests/reactive_managed_reauth_tests.rs"] #[path = "acp_session_tests/reminder_policy_tests.rs"]
mod reactive_managed_reauth_tests; mod reminder_policy_tests;
#[cfg(test)] #[cfg(test)]
#[path = "acp_session_tests/session_thread_tests.rs"] #[path = "acp_session_tests/session_thread_tests.rs"]
mod session_thread_tests; mod session_thread_tests;
#[cfg(test)] #[cfg(test)]
#[path = "acp_session_tests/tool_auth_retry_tests.rs"]
mod tool_auth_retry_tests;
#[cfg(test)]
#[path = "acp_session_tests/turn/turn_end_guard_tests.rs"] #[path = "acp_session_tests/turn/turn_end_guard_tests.rs"]
mod turn_end_guard_tests; mod turn_end_guard_tests;
#[cfg(test)] #[cfg(test)]
@@ -1809,215 +1477,3 @@ mod wait_for_mcp_prefix_tests;
#[cfg(test)] #[cfg(test)]
#[path = "acp_session_tests/web_search_e2e_tests.rs"] #[path = "acp_session_tests/web_search_e2e_tests.rs"]
mod web_search_e2e_tests; mod web_search_e2e_tests;
#[cfg(test)]
mod managed_gateway_tool_tests {
use super::*;
use kigi_tools::types::output::{MCPOutput, ToolOutput};
use kigi_tools::types::tool::{ToolKind, ToolNamespace};
use kigi_tools::types::tool_metadata::ToolMetadata;
#[derive(Debug)]
struct FixtureMcpTool;
impl ToolMetadata for FixtureMcpTool {
fn kind(&self) -> ToolKind {
ToolKind::Other
}
fn tool_namespace(&self) -> ToolNamespace {
ToolNamespace::MCP
}
fn description_template(&self) -> &str {
"fixture"
}
}
impl kigi_tool_runtime::Tool for FixtureMcpTool {
type Args = serde_json::Value;
type Output = ToolOutput;
fn id(&self) -> kigi_tool_protocol::ToolId {
kigi_tool_protocol::ToolId::new("server__tool").expect("valid")
}
fn description(
&self,
_ctx: &::kigi_tool_runtime::ListToolsContext,
) -> kigi_tool_types::ToolDescription {
kigi_tool_types::ToolDescription::new("server__tool", "fixture")
}
async fn run(
&self,
_ctx: kigi_tool_runtime::ToolCallContext,
_args: serde_json::Value,
) -> Result<ToolOutput, kigi_tool_runtime::ToolError> {
Ok(ToolOutput::MCP(MCPOutput::okay_output(
"server__tool".to_string(),
"server".to_string(),
"ok".to_string(),
)))
}
}
#[tokio::test]
async fn refresh_snapshot_seeds_only_admitted_gateway_catalog_entries() {
let bridge = Arc::new(crate::tools::bridge::ToolBridge::for_test());
bridge
.register_mcp_tools(
"server__tool".to_string(),
FixtureMcpTool,
Some(serde_json::json!({ "type" : "object" })),
)
.await
.expect("local fixture registration succeeds");
let mcp_state = Arc::new(TokioMutex::new(McpState::new(vec![])));
let managed = crate::session::managed_mcp::ManagedMcpStateHandle::default();
{
let mut state = managed.lock().await;
state.enable_gateway_tools();
let epoch = state.start_gateway_tool_fetch().unwrap();
assert!(state.complete_gateway_tool_fetch(
epoch,
crate::session::managed_mcp::GatewayToolCatalog {
tools: vec![
crate::session::managed_mcp::GatewayTool {
connector_id: "server".to_string(),
connector_name: "Gateway Collision".to_string(),
tool_id: "tool".to_string(),
tool_name: "Collision".to_string(),
call_id: "gateway.collision".to_string(),
description: "Gateway collision".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
crate::session::managed_mcp::GatewayTool {
connector_id: "gateway".to_string(),
connector_name: "Gateway".to_string(),
tool_id: "search".to_string(),
tool_name: "Search".to_string(),
call_id: "gateway.search".to_string(),
description: "Gateway search".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
],
total_tools: 2,
connectors_needing_reauth: vec![],
}
));
}
let snapshot = Arc::new(std::sync::Mutex::new(
crate::session::tool_index::ToolMetadataSnapshot::default(),
));
refresh_mcp_snapshot_for_test(bridge.clone(), mcp_state, managed, snapshot.clone()).await;
let catalog = bridge
.read_resource::<kigi_tools::types::resources::ManagedGatewayToolCatalog>()
.await
.expect("catalog resource should be seeded");
assert!(catalog.get("gateway__search").is_some());
assert!(
catalog.get("server__tool").is_none(),
"gateway catalog resource must match admitted snapshot and skip local collisions"
);
}
#[tokio::test]
async fn refresh_snapshot_excludes_disabled_gateway_tools_and_connectors() {
let bridge = Arc::new(crate::tools::bridge::ToolBridge::for_test());
let mcp_state = Arc::new(TokioMutex::new(McpState::new(vec![])));
let managed = crate::session::managed_mcp::ManagedMcpStateHandle::default();
{
let mut state = managed.lock().await;
state.enable_gateway_tools();
let epoch = state.start_gateway_tool_fetch().unwrap();
assert!(state.complete_gateway_tool_fetch(
epoch,
crate::session::managed_mcp::GatewayToolCatalog {
tools: vec![
crate::session::managed_mcp::GatewayTool {
connector_id: "linear".to_string(),
connector_name: "Linear".to_string(),
tool_id: "list_issues".to_string(),
tool_name: "List".to_string(),
call_id: "linear.list_issues".to_string(),
description: "List issues".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
crate::session::managed_mcp::GatewayTool {
connector_id: "linear".to_string(),
connector_name: "Linear".to_string(),
tool_id: "create_issue".to_string(),
tool_name: "Create".to_string(),
call_id: "linear.create_issue".to_string(),
description: "Create issue".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
crate::session::managed_mcp::GatewayTool {
connector_id: "slack".to_string(),
connector_name: "Slack".to_string(),
tool_id: "search".to_string(),
tool_name: "Search".to_string(),
call_id: "slack.search".to_string(),
description: "Search Slack".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
],
total_tools: 3,
connectors_needing_reauth: vec![],
}
));
}
let snapshot = Arc::new(std::sync::Mutex::new(
crate::session::tool_index::ToolMetadataSnapshot::default(),
));
let disabled: std::collections::HashMap<String, std::collections::HashSet<String>> =
std::collections::HashMap::from([
(
"linear".to_string(),
std::collections::HashSet::from(["linear__create_issue".to_string()]),
),
(
crate::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY.to_string(),
std::collections::HashSet::from(["slack".to_string()]),
),
]);
refresh_mcp_snapshot_for_test_with_disabled(
bridge.clone(),
mcp_state,
managed,
snapshot.clone(),
&disabled,
)
.await;
let catalog = bridge
.read_resource::<kigi_tools::types::resources::ManagedGatewayToolCatalog>()
.await
.expect("catalog resource should be seeded");
assert!(catalog.get("linear__list_issues").is_some());
assert!(catalog.get("linear__create_issue").is_none());
assert!(catalog.get("slack__search").is_none());
let snapshot = snapshot.lock().unwrap();
let names: std::collections::HashSet<&str> = snapshot
.tools
.iter()
.map(|tool| tool.qualified_name.as_str())
.collect();
assert!(names.contains("linear__list_issues"));
assert!(!names.contains("linear__create_issue"));
assert!(!names.contains("slack__search"));
}
}
#[cfg(test)]
#[path = "acp_session_tests/goal/goal_backoff_tests.rs"]
mod goal_backoff_tests;
#[cfg(test)]
#[path = "acp_session_tests/goal/goal_classifier_e2e_tests.rs"]
mod goal_classifier_e2e_tests;
#[cfg(test)]
#[path = "acp_session_tests/goal/goal_planner_e2e_tests.rs"]
mod goal_planner_e2e_tests;
#[cfg(test)]
#[path = "acp_session_tests/goal/goal_strategist_e2e_tests.rs"]
mod goal_strategist_e2e_tests;
#[cfg(test)]
#[path = "acp_session_tests/goal/goal_summarizer_e2e_tests.rs"]
mod goal_summarizer_e2e_tests;
#[cfg(test)]
#[path = "acp_session_tests/interjection_tests.rs"]
mod interjection_tests;
#[cfg(test)]
#[path = "acp_session_tests/recap_display_only_tests.rs"]
mod recap_display_only_tests;
#[cfg(test)]
#[path = "acp_session_tests/reminder_policy_tests.rs"]
mod reminder_policy_tests;
@@ -900,17 +900,9 @@ impl SessionActor {
// the order-sensitive `update_configs` would cause (merge order is // the order-sensitive `update_configs` would cause (merge order is
// non-deterministic). Mirrors the `UpdateMcpServers` command handler. // non-deterministic). Mirrors the `UpdateMcpServers` command handler.
let t_mcp = std::time::Instant::now(); let t_mcp = std::time::Instant::now();
let managed_configs = {
let mcp_handle = self.managed_mcp_handle.lock().await;
match &mcp_handle.cache {
crate::session::managed_mcp::ManagedMcpCache::Ready(configs) => configs.clone(),
_ => vec![],
}
};
let new_mcp_servers = crate::session::managed_mcp::merge_managed_mcp_servers( let new_mcp_servers = crate::session::managed_mcp::merge_managed_mcp_servers(
self.initial_client_mcp_servers.clone(), self.initial_client_mcp_servers.clone(),
session_cwd, session_cwd,
&managed_configs,
new_registry_snapshot.as_deref(), new_registry_snapshot.as_deref(),
&self.rebuild_spec.compat, &self.rebuild_spec.compat,
); );
@@ -17,230 +17,6 @@ impl SessionActor {
} }
self.ensure_mcp_tools_initialized().await; self.ensure_mcp_tools_initialized().await;
} }
/// If managed tokens are near expiry, swap clients using the agent-level cache.
pub(super) async fn refresh_managed_mcp_if_stale(&self) {
use crate::session::managed_mcp::ManagedMcpCache;
let is_stale = crate::session::managed_mcp::managed_token_is_stale(
*self.managed_mcp_expires_at.lock().unwrap(),
chrono::Utc::now(),
);
if !is_stale {
return;
}
let fresh_configs = match &self.managed_mcp_handle.lock().await.cache {
ManagedMcpCache::Ready(c) if !c.is_empty() => c.clone(),
_ => return,
};
self.mcp_state.lock().await.refresh_managed_clients(
fresh_configs
.iter()
.map(|c| (c.endpoint.as_str(), &c.headers)),
);
*self.managed_mcp_expires_at.lock().unwrap() = fresh_configs
.iter()
.filter_map(|c| c.token_expires_at)
.min();
}
/// Reactive managed re-auth: cache-bypassing config re-fetch, swap in fresh
/// headers, re-handshake once. `Ok(())` if the server is `Ready` afterward.
///
/// Owner-scoped — only the session owning the client in `owned_clients` swaps
/// (`refresh_managed_clients` ignores a subagent's shared Arc, which recovers
/// via the leader on its next pool snapshot). The shared `ManagedMcpState`
/// cooldown coalesces concurrent failures and bounds a revoked connector.
/// Shared by both reactive entry points; locks are taken sequentially, never
/// nested (acquire `managed_mcp_handle` before `mcp_state` if that changes).
///
/// Emits `metrics.mcp.managed.reauth.{triggered,outcome,cooldown_terminal}`
/// tracing targets, plus an `mcp.server_connection` span (`connected` with
/// `tool_count` on recovery, `failed`/`error_type=auth` on terminal exhaustion).
pub(super) async fn reactive_managed_reauth(&self, server_name: &str) -> Result<(), String> {
if !self
.mcp_state
.lock()
.await
.owned_clients
.contains_key(server_name)
{
return Err(format!(
"session does not own managed client '{server_name}'"
));
}
let now = chrono::Utc::now();
if !self
.managed_mcp_handle
.lock()
.await
.reauth_allowed(server_name, now)
{
return Err(format!(
"managed reactive re-auth for '{server_name}' is in cooldown"
));
}
tracing::info!(
target : "metrics.mcp.managed.reauth.triggered", server = % server_name
);
tracing::info!(
server = % server_name,
"managed MCP auth rejection detected, attempting reactive re-fetch"
);
let scope = || {
crate::util::config::mcp_server_scope(
server_name,
std::path::Path::new(self.session_info.cwd.as_str()),
)
};
let started = std::time::Instant::now();
match self.reactive_managed_reauth_inner(server_name).await {
Ok(tool_count) => {
let elapsed_ms = started.elapsed().as_millis() as i64;
self.managed_mcp_handle
.lock()
.await
.record_reauth_success(server_name);
tracing::info!(
target : "metrics.mcp.managed.reauth.outcome", server = %
server_name, result = "recovered",
);
tracing::info!(
server = % server_name, "managed MCP reactive re-auth recovered"
);
crate::session::telemetry::emit_mcp_connection_span(
"connected",
server_name,
"http",
scope(),
Some(elapsed_ms),
Some(tool_count as i64),
None,
);
let payload = crate::session::mcp_dispatcher::McpServerStatusPayload {
session_id: self.session_id_string(),
name: server_name.to_string(),
source: crate::session::mcp_dispatcher::classify_source(server_name),
status: crate::session::mcp_dispatcher::McpServerStatus::Ready,
reason:
crate::session::mcp_dispatcher::McpServerStatusReason::ManagedTokenRefreshed,
detail: None,
tools: None,
};
crate::session::mcp_restart::forward_status(&self.notifications.gateway, &payload);
Ok(())
}
Err(e) => {
let elapsed_ms = started.elapsed().as_millis() as i64;
let failed_at = chrono::Utc::now();
let terminal = {
let mut st = self.managed_mcp_handle.lock().await;
st.record_reauth_failure(server_name, failed_at);
st.reauth_is_terminal(server_name)
};
if terminal {
self.mcp_state
.lock()
.await
.record_init_failure(server_name, true, None);
let payload = crate::session::mcp_dispatcher::McpServerStatusPayload {
session_id: self.session_id_string(),
name: server_name.to_string(),
source: crate::session::mcp_dispatcher::classify_source(server_name),
status: crate::session::mcp_dispatcher::McpServerStatus::NeedsAuth,
reason: crate::session::mcp_dispatcher::McpServerStatusReason::AuthExpired,
detail: None,
tools: None,
};
crate::session::mcp_restart::forward_status(
&self.notifications.gateway,
&payload,
);
tracing::warn!(
target : "metrics.mcp.managed.reauth.cooldown_terminal", server =
% server_name,
);
tracing::warn!(
server = % server_name,
"managed MCP reactive re-auth exhausted; surfacing NeedsAuth"
);
crate::session::telemetry::emit_mcp_connection_span(
"failed",
server_name,
"http",
scope(),
Some(elapsed_ms),
None,
Some("auth"),
);
self.unregister_server_tools(server_name);
self.refresh_mcp_snapshot_and_schedule_reminder().await;
}
tracing::info!(
target : "metrics.mcp.managed.reauth.outcome", server = %
server_name, result = if terminal { "failed" } else { "cooldown" },
);
Err(e)
}
}
}
/// Inner half of [`Self::reactive_managed_reauth`]: snapshot inputs, force
/// a fresh proxy fetch, swap clients, re-handshake, and register tools.
/// Split out so the caller owns the cooldown gate and the status push.
/// Returns the number of tools registered after recovery, so the caller can
/// stamp the recovered `mcp.server_connection` span with a `tool_count`.
async fn reactive_managed_reauth_inner(&self, server_name: &str) -> Result<usize, String> {
let auth_manager = self.auth_manager.clone();
let proxy_url = self.models_manager.endpoints().proxy_url();
let Some(am) = auth_manager else {
return Err("no auth manager; cannot re-fetch managed configs".into());
};
crate::session::managed_mcp::invalidate_cache(&self.managed_mcp_handle).await;
let fresh_configs = crate::session::managed_mcp::fetch_managed_mcp_configs(
&self.managed_mcp_handle,
&proxy_url,
&am,
)
.await;
if fresh_configs.is_empty() {
return Err("managed re-fetch returned no configs".into());
}
{
let mut st = self.mcp_state.lock().await;
st.refresh_managed_clients(
fresh_configs
.iter()
.map(|c| (c.endpoint.as_str(), &c.headers)),
);
}
*self.managed_mcp_expires_at.lock().unwrap() = fresh_configs
.iter()
.filter_map(|c| c.token_expires_at)
.min();
let client = {
let st = self.mcp_state.lock().await;
st.get_client(server_name)
.cloned()
.ok_or_else(|| format!("client '{server_name}' missing after refresh"))?
};
let registrations = client
.get_tool_registrations(self.mcp_state.clone())
.await
.map_err(|e| format!("re-handshake failed: {e}"))?;
let mut mcp_state = self.mcp_state.lock().await;
mcp_state.auth_required.remove(server_name);
mcp_state.clear_init_failed(server_name);
let tool_count = registrations.len();
let mut ui_tools: std::collections::HashMap<
String,
Vec<crate::extensions::mcp::McpToolEntry>,
> = std::collections::HashMap::new();
for reg in registrations {
self.register_mcp_tool(server_name, reg, &mut mcp_state, &mut ui_tools)
.await;
}
drop(mcp_state);
self.refresh_mcp_snapshot_and_schedule_reminder().await;
self.emit_mcp_tools_changed_notifications(ui_tools);
Ok(tool_count)
}
/// Register tools from shared (inherited) MCP clients on this session's ToolBridge. /// Register tools from shared (inherited) MCP clients on this session's ToolBridge.
/// ///
/// Shared clients are already connected (Arc-shared from parent), so /// Shared clients are already connected (Arc-shared from parent), so
@@ -411,9 +187,6 @@ impl SessionActor {
/// Runs force_reauth (browser flow), then re-initializes the server /// Runs force_reauth (browser flow), then re-initializes the server
/// and registers its tools. /// and registers its tools.
pub(super) async fn handle_mcp_auth_trigger(&self, server_name: &str) -> Result<(), String> { pub(super) async fn handle_mcp_auth_trigger(&self, server_name: &str) -> Result<(), String> {
if server_name.starts_with(crate::session::managed_mcp::MANAGED_MCP_PREFIX) {
return Err("To authenticate, visit grok.com".to_string());
}
let client = { let client = {
let state = self.mcp_state.lock().await; let state = self.mcp_state.lock().await;
state state
@@ -527,28 +300,13 @@ impl SessionActor {
/// `<system-reminder>` at a turn boundary. The `search_tool` description /// `<system-reminder>` at a turn boundary. The `search_tool` description
/// itself stays static (cacheable). /// itself stays static (cacheable).
pub(super) async fn refresh_mcp_snapshot_and_schedule_reminder(&self) { pub(super) async fn refresh_mcp_snapshot_and_schedule_reminder(&self) {
let disabled_gateway_tools = crate::util::config::get_all_mcp_disabled_tools(
std::path::Path::new(&self.session_info.cwd),
);
self.refresh_mcp_snapshot_and_schedule_reminder_with_disabled(&disabled_gateway_tools)
.await;
}
pub(super) async fn refresh_mcp_snapshot_and_schedule_reminder_with_disabled(
&self,
disabled_gateway_tools: &std::collections::HashMap<
String,
std::collections::HashSet<String>,
>,
) {
let mcp_initialized = self.mcp_state.lock().await.is_initialized(); let mcp_initialized = self.mcp_state.lock().await.is_initialized();
refresh_mcp_snapshot_and_schedule_reminder_with( refresh_mcp_snapshot_and_schedule_reminder_with(
self.agent.borrow().tool_bridge().clone(), self.agent.borrow().tool_bridge().clone(),
Arc::clone(&self.mcp_state), Arc::clone(&self.mcp_state),
self.managed_mcp_handle.clone(),
self.tool_metadata_snapshot.clone(), self.tool_metadata_snapshot.clone(),
Arc::clone(&self.mcp_reminder_dirty), Arc::clone(&self.mcp_reminder_dirty),
mcp_initialized, mcp_initialized,
disabled_gateway_tools,
self.cursor_mcps_root(), self.cursor_mcps_root(),
) )
.await; .await;
@@ -638,8 +396,7 @@ impl SessionActor {
"connection failed".to_string() "connection failed".to_string()
}; };
let retries_on_use = !mcp_state.auth_required.contains(name) let retries_on_use = !mcp_state.auth_required.contains(name)
&& matches!(cfg, acp::McpServer::Http(_) | acp::McpServer::Sse(_)) && matches!(cfg, acp::McpServer::Http(_) | acp::McpServer::Sse(_));
&& !name.starts_with(crate::session::managed_mcp::MANAGED_MCP_PREFIX);
let reason = if retries_on_use { let reason = if retries_on_use {
format!("{base} — retries automatically on next tool call") format!("{base} — retries automatically on next tool call")
} else { } else {
@@ -727,14 +484,9 @@ impl SessionActor {
!disabled.contains(server) !disabled.contains(server)
} }
/// HTTP analog of [`Self::is_stdio_server_configured`]: `true` iff /// HTTP analog of [`Self::is_stdio_server_configured`]: `true` iff
/// `server` has an enabled, **non-managed** `Http` / `Sse` config entry. /// `server` has an enabled `Http` / `Sse` config entry.
/// Gates [`crate::session::mcp_restart::maybe_schedule_http_recovery`]. /// Gates [`crate::session::mcp_restart::maybe_schedule_http_recovery`].
/// Managed connectors (`MANAGED_MCP_PREFIX`) are excluded — out of scope
/// for in-place recovery; this mirrors the dispatcher's filter.
pub(crate) async fn is_http_server_configured(&self, server: &str) -> bool { pub(crate) async fn is_http_server_configured(&self, server: &str) -> bool {
if server.starts_with(crate::session::managed_mcp::MANAGED_MCP_PREFIX) {
return false;
}
let mcp_state = self.mcp_state.lock().await; let mcp_state = self.mcp_state.lock().await;
let is_http_in_configs = mcp_state.configs.iter().any(|c| { let is_http_in_configs = mcp_state.configs.iter().any(|c| {
matches!( matches!(
@@ -1076,17 +828,6 @@ impl SessionActor {
&mcp_server_configs, &mcp_server_configs,
cwd, cwd,
)); ));
let managed_count = mcp_server_configs
.iter()
.filter(|c| {
mcp_server_name(c).starts_with(crate::session::managed_mcp::MANAGED_MCP_PREFIX)
})
.count() as u32;
self.events
.emit(kigi_file_utils::events::Event::McpManagedConfigResult {
server_count: managed_count,
error: None,
});
} }
let configs_to_start: Vec<_> = mcp_server_configs let configs_to_start: Vec<_> = mcp_server_configs
.iter() .iter()
@@ -1264,14 +1005,10 @@ impl SessionActor {
let tool_bridge = self.agent.borrow().tool_bridge().clone(); let tool_bridge = self.agent.borrow().tool_bridge().clone();
let gateway = self.notifications.gateway.clone(); let gateway = self.notifications.gateway.clone();
let tool_snapshot = self.tool_metadata_snapshot.clone(); let tool_snapshot = self.tool_metadata_snapshot.clone();
let managed_mcp_handle = self.managed_mcp_handle.clone();
let mcp_reminder_dirty = Arc::clone(&self.mcp_reminder_dirty); let mcp_reminder_dirty = Arc::clone(&self.mcp_reminder_dirty);
let mcp_handshakes_done = Arc::clone(&self.mcp_handshakes_done); let mcp_handshakes_done = Arc::clone(&self.mcp_handshakes_done);
let session_id_owned = self.session_info.id.0.clone(); let session_id_owned = self.session_info.id.0.clone();
let mcps_root_bg = self.cursor_mcps_root(); let mcps_root_bg = self.cursor_mcps_root();
let disabled_gateway_tools_bg = crate::util::config::get_all_mcp_disabled_tools(
std::path::Path::new(&self.session_info.cwd),
);
let server_transport_map: std::collections::HashMap<String, &'static str> = let server_transport_map: std::collections::HashMap<String, &'static str> =
mcp_server_configs mcp_server_configs
.iter() .iter()
@@ -1348,13 +1085,7 @@ impl SessionActor {
Ok((server_name, handles, server_start.elapsed(), timeout_sec)) Ok((server_name, handles, server_start.elapsed(), timeout_sec))
} }
Err(e) => { Err(e) => {
let needs_auth = if server_name let needs_auth = client.has_auth();
.starts_with(crate::session::managed_mcp::MANAGED_MCP_PREFIX)
{
e.is_auth_rejection()
} else {
client.has_auth()
};
tracing::warn!( tracing::warn!(
server = server_name.as_str(), elapsed_ms = server_start server = server_name.as_str(), elapsed_ms = server_start
.elapsed().as_millis() as u64, timeout_sec, error = % e, .elapsed().as_millis() as u64, timeout_sec, error = % e,
@@ -1664,11 +1395,9 @@ impl SessionActor {
refresh_mcp_snapshot_and_schedule_reminder_with( refresh_mcp_snapshot_and_schedule_reminder_with(
tool_bridge.clone(), tool_bridge.clone(),
Arc::clone(&mcp_state_bg), Arc::clone(&mcp_state_bg),
managed_mcp_handle.clone(),
tool_snapshot, tool_snapshot,
mcp_reminder_dirty, mcp_reminder_dirty,
true, true,
&disabled_gateway_tools_bg,
mcps_root_bg, mcps_root_bg,
) )
.await; .await;
@@ -19,27 +19,12 @@ impl McpReminderMode {
} }
} }
pub(super) fn gateway_tool_is_disabled(
tool: &crate::session::managed_mcp::GatewayTool,
disabled_gateway_tools: &std::collections::HashMap<String, std::collections::HashSet<String>>,
) -> bool {
let qualified_name = tool.qualified_name();
disabled_gateway_tools
.get(crate::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY)
.is_some_and(|set| set.contains(&tool.connector_id))
|| disabled_gateway_tools
.get(&tool.connector_id)
.is_some_and(|set| set.contains(&qualified_name))
}
pub(super) async fn refresh_mcp_snapshot_and_schedule_reminder_with( pub(super) async fn refresh_mcp_snapshot_and_schedule_reminder_with(
tool_bridge: Arc<crate::tools::bridge::ToolBridge>, tool_bridge: Arc<crate::tools::bridge::ToolBridge>,
mcp_state: Arc<TokioMutex<McpState>>, mcp_state: Arc<TokioMutex<McpState>>,
managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle,
tool_metadata_snapshot: Arc<std::sync::Mutex<crate::session::tool_index::ToolMetadataSnapshot>>, tool_metadata_snapshot: Arc<std::sync::Mutex<crate::session::tool_index::ToolMetadataSnapshot>>,
mcp_reminder_dirty: Arc<std::sync::atomic::AtomicBool>, mcp_reminder_dirty: Arc<std::sync::atomic::AtomicBool>,
mcp_initialized: bool, mcp_initialized: bool,
disabled_gateway_tools: &std::collections::HashMap<String, std::collections::HashSet<String>>,
// External harness only: per-workspace `mcps/` descriptor root. `Some` makes // External harness only: per-workspace `mcps/` descriptor root. `Some` makes
// this refresh also update the on-disk descriptor mirror so late-connecting // this refresh also update the on-disk descriptor mirror so late-connecting
// servers become discoverable; `None` for other agent types (no-op). // servers become discoverable; `None` for other agent types (no-op).
@@ -51,7 +36,7 @@ pub(super) async fn refresh_mcp_snapshot_and_schedule_reminder_with(
let all_defs = tool_bridge.tool_definitions().await; let all_defs = tool_bridge.tool_definitions().await;
let mut seen_tools = std::collections::HashSet::new(); let mut seen_tools = std::collections::HashSet::new();
let mut mcp_tools: Vec<ToolMetadata> = all_defs let mcp_tools: Vec<ToolMetadata> = all_defs
.iter() .iter()
.filter(|d| d.function.name.contains("__")) .filter(|d| d.function.name.contains("__"))
.filter(|d| seen_tools.insert(d.function.name.clone())) .filter(|d| seen_tools.insert(d.function.name.clone()))
@@ -68,61 +53,6 @@ pub(super) async fn refresh_mcp_snapshot_and_schedule_reminder_with(
}) })
.collect(); .collect();
let (gateway_catalog, mut gateway_connectors) = {
let state = managed_mcp_handle.lock().await;
let catalog = if state.gateway_tools_active {
match &state.gateway_tool_cache {
crate::session::managed_mcp::GatewayToolCatalogCache::Ready(catalog) => {
Some(catalog.clone())
}
_ => None,
}
} else {
None
};
let connectors: Vec<String> = state.gateway_tool_connectors_seen.iter().cloned().collect();
(catalog, connectors)
};
if let Some(catalog) = gateway_catalog.as_ref() {
gateway_connectors.extend(catalog.tools.iter().map(|tool| tool.connector_id.clone()));
}
gateway_connectors.sort_unstable();
gateway_connectors.dedup();
let mut gateway_resource_entries = Vec::new();
if let Some(catalog) = gateway_catalog.as_ref() {
for tool in &catalog.tools {
let qualified_name = tool.qualified_name();
if gateway_tool_is_disabled(tool, disabled_gateway_tools) {
continue;
}
if !seen_tools.insert(qualified_name.clone()) {
continue;
}
gateway_resource_entries.push((
qualified_name.clone(),
kigi_tools::types::resources::ManagedGatewayToolSource {
connector_id: tool.connector_id.clone(),
connector_name: tool.connector_name.clone(),
tool_id: tool.tool_id.clone(),
tool_name: tool.tool_name.clone(),
call_id: tool.call_id.clone(),
},
));
// Gateway ids are the model/search contract. Display labels stay
// out of ToolMetadata so search_tool and permissions use stable ids:
// connector_id/tool_id here, connector_name/tool_name in UI only.
mcp_tools.push(ToolMetadata {
qualified_name,
server_name: tool.connector_id.clone(),
tool_name: tool.tool_id.clone(),
description: tool.description.clone(),
parameters: extract_parameter_names(&tool.json_schema),
input_schema: tool.json_schema.clone(),
});
}
}
let servers_with_tools: std::collections::HashSet<&str> = let servers_with_tools: std::collections::HashSet<&str> =
mcp_tools.iter().map(|t| t.server_name.as_str()).collect(); mcp_tools.iter().map(|t| t.server_name.as_str()).collect();
@@ -150,12 +80,6 @@ pub(super) async fn refresh_mcp_snapshot_and_schedule_reminder_with(
snapshot.mcp_initialized = mcp_initialized; snapshot.mcp_initialized = mcp_initialized;
} }
tool_bridge
.update_resource(kigi_tools::types::resources::ManagedGatewayToolCatalog(
gateway_resource_entries.into_iter().collect(),
))
.await;
mcp_reminder_dirty.store(true, std::sync::atomic::Ordering::Relaxed); mcp_reminder_dirty.store(true, std::sync::atomic::Ordering::Relaxed);
tracing::debug!("MCP snapshot updated, reminder marked dirty"); tracing::debug!("MCP snapshot updated, reminder marked dirty");
@@ -170,30 +94,8 @@ pub(super) async fn refresh_mcp_snapshot_and_schedule_reminder_with(
.map(|(n, c)| (n.clone(), Arc::clone(c))) .map(|(n, c)| (n.clone(), Arc::clone(c)))
.collect() .collect()
}; };
let protected_connectors = clients.iter().map(|(name, _)| name.clone()).collect();
let mut gateway_descriptors = Vec::new();
if let Some(catalog) = gateway_catalog.as_ref() {
for tool in &catalog.tools {
if gateway_tool_is_disabled(tool, disabled_gateway_tools) {
continue;
}
gateway_descriptors.push(crate::session::mcp_descriptors::GatewayToolDescriptor {
connector_id: tool.connector_id.clone(),
tool_id: tool.tool_id.clone(),
description: tool.description.clone(),
json_schema: tool.json_schema.clone(),
});
}
}
crate::session::mcp_descriptors::materialize_descriptors_for_clients(&mcps_root, clients) crate::session::mcp_descriptors::materialize_descriptors_for_clients(&mcps_root, clients)
.await; .await;
crate::session::mcp_descriptors::materialize_descriptors_for_gateway_tools(
&mcps_root,
gateway_descriptors,
gateway_connectors,
protected_connectors,
)
.await;
} }
} }
@@ -201,35 +103,14 @@ pub(super) async fn refresh_mcp_snapshot_and_schedule_reminder_with(
pub(crate) async fn refresh_mcp_snapshot_for_test( pub(crate) async fn refresh_mcp_snapshot_for_test(
tool_bridge: Arc<crate::tools::bridge::ToolBridge>, tool_bridge: Arc<crate::tools::bridge::ToolBridge>,
mcp_state: Arc<TokioMutex<McpState>>, mcp_state: Arc<TokioMutex<McpState>>,
managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle,
tool_metadata_snapshot: Arc<std::sync::Mutex<crate::session::tool_index::ToolMetadataSnapshot>>, tool_metadata_snapshot: Arc<std::sync::Mutex<crate::session::tool_index::ToolMetadataSnapshot>>,
) {
refresh_mcp_snapshot_for_test_with_disabled(
tool_bridge,
mcp_state,
managed_mcp_handle,
tool_metadata_snapshot,
&Default::default(),
)
.await;
}
#[cfg(test)]
pub(crate) async fn refresh_mcp_snapshot_for_test_with_disabled(
tool_bridge: Arc<crate::tools::bridge::ToolBridge>,
mcp_state: Arc<TokioMutex<McpState>>,
managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle,
tool_metadata_snapshot: Arc<std::sync::Mutex<crate::session::tool_index::ToolMetadataSnapshot>>,
disabled_gateway_tools: &std::collections::HashMap<String, std::collections::HashSet<String>>,
) { ) {
refresh_mcp_snapshot_and_schedule_reminder_with( refresh_mcp_snapshot_and_schedule_reminder_with(
tool_bridge, tool_bridge,
mcp_state, mcp_state,
managed_mcp_handle,
tool_metadata_snapshot, tool_metadata_snapshot,
Arc::new(std::sync::atomic::AtomicBool::new(false)), Arc::new(std::sync::atomic::AtomicBool::new(false)),
false, false,
disabled_gateway_tools,
None, None,
) )
.await; .await;
@@ -196,9 +196,6 @@ impl SessionActor {
std::sync::Arc::new(tool_index), std::sync::Arc::new(tool_index),
)) ))
.await; .await;
if let Some(client) = self.rebuild_spec.managed_gateway_tool_client.clone() {
bridge.update_resource(client).await;
}
let plan_path = self.plan_mode.lock().plan_file_path().to_path_buf(); let plan_path = self.plan_mode.lock().plan_file_path().to_path_buf();
bridge bridge
.update_resource(kigi_tools::types::resources::PlanFilePath(plan_path)) .update_resource(kigi_tools::types::resources::PlanFilePath(plan_path))
@@ -509,72 +509,9 @@ impl SessionActor {
folder_path: server_dir.map(|d| d.to_string_lossy().to_string()), folder_path: server_dir.map(|d| d.to_string_lossy().to_string()),
}); });
} }
let gateway_entries = self.gather_gateway_mcp_servers(mcps_root.as_deref()).await;
entries.extend(gateway_entries);
entries.sort_by(|a, b| a.name.cmp(&b.name)); entries.sort_by(|a, b| a.name.cmp(&b.name));
entries entries
} }
async fn gather_gateway_mcp_servers(
&self,
mcps_root: Option<&std::path::Path>,
) -> Vec<kigi_agent::prompt::user_message::McpServerEntry> {
use kigi_agent::prompt::user_message::McpServerEntry;
let disabled_gateway_tools = crate::util::config::get_all_mcp_disabled_tools(
std::path::Path::new(&self.session_info.cwd),
);
let catalog = {
let state = self.managed_mcp_handle.lock().await;
if state.gateway_tools_active {
match &state.gateway_tool_cache {
crate::session::managed_mcp::GatewayToolCatalogCache::Ready(catalog) => {
Some(catalog.clone())
}
_ => None,
}
} else {
None
}
};
let Some(catalog) = catalog else {
return Vec::new();
};
let mut connectors = std::collections::BTreeMap::<String, String>::new();
let mut gateway_connectors: Vec<String> = catalog
.tools
.iter()
.map(|tool| tool.connector_id.clone())
.collect();
gateway_connectors.sort_unstable();
gateway_connectors.dedup();
let mut descriptors = Vec::new();
for tool in &catalog.tools {
if gateway_tool_is_disabled(tool, &disabled_gateway_tools) {
continue;
}
connectors
.entry(tool.connector_id.clone())
.or_insert_with(|| tool.connector_name.clone());
descriptors.push(crate::session::mcp_descriptors::GatewayToolDescriptor {
connector_id: tool.connector_id.clone(),
tool_id: tool.tool_id.clone(),
description: tool.description.clone(),
json_schema: tool.json_schema.clone(),
});
}
connectors
.into_iter()
.map(|(connector_id, connector_name)| McpServerEntry {
folder_path: mcps_root.map(|root| {
crate::session::mcp_descriptors::server_descriptor_dir(root, &connector_id)
.to_string_lossy()
.to_string()
}),
name: connector_id,
server_use_instructions: (!connector_name.trim().is_empty())
.then_some(connector_name),
})
.collect()
}
/// Build a `PathRewriter` for sanitizing overlay paths in model-facing text. /// Build a `PathRewriter` for sanitizing overlay paths in model-facing text.
/// ///
/// Returns `None` when `display_cwd` is unset (no rewriting needed). Used /// Returns `None` when `display_cwd` is unset (no rewriting needed). Used
@@ -559,38 +559,7 @@ pub(super) async fn run_session(
tracing::warn!(server = sname.as_str(), error = % e, tracing::warn!(server = sname.as_str(), error = % e,
"Failed to persist server enabled state to config"); } let _ = respond_to "Failed to persist server enabled state to config"); } let _ = respond_to
.send(Ok(())); }); } SessionCommand::ToggleMcpTool { server_name, tool_name, .send(Ok(())); }); } SessionCommand::ToggleMcpTool { server_name, tool_name,
enabled, is_managed_gateway, respond_to } => { if is_managed_gateway { let enabled, respond_to } => { let
mut disabled_tools = crate
::util::config::get_all_mcp_disabled_tools(std::path::Path::new(& session
.session_info.cwd)); if tool_name.is_empty() { let set = disabled_tools
.entry(crate ::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY
.to_string()).or_default(); if enabled { set.remove(& server_name); } else {
set.insert(server_name.clone()); }
if set.is_empty() { disabled_tools
.remove(crate ::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY); } }
else if enabled { if let Some(set) = disabled_tools.get_mut(& server_name) {
set.remove(& tool_name); if set.is_empty() { disabled_tools.remove(&
server_name); } } } else { disabled_tools.entry(server_name.clone())
.or_default().insert(tool_name.clone()); } session
.refresh_mcp_snapshot_and_schedule_reminder_with_disabled(& disabled_tools,).
await; session.refresh_goal_harness_enabled(). await; let disabled_vec : Vec
< String > = if tool_name.is_empty() { disabled_tools.get(crate
::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY).map(| s | s.iter()
.cloned().collect()).unwrap_or_default() } else { disabled_tools.get(&
server_name).map(| s | s.iter().cloned().collect()).unwrap_or_default() };
let notifications = session.notifications.gateway.clone(); let session_id =
session.session_info.id.0.clone(); let server_for_persist = if tool_name
.is_empty() { crate ::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY
.to_string() } else { server_name.clone() }; tokio::task::spawn_local(async
move { if let Err(e) = crate ::util::config::save_mcp_disabled_tools(&
server_for_persist, & disabled_vec,). await { tracing::warn!(server =
server_for_persist.as_str(), error = % e,
"Failed to persist disabled_tools to config"); } let payload = crate
::extensions::mcp::McpToolsChanged { session_id : session_id.to_string(),
server_name : String::new(), tools : Vec::new(), }; if let Ok(params) =
serde_json::value::to_raw_value(& payload) { notifications
.forward_fire_and_forget(acp::ExtNotification::new("x.ai/mcp/tools_changed",
params.into())); } let _ = respond_to.send(Ok(())); }); continue; } let
qualified = format!("{}{}{}", server_name, crate qualified = format!("{}{}{}", server_name, crate
::session::mcp_servers::MCP_TOOL_NAME_DELIMITER, tool_name,); let mut ::session::mcp_servers::MCP_TOOL_NAME_DELIMITER, tool_name,); let mut
mcp_state = session.mcp_state.lock(). await; if enabled { if let Some(set) = mcp_state = session.mcp_state.lock(). await; if enabled { if let Some(set) =
@@ -663,11 +632,7 @@ pub(super) async fn run_session(
SessionCommand::McpAuthTrigger { server_name, respond_to } => { let s = SessionCommand::McpAuthTrigger { server_name, respond_to } => { let s =
session.clone(); tokio::task::spawn_local(async move { let result = s session.clone(); tokio::task::spawn_local(async move { let result = s
.handle_mcp_auth_trigger(& server_name). await; let _ = respond_to .handle_mcp_auth_trigger(& server_name). await; let _ = respond_to
.send(result); }); } SessionCommand::GetManagedGatewayDisabledTools { .send(result); }); } SessionCommand::RetryAuthRequiredServers { respond_to } => { let s = session
respond_to } => { let disabled_tools = crate
::util::config::get_all_mcp_disabled_tools(std::path::Path::new(& session
.session_info.cwd),); let _ = respond_to.send(disabled_tools); }
SessionCommand::RetryAuthRequiredServers { respond_to } => { let s = session
.clone(); tokio::task::spawn_local(async move { s .clone(); tokio::task::spawn_local(async move { s
.retry_auth_required_servers(). await; let _ = respond_to.send(()); }); } .retry_auth_required_servers(). await; let _ = respond_to.send(()); }); }
SessionCommand::RefreshMcpSearchIndex => { session SessionCommand::RefreshMcpSearchIndex => { session
@@ -139,9 +139,6 @@ pub(crate) async fn spawn_session_actor(
persisted_announcement_state: Option<crate::session::announcement_state::AnnouncementState>, persisted_announcement_state: Option<crate::session::announcement_state::AnnouncementState>,
memory_config: Option<crate::config::MemoryConfig>, memory_config: Option<crate::config::MemoryConfig>,
feedback_flags: crate::session::feedback_manager::FeedbackFlags, feedback_flags: crate::session::feedback_manager::FeedbackFlags,
managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle,
managed_mcp_expires_at: Option<chrono::DateTime<chrono::Utc>>,
managed_mcp_proxy_base_url: String,
session_model_id: acp::ModelId, session_model_id: acp::ModelId,
session_yolo_mode: bool, session_yolo_mode: bool,
session_auto_mode: bool, session_auto_mode: bool,
@@ -150,8 +147,6 @@ pub(crate) async fn spawn_session_actor(
max_retries: Option<u32>, max_retries: Option<u32>,
web_search_config: kigi_tools::implementations::WebSearchConfig, web_search_config: kigi_tools::implementations::WebSearchConfig,
web_fetch_config: kigi_tools::implementations::grok_build::web_fetch::WebFetchConfig, web_fetch_config: kigi_tools::implementations::grok_build::web_fetch::WebFetchConfig,
image_gen_config: kigi_tools::implementations::grok_build::image_gen::ImageGenConfig,
video_gen_config: kigi_tools::implementations::grok_build::video_gen::VideoGenConfig,
app_builder_deployer_config: kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig, app_builder_deployer_config: kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig,
write_file_enabled: bool, write_file_enabled: bool,
goal_enabled: bool, goal_enabled: bool,
@@ -660,14 +655,6 @@ pub(crate) async fn spawn_session_actor(
let context_window_tokens = context_window_override let context_window_tokens = context_window_override
.map(|c| c.get()) .map(|c| c.get())
.unwrap_or(sampling_config.context_window); .unwrap_or(sampling_config.context_window);
let managed_gateway_tool_client = auth_manager.as_ref().map(|am| {
kigi_tools::types::resources::ManagedGatewayToolClient(Arc::new(
ShellManagedGatewayToolClient {
proxy_base_url: managed_mcp_proxy_base_url.clone(),
auth_manager: am.clone(),
},
))
});
let mcp_state = { let mcp_state = {
let mut state = McpState::new_with_meta(mcp_servers.clone(), mcp_meta_config_map); let mut state = McpState::new_with_meta(mcp_servers.clone(), mcp_meta_config_map);
if let Some(ref pool) = parent_mcp_pool { if let Some(ref pool) = parent_mcp_pool {
@@ -711,8 +698,6 @@ pub(crate) async fn spawn_session_actor(
web_search_config: web_search_config.clone(), web_search_config: web_search_config.clone(),
backend_search: backend_tools_enabled, backend_search: backend_tools_enabled,
web_fetch_config: web_fetch_config.clone(), web_fetch_config: web_fetch_config.clone(),
image_gen_config: image_gen_config.clone(),
video_gen_config: video_gen_config.clone(),
app_builder_deployer_config: app_builder_deployer_config.clone(), app_builder_deployer_config: app_builder_deployer_config.clone(),
write_file_enabled, write_file_enabled,
subagents_enabled, subagents_enabled,
@@ -739,7 +724,6 @@ pub(crate) async fn spawn_session_actor(
respect_gitignore, respect_gitignore,
path_not_found_hints, path_not_found_hints,
mcp_state: mcp_state.clone(), mcp_state: mcp_state.clone(),
managed_gateway_tool_client: managed_gateway_tool_client.clone(),
is_non_interactive: startup_hints.non_interactive, is_non_interactive: startup_hints.non_interactive,
system_prompt_label, system_prompt_label,
owner_session_id: Some(session_info.id.0.to_string()), owner_session_id: Some(session_info.id.0.to_string()),
@@ -1125,8 +1109,6 @@ pub(crate) async fn spawn_session_actor(
goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false),
pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()),
goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false), goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false),
managed_mcp_handle,
managed_mcp_expires_at: std::sync::Mutex::new(managed_mcp_expires_at),
tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())),
mcp_announced_servers: Mutex::new( mcp_announced_servers: Mutex::new(
persisted_announcement_state persisted_announcement_state
@@ -1224,14 +1206,6 @@ pub(crate) async fn spawn_session_actor(
)) ))
.await; .await;
} }
if let Some(client) = managed_gateway_tool_client.clone() {
session
.agent
.borrow()
.tool_bridge()
.update_resource(client)
.await;
}
{ {
let plan_path = session.plan_mode.lock().plan_file_path().to_path_buf(); let plan_path = session.plan_mode.lock().plan_file_path().to_path_buf();
session session
@@ -1448,7 +1422,6 @@ pub(crate) async fn spawn_session_actor(
permission_handle: permissions_for_handle, permission_handle: permissions_for_handle,
attribution_callback: attribution_callback_for_handle, attribution_callback: attribution_callback_for_handle,
agent_name: agent_name_for_handle, agent_name: agent_name_for_handle,
managed_mcp_proxy_base_url,
session_default_agent_profile, session_default_agent_profile,
allowed_subagent_types: allowed_subagent_types_for_handle, allowed_subagent_types: allowed_subagent_types_for_handle,
hook_registry: hook_registry_for_handle, hook_registry: hook_registry_for_handle,
@@ -1541,9 +1514,6 @@ pub(crate) async fn spawn_session_on_thread(
persisted_announcement_state: Option<crate::session::announcement_state::AnnouncementState>, persisted_announcement_state: Option<crate::session::announcement_state::AnnouncementState>,
memory_config: Option<crate::config::MemoryConfig>, memory_config: Option<crate::config::MemoryConfig>,
feedback_flags: crate::session::feedback_manager::FeedbackFlags, feedback_flags: crate::session::feedback_manager::FeedbackFlags,
managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle,
managed_mcp_expires_at: Option<chrono::DateTime<chrono::Utc>>,
managed_mcp_proxy_base_url: String,
session_model_id: acp::ModelId, session_model_id: acp::ModelId,
session_yolo_mode: bool, session_yolo_mode: bool,
session_auto_mode: bool, session_auto_mode: bool,
@@ -1552,8 +1522,6 @@ pub(crate) async fn spawn_session_on_thread(
max_retries: Option<u32>, max_retries: Option<u32>,
web_search_config: kigi_tools::implementations::WebSearchConfig, web_search_config: kigi_tools::implementations::WebSearchConfig,
web_fetch_config: kigi_tools::implementations::grok_build::web_fetch::WebFetchConfig, web_fetch_config: kigi_tools::implementations::grok_build::web_fetch::WebFetchConfig,
image_gen_config: kigi_tools::implementations::grok_build::image_gen::ImageGenConfig,
video_gen_config: kigi_tools::implementations::grok_build::video_gen::VideoGenConfig,
app_builder_deployer_config: kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig, app_builder_deployer_config: kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig,
write_file_enabled: bool, write_file_enabled: bool,
goal_enabled: bool, goal_enabled: bool,
@@ -1688,9 +1656,6 @@ pub(crate) async fn spawn_session_on_thread(
persisted_announcement_state, persisted_announcement_state,
memory_config, memory_config,
feedback_flags, feedback_flags,
managed_mcp_handle,
managed_mcp_expires_at,
managed_mcp_proxy_base_url,
session_model_id, session_model_id,
session_yolo_mode, session_yolo_mode,
session_auto_mode, session_auto_mode,
@@ -1699,8 +1664,6 @@ pub(crate) async fn spawn_session_on_thread(
max_retries, max_retries,
web_search_config, web_search_config,
web_fetch_config, web_fetch_config,
image_gen_config,
video_gen_config,
app_builder_deployer_config, app_builder_deployer_config,
write_file_enabled, write_file_enabled,
goal_enabled, goal_enabled,
@@ -480,31 +480,13 @@ impl SessionActor {
} }
}); });
let _drainer_guard = crate::util::AbortOnDrop(drainer); let _drainer_guard = crate::util::AbortOnDrop(drainer);
while let Some((idx, mut result)) = dispatch_rx.recv().await { while let Some((idx, result)) = dispatch_rx.recv().await {
let prepared = approved_slots[idx] let prepared = approved_slots[idx]
.take() .take()
.expect("dispatch index should match an approved slot exactly once"); .expect("dispatch index should match an approved slot exactly once");
self.signals_handle().record_tool_call(&prepared.tool_name); self.signals_handle().record_tool_call(&prepared.tool_name);
let tool_start = self.events.tool_started(prepared.tool_name.clone()); let tool_start = self.events.tool_started(prepared.tool_name.clone());
let mut post_tool_use_result: Option<serde_json::Value> = None; let mut post_tool_use_result: Option<serde_json::Value> = None;
if let Some((server, _)) =
crate::session::mcp_servers::parse_mcp_tool_name(&prepared.tool_name)
&& server.starts_with(crate::session::managed_mcp::MANAGED_MCP_PREFIX)
{
let auth_rejected = match &result {
Err(err) => kigi_mcp::servers::is_auth_rejection_message(&err.to_string()),
Ok(tool_result) => {
tool_result.output.is_error()
&& kigi_mcp::servers::is_auth_rejection_message(
&tool_result.prompt_text,
)
}
};
if auth_rejected && self.reactive_managed_reauth(&server).await.is_ok() {
result = dispatch_tool(&self.workspace_ops, &prepared, &self.session_info.id.0)
.await;
}
}
let tool_result_size_bytes = match &result { let tool_result_size_bytes = match &result {
Ok(tool_result) => tool_result.prompt_text.len() as i64, Ok(tool_result) => tool_result.prompt_text.len() as i64,
Err(_) => 0, Err(_) => 0,
@@ -745,12 +727,6 @@ impl SessionActor {
} }
let mcp_parts = parse_mcp_tool_name(&call.function.name); let mcp_parts = parse_mcp_tool_name(&call.function.name);
let is_mcp_tool = mcp_parts.is_some(); let is_mcp_tool = mcp_parts.is_some();
if let Some((ref server, _)) = mcp_parts
&& server.starts_with(crate::session::managed_mcp::MANAGED_MCP_PREFIX)
{
let _span = tracing::info_span!("tool.refresh_managed_mcp").entered();
self.refresh_managed_mcp_if_stale().await;
}
if is_mcp_tool && !self.mcp_state.lock().await.is_initialized() { if is_mcp_tool && !self.mcp_state.lock().await.is_initialized() {
match self.mcp_strategy { match self.mcp_strategy {
McpInitStrategy::Blocking => { McpInitStrategy::Blocking => {
@@ -1483,33 +1459,6 @@ impl SessionActor {
vec![], vec![],
vec![], vec![],
), ),
ToolInput::ImageGen(ig) => (
format!("imagine: {}", ig.prompt),
acp::ToolKind::Other,
vec![],
vec![],
),
ToolInput::ImageEdit(ie) => (
format!("imagine-edit: {}", ie.prompt),
acp::ToolKind::Other,
vec![],
vec![],
),
ToolInput::ImageToVideo(i2v) => (
format!(
"image-to-video: {}",
i2v.prompt.as_deref().unwrap_or(&i2v.image)
),
acp::ToolKind::Other,
vec![],
vec![],
),
ToolInput::ReferenceToVideo(r2v) => (
format!("reference-to-video: {}", r2v.prompt),
acp::ToolKind::Other,
vec![],
vec![],
),
ToolInput::MCPTool(mcp_tool) => ( ToolInput::MCPTool(mcp_tool) => (
mcp_tool.tool_name.to_owned(), mcp_tool.tool_name.to_owned(),
acp::ToolKind::Other, acp::ToolKind::Other,
@@ -128,7 +128,7 @@ pub(super) fn resolve_session_shell() -> String {
} }
/// Key in `ToolError::details` that carries the HTTP status code. /// Key in `ToolError::details` that carries the HTTP status code.
/// Used by both error producers (image_gen, video_gen, test helpers) and /// Used by both error producers (tool HTTP clients, test helpers) and
/// the `is_auth_tool_error` classifier to avoid accidental key mismatch. /// the `is_auth_tool_error` classifier to avoid accidental key mismatch.
pub(crate) const HTTP_STATUS_DETAILS_KEY: &str = "status"; pub(crate) const HTTP_STATUS_DETAILS_KEY: &str = "status";
@@ -574,16 +574,6 @@ impl SessionActor {
.data(format!("failed to save user images to assets dir: {e}")) .data(format!("failed to save user images to assets dir: {e}"))
})? })?
}; };
let attached_image_refs = if self.is_cursor_harness() {
Vec::new()
} else {
crate::session::placeholder_images::attached_image_references(&user_images)
};
self.tool_bridge_handle()
.update_resource(kigi_tools::types::resources::AttachedImages(
attached_image_refs,
))
.await;
let prompt_text_for_hook = user_message.clone(); let prompt_text_for_hook = user_message.clone();
{ {
let origin = super::super::PromptOrigin::from_prompt_id(prompt_id); let origin = super::super::PromptOrigin::from_prompt_id(prompt_id);
@@ -232,8 +232,6 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false),
pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()),
goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false), goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false),
managed_mcp_handle: Default::default(),
managed_mcp_expires_at: std::sync::Mutex::new(None),
initial_client_mcp_servers: vec![], initial_client_mcp_servers: vec![],
tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())),
mcp_announced_servers: Mutex::new(HashMap::new()), mcp_announced_servers: Mutex::new(HashMap::new()),
@@ -670,8 +668,6 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false),
pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()),
goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false), goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false),
managed_mcp_handle: Default::default(),
managed_mcp_expires_at: std::sync::Mutex::new(None),
initial_client_mcp_servers: vec![], initial_client_mcp_servers: vec![],
tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())),
mcp_announced_servers: Mutex::new(HashMap::new()), mcp_announced_servers: Mutex::new(HashMap::new()),
@@ -917,8 +913,6 @@ async fn cancel_running_task_teardown_clears_running_and_pending_work() {
goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false),
pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()),
goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false), goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false),
managed_mcp_handle: Default::default(),
managed_mcp_expires_at: std::sync::Mutex::new(None),
initial_client_mcp_servers: vec![], initial_client_mcp_servers: vec![],
tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())),
mcp_announced_servers: Mutex::new(HashMap::new()), mcp_announced_servers: Mutex::new(HashMap::new()),
@@ -1897,8 +1891,6 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() {
goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false),
pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()),
goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false), goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false),
managed_mcp_handle: Default::default(),
managed_mcp_expires_at: std::sync::Mutex::new(None),
initial_client_mcp_servers: vec![], initial_client_mcp_servers: vec![],
tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())),
mcp_announced_servers: Mutex::new(HashMap::new()), mcp_announced_servers: Mutex::new(HashMap::new()),
@@ -58,8 +58,6 @@ async fn tool_bridge_routes_writes_through_injected_fs() {
web_search_config: Default::default(), web_search_config: Default::default(),
web_fetch_config: Default::default(), web_fetch_config: Default::default(),
lsp: None, lsp: None,
image_gen_config: Default::default(),
video_gen_config: Default::default(),
app_builder_deployer_config: Default::default(), app_builder_deployer_config: Default::default(),
api_key_provider: None, api_key_provider: None,
attribution_callback: None, attribution_callback: None,
@@ -261,8 +261,6 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false),
pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()),
goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false), goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false),
managed_mcp_handle: Default::default(),
managed_mcp_expires_at: std::sync::Mutex::new(None),
initial_client_mcp_servers: vec![], initial_client_mcp_servers: vec![],
tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())),
mcp_announced_servers: Mutex::new(HashMap::new()), mcp_announced_servers: Mutex::new(HashMap::new()),
@@ -190,8 +190,6 @@ async fn create_test_actor(
goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false),
pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()),
goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false), goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false),
managed_mcp_handle: Default::default(),
managed_mcp_expires_at: std::sync::Mutex::new(None),
initial_client_mcp_servers: vec![], initial_client_mcp_servers: vec![],
tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())),
mcp_announced_servers: Mutex::new(HashMap::new()), mcp_announced_servers: Mutex::new(HashMap::new()),
@@ -630,8 +628,6 @@ async fn create_test_actor_with_memory(
goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false),
pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()),
goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false), goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false),
managed_mcp_handle: Default::default(),
managed_mcp_expires_at: std::sync::Mutex::new(None),
initial_client_mcp_servers: vec![], initial_client_mcp_servers: vec![],
tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())),
mcp_announced_servers: Mutex::new(HashMap::new()), mcp_announced_servers: Mutex::new(HashMap::new()),
@@ -1384,8 +1380,6 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false),
pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()),
goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false), goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false),
managed_mcp_handle: Default::default(),
managed_mcp_expires_at: std::sync::Mutex::new(None),
initial_client_mcp_servers: vec![], initial_client_mcp_servers: vec![],
tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())),
mcp_announced_servers: Mutex::new(HashMap::new()), mcp_announced_servers: Mutex::new(HashMap::new()),
@@ -252,8 +252,6 @@ async fn create_test_actor_with_memory(
goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false),
pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()),
goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false), goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false),
managed_mcp_handle: Default::default(),
managed_mcp_expires_at: std::sync::Mutex::new(None),
initial_client_mcp_servers: vec![], initial_client_mcp_servers: vec![],
tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())),
mcp_announced_servers: Mutex::new(HashMap::new()), mcp_announced_servers: Mutex::new(HashMap::new()),
@@ -1,476 +0,0 @@
//! End-to-end coverage for the reactive managed-MCP re-auth flow that the
//! sibling `reactive_managed_reauth_tests.rs` only exercises at the guard-rail
//! level (owner-scope + cooldown). Here we drive the full
//! `invalidate_cache → get_or_fetch → refresh_managed_clients → re-handshake`
//! loop against a real in-process HTTP MCP server and assert the wire-visible
//! `x.ai/mcp/server_status` pushes that clients consuming only `server_status`
//! (not the `mcp/list` snapshot) depend on.
//!
//! Unlike the unit harness, these tests KEEP `gw_rx` so the forwarded
//! `McpServerStatusPayload`s (`ready`/`managed_token_refreshed` on recovery,
//! `needsauth`/`auth_expired` on terminal exhaustion) can be asserted.
//!
//! The mock is a hand-rolled axum streamable-HTTP server (same shape proven to
//! handshake against rmcp 2.1 in `kigi-mcp/tests/repro_sse_flood.rs`): a
//! `POST` that answers `initialize` + `tools/list` while `reject == false` and
//! `401`s while `reject == true`, plus a standing-GET SSE stream. A separate
//! `GET /mcp/configs` route stands in for the cli-chat-proxy managed-config backend fetch, so
//! the re-auth loop's proxy round-trip is real too.
use crate::session::acp_session::support::*;
use crate::session::acp_session::*;
use crate::session::managed_mcp::{MANAGED_MCP_PREFIX, ManagedMcpConfig};
use crate::session::mcp_dispatcher::{
McpServerStatus, McpServerStatusPayload, McpServerStatusReason, SERVER_STATUS_METHOD,
};
use agent_client_protocol as acp;
use axum::body::Body;
use axum::extract::State;
use axum::http::{StatusCode, header};
use axum::response::{IntoResponse, Response};
use axum::routing::{get, post};
use chrono::Utc;
use kigi_mcp::servers::{ClientStateKind, HttpConfig, McpClient};
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
const MANAGED: &str = "grok_com_testconnector";
// ── Mock cli-chat-proxy + MCP server ──────────────────────────────────────
#[derive(Clone)]
struct MockState {
/// While `true`, the MCP `POST` 401s — a silently-revoked token.
reject: Arc<AtomicBool>,
/// Count of managed-config fetches (`GET /mcp/configs`) so a cooldown-gated
/// attempt can be proven to skip the network entirely.
config_fetches: Arc<AtomicUsize>,
/// The server's own MCP endpoint, echoed back in the managed config so the
/// re-fetched config points the client at this same mock.
mcp_url: String,
}
/// Stand-in for cli-chat-proxy `GET /v1/mcp/configs`: always succeeds (a
/// revoked *connector* token still re-fetches a fresh *proxy* token), so the
/// re-handshake outcome is governed solely by the `reject` flag.
async fn handle_configs(State(s): State<MockState>) -> Response {
s.config_fetches.fetch_add(1, Ordering::Relaxed);
let body = serde_json::json!({
"mcp_servers": [{
"name": "testconnector",
"endpoint": s.mcp_url,
"headers": {"Authorization": "Bearer fresh"},
"token_expires_at": (Utc::now() + chrono::Duration::hours(1)).to_rfc3339(),
"scope": "workspace",
}]
});
axum::Json(body).into_response()
}
/// MCP streamable-HTTP `POST`. 401s while `reject`; otherwise a minimal valid
/// `initialize` result (with the mandatory `mcp-session-id` header) and an
/// empty `tools/list`, so a re-handshake succeeds.
async fn handle_mcp_post(
State(s): State<MockState>,
axum::Json(req): axum::Json<serde_json::Value>,
) -> Response {
if s.reject.load(Ordering::Relaxed) {
return StatusCode::UNAUTHORIZED.into_response();
}
match req["method"].as_str() {
Some("initialize") => {
let result = serde_json::json!({
"jsonrpc": "2.0",
"id": req["id"],
"result": {
"protocolVersion": req["params"]["protocolVersion"],
"capabilities": {},
"serverInfo": {"name": "mock", "version": "0.0.0"},
},
});
([("mcp-session-id", "mock-session-1")], axum::Json(result)).into_response()
}
Some("tools/list") => {
let result = serde_json::json!({
"jsonrpc": "2.0",
"id": req["id"],
"result": {"tools": []},
});
axum::Json(result).into_response()
}
// notifications/initialized and anything else.
_ => StatusCode::ACCEPTED.into_response(),
}
}
/// Standing-GET SSE stream that stays open (a healthy server) — rmcp opens it
/// after `initialize`; a pending body keeps it alive without reconnect churn.
async fn handle_mcp_get() -> Response {
(
[(header::CONTENT_TYPE, "text/event-stream")],
Body::from_stream(futures::stream::pending::<Result<String, std::io::Error>>()),
)
.into_response()
}
/// Bind on an ephemeral port and serve the mock. Returns the proxy base URL
/// (`http://addr` — the managed fetch appends `/mcp/configs`), the MCP endpoint
/// URL, and the fetch counter.
async fn spawn_mock(reject: Arc<AtomicBool>) -> (String, String, Arc<AtomicUsize>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind mock");
let addr = listener.local_addr().expect("addr");
let proxy_base = format!("http://{addr}");
let mcp_url = format!("http://{addr}/mcp");
let config_fetches = Arc::new(AtomicUsize::new(0));
let state = MockState {
reject,
config_fetches: config_fetches.clone(),
mcp_url: mcp_url.clone(),
};
let app = axum::Router::new()
.route("/mcp/configs", get(handle_configs))
.route("/mcp", post(handle_mcp_post).get(handle_mcp_get))
.with_state(state);
tokio::spawn(async move {
let _ = axum::serve(listener, app).await;
});
(proxy_base, mcp_url, config_fetches)
}
// ── Test wiring helpers ────────────────────────────────────────────────────
/// Build an actor with a live (disk-backed, env-free) `AuthManager` holding a
/// valid token and a `ModelsManager` whose cli-chat-proxy points at `proxy_base`
/// — the two pieces the default test actor lacks, both required for the inner
/// re-fetch to reach the mock instead of the real proxy.
async fn actor_with_proxy(
proxy_base: &str,
gw_tx: tokio::sync::mpsc::UnboundedSender<kigi_acp_lib::AcpClientMessage>,
) -> (SessionActor, tempfile::TempDir) {
let (persist_tx, _persist_rx) = tokio::sync::mpsc::unbounded_channel();
let mut actor = create_test_actor(100, 128_000, 80, gw_tx, persist_tx).await;
let home = tempfile::tempdir().expect("tempdir");
let auth_manager = Arc::new(crate::auth::AuthManager::new(
home.path(),
crate::auth::KimiCodeConfig::default(),
));
// Valid (1h) token in-memory only — `auth()` fast-paths it without network.
auth_manager.hot_swap(crate::auth::KimiAuth {
expires_at: Some(Utc::now() + chrono::Duration::hours(1)),
..crate::auth::KimiAuth::test_default()
});
let cfg = crate::agent::config::Config {
endpoints: crate::agent::config::EndpointsConfig {
coding_api_base_url: Some(proxy_base.to_string()),
..Default::default()
},
..Default::default()
};
actor.models_manager = crate::agent::models::ModelsManager::new(
None,
Default::default(),
acp::ModelId::new("default"),
auth_manager.clone(),
cfg,
);
actor.auth_manager = Some(auth_manager);
(actor, home)
}
/// Seed `actor` so it owns a managed client pointed at `mcp_url` with a STALE
/// token, plus the matching config entry `refresh_managed_clients` keys the
/// in-place swap on, plus a `Ready` managed cache.
async fn seed_managed(actor: &SessionActor, mcp_url: &str) {
{
let mut st = actor.mcp_state.lock().await;
// `refresh_managed_clients` matches the owned client to a fresh config
// by looking up `configs` for an Http server with the same name.
st.configs = vec![acp::McpServer::Http(
acp::McpServerHttp::new(MANAGED.to_string(), mcp_url.to_string()).headers(vec![]),
)];
st.owned_clients.insert(
MANAGED.to_string(),
Arc::new(McpClient::new_http(
MANAGED.to_string(),
HttpConfig {
url: mcp_url.to_string(),
headers: vec![("Authorization".into(), "Bearer stale".into())],
},
None,
None,
)),
);
}
let handle = actor.managed_mcp_handle.clone();
let mut configs = HashMap::new();
configs.insert("Authorization".to_string(), "Bearer fresh".to_string());
handle.lock().await.complete_fetch(
vec![ManagedMcpConfig {
name: "testconnector".to_string(),
endpoint: mcp_url.to_string(),
headers: configs,
token_expires_at: Some(Utc::now() + chrono::Duration::hours(1)),
scope: Some("workspace".to_string()),
scope_id: None,
scope_name: None,
}],
&handle,
None,
);
}
/// Drain all `x.ai/mcp/server_status` pushes currently queued on `gw_rx`.
fn drain_status_pushes(
gw_rx: &mut tokio::sync::mpsc::UnboundedReceiver<kigi_acp_lib::AcpClientMessage>,
) -> Vec<McpServerStatusPayload> {
let mut out = Vec::new();
while let Ok(msg) = gw_rx.try_recv() {
if let kigi_acp_lib::AcpClientMessage::ExtNotification(args) = msg
&& args.request.method.as_ref() == SERVER_STATUS_METHOD
&& let Ok(payload) =
serde_json::from_str::<McpServerStatusPayload>(args.request.params.get())
{
out.push(payload);
}
}
out
}
// ── Case 1: recover-on-second-fetch (the happy path) ───────────────────────
/// A managed token rejected on the first re-handshake (`reject = true`) but
/// fixed before the next attempt recovers end-to-end: the client lands `Ready`
/// and a `ready`/`managed_token_refreshed` status push hits the wire.
///
/// The cooldown clear between attempts stands in for the proactive refresh's
/// `clear_reauth_cooldowns` (the real path that re-enables a re-authorized
/// connector) — driving real wall-clock backoff would be slow and flaky.
#[tokio::test(flavor = "multi_thread")]
async fn recovers_on_second_attempt_and_pushes_managed_token_refreshed() {
let reject = Arc::new(AtomicBool::new(true));
let (proxy_base, mcp_url, _fetches) = spawn_mock(reject.clone()).await;
let (gw_tx, mut gw_rx) = tokio::sync::mpsc::unbounded_channel();
let local = tokio::task::LocalSet::new();
local
.run_until(async move {
// `create_test_actor` spawns a local terminal task, so the actor
// must be built inside the `LocalSet`.
let (actor, _home) = actor_with_proxy(&proxy_base, gw_tx).await;
seed_managed(&actor, &mcp_url).await;
// Attempt 1: the connector still 401s, so the re-handshake fails.
let first = actor.reactive_managed_reauth(MANAGED).await;
assert!(first.is_err(), "first attempt must fail while rejecting");
// Below the terminal cap, so no NeedsAuth push yet.
assert!(
drain_status_pushes(&mut gw_rx).is_empty(),
"a single non-terminal failure must not push a status",
);
// Connector re-authorized; clear the cooldown (proactive-refresh
// analog) and retry.
reject.store(false, Ordering::Relaxed);
actor
.managed_mcp_handle
.lock()
.await
.clear_reauth_cooldowns();
actor
.reactive_managed_reauth(MANAGED)
.await
.expect("second attempt must recover once the connector accepts");
// Client is Ready and no longer parked needs-auth.
let client = actor
.mcp_state
.lock()
.await
.get_client(MANAGED)
.cloned()
.expect("managed client present");
assert_eq!(
client.state_kind().await,
ClientStateKind::Ready,
"recovered client must be Ready",
);
assert!(
!actor.mcp_state.lock().await.auth_required.contains(MANAGED),
"recovered server must not be in auth_required",
);
let pushes = drain_status_pushes(&mut gw_rx);
assert!(
pushes.iter().any(|p| p.name == MANAGED
&& p.status == McpServerStatus::Ready
&& p.reason == McpServerStatusReason::ManagedTokenRefreshed),
"expected a ready/managed_token_refreshed push, got: {pushes:?}",
);
})
.await;
}
// ── Case 2: terminal after 3 failures ──────────────────────────────────────
/// Three consecutive failed re-auths park the connector in the terminal
/// `auth_required` state, push `needsauth`/`auth_expired`, and gate the 4th
/// immediate attempt (no extra config fetch).
///
/// Deterministic without sleeps: the default test actor has no `auth_manager`,
/// so the inner re-fetch fails fast (no network). Two prior failures are seeded
/// with an injected past `now` so their backoff windows are already elapsed
/// (the cooldown API takes `now`); the third — the real
/// `reactive_managed_reauth` call — is what crosses the terminal cap and fires
/// the NeedsAuth push.
#[tokio::test(flavor = "current_thread")]
async fn terminal_after_three_failures_pushes_needsauth_then_gates() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (gw_tx, mut gw_rx) = tokio::sync::mpsc::unbounded_channel();
let (persist_tx, _persist_rx) = tokio::sync::mpsc::unbounded_channel();
let actor = create_test_actor(100, 128_000, 80, gw_tx, persist_tx).await;
actor
.mcp_state
.lock()
.await
.owned_clients
.insert(MANAGED.to_string(), Arc::new(McpClient::stub(MANAGED)));
// Seed two prior failures whose backoff windows are already in the
// past (failures = 2, below the cap of 3, so still eligible).
let past = Utc::now() - chrono::Duration::hours(1);
{
let mut h = actor.managed_mcp_handle.lock().await;
h.record_reauth_failure(MANAGED, past);
h.record_reauth_failure(MANAGED, past);
assert!(
h.reauth_allowed(MANAGED, Utc::now()),
"elapsed window + below cap must be eligible",
);
assert!(!h.reauth_is_terminal(MANAGED));
}
// Third (real) attempt: inner fails fast (no auth manager) and
// records the terminal failure.
let err = actor
.reactive_managed_reauth(MANAGED)
.await
.expect_err("third attempt must fail");
assert!(err.contains("auth manager"), "got: {err}");
// Terminal: parked auth_required + NeedsAuth/auth_expired push.
assert!(
actor.mcp_state.lock().await.auth_required.contains(MANAGED),
"exhausted connector must be parked in auth_required",
);
assert!(
actor
.managed_mcp_handle
.lock()
.await
.reauth_is_terminal(MANAGED),
);
let pushes = drain_status_pushes(&mut gw_rx);
assert!(
pushes.iter().any(|p| p.name == MANAGED
&& p.status == McpServerStatus::NeedsAuth
&& p.reason == McpServerStatusReason::AuthExpired),
"expected a needsauth/auth_expired push, got: {pushes:?}",
);
// Fourth immediate attempt is cooldown-gated by the terminal state.
let err4 = actor
.reactive_managed_reauth(MANAGED)
.await
.expect_err("fourth attempt must be gated");
assert!(err4.contains("cooldown"), "got: {err4}");
assert!(
drain_status_pushes(&mut gw_rx).is_empty(),
"a gated attempt must not push another status",
);
})
.await;
}
// ── Case 3: entry-B classification ─────────────────────────────────────────
/// The mid-session entry-B gate routes an auth-rejection on a managed tool into
/// `reactive_managed_reauth` (observable as an armed cooldown) but leaves a
/// non-auth `Ok(is_error)` body (e.g. a 403 policy denial) untouched. Mirrors
/// the exact classifier (`is_auth_rejection_message`) + managed-prefix gate the
/// loop in `tool_calls.rs` keys on, plus the resulting side effect.
#[tokio::test(flavor = "current_thread")]
async fn entry_b_routes_auth_rejection_but_not_policy_denial() {
use kigi_mcp::servers::{is_auth_rejection_message, parse_mcp_tool_name};
// Managed-prefix gate: only `grok_com_*` tools enter entry-B.
let (managed_server, _) =
parse_mcp_tool_name(&format!("{MANAGED}__create_issue")).expect("qualified name");
assert!(managed_server.starts_with(MANAGED_MCP_PREFIX));
let (local_server, _) = parse_mcp_tool_name("github__create_issue").expect("qualified name");
assert!(!local_server.starts_with(MANAGED_MCP_PREFIX));
// Classification of both failure shapes entry-B sees:
// Err(ToolError) -> err.to_string()
// Ok(is_error) -> tool_result.prompt_text
assert!(
is_auth_rejection_message("MCP error: HTTP 401 Unauthorized"),
"401 error string must route into re-auth",
);
assert!(
is_auth_rejection_message("authentication required"),
"auth wording must route into re-auth",
);
assert!(
!is_auth_rejection_message("403 Forbidden: policy denied for this connector"),
"a 403 policy denial must NOT route into re-auth",
);
// Side-effect check: only the auth path arms the cooldown. An owner stub
// makes the inner re-fetch fail fast (no auth manager), arming the gate.
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (gw_tx, _gw_rx) = tokio::sync::mpsc::unbounded_channel();
let (persist_tx, _persist_rx) = tokio::sync::mpsc::unbounded_channel();
let actor = create_test_actor(100, 128_000, 80, gw_tx, persist_tx).await;
actor
.mcp_state
.lock()
.await
.owned_clients
.insert(MANAGED.to_string(), Arc::new(McpClient::stub(MANAGED)));
// Non-auth body classified false -> entry-B does NOT call re-auth,
// so the cooldown stays pristine.
assert!(
actor
.managed_mcp_handle
.lock()
.await
.reauth_allowed(MANAGED, Utc::now()),
"no re-auth call yet: cooldown must be clean",
);
// Auth body classified true -> entry-B calls re-auth; the failed
// attempt arms the cooldown window.
let _ = actor.reactive_managed_reauth(MANAGED).await;
assert!(
!actor
.managed_mcp_handle
.lock()
.await
.reauth_allowed(MANAGED, Utc::now()),
"the auth path must have armed the cooldown",
);
})
.await;
}
@@ -1,97 +0,0 @@
//! Focused tests for the reactive managed re-auth routine's deterministic guard
//! rails: owner-scoping and the per-server cooldown gate. The
//! full re-fetch + swap + re-handshake loop is covered at the unit level by the
//! `kigi-mcp` and `managed_mcp` tests; here we assert the `SessionActor`
//! wiring around those primitives.
use crate::session::acp_session::support::*;
use crate::session::acp_session::*;
use kigi_mcp::servers::McpClient;
use std::sync::Arc;
const MANAGED: &str = "grok_com_testconnector";
async fn make_actor() -> SessionActor {
let (gw_tx, _gw_rx) = tokio::sync::mpsc::unbounded_channel();
let (persist_tx, _persist_rx) = tokio::sync::mpsc::unbounded_channel();
create_test_actor(100, 128_000, 80, gw_tx, persist_tx).await
}
/// Owner-scoping: a session that does NOT own the managed client in
/// `owned_clients` must refuse the in-place swap (a subagent holds the
/// client as a shared Arc and recovers via the leader instead).
#[tokio::test(flavor = "current_thread")]
async fn reactive_managed_reauth_skips_non_owned_client() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let actor = make_actor().await;
let err = actor
.reactive_managed_reauth(MANAGED)
.await
.expect_err("non-owned client must not be re-auth'd");
assert!(
err.contains("does not own"),
"expected owner-scope rejection, got: {err}",
);
})
.await;
}
/// With the client owned but no `auth_manager` available, the inner
/// re-fetch fails fast (no network) and the outer routine records exactly
/// one cooldown failure — the next immediate attempt is gated, and the
/// server is not yet parked in the terminal `auth_required` state (one
/// failure is below the attempt cap).
#[tokio::test(flavor = "current_thread")]
async fn reactive_managed_reauth_records_cooldown_on_failure() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let actor = make_actor().await;
actor
.mcp_state
.lock()
.await
.owned_clients
.insert(MANAGED.to_string(), Arc::new(McpClient::stub(MANAGED)));
// First attempt: owner + cooldown gates pass, inner fails on the
// missing auth manager.
let err = actor
.reactive_managed_reauth(MANAGED)
.await
.expect_err("no auth manager → inner re-fetch must fail");
assert!(
err.contains("auth manager"),
"expected auth-manager failure, got: {err}",
);
// The failure armed the cooldown window, so an immediate second
// attempt is refused by the gate (not retried).
let err2 = actor
.reactive_managed_reauth(MANAGED)
.await
.expect_err("cooldown must gate the immediate retry");
assert!(
err2.contains("cooldown"),
"expected cooldown rejection, got: {err2}",
);
// A single failure is below the terminal attempt cap, so the
// server is not yet surfaced as needs-auth.
assert!(
!actor.mcp_state.lock().await.auth_required.contains(MANAGED),
"one failure must not park the server in auth_required",
);
assert!(
!actor
.managed_mcp_handle
.lock()
.await
.reauth_is_terminal(MANAGED),
"one failure must not be terminal",
);
})
.await;
}
@@ -198,8 +198,6 @@ pub(super) async fn make_replay_send_update_fixture() -> ReplaySendUpdateFixture
goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false),
pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()),
goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false), goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false),
managed_mcp_handle: Default::default(),
managed_mcp_expires_at: std::sync::Mutex::new(None),
initial_client_mcp_servers: vec![], initial_client_mcp_servers: vec![],
tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())),
mcp_announced_servers: Mutex::new(HashMap::new()), mcp_announced_servers: Mutex::new(HashMap::new()),
@@ -89,8 +89,6 @@ async fn test_agent_from_config(
web_search_config: Default::default(), web_search_config: Default::default(),
web_fetch_config: Default::default(), web_fetch_config: Default::default(),
lsp: None, lsp: None,
image_gen_config: Default::default(),
video_gen_config: Default::default(),
app_builder_deployer_config: Default::default(), app_builder_deployer_config: Default::default(),
api_key_provider: None, api_key_provider: None,
attribution_callback: None, attribution_callback: None,
@@ -305,8 +303,6 @@ pub(crate) async fn create_test_actor_ex(
goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false),
pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()),
goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false), goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false),
managed_mcp_handle: Default::default(),
managed_mcp_expires_at: std::sync::Mutex::new(None),
initial_client_mcp_servers: vec![], initial_client_mcp_servers: vec![],
tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())),
mcp_announced_servers: Mutex::new(HashMap::new()), mcp_announced_servers: Mutex::new(HashMap::new()),
@@ -77,7 +77,7 @@ fn err(msg: &str) -> Result<ToolRunResult, kigi_tool_runtime::ToolError> {
)) ))
} }
/// Production-shaped HTTP failure (image_gen / video_gen emit this on /// Production-shaped HTTP failure (tool HTTP clients emit this on
/// any non-success status). Use for retry tests that should exercise /// any non-success status). Use for retry tests that should exercise
/// the structured status-code path rather than the string fallback. /// the structured status-code path rather than the string fallback.
fn http_err(status: u16, msg: &str) -> Result<ToolRunResult, kigi_tool_runtime::ToolError> { fn http_err(status: u16, msg: &str) -> Result<ToolRunResult, kigi_tool_runtime::ToolError> {
@@ -98,14 +98,14 @@ fn is_auth_tool_error_classification() {
// (expected, error) — covers every branch + a sample of negatives // (expected, error) — covers every branch + a sample of negatives
// a careless edit could plausibly break. // a careless edit could plausibly break.
let cases: Vec<(bool, kigi_tool_runtime::ToolError)> = vec![ let cases: Vec<(bool, kigi_tool_runtime::ToolError)> = vec![
// Primary path: image_gen / video_gen now surface 401s as // Primary path: tool HTTP clients surface 401s as
// structured custom errors with status in details; classifier // structured custom errors with status in details; classifier
// matches the status code, not the rendered string. // matches the status code, not the rendered string.
( (
true, true,
kigi_tool_runtime::ToolError::new( kigi_tool_runtime::ToolError::new(
kigi_tool_runtime::ToolErrorKind::Custom, kigi_tool_runtime::ToolErrorKind::Custom,
"Image generation failed with HTTP 401 Unauthorized: missing token", "Tool call failed with HTTP 401 Unauthorized: missing token",
) )
.with_details( .with_details(
serde_json::json!({"code": "http_failure", HTTP_STATUS_DETAILS_KEY: 401}), serde_json::json!({"code": "http_failure", HTTP_STATUS_DETAILS_KEY: 401}),
@@ -175,14 +175,14 @@ fn is_auth_tool_error_classification() {
// Negative: transport failure must not trigger a token refresh. // Negative: transport failure must not trigger a token refresh.
( (
false, false,
kigi_tool_runtime::ToolError::invalid_arguments("Image generation timed out after 60s"), kigi_tool_runtime::ToolError::invalid_arguments("tool call timed out after 60s"),
), ),
// Negative: structural not-found error; not a network response. // Negative: structural not-found error; not a network response.
( (
false, false,
kigi_tool_runtime::ToolError::not_found( kigi_tool_runtime::ToolError::not_found(
kigi_tool_protocol::ToolId::new("image_gen").expect("valid"), kigi_tool_protocol::ToolId::new("no_such_tool").expect("valid"),
"Tool not found: image_gen", "Tool not found: no_such_tool",
), ),
), ),
// Negative: bare digits embedded in a request id must not trigger // Negative: bare digits embedded in a request id must not trigger
@@ -240,11 +240,11 @@ async fn auth_error_with_successful_refresh_retries() {
let am = succeeding_am(); let am = succeeding_am();
let calls = AtomicUsize::new(0); let calls = AtomicUsize::new(0);
let r = call_with_auth_retry(Some(&am), None, "image_gen", || { let r = call_with_auth_retry(Some(&am), None, "test_tool", || {
let n = calls.fetch_add(1, Ordering::SeqCst); let n = calls.fetch_add(1, Ordering::SeqCst);
async move { async move {
if n == 0 { if n == 0 {
http_err(401, "Image generation failed with HTTP 401 Unauthorized: x") http_err(401, "Tool call failed with HTTP 401 Unauthorized: x")
} else { } else {
ok_result("retried-ok") ok_result("retried-ok")
} }
@@ -299,7 +299,7 @@ async fn retry_is_bounded_at_one_even_if_retry_also_fails_with_auth() {
let r = call_with_auth_retry(Some(&am), None, "test_tool", || { let r = call_with_auth_retry(Some(&am), None, "test_tool", || {
calls.fetch_add(1, Ordering::SeqCst); calls.fetch_add(1, Ordering::SeqCst);
async { http_err(401, "Image generation failed with HTTP 401 Unauthorized: x") } async { http_err(401, "Tool call failed with HTTP 401 Unauthorized: x") }
}) })
.await; .await;
@@ -37,8 +37,6 @@ async fn web_search_errors_when_disabled() {
web_search_config: kigi_tools::implementations::web_search::WebSearchConfig::Disabled, web_search_config: kigi_tools::implementations::web_search::WebSearchConfig::Disabled,
web_fetch_config: Default::default(), web_fetch_config: Default::default(),
lsp: None, lsp: None,
image_gen_config: Default::default(),
video_gen_config: Default::default(),
app_builder_deployer_config: Default::default(), app_builder_deployer_config: Default::default(),
api_key_provider: None, api_key_provider: None,
attribution_callback: None, attribution_callback: None,
@@ -46,11 +46,9 @@ use kigi_agent::{Agent, AgentBuilder, CompactionPolicy, ReminderPolicy};
use kigi_tools::computer::types::{AsyncFileSystem, TerminalBackend}; use kigi_tools::computer::types::{AsyncFileSystem, TerminalBackend};
use kigi_tools::implementations::grok_build::ask_user_question::types::UserQuestionRequest; use kigi_tools::implementations::grok_build::ask_user_question::types::UserQuestionRequest;
use kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig; use kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig;
use kigi_tools::implementations::grok_build::image_gen::ImageGenConfig;
use kigi_tools::implementations::grok_build::task::types::{ use kigi_tools::implementations::grok_build::task::types::{
MonitorEventBuffer, SubagentEvent, TaskModelValidator, MonitorEventBuffer, SubagentEvent, TaskModelValidator,
}; };
use kigi_tools::implementations::grok_build::video_gen::VideoGenConfig;
use kigi_tools::implementations::grok_build::web_fetch::WebFetchConfig; use kigi_tools::implementations::grok_build::web_fetch::WebFetchConfig;
use kigi_tools::implementations::lsp::LspBackend; use kigi_tools::implementations::lsp::LspBackend;
use kigi_tools::implementations::web_search::WebSearchConfig; use kigi_tools::implementations::web_search::WebSearchConfig;
@@ -95,8 +93,6 @@ pub(crate) struct AgentRebuildSpec {
pub web_search_config: WebSearchConfig, pub web_search_config: WebSearchConfig,
pub backend_search: bool, pub backend_search: bool,
pub web_fetch_config: WebFetchConfig, pub web_fetch_config: WebFetchConfig,
pub image_gen_config: ImageGenConfig,
pub video_gen_config: VideoGenConfig,
pub app_builder_deployer_config: AppBuilderDeployerConfig, pub app_builder_deployer_config: AppBuilderDeployerConfig,
pub write_file_enabled: bool, pub write_file_enabled: bool,
pub subagents_enabled: bool, pub subagents_enabled: bool,
@@ -125,7 +121,6 @@ pub(crate) struct AgentRebuildSpec {
pub respect_gitignore: bool, pub respect_gitignore: bool,
pub path_not_found_hints: bool, pub path_not_found_hints: bool,
pub mcp_state: Arc<tokio::sync::Mutex<crate::session::mcp_servers::McpState>>, pub mcp_state: Arc<tokio::sync::Mutex<crate::session::mcp_servers::McpState>>,
pub managed_gateway_tool_client: Option<kigi_tools::types::resources::ManagedGatewayToolClient>,
pub is_non_interactive: bool, pub is_non_interactive: bool,
pub system_prompt_label: String, pub system_prompt_label: String,
pub owner_session_id: Option<String>, pub owner_session_id: Option<String>,
@@ -190,8 +185,6 @@ impl AgentRebuildSpec {
web_search_config, web_search_config,
backend_search, backend_search,
web_fetch_config, web_fetch_config,
image_gen_config,
video_gen_config,
app_builder_deployer_config, app_builder_deployer_config,
write_file_enabled, write_file_enabled,
subagents_enabled, subagents_enabled,
@@ -218,7 +211,6 @@ impl AgentRebuildSpec {
respect_gitignore, respect_gitignore,
path_not_found_hints, path_not_found_hints,
mcp_state, mcp_state,
managed_gateway_tool_client,
is_non_interactive, is_non_interactive,
system_prompt_label, system_prompt_label,
owner_session_id, owner_session_id,
@@ -244,8 +236,6 @@ impl AgentRebuildSpec {
.with_state_path(bridge_state_path.clone()) .with_state_path(bridge_state_path.clone())
.with_web_search_config(web_search_config.clone()) .with_web_search_config(web_search_config.clone())
.with_backend_search(*backend_search) .with_backend_search(*backend_search)
.with_image_gen_config(image_gen_config.clone())
.with_video_gen_config(video_gen_config.clone())
.with_app_builder_deployer_config(app_builder_deployer_config.clone()) .with_app_builder_deployer_config(app_builder_deployer_config.clone())
.with_web_fetch_config(web_fetch_config.clone()) .with_web_fetch_config(web_fetch_config.clone())
.with_write_file_enabled(*write_file_enabled) .with_write_file_enabled(*write_file_enabled)
@@ -351,9 +341,6 @@ impl AgentRebuildSpec {
*path_not_found_hints, *path_not_found_hints,
)) ))
.await; .await;
if let Some(client) = managed_gateway_tool_client.clone() {
agent.tool_bridge().update_resource(client).await;
}
{ {
use kigi_tools::implementations::grok_build::ask_user_question::UserQuestionSender; use kigi_tools::implementations::grok_build::ask_user_question::UserQuestionSender;
agent agent
@@ -392,8 +379,6 @@ pub(crate) fn test_rebuild_spec_default() -> Arc<AgentRebuildSpec> {
web_search_config: WebSearchConfig::default(), web_search_config: WebSearchConfig::default(),
backend_search: false, backend_search: false,
web_fetch_config: WebFetchConfig::Disabled, web_fetch_config: WebFetchConfig::Disabled,
image_gen_config: ImageGenConfig::default(),
video_gen_config: VideoGenConfig::default(),
app_builder_deployer_config: AppBuilderDeployerConfig::default(), app_builder_deployer_config: AppBuilderDeployerConfig::default(),
write_file_enabled: true, write_file_enabled: true,
subagents_enabled: false, subagents_enabled: false,
@@ -422,7 +407,6 @@ pub(crate) fn test_rebuild_spec_default() -> Arc<AgentRebuildSpec> {
mcp_state: Arc::new(tokio::sync::Mutex::new( mcp_state: Arc::new(tokio::sync::Mutex::new(
crate::session::mcp_servers::McpState::new(vec![]), crate::session::mcp_servers::McpState::new(vec![]),
)), )),
managed_gateway_tool_client: None,
is_non_interactive: false, is_non_interactive: false,
system_prompt_label: kigi_agent::DEFAULT_SYSTEM_PROMPT_LABEL.to_string(), system_prompt_label: kigi_agent::DEFAULT_SYSTEM_PROMPT_LABEL.to_string(),
owner_session_id: Some("test-session".to_string()), owner_session_id: Some("test-session".to_string()),
@@ -346,7 +346,7 @@ pub enum SessionCommand {
server_name: String, server_name: String,
enabled: bool, enabled: bool,
/// Fully-formed server config to add when re-enabling. Built by the /// Fully-formed server config to add when re-enabling. Built by the
/// caller via `merge_managed_mcp_servers` (with OAuth headers injected). /// caller via `merge_managed_mcp_servers`.
/// `None` when disabling. /// `None` when disabling.
server_config: Option<acp::McpServer>, server_config: Option<acp::McpServer>,
respond_to: oneshot::Sender<Result<(), acp::Error>>, respond_to: oneshot::Sender<Result<(), acp::Error>>,
@@ -357,17 +357,12 @@ pub enum SessionCommand {
server_name: String, server_name: String,
tool_name: String, tool_name: String,
enabled: bool, enabled: bool,
is_managed_gateway: bool,
respond_to: oneshot::Sender<Result<(), acp::Error>>, respond_to: oneshot::Sender<Result<(), acp::Error>>,
}, },
/// Read MCP status: which servers are configured, which clients are healthy, what tools. /// Read MCP status: which servers are configured, which clients are healthy, what tools.
GetMcpStatus { GetMcpStatus {
respond_to: oneshot::Sender<crate::extensions::mcp::McpStatusSnapshot>, respond_to: oneshot::Sender<crate::extensions::mcp::McpStatusSnapshot>,
}, },
GetManagedGatewayDisabledTools {
respond_to:
oneshot::Sender<std::collections::HashMap<String, std::collections::HashSet<String>>>,
},
/// Snapshot the session's live MCP client pool for subagent inheritance. /// Snapshot the session's live MCP client pool for subagent inheritance.
SnapshotMcpPool { SnapshotMcpPool {
respond_to: oneshot::Sender<Option<crate::session::mcp_servers::SharedMcpPool>>, respond_to: oneshot::Sender<Option<crate::session::mcp_servers::SharedMcpPool>>,
@@ -2288,8 +2288,6 @@ mod inline_auto_compact_flow_tests {
std::collections::VecDeque::new(), std::collections::VecDeque::new(),
), ),
goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false), goal_classifier_in_flight: std::sync::atomic::AtomicBool::new(false),
managed_mcp_handle: Default::default(),
managed_mcp_expires_at: std::sync::Mutex::new(None),
initial_client_mcp_servers: vec![], initial_client_mcp_servers: vec![],
tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())),
mcp_announced_servers: parking_lot::Mutex::new(std::collections::HashMap::new()), mcp_announced_servers: parking_lot::Mutex::new(std::collections::HashMap::new()),
@@ -130,7 +130,6 @@ pub struct SessionHandle {
pub attribution_callback: Option<kigi_sampler::SharedAttributionCallback>, pub attribution_callback: Option<kigi_sampler::SharedAttributionCallback>,
/// The agent definition name for this session. /// The agent definition name for this session.
pub agent_name: String, pub agent_name: String,
pub managed_mcp_proxy_base_url: String,
pub session_default_agent_profile: Option<String>, pub session_default_agent_profile: Option<String>,
/// Subagent types this agent can spawn (from Agent(t1, t2) in tools). /// Subagent types this agent can spawn (from Agent(t1, t2) in tools).
pub allowed_subagent_types: Option<Vec<String>>, pub allowed_subagent_types: Option<Vec<String>>,
@@ -398,25 +397,6 @@ impl SessionHandle {
server_name: String, server_name: String,
tool_name: String, tool_name: String,
enabled: bool, enabled: bool,
) -> Result<(), agent_client_protocol::Error> {
self.toggle_mcp_tool_with_source(server_name, tool_name, enabled, false)
.await
}
pub async fn toggle_managed_gateway_tool(
&self,
server_name: String,
tool_name: String,
enabled: bool,
) -> Result<(), agent_client_protocol::Error> {
self.toggle_mcp_tool_with_source(server_name, tool_name, enabled, true)
.await
}
async fn toggle_mcp_tool_with_source(
&self,
server_name: String,
tool_name: String,
enabled: bool,
is_managed_gateway: bool,
) -> Result<(), agent_client_protocol::Error> { ) -> Result<(), agent_client_protocol::Error> {
let (tx, rx) = oneshot::channel(); let (tx, rx) = oneshot::channel();
if self if self
@@ -425,7 +405,6 @@ impl SessionHandle {
server_name, server_name,
tool_name, tool_name,
enabled, enabled,
is_managed_gateway,
respond_to: tx, respond_to: tx,
}) })
.is_err() .is_err()
@@ -435,17 +414,6 @@ impl SessionHandle {
rx.await rx.await
.map_err(|_| agent_client_protocol::Error::internal_error().data("session closed"))? .map_err(|_| agent_client_protocol::Error::internal_error().data("session closed"))?
} }
pub async fn managed_gateway_disabled_tool_names(&self) -> HashMap<String, HashSet<String>> {
let (tx, rx) = oneshot::channel();
if self
.cmd_tx
.send(SessionCommand::GetManagedGatewayDisabledTools { respond_to: tx })
.is_err()
{
return HashMap::new();
}
rx.await.unwrap_or_default()
}
pub async fn retry_auth_required_servers(&self) { pub async fn retry_auth_required_servers(&self) {
let (tx, rx) = oneshot::channel(); let (tx, rx) = oneshot::channel();
if self if self
@@ -1,6 +1,6 @@
//! Shell-side managed MCP: merges MCP server sources, then injects managed //! Shell-side MCP server merging plus local managed-settings policy
//! OAuth headers, and binds the extracted credential/catalog machinery to //! (`managed-settings.json` allow/deny lists — local policy files, not a
//! shell's auth manager. //! remote service).
//! //!
//! Merge layers are applied in order; later `insert()` beats earlier //! Merge layers are applied in order; later `insert()` beats earlier
//! `or_insert()`: //! `or_insert()`:
@@ -9,59 +9,14 @@
//! - ~/.claude.json — `or_insert` (imported user/local MCP servers) //! - ~/.claude.json — `or_insert` (imported user/local MCP servers)
//! - `.mcp.json` — `or_insert` (team baseline) //! - `.mcp.json` — `or_insert` (team baseline)
//! - Client — `insert` (always wins) //! - Client — `insert` (always wins)
//! - Managed — header injection + auto-create missing connectors
//!
//! The transport/cache/injection core lives in
//! `kigi_shell_session_support::managed_mcp` and is re-exported here so
//! `crate::session::managed_mcp::…` paths keep resolving unchanged.
pub use kigi_shell_session_support::managed_mcp::*;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::Arc;
use agent_client_protocol as acp; use agent_client_protocol as acp;
/// Build a [`RefreshContext`] whose token provider resolves fresh tokens from /// Normalize a URL for dedup purposes (trailing slash dropped).
/// `auth_manager`; the extracted refresh task never sees the auth manager pub fn normalize_url(url: &str) -> String {
/// itself, only the closure. url.trim_end_matches('/').to_string()
fn refresh_context(
proxy_base_url: String,
auth_manager: Arc<crate::auth::AuthManager>,
) -> RefreshContext {
RefreshContext {
proxy_base_url,
token_provider: Arc::new(move || -> TokenFuture {
let auth_manager = auth_manager.clone();
Box::pin(async move { auth_manager.get_valid_token().await.ok() })
}),
}
}
/// Resolve an auth key from `auth_manager` then [`get_or_fetch`] the managed MCP
/// configs (with a [`RefreshContext`] for proactive refresh). Single source for
/// the auth-key dance across every managed-config fetch —
/// [`crate::agent::MvpAgent::get_managed_mcp_configs`], the interactive
/// folder-trust grant reload, agent-init MCP setup, and the reactive re-auth
/// re-fetch — so the copies can't drift.
/// Callers gate on `can_fetch_managed_mcps`/auth before calling.
pub(crate) async fn fetch_managed_mcp_configs(
handle: &ManagedMcpStateHandle,
proxy_url: &str,
auth_manager: &Arc<crate::auth::AuthManager>,
) -> Vec<ManagedMcpConfig> {
let auth_key = auth_manager
.get_valid_token()
.await
.ok()
.or_else(|| auth_manager.current_or_expired().map(|a| a.key));
get_or_fetch(
handle,
proxy_url,
auth_key.as_deref(),
Some(refresh_context(proxy_url.to_string(), auth_manager.clone())),
)
.await
} }
/// Dedup key for the merge map: normalized URL for Http/Sse, name for Stdio. /// Dedup key for the merge map: normalized URL for Http/Sse, name for Stdio.
@@ -88,17 +43,10 @@ pub(crate) fn mcp_server_name(s: &acp::McpServer) -> &str {
pub fn merge_managed_mcp_servers( pub fn merge_managed_mcp_servers(
client_mcp_servers: Vec<acp::McpServer>, client_mcp_servers: Vec<acp::McpServer>,
cwd: &std::path::Path, cwd: &std::path::Path,
managed_configs: &[ManagedMcpConfig],
plugin_registry: Option<&kigi_agent::plugins::PluginRegistry>, plugin_registry: Option<&kigi_agent::plugins::PluginRegistry>,
compat: &kigi_tools::types::compat::CompatConfig, compat: &kigi_tools::types::compat::CompatConfig,
) -> Vec<acp::McpServer> { ) -> Vec<acp::McpServer> {
merge_managed_mcp_servers_with_policy( merge_managed_mcp_servers_with_policy(client_mcp_servers, cwd, plugin_registry, compat)
client_mcp_servers,
cwd,
managed_configs,
plugin_registry,
compat,
)
.into_iter() .into_iter()
.filter(|s| s.disabled_reason.is_none()) .filter(|s| s.disabled_reason.is_none())
.map(|s| s.server) .map(|s| s.server)
@@ -108,7 +56,6 @@ pub fn merge_managed_mcp_servers(
pub fn merge_managed_mcp_servers_with_policy( pub fn merge_managed_mcp_servers_with_policy(
client_mcp_servers: Vec<acp::McpServer>, client_mcp_servers: Vec<acp::McpServer>,
cwd: &std::path::Path, cwd: &std::path::Path,
managed_configs: &[ManagedMcpConfig],
plugin_registry: Option<&kigi_agent::plugins::PluginRegistry>, plugin_registry: Option<&kigi_agent::plugins::PluginRegistry>,
compat: &kigi_tools::types::compat::CompatConfig, compat: &kigi_tools::types::compat::CompatConfig,
) -> Vec<McpServerWithPolicy> { ) -> Vec<McpServerWithPolicy> {
@@ -125,8 +72,6 @@ pub fn merge_managed_mcp_servers_with_policy(
let disabled = crate::util::config::disabled_mcp_server_names(cwd); let disabled = crate::util::config::disabled_mcp_server_names(cwd);
let mut merged: Vec<acp::McpServer> = servers.into_values().collect(); let mut merged: Vec<acp::McpServer> = servers.into_values().collect();
inject_managed_headers(&mut merged, managed_configs);
auto_inject_managed_servers_with_disabled(&mut merged, managed_configs, &disabled);
// Deterministic order: this list is collected from a HashMap (random // Deterministic order: this list is collected from a HashMap (random
// iteration order). Downstream equality checks (`mcp_servers_equal`, used // iteration order). Downstream equality checks (`mcp_servers_equal`, used
// by both `update_configs` and the `update_configs_diff` short-circuit) are // by both `update_configs` and the `update_configs_diff` short-circuit) are
@@ -349,59 +294,6 @@ pub fn merge_managed_mcp_servers_sourced(
servers.into_values().collect() servers.into_values().collect()
} }
/// Auto-create `grok_com_*` entries for managed configs not already in `merged`.
/// Dedup by display name (first scope wins). Skips names in `disabled_names`.
pub(crate) fn auto_inject_managed_servers_with_disabled(
merged: &mut Vec<acp::McpServer>,
managed_configs: &[ManagedMcpConfig],
disabled_names: &std::collections::HashSet<String>,
) {
if managed_configs.is_empty() {
return;
}
let existing_names: std::collections::HashSet<String> = merged
.iter()
.map(|s| mcp_server_name(s).to_owned())
.collect();
let mut seen_display_names: std::collections::HashSet<String> =
std::collections::HashSet::new();
let mut count = 0usize;
for config in managed_configs {
if config.headers.is_empty() {
continue;
}
let name = to_managed_name(&config.name);
if existing_names.contains(&name) {
continue;
}
if disabled_names.contains(&name) {
tracing::debug!(server_name = %name, "Auto-inject skipped: disabled in config.toml");
continue;
}
if !seen_display_names.insert(config.name.to_lowercase()) {
continue;
}
let headers = config
.headers
.iter()
.map(|(k, v)| acp::HttpHeader::new(k.clone(), v.clone()))
.collect();
merged.push(acp::McpServer::Http(
acp::McpServerHttp::new(name, config.endpoint.clone()).headers(headers),
));
count += 1;
}
if count > 0 {
tracing::info!(count, "Auto-injected managed MCP connectors");
}
}
fn load_plugin_mcp_servers( fn load_plugin_mcp_servers(
mcp_path: &std::path::Path, mcp_path: &std::path::Path,
plugin_name: &str, plugin_name: &str,
@@ -496,41 +388,10 @@ pub fn merge_plugin_oauth_into(
mod tests { mod tests {
use super::*; use super::*;
fn make_managed(name: &str, endpoint: &str, scope: &str) -> ManagedMcpConfig {
ManagedMcpConfig {
name: name.to_string(),
endpoint: endpoint.to_string(),
headers: HashMap::from([("Authorization".into(), "Bearer tok".into())]),
token_expires_at: None,
scope: Some(scope.to_string()),
scope_id: Some(format!("{scope}-id-123")),
scope_name: None,
}
}
fn empty_cwd() -> tempfile::TempDir { fn empty_cwd() -> tempfile::TempDir {
tempfile::tempdir().unwrap() tempfile::tempdir().unwrap()
} }
#[test]
fn auto_inject_creates_server_for_unmatched_managed_config() {
let managed = vec![make_managed("Slack", "https://mcp.slack.com/sse", "user")];
let cwd = empty_cwd();
let compat = kigi_tools::types::compat::CompatConfig::default();
let merged = merge_managed_mcp_servers(vec![], cwd.path(), &managed, None, &compat);
let slack = merged
.iter()
.find(|s| matches!(s, acp::McpServer::Http(acp::McpServerHttp { name, .. }) if name == "grok_com_slack"));
let slack = slack.expect("should have auto-injected grok_com_slack");
match slack {
acp::McpServer::Http(acp::McpServerHttp { url, headers, .. }) => {
assert_eq!(url, "https://mcp.slack.com/sse");
assert!(headers.iter().any(|h| h.name == "Authorization"));
}
other => panic!("expected Http server, got {:?}", other),
}
}
/// A client-provided server (e.g. a client session binding injected at /// A client-provided server (e.g. a client session binding injected at
/// `session/new`) exists in no on-disk config and no managed catalog — /// `session/new`) exists in no on-disk config and no managed catalog —
/// the merge must keep it. Config hot-reload handlers /// the merge must keep it. Config hot-reload handlers
@@ -549,7 +410,7 @@ mod tests {
)]; )];
let cwd = empty_cwd(); let cwd = empty_cwd();
let compat = kigi_tools::types::compat::CompatConfig::default(); let compat = kigi_tools::types::compat::CompatConfig::default();
let merged = merge_managed_mcp_servers(client, cwd.path(), &[], None, &compat); let merged = merge_managed_mcp_servers(client, cwd.path(), None, &compat);
assert!( assert!(
merged.iter().any(|s| matches!( merged.iter().any(|s| matches!(
s, s,
@@ -688,112 +549,6 @@ mod tests {
assert_eq!(surviving, ["slackbot"]); assert_eq!(surviving, ["slackbot"]);
} }
/// Drive the expectation through [`to_managed_name`] (not a hand-written
/// literal) so this fails if policy/runtime name normalization ever drifts.
#[test]
fn policy_server_name_matches_to_managed_name_transform() {
use kigi_workspace::permission::resolution::{AllowedMcpServer, McpServerAllowlist};
let managed_server = |runtime: &str| {
vec![acp::McpServer::Http(
acp::McpServerHttp::new(runtime.to_string(), "https://mcp.example.com/sse")
.headers(vec![]),
)]
};
let name_entry = |display: &str| AllowedMcpServer::Name {
name: display.to_string(),
};
let source = || Some(std::path::PathBuf::from("/test/managed-settings.json"));
for display in ["Slack", "My Server"] {
let runtime = to_managed_name(display);
let deny = McpServerAllowlist::new(vec![], vec![name_entry(display)], source());
let tagged = apply_mcp_server_policy(
managed_server(&runtime),
&std::collections::HashSet::new(),
&deny,
);
assert!(
matches!(
tagged[0].disabled_reason,
Some(McpDisabledReason::Denylist { .. })
),
"deny serverName {display:?} must block runtime {runtime:?}, got {:?}",
tagged[0].disabled_reason
);
let allow = McpServerAllowlist::new(vec![name_entry(display)], vec![], source());
let tagged = apply_mcp_server_policy(
managed_server(&runtime),
&std::collections::HashSet::new(),
&allow,
);
assert!(
tagged[0].disabled_reason.is_none(),
"allow serverName {display:?} must keep runtime {runtime:?}, got {:?}",
tagged[0].disabled_reason
);
}
}
#[test]
fn auto_inject_dedup_by_display_name_first_scope_wins() {
let managed = vec![
make_managed("Linear", "https://mcp.linear.app", "user"),
make_managed("Linear", "https://mcp.linear.app", "team"),
];
let cwd = empty_cwd();
let compat = kigi_tools::types::compat::CompatConfig::default();
let merged = merge_managed_mcp_servers(vec![], cwd.path(), &managed, None, &compat);
let linear_count = merged
.iter()
.filter(|s| matches!(s, acp::McpServer::Http(acp::McpServerHttp { name, .. }) if name == "grok_com_linear"))
.count();
assert_eq!(linear_count, 1, "should dedup by display name");
}
#[test]
fn auto_inject_skips_existing_server() {
let managed = vec![make_managed("Slack", "https://mcp.slack.com/sse", "user")];
let client = vec![acp::McpServer::Http(
acp::McpServerHttp::new(
"grok_com_slack".to_string(),
"https://mcp.slack.com/sse".to_string(),
)
.headers(vec![]),
)];
let cwd = empty_cwd();
let compat = kigi_tools::types::compat::CompatConfig::default();
let merged = merge_managed_mcp_servers(client, cwd.path(), &managed, None, &compat);
let slack_count = merged
.iter()
.filter(|s| matches!(s, acp::McpServer::Http(acp::McpServerHttp { name, .. }) if name == "grok_com_slack"))
.count();
assert_eq!(slack_count, 1, "should not duplicate existing server");
}
#[test]
fn auto_inject_skips_disabled() {
let managed = vec![
make_managed("Slack", "https://mcp.slack.com/sse", "user"),
make_managed("Linear", "https://mcp.linear.app", "user"),
];
let disabled: std::collections::HashSet<String> =
["grok_com_slack".to_string()].into_iter().collect();
let mut merged = vec![];
auto_inject_managed_servers_with_disabled(&mut merged, &managed, &disabled);
let has_slack = merged
.iter()
.any(|s| matches!(s, acp::McpServer::Http(acp::McpServerHttp { name, .. }) if name == "grok_com_slack"));
let has_linear = merged
.iter()
.any(|s| matches!(s, acp::McpServer::Http(acp::McpServerHttp { name, .. }) if name == "grok_com_linear"));
assert!(!has_slack, "disabled connector should be skipped");
assert!(has_linear, "non-disabled connector should be injected");
}
#[test] #[test]
fn lower_precedence_http_servers_are_blocked_by_toml_name_claims() { fn lower_precedence_http_servers_are_blocked_by_toml_name_claims() {
let cwd = tempfile::tempdir().unwrap(); let cwd = tempfile::tempdir().unwrap();
@@ -821,7 +576,7 @@ enabled = false
.unwrap(); .unwrap();
let compat = kigi_tools::types::compat::CompatConfig::default(); let compat = kigi_tools::types::compat::CompatConfig::default();
let merged = merge_managed_mcp_servers(vec![], cwd.path(), &[], None, &compat); let merged = merge_managed_mcp_servers(vec![], cwd.path(), None, &compat);
assert!( assert!(
!merged.iter().any(|server| matches!( !merged.iter().any(|server| matches!(
server, server,
@@ -859,7 +614,7 @@ enabled = false
) )
.headers(vec![]), .headers(vec![]),
)]; )];
let merged = merge_managed_mcp_servers(client, untrusted.path(), &[], None, &compat); let merged = merge_managed_mcp_servers(client, untrusted.path(), None, &compat);
assert!( assert!(
!merged.iter().any(|s| mcp_server_name(s) == "projsrv"), !merged.iter().any(|s| mcp_server_name(s) == "projsrv"),
"untrusted workspace must drop its repo-local MCP server" "untrusted workspace must drop its repo-local MCP server"
@@ -871,7 +626,7 @@ enabled = false
let trusted = repo_with_project_server(); let trusted = repo_with_project_server();
crate::agent::folder_trust::record_for_test(trusted.path(), true); crate::agent::folder_trust::record_for_test(trusted.path(), true);
let merged = merge_managed_mcp_servers(vec![], trusted.path(), &[], None, &compat); let merged = merge_managed_mcp_servers(vec![], trusted.path(), None, &compat);
assert!( assert!(
merged.iter().any(|s| mcp_server_name(s) == "projsrv"), merged.iter().any(|s| mcp_server_name(s) == "projsrv"),
"trusted workspace must keep its repo-local MCP server" "trusted workspace must keep its repo-local MCP server"
@@ -940,8 +695,7 @@ enabled = false
let disabled: std::collections::HashSet<String> = let disabled: std::collections::HashSet<String> =
["test-server".to_string()].into_iter().collect(); ["test-server".to_string()].into_iter().collect();
// auto_inject_managed_servers_with_disabled is for managed servers; // For plugin servers, the disabled check happens during merge.
// for plugin servers, the disabled check happens during merge.
// Verify the server name matches what would be checked. // Verify the server name matches what would be checked.
assert!( assert!(
disabled.contains("test-server"), disabled.contains("test-server"),
@@ -7,25 +7,14 @@
//! Local MCP writes are upsert-only — folders for servers removed mid-session are //! Local MCP writes are upsert-only — folders for servers removed mid-session are
//! not pruned (cleaned on the next session's first-turn build); pruning against //! not pruned (cleaned on the next session's first-turn build); pruning against
//! an async-changing client set risks deleting a just-connected server's folder. //! an async-changing client set risks deleting a just-connected server's folder.
//! Managed gateway writes converge to the admitted gateway catalog so disabled
//! gateway tools are not discoverable from stale descriptor files.
//! //!
//! Owning the descriptor I/O here keeps `acp_session.rs` thin. //! Owning the descriptor I/O here keeps `acp_session.rs` thin.
use std::collections::{BTreeMap, BTreeSet, HashSet};
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use std::sync::Arc; use std::sync::Arc;
use crate::session::mcp_servers::{McpClient, sanitize_descriptor_segment}; use crate::session::mcp_servers::{McpClient, sanitize_descriptor_segment};
#[derive(Debug, Clone)]
pub(crate) struct GatewayToolDescriptor {
pub(crate) connector_id: String,
pub(crate) tool_id: String,
pub(crate) description: String,
pub(crate) json_schema: serde_json::Value,
}
/// Per-server descriptor folder: `<mcps_root>/<sanitized server name>`. Uses the /// Per-server descriptor folder: `<mcps_root>/<sanitized server name>`. Uses the
/// sanitizer shared with `kigi-mcp` so the advertised folder matches disk. /// sanitizer shared with `kigi-mcp` so the advertised folder matches disk.
pub(crate) fn server_descriptor_dir(mcps_root: &Path, server_name: &str) -> PathBuf { pub(crate) fn server_descriptor_dir(mcps_root: &Path, server_name: &str) -> PathBuf {
@@ -55,192 +44,12 @@ pub(crate) async fn materialize_descriptors_for_clients(
} }
} }
pub(crate) async fn materialize_descriptors_for_gateway_tools(
mcps_root: &Path,
tools: Vec<GatewayToolDescriptor>,
gateway_connectors: Vec<String>,
protected_connectors: HashSet<String>,
) {
let mut files_by_connector: BTreeMap<String, Vec<(String, Vec<u8>)>> = BTreeMap::new();
for tool in tools {
let descriptor = serde_json::json!({
"name": tool.tool_id,
"description": tool.description,
"inputSchema": tool.json_schema,
});
match serde_json::to_vec_pretty(&descriptor) {
Ok(bytes) => {
let file_name = format!("{}.json", sanitize_descriptor_segment(&tool.tool_id));
files_by_connector
.entry(tool.connector_id)
.or_default()
.push((file_name, bytes));
}
Err(e) => tracing::warn!(
connector = %tool.connector_id,
tool = %tool.tool_id,
error = %e,
"failed to serialize managed gateway tool descriptor",
),
}
}
let connectors: BTreeSet<String> = gateway_connectors.into_iter().collect();
let mcps_root = mcps_root.to_path_buf();
if let Err(e) = tokio::task::spawn_blocking(move || {
for connector_id in connectors {
let server_dir = server_descriptor_dir(&mcps_root, &connector_id);
let tools_dir = server_dir.join("tools");
let files = files_by_connector.remove(&connector_id).unwrap_or_default();
let admitted: BTreeSet<String> = files.iter().map(|(name, _)| name.clone()).collect();
let protected = protected_connectors.contains(&connector_id);
if protected {
if let Err(e) = remove_gateway_owned_tool_descriptors(&tools_dir)
&& e.kind() != std::io::ErrorKind::NotFound
{
tracing::warn!(
connector = %connector_id,
path = %tools_dir.display(),
error = %e,
"failed to remove managed gateway descriptors from protected connector dir",
);
}
continue;
}
if files.is_empty() {
if let Err(e) = std::fs::remove_dir_all(&server_dir)
&& e.kind() != std::io::ErrorKind::NotFound
{
tracing::warn!(
connector = %connector_id,
path = %server_dir.display(),
error = %e,
"failed to remove stale managed gateway descriptor dir",
);
}
continue;
}
if let Err(e) = std::fs::create_dir_all(&tools_dir) {
tracing::warn!(
connector = %connector_id,
path = %tools_dir.display(),
error = %e,
"failed to create managed gateway tools descriptor dir",
);
continue;
}
for (file_name, bytes) in files {
let path = tools_dir.join(&file_name);
let write_result =
tempfile::NamedTempFile::new_in(&tools_dir).and_then(|mut tmp| {
std::io::Write::write_all(&mut tmp, &bytes)?;
tmp.persist(&path).map_err(|e| e.error)
});
if let Err(e) = write_result {
tracing::warn!(
connector = %connector_id,
path = %path.display(),
error = %e,
"failed to write managed gateway tool descriptor",
);
}
}
prune_stale_gateway_tool_descriptors(&tools_dir, &admitted, &connector_id);
}
})
.await
{
tracing::warn!(
error = %e,
"managed gateway descriptor write task panicked",
);
}
}
fn remove_gateway_owned_tool_descriptors(tools_dir: &Path) -> std::io::Result<()> {
let entries = std::fs::read_dir(tools_dir)?;
for entry in entries.flatten() {
let path = entry.path();
if path.extension().and_then(|ext| ext.to_str()) != Some("json") {
continue;
}
if is_gateway_owned_descriptor(&path) {
std::fs::remove_file(&path)?;
}
}
Ok(())
}
fn is_gateway_owned_descriptor(path: &Path) -> bool {
let Ok(contents) = std::fs::read(path) else {
return false;
};
let Ok(value) = serde_json::from_slice::<serde_json::Value>(&contents) else {
return false;
};
value
.get("x-grok-managed-gateway")
.and_then(|v| v.as_bool())
== Some(true)
}
fn prune_stale_gateway_tool_descriptors(
tools_dir: &Path,
admitted: &BTreeSet<String>,
connector_id: &str,
) {
let entries = match std::fs::read_dir(tools_dir) {
Ok(entries) => entries,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return,
Err(e) => {
tracing::warn!(
connector = %connector_id,
path = %tools_dir.display(),
error = %e,
"failed to read managed gateway tools descriptor dir for pruning",
);
return;
}
};
for entry in entries.flatten() {
let path = entry.path();
if path.extension().and_then(|ext| ext.to_str()) != Some("json") {
continue;
}
let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else {
continue;
};
if admitted.contains(file_name) {
continue;
}
if let Err(e) = std::fs::remove_file(&path) {
tracing::warn!(
connector = %connector_id,
path = %path.display(),
error = %e,
"failed to remove stale managed gateway tool descriptor",
);
}
}
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
#[test] #[test]
fn sanitize_replaces_unsafe_chars_and_never_empty() { fn sanitize_replaces_unsafe_chars_and_never_empty() {
assert_eq!(
sanitize_descriptor_segment("grok_com_linear"),
"grok_com_linear"
);
assert_eq!(sanitize_descriptor_segment("a/b:c d"), "a_b_c_d"); assert_eq!(sanitize_descriptor_segment("a/b:c d"), "a_b_c_d");
assert_eq!(sanitize_descriptor_segment(""), "_"); assert_eq!(sanitize_descriptor_segment(""), "_");
assert_eq!(sanitize_descriptor_segment("keep-1.2_x"), "keep-1.2_x"); assert_eq!(sanitize_descriptor_segment("keep-1.2_x"), "keep-1.2_x");
@@ -251,167 +60,4 @@ mod tests {
let root = Path::new("/home/u/.kigi/projects/enc/mcps"); let root = Path::new("/home/u/.kigi/projects/enc/mcps");
assert_eq!(server_descriptor_dir(root, "vercel"), root.join("vercel")); assert_eq!(server_descriptor_dir(root, "vercel"), root.join("vercel"));
} }
#[tokio::test]
async fn materializes_gateway_tool_descriptors_by_connector_and_tool() {
let temp = tempfile::tempdir().unwrap();
materialize_descriptors_for_gateway_tools(
temp.path(),
vec![
GatewayToolDescriptor {
connector_id: "linear".into(),
tool_id: "list_issues".into(),
description: "List issues".into(),
json_schema: serde_json::json!({
"type": "object",
"properties": {"limit": {"type": "number"}}
}),
},
GatewayToolDescriptor {
connector_id: "slack/team".into(),
tool_id: "search messages".into(),
description: "Search Slack".into(),
json_schema: serde_json::json!({
"type": "object",
"properties": {"query": {"type": "string"}}
}),
},
],
vec!["linear".into(), "slack/team".into()],
HashSet::new(),
)
.await;
let linear: serde_json::Value = serde_json::from_str(
&std::fs::read_to_string(temp.path().join("linear/tools/list_issues.json")).unwrap(),
)
.unwrap();
assert_eq!(linear["name"], "list_issues");
assert_eq!(linear["description"], "List issues");
assert_eq!(
linear["inputSchema"]["properties"]["limit"]["type"],
"number"
);
let slack: serde_json::Value = serde_json::from_str(
&std::fs::read_to_string(temp.path().join("slack_team/tools/search_messages.json"))
.unwrap(),
)
.unwrap();
assert_eq!(slack["name"], "search messages");
assert_eq!(slack["description"], "Search Slack");
assert_eq!(
slack["inputSchema"]["properties"]["query"]["type"],
"string"
);
}
#[tokio::test]
async fn gateway_descriptor_collision_prunes_only_gateway_owned_files() {
let temp = tempfile::tempdir().unwrap();
let tools_dir = temp.path().join("linear/tools");
std::fs::create_dir_all(&tools_dir).unwrap();
std::fs::write(
tools_dir.join("local_tool.json"),
serde_json::to_vec_pretty(&serde_json::json!({
"name": "local_tool",
"description": "Local tool",
"inputSchema": {"type": "object"}
}))
.unwrap(),
)
.unwrap();
std::fs::write(
tools_dir.join("gateway_tool.json"),
serde_json::to_vec_pretty(&serde_json::json!({
"name": "gateway_tool",
"description": "Gateway tool",
"inputSchema": {"type": "object"},
"x-grok-managed-gateway": true,
}))
.unwrap(),
)
.unwrap();
assert!(tools_dir.join("gateway_tool.json").exists());
materialize_descriptors_for_gateway_tools(
temp.path(),
Vec::new(),
vec!["linear".into()],
HashSet::from(["linear".into()]),
)
.await;
assert!(tools_dir.join("local_tool.json").exists());
assert!(!tools_dir.join("gateway_tool.json").exists());
}
#[tokio::test]
async fn protected_gateway_connector_skips_writing_into_local_server_dir() {
let temp = tempfile::tempdir().unwrap();
materialize_descriptors_for_gateway_tools(
temp.path(),
vec![GatewayToolDescriptor {
connector_id: "linear".into(),
tool_id: "gateway_tool".into(),
description: "Gateway tool".into(),
json_schema: serde_json::json!({"type": "object"}),
}],
vec!["linear".into()],
HashSet::from(["linear".into()]),
)
.await;
assert!(!temp.path().join("linear/tools/gateway_tool.json").exists());
}
#[tokio::test]
async fn prunes_stale_gateway_tool_descriptors() {
let temp = tempfile::tempdir().unwrap();
materialize_descriptors_for_gateway_tools(
temp.path(),
vec![
GatewayToolDescriptor {
connector_id: "linear".into(),
tool_id: "list_issues".into(),
description: "List issues".into(),
json_schema: serde_json::json!({"type": "object"}),
},
GatewayToolDescriptor {
connector_id: "slack".into(),
tool_id: "search".into(),
description: "Search Slack".into(),
json_schema: serde_json::json!({"type": "object"}),
},
],
vec!["linear".into(), "slack".into()],
HashSet::new(),
)
.await;
assert!(temp.path().join("linear/tools/list_issues.json").exists());
assert!(temp.path().join("slack/tools/search.json").exists());
materialize_descriptors_for_gateway_tools(
temp.path(),
vec![GatewayToolDescriptor {
connector_id: "linear".into(),
tool_id: "list_issues".into(),
description: "List issues".into(),
json_schema: serde_json::json!({"type": "object"}),
}],
vec!["linear".into(), "slack".into()],
HashSet::new(),
)
.await;
assert!(temp.path().join("linear/tools/list_issues.json").exists());
assert!(!temp.path().join("slack/tools/search.json").exists());
assert!(!temp.path().join("slack").exists());
materialize_descriptors_for_gateway_tools(
temp.path(),
Vec::new(),
vec!["linear".into()],
HashSet::new(),
)
.await;
assert!(!temp.path().join("linear").exists());
}
} }
@@ -50,7 +50,6 @@ use tokio::sync::Mutex as TokioMutex;
use tokio::sync::mpsc::UnboundedReceiver; use tokio::sync::mpsc::UnboundedReceiver;
use crate::extensions::mcp::McpServerSource; use crate::extensions::mcp::McpServerSource;
use crate::session::managed_mcp::MANAGED_MCP_PREFIX;
/// Tumbling-window coalescing period. See module doc. /// Tumbling-window coalescing period. See module doc.
pub const COALESCE_WINDOW: Duration = Duration::from_millis(50); pub const COALESCE_WINDOW: Duration = Duration::from_millis(50);
@@ -65,10 +64,9 @@ pub const SERVER_STATUS_METHOD: &str = "x.ai/mcp/server_status";
pub struct McpServerStatusPayload { pub struct McpServerStatusPayload {
/// Owning session id. /// Owning session id.
pub session_id: String, pub session_id: String,
/// MCP server name (`grok_com_linear`, `github`, ...). /// MCP server name (`github`, ...).
pub name: String, pub name: String,
/// `managed` (sourced from cli-chat-proxy / `grok_com_` prefix) /// Always `local` (user `.kigi/config.toml` and friends).
/// or `local` (user `.kigi/config.toml`).
pub source: McpServerSource, pub source: McpServerSource,
/// Current status — see [`McpServerStatus`]. /// Current status — see [`McpServerStatus`].
pub status: McpServerStatus, pub status: McpServerStatus,
@@ -131,22 +129,12 @@ pub enum McpServerStatusReason {
RestartSucceeded, RestartSucceeded,
/// The auto-restart path exhausted retries. /// The auto-restart path exhausted retries.
RestartFailed, RestartFailed,
/// A managed connector's reactive re-auth re-fetched a fresh token,
/// swapped the client, and re-handshook successfully. Distinct from
/// `RestartSucceeded` (reserved for the transport-close auto-restart
/// path) so a recovered managed token is observable on the wire.
ManagedTokenRefreshed,
} }
/// Build [`McpServerSource`] from a server name. Mirrors the /// Build [`McpServerSource`] from a server name. All servers are
/// existing convention used by `build_mcp_catalog` and friends: /// locally configured.
/// names with the `MANAGED_MCP_PREFIX` prefix are managed. pub fn classify_source(_name: &str) -> McpServerSource {
pub fn classify_source(name: &str) -> McpServerSource {
if name.starts_with(MANAGED_MCP_PREFIX) {
McpServerSource::Managed
} else {
McpServerSource::Local McpServerSource::Local
}
} }
/// State for the dispatcher's "intentional teardown" tracking. /// State for the dispatcher's "intentional teardown" tracking.
@@ -392,22 +380,6 @@ pub fn build_payload(
McpServerStatusReason::TransportClosed, McpServerStatusReason::TransportClosed,
None, None,
), ),
// A managed connector rejected for auth reasons surfaces as
// NeedsAuth ("visit grok.com"), not a generic Unavailable, so a
// client consuming only `server_status` (not the `mcp/list`
// `auth_required` boolean) shows the correct terminal state. Uses
// the same `is_auth_rejection_message` classifier the reroute and
// the reactive recovery path key on, so they cannot drift.
(McpClientEventKind::HandshakeFailed, McpClientEvent::HandshakeFailed { reason, .. })
if source == McpServerSource::Managed
&& kigi_mcp::servers::is_auth_rejection_message(reason) =>
{
(
McpServerStatus::NeedsAuth,
McpServerStatusReason::AuthExpired,
Some(reason.clone()),
)
}
(McpClientEventKind::HandshakeFailed, McpClientEvent::HandshakeFailed { reason, .. }) => { (McpClientEventKind::HandshakeFailed, McpClientEvent::HandshakeFailed { reason, .. }) => {
let detail = reason.clone(); let detail = reason.clone();
( (
@@ -586,7 +558,7 @@ pub fn recoverable_http_servers(
) )
}) })
.map(|c| mcp_server_name(c).to_string()) .map(|c| mcp_server_name(c).to_string())
.filter(|name| !name.starts_with(MANAGED_MCP_PREFIX) && !disabled.contains(name)) .filter(|name| !disabled.contains(name))
.collect() .collect()
} }
@@ -984,58 +956,9 @@ mod tests {
); );
} }
/// Contract: a managed connector whose handshake is rejected for /// A local server whose handshake error happens to contain auth
/// auth reasons surfaces as `NeedsAuth`/`auth_expired` ("visit /// wording stays `Unavailable` (local auth recovery is the OAuth
/// grok.com"), NOT a generic `Unavailable`. Keys on the shared /// path, not this one).
/// `is_auth_rejection_message` classifier.
#[test]
fn managed_handshake_auth_rejection_maps_to_needs_auth() {
let key = (
"grok_com_notion".to_string(),
McpClientEventKind::HandshakeFailed,
);
let ev = McpClientEvent::HandshakeFailed {
server: "grok_com_notion".to_string(),
reason: "Auth required, when send initialize request".to_string(),
};
let payload = build_payload("sess1", &key, &ev);
assert_eq!(payload.source, McpServerSource::Managed);
assert_eq!(payload.status, McpServerStatus::NeedsAuth);
assert_eq!(payload.reason, McpServerStatusReason::AuthExpired);
let json = serde_json::to_value(&payload).unwrap();
// `McpServerStatus` serializes lowercase (no underscore); the
// reason enum serializes snake_case.
assert_eq!(json["status"], "needsauth");
assert_eq!(json["reason"], "auth_expired");
}
/// A managed handshake failure that is NOT an auth rejection (e.g. a
/// 403 policy denial or a 502) must stay `Unavailable` — the
/// `NeedsAuth` arm is auth-only.
#[test]
fn managed_handshake_non_auth_stays_unavailable() {
for reason in ["403 Forbidden", "cli-chat-proxy returned 502"] {
let key = (
"grok_com_slack".to_string(),
McpClientEventKind::HandshakeFailed,
);
let ev = McpClientEvent::HandshakeFailed {
server: "grok_com_slack".to_string(),
reason: reason.to_string(),
};
let payload = build_payload("sess1", &key, &ev);
assert_eq!(
payload.status,
McpServerStatus::Unavailable,
"non-auth managed failure must stay Unavailable: {reason}",
);
assert_eq!(payload.reason, McpServerStatusReason::HandshakeFailed);
}
}
/// The `NeedsAuth` arm is managed-only: a local (non-managed) server
/// whose handshake error happens to contain auth wording stays
/// `Unavailable` (local auth recovery is the OAuth path, not this one).
#[test] #[test]
fn local_handshake_auth_rejection_stays_unavailable() { fn local_handshake_auth_rejection_stays_unavailable() {
let key = ("github".to_string(), McpClientEventKind::HandshakeFailed); let key = ("github".to_string(), McpClientEventKind::HandshakeFailed);
@@ -1049,32 +972,6 @@ mod tests {
assert_eq!(payload.reason, McpServerStatusReason::HandshakeFailed); assert_eq!(payload.reason, McpServerStatusReason::HandshakeFailed);
} }
/// Wire contract: the new `ManagedTokenRefreshed` reason (emitted by
/// the reactive re-auth success push) serializes to snake_case.
#[test]
fn managed_token_refreshed_reason_serializes() {
let payload = McpServerStatusPayload {
session_id: "sess1".to_string(),
name: "grok_com_linear".to_string(),
source: McpServerSource::Managed,
status: McpServerStatus::Ready,
reason: McpServerStatusReason::ManagedTokenRefreshed,
detail: None,
tools: None,
};
let json = serde_json::to_value(&payload).unwrap();
assert_eq!(json["status"], "ready");
assert_eq!(json["reason"], "managed_token_refreshed");
}
/// Contract: managed server names (starting with `grok_com_`)
/// are classified as `Managed`; everything else as `Local`.
#[test]
fn classify_source_uses_managed_prefix() {
assert_eq!(classify_source("grok_com_linear"), McpServerSource::Managed);
assert_eq!(classify_source("github"), McpServerSource::Local);
}
/// Snapshot of the wire shape for one TransportClosed status push. /// Snapshot of the wire shape for one TransportClosed status push.
/// ///
/// Locks the camelCase field naming, lowercase enum values, and /// Locks the camelCase field naming, lowercase enum values, and
@@ -1393,13 +1290,12 @@ mod tests {
) )
} }
/// `recoverable_http_servers` keeps only non-managed, non-disabled /// `recoverable_http_servers` keeps only non-disabled
/// HTTP/SSE entries — the same predicate as the recovery gate. /// HTTP/SSE entries — the same predicate as the recovery gate.
#[test] #[test]
fn recoverable_http_servers_excludes_managed_stdio_and_disabled() { fn recoverable_http_servers_excludes_stdio_and_disabled() {
let configs = vec![ let configs = vec![
http_cfg("http-mcp-server"), http_cfg("http-mcp-server"),
http_cfg("grok_com_slack"), // managed
http_cfg("admin_off"), // disabled http_cfg("admin_off"), // disabled
stdio_cfg("local_stdio"), // stdio stdio_cfg("local_stdio"), // stdio
]; ];
@@ -50,14 +50,7 @@ pub fn build_config_resolved_event(
.map(|c| kigi_file_utils::events::McpConfigServer { .map(|c| kigi_file_utils::events::McpConfigServer {
name: inner::mcp_server_name(c).to_string(), name: inner::mcp_server_name(c).to_string(),
transport: inner::mcp_transport_str(c).to_string(), transport: inner::mcp_transport_str(c).to_string(),
source: if inner::mcp_server_name(c) source: "local".to_string(),
.starts_with(crate::session::managed_mcp::MANAGED_MCP_PREFIX)
{
"managed"
} else {
"local"
}
.to_string(),
}) })
.collect(); .collect();
kigi_file_utils::events::Event::McpConfigResolved { servers, disabled } kigi_file_utils::events::Event::McpConfigResolved { servers, disabled }
@@ -1697,11 +1697,11 @@ mod tests {
#[test] #[test]
fn build_tools_meta_serialises_tool_names() { fn build_tools_meta_serialises_tool_names() {
let names = vec!["scheduler_create".to_string(), "image_gen".to_string()]; let names = vec!["scheduler_create".to_string(), "web_search".to_string()];
let v = build_tools_meta(&names); let v = build_tools_meta(&names);
assert_eq!( assert_eq!(
serde_json::Value::Object(v), serde_json::Value::Object(v),
serde_json::json!({"tools": ["scheduler_create", "image_gen"]}) serde_json::json!({"tools": ["scheduler_create", "web_search"]})
); );
} }
@@ -83,8 +83,6 @@ pub(crate) fn ctx_with_toggle(toggle: HashMap<String, bool>) -> SubagentSpawnCon
memory_config: None, memory_config: None,
web_search_config: Default::default(), web_search_config: Default::default(),
web_fetch_config: Default::default(), web_fetch_config: Default::default(),
image_gen_config: Default::default(),
video_gen_config: Default::default(),
app_builder_deployer_config: Default::default(), app_builder_deployer_config: Default::default(),
write_file_enabled: true, write_file_enabled: true,
goal_enabled: false, goal_enabled: false,
@@ -128,8 +126,6 @@ pub(crate) fn ctx_with_toggle(toggle: HashMap<String, bool>) -> SubagentSpawnCon
parent_model_agent_type: None, parent_model_agent_type: None,
allowed_subagent_types: None, allowed_subagent_types: None,
parent_mcp_configs: vec![], parent_mcp_configs: vec![],
managed_mcp_state: crate::session::managed_mcp::ManagedMcpStateHandle::default(),
managed_mcp_proxy_base_url: String::new(),
parent_mcp_pool: None, parent_mcp_pool: None,
parent_tool_snapshot: None, parent_tool_snapshot: None,
parent_skills: None, parent_skills: None,
@@ -328,12 +328,10 @@ pub fn all_toml_mcp_server_names(cwd: &std::path::Path) -> std::collections::Has
load_all_mcp_configs(cwd).keys().cloned().collect() load_all_mcp_configs(cwd).keys().cloned().collect()
} }
pub const MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY: &str = "__managed_gateway_connectors";
/// Persist `disabled_tools` for a server under `[disabled_mcp_tools]` in config.toml. /// Persist `disabled_tools` for a server under `[disabled_mcp_tools]` in config.toml.
/// ///
/// Uses a dedicated top-level section (not `[mcp_servers]`) to avoid creating /// Uses a dedicated top-level section (not `[mcp_servers]`) to avoid creating
/// incomplete server entries that fail to deserialize for managed servers. /// incomplete server entries that fail to deserialize.
pub async fn save_mcp_disabled_tools(server_name: &str, disabled_tools: &[String]) -> Result<()> { pub async fn save_mcp_disabled_tools(server_name: &str, disabled_tools: &[String]) -> Result<()> {
let path = config_path(); let path = config_path();
let mut root: TomlValue = match tokio::fs::read_to_string(&path).await { let mut root: TomlValue = match tokio::fs::read_to_string(&path).await {
@@ -1,6 +1,5 @@
//! Canonical slash-command wording (`/loop`, `/imagine`, `/imagine-video`, `/goal`), //! Canonical slash-command wording (`/loop`, `/goal`), shared by every
//! shared by every front-end (Grok Build shell/pager and other hosts) so //! front-end (shell/pager and other hosts) so expansions cannot drift.
//! expansions cannot drift.
/// Canonical tool name advertised by the scheduler create tool. Gating code /// Canonical tool name advertised by the scheduler create tool. Gating code
/// (shell `CommandAvailability`, pager `required_tools`, host command lists) /// (shell `CommandAvailability`, pager `required_tools`, host command lists)
@@ -44,91 +43,6 @@ pub fn loop_schedule_instruction(args: &str) -> String {
) )
} }
/// Canonical name of the image generation tool; gates `/imagine`.
pub const IMAGE_GEN_TOOL_NAME: &str = "image_gen";
/// Advertised name of the /imagine command.
pub const IMAGINE_COMMAND_NAME: &str = "imagine";
/// Canonical name of the image-to-video tool; gates `/imagine-video`.
pub const IMAGE_TO_VIDEO_TOOL_NAME: &str = "image_to_video";
/// Advertised name of the /imagine-video command.
pub const IMAGINE_VIDEO_COMMAND_NAME: &str = "imagine-video";
/// Usage hint shown when `/imagine` is invoked with no arguments.
pub fn imagine_usage_message() -> &'static str {
"Usage: /imagine <description>\n\
Provide a text description to generate an image."
}
/// Build the model instruction that `/imagine` expands into for `prompt`.
pub fn imagine_instruction(prompt: &str) -> String {
format!(
"Call the image_gen tool immediately, passing the user's prompt below \
verbatim do not rewrite, embellish, or expand it. \
After the tool completes, briefly acknowledge and mention \
where the image was saved.\n\n\
Prompt: {prompt}"
)
}
/// Usage hint shown when `/imagine-video` is invoked with no arguments.
pub fn imagine_video_usage_message() -> &'static str {
"Usage: /imagine-video <description>\n\
Provide a text description to generate a video."
}
/// Build the model instruction that `/imagine-video` expands into for `prompt`.
pub fn imagine_video_instruction(prompt: &str) -> String {
format!(
"{IMAGINE_VIDEO_SKILL}\n\n\
User prompt: {prompt}"
)
}
/// Video workflow guidance injected by `/imagine-video`.
const IMAGINE_VIDEO_SKILL: &str = "\
# Imagine Video
Video starts from an image there is no text-to-video tool. \
Default to `image_to_video`; use `reference_to_video` only when the user \
explicitly asks for it or a shot genuinely needs multiple reference images.
## Default: single clip
Unless the user asks for a long video, multiple scenes, or a multi-shot sequence, \
generate **one** video:
1. Create a source image with `image_gen` that stages the first frame \
(composition, subject, lighting).
2. Call `image_to_video` with that image and a short prompt describing the motion \
or camera move (12 sentences, present tense).
3. After the tool completes, mention the saved file path so the user can find it.
## Longer / multi-shot videos
When the user requests a longer video, multiple scenes, or a narrative sequence:
1. **Plan the story as shots** break the idea into distinct shots, one beat each.
2. **Favor frequent, short shots** prefer more 6s clips over fewer long ones; more cuts keep it dynamic.
3. **Create each shot's source image** with `image_gen` (or `image_edit` to combine references), keeping characters and settings consistent across shots.
4. **Animate each shot with `image_to_video`** the source image becomes frame 1.
5. **Assemble with FFmpeg** using stream copy (`ffmpeg -f concat ... -c copy` never re-encode). \
Keep every shot at the same resolution and frame rate so the concat works. \
After assembly, mention the final output path.
## Shot guidance
- **Prompt-craft:** one short, vivid moment in present tense with a clear camera movement, in 12 sentences.
- **Minimal but interesting:** one clear subject, one simple motion or camera move per shot. Avoid complex multi-action animation; make the shot compelling through composition, lighting, and a strong moment.
- **Complex source image?** Intricate frames (busy geometry, fine detail, heavy reflections) warp when animated. Keep the subject fixed and move only the camera (slow push-in, orbit, or parallax), or break into simpler shots. For new shots, generate a simpler, animation-friendly base image rather than animating a busy one.
- **`image_to_video` animates from frame 1** stage the first frame with `image_gen`/`image_edit` before animating.
- **Aspect ratio:** set it on the source image (`image_gen` `aspect_ratio`); don't re-crop an existing video.
- **Duration:** 6s or 10s only (prefer 6s); round to the nearest.
- **Real people:** reference-first drive the video from a verified reference image; never animate a named person without one.
- Don't loop the same clip unless asked.";
pub const UPDATE_GOAL_TOOL_NAME: &str = "update_goal"; pub const UPDATE_GOAL_TOOL_NAME: &str = "update_goal";
pub const GOAL_COMMAND_NAME: &str = "goal"; pub const GOAL_COMMAND_NAME: &str = "goal";
@@ -170,22 +84,6 @@ pub fn goal_instruction(objective: &str) -> String {
mod tests { mod tests {
use super::*; use super::*;
#[test]
fn imagine_instruction_carries_prompt_verbatim() {
let text = imagine_instruction("a golden sunset");
assert!(text.contains("a golden sunset"));
assert!(text.contains("image_gen"));
assert!(text.contains("verbatim"));
}
#[test]
fn imagine_video_instruction_carries_prompt_and_workflow() {
let text = imagine_video_instruction("a cat playing piano");
assert!(text.contains("a cat playing piano"));
assert!(text.contains("image_to_video"));
assert!(text.contains("FFmpeg"));
}
#[test] #[test]
fn instruction_carries_args_and_contract_tokens() { fn instruction_carries_args_and_contract_tokens() {
let text = loop_schedule_instruction("every 30 minutes do x"); let text = loop_schedule_instruction("every 30 minutes do x");
@@ -36,7 +36,7 @@
], ],
"definitions": { "definitions": {
"ToolKind": { "ToolKind": {
"description": "Categorizes what a tool does at a high level. Open set — consumers must tolerate unknown values (Rust deserializes them to `other` via `#[serde(other)]`). Known values: `read`, `edit`, `delete`, `list_dir`, `write`, `move`, `search`, `lsp`, `execute`, `plan`, `web_search`, `web_fetch`, `background_task_action`, `wait_tasks_action`, `kill_task_action`, `list`, `skill`, `memory_search`, `memory_get`, `task`, `enter_plan`, `exit_plan`, `ask_user`, `image_gen`, `video_gen`, `image_to_video`, `reference_to_video`, `deploy_app`, `search_tool`, `use_tool`, `monitor`, `goal_update`, `other`.", "description": "Categorizes what a tool does at a high level. Open set — consumers must tolerate unknown values (Rust deserializes them to `other` via `#[serde(other)]`). Known values: `read`, `edit`, `delete`, `list_dir`, `write`, `move`, `search`, `lsp`, `execute`, `plan`, `web_search`, `web_fetch`, `background_task_action`, `wait_tasks_action`, `kill_task_action`, `list`, `skill`, `memory_search`, `memory_get`, `task`, `enter_plan`, `exit_plan`, `ask_user`, `deploy_app`, `search_tool`, `use_tool`, `monitor`, `goal_update`, `other`.",
"type": "string" "type": "string"
}, },
"ToolNamespace": { "ToolNamespace": {
@@ -8,18 +8,12 @@ pub const SENT_BEARER_PREFIX_LEN: usize = 12;
/// Which tool endpoint produced the 401. /// Which tool endpoint produced the 401.
#[derive(Debug, Clone, Copy, PartialEq, Eq)] #[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolConsumer { pub enum ToolConsumer {
ImageGen,
VideoGenStart,
VideoGenPoll,
WebSearch, WebSearch,
} }
impl ToolConsumer { impl ToolConsumer {
pub fn as_str(self) -> &'static str { pub fn as_str(self) -> &'static str {
match self { match self {
Self::ImageGen => "ImageGen",
Self::VideoGenStart => "VideoGen.start",
Self::VideoGenPoll => "VideoGen.poll",
Self::WebSearch => "WebSearch", Self::WebSearch => "WebSearch",
} }
} }
@@ -95,9 +89,6 @@ mod tests {
#[test] #[test]
fn tool_consumer_as_str_stable_identifiers() { fn tool_consumer_as_str_stable_identifiers() {
assert_eq!(ToolConsumer::ImageGen.as_str(), "ImageGen");
assert_eq!(ToolConsumer::VideoGenStart.as_str(), "VideoGen.start");
assert_eq!(ToolConsumer::VideoGenPoll.as_str(), "VideoGen.poll");
assert_eq!(ToolConsumer::WebSearch.as_str(), "WebSearch"); assert_eq!(ToolConsumer::WebSearch.as_str(), "WebSearch");
} }
} }
@@ -1,704 +0,0 @@
//! `image_edit` tool — edits or transforms images via the xAI Imagine
//! `/images/edits` endpoint using one or more reference images.
//!
//! Use cases include likeness preservation, style transfer, subject lock,
//! remixing, and general image-to-image editing. The model chooses this
//! tool (instead of `image_gen`) when the user provides reference photos.
//!
//! Reference images are specified as filesystem paths or
//! `data:image/...;base64,...` URLs. The tool reads the bytes, compresses
//! them to fit API limits, and POSTs to the edit endpoint.
//!
//! Shares the same [`ImageGenClient`] and session credentials as
//! `image_gen` — no additional configuration is needed.
use std::io::Cursor;
use base64::Engine as _;
use image::ImageReader;
use reqwest::header::AUTHORIZATION;
use crate::attribution::ToolConsumer;
use crate::implementations::grok_build::image_gen::{ImageGenClient, ImageGenResponse};
use crate::types::output::{MediaGenOutput, ToolOutput};
use crate::types::requirements::{Expr, ToolRequirement};
use crate::types::resources::SessionFolder;
use crate::types::tool::{ToolKind, ToolNamespace};
use crate::util::image_compress::{FilterType, ReEncodeParams, re_encode_under_limit};
const XAI_IMAGINE_MODEL: &str = "grok-imagine-image-quality";
/// Size/dimension limits for reference images sent to the Imagine API.
/// Tighter than the vision path; the backend returns 400 when exceeded.
const MAX_REF_RAW_BYTES: usize = 400 * 1024;
const MAX_REF_DIMENSION: u32 = 768;
const MIN_REF_DIMENSION: u32 = 256;
const REF_QUALITY_STEPS: &[u8] = &[80, 65, 50, 35];
const MAX_REF_DECODE_PIXELS: u64 = 12_000_000;
pub const IMAGE_EDIT_TOOL_NAME: &str = "image_edit";
// ---------------------------------------------------------------------------
// Compression
// ---------------------------------------------------------------------------
/// Compress a reference image to fit within Imagine API limits.
///
/// Returns `(bytes, mime)`. Small JPEG/PNG inputs pass through unchanged.
fn compress_reference(
raw_bytes: Vec<u8>,
) -> Result<(Vec<u8>, &'static str), kigi_tool_runtime::ToolError> {
// Fast path: small JPEG/PNG passes through unchanged. Other formats
// (WebP, GIF, etc.) always re-encode to guarantee API-compatible output.
if raw_bytes.len() <= MAX_REF_RAW_BYTES
&& let Some(kind) = infer::get(&raw_bytes)
{
match kind.mime_type() {
"image/jpeg" => return Ok((raw_bytes, "image/jpeg")),
"image/png" => return Ok((raw_bytes, "image/png")),
_ => {}
}
}
// Refuse to decode absurdly large images.
let reader = ImageReader::new(Cursor::new(&raw_bytes))
.with_guessed_format()
.map_err(|_| {
kigi_tool_runtime::ToolError::invalid_arguments(
"could not detect image format for reference",
)
})?;
if let Ok((w, h)) = reader.into_dimensions()
&& (w as u64) * (h as u64) > MAX_REF_DECODE_PIXELS
{
return Err(kigi_tool_runtime::ToolError::invalid_arguments(format!(
"image reference is too large to process ({w}\u{00d7}{h} pixels)",
)));
}
// `into_dimensions` consumed the reader; re-open to decode.
let img = ImageReader::new(Cursor::new(&raw_bytes))
.with_guessed_format()
.ok()
.and_then(|r| r.decode().ok())
.ok_or_else(|| {
kigi_tool_runtime::ToolError::invalid_arguments("failed to decode image reference")
})?;
let params = ReEncodeParams {
max_bytes: MAX_REF_RAW_BYTES,
max_side_px: MAX_REF_DIMENSION,
// Imagine backend limits are side-based; no pixel-area cap applies.
max_pixels: u64::MAX,
min_side_px: MIN_REF_DIMENSION,
quality_steps: REF_QUALITY_STEPS,
filter: FilterType::Lanczos3,
};
let (buf, _w, _h, mime) = re_encode_under_limit(&img, &params).map_err(|e| {
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"could not compress image reference small enough for Imagine API: {e}"
))
})?;
Ok((buf, mime))
}
// ---------------------------------------------------------------------------
// Reference resolution
// ---------------------------------------------------------------------------
/// Resolve a reference (filesystem path or `data:image/...;base64,...` URL)
/// into a compressed data URL for the Imagine API.
async fn resolve_to_data_url(value: &str) -> Result<String, kigi_tool_runtime::ToolError> {
let value = value.trim();
// Accept `file://` URIs (e.g. an attachment's durable URI) by reading
// the underlying path. Data URLs and bare paths are untouched.
let value = value.strip_prefix("file://").unwrap_or(value);
let raw_bytes = if value.starts_with("data:image/") {
let comma = value.find(',').ok_or_else(|| {
kigi_tool_runtime::ToolError::invalid_arguments("malformed data URL in image reference")
})?;
if !value[..comma].contains(";base64") {
return Err(kigi_tool_runtime::ToolError::invalid_arguments(
"image references only support base64 data URLs",
));
}
base64::engine::general_purpose::STANDARD
.decode(&value[comma + 1..])
.map_err(|e| {
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"invalid base64 in image reference: {e}"
))
})?
} else {
tokio::fs::read(value).await.map_err(|e| {
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"image reference not readable: {value} ({e})"
))
})?
};
if raw_bytes.is_empty() {
return Err(kigi_tool_runtime::ToolError::invalid_arguments(
"image reference contained no data",
));
}
let (compressed, mime) = compress_reference(raw_bytes)?;
let b64 = base64::engine::general_purpose::STANDARD.encode(&compressed);
Ok(format!("data:{mime};base64,{b64}"))
}
// ---------------------------------------------------------------------------
// Attachment reference resolution
// ---------------------------------------------------------------------------
/// Parse an attached-image reference token into its 1-based display number.
///
/// Accepts the forms the model naturally produces for an image the user
/// attached to the conversation: `[Image #1]`, `Image #1`, `image #1`, or
/// a bare `#1`. Returns `None` for anything else — filesystem paths and
/// `data:` / `file://` URLs fall through to direct resolution.
fn parse_attachment_token(value: &str) -> Option<usize> {
let trimmed = value.trim();
// Strip optional surrounding brackets: `[…]`.
let inner = trimmed
.strip_prefix('[')
.and_then(|s| s.strip_suffix(']'))
.unwrap_or(trimmed)
.trim();
// Strip an optional leading `image` label (case-insensitive). The
// 5-byte prefix is ASCII, so slicing at byte 5 stays on a boundary.
let rest = match inner.get(..5).map(str::to_ascii_lowercase).as_deref() {
Some("image") => inner[5..].trim_start(),
_ => inner,
};
// Require the `#` sigil followed by a bare positive integer.
let digits = rest.strip_prefix('#')?.trim();
match digits.parse::<usize>() {
Ok(n) if n >= 1 => Some(n),
_ => None,
}
}
/// Resolve a single `image` argument to a reference `resolve_to_data_url`
/// can read.
///
/// Attachment tokens (`[Image #N]`) are mapped to the durable reference
/// the shell recorded for the current turn; everything else (filesystem
/// paths, `data:` / `file://` URLs) passes through unchanged.
fn resolve_attachment_reference(
reference: &str,
attached: Option<&crate::types::resources::AttachedImages>,
) -> Result<String, kigi_tool_runtime::ToolError> {
let Some(n) = parse_attachment_token(reference) else {
return Ok(reference.to_owned());
};
let registry = attached.filter(|a| !a.0.is_empty()).ok_or_else(|| {
// Tokens only resolve against the current message's attachments. An
// empty registry usually means the image was attached in an earlier
// message (cross-turn editing isn't supported yet), so steer the
// model to ask for a re-attach rather than retry the dead token.
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"image reference {reference:?} matches no image attached to this message. If it was \
attached earlier in the conversation, ask the user to re-attach it here; otherwise \
pass an absolute filesystem path or a data: URL."
))
})?;
registry.reference_for(n).map(str::to_owned).ok_or_else(|| {
let available: Vec<String> = registry
.0
.iter()
.map(|(num, _)| format!("[Image #{num}]"))
.collect();
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"image reference {reference:?} does not match any attached image. Available: {}.",
available.join(", ")
))
})
}
// ---------------------------------------------------------------------------
// Tool input / schema
// ---------------------------------------------------------------------------
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
pub struct ImageEditInput {
#[schemars(
description = "A text description of the desired edit or transformation. Describe what the output image should look like, referencing the input image(s)."
)]
pub prompt: String,
#[schemars(
description = "Reference image(s) to condition the edit on. Each is one reference, in priority order: (1) a user attachment — its placeholder token, e.g. \"[Image #1]\" (attachments have no path you can see, so never invent one); (2) an absolute filesystem path the user gave you; (3) a `data:image/...;base64,...` URL."
)]
pub image: Vec<String>,
#[serde(default = "default_aspect_ratio")]
#[schemars(
description = "The aspect ratio of the output image. For single-image edits this is ignored — the output matches the input image's aspect ratio. For multi-image edits, defaults to 'auto'. Supported values: 1:1, 16:9, 9:16, 4:3, 3:4, 3:2, 2:3, 2:1, 1:2, 19.5:9, 9:19.5, 20:9, 9:20, auto."
)]
pub aspect_ratio: String,
}
fn default_aspect_ratio() -> String {
"auto".to_owned()
}
// ---------------------------------------------------------------------------
// Tool implementation
// ---------------------------------------------------------------------------
#[derive(Debug, Default)]
pub struct ImageEditTool;
impl crate::types::tool_metadata::ToolMetadata for ImageEditTool {
fn kind(&self) -> ToolKind {
ToolKind::ImageGen
}
fn tool_namespace(&self) -> ToolNamespace {
ToolNamespace::GrokBuild
}
fn description_template(&self) -> &str {
r##"Edit or transform existing image(s) via the xAI Imagine API; use instead of image_gen for image-to-image work (preserve likeness, transfer style, remix). Returns the saved image's absolute path. When telling the user where it was saved, refer to it by its short session-relative path (e.g. `images/1.jpg`) rather than the absolute path, so it renders as a clickable link that opens the image. Each required `image` is one reference — a user-attachment token (e.g. "[Image #1]"), an absolute filesystem path, or a `data:image/...;base64,...` URL (see the `image` parameter for the resolution order and details)."##
}
fn requires_expr(&self) -> Expr<ToolRequirement> {
Expr::True
}
}
impl kigi_tool_runtime::Tool for ImageEditTool {
type Args = ImageEditInput;
type Output = ToolOutput;
fn id(&self) -> kigi_tool_protocol::ToolId {
kigi_tool_protocol::ToolId::new("image_edit").expect("valid tool id")
}
fn description(
&self,
_ctx: &::kigi_tool_runtime::ListToolsContext,
) -> kigi_tool_types::ToolDescription {
kigi_tool_types::ToolDescription::new(
"image_edit",
crate::types::tool_metadata::ToolMetadata::description_template(self),
)
}
fn capabilities(&self) -> kigi_tool_protocol::ToolCapabilities {
kigi_tool_protocol::ToolCapabilities {
is_read_only: false,
tool_scope: Some(kigi_tool_protocol::ToolScope::Write),
..Default::default()
}
}
#[tracing::instrument(
name = "tool.image_edit",
skip_all,
fields(prompt_len = input.prompt.len(), num_images = input.image.len(), aspect_ratio = %input.aspect_ratio)
)]
async fn run(
&self,
ctx: kigi_tool_runtime::ToolCallContext,
input: ImageEditInput,
) -> Result<ToolOutput, kigi_tool_runtime::ToolError> {
use crate::types::tool_metadata::shared_resources;
let resources = shared_resources(&ctx)?;
if input.image.is_empty() {
return Err(kigi_tool_runtime::ToolError::invalid_arguments(
"image_edit requires at least one reference image. \
Use image_gen for text-only generation.",
));
}
let client = {
let res = resources.lock().await;
res.require::<ImageGenClient>()?.clone()
};
// Free / X Basic users are zero-limited on Imagine server-side; return
// the upsell prose instead of a doomed request (shares `image_gen`'s
// message and short-circuits before resolving any attachments).
if client.is_tier_restricted() {
return Ok(ToolOutput::Text(
super::image_gen::TIER_RESTRICTED_UPSELL.into(),
));
}
// Snapshot the per-turn attachment registry so `[Image #N]` tokens
// resolve to the real attachment (see `resolve_attachment_reference`).
let attached_images = {
let res = resources.lock().await;
res.get::<crate::types::resources::AttachedImages>()
.cloned()
};
// Resolve all references to compressed data URLs.
let mut data_urls = Vec::with_capacity(input.image.len());
for r in &input.image {
let resolved = resolve_attachment_reference(r, attached_images.as_ref())?;
data_urls.push(resolve_to_data_url(&resolved).await?);
}
tracing::info!(count = data_urls.len(), "resolved image references");
let base = client.base_url().trim_end_matches('/');
let url = format!("{base}/images/edits");
let mut payload = serde_json::json!({
"model": XAI_IMAGINE_MODEL,
"prompt": input.prompt,
"n": 1,
"resolution": "1k",
"response_format": "b64_json",
});
// API: single ref → "image" object; multiple → "images" array.
// For single-image edits the API auto-detects aspect ratio from the
// input image and ignores the `aspect_ratio` field. Only send it
// for multi-image edits where the API needs an explicit ratio.
let mut imgs: Vec<serde_json::Value> = data_urls
.iter()
.map(|u| serde_json::json!({ "url": u }))
.collect();
if imgs.len() == 1 {
payload["image"] = imgs.pop().unwrap();
} else {
payload["images"] = serde_json::Value::Array(imgs);
payload["aspect_ratio"] = serde_json::json!(input.aspect_ratio);
}
let sent_bearer = client.current_bearer().await;
let mut req = client.http().post(&url).json(&payload);
if let Some(ref key) = sent_bearer {
req = req.header(AUTHORIZATION, format!("Bearer {key}"));
}
let response = req.send().await.map_err(|e| {
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"Image edit API request failed: {e}"
))
})?;
let status = response.status();
if status == reqwest::StatusCode::UNAUTHORIZED {
client.record_401_attribution(ToolConsumer::ImageGen, sent_bearer.as_deref());
}
if !status.is_success() {
let body = response.text().await.unwrap_or_default();
let truncated: String = body.chars().take(200).collect();
tracing::warn!(http_status = %status, "Imagine edit API error: {truncated}");
return Err(kigi_tool_runtime::ToolError::new(
kigi_tool_runtime::ToolErrorKind::Custom,
format!("Image edit failed with HTTP {status}: {truncated}"),
)
.with_details(serde_json::json!({"code": "http_failure", "status": status.as_u16()})));
}
let body = response.text().await.map_err(|e| {
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"Failed to read image edit response body: {e}"
))
})?;
let resp_json: ImageGenResponse = serde_json::from_str(&body).map_err(|e| {
let preview: String = body.chars().take(500).collect();
tracing::warn!("Imagine edit API returned unparseable body: {preview}");
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"Failed to parse image edit response: {e} — body preview: {preview}"
))
})?;
let b64_data = resp_json.b64_data().unwrap_or("");
if b64_data.is_empty() {
return Err(kigi_tool_runtime::ToolError::invalid_arguments(
"Image edit returned no image data.",
));
}
let image_bytes = base64::engine::general_purpose::STANDARD
.decode(b64_data)
.map_err(|e| {
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"Failed to decode base64 image data: {e}"
))
})?;
let session_folder = {
let res = resources.lock().await;
res.require::<SessionFolder>()?.0.clone()
};
let absolute_path = client
.writer()
.save(&session_folder, &image_bytes, None)
.await
.map_err(|e| kigi_tool_runtime::ToolError::invalid_arguments(e.to_string()))?;
tracing::info!(
path = %absolute_path.display(),
bytes = image_bytes.len(),
"edited image saved to disk"
);
Ok(ToolOutput::ImageEdit(MediaGenOutput::new(absolute_path)))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::tool_metadata::test_ctx_with_call_id;
#[test]
fn tool_name_and_description() {
let tool = ImageEditTool;
assert_eq!(kigi_tool_runtime::Tool::id(&tool).as_str(), "image_edit");
let desc = crate::types::tool_metadata::ToolMetadata::description_template(&tool);
assert!(desc.contains("Edit or transform"));
}
#[test]
fn input_deserialization() {
let input: ImageEditInput =
serde_json::from_str(r#"{"prompt": "anime style", "image": ["/Users/me/photo.jpg"]}"#)
.unwrap();
assert_eq!(input.prompt, "anime style");
assert_eq!(input.image, vec!["/Users/me/photo.jpg"]);
assert_eq!(input.aspect_ratio, "auto");
}
#[test]
fn input_requires_image() {
// image field is required by schema — empty array is a runtime check.
let input: ImageEditInput =
serde_json::from_str(r#"{"prompt": "test", "image": []}"#).unwrap();
assert!(input.image.is_empty());
}
#[tokio::test]
async fn rejects_empty_image_array() {
let tool = ImageEditTool;
let resources = crate::types::resources::Resources::new();
let result = kigi_tool_runtime::Tool::run(
&tool,
test_ctx_with_call_id(resources.into_shared(), "test-call"),
ImageEditInput {
prompt: "test".into(),
image: vec![],
aspect_ratio: "auto".into(),
},
)
.await;
let err = result.unwrap_err().to_string();
assert!(err.contains("at least one reference image"), "got: {err}");
}
#[tokio::test]
async fn errors_when_client_missing() {
let tool = ImageEditTool;
let resources = crate::types::resources::Resources::new();
let result = kigi_tool_runtime::Tool::run(
&tool,
test_ctx_with_call_id(resources.into_shared(), "test-call"),
ImageEditInput {
prompt: "test".into(),
image: vec!["/some/path.jpg".into()],
aspect_ratio: "auto".into(),
},
)
.await;
let err = result.unwrap_err().to_string();
assert!(err.contains("missing required resource"), "got: {err}");
}
// ── compress_reference ───────────────────────────────────────────
fn tiny_jpeg() -> Vec<u8> {
use image::{DynamicImage, RgbImage};
let img = DynamicImage::ImageRgb8(RgbImage::new(2, 2));
let mut buf = Vec::new();
img.write_to(
&mut std::io::Cursor::new(&mut buf),
image::ImageFormat::Jpeg,
)
.unwrap();
buf
}
fn tiny_png() -> Vec<u8> {
use image::{DynamicImage, RgbaImage};
let img = DynamicImage::ImageRgba8(RgbaImage::new(2, 2));
let mut buf = Vec::new();
img.write_to(&mut std::io::Cursor::new(&mut buf), image::ImageFormat::Png)
.unwrap();
buf
}
#[test]
fn compress_small_jpeg_passthrough() {
let jpeg = tiny_jpeg();
let (out, mime) = compress_reference(jpeg.clone()).unwrap();
assert_eq!(out, jpeg);
assert_eq!(mime, "image/jpeg");
}
#[test]
fn compress_small_png_passthrough() {
let png = tiny_png();
let (out, mime) = compress_reference(png.clone()).unwrap();
assert_eq!(out, png);
assert_eq!(mime, "image/png");
}
#[test]
fn compress_oversized_shrinks() {
use image::{DynamicImage, RgbImage};
let mut img = RgbImage::new(1600, 1600);
for (i, px) in img.pixels_mut().enumerate() {
let v = (i * 37 + 13) as u8;
*px = image::Rgb([v, v.wrapping_add(80), v.wrapping_add(160)]);
}
let mut buf = Vec::new();
let enc = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut buf, 100);
DynamicImage::ImageRgb8(img)
.write_with_encoder(enc)
.unwrap();
assert!(buf.len() > MAX_REF_RAW_BYTES);
let (out, mime) = compress_reference(buf).unwrap();
assert!(out.len() <= MAX_REF_RAW_BYTES);
assert!(mime == "image/jpeg" || mime == "image/png");
}
// ── resolve_to_data_url ──────────────────────────────────────────
#[tokio::test]
async fn resolve_filesystem_path() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test.jpg");
std::fs::write(&path, tiny_jpeg()).unwrap();
let url = resolve_to_data_url(path.to_str().unwrap()).await.unwrap();
assert!(url.starts_with("data:image/jpeg;base64,"));
}
#[tokio::test]
async fn resolve_data_url_roundtrip() {
let jpeg = tiny_jpeg();
let b64 = base64::engine::general_purpose::STANDARD.encode(&jpeg);
let input = format!("data:image/jpeg;base64,{b64}");
let url = resolve_to_data_url(&input).await.unwrap();
assert!(url.starts_with("data:image/jpeg;base64,"));
}
#[tokio::test]
async fn resolve_missing_file_errors() {
assert!(resolve_to_data_url("/nonexistent/image.jpg").await.is_err());
}
#[tokio::test]
async fn resolve_malformed_data_url_errors() {
assert!(resolve_to_data_url("data:image/jpeg").await.is_err());
}
#[tokio::test]
async fn resolve_file_uri_reads_path() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test.jpg");
std::fs::write(&path, tiny_jpeg()).unwrap();
let uri = format!("file://{}", path.display());
let url = resolve_to_data_url(&uri).await.unwrap();
assert!(url.starts_with("data:image/jpeg;base64,"));
}
// ── parse_attachment_token ───────────────────────────────────────
#[test]
fn parse_attachment_token_accepts_known_forms() {
assert_eq!(parse_attachment_token("[Image #1]"), Some(1));
assert_eq!(parse_attachment_token("Image #2"), Some(2));
assert_eq!(parse_attachment_token("image #3"), Some(3));
assert_eq!(parse_attachment_token("[image #4]"), Some(4));
assert_eq!(parse_attachment_token("Image#5"), Some(5));
assert_eq!(parse_attachment_token("#6"), Some(6));
assert_eq!(parse_attachment_token(" [Image #7] "), Some(7));
}
#[test]
fn parse_attachment_token_rejects_non_tokens() {
assert_eq!(parse_attachment_token("/Users/me/photo.jpg"), None);
assert_eq!(parse_attachment_token("data:image/png;base64,AAAA"), None);
assert_eq!(parse_attachment_token("file:///tmp/x.png"), None);
assert_eq!(parse_attachment_token("[Image #0]"), None);
assert_eq!(parse_attachment_token("[Image #]"), None);
assert_eq!(parse_attachment_token("Image one"), None);
assert_eq!(parse_attachment_token(""), None);
}
// ── resolve_attachment_reference ─────────────────────────────────
#[test]
fn resolve_reference_passes_through_non_tokens() {
let resolved = resolve_attachment_reference("/Users/me/photo.jpg", None).unwrap();
assert_eq!(resolved, "/Users/me/photo.jpg");
}
#[test]
fn resolve_reference_maps_token_to_registry() {
let attached = crate::types::resources::AttachedImages(vec![
(1, "/tmp/a.png".to_owned()),
(2, "/tmp/b.png".to_owned()),
]);
assert_eq!(
resolve_attachment_reference("[Image #1]", Some(&attached)).unwrap(),
"/tmp/a.png"
);
assert_eq!(
resolve_attachment_reference("Image #2", Some(&attached)).unwrap(),
"/tmp/b.png"
);
}
#[test]
fn resolve_reference_maps_by_number_not_position() {
// After a mid-compose chip removal the surviving numbers are
// non-contiguous (`#1`, `#3`). Resolution must key on the number,
// not the list position, or `[Image #3]` would resolve to the wrong
// file (or wrongly error).
let attached = crate::types::resources::AttachedImages(vec![
(1, "/tmp/first.png".to_owned()),
(3, "/tmp/third.png".to_owned()),
]);
assert_eq!(
resolve_attachment_reference("[Image #3]", Some(&attached)).unwrap(),
"/tmp/third.png"
);
// `[Image #2]` was removed → no match.
assert!(resolve_attachment_reference("[Image #2]", Some(&attached)).is_err());
}
#[test]
fn resolve_reference_token_without_registry_errors() {
let err = resolve_attachment_reference("[Image #1]", None)
.unwrap_err()
.to_string();
assert!(err.contains("re-attach"), "got: {err}");
}
#[test]
fn resolve_reference_unmatched_number_errors() {
let attached = crate::types::resources::AttachedImages(vec![(1, "/tmp/a.png".to_owned())]);
let err = resolve_attachment_reference("[Image #2]", Some(&attached))
.unwrap_err()
.to_string();
assert!(err.contains("does not match"), "got: {err}");
assert!(err.contains("[Image #1]"), "should list available: {err}");
}
}
@@ -1,585 +0,0 @@
//! `image_gen` tool — generates images via the xAI Imagine API and saves
//! them to the local filesystem so the model can reference them in code
//! (e.g. `<img src="images/hero.jpg">`).
//!
//! Architecture follows the same pattern as `web_search`:
//!
//! - [`ImageGenConfig`] is built from session credentials by the host and
//! injected into the tool registry.
//! - When `Enabled`, an [`ImageGenClient`] is constructed once and injected
//! into `Resources`. The tool reads it at runtime via `resources.require()`.
//! - When `Disabled`, the tool is not registered so the model never sees it.
//!
//! The generated image is written to `<session_folder>/images/<n>.jpg`
//! where `<n>` is a session-scoped counter (1, 2, 3, ... — 1 token each).
//! The tool returns the absolute path so the model can copy or move the
//! image into the project working directory when it needs a persistent asset.
use base64::Engine as _;
use reqwest::header::{AUTHORIZATION, CONTENT_TYPE, HeaderValue};
use crate::attribution::{SharedAttributionCallback, ToolConsumer};
use crate::types::SharedApiKeyProvider;
use crate::types::output::{MediaGenOutput, ToolOutput};
use crate::types::requirements::{Expr, ToolRequirement};
use crate::types::resources::SessionFolder;
use crate::types::tool::{ToolKind, ToolNamespace};
/// Default Imagine model for `image_gen`. Used unless an explicit
/// `model_override` is supplied via `ImageGenConfig::Enabled`.
const XAI_IMAGINE_MODEL: &str = "grok-imagine-image-quality";
// Some Imagine models (e.g. `grok-imagine-image`, selectable via `model_override`)
// expand the prompt then generate, and the proxy buffers
// the whole image before sending any bytes — so the client may receive nothing
// for well over a minute. Keep these generous so a slow-but-progressing
// generation isn't cut off.
const IMAGE_GEN_TIMEOUT_SECS: u64 = 300;
const IMAGE_GEN_READ_TIMEOUT_SECS: u64 = 240;
const DEFAULT_IMAGE_DIR: &str = "images";
pub use kigi_tools_api::slash_commands::{
IMAGE_GEN_TOOL_NAME, IMAGINE_COMMAND_NAME, imagine_instruction, imagine_usage_message,
};
/// Prose returned to the model (as a normal, successful tool result) when a
/// free / X Basic user calls `image_gen` or `image_edit`. The model relays it
/// to the user. The deliberate `/imagine` slash command shows the richer
/// SuperGrok upsell modal instead; this covers the natural-language path.
pub(crate) const TIER_RESTRICTED_UPSELL: &str = "Image generation is a SuperGrok feature and isn't available on the free or X Basic tier. Let the user know they can unlock image and video generation by upgrading to SuperGrok: https://grok.com/supergrok?referrer=grok-build. Do not retry this tool.";
/// HTTP client for xAI Imagine API. Cloned per-request; shares `Arc` state.
#[derive(Clone)]
pub struct ImageGenClient {
http: reqwest::Client,
base_url: String,
/// Imagine model slug used by `generate()`. Selected at construction
/// from `ImageGenConfig::model_override` (falling back to
/// [`XAI_IMAGINE_MODEL`]). `image_edit` uses its own model and is
/// unaffected.
model: String,
writer: super::storage::SessionFileWriter,
api_key_provider: Option<SharedApiKeyProvider>,
/// Optional 401-attribution hook. Hosts wire this so a 401 from the
/// Imagine API emits an `auth_401_attribution` event with
/// `consumer == "ImageGen"` for unified auth-failure telemetry.
attribution_callback: Option<SharedAttributionCallback>,
/// When `true`, the user is on a tier the Imagine server zero-limits
/// (free / X Basic). `image_gen` / `image_edit` short-circuit before any
/// HTTP call and return the SuperGrok upsell prose instead. See
/// [`ImageGenClient::is_tier_restricted`].
tier_restricted: bool,
}
impl ImageGenClient {
pub fn new(
config: &ImageGenConfig,
api_key_provider: Option<SharedApiKeyProvider>,
) -> Result<Self, kigi_tool_runtime::ToolError> {
let ImageGenConfig::Enabled {
api_key,
base_url,
extra_headers,
model_override,
tier_restricted,
..
} = config
else {
return Err(kigi_tool_runtime::ToolError::invalid_arguments(
"Cannot create ImageGenClient from disabled config",
));
};
let model = model_override
.clone()
.filter(|m| !m.trim().is_empty())
.unwrap_or_else(|| XAI_IMAGINE_MODEL.to_owned());
let mut headers = reqwest::header::HeaderMap::new();
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
// Always bake the static api_key as the default Authorization header.
// The dynamic provider overrides per-request; this is the fallback.
headers.insert(
AUTHORIZATION,
HeaderValue::from_str(&format!("Bearer {api_key}")).map_err(|e| {
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"Invalid API key for header: {e}"
))
})?,
);
extra_headers.into_iter().try_for_each(|(key, value)| {
let header_name =
reqwest::header::HeaderName::from_bytes(key.as_bytes()).map_err(|e| {
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"Invalid header name '{key}': {e}"
))
})?;
let header_value = HeaderValue::from_str(value).map_err(|e| {
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"Invalid header value for '{key}': {e}"
))
})?;
headers.insert(header_name, header_value);
Ok::<(), kigi_tool_runtime::ToolError>(())
})?;
let http = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(IMAGE_GEN_TIMEOUT_SECS))
.read_timeout(std::time::Duration::from_secs(IMAGE_GEN_READ_TIMEOUT_SECS))
.default_headers(headers)
.build()
.map_err(|e| {
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"Failed to build HTTP client: {e}"
))
})?;
Ok(Self {
http,
base_url: base_url.clone(),
model,
writer: super::storage::SessionFileWriter::new(DEFAULT_IMAGE_DIR, "jpg"),
api_key_provider,
attribution_callback: None,
tier_restricted: *tier_restricted,
})
}
/// Whether the current user's tier (free / X Basic) is zero-limited on
/// Imagine server-side. `image_gen` / `image_edit` use this to short-circuit
/// with the SuperGrok upsell instead of issuing a doomed request.
pub(crate) fn is_tier_restricted(&self) -> bool {
self.tier_restricted
}
/// Wire a 401-attribution callback into this client. Idempotent;
/// safe to call before or after the first request. Builder-style
/// so `new()` callers that don't care can ignore it.
pub fn with_attribution_callback(
mut self,
callback: Option<SharedAttributionCallback>,
) -> Self {
self.attribution_callback = callback;
self
}
pub(crate) async fn current_bearer(&self) -> Option<String> {
crate::types::api_key_provider::resolve_bearer(self.api_key_provider.as_ref()).await
}
pub(crate) fn record_401_attribution(&self, consumer: ToolConsumer, sent_bearer: Option<&str>) {
crate::attribution::emit_401(self.attribution_callback.as_ref(), consumer, sent_bearer);
}
pub(crate) fn base_url(&self) -> &str {
&self.base_url
}
pub(crate) fn http(&self) -> &reqwest::Client {
&self.http
}
pub(crate) fn writer(&self) -> &super::storage::SessionFileWriter {
&self.writer
}
pub async fn generate(
&self,
prompt: &str,
aspect_ratio: &str,
) -> Result<Vec<u8>, kigi_tool_runtime::ToolError> {
let url = format!("{}/images/generations", self.base_url.trim_end_matches('/'));
let payload = serde_json::json!({
"model": self.model,
"prompt": prompt,
"n": 1,
"aspect_ratio": aspect_ratio,
"resolution": "1k",
"response_format": "b64_json",
});
// Capture the bearer once so the request and the 401-attribution
// emit see the same value (even if the provider rotates between
// the send and the response handling).
let sent_bearer = self.current_bearer().await;
let mut req = self.http.post(&url).json(&payload);
if let Some(ref key) = sent_bearer {
req = req.header(AUTHORIZATION, format!("Bearer {key}"));
}
let response = req.send().await.map_err(|e| {
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"Image generation API request failed: {e}"
))
})?;
let status = response.status();
if status == reqwest::StatusCode::UNAUTHORIZED {
self.record_401_attribution(ToolConsumer::ImageGen, sent_bearer.as_deref());
}
if !status.is_success() {
let body = response.text().await.unwrap_or_default();
let truncated: String = body.chars().take(200).collect();
tracing::warn!(http_status = %status, "Imagine API error: {truncated}");
return Err(kigi_tool_runtime::ToolError::new(
kigi_tool_runtime::ToolErrorKind::Custom,
format!("Image generation failed with HTTP {status}: {truncated}"),
)
.with_details(serde_json::json!({"code": "http_failure", "status": status.as_u16()})));
}
let body = response.text().await.map_err(|e| {
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"Failed to read image generation response body: {e}"
))
})?;
let resp_json: ImageGenResponse = serde_json::from_str(&body).map_err(|e| {
let preview: String = body.chars().take(500).collect();
tracing::warn!("Imagine API returned unparseable body: {preview}");
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"Failed to parse image generation response: {e} — body preview: {preview}"
))
})?;
let b64_data = resp_json.b64_data().unwrap_or("");
if b64_data.is_empty() {
return Err(kigi_tool_runtime::ToolError::invalid_arguments(
"Image generation returned no image data.",
));
}
base64::engine::general_purpose::STANDARD
.decode(b64_data)
.map_err(|e| {
kigi_tool_runtime::ToolError::invalid_arguments(format!(
"Failed to decode base64 image data: {e}"
))
})
}
}
/// `Enabled` means credentials are present; each tool has its own gate.
#[derive(Debug, Clone, Default)]
pub enum ImageGenConfig {
#[default]
Disabled,
Enabled {
api_key: String,
base_url: String,
extra_headers: indexmap::IndexMap<String, String>,
image_gen_enabled: bool,
image_edit_enabled: bool,
/// Optional Imagine model override for `image_gen`. When `Some(non-empty)`,
/// `image_gen` calls that model instead of the default quality model
/// ([`XAI_IMAGINE_MODEL`]). Driven by the remote
/// `image_gen_model_override` config flag. `image_edit` is unaffected.
model_override: Option<String>,
/// `true` when the user is on a tier the Imagine server zero-limits
/// (free / X Basic). The tools stay advertised to the model, but
/// `image_gen` / `image_edit` short-circuit at call time with the
/// SuperGrok upsell prose instead of a doomed request. Set by the
/// host from the subscription tier; always `false` for team /
/// API-key / workspace callers.
tier_restricted: bool,
},
}
impl ImageGenConfig {
/// Credentials present — required to construct any of the clients.
pub fn has_credentials(&self) -> bool {
matches!(self, Self::Enabled { .. })
}
pub fn image_gen_enabled(&self) -> bool {
matches!(
self,
Self::Enabled {
image_gen_enabled: true,
..
}
)
}
pub fn image_edit_enabled(&self) -> bool {
matches!(
self,
Self::Enabled {
image_edit_enabled: true,
..
}
)
}
/// The configured `image_gen` model override, if any. `None` means the
/// default quality model ([`XAI_IMAGINE_MODEL`]) is used.
pub fn model_override(&self) -> Option<&str> {
match self {
Self::Enabled { model_override, .. } => {
model_override.as_deref().filter(|m| !m.trim().is_empty())
}
Self::Disabled => None,
}
}
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
pub struct ImageGenInput {
#[schemars(description = "Text description of the image to generate.")]
pub prompt: String,
#[serde(default = "default_aspect_ratio")]
#[schemars(
description = "Aspect ratio of the generated image, decide it based on the user's request. Defaults to 'auto'. 1:1 for square (icons, profiles), 16:9 for wide (landscapes, cinematic), 9:16 for tall (phone wallpapers, stories), 3:2 for horizontal photos, 2:3 for vertical (portraits, posters)."
)]
pub aspect_ratio: String,
}
fn default_aspect_ratio() -> String {
"auto".to_owned()
}
#[derive(Debug, serde::Deserialize)]
pub struct ImageGenResponse {
#[serde(default)]
data: Vec<ImageGenData>,
}
impl ImageGenResponse {
pub fn b64_data(&self) -> Option<&str> {
self.data.first().and_then(|d| d.b64_json.as_deref())
}
}
#[derive(Debug, serde::Deserialize)]
struct ImageGenData {
b64_json: Option<String>,
}
#[derive(Debug, Default)]
pub struct ImageGenTool;
impl crate::types::tool_metadata::ToolMetadata for ImageGenTool {
fn kind(&self) -> ToolKind {
ToolKind::ImageGen
}
fn tool_namespace(&self) -> ToolNamespace {
ToolNamespace::GrokBuild
}
fn description_template(&self) -> &str {
"Generate a new image from a text description using Imagine; returns the saved image's absolute path. When telling the user where it was saved, refer to it by its short session-relative path (e.g. `images/1.jpg`) rather than the absolute path, so it renders as a clickable link that opens the image. To produce multiple images, emit multiple tool calls with distinct prompts."
}
fn requires_expr(&self) -> Expr<ToolRequirement> {
Expr::True
}
}
impl kigi_tool_runtime::Tool for ImageGenTool {
type Args = ImageGenInput;
type Output = ToolOutput;
fn id(&self) -> kigi_tool_protocol::ToolId {
kigi_tool_protocol::ToolId::new("image_gen").expect("valid tool id")
}
fn description(
&self,
_ctx: &::kigi_tool_runtime::ListToolsContext,
) -> kigi_tool_types::ToolDescription {
kigi_tool_types::ToolDescription::new(
"image_gen",
crate::types::tool_metadata::ToolMetadata::description_template(self),
)
}
fn capabilities(&self) -> kigi_tool_protocol::ToolCapabilities {
kigi_tool_protocol::ToolCapabilities {
is_read_only: false,
tool_scope: Some(kigi_tool_protocol::ToolScope::Write),
..Default::default()
}
}
#[tracing::instrument(
name = "tool.image_gen",
skip_all,
fields(prompt_len = input.prompt.len(), aspect_ratio = %input.aspect_ratio)
)]
async fn run(
&self,
ctx: kigi_tool_runtime::ToolCallContext,
input: ImageGenInput,
) -> Result<ToolOutput, kigi_tool_runtime::ToolError> {
use crate::types::tool_metadata::shared_resources;
let resources = shared_resources(&ctx)?;
let client = {
let res = resources.lock().await;
res.require::<ImageGenClient>()?.clone()
};
// Free / X Basic users are zero-limited on Imagine server-side; return
// the upsell prose instead of a doomed request (the tool stays
// advertised so the model can surface the nudge in-conversation).
if client.is_tier_restricted() {
return Ok(ToolOutput::Text(TIER_RESTRICTED_UPSELL.into()));
}
let image_bytes = client.generate(&input.prompt, &input.aspect_ratio).await?;
let session_folder = {
let res = resources.lock().await;
res.require::<SessionFolder>()?.0.clone()
};
let absolute_path = client
.writer
.save(&session_folder, &image_bytes, None)
.await
.map_err(|e| kigi_tool_runtime::ToolError::invalid_arguments(e.to_string()))?;
tracing::info!(
path = %absolute_path.display(),
bytes = image_bytes.len(),
"image saved to disk"
);
Ok(ToolOutput::ImageGen(MediaGenOutput::new(absolute_path)))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::tool_metadata::test_ctx_with_call_id;
#[test]
fn tool_name_and_description() {
let tool = ImageGenTool;
assert_eq!(kigi_tool_runtime::Tool::id(&tool).as_str(), "image_gen");
assert!(
crate::types::tool_metadata::ToolMetadata::description_template(&tool)
.contains("Generate a new image from a text description")
);
}
#[test]
fn default_aspect_ratio_is_auto() {
let input: ImageGenInput = serde_json::from_str(r#"{"prompt": "test"}"#).unwrap();
assert_eq!(input.aspect_ratio, "auto");
}
#[test]
fn per_tool_gates_are_independent() {
let cfg = ImageGenConfig::Enabled {
api_key: "k".into(),
base_url: "https://api.x.ai/v1".into(),
extra_headers: indexmap::IndexMap::new(),
image_gen_enabled: false,
image_edit_enabled: true,
model_override: Some("grok-imagine-image".into()),
tier_restricted: false,
};
assert!(cfg.has_credentials());
assert!(!cfg.image_gen_enabled());
assert!(cfg.image_edit_enabled());
assert_eq!(cfg.model_override(), Some("grok-imagine-image"));
assert!(!ImageGenConfig::Disabled.has_credentials());
}
#[test]
fn client_selects_model_from_override() {
let mk = |model_override: Option<&str>| ImageGenConfig::Enabled {
api_key: "k".into(),
base_url: "https://api.x.ai/v1".into(),
extra_headers: indexmap::IndexMap::new(),
image_gen_enabled: true,
image_edit_enabled: true,
model_override: model_override.map(String::from),
tier_restricted: false,
};
// No override → default quality model.
assert_eq!(
ImageGenClient::new(&mk(None), None).unwrap().model,
XAI_IMAGINE_MODEL
);
// Empty override → treated as no override.
assert_eq!(
ImageGenClient::new(&mk(Some("")), None).unwrap().model,
XAI_IMAGINE_MODEL
);
// Override → that exact model slug.
assert_eq!(
ImageGenClient::new(&mk(Some("grok-imagine-image")), None)
.unwrap()
.model,
"grok-imagine-image"
);
}
#[tokio::test]
async fn errors_when_client_missing() {
let tool = ImageGenTool;
let resources = crate::types::resources::Resources::new();
let result = kigi_tool_runtime::Tool::run(
&tool,
test_ctx_with_call_id(resources.into_shared(), "test-call"),
ImageGenInput {
prompt: "a test image".into(),
aspect_ratio: "auto".into(),
},
)
.await;
assert!(result.is_err());
let err_msg = result.unwrap_err().to_string();
assert!(
err_msg.contains("missing required resource"),
"Expected MissingResource error, got: {err_msg}"
);
}
#[tokio::test]
async fn tier_restricted_short_circuits_with_upsell() {
// A free / X Basic user's image_gen call returns the SuperGrok upsell
// prose as a normal result (no HTTP, no error card) so the model can
// relay it. Only the client is inserted — the short-circuit returns
// before any other resource (e.g. SessionFolder) is required.
let cfg = ImageGenConfig::Enabled {
api_key: "k".into(),
base_url: "https://api.x.ai/v1".into(),
extra_headers: indexmap::IndexMap::new(),
image_gen_enabled: true,
image_edit_enabled: true,
model_override: None,
tier_restricted: true,
};
let mut resources = crate::types::resources::Resources::new();
resources.insert(ImageGenClient::new(&cfg, None).unwrap());
let result = kigi_tool_runtime::Tool::run(
&ImageGenTool,
test_ctx_with_call_id(resources.into_shared(), "test-call"),
ImageGenInput {
prompt: "a cat".into(),
aspect_ratio: "auto".into(),
},
)
.await
.expect("tier-restricted call must succeed with upsell prose");
match result {
ToolOutput::Text(t) => {
assert!(t.text.contains("SuperGrok"), "got: {}", t.text);
assert!(t.text.contains("supergrok?referrer=grok-build"));
}
other => panic!("expected Text upsell, got {other:?}"),
}
}
}
@@ -15,8 +15,6 @@ pub mod deploy_app;
pub mod enter_plan_mode; pub mod enter_plan_mode;
pub mod exit_plan_mode; pub mod exit_plan_mode;
pub mod grep; pub mod grep;
pub mod image_edit;
pub mod image_gen;
pub mod kill_task; pub mod kill_task;
pub mod list_dir; pub mod list_dir;
pub mod lsp; pub mod lsp;
@@ -29,7 +27,6 @@ pub mod task;
pub mod task_output; pub mod task_output;
pub mod todo; pub mod todo;
pub mod update_goal; pub mod update_goal;
pub mod video_gen;
pub mod web_fetch; pub mod web_fetch;
pub mod web_search; pub mod web_search;
pub use ask_user_question::AskUserQuestionTool; pub use ask_user_question::AskUserQuestionTool;
@@ -38,11 +35,6 @@ pub use deploy_app::{AppBuilderDeployerConfig, DEPLOY_APP_TOOL_NAME};
pub use enter_plan_mode::EnterPlanModeTool; pub use enter_plan_mode::EnterPlanModeTool;
pub use exit_plan_mode::ExitPlanModeTool; pub use exit_plan_mode::ExitPlanModeTool;
pub use grep::GrepTool; pub use grep::GrepTool;
pub use image_edit::{IMAGE_EDIT_TOOL_NAME, ImageEditTool};
pub use image_gen::{
IMAGE_GEN_TOOL_NAME, IMAGINE_COMMAND_NAME, ImageGenTool, imagine_instruction,
imagine_usage_message,
};
pub use kill_task::{KillTaskTool, KillTerminalCommandTool}; pub use kill_task::{KillTaskTool, KillTerminalCommandTool};
pub use list_dir::ListDirTool; pub use list_dir::ListDirTool;
pub use lsp::LspTool; pub use lsp::LspTool;
@@ -58,10 +50,5 @@ pub use task::TaskTool;
pub use task_output::{GetTerminalCommandOutputTool, TaskOutputTool, WaitTasksTool}; pub use task_output::{GetTerminalCommandOutputTool, TaskOutputTool, WaitTasksTool};
pub use todo::TodoWriteTool; pub use todo::TodoWriteTool;
pub use update_goal::{UPDATE_GOAL_TOOL_NAME, UpdateGoalTool}; pub use update_goal::{UPDATE_GOAL_TOOL_NAME, UpdateGoalTool};
pub use video_gen::{
IMAGE_TO_VIDEO_TOOL_NAME, IMAGINE_VIDEO_COMMAND_NAME, ImageToVideoTool,
REFERENCE_TO_VIDEO_TOOL_NAME, ReferenceToVideoTool, imagine_video_instruction,
imagine_video_usage_message,
};
pub use web_fetch::{WebFetchClient, WebFetchConfig, WebFetchParams, WebFetchTool}; pub use web_fetch::{WebFetchClient, WebFetchConfig, WebFetchParams, WebFetchTool};
pub use web_search::WebSearchTool; pub use web_search::WebSearchTool;
@@ -237,10 +237,6 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode {
ToolKind::MemoryGet, ToolKind::MemoryGet,
ToolKind::WebSearch, ToolKind::WebSearch,
ToolKind::WebFetch, ToolKind::WebFetch,
ToolKind::ImageGen,
ToolKind::VideoGen,
ToolKind::ImageToVideo,
ToolKind::ReferenceToVideo,
ToolKind::BackgroundTaskAction, ToolKind::BackgroundTaskAction,
ToolKind::KillTaskAction, ToolKind::KillTaskAction,
ToolKind::Task, ToolKind::Task,
@@ -285,10 +281,6 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode {
ToolKind::MemoryGet, ToolKind::MemoryGet,
ToolKind::WebSearch, ToolKind::WebSearch,
ToolKind::WebFetch, ToolKind::WebFetch,
ToolKind::ImageGen,
ToolKind::VideoGen,
ToolKind::ImageToVideo,
ToolKind::ReferenceToVideo,
ToolKind::BackgroundTaskAction, ToolKind::BackgroundTaskAction,
ToolKind::KillTaskAction, ToolKind::KillTaskAction,
ToolKind::Task, ToolKind::Task,
File diff suppressed because it is too large Load Diff
@@ -3,7 +3,7 @@
use schemars::JsonSchema; use schemars::JsonSchema;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use crate::types::output::{MCPOutput, ToolOutput}; use crate::types::output::ToolOutput;
use crate::types::tool::{ToolKind, ToolNamespace}; use crate::types::tool::{ToolKind, ToolNamespace};
use crate::util::mcp_truncate::{McpTruncateContext, truncate_tool_output}; use crate::util::mcp_truncate::{McpTruncateContext, truncate_tool_output};
@@ -94,48 +94,6 @@ async fn dispatch_local_mcp(
.map_err(|e| kigi_tool_runtime::ToolError::custom("output_decoding", e.to_string())) .map_err(|e| kigi_tool_runtime::ToolError::custom("output_decoding", e.to_string()))
} }
fn gateway_result_is_error(result: &serde_json::Value) -> bool {
result
.get("isError")
.or_else(|| result.get("is_error"))
.and_then(|v| v.as_bool())
.unwrap_or(false)
}
fn gateway_result_to_text(result: serde_json::Value) -> String {
if let Some(content) = result.get("content").and_then(|v| v.as_array()) {
let parts: Vec<String> = content
.iter()
.filter_map(|item| {
if item.get("type").and_then(|v| v.as_str()) == Some("text") {
item.get("text").and_then(|v| v.as_str()).map(str::to_owned)
} else if item.get("type").and_then(|v| v.as_str()) == Some("image") {
let mime = item
.get("mimeType")
.or_else(|| item.get("mime_type"))
.and_then(|v| v.as_str())
.unwrap_or("image/png");
item.get("data")
.and_then(|v| v.as_str())
.map(|data| format!("data:{mime};base64,{data}"))
} else if item.get("type").and_then(|v| v.as_str()) == Some("resource") {
serde_json::to_string(item).ok()
} else {
None
}
})
.collect();
if !parts.is_empty() {
return parts.join("\n");
}
}
match result {
serde_json::Value::String(s) => s,
other => serde_json::to_string_pretty(&other).unwrap_or_default(),
}
}
fn normalize_mcp_arguments(input: serde_json::Value) -> serde_json::Value { fn normalize_mcp_arguments(input: serde_json::Value) -> serde_json::Value {
match input { match input {
serde_json::Value::String(s) => match serde_json::from_str(&s) { serde_json::Value::String(s) => match serde_json::from_str(&s) {
@@ -147,54 +105,6 @@ fn normalize_mcp_arguments(input: serde_json::Value) -> serde_json::Value {
} }
} }
fn is_local_tool_id_rejection(err: &kigi_tool_runtime::ToolError, tool_name: &str) -> bool {
err.kind == kigi_tool_runtime::ToolErrorKind::InvalidArguments
&& err.detail == format!("invalid tool name: '{tool_name}'")
}
async fn gateway_lookup(
ctx: &kigi_tool_runtime::ToolCallContext,
tool_name: &str,
) -> (
Option<crate::types::resources::ManagedGatewayToolSource>,
Option<crate::types::resources::ManagedGatewayToolClient>,
) {
let Some(resources) = crate::types::tool_metadata::shared_resources(ctx).ok() else {
return (None, None);
};
let guard = resources.lock().await;
let source = guard
.get::<crate::types::resources::ManagedGatewayToolCatalog>()
.and_then(|catalog| catalog.get(tool_name).cloned());
let client = guard
.get::<crate::types::resources::ManagedGatewayToolClient>()
.cloned()
.filter(|_| source.is_some());
(source, client)
}
fn gateway_response_to_output(
tool_name: &str,
source: crate::types::resources::ManagedGatewayToolSource,
result: serde_json::Value,
) -> ToolOutput {
let is_error = gateway_result_is_error(&result);
let text = gateway_result_to_text(result);
if is_error {
ToolOutput::MCP(MCPOutput::errored(
tool_name.to_owned(),
source.connector_name,
text,
))
} else {
ToolOutput::MCP(MCPOutput::okay_output(
tool_name.to_owned(),
source.connector_name,
text,
))
}
}
pub async fn dispatch_mcp_tool( pub async fn dispatch_mcp_tool(
ctx: &kigi_tool_runtime::ToolCallContext, ctx: &kigi_tool_runtime::ToolCallContext,
tool_name: &str, tool_name: &str,
@@ -202,71 +112,16 @@ pub async fn dispatch_mcp_tool(
caller: &str, caller: &str,
) -> Result<ToolOutput, kigi_tool_runtime::ToolError> { ) -> Result<ToolOutput, kigi_tool_runtime::ToolError> {
let tool_input = normalize_mcp_arguments(tool_input); let tool_input = normalize_mcp_arguments(tool_input);
let (gateway_source, gateway_client) = gateway_lookup(ctx, tool_name).await; let Some(dispatch) = ctx
let dispatch = ctx
.extensions .extensions
.get::<crate::types::resources::InnerDispatch>(); .get::<crate::types::resources::InnerDispatch>()
else {
if gateway_source.is_none() && dispatch.is_none() {
return Err(kigi_tool_runtime::ToolError::invalid_arguments(format!( return Err(kigi_tool_runtime::ToolError::invalid_arguments(format!(
"{caller} called outside of tool execution context. inner_dispatch not set -- this is a bug." "{caller} called outside of tool execution context. inner_dispatch not set -- this is a bug."
))); )));
}
if let Some(source) = gateway_source {
// A gateway-catalog name can collide with a local `server__tool` MCP
// tool. Local wins on a name clash: probe local dispatch first and only
// fall through to the gateway when the local side reports the tool as
// not found, or rejects the catalog-derived name as an invalid local
// ToolId. A real error from a local tool that actually dispatched
// propagates instead of silently retrying against the gateway.
if tool_name.contains("__")
&& let Some(dispatch) = dispatch.clone()
{
match dispatch_local_mcp(dispatch, tool_name, tool_input.clone(), ctx.clone()).await {
Ok(local_output) => return Ok(local_output),
Err(err)
if err.kind != kigi_tool_runtime::ToolErrorKind::NotFound
&& !is_local_tool_id_rejection(&err, tool_name) =>
{
return Err(err);
}
Err(_) => {}
}
}
let Some(client) = gateway_client else {
return Err(kigi_tool_runtime::ToolError::custom(
"managed_gateway_unavailable",
format!(
"Managed MCP gateway tool '{}' is indexed but no gateway client is available.",
tool_name
),
));
}; };
let response = client
.0
.call_tool(&source.call_id, tool_input, caller)
.await?;
tracing::debug!(
tool_name = %tool_name,
reauth = response.connectors_needing_reauth.len(),
"Managed MCP gateway tool call completed"
);
return Ok(gateway_response_to_output(
tool_name,
source,
response.result,
));
}
dispatch_local_mcp( dispatch_local_mcp(dispatch, tool_name, tool_input, ctx.clone()).await
dispatch.expect("dispatch is set for local MCP path"),
tool_name,
tool_input,
ctx.clone(),
)
.await
} }
impl crate::types::tool_metadata::ToolMetadata for UseTool { impl crate::types::tool_metadata::ToolMetadata for UseTool {
@@ -316,15 +171,11 @@ impl kigi_tool_runtime::Tool for UseTool {
ctx: kigi_tool_runtime::ToolCallContext, ctx: kigi_tool_runtime::ToolCallContext,
input: UseToolInput, input: UseToolInput,
) -> Result<ToolOutput, kigi_tool_runtime::ToolError> { ) -> Result<ToolOutput, kigi_tool_runtime::ToolError> {
use crate::types::resources::{EnabledNativeToolNames, ManagedGatewayToolCatalog, Params}; use crate::types::resources::{EnabledNativeToolNames, Params};
let resources = crate::types::tool_metadata::shared_resources(&ctx).ok(); let resources = crate::types::tool_metadata::shared_resources(&ctx).ok();
let (gateway_source, is_native, search_tool_name) = let (is_native, search_tool_name) = if let Some(resources) = resources.as_ref() {
if let Some(resources) = resources.as_ref() {
let guard = resources.lock().await; let guard = resources.lock().await;
let gateway_source = guard
.get::<ManagedGatewayToolCatalog>()
.and_then(|catalog| catalog.get(&input.tool_name).cloned());
let correction_enabled = guard let correction_enabled = guard
.get::<Params<UseToolParams>>() .get::<Params<UseToolParams>>()
.is_none_or(|p| p.0.native_tool_correction); .is_none_or(|p| p.0.native_tool_correction);
@@ -337,12 +188,12 @@ impl kigi_tool_runtime::Tool for UseTool {
.and_then(|r| r.tool_for_kind(ToolKind::SearchTool)) .and_then(|r| r.tool_for_kind(ToolKind::SearchTool))
.map(str::to_string) .map(str::to_string)
.unwrap_or_else(|| "search_tool".to_string()); .unwrap_or_else(|| "search_tool".to_string());
(gateway_source, native, st) (native, st)
} else { } else {
(None, false, "search_tool".to_string()) (false, "search_tool".to_string())
}; };
if !input.tool_name.contains("__") && gateway_source.is_none() { if !input.tool_name.contains("__") {
return Err(if is_native { return Err(if is_native {
// Native tool wrongly routed through use_tool. Tell the model // Native tool wrongly routed through use_tool. Tell the model
// to call it directly. Strategy chosen via offline eval over // to call it directly. Strategy chosen via offline eval over
@@ -448,39 +299,6 @@ mod tests {
(ctx, args) (ctx, args)
} }
struct NotFoundDispatch;
struct InvalidArgumentsDispatch;
#[async_trait::async_trait]
impl kigi_tool_runtime::ToolDispatch for NotFoundDispatch {
async fn call(
&self,
tool_id: kigi_tool_protocol::ToolId,
_args: serde_json::Value,
_ctx: kigi_tool_runtime::ToolCallContext,
) -> kigi_tool_runtime::ToolStream<kigi_tool_runtime::TypedToolOutput> {
kigi_tool_runtime::terminal_only(Err(kigi_tool_runtime::ToolError::not_found(
tool_id,
"Tool not found",
)))
}
}
#[async_trait::async_trait]
impl kigi_tool_runtime::ToolDispatch for InvalidArgumentsDispatch {
async fn call(
&self,
_tool_id: kigi_tool_protocol::ToolId,
_args: serde_json::Value,
_ctx: kigi_tool_runtime::ToolCallContext,
) -> kigi_tool_runtime::ToolStream<kigi_tool_runtime::TypedToolOutput> {
kigi_tool_runtime::terminal_only(Err(kigi_tool_runtime::ToolError::invalid_arguments(
"local validation failed",
)))
}
}
/// Mock dispatch that always returns an error. /// Mock dispatch that always returns an error.
struct ErrorToolDispatch { struct ErrorToolDispatch {
error: String, error: String,
@@ -604,238 +422,6 @@ mod tests {
assert!(err.detail.contains("bad__tool")); assert!(err.detail.contains("bad__tool"));
} }
#[derive(Clone)]
struct MockGatewayCaller {
captured: SharedArgs,
result: serde_json::Value,
expected_call_id: Option<&'static str>,
}
#[async_trait::async_trait]
impl crate::types::resources::ManagedGatewayToolCaller for MockGatewayCaller {
async fn call_tool(
&self,
call_id: &str,
arguments: serde_json::Value,
_caller: &str,
) -> Result<
crate::types::resources::ManagedGatewayToolCallResponse,
kigi_tool_runtime::ToolError,
> {
if let Some(expected) = self.expected_call_id {
assert_eq!(call_id, expected);
}
*self.captured.lock().unwrap() = Some(arguments);
Ok(crate::types::resources::ManagedGatewayToolCallResponse {
result: self.result.clone(),
connectors_needing_reauth: vec![],
})
}
}
fn gateway_resources(
captured: SharedArgs,
result: serde_json::Value,
) -> crate::types::resources::SharedResources {
gateway_resources_with_expected_call_id(captured, result, Some("grafana.searchDashboards"))
}
fn gateway_resources_with_expected_call_id(
captured: SharedArgs,
result: serde_json::Value,
expected_call_id: Option<&'static str>,
) -> crate::types::resources::SharedResources {
use crate::types::resources::{
ManagedGatewayToolCatalog, ManagedGatewayToolClient, ManagedGatewayToolSource,
Resources,
};
let mut resources = Resources::new();
resources.insert(ManagedGatewayToolCatalog(std::collections::HashMap::from(
[
(
"grafana__search_dashboards".to_string(),
ManagedGatewayToolSource {
connector_id: "grafana".to_string(),
connector_name: "Grafana".to_string(),
tool_id: "search_dashboards".to_string(),
tool_name: "Search Dashboards".to_string(),
call_id: "grafana.searchDashboards".to_string(),
},
),
(
"server__tool".to_string(),
ManagedGatewayToolSource {
connector_id: "server".to_string(),
connector_name: "Gateway Collision".to_string(),
tool_id: "tool".to_string(),
tool_name: "Tool".to_string(),
call_id: "gateway.collision".to_string(),
},
),
(
"connector__bad/id".to_string(),
ManagedGatewayToolSource {
connector_id: "connector".to_string(),
connector_name: "Gateway Invalid Local".to_string(),
tool_id: "bad/id".to_string(),
tool_name: "Bad ID".to_string(),
call_id: "gateway.invalidLocal".to_string(),
},
),
],
)));
resources.insert(ManagedGatewayToolClient(Arc::new(MockGatewayCaller {
captured,
result,
expected_call_id,
})));
resources.into_shared()
}
#[tokio::test]
async fn gateway_tool_dispatches_to_gateway_call_id() {
let captured: SharedArgs = Arc::new(std::sync::Mutex::new(None));
let ctx = ctx_with_dispatch_and_resources(
NotFoundDispatch,
gateway_resources(
Arc::clone(&captured),
serde_json::json!({"content": [{"type": "text", "text": "dashboards"}]}),
),
);
let result = kigi_tool_runtime::Tool::run(
&UseTool,
ctx,
UseToolInput {
tool_name: "grafana__search_dashboards".into(),
tool_input: serde_json::json!({"query": "prod"}),
},
)
.await
.unwrap();
assert_eq!(captured.lock().unwrap().clone().unwrap()["query"], "prod");
if let ToolOutput::MCP(mcp) = result {
match mcp.output() {
crate::types::output::MCPOutputDetails::OkayOutput(text) => {
assert_eq!(text, "dashboards")
}
_ => panic!("expected okay output"),
}
} else {
panic!("expected gateway result to map to MCP output");
}
}
#[tokio::test]
async fn gateway_error_result_maps_to_mcp_error() {
let captured: SharedArgs = Arc::new(std::sync::Mutex::new(None));
let ctx = ctx_with_dispatch_and_resources(
NotFoundDispatch,
gateway_resources(
Arc::clone(&captured),
serde_json::json!({
"isError": true,
"content": [{"type": "text", "text": "remote failed"}]
}),
),
);
let result = kigi_tool_runtime::Tool::run(
&UseTool,
ctx,
UseToolInput {
tool_name: "grafana__search_dashboards".into(),
tool_input: serde_json::json!({}),
},
)
.await
.unwrap();
assert!(result.is_error());
assert!(
result
.to_prompt_format()
.contains("Failed to call grafana__search_dashboards: remote failed")
);
}
#[tokio::test]
async fn gateway_snake_case_error_result_maps_to_mcp_error() {
let captured: SharedArgs = Arc::new(std::sync::Mutex::new(None));
let ctx = ctx_with_dispatch_and_resources(
NotFoundDispatch,
gateway_resources(
Arc::clone(&captured),
serde_json::json!({
"is_error": true,
"content": [{"type": "text", "text": "snake failed"}]
}),
),
);
let result = kigi_tool_runtime::Tool::run(
&UseTool,
ctx,
UseToolInput {
tool_name: "grafana__search_dashboards".into(),
tool_input: serde_json::json!({}),
},
)
.await
.unwrap();
assert!(result.is_error());
assert!(result.to_prompt_format().contains("snake failed"));
}
#[tokio::test]
async fn gateway_call_result_converts_to_model_visible_output() {
let captured: SharedArgs = Arc::new(std::sync::Mutex::new(None));
let ctx = ctx_with_dispatch_and_resources(
NotFoundDispatch,
gateway_resources(Arc::clone(&captured), serde_json::json!({"ok": true})),
);
let result = kigi_tool_runtime::Tool::run(
&UseTool,
ctx,
UseToolInput {
tool_name: "grafana__search_dashboards".into(),
tool_input: serde_json::json!({}),
},
)
.await
.unwrap();
assert!(result.to_prompt_format().contains("\"ok\": true"));
}
#[tokio::test]
async fn gateway_null_arguments_default_to_object() {
let captured: SharedArgs = Arc::new(std::sync::Mutex::new(None));
let ctx = ctx_with_dispatch_and_resources(
NotFoundDispatch,
gateway_resources(Arc::clone(&captured), serde_json::json!("ok")),
);
kigi_tool_runtime::Tool::run(
&UseTool,
ctx,
UseToolInput {
tool_name: "grafana__search_dashboards".into(),
tool_input: serde_json::Value::Null,
},
)
.await
.unwrap();
assert_eq!(
captured.lock().unwrap().clone().unwrap(),
serde_json::json!({})
);
}
#[tokio::test] #[tokio::test]
async fn normalizes_string_encoded_tool_input() { async fn normalizes_string_encoded_tool_input() {
let tool = UseTool; let tool = UseTool;
@@ -902,88 +488,6 @@ mod tests {
assert_eq!(captured, serde_json::Value::String("not json".into())); assert_eq!(captured, serde_json::Value::String("not json".into()));
} }
#[tokio::test]
async fn gateway_tool_with_invalid_local_tool_id_falls_back_to_gateway() {
let gateway_captured: SharedArgs = Arc::new(std::sync::Mutex::new(None));
let ctx = ctx_with_dispatch_and_resources(
NotFoundDispatch,
gateway_resources_with_expected_call_id(
Arc::clone(&gateway_captured),
serde_json::json!("gateway ran"),
Some("gateway.invalidLocal"),
),
);
let result = kigi_tool_runtime::Tool::run(
&UseTool,
ctx,
UseToolInput {
tool_name: "connector__bad/id".into(),
tool_input: serde_json::json!({"q": "x"}),
},
)
.await
.unwrap();
assert_eq!(gateway_captured.lock().unwrap().clone().unwrap()["q"], "x");
assert!(matches!(result, ToolOutput::MCP(_)));
}
#[tokio::test]
async fn gateway_catalog_collision_propagates_local_non_not_found_error() {
let gateway_captured: SharedArgs = Arc::new(std::sync::Mutex::new(None));
let ctx = ctx_with_dispatch_and_resources(
InvalidArgumentsDispatch,
gateway_resources(
Arc::clone(&gateway_captured),
serde_json::json!("gateway should not run"),
),
);
let result = kigi_tool_runtime::Tool::run(
&UseTool,
ctx,
UseToolInput {
tool_name: "server__tool".into(),
tool_input: serde_json::json!({"local": true}),
},
)
.await;
let err = result.unwrap_err();
assert_eq!(err.kind, kigi_tool_runtime::ToolErrorKind::InvalidArguments);
assert!(err.detail.contains("local validation failed"));
assert!(gateway_captured.lock().unwrap().is_none());
}
#[tokio::test]
async fn gateway_catalog_collision_prefers_local_dispatch_for_server_tool() {
let captured: SharedArgs = Arc::new(std::sync::Mutex::new(None));
let ctx = ctx_with_dispatch_and_resources(
CapturingDispatch {
captured_args: Arc::clone(&captured),
},
gateway_resources(
Arc::new(std::sync::Mutex::new(None)),
serde_json::json!("gateway should not run"),
),
);
let result = kigi_tool_runtime::Tool::run(
&UseTool,
ctx,
UseToolInput {
tool_name: "server__tool".into(),
tool_input: serde_json::json!({"local": true}),
},
)
.await;
assert!(result.is_ok());
let captured = captured.lock().unwrap().clone().unwrap();
assert_eq!(captured, serde_json::json!({"local": true}));
}
#[tokio::test] #[tokio::test]
async fn local_server_tool_still_uses_local_dispatch_path() { async fn local_server_tool_still_uses_local_dispatch_path() {
let tool = UseTool; let tool = UseTool;
@@ -105,10 +105,6 @@ pub fn canonical_input(input: &ToolInput) -> Option<serde_json::Value> {
| ToolInput::KillTask(_) | ToolInput::KillTask(_)
| ToolInput::Task(_) | ToolInput::Task(_)
| ToolInput::WebSearch(_) | ToolInput::WebSearch(_)
| ToolInput::ImageGen(_)
| ToolInput::ImageEdit(_)
| ToolInput::ImageToVideo(_)
| ToolInput::ReferenceToVideo(_)
| ToolInput::WebFetch(_) | ToolInput::WebFetch(_)
| ToolInput::ApplyPatch(_) | ToolInput::ApplyPatch(_)
| ToolInput::HashlineEdit(_) | ToolInput::HashlineEdit(_)
@@ -263,16 +263,6 @@ pub struct SessionContext {
/// passed to every session. Same pattern as `fs` and `backend`. /// passed to every session. Same pattern as `fs` and `backend`.
/// When `Some`, inserted into `Resources` so `LspTool` can use it. /// When `Some`, inserted into `Resources` so `LspTool` can use it.
pub lsp: Option<std::sync::Arc<dyn crate::implementations::lsp::LspBackend>>, pub lsp: Option<std::sync::Arc<dyn crate::implementations::lsp::LspBackend>>,
/// Optional image generation configuration. When `Enabled`, an `ImageGenClient`
/// is created and injected into `Resources` so the `image_gen` tool can
/// call the xAI Imagine API. When `Disabled` (default), the tool is not
/// registered and image generation is unavailable.
pub image_gen_config: crate::implementations::grok_build::image_gen::ImageGenConfig,
/// Optional video generation configuration. When `Enabled`, a `VideoGenClient`
/// is created and injected into `Resources` so the `video_gen` tool can
/// call the xAI Video Generation API. When `Disabled` (default), the tool is not
/// registered and video generation is unavailable.
pub video_gen_config: crate::implementations::grok_build::video_gen::VideoGenConfig,
/// Optional deploy service configuration. When enabled, the /// Optional deploy service configuration. When enabled, the
/// `deploy_app` tool connects to the service at call time using the shared /// `deploy_app` tool connects to the service at call time using the shared
/// API key provider. /// API key provider.
@@ -284,7 +274,7 @@ pub struct SessionContext {
/// Prevents 401 failures when a session outlives the initial token lifetime. /// Prevents 401 failures when a session outlives the initial token lifetime.
pub api_key_provider: Option<crate::types::SharedApiKeyProvider>, pub api_key_provider: Option<crate::types::SharedApiKeyProvider>,
/// Optional 401-attribution callback for tool HTTP clients. When /// Optional 401-attribution callback for tool HTTP clients. When
/// set, a 401 from `image_gen` / `video_gen` / `web_search` /// set, a 401 from `web_search`
/// emits an `auth_401_attribution` event via this hook. Hosts can /// emits an `auth_401_attribution` event via this hook. Hosts can
/// wire this to the same attribution sink used for inference-side /// wire this to the same attribution sink used for inference-side
/// 401s so tool and chat auth failures share one telemetry path. /// 401s so tool and chat auth failures share one telemetry path.
@@ -673,10 +663,6 @@ impl ToolRegistryBuilder {
b.register::<grok_build::WebSearchTool>(); b.register::<grok_build::WebSearchTool>();
b.register_with_params::<grok_build::WebFetchTool, grok_build::web_fetch::WebFetchParams>(); b.register_with_params::<grok_build::WebFetchTool, grok_build::web_fetch::WebFetchParams>();
b.register::<grok_build::LspTool>(); b.register::<grok_build::LspTool>();
b.register::<grok_build::ImageGenTool>();
b.register::<grok_build::ImageEditTool>();
b.register::<grok_build::ImageToVideoTool>();
b.register::<grok_build::ReferenceToVideoTool>();
b.register::<grok_build::EnterPlanModeTool>(); b.register::<grok_build::EnterPlanModeTool>();
b.register::<grok_build::ExitPlanModeTool>(); b.register::<grok_build::ExitPlanModeTool>();
b.register_with_params::< b.register_with_params::<
@@ -994,34 +980,6 @@ impl ToolRegistryBuilder {
if let Some(lsp) = ctx.lsp { if let Some(lsp) = ctx.lsp {
resources.insert(lsp); resources.insert(lsp);
} }
if ctx.image_gen_config.has_credentials() {
match crate::implementations::grok_build::image_gen::ImageGenClient::new(
&ctx.image_gen_config,
ctx.api_key_provider.clone(),
) {
Ok(client) => {
let client = client.with_attribution_callback(ctx.attribution_callback.clone());
resources.insert(client);
}
Err(e) => {
tracing::warn!("Failed to create ImageGenClient: {e}");
}
}
}
if ctx.video_gen_config.is_enabled() {
match crate::implementations::grok_build::video_gen::VideoGenClient::new(
&ctx.video_gen_config,
ctx.api_key_provider.clone(),
) {
Ok(client) => {
let client = client.with_attribution_callback(ctx.attribution_callback.clone());
resources.insert(client);
}
Err(e) => {
tracing::warn!("Failed to create VideoGenClient: {e}");
}
}
}
if let crate::implementations::grok_build::web_fetch::WebFetchConfig::Enabled { params } = if let crate::implementations::grok_build::web_fetch::WebFetchConfig::Enabled { params } =
&ctx.web_fetch_config &ctx.web_fetch_config
{ {
@@ -1999,10 +1957,6 @@ mod tests {
web_fetch_config: web_fetch_config:
crate::implementations::grok_build::web_fetch::WebFetchConfig::default(), crate::implementations::grok_build::web_fetch::WebFetchConfig::default(),
lsp: None, lsp: None,
image_gen_config:
crate::implementations::grok_build::image_gen::ImageGenConfig::default(),
video_gen_config:
crate::implementations::grok_build::video_gen::VideoGenConfig::default(),
app_builder_deployer_config: app_builder_deployer_config:
crate::implementations::grok_build::deploy_app::AppBuilderDeployerConfig::default(), crate::implementations::grok_build::deploy_app::AppBuilderDeployerConfig::default(),
api_key_provider: None, api_key_provider: None,
@@ -2148,7 +2102,6 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn full_toolset_descriptions_render_cleanly() { async fn full_toolset_descriptions_render_cleanly() {
use crate::implementations::grok_build::{ use crate::implementations::grok_build::{
IMAGE_GEN_TOOL_NAME, IMAGE_TO_VIDEO_TOOL_NAME, REFERENCE_TO_VIDEO_TOOL_NAME,
SCHEDULER_CREATE_TOOL_NAME, SCHEDULER_DELETE_TOOL_NAME, SCHEDULER_CREATE_TOOL_NAME, SCHEDULER_DELETE_TOOL_NAME,
}; };
let builder = ToolRegistryBuilder::new(); let builder = ToolRegistryBuilder::new();
@@ -2169,9 +2122,6 @@ mod tests {
"web_search", "web_search",
"web_fetch", "web_fetch",
"lsp", "lsp",
IMAGE_GEN_TOOL_NAME,
IMAGE_TO_VIDEO_TOOL_NAME,
REFERENCE_TO_VIDEO_TOOL_NAME,
"monitor", "monitor",
SCHEDULER_CREATE_TOOL_NAME, SCHEDULER_CREATE_TOOL_NAME,
SCHEDULER_DELETE_TOOL_NAME, SCHEDULER_DELETE_TOOL_NAME,
@@ -597,10 +597,6 @@ pub fn consumed_completion_ids(output: &ToolOutput) -> Vec<&str> {
| ToolOutput::SchedulerDelete(_) | ToolOutput::SchedulerDelete(_)
| ToolOutput::SchedulerList(_) | ToolOutput::SchedulerList(_)
| ToolOutput::UpdateGoal(_) | ToolOutput::UpdateGoal(_)
| ToolOutput::ImageGen(_)
| ToolOutput::ImageToVideo(_)
| ToolOutput::ReferenceToVideo(_)
| ToolOutput::ImageEdit(_)
| ToolOutput::Dynamic(_) => {} | ToolOutput::Dynamic(_) => {}
} }
ids ids
@@ -59,10 +59,6 @@ impl ToolKind {
ToolKind::EnterPlan => "Enter Plan Mode", ToolKind::EnterPlan => "Enter Plan Mode",
ToolKind::ExitPlan => "Exit Plan Mode", ToolKind::ExitPlan => "Exit Plan Mode",
ToolKind::AskUser => "Ask User", ToolKind::AskUser => "Ask User",
ToolKind::ImageGen => "Generate Image",
ToolKind::VideoGen => "Generate Video",
ToolKind::ImageToVideo => "Generate Video",
ToolKind::ReferenceToVideo => "Generate Video",
ToolKind::DeployApp => "Deploy App", ToolKind::DeployApp => "Deploy App",
ToolKind::SearchTool => "Search Tools", ToolKind::SearchTool => "Search Tools",
ToolKind::UseTool => "Use Tool", ToolKind::UseTool => "Use Tool",
@@ -100,10 +96,6 @@ impl ToolKind {
| ToolKind::KillTaskAction | ToolKind::KillTaskAction
| ToolKind::Skill | ToolKind::Skill
| ToolKind::Task | ToolKind::Task
| ToolKind::ImageGen
| ToolKind::VideoGen
| ToolKind::ImageToVideo
| ToolKind::ReferenceToVideo
| ToolKind::DeployApp | ToolKind::DeployApp
| ToolKind::SearchTool | ToolKind::SearchTool
| ToolKind::UseTool | ToolKind::UseTool
@@ -55,73 +55,6 @@ impl From<serde_json::Value> for DynamicOutput {
Self { value } Self { value }
} }
} }
/// Typed saved path for the media tools (`image_gen` / `video_gen` /
/// `image_edit`), so consumers read it directly instead of scraping the prose.
/// A struct (not a bare `PathBuf`) is required: `ToolOutput` is internally
/// tagged and only accepts map payloads.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MediaGenOutput {
/// Absolute path to the saved media file. Empty for [`Self::uploaded`].
pub path: PathBuf,
/// Basename of the saved media file (for example, `8.jpg`).
#[serde(default)]
pub filename: String,
/// Session-relative media directory name (for example, `images` or `videos`).
#[serde(default)]
pub session_folder: String,
/// Set when the media was uploaded to a remote presigned URL (ZDR video
/// output) and is not available locally; omitted otherwise.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub uploaded_url: Option<String>,
}
impl MediaGenOutput {
pub fn new(path: PathBuf) -> Self {
let filename = path
.file_name()
.map(|name| name.to_string_lossy().into_owned())
.unwrap_or_default();
let session_folder = path
.parent()
.and_then(|parent| parent.file_name())
.map(|name| name.to_string_lossy().into_owned())
.unwrap_or_default();
Self {
path,
filename,
session_folder,
uploaded_url: None,
}
}
/// Media uploaded to a remote presigned URL and not available locally
/// (ZDR video output). No local path/filename/session folder.
pub fn uploaded(url: String) -> Self {
Self {
path: PathBuf::new(),
filename: String::new(),
session_folder: String::new(),
uploaded_url: Some(url),
}
}
/// Model-facing prose. `action` is the variant's lead-in
/// ("Image generated" / "Video generated" / "Image edited"); the trailing
/// guidance stops the model re-reading or narrating the result.
pub fn prompt_text(&self, action: &str) -> String {
if let Some(url) = &self.uploaded_url {
return format!(
"{action} and uploaded to {url}. The file is not available locally — reference it by this URL. Do not read or re-display it, and do not describe how it appears to the user."
);
}
let path = self.path.to_string_lossy().to_string();
let message = format!(
"{action} and saved to {path}. Do not read or re-display it, and do not describe how it appears to the user."
);
serde_json::json!(
{ "path" : path, "filename" : & self.filename, "session_folder" : & self
.session_folder, "message" : message, }
)
.to_string()
}
}
use crate::implementations::grok_build::todo::{TodoItem, TodoState}; use crate::implementations::grok_build::todo::{TodoItem, TodoState};
use crate::implementations::skills::skill::SkillOutput; use crate::implementations::skills::skill::SkillOutput;
use crate::util::truncate::{DEFAULT_SOFT_WRAP_WIDTH, soft_wrap_lines}; use crate::util::truncate::{DEFAULT_SOFT_WRAP_WIDTH, soft_wrap_lines};
@@ -653,14 +586,6 @@ pub enum ToolOutput {
/// (e.g., memory_search, memory_get). The string is the pre-formatted /// (e.g., memory_search, memory_get). The string is the pre-formatted
/// prompt text — no additional rendering is needed. /// prompt text — no additional rendering is needed.
Text(TextOutput), Text(TextOutput),
#[from(skip)]
ImageGen(MediaGenOutput),
#[from(skip)]
ImageToVideo(MediaGenOutput),
#[from(skip)]
ReferenceToVideo(MediaGenOutput),
#[from(skip)]
ImageEdit(MediaGenOutput),
} }
impl ToolOutput { impl ToolOutput {
/// Whether this output is a logical tool failure, for `tool.execution`'s /// Whether this output is a logical tool failure, for `tool.execution`'s
@@ -977,10 +902,6 @@ impl ToolOutput {
ToolOutput::UpdateGoal(o) => o.summary.clone(), ToolOutput::UpdateGoal(o) => o.summary.clone(),
ToolOutput::Dynamic(v) => serde_json::to_string_pretty(&v.value).unwrap_or_default(), ToolOutput::Dynamic(v) => serde_json::to_string_pretty(&v.value).unwrap_or_default(),
ToolOutput::Text(text) => text.text.clone(), ToolOutput::Text(text) => text.text.clone(),
ToolOutput::ImageGen(m) => m.prompt_text("Image generated"),
ToolOutput::ImageToVideo(m) => m.prompt_text("Video generated"),
ToolOutput::ReferenceToVideo(m) => m.prompt_text("Video generated"),
ToolOutput::ImageEdit(m) => m.prompt_text("Image edited"),
} }
} }
} }
@@ -1320,91 +1241,6 @@ mod tests {
); );
} }
#[test] #[test]
fn media_gen_output() {
let cases = [
(
ToolOutput::ImageGen(MediaGenOutput::new("/tmp/images/1.jpg".into())),
"ImageGen",
"/tmp/images/1.jpg",
"1.jpg",
"images",
"Image generated and saved to /tmp/images/1.jpg. Do not read or re-display it, and do not describe how it appears to the user.",
),
(
ToolOutput::ImageToVideo(MediaGenOutput::new("/tmp/videos/2.mp4".into())),
"ImageToVideo",
"/tmp/videos/2.mp4",
"2.mp4",
"videos",
"Video generated and saved to /tmp/videos/2.mp4. Do not read or re-display it, and do not describe how it appears to the user.",
),
(
ToolOutput::ReferenceToVideo(MediaGenOutput::new("/tmp/videos/3.mp4".into())),
"ReferenceToVideo",
"/tmp/videos/3.mp4",
"3.mp4",
"videos",
"Video generated and saved to /tmp/videos/3.mp4. Do not read or re-display it, and do not describe how it appears to the user.",
),
(
ToolOutput::ImageEdit(MediaGenOutput::new("/tmp/images/2.jpg".into())),
"ImageEdit",
"/tmp/images/2.jpg",
"2.jpg",
"images",
"Image edited and saved to /tmp/images/2.jpg. Do not read or re-display it, and do not describe how it appears to the user.",
),
];
for (output, ty, path, filename, session_folder, message) in cases {
let prompt_json: serde_json::Value =
serde_json::from_str(&output.to_prompt_format()).unwrap();
assert_eq!(prompt_json["path"], path);
assert_eq!(prompt_json["filename"], filename);
assert_eq!(prompt_json["session_folder"], session_folder);
assert_eq!(prompt_json["message"], message);
let json = to_json(output);
assert_eq!(json["type"], ty);
assert_eq!(json["path"], path);
assert_eq!(json["filename"], filename);
assert_eq!(json["session_folder"], session_folder);
let (ToolOutput::ImageGen(m)
| ToolOutput::ImageToVideo(m)
| ToolOutput::ReferenceToVideo(m)
| ToolOutput::ImageEdit(m)) = serde_json::from_value(json).unwrap()
else {
panic!("unexpected variant");
};
assert_eq!(m.path, PathBuf::from(path));
assert_eq!(m.filename, filename);
assert_eq!(m.session_folder, session_folder);
}
}
#[test]
fn media_gen_output_uploaded() {
let url = "https://files.example.com/team/video-abc.mp4";
let output = ToolOutput::ImageToVideo(MediaGenOutput::uploaded(url.to_string()));
let prompt = output.to_prompt_format();
assert!(prompt.contains(url), "prompt must include the upload URL");
assert!(
prompt.contains("not available locally"),
"prompt must tell the model the file is remote-only"
);
assert!(
prompt.contains("Do not read or re-display"),
"prompt must include re-display guard"
);
let json = to_json(output);
assert_eq!(json["uploaded_url"], url);
assert!(
json.get("path").is_some(),
"path field must be present (empty for uploaded)"
);
let ToolOutput::ImageToVideo(m) = serde_json::from_value(json).unwrap() else {
panic!("unexpected variant");
};
assert_eq!(m, MediaGenOutput::uploaded(url.to_string()));
}
#[test]
fn read_file_not_found_json() { fn read_file_not_found_json() {
let json = let json =
to_json(ReadFileOutput::FileNotFound("Error: /tmp/x does not exist.".into()).into()); to_json(ReadFileOutput::FileNotFound("Error: /tmp/x does not exist.".into()).into());
@@ -544,42 +544,6 @@ pub fn display_cwd_or_cwd(cwd: &std::path::Path, display_cwd: Option<&std::path:
/// through the outer `ToolBridge` (which would deadlock). /// through the outer `ToolBridge` (which would deadlock).
#[derive(Clone)] #[derive(Clone)]
pub struct InnerDispatch(pub std::sync::Arc<dyn kigi_tool_runtime::ToolDispatch>); pub struct InnerDispatch(pub std::sync::Arc<dyn kigi_tool_runtime::ToolDispatch>);
#[derive(Debug, Clone)]
pub struct ManagedGatewayToolSource {
pub connector_id: String,
pub connector_name: String,
pub tool_id: String,
pub tool_name: String,
pub call_id: String,
}
#[derive(Debug, Clone, Default)]
pub struct ManagedGatewayToolCatalog(pub HashMap<String, ManagedGatewayToolSource>);
impl ManagedGatewayToolCatalog {
pub fn get(&self, name: &str) -> Option<&ManagedGatewayToolSource> {
self.0.get(name)
}
}
#[derive(Debug, Clone)]
pub struct ManagedGatewayToolCallResponse {
pub result: serde_json::Value,
pub connectors_needing_reauth: Vec<String>,
}
#[async_trait::async_trait]
pub trait ManagedGatewayToolCaller: Send + Sync {
async fn call_tool(
&self,
call_id: &str,
arguments: serde_json::Value,
caller: &str,
) -> Result<ManagedGatewayToolCallResponse, kigi_tool_runtime::ToolError>;
}
#[derive(Clone)]
pub struct ManagedGatewayToolClient(pub Arc<dyn ManagedGatewayToolCaller>);
impl std::fmt::Debug for ManagedGatewayToolClient {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ManagedGatewayToolClient").finish()
}
}
/// Whether streaming output is enabled for this invocation. /// Whether streaming output is enabled for this invocation.
#[derive(Debug, Clone, Copy)] #[derive(Debug, Clone, Copy)]
pub struct StreamEnabled(pub bool); pub struct StreamEnabled(pub bool);
@@ -753,33 +717,6 @@ impl AvailableSkills {
/// Session folder for logs and output files. /// Session folder for logs and output files.
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct SessionFolder(pub PathBuf); pub struct SessionFolder(pub PathBuf);
/// Per-turn registry mapping each attached image's `[Image #N]` display
/// number to a reference `image_edit` can resolve.
///
/// The model sees attachments inline (as pixels) and only the `[Image #N]`
/// token in text — never a path — so this lets `image_edit` resolve that
/// token instead of fabricating a filesystem path it can't know.
///
/// Keyed by display **number**, not list position: numbers are not
/// renumbered when a chip is removed mid-compose (`#1` and `#3` survive
/// after `#2`) and images may be dropped during normalization, so the two
/// diverge. Each reference is a bare filesystem path (the durable
/// `session_image_path`) or a `data:<mime>;base64,<data>` URL fallback.
///
/// Replaced wholesale each turn (empty when there are no attachments) so a
/// stale registry never resolves to a prior turn's image. Ephemeral — not
/// persisted, not serde-registered.
#[derive(Debug, Clone, Default)]
pub struct AttachedImages(pub Vec<(usize, String)>);
impl AttachedImages {
/// Resolve an `[Image #N]` display number to its reference string.
pub fn reference_for(&self, display_number: usize) -> Option<&str> {
self.0
.iter()
.find(|(n, _)| *n == display_number)
.map(|(_, reference)| reference.as_str())
}
}
/// Notification handle for streaming tool output. /// Notification handle for streaming tool output.
#[derive(Clone)] #[derive(Clone)]
pub struct NotificationHandle(pub ToolNotificationHandle); pub struct NotificationHandle(pub ToolNotificationHandle);
@@ -91,10 +91,6 @@ pub enum ToolKind {
EnterPlan, EnterPlan,
ExitPlan, ExitPlan,
AskUser, AskUser,
ImageGen,
VideoGen,
ImageToVideo,
ReferenceToVideo,
DeployApp, DeployApp,
SearchTool, SearchTool,
UseTool, UseTool,
@@ -19,14 +19,11 @@ use crate::implementations::grok_build::ask_user_question::AskUserQuestionInput;
use crate::implementations::grok_build::enter_plan_mode::EnterPlanModeInput; use crate::implementations::grok_build::enter_plan_mode::EnterPlanModeInput;
use crate::implementations::grok_build::exit_plan_mode::ExitPlanModeInput; use crate::implementations::grok_build::exit_plan_mode::ExitPlanModeInput;
use crate::implementations::grok_build::grep::GrepSearchInput; use crate::implementations::grok_build::grep::GrepSearchInput;
use crate::implementations::grok_build::image_edit::ImageEditInput;
use crate::implementations::grok_build::image_gen::ImageGenInput;
use crate::implementations::grok_build::list_dir::ListDirInput; use crate::implementations::grok_build::list_dir::ListDirInput;
use crate::implementations::grok_build::read_file::ReadFileInput; use crate::implementations::grok_build::read_file::ReadFileInput;
use crate::implementations::grok_build::search_replace::SearchReplaceInput; use crate::implementations::grok_build::search_replace::SearchReplaceInput;
use crate::implementations::grok_build::todo::TodoWriteInput; use crate::implementations::grok_build::todo::TodoWriteInput;
use crate::implementations::grok_build::update_goal::UpdateGoalInput; use crate::implementations::grok_build::update_goal::UpdateGoalInput;
use crate::implementations::grok_build::video_gen::{ImageToVideoInput, ReferenceToVideoInput};
use crate::implementations::grok_build::web_fetch::WebFetchInput; use crate::implementations::grok_build::web_fetch::WebFetchInput;
use crate::implementations::grok_build::web_search::WebSearchInput; use crate::implementations::grok_build::web_search::WebSearchInput;
use crate::implementations::lsp::LspToolInput; use crate::implementations::lsp::LspToolInput;
@@ -71,10 +68,6 @@ pub enum ToolInput {
KillTask(KillTaskToolInput), KillTask(KillTaskToolInput),
Task(TaskToolInput), Task(TaskToolInput),
WebSearch(WebSearchInput), WebSearch(WebSearchInput),
ImageGen(ImageGenInput),
ImageEdit(ImageEditInput),
ImageToVideo(ImageToVideoInput),
ReferenceToVideo(ReferenceToVideoInput),
WebFetch(WebFetchInput), WebFetch(WebFetchInput),
Write(WriteInput), Write(WriteInput),
ApplyPatch(ApplyPatchInput), ApplyPatch(ApplyPatchInput),
@@ -321,26 +321,6 @@ Open the extensions modal on the Skills tab to view installed skills.
--- ---
## Media Generation
### `/imagine <description>`
Generate an image from a text description.
```
/imagine a golden sunset over a calm ocean with silhouetted palm trees
```
### `/imagine-video <description>`
Generate a video from an image or text description. Plans shots, generates source images, and animates them with `image_to_video`.
```
/imagine-video a cat playing piano in a jazz club
```
---
## Scheduling ## Scheduling
### `/loop [interval] <prompt>` ### `/loop [interval] <prompt>`
+1 -155
View File
@@ -1897,13 +1897,6 @@ fn tool_call_to_block(tc: &acp::ToolCall, session_cwd: Option<&Path>) -> RenderB
} }
RenderBlock::ToolCall(ToolCallBlock::UseTool(block)) RenderBlock::ToolCall(ToolCallBlock::UseTool(block))
} }
_ if matches!(
extract_raw_field(tc, "variant").as_deref(),
Some("ImageGen") | Some("ImageToVideo") | Some("ReferenceToVideo") | Some("ImageEdit")
) =>
{
media_gen_block(tc, success)
}
_ if tc.title.starts_with("Memory search:") => { _ if tc.title.starts_with("Memory search:") => {
let query = tc let query = tc
.title .title
@@ -2007,45 +2000,6 @@ fn tool_call_title(tc: &acp::ToolCall) -> Cow<'_, str> {
Cow::Borrowed(&tc.title) Cow::Borrowed(&tc.title)
} }
} }
/// Build the media block from the typed `raw_output` path.
fn media_gen_block(tc: &acp::ToolCall, success: bool) -> RenderBlock {
let mut block = OtherToolCallBlock::new(tool_call_title(tc), String::new());
if !success {
let err = content_text(tc);
block.error = Some(if err.is_empty() { "Failed".into() } else { err });
} else if let Some((path, is_video)) = media_gen_ref(tc) {
block = block.with_media_ref(path, is_video);
} else if let Some(text) = media_gen_text(tc) {
block.set_output_text(text);
}
RenderBlock::ToolCall(ToolCallBlock::Other(block))
}
/// Plain-text body of a media-variant tool that returned `ToolOutput::Text`
/// rather than a media file (the free / X Basic SuperGrok-upsell short-circuit).
/// `None` for real media outputs — including ZDR upload-only results — so their
/// typed rendering is untouched.
fn media_gen_text(tc: &acp::ToolCall) -> Option<String> {
match serde_json::from_value::<ToolOutput>(tc.raw_output.clone()?).ok()? {
ToolOutput::Text(t) => (!t.text.is_empty()).then_some(t.text),
_ => None,
}
}
/// Local `(path, is_video)` from typed `raw_output`.
///
/// Returns `None` when `raw_output` is missing/unparseable, not a media
/// variant, or has no openable local file (ZDR `uploaded_url` / empty path).
fn media_gen_ref(tc: &acp::ToolCall) -> Option<(std::path::PathBuf, bool)> {
let (media, is_video) =
match serde_json::from_value::<ToolOutput>(tc.raw_output.clone()?).ok()? {
ToolOutput::ImageGen(m) | ToolOutput::ImageEdit(m) => (m, false),
ToolOutput::ImageToVideo(m) | ToolOutput::ReferenceToVideo(m) => (m, true),
_ => return None,
};
if media.uploaded_url.is_some() || media.path.as_os_str().is_empty() {
return None;
}
Some((media.path, is_video))
}
/// Extract text content from a ContentBlock. /// Extract text content from a ContentBlock.
fn extract_text_from_content(content: &acp::ContentBlock) -> String { fn extract_text_from_content(content: &acp::ContentBlock) -> String {
match content { match content {
@@ -5536,7 +5490,7 @@ mod tests {
} }
#[test] #[test]
fn build_and_parse_tools_meta_round_trip() { fn build_and_parse_tools_meta_round_trip() {
let names = vec!["scheduler_create".to_string(), "image_gen".to_string()]; let names = vec!["scheduler_create".to_string(), "web_search".to_string()];
let wire = serde_json::json!({ "tools" : names }); let wire = serde_json::json!({ "tools" : names });
assert_eq!(parse_tools_meta(wire.as_object()), Some(names)); assert_eq!(parse_tools_meta(wire.as_object()), Some(names));
} }
@@ -6421,112 +6375,4 @@ mod tests {
); );
} }
} }
/// Every video ToolInput variant must route through `media_gen_block` so
/// `[Open Video]` uses the typed `MediaGenOutput.path` (not a regex scrape
/// of the JSON prompt text — fragile on Windows with %-encoded session dirs).
#[test]
fn video_tool_variants_use_typed_path_not_generic_scrape() {
use crate::scrollback::block::BlockContent;
let dir = tempfile::tempdir().unwrap();
let video_path = dir.path().join("1.mp4");
std::fs::write(&video_path, b"fake-mp4").unwrap();
let cases: &[(&str, ToolOutput)] = &[
(
"ImageToVideo",
ToolOutput::ImageToVideo(kigi_tools::types::output::MediaGenOutput::new(
video_path.clone(),
)),
),
(
"ReferenceToVideo",
ToolOutput::ReferenceToVideo(kigi_tools::types::output::MediaGenOutput::new(
video_path.clone(),
)),
),
];
for (variant, output) in cases {
let tc = acp::ToolCall::new(
acp::ToolCallId::new(Arc::from(format!("media-{variant}"))),
variant.to_string(),
)
.kind(acp::ToolKind::Other)
.status(acp::ToolCallStatus::Completed)
.content(vec![])
.raw_input(Some(serde_json::json!({ "variant" : variant })))
.raw_output(serde_json::to_value(output).ok())
.locations(vec![]);
let block = tool_call_to_block(&tc, None);
let open_path = block
.inline_open_button()
.map(|(p, is_video)| {
assert!(is_video, "{variant}: expected video open button");
p
})
.or_else(|| block.video_references().first().map(|r| r.path.clone()))
.unwrap_or_else(|| panic!("{variant}: missing media ref / open button"));
assert_eq!(
open_path, video_path,
"{variant}: open path must be the typed MediaGenOutput.path"
);
}
}
#[test]
fn media_gen_ref_skips_uploaded_only_video() {
let output = ToolOutput::ImageToVideo(kigi_tools::types::output::MediaGenOutput::uploaded(
"https://bucket.example/videos/x.mp4".into(),
));
let tc = acp::ToolCall::new(
acp::ToolCallId::new(Arc::from("zdr-upload")),
"image_to_video",
)
.kind(acp::ToolKind::Other)
.status(acp::ToolCallStatus::Completed)
.content(vec![])
.raw_input(Some(serde_json::json!({ "variant" : "ImageToVideo" })))
.raw_output(serde_json::to_value(output).ok())
.locations(vec![]);
assert!(
media_gen_ref(&tc).is_none(),
"uploaded_url-only media must not claim a local open path"
);
}
/// A tier-restricted (free / X Basic) imagine call short-circuits with the
/// SuperGrok upsell as `ToolOutput::Text` on a `Completed` status. The media
/// renderer has no file to open, so it must surface the upsell text in the
/// card body (not a bare title) and must NOT mark the card as an error.
#[test]
fn tier_restricted_media_shows_upsell_text_not_error() {
let upsell = "Image generation is a SuperGrok feature. Upgrade at \
https://grok.com/supergrok?referrer=grok-build";
let output = ToolOutput::Text(kigi_tools::types::output::TextOutput::from(upsell));
let tc = acp::ToolCall::new(
acp::ToolCallId::new(Arc::from("tier-restricted-img")),
"image_gen",
)
.kind(acp::ToolKind::Other)
.status(acp::ToolCallStatus::Completed)
.content(vec![acp::ToolCallContent::Content(acp::Content::new(
acp::ContentBlock::Text(acp::TextContent::new(upsell)),
))])
.raw_input(Some(serde_json::json!({ "variant" : "ImageGen" })))
.raw_output(serde_json::to_value(output).ok())
.locations(vec![]);
let RenderBlock::ToolCall(ToolCallBlock::Other(block)) = tool_call_to_block(&tc, None)
else {
panic!("expected an Other tool-call block");
};
assert!(
block.is_success(),
"the upsell is a successful result, not an error"
);
assert!(
block
.output
.as_deref()
.unwrap_or_default()
.contains("SuperGrok"),
"upsell text must be shown in the card body, got: {:?}",
block.output
);
}
} }
@@ -50,9 +50,7 @@
tools: Vec::new(), tools: Vec::new(),
enabled: true, enabled: true,
source: "local".into(), source: "local".into(),
wire_source: crate::views::mcps_modal::McpWireSource::Local,
plugin_name: None, plugin_name: None,
is_managed_gateway: false,
}, },
])); ]));
} }
@@ -520,9 +518,7 @@
tools: Vec::new(), tools: Vec::new(),
enabled: true, enabled: true,
source: "local".into(), source: "local".into(),
wire_source: crate::views::mcps_modal::McpWireSource::Local,
plugin_name: None, plugin_name: None,
is_managed_gateway: false,
}, },
])); ]));
} }
@@ -1818,7 +1818,7 @@ pub(super) fn make_mcps_modal_with_servers(
state state
} }
pub(super) fn seed_owner_agent_with_open_modal(app: &mut AppView) { pub(super) fn seed_owner_agent_with_open_modal(app: &mut AppView) {
use crate::views::mcps_modal::{McpServerDisplayStatus, McpServerInfo, McpWireSource}; use crate::views::mcps_modal::{McpServerDisplayStatus, McpServerInfo};
let owner = app.agents.get_mut(&AgentId(0)).expect("owner present"); let owner = app.agents.get_mut(&AgentId(0)).expect("owner present");
owner.extensions_modal = Some( owner.extensions_modal = Some(
make_mcps_modal_with_servers( make_mcps_modal_with_servers(
@@ -1826,8 +1826,7 @@ pub(super) fn seed_owner_agent_with_open_modal(app: &mut AppView) {
McpServerInfo { name : "alpha".into(), display_name : None, status : McpServerInfo { name : "alpha".into(), display_name : None, status :
McpServerDisplayStatus::Initializing, tool_count : 0, auth_required : McpServerDisplayStatus::Initializing, tool_count : 0, auth_required :
false, tools : Vec::new(), enabled : true, source : "local".into(), false, tools : Vec::new(), enabled : true, source : "local".into(),
wire_source : McpWireSource::Local, plugin_name : None, plugin_name : None, }
is_managed_gateway : false, }
], ],
), ),
); );
@@ -58,8 +58,6 @@ pub enum Action {
ExitSessionConfirmed, ExitSessionConfirmed,
/// Open an arbitrary URL in the system browser (with scheme validation). /// Open an arbitrary URL in the system browser (with scheme validation).
OpenUrl(String), OpenUrl(String),
/// Open grok.com managed connectors, appending session teamId when set.
OpenManagedConnectors,
/// Cycle to the next visible link (or highlight the first if none selected). /// Cycle to the next visible link (or highlight the first if none selected).
OpenNextLink, OpenNextLink,
/// Cycle to the previous visible link. /// Cycle to the previous visible link.
+3 -3
View File
@@ -100,7 +100,7 @@ impl QueuedPrompt {
/// `true` for plain rows (no `wire_blocks`) and for raw skill slash rows /// `true` for plain rows (no `wire_blocks`) and for raw skill slash rows
/// (`/find-session args` — a single Text block equal to `text`, expanded /// (`/find-session args` — a single Text block equal to `text`, expanded
/// shell-side at delivery), so interjecting `text` loses nothing. `false` /// shell-side at delivery), so interjecting `text` loses nothing. `false`
/// when the payload was expanded client-side (`/imagine`, `/loop`): /// when the payload was expanded client-side (`/loop`):
/// interjecting those by `text` would drop the expansion, and by payload /// interjecting those by `text` would drop the expansion, and by payload
/// would render the raw instruction. /// would render the raw instruction.
pub fn wire_matches_display(&self) -> bool { pub fn wire_matches_display(&self) -> bool {
@@ -1272,7 +1272,7 @@ mod tests {
} }
/// `wire_matches_display` splits interjectable rows (no payload, or a raw /// `wire_matches_display` splits interjectable rows (no payload, or a raw
/// skill slash payload equal to the display text) from client-expanded /// skill slash payload equal to the display text) from client-expanded
/// payloads (`/imagine`, `/loop`) that must run as their own turn. /// payloads (`/loop`) that must run as their own turn.
#[test] #[test]
fn wire_matches_display_classifies_payload_shapes() { fn wire_matches_display_classifies_payload_shapes() {
let text_block = |t: &str| acp::ContentBlock::Text(acp::TextContent::new(t.to_string())); let text_block = |t: &str| acp::ContentBlock::Text(acp::TextContent::new(t.to_string()));
@@ -1285,7 +1285,7 @@ mod tests {
assert!(raw_skill.wire_matches_display(), "raw slash payload"); assert!(raw_skill.wire_matches_display(), "raw slash payload");
let expanded = QueuedPrompt { let expanded = QueuedPrompt {
wire_blocks: Some(vec![text_block("<skill>body</skill>")]), wire_blocks: Some(vec![text_block("<skill>body</skill>")]),
..QueuedPrompt::plain(3, "/imagine cat", QueueEntryKind::Prompt) ..QueuedPrompt::plain(3, "/loop cat", QueueEntryKind::Prompt)
}; };
assert!(!expanded.wire_matches_display(), "expanded payload"); assert!(!expanded.wire_matches_display(), "expanded payload");
let multi_block = QueuedPrompt { let multi_block = QueuedPrompt {
@@ -7,7 +7,7 @@ use super::test_fixtures;
use crate::app::actions::Action; use crate::app::actions::Action;
use crate::app::app_view::InputOutcome; use crate::app::app_view::InputOutcome;
use crate::views::file_search::line_viewer::LineViewerState; use crate::views::file_search::line_viewer::LineViewerState;
use crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; use crossterm::event::{KeyCode, KeyEvent};
impl AgentView { impl AgentView {
// -- Agents modal input handling -- // -- Agents modal input handling --
@@ -227,19 +227,6 @@ impl AgentView {
return self.handle_modal_input_key(key); return self.handle_modal_input_key(key);
} }
// Ctrl+O opens grok.com connectors on the MCP servers tab.
if key.code == KeyCode::Char('o')
&& key.modifiers == KeyModifiers::CONTROL
&& self.extensions_modal.as_ref().is_some_and(|s| {
s.active_tab == crate::views::extensions_modal::ExtensionsTab::McpServers
&& !s.picker_state.search_active
})
{
return self.execute_modal_button_action(
crate::views::extensions_modal::ButtonAction::OpenManagedConnectors,
);
}
// Route chrome keys through ModalWindow first (mirrors the mouse path). // Route chrome keys through ModalWindow first (mirrors the mouse path).
// Handles Esc -> CloseRequested and h/l (or L/R when not tabs-focused) // Handles Esc -> CloseRequested and h/l (or L/R when not tabs-focused)
// -> fold outcomes when FoldInfo provided. // -> fold outcomes when FoldInfo provided.
@@ -768,22 +755,6 @@ impl AgentView {
&config, &config,
); );
// Open the connectors URL on mouse-down (parity with Ctrl+O). A section-row
// click routes as Selected or NonSelectableClick, so intercept both here.
let clicked_entry = match &outcome {
crate::views::picker::PickerOutcome::Selected(i)
| crate::views::picker::PickerOutcome::Expand(i)
| crate::views::picker::PickerOutcome::NonSelectableClick(i) => Some(*i),
_ => None,
};
if let Some(idx) = clicked_entry
&& self.extensions_modal_click_opens_connectors(idx, mouse.row)
{
return self.execute_modal_button_action(
crate::views::extensions_modal::ButtonAction::OpenManagedConnectors,
);
}
// Hover states are managed by ModalWindow (close) and picker (filter). // Hover states are managed by ModalWindow (close) and picker (filter).
match outcome { match outcome {
@@ -864,20 +835,6 @@ impl AgentView {
state.picker_state.scroll_offset = None; state.picker_state.scroll_offset = None;
} }
/// Whether a click at `mouse_row` on entry `entry_idx` hit the connectors URL
/// link band recorded at last paint (opens the URL instead of folding).
fn extensions_modal_click_opens_connectors(&self, entry_idx: usize, mouse_row: u16) -> bool {
self.extensions_modal.as_ref().is_some_and(|state| {
// Parity with the Ctrl+O guard: don't open while the search bar has focus.
!state.picker_state.search_active
&& state
.picker_state
.link_band
.as_ref()
.is_some_and(|(idx, band)| *idx == entry_idx && band.contains(&mouse_row))
})
}
/// Non-selectable mask for the extensions modal picker (from last render). /// Non-selectable mask for the extensions modal picker (from last render).
fn extensions_modal_non_selectable_mask( fn extensions_modal_non_selectable_mask(
state: &crate::views::extensions_modal::ExtensionsModalState, state: &crate::views::extensions_modal::ExtensionsModalState,
@@ -1094,9 +1051,6 @@ impl AgentView {
&& let Some(idx) = state.selected_data_index() && let Some(idx) = state.selected_data_index()
&& let Some(server) = servers.get(idx) && let Some(server) = servers.get(idx)
{ {
if server.is_managed_gateway {
return InputOutcome::Action(Action::OpenManagedConnectors);
}
// Drop repeats while an action is in flight on the same // Drop repeats while an action is in flight on the same
// row to avoid double-spawning the OAuth browser flow. // row to avoid double-spawning the OAuth browser flow.
let sel = state.picker_state.selected; let sel = state.picker_state.selected;
@@ -1120,9 +1074,6 @@ impl AgentView {
InputOutcome::Action(Action::ReloadSkills) InputOutcome::Action(Action::ReloadSkills)
} }
ButtonAction::RefreshMcpList => InputOutcome::Action(Action::RefreshMcpList), ButtonAction::RefreshMcpList => InputOutcome::Action(Action::RefreshMcpList),
ButtonAction::OpenManagedConnectors => {
InputOutcome::Action(Action::OpenManagedConnectors)
}
ButtonAction::ToggleSelectedMcpServer => { ButtonAction::ToggleSelectedMcpServer => {
if let Some(ref mut state) = self.extensions_modal { if let Some(ref mut state) = self.extensions_modal {
use crate::views::extensions_modal::TabDataState; use crate::views::extensions_modal::TabDataState;
@@ -1179,29 +1130,15 @@ impl AgentView {
ButtonAction::RemoveSelectedMcpServer => { ButtonAction::RemoveSelectedMcpServer => {
let resolved = self.extensions_modal.as_ref().and_then(|state| { let resolved = self.extensions_modal.as_ref().and_then(|state| {
use crate::views::extensions_modal::TabDataState; use crate::views::extensions_modal::TabDataState;
use crate::views::mcps_modal::is_removable;
let TabDataState::Loaded(ref servers) = state.mcps_data else { let TabDataState::Loaded(ref servers) = state.mcps_data else {
return None; return None;
}; };
let idx = state.selected_data_index()?; let idx = state.selected_data_index()?;
let server = servers.get(idx)?; let server = servers.get(idx)?;
if is_removable(server) { Some(server.name.clone())
Some(Ok(server.name.clone()))
} else {
Some(Err(server.name.clone()))
}
}); });
match resolved { match resolved {
Some(Err(name)) => { Some(server_name) => {
if let Some(ref mut s) = self.extensions_modal {
s.modal_message =
Some(crate::views::extensions_modal::ModalMessage::Error(
format!("Cannot remove managed server '{name}'"),
));
}
InputOutcome::Changed
}
Some(Ok(server_name)) => {
if let Some(ref mut s) = self.extensions_modal { if let Some(ref mut s) = self.extensions_modal {
s.pending_action = Some("removing...".into()); s.pending_action = Some("removing...".into());
s.pending_entry_index = Some(s.picker_state.selected); s.pending_entry_index = Some(s.picker_state.selected);
@@ -1629,148 +1566,3 @@ mod extensions_modal_search_key_tests {
assert!(state.picker_state.search_active); assert!(state.picker_state.search_active);
} }
} }
#[cfg(test)]
mod connectors_url_click_tests {
use super::AgentView;
use crate::app::actions::Action;
use crate::app::app_view::InputOutcome;
use crate::views::extensions_modal::{
ExtensionsModalState, ExtensionsTab, TabDataState, render_extensions_modal,
};
use crate::views::mcps_modal::{McpServerDisplayStatus, McpServerInfo, McpWireSource};
use crossterm::event::{KeyModifiers, MouseButton, MouseEvent, MouseEventKind};
use ratatui::buffer::Buffer;
use ratatui::layout::Rect;
fn managed_server() -> McpServerInfo {
McpServerInfo {
name: "grok_com_linear".into(),
display_name: None,
status: McpServerDisplayStatus::Ready,
tool_count: 0,
auth_required: false,
tools: vec![],
enabled: true,
source: "managed".into(),
wire_source: McpWireSource::Managed,
plugin_name: None,
is_managed_gateway: false,
}
}
// Build an agent whose extensions modal shows an expanded Managed section,
// then paint it so `hit_areas` + `link_band` reflect the real layout.
fn rendered_agent() -> AgentView {
let mut agent = super::test_fixtures::make_agent();
let mut state = ExtensionsModalState::new(ExtensionsTab::McpServers);
state.mcps_data = TabDataState::Loaded(vec![managed_server()]);
agent.extensions_modal = Some(state);
let area = Rect::new(0, 0, 100, 40);
let mut buf = Buffer::empty(area);
render_extensions_modal(
&mut buf,
area,
agent.extensions_modal.as_mut().unwrap(),
None,
false,
0,
);
agent
}
fn left_down(column: u16, row: u16) -> MouseEvent {
MouseEvent {
kind: MouseEventKind::Down(MouseButton::Left),
column,
row,
modifiers: KeyModifiers::NONE,
}
}
// (column inside the Managed row, its recorded URL band) from the last paint.
fn managed_url_hit(agent: &AgentView) -> (u16, std::ops::Range<u16>) {
let state = agent.extensions_modal.as_ref().unwrap();
let (entry_idx, band) = state
.picker_state
.link_band
.clone()
.expect("expanded Managed section records a connectors URL band");
let hit = state.picker_state.hit_areas.as_ref().unwrap();
let pos = hit
.entry_indices
.iter()
.position(|&e| e == entry_idx)
.unwrap();
(hit.item_rects[pos].x + 2, band)
}
#[test]
fn mouse_down_on_url_row_opens_connectors() {
let mut agent = rendered_agent();
let (col, band) = managed_url_hit(&agent);
let outcome = agent.handle_extensions_modal_mouse(&left_down(col, band.start));
assert!(matches!(
outcome,
InputOutcome::Action(Action::OpenManagedConnectors)
));
// The section stays expanded (opened, did not fold).
assert!(
!agent
.extensions_modal
.as_ref()
.unwrap()
.mcps_collapsed_sections
.contains("mcp-section:managed")
);
}
#[test]
fn mouse_down_on_label_row_folds_not_opens() {
let mut agent = rendered_agent();
let (col, label_row) = {
let state = agent.extensions_modal.as_ref().unwrap();
let (entry_idx, _band) = state.picker_state.link_band.clone().unwrap();
let hit = state.picker_state.hit_areas.as_ref().unwrap();
let pos = hit
.entry_indices
.iter()
.position(|&e| e == entry_idx)
.unwrap();
let rect = hit.item_rects[pos];
(rect.x + 2, rect.y) // first row of the item rect is the fold-toggle label
};
let outcome = agent.handle_extensions_modal_mouse(&left_down(col, label_row));
assert!(!matches!(
outcome,
InputOutcome::Action(Action::OpenManagedConnectors)
));
// Fold happened: the Managed section is now collapsed.
assert!(
agent
.extensions_modal
.as_ref()
.unwrap()
.mcps_collapsed_sections
.contains("mcp-section:managed")
);
}
#[test]
fn mouse_down_on_url_row_while_searching_does_not_open() {
// Parity with the Ctrl+O guard: opening is suppressed while search is active.
let mut agent = rendered_agent();
let (col, band) = managed_url_hit(&agent);
agent
.extensions_modal
.as_mut()
.unwrap()
.picker_state
.search_active = true;
let outcome = agent.handle_extensions_modal_mouse(&left_down(col, band.start));
assert!(!matches!(
outcome,
InputOutcome::Action(Action::OpenManagedConnectors)
));
}
}
@@ -498,7 +498,7 @@ impl AgentView {
/// can't be resolved. Prompt-like rows may interject: plain prompts, plus /// can't be resolved. Prompt-like rows may interject: plain prompts, plus
/// raw skill slash rows (`/find-session args`) whose wire payload IS the /// raw skill slash rows (`/find-session args`) whose wire payload IS the
/// display text — the shell expands those at the interjection drain. Rows /// display text — the shell expands those at the interjection drain. Rows
/// with a client-expanded payload (`/imagine`, `/loop`) and non-prompt /// with a client-expanded payload (`/loop`) and non-prompt
/// kinds stay queued: interjecting them would send the display text, not /// kinds stay queued: interjecting them would send the display text, not
/// the payload. /// the payload.
pub(in crate::app) fn queue_row_prompt_like(&self, id: u64) -> Option<bool> { pub(in crate::app) fn queue_row_prompt_like(&self, id: u64) -> Option<bool> {
@@ -1309,13 +1309,13 @@ mod queue_edit_routing_tests {
); );
} }
/// A client-expanded row (`/imagine`-shaped: wire payload != display /// A client-expanded row (`/loop`-shaped: wire payload != display
/// text) stays queued — interjecting it would send the display text, /// text) stays queued — interjecting it would send the display text,
/// not the payload. /// not the payload.
#[test] #[test]
fn force_interject_local_expanded_row_keeps_it_queued() { fn force_interject_local_expanded_row_keeps_it_queued() {
let mut agent = let mut agent =
running_agent_with_local_skill("/imagine a cat", "<expanded imagine instructions>"); running_agent_with_local_skill("/loop 5m check", "<expanded loop instructions>");
let registry = non_vscode_registry(); let registry = non_vscode_registry();
let ids = agent.queue.entry_ids(); let ids = agent.queue.entry_ids();
@@ -869,12 +869,6 @@ pub(crate) fn dispatch(action: Action, app: &mut AppView) -> Vec<Effect> {
} }
vec![] vec![]
} }
Action::OpenManagedConnectors => {
use crate::terminal::hyperlinks::SchemeFilter;
let url = crate::views::mcps_modal::managed_connectors_url(None);
crate::app::link_opener::open_url_if_safe(&url, SchemeFilter::Standard);
vec![]
}
Action::OpenNextLink => { Action::OpenNextLink => {
with_active_agent(app, |agent| agent.cycle_highlighted_link(true)); with_active_agent(app, |agent| agent.cycle_highlighted_link(true));
vec![] vec![]
@@ -130,7 +130,7 @@ fn unknown_non_restricted_command_still_passes_through() {
app.agents app.agents
.get_mut(&id) .get_mut(&id)
.unwrap() .unwrap()
.set_restricted_commands(&["imagine".to_string()]); .set_restricted_commands(&["loop".to_string()]);
let effects = dispatch(Action::SendPrompt("/frobnicate arg".into()), &mut app); let effects = dispatch(Action::SendPrompt("/frobnicate arg".into()), &mut app);
@@ -1212,7 +1212,6 @@ impl AgentView {
badge: "", badge: "",
badge_color: None, badge_color: None,
collapsible: false, collapsible: false,
underline_last_desc: false,
fields: &[], fields: &[],
description_lines: &[], description_lines: &[],
summary_lines: &[], summary_lines: &[],
@@ -1644,7 +1643,6 @@ impl AgentView {
badge: "", badge: "",
badge_color: None, badge_color: None,
collapsible: false, collapsible: false,
underline_last_desc: false,
}) })
} }
}) })
@@ -1716,7 +1714,6 @@ impl AgentView {
badge: "", badge: "",
badge_color: None, badge_color: None,
collapsible: false, collapsible: false,
underline_last_desc: false,
}) })
}) })
.collect(); .collect();
@@ -2008,7 +2005,6 @@ impl AgentView {
badge: if has_snippet { "match" } else { "" }, badge: if has_snippet { "match" } else { "" },
badge_color: Some(theme.accent_user), badge_color: Some(theme.accent_user),
collapsible: true, collapsible: true,
underline_last_desc: false,
})); }));
non_sel_flags.push(false); non_sel_flags.push(false);
} }

Some files were not shown because too many files have changed in this diff Show More