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/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 {