bingshuo/grok-build-upstream-mirror
Watch
1
0
Fork
You've already forked grok-build-upstream-mirror
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Synced from monorepo

Browse source 
Changes:
- Classify clipboard delivery confidence
- Add durable session update append
- Scope the xAI session bearer to first-party memory embedding endpoints
- Persist subagent outputs to disk and bound long-lived agent state
- Add MiniSweAgent:bash for mini-swe-agent parity
- Revert taking local sessions off the persistent shell
- Contextual tip recommending grok wrap on SSH sessions
- Voice STT bearer from model BYOK env_key/api_key
- Define exact website policies for sandbox
- Gate unsafe shell environments
- Shared pin hoist; single require_sha gate for marketplace plugins
- Server-signed is-managed claim (closes sidecar-removal downgrade)
- Optional require_sha pin for remote plugin installs
- Show session title and last exchange in the exit resume hint
- Gate shell output redirects
- Warn when fail_closed is present but not a boolean
- Add canonical text editing core (ratatui-textarea)
- Keep execution state out of goal scratch
- Add acknowledged persistence primitives
- Inherit child network restrictions in sandbox
- Fail closed when hook matchers fail to recompile
- Add MCP setup preferences for plugin MCPs
- Gate sourced shell scripts
- Gate file-typed project hooks
- grok wrap: restore terminal modes on child death
- Harden owner-only permissions on auth and MCP credentials
- Create crash dump files with owner-only permissions
- Write the agent_id cache owner-only (0600)
- SessionMetrics mode skips Mixpanel profile sync
- Dashboard: slim live-tail peek
- Yank full queued prompt text, not (+N lines)
- Defeat clock-rollback on the signed managed-config cache
- Stop early session/cancel from overtaking the prompt and wedging the turn slot
- Self-heal a diverged agent entrypoint on startup
- Add matched inference expectations in test-support
- Add AuthSingleFlight cancel/successor gap tests
- Remove consumer from external OTEL allowlist and pin scrub coverage
- Enable /copy in minimal mode
- Surface capacity and API-key detail on 429 errors
- Single-flight interactive auth
- Fix PageUp/PageDown skipping lines behind sticky prompt header
This commit is contained in:
grokkybara[bot] 2026-07-17 14:19:50 +01:00
commit 98c3b2438a
225 changed files with 18836 additions and 7156 deletions
Download patch file Download diff file Expand all files Collapse all files

3382
crates/codegen/xai-grok-shell/CHANGELOG.md
View file

File diff suppressed because it is too large Load diff

2
crates/codegen/xai-grok-shell/Cargo.toml
View file

@ -1,7 +1,7 @@
[package]
license = "Apache-2.0"
name = "xai-grok-shell"
version = "0.2.101"
version = "0.2.102"
edition.workspace = true
[features]

187
crates/codegen/xai-grok-shell/changelogs/0.2.102.json Normal file
View file

@ -0,0 +1,187 @@
[
{
"category": "breaking",
"description": "**--minimal** and **--fullscreen** flags now apply only to the current session.",
"breaking_change": true
},
{
"category": "features",
"description": "**New /jump slash command** lets you quickly jump to any previous turn in the conversation.",
"breaking_change": false
},
{
"category": "features",
"description": "**New /timeline sidebar** shows a clickable tick rail for fast navigation between conversation turns.",
"breaking_change": false
},
{
"category": "features",
"description": "**grok login** now requests Grok Projects scopes so workspace listing works after consent.",
"breaking_change": false
},
{
"category": "features",
"description": "**Permission mode** can now be set fleet-wide via remote config when no local setting exists.",
"breaking_change": false
},
{
"category": "features",
"description": "**Edit tool output** has a setting to show a compact one-line summary instead of always-expanded diffs.",
"breaking_change": false
},
{
"category": "features",
"description": "**Tab completion** in !bash mode now works like a normal terminal (prefix fill, dropdown, directory drill-down).",
"breaking_change": false
},
{
"category": "features",
"description": "**Enterprise deployments** can now disable voice dictation via `requirements.toml` so `/voice` and Ctrl+Space are hidden for everyone.",
"breaking_change": false
},
{
"category": "features",
"description": "**User prompts** now appear bold only in `--minimal` mode; fullscreen keeps normal weight.",
"breaking_change": false
},
{
"category": "features",
"description": "**`grok plugin install`** now accepts a marketplace's registered name as a qualifier.",
"breaking_change": false
},
{
"category": "features",
"description": "Consecutive edits to the same file now collapse into a single scrollback row when collapsed edit blocks are enabled.",
"breaking_change": false
},
{
"category": "features",
"description": "Local sessions now inherit your shell environment variables and keep the current directory across commands.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Login and re-login** no longer stack multiple device-code polls or leave stale flows running.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Background task tools** now render with correct icons and titles instead of the generic MCP wrench.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Task tool** now correctly validates and displays allowed model slugs for subagents.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Rewind** now correctly handles bash transcripts, permission follow-ups, and sessions that mix old and new prompt markers.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Re-login** during a session now immediately uses the new token instead of requiring a new session.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Terminal commands** using globs now behave the same on zsh as on bash and no longer fail with shell errors.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Installer** no longer replaces stowed shell configuration symlinks with plain files on upgrade.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Voice transcription** now works with enterprise API bases and API-key authentication.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Fixed crashes** on some network-mounted home directories by using a safer SQLite journal mode.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Home and End keys** now move to the ends of the current wrapped line in the prompt.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Arrow keys and Esc** now work correctly inside viewers opened from the dashboard.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Warns at startup** when user and project sandbox profiles define the same name differently.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Billing upgrade links** now show the full URL in the transcript (and copy it) when a browser cannot be opened.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Fixed Ctrl+Y yank** no longer working after sending a prompt.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**No longer shows permission prompts** seconds after a turn was cancelled with Esc or Ctrl+C.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Page Up and Page Down** now move the highlighted entry to the top or bottom of the visible scrollback area.",
"breaking_change": false
},
{
"category": "fixes",
"description": "Conflicting project and user sandbox profiles now show a clear warning on the welcome screen.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**OAuth login URLs** no longer contain duplicate referrer parameters.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**File links** in official VS Code Remote-SSH terminals now use VS Code's native path handling.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Minimal mode** now shows the folder-trust prompt after sign-in when required.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Skills** whose names collide with built-in slash commands are now reachable via qualified names.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Fixed background task tracking** when using grok -p --no-wait-for-background so tasks are properly reaped on exit.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Rate limit errors (429)** now show specific server messages (capacity, team limits, free-usage) instead of generic upgrade prompts, with correct copy based on auth type.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**`/copy` slash command** is now available in minimal mode.",
"breaking_change": false
},
{
"category": "performance",
"description": "**Improved recap and compaction** behavior.",
"breaking_change": false
}
]

50
crates/codegen/xai-grok-shell/changelogs/0.2.102.md Normal file
View file

@ -0,0 +1,50 @@
# 0.2.102 — 2026-07-16
## Breaking Changes
- **--minimal** and **--fullscreen** flags now apply only to the current session.
## Features
- **New /jump slash command** lets you quickly jump to any previous turn in the conversation.
- **New /timeline sidebar** shows a clickable tick rail for fast navigation between conversation turns.
- **grok login** now requests Grok Projects scopes so workspace listing works after consent.
- **Permission mode** can now be set fleet-wide via remote config when no local setting exists.
- **Edit tool output** has a setting to show a compact one-line summary instead of always-expanded diffs.
- **Tab completion** in !bash mode now works like a normal terminal (prefix fill, dropdown, directory drill-down).
- **Enterprise deployments** can now disable voice dictation via `requirements.toml` so `/voice` and Ctrl+Space are hidden for everyone.
- **User prompts** now appear bold only in `--minimal` mode; fullscreen keeps normal weight.
- **`grok plugin install`** now accepts a marketplace's registered name as a qualifier.
- Consecutive edits to the same file now collapse into a single scrollback row when collapsed edit blocks are enabled.
- Local sessions now inherit your shell environment variables and keep the current directory across commands.
## Bug Fixes
- **Login and re-login** no longer stack multiple device-code polls or leave stale flows running.
- **Background task tools** now render with correct icons and titles instead of the generic MCP wrench.
- **Task tool** now correctly validates and displays allowed model slugs for subagents.
- **Rewind** now correctly handles bash transcripts, permission follow-ups, and sessions that mix old and new prompt markers.
- **Re-login** during a session now immediately uses the new token instead of requiring a new session.
- **Terminal commands** using globs now behave the same on zsh as on bash and no longer fail with shell errors.
- **Installer** no longer replaces stowed shell configuration symlinks with plain files on upgrade.
- **Voice transcription** now works with enterprise API bases and API-key authentication.
- **Fixed crashes** on some network-mounted home directories by using a safer SQLite journal mode.
- **Home and End keys** now move to the ends of the current wrapped line in the prompt.
- **Arrow keys and Esc** now work correctly inside viewers opened from the dashboard.
- **Warns at startup** when user and project sandbox profiles define the same name differently.
- **Billing upgrade links** now show the full URL in the transcript (and copy it) when a browser cannot be opened.
- **Fixed Ctrl+Y yank** no longer working after sending a prompt.
- **No longer shows permission prompts** seconds after a turn was cancelled with Esc or Ctrl+C.
- **Page Up and Page Down** now move the highlighted entry to the top or bottom of the visible scrollback area.
- Conflicting project and user sandbox profiles now show a clear warning on the welcome screen.
- **OAuth login URLs** no longer contain duplicate referrer parameters.
- **File links** in official VS Code Remote-SSH terminals now use VS Code's native path handling.
- **Minimal mode** now shows the folder-trust prompt after sign-in when required.
- **Skills** whose names collide with built-in slash commands are now reachable via qualified names.
- **Fixed background task tracking** when using grok -p --no-wait-for-background so tasks are properly reaped on exit.
- **Rate limit errors (429)** now show specific server messages (capacity, team limits, free-usage) instead of generic upgrade prompts, with correct copy based on auth type.
- **`/copy` slash command** is now available in minimal mode.
## Performance
- **Improved recap and compaction** behavior.

10
crates/codegen/xai-grok-shell/src/agent/config.rs
View file

