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:
- Gate session-lifecycle heap steady state with a dhat soak
- Unbreak merge lifecycle e2e after default model → grok-4.5
- Scan home-scope rules dirs at <root>/rules
- Complete text-input paste and terminal parity
- Gate project roles and personas
- Use canonical editing in dialogs
- Use canonical editing in search bars
- Reject ambiguous MCP tool IDs
- Harden Git operands for plugins
- Simplify queue drain API
- Pass RFC 9207 iss through MCP OAuth token exchange
- Show leader roster when local agents map is empty
- Use canonical editing in Persona views
- Remove marketplace default-skills auto-install and purge old installs
- Use canonical editing in extension forms
- Add canonical dashboard text editing
- Use canonical editing in settings
- Add /summarize as a /recap alias
- Restore previous agent when exiting dashboard
- Use tool_choice auto for compaction
- Settings toggle for snap-prompt-to-top on send
- Update default models to grok-4.5
- Source login shell once for local bash (env + alias/function snapshot)
- Template hardcoded param names in server-native tool descriptions
- Fix System-Reminder XML tag injection in CLAUDE.md via agents_md
- Fix remote workspace-server hardcoding LSP trust (repo code execution risk)
- Clear orphaned tool-call updates at turn end
- Suppress task wake after cancel
- Send x-grok-client-identifier on direct API tool calls
- Harden dashboard peek lease transitions
- Host /btw side panel in live region (minimal mode)
- Bound scroll presentation latency
- Highlight multi-line constructs correctly in diffs and the file viewer
- Block web_fetch non-public IPs; local opt-in is explicit-host only
- Seed coding_data_retention_opt_out=false for OAuth e2es in pty-harness
- Follow up clipboard delivery feedback
- Use canonical editing in pickers
- Route TextArea through canonical editor
- Persistent "watching" status row; quieter turn markers
- Gate sensitive edit targets
- Expose agent registry counts and gate session churn on them
- Default coding data sharing to opt-out until server preference applies
- Wire chat attachment ids through gateway prompts
- On auth refresh failure, issue retry
- Forward preview provenance and computer lifecycle state
- Document independent privacy controls and scope /privacy output
- Strip SamplingError Display prefix on rate-limit UI copy
- Stop dumping Cloudflare HTML into Retry failed
- Disable in-place prompt edit (scroll jank on enter)
- Strip forced ANSI color from gh pr view JSON
- Plumb bash tool description onto ToolUsageCard wire
This commit is contained in:
grokkybara[bot] 2026-07-18 19:48:28 +01:00
commit 7cfcb20d2b
292 changed files with 23315 additions and 9209 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -630,7 +630,6 @@ pub struct Requirements {
pub struct RuntimeResolutionContext<'a> {
pub raw_config: &'a toml::Value,
pub remote_settings: Option<&'a crate::util::config::RemoteSettings>,
pub cwd: Option<&'a std::path::Path>,
pub is_headless: bool,
/// `Some(true)` = CLI explicitly enabled, `None` = defer to config/env/remote.
pub cli_subagents: Option<bool>,
@ -1151,6 +1150,9 @@ pub struct MarketplaceConfig {
/// Written/read out-of-band by `extensions::marketplace`, opaque so a wrong-typed value can't fail load.
#[serde(default)]
pub official_marketplace_auto_installed: Option<toml::Value>,
/// Written/read out-of-band by `extensions::marketplace`, opaque so a wrong-typed value can't fail load.
#[serde(default)]
pub default_skills_installs_purged: Option<toml::Value>,
}
/// A single `[[marketplace.sources]]` entry.
#[derive(Clone, Debug, Deserialize)]
@ -1455,11 +1457,11 @@ pub struct Config {
/// Keys are agent names, values are booleans. Omitted agents default to enabled.
#[serde(skip)]
pub subagent_toggle: std::collections::HashMap<String, bool>,
/// Per-subagent role definitions from `[subagents.roles]` in config.toml
/// and `.grok/roles/*.toml` file discovery.
/// Trust-independent roles from inline, user, and bundled sources.
#[serde(skip)]
pub subagent_roles:
std::collections::HashMap<String, xai_grok_subagent_resolution::config::SubagentRole>,
/// Trust-independent personas from inline, user, and bundled sources.
#[serde(skip)]
pub subagent_personas:
std::collections::HashMap<String, xai_grok_subagent_resolution::config::SubagentPersona>,
@ -1885,18 +1887,13 @@ impl Config {
config.apply_env_overrides();
Ok(config)
}
/// Populate `#[serde(skip)]` subagent fields from `SubagentsConfig::resolve()`.
/// Populate trust-independent `#[serde(skip)]` subagent base fields.
///
/// Must be called after `new_from_toml_cfg` on the **primary startup path**
/// before the config is handed to `MvpAgent`. Model-reload and API-key-reload
/// paths only read model/key fields and do not need this call.
pub fn resolve_subagents(
&mut self,
cli_flag: bool,
raw_config: &toml::Value,
cwd: Option<&std::path::Path>,
) {
let sa = crate::config::SubagentsConfig::resolve(cli_flag, raw_config, cwd);
/// before the config is handed to `MvpAgent`. Project definitions are overlaid
/// per cwd after that cwd's authoritative folder-trust resolve.
pub fn resolve_subagents(&mut self, cli_flag: bool, raw_config: &toml::Value) {
let sa = crate::config::SubagentsConfig::resolve(cli_flag, raw_config);
self.subagents_enabled = sa.enabled;
self.subagent_model_overrides = sa.models;
self.subagent_toggle = sa.toggle;
@ -1906,7 +1903,7 @@ impl Config {
/// Resolve all `#[serde(skip)]` runtime fields that have resolver functions.
///
/// Call immediately after `new_from_toml_cfg()`. Fields resolved:
/// - subagents (6 fields) via `SubagentsConfig::resolve`
/// - subagents base layers (6 fields) via `SubagentsConfig::resolve`
/// - respect_gitignore via `ToolsConfig::resolve`
/// - disable_zdr_incompatible_tools via `ToolsConfig::resolve`
/// - managed_mcps_enabled via `ManagedMcpsConfig::resolve`
@ -1924,7 +1921,7 @@ impl Config {
self.web_search_model_override = ctx.cli_web_search_model.map(|s| s.to_owned());
self.session_summary_model_override = ctx.cli_session_summary_model.map(|s| s.to_owned());
let cli_flag = ctx.cli_subagents.unwrap_or(false);
self.resolve_subagents(cli_flag, ctx.raw_config, ctx.cwd);
self.resolve_subagents(cli_flag, ctx.raw_config);
let tools = crate::config::ToolsConfig::resolve(ctx.raw_config);
self.respect_gitignore = match self.requirements.respect_gitignore.pinned() {
Some(pinned) => pinned,
@ -1975,16 +1972,11 @@ impl Config {
self.compat_resolved = resolve_compat_config(&self.compat, ctx.remote_settings);
}
/// Re-resolve eagerly-resolved runtime fields using the current `Config`
/// state and fresh `raw_config` + `cwd`. Builds a
/// [`RuntimeResolutionContext`] from the CLI flags already stored on this
/// `Config` so callers don't need to manually extract each field.
/// state and fresh `raw_config`. Builds a [`RuntimeResolutionContext`] from
/// the CLI flags already stored on this `Config`.
///
/// Integration test coverage: `tests/test_settings_refresh.rs`.
pub fn re_resolve_runtime_fields(
&mut self,
raw_config: &toml::Value,
cwd: Option<&std::path::Path>,
) {
pub fn re_resolve_runtime_fields(&mut self, raw_config: &toml::Value) {
let remote_settings = self.remote_settings.clone();
let cli_web_search_model = self.web_search_model_override.clone();
let cli_session_summary_model = self.session_summary_model_override.clone();
@ -1992,7 +1984,6 @@ impl Config {
let ctx = RuntimeResolutionContext {
raw_config,
remote_settings: remote_settings.as_ref(),
cwd,
is_headless: self.mode == AgentMode::Headless,
cli_subagents: self.cli_subagents,
cli_web_search_model: cli_web_search_model.as_deref(),
@ -2599,6 +2590,17 @@ impl Config {
.resolve()
.value
}
pub(crate) fn resolve_compaction_tool_choice(
&self,
) -> crate::util::config::CompactionToolChoice {
crate::util::config::resolve_compaction_tool_choice_from(
env_string(crate::util::config::ENV_COMPACTION_TOOL_CHOICE).as_deref(),
self.features.compaction_tool_choice.as_deref(),
self.remote_settings
.as_ref()
.and_then(|r| r.compaction_tool_choice.as_deref()),
)
}
/// Precedence: env `GROK_COMPACTION_DETAIL`, then config
/// `features.compaction_detail`, then remote settings
/// `remote_settings.compaction_detail`, then default (`verbose`). Drives the
@ -3368,6 +3370,10 @@ struct DefaultModelJson {
compaction_at_tokens: Option<CompactionAtTokens>,
#[serde(default)]
show_model_fingerprint: bool,
#[serde(default)]
auto_compact_threshold_percent: Option<u8>,
#[serde(default)]
system_prompt_label: Option<String>,
}
fn default_models(endpoints: &EndpointsConfig) -> IndexMap<String, ModelEntryConfig> {
let root: serde_json::Value = serde_json::from_str(crate::models::DEFAULT_MODELS_JSON)
@ -3402,8 +3408,8 @@ fn default_models(endpoints: &EndpointsConfig) -> IndexMap<String, ModelEntryCon
name: m.name,
description: m.description,
context_window,
auto_compact_threshold_percent: None,
system_prompt_label: None,
auto_compact_threshold_percent: m.auto_compact_threshold_percent,
system_prompt_label: m.system_prompt_label,
temperature: m.temperature,
top_p: m.top_p,
max_completion_tokens: m.max_completion_tokens,
@ -4192,6 +4198,8 @@ pub struct Features {
/// Feed the summarizer the verbatim conversation instead of the lossy rewrite; `None` = defer to env/remote settings/default (true).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub compaction_verbatim_input: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub compaction_tool_choice: Option<String>,
/// Snapshot a completed subagent's isolated worktree into a durable git ref
/// and delete its directory (resume rehydrates from the ref). This is the
/// per-deployment rollout lever (set in managed_config.toml `[features]`).
@ -5153,7 +5161,6 @@ reasoning_effort = "low"
RuntimeResolutionContext {
raw_config: raw,
remote_settings: None,
cwd: None,
is_headless: false,
cli_subagents: None,
cli_web_search_model: None,
@ -5182,7 +5189,6 @@ reasoning_effort = "low"
RuntimeResolutionContext {
raw_config: raw,
remote_settings: None,
cwd: None,
is_headless: true,
cli_subagents: None,
cli_web_search_model: None,
@ -9215,6 +9221,7 @@ agent_type = "cursor"
url = "https://mcp.test.com"
[toolset.bash]
timeout_secs = 120
login_shell_capture = true
[shortcuts]
ctrl_k = "search"
[grok_com_config]
@ -9325,6 +9332,7 @@ agent_type = "cursor"
let toml_str = r#"
[marketplace]
official_marketplace_auto_installed = "yes"
default_skills_installs_purged = "yes"
"#;
let unused = unused_keys_from_toml(toml_str);
assert!(unused.is_empty(), "got: {unused:?}");
@ -9342,6 +9350,7 @@ agent_type = "cursor"
deny = ["Bash(rm *)"]
[marketplace]
official_marketplace_auto_installed = true
default_skills_installs_purged = true
[ui]
yollo = true
"#,
@ -10108,7 +10117,6 @@ hooks = true
config.resolve_runtime_fields(&RuntimeResolutionContext {
raw_config: &raw,
remote_settings: Some(&remote),
cwd: None,
is_headless: false,
cli_subagents: None,
cli_web_search_model: None,
@ -10134,7 +10142,6 @@ hooks = true
cfg.resolve_runtime_fields(&RuntimeResolutionContext {
raw_config: &raw,
remote_settings: None,
cwd: None,
is_headless: false,
cli_subagents: None,
cli_web_search_model: None,
@ -10170,7 +10177,6 @@ hooks = true
cfg.resolve_runtime_fields(&RuntimeResolutionContext {
raw_config: &raw,
remote_settings: None,
cwd: None,
is_headless: true,
cli_subagents: None,
cli_web_search_model: None,
@ -10202,7 +10208,6 @@ hooks = true
cfg.resolve_runtime_fields(&RuntimeResolutionContext {
raw_config: &raw,
remote_settings: Some(&remote),
cwd: None,
is_headless: false,
cli_subagents: None,
cli_web_search_model: None,
@ -10225,7 +10230,6 @@ hooks = true
cfg.resolve_runtime_fields(&RuntimeResolutionContext {
raw_config: &raw,
remote_settings: None,
cwd: None,
is_headless: false,
cli_subagents: None,
cli_web_search_model: None,
@ -10248,7 +10252,6 @@ hooks = true
cfg.resolve_runtime_fields(&RuntimeResolutionContext {
raw_config: &raw,
remote_settings: None,
cwd: None,
is_headless: false,
cli_subagents: Some(true),
cli_web_search_model: None,
@ -10272,7 +10275,6 @@ hooks = true
cfg.resolve_runtime_fields(&RuntimeResolutionContext {
raw_config: &raw,
remote_settings: None,
cwd: None,
is_headless: false,
cli_subagents: None,
cli_web_search_model: None,
@ -10296,7 +10298,6 @@ hooks = true
cfg.resolve_runtime_fields(&RuntimeResolutionContext {
raw_config: &raw,
remote_settings: None,
cwd: None,
is_headless: false,
cli_subagents: None,
cli_web_search_model: Some("custom-ws"),
@ -10324,7 +10325,6 @@ hooks = true
cfg.resolve_runtime_fields(&RuntimeResolutionContext {
raw_config: &raw,
remote_settings: Some(&remote),
cwd: None,
is_headless: false,
cli_subagents: None,
cli_web_search_model: None,
@ -10347,7 +10347,6 @@ hooks = true
let ctx = RuntimeResolutionContext {
raw_config: &raw,
remote_settings: None,
cwd: None,
is_headless: false,
cli_subagents: None,
cli_web_search_model: None,
@ -10872,12 +10871,13 @@ default = "grok-4.5"
#[test]
fn resolve_model_list_inherits_context_window_from_default_when_prefetched_has_fallback() {
let cfg = Config::default();
let dm = crate::models::default_model();
let default_cw = DEFAULT_CONTEXT_WINDOW;
let entry = prefetch_model_entry("grok-build", default_cw, ApiBackend::default());
let entry = prefetch_model_entry(dm, default_cw, ApiBackend::default());
let mut prefetched = IndexMap::new();
prefetched.insert("grok-build".to_owned(), entry);
prefetched.insert(dm.to_owned(), entry);
let resolved = resolve_model_list(&cfg, Some(prefetched));
let entry = resolved.get("grok-build").expect("model must exist");
let entry = resolved.get(dm).expect("model must exist");
assert_ne!(
entry.info.context_window.get(),
default_cw,
@ -10887,12 +10887,13 @@ default = "grok-4.5"
#[test]
fn resolve_model_list_does_not_override_explicitly_set_context_window() {
let cfg = Config::default();
let dm = crate::models::default_model();
let explicit_cw = 65_536;
let entry = prefetch_model_entry("grok-build", explicit_cw, ApiBackend::default());
let entry = prefetch_model_entry(dm, explicit_cw, ApiBackend::default());
let mut prefetched = IndexMap::new();
prefetched.insert("grok-build".to_owned(), entry);
prefetched.insert(dm.to_owned(), entry);
let resolved = resolve_model_list(&cfg, Some(prefetched));
let entry = resolved.get("grok-build").expect("model must exist");
let entry = resolved.get(dm).expect("model must exist");
assert_eq!(
entry.info.context_window.get(),
explicit_cw,
@ -10902,14 +10903,15 @@ default = "grok-4.5"
#[test]
fn resolve_model_list_inherits_agent_type_and_api_backend() {
let cfg = Config::default();
let dm = crate::models::default_model();
let default_cw = DEFAULT_CONTEXT_WINDOW;
let entry = prefetch_model_entry("grok-build", default_cw, ApiBackend::default());
let entry = prefetch_model_entry(dm, default_cw, ApiBackend::default());
let mut prefetched = IndexMap::new();
prefetched.insert("grok-build".to_owned(), entry);
prefetched.insert(dm.to_owned(), entry);
let resolved = resolve_model_list(&cfg, Some(prefetched));
let entry = resolved.get("grok-build").expect("model must exist");
let entry = resolved.get(dm).expect("model must exist");
let defaults = default_model_entries(&EndpointsConfig::default());
if let Some(default) = defaults.get("grok-build") {
if let Some(default) = defaults.get(dm) {
if default.info.agent_type != DEFAULT_AGENT_TYPE {
assert_eq!(
entry.info.agent_type, default.info.agent_type,
@ -10947,23 +10949,25 @@ default = "grok-4.5"
#[test]
fn resolve_model_list_prunes_bundled_entries_not_in_prefetch() {
let cfg = Config::default();
let dm = crate::models::default_model();
let mut defs = default_model_entries(&EndpointsConfig::default());
let mut p = IndexMap::new();
if let Some(e) = defs.shift_remove("grok-build") {
p.insert("grok-build".to_string(), e);
if let Some(e) = defs.shift_remove(dm) {
p.insert(dm.to_string(), e);
}
let resolved = resolve_model_list(&cfg, Some(p));
assert!(resolved.contains_key("grok-build"));
assert!(resolved.contains_key(dm));
let no_p = resolve_model_list(&cfg, None);
assert!(no_p.contains_key("grok-build"));
assert!(no_p.contains_key(dm));
}
#[test]
fn resolve_model_list_prefetch_visibility_matches_auth_and_server_list() {
let cfg = Config::default();
let dm = crate::models::default_model();
let mut defs = default_model_entries(&EndpointsConfig::default());
let mut p = IndexMap::new();
if let Some(e) = defs.shift_remove("grok-build") {
p.insert("grok-build".to_string(), e);
if let Some(e) = defs.shift_remove(dm) {
p.insert(dm.to_string(), e);
}
let resolved = resolve_model_list(&cfg, Some(p));
let sess: Vec<_> = resolved
@ -10975,27 +10979,29 @@ default = "grok-4.5"
.filter(|e| e.visible_for_auth(false))
.collect();
assert_eq!(sess.len(), 1);
assert!(api.is_empty());
assert_eq!(api.len(), 1);
}
#[test]
fn resolve_model_list_keeps_prefetch_only_entries_and_prunes_defaults() {
let cfg = Config::default();
let dm = crate::models::default_model();
let mut p = IndexMap::new();
let e = prefetch_model_entry("secret-xyz", 200000, ApiBackend::default());
p.insert("secret-xyz".to_string(), e);
let resolved = resolve_model_list(&cfg, Some(p));
assert!(resolved.contains_key("secret-xyz"));
assert!(!resolved.contains_key("grok-build"));
assert!(!resolved.contains_key(dm));
}
#[test]
fn resolve_model_list_prefetch_replaces_bundled_entirely() {
let cfg = Config::default();
let dm = crate::models::default_model();
let mut p = IndexMap::new();
let e = prefetch_model_entry("grok-4.5", 500_000, ApiBackend::Responses);
p.insert("grok-4.5".to_string(), e);
let e = prefetch_model_entry("other-model", 500_000, ApiBackend::Responses);
p.insert("other-model".to_string(), e);
let resolved = resolve_model_list(&cfg, Some(p));
assert!(resolved.contains_key("grok-4.5"));
assert!(!resolved.contains_key("grok-build"));
assert!(resolved.contains_key("other-model"));
assert!(!resolved.contains_key(dm));
}
#[test]
fn resolve_model_list_empty_prefetch_yields_empty_base() {
@ -11003,14 +11009,14 @@ default = "grok-4.5"
let resolved = resolve_model_list(&cfg, Some(IndexMap::new()));
assert!(resolved.is_empty());
}
/// Regression: enterprise managed config aliases grok-build to their own
/// endpoint with env_key. The bundled grok-build has supported_in_api=false.
/// The config overlay must be visible to API-key users (env_key = BYOK).
/// Regression: enterprise managed config overlays env_key on an oauth-only
/// catalog entry. BYOK must force visibility for API-key users so a
/// base `supported_in_api: false` does not leak into the overlay.
#[test]
fn byok_config_overlay_visible_to_api_key_users() {
let raw: toml::Value = toml::from_str(
r#"
[model.grok-build]
[model.enterprise-alias]
model = "grok-4.5"
base_url = "https://inference.company.com/v1"
env_key = "COMPANY_TOKEN"
@ -11018,31 +11024,48 @@ default = "grok-4.5"
)
.unwrap();
let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse");
let resolved = resolve_model_list(&cfg, None);
let entry = resolved.get("grok-build").expect("grok-build must exist");
let mut base = prefetch_model_entry("enterprise-alias", 200_000, ApiBackend::default());
base.info.supported_in_api = false;
let mut prefetched = IndexMap::new();
prefetched.insert("enterprise-alias".to_owned(), base);
let resolved = resolve_model_list(&cfg, Some(prefetched));
let entry = resolved
.get("enterprise-alias")
.expect("enterprise-alias must exist");
assert!(
entry.visible_for_auth(false),
"BYOK config entry must be visible to API-key users — \
bundled supported_in_api=false must not leak into credentialed overlays"
env_key must override base supported_in_api=false"
);
}
/// Guard: config overlay WITHOUT credentials must NOT override the
/// bundled supported_in_api flag. Only BYOK triggers the override.
/// Guard: config overlay WITHOUT credentials must NOT flip the
/// bundled supported_in_api flag. Only BYOK triggers that override.
#[test]
fn plain_config_overlay_preserves_bundled_visibility() {
let raw: toml::Value = toml::from_str(
let dm = crate::models::default_model();
let bundled = default_model_entries(&EndpointsConfig::default())
.get(dm)
.expect("bundled default must exist")
.clone();
let raw: toml::Value = toml::from_str(&format!(
r#"
[model.grok-build]
[model."{dm}"]
context_window = 300000
"#,
)
"#
))
.unwrap();
let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse");
let resolved = resolve_model_list(&cfg, None);
let entry = resolved.get("grok-build").expect("grok-build must exist");
assert!(
!entry.visible_for_auth(false),
"non-BYOK config overlay must preserve bundled supported_in_api=false"
let entry = resolved.get(dm).expect("bundled default must exist");
assert_eq!(
entry.visible_for_auth(false),
bundled.visible_for_auth(false),
"non-BYOK config overlay must preserve bundled supported_in_api"
);
assert_eq!(
entry.visible_for_auth(true),
bundled.visible_for_auth(true),
"non-BYOK config overlay must preserve bundled OAuth visibility"
);
}
#[test]

2
crates/codegen/xai-grok-shell/src/agent/init.rs
View file

@ -137,6 +137,8 @@ fn init_process(cfg: &AgentConfig, auth_manager: &AuthManager) {
let grok_home = crate::util::grok_home::grok_home();
crate::builtin::extract_bundled_files(&grok_home);
crate::extensions::marketplace::purge_default_skills_installs(&grok_home);
// Auto-register is gated (default off; env/remote settings enables). Kept out
// of extract_bundled_files so the gate can read the resolved
// remote_settings, which resolve_config has populated by now.

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

@ -2314,6 +2314,7 @@ impl acp::Agent for MvpAgent {
traceparent: xai_file_utils::trace_context::current_traceparent(),
json_schema,
send_now,
admission: None,
respond_to: tx,
persist_ack: None,
parsed_prompt_tx,

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

@ -873,7 +873,6 @@ impl MvpAgent {
auth: &crate::auth::GrokAuth,
) {
self.refresh_remote_settings(auth).await;
let cwd = std::env::current_dir().ok();
{
let mut cfg = self.cfg.borrow_mut();
crate::util::config::sync_campaign_fields(&mut cfg);
@ -884,7 +883,7 @@ impl MvpAgent {
);
toml::Value::Table(toml::map::Map::new())
});
cfg.re_resolve_runtime_fields(&raw_config, cwd.as_deref());
cfg.re_resolve_runtime_fields(&raw_config);
}
self.sync_collection_config_gate();
self.emit_settings_update_notification();
@ -1401,6 +1400,7 @@ impl MvpAgent {
/// Params resolution (TOML > env > remote settings > default):
/// - `proxy_endpoint`: `[toolset.web_fetch] proxy_endpoint` > `GROK_WEB_FETCH_PROXY` > remote settings > None
/// - `allowed_domains`: `[toolset.web_fetch] allowed_domains` > remote settings > built-in defaults
/// - `allow_local`: `[toolset.web_fetch] allow_local` > `GROK_WEB_FETCH_ALLOW_LOCAL` > false
pub(super) fn prepare_web_fetch_config(
&self,
) -> xai_grok_tools::implementations::grok_build::web_fetch::WebFetchConfig {
@ -1507,10 +1507,6 @@ impl MvpAgent {
resident_roster_titles: RefCell::new(HashMap::new()),
initialize_request: OnceLock::new(),
gateway,
subagent_model_overrides: cfg.subagent_model_overrides.clone(),
subagent_toggle: cfg.subagent_toggle.clone(),
subagent_roles: cfg.subagent_roles.clone(),
subagent_personas: cfg.subagent_personas.clone(),
launch_cwd: std::env::current_dir()
.unwrap_or_else(|_| std::path::PathBuf::from(".")),
launch_dir_trust: std::cell::OnceCell::new(),
@ -1519,11 +1515,6 @@ impl MvpAgent {
cfg.plugins.cli_plugin_dirs.clone(),
),
plugin_registry_initialized: std::cell::Cell::new(false),
persona_io_summaries: cfg
.subagent_personas
.iter()
.map(|(name, p)| p.render_io_summary(name))
.collect(),
models_manager,
chat_modes: {
let chat_modes = crate::agent::chat_modes::ChatModesManager::new(
@ -3182,6 +3173,7 @@ impl MvpAgent {
.cfg
.borrow()
.resolve_compaction_verbatim_input();
let compaction_tool_choice = self.cfg.borrow().resolve_compaction_tool_choice();
let two_pass_enabled = self.cfg.borrow().is_two_pass_compaction_enabled();
let auto_update = self.cfg.borrow().cli.auto_update;
let client_type = *self.client_type.borrow();
@ -3384,7 +3376,7 @@ impl MvpAgent {
let laziness_debug_log_for_spawn = self.cfg.borrow().laziness_debug_log.clone();
let respect_gitignore = self.cfg.borrow().respect_gitignore;
let path_not_found_hints = self.cfg.borrow().path_not_found_hints;
let subagent_toggle = self.subagent_toggle.clone();
let subagent_toggle = self.cfg.borrow().subagent_toggle.clone();
let handle_display_cwd = prompt_display_cwd.clone();
let auth_manager = Some(self.auth_manager.clone());
let bash_params_json = {
@ -3527,6 +3519,7 @@ impl MvpAgent {
.as_ref()
.and_then(|m| m.get("x.ai/gitHeadChanged"))
.and_then(|v| v.as_bool());
let session_cwd = std::path::Path::new(&session_info.cwd);
let fs_watch_caps = crate::session::fs_watch::FsWatchCapabilities::resolve(crate::session::fs_watch::CapabilityInputs {
client_notify: fs_notify_config.is_some(),
hunk_tracking: hunk_plan.enabled(),
@ -3561,6 +3554,7 @@ impl MvpAgent {
system_prompt_label,
compaction_mode,
compaction_verbatim_input,
compaction_tool_choice,
two_pass_enabled,
buffering_settings,
origin_client.clone(),
@ -3608,7 +3602,7 @@ impl MvpAgent {
client_hooks,
prompt_display_cwd,
subagent_toggle,
self.persona_io_summaries.clone(),
Vec::new(),
xai_grok_agent::prompt::context::PromptAudience::Primary,
None,
None,
@ -3618,7 +3612,6 @@ impl MvpAgent {
path_not_found_hints,
tool_params_json,
{
let session_cwd = std::path::Path::new(&session_info.cwd);
let disk_cfg = crate::config::resolve_effective_plugins_config(
session_cwd,
)

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

@ -794,20 +794,6 @@ pub struct MvpAgent {
/// notification has `Next` priority. Drained by the session turn loop
/// (`inject_pending_monitor_events`) into a hidden synthetic user message.
monitor_event_buffer: xai_grok_tools::implementations::grok_build::task::types::MonitorEventBuffer,
/// Per-subagent model ID overrides from config.toml `[subagents.models]`.
/// Populated from `SubagentsConfig.models` during `with_models()`.
subagent_model_overrides: std::collections::HashMap<String, String>,
/// Per-subagent enable/disable toggles from config.toml `[subagents.toggle]`.
/// Populated from `SubagentsConfig.toggle` during `with_models()`.
subagent_toggle: std::collections::HashMap<String, bool>,
subagent_roles: std::collections::HashMap<
String,
xai_grok_subagent_resolution::config::SubagentRole,
>,
subagent_personas: std::collections::HashMap<
String,
xai_grok_subagent_resolution::config::SubagentPersona,
>,
/// The process launch directory, captured once at construction so the
/// deferred launch-dir init paths share one source of truth instead of each
/// re-calling `std::env::current_dir()` (which could drift if the process
@ -827,7 +813,6 @@ pub struct MvpAgent {
/// the first session-creating call via [`Self::ensure_plugin_registry`];
/// this flag keeps that to a single discovery walk.
plugin_registry_initialized: std::cell::Cell<bool>,
persona_io_summaries: Vec<String>,
/// Single-flight guard for the proactive bundle sync background task.
///
/// `maybe_sync_bundle_in_background` is invoked from each post-auth path
@ -1176,6 +1161,9 @@ fn inject_proxy_headers(
.map(String::from)
.unwrap_or_else(|| xai_grok_version::VERSION.to_string())
});
headers
.entry("x-grok-client-identifier".to_string())
.or_insert_with(crate::http::process_client_identifier);
if crate::util::is_cli_chat_proxy_url(base_url) {
headers
.entry("X-XAI-Token-Auth".to_string())
@ -1264,6 +1252,7 @@ mod session_lifecycle;
mod subagent_coordinator;
mod agent_ops;
mod acp_agent;
pub(crate) use session_lifecycle::RegistrySnapshot;
pub(super) use super::ext_parsers;
/// Emit the `auth.lifecycle` login span with optional user id and error
/// category. Named `auth.lifecycle` (not `auth`) to avoid colliding with the

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

@ -402,4 +402,47 @@ impl MvpAgent {
.await
.unwrap_or(true)
}
/// Entry counts for every collection [`Self::remove_session`] drains,
/// plus the workspace binding and subagent maps.
pub(crate) fn registry_snapshot(&self) -> RegistrySnapshot {
let (subagent_pending, subagent_active, subagent_completed) =
self.subagent_coordinator.borrow().registry_snapshot();
RegistrySnapshot {
sessions: self.sessions.borrow().len(),
session_threads: self.session_threads.borrow().len(),
dispatch_locks: self.dispatch_locks.borrow().len(),
session_turn_numbers: self.session_turn_numbers.borrow().len(),
permission_event_receivers: self.permission_event_receivers.borrow().len(),
model_unavailable_sessions: self.model_unavailable_sessions.borrow().len(),
session_live_state: self.session_live_state.borrow().len(),
session_index_claims: self.session_index_claims.borrow().len(),
require_gateway_sessions: self.require_gateway_sessions.borrow().len(),
subagent_pending,
subagent_active,
subagent_completed,
workspace_bindings: self
.workspace_ops
.borrow()
.as_ref()
.and_then(|ops| ops.workspace_handle().map(|h| h.session_count())),
}
}
}
/// Field names are the wire contract of `x.ai/debug/agent`'s `registries`
/// object; each maps to the same-named registry.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct RegistrySnapshot {
pub sessions: usize,
pub session_threads: usize,
pub dispatch_locks: usize,
pub session_turn_numbers: usize,
pub permission_event_receivers: usize,
pub model_unavailable_sessions: usize,
pub session_live_state: usize,
pub session_index_claims: usize,
pub require_gateway_sessions: usize,
pub subagent_pending: usize,
pub subagent_active: usize,
pub subagent_completed: usize,
pub workspace_bindings: Option<usize>,
}

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

@ -268,14 +268,17 @@ impl MvpAgent {
ps.and_then(|h| h.allowed_subagent_types.clone()),
)
};
let cli_agent_names: Vec<String> = {
let (cli_agent_names, subagent_toggle) = {
let cfg = self.cfg.borrow();
cfg.cli_agents.iter().map(|d| d.name.clone()).collect()
(
cfg.cli_agents.iter().map(|d| d.name.clone()).collect(),
cfg.subagent_toggle.clone(),
)
};
crate::agent::subagent::SubagentValidationContext {
parent_cwd,
plugin_registry: self.plugin_registry_handle.snapshot(),
subagent_toggle: self.subagent_toggle.clone(),
subagent_toggle,
allowed_subagent_types,
cli_agent_names,
}
@ -432,6 +435,23 @@ impl MvpAgent {
}
None => (None, None),
};
let project_trusted = crate::agent::folder_trust::project_scope_allowed(&parent_cwd);
let (base_roles, base_personas, subagent_model_overrides, subagent_toggle) = {
let cfg = self.cfg.borrow();
(
cfg.subagent_roles.clone(),
cfg.subagent_personas.clone(),
cfg.subagent_model_overrides.clone(),
cfg.subagent_toggle.clone(),
)
};
let (subagent_roles, subagent_personas) =
crate::config::SubagentsConfig::effective_definition_maps(
&base_roles,
&base_personas,
&parent_cwd,
project_trusted,
);
Some(crate::agent::subagent::SubagentSpawnContext {
lsp: parent_lsp,
gateway: self.gateway.clone(),
@ -484,11 +504,10 @@ impl MvpAgent {
parent_chat_state,
parent_max_turns,
available_models,
subagent_model_overrides: self.subagent_model_overrides.clone(),
subagent_toggle: self.subagent_toggle.clone(),
subagent_roles: self.subagent_roles.clone(),
subagent_personas: self.subagent_personas.clone(),
persona_io_summaries: self.persona_io_summaries.clone(),
subagent_model_overrides,
subagent_toggle,
subagent_roles,
subagent_personas,
disable_web_search: self.cfg.borrow().disable_web_search,
todo_gate: self.cfg.borrow().todo_gate,
remote_settings: self.cfg.borrow().remote_settings.clone(),
@ -549,11 +568,11 @@ impl MvpAgent {
parent_skills: None,
parent_skills_config: self.cfg.borrow().skills.clone(),
parent_compat: self.cfg.borrow().compat_resolved,
auto_wake_delivered: {
task_completion_reservations: {
let sessions = self.sessions.borrow();
sessions
.get(&parent_sid)
.and_then(|h| h.tool_context.auto_wake_delivered.clone())
.and_then(|h| h.tool_context.task_completion_reservations.clone())
},
synthetic_trace_tx: {
let sessions = self.sessions.borrow();

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

@ -2517,6 +2517,45 @@ async fn prepare_image_gen_config_fails_open_without_auth() {
"no resolved auth ⇒ fail open (tools not tier-restricted)"
);
}
/// The imagine tools bypass cli-chat-proxy (direct API calls), so the server
/// can only scope the coding data-retention opt-out (`/privacy opt-out`) to
/// Build traffic via the `x-grok-client-identifier` header. If this header is
/// dropped, opted-out users' imagine prompts are logged/retained server-side.
#[tokio::test(flavor = "current_thread")]
async fn prepare_image_gen_config_sends_client_identifier_header() {
use xai_grok_tools::implementations::grok_build::image_gen::ImageGenConfig;
let agent = build_minimal_agent_for_tests();
agent.sampling_config.borrow_mut().api_key = Some("test-key".to_string());
let ImageGenConfig::Enabled { extra_headers, .. } = agent.prepare_image_gen_config() else {
panic!("expected Enabled");
};
assert_eq!(
extra_headers
.get("x-grok-client-identifier")
.map(String::as_str),
Some(crate::http::process_client_identifier().as_str()),
"imagine API calls must carry the client identifier so the server \
applies the coding ZDR opt-out to Build traffic"
);
}
/// Same contract for video generation (also a direct API call).
#[tokio::test(flavor = "current_thread")]
async fn prepare_video_gen_config_sends_client_identifier_header() {
use xai_grok_tools::implementations::grok_build::video_gen::VideoGenConfig;
let agent = build_minimal_agent_for_tests();
agent.sampling_config.borrow_mut().api_key = Some("test-key".to_string());
let VideoGenConfig::Enabled { extra_headers, .. } = agent.prepare_video_gen_config() else {
panic!("expected Enabled");
};
assert_eq!(
extra_headers
.get("x-grok-client-identifier")
.map(String::as_str),
Some(crate::http::process_client_identifier().as_str()),
"video gen API calls must carry the client identifier so the server \
applies the coding ZDR opt-out to Build traffic"
);
}
#[tokio::test]
async fn data_collection_enabled_for_normal_user() {
let agent = build_agent_with_auth(crate::auth::GrokAuth::test_default());
@ -3642,12 +3681,148 @@ fn repo_with_project_mcp_server() -> tempfile::TempDir {
.unwrap();
tmp
}
fn write_project_subagent_definitions(cwd: &std::path::Path) {
let roles = cwd.join(".grok/roles");
let personas = cwd.join(".grok/personas");
std::fs::create_dir_all(&roles).unwrap();
std::fs::create_dir_all(&personas).unwrap();
std::fs::write(roles.join("probe.toml"), "description = \"Project role\"").unwrap();
std::fs::write(
personas.join("probe.toml"),
"instructions = \"Project persona\"",
)
.unwrap();
}
fn folder_trust_on() -> crate::util::config::RemoteSettings {
crate::util::config::RemoteSettings {
folder_trust_enabled: Some(true),
..Default::default()
}
}
#[test]
#[serial_test::serial]
fn subagent_spawn_context_reloads_project_definitions_after_trust_changes() {
let repo = tempfile::tempdir().unwrap();
git2::Repository::init(repo.path()).unwrap();
write_project_subagent_definitions(repo.path());
run_local_for_bridge_test(|| async {
let (agent, _rx) = build_agent_with_gateway_rx();
let sid = acp::SessionId::new("roles-personas-trust-transition");
let (mut handle, _tx, _cmd_rx) = make_live_session_handle(&sid, None);
handle.info.cwd = repo.path().display().to_string();
agent.sessions.borrow_mut().insert(sid.clone(), handle);
{
let mut cfg = agent.cfg.borrow_mut();
cfg.subagent_roles.insert(
"refreshed".into(),
xai_grok_subagent_resolution::config::SubagentRole {
description: "Refreshed user role".into(),
source_dir: Some(repo.path().join("user-roles")),
..Default::default()
},
);
cfg.subagent_model_overrides
.insert("probe".into(), "refreshed-model".into());
cfg.subagent_toggle.insert("probe".into(), false);
}
crate::agent::folder_trust::record_for_test(repo.path(), false);
let untrusted = agent.build_subagent_spawn_context(sid.0.as_ref());
assert!(!untrusted.subagent_roles.contains_key("probe"));
assert!(!untrusted.subagent_personas.contains_key("probe"));
assert_eq!(
untrusted
.subagent_roles
.get("refreshed")
.map(|role| role.description.as_str()),
Some("Refreshed user role")
);
assert_eq!(
untrusted
.subagent_model_overrides
.get("probe")
.map(String::as_str),
Some("refreshed-model")
);
assert_eq!(untrusted.subagent_toggle.get("probe"), Some(&false));
crate::agent::folder_trust::record_for_test(repo.path(), true);
let trusted = agent.build_subagent_spawn_context(sid.0.as_ref());
assert_eq!(
trusted
.subagent_roles
.get("probe")
.map(|role| role.description.as_str()),
Some("Project role")
);
assert!(trusted.subagent_personas.contains_key("probe"));
crate::agent::folder_trust::record_for_test(repo.path(), false);
let revoked = agent.build_subagent_spawn_context(sid.0.as_ref());
assert!(!revoked.subagent_roles.contains_key("probe"));
assert!(!revoked.subagent_personas.contains_key("probe"));
});
}
/// End-to-end gate wiring: project `.grok/roles` / `personas` alone must drive
/// real `resolve_and_record` untrusted (not a forced `record_for_test` verdict),
/// keep project defs out of Task spawn context, then re-admit them after grant.
#[test]
#[serial_test::serial]
fn project_roles_personas_gated_via_resolve_and_record_chain() {
use xai_grok_test_support::EnvGuard;
let home = tempfile::tempdir().unwrap();
let _env = EnvGuard::set("GROK_HOME", home.path());
let _sim = EnvGuard::set(xai_grok_version::TEST_VERSION_ENV, "0.0-sim");
let _flag = EnvGuard::unset("GROK_FOLDER_TRUST");
let repo = tempfile::tempdir().unwrap();
git2::Repository::init(repo.path()).unwrap();
write_project_subagent_definitions(repo.path());
run_local_for_bridge_test(|| async {
let (agent, _rx) = build_agent_with_gateway_rx();
let sid = acp::SessionId::new("roles-personas-resolve-chain");
let (mut handle, _tx, _cmd_rx) = make_live_session_handle(&sid, None);
handle.info.cwd = repo.path().display().to_string();
agent.sessions.borrow_mut().insert(sid.clone(), handle);
let allowed = crate::agent::folder_trust::resolve_and_record(
repo.path(),
Some(&folder_trust_on()),
false,
);
assert!(
!allowed,
"roles/personas markers alone must resolve untrusted without a grant"
);
assert!(
!crate::agent::folder_trust::project_scope_allowed(repo.path()),
"cached verdict after resolve_and_record must stay untrusted"
);
let untrusted = agent.build_subagent_spawn_context(sid.0.as_ref());
assert!(
!untrusted.subagent_roles.contains_key("probe"),
"untrusted: project role must stay out of spawn context"
);
assert!(
!untrusted.subagent_personas.contains_key("probe"),
"untrusted: project persona must stay out of spawn context"
);
crate::agent::folder_trust::grant_folder_trust(repo.path());
let allowed = crate::agent::folder_trust::resolve_and_record(
repo.path(),
Some(&folder_trust_on()),
false,
);
assert!(allowed, "store-granted folder must resolve trusted");
let trusted = agent.build_subagent_spawn_context(sid.0.as_ref());
assert_eq!(
trusted
.subagent_roles
.get("probe")
.map(|role| role.description.as_str()),
Some("Project role")
);
assert!(
trusted.subagent_personas.contains_key("probe"),
"trusted: project persona must enter spawn context after grant"
);
});
}
/// Pull the next `x.ai/folder_trust/request` reverse-request off the gateway and
/// answer it with `outcome`. Returns the request's decoded params.
async fn answer_folder_trust_request(
@ -4559,3 +4734,5 @@ mod soft_default_settings_emit {
.await;
}
}
#[cfg(feature = "dhat-heap")]
mod dhat_soak;

101
crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests/dhat_soak.rs Normal file
View file

@ -0,0 +1,101 @@
//! Heap-leak test for the session lifecycle: create and remove many sessions,
//! then fail if heap memory grows per session. Run:
//! leader_session_lifecycle_heap_steady_state -- --ignored --nocapture
use super::*;
use xai_grok_workspace::permission::PermissionEvent;
const MAX_BLOCKS_PER_SESSION: f64 = 0.5;
const MAX_BYTES_PER_SESSION: f64 = 1024.0;
/// Creates the per-session state that `remove_session` must clean up, then
/// removes the session. A full `SessionHandle` would allocate so much
/// unrelated memory that a small leak would be lost in the noise.
fn populate_and_evict(agent: &MvpAgent, i: usize) {
let sid = acp::SessionId::new(format!("soak-{i}"));
{
let ops = agent.workspace_ops.borrow();
let ops = ops.as_ref().expect("test installs workspace ops");
let toolset = std::sync::Arc::new(
xai_grok_tools::registry::types::FinalizedToolset::empty_for_test(),
);
ops.bind_local_session(
sid.0.as_ref(),
std::env::temp_dir(),
xai_hunk_tracker::HunkTrackerHandle::noop(),
toolset,
None,
)
.expect("bind_local_session must succeed");
}
let (_ptx, prx) = tokio::sync::mpsc::unbounded_channel::<PermissionEvent>();
agent
.permission_event_receivers
.borrow_mut()
.insert(sid.clone(), prx);
agent
.session_turn_numbers
.borrow_mut()
.insert(sid.clone(), i as u64);
agent.model_unavailable_sessions.borrow_mut().insert(
sid.0.to_string(),
acp::ModelId::new(std::sync::Arc::from("gone-model")),
);
agent.remove_session(&sid);
}
/// Waits for background tasks to finish before reading heap stats.
async fn quiesce() {
const YIELD_ROUNDS: usize = 50;
const SETTLE: std::time::Duration = std::time::Duration::from_millis(50);
for _ in 0..YIELD_ROUNDS {
tokio::task::yield_now().await;
}
tokio::time::sleep(SETTLE).await;
for _ in 0..YIELD_ROUNDS {
tokio::task::yield_now().await;
}
}
/// Creating and removing N sessions must not grow the heap.
///
/// Only one `dhat::Profiler` can exist at a time, and the test harness runs
/// tests in parallel, so keep this the only test that creates one.
#[test]
fn leader_session_lifecycle_heap_steady_state() {
run_local_for_bridge_test(|| async {
let agent = build_minimal_agent_for_tests();
*agent.workspace_ops.borrow_mut() = Some(xai_grok_workspace::WorkspaceOps::for_test());
let _profiler = dhat::Profiler::builder().testing().build();
const WARMUP: usize = 16;
const MEASURE: usize = 256;
for i in 0..WARMUP {
populate_and_evict(&agent, i);
}
quiesce().await;
let before = dhat::HeapStats::get();
for i in WARMUP..(WARMUP + MEASURE) {
populate_and_evict(&agent, i);
}
quiesce().await;
let after = dhat::HeapStats::get();
let d_blocks = after.curr_blocks as i64 - before.curr_blocks as i64;
let d_bytes = after.curr_bytes as i64 - before.curr_bytes as i64;
let blocks_per = d_blocks as f64 / MEASURE as f64;
let bytes_per = d_bytes as f64 / MEASURE as f64;
eprintln!(
"DHAT_SOAK_SUMMARY {}",
serde_json::json!({ "warmup_sessions" : WARMUP,
"measured_sessions" : MEASURE, "before_blocks" : before.curr_blocks,
"before_bytes" : before.curr_bytes, "after_blocks" : after.curr_blocks,
"after_bytes" : after.curr_bytes, "blocks_per_session" : blocks_per,
"bytes_per_session" : bytes_per, "max_blocks_per_session" :
MAX_BLOCKS_PER_SESSION, "max_bytes_per_session" : MAX_BYTES_PER_SESSION,
"pass" : blocks_per < MAX_BLOCKS_PER_SESSION && bytes_per <
MAX_BYTES_PER_SESSION })
);
assert!(
blocks_per < MAX_BLOCKS_PER_SESSION,
"block-count leak: {blocks_per:.3} blocks/session retained ({d_blocks} over {MEASURE} cycles) exceeds the {MAX_BLOCKS_PER_SESSION} gate"
);
assert!(
bytes_per < MAX_BYTES_PER_SESSION,
"byte leak: {bytes_per:.1} bytes/session retained ({d_bytes} over {MEASURE} cycles) exceeds the {MAX_BYTES_PER_SESSION} gate"
);
});
}

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

@ -302,6 +302,10 @@ impl SubagentCoordinator {
will_wake: false,
})
}
/// Lifecycle-map entry counts as `(pending, active, completed)`.
pub(crate) fn registry_snapshot(&self) -> (usize, usize, usize) {
(self.pending.len(), self.active.len(), self.completed.len())
}
/// 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 {

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

@ -1081,6 +1081,7 @@ pub(crate) async fn handle_subagent_request(
xai_grok_agent::DEFAULT_SYSTEM_PROMPT_LABEL.to_string(),
xai_chat_state::CompactionMode::Summary,
ctx.resolve_compaction_verbatim_input(),
ctx.resolve_compaction_tool_choice(),
false,
None,
None,
@ -1156,7 +1157,7 @@ pub(crate) async fn handle_subagent_request(
ctx.client_hooks.clone(),
None,
std::collections::HashMap::new(),
ctx.persona_io_summaries.clone(),
Vec::new(),
xai_grok_agent::prompt::context::PromptAudience::Subagent,
effective_runtime.role_prompt.clone(),
None,
@ -1304,6 +1305,7 @@ pub(crate) async fn handle_subagent_request(
traceparent: xai_file_utils::trace_context::current_traceparent(),
json_schema: None,
send_now: false,
admission: None,
respond_to: prompt_tx,
persist_ack: None,
parsed_prompt_tx: None,
@ -2005,7 +2007,7 @@ pub(crate) async fn handle_subagent_request(
&request.id,
&result,
&request,
&ctx.auto_wake_delivered,
&ctx.task_completion_reservations,
ctx.parent_cmd_tx.as_ref(),
&ctx.task_output_tool_name,
&ctx.synthetic_trace_tx,

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

@ -243,9 +243,6 @@ pub(crate) struct SubagentSpawnContext {
/// Subagent personas config for persona/SOUL layering.
pub subagent_personas:
std::collections::HashMap<String, xai_grok_subagent_resolution::config::SubagentPersona>,
/// Pre-rendered persona IO summaries for the task tool description.
/// Threaded through to child sessions for recursive persona discovery.
pub persona_io_summaries: Vec<String>,
/// Parent session's ChatStateHandle — used to read the actual live
/// sampling config and credentials from the parent session actor (async).
/// Cheap Clone (mpsc sender). `None` when parent SessionHandle not found.
@ -348,9 +345,9 @@ pub(crate) struct SubagentSpawnContext {
/// Parent's resolved vendor-compat config, inherited by the child so its
/// skills / rules / AGENTS.md discovery honors the same vendor toggles.
pub parent_compat: xai_grok_tools::types::compat::CompatConfig,
/// Shared set of IDs delivered via auto-wake synthetic prompts.
pub auto_wake_delivered:
Option<xai_grok_tools::reminders::task_completion::AutoWakeDeliveredIds>,
/// Shared completion reservations held by auto-wake prompts.
pub task_completion_reservations:
Option<xai_grok_tools::reminders::task_completion::TaskCompletionReservations>,
/// Channel for requesting trace uploads for synthetic auto-wake turns.
pub synthetic_trace_tx:
Option<tokio::sync::mpsc::UnboundedSender<crate::upload::turn::SyntheticTurnTraceRequest>>,
@ -422,6 +419,18 @@ impl SubagentSpawnContext {
.resolve()
.value
}
pub fn resolve_compaction_tool_choice(&self) -> crate::util::config::CompactionToolChoice {
crate::util::config::resolve_compaction_tool_choice_from(
crate::agent::config::env_string(crate::util::config::ENV_COMPACTION_TOOL_CHOICE)
.as_deref(),
self.agent_config
.as_ref()
.and_then(|c| c.features.compaction_tool_choice.as_deref()),
self.remote_settings
.as_ref()
.and_then(|r| r.compaction_tool_choice.as_deref()),
)
}
/// Whether a completed subagent's worktree is snapshotted into a durable ref
/// and its directory deleted. Resolution mirrors the other subagent gates
/// (env > config > remote settings > default). Default `false` so it ships dark;
@ -1955,7 +1964,7 @@ fn cancellation_error_message(
/// result has not already been consumed (via block-wait or explicit kill).
/// Also suppressed while the parent's goal loop is active (mirrors the bash
/// gate in `notification_bridge`); skipping the inject also skips the
/// `auto_wake_delivered.insert`, leaving surfaces 2/3 free to drain it.
/// the completion reservation, leaving surfaces 2/3 free to drain it.
/// `parent_channel_open` folds `inject_subagent_completed_prompt`'s own
/// no-channel bail into the decision, so the `will_wake` stamped on the
/// completion notification can never promise a wake the inject won't do.
@ -1992,7 +2001,9 @@ fn inject_subagent_completed_prompt(
subagent_id: &str,
result: &SubagentResult,
request: &SubagentRequest,
auto_wake_delivered: &Option<xai_grok_tools::reminders::task_completion::AutoWakeDeliveredIds>,
task_completion_reservations: &Option<
xai_grok_tools::reminders::task_completion::TaskCompletionReservations,
>,
parent_cmd_tx: Option<&mpsc::UnboundedSender<SessionCommand>>,
task_output_tool_name: &str,
synthetic_trace_tx: &Option<
@ -2002,8 +2013,8 @@ fn inject_subagent_completed_prompt(
let Some(cmd_tx) = parent_cmd_tx else {
return;
};
if let Some(auto_wake) = auto_wake_delivered {
auto_wake.insert(subagent_id.to_string());
if let Some(reservations) = task_completion_reservations {
reservations.reserve(subagent_id.to_string());
}
let summary = SubagentCompletionSummary {
subagent_id: subagent_id.to_string(),
@ -2032,21 +2043,30 @@ fn inject_subagent_completed_prompt(
};
let (respond_to, completion_rx) = tokio::sync::oneshot::channel();
let prompt_blocks = vec![acp::ContentBlock::Text(acp::TextContent::new(wrapped))];
let _ = cmd_tx.send(SessionCommand::Prompt {
prompt_id: prompt_id.clone(),
prompt_blocks,
prompt_mode: crate::session::plan_mode::PromptMode::Agent,
artifact_upload_ctx: None,
client_identifier: None,
screen_mode: None,
verbatim: true,
traceparent: None,
json_schema: None,
send_now: false,
respond_to,
persist_ack: None,
parsed_prompt_tx: None,
});
if cmd_tx
.send(SessionCommand::Prompt {
prompt_id: prompt_id.clone(),
prompt_blocks,
prompt_mode: crate::session::plan_mode::PromptMode::Agent,
artifact_upload_ctx: None,
client_identifier: None,
screen_mode: None,
verbatim: true,
traceparent: None,
json_schema: None,
send_now: false,
admission: None,
respond_to,
persist_ack: None,
parsed_prompt_tx: None,
})
.is_err()
{
if let Some(reservations) = task_completion_reservations {
reservations.release(subagent_id);
}
return;
}
if let Some(trace_tx) = synthetic_trace_tx {
let _ = trace_tx.send(crate::upload::turn::SyntheticTurnTraceRequest {
session_id: acp::SessionId::new(request.parent_session_id.clone()),

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

@ -556,7 +556,7 @@ fn should_auto_wake_subagent_suppressed_by_block_waited_or_killed() {
}
/// A goal loop active in the parent suppresses the subagent
/// auto-wake synthetic prompt — the structural sibling of the bash gate.
/// Skipping the inject here also skips `auto_wake_delivered.insert`, so the
/// Skipping the inject here also skips its completion reservation, so the
/// per-tool-call / between-turn surfaces stay free to drain the completion.
#[test]
fn should_auto_wake_subagent_suppressed_by_goal_loop() {
@ -587,7 +587,7 @@ fn auto_wake_test_request(id: &str) -> SubagentRequest {
}
/// Behavior-level: the action half of the subagent auto-wake.
/// When the gate lets it run, `inject_subagent_completed_prompt` sends the
/// synthetic `Prompt` to the parent AND marks the id auto-wake-delivered.
/// synthetic `Prompt` to the parent and reserves its completion ID.
/// Paired with `should_auto_wake_subagent_suppressed_by_goal_loop`, this
/// proves the full Gap-1 contract on the subagent surface: goal active →
/// gate false → this never runs (no prompt, not marked, so surfaces 2/3
@ -595,7 +595,7 @@ fn auto_wake_test_request(id: &str) -> SubagentRequest {
#[test]
fn inject_subagent_completed_prompt_sends_prompt_and_marks_delivered() {
let (cmd_tx, mut cmd_rx) = mpsc::unbounded_channel::<SessionCommand>();
let auto_wake = xai_grok_tools::reminders::task_completion::AutoWakeDeliveredIds::default();
let reservations = xai_grok_tools::reminders::task_completion::TaskCompletionReservations::default();
let request = auto_wake_test_request("sa-1");
let result = SubagentResult {
success: true,
@ -607,7 +607,7 @@ fn inject_subagent_completed_prompt_sends_prompt_and_marks_delivered() {
"sa-1",
&result,
&request,
&Some(auto_wake.clone()),
&Some(reservations.clone()),
Some(&cmd_tx),
"get_command_or_subagent_output",
&None,
@ -619,7 +619,36 @@ fn inject_subagent_completed_prompt_sends_prompt_and_marks_delivered() {
}
_ => panic!("expected SessionCommand::Prompt"),
}
assert_eq!(auto_wake.snapshot(), vec!["sa-1".to_string()]);
assert_eq!(reservations.snapshot(), vec!["sa-1".to_string()]);
}
#[test]
fn inject_subagent_completed_prompt_releases_reservation_when_parent_closed() {
let (cmd_tx, cmd_rx) = mpsc::unbounded_channel::<SessionCommand>();
drop(cmd_rx);
let reservations = xai_grok_tools::reminders::task_completion::TaskCompletionReservations::default();
reservations.reserve("sa-closed".into());
let (trace_tx, mut trace_rx) = mpsc::unbounded_channel();
inject_subagent_completed_prompt(
"sa-closed",
&SubagentResult {
success: true,
subagent_id: "sa-closed".into(),
child_session_id: "sa-closed".into(),
..Default::default()
},
&auto_wake_test_request("sa-closed"),
&Some(reservations.clone()),
Some(&cmd_tx),
"get_command_or_subagent_output",
&Some(trace_tx),
);
assert!(
reservations.contains("sa-closed"),
"send failure must release only the reservation acquired by this attempt"
);
reservations.release("sa-closed");
assert!(! reservations.contains("sa-closed"));
assert!(trace_rx.try_recv().is_err());
}
#[test]
fn mark_explicitly_killed_sets_flag_on_completed() {

2
crates/codegen/xai-grok-shell/src/auth/device_code.rs
View file

@ -482,7 +482,7 @@ async fn build_auth(
team_role: None,
user_blocked_reason: None,
team_blocked_reasons: vec![],
coding_data_retention_opt_out: false,
coding_data_retention_opt_out: crate::auth::default_coding_data_retention_opt_out(),
has_grok_code_access: None,
refresh_token: tokens.refresh_token.clone(),
expires_at: tokens.expires_in.map(|s| now + Duration::seconds(s)),

2
crates/codegen/xai-grok-shell/src/auth/external_auth.rs
View file

@ -72,7 +72,7 @@ pub(crate) fn parse_output(output: &std::process::Output) -> anyhow::Result<Grok
organization_role: None,
user_blocked_reason: None,
team_blocked_reasons: vec![],
coding_data_retention_opt_out: false,
coding_data_retention_opt_out: crate::auth::default_coding_data_retention_opt_out(),
has_grok_code_access: None,
refresh_token,
expires_at,

22
crates/codegen/xai-grok-shell/src/auth/meta.rs
View file

@ -11,7 +11,7 @@ pub struct GateInfo {
}
/// Typed auth metadata passed from the shell to the pager via ACP.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuthMeta {
#[serde(default)]
pub email: Option<String>,
@ -26,7 +26,8 @@ pub struct AuthMeta {
pub is_zdr: bool,
#[serde(default)]
pub team_role: Option<String>,
#[serde(default)]
/// Defaults to opted-out (safer) until auth meta is populated.
#[serde(default = "crate::auth::default_coding_data_retention_opt_out")]
pub coding_data_retention_opt_out: bool,
#[serde(default)]
pub show_resolved_model: Option<bool>,
@ -38,3 +39,20 @@ pub struct AuthMeta {
#[serde(default)]
pub subscription_tier: Option<String>,
}
impl Default for AuthMeta {
fn default() -> Self {
Self {
email: None,
auth_mode: None,
team_id: None,
team_name: None,
is_zdr: false,
team_role: None,
coding_data_retention_opt_out: crate::auth::default_coding_data_retention_opt_out(),
show_resolved_model: None,
gate: None,
subscription_tier: None,
}
}
}

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

@ -37,7 +37,9 @@ pub use error::{AuthError, RefreshTokenError, RefreshTokenFailedReason};
pub use manager::{AuthManager, shared_api_key_provider};
pub use meta::{AuthMeta, GateInfo};
pub use model::{AuthMode, GrokAuth, lookup_auth};
pub(crate) use model::{TOKEN_TTL, UserInfo, is_expired, token_suffix};
pub(crate) use model::{
TOKEN_TTL, UserInfo, default_coding_data_retention_opt_out, is_expired, token_suffix,
};
pub(crate) use refresh::DiagnosticUploader;
pub use storage::{
clear_api_key, read_api_key, read_auth_json, read_token_by_scope, store_api_key,

35
crates/codegen/xai-grok-shell/src/auth/model.rs
View file

@ -16,6 +16,13 @@ pub const API_KEY_SCOPE: &str = "xai::api_key";
const BLOCKED_REASON_NO_LOGS: &str = "BLOCKED_REASON_NO_LOGS";
const BLOCKED_REASON_NO_LOGS_MODERATED: &str = "BLOCKED_REASON_NO_LOGS_MODERATED";
/// Fresh-credential / missing-field default: opted out until the user or
/// server enrichment opts in. Single source for `GrokAuth`, `AuthMeta`, and
/// every login-path constructor so the sides cannot drift.
pub(crate) fn default_coding_data_retention_opt_out() -> bool {
true
}
/// Token provenance (debugging/auth.json only -- no code branches on this).
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
@ -69,7 +76,9 @@ pub struct GrokAuth {
pub user_blocked_reason: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub team_blocked_reasons: Vec<String>,
#[serde(default)]
/// Defaults to `true` (opted out) for safer consumer privacy until the
/// user explicitly shares or server enrichment sets the team preference.
#[serde(default = "default_coding_data_retention_opt_out")]
pub coding_data_retention_opt_out: bool,
/// Deprecated. Kept for deserializing existing auth.json files.
@ -218,7 +227,7 @@ impl Default for GrokAuth {
organization_role: None,
user_blocked_reason: None,
team_blocked_reasons: vec![],
coding_data_retention_opt_out: false,
coding_data_retention_opt_out: default_coding_data_retention_opt_out(),
has_grok_code_access: None,
refresh_token: None,
expires_at: None,
@ -239,6 +248,9 @@ impl GrokAuth {
Self {
key: "test-key".into(),
user_id: "test-user".into(),
// Tests that exercise collection gates need sharing enabled by
// default; opt out explicitly when asserting the privacy path.
coding_data_retention_opt_out: false,
..Default::default()
}
}
@ -486,4 +498,23 @@ mod tests {
let info: UserInfo = serde_json::from_str(json).unwrap();
assert_eq!(info.subscription_tier.as_deref(), Some(""));
}
/// Pre-default auth.json (no coding_data_retention_opt_out key) must
/// deserialize as opted-out, not the old fail-open false.
#[test]
fn missing_coding_data_retention_opt_out_deserializes_opted_out() {
let json = r#"{
"key": "k",
"auth_mode": "oidc",
"create_time": "2020-01-01T00:00:00Z",
"user_id": "u"
}"#;
let auth: GrokAuth = serde_json::from_str(json).unwrap();
assert!(
auth.coding_data_retention_opt_out,
"missing field must default to opted-out"
);
assert!(default_coding_data_retention_opt_out());
assert!(GrokAuth::default().coding_data_retention_opt_out);
}
}

4
crates/codegen/xai-grok-shell/src/auth/oidc/protocol.rs
View file

@ -689,7 +689,7 @@ pub(super) async fn validate_and_extract_user_info(
organization_role: None,
user_blocked_reason: None,
team_blocked_reasons: vec![],
coding_data_retention_opt_out: false,
coding_data_retention_opt_out: crate::auth::default_coding_data_retention_opt_out(),
})
}
pub(super) async fn extract_user_info(
@ -720,7 +720,7 @@ pub(super) async fn extract_user_info(
organization_role: None,
user_blocked_reason: None,
team_blocked_reasons: vec![],
coding_data_retention_opt_out: false,
coding_data_retention_opt_out: crate::auth::default_coding_data_retention_opt_out(),
});
}
let token = id_token.ok_or_else(|| anyhow::Error::new(OidcError::MissingIdToken))?;

18
crates/codegen/xai-grok-shell/src/auth/refresh/external_refresher.rs
View file

@ -61,6 +61,9 @@ impl TokenRefresher for ExternalBinaryRefresher {
.await
{
Err(_elapsed) => {
// Transient: a hard-expired access token after idle must still
// allow 401 / pre-flight retry. Mapping timeout to permanent
// failure poisoned recovery for PERMANENT_FAILURE_TTL.
tracing::warn!(
timeout_ms,
"auth: external binary refresh timed out (thread leaked)"
@ -70,7 +73,7 @@ impl TokenRefresher for ExternalBinaryRefresher {
None,
Some(serde_json::json!({ "timeout_ms": timeout_ms })),
);
self.record_failure(format!("external binary timed out after {timeout_ms}ms"))
RefreshOutcome::transient(format!("external binary timed out after {timeout_ms}ms"))
}
Ok(Ok(Some(auth))) => {
crate::unified_log::info("auth: external binary refresh succeeded", None, None);
@ -131,10 +134,8 @@ mod tests {
}
}
/// A binary that outlives the (test-shortened) timeout hits the `Elapsed`
/// arm and maps to the same non-sticky `Other` permanent failure.
#[tokio::test]
async fn external_binary_timeout_is_non_sticky_permanent() {
async fn external_binary_timeout_is_transient() {
struct SlowRunner;
impl ExternalCommandRunner for SlowRunner {
fn run_external_command(&self, _command: &str) -> Option<GrokAuth> {
@ -145,14 +146,13 @@ mod tests {
let refresher = ExternalBinaryRefresher::new(Arc::new(SlowRunner), "auth-binary".into())
.with_timeout(std::time::Duration::from_millis(5));
match refresher.refresh(RefreshReason::ServerRejected).await {
RefreshOutcome::PermanentFailure { error, .. } => {
assert_eq!(error.reason, RefreshTokenFailedReason::Other);
RefreshOutcome::TransientFailure { message } => {
assert!(
!error.reason.is_sticky(),
"timeout must age out, not strand"
message.contains("timed out"),
"timeout message must be greppable, got {message}"
);
}
other => panic!("a timed-out binary must be a permanent Other failure, got {other:?}"),
other => panic!("a timed-out binary must be TransientFailure, got {other:?}"),
}
}

5
crates/codegen/xai-grok-shell/src/bundle.rs
View file

@ -84,10 +84,7 @@ struct BundleFile<'a> {
}
pub fn bundled_root() -> PathBuf {
dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".grok")
.join(BUNDLED_DIR_NAME)
xai_grok_config::grok_home().join(BUNDLED_DIR_NAME)
}
pub fn read_cached_manifest(root: &Path) -> Result<Option<BundleManifest>> {

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

@ -444,9 +444,23 @@ impl SubagentsConfig {
/// intent (CLI flag, `GROK_SUBAGENTS`, `[subagents] enabled`) changes
/// the default.
///
/// When `cwd` is provided, file-based roles are discovered from
/// `{cwd}/.grok/roles/*.toml` and merged (inline config takes precedence).
pub fn resolve(cli_flag: bool, config: &toml::Value, cwd: Option<&std::path::Path>) -> Self {
/// Project files are excluded from this trust-independent base; Task
/// boundaries overlay them using the parent cwd's authoritative trust verdict.
pub fn resolve(cli_flag: bool, config: &toml::Value) -> Self {
let user_grok_root = xai_grok_config::user_grok_home();
Self::resolve_base_with_sources(
cli_flag,
config,
user_grok_root.as_deref(),
&bundle::bundled_root(),
)
}
pub(crate) fn resolve_base_with_sources(
cli_flag: bool,
config: &toml::Value,
user_grok_root: Option<&std::path::Path>,
bundled_root: &std::path::Path,
) -> Self {
let mut result: Self = config
.get("subagents")
.and_then(|v| v.clone().try_into().ok())
@ -460,19 +474,40 @@ impl SubagentsConfig {
true,
);
result.enabled = resolved.value;
if let Some(cwd) = cwd {
result.discover_roles(cwd);
result.discover_personas(cwd);
if let Some(root) = user_grok_root {
result.discover_roles_in_dir(&root.join("roles"));
result.discover_personas_in_dir(&root.join("personas"));
}
if let Some(home) = dirs::home_dir() {
result.discover_roles(&home);
result.discover_personas(&home);
}
let bundled_root = bundle::bundled_root();
result.discover_roles_in_dir(&bundled_root.join("roles"));
result.discover_personas_in_dir(&bundled_root.join("personas"));
result
}
pub(crate) fn effective_definition_maps(
roles: &std::collections::HashMap<String, SubagentRole>,
personas: &std::collections::HashMap<String, SubagentPersona>,
cwd: &std::path::Path,
project_trusted: bool,
) -> (
std::collections::HashMap<String, SubagentRole>,
std::collections::HashMap<String, SubagentPersona>,
) {
let mut project = Self::default();
if project_trusted {
project.discover_roles(cwd);
project.discover_personas(cwd);
}
for (name, role) in roles {
if role.source_dir.is_none() || !project.roles.contains_key(name) {
project.roles.insert(name.clone(), role.clone());
}
}
for (name, persona) in personas {
if persona.source_path.is_none() || !project.personas.contains_key(name) {
project.personas.insert(name.clone(), persona.clone());
}
}
(project.roles, project.personas)
}
}
/// Managed MCP connector fetching config (`[managed_mcps]` in config.toml).
///

392
crates/codegen/xai-grok-shell/src/config/tests.rs
View file

@ -975,7 +975,7 @@ fn with_grok_subagents<T>(value: &str, f: impl FnOnce() -> T) -> T {
fn subagents_config_default_enabled() {
without_grok_subagents(|| {
let config = toml::Value::Table(toml::map::Map::new());
let sa = SubagentsConfig::resolve(false, &config, None);
let sa = SubagentsConfig::resolve(false, &config);
assert!(sa.enabled);
});
}
@ -983,7 +983,7 @@ fn subagents_config_default_enabled() {
fn subagents_config_cli_flag_enables() {
without_grok_subagents(|| {
let config = toml::Value::Table(toml::map::Map::new());
let sa = SubagentsConfig::resolve(true, &config, None);
let sa = SubagentsConfig::resolve(true, &config);
assert!(sa.enabled);
});
}
@ -993,7 +993,7 @@ fn subagents_config_env_var_enables() {
"1",
|| {
let config = toml::Value::Table(toml::map::Map::new());
let sa = SubagentsConfig::resolve(false, &config, None);
let sa = SubagentsConfig::resolve(false, &config);
assert!(sa.enabled);
},
);
@ -1005,7 +1005,7 @@ fn subagents_config_env_var_disables() {
|| {
let config: toml::Value = toml::from_str("[subagents]\nenabled = true")
.unwrap();
let sa = SubagentsConfig::resolve(false, &config, None);
let sa = SubagentsConfig::resolve(false, &config);
assert!(! sa.enabled, "GROK_SUBAGENTS=0 should override config file");
},
);
@ -1014,7 +1014,7 @@ fn subagents_config_env_var_disables() {
fn subagents_config_toml_enables() {
without_grok_subagents(|| {
let config: toml::Value = toml::from_str("[subagents]\nenabled = true").unwrap();
let sa = SubagentsConfig::resolve(false, &config, None);
let sa = SubagentsConfig::resolve(false, &config);
assert!(sa.enabled);
});
}
@ -1023,7 +1023,7 @@ fn subagents_config_local_disabled_wins() {
without_grok_subagents(|| {
let config: toml::Value = toml::from_str("[subagents]\nenabled = false")
.unwrap();
let sa = SubagentsConfig::resolve(false, &config, None);
let sa = SubagentsConfig::resolve(false, &config);
assert!(! sa.enabled, "local [subagents] enabled=false should win");
});
}
@ -1033,7 +1033,7 @@ fn subagents_config_env_var_disables_default() {
"0",
|| {
let config = toml::Value::Table(toml::map::Map::new());
let sa = SubagentsConfig::resolve(false, &config, None);
let sa = SubagentsConfig::resolve(false, &config);
assert!(
! sa.enabled, "GROK_SUBAGENTS=0 should override the enabled default"
);
@ -1050,7 +1050,7 @@ fn subagents_config_remote_settings_key_is_ignored() {
)
.expect("unknown subagents_enabled key must not break parsing");
let config = toml::Value::Table(toml::map::Map::new());
let sa = SubagentsConfig::resolve(false, &config, None);
let sa = SubagentsConfig::resolve(false, &config);
assert!(sa.enabled);
});
}
@ -1060,7 +1060,7 @@ fn subagents_config_cli_flag_overrides_env_var() {
"0",
|| {
let config = toml::Value::Table(toml::map::Map::new());
let sa = SubagentsConfig::resolve(true, &config, None);
let sa = SubagentsConfig::resolve(true, &config);
assert!(sa.enabled, "--subagents CLI flag should override GROK_SUBAGENTS=0");
},
);
@ -1079,7 +1079,7 @@ fn subagents_config_models_parsed() {
"#,
)
.unwrap();
let sa = SubagentsConfig::resolve(false, &config, None);
let sa = SubagentsConfig::resolve(false, &config);
assert!(sa.enabled);
assert_eq!(sa.models.len(), 2);
assert_eq!(sa.models.get("explore").unwrap(), "grok-3-fast");
@ -1090,7 +1090,7 @@ fn subagents_config_models_parsed() {
fn subagents_config_models_empty_when_missing() {
without_grok_subagents(|| {
let config: toml::Value = toml::from_str("[subagents]\nenabled = true").unwrap();
let sa = SubagentsConfig::resolve(false, &config, None);
let sa = SubagentsConfig::resolve(false, &config);
assert!(sa.enabled);
assert!(sa.models.is_empty());
});
@ -1105,7 +1105,7 @@ fn subagents_config_models_without_enabled() {
"#,
)
.unwrap();
let sa = SubagentsConfig::resolve(false, &config, None);
let sa = SubagentsConfig::resolve(false, &config);
assert!(
! sa.enabled, "explicit [subagents] section without enabled should be false"
);
@ -1125,7 +1125,7 @@ fn subagents_config_models_with_env_var_enables() {
"#,
)
.unwrap();
let sa = SubagentsConfig::resolve(false, &config, None);
let sa = SubagentsConfig::resolve(false, &config);
assert!(sa.enabled, "GROK_SUBAGENTS=1 should enable");
assert_eq!(sa.models.get("explore").unwrap(), "grok-3-fast");
},
@ -1147,7 +1147,7 @@ fn subagents_config_toggle_mixed_values() {
"#,
)
.unwrap();
let sa = SubagentsConfig::resolve(false, &config, None);
let sa = SubagentsConfig::resolve(false, &config);
assert!(sa.enabled);
assert_eq!(sa.toggle.len(), 4);
assert_eq!(sa.toggle.get("explore").copied(), Some(true));
@ -1160,7 +1160,7 @@ fn subagents_config_toggle_mixed_values() {
fn subagents_config_toggle_missing_defaults_to_empty() {
without_grok_subagents(|| {
let config: toml::Value = toml::from_str("[subagents]\nenabled = true").unwrap();
let sa = SubagentsConfig::resolve(false, &config, None);
let sa = SubagentsConfig::resolve(false, &config);
assert!(sa.enabled);
assert!(
sa.toggle.is_empty(),
@ -2254,6 +2254,122 @@ fn discover_personas_inline_takes_precedence() {
Some("Inline strict"),
);
}
fn write_subagent_definitions(root: &std::path::Path, definitions: &[(&str, &str)]) {
let roles = root.join("roles");
let personas = root.join("personas");
std::fs::create_dir_all(&roles).unwrap();
std::fs::create_dir_all(&personas).unwrap();
for (name, source) in definitions {
std::fs::write(
roles.join(format!("{name}.toml")),
format!("description = \"{source} role\""),
)
.unwrap();
std::fs::write(
personas.join(format!("{name}.toml")),
format!("instructions = \"{source} persona\""),
)
.unwrap();
}
}
#[test]
fn project_overlay_preserves_source_precedence() {
let tmp = tempfile::TempDir::new().unwrap();
let project = tmp.path().join("project");
let home = tmp.path().join("home");
let bundled = tmp.path().join("bundled");
write_subagent_definitions(
&project.join(".grok"),
&[
("shadowed", "Project"),
("bundled-shadowed", "Project"),
("inline", "Project"),
("project-only", "Project"),
],
);
write_subagent_definitions(
&home.join(".grok"),
&[("shadowed", "User"), ("user-only", "User")],
);
write_subagent_definitions(
&bundled,
&[("bundled-shadowed", "Bundled"), ("bundled-only", "Bundled")],
);
let config = toml::from_str::<
toml::Value,
>(
r#"
[subagents]
enabled = true
[subagents.roles.inline]
description = "Inline role"
[subagents.personas.inline]
instructions = "Inline persona"
"#,
)
.unwrap();
let base = SubagentsConfig::resolve_base_with_sources(
false,
&config,
Some(&home.join(".grok")),
&bundled,
);
let resolve = |project_trusted| {
let (roles, personas) = SubagentsConfig::effective_definition_maps(
&base.roles,
&base.personas,
&project,
project_trusted,
);
SubagentsConfig {
roles,
personas,
..Default::default()
}
};
let untrusted = resolve(false);
assert_eq!(untrusted.get_role("shadowed").unwrap().description, "User role");
assert_eq!(
untrusted.get_persona("shadowed").and_then(| persona | persona.instructions
.as_deref()), Some("User persona")
);
assert!(untrusted.get_role("project-only").is_none());
assert!(untrusted.get_persona("project-only").is_none());
assert!(untrusted.get_role("user-only").is_some());
assert!(untrusted.get_persona("user-only").is_some());
assert!(untrusted.get_role("bundled-only").is_some());
assert!(untrusted.get_persona("bundled-only").is_some());
assert_eq!(
untrusted.get_role("bundled-shadowed").unwrap().description, "Bundled role"
);
assert_eq!(
untrusted.get_persona("bundled-shadowed").and_then(| persona | persona
.instructions.as_deref()), Some("Bundled persona")
);
let trusted = resolve(true);
assert_eq!(trusted.get_role("shadowed").unwrap().description, "Project role");
assert_eq!(
trusted.get_persona("shadowed").and_then(| persona | persona.instructions
.as_deref()), Some("Project persona")
);
assert_eq!(
trusted.get_role("bundled-shadowed").unwrap().description, "Project role"
);
assert_eq!(
trusted.get_persona("bundled-shadowed").and_then(| persona | persona.instructions
.as_deref()), Some("Project persona")
);
assert_eq!(trusted.get_role("inline").unwrap().description, "Inline role");
assert_eq!(
trusted.get_persona("inline").and_then(| persona | persona.instructions
.as_deref()), Some("Inline persona")
);
let denied_again = resolve(false);
assert_eq!(denied_again.get_role("shadowed").unwrap().description, "User role");
assert!(denied_again.get_role("project-only").is_none());
}
#[test]
fn bundled_personas_and_roles_have_lowest_priority_in_resolve_order() {
let tmp = tempfile::TempDir::new().unwrap();
@ -2296,84 +2412,104 @@ fn bundled_personas_and_roles_have_lowest_priority_in_resolve_order() {
r#"instructions = "Project persona""#,
)
.unwrap();
with_env_var(
"HOME",
home.to_str().unwrap(),
|| {
let config = toml::from_str::<
toml::Value,
>(
r#"
[subagents]
enabled = true
let config = toml::from_str::<
toml::Value,
>(
r#"
[subagents]
enabled = true
[subagents.roles.reviewer]
description = "Inline reviewer"
[subagents.roles.reviewer]
description = "Inline reviewer"
[subagents.personas.reviewer]
instructions = "Inline persona"
"#,
)
.unwrap();
let resolved = SubagentsConfig::resolve(true, &config, Some(&workspace));
assert_eq!(
resolved.get_role("reviewer").unwrap().description, "Inline reviewer"
);
assert_eq!(
resolved.get_persona("reviewer").unwrap().instructions.as_deref(),
Some("Inline persona")
);
},
[subagents.personas.reviewer]
instructions = "Inline persona"
"#,
)
.unwrap();
let base = SubagentsConfig::resolve_base_with_sources(
true,
&config,
Some(&home.join(".grok")),
&bundled,
);
let (roles, personas) = SubagentsConfig::effective_definition_maps(
&base.roles,
&base.personas,
&workspace,
true,
);
let resolved = SubagentsConfig {
roles,
personas,
..Default::default()
};
assert_eq!(resolved.get_role("reviewer").unwrap().description, "Inline reviewer");
assert_eq!(
resolved.get_persona("reviewer").unwrap().instructions.as_deref(),
Some("Inline persona")
);
std::fs::remove_file(workspace.join(".grok/roles/reviewer.toml")).unwrap();
std::fs::remove_file(workspace.join(".grok/personas/reviewer.toml")).unwrap();
with_env_var(
"HOME",
home.to_str().unwrap(),
|| {
let config = toml::from_str::<
toml::Value,
>(
r#"
[subagents]
enabled = true
"#,
)
.unwrap();
let resolved = SubagentsConfig::resolve(true, &config, Some(&workspace));
assert_eq!(
resolved.get_role("reviewer").unwrap().description, "User reviewer"
);
assert_eq!(
resolved.get_persona("reviewer").unwrap().instructions.as_deref(),
Some("User persona")
);
},
let config = toml::from_str::<
toml::Value,
>(r#"
[subagents]
enabled = true
"#)
.unwrap();
let base = SubagentsConfig::resolve_base_with_sources(
true,
&config,
Some(&home.join(".grok")),
&bundled,
);
let (roles, personas) = SubagentsConfig::effective_definition_maps(
&base.roles,
&base.personas,
&workspace,
true,
);
let resolved = SubagentsConfig {
roles,
personas,
..Default::default()
};
assert_eq!(resolved.get_role("reviewer").unwrap().description, "User reviewer");
assert_eq!(
resolved.get_persona("reviewer").unwrap().instructions.as_deref(),
Some("User persona")
);
std::fs::remove_file(home.join(".grok/roles/reviewer.toml")).unwrap();
std::fs::remove_file(home.join(".grok/personas/reviewer.toml")).unwrap();
with_env_var(
"HOME",
home.to_str().unwrap(),
|| {
let config = toml::from_str::<
toml::Value,
>(
r#"
[subagents]
enabled = true
"#,
)
.unwrap();
let resolved = SubagentsConfig::resolve(true, &config, Some(&workspace));
assert_eq!(
resolved.get_role("reviewer").unwrap().description, "Bundled reviewer"
);
assert_eq!(
resolved.get_persona("reviewer").unwrap().instructions.as_deref(),
Some("Bundled persona")
);
},
let config = toml::from_str::<
toml::Value,
>(r#"
[subagents]
enabled = true
"#)
.unwrap();
let base = SubagentsConfig::resolve_base_with_sources(
true,
&config,
Some(&home.join(".grok")),
&bundled,
);
let (roles, personas) = SubagentsConfig::effective_definition_maps(
&base.roles,
&base.personas,
&workspace,
true,
);
let resolved = SubagentsConfig {
roles,
personas,
..Default::default()
};
assert_eq!(resolved.get_role("reviewer").unwrap().description, "Bundled reviewer");
assert_eq!(
resolved.get_persona("reviewer").unwrap().instructions.as_deref(),
Some("Bundled persona")
);
}
#[test]
@ -2952,6 +3088,92 @@ fn managed_settings_does_not_override_user_yolo() {
fn simulate_release_build() -> xai_grok_test_support::EnvGuard {
xai_grok_test_support::EnvGuard::set(xai_grok_version::TEST_VERSION_ENV, "0.0.0-sim")
}
#[test]
fn project_overlay_tracks_authoritative_trust_transitions() {
let source_root = tempfile::tempdir().unwrap();
let repo = tempfile::tempdir().unwrap();
git2::Repository::init(repo.path()).unwrap();
write_subagent_definitions(
&repo.path().join(".grok"),
&[("shared", "Project"), ("project-only", "Project")],
);
let mut base = SubagentsConfig::default();
base.roles
.insert(
"shared".into(),
SubagentRole {
description: "User role".into(),
source_dir: Some(source_root.path().join("roles")),
..Default::default()
},
);
base.personas
.insert(
"shared".into(),
SubagentPersona {
instructions: Some("User persona".into()),
source_path: Some(
source_root.path().join("personas/shared.toml").display().to_string(),
),
..Default::default()
},
);
let (untrusted_roles, _) = SubagentsConfig::effective_definition_maps(
&base.roles,
&base.personas,
repo.path(),
false,
);
assert_eq!(untrusted_roles["shared"].description, "User role");
assert!(! untrusted_roles.contains_key("project-only"));
let (trusted_roles, trusted_personas) = SubagentsConfig::effective_definition_maps(
&base.roles,
&base.personas,
repo.path(),
true,
);
assert_eq!(trusted_roles["shared"].description, "Project role");
assert!(trusted_personas.contains_key("project-only"));
let (revoked_roles, _) = SubagentsConfig::effective_definition_maps(
&base.roles,
&base.personas,
repo.path(),
false,
);
assert_eq!(revoked_roles["shared"].description, "User role");
assert!(! revoked_roles.contains_key("project-only"));
}
#[test]
fn base_resolver_without_project_cwd_keeps_project_files_out() {
let tmp = tempfile::tempdir().unwrap();
write_subagent_definitions(&tmp.path().join(".grok"), &[("project", "Project")]);
let base = SubagentsConfig::resolve_base_with_sources(
false,
&toml::Value::Table(Default::default()),
None,
&tmp.path().join("bundled"),
);
assert!(base.get_role("project").is_none());
assert!(base.get_persona("project").is_none());
}
#[test]
fn explicit_grok_root_is_the_only_user_source() {
let tmp = tempfile::tempdir().unwrap();
let ambient = tmp.path().join("ambient-home/.grok");
let configured = tmp.path().join("configured-grok-home");
write_subagent_definitions(&ambient, &[("ambient", "Ambient")]);
write_subagent_definitions(&configured, &[("configured", "Configured")]);
let base = SubagentsConfig::resolve_base_with_sources(
false,
&toml::Value::Table(Default::default()),
Some(&configured),
&configured.join("bundled"),
);
assert!(base.get_role("ambient").is_none());
assert!(base.get_persona("ambient").is_none());
assert!(base.get_role("configured").is_some());
assert!(base.get_persona("configured").is_some());
}
/// SECURITY (plugin-RCE): a PROJECT-declared `[plugins].paths` loads as an
/// auto-enabled, auto-trusted ConfigPath plugin, so it must merge into the
/// effective config ONLY when the folder is trusted; project

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

@ -242,8 +242,11 @@ fn handle_info(agent: &MvpAgent) -> ExtResult {
.as_ref()
.map(|a| a.team_blocked_reasons.clone())
.unwrap_or_default(),
// No credential ⇒ unknown privacy state: report opted-out (fail closed),
// matching `AuthManager::allows_data_collection` / GrokAuth Default.
coding_data_retention_opt_out: auth
.as_ref()
.is_some_and(|a| a.coding_data_retention_opt_out),
.map(|a| a.coding_data_retention_opt_out)
.unwrap_or_else(crate::auth::default_coding_data_retention_opt_out),
})
}

127
crates/codegen/xai-grok-shell/src/extensions/bundle.rs
View file

@ -432,28 +432,6 @@ mod tests {
use serial_test::serial;
use std::sync::{Arc, Mutex};
use tempfile::TempDir;
struct HomeGuard {
previous: Option<std::ffi::OsString>,
}
impl Drop for HomeGuard {
fn drop(&mut self) {
match self.previous.take() {
Some(previous) => unsafe {
std::env::set_var("HOME", previous);
},
None => unsafe {
std::env::remove_var("HOME");
},
}
}
}
fn with_bundled_home(tmp: &TempDir) -> HomeGuard {
let previous = std::env::var_os("HOME");
unsafe {
std::env::set_var("HOME", tmp.path());
}
HomeGuard { previous }
}
fn sample_bundle() -> SubagentBundle {
let mut bundle = SubagentBundle::empty("bundle-v1");
bundle.personas.insert(
@ -588,8 +566,8 @@ mod tests {
#[serial]
fn status_reports_no_cache_when_manifest_missing() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let status = status_bundle_at(&bundle::bundled_root()).unwrap();
let root = tmp.path().join("bundled");
let status = status_bundle_at(&root).unwrap();
assert_eq!(
status,
BundleStatusResult {
@ -608,15 +586,14 @@ mod tests {
#[serial]
fn status_reports_cached_entries_from_manifest_and_disk() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = bundle::bundled_root();
let root = tmp.path().join("bundled");
bundle::write_bundle_to_cache(&root, &sample_bundle()).unwrap();
std::fs::write(
root.join("personas/local-only.toml"),
"instructions = \"ignore\"",
)
.unwrap();
let status = status_bundle_at(&bundle::bundled_root()).unwrap();
let status = status_bundle_at(&root).unwrap();
assert!(status.has_cache);
assert_eq!(status.version.as_deref(), Some("bundle-v1"));
assert_eq!(status.personas, vec!["researcher"]);
@ -628,14 +605,13 @@ mod tests {
#[serial]
async fn sync_success_writes_cache_and_returns_counts() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = tmp.path().join("bundled");
let bundle = sample_bundle();
let (proxy_base_url, _seen_headers, server) = start_bundle_server(
axum::http::StatusCode::OK,
serde_json::to_value(&bundle).unwrap(),
)
.await;
let root = bundle::bundled_root();
let am = test_auth_manager();
let result = sync_bundle_to_root(&root, &proxy_base_url, Some(&am), None, None, false)
.await
@ -654,11 +630,10 @@ mod tests {
#[serial]
async fn sync_force_true_has_same_write_semantics() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = tmp.path().join("bundled");
let bundle = sample_bundle();
let (proxy_base_url, _seen_headers, server) =
start_bundle_server(StatusCode::OK, serde_json::to_value(&bundle).unwrap()).await;
let root = bundle::bundled_root();
let am = test_auth_manager();
let normal = sync_bundle_to_root(&root, &proxy_base_url, Some(&am), None, None, false)
.await
@ -673,23 +648,16 @@ mod tests {
#[serial]
async fn sync_http_failure_surfaces_error() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = tmp.path().join("bundled");
let (proxy_base_url, _seen_headers, server) = start_bundle_server(
StatusCode::UNAUTHORIZED,
serde_json::json!({ "error" : "unauthorized" }),
)
.await;
let am = test_auth_manager();
let error = sync_bundle_to_root(
&bundle::bundled_root(),
&proxy_base_url,
Some(&am),
None,
None,
false,
)
.await
.unwrap_err();
let error = sync_bundle_to_root(&root, &proxy_base_url, Some(&am), None, None, false)
.await
.unwrap_err();
assert!(error.to_string().contains("401"));
server.abort();
}
@ -697,13 +665,13 @@ mod tests {
#[serial]
async fn sync_uses_deployment_key_auth_mode() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = tmp.path().join("bundled");
let bundle = sample_bundle();
let (proxy_base_url, seen_headers, server) =
start_bundle_server(StatusCode::OK, serde_json::to_value(&bundle).unwrap()).await;
let am = test_auth_manager();
let result = sync_bundle_to_root(
&bundle::bundled_root(),
&root,
&proxy_base_url,
Some(&am),
Some("deploy-key"),
@ -726,8 +694,7 @@ mod tests {
#[serial]
fn status_only_reports_bundled_cache_not_higher_priority_sources() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = bundle::bundled_root();
let root = tmp.path().join("bundled");
bundle::write_bundle_to_cache(&root, &sample_bundle()).unwrap();
let project_root = tmp.path().join("workspace");
std::fs::create_dir_all(project_root.join(".grok/personas")).unwrap();
@ -742,23 +709,26 @@ mod tests {
"description = \"project role\"\n",
)
.unwrap();
let config = crate::config::SubagentsConfig::resolve(
let base = crate::config::SubagentsConfig::resolve_base_with_sources(
false,
&toml::Value::Table(Default::default()),
Some(&project_root),
None,
&root,
);
let (roles, personas) = crate::config::SubagentsConfig::effective_definition_maps(
&base.roles,
&base.personas,
&project_root,
true,
);
assert_eq!(
config
.personas
personas
.get("researcher")
.and_then(|persona| persona.instructions.as_deref()),
Some("project persona")
);
assert_eq!(
config
.roles
.get("reviewer")
.map(|role| role.description.as_str()),
roles.get("reviewer").map(|role| role.description.as_str()),
Some("project role")
);
let status = status_bundle_at(&root).unwrap();
@ -771,9 +741,9 @@ mod tests {
#[serial]
fn sync_requires_auth_or_deployment_key() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = tmp.path().join("bundled");
let error = futures::executor::block_on(sync_bundle_to_root(
&bundle::bundled_root(),
&root,
"http://127.0.0.1:1/v1",
None,
None,
@ -790,8 +760,7 @@ mod tests {
#[serial]
fn get_entry_reads_persona_file() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = bundle::bundled_root();
let root = tmp.path().join("bundled");
bundle::write_bundle_to_cache(&root, &sample_bundle()).unwrap();
let result = get_entry_at(&root, "persona", "researcher").unwrap();
assert_eq!(result.kind, "persona");
@ -802,8 +771,7 @@ mod tests {
#[serial]
fn get_entry_unknown_kind_returns_error() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = bundle::bundled_root();
let root = tmp.path().join("bundled");
let err = get_entry_at(&root, "widget", "foo").unwrap_err();
assert!(err.to_string().contains("unknown entry kind: widget"));
}
@ -811,8 +779,7 @@ mod tests {
#[serial]
fn get_entry_missing_file_returns_error() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = bundle::bundled_root();
let root = tmp.path().join("bundled");
bundle::write_bundle_to_cache(&root, &sample_bundle()).unwrap();
let err = get_entry_at(&root, "persona", "nonexistent").unwrap_err();
assert!(err.to_string().contains("not found in bundle cache"));
@ -840,8 +807,7 @@ mod tests {
#[serial]
fn status_includes_persona_and_role_details() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = bundle::bundled_root();
let root = tmp.path().join("bundled");
bundle::write_bundle_to_cache(&root, &sample_bundle()).unwrap();
let status = status_bundle_at(&root).unwrap();
assert_eq!(status.persona_details.len(), 1);
@ -862,8 +828,7 @@ mod tests {
#[serial]
fn status_without_toml_files_returns_empty_details() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = bundle::bundled_root();
let root = tmp.path().join("bundled");
let mut bundle = SubagentBundle::empty("v1");
bundle
.personas
@ -912,12 +877,12 @@ mod tests {
#[serial]
async fn sync_with_skills_reports_skills_count() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = tmp.path().join("bundled");
let bundle = sample_bundle_with_skills();
let (proxy_base_url, _seen_headers, server) =
start_bundle_server(StatusCode::OK, serde_json::to_value(&bundle).unwrap()).await;
let result = sync_bundle_to_root(
&bundle::bundled_root(),
&root,
&proxy_base_url,
Some(&test_auth_manager()),
None,
@ -936,8 +901,7 @@ mod tests {
#[serial]
fn status_lists_skill_names_from_manifest() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = bundle::bundled_root();
let root = tmp.path().join("bundled");
bundle::write_bundle_to_cache(&root, &sample_bundle_with_skills()).unwrap();
let status = status_bundle_at(&root).unwrap();
assert!(status.has_cache);
@ -948,8 +912,7 @@ mod tests {
#[serial]
fn status_skills_only_lists_files_present_on_disk() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = bundle::bundled_root();
let root = tmp.path().join("bundled");
bundle::write_bundle_to_cache(&root, &sample_bundle_with_skills()).unwrap();
std::fs::remove_file(root.join("skills/commit/SKILL.md")).unwrap();
let status = status_bundle_at(&root).unwrap();
@ -981,7 +944,7 @@ mod tests {
#[serial]
async fn sync_with_archive_endpoint_extracts_and_reports_counts() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = tmp.path().join("bundled");
let archive = make_test_archive(&[
("bundle.json", br#"{"version":"archive-v1"}"#),
(
@ -993,7 +956,7 @@ mod tests {
]);
let (proxy_base_url, server) = start_archive_bundle_server(archive).await;
let result = sync_bundle_to_root(
&bundle::bundled_root(),
&root,
&proxy_base_url,
Some(&test_auth_manager()),
None,
@ -1007,28 +970,20 @@ mod tests {
assert_eq!(result.roles_count, 1);
assert_eq!(result.agents_count, 0);
assert_eq!(result.skills_count, 1);
assert!(
bundle::bundled_root()
.join("personas/researcher.toml")
.exists()
);
assert!(
bundle::bundled_root()
.join("skills/commit/SKILL.md")
.exists()
);
assert!(root.join("personas/researcher.toml").exists());
assert!(root.join("skills/commit/SKILL.md").exists());
server.abort();
}
#[tokio::test(flavor = "current_thread")]
#[serial]
async fn sync_falls_back_to_legacy_when_archive_unavailable() {
let tmp = TempDir::new().unwrap();
let _home = with_bundled_home(&tmp);
let root = tmp.path().join("bundled");
let bundle = sample_bundle_with_skills();
let (proxy_base_url, _seen_headers, server) =
start_bundle_server(StatusCode::OK, serde_json::to_value(&bundle).unwrap()).await;
let result = sync_bundle_to_root(
&bundle::bundled_root(),
&root,
&proxy_base_url,
Some(&test_auth_manager()),
None,

8
crates/codegen/xai-grok-shell/src/extensions/debug.rs
View file

@ -7,6 +7,7 @@
//! - `trigger_feedback`: fire a synthetic `FeedbackRequestNotification`.
//! - `arm_auto_compact`: arm the next turn to unconditionally trigger
//! auto-compaction, regardless of context window usage.
//! - `agent`: agent-process diagnostics (registry counts).
use agent_client_protocol as acp;
@ -22,10 +23,17 @@ pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
handle_trigger_feedback(agent, args).await
}
"x.ai/debug/arm_auto_compact" => handle_arm_auto_compact(agent, args),
"x.ai/debug/agent" => handle_agent(agent),
_ => Err(acp::Error::method_not_found()),
}
}
fn handle_agent(agent: &MvpAgent) -> ExtResult {
ExtMethodResult::success(serde_json::json!({ "registries": agent.registry_snapshot() }))
.to_ext_response()
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
}
async fn handle_trigger_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
use crate::session::feedback::{FeedbackMode, FeedbackTier};

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

@ -19,13 +19,13 @@ fn load_filtered_marketplace_sources() -> Vec<xai_grok_plugin_marketplace::Marke
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
match args.method.as_ref() {
"x.ai/marketplace/list" => handle_list(agent, args).await,
"x.ai/marketplace/list" => handle_list().await,
"x.ai/marketplace/action" => handle_action(agent, args).await,
_ => Err(acp::Error::method_not_found()),
}
}
async fn handle_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
async fn handle_list() -> ExtResult {
let t0 = std::time::Instant::now();
let sources = load_filtered_marketplace_sources();
let source_names: Vec<String> = sources
@ -94,15 +94,10 @@ async fn handle_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
results.push(scan);
}
// Auto-install default-skills entries that aren't already installed.
let t_auto = std::time::Instant::now();
let session_id = super::parse_session_id(args);
auto_install_defaults(agent, &sources, &results, session_id.as_ref(), false).await;
xai_grok_telemetry::unified_log::info(
"marketplace handle_list: complete",
None,
Some(serde_json::json!({
"auto_install_ms": t_auto.elapsed().as_millis() as u64,
"total_ms": t0.elapsed().as_millis() as u64,
})),
);
@ -146,10 +141,6 @@ async fn handle_action(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
refreshed += 1;
}
// Refresh default-skills from all sources.
let scan_results: Vec<_> = sources.iter().map(|source| scan_source(source).0).collect();
auto_install_defaults(agent, &sources, &scan_results, Some(&sid), true).await;
let msg = if errors.is_empty() {
format!("Refreshed {refreshed} source(s).")
} else {
@ -635,115 +626,6 @@ async fn handle_uninstall(
}
}
/// Auto-install or refresh `default-skills` entries from marketplace sources.
///
/// When `force_refresh` is false, only installs if not already present.
/// When `force_refresh` is true, removes and re-copies from source.
async fn auto_install_defaults(
agent: &MvpAgent,
sources: &[xai_grok_plugin_marketplace::MarketplaceSource],
results: &[MarketplaceScanResult],
session_id: Option<&acp::SessionId>,
force_refresh: bool,
) {
use xai_grok_plugin_marketplace::installer;
let mut any_changed = false;
for (source, scan) in sources.iter().zip(results.iter()) {
// Find the default-skills entry.
let default_entry = scan
.plugins
.iter()
.find(|p| p.relative_path == "default-skills");
let Some(entry) = default_entry else {
continue;
};
// Skip if no components.
if entry.skill_count == 0 && !entry.has_hooks && !entry.has_agents && !entry.has_mcp {
continue;
}
// Resolve marketplace root.
let marketplace_lease;
let marketplace_root = match &source.kind {
xai_grok_plugin_marketplace::SourceKind::Local { path } => {
marketplace_lease = None;
path.clone()
}
xai_grok_plugin_marketplace::SourceKind::Git { url, branch } => {
let cache_root = xai_grok_plugin_marketplace::git::default_cache_root();
match xai_grok_plugin_marketplace::git::sync_source_cache_with_mode(
url,
branch.as_deref(),
&cache_root,
if force_refresh {
xai_grok_plugin_marketplace::git::SyncMode::Force
} else {
xai_grok_plugin_marketplace::git::SyncMode::UseTtl
},
) {
Ok(lease) => {
let cached_path = lease.path.clone();
marketplace_lease = Some(lease);
cached_path
}
Err(_) => continue,
}
}
};
// Check if already installed.
let mut reg = xai_grok_agent::plugins::install_registry::InstallRegistry::load();
let existing = installer::find_installed_marketplace_plugin(
&reg,
&scan.source_url_or_path,
"default-skills",
);
if existing.is_some() && !force_refresh {
// Already installed and not forcing refresh — skip.
continue;
}
// Remove old copy if present (refresh or reinstall).
if let Some((existing_key, _)) = existing {
let old_dir = reg.install_dir().join(&existing_key);
let _ = std::fs::remove_dir_all(&old_dir);
reg.remove(&existing_key);
let _ = reg.save();
reg = xai_grok_agent::plugins::install_registry::InstallRegistry::load();
}
let provenance = xai_grok_agent::plugins::install_registry::MarketplaceProvenance {
source_url_or_path: scan.source_url_or_path.clone(),
source_display_name: source.name.clone(),
plugin_subdir: "default-skills".to_string(),
};
let install_result = installer::install_from_marketplace(
&marketplace_root,
"default-skills",
provenance,
&mut reg,
);
drop(marketplace_lease);
if let Ok(installer::MarketplaceInstallResult::Installed { repo_key }) = install_result {
tracing::info!(
source = %source.name,
repo_key = %repo_key,
"auto-installed/refreshed default-skills from marketplace"
);
any_changed = true;
}
}
// Trigger plugin reload if we auto-installed anything.
if any_changed && let Some(sid) = session_id {
let _ = agent
.execute_plugins_action(sid, xai_hooks_plugins_types::PluginsAction::Reload)
.await;
}
}
fn scan_source(
source: &xai_grok_plugin_marketplace::MarketplaceSource,
) -> (MarketplaceScanResult, bool) {
@ -1209,9 +1091,7 @@ fn remove_source_locked(source_url_or_path: &str) -> xai_hooks_plugins_types::Ac
}
}
/// Set the official auto-installed flag on a TOML document string (pure, no I/O)
/// so callers can fold it into a single atomic write. Preserves formatting.
fn set_official_flag_in_toml(content: &str) -> std::io::Result<String> {
fn set_marketplace_bool_flag_in_toml(content: &str, key: &str) -> std::io::Result<String> {
let mut doc = content.parse::<toml_edit::DocumentMut>().map_err(|e| {
std::io::Error::new(
std::io::ErrorKind::InvalidData,
@ -1229,23 +1109,21 @@ fn set_official_flag_in_toml(content: &str) -> std::io::Result<String> {
"[marketplace] is not a table",
)
})?;
marketplace["official_marketplace_auto_installed"] = toml_edit::value(true);
marketplace[key] = toml_edit::value(true);
Ok(doc.to_string())
}
/// Set the official auto-installed flag in `config.toml` (atomic write).
fn set_official_marketplace_auto_installed(config_path: &std::path::Path) -> std::io::Result<()> {
fn set_marketplace_bool_flag(config_path: &std::path::Path, key: &str) -> std::io::Result<()> {
if let Some(parent) = config_path.parent() {
let _ = std::fs::create_dir_all(parent);
}
let existing = crate::util::config::read_to_string_or_empty(config_path)?;
let updated = set_official_flag_in_toml(&existing)?;
let updated = set_marketplace_bool_flag_in_toml(&existing, key)?;
crate::util::config::atomic_write_string(config_path, &updated)
}
/// Read the `official_marketplace_auto_installed` flag; `false` on any failure.
fn read_official_marketplace_auto_installed(config_path: &std::path::Path) -> bool {
fn read_marketplace_bool_flag(config_path: &std::path::Path, key: &str) -> bool {
let raw = match std::fs::read_to_string(config_path) {
Ok(s) => s,
Err(_) => return false,
@ -1256,11 +1134,23 @@ fn read_official_marketplace_auto_installed(config_path: &std::path::Path) -> bo
};
parsed
.get("marketplace")
.and_then(|m| m.get("official_marketplace_auto_installed"))
.and_then(|m| m.get(key))
.and_then(|v| v.as_bool())
.unwrap_or(false)
}
fn set_official_flag_in_toml(content: &str) -> std::io::Result<String> {
set_marketplace_bool_flag_in_toml(content, "official_marketplace_auto_installed")
}
fn set_official_marketplace_auto_installed(config_path: &std::path::Path) -> std::io::Result<()> {
set_marketplace_bool_flag(config_path, "official_marketplace_auto_installed")
}
fn read_official_marketplace_auto_installed(config_path: &std::path::Path) -> bool {
read_marketplace_bool_flag(config_path, "official_marketplace_auto_installed")
}
/// Acquire an advisory exclusive `flock` on `<grok_home>/.config-init.lock`,
/// retrying briefly under contention, to serialize first-run auto-register
/// across processes. Only `WouldBlock` retries; other I/O errors return early.
@ -1290,6 +1180,134 @@ fn acquire_init_lock(grok_home: &std::path::Path) -> std::io::Result<std::fs::Fi
))
}
fn is_default_skills_plugin_subdir(plugin_subdir: &str) -> bool {
plugin_subdir == "default-skills"
}
fn default_skills_repo_keys<'a>(
repos: impl IntoIterator<
Item = (
&'a str,
&'a xai_grok_agent::plugins::install_registry::InstalledRepo,
),
>,
) -> Vec<&'a str> {
repos
.into_iter()
.filter_map(|(key, repo)| {
repo.marketplace
.as_ref()
.filter(|mp| is_default_skills_plugin_subdir(&mp.plugin_subdir))
.map(|_| key)
})
.collect()
}
fn set_default_skills_installs_purged(config_path: &std::path::Path) -> std::io::Result<()> {
set_marketplace_bool_flag(config_path, "default_skills_installs_purged")
}
fn read_default_skills_installs_purged(config_path: &std::path::Path) -> bool {
read_marketplace_bool_flag(config_path, "default_skills_installs_purged")
}
/// One-shot purge of legacy marketplace `default-skills` installs.
///
/// Gated by sticky `default_skills_installs_purged` in config.toml. Best-effort:
/// errors are logged and never block startup.
pub fn purge_default_skills_installs(grok_home: &std::path::Path) {
purge_default_skills_installs_impl(grok_home, || {
xai_grok_agent::plugins::install_registry::InstallRegistry::try_load_from(
xai_grok_agent::plugins::install_registry::InstallRegistry::resolve_install_dir(),
)
});
}
fn purge_default_skills_installs_impl(
grok_home: &std::path::Path,
load_registry: impl FnOnce() -> Result<
xai_grok_agent::plugins::install_registry::InstallRegistry,
xai_grok_agent::plugins::install_registry::InstallError,
>,
) {
let config_path = grok_home.join("config.toml");
if read_default_skills_installs_purged(&config_path) {
return;
}
let _lock = match acquire_init_lock(grok_home) {
Ok(f) => f,
Err(e) => {
tracing::warn!(
error = %e,
path = %grok_home.join(".config-init.lock").display(),
"skipping default-skills purge: failed to acquire init lock"
);
return;
}
};
if read_default_skills_installs_purged(&config_path) {
return;
}
let mut registry = match load_registry() {
Ok(reg) => reg,
Err(e) => {
tracing::warn!(
error = %e,
"skipping default-skills purge: failed to load install registry"
);
return;
}
};
let keys: Vec<String> = default_skills_repo_keys(registry.list())
.into_iter()
.map(|k| k.to_string())
.collect();
for key in &keys {
let path = registry
.get_repo(key)
.map(|r| r.path.clone())
.unwrap_or_else(|| registry.install_dir().join(key));
if path.exists()
&& let Err(e) = std::fs::remove_dir_all(&path)
{
let _ = std::fs::remove_file(&path);
if path.exists() {
tracing::warn!(
error = %e,
path = %path.display(),
repo_key = %key,
"failed to remove default-skills install dir"
);
}
}
registry.remove(key);
}
if !keys.is_empty() {
if let Err(e) = registry.save() {
tracing::warn!(error = %e, "failed to save registry after default-skills purge");
return;
}
tracing::info!(
count = keys.len(),
"purged legacy default-skills marketplace installs"
);
}
if let Err(e) = set_default_skills_installs_purged(&config_path) {
tracing::warn!(
error = %e,
path = %config_path.display(),
"failed to set default_skills_installs_purged flag"
);
}
}
/// Auto-register the official xAI marketplace source on first run.
///
/// Gated by the caller (`init_process`); see
@ -1671,6 +1689,184 @@ mod official_source_tests {
}
}
#[cfg(test)]
mod default_skills_purge_tests {
use super::*;
use xai_grok_agent::plugins::install_registry::{
InstallKind, InstallRegistry, InstalledRepo, MarketplaceProvenance, RepoPlugin,
};
fn repo_at(path: &std::path::Path, plugin_subdir: Option<&str>) -> InstalledRepo {
InstalledRepo {
kind: InstallKind::Local {
source_path: path.to_path_buf(),
subdir: None,
},
installed_at: String::new(),
updated_at: String::new(),
path: path.to_path_buf(),
plugins: std::collections::HashMap::from([(
"p".into(),
RepoPlugin {
subdir: None,
version: None,
},
)]),
marketplace: plugin_subdir.map(|subdir| MarketplaceProvenance {
source_url_or_path: "https://example.com/market.git".into(),
source_display_name: "Test".into(),
plugin_subdir: subdir.into(),
}),
}
}
#[test]
fn match_is_exact_plugin_subdir_only() {
assert!(is_default_skills_plugin_subdir("default-skills"));
assert!(!is_default_skills_plugin_subdir("plugins/default-skills"));
assert!(!is_default_skills_plugin_subdir("default-skills/extra"));
assert!(!is_default_skills_plugin_subdir("defaults-skills"));
assert!(!is_default_skills_plugin_subdir(""));
}
#[test]
fn collects_only_default_skills_repo_keys() {
let default_skills = repo_at(std::path::Path::new("/tmp/ds"), Some("default-skills"));
let other = repo_at(std::path::Path::new("/tmp/office"), Some("plugins/office"));
let no_marketplace = repo_at(std::path::Path::new("/tmp/local"), None);
let keys = default_skills_repo_keys([
("ds-aaaa", &default_skills),
("office-bbbb", &other),
("local-cccc", &no_marketplace),
]);
assert_eq!(keys, vec!["ds-aaaa"]);
}
#[test]
fn purged_flag_toml_preserves_other_content() {
let content =
"[ui]\ntheme = \"dark\"\n[marketplace]\nofficial_marketplace_auto_installed = true\n";
let out =
set_marketplace_bool_flag_in_toml(content, "default_skills_installs_purged").unwrap();
assert!(out.contains("theme = \"dark\""), "{out}");
assert!(
out.contains("official_marketplace_auto_installed = true"),
"{out}"
);
assert!(
out.contains("default_skills_installs_purged = true"),
"{out}"
);
}
#[test]
fn read_purged_flag_false_when_missing_or_wrong_type() {
let tmp = tempfile::tempdir().unwrap();
let path = tmp.path().join("config.toml");
assert!(!read_default_skills_installs_purged(&path));
std::fs::write(
&path,
"[marketplace]\ndefault_skills_installs_purged = \"yes\"\n",
)
.unwrap();
assert!(!read_default_skills_installs_purged(&path));
std::fs::write(
&path,
"[marketplace]\ndefault_skills_installs_purged = true\n",
)
.unwrap();
assert!(read_default_skills_installs_purged(&path));
}
#[test]
fn purge_sets_flag_when_nothing_to_delete() {
let tmp = tempfile::tempdir().unwrap();
let home = tmp.path();
let install_dir = home.join("installed-plugins");
purge_default_skills_installs_impl(home, || {
Ok(InstallRegistry::empty(install_dir.clone()))
});
let config_path = home.join("config.toml");
assert!(read_default_skills_installs_purged(&config_path));
let after_first = std::fs::read_to_string(&config_path).unwrap();
purge_default_skills_installs_impl(home, || Ok(InstallRegistry::empty(install_dir)));
let after_second = std::fs::read_to_string(&config_path).unwrap();
assert_eq!(after_first, after_second);
}
#[test]
fn purge_skips_flag_when_registry_load_fails() {
let tmp = tempfile::tempdir().unwrap();
let home = tmp.path();
let install_dir = home.join("installed-plugins");
std::fs::create_dir_all(&install_dir).unwrap();
std::fs::write(install_dir.join("registry.json"), "{not-json").unwrap();
purge_default_skills_installs_impl(home, || {
InstallRegistry::try_load_from(install_dir.clone())
});
assert!(!read_default_skills_installs_purged(
&home.join("config.toml")
));
}
#[test]
fn purge_removes_default_skills_retains_others_and_sets_flag() {
let tmp = tempfile::tempdir().unwrap();
let home = tmp.path();
let install_dir = home.join("installed-plugins");
std::fs::create_dir_all(&install_dir).unwrap();
let ds_path = install_dir.join("ds-aaaa");
std::fs::create_dir_all(&ds_path).unwrap();
std::fs::write(ds_path.join("marker"), "ds").unwrap();
let other_path = install_dir.join("office-bbbb");
std::fs::create_dir_all(&other_path).unwrap();
std::fs::write(other_path.join("marker"), "office").unwrap();
let mut registry = InstallRegistry::empty(install_dir.clone());
registry.insert("ds-aaaa".into(), repo_at(&ds_path, Some("default-skills")));
registry.insert(
"office-bbbb".into(),
repo_at(&other_path, Some("plugins/office")),
);
registry.save().unwrap();
let install_dir_for_load = install_dir.clone();
purge_default_skills_installs_impl(home, move || {
InstallRegistry::try_load_from(install_dir_for_load)
});
assert!(
!ds_path.exists(),
"default-skills install dir must be removed"
);
assert!(other_path.exists(), "non-matching install must be retained");
let reloaded = InstallRegistry::load_from(install_dir.clone());
assert!(reloaded.get_repo("ds-aaaa").is_none());
assert!(reloaded.get_repo("office-bbbb").is_some());
let config_path = home.join("config.toml");
assert!(read_default_skills_installs_purged(&config_path));
let after_first = std::fs::read_to_string(&config_path).unwrap();
let install_dir_for_reload = install_dir;
purge_default_skills_installs_impl(home, move || {
InstallRegistry::try_load_from(install_dir_for_reload)
});
let after_second = std::fs::read_to_string(&config_path).unwrap();
assert_eq!(after_first, after_second);
assert!(other_path.exists());
}
}
#[cfg(test)]
mod conversion_tests {
use super::*;

16
crates/codegen/xai-grok-shell/src/extensions/notification.rs
View file

@ -543,10 +543,10 @@ pub enum SessionUpdate {
/// Task completed notification
TaskCompleted {
task_snapshot: TaskSnapshot,
/// Whether an auto-wake prompt follows this completion. The pager
/// skips its between-turns status line when set — the wake turn's
/// end marker carries the fresh counts instead. Missing (old
/// shells) reads as `false`: emit the line.
/// Advisory: an auto-wake prompt follows this completion. The
/// first-party TUI no longer consumes it (remaining background work
/// is surfaced by its persistent "watching" status row); kept for
/// wire compatibility and other clients. Missing reads as `false`.
#[serde(default)]
will_wake: bool,
},
@ -648,10 +648,10 @@ pub enum SessionUpdate {
/// Final output text from the subagent (if completed).
#[serde(default, skip_serializing_if = "Option::is_none")]
output: Option<String>,
/// Whether an auto-wake prompt follows this completion. The pager
/// skips its between-turns status line when set — the wake turn's
/// end marker carries the fresh counts instead. Missing (old
/// shells) reads as `false`: emit the line.
/// Advisory: an auto-wake prompt follows this completion. The
/// first-party TUI no longer consumes it (remaining background work
/// is surfaced by its persistent "watching" status row); kept for
/// wire compatibility and other clients. Missing reads as `false`.
#[serde(default)]
will_wake: bool,
},

25
crates/codegen/xai-grok-shell/src/extensions/pr.rs
View file

@ -84,13 +84,21 @@ async fn gh_pr_view_by_branch(cwd: &str, branch: &str) -> Option<PrData> {
.stdin(std::process::Stdio::null());
xai_grok_tools::util::detach_command(&mut cmd);
cmd.envs(xai_grok_tools::util::pager_env());
// gh colorizes even piped --json output under CLICOLOR_FORCE or
// GH_FORCE_TTY (inherited from terminal-launched dev environments), and
// forcing beats NO_COLOR in gh's precedence; there is no --no-color flag
// (cli/cli#9436). CLICOLOR_FORCE=0 is gh's documented off-switch.
cmd.env("NO_COLOR", "1");
cmd.env("CLICOLOR_FORCE", "0");
cmd.env_remove("GH_FORCE_TTY");
let output = cmd.output().await.ok()?;
if !output.status.success() {
return None;
}
let parsed = serde_json::from_slice::<GhPrViewResponse>(&output.stdout).ok()?;
let parsed =
serde_json::from_slice::<GhPrViewResponse>(&strip_ansi_csi(&output.stdout)).ok()?;
let url = parsed.url?;
let state = match parsed
.state
@ -131,7 +139,10 @@ async fn gh_pr_is_in_merge_queue(cwd: &str, pr_url: &str) -> bool {
.stdin(std::process::Stdio::null());
xai_grok_tools::util::detach_command(&mut cmd);
cmd.envs(xai_grok_tools::util::pager_env());
// Forcing (CLICOLOR_FORCE/GH_FORCE_TTY) beats NO_COLOR in gh's precedence.
cmd.env("NO_COLOR", "1");
cmd.env("CLICOLOR_FORCE", "0");
cmd.env_remove("GH_FORCE_TTY");
let output = match cmd.output().await {
Ok(output) => output,
Err(_) => return false,
@ -187,6 +198,18 @@ fn strip_ansi_csi(bytes: &[u8]) -> Vec<u8> {
mod tests {
use super::*;
#[test]
fn gh_pr_view_json_parses_after_stripping_forced_color() {
let stdout = b"\x1b[1;37m{\x1b[m\n \x1b[1;34m\"isDraft\"\x1b[m\x1b[1;37m:\x1b[m \x1b[33mfalse\x1b[m\x1b[1;37m,\x1b[m\n \x1b[1;34m\"number\"\x1b[m\x1b[1;37m:\x1b[m 242682\x1b[1;37m,\x1b[m\n \x1b[1;34m\"state\"\x1b[m\x1b[1;37m:\x1b[m \x1b[32m\"OPEN\"\x1b[m\x1b[1;37m,\x1b[m\n \x1b[1;34m\"title\"\x1b[m\x1b[1;37m:\x1b[m \x1b[32m\"t\"\x1b[m\x1b[1;37m,\x1b[m\n \x1b[1;34m\"url\"\x1b[m\x1b[1;37m:\x1b[m \x1b[32m\"https://github.com/xai-org/xai/pull/242682\"\x1b[m\n\x1b[1;37m}\x1b[m\n";
let parsed = serde_json::from_slice::<GhPrViewResponse>(&strip_ansi_csi(stdout)).unwrap();
assert_eq!(parsed.number, Some(242682));
assert_eq!(parsed.state.as_deref(), Some("OPEN"));
assert_eq!(
parsed.url.as_deref(),
Some("https://github.com/xai-org/xai/pull/242682")
);
}
#[test]
fn parse_is_in_merge_queue_true() {
let stdout = br#"{"data":{"resource":{"isInMergeQueue":true}}}"#;

154
crates/codegen/xai-grok-shell/src/inspect/mod.rs
View file

@ -24,7 +24,7 @@ use xai_grok_tools::util::truncate::estimate_tokens;
const TREE: &str = "\u{2514}";
/// Coarse scope label for project instructions and plugin entries.
#[derive(Debug, Clone, Copy, Serialize)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum Scope {
Project,
@ -440,13 +440,35 @@ fn has_rules_directory(file_path: &str, config_dir: &str) -> bool {
false
}
fn instruction_scope(
file_path: &str,
grok_home: &Path,
vendor_homes: &[(PathBuf, bool)],
workspace_root: &Path,
) -> Scope {
if crate::util::is_user_instruction_path(
Path::new(file_path),
grok_home,
vendor_homes,
Some(workspace_root),
) {
Scope::Global
} else {
Scope::Project
}
}
fn instruction_file_type(
file_path: &str,
grok_home: &Path,
claude_imported: bool,
extra_rule_prefixes: &[PathBuf],
) -> &'static str {
let path = Path::new(file_path);
if has_rules_directory(file_path, ".grok")
if path
.parent()
.is_some_and(|parent| parent == grok_home.join("rules"))
|| has_rules_directory(file_path, ".grok")
|| has_rules_directory(file_path, ".cursor")
|| (!claude_imported && has_rules_directory(file_path, ".claude"))
|| extra_rule_prefixes
@ -468,7 +490,19 @@ async fn list_instructions(cwd: &Path) -> Vec<InstructionFile> {
)
.await;
let grok_home = Some(crate::util::grok_home::grok_home());
let grok_home = crate::util::grok_home::grok_home();
let vendor_homes = dirs::home_dir()
.map(|home_dir| {
vec![
(home_dir.join(".claude"), true),
(home_dir.join(".cursor"), true),
]
})
.unwrap_or_default();
let workspace_root = git2::Repository::discover(cwd)
.ok()
.and_then(|repo| repo.workdir().map(Path::to_path_buf))
.unwrap_or_else(|| cwd.to_path_buf());
// Phase 2 cutoff: when imported, stop classifying `.claude/rules/` paths
// as rules. Equivalent dirs come in via `[paths] extra_rule_dirs`.
@ -494,15 +528,9 @@ async fn list_instructions(cwd: &Path) -> Vec<InstructionFile> {
configs
.into_iter()
.map(|c| {
let file_type = instruction_file_type(&c.file_path, imported, &extra_rule_prefixes);
let scope = if grok_home
.as_deref()
.is_some_and(|home| Path::new(&c.file_path).starts_with(home))
{
Scope::Global
} else {
Scope::Project
};
let file_type =
instruction_file_type(&c.file_path, &grok_home, imported, &extra_rule_prefixes);
let scope = instruction_scope(&c.file_path, &grok_home, &vendor_homes, &workspace_root);
let size = c.content.len();
let vendor = derive_vendor(&c.file_path).map(String::from);
InstructionFile {
@ -1604,7 +1632,7 @@ mod tests {
("claude", "/repo/.claude/rules/team.md"),
("claude", r"C:\repo\.claude\rules\team.md"),
] {
let file_type = instruction_file_type(path, false, &[]);
let file_type = instruction_file_type(path, Path::new("/home/user/.grok"), false, &[]);
assert_eq!(file_type, "rules");
assert_eq!(
instruction_compat_status(&Some(vendor.to_owned()), file_type, &report),
@ -1613,19 +1641,25 @@ mod tests {
}
for path in ["/repo/.grok/rules/team.md", r"C:\repo\.grok\rules\team.md"] {
assert_eq!(instruction_file_type(path, false, &[]), "rules");
assert_eq!(
instruction_file_type(path, Path::new("/home/user/.grok"), false, &[]),
"rules"
);
}
for path in [
"/repo/.cursor/rules/team.md",
r"C:\repo\.cursor\rules\team.md",
] {
assert_eq!(instruction_file_type(path, true, &[]), "rules");
assert_eq!(
instruction_file_type(path, Path::new("/home/user/.grok"), true, &[]),
"rules"
);
}
for path in [
"/repo/.claude/rules/team.md",
r"C:\repo\.claude\rules\team.md",
] {
let file_type = instruction_file_type(path, true, &[]);
let file_type = instruction_file_type(path, Path::new("/home/user/.grok"), true, &[]);
assert_eq!(file_type, "agents_md");
assert_eq!(
instruction_compat_status(&Some("claude".to_owned()), file_type, &report),
@ -1636,10 +1670,96 @@ mod tests {
"/repo/not.cursor/rules/team.md",
r"C:\repo\.cursor\ruleset\team.md",
] {
assert_eq!(instruction_file_type(path, false, &[]), "agents_md");
assert_eq!(
instruction_file_type(path, Path::new("/home/user/.grok"), false, &[]),
"agents_md"
);
}
}
#[test]
fn grok_home_nested_in_workspace_keeps_direct_surfaces_global() {
let grok_home = Path::new("/repo/config");
let workspace = Path::new("/repo");
for path in ["/repo/config/AGENTS.md", "/repo/config/rules/global.md"] {
assert!(matches!(
instruction_scope(path, grok_home, &[], workspace),
Scope::Global
));
}
for path in [
"/repo/config/.grok/rules/project.md",
"/repo/config/src/AGENTS.md",
] {
assert!(matches!(
instruction_scope(path, grok_home, &[], workspace),
Scope::Project
));
}
}
#[test]
fn vendor_home_nested_in_workspace_keeps_direct_surfaces_global() {
let vendor_homes = vec![(Path::new("/repo/.claude").to_path_buf(), true)];
let workspace = Path::new("/repo");
for path in ["/repo/.claude/rules/global.md", "/repo/.claude/CLAUDE.md"] {
assert!(matches!(
instruction_scope(path, Path::new("/other/grok"), &vendor_homes, workspace),
Scope::Global
));
}
for path in [
"/repo/.claude/.claude/rules/project.md",
"/repo/.claude/src/AGENTS.md",
] {
assert!(matches!(
instruction_scope(path, Path::new("/other/grok"), &vendor_homes, workspace),
Scope::Project
));
}
}
#[test]
fn workspace_scope_wins_inside_grok_home() {
let grok_home = Path::new("/custom/grok");
let workspace = Path::new("/custom/grok/worktrees/repo");
for path in [
"/custom/grok/worktrees/repo/.cursor/rules/project.md",
"/custom/grok/worktrees/repo/src/AGENTS.md",
] {
assert!(matches!(
instruction_scope(path, grok_home, &[], workspace),
Scope::Project
));
}
assert!(matches!(
instruction_scope("/custom/grok/rules/global.md", grok_home, &[], workspace,),
Scope::Global
));
}
#[test]
fn custom_grok_home_rules_are_classified_as_rules() {
assert_eq!(
instruction_file_type(
"/custom/config/rules/team.md",
Path::new("/custom/config"),
false,
&[],
),
"rules"
);
assert_eq!(
instruction_file_type(
"/custom/config/AGENTS.md",
Path::new("/custom/config"),
false,
&[],
),
"agents_md"
);
}
#[test]
fn describe_config_file_flags_empty_and_parse_error() {
let dir = tempfile::tempdir().unwrap();

1184
crates/codegen/xai-grok-shell/src/leader/server.rs
View file

File diff suppressed because it is too large Load diff

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

@ -5,6 +5,9 @@
unreachable_code,
dead_code
)]
#[cfg(all(test, feature = "dhat-heap"))]
#[global_allocator]
static DHAT_ALLOC: dhat::Alloc = dhat::Alloc;
pub(crate) use xai_grok_telemetry::unified_log;
pub use xai_tracing_macros::{teprintln, timed, tprintln};
pub mod active_sessions;

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

@ -13,11 +13,9 @@ use agent_client_protocol as acp;
/// Uses the JSON-RPC implementation-defined server error range (-32000 to -32099).
///
/// 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.
/// Clients derive user-facing text via [`format_rate_limited_user_message`].
/// The desktop path is unchanged: `prompt_complete_fields` still reports the
/// stop reason with no detail.
pub const RATE_LIMITED_ERROR_CODE: i32 = -32003;
/// OAuth / session rate-limit copy (personal plan upgrade path).
@ -29,20 +27,6 @@ pub const RATE_LIMITED_USER_MESSAGE_OAUTH: &str =
/// See https://docs.x.ai/developers/rate-limits#rate-limit-tiers
pub const RATE_LIMITED_USER_MESSAGE_API_KEY: &str = "You\u{2019}ve hit your team\u{2019}s API rate limit. Ask a team admin to purchase more credits for higher limits, or try again later. See https://docs.x.ai/developers/rate-limits#rate-limit-tiers";
/// Pick rate-limit copy from the *active* auth method.
///
/// Pass the real `is_api_key_auth` flag (pager `AppView`, `AuthMethodKind::is_api_key`
/// for the selected method). Do **not** decide from `has_xai_api_key_env()` alone:
/// when both an env key and a cached OAuth session exist, auth prefers the
/// cached session over the API key.
pub fn rate_limited_user_message(is_api_key_auth: bool) -> &'static str {
if is_api_key_auth {
RATE_LIMITED_USER_MESSAGE_API_KEY
} else {
RATE_LIMITED_USER_MESSAGE_OAUTH
}
}
/// 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
@ -62,29 +46,45 @@ pub fn is_free_usage_exhausted_error(detail: &str) -> bool {
/// 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.
/// Free-usage code first (consumer-only; intentional before API-key rewrite).
/// API-key + personal SuperGrok upsell → team credits copy. Else the body
/// after stripping `API error (status …):` (SamplingError Display prefix).
/// Empty → OAuth vs API-key fallback. Callers that show this in UI should
/// still run their usual sanitizer (scrub/cap).
pub fn format_rate_limited_user_message(
server_detail: Option<&str>,
is_api_key_auth: bool,
) -> String {
// Free-usage sniff works on the prefixed wire string (`contains` the code).
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()) {
let detail = strip_sampling_api_error_prefix(detail);
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()
if is_api_key_auth {
RATE_LIMITED_USER_MESSAGE_API_KEY
} else {
RATE_LIMITED_USER_MESSAGE_OAUTH
}
.to_string()
}
/// Drop `SamplingError::Api`'s Display prefix so users see the IC body, not
/// `API error (status 429 Too Many Requests): …`.
fn strip_sampling_api_error_prefix(detail: &str) -> &str {
const PREFIX: &str = "API error (status ";
const SEP: &str = "): ";
if let Some(rest) = detail.strip_prefix(PREFIX)
&& let Some(idx) = rest.find(SEP)
{
return rest[idx + SEP.len()..].trim();
}
detail.trim()
}
/// IC sometimes reuses OAuth free-tier upsell copy on 429s ("upgrade to a Grok
@ -371,13 +371,13 @@ mod tests {
}
#[test]
fn rate_limited_user_message_oauth_vs_api_key() {
fn rate_limited_fallback_oauth_vs_api_key() {
assert_eq!(
rate_limited_user_message(false),
format_rate_limited_user_message(None, false),
RATE_LIMITED_USER_MESSAGE_OAUTH
);
assert_eq!(
rate_limited_user_message(true),
format_rate_limited_user_message(None, true),
RATE_LIMITED_USER_MESSAGE_API_KEY
);
assert!(RATE_LIMITED_USER_MESSAGE_OAUTH.contains("Upgrade your account"));
@ -392,19 +392,19 @@ mod tests {
#[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
);
let body = "The service is temporarily at capacity. Please retry your request shortly.";
// Production detail is SamplingError::Api Display (prefixed).
let wire = format!("API error (status 429 Too Many Requests): {body}");
assert_eq!(format_rate_limited_user_message(Some(&wire), false), body);
assert_eq!(format_rate_limited_user_message(Some(&wire), true), body);
// 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);
let team_wire = format!("API error (status 429 Too Many Requests): {team}");
assert_eq!(
format_rate_limited_user_message(Some(&team_wire), true),
team
);
assert_eq!(
format_rate_limited_user_message(Some("slow down"), false),
"slow down"
@ -413,21 +413,25 @@ mod tests {
#[test]
fn format_rate_limited_api_key_rewrites_consumer_subscription_upsell() {
let rpm = "Some resource has been exhausted: You are sending requests too quickly. \
let body = "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";
let wire = format!("API error (status 429 Too Many Requests): {body}");
// OAuth keeps the IC body (personal plan upgrade is correct).
assert_eq!(format_rate_limited_user_message(Some(rpm), false), rpm);
assert_eq!(format_rate_limited_user_message(Some(&wire), false), body);
// API key must not push grok.com SuperGrok — team credits / rate-limit tiers.
assert_eq!(
format_rate_limited_user_message(Some(rpm), true),
format_rate_limited_user_message(Some(&wire), 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 format_rate_limited_strips_api_error_display_prefix() {
let body = "The service is temporarily at capacity.";
let wire = format!("API error (status 429 Too Many Requests): {body}");
assert_eq!(format_rate_limited_user_message(Some(&wire), false), body);
assert!(!format_rate_limited_user_message(Some(&wire), false).contains("API error"));
}
#[test]
@ -446,13 +450,15 @@ mod tests {
#[test]
fn format_rate_limited_free_usage_uses_paywall_copy() {
let detail = "subscription:free-usage-exhausted: You have used all your free usage.";
let wire = "API error (status 429 Too Many Requests): \
subscription:free-usage-exhausted: You have used all your free usage.";
assert_eq!(
format_rate_limited_user_message(Some(detail), false),
format_rate_limited_user_message(Some(wire), false),
FREE_USAGE_USER_MESSAGE
);
// Free-usage code is consumer-only; still wins for API-key callers.
assert_eq!(
format_rate_limited_user_message(Some(detail), true),
format_rate_limited_user_message(Some(wire), true),
FREE_USAGE_USER_MESSAGE
);
}

16
crates/codegen/xai-grok-shell/src/session/acp_session.rs
View file

@ -10,7 +10,7 @@
//!
use super::commands::{
ParsedPromptInfo, PromptCompletionKind, PromptTurnOk, PromptTurnResult, SessionCommand,
ok_end_turn,
TaskWakeAdmission, TaskWakeFallback, ok_end_turn,
};
use super::handle::SessionHandle;
use super::notifications::NotificationSender;
@ -196,6 +196,9 @@ pub(crate) struct InputItem {
pub(crate) json_schema: Option<serde_json::Value>,
/// Who originated this prompt — user or auto-wake system.
pub(crate) origin: super::PromptOrigin,
/// Typed deferred completion retained while an admitted task wake is queued.
/// Consumed by Ctrl+C if it removes the wake before the turn starts.
pub(crate) task_wake_fallback: Option<TaskWakeFallback>,
pub(crate) respond_to: oneshot::Sender<PromptTurnResult>,
/// Fired after the user message is in chat history and a persistence flush
/// barrier has completed (see `SessionCommand::Prompt::persist_ack`).
@ -271,8 +274,8 @@ pub(crate) struct State {
pub(crate) running_task: Option<AgentTask>,
pub(crate) pending_inputs: VecDeque<InputItem>,
pub(crate) pending_notifications: Vec<PendingNotification>,
/// When true, notifications are buffered but not drained until the next
/// user-initiated prompt arrives. Set on cancel, cleared on user Prompt.
/// When true, notifications are buffered but not drained until genuine
/// user re-engagement. Set by interactive Ctrl+C, cleared by a user prompt.
pub(crate) notifications_suppressed: bool,
/// Active prompt is still rewindable until the first outbound prompt-scoped
/// event is emitted.
@ -300,13 +303,13 @@ impl State {
}
/// Sweep `pending_inputs`, removing entries matching `drop_if` EXCEPT the
/// running turn's own slot, and return the removed items (callers harvest
/// them for telemetry counts / `auto_wake_delivered` un-marks).
/// them for telemetry counts / reservation releases).
///
/// Returned items still carry live `respond_to` senders that this helper
/// does NOT resolve — dropping them unfulfilled is correct only for
/// synthetic items (no client RPC awaits them, the current callers); a
/// caller whose predicate can match user-originated items must resolve
/// each returned item (see `respond_removed_queued_prompt`) or the
/// each returned item (see `respond_removed_prompt`) or the
/// client's `session/prompt` hangs and fails spuriously.
///
/// The guard is the safety invariant every sweep must inherit: the
@ -343,7 +346,8 @@ impl State {
/// so they share one definition of idleness, with no drift between them.
///
/// Returns `true` exactly when: no turn is running, no user prompt is
/// queued, and notifications haven't been suppressed by a cancel.
/// queued, and interactive Ctrl+C has not suppressed notifications pending
/// genuine user re-engagement.
pub(crate) fn is_session_idle_for_injection(state: &State) -> bool {
state.running_task.is_none()
&& state.pending_inputs.is_empty()

1
crates/codegen/xai-grok-shell/src/session/acp_session_impl/goal.rs
View file

@ -2202,6 +2202,7 @@ impl SessionActor {
verbatim: true,
json_schema: None,
origin: super::super::PromptOrigin::GoalSummary,
task_wake_fallback: None,
respond_to,
persist_ack: None,
parsed_prompt_tx: None,

1
crates/codegen/xai-grok-shell/src/session/acp_session_impl/interjection.rs
View file

@ -73,6 +73,7 @@ impl SessionActor {
verbatim: false,
json_schema: None,
origin: super::super::PromptOrigin::User,
task_wake_fallback: None,
respond_to,
persist_ack: None,
parsed_prompt_tx: None,

6
crates/codegen/xai-grok-shell/src/session/acp_session_impl/model_switch.rs
View file

@ -216,6 +216,12 @@ impl SessionActor {
),
)
.await;
if let Some(reservations) = self.tool_context.task_completion_reservations.clone() {
bridge.update_resource(reservations).await;
}
if let Some(gate) = self.tool_context.task_wake_suppressed.clone() {
bridge.update_resource(gate).await;
}
self.inject_deny_read_globs().await;
}
{

201
crates/codegen/xai-grok-shell/src/session/acp_session_impl/notification_drain.rs
View file

@ -19,6 +19,99 @@ pub(crate) struct PendingNotification {
}
impl SessionActor {
pub(super) fn push_pending_notification(state: &mut State, notification: PendingNotification) {
state.pending_notifications.push(notification);
let excess = state
.pending_notifications
.len()
.saturating_sub(MAX_PENDING_NOTIFICATIONS);
if excess > 0 {
state.pending_notifications.drain(..excess);
tracing::warn!(
dropped = excess,
"Dropped oldest pending notifications (exceeded cap of {})",
MAX_PENDING_NOTIFICATIONS,
);
}
}
pub(super) fn push_task_wake_fallback(state: &mut State, fallback: TaskWakeFallback) {
Self::push_pending_notification(
state,
PendingNotification {
prompt_id: fallback.prompt_id,
prompt_blocks: fallback.prompt_blocks,
priority: NotificationPriority::Later,
source: fallback.source,
},
);
}
pub(super) async fn consume_deferred_completions(&self) -> Vec<String> {
let mut state = self.state.lock().await;
self.sweep_monitor_buffer_into_pending(&mut state, "monitor-user-start-drain");
let mut completion_ids: Vec<String> = state
.pending_notifications
.iter()
.filter_map(|notification| match &notification.source {
NotificationSource::BashTaskCompleted { task_id }
| NotificationSource::MonitorCompleted { task_id } => Some(task_id.clone()),
NotificationSource::MonitorEvent { .. } => None,
})
.collect();
completion_ids.sort();
completion_ids.dedup();
let deferred_ids: std::collections::HashSet<&str> =
completion_ids.iter().map(String::as_str).collect();
let notifications = std::mem::take(&mut state.pending_notifications);
let mut deferred = Vec::new();
let mut retained = Vec::new();
for notification in notifications {
let consume = match &notification.source {
NotificationSource::BashTaskCompleted { .. }
| NotificationSource::MonitorCompleted { .. } => true,
NotificationSource::MonitorEvent { task_id } => {
deferred_ids.contains(task_id.as_str())
}
};
if consume {
deferred.push(notification);
} else {
retained.push(notification);
}
}
state.pending_notifications = retained;
let completion_blocks =
Self::notification_blocks(&deferred, &self.tool_context.task_output_tool_name);
drop(state);
let completion_text = completion_blocks
.into_iter()
.filter_map(|block| match block {
acp::ContentBlock::Text(text) => Some(text.text),
_ => None,
})
.collect::<Vec<_>>()
.join("\n");
if !completion_text.is_empty() {
self.push_system_reminder(&completion_text);
}
let completion_id_refs: Vec<&str> = completion_ids.iter().map(String::as_str).collect();
self.mark_completions_reported(&completion_id_refs).await;
completion_ids
}
pub(super) async fn consume_deferred_completions_for_user_turn(&self) {
let consumed = self.consume_deferred_completions().await;
if let Some(reservations) = &self.tool_context.task_completion_reservations {
for task_id in consumed {
reservations.release(&task_id);
}
}
}
pub(super) async fn maybe_start_running_task(
self: Arc<Self>,
completion_tx: mpsc::UnboundedSender<(String, PromptTurnResult)>,
@ -70,6 +163,7 @@ impl SessionActor {
screen_mode,
verbatim,
json_schema,
origin,
) = {
let Some(front) = state.pending_inputs.front_mut() else {
return;
@ -86,8 +180,20 @@ impl SessionActor {
front.screen_mode.clone(),
front.verbatim,
front.json_schema.clone(),
front.origin.clone(),
)
};
if matches!(origin, super::PromptOrigin::User) {
if let Some(gate) = &self.tool_context.task_wake_suppressed {
gate.set(false);
}
state.notifications_suppressed = false;
xai_grok_telemetry::unified_log::info(
"shell.task_wake.gate_cleared",
Some(self.session_info.id.0.as_ref()),
Some(serde_json::json!({ "reason": "queued_user_promotion" })),
);
}
{
let mut current_prompt_id = self
.current_prompt_id
@ -262,16 +368,19 @@ impl SessionActor {
buffer,
Some(self.session_info.id.0.as_ref()),
) {
state.pending_notifications.push(PendingNotification {
prompt_id: format!("{prompt_id_prefix}-{}", uuid::Uuid::now_v7()),
prompt_blocks: vec![acp::ContentBlock::Text(acp::TextContent::new(
event.event_text,
))],
priority: NotificationPriority::Next,
source: NotificationSource::MonitorEvent {
task_id: event.task_id,
Self::push_pending_notification(
state,
PendingNotification {
prompt_id: format!("{prompt_id_prefix}-{}", uuid::Uuid::now_v7()),
prompt_blocks: vec![acp::ContentBlock::Text(acp::TextContent::new(
event.event_text,
))],
priority: NotificationPriority::Next,
source: NotificationSource::MonitorEvent {
task_id: event.task_id,
},
},
});
);
}
}
@ -303,35 +412,34 @@ impl SessionActor {
(to_surface, dropped)
}
/// Build the merged `NotificationDrain` `InputItem` from `notifications`
/// and push it onto `state.pending_inputs`. Always returns `true` so the
/// caller starts the running task.
///
/// Monitor-event notifications are collapsed into ONE
/// `format_monitor_events` block (same batched/deduped shape as the
/// mid-turn injection — `<monitor task_id=…>` groups with `[label N]`
/// lines), placed at the position of the first monitor entry. Other
/// notifications keep their raw blocks. Everything joins with `---`.
pub(super) fn drain_notifications_into_turn(
state: &mut State,
notifications: Vec<PendingNotification>,
fn notification_blocks(
notifications: &[PendingNotification],
task_output_tool_name: &str,
) -> bool {
) -> Vec<acp::ContentBlock> {
use xai_grok_tools::implementations::grok_build::task::types::MonitorEventNotification;
// Collapse monitor entries: collect their text into events, remember
// where the first one sat so the batch lands in arrival position.
let completion_task_ids: std::collections::HashSet<&str> = notifications
.iter()
.filter_map(|notification| match &notification.source {
NotificationSource::MonitorCompleted { task_id } => Some(task_id.as_str()),
NotificationSource::MonitorEvent { .. }
| NotificationSource::BashTaskCompleted { .. } => None,
})
.collect();
let mut monitor_events: Vec<MonitorEventNotification> = Vec::new();
let mut sections: Vec<Vec<acp::ContentBlock>> = Vec::new();
let mut monitor_section_idx: Option<usize> = None;
for notif in &notifications {
match &notif.source {
for notification in notifications {
match &notification.source {
NotificationSource::MonitorEvent { task_id } => {
let event_text = notif
if completion_task_ids.contains(task_id.as_str()) {
continue;
}
let event_text = notification
.prompt_blocks
.iter()
.filter_map(|b| match b {
acp::ContentBlock::Text(t) => Some(t.text.as_str()),
.filter_map(|block| match block {
acp::ContentBlock::Text(text) => Some(text.text.as_str()),
_ => None,
})
.collect::<Vec<_>>()
@ -343,33 +451,42 @@ impl SessionActor {
});
if monitor_section_idx.is_none() {
monitor_section_idx = Some(sections.len());
sections.push(Vec::new()); // placeholder, filled below
sections.push(Vec::new());
}
}
NotificationSource::BashTaskCompleted { .. } => {
sections.push(notif.prompt_blocks.clone());
NotificationSource::MonitorCompleted { .. }
| NotificationSource::BashTaskCompleted { .. } => {
sections.push(notification.prompt_blocks.clone());
}
}
}
if let (Some(idx), Some(batch)) = (
if let (Some(index), Some(batch)) = (
monitor_section_idx,
xai_grok_tools::reminders::task_completion::format_monitor_events(
&monitor_events,
Some(task_output_tool_name),
),
) {
sections[idx] = vec![acp::ContentBlock::Text(acp::TextContent::new(batch))];
sections[index] = vec![acp::ContentBlock::Text(acp::TextContent::new(batch))];
}
let mut merged_blocks: Vec<acp::ContentBlock> = Vec::new();
for (i, section) in sections.iter().enumerate() {
if i > 0 {
merged_blocks.push(acp::ContentBlock::Text(acp::TextContent::new(
"---".to_string(),
)));
let mut blocks = Vec::new();
for (index, section) in sections.iter().enumerate() {
if index > 0 {
blocks.push(acp::ContentBlock::Text(acp::TextContent::new("---")));
}
merged_blocks.extend(section.iter().cloned());
blocks.extend(section.iter().cloned());
}
blocks
}
/// Merge notifications into one queued `NotificationDrain` turn.
pub(super) fn drain_notifications_into_turn(
state: &mut State,
notifications: Vec<PendingNotification>,
task_output_tool_name: &str,
) -> bool {
let merged_blocks = Self::notification_blocks(&notifications, task_output_tool_name);
let merged_prompt_id = format!("notifications-{}", uuid::Uuid::now_v7());
@ -389,6 +506,7 @@ impl SessionActor {
verbatim: true,
json_schema: None,
origin: super::PromptOrigin::NotificationDrain,
task_wake_fallback: None,
respond_to,
persist_ack: None,
parsed_prompt_tx: None,
@ -402,6 +520,7 @@ impl SessionActor {
later_count = notifications.iter().filter(|n| n.priority == NotificationPriority::Later).count(),
sources = %notifications.iter().map(|n| match &n.source {
NotificationSource::MonitorEvent { task_id } => format!("monitor:{task_id}"),
NotificationSource::MonitorCompleted { task_id } => format!("monitor-completed:{task_id}"),
NotificationSource::BashTaskCompleted { task_id } => format!("bash:{task_id}"),
}).collect::<Vec<_>>().join(","),
"Drained pending notifications into single batched turn"

178
crates/codegen/xai-grok-shell/src/session/acp_session_impl/prompt_build.rs
View file

@ -3,11 +3,6 @@
//! payload preparation.
#![allow(clippy::items_after_test_module)]
use super::*;
/// Partition the AGENTS.md / Claude.md / `.grok/rules/*.md` files returned
/// by `read_agents_config_with_paths` into "workspace" (cwd / repo root /
/// extra workspace user dir) and "user" (`~/.grok/`, `~/.claude/`,
/// `~/.grok/bundled/`) buckets, mirroring the split between
/// `<always_applied_workspace_rules>` and `<user_rules>`.
/// Normalize a free-form name (e.g. an MCP server identifier) into a
/// single safe filesystem segment.
///
@ -41,20 +36,24 @@ pub(super) fn pick_user_image_url(image: &agent_client_protocol::ImageContent) -
}
fn partition_rules_by_scope(
files: Vec<xai_grok_agent::prompt::agents_md::AgentConfigFile>,
grok_home: &std::path::Path,
vendor_homes: &[(std::path::PathBuf, bool)],
workspace_root: Option<&std::path::Path>,
) -> (
Vec<xai_grok_agent::prompt::user_message::RuleEntry>,
Vec<xai_grok_agent::prompt::user_message::RuleEntry>,
) {
let home = dirs::home_dir().map(|p| p.to_string_lossy().to_string());
let user_prefixes: Vec<String> = match home {
Some(h) => vec![format!("{h}/.grok/"), format!("{h}/.claude/")],
None => vec![],
};
let mut workspace = Vec::new();
let mut user = Vec::new();
for f in files {
let entry = xai_grok_agent::prompt::user_message::RuleEntry::from(f);
if user_prefixes.iter().any(|p| entry.path.starts_with(p)) {
for file in files {
let is_user_rule = crate::util::is_user_instruction_path(
std::path::Path::new(&file.file_path),
grok_home,
vendor_homes,
workspace_root,
);
let entry = xai_grok_agent::prompt::user_message::RuleEntry::from(file);
if is_user_rule {
user.push(entry);
} else {
workspace.push(entry);
@ -62,6 +61,133 @@ fn partition_rules_by_scope(
}
(workspace, user)
}
#[cfg(test)]
mod partition_rules_by_scope_tests {
use super::partition_rules_by_scope;
use std::path::Path;
use xai_grok_agent::prompt::agents_md::AgentConfigFile;
fn file(path: &str) -> AgentConfigFile {
AgentConfigFile {
file_name: Path::new(path)
.file_name()
.unwrap()
.to_string_lossy()
.into_owned(),
file_path: path.to_string(),
content: path.to_string(),
}
}
fn paths(entries: &[xai_grok_agent::prompt::user_message::RuleEntry]) -> Vec<&str> {
entries.iter().map(|entry| entry.content.as_str()).collect()
}
#[test]
fn partitions_custom_grok_and_vendor_home_rules_as_user_scope() {
let files = vec![
file("/custom/config/rules/a.md"),
file("/home/user/.cursor/rules/b.md"),
file("/repo/.cursor/rules/c.md"),
file("/repo/src/AGENTS.md"),
file("/custom/config/rules/d.md"),
];
let vendor_homes = vec![
(Path::new("/home/user/.claude").to_path_buf(), true),
(Path::new("/home/user/.cursor").to_path_buf(), true),
];
let (workspace, user) = partition_rules_by_scope(
files,
Path::new("/custom/config"),
&vendor_homes,
Some(Path::new("/repo")),
);
assert_eq!(
paths(&user),
vec![
"/custom/config/rules/a.md",
"/home/user/.cursor/rules/b.md",
"/custom/config/rules/d.md",
]
);
assert_eq!(
paths(&workspace),
vec!["/repo/.cursor/rules/c.md", "/repo/src/AGENTS.md"]
);
}
#[test]
fn grok_home_nested_in_workspace_keeps_direct_surfaces_user_scoped() {
let files = vec![
file("/repo/config/AGENTS.md"),
file("/repo/config/rules/global.md"),
file("/repo/config/.grok/rules/project.md"),
file("/repo/config/src/AGENTS.md"),
];
let (workspace, user) = partition_rules_by_scope(
files,
Path::new("/repo/config"),
&[],
Some(Path::new("/repo")),
);
assert_eq!(
paths(&user),
vec!["/repo/config/AGENTS.md", "/repo/config/rules/global.md"]
);
assert_eq!(
paths(&workspace),
vec![
"/repo/config/.grok/rules/project.md",
"/repo/config/src/AGENTS.md",
]
);
}
#[test]
fn vendor_home_nested_in_workspace_keeps_direct_surfaces_user_scoped() {
let files = vec![
file("/repo/.claude/rules/global.md"),
file("/repo/.claude/CLAUDE.md"),
file("/repo/.claude/.claude/rules/project.md"),
file("/repo/.claude/src/AGENTS.md"),
];
let vendor_homes = vec![(Path::new("/repo/.claude").to_path_buf(), true)];
let (workspace, user) = partition_rules_by_scope(
files,
Path::new("/other/grok"),
&vendor_homes,
Some(Path::new("/repo")),
);
assert_eq!(
paths(&user),
vec!["/repo/.claude/rules/global.md", "/repo/.claude/CLAUDE.md"]
);
assert_eq!(
paths(&workspace),
vec![
"/repo/.claude/.claude/rules/project.md",
"/repo/.claude/src/AGENTS.md",
]
);
}
#[test]
fn nested_grok_home_workspace_files_stay_workspace_scoped() {
let files = vec![
file("/custom/grok/rules/global.md"),
file("/custom/grok/worktrees/repo/.cursor/rules/project.md"),
file("/custom/grok/worktrees/repo/src/AGENTS.md"),
];
let (workspace, user) = partition_rules_by_scope(
files,
Path::new("/custom/grok"),
&[],
Some(Path::new("/custom/grok/worktrees/repo")),
);
assert_eq!(paths(&user), vec!["/custom/grok/rules/global.md"]);
assert_eq!(
paths(&workspace),
vec![
"/custom/grok/worktrees/repo/.cursor/rules/project.md",
"/custom/grok/worktrees/repo/src/AGENTS.md",
]
);
}
}
/// True iff `conversation` already contains a project-instructions reminder,
/// either tagged [`SyntheticReason::ProjectInstructions`] or a legacy untagged
/// copy whose first text part starts with [`LEGACY_AGENTS_MD_REMINDER_PREFIX`].
@ -393,7 +519,31 @@ impl SessionActor {
let bridge = self.agent.borrow().tool_bridge().clone();
let (vcs_root, vcs_status) = self.gather_vcs_for_prefix(cwd).await;
let agents_files = read_agents_config_with_paths(&cwd_str, self.rebuild_spec.compat).await;
let (workspace_rules, user_rules) = partition_rules_by_scope(agents_files);
let grok_home = xai_grok_config::grok_home();
let vendor_homes = dirs::home_dir()
.map(|home_dir| {
vec![
(
home_dir.join(".claude"),
self.rebuild_spec.compat.claude.agents,
),
(
home_dir.join(".cursor"),
self.rebuild_spec.compat.cursor.agents,
),
]
})
.unwrap_or_default();
let workspace_root = git2::Repository::discover(cwd)
.ok()
.and_then(|repo| repo.workdir().map(std::path::Path::to_path_buf))
.unwrap_or_else(|| cwd.to_path_buf());
let (workspace_rules, user_rules) = partition_rules_by_scope(
agents_files,
&grok_home,
&vendor_homes,
Some(&workspace_root),
);
let skills = bridge.slash_skills().await;
let mcp_servers = self.gather_mcp_servers(cwd).await;
let shell = resolve_session_shell();

35
crates/codegen/xai-grok-shell/src/session/acp_session_impl/prompt_queue.rs
View file

@ -18,6 +18,7 @@ impl SessionActor {
verbatim: bool,
json_schema: Option<serde_json::Value>,
send_now: bool,
task_wake_fallback: Option<TaskWakeFallback>,
respond_to: oneshot::Sender<PromptTurnResult>,
persist_ack: Option<oneshot::Sender<()>>,
parsed_prompt_tx: Option<oneshot::Sender<ParsedPromptInfo>>,
@ -84,10 +85,8 @@ impl SessionActor {
(trace_gcs_config, artifact_tracker)
};
// Pre-mark auto-wake completion IDs so TaskCompletionReminder doesn't
// duplicate the notification the model already sees in the synthetic prompt.
if let Some(id) = origin.completion_id() {
self.mark_completions_reported(&[id]).await;
if let crate::session::PromptOrigin::SubagentCompleted { subagent_id } = &origin {
self.mark_completions_reported(&[subagent_id]).await;
}
// For synthetic prompts, derive trace config from the template
@ -133,17 +132,18 @@ impl SessionActor {
});
if preempt_armed {
let dropped = state.sweep_pending_inputs(|i| i.origin.is_synthetic());
if let Some(reservations) = &self.tool_context.task_completion_reservations {
for task_id in dropped
.iter()
.filter_map(|item| item.origin.completion_id())
{
reservations.release(task_id);
}
}
tracing::info!(
dropped_count = dropped.len(),
"auto-wake: dropping pending synthetic prompts (user prompt has priority)"
);
// Un-mark preempted IDs so TaskCompletionReminder can report
// them on the next pass instead of permanently suppressing them.
if let Some(ref auto_wake) = self.tool_context.auto_wake_delivered {
for id in dropped.iter().filter_map(|i| i.origin.completion_id()) {
auto_wake.remove(id);
}
}
}
}
@ -190,6 +190,7 @@ impl SessionActor {
verbatim,
json_schema,
origin,
task_wake_fallback,
respond_to,
persist_ack,
parsed_prompt_tx,
@ -392,7 +393,7 @@ impl SessionActor {
/// benign no-op — the actor still re-broadcasts so the client reconciles.
/// The in-flight turn is never removed. `owner` (when `Some`) scopes the
/// edit to the requesting client's own items.
/// Resolve a removed/cleared queued prompt's in-flight `session/prompt` RPC
/// Resolve a removed prompt's in-flight `session/prompt` RPC
/// before its [`InputItem`] is dropped.
///
/// A queued prompt still has a client awaiting its `respond_to` oneshot (the
@ -412,8 +413,8 @@ impl SessionActor {
/// delta, so other attached clients (leader mode) don't see the running
/// turn spuriously end. Token count is `0` — a removed queued prompt never
/// ran (and the value is discarded by the gate regardless).
fn respond_removed_queued_prompt(item: InputItem) {
let _ = item.respond_to.send(Ok(PromptTurnOk {
pub(super) fn respond_removed_prompt(respond_to: oneshot::Sender<PromptTurnResult>) {
let _ = respond_to.send(Ok(PromptTurnOk {
stop_reason: acp::StopReason::Cancelled,
total_tokens: 0,
turn_snapshot: None,
@ -441,7 +442,7 @@ impl SessionActor {
})
{
if let Some(item) = state.pending_inputs.remove(pos) {
Self::respond_removed_queued_prompt(item);
Self::respond_removed_prompt(item.respond_to);
}
removed = true;
}
@ -619,7 +620,7 @@ impl SessionActor {
let mut state = self.state.lock().await;
// Partition rather than `retain`: each cleared user prompt still has a
// client awaiting its `respond_to`, so it must be resolved with
// `Cancelled` (see [`respond_removed_queued_prompt`]) instead of being
// `Cancelled` (see [`respond_removed_prompt`]) instead of being
// dropped — a bare drop surfaces as "session failed to respond" and a
// spurious "Turn failed" on the running turn.
let running_id = state.running_prompt_id().map(str::to_string);
@ -638,7 +639,7 @@ impl SessionActor {
if keep {
kept.push_back(item);
} else {
Self::respond_removed_queued_prompt(item);
Self::respond_removed_prompt(item.respond_to);
}
}
state.pending_inputs = kept;

16
crates/codegen/xai-grok-shell/src/session/acp_session_impl/reminders.rs
View file

@ -249,8 +249,8 @@ impl SessionActor {
/// Mark completion IDs as reported in the shared
/// `ReportedTaskCompletions` state so the per-tool-call
/// `TaskCompletionReminder` won't (re-)surface them. Used both to dedupe
/// completions the model already saw (notification-drain / auto-wake
/// prompts) and to drop them during the goal loop (between-turn drain).
/// completions the model actually saw (notification-drain / started
/// auto-wake prompts) and to drop them during the goal loop (between-turn drain).
/// No-op on an empty list.
pub(super) async fn mark_completions_reported(&self, ids: &[&str]) {
if ids.is_empty() {
@ -283,7 +283,13 @@ impl SessionActor {
pub(super) async fn drain_between_turn_completions(&self) {
let goal_loop_active = self.goal_loop_active();
let bridge = self.agent.borrow().tool_bridge().clone();
let bash_completions = bridge.drain_between_turn_bash_completions().await;
let reserved = self
.tool_context
.task_completion_reservations
.as_ref()
.map(|reservations| reservations.snapshot())
.unwrap_or_default();
let bash_completions = bridge.drain_between_turn_bash_completions(&reserved).await;
if !bash_completions.is_empty() {
let ids: Vec<&str> = bash_completions
.iter()
@ -324,9 +330,9 @@ impl SessionActor {
};
let suppress_ids = self
.tool_context
.auto_wake_delivered
.task_completion_reservations
.as_ref()
.map(|a| a.snapshot())
.map(|reservations| reservations.snapshot())
.unwrap_or_default();
let (respond_to, rx) = tokio::sync::oneshot::channel();
if tx

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

@ -30,6 +30,58 @@ mod yolo_toggle_report_tests {
/// Best-effort removal of this session's per-session scratch staging on
/// teardown. A no-op in builds without a scratch producer.
fn cleanup_session_scratch(_session: &SessionActor) {}
impl SessionActor {
/// Serialize terminal task-wake admission with interactive cancellation.
pub(super) async fn admit_task_completion_wake(
&self,
origin: &super::PromptOrigin,
admission: TaskWakeAdmission,
) -> Option<TaskWakeFallback> {
let TaskWakeAdmission {
respond_to,
fallback,
} = admission;
let super::PromptOrigin::TaskCompleted { task_id } = origin else {
return respond_to.send(true).is_ok().then_some(fallback);
};
let gate_suppressed = self
.tool_context
.task_wake_suppressed
.as_ref()
.is_some_and(|gate| gate.get());
let mut state = self.state.lock().await;
let state_suppressed = state.notifications_suppressed;
let admitted = !gate_suppressed && !state_suppressed;
if !admitted {
Self::push_task_wake_fallback(&mut state, fallback);
drop(state);
xai_grok_telemetry::unified_log::info(
"shell.task_wake.actor_admission",
Some(self.session_info.id.0.as_ref()),
Some(serde_json::json!(
{ "task_id" : task_id, "gate" : gate_suppressed, "state" :
state_suppressed, "admitted" : false, }
)),
);
let _ = respond_to.send(false);
return None;
}
if respond_to.send(true).is_err() {
Self::push_task_wake_fallback(&mut state, fallback);
return None;
}
drop(state);
xai_grok_telemetry::unified_log::info(
"shell.task_wake.actor_admission",
Some(self.session_info.id.0.as_ref()),
Some(serde_json::json!(
{ "task_id" : task_id, "gate" : gate_suppressed, "state" :
state_suppressed, "admitted" : true, }
)),
);
Some(fallback)
}
}
pub(super) async fn run_session(
session: Arc<SessionActor>,
mut cmd_rx: mpsc::UnboundedReceiver<SessionCommand>,
@ -280,11 +332,18 @@ pub(super) async fn run_session(
move { s.resume_plan_approval(completion_tx). await; }); }
SessionCommand::Prompt { prompt_id, prompt_blocks, prompt_mode,
artifact_upload_ctx, client_identifier, screen_mode, verbatim, traceparent,
json_schema, send_now, respond_to, persist_ack, parsed_prompt_tx } => {
session.ensure_prefix_ready(). await; let origin =
super::PromptOrigin::from_prompt_id(& prompt_id); if ! origin.is_synthetic()
{ let mut state = session.state.lock(). await; state.notifications_suppressed
= false; session.user_input_generation.fetch_add(1,
json_schema, send_now, admission, respond_to, persist_ack, parsed_prompt_tx }
=> { let origin = super::PromptOrigin::from_prompt_id(& prompt_id); let
(actor_admitted, task_wake_fallback) = match admission { Some(admission) => {
let fallback = session.admit_task_completion_wake(& origin, admission).
await; (fallback.is_some(), fallback) } None => (true, None), }; if !
actor_admitted { SessionActor::respond_removed_prompt(respond_to); continue;
} session.ensure_prefix_ready(). await; if ! origin.is_synthetic() { if let
Some(gate) = & session.tool_context.task_wake_suppressed { gate.set(false); }
let mut state = session.state.lock(). await; state.notifications_suppressed =
false; xai_grok_telemetry::unified_log::info("shell.task_wake.gate_cleared",
Some(session.session_info.id.0.as_ref()), Some(serde_json::json!({ "reason" :
"user_intake" })),); session.user_input_generation.fetch_add(1,
std::sync::atomic::Ordering::AcqRel); } if origin.is_synthetic() { let state
= session.state.lock(). await; let has_running = state.running_task
.is_some(); let queue_depth = state.pending_inputs.len(); drop(state);
@ -297,8 +356,8 @@ pub(super) async fn run_session(
=> (Some(tu.gcs_config), Some(tu.artifact_tracker)), None => (None, None), };
let cancel_for_send_now = session.queue_input(prompt_blocks, prompt_id,
prompt_mode, trace_gcs_config, artifact_tracker, client_identifier,
screen_mode, verbatim, json_schema, send_now, respond_to, persist_ack,
parsed_prompt_tx). await; if cancel_for_send_now { session
screen_mode, verbatim, json_schema, send_now, task_wake_fallback, respond_to,
persist_ack, parsed_prompt_tx). await; if cancel_for_send_now { session
.cancel_turn_for_send_now(& mut replay_buffer). await; }
SessionActor::maybe_start_running_task(session.clone(), completion_tx
.clone()). await; } SessionCommand::SessionMode { session_mode, responds_to }
@ -395,14 +454,9 @@ pub(super) async fn run_session(
.session_info.id.0.to_string(),), }, MAX_BUFFER_EVENTS,);
tracing::debug!(task_id = % task_id,
"Routed monitor event to mid-turn buffer"); } } else { { let mut state =
session.state.lock(). await; state.pending_notifications
.push(PendingNotification { prompt_id, prompt_blocks, priority, source, });
if state.pending_notifications.len() > MAX_PENDING_NOTIFICATIONS { let excess
= state.pending_notifications.len() - MAX_PENDING_NOTIFICATIONS; state
.pending_notifications.drain(..excess); tracing::warn!(dropped = excess,
"Dropped oldest pending notifications (exceeded cap of {})",
MAX_PENDING_NOTIFICATIONS,); } }
SessionActor::maybe_drain_notifications(session.clone(), completion_tx
session.state.lock(). await; SessionActor::push_pending_notification(& mut
state, PendingNotification { prompt_id, prompt_blocks, priority, source, },);
} SessionActor::maybe_drain_notifications(session.clone(), completion_tx
.clone()). await; } } SessionCommand::RecordGoalTurnTaskIds { task_ids } => {
session.record_reparented_goal_turn_task_ids(task_ids); }
SessionCommand::RemoveQueuedPrompt { id, expected_version, owner } => {
@ -421,16 +475,19 @@ pub(super) async fn run_session(
kill_background_tasks, rewind_if_pristine, trigger, } => { if let
Some(notification) = replay_buffer.flush() { session
.emit_buffered(notification). await; } session.pending_interjections.clear();
session.cancel_running_task(cancel_subagents, kill_background_tasks,
let suppress_task_wakes = trigger.as_deref() == Some("ctrl_c"); session
.cancel_running_task(cancel_subagents, kill_background_tasks,
rewind_if_pristine, trigger,). await; session.auto_pause_goal_if_active(crate
::session::goal_tracker::GoalPauseReason::User,). await;
SessionActor::maybe_start_running_task(session.clone(), completion_tx
.clone()). await; } SessionCommand::CompactSession { user_context, respond_to
} => { let s = session.clone(); tokio::task::spawn_local(async move { let
compact_session = s.run_compact(user_context). await; let _ = respond_to
.send(compact_session); }); } SessionCommand::ReloadPlugins { registry } => {
if ! session.startup_hints.is_subagent { let registry = session
.preserve_session_plugin_dirs(registry); session
.clone()). await; if ! suppress_task_wakes {
SessionActor::maybe_drain_notifications(session.clone(), completion_tx
.clone(),). await; } } SessionCommand::CompactSession { user_context,
respond_to } => { let s = session.clone(); tokio::task::spawn_local(async
move { let compact_session = s.run_compact(user_context). await; let _ =
respond_to.send(compact_session); }); } SessionCommand::ReloadPlugins {
registry } => { if ! session.startup_hints.is_subagent { let registry =
session.preserve_session_plugin_dirs(registry); session
.apply_plugin_registry_snapshot(registry). await; } }
SessionCommand::ReloadHooks => { if ! session.startup_hints.is_subagent { let
_ = session.reload_hooks_impl(). await; } }
@ -751,19 +808,20 @@ pub(super) async fn run_session(
prompt_blocks, prompt_mode : crate ::session::plan_mode::PromptMode::Agent,
trace_gcs_config : None, artifact_tracker : None, client_identifier : None,
screen_mode : None, verbatim : true, json_schema : None, origin :
super::PromptOrigin::GoalSummary, respond_to, persist_ack : None,
parsed_prompt_tx : None, queue_meta : None, send_now : false, }); }
SessionActor::maybe_start_running_task(session.clone(), completion_tx
.clone()). await; } SessionCommand::TakeTurnMessages { respond_to } => { let
result = session.chat_state_handle.take_turn_messages(). await; let _ =
respond_to.send(result); } SessionCommand::TakeHarnessTraceTurns { respond_to
} => { let result = session.chat_state_handle.take_harness_trace_turns().
await; let _ = respond_to.send(result); }
SessionCommand::TakeStreamingCapture { prompt_id, respond_to } => { let taken
= { let mut cap = session.streaming_turn_capture.lock(); if cap.prompt_id
.as_deref() == Some(prompt_id.as_str()) { Some(std::mem::take(& mut * cap)) }
else { if ! cap.is_empty() { tracing::warn!(requested_prompt_id = %
prompt_id, slot_prompt_id = ? cap.prompt_id,
super::PromptOrigin::GoalSummary, task_wake_fallback : None, respond_to,
persist_ack : None, parsed_prompt_tx : None, queue_meta : None, send_now :
false, }); } SessionActor::maybe_start_running_task(session.clone(),
completion_tx.clone()). await; } SessionCommand::TakeTurnMessages {
respond_to } => { let result = session.chat_state_handle.take_turn_messages()
. await; let _ = respond_to.send(result); }
SessionCommand::TakeHarnessTraceTurns { respond_to } => { let result =
session.chat_state_handle.take_harness_trace_turns(). await; let _ =
respond_to.send(result); } SessionCommand::TakeStreamingCapture { prompt_id,
respond_to } => { let taken = { let mut cap = session.streaming_turn_capture
.lock(); if cap.prompt_id.as_deref() == Some(prompt_id.as_str()) {
Some(std::mem::take(& mut * cap)) } else { if ! cap.is_empty() {
tracing::warn!(requested_prompt_id = % prompt_id, slot_prompt_id = ? cap
.prompt_id,
"streaming_capture race: live slot belongs to a different prompt; \
dropping streaming_partial.json for the requested turn",);
} None } }; let result = taken.and_then(| mut cap | { cap

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

@ -913,6 +913,14 @@ impl SessionActor {
}
}
/// Proactively refresh the auth token if near expiry.
///
/// Session-token path is best-effort: on success, update credentials and
/// return. On failure, do **not** fall through to the JWT/config.toml
/// branch when the session gate was active — that path is for BYOK JWTs
/// only. Falling through after a failed session refresh left hard-expired
/// opaque tokens (External/OIDC) on the wire and guaranteed a 401.
/// Soft failures with a still-usable access token still return here
/// (grace / optimistic send); 401 recovery remains the safety net.
pub(super) async fn refresh_token_if_expired(&self) {
if let Some(ref am) = self.auth_manager {
let creds = self.chat_state_handle.get_credentials().await;
@ -922,15 +930,33 @@ impl SessionActor {
.await
.map(|c| (c.model, c.base_url))
.unwrap_or_default();
if self.auth_gate(&model_id, &base_url).active()
&& let Ok(key) = am.get_valid_token().await
{
if creds.api_key.as_deref() != Some(&key) {
let mut creds = creds;
creds.api_key = Some(key);
self.chat_state_handle.update_credentials(creds);
if self.auth_gate(&model_id, &base_url).active() {
match am.get_valid_token().await {
Ok(key) => {
if creds.api_key.as_deref() != Some(&key) {
let mut creds = creds;
creds.api_key = Some(key);
self.chat_state_handle.update_credentials(creds);
}
return;
}
Err(e) => {
let hard_expired = !am.has_usable_token();
tracing::warn!(
error = % e, hard_expired, model = % model_id,
"auth: preflight get_valid_token failed"
);
xai_grok_telemetry::unified_log::warn(
"auth.preflight.refresh_failed",
Some(self.session_info.id.0.as_ref()),
Some(serde_json::json!(
{ "error" : format!("{e}"), "hard_expired" : hard_expired,
"model" : model_id, }
)),
);
return;
}
}
return;
}
} else {
xai_grok_telemetry::unified_log::debug(

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

@ -118,6 +118,7 @@ pub(crate) async fn spawn_session_actor(
system_prompt_label: String,
compaction_mode: xai_chat_state::CompactionMode,
compaction_verbatim_input: bool,
compaction_tool_choice: crate::util::config::CompactionToolChoice,
two_pass_enabled: bool,
buffering_settings: Option<BufferingSettings>,
origin_client: Option<crate::http::OriginClientInfo>,
@ -464,9 +465,12 @@ pub(crate) async fn spawn_session_actor(
None => FileStateTracker::new(),
});
let file_state_handle = FileStateHandle::new(file_state_tracker.clone());
let auto_wake_delivered =
xai_grok_tools::reminders::task_completion::AutoWakeDeliveredIds::default();
tool_context.auto_wake_delivered = Some(auto_wake_delivered.clone());
let task_completion_reservations =
xai_grok_tools::reminders::task_completion::TaskCompletionReservations::default();
let task_wake_suppressed =
xai_grok_tools::reminders::task_completion::TaskWakeSuppressed::default();
tool_context.task_completion_reservations = Some(task_completion_reservations.clone());
tool_context.task_wake_suppressed = Some(task_wake_suppressed.clone());
let synthetic_trace_tx_shared: std::sync::Arc<
std::sync::Mutex<
Option<
@ -522,7 +526,8 @@ pub(crate) async fn spawn_session_actor(
current_prompt_mode: current_prompt_mode.clone(),
turn_prompt_mode: turn_prompt_mode.clone(),
session_cmd_tx: cmd_tx.clone(),
auto_wake_delivered: auto_wake_delivered.clone(),
task_completion_reservations: task_completion_reservations.clone(),
task_wake_suppressed: task_wake_suppressed.clone(),
synthetic_trace_tx: synthetic_trace_tx_shared.clone(),
task_output_tool_name: task_output_tool_name.clone(),
read_tool_name: read_tool_name.clone(),
@ -545,19 +550,46 @@ pub(crate) async fn spawn_session_actor(
grep_ugrep,
}
};
let cursor_harness = false;
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(),
cursor_harness,
);
let terminal_backend: std::sync::Arc<dyn xai_grok_tools::computer::types::TerminalBackend> =
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
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 => {
let login_shell_capture = crate::util::config::resolve_login_shell_capture(
remote_settings.as_ref().and_then(|r| r.login_shell_capture),
);
std::sync::Arc::new(LocalTerminalBackend::new_local_with_login_shell_capture(
resolve_search_shadows(),
login_shell_capture,
))
}
};
if matches!(
terminal_backend_kind,
TerminalBackendKind::LocalPersistent | TerminalBackendKind::LocalNonPersistent
) {
terminal_backend
.warm_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(
@ -866,6 +898,14 @@ pub(crate) async fn spawn_session_actor(
);
e
})?;
agent
.tool_bridge()
.update_resource(task_completion_reservations.clone())
.await;
agent
.tool_bridge()
.update_resource(task_wake_suppressed)
.await;
let resolved_task_output =
xai_grok_tools::reminders::task_completion::resolve_task_output_tool_name(
agent.tool_bridge(),
@ -1143,6 +1183,7 @@ pub(crate) async fn spawn_session_actor(
previous_model: std::cell::Cell::new(None),
compaction_mode,
verbatim_input: compaction_verbatim_input,
tool_choice: compaction_tool_choice,
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@ -1686,6 +1727,7 @@ pub(crate) async fn spawn_session_on_thread(
system_prompt_label: String,
compaction_mode: xai_chat_state::CompactionMode,
compaction_verbatim_input: bool,
compaction_tool_choice: crate::util::config::CompactionToolChoice,
two_pass_enabled: bool,
buffering_settings: Option<BufferingSettings>,
origin_client: Option<crate::http::OriginClientInfo>,
@ -1848,6 +1890,7 @@ pub(crate) async fn spawn_session_on_thread(
system_prompt_label,
compaction_mode,
compaction_verbatim_input,
compaction_tool_choice,
two_pass_enabled,
buffering_settings,
origin_client,
@ -2017,3 +2060,78 @@ 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,
cursor_harness: bool,
) -> TerminalBackendKind {
if is_subagent && has_parent_backend {
TerminalBackendKind::ReuseParent
} else if client_terminal_capable && has_gateway {
TerminalBackendKind::AcpClient
} else if cursor_harness {
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_cursor_harness_selects_persistent_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
);
}
}

83
crates/codegen/xai-grok-shell/src/session/acp_session_impl/tasks_cancel.rs
View file

@ -216,7 +216,15 @@ impl SessionActor {
self.cancel_running_task(false, false, false, Some("send_now".to_string()))
.await;
// Re-enable notification drains: unlike Ctrl+C, a send-now means the user is re-engaged.
if let Some(gate) = &self.tool_context.task_wake_suppressed {
gate.set(false);
}
self.state.lock().await.notifications_suppressed = false;
xai_grok_telemetry::unified_log::info(
"shell.task_wake.gate_cleared",
Some(self.session_info.id.0.as_ref()),
Some(serde_json::json!({ "reason": "send_now" })),
);
}
pub(super) async fn cancel_running_task(
@ -226,6 +234,32 @@ impl SessionActor {
rewind_if_pristine: bool,
trigger: Option<String>,
) {
let suppress_task_wakes = trigger.as_deref() == Some("ctrl_c");
if suppress_task_wakes {
if let Some(gate) = &self.tool_context.task_wake_suppressed {
gate.set(true);
}
let mut state = self.state.try_lock().expect("session state is actor-owned");
state.notifications_suppressed = true;
xai_grok_telemetry::unified_log::info(
"shell.task_wake.cancel_barrier",
Some(self.session_info.id.0.as_ref()),
Some(serde_json::json!({
"ctrl_c": true,
"gate": self
.tool_context
.task_wake_suppressed
.as_ref()
.is_some_and(|gate| gate.get()),
"state": state.notifications_suppressed,
})),
);
drop(state);
if let Some(is_turn_active) = &self.tool_context.is_turn_active {
is_turn_active.store(false, std::sync::atomic::Ordering::Relaxed);
}
}
// Unified-log processing marker (counterpart of `shell.cancel.received`
// in `MvpAgent::cancel`): records which prompt the cancel lands on so
// a stuck "Cancelling…" can be attributed to delivery vs. processing.
@ -319,8 +353,8 @@ impl SessionActor {
// Closes the race between abort() and TurnActiveGuard drop:
// is_turn_active may still be true, causing InjectNotification
// to route Next-priority events to the buffer instead of
// pending_notifications. Moving them here ensures they survive
// in the notification queue (suppressed until next user prompt).
// pending_notifications. Moving them here ensures they survive in
// the queue; Ctrl+C defers their drain, while other cancels do not.
self.sweep_monitor_buffer_into_pending(&mut state, "monitor-cancel-drain");
// When killing all background tasks, also clear their pending
@ -333,7 +367,15 @@ impl SessionActor {
if let Some(task) = state.running_task.take() {
task.abort();
}
if let Some(gate) = &self.tool_context.task_wake_suppressed {
gate.set(false);
}
state.notifications_suppressed = false;
xai_grok_telemetry::unified_log::info(
"shell.task_wake.gate_cleared",
Some(self.session_info.id.0.as_ref()),
Some(serde_json::json!({ "reason": "rewind" })),
);
state.rewindable = false;
state.pending_inputs.pop_front()
} else {
@ -342,9 +384,7 @@ impl SessionActor {
let running_task = if rewound_input.is_some() {
None
} else {
let running_task = state.running_task.take();
state.notifications_suppressed = true;
running_task
state.running_task.take()
};
// Decide which queued inputs get resolved with `Cancelled` now vs.
@ -355,10 +395,10 @@ impl SessionActor {
// path that sends `Shutdown` next): drain the WHOLE queue — there
// is no point starting the next prompt and draining resolves every
// queued input's `respond_to` cleanly.
// * normal interactive cancel: remove ONLY the running turn,
// PRESERVING every queued prompt so the `Cancel` handler's
// follow-up `maybe_start_running_task` promotes the new front (the
// user's next queued prompt) and rebroadcasts `x.ai/queue/changed`.
// * normal cancel: remove the running turn; only Ctrl+C also removes
// queued terminal task-completion wakes. Preserve real user prompts
// and unrelated synthetic entries so `maybe_start_running_task` can
// promote the next genuine user turn.
// The cancelling client does not pull any prompt back into its
// input — the server queue is the single source of truth for what
// runs next. Previously every cancel did `std::mem::take`,
@ -391,6 +431,13 @@ impl SessionActor {
let is_running_turn = idx == 0;
if is_running_turn {
cancelled.push_back(item);
} else if suppress_task_wakes
&& matches!(&item.origin, super::PromptOrigin::TaskCompleted { .. })
{
if let Some(fallback) = item.task_wake_fallback {
Self::push_task_wake_fallback(&mut state, fallback);
}
Self::respond_removed_prompt(item.respond_to);
} else {
kept.push_back(item);
}
@ -490,6 +537,9 @@ impl SessionActor {
if let Some(running_task) = running_task {
running_task.abort();
}
if let Some(is_turn_active) = &self.tool_context.is_turn_active {
is_turn_active.store(false, std::sync::atomic::Ordering::Relaxed);
}
// The aborted turn's `BlockingWaitGuard`s drop asynchronously (they
// live in tool futures owned by the drainer task / subagent spawn
// task). Until they do, `queue_input` would read a stale depth > 0 and
@ -564,19 +614,14 @@ impl SessionActor {
return;
}
// Un-mark cancelled synthetic IDs so TaskCompletionReminder can
// report them on the next pass instead of permanently suppressing them.
if let Some(ref auto_wake) = self.tool_context.auto_wake_delivered {
for input in pending_inputs.iter() {
if let Some(id) = input.origin.completion_id() {
auto_wake.remove(id);
}
}
}
for (idx, input) in pending_inputs.into_iter().enumerate() {
// Running turn is idx 0; queued prompts never spent tokens.
let is_running_turn = idx == 0;
if let Some(task_id) = input.origin.completion_id()
&& let Some(reservations) = &self.tool_context.task_completion_reservations
{
reservations.release(task_id);
}
let _ = input
.respond_to
.send(Ok(PromptTurnOk {

60
crates/codegen/xai-grok-shell/src/session/acp_session_impl/tool_calls.rs
View file

@ -1064,6 +1064,15 @@ impl SessionActor {
self.permissions.set_classifier_transcript(turns);
}
}
let edit_path_context = matches!(&access_kind, AccessKind::Edit(_)).then(|| {
xai_grok_workspace::permission::types::EditPathContext {
real_cwd: std::path::PathBuf::from(self.session_info.cwd.as_str()),
display_cwd: self
.display_cwd
.get()
.map(|cwd| std::path::PathBuf::from(cwd.as_str())),
}
});
let decision = {
let _pending_guard =
crate::session::pending_interaction::PendingInteractionGuard::new(
@ -1074,9 +1083,10 @@ impl SessionActor {
crate::session::pending_interaction::PendingKind::Permission,
);
self.permissions
.request(
.request_with_edit_path_context(
access_kind.clone(),
tool_call_update,
edit_path_context,
Some(self.session_info.id.0.to_string()),
None,
None,
@ -1520,6 +1530,7 @@ impl SessionActor {
false,
None,
false,
None,
respond_to,
None,
None,
@ -1887,9 +1898,8 @@ impl SessionActor {
/// which is the same predicate used by `TaskCompletionReminder` —
/// they cannot drift because they share the function.
///
/// `AutoWakeDeliveredIds` is **deliberately not unmarked** here
/// (unlike `queue_input`'s preempt path). The tool result that
/// triggered this sweep IS the canonical consumption surface, and
/// Reservations are deliberately not released here because the tool result
/// that triggered this sweep is the canonical consumption surface, and
/// `TaskCompletionReminder` already suppresses the per-tool-call
/// reminder for these IDs via its own suppress list (also derived
/// from `consumed_completion_ids`). Un-marking here would risk a
@ -1905,18 +1915,26 @@ impl SessionActor {
return;
}
let mut state = self.state.lock().await;
let dropped_inputs = state
.sweep_pending_inputs(|i| {
i.origin
.completion_id()
.is_some_and(|id| consumed_ids.contains(&id))
})
.len();
let dropped = state.sweep_pending_inputs(|i| {
i.origin
.completion_id()
.is_some_and(|id| consumed_ids.contains(&id))
});
let dropped_inputs = dropped.len();
let before_notifications = state.pending_notifications.len();
state
.pending_notifications
.retain(|n| !consumed_ids.contains(&n.source.task_id()));
let dropped_notifications = before_notifications - state.pending_notifications.len();
drop(state);
if let Some(reservations) = &self.tool_context.task_completion_reservations {
for task_id in dropped
.iter()
.filter_map(|input| input.origin.completion_id())
{
reservations.release(task_id);
}
}
if dropped_inputs > 0 || dropped_notifications > 0 {
tracing::info!(
dropped_inputs, dropped_notifications, consumed_ids = ? consumed_ids,
@ -1937,8 +1955,26 @@ impl SessionActor {
/// returns. Real user inputs are preserved.
pub(super) async fn drop_pending_synthetic_items(&self) {
let mut state = self.state.lock().await;
state.pending_inputs.retain(|i| !i.origin.is_synthetic());
let mut kept = VecDeque::with_capacity(state.pending_inputs.len());
let mut dropped = Vec::new();
for input in std::mem::take(&mut state.pending_inputs) {
if input.origin.is_synthetic() {
dropped.push(input);
} else {
kept.push_back(input);
}
}
state.pending_inputs = kept;
state.pending_notifications.clear();
drop(state);
if let Some(reservations) = &self.tool_context.task_completion_reservations {
for task_id in dropped
.iter()
.filter_map(|input| input.origin.completion_id())
{
reservations.release(task_id);
}
}
}
/// Record git/PR ops from a successful tool result into session signals
/// (`turn_result.json`) and telemetry. Detection runs here at the shell's

20
crates/codegen/xai-grok-shell/src/session/acp_session_impl/turn.rs
View file

@ -238,7 +238,14 @@ impl SessionActor {
{ "prompt_id" : prompt_id, "block_count" : prompt_blocks.len(), }
)),
);
if !super::super::PromptOrigin::from_prompt_id(prompt_id).is_synthetic() {
let origin = super::super::PromptOrigin::from_prompt_id(prompt_id);
if let Some(completion_id) = origin.completion_id() {
self.mark_completions_reported(&[completion_id]).await;
if let Some(reservations) = &self.tool_context.task_completion_reservations {
reservations.release(completion_id);
}
}
if !origin.is_synthetic() {
self.cancel_pending_recap_for_new_prompt();
}
*self.turn_start_prompt_mode.lock() = prompt_mode;
@ -623,6 +630,17 @@ impl SessionActor {
self.maybe_inject_date_rollover_reminder().await;
self.inject_plan_mode_reminders().await;
self.inject_resumed_tasks_reminder();
if matches!(&origin, super::super::PromptOrigin::User) {
if let Some(gate) = &self.tool_context.task_wake_suppressed {
gate.set(false);
}
xai_grok_telemetry::unified_log::info(
"shell.task_wake.gate_cleared",
Some(self.session_info.id.0.as_ref()),
Some(serde_json::json!({ "reason" : "handle_prompt_user_start" })),
);
self.consume_deferred_completions_for_user_turn().await;
}
self.drain_between_turn_completions().await;
let user_message = if user_images.is_empty() {
user_message

95
crates/codegen/xai-grok-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs
View file

@ -285,6 +285,101 @@ async fn pre_flight_refresh_skips_api_key_auth_type() {
.await;
}
/// Hard-expired session token: pre-flight must call the refresher and must
/// not leave credentials stuck while pretending the JWT/config path applies.
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial(attribution_emit_count)]
async fn pre_flight_refreshes_hard_expired_session_token() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let called = Arc::new(AtomicBool::new(false));
let refresher: Arc<dyn crate::auth::refresh::TokenRefresher> =
Arc::new(AlwaysSucceedRefresher {
called: called.clone(),
});
let (_dir, am) = auth_manager_with_refresher(refresher);
assert!(
!am.has_usable_token(),
"precondition: access token is hard-expired"
);
let (actor, _rx) = make_actor_with_auth_manager(Some(am.clone())).await;
actor.refresh_token_if_expired().await;
assert!(
called.load(Ordering::SeqCst),
"pre-flight must invoke the refresher for a hard-expired session token"
);
assert_eq!(
actor
.chat_state_handle
.get_credentials()
.await
.api_key
.as_deref(),
Some("refreshed-test-token"),
"credentials must be updated to the refreshed bearer"
);
assert!(am.has_usable_token());
})
.await;
}
/// Hard-expired + failed refresh: do not fall through to JWT/config.toml;
/// leave credentials unchanged so 401 recovery remains the safety net.
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial(attribution_emit_count)]
async fn pre_flight_hard_expired_refresh_failure_skips_jwt_fallthrough() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let call_count = Arc::new(std::sync::atomic::AtomicU32::new(0));
let refresher: Arc<dyn crate::auth::refresh::TokenRefresher> = Arc::new({
struct AlwaysFail(Arc<std::sync::atomic::AtomicU32>);
#[async_trait::async_trait]
impl crate::auth::refresh::TokenRefresher for AlwaysFail {
async fn refresh(
&self,
_: crate::auth::refresh::RefreshReason,
) -> crate::auth::refresh::RefreshOutcome {
self.0.fetch_add(1, Ordering::SeqCst);
crate::auth::refresh::RefreshOutcome::transient("refresh failed")
}
}
AlwaysFail(call_count.clone())
});
let (_dir, am) = auth_manager_with_refresher(refresher);
let (actor, _rx) = make_actor_with_auth_manager(Some(am.clone())).await;
actor.refresh_token_if_expired().await;
assert!(
call_count.load(Ordering::SeqCst) >= 1,
"pre-flight must attempt refresh"
);
assert_eq!(
actor
.chat_state_handle
.get_credentials()
.await
.api_key
.as_deref(),
Some("initial-test-key"),
"failed hard-expired pre-flight must not invent a JWT/config bearer"
);
assert!(
!am.has_usable_token(),
"token remains hard-expired after failed refresh"
);
assert!(
am.permanent_failure().is_none(),
"transient refresh failure must not poison permanent_failure"
);
})
.await;
}
/// Proactive refresh keeps the cache hot so `refresh_token_if_expired`
/// (per-turn pre-flight) is a cache hit — the refresher fires once
/// (proactive), then the per-turn call sees the fresh token without

600
crates/codegen/xai-grok-shell/src/session/acp_session_tests/auto_wake_suppression_tests.rs
View file

@ -60,6 +60,38 @@ fn bash_completed_notification(task_id: &str) -> PendingNotification {
},
}
}
fn monitor_completed_notification(task_id: &str) -> PendingNotification {
PendingNotification {
prompt_id: format!("monitor-completed-{task_id}"),
prompt_blocks: vec![],
priority: NotificationPriority::Later,
source: NotificationSource::MonitorCompleted {
task_id: task_id.to_string(),
},
}
}
fn task_wake_admission(
task_id: &str,
source: NotificationSource,
) -> (
crate::session::commands::TaskWakeAdmission,
oneshot::Receiver<bool>,
) {
let (respond_to, response_rx) = oneshot::channel();
(
crate::session::commands::TaskWakeAdmission {
respond_to,
fallback: crate::session::commands::TaskWakeFallback {
prompt_id: format!("deferred-{task_id}"),
prompt_blocks: vec![acp::ContentBlock::Text(acp::TextContent::new(format!(
"completion {task_id}"
)))],
source,
},
},
response_rx,
)
}
fn monitor_event_notification(task_id: &str) -> PendingNotification {
PendingNotification {
prompt_id: format!("monitor-{task_id}"),
@ -73,6 +105,30 @@ fn monitor_event_notification(task_id: &str) -> PendingNotification {
/// Monitor notifications in the idle drain collapse into ONE
/// `format_monitor_events` block (same shape as the mid-turn injection);
/// non-monitor notifications keep their raw blocks, `---`-separated.
#[test]
fn pending_notification_cap_keeps_newest_entries() {
let mut state = State {
running_task: None,
pending_inputs: std::collections::VecDeque::new(),
pending_notifications: Vec::new(),
notifications_suppressed: true,
rewindable: false,
nudges_used_this_session: 0,
};
for index in 0..(MAX_PENDING_NOTIFICATIONS + 3) {
SessionActor::push_pending_notification(
&mut state,
bash_completed_notification(&format!("task-{index}")),
);
}
assert_eq!(state.pending_notifications.len(), MAX_PENDING_NOTIFICATIONS);
assert_eq!(state.pending_notifications[0].source.task_id(), "task-3");
let newest = format!("task-{}", MAX_PENDING_NOTIFICATIONS + 2);
assert_eq!(
state.pending_notifications.last().unwrap().source.task_id(),
newest
);
}
#[tokio::test(flavor = "current_thread")]
async fn drain_batches_monitor_notifications_into_formatted_block() {
let local = tokio::task::LocalSet::new();
@ -144,6 +200,514 @@ async fn drain_batches_monitor_notifications_into_formatted_block() {
})
.await;
}
#[tokio::test(flavor = "current_thread")]
async fn cancel_barrier_rejects_task_completion_wake_without_reporting_it() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (gateway_tx, _) =
tokio::sync::mpsc::unbounded_channel::<xai_acp_lib::AcpClientMessage>();
let (persistence_tx, _) = tokio::sync::mpsc::unbounded_channel::<PersistenceMsg>();
let actor = create_test_actor(0, 256_000, 85, gateway_tx, persistence_tx).await;
let reservations = actor
.tool_context
.task_completion_reservations
.clone()
.expect("completion reservations");
reservations.reserve("bg-suppressed".to_string());
actor.state.lock().await.notifications_suppressed = true;
let gate = actor
.tool_context
.task_wake_suppressed
.clone()
.expect("task-wake gate");
gate.set(true);
let resources = actor
.agent
.borrow()
.tool_bridge()
.clone()
.shared_resources()
.await;
{
let mut resources = resources.lock().await;
resources.insert(reservations.clone());
resources.insert(gate.clone());
}
let origin = crate::session::PromptOrigin::TaskCompleted {
task_id: "bg-suppressed".to_string(),
};
let (admission, response_rx) = task_wake_admission(
"bg-suppressed",
NotificationSource::BashTaskCompleted {
task_id: "bg-suppressed".to_string(),
},
);
assert!(
actor
.admit_task_completion_wake(&origin, admission)
.await
.is_none()
);
assert_eq!(response_rx.await, Ok(false));
assert!(gate.get());
let state = actor.state.lock().await;
assert!(state.running_task.is_none());
assert!(state.pending_inputs.is_empty());
assert!(
matches!(state.pending_notifications.as_slice(), [PendingNotification {
source : NotificationSource::BashTaskCompleted { task_id }, .. }] if
task_id == "bg-suppressed")
);
drop(state);
assert!(reservations.contains("bg-suppressed"));
let res = resources.lock().await;
assert!(
res.get::<xai_grok_tools::types::resources::State<
xai_grok_tools::reminders::task_completion::ReportedTaskCompletions,
>>()
.is_none(),
"declined admission must not report before user re-engagement"
);
drop(res);
let reminder = xai_grok_tools::reminders::TaskCompletionReminder;
let reminders = xai_grok_tools::types::tool::Reminder::collect_reminders(
&reminder,
resources,
&ToolOutput::Dynamic(serde_json::Value::Null.into()),
)
.await;
assert!(reminders.is_empty());
assert!(reservations.contains("bg-suppressed"));
reservations.release("bg-suppressed");
})
.await;
}
#[tokio::test(flavor = "current_thread")]
async fn closed_admission_ack_stores_fallback_before_prompt_rejection() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (gateway_tx, _) =
tokio::sync::mpsc::unbounded_channel::<xai_acp_lib::AcpClientMessage>();
let (persistence_tx, _) = tokio::sync::mpsc::unbounded_channel::<PersistenceMsg>();
let actor = create_test_actor(0, 256_000, 85, gateway_tx, persistence_tx).await;
let origin = crate::session::PromptOrigin::TaskCompleted {
task_id: "mon-timeout".to_string(),
};
let (admission, response_rx) = task_wake_admission(
"mon-timeout",
NotificationSource::MonitorCompleted {
task_id: "mon-timeout".to_string(),
},
);
drop(response_rx);
assert!(
actor
.admit_task_completion_wake(&origin, admission)
.await
.is_none()
);
let state = actor.state.lock().await;
assert!(
matches!(state.pending_notifications.as_slice(), [PendingNotification {
source : NotificationSource::MonitorCompleted { task_id }, .. }] if
task_id == "mon-timeout")
);
})
.await;
}
#[tokio::test(flavor = "current_thread")]
async fn non_task_prompt_is_not_subject_to_task_wake_barrier() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (gateway_tx, _) =
tokio::sync::mpsc::unbounded_channel::<xai_acp_lib::AcpClientMessage>();
let (persistence_tx, _) = tokio::sync::mpsc::unbounded_channel::<PersistenceMsg>();
let actor = create_test_actor(0, 256_000, 85, gateway_tx, persistence_tx).await;
actor.state.lock().await.notifications_suppressed = true;
let (admission, response_rx) = task_wake_admission(
"sub-1",
NotificationSource::BashTaskCompleted {
task_id: "sub-1".to_string(),
},
);
assert!(
actor
.admit_task_completion_wake(
&crate::session::PromptOrigin::SubagentCompleted {
subagent_id: "sub-1".to_string(),
},
admission,
)
.await
.is_some(),
"subagent completion is outside terminal task-wake suppression scope"
);
assert_eq!(response_rx.await, Ok(true));
})
.await;
}
#[tokio::test(flavor = "current_thread")]
async fn task_completion_wake_is_admitted_without_cancel_barrier() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (gateway_tx, _) =
tokio::sync::mpsc::unbounded_channel::<xai_acp_lib::AcpClientMessage>();
let (persistence_tx, _) = tokio::sync::mpsc::unbounded_channel::<PersistenceMsg>();
let actor = std::sync::Arc::new(
create_test_actor(0, 256_000, 85, gateway_tx, persistence_tx).await,
);
let origin = crate::session::PromptOrigin::TaskCompleted {
task_id: "bg-normal".to_string(),
};
actor
.tool_context
.task_completion_reservations
.as_ref()
.expect("completion reservations")
.reserve("bg-normal".to_string());
let (admission, response_rx) = task_wake_admission(
"bg-normal",
NotificationSource::BashTaskCompleted {
task_id: "bg-normal".to_string(),
},
);
let fallback = actor
.admit_task_completion_wake(&origin, admission)
.await
.expect("normal task wake should be admitted");
assert_eq!(response_rx.await, Ok(true));
let (respond_to, _rx) = oneshot::channel();
actor
.queue_input(
vec![],
"task-completed-bg-normal".to_string(),
PromptMode::Agent,
None,
None,
None,
None,
true,
None,
false,
Some(fallback),
respond_to,
None,
None,
)
.await;
let state = actor.state.lock().await;
assert_eq!(state.pending_inputs.len(), 1);
assert!(
matches!(state.pending_inputs.front().map(| item | & item.origin),
Some(crate ::session::PromptOrigin::TaskCompleted { task_id }) if task_id
== "bg-normal")
);
drop(state);
let resources = actor
.agent
.borrow()
.tool_bridge()
.clone()
.shared_resources()
.await;
assert!(
resources
.lock()
.await
.get::<xai_grok_tools::types::resources::State<
xai_grok_tools::reminders::task_completion::ReportedTaskCompletions,
>>()
.is_none(),
"queue acceptance alone must not mark the completion reported"
);
let actor_for_turn = actor.clone();
let turn = tokio::task::spawn_local(async move {
actor_for_turn
.handle_prompt(
"task-completed-bg-normal",
vec![acp::ContentBlock::Text(acp::TextContent::new("done"))],
PromptMode::Agent,
None,
None,
None,
None,
true,
None,
None,
None,
)
.await
});
tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
if already_reported(&actor, "bg-normal").await {
break;
}
tokio::task::yield_now().await;
}
})
.await
.expect("synthetic turn marked completion reported");
turn.abort();
assert!(
already_reported(&actor, "bg-normal").await,
"actual synthetic turn start must mark the completion reported"
);
assert!(
actor
.tool_context
.task_completion_reservations
.as_ref()
.is_none_or(|ids| !ids.contains("bg-normal"))
);
})
.await;
}
#[tokio::test(flavor = "current_thread")]
async fn genuine_user_start_consumes_deferred_completions_without_notification_turn() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (gateway_tx, _) =
tokio::sync::mpsc::unbounded_channel::<xai_acp_lib::AcpClientMessage>();
let (persistence_tx, _) = tokio::sync::mpsc::unbounded_channel::<PersistenceMsg>();
let actor = std::sync::Arc::new(
create_test_actor(0, 256_000, 85, gateway_tx, persistence_tx).await,
);
let body = xai_grok_tools::reminders::task_completion::format_monitor_completion(
&xai_grok_tools::types::TaskSnapshot {
task_id: "mon-quiet".to_string(),
command: "tail -f quiet.log".to_string(),
display_command: Some("[monitor] quiet logs".to_string()),
cwd: String::new(),
start_time: std::time::SystemTime::now(),
end_time: Some(std::time::SystemTime::now()),
output: String::new(),
output_file: std::path::PathBuf::new(),
truncated: false,
exit_code: Some(0),
signal: None,
completed: true,
kind: xai_grok_tools::computer::types::TaskKind::Monitor,
block_waited: false,
explicitly_killed: false,
owner_session_id: None,
},
Some("get_command_or_subagent_output"),
);
{
let mut state = actor.state.lock().await;
state.notifications_suppressed = true;
state
.pending_notifications
.push(monitor_event_notification("mon-quiet"));
let mut monitor_completion = monitor_completed_notification("mon-quiet");
monitor_completion.prompt_blocks =
vec![acp::ContentBlock::Text(acp::TextContent::new(body))];
state.pending_notifications.push(monitor_completion);
let mut bash_completion = bash_completed_notification("bash-deferred");
bash_completion.prompt_blocks = vec![acp::ContentBlock::Text(
acp::TextContent::new("Background task bash-deferred completed."),
)];
state.pending_notifications.push(bash_completion);
}
let reservations = actor
.tool_context
.task_completion_reservations
.as_ref()
.expect("completion reservations");
reservations.reserve("mon-quiet".to_string());
reservations.reserve("bash-deferred".to_string());
actor
.tool_context
.task_wake_suppressed
.as_ref()
.expect("task-wake gate")
.set(true);
let actor_for_turn = actor.clone();
let turn = tokio::task::spawn_local(async move {
actor_for_turn
.handle_prompt(
"user-deferred-completions",
vec![acp::ContentBlock::Text(acp::TextContent::new("continue"))],
PromptMode::Agent,
None,
None,
None,
None,
false,
None,
None,
None,
)
.await
});
tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
if actor.state.lock().await.pending_notifications.is_empty() {
break;
}
tokio::task::yield_now().await;
}
})
.await
.expect("user turn consumed deferred completions");
turn.abort();
tokio::task::yield_now().await;
let state = actor.state.lock().await;
assert!(state.notifications_suppressed);
assert!(state.pending_notifications.is_empty());
assert!(state.pending_inputs.iter().all(|input| !matches!(
input.origin,
crate::session::PromptOrigin::NotificationDrain
)));
drop(state);
let (completion_tx, _completion_rx) = tokio::sync::mpsc::unbounded_channel();
SessionActor::maybe_drain_notifications(actor.clone(), completion_tx).await;
let state = actor.state.lock().await;
assert!(state.pending_inputs.iter().all(|input| !matches!(
input.origin,
crate::session::PromptOrigin::NotificationDrain
)));
drop(state);
assert!(!reservations.contains("mon-quiet"));
assert!(!reservations.contains("bash-deferred"));
assert!(
!actor
.tool_context
.task_wake_suppressed
.as_ref()
.expect("task-wake gate")
.get()
);
let conversation = actor.chat_state_handle.get_conversation().await;
let text = conversation
.iter()
.map(|item| item.text_content())
.collect::<Vec<_>>()
.join("\n");
assert!(text.contains("Monitor \"mon-quiet\" ended"));
assert!(text.contains("Background task bash-deferred completed."));
assert!(!text.contains("<monitor-event"));
})
.await;
}
#[tokio::test(flavor = "current_thread")]
async fn accepted_reservation_survives_user_start() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (gateway_tx, _) =
tokio::sync::mpsc::unbounded_channel::<xai_acp_lib::AcpClientMessage>();
let (persistence_tx, _) = tokio::sync::mpsc::unbounded_channel::<PersistenceMsg>();
let actor = std::sync::Arc::new(
create_test_actor(0, 256_000, 85, gateway_tx, persistence_tx).await,
);
let reservations = actor
.tool_context
.task_completion_reservations
.as_ref()
.expect("completion reservations");
reservations.reserve("accepted-race".to_string());
actor
.tool_context
.task_wake_suppressed
.as_ref()
.expect("task-wake gate")
.set(true);
let actor_for_turn = actor.clone();
let turn = tokio::task::spawn_local(async move {
actor_for_turn
.handle_prompt(
"user-accepted-race",
vec![acp::ContentBlock::Text(acp::TextContent::new("continue"))],
PromptMode::Agent,
None,
None,
None,
None,
false,
None,
None,
None,
)
.await
});
tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
if actor
.tool_context
.task_wake_suppressed
.as_ref()
.is_none_or(|gate| !gate.get())
{
break;
}
tokio::task::yield_now().await;
}
})
.await
.expect("user turn started");
assert!(reservations.contains("accepted-race"));
turn.abort();
reservations.release("accepted-race");
})
.await;
}
#[tokio::test(flavor = "current_thread")]
async fn same_id_bash_completion_does_not_suppress_monitor_event() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (gateway_tx, _) = tokio::sync::mpsc::unbounded_channel::<
xai_acp_lib::AcpClientMessage,
>();
let (persistence_tx, _) = tokio::sync::mpsc::unbounded_channel::<
PersistenceMsg,
>();
let actor = create_test_actor(0, 256_000, 85, gateway_tx, persistence_tx)
.await;
let monitor = PendingNotification {
prompt_id: "monitor-shared".to_string(),
prompt_blocks: vec![
acp::ContentBlock::Text(acp::TextContent::new("<monitor-event description=\"watch\" task_id=\"shared\">\nstdout\n</monitor-event>",))
],
priority: NotificationPriority::Next,
source: NotificationSource::MonitorEvent {
task_id: "shared".to_string(),
},
};
let mut bash = bash_completed_notification("shared");
bash.prompt_blocks = vec![
acp::ContentBlock::Text(acp::TextContent::new("Background task shared completed.",))
];
let mut state = actor.state.lock().await;
SessionActor::drain_notifications_into_turn(
&mut state,
vec![monitor, bash],
"get_command_or_subagent_output",
);
let text = state
.pending_inputs
.back()
.expect("drained turn")
.prompt_blocks
.iter()
.filter_map(|block| match block {
acp::ContentBlock::Text(text) => Some(text.text.as_str()),
_ => None,
})
.collect::<Vec<_>>()
.join("\n");
assert!(text.contains("<monitor-event task_id=\"shared\">"));
assert!(text.contains("Background task shared completed."));
})
.await;
}
/// Fix 1, TaskOutput(completed) — the matching pending `task-completed-{id}`
/// input must be dropped; any non-matching synthetic prompt must survive.
#[tokio::test(flavor = "current_thread")]
@ -248,6 +812,12 @@ async fn user_prompt_preempt_keeps_running_synthetic_slot() {
tokio::sync::mpsc::unbounded_channel::<xai_acp_lib::AcpClientMessage>();
let (persistence_tx, _) = tokio::sync::mpsc::unbounded_channel::<PersistenceMsg>();
let actor = create_test_actor(0, 256_000, 85, gateway_tx, persistence_tx).await;
let reservations = actor
.tool_context
.task_completion_reservations
.as_ref()
.expect("completion reservations");
reservations.reserve("bg-other".to_string());
{
let mut state = actor.state.lock().await;
state.running_task = Some(running_task_stub("task-completed-bg-target"));
@ -271,6 +841,7 @@ async fn user_prompt_preempt_keeps_running_synthetic_slot() {
false,
None,
false,
None,
respond_to,
None,
None,
@ -288,6 +859,10 @@ async fn user_prompt_preempt_keeps_running_synthetic_slot() {
"the running synthetic turn's slot must survive the user-priority \
preempt; only the queued non-running synthetic is dropped"
);
assert!(
!reservations.contains("bg-other"),
"ordinary user-priority preemption releases ownership immediately"
);
})
.await;
}
@ -902,16 +1477,15 @@ async fn reparented_record_is_noop_without_goal_harness() {
}
/// Regression: the between-turn completion drain must suppress subagent
/// completions already delivered to the model via auto-wake synthetic
/// prompts. Without `auto_wake_delivered` feeding `suppress_ids`, the same
/// prompts. Without completion reservations feeding `suppress_ids`, the same
/// completion is reported twice — once as the auto-wake "Background subagent
/// … completed" prompt and again as the "While you were idle, N background
/// subagent(s) completed" reminder.
#[tokio::test(flavor = "current_thread")]
async fn between_turn_drain_suppresses_auto_wake_delivered_subagents() {
async fn between_turn_drain_suppresses_reserved_subagents() {
use xai_grok_tools::implementations::grok_build::task::types::{
SubagentCompletionSummary, SubagentEvent,
};
use xai_grok_tools::reminders::task_completion::AutoWakeDeliveredIds;
let local = tokio::task::LocalSet::new();
local
.run_until(async {
@ -919,9 +1493,12 @@ async fn between_turn_drain_suppresses_auto_wake_delivered_subagents() {
tokio::sync::mpsc::unbounded_channel::<xai_acp_lib::AcpClientMessage>();
let (persistence_tx, _) = tokio::sync::mpsc::unbounded_channel::<PersistenceMsg>();
let mut actor = create_test_actor(0, 256_000, 85, gateway_tx, persistence_tx).await;
let auto_wake = AutoWakeDeliveredIds::default();
auto_wake.insert("sa-autowake".to_string());
actor.tool_context.auto_wake_delivered = Some(auto_wake);
actor
.tool_context
.task_completion_reservations
.as_ref()
.expect("completion reservations")
.reserve("sa-autowake".to_string());
let captured: std::sync::Arc<std::sync::Mutex<Vec<String>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<SubagentEvent>();
@ -951,7 +1528,7 @@ async fn between_turn_drain_suppresses_auto_wake_delivered_subagents() {
let suppress = captured.lock().unwrap().clone();
assert!(
suppress.contains(&"sa-autowake".to_string()),
"between-turn drain must pass auto-wake-delivered ids as suppress_ids: \
"between-turn drain must pass reserved ids as suppress_ids: \
{suppress:?}",
);
let conversation = actor.chat_state_handle.get_conversation().await;
@ -966,7 +1543,14 @@ async fn between_turn_drain_suppresses_auto_wake_delivered_subagents() {
);
assert!(
!texts.contains("sa-autowake"),
"auto-wake-delivered completion must NOT be re-surfaced: {texts}",
"reserved completion must NOT be re-surfaced: {texts}",
);
assert!(
actor
.tool_context
.task_completion_reservations
.as_ref()
.is_some_and(|ids| ids.contains("sa-autowake"))
);
})
.await;

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

@ -157,6 +157,7 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
previous_model: std::cell::Cell::new(None),
compaction_mode: xai_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: crate::util::config::CompactionToolChoice::Auto,
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@ -609,6 +610,7 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
previous_model: std::cell::Cell::new(None),
compaction_mode: xai_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: crate::util::config::CompactionToolChoice::Auto,
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@ -882,6 +884,7 @@ async fn cancel_running_task_teardown_clears_running_and_pending_work() {
previous_model: std::cell::Cell::new(None),
compaction_mode: xai_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: crate::util::config::CompactionToolChoice::Auto,
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@ -1052,6 +1055,7 @@ async fn cancel_running_task_teardown_clears_running_and_pending_work() {
verbatim: false,
json_schema: None,
origin: crate::session::PromptOrigin::User,
task_wake_fallback: None,
respond_to: tx,
persist_ack: None,
parsed_prompt_tx: None,
@ -1488,6 +1492,7 @@ async fn cancel_running_task_interactive_preserves_queued_work() {
verbatim: false,
json_schema: None,
origin: crate::session::PromptOrigin::User,
task_wake_fallback: None,
respond_to,
persist_ack: None,
parsed_prompt_tx: None,
@ -1650,6 +1655,191 @@ async fn cancel_after_own_completion_sweep_preserves_queued_user_prompt() {
})
.await;
}
#[tokio::test(flavor = "current_thread")]
async fn interactive_cancel_drops_queued_task_wakes_and_promotes_user() {
use tokio::sync::oneshot::error::TryRecvError;
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (gateway_tx, _gateway_rx) =
tokio::sync::mpsc::unbounded_channel::<xai_acp_lib::AcpClientMessage>();
let (persistence_tx, _persistence_rx) =
tokio::sync::mpsc::unbounded_channel::<PersistenceMsg>();
let actor = create_test_actor(0, 256_000, 85, gateway_tx, persistence_tx).await;
let reservations = actor
.tool_context
.task_completion_reservations
.clone()
.expect("completion reservations");
reservations.reserve("bg-queued".to_string());
let actor = Arc::new(actor);
let (running_item, mut running_rx) =
input_with_origin_rx("user-running", crate::session::PromptOrigin::User);
let (mut wake_item, mut wake_rx) = input_with_origin_rx(
"task-completed-bg-queued",
crate::session::PromptOrigin::TaskCompleted {
task_id: "bg-queued".to_string(),
},
);
wake_item.task_wake_fallback = Some(crate::session::commands::TaskWakeFallback {
prompt_id: "bash-completed-bg-queued".to_string(),
prompt_blocks: vec![acp::ContentBlock::Text(acp::TextContent::new(
"completion bg-queued",
))],
source: NotificationSource::BashTaskCompleted {
task_id: "bg-queued".to_string(),
},
});
let (queued_user, mut queued_user_rx) =
input_with_origin_rx("user-next", crate::session::PromptOrigin::User);
{
let mut state = actor.state.lock().await;
state.running_task = Some(running_task_stub("user-running"));
state.pending_inputs.push_back(running_item);
state.pending_inputs.push_back(wake_item);
state.pending_inputs.push_back(queued_user);
}
let cancel = actor.cancel_running_task(true, false, false, Some("ctrl_c".to_string()));
tokio::pin!(cancel);
tokio::select! {
_ = & mut cancel => {} _ = tokio::task::yield_now() => { assert!(actor
.state.try_lock().expect("state lock").notifications_suppressed,
"Ctrl+C must arm actor suppression before the first await");
assert!(actor.tool_context.task_wake_suppressed.as_ref().is_some_and(|
gate | gate.get()),
"Ctrl+C must arm the reminder gate before the first await"); cancel.
await; }
}
assert!(
actor
.tool_context
.task_wake_suppressed
.as_ref()
.is_some_and(|gate| gate.get()),
"Ctrl+C must synchronously arm the reminder gate"
);
{
let state = actor.state.lock().await;
let remaining: Vec<&str> = state
.pending_inputs
.iter()
.map(|item| item.prompt_id.as_str())
.collect();
assert_eq!(remaining, vec!["user-next"]);
assert!(
matches!(state.pending_notifications.as_slice(), [PendingNotification
{ source : NotificationSource::BashTaskCompleted { task_id }, .. }]
if task_id == "bg-queued")
);
assert!(state.notifications_suppressed);
}
assert!(matches!(running_rx.try_recv(), Ok(Ok(_))));
assert!(matches!(
wake_rx.try_recv(),
Ok(Ok(crate::session::commands::PromptTurnOk {
completion_kind: PromptCompletionKind::RemovedFromQueue,
..
}))
));
assert!(matches!(
queued_user_rx.try_recv(),
Err(TryRecvError::Empty)
));
assert!(reservations.contains("bg-queued"));
actor.consume_deferred_completions_for_user_turn().await;
{
let state = actor.state.lock().await;
assert!(
state.pending_notifications.is_empty(),
"the genuine user turn must consume the parked fallback exactly once"
);
}
assert!(!reservations.contains("bg-queued"));
actor.consume_deferred_completions_for_user_turn().await;
assert!(
actor.state.lock().await.pending_notifications.is_empty(),
"a second user-start drain must not rediscover the completion"
);
})
.await;
}
#[tokio::test(flavor = "current_thread")]
async fn ctrl_c_clears_turn_active_before_background_completion_routes() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (gateway_tx, _gateway_rx) =
tokio::sync::mpsc::unbounded_channel::<xai_acp_lib::AcpClientMessage>();
let (persistence_tx, _persistence_rx) =
tokio::sync::mpsc::unbounded_channel::<PersistenceMsg>();
let mut actor = create_test_actor(0, 256_000, 85, gateway_tx, persistence_tx).await;
let is_turn_active = Arc::new(std::sync::atomic::AtomicBool::new(true));
actor.tool_context.is_turn_active = Some(is_turn_active.clone());
let (running_item, _running_rx) =
input_with_origin_rx("user-running", crate::session::PromptOrigin::User);
{
let mut state = actor.state.lock().await;
state.running_task = Some(running_task_stub("user-running"));
state.pending_inputs.push_back(running_item);
}
actor
.cancel_running_task(true, false, false, Some("ctrl_c".to_string()))
.await;
assert!(!is_turn_active.load(std::sync::atomic::Ordering::Relaxed));
assert!(actor.state.lock().await.notifications_suppressed);
})
.await;
}
#[tokio::test(flavor = "current_thread")]
async fn non_ctrl_c_cancel_preserves_queued_task_wakes_and_does_not_arm_barrier() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
for trigger in [Some("esc"), Some("mouse"), Some("dashboard"), None] {
let (gateway_tx, _gateway_rx) =
tokio::sync::mpsc::unbounded_channel::<xai_acp_lib::AcpClientMessage>();
let (persistence_tx, _persistence_rx) =
tokio::sync::mpsc::unbounded_channel::<PersistenceMsg>();
let actor = create_test_actor(0, 256_000, 85, gateway_tx, persistence_tx).await;
let (running_item, _running_rx) =
input_with_origin_rx("user-running", crate::session::PromptOrigin::User);
let (wake_item, _wake_rx) = input_with_origin_rx(
"task-completed-bg-preserved",
crate::session::PromptOrigin::TaskCompleted {
task_id: "bg-preserved".to_string(),
},
);
let (queued_user, _queued_user_rx) =
input_with_origin_rx("user-next", crate::session::PromptOrigin::User);
{
let mut state = actor.state.lock().await;
state.running_task = Some(running_task_stub("user-running"));
state.pending_inputs.push_back(running_item);
state.pending_inputs.push_back(wake_item);
state.pending_inputs.push_back(queued_user);
}
actor
.cancel_running_task(true, false, false, trigger.map(str::to_string))
.await;
let state = actor.state.lock().await;
let remaining: Vec<&str> = state
.pending_inputs
.iter()
.map(|item| item.prompt_id.as_str())
.collect();
assert_eq!(
remaining,
vec!["task-completed-bg-preserved", "user-next"],
"non-Ctrl+C cancel {trigger:?} must preserve the queued task wake"
);
assert!(
!state.notifications_suppressed,
"non-Ctrl+C cancel {trigger:?} must not arm task-wake suppression"
);
}
})
.await;
}
/// Regression for the cancel-spinner hang: an interactive cancel must resolve
/// the in-flight front prompt's `respond_to` with `Cancelled` even when
/// `state.running_task` is `None`.
@ -1682,6 +1872,7 @@ async fn cancel_resolves_front_when_running_task_is_none() {
verbatim: false,
json_schema: None,
origin: crate::session::PromptOrigin::User,
task_wake_fallback: None,
respond_to,
persist_ack: None,
parsed_prompt_tx: None,
@ -1926,6 +2117,7 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() {
previous_model: std::cell::Cell::new(None),
compaction_mode: xai_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: crate::util::config::CompactionToolChoice::Auto,
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@ -2209,6 +2401,7 @@ async fn cancel_keeps_remaining_queued_prompts_visible_to_clients() {
verbatim: false,
json_schema: None,
origin: crate::session::PromptOrigin::User,
task_wake_fallback: None,
respond_to,
persist_ack: None,
parsed_prompt_tx: None,

2
crates/codegen/xai-grok-shell/src/session/acp_session_tests/goal/goal_backoff_tests.rs
View file

@ -193,6 +193,7 @@ async fn seed_pending_classifier_nudge(actor: &SessionActor) {
verbatim: true,
json_schema: None,
origin: crate::session::PromptOrigin::GoalClassifierNudge,
task_wake_fallback: None,
respond_to,
persist_ack: None,
parsed_prompt_tx: None,
@ -2937,6 +2938,7 @@ async fn idempotency_matcher_suppresses_goal_summary_when_classifier_nudge_pendi
verbatim: true,
json_schema: None,
origin: crate::session::PromptOrigin::GoalClassifierNudge,
task_wake_fallback: None,
respond_to,
persist_ack: None,
parsed_prompt_tx: None,

1
crates/codegen/xai-grok-shell/src/session/acp_session_tests/goal/goal_classifier_e2e_tests.rs
View file

@ -1683,6 +1683,7 @@ async fn goal_classifier_nudge_suppresses_subsequent_goal_summary() {
verbatim: true,
json_schema: None,
origin: PromptOrigin::GoalClassifierNudge,
task_wake_fallback: None,
respond_to,
persist_ack: None,
parsed_prompt_tx: None,

1
crates/codegen/xai-grok-shell/src/session/acp_session_tests/idle_resume_tests.rs
View file

@ -181,6 +181,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
previous_model: std::cell::Cell::new(None),
compaction_mode: xai_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: crate::util::config::CompactionToolChoice::Auto,
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},

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

@ -111,6 +111,7 @@ async fn create_test_actor(
previous_model: std::cell::Cell::new(None),
compaction_mode: xai_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: crate::util::config::CompactionToolChoice::Auto,
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@ -541,6 +542,7 @@ async fn create_test_actor_with_memory(
previous_model: std::cell::Cell::new(None),
compaction_mode: xai_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: crate::util::config::CompactionToolChoice::Auto,
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@ -1309,6 +1311,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
previous_model: std::cell::Cell::new(None),
compaction_mode: xai_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: crate::util::config::CompactionToolChoice::Auto,
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},

1
crates/codegen/xai-grok-shell/src/session/acp_session_tests/laziness/laziness_integration_tests.rs
View file

@ -336,6 +336,7 @@ async fn idle_recheck_after_sleep_short_circuits_silently() {
verbatim: true,
json_schema: None,
origin: crate::session::PromptOrigin::User,
task_wake_fallback: None,
respond_to,
persist_ack: None,
parsed_prompt_tx: None,

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

@ -161,6 +161,7 @@ async fn create_test_actor_with_memory(
previous_model: std::cell::Cell::new(None),
compaction_mode: xai_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: crate::util::config::CompactionToolChoice::Auto,
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},

11
crates/codegen/xai-grok-shell/src/session/acp_session_tests/prompt_queue_actor_tests.rs
View file

@ -842,6 +842,7 @@ async fn queue_input_send_now_inserts_behind_running_front_and_requests_cancel()
false,
None,
/* send_now */ true,
None,
respond_to,
None,
None,
@ -902,6 +903,7 @@ async fn queue_input_stacked_send_now_prompts_insert_fifo_during_goal_turn() {
false,
None,
/* send_now */ true,
None,
respond_to,
None,
None,
@ -954,6 +956,7 @@ async fn queue_input_auto_send_now_only_inside_wait_window() {
false,
None,
false,
None,
respond_to,
None,
None,
@ -978,6 +981,7 @@ async fn queue_input_auto_send_now_only_inside_wait_window() {
false,
None,
false,
None,
respond_to,
None,
None,
@ -1036,6 +1040,7 @@ async fn queue_input_auto_send_now_when_wait_and_held_queue_empty() {
false,
None,
false,
None,
respond_to,
None,
None,
@ -1073,6 +1078,7 @@ async fn queue_input_auto_send_now_when_wait_and_held_queue_empty() {
false,
None,
false,
None,
respond_to,
None,
None,
@ -1140,6 +1146,7 @@ async fn queue_input_auto_send_now_during_foreground_subagent_await_window() {
false,
None,
false,
None,
respond_to,
None,
None,
@ -1170,6 +1177,7 @@ async fn queue_input_auto_send_now_during_foreground_subagent_await_window() {
false,
None,
false,
None,
respond_to,
None,
None,
@ -1229,6 +1237,7 @@ async fn queue_input_send_now_exempts_synthetic_and_goal_turns() {
true,
None,
false,
None,
respond_to,
None,
None,
@ -1253,6 +1262,7 @@ async fn queue_input_send_now_exempts_synthetic_and_goal_turns() {
false,
None,
true,
None,
respond_to,
None,
None,
@ -1380,6 +1390,7 @@ async fn queue_input_send_now_pins_front_on_running_task_identity() {
false,
None,
/* send_now */ true,
None,
respond_to,
None,
None,

2
crates/codegen/xai-grok-shell/src/session/acp_session_tests/recap_display_only_tests.rs
View file

@ -63,6 +63,7 @@ async fn queue_input_user_prompt_bumps_recap_epoch() {
false,
None,
false,
None,
respond_to,
None,
None,
@ -101,6 +102,7 @@ async fn queue_input_synthetic_does_not_bump_recap_epoch() {
false,
None,
false,
None,
respond_to,
None,
None,

1
crates/codegen/xai-grok-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs
View file

@ -115,6 +115,7 @@ pub(super) async fn make_replay_send_update_fixture() -> ReplaySendUpdateFixture
previous_model: std::cell::Cell::new(None),
compaction_mode: xai_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: crate::util::config::CompactionToolChoice::Auto,
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},

51
crates/codegen/xai-grok-shell/src/session/acp_session_tests/rewrite_zero_turn_prefix_tests.rs
View file

@ -98,6 +98,57 @@ async fn rebuild_reinjects_goal_update_handle() {
})
.await;
}
#[tokio::test(flavor = "current_thread")]
async fn rebuild_reinjects_task_completion_resource_identity() {
use xai_grok_tools::reminders::task_completion::{
TaskCompletionReservations, TaskWakeSuppressed,
};
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (gw_tx, _gw_rx) = tokio::sync::mpsc::unbounded_channel();
let (persist_tx, _persist_rx) = tokio::sync::mpsc::unbounded_channel();
let actor = create_test_actor(0, 256_000, 85, gw_tx, persist_tx).await;
let session_reservations = actor
.tool_context
.task_completion_reservations
.clone()
.expect("session completion reservations");
let session_gate = actor
.tool_context
.task_wake_suppressed
.clone()
.expect("session task-wake gate");
session_reservations.reserve("before-rebuild".to_string());
session_gate.set(true);
actor
.handle_rebuild_agent_for_definition(
xai_grok_agent::AgentDefinition::default_grok_build(),
)
.await
.expect("zero-turn rebuild should succeed");
let bridge = actor.agent.borrow().tool_bridge().clone();
let resources = bridge.shared_resources().await;
let guard = resources.lock().await;
let rebuilt_reservations = guard
.get::<TaskCompletionReservations>()
.expect("rebuilt bridge completion reservations");
let rebuilt_gate = guard
.get::<TaskWakeSuppressed>()
.expect("rebuilt bridge task-wake gate");
assert!(rebuilt_reservations.contains("before-rebuild"));
assert!(rebuilt_gate.get());
session_reservations.release("before-rebuild");
session_gate.set(false);
assert!(!rebuilt_reservations.contains("before-rebuild"));
assert!(!rebuilt_gate.get());
rebuilt_reservations.reserve("from-rebuilt-bridge".to_string());
rebuilt_gate.set(true);
assert!(session_reservations.contains("from-rebuilt-bridge"));
assert!(session_gate.get());
})
.await;
}
/// The seeded skill used by the rebuild skill-reminder tests. A non-plugin
/// Local skill is always listable, so it renders into the grok markdown skill
/// catalog when the pending baseline is drained for a different agent.

26
crates/codegen/xai-grok-shell/src/session/acp_session_tests/support.rs
View file

@ -159,7 +159,12 @@ pub(crate) async fn create_test_actor_ex(
xai_hunk_tracker::TrackingMode::AgentOnly,
tokio_util::sync::CancellationToken::new(),
);
let tool_context = ToolContext::new(cwd.clone(), None, None, fs, terminal, hunk_tracker_handle);
let mut tool_context =
ToolContext::new(cwd.clone(), None, None, fs, terminal, hunk_tracker_handle);
tool_context.task_completion_reservations =
Some(xai_grok_tools::reminders::task_completion::TaskCompletionReservations::default());
tool_context.task_wake_suppressed =
Some(xai_grok_tools::reminders::task_completion::TaskWakeSuppressed::default());
let state = TokioMutex::new(State {
running_task: None,
pending_inputs: VecDeque::new(),
@ -235,6 +240,7 @@ pub(crate) async fn create_test_actor_ex(
previous_model: std::cell::Cell::new(None),
compaction_mode: xai_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: crate::util::config::CompactionToolChoice::Auto,
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@ -358,6 +364,22 @@ pub(crate) async fn create_test_actor_ex(
workspace_ops: xai_grok_workspace::WorkspaceOps::for_test(),
trace_config_template: std::cell::RefCell::new(None),
};
if let Some(reservations) = actor.tool_context.task_completion_reservations.clone() {
actor
.agent
.borrow()
.tool_bridge()
.update_resource(reservations)
.await;
}
if let Some(gate) = actor.tool_context.task_wake_suppressed.clone() {
actor
.agent
.borrow()
.tool_bridge()
.update_resource(gate)
.await;
}
(actor, event_rx)
}
#[cfg(test)]
@ -399,6 +421,7 @@ pub(crate) fn user_item_with_rx(
verbatim: false,
json_schema: None,
origin: crate::session::PromptOrigin::User,
task_wake_fallback: None,
respond_to,
persist_ack: None,
parsed_prompt_tx: None,
@ -438,6 +461,7 @@ pub(crate) fn input_with_origin_rx(
verbatim,
json_schema: None,
origin,
task_wake_fallback: None,
respond_to,
persist_ack: None,
parsed_prompt_tx: None,

1
crates/codegen/xai-grok-shell/src/session/acp_session_tests/turn_completion_emit_tests.rs
View file

@ -70,6 +70,7 @@ fn pending_input(prompt_id: &str) -> (InputItem, oneshot::Receiver<PromptTurnRes
verbatim: false,
json_schema: None,
origin: crate::session::PromptOrigin::User,
task_wake_fallback: None,
respond_to,
persist_ack: None,
parsed_prompt_tx: None,

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

@ -40,7 +40,7 @@ pub enum PromptCompletionKind {
/// every attached leader-mode client the *running* turn ended) and the
/// roster `Idle` delta (which would flip the dashboard off `Working` while
/// the real turn is still in flight) must be skipped. See
/// `MvpAgent::prompt`'s short-circuit and `respond_removed_queued_prompt`.
/// `MvpAgent::prompt`'s short-circuit and `respond_removed_prompt`.
RemovedFromQueue,
}
/// Successful prompt/turn payload returned to the ACP layer and trace uploaders.
@ -94,15 +94,29 @@ pub enum NotificationPriority {
#[derive(Debug, Clone)]
pub enum NotificationSource {
MonitorEvent { task_id: String },
MonitorCompleted { task_id: String },
BashTaskCompleted { task_id: String },
}
impl NotificationSource {
pub fn task_id(&self) -> &str {
match self {
Self::MonitorEvent { task_id } | Self::BashTaskCompleted { task_id } => task_id,
Self::MonitorEvent { task_id }
| Self::MonitorCompleted { task_id }
| Self::BashTaskCompleted { task_id } => task_id,
}
}
}
#[derive(Debug)]
pub struct TaskWakeFallback {
pub prompt_id: String,
pub prompt_blocks: Vec<acp::ContentBlock>,
pub source: NotificationSource,
}
#[derive(Debug)]
pub struct TaskWakeAdmission {
pub respond_to: oneshot::Sender<bool>,
pub fallback: TaskWakeFallback,
}
pub enum SessionCommand {
Initialize {
system_prompt: String,
@ -142,6 +156,8 @@ pub enum SessionCommand {
/// Also derived server-side during an interruptible wait (see
/// [`SessionActor::queue_input`]).
send_now: bool,
/// Actor-authoritative admission and deferred fallback for terminal task wakes.
admission: Option<TaskWakeAdmission>,
respond_to: oneshot::Sender<PromptTurnResult>,
/// Optional oneshot fired after the user message has been appended to
/// chat history and a persistence flush barrier has completed, before
@ -567,10 +583,10 @@ pub enum SessionCommand {
},
/// Cancel the running turn. `kill_background_tasks` distinguishes a hard
/// teardown (subagent shutdown — drains the whole queue) from a normal
/// interactive cancel (Ctrl+C — preserves the queued prompts so the next
/// one auto-runs). On an interactive cancel only the running turn (the front
/// of `pending_inputs`) is torn down; the follow-up `maybe_start_running_task`
/// promotes the new front so the user's next queued prompt auto-runs.
/// interactive cancel (Ctrl+C — preserves queued user prompts so the next
/// one auto-runs). Ctrl+C tears down the running turn and queued terminal
/// task-completion wakes; other cancel triggers tear down only the running
/// turn. The follow-up `maybe_start_running_task` promotes the next item.
Cancel {
cancel_subagents: bool,
kill_background_tasks: bool,

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

@ -206,6 +206,7 @@ impl SessionActor {
&sampling_config,
self.inference_idle_timeout,
wall_clock_budget_secs,
self.compaction.tool_choice,
)
.await
{
@ -999,6 +1000,7 @@ impl SessionActor {
sampling_config.clone(),
self.inference_idle_timeout,
wall_clock_budget_secs,
self.compaction.tool_choice,
);
let observer =
crate::session::helpers::full_replace_compaction::ShellFullReplaceObserver::new(
@ -2238,6 +2240,7 @@ mod inline_auto_compact_flow_tests {
previous_model: std::cell::Cell::new(None),
compaction_mode: xai_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: crate::util::config::CompactionToolChoice::Auto,
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},

1
crates/codegen/xai-grok-shell/src/session/compaction_config.rs
View file

@ -147,6 +147,7 @@ pub struct CompactionConfig {
pub compaction_mode: xai_chat_state::CompactionMode,
/// When `true`, feed the summarizer the verbatim conversation instead of the lossy rewrite (the retry loop may still fall back).
pub verbatim_input: bool,
pub tool_choice: crate::util::config::CompactionToolChoice,
/// Prefire two-pass state (background NOTE₁ cache + in-flight guard).
/// `Default` (empty cache, not in-flight).
pub prefire: PrefireState,

4
crates/codegen/xai-grok-shell/src/session/helpers/full_replace_compaction.rs
View file

@ -71,6 +71,7 @@ pub(crate) struct ShellCompactionSampler {
/// Wall-clock budget (secs) forwarded to `generate_session_compact` as the
/// reasoning-runaway backstop; `0` disables it.
wall_clock_budget_secs: u64,
tool_choice: crate::util::config::CompactionToolChoice,
/// Full output of the most recent successful sample (for L5 telemetry).
last_success: Mutex<Option<CompactOutput>>,
}
@ -87,6 +88,7 @@ impl ShellCompactionSampler {
sampling_config: SamplingConfig,
idle_timeout: Duration,
wall_clock_budget_secs: u64,
tool_choice: crate::util::config::CompactionToolChoice,
) -> Self {
Self {
use_short_prompt,
@ -98,6 +100,7 @@ impl ShellCompactionSampler {
sampling_config,
idle_timeout,
wall_clock_budget_secs,
tool_choice,
last_success: Mutex::new(None),
}
}
@ -136,6 +139,7 @@ impl CompactionSampler for ShellCompactionSampler {
&self.sampling_config,
self.idle_timeout,
self.wall_clock_budget_secs,
self.tool_choice,
)
.await
{

178
crates/codegen/xai-grok-shell/src/session/helpers/session_compact.rs
View file

@ -339,9 +339,7 @@ where
/// prefix and force a full prefill on the summarizer call — attaching them
/// keeps the request prefix byte-identical to the turn requests so the
/// engine reuses the session's KV cache (the whole point of the verbatim
/// input path). Tool *use* is forbidden via `tool_choice: none` where the
/// backend can express it (ChatCompletions, Responses); the Messages wire
/// enum has no `none`, so that path relies on the prompt instruction alone.
/// input path).
///
/// Errors carry a [`CompactFailure`] classification so the caller can
/// short-circuit retries on deterministic failures (4xx schema violations,
@ -356,8 +354,17 @@ pub(crate) async fn generate_session_compact(
sampling_config: &SamplingConfig,
idle_timeout: std::time::Duration,
wall_clock_budget_secs: u64,
tool_choice: crate::util::config::CompactionToolChoice,
) -> Result<CompactOutput, CompactFailure> {
let num_messages = chat_history.len();
let wire_tool_choice = match tool_choice {
crate::util::config::CompactionToolChoice::Auto => ToolChoice::auto(),
crate::util::config::CompactionToolChoice::None => ToolChoice::none(),
};
let conversation_tool_choice = match tool_choice {
crate::util::config::CompactionToolChoice::Auto => ConversationToolChoice::Auto,
crate::util::config::CompactionToolChoice::None => ConversationToolChoice::None,
};
let output = match sampling_config.api_backend {
ApiBackend::ChatCompletions => {
let chat_messages: Vec<ChatRequestMessage> =
@ -373,7 +380,7 @@ pub(crate) async fn generate_session_compact(
.map(|t| ToolDefinition::function(t.name, t.description, t.parameters))
.collect(),
)
.with_tool_choice(ToolChoice::none());
.with_tool_choice(wire_tool_choice);
}
let sid = session_id.to_string();
message.x_grok_conv_id = Some(sid.clone());
@ -467,7 +474,7 @@ pub(crate) async fn generate_session_compact(
ApiBackend::Responses => {
let request = ConversationRequest {
items: chat_history,
tool_choice: (!tools.is_empty()).then_some(ConversationToolChoice::None),
tool_choice: (!tools.is_empty()).then_some(conversation_tool_choice),
tools,
hosted_tools,
model: Some(sampling_config.model.to_owned()),
@ -1578,6 +1585,7 @@ mod reasoning_compaction_regression_tests {
&config,
std::time::Duration::from_secs(30),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await
.unwrap_or_else(|_| panic!("compaction must succeed"));
@ -1666,6 +1674,7 @@ mod reasoning_compaction_regression_tests {
&config,
std::time::Duration::from_secs(30),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await;
let output = result
@ -1673,12 +1682,8 @@ mod reasoning_compaction_regression_tests {
assert_eq!(output.content, "<summary>ok</summary>");
let _ = shutdown_tx.send(());
}
/// The compaction request must carry the turn loop's tool definitions
/// (prompt-prefix/KV-cache alignment) with `tool_choice: "none"`, and
/// must omit both keys when no tools are passed (Chat Completions rejects a bare
/// `tool_choice`).
#[tokio::test]
async fn chat_completions_compaction_attaches_tools_with_tool_choice_none() {
async fn chat_completions_compaction_attaches_tools_with_tool_choice_auto() {
use std::sync::{Arc, Mutex};
let captured: Arc<Mutex<Vec<serde_json::Value>>> = Arc::new(Mutex::new(Vec::new()));
let cap = captured.clone();
@ -1731,6 +1736,7 @@ mod reasoning_compaction_regression_tests {
&config,
std::time::Duration::from_secs(30),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await
.unwrap_or_else(|_| panic!("compaction with tools must succeed"));
@ -1744,6 +1750,7 @@ mod reasoning_compaction_regression_tests {
&config,
std::time::Duration::from_secs(30),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await
.unwrap_or_else(|_| panic!("compaction without tools must succeed"));
@ -1752,8 +1759,8 @@ mod reasoning_compaction_regression_tests {
let with_tools = &bodies[0];
assert_eq!(
with_tools["tool_choice"],
json!("none"),
"tool use must be disabled at decode time"
json!("auto"),
"default compaction tool_choice is auto"
);
let sent_tools = with_tools["tools"]
.as_array()
@ -1771,6 +1778,149 @@ mod reasoning_compaction_regression_tests {
);
let _ = shutdown_tx.send(());
}
fn responses_summary_stream() -> Vec<Event> {
vec![
Event::default().data(
json!({ "type" : "response.created", "sequence_number"
: 0, "response" : { "id" : "resp_test", "object" : "response", "created_at" :
1234567890, "model" : "test-model", "status" : "in_progress", "output" : [] }
})
.to_string(),
),
Event::default().data(
json!({ "type" :
"response.output_text.delta", "sequence_number" : 1, "item_id" : "msg_test",
"output_index" : 0, "content_index" : 0, "delta" : "<summary>ok</summary>" })
.to_string(),
),
Event::default().data(
json!({ "type" : "response.completed",
"sequence_number" : 2, "response" : { "id" : "resp_test", "object" :
"response", "created_at" : 1234567890, "model" : "test-model", "status" :
"completed", "output" : [] } })
.to_string(),
),
]
}
fn test_config_responses(base_url: &str) -> SamplerConfig {
let mut config = test_config(base_url);
config.api_backend = ApiBackend::Responses;
config
}
#[tokio::test]
async fn responses_compaction_attaches_tools_with_tool_choice_auto() {
use std::sync::{Arc, Mutex};
let captured: Arc<Mutex<Vec<serde_json::Value>>> = Arc::new(Mutex::new(Vec::new()));
let cap = captured.clone();
let app = Router::new().route(
"/v1/responses",
post(move |body: axum::Json<serde_json::Value>| {
let cap = cap.clone();
async move {
cap.lock().unwrap().push(body.0);
let stream = stream::iter(
responses_summary_stream()
.into_iter()
.map(Ok::<_, std::convert::Infallible>),
);
Sse::new(stream).keep_alive(KeepAlive::default())
}
}),
);
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
tokio::spawn(async move {
axum::serve(listener, app)
.with_graceful_shutdown(async {
let _ = shutdown_rx.await;
})
.await
.unwrap();
});
let base_url = format!("http://{addr}/v1");
let config = test_config_responses(&base_url);
let chat_history = vec![
ConversationItem::system("You are a helpful assistant."),
ConversationItem::user("<user_query>\nfix the bug\n</user_query>"),
ConversationItem::assistant("I fixed it."),
ConversationItem::user("Summarize the conversation so far."),
];
let tools = vec![ToolSpec {
name: "read_file".to_string(),
description: Some("Reads a file".to_string()),
parameters: json!({ "type" : "object", "properties" : {} }),
}];
let hosted = vec![HostedTool::WebSearch {
allowed_domains: None,
}];
let client = Client::new(config.clone()).unwrap();
generate_session_compact(
chat_history.clone(),
tools,
hosted,
client,
acp::SessionId::new("test-session"),
&config,
std::time::Duration::from_secs(30),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await
.unwrap_or_else(|_| panic!("Responses compaction with tools must succeed"));
let client = Client::new(config.clone()).unwrap();
generate_session_compact(
chat_history,
vec![],
vec![],
client,
acp::SessionId::new("test-session"),
&config,
std::time::Duration::from_secs(30),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await
.unwrap_or_else(|_| panic!("Responses compaction without tools must succeed"));
let bodies = captured.lock().unwrap();
assert_eq!(bodies.len(), 2, "mock must have served both requests");
let with_tools = &bodies[0];
assert_eq!(
with_tools["tool_choice"],
json!("auto"),
"default Responses compaction tool_choice is auto"
);
let sent_tools = with_tools["tools"]
.as_array()
.expect("tools must be attached for prefix-cache alignment");
let has_read_file = sent_tools.iter().any(|t| {
t.get("name") == Some(&json!("read_file"))
|| t.pointer("/name") == Some(&json!("read_file"))
});
assert!(
has_read_file,
"client function tool must be present: {sent_tools:?}"
);
assert!(
sent_tools
.iter()
.any(|t| t.get("type") == Some(&json!("web_search"))),
"hosted web_search must be present for prefix alignment: {sent_tools:?}"
);
let without_tools = &bodies[1];
assert!(
without_tools
.get("tools")
.map(|t| t.as_array().is_none_or(|a| a.is_empty()))
.unwrap_or(true),
"no tools when none are passed"
);
assert!(
without_tools.get("tool_choice").is_none(),
"tool_choice without tools should be omitted"
);
let _ = shutdown_tx.send(());
}
#[tokio::test]
async fn stalled_compaction_stream_times_out_as_transient() {
let app = Router::new().route(
@ -1807,6 +1957,7 @@ mod reasoning_compaction_regression_tests {
&config,
std::time::Duration::from_millis(150),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await;
match result {
@ -1874,6 +2025,7 @@ mod reasoning_compaction_regression_tests {
&config,
std::time::Duration::from_millis(150),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await;
match result {
@ -1945,6 +2097,7 @@ mod reasoning_compaction_regression_tests {
&config,
std::time::Duration::from_millis(150),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await;
match result {
@ -2013,6 +2166,7 @@ mod reasoning_compaction_regression_tests {
&config,
std::time::Duration::from_millis(150),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await;
match result {

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

@ -137,7 +137,7 @@ impl PromptOrigin {
}
}
/// If this is an auto-wake prompt, returns the inner completion ID
/// (task or subagent ID). Used to pre-mark reported completions.
/// (task or subagent ID). Used by queue preemption and cancellation cleanup.
pub fn completion_id(&self) -> Option<&str> {
match self {
Self::TaskCompleted { task_id } => Some(task_id),

3
crates/codegen/xai-grok-shell/src/test_support/lsp_runtime.rs
View file

@ -97,7 +97,6 @@ pub(crate) fn ctx_with_toggle(toggle: HashMap<String, bool>) -> SubagentSpawnCon
parent_session_info: None,
subagent_roles: HashMap::new(),
subagent_personas: HashMap::new(),
persona_io_summaries: Vec::new(),
parent_chat_state: None,
available_models: indexmap::IndexMap::new(),
subagent_model_overrides: HashMap::new(),
@ -141,7 +140,7 @@ pub(crate) fn ctx_with_toggle(toggle: HashMap<String, bool>) -> SubagentSpawnCon
parent_skills: None,
parent_skills_config: xai_grok_agent::prompt::skills::SkillsConfig::default(),
parent_compat: xai_grok_tools::types::compat::CompatConfig::default(),
auto_wake_delivered: None,
task_completion_reservations: None,
synthetic_trace_tx: None,
task_output_tool_name: xai_grok_tools::reminders::task_completion::DEFAULT_TASK_OUTPUT_TOOL
.to_string(),

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

@ -35,6 +35,7 @@ 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>,
pub login_shell_capture: Option<bool>,
}
impl BashToolConfig {
@ -110,6 +111,10 @@ pub struct WebFetchToolConfig {
/// default allowlist. An explicit empty list blocks all fetches.
/// Resolution: TOML > remote settings > built-in defaults.
pub allowed_domains: Option<Vec<String>>,
/// Allow fetches to explicit loopback hosts only (`localhost` / `127.0.0.0/8`
/// / `::1`). Private and metadata ranges stay blocked. Default off.
/// Resolution: TOML > `GROK_WEB_FETCH_ALLOW_LOCAL` env > false.
pub allow_local: Option<bool>,
}
impl WebFetchToolConfig {
@ -139,10 +144,15 @@ impl WebFetchToolConfig {
.cloned()
.or_else(|| remote_domains.map(|d| d.to_vec()));
let allow_local = self
.allow_local
.or_else(|| xai_grok_config::env_bool("GROK_WEB_FETCH_ALLOW_LOCAL"));
xai_grok_tools::implementations::grok_build::web_fetch::WebFetchParams {
proxy_endpoint,
allowed_domains,
context_window_tokens,
allow_local,
..Default::default()
}
}
@ -576,6 +586,7 @@ mod tests {
let local = WebFetchToolConfig {
proxy_endpoint: Some("https://toml-proxy.example.com".to_owned()),
allowed_domains: Some(vec!["toml.example.com".to_owned()]),
allow_local: Some(true),
};
let params = local.resolve_params(
Some("https://remote-proxy.example.com"),
@ -590,6 +601,8 @@ mod tests {
params.allowed_domains,
Some(vec!["toml.example.com".to_owned()])
);
assert_eq!(params.allow_local, Some(true));
assert!(params.allow_local());
}
#[test]
@ -608,6 +621,7 @@ mod tests {
params.allowed_domains,
Some(vec!["remote.example.com".to_owned()])
);
assert!(!params.allow_local());
}
#[test]
@ -616,6 +630,7 @@ mod tests {
let params = local.resolve_params(None, None, None);
assert!(params.proxy_endpoint.is_none());
assert!(params.allowed_domains.is_none());
assert!(!params.allow_local());
}
#[test]
@ -623,6 +638,7 @@ mod tests {
let local = WebFetchToolConfig {
proxy_endpoint: None,
allowed_domains: Some(vec![]),
allow_local: None,
};
let params = local.resolve_params(None, Some(&["remote.example.com".to_owned()]), None);
assert_eq!(params.allowed_domains, Some(vec![]));

569
crates/codegen/xai-grok-shell/src/tools/notification_bridge.rs
View file

@ -17,6 +17,8 @@ use crate::session::commands::{NotificationPriority, NotificationSource};
use crate::session::persistence::PersistenceMsg;
use xai_grok_workspace::session::file_state::FileStateTracker;
const TASK_WAKE_ADMISSION_TIMEOUT: std::time::Duration = std::time::Duration::from_millis(250);
/// Configuration for the notification bridge.
pub struct NotificationBridgeConfig {
/// ACP gateway for sending streaming updates to TUI
@ -55,9 +57,9 @@ pub struct NotificationBridgeConfig {
pub turn_prompt_mode: Arc<parking_lot::Mutex<crate::session::plan_mode::PromptMode>>,
/// Session command channel for monitor events and task-completed injections.
pub session_cmd_tx: mpsc::UnboundedSender<SessionCommand>,
/// Shared set of IDs delivered via auto-wake, used to suppress duplicate
/// `TaskCompletionReminder` entries for the same task/subagent.
pub auto_wake_delivered: xai_grok_tools::reminders::task_completion::AutoWakeDeliveredIds,
pub task_completion_reservations:
xai_grok_tools::reminders::task_completion::TaskCompletionReservations,
pub task_wake_suppressed: xai_grok_tools::reminders::task_completion::TaskWakeSuppressed,
/// Channel for requesting trace uploads for synthetic auto-wake turns.
/// Wrapped in `Arc<Mutex<..>>` because the coordinator creates the channel
/// after the notification bridge is spawned — the bridge reads the latest
@ -331,20 +333,11 @@ async fn handle_notification(
.goal_loop_active
.load(std::sync::atomic::Ordering::Relaxed);
// Block-waited / explicitly-killed: the model already has the result
// (blocking wait return or kill_task tool response). Skip auto-wake
// for both bash and monitors — a redundant synthetic prompt is noise.
//
// Natural monitor exit (including exit code 0) MUST auto-wake the
// same way bash does. Relying only on the pipeline's terminal
// `MonitorEvent` + idle-gated `InjectNotification` was easy to miss
// when the agent was idle and the monitor produced no further
// stdout. The pager still receives x.ai/task_completed below for UI.
// Stamped on the completion notification below so the pager knows
// whether a wake response follows the chip.
// Natural monitor exit uses the same immediate wake path as bash;
// x.ai/task_completed still drives the pager UI in every branch.
let mut will_wake = false;
if task_snapshot.block_waited || task_snapshot.explicitly_killed {
// no auto-wake
// The blocking wait or kill result already reports completion.
} else if goal_loop_active {
// Goal loop active: suppress the wake (synthetic prompt + the
// idle-gated fallback); surfaces 2/3 drain it. See
@ -355,21 +348,7 @@ async fn handle_notification(
"auto-wake: suppressed completion (goal loop active)"
);
} else if config.auto_wake_enabled {
// Mark delivered so `TaskCompletionReminder` suppresses the
// duplicate on the next tool call (bash and monitor alike).
config.auto_wake_delivered.insert(task_id.clone());
// Monitor exit: the TaskCompleted Prompt is the sole model-facing
// wake. Drop any already-queued pipeline MonitorEvents for this
// task (stdout lines + terminal ended) so they do not start a
// second NotificationDrain turn after the wake.
if is_monitor {
let _ = config
.session_cmd_tx
.send(SessionCommand::DropMonitorNotifications {
task_id: task_id.clone(),
});
}
config.task_completion_reservations.reserve(task_id.clone());
let tool_name = resolved_tool_name(&config.task_output_tool_name);
let read_name = resolved_tool_name(&config.read_tool_name);
@ -389,25 +368,20 @@ async fn handle_notification(
let prompt_id = format!("task-completed-{task_id}");
let prompt_blocks = vec![acp::ContentBlock::Text(acp::TextContent::new(message))];
// Capture a pre-prompt session snapshot for the trace upload path.
let (before_copy_tx, before_copy_rx) = tokio::sync::oneshot::channel();
let _ = config.session_cmd_tx.send(SessionCommand::CopyFile {
respond_to: before_copy_tx,
});
let synthetic_trace_tx = config
.synthetic_trace_tx
.lock()
.unwrap_or_else(|e| e.into_inner())
.clone();
let (respond_to, completion_rx) = tokio::sync::oneshot::channel();
let (admission_tx, admission_rx) = tokio::sync::oneshot::channel();
tracing::info!(
task_id = %task_id,
prompt_id = %prompt_id,
is_monitor,
"auto-wake: injecting synthetic prompt for completed background task"
"auto-wake: requesting synthetic prompt admission for completed background task"
);
// Stamp from the actual enqueue: `will_wake` on the completion
// notification must never promise a wake this send didn't queue
// (mirrors `parent_channel_open` in `should_auto_wake_subagent`).
// The channel is unbounded, so this only fails when the session
// actor is already gone.
will_wake = config
let enqueued = config
.session_cmd_tx
.send(SessionCommand::Prompt {
prompt_id: prompt_id.clone(),
@ -420,32 +394,99 @@ async fn handle_notification(
traceparent: xai_file_utils::trace_context::current_traceparent(),
json_schema: None,
send_now: false,
admission: Some(crate::session::commands::TaskWakeAdmission {
respond_to: admission_tx,
fallback: crate::session::commands::TaskWakeFallback {
prompt_id: if is_monitor {
format!("monitor-completed-{task_id}")
} else {
format!("bash-completed-{task_id}")
},
prompt_blocks: vec![acp::ContentBlock::Text(
acp::TextContent::new(body.clone()),
)],
source: if is_monitor {
NotificationSource::MonitorCompleted {
task_id: task_id.clone(),
}
} else {
NotificationSource::BashTaskCompleted {
task_id: task_id.clone(),
}
},
},
}),
respond_to,
persist_ack: None,
parsed_prompt_tx: None,
})
.is_ok();
if let Some(ref trace_tx) = *config
.synthetic_trace_tx
.lock()
.unwrap_or_else(|e| e.into_inner())
{
tracing::info!(
task_id = %task_id,
"auto-wake: sending synthetic turn trace request"
);
let _ = trace_tx.send(crate::upload::turn::SyntheticTurnTraceRequest {
session_id: config.session_id.clone(),
prompt_id,
completion_rx,
before_session_copy_rx: before_copy_rx,
});
if !enqueued {
config.task_completion_reservations.release(&task_id);
}
let admitted = if enqueued {
tokio::time::timeout(TASK_WAKE_ADMISSION_TIMEOUT, admission_rx)
.await
.ok()
.and_then(Result::ok)
.unwrap_or(false)
} else {
tracing::debug!(
task_id = %task_id,
"auto-wake: no synthetic_trace_tx, skipping trace request"
);
false
};
will_wake = admitted;
xai_grok_telemetry::unified_log::info(
"shell.task_wake.bridge_admission",
Some(config.session_id.0.as_ref()),
Some(serde_json::json!({
"task_id": &task_id,
"monitor": is_monitor,
"enqueued": enqueued,
"admitted": admitted,
"gate": config.task_wake_suppressed.get(),
})),
);
if will_wake {
if is_monitor {
let _ =
config
.session_cmd_tx
.send(SessionCommand::DropMonitorNotifications {
task_id: task_id.clone(),
});
}
if let Some(trace_tx) = synthetic_trace_tx {
let (before_copy_tx, before_session_copy_rx) =
tokio::sync::oneshot::channel();
let copy_requested = config
.session_cmd_tx
.send(SessionCommand::CopyFile {
respond_to: before_copy_tx,
})
.is_ok();
if copy_requested {
tracing::info!(
task_id = %task_id,
"auto-wake: sending synthetic turn trace request"
);
let _ = trace_tx.send(crate::upload::turn::SyntheticTurnTraceRequest {
session_id: config.session_id.clone(),
prompt_id,
completion_rx,
before_session_copy_rx,
});
} else {
tracing::debug!(
task_id = %task_id,
"auto-wake: session snapshot request failed, skipping trace request"
);
}
} else {
tracing::debug!(
task_id = %task_id,
"auto-wake: no synthetic trace consumer, skipping trace request"
);
}
}
} else {
// Auto-wake disabled — fall back to idle-gated notification drain.
@ -463,17 +504,28 @@ async fn handle_notification(
read_name,
)
};
let prompt_id = format!("bash-completed-{task_id}");
let prompt_blocks = vec![acp::ContentBlock::Text(acp::TextContent::new(message))];
let source = if is_monitor {
NotificationSource::MonitorCompleted {
task_id: task_id.clone(),
}
} else {
NotificationSource::BashTaskCompleted {
task_id: task_id.clone(),
}
};
let _ = config
.session_cmd_tx
.send(SessionCommand::InjectNotification {
prompt_id,
prompt_blocks,
priority: NotificationPriority::Later,
source: NotificationSource::BashTaskCompleted {
task_id: task_id.clone(),
prompt_id: if is_monitor {
format!("monitor-completed-{task_id}")
} else {
format!("bash-completed-{task_id}")
},
prompt_blocks: vec![acp::ContentBlock::Text(acp::TextContent::new(
message,
))],
priority: NotificationPriority::Later,
source,
});
}
@ -702,7 +754,7 @@ async fn handle_notification(
// If this monitor already auto-woke via TaskCompleted, do not inject
// model-facing notifications (avoids a second NotificationDrain turn
// with the same ended signal). Pager UI still got the event above.
if config.auto_wake_delivered.contains(&event.task_id) {
if config.task_completion_reservations.contains(&event.task_id) {
tracing::debug!(
task_id = %event.task_id,
"skipping model inject for monitor event: task already auto-woke via TaskCompleted"
@ -808,6 +860,38 @@ mod tests {
use xai_grok_tools::computer::types::TaskKind;
use xai_grok_tools::types::TaskSnapshot;
/// Drive the admission handshake inline so receiver assertions observe the
/// bridge's command order without racing a detached proxy task.
async fn handle_notification_with_admission(
config: &NotificationBridgeConfig,
notification: ToolNotification,
offsets: &mut HashMap<String, usize>,
cmd_rx: &mut mpsc::UnboundedReceiver<SessionCommand>,
accepted: bool,
) {
let notification = handle_notification(config, notification, offsets);
tokio::pin!(notification);
let mut command = tokio::select! {
_ = &mut notification => panic!("notification completed before requesting admission"),
command = cmd_rx.recv() => command.expect("expected task-wake prompt"),
};
let SessionCommand::Prompt { admission, .. } = &mut command else {
panic!("expected task-wake prompt");
};
admission
.take()
.expect("expected task-wake admission request")
.respond_to
.send(accepted)
.expect("notification must still be awaiting admission");
config
.session_cmd_tx
.send(command)
.expect("test command receiver must remain open");
notification.await;
}
fn make_test_config() -> (
NotificationBridgeConfig,
mpsc::UnboundedReceiver<SessionCommand>,
@ -816,15 +900,22 @@ mod tests {
(config, session_cmd_rx)
}
/// Same as [`make_test_config`] but also returns the gateway and
/// persistence receivers so a test can inspect the notifications
/// emitted by `handle_notification`. Use this for plan-mode tests.
#[allow(clippy::type_complexity)]
fn make_test_config_full() -> (
NotificationBridgeConfig,
mpsc::UnboundedReceiver<xai_acp_lib::AcpClientMessage>,
mpsc::UnboundedReceiver<PersistenceMsg>,
mpsc::UnboundedReceiver<SessionCommand>,
) {
make_test_config_full_raw()
}
#[allow(clippy::type_complexity)]
fn make_test_config_full_raw() -> (
NotificationBridgeConfig,
mpsc::UnboundedReceiver<xai_acp_lib::AcpClientMessage>,
mpsc::UnboundedReceiver<PersistenceMsg>,
mpsc::UnboundedReceiver<SessionCommand>,
) {
let (gateway_tx, gateway_rx) = mpsc::unbounded_channel();
let gateway = xai_acp_lib::AcpAgentGatewaySender::new(gateway_tx);
@ -850,8 +941,10 @@ mod tests {
crate::session::plan_mode::PromptMode::Agent,
)),
session_cmd_tx,
auto_wake_delivered:
xai_grok_tools::reminders::task_completion::AutoWakeDeliveredIds::default(),
task_completion_reservations:
xai_grok_tools::reminders::task_completion::TaskCompletionReservations::default(),
task_wake_suppressed:
xai_grok_tools::reminders::task_completion::TaskWakeSuppressed::default(),
synthetic_trace_tx: Arc::new(std::sync::Mutex::new(None)),
task_output_tool_name: Arc::new(std::sync::OnceLock::new()),
read_tool_name: Arc::new(std::sync::OnceLock::new()),
@ -896,14 +989,11 @@ mod tests {
let notification = ToolNotification::TaskCompleted(snapshot);
let mut offsets = HashMap::new();
handle_notification(&config, notification, &mut offsets).await;
handle_notification_with_admission(&config, notification, &mut offsets, &mut cmd_rx, true)
.await;
// Auto-wake sends CopyFile first, then Prompt (not InjectNotification).
let cmd1 = cmd_rx.try_recv().expect("expected CopyFile");
assert!(matches!(cmd1, SessionCommand::CopyFile { .. }));
let cmd2 = cmd_rx.try_recv().expect("expected Prompt");
match cmd2 {
let command = cmd_rx.try_recv().expect("expected Prompt");
match command {
SessionCommand::Prompt {
prompt_id,
prompt_blocks,
@ -946,7 +1036,7 @@ mod tests {
/// Gap 1: while a goal loop is active, a completed background bash task
/// must NOT fire the synthetic auto-wake prompt — an async "task completed"
/// wake mid-goal derails a weak model. It must also NOT be marked
/// auto-wake-delivered (so surface 2's `TaskCompletionReminder` is free to
/// reserved (so surface 2's `TaskCompletionReminder` is free to
/// drain it). The pager's `x.ai/task_completed` notification still fires.
#[tokio::test]
async fn bash_task_completed_suppresses_auto_wake_during_goal_loop() {
@ -983,10 +1073,10 @@ mod tests {
cmd_rx.try_recv().is_err(),
"goal-loop-active bash completion must not inject auto-wake commands"
);
// Not marked auto-wake-delivered: surface 2 must be free to drain it.
// Not marked reserved: surface 2 must be free to drain it.
assert!(
config.auto_wake_delivered.snapshot().is_empty(),
"goal-loop-active completion must not be marked auto-wake-delivered"
config.task_completion_reservations.snapshot().is_empty(),
"goal-loop-active completion must not be marked reserved"
);
// The pager UI notification must still be emitted.
let mut found_ext = false;
@ -1005,31 +1095,27 @@ mod tests {
/// Gap 1 (preserve non-goal behavior): with the goal loop inactive — the
/// default for a normal session — a completed bash task DOES fire the
/// synthetic auto-wake prompt AND is marked auto-wake-delivered so surface
/// synthetic auto-wake prompt AND is marked reserved so surface
/// 2 suppresses the duplicate reminder.
#[tokio::test]
async fn bash_task_completed_auto_wakes_and_marks_delivered_without_goal_loop() {
async fn bash_task_completed_auto_wakes_and_reserves_without_goal_loop() {
let (config, mut cmd_rx) = make_test_config();
config
.task_output_tool_name
.set(Some("get_command_or_subagent_output".to_string()))
.expect("slot is fresh in this test fixture");
// goal_loop_active defaults to false (normal session).
let snapshot = make_task_snapshot("bg-normal", TaskKind::Bash);
let mut offsets = HashMap::new();
handle_notification(
handle_notification_with_admission(
&config,
ToolNotification::TaskCompleted(snapshot),
&mut offsets,
&mut cmd_rx,
true,
)
.await;
// Auto-wake sends CopyFile, then the synthetic Prompt.
assert!(matches!(
cmd_rx.try_recv(),
Ok(SessionCommand::CopyFile { .. })
));
assert!(matches!(
cmd_rx.try_recv(),
Ok(SessionCommand::Prompt { .. })
@ -1038,14 +1124,12 @@ mod tests {
cmd_rx.try_recv(),
Ok(SessionCommand::DispatchNotificationHook { .. })
));
// And the task IS marked auto-wake-delivered.
assert_eq!(
config.auto_wake_delivered.snapshot(),
config.task_completion_reservations.snapshot(),
vec!["bg-normal".to_string()],
);
}
/// `will_wake` off the emitted `x.ai/task_completed` params.
fn task_completed_will_wake(
gateway_rx: &mut mpsc::UnboundedReceiver<xai_acp_lib::AcpClientMessage>,
) -> Option<bool> {
@ -1065,56 +1149,209 @@ mod tests {
/// follows, emit when nothing else will mark the moment).
#[tokio::test]
async fn task_completed_notification_stamps_will_wake() {
// Wake leg: auto-wake enabled, no suppression.
let (config, mut gateway_rx, _persistence_rx, _cmd_rx) = make_test_config_full();
let (config, mut gateway_rx, _persistence_rx, mut cmd_rx) = make_test_config_full();
config
.task_output_tool_name
.set(Some("get_command_or_subagent_output".to_string()))
.expect("slot is fresh in this test fixture");
let (trace_tx, mut trace_rx) = mpsc::unbounded_channel();
*config
.synthetic_trace_tx
.lock()
.unwrap_or_else(|e| e.into_inner()) = Some(trace_tx);
let mut offsets = HashMap::new();
handle_notification(
handle_notification_with_admission(
&config,
ToolNotification::TaskCompleted(make_task_snapshot("bg-wake", TaskKind::Bash)),
&mut offsets,
&mut cmd_rx,
true,
)
.await;
assert!(matches!(
cmd_rx.recv().await,
Some(SessionCommand::Prompt { .. })
));
match cmd_rx.recv().await {
Some(SessionCommand::CopyFile { respond_to }) => drop(respond_to),
_ => panic!("trace copy must follow accepted prompt admission"),
}
assert_eq!(
task_completed_will_wake(&mut gateway_rx),
Some(true),
"an auto-woken completion must stamp will_wake: true"
);
assert!(
trace_rx.try_recv().is_ok(),
"accepted admission must request a synthetic-turn trace"
);
// Suppressed leg: goal loop active — no wake follows the chip.
let (config, mut gateway_rx, _persistence_rx, _cmd_rx) = make_test_config_full();
config
.goal_loop_active
.store(true, std::sync::atomic::Ordering::Relaxed);
let (config, mut gateway_rx, mut persistence_rx, mut cmd_rx) = make_test_config_full();
let (trace_tx, mut trace_rx) = mpsc::unbounded_channel();
*config
.synthetic_trace_tx
.lock()
.unwrap_or_else(|e| e.into_inner()) = Some(trace_tx);
let mut offsets = HashMap::new();
handle_notification(
handle_notification_with_admission(
&config,
ToolNotification::TaskCompleted(make_task_snapshot("bg-goal", TaskKind::Bash)),
ToolNotification::TaskCompleted(make_task_snapshot("bg-declined", TaskKind::Bash)),
&mut offsets,
&mut cmd_rx,
false,
)
.await;
assert_eq!(
task_completed_will_wake(&mut gateway_rx),
Some(false),
"a suppressed completion must stamp will_wake: false"
"an actor-declined completion must stamp will_wake: false"
);
assert!(
config.task_completion_reservations.contains("bg-declined"),
"the actor owns reservation release after queuing the deferred fallback"
);
assert!(
trace_rx.try_recv().is_err(),
"declined admission must not request a synthetic-turn trace"
);
assert!(matches!(
cmd_rx.try_recv(),
Ok(SessionCommand::Prompt { .. })
));
assert!(matches!(
cmd_rx.try_recv(),
Ok(SessionCommand::DispatchNotificationHook { .. })
));
let mut persisted = false;
while let Ok(message) = persistence_rx.try_recv() {
if let PersistenceMsg::Update(crate::session::storage::SessionUpdate::Xai(update)) =
message
&& matches!(
&update.update,
crate::extensions::notification::SessionUpdate::TaskCompleted { .. }
)
{
persisted = true;
}
}
assert!(
persisted,
"declined admission must still persist x.ai/task_completed"
);
}
#[tokio::test(start_paused = true)]
async fn stalled_admission_is_bounded_and_task_completion_still_emits() {
let (config, mut gateway_rx, mut persistence_rx, mut cmd_rx) = make_test_config_full_raw();
config
.task_output_tool_name
.set(Some("get_command_or_subagent_output".to_string()))
.expect("slot is fresh in this test fixture");
let mut offsets = HashMap::new();
let notification = handle_notification(
&config,
ToolNotification::TaskCompleted(make_task_snapshot("bg-stalled", TaskKind::Bash)),
&mut offsets,
);
tokio::pin!(notification);
tokio::select! {
_ = &mut notification => panic!("admission should still be waiting"),
command = cmd_rx.recv() => assert!(matches!(command, Some(SessionCommand::Prompt { .. }))),
}
tokio::time::advance(TASK_WAKE_ADMISSION_TIMEOUT + std::time::Duration::from_millis(1))
.await;
tokio::task::yield_now().await;
notification.await;
assert_eq!(task_completed_will_wake(&mut gateway_rx), Some(false));
assert!(
config.task_completion_reservations.contains("bg-stalled"),
"a timed-out admission may still be handled and deferred by the actor"
);
let mut persisted_completion = false;
while let Ok(message) = persistence_rx.try_recv() {
if let PersistenceMsg::Update(crate::session::storage::SessionUpdate::Xai(update)) =
message
&& matches!(
&update.update,
crate::extensions::notification::SessionUpdate::TaskCompleted { .. }
)
{
persisted_completion = true;
}
}
assert!(persisted_completion);
}
#[tokio::test(start_paused = true)]
async fn timed_out_monitor_admission_queues_one_fallback_and_late_actor_drops_prompt() {
let (config, mut gateway_rx, _persistence_rx, mut cmd_rx) = make_test_config_full_raw();
config
.task_output_tool_name
.set(Some("get_command_or_subagent_output".to_string()))
.expect("slot is fresh in this test fixture");
let mut offsets = HashMap::new();
let notification = handle_notification(
&config,
ToolNotification::TaskCompleted(make_task_snapshot("mon-timeout", TaskKind::Monitor)),
&mut offsets,
);
tokio::pin!(notification);
let prompt = tokio::select! {
_ = &mut notification => panic!("admission should still be waiting"),
command = cmd_rx.recv() => command.expect("prompt command"),
};
tokio::time::advance(TASK_WAKE_ADMISSION_TIMEOUT + std::time::Duration::from_millis(1))
.await;
tokio::task::yield_now().await;
notification.await;
let SessionCommand::Prompt {
admission: Some(admission),
respond_to,
..
} = prompt
else {
panic!("expected task wake prompt");
};
assert!(matches!(
admission.fallback.source,
NotificationSource::MonitorCompleted { ref task_id } if task_id == "mon-timeout"
));
assert!(admission.respond_to.send(true).is_err());
let _ = respond_to.send(Ok(crate::session::commands::PromptTurnOk {
stop_reason: acp::StopReason::Cancelled,
total_tokens: 0,
turn_snapshot: None,
completion_kind: crate::session::commands::PromptCompletionKind::RemovedFromQueue,
structured_output: None,
usage: None,
}));
assert!(matches!(
cmd_rx.try_recv(),
Ok(SessionCommand::DispatchNotificationHook { .. })
));
assert!(cmd_rx.try_recv().is_err());
assert_eq!(task_completed_will_wake(&mut gateway_rx), Some(false));
assert!(
config.task_completion_reservations.contains("mon-timeout"),
"the late actor fallback retains the reservation until user delivery"
);
}
/// Dead session actor: the synthetic Prompt enqueue fails, so no wake will
/// ever run — the notification must stamp `will_wake: false`, not promise a
/// wake the send didn't queue (the pager would suppress its between-turns
/// status line for a wake that never comes).
#[tokio::test]
async fn task_completed_stamps_will_wake_false_when_session_channel_closed() {
let (config, mut gateway_rx, _persistence_rx, cmd_rx) = make_test_config_full();
let (config, mut gateway_rx, _persistence_rx, cmd_rx) = make_test_config_full_raw();
config
.task_output_tool_name
.set(Some("get_command_or_subagent_output".to_string()))
.expect("slot is fresh in this test fixture");
drop(cmd_rx);
config
.task_completion_reservations
.reserve("bg-dead".into());
let mut offsets = HashMap::new();
handle_notification(
&config,
@ -1127,6 +1364,9 @@ mod tests {
Some(false),
"a completion whose wake prompt could not be enqueued must stamp will_wake: false"
);
assert!(config.task_completion_reservations.contains("bg-dead"));
config.task_completion_reservations.release("bg-dead");
assert!(!config.task_completion_reservations.contains("bg-dead"));
}
/// Gap 1 (adjacent branch): the goal-loop arm sits BEFORE the
@ -1165,8 +1405,7 @@ mod tests {
cmd_rx.try_recv().is_err(),
"goal-loop-active completion must not InjectNotification with auto-wake disabled"
);
// And not marked auto-wake-delivered.
assert!(config.auto_wake_delivered.snapshot().is_empty());
assert!(config.task_completion_reservations.snapshot().is_empty());
}
/// Natural monitor exit (including exit code 0) must immediate-auto-wake
@ -1186,27 +1425,15 @@ mod tests {
snapshot.exit_code = Some(0);
let mut offsets = HashMap::new();
handle_notification(
handle_notification_with_admission(
&config,
ToolNotification::TaskCompleted(snapshot),
&mut offsets,
&mut cmd_rx,
true,
)
.await;
// Drop queued pipeline events first (sole-wake guarantee).
match cmd_rx
.try_recv()
.expect("expected DropMonitorNotifications")
{
SessionCommand::DropMonitorNotifications { task_id } => {
assert_eq!(task_id, "mon-456");
}
_ => panic!("expected DropMonitorNotifications before auto-wake Prompt"),
}
assert!(matches!(
cmd_rx.try_recv(),
Ok(SessionCommand::CopyFile { .. })
));
let cmd = cmd_rx.try_recv().expect("expected Prompt auto-wake");
match cmd {
SessionCommand::Prompt {
@ -1236,22 +1463,79 @@ mod tests {
}
_ => panic!("expected Prompt auto-wake for natural monitor exit"),
}
match cmd_rx
.try_recv()
.expect("accepted monitor wake must drop pipeline notifications")
{
SessionCommand::DropMonitorNotifications { task_id } => {
assert_eq!(task_id, "mon-456");
}
_ => panic!("expected DropMonitorNotifications after accepted Prompt"),
}
assert!(matches!(
cmd_rx.try_recv(),
Ok(SessionCommand::DispatchNotificationHook { .. })
));
assert_eq!(
config.auto_wake_delivered.snapshot(),
config.task_completion_reservations.snapshot(),
vec!["mon-456".to_string()],
);
}
/// After TaskCompleted auto-wake marked the task delivered, late pipeline
#[tokio::test]
async fn declined_quiet_monitor_wake_queues_canonical_deferred_completion() {
let (config, _gateway_rx, mut persistence_rx, mut cmd_rx) = make_test_config_full();
config
.task_output_tool_name
.set(Some("get_command_or_subagent_output".to_string()))
.expect("slot is fresh in this test fixture");
let mut offsets = HashMap::new();
handle_notification_with_admission(
&config,
ToolNotification::TaskCompleted(make_task_snapshot("mon-declined", TaskKind::Monitor)),
&mut offsets,
&mut cmd_rx,
false,
)
.await;
assert!(matches!(
cmd_rx.try_recv(),
Ok(SessionCommand::Prompt { .. })
));
assert!(matches!(
cmd_rx.try_recv(),
Ok(SessionCommand::DispatchNotificationHook { .. })
));
assert!(cmd_rx.try_recv().is_err());
let mut persisted_completion = false;
while let Ok(message) = persistence_rx.try_recv() {
if let PersistenceMsg::Update(crate::session::storage::SessionUpdate::Xai(update)) =
message
&& matches!(
&update.update,
crate::extensions::notification::SessionUpdate::TaskCompleted { .. }
)
{
persisted_completion = true;
}
}
assert!(persisted_completion);
assert!(
config.task_completion_reservations.contains("mon-declined"),
"the actor owns reservation release after queuing the deferred fallback"
);
}
/// After TaskCompleted auto-wake reserves the task, late pipeline
/// MonitorEvents must not inject another model-facing notification.
#[tokio::test]
async fn monitor_event_skipped_after_task_completed_auto_wake() {
let (config, mut cmd_rx) = make_test_config();
config.auto_wake_delivered.insert("mon-done".into());
config
.task_completion_reservations
.reserve("mon-done".into());
let mut offsets = HashMap::new();
handle_notification(
@ -1303,7 +1587,7 @@ mod tests {
cmd_rx.try_recv().is_err(),
"explicitly-killed monitor must not auto-wake"
);
assert!(config.auto_wake_delivered.snapshot().is_empty());
assert!(config.task_completion_reservations.snapshot().is_empty());
}
/// Goal-loop suppression applies to monitor completions too.
@ -1336,7 +1620,7 @@ mod tests {
cmd_rx.try_recv().is_err(),
"goal-loop-active monitor completion must not auto-wake"
);
assert!(config.auto_wake_delivered.snapshot().is_empty());
assert!(config.task_completion_reservations.snapshot().is_empty());
}
#[tokio::test]
@ -1808,16 +2092,14 @@ mod tests {
#[tokio::test]
async fn bash_completion_uses_single_task_id_clone() {
// Verify the task_id appears in the prompt_id (auto-wake path).
let (config, mut cmd_rx) = make_test_config();
let snapshot = make_task_snapshot("unique-id-789", TaskKind::Bash);
let notification = ToolNotification::TaskCompleted(snapshot);
let mut offsets = HashMap::new();
handle_notification(&config, notification, &mut offsets).await;
handle_notification_with_admission(&config, notification, &mut offsets, &mut cmd_rx, true)
.await;
// Skip CopyFile
let _ = cmd_rx.try_recv().unwrap();
let cmd = cmd_rx.try_recv().unwrap();
if let SessionCommand::Prompt { prompt_id, .. } = cmd {
assert_eq!(prompt_id, "task-completed-unique-id-789");
@ -2044,7 +2326,6 @@ mod tests {
/// Extract the auto-wake prompt text emitted on the session command channel.
fn auto_wake_prompt_text(cmd_rx: &mut mpsc::UnboundedReceiver<SessionCommand>) -> String {
let _ = cmd_rx.try_recv().expect("expected CopyFile");
let cmd = cmd_rx.try_recv().expect("expected Prompt");
match cmd {
SessionCommand::Prompt { prompt_blocks, .. } => match &prompt_blocks[0] {
@ -2087,10 +2368,12 @@ mod tests {
.expect("fresh slot");
let snapshot = make_large_bash_snapshot("bg-disk-1", output_file.clone());
let mut offsets = HashMap::new();
handle_notification(
handle_notification_with_admission(
&config_auto,
ToolNotification::TaskCompleted(snapshot),
&mut offsets,
&mut cmd_rx_auto,
true,
)
.await;
let prompt = auto_wake_prompt_text(&mut cmd_rx_auto);

17
crates/codegen/xai-grok-shell/src/tools/tool_context.rs
View file

@ -77,10 +77,10 @@ pub struct ToolContext {
/// synthetic user message before the next sampling step.
pub monitor_event_buffer:
Option<xai_grok_tools::implementations::grok_build::task::types::MonitorEventBuffer>,
/// Shared set of IDs delivered via auto-wake synthetic prompts.
/// Used by `TaskCompletionReminder` to suppress duplicate reminders.
pub auto_wake_delivered:
Option<xai_grok_tools::reminders::task_completion::AutoWakeDeliveredIds>,
pub task_completion_reservations:
Option<xai_grok_tools::reminders::task_completion::TaskCompletionReservations>,
pub task_wake_suppressed:
Option<xai_grok_tools::reminders::task_completion::TaskWakeSuppressed>,
/// Channel for requesting trace uploads for synthetic auto-wake turns.
pub(crate) synthetic_trace_tx:
Option<tokio::sync::mpsc::UnboundedSender<crate::upload::turn::SyntheticTurnTraceRequest>>,
@ -143,7 +143,8 @@ impl ToolContext {
lsp_server_names: Vec::new(),
is_turn_active: None,
monitor_event_buffer: None,
auto_wake_delivered: None,
task_completion_reservations: None,
task_wake_suppressed: None,
synthetic_trace_tx: None,
synthetic_trace_tx_shared: None,
task_output_tool_name:
@ -179,7 +180,8 @@ impl ToolContext {
lsp_server_names: Vec::new(),
is_turn_active: None,
monitor_event_buffer: None,
auto_wake_delivered: None,
task_completion_reservations: None,
task_wake_suppressed: None,
synthetic_trace_tx: None,
synthetic_trace_tx_shared: None,
task_output_tool_name:
@ -235,7 +237,8 @@ mod tests {
lsp_server_names: Vec::new(),
is_turn_active: None,
monitor_event_buffer: None,
auto_wake_delivered: None,
task_completion_reservations: None,
task_wake_suppressed: None,
synthetic_trace_tx: None,
synthetic_trace_tx_shared: None,
task_output_tool_name:

77
crates/codegen/xai-grok-shell/src/util/config/resolve/compaction.rs
View file

@ -1,6 +1,38 @@
/// Default auto-compact threshold (% of context window) when no source sets it.
pub const DEFAULT_AUTO_COMPACT_THRESHOLD_PERCENT: u8 = 85;
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum CompactionToolChoice {
#[default]
Auto,
None,
}
impl std::str::FromStr for CompactionToolChoice {
type Err = ();
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.trim().to_ascii_lowercase().as_str() {
"auto" => Ok(Self::Auto),
"none" => Ok(Self::None),
_ => Err(()),
}
}
}
pub(crate) const ENV_COMPACTION_TOOL_CHOICE: &str = "GROK_COMPACTION_TOOL_CHOICE";
pub fn resolve_compaction_tool_choice_from(
env: Option<&str>,
config: Option<&str>,
remote: Option<&str>,
) -> CompactionToolChoice {
env.and_then(|s| s.parse().ok())
.or_else(|| config.and_then(|s| s.parse().ok()))
.or_else(|| remote.and_then(|s| s.parse().ok()))
.unwrap_or_default()
}
/// Env-var override for `auto_compact_threshold_percent`. Parsed as `u8`;
/// out-of-range or unparseable values are ignored.
pub(crate) const ENV_AUTO_COMPACT_THRESHOLD_PERCENT: &str = "GROK_AUTO_COMPACT_THRESHOLD_PERCENT";
@ -140,3 +172,48 @@ mod compaction_wall_clock_budget_tests {
assert_eq!(resolve(Some(5)), 5); // low values pass through (warned, not clamped)
}
}
#[cfg(test)]
mod compaction_tool_choice_tests {
use super::{CompactionToolChoice, resolve_compaction_tool_choice_from as resolve};
#[test]
fn default_is_auto() {
assert_eq!(resolve(None, None, None), CompactionToolChoice::Auto);
}
#[test]
fn precedence_env_over_config_over_remote() {
assert_eq!(
resolve(Some("none"), Some("auto"), Some("auto")),
CompactionToolChoice::None
);
assert_eq!(
resolve(None, Some("none"), Some("auto")),
CompactionToolChoice::None
);
assert_eq!(
resolve(None, None, Some("none")),
CompactionToolChoice::None
);
}
#[test]
fn garbage_falls_through() {
assert_eq!(
resolve(Some("garbage"), None, Some("none")),
CompactionToolChoice::None
);
assert_eq!(
resolve(Some("garbage"), Some("also-bad"), None),
CompactionToolChoice::Auto
);
}
#[test]
fn from_str_case_insensitive() {
assert_eq!("AUTO".parse(), Ok(CompactionToolChoice::Auto));
assert_eq!(" None ".parse(), Ok(CompactionToolChoice::None));
assert!("required".parse::<CompactionToolChoice>().is_err());
}
}

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

@ -62,6 +62,137 @@ fn resolve_search_tool_enabled(
env.or(config).or(managed).unwrap_or(true)
}
const ENV_LOGIN_SHELL_CAPTURE: &str = "GROK_LOGIN_ENV";
fn login_shell_capture_from_toml(v: Option<&TomlValue>) -> Option<bool> {
v?.get("toolset")?
.get("bash")?
.get("login_shell_capture")?
.as_bool()
}
pub fn resolve_login_shell_capture(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, "login_shell_capture: failed to load config layers");
None
}
};
resolve_login_shell_capture_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_login_shell_capture_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_LOGIN_SHELL_CAPTURE)
.requirement(login_shell_capture_from_toml(requirements))
.config(login_shell_capture_from_toml(user))
.managed(
login_shell_capture_from_toml(managed)
.or_else(|| login_shell_capture_from_toml(system_managed)),
)
.feature_flag(remote)
.default(true)
.resolve()
.value
}
#[cfg(test)]
mod login_shell_capture_tests {
use super::{ENV_LOGIN_SHELL_CAPTURE, resolve_login_shell_capture_tiers};
use toml::Value as TomlValue;
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_LOGIN_SHELL_CAPTURE) };
g
}
fn cfg(enabled: bool) -> TomlValue {
toml::from_str(&format!(
"[toolset.bash]\nlogin_shell_capture = {enabled}\n"
))
.unwrap()
}
#[test]
fn defaults_on() {
let _g = guard();
assert!(resolve_login_shell_capture_tiers(
None, None, None, None, None
));
}
#[test]
fn remote_flag_can_disable() {
let _g = guard();
assert!(!resolve_login_shell_capture_tiers(
None,
None,
None,
None,
Some(false)
));
}
#[test]
fn user_config_beats_remote() {
let _g = guard();
assert!(resolve_login_shell_capture_tiers(
None,
Some(&cfg(true)),
None,
None,
Some(false)
));
assert!(!resolve_login_shell_capture_tiers(
None,
Some(&cfg(false)),
None,
None,
Some(true)
));
}
#[test]
fn env_beats_config_and_remote() {
let _g = guard();
unsafe { std::env::set_var(ENV_LOGIN_SHELL_CAPTURE, "0") };
let off = resolve_login_shell_capture_tiers(None, Some(&cfg(true)), None, None, Some(true));
unsafe { std::env::remove_var(ENV_LOGIN_SHELL_CAPTURE) };
assert!(!off);
}
#[test]
fn requirements_win_outright() {
let _g = guard();
unsafe { std::env::set_var(ENV_LOGIN_SHELL_CAPTURE, "1") };
let off = resolve_login_shell_capture_tiers(
Some(&cfg(false)),
Some(&cfg(true)),
None,
None,
Some(true),
);
unsafe { std::env::remove_var(ENV_LOGIN_SHELL_CAPTURE) };
assert!(!off);
}
}
/// 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.

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

@ -23,6 +23,10 @@ pub async fn set_show_timeline(value: bool) -> Result<()> {
update_config(|cfg| cfg.ui.show_timeline = Some(value)).await
}
pub async fn set_page_flip_on_send(value: bool) -> Result<()> {
update_config(|cfg| cfg.ui.page_flip_on_send = Some(value)).await
}
/// Persist `[ui].simple_mode` via `update_config`. Same `Option<bool>`
/// shape as `show_timestamps`.
pub async fn set_simple_mode(value: bool) -> Result<()> {

59
crates/codegen/xai-grok-shell/src/util/mod.rs
View file

@ -8,6 +8,33 @@ pub mod hooks;
// unchanged.
pub use xai_grok_shell_base::util::*;
pub(crate) fn is_user_instruction_path(
path: &std::path::Path,
grok_home: &std::path::Path,
vendor_homes: &[(std::path::PathBuf, bool)],
workspace_root: Option<&std::path::Path>,
) -> bool {
let parent = path.parent();
let grok_rules = grok_home.join("rules");
let is_exact_home_surface = parent
.is_some_and(|parent| parent == grok_home || parent == grok_rules)
|| vendor_homes.iter().any(|(vendor_home, named_enabled)| {
parent.is_some_and(|parent| {
(*named_enabled && parent == vendor_home) || parent == vendor_home.join("rules")
})
});
if is_exact_home_surface {
return true;
}
if workspace_root.is_some_and(|root| path.starts_with(root)) {
return false;
}
path.starts_with(grok_home)
|| vendor_homes
.iter()
.any(|(vendor_home, _)| path.starts_with(vendor_home))
}
/// Aborts the wrapped tokio task when dropped.
///
/// Use to tie a spawned helper task's lifetime to an async scope so that
@ -22,3 +49,35 @@ impl Drop for AbortOnDrop {
self.0.abort();
}
}
#[cfg(test)]
mod is_user_instruction_path_tests {
use super::is_user_instruction_path;
use std::path::Path;
#[test]
fn grok_home_named_file_nested_in_workspace_is_user_scoped() {
assert!(is_user_instruction_path(
Path::new("/repo/config/AGENTS.md"),
Path::new("/repo/config"),
&[],
Some(Path::new("/repo")),
));
assert!(!is_user_instruction_path(
Path::new("/repo/config/src/AGENTS.md"),
Path::new("/repo/config"),
&[],
Some(Path::new("/repo")),
));
}
#[test]
fn workspace_descendants_under_grok_home_stay_project_scoped() {
assert!(!is_user_instruction_path(
Path::new("/custom/grok/worktrees/repo/src/AGENTS.md"),
Path::new("/custom/grok"),
&[],
Some(Path::new("/custom/grok/worktrees/repo")),
));
}
}