From 5e4e24db99b91de8a09434a9730913bb0f66bed6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Fri, 17 Jul 2026 23:45:05 -0400 Subject: [PATCH] M2 audit: excise managed connectors and xAI media-gen tools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- Cargo.lock | 22 - Cargo.toml | 2 - crates/codegen/kigi-agent/src/builder.rs | 58 +- crates/codegen/kigi-agent/src/config.rs | 8 +- .../kigi-agent/src/plugins/hooks_adapter.rs | 2 +- crates/codegen/kigi-config-types/src/lib.rs | 36 - crates/codegen/kigi-mcp/src/servers.rs | 20 +- .../codegen/kigi-pager-minimal/src/panel.rs | 5 +- .../kigi-pager-render/src/gboom/mod.rs | 2 +- .../kigi-pager-render/src/link_opener.rs | 73 +- .../kigi-pager-render/src/prompt_images.rs | 2 +- .../kigi-pager-render/src/render/osc8.rs | 6 +- .../kigi-shared/src/placeholder_images.rs | 94 +- .../kigi-shell-session-support/Cargo.toml | 32 - .../kigi-shell-session-support/src/lib.rs | 11 - .../src/managed_mcp.rs | 1467 ---------------- crates/codegen/kigi-shell/Cargo.toml | 1 - .../kigi-shell/skills/imagine/SKILL.md | 131 -- crates/codegen/kigi-shell/src/agent/app.rs | 17 - crates/codegen/kigi-shell/src/agent/config.rs | 232 --- .../src/agent/mvp_agent/acp_agent.rs | 20 +- .../src/agent/mvp_agent/agent_ops.rs | 232 +-- .../agent/mvp_agent/folder_trust_prompt.rs | 33 +- .../kigi-shell/src/agent/mvp_agent/mod.rs | 4 - .../agent/mvp_agent/subagent_coordinator.rs | 7 - .../kigi-shell/src/agent/mvp_agent/tests.rs | 58 +- .../src/agent/subagent/handle_request.rs | 5 - .../kigi-shell/src/agent/subagent/mod.rs | 7 - .../src/agent/subagent/tests/mod.rs | 11 - .../kigi-shell/src/auth/attribution.rs | 48 +- crates/codegen/kigi-shell/src/auth/manager.rs | 2 +- crates/codegen/kigi-shell/src/builtin.rs | 13 +- crates/codegen/kigi-shell/src/config/mod.rs | 120 +- crates/codegen/kigi-shell/src/config/tests.rs | 253 +-- .../codegen/kigi-shell/src/extensions/mcp.rs | 761 +------- .../src/extensions/session_admin.rs | 19 +- .../codegen/kigi-shell/src/leader/server.rs | 10 +- crates/codegen/kigi-shell/src/mcp_doctor.rs | 87 +- .../kigi-shell/src/session/acp_conversion.rs | 41 - .../kigi-shell/src/session/acp_session.rs | 612 +------ .../session/acp_session_impl/hooks_plugins.rs | 8 - .../src/session/acp_session_impl/mcp.rs | 277 +-- .../session/acp_session_impl/mcp_snapshot.rs | 121 +- .../session/acp_session_impl/model_switch.rs | 3 - .../session/acp_session_impl/prompt_build.rs | 63 - .../src/session/acp_session_impl/run_loop.rs | 39 +- .../src/session/acp_session_impl/spawn.rs | 37 - .../session/acp_session_impl/tool_calls.rs | 53 +- .../session/acp_session_impl/tool_dispatch.rs | 2 +- .../src/session/acp_session_impl/turn.rs | 10 - .../cancel_running_task_tests.rs | 8 - .../fs_injection_regression_tests.rs | 2 - .../acp_session_tests/idle_resume_tests.rs | 2 - .../inline_auto_compact_flow_tests.rs | 6 - .../acp_session_tests/memory_config_tests.rs | 2 - .../reactive_managed_reauth_e2e_tests.rs | 476 ----- .../reactive_managed_reauth_tests.rs | 97 -- .../replay_buffer_send_update_tests.rs | 2 - .../src/session/acp_session_tests/support.rs | 4 - ...etry_tests.rs => tool_auth_retry_tests.rs} | 18 +- .../acp_session_tests/web_search_e2e_tests.rs | 2 - .../kigi-shell/src/session/agent_rebuild.rs | 16 - .../kigi-shell/src/session/commands.rs | 7 +- .../kigi-shell/src/session/compaction.rs | 2 - .../codegen/kigi-shell/src/session/handle.rs | 32 - .../kigi-shell/src/session/managed_mcp.rs | 278 +-- .../kigi-shell/src/session/mcp_descriptors.rs | 354 ---- .../kigi-shell/src/session/mcp_dispatcher.rs | 132 +- .../kigi-shell/src/session/mcp_servers.rs | 9 +- .../kigi-shell/src/session/slash_commands.rs | 4 +- .../src/test_support/lsp_runtime.rs | 4 - .../codegen/kigi-shell/src/util/config/mcp.rs | 4 +- .../kigi-tools-api/src/slash_commands.rs | 106 +- .../kigi-tools/schema/tool_meta.schema.json | 2 +- crates/codegen/kigi-tools/src/attribution.rs | 9 - .../grok_build/image_edit/mod.rs | 704 -------- .../grok_build/image_gen/mod.rs | 585 ------- .../src/implementations/grok_build/mod.rs | 13 - .../implementations/grok_build/task/types.rs | 8 - .../grok_build/video_gen/mod.rs | 1536 ----------------- .../src/implementations/use_tool/mod.rs | 548 +----- .../codegen/kigi-tools/src/normalization.rs | 4 - .../codegen/kigi-tools/src/registry/types.rs | 52 +- .../src/reminders/task_completion.rs | 4 - .../codegen/kigi-tools/src/tool_taxonomy.rs | 8 - crates/codegen/kigi-tools/src/types/output.rs | 164 -- .../codegen/kigi-tools/src/types/resources.rs | 63 - crates/codegen/kigi-tools/src/types/tool.rs | 4 - .../codegen/kigi-tools/src/types/tool_io.rs | 7 - .../docs/user-guide/04-slash-commands.md | 20 - crates/codegen/kigi-tui/src/acp/tracker.rs | 156 +- .../kigi-tui/src/app/acp_handler/tests/mcp.rs | 4 - .../kigi-tui/src/app/acp_handler/tests/mod.rs | 5 +- crates/codegen/kigi-tui/src/app/actions.rs | 2 - crates/codegen/kigi-tui/src/app/agent.rs | 6 +- .../kigi-tui/src/app/agent_view/modals.rs | 214 +-- .../kigi-tui/src/app/agent_view/queue.rs | 6 +- .../kigi-tui/src/app/dispatch/router.rs | 6 - .../src/app/dispatch/tests/billing.rs | 2 +- crates/codegen/kigi-tui/src/app/modals.rs | 4 - .../src/scrollback/blocks/tool/other.rs | 2 +- .../codegen/kigi-tui/src/scrollback/render.rs | 6 +- .../kigi-tui/src/scrollback/state/layout.rs | 2 +- .../kigi-tui/src/slash/commands/imagine.rs | 108 -- .../src/slash/commands/imagine_video.rs | 116 -- .../kigi-tui/src/slash/commands/mod.rs | 4 - crates/codegen/kigi-tui/src/slash/mod.rs | 8 +- .../kigi-tui/src/views/dashboard/render.rs | 1 - .../kigi-tui/src/views/extensions_modal.rs | 116 +- .../codegen/kigi-tui/src/views/mcps_modal.rs | 233 +-- crates/codegen/kigi-tui/src/views/modal.rs | 1 - crates/codegen/kigi-tui/src/views/picker.rs | 121 +- .../kigi-tui/src/views/session_picker.rs | 1 - .../kigi-tui/src/views/shortcuts_help.rs | 3 - .../codegen/kigi-tui/src/views/welcome/mod.rs | 2 - .../kigi-workspace-types/src/rpc/fs.rs | 2 +- .../codegen/kigi-workspace/src/capability.rs | 7 +- .../kigi-workspace/src/session/tool_config.rs | 6 - .../tests/serde_roundtrip.rs | 6 +- crates/common/kigi-tool-runtime/src/render.rs | 3 - 120 files changed, 301 insertions(+), 11327 deletions(-) delete mode 100644 crates/codegen/kigi-shell-session-support/Cargo.toml delete mode 100644 crates/codegen/kigi-shell-session-support/src/lib.rs delete mode 100644 crates/codegen/kigi-shell-session-support/src/managed_mcp.rs delete mode 100644 crates/codegen/kigi-shell/skills/imagine/SKILL.md delete mode 100644 crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_e2e_tests.rs delete mode 100644 crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_tests.rs rename crates/codegen/kigi-shell/src/session/acp_session_tests/{media_gen_auth_retry_tests.rs => tool_auth_retry_tests.rs} (93%) delete mode 100644 crates/codegen/kigi-tools/src/implementations/grok_build/image_edit/mod.rs delete mode 100644 crates/codegen/kigi-tools/src/implementations/grok_build/image_gen/mod.rs delete mode 100644 crates/codegen/kigi-tools/src/implementations/grok_build/video_gen/mod.rs delete mode 100644 crates/codegen/kigi-tui/src/slash/commands/imagine.rs delete mode 100644 crates/codegen/kigi-tui/src/slash/commands/imagine_video.rs diff --git a/Cargo.lock b/Cargo.lock index 9a6710f..f86df6e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6293,7 +6293,6 @@ dependencies = [ "kigi-sandbox", "kigi-shared", "kigi-shell-base", - "kigi-shell-session-support", "kigi-sqlite-journal", "kigi-subagent-resolution", "kigi-system-power", @@ -6389,27 +6388,6 @@ dependencies = [ "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]] name = "kigi-sqlite-journal" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 0d44d4c..31d20ed 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -44,7 +44,6 @@ members = [ "crates/codegen/kigi-shared", "crates/codegen/kigi-shell", "crates/codegen/kigi-shell-base", - "crates/codegen/kigi-shell-session-support", "crates/codegen/kigi-subagent-resolution", "crates/codegen/kigi-test-support", "crates/codegen/kigi-tools", @@ -274,7 +273,6 @@ kigi-secrets = { path = "crates/codegen/kigi-secrets" } kigi-shared = { path = "crates/codegen/kigi-shared" } kigi-shell = { path = "crates/codegen/kigi-shell" } 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-test-support = { path = "crates/codegen/kigi-test-support" } kigi-tools = { path = "crates/codegen/kigi-tools" } diff --git a/crates/codegen/kigi-agent/src/builder.rs b/crates/codegen/kigi-agent/src/builder.rs index 29195b5..40109fa 100644 --- a/crates/codegen/kigi-agent/src/builder.rs +++ b/crates/codegen/kigi-agent/src/builder.rs @@ -93,8 +93,6 @@ pub struct AgentBuilder { backend_search: bool, web_fetch_config: kigi_tools::implementations::grok_build::web_fetch::WebFetchConfig, lsp: Option>, - 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, write_file_enabled: bool, @@ -218,8 +216,6 @@ impl AgentBuilder { backend_search: false, web_fetch_config: Default::default(), lsp: None, - image_gen_config: Default::default(), - video_gen_config: Default::default(), app_builder_deployer_config: Default::default(), write_file_enabled: true, subagents_enabled: false, @@ -451,33 +447,6 @@ impl AgentBuilder { self.lsp = Some(handle); 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. pub fn with_app_builder_deployer_config( mut self, @@ -495,10 +464,9 @@ impl AgentBuilder { self } /// Set the 401-attribution callback for tool HTTP clients - /// (`image_gen`, `video_gen`, `web_search`). When set, a 401 - /// from any of those tools emits an `auth_401_attribution` - /// event with `consumer` of `"ImageGen"` / `"VideoGen.start"` / - /// `"VideoGen.poll"` / `"WebSearch"`. Callers should pass the + /// (`web_search`). When set, a 401 from such a tool emits an + /// `auth_401_attribution` event with `consumer` of + /// `"WebSearch"`. Callers should pass the /// same `ShellAttribution` instance they wire into /// `kigi_sampler::SamplerConfig::attribution_callback` so /// all 401s share the same `AuthManager` reference and land in @@ -712,24 +680,6 @@ impl AgentBuilder { .tools .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 .tools .iter() @@ -1024,8 +974,6 @@ impl AgentBuilder { web_search_config: self.web_search_config, web_fetch_config: self.web_fetch_config, 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, api_key_provider: self.api_key_provider, attribution_callback: self.attribution_callback, diff --git a/crates/codegen/kigi-agent/src/config.rs b/crates/codegen/kigi-agent/src/config.rs index 44ce8b5..2d64c5e 100644 --- a/crates/codegen/kigi-agent/src/config.rs +++ b/crates/codegen/kigi-agent/src/config.rs @@ -180,9 +180,6 @@ pub fn workspace_grok_build_toolset() -> ToolServerConfig { tools.push((&grok_build::ExitPlanModeTool).into()); tools.push((&grok_build::AskUserQuestionTool).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((&memory::search_tool::MemorySearchImpl).into()); tools.push((&memory::get_tool::MemoryGetImpl).into()); @@ -447,9 +444,6 @@ fn orchestrator_toolset() -> ToolServerConfig { (&grok_build::MonitorTool).into(), (&grok_build::WebSearchTool).into(), (&grok_build::WebFetchTool).into(), - (&grok_build::ImageGenTool).into(), - (&grok_build::ImageToVideoTool).into(), - (&grok_build::ReferenceToVideoTool).into(), (&memory::MemorySearchImpl).into(), (&memory::MemoryGetImpl).into(), ], @@ -744,7 +738,7 @@ pub struct AgentDefinition { pub agents_md: bool, /// When true (the default), the AgentBuilder layers session-level optional /// 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. /// /// Set this to `false` for harnesses that need an exact, minimal toolset diff --git a/crates/codegen/kigi-agent/src/plugins/hooks_adapter.rs b/crates/codegen/kigi-agent/src/plugins/hooks_adapter.rs index 53e55d9..3606301 100644 --- a/crates/codegen/kigi-agent/src/plugins/hooks_adapter.rs +++ b/crates/codegen/kigi-agent/src/plugins/hooks_adapter.rs @@ -171,7 +171,7 @@ fn process_hooks_content( // Substitute plugin env vars in command paths at config-load time so // that hooks like `${CLAUDE_PLUGIN_ROOT}/hooks/foo.sh` resolve to the // 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 { let cmd_str = cmd.to_string_lossy(); // Mirror what `managed_mcp::load_plugin_mcp_servers_from_config` diff --git a/crates/codegen/kigi-config-types/src/lib.rs b/crates/codegen/kigi-config-types/src/lib.rs index fd5f922..f6227a8 100644 --- a/crates/codegen/kigi-config-types/src/lib.rs +++ b/crates/codegen/kigi-config-types/src/lib.rs @@ -445,11 +445,6 @@ pub struct RemoteSettings { skip_serializing_if = "Vec::is_empty" )] pub goal_skeptic_models: Vec, - /// Remote fallback for managed MCP connector fetching. - #[serde(default)] - pub managed_mcps_enabled: Option, - #[serde(default)] - pub managed_mcp_gateway_tools_enabled: Option, /// Enable user-facing feedback (heuristic popups, `/feedback` command). #[serde(default)] pub feedback_enabled: Option, @@ -517,20 +512,6 @@ pub struct RemoteSettings { /// is set in config.toml. Absent → default (**disabled** — ships dark). #[serde(default)] pub subagent_worktree_snapshot_enabled: Option, - /// 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, - /// 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, - /// 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, /// When `Some(true)`, enable the process-wide image normalize cache that /// amortises decode + integrity-check + re-encode work across SessionActors. /// 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). #[serde(default)] pub compaction_verbatim_input: Option, - /// 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>, /// Master switch for jemalloc heap sampling + threshold dumps. /// `Some(true)` enables, `Some(false)` kill-switch, `None` = client default off. #[serde(default)] @@ -706,16 +680,6 @@ pub struct RemoteSettings { #[serde(default)] pub jemalloc_heap_profile_poll_interval_secs: Option, } -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 /// `[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 diff --git a/crates/codegen/kigi-mcp/src/servers.rs b/crates/codegen/kigi-mcp/src/servers.rs index 8b4c6a7..13ccde5 100644 --- a/crates/codegen/kigi-mcp/src/servers.rs +++ b/crates/codegen/kigi-mcp/src/servers.rs @@ -46,7 +46,7 @@ pub use kigi_workspace_types::MCP_TOOL_NAME_DELIMITER; /// Normalize an MCP server URL for comparison: strip trailing slashes. /// 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. fn normalize_url(url: &str) -> String { url.trim_end_matches('/').to_string() @@ -112,7 +112,7 @@ pub struct McpConfigDiff { pub retained: Vec, } -/// 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; /// 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") { 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) .ok() .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() { 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://figma.com/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("grok_com_linear", "https://example.com/mcp")); assert!(!is_figma_mcp("linear", "not-a-url")); assert!(!is_figma_mcp("linear", "https://notfigma.com/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() { // The verbatim string captured in production for a managed handshake. 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("unauthorized")); diff --git a/crates/codegen/kigi-pager-minimal/src/panel.rs b/crates/codegen/kigi-pager-minimal/src/panel.rs index 0fe8b46..e61aa1d 100644 --- a/crates/codegen/kigi-pager-minimal/src/panel.rs +++ b/crates/codegen/kigi-pager-minimal/src/panel.rs @@ -461,7 +461,6 @@ fn render_mcps( badge: badges[i].as_str(), badge_color: badge_colors[i], collapsible: collapsibles[i], - underline_last_desc: false, }) }) .collect(); @@ -516,7 +515,7 @@ fn measure_entries(entries: &[PickerEntry<'_>]) -> u16 { mod tests { use super::*; 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; fn agent() -> AgentView { @@ -533,9 +532,7 @@ mod tests { tools: Vec::new(), enabled: true, source: "local".to_string(), - wire_source: McpWireSource::Local, plugin_name: None, - is_managed_gateway: false, } } diff --git a/crates/codegen/kigi-pager-render/src/gboom/mod.rs b/crates/codegen/kigi-pager-render/src/gboom/mod.rs index d1aceb5..992f24e 100644 --- a/crates/codegen/kigi-pager-render/src/gboom/mod.rs +++ b/crates/codegen/kigi-pager-render/src/gboom/mod.rs @@ -2,7 +2,7 @@ //! the terminal via the kitty graphics protocol. //! //! 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 //! simulation steps with wall-clock `dt`, so gameplay speed is independent //! of the achieved frame rate. diff --git a/crates/codegen/kigi-pager-render/src/link_opener.rs b/crates/codegen/kigi-pager-render/src/link_opener.rs index b8d861e..8e49133 100644 --- a/crates/codegen/kigi-pager-render/src/link_opener.rs +++ b/crates/codegen/kigi-pager-render/src/link_opener.rs @@ -92,7 +92,7 @@ fn build_open_path_command(path: &std::path::Path) -> std::process::Command { /// - **Windows**: `explorer.exe /select,` reveals + highlights the file /// in Explorer. We deliberately avoid `cmd /c start`, whose `%VAR%` /// 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. pub fn open_path(path: &std::path::Path) -> bool { // 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)] mod tests { 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] fn fallback_scheme_case_insensitive() { // Uppercase scheme that url::Url::parse rejects triggers fallback path; diff --git a/crates/codegen/kigi-pager-render/src/prompt_images.rs b/crates/codegen/kigi-pager-render/src/prompt_images.rs index ea35aa2..1b147ae 100644 --- a/crates/codegen/kigi-pager-render/src/prompt_images.rs +++ b/crates/codegen/kigi-pager-render/src/prompt_images.rs @@ -1879,7 +1879,7 @@ impl ScrollbackVideoRef { pub fn extract_video_refs(text: &str) -> Vec { 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 = LazyLock::new(|| regex::Regex::new(MARKDOWN_IMAGE_REF_PATTERN).unwrap()); diff --git a/crates/codegen/kigi-pager-render/src/render/osc8.rs b/crates/codegen/kigi-pager-render/src/render/osc8.rs index 443a87e..a0a2fe6 100644 --- a/crates/codegen/kigi-pager-render/src/render/osc8.rs +++ b/crates/codegen/kigi-pager-render/src/render/osc8.rs @@ -254,7 +254,7 @@ struct RowSegment { /// joiner to the *previous* row (see `BlockLine::joiner`): `None` = hard /// break, `Some("")` = mid-word wrap, `Some(" ")` = word wrap. Consecutive /// 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 /// narrow panes) is detected whole and each row's fragment gets its own /// clickable overlay region. Spans within a row are likewise concatenated so @@ -858,7 +858,7 @@ mod tests { #[test] fn scan_detects_grok_session_media_path() { // 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 mut overlay = LinkOverlay::new(); scan_unjoined(std::iter::once((0, &line)), 0, &[], &mut overlay); @@ -873,7 +873,7 @@ mod tests { #[test] 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 // each row was scanned in isolation, so only the `/Users/alice` // fragment on the first row matched and became clickable. diff --git a/crates/codegen/kigi-shared/src/placeholder_images.rs b/crates/codegen/kigi-shared/src/placeholder_images.rs index 2dc5202..32184af 100644 --- a/crates/codegen/kigi-shared/src/placeholder_images.rs +++ b/crates/codegen/kigi-shared/src/placeholder_images.rs @@ -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 /// is recorded on its ACP image block, so the server can resolve /// `[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"; /// 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()) } -/// 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:;base64,` 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. /// /// 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.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()), - ] - ); - } } diff --git a/crates/codegen/kigi-shell-session-support/Cargo.toml b/crates/codegen/kigi-shell-session-support/Cargo.toml deleted file mode 100644 index d92ea7d..0000000 --- a/crates/codegen/kigi-shell-session-support/Cargo.toml +++ /dev/null @@ -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 diff --git a/crates/codegen/kigi-shell-session-support/src/lib.rs b/crates/codegen/kigi-shell-session-support/src/lib.rs deleted file mode 100644 index b50be2e..0000000 --- a/crates/codegen/kigi-shell-session-support/src/lib.rs +++ /dev/null @@ -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; diff --git a/crates/codegen/kigi-shell-session-support/src/managed_mcp.rs b/crates/codegen/kigi-shell-session-support/src/managed_mcp.rs deleted file mode 100644 index fabed36..0000000 --- a/crates/codegen/kigi-shell-session-support/src/managed_mcp.rs +++ /dev/null @@ -1,1467 +0,0 @@ -//! Managed MCP credential resolution via cli-chat-proxy. -//! -//! For remote MCP servers where the user has completed OAuth enrollment, -//! this module resolves credentials at agent init (cached across sessions) -//! and proactively refreshes them before token expiry. -//! -//! Config-file/plugin merge layering (which reads shell's config system) lives -//! in shell's `session::managed_mcp`, which re-exports everything here. - -use std::collections::{HashMap, HashSet}; -use std::sync::Arc; - -use agent_client_protocol as acp; -use chrono::{DateTime, Utc}; -use tokio_util::sync::CancellationToken; - -/// Agent-level cache for managed MCP configs. -/// -/// Explicit tri-state prevents concurrent double-fetches: -/// only the first caller transitions `NotFetched → Fetching`; -/// subsequent callers see `Fetching` and wait for the in-flight fetch. -pub enum ManagedMcpCache { - NotFetched, - /// Fetch in progress — callers should wait on `fetch_notify`. - Fetching, - /// May be empty if no managed servers are configured for this user. - Ready(Vec), -} - -/// Consecutive failed reactive re-auth attempts before a managed server is -/// parked in a terminal needs-auth state. Once reached, the cooldown gate -/// refuses further attempts until a successful proactive fetch clears it. -const MAX_REACTIVE_REAUTH_ATTEMPTS: u32 = 3; - -/// Defensive upper bound (seconds) on the exponential backoff between reactive -/// re-auth attempts. The terminal attempt cap (`2^3 = 8s`) means the reactive -/// path never reaches this ceiling today. -const REACTIVE_REAUTH_BACKOFF_CAP_SECS: u64 = 64; - -/// Per-server cooldown for the reactive managed re-auth path: a genuinely -/// revoked connector keeps returning a bad token, so each failed attempt pushes -/// the next eligible instant out by capped exponential backoff and the server -/// goes terminal after `MAX_REACTIVE_REAUTH_ATTEMPTS`. -#[derive(Debug, Clone)] -struct ManagedReauthState { - consecutive_failures: u32, - next_allowed_at: DateTime, -} - -impl Default for ManagedReauthState { - fn default() -> Self { - // No backoff window yet — the first attempt is always eligible. - Self { - consecutive_failures: 0, - next_allowed_at: DateTime::::MIN_UTC, - } - } -} - -impl ManagedReauthState { - /// Terminal once the attempt cap is hit: no further reactive attempts until - /// a successful fetch clears the entry. - fn is_terminal(&self) -> bool { - self.consecutive_failures >= MAX_REACTIVE_REAUTH_ATTEMPTS - } - - /// Eligible when the backoff window has elapsed and the cap is not reached. - fn is_eligible(&self, now: DateTime) -> bool { - !self.is_terminal() && now >= self.next_allowed_at - } - - /// Bump the failure count and push the next eligible instant out by - /// `min(2^failures, REACTIVE_REAUTH_BACKOFF_CAP_SECS)` seconds. - fn record_failure(&mut self, now: DateTime) { - self.consecutive_failures = self.consecutive_failures.saturating_add(1); - let backoff_secs = 2u64 - .saturating_pow(self.consecutive_failures) - .min(REACTIVE_REAUTH_BACKOFF_CAP_SECS); - self.next_allowed_at = now + chrono::Duration::seconds(backoff_secs as i64); - } -} - -/// Agent-level cache for managed MCP gateway tool catalogs. -pub enum GatewayToolCatalogCache { - NotFetched, - /// Fetch in progress for the recorded gateway tool epoch. - Fetching(u64), - /// May be empty if the user has no gateway-exposed tools. - Ready(GatewayToolCatalog), -} - -enum FetchDecision { - Ready(T), - Wait(tokio::sync::futures::OwnedNotified), - Fetch, -} - -trait ThreeStateFetchCache { - type Value: Clone; - - fn claim_fetch(&mut self, notify: Arc) -> FetchDecision; -} - -impl ThreeStateFetchCache for ManagedMcpCache { - type Value = Vec; - - fn claim_fetch(&mut self, notify: Arc) -> FetchDecision { - match self { - Self::Ready(configs) => FetchDecision::Ready(configs.clone()), - Self::Fetching => FetchDecision::Wait(notify.notified_owned()), - Self::NotFetched => { - *self = Self::Fetching; - FetchDecision::Fetch - } - } - } -} - -pub struct ManagedMcpState { - pub cache: ManagedMcpCache, - pub fetch_notify: Arc, - pub gateway_tools_active: bool, - pub gateway_tool_epoch: u64, - pub gateway_tool_cache: GatewayToolCatalogCache, - pub gateway_tool_fetch_notify: Arc, - /// Retained across gateway disable/cache invalidation so the on-disk - /// MCP descriptor mirror can remove stale gateway connector directories when - /// the current catalog is empty or absent. - pub gateway_tool_connectors_seen: HashSet, - pub refresh_task_spawned: bool, - /// Cancels the background refresh task on drop. - refresh_cancel: CancellationToken, - /// Per-server reactive re-auth cooldown, keyed by MCP server name: one - /// backoff entry per connector. Coalescing of concurrent attempts is - /// best-effort — the caller takes the mutex sequentially for the - /// `reauth_allowed` check and the later `record_reauth_failure`, not across - /// the network attempt in between, so two simultaneously in-flight tool - /// calls can each record a failure and reach the cap in fewer real re-auth - /// rounds. Acceptable because the terminal state is cleared by the next - /// proactive `clear_reauth_cooldowns`. - reauth_cooldown: HashMap, -} - -impl Default for ManagedMcpState { - fn default() -> Self { - Self { - cache: ManagedMcpCache::NotFetched, - fetch_notify: Arc::new(tokio::sync::Notify::new()), - gateway_tools_active: false, - gateway_tool_epoch: 0, - gateway_tool_cache: GatewayToolCatalogCache::NotFetched, - gateway_tool_fetch_notify: Arc::new(tokio::sync::Notify::new()), - gateway_tool_connectors_seen: HashSet::new(), - refresh_task_spawned: false, - refresh_cancel: CancellationToken::new(), - reauth_cooldown: HashMap::new(), - } - } -} - -impl Drop for ManagedMcpState { - fn drop(&mut self) { - self.refresh_cancel.cancel(); - } -} - -async fn wait_for_fetch_slot( - handle: &ManagedMcpStateHandle, - claim: impl Fn(&mut ManagedMcpState) -> FetchDecision, -) -> Option { - loop { - let decision = { - let mut state = handle.lock().await; - claim(&mut state) - }; - match decision { - FetchDecision::Ready(value) => return Some(value), - FetchDecision::Wait(notified) => notified.await, - FetchDecision::Fetch => return None, - } - } -} - -async fn get_authenticated_json( - url: &str, - auth_key: &str, - unavailable_message: &'static str, - fetch_failed_message: &'static str, - parse_error_message: &'static str, -) -> Result { - let resp = match kigi_http::shared_client() - .get(url) - .timeout(std::time::Duration::from_secs(10)) - .header("Authorization", format!("Bearer {}", auth_key)) - .header("X-XAI-Token-Auth", "xai-grok-cli") - .header("x-grok-client-version", kigi_version::VERSION) - .send() - .await - { - Ok(r) if r.status().is_success() => r, - Ok(r) => { - let status = r.status(); - tracing::warn!(status = %status, "{}", unavailable_message); - return Err(ManagedMcpFetchError::Status { - status, - message: format!("HTTP {status}"), - }); - } - Err(e) => { - tracing::warn!(error = %e, "{}", fetch_failed_message); - return Err(e.into()); - } - }; - - match resp.json::().await { - Ok(value) => Ok(value), - Err(e) => { - tracing::debug!(error = %e, "{}", parse_error_message); - Err(e.into()) - } - } -} - -impl ManagedMcpState { - /// Store fetched configs, optionally spawn the proactive refresh task, and - /// wake any concurrent callers that were waiting on this fetch. - pub fn complete_fetch( - &mut self, - configs: Vec, - state_handle: &ManagedMcpStateHandle, - refresh_ctx: Option, - ) { - let should_refresh = configs.iter().any(|c| c.token_expires_at.is_some()); - self.cache = ManagedMcpCache::Ready(configs); - - if should_refresh - && !self.refresh_task_spawned - && let Some(ctx) = refresh_ctx - { - spawn_cache_refresh_task(state_handle.clone(), ctx, self.refresh_cancel.clone()); - self.refresh_task_spawned = true; - } - - self.fetch_notify.notify_waiters(); - } - - /// Record a failed fetch: roll the cache back to `NotFetched` so the next - /// caller retries, and wake concurrent waiters so they don't hang on a - /// fetch that will never complete. - /// - /// Deliberately NOT `Ready(vec![])`: the agent cache has no TTL, so - /// committing a transient failure (expired token at leader startup, proxy - /// blip) as an empty catalog would erase managed connectors for the whole - /// process lifetime — days, for a leader. - pub fn fail_fetch(&mut self) { - self.cache = ManagedMcpCache::NotFetched; - self.fetch_notify.notify_waiters(); - } - - /// True if a reactive re-auth attempt for `server` is permitted at `now`: - /// no prior cooldown entry, or the backoff window elapsed and the terminal - /// attempt cap is not reached. - pub fn reauth_allowed(&self, server: &str, now: DateTime) -> bool { - self.reauth_cooldown - .get(server) - .is_none_or(|state| state.is_eligible(now)) - } - - /// True once `server` exhausted `MAX_REACTIVE_REAUTH_ATTEMPTS` — the - /// terminal needs-auth state that holds until a proactive refresh clears the - /// cooldown or a reactive re-auth succeeds. - pub fn reauth_is_terminal(&self, server: &str) -> bool { - self.reauth_cooldown - .get(server) - .is_some_and(ManagedReauthState::is_terminal) - } - - /// Record a failed reactive re-auth for `server`: bump the failure count and - /// extend the backoff window. - pub fn record_reauth_failure(&mut self, server: &str, now: DateTime) { - self.reauth_cooldown - .entry(server.to_string()) - .or_default() - .record_failure(now); - } - - /// Reset `server`'s cooldown after a successful reactive re-auth. - pub fn record_reauth_success(&mut self, server: &str) { - self.reauth_cooldown.remove(server); - } - - /// Clear every server's reactive re-auth cooldown. Invoked only by the - /// proactive background refresh after a fresh fetch, so a parked (terminal) - /// connector re-authorized on grok.com can retry. The reactive path must NOT - /// trigger this: a still-rejected token would reset its own attempt cap each - /// attempt and loop instead of going terminal. - pub fn clear_reauth_cooldowns(&mut self) { - self.reauth_cooldown.clear(); - } - - pub fn enable_gateway_tools(&mut self) -> u64 { - if !self.gateway_tools_active { - self.gateway_tool_epoch = self.gateway_tool_epoch.wrapping_add(1); - } - self.gateway_tools_active = true; - self.gateway_tool_epoch - } - - pub fn start_gateway_tool_fetch(&mut self) -> Option { - if !self.gateway_tools_active { - return None; - } - self.gateway_tool_cache = GatewayToolCatalogCache::Fetching(self.gateway_tool_epoch); - Some(self.gateway_tool_epoch) - } - - pub fn complete_gateway_tool_fetch(&mut self, epoch: u64, catalog: GatewayToolCatalog) -> bool { - if !self.gateway_tools_active || self.gateway_tool_epoch != epoch { - self.gateway_tool_fetch_notify.notify_waiters(); - return false; - } - self.gateway_tool_connectors_seen - .extend(catalog.tools.iter().map(|tool| tool.connector_id.clone())); - self.gateway_tool_cache = GatewayToolCatalogCache::Ready(catalog); - self.gateway_tool_fetch_notify.notify_waiters(); - true - } - - pub fn fail_gateway_tool_fetch(&mut self, epoch: u64) { - if self.gateway_tools_active - && self.gateway_tool_epoch == epoch - && matches!(self.gateway_tool_cache, GatewayToolCatalogCache::Fetching(fetch_epoch) if fetch_epoch == epoch) - { - self.gateway_tool_cache = GatewayToolCatalogCache::NotFetched; - } - self.gateway_tool_fetch_notify.notify_waiters(); - } - - pub fn disable_gateway_tools(&mut self) { - self.gateway_tools_active = false; - self.gateway_tool_epoch = self.gateway_tool_epoch.wrapping_add(1); - self.gateway_tool_cache = GatewayToolCatalogCache::NotFetched; - self.gateway_tool_fetch_notify.notify_waiters(); - } -} - -pub type ManagedMcpStateHandle = Arc>; - -#[derive(Debug, Clone, serde::Deserialize)] -pub struct ManagedMcpConfig { - /// Human-readable connector name (e.g. "Slack", "Linear"). - #[serde(default)] - pub name: String, - pub endpoint: String, - #[serde(default)] - pub headers: HashMap, - pub token_expires_at: Option>, - #[serde(default)] - pub scope: Option, - #[serde(default)] - pub scope_id: Option, - #[serde(default)] - pub scope_name: Option, -} - -#[derive(Debug, Clone, serde::Deserialize)] -struct McpConfigsResponse { - mcp_servers: Vec, -} - -#[derive(Debug, Clone, serde::Serialize)] -pub struct GatewayToolCallRequest { - pub call_id: String, - pub arguments: serde_json::Value, -} - -#[derive(Debug, Clone, serde::Deserialize)] -pub struct GatewayToolCallResponse { - pub result: serde_json::Value, - #[serde(default)] - pub connectors_needing_reauth: Vec, -} - -#[derive(Debug, Clone, serde::Deserialize)] -pub struct GatewayToolCatalog { - #[serde(default)] - pub tools: Vec, - #[serde(default)] - pub total_tools: u32, - #[serde(default)] - pub connectors_needing_reauth: Vec, -} - -#[derive(Debug, Clone, serde::Deserialize)] -pub struct GatewayTool { - pub connector_id: String, - pub connector_name: String, - pub tool_id: String, - pub tool_name: String, - pub call_id: String, - pub description: String, - pub json_schema: serde_json::Value, -} - -impl GatewayTool { - pub fn qualified_name(&self) -> String { - format!("{}__{}", self.connector_id, self.tool_id) - } -} - -/// Why a managed-MCP config fetch failed. Distinguishes "fetch failed" from -/// the legitimate "fetched, zero connectors configured" (`Ok(vec![])`) so the -/// agent cache never commits a transient failure as a permanent empty catalog. -#[derive(Debug, thiserror::Error)] -pub enum ManagedMcpFetchError { - #[error("HTTP {status}: {message}")] - Status { - status: reqwest::StatusCode, - message: String, - }, - #[error("transport: {0}")] - Transport(#[from] reqwest::Error), - /// No usable auth token at fetch time. - #[error("no auth token available")] - NoAuth, -} - -/// Fetch managed MCP configs from cli-chat-proxy (`GET /v1/mcp/configs`). -/// -/// `Ok(vec![])` means the server answered and the user genuinely has no -/// managed connectors. `Err(_)` means we don't know (HTTP error, transport -/// failure, parse error) — callers must NOT cache the result as empty. -pub async fn fetch_managed_configs( - proxy_base_url: &str, - auth_key: &str, -) -> Result, ManagedMcpFetchError> { - let url = format!("{}/mcp/configs", proxy_base_url); - - let response: McpConfigsResponse = get_authenticated_json( - &url, - auth_key, - "Managed MCP configs unavailable", - "Managed MCP configs fetch failed", - "Managed MCP configs parse error", - ) - .await?; - tracing::info!( - count = response.mcp_servers.len(), - "Fetched managed MCP configs" - ); - Ok(response.mcp_servers) -} - -const GATEWAY_TOOL_CALL_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); - -pub async fn call_gateway_tool( - proxy_base_url: &str, - auth_key: &str, - call_id: &str, - arguments: serde_json::Value, -) -> Result { - let url = format!("{}/mcp/tools/call", proxy_base_url); - let arguments = if arguments.is_null() { - serde_json::json!({}) - } else { - arguments - }; - let request = GatewayToolCallRequest { - call_id: call_id.to_owned(), - arguments, - }; - - let resp = match kigi_http::shared_client() - .post(&url) - .timeout(GATEWAY_TOOL_CALL_TIMEOUT) - .header("Authorization", format!("Bearer {}", auth_key)) - .header("X-XAI-Token-Auth", "xai-grok-cli") - .header("x-grok-client-version", kigi_version::VERSION) - .json(&request) - .send() - .await - { - Ok(r) if r.status().is_success() => r, - Ok(r) => { - let status = r.status(); - let message = gateway_error_message(status, r).await; - tracing::warn!( - call_id = %call_id, - "Managed MCP gateway tool call unavailable: HTTP {status}" - ); - return Err(ManagedMcpFetchError::Status { status, message }); - } - Err(e) => { - tracing::warn!( - call_id = %call_id, - "Managed MCP gateway tool call failed: {}", - e - ); - return Err(e.into()); - } - }; - - match resp.json::().await { - Ok(response) => Ok(response), - Err(e) => { - tracing::debug!( - call_id = %call_id, - "Managed MCP gateway tool call parse error: {}", - e - ); - Err(e.into()) - } - } -} - -async fn gateway_error_message(status: reqwest::StatusCode, response: reqwest::Response) -> String { - let fallback = format!("HTTP {status}"); - let Ok(body) = response.text().await else { - return fallback; - }; - if body.trim().is_empty() { - return fallback; - } - match serde_json::from_str::(&body) { - Ok(value) => value - .get("error") - .and_then(|v| v.as_str()) - .map(str::to_owned) - .unwrap_or(fallback), - Err(_) => fallback, - } -} - -/// Fetch the managed MCP gateway tool catalog from cli-chat-proxy -/// (`GET /v1/mcp/tools/list`). -/// -/// `Ok(catalog)` means the server answered and the catalog contents are -/// authoritative for this fetch, even when empty. `Err(_)` means freshness is -/// unknown and callers must leave any cache retryable rather than committing an -/// empty catalog. -pub async fn fetch_gateway_tool_catalog( - proxy_base_url: &str, - auth_key: &str, -) -> Result { - let url = format!("{}/mcp/tools/list", proxy_base_url); - - let catalog: GatewayToolCatalog = get_authenticated_json( - &url, - auth_key, - "Managed MCP gateway tools unavailable", - "Managed MCP gateway tools fetch failed", - "Managed MCP gateway tools parse error", - ) - .await?; - tracing::info!( - count = catalog.tools.len(), - total_tools = catalog.total_tools, - reauth = catalog.connectors_needing_reauth.len(), - "Fetched managed MCP gateway tool catalog" - ); - Ok(catalog) -} - -/// Invalidate all managed MCP caches so the next caller refetches both legacy -/// managed configs and gateway tools. -pub async fn invalidate_cache(handle: &ManagedMcpStateHandle) { - let mut state = handle.lock().await; - state.cache = ManagedMcpCache::NotFetched; - state.gateway_tool_cache = GatewayToolCatalogCache::NotFetched; -} - -/// Invalidate only the gateway tool catalog so the next gateway-aware caller -/// refetches `/v1/mcp/tools/list`. -pub async fn invalidate_gateway_tool_cache(handle: &ManagedMcpStateHandle) { - let mut state = handle.lock().await; - state.gateway_tool_cache = GatewayToolCatalogCache::NotFetched; -} - -/// Fetch-or-wait: returns cached configs if ready, otherwise fetches once -/// and wakes any concurrent waiters. Callers provide credentials; this -/// function owns the tri-state lifecycle. -/// -/// Only a successful fetch (including a genuine zero-connector response) is -/// committed to the cache. A failed fetch — or a missing auth token — rolls -/// back to `NotFetched` and returns `vec![]` for this caller only, so the -/// next caller retries instead of inheriting a poisoned empty catalog. -/// Woken waiters re-enter the loop, observe `NotFetched`, and become the -/// next fetcher (bounded: each caller performs at most one fetch). -pub async fn get_or_fetch( - handle: &ManagedMcpStateHandle, - proxy_url: &str, - auth_key: Option<&str>, - refresh_ctx: Option, -) -> Vec { - let fetch = wait_for_fetch_slot(handle, |state| { - state.cache.claim_fetch(state.fetch_notify.clone()) - }) - .await; - if let Some(configs) = fetch { - return configs; - } - - let result = match auth_key { - Some(key) => fetch_managed_configs(proxy_url, key).await, - None => Err(ManagedMcpFetchError::NoAuth), - }; - - match result { - Ok(configs) => { - handle - .lock() - .await - .complete_fetch(configs.clone(), handle, refresh_ctx); - configs - } - Err(e) => { - tracing::warn!( - error = %e, - "Managed MCP fetch failed; leaving cache unpopulated for retry" - ); - handle.lock().await.fail_fetch(); - vec![] - } - } -} - -/// Fetch-or-wait for the managed MCP gateway tool catalog. -/// -/// Returns `Some(catalog)` for either a cached catalog or a successful fresh -/// fetch, including a genuine empty catalog. Returns `None` when gateway tools -/// are disabled by the caller, auth is unavailable, or the fetch failed. Failed -/// fetches roll back to `NotFetched`, so a later caller can retry. -pub async fn get_or_fetch_gateway_tool_catalog( - handle: &ManagedMcpStateHandle, - proxy_url: &str, - auth_key: Option<&str>, -) -> Option { - let fetch_epoch = loop { - let maybe_notify = { - let mut state = handle.lock().await; - if !state.gateway_tools_active { - return None; - } - match &state.gateway_tool_cache { - GatewayToolCatalogCache::Ready(catalog) => return Some(catalog.clone()), - GatewayToolCatalogCache::Fetching(_) => { - Some(state.gateway_tool_fetch_notify.clone().notified_owned()) - } - GatewayToolCatalogCache::NotFetched => { - let epoch = state.start_gateway_tool_fetch()?; - break epoch; - } - } - }; - - if let Some(notified) = maybe_notify { - notified.await; - continue; - } - }; - - let result = match auth_key { - Some(key) => fetch_gateway_tool_catalog(proxy_url, key).await, - None => Err(ManagedMcpFetchError::NoAuth), - }; - - match result { - Ok(catalog) => { - let committed = handle - .lock() - .await - .complete_gateway_tool_fetch(fetch_epoch, catalog.clone()); - committed.then_some(catalog) - } - Err(e) => { - tracing::warn!( - error = %e, - "Managed MCP gateway tool fetch failed; leaving cache unpopulated for retry" - ); - handle.lock().await.fail_gateway_tool_fetch(fetch_epoch); - None - } - } -} - -/// Namespace prefix for managed MCP servers. -/// -/// Servers with names starting with this prefix are managed by grok.com — -/// their OAuth credentials are stored server-side. -/// Servers without this prefix are user-managed (local keychain, config.toml headers, etc.). -/// -/// Examples: -/// `grok_com_linear` → managed by grok.com -/// `grok_com_slack` → managed by grok.com -/// `my_company_api` → user-managed (local) -/// -/// Single source of truth lives in `kigi-workspace` (which matches policy -/// `serverName`s against it); re-exported here so the two never drift. -pub use kigi_workspace::permission::resolution::MANAGED_MCP_PREFIX; - -/// `"Linear"` -> `"grok_com_linear"`. Shares normalization and the -/// `MANAGED_MCP_NAME_MAX_CHARS` cap with policy matching (`mcp_name_matches`) so -/// the runtime name and a policy `serverName` never drift. -pub fn to_managed_name(display_name: &str) -> String { - use kigi_workspace::permission::resolution::{ - MANAGED_MCP_NAME_MAX_CHARS, normalize_managed_name, - }; - let raw = format!( - "{MANAGED_MCP_PREFIX}{}", - normalize_managed_name(display_name) - ); - kigi_shell_base::util::truncate(&raw, MANAGED_MCP_NAME_MAX_CHARS).to_string() -} - -/// Minutes before token expiry to refresh credentials. -pub const TOKEN_EXPIRY_BUFFER_MINUTES: i64 = 5; - -/// Whether a managed token should be treated as stale (eligible for a swap). -/// -/// A `None` expiry carries no TTL to reason about, so we conservatively treat -/// it as stale; otherwise a tokenless connector would never become eligible for -/// a swap. The downstream rebuild remains gated on actual header changes. -pub fn managed_token_is_stale( - expires_at: Option>, - now: chrono::DateTime, -) -> bool { - match expires_at { - Some(exp) => now > exp - chrono::Duration::minutes(TOKEN_EXPIRY_BUFFER_MINUTES), - None => true, - } -} - -/// Returns `true` if this server should use server-side managed credentials. -/// -/// Both conditions must hold: -/// 1. Server name starts with `MANAGED_MCP_PREFIX` ("grok_com_") -/// 2. Server URL matches a managed config endpoint -/// -/// This prevents false injection if a user accidentally names a server `grok_com_*` -/// but it's not actually in the catalog, and prevents injecting into servers -/// that happen to share a URL but aren't opted in to managed auth. -pub fn should_inject_managed_auth( - server_name: &str, - server_url: &str, - managed_by_url: &HashMap, -) -> bool { - server_name.starts_with(MANAGED_MCP_PREFIX) - && managed_by_url.contains_key(&normalize_url(server_url)) -} - -pub fn normalize_url(url: &str) -> String { - url.trim_end_matches('/').to_string() -} - -/// Key for managed config lookup: (normalized_url, scope, scope_id). -type ManagedConfigKey = (String, Option, Option); - -/// Inject managed OAuth headers into `grok_com_`-prefixed MCP servers. -/// -/// Matches on endpoint URL + scope from `X-Connector-Scope` headers. -/// Falls back to URL-only match when no scope headers are present (backward compat). -/// Existing headers are preserved; managed headers are appended. -/// Non-prefixed servers are left untouched. -pub fn inject_managed_headers(servers: &mut [acp::McpServer], managed: &[ManagedMcpConfig]) { - tracing::debug!( - servers = servers.len(), - managed = managed.len(), - "Injecting managed MCP credentials" - ); - if managed.is_empty() { - return; - } - - let managed_by_key: HashMap = managed - .iter() - .map(|m| { - let key = ( - normalize_url(&m.endpoint), - m.scope.clone(), - m.scope_id.clone(), - ); - (key, m) - }) - .collect(); - - // URL-only fallback for backward compat - let managed_by_url: HashMap = managed - .iter() - .map(|m| (normalize_url(&m.endpoint), m)) - .collect(); - - let mut injected = 0usize; - let mut skipped_no_prefix = 0usize; - let mut skipped_no_match = 0usize; - let mut skipped_no_headers = 0usize; - - for server in servers.iter_mut() { - let (name, url, headers) = match server { - acp::McpServer::Http(acp::McpServerHttp { - name, url, headers, .. - }) - | acp::McpServer::Sse(acp::McpServerSse { - name, url, headers, .. - }) => (name.as_str(), url.as_str(), headers), - _ => continue, - }; - - let normalized_url = normalize_url(url); - - if !name.starts_with(MANAGED_MCP_PREFIX) { - if managed_by_url.contains_key(&normalized_url) { - skipped_no_prefix += 1; - tracing::debug!( - server_name = %name, - server_url = %url, - "Skipping managed injection: URL matches but name lacks '{}' prefix", - MANAGED_MCP_PREFIX, - ); - } - continue; - } - - let scope = headers - .iter() - .find(|h| h.name.eq_ignore_ascii_case("x-connector-scope")) - .map(|h| h.value.clone()); - let scope_id = headers - .iter() - .find(|h| h.name.eq_ignore_ascii_case("x-connector-scope-id")) - .map(|h| h.value.clone()); - - let config = match (&scope, &scope_id) { - (Some(s), Some(id)) => { - let key = (normalized_url.clone(), Some(s.clone()), Some(id.clone())); - managed_by_key.get(&key).copied() - } - _ => managed_by_url.get(&normalized_url).copied(), - }; - - let Some(config) = config else { - skipped_no_match += 1; - tracing::debug!( - server_name = %name, - server_url = %url, - scope = ?scope, - scope_id = ?scope_id, - "Skipping managed injection: no matching managed config" - ); - continue; - }; - - if config.headers.is_empty() { - skipped_no_headers += 1; - tracing::debug!( - server_name = %name, - server_url = %url, - "Skipping managed injection: managed config matched but has no headers", - ); - continue; - } - - let managed_keys: std::collections::HashSet<&str> = - config.headers.keys().map(|k| k.as_str()).collect(); - headers.retain(|h| { - !managed_keys.contains(h.name.as_str()) - && !h.name.eq_ignore_ascii_case("x-connector-scope") - && !h.name.eq_ignore_ascii_case("x-connector-scope-id") - }); - headers.extend( - config - .headers - .iter() - .map(|(k, v)| acp::HttpHeader::new(k.clone(), v.clone())), - ); - - injected += 1; - } - - if injected > 0 { - tracing::info!(count = injected, "Injected managed MCP credentials"); - } - if skipped_no_prefix > 0 { - tracing::info!( - count = skipped_no_prefix, - "Skipped servers with matching URLs but missing '{}' prefix", - MANAGED_MCP_PREFIX, - ); - } - if skipped_no_match > 0 { - tracing::info!( - count = skipped_no_match, - "Skipped servers with '{}' prefix but no matching managed config (URL+scope)", - MANAGED_MCP_PREFIX, - ); - } - if skipped_no_headers > 0 { - tracing::info!( - count = skipped_no_headers, - "Skipped servers with matching managed config but empty headers", - ); - } -} - -/// Boxed future returned by a [`TokenProvider`] call. -pub type TokenFuture = std::pin::Pin> + Send>>; - -/// Resolves a fresh bearer token for each proactive refresh attempt; `None` -/// when no valid token is available. Injected by the caller so this crate -/// stays independent of shell's auth manager. -pub type TokenProvider = Arc TokenFuture + Send + Sync>; - -/// Context needed for the proactive refresh background task. -pub struct RefreshContext { - pub proxy_base_url: String, - /// Per-attempt token resolution (backed by shell's live auth manager). - pub token_provider: TokenProvider, -} - -/// Proactive refresh: sleep until ~5 min before earliest token expiry, -/// re-fetch configs, and update the agent-level cache so new sessions -/// (and re-connected MCP clients) get fresh credentials. -pub fn spawn_cache_refresh_task( - state: ManagedMcpStateHandle, - ctx: RefreshContext, - cancel: CancellationToken, -) -> tokio::task::JoinHandle<()> { - tokio::spawn(async move { - let mut consecutive_failures: u32 = 0; - loop { - let iteration = async { - let wake_at = { - let state_ref = state.lock().await; - match &state_ref.cache { - ManagedMcpCache::Ready(configs) => configs - .iter() - .filter_map(|c| c.token_expires_at) - .min() - .map(|exp| exp - chrono::Duration::minutes(TOKEN_EXPIRY_BUFFER_MINUTES)) - .unwrap_or_else(|| Utc::now() + chrono::Duration::minutes(30)), - _ => Utc::now() + chrono::Duration::minutes(30), - } - }; - - let sleep_dur = (wake_at - Utc::now()) - .to_std() - .unwrap_or(std::time::Duration::from_secs(60)); - - tracing::debug!("MCP refresh: sleeping {sleep_dur:?} until next token refresh"); - tokio::time::sleep(sleep_dur).await; - - let Some(auth_key) = (ctx.token_provider)().await else { - consecutive_failures = consecutive_failures.saturating_add(1); - let backoff_secs = 2u64.pow(consecutive_failures.min(5)); - tracing::debug!( - failures = consecutive_failures, - backoff_secs, - "MCP refresh: auth unavailable, backing off before retry" - ); - tokio::time::sleep(std::time::Duration::from_secs(backoff_secs)).await; - return; - }; - - // Empty success is treated like failure here on purpose: the - // refresh task only runs when a previous fetch returned - // expiring configs, so keep the old (still usable) cache and - // back off rather than wiping it on a suspicious response. - let configs = match fetch_managed_configs(&ctx.proxy_base_url, &auth_key).await { - Ok(configs) if !configs.is_empty() => configs, - Ok(_) | Err(_) => { - consecutive_failures = consecutive_failures.saturating_add(1); - let backoff_secs = 2u64.pow(consecutive_failures.min(5)); - tracing::debug!( - failures = consecutive_failures, - backoff_secs, - "MCP refresh: failed or empty response, backing off before retry" - ); - tokio::time::sleep(std::time::Duration::from_secs(backoff_secs)).await; - return; - } - }; - - consecutive_failures = 0; - tracing::info!( - count = configs.len(), - "MCP refresh: updated managed config cache" - ); - { - let mut s = state.lock().await; - s.complete_fetch(configs, &state, None); - // A proactive refresh pulled fresh configs, so let a parked - // (terminal) connector retry reactively on its next tool call. - s.clear_reauth_cooldowns(); - } - }; - - tokio::select! { - _ = cancel.cancelled() => { - tracing::debug!("MCP refresh: task cancelled"); - return; - } - _ = iteration => {} - } - } - }) -} - -#[cfg(test)] -mod tests { - 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, - } - } - - #[test] - fn managed_token_none_expiry_is_stale() { - let now = chrono::Utc::now(); - // No TTL info → conservatively stale. - assert!(managed_token_is_stale(None, now)); - // Well inside the buffer window → not yet stale. - assert!(!managed_token_is_stale( - Some(now + chrono::Duration::hours(1)), - now - )); - // Within the buffer window → stale. - assert!(managed_token_is_stale( - Some(now + chrono::Duration::minutes(TOKEN_EXPIRY_BUFFER_MINUTES - 1)), - now - )); - // Already expired → stale. - assert!(managed_token_is_stale( - Some(now - chrono::Duration::minutes(1)), - now - )); - } - - /// A failed fetch (here: no auth token) must NOT be committed to the - /// cache as `Ready([])`. The cache has no TTL, so caching a transient - /// failure as an empty catalog would erase managed connectors for the - /// process lifetime — days, for a leader. It must roll back to - /// `NotFetched` so the next caller retries. - #[tokio::test] - async fn failed_fetch_is_not_cached_as_ready_empty() { - let handle = ManagedMcpStateHandle::default(); - let configs = get_or_fetch(&handle, "http://127.0.0.1:0", None, None).await; - assert!(configs.is_empty()); - assert!( - matches!(handle.lock().await.cache, ManagedMcpCache::NotFetched), - "failed fetch must roll back to NotFetched, not poison the cache as Ready([])" - ); - } - - #[test] - fn gateway_tool_catalog_deserializes() { - let catalog: GatewayToolCatalog = serde_json::from_str( - r#"{ - "tools": [ - { - "connector_id": "gmail", - "connector_name": "Gmail", - "tool_id": "search", - "tool_name": "Search Gmail", - "call_id": "gmail_search", - "description": "Search email by query", - "json_schema": { - "type": "object", - "properties": { - "query": { "type": "string" } - }, - "required": ["query"] - } - } - ], - "total_tools": 1, - "connectors_needing_reauth": ["Slack"] - }"#, - ) - .unwrap(); - - assert_eq!(1, catalog.total_tools); - let without_total_tools: GatewayToolCatalog = serde_json::from_str( - r#"{ - "tools": [], - "connectors_needing_reauth": [] - }"#, - ) - .unwrap(); - assert_eq!(0, without_total_tools.total_tools); - assert_eq!(vec!["Slack"], catalog.connectors_needing_reauth); - assert_eq!("gmail_search", catalog.tools[0].call_id); - assert_eq!("gmail__search", catalog.tools[0].qualified_name()); - assert_eq!("gmail", catalog.tools[0].connector_id); - assert_eq!("Gmail", catalog.tools[0].connector_name); - assert_eq!("search", catalog.tools[0].tool_id); - assert_eq!("Search Gmail", catalog.tools[0].tool_name); - assert_eq!( - Some("string"), - catalog.tools[0] - .json_schema - .pointer("/properties/query/type") - .and_then(|v| v.as_str()) - ); - } - - #[tokio::test] - async fn gateway_tool_call_error_preserves_proxy_message() { - use axum::Router; - use axum::routing::post; - use tokio::net::TcpListener; - - let app = Router::new().route( - "/mcp/tools/call", - post(|| async { - ( - axum::http::StatusCode::BAD_REQUEST, - axum::Json(serde_json::json!({ - "code": "Client specified an invalid argument", - "error": "Invalid arguments for google_calendar_availability: missing field `calendars`" - })), - ) - }), - ); - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - - let err = call_gateway_tool( - &format!("http://{addr}"), - "token", - "google_calendar_availability", - serde_json::json!({}), - ) - .await - .unwrap_err(); - - match err { - ManagedMcpFetchError::Status { status, message } => { - assert_eq!(reqwest::StatusCode::BAD_REQUEST, status); - assert_eq!( - "Invalid arguments for google_calendar_availability: missing field `calendars`", - message - ); - } - other => panic!("expected status error, got {other:?}"), - } - } - - #[test] - fn disable_gateway_tools_clears_cached_catalog() { - let mut state = ManagedMcpState::default(); - state.enable_gateway_tools(); - let epoch = state.start_gateway_tool_fetch().unwrap(); - assert!(state.complete_gateway_tool_fetch( - epoch, - GatewayToolCatalog { - tools: vec![], - total_tools: 0, - connectors_needing_reauth: vec![], - } - )); - assert!(state.gateway_tools_active); - assert!(matches!( - state.gateway_tool_cache, - GatewayToolCatalogCache::Ready(_) - )); - - state.disable_gateway_tools(); - assert!(!state.gateway_tools_active); - assert!(matches!( - state.gateway_tool_cache, - GatewayToolCatalogCache::NotFetched - )); - } - - #[test] - fn stale_gateway_tool_fetch_success_does_not_commit_after_disable() { - let mut state = ManagedMcpState::default(); - state.enable_gateway_tools(); - let epoch = state.start_gateway_tool_fetch().unwrap(); - state.disable_gateway_tools(); - - let committed = state.complete_gateway_tool_fetch( - epoch, - GatewayToolCatalog { - tools: vec![], - total_tools: 0, - connectors_needing_reauth: vec![], - }, - ); - - assert!(!committed); - assert!(!state.gateway_tools_active); - assert!(matches!( - state.gateway_tool_cache, - GatewayToolCatalogCache::NotFetched - )); - } - - #[tokio::test] - async fn gateway_tool_waiter_woken_by_disable_does_not_reenable() { - let handle = ManagedMcpStateHandle::default(); - { - let mut state = handle.lock().await; - state.enable_gateway_tools(); - state.start_gateway_tool_fetch().unwrap(); - } - let waiter_handle = handle.clone(); - let waiter = tokio::spawn(async move { - get_or_fetch_gateway_tool_catalog(&waiter_handle, "http://127.0.0.1:0", Some("token")) - .await - }); - - tokio::task::yield_now().await; - handle.lock().await.disable_gateway_tools(); - let catalog = tokio::time::timeout(std::time::Duration::from_secs(1), waiter) - .await - .expect("waiter must wake after disable") - .expect("waiter task should not panic"); - assert!(catalog.is_none()); - let state = handle.lock().await; - assert!(!state.gateway_tools_active); - assert!(matches!( - state.gateway_tool_cache, - GatewayToolCatalogCache::NotFetched - )); - } - - #[tokio::test] - async fn failed_gateway_tool_fetch_is_not_cached_as_ready_empty() { - let handle = ManagedMcpStateHandle::default(); - let catalog = get_or_fetch_gateway_tool_catalog(&handle, "http://127.0.0.1:0", None).await; - assert!(catalog.is_none()); - assert!( - matches!( - handle.lock().await.gateway_tool_cache, - GatewayToolCatalogCache::NotFetched - ), - "failed gateway tool fetch must roll back to NotFetched, not poison the cache as Ready(empty)" - ); - } - - #[test] - fn failed_gateway_tool_fetch_does_not_clear_ready_catalog_from_same_epoch() { - let mut state = ManagedMcpState::default(); - state.enable_gateway_tools(); - let epoch = state.start_gateway_tool_fetch().unwrap(); - assert!(state.complete_gateway_tool_fetch( - epoch, - GatewayToolCatalog { - tools: vec![], - total_tools: 0, - connectors_needing_reauth: vec![], - }, - )); - - state.fail_gateway_tool_fetch(epoch); - assert!(matches!( - state.gateway_tool_cache, - GatewayToolCatalogCache::Ready(_) - )); - } - - #[tokio::test] - async fn successful_gateway_tool_fetch_is_cached_ready() { - let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let app_calls = calls.clone(); - let app = axum::Router::new().route( - "/mcp/tools/list", - axum::routing::get(move || { - app_calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - async { - axum::Json(serde_json::json!({ - "tools": [ - { - "connector_id": "gmail", - "connector_name": "Gmail", - "tool_id": "search", - "tool_name": "Search Gmail", - "call_id": "gmail_search", - "description": "Search email by query", - "json_schema": { "type": "object" } - } - ], - "total_tools": 1, - "connectors_needing_reauth": [] - })) - } - }), - ); - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base_url = format!("http://{}", listener.local_addr().unwrap()); - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - - let handle = ManagedMcpStateHandle::default(); - handle.lock().await.enable_gateway_tools(); - let catalog = get_or_fetch_gateway_tool_catalog(&handle, &base_url, Some("token")) - .await - .expect("gateway catalog fetch should succeed"); - assert_eq!("gmail__search", catalog.tools[0].qualified_name()); - assert!(matches!( - handle.lock().await.gateway_tool_cache, - GatewayToolCatalogCache::Ready(_) - )); - - let cached = - get_or_fetch_gateway_tool_catalog(&handle, "http://127.0.0.1:0", Some("token")) - .await - .expect("second call should use cached catalog"); - assert_eq!("gmail_search", cached.tools[0].call_id); - assert_eq!(1, calls.load(std::sync::atomic::Ordering::SeqCst)); - server.abort(); - } - - #[tokio::test] - async fn gateway_tool_fetch_waiter_survives_notify_before_await() { - let handle = ManagedMcpStateHandle::default(); - let (epoch, registered) = { - let mut state = handle.lock().await; - state.enable_gateway_tools(); - let epoch = state.start_gateway_tool_fetch().unwrap(); - ( - epoch, - state.gateway_tool_fetch_notify.clone().notified_owned(), - ) - }; - handle.lock().await.fail_gateway_tool_fetch(epoch); - tokio::time::timeout(std::time::Duration::from_secs(1), registered) - .await - .expect("registered gateway catalog waiter must observe notify_waiters"); - } - - /// Concurrent callers must not hang when the in-flight fetch fails: - /// `fail_fetch` wakes waiters, which re-enter the loop and retry (or - /// fail) themselves instead of waiting forever on `fetch_notify`. - #[tokio::test] - async fn concurrent_callers_do_not_hang_on_failed_fetch() { - let handle = ManagedMcpStateHandle::default(); - let h2 = handle.clone(); - let (a, b) = tokio::time::timeout(std::time::Duration::from_secs(5), async { - tokio::join!( - get_or_fetch(&handle, "http://127.0.0.1:0", None, None), - get_or_fetch(&h2, "http://127.0.0.1:0", None, None), - ) - }) - .await - .expect("concurrent get_or_fetch must not hang on failure"); - assert!(a.is_empty() && b.is_empty()); - assert!(matches!( - handle.lock().await.cache, - ManagedMcpCache::NotFetched - )); - } - - /// A fresh server with no cooldown entry is always eligible, and its first - /// failure escalates the backoff window without going terminal yet. - #[test] - fn reauth_first_attempt_allowed_then_backs_off() { - let mut state = ManagedMcpState::default(); - let now = Utc::now(); - assert!(state.reauth_allowed("grok_com_slack", now)); - assert!(!state.reauth_is_terminal("grok_com_slack")); - - state.record_reauth_failure("grok_com_slack", now); - // 2^1 = 2s backoff: not eligible now, eligible after the window. - assert!(!state.reauth_allowed("grok_com_slack", now)); - assert!(state.reauth_allowed("grok_com_slack", now + chrono::Duration::seconds(2))); - assert!(!state.reauth_is_terminal("grok_com_slack")); - } - - /// Backoff escalates per failure and is capped at - /// `REACTIVE_REAUTH_BACKOFF_CAP_SECS`; the cap is observed by forcing the - /// failure count high enough that `2^n` would exceed it. - #[test] - fn reauth_backoff_escalates_and_caps() { - let mut state = ManagedReauthState::default(); - let now = Utc::now(); - - state.record_failure(now); - assert_eq!(state.next_allowed_at, now + chrono::Duration::seconds(2)); - state.record_failure(now); - assert_eq!(state.next_allowed_at, now + chrono::Duration::seconds(4)); - - // Drive failures past the cap exponent; the window must clamp to 64s. - for _ in 0..10 { - state.record_failure(now); - } - assert_eq!( - state.next_allowed_at, - now + chrono::Duration::seconds(REACTIVE_REAUTH_BACKOFF_CAP_SECS as i64) - ); - } - - /// After `MAX_REACTIVE_REAUTH_ATTEMPTS` consecutive failures the server is - /// terminal and never eligible again — even past the backoff window — until - /// the cooldown is cleared. - #[test] - fn reauth_goes_terminal_after_max_attempts() { - let mut state = ManagedMcpState::default(); - let now = Utc::now(); - for _ in 0..MAX_REACTIVE_REAUTH_ATTEMPTS { - state.record_reauth_failure("grok_com_slack", now); - } - assert!(state.reauth_is_terminal("grok_com_slack")); - // Even far past any backoff window, a terminal server stays ineligible. - assert!(!state.reauth_allowed("grok_com_slack", now + chrono::Duration::seconds(3600))); - } - - /// A successful reactive re-auth resets that server's cooldown so the next - /// failure starts a fresh backoff. - #[test] - fn reauth_success_resets_cooldown() { - let mut state = ManagedMcpState::default(); - let now = Utc::now(); - for _ in 0..MAX_REACTIVE_REAUTH_ATTEMPTS { - state.record_reauth_failure("grok_com_slack", now); - } - assert!(state.reauth_is_terminal("grok_com_slack")); - - state.record_reauth_success("grok_com_slack"); - assert!(state.reauth_allowed("grok_com_slack", now)); - assert!(!state.reauth_is_terminal("grok_com_slack")); - } - - /// `complete_fetch` must NOT clear the reactive cooldown: the reactive path - /// re-fetches through it, so clearing there would reset a still-rejected - /// connector's attempt cap every attempt and loop. Only the explicit - /// `clear_reauth_cooldowns` (invoked by the proactive refresh) clears it. - #[test] - fn complete_fetch_preserves_cooldown_clear_resets_it() { - let handle = ManagedMcpStateHandle::default(); - let now = Utc::now(); - { - let mut state = handle.blocking_lock(); - for _ in 0..MAX_REACTIVE_REAUTH_ATTEMPTS { - state.record_reauth_failure("grok_com_slack", now); - state.record_reauth_failure("grok_com_linear", now); - } - assert!(state.reauth_is_terminal("grok_com_slack")); - assert!(state.reauth_is_terminal("grok_com_linear")); - - // A fetch (the reactive path's re-fetch) must leave the cooldown intact. - // refresh_ctx = None so no background task is spawned in the test. - state.complete_fetch( - vec![make_managed("Slack", "https://mcp.slack.com/sse", "user")], - &handle, - None, - ); - assert!(state.reauth_is_terminal("grok_com_slack")); - assert!(state.reauth_is_terminal("grok_com_linear")); - - // The explicit proactive-refresh clear resets every server. - state.clear_reauth_cooldowns(); - assert!(state.reauth_allowed("grok_com_slack", now)); - assert!(state.reauth_allowed("grok_com_linear", now)); - assert!(!state.reauth_is_terminal("grok_com_slack")); - assert!(!state.reauth_is_terminal("grok_com_linear")); - } - } -} diff --git a/crates/codegen/kigi-shell/Cargo.toml b/crates/codegen/kigi-shell/Cargo.toml index d7446b9..ff6581e 100644 --- a/crates/codegen/kigi-shell/Cargo.toml +++ b/crates/codegen/kigi-shell/Cargo.toml @@ -115,7 +115,6 @@ prost = { workspace = true, optional = true } kigi-paths = { path = "../kigi-paths" } kigi-shell-base = { workspace = true } kigi-env = { workspace = true } -kigi-shell-session-support = { workspace = true } kigi-shared = { workspace = true } kigi-tracing-macros = { path = "../kigi-tracing-macros" } kigi-fsnotify = { path = "../kigi-fsnotify" } diff --git a/crates/codegen/kigi-shell/skills/imagine/SKILL.md b/crates/codegen/kigi-shell/skills/imagine/SKILL.md deleted file mode 100644 index b59ea56..0000000 --- a/crates/codegen/kigi-shell/skills/imagine/SKILL.md +++ /dev/null @@ -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. \ No newline at end of file diff --git a/crates/codegen/kigi-shell/src/agent/app.rs b/crates/codegen/kigi-shell/src/agent/app.rs index bc9a49e..5d35176 100644 --- a/crates/codegen/kigi-shell/src/agent/app.rs +++ b/crates/codegen/kigi-shell/src/agent/app.rs @@ -917,15 +917,6 @@ pub async fn run_leader( } ConfigUpdate::AuthCleared => { 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; kigi_log::unified_log::warn( "auth cleared from disk", @@ -1144,14 +1135,6 @@ mod tests { ); let msg: serde_json::Value = serde_json::from_str(line.trim_end()).unwrap(); 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] diff --git a/crates/codegen/kigi-shell/src/agent/config.rs b/crates/codegen/kigi-shell/src/agent/config.rs index f3ff708..e173ec0 100644 --- a/crates/codegen/kigi-shell/src/agent/config.rs +++ b/crates/codegen/kigi-shell/src/agent/config.rs @@ -301,9 +301,6 @@ pub struct Requirements { pub tool_search: Constrained, pub web_fetch: Constrained, pub ask_user_question: Constrained, - pub image_gen: Constrained, - pub image_edit: Constrained, - pub video_gen: Constrained, pub write_file: Constrained, /// Voice dictation (STT). Pin via requirements/managed `[features] voice_mode`. pub voice_mode: Constrained, @@ -1075,8 +1072,6 @@ pub struct Config { pub memory: crate::config::MemoryConfig, #[serde(default, skip_serializing)] pub compaction: CompactionConfig, - #[serde(default, skip_serializing)] - pub managed_mcps: crate::config::ManagedMcpsConfig, /// `[auth]` alias — consumed by `expand_auth_alias` before serde. /// Typed as `KimiCodeConfig` (same schema) so sub-field typos are caught. #[serde(default, skip_serializing)] @@ -1206,18 +1201,6 @@ pub struct Config { /// Resolved by [`crate::config::ToolsConfig::resolve`]. #[serde(skip)] 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, /// Whether to enrich path-not-found errors with CWD reminders, /// "dropped repo folder" correction, and similar-name suggestions. /// Default `false`. Enabled via remote settings. @@ -1225,13 +1208,6 @@ pub struct Config { /// which sessions had path-not-found hints active. #[serde(default)] 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 /// completes, immediately inject a synthetic prompt instead of waiting /// for the idle-gated notification drain. @@ -1446,7 +1422,6 @@ impl Default for Config { subagents: crate::config::SubagentsConfig::default(), memory: crate::config::MemoryConfig::default(), compaction: CompactionConfig::default(), - managed_mcps: crate::config::ManagedMcpsConfig::default(), auth: None, desktop: None, tips: None, @@ -1476,15 +1451,11 @@ impl Default for Config { todo_gate: false, laziness_debug_log: None, respect_gitignore: false, - disable_zdr_incompatible_tools: false, - zdr_video_output_s3: None, path_not_found_hints: false, cli_experimental_memory: false, cli_no_memory: false, cli_subagents: None, memory_config: None, - managed_mcps_enabled: true, - managed_mcp_gateway_tools_enabled: false, auto_wake_enabled: true, compat_resolved: CompatConfig::default(), requirements: Requirements::default(), @@ -1603,8 +1574,6 @@ impl Config { /// Call immediately after `new_from_toml_cfg()`. Fields resolved: /// - subagents (6 fields) via `SubagentsConfig::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 / /// prompt_suggest_model_pin via `ModelOverrideConfig::resolve` /// - memory_config via `MemoryConfig::resolve` @@ -1624,15 +1593,6 @@ impl Config { Some(pinned) => pinned, 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( ctx.cli_session_summary_model, ctx.raw_config, @@ -1867,51 +1827,6 @@ impl Config { .default(true) .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 { - 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 { - 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 { - 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 /// reach cli-chat-proxy `/v1/settings` (custom `models_base_url`, external /// `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`). #[serde(default, skip_serializing_if = "Option::is_none")] pub two_pass_compaction: Option, - /// Video generation tool. `None` = defer to remote settings / env / default (false). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub video_gen: Option, - /// `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, /// Write file tool. `None` = defer to remote settings / env / default (true). #[serde(default, skip_serializing_if = "Option::is_none")] pub write_file: Option, @@ -4472,34 +4380,6 @@ reasoning_effort = "low" assert_eq!(cfg.toolset.bash.timeout_secs, Some(30.5)); } #[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 ctx(raw: &toml::Value, disable_web_search: bool) -> RuntimeResolutionContext<'_> { RuntimeResolutionContext { @@ -7215,51 +7095,6 @@ reasoning_effort = "low" assert_eq!(r.source, ConfigSource::Remote); 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 /// start from a known baseline regardless of run order. fn clear_goal_envs() { @@ -8103,8 +7938,6 @@ agent_type = "cursor" secret = "value" [worktree_pool] pool_size = 4 - [managed_mcps] - enabled = true [mcp_servers.test] url = "https://mcp.test.com" [toolset.bash] @@ -8239,12 +8072,6 @@ agent_type = "cursor" 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 { COMPAT_CELLS .into_iter() @@ -8531,7 +8358,6 @@ hooks = true #[serial] fn resolve_runtime_fields_interactive_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 { @@ -8550,8 +8376,6 @@ hooks = true }); assert!(cfg.subagents_enabled); assert!(!cfg.respect_gitignore); - assert!(cfg.managed_mcps_enabled); - assert!(!cfg.managed_mcp_gateway_tools_enabled); assert_eq!( cfg.session_summary_model, Some(crate::models::default_session_summary_model().to_owned()) @@ -8560,60 +8384,6 @@ hooks = true } #[test] #[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() { clear_runtime_env_vars(); let raw: toml::Value = toml::from_str("[subagents]\nenabled = true").unwrap(); @@ -8751,12 +8521,10 @@ hooks = true cfg.resolve_runtime_fields(&ctx); let first_subagents = cfg.subagents_enabled; let first_gitignore = cfg.respect_gitignore; - let first_mcps = cfg.managed_mcps_enabled; let first_ss = cfg.session_summary_model.clone(); cfg.resolve_runtime_fields(&ctx); assert_eq!(cfg.subagents_enabled, first_subagents); assert_eq!(cfg.respect_gitignore, first_gitignore); - assert_eq!(cfg.managed_mcps_enabled, first_mcps); assert_eq!(cfg.session_summary_model, first_ss); } #[test] diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs index 8f62131..09957bb 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs @@ -270,13 +270,7 @@ impl acp::Agent for MvpAgent { let current_working_directory = self.launch_cwd.clone(); let hostname = gethostname::gethostname(); let mcp_servers: Vec = Vec::new(); - let fetch_managed_mcps = self.cfg.borrow().managed_mcps_enabled - && 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(); + self.spawn_initialize_launch_mcp_setup(); let init_model_state = self.model_state(None); Ok( acp::InitializeResponse::new(acp::ProtocolVersion::V1) @@ -545,9 +539,7 @@ impl acp::Agent for MvpAgent { let remote_settings = self.cfg.borrow().remote_settings.clone(); folder_trust::resolve_and_record(cwd.as_path(), remote_settings.as_ref(), false); let initial_client_mcp_servers = arguments.mcp_servers.clone(); - let (mcp_servers, managed_mcp_expires_at) = self - .resolve_mcp_servers(arguments.mcp_servers, cwd.as_path()) - .await; + let mcp_servers = self.resolve_mcp_servers(arguments.mcp_servers, cwd.as_path()); let mcp_meta_config_map = parse_mcp_meta_config(arguments.meta.as_ref()); let client_session_id = arguments .meta @@ -756,7 +748,6 @@ impl acp::Agent for MvpAgent { persisted_goal_mode: None, persisted_announcement_state: None, session_meta: arguments.meta.as_ref(), - managed_mcp_expires_at, model_agent_type: model_agent_type.as_deref(), session_model_id, session_yolo_mode, @@ -902,9 +893,7 @@ impl acp::Agent for MvpAgent { let remote_settings = self.cfg.borrow().remote_settings.clone(); folder_trust::resolve_and_record(cwd.as_path(), remote_settings.as_ref(), false); let initial_client_mcp_servers = client_mcp_servers.clone(); - let (mcp_servers, managed_mcp_expires_at) = self - .resolve_mcp_servers(client_mcp_servers, cwd.as_path()) - .await; + let mcp_servers = self.resolve_mcp_servers(client_mcp_servers, cwd.as_path()); let mcp_meta_config_map = parse_mcp_meta_config(request_meta.as_ref()); let mut load_timer = crate::instrumentation_timer!("session.load_session"); 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_announcement_state, session_meta: request_meta.as_ref(), - managed_mcp_expires_at, model_agent_type: persisted_agent_name.as_deref(), session_model_id: summary.current_model_id.clone(), session_yolo_mode, @@ -2083,7 +2071,7 @@ impl acp::Agent for MvpAgent { | "x.ai/internal/reload_all_mcp_servers" | "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/auth_cleared" | "x.ai/plugins/reload" + | "x.ai/plugins/reload" | "x.ai/commands/list" => { crate::extensions::session_admin::handle(self, &args).await } diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs index b6e97ca..6db8958 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs @@ -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 { - 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 { - 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>, - ) { - 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 /// with its path. /// @@ -226,15 +109,12 @@ impl MvpAgent { /// Resolve folder trust and load launch-dir MCP configs after `initialize` /// returns. The walks are synchronous and expensive in large monorepos; they /// 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 compat = self.cfg.borrow().compat_resolved; let remote_settings = self.cfg.borrow().remote_settings.clone(); let gateway = self.gateway.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 { let local_mcp_servers = match tokio::task::spawn_blocking(move || { let local = crate::util::config::load_mcp_servers(&cwd, &compat); @@ -258,27 +138,9 @@ impl MvpAgent { } crate::extensions::mcp::notify_servers_updated( &gateway, - &[], &local_mcp_servers, ) .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( @@ -311,23 +173,19 @@ impl MvpAgent { plugin_count = count, "lazily populated plugin registry snapshot" ); } - /// Fetch managed configs, merge with client servers, return merged list + earliest expiry. - pub(super) async fn resolve_mcp_servers( + /// Merge on-disk/plugin MCP servers with client servers. + pub(super) fn resolve_mcp_servers( &self, client_servers: Vec, cwd: &std::path::Path, - ) -> (Vec, Option>) { + ) -> Vec { self.ensure_plugin_registry(); - let managed = self.get_managed_mcp_configs().await; - let expires_at = managed.iter().filter_map(|c| c.token_expires_at).min(); - let merged = crate::session::managed_mcp::merge_managed_mcp_servers( + crate::session::managed_mcp::merge_managed_mcp_servers( client_servers, cwd, - &managed, self.plugin_registry_handle.snapshot().as_deref(), &self.cfg.borrow().compat_resolved, - ); - (merged, expires_at) + ) } /// Set the memory configuration (called from TUI after config resolution). pub fn set_memory_config(&mut self, config: crate::config::MemoryConfig) { @@ -760,39 +618,6 @@ impl MvpAgent { ); (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. pub(super) fn prepare_app_builder_deployer_config( &self, @@ -800,40 +625,6 @@ impl MvpAgent { use kigi_tools::implementations::grok_build::deploy_app::AppBuilderDeployerConfig; 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 /// the Kimi Code subscription channel (`POST {coding_base}/search`, /// kimi-cli `auth/platforms.py`), so this is `Enabled` only for OAuth @@ -1016,7 +807,6 @@ impl MvpAgent { worktree_type, restore_code, session_registry_local, - managed_mcp_cache: Default::default(), agent_mcp_state: std::sync::Arc::new( tokio::sync::Mutex::new( crate::session::mcp_servers::McpState::new(vec![]), @@ -1829,7 +1619,6 @@ impl MvpAgent { persisted_goal_mode, persisted_announcement_state, session_meta, - managed_mcp_expires_at, model_agent_type, session_model_id, session_yolo_mode, @@ -2258,8 +2047,6 @@ impl MvpAgent { .and_then(|entry| entry.info.max_retries); let origin_client = self.origin_client_info_from_meta(init.meta.as_ref()); 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 web_fetch_config = self.prepare_web_fetch_config(); let write_file_enabled = self.cfg.borrow().resolve_write_file().value; @@ -2309,7 +2096,6 @@ impl MvpAgent { let cfg = self.cfg.borrow(); cfg.resolve_backend_tools().value }; - let managed_mcp_proxy_url = self.cfg.borrow().endpoints.proxy_url(); let init_meta = self .initialize_request .get() @@ -2471,9 +2257,6 @@ impl MvpAgent { persisted_announcement_state, self.memory_config.clone(), feedback_flags, - self.managed_mcp_cache.clone(), - managed_mcp_expires_at, - managed_mcp_proxy_url, session_model_id, session_yolo_mode, session_auto_mode, @@ -2482,8 +2265,6 @@ impl MvpAgent { model_max_retries, web_search_config, web_fetch_config, - image_gen_config, - video_gen_config, app_builder_deployer_config, write_file_enabled, goal_enabled, @@ -2617,7 +2398,6 @@ impl MvpAgent { self.notify_session_cwd_for_watch(std::path::Path::new(&session_info.cwd)); self.activity.register_session(&session_info.id.0, &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(); tokio::spawn(async move { crate::session::prompt_history::truncate_if_needed_async(cwd_for_maintenance) diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/folder_trust_prompt.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/folder_trust_prompt.rs index 0253e0e..21243b3 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/folder_trust_prompt.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/folder_trust_prompt.rs @@ -150,10 +150,6 @@ impl MvpAgent { let gateway = self.gateway.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 remote = remote.cloned(); let cwd = cwd.to_path_buf(); @@ -258,10 +254,6 @@ impl MvpAgent { gateway: &gateway, targets, plugin_handle: &plugin_handle, - managed_mcp_cache: &managed_mcp_cache, - auth_manager: &auth_manager, - can_fetch_managed, - proxy_url: &proxy_url, compat: &compat, prompt_cwd: &cwd, }) @@ -291,10 +283,6 @@ struct ReloadAfterGrant<'a> { /// Every session sharing the granted workspace, each with its own cwd. targets: Vec, plugin_handle: &'a kigi_agent::plugins::SharedPluginRegistryHandle, - managed_mcp_cache: &'a crate::session::managed_mcp::ManagedMcpStateHandle, - auth_manager: &'a std::sync::Arc, - can_fetch_managed: bool, - proxy_url: &'a str, compat: &'a kigi_tools::types::CompatConfig, /// The prompting session's cwd — used only for the client catalog push. 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 /// (no restart), driving the canonical primitives the normal spawn/reload paths /// 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` /// (`SessionCommand::ReloadPlugins`), and `reload_hooks_impl` /// (`SessionCommand::ReloadHooks`), then push the refreshed MCP catalog. LSP is /// spawn-baked and applies on the next session open (see module docs). Caller /// must have granted + recorded trust first. async fn reload_project_servers_after_grant(ctx: ReloadAfterGrant<'_>) { - // Managed (gateway/Toolbox) servers must survive the re-merge; fetch them once - // (cwd-independent) via the shared helper (single-sources the auth-key dance - // 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![] - }; + // The plugin MCP snapshot is global, so it is fine to reuse across cwds + // for the merge. let plugin_snapshot = ctx.plugin_handle.snapshot(); 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( target.initial_client_mcp_servers, session_cwd, - &managed, plugin_snapshot.as_deref(), ctx.compat, ); @@ -375,7 +350,7 @@ async fn reload_project_servers_after_grant(ctx: ReloadAfterGrant<'_>) { ctx.prompt_cwd, 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)] diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs index 17d5352..9b57509 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs @@ -124,7 +124,6 @@ pub(crate) struct SessionSpawnOptions<'a> { crate::session::announcement_state::AnnouncementState, >, pub session_meta: Option<&'a acp::Meta>, - pub managed_mcp_expires_at: Option>, pub model_agent_type: Option<&'a str>, pub session_model_id: acp::ModelId, pub session_yolo_mode: bool, @@ -260,7 +259,6 @@ pub(crate) fn chat_session_spawn_options<'a>( persisted_goal_mode: None, persisted_announcement_state: None, session_meta, - managed_mcp_expires_at: None, model_agent_type, session_model_id, session_yolo_mode, @@ -591,8 +589,6 @@ pub struct MvpAgent { /// Local config.toml override for session registry (`[cli] session_registry`). /// `Some(true)` enables, `Some(false)` disables, `None` defers to remote settings. session_registry_local: Option, - /// 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-scoped, not per-client. agent_mcp_state: std::sync::Arc< diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/subagent_coordinator.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/subagent_coordinator.rs index a0a1aa7..79e93a2 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/subagent_coordinator.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/subagent_coordinator.rs @@ -304,7 +304,6 @@ impl MvpAgent { session_env, parent_attribution_callback, parent_agent_name, - parent_managed_mcp_proxy_base_url, ) = { let sessions = self.sessions.borrow(); let ps = sessions.get(&parent_sid); @@ -339,7 +338,6 @@ impl MvpAgent { .unwrap_or_else(|| std::sync::Arc::new(std::collections::HashMap::new())), ps.and_then(|h| h.attribution_callback.clone()), ps.map(|h| h.agent_name.clone()), - ps.map(|h| h.managed_mcp_proxy_base_url.clone()), ) }; let ( @@ -401,8 +399,6 @@ impl MvpAgent { gateway: self.gateway.clone(), client_hooks: Default::default(), 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(), auth_method_id: self .auth_method_id @@ -429,8 +425,6 @@ impl MvpAgent { memory_config: self.memory_config.clone(), web_search_config: self.prepare_web_search_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(), write_file_enabled: self.cfg.borrow().resolve_write_file().value, goal_enabled: self.cfg.borrow().resolve_goal().value, @@ -505,7 +499,6 @@ impl MvpAgent { .map(|h| h.mcp_servers.clone()) .unwrap_or_default() }, - managed_mcp_state: self.managed_mcp_cache.clone(), parent_mcp_pool: None, parent_tool_snapshot: None, parent_skills: None, diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs index 7994be5..8295293 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs @@ -648,8 +648,6 @@ async fn file_toolset_override_e2e_to_finalized_toolset() { web_search_config: kigi_tools::implementations::web_search::WebSearchConfig::default(), web_fetch_config: Default::default(), 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(), api_key_provider: None, attribution_callback: None, @@ -744,7 +742,6 @@ fn make_test_handle( permission_handle: kigi_workspace::permission::PermissionHandle::allow_all(), attribution_callback: None, agent_name: "grok-build".to_string(), - managed_mcp_proxy_base_url: String::new(), session_default_agent_profile: None, allowed_subagent_types: None, hook_registry: None, @@ -1136,38 +1133,6 @@ fn test_sessionless_request_requires_session_id() { "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(¶ms).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. fn build_minimal_agent_for_tests() -> MvpAgent { 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` /// regardless of live-token state -- BYOK sessions have nothing to refresh, /// 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")] async fn auth_type_xai_api_key_no_current_returns_api_key() { 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", ); } -/// 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 /// disk-backed Build path. Anything but an explicit `kind: "chat"` is Build. #[test] @@ -2107,7 +2052,6 @@ fn chat_session_spawn_options_matches_thin_profile() { assert!(!opts.client_fs_read); assert!(!opts.client_fs_write); assert!(opts.chat_history.is_empty()); - assert!(opts.managed_mcp_expires_at.is_none()); assert!(!opts.session_auto_mode); assert!( opts.persistence.is_noop(), diff --git a/crates/codegen/kigi-shell/src/agent/subagent/handle_request.rs b/crates/codegen/kigi-shell/src/agent/subagent/handle_request.rs index ce46d37..72de31d 100644 --- a/crates/codegen/kigi-shell/src/agent/subagent/handle_request.rs +++ b/crates/codegen/kigi-shell/src/agent/subagent/handle_request.rs @@ -1054,9 +1054,6 @@ pub(crate) async fn handle_subagent_request( ctx.memory_config.clone() }, Default::default(), - ctx.managed_mcp_state.clone(), - None, - ctx.managed_mcp_proxy_base_url.clone(), effective_model_id, ctx.yolo_mode || matches!( @@ -1069,8 +1066,6 @@ pub(crate) async fn handle_subagent_request( None, ctx.web_search_config.clone(), ctx.web_fetch_config.clone(), - ctx.image_gen_config.clone(), - ctx.video_gen_config.clone(), ctx.app_builder_deployer_config.clone(), ctx.write_file_enabled, ctx.goal_enabled, diff --git a/crates/codegen/kigi-shell/src/agent/subagent/mod.rs b/crates/codegen/kigi-shell/src/agent/subagent/mod.rs index bcc39d2..4a83e01 100644 --- a/crates/codegen/kigi-shell/src/agent/subagent/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/subagent/mod.rs @@ -144,7 +144,6 @@ pub(crate) struct SubagentSpawnContext { /// context is built (an async snapshot from the parent session actor). pub client_hooks: crate::extensions::hooks::ClientHooks, pub sampling_config: kigi_sampler::SamplerConfig, - pub managed_mcp_proxy_base_url: String, /// The staging auth header value propagated from the parent. Used /// when materialising subagent `SamplerConfig`s for auth-flow tracking /// 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, /// Resolved config for web fetch. 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. pub app_builder_deployer_config: 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) /// will not appear here. Named references only resolve against the initial config. pub parent_mcp_configs: Vec, - /// 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. pub parent_mcp_pool: Option, /// Snapshot of the parent session's resolved tool schema at spawn time. diff --git a/crates/codegen/kigi-shell/src/agent/subagent/tests/mod.rs b/crates/codegen/kigi-shell/src/agent/subagent/tests/mod.rs index 5ba5190..60cdf2d 100644 --- a/crates/codegen/kigi-shell/src/agent/subagent/tests/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/subagent/tests/mod.rs @@ -143,16 +143,6 @@ fn subagent_inherits_parent_lsp_via_context() { ); } #[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() { let ctx = ctx_with_toggle(HashMap::new()); assert!(ctx.lsp.is_none()); @@ -1110,7 +1100,6 @@ fn dummy_tracker( permission_handle: kigi_workspace::permission::PermissionHandle::allow_all(), attribution_callback: None, agent_name: "grok-build".to_string(), - managed_mcp_proxy_base_url: String::new(), session_default_agent_profile: None, allowed_subagent_types: None, hook_registry: None, diff --git a/crates/codegen/kigi-shell/src/auth/attribution.rs b/crates/codegen/kigi-shell/src/auth/attribution.rs index ea5b2d1..c4da46f 100644 --- a/crates/codegen/kigi-shell/src/auth/attribution.rs +++ b/crates/codegen/kigi-shell/src/auth/attribution.rs @@ -123,7 +123,7 @@ impl ShellAttribution { /// Tool-side counterpart of [`Self::new`]: returns /// `Arc` for the /// `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 /// same `auth_401_attribution` event format -- only the trait /// 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) -/// in `kigi-tools` emits a 401 attribution event through this -/// trait when its HTTP request returns UNAUTHORIZED. Same shape as -/// the sampler-side impl 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. +/// Tool-side hook: each tool client (web_search) in `kigi-tools` +/// emits a 401 attribution event through this trait when its HTTP +/// request returns UNAUTHORIZED. Same shape as the sampler-side impl +/// above; routes to the same pair of sinks. impl ToolAuth401AttributionCallback for ShellAttribution { fn record_401(&self, consumer: ToolConsumer, sent_bearer_prefix: Option<&str>) { let (kind, op) = match consumer { - ToolConsumer::ImageGen => (ConsumerKind::ImageGen, ""), - ToolConsumer::VideoGenStart => (ConsumerKind::VideoGen, "start"), - ToolConsumer::VideoGenPoll => (ConsumerKind::VideoGen, "poll"), ToolConsumer::WebSearch => (ConsumerKind::WebSearch, ""), }; record_consumer_401( @@ -207,15 +200,6 @@ pub(crate) enum ConsumerKind { /// No per-op discriminator -- the consumer string is just /// `"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 /// `POST /responses` with a `WebSearch` tool. No per-op /// discriminator; consumer string is just `"WebSearch"`. @@ -230,8 +214,6 @@ impl ConsumerKind { Self::FeedbackClient => "FeedbackClient", Self::SessionRegistryClient => "SessionRegistryClient", Self::IdleResumeModelRefresh => "IdleResumeModelRefresh", - Self::ImageGen => "ImageGen", - Self::VideoGen => "VideoGen", Self::WebSearch => "WebSearch", } } @@ -239,13 +221,10 @@ impl ConsumerKind { /// `true` for variants that take a per-operation discriminator /// appended as `.`. `false` for variants whose /// `consumer` string is just the prefix - /// (`IdleResumeModelRefresh`, `ImageGen`, `WebSearch` -- each is + /// (`IdleResumeModelRefresh`, `WebSearch` -- each is /// a single endpoint with no sub-operation). fn takes_op(self) -> bool { - !matches!( - self, - Self::IdleResumeModelRefresh | Self::ImageGen | Self::WebSearch - ) + !matches!(self, Self::IdleResumeModelRefresh | 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. /// /// 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 /// per-client `record_401_attribution` wrappers in /// `agent/feedback_client.rs` and `agent/session_registry_client.rs` each @@ -606,10 +585,6 @@ mod tests { "ignored", "IdleResumeModelRefresh", ), - (ConsumerKind::ImageGen, "", "ImageGen"), - (ConsumerKind::ImageGen, "ignored", "ImageGen"), - (ConsumerKind::VideoGen, "start", "VideoGen.start"), - (ConsumerKind::VideoGen, "poll", "VideoGen.poll"), (ConsumerKind::WebSearch, "", "WebSearch"), (ConsumerKind::WebSearch, "ignored", "WebSearch"), ]; @@ -647,12 +622,7 @@ mod tests { let cb: Arc = ShellAttribution::new_tool_callback(am_arc.clone(), Some("sid-tool".into())); - let cases = [ - (ToolConsumer::ImageGen, "ImageGen"), - (ToolConsumer::VideoGenStart, "VideoGen.start"), - (ToolConsumer::VideoGenPoll, "VideoGen.poll"), - (ToolConsumer::WebSearch, "WebSearch"), - ]; + let cases = [(ToolConsumer::WebSearch, "WebSearch")]; for (consumer, expected_consumer_str) in cases { cb.record_401(consumer, Some("bearer-1234567890")); diff --git a/crates/codegen/kigi-shell/src/auth/manager.rs b/crates/codegen/kigi-shell/src/auth/manager.rs index 9ad92b5..9ea253c 100644 --- a/crates/codegen/kigi-shell/src/auth/manager.rs +++ b/crates/codegen/kigi-shell/src/auth/manager.rs @@ -1877,7 +1877,7 @@ impl AuthManager { } /// Bridges `Arc` 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. pub(crate) struct SharedAuthKeyProvider(pub Arc); diff --git a/crates/codegen/kigi-shell/src/builtin.rs b/crates/codegen/kigi-shell/src/builtin.rs index 001e818..ef2d9e7 100644 --- a/crates/codegen/kigi-shell/src/builtin.rs +++ b/crates/codegen/kigi-shell/src/builtin.rs @@ -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 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 IMAGINE_SKILL_MD: &str = include_str!("../skills/imagine/SKILL.md"); /// 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"); /// 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. /// - After the rename has had time to propagate, you **may** clean old /// 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 /// 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), ("create-skill", CREATE_SKILL_MD), ("code-review", CODE_REVIEW_SKILL_MD), - ("imagine", IMAGINE_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("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(".metadata_version"), "0.0.0-stale").unwrap(); // 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::write( 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(), "old" ); - for name in ["check-work", "imagine", "code-review"] { + for name in ["check-work", "code-review"] { assert_ne!( std::fs::read_to_string(home.join(format!("skills/{name}/SKILL.md"))).unwrap(), "old", @@ -248,7 +247,7 @@ mod tests { // Legacy skill directories must have been removed (the key part of // 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!( !home.join(format!("skills/{name}")).exists(), "legacy '{name}' skill directory should have been deleted during version bump" diff --git a/crates/codegen/kigi-shell/src/config/mod.rs b/crates/codegen/kigi-shell/src/config/mod.rs index 44db347..4bf8c38 100644 --- a/crates/codegen/kigi-shell/src/config/mod.rs +++ b/crates/codegen/kigi-shell/src/config/mod.rs @@ -474,62 +474,6 @@ impl SubagentsConfig { 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]`. #[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)] #[serde(default)] @@ -656,8 +600,7 @@ impl ModelOverrideConfig { /// /// ```toml /// [tools] -/// disable_zdr_incompatible_tools = true -/// # [tools.zdr_video_output_s3] — see ZdrVideoOutputS3Config +/// respect_gitignore = true /// ``` #[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)] #[serde(default)] @@ -665,29 +608,12 @@ pub struct ToolsConfig { /// When `true`, all tools (including `read_file`) filter gitignored /// files. When `false` (default), each tool picks its own default. 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, } impl ToolsConfig { /// Resolve the final tools config, in priority order: - /// 1. Env vars `KIGI_RESPECT_GITIGNORE` and - /// `KIGI_DISABLE_ZDR_INCOMPATIBLE_TOOLS` (`0`/`false` off, - /// `1`/`true` on). + /// 1. Env var `KIGI_RESPECT_GITIGNORE` (`0`/`false` off, `1`/`true` on). /// 2. `[tools]` block from the merged effective config. - /// 3. Defaults (both `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. + /// 3. Default (`false`). pub fn resolve(config: &toml::Value) -> Self { let tools = config.get("tools"); let mut result = Self { @@ -695,33 +621,6 @@ impl ToolsConfig { .and_then(|t| t.get("respect_gitignore")) .and_then(|v| v.as_bool()) .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() { 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 } } @@ -1009,9 +899,6 @@ fn apply_requirements_inner( pin_feature!(tool_search); pin_feature!(web_fetch); 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!(voice_mode); pin_requirement_only!(remote_fetch); @@ -1020,7 +907,6 @@ fn apply_requirements_inner( enforce_opt!("cli", "show_tips", config.cli.show_tips); enforce_val!("memory", "enabled", config.memory.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") { config .requirements diff --git a/crates/codegen/kigi-shell/src/config/tests.rs b/crates/codegen/kigi-shell/src/config/tests.rs index cdf8fb3..3884597 100644 --- a/crates/codegen/kigi-shell/src/config/tests.rs +++ b/crates/codegen/kigi-shell/src/config/tests.rs @@ -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"); } -fn with_managed_mcp_env( - 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( ss: 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 -/// `ToolsConfig::resolve`, so tests across both fields can't race. +/// Lock shared by every test that touches the env var read by +/// `ToolsConfig::resolve`, so tests can't race. 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( - 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(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(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] fn tools_config_default_disabled() { @@ -1858,87 +1698,6 @@ fn tools_config_env_false_overrides_toml_true() { ); } #[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() { let toml_str = r#" [roles.researcher] diff --git a/crates/codegen/kigi-shell/src/extensions/mcp.rs b/crates/codegen/kigi-shell/src/extensions/mcp.rs index 8e5fd91..97a7fd8 100644 --- a/crates/codegen/kigi-shell/src/extensions/mcp.rs +++ b/crates/codegen/kigi-shell/src/extensions/mcp.rs @@ -2,7 +2,7 @@ //! //! - `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/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 //! `StatusDispatcher` (transport-closed pollers, handshake failures, //! 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"; } use crate::agent::MvpAgent; -use crate::session::managed_mcp::MANAGED_MCP_PREFIX; use crate::session::mcp_servers::{MCP_TOOL_NAME_DELIMITER, McpClient, McpServerName, McpState}; // ── Wire types: mcp/list ──────────────────────────────────────────── @@ -58,8 +57,9 @@ use crate::session::mcp_servers::{MCP_TOOL_NAME_DELIMITER, McpClient, McpServerN pub struct McpListRequest { #[serde(default)] pub session_id: Option, - /// When false, bypasses the managed MCP config cache and fetches fresh - /// from cli-chat-proxy. Set this after OAuth enrollment or disconnect. + /// When false, forces an OAuth retry pass on servers stuck in + /// auth-required state before building the list. Set this after OAuth + /// enrollment or disconnect. #[serde(default = "default_true")] pub cache: bool, } @@ -92,21 +92,13 @@ pub struct McpServerEntry { /// MCP server config for the `mcp/list` catalog response. /// /// 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 #[derive(Debug, Clone, Serialize)] #[serde(tag = "type", rename_all = "camelCase")] pub enum McpServerConfig { #[serde(rename = "http")] - Http { - url: String, - #[serde(skip_serializing_if = "Option::is_none")] - scope: Option, - #[serde(rename = "scopeId", skip_serializing_if = "Option::is_none")] - scope_id: Option, - #[serde(rename = "scopeName", skip_serializing_if = "Option::is_none")] - scope_name: Option, - }, + Http { url: String }, #[serde(rename = "stdio")] Stdio { command: std::path::PathBuf, @@ -115,8 +107,6 @@ pub enum McpServerConfig { #[serde(default, skip_serializing_if = "Vec::is_empty")] env: Vec, }, - #[serde(rename = "managedGateway")] - ManagedGateway, } #[derive(Debug, Clone, Serialize)] @@ -128,7 +118,6 @@ pub struct McpEnvVar { #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] #[serde(rename_all = "lowercase")] pub enum McpServerSource { - Managed, Local, } @@ -306,14 +295,13 @@ pub struct McpReadResourceContent { pub meta: Option, } -/// Push the full MCP catalog to the client. Called in the background after -/// managed configs resolve so `initialize()` isn't blocked by the network fetch. +/// Push the full MCP catalog to the client. Called in the background so +/// `initialize()` isn't blocked by catalog construction. pub async fn notify_servers_updated( gateway: &kigi_acp_lib::AcpAgentGatewaySender, - managed_configs: &[crate::session::managed_mcp::ManagedMcpConfig], local_servers: &[acp::McpServer], ) { - let catalog = build_mcp_catalog(managed_configs, local_servers); + let catalog = build_mcp_catalog(local_servers); let payload = McpServersUpdated { 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`, /// and `mcp/servers_updated`. -pub fn build_mcp_catalog( - managed_configs: &[crate::session::managed_mcp::ManagedMcpConfig], - local_servers: &[acp::McpServer], -) -> Vec { - 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>, -) -> Vec { +pub fn build_mcp_catalog(local_servers: &[acp::McpServer]) -> Vec { let mut servers: Vec = Vec::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) for server in local_servers { let name = crate::session::mcp_servers::mcp_server_name(server).to_string(); if seen.insert(name.clone()) { - let source = if name.starts_with(MANAGED_MCP_PREFIX) { - McpServerSource::Managed - } else { - McpServerSource::Local - }; let config = match server { acp::McpServer::Http(acp::McpServerHttp { url, .. }) - | acp::McpServer::Sse(acp::McpServerSse { url, .. }) => McpServerConfig::Http { - url: url.clone(), - scope: None, - scope_id: None, - scope_name: None, - }, + | acp::McpServer::Sse(acp::McpServerSse { url, .. }) => { + McpServerConfig::Http { url: url.clone() } + } acp::McpServer::Stdio(acp::McpServerStdio { command, args, env, .. }) => McpServerConfig::Stdio { @@ -517,7 +412,7 @@ pub fn build_mcp_catalog_with_gateway_tools( servers.push(McpServerEntry { name, display_name: None, - source, + source: McpServerSource::Local, config, source_label: None, session: None, @@ -528,51 +423,17 @@ pub fn build_mcp_catalog_with_gateway_tools( 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, - gateway_tools_enabled: bool, -) -> bool { - if catalog_names.contains(name) { - return false; - } - !gateway_tools_enabled -} - fn disabled_server_placeholder_entry(name: &str) -> McpServerEntry { - let is_managed_gateway = name.starts_with(MANAGED_GATEWAY_ENTRY_PREFIX); - let source = if is_managed_gateway || name.starts_with(MANAGED_MCP_PREFIX) { - McpServerSource::Managed - } else { - McpServerSource::Local - }; - let config = if is_managed_gateway { - McpServerConfig::ManagedGateway - } else { - McpServerConfig::Stdio { + McpServerEntry { + name: name.to_owned(), + display_name: None, + source: McpServerSource::Local, + source_label: None, + config: McpServerConfig::Stdio { command: std::path::PathBuf::new(), args: 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 { enabled: false, status: None, @@ -876,13 +737,6 @@ pub async fn call_mcp_tool( // ── mcp/list handler ──────────────────────────────────────────────── 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::(args)?; 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()))) .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 acp_id = acp::SessionId::new(sid.clone()); 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_state_fut = async { - let handle = session_handle.as_ref()?; - // Auth retries belong on explicit refresh: skipping them on cached - // opens saves ~500ms when multiple OAuth servers are configured. - if !cache { - handle.retry_auth_required_servers().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 + let session_snapshot = match session_handle.as_ref() { + Some(handle) => { + // Auth retries belong on explicit refresh: skipping them on + // cached opens saves ~500ms when multiple OAuth servers are + // configured. + if !req.cache { + handle.retry_auth_required_servers().await; } - }, - session_state_fut - ); + Some(handle.get_mcp_status().await) + } + None => None, + }; let local_servers = 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_with_gateway_tools( - &managed_configs, - &local_servers, - gateway_catalog.as_ref(), - &disabled_tools, - ); + let mut servers = build_mcp_catalog(&local_servers); // Include disabled servers from config so they appear in the list // 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 = 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) { + if !catalog_names.contains(name) { servers.push(disabled_server_placeholder_entry(name)); } } 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, // which requires `req.session_id` to have been `Some`. Rather than // assert that non-local invariant with `expect` (which a future - // refactor of `session_state_fut` could silently turn into a panic - // in a request handler), use a local `if let` guard around the only - // consumer — the debug log. We emit `%sid` (Display) to match the - // sibling "session not found" log; `?req.session_id` would wrap the - // bare string as `Some("...")` and diverge from the earlier format. + // refactor could silently turn into a panic in a request handler), + // use a local `if let` guard around the only consumer — the debug + // log. if let Some(sid) = req.session_id.as_ref() { 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. 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 .configs .iter() @@ -1444,38 +1231,13 @@ async fn handle_toggle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { .get_session_handle(&acp_id) .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 (may need async I/O - // for managed config fetch). The actual config mutation happens atomically - // inside the session actor via ToggleMcpServer. + // Build the server config outside the session actor. The actual config + // mutation happens atomically inside the session actor via + // ToggleMcpServer. let server_config = if req.enabled { let cwd = agent .get_session_cwd(&acp_id) .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 { tracing::warn!( 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( vec![], &cwd, - &managed_configs, agent.plugin_registry_snapshot().as_deref(), &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); match found { 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. - let mut chars = display.chars(); + let mut chars = req.server_name.chars(); let capitalized: String = match chars.next() { Some(c) => c.to_uppercase().chain(chars).collect(), - None => display.to_string(), + None => req.server_name.clone(), }; let path = match &s.disabled_reason { Some( @@ -1525,12 +1282,6 @@ async fn handle_toggle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { _ => {} } 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 { None }; @@ -1561,25 +1312,10 @@ async fn handle_toggle_tool(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResu .get_session_handle(&acp_id) .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 - .toggle_mcp_tool(req.server_name, req.tool_name, req.enabled) - .await - } - .map_err(|e| acp::Error::internal_error().data(e.to_string()))?; + handle + .toggle_mcp_tool(req.server_name, req.tool_name, req.enabled) + .await + .map_err(|e| acp::Error::internal_error().data(e.to_string()))?; to_ext_response(Ok(McpToggleResponse { ok: true })) } @@ -1687,149 +1423,16 @@ mod tests { 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] fn test_mcp_list_response_serialization() { let resp = McpListResponse { servers: vec![ McpServerEntry { - name: "grok_com_linear".to_string(), + name: "linear".to_string(), display_name: None, - source: McpServerSource::Managed, + source: McpServerSource::Local, config: McpServerConfig::Http { 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, session: None, @@ -1860,34 +1463,11 @@ mod tests { ], }; let json = serde_json::to_value(&resp).unwrap(); - // [0] managed HTTP - assert_eq!(json["servers"][0]["source"], "managed"); + // [0] local HTTP + assert_eq!(json["servers"][0]["source"], "local"); assert_eq!(json["servers"][0]["type"], "http"); 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()); - // 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 assert_eq!(json["servers"][1]["source"], "local"); 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, - disabled_names: &[&str], - gateway_tools_enabled: bool, - ) { - let catalog_names: std::collections::HashSet = - 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> = 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] fn test_mcp_call_response_serialization() { let resp = McpCallResponse { @@ -2192,15 +1542,12 @@ mod tests { #[test] fn test_disabled_session_state_serialization() { let entry = McpServerEntry { - name: "grok_com_slack".to_string(), + name: "slack".to_string(), display_name: None, - source: McpServerSource::Managed, + source: McpServerSource::Local, source_label: None, config: McpServerConfig::Http { 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 { enabled: false, @@ -2211,8 +1558,6 @@ mod tests { }; let json = serde_json::to_value(&entry).unwrap(); assert_eq!(json["type"], "http"); - assert_eq!(json["scope"], "user"); - assert_eq!(json["scopeId"], "user-uuid-456"); assert_eq!(json["session"]["enabled"], false); assert!(json["session"].get("status").is_none()); assert!(json["session"].get("tools").is_none()); diff --git a/crates/codegen/kigi-shell/src/extensions/session_admin.rs b/crates/codegen/kigi-shell/src/extensions/session_admin.rs index f09b8c3..9e434a9 100644 --- a/crates/codegen/kigi-shell/src/extensions/session_admin.rs +++ b/crates/codegen/kigi-shell/src/extensions/session_admin.rs @@ -13,7 +13,6 @@ //! - `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_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/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_models" => handle_reload_models(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/commands/list" => handle_commands_list(agent, args).await, _ => Err(acp::Error::method_not_found()), @@ -219,11 +217,9 @@ async fn handle_update_mcp_servers(agent: &MvpAgent, args: &acp::ExtRequest) -> (h, cwd) }; - let managed = agent.get_managed_mcp_configs().await; let merged = crate::session::managed_mcp::merge_managed_mcp_servers( params.mcp_servers.clone(), &cwd, - &managed, agent.plugin_registry_handle().snapshot().as_deref(), &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())); } - let managed = agent.get_managed_mcp_configs().await; let mut updated = 0u32; for session_id in &session_ids { 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 compat = agent.cfg.borrow().compat_resolved; // Re-seed the merge with the session's original client-provided MCP - // servers (e.g. a managed connector injected at `session/new` by a - // client session binding). `merge_managed_mcp_servers` already + // servers (e.g. a server injected at `session/new` by a client + // session binding). `merge_managed_mcp_servers` already // re-reads every disk source (config.toml, plugins, ~/.claude.json, // ~/.cursor/mcp.json, .mcp.json) internally, so passing // `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( handle.initial_client_mcp_servers.clone(), &cwd, - &managed, agent.plugin_registry_handle().snapshot().as_deref(), &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())); } - let managed = agent.get_managed_mcp_configs().await; let mut updated = 0u32; for (session_id, cwd) in &session_ids { 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( handle.initial_client_mcp_servers.clone(), cwd, - &managed, agent.plugin_registry_handle().snapshot().as_deref(), &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())) } -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 async fn handle_plugins_reload(agent: &MvpAgent) -> ExtResult { diff --git a/crates/codegen/kigi-shell/src/leader/server.rs b/crates/codegen/kigi-shell/src/leader/server.rs index 0aa1735..ab0b520 100644 --- a/crates/codegen/kigi-shell/src/leader/server.rs +++ b/crates/codegen/kigi-shell/src/leader/server.rs @@ -5297,20 +5297,16 @@ mod tests { 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; 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(); let got_a = next_acp_payload(&mut reader_a).await; let got_b = next_acp_payload(&mut reader_b).await; assert!( - got_a - .as_deref() - .is_some_and(|p| p.contains("grok_com_slack")), + got_a.as_deref().is_some_and(|p| p.contains("team_slack")), "client A must receive the MCP catalog broadcast, got {got_a:?}" ); assert!( - got_b - .as_deref() - .is_some_and(|p| p.contains("grok_com_slack")), + got_b.as_deref().is_some_and(|p| p.contains("team_slack")), "client B must receive the MCP catalog broadcast, got {got_b:?}" ); cancel.cancel(); diff --git a/crates/codegen/kigi-shell/src/mcp_doctor.rs b/crates/codegen/kigi-shell/src/mcp_doctor.rs index c4d786c..abb438b 100644 --- a/crates/codegen/kigi-shell/src/mcp_doctor.rs +++ b/crates/codegen/kigi-shell/src/mcp_doctor.rs @@ -7,8 +7,6 @@ use std::sync::Arc; use kigi_tools::types::config_source::ConfigSource; use serde::Serialize; -use crate::auth::KimiCodeConfig; -use crate::session::managed_mcp; use crate::session::mcp_servers; // ── Report types ──────────────────────────────────────────────── @@ -238,85 +236,6 @@ fn discover_servers(cwd: &Path) -> (Vec, Vec) -> (ConfigSourceStatus, Vec) { - ( - ConfigSourceStatus { - path: MANAGED_SOURCE_LABEL.to_string(), - status: ConfigSourceState::Skipped { - reason: reason.into(), - }, - }, - vec![], - ) -} - -fn managed_found( - count: usize, - servers: Vec, -) -> (ConfigSourceStatus, Vec) { - ( - 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) { - 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 = 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 = servers - .into_iter() - .map(|server| DiscoveredServer { - server, - source: source.clone(), - }) - .collect(); - - managed_found(discovered.len(), discovered) -} - // ── Check functions ───────────────────────────────────────────── fn resolve_command(command: &str) -> Option { @@ -558,11 +477,7 @@ pub async fn run_auth(cwd: &Path, name: &str) -> Result { } pub async fn run_doctor(cwd: &Path, name_filter: Option<&str>) -> DoctorReport { - let (mut sources, mut discovered) = discover_servers(cwd); - - let (managed_source, managed_servers) = try_discover_managed_servers().await; - sources.push(managed_source); - discovered.extend(managed_servers); + let (mut sources, discovered) = discover_servers(cwd); let allowlist = &kigi_workspace::permission::resolution::managed_settings().mcp_allowlist; if allowlist.is_restricted() { diff --git a/crates/codegen/kigi-shell/src/session/acp_conversion.rs b/crates/codegen/kigi-shell/src/session/acp_conversion.rs index cdf5524..205678e 100644 --- a/crates/codegen/kigi-shell/src/session/acp_conversion.rs +++ b/crates/codegen/kigi-shell/src/session/acp_conversion.rs @@ -533,19 +533,6 @@ pub fn acp_tool_update( )])) .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) => { // Text includes resume handle for discoverability + meta for TUI. // 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] fn test_path_rewriter_rewrites_list_dir_raw_output() { let rw = PathRewriter::new( diff --git a/crates/codegen/kigi-shell/src/session/acp_session.rs b/crates/codegen/kigi-shell/src/session/acp_session.rs index b878ded..2a7339e 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session.rs @@ -353,166 +353,6 @@ pub(crate) fn state_is_busy(state: &State) -> bool { state.running_task.is_some() || !state.pending_inputs.is_empty() } use crate::auth::AuthManager; -#[derive(Clone)] -struct ShellManagedGatewayToolClient { - proxy_base_url: String, - auth_manager: Arc, -} -#[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. #[derive(Debug, Clone)] pub(crate) struct PreparedToolCall { @@ -858,10 +698,6 @@ pub(crate) struct SessionActor { /// flag is set short-circuits through /// [`Self::account_not_achieved_without_sampler`]. 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>>, /// Original client-provided MCP servers from session creation. /// Retained for re-merge during plugin reload. pub(crate) initial_client_mcp_servers: Vec, @@ -1315,200 +1151,6 @@ fn load_prompt_context_from_dir( #[cfg(test)] #[path = "acp_session_tests/client_hooks_tests.rs"] 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 { - 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> = - 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, /// 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. @@ -1519,7 +1161,6 @@ mod fs_injection_regression_tests; #[path = "acp_session_tests/interjection_actor_tests.rs"] mod interjection_actor_tests; #[cfg(test)] -#[cfg(test)] #[path = "acp_session_tests/permission_auto_mode_tests.rs"] mod permission_auto_mode_tests; /// 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"] mod record_response_token_usage_tests; #[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"] mod replay_buffer_send_update_tests; #[cfg(test)] @@ -1578,9 +1222,15 @@ mod rewrite_zero_turn_prefix_tests; #[path = "acp_session_tests/subagent_usage_fold_tests.rs"] mod subagent_usage_fold_tests; #[cfg(test)] +#[path = "acp_session_tests/support.rs"] +mod support; +#[cfg(test)] #[path = "acp_session_tests/turn_completion_emit_tests.rs"] mod turn_completion_emit_tests; #[cfg(test)] +#[path = "acp_session_tests/usage_categories_tests.rs"] +mod usage_categories_tests; +#[cfg(test)] mod tool_meta_stamp_tests { //! Pin the `x.ai/tool` stamps on the harness emission paths: the early //! 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"] mod feedback_turn_lookup_tests; #[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"] mod goal_reminder_subagent_rules_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/idle_resume_tests.rs"] mod idle_resume_tests; #[cfg(test)] #[path = "acp_session_tests/inline_auto_compact_flow_tests.rs"] mod inline_auto_compact_flow_tests; #[cfg(test)] +#[path = "acp_session_tests/interjection_tests.rs"] +mod interjection_tests; +#[cfg(test)] #[path = "acp_session_tests/laziness/laziness_debug_tests.rs"] mod laziness_debug_tests; #[cfg(test)] @@ -1780,9 +1448,6 @@ mod laziness_integration_tests; #[path = "acp_session_tests/load_user_prompts_tests.rs"] mod load_user_prompts_tests; #[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"] mod memory_config_tests; #[cfg(test)] @@ -1792,15 +1457,18 @@ mod parallel_dispatch_tests; #[path = "acp_session_tests/prompt_context_persistence_tests.rs"] mod prompt_context_persistence_tests; #[cfg(test)] -#[path = "acp_session_tests/reactive_managed_reauth_e2e_tests.rs"] -mod reactive_managed_reauth_e2e_tests; +#[path = "acp_session_tests/recap_display_only_tests.rs"] +mod recap_display_only_tests; #[cfg(test)] -#[path = "acp_session_tests/reactive_managed_reauth_tests.rs"] -mod reactive_managed_reauth_tests; +#[path = "acp_session_tests/reminder_policy_tests.rs"] +mod reminder_policy_tests; #[cfg(test)] #[path = "acp_session_tests/session_thread_tests.rs"] mod session_thread_tests; #[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"] mod turn_end_guard_tests; #[cfg(test)] @@ -1809,215 +1477,3 @@ mod wait_for_mcp_prefix_tests; #[cfg(test)] #[path = "acp_session_tests/web_search_e2e_tests.rs"] 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 { - 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::() - .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> = - 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::() - .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; diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/hooks_plugins.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/hooks_plugins.rs index d58ff96..f7690d5 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/hooks_plugins.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/hooks_plugins.rs @@ -900,17 +900,9 @@ impl SessionActor { // the order-sensitive `update_configs` would cause (merge order is // non-deterministic). Mirrors the `UpdateMcpServers` command handler. 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( self.initial_client_mcp_servers.clone(), session_cwd, - &managed_configs, new_registry_snapshot.as_deref(), &self.rebuild_spec.compat, ); diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/mcp.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/mcp.rs index 7aac4bf..c93be26 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/mcp.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/mcp.rs @@ -17,230 +17,6 @@ impl SessionActor { } 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 { - 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, - > = 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. /// /// 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 /// and registers its tools. 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 state = self.mcp_state.lock().await; state @@ -527,28 +300,13 @@ impl SessionActor { /// `` at a turn boundary. The `search_tool` description /// itself stays static (cacheable). 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, - >, - ) { let mcp_initialized = self.mcp_state.lock().await.is_initialized(); refresh_mcp_snapshot_and_schedule_reminder_with( self.agent.borrow().tool_bridge().clone(), Arc::clone(&self.mcp_state), - self.managed_mcp_handle.clone(), self.tool_metadata_snapshot.clone(), Arc::clone(&self.mcp_reminder_dirty), mcp_initialized, - disabled_gateway_tools, self.cursor_mcps_root(), ) .await; @@ -638,8 +396,7 @@ impl SessionActor { "connection failed".to_string() }; let retries_on_use = !mcp_state.auth_required.contains(name) - && matches!(cfg, acp::McpServer::Http(_) | acp::McpServer::Sse(_)) - && !name.starts_with(crate::session::managed_mcp::MANAGED_MCP_PREFIX); + && matches!(cfg, acp::McpServer::Http(_) | acp::McpServer::Sse(_)); let reason = if retries_on_use { format!("{base} — retries automatically on next tool call") } else { @@ -727,14 +484,9 @@ impl SessionActor { !disabled.contains(server) } /// 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`]. - /// 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 { - if server.starts_with(crate::session::managed_mcp::MANAGED_MCP_PREFIX) { - return false; - } let mcp_state = self.mcp_state.lock().await; let is_http_in_configs = mcp_state.configs.iter().any(|c| { matches!( @@ -1076,17 +828,6 @@ impl SessionActor { &mcp_server_configs, 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 .iter() @@ -1264,14 +1005,10 @@ impl SessionActor { let tool_bridge = self.agent.borrow().tool_bridge().clone(); let gateway = self.notifications.gateway.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_handshakes_done = Arc::clone(&self.mcp_handshakes_done); let session_id_owned = self.session_info.id.0.clone(); 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 = mcp_server_configs .iter() @@ -1348,13 +1085,7 @@ impl SessionActor { Ok((server_name, handles, server_start.elapsed(), timeout_sec)) } Err(e) => { - let needs_auth = if server_name - .starts_with(crate::session::managed_mcp::MANAGED_MCP_PREFIX) - { - e.is_auth_rejection() - } else { - client.has_auth() - }; + let needs_auth = client.has_auth(); tracing::warn!( server = server_name.as_str(), elapsed_ms = server_start .elapsed().as_millis() as u64, timeout_sec, error = % e, @@ -1664,11 +1395,9 @@ impl SessionActor { refresh_mcp_snapshot_and_schedule_reminder_with( tool_bridge.clone(), Arc::clone(&mcp_state_bg), - managed_mcp_handle.clone(), tool_snapshot, mcp_reminder_dirty, true, - &disabled_gateway_tools_bg, mcps_root_bg, ) .await; diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/mcp_snapshot.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/mcp_snapshot.rs index c8ec8f5..5820951 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/mcp_snapshot.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/mcp_snapshot.rs @@ -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>, -) -> 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( tool_bridge: Arc, mcp_state: Arc>, - managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle, tool_metadata_snapshot: Arc>, mcp_reminder_dirty: Arc, mcp_initialized: bool, - disabled_gateway_tools: &std::collections::HashMap>, // External harness only: per-workspace `mcps/` descriptor root. `Some` makes // this refresh also update the on-disk descriptor mirror so late-connecting // 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 mut seen_tools = std::collections::HashSet::new(); - let mut mcp_tools: Vec = all_defs + let mcp_tools: Vec = all_defs .iter() .filter(|d| d.function.name.contains("__")) .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(); - 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 = 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> = 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; } - 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); 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))) .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) .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( tool_bridge: Arc, mcp_state: Arc>, - managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle, tool_metadata_snapshot: Arc>, -) { - 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, - mcp_state: Arc>, - managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle, - tool_metadata_snapshot: Arc>, - disabled_gateway_tools: &std::collections::HashMap>, ) { refresh_mcp_snapshot_and_schedule_reminder_with( tool_bridge, mcp_state, - managed_mcp_handle, tool_metadata_snapshot, Arc::new(std::sync::atomic::AtomicBool::new(false)), false, - disabled_gateway_tools, None, ) .await; diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs index 199aa52..4261017 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs @@ -196,9 +196,6 @@ impl SessionActor { std::sync::Arc::new(tool_index), )) .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(); bridge .update_resource(kigi_tools::types::resources::PlanFilePath(plan_path)) diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs index 3f2d587..762086a 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs @@ -509,72 +509,9 @@ impl SessionActor { 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 } - async fn gather_gateway_mcp_servers( - &self, - mcps_root: Option<&std::path::Path>, - ) -> Vec { - 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::::new(); - let mut gateway_connectors: Vec = 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. /// /// Returns `None` when `display_cwd` is unset (no rewriting needed). Used diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/run_loop.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/run_loop.rs index e14c114..3002c64 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/run_loop.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/run_loop.rs @@ -559,38 +559,7 @@ pub(super) async fn run_session( tracing::warn!(server = sname.as_str(), error = % e, "Failed to persist server enabled state to config"); } let _ = respond_to .send(Ok(())); }); } SessionCommand::ToggleMcpTool { server_name, tool_name, - enabled, is_managed_gateway, respond_to } => { if is_managed_gateway { 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 + enabled, respond_to } => { let qualified = format!("{}{}{}", server_name, crate ::session::mcp_servers::MCP_TOOL_NAME_DELIMITER, tool_name,); let mut 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 = session.clone(); tokio::task::spawn_local(async move { let result = s .handle_mcp_auth_trigger(& server_name). await; let _ = respond_to - .send(result); }); } SessionCommand::GetManagedGatewayDisabledTools { - 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 + .send(result); }); } SessionCommand::RetryAuthRequiredServers { respond_to } => { let s = session .clone(); tokio::task::spawn_local(async move { s .retry_auth_required_servers(). await; let _ = respond_to.send(()); }); } SessionCommand::RefreshMcpSearchIndex => { session diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs index a005e8a..7bb4355 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs @@ -139,9 +139,6 @@ pub(crate) async fn spawn_session_actor( persisted_announcement_state: Option, memory_config: Option, feedback_flags: crate::session::feedback_manager::FeedbackFlags, - managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle, - managed_mcp_expires_at: Option>, - managed_mcp_proxy_base_url: String, session_model_id: acp::ModelId, session_yolo_mode: bool, session_auto_mode: bool, @@ -150,8 +147,6 @@ pub(crate) async fn spawn_session_actor( max_retries: Option, web_search_config: kigi_tools::implementations::WebSearchConfig, 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, write_file_enabled: bool, goal_enabled: bool, @@ -660,14 +655,6 @@ pub(crate) async fn spawn_session_actor( let context_window_tokens = context_window_override .map(|c| c.get()) .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 mut state = McpState::new_with_meta(mcp_servers.clone(), mcp_meta_config_map); 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(), backend_search: backend_tools_enabled, 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(), write_file_enabled, subagents_enabled, @@ -739,7 +724,6 @@ pub(crate) async fn spawn_session_actor( respect_gitignore, path_not_found_hints, mcp_state: mcp_state.clone(), - managed_gateway_tool_client: managed_gateway_tool_client.clone(), is_non_interactive: startup_hints.non_interactive, system_prompt_label, 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), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), 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())), mcp_announced_servers: Mutex::new( persisted_announcement_state @@ -1224,14 +1206,6 @@ pub(crate) async fn spawn_session_actor( )) .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(); session @@ -1448,7 +1422,6 @@ pub(crate) async fn spawn_session_actor( permission_handle: permissions_for_handle, attribution_callback: attribution_callback_for_handle, agent_name: agent_name_for_handle, - managed_mcp_proxy_base_url, session_default_agent_profile, allowed_subagent_types: allowed_subagent_types_for_handle, hook_registry: hook_registry_for_handle, @@ -1541,9 +1514,6 @@ pub(crate) async fn spawn_session_on_thread( persisted_announcement_state: Option, memory_config: Option, feedback_flags: crate::session::feedback_manager::FeedbackFlags, - managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle, - managed_mcp_expires_at: Option>, - managed_mcp_proxy_base_url: String, session_model_id: acp::ModelId, session_yolo_mode: bool, session_auto_mode: bool, @@ -1552,8 +1522,6 @@ pub(crate) async fn spawn_session_on_thread( max_retries: Option, web_search_config: kigi_tools::implementations::WebSearchConfig, 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, write_file_enabled: bool, goal_enabled: bool, @@ -1688,9 +1656,6 @@ pub(crate) async fn spawn_session_on_thread( persisted_announcement_state, memory_config, feedback_flags, - managed_mcp_handle, - managed_mcp_expires_at, - managed_mcp_proxy_base_url, session_model_id, session_yolo_mode, session_auto_mode, @@ -1699,8 +1664,6 @@ pub(crate) async fn spawn_session_on_thread( max_retries, web_search_config, web_fetch_config, - image_gen_config, - video_gen_config, app_builder_deployer_config, write_file_enabled, goal_enabled, diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/tool_calls.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/tool_calls.rs index da80b97..b372725 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/tool_calls.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/tool_calls.rs @@ -480,31 +480,13 @@ impl SessionActor { } }); 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] .take() .expect("dispatch index should match an approved slot exactly once"); self.signals_handle().record_tool_call(&prepared.tool_name); let tool_start = self.events.tool_started(prepared.tool_name.clone()); let mut post_tool_use_result: Option = 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 { Ok(tool_result) => tool_result.prompt_text.len() as i64, Err(_) => 0, @@ -745,12 +727,6 @@ impl SessionActor { } let mcp_parts = parse_mcp_tool_name(&call.function.name); 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() { match self.mcp_strategy { McpInitStrategy::Blocking => { @@ -1483,33 +1459,6 @@ impl SessionActor { 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) => ( mcp_tool.tool_name.to_owned(), acp::ToolKind::Other, diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/tool_dispatch.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/tool_dispatch.rs index 5adde16..1475a05 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/tool_dispatch.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/tool_dispatch.rs @@ -128,7 +128,7 @@ pub(super) fn resolve_session_shell() -> String { } /// 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. pub(crate) const HTTP_STATUS_DETAILS_KEY: &str = "status"; diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/turn.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/turn.rs index 3a8664e..ec7b649 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/turn.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/turn.rs @@ -574,16 +574,6 @@ impl SessionActor { .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 origin = super::super::PromptOrigin::from_prompt_id(prompt_id); diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs index 72afcf8..94d8459 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs @@ -232,8 +232,6 @@ async fn persist_ack_waits_for_disk_flush_before_success() { goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), 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![], tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), 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), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), 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![], tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), 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), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), 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![], tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), 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), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), 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![], tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), mcp_announced_servers: Mutex::new(HashMap::new()), diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/fs_injection_regression_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/fs_injection_regression_tests.rs index 0cfabb9..ef3730d 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/fs_injection_regression_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/fs_injection_regression_tests.rs @@ -58,8 +58,6 @@ async fn tool_bridge_routes_writes_through_injected_fs() { web_search_config: Default::default(), web_fetch_config: Default::default(), lsp: None, - image_gen_config: Default::default(), - video_gen_config: Default::default(), app_builder_deployer_config: Default::default(), api_key_provider: None, attribution_callback: None, diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs index ec126e6..7662c62 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs @@ -261,8 +261,6 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() { goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), 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![], tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), mcp_announced_servers: Mutex::new(HashMap::new()), diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs index f466d3f..e25fa30 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs @@ -190,8 +190,6 @@ async fn create_test_actor( goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), 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![], tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), 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), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), 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![], tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), 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), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), 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![], tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), mcp_announced_servers: Mutex::new(HashMap::new()), diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs index 057ae00..f2f7d49 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs @@ -252,8 +252,6 @@ async fn create_test_actor_with_memory( goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), 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![], tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), mcp_announced_servers: Mutex::new(HashMap::new()), diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_e2e_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_e2e_tests.rs deleted file mode 100644 index 6a1e885..0000000 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_e2e_tests.rs +++ /dev/null @@ -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, - /// Count of managed-config fetches (`GET /mcp/configs`) so a cooldown-gated - /// attempt can be proven to skip the network entirely. - config_fetches: Arc, - /// 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) -> 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, - axum::Json(req): axum::Json, -) -> 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::>()), - ) - .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) -> (String, String, Arc) { - 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, -) -> (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, -) -> Vec { - 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::(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; -} diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_tests.rs deleted file mode 100644 index 20f4a31..0000000 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_tests.rs +++ /dev/null @@ -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; -} diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs index 2c310a2..3513396 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs @@ -198,8 +198,6 @@ pub(super) async fn make_replay_send_update_fixture() -> ReplaySendUpdateFixture goal_plan_reconciled: std::sync::atomic::AtomicBool::new(false), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), 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![], tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), mcp_announced_servers: Mutex::new(HashMap::new()), diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs index f841c97..191b288 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs @@ -89,8 +89,6 @@ async fn test_agent_from_config( web_search_config: Default::default(), web_fetch_config: Default::default(), lsp: None, - image_gen_config: Default::default(), - video_gen_config: Default::default(), app_builder_deployer_config: Default::default(), api_key_provider: 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), pending_classifier_completions: parking_lot::Mutex::new(VecDeque::new()), 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![], tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), mcp_announced_servers: Mutex::new(HashMap::new()), diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/media_gen_auth_retry_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/tool_auth_retry_tests.rs similarity index 93% rename from crates/codegen/kigi-shell/src/session/acp_session_tests/media_gen_auth_retry_tests.rs rename to crates/codegen/kigi-shell/src/session/acp_session_tests/tool_auth_retry_tests.rs index 85f8096..296678a 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/media_gen_auth_retry_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/tool_auth_retry_tests.rs @@ -77,7 +77,7 @@ fn err(msg: &str) -> Result { )) } -/// 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 /// the structured status-code path rather than the string fallback. fn http_err(status: u16, msg: &str) -> Result { @@ -98,14 +98,14 @@ fn is_auth_tool_error_classification() { // (expected, error) — covers every branch + a sample of negatives // a careless edit could plausibly break. 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 // matches the status code, not the rendered string. ( true, kigi_tool_runtime::ToolError::new( 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( 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. ( 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. ( false, kigi_tool_runtime::ToolError::not_found( - kigi_tool_protocol::ToolId::new("image_gen").expect("valid"), - "Tool not found: image_gen", + kigi_tool_protocol::ToolId::new("no_such_tool").expect("valid"), + "Tool not found: no_such_tool", ), ), // 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 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); async move { 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 { 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", || { 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; diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/web_search_e2e_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/web_search_e2e_tests.rs index 67e55cd..80b58b2 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/web_search_e2e_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/web_search_e2e_tests.rs @@ -37,8 +37,6 @@ async fn web_search_errors_when_disabled() { web_search_config: kigi_tools::implementations::web_search::WebSearchConfig::Disabled, web_fetch_config: Default::default(), lsp: None, - image_gen_config: Default::default(), - video_gen_config: Default::default(), app_builder_deployer_config: Default::default(), api_key_provider: None, attribution_callback: None, diff --git a/crates/codegen/kigi-shell/src/session/agent_rebuild.rs b/crates/codegen/kigi-shell/src/session/agent_rebuild.rs index 858bd7a..feba196 100644 --- a/crates/codegen/kigi-shell/src/session/agent_rebuild.rs +++ b/crates/codegen/kigi-shell/src/session/agent_rebuild.rs @@ -46,11 +46,9 @@ use kigi_agent::{Agent, AgentBuilder, CompactionPolicy, ReminderPolicy}; use kigi_tools::computer::types::{AsyncFileSystem, TerminalBackend}; 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::image_gen::ImageGenConfig; use kigi_tools::implementations::grok_build::task::types::{ 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::lsp::LspBackend; use kigi_tools::implementations::web_search::WebSearchConfig; @@ -95,8 +93,6 @@ pub(crate) struct AgentRebuildSpec { pub web_search_config: WebSearchConfig, pub backend_search: bool, pub web_fetch_config: WebFetchConfig, - pub image_gen_config: ImageGenConfig, - pub video_gen_config: VideoGenConfig, pub app_builder_deployer_config: AppBuilderDeployerConfig, pub write_file_enabled: bool, pub subagents_enabled: bool, @@ -125,7 +121,6 @@ pub(crate) struct AgentRebuildSpec { pub respect_gitignore: bool, pub path_not_found_hints: bool, pub mcp_state: Arc>, - pub managed_gateway_tool_client: Option, pub is_non_interactive: bool, pub system_prompt_label: String, pub owner_session_id: Option, @@ -190,8 +185,6 @@ impl AgentRebuildSpec { web_search_config, backend_search, web_fetch_config, - image_gen_config, - video_gen_config, app_builder_deployer_config, write_file_enabled, subagents_enabled, @@ -218,7 +211,6 @@ impl AgentRebuildSpec { respect_gitignore, path_not_found_hints, mcp_state, - managed_gateway_tool_client, is_non_interactive, system_prompt_label, owner_session_id, @@ -244,8 +236,6 @@ impl AgentRebuildSpec { .with_state_path(bridge_state_path.clone()) .with_web_search_config(web_search_config.clone()) .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_web_fetch_config(web_fetch_config.clone()) .with_write_file_enabled(*write_file_enabled) @@ -351,9 +341,6 @@ impl AgentRebuildSpec { *path_not_found_hints, )) .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; agent @@ -392,8 +379,6 @@ pub(crate) fn test_rebuild_spec_default() -> Arc { web_search_config: WebSearchConfig::default(), backend_search: false, web_fetch_config: WebFetchConfig::Disabled, - image_gen_config: ImageGenConfig::default(), - video_gen_config: VideoGenConfig::default(), app_builder_deployer_config: AppBuilderDeployerConfig::default(), write_file_enabled: true, subagents_enabled: false, @@ -422,7 +407,6 @@ pub(crate) fn test_rebuild_spec_default() -> Arc { mcp_state: Arc::new(tokio::sync::Mutex::new( crate::session::mcp_servers::McpState::new(vec![]), )), - managed_gateway_tool_client: None, is_non_interactive: false, system_prompt_label: kigi_agent::DEFAULT_SYSTEM_PROMPT_LABEL.to_string(), owner_session_id: Some("test-session".to_string()), diff --git a/crates/codegen/kigi-shell/src/session/commands.rs b/crates/codegen/kigi-shell/src/session/commands.rs index 8ff8b1f..068b888 100644 --- a/crates/codegen/kigi-shell/src/session/commands.rs +++ b/crates/codegen/kigi-shell/src/session/commands.rs @@ -346,7 +346,7 @@ pub enum SessionCommand { server_name: String, enabled: bool, /// 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. server_config: Option, respond_to: oneshot::Sender>, @@ -357,17 +357,12 @@ pub enum SessionCommand { server_name: String, tool_name: String, enabled: bool, - is_managed_gateway: bool, respond_to: oneshot::Sender>, }, /// Read MCP status: which servers are configured, which clients are healthy, what tools. GetMcpStatus { respond_to: oneshot::Sender, }, - GetManagedGatewayDisabledTools { - respond_to: - oneshot::Sender>>, - }, /// Snapshot the session's live MCP client pool for subagent inheritance. SnapshotMcpPool { respond_to: oneshot::Sender>, diff --git a/crates/codegen/kigi-shell/src/session/compaction.rs b/crates/codegen/kigi-shell/src/session/compaction.rs index 2ecec33..54746b4 100644 --- a/crates/codegen/kigi-shell/src/session/compaction.rs +++ b/crates/codegen/kigi-shell/src/session/compaction.rs @@ -2288,8 +2288,6 @@ mod inline_auto_compact_flow_tests { std::collections::VecDeque::new(), ), 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![], tool_metadata_snapshot: Arc::new(std::sync::Mutex::new(Default::default())), mcp_announced_servers: parking_lot::Mutex::new(std::collections::HashMap::new()), diff --git a/crates/codegen/kigi-shell/src/session/handle.rs b/crates/codegen/kigi-shell/src/session/handle.rs index eea3d30..72d09b1 100644 --- a/crates/codegen/kigi-shell/src/session/handle.rs +++ b/crates/codegen/kigi-shell/src/session/handle.rs @@ -130,7 +130,6 @@ pub struct SessionHandle { pub attribution_callback: Option, /// The agent definition name for this session. pub agent_name: String, - pub managed_mcp_proxy_base_url: String, pub session_default_agent_profile: Option, /// Subagent types this agent can spawn (from Agent(t1, t2) in tools). pub allowed_subagent_types: Option>, @@ -398,25 +397,6 @@ impl SessionHandle { server_name: String, tool_name: String, 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> { let (tx, rx) = oneshot::channel(); if self @@ -425,7 +405,6 @@ impl SessionHandle { server_name, tool_name, enabled, - is_managed_gateway, respond_to: tx, }) .is_err() @@ -435,17 +414,6 @@ impl SessionHandle { rx.await .map_err(|_| agent_client_protocol::Error::internal_error().data("session closed"))? } - pub async fn managed_gateway_disabled_tool_names(&self) -> HashMap> { - 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) { let (tx, rx) = oneshot::channel(); if self diff --git a/crates/codegen/kigi-shell/src/session/managed_mcp.rs b/crates/codegen/kigi-shell/src/session/managed_mcp.rs index 530f4ad..4b769aa 100644 --- a/crates/codegen/kigi-shell/src/session/managed_mcp.rs +++ b/crates/codegen/kigi-shell/src/session/managed_mcp.rs @@ -1,6 +1,6 @@ -//! Shell-side managed MCP: merges MCP server sources, then injects managed -//! OAuth headers, and binds the extracted credential/catalog machinery to -//! shell's auth manager. +//! Shell-side MCP server merging plus local managed-settings policy +//! (`managed-settings.json` allow/deny lists — local policy files, not a +//! remote service). //! //! Merge layers are applied in order; later `insert()` beats earlier //! `or_insert()`: @@ -9,59 +9,14 @@ //! - ~/.claude.json — `or_insert` (imported user/local MCP servers) //! - `.mcp.json` — `or_insert` (team baseline) //! - 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::sync::Arc; use agent_client_protocol as acp; -/// Build a [`RefreshContext`] whose token provider resolves fresh tokens from -/// `auth_manager`; the extracted refresh task never sees the auth manager -/// itself, only the closure. -fn refresh_context( - proxy_base_url: String, - auth_manager: Arc, -) -> 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, -) -> Vec { - 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 +/// Normalize a URL for dedup purposes (trailing slash dropped). +pub fn normalize_url(url: &str) -> String { + url.trim_end_matches('/').to_string() } /// Dedup key for the merge map: normalized URL for Http/Sse, name for Stdio. @@ -88,27 +43,19 @@ pub(crate) fn mcp_server_name(s: &acp::McpServer) -> &str { pub fn merge_managed_mcp_servers( client_mcp_servers: Vec, cwd: &std::path::Path, - managed_configs: &[ManagedMcpConfig], plugin_registry: Option<&kigi_agent::plugins::PluginRegistry>, compat: &kigi_tools::types::compat::CompatConfig, ) -> Vec { - merge_managed_mcp_servers_with_policy( - client_mcp_servers, - cwd, - managed_configs, - plugin_registry, - compat, - ) - .into_iter() - .filter(|s| s.disabled_reason.is_none()) - .map(|s| s.server) - .collect() + merge_managed_mcp_servers_with_policy(client_mcp_servers, cwd, plugin_registry, compat) + .into_iter() + .filter(|s| s.disabled_reason.is_none()) + .map(|s| s.server) + .collect() } pub fn merge_managed_mcp_servers_with_policy( client_mcp_servers: Vec, cwd: &std::path::Path, - managed_configs: &[ManagedMcpConfig], plugin_registry: Option<&kigi_agent::plugins::PluginRegistry>, compat: &kigi_tools::types::compat::CompatConfig, ) -> Vec { @@ -125,8 +72,6 @@ pub fn merge_managed_mcp_servers_with_policy( let disabled = crate::util::config::disabled_mcp_server_names(cwd); let mut merged: Vec = 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 // iteration order). Downstream equality checks (`mcp_servers_equal`, used // 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() } -/// 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, - managed_configs: &[ManagedMcpConfig], - disabled_names: &std::collections::HashSet, -) { - if managed_configs.is_empty() { - return; - } - - let existing_names: std::collections::HashSet = merged - .iter() - .map(|s| mcp_server_name(s).to_owned()) - .collect(); - let mut seen_display_names: std::collections::HashSet = - 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( mcp_path: &std::path::Path, plugin_name: &str, @@ -496,41 +388,10 @@ pub fn merge_plugin_oauth_into( mod tests { 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 { 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 /// `session/new`) exists in no on-disk config and no managed catalog — /// the merge must keep it. Config hot-reload handlers @@ -549,7 +410,7 @@ mod tests { )]; let cwd = empty_cwd(); 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!( merged.iter().any(|s| matches!( s, @@ -688,112 +549,6 @@ mod tests { 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 = - ["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] fn lower_precedence_http_servers_are_blocked_by_toml_name_claims() { let cwd = tempfile::tempdir().unwrap(); @@ -821,7 +576,7 @@ enabled = false .unwrap(); 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!( !merged.iter().any(|server| matches!( server, @@ -859,7 +614,7 @@ enabled = false ) .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!( !merged.iter().any(|s| mcp_server_name(s) == "projsrv"), "untrusted workspace must drop its repo-local MCP server" @@ -871,7 +626,7 @@ enabled = false let trusted = repo_with_project_server(); 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!( merged.iter().any(|s| mcp_server_name(s) == "projsrv"), "trusted workspace must keep its repo-local MCP server" @@ -940,8 +695,7 @@ enabled = false let disabled: std::collections::HashSet = ["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. assert!( disabled.contains("test-server"), diff --git a/crates/codegen/kigi-shell/src/session/mcp_descriptors.rs b/crates/codegen/kigi-shell/src/session/mcp_descriptors.rs index 4cb546d..0e3f44a 100644 --- a/crates/codegen/kigi-shell/src/session/mcp_descriptors.rs +++ b/crates/codegen/kigi-shell/src/session/mcp_descriptors.rs @@ -7,25 +7,14 @@ //! 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 //! 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. -use std::collections::{BTreeMap, BTreeSet, HashSet}; use std::path::{Path, PathBuf}; use std::sync::Arc; 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: `/`. Uses the /// sanitizer shared with `kigi-mcp` so the advertised folder matches disk. 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, - gateway_connectors: Vec, - protected_connectors: HashSet, -) { - let mut files_by_connector: BTreeMap)>> = 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 = 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 = 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::(&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, - 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)] mod tests { use super::*; #[test] 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(""), "_"); 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"); 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()); - } } diff --git a/crates/codegen/kigi-shell/src/session/mcp_dispatcher.rs b/crates/codegen/kigi-shell/src/session/mcp_dispatcher.rs index a574d92..f7ef3a2 100644 --- a/crates/codegen/kigi-shell/src/session/mcp_dispatcher.rs +++ b/crates/codegen/kigi-shell/src/session/mcp_dispatcher.rs @@ -50,7 +50,6 @@ use tokio::sync::Mutex as TokioMutex; use tokio::sync::mpsc::UnboundedReceiver; use crate::extensions::mcp::McpServerSource; -use crate::session::managed_mcp::MANAGED_MCP_PREFIX; /// Tumbling-window coalescing period. See module doc. 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 { /// Owning session id. pub session_id: String, - /// MCP server name (`grok_com_linear`, `github`, ...). + /// MCP server name (`github`, ...). pub name: String, - /// `managed` (sourced from cli-chat-proxy / `grok_com_` prefix) - /// or `local` (user `.kigi/config.toml`). + /// Always `local` (user `.kigi/config.toml` and friends). pub source: McpServerSource, /// Current status — see [`McpServerStatus`]. pub status: McpServerStatus, @@ -131,22 +129,12 @@ pub enum McpServerStatusReason { RestartSucceeded, /// The auto-restart path exhausted retries. 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 -/// existing convention used by `build_mcp_catalog` and friends: -/// names with the `MANAGED_MCP_PREFIX` prefix are managed. -pub fn classify_source(name: &str) -> McpServerSource { - if name.starts_with(MANAGED_MCP_PREFIX) { - McpServerSource::Managed - } else { - McpServerSource::Local - } +/// Build [`McpServerSource`] from a server name. All servers are +/// locally configured. +pub fn classify_source(_name: &str) -> McpServerSource { + McpServerSource::Local } /// State for the dispatcher's "intentional teardown" tracking. @@ -392,22 +380,6 @@ pub fn build_payload( McpServerStatusReason::TransportClosed, 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, .. }) => { let detail = reason.clone(); ( @@ -586,7 +558,7 @@ pub fn recoverable_http_servers( ) }) .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() } @@ -984,58 +956,9 @@ mod tests { ); } - /// Contract: a managed connector whose handshake is rejected for - /// auth reasons surfaces as `NeedsAuth`/`auth_expired` ("visit - /// grok.com"), NOT a generic `Unavailable`. Keys on the shared - /// `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). + /// A local server whose handshake error happens to contain auth + /// wording stays `Unavailable` (local auth recovery is the OAuth + /// path, not this one). #[test] fn local_handshake_auth_rejection_stays_unavailable() { let key = ("github".to_string(), McpClientEventKind::HandshakeFailed); @@ -1049,32 +972,6 @@ mod tests { 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. /// /// Locks the camelCase field naming, lowercase enum values, and @@ -1393,15 +1290,14 @@ 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. #[test] - fn recoverable_http_servers_excludes_managed_stdio_and_disabled() { + fn recoverable_http_servers_excludes_stdio_and_disabled() { let configs = vec![ http_cfg("http-mcp-server"), - http_cfg("grok_com_slack"), // managed - http_cfg("admin_off"), // disabled - stdio_cfg("local_stdio"), // stdio + http_cfg("admin_off"), // disabled + stdio_cfg("local_stdio"), // stdio ]; let disabled: HashSet = ["admin_off".to_string()].into_iter().collect(); let got = recoverable_http_servers(&configs, &disabled); diff --git a/crates/codegen/kigi-shell/src/session/mcp_servers.rs b/crates/codegen/kigi-shell/src/session/mcp_servers.rs index 7646f42..082fbad 100644 --- a/crates/codegen/kigi-shell/src/session/mcp_servers.rs +++ b/crates/codegen/kigi-shell/src/session/mcp_servers.rs @@ -50,14 +50,7 @@ pub fn build_config_resolved_event( .map(|c| kigi_file_utils::events::McpConfigServer { name: inner::mcp_server_name(c).to_string(), transport: inner::mcp_transport_str(c).to_string(), - source: if inner::mcp_server_name(c) - .starts_with(crate::session::managed_mcp::MANAGED_MCP_PREFIX) - { - "managed" - } else { - "local" - } - .to_string(), + source: "local".to_string(), }) .collect(); kigi_file_utils::events::Event::McpConfigResolved { servers, disabled } diff --git a/crates/codegen/kigi-shell/src/session/slash_commands.rs b/crates/codegen/kigi-shell/src/session/slash_commands.rs index 740af79..be1090d 100644 --- a/crates/codegen/kigi-shell/src/session/slash_commands.rs +++ b/crates/codegen/kigi-shell/src/session/slash_commands.rs @@ -1697,11 +1697,11 @@ mod tests { #[test] 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); assert_eq!( serde_json::Value::Object(v), - serde_json::json!({"tools": ["scheduler_create", "image_gen"]}) + serde_json::json!({"tools": ["scheduler_create", "web_search"]}) ); } diff --git a/crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs b/crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs index bfe24b0..cfe322c 100644 --- a/crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs +++ b/crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs @@ -83,8 +83,6 @@ pub(crate) fn ctx_with_toggle(toggle: HashMap) -> SubagentSpawnCon memory_config: None, web_search_config: Default::default(), web_fetch_config: Default::default(), - image_gen_config: Default::default(), - video_gen_config: Default::default(), app_builder_deployer_config: Default::default(), write_file_enabled: true, goal_enabled: false, @@ -128,8 +126,6 @@ pub(crate) fn ctx_with_toggle(toggle: HashMap) -> SubagentSpawnCon parent_model_agent_type: None, allowed_subagent_types: None, 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_tool_snapshot: None, parent_skills: None, diff --git a/crates/codegen/kigi-shell/src/util/config/mcp.rs b/crates/codegen/kigi-shell/src/util/config/mcp.rs index c5d4fad..b62afa5 100644 --- a/crates/codegen/kigi-shell/src/util/config/mcp.rs +++ b/crates/codegen/kigi-shell/src/util/config/mcp.rs @@ -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() } -pub const MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY: &str = "__managed_gateway_connectors"; - /// 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 -/// 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<()> { let path = config_path(); let mut root: TomlValue = match tokio::fs::read_to_string(&path).await { diff --git a/crates/codegen/kigi-tools-api/src/slash_commands.rs b/crates/codegen/kigi-tools-api/src/slash_commands.rs index 8bfdb8f..4dca7df 100644 --- a/crates/codegen/kigi-tools-api/src/slash_commands.rs +++ b/crates/codegen/kigi-tools-api/src/slash_commands.rs @@ -1,6 +1,5 @@ -//! Canonical slash-command wording (`/loop`, `/imagine`, `/imagine-video`, `/goal`), -//! shared by every front-end (Grok Build shell/pager and other hosts) so -//! expansions cannot drift. +//! Canonical slash-command wording (`/loop`, `/goal`), shared by every +//! front-end (shell/pager and other hosts) so expansions cannot drift. /// Canonical tool name advertised by the scheduler create tool. Gating code /// (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 \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 \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 (1–2 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 1–2 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 GOAL_COMMAND_NAME: &str = "goal"; @@ -170,22 +84,6 @@ pub fn goal_instruction(objective: &str) -> String { mod tests { 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] fn instruction_carries_args_and_contract_tokens() { let text = loop_schedule_instruction("every 30 minutes do x"); diff --git a/crates/codegen/kigi-tools/schema/tool_meta.schema.json b/crates/codegen/kigi-tools/schema/tool_meta.schema.json index 405f835..d06fa94 100644 --- a/crates/codegen/kigi-tools/schema/tool_meta.schema.json +++ b/crates/codegen/kigi-tools/schema/tool_meta.schema.json @@ -36,7 +36,7 @@ ], "definitions": { "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" }, "ToolNamespace": { diff --git a/crates/codegen/kigi-tools/src/attribution.rs b/crates/codegen/kigi-tools/src/attribution.rs index 50fa901..0015d56 100644 --- a/crates/codegen/kigi-tools/src/attribution.rs +++ b/crates/codegen/kigi-tools/src/attribution.rs @@ -8,18 +8,12 @@ pub const SENT_BEARER_PREFIX_LEN: usize = 12; /// Which tool endpoint produced the 401. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ToolConsumer { - ImageGen, - VideoGenStart, - VideoGenPoll, WebSearch, } impl ToolConsumer { pub fn as_str(self) -> &'static str { match self { - Self::ImageGen => "ImageGen", - Self::VideoGenStart => "VideoGen.start", - Self::VideoGenPoll => "VideoGen.poll", Self::WebSearch => "WebSearch", } } @@ -95,9 +89,6 @@ mod tests { #[test] 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"); } } diff --git a/crates/codegen/kigi-tools/src/implementations/grok_build/image_edit/mod.rs b/crates/codegen/kigi-tools/src/implementations/grok_build/image_edit/mod.rs deleted file mode 100644 index b8505a3..0000000 --- a/crates/codegen/kigi-tools/src/implementations/grok_build/image_edit/mod.rs +++ /dev/null @@ -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, -) -> Result<(Vec, &'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, ¶ms).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 { - 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 { - 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::() { - 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 { - 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 = 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, - - #[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 { - 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 { - 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::()?.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::() - .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 = 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::()?.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 { - 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 { - 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}"); - } -} diff --git a/crates/codegen/kigi-tools/src/implementations/grok_build/image_gen/mod.rs b/crates/codegen/kigi-tools/src/implementations/grok_build/image_gen/mod.rs deleted file mode 100644 index d045b53..0000000 --- a/crates/codegen/kigi-tools/src/implementations/grok_build/image_gen/mod.rs +++ /dev/null @@ -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. ``). -//! -//! 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 `/images/.jpg` -//! where `` 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, - /// 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, - /// 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, - ) -> Result { - 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, - ) -> Self { - self.attribution_callback = callback; - self - } - - pub(crate) async fn current_bearer(&self) -> Option { - 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, 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, - 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, - /// `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, -} - -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, -} - -#[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 { - 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 { - use crate::types::tool_metadata::shared_resources; - let resources = shared_resources(&ctx)?; - - let client = { - let res = resources.lock().await; - res.require::()?.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::()?.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:?}"), - } - } -} diff --git a/crates/codegen/kigi-tools/src/implementations/grok_build/mod.rs b/crates/codegen/kigi-tools/src/implementations/grok_build/mod.rs index 294823a..4e0eff7 100644 --- a/crates/codegen/kigi-tools/src/implementations/grok_build/mod.rs +++ b/crates/codegen/kigi-tools/src/implementations/grok_build/mod.rs @@ -15,8 +15,6 @@ pub mod deploy_app; pub mod enter_plan_mode; pub mod exit_plan_mode; pub mod grep; -pub mod image_edit; -pub mod image_gen; pub mod kill_task; pub mod list_dir; pub mod lsp; @@ -29,7 +27,6 @@ pub mod task; pub mod task_output; pub mod todo; pub mod update_goal; -pub mod video_gen; pub mod web_fetch; pub mod web_search; 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 exit_plan_mode::ExitPlanModeTool; 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 list_dir::ListDirTool; pub use lsp::LspTool; @@ -58,10 +50,5 @@ pub use task::TaskTool; pub use task_output::{GetTerminalCommandOutputTool, TaskOutputTool, WaitTasksTool}; pub use todo::TodoWriteTool; 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_search::WebSearchTool; diff --git a/crates/codegen/kigi-tools/src/implementations/grok_build/task/types.rs b/crates/codegen/kigi-tools/src/implementations/grok_build/task/types.rs index 4a8e259..ed04108 100644 --- a/crates/codegen/kigi-tools/src/implementations/grok_build/task/types.rs +++ b/crates/codegen/kigi-tools/src/implementations/grok_build/task/types.rs @@ -237,10 +237,6 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode { ToolKind::MemoryGet, ToolKind::WebSearch, ToolKind::WebFetch, - ToolKind::ImageGen, - ToolKind::VideoGen, - ToolKind::ImageToVideo, - ToolKind::ReferenceToVideo, ToolKind::BackgroundTaskAction, ToolKind::KillTaskAction, ToolKind::Task, @@ -285,10 +281,6 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode { ToolKind::MemoryGet, ToolKind::WebSearch, ToolKind::WebFetch, - ToolKind::ImageGen, - ToolKind::VideoGen, - ToolKind::ImageToVideo, - ToolKind::ReferenceToVideo, ToolKind::BackgroundTaskAction, ToolKind::KillTaskAction, ToolKind::Task, diff --git a/crates/codegen/kigi-tools/src/implementations/grok_build/video_gen/mod.rs b/crates/codegen/kigi-tools/src/implementations/grok_build/video_gen/mod.rs deleted file mode 100644 index 649a7e1..0000000 --- a/crates/codegen/kigi-tools/src/implementations/grok_build/video_gen/mod.rs +++ /dev/null @@ -1,1536 +0,0 @@ -//! Video generation module. Hosts the shared [`VideoGenClient`] and the -//! `image_to_video` and `reference_to_video` tools, which generate videos via -//! the xAI Video Generation API and save them to the local filesystem so the -//! model can reference them in code (e.g. `