@ -3917,10 +3917,9 @@ impl ModelEntry {
api_base_url: entry.api_base_url.clone(),
}
}
/// The model's own (BYOK) credential: a non-empty `api_key`, else the first
/// set, non-empty `env_key` value. `None` means the model has no usable own
/// credential and resolution should fall through to the session / global key.
fn own_credential(&self) -> Option<String> {
/// Non-empty `api_key`, else first non-empty resolved `env_key`.
/// `None` → fall through to session / global key.
pub(crate) fn own_credential(&self) -> Option<String> {
first_own_credential(self.api_key.as_deref(), self.env_key.as_ref())
}
/// `true` when the model has a non-empty `api_key` or an `env_key` that
@ -4360,7 +4359,7 @@ pub fn enforce_disable_api_key_auth(
) {
if disable_api_key_auth
&& creds.auth_type == xai_chat_state::AuthType::ApiKey
&& crate::util::is_first_party_xai_url(&creds.base_url)
&& crate::util::is_xai_api_url(&creds.base_url)
{
creds.auth_type = xai_chat_state::AuthType::SessionToken;
creds.api_key = session_key.map(str::to_owned);
@ -9216,7 +9215,6 @@ agent_type = "cursor"
url = "https://mcp.test.com"
[toolset.bash]
timeout_secs = 120
persistent_shell = true
[shortcuts]
ctrl_k = "search"
[grok_com_config]

3
crates/codegen/xai-grok-shell/src/agent/handlers/model_switch.rs
View file

@ -223,11 +223,12 @@ pub(crate) async fn apply(
current_agent_type: None,
});
if agent.cfg.borrow().mode != config::AgentMode::Leader {
agent.models_manager.set_current_model_id(model_id);
agent.models_manager.set_current_model_id(model_id.clone());
agent
.models_manager
.set_current_reasoning_effort(applied_effort);
}
agent.sync_process_static_api_key(Some(model_id.0.as_ref()));
Ok(acp::SetSessionModelResponse::new().meta(
serde_json::json!({ "model" : updated_model, })
.as_object()

78
crates/codegen/xai-grok-shell/src/agent/mvp_agent/acp_agent.rs
View file

@ -367,6 +367,7 @@ impl acp::Agent for MvpAgent {
);
self.set_auth_method(default_id);
}
self.sync_process_static_api_key(None);
let current_working_directory = self.launch_cwd.clone();
let hostname = gethostname::gethostname();
let mcp_servers: Vec<crate::extensions::mcp::McpServerEntry> = Vec::new();
@ -513,6 +514,7 @@ impl acp::Agent for MvpAgent {
}
}
self.set_auth_method(arguments.method_id.clone());
self.sync_process_static_api_key(None);
self.ensure_telemetry_client();
if crate::agent::chat_modes::process_chat_mode_enabled() {
self.chat_modes.warm_in_background();
@ -718,44 +720,52 @@ impl acp::Agent for MvpAgent {
),
);
let login_override = auth_meta.login_override();
let (auth, _did_auth) = if !auth_meta.headless {
let mut cancelled = false;
let client_seq = auth_meta.request_seq;
let auth_result = if !auth_meta.headless {
let (url_tx, url_rx) = tokio::sync::oneshot::channel();
let (code_tx, code_rx) = tokio::sync::mpsc::channel(1);
*self.auth_code_tx.borrow_mut() = Some(code_tx);
*self.auth_url_rx.borrow_mut() = Some(url_rx);
let result = crate::auth::run_auth_flow_with_stderr_bridge(
&self.auth_manager,
grok_ctx,
crate::auth::AuthChannels {
url_tx: Some(url_tx),
code_rx,
},
auth_meta.reauth,
auth_meta.force_interactive,
login_override,
)
.await;
*self.auth_code_tx.borrow_mut() = None;
*self.auth_url_rx.borrow_mut() = None;
result
let (cancel, _guard) = self
.interactive_auth
.begin(
Some(
crate::auth::single_flight::AttemptChannels::new(
code_tx,
url_rx,
),
),
client_seq,
);
tokio::select! {
biased; _ = cancel.cancelled() => { cancelled = true;
Err(anyhow::anyhow!("Authentication cancelled")) } r = crate
::auth::run_auth_flow_with_stderr_bridge(& self.auth_manager,
grok_ctx, crate ::auth::AuthChannels { url_tx : Some(url_tx),
code_rx, }, auth_meta.reauth, auth_meta.force_interactive,
login_override,) => r,
}
} else {
crate::auth::run_auth_flow(
&self.auth_manager,
grok_ctx,
auth_meta.reauth,
None,
None,
None,
login_override,
)
.await
}
let (cancel, _guard) = self.interactive_auth.begin(None, client_seq);
tokio::select! {
biased; _ = cancel.cancelled() => { cancelled = true;
Err(anyhow::anyhow!("Authentication cancelled")) } r = crate
::auth::run_auth_flow(& self.auth_manager, grok_ctx, auth_meta
.reauth, None, None, None, login_override,) => r,
}
};
let (auth, _did_auth) = auth_result
.map_err(|e| {
emit_login_span(
false,
arguments.method_id.0.as_ref(),
None,
Some("login_flow_failed"),
Some(
if cancelled {
"login_cancelled"
} else {
"login_flow_failed"
},
),
);
let mut err = acp::Error::auth_required();
err.message = e.to_string();
@ -2081,8 +2091,8 @@ impl acp::Agent for MvpAgent {
return Ok(acp::PromptResponse::new(acp::StopReason::EndTurn));
}
}
let intake_lock = self.prompt_intake_lock(&arguments.session_id);
let intake_guard = intake_lock.lock().await;
let dispatch_lock = self.dispatch_lock(&arguments.session_id);
let dispatch_guard = dispatch_lock.lock().await;
let meta_prompt_mode = arguments
.meta
.as_ref()
@ -2312,7 +2322,7 @@ impl acp::Agent for MvpAgent {
acp::Error::internal_error()
.data(format!("failed to dispatch prompt to session: {e}"))
})?;
drop(intake_guard);
drop(dispatch_guard);
self.push_roster_activity_delta(
&arguments.session_id,
crate::agent::roster::RosterActivity::Working,
@ -3081,6 +3091,8 @@ impl acp::Agent for MvpAgent {
.and_then(|m| m.get("rewindIfPristine"))
.and_then(|v| v.as_bool())
.unwrap_or(false);
let dispatch_lock = self.dispatch_lock(&args.session_id);
let _dispatch_guard = dispatch_lock.lock().await;
let _ = handle
.cmd_tx
.send(SessionCommand::Cancel {

39
crates/codegen/xai-grok-shell/src/agent/mvp_agent/agent_ops.rs
View file

@ -3,6 +3,18 @@
//! Inherent [`MvpAgent`] helpers (MCP/clients/gateway, settings/models, session ops, spawn).
//! Co-located child of `mvp_agent` (`use super::*`).
use super::*;
/// `preferred` model, else catalog `current`, else first with own credentials.
fn byok_from_models(
models: &indexmap::IndexMap<String, ModelEntry>,
preferred: Option<&str>,
current: &str,
) -> Option<String> {
preferred
.and_then(|id| models.get(id))
.and_then(|m| m.own_credential())
.or_else(|| models.get(current).and_then(|m| m.own_credential()))
.or_else(|| models.values().find_map(|m| m.own_credential()))
}
impl MvpAgent {
pub(super) fn resolve_image_description_model(&self) -> String {
self.cfg
@ -78,6 +90,20 @@ impl MvpAgent {
pub(super) fn set_auth_method(&self, id: acp::AuthMethodId) {
self.auth_method_id.store(Some(std::sync::Arc::new(id)));
}
/// Publish model-owned credentials for voice/tools static fallthrough.
/// Only [`ModelEntry::own_credential`] — not `sampling_config.api_key` (may be a session JWT).
pub(crate) fn sync_process_static_api_key(&self, preferred_model_id: Option<&str>) {
if self.cfg.borrow().grok_com_config.api_key_auth_disabled() {
self.auth_manager.set_process_static_api_key(None);
return;
}
let models = self.models_manager.models();
let current = self.models_manager.current_model_id();
self.auth_manager
.set_process_static_api_key(
byok_from_models(&models, preferred_model_id, current.0.as_ref()),
);
}
/// Return auth for sync config construction.
pub(super) fn current_or_buffered_auth(&self) -> Option<crate::auth::GrokAuth> {
self.auth_manager
@ -1414,6 +1440,14 @@ impl MvpAgent {
) -> Self {
models_manager.set_gateway(gateway.clone());
let sampling_config = models_manager.sampling_config();
if !cfg.grok_com_config.api_key_auth_disabled() {
let models = models_manager.models();
let current = models_manager.current_model_id();
auth_manager
.set_process_static_api_key(
byok_from_models(&models, None, current.0.as_ref()),
);
}
crate::upload::trace::spawn_purge_stale_upload_scratch();
let storage_mode = cfg.storage_mode;
let default_yolo_mode = cfg.default_yolo_mode;
@ -1468,7 +1502,7 @@ impl MvpAgent {
sessions: RefCell::new(HashMap::new()),
activity,
loading_sessions: RefCell::new(HashMap::new()),
prompt_intake_locks: RefCell::new(HashMap::new()),
dispatch_locks: RefCell::new(HashMap::new()),
session_threads: RefCell::new(HashMap::new()),
resident_roster_titles: RefCell::new(HashMap::new()),
initialize_request: OnceLock::new(),
@ -1504,8 +1538,7 @@ impl MvpAgent {
auth_method_id: crate::agent::auth_method::new_shared_auth_method_id(None),
sampling_config: RefCell::new(sampling_config),
auth_manager,
auth_code_tx: RefCell::new(None),
auth_url_rx: RefCell::new(None),
interactive_auth: Default::default(),
client_type: RefCell::new(ClientType::default()),
code_nav_enabled: std::cell::Cell::new(false),
interactive_trust_client: std::cell::Cell::new(false),

36
crates/codegen/xai-grok-shell/src/agent/mvp_agent/mod.rs
View file

@ -62,7 +62,7 @@ use xai_grok_sampling_types::{
supports_reasoning_effort_meta,
};
use crate::agent::update_chunk_merge;
use crate::auth::{AuthManager, AuthUrlInfo};
use crate::auth::AuthManager;
use crate::config::StorageMode;
use crate::extensions::notification::{SessionNotification, SessionUpdate};
use xai_grok_telemetry::id::{agent_id, agent_instance_id};
@ -614,19 +614,14 @@ pub struct MvpAgent {
loading_sessions: RefCell<
HashMap<acp::SessionId, tokio::sync::watch::Receiver<bool>>,
>,
/// Per-session prompt-intake serialization lock. LEADER-SAFE(per-session):
/// keyed by SessionId, mirrors `sessions` lifecycle.
///
/// Each incoming `session/prompt` RPC is dispatched as its own task by the
/// ACP message loop, and [`Self::prompt`] runs an async preamble (prompt-mode
/// query, trace context, model lookup) BEFORE it enqueues
/// `SessionCommand::Prompt` onto the actor's FIFO mailbox. Without
/// serialization those preambles interleave across tasks, so the mailbox —
/// and therefore the authoritative prompt queue — receives prompts out of
/// submission order. `prompt()` holds this lock across the preamble and
/// releases it immediately after the enqueue (the turn itself runs unlocked),
/// which makes intake order match arrival order.
prompt_intake_locks: RefCell<
/// Per-session lock ordering dispatch onto the actor's mailbox:
/// [`Self::prompt`] holds it across its intake preamble and
/// [`Self::cancel`] around its `Cancel` send, so prompts land in
/// submission order and a cancel cannot overtake the prompt it targets
/// (see `cancel_never_overtakes_in_flight_prompt_intake`). Cancels wait
/// out preambles held ahead of them — keep preambles lean; bridge cancels
/// are unordered. LEADER-SAFE(per-session): mirrors `sessions` lifecycle.
dispatch_locks: RefCell<
HashMap<acp::SessionId, std::rc::Rc<tokio::sync::Mutex<()>>>,
>,
/// LEADER-SAFE(per-session): keyed by SessionId. Mirrors `sessions` lifecycle.
@ -654,10 +649,11 @@ pub struct MvpAgent {
/// grok.com chat-product catalog (`/rest/modes`) for chat sessions; distinct
/// from `models_manager` (the build `/v1/models` catalog).
pub(crate) chat_modes: crate::agent::chat_modes::ChatModesManager,
/// Forwards pasted codes from `handle_auth_submit_code` to the auth flow.
pub(crate) auth_code_tx: RefCell<Option<tokio::sync::mpsc::Sender<String>>>,
/// Receives the auth URL from the auth flow; read by `handle_auth_get_url`.
pub(crate) auth_url_rx: RefCell<Option<tokio::sync::oneshot::Receiver<AuthUrlInfo>>>,
/// Single-flight guard for interactive login (device poll / loopback
/// wait). Owns the active attempt's cancel token and its code/url
/// channels; a new `authenticate` or `x.ai/auth/cancel` cancels the
/// prior attempt.
pub(crate) interactive_auth: crate::auth::single_flight::AuthSingleFlight,
/// Client type. LEADER-SAFE(init-once): set once during `initialize` from
/// `_meta.clientIdentifier` (injected by the IPC server in leader mode).
///
@ -1127,6 +1123,10 @@ struct AuthRequestMeta {
/// user abandons the browser flow, the current session continues.
#[serde(default)]
force_interactive: bool,
/// Pager auth `request_seq` for this attempt. Scopes `x.ai/auth/cancel`
/// so a delayed cancel cannot tear down a successor login.
#[serde(default)]
request_seq: Option<u64>,
}
impl AuthRequestMeta {
/// `--oauth` → force loopback; otherwise default (loopback).

19
crates/codegen/xai-grok-shell/src/agent/mvp_agent/session_lifecycle.rs
View file

@ -30,14 +30,10 @@ impl MvpAgent {
});
}
}
/// Remove a session and its thread handle without finalizing the cloud
/// replica; the conversation stays resumable on disk. Reached by
/// dead-actor reaping and the terminal close and delete paths. Idle
/// unload does not route here: `handle_evict_sessions` removes its
/// handle inline and keeps the thread for reconnect.
/// Remove a session without finalizing; it stays resumable on disk.
pub(crate) fn remove_session(&self, id: &acp::SessionId) {
self.sessions.borrow_mut().remove(id);
self.prompt_intake_locks.borrow_mut().remove(id);
self.dispatch_locks.borrow_mut().remove(id);
self.session_threads.borrow_mut().remove(id);
self.session_index_claims.borrow_mut().remove(id);
self.require_gateway_sessions.borrow_mut().remove(id);
@ -51,13 +47,10 @@ impl MvpAgent {
ops.end_local_session(id.0.as_ref());
}
}
/// Get-or-create the per-session prompt-intake lock (see
/// [`Self::prompt_intake_locks`]). Cheap clone of the shared `Rc`.
pub(super) fn prompt_intake_lock(
&self,
id: &acp::SessionId,
) -> std::rc::Rc<tokio::sync::Mutex<()>> {
self.prompt_intake_locks
/// Get-or-create the per-session dispatch lock (see
/// [`Self::dispatch_locks`]). Cheap clone of the shared `Rc`.
pub(super) fn dispatch_lock(&self, id: &acp::SessionId) -> std::rc::Rc<tokio::sync::Mutex<()>> {
self.dispatch_locks
.borrow_mut()
.entry(id.clone())
.or_default()

3
crates/codegen/xai-grok-shell/src/agent/mvp_agent/subagent_coordinator.rs
View file

@ -74,9 +74,6 @@ impl MvpAgent {
.borrow_mut()
.register_block_wait(&subagent_id, slot.clone());
}
this.subagent_coordinator
.borrow_mut()
.evict_stale_completed();
result
};
let snapshot = resolve_snapshot(lookup).await;

52
crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests.rs
View file

@ -3000,6 +3000,58 @@ fn cancel_does_not_forward_to_bridge_in_local_mode() {
);
});
}
/// Regression (post-cancel slot hang, first bad release 0.2.101; see
/// `dispatch_locks`). SDK e2e shape:
/// `test_cancel_ends_in_flight_turn_and_frees_slot` (grok-agent-sdk).
#[test]
fn cancel_never_overtakes_in_flight_prompt_intake() {
use crate::session::SessionCommand;
use acp::Agent as _;
run_local_for_bridge_test(|| async {
let agent = build_minimal_agent_for_tests();
let sid = acp::SessionId::new("sess-cancel-intake-race");
let (handle, _tx, mut cmd_rx) = make_live_session_handle(&sid, None);
agent.sessions.borrow_mut().insert(sid.clone(), handle);
let order: std::rc::Rc<std::cell::RefCell<Vec<&'static str>>> =
std::rc::Rc::new(std::cell::RefCell::new(Vec::new()));
let (intake_parked_tx, intake_parked_rx) = tokio::sync::oneshot::channel::<()>();
let driver_order = order.clone();
tokio::task::spawn_local(async move {
let mut intake_parked_tx = Some(intake_parked_tx);
while let Some(cmd) = cmd_rx.recv().await {
match cmd {
SessionCommand::GetCurrentPromptMode { .. } => {
if let Some(tx) = intake_parked_tx.take() {
let _ = tx.send(());
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
}
SessionCommand::Prompt { .. } => driver_order.borrow_mut().push("prompt"),
SessionCommand::Cancel { .. } => driver_order.borrow_mut().push("cancel"),
_ => {}
}
}
});
let prompt_fut = agent.prompt(acp::PromptRequest::new(
sid.clone(),
vec![acp::ContentBlock::from("hi")],
));
let cancel_fut = async {
intake_parked_rx
.await
.expect("prompt intake reaches the fake actor");
let _ = agent
.cancel(acp::CancelNotification::new(sid.clone()))
.await;
};
let _ = futures::join!(prompt_fut, cancel_fut);
assert_eq!(
order.borrow().as_slice(),
["prompt", "cancel"],
"cancel must land on the actor mailbox after the prompt it targets"
);
});
}
use crate::session::SessionCommand as TestSessionCommand;
/// Build a session handle wired to a *live* command channel. Returns the
/// handle (move into `sessions`) plus a probe `cmd_tx`/`cmd_rx` so a test

52
crates/codegen/xai-grok-shell/src/agent/subagent/coordinator_lifecycle.rs
View file

@ -323,8 +323,10 @@ impl SubagentCoordinator {
effective_model_id: String::new(),
block_waited: false,
explicitly_killed: false,
persisted_output_dir: None,
},
);
self.enforce_completed_cap();
if surface_completion {
self.pending_completions
.push(SubagentCompletionSummary {
@ -353,6 +355,7 @@ impl SubagentCoordinator {
description: String,
subagent_type: String,
result: SubagentResult,
persisted_output_dir: Option<PathBuf>,
) -> Option<SubagentTracker> {
let tracker = self.active.remove(id);
self.sync_running_gauge();
@ -383,30 +386,26 @@ impl SubagentCoordinator {
let block_waited = tracker.as_ref().is_some_and(|t| t.block_waited);
let explicitly_killed = tracker.as_ref().is_some_and(|t| t.explicitly_killed);
let surface_completion = tracker.as_ref().is_none_or(|t| t.surface_completion);
self.completed
.insert(
id.to_string(),
CompletedSubagent {
subagent_id: id.to_string(),
parent_session_id,
parent_prompt_id,
child_session_id,
description,
subagent_type,
persona,
started_at,
completed_at: std::time::Instant::now(),
result,
resumed_from,
child_cwd,
worktree_path,
snapshot_ref: None,
effective_model_id,
block_waited,
explicitly_killed,
},
);
let completed = self.completed.get(id).expect("just inserted");
let mut completed = CompletedSubagent {
subagent_id: id.to_string(),
parent_session_id,
parent_prompt_id,
child_session_id,
description,
subagent_type,
persona,
started_at,
completed_at: std::time::Instant::now(),
result,
resumed_from,
child_cwd,
worktree_path,
snapshot_ref: None,
effective_model_id,
block_waited,
explicitly_killed,
persisted_output_dir,
};
let success = completed.result.success && !completed.result.cancelled;
{
let preview = crate::util::truncate(&completed.result.output, 200);
@ -444,6 +443,11 @@ impl SubagentCoordinator {
output: completed.result.output.clone(),
});
}
if completed.persisted_output_dir.is_some() {
completed.result.output = Arc::from("");
}
self.completed.insert(id.to_string(), completed);
self.enforce_completed_cap();
self.completion_notify.notify_waiters();
tracker
}

32
crates/codegen/xai-grok-shell/src/agent/subagent/coordinator_query.rs
View file

@ -55,8 +55,17 @@ impl SubagentCoordinator {
reason: completed.result.error.clone(),
}
} else if completed.result.success {
let output = match &completed.persisted_output_dir {
Some(dir) => {
read_subagent_output(dir)
.unwrap_or_else(|| {
OUTPUT_UNAVAILABLE_PLACEHOLDER.to_string()
})
}
None => completed.result.output.to_string(),
};
SubagentSnapshotStatus::Completed {
output: completed.result.output.to_string(),
output,
tool_calls: completed.result.tool_calls,
turns: completed.result.turns,
worktree_path: completed.result.worktree_path.clone(),
@ -209,7 +218,7 @@ impl SubagentCoordinator {
/// to a different parent session (prevents cross-session context bleed).
///
/// Fast path: checks the in-memory `completed` map first. When that
/// misses (e.g. after TTL eviction), falls back to on-disk metadata
/// misses (e.g. after cap eviction), falls back to on-disk metadata
/// in `{parent_session_dir}/subagents/{id}/meta.json`.
pub(crate) fn resumable_source_for(
&self,
@ -293,10 +302,21 @@ impl SubagentCoordinator {
will_wake: false,
})
}
/// TTL cleanup: remove completed entries older than 30 minutes.
pub fn evict_stale_completed(&mut self) {
let cutoff = std::time::Duration::from_secs(30 * 60);
self.completed.retain(|_, entry| entry.completed_at.elapsed() < cutoff);
/// Oldest completions are evicted first; their `output.json` stays on disk.
pub fn enforce_completed_cap(&mut self) {
if self.completed.len() <= MAX_COMPLETED_ENTRIES {
return;
}
let excess = self.completed.len() - MAX_COMPLETED_ENTRIES;
let mut by_age: Vec<(std::time::Instant, String)> = self
.completed
.iter()
.map(|(id, e)| (e.completed_at, id.clone()))
.collect();
by_age.sort_unstable_by_key(|(completed_at, _)| *completed_at);
for (_, id) in by_age.into_iter().take(excess) {
self.completed.remove(&id);
}
}
/// Snapshot all currently-running subagents for compaction state context.
///

4
crates/codegen/xai-grok-shell/src/agent/subagent/handle_request.rs
View file

@ -1765,7 +1765,8 @@ pub(crate) async fn handle_subagent_request(
}
}
}
update_subagent_meta_completed(&subagent_meta_dir, &result, &gcs_upload_ctx);
let persisted_output_dir = persist_subagent_output(&subagent_meta_dir, &result);
persist_subagent_completion(&subagent_meta_dir, &result, &gcs_upload_ctx);
let final_status = result.status().to_string();
let snapshot_dispose_enabled = ctx.resolve_subagent_worktree_snapshot_enabled();
let telemetry_tokens = if result.tool_calls > 0 || result.success {
@ -1994,6 +1995,7 @@ pub(crate) async fn handle_subagent_request(
request.description.clone(),
request.subagent_type.clone(),
result.clone(),
persisted_output_dir,
);
if let Some(snapshot_ref) = disposed_snapshot_ref {
coordinator.borrow_mut().set_completed_snapshot_ref(&request.id, snapshot_ref);

96
crates/codegen/xai-grok-shell/src/agent/subagent/mod.rs
View file

@ -474,7 +474,7 @@ pub(crate) struct CompletedSubagent {
pub subagent_type: String,
pub persona: Option<String>,
pub started_at: std::time::Instant,
/// When the subagent moved to the completed map. Used for TTL eviction.
/// When the subagent moved to the completed map. Orders cap eviction.
pub completed_at: std::time::Instant,
pub result: SubagentResult,
/// ID of the source subagent this session was resumed from.
@ -492,6 +492,12 @@ pub(crate) struct CompletedSubagent {
pub block_waited: bool,
/// Set when the model explicitly killed this subagent via the kill tool.
pub explicitly_killed: bool,
/// Directory whose `output.json` holds the output text; when set, the
/// stored `result.output` is cleared and `lookup` reads from disk.
/// `None` (failures, empty outputs, failed writes) serves from memory.
/// Process-scoped and local-only: resume survives a restart via
/// `meta.json`, and trace upload carries the text to GCS.
pub persisted_output_dir: Option<PathBuf>,
}
/// Lightweight entry for subagents that have been requested but are still
/// initializing (creating worktree, resolving config, spawning session).
@ -583,6 +589,11 @@ pub(crate) struct SubagentCoordinator {
/// Cleared on freeze/cancel. See AGENTS.md rule 3 for the completeness model.
subagent_usage_not_applied_prompts: std::collections::HashSet<String>,
}
/// Cap on the completed map (entries are small: identity, counts, and an
/// error string; successful output text lives in `output.json`).
pub(crate) const MAX_COMPLETED_ENTRIES: usize = 1024;
/// Served when an entry's `output.json` cannot be read back.
pub(crate) const OUTPUT_UNAVAILABLE_PLACEHOLDER: &str = "[subagent output no longer available]";
fn tracker_to_summary(t: &SubagentTracker) -> ActiveSubagentSummary {
ActiveSubagentSummary {
subagent_id: t.subagent_id.clone(),
@ -2131,7 +2142,7 @@ fn fail_subagent(
duration_ms,
..Default::default()
};
update_subagent_meta_completed(subagent_meta_dir, &result, gcs_ctx);
persist_subagent_completion(subagent_meta_dir, &result, gcs_ctx);
emit_subagent_notification(
gateway,
parent_session_id,
@ -2190,7 +2201,7 @@ async fn cancel_pending_subagent_at_promote(
duration_ms,
..Default::default()
};
update_subagent_meta_completed(subagent_meta_dir, &result, gcs_ctx);
persist_subagent_completion(subagent_meta_dir, &result, gcs_ctx);
emit_subagent_notification(
gateway,
parent_session_id,
@ -2552,14 +2563,21 @@ impl SubagentSessionMetadata {
}
}
}
/// Write via a same-directory temp file and rename, so a crash mid-write
/// cannot leave a torn `meta.json` or `output.json`.
fn atomic_write(path: &Path, contents: &str) -> std::io::Result<()> {
let parent = path.parent().ok_or_else(|| {
std::io::Error::new(std::io::ErrorKind::InvalidInput, "path has no parent")
})?;
std::fs::create_dir_all(parent)?;
let tmp = tempfile::NamedTempFile::new_in(parent)?;
std::fs::write(tmp.path(), contents)?;
tmp.persist(path)?;
Ok(())
}
/// Write `meta.json`. Returns `true` on success so callers on the resume-pointer
/// path can gate worktree disposal on a durable write.
fn write_subagent_meta(dir: &Path, meta: &SubagentMeta) -> bool {
if let Err(e) = std::fs::create_dir_all(dir) {
tracing::warn!(error = % e, "failed to create subagent meta dir");
return false;
}
let meta_path = dir.join("meta.json");
let json = match serde_json::to_string_pretty(meta) {
Ok(json) => json,
Err(e) => {
@ -2567,12 +2585,63 @@ fn write_subagent_meta(dir: &Path, meta: &SubagentMeta) -> bool {
return false;
}
};
if let Err(e) = std::fs::write(&meta_path, json) {
if let Err(e) = atomic_write(&dir.join("meta.json"), &json) {
tracing::warn!(error = % e, "failed to write subagent meta");
return false;
}
true
}
/// On-disk schema of `output.json`, written beside `meta.json`.
#[derive(serde::Deserialize)]
struct SubagentOutputFile {
schema_version: u32,
output: String,
}
/// Borrowed twin of [`SubagentOutputFile`] so serialization does not copy
/// the output text.
#[derive(serde::Serialize)]
struct SubagentOutputFileRef<'a> {
schema_version: u32,
output: &'a str,
}
const SUBAGENT_OUTPUT_SCHEMA_VERSION: u32 = 1;
fn write_subagent_output(dir: &Path, output: &str) -> bool {
let file = SubagentOutputFileRef {
schema_version: SUBAGENT_OUTPUT_SCHEMA_VERSION,
output,
};
let json = match serde_json::to_string(&file) {
Ok(json) => json,
Err(e) => {
tracing::warn!(error = % e, "failed to serialize subagent output");
return false;
}
};
if let Err(e) = atomic_write(&dir.join("output.json"), &json) {
tracing::warn!(error = % e, "failed to write subagent output");
return false;
}
true
}
/// Read back `output.json`. `None` on any read or parse failure.
pub(crate) fn read_subagent_output(dir: &Path) -> Option<String> {
let data = std::fs::read_to_string(dir.join("output.json")).ok()?;
let file: SubagentOutputFile = match serde_json::from_str(&data) {
Ok(file) => file,
Err(e) => {
tracing::warn!(error = % e, "failed to parse subagent output.json");
return None;
}
};
if file.schema_version != SUBAGENT_OUTPUT_SCHEMA_VERSION {
tracing::warn!(
found = file.schema_version,
expected = SUBAGENT_OUTPUT_SCHEMA_VERSION,
"unexpected output.json schema version"
);
}
Some(file.output)
}
/// Extra runtime context for GCS artifact upload. `SubagentMeta` doesn't
/// persist these fields, so they're carried from the spawn site.
#[derive(Clone)]
@ -2595,7 +2664,7 @@ struct GcsUploadContext {
/// any read/parse/write failure is `warn!`-logged (this is the critical resume
/// pointer) so the caller keeps the worktree rather than removing it without a
/// recoverable ref. Also re-asserts the terminal `status` so a failed
/// `update_subagent_meta_completed` write can't leave a non-terminal record that
/// `persist_subagent_completion` write can't leave a non-terminal record that
/// `resumable_source_for` rejects after the worktree is removed.
fn update_subagent_meta_snapshot_ref(dir: &Path, snapshot_ref: &str, status: &str) -> bool {
let meta_path = dir.join("meta.json");
@ -2622,7 +2691,12 @@ fn update_subagent_meta_snapshot_ref(dir: &Path, snapshot_ref: &str, status: &st
meta.status = status.to_string();
write_subagent_meta(dir, &meta)
}
fn update_subagent_meta_completed(dir: &Path, result: &SubagentResult, gcs_ctx: &GcsUploadContext) {
#[must_use]
fn persist_subagent_output(dir: &Path, result: &SubagentResult) -> Option<PathBuf> {
(result.success && !result.output.is_empty() && write_subagent_output(dir, &result.output))
.then(|| dir.to_path_buf())
}
fn persist_subagent_completion(dir: &Path, result: &SubagentResult, gcs_ctx: &GcsUploadContext) {
let meta_path = dir.join("meta.json");
if let Ok(data) = std::fs::read_to_string(&meta_path)
&& let Ok(mut meta) = serde_json::from_str::<SubagentMeta>(&data)

196
crates/codegen/xai-grok-shell/src/agent/subagent/tests/mod.rs
View file

@ -208,6 +208,7 @@ fn lookup_returns_ready_for_completed_subagent() {
duration_ms: 1234,
..Default::default()
},
None,
);
let lookup = coordinator.lookup("sub-1");
assert!(lookup.is_some());
@ -384,6 +385,7 @@ async fn running_gauge_tracks_pending_and_active() {
"gauge task".into(),
"general-purpose".into(),
SubagentResult::default(),
None,
);
assert_eq!(gauge.load(Ordering::Relaxed), 0, "completed does not count");
coordinator
@ -435,6 +437,7 @@ fn mark_block_waited_sets_flag_on_completed() {
child_session_id: "sub-bw".into(),
..Default::default()
},
None,
);
assert!(! coordinator.is_block_waited("sub-bw"));
coordinator.mark_block_waited("sub-bw");
@ -525,6 +528,7 @@ async fn mark_explicitly_killed_active_then_propagates_to_completed() {
child_session_id: "sub-ek".into(),
..Default::default()
},
None,
);
assert!(
coordinator.is_explicitly_killed("sub-ek"),
@ -631,6 +635,7 @@ fn mark_explicitly_killed_sets_flag_on_completed() {
child_session_id: "sub-ek-c".into(),
..Default::default()
},
None,
);
assert!(! coordinator.is_explicitly_killed("sub-ek-c"));
coordinator.mark_explicitly_killed("sub-ek-c");
@ -658,6 +663,7 @@ async fn block_waited_propagates_through_move_to_completed() {
child_session_id: "sub-prop".into(),
..Default::default()
},
None,
);
assert!(coordinator.is_block_waited("sub-prop"));
}
@ -676,6 +682,7 @@ fn complete_dummy(coordinator: &mut SubagentCoordinator, id: &str, surface: bool
child_session_id: id.into(),
..Default::default()
},
None,
);
}
#[tokio::test]
@ -861,44 +868,173 @@ fn move_pending_to_cancelled_creates_cancelled_entry() {
}
}
}
fn completed_with_output(
id: &str,
text: &str,
persisted_output_dir: Option<PathBuf>,
) -> CompletedSubagent {
CompletedSubagent {
subagent_id: id.into(),
parent_session_id: String::new(),
parent_prompt_id: None,
child_session_id: String::new(),
description: "task".into(),
subagent_type: "explore".into(),
persona: None,
started_at: std::time::Instant::now(),
completed_at: std::time::Instant::now(),
result: SubagentResult {
success: true,
output: std::sync::Arc::from(text),
..Default::default()
},
resumed_from: None,
child_cwd: String::new(),
worktree_path: None,
snapshot_ref: None,
effective_model_id: String::new(),
block_waited: false,
explicitly_killed: false,
persisted_output_dir,
}
}
fn lookup_output(coordinator: &SubagentCoordinator, id: &str) -> String {
match coordinator.lookup(id) {
Some(SnapshotLookup::Ready(snap)) => {
match snap.status {
SubagentSnapshotStatus::Completed { output, .. } => output,
other => panic!("expected Completed status, got {other:?}"),
}
}
other => {
panic!(
"expected Ready lookup, got {:?}", other.map(| _ | "NeedsSignals/other")
)
}
}
}
#[test]
fn evict_stale_completed_uses_completion_time() {
fn lookup_degrades_to_placeholder_when_output_file_is_missing() {
let dir = tempfile::tempdir().expect("tempdir");
let mut coordinator = SubagentCoordinator::new();
coordinator
.completed
.insert(
"sub-recent".to_string(),
CompletedSubagent {
subagent_id: "sub-recent".into(),
parent_session_id: String::new(),
parent_prompt_id: None,
child_session_id: String::new(),
description: "long-running".into(),
subagent_type: "explore".into(),
persona: None,
started_at: std::time::Instant::now()
- std::time::Duration::from_secs(31 * 60),
completed_at: std::time::Instant::now(),
result: SubagentResult {
success: true,
..Default::default()
},
resumed_from: None,
child_cwd: String::new(),
worktree_path: None,
snapshot_ref: None,
effective_model_id: String::new(),
block_waited: false,
explicitly_killed: false,
},
"sub-gone".to_string(),
completed_with_output("sub-gone", "", Some(dir.path().to_path_buf())),
);
coordinator.evict_stale_completed();
assert!(
coordinator.completed.contains_key("sub-recent"),
"recently completed subagent should not be evicted"
assert_eq!(
lookup_output(& coordinator, "sub-gone"), OUTPUT_UNAVAILABLE_PLACEHOLDER,
"an entry whose output.json is gone must degrade, not fail the query"
);
}
#[test]
fn lookup_serves_unpersisted_output_from_memory() {
let mut coordinator = SubagentCoordinator::new();
coordinator
.completed
.insert("sub-mem".to_string(), completed_with_output("sub-mem", "output", None));
assert_eq!(
lookup_output(& coordinator, "sub-mem"), "output",
"an entry with nothing on disk must serve its in-memory output"
);
}
#[test]
fn completed_entries_are_capped_oldest_first() {
let mut coordinator = SubagentCoordinator::new();
let base = std::time::Instant::now();
for i in 0..(MAX_COMPLETED_ENTRIES + 2) {
let mut entry = completed_with_output(
&format!("sub-{i}"),
"",
Some(std::path::PathBuf::from("/nonexistent")),
);
entry.completed_at = base + std::time::Duration::from_millis(i as u64);
coordinator.completed.insert(format!("sub-{i}"), entry);
}
coordinator.enforce_completed_cap();
assert_eq!(
coordinator.completed.len(), MAX_COMPLETED_ENTRIES,
"the completed map must be capped at MAX_COMPLETED_ENTRIES"
);
assert!(
! coordinator.completed.contains_key("sub-0") && ! coordinator.completed
.contains_key("sub-1"), "the oldest completions must be evicted first"
);
assert!(
coordinator.completed.contains_key("sub-2"),
"entries within the cap must survive"
);
}
#[test]
fn move_to_completed_clears_persisted_output_after_the_summary_clone() {
let dir = tempfile::tempdir().expect("tempdir");
let full_output = "final report".repeat(100);
assert!(write_subagent_output(dir.path(), & full_output));
let mut coordinator = SubagentCoordinator::new();
coordinator
.move_to_completed(
"sub-e2e",
"task".into(),
"explore".into(),
SubagentResult {
success: true,
output: std::sync::Arc::from(full_output.as_str()),
subagent_id: "sub-e2e".into(),
child_session_id: "sub-e2e".into(),
..Default::default()
},
Some(dir.path().to_path_buf()),
);
let entry = coordinator.completed.get("sub-e2e").expect("entry inserted");
assert!(
entry.result.output.is_empty(),
"a persisted entry must not keep the output in memory"
);
assert_eq!(
lookup_output(& coordinator, "sub-e2e"), full_output,
"lookup must serve the persisted output from disk"
);
let summaries = coordinator.drain_pending_completions();
assert_eq!(
&* summaries[0].output, full_output,
"the completion summary must carry the full output"
);
}
#[test]
fn persist_gate_only_persists_successful_nonempty_outputs() {
let dir = tempfile::tempdir().expect("tempdir");
let ok = SubagentResult {
success: true,
output: std::sync::Arc::from("text"),
..Default::default()
};
assert_eq!(
persist_subagent_output(dir.path(), & ok), Some(dir.path().to_path_buf())
);
let empty = SubagentResult {
success: true,
..Default::default()
};
assert_eq!(persist_subagent_output(dir.path(), & empty), None);
let failed = SubagentResult {
success: false,
output: std::sync::Arc::from("partial"),
..Default::default()
};
assert_eq!(persist_subagent_output(dir.path(), & failed), None);
}
#[test]
fn subagent_output_roundtrips_through_output_json() {
let dir = tempfile::tempdir().expect("tempdir");
let output = "line one\nline two with unicode ✓";
assert!(write_subagent_output(dir.path(), output));
assert_eq!(read_subagent_output(dir.path()).as_deref(), Some(output));
assert_eq!(read_subagent_output(& dir.path().join("missing")), None);
std::fs::write(dir.path().join("output.json"), "not json").expect("corrupt file");
assert_eq!(read_subagent_output(dir.path()), None);
}
#[test]
fn cancel_with_outcome_fires_pending_token() {
let mut coordinator = SubagentCoordinator::new();
let token = CancellationToken::new();
@ -945,6 +1081,7 @@ async fn cancel_with_outcome_returns_variant_for_active_finished_unknown() {
subagent_id: "sub-done".to_string(),
..Default::default()
},
None,
);
assert!(
matches!(coordinator.cancel_with_outcome("sub-done"),
@ -1026,6 +1163,7 @@ fn completed_takes_precedence_over_pending_in_lookup() {
child_session_id: "child-dup".to_string(),
..Default::default()
},
None,
);
let lookup = coordinator.lookup("sub-dup");
assert!(

26
crates/codegen/xai-grok-shell/src/agent/subagent/tests/rest.rs
View file

@ -346,6 +346,7 @@ fn resumable_source_returns_info_for_completed_subagent() {
effective_model_id: "grok-3".into(),
block_waited: false,
explicitly_killed: false,
persisted_output_dir: None,
},
);
let info = coordinator
@ -373,6 +374,7 @@ fn resumable_source_survives_move_to_completed_with_metadata() {
child_session_id: "sub-moved".into(),
..Default::default()
},
None,
);
let info = coordinator
.resumable_source_for("sub-moved", "", Path::new("/tmp"))
@ -593,7 +595,7 @@ fn snapshot_ref_write_promotes_nonterminal_status_to_terminal() {
assert_eq!("completed", reread.status);
}
/// The coordinator setter stamps the snapshot ref onto the in-memory
/// completed entry so `resume_from` can rehydrate before TTL eviction.
/// completed entry so `resume_from` can rehydrate before cap eviction.
#[tokio::test]
async fn set_completed_snapshot_ref_updates_in_memory_entry() {
let mut coordinator = SubagentCoordinator::new();
@ -609,6 +611,7 @@ async fn set_completed_snapshot_ref_updates_in_memory_entry() {
child_session_id: "sa-mem".into(),
..Default::default()
},
None,
);
let before = coordinator
.resumable_source_for("sa-mem", "session-A", Path::new("/tmp"))
@ -621,7 +624,7 @@ async fn set_completed_snapshot_ref_updates_in_memory_entry() {
.unwrap();
assert_eq!(after.snapshot_ref.as_deref(), Some("refs/grok/subagents/sa-mem"));
}
/// Unknown id is a no-op (entry already TTL-evicted; meta.json still holds it).
/// Unknown id is a no-op (entry already cap-evicted; meta.json still holds it).
#[test]
fn set_completed_snapshot_ref_unknown_id_is_noop() {
let mut coordinator = SubagentCoordinator::new();
@ -704,6 +707,7 @@ fn coordinator_with_completed(id: &str) -> SubagentCoordinator {
child_session_id: id.into(),
..Default::default()
},
None,
);
coordinator
}
@ -775,7 +779,8 @@ async fn gate_on_completion_clears_model_facing_worktree_path_but_resume_retains
if worktree_removed {
result.worktree_path = None;
}
coordinator.move_to_completed("disp-1", "task".into(), "explore".into(), result);
coordinator
.move_to_completed("disp-1", "task".into(), "explore".into(), result, None);
coordinator
.set_completed_snapshot_ref("disp-1", "refs/grok/subagents/disp-1".into());
let listed = coordinator.completed.get("disp-1").expect("completed entry");
@ -805,7 +810,8 @@ async fn gate_on_completion_retains_worktree_path_when_not_removed() {
if worktree_removed {
result.worktree_path = None;
}
coordinator.move_to_completed("keep-1", "task".into(), "explore".into(), result);
coordinator
.move_to_completed("keep-1", "task".into(), "explore".into(), result, None);
let entry = coordinator.completed.get("keep-1").expect("completed entry");
assert_eq!(Some(wt.to_string_lossy().into_owned()), entry.result.worktree_path);
}
@ -856,6 +862,7 @@ async fn disposal_completes_before_subagent_is_observable() {
child_session_id: "order-1".into(),
..Default::default()
},
None,
);
if let Some(r) = disposed_snapshot_ref {
coordinator.set_completed_snapshot_ref("order-1", r);
@ -1351,6 +1358,7 @@ fn resumable_source_rejects_cross_session_lookup() {
effective_model_id: String::new(),
block_waited: false,
explicitly_killed: false,
persisted_output_dir: None,
},
);
assert!(
@ -1803,6 +1811,7 @@ async fn reconcile_reemits_rewound_finish_even_when_id_still_in_completed_regist
success: true,
..Default::default()
},
None,
);
let (cmd_tx, mut cmd_rx) = mpsc::unbounded_channel();
let unfinished = vec![(id.to_string(), format!("child-{id}"))];
@ -1845,6 +1854,7 @@ async fn reconcile_reemits_real_outcome_for_completed_with_running_meta() {
success: true,
..Default::default()
},
None,
);
let (cmd_tx, mut cmd_rx) = mpsc::unbounded_channel();
let unfinished = vec![(id.to_string(), format!("child-{id}"))];
@ -2166,6 +2176,7 @@ fn completed_subagent_propagates_resumed_from() {
effective_model_id: "grok-3".into(),
block_waited: false,
explicitly_killed: false,
persisted_output_dir: None,
},
);
let refs = coordinator.spawned_refs_for_prompt("prompt-1");
@ -2193,6 +2204,7 @@ async fn completion_notify_fires_on_move_to_completed() {
duration_ms: 100,
..Default::default()
},
None,
);
tokio::time::timeout(std::time::Duration::from_millis(50), notified)
.await
@ -2216,6 +2228,7 @@ fn drain_pending_completions_returns_and_clears() {
duration_ms: 500,
..Default::default()
},
None,
);
coordinator
.move_to_completed(
@ -2231,6 +2244,7 @@ fn drain_pending_completions_returns_and_clears() {
duration_ms: 200,
..Default::default()
},
None,
);
let summaries = coordinator.drain_pending_completions();
assert_eq!(summaries.len(), 2);
@ -2262,6 +2276,7 @@ fn drain_pending_completions_cancelled_is_not_success() {
child_session_id: "sub-c1".to_string(),
..Default::default()
},
None,
);
let summaries = coordinator.drain_pending_completions();
assert_eq!(summaries.len(), 1);
@ -2315,6 +2330,7 @@ async fn outstanding_for_prompt_excludes_completed() {
child_session_id: "sub-done".to_string(),
..Default::default()
},
None,
);
let outstanding = coordinator.outstanding_for_prompt("prompt-X");
assert!(
@ -2369,6 +2385,7 @@ async fn subagent_usage_not_applied_sticky_after_completion_and_is_prompt_scoped
child_session_id: "sub-1".to_string(),
..Default::default()
},
None,
);
assert!(coordinator.outstanding_for_prompt("p-1").is_empty());
assert!(coordinator.subagent_usage_not_applied("p-1"));
@ -2444,6 +2461,7 @@ fn completions_buffered_while_turn_inactive_drained_later() {
child_session_id: "sub-idle".to_string(),
..Default::default()
},
None,
);
let drained = coordinator.drain_pending_completions();
assert_eq!(drained.len(), 1);

43
crates/codegen/xai-grok-shell/src/auth/credential_provider.rs
View file

@ -87,6 +87,24 @@ impl AuthCredentialProvider for ShellAuthCredentialProvider {
self.static_credentials.deployment_key.is_none()
}
}
/// Resolves the embedding credentials for `embed_base_url`, attaching the xAI
/// session credential only to xAI-operated endpoints over `https`.
pub(crate) fn embedding_session_credentials(
embed_base_url: &str,
auth_manager: Option<&Arc<AuthManager>>,
api_key_provider: Option<xai_grok_tools::types::SharedApiKeyProvider>,
) -> xai_grok_memory::EndpointScopedCredentials {
let auth_credentials = auth_manager.map(|am| {
Arc::new(ShellAuthCredentialProvider::new(am.clone(), None, None))
as Arc<dyn AuthCredentialProvider>
});
xai_grok_memory::EndpointScopedCredentials::for_endpoint(
embed_base_url,
crate::util::is_xai_api_bearer_url,
auth_credentials,
api_key_provider,
)
}
/// Build a `StorageClient` for proxy uploads (including the high-volume
/// `batch_upload` used for repo context / `repo_changes_dedup`).
///
@ -568,6 +586,31 @@ mod tests {
"snapshot must reflect refreshed token for subsequent apply() calls"
);
}
#[test]
fn embedding_session_credentials_scopes_to_first_party() {
let _guard = EarlyInvalidationGuard::pin_to_default();
let dir = tempfile::tempdir().unwrap();
let mgr = make_manager(
&dir,
Some(make_auth("xai-session-token", ChronoDuration::hours(1))),
);
let api_key_provider: xai_grok_tools::types::SharedApiKeyProvider =
Arc::new(crate::auth::manager::SharedAuthKeyProvider(mgr.clone()));
for denied in ["https://byok.attacker.example/v1", "http://api.x.ai/v1"] {
let resolved =
embedding_session_credentials(denied, Some(&mgr), Some(api_key_provider.clone()));
assert!(
resolved.is_empty(),
"session credentials must not reach {denied}"
);
}
let resolved = embedding_session_credentials(
"https://api.x.ai/v1",
Some(&mgr),
Some(api_key_provider),
);
assert!(!resolved.is_empty());
}
/// Deployment-key path has no recovery (operator owns the bearer).
#[tokio::test]
async fn refresh_after_unauthorized_is_noop_for_deployment_key() {

20
crates/codegen/xai-grok-shell/src/auth/manager.rs
View file

@ -164,6 +164,9 @@ pub struct AuthManager {
disk_state: RwLock<Option<DiskAuthState>>,
/// See [`Self::cached_disk_api_key`].
static_key_cache: parking_lot::Mutex<Option<StaticKeyCacheEntry>>,
/// Model `api_key` / resolved `env_key` for voice/tools without a session.
/// Not a session token (those live on `inner`). Prefers over disk; env wins.
process_static_api_key: parking_lot::RwLock<Option<String>>,
sleep_gate: SleepGate,
/// Count of in-flight IdP refreshes (the network call only), so a
/// sleep-imminent transition can wait for a refresh straddling suspend to
@ -405,6 +408,7 @@ impl AuthManager {
refresh_notify: Arc::new(tokio::sync::Notify::new()),
disk_state: RwLock::new(disk_state),
static_key_cache: parking_lot::Mutex::new(None),
process_static_api_key: parking_lot::RwLock::new(None),
sleep_gate: SleepGate::default(),
refresh_in_flight: std::sync::atomic::AtomicU32::new(0),
refresh_drain_lock: parking_lot::Mutex::new(()),
@ -2221,11 +2225,8 @@ pub(crate) fn compute_proactive_sleep(this: &AuthManager) -> StdDuration {
}
}
/// Tools + pager voice: session token first, then static API key.
///
/// Static fallthrough (`XAI_API_KEY` / `auth.json` `xai::api_key`) makes voice
/// work on API-key-only setups without OAuth. API-key login already persists
/// the env key to disk.
/// Tools + pager voice bearer. Static: env → process model key → disk.
/// Kill-switch / `preferred_method = oidc` block static keys.
pub(crate) struct SharedAuthKeyProvider(pub Arc<AuthManager>);
impl xai_grok_tools::types::ApiKeyProvider for SharedAuthKeyProvider {
@ -2266,7 +2267,7 @@ fn prefers_static_api_key(am: &AuthManager) -> bool {
)
}
/// Env → `auth.json` `xai::api_key`. Off under kill-switch or `preferred_method = oidc`.
/// Env → process model key → disk. Off under kill-switch / oidc pin.
fn resolve_static_api_key(am: &AuthManager) -> Option<String> {
if am.grok_com_config.api_key_auth_disabled() {
return None;
@ -2278,6 +2279,7 @@ fn resolve_static_api_key(am: &AuthManager) -> Option<String> {
return None;
}
non_empty_key(crate::agent::auth_method::read_xai_api_key_env().ok())
.or_else(|| non_empty_key(am.process_static_api_key.read().clone()))
.or_else(|| am.cached_disk_api_key())
}
@ -2328,6 +2330,12 @@ impl AuthManager {
}
}
}
/// Set the process model key (empty clears). Not for session tokens.
pub fn set_process_static_api_key(&self, key: Option<String>) {
let key = key.map(|k| k.trim().to_string()).filter(|k| !k.is_empty());
*self.process_static_api_key.write() = key;
}
}
fn non_empty_key(key: Option<String>) -> Option<String> {

99
crates/codegen/xai-grok-shell/src/auth/manager_tests.rs
View file

@ -3586,6 +3586,105 @@ async fn shared_api_key_provider_disk_memo_follows_rewrites() {
assert_eq!(provider.current_api_key_async().await, None);
}
#[tokio::test]
#[serial_test::serial]
async fn process_key_from_model_env_key() {
use crate::agent::config::{Config, resolve_model_list};
use xai_grok_test_support::EnvGuard;
const ENV: &str = "TEST_MODEL_ENV_KEY";
const TOKEN: &str = "model-env-token";
let _xai = EnvGuard::unset("XAI_API_KEY");
let _legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY");
let _tok = EnvGuard::set(ENV, TOKEN);
let dm = crate::models::default_model();
let cfg = Config::new_from_toml_cfg(
&toml::from_str(&format!(
r#"
[model."{dm}"]
model = "{dm}"
env_key = "{ENV}"
"#
))
.unwrap(),
)
.unwrap();
let key = resolve_model_list(&cfg, None)
.get(dm)
.and_then(|m| m.own_credential())
.unwrap();
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
assert!(mgr.current().is_none());
mgr.set_process_static_api_key(Some(key));
assert_eq!(
shared_api_key_provider(mgr)
.current_api_key_async()
.await
.as_deref(),
Some(TOKEN)
);
}
#[tokio::test]
#[serial_test::serial]
async fn process_key_precedence() {
use xai_grok_test_support::EnvGuard;
let _xai = EnvGuard::unset("XAI_API_KEY");
let _legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY");
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let provider = shared_api_key_provider(mgr.clone());
assert_eq!(provider.current_api_key_async().await, None);
crate::auth::store_api_key(dir.path(), "disk").unwrap();
assert_eq!(
provider.current_api_key_async().await.as_deref(),
Some("disk")
);
mgr.set_process_static_api_key(Some(" process ".into()));
assert_eq!(
provider.current_api_key_async().await.as_deref(),
Some("process")
);
{
let _key = EnvGuard::set("XAI_API_KEY", "env");
assert_eq!(
provider.current_api_key_async().await.as_deref(),
Some("env")
);
}
mgr.set_process_static_api_key(None);
assert_eq!(
provider.current_api_key_async().await.as_deref(),
Some("disk")
);
let dir_blocked = tempfile::tempdir().unwrap();
let blocked = Arc::new(AuthManager::new(
dir_blocked.path(),
GrokComConfig {
disable_api_key_auth: Some(true),
..GrokComConfig::default()
},
));
blocked.set_process_static_api_key(Some("ignored".into()));
assert_eq!(
shared_api_key_provider(blocked)
.current_api_key_async()
.await,
None
);
}
fn expired_oidc() -> GrokAuth {
GrokAuth {
key: "expired-key".into(),

1
crates/codegen/xai-grok-shell/src/auth/mod.rs
View file

@ -13,6 +13,7 @@ mod model;
pub mod oidc;
pub(crate) mod recovery;
pub(crate) mod refresh;
pub(crate) mod single_flight;
mod storage;
pub(crate) mod token_type;
pub(crate) use config::LEGACY_AUTH_SCOPE;

359
crates/codegen/xai-grok-shell/src/auth/single_flight.rs Normal file
View file

@ -0,0 +1,359 @@
//! Single-flight guard for interactive login.
//!
//! At most one device-code / loopback wait runs at a time: starting a new
//! attempt (or an explicit `x.ai/auth/cancel`) cancels the previous one, so
//! remint/retry cannot stack device-code mints.
//!
//! The attempt owns **all** attempt-scoped state — the cancellation token and
//! the code/url channels — so replacing an attempt swaps everything
//! atomically, and a cancelled predecessor that finishes late structurally
//! cannot touch its successor's channels. Generations guard `end()` the same
//! way: a stale finisher must not clear a newer attempt. Client `request_seq`
//! scopes explicit cancels so a delayed `x.ai/auth/cancel` cannot tear down a
//! successor login.
use std::cell::{Cell, RefCell};
use tokio_util::sync::CancellationToken;
use super::flow::AuthUrlInfo;
/// Channels wired between the ACP ext handlers and one interactive auth flow.
/// `None` for headless attempts (no URL to show, no code to paste).
pub(crate) struct AttemptChannels {
/// Forwards pasted codes from `x.ai/auth/submit_code` to the flow.
code_tx: tokio::sync::mpsc::Sender<String>,
/// Yields the auth URL to `x.ai/auth/get_url`. `Option` so
/// [`AuthSingleFlight::take_url_rx`] can move it out while the attempt
/// lives on (one-shot read).
url_rx: Option<tokio::sync::oneshot::Receiver<AuthUrlInfo>>,
}
struct Attempt {
token: CancellationToken,
channels: Option<AttemptChannels>,
/// Pager `request_seq` for this attempt (scopes delayed cancel RPCs).
client_seq: Option<u64>,
}
/// Why [`AuthSingleFlight::submit_code`] failed.
#[derive(Debug)]
pub(crate) enum SubmitCodeError {
/// No interactive attempt is waiting for a code (idle or headless).
NoPendingAttempt,
/// Channel send failed (attempt channels already closed).
SendFailed(tokio::sync::mpsc::error::TrySendError<String>),
}
#[derive(Default)]
pub(crate) struct AuthSingleFlight {
active: RefCell<Option<Attempt>>,
generation: Cell<u64>,
}
/// RAII end for a [`AuthSingleFlight::begin`] generation: calls [`AuthSingleFlight::end`]
/// on drop so an aborted authenticate future cannot leak attempt state.
pub(crate) struct AuthAttemptGuard<'a> {
sf: &'a AuthSingleFlight,
generation: u64,
ended: Cell<bool>,
}
impl AuthAttemptGuard<'_> {
/// Explicit end (same as drop). Idempotent.
pub(crate) fn end(&self) {
if !self.ended.replace(true) {
self.sf.end(self.generation);
}
}
#[cfg(test)]
pub(crate) fn generation(&self) -> u64 {
self.generation
}
}
impl Drop for AuthAttemptGuard<'_> {
fn drop(&mut self) {
self.end();
}
}
impl AuthSingleFlight {
/// Start a new attempt, cancelling any prior in-flight one. Returns the
/// new attempt's token and an [`AuthAttemptGuard`] that ends this generation
/// on drop (pass no separate `end` — the guard is the only closer).
///
/// `client_seq` is the pager auth `request_seq` (when known); used by
/// [`Self::cancel_for_client_seq`] so a delayed cancel cannot kill a
/// successor attempt.
pub(crate) fn begin(
&self,
channels: Option<AttemptChannels>,
client_seq: Option<u64>,
) -> (CancellationToken, AuthAttemptGuard<'_>) {
let generation = self.generation.get().wrapping_add(1);
self.generation.set(generation);
let token = CancellationToken::new();
if let Some(prev) = self.active.borrow_mut().replace(Attempt {
token: token.clone(),
channels,
client_seq,
}) {
tracing::info!("auth: cancelling prior interactive auth for single-flight");
prev.token.cancel();
}
(
token,
AuthAttemptGuard {
sf: self,
generation,
ended: Cell::new(false),
},
)
}
/// Finish an attempt: drops its token *and channels* only if `generation`
/// is still the active one (a stale finisher must not clear a newer
/// attempt's state).
pub(crate) fn end(&self, generation: u64) {
if self.generation.get() == generation {
*self.active.borrow_mut() = None;
}
}
/// Cancel the active attempt, if any. Idempotent. Prefer
/// [`Self::cancel_for_client_seq`] when the caller has a pager `request_seq`
/// so a delayed cancel cannot tear down a newer login.
pub(crate) fn cancel(&self) {
if let Some(prev) = self.active.borrow_mut().take() {
tracing::info!("auth: interactive auth cancelled");
prev.token.cancel();
}
}
/// Cancel only if the active attempt was started for `client_seq`. A stale
/// cancel (successor already began) is a no-op.
pub(crate) fn cancel_for_client_seq(&self, client_seq: u64) {
let mut active = self.active.borrow_mut();
match active.as_ref() {
Some(a) if a.client_seq == Some(client_seq) => {
if let Some(prev) = active.take() {
tracing::info!(
client_seq,
"auth: interactive auth cancelled for client request_seq"
);
prev.token.cancel();
}
}
Some(a) => {
tracing::debug!(
client_seq,
active_client_seq = ?a.client_seq,
"auth: ignoring stale cancel for superseded request_seq"
);
}
None => {
tracing::debug!(
client_seq,
"auth: cancel_for_client_seq with no active attempt"
);
}
}
}
/// Forward a pasted code to the active attempt's flow.
pub(crate) fn submit_code(&self, code: String) -> Result<(), SubmitCodeError> {
match self
.active
.borrow()
.as_ref()
.and_then(|a| a.channels.as_ref())
{
Some(ch) => ch
.code_tx
.try_send(code)
.map_err(SubmitCodeError::SendFailed),
None => Err(SubmitCodeError::NoPendingAttempt),
}
}
/// Take the active attempt's URL receiver (one-shot; subsequent calls
/// return `None`, as does an idle or headless attempt).
pub(crate) fn take_url_rx(&self) -> Option<tokio::sync::oneshot::Receiver<AuthUrlInfo>> {
self.active
.borrow_mut()
.as_mut()
.and_then(|a| a.channels.as_mut().and_then(|ch| ch.url_rx.take()))
}
}
impl AttemptChannels {
pub(crate) fn new(
code_tx: tokio::sync::mpsc::Sender<String>,
url_rx: tokio::sync::oneshot::Receiver<AuthUrlInfo>,
) -> Self {
Self {
code_tx,
url_rx: Some(url_rx),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn channels() -> (AttemptChannels, tokio::sync::mpsc::Receiver<String>) {
let (code_tx, code_rx) = tokio::sync::mpsc::channel(1);
let (_url_tx, url_rx) = tokio::sync::oneshot::channel();
(AttemptChannels::new(code_tx, url_rx), code_rx)
}
#[test]
fn begin_cancels_prior_attempt() {
let sf = AuthSingleFlight::default();
let (first, _g1) = sf.begin(None, None);
let (second, _g2) = sf.begin(None, None);
assert!(first.is_cancelled(), "prior attempt must be cancelled");
assert!(!second.is_cancelled(), "new attempt must be live");
}
#[test]
fn cancel_stops_active_attempt_and_is_idempotent() {
let sf = AuthSingleFlight::default();
let (token, _g) = sf.begin(None, None);
sf.cancel();
assert!(token.is_cancelled());
sf.cancel(); // no active attempt — must not panic
}
#[test]
fn stale_end_does_not_clear_newer_attempt() {
let sf = AuthSingleFlight::default();
let (_first, first_guard) = sf.begin(None, None);
let first_gen = first_guard.generation();
// Keep first_guard alive but end via generation (stale after second begin).
let (second, _second_guard) = sf.begin(None, None);
sf.end(first_gen); // stale finisher
sf.cancel(); // must still cancel the second attempt's token
assert!(
second.is_cancelled(),
"stale end() must not have cleared the active token"
);
}
#[test]
fn current_end_drops_the_stored_attempt() {
let sf = AuthSingleFlight::default();
let (token, guard) = sf.begin(None, None);
guard.end();
sf.cancel(); // nothing active — must not cancel the finished attempt
assert!(!token.is_cancelled());
}
/// The race the attempt object exists to prevent: a cancelled
/// predecessor finishing late must not drop the successor's channels.
#[test]
fn stale_end_leaves_successor_channels_intact() {
let sf = AuthSingleFlight::default();
let (_first, first_guard) = sf.begin(None, None);
let first_gen = first_guard.generation();
let (ch, mut code_rx) = channels();
let (_second, _g2) = sf.begin(Some(ch), Some(2));
sf.end(first_gen); // stale finisher (attempt #1's cleanup)
sf.submit_code("1234".into())
.expect("successor's code channel must still be wired");
assert_eq!(code_rx.try_recv().as_deref(), Ok("1234"));
assert!(
sf.take_url_rx().is_some(),
"successor's url receiver must still be present"
);
}
#[test]
fn submit_code_and_url_rx_absent_when_idle_or_headless() {
let sf = AuthSingleFlight::default();
assert!(
matches!(
sf.submit_code("x".into()),
Err(SubmitCodeError::NoPendingAttempt)
),
"idle: no attempt is waiting for a code"
);
assert!(sf.take_url_rx().is_none());
let _g = sf.begin(None, None); // headless attempt: token only
assert!(
matches!(
sf.submit_code("x".into()),
Err(SubmitCodeError::NoPendingAttempt)
),
"headless: no channels"
);
assert!(sf.take_url_rx().is_none());
}
#[test]
fn cancel_for_client_seq_ignores_stale_seq() {
let sf = AuthSingleFlight::default();
let (first, _g1) = sf.begin(None, Some(1));
let (second, _g2) = sf.begin(None, Some(2));
assert!(first.is_cancelled());
sf.cancel_for_client_seq(1); // delayed cancel for attempt 1
assert!(
!second.is_cancelled(),
"stale cancel must not tear down the successor"
);
sf.cancel_for_client_seq(2);
assert!(second.is_cancelled());
}
#[test]
fn attempt_guard_ends_on_drop() {
let sf = AuthSingleFlight::default();
let (token, guard) = sf.begin(None, Some(7));
drop(guard);
sf.cancel(); // nothing active
assert!(!token.is_cancelled());
assert!(matches!(
sf.submit_code("x".into()),
Err(SubmitCodeError::NoPendingAttempt)
));
}
/// Headless (and interactive) authenticate `select!`s on this token —
/// cancel must interrupt a long wait rather than leaving it racing
/// (logout / unscoped cancel path).
#[tokio::test]
async fn cancel_interrupts_waiting_select() {
let sf = AuthSingleFlight::default();
let (cancel, _guard) = sf.begin(None, Some(42)); // headless: no channels
let waiter = tokio::spawn(async move {
tokio::select! {
biased;
_ = cancel.cancelled() => "cancelled",
_ = tokio::time::sleep(std::time::Duration::from_secs(30)) => "timeout",
}
});
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
sf.cancel(); // same as handle_logout / unscoped cancel
assert_eq!(waiter.await.expect("join"), "cancelled");
}
/// Logout-style unscoped cancel, then a new begin must not see prior channels.
#[test]
fn cancel_then_begin_is_clean_for_successor() {
let sf = AuthSingleFlight::default();
let (ch, mut code_rx) = channels();
let (old, _g) = sf.begin(Some(ch), Some(1));
sf.cancel();
assert!(old.is_cancelled());
let (ch2, mut code_rx2) = channels();
let (new, _g2) = sf.begin(Some(ch2), Some(2));
assert!(!new.is_cancelled());
sf.submit_code("ok".into()).expect("successor wired");
assert_eq!(code_rx2.try_recv().as_deref(), Ok("ok"));
assert!(code_rx.try_recv().is_err(), "prior channel must be dead");
}
}

80
crates/codegen/xai-grok-shell/src/auth/storage.rs
View file

@ -52,6 +52,16 @@ pub fn read_auth_json(auth_file: &Path) -> std::io::Result<AuthStore> {
let mut contents = String::new();
file.read_to_string(&mut contents)?;
// Tighten world-readable copies (hand-restored, umask edge cases, etc.).
// Best-effort: a chmod failure must not block login/read paths.
if let Err(e) = crate::util::secure_file::ensure_owner_only_permissions(auth_file) {
tracing::warn!(
path = %auth_file.display(),
error = %e,
"auth: failed to enforce owner-only permissions on auth.json"
);
}
// Empty files are valid (recover from prior crash/partial write).
let trimmed = contents.trim();
if trimmed.is_empty() {
@ -114,6 +124,8 @@ pub(crate) fn backup_corrupt_auth_file(path: &Path) -> Option<PathBuf> {
match std::fs::rename(path, &backup) {
Ok(()) => {
// Corrupt backups still hold token material — keep them owner-only.
let _ = crate::util::secure_file::ensure_owner_only_permissions(&backup);
tracing::warn!(
original = %path.display(),
backup = %backup.display(),
@ -246,9 +258,17 @@ fn write_store_to(path: &Path, auth_store: &AuthStore) -> std::io::Result<()> {
.into_inner()
.map_err(|e| e.into_error())?
.sync_all()?;
#[cfg(windows)]
{
crate::util::secure_file::set_windows_secure_permissions(path)?;
// `open_secure_file` mode bits apply only on create; tighten existing paths.
// Best-effort after durable content: a chmod-only failure must not look
// like a failed write. The in-place fallback restores the prior snapshot
// on any `write_store_to` Err, which would discard freshly written tokens.
// Load path re-tightens on next read.
if let Err(e) = crate::util::secure_file::ensure_owner_only_permissions(path) {
tracing::warn!(
error = %e,
path = %path.display(),
"auth: failed to ensure owner-only permissions after write"
);
}
Ok(())
}
@ -263,6 +283,15 @@ fn write_auth_json_atomic(auth_file: &Path, auth_store: &AuthStore) -> std::io::
let _ = std::fs::remove_file(auth_file);
}
std::fs::rename(&tmp, auth_file)?;
// Re-assert on the final path (covers rename edge cases / FS quirks).
// Best-effort: rename already published the new tokens.
if let Err(e) = crate::util::secure_file::ensure_owner_only_permissions(auth_file) {
tracing::warn!(
error = %e,
path = %auth_file.display(),
"auth: failed to ensure owner-only permissions after rename"
);
}
Ok(())
}
@ -317,10 +346,7 @@ fn restore_prior_bytes(auth_file: &Path, bytes: &[u8]) -> std::io::Result<()> {
let mut file = open_secure_file(auth_file)?;
file.write_all(bytes)?;
file.sync_all()?;
#[cfg(windows)]
{
crate::util::secure_file::set_windows_secure_permissions(auth_file)?;
}
crate::util::secure_file::ensure_owner_only_permissions(auth_file)?;
Ok(())
}
@ -433,6 +459,46 @@ mod write_fallback_tests {
assert_eq!(mode & 0o777, 0o600, "in-place write must stay 0o600");
}
#[cfg(unix)]
#[test]
fn write_tightens_preexisting_world_readable_auth_json() {
use std::os::unix::fs::PermissionsExt;
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("auth.json");
std::fs::write(&path, b"{}").unwrap();
let mut loose = std::fs::metadata(&path).unwrap().permissions();
loose.set_mode(0o644);
std::fs::set_permissions(&path, loose).unwrap();
write_auth_json(&path, &sample_store()).unwrap();
let mode = std::fs::metadata(&path).unwrap().permissions().mode();
assert_eq!(
mode & 0o777,
0o600,
"rewrite must tighten preexisting open perms"
);
}
#[cfg(unix)]
#[test]
fn read_tightens_world_readable_auth_json() {
use std::os::unix::fs::PermissionsExt;
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("auth.json");
write_auth_json(&path, &sample_store()).unwrap();
let mut loose = std::fs::metadata(&path).unwrap().permissions();
loose.set_mode(0o644);
std::fs::set_permissions(&path, loose).unwrap();
let _ = read_auth_json(&path).unwrap();
let mode = std::fs::metadata(&path).unwrap().permissions().mode();
assert_eq!(
mode & 0o777,
0o600,
"load must tighten open auth.json perms"
);
}
/// A `StorageFull` (ENOSPC) failure on the atomic path must fall back to
/// the in-place write so the credential still lands on disk.
#[test]

11
crates/codegen/xai-grok-shell/src/config/mod.rs
View file

@ -818,12 +818,11 @@ pub use xai_grok_config::ConfigLayers;
pub use xai_grok_config::{
MDM_REQUIREMENTS_SOURCE, RequirementsLayer, RequirementsSource, ServingIdentity, SyncMarker,
claude_managed_settings_probe_path, confirmed_team_switch, confirmed_team_switch_at,
fail_closed_flag_from_str, is_managed_config_hard_stale_for, is_managed_config_stale_for,
load_config_file, load_from_disk, load_managed_config, load_merged_requirements,
load_system_managed_config, load_toml_file, managed_config_identity_changed_at,
managed_deployment_id, managed_policy_compromised_for, mark_managed_config_synced,
mark_managed_config_synced_at, normalize_identity, requirements_layers, system_config_dir,
user_grok_home,
is_managed_config_hard_stale_for, is_managed_config_stale_for, load_config_file,
load_from_disk, load_managed_config, load_merged_requirements, load_system_managed_config,
load_toml_file, managed_config_identity_changed_at, managed_deployment_id,
managed_policy_compromised_for, mark_managed_config_synced, mark_managed_config_synced_at,
normalize_identity, requirements_layers, system_config_dir, user_grok_home,
};
/// Map of "dotted.path" to which config file the value came from.
pub fn config_origins(

42
crates/codegen/xai-grok-shell/src/extensions/auth.rs
View file

@ -19,6 +19,7 @@ pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
"x.ai/setApiKey" => handle_set_api_key(args),
"x.ai/auth/submit_code" => handle_submit_code(agent, args),
"x.ai/auth/get_url" => handle_get_url(agent).await,
"x.ai/auth/cancel" => handle_cancel(agent, args),
"x.ai/auth/logout" => handle_logout(agent, args).await,
"x.ai/auth/info" => handle_info(agent),
"x.ai/auth/check_subscription" => handle_check_subscription(agent).await,
@ -26,6 +27,26 @@ pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
}
}
/// Stop an in-flight interactive login (device poll / loopback wait).
/// Idempotent: no-op when nothing is waiting.
///
/// When `request_seq` is present, only that attempt is cancelled — a delayed
/// cancel cannot tear down a successor login that already replaced it.
fn handle_cancel(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
#[derive(Deserialize)]
struct CancelParams {
#[serde(default)]
request_seq: Option<u64>,
}
let params: CancelParams =
serde_json::from_str(args.params.get()).unwrap_or(CancelParams { request_seq: None });
match params.request_seq {
Some(seq) => agent.interactive_auth.cancel_for_client_seq(seq),
None => agent.interactive_auth.cancel(),
}
to_raw_response(&serde_json::json!({ "cancelled": true }))
}
async fn handle_get_bearer_token(agent: &MvpAgent) -> ExtResult {
let token = match agent.auth_manager.get_valid_token().await {
Ok(token) => Some(token),
@ -85,20 +106,20 @@ fn handle_submit_code(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
let params: SubmitCodeParams = serde_json::from_str(args.params.get())
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
let auth_code_tx = agent.auth_code_tx.borrow();
if let Some(ref tx) = *auth_code_tx {
tx.try_send(params.code).map_err(|e| {
acp::Error::internal_error().data(format!("failed to submit auth code: {e}"))
})?;
to_raw_response(&serde_json::json!({ "submitted": true }))
} else {
Err(acp::Error::invalid_params().data("no pending auth session"))
match agent.interactive_auth.submit_code(params.code) {
Ok(()) => to_raw_response(&serde_json::json!({ "submitted": true })),
Err(crate::auth::single_flight::SubmitCodeError::SendFailed(e)) => {
Err(acp::Error::internal_error().data(format!("failed to submit auth code: {e}")))
}
Err(crate::auth::single_flight::SubmitCodeError::NoPendingAttempt) => {
Err(acp::Error::invalid_params().data("no pending auth session"))
}
}
}
/// Awaits the auth URL from the oneshot channel (blocks until ready).
async fn handle_get_url(agent: &MvpAgent) -> ExtResult {
let rx = agent.auth_url_rx.borrow_mut().take();
let rx = agent.interactive_auth.take_url_rx();
// `None` when no URL was sent (cached creds, early error, second poll):
// report mode as `null` rather than mislabeling it `loopback`.
let (auth_url, mode) = match rx {
@ -125,6 +146,9 @@ async fn handle_logout(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
let params: LogoutParams = serde_json::from_str(args.params.get())
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
// Stop any in-flight login so it cannot write credentials back after logout.
agent.interactive_auth.cancel();
let result = crate::auth::perform_logout(&agent.auth_manager, params.scope.as_deref())
.map_err(|e| acp::Error::internal_error().data(format!("failed to logout: {e}")))?;
// `auth.lifecycle` (not `auth`) avoids colliding with the pre-existing

4
crates/codegen/xai-grok-shell/src/extensions/marketplace.rs
View file

@ -288,11 +288,13 @@ async fn handle_update(
plugin_subdir: plugin_relative_path.to_string(),
};
let mut registry = xai_grok_agent::plugins::install_registry::InstallRegistry::load();
let require_sha = crate::plugin::marketplace_require_sha();
let update_result = installer::update_from_marketplace_entry_transactional(
&marketplace_root,
&entry,
provenance,
&mut registry,
require_sha,
);
drop(marketplace_lease);
@ -398,6 +400,7 @@ async fn handle_install(
plugin_subdir: plugin_relative_path.to_string(),
};
let mut registry = xai_grok_agent::plugins::install_registry::InstallRegistry::load();
let require_sha = crate::plugin::marketplace_require_sha();
match installer::install_from_remote_url(
&remote_url,
remote_ref.as_deref(),
@ -406,6 +409,7 @@ async fn handle_install(
plugin_relative_path,
provenance,
&mut registry,
require_sha,
) {
Ok(installer::MarketplaceInstallResult::Installed { repo_key }) => {
// Auto-enable installed plugin so it's active after reload.

312
crates/codegen/xai-grok-shell/src/extensions/mcp.rs
View file

@ -38,6 +38,7 @@ pub mod mcp_methods {
pub const READ_RESOURCE: &str = "x.ai/mcp/read_resource";
pub const AUTH_STATUS: &str = "x.ai/mcp/auth_status";
pub const AUTH_TRIGGER: &str = "x.ai/mcp/auth_trigger";
pub const SETUP: &str = "x.ai/mcp/setup";
pub const TOGGLE: &str = "x.ai/mcp/toggle";
pub const TOGGLE_TOOL: &str = "x.ai/mcp/toggle_tool";
pub const UPSERT: &str = "x.ai/mcp/upsert";
@ -83,6 +84,10 @@ pub struct McpServerEntry {
pub source: McpServerSource,
#[serde(skip_serializing_if = "Option::is_none")]
pub source_label: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub setup: Option<crate::util::config::McpSetupConfig>,
#[serde(skip_serializing_if = "Option::is_none")]
pub setup_values: Option<HashMap<String, String>>,
#[serde(flatten)]
pub config: McpServerConfig,
#[serde(skip_serializing_if = "Option::is_none")]
@ -142,6 +147,8 @@ pub struct McpServerSessionState {
pub tools: Vec<McpToolEntry>,
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
pub auth_required: bool,
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
pub setup_required: bool,
}
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
@ -149,6 +156,7 @@ pub struct McpServerSessionState {
pub enum McpSessionStatus {
Ready,
Initializing,
SetupRequired,
Unavailable,
}
@ -339,6 +347,7 @@ enum McpRoute {
ReadResource,
AuthStatus,
AuthTrigger,
Setup,
Toggle,
ToggleTool,
Upsert,
@ -352,6 +361,7 @@ fn route_mcp_method(method: &str) -> Option<McpRoute> {
mcp_methods::READ_RESOURCE => McpRoute::ReadResource,
mcp_methods::AUTH_STATUS => McpRoute::AuthStatus,
mcp_methods::AUTH_TRIGGER => McpRoute::AuthTrigger,
mcp_methods::SETUP => McpRoute::Setup,
mcp_methods::TOGGLE => McpRoute::Toggle,
mcp_methods::TOGGLE_TOOL => McpRoute::ToggleTool,
mcp_methods::UPSERT => McpRoute::Upsert,
@ -368,6 +378,7 @@ pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
Some(McpRoute::ReadResource) => handle_read_resource(agent, args).await,
Some(McpRoute::AuthStatus) => handle_auth_status(agent, args).await,
Some(McpRoute::AuthTrigger) => handle_auth_trigger(agent, args).await,
Some(McpRoute::Setup) => handle_setup(agent, args).await,
Some(McpRoute::Toggle) => handle_toggle(agent, args).await,
Some(McpRoute::ToggleTool) => handle_toggle_tool(agent, args).await,
Some(McpRoute::Upsert) => handle_upsert(agent, args).await,
@ -423,6 +434,8 @@ pub fn build_mcp_catalog_with_gateway_tools(
scope_name: config.scope_name.clone(),
},
source_label: None,
setup: None,
setup_values: None,
session: None,
});
}
@ -459,6 +472,8 @@ pub fn build_mcp_catalog_with_gateway_tools(
source: McpServerSource::Managed,
config: McpServerConfig::ManagedGateway,
source_label: None,
setup: None,
setup_values: None,
session: Some(McpServerSessionState {
enabled: !server_disabled,
status: (!auth_required && !server_disabled).then_some(McpSessionStatus::Ready),
@ -476,6 +491,7 @@ pub fn build_mcp_catalog_with_gateway_tools(
})
.collect(),
auth_required,
setup_required: false,
}),
});
}
@ -520,6 +536,8 @@ pub fn build_mcp_catalog_with_gateway_tools(
source,
config,
source_label: None,
setup: None,
setup_values: None,
session: None,
});
}
@ -572,12 +590,15 @@ fn disabled_server_placeholder_entry(name: &str) -> McpServerEntry {
.map(str::to_owned),
source,
source_label: None,
setup: None,
setup_values: None,
config,
session: Some(McpServerSessionState {
enabled: false,
status: None,
tools: vec![],
auth_required: false,
setup_required: false,
}),
}
}
@ -939,8 +960,9 @@ async fn handle_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
session_state_fut
);
let local_servers =
crate::util::config::load_mcp_servers(&cwd, &agent.cfg.borrow().compat_resolved);
let compat = agent.cfg.borrow().compat_resolved;
let plugin_registry_snapshot = agent.plugin_registry_snapshot();
let local_servers = crate::util::config::load_mcp_servers(&cwd, &compat);
let disabled_tools = crate::util::config::get_all_mcp_disabled_tools(&cwd);
let mut servers = build_mcp_catalog_with_gateway_tools(
&managed_configs,
@ -948,10 +970,65 @@ async fn handle_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
gateway_catalog.as_ref(),
&disabled_tools,
);
let disabled_names = crate::util::config::disabled_mcp_server_names(&cwd);
let setup_entries = crate::util::config::collect_mcp_setup_configs(
&cwd,
plugin_registry_snapshot.as_deref(),
&compat,
);
let preferences = crate::util::config::load_mcp_preferences().file();
for (name, setup_entry) in setup_entries {
if servers.iter().any(|entry| entry.name == name) {
continue;
}
let enabled = !disabled_names.contains(&name);
let setup_schema = setup_entry.config.setup.clone();
let (setup, setup_required, status) = match setup_entry
.config
.resolve_setup(preferences.servers.get(&name))
{
crate::util::config::McpSetupResolution::Required(setup) => {
(Some(setup), true, Some(McpSessionStatus::SetupRequired))
}
// Surface schema/template breakage instead of dropping the row.
crate::util::config::McpSetupResolution::Invalid(_) => {
(setup_schema, true, Some(McpSessionStatus::SetupRequired))
}
crate::util::config::McpSetupResolution::Resolved(_) => continue,
};
let values = preferences
.servers
.get(&name)
.map(|prefs| prefs.values.clone());
servers.push(McpServerEntry {
name: name.clone(),
display_name: None,
source: McpServerSource::Local,
source_label: setup_entry
.source
.plugin
.as_ref()
.map(|plugin| format!("plugin: {plugin}")),
setup,
setup_values: values,
config: McpServerConfig::Http {
url: String::new(),
scope: None,
scope_id: None,
scope_name: None,
},
session: Some(McpServerSessionState {
enabled,
status,
tools: vec![],
auth_required: false,
setup_required,
}),
});
}
// Include disabled servers from config so they appear in the list
// with enabled=false and can be re-enabled by the user.
let disabled_names = crate::util::config::disabled_mcp_server_names(&cwd);
let catalog_names: std::collections::HashSet<String> =
servers.iter().map(|s| s.name.clone()).collect();
for name in &disabled_names {
@ -999,6 +1076,13 @@ async fn handle_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
// Annotate catalog entries with session state.
for entry in &mut servers {
if entry
.session
.as_ref()
.is_some_and(|session| session.setup_required)
{
continue;
}
let managed_gateway_session = entry.source == McpServerSource::Managed
&& matches!(&entry.config, McpServerConfig::ManagedGateway);
if managed_gateway_session {
@ -1027,6 +1111,7 @@ async fn handle_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
status,
tools,
auth_required: snapshot.auth_required.contains(&entry.name),
setup_required: false,
});
}
@ -1038,6 +1123,8 @@ async fn handle_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
display_name: None,
source: McpServerSource::Local,
source_label: None,
setup: None,
setup_values: None,
config: McpServerConfig::Stdio {
command: std::path::PathBuf::new(),
args: Vec::new(),
@ -1048,6 +1135,7 @@ async fn handle_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
status: Some(client_status.status.clone()),
tools: client_status.tools.clone(),
auth_required: snapshot.auth_required.contains(&client_status.name),
setup_required: false,
}),
});
}
@ -1056,9 +1144,11 @@ async fn handle_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
// Tag servers with the owning plugin (covers both a plugin's .mcp.json and
// its inline plugin.json mcpServers via the registry's deduped owner map).
if let Some(registry) = agent.plugin_registry_snapshot() {
if let Some(registry) = plugin_registry_snapshot.as_ref() {
for entry in &mut servers {
if let Some(plugin_name) = registry.mcp_server_owner(&entry.name) {
if entry.source_label.is_none()
&& let Some(plugin_name) = registry.mcp_server_owner(&entry.name)
{
entry.source_label = Some(format!("plugin: {plugin_name}"));
}
}
@ -1397,6 +1487,8 @@ struct McpAuthTriggerRequest {
#[derive(serde::Serialize)]
struct McpAuthTriggerResponse {
status: &'static str,
#[serde(skip_serializing_if = "Option::is_none")]
setup: Option<crate::util::config::McpSetupConfig>,
/// Descriptive failure reason from the shell. `None` on success and on
/// failures with no detail; surfaced verbatim by the TUI.
#[serde(skip_serializing_if = "Option::is_none")]
@ -1409,21 +1501,174 @@ async fn handle_auth_trigger(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtRes
let handle = agent
.get_session_handle(&acp_id)
.ok_or_else(|| acp::Error::invalid_params().data("session not found"))?;
let cwd = agent
.get_session_cwd(&acp_id)
.unwrap_or_else(|| std::env::current_dir().unwrap_or_default());
let setup_entries = crate::util::config::collect_mcp_setup_configs(
&cwd,
agent.plugin_registry_snapshot().as_deref(),
&agent.cfg.borrow().compat_resolved,
);
let preferences = crate::util::config::load_mcp_preferences().file();
if let Some(entry) = setup_entries.get(&req.server_name) {
match entry
.config
.resolve_setup(preferences.servers.get(&req.server_name))
{
crate::util::config::McpSetupResolution::Required(setup) => {
return to_ext_response(Ok(McpAuthTriggerResponse {
status: "setup_required",
setup: Some(setup),
error: None,
}));
}
crate::util::config::McpSetupResolution::Invalid(reason) => {
return to_ext_response(Ok(McpAuthTriggerResponse {
status: "setup_required",
setup: entry.config.setup.clone(),
error: Some(reason),
}));
}
crate::util::config::McpSetupResolution::Resolved(_) => {}
}
}
match handle.mcp_auth_trigger(req.server_name).await {
Ok(()) => to_ext_response(Ok(McpAuthTriggerResponse {
status: "authenticated",
setup: None,
error: None,
})),
Err(e) => {
tracing::warn!(%e, "MCP auth trigger failed");
to_ext_response(Ok(McpAuthTriggerResponse {
status: "failed",
setup: None,
error: Some(e),
}))
}
}
}
#[derive(serde::Deserialize)]
#[serde(rename_all = "camelCase")]
struct McpSetupRequest {
session_id: String,
server_name: String,
values: HashMap<String, String>,
}
#[derive(serde::Serialize)]
#[serde(rename_all = "camelCase")]
struct McpSetupResponse {
ok: bool,
}
async fn handle_setup(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
let req = parse_params::<McpSetupRequest>(args)?;
let acp_id = acp::SessionId::new(req.session_id.clone());
let handle = agent
.get_session_handle(&acp_id)
.ok_or_else(|| acp::Error::invalid_params().data("session not found"))?;
let cwd = agent
.get_session_cwd(&acp_id)
.unwrap_or_else(|| std::env::current_dir().unwrap_or_default());
let setup_entries = crate::util::config::collect_mcp_setup_configs(
&cwd,
agent.plugin_registry_snapshot().as_deref(),
&agent.cfg.borrow().compat_resolved,
);
let entry = setup_entries
.get(&req.server_name)
.ok_or_else(|| acp::Error::invalid_params().data("server setup not found"))?;
let setup = entry
.config
.setup
.as_ref()
.ok_or_else(|| acp::Error::invalid_params().data("server setup not found"))?;
// Only schema field ids (never arbitrary client keys).
let filtered_values: HashMap<String, String> = setup
.fields
.iter()
.filter_map(|field| {
req.values
.get(&field.id)
.map(|value| (field.id.clone(), value.clone()))
})
.collect();
let pending_preferences = crate::util::config::McpServerPreferences {
values: filtered_values,
source: Some(entry.source.clone()),
updated_at: Some(chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true)),
};
match entry.config.resolve_setup(Some(&pending_preferences)) {
crate::util::config::McpSetupResolution::Resolved(_) => {}
crate::util::config::McpSetupResolution::Required(_) => {
return Err(acp::Error::invalid_params().data("setup values incomplete"));
}
crate::util::config::McpSetupResolution::Invalid(reason) => {
return Err(acp::Error::invalid_params().data(reason));
}
}
let load = crate::util::config::load_mcp_preferences();
if !load.is_writable() {
return Err(acp::Error::internal_error().data(
"MCP preferences file is unreadable; fix or remove mcp_preferences.json before saving",
));
}
let mut prefs = load.file();
let previous_entry = prefs.servers.get(&req.server_name).cloned();
prefs
.servers
.insert(req.server_name.clone(), pending_preferences);
crate::util::config::save_mcp_preferences(&prefs)
.await
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
let rollback = || async {
let _ = crate::util::config::restore_mcp_preference_server(
&req.server_name,
previous_entry.clone(),
)
.await;
};
let managed_configs = agent.get_managed_mcp_configs().await;
let all_servers_with_policy =
crate::session::managed_mcp::merge_managed_mcp_servers_with_policy(
vec![],
&cwd,
&managed_configs,
agent.plugin_registry_snapshot().as_deref(),
&agent.cfg.borrow().compat_resolved,
);
let found = match all_servers_with_policy
.into_iter()
.find(|s| crate::session::mcp_servers::mcp_server_name(&s.server) == req.server_name)
{
Some(found) => found,
None => {
rollback().await;
return Err(acp::Error::internal_error().data("server did not resolve after setup"));
}
};
if let Some(reason) = found.disabled_reason {
rollback().await;
return Err(acp::Error::invalid_params().data(reason.to_string()));
}
if let Err(e) = handle
.toggle_mcp_server(req.server_name.clone(), true, Some(found.server))
.await
{
rollback().await;
return Err(acp::Error::internal_error().data(e.to_string()));
}
to_ext_response(Ok(McpSetupResponse { ok: true }))
}
// ── mcp/toggle handler ───────────────────────────────────────────────
#[derive(serde::Deserialize)]
@ -1834,6 +2079,8 @@ mod tests {
scope_name: Some("Grok CLI".to_string()),
},
source_label: None,
setup: None,
setup_values: None,
session: None,
},
McpServerEntry {
@ -1841,6 +2088,8 @@ mod tests {
display_name: None,
source: McpServerSource::Local,
source_label: None,
setup: None,
setup_values: None,
config: McpServerConfig::Stdio {
command: "/usr/bin/mcp-filesystem".into(),
args: vec!["--root".to_string(), "/home".to_string()],
@ -1850,6 +2099,7 @@ mod tests {
enabled: true,
status: Some(McpSessionStatus::Ready),
auth_required: false,
setup_required: false,
tools: vec![McpToolEntry {
name: "read_file".to_string(),
display_name: None,
@ -1876,12 +2126,15 @@ mod tests {
display_name: Some("linear".to_string()),
source: McpServerSource::Managed,
source_label: None,
setup: None,
setup_values: None,
config: McpServerConfig::ManagedGateway,
session: Some(McpServerSessionState {
enabled: true,
status: Some(McpSessionStatus::Ready),
tools: vec![],
auth_required: false,
setup_required: false,
}),
})
.unwrap();
@ -2152,10 +2405,54 @@ mod tests {
assert_eq!(json["isError"], false);
}
#[test]
fn test_mcp_list_setup_required_serialization() {
let entry = McpServerEntry {
name: "acme".to_string(),
display_name: None,
source: McpServerSource::Local,
source_label: Some("plugin: acme".to_string()),
setup: Some(crate::util::config::McpSetupConfig {
fields: vec![crate::util::config::McpSetupField {
id: "site".to_string(),
label: "Site".to_string(),
field_type: crate::util::config::McpSetupFieldType::Select,
required: true,
default: Some("us1".to_string()),
options: vec![crate::util::config::McpSetupOption {
label: "US5".to_string(),
value: "us5".to_string(),
}],
}],
variables: HashMap::new(),
}),
setup_values: Some(HashMap::from([("site".to_string(), "us5".to_string())])),
config: McpServerConfig::Http {
url: String::new(),
scope: None,
scope_id: None,
scope_name: None,
},
session: Some(McpServerSessionState {
enabled: true,
status: Some(McpSessionStatus::SetupRequired),
tools: vec![],
auth_required: false,
setup_required: true,
}),
};
let json = serde_json::to_value(&entry).unwrap();
assert_eq!(json["session"]["status"], "setuprequired");
assert_eq!(json["session"]["setupRequired"], true);
assert_eq!(json["setup"]["fields"][0]["id"], "site");
assert_eq!(json["setupValues"]["site"], "us5");
}
#[test]
fn test_mcp_auth_trigger_response_success_no_error_field() {
let resp = McpAuthTriggerResponse {
status: "authenticated",
setup: None,
error: None,
};
let json = serde_json::to_value(&resp).unwrap();
@ -2170,6 +2467,7 @@ mod tests {
fn test_mcp_auth_trigger_response_failure_carries_error() {
let resp = McpAuthTriggerResponse {
status: "failed",
setup: None,
error: Some("MCP server 'linear' does not use OAuth".to_string()),
};
let json = serde_json::to_value(&resp).unwrap();
@ -2184,6 +2482,7 @@ mod tests {
fn test_mcp_auth_trigger_response_failure_omits_error_when_none() {
let resp = McpAuthTriggerResponse {
status: "failed",
setup: None,
error: None,
};
let json = serde_json::to_value(&resp).unwrap();
@ -2198,6 +2497,8 @@ mod tests {
display_name: None,
source: McpServerSource::Managed,
source_label: None,
setup: None,
setup_values: None,
config: McpServerConfig::Http {
url: "https://mcp.slack.com".to_string(),
scope: Some("user".to_string()),
@ -2209,6 +2510,7 @@ mod tests {
status: None,
tools: vec![],
auth_required: false,
setup_required: false,
}),
};
let json = serde_json::to_value(&entry).unwrap();

2
crates/codegen/xai-grok-shell/src/extensions/session_admin.rs
View file

@ -586,6 +586,7 @@ fn handle_reload_models(agent: &MvpAgent) -> ExtResult {
let merged_config = agent.cfg.borrow().clone();
agent.models_manager.apply_config(merged_config);
agent.sync_process_static_api_key(None);
let count = agent.models_manager.models().len();
tracing::info!(count, "model list reloaded from config.toml");
@ -608,6 +609,7 @@ fn handle_reload_models(agent: &MvpAgent) -> ExtResult {
/// rather than rebuilding the catalog and notifying clients mid-flight.
fn handle_reload_models_cache(agent: &MvpAgent) -> ExtResult {
agent.models_manager.reload_from_disk_cache();
agent.sync_process_static_api_key(None);
ExtMethodResult::success(serde_json::json!({ "reloaded": true }))
.to_ext_response()
.map_err(|e| acp::Error::internal_error().data(e.to_string()))

82
crates/codegen/xai-grok-shell/src/managed_config.rs
View file

@ -9,10 +9,11 @@ use response::{ApplyOutcome, ManagedConfigResponse, ManagedConfigSource, verify_
/// Server-synced policy artifacts. Excludes the sync marker ([`remove_managed_config_files`]
/// removes that last, only on full success).
pub const MANAGED_ARTIFACT_FILES: [&str; 3] = [
pub const MANAGED_ARTIFACT_FILES: [&str; 4] = [
xai_grok_config::MANAGED_CONFIG_FILENAME,
xai_grok_config::REQUIREMENTS_FILENAME,
xai_grok_config::signed_policy::SIGNATURE_SIDECAR_FILE,
xai_grok_config::signed_policy::MANAGED_IDENTITY_SIDECAR_FILE,
];
/// Delete server-synced files then the marker (never `config.toml`).
@ -37,6 +38,10 @@ fn remove_managed_config_files(home: &std::path::Path) {
"{}.",
xai_grok_config::signed_policy::SIGNATURE_SIDECAR_FILE
),
format!(
"{}.",
xai_grok_config::signed_policy::MANAGED_IDENTITY_SIDECAR_FILE
),
];
if let Ok(entries) = std::fs::read_dir(home) {
for entry in entries.flatten() {
@ -374,7 +379,7 @@ fn managed_config_sync_interval() -> std::time::Duration {
}
/// Periodically sync managed config in the background. Best-effort.
pub(crate) fn spawn_sync(cancel: tokio_util::sync::CancellationToken) {
pub fn spawn_sync(cancel: tokio_util::sync::CancellationToken) {
tokio::spawn(async move {
let mut interval = tokio::time::interval(managed_config_sync_interval());
interval.tick().await; // skip immediate first tick
@ -388,6 +393,10 @@ pub(crate) fn spawn_sync(cancel: tokio_util::sync::CancellationToken) {
// Clear a logged-out team's files before deciding to fetch, so
// stale enforced policy never outlives the tick.
clear_orphan();
// Raise the floor each tick so a long offline session keeps recording
// observed time; otherwise a later rollback could make an expired policy
// read valid.
bump_managed_rollback_floor();
if !crate::config::is_managed_config_stale_for(&current_serving_identity())
|| !is_fetch_enabled()
@ -652,6 +661,15 @@ fn apply_fetched(
if let Some(verified) = verified {
clear_squatting_dir(&home.join(xai_grok_config::signed_policy::SIGNATURE_SIDECAR_FILE));
xai_grok_config::signed_policy::write_sidecar(&home, &verified.sidecar)?;
// Disk errors are fatal, like the policy sidecar's.
if let Some(claim_sidecar) =
verified_claim_sidecar(body, served_principal_of(&verified.payload))
{
clear_squatting_dir(
&home.join(xai_grok_config::signed_policy::MANAGED_IDENTITY_SIDECAR_FILE),
);
xai_grok_config::signed_policy::write_managed_identity_sidecar(&home, &claim_sidecar)?;
}
}
// Marker last, still under the lock: written post-release, a concurrent purge could
// delete the files it describes. A squatting dir would fail the atomic rename forever.
@ -675,8 +693,48 @@ fn apply_fetched(
Ok(ApplyOutcome::Applied { wrote })
}
/// The principal a verified payload binds: `deployment_id`, else `team_id` (server parity).
fn served_principal_of(payload: &xai_grok_config::signed_policy::SignedPayload) -> Option<&str> {
payload
.deployment_id
.as_deref()
.or(payload.team_id.as_deref())
}
/// The fetched claim envelope, if it verifies and binds to the served principal.
/// `None` skips (old server / unverifiable / foreign): a bad claim must not fail
/// the apply — it only hardens the policy sidecar.
fn verified_claim_sidecar(
body: &ManagedConfigResponse,
served_principal: Option<&str>,
) -> Option<xai_grok_config::signed_policy::SignatureEnvelope> {
use xai_grok_config::signed_policy::now_unix;
let sidecar = body.managed_identity_sidecar()?;
// Unclamped wall clock, like the policy verify: a fresh claim heals an inflated floor.
let claim = match xai_grok_config::signed_policy::verify_fetched_claim(&sidecar, now_unix()) {
Ok(claim) => claim,
Err(e) => {
tracing::debug!("is-managed claim did not verify; not persisting it: {e}");
return None;
}
};
if !claim_binds_to(&claim, served_principal) {
tracing::debug!("is-managed claim is bound to a different principal; not persisting it");
return None;
}
Some(sidecar)
}
/// The persist rule: a verified claim persists only when bound to the served principal.
fn claim_binds_to(
claim: &xai_grok_config::signed_policy::ManagedIdentityClaim,
served_principal: Option<&str>,
) -> bool {
served_principal == Some(claim.principal.as_str())
}
/// Evict the prior principal's policy artifacts on a confirmed switch; this apply then
/// writes the new set and rebinds the marker. Includes the sidecar — a verification-inactive
/// writes the new set and rebinds the marker. Includes the sidecars — a verification-inactive
/// build must not leave the prior tenant's sidecar to read foreign-bound on a signing build.
fn evict_prior_managed_config(home: &std::path::Path) {
for name in MANAGED_ARTIFACT_FILES {
@ -880,6 +938,8 @@ pub fn managed_policy_gate() -> Result<(), String> {
}
// Purge first so an offline team switch isn't misread as a substituted cache.
purge_prior_tenant_on_identity_change();
// Raise the floor after the purge so a purged marker stays absent.
bump_managed_rollback_floor();
managed_policy_gate_decision(
managed_principal_present(),
// Expiry-ignoring: a backdated auth.json must not resolve Team→None and relax binding.
@ -919,6 +979,22 @@ fn purge_prior_tenant_on_identity_change() {
}
}
/// Floor tick (session start + background sync tick), best-effort under the
/// managed-config lock — a failed tick must not refuse a session.
fn bump_managed_rollback_floor() {
// Re-checked inside `bump_rollback_floor`; this early-out skips the lock I/O when dark.
if !xai_grok_config::signed_policy::verification_active() {
return;
}
let home = crate::util::grok_home::grok_home();
match try_lock_managed_config(&home) {
Some(_lock) => {
xai_grok_config::bump_rollback_floor(&home);
}
None => tracing::debug!("managed-config lock contended; skipping the floor tick"),
}
}
/// Pure decision behind [`managed_policy_gate`]: fail closed only when a managed principal is active AND its policy is compromised.
fn managed_policy_gate_decision(
managed_principal_present: bool,

85
crates/codegen/xai-grok-shell/src/managed_config/response.rs
View file

@ -91,6 +91,11 @@ pub(super) struct ManagedConfigResponse {
/// signed payload's policy is the trusted copy when verification is on.
#[serde(default)]
pub(super) signatures: Option<Vec<xai_grok_config::signed_policy::SignatureEnvelope>>,
/// The is-managed claim envelopes (additive; absent from old servers), same
/// rotation shape as `signatures`, persisted as their own sidecar.
#[serde(default)]
pub(super) managed_identity_signatures:
Option<Vec<xai_grok_config::signed_policy::SignatureEnvelope>>,
}
impl ManagedConfigResponse {
@ -105,21 +110,20 @@ impl ManagedConfigResponse {
pub(super) fn signature_sidecar(
&self,
) -> Option<xai_grok_config::signed_policy::SignatureEnvelope> {
self.signature_sidecar_with(xai_grok_config::signed_policy::embedded_key_id_trusted)
pick_trusted_envelope(
self.signatures.as_deref(),
xai_grok_config::signed_policy::embedded_key_id_trusted,
)
}
/// Predicate-injected core of [`Self::signature_sidecar`] so tests can pick
/// without a compiled-in key set.
fn signature_sidecar_with(
/// The claim envelope to verify — same picking rule as [`Self::signature_sidecar`].
pub(super) fn managed_identity_sidecar(
&self,
key_id_trusted: impl Fn(&str) -> bool,
) -> Option<xai_grok_config::signed_policy::SignatureEnvelope> {
let envelopes = self.signatures.as_deref()?;
envelopes
.iter()
.find(|e| key_id_trusted(&e.key_id))
.or_else(|| envelopes.first())
.cloned()
pick_trusted_envelope(
self.managed_identity_signatures.as_deref(),
xai_grok_config::signed_policy::embedded_key_id_trusted,
)
}
/// Non-empty served content, recorded in the marker so staleness can later detect a deleted file.
@ -133,12 +137,23 @@ impl ManagedConfigResponse {
self.requirements.as_deref().is_some_and(|s| !s.is_empty())
}
/// The served opt-in (`fail_closed`), read from the payload not disk, so it's authoritative even when
/// the on-disk apply is skipped under lock contention.
/// Served `fail_closed` from the payload (not disk). Non-bool → warn once, treat as false.
pub(super) fn requirements_fail_closed(&self) -> bool {
self.requirements
.as_deref()
.is_some_and(crate::config::fail_closed_flag_from_str)
let Some(req) = self.requirements.as_deref() else {
return false;
};
use prod_mc_cli_chat_proxy_types::{FailClosedFlag, fail_closed_flag_status};
let status = fail_closed_flag_status(req);
if matches!(status, FailClosedFlag::Invalid) {
static WARN_ONCE: std::sync::Once = std::sync::Once::new();
WARN_ONCE.call_once(|| {
tracing::warn!(
"served requirements fail_closed is present but not a boolean \
(e.g. fail_closed = \"true\"); treating as false - use fail_closed = true"
);
});
}
status.is_enabled()
}
}
@ -166,6 +181,19 @@ impl ApplyOutcome {
}
}
/// Pick the envelope whose (hint-only) key_id is trusted, else the first.
fn pick_trusted_envelope(
envelopes: Option<&[xai_grok_config::signed_policy::SignatureEnvelope]>,
key_id_trusted: impl Fn(&str) -> bool,
) -> Option<xai_grok_config::signed_policy::SignatureEnvelope> {
let envelopes = envelopes?;
envelopes
.iter()
.find(|e| key_id_trusted(&e.key_id))
.or_else(|| envelopes.first())
.cloned()
}
/// A fetched envelope that passed verification: the sidecar to persist, plus its
/// parsed (now-trusted) payload.
pub(super) struct VerifiedEnvelope {
@ -184,6 +212,7 @@ pub(super) fn verify_signed_envelope(
let sidecar = body.signature_sidecar().ok_or_else(|| {
"managed policy is required but the server returned no signature".to_owned()
})?;
// Unclamped wall clock: a fresh envelope must heal an inflated floor, not be refused by it.
let payload = signed_policy::verify_fetched(&sidecar, active_team_id, now_unix())
.map_err(|e| e.to_string())?;
if body.managed_config != payload.managed_config || body.requirements != payload.requirements {
@ -196,39 +225,33 @@ pub(super) fn verify_signed_envelope(
mod tests {
use super::*;
/// Picking: the first trusted-key_id entry wins; no trusted entry → the first entry
/// (picking must not invent absence); no array (old/unsigned server) → None.
/// Picking (shared by the policy and claim carriers): the first trusted-key_id
/// entry wins; no trusted entry → the first entry (picking must not invent
/// absence); no array (old/unsigned server) → None.
#[test]
fn signature_sidecar_picks_trusted_envelope_then_falls_back() {
fn pick_trusted_envelope_prefers_trusted_then_falls_back() {
use xai_grok_config::signed_policy::SignatureEnvelope;
let envelope = |kid: &str| SignatureEnvelope {
signed_payload: format!("payload-{kid}"),
signature: format!("sig-{kid}"),
key_id: kid.to_owned(),
};
let body = ManagedConfigResponse {
signatures: Some(vec![envelope("v1"), envelope("v2")]),
..Default::default()
};
let envelopes = vec![envelope("v1"), envelope("v2")];
// A rotated client trusting only v2 picks the v2 envelope from the array.
let picked = body.signature_sidecar_with(|id| id == "v2").unwrap();
let picked = pick_trusted_envelope(Some(&envelopes), |id| id == "v2").unwrap();
assert_eq!(picked.key_id, "v2");
assert_eq!(picked.signed_payload, "payload-v2");
// Trusting v1 picks the primary entry (first in the array).
let picked = body.signature_sidecar_with(|id| id == "v1").unwrap();
let picked = pick_trusted_envelope(Some(&envelopes), |id| id == "v1").unwrap();
assert_eq!(picked.key_id, "v1");
// No trusted id → the first entry, so verification reports UnknownKeyId.
let picked = body.signature_sidecar_with(|_| false).unwrap();
let picked = pick_trusted_envelope(Some(&envelopes), |_| false).unwrap();
assert_eq!(picked.key_id, "v1");
// Nothing signed at all → None.
assert!(
ManagedConfigResponse::default()
.signature_sidecar_with(|_| true)
.is_none()
);
assert!(pick_trusted_envelope(None, |_| true).is_none());
}
}

52
crates/codegen/xai-grok-shell/src/managed_config/tests.rs
View file

@ -338,7 +338,8 @@ fn marker_is_not_a_managed_artifact() {
[
"managed_config.toml",
"requirements.toml",
"managed_config.sig.json"
"managed_config.sig.json",
"managed_identity.sig.json"
],
"the artifact list is load-bearing for every derived loop; change it deliberately"
);
@ -393,3 +394,52 @@ fn purge_keeps_marker_when_an_artifact_removal_fails() {
"with every artifact removed, the marker goes last"
);
}
// --- The is-managed claim persist rules ---
/// Deployment id wins over team id (server parity).
#[test]
fn served_principal_prefers_deployment_id() {
use xai_grok_config::signed_policy::SignedPayload;
let payload = |dep: Option<&str>, team: Option<&str>| SignedPayload {
typ: xai_grok_config::signed_policy::MANAGED_POLICY_TYP.into(),
version: 1,
deployment_id: dep.map(Into::into),
team_id: team.map(Into::into),
managed_config: None,
requirements: None,
fail_closed: false,
expires_at: 0,
key_id: "v1".into(),
};
assert_eq!(
served_principal_of(&payload(Some("dep-1"), Some("team-007"))),
Some("dep-1")
);
assert_eq!(
served_principal_of(&payload(None, Some("team-007"))),
Some("team-007")
);
assert_eq!(served_principal_of(&payload(None, None)), None);
}
/// A verified claim persists ONLY when bound to the served principal.
#[test]
fn claim_persists_only_when_bound_to_served_principal() {
let claim = |principal: &str| xai_grok_config::signed_policy::ManagedIdentityClaim {
typ: xai_grok_config::signed_policy::MANAGED_IDENTITY_TYP.into(),
principal: principal.into(),
fail_closed: true,
expires_at: 4_000_000_000,
key_id: "v1".into(),
};
assert!(claim_binds_to(&claim("team-007"), Some("team-007")));
assert!(!claim_binds_to(&claim("team-evil"), Some("team-007")));
assert!(!claim_binds_to(&claim("team-007"), None));
}
/// Old server, no claim envelopes: nothing persists, nothing errors.
#[test]
fn absent_claim_is_skipped() {
assert!(verified_claim_sidecar(&ManagedConfigResponse::default(), Some("team-007")).is_none());
}

28
crates/codegen/xai-grok-shell/src/plugin.rs
View file

@ -53,7 +53,8 @@ pub fn install_plugin(source: &str, cwd: &Path) -> Result<InstallOutcome, Instal
let is_local = matches!(install_source, git_install::InstallSource::Local { .. });
let mut registry = InstallRegistry::load();
let result = git_install::install_from_source(&install_source, &registry)?;
let result =
git_install::install_from_source(&install_source, &registry, marketplace_require_sha())?;
let repo = git_install::build_installed_repo(&result, &install_source);
registry.insert(result.repo_key.clone(), repo);
@ -281,11 +282,13 @@ fn update_marketplace_repo(
name: provenance.plugin_subdir.clone(),
})?;
let require_sha = crate::plugin::marketplace_require_sha();
installer::update_from_marketplace_entry_transactional(
&marketplace_root.path,
&entry,
provenance,
registry,
require_sha,
)
}
@ -398,7 +401,7 @@ pub fn update_plugins_by_selector(
},
}
} else {
match git_install::update_repo(repo_key, repo) {
match git_install::update_repo(repo_key, repo, marketplace_require_sha()) {
Ok(UpdateStatus::Updated(result)) if result.changed => {
apply_update_to_registry(&mut registry, repo_key, &result);
RepoUpdateOutcome::Updated {
@ -527,6 +530,7 @@ pub fn classify_install_error(err: &InstallError) -> String {
InstallError::Json { .. } => "json",
InstallError::PluginNotFound { .. } => "not_found",
InstallError::ShaMismatch { .. } => "sha_mismatch",
InstallError::UnpinnedRemoteRefused { .. } => "unpinned_remote_refused",
InstallError::InstallFailed { .. } => "install_failed",
}
.to_string()
@ -699,6 +703,15 @@ fn bullet_list(items: &[String]) -> String {
.join("\n")
}
/// The require-sha pin policy for remote plugin code. Disk-only config + env,
/// both tighten-only: a remote campaign overlay must not be able to relax a
/// local security policy, and an unreadable config falls back to the env knob.
pub fn marketplace_require_sha() -> bool {
xai_grok_config::load_effective_config_disk_only()
.map(|c| xai_grok_plugin_marketplace::load_require_sha(&c))
.unwrap_or_else(|_| xai_grok_plugin_marketplace::env_require_sha())
}
/// Marketplace sources from config.toml + settings JSON, unfiltered.
pub fn load_marketplace_sources() -> Vec<MarketplaceSource> {
let config = crate::config::load_effective_config()
@ -1059,6 +1072,7 @@ fn install_marketplace_entry(
};
let result = if let Some(remote_url) = entry.remote_url.as_deref() {
let require_sha = crate::plugin::marketplace_require_sha();
installer::install_from_remote_url(
remote_url,
entry.remote_ref.as_deref(),
@ -1067,6 +1081,7 @@ fn install_marketplace_entry(
&plugin_subdir,
provenance,
registry,
require_sha,
)
} else {
installer::install_from_marketplace(marketplace_root, &plugin_subdir, provenance, registry)
@ -1429,6 +1444,13 @@ mod tests {
}),
"sha_mismatch"
);
assert_eq!(
classify_install_error(&InstallError::UnpinnedRemoteRefused {
plugin: "p".into(),
url: "u".into()
}),
"unpinned_remote_refused"
);
assert_eq!(
classify_install_error(&InstallError::InstallFailed { detail: "x".into() }),
"install_failed"
@ -1462,7 +1484,7 @@ mod tests {
plugins: HashMap::new(),
marketplace: None,
};
let status = git_install::update_repo("local", &repo).unwrap();
let status = git_install::update_repo("local", &repo, false).unwrap();
assert!(matches!(status, UpdateStatus::LiveLocal));
}

149
crates/codegen/xai-grok-shell/src/sampling/error.rs
View file

@ -12,9 +12,12 @@ use agent_client_protocol as acp;
/// ACP error code for rate-limited requests (HTTP 429).
/// Uses the JSON-RPC implementation-defined server error range (-32000 to -32099).
///
/// Contract: this code must only be set for actual HTTP 429 responses from the
/// sampling client. Clients (desktop, pager) suppress error detail when they
/// see this code and show a user-friendly upgrade message instead.
/// Contract: set only for actual HTTP 429 responses from the sampling client.
/// User-facing text is produced by [`format_rate_limited_user_message`] (free-usage
/// paywall rewrite, else server body, else a generic fallback). Pager/headless
/// use that helper. Desktop may still special-case `stopReason: rate_limit` with
/// its own UI and ignore the body — that is a client choice, not a shell
/// requirement to suppress detail.
pub const RATE_LIMITED_ERROR_CODE: i32 = -32003;
/// OAuth / session rate-limit copy (personal plan upgrade path).
@ -40,6 +43,59 @@ pub fn rate_limited_user_message(is_api_key_auth: bool) -> &'static str {
}
}
/// Well-known free-usage exhaustion code CCP returns on HTTP 429.
/// Matches `prod_util_well_known_errors::SUBSCRIPTION_FREE_USAGE_EXHAUSTED`.
/// sampling-types' `parse_error_bytes` prepends the flat `code` to the
/// flattened message, so this reaches clients embedded in error detail.
pub const FREE_USAGE_EXHAUSTED_ERROR_CODE: &str = "subscription:free-usage-exhausted";
/// User-facing free-usage exhaustion copy (paywall). Deliberately promises no
/// reset duration — the quota window is backend-config-driven.
pub const FREE_USAGE_USER_MESSAGE: &str = "You\u{2019}ve reached your free Grok Build usage limit for now. Get SuperGrok for much higher limits, or try again later: https://grok.com/supergrok?referrer=grok-build";
/// Whether flattened server detail is free-usage-quota exhaustion (paywall),
/// not transient throttling. Sniffs the well-known code embedded by
/// `parse_error_bytes`.
pub fn is_free_usage_exhausted_error(detail: &str) -> bool {
detail.contains(FREE_USAGE_EXHAUSTED_ERROR_CODE)
}
/// User-facing text for an ACP -32003 rate-limit error.
///
/// 1. Free-usage well-known code → [`FREE_USAGE_USER_MESSAGE`] (OAuth product paywall).
/// 2. Non-empty body that pushes a **personal** Grok subscription, when the
/// caller is on **API key** auth → [`RATE_LIMITED_USER_MESSAGE_API_KEY`]
/// (team credits / console rate-limit tiers; not grok.com SuperGrok).
/// 3. Any other non-empty body → shown as-is (capacity, team RPS limits, etc.).
/// 4. Empty body → [`rate_limited_user_message`].
///
/// Pass the real active auth method (`AppView.is_api_key_auth` /
/// `AuthMethodKind::is_api_key`), not env-only key presence.
pub fn format_rate_limited_user_message(
server_detail: Option<&str>,
is_api_key_auth: bool,
) -> String {
if server_detail.is_some_and(is_free_usage_exhausted_error) {
return FREE_USAGE_USER_MESSAGE.to_string();
}
if let Some(detail) = server_detail.map(str::trim).filter(|s| !s.is_empty()) {
if is_api_key_auth && pushes_consumer_subscription_upsell(detail) {
return RATE_LIMITED_USER_MESSAGE_API_KEY.to_string();
}
return detail.to_string();
}
rate_limited_user_message(is_api_key_auth).to_string()
}
/// IC sometimes reuses OAuth free-tier upsell copy on 429s ("upgrade to a Grok
/// subscription" / grok.com/supergrok). That is wrong for API-key / team auth:
/// higher limits come from credits and spend-based rate-limit tiers, not a
/// personal SuperGrok plan.
fn pushes_consumer_subscription_upsell(detail: &str) -> bool {
let d = detail.to_ascii_lowercase();
d.contains("grok.com/supergrok") || d.contains("upgrade to a grok subscription")
}
/// Map a `SamplingError` to an ACP `Error` for client-facing responses.
/// This stays in xai-grok-shell because it depends on `agent_client_protocol::Error`.
pub fn map_sampling_err_to_acp(err: SamplingError) -> acp::Error {
@ -334,6 +390,93 @@ mod tests {
assert!(!RATE_LIMITED_USER_MESSAGE_API_KEY.contains("Upgrade your account"));
}
#[test]
fn format_rate_limited_surfaces_nonempty_server_detail() {
let service = "The service is temporarily at capacity. Please retry your request shortly.";
assert_eq!(
format_rate_limited_user_message(Some(service), false),
service
);
assert_eq!(
format_rate_limited_user_message(Some(service), true),
service
);
// Team console rate-limit copy has no personal SuperGrok upsell — surface as-is.
let team = "resource-exhausted: Too many requests for team abc. See https://console.x.ai/team/default/rate-limits.";
assert_eq!(format_rate_limited_user_message(Some(team), true), team);
assert_eq!(
format_rate_limited_user_message(Some("slow down"), false),
"slow down"
);
}
#[test]
fn format_rate_limited_api_key_rewrites_consumer_subscription_upsell() {
let rpm = "Some resource has been exhausted: You are sending requests too quickly. \
Please slow down, or upgrade to a Grok subscription for higher limits: \
https://grok.com/supergrok";
// OAuth keeps the IC body (personal plan upgrade is correct).
assert_eq!(format_rate_limited_user_message(Some(rpm), false), rpm);
// API key must not push grok.com SuperGrok — team credits / rate-limit tiers.
assert_eq!(
format_rate_limited_user_message(Some(rpm), true),
RATE_LIMITED_USER_MESSAGE_API_KEY
);
assert!(
RATE_LIMITED_USER_MESSAGE_API_KEY
.contains("https://docs.x.ai/developers/rate-limits#rate-limit-tiers")
);
assert!(!RATE_LIMITED_USER_MESSAGE_API_KEY.contains("grok.com/supergrok"));
}
#[test]
fn is_free_usage_exhausted_error_sniffs_well_known_code() {
assert!(is_free_usage_exhausted_error(
"subscription:free-usage-exhausted: You have used all your free usage."
));
assert!(is_free_usage_exhausted_error(
"API error (status 429): subscription:free-usage-exhausted quota hit"
));
assert!(!is_free_usage_exhausted_error("throttled"));
assert!(!is_free_usage_exhausted_error(
"The service is temporarily at capacity."
));
}
#[test]
fn format_rate_limited_free_usage_uses_paywall_copy() {
let detail = "subscription:free-usage-exhausted: You have used all your free usage.";
assert_eq!(
format_rate_limited_user_message(Some(detail), false),
FREE_USAGE_USER_MESSAGE
);
assert_eq!(
format_rate_limited_user_message(Some(detail), true),
FREE_USAGE_USER_MESSAGE
);
}
#[test]
fn format_rate_limited_empty_detail_uses_auth_aware_fallback() {
assert_eq!(
format_rate_limited_user_message(None, false),
RATE_LIMITED_USER_MESSAGE_OAUTH
);
assert_eq!(
format_rate_limited_user_message(Some(""), false),
RATE_LIMITED_USER_MESSAGE_OAUTH
);
assert_eq!(
format_rate_limited_user_message(None, true),
RATE_LIMITED_USER_MESSAGE_API_KEY
);
assert_eq!(
format_rate_limited_user_message(Some(" "), true),
RATE_LIMITED_USER_MESSAGE_API_KEY
);
}
#[test]
fn rate_limit_error_uses_dedicated_code() {
let err = SamplingError::Api {

7
crates/codegen/xai-grok-shell/src/session/acp_session_impl/hooks_plugins.rs
View file

@ -311,6 +311,7 @@ impl SessionActor {
match xai_grok_agent::plugins::git_install::install_from_source(
&install_source,
&registry,
crate::plugin::marketplace_require_sha(),
) {
Ok(result) => {
let repo = xai_grok_agent::plugins::git_install::build_installed_repo(
@ -448,7 +449,11 @@ impl SessionActor {
let mut messages = Vec::new();
let mut any_updated = false;
for (key, repo) in &repos_to_update {
match xai_grok_agent::plugins::git_install::update_repo(key, repo) {
match xai_grok_agent::plugins::git_install::update_repo(
key,
repo,
crate::plugin::marketplace_require_sha(),
) {
Ok(status) => {
use xai_grok_agent::plugins::git_install::UpdateStatus;
match status {

80
crates/codegen/xai-grok-shell/src/session/acp_session_impl/mcp.rs
View file

@ -414,16 +414,10 @@ impl SessionActor {
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
.get_client(server_name)
.cloned()
.ok_or_else(|| format!("MCP server '{}' not found", server_name))?
let client = match self.mcp_state.lock().await.get_client(server_name).cloned() {
Some(c) if c.has_auth() => c,
_ => self.recreate_http_client_with_oauth(server_name).await?,
};
if !client.has_auth() {
return Err(format!("MCP server '{}' does not use OAuth", server_name));
}
if !client.force_reauth(true).await {
return Err(format!(
"Authentication failed for MCP server '{}'",
@ -437,6 +431,7 @@ impl SessionActor {
.map_err(|e| format!("Failed to get tools after auth: {}", e))?;
let mut mcp_state = self.mcp_state.lock().await;
mcp_state.auth_required.remove(server_name);
mcp_state.init_failed.remove(server_name);
let mut ui_tools: std::collections::HashMap<
String,
Vec<crate::extensions::mcp::McpToolEntry>,
@ -455,6 +450,73 @@ impl SessionActor {
);
Ok(())
}
/// Rebuild an HTTP MCP client with Interactive OAuth discovery and swap it
/// into session state. Used when auth is requested for a client that was
/// previously started without an `AuthorizationManager`.
async fn recreate_http_client_with_oauth(
&self,
server_name: &str,
) -> Result<std::sync::Arc<crate::session::mcp_servers::McpClient>, String> {
let (server_config, meta_config, event_tx) = {
let mcp_state = self.mcp_state.lock().await;
let server_config = mcp_state
.configs
.iter()
.find(|c| crate::session::mcp_servers::mcp_server_name(c) == server_name)
.cloned()
.ok_or_else(|| format!("MCP server '{}' not found in config", server_name))?;
match &server_config {
acp::McpServer::Http(_) | acp::McpServer::Sse(_) => {}
_ => {
return Err(format!("MCP server '{}' does not use OAuth", server_name));
}
}
let meta_config = mcp_state.meta_config_map.get(server_name).cloned();
let event_tx = mcp_state.client_event_tx();
(server_config, meta_config, event_tx)
};
let cwd = std::path::Path::new(&self.session_info.cwd);
let session_id = self.session_info.id.0.as_ref();
let (_, oauth_config_map) =
crate::util::config::load_mcp_servers_with_oauth(cwd, &self.rebuild_spec.compat);
let byo_config = oauth_config_map.get(server_name).cloned();
let event_writer = self.events.writer();
let mode = crate::session::mcp_servers::OauthInteractivity::Interactive;
let new_client = crate::session::mcp_servers::start_mcp_server(
server_config,
Some(session_id),
Some(cwd),
meta_config.as_ref(),
byo_config.as_ref(),
&event_writer,
mode,
)
.await
.map_err(|e| format!("Failed to prepare OAuth for '{}': {}", server_name, e))?;
if !new_client.has_auth() {
return Err(format!(
"MCP server '{}' does not support OAuth (discovery found no authorization support)",
server_name
));
}
if let Some(tx) = event_tx {
new_client.set_event_tx(Some(tx));
}
let arc = std::sync::Arc::new(new_client);
{
let mut mcp_state = self.mcp_state.lock().await;
mcp_state
.owned_clients
.insert(server_name.to_string(), arc.clone());
mcp_state.auth_required.insert(server_name.to_string());
mcp_state.init_failed.remove(server_name);
}
tracing::info!(
server = server_name,
"Rebuilt MCP HTTP client with OAuth manager for auth_trigger"
);
Ok(arc)
}
/// Attempt to re-initialize MCP servers stuck in `auth_required`.
///
/// For each server, tries `try_reauth_from_disk` which checks the credential

11
crates/codegen/xai-grok-shell/src/session/acp_session_impl/run_loop.rs
View file

@ -536,12 +536,13 @@ pub(super) async fn run_session(
respond_to } => { session.events
.emit(xai_file_utils::events::Event::McpServerToggled { server_name :
server_name.clone(), enabled, }); let mut mcp_state = session.mcp_state
.lock(). await; let mut configs = mcp_state.configs.clone(); if enabled { let
already_present = configs.iter().any(| c | crate
::session::mcp_servers::mcp_server_name(c) == server_name); if
.lock(). await; let mut configs = mcp_state.configs.clone(); if enabled { if
let Some(config) = server_config { configs.retain(| c | { crate
::session::mcp_servers::mcp_server_name(c) != server_name }); configs
.push(config); } else { let already_present = configs.iter().any(| c | {
crate ::session::mcp_servers::mcp_server_name(c) == server_name }); if
already_present { drop(mcp_state); let _ = respond_to.send(Ok(())); continue;
} if let Some(config) = server_config { configs.push(config); } else {
drop(mcp_state); let _ = respond_to.send(Err(acp::Error::invalid_params()
} drop(mcp_state); let _ = respond_to.send(Err(acp::Error::invalid_params()
.data(format!("server '{}' not found in config", server_name)))); continue; }
} else { configs.retain(| c | crate
::session::mcp_servers::mcp_server_name(c) != server_name); } let diff =

2
crates/codegen/xai-grok-shell/src/session/acp_session_impl/sampler_turn.rs
View file

@ -49,7 +49,7 @@ impl SessionTokenAuthGate {
is_session_based: auth_method_id
.is_some_and(crate::agent::auth_method::is_session_based_method),
model_byok,
endpoint_is_first_party: crate::util::is_first_party_xai_url(base_url),
endpoint_is_first_party: crate::util::is_xai_api_url(base_url),
}
}
fn active(self) -> bool {

135
crates/codegen/xai-grok-shell/src/session/acp_session_impl/spawn.rs
View file

@ -545,41 +545,19 @@ pub(crate) async fn spawn_session_actor(
grep_ugrep,
}
};
let persistent_local_shell = crate::util::config::resolve_persistent_local_shell(
remote_settings
.as_ref()
.and_then(|r| r.persistent_local_shell),
);
let terminal_backend_kind = select_terminal_backend_kind(
startup_hints.is_subagent,
parent_terminal_backend.is_some(),
client_terminal_capable,
tool_context.gateway.is_some(),
persistent_local_shell,
);
let terminal_backend: std::sync::Arc<dyn xai_grok_tools::computer::types::TerminalBackend> =
match terminal_backend_kind {
TerminalBackendKind::ReuseParent => parent_terminal_backend
.expect("ReuseParent is only selected when a parent backend is present"),
TerminalBackendKind::AcpClient => {
std::sync::Arc::new(crate::terminal::AcpTerminalAdapter::new(
tool_context.gateway.clone().unwrap(),
tool_context.session_id.clone().unwrap(),
))
as std::sync::Arc<dyn xai_grok_tools::computer::types::TerminalBackend>
}
TerminalBackendKind::LocalPersistent => std::sync::Arc::new(
LocalTerminalBackend::new_local_with_persistent_shell(resolve_search_shadows()),
),
TerminalBackendKind::LocalNonPersistent => {
std::sync::Arc::new(LocalTerminalBackend::new_local(resolve_search_shadows()))
}
if let Some(parent_tb) = parent_terminal_backend.filter(|_| startup_hints.is_subagent) {
parent_tb
} else if client_terminal_capable && tool_context.gateway.is_some() {
std::sync::Arc::new(crate::terminal::AcpTerminalAdapter::new(
tool_context.gateway.clone().unwrap(),
tool_context.session_id.clone().unwrap(),
)) as std::sync::Arc<dyn xai_grok_tools::computer::types::TerminalBackend>
} else {
let backend: std::sync::Arc<dyn xai_grok_tools::computer::types::TerminalBackend> =
std::sync::Arc::new(LocalTerminalBackend::new_local(resolve_search_shadows()));
backend
};
if terminal_backend_kind == TerminalBackendKind::LocalPersistent {
terminal_backend
.warm_persistent_shell(tool_context.cwd.as_path())
.await;
}
let fs_backend: std::sync::Arc<dyn xai_grok_tools::computer::types::AsyncFileSystem> =
if client_fs_capable && tool_context.gateway.is_some() {
std::sync::Arc::new(xai_grok_workspace::file_system::AcpFsAdapter::new(
@ -712,6 +690,11 @@ pub(crate) async fn spawn_session_actor(
} else {
None
};
let embed_credentials = crate::auth::credential_provider::embedding_session_credentials(
&embed_base_url,
auth_manager.as_ref(),
api_key_provider.clone(),
);
let params = crate::session::memory::MemoryBackendParams {
session_id: session_info.id.to_string(),
embed_config: memory_config.as_ref().map(|mc| mc.embedding.clone()),
@ -723,16 +706,7 @@ pub(crate) async fn spawn_session_actor(
watcher,
stale_claim_secs: watcher_config.stale_claim_secs,
search_source: "tool",
api_key_provider: api_key_provider.clone(),
auth_credentials: auth_manager.as_ref().map(|am| {
std::sync::Arc::new(
crate::auth::credential_provider::ShellAuthCredentialProvider::new(
am.clone(),
None,
None,
),
) as std::sync::Arc<dyn xai_grok_auth::AuthCredentialProvider>
}),
embedding_credentials: embed_credentials,
};
let backend = crate::session::memory::MemoryBackendImpl::from_session_params(
storage.clone(),
@ -2043,78 +2017,3 @@ impl crate::session::mcp_restart::RestartActions for SessionRestartActions {
.end_restart(server);
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum TerminalBackendKind {
ReuseParent,
AcpClient,
LocalPersistent,
LocalNonPersistent,
}
fn select_terminal_backend_kind(
is_subagent: bool,
has_parent_backend: bool,
client_terminal_capable: bool,
has_gateway: bool,
local_persistent: bool,
) -> TerminalBackendKind {
if is_subagent && has_parent_backend {
TerminalBackendKind::ReuseParent
} else if client_terminal_capable && has_gateway {
TerminalBackendKind::AcpClient
} else if local_persistent {
TerminalBackendKind::LocalPersistent
} else {
TerminalBackendKind::LocalNonPersistent
}
}
#[cfg(test)]
mod terminal_backend_select_tests {
use super::{TerminalBackendKind, select_terminal_backend_kind};
#[test]
fn subagent_with_parent_reuses_parent() {
assert_eq!(
select_terminal_backend_kind(true, true, true, true, true),
TerminalBackendKind::ReuseParent
);
}
#[test]
fn subagent_without_parent_falls_through() {
assert_eq!(
select_terminal_backend_kind(true, false, true, true, true),
TerminalBackendKind::AcpClient
);
assert_eq!(
select_terminal_backend_kind(true, false, false, true, true),
TerminalBackendKind::LocalPersistent
);
}
#[test]
fn non_subagent_never_reuses_parent() {
assert_eq!(
select_terminal_backend_kind(false, true, false, false, true),
TerminalBackendKind::LocalPersistent
);
}
#[test]
fn client_terminal_uses_acp_only_with_gateway() {
assert_eq!(
select_terminal_backend_kind(false, false, true, true, true),
TerminalBackendKind::AcpClient
);
assert_eq!(
select_terminal_backend_kind(false, false, true, false, true),
TerminalBackendKind::LocalPersistent
);
}
#[test]
fn local_session_persistent_flag_selects_backend() {
assert_eq!(
select_terminal_backend_kind(false, false, false, false, true),
TerminalBackendKind::LocalPersistent
);
assert_eq!(
select_terminal_backend_kind(false, false, false, false, false),
TerminalBackendKind::LocalNonPersistent
);
}
}

3
crates/codegen/xai-grok-shell/src/session/acp_session_tests/cancel_running_task_tests.rs
View file

@ -555,8 +555,7 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
watcher: None,
stale_claim_secs: 60,
search_source: "tool",
api_key_provider: None,
auth_credentials: None,
embedding_credentials: crate::session::memory::EndpointScopedCredentials::none(),
};
let (event_tx, _event_rx) = tokio::sync::mpsc::unbounded_channel::<SessionEvent>();
let actor = Arc::new(SessionActor {

8
crates/codegen/xai-grok-shell/src/session/acp_session_tests/goal/goal_reminder_subagent_rules_tests.rs
View file

@ -942,7 +942,10 @@ fn render_goal_rules_substitutes_all_placeholders_in_slim_template() {
"slim template must not publish a per-goal verdict file path:\n{body}",
);
assert!(body.contains("/tmp/grok-goal-x/implementer"));
assert!(body.contains("`{SCRATCH}` placeholder resolves to"));
assert!(body.contains("`{SCRATCH}` placeholder"));
assert!(body.contains("Use existing\nuser, system, or project defaults"));
assert!(body.contains("`CARGO_HOME`, `RUSTUP_HOME`"));
assert!(body.contains("the scratch dir is deleted when the goal ends"));
for placeholder in [
"{OBJECTIVE}",
"{GOAL_TOOL}",
@ -1478,6 +1481,9 @@ fn render_goal_continuation_directive_substitutes_all_placeholders() {
body.contains("`{SCRATCH}`"),
"continuation directive must reference the {{SCRATCH}} placeholder:\n{body}",
);
assert!(body.contains("existing user, system, or project defaults"));
assert!(body.contains("`CARGO_HOME`, `RUSTUP_HOME`"));
assert!(body.contains("deleted when the goal ends"));
const AUDIT_CONTRACT: &str = "AUDITS your committed tests";
assert!(
GOAL_CONTINUATION_DIRECTIVE_TEMPLATE.contains(AUDIT_CONTRACT),

6
crates/codegen/xai-grok-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs
View file

@ -398,8 +398,7 @@ fn initial_injection_backend_params_use_override_min_score() {
watcher: None,
stale_claim_secs: 60,
search_source: "tool",
api_key_provider: None,
auth_credentials: None,
embedding_credentials: crate::session::memory::EndpointScopedCredentials::none(),
};
let initial_injection = crate::config::MemoryInitialInjectionConfig {
enabled: true,
@ -427,8 +426,7 @@ fn initial_injection_backend_params_preserve_default_zero_min_score() {
watcher: None,
stale_claim_secs: 60,
search_source: "tool",
api_key_provider: None,
auth_credentials: None,
embedding_credentials: crate::session::memory::EndpointScopedCredentials::none(),
};
let (adjusted, effective_min_score) = build_initial_injection_backend_params(
&params,

9
crates/codegen/xai-grok-shell/src/session/acp_session_tests/memory_config_tests.rs
View file

@ -18,8 +18,7 @@ fn initial_injection_backend_params_use_override_min_score() {
watcher: None,
stale_claim_secs: 60,
search_source: "tool",
api_key_provider: None,
auth_credentials: None,
embedding_credentials: crate::session::memory::EndpointScopedCredentials::none(),
};
let initial_injection = crate::config::MemoryInitialInjectionConfig {
enabled: true,
@ -47,8 +46,7 @@ fn initial_injection_backend_params_preserve_default_zero_min_score() {
watcher: None,
stale_claim_secs: 60,
search_source: "tool",
api_key_provider: None,
auth_credentials: None,
embedding_credentials: crate::session::memory::EndpointScopedCredentials::none(),
};
let (adjusted, effective_min_score) = build_initial_injection_backend_params(
&params,
@ -521,8 +519,7 @@ async fn create_injection_ready_actor(
watcher: None,
stale_claim_secs: 60,
search_source: "tool",
api_key_provider: None,
auth_credentials: None,
embedding_credentials: crate::session::memory::EndpointScopedCredentials::none(),
});
actor
.chat_state_handle

6
crates/codegen/xai-grok-shell/src/session/memory/mod.rs
View file

@ -12,7 +12,7 @@
pub mod hooks;
pub use xai_grok_memory::{
MemoryBackendImpl, MemoryBackendParams, MemoryIndex, MemoryScope, MemoryStorage, archive,
backend, chunker, dream, dream_lock, embed_missing_chunks, embedding, index, init_sqlite_vec,
mmr, query_expansion, schema, search, storage, text_utils, watcher,
EndpointScopedCredentials, MemoryBackendImpl, MemoryBackendParams, MemoryIndex, MemoryScope,
MemoryStorage, archive, backend, chunker, dream, dream_lock, embed_missing_chunks, embedding,
index, init_sqlite_vec, mmr, query_expansion, schema, search, storage, text_utils, watcher,
};

117
crates/codegen/xai-grok-shell/src/session/persistence.rs
View file

@ -306,6 +306,10 @@ pub struct SessionStateCopy {
pub enum PersistenceMsg {
/// A session update (ACP update or xAI extension update)
Update(SessionUpdate),
AppendUpdateDurablyAndAck {
update: SessionUpdate,
respond_to: tokio::sync::oneshot::Sender<io::Result<()>>,
},
ContentChunk(PersistenceContentChunk),
Chat(ConversationItem),
/// Replace the entire chat history (used for compaction)
@ -1494,32 +1498,69 @@ impl SessionPersistence {
}
}
async fn write_update(&mut self, update: &SessionUpdate) {
if let Err(e) = self.storage.append_update(&self.info, update).await {
tracing::warn!(?e, "failed to write update");
async fn write_update(
&self,
update: &SessionUpdate,
) -> Result<(), crate::session::storage::AppendUpdateError> {
self.storage
.append_update_commit_aware(&self.info, update)
.await
}
fn queue_acp_sync(&self, notification: acp::SessionNotification) {
if let Some(sync) = &self.remote_sync {
sync.queue(notification.clone());
}
if let Some(relay) = &self.relay_sync {
relay.queue(notification);
}
}
fn finish_pending_append(
pending: &mut Option<acp::SessionNotification>,
notification: acp::SessionNotification,
result: Result<(), crate::session::storage::AppendUpdateError>,
) -> Result<acp::SessionNotification, io::Error> {
match result {
Ok(()) => Ok(notification),
Err(crate::session::storage::AppendUpdateError::NotCommitted(error)) => {
*pending = Some(notification);
Err(error)
}
Err(crate::session::storage::AppendUpdateError::Committed(error)) => Err(error),
}
}
async fn drain_pending(&mut self) -> io::Result<()> {
if let Some(notification) = self.pending_notification.take() {
let result = self
.write_update(&SessionUpdate::Acp(Box::new(notification.clone())))
.await;
match Self::finish_pending_append(
&mut self.pending_notification,
notification.clone(),
result,
) {
Ok(notification) => self.queue_acp_sync(notification),
Err(error) => {
if self.pending_notification.is_none() {
self.queue_acp_sync(notification);
}
return Err(error);
}
}
}
Ok(())
}
/// Flush any pending merged ACP notification to disk and remote sync.
async fn flush_pending(&mut self) {
// Write any pending merged ACP notification
if let Some(notification) = self.pending_notification.take() {
self.write_update(&SessionUpdate::Acp(Box::new(notification.clone())))
.await;
// HTTP-based remote sync (Writeback mode)
if let Some(sync) = &self.remote_sync {
sync.queue(notification.clone());
}
// WebSocket-based relay sync (real-time sharing)
if let Some(relay) = &self.relay_sync {
relay.queue(notification);
}
if let Err(error) = self.drain_pending().await {
tracing::warn!(?error, "failed to write pending update");
}
// Flush HTTP sync
if let Some(sync) = &self.remote_sync {
sync.flush();
}
// Flush WebSocket relay
if let Some(relay) = &self.relay_sync {
relay.flush();
}
@ -1560,24 +1601,42 @@ impl SessionPersistence {
SessionUpdate::Acp(notification) => {
// ACP notifications use merging to coalesce consecutive text chunks
if let Some(to_write) = self.maybe_merge_notification(&notification) {
self.write_update(&SessionUpdate::Acp(Box::new(to_write.clone())))
.await;
// HTTP-based remote sync (Writeback mode)
if let Some(sync) = &self.remote_sync {
sync.queue(to_write.clone());
}
// WebSocket-based relay sync (real-time sharing)
if let Some(relay) = &self.relay_sync {
relay.queue(to_write);
match self
.write_update(&SessionUpdate::Acp(Box::new(to_write.clone())))
.await
{
Ok(())
| Err(crate::session::storage::AppendUpdateError::Committed(
_,
)) => {
self.queue_acp_sync(to_write);
}
Err(error) => tracing::warn!(%error, "failed to write update"),
}
}
}
SessionUpdate::Xai(_) => {
// xAI notifications are written directly without merging
self.write_update(&update).await;
if let Err(error) = self.write_update(&update).await {
tracing::warn!(%error, "failed to write update");
}
}
}
}
PersistenceMsg::AppendUpdateDurablyAndAck { update, respond_to } => {
let result = async {
self.drain_pending().await?;
self.storage
.append_update_durable(&self.info, &update)
.await?;
if let SessionUpdate::Acp(notification) = update {
self.queue_acp_sync(*notification);
}
Ok(())
}
.await;
let _ = respond_to.send(result);
}
PersistenceMsg::Chat(chat_msg) => {
if let Err(e) = self
.storage
@ -2526,6 +2585,10 @@ fn classify_remote_delete(
}
}
#[cfg(test)]
#[path = "persistence_tests.rs"]
mod durable_update_tests;
#[cfg(test)]
mod delete_session_history_tests {
use super::{DeleteSessionError, SessionDeletion, classify_remote_delete};

151
crates/codegen/xai-grok-shell/src/session/persistence_tests.rs Normal file
View file

@ -0,0 +1,151 @@
use super::*;
struct ActorGuard {
handle: PersistenceHandle,
task: tokio::task::JoinHandle<()>,
}
impl ActorGuard {
async fn stop(self) {
self.task.abort();
let _ = self.task.await;
}
}
fn test_actor(info: Info, storage: Arc<dyn StorageAdapter>) -> ActorGuard {
let (tx, rx) = mpsc::unbounded_channel();
let summary_tx = tx.clone();
let sampling_client = OaiCompatClient::new(xai_grok_sampler::SamplerConfig::default()).unwrap();
let task = tokio::spawn(
SessionPersistence {
info,
storage,
pending_notification: None,
rx,
remote_sync: None,
relay_sync: None,
summary: crate::session::summary::SummaryGenerator::new(
crate::session::summary::SummaryConfig {
sampling_client,
model: String::new(),
persistence_tx: summary_tx,
},
),
registry_title_sync: None,
gateway: None,
}
.run(),
);
ActorGuard {
handle: PersistenceHandle { tx, noop: false },
task,
}
}
fn notification(info: &Info, text: &str) -> acp::SessionNotification {
acp::SessionNotification::new(
info.id.clone(),
acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new(acp::ContentBlock::Text(
acp::TextContent::new(text),
))),
)
}
fn neutral_update(info: &Info, text: &str) -> SessionUpdate {
SessionUpdate::Acp(Box::new(notification(info, text)))
}
#[test]
fn committed_error_does_not_restore_pending_notification() {
let notification = notification(
&Info {
id: acp::SessionId::new("committed-update"),
cwd: "/test".into(),
},
"committed",
);
let mut pending = None;
let result = SessionPersistence::finish_pending_append(
&mut pending,
notification,
Err(crate::session::storage::AppendUpdateError::Committed(
io::Error::other("summary patch failed"),
)),
);
assert_eq!(result.unwrap_err().to_string(), "summary patch failed");
assert!(pending.is_none());
}
#[test]
fn uncommitted_error_restores_pending_notification() {
let notification = notification(
&Info {
id: acp::SessionId::new("uncommitted-update"),
cwd: "/test".into(),
},
"pending",
);
let mut pending = None;
let result = SessionPersistence::finish_pending_append(
&mut pending,
notification,
Err(crate::session::storage::AppendUpdateError::NotCommitted(
io::Error::other("append failed"),
)),
);
assert!(result.is_err());
assert!(pending.is_some());
}
#[tokio::test]
async fn durable_ack_drains_pending_update_in_fifo_order() {
let dir = tempfile::tempdir().unwrap();
let info = Info {
id: acp::SessionId::new("durable-update"),
cwd: dir.path().to_string_lossy().into_owned(),
};
let storage = Arc::new(JsonlStorageAdapter::with_explicit_session_dir(
dir.path().to_path_buf(),
));
storage
.init_session(&info, default_model_id())
.await
.unwrap();
let actor = test_actor(info.clone(), storage.clone());
actor
.handle
.tx
.send(PersistenceMsg::Update(neutral_update(&info, "before")))
.unwrap();
let (respond_to, response) = tokio::sync::oneshot::channel();
actor
.handle
.tx
.send(PersistenceMsg::AppendUpdateDurablyAndAck {
update: neutral_update(&info, "durable"),
respond_to,
})
.unwrap();
response.await.unwrap().unwrap();
let summary = storage.load_summary(&info).await.unwrap();
assert_eq!(summary.num_messages, 2);
let updates = storage.load_session(&info).await.unwrap().updates;
let texts = updates
.iter()
.filter_map(|update| {
let SessionUpdate::Acp(notification) = update else {
return None;
};
let acp::SessionUpdate::AgentMessageChunk(chunk) = &notification.update else {
return None;
};
let acp::ContentBlock::Text(text) = &chunk.content else {
return None;
};
Some(text.text.clone())
})
.collect::<Vec<_>>();
assert_eq!(texts, ["before", "durable"]);
actor.stop().await;
}

173
crates/codegen/xai-grok-shell/src/session/storage/jsonl/durable_tests.rs Normal file
View file

@ -0,0 +1,173 @@
use super::*;
use crate::session::info::Info;
use crate::session::persistence::default_model_id;
use crate::session::storage::{SessionUpdate, StorageAdapter};
fn info() -> Info {
Info {
id: acp::SessionId::new("durable-jsonl"),
cwd: "/test".into(),
}
}
fn update(info: &Info, text: String) -> SessionUpdate {
SessionUpdate::Acp(Box::new(acp::SessionNotification::new(
info.id.clone(),
acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new(acp::ContentBlock::Text(
acp::TextContent::new(text),
))),
)))
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn ordinary_and_durable_appends_keep_every_physical_line_parseable() {
const N: usize = 100;
let dir = tempfile::tempdir().unwrap();
let info = info();
let adapter = JsonlStorageAdapter::with_explicit_session_dir(dir.path().to_path_buf());
adapter
.init_session(&info, default_model_id())
.await
.unwrap();
let ordinary = adapter.clone();
let durable = adapter.clone();
let info_a = info.clone();
let info_b = info.clone();
let ordinary = tokio::spawn(async move {
for index in 0..N {
ordinary
.append_update(&info_a, &update(&info_a, format!("ordinary-{index}")))
.await
.unwrap();
}
});
let durable = tokio::spawn(async move {
for index in 0..N {
durable
.append_update_durable(&info_b, &update(&info_b, format!("durable-{index}")))
.await
.unwrap();
}
});
ordinary.await.unwrap();
durable.await.unwrap();
let bytes = std::fs::read(dir.path().join("updates.jsonl")).unwrap();
let parsed = bytes
.split(|byte| *byte == b'\n')
.filter(|line| !line.is_empty())
.map(serde_json::from_slice::<SessionUpdateEnvelope>)
.collect::<Result<Vec<_>, _>>()
.unwrap();
assert_eq!(parsed.len(), N * 2);
}
#[tokio::test]
async fn append_commit_is_reported_when_bookkeeping_fails() {
let dir = tempfile::tempdir().unwrap();
let info = info();
let adapter = JsonlStorageAdapter::with_explicit_session_dir(dir.path().to_path_buf());
adapter
.init_session(&info, default_model_id())
.await
.unwrap();
let result = adapter
.append_update_with_bookkeeping(&info, &update(&info, "committed".into()), async {
Err(io::Error::other("summary patch failed"))
})
.await;
assert!(matches!(
result,
Err(crate::session::storage::AppendUpdateError::Committed(_))
));
let bytes = std::fs::read(dir.path().join("updates.jsonl")).unwrap();
let parsed = bytes
.split(|byte| *byte == b'\n')
.filter(|line| !line.is_empty())
.map(serde_json::from_slice::<SessionUpdateEnvelope>)
.collect::<Result<Vec<_>, _>>()
.unwrap();
assert_eq!(parsed.len(), 1);
}
#[test]
fn lock_serializes_tail_heal_and_complete_record() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("updates.jsonl");
std::fs::write(&path, b"torn").unwrap();
JsonlStorageAdapter::append_jsonl_line_sync(
&path,
b"{\"record\":1}\n".to_vec(),
AppendDurability::Buffered,
)
.unwrap();
assert_eq!(
std::fs::read_to_string(path).unwrap(),
"torn\n{\"record\":1}\n"
);
}
#[test]
fn directory_barrier_failure_is_retried_even_after_file_exists() {
use std::sync::atomic::{AtomicUsize, Ordering};
static ATTEMPTS: AtomicUsize = AtomicUsize::new(0);
static TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
let _guard = TEST_LOCK.lock().unwrap();
fn sync_file(file: &std::fs::File) -> io::Result<()> {
file.sync_all()
}
fn flaky_parent(_path: &Path) -> io::Result<()> {
if ATTEMPTS.fetch_add(1, Ordering::SeqCst) == 0 {
Err(io::Error::other("directory barrier failed"))
} else {
Ok(())
}
}
ATTEMPTS.store(0, Ordering::SeqCst);
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("updates.jsonl");
assert!(
JsonlStorageAdapter::append_jsonl_line_sync_with(
&path,
b"{\"record\":1}\n".to_vec(),
AppendDurability::Durable,
sync_file,
flaky_parent,
)
.is_err()
);
JsonlStorageAdapter::append_jsonl_line_sync_with(
&path,
b"{\"record\":1}\n".to_vec(),
AppendDurability::Durable,
sync_file,
flaky_parent,
)
.unwrap();
assert_eq!(ATTEMPTS.load(Ordering::SeqCst), 2);
}
#[test]
fn file_barrier_error_propagates() {
fn fail(_file: &std::fs::File) -> io::Result<()> {
Err(io::Error::other("file barrier failed"))
}
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("updates.jsonl");
let error = JsonlStorageAdapter::append_jsonl_line_sync_with(
&path,
b"{\"record\":1}\n".to_vec(),
AppendDurability::Durable,
fail,
|_| Ok(()),
)
.unwrap_err();
assert_eq!(error.to_string(), "file barrier failed");
}
#[cfg(target_os = "macos")]
#[test]
fn darwin_fullfsync_seam_reports_invalid_descriptor() {
assert!(JsonlStorageAdapter::fullfsync_raw(-1).is_err());
}

197
crates/codegen/xai-grok-shell/src/session/storage/jsonl/mod.rs
View file

@ -8,10 +8,12 @@ use crate::session::persistence::{CHAT_FORMAT_VERSION, Summary};
use crate::tools::todo::TodoState;
use agent_client_protocol as acp;
use async_trait::async_trait;
use fs2::FileExt;
use std::fs::OpenOptions;
use std::io::{self, Read};
use std::io::{self, Read, Seek, Write};
#[cfg(target_os = "macos")]
use std::os::fd::AsRawFd;
use std::path::{Path, PathBuf};
use tokio::io::AsyncWriteExt;
use xai_grok_workspace::session::file_state::RewindPoint;
/// How the adapter resolves the session directory on disk.
///
@ -25,6 +27,10 @@ enum SessionDirMode {
/// New: use this directory directly (for subagent children).
Explicit(PathBuf),
}
pub(super) enum AppendDurability {
Buffered,
Durable,
}
/// JSONL-based storage adapter (legacy format)
/// Stores sessions in {root}/sessions/{url_encoded_cwd}/{session_id}/
#[derive(Clone)]
@ -75,7 +81,7 @@ impl JsonlStorageAdapter {
SessionDirMode::Explicit(dir) => dir.clone(),
}
}
fn updates_file(&self, info: &Info) -> PathBuf {
pub(super) fn updates_file(&self, info: &Info) -> PathBuf {
self.session_dir(info).join("updates.jsonl")
}
fn chat_file(&self, info: &Info) -> PathBuf {
@ -246,33 +252,125 @@ impl JsonlStorageAdapter {
/// the torn record is terminated as its own (single) corrupt line. This
/// bounds the damage of any torn write to exactly one record, which the
/// lenient readers (e.g. [`Self::read_chat_history_sync`]) then skip.
async fn append_jsonl_line(&self, path: PathBuf, mut line: Vec<u8>) -> io::Result<()> {
async fn append_jsonl_line(&self, path: PathBuf, line: Vec<u8>) -> io::Result<()> {
Self::append_jsonl_line_locked(path, line, AppendDurability::Buffered).await
}
async fn append_jsonl_line_locked(
path: PathBuf,
line: Vec<u8>,
durability: AppendDurability,
) -> io::Result<()> {
tokio::task::spawn_blocking(move || Self::append_jsonl_line_sync(&path, line, durability))
.await
.map_err(io::Error::other)?
}
fn append_jsonl_line_sync(
path: &Path,
line: Vec<u8>,
durability: AppendDurability,
) -> io::Result<()> {
Self::append_jsonl_line_sync_with(
path,
line,
durability,
Self::sync_file_durable,
Self::sync_parent_directory,
)
}
fn append_jsonl_line_sync_with(
path: &Path,
mut line: Vec<u8>,
durability: AppendDurability,
sync_file: fn(&std::fs::File) -> io::Result<()>,
sync_parent: fn(&Path) -> io::Result<()>,
) -> io::Result<()> {
debug_assert!(line.ends_with(b"\n"), "JSONL record must end with \\n");
let mut file = tokio::fs::OpenOptions::new()
let lock_path = path.with_extension("jsonl.lock");
let lock = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.append(true)
.open(&path)
.await?;
let len = file.metadata().await?.len();
if len > 0 {
use tokio::io::{AsyncReadExt as _, AsyncSeekExt as _};
file.seek(io::SeekFrom::Start(len - 1)).await?;
let mut last = [0u8; 1];
file.read_exact(&mut last).await?;
if last[0] != b'\n' {
tracing::warn!(
path = % path.display(),
"jsonl file has a torn trailing line (previous append crashed \
mid-write?); terminating it before appending"
);
line.insert(0, b'\n');
.truncate(false)
.open(lock_path)?;
lock.lock_exclusive()?;
let result = (|| {
let mut file = OpenOptions::new()
.read(true)
.create(true)
.append(true)
.open(path)?;
let len = file.metadata()?.len();
if len > 0 {
file.seek(io::SeekFrom::Start(len - 1))?;
let mut last = [0u8; 1];
file.read_exact(&mut last)?;
if last[0] != b'\n' {
tracing::warn!(
path = % path.display(), "terminating torn jsonl tail"
);
line.insert(0, b'\n');
}
}
file.write_all(&line)?;
file.flush()?;
if matches!(durability, AppendDurability::Durable) {
sync_file(&file)?;
drop(file);
sync_parent(path)?;
} else {
drop(file);
}
Ok(())
})();
let _ = lock.unlock();
result
}
#[cfg(target_os = "macos")]
fn sync_file_durable(file: &std::fs::File) -> io::Result<()> {
file.sync_all()?;
Self::fullfsync_raw(file.as_raw_fd())
}
#[cfg(target_os = "macos")]
fn fullfsync_raw(fd: std::os::fd::RawFd) -> io::Result<()> {
let result = unsafe { libc::fcntl(fd, libc::F_FULLFSYNC) };
if result == -1 {
return Err(io::Error::last_os_error());
}
file.write_all(&line).await?;
file.flush().await?;
Ok(())
}
#[cfg(all(unix, not(target_os = "macos")))]
fn sync_file_durable(file: &std::fs::File) -> io::Result<()> {
file.sync_all()
}
#[cfg(windows)]
fn sync_file_durable(file: &std::fs::File) -> io::Result<()> {
file.sync_all()
}
#[cfg(not(any(unix, windows)))]
fn sync_file_durable(_file: &std::fs::File) -> io::Result<()> {
Err(io::Error::new(
io::ErrorKind::Unsupported,
"durable file sync is unsupported on this platform",
))
}
#[cfg(unix)]
fn sync_parent_directory(path: &Path) -> io::Result<()> {
let parent = path
.parent()
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "update has no parent"))?;
std::fs::File::open(parent)?.sync_all()
}
#[cfg(windows)]
fn sync_parent_directory(_path: &Path) -> io::Result<()> {
Ok(())
}
#[cfg(not(any(unix, windows)))]
fn sync_parent_directory(_path: &Path) -> io::Result<()> {
Err(io::Error::new(
io::ErrorKind::Unsupported,
"durable directory sync is unsupported on this platform",
))
}
/// Write a full JSONL file (rewriting all items), crash-atomically: serialize
/// to a temp file then rename over the target, so a crash / `ENOSPC` mid-write
/// can't truncate the existing file (e.g. lose `rewind_points.jsonl` history).
@ -307,17 +405,34 @@ impl JsonlStorageAdapter {
Ok(items)
}
/// Append a session update to the updates.jsonl file, wrapping it in an envelope with timestamp.
async fn append_update_to_file(
pub(super) async fn append_update_to_file(
&self,
path: PathBuf,
update: &super::SessionUpdate,
durability: AppendDurability,
) -> io::Result<()> {
let envelope = SessionUpdateEnvelope::from_update(update)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
let mut line = serde_json::to_vec(&envelope)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
line.push(b'\n');
self.append_jsonl_line(path, line).await
Self::append_jsonl_line_locked(path, line, durability).await
}
pub(crate) async fn append_update_with_bookkeeping<F>(
&self,
info: &Info,
update: &super::SessionUpdate,
bookkeeping: F,
) -> Result<(), super::AppendUpdateError>
where
F: std::future::Future<Output = io::Result<()>>,
{
self.append_update_to_file(self.updates_file(info), update, AppendDurability::Buffered)
.await
.map_err(super::AppendUpdateError::NotCommitted)?;
bookkeeping
.await
.map_err(super::AppendUpdateError::Committed)
}
/// Read session updates from an updates.jsonl file, handling both envelope and legacy formats.
///
@ -951,7 +1066,35 @@ impl StorageAdapter for JsonlStorageAdapter {
.await
}
async fn append_update(&self, info: &Info, update: &super::SessionUpdate) -> io::Result<()> {
self.append_update_to_file(self.updates_file(info), update)
self.append_update_commit_aware(info, update)
.await
.map_err(super::AppendUpdateError::into_io_error)
}
async fn append_update_commit_aware(
&self,
info: &Info,
update: &super::SessionUpdate,
) -> Result<(), super::AppendUpdateError> {
self.append_update_with_bookkeeping(
info,
update,
self.apply_summary_patch(
info,
super::summary_write::SummaryPatch {
record_activity: true,
messages: Some(super::summary_write::CounterOp::Increment(1)),
..Default::default()
},
),
)
.await
}
async fn append_update_durable(
&self,
info: &Info,
update: &super::SessionUpdate,
) -> io::Result<()> {
self.append_update_to_file(self.updates_file(info), update, AppendDurability::Durable)
.await?;
self.apply_summary_patch(
info,
@ -1526,4 +1669,6 @@ fn is_valid_data_uri_image(url: &str) -> bool {
}
}
#[cfg(test)]
mod durable_tests;
#[cfg(test)]
mod tests;

44
crates/codegen/xai-grok-shell/src/session/storage/mod.rs
View file

@ -484,6 +484,28 @@ pub fn updates_truncate_for_prompt(updates: &[SessionUpdate], target_prompt_inde
updates.len()
}
#[derive(Debug)]
pub enum AppendUpdateError {
NotCommitted(io::Error),
Committed(io::Error),
}
impl AppendUpdateError {
pub fn into_io_error(self) -> io::Error {
match self {
Self::NotCommitted(error) | Self::Committed(error) => error,
}
}
}
impl std::fmt::Display for AppendUpdateError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::NotCommitted(error) | Self::Committed(error) => error.fmt(formatter),
}
}
}
/// Storage adapter trait for session persistence
/// Abstracts over different storage backends (JSONL, SQLite, etc.)
#[async_trait]
@ -512,6 +534,28 @@ pub trait StorageAdapter: Send + Sync {
/// Append a session update (ACP update or xAI extension update) and increment counter
async fn append_update(&self, info: &Info, update: &SessionUpdate) -> io::Result<()>;
/// Append one update and report whether the replay record was committed before an error.
async fn append_update_commit_aware(
&self,
info: &Info,
update: &SessionUpdate,
) -> Result<(), AppendUpdateError> {
self.append_update(info, update)
.await
.map_err(AppendUpdateError::NotCommitted)
}
/// Append one update with the ordinary bookkeeping and a durable log barrier.
///
/// Adapters without this capability return `Unsupported`; callers must tolerate a duplicate
/// record when retrying an error that occurred after the append reached storage.
async fn append_update_durable(&self, _info: &Info, _update: &SessionUpdate) -> io::Result<()> {
Err(io::Error::new(
io::ErrorKind::Unsupported,
"durable session update append is unsupported",
))
}
/// Append a chat message and increment counter
async fn append_chat_message(&self, info: &Info, message: &ConversationItem) -> io::Result<()>;

14
crates/codegen/xai-grok-shell/src/session/templates/goal_continuation_directive.md
View file

@ -12,11 +12,15 @@ Keep your {todo_tool} list current (≥1 `in_progress`, descriptive
`activeForm`). Run targeted tests after every change you make, not
just at the end. Tests must drive the SHIPPED code on the real path — no
hard-coded values, no starting past the thing under test, no
re-implementing it. Save captured test output and artifacts to your
scratch dir {scratch_dir} {scratch_status}, never shared `/tmp/...`;
the plan's `{SCRATCH}` placeholder resolves there. The verifier AUDITS your committed tests and
saved evidence rather than rebuilding them — leave honest proof or you
WILL be refuted.
re-implementing it. Use your scratch dir {scratch_dir} {scratch_status} only for
captured test output, temp scripts, and throwaway artifacts, never shared
`/tmp/...`. Use existing user, system, or project defaults for execution
dependencies and environment state. NEVER set `HOME`, `CARGO_HOME`, `RUSTUP_HOME`,
package-manager homes, virtualenvs, caches, or config dirs to scratch, or persist
references to scratch, which is deleted when the goal ends.
The plan's `{SCRATCH}` placeholder resolves there. The verifier AUDITS your committed tests
and saved evidence rather than rebuilding them — leave honest
proof or you WILL be refuted.
Before calling `{goal_tool}(completed: true)`, run the
plan's `## Verification plan` steps yourself and confirm the observations
it lists hold — the harness re-checks against those SAME steps each attempt

14
crates/codegen/xai-grok-shell/src/session/templates/goal_rules.md
View file

@ -22,11 +22,15 @@ worse than none.
VERIFY AS YOU GO: run each change. If output is visual, capture and inspect it;
for data/config, validate programmatically.
SCRATCH: write captured test output, temp scripts, and throwaway artifacts to
your private scratch dir {SCRATCH_DIR} — never to shared `/tmp/...` (skeptics and
concurrent goals collide there). {SCRATCH_STATUS} The plan's
`{SCRATCH}` placeholder resolves to it. The verifier AUDITS your committed tests and saved evidence instead of
rebuilding them, so honest, durable proof is what passes.
SCRATCH: use your private scratch dir {SCRATCH_DIR} only for captured test
output, temp scripts, and throwaway artifacts — never shared `/tmp/...` paths
(skeptics and concurrent goals collide there). {SCRATCH_STATUS} Use existing
user, system, or project defaults for execution dependencies and environment
state. NEVER set `HOME`, `CARGO_HOME`, `RUSTUP_HOME`, package-manager homes,
virtualenvs, caches, or config dirs to scratch, or write persistent config that
references scratch; the scratch dir is deleted when the goal ends. The plan's `{SCRATCH}` placeholder
resolves to it. The verifier AUDITS your committed tests and saved evidence
instead of rebuilding them, so honest, durable proof is what passes.
TEST PROACTIVELY: run targeted tests after every change, not just at the end.
Before calling `{GOAL_TOOL}(completed: true)`, run the test suite relevant to

3
crates/codegen/xai-grok-shell/src/tools/config.rs
View file

@ -35,9 +35,6 @@ pub struct BashToolConfig {
/// Whether to allow a background `&` operator in foreground commands
/// (default: `true`). Resolution: config.toml (this) > remote settings > `true`.
pub allow_background_operator: Option<bool>,
/// Declared so the unknown-key scan accepts `[toolset.bash] persistent_shell`;
/// the effective value is resolved (layered) by `resolve_persistent_local_shell`.
pub persistent_shell: Option<bool>,
}
impl BashToolConfig {

20
crates/codegen/xai-grok-shell/src/util/config/hints.rs
View file

@ -108,6 +108,7 @@ pub struct ResolvedContextualHints {
pub send_now: bool,
pub small_screen: bool,
pub word_select: bool,
pub ssh_wrap: bool,
}
impl Default for ResolvedContextualHints {
@ -119,6 +120,7 @@ impl Default for ResolvedContextualHints {
send_now: true,
small_screen: true,
word_select: true,
ssh_wrap: true,
}
}
}
@ -148,6 +150,7 @@ pub fn resolve_contextual_hints(
send_now: resolve_tip(ui.send_now, remote.and_then(|r| r.send_now)),
small_screen: resolve_tip(ui.small_screen, remote.and_then(|r| r.small_screen)),
word_select: resolve_tip(ui.word_select, remote.and_then(|r| r.word_select)),
ssh_wrap: resolve_tip(ui.ssh_wrap, remote.and_then(|r| r.ssh_wrap)),
}
}
@ -255,6 +258,7 @@ mod tests {
send_now,
small_screen: None,
word_select,
ssh_wrap: None,
}
}
@ -268,6 +272,7 @@ mod tests {
assert!(resolved.send_now, "send_now defaults ON");
assert!(resolved.small_screen, "small_screen defaults ON");
assert!(resolved.word_select, "word_select defaults ON");
assert!(resolved.ssh_wrap, "ssh_wrap defaults ON");
}
#[test]
@ -285,19 +290,26 @@ mod tests {
assert!(resolved.send_now);
assert!(resolved.small_screen);
assert!(resolved.word_select);
assert!(resolved.ssh_wrap);
}
#[test]
fn contextual_hints_remote_tier_controls_default_per_tip() {
let _g = contextual_hints_guard();
// Remote disables plan_mode; absent tips fall through to default ON.
let r = remote(None, Some(false), None, None, None);
// Remote disables plan_mode + ssh_wrap; absent tips fall through to
// default ON. Setting two distinct fields also catches a cross-wired
// resolver line (reading one remote field into another's gate).
let r = ContextualHintsRemote {
ssh_wrap: Some(false),
..remote(None, Some(false), None, None, None)
};
let resolved = resolve_contextual_hints(&ContextualHints::default(), Some(&r));
assert!(resolved.undo, "absent remote tip → default ON");
assert!(!resolved.plan_mode, "remote `false` soft-disables");
assert!(resolved.image_input);
assert!(resolved.send_now);
assert!(resolved.word_select);
assert!(!resolved.ssh_wrap, "remote `false` soft-disables ssh_wrap");
}
#[test]
@ -328,6 +340,7 @@ mod tests {
send_now: Some(false),
small_screen: Some(false),
word_select: Some(false),
ssh_wrap: Some(false),
};
let r = remote(
Some(false),
@ -344,6 +357,7 @@ mod tests {
&& resolved.send_now
&& resolved.small_screen
&& resolved.word_select
&& resolved.ssh_wrap
);
unsafe { std::env::remove_var(ENV_CONTEXTUAL_HINTS) };
}
@ -360,6 +374,7 @@ mod tests {
send_now: Some(true),
small_screen: Some(true),
word_select: Some(true),
ssh_wrap: Some(true),
};
let r = remote(Some(true), Some(true), Some(true), Some(true), Some(true));
let resolved = resolve_contextual_hints(&ui, Some(&r));
@ -370,6 +385,7 @@ mod tests {
&& !resolved.send_now
&& !resolved.small_screen
&& !resolved.word_select
&& !resolved.ssh_wrap
);
unsafe { std::env::remove_var(ENV_CONTEXTUAL_HINTS) };
}

325
crates/codegen/xai-grok-shell/src/util/config/mcp.rs
View file

@ -12,7 +12,11 @@ use xai_grok_tools::types::compat::{CompatConfig, CompatConfigToml};
pub use xai_grok_mcp::oauth_config::{McpOAuthConfig, McpOAuthConfigMap};
// MCP server config value types extracted to `xai-grok-config-types` (config
// dependency inversion); re-exported so `crate::util::config::*` paths keep working.
pub use xai_grok_config_types::{McpJsonOAuthBlock, McpServerConfig, McpServerTransportConfig};
pub use xai_grok_config_types::{
McpJsonOAuthBlock, McpPreferenceSource, McpPreferencesFile, McpServerConfig,
McpServerPreferences, McpServerTransportConfig, McpSetupConfig, McpSetupDerivedValue,
McpSetupField, McpSetupFieldType, McpSetupOption, McpSetupResolution,
};
// Permission-policy value types likewise extracted; re-exported to keep paths stable.
pub use xai_grok_config_types::{
PatternMode, PermissionConfig, PermissionRule, RuleAction, ToolFilter,
@ -141,8 +145,17 @@ pub fn load_mcp_servers_with_oauth(
let mut oauth_configs = McpOAuthConfigMap::new();
let mut acp_servers = Vec::new();
let preferences = load_mcp_preferences().file();
let sub = &crate::config::expand_env_vars_in_string;
for (name, mut config) in servers_map {
for (name, config) in servers_map {
let mut config = match config.resolve_setup(preferences.servers.get(&name)) {
McpSetupResolution::Resolved(config) => config,
McpSetupResolution::Required(_) => continue,
McpSetupResolution::Invalid(reason) => {
tracing::warn!(server = %name, error = %reason, "MCP setup config is invalid");
continue;
}
};
config.expand_strings(sub);
if let Some(oauth) = config.oauth_config() {
oauth_configs.insert(name.clone(), oauth);
@ -195,10 +208,19 @@ pub fn load_mcp_servers(cwd: &std::path::Path, compat: &CompatConfig) -> Vec<acp
/// tracking. Using [`load_mcp_servers`] there would cause all entries to be
/// tagged as `ConfigSource::ConfigToml`, hiding the true origin.
pub(crate) fn load_mcp_servers_toml_only(cwd: &std::path::Path) -> Vec<acp::McpServer> {
let preferences = load_mcp_preferences().file();
let sub = &crate::config::expand_env_vars_in_string;
load_all_mcp_configs(cwd)
.into_iter()
.filter_map(|(name, mut config)| {
.filter_map(|(name, config)| {
let mut config = match config.resolve_setup(preferences.servers.get(&name)) {
McpSetupResolution::Resolved(config) => config,
McpSetupResolution::Required(_) => return None,
McpSetupResolution::Invalid(reason) => {
tracing::warn!(server = %name, error = %reason, "MCP setup config is invalid");
return None;
}
};
config.expand_strings(sub);
config.to_acp_mcp_server(name)
})
@ -267,10 +289,19 @@ pub(crate) fn reload_mcp_servers_merged(
servers.entry(name).or_insert(config);
}
let preferences = load_mcp_preferences().file();
let sub = &crate::config::expand_env_vars_in_string;
servers
.into_iter()
.filter_map(|(name, mut config)| {
.filter_map(|(name, config)| {
let mut config = match config.resolve_setup(preferences.servers.get(&name)) {
McpSetupResolution::Resolved(config) => config,
McpSetupResolution::Required(_) => return None,
McpSetupResolution::Invalid(reason) => {
tracing::warn!(server = %name, error = %reason, "MCP setup config is invalid");
return None;
}
};
config.expand_strings(sub);
config.to_acp_mcp_server(name)
})
@ -317,6 +348,230 @@ pub fn all_toml_mcp_server_names(cwd: &std::path::Path) -> std::collections::Has
load_all_mcp_configs(cwd).keys().cloned().collect()
}
pub fn mcp_preferences_path() -> PathBuf {
xai_grok_config::grok_home().join("mcp_preferences.json")
}
/// Result of loading prefs. Corrupt files are readable as empty for resolution
/// but must not be overwritten (would clobber other servers).
#[derive(Debug, Clone)]
pub enum McpPreferencesLoad {
Ok(McpPreferencesFile),
Missing,
Corrupt,
}
impl McpPreferencesLoad {
pub fn file(&self) -> McpPreferencesFile {
match self {
Self::Ok(f) => f.clone(),
Self::Missing | Self::Corrupt => McpPreferencesFile::default(),
}
}
pub fn is_writable(&self) -> bool {
!matches!(self, Self::Corrupt)
}
}
pub fn load_mcp_preferences() -> McpPreferencesLoad {
load_mcp_preferences_from(&mcp_preferences_path())
}
pub fn load_mcp_preferences_from(path: &std::path::Path) -> McpPreferencesLoad {
let content = match std::fs::read_to_string(path) {
Ok(content) => content,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return McpPreferencesLoad::Missing,
Err(e) => {
tracing::warn!(path = %path.display(), error = %e, "failed to read MCP preferences");
return McpPreferencesLoad::Corrupt;
}
};
match serde_json::from_str(&content) {
Ok(file) => McpPreferencesLoad::Ok(file),
Err(e) => {
tracing::warn!(path = %path.display(), error = %e, "failed to parse MCP preferences");
McpPreferencesLoad::Corrupt
}
}
}
pub async fn save_mcp_preferences(prefs: &McpPreferencesFile) -> Result<()> {
save_mcp_preferences_to(&mcp_preferences_path(), prefs).await
}
pub async fn save_mcp_preferences_to(
path: &std::path::Path,
prefs: &McpPreferencesFile,
) -> Result<()> {
if matches!(load_mcp_preferences_from(path), McpPreferencesLoad::Corrupt) {
anyhow::bail!(
"refusing to overwrite unreadable MCP preferences at {}",
path.display()
);
}
let json = serde_json::to_string_pretty(prefs)?;
if let Some(parent) = path.parent() {
tokio::fs::create_dir_all(parent).await?;
}
let tmp = path.with_extension(format!(
"json.tmp.{}{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0)
));
tokio::fs::write(&tmp, &json).await?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
tokio::fs::set_permissions(&tmp, std::fs::Permissions::from_mode(0o600))
.await
.map_err(|e| anyhow::anyhow!("failed to set mcp preferences permissions: {e}"))?;
}
tokio::fs::rename(&tmp, path).await?;
Ok(())
}
/// Restore a single server key after a failed setup (best-effort).
pub async fn restore_mcp_preference_server(
server_name: &str,
previous: Option<McpServerPreferences>,
) -> Result<()> {
let load = load_mcp_preferences();
if !load.is_writable() {
return Ok(());
}
let mut prefs = load.file();
match previous {
Some(entry) => {
prefs.servers.insert(server_name.to_string(), entry);
}
None => {
prefs.servers.remove(server_name);
}
}
save_mcp_preferences(&prefs).await
}
/// Unresolved setup-bearing MCP config collected for `/mcps` list and auth.
#[derive(Debug, Clone)]
pub struct McpSetupServerEntry {
pub name: String,
pub config: McpServerConfig,
pub source: McpPreferenceSource,
}
/// Collect MCP configs that declare a `setup` schema from config and plugins.
/// Used to surface setup-required rows and drive `x.ai/mcp/setup`.
pub fn collect_mcp_setup_configs(
cwd: &std::path::Path,
plugin_registry: Option<&xai_grok_agent::plugins::PluginRegistry>,
compat: &CompatConfig,
) -> IndexMap<String, McpSetupServerEntry> {
let mut result = IndexMap::new();
for (name, (config, scope)) in load_mcp_server_configs_with_project(cwd) {
if !config.enabled || config.setup.is_none() {
continue;
}
result.insert(
name.clone(),
McpSetupServerEntry {
name,
config,
source: McpPreferenceSource {
kind: "config".to_string(),
plugin: None,
scope: Some(scope.to_string()),
},
},
);
}
if !crate::claude_import::is_claude_import_marked_with_log("collect_mcp_setup_configs") {
for (name, config) in load_claude_json_mcp_servers_as_configs(cwd, compat) {
if !config.enabled || config.setup.is_none() {
continue;
}
result.entry(name.clone()).or_insert(McpSetupServerEntry {
name,
config,
source: McpPreferenceSource {
kind: "config".to_string(),
plugin: None,
scope: Some(MCP_SCOPE_USER.to_string()),
},
});
}
for (name, config) in load_cursor_mcp_servers_as_configs(cwd, compat) {
if !config.enabled || config.setup.is_none() {
continue;
}
result.entry(name.clone()).or_insert(McpSetupServerEntry {
name,
config,
source: McpPreferenceSource {
kind: "config".to_string(),
plugin: None,
scope: Some(MCP_SCOPE_USER.to_string()),
},
});
}
for (name, config) in load_mcp_json_servers_as_configs(cwd) {
if !config.enabled || config.setup.is_none() {
continue;
}
result.entry(name.clone()).or_insert(McpSetupServerEntry {
name,
config,
source: McpPreferenceSource {
kind: "config".to_string(),
plugin: None,
scope: Some(MCP_SCOPE_PROJECT.to_string()),
},
});
}
}
if let Some(registry) = plugin_registry {
let toml_claimed_names = all_toml_mcp_server_names(cwd);
for plugin in registry.active_plugins() {
// File first, then inline; first-wins matches runtime plugin load.
let mut plugin_configs = IndexMap::new();
if let Some(ref mcp_path) = plugin.mcp_config_path
&& let Some(config) = read_mcp_json(mcp_path)
{
for (name, server) in config.mcp_servers {
plugin_configs.entry(name).or_insert(server);
}
}
if let Some(ref inline_value) = plugin.inline_mcp_servers {
let normalized =
xai_grok_agent::plugins::manifest::normalize_inline_mcp_servers(inline_value);
if let Ok(config) = serde_json::from_value::<McpConfig>(normalized) {
for (name, server) in config.mcp_servers {
plugin_configs.entry(name).or_insert(server);
}
}
}
for (name, config) in plugin_configs {
if toml_claimed_names.contains(&name) || !config.enabled || config.setup.is_none() {
continue;
}
result.entry(name.clone()).or_insert(McpSetupServerEntry {
name,
config,
source: McpPreferenceSource {
kind: "plugin".to_string(),
plugin: Some(plugin.name.clone()),
scope: None,
},
});
}
}
}
result
}
pub const MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY: &str = "__managed_gateway_connectors";
/// Persist `disabled_tools` for a server under `[disabled_mcp_tools]` in config.toml.
@ -681,10 +936,23 @@ pub(crate) fn parse_mcp_config_with_oauth(
source_label: &str,
sub: &dyn Fn(&str) -> String,
) -> (Vec<acp::McpServer>, McpOAuthConfigMap) {
let preferences = load_mcp_preferences().file();
let mut servers = Vec::new();
let mut oauth_configs = McpOAuthConfigMap::new();
for (name, server_config) in &config.mcp_servers {
let mut server_config = server_config.clone();
let mut server_config = match server_config.resolve_setup(preferences.servers.get(name)) {
McpSetupResolution::Resolved(config) => config,
McpSetupResolution::Required(_) => continue,
McpSetupResolution::Invalid(reason) => {
tracing::warn!(
source = source_label,
server = %name,
error = %reason,
"MCP setup config is invalid"
);
continue;
}
};
server_config.expand_strings(sub);
if let Some(oauth) = server_config.oauth_config() {
oauth_configs.insert(name.clone(), oauth);
@ -1651,5 +1919,52 @@ enabled = false
);
}
#[tokio::test]
async fn mcp_preferences_missing_malformed_and_save_round_trip() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("mcp_preferences.json");
assert!(matches!(
load_mcp_preferences_from(&path),
McpPreferencesLoad::Missing
));
assert!(load_mcp_preferences_from(&path).file().servers.is_empty());
std::fs::write(&path, "not json").unwrap();
assert!(matches!(
load_mcp_preferences_from(&path),
McpPreferencesLoad::Corrupt
));
let prefs = McpPreferencesFile {
version: 1,
servers: HashMap::from([(
"acme".to_string(),
McpServerPreferences {
values: HashMap::from([("site".to_string(), "us5".to_string())]),
source: Some(McpPreferenceSource {
kind: "plugin".to_string(),
plugin: Some("acme".to_string()),
scope: None,
}),
updated_at: Some("2026-06-19T00:00:00Z".to_string()),
},
)]),
};
assert!(save_mcp_preferences_to(&path, &prefs).await.is_err());
std::fs::remove_file(&path).unwrap();
save_mcp_preferences_to(&path, &prefs).await.unwrap();
let loaded = load_mcp_preferences_from(&path).file();
assert_eq!(loaded.servers["acme"].values["site"], "us5");
assert_eq!(
loaded.servers["acme"]
.source
.as_ref()
.unwrap()
.plugin
.as_deref(),
Some("acme")
);
}
// === merge_section tests ===
}

149
crates/codegen/xai-grok-shell/src/util/config/resolve/toolset.rs
View file

@ -62,155 +62,6 @@ fn resolve_search_tool_enabled(
env.or(config).or(managed).unwrap_or(true)
}
const ENV_PERSISTENT_SHELL: &str = "GROK_PERSISTENT_SHELL";
fn persistent_shell_from_toml(v: Option<&TomlValue>) -> Option<bool> {
v?.get("toolset")?
.get("bash")?
.get("persistent_shell")?
.as_bool()
}
pub fn resolve_persistent_local_shell(remote: Option<bool>) -> bool {
let requirements = crate::config::load_merged_requirements();
let layers = match crate::config::ConfigLayers::load() {
Ok(l) => Some(l),
Err(e) => {
tracing::warn!(error = %e, "persistent_local_shell: failed to load config layers");
None
}
};
resolve_persistent_local_shell_tiers(
requirements.as_ref(),
layers.as_ref().map(|l| &l.user),
layers.as_ref().map(|l| &l.managed),
layers.as_ref().map(|l| &l.system_managed),
remote,
)
}
fn resolve_persistent_local_shell_tiers(
requirements: Option<&TomlValue>,
user: Option<&TomlValue>,
managed: Option<&TomlValue>,
system_managed: Option<&TomlValue>,
remote: Option<bool>,
) -> bool {
use crate::agent::config::BoolFlag;
BoolFlag::env(ENV_PERSISTENT_SHELL)
.requirement(persistent_shell_from_toml(requirements))
.config(persistent_shell_from_toml(user))
.managed(
persistent_shell_from_toml(managed)
.or_else(|| persistent_shell_from_toml(system_managed)),
)
.feature_flag(remote)
.default(true)
.resolve()
.value
}
#[cfg(test)]
mod persistent_local_shell_tests {
use super::{ENV_PERSISTENT_SHELL, resolve_persistent_local_shell_tiers};
use toml::Value as TomlValue;
// GROK_PERSISTENT_SHELL is process-global (the documented kill-switch a dev
// may export); serialize and force it unset so these tests can't go flaky.
static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
fn guard() -> std::sync::MutexGuard<'static, ()> {
let g = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
unsafe { std::env::remove_var(ENV_PERSISTENT_SHELL) };
g
}
fn cfg(persistent: bool) -> TomlValue {
toml::from_str(&format!(
"[toolset.bash]\npersistent_shell = {persistent}\n"
))
.unwrap()
}
#[test]
fn default_is_true() {
let _g = guard();
assert!(resolve_persistent_local_shell_tiers(
None, None, None, None, None
));
}
#[test]
fn remote_false_rolls_back() {
let _g = guard();
assert!(!resolve_persistent_local_shell_tiers(
None,
None,
None,
None,
Some(false)
));
}
#[test]
fn config_false_rolls_back() {
let _g = guard();
let off = cfg(false);
assert!(!resolve_persistent_local_shell_tiers(
None,
Some(&off),
None,
None,
None
));
}
#[test]
fn config_beats_remote() {
let _g = guard();
let on = cfg(true);
assert!(resolve_persistent_local_shell_tiers(
None,
Some(&on),
None,
None,
Some(false)
));
}
#[test]
fn requirement_overrides_remote() {
let _g = guard();
let on = cfg(true);
assert!(resolve_persistent_local_shell_tiers(
Some(&on),
None,
None,
None,
Some(false)
));
}
#[test]
fn managed_and_system_managed_apply_below_config() {
let _g = guard();
let off = cfg(false);
assert!(!resolve_persistent_local_shell_tiers(
None,
None,
Some(&off),
None,
None
));
assert!(!resolve_persistent_local_shell_tiers(
None,
None,
None,
Some(&off),
None
));
}
}
/// Env override for `[toolset.ask_user_question] timeout_enabled` (parsed by
/// the shared [`xai_grok_config::env_bool`] via `BoolFlag`). The secs env var
/// lives in the tools crate (`RESPONSE_TIMEOUT_ENV`), parsed once there.

5
crates/codegen/xai-grok-shell/src/util/config/settings_writes.rs
View file

@ -60,6 +60,11 @@ pub async fn set_contextual_hint_word_select(value: bool) -> Result<()> {
update_config(|cfg| cfg.ui.contextual_hints.word_select = Some(value)).await
}
/// Persist `[ui.contextual_hints].ssh_wrap` via `update_config`.
pub async fn set_contextual_hint_ssh_wrap(value: bool) -> Result<()> {
update_config(|cfg| cfg.ui.contextual_hints.ssh_wrap = Some(value)).await
}
/// Persist `[ui].theme` via `update_config`. Caller must pass the
/// canonical theme name (`groknight`, `tokyonight`, `auto`, etc.).
pub async fn set_theme(value: String) -> Result<()> {

1
crates/codegen/xai-grok-shell/tests/test_trusted_local_plugin_refresh_e2e.rs
View file

@ -45,6 +45,7 @@ fn register_local_install(registry: &mut InstallRegistry, source: &Path) -> Inst
subdir: None,
},
registry,
false,
)
.expect("install local plugin");
let plugins = installed