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

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

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

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

121 files changed, 8 deleted. Gates: workspace check/clippy 0/0, fmt,
deny ok; suites green (tools 2554, shell 4862, tui 6608, workspace
1042). Remaining grok.com strings live only in the auth-method ids and
changelog archives (§9/M3 sweep).
This commit is contained in:
ZacharyZhang-NY
2026-07-17 23:45:05 -04:00
parent fa75eb139a
commit 5e4e24db99
120 changed files with 301 additions and 11327 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/codegen/kigi-shell/Cargo.toml
-1
View File
@@ -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" }
crates/codegen/kigi-shell/skills/imagine/SKILL.md
-131
View File
@@ -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.
crates/codegen/kigi-shell/src/agent/app.rs
-17
View File
@@ -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]
crates/codegen/kigi-shell/src/agent/config.rs
-232
View File
@@ -301,9 +301,6 @@ pub struct Requirements {
pub tool_search: Constrained<bool>,
pub web_fetch: Constrained<bool>,
pub ask_user_question: Constrained<bool>,
pub image_gen: Constrained<bool>,
pub image_edit: Constrained<bool>,
pub video_gen: Constrained<bool>,
pub write_file: Constrained<bool>,
/// Voice dictation (STT). Pin via requirements/managed `[features] voice_mode`.
pub voice_mode: Constrained<bool>,
@@ -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<kigi_tools::implementations::grok_build::video_gen::ZdrVideoOutputS3Config>,
/// 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<bool> {
BoolFlag::env("KIGI_IMAGE_GEN")
.requirement(self.requirements.image_gen.pinned())
.default(true)
.resolve()
}
/// `image_edit` tool gate.
///
/// The remote settings `imagine_tools_disabled` denylist is authoritative:
/// when it lists `image_edit`, the tool is force-removed and local
/// env/config can't re-enable it. A managed requirement pin still outranks
/// it; otherwise the tool defaults on and is overridable via
/// `KIGI_IMAGE_EDIT`.
pub(crate) fn resolve_image_edit(&self) -> Resolved<bool> {
use kigi_tools::implementations::grok_build::IMAGE_EDIT_TOOL_NAME;
if let Some(pinned) = self.requirements.image_edit.pinned() {
return Resolved::new(pinned, ConfigSource::Requirement);
}
if self
.remote_settings
.as_ref()
.is_some_and(|s| s.imagine_tool_disabled(IMAGE_EDIT_TOOL_NAME))
{
return Resolved::new(false, ConfigSource::Remote);
}
BoolFlag::env("KIGI_IMAGE_EDIT").default(true).resolve()
}
/// Optional Imagine model override for `image_gen`. When set (non-empty),
/// `image_gen` calls this model slug instead of the default quality model.
/// Precedence: env `KIGI_IMAGE_GEN_MODEL_OVERRIDE` > `[features]
/// image_gen_model_override` config > remote settings `image_gen_model_override`.
/// `None` → default model (`grok-imagine-image-quality`).
pub(crate) fn resolve_image_gen_model_override(&self) -> Option<String> {
resolve_string_flag(
None,
"KIGI_IMAGE_GEN_MODEL_OVERRIDE",
self.features.image_gen_model_override.as_deref(),
self.remote_settings
.as_ref()
.and_then(|s| s.image_gen_model_override.as_deref()),
)
.map(|r| r.value)
}
/// Goal mode (`/goal`) master switch. Default ON: deployments that can't
/// 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<bool>,
/// Video generation tool. `None` = defer to remote settings / env / default (false).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub video_gen: Option<bool>,
/// `image_gen` Imagine model override. `None`/empty = defer to remote settings
/// (`image_gen_model_override`) / env / default (`grok-imagine-image-quality`).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub image_gen_model_override: Option<String>,
/// Write file tool. `None` = defer to remote settings / env / default (true).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub write_file: Option<bool>,
@@ -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<EnvGuard> {
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]
crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs
+4 -16
View File
@@ -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<crate::extensions::mcp::McpServerEntry> = 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
}
crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs
+6 -226
View File
@@ -84,123 +84,6 @@ impl MvpAgent {
}
})
}
fn has_managed_mcp_auth(&self) -> bool {
self.auth_manager
.current_or_expired()
.is_some_and(|a| a.is_session_auth())
}
/// Requires feature flag AND xAI authentication (OIDC or legacy WebLogin).
pub(super) fn can_fetch_managed_mcps(&self) -> bool {
let cfg = self.cfg.borrow();
cfg.managed_mcps_enabled && !cfg.managed_mcp_gateway_tools_enabled
&& self.has_managed_mcp_auth()
}
fn can_fetch_managed_mcp_gateway_tools(&self) -> bool {
self.cfg.borrow().managed_mcp_gateway_tools_enabled
&& self.has_managed_mcp_auth()
}
pub async fn get_managed_mcp_configs(
&self,
) -> Vec<crate::session::managed_mcp::ManagedMcpConfig> {
if !self.can_fetch_managed_mcps() {
return vec![];
}
let proxy_url = self.cfg.borrow().endpoints.proxy_url();
crate::session::managed_mcp::fetch_managed_mcp_configs(
&self.managed_mcp_cache,
&proxy_url,
&self.auth_manager,
)
.await
}
pub async fn get_managed_mcp_gateway_tool_catalog(
&self,
) -> Option<crate::session::managed_mcp::GatewayToolCatalog> {
if !self.can_fetch_managed_mcp_gateway_tools() {
self.managed_mcp_cache.lock().await.disable_gateway_tools();
return None;
}
self.managed_mcp_cache.lock().await.enable_gateway_tools();
let proxy_url = self.cfg.borrow().endpoints.proxy_url();
let auth_key = self
.auth_manager
.get_valid_token()
.await
.ok()
.or_else(|| self.auth_manager.current_or_expired().map(|a| a.key));
crate::session::managed_mcp::get_or_fetch_gateway_tool_catalog(
&self.managed_mcp_cache,
&proxy_url,
auth_key.as_deref(),
)
.await
}
pub fn managed_mcp_cache(
&self,
) -> &crate::session::managed_mcp::ManagedMcpStateHandle {
&self.managed_mcp_cache
}
pub(crate) fn disable_managed_gateway_tools_and_refresh_sessions(&self) {
self.disable_managed_gateway_tools_and_refresh_sessions_with_txs(
self.sessions.borrow().values().map(|handle| handle.cmd_tx.clone()).collect(),
);
}
fn disable_managed_gateway_tools_and_refresh_sessions_with_txs(
&self,
session_txs: Vec<tokio::sync::mpsc::UnboundedSender<SessionCommand>>,
) {
let cache = self.managed_mcp_cache.clone();
tokio::task::spawn_local(async move {
cache.lock().await.disable_gateway_tools();
for tx in session_txs {
let _ = tx.send(SessionCommand::RefreshMcpSearchIndex);
}
});
}
pub(crate) fn spawn_managed_gateway_tool_catalog_fetch(&self) {
let session_txs: Vec<_> = self
.sessions
.borrow()
.values()
.map(|handle| handle.cmd_tx.clone())
.collect();
if !self.can_fetch_managed_mcp_gateway_tools() {
self.disable_managed_gateway_tools_and_refresh_sessions_with_txs(
session_txs,
);
return;
}
let cache = self.managed_mcp_cache.clone();
let proxy_url = self.cfg.borrow().endpoints.proxy_url();
let auth_manager = self.auth_manager.clone();
tokio::task::spawn_local(async move {
let auth_key = auth_manager
.get_valid_token()
.await
.ok()
.or_else(|| auth_manager.current_or_expired().map(|a| a.key));
if !auth_manager
.current_or_expired()
.is_some_and(|a| a.is_session_auth())
{
cache.lock().await.disable_gateway_tools();
for tx in session_txs {
let _ = tx.send(SessionCommand::RefreshMcpSearchIndex);
}
return;
}
cache.lock().await.enable_gateway_tools();
crate::session::managed_mcp::get_or_fetch_gateway_tool_catalog(
&cache,
&proxy_url,
auth_key.as_deref(),
)
.await;
for tx in session_txs {
let _ = tx.send(SessionCommand::RefreshMcpSearchIndex);
}
});
}
/// Resolve the launch dir's project-scope trust verdict ONCE and return it
/// 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<acp::McpServer>,
cwd: &std::path::Path,
) -> (Vec<acp::McpServer>, Option<chrono::DateTime<chrono::Utc>>) {
) -> Vec<acp::McpServer> {
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)
crates/codegen/kigi-shell/src/agent/mvp_agent/folder_trust_prompt.rs
+4 -29
View File
@@ -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<ReloadTarget>,
plugin_handle: &'a kigi_agent::plugins::SharedPluginRegistryHandle,
managed_mcp_cache: &'a crate::session::managed_mcp::ManagedMcpStateHandle,
auth_manager: &'a std::sync::Arc<AuthManager>,
can_fetch_managed: bool,
proxy_url: &'a str,
compat: &'a kigi_tools::types::CompatConfig,
/// 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)]
crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs
-4
View File
@@ -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<chrono::DateTime<chrono::Utc>>,
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<bool>,
/// Managed MCP configs and gateway tool catalog; lazily fetched.
managed_mcp_cache: crate::session::managed_mcp::ManagedMcpStateHandle,
/// Agent-level MCP server state. LEADER-SAFE(shared): MCP servers are
/// agent-scoped, not per-client.
agent_mcp_state: std::sync::Arc<
crates/codegen/kigi-shell/src/agent/mvp_agent/subagent_coordinator.rs
-7
View File
@@ -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,
crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs
+1 -57
View File
@@ -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(&params).unwrap()),
))
.await
.expect("auth_cleared must route through session-admin");
let cmd = tokio::time::timeout(std::time::Duration::from_secs(1), cmd_rx.recv())
.await
.expect("refresh command should be sent")
.expect("channel should stay open until command is received");
assert!(matches!(cmd, SessionCommand::RefreshMcpSearchIndex));
assert!(!agent.managed_mcp_cache.lock().await.gateway_tools_active);
})
.await;
}
/// Build a minimal MvpAgent suitable for testing extension methods.
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(),
crates/codegen/kigi-shell/src/agent/subagent/handle_request.rs
-5
View File
@@ -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,
crates/codegen/kigi-shell/src/agent/subagent/mod.rs
-7
View File
@@ -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<agent_client_protocol::McpServer>,
/// Parent's managed MCP state handle (Arc-shared, no re-fetch).
pub managed_mcp_state: crate::session::managed_mcp::ManagedMcpStateHandle,
/// Snapshot of the parent session's MCP client pool at spawn time.
pub parent_mcp_pool: Option<crate::session::mcp_servers::SharedMcpPool>,
/// Snapshot of the parent session's resolved tool schema at spawn time.
crates/codegen/kigi-shell/src/agent/subagent/tests/mod.rs
-11
View File
@@ -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,
crates/codegen/kigi-shell/src/auth/attribution.rs
+9 -39
View File
@@ -123,7 +123,7 @@ impl ShellAttribution {
/// Tool-side counterpart of [`Self::new`]: returns
/// `Arc<dyn kigi_tools::Auth401AttributionCallback>` 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 `<prefix>.<op>`. `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<dyn ToolAuth401AttributionCallback> =
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"));
crates/codegen/kigi-shell/src/auth/manager.rs
+1 -1
View File
@@ -1877,7 +1877,7 @@ impl AuthManager {
}
/// Bridges `Arc<AuthManager>` into the `ApiKeyProvider` trait used by
/// tool clients (image_gen, video_gen, web_search, embedding). Sync callers
/// tool clients (web_search, embedding). Sync callers
/// get the buffered snapshot; async callers drive the refresh chain.
pub(crate) struct SharedAuthKeyProvider(pub Arc<AuthManager>);
crates/codegen/kigi-shell/src/builtin.rs
+6 -7
View File
@@ -5,7 +5,6 @@ const BUNDLED_FILES: &[(&str, &str)] = &[("README.md", include_str!("../README.m
const HELP_SKILL_MD: &str = include_str!("../skills/help/SKILL.md");
const 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"
crates/codegen/kigi-shell/src/config/mod.rs
+3 -117
View File
@@ -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<kigi_tools::implementations::grok_build::video_gen::ZdrVideoOutputS3Config>,
}
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
crates/codegen/kigi-shell/src/config/tests.rs
+6 -247
View File
@@ -1201,154 +1201,6 @@ fn subagents_config_is_subagent_enabled_false_when_toggled_off() {
);
assert!(sa.is_subagent_enabled("explore"), "explore = true should return enabled");
}
fn with_managed_mcp_env<T>(
managed_mcps: Option<&str>,
gateway_tools: Option<&str>,
f: impl FnOnce() -> T,
) -> T {
static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
let _guard = LOCK.lock().unwrap_or_else(|e| e.into_inner());
with_env_var_opt(
"KIGI_MANAGED_MCPS_ENABLED",
managed_mcps,
|| with_env_var_opt("KIGI_MANAGED_MCP_GATEWAY_TOOLS_ENABLED", gateway_tools, f),
)
}
#[test]
#[serial_test::serial]
fn managed_mcps_interactive_default_enabled() {
with_managed_mcp_env(
None,
None,
|| {
let empty = toml::Value::Table(toml::map::Map::new());
let cfg = ManagedMcpsConfig::resolve(&empty, None, false);
assert!(cfg.enabled);
},
);
}
#[test]
#[serial_test::serial]
fn managed_mcps_headless_default_disabled() {
with_managed_mcp_env(
None,
None,
|| {
let empty = toml::Value::Table(toml::map::Map::new());
let cfg = ManagedMcpsConfig::resolve(&empty, None, true);
assert!(! cfg.enabled);
},
);
}
#[test]
#[serial_test::serial]
fn managed_mcp_gateway_tools_default_disabled() {
with_managed_mcp_env(
None,
None,
|| {
let empty = toml::Value::Table(toml::map::Map::new());
let cfg = ManagedMcpsConfig::resolve(&empty, None, false);
assert!(! cfg.gateway_tools_enabled);
},
);
}
#[test]
#[serial_test::serial]
fn managed_mcp_gateway_tools_require_managed_master() {
with_managed_mcp_env(
None,
None,
|| {
let config: toml::Value = toml::from_str(
r#"
[managed_mcps]
gateway_tools_enabled = true
"#,
)
.unwrap();
let remote = crate::util::config::RemoteSettings {
managed_mcps_enabled: Some(false),
..Default::default()
};
let cfg = ManagedMcpsConfig::resolve(&config, Some(&remote), true);
assert!(! cfg.enabled);
assert!(! cfg.gateway_tools_enabled);
},
);
}
#[test]
#[serial_test::serial]
fn managed_mcp_gateway_tools_remote_enabled() {
with_managed_mcp_env(
None,
None,
|| {
let empty = toml::Value::Table(toml::map::Map::new());
let remote = crate::util::config::RemoteSettings {
managed_mcp_gateway_tools_enabled: Some(true),
..Default::default()
};
let cfg = ManagedMcpsConfig::resolve(&empty, Some(&remote), false);
assert!(cfg.gateway_tools_enabled);
},
);
}
#[test]
#[serial_test::serial]
fn managed_mcp_gateway_tools_env_overrides_remote() {
with_managed_mcp_env(
None,
Some("0"),
|| {
let empty = toml::Value::Table(toml::map::Map::new());
let remote = crate::util::config::RemoteSettings {
managed_mcp_gateway_tools_enabled: Some(true),
..Default::default()
};
let cfg = ManagedMcpsConfig::resolve(&empty, Some(&remote), false);
assert!(! cfg.gateway_tools_enabled);
},
);
}
#[test]
#[serial_test::serial]
fn managed_mcp_gateway_tools_env_on_overrides_remote_off() {
with_managed_mcp_env(
None,
Some("1"),
|| {
let empty = toml::Value::Table(toml::map::Map::new());
let remote = crate::util::config::RemoteSettings {
managed_mcp_gateway_tools_enabled: Some(false),
..Default::default()
};
let cfg = ManagedMcpsConfig::resolve(&empty, Some(&remote), false);
assert!(cfg.gateway_tools_enabled);
},
);
}
#[test]
#[serial_test::serial]
fn managed_mcp_gateway_tools_enabled_with_managed_master() {
with_managed_mcp_env(
None,
None,
|| {
let config: toml::Value = toml::from_str(
r#"
[managed_mcps]
enabled = true
gateway_tools_enabled = true
"#,
)
.unwrap();
let cfg = ManagedMcpsConfig::resolve(&config, None, false);
assert!(cfg.enabled);
assert!(cfg.gateway_tools_enabled);
},
);
}
fn with_model_overrides_env_full<T>(
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<T>(
respect_gitignore: Option<&str>,
disable_zdr: Option<&str>,
f: impl FnOnce() -> T,
) -> T {
let _guard = TOOLS_ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
with_env_var_opt(
"KIGI_RESPECT_GITIGNORE",
respect_gitignore,
|| with_env_var_opt("KIGI_DISABLE_ZDR_INCOMPATIBLE_TOOLS", disable_zdr, f),
)
}
fn without_grok_respect_gitignore<T>(f: impl FnOnce() -> T) -> T {
with_tools_env(None, None, f)
let _guard = TOOLS_ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
with_env_var_opt("KIGI_RESPECT_GITIGNORE", None, f)
}
fn with_grok_respect_gitignore<T>(value: &str, f: impl FnOnce() -> T) -> T {
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]
crates/codegen/kigi-shell/src/extensions/mcp.rs
+53 -708
View File
@@ -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<String>,
/// 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<String>,
#[serde(rename = "scopeId", skip_serializing_if = "Option::is_none")]
scope_id: Option<String>,
#[serde(rename = "scopeName", skip_serializing_if = "Option::is_none")]
scope_name: Option<String>,
},
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<McpEnvVar>,
},
#[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<serde_json::Value>,
}
/// 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<McpServerEntry> {
build_mcp_catalog_with_gateway_tools(managed_configs, local_servers, None, &Default::default())
}
pub fn build_mcp_catalog_with_gateway_tools(
managed_configs: &[crate::session::managed_mcp::ManagedMcpConfig],
local_servers: &[acp::McpServer],
gateway_catalog: Option<&crate::session::managed_mcp::GatewayToolCatalog>,
disabled_tools: &HashMap<String, HashSet<String>>,
) -> Vec<McpServerEntry> {
pub fn build_mcp_catalog(local_servers: &[acp::McpServer]) -> Vec<McpServerEntry> {
let mut servers: Vec<McpServerEntry> = 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<String>,
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::<McpListRequest>(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<String> =
servers.iter().map(|s| s.name.clone()).collect();
for name in &disabled_names {
if should_append_disabled_mcp_placeholder(name, &catalog_names, gateway_tools_enabled) {
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<McpServerEntry>,
disabled_names: &[&str],
gateway_tools_enabled: bool,
) {
let catalog_names: std::collections::HashSet<String> =
servers.iter().map(|s| s.name.clone()).collect();
for name in disabled_names {
if should_append_disabled_mcp_placeholder(name, &catalog_names, gateway_tools_enabled) {
servers.push(disabled_server_placeholder_entry(name));
}
}
}
#[test]
fn disabled_placeholders_match_handle_list_catalog_relationships() {
// Empty loads (gateway on, nothing in catalog) + orphan legacy disables
// only in disabled_mcp_servers — production ghost-stub regression.
let mut servers = build_mcp_catalog_with_gateway_tools(&[], &[], None, &Default::default());
append_disabled_like_handle_list(
&mut servers,
&["grok_com_slack", "grok_mcp_linear"],
true,
);
assert!(
servers.is_empty(),
"gateway on + no catalog rows → no stubs for orphan disables"
);
// Same orphans with gateway off → still placeholders (legacy UX).
let mut servers = build_mcp_catalog_with_gateway_tools(&[], &[], None, &Default::default());
append_disabled_like_handle_list(&mut servers, &["grok_com_slack"], false);
assert_eq!(servers.len(), 1);
assert_eq!(servers[0].name, "grok_com_slack");
assert!(!servers[0].session.as_ref().unwrap().enabled);
// Name already in catalog (gateway row) → never double-append.
let gateway = crate::session::managed_mcp::GatewayToolCatalog {
tools: vec![gateway_tool(
"linear",
"Linear",
"list_issues",
"List issues",
"linear.list_issues",
"List Linear issues",
)],
total_tools: 1,
connectors_needing_reauth: vec![],
};
let mut servers =
build_mcp_catalog_with_gateway_tools(&[], &[], Some(&gateway), &Default::default());
let gateway_entry = managed_gateway_entry_name("linear");
assert!(servers.iter().any(|s| s.name == gateway_entry));
let before = servers.len();
append_disabled_like_handle_list(&mut servers, &[gateway_entry.as_str()], true);
append_disabled_like_handle_list(&mut servers, &[gateway_entry.as_str()], false);
assert_eq!(
servers.len(),
before,
"disabled name already in catalog must not add a second row"
);
assert_eq!(
servers.iter().filter(|s| s.name == gateway_entry).count(),
1
);
}
#[test]
fn gateway_catalog_honors_disabled_connectors_and_tools() {
let catalog = crate::session::managed_mcp::GatewayToolCatalog {
tools: vec![
gateway_tool(
"linear",
"Linear",
"list_issues",
"List issues",
"linear.list_issues",
"List Linear issues",
),
gateway_tool(
"linear",
"Linear",
"create_issue",
"Create issue",
"linear.create_issue",
"Create a Linear issue",
),
],
total_tools: 2,
connectors_needing_reauth: vec![],
};
let disabled: HashMap<String, HashSet<String>> = HashMap::from([
(
crate::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY.to_string(),
HashSet::from(["linear".to_string()]),
),
(
"linear".to_string(),
HashSet::from(["linear__create_issue".to_string()]),
),
]);
let servers = build_mcp_catalog_with_gateway_tools(&[], &[], Some(&catalog), &disabled);
let session = servers[0].session.as_ref().unwrap();
assert!(!session.enabled);
assert!(session.status.is_none());
assert!(session.tools[0].enabled);
assert!(!session.tools[1].enabled);
}
#[test]
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());
crates/codegen/kigi-shell/src/extensions/session_admin.rs
+2 -17
View File
@@ -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 {
crates/codegen/kigi-shell/src/leader/server.rs
+3 -7
View File
@@ -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();
crates/codegen/kigi-shell/src/mcp_doctor.rs
+1 -86
View File
@@ -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<ConfigSourceStatus>, Vec<DiscoveredServe
(sources, servers)
}
// ── Managed (grok.com) server discovery ─────────────────────────
const MANAGED_SOURCE_LABEL: &str = "grok.com";
fn managed_skipped(reason: impl Into<String>) -> (ConfigSourceStatus, Vec<DiscoveredServer>) {
(
ConfigSourceStatus {
path: MANAGED_SOURCE_LABEL.to_string(),
status: ConfigSourceState::Skipped {
reason: reason.into(),
},
},
vec![],
)
}
fn managed_found(
count: usize,
servers: Vec<DiscoveredServer>,
) -> (ConfigSourceStatus, Vec<DiscoveredServer>) {
(
ConfigSourceStatus {
path: MANAGED_SOURCE_LABEL.to_string(),
status: ConfigSourceState::Found {
server_count: count,
},
},
servers,
)
}
/// Discover managed `grok_com_*` servers if the user has xAI auth on disk.
async fn try_discover_managed_servers() -> (ConfigSourceStatus, Vec<DiscoveredServer>) {
let kigi_home = kigi_tools::util::kigi_home::kigi_home();
let kimi_code_config = KimiCodeConfig::default();
let auth_manager = Arc::new(crate::auth::AuthManager::new(&kigi_home, kimi_code_config));
let Some(snapshot) = auth_manager.current_or_expired() else {
return managed_skipped("not logged in");
};
if !snapshot.is_session_auth() {
return managed_skipped(format!("{:?} auth (not xAI OIDC)", snapshot.auth_mode));
}
let token = match auth_manager.get_valid_token().await {
Ok(key) => key,
Err(_) => return managed_skipped("auth expired — run `kigi login`"),
};
let proxy_url = crate::agent::config::EndpointsConfig::from_effective_config().proxy_url();
let configs = match managed_mcp::fetch_managed_configs(&proxy_url, &token).await {
Ok(configs) => configs,
Err(e) => return managed_skipped(format!("fetch failed: {e}")),
};
if configs.is_empty() {
return managed_found(0, vec![]);
}
let mut servers: Vec<agent_client_protocol::McpServer> = vec![];
managed_mcp::auto_inject_managed_servers_with_disabled(
&mut servers,
&configs,
&Default::default(),
);
managed_mcp::inject_managed_headers(&mut servers, &configs);
let source = ConfigSource::Managed { path: None };
let discovered: Vec<DiscoveredServer> = servers
.into_iter()
.map(|server| DiscoveredServer {
server,
source: source.clone(),
})
.collect();
managed_found(discovered.len(), discovered)
}
// ── Check functions ─────────────────────────────────────────────
fn resolve_command(command: &str) -> Option<String> {
@@ -558,11 +477,7 @@ pub async fn run_auth(cwd: &Path, name: &str) -> Result<usize, String> {
}
pub async fn run_doctor(cwd: &Path, name_filter: Option<&str>) -> DoctorReport {
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() {
crates/codegen/kigi-shell/src/session/acp_conversion.rs
-41
View File
@@ -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(
crates/codegen/kigi-shell/src/session/acp_session.rs
+34 -578
View File
@@ -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<AuthManager>,
}
#[async_trait::async_trait]
impl kigi_tools::types::resources::ManagedGatewayToolCaller for ShellManagedGatewayToolClient {
async fn call_tool(
&self,
call_id: &str,
arguments: serde_json::Value,
caller: &str,
) -> Result<
kigi_tools::types::resources::ManagedGatewayToolCallResponse,
kigi_tool_runtime::ToolError,
> {
let auth_key = self
.auth_manager
.get_valid_token()
.await
.ok()
.or_else(|| self.auth_manager.current_or_expired().map(|a| a.key))
.ok_or_else(|| kigi_tool_runtime::ToolError::unauthorized("no auth token available"))?;
let response = crate::session::managed_mcp::call_gateway_tool(
&self.proxy_base_url,
&auth_key,
call_id,
arguments,
)
.await
.map_err(|error| managed_gateway_error_to_tool_error(error, caller))?;
Ok(
kigi_tools::types::resources::ManagedGatewayToolCallResponse {
result: response.result,
connectors_needing_reauth: response.connectors_needing_reauth,
},
)
}
}
fn managed_gateway_error_to_tool_error(
error: crate::session::managed_mcp::ManagedMcpFetchError,
caller: &str,
) -> kigi_tool_runtime::ToolError {
match error {
crate::session::managed_mcp::ManagedMcpFetchError::Status { status, message } => {
let detail = format!("Managed MCP gateway tool call failed: {message}");
let mut err = if status == reqwest::StatusCode::UNAUTHORIZED {
kigi_tool_runtime::ToolError::unauthorized(detail)
} else if status == reqwest::StatusCode::FORBIDDEN {
kigi_tool_runtime::ToolError::permission_denied(detail)
} else {
let tool_id = kigi_tool_protocol::ToolId::new(caller).unwrap_or_else(|_| {
kigi_tool_protocol::ToolId::new("use_tool").expect("valid")
});
kigi_tool_runtime::ToolError::execution(tool_id, detail)
};
match err.details.as_mut() {
Some(serde_json::Value::Object(map)) => {
map.insert(
HTTP_STATUS_DETAILS_KEY.to_string(),
serde_json::json!(status.as_u16()),
);
}
_ => {
err.details =
Some(serde_json::json!({ HTTP_STATUS_DETAILS_KEY : status.as_u16(), }));
}
}
err
}
crate::session::managed_mcp::ManagedMcpFetchError::Transport(e) => {
kigi_tool_runtime::ToolError::network_error(format!(
"Managed MCP gateway tool call failed: {}",
e.without_url()
))
}
crate::session::managed_mcp::ManagedMcpFetchError::NoAuth => {
kigi_tool_runtime::ToolError::unauthorized("no auth token available")
}
}
}
#[cfg(test)]
mod managed_gateway_error_tests {
use super::*;
fn status_error(code: u16, message: &str) -> crate::session::managed_mcp::ManagedMcpFetchError {
crate::session::managed_mcp::ManagedMcpFetchError::Status {
status: reqwest::StatusCode::from_u16(code).unwrap(),
message: message.to_string(),
}
}
#[test]
fn unauthorized_status_maps_to_unauthorized_and_carries_status() {
let err = managed_gateway_error_to_tool_error(status_error(401, "expired"), "use_tool");
assert_eq!(err.kind, kigi_tool_runtime::ToolErrorKind::Unauthorized);
assert!(err.detail.contains("expired"));
let details = err.details.as_ref().unwrap();
assert_eq!(
details.get(HTTP_STATUS_DETAILS_KEY),
Some(&serde_json::json!(401))
);
}
#[test]
fn forbidden_status_maps_to_permission_denied_and_carries_status() {
let err = managed_gateway_error_to_tool_error(status_error(403, "denied"), "use_tool");
assert_eq!(err.kind, kigi_tool_runtime::ToolErrorKind::PermissionDenied);
let details = err.details.as_ref().unwrap();
assert_eq!(
details.get(HTTP_STATUS_DETAILS_KEY),
Some(&serde_json::json!(403))
);
}
#[test]
fn general_status_maps_to_execution_with_caller_tool_id() {
let err = managed_gateway_error_to_tool_error(status_error(500, "boom"), "CallMcpTool");
assert_eq!(err.kind, kigi_tool_runtime::ToolErrorKind::Execution);
let details = err.details.as_ref().unwrap();
assert_eq!(
details.get(HTTP_STATUS_DETAILS_KEY),
Some(&serde_json::json!(500))
);
assert_eq!(
details.get("tool_id"),
Some(&serde_json::json!("CallMcpTool"))
);
}
#[test]
fn general_status_falls_back_to_use_tool_for_unknown_caller() {
let err = managed_gateway_error_to_tool_error(status_error(500, "boom"), "not a tool id");
assert_eq!(err.kind, kigi_tool_runtime::ToolErrorKind::Execution);
let details = err.details.as_ref().unwrap();
assert_eq!(details.get("tool_id"), Some(&serde_json::json!("use_tool")));
}
#[test]
fn no_auth_maps_to_unauthorized() {
let err = managed_gateway_error_to_tool_error(
crate::session::managed_mcp::ManagedMcpFetchError::NoAuth,
"use_tool",
);
assert_eq!(err.kind, kigi_tool_runtime::ToolErrorKind::Unauthorized);
}
#[tokio::test]
async fn transport_error_maps_to_network_error_without_url() {
let transport = reqwest::Client::new()
.post("http://127.0.0.1:1/mcp/tools/call")
.send()
.await
.expect_err("connection to a dead port should fail");
let err = managed_gateway_error_to_tool_error(
crate::session::managed_mcp::ManagedMcpFetchError::Transport(transport),
"use_tool",
);
assert_eq!(err.kind, kigi_tool_runtime::ToolErrorKind::NetworkError);
assert!(err.detail.contains("Managed MCP gateway tool call failed"));
assert!(
!err.detail.contains("http://"),
"transport detail must not leak the proxy URL: {}",
err.detail
);
}
}
/// Data carried from prepare_tool_call → dispatch_tool → finalize.
#[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<Option<chrono::DateTime<chrono::Utc>>>,
/// Original client-provided MCP servers from session creation.
/// Retained for re-merge during plugin reload.
pub(crate) initial_client_mcp_servers: Vec<acp::McpServer>,
@@ -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<ToolOutput, kigi_tool_runtime::ToolError> {
Ok(ToolOutput::MCP(MCPOutput::okay_output(
"server__tool".to_string(),
"server".to_string(),
"ok".to_string(),
)))
}
}
#[tokio::test]
async fn refresh_snapshot_indexes_only_admitted_gateway_tools() {
let bridge = Arc::new(crate::tools::bridge::ToolBridge::for_test());
bridge
.register_mcp_tools(
"server__tool".to_string(),
FixtureMcpTool,
Some(serde_json::json!({ "type" : "object" })),
)
.await
.expect("local fixture registration succeeds");
let mcp_state = Arc::new(TokioMutex::new(McpState::new(vec![])));
let managed = crate::session::managed_mcp::ManagedMcpStateHandle::default();
{
let mut state = managed.lock().await;
state.enable_gateway_tools();
let epoch = state.start_gateway_tool_fetch().unwrap();
assert!(state.complete_gateway_tool_fetch(
epoch,
crate::session::managed_mcp::GatewayToolCatalog {
tools: vec![
crate::session::managed_mcp::GatewayTool {
connector_id: "server".to_string(),
connector_name: "Gateway Collision".to_string(),
tool_id: "tool".to_string(),
tool_name: "Collision".to_string(),
call_id: "gateway.collision".to_string(),
description: "Gateway collision".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
crate::session::managed_mcp::GatewayTool {
connector_id: "gateway".to_string(),
connector_name: "Gateway".to_string(),
tool_id: "search".to_string(),
tool_name: "Search".to_string(),
call_id: "gateway.search".to_string(),
description: "Gateway search".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
],
total_tools: 2,
connectors_needing_reauth: vec![],
}
));
}
let snapshot = Arc::new(std::sync::Mutex::new(
crate::session::tool_index::ToolMetadataSnapshot::default(),
));
refresh_mcp_snapshot_for_test(bridge, mcp_state, managed, snapshot.clone()).await;
let snapshot = snapshot.lock().unwrap();
let names: std::collections::HashSet<&str> = snapshot
.tools
.iter()
.map(|tool| tool.qualified_name.as_str())
.collect();
assert!(names.contains("gateway__search"));
let server_tool = snapshot
.tools
.iter()
.find(|tool| tool.qualified_name == "server__tool")
.expect("local MCP tool remains indexed");
assert_eq!(server_tool.server_name, "server");
assert_eq!(server_tool.description, "fixture");
}
#[tokio::test]
async fn refresh_snapshot_excludes_disabled_gateway_tools_and_connectors() {
let bridge = Arc::new(crate::tools::bridge::ToolBridge::for_test());
let mcp_state = Arc::new(TokioMutex::new(McpState::new(vec![])));
let managed = crate::session::managed_mcp::ManagedMcpStateHandle::default();
{
let mut state = managed.lock().await;
state.enable_gateway_tools();
let epoch = state.start_gateway_tool_fetch().unwrap();
assert!(state.complete_gateway_tool_fetch(
epoch,
crate::session::managed_mcp::GatewayToolCatalog {
tools: vec![
crate::session::managed_mcp::GatewayTool {
connector_id: "linear".to_string(),
connector_name: "Linear".to_string(),
tool_id: "list_issues".to_string(),
tool_name: "List".to_string(),
call_id: "linear.list_issues".to_string(),
description: "List issues".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
crate::session::managed_mcp::GatewayTool {
connector_id: "linear".to_string(),
connector_name: "Linear".to_string(),
tool_id: "create_issue".to_string(),
tool_name: "Create".to_string(),
call_id: "linear.create_issue".to_string(),
description: "Create issue".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
crate::session::managed_mcp::GatewayTool {
connector_id: "slack".to_string(),
connector_name: "Slack".to_string(),
tool_id: "search".to_string(),
tool_name: "Search".to_string(),
call_id: "slack.search".to_string(),
description: "Search Slack".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
],
total_tools: 3,
connectors_needing_reauth: vec![],
}
));
}
let snapshot = Arc::new(std::sync::Mutex::new(
crate::session::tool_index::ToolMetadataSnapshot::default(),
));
let disabled: std::collections::HashMap<String, std::collections::HashSet<String>> =
std::collections::HashMap::from([
(
"linear".to_string(),
std::collections::HashSet::from(["linear__create_issue".to_string()]),
),
(
crate::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY.to_string(),
std::collections::HashSet::from(["slack".to_string()]),
),
]);
refresh_mcp_snapshot_for_test_with_disabled(
bridge,
mcp_state,
managed,
snapshot.clone(),
&disabled,
)
.await;
let snapshot = snapshot.lock().unwrap();
let names: std::collections::HashSet<&str> = snapshot
.tools
.iter()
.map(|tool| tool.qualified_name.as_str())
.collect();
assert!(names.contains("linear__list_issues"));
assert!(!names.contains("linear__create_issue"));
assert!(!names.contains("slack__search"));
}
}
/// ToolBridge must route file operations through the injected FileSystem,
/// 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<ToolOutput, kigi_tool_runtime::ToolError> {
Ok(ToolOutput::MCP(MCPOutput::okay_output(
"server__tool".to_string(),
"server".to_string(),
"ok".to_string(),
)))
}
}
#[tokio::test]
async fn refresh_snapshot_seeds_only_admitted_gateway_catalog_entries() {
let bridge = Arc::new(crate::tools::bridge::ToolBridge::for_test());
bridge
.register_mcp_tools(
"server__tool".to_string(),
FixtureMcpTool,
Some(serde_json::json!({ "type" : "object" })),
)
.await
.expect("local fixture registration succeeds");
let mcp_state = Arc::new(TokioMutex::new(McpState::new(vec![])));
let managed = crate::session::managed_mcp::ManagedMcpStateHandle::default();
{
let mut state = managed.lock().await;
state.enable_gateway_tools();
let epoch = state.start_gateway_tool_fetch().unwrap();
assert!(state.complete_gateway_tool_fetch(
epoch,
crate::session::managed_mcp::GatewayToolCatalog {
tools: vec![
crate::session::managed_mcp::GatewayTool {
connector_id: "server".to_string(),
connector_name: "Gateway Collision".to_string(),
tool_id: "tool".to_string(),
tool_name: "Collision".to_string(),
call_id: "gateway.collision".to_string(),
description: "Gateway collision".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
crate::session::managed_mcp::GatewayTool {
connector_id: "gateway".to_string(),
connector_name: "Gateway".to_string(),
tool_id: "search".to_string(),
tool_name: "Search".to_string(),
call_id: "gateway.search".to_string(),
description: "Gateway search".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
],
total_tools: 2,
connectors_needing_reauth: vec![],
}
));
}
let snapshot = Arc::new(std::sync::Mutex::new(
crate::session::tool_index::ToolMetadataSnapshot::default(),
));
refresh_mcp_snapshot_for_test(bridge.clone(), mcp_state, managed, snapshot.clone()).await;
let catalog = bridge
.read_resource::<kigi_tools::types::resources::ManagedGatewayToolCatalog>()
.await
.expect("catalog resource should be seeded");
assert!(catalog.get("gateway__search").is_some());
assert!(
catalog.get("server__tool").is_none(),
"gateway catalog resource must match admitted snapshot and skip local collisions"
);
}
#[tokio::test]
async fn refresh_snapshot_excludes_disabled_gateway_tools_and_connectors() {
let bridge = Arc::new(crate::tools::bridge::ToolBridge::for_test());
let mcp_state = Arc::new(TokioMutex::new(McpState::new(vec![])));
let managed = crate::session::managed_mcp::ManagedMcpStateHandle::default();
{
let mut state = managed.lock().await;
state.enable_gateway_tools();
let epoch = state.start_gateway_tool_fetch().unwrap();
assert!(state.complete_gateway_tool_fetch(
epoch,
crate::session::managed_mcp::GatewayToolCatalog {
tools: vec![
crate::session::managed_mcp::GatewayTool {
connector_id: "linear".to_string(),
connector_name: "Linear".to_string(),
tool_id: "list_issues".to_string(),
tool_name: "List".to_string(),
call_id: "linear.list_issues".to_string(),
description: "List issues".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
crate::session::managed_mcp::GatewayTool {
connector_id: "linear".to_string(),
connector_name: "Linear".to_string(),
tool_id: "create_issue".to_string(),
tool_name: "Create".to_string(),
call_id: "linear.create_issue".to_string(),
description: "Create issue".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
crate::session::managed_mcp::GatewayTool {
connector_id: "slack".to_string(),
connector_name: "Slack".to_string(),
tool_id: "search".to_string(),
tool_name: "Search".to_string(),
call_id: "slack.search".to_string(),
description: "Search Slack".to_string(),
json_schema: serde_json::json!({ "type" : "object" }),
},
],
total_tools: 3,
connectors_needing_reauth: vec![],
}
));
}
let snapshot = Arc::new(std::sync::Mutex::new(
crate::session::tool_index::ToolMetadataSnapshot::default(),
));
let disabled: std::collections::HashMap<String, std::collections::HashSet<String>> =
std::collections::HashMap::from([
(
"linear".to_string(),
std::collections::HashSet::from(["linear__create_issue".to_string()]),
),
(
crate::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY.to_string(),
std::collections::HashSet::from(["slack".to_string()]),
),
]);
refresh_mcp_snapshot_for_test_with_disabled(
bridge.clone(),
mcp_state,
managed,
snapshot.clone(),
&disabled,
)
.await;
let catalog = bridge
.read_resource::<kigi_tools::types::resources::ManagedGatewayToolCatalog>()
.await
.expect("catalog resource should be seeded");
assert!(catalog.get("linear__list_issues").is_some());
assert!(catalog.get("linear__create_issue").is_none());
assert!(catalog.get("slack__search").is_none());
let snapshot = snapshot.lock().unwrap();
let names: std::collections::HashSet<&str> = snapshot
.tools
.iter()
.map(|tool| tool.qualified_name.as_str())
.collect();
assert!(names.contains("linear__list_issues"));
assert!(!names.contains("linear__create_issue"));
assert!(!names.contains("slack__search"));
}
}
#[cfg(test)]
#[path = "acp_session_tests/goal/goal_backoff_tests.rs"]
mod goal_backoff_tests;
#[cfg(test)]
#[path = "acp_session_tests/goal/goal_classifier_e2e_tests.rs"]
mod goal_classifier_e2e_tests;
#[cfg(test)]
#[path = "acp_session_tests/goal/goal_planner_e2e_tests.rs"]
mod goal_planner_e2e_tests;
#[cfg(test)]
#[path = "acp_session_tests/goal/goal_strategist_e2e_tests.rs"]
mod goal_strategist_e2e_tests;
#[cfg(test)]
#[path = "acp_session_tests/goal/goal_summarizer_e2e_tests.rs"]
mod goal_summarizer_e2e_tests;
#[cfg(test)]
#[path = "acp_session_tests/interjection_tests.rs"]
mod interjection_tests;
#[cfg(test)]
#[path = "acp_session_tests/recap_display_only_tests.rs"]
mod recap_display_only_tests;
#[cfg(test)]
#[path = "acp_session_tests/reminder_policy_tests.rs"]
mod reminder_policy_tests;
crates/codegen/kigi-shell/src/session/acp_session_impl/hooks_plugins.rs
-8
View File
@@ -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,
);
crates/codegen/kigi-shell/src/session/acp_session_impl/mcp.rs
+3 -274
View File
@@ -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<usize, String> {
let auth_manager = self.auth_manager.clone();
let proxy_url = self.models_manager.endpoints().proxy_url();
let Some(am) = auth_manager else {
return Err("no auth manager; cannot re-fetch managed configs".into());
};
crate::session::managed_mcp::invalidate_cache(&self.managed_mcp_handle).await;
let fresh_configs = crate::session::managed_mcp::fetch_managed_mcp_configs(
&self.managed_mcp_handle,
&proxy_url,
&am,
)
.await;
if fresh_configs.is_empty() {
return Err("managed re-fetch returned no configs".into());
}
{
let mut st = self.mcp_state.lock().await;
st.refresh_managed_clients(
fresh_configs
.iter()
.map(|c| (c.endpoint.as_str(), &c.headers)),
);
}
*self.managed_mcp_expires_at.lock().unwrap() = fresh_configs
.iter()
.filter_map(|c| c.token_expires_at)
.min();
let client = {
let st = self.mcp_state.lock().await;
st.get_client(server_name)
.cloned()
.ok_or_else(|| format!("client '{server_name}' missing after refresh"))?
};
let registrations = client
.get_tool_registrations(self.mcp_state.clone())
.await
.map_err(|e| format!("re-handshake failed: {e}"))?;
let mut mcp_state = self.mcp_state.lock().await;
mcp_state.auth_required.remove(server_name);
mcp_state.clear_init_failed(server_name);
let tool_count = registrations.len();
let mut ui_tools: std::collections::HashMap<
String,
Vec<crate::extensions::mcp::McpToolEntry>,
> = std::collections::HashMap::new();
for reg in registrations {
self.register_mcp_tool(server_name, reg, &mut mcp_state, &mut ui_tools)
.await;
}
drop(mcp_state);
self.refresh_mcp_snapshot_and_schedule_reminder().await;
self.emit_mcp_tools_changed_notifications(ui_tools);
Ok(tool_count)
}
/// Register tools from shared (inherited) MCP clients on this session's ToolBridge.
///
/// 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 {
/// `<system-reminder>` 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<String>,
>,
) {
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<String, &'static str> =
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;
crates/codegen/kigi-shell/src/session/acp_session_impl/mcp_snapshot.rs
+1 -120
View File
@@ -19,27 +19,12 @@ impl McpReminderMode {
}
}
pub(super) fn gateway_tool_is_disabled(
tool: &crate::session::managed_mcp::GatewayTool,
disabled_gateway_tools: &std::collections::HashMap<String, std::collections::HashSet<String>>,
) -> bool {
let qualified_name = tool.qualified_name();
disabled_gateway_tools
.get(crate::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY)
.is_some_and(|set| set.contains(&tool.connector_id))
|| disabled_gateway_tools
.get(&tool.connector_id)
.is_some_and(|set| set.contains(&qualified_name))
}
pub(super) async fn refresh_mcp_snapshot_and_schedule_reminder_with(
tool_bridge: Arc<crate::tools::bridge::ToolBridge>,
mcp_state: Arc<TokioMutex<McpState>>,
managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle,
tool_metadata_snapshot: Arc<std::sync::Mutex<crate::session::tool_index::ToolMetadataSnapshot>>,
mcp_reminder_dirty: Arc<std::sync::atomic::AtomicBool>,
mcp_initialized: bool,
disabled_gateway_tools: &std::collections::HashMap<String, std::collections::HashSet<String>>,
// 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<ToolMetadata> = all_defs
let mcp_tools: Vec<ToolMetadata> = 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<String> = state.gateway_tool_connectors_seen.iter().cloned().collect();
(catalog, connectors)
};
if let Some(catalog) = gateway_catalog.as_ref() {
gateway_connectors.extend(catalog.tools.iter().map(|tool| tool.connector_id.clone()));
}
gateway_connectors.sort_unstable();
gateway_connectors.dedup();
let mut gateway_resource_entries = Vec::new();
if let Some(catalog) = gateway_catalog.as_ref() {
for tool in &catalog.tools {
let qualified_name = tool.qualified_name();
if gateway_tool_is_disabled(tool, disabled_gateway_tools) {
continue;
}
if !seen_tools.insert(qualified_name.clone()) {
continue;
}
gateway_resource_entries.push((
qualified_name.clone(),
kigi_tools::types::resources::ManagedGatewayToolSource {
connector_id: tool.connector_id.clone(),
connector_name: tool.connector_name.clone(),
tool_id: tool.tool_id.clone(),
tool_name: tool.tool_name.clone(),
call_id: tool.call_id.clone(),
},
));
// Gateway ids are the model/search contract. Display labels stay
// out of ToolMetadata so search_tool and permissions use stable ids:
// connector_id/tool_id here, connector_name/tool_name in UI only.
mcp_tools.push(ToolMetadata {
qualified_name,
server_name: tool.connector_id.clone(),
tool_name: tool.tool_id.clone(),
description: tool.description.clone(),
parameters: extract_parameter_names(&tool.json_schema),
input_schema: tool.json_schema.clone(),
});
}
}
let servers_with_tools: std::collections::HashSet<&str> =
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<crate::tools::bridge::ToolBridge>,
mcp_state: Arc<TokioMutex<McpState>>,
managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle,
tool_metadata_snapshot: Arc<std::sync::Mutex<crate::session::tool_index::ToolMetadataSnapshot>>,
) {
refresh_mcp_snapshot_for_test_with_disabled(
tool_bridge,
mcp_state,
managed_mcp_handle,
tool_metadata_snapshot,
&Default::default(),
)
.await;
}
#[cfg(test)]
pub(crate) async fn refresh_mcp_snapshot_for_test_with_disabled(
tool_bridge: Arc<crate::tools::bridge::ToolBridge>,
mcp_state: Arc<TokioMutex<McpState>>,
managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle,
tool_metadata_snapshot: Arc<std::sync::Mutex<crate::session::tool_index::ToolMetadataSnapshot>>,
disabled_gateway_tools: &std::collections::HashMap<String, std::collections::HashSet<String>>,
) {
refresh_mcp_snapshot_and_schedule_reminder_with(
tool_bridge,
mcp_state,
managed_mcp_handle,
tool_metadata_snapshot,
Arc::new(std::sync::atomic::AtomicBool::new(false)),
false,
disabled_gateway_tools,
None,
)
.await;
crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs
-3
View File
@@ -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))
crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs
-63
View File
@@ -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<kigi_agent::prompt::user_message::McpServerEntry> {
use kigi_agent::prompt::user_message::McpServerEntry;
let disabled_gateway_tools = crate::util::config::get_all_mcp_disabled_tools(
std::path::Path::new(&self.session_info.cwd),
);
let catalog = {
let state = self.managed_mcp_handle.lock().await;
if state.gateway_tools_active {
match &state.gateway_tool_cache {
crate::session::managed_mcp::GatewayToolCatalogCache::Ready(catalog) => {
Some(catalog.clone())
}
_ => None,
}
} else {
None
}
};
let Some(catalog) = catalog else {
return Vec::new();
};
let mut connectors = std::collections::BTreeMap::<String, String>::new();
let mut gateway_connectors: Vec<String> = catalog
.tools
.iter()
.map(|tool| tool.connector_id.clone())
.collect();
gateway_connectors.sort_unstable();
gateway_connectors.dedup();
let mut descriptors = Vec::new();
for tool in &catalog.tools {
if gateway_tool_is_disabled(tool, &disabled_gateway_tools) {
continue;
}
connectors
.entry(tool.connector_id.clone())
.or_insert_with(|| tool.connector_name.clone());
descriptors.push(crate::session::mcp_descriptors::GatewayToolDescriptor {
connector_id: tool.connector_id.clone(),
tool_id: tool.tool_id.clone(),
description: tool.description.clone(),
json_schema: tool.json_schema.clone(),
});
}
connectors
.into_iter()
.map(|(connector_id, connector_name)| McpServerEntry {
folder_path: mcps_root.map(|root| {
crate::session::mcp_descriptors::server_descriptor_dir(root, &connector_id)
.to_string_lossy()
.to_string()
}),
name: connector_id,
server_use_instructions: (!connector_name.trim().is_empty())
.then_some(connector_name),
})
.collect()
}
/// Build a `PathRewriter` for sanitizing overlay paths in model-facing text.
///
/// Returns `None` when `display_cwd` is unset (no rewriting needed). Used
crates/codegen/kigi-shell/src/session/acp_session_impl/run_loop.rs
+2 -37
View File
@@ -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
crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs
-37
View File
@@ -139,9 +139,6 @@ pub(crate) async fn spawn_session_actor(
persisted_announcement_state: Option<crate::session::announcement_state::AnnouncementState>,
memory_config: Option<crate::config::MemoryConfig>,
feedback_flags: crate::session::feedback_manager::FeedbackFlags,
managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle,
managed_mcp_expires_at: Option<chrono::DateTime<chrono::Utc>>,
managed_mcp_proxy_base_url: String,
session_model_id: acp::ModelId,
session_yolo_mode: bool,
session_auto_mode: bool,
@@ -150,8 +147,6 @@ pub(crate) async fn spawn_session_actor(
max_retries: Option<u32>,
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<crate::session::announcement_state::AnnouncementState>,
memory_config: Option<crate::config::MemoryConfig>,
feedback_flags: crate::session::feedback_manager::FeedbackFlags,
managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle,
managed_mcp_expires_at: Option<chrono::DateTime<chrono::Utc>>,
managed_mcp_proxy_base_url: String,
session_model_id: acp::ModelId,
session_yolo_mode: bool,
session_auto_mode: bool,
@@ -1552,8 +1522,6 @@ pub(crate) async fn spawn_session_on_thread(
max_retries: Option<u32>,
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,
crates/codegen/kigi-shell/src/session/acp_session_impl/tool_calls.rs
+1 -52
View File
@@ -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<serde_json::Value> = None;
if let Some((server, _)) =
crate::session::mcp_servers::parse_mcp_tool_name(&prepared.tool_name)
&& server.starts_with(crate::session::managed_mcp::MANAGED_MCP_PREFIX)
{
let auth_rejected = match &result {
Err(err) => kigi_mcp::servers::is_auth_rejection_message(&err.to_string()),
Ok(tool_result) => {
tool_result.output.is_error()
&& kigi_mcp::servers::is_auth_rejection_message(
&tool_result.prompt_text,
)
}
};
if auth_rejected && self.reactive_managed_reauth(&server).await.is_ok() {
result = dispatch_tool(&self.workspace_ops, &prepared, &self.session_info.id.0)
.await;
}
}
let tool_result_size_bytes = match &result {
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,
crates/codegen/kigi-shell/src/session/acp_session_impl/tool_dispatch.rs
+1 -1
View File
@@ -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";
crates/codegen/kigi-shell/src/session/acp_session_impl/turn.rs
-10
View File
@@ -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);
crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs
-8
View File
@@ -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()),
crates/codegen/kigi-shell/src/session/acp_session_tests/fs_injection_regression_tests.rs
-2
View File
@@ -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,
crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs
-2
View File
@@ -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()),
crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs
-6
View File
@@ -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()),
crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs
-2
View File
@@ -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()),
crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_e2e_tests.rs
-476
View File
@@ -1,476 +0,0 @@
//! End-to-end coverage for the reactive managed-MCP re-auth flow that the
//! sibling `reactive_managed_reauth_tests.rs` only exercises at the guard-rail
//! level (owner-scope + cooldown). Here we drive the full
//! `invalidate_cache → get_or_fetch → refresh_managed_clients → re-handshake`
//! loop against a real in-process HTTP MCP server and assert the wire-visible
//! `x.ai/mcp/server_status` pushes that clients consuming only `server_status`
//! (not the `mcp/list` snapshot) depend on.
//!
//! Unlike the unit harness, these tests KEEP `gw_rx` so the forwarded
//! `McpServerStatusPayload`s (`ready`/`managed_token_refreshed` on recovery,
//! `needsauth`/`auth_expired` on terminal exhaustion) can be asserted.
//!
//! The mock is a hand-rolled axum streamable-HTTP server (same shape proven to
//! handshake against rmcp 2.1 in `kigi-mcp/tests/repro_sse_flood.rs`): a
//! `POST` that answers `initialize` + `tools/list` while `reject == false` and
//! `401`s while `reject == true`, plus a standing-GET SSE stream. A separate
//! `GET /mcp/configs` route stands in for the cli-chat-proxy managed-config backend fetch, so
//! the re-auth loop's proxy round-trip is real too.
use crate::session::acp_session::support::*;
use crate::session::acp_session::*;
use crate::session::managed_mcp::{MANAGED_MCP_PREFIX, ManagedMcpConfig};
use crate::session::mcp_dispatcher::{
McpServerStatus, McpServerStatusPayload, McpServerStatusReason, SERVER_STATUS_METHOD,
};
use agent_client_protocol as acp;
use axum::body::Body;
use axum::extract::State;
use axum::http::{StatusCode, header};
use axum::response::{IntoResponse, Response};
use axum::routing::{get, post};
use chrono::Utc;
use kigi_mcp::servers::{ClientStateKind, HttpConfig, McpClient};
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
const MANAGED: &str = "grok_com_testconnector";
// ── Mock cli-chat-proxy + MCP server ──────────────────────────────────────
#[derive(Clone)]
struct MockState {
/// While `true`, the MCP `POST` 401s — a silently-revoked token.
reject: Arc<AtomicBool>,
/// Count of managed-config fetches (`GET /mcp/configs`) so a cooldown-gated
/// attempt can be proven to skip the network entirely.
config_fetches: Arc<AtomicUsize>,
/// The server's own MCP endpoint, echoed back in the managed config so the
/// re-fetched config points the client at this same mock.
mcp_url: String,
}
/// Stand-in for cli-chat-proxy `GET /v1/mcp/configs`: always succeeds (a
/// revoked *connector* token still re-fetches a fresh *proxy* token), so the
/// re-handshake outcome is governed solely by the `reject` flag.
async fn handle_configs(State(s): State<MockState>) -> Response {
s.config_fetches.fetch_add(1, Ordering::Relaxed);
let body = serde_json::json!({
"mcp_servers": [{
"name": "testconnector",
"endpoint": s.mcp_url,
"headers": {"Authorization": "Bearer fresh"},
"token_expires_at": (Utc::now() + chrono::Duration::hours(1)).to_rfc3339(),
"scope": "workspace",
}]
});
axum::Json(body).into_response()
}
/// MCP streamable-HTTP `POST`. 401s while `reject`; otherwise a minimal valid
/// `initialize` result (with the mandatory `mcp-session-id` header) and an
/// empty `tools/list`, so a re-handshake succeeds.
async fn handle_mcp_post(
State(s): State<MockState>,
axum::Json(req): axum::Json<serde_json::Value>,
) -> Response {
if s.reject.load(Ordering::Relaxed) {
return StatusCode::UNAUTHORIZED.into_response();
}
match req["method"].as_str() {
Some("initialize") => {
let result = serde_json::json!({
"jsonrpc": "2.0",
"id": req["id"],
"result": {
"protocolVersion": req["params"]["protocolVersion"],
"capabilities": {},
"serverInfo": {"name": "mock", "version": "0.0.0"},
},
});
([("mcp-session-id", "mock-session-1")], axum::Json(result)).into_response()
}
Some("tools/list") => {
let result = serde_json::json!({
"jsonrpc": "2.0",
"id": req["id"],
"result": {"tools": []},
});
axum::Json(result).into_response()
}
// notifications/initialized and anything else.
_ => StatusCode::ACCEPTED.into_response(),
}
}
/// Standing-GET SSE stream that stays open (a healthy server) — rmcp opens it
/// after `initialize`; a pending body keeps it alive without reconnect churn.
async fn handle_mcp_get() -> Response {
(
[(header::CONTENT_TYPE, "text/event-stream")],
Body::from_stream(futures::stream::pending::<Result<String, std::io::Error>>()),
)
.into_response()
}
/// Bind on an ephemeral port and serve the mock. Returns the proxy base URL
/// (`http://addr` — the managed fetch appends `/mcp/configs`), the MCP endpoint
/// URL, and the fetch counter.
async fn spawn_mock(reject: Arc<AtomicBool>) -> (String, String, Arc<AtomicUsize>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind mock");
let addr = listener.local_addr().expect("addr");
let proxy_base = format!("http://{addr}");
let mcp_url = format!("http://{addr}/mcp");
let config_fetches = Arc::new(AtomicUsize::new(0));
let state = MockState {
reject,
config_fetches: config_fetches.clone(),
mcp_url: mcp_url.clone(),
};
let app = axum::Router::new()
.route("/mcp/configs", get(handle_configs))
.route("/mcp", post(handle_mcp_post).get(handle_mcp_get))
.with_state(state);
tokio::spawn(async move {
let _ = axum::serve(listener, app).await;
});
(proxy_base, mcp_url, config_fetches)
}
// ── Test wiring helpers ────────────────────────────────────────────────────
/// Build an actor with a live (disk-backed, env-free) `AuthManager` holding a
/// valid token and a `ModelsManager` whose cli-chat-proxy points at `proxy_base`
/// — the two pieces the default test actor lacks, both required for the inner
/// re-fetch to reach the mock instead of the real proxy.
async fn actor_with_proxy(
proxy_base: &str,
gw_tx: tokio::sync::mpsc::UnboundedSender<kigi_acp_lib::AcpClientMessage>,
) -> (SessionActor, tempfile::TempDir) {
let (persist_tx, _persist_rx) = tokio::sync::mpsc::unbounded_channel();
let mut actor = create_test_actor(100, 128_000, 80, gw_tx, persist_tx).await;
let home = tempfile::tempdir().expect("tempdir");
let auth_manager = Arc::new(crate::auth::AuthManager::new(
home.path(),
crate::auth::KimiCodeConfig::default(),
));
// Valid (1h) token in-memory only — `auth()` fast-paths it without network.
auth_manager.hot_swap(crate::auth::KimiAuth {
expires_at: Some(Utc::now() + chrono::Duration::hours(1)),
..crate::auth::KimiAuth::test_default()
});
let cfg = crate::agent::config::Config {
endpoints: crate::agent::config::EndpointsConfig {
coding_api_base_url: Some(proxy_base.to_string()),
..Default::default()
},
..Default::default()
};
actor.models_manager = crate::agent::models::ModelsManager::new(
None,
Default::default(),
acp::ModelId::new("default"),
auth_manager.clone(),
cfg,
);
actor.auth_manager = Some(auth_manager);
(actor, home)
}
/// Seed `actor` so it owns a managed client pointed at `mcp_url` with a STALE
/// token, plus the matching config entry `refresh_managed_clients` keys the
/// in-place swap on, plus a `Ready` managed cache.
async fn seed_managed(actor: &SessionActor, mcp_url: &str) {
{
let mut st = actor.mcp_state.lock().await;
// `refresh_managed_clients` matches the owned client to a fresh config
// by looking up `configs` for an Http server with the same name.
st.configs = vec![acp::McpServer::Http(
acp::McpServerHttp::new(MANAGED.to_string(), mcp_url.to_string()).headers(vec![]),
)];
st.owned_clients.insert(
MANAGED.to_string(),
Arc::new(McpClient::new_http(
MANAGED.to_string(),
HttpConfig {
url: mcp_url.to_string(),
headers: vec![("Authorization".into(), "Bearer stale".into())],
},
None,
None,
)),
);
}
let handle = actor.managed_mcp_handle.clone();
let mut configs = HashMap::new();
configs.insert("Authorization".to_string(), "Bearer fresh".to_string());
handle.lock().await.complete_fetch(
vec![ManagedMcpConfig {
name: "testconnector".to_string(),
endpoint: mcp_url.to_string(),
headers: configs,
token_expires_at: Some(Utc::now() + chrono::Duration::hours(1)),
scope: Some("workspace".to_string()),
scope_id: None,
scope_name: None,
}],
&handle,
None,
);
}
/// Drain all `x.ai/mcp/server_status` pushes currently queued on `gw_rx`.
fn drain_status_pushes(
gw_rx: &mut tokio::sync::mpsc::UnboundedReceiver<kigi_acp_lib::AcpClientMessage>,
) -> Vec<McpServerStatusPayload> {
let mut out = Vec::new();
while let Ok(msg) = gw_rx.try_recv() {
if let kigi_acp_lib::AcpClientMessage::ExtNotification(args) = msg
&& args.request.method.as_ref() == SERVER_STATUS_METHOD
&& let Ok(payload) =
serde_json::from_str::<McpServerStatusPayload>(args.request.params.get())
{
out.push(payload);
}
}
out
}
// ── Case 1: recover-on-second-fetch (the happy path) ───────────────────────
/// A managed token rejected on the first re-handshake (`reject = true`) but
/// fixed before the next attempt recovers end-to-end: the client lands `Ready`
/// and a `ready`/`managed_token_refreshed` status push hits the wire.
///
/// The cooldown clear between attempts stands in for the proactive refresh's
/// `clear_reauth_cooldowns` (the real path that re-enables a re-authorized
/// connector) — driving real wall-clock backoff would be slow and flaky.
#[tokio::test(flavor = "multi_thread")]
async fn recovers_on_second_attempt_and_pushes_managed_token_refreshed() {
let reject = Arc::new(AtomicBool::new(true));
let (proxy_base, mcp_url, _fetches) = spawn_mock(reject.clone()).await;
let (gw_tx, mut gw_rx) = tokio::sync::mpsc::unbounded_channel();
let local = tokio::task::LocalSet::new();
local
.run_until(async move {
// `create_test_actor` spawns a local terminal task, so the actor
// must be built inside the `LocalSet`.
let (actor, _home) = actor_with_proxy(&proxy_base, gw_tx).await;
seed_managed(&actor, &mcp_url).await;
// Attempt 1: the connector still 401s, so the re-handshake fails.
let first = actor.reactive_managed_reauth(MANAGED).await;
assert!(first.is_err(), "first attempt must fail while rejecting");
// Below the terminal cap, so no NeedsAuth push yet.
assert!(
drain_status_pushes(&mut gw_rx).is_empty(),
"a single non-terminal failure must not push a status",
);
// Connector re-authorized; clear the cooldown (proactive-refresh
// analog) and retry.
reject.store(false, Ordering::Relaxed);
actor
.managed_mcp_handle
.lock()
.await
.clear_reauth_cooldowns();
actor
.reactive_managed_reauth(MANAGED)
.await
.expect("second attempt must recover once the connector accepts");
// Client is Ready and no longer parked needs-auth.
let client = actor
.mcp_state
.lock()
.await
.get_client(MANAGED)
.cloned()
.expect("managed client present");
assert_eq!(
client.state_kind().await,
ClientStateKind::Ready,
"recovered client must be Ready",
);
assert!(
!actor.mcp_state.lock().await.auth_required.contains(MANAGED),
"recovered server must not be in auth_required",
);
let pushes = drain_status_pushes(&mut gw_rx);
assert!(
pushes.iter().any(|p| p.name == MANAGED
&& p.status == McpServerStatus::Ready
&& p.reason == McpServerStatusReason::ManagedTokenRefreshed),
"expected a ready/managed_token_refreshed push, got: {pushes:?}",
);
})
.await;
}
// ── Case 2: terminal after 3 failures ──────────────────────────────────────
/// Three consecutive failed re-auths park the connector in the terminal
/// `auth_required` state, push `needsauth`/`auth_expired`, and gate the 4th
/// immediate attempt (no extra config fetch).
///
/// Deterministic without sleeps: the default test actor has no `auth_manager`,
/// so the inner re-fetch fails fast (no network). Two prior failures are seeded
/// with an injected past `now` so their backoff windows are already elapsed
/// (the cooldown API takes `now`); the third — the real
/// `reactive_managed_reauth` call — is what crosses the terminal cap and fires
/// the NeedsAuth push.
#[tokio::test(flavor = "current_thread")]
async fn terminal_after_three_failures_pushes_needsauth_then_gates() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (gw_tx, mut gw_rx) = tokio::sync::mpsc::unbounded_channel();
let (persist_tx, _persist_rx) = tokio::sync::mpsc::unbounded_channel();
let actor = create_test_actor(100, 128_000, 80, gw_tx, persist_tx).await;
actor
.mcp_state
.lock()
.await
.owned_clients
.insert(MANAGED.to_string(), Arc::new(McpClient::stub(MANAGED)));
// Seed two prior failures whose backoff windows are already in the
// past (failures = 2, below the cap of 3, so still eligible).
let past = Utc::now() - chrono::Duration::hours(1);
{
let mut h = actor.managed_mcp_handle.lock().await;
h.record_reauth_failure(MANAGED, past);
h.record_reauth_failure(MANAGED, past);
assert!(
h.reauth_allowed(MANAGED, Utc::now()),
"elapsed window + below cap must be eligible",
);
assert!(!h.reauth_is_terminal(MANAGED));
}
// Third (real) attempt: inner fails fast (no auth manager) and
// records the terminal failure.
let err = actor
.reactive_managed_reauth(MANAGED)
.await
.expect_err("third attempt must fail");
assert!(err.contains("auth manager"), "got: {err}");
// Terminal: parked auth_required + NeedsAuth/auth_expired push.
assert!(
actor.mcp_state.lock().await.auth_required.contains(MANAGED),
"exhausted connector must be parked in auth_required",
);
assert!(
actor
.managed_mcp_handle
.lock()
.await
.reauth_is_terminal(MANAGED),
);
let pushes = drain_status_pushes(&mut gw_rx);
assert!(
pushes.iter().any(|p| p.name == MANAGED
&& p.status == McpServerStatus::NeedsAuth
&& p.reason == McpServerStatusReason::AuthExpired),
"expected a needsauth/auth_expired push, got: {pushes:?}",
);
// Fourth immediate attempt is cooldown-gated by the terminal state.
let err4 = actor
.reactive_managed_reauth(MANAGED)
.await
.expect_err("fourth attempt must be gated");
assert!(err4.contains("cooldown"), "got: {err4}");
assert!(
drain_status_pushes(&mut gw_rx).is_empty(),
"a gated attempt must not push another status",
);
})
.await;
}
// ── Case 3: entry-B classification ─────────────────────────────────────────
/// The mid-session entry-B gate routes an auth-rejection on a managed tool into
/// `reactive_managed_reauth` (observable as an armed cooldown) but leaves a
/// non-auth `Ok(is_error)` body (e.g. a 403 policy denial) untouched. Mirrors
/// the exact classifier (`is_auth_rejection_message`) + managed-prefix gate the
/// loop in `tool_calls.rs` keys on, plus the resulting side effect.
#[tokio::test(flavor = "current_thread")]
async fn entry_b_routes_auth_rejection_but_not_policy_denial() {
use kigi_mcp::servers::{is_auth_rejection_message, parse_mcp_tool_name};
// Managed-prefix gate: only `grok_com_*` tools enter entry-B.
let (managed_server, _) =
parse_mcp_tool_name(&format!("{MANAGED}__create_issue")).expect("qualified name");
assert!(managed_server.starts_with(MANAGED_MCP_PREFIX));
let (local_server, _) = parse_mcp_tool_name("github__create_issue").expect("qualified name");
assert!(!local_server.starts_with(MANAGED_MCP_PREFIX));
// Classification of both failure shapes entry-B sees:
// Err(ToolError) -> err.to_string()
// Ok(is_error) -> tool_result.prompt_text
assert!(
is_auth_rejection_message("MCP error: HTTP 401 Unauthorized"),
"401 error string must route into re-auth",
);
assert!(
is_auth_rejection_message("authentication required"),
"auth wording must route into re-auth",
);
assert!(
!is_auth_rejection_message("403 Forbidden: policy denied for this connector"),
"a 403 policy denial must NOT route into re-auth",
);
// Side-effect check: only the auth path arms the cooldown. An owner stub
// makes the inner re-fetch fail fast (no auth manager), arming the gate.
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (gw_tx, _gw_rx) = tokio::sync::mpsc::unbounded_channel();
let (persist_tx, _persist_rx) = tokio::sync::mpsc::unbounded_channel();
let actor = create_test_actor(100, 128_000, 80, gw_tx, persist_tx).await;
actor
.mcp_state
.lock()
.await
.owned_clients
.insert(MANAGED.to_string(), Arc::new(McpClient::stub(MANAGED)));
// Non-auth body classified false -> entry-B does NOT call re-auth,
// so the cooldown stays pristine.
assert!(
actor
.managed_mcp_handle
.lock()
.await
.reauth_allowed(MANAGED, Utc::now()),
"no re-auth call yet: cooldown must be clean",
);
// Auth body classified true -> entry-B calls re-auth; the failed
// attempt arms the cooldown window.
let _ = actor.reactive_managed_reauth(MANAGED).await;
assert!(
!actor
.managed_mcp_handle
.lock()
.await
.reauth_allowed(MANAGED, Utc::now()),
"the auth path must have armed the cooldown",
);
})
.await;
}
crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_tests.rs
-97
View File
@@ -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;
}
crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs
-2
View File
@@ -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()),
crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs
-4
View File
@@ -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()),
crates/codegen/kigi-shell/src/session/acp_session_tests/media_gen_auth_retry_tests.rs → crates/codegen/kigi-shell/src/session/acp_session_tests/tool_auth_retry_tests.rs
+9 -9
View File
@@ -77,7 +77,7 @@ fn err(msg: &str) -> Result<ToolRunResult, kigi_tool_runtime::ToolError> {
))
}
/// Production-shaped HTTP failure (image_gen / video_gen emit this on
/// Production-shaped HTTP failure (tool HTTP clients emit this on
/// any non-success status). Use for retry tests that should exercise
/// the structured status-code path rather than the string fallback.
fn http_err(status: u16, msg: &str) -> Result<ToolRunResult, kigi_tool_runtime::ToolError> {
@@ -98,14 +98,14 @@ fn is_auth_tool_error_classification() {
// (expected, error) — covers every branch + a sample of negatives
// 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;
crates/codegen/kigi-shell/src/session/acp_session_tests/web_search_e2e_tests.rs
-2
View File
@@ -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,
crates/codegen/kigi-shell/src/session/agent_rebuild.rs
-16
View File
@@ -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<tokio::sync::Mutex<crate::session::mcp_servers::McpState>>,
pub managed_gateway_tool_client: Option<kigi_tools::types::resources::ManagedGatewayToolClient>,
pub is_non_interactive: bool,
pub system_prompt_label: String,
pub owner_session_id: Option<String>,
@@ -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<AgentRebuildSpec> {
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<AgentRebuildSpec> {
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()),
crates/codegen/kigi-shell/src/session/commands.rs
+1 -6
View File
@@ -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<acp::McpServer>,
respond_to: oneshot::Sender<Result<(), acp::Error>>,
@@ -357,17 +357,12 @@ pub enum SessionCommand {
server_name: String,
tool_name: String,
enabled: bool,
is_managed_gateway: bool,
respond_to: oneshot::Sender<Result<(), acp::Error>>,
},
/// Read MCP status: which servers are configured, which clients are healthy, what tools.
GetMcpStatus {
respond_to: oneshot::Sender<crate::extensions::mcp::McpStatusSnapshot>,
},
GetManagedGatewayDisabledTools {
respond_to:
oneshot::Sender<std::collections::HashMap<String, std::collections::HashSet<String>>>,
},
/// Snapshot the session's live MCP client pool for subagent inheritance.
SnapshotMcpPool {
respond_to: oneshot::Sender<Option<crate::session::mcp_servers::SharedMcpPool>>,
crates/codegen/kigi-shell/src/session/compaction.rs
-2
View File
@@ -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()),
crates/codegen/kigi-shell/src/session/handle.rs
-32
View File
@@ -130,7 +130,6 @@ pub struct SessionHandle {
pub attribution_callback: Option<kigi_sampler::SharedAttributionCallback>,
/// The agent definition name for this session.
pub agent_name: String,
pub managed_mcp_proxy_base_url: String,
pub session_default_agent_profile: Option<String>,
/// Subagent types this agent can spawn (from Agent(t1, t2) in tools).
pub allowed_subagent_types: Option<Vec<String>>,
@@ -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<String, HashSet<String>> {
let (tx, rx) = oneshot::channel();
if self
.cmd_tx
.send(SessionCommand::GetManagedGatewayDisabledTools { respond_to: tx })
.is_err()
{
return HashMap::new();
}
rx.await.unwrap_or_default()
}
pub async fn retry_auth_required_servers(&self) {
let (tx, rx) = oneshot::channel();
if self
crates/codegen/kigi-shell/src/session/managed_mcp.rs
+16 -262
View File
@@ -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<crate::auth::AuthManager>,
) -> RefreshContext {
RefreshContext {
proxy_base_url,
token_provider: Arc::new(move || -> TokenFuture {
let auth_manager = auth_manager.clone();
Box::pin(async move { auth_manager.get_valid_token().await.ok() })
}),
}
}
/// Resolve an auth key from `auth_manager` then [`get_or_fetch`] the managed MCP
/// configs (with a [`RefreshContext`] for proactive refresh). Single source for
/// the auth-key dance across every managed-config fetch —
/// [`crate::agent::MvpAgent::get_managed_mcp_configs`], the interactive
/// folder-trust grant reload, agent-init MCP setup, and the reactive re-auth
/// re-fetch — so the copies can't drift.
/// Callers gate on `can_fetch_managed_mcps`/auth before calling.
pub(crate) async fn fetch_managed_mcp_configs(
handle: &ManagedMcpStateHandle,
proxy_url: &str,
auth_manager: &Arc<crate::auth::AuthManager>,
) -> Vec<ManagedMcpConfig> {
let auth_key = auth_manager
.get_valid_token()
.await
.ok()
.or_else(|| auth_manager.current_or_expired().map(|a| a.key));
get_or_fetch(
handle,
proxy_url,
auth_key.as_deref(),
Some(refresh_context(proxy_url.to_string(), auth_manager.clone())),
)
.await
/// 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<acp::McpServer>,
cwd: &std::path::Path,
managed_configs: &[ManagedMcpConfig],
plugin_registry: Option<&kigi_agent::plugins::PluginRegistry>,
compat: &kigi_tools::types::compat::CompatConfig,
) -> Vec<acp::McpServer> {
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<acp::McpServer>,
cwd: &std::path::Path,
managed_configs: &[ManagedMcpConfig],
plugin_registry: Option<&kigi_agent::plugins::PluginRegistry>,
compat: &kigi_tools::types::compat::CompatConfig,
) -> Vec<McpServerWithPolicy> {
@@ -125,8 +72,6 @@ pub fn merge_managed_mcp_servers_with_policy(
let disabled = crate::util::config::disabled_mcp_server_names(cwd);
let mut merged: Vec<acp::McpServer> = servers.into_values().collect();
inject_managed_headers(&mut merged, managed_configs);
auto_inject_managed_servers_with_disabled(&mut merged, managed_configs, &disabled);
// Deterministic order: this list is collected from a HashMap (random
// 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<acp::McpServer>,
managed_configs: &[ManagedMcpConfig],
disabled_names: &std::collections::HashSet<String>,
) {
if managed_configs.is_empty() {
return;
}
let existing_names: std::collections::HashSet<String> = merged
.iter()
.map(|s| mcp_server_name(s).to_owned())
.collect();
let mut seen_display_names: std::collections::HashSet<String> =
std::collections::HashSet::new();
let mut count = 0usize;
for config in managed_configs {
if config.headers.is_empty() {
continue;
}
let name = to_managed_name(&config.name);
if existing_names.contains(&name) {
continue;
}
if disabled_names.contains(&name) {
tracing::debug!(server_name = %name, "Auto-inject skipped: disabled in config.toml");
continue;
}
if !seen_display_names.insert(config.name.to_lowercase()) {
continue;
}
let headers = config
.headers
.iter()
.map(|(k, v)| acp::HttpHeader::new(k.clone(), v.clone()))
.collect();
merged.push(acp::McpServer::Http(
acp::McpServerHttp::new(name, config.endpoint.clone()).headers(headers),
));
count += 1;
}
if count > 0 {
tracing::info!(count, "Auto-injected managed MCP connectors");
}
}
fn load_plugin_mcp_servers(
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<String> =
["grok_com_slack".to_string()].into_iter().collect();
let mut merged = vec![];
auto_inject_managed_servers_with_disabled(&mut merged, &managed, &disabled);
let has_slack = merged
.iter()
.any(|s| matches!(s, acp::McpServer::Http(acp::McpServerHttp { name, .. }) if name == "grok_com_slack"));
let has_linear = merged
.iter()
.any(|s| matches!(s, acp::McpServer::Http(acp::McpServerHttp { name, .. }) if name == "grok_com_linear"));
assert!(!has_slack, "disabled connector should be skipped");
assert!(has_linear, "non-disabled connector should be injected");
}
#[test]
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<String> =
["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"),
crates/codegen/kigi-shell/src/session/mcp_descriptors.rs
-354
View File
@@ -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: `<mcps_root>/<sanitized server name>`. 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<GatewayToolDescriptor>,
gateway_connectors: Vec<String>,
protected_connectors: HashSet<String>,
) {
let mut files_by_connector: BTreeMap<String, Vec<(String, Vec<u8>)>> = BTreeMap::new();
for tool in tools {
let descriptor = serde_json::json!({
"name": tool.tool_id,
"description": tool.description,
"inputSchema": tool.json_schema,
});
match serde_json::to_vec_pretty(&descriptor) {
Ok(bytes) => {
let file_name = format!("{}.json", sanitize_descriptor_segment(&tool.tool_id));
files_by_connector
.entry(tool.connector_id)
.or_default()
.push((file_name, bytes));
}
Err(e) => tracing::warn!(
connector = %tool.connector_id,
tool = %tool.tool_id,
error = %e,
"failed to serialize managed gateway tool descriptor",
),
}
}
let connectors: BTreeSet<String> = gateway_connectors.into_iter().collect();
let mcps_root = mcps_root.to_path_buf();
if let Err(e) = tokio::task::spawn_blocking(move || {
for connector_id in connectors {
let server_dir = server_descriptor_dir(&mcps_root, &connector_id);
let tools_dir = server_dir.join("tools");
let files = files_by_connector.remove(&connector_id).unwrap_or_default();
let admitted: BTreeSet<String> = files.iter().map(|(name, _)| name.clone()).collect();
let protected = protected_connectors.contains(&connector_id);
if protected {
if let Err(e) = remove_gateway_owned_tool_descriptors(&tools_dir)
&& e.kind() != std::io::ErrorKind::NotFound
{
tracing::warn!(
connector = %connector_id,
path = %tools_dir.display(),
error = %e,
"failed to remove managed gateway descriptors from protected connector dir",
);
}
continue;
}
if files.is_empty() {
if let Err(e) = std::fs::remove_dir_all(&server_dir)
&& e.kind() != std::io::ErrorKind::NotFound
{
tracing::warn!(
connector = %connector_id,
path = %server_dir.display(),
error = %e,
"failed to remove stale managed gateway descriptor dir",
);
}
continue;
}
if let Err(e) = std::fs::create_dir_all(&tools_dir) {
tracing::warn!(
connector = %connector_id,
path = %tools_dir.display(),
error = %e,
"failed to create managed gateway tools descriptor dir",
);
continue;
}
for (file_name, bytes) in files {
let path = tools_dir.join(&file_name);
let write_result =
tempfile::NamedTempFile::new_in(&tools_dir).and_then(|mut tmp| {
std::io::Write::write_all(&mut tmp, &bytes)?;
tmp.persist(&path).map_err(|e| e.error)
});
if let Err(e) = write_result {
tracing::warn!(
connector = %connector_id,
path = %path.display(),
error = %e,
"failed to write managed gateway tool descriptor",
);
}
}
prune_stale_gateway_tool_descriptors(&tools_dir, &admitted, &connector_id);
}
})
.await
{
tracing::warn!(
error = %e,
"managed gateway descriptor write task panicked",
);
}
}
fn remove_gateway_owned_tool_descriptors(tools_dir: &Path) -> std::io::Result<()> {
let entries = std::fs::read_dir(tools_dir)?;
for entry in entries.flatten() {
let path = entry.path();
if path.extension().and_then(|ext| ext.to_str()) != Some("json") {
continue;
}
if is_gateway_owned_descriptor(&path) {
std::fs::remove_file(&path)?;
}
}
Ok(())
}
fn is_gateway_owned_descriptor(path: &Path) -> bool {
let Ok(contents) = std::fs::read(path) else {
return false;
};
let Ok(value) = serde_json::from_slice::<serde_json::Value>(&contents) else {
return false;
};
value
.get("x-grok-managed-gateway")
.and_then(|v| v.as_bool())
== Some(true)
}
fn prune_stale_gateway_tool_descriptors(
tools_dir: &Path,
admitted: &BTreeSet<String>,
connector_id: &str,
) {
let entries = match std::fs::read_dir(tools_dir) {
Ok(entries) => entries,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return,
Err(e) => {
tracing::warn!(
connector = %connector_id,
path = %tools_dir.display(),
error = %e,
"failed to read managed gateway tools descriptor dir for pruning",
);
return;
}
};
for entry in entries.flatten() {
let path = entry.path();
if path.extension().and_then(|ext| ext.to_str()) != Some("json") {
continue;
}
let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else {
continue;
};
if admitted.contains(file_name) {
continue;
}
if let Err(e) = std::fs::remove_file(&path) {
tracing::warn!(
connector = %connector_id,
path = %path.display(),
error = %e,
"failed to remove stale managed gateway tool descriptor",
);
}
}
}
#[cfg(test)]
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());
}
}
crates/codegen/kigi-shell/src/session/mcp_dispatcher.rs
+14 -118
View File
@@ -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<String> = ["admin_off".to_string()].into_iter().collect();
let got = recoverable_http_servers(&configs, &disabled);
crates/codegen/kigi-shell/src/session/mcp_servers.rs
+1 -8
View File
@@ -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 }
crates/codegen/kigi-shell/src/session/slash_commands.rs
+2 -2
View File
@@ -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"]})
);
}
crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs
-4
View File
@@ -83,8 +83,6 @@ pub(crate) fn ctx_with_toggle(toggle: HashMap<String, bool>) -> 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<String, bool>) -> 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,
crates/codegen/kigi-shell/src/util/config/mcp.rs
+1 -3
View File
@@ -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 {