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/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"]})
);
}