会话生命周期与 MCP 稳定性

7impl acp::Agent for MvpAgent {
8    /// In the meta, we provide
9    ///   - model_state: the model state, useful for the client to display available models and the default model.
10    ///
11    /// SINGLE-CALL INVARIANT: this method is the sole writer of
12    /// `self.auth_method_id` during initialization. It is called exactly once
13    /// per agent process by the ACP server before any session-creating
14    /// requests, while `auth_method_id` is still `None` (initialized at
15    /// `MvpAgent::new`). The auth-method block below relies on that
16    /// invariant when it unconditionally writes the default id returned by
17    /// `auth_method::build_auth_methods`. If you ever need to call
18    /// `initialize()` more than once, restore an `is_none()` guard around
19    /// the `auth_method_id` write at the call site so a re-init doesn't
20    /// silently downgrade an api-key user to a session-token user.
21    async fn initialize(
22        &self,
23        arguments: acp::InitializeRequest,
24    ) -> Result<acp::InitializeResponse, acp::Error> {
25        tracing::debug!(target : "sampling_log", "Received initialize request");
26        xai_grok_telemetry::unified_log::info("agent initialized", None, None);
27        self.start_subagent_coordinator();
28        tokio::task::spawn_blocking(|| {
29            crate::session::worktree_pool::cleanup_stale_pool_worktrees(None);
30        });
31        tokio::task::spawn_blocking(|| {
32            crate::session::persistence::cleanup_stale_sessions(None);
33        });
34        {
35            let root = crate::util::grok_home::grok_home();
36            crate::session::storage::search::SEARCH_INDEX_MANAGER.bootstrap_once(root);
37        }
38        const PERMISSION_CLEANUP_TTL_DAYS: u64 = 30;
39        static CLEANUP_PERMISSIONS_ONCE: std::sync::Once = std::sync::Once::new();
40        CLEANUP_PERMISSIONS_ONCE
41            .call_once(|| {
42                tokio::task::spawn(
43                    xai_grok_workspace::permission::cleanup_stale_permission_state(
44                        std::time::Duration::from_secs(
45                            PERMISSION_CLEANUP_TTL_DAYS * 24 * 60 * 60,
46                        ),
47                    ),
48                );
49            });
50        xai_grok_workspace::trust::migrate_legacy_hook_trust();

77        let mut client_type = arguments
78            .meta
79            .as_ref()
80            .and_then(|m| m.get("clientType"))
81            .and_then(|v| serde_json::from_value::<ClientType>(v.clone()).ok())
82            .unwrap_or_default();
83        let client_identifier = arguments
84            .meta
85            .as_ref()
86            .and_then(|m| m.get("clientIdentifier"))
87            .and_then(|v| v.as_str())
88            .map(|s| s.to_string());
89        if let Some(ref id) = client_identifier {
90            tracing::info!("Client identifier set to: {}", id);
91        }
92        if client_type == ClientType::Generic {
93            match client_identifier.as_deref() {
94                Some("grok-web") => client_type = ClientType::GrokWeb,
95                Some("nebula") => client_type = ClientType::Nebula,
96                Some("grok-code-extension") => client_type = ClientType::Extension,
97                Some("grok-desktop") => client_type = ClientType::Desktop,
98                _ => {}
99            }
100        }
101        *self.client_type.borrow_mut() = client_type;
102        tracing::info!("Client type set to: {:?}", client_type);
103        let code_nav_enabled = Self::parse_code_nav_capability(&arguments);
104        self.code_nav_enabled.set(code_nav_enabled);
105        tracing::info!(
106            code_nav_enabled, client_type = ? client_type, event =
107            "code_nav_capability_parsed",
108            "code-nav capability initialized from initialize request; \
109             index will start lazily on first x.ai/code/* request if eligible"
110        );
111        let interactive_trust_client = Self::parse_interactive_trust_capability(
112            &arguments,
113        );
114        self.interactive_trust_client.set(interactive_trust_client);
115        let client_supports_mcp_apps = arguments
116            .meta
117            .as_ref()
118            .and_then(|m| m.get("mcpApps"))
119            .and_then(|v| v.as_bool())
120            .unwrap_or(false);
121        if client_supports_mcp_apps {
122            tracing::info!("Client supports MCP Apps");
123        }
124        let buffering_settings = arguments
125            .meta
126            .as_ref()
127            .and_then(|m| m.get("bufferingSettings"))
128            .map(|value| serde_json::from_value::<
129                update_chunk_merge::BufferingSettings,
130            >(value.clone()))
131            .transpose()
132            .map_err(|err| {
133                tracing::warn!(
134                    error = ? err, "Failed to parse buffering settings from init meta"
135                );
136                err
137            })
138            .unwrap_or(None);
139        tracing::info!(? buffering_settings, "Buffering settings from init");
140        *self.buffering_settings.borrow_mut() = buffering_settings;
141        if self.initialize_request.set(arguments).is_err() {
142            tracing::info!("Initialize called on reconnect (already initialized)");
143        }

101    fn has_managed_mcp_auth(&self) -> bool {
102        self.auth_manager
103            .current_or_expired()
104            .is_some_and(|a| a.is_managed_mcp_eligible())
105    }
106    /// Requires feature flag AND xAI authentication (OIDC or legacy WebLogin).
107    pub(super) fn can_fetch_managed_mcps(&self) -> bool {
108        let cfg = self.cfg.borrow();
109        cfg.managed_mcps_enabled && !cfg.managed_mcp_gateway_tools_enabled
110            && self.has_managed_mcp_auth()
111    }
112    fn can_fetch_managed_mcp_gateway_tools(&self) -> bool {
113        self.cfg.borrow().managed_mcp_gateway_tools_enabled
114            && self.has_managed_mcp_auth()
115    }
116    pub async fn get_managed_mcp_configs(
117        &self,
118    ) -> Vec<crate::session::managed_mcp::ManagedMcpConfig> {
119        if !self.can_fetch_managed_mcps() {
120            return vec![];
121        }
122        let proxy_url = self.cfg.borrow().endpoints.proxy_url();
123        crate::session::managed_mcp::fetch_managed_mcp_configs(
124                &self.managed_mcp_cache,
125                &proxy_url,
126                &self.auth_manager,
127            )
128            .await
129    }
130    pub async fn get_managed_mcp_gateway_tool_catalog(
131        &self,
132    ) -> Option<crate::session::managed_mcp::GatewayToolCatalog> {
133        if !self.can_fetch_managed_mcp_gateway_tools() {
134            self.managed_mcp_cache.lock().await.disable_gateway_tools();
135            return None;
136        }
137        self.managed_mcp_cache.lock().await.enable_gateway_tools();
138        let proxy_url = self.cfg.borrow().endpoints.proxy_url();
139        let auth_key = self
140            .auth_manager
141            .get_valid_token()
142            .await
143            .ok()
144            .or_else(|| self.auth_manager.current_or_expired().map(|a| a.key));
145        crate::session::managed_mcp::get_or_fetch_gateway_tool_catalog(
146                &self.managed_mcp_cache,
147                &proxy_url,
148                auth_key.as_deref(),
149            )
150            .await
151    }

174    pub(crate) fn spawn_managed_gateway_tool_catalog_fetch(&self) {
175        let session_txs: Vec<_> = self
176            .sessions
177            .borrow()
178            .values()
179            .map(|handle| handle.cmd_tx.clone())
180            .collect();
181        if !self.can_fetch_managed_mcp_gateway_tools() {
182            self.disable_managed_gateway_tools_and_refresh_sessions_with_txs(
183                session_txs,
184            );
185            return;
186        }
187        let cache = self.managed_mcp_cache.clone();
188        let proxy_url = self.cfg.borrow().endpoints.proxy_url();
189        let auth_manager = self.auth_manager.clone();
190        tokio::task::spawn_local(async move {
191            let auth_key = auth_manager
192                .get_valid_token()
193                .await
194                .ok()
195                .or_else(|| auth_manager.current_or_expired().map(|a| a.key));
196            if !auth_manager
197                .current_or_expired()
198                .is_some_and(|a| a.is_managed_mcp_eligible())
199            {
200                cache.lock().await.disable_gateway_tools();
201                for tx in session_txs {
202                    let _ = tx.send(SessionCommand::RefreshMcpSearchIndex);
203                }
204                return;
205            }
206            cache.lock().await.enable_gateway_tools();
207            crate::session::managed_mcp::get_or_fetch_gateway_tool_catalog(
208                    &cache,
209                    &proxy_url,
210                    auth_key.as_deref(),
211                )
212                .await;
213            for tx in session_txs {
214                let _ = tx.send(SessionCommand::RefreshMcpSearchIndex);
215            }
216        });
217    }
218    /// Resolve the launch dir's project-scope trust verdict ONCE and return it
219    /// with its path.
220    ///
221    /// Memoizes the single [`folder_trust::resolve_launch_dir_trust`] gather (see
222    /// it for the dedup + TOCTOU contract) so the two one-shot init helpers
223    /// (`ensure_plugin_registry` and `ensure_local_workspace_ops`) share it
224    /// instead of each re-scanning. They share a single point-in-time verdict
225    /// rather than two independent re-scans; the sub-millisecond, startup-only
226    /// window between them is intentional (the cross-session TOCTOU re-scan is
227    /// preserved per the contract).
228    fn prime_launch_dir_trust(&self) -> (&std::path::Path, bool) {
229        let trust = *self
230            .launch_dir_trust
231            .get_or_init(|| {
232                let remote_settings = self.cfg.borrow().remote_settings.clone();
233                folder_trust::resolve_launch_dir_trust(
234                    &self.launch_cwd,
235                    remote_settings.as_ref(),
236                )
237            });
238        (&self.launch_cwd, trust)
239    }

1//! Single source of truth for the `x.ai/mcp/*` ACP wire strings.
2//!
3//! These method/`_meta` keys are part of the cross-language MCP-over-ACP
4//! protocol the SDK speaks (mirrors the SDK's `_mcp_wire.py` / `mcpWire.ts`).
5//! Reference these constants instead of re-typing the literals so the agent and
6//! SDK can't drift apart.
7
8/// Forward tool-invocation method (client -> agent): `x.ai/mcp/call`.
9///
10/// The pager/client asks the agent to invoke an MCP tool on a server the agent is
11/// connected to, outside the LLM loop. See `extensions::mcp::handle_call`.
12pub const MCP_CALL: &str = "x.ai/mcp/call";
13
14/// Reverse zero-IPC tool-invocation method (agent -> client): `x.ai/mcp/sdk_call`.
15///
16/// The agent invokes a tool that lives in the SDK's in-process MCP server by sending
17/// the MCP JSON-RPC message back to the client over the ACP reverse channel. Distinct
18/// from [`MCP_CALL`] so the two disjoint schemas don't share a method string for
19/// metrics/tracing. See the agent-side ACP invoker that handles this method.
20pub const MCP_SDK_CALL: &str = "x.ai/mcp/sdk_call";
21
22/// `session/new` `_meta` key listing in-process SDK MCP servers: `x.ai/mcp/servers`.
23pub const MCP_SERVERS: &str = "x.ai/mcp/servers";
24
25/// `initialize` `_meta` capability flag advertising in-process SDK MCP support
26/// (enables the SDK's `transport="acp"`): `x.ai/mcp/sdk`.
27pub const MCP_SDK: &str = "x.ai/mcp/sdk";

6    /// Ask a live session actor to shut down.
7    pub(crate) fn request_session_shutdown(&self, id: &acp::SessionId) {
8        if let Some(handle) = self.sessions.borrow().get(id) {
9            let _ = handle.cmd_tx.send(SessionCommand::Shutdown);
10        }
11    }
12    /// Finalize the cloud session replica (fire-and-forget, "Hook 4").
13    ///
14    /// Marks the session **done** upstream, so this MUST only run on a genuine
15    /// session end — a terminal/explicit close (`x.ai/session/close`). It must
16    /// NOT run on a mere client disconnect or a dead-actor reap: those leave the
17    /// conversation resumable on disk, and finalizing would wrongly mark a still
18    /// running/resumable session "done".
19    pub(super) fn finalize_session_replica(&self, id: &acp::SessionId) {
20        #[cfg(test)]
21        self.finalize_spy.borrow_mut().push(id.0.to_string());
22        if let Some(client) = self.session_registry_client() {
23            let sid = id.0.to_string();
24            tokio::spawn(async move {
25                if let Err(e) = client.finalize(&sid).await {
26                    tracing::warn!(
27                        error = % e, "session registry finalize failed (non-fatal)"
28                    );
29                }
30            });
31        }
32    }
33    /// Remove a session and its thread handle **without** finalizing the cloud
34    /// replica.
35    ///
36    /// Used for dead-actor reaping and idle-unload: the conversation stays
37    /// resumable on disk, so it must NOT be marked "done" upstream. Genuine
38    /// terminal closes go through [`MvpAgent::close_session_explicit`]. Also
39    /// drops the `session_live_state` entry so that map stays bounded.
40    pub(crate) fn remove_session(&self, id: &acp::SessionId) {
41        self.sessions.borrow_mut().remove(id);
42        self.prompt_intake_locks.borrow_mut().remove(id);
43        self.session_threads.borrow_mut().remove(id);
44        self.session_index_claims.borrow_mut().remove(id);
45        self.require_gateway_sessions.borrow_mut().remove(id);
46        self.session_live_state.borrow_mut().remove(id);
47    }
48    /// Get-or-create the per-session prompt-intake lock (see
49    /// [`Self::prompt_intake_locks`]). Cheap clone of the shared `Rc`.
50    pub(super) fn prompt_intake_lock(
51        &self,
52        id: &acp::SessionId,
53    ) -> std::rc::Rc<tokio::sync::Mutex<()>> {
54        self.prompt_intake_locks
55            .borrow_mut()
56            .entry(id.clone())
57            .or_default()
58            .clone()
59    }
60    /// Close a session in response to an **explicit** terminal close
61    /// (`x.ai/session/close`). Finalizes the cloud replica (genuine session
62    /// end), then removes the session terminally as `Completed`.
63    pub(crate) fn close_session_explicit(&self, id: &acp::SessionId) {
64        self.finalize_session_replica(id);
65        self.remove_session_terminal(id, SessionLiveState::Completed);
66    }

78    /// Roster-delta hook for a terminally removed session. Broadcasts an
79    /// `x.ai/sessions/changed` notification with the session in `removed` so
80    /// every attached dashboard drops the row promptly. Also
81    /// records the call site (and the terminal state) for test observability,
82    /// since the `session_live_state` entry is dropped on removal.
83    pub(super) fn record_roster_delta(&self, id: &acp::SessionId, final_state: SessionLiveState) {
84        #[cfg(test)]
85        self.roster_delta_spy
86            .borrow_mut()
87            .push((id.0.to_string(), final_state));
88        tracing::debug!(
89            session_id = % id.0, ? final_state, "roster delta: session removed"
90        );
91        self.emit_roster_changed(Vec::new(), vec![id.0.to_string()]);
92    }
93    /// Roster-delta hook for a newly-resident / changed session. Broadcasts an
94    /// `x.ai/sessions/changed` notification with the current entry in
95    /// `upserted` so dashboards add/refresh the row.
96    pub(crate) fn push_roster_delta_upserted(&self, id: &acp::SessionId) {
97        if let Some(entry) = self.resident_roster_entry(id) {
98            self.emit_roster_changed(vec![entry], Vec::new());
99        }
100    }
101    /// Emit an `x.ai/sessions/changed` upsert for a resident session with an
102    /// explicit `activity`, so every attached dashboard reflects a
103    /// turn-boundary transition (Working / Idle / NeedsInput) *immediately*
104    /// rather than waiting for the ≤1s roster poll (deltas are emitted
105    /// at turn-start/turn-end). Without this, a viewer client that holds no
106    /// local `AgentView` for the session only learns its activity from the
107    /// poll, so a turn driven by another client shows as `Idle` for up to a
108    /// poll interval — and not at all while that viewer's poll is dormant.
109    ///
110    /// The `activity` is supplied by the caller rather than read from
111    /// `resident_activity` because at turn-start the actor may not have
112    /// published `current_prompt_id` yet (it is set asynchronously once the
113    /// actor dequeues the `SessionCommand::Prompt`), so a natural read would
114    /// still observe `Idle`. The authoritative entry (cwd / worktree / model /
115    /// yolo) is built by `resident_roster_entry`, so it never diverges from
116    /// the polled entry; only the `activity` field is overridden.
117    pub(super) fn push_roster_activity_delta(
118        &self,
119        id: &acp::SessionId,
120        activity: crate::agent::roster::RosterActivity,
121    ) {
122        if let Some(mut entry) = self.resident_roster_entry(id) {
123            entry.activity = activity;
124            self.emit_roster_changed(vec![entry], Vec::new());
125        }
126    }
127    /// Fan an `x.ai/sessions/changed` delta out to every attached client.
128    ///
129    /// This is a roster-wide notification (no `sessionId`), so the leader IPC
130    /// server broadcasts it to all clients rather than routing by session (see
131    /// the `x.ai/sessions/changed` special-case in `leader/server.rs`).
132    pub(super) fn emit_roster_changed(
133        &self,
134        upserted: Vec<crate::agent::roster::RosterEntry>,
135        removed: Vec<String>,
136    ) {
137        if upserted.is_empty() && removed.is_empty() {
138            return;
139        }
140        let payload = crate::agent::roster::RosterChanged { upserted, removed };
141        if let Ok(params) = serde_json::value::to_raw_value(&payload) {
142            self.gateway
143                .forward_fire_and_forget(acp::ExtNotification::new(
144                    crate::agent::roster::SESSIONS_CHANGED_METHOD,
145                    params.into(),
146                ));
147        }
148    }
149    /// Coarse activity of a resident session for the dashboard status column.
150    ///
151    /// Precedence: a non-empty pending-interaction map →
152    /// `NeedsInput` (wins even over a running turn — a session awaiting a
153    /// permission *mid-turn* is "needs input"); else a running turn →
154    /// `Working`; else map the coarse `SessionLiveState`.
155    pub(super) fn resident_activity(
156        &self,
157        id: &acp::SessionId,
158    ) -> crate::agent::roster::RosterActivity {
159        use crate::agent::roster::RosterActivity;
160        let (needs_input, turn_running) = self
161            .sessions
162            .borrow()
163            .get(id)
164            .map(|h| {
165                let needs_input = h
166                    .pending_interactions
167                    .lock()
168                    .map(|g| !g.is_empty())
169                    .unwrap_or(false);
170                let turn_running = h
171                    .current_prompt_id
172                    .lock()
173                    .map(|g| g.is_some())
174                    .unwrap_or(false);
175                (needs_input, turn_running)
176            })
177            .unwrap_or((false, false));
178        if needs_input {
179            return RosterActivity::NeedsInput;
180        }
181        if turn_running {
182            return RosterActivity::Working;
183        }
184        match self.session_live_state.borrow().get(id).copied() {
185            Some(SessionLiveState::Completed) => RosterActivity::Completed,
186            Some(SessionLiveState::DeadFailed) => RosterActivity::Dead,
187            Some(SessionLiveState::Dormant) => RosterActivity::Dormant,
188            _ => RosterActivity::Idle,
189        }
190    }

110/// The flood: body death on the standing GET -> zero-backoff reconnect loop,
111/// one WARN per GET (set RUST_LOG=rmcp=warn to see them).
112#[tokio::test(flavor = "multi_thread")]
113async fn repro_zero_backoff_reconnect_flood() {
114    let (url, gets) = spawn_fake_server(GetBehavior::AbnormalBodyDeath).await;
115    let transport = StreamableHttpClientTransport::from_uri(url.as_str());
116    let client = ().serve(transport).await.expect("handshake against fake server should succeed");
117
118    const OBSERVE: Duration = Duration::from_secs(3);
119    tokio::time::sleep(OBSERVE).await;
120    let n = gets.load(Ordering::Relaxed);
121    let _ = client.cancel().await;
122
123    eprintln!(
124        "[repro] {n} standing-GET reconnect attempts (= WARN log lines) in {OBSERVE:?} \
125         = {:.0}/sec",
126        n as f64 / OBSERVE.as_secs_f64()
127    );
128    // A backoff-respecting client would attempt ~3-5 in 3s; the bug produces
129    // hundreds-to-thousands.
130    assert!(
131        n > 20,
132        "expected a zero-backoff reconnect flood (>20 GETs in {OBSERVE:?}), got {n}; \
133         if this FAILS with a small n, the rmcp loop got fixed - delete mcp_http_client.rs"
134    );
135}
136
137/// The fix: same body-killing server, client wrapped in `McpHttpClient` —
138/// the backoff schedule allows only a handful of reconnects.
139#[tokio::test(flavor = "multi_thread")]
140async fn throttled_client_bounds_the_flood() {
141    let (url, gets) = spawn_fake_server(GetBehavior::AbnormalBodyDeath).await;
142    let throttled = McpHttpClient::new(
143        reqwest::Client::default(),
144        "fake-server",
145        WarnBudget::default(),
146    );
147    let transport = StreamableHttpClientTransport::with_client(
148        throttled,
149        StreamableHttpClientTransportConfig::with_uri(url.as_str()),
150    );
151    let client = ().serve(transport).await.expect("handshake against fake server should succeed");
152
153    // Checkpoint: backoff must engage early (instant, instant, 0.5s => 3-4
154    // GETs by 1.2s; an unthrottled client would be in the hundreds).
155    tokio::time::sleep(Duration::from_millis(1200)).await;
156    let early = gets.load(Ordering::Relaxed);
157    assert!(
158        early <= 4,
159        "backoff must engage early; got {early} GETs by 1.2s"
160    );
161
162    const OBSERVE: Duration = Duration::from_secs(4);
163    tokio::time::sleep(OBSERVE - Duration::from_millis(1200)).await;
164    let n = gets.load(Ordering::Relaxed);
165    let _ = client.cancel().await;
166
167    eprintln!("[repro] {n} throttled reconnect attempts in {OBSERVE:?} (schedule allows ~5-6)");
168    assert!(
169        (3..=8).contains(&n),
170        "throttled reconnects should follow the backoff schedule (~5-6 in {OBSERVE:?}), got {n}"