Grok Build 终端原生 Coding Agent 运行时

1456fn main() {
1457    xai_grok_pager_minimal::install();
1458    #[cfg(all(feature = "jemalloc", unix))]
1459    xai_grok_pager::memory_release::install_release_hook(purge_jemalloc_retained_pages);
1460    #[cfg(all(feature = "jemalloc", unix))]
1461    {
1462        xai_grok_pager::memory_trace::install_allocator_stats_provider(jemalloc_allocator_stats);
1463        xai_grok_pager::memory_trace::install_allocator_dump_provider(jemalloc_stats_dump);
1464    }
1465    #[cfg(all(feature = "jemalloc", unix))]
1466    install_heap_profile_hooks();
1467    if let Some(code) = xai_grok_pager::app::mermaid_worker::maybe_run_render_subprocess() {
1468        std::process::exit(code);
1469    }
1470    xai_grok_pager::memory_trace::start(
1471        xai_grok_shell::util::grok_home::grok_home().join("memtrace"),
1472    );
1473    raise_fd_limit();
1474    if let Err(e) = xai_grok_config::validate_requirements() {
1475        eprintln!("Couldn't start Grok: {e}");
1476        eprintln!();
1477        eprintln!(
1478            "Update Grok to a version the policy allows, or ask your administrator \
1479             to fix the managed requirements."
1480        );
1481        std::process::exit(2);
1482    }
1483    let _sentry_guard = xai_grok_telemetry::sentry::init(xai_grok_telemetry::sentry::Config {
1484        client: "grok-pager",
1485        client_version: PAGER_CLIENT_VERSION,
1486        release: env!("VERSION_WITH_COMMIT"),
1487        disabled: xai_grok_shell::agent::config::is_error_reporting_disabled_sync(),
1488    });

1777    let headless_prompt = xai_grok_pager::headless::HeadlessPrompt::from_args(
1778        args.single.as_deref(),
1779        args.prompt_json.as_deref(),
1780        args.prompt_file.as_deref(),
1781    )?;
1782    if let Some(prompt) = headless_prompt {
1783        init_tracing_simple(HEADLESS_ENTRYPOINT);
1784        let _otel_guard = xai_grok_telemetry::otel_layer::otel_guard();
1785        enforce_minimum_version_or_exit(&update_config).await;
1786        let launch_yolo = xai_grok_shell::util::config::effective_yolo_for_launch(
1787            args.yolo,
1788            args.permission_mode_flag.as_deref(),
1789            None,
1790        );
1791        if let Some(warning) = launch_yolo.blocked_warning {
1792            eprintln!("grok: {warning}");
1793        }
1794        let json_schema = args
1795            .json_schema
1796            .as_deref()
1797            .map(xai_grok_pager::headless::parse_json_schema)
1798            .transpose()?;
1799        if json_schema.is_some() {
1800            if args.output_format == xai_grok_pager::headless::OutputFormat::Plain {
1801                args.output_format = xai_grok_pager::headless::OutputFormat::Json;
1802            }
1803            if args.self_verify {
1804                anyhow::bail!(
1805                    "--json-schema and --self-verify cannot be used together: \
1806                     verification output would corrupt the structured response"
1807                );
1808            }
1809        }
1810        return xai_grok_pager::headless::run_single_turn(
1811            prompt,
1812            args.verbatim,
1813            xai_grok_pager::headless::HeadlessOptions {
1814                session_id: args.session_id.clone(),
1815                resume: args.resume_session.or(args.load_session),
1816                cwd: args.cwd,
1817                yolo: launch_yolo.yolo,
1818                trust: args.trust,
1819                output_format: args.output_format,
1820                json_schema,
1821                model: args.model,
1822                rules: args.rules,
1823                system_prompt_override: args.system_prompt_override.clone(),
1824                continue_last_session: args.continue_last_session,
1825                fork_session: args.fork_session,

97fn plan_reconnect_load(
98    agent: &super::agent_view::AgentView,
99    fallback_cwd: &std::path::Path,
100) -> Option<ReconnectLoadPlan> {
101    let session_id = agent.session.session_id.clone()?;
102    let cwd = if agent.session.cwd.as_os_str().is_empty() {
103        fallback_cwd.to_path_buf()
104    } else {
105        agent.session.cwd.clone()
106    };
107    let yolo = agent.session.is_yolo();
108    // Set BOTH yoloMode and autoMode explicitly. The leader's capability injection
109    // only fills ABSENT keys, so omitting autoMode here lets a stale launch-time
110    // `ClientCapabilities.auto_mode` re-enable Auto after the user left it (e.g.
111    // Shift+Tab to Ask). Auto is per-agent (symmetric with yolo) — derive it from
112    // this agent's own `auto_mode` so a background tab reconnects with ITS mode,
113    // not the active tab's global `current_ui` mirror.
114    let auto = super::dispatch::effective_auto(yolo, agent.session.is_auto());
115    let mut meta = serde_json::json!({ "yoloMode": yolo, "autoMode": auto });
116    if let Some(ref cursor) = agent.last_seen_event_id {
117        meta["cursor"] = serde_json::Value::String(cursor.clone());
118    }
119    Some(ReconnectLoadPlan {
120        session_id,
121        cwd,
122        meta,
123    })
124}
125
126/// Resolve the two post-reconnect restore outcomes from the per-agent
127/// `session/load` results.
128///
129/// - `all_restored` (AND across every reloaded tab, plus `init_ok`) drives the
130///   user-facing toast: it reports whether the WHOLE reconnect came back.
131/// - `active_restored` is per-agent: the ACTIVE tab's OWN reload succeeded. It
132///   gates that tab's post-reconnect queue drain. Gating the drain on
133///   `all_restored` would let one failed background tab strand prompts queued
134///   on a healthy active tab — the drain (`dispatch_drain_queue`) only ever
135///   touches the active agent, so a background failure has no bearing on it.
136///
137/// `loads` maps each reloaded agent to `(success, running_prompt_id)`; an agent
138/// in `pending_agent_ids` but absent from `loads` is treated as failed
139/// (mirrors the `unwrap_or((false, _))` at the finalize site).
140fn reconnect_restore_outcome(
141    init_ok: bool,
142    pending_agent_ids: &[super::agent::AgentId],
143    loads: &std::collections::HashMap<super::agent::AgentId, (bool, Option<String>)>,
144    active_agent_id: Option<super::agent::AgentId>,
145) -> (bool, bool) {
146    let load_ok = |id: &super::agent::AgentId| -> bool { loads.get(id).is_some_and(|(ok, _)| *ok) };
147    let all_restored = init_ok && pending_agent_ids.iter().all(load_ok);
148    let active_restored = init_ok
149        && active_agent_id.is_some_and(|aid| pending_agent_ids.contains(&aid) && load_ok(&aid));
150    (all_restored, active_restored)

17/// Unique local identifier for an agent.
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
19pub struct AgentId(pub usize);
20/// Whether a queue entry is a regular prompt or a slash command.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum QueueEntryKind {
23    /// Regular user prompt — sent via `PromptRequest`.
24    Prompt,
25    /// Slash command (e.g., `/compact`) — dispatched as `ExtRequest` or local action.
26    Command,
27    /// Direct bash command — bypasses agent loop, executed by shell directly.
28    BashCommand,
29    /// Scheduled (cron) prompt -- injected by the scheduler via ACP notification.
30    Cron,
31}
32impl QueueEntryKind {
33    /// Short, stable label for telemetry / profiling logs.
34    pub fn as_label(&self) -> &'static str {
35        match self {
36            Self::Prompt => "prompt",
37            Self::Command => "command",
38            Self::BashCommand => "bash_command",
39            Self::Cron => "cron",
40        }
41    }
42}
43/// An entry waiting in the queue to be sent to the agent.
44///
45/// Each entry gets a monotonically increasing `id` for stable tracking
46/// (e.g., when editing a queued prompt whose positional index shifts as
47/// earlier prompts drain). The user-facing display uses the 1-based
48/// positional index (`#1`, `#2`, …), never the internal `id`.
49#[derive(Debug, Clone)]
50pub struct QueuedPrompt {
51    /// Monotonic ID, unique within this agent's session. Never reused.
52    pub id: u64,
53    /// The prompt text (or command text, e.g. "/compact").
54    pub text: String,
55    /// Whether this is a prompt or a slash command.
56    pub kind: QueueEntryKind,
57    /// Optional separate payload for the wire. When `Some`, this is sent
58    /// instead of `text`. Used for skill injection where the display
59    /// shows `/commit args` but the wire carries the skill XML content.
60    pub wire_blocks: Option<Vec<acp::ContentBlock>>,
61    /// Images attached to this prompt. Drained from `PromptWidget` at
62    /// submission time. Preserved across queue text edits.
63    pub images: Vec<crate::prompt_images::PastedImage>,
64    /// Whether this prompt should display as a skill invocation (teal accent).
65    /// Only meaningful when `wire_blocks` is `Some`.
66    pub display_as_skill: bool,
67    /// Recognized slash-token byte ranges into `text`, captured from the
68    /// composer at submit time; empty = no token styling.
69    pub skill_token_ranges: Vec<std::ops::Range<usize>>,
70    /// Scheduler task ID for cron prompts. Used for per-task dedup.
71    pub task_id: Option<String>,
72    /// Human-readable schedule (e.g. "every 5 minutes") for cron prompts.
73    /// Threaded into the system-reminder framing sent to the model.
74    pub human_schedule: Option<String>,
75    /// All chip elements captured from the textarea at send time.
76    /// Threaded into `InFlightPrompt` so rewind restores collapsed chips.
77    pub chip_elements: Vec<ChipElement>,
78}

161/// Maximum in-memory stdout per background task (10 MB).
162pub const BG_TASK_MAX_STDOUT: usize = 10 * 1024 * 1024;
163/// How long to wait for a kill response before auto-clearing `pending_kill`
164/// so the user can retry. Applied to both bg tasks and subagents.
165pub const PENDING_KILL_TIMEOUT_SECS: u64 = 10;
166/// Status of a background task.
167#[derive(Debug, Clone, Copy, PartialEq, Eq)]
168pub enum BgTaskStatus {
169    /// Currently running.
170    Running,
171    /// Completed successfully (exit 0).
172    Done,
173    /// Failed (non-zero exit, signal, timeout, OOM, etc.).
174    Failed,
175}
176/// Central state for a single background task.
177///
178/// Stored in `AgentSession::bg_tasks` keyed by `task_id`.
179/// Both the scrollback `BgTaskBlock` and the bg task pane read from this.
180#[derive(Debug, Clone)]
181pub struct BgTaskState {
182    pub task_id: String,
183    pub tool_call_id: String,
184    pub command: String,
185    pub description: Option<String>,
186    pub cwd: String,
187    pub output_file: String,
188    pub status: BgTaskStatus,
189    pub start_time: SystemTime,
190    pub end_time: Option<SystemTime>,
191    pub exit_code: Option<i32>,
192    pub signal: Option<String>,
193    /// Accumulated stdout (full cumulative buffer from shell, max BG_TASK_MAX_STDOUT).
194    ///
195    /// Mutate via [`BgTaskState::set_stdout`] / [`BgTaskState::append_stdout`]
196    /// so `stdout_line_count` and `truncated` stay in sync.
197    pub stdout: String,
198    /// Cached `stdout.lines().count()`. Kept in sync by [`Self::set_stdout`]
199    /// and [`Self::append_stdout`] so the tasks-pane overlay doesn't have to
200    /// memchr-scan up to `BG_TASK_MAX_STDOUT` bytes per visible task per
201    /// render frame.
202    pub stdout_line_count: usize,
203    /// Whether the rolling buffer has dropped data. Either the shell-side
204    /// `BashOutput.truncated` flag was set when the chunk arrived, or the
205    /// TUI itself trimmed the buffer to stay under `BG_TASK_MAX_STDOUT` (in
206    /// `set_stdout` / `append_stdout`). Once `true`, stays `true` — the
207    /// real line count is at least `stdout_line_count`, hence the `(N+)`
208    /// badge.
209    pub truncated: bool,
210    /// Kill request sent, awaiting task_completed.
211    pub pending_kill: bool,
212    /// When the kill request was sent. Used to auto-clear `pending_kill`
213    /// after a timeout so the user can retry if the response is lost.
214    pub kill_requested_at: Option<Instant>,
215    /// Scrollback entry ID for the "Task started" block (for finish_running).
216    pub scrollback_entry_id: Option<crate::scrollback::entry::EntryId>,
217    /// True when this background task is a monitor (the `monitor` tool). The
218    /// tasks pane renders monitors with a blue "Monitor" tag + neutral text
219    /// (mirroring scheduled `/loop` rows) instead of bash-highlighting the
220    /// command. Set from the `monitor_description` field of the
221    /// `TaskBackgrounded` notification.
222    pub is_monitor: bool,

195/// Suspend the inline TUI, run a blocking child that takes over the tty
196/// (`$EDITOR`, `$PAGER`, …), then restore. Shared by the editor and transcript
197/// suspend paths so the subtle reader-park + raw-mode + alt-screen handoff lives
198/// in one place.
199///
200/// The reader thread is parked first so the child (which inherits this tty)
201/// keeps every keystroke instead of racing the reader; on return, buffered
202/// terminal query replies and any pre-park keystroke are drained before the
203/// reader resumes.
204///
205/// The frame **writer** thread is then drained (bounded) before the child
206/// starts: frames are written to the tty asynchronously, and the frame that
207/// armed this suspend (e.g. minimal's final `/transcript` pump slice) is
208/// typically queued microseconds before we get here. Un-drained, those bytes
209/// race the child's own output — they can land on the child's alternate
210/// screen (so the main screen never receives them; on an inline viewport a
211/// commit's scroll then leaves every following row off by its height) or tear
212/// around the alt-screen switch, printing escape fragments (`[`…) that the
213/// renderer's diff can't see and thus never repairs.
214///
215/// In minimal mode the physical cursor is probed (`ESC[6n`) right before the
216/// child runs and right after it exits — while the reader thread is still
217/// parked, so the replies can't be stolen. The returned "cursor after the
218/// child, iff it moved" tells the caller whether the child *printed to the
219/// main screen* (cat-style pager: cursor left below its output) or *restored
220/// it* (less-style alt-screen pager: `rmcup` puts the cursor back exactly
221/// where it was): the caller re-anchors the live region below the new output
222/// in the first case and repaints in place in the second.
223fn suspend_for_child(
224    screen_mode: crate::app::ScreenMode,
225    writer_sync: &crate::render::draw::WriterSync,
226    input_paused: &std::sync::atomic::AtomicBool,
227    reader_parked: &std::sync::atomic::AtomicBool,
228    input_rx: &mut tokio::sync::mpsc::UnboundedReceiver<crossterm::event::Event>,
229    run_child: impl FnOnce(),
230) -> Option<(u16, u16)> {
231    use std::sync::atomic::Ordering;
232    // Pause the reader thread, then wait for a FRESH park (reader provably out of
233    // crossterm) so the main thread is the sole poll/read caller; bounded so a
234    // dead reader can't hang us.
235    input_paused.store(true, Ordering::Release);
236    reader_parked.store(false, Ordering::Release);
237    let park_deadline = std::time::Instant::now() + Duration::from_millis(500);
238    while !reader_parked.load(Ordering::Acquire) && std::time::Instant::now() < park_deadline {
239        std::thread::sleep(Duration::from_millis(5));
240    }
241    // Every queued frame must be ON the tty before the child takes it (and, in
242    // fullscreen, before LeaveAlternateScreen below — mirroring teardown's
243    // "no late frame after LeaveAlternateScreen" drain). Bounded like the
244    // reader park so a wedged pty can't hang the suspend.
245    if !writer_sync.wait_drained(Duration::from_millis(750)) {
246        tracing::warn!("suspend: frame writer not drained within 750ms; proceeding");
247    }
248    // Pre-child cursor probe (minimal only — minimal's startup already proved
249    // this terminal answers CPR). Reader is parked, so the reply is ours.
250    let pre_cursor = screen_mode
251        .is_minimal()
252        .then(|| crossterm::cursor::position().ok())
253        .flatten();
254    if screen_mode.is_fullscreen() {
255        xai_grok_shell::util::with_locked_stderr(|stderr| {
256            let _ = crossterm::execute!(stderr, crossterm::terminal::LeaveAlternateScreen);
257        });
258    }
259    let _ = crossterm::terminal::disable_raw_mode();
260    run_child();

6/// Capability mode applied to a session's toolset.
7///
8/// A partial order is defined via [`CapabilityMode::is_subset_of`]:
9/// `ReadOnly < ReadWrite < All` and `ReadOnly < Execute < All`.
10/// `ReadWrite` and `Execute` are *incomparable* (neither is a subset
11/// of the other). `fork_session` enforces `child <= parent`.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
13#[serde(rename_all = "snake_case")]
14pub enum CapabilityMode {
15    /// Reading and searching only. No edits, no shell, no background tasks.
16    ReadOnly,
17    /// Read + edit. No shell execution.
18    ReadWrite,
19    /// Read + shell execution + background-task control. No edits.
20    Execute,
21    /// Every tool kind allowed.
22    All,
23}
24
25impl Default for CapabilityMode {
26    /// Defaults to [`CapabilityMode::ReadWrite`] (subagent default;
27    /// the main/root session is always `All`).
28    fn default() -> Self {
29        Self::ReadWrite
30    }
31}
32
33impl CapabilityMode {
34    /// Filter `config.tools` by capability mode, returning a copy with
35    /// disallowed tools dropped.
36    ///
37    /// Tools whose `kind` is `None` (baseline, e.g. ad-hoc tools
38    /// declared via `ToolConfig::simple`) are preserved across all
39    /// modes. **MCP-origin** `kind: None` tools are NOT preserved by
40    /// this method; see `resolve_session_toolset` for the asymmetric
41    /// handling.
42    pub fn filter(self, config: &ToolServerConfig) -> ToolServerConfig {
43        let kept: Vec<ToolConfig> = config
44            .tools
45            .iter()
46            .filter(|tool| match tool.kind {
47                None => true,
48                Some(kind) => kind_allowed(self, kind),
49            })
50            .cloned()
51            .collect();
52        ToolServerConfig {
53            tools: kept,
54            behavior_preset: config.behavior_preset.clone(),
55        }
56    }
57
58    /// Whether every kind allowed by `self` is also allowed by `other`.
59    /// Used by `fork_session` to reject capability widening.
60    pub fn is_subset_of(self, other: CapabilityMode) -> bool {
61        for kind in ALL_TOOL_KINDS {
62            if kind_allowed(self, *kind) && !kind_allowed(other, *kind) {
63                return false;
64            }
65        }
66        true
67    }
68}

1//! New-architecture tool implementations (NewTool trait).
2//!
3//! Each sub-module here contains a tool that implements `NewTool` instead
4//! of the old `Tool` trait. During migration, old implementations live in
5//! `implementations/<tool>/` and new implementations live in
6//! `implementations/grok_build/<tool>/`.
7//!
8//! The [`register_all()`] function is the single entry-point for wiring up
9//! the standard toolset. It inserts shared resources (`Terminal`,
10//! `AvailableSkills`, `BashParams`) and registers every built-in tool.
11pub mod ask_user_question;
12pub mod bash;
13#[path = "deploy_app_stub.rs"]
14pub mod deploy_app;
15pub mod enter_plan_mode;
16pub mod exit_plan_mode;
17pub mod grep;
18pub mod image_edit;
19pub mod image_gen;
20pub mod kill_task;
21pub mod list_dir;
22pub mod lsp;
23pub mod monitor;
24pub mod read_file;
25pub mod scheduler;
26pub mod search_replace;
27pub(crate) mod storage;
28pub mod task;
29pub mod task_output;
30pub mod todo;
31pub mod update_goal;
32pub mod video_gen;
33pub mod web_fetch;
34pub mod web_search;
35pub use ask_user_question::AskUserQuestionTool;
36pub use bash::BashTool;
37pub use deploy_app::{AppBuilderDeployerConfig, DEPLOY_APP_TOOL_NAME};
38pub use enter_plan_mode::EnterPlanModeTool;
39pub use exit_plan_mode::ExitPlanModeTool;
40pub use grep::GrepTool;
41pub use image_edit::{IMAGE_EDIT_TOOL_NAME, ImageEditTool};
42pub use image_gen::{
43    IMAGE_GEN_TOOL_NAME, IMAGINE_COMMAND_NAME, ImageGenTool, imagine_instruction,
44    imagine_usage_message,
45};
46pub use kill_task::{KillTaskTool, KillTerminalCommandTool};
47pub use list_dir::ListDirTool;
48pub use lsp::LspTool;
49pub use monitor::tool::MonitorTool;
50pub use read_file::ReadFileTool;
51pub use scheduler::create::{
52    SCHEDULER_CREATE_TOOL_NAME, SchedulerCreateTool, loop_schedule_instruction, loop_usage_message,
53};
54pub use scheduler::delete::{SCHEDULER_DELETE_TOOL_NAME, SchedulerDeleteTool};
55pub use scheduler::list::SchedulerListTool;
56pub use search_replace::SearchReplaceTool;
57pub use task::TaskTool;
58pub use task_output::{GetTerminalCommandOutputTool, TaskOutputTool, WaitTasksTool};
59pub use todo::TodoWriteTool;
60pub use update_goal::{UPDATE_GOAL_TOOL_NAME, UpdateGoalTool};
61pub use video_gen::{
62    IMAGE_TO_VIDEO_TOOL_NAME, IMAGINE_VIDEO_COMMAND_NAME, ImageToVideoTool,
63    REFERENCE_TO_VIDEO_TOOL_NAME, ReferenceToVideoTool, imagine_video_instruction,
64    imagine_video_usage_message,
65};
66pub use web_fetch::{WebFetchClient, WebFetchConfig, WebFetchParams, WebFetchTool};
67pub use web_search::WebSearchTool;

1375    /// Execute a tool, returning only its raw output.
1376    ///
1377    /// Unlike [`call()`], this skips reminders and persistence. Used by
1378    /// `InnerDispatchForToolset` so that `use_tool`'s
1379    /// dispatch to a target tool does not double-run post-processing: the
1380    /// outer `call("use_tool")` does one round of post-processing over the
1381    /// target's output.
1382    ///
1383    /// Inner dispatch is intentionally **not** populated in the forwarded
1384    /// context. MCP tools (the targets of `use_tool`) are passthrough
1385    /// implementations that never call `use_tool` themselves, so they do
1386    /// not need inner dispatch.
1387    ///
1388    /// The `parent_ctx` carries call-id, cwd, resources, etc. from the
1389    /// outer call. A fresh child context is built from it — stripping
1390    /// `InnerDispatch` to prevent recursion.
1391    async fn call_raw(
1392        &self,
1393        tool_name: &str,
1394        tool_args: serde_json::Value,
1395        parent_ctx: xai_tool_runtime::ToolCallContext,
1396    ) -> Result<crate::types::output::ToolOutput, xai_tool_runtime::ToolError> {
1397        let (registry_id, output_converter, reverse_params) = {
1398            let tools = self.tools.read();
1399            let entry = tools
1400                .iter()
1401                .find(|t| t.client_name == tool_name)
1402                .ok_or_else(|| Self::tool_not_found_error(tool_name))?;
1403            (
1404                entry.registry_id.clone(),
1405                entry.output_converter.clone(),
1406                entry.reverse_params.clone(),
1407            )
1408        };
1409        let canonical_params = if reverse_params.is_empty() {
1410            tool_args
1411        } else {
1412            remap_json_keys(tool_args, &reverse_params)
1413        };
1414        let mut ctx = xai_tool_runtime::ToolCallContext::new(parent_ctx.call_id.clone());
1415        ctx.extensions.insert(self.resources.clone());
1416        ctx.extensions.insert_arc(Arc::clone(&self.renderer));
1417        if let Some(cwd) = parent_ctx.extensions.get::<xai_tool_runtime::Cwd>() {
1418            ctx.extensions.insert((*cwd).clone());
1419        }
1420        let tool_id = xai_tool_protocol::ToolId::new(&registry_id)
1421            .unwrap_or_else(|_| xai_tool_protocol::ToolId::new("unknown").expect("valid"));
1422        let lr_handle = self.local_registry.find(&tool_id).ok_or_else(|| {
1423            xai_tool_runtime::ToolError::not_found(
1424                tool_id,
1425                format!("Tool not found in LocalRegistry: {registry_id}"),
1426            )
1427        })?;
1428        let stream = lr_handle.execute(ctx, canonical_params).await;
1429        let value = drain_value_stream(stream).await?;
1430        (output_converter)(value)
1431            .map_err(|e| xai_tool_runtime::ToolError::custom("output_decoding", e.to_string()))
1432    }

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();

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    }

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";

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}"

182    #[test]
183    fn capability_mode_filter_table_is_exhaustive_per_kind() {
184        for &mode in &[
185            CapabilityMode::ReadOnly,
186            CapabilityMode::ReadWrite,
187            CapabilityMode::Execute,
188            CapabilityMode::All,
189        ] {
190            for &kind in ALL_TOOL_KINDS {
191                let id = format!("kind_{kind:?}");
192                let cfg = make_cfg(vec![test_support::tc(&id, Some(kind))]);
193                let out = mode.filter(&cfg);
194                let expected_present = kind_allowed(mode, kind);
195                let actually_present = out.tools.iter().any(|t| t.id == id);
196                assert_eq!(
197                    actually_present, expected_present,
198                    "({mode:?}, {kind:?}): expected present={expected_present}, got {actually_present}"
199                );
200            }
201        }
202    }
203
204    #[test]
205    fn capability_mode_filter_anchored_membership() {
206        let cfg = make_cfg(vec![
207            test_support::tc("read", Some(ToolKind::Read)),
208            test_support::tc("search", Some(ToolKind::Search)),
209            test_support::tc("inspect", Some(ToolKind::Lsp)),
210            test_support::tc("edit", Some(ToolKind::Edit)),
211            test_support::tc("write", Some(ToolKind::Write)),
212            test_support::tc("bash", Some(ToolKind::Execute)),
213            test_support::tc("bg", Some(ToolKind::BackgroundTaskAction)),
214            test_support::tc("plan", Some(ToolKind::Plan)),
215            test_support::tc("ask", Some(ToolKind::AskUser)),
216            test_support::tc("other", Some(ToolKind::Other)),
217        ]);
218
219        let names = |c: &ToolServerConfig| -> Vec<String> {
220            c.tools.iter().map(|t| t.id.clone()).collect()
221        };
222
223        let ro = CapabilityMode::ReadOnly.filter(&cfg);
224        assert_eq!(names(&ro), vec!["read", "search", "inspect", "plan", "ask"]);
225
226        let rw = CapabilityMode::ReadWrite.filter(&cfg);
227        assert_eq!(
228            names(&rw),
229            vec!["read", "search", "inspect", "edit", "write", "plan", "ask"]
230        );
231
232        let ex = CapabilityMode::Execute.filter(&cfg);
233        assert_eq!(
234            names(&ex),
235            vec!["read", "search", "inspect", "bash", "bg", "plan", "ask"]
236        );
237
238        let all = CapabilityMode::All.filter(&cfg);
239        assert_eq!(
240            names(&all),
241            vec![
242                "read", "search", "inspect", "edit", "write", "bash", "bg", "plan", "ask", "other"
243            ]
244        );
245    }