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:
- Detect the herdr multiplexer
- Mark /gboom as non-production code
- Bound peak memory when loading a large session
- Add a subagent lifecycle soak bounding threads, fds, and heap
- Stream inherited replay to bound fork memory
- Copy full plan from plan approval with y
- Stop armed signature verification from deleting the managed-deny smoke policy
- Add source-tagged terminal version telemetry
- Show the UI instantly and fetch models and settings in the background
- Session test helpers
- computer_reason on the ConversationHistoryDone trailer
This commit is contained in:
grokkybara[bot] 2026-07-26 20:03:03 +01:00
commit b41c75a578
92 changed files with 9410 additions and 3788 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -8,7 +8,11 @@ edition.workspace = true
default = []
unstable = []
dhat-heap = ["dep:dhat"]
default-bazel = []
# Session synthesis + in-process e2e harness (`session::testkit`) for soak,
# load, and bench tests. Off by default; the tests/benches that use it declare
# it via `required-features`.
test-support = []
default-bazel = ["test-support"]
[dependencies]
dunce = { workspace = true }
@ -215,6 +219,24 @@ harness = false
[[bench]]
name = "fork_copy"
harness = false
required-features = ["test-support"]
# Consume `session::testkit`, so they need the gate (on by default under Bazel).
[[test]]
name = "test_session_load_memory"
required-features = ["test-support"]
[[test]]
name = "session_fork_replay_memory"
required-features = ["test-support"]
[[test]]
name = "session_load_perf"
required-features = ["test-support"]
[[test]]
name = "testkit_synth_roundtrip"
required-features = ["test-support"]
[lints]
workspace = true

70
crates/codegen/xai-grok-shell/benches/fork_copy.rs
View file

@ -1,8 +1,8 @@
//! Fork-path benchmark and profiling workbench.
//!
//! Synthesizes a session whose `updates.jsonl` matches a configurable target
//! size (realistic mixed update shapes: user/agent chunks, tool calls with
//! large results), then measures `StorageAdapter::copy_session_data` — the
//! size (production byte/line shape: user and agent chunks plus one bulky
//! trailing chunk), then measures `StorageAdapter::copy_session_data`, the
//! path that materializes the whole file and produced multi-GB RSS spikes on
//! large production sessions. Also the substrate for allocation/CPU profiling
//! (`cargo flamegraph --bench fork_copy`, dhat) and future peak-RSS bounds.
@ -13,74 +13,14 @@
use std::hint::black_box;
use std::time::Duration;
use acp::{ContentBlock, ContentChunk, TextContent};
use agent_client_protocol as acp;
use criterion::{
BenchmarkId, Criterion, SamplingMode, Throughput, criterion_group, criterion_main,
};
use tempfile::TempDir;
use xai_grok_shell::session::info::Info;
use xai_grok_shell::session::storage::{
CopySessionOptions, JsonlStorageAdapter, SessionUpdate, StorageAdapter,
};
/// One synthetic "turn": a user chunk, agent chunks, and a bulky tool result,
/// so line-size distribution and parse cost resemble production sessions.
fn turn_updates(info: &Info, turn: usize) -> Vec<SessionUpdate> {
let text = |s: String| ContentChunk::new(ContentBlock::Text(TextContent::new(s)));
let notify =
|u| SessionUpdate::Acp(Box::new(acp::SessionNotification::new(info.id.clone(), u)));
let mut updates = vec![notify(acp::SessionUpdate::UserMessageChunk(text(format!(
"prompt {turn}: check the build and summarize failures"
))))];
for i in 0..8 {
updates.push(notify(acp::SessionUpdate::AgentMessageChunk(text(format!(
"agent chunk {turn}/{i}: analyzing module {i} for regressions and drafting a fix plan"
)))));
}
// ~4 KB tool-result payload: the dominant byte source in real sessions.
updates.push(notify(acp::SessionUpdate::AgentMessageChunk(text(
format!("tool result {turn}: {}", "x".repeat(4096)),
))));
updates
}
/// Build a session dir whose `updates.jsonl` is at least `target_bytes`.
fn synthesize_session(root: &TempDir, target_bytes: u64) -> Info {
let adapter = JsonlStorageAdapter::with_root(root.path().to_path_buf());
let info = Info {
id: acp::SessionId::new("fork-bench-src"),
cwd: "/bench/workspace".to_string(),
};
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("bench runtime");
rt.block_on(async {
adapter
.init_session(&info, acp::ModelId::new("bench-model"))
.await
.expect("init session");
let updates_path = adapter.updates_file_path(&info).expect("updates path");
let mut turn = 0usize;
loop {
for update in turn_updates(&info, turn) {
adapter.append_update(&info, &update).await.expect("append");
}
turn += 1;
// Stat every 32 turns; sizes only grow.
if turn % 32 == 0
&& std::fs::metadata(&updates_path)
.map(|m| m.len())
.unwrap_or(0)
>= target_bytes
{
break;
}
}
});
info
}
use xai_grok_shell::session::storage::{CopySessionOptions, JsonlStorageAdapter, StorageAdapter};
use xai_grok_shell::session::testkit::synth::synthesize_to_target_bytes;
fn bench_fork_copy(c: &mut Criterion) {
let target_mb: u64 = std::env::var("FORK_BENCH_MB")
@ -88,7 +28,7 @@ fn bench_fork_copy(c: &mut Criterion) {
.and_then(|v| v.parse().ok())
.unwrap_or(16);
let root = TempDir::new().expect("tempdir");
let source = synthesize_session(&root, target_mb * 1024 * 1024);
let source = synthesize_to_target_bytes(root.path(), target_mb * 1024 * 1024);
let adapter = JsonlStorageAdapter::with_root(root.path().to_path_buf());
let updates_len = std::fs::metadata(adapter.updates_file_path(&source).expect("updates path"))
.expect("updates.jsonl")

399
crates/codegen/xai-grok-shell/src/agent/app.rs
View file

@ -20,7 +20,7 @@ use crate::agent::config::{Config as AgentConfig, ModelEntry};
use crate::agent::init::{bootstrap, exit_on_config_error};
use crate::agent::models::{ModelFetchAuth, prefetch_models_blocking};
use crate::agent::mvp_agent::MvpAgent;
use crate::auth::{AuthManager, AuthMode, GrokAuth, run_auth_flow};
use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig, run_auth_flow};
use crate::util::grok_home;
use dirs;
@ -176,24 +176,6 @@ pub(crate) async fn run_auto_update_checker(
}
}
/// Prefetch models from the API (must be called outside LocalSet).
async fn prefetch_models(agent_config: &AgentConfig) -> Option<IndexMap<String, ModelEntry>> {
let auth = agent_config.create_auth_manager().current();
let endpoints = agent_config.endpoints.clone();
let fetch_auth = ModelFetchAuth::resolve(&endpoints, auth.is_some());
if auth.is_some() || endpoints.has_custom_endpoint() || fetch_auth != ModelFetchAuth::Session {
tokio::task::spawn_blocking(move || {
prefetch_models_blocking(&endpoints, auth.as_ref(), fetch_auth)
})
.await
.ok()
.flatten()
} else {
None
}
}
/// Spawn the agent inside a LocalSet and return a handle to the I/O future.
fn spawn_agent_local(
agent_config: AgentConfig,
@ -207,6 +189,8 @@ fn spawn_agent_local(
let gateway = GatewaySender::new(gw_tx);
let mut agent = MvpAgent::new(gateway, &agent_config, auth_manager, prefetched_models)
.unwrap_or_else(exit_on_config_error);
// Background the catalog refresh so readiness never blocks on the network.
agent.models_manager.spawn_background_refresh();
if let Some(mc) = memory_config {
agent.set_memory_config(mc);
}
@ -333,12 +317,7 @@ pub async fn run_stdio_agent(
let _total_timer = crate::instrumentation_timer!("startup.stdio_agent_total");
let outgoing = tokio::io::stdout().compat_write();
let prefetched_models = if prefetched_models.is_some() {
prefetched_models
} else {
let _timer = crate::instrumentation_timer!("startup.stdio_prefetch_models");
prefetch_models(agent_config).await
};
// Non-blocking boot: catalog refreshes in the background, not before readiness.
let agent_config = agent_config.clone();
// Use a simplex intermediary between stdin and the agent so we can
@ -400,6 +379,9 @@ pub async fn run_stdio_agent(
// Restore managed policy right before bootstrap reads it (no stale window after prefetch).
crate::managed_config::ensure_managed_policy_present(&auth_manager).await;
// Fail-closed external-OTEL gate: suppress until settings resolve,
// opening now only for a pure env-API-key user (no remote policy).
apply_otel_config(&auth_manager, &agent_config.grok_com_config);
let handle_io = spawn_agent_local(
agent_config,
auth_manager,
@ -494,8 +476,8 @@ async fn run_headless_inner(
)
.await?
} else {
// Don't pre-resolve via try_ensure_session_noninteractive: run_auth_flow below
// already mints external/devbox creds, so it would run the provider twice.
// Don't pre-resolve auth here: run_auth_flow below already mints
// external/devbox creds, so it would run the provider twice.
let auth_manager = Arc::new(AuthManager::new(&grok_home::grok_home(), ctx.clone()));
if crate::agent::auth_method::has_xai_api_key_env()
&& ctx.auth_provider_command.is_none()
@ -966,6 +948,33 @@ impl DeferredRelayArm {
}
}
/// Close the external-OTEL gate before telemetry init; see
/// [`crate::agent::otel_gate`].
pub fn suppress_otel() {
crate::agent::otel_gate::suppress();
}
/// Startup external-OTEL gate for an in-process (embedded) agent. Mirrors the
/// leader startup gate so the pager process is fail-closed by construction at the
/// agent boundary: suppress until the agent's first settings outcome, except a
/// pure env-API-key user (no session now, none minting) whose stream has no
/// remote policy and may emit immediately.
pub fn apply_otel_config(auth_manager: &AuthManager, grok_com_config: &GrokComConfig) {
suppress_otel();
// Session presence is disk-based (valid or expired), not refresh success: an
// expired session the refresher will renew still has a remote policy, so it
// must keep the gate closed.
let has_session = auth_manager.current().is_some() || auth_manager.read_disk_auth().is_some();
if crate::agent::otel_gate::should_open_at_startup(crate::agent::otel_gate::StartupGate {
has_session,
has_api_key_env: crate::agent::auth_method::has_xai_api_key_env(),
session_pending: crate::agent::otel_gate::is_session_pending(has_session, grok_com_config),
remote_fetch_enabled: crate::util::config::resolve_remote_fetch_enabled(),
}) {
crate::agent::otel_gate::open_at_startup();
}
}
/// Run the agent in leader mode, accepting IPC connections from multiple clients.
/// When a grok.com session is present, the leader connects to the websocket relay
/// after startup (post-auth, post-prefetch); BYOK / no-session leaders start
@ -979,11 +988,12 @@ impl DeferredRelayArm {
/// 2. Socket cleanup, channel + readiness-watch creation.
/// 3. IPC server started (`tokio::spawn`) — socket bound HERE, before auth.
/// 4. Wait for socket to appear (fast: < 100 ms).
/// 5. Auth + model prefetch (slow path, but socket already available to clients).
/// - Auth resolves non-interactively; `None` (BYOK / no session) is not an
/// error — the relay is gated off and login is deferred to ACP.
/// 6. `ready_tx.send(true)` — unblocks ACP forwarding in the IPC server.
/// 7. LocalSet: agent, IPC↔agent bridges, WS↔agent bridges, relay, config watcher.
/// 5. Lock handoff with spawner (if launched via connect_or_spawn).
/// 6. Bounded non-interactive auth (no blocking model/settings prefetch; those
/// stream in after readiness). `None` (BYOK / no session) is not an error:
/// the relay stays off and a background cold-mint / re-login can start it later.
/// 7. `ready_tx.send(true)` — unblocks ACP forwarding in the IPC server.
/// 8. LocalSet: agent, IPC↔agent bridges, WS↔agent bridges, relay, config watcher.
///
/// # Arguments
///
@ -991,7 +1001,7 @@ impl DeferredRelayArm {
/// * `no_exit_on_disconnect` - If true, the leader will not exit when all clients disconnect
/// * `relay_on_demand` - If true, defer the grok.com relay WebSocket until the
/// first headless IPC client registers; if false (default), connect eagerly at
/// startup. See [`spawn_leader_relay`].
/// startup; a session acquired later arms it via [`DeferredRelayArm`].
pub async fn run_leader(
agent_config: &AgentConfig,
no_exit_on_disconnect: bool,
@ -999,7 +1009,6 @@ pub async fn run_leader(
auto_update_check: Option<LeaderAutoUpdateConfig>,
memory_config: Option<crate::config::MemoryConfig>,
) -> anyhow::Result<()> {
use crate::agent::relay::RelayConfig;
use crate::leader::{
LeaderLock, LeaderServerControlState, LeaderServerMetadata, LockError, ShutdownReason,
compute_ws_url_suffix, run_leader_server,
@ -1125,7 +1134,8 @@ pub async fn run_leader(
// Relay demand watch: the IPC server flips this to `true` when the first
// headless client registers. Only consulted when `relay_on_demand` is set
// (leaders auto-spawned by interactive clients); an eager leader connects
// the relay at startup and ignores it. See `spawn_leader_relay`.
// the relay once a session is present and ignores it. See
// the config-update loop's `DeferredRelayArm`.
let (relay_demand_tx, relay_demand_rx) = watch::channel(false);
let client_count = Arc::new(AtomicUsize::new(0));
@ -1202,51 +1212,56 @@ pub async fn run_leader(
// messages during this window receive a `leader_starting` error and can retry.
let ctx = &agent_config.grok_com_config;
// Never interactive: a detached leader has no TTY (forcing OAuth here hung BYOK).
let auth: Option<GrokAuth> = crate::auth::try_ensure_session_noninteractive(ctx).await;
suppress_otel(); // idempotent re-assert
// No-mint on the readiness path: a cached/expired session + a bounded
// (~5s) refresh only. A session-less-but-mintable leader is minted by the
// post-readiness background task below, so readiness never blocks on the
// provider command (which could take up to STARTUP_AUTH_TIMEOUT ~60s).
let auth: Option<GrokAuth> = crate::auth::try_noninteractive_auth_no_mint(ctx).await;
// ── Phase 6b: Legacy devbox auth migration ─────────────────────────────
let auth: Option<GrokAuth> = migrate_devbox_auth_if_legacy(auth, &agent_config).await;
let auth_for_prefetch: Option<GrokAuth> = auth.clone();
let endpoints_for_prefetch = agent_config.endpoints.clone();
let fetch_auth_for_prefetch = ModelFetchAuth::resolve(&endpoints_for_prefetch, auth.is_some());
// The shared pair helper owns the remote_fetch gate for both halves, so a
// disabled knob cannot block leader readiness on settings retries.
let (prefetched_models, remote_settings) = tokio::task::spawn_blocking(move || {
crate::agent::models::prefetch_models_and_settings_blocking(
&endpoints_for_prefetch,
auth_for_prefetch.as_ref(),
fetch_auth_for_prefetch,
)
})
.await
.unwrap_or((None, None));
// A session-less leader that can still mint one (auth provider / devbox) will
// acquire a grok.com session post-readiness whose fleet policy governs
// external OTEL; see the background cold-mint below.
// Disk presence, not the is_xai-filtered no-mint result: an enterprise
// session has remote policy and must keep the gate closed.
let has_session = auth.is_some()
|| agent_config
.create_auth_manager()
.read_disk_auth()
.is_some();
let session_pending =
crate::agent::otel_gate::is_session_pending(has_session, &agent_config.grok_com_config);
if crate::agent::otel_gate::should_open_at_startup(crate::agent::otel_gate::StartupGate {
has_session,
has_api_key_env: crate::agent::auth_method::has_xai_api_key_env(),
session_pending,
remote_fetch_enabled: crate::util::config::resolve_remote_fetch_enabled(),
}) {
info!("Pure env-API-key leader; opening external-OTEL gate (no remote policy applies)");
crate::agent::otel_gate::open_at_startup();
}
// Process-wide image normalize cache: off by default, toggled here from
// `RemoteSettings.image_normalize_cache_enabled` once at startup.
let image_normalize_cache_enabled = remote_settings
.as_ref()
.and_then(|r| r.image_normalize_cache_enabled)
.unwrap_or(false);
crate::session::normalize_cache::NormalizeCache::global()
.set_enabled(image_normalize_cache_enabled);
tracing::debug!(
enabled = image_normalize_cache_enabled,
"image normalize cache toggle resolved from remote settings"
);
// Non-blocking boot: nothing is prefetched; the catalog and remote settings
// stream in after readiness via the background refreshes below.
let prefetched_models: Option<_> = None;
let remote_settings: Option<_> = None;
// ── Phase 7: Signal readiness ─────────────────────────────────────────────
//
// Unblocks ACP forwarding inside the IPC server. From this point on, client
// ACP messages are forwarded to the agent as normal.
let _ = ready_tx.send(true);
info!("Leader ready: auth and model prefetch complete, ACP forwarding enabled");
info!(
"Leader ready: local-only boot (model/settings refresh runs in background), ACP forwarding enabled"
);
// ── Phase 8: LocalSet — agent, bridges, relay, config watcher ────────────
let local_set = tokio::task::LocalSet::new();
let remote_settings_for_reloader = remote_settings.clone();
let mut agent_config_for_spawn = agent_config.clone();
agent_config_for_spawn.remote_settings = remote_settings;
crate::util::config::sync_campaign_fields(&mut agent_config_for_spawn);
@ -1260,19 +1275,24 @@ pub async fn run_leader(
// process so a refresh can't straddle a suspend.
shared_auth_manager.start_system_power_listener();
// Resolved from startup auth here; when this is `None` (leader booted
// without auth) the relay is NOT permanently off — the config-update loop
// arms it later via `DeferredRelayArm` when the watcher hot-reloads a
// relay-eligible token.
// The refresher lands on `shared_auth_manager` during `MvpAgent`
// construction below; a relay 401 in the window before that surfaces as
// a transient recovery failure and is retried, not a dead end.
let relay_config: Option<RelayConfig> =
relay_config_for_session(auth.as_ref(), &agent_config, &shared_auth_manager);
// Seed the startup-resolved session into the shared manager so per-request
// `auth()` and relay eligibility read it as the single source.
if let Some(session) = auth.as_ref()
&& should_seed_shared_session(shared_auth_manager.current_or_expired().as_ref(), session)
{
shared_auth_manager.hot_swap(session.clone());
}
// Relay start policy from startup auth; `None` (session-less boot) is not
// permanent — the background cold-mint's auth.json write drives the
// config-update loop to arm the relay via `DeferredRelayArm`.
let relay_config = relay_config_for_session(auth.as_ref(), &agent_config, &shared_auth_manager);
// Same manager as the leader, so the exposure never writes auth.json itself.
workspace_control.set_auth_manager(shared_auth_manager.clone());
let auth_manager_for_agent = shared_auth_manager.clone();
let auth_manager_for_config = shared_auth_manager;
let auth_manager_for_config = shared_auth_manager.clone();
let auth_manager_for_mint = shared_auth_manager.clone();
// Restore managed policy right before bootstrap reads it (no stale window after the long auth/prefetch phase).
crate::managed_config::ensure_managed_policy_present(&auth_manager_for_agent).await;
@ -1283,6 +1303,9 @@ pub async fn run_leader(
prefetched_models,
)
.unwrap_or_else(exit_on_config_error);
shared_models_manager.spawn_background_refresh();
let models_manager_for_agent = shared_models_manager.clone();
let models_manager_for_config = shared_models_manager;
@ -1422,6 +1445,33 @@ pub async fn run_leader(
}
});
// Re-run the minter off the readiness path: the startup attempt is
// no-mint, so a mintable leader (auth provider / devbox) acquires
// its session here. Runs on the LocalSet (the external-provider
// flow is `!Send`); on success `mint_session_noninteractive`
// persists to auth.json, which the config-update loop below picks up
// to heal `auth()` and arm the relay via `DeferredRelayArm`.
if session_pending {
let mint_auth_manager = auth_manager_for_mint;
let mint_cancel = cancel_clone.clone();
tokio::task::spawn_local(async move {
tokio::select! {
biased;
_ = mint_cancel.cancelled() => {}
minted = crate::auth::mint_session_noninteractive(&mint_auth_manager)
=> match minted {
Some(session) => info!(
is_xai = session.is_xai_auth(),
"background cold-mint acquired a session post-readiness"
),
None => warn!(
"background cold-mint found no session; leader remains session-less"
),
},
}
});
}
// Start (or arm) the grok.com relay. Eager by default — a bare
// `grok agent leader` (devbox / systemd) has no local IPC clients
// and receives remote prompts *through* the relay, so it must
@ -1443,11 +1493,10 @@ pub async fn run_leader(
cancel_clone.clone(),
);
} else {
// No relay-eligible auth at startup (BYOK / local-only — or a
// devbox whose initial mint failed transiently). Don't decide
// "relay off" forever: park the parts so the config-update
// loop below arms the relay when the watcher hot-reloads a
// relay-eligible token. See `DeferredRelayArm`.
// No relay-eligible auth at startup (BYOK / local-only, or a
// devbox whose initial mint is still pending). Park the parts so
// the config-update loop arms the relay once the background
// cold-mint (or a re-login) writes a relay-eligible token.
info!(
"Relay not started: no grok.com session token \
(BYOK / local-only leader); will arm if an eligible \
@ -1555,7 +1604,7 @@ pub async fn run_leader(
initial_auth_key_hash,
initial_config,
auth_scope,
remote_settings_for_reloader,
None, // settings stream in after readiness via background refresh
config_update_tx,
agent_config.cli_experimental_memory,
agent_config.cli_no_memory,
@ -1872,7 +1921,7 @@ mod tests {
));
}
// ===== spawn_leader_relay start-policy tests =====
// ===== relay supervisor start-invariant tests =====
/// Mock relay WS server: counts accepted WebSocket connections and holds
/// each open so the relay loop doesn't immediately reconnect.
@ -1900,8 +1949,8 @@ mod tests {
(addr, count)
}
/// Relay config pointing at the mock server, built through the only
/// constructor (`for_session`) with a relay-eligible x.ai OIDC session.
/// A `RelayConfig` built via the production constructor (`for_session`) with
/// a relay-eligible x.ai OIDC session.
fn test_relay_config(addr: std::net::SocketAddr) -> crate::agent::relay::RelayConfig {
let auth = GrokAuth {
auth_mode: AuthMode::Oidc,
@ -1917,6 +1966,105 @@ mod tests {
.expect("x.ai OIDC session must be relay-eligible")
}
/// The external-OTEL gate opens at startup only for a pure env-API-key leader:
/// env key set, no session, no pending mint. Any session (resolved of any
/// credential type, or about to be minted) makes it wait for the fetch.
#[test]
fn otel_gate_opens_only_for_pure_env_api_key_leader() {
use crate::agent::otel_gate::{StartupGate, should_open_at_startup};
let opens = |has_session, has_api_key_env, session_pending| {
should_open_at_startup(StartupGate {
has_session,
has_api_key_env,
session_pending,
remote_fetch_enabled: true,
})
};
// (has_session, has_api_key_env, session_pending)
assert!(opens(false, true, false), "pure env API key → opens");
assert!(!opens(true, true, false), "any resolved session → waits");
assert!(!opens(true, false, false), "session, no env key → waits");
assert!(
!opens(false, true, true),
"pending mint → session coming, waits"
);
assert!(
!opens(false, false, false),
"no env key, no session → waits"
);
}
/// The embedded startup gate (every pager `--no-leader` / fallback path) must be
/// fail-closed by construction: a session user stays closed until the agent
/// resolves settings, even when an env API key is also present (the key must
/// not bypass the session's remote policy). The pure env-API-key open path
/// is covered by `otel_gate_opens_only_for_pure_env_api_key_leader`, since
/// `is_session_pending` is environment-dependent (true in a devbox/CI pod).
#[test]
#[serial_test::serial]
fn embedded_otel_gate_keeps_a_session_user_fail_closed() {
use crate::agent::auth_method::{LEGACY_XAI_API_KEY_ENV_VAR, XAI_API_KEY_ENV_VAR};
use xai_grok_telemetry::external::{
is_settings_gate_open, mark_external_otel_settings_resolved,
};
unsafe fn set_or_clear(key: &str, value: Option<std::ffi::OsString>) {
match value {
Some(v) => unsafe { std::env::set_var(key, v) },
None => unsafe { std::env::remove_var(key) },
}
}
/// Restores the api-key env and reopens the gate on drop so no state leaks.
struct Restore {
key: Option<std::ffi::OsString>,
legacy: Option<std::ffi::OsString>,
}
impl Drop for Restore {
fn drop(&mut self) {
// SAFETY: serialized by `#[serial]`.
unsafe {
set_or_clear(XAI_API_KEY_ENV_VAR, self.key.take());
set_or_clear(LEGACY_XAI_API_KEY_ENV_VAR, self.legacy.take());
}
mark_external_otel_settings_resolved();
}
}
let _restore = Restore {
key: std::env::var_os(XAI_API_KEY_ENV_VAR),
legacy: std::env::var_os(LEGACY_XAI_API_KEY_ENV_VAR),
};
let cfg = GrokComConfig::default();
// SAFETY: serialized by `#[serial]`.
unsafe {
std::env::set_var(XAI_API_KEY_ENV_VAR, "test-key");
std::env::remove_var(LEGACY_XAI_API_KEY_ENV_VAR);
}
let session = GrokAuth {
// Far-future expiry so `current()` accepts it regardless of clock skew
// or a leaked `GROK_AUTH_EARLY_INVALIDATION_SECS` buffer; the gate reads
// session presence via `current()`, which filters expired tokens.
expires_at: chrono::DateTime::from_timestamp(9_999_999_999, 0),
auth_mode: AuthMode::Oidc,
oidc_issuer: Some(crate::auth::XAI_OAUTH2_ISSUER.to_string()),
..GrokAuth::test_default()
};
let with_session = {
let dir = tempfile::tempdir().unwrap();
let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
am.hot_swap(session);
am
};
apply_otel_config(&with_session, &cfg);
assert!(
!is_settings_gate_open(),
"a session user must boot fail-closed even with an env key set"
);
}
/// Wait until at least one relay connection is accepted, or panic.
async fn wait_for_connection(count: &Arc<AtomicU32>, context: &str) {
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
@ -2098,6 +2246,89 @@ mod tests {
cancel.cancel();
}
/// End-to-end for the merge reconciliation: a background cold-mint persists
/// a relay-eligible session to auth.json, the config watcher emits
/// `ConfigUpdate::Auth`, and that arms the deferred relay.
#[tokio::test]
async fn cold_mint_auth_write_arms_deferred_relay() {
use crate::config::reloader::{ConfigReloader, ConfigUpdate, hash_auth_key};
let (addr, _count) = spawn_mock_relay_server().await;
let grok_com_config = crate::auth::GrokComConfig {
grok_ws_url: format!("ws://{addr}"),
grok_ws_origin: format!("http://{addr}"),
..Default::default()
};
let tmp = tempfile::tempdir().unwrap();
let scope = "https://test.example.com".to_string();
let session = GrokAuth {
auth_mode: AuthMode::Oidc,
oidc_issuer: Some(crate::auth::XAI_OAUTH2_ISSUER.to_string()),
..GrokAuth::test_default()
};
let mut store = std::collections::BTreeMap::new();
store.insert(scope.clone(), session);
std::fs::write(
tmp.path().join("auth.json"),
serde_json::to_string_pretty(&store).unwrap(),
)
.unwrap();
let (tx, mut rx) = mpsc::unbounded_channel();
let mut reloader = ConfigReloader::new(
tmp.path().to_path_buf(),
hash_auth_key("sessionless-boot"),
toml::Value::Table(Default::default()),
scope,
None,
tx,
false,
false,
);
reloader.reload_auth().unwrap();
let ConfigUpdate::Auth(minted) = rx
.try_recv()
.expect("cold-mint auth.json write must emit ConfigUpdate::Auth")
else {
panic!("expected ConfigUpdate::Auth");
};
let auth_manager = Arc::new(AuthManager::new(tmp.path(), grok_com_config.clone()));
let (ws_to_agent_tx, _ws_to_agent_rx) = mpsc::unbounded_channel();
let agent_to_ws_tx: Rc<Mutex<Option<mpsc::UnboundedSender<String>>>> =
Rc::new(Mutex::new(None));
let agent_to_ws_tx_probe = agent_to_ws_tx.clone();
let (_demand_tx, demand_rx) = watch::channel(false);
let slot = Rc::new(std::cell::RefCell::new(None));
let cancel = CancellationToken::new();
let arm = DeferredRelayArm {
relay_on_demand: false,
relay_demand_rx: demand_rx,
ws_to_agent_tx,
agent_to_ws_tx,
cancel: cancel.clone(),
slot: slot.clone(),
grok_com_config,
alpha_test_key: None,
};
let local = tokio::task::LocalSet::new();
local
.run_until(async {
assert!(
arm.arm_if_eligible(&minted, &auth_manager).is_none(),
"a cold-minted relay-eligible session must arm the relay"
);
assert!(slot.borrow().is_some(), "relay handle must be parked");
assert!(
agent_to_ws_tx_probe.lock().is_some(),
"outbound relay sender must be installed"
);
})
.await;
cancel.cancel();
}
/// The watcher-injected internal reload requests must carry the ACP
/// wire-level `_` extension prefix. `agent-client-protocol`'s inbound
/// decoder routes non-built-in methods to `ext_method` only when

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

@ -3453,6 +3453,17 @@ pub fn apply_external_otel_remote_policy(settings: Option<&crate::util::config::
}
}
/// Seed free-function remote caches after writing `Config.remote_settings`.
///
/// Called from `init.rs` at boot and from the agent when backgrounded settings
/// arrive later, so every side effect here must be idempotent and safe to
/// re-apply. The emission-gate flip is owned by
/// [`crate::agent::otel_gate::OtelGate`], not here.
///
/// The `force_disable` write here is `Relaxed`; the synchronizing publish is
/// `OtelGate::apply_and_open`, which applies the same tighten-only policy and then
/// opens the gate with a `Release` swap. Removing that second application to
/// deduplicate would leave only the `Relaxed` store and reopen an ARM
/// visibility hole.
pub fn apply_remote_settings_side_effects(settings: Option<&crate::util::config::RemoteSettings>) {
if let Some(s) = settings {
let origin_trusted = crate::util::is_prod_cli_chat_proxy_url(
@ -3477,6 +3488,11 @@ pub fn apply_remote_settings_side_effects(settings: Option<&crate::util::config:
settings.and_then(|s| s.crash_handler_enabled),
);
apply_external_otel_remote_policy(settings);
let image_normalize_cache_enabled = settings
.and_then(|r| r.image_normalize_cache_enabled)
.unwrap_or(false);
crate::session::normalize_cache::NormalizeCache::global()
.set_enabled(image_normalize_cache_enabled);
}
/// Read `env.<key>` from Claude-compat `managed_settings.json`. `Some(true)`
/// indicates a force-off signal from a Mac-MDM-style admin policy.

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

@ -119,15 +119,15 @@ fn resolve_config(cfg: &AgentConfig, auth_manager: &AuthManager) -> AgentConfig
crate::util::config::sync_campaign_fields(&mut cfg);
// env var > remote settings > Local. Skip remote settings for Generic (grok -p, subagents).
let has_xai_auth = auth_manager.current().is_some_and(|a| a.is_xai_auth());
if cfg.storage_mode == StorageMode::Local
&& cfg.mode != crate::agent::config::AgentMode::Generic
{
cfg.storage_mode = StorageMode::resolve(None, cfg.remote_settings.as_ref());
cfg.storage_mode =
StorageMode::from_remote_gated(cfg.remote_settings.as_ref(), has_xai_auth);
}
// Writeback talks to the code backend; requires grok.com auth.
if cfg.storage_mode == StorageMode::Writeback
&& !auth_manager.current().is_some_and(|a| a.is_xai_auth())
{
// A CLI/env-set Writeback still requires grok.com auth.
if cfg.storage_mode == StorageMode::Writeback && !has_xai_auth {
tracing::info!("Writeback is disabled: requires auth with grok.com");
cfg.storage_mode = StorageMode::Local;
}
@ -165,9 +165,9 @@ fn init_process(cfg: &AgentConfig, auth_manager: &AuthManager) {
crate::extensions::marketplace::purge_default_skills_installs(&grok_home);
// Auto-register is gated (default off; env/remote settings enables). Kept out
// of built-in extraction so the gate can read the resolved
// remote_settings, which resolve_config has populated by now.
// At boot remote_settings may still be None (fetches are backgrounded),
// so only an env opt-in fires here; the gate is re-evaluated once
// settings arrive (see `MvpAgent::reapply_official_marketplace`).
if cfg.resolve_official_marketplace_auto_register().value {
crate::extensions::marketplace::ensure_official_marketplace_source(&grok_home);
}

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

@ -12,6 +12,7 @@ pub mod init;
pub mod model_providers;
pub mod models;
pub mod mvp_agent;
pub(crate) mod otel_gate;
pub(crate) mod proxy;
pub mod relay;
pub(crate) mod restore_code;

3117
crates/codegen/xai-grok-shell/src/agent/models.rs
View file

File diff suppressed because it is too large Load diff

211
crates/codegen/xai-grok-shell/src/agent/models/cache.rs Normal file
View file

@ -0,0 +1,211 @@
use super::*;
// ── Disk cache ──────────────────────────────────────────────────────────────
pub(crate) const MODELS_CACHE_FILE: &str = "models_cache.json";
pub(crate) const CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(300);
#[derive(serde::Serialize, serde::Deserialize)]
pub(crate) struct ModelsCache {
pub(crate) fetched_at: DateTime<Utc>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) grok_version: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) auth_method: Option<CacheAuthMethod>,
/// Models-list URL this catalog was fetched from; compared on load so a cache written against another backend is a miss.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) origin: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) etag: Option<String>,
pub(crate) models: IndexMap<String, ModelEntry>,
}
impl ModelsCache {
fn is_fresh(&self, ttl: std::time::Duration) -> bool {
let Ok(ttl) = ChronoDuration::from_std(ttl) else {
return false;
};
let age = Utc::now().signed_duration_since(self.fetched_at);
age >= ChronoDuration::zero() && age < ttl
}
}
pub(crate) struct CacheResult {
pub(crate) models: IndexMap<String, ModelEntry>,
pub(crate) etag: Option<String>,
}
pub(crate) struct ModelsCacheManager {
pub(crate) path: std::path::PathBuf,
pub(crate) ttl: std::time::Duration,
}
impl ModelsCacheManager {
pub(crate) fn new() -> Self {
Self {
path: crate::util::grok_home::grok_home().join(MODELS_CACHE_FILE),
ttl: CACHE_TTL,
}
}
pub(crate) fn load_fresh(
&self,
expected_auth: &CacheAuthMethod,
expected_origin: &str,
) -> Option<CacheResult> {
let data = std::fs::read(&self.path).ok()?;
let cache: ModelsCache = serde_json::from_slice(&data).ok()?;
if cache.grok_version.as_deref() != Some(xai_grok_version::VERSION) {
tracing::debug!("models cache version mismatch");
return None;
}
if cache.auth_method.as_ref() != Some(expected_auth) {
tracing::debug!("models cache auth method mismatch");
return None;
}
if cache.origin.as_deref() != Some(expected_origin) {
tracing::debug!(
cached = ?cache.origin,
expected = expected_origin,
"models cache origin mismatch"
);
return None;
}
if !cache.is_fresh(self.ttl) {
tracing::debug!("models cache is stale");
return None;
}
tracing::debug!(count = cache.models.len(), "loaded models from disk cache");
Some(CacheResult {
models: cache.models,
etag: cache.etag,
})
}
pub(crate) fn persist(
&self,
models: &IndexMap<String, ModelEntry>,
etag: Option<&str>,
auth_method: CacheAuthMethod,
origin: &str,
) {
let cache = ModelsCache {
fetched_at: Utc::now(),
grok_version: Some(xai_grok_version::VERSION.to_string()),
auth_method: Some(auth_method),
origin: Some(origin.to_string()),
etag: etag.map(|s| s.to_string()),
models: models.clone(),
};
self.atomic_write(&cache);
}
pub(crate) async fn renew_ttl(&self, expected_auth: &CacheAuthMethod, expected_origin: &str) {
let data = match tokio::fs::read(&self.path).await {
Ok(data) => data,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return,
Err(e) => {
tracing::warn!(error = %e, "models cache TTL renewal: read failed");
return;
}
};
let Ok(mut cache) = serde_json::from_slice::<ModelsCache>(&data) else {
return;
};
if cache.auth_method.as_ref() != Some(expected_auth) {
tracing::debug!("models cache TTL renewal skipped: auth method mismatch");
return;
}
if cache.origin.as_deref() != Some(expected_origin) {
tracing::debug!("models cache TTL renewal skipped: origin mismatch");
return;
}
cache.fetched_at = Utc::now();
self.atomic_write_async(&cache).await;
tracing::debug!("models cache TTL renewed");
}
pub(crate) fn invalidate(&self) {
match std::fs::remove_file(&self.path) {
Ok(()) => tracing::info!("models disk cache invalidated"),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
Err(e) => tracing::warn!(error = %e, "failed to invalidate models disk cache"),
}
}
/// Per-writer temp path: `~/.grok` is shared across concurrent CLI
fn unique_tmp_path(&self) -> std::path::PathBuf {
static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
let n = SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
self.path
.with_extension(format!("json.tmp.{}.{n}", std::process::id()))
}
/// Best-effort removal of temp files a crash left in the write→rename window; only sweeps entries older than the TTL.
fn sweep_stale_tmp(&self) {
let (Some(parent), Some(stem)) = (
self.path.parent(),
self.path.file_name().and_then(|s| s.to_str()),
) else {
return;
};
let prefix = format!("{stem}.tmp.");
let Ok(entries) = std::fs::read_dir(parent) else {
return;
};
let now = std::time::SystemTime::now();
for entry in entries.flatten() {
let Some(name) = entry.file_name().to_str().map(str::to_owned) else {
continue;
};
if !name.starts_with(&prefix) {
continue;
}
let is_stale = entry
.metadata()
.and_then(|m| m.modified())
.ok()
.and_then(|t| now.duration_since(t).ok())
.is_some_and(|age| age > self.ttl);
if is_stale {
let _ = std::fs::remove_file(entry.path());
}
}
}
pub(crate) fn atomic_write(&self, cache: &ModelsCache) {
if let Some(parent) = self.path.parent() {
let _ = std::fs::create_dir_all(parent);
}
self.sweep_stale_tmp();
let Ok(json) = serde_json::to_vec_pretty(cache) else {
return;
};
let tmp = self.unique_tmp_path();
if std::fs::write(&tmp, &json).is_ok() {
if std::fs::rename(&tmp, &self.path).is_err() {
let _ = std::fs::remove_file(&tmp);
}
} else {
let _ = std::fs::remove_file(&tmp);
}
}
pub(crate) async fn atomic_write_async(&self, cache: &ModelsCache) {
if let Some(parent) = self.path.parent() {
let _ = tokio::fs::create_dir_all(parent).await;
}
self.sweep_stale_tmp();
let Ok(json) = serde_json::to_vec_pretty(cache) else {
return;
};
let tmp = self.unique_tmp_path();
if tokio::fs::write(&tmp, &json).await.is_ok() {
if tokio::fs::rename(&tmp, &self.path).await.is_err() {
let _ = tokio::fs::remove_file(&tmp).await;
}
} else {
let _ = tokio::fs::remove_file(&tmp).await;
}
}
}

41
crates/codegen/xai-grok-shell/src/agent/models/endpoint.rs Normal file
View file

@ -0,0 +1,41 @@
use super::*;
/// Boxed future returned by [`ModelsEndpoint::fetch_models`].
pub(crate) type ModelsFetchFuture =
Pin<Box<dyn Future<Output = Option<IndexMap<String, ModelEntry>>> + Send>>;
/// Injectable `/v1/models` transport; tests inject a fake.
pub(crate) trait ModelsEndpoint: Send + Sync {
fn fetch_models(
&self,
endpoints: config::EndpointsConfig,
auth: Option<GrokAuth>,
fetch_auth: ModelFetchAuth,
) -> ModelsFetchFuture;
}
/// Default transport: the real `/v1/models` fetch.
pub(crate) struct HttpModelsEndpoint;
impl ModelsEndpoint for HttpModelsEndpoint {
fn fetch_models(
&self,
endpoints: config::EndpointsConfig,
auth: Option<GrokAuth>,
fetch_auth: ModelFetchAuth,
) -> ModelsFetchFuture {
Box::pin(fetch_models_async(endpoints, auth, fetch_auth))
}
}
pub(crate) async fn fetch_models_async(
endpoints: config::EndpointsConfig,
auth: Option<GrokAuth>,
fetch_auth: ModelFetchAuth,
) -> Option<IndexMap<String, ModelEntry>> {
tokio::task::spawn_blocking(move || {
prefetch_models_blocking(&endpoints, auth.as_ref(), fetch_auth)
})
.await
.unwrap_or(None)
}

249
crates/codegen/xai-grok-shell/src/agent/models/fetch.rs Normal file
View file

@ -0,0 +1,249 @@
use super::*;
// ── Fetch ───────────────────────────────────────────────────────────────────
/// Build the prefetched model map from a flat list of entries.
pub(crate) fn build_prefetched_map(
models: Vec<config::ModelEntryConfig>,
api_base_url_override: Option<String>,
) -> IndexMap<String, ModelEntry> {
let mut map: IndexMap<String, ModelEntry> = IndexMap::with_capacity(models.len());
for m in models {
let key = m.id.clone().unwrap_or_else(|| m.model.clone());
let info = config::ModelInfo::from_config(&m);
let entry = ModelEntry {
info,
api_key: None,
env_key: None,
auth_provider: None,
api_base_url: m.api_base_url.clone().or(api_base_url_override.clone()),
};
map.insert(key, entry);
}
map
}
/// Fetch remote models. Checks disk cache first; persists after fetch.
pub(crate) fn prefetch_models_blocking(
endpoints: &config::EndpointsConfig,
auth: Option<&GrokAuth>,
fetch_auth: ModelFetchAuth,
) -> Option<IndexMap<String, ModelEntry>> {
prefetch_models_blocking_gated(
endpoints,
auth,
fetch_auth,
crate::util::config::resolve_remote_fetch_enabled(),
)
}
/// Blocking models + `/v1/settings` prefetch pair, shared by the early
pub(crate) fn prefetch_models_and_settings_blocking(
endpoints: &config::EndpointsConfig,
auth: Option<&GrokAuth>,
fetch_auth: ModelFetchAuth,
) -> (
Option<IndexMap<String, ModelEntry>>,
Option<crate::util::config::RemoteSettings>,
) {
let remote_fetch_enabled = crate::util::config::resolve_remote_fetch_enabled();
let models = prefetch_models_blocking_gated(endpoints, auth, fetch_auth, remote_fetch_enabled);
let settings = match auth {
Some(auth) if remote_fetch_enabled => {
let _timer = crate::instrumentation_timer!("startup.early_settings_fetch");
crate::remote::fetch_settings_blocking(
&endpoints.proxy_url(),
auth,
endpoints.alpha_test_key.as_deref(),
)
.into_option()
}
_ => None,
};
(models, settings)
}
/// `remote_fetch_enabled` is a parameter so the pair helper above resolves the knob once for both halves.
fn prefetch_models_blocking_gated(
endpoints: &config::EndpointsConfig,
auth: Option<&GrokAuth>,
fetch_auth: ModelFetchAuth,
remote_fetch_enabled: bool,
) -> Option<IndexMap<String, ModelEntry>> {
let cache_auth = fetch_auth.cache_auth_method();
let cache_origin = crate::remote::models_list_url(endpoints, fetch_auth);
let cache = ModelsCacheManager::new();
if let Some(cached) = cache.load_fresh(&cache_auth, &cache_origin) {
return Some(cached.models);
}
if !remote_fetch_enabled {
tracing::info!("models fetch skipped: remote_fetch disabled");
return None;
}
let _timer = crate::instrumentation_timer!("startup.fetch_models_blocking");
match fetch_models_blocking(endpoints, auth, fetch_auth) {
Ok(FetchModelsResult { models, etag }) if !models.is_empty() => {
let api_base_url_override = match fetch_auth {
ModelFetchAuth::ApiKey => Some(endpoints.xai_api_base_url.clone()),
_ => None,
};
let map = build_prefetched_map(models, api_base_url_override);
tracing::info!(count = map.len(), etag = ?etag, "Prefetched models");
cache.persist(&map, etag.as_deref(), cache_auth, &cache_origin);
Some(map)
}
Ok(FetchModelsResult { .. }) => {
tracing::warn!("Models endpoint returned empty list");
None
}
Err(e) => {
tracing::warn!("Failed to fetch models: {:?}", e);
None
}
}
}
/// Startup prefetch result: models + remote settings.
pub struct EarlyPrefetchResult {
pub models: Option<IndexMap<String, ModelEntry>>,
pub settings: Option<crate::util::config::RemoteSettings>,
}
/// Handle for a startup prefetch thread.
pub type EarlyPrefetchHandle = std::thread::JoinHandle<EarlyPrefetchResult>;
pub(crate) struct PrefetchEnv {
pub(crate) auth: Option<GrokAuth>,
pub(crate) endpoints: config::EndpointsConfig,
pub(crate) model_fetch_auth: ModelFetchAuth,
}
/// Effective startup endpoints, resolved config-aware (not env-only) so the prefetch can't leak the bearer to api.x.ai.
fn resolve_startup_endpoints() -> config::EndpointsConfig {
let mut endpoints = config::EndpointsConfig::from_effective_config();
if endpoints.deployment_key.is_none() {
endpoints.deployment_key = crate::managed_config::resolve_deployment_key();
}
endpoints
}
/// Decision core of the startup prefetch gate, split from the config loading
pub(crate) fn resolve_prefetch_env_from_parts(
auth: Option<GrokAuth>,
endpoints: config::EndpointsConfig,
remote_fetch_enabled: bool,
) -> Option<PrefetchEnv> {
if !remote_fetch_enabled {
tracing::info!("startup model/settings prefetch skipped: remote_fetch disabled");
return None;
}
let model_fetch_auth = ModelFetchAuth::resolve(&endpoints, auth.is_some());
if auth.is_none()
&& !endpoints.has_custom_endpoint()
&& model_fetch_auth == ModelFetchAuth::Session
{
return None;
}
Some(PrefetchEnv {
auth,
endpoints,
model_fetch_auth,
})
}
/// Start model + settings prefetch on a background thread using pre-resolved auth.
pub fn start_early_prefetch_with_auth(auth: Option<GrokAuth>) -> Option<EarlyPrefetchHandle> {
start_early_prefetch_with_auth_gated(auth, true)
}
/// `sync_managed = false` skips the managed-config sync, so a remote kill-switch
/// can apply on cold start before the fail-closed managed-policy gate without an
/// online sync healing a tampered on-disk policy first.
fn start_early_prefetch_with_auth_gated(
auth: Option<GrokAuth>,
sync_managed: bool,
) -> Option<EarlyPrefetchHandle> {
let _timer = crate::instrumentation_timer!("startup.early_prefetch_launch");
let endpoints = resolve_startup_endpoints();
if sync_managed {
spawn_managed_config_sync_if_stale(&endpoints);
}
let env = resolve_prefetch_env_from_parts(
auth,
endpoints,
crate::util::config::resolve_remote_fetch_enabled(),
)?;
Some(spawn_prefetch_thread(env))
}
/// Start model + settings prefetch on a background thread.
pub fn start_early_prefetch(grok_com_config: Option<GrokComConfig>) -> Option<EarlyPrefetchHandle> {
start_early_prefetch_impl(grok_com_config, true)
}
/// Prefetch models + remote settings only — no managed-config sync. Used before
/// the managed-policy gate (see `start_early_prefetch_with_auth_gated`).
pub fn start_early_prefetch_settings_only(
grok_com_config: Option<GrokComConfig>,
) -> Option<EarlyPrefetchHandle> {
start_early_prefetch_impl(grok_com_config, false)
}
fn start_early_prefetch_impl(
grok_com_config: Option<GrokComConfig>,
sync_managed: bool,
) -> Option<EarlyPrefetchHandle> {
let grok_home = crate::util::grok_home::grok_home();
let auth = AuthManager::new(&grok_home, grok_com_config.unwrap_or_default()).current();
start_early_prefetch_with_auth_gated(auth, sync_managed)
}
fn spawn_prefetch_thread(env: PrefetchEnv) -> EarlyPrefetchHandle {
std::thread::spawn(move || {
let mut timer = crate::instrumentation_timer!("startup.early_prefetch");
let proxy_endpoint = env.endpoints.proxy_url();
timer.with_field("endpoint", proxy_endpoint.as_str());
let (models, settings) = prefetch_models_and_settings_blocking(
&env.endpoints,
env.auth.as_ref(),
env.model_fetch_auth,
);
EarlyPrefetchResult { models, settings }
})
}
/// Best-effort, bounded managed-config sync on a detached thread, off the readiness path (syncs at launch; the interval task covers steady state).
fn spawn_managed_config_sync_if_stale(endpoints: &config::EndpointsConfig) {
let should_sync = (endpoints.deployment_key.is_some()
|| crate::managed_config::has_active_team_auth())
&& crate::config::is_managed_config_stale_for(
&crate::managed_config::current_serving_identity(),
)
&& crate::managed_config::is_fetch_enabled();
if !should_sync {
return;
}
std::thread::spawn(|| {
let Ok(rt) = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
else {
return;
};
crate::managed_config::clear_orphan();
// tokio timer outside a runtime context panics ("no reactor running").
let _ = rt.block_on(async {
tokio::time::timeout(
crate::http::STARTUP_FETCH_TIMEOUT,
crate::managed_config::sync(),
)
.await
});
});
}

321
crates/codegen/xai-grok-shell/src/agent/models/resolution.rs Normal file
View file

@ -0,0 +1,321 @@
use super::*;
/// Map a model id (catalog key or routing slug) to its catalog key.
pub(crate) fn resolve_catalog_key(
models: &IndexMap<String, ModelEntry>,
id: &acp::ModelId,
) -> Option<acp::ModelId> {
let id_str = id.0.as_ref();
if models.contains_key(id_str) {
return Some(id.clone());
}
models
.iter()
.rev()
.find(|(_, entry)| entry.info.model == id_str)
.map(|(key, _)| acp::ModelId::new(key.clone()))
}
/// Catalog key for a persisted session model id, restricted to **selectable**
pub(crate) fn selectable_catalog_key_for_persisted(
models: &IndexMap<String, ModelEntry>,
available: &IndexMap<acp::ModelId, acp::ModelInfo>,
id: &acp::ModelId,
) -> Option<acp::ModelId> {
if available.contains_key(id) {
return Some(id.clone());
}
let id_str = id.0.as_ref();
if let Some((key, _)) = models.iter().rev().find(|(key, entry)| {
available.contains_key(&acp::ModelId::new((*key).clone())) && entry.info.model == id_str
}) {
return Some(acp::ModelId::new(key.clone()));
}
resolve_catalog_key(models, id).filter(|key| available.contains_key(key))
}
/// A "campaign-only" preferred flip: the default changed and either side's value
pub(crate) fn is_campaign_only_flip(
old_preferred: &Option<String>,
new_preferred: &Option<String>,
campaign_defaults: &std::collections::HashSet<String>,
) -> bool {
if new_preferred == old_preferred || new_preferred.is_none() {
return false;
}
new_preferred
.as_ref()
.is_some_and(|p| campaign_defaults.contains(p))
|| old_preferred
.as_ref()
.is_some_and(|p| campaign_defaults.contains(p))
}
/// Pick the default model: CLI > env > config > remote-settings hint, falling
pub(crate) fn resolve_default_model(
cfg: &config::Config,
catalog: &IndexMap<String, ModelEntry>,
is_session_auth: bool,
) -> (String, ModelEntry, config::ConfigSource) {
let visible: IndexMap<String, ModelEntry> = catalog
.iter()
.filter(|(_, e)| e.info.visible_for_auth(is_session_auth) && e.info.user_selectable)
.map(|(k, v)| (k.clone(), v.clone()))
.collect();
let model_pref = config::resolve_string_flag(
cfg.default_model_override.as_deref(),
"GROK_DEFAULT_MODEL",
cfg.models.default.as_deref(),
cfg.remote_settings
.as_ref()
.and_then(|rs| rs.default_model.as_deref()),
);
let first_or_fallback = || -> (String, ModelEntry) {
if let Some((key, first)) = visible.first() {
return (key.clone(), first.clone());
}
if let Some((key, entry)) = catalog.iter().find(|(_, e)| e.info.user_selectable) {
tracing::warn!("no auth-visible selectable model; using first selectable entry");
return (key.clone(), entry.clone());
}
tracing::warn!("no selectable models; falling back to bundled default (pre-catalog)");
let default_id = crate::models::default_model().to_string();
let mut entry = ModelEntry::fallback(&default_id, &cfg.endpoints);
entry.info.user_selectable = match ModelGlobSet::compile(cfg.models.allowed_models.as_ref())
{
Ok(None) => true,
Ok(Some(set)) => set.matches(&default_id, &default_id),
Err(_) => false,
};
(default_id, entry)
};
match &model_pref {
None => {
let (key, first) = first_or_fallback();
(key, first, config::ConfigSource::Default)
}
Some(pref) => {
let found = visible
.get_key_value(&pref.value)
.or_else(|| visible.iter().find(|(_, m)| m.model == pref.value));
if let Some((key, entry)) = found {
(key.clone(), entry.clone(), pref.source)
} else {
let is_explicit = matches!(
pref.source,
config::ConfigSource::Cli
| config::ConfigSource::Env
| config::ConfigSource::Config
);
if is_explicit {
tracing::warn!(
model_id = %pref.value, source = %pref.source,
"preferred model not in available models, falling back"
);
} else {
tracing::debug!(
model_id = %pref.value, source = %pref.source,
"remote default_model not in available models, skipping"
);
}
let campaign_pref_missing = cfg.models.default_is_campaign_driven
&& matches!(pref.source, config::ConfigSource::Config);
if campaign_pref_missing
&& let Some(prev) = cfg
.models
.pre_campaign_default
.as_deref()
.filter(|s| !s.is_empty())
&& let Some((key, entry)) = visible
.get_key_value(prev)
.or_else(|| visible.iter().find(|(_, m)| m.model == prev))
{
tracing::info!(
unavailable = %pref.value, fallback = %prev,
"campaign-driven default unavailable in catalog; recovering the pre-campaign default"
);
return (key.clone(), entry.clone(), config::ConfigSource::Config);
}
let (key, first) = first_or_fallback();
(key, first, config::ConfigSource::Default)
}
}
}
}
/// Filter hidden and auth-gated entries out of `catalog` and convert to ACP wire format.
pub fn available_models(
catalog: &IndexMap<String, ModelEntry>,
is_session_auth: bool,
) -> IndexMap<acp::ModelId, acp::ModelInfo> {
let visible: IndexMap<String, ModelEntry> = catalog
.iter()
.filter(|(_, e)| e.info.visible_for_auth(is_session_auth))
.map(|(k, v)| (k.clone(), v.clone()))
.collect();
config::to_acp_model_info(&visible)
}
/// Compiled glob matcher shared by `allowed_models`, `disabled_models`, and `hidden_models` (matched against catalog key or model id).
pub(crate) struct ModelGlobSet(GlobSet);
impl ModelGlobSet {
/// Compile a filter list (`Ok(None)` for `None`/empty). Fails **closed**: an invalid pattern returns `Err` listing every bad one.
pub(crate) fn compile(patterns: Option<&Vec<String>>) -> Result<Option<Self>, Vec<String>> {
let patterns = match patterns {
Some(p) if !p.is_empty() => p,
_ => return Ok(None),
};
let mut builder = GlobSetBuilder::new();
let mut invalid = Vec::new();
for pat in patterns {
match Glob::new(pat) {
Ok(glob) => {
builder.add(glob);
}
Err(_) => invalid.push(pat.clone()),
}
}
if !invalid.is_empty() {
return Err(invalid);
}
builder
.build()
.map(|set| Some(Self(set)))
.map_err(|e| vec![e.to_string()])
}
fn matches(&self, key: &str, model: &str) -> bool {
self.0.is_match(key) || self.0.is_match(model)
}
}
/// Single source of truth for the catalog. Applies, in order: `disabled_models`
pub fn resolve_model_catalog(
cfg: &config::Config,
prefetched: Option<IndexMap<String, ModelEntry>>,
) -> IndexMap<String, ModelEntry> {
let mut catalog: IndexMap<String, ModelEntry> = config::resolve_model_list(cfg, prefetched);
if let Ok(Some(disabled)) = ModelGlobSet::compile(cfg.models.disabled_models.as_ref()) {
let before = catalog.len();
catalog.retain(|key, entry| !disabled.matches(key, &entry.model));
let removed = before - catalog.len();
if removed > 0 {
tracing::info!(count = removed, "disabled_models: removed from catalog");
}
}
match ModelGlobSet::compile(cfg.models.allowed_models.as_ref()) {
Ok(None) => {
for entry in catalog.values_mut() {
entry.info.user_selectable = true;
}
}
Ok(Some(allowed)) => {
for (key, entry) in catalog.iter_mut() {
entry.info.user_selectable = allowed.matches(key, &entry.model);
}
}
Err(bad) => {
tracing::error!(patterns = ?bad, "allowed_models: invalid glob(s); marking nothing selectable");
for entry in catalog.values_mut() {
entry.info.user_selectable = false;
}
}
}
if let Ok(Some(hidden)) = ModelGlobSet::compile(cfg.models.hidden_models.as_ref()) {
for (key, entry) in catalog.iter_mut() {
if hidden.matches(key, &entry.model) {
entry.info.hidden = true;
}
}
}
if let Some(effort) = cfg.models.default_reasoning_effort
&& let Some(default_id) = cfg.models.default.as_deref()
&& let Some(entry) = catalog.get_mut(default_id)
&& entry.info.supports_reasoning_effort
{
entry.info.reasoning_effort = Some(effort);
}
if let Some(effort) = cfg.reasoning_effort_override {
for entry in catalog.values_mut() {
if model_offers_reasoning_effort(&entry.info, effort) {
entry.info.reasoning_effort = Some(effort);
}
}
}
catalog
}
/// Whether `effort` is a value this model will accept on the wire.
fn model_offers_reasoning_effort(info: &config::ModelInfo, effort: ReasoningEffort) -> bool {
if !info.supports_reasoning_effort {
return false;
}
if info.reasoning_efforts.is_empty() {
matches!(
effort,
ReasoningEffort::Low
| ReasoningEffort::Medium
| ReasoningEffort::High
| ReasoningEffort::Xhigh
)
} else {
info.reasoning_efforts.iter().any(|opt| opt.value == effort)
}
}
/// True when an active `allowed_models` allowlist leaves no selectable model.
pub(crate) fn allowlist_matches_nothing(
cfg: &config::Config,
catalog: &IndexMap<String, ModelEntry>,
) -> bool {
cfg.models
.allowed_models
.as_ref()
.is_some_and(|a| !a.is_empty())
&& !catalog.values().any(|e| e.info.user_selectable)
}
/// Reject an `allowed_models` allowlist that leaves no selectable model, or excludes an explicitly configured default; run only against a real catalog.
pub(crate) fn validate_selectable(
cfg: &config::Config,
catalog: &IndexMap<String, ModelEntry>,
) -> Result<(), String> {
let Some(allowed) = cfg.models.allowed_models.as_ref().filter(|a| !a.is_empty()) else {
return Ok(());
};
let patterns = allowed.join(", ");
if !catalog.values().any(|e| e.info.user_selectable) {
return Err(format!(
"None of your available models match allowed_models ({patterns}). \
Broaden the patterns or remove allowed_models, then try again."
));
}
for (src, id) in [
("default", cfg.models.default.as_deref()),
("-m flag", cfg.default_model_override.as_deref()),
] {
if let Some(id) = id
&& let Some(entry) = catalog
.get(id)
.or_else(|| catalog.values().find(|e| e.model == id))
&& !entry.info.user_selectable
{
return Err(format!(
"\"{id}\" (your {src}) isn't allowed by allowed_models ({patterns}). \
Add it to allowed_models, or set a different model."
));
}
}
Ok(())
}

1952
crates/codegen/xai-grok-shell/src/agent/models/tests.rs Normal file
View file

File diff suppressed because it is too large Load diff

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

@ -77,6 +77,9 @@ impl acp::Agent for MvpAgent {
tracing::debug!(target: "sampling_log", "Received initialize request");
xai_grok_telemetry::unified_log::info("agent initialized", None, None);
self.start_subagent_coordinator();
if self.cfg.borrow().remote_settings.is_none() {
self.spawn_settings_reapply();
}
let (auto_gc_policy, run_auto_gc) = {
let cfg = self.cfg.borrow();
let has_remote = cfg.remote_settings.is_some();
@ -315,7 +318,14 @@ impl acp::Agent for MvpAgent {
);
let mut has_cached_token = init_has_current;
if !init_has_current && init_is_expired {
let refreshed = self.auth_manager.auth().await.is_ok();
let refreshed = matches!(
tokio::time::timeout(
crate::http::STARTUP_AUTH_REFRESH_TIMEOUT,
self.auth_manager.auth(),
)
.await,
Ok(Ok(_))
);
if refreshed {
tracing::debug!(
auth_type = ?self.auth_type(),
@ -741,10 +751,9 @@ impl acp::Agent for MvpAgent {
.authenticate_after_cached_token_unavailable(arguments)
.await;
}
self.refresh_remote_settings(&auth).await;
self.emit_settings_update_notification();
self.enforce_grok_code_access(&auth).await;
self.maybe_sync_bundle_in_background(false);
let auth_for_settings = auth.clone();
{
let mut sampling_config = self.sampling_config.borrow_mut();
sampling_config.api_key = Some(auth.key);
@ -766,7 +775,7 @@ impl acp::Agent for MvpAgent {
auth_method: "cached_token".to_string(),
user_id: uid,
});
self.maybe_fetch_post_auth_settings().await;
self.spawn_post_auth_settings(auth_for_settings);
Ok(self.auth_response_with_meta())
}
auth_method::GROK_COM_METHOD_ID | auth_method::OIDC_METHOD_ID => {
@ -890,8 +899,6 @@ impl acp::Agent for MvpAgent {
);
}
self.auth_manager.hot_swap(auth.clone());
self.refresh_remote_settings(&auth).await;
self.emit_settings_update_notification();
self.enforce_grok_code_access(&auth).await;
self.maybe_sync_bundle_in_background(false);
tokio::task::spawn_local(
@ -912,7 +919,7 @@ impl acp::Agent for MvpAgent {
auth_method: arguments.method_id.0.as_ref().to_string(),
user_id: Some(auth.user_id.clone()),
});
self.maybe_fetch_post_auth_settings().await;
self.spawn_post_auth_settings(auth);
Ok(self.auth_response_with_meta())
}
_ => {
@ -941,9 +948,7 @@ impl acp::Agent for MvpAgent {
.data("initialize must be called before new_session")
})?;
self.seed_client_config_auth_if_available();
if let Ok(auth) = self.auth_manager.auth().await {
self.refresh_settings_and_reapply(&auth).await;
}
self.spawn_settings_reapply();
let cwd = AbsPathBuf::new(arguments.cwd.clone())
.map_err(|e| acp::Error::invalid_params().data(e.to_string()))?;
let remote_settings = self.cfg.borrow().remote_settings.clone();
@ -1025,7 +1030,29 @@ impl acp::Agent for MvpAgent {
let mut disallowed_custom: Option<String> = None;
let session_initial_model = chat_initial_model(is_chat_kind, custom_model_id);
let build_custom_model_id = if is_chat_kind { None } else { custom_model_id };
let campaign_nudge = if is_chat_kind {
None
} else {
crate::util::config::campaign_driven_models_default()
.filter(|c| {
build_custom_model_id.is_none()
|| build_custom_model_id == c.pre_campaign.as_deref()
|| build_custom_model_id == Some(c.value.as_str())
})
};
let campaign_nudged = campaign_nudge.is_some();
if let Some(c) = &campaign_nudge {
tracing::info!(
model = %c.value,
requested = ?custom_model_id,
"new_session: applying campaign-driven default model"
);
}
let build_custom_model_id: Option<String> = campaign_nudge
.map(|c| c.value)
.or_else(|| build_custom_model_id.map(str::to_owned));
let resolved_custom_model = build_custom_model_id
.as_deref()
.and_then(|custom_model| match self
.resolve_model_id(&acp::ModelId::new(custom_model))
{
@ -1043,7 +1070,9 @@ impl acp::Agent for MvpAgent {
requested_model = custom_model,
"Requested model not allowed by allowed_models; falling back to current default model"
);
disallowed_custom = Some(custom_model.to_string());
if !campaign_nudged {
disallowed_custom = Some(custom_model.to_string());
}
None
}
Err(_) => {
@ -1124,7 +1153,7 @@ impl acp::Agent for MvpAgent {
&session_info,
model_id,
summary_client,
self.storage_mode,
self.storage_mode.get(),
Some(self.auth_manager.clone()),
relay_sync,
Some(self.gateway.clone()),
@ -1463,7 +1492,7 @@ impl acp::Agent for MvpAgent {
let (persistence_info, persistence) = crate::session::persistence::load_light(
&session_info,
summary_client,
self.storage_mode,
self.storage_mode.get(),
Some(self.auth_manager.clone()),
backend.as_ref(),
relay_sync,

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

@ -715,7 +715,7 @@ impl MvpAgent {
}
/// `true` when the agent runs in writeback storage mode.
pub(crate) fn is_writeback_storage(&self) -> bool {
matches!(self.storage_mode, StorageMode::Writeback)
matches!(self.storage_mode.get(), StorageMode::Writeback)
}
/// Resolved cli-chat-proxy base for session features (via
/// `proxy_url`). Not for the deployment-config fetch.
@ -885,19 +885,178 @@ impl MvpAgent {
pub(crate) fn deployment_key(&self) -> Option<String> {
self.cfg.borrow().endpoints.deployment_key.clone()
}
/// Re-fetch remote settings and re-init the telemetry client.
/// Apply settings side effects + push `x.ai/settings/update` to clients.
/// Shared tail for every settings-arrival site.
pub(super) fn on_remote_settings_changed(&self) {
crate::agent::config::apply_remote_settings_side_effects(
self.cfg.borrow().remote_settings.as_ref(),
);
if let Some(identity) = self
.auth_manager
.current_or_expired()
.filter(|a| a.is_xai_auth())
.map(|a| a.user_id)
{
self.tier_allowed
.set(
super::settings_allow_access(
self.cfg.borrow().remote_settings.as_ref(),
),
);
*self.allow_access_resolved_for.borrow_mut() = Some(identity);
}
self.reapply_storage_mode();
self.reapply_official_marketplace();
{
let cfg_snapshot = self.cfg.borrow().clone();
if self.sessions.borrow().is_empty() {
self.models_manager.apply_config_reselecting_default(cfg_snapshot);
} else {
self.models_manager.apply_config(cfg_snapshot);
}
}
self.sync_collection_config_gate();
self.emit_settings_update_notification();
self.emit_announcements(AnnouncementsPushMode::IfChanged);
self.reconfigure_heap_profile_monitor();
}
/// Re-evaluates the official-marketplace auto-register gate now that
/// remote settings exist. `init_process` ran the same gate at boot without
/// them, so a settings-targeted (not env-set) team would otherwise never
/// register. Idempotent: a no-op once installed.
fn reapply_official_marketplace(&self) {
if self.cfg.borrow().resolve_official_marketplace_auto_register().value {
crate::extensions::marketplace::ensure_official_marketplace_source(
&crate::util::grok_home::grok_home(),
);
}
}
/// Upgrade storage mode from newly-arrived remote settings. Mirrors the
/// `resolve_config` gate: only upgrades from `Local`, writeback needs xai auth.
fn reapply_storage_mode(&self) {
if self.storage_mode.get() != StorageMode::Local {
return;
}
let resolved_mode = {
let cfg = self.cfg.borrow();
if cfg.mode == crate::agent::config::AgentMode::Generic {
return;
}
let has_xai_auth = self
.auth_manager
.current_or_expired()
.is_some_and(|a| a.is_xai_auth());
StorageMode::from_remote_gated(cfg.remote_settings.as_ref(), has_xai_auth)
};
if resolved_mode == self.storage_mode.get() {
return;
}
tracing::info!(?resolved_mode, "storage mode upgraded from remote settings");
self.storage_mode.set(resolved_mode);
if resolved_mode == StorageMode::Writeback {
for handle in self.sessions.borrow().values() {
let _ = handle
.persistence_tx
.send(crate::session::persistence::PersistenceMsg::UpgradeToWriteback {
auth_manager: self.auth_manager.clone(),
});
}
}
}
/// Run the blocking `/settings` fetch for `auth` off the runtime thread.
async fn fetch_settings(
&self,
auth: &crate::auth::GrokAuth,
) -> crate::remote::SettingsFetch {
let (base_url, alpha) = {
let cfg = self.cfg.borrow();
(cfg.endpoints.proxy_url(), cfg.endpoints.alpha_test_key.clone())
};
let auth = auth.clone();
match tokio::task::spawn_blocking(move || crate::remote::fetch_settings_blocking(
&base_url,
&auth,
alpha.as_deref(),
))
.await
{
Ok(outcome) => outcome,
Err(e) => {
tracing::warn!(error = %e, "settings fetch task panicked");
crate::remote::SettingsFetch::Retry
}
}
}
/// Fetch remote settings for `auth` and drive the external-OTEL gate from
/// the outcome. Re-closes the gate first only on an account switch, then
/// hands the outcome to [`OtelGate::resolve`], which returns the settings
/// only on a successful fetch for the still-live identity. Single seam for
/// both post-auth callers.
///
/// Called unconditionally from both auth handlers so that:
/// - First install / expired OIDC token: settings are fetched for
/// the first time (the early prefetch had no auth to use).
/// - Reauth / account switch: settings are refreshed to reflect
/// the new user's remote settings targeting attributes.
/// [`OtelGate::resolve`]: crate::agent::otel_gate::OtelGate::resolve
pub(super) async fn fetch_settings_resolving_gate(
&self,
auth: &crate::auth::GrokAuth,
) -> Option<crate::util::config::RemoteSettings> {
let identity = auth.user_id.clone();
self.otel_gate.rearm_on_switch(&identity);
let outcome = self.fetch_settings_self_healing_401(auth).await;
let live = self.auth_manager.current_or_expired().map(|a| a.user_id);
self.otel_gate.resolve(&identity, outcome, live.as_deref())
}
/// Fetch settings; on a `401` try one self-healing [`AuthManager::auth`]
/// refresh and re-fetch if it yields a *different* token (recovers a 401
/// from a token that expired mid-fetch). The refresh is bounded by
/// `STARTUP_AUTH_REFRESH_TIMEOUT` so a wedged IdP can't hang the caller; on
/// timeout or error the original `Rejected` stands.
async fn fetch_settings_self_healing_401(
&self,
auth: &crate::auth::GrokAuth,
) -> crate::remote::SettingsFetch {
let outcome = self.fetch_settings(auth).await;
if matches!(outcome, crate::remote::SettingsFetch::Rejected)
&& let Ok(Ok(fresh)) = tokio::time::timeout(
crate::http::STARTUP_AUTH_REFRESH_TIMEOUT,
self.auth_manager.auth(),
)
.await && fresh.key != auth.key
{
return self.fetch_settings(&fresh).await;
}
outcome
}
/// Writes remote settings into `cfg` along with the fields derived from
/// them, so no derived field drifts between post-fetch callers.
pub(super) fn store_remote_settings(
&self,
settings: crate::util::config::RemoteSettings,
) {
let mut cfg = self.cfg.borrow_mut();
cfg.remote_settings = Some(settings);
crate::util::config::sync_campaign_fields(&mut cfg);
if let Some(v) = cfg
.remote_settings
.as_ref()
.and_then(|s| s.path_not_found_hints)
{
cfg.path_not_found_hints = v;
}
}
/// Stores settings and fans out side effects via
/// [`Self::on_remote_settings_changed`]. Shared tail for callers that do
/// not also re-init the telemetry client (those use
/// [`Self::refresh_remote_settings`]).
pub(super) fn install_remote_settings(
&self,
settings: crate::util::config::RemoteSettings,
) {
self.store_remote_settings(settings);
self.on_remote_settings_changed();
}
/// Re-fetch remote settings, re-init the telemetry client, apply side
/// effects, and push `x.ai/settings/update` to clients. Called from both
/// auth handlers (first install + reauth/account switch).
///
/// This only refreshes `cfg.remote_settings` and re-inits the
/// telemetry client (the only global static). Other settings
/// derived from `remote_settings` (`is_trace_upload_enabled`,
/// `web_fetch_enabled`, etc.) are resolved lazily per-turn from
/// `cfg` and pick up the new values automatically.
/// Agent-level fields materialised at startup (`worktree_type`,
/// `restore_code`) are NOT re-resolved here; that requires a
/// broader refactor of the init path.
@ -910,11 +1069,11 @@ impl MvpAgent {
let user_id = auth.user_id.clone();
let team_id = auth.team_id.clone();
let remote_was_absent = self.cfg.borrow().remote_settings.is_none();
let Some(settings) = self.fetch_remote_settings(auth.clone()).await else {
tracing::warn!("post-auth settings refresh failed (HTTP or parse error)");
let Some(settings) = self.fetch_settings_resolving_gate(auth).await else {
return;
};
tracing::info!("post-auth settings refreshed");
self.store_remote_settings(settings);
let (
telemetry_config,
telemetry_mode,
@ -923,11 +1082,9 @@ impl MvpAgent {
deployment_key,
subscription_tier,
) = {
let mut cfg = self.cfg.borrow_mut();
cfg.remote_settings = Some(settings);
crate::util::config::sync_campaign_fields(&mut cfg);
crate::agent::config::apply_remote_settings_side_effects(
cfg.remote_settings.as_ref(),
let cfg = self.cfg.borrow();
crate::util::config::cache_remote_mcp_startup_timeout_secs(
cfg.remote_settings.as_ref().and_then(|s| s.mcp_startup_timeout_secs),
);
let telemetry_mode = cfg.resolve_telemetry_mode();
let trace_upload = cfg.resolve_trace_upload();
@ -953,7 +1110,6 @@ impl MvpAgent {
subscription_tier_display,
)
};
self.sync_collection_config_gate();
let subscription_tier = resolve_subscription_tier_for_telemetry(
subscription_tier,
self.auth_manager.current_or_expired().as_ref(),
@ -970,8 +1126,7 @@ impl MvpAgent {
crate::http::shared_client(),
);
crate::auth::credential_provider::sync_external_otel_identity();
self.emit_announcements(AnnouncementsPushMode::IfChanged);
self.reconfigure_heap_profile_monitor();
self.on_remote_settings_changed();
if remote_was_absent {
self.spawn_auto_worktree_gc();
}
@ -1004,6 +1159,82 @@ impl MvpAgent {
self.emit_announcements(AnnouncementsPushMode::Force);
self.reconfigure_heap_profile_monitor();
}
/// Spawns a background task coalesced on `in_flight`: a request while one
/// is in flight is dropped. The task is bounded by
/// `SETTINGS_REAPPLY_TIMEOUT`. Returns whether a task was spawned.
fn spawn_coalesced_settings_task(
&self,
in_flight: &std::rc::Rc<std::cell::Cell<bool>>,
task: impl std::future::Future<Output = ()> + 'static,
) -> bool {
if in_flight.replace(true) {
return false;
}
let in_flight = in_flight.clone();
tokio::task::spawn_local(async move {
struct ClearOnDrop(std::rc::Rc<std::cell::Cell<bool>>);
impl Drop for ClearOnDrop {
fn drop(&mut self) {
self.0.set(false);
}
}
let _clear = ClearOnDrop(in_flight);
let _ = tokio::time::timeout(crate::http::SETTINGS_REAPPLY_TIMEOUT, task)
.await;
});
true
}
/// Fire-and-forget remote settings refresh for new sessions (at most one
/// in flight).
pub(super) fn spawn_settings_reapply(&self) {
let agent_ref = LocalRef::new(self);
let auth_manager = self.auth_manager.clone();
let _spawned = self
.spawn_coalesced_settings_task(
&self.settings_reapply_in_flight,
async move {
let auth_result = tokio::time::timeout(
crate::http::STARTUP_FETCH_TIMEOUT,
auth_manager.auth(),
)
.await;
if let Ok(Ok(auth)) = auth_result {
let agent = agent_ref.get();
if agent.post_auth_settings_in_flight.get() {
return;
}
agent.refresh_settings_and_reapply(&auth).await;
}
},
);
#[cfg(test)]
if _spawned {
self.settings_reapply_spawn_count
.set(self.settings_reapply_spawn_count.get() + 1);
}
}
/// Resolve post-auth remote settings in the background so a slow or hung
/// `/settings` can't gate `authenticate` (and thus the client's first draw).
/// The external-OTEL gate stays fail-closed until this resolves; the result
/// reaches clients via `x.ai/settings/update`. Its own guard keeps an
/// in-flight reapply from coalescing away the authenticated identity.
pub(super) fn spawn_post_auth_settings(&self, auth: crate::auth::GrokAuth) {
let agent_ref = LocalRef::new(self);
let _spawned = self
.spawn_coalesced_settings_task(
&self.post_auth_settings_in_flight,
async move {
let agent = agent_ref.get();
agent.refresh_remote_settings(&auth).await;
agent.maybe_fetch_post_auth_settings().await;
},
);
#[cfg(test)]
if _spawned {
self.post_auth_settings_spawn_count
.set(self.post_auth_settings_spawn_count.get() + 1);
}
}
/// Spawn the periodic remote-settings poll that pushes mid-session
/// announcement changes to connected clients. Idempotent; plain loop (no
/// cancellation) like `ensure_session_supervisor` — the LocalSet drop at
@ -1165,7 +1396,7 @@ impl MvpAgent {
))
.await
{
Ok(settings) => settings,
Ok(outcome) => outcome.into_option(),
Err(e) => {
tracing::warn!(error = %e, "settings fetch task panicked");
None
@ -1657,7 +1888,9 @@ impl MvpAgent {
RefCell::new(std::collections::HashSet::new()),
),
tier_allowed: std::cell::Cell::new(true),
storage_mode,
allow_access_resolved_for: std::cell::RefCell::new(None),
storage_mode: std::cell::Cell::new(storage_mode),
otel_gate: crate::agent::otel_gate::OtelGate::default(),
default_yolo_mode,
default_auto_mode,
trace_upload_live: Arc::new(
@ -1700,6 +1933,8 @@ impl MvpAgent {
),
session_live_state: RefCell::new(HashMap::new()),
supervisor_started: std::cell::Cell::new(false),
settings_reapply_in_flight: std::rc::Rc::new(std::cell::Cell::new(false)),
post_auth_settings_in_flight: std::rc::Rc::new(std::cell::Cell::new(false)),
announcements_gen: std::cell::Cell::new(0),
last_emitted_announcements: RefCell::new(Vec::new()),
announcements_refresh_started: std::cell::Cell::new(false),
@ -1713,6 +1948,10 @@ impl MvpAgent {
roster_delta_spy: RefCell::new(Vec::new()),
#[cfg(test)]
supervisor_spawn_count: std::cell::Cell::new(0),
#[cfg(test)]
settings_reapply_spawn_count: std::cell::Cell::new(0),
#[cfg(test)]
post_auth_settings_spawn_count: std::cell::Cell::new(0),
};
instance
.auth_manager
@ -1930,7 +2169,7 @@ impl MvpAgent {
}
/// Returns the storage mode configured for this agent
pub fn storage_mode(&self) -> StorageMode {
self.storage_mode
self.storage_mode.get()
}
/// Returns the background copy context for managing background file copy tasks.
pub fn background_copy_context(&self) -> BackgroundCopyContext {

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

@ -502,6 +502,13 @@ struct SettingsUpdateNotification {
tips: Option<Vec<String>>,
slash_command_tags: Option<std::collections::BTreeMap<String, String>>,
announcements: Option<Vec<xai_grok_announcements::RemoteAnnouncement>>,
/// Remote campaigns snapshot for the client's process-global campaign
/// cache. `Some` whenever settings exist (empty means campaigns were
/// withdrawn); `None` when the agent has no settings yet, which clients
/// treat as "leave the cache alone". In leader mode this push is the only
/// seam that seeds the TUI process, so a `/model` pick can record a remote
/// campaign's dismissal even when the TUI's own startup prefetch missed.
campaigns: Option<Vec<crate::util::config::CampaignOverride>>,
gate_message: Option<String>,
gate_url: Option<String>,
gate_label: Option<String>,
@ -702,8 +709,17 @@ pub struct MvpAgent {
/// external-auth users bypass the check). When `false`, the pager shows a
/// gate CTA instead of the prompt.
tier_allowed: std::cell::Cell<bool>,
/// Storage mode - determines whether to sync to backend (writeback) or local only
storage_mode: StorageMode,
/// The `user_id` the current `tier_allowed` verdict was resolved for.
/// `cfg.remote_settings` isn't reset on account switch, so a mismatch here
/// means "unknown" (provisional open), like `OtelGate::rearm_on_switch`.
allow_access_resolved_for: std::cell::RefCell<Option<String>>,
/// Writeback vs local. `Cell` so [`Self::reapply_storage_mode`] can
/// upgrade it when remote settings land; persistence reads the live value.
/// Authoritative post-construction — `Config.storage_mode` is only the
/// boot seed.
storage_mode: std::cell::Cell<StorageMode>,
/// External-OTEL emission gate; see [`crate::agent::otel_gate`].
otel_gate: crate::agent::otel_gate::OtelGate,
/// Default YOLO mode - when true, sessions start with auto-approve enabled.
/// Per-session YOLO tracking lives in SessionHandle.yolo_mode.
default_yolo_mode: bool,
@ -868,6 +884,13 @@ pub struct MvpAgent {
/// once (on the first `spawn_and_register_session`). See
/// `ensure_session_supervisor`.
supervisor_started: std::cell::Cell<bool>,
/// Dedup guard for `spawn_settings_reapply`; at most one task in flight.
/// `Rc` so the drop-guard owns a clone without dereferencing the agent.
settings_reapply_in_flight: std::rc::Rc<std::cell::Cell<bool>>,
/// Separate dedup guard for `spawn_post_auth_settings`, so an in-flight
/// reapply can't coalesce away a freshly authenticated identity's gate and
/// settings resolution.
post_auth_settings_in_flight: std::rc::Rc<std::cell::Cell<bool>>,
/// Last value handed out by `next_announcements_gen` (single-threaded
/// LocalSet, so a plain `Cell` suffices). LEADER-SAFE(shared): one
/// agent-wide push stream.
@ -903,18 +926,18 @@ pub struct MvpAgent {
/// actually spawned. Asserts `ensure_session_supervisor` is idempotent.
#[cfg(test)]
supervisor_spawn_count: std::cell::Cell<usize>,
/// Test-only: counts `spawn_settings_reapply` tasks spawned past the
/// in-flight guard.
#[cfg(test)]
settings_reapply_spawn_count: std::cell::Cell<usize>,
/// Test-only: counts `spawn_post_auth_settings` tasks spawned past its
/// own guard.
#[cfg(test)]
post_auth_settings_spawn_count: std::cell::Cell<usize>,
}
/// Kick off background warmup of the async shared HTTP client.
///
/// Building a `reqwest::Client` is expensive (~95ms) because it loads TLS
/// root certificates. This function spawns a thread to initialize both
/// the shared client and a throwaway sampling client concurrently so
/// that TLS roots are cached before the first session needs them.
///
/// Safe to call multiple times — the underlying `OnceLock` ensures only
/// the first initialization does real work for `shared_client()`. The
/// sampling client is discarded, but the TLS root certificates it loads
/// are cached at the process level by `rustls-native-certs`.
/// Spawn a thread to warm the shared async HTTP client (`OnceLock`-cached).
/// Loading TLS root certs is ~95ms; doing it here avoids a cold-start hit
/// on the first request. Idempotent.
pub fn warm_async_http_client() {
std::thread::spawn(|| {
let _timer = crate::instrumentation_timer!("startup.async_http_warmup");
@ -1769,15 +1792,27 @@ impl MvpAgent {
/// Check whether the user has access via remote settings `allow_access`.
///
/// Non-xAI auth (API keys, enterprise) always passes. For xAI OAuth2
/// users, reads `allow_access` from remote settings. Defaults to
/// `false` (blocked) when remote settings are unavailable.
/// users, reads `allow_access` from remote settings. When settings exist
/// but the field is absent/false, defaults to `false` (blocked); when
/// settings have not arrived yet (background fetch pending) the gate is
/// provisionally open and re-resolved on arrival.
pub(super) async fn enforce_grok_code_access(&self, auth: &crate::auth::GrokAuth) {
if !auth.is_xai_auth() {
self.tier_allowed.set(true);
return;
}
let settings_for_this_identity = self.cfg.borrow().remote_settings.is_some()
&& self.allow_access_resolved_for.borrow().as_deref()
== Some(auth.user_id.as_str());
if !settings_for_this_identity
&& crate::util::config::resolve_remote_fetch_enabled()
{
self.tier_allowed.set(true);
return;
}
let allow = settings_allow_access(self.cfg.borrow().remote_settings.as_ref());
self.tier_allowed.set(allow);
*self.allow_access_resolved_for.borrow_mut() = Some(auth.user_id.clone());
if !allow {
tracing::info!(
"auth: user blocked by allow_access (remote settings grok_build_access_gate)"
@ -1832,18 +1867,11 @@ impl MvpAgent {
}),
),
);
if let Some(settings) = unblocked.settings {
let remote_was_absent = self.cfg.borrow().remote_settings.is_none();
{
let mut cfg = self.cfg.borrow_mut();
cfg.remote_settings = Some(settings);
crate::agent::config::apply_remote_settings_side_effects(
cfg.remote_settings.as_ref(),
);
}
self.sync_collection_config_gate();
self.emit_announcements(AnnouncementsPushMode::IfChanged);
self.reconfigure_heap_profile_monitor();
let remote_was_absent = self.cfg.borrow().remote_settings.is_none();
if let Some(auth) = self.auth_manager.current()
&& let Some(settings) = self.fetch_settings_resolving_gate(&auth).await
{
self.install_remote_settings(settings);
if remote_was_absent {
self.spawn_auto_worktree_gc();
}
@ -2029,43 +2057,17 @@ impl MvpAgent {
if self.cfg.borrow().remote_settings.is_some() {
return;
}
if !crate::util::config::resolve_remote_fetch_enabled() {
return;
}
let Some(auth) = self.auth_manager.current() else {
return;
};
let is_xai_auth = auth.is_xai_auth();
let Some(settings) = self.fetch_remote_settings(auth).await else {
let Some(settings) = self.fetch_settings_resolving_gate(&auth).await else {
return;
};
tracing::info!("post-auth remote_settings fetch succeeded");
{
let mut cfg = self.cfg.borrow_mut();
cfg.remote_settings = Some(settings);
crate::agent::config::apply_remote_settings_side_effects(
cfg.remote_settings.as_ref(),
);
if cfg.storage_mode == StorageMode::Local
&& cfg.mode != crate::agent::config::AgentMode::Generic
{
cfg.storage_mode = StorageMode::resolve(
None,
cfg.remote_settings.as_ref(),
);
if cfg.storage_mode == StorageMode::Writeback && !is_xai_auth {
cfg.storage_mode = StorageMode::Local;
}
}
if let Some(v) = cfg
.remote_settings
.as_ref()
.and_then(|s| s.path_not_found_hints)
{
cfg.path_not_found_hints = v;
}
}
self.sync_collection_config_gate();
self.emit_settings_update_notification();
self.emit_announcements(AnnouncementsPushMode::IfChanged);
self.reconfigure_heap_profile_monitor();
self.install_remote_settings(settings);
self.spawn_auto_worktree_gc();
}
/// Resolve current auto-GC policy and run it on the blocking pool.
@ -2095,6 +2097,7 @@ impl MvpAgent {
tips: rs.and_then(|s| s.tips.clone()),
slash_command_tags: rs.and_then(|s| s.slash_command_tags.clone()),
announcements: rs.and_then(|s| s.announcements.clone()),
campaigns: rs.map(|s| s.campaigns.clone()),
gate_message: rs.and_then(|s| s.gate_message.clone()),
gate_url: rs.and_then(|s| s.gate_url.clone()),
gate_label: rs.and_then(|s| s.gate_label.clone()),

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

@ -3759,6 +3759,400 @@ fn supervisor_reaps_panicked_resident_actor() {
);
});
}
/// Regression: writeback must self-correct once remote settings arrive
/// (the field used to be frozen at construction).
#[tokio::test]
#[serial_test::serial]
async fn storage_mode_self_corrects_to_writeback_when_settings_arrive() {
let _env = crate::env::EnvVarGuard::remove("GROK_STORAGE_MODE");
let auth = crate::auth::GrokAuth {
auth_mode: crate::auth::AuthMode::Oidc,
oidc_issuer: Some("https://auth.x.ai".to_string()),
key: "test-token".to_string(),
..Default::default()
};
let agent = build_agent_with_auth(auth);
agent.cfg.borrow_mut().mode = crate::agent::config::AgentMode::Leader;
assert_eq!(agent.storage_mode(), StorageMode::Local);
agent.cfg.borrow_mut().remote_settings = Some(crate::util::config::RemoteSettings {
writeback_enabled: Some(true),
..Default::default()
});
agent.on_remote_settings_changed();
assert_eq!(agent.storage_mode(), StorageMode::Writeback);
}
/// `spawn_settings_reapply` coalesces: while one reapply is in flight,
/// repeated calls (boot + rapid `/new`) do not spawn overlapping tasks.
#[test]
fn spawn_settings_reapply_coalesces_while_in_flight() {
run_local_for_bridge_test(|| async {
let agent = build_minimal_agent_for_tests();
assert_eq!(agent.settings_reapply_spawn_count.get(), 0);
agent.spawn_settings_reapply();
agent.spawn_settings_reapply();
agent.spawn_settings_reapply();
assert_eq!(
agent.settings_reapply_spawn_count.get(),
1,
"overlapping settings reapplies must coalesce to a single task"
);
assert!(agent.settings_reapply_in_flight.get());
});
}
/// The in-flight guard clears on task completion (via the `ClearOnDrop`
/// guard, so it also clears on panic), allowing a later reapply to re-spawn.
#[test]
fn spawn_settings_reapply_clears_flag_after_completion() {
run_local_for_bridge_test(|| async {
let agent = build_minimal_agent_for_tests();
agent.spawn_settings_reapply();
assert_eq!(agent.settings_reapply_spawn_count.get(), 1);
assert!(agent.settings_reapply_in_flight.get());
let mut cleared = false;
for _ in 0..40 {
if !agent.settings_reapply_in_flight.get() {
cleared = true;
break;
}
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
}
assert!(
cleared,
"in-flight flag must clear after the task completes"
);
agent.spawn_settings_reapply();
assert_eq!(
agent.settings_reapply_spawn_count.get(),
2,
"a reapply after completion must spawn again"
);
});
}
/// The post-auth fetch has its own guard, so an in-flight settings reapply
/// cannot coalesce away a freshly authenticated identity's gate and settings
/// resolution.
#[test]
fn post_auth_settings_not_coalesced_by_in_flight_reapply() {
run_local_for_bridge_test(|| async {
let agent = build_minimal_agent_for_tests();
agent.spawn_settings_reapply();
assert!(agent.settings_reapply_in_flight.get());
agent.spawn_post_auth_settings(crate::auth::GrokAuth::test_default());
assert_eq!(
agent.post_auth_settings_spawn_count.get(),
1,
"post-auth must spawn on its own guard despite an in-flight reapply"
);
assert!(agent.post_auth_settings_in_flight.get());
});
}
/// Agent with pre-loaded auth, a gateway receiver (to assert emitted
/// notifications), and the proxy URL pointed at a mock `/v1/settings`.
fn build_agent_with_auth_and_proxy(
auth: crate::auth::GrokAuth,
proxy_url: String,
mode: crate::agent::config::AgentMode,
) -> (
MvpAgent,
tokio::sync::mpsc::UnboundedReceiver<xai_acp_lib::AcpClientMessage>,
) {
use crate::agent::config::Config as AgentConfig;
use crate::auth::{AuthManager, GrokComConfig};
let temp_dir = tempfile::tempdir().unwrap();
let auth_manager =
std::sync::Arc::new(AuthManager::new(temp_dir.path(), GrokComConfig::default()));
auth_manager.hot_swap(auth);
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
let gateway = GatewaySender::new(tx);
let mut cfg = AgentConfig {
mode,
..Default::default()
};
cfg.endpoints.cli_chat_proxy_base_url = Some(proxy_url);
let agent = MvpAgent::new(gateway, &cfg, auth_manager, None).expect("valid test config");
(agent, rx)
}
/// Drain the gateway, returning `true` if any `x.ai/settings/update`
/// notification was emitted (and acking each so the sender doesn't warn).
fn drained_settings_update(
rx: &mut tokio::sync::mpsc::UnboundedReceiver<xai_acp_lib::AcpClientMessage>,
) -> bool {
let mut found = false;
while let Ok(msg) = rx.try_recv() {
if let xai_acp_lib::AcpClientMessage::ExtNotification(args) = msg {
if &*args.request.method == "x.ai/settings/update" {
found = true;
}
let _ = args.response_tx.send(Ok(()));
}
}
found
}
/// Re-open the process-global external-OTEL gate on drop so a closed gate
/// never leaks into another test.
struct RestoreOtelGate;
impl Drop for RestoreOtelGate {
fn drop(&mut self) {
xai_grok_telemetry::external::mark_external_otel_settings_resolved();
}
}
/// Regression: `cfg.remote_settings` is not reset on an account switch, so the
/// access gate must not read a previous identity's cached `allow_access`. A
/// mismatched identity stays provisionally open (unknown), like the OTEL gate's
/// `rearm_on_switch`.
#[tokio::test]
async fn access_gate_does_not_leak_verdict_across_identities() {
use crate::agent::config::AgentMode;
use crate::auth::{GrokAuth, XAI_OAUTH2_ISSUER};
let auth_a = GrokAuth {
oidc_issuer: Some(XAI_OAUTH2_ISSUER.to_string()),
user_id: "user-a".into(),
..GrokAuth::test_default()
};
let (agent, _rx) = build_agent_with_auth_and_proxy(
auth_a,
"http://127.0.0.1:1/".to_string(),
AgentMode::Leader,
);
{
let mut cfg = agent.cfg.borrow_mut();
cfg.remote_settings = Some(crate::util::config::RemoteSettings {
allow_access: Some(false),
..Default::default()
});
}
*agent.allow_access_resolved_for.borrow_mut() = Some("user-a".to_string());
let auth_b = GrokAuth {
oidc_issuer: Some(XAI_OAUTH2_ISSUER.to_string()),
user_id: "user-b".into(),
..GrokAuth::test_default()
};
assert!(auth_b.is_xai_auth(), "precondition: first-party xAI auth");
agent.enforce_grok_code_access(&auth_b).await;
assert!(
agent.tier_allowed.get(),
"identity B must not inherit identity A's denied allow_access verdict",
);
}
/// First-party xAI auth + `writeback_enabled` settings → storage upgrades to
/// Writeback; the settings arrival also emits `x.ai/settings/update` and opens
/// the external-OTEL gate.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial]
async fn post_auth_settings_xai_upgrades_writeback_emits_and_opens_gate() {
use crate::agent::config::AgentMode;
use crate::auth::{GrokAuth, XAI_OAUTH2_ISSUER};
let _restore = RestoreOtelGate;
let _storage_env = crate::env::EnvVarGuard::remove("GROK_STORAGE_MODE");
let server = xai_grok_test_support::MockInferenceServer::start()
.await
.unwrap();
server.set_settings(serde_json::json!({
"writeback_enabled": true,
"allow_access": true,
}));
let xai_auth = GrokAuth {
oidc_issuer: Some(XAI_OAUTH2_ISSUER.to_string()),
..GrokAuth::test_default()
};
assert!(xai_auth.is_xai_auth(), "precondition: first-party xAI auth");
let (agent, mut rx) =
build_agent_with_auth_and_proxy(xai_auth, server.url(), AgentMode::Leader);
assert_eq!(
agent.storage_mode(),
StorageMode::Local,
"precondition: leader boots in Local storage mode"
);
xai_grok_telemetry::external::suppress_external_otel_until_settings();
assert!(!xai_grok_telemetry::external::is_settings_gate_open());
agent.maybe_fetch_post_auth_settings().await;
assert_eq!(
agent.storage_mode(),
StorageMode::Writeback,
"xai auth + writeback_enabled settings must upgrade storage to Writeback"
);
assert!(
xai_grok_telemetry::external::is_settings_gate_open(),
"a settings response must open the external-OTEL gate"
);
assert!(
drained_settings_update(&mut rx),
"settings arrival must push x.ai/settings/update to clients"
);
}
/// BYOK auth must not be upgraded to `Writeback` even when the server
/// advertises it; the push and gate still fire.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial]
async fn post_auth_settings_non_xai_keeps_local_but_still_emits() {
use crate::agent::config::AgentMode;
use crate::auth::{AuthMode, GrokAuth};
let _restore = RestoreOtelGate;
let server = xai_grok_test_support::MockInferenceServer::start()
.await
.unwrap();
server.set_settings(serde_json::json!({
"writeback_enabled": true,
"allow_access": true,
}));
let api_auth = GrokAuth {
auth_mode: AuthMode::ApiKey,
..GrokAuth::test_default()
};
assert!(
!api_auth.is_xai_auth(),
"precondition: non-first-party auth"
);
let (agent, mut rx) =
build_agent_with_auth_and_proxy(api_auth, server.url(), AgentMode::Leader);
xai_grok_telemetry::external::suppress_external_otel_until_settings();
agent.maybe_fetch_post_auth_settings().await;
assert_eq!(
agent.storage_mode(),
StorageMode::Local,
"non-xai auth must stay Local even when writeback is advertised remotely"
);
assert!(
xai_grok_telemetry::external::is_settings_gate_open(),
"a settings response must open the gate regardless of auth kind"
);
assert!(
drained_settings_update(&mut rx),
"settings arrival must push x.ai/settings/update for non-xai auth too"
);
}
/// A failed post-auth fetch must re-close the gate and leave it closed. Guards
/// two behaviors a passing-on-`Fetched` test can't: the account-switch
/// re-suppress fires (gate was open, identity not yet resolved), and a
/// transient/4xx outcome (`Retry`) does not reopen it.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial]
async fn post_auth_settings_retry_re_suppresses_and_stays_closed() {
use crate::agent::config::AgentMode;
use crate::auth::{GrokAuth, XAI_OAUTH2_ISSUER};
let _restore = RestoreOtelGate;
let server = xai_grok_test_support::MockInferenceServer::start()
.await
.unwrap();
let xai_auth = GrokAuth {
oidc_issuer: Some(XAI_OAUTH2_ISSUER.to_string()),
..GrokAuth::test_default()
};
let (agent, _rx) = build_agent_with_auth_and_proxy(xai_auth, server.url(), AgentMode::Leader);
xai_grok_telemetry::external::mark_external_otel_settings_resolved();
assert!(xai_grok_telemetry::external::is_settings_gate_open());
agent.maybe_fetch_post_auth_settings().await;
assert!(
!xai_grok_telemetry::external::is_settings_gate_open(),
"a Retry (failed) post-auth fetch must re-close the gate and keep it closed"
);
}
/// A same-credential refresh must NOT re-suppress a gate already resolved for
/// that credential; the reason `OtelGate` remembers the identity. With the
/// gate resolved-open for this identity, a later failing (`Retry`) refresh
/// leaves it OPEN (regressing the identity guard would re-close it forever).
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial]
async fn same_credential_refresh_does_not_flap_resolved_gate() {
use crate::agent::config::AgentMode;
use crate::auth::{GrokAuth, XAI_OAUTH2_ISSUER};
let _restore = RestoreOtelGate;
let server = xai_grok_test_support::MockInferenceServer::start()
.await
.unwrap();
let xai_auth = GrokAuth {
oidc_issuer: Some(XAI_OAUTH2_ISSUER.to_string()),
..GrokAuth::test_default()
};
let (agent, _rx) =
build_agent_with_auth_and_proxy(xai_auth.clone(), server.url(), AgentMode::Leader);
agent.otel_gate.set_resolved_for(&xai_auth.user_id);
xai_grok_telemetry::external::mark_external_otel_settings_resolved();
assert!(xai_grok_telemetry::external::is_settings_gate_open());
agent.refresh_remote_settings(&xai_auth).await;
assert!(
xai_grok_telemetry::external::is_settings_gate_open(),
"a same-credential refresh must not flap a gate already resolved for it"
);
}
/// A `/settings` 401 from a token that rotated mid-flight must self-heal:
/// refresh once and, if the token changed, re-fetch with it. Without the
/// re-fetch the stale 401 fails OPEN (no remote policy).
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial]
async fn settings_self_heal_refetches_after_token_rotation() {
use crate::agent::config::AgentMode;
use crate::auth::refresh::{RefreshOutcome, TokenRefresher};
use crate::auth::{GrokAuth, XAI_OAUTH2_ISSUER};
let _restore = RestoreOtelGate;
let server = xai_grok_test_support::MockInferenceServer::start_with_required_auth(
vec![xai_grok_test_support::MockModelEntry::new("grok-build")],
"rotated-key",
)
.await
.unwrap();
server.set_settings(serde_json::json!({ "allow_access": true }));
struct RotatingRefresher;
#[async_trait::async_trait]
impl TokenRefresher for RotatingRefresher {
async fn refresh(&self, _r: crate::auth::manager::RefreshReason) -> RefreshOutcome {
RefreshOutcome::Success(Box::new(GrokAuth {
key: "rotated-key".into(),
oidc_issuer: Some(XAI_OAUTH2_ISSUER.to_string()),
refresh_token: Some("rt".into()),
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
..GrokAuth::test_default()
}))
}
}
let stale = GrokAuth {
key: "stale-key".into(),
oidc_issuer: Some(XAI_OAUTH2_ISSUER.to_string()),
refresh_token: Some("rt".into()),
expires_at: Some(chrono::Utc::now() - chrono::Duration::hours(1)),
..GrokAuth::test_default()
};
let (agent, _rx) =
build_agent_with_auth_and_proxy(stale.clone(), server.url(), AgentMode::Leader);
agent
.auth_manager
.set_refresher(std::sync::Arc::new(RotatingRefresher));
xai_grok_telemetry::external::suppress_external_otel_until_settings();
agent.refresh_remote_settings(&stale).await;
assert!(
xai_grok_telemetry::external::is_settings_gate_open(),
"the rotated-token re-fetch must land settings and open the gate"
);
assert!(
agent.cfg.borrow().remote_settings.is_some(),
"the re-fetched settings must be stored"
);
}
/// A logout can land while the detached post-auth fetch is in flight; the
/// result must not be cached for the logged-out identity.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial]
async fn settings_not_cached_when_identity_logs_out_during_fetch() {
use crate::agent::config::AgentMode;
use crate::auth::{GrokAuth, XAI_OAUTH2_ISSUER};
let _restore = RestoreOtelGate;
let server = xai_grok_test_support::MockInferenceServer::start()
.await
.unwrap();
server.set_settings(serde_json::json!({ "allow_access": true }));
let xai_auth = GrokAuth {
oidc_issuer: Some(XAI_OAUTH2_ISSUER.to_string()),
..GrokAuth::test_default()
};
let (agent, _rx) =
build_agent_with_auth_and_proxy(xai_auth.clone(), server.url(), AgentMode::Leader);
agent.auth_manager.clear_in_memory();
agent.refresh_remote_settings(&xai_auth).await;
assert!(
agent.cfg.borrow().remote_settings.is_none(),
"settings fetched for a logged-out identity must not be cached"
);
}
/// `ensure_session_supervisor` is idempotent: calling it repeatedly spawns
/// the sweeper loop exactly once.
#[test]

209
crates/codegen/xai-grok-shell/src/agent/otel_gate.rs Normal file
View file

@ -0,0 +1,209 @@
//! External-OTEL emission gate.
//!
//! Single owner of the fail-closed gate that decides whether customer-owned
//! OTEL telemetry may ship, over the process-global flag in
//! [`xai_grok_telemetry::external`]:
//!
//! 1. Startup (no leader instance yet): [`suppress`] closes the gate before
//! telemetry init; [`open_at_startup`] re-opens it only for a pure
//! env-API-key leader ([`should_open_at_startup`]), which has no remote
//! policy to fetch.
//! 2. Post-auth/refresh (per-leader): [`OtelGate::resolve`] drives the gate
//! from the [`SettingsFetch`] outcome for the still-live identity.
//!
//! A leader that never authenticates keeps the gate closed for life: the gate
//! fails safe by dropping telemetry, never by shipping it early.
use crate::remote::SettingsFetch;
use crate::util::config::RemoteSettings;
/// Closes the gate. Process-global and idempotent; callable before any
pub(crate) fn suppress() {
xai_grok_telemetry::external::suppress_external_otel_until_settings();
}
/// Inputs to [`should_open_at_startup`]. Named fields prevent transposed
pub(crate) struct StartupGate {
pub(crate) has_session: bool,
pub(crate) has_api_key_env: bool,
pub(crate) session_pending: bool,
/// When false, no remote fleet policy can arrive, so the gate fails open to the leader's local telemetry decision.
pub(crate) remote_fetch_enabled: bool,
}
/// Returns whether a leader opens the gate at startup: only a pure
pub(crate) fn should_open_at_startup(gate: StartupGate) -> bool {
// No remote policy will arrive with `remote_fetch` off, so fail open to
if !gate.remote_fetch_enabled {
return true;
}
!gate.has_session && gate.has_api_key_env && !gate.session_pending
}
/// Returns whether a session-less startup is about to mint a grok.com session
pub(crate) fn is_session_pending(
has_session: bool,
grok_com_config: &crate::auth::GrokComConfig,
) -> bool {
!has_session
&& (grok_com_config.auth_provider_command.is_some()
|| crate::auth::devbox_login::is_devbox_environment())
}
/// Opens the gate at startup once [`should_open_at_startup`] holds; a later session re-resolves via [`OtelGate::resolve`].
pub(crate) fn open_at_startup() {
xai_grok_telemetry::external::mark_external_otel_settings_resolved();
}
/// Per-leader memory over the process-global external-OTEL gate: the credential
#[derive(Default)]
pub(crate) struct OtelGate {
resolved_for: std::cell::RefCell<Option<String>>,
}
impl OtelGate {
/// Re-closes the gate before fetching a different identity's policy, so a stale open can't leak across an account switch.
pub(crate) fn rearm_on_switch(&self, identity: &str) {
if identity.is_empty() || self.resolved_for.borrow().as_deref() != Some(identity) {
xai_grok_telemetry::external::suppress_external_otel_until_settings();
}
}
/// Drives the gate from a settings-fetch `outcome` for `identity`: fail-closed on transient outcomes, opens on a definitive one. Returns settings only when fetched.
pub(crate) fn resolve(
&self,
identity: &str,
outcome: SettingsFetch,
live_identity: Option<&str>,
) -> Option<RemoteSettings> {
if live_identity != Some(identity) {
return None;
}
match outcome {
SettingsFetch::Fetched(settings) => {
self.apply_and_open(identity, Some(&settings));
Some(*settings)
}
SettingsFetch::Rejected => {
self.apply_and_open(identity, None);
None
}
SettingsFetch::Retry => None,
}
}
/// Applies the tighten-only fleet policy from `settings` (`None` on a `401`), then opens the gate and records `identity` (policy before open).
fn apply_and_open(&self, identity: &str, settings: Option<&RemoteSettings>) {
crate::agent::config::apply_external_otel_remote_policy(settings);
xai_grok_telemetry::external::mark_external_otel_settings_resolved();
*self.resolved_for.borrow_mut() = Some(identity.to_owned());
}
#[cfg(test)]
pub(crate) fn set_resolved_for(&self, identity: &str) {
*self.resolved_for.borrow_mut() = Some(identity.to_owned());
}
}
#[cfg(test)]
mod tests {
use super::*;
use xai_grok_telemetry::external::{
is_settings_gate_open, mark_external_otel_settings_resolved,
suppress_external_otel_until_settings,
};
/// Restore the process-global gate open on exit so a closed gate never leaks.
struct RestoreGate;
impl Drop for RestoreGate {
fn drop(&mut self) {
mark_external_otel_settings_resolved();
}
}
fn fetched() -> SettingsFetch {
SettingsFetch::Fetched(Box::default())
}
#[test]
fn startup_gate_fails_open_when_remote_fetch_disabled() {
// remote_fetch off => no remote policy will ever arrive => fail open,
assert!(should_open_at_startup(StartupGate {
has_session: true,
has_api_key_env: false,
session_pending: false,
remote_fetch_enabled: false,
}));
assert!(!should_open_at_startup(StartupGate {
has_session: true,
has_api_key_env: false,
session_pending: false,
remote_fetch_enabled: true,
}));
}
#[test]
#[serial_test::serial]
fn resolve_opens_only_on_definitive_outcome_for_live_identity() {
let _restore = RestoreGate;
let gate = OtelGate::default();
suppress_external_otel_until_settings();
assert!(
gate.resolve("alice", SettingsFetch::Retry, Some("alice"))
.is_none()
);
assert!(
!is_settings_gate_open(),
"a transient outcome stays fail-closed"
);
assert!(
gate.resolve("alice", SettingsFetch::Rejected, Some("alice"))
.is_none()
);
assert!(
is_settings_gate_open(),
"a rejected credential opens the gate"
);
suppress_external_otel_until_settings();
assert!(gate.resolve("alice", fetched(), Some("alice")).is_some());
assert!(
is_settings_gate_open(),
"a fetched outcome opens for the live identity"
);
}
#[test]
#[serial_test::serial]
fn resolve_skips_open_for_a_stale_identity() {
let _restore = RestoreGate;
let gate = OtelGate::default();
suppress_external_otel_until_settings();
assert!(
gate.resolve("alice", fetched(), Some("bob")).is_none(),
"a stale identity must not return settings"
);
assert!(
!is_settings_gate_open(),
"a stale identity must not open the gate"
);
}
#[test]
#[serial_test::serial]
fn rearm_re_closes_for_an_empty_identity() {
let _restore = RestoreGate;
let gate = OtelGate::default();
gate.set_resolved_for("");
mark_external_otel_settings_resolved();
gate.rearm_on_switch("");
assert!(
!is_settings_gate_open(),
"an empty identity must always re-close (cannot prove same credential)"
);
}
}

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

@ -307,6 +307,10 @@ async fn run_persistent_agent(
// Restore managed policy right before bootstrap reads it — the agent is created lazily here,
// so an earlier restore could go stale before the gate.
crate::managed_config::ensure_managed_policy_present(&auth_manager).await;
// Fail-closed external-OTEL gate: suppress until settings resolve, opening
// now only for a pure env-API-key user (no remote policy). Matches the
// stdio/leader boot; per-connection settings reopen it via `initialize`.
crate::agent::app::apply_otel_config(&auth_manager, &agent_config.grok_com_config);
let agent = Rc::new(
MvpAgent::new(gateway, &agent_config, auth_manager, prefetched_models)
.unwrap_or_else(crate::agent::init::exit_on_config_error),

29
crates/codegen/xai-grok-shell/src/agent/subscription_check.rs
View file

@ -3,8 +3,8 @@
//! Provides `single_check()` which queries `GET /user?include=subscription`
//! for the live subscription tier from the backend, independent of the JWT.
//! If a qualifying tier is detected, does a best-effort JWT refresh and
//! settings re-fetch, then returns an `UnblockResult` so the agent can
//! lift the gate.
//! returns an `UnblockResult` so the agent can re-fetch settings and lift
//! the gate through its own settings seam.
//!
//! The pager drives the polling via `x.ai/auth/check_subscription`: the 5s
//! paywall chain, the free-tier watch, the refocus check, and
@ -27,11 +27,9 @@ const QUALIFYING_TIERS: &[&str] = &[
"XPremium",
"XBasic",
];
/// Successful subscription check result: confirmed qualifying tier +
/// optionally refreshed settings.
/// Successful subscription check result: a confirmed qualifying tier.
pub(crate) struct UnblockResult {
pub(crate) new_tier: String,
pub(crate) settings: Option<crate::util::config::RemoteSettings>,
}
/// Fetch `/user?include=subscription` and return the parsed `UserInfo`.
async fn fetch_user_info(
@ -68,9 +66,9 @@ async fn fetch_user_info(
/// the paywall is shown (`x.ai/auth/check_subscription`).
///
/// Queries `/user?include=subscription` for the live tier. If a qualifying
/// tier is found, does a best-effort JWT refresh + settings re-fetch and
/// returns `Some(UnblockResult)`. Returns `None` if no qualifying
/// subscription exists or the request fails.
/// tier is found, does a best-effort JWT refresh and returns
/// `Some(UnblockResult)`. Returns `None` if no qualifying subscription
/// exists or the request fails.
#[tracing::instrument(name = "paywall_check", skip_all, fields(user_id = %user_id))]
pub(crate) async fn single_check(
auth_manager: Arc<AuthManager>,
@ -137,25 +135,12 @@ pub(crate) async fn single_check(
})),
);
}
let settings = if crate::util::config::resolve_remote_fetch_enabled() {
let base_url = proxy_base_url.to_string();
let auth_for_settings = auth_manager.current().unwrap_or(auth);
let atk = alpha_test_key.map(str::to_string);
tokio::task::spawn_blocking(move || {
crate::remote::fetch_settings_blocking(&base_url, &auth_for_settings, atk.as_deref())
})
.await
.ok()
.flatten()
} else {
None
};
xai_grok_telemetry::unified_log::info(
"paywall_check_unblocked",
None,
Some(serde_json::json!({ "user_id": user_id, "new_tier": new_tier })),
);
Some(UnblockResult { new_tier, settings })
Some(UnblockResult { new_tier })
}
#[cfg(test)]
mod tests {

146
crates/codegen/xai-grok-shell/src/auth/flow.rs
View file

@ -220,6 +220,9 @@ async fn run_external_auth_provider(
.stdin(std::process::Stdio::null())
.stdout(std::process::Stdio::piped())
.kill_on_drop(true);
// TODO: `kill_on_drop` SIGKILLs only the direct `sh` child; a provider that
// backgrounds work (setsid / `&`) leaks the grandchild on shutdown-cancel.
// Proper fix: pgid-kill via xai-tty-utils.
// TUI: pipe stderr and forward via callback — inherit would corrupt the
// alternate screen. CLI / headless: inherit so URLs and progress appear in
@ -675,12 +678,23 @@ async fn run_auth_flow_inner(
///
/// Returns `None` when no valid credentials can be obtained non-interactively.
pub async fn try_ensure_fresh_auth(grok_com_config: &GrokComConfig) -> Option<GrokAuth> {
let grok_home = grok_home::grok_home();
let auth_manager = std::sync::Arc::new(AuthManager::new(&grok_home, grok_com_config.clone()));
try_ensure_fresh_auth_with(&build_startup_auth_manager(grok_com_config)).await
}
// auth() handles cached-valid (fast path), OIDC refresh, external
// binary -- all through refresh_chain (single mutation point).
/// Builds and configures the startup `AuthManager`; the policy helpers below
/// take it injected so tests can substitute their own.
fn build_startup_auth_manager(grok_com_config: &GrokComConfig) -> Arc<AuthManager> {
let auth_manager = Arc::new(AuthManager::new(
&grok_home::grok_home(),
grok_com_config.clone(),
));
// auth()'s OIDC/external refresh needs the refresher configured first.
auth_manager.configure_refresher(grok_com_config.auth_provider_command.clone(), None);
auth_manager
}
/// Policy: cached-valid creds, else silent refresh (no interactive login).
async fn try_ensure_fresh_auth_with(auth_manager: &Arc<AuthManager>) -> Option<GrokAuth> {
match auth_manager.auth().await {
Ok(auth) => Some(auth),
Err(e) => {
@ -690,24 +704,37 @@ pub async fn try_ensure_fresh_auth(grok_com_config: &GrokComConfig) -> Option<Gr
}
}
/// Like `try_ensure_fresh_auth` but also mints on cold start (external provider /
/// devbox, never a browser; may take up to ~300s). For detached modes only.
pub(crate) async fn try_ensure_session_noninteractive(
/// Readiness-path auth: a bounded refresh plus the expired-but-refreshable
/// cached session, but no cold mint (which can run a provider command up to
/// `STARTUP_AUTH_TIMEOUT`). Minting is deferred to the post-readiness
/// background task, so readiness waits at most `STARTUP_AUTH_REFRESH_TIMEOUT`.
pub(crate) async fn try_noninteractive_auth_no_mint(
grok_com_config: &GrokComConfig,
) -> Option<GrokAuth> {
if let Some(auth) = try_ensure_fresh_auth(grok_com_config).await {
return Some(auth);
}
let grok_home = grok_home::grok_home();
let auth_manager = Arc::new(AuthManager::new(&grok_home, grok_com_config.clone()));
try_noninteractive_auth_no_mint_with(&build_startup_auth_manager(grok_com_config)).await
}
// Transient refresh failure: credentials remain (usable on 401 recovery).
// Permanent failure already discarded them.
if let Some(expired) = expired_refreshable_session(&auth_manager) {
return Some(expired);
/// Policy behind [`try_noninteractive_auth_no_mint`], with the `AuthManager`
/// injected for tests.
async fn try_noninteractive_auth_no_mint_with(auth_manager: &Arc<AuthManager>) -> Option<GrokAuth> {
match tokio::time::timeout(
crate::http::STARTUP_AUTH_REFRESH_TIMEOUT,
try_ensure_fresh_auth_with(auth_manager),
)
.await
{
Ok(Some(auth)) => return Some(auth),
Ok(None) => {}
Err(_elapsed) => {
tracing::warn!(
timeout_secs = crate::http::STARTUP_AUTH_REFRESH_TIMEOUT.as_secs(),
"boot auth refresh timed out; using cached/expired session (mint deferred to background)"
);
}
}
mint_session_noninteractive(&auth_manager, grok_com_config).await
// Expired-but-refreshable cached session self-heals on the first 401; no
// cold mint on the readiness path.
expired_refreshable_session(auth_manager)
}
/// A cached, refreshable session (not BYOK/ApiKey). Reached only after fresh
@ -719,11 +746,14 @@ fn expired_refreshable_session(auth_manager: &AuthManager) -> Option<GrokAuth> {
}
/// Cold-start mint via non-interactive providers (external command, devbox);
/// `None` when none is available.
async fn mint_session_noninteractive(
/// `None` when none is available. Persists the result into `auth_manager` (disk
/// and in-memory) so per-request `auth()` self-heals. Carries no timeout of its
/// own: the readiness-path caller imposes `STARTUP_AUTH_TIMEOUT`, while the
/// leader's background re-mint runs uncapped (only the provider's ~300s ceiling).
pub(crate) async fn mint_session_noninteractive(
auth_manager: &Arc<AuthManager>,
grok_com_config: &GrokComConfig,
) -> Option<GrokAuth> {
let grok_com_config = auth_manager.grok_com_config();
// preferred_method=api_key: never auto-mint OIDC (fail-closed).
if grok_com_config.blocks_automatic_oidc() {
tracing::debug!(
@ -1157,7 +1187,7 @@ mod tests {
AuthManager::new(dir.path(), cfg.clone()).with_proxy_base_url(&dead_proxy_url()),
);
let auth = mint_session_noninteractive(&mgr, &cfg).await;
let auth = mint_session_noninteractive(&mgr).await;
assert_eq!(auth.map(|a| a.key), Some("xai-ext-token".to_string()));
}
@ -1967,4 +1997,76 @@ mod tests {
"wrong-team auth.json must be cleared, forcing a compliant re-login"
);
}
/// Mock OIDC IdP whose `/token` endpoint never responds, so a refresh
/// attempt hangs until the caller bounds it.
async fn start_hanging_oidc_idp() -> (String, tokio::task::JoinHandle<()>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port());
let b = base.clone();
let app = axum::Router::new()
.route(
"/.well-known/openid-configuration",
axum::routing::get(move || {
let b = b.clone();
async move {
axum::Json(serde_json::json!({
"authorization_endpoint": format!("{b}/authorize"),
"token_endpoint": format!("{b}/token"),
}))
}
}),
)
.route(
"/token",
axum::routing::post(|| async {
// Never responds: the caller must bound the refresh.
tokio::time::sleep(std::time::Duration::from_secs(3600)).await;
axum::Json(serde_json::json!({}))
}),
);
let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
(base, handle)
}
/// The readiness-path `_no_mint` variant bounds the refresh (~5s) and never
/// engages the cold-mint fallback, so leader readiness can't block on a
/// provider command up to the 60s `STARTUP_AUTH_TIMEOUT` cap.
#[tokio::test]
async fn no_mint_readiness_auth_is_bounded() {
let (idp_base, server) = start_hanging_oidc_idp().await;
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let am = Arc::new(AuthManager::new(dir.path(), cfg.clone()));
am.configure_refresher(cfg.auth_provider_command.clone(), None);
am.hot_swap(GrokAuth {
key: "expired".into(),
auth_mode: AuthMode::Oidc,
oidc_issuer: Some(idp_base.clone()),
oidc_client_id: Some("test-client".into()),
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() - chrono::Duration::hours(1)),
..GrokAuth::test_default()
});
let started = std::time::Instant::now();
let result = try_noninteractive_auth_no_mint_with(&am).await;
let elapsed = started.elapsed();
assert!(
elapsed >= crate::http::STARTUP_AUTH_REFRESH_TIMEOUT,
"expected a bounded refresh attempt (elapsed {elapsed:?})"
);
assert!(
elapsed < crate::http::STARTUP_AUTH_TIMEOUT,
"no-mint readiness auth must not engage the 60s cold-mint cap (elapsed {elapsed:?}); readiness would block on a provider command"
);
assert!(
result.is_none(),
"a non-xAI expired session is no first-party fallback and no mint runs on this path, so no auth is produced"
);
server.abort();
}
}

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

@ -124,6 +124,14 @@ const PERMANENT_FAILURE_TTL: StdDuration = StdDuration::from_secs(300);
/// `attempted_verdict_key`, when a verdict is stored), never co-held. Never hold
/// a `parking_lot` guard across `.await`. Refreshers return [`RefreshOutcome`]
/// for `refresh_chain` to apply.
/// Redacted `Debug` so `AuthManager` (held via `Arc` inside `Debug`-derived
/// types like `PersistenceMsg`) never leaks credentials into logs or panics.
impl std::fmt::Debug for AuthManager {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AuthManager").finish_non_exhaustive()
}
}
pub struct AuthManager {
/// In-memory bearer. Mutate via [`Self::with_inner_write`] or
/// [`Self::refresh_chain`]; the closure helpers' sync return type

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

@ -865,9 +865,9 @@ async fn verdict_not_keyed_on_in_mem_bearer() {
/// but cannot write it to disk must surface `Transient` AND still swap the
/// in-memory bearer to the fresh token (the "always update in-memory even if the
/// disk write failed" invariant — without it a disk hiccup strands the session).
/// The write is failed deterministically (root-safe) by planting a *directory*
/// at the atomic-write temp path so `open_secure_file` hits `EISDIR`; the
/// auth.json read (file absent) and the file lock still succeed.
/// The write is failed deterministically (root-safe) via the path-scoped
/// `WRITE_FAULT_PATH` injection in `storage.rs`; the auth.json read (file
/// absent) and the file lock still succeed.
#[tokio::test]
async fn refresh_persist_failure_is_transient_but_swaps_in_memory() {
let dir = tempfile::tempdir().unwrap();
@ -882,14 +882,20 @@ async fn refresh_persist_failure_is_transient_but_swaps_in_memory() {
..GrokAuth::test_default()
});
// `write_auth_json_atomic` writes `auth.json.<pid>.tmp` then renames; a
// directory there makes the temp-file open fail with EISDIR (enforced even
// for root), so the persist fails while the read/lock paths are unaffected.
std::fs::create_dir(
dir.path()
.join(format!("auth.json.{}.tmp", std::process::id())),
)
.unwrap();
// Fail every atomic write to THIS tempdir's auth.json (path-scoped, so
// parallel tests are unaffected). Cleared on drop.
struct FaultGuard;
impl Drop for FaultGuard {
fn drop(&mut self) {
*crate::auth::storage::WRITE_FAULT_PATH
.lock()
.unwrap_or_else(|e| e.into_inner()) = None;
}
}
let _fault = FaultGuard;
*crate::auth::storage::WRITE_FAULT_PATH
.lock()
.unwrap_or_else(|e| e.into_inner()) = Some(dir.path().join("auth.json"));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: Arc::new(AtomicU32::new(0)),

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

@ -31,8 +31,8 @@ pub use config::{
};
pub(crate) use external_auth::{parse_output, refresh_with_command};
pub(crate) use flow::{
AuthChannels, run_auth_flow, run_auth_flow_with_stderr_bridge,
try_ensure_session_noninteractive,
AuthChannels, mint_session_noninteractive, run_auth_flow, run_auth_flow_with_stderr_bridge,
try_noninteractive_auth_no_mint,
};
pub use flow::{
AuthUrlInfo, AuthUrlMode, LoginTransportOverride, LogoutResult, ensure_authenticated,

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

@ -1,6 +1,7 @@
use std::fs::File;
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use super::model::{API_KEY_SCOPE, AuthMode, AuthStore, GrokAuth, lookup_auth};
@ -273,16 +274,55 @@ fn write_store_to(path: &Path, auth_store: &AuthStore) -> std::io::Result<()> {
Ok(())
}
/// Test-only, path-scoped write fault: `write_auth_json_atomic` fails with
/// `Unsupported` for exactly this `auth.json` path. Path-scoped so parallel
/// tests in the same process do not sabotage each other.
#[cfg(test)]
pub(super) static WRITE_FAULT_PATH: std::sync::Mutex<Option<PathBuf>> = std::sync::Mutex::new(None);
/// Atomic write: tmp + rename. Unix `rename(2)` replaces atomically;
/// Windows `rename` requires removing the target first.
fn write_auth_json_atomic(auth_file: &Path, auth_store: &AuthStore) -> std::io::Result<()> {
let tmp = auth_file.with_extension(format!("json.{}.tmp", std::process::id()));
#[cfg(test)]
if WRITE_FAULT_PATH
.lock()
.unwrap_or_else(|e| e.into_inner())
.as_deref()
== Some(auth_file)
{
return Err(std::io::Error::new(
std::io::ErrorKind::Unsupported,
"injected write fault (WRITE_FAULT_PATH)",
));
}
// Unique per write (pid + monotonic seq): two concurrent in-process writers
// (e.g. background mint + proactive refresher) must not share one tmp path.
static TMP_SEQ: AtomicU64 = AtomicU64::new(0);
let tmp = auth_file.with_extension(format!(
"json.{}.{}.tmp",
std::process::id(),
TMP_SEQ.fetch_add(1, Ordering::Relaxed)
));
// Reclaim the temp file on any early return (write/sync/rename failure); the
// unique name otherwise accumulates one orphan per failed write.
struct TmpReclaim<'a>(Option<&'a Path>);
impl Drop for TmpReclaim<'_> {
fn drop(&mut self) {
if let Some(p) = self.0 {
let _ = std::fs::remove_file(p);
}
}
}
let mut tmp_reclaim = TmpReclaim(Some(&tmp));
write_store_to(&tmp, auth_store)?;
#[cfg(windows)]
{
let _ = std::fs::remove_file(auth_file);
}
std::fs::rename(&tmp, auth_file)?;
tmp_reclaim.0 = None; // renamed into place; nothing to reclaim
// Re-assert on the final path (covers rename edge cases / FS quirks).
// Best-effort: rename already published the new tokens.
if let Err(e) = crate::util::secure_file::ensure_owner_only_permissions(auth_file) {
@ -533,6 +573,31 @@ mod write_fallback_tests {
assert_eq!(read_key(&path).as_deref(), Some("secret-key"));
}
/// On a failed atomic write, the `TmpReclaim` guard must remove the temp
/// file so no orphan accumulates. Here `auth.json` is a directory, so the
/// `rename` fails after the temp file is written.
#[test]
fn atomic_write_reclaims_tmp_on_failure() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("auth.json");
std::fs::create_dir(&path).unwrap();
assert!(
write_auth_json_atomic(&path, &sample_store()).is_err(),
"rename onto a directory must fail"
);
let orphans: Vec<_> = std::fs::read_dir(dir.path())
.unwrap()
.filter_map(Result::ok)
.filter(|e| e.file_name().to_string_lossy().contains(".tmp"))
.collect();
assert!(
orphans.is_empty(),
"TmpReclaim must remove the temp file on failure: {orphans:?}"
);
}
/// A fallback write that truncates then fails must roll back to the prior
/// bytes instead of leaving an empty/torn file — otherwise a second
/// disk-full failure would destroy a previously-valid credential.

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

@ -840,6 +840,19 @@ impl StorageMode {
}
Self::Local
}
/// Resolve from remote settings, enforcing the rule that `Writeback`
/// requires grok.com auth (it syncs to grok-code-backend). This is the
/// single home for that gate, used at boot ([`crate::agent::init`]) and by
/// the post-readiness self-heal (`MvpAgent::reapply_storage_mode`).
pub fn from_remote_gated(
remote: Option<&crate::util::config::RemoteSettings>,
has_xai_auth: bool,
) -> Self {
match Self::resolve(None, remote) {
Self::Writeback if !has_xai_auth => Self::Local,
mode => mode,
}
}
/// Returns true if this mode syncs to the backend.
pub fn is_writeback(&self) -> bool {
matches!(self, Self::Writeback)

2
crates/codegen/xai-grok-shell/src/config/reloader.rs
View file

@ -273,7 +273,7 @@ impl ConfigReloader {
}
}
fn reload_auth(&mut self) -> anyhow::Result<()> {
pub(crate) fn reload_auth(&mut self) -> anyhow::Result<()> {
let auth_path = self.grok_home.join("auth.json");
let store = read_auth_json(&auth_path)?;

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

@ -3615,3 +3615,26 @@ fn kill_switched_cold_cwd_stays_allowed_through_plugins_config_read() {
"gate must still allow the kill-switched folder after the config read"
);
}
/// Writeback requires grok.com auth: remote may advertise it, but a non-xai
/// credential is downgraded to `Local`.
#[test]
#[serial_test::serial]
fn from_remote_gated_requires_xai_auth_for_writeback() {
let _env = crate::env::EnvVarGuard::remove("GROK_STORAGE_MODE");
let writeback = crate::util::config::RemoteSettings {
writeback_enabled: Some(true),
..Default::default()
};
assert_eq!(
StorageMode::from_remote_gated(Some(&writeback), true),
StorageMode::Writeback
);
assert_eq!(
StorageMode::from_remote_gated(Some(&writeback), false),
StorageMode::Local,
);
assert_eq!(
StorageMode::from_remote_gated(None, true),
StorageMode::Local
);
}

11
crates/codegen/xai-grok-shell/src/leader/client.rs
View file

@ -26,12 +26,13 @@ const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30);
/// Timeout for receiving registration response from server.
/// This prevents indefinite hangs if the server doesn't respond.
const REGISTRATION_RESPONSE_TIMEOUT: Duration = Duration::from_secs(10);
/// Timeout for waiting for `LeaderReady` after a `Registered { ready: false }` response.
/// Timeout for waiting for `LeaderReady` after a `Registered { ready: false }`.
///
/// Auth + model prefetch can take significant time (network calls, potential browser
/// OAuth flow). 5 minutes is generous enough to cover all practical scenarios; if the
/// leader fails it will close the connection first anyway.
const LEADER_READY_TIMEOUT: Duration = Duration::from_secs(300);
/// The leader signals readiness right after its bounded sign-in
/// (`STARTUP_AUTH_TIMEOUT`); model/settings prefetch runs off the readiness path
/// and the leader never opens a browser OAuth flow. This therefore only needs to
/// cover that bounded auth plus margin, matching the client connect ceiling.
const LEADER_READY_TIMEOUT: Duration = crate::http::MIN_CLIENT_CONNECT_TIMEOUT;
/// Reason the client disconnected from the leader server.
///

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

@ -849,7 +849,7 @@ fn make_leader_starting_error(json: &serde_json::Value) -> Option<String> {
"error": {
"code": -32002,
"message": "leader_starting",
"data": "Leader is still initializing (auth/prefetch in progress). Retry shortly."
"data": "Leader is still initializing (auth in progress). Retry shortly."
}
});
Some(response.to_string())
@ -1480,7 +1480,8 @@ fn make_version_mismatch_notification(
/// JSON-RPC error so the client can retry rather than hang.
/// - ACP notifications (no `id`) are dropped with a trace log.
///
/// Once `ready_rx` is signaled `true` (auth + prefetch complete), all subsequent
/// Once `ready_rx` is signaled `true` (socket bound + bounded auth complete; the
/// model catalog and remote settings stream in afterward), all subsequent
/// ACP traffic is forwarded to the agent as normal.
///
/// # Arguments
@ -2567,14 +2568,16 @@ pub struct ServerHandle {
pub client_count: Arc<AtomicUsize>,
/// Atomic flag: `true` while the agent has pending (in-flight) requests
pub agent_busy: Arc<AtomicBool>,
/// Signal the IPC server that the leader is fully ready (auth + prefetch complete).
/// Signal the IPC server that the leader is fully ready (socket bound + bounded auth;
/// catalog/settings refresh runs in the background).
///
/// Send `true` once the leader has finished initializing. Until then, ACP requests
/// receive a `leader_starting` error and ACP notifications are dropped.
///
/// `spawn_leader_server` sends `true` immediately so that callers that do not need
/// staged startup (e.g. tests, in-process use) get a fully-ready server out of the box.
/// Production leader startup (`run_leader`) holds this back until auth + prefetch succeed.
/// Production leader startup (`run_leader`) holds this back until bounded auth completes
/// (catalog/settings are no longer prefetched; they refresh in the background).
pub ready_tx: watch::Sender<bool>,
/// Set the shutdown reason before cancelling so clients receive the correct `ShuttingDown`
/// reason. The default value is [`ShutdownReason::Manual`]; send

134
crates/codegen/xai-grok-shell/src/remote/client.rs
View file

@ -552,24 +552,60 @@ impl BackendClient {
Ok(())
}
}
/// Fetch remote settings from cli-chat-proxy `GET /v1/settings`.
///
/// This is a blocking call intended for use in the early prefetch thread
/// (`std::thread::spawn`, no tokio runtime). Returns `None` on any error
/// so startup is never blocked by a settings fetch failure.
///
/// Retries up to 2 times (3 attempts total) on transient errors (5xx,
/// network). 4xx and parse errors are not retried.
/// Outcome of a blocking settings fetch. Distinguishes the three cases the
/// external-OTEL gate cares about (see [`crate::agent::mvp_agent`]).
#[derive(Debug)]
#[must_use]
#[non_exhaustive]
pub enum SettingsFetch {
/// Settings fetched and parsed; carries the policy that resolves the gate.
/// Boxed because `RemoteSettings` is large and the other variants are unit-sized.
Fetched(Box<crate::util::config::RemoteSettings>),
/// Credential unambiguously rejected (401): the remote policy will never reach
/// this leader, so the gate may open without waiting.
Rejected,
/// Transient/ambiguous (network, 5xx exhausted, 403/429/other 4xx, unparseable
/// 2xx): outcome unknown. Leave the gate closed (fail-closed), retry later.
Retry,
}
impl SettingsFetch {
/// For callers that only want the settings and treat every failure alike.
pub fn into_option(self) -> Option<crate::util::config::RemoteSettings> {
match self {
SettingsFetch::Fetched(s) => Some(*s),
SettingsFetch::Rejected | SettingsFetch::Retry => None,
}
}
}
/// Blocking settings fetch; makes up to
/// [`crate::http::SETTINGS_FETCH_MAX_ATTEMPTS`] attempts on transient failures.
pub fn fetch_settings_blocking(
cli_chat_proxy_base_url: &str,
auth: &GrokAuth,
alpha_test_key: Option<&str>,
) -> Option<crate::util::config::RemoteSettings> {
let client = crate::http::shared_blocking_client();
let url = format!("{}/settings", cli_chat_proxy_base_url);
for attempt in 0u64..3 {
) -> SettingsFetch {
fetch_settings_blocking_with_attempts(
cli_chat_proxy_base_url,
auth,
alpha_test_key,
crate::http::SETTINGS_FETCH_MAX_ATTEMPTS,
)
}
/// Settings-fetch core with a caller-chosen attempt budget. Private so the
/// attempt count stays out of the public API; tests use it to skip retry
/// backoff on the transient-failure paths.
fn fetch_settings_blocking_with_attempts(
cli_chat_proxy_base_url: &str,
auth: &GrokAuth,
alpha_test_key: Option<&str>,
max_attempts: u32,
) -> SettingsFetch {
let client = crate::http::shared_startup_blocking_client();
let url = format!("{cli_chat_proxy_base_url}/settings");
let max_attempts = max_attempts.max(1);
for attempt in 0u32..max_attempts {
if attempt > 0 {
std::thread::sleep(std::time::Duration::from_millis(500 * attempt));
std::thread::sleep(std::time::Duration::from_millis(500 * u64::from(attempt)));
}
let request =
add_cli_chat_proxy_headers_blocking(client.get(&url), auth, alpha_test_key, &url);
@ -577,11 +613,11 @@ pub fn fetch_settings_blocking(
Ok(resp) if resp.status().is_success() => match resp.json() {
Ok(settings) => {
tracing::debug!("Fetched remote settings from cli-chat-proxy");
return Some(settings);
return SettingsFetch::Fetched(Box::new(settings));
}
Err(e) => {
tracing::warn!(attempt, "Failed to parse settings response: {e}");
return None;
return SettingsFetch::Retry;
}
},
Ok(resp) if resp.status().is_server_error() => {
@ -592,9 +628,19 @@ pub fn fetch_settings_blocking(
);
continue;
}
Ok(resp) if resp.status() == reqwest::StatusCode::UNAUTHORIZED => {
tracing::warn!(
status = resp.status().as_u16(),
"Settings fetch rejected (401)"
);
return SettingsFetch::Rejected;
}
Ok(resp) => {
tracing::warn!(status = resp.status().as_u16(), "Failed to fetch settings");
return None;
tracing::warn!(
status = resp.status().as_u16(),
"Settings fetch failed (non-2xx)"
);
return SettingsFetch::Retry;
}
Err(e) => {
tracing::warn!(attempt, "Settings fetch network error: {e}");
@ -602,8 +648,8 @@ pub fn fetch_settings_blocking(
}
}
}
tracing::error!("Settings fetch failed after 3 attempts");
None
tracing::error!(max_attempts, "Settings fetch failed");
SettingsFetch::Retry
}
#[derive(Deserialize)]
struct LoginConfigResponse {
@ -720,7 +766,7 @@ pub(crate) fn fetch_models_blocking(
auth: Option<&GrokAuth>,
fetch_auth: crate::agent::models::ModelFetchAuth,
) -> Result<FetchModelsResult, BackendError> {
let client = crate::http::shared_blocking_client();
let client = crate::http::shared_startup_blocking_client();
let source = ListModelsEndpoint::from_endpoints(endpoints, fetch_auth);
let inference_base_url = endpoints.resolve_inference_base_url();
tracing::info!("Fetching models from {}", source.url);
@ -1177,6 +1223,54 @@ mod tests {
assert_eq!(h.user_id, None, "must not send x-userid");
assert_eq!(h.email, None, "must not send x-email");
}
/// Mock cli-chat-proxy serving `GET /settings` with a fixed status + body.
async fn start_settings_server(
status: StatusCode,
body: String,
) -> (String, tokio::task::JoinHandle<()>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port());
let app = Router::new().route(
"/settings",
get(move || {
let body = body.clone();
async move { (status, body) }
}),
);
let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
(base, handle)
}
/// `fetch_settings_blocking` maps each HTTP outcome to the [`SettingsFetch`]
/// variant the external-OTEL gate relies on; 401 is the only outcome that
/// yields `Rejected`, everything else non-2xx fails closed as `Retry`.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn settings_fetch_maps_status_to_outcome() {
let auth = GrokAuth::test_default();
let cases: [(StatusCode, &str, &str); 6] = [
(StatusCode::OK, "{}", "Fetched"),
(StatusCode::UNAUTHORIZED, "{}", "Rejected"),
(StatusCode::FORBIDDEN, "{}", "Retry"),
(StatusCode::TOO_MANY_REQUESTS, "{}", "Retry"),
(StatusCode::INTERNAL_SERVER_ERROR, "{}", "Retry"),
(StatusCode::OK, "not json", "Retry"),
];
for (status, body, expected) in cases {
let (base, server) = start_settings_server(status, body.to_string()).await;
let a = auth.clone();
let outcome = tokio::task::spawn_blocking(move || {
fetch_settings_blocking_with_attempts(&base, &a, None, 1)
})
.await
.unwrap();
server.abort();
let got = match outcome {
SettingsFetch::Fetched(_) => "Fetched",
SettingsFetch::Rejected => "Rejected",
SettingsFetch::Retry => "Retry",
};
assert_eq!(got, expected, "status {status}, body {body:?}");
}
}
#[derive(Debug, Default, Clone)]
struct SeenHeaders {
authorization: Option<String>,

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

@ -24,7 +24,7 @@ pub use chat_models_client::{
ChatModelsClient, ChatModelsError, ListModesResponse, Mode, ModeAvailability,
};
pub use client::{
BackendClient, BackendError, FetchModelsResult, FetchedBundle, fetch_bundle,
BackendClient, BackendError, FetchModelsResult, FetchedBundle, SettingsFetch, fetch_bundle,
fetch_login_device_flow, fetch_settings_blocking, fetch_subagent_bundle, share_url,
};
pub(crate) use client::{DEFAULT_CONTEXT_WINDOW, fetch_models_blocking, models_list_url};

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

@ -355,6 +355,8 @@ pub mod storage;
pub(crate) mod streaming_capture;
pub(crate) mod summary;
pub(crate) mod telemetry;
#[cfg(feature = "test-support")]
pub mod testkit;
pub mod tool_index;
pub(crate) mod turn_completion;
pub mod unified_list;

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

@ -372,6 +372,11 @@ pub enum PersistenceMsg {
/// Routed back through the persistence channel so the storage write
/// stays sequential with other summary.json mutations.
GeneratedTitle(String),
/// Enable remote writeback for a session created `Local` before remote
/// settings resolved (non-blocking startup); backfills its local history.
UpgradeToWriteback {
auth_manager: Arc<crate::auth::AuthManager>,
},
Flush,
/// Flush all pending writes, then signal the caller once the flush is complete.
/// Unlike `Flush` (fire-and-forget), this is a **sync barrier**: the caller's
@ -1591,6 +1596,9 @@ struct SessionPersistence {
pending_notification: Option<acp::SessionNotification>,
rx: mpsc::UnboundedReceiver<PersistenceMsg>,
remote_sync: Option<RemoteSync>,
/// True only for sessions created this run (not resumed); gates the
/// writeback backfill so a resumed, already-synced session isn't re-sent.
created_fresh: bool,
/// WebSocket-based relay sync for real-time session sharing.
/// This streams updates to the relay backend in addition to local persistence.
relay_sync: Option<crate::relay::RelaySync>,
@ -1710,6 +1718,53 @@ impl SessionPersistence {
}
}
/// Enable writeback for a session created `Local` before settings resolved:
/// build the sync and (for a fresh session) backfill its local-only history.
/// No-op once syncing, so a repeat upgrade is harmless.
async fn upgrade_to_writeback(&mut self, auth_manager: Arc<crate::auth::AuthManager>) {
if self.remote_sync.is_some() {
return;
}
// Flush the merge-pending notification so the backfill re-reads it.
self.flush_pending().await;
let persisted = match self.storage.load_session(&self.info).await {
Ok(persisted) => persisted,
Err(error) => {
tracing::warn!(%error, "writeback upgrade: failed to load session for backfill");
return;
}
};
let remote_sync = match init_remote_sync(
&persisted.summary,
StorageMode::Writeback,
Some(auth_manager),
) {
Ok(Some(remote_sync)) => remote_sync,
// ZDR team, or nothing to do: leave the session local-only.
Ok(None) => return,
Err(error) => {
tracing::warn!(%error, "writeback upgrade: remote sync init failed");
return;
}
};
// Fresh-only backfill; see `backfill_updates_to_sync`.
let backfilled =
backfill_updates_to_sync(self.created_fresh, persisted.updates, &remote_sync);
if self.created_fresh {
tracing::info!(
session_id = %self.info.id,
backfilled,
"writeback enabled after settings arrival; backfilled local-only history",
);
} else {
tracing::info!(
session_id = %self.info.id,
"writeback enabled for resumed session; forward-only, no backfill",
);
}
self.remote_sync = Some(remote_sync);
}
fn finish_pending_append(
notification: acp::SessionNotification,
result: Result<(), crate::session::storage::AppendUpdateError>,
@ -1806,6 +1861,9 @@ impl SessionPersistence {
spawn_worktree_touch(&self.info);
}
match msg {
PersistenceMsg::UpgradeToWriteback { auth_manager } => {
self.upgrade_to_writeback(auth_manager).await;
}
PersistenceMsg::Flush => {
self.flush_pending().await;
}
@ -2239,6 +2297,29 @@ fn collect_session_files_recursive(base: &Path, dir: &Path, files: &mut Vec<Copi
}
}
/// Queue a fresh session's local-only ACP history to `remote_sync` (xAI updates
/// are never synced), returning the count. Resumed sessions are forward-only:
/// their prior history may already be on the backend (which appends by content,
/// no per-message id), so re-sending would duplicate.
fn backfill_updates_to_sync(
created_fresh: bool,
updates: Vec<SessionUpdate>,
remote_sync: &RemoteSync,
) -> usize {
if !created_fresh {
return 0;
}
let mut backfilled = 0usize;
for update in updates {
if let SessionUpdate::Acp(notification) = update {
remote_sync.queue(*notification);
backfilled += 1;
}
}
remote_sync.flush();
backfilled
}
fn init_remote_sync(
summary: &Summary,
storage_mode: StorageMode,
@ -2432,6 +2513,7 @@ pub(crate) async fn new(
pending_notification: None,
rx,
remote_sync: remote_sync.clone(),
created_fresh: true,
relay_sync,
summary: crate::session::summary::SummaryGenerator::new(
crate::session::summary::SummaryConfig {
@ -2502,6 +2584,7 @@ pub async fn new_with_explicit_dir(
pending_notification: None,
rx,
remote_sync: None,
created_fresh: false,
relay_sync: None,
summary: crate::session::summary::SummaryGenerator::new(
crate::session::summary::SummaryConfig {
@ -2630,6 +2713,7 @@ pub(crate) async fn load(
pending_notification: None,
rx,
remote_sync: remote_sync.clone(),
created_fresh: false,
relay_sync,
summary: summary_gen,
registry_title_sync,
@ -2716,6 +2800,7 @@ pub(crate) async fn load_light(
pending_notification: None,
rx,
remote_sync: remote_sync.clone(),
created_fresh: false,
relay_sync,
summary: summary_gen,
registry_title_sync,

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

@ -32,6 +32,8 @@ fn test_actor_with_remote_sync(
pending_notification: None,
rx,
remote_sync,
// Resumed-style actor for these tests; upgrade backfill is fresh-only.
created_fresh: false,
relay_sync: None,
summary: crate::session::summary::SummaryGenerator::new(
crate::session::summary::SummaryConfig {
@ -64,6 +66,37 @@ fn neutral_update(info: &Info, text: &str) -> SessionUpdate {
SessionUpdate::Acp(Box::new(notification(info, text)))
}
#[tokio::test]
async fn writeback_backfill_is_fresh_only_and_acp_only() {
let info = Info {
id: acp::SessionId::new("wb-backfill"),
cwd: "/test".into(),
};
// Fresh session: every ACP update is queued to the writeback sync.
let (sync, mut observed) = RemoteSync::test_observer();
let updates = vec![neutral_update(&info, "a"), neutral_update(&info, "b")];
let n = backfill_updates_to_sync(true, updates, &sync);
assert_eq!(n, 2, "a fresh session backfills its full local ACP history");
for _ in 0..2 {
tokio::time::timeout(std::time::Duration::from_secs(1), observed.recv())
.await
.expect("backfilled notification not observed within 1s")
.expect("observer channel closed unexpectedly");
}
// Resumed session: nothing is backfilled (prior history may already be synced).
let (sync2, mut observed2) = RemoteSync::test_observer();
let n2 = backfill_updates_to_sync(false, vec![neutral_update(&info, "a")], &sync2);
assert_eq!(n2, 0, "a resumed session is forward-only, no backfill");
assert!(
tokio::time::timeout(std::time::Duration::from_millis(200), observed2.recv())
.await
.is_err(),
"resumed session must not re-send any prior history",
);
}
fn break_summary_writes(dir: &std::path::Path) {
let summary = dir.join("summary.json");
std::fs::remove_file(&summary).unwrap();

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

@ -1308,81 +1308,119 @@ pub(crate) struct RawChunkMetaPeek {
pub host_turn: Option<bool>,
}
/// Filter rewind dead branches from raw JSONL lines.
/// Skips parsing entirely when no rewind markers are present.
///
/// This is the canonical implementation of rewind dead-branch filtering,
/// used by both the initial replay and delta replay paths.
pub(crate) fn filter_rewind_lines<'a>(lines: Vec<&'a str>) -> Vec<&'a str> {
let has_rewinds = lines.iter().any(|l| l.contains(&*REWIND_MARKER));
if !has_rewinds {
return lines;
}
/// Role of one item in the rewind timeline, as seen by [`filter_rewind_by`].
enum RewindStep {
/// Rewind marker: truncate survivors back to `target`'s prompt boundary.
Rewind { target: usize },
/// User-message chunk opening (or continuing) a prompt run.
UserChunk { prompt_index: Option<usize> },
/// Anything else: kept, but ends the current user run.
Other,
}
let mut result: Vec<&str> = Vec::with_capacity(lines.len());
/// Shared rewind dead-branch filter. `classify` maps each item to its
/// [`RewindStep`]; the driver tracks prompt boundaries and, on a marker,
/// truncates survivors back to the target prompt. [`filter_rewind_lines`] and
/// [`filter_rewind_updates`] wrap this over raw JSONL and typed updates so the
/// two paths share one algorithm.
fn filter_rewind_by<T>(items: Vec<T>, classify: impl Fn(&T) -> RewindStep) -> Vec<T> {
let mut result: Vec<T> = Vec::with_capacity(items.len());
let mut prompt_starts: Vec<usize> = Vec::new();
let mut tracker = UserRunTurnTracker::new();
for line in &lines {
let (raw_params, is_xai) = if let Ok(env) = serde_json::from_str::<RawLinePeek<'_>>(line) {
let raw = env.params.map(|p| p.get()).unwrap_or(line);
let xai = env.method == Some(XAI_SESSION_UPDATE_METHOD);
(raw, xai)
} else {
(*line, false)
};
let peek = serde_json::from_str::<RawParamsPeek<'_>>(raw_params)
.ok()
.and_then(|p| p.update);
let tag = peek
.as_ref()
.map(|u| (u.session_update, u.target_prompt_index));
if is_xai
&& let Some((s, Some(target))) = tag.as_ref().map(|(s, t)| (*s, *t))
&& s == *REWIND_MARKER
{
let trunc = prompt_starts.get(target).copied().unwrap_or(result.len());
result.truncate(trunc);
prompt_starts.truncate(target);
tracker.on_non_user();
continue;
}
let is_host_turn = peek
.as_ref()
.and_then(|u| u.meta.as_ref())
.and_then(|m| m.host_turn)
.unwrap_or(false);
let is_user_chunk = !is_xai
&& !is_host_turn
&& tag
.as_ref()
.map(|(s, _)| *s == *USER_MESSAGE_CHUNK)
.unwrap_or(false);
if is_user_chunk {
let pi = peek.as_ref().and_then(|u| {
u.meta
.as_ref()
.and_then(|m| m.prompt_index.map(|v| v as usize))
});
if tracker.on_user_chunk(pi) {
prompt_starts.push(result.len());
for item in items {
match classify(&item) {
RewindStep::Rewind { target } => {
// Out-of-range target keeps every survivor: fold to `result.len()`.
let trunc = prompt_starts.get(target).copied().unwrap_or(result.len());
result.truncate(trunc);
prompt_starts.truncate(target);
tracker.on_non_user();
continue;
}
} else {
tracker.on_non_user();
RewindStep::UserChunk { prompt_index } => {
if tracker.on_user_chunk(prompt_index) {
prompt_starts.push(result.len());
}
}
RewindStep::Other => tracker.on_non_user(),
}
result.push(line);
result.push(item);
}
result
}
/// Classify a raw JSONL line by peeking at its tag and `_meta` without fully
/// deserializing the payload.
fn rewind_step_for_line(line: &str) -> RewindStep {
let (raw_params, is_xai) = if let Ok(env) = serde_json::from_str::<RawLinePeek<'_>>(line) {
let raw = env.params.map(|p| p.get()).unwrap_or(line);
(raw, env.method == Some(XAI_SESSION_UPDATE_METHOD))
} else {
(line, false)
};
let Some(u) = serde_json::from_str::<RawParamsPeek<'_>>(raw_params)
.ok()
.and_then(|p| p.update)
else {
return RewindStep::Other;
};
if is_xai
&& u.session_update == *REWIND_MARKER
&& let Some(target) = u.target_prompt_index
{
return RewindStep::Rewind { target };
}
let is_host_turn = u.meta.as_ref().and_then(|m| m.host_turn).unwrap_or(false);
if !is_xai && !is_host_turn && u.session_update == *USER_MESSAGE_CHUNK {
let prompt_index = u
.meta
.as_ref()
.and_then(|m| m.prompt_index.map(|v| v as usize));
return RewindStep::UserChunk { prompt_index };
}
RewindStep::Other
}
/// Classify a typed `SessionUpdate`.
fn rewind_step_for_update(update: &SessionUpdate) -> RewindStep {
if let SessionUpdate::Xai(n) = update
&& let crate::extensions::notification::SessionUpdate::RewindMarker {
target_prompt_index,
..
} = &n.update
{
return RewindStep::Rewind {
target: *target_prompt_index,
};
}
if is_acp_user_message_chunk(update) && !is_host_turn_update(update) {
return RewindStep::UserChunk {
prompt_index: acp_user_chunk_prompt_index(update),
};
}
RewindStep::Other
}
/// Filter rewind dead branches from raw JSONL lines.
///
/// Canonical raw-line rewind filter used by the initial and delta replay paths.
/// Skips parsing entirely when no rewind markers are present.
pub(crate) fn filter_rewind_lines(lines: Vec<&str>) -> Vec<&str> {
if !lines.iter().any(|l| l.contains(&*REWIND_MARKER)) {
return lines;
}
filter_rewind_by(lines, |line| rewind_step_for_line(line))
}
/// Filter rewind dead branches from typed `SessionUpdate` values.
///
/// This is the typed equivalent of [`filter_rewind_lines`] — same algorithm
/// (prompt-boundary tracking + truncation on `RewindMarker`) but operates on
/// fully-deserialized updates instead of raw JSON strings.
/// Typed equivalent of [`filter_rewind_lines`] over the same
/// [`filter_rewind_by`] driver, operating on fully-deserialized updates.
pub fn filter_rewind_updates(updates: Vec<SessionUpdate>) -> Vec<SessionUpdate> {
let has_rewinds = updates.iter().any(|u| {
matches!(
@ -1396,39 +1434,7 @@ pub fn filter_rewind_updates(updates: Vec<SessionUpdate>) -> Vec<SessionUpdate>
if !has_rewinds {
return updates;
}
let mut result: Vec<SessionUpdate> = Vec::with_capacity(updates.len());
let mut prompt_starts: Vec<usize> = Vec::new();
let mut tracker = UserRunTurnTracker::new();
for update in updates {
// Check for rewind marker — truncate back to the target prompt.
if let SessionUpdate::Xai(ref n) = update
&& let crate::extensions::notification::SessionUpdate::RewindMarker {
target_prompt_index,
..
} = &n.update
{
let trunc = prompt_starts
.get(*target_prompt_index)
.copied()
.unwrap_or(result.len());
result.truncate(trunc);
prompt_starts.truncate(*target_prompt_index);
tracker.on_non_user();
continue;
}
if is_acp_user_message_chunk(&update) && !is_host_turn_update(&update) {
if tracker.on_user_chunk(acp_user_chunk_prompt_index(&update)) {
prompt_starts.push(result.len());
}
} else {
tracker.on_non_user();
}
result.push(update);
}
result
filter_rewind_by(updates, rewind_step_for_update)
}
/// Strip `<fork-context>` and `<resume-context>` XML wrappers from user
@ -1457,21 +1463,14 @@ pub fn strip_context_wrappers(update: acp::SessionUpdate) -> acp::SessionUpdate
acp::SessionUpdate::UserMessageChunk(chunk)
}
/// Load session updates from disk, ready for replay or export.
///
/// This is the canonical way to get replay-ready typed updates from a session
/// ID. It:
/// 1. Locates the session directory via [`find_session_dir_by_id`]
/// 2. Opens `updates.jsonl` via [`UpdatesIterator`]
/// 3. Collects all parseable updates (skipping malformed lines)
/// 4. Filters rewind dead branches via [`filter_rewind_updates`]
/// 5. Strips `<fork-context>` / `<resume-context>` wrappers from user messages
/// via [`strip_context_wrappers`]
///
/// Returns `None` if the session is not found or has no `updates.jsonl`.
/// Returns only `SessionUpdate::Acp` updates (xAI-extension updates like
/// rewind markers and compaction signals are consumed by the filter and not
/// included in the output).
// Replay-loader family, all resolving through `replay_updates_path_in_dir` and
// reading through `for_each_replay_update_in_file`. Pick by need:
// - production, current grok home: `load_updates_for_replay`
// - production, streaming (bounded): `stream_replay_updates_at`
// - tests, explicit grok home: `load_updates_for_replay_at` (typed reference)
/// Load replay-ready typed ACP updates for a session, or `None` when the
/// session or its `updates.jsonl` is missing.
pub fn load_updates_for_replay(
session_id: &str,
) -> std::io::Result<Option<Vec<acp::SessionUpdate>>> {
@ -1480,14 +1479,53 @@ pub fn load_updates_for_replay(
else {
return Ok(None);
};
load_updates_for_replay_from_dir(&session_dir)
let Some(updates_path) = replay_updates_path_in_dir(&session_dir) else {
return Ok(None);
};
Ok(Some(collect_replay_updates(&updates_path)?))
}
/// Like [`load_updates_for_replay`], but resolves the session under a specific grok home.
/// Like [`load_updates_for_replay`], but resolves the session under a specific
/// grok home. Typed, materialize-all replay reader: collects every update into
/// owned `Vec`s. Production forwards replay through [`stream_replay_updates_at`]
/// to bound peak memory, so this has no production caller and is compiled only
/// for tests: the `testkit_synth_roundtrip` and `session_load_perf` parity
/// references and the in-crate relocation tests.
#[cfg(any(test, feature = "test-support"))]
pub fn load_updates_for_replay_at(
session_id: &str,
grok_home: &std::path::Path,
) -> std::io::Result<Option<Vec<acp::SessionUpdate>>> {
let Some(updates_path) = resolve_replay_updates_path(session_id, grok_home)? else {
return Ok(None);
};
Ok(Some(collect_replay_updates(&updates_path)?))
}
/// The session dir's `updates.jsonl` path if it exists, else `None`. Sole owner
/// of the "does this dir have a replayable updates file" gate.
fn replay_updates_path_in_dir(session_dir: &std::path::Path) -> Option<std::path::PathBuf> {
let updates_path = session_dir.join(UPDATES_FILE);
updates_path.exists().then_some(updates_path)
}
/// Collect every replay-ready ACP update from `updates_path` into a `Vec`, the
/// materializing counterpart of the streaming [`for_each_replay_update_in_file`].
fn collect_replay_updates(
updates_path: &std::path::Path,
) -> std::io::Result<Vec<acp::SessionUpdate>> {
let mut acp_updates: Vec<acp::SessionUpdate> = Vec::new();
for_each_replay_update_in_file(updates_path, |u| acp_updates.push(u))?;
Ok(acp_updates)
}
/// Resolve `updates.jsonl` for `session_id` under `grok_home`, or `None` when
/// the session directory or the file is missing. Shared by the typed
/// `load_updates_for_replay_at` and the streaming [`stream_replay_updates_at`].
fn resolve_replay_updates_path(
session_id: &str,
grok_home: &std::path::Path,
) -> std::io::Result<Option<std::path::PathBuf>> {
let sessions_root = grok_home.join("sessions");
let Some(session_dir) =
crate::session::persistence::find_persisted_session_dir_by_id_in_root_result(
@ -1497,45 +1535,103 @@ pub fn load_updates_for_replay_at(
else {
return Ok(None);
};
load_updates_for_replay_from_dir(&session_dir)
Ok(replay_updates_path_in_dir(&session_dir))
}
fn load_updates_for_replay_from_dir(
session_dir: &std::path::Path,
) -> std::io::Result<Option<Vec<acp::SessionUpdate>>> {
let updates_path = session_dir.join(UPDATES_FILE);
let Some(iter) = UpdatesIterator::open(&updates_path)? else {
return Ok(None);
/// Whether a replay stream forwarded any update. Gates the caller's
/// post-replay memory purge: `Empty` means nothing was reclaimable.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[must_use]
pub enum ReplayEmission {
Emitted,
Empty,
}
/// Invoke `f` once per replay-ready ACP update for a session under `grok_home`,
/// never building the full typed `Vec`. Reads the session's JSONL transcript
/// directly; a non-JSONL backend would need its own bounded replay.
///
/// Forking or resuming replays the inherited transcript. The typed load parsed
/// the whole file and copied it several times, so a large session briefly held
/// several times its size in live heap and a per-user memory cgroup OOM-killed
/// it. Streaming holds one typed update at a time, so peak drops to about the
/// file size.
///
/// `Empty` folds the missing-session, missing-file, and no-ACP-updates cases;
/// the typed `load_updates_for_replay_at` keeps them distinct (`Ok(None)` vs
/// `Ok(Some(vec![]))`) since it returns the parsed contents rather than a purge
/// signal.
///
/// The sink is infallible by design: replay only rehydrates UI scrollback, a
/// best-effort step, so failing to apply one update must neither abort the
/// stream nor surface an error. I/O errors from reading the file still
/// propagate via the `Result`.
pub fn stream_replay_updates_at<F: FnMut(acp::SessionUpdate)>(
session_id: &str,
grok_home: &std::path::Path,
f: F,
) -> std::io::Result<ReplayEmission> {
let Some(updates_path) = resolve_replay_updates_path(session_id, grok_home)? else {
return Ok(ReplayEmission::Empty);
};
let all: Vec<SessionUpdate> = iter.filter_map(|r| r.ok()).collect();
let filtered = filter_rewind_updates(all);
let acp_updates: Vec<acp::SessionUpdate> = filtered
.into_iter()
.filter_map(|u| match u {
SessionUpdate::Acp(notif) => Some(strip_context_wrappers(notif.update)),
SessionUpdate::Xai(_) => None,
})
.collect();
Ok(Some(acp_updates))
Ok(if for_each_replay_update_in_file(&updates_path, f)? {
ReplayEmission::Emitted
} else {
ReplayEmission::Empty
})
}
pub(crate) struct PreparedReplay<'a> {
// Rewind can drop earlier lines, so surviving lines are held until the end of
// the file; one `String` plus `&str` slices keeps that minimal. Output matches
// the typed load. Returns whether any ACP update was forwarded.
fn for_each_replay_update_in_file<F: FnMut(acp::SessionUpdate)>(
updates_path: &std::path::Path,
mut f: F,
) -> std::io::Result<bool> {
// Whole-file read is bounded by file size; only the forwarding is streamed.
let raw_contents = std::fs::read_to_string(updates_path)?;
let live: Vec<&str> = filter_rewind_lines(
raw_contents
.lines()
.filter(|l| !l.trim().is_empty())
.collect(),
);
let mut forwarded = false;
for line in live {
match SessionUpdateEnvelope::from_str(line) {
// Only ACP updates replay.
Ok(SessionUpdate::Acp(notif)) => {
forwarded = true;
f(strip_context_wrappers(notif.update));
}
// Xai extensions (rewind markers, compaction signals) are consumed
// by the filter and intentionally dropped (matching the typed load).
Ok(SessionUpdate::Xai(_)) => {}
// Best-effort: an unparseable line (e.g. a partially written trailing
// line) is skipped rather than aborting replay; the typed load drops
// it too. Logged for diagnostics.
Err(e) => tracing::debug!(error = %e, "skipping unparseable replay line"),
}
}
Ok(forwarded)
}
#[doc(hidden)]
pub struct PreparedReplay<'a> {
/// Rewind-filtered replay lines, each borrowed from the input transcript.
pub lines: Vec<&'a str>,
pub mark_replay: bool,
pub last_tokens: u64,
pub(crate) mark_replay: bool,
pub(crate) last_tokens: u64,
/// Highest `eventId` counter across all live (rewind-filtered) lines, used
/// to re-seed the process-global event counter on resume so post-load live
/// events keep monotonically increasing ids (see
/// [`crate::util::event_id::ensure_event_counter_at_least`]). `None` when no
/// line carried a parseable `eventId` (older shell).
pub max_event_seq: Option<u64>,
pub total_live: usize,
/// Replayed spawns with no matching finish (a rewind can drop the finish) —
pub(crate) max_event_seq: Option<u64>,
pub(crate) total_live: usize,
/// Replayed spawns with no matching finish (a rewind can drop the finish):
/// `(subagent_id, child_session_id)`, reconciled on load.
pub unfinished_subagents: Vec<(String, String)>,
pub(crate) unfinished_subagents: Vec<(String, String)>,
}
/// Unpaired spawns across the rewind-filtered timeline. Substring pre-filter
@ -1659,23 +1755,18 @@ fn line_has_event_id(line: &str, cursor_id: &str) -> bool {
line_event_id(line).as_deref() == Some(cursor_id)
}
/// Rewind-filter, resolve the reconnect cursor, drop redundant command catalogs,
/// and scan `totalTokens`. Pure data processing — no gateway, no async.
/// Rewind-filter, resolve the reconnect cursor, drop redundant command
/// catalogs, and scan `totalTokens`. Pure data processing, no I/O.
///
/// The cursor is resolved BEFORE dropping ACUs: ACUs carry `_meta.eventId` and the
/// post-load re-advertise is usually the *last* persisted event, so an idle client
/// commonly reconnects with an ACU's eventId as its cursor. Resolving against the
/// ACU-inclusive set keeps incremental reconnect cheap instead of a full replay.
pub(crate) fn prepare_replay_lines<'a>(
raw_contents: &'a str,
cursor: Option<&str>,
) -> PreparedReplay<'a> {
let filtered = filter_rewind_lines(
raw_contents
.lines()
.filter(|l| !l.trim().is_empty())
.collect(),
);
/// The cursor is resolved before dropping ACUs, because an idle client often
/// reconnects with an ACU's `eventId` as its cursor; resolving against the
/// ACU-inclusive set keeps reconnect incremental instead of a full replay.
///
/// `#[doc(hidden)] pub` (not stable API): production replay uses it, and the
/// session-load memory test drives it to check the peek stays zero-copy.
#[doc(hidden)]
pub fn prepare_replay_lines<'a>(contents: &'a str, cursor: Option<&str>) -> PreparedReplay<'a> {
let filtered = filter_rewind_lines(contents.lines().filter(|l| !l.trim().is_empty()).collect());
// Highest `eventId` counter across all live (rewind-filtered) lines, used to
// re-seed the process-global event counter on resume so post-load live events
@ -2994,6 +3085,178 @@ mod tests {
assert!(result[2].contains("final"));
}
/// The raw-line filter and the typed filter must truncate an identical
/// rewind timeline to the same surviving updates, in the same order.
#[test]
fn filter_rewind_lines_and_updates_agree() {
let u1 = acp_envelope(
r#"{"sessionUpdate":"user_message_chunk","content":{"type":"text","text":"p1"}}"#,
);
let a1 = acp_envelope(
r#"{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"r1"}}"#,
);
let u2 = acp_envelope(
r#"{"sessionUpdate":"user_message_chunk","content":{"type":"text","text":"p2"}}"#,
);
let a2 = acp_envelope(
r#"{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"r2"}}"#,
);
let rw1 = xai_envelope(
r#"{"sessionUpdate":"rewind_marker","target_prompt_index":2,"created_at":"2024-01-01"}"#,
);
let u3 = acp_envelope(
r#"{"sessionUpdate":"user_message_chunk","content":{"type":"text","text":"p3"}}"#,
);
let a3 = acp_envelope(
r#"{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"r3"}}"#,
);
let rw2 = xai_envelope(
r#"{"sessionUpdate":"rewind_marker","target_prompt_index":1,"created_at":"2024-01-01"}"#,
);
let u4 = acp_envelope(
r#"{"sessionUpdate":"user_message_chunk","content":{"type":"text","text":"final"}}"#,
);
let lines = vec![
u1.as_str(),
a1.as_str(),
u2.as_str(),
a2.as_str(),
rw1.as_str(),
u3.as_str(),
a3.as_str(),
rw2.as_str(),
u4.as_str(),
];
let ser = |u: &SessionUpdate| serde_json::to_string(u).unwrap();
let via_lines: Vec<String> = filter_rewind_lines(lines.clone())
.iter()
.map(|l| ser(&SessionUpdateEnvelope::from_str(l).unwrap()))
.collect();
let typed: Vec<SessionUpdate> = lines
.iter()
.map(|l| SessionUpdateEnvelope::from_str(l).unwrap())
.collect();
let via_updates: Vec<String> = filter_rewind_updates(typed).iter().map(ser).collect();
assert_eq!(via_lines, via_updates);
}
/// An out-of-range rewind target folds to `result.len()` (the
/// `unwrap_or(result.len())` branch in `filter_rewind_by`), so truncation is
/// a no-op and every survivor is kept.
#[test]
fn filter_rewind_out_of_range_target_keeps_all() {
let u1 = acp_envelope(
r#"{"sessionUpdate":"user_message_chunk","content":{"type":"text","text":"p1"}}"#,
);
let a1 = acp_envelope(
r#"{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"r1"}}"#,
);
// Only prompt index 0 exists; target 5 is out of range.
let rw = xai_envelope(
r#"{"sessionUpdate":"rewind_marker","target_prompt_index":5,"created_at":"2024-01-01"}"#,
);
let u2 = acp_envelope(
r#"{"sessionUpdate":"user_message_chunk","content":{"type":"text","text":"p2"}}"#,
);
let lines = vec![u1.as_str(), a1.as_str(), rw.as_str(), u2.as_str()];
let result = filter_rewind_lines(lines);
// Marker is dropped; the three ACP survivors remain in order.
assert_eq!(result.len(), 3);
assert!(result[0].contains("p1"));
assert!(result[1].contains("r1"));
assert!(result[2].contains("p2"));
}
/// A session with no `updates.jsonl` streams nothing, so the emission gate
/// reports `Empty` and forwards no updates.
#[test]
fn stream_replay_updates_at_missing_session_is_empty() {
let grok_home = tempfile::tempdir().unwrap();
std::fs::create_dir_all(grok_home.path().join("sessions")).unwrap();
let mut count = 0usize;
let emission =
stream_replay_updates_at("does-not-exist", grok_home.path(), |_| count += 1).unwrap();
assert_eq!(emission, ReplayEmission::Empty);
assert_eq!(count, 0);
}
/// A resolvable session whose `updates.jsonl` cannot be read surfaces the
/// error rather than folding to `Empty`, so the caller logs a real fault
/// instead of mistaking it for an absent transcript. (The path is a
/// directory, which `read_to_string` rejects.)
#[test]
fn stream_replay_updates_at_surfaces_read_errors() {
let grok_home = tempfile::tempdir().unwrap();
let session_dir = grok_home.path().join("sessions").join("cwd").join("sess");
std::fs::create_dir_all(&session_dir).unwrap();
std::fs::write(session_dir.join(SUMMARY_FILE), "{}").unwrap();
std::fs::create_dir(session_dir.join(UPDATES_FILE)).unwrap();
let result = stream_replay_updates_at("sess", grok_home.path(), |_| {});
assert!(
result.is_err(),
"read fault must surface, not fold to Empty: {result:?}"
);
}
/// End-to-end: the streaming core (`for_each_replay_update_in_file`, what
/// `stream_replay_updates_at` wraps) applies rewind over a real file and
/// yields the same survivors as the typed parse-all path.
#[test]
fn streaming_replay_applies_rewind_like_the_typed_path() {
let u1 = acp_envelope(
r#"{"sessionUpdate":"user_message_chunk","content":{"type":"text","text":"p1"}}"#,
);
let a1 = acp_envelope(
r#"{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"r1"}}"#,
);
let u2 = acp_envelope(
r#"{"sessionUpdate":"user_message_chunk","content":{"type":"text","text":"p2"}}"#,
);
// Rewind to prompt 1 drops p2.
let rw = xai_envelope(
r#"{"sessionUpdate":"rewind_marker","target_prompt_index":1,"created_at":"2024-01-01"}"#,
);
let u3 = acp_envelope(
r#"{"sessionUpdate":"user_message_chunk","content":{"type":"text","text":"final"}}"#,
);
let raw = format!("{u1}\n{a1}\n{u2}\n{rw}\n{u3}\n");
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join(UPDATES_FILE);
std::fs::write(&path, &raw).unwrap();
let mut streamed = Vec::new();
let forwarded = for_each_replay_update_in_file(&path, |u| streamed.push(u)).unwrap();
assert!(forwarded);
// Typed reference: parse all, rewind-filter, map ACP survivors.
let typed: Vec<SessionUpdate> = raw
.lines()
.map(|l| SessionUpdateEnvelope::from_str(l).unwrap())
.collect();
let reference: Vec<acp::SessionUpdate> = filter_rewind_updates(typed)
.into_iter()
.filter_map(|u| match u {
SessionUpdate::Acp(notif) => Some(strip_context_wrappers(notif.update)),
SessionUpdate::Xai(_) => None,
})
.collect();
let ser = |u: &acp::SessionUpdate| serde_json::to_string(u).unwrap();
assert_eq!(
streamed.iter().map(ser).collect::<Vec<_>>(),
reference.iter().map(ser).collect::<Vec<_>>(),
);
}
// ── prepare_replay_lines tests ───────────────────────────────────────────
/// Envelope with _meta at the params level (where the real agent puts it).

127
crates/codegen/xai-grok-shell/src/session/testkit/e2e.rs Normal file
View file

@ -0,0 +1,127 @@
//! In-process `session/load` harness: a real `MvpAgent` wired to a client over
//! ACP duplex pipes, so a test can time a real load round-trip without a
//! subprocess.
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};
use agent_client_protocol::{self as acp};
use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
use xai_acp_lib::{
AcpAgentGatewayReceiver as GatewayReceiver, AcpAgentGatewaySender as GatewaySender,
LineBufferedRead,
};
use crate::agent::config::Config as AgentConfig;
use crate::agent::mvp_agent::MvpAgent;
const DUPLEX_BUFFER_BYTES: usize = 16 * 1024 * 1024;
const INIT_TIMEOUT: Duration = Duration::from_secs(60);
const LOAD_TIMEOUT: Duration = Duration::from_secs(180);
/// A completed `session/load` over the shared in-process harness. `client_conn`
/// is returned so the caller keeps the connection alive for any post-load
/// notifications (e.g. the re-advertise) it still wants to observe.
pub struct LoadedAgent {
pub client_conn: acp::ClientSideConnection,
pub load_started: Instant,
pub load_elapsed: Duration,
}
/// Stand up a real `MvpAgent` over in-process ACP pipes wired to `client`, run
/// the initialize and authenticate handshake, then time one `session/load`
/// round-trip. Must run inside a `LocalSet`, since it spawns local tasks.
pub async fn load_session_via_agent<C: acp::Client + 'static>(
client: C,
client_type: &str,
session_id: acp::SessionId,
cwd: PathBuf,
) -> LoadedAgent {
let agent_config = AgentConfig::default();
let auth_manager = Arc::new(agent_config.create_auth_manager());
let (gw_tx, gw_rx) = tokio::sync::mpsc::unbounded_channel();
let gateway = GatewaySender::new(gw_tx);
let agent = MvpAgent::new(gateway, &agent_config, auth_manager, None).expect("valid config");
let (c2a_a, c2a_b) = tokio::io::duplex(DUPLEX_BUFFER_BYTES);
let (a2c_a, a2c_b) = tokio::io::duplex(DUPLEX_BUFFER_BYTES);
// Agent side.
let agent_incoming = LineBufferedRead::spawn_local(c2a_b.compat());
let (agent_conn, agent_io) =
acp::AgentSideConnection::new(agent, a2c_a.compat_write(), agent_incoming, |fut| {
tokio::task::spawn_local(fut);
});
tokio::task::spawn_local(
GatewayReceiver::new(gw_rx, agent_conn)
.with_on_meta(xai_file_utils::trace_context::span_from_meta_traceparent)
.run(),
);
tokio::task::spawn_local(agent_io);
// Client side.
let client_incoming = LineBufferedRead::spawn_local(a2c_b.compat());
let (client_conn, client_io) =
acp::ClientSideConnection::new(client, c2a_a.compat_write(), client_incoming, |fut| {
tokio::task::spawn_local(fut);
});
tokio::task::spawn_local(client_io);
use acp::Agent as _;
let init = tokio::time::timeout(
INIT_TIMEOUT,
client_conn.initialize(
acp::InitializeRequest::new(acp::ProtocolVersion::V1)
.client_capabilities(
acp::ClientCapabilities::new()
.fs(acp::FileSystemCapabilities::new())
.terminal(false),
)
.meta(
serde_json::json!({
"startupHints": { "nonInteractive": true, "skipGitStatus": true, "skipProjectLayout": true },
"clientType": client_type,
"clientVersion": "0.0-test",
})
.as_object()
.cloned(),
),
),
)
.await
.expect("initialize timed out")
.expect("initialize failed");
// Best-effort auth: the mock backend accepts anything and `load` does not
// require a prior success, so ignore any failure here.
if let Some(method) = init
.auth_methods
.iter()
.find(|m| &*m.id().0 == "xai.api_key")
{
let _ = client_conn
.authenticate(
acp::AuthenticateRequest::new(method.id().clone())
.meta(serde_json::json!({ "headless": true }).as_object().cloned()),
)
.await;
}
let load_started = Instant::now();
tokio::time::timeout(
LOAD_TIMEOUT,
client_conn.load_session(acp::LoadSessionRequest::new(session_id, cwd)),
)
.await
.expect("session/load timed out (>180s)")
.expect("session/load failed");
let load_elapsed = load_started.elapsed();
LoadedAgent {
client_conn,
load_started,
load_elapsed,
}
}

9
crates/codegen/xai-grok-shell/src/session/testkit/mod.rs Normal file
View file

@ -0,0 +1,9 @@
//! Session synthesis and in-process e2e harness for the load-perf and fork
//! bench tests.
//!
//! Lives in `xai-grok-shell` (feature `test-support`) rather than
//! `xai-grok-test-support` because synthesis drives the real
//! `JsonlStorageAdapter`; the reverse dependency would be circular.
pub mod e2e;
pub mod synth;

77
crates/codegen/xai-grok-shell/src/session/testkit/synth/bench.rs Normal file
View file

@ -0,0 +1,77 @@
//! Adapter-driven session synthesis for benches: appends realistic turns through
//! the real `JsonlStorageAdapter` until `updates.jsonl` reaches a byte target,
//! so fork/copy benchmarks measure production-shaped data.
use std::path::Path;
use agent_client_protocol::{self as acp};
use crate::session::info::Info;
use crate::session::storage::{JsonlStorageAdapter, SessionUpdate, StorageAdapter};
const AGENT_CHUNKS_PER_TURN: usize = 8;
/// Stands in for a large tool result, the dominant byte source in real
/// sessions. Emitted as an agent message chunk so the byte and line shape match
/// production rather than the `ToolCall` kind.
const BULKY_CHUNK_BYTES: usize = 4096;
fn turn_updates(info: &Info, turn: usize) -> Vec<SessionUpdate> {
let text =
|s: String| acp::ContentChunk::new(acp::ContentBlock::Text(acp::TextContent::new(s)));
let notify =
|u| SessionUpdate::Acp(Box::new(acp::SessionNotification::new(info.id.clone(), u)));
let mut updates = vec![notify(acp::SessionUpdate::UserMessageChunk(text(format!(
"prompt {turn}: check the build and summarize failures"
))))];
for i in 0..AGENT_CHUNKS_PER_TURN {
updates.push(notify(acp::SessionUpdate::AgentMessageChunk(text(format!(
"agent chunk {turn}/{i}: analyzing module {i} for regressions and drafting a fix plan"
)))));
}
updates.push(notify(acp::SessionUpdate::AgentMessageChunk(text(
format!("bulky chunk {turn}: {}", "x".repeat(BULKY_CHUNK_BYTES)),
))));
updates
}
/// Build a session dir under `root` whose `updates.jsonl` is at least
/// `target_bytes`, appending realistic mixed updates through the real adapter.
///
/// Synchronous (drives the async adapter on its own current-thread runtime) so
/// Criterion benches can call it directly outside an async context.
pub fn synthesize_to_target_bytes(root: &Path, target_bytes: u64) -> Info {
let adapter = JsonlStorageAdapter::with_root(root.to_path_buf());
let info = Info {
id: acp::SessionId::new("fork-bench-src"),
cwd: "/bench/workspace".to_string(),
};
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("bench runtime");
rt.block_on(async {
adapter
.init_session(&info, acp::ModelId::new("bench-model"))
.await
.expect("init session");
let updates_path = adapter.updates_file_path(&info).expect("updates path");
let mut turn = 0usize;
loop {
for update in turn_updates(&info, turn) {
adapter.append_update(&info, &update).await.expect("append");
}
turn += 1;
// Stat every 32 turns; sizes only grow. A persistent stat failure
// panics here rather than spinning the append loop forever.
if turn.is_multiple_of(32)
&& std::fs::metadata(&updates_path)
.expect("stat updates.jsonl")
.len()
>= target_bytes
{
break;
}
}
});
info
}

10
crates/codegen/xai-grok-shell/src/session/testkit/synth/mod.rs Normal file
View file

@ -0,0 +1,10 @@
//! On-disk session synthesis. [`replay`] writes
//! `updates.jsonl`/`rewind_points.jsonl` envelopes directly for exact
//! ACU/rewind control; [`bench`] appends through the real storage adapter up to
//! a byte target for fork/copy benchmarks.
pub mod bench;
pub mod replay;
pub use bench::synthesize_to_target_bytes;
pub use replay::{SessionSpec, locate_session_dir, prepare_session, sid, write_rewind_jsonl};

270
crates/codegen/xai-grok-shell/src/session/testkit/synth/replay.rs Normal file
View file

@ -0,0 +1,270 @@
//! Session synthesis for replay/load tests: writes
//! `updates.jsonl`/`rewind_points.jsonl` directly, for exact control over ACU
//! redundancy and rewind points.
use std::path::{Path, PathBuf};
use agent_client_protocol::{self as acp};
use xai_grok_workspace::session::file_state::{FileSnapshot, FlexiblePath, RewindPoint};
use crate::session::info::Info;
use crate::session::storage::{JsonlStorageAdapter, StorageAdapter};
fn parse_or<T: std::str::FromStr>(key: &str, found: Option<String>, default: T) -> T {
let Some(text) = found else {
return default;
};
match text.parse() {
Ok(value) => value,
Err(_) => {
eprintln!("[testkit] ignoring unparseable {key}={text:?}; using default");
default
}
}
}
/// Generation parameters; fields double as the defaults for
/// [`SessionSpec::from_env_prefixed`].
pub struct SessionSpec {
pub turns: usize,
/// `available_commands_update`s persisted per turn: the redundant catalog
/// a real session re-advertises on every skill/subagent boundary.
pub acu_per_turn: usize,
pub catalog_commands: usize,
pub catalog_desc_len: usize,
pub agent_chunks_per_turn: usize,
pub agent_chunk_len: usize,
pub rewind_points: usize,
pub files_per_rewind: usize,
pub file_content_len: usize,
}
impl Default for SessionSpec {
/// Baseline yielding a ~20 MB `updates.jsonl`; callers override the knobs
/// they scale up.
fn default() -> Self {
Self {
turns: 60,
acu_per_turn: 15,
catalog_commands: 64,
catalog_desc_len: 320,
agent_chunks_per_turn: 8,
agent_chunk_len: 2000,
rewind_points: 20,
files_per_rewind: 20,
file_content_len: 4000,
}
}
}
impl SessionSpec {
/// Read `<prefix>_*` env overrides on top of `defaults`, scaling `turns` and
/// `rewind_points` by `<prefix>_SCALE`.
pub fn from_env_prefixed(prefix: &str, defaults: Self) -> Self {
Self::from_lookup(prefix, defaults, |key| std::env::var(key).ok())
}
/// [`from_env_prefixed`] with an injectable lookup, so the override and
/// scale arithmetic is unit-testable without touching process env.
fn from_lookup(prefix: &str, defaults: Self, get: impl Fn(&str) -> Option<String>) -> Self {
let val = |name: &str, default| {
let key = format!("{prefix}_{name}");
parse_or(&key, get(&key), default)
};
let scale = val("SCALE", 1usize).max(1);
Self {
turns: val("TURNS", defaults.turns) * scale,
acu_per_turn: val("ACU_PER_TURN", defaults.acu_per_turn),
catalog_commands: val("CATALOG_COMMANDS", defaults.catalog_commands),
catalog_desc_len: val("CATALOG_DESC_LEN", defaults.catalog_desc_len),
agent_chunks_per_turn: val("AGENT_CHUNKS_PER_TURN", defaults.agent_chunks_per_turn),
agent_chunk_len: val("AGENT_CHUNK_LEN", defaults.agent_chunk_len),
rewind_points: val("REWIND_POINTS", defaults.rewind_points) * scale,
files_per_rewind: val("FILES_PER_REWIND", defaults.files_per_rewind),
file_content_len: val("FILE_CONTENT_LEN", defaults.file_content_len),
}
}
}
fn filler(n: usize) -> String {
const WORDS: &[&str] = &[
"alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "hotel", "india",
"juliet", "kilo", "lima", "mike", "november", "oscar", "papa", "quebec", "romeo",
];
let mut s = String::with_capacity(n + 8);
let mut i = 0usize;
while s.len() < n {
s.push_str(WORDS[i % WORDS.len()]);
s.push(' ');
i += 1;
}
s.truncate(n);
s
}
pub fn sid(session_id: &str) -> acp::SessionId {
acp::SessionId::new(session_id.to_string())
}
fn text_chunk(text: String) -> acp::ContentChunk {
acp::ContentChunk::new(acp::ContentBlock::Text(acp::TextContent::new(text)))
}
/// One large `AvailableCommandsUpdate`: the redundant catalog re-persisted
/// thousands of times in the real session.
fn available_commands_update(spec: &SessionSpec) -> acp::SessionUpdate {
let desc = filler(spec.catalog_desc_len);
let commands: Vec<acp::AvailableCommand> = (0..spec.catalog_commands)
.map(|i| {
acp::AvailableCommand::new(format!("command-number-{i:03}"), desc.clone()).input(Some(
acp::AvailableCommandInput::Unstructured(acp::UnstructuredCommandInput::new(
"[optional arguments here]".to_string(),
)),
))
})
.collect();
acp::SessionUpdate::AvailableCommandsUpdate(acp::AvailableCommandsUpdate::new(commands))
}
fn envelope_line(session_id: &str, update: acp::SessionUpdate) -> String {
let update_val = serde_json::to_value(&update).expect("serialize update");
let params = serde_json::json!({
"sessionId": session_id,
"update": update_val,
});
let envelope = serde_json::json!({
"timestamp": 0u64,
"method": "session/update",
"params": params,
});
serde_json::to_string(&envelope).expect("serialize envelope")
}
fn write_updates_jsonl(path: &Path, session_id: &str, spec: &SessionSpec) {
let mut out = String::new();
for turn in 0..spec.turns {
out.push_str(&envelope_line(
session_id,
acp::SessionUpdate::UserMessageChunk(text_chunk(format!(
"user prompt for turn {turn}"
))),
));
out.push('\n');
for _ in 0..spec.acu_per_turn {
out.push_str(&envelope_line(session_id, available_commands_update(spec)));
out.push('\n');
}
for _ in 0..spec.agent_chunks_per_turn {
out.push_str(&envelope_line(
session_id,
acp::SessionUpdate::AgentMessageChunk(text_chunk(filler(spec.agent_chunk_len))),
));
out.push('\n');
}
}
std::fs::write(path, out).expect("write updates.jsonl");
}
pub fn write_rewind_jsonl(path: &Path, spec: &SessionSpec) {
let mut out = String::new();
for p in 0..spec.rewind_points {
let mut rp = RewindPoint::new(p);
for f in 0..spec.files_per_rewind {
let fp =
FlexiblePath::Absolute(PathBuf::from(format!("/repo/src/module_{p}/file_{f}.rs")));
rp.add_snapshot(FileSnapshot::new_flexible(
fp.clone(),
Some(filler(spec.file_content_len)),
));
rp.set_after_snapshot(FileSnapshot::new_flexible(
fp,
Some(filler(spec.file_content_len + 64)),
));
}
out.push_str(&serde_json::to_string(&rp).expect("serialize rewind point"));
out.push('\n');
}
std::fs::write(path, out).expect("write rewind_points.jsonl");
}
/// Locate `<root>/sessions/<encoded cwd>/<id>` without the crate-internal cwd encoder.
pub fn locate_session_dir(root: &Path, id: &str) -> PathBuf {
let sessions = root.join("sessions");
for entry in std::fs::read_dir(&sessions).expect("read sessions dir") {
let entry = entry.expect("read sessions dir entry");
let candidate = entry.path().join(id);
if candidate.is_dir() {
return candidate;
}
}
panic!(
"could not locate session dir for {id} under {}",
sessions.display()
);
}
/// Synthesize a session on disk under `root` for working dir `cwd`: the summary
/// through the production storage adapter, then `updates.jsonl` and
/// `rewind_points.jsonl` written directly. Returns the `Info` and its directory.
pub async fn prepare_session(root: &Path, cwd: &Path, spec: &SessionSpec) -> (Info, PathBuf) {
let adapter = JsonlStorageAdapter::with_root(root.to_path_buf());
let id = uuid::Uuid::new_v4().to_string();
let info = Info {
id: sid(&id),
cwd: cwd.to_string_lossy().to_string(),
};
adapter
.init_session(&info, acp::ModelId::new("test-model"))
.await
.expect("init_session");
let dir = locate_session_dir(root, &id);
write_updates_jsonl(&dir.join("updates.jsonl"), &id, spec);
write_rewind_jsonl(&dir.join("rewind_points.jsonl"), spec);
(info, dir)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn filler_is_exactly_n_bytes() {
assert_eq!(filler(0).len(), 0);
assert_eq!(filler(100).len(), 100);
assert_eq!(filler(4096).len(), 4096);
}
#[test]
fn from_lookup_falls_back_to_defaults_when_absent() {
let d = SessionSpec::default();
let spec = SessionSpec::from_lookup("P", SessionSpec::default(), |_| None);
assert_eq!(spec.turns, d.turns);
assert_eq!(spec.rewind_points, d.rewind_points);
assert_eq!(spec.agent_chunks_per_turn, d.agent_chunks_per_turn);
}
#[test]
fn from_lookup_applies_overrides_and_scale() {
let env = std::collections::HashMap::from([
("P_TURNS".to_string(), "10".to_string()),
("P_SCALE".to_string(), "3".to_string()),
("P_FILES_PER_REWIND".to_string(), "not_a_number".to_string()),
]);
let d = SessionSpec::default();
let spec = SessionSpec::from_lookup("P", SessionSpec::default(), |k| env.get(k).cloned());
assert_eq!(spec.turns, 10 * 3, "override is multiplied by SCALE");
assert_eq!(
spec.rewind_points,
d.rewind_points * 3,
"SCALE also scales rewind_points"
);
assert_eq!(
spec.acu_per_turn, d.acu_per_turn,
"an unscaled field keeps its default"
);
assert_eq!(
spec.files_per_rewind, d.files_per_rewind,
"an unparseable override falls back to the (unscaled) default"
);
}
}

95
crates/codegen/xai-grok-shell/src/util/config/campaigns.rs
View file

@ -237,6 +237,59 @@ pub fn load_effective_config_disk_only() -> std::io::Result<toml::Value> {
Ok(ConfigLayers::load()?.effective_config_disk_only())
}
/// The effective `models.default` while an **active** campaign drives it, plus
/// the pre-campaign base value it overrode.
pub struct CampaignModelsDefault {
/// The campaign-nudged default model.
pub value: String,
/// The pre-campaign base `models.default` (`None` when the user had none).
pub pre_campaign: Option<String>,
}
/// Resolve [`CampaignModelsDefault`] fresh from the config layers, the remote
/// campaign cache, and the on-disk dismiss state.
///
/// `None` unless an active (non-dismissed, kill-switch-respecting,
/// requirements-losing) campaign changes the effective `models.default`.
/// Session creation uses this to apply a campaign to `/new` even when remote
/// settings arrived only after boot: the `ModelsManager`'s `current_model_id`
/// was resolved pre-campaign, and a campaign-only flip deliberately never
/// re-targets it (see `ModelsManager::apply_config`), so `/new` re-evaluates
/// here instead.
///
/// Reading the dismiss state fresh makes a `/model` pick win instantly:
/// [`persist_user_choice`] records the dismissal before the config write, so
/// the very next `/new` resolves campaign-free.
pub fn campaign_driven_models_default() -> Option<CampaignModelsDefault> {
let layers = ConfigLayers::load().ok()?;
campaign_driven_models_default_from(&layers, &cached_remote_campaigns(), &load_dismissed_ids())
}
/// Env-free resolution core of [`campaign_driven_models_default`] (unit-testable
/// without touching `GROK_HOME` / the process-global cache).
fn campaign_driven_models_default_from(
layers: &ConfigLayers,
remote: &[CampaignEntry],
dismissed: &HashSet<String>,
) -> Option<CampaignModelsDefault> {
let base = layers.effective_config_base();
let active = resolve_active_campaigns_from_layers(layers, &base, remote, dismissed);
if active.is_empty() {
return None;
}
let mut effective = base.clone();
layers.apply_campaign_overrides(&mut effective, &active);
let base_value = read_path(&base, MODELS_DEFAULT_PATH);
let value = read_path(&effective, MODELS_DEFAULT_PATH);
if value == base_value {
return None;
}
Some(CampaignModelsDefault {
value: as_string(value)?,
pre_campaign: as_string(base_value),
})
}
/// Read the value at `path` from an effective-config tree.
fn read_path(tree: &toml::Value, path: PatchPath) -> Option<toml::Value> {
let mut cur = tree;
@ -513,6 +566,48 @@ mod tests {
);
}
/// `campaign_driven_models_default_from` tracks remote entries and
/// dismissals: `Some` while the campaign is active, `None` the instant its
/// dismissal lands, so a `/new` right after a `/model` pick never re-nudges.
#[test]
#[serial]
fn campaign_driven_models_default_tracks_remote_and_dismissals() {
let _over = EnvGuard::unset("GROK_CAMPAIGNS_OVERRIDE");
let _kill = EnvGuard::unset("GROK_CAMPAIGNS");
let layers = ConfigLayers {
user: toml::from_str("[models]\ndefault = \"config-model\"\n").unwrap(),
..Default::default()
};
let remote = vec![CampaignEntry {
id: "t-models-nudge".into(),
patch: models_default_patch("campaign-model"),
}];
let nudge = campaign_driven_models_default_from(&layers, &remote, &HashSet::new())
.expect("active campaign drives the default");
assert_eq!(nudge.value, "campaign-model");
assert_eq!(nudge.pre_campaign.as_deref(), Some("config-model"));
// A dismissal (what a `/model` pick records first) deactivates the
// nudge for the very next resolution.
let dismissed: HashSet<String> = ["t-models-nudge".to_string()].into_iter().collect();
assert!(
campaign_driven_models_default_from(&layers, &remote, &dismissed).is_none(),
"a dismissed campaign must not nudge"
);
// A campaign that loses to a requirements pin never reports
// campaign-driven.
let mut pinned = layers.clone();
pinned.user_requirements =
Some(toml::from_str("[models]\ndefault = \"config-model\"\n").unwrap());
assert!(
campaign_driven_models_default_from(&pinned, &remote, &HashSet::new()).is_none(),
"a requirements-pinned default must not report campaign-driven"
);
}
/// `GROK_CAMPAIGNS_OVERRIDE="[]"` replaces all sources with nothing — even
/// layer + remote campaigns resolve to empty.
#[test]

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

@ -14,8 +14,9 @@ mod worktree;
pub use announcements::*;
pub use campaigns::{
load_effective_config, load_effective_config_disk_only, persist_models_default,
remote_campaigns_from_settings, set_remote_campaigns_from_settings, sync_campaign_fields,
CampaignModelsDefault, campaign_driven_models_default, load_effective_config,
load_effective_config_disk_only, persist_models_default, remote_campaigns_from_settings,
set_remote_campaigns_from_settings, sync_campaign_fields,
};
pub use hints::*;
pub use load::*;

186
crates/codegen/xai-grok-shell/tests/session_fork_replay_memory.rs Normal file
View file

@ -0,0 +1,186 @@
//! Regression guard for the fork/resume replay OOM: a counting allocator checks
//! the streaming load ([`stream_replay_updates_at`]) peaks far below the old
//! parse-all load and forwards identical content.
#![cfg(unix)]
use std::alloc::{GlobalAlloc, Layout, System};
use std::path::Path;
use std::sync::atomic::{AtomicUsize, Ordering};
use agent_client_protocol as acp;
use tempfile::TempDir;
use xai_grok_shell::session::storage::{
ReplayEmission, SessionUpdate, UpdatesIterator, filter_rewind_updates,
stream_replay_updates_at, strip_context_wrappers,
};
use xai_grok_shell::session::testkit::synth::{self, SessionSpec};
// `Relaxed` suffices: single-threaded high-water counters read after the
// measured section, with no ordering dependency on other memory.
struct CountingAlloc;
static LIVE: AtomicUsize = AtomicUsize::new(0);
static PEAK: AtomicUsize = AtomicUsize::new(0);
unsafe impl GlobalAlloc for CountingAlloc {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
let ptr = unsafe { System.alloc(layout) };
if !ptr.is_null() {
let now = LIVE.fetch_add(layout.size(), Ordering::Relaxed) + layout.size();
PEAK.fetch_max(now, Ordering::Relaxed);
}
ptr
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
LIVE.fetch_sub(layout.size(), Ordering::Relaxed);
unsafe { System.dealloc(ptr, layout) }
}
}
#[global_allocator]
static ALLOC: CountingAlloc = CountingAlloc;
fn begin_measure() -> usize {
let base = LIVE.load(Ordering::Relaxed);
PEAK.store(base, Ordering::Relaxed);
base
}
fn peak() -> usize {
PEAK.load(Ordering::Relaxed)
}
/// Heavy ACU catalog, no rewind points: the parity assert covers the
/// typed-to-string swap over this transcript, not rewind-marker filtering
/// (`synth` emits no `rewind_marker` envelopes on the replay path).
fn fork_replay_spec() -> SessionSpec {
SessionSpec::from_env_prefixed(
"FORK_REPLAY",
SessionSpec {
turns: 60,
acu_per_turn: 8,
catalog_commands: 200,
catalog_desc_len: 48,
agent_chunks_per_turn: 4,
agent_chunk_len: 1500,
rewind_points: 0,
files_per_rewind: 0,
file_content_len: 0,
},
)
}
fn reference_load_all(updates_path: &Path) -> Vec<acp::SessionUpdate> {
let iter = UpdatesIterator::open(updates_path)
.expect("open updates")
.expect("updates file exists");
let all: Vec<SessionUpdate> = iter.filter_map(|r| r.ok()).collect();
let filtered = filter_rewind_updates(all);
filtered
.into_iter()
.filter_map(|u| match u {
SessionUpdate::Acp(notif) => Some(strip_context_wrappers(notif.update)),
SessionUpdate::Xai(_) => None,
})
.collect()
}
fn serialize(u: &acp::SessionUpdate) -> String {
serde_json::to_string(u).expect("serialize replayed update")
}
/// Streaming holds one `read_to_string` copy (~1x) plus transient per-update
/// structs; the old parse-all path held several multiples. Headroom over 1x.
const MAX_STREAM_PEAK_TO_DISK_RATIO: f64 = 2.0;
// Serial: the process-global counters are only valid when no other test
// allocates concurrently.
#[test]
#[serial_test::serial]
fn fork_replay_stream_is_bounded_and_faithful() {
let root = TempDir::new().unwrap();
let cwd = TempDir::new().unwrap();
let opts = fork_replay_spec();
let (id, updates_path) = {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
let out = rt.block_on(async {
let (info, dir) = synth::prepare_session(root.path(), cwd.path(), &opts).await;
(info.id.0.to_string(), dir.join("updates.jsonl"))
});
drop(rt);
out
};
let on_disk = std::fs::metadata(&updates_path).unwrap().len() as usize;
// The ratio bound is only meaningful once fixed overhead is dwarfed by
// content; guard against a shrunk spec making the assert trivial.
assert!(
on_disk > 256 * 1024,
"fork-replay fixture must be sizeable to bound the peak ratio, got {on_disk} B"
);
let base_old = begin_measure();
let reference = reference_load_all(&updates_path);
let old_peak = peak() - base_old;
let ref_count = reference.len();
let ref_serialized: Vec<String> = reference.iter().map(serialize).collect();
drop(reference);
let mut stream_count = 0usize;
let base_new = begin_measure();
let outcome = stream_replay_updates_at(&id, root.path(), |_update| {
stream_count += 1;
})
.expect("stream_replay_updates_at");
let new_peak = peak() - base_new;
// Serializing during the measured pass would inflate its peak, so replay a
// second time purely to collect the parity data.
let mut streamed_serialized: Vec<String> = Vec::new();
let _ = stream_replay_updates_at(&id, root.path(), |u| {
streamed_serialized.push(serialize(&u))
})
.expect("stream again");
let new_ratio = new_peak as f64 / on_disk.max(1) as f64;
let old_ratio = old_peak as f64 / on_disk.max(1) as f64;
eprintln!(
"FORK_REPLAY_MEMORY {}",
serde_json::json!({
"on_disk_mb": on_disk as f64 / 1e6,
"old_parse_all_peak_mb": old_peak as f64 / 1e6,
"old_ratio": old_ratio,
"new_stream_peak_mb": new_peak as f64 / 1e6,
"new_ratio": new_ratio,
"reduction_x": old_peak as f64 / new_peak.max(1) as f64,
"count": stream_count,
})
);
assert!(
outcome == ReplayEmission::Emitted && stream_count > 0,
"expected a non-empty replay"
);
assert_eq!(
streamed_serialized, ref_serialized,
"streamed updates must match the typed parse-all path byte-for-byte, in order"
);
assert_eq!(
stream_count, ref_count,
"streamed update count must equal the typed path count"
);
assert!(
new_peak < old_peak,
"streaming peak ({new_peak} B) must be below the parse-all peak ({old_peak} B)"
);
assert!(
new_ratio < MAX_STREAM_PEAK_TO_DISK_RATIO,
"streaming peak {new_ratio:.2}x on-disk must stay near the file size \
(< {MAX_STREAM_PEAK_TO_DISK_RATIO}x); the whole transcript is no longer \
materialized as typed structs"
);
}

462
crates/codegen/xai-grok-shell/tests/session_load_perf.rs
View file

@ -1,10 +1,10 @@
//! End-to-end measurement of why resuming a large session is slow — the time
//! End-to-end measurement of why resuming a large session is slow: the time
//! spent before the client can render anything.
//!
//! The pager resumes via `session/load` and blocks on the response. The shell
//! answers by (1) `load_light` (chat history; rewind points now load lazily) and
//! (2) `replay_session_updates` — reading `updates.jsonl`, filtering it, typed-
//! parsing every line, and forwarding each as a `session/update`. All of that
//! (2) `replay_session_updates`, which reads `updates.jsonl`, filters it, typed
//! parses every line, and forwards each as a `session/update`. All of that
//! happens while the client waits; both tests drive the real production code.
//!
//! * [`phase_breakdown_real_functions`] drives the exact load-path functions
@ -12,20 +12,23 @@
//! wall-clock to rewind load, chat+summary load, and updates read+parse+filter,
//! then prints a per-`sessionUpdate`-kind byte breakdown of `updates.jsonl`.
//! * [`full_session_load_e2e`] stands up a real `MvpAgent` over in-process ACP
//! pipes; times `session/load` end-to-end, counts replayed notifications, and
//! dumps the shell's own per-phase `instrumentation_timer!` events.
//! pipes (via [`load_session_via_agent`]); times `session/load` end-to-end,
//! counts replayed notifications, and dumps the shell's own per-phase
//! `instrumentation_timer!` events.
//!
//! Session data (both tests): a synthetic session mirroring the pathological real
//! one (redundant `available_commands_update` + big rewind snapshots; size knobs
//! via env, see [`GenOpts::from_env`]), or a real session dir via
//! `GROK_PERF_SESSION_SRC=/path/to/<session-dir>`.
//! Session data (both tests): a synthetic session from the shared
//! [`synth`](xai_grok_shell::session::testkit::synth) generator (redundant
//! `available_commands_update` + big rewind snapshots; size knobs via
//! `GROK_PERF_*`), or a real session dir via `GROK_PERF_SESSION_SRC=<session-dir>`.
//!
//! Run:
//! cargo test -p xai-grok-shell --test session_load_perf -- --nocapture
//! cargo test -p xai-grok-shell --test session_load_perf full_session_load_e2e -- --ignored --nocapture
//! Run (needs the `test-support` feature; on by default under Bazel):
//! cargo test -p xai-grok-shell --features test-support --test session_load_perf -- --nocapture
//! cargo test -p xai-grok-shell --features test-support --test session_load_perf full_session_load_e2e -- --ignored --nocapture
use std::cell::RefCell;
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::rc::Rc;
use std::time::{Duration, Instant};
use agent_client_protocol::{self as acp};
@ -35,195 +38,30 @@ use xai_grok_shell::session::info::Info;
use xai_grok_shell::session::storage::{
JsonlStorageAdapter, StorageAdapter, load_updates_for_replay_at,
};
use xai_grok_workspace::session::file_state::{FileSnapshot, FlexiblePath, RewindPoint};
use xai_grok_shell::session::testkit::e2e::load_session_via_agent;
use xai_grok_shell::session::testkit::synth::{self, SessionSpec};
// ───────────────────────── size knobs ─────────────────────────
// ───────────────────────── session spec ─────────────────────────
/// Generation parameters. Defaults produce a session large enough that the
/// per-phase costs are clearly measurable (tens of MB) while still finishing
/// in a few seconds. Scale up via env to approach a real heavy session.
struct GenOpts {
turns: usize,
/// `available_commands_update`s persisted per turn. The real session had
/// ~12.5 of these per turn — the slash-command catalog re-advertised on
/// every skill discovery / subagent boundary.
acu_per_turn: usize,
catalog_commands: usize,
catalog_desc_len: usize,
agent_chunks_per_turn: usize,
agent_chunk_len: usize,
rewind_points: usize,
files_per_rewind: usize,
file_content_len: usize,
}
impl GenOpts {
fn from_env() -> Self {
fn g(key: &str, default: usize) -> usize {
std::env::var(key)
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(default)
}
// A single multiplier for quick scaling of the dominant contributors.
let scale = g("GROK_PERF_SCALE", 1).max(1);
Self {
turns: g("GROK_PERF_TURNS", 80) * scale,
acu_per_turn: g("GROK_PERF_ACU_PER_TURN", 15),
catalog_commands: g("GROK_PERF_CATALOG_COMMANDS", 64),
catalog_desc_len: g("GROK_PERF_CATALOG_DESC_LEN", 320),
agent_chunks_per_turn: g("GROK_PERF_AGENT_CHUNKS_PER_TURN", 8),
agent_chunk_len: g("GROK_PERF_AGENT_CHUNK_LEN", 2000),
rewind_points: g("GROK_PERF_REWIND_POINTS", 60) * scale,
files_per_rewind: g("GROK_PERF_FILES_PER_REWIND", 40),
file_content_len: g("GROK_PERF_FILE_CONTENT_LEN", 8000),
}
}
}
// ───────────────────────── filler ─────────────────────────
/// Deterministic, non-trivially-compressible-ish filler of `n` bytes. Uses a
/// rotating word list so serde has real strings to allocate (not one repeated
/// byte), matching the cost profile of real prose/code content.
fn filler(n: usize) -> String {
const WORDS: &[&str] = &[
"alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "hotel", "india",
"juliet", "kilo", "lima", "mike", "november", "oscar", "papa", "quebec", "romeo",
];
let mut s = String::with_capacity(n + 8);
let mut i = 0usize;
while s.len() < n {
s.push_str(WORDS[i % WORDS.len()]);
s.push(' ');
i += 1;
}
s.truncate(n);
s
}
// ───────────────────────── update synthesis ─────────────────────────
fn sid(session_id: &str) -> acp::SessionId {
acp::SessionId::new(session_id.to_string())
}
fn text_chunk(text: String) -> acp::ContentChunk {
acp::ContentChunk::new(acp::ContentBlock::Text(acp::TextContent::new(text)))
}
/// Build one large `AvailableCommandsUpdate` — the redundant catalog that the
/// real session re-persisted thousands of times.
fn available_commands_update(opts: &GenOpts) -> acp::SessionUpdate {
let desc = filler(opts.catalog_desc_len);
let commands: Vec<acp::AvailableCommand> = (0..opts.catalog_commands)
.map(|i| {
acp::AvailableCommand::new(format!("command-number-{i:03}"), desc.clone()).input(Some(
acp::AvailableCommandInput::Unstructured(acp::UnstructuredCommandInput::new(
"[optional arguments here]".to_string(),
)),
))
})
.collect();
acp::SessionUpdate::AvailableCommandsUpdate(acp::AvailableCommandsUpdate::new(commands))
}
/// Serialize one notification into the exact on-disk `updates.jsonl` envelope:
/// `{"timestamp":..,"method":"session/update","params":<SessionNotification>}`.
///
/// Params are plain JSON (not the typed `acp::SessionNotification`) so generation
/// doesn't depend on the acp crate's `_meta` field type; the production replay
/// still parses it back into a typed notification — the cost we're measuring.
fn envelope_line(session_id: &str, update: acp::SessionUpdate) -> String {
let update_val = serde_json::to_value(&update).expect("serialize update");
let params = serde_json::json!({
"sessionId": session_id,
"update": update_val,
});
let envelope = serde_json::json!({
"timestamp": 0u64,
"method": "session/update",
"params": params,
});
serde_json::to_string(&envelope).expect("serialize envelope")
}
/// Per-kind statistics for the generated/loaded updates file.
#[derive(Default)]
struct KindStats {
count: BTreeMap<String, u64>,
bytes: BTreeMap<String, u64>,
}
fn generate_updates_jsonl(path: &Path, session_id: &str, opts: &GenOpts) {
let mut out = String::new();
for turn in 0..opts.turns {
out.push_str(&envelope_line(
session_id,
acp::SessionUpdate::UserMessageChunk(text_chunk(format!(
"user prompt for turn {turn}"
))),
));
out.push('\n');
for _ in 0..opts.acu_per_turn {
out.push_str(&envelope_line(session_id, available_commands_update(opts)));
out.push('\n');
}
for _ in 0..opts.agent_chunks_per_turn {
out.push_str(&envelope_line(
session_id,
acp::SessionUpdate::AgentMessageChunk(text_chunk(filler(opts.agent_chunk_len))),
));
out.push('\n');
}
}
std::fs::write(path, out).expect("write updates.jsonl");
}
fn generate_rewind_jsonl(path: &Path, opts: &GenOpts) {
let mut out = String::new();
for p in 0..opts.rewind_points {
let mut rp = RewindPoint::new(p);
for f in 0..opts.files_per_rewind {
let fp =
FlexiblePath::Absolute(PathBuf::from(format!("/repo/src/module_{p}/file_{f}.rs")));
rp.add_snapshot(FileSnapshot::new_flexible(
fp.clone(),
Some(filler(opts.file_content_len)),
));
rp.set_after_snapshot(FileSnapshot::new_flexible(
fp,
Some(filler(opts.file_content_len + 64)),
));
}
out.push_str(&serde_json::to_string(&rp).expect("serialize rewind point"));
out.push('\n');
}
std::fs::write(path, out).expect("write rewind_points.jsonl");
/// Perf-tool defaults over the shared [`SessionSpec`], tuned to the pathological
/// real session; scale/override via `GROK_PERF_*` (e.g. `GROK_PERF_TURNS`,
/// `GROK_PERF_SCALE`), or point `GROK_PERF_SESSION_SRC` at a real session dir.
fn perf_spec() -> SessionSpec {
SessionSpec::from_env_prefixed(
"GROK_PERF",
SessionSpec {
turns: 80,
rewind_points: 60,
files_per_rewind: 40,
file_content_len: 8000,
..SessionSpec::default()
},
)
}
// ───────────────────────── session setup ─────────────────────────
/// Find `<root>/sessions/<enc-cwd>/<id>` without depending on the (internal)
/// cwd encoder: scan the one level of cwd dirs for a child named `<id>`.
fn locate_session_dir(root: &Path, id: &str) -> PathBuf {
let sessions = root.join("sessions");
for entry in std::fs::read_dir(&sessions)
.expect("read sessions dir")
.flatten()
{
let candidate = entry.path().join(id);
if candidate.is_dir() {
return candidate;
}
}
panic!(
"could not locate session dir for {id} under {}",
sessions.display()
);
}
/// Recursively copy a directory tree.
/// Recursively copy a directory tree (real-session overlay only).
fn copy_tree(src: &Path, dst: &Path) {
std::fs::create_dir_all(dst).unwrap();
for entry in std::fs::read_dir(src).unwrap().flatten() {
@ -237,66 +75,58 @@ fn copy_tree(src: &Path, dst: &Path) {
}
}
/// Prepare a session on disk under `root` for working dir `cwd`. Returns the
/// `Info` and the session directory path. Uses `GROK_PERF_SESSION_SRC` if set
/// (copies a real session), otherwise synthesizes one via the production
/// storage adapter (summary) + raw envelope writes (updates/rewind).
async fn prepare_session(root: &Path, cwd: &Path, opts: &GenOpts) -> (Info, PathBuf) {
/// Prepare a session on disk under `root` for working dir `cwd`. With
/// `GROK_PERF_SESSION_SRC` set, copy a real session over a registered stub
/// (keeping our `summary.json`); otherwise synthesize one via
/// [`synth::prepare_session`].
async fn prepare_session(root: &Path, cwd: &Path, spec: &SessionSpec) -> (Info, PathBuf) {
let Ok(src) = std::env::var("GROK_PERF_SESSION_SRC") else {
return synth::prepare_session(root, cwd, spec).await;
};
let adapter = JsonlStorageAdapter::with_root(root.to_path_buf());
if let Ok(src) = std::env::var("GROK_PERF_SESSION_SRC") {
// Real session: create a registered session shell to get the encoded
// cwd dir + a valid summary, then overlay the real files on top.
let id = uuid::Uuid::new_v4().to_string();
let info = Info {
id: sid(&id),
cwd: cwd.to_string_lossy().to_string(),
};
adapter
.init_session(&info, acp::ModelId::new("test-model"))
.await
.expect("init_session");
let dir = locate_session_dir(root, &id);
// Copy real session files (updates/rewind/chat/etc.) over the stub,
// but keep our freshly-written summary.json (correct id + cwd + model).
for name in ["updates.jsonl", "rewind_points.jsonl", "chat_history.jsonl"] {
let from = Path::new(&src).join(name);
if from.exists() {
std::fs::copy(&from, dir.join(name)).unwrap();
}
}
// Compaction checkpoints may be referenced by replay; copy if present.
let ckpt = Path::new(&src).join("compaction_checkpoints");
if ckpt.is_dir() {
copy_tree(&ckpt, &dir.join("compaction_checkpoints"));
}
eprintln!("[perf] using REAL session copied from {src}");
return (info, dir);
}
let id = uuid::Uuid::new_v4().to_string();
let info = Info {
id: sid(&id),
id: synth::sid(&id),
cwd: cwd.to_string_lossy().to_string(),
};
adapter
.init_session(&info, acp::ModelId::new("test-model"))
.await
.expect("init_session");
let dir = locate_session_dir(root, &id);
let t = Instant::now();
generate_updates_jsonl(&dir.join("updates.jsonl"), &id, opts);
generate_rewind_jsonl(&dir.join("rewind_points.jsonl"), opts);
eprintln!(
"[perf] generated synthetic session in {} ms (turns={}, acu/turn={})",
t.elapsed().as_millis(),
opts.turns,
opts.acu_per_turn
);
let dir = synth::locate_session_dir(root, &id);
for name in ["updates.jsonl", "rewind_points.jsonl", "chat_history.jsonl"] {
let from = Path::new(&src).join(name);
if from.exists() {
std::fs::copy(&from, dir.join(name)).unwrap();
}
}
let ckpt = Path::new(&src).join("compaction_checkpoints");
if ckpt.is_dir() {
copy_tree(&ckpt, &dir.join("compaction_checkpoints"));
}
eprintln!("[perf] using REAL session copied from {src}");
(info, dir)
}
/// Re-create the rewind file after the isolation step deletes it (synthetic
/// case). For a real session copy we cannot regenerate; leave it absent.
fn generate_or_restore_rewind(path: &Path, spec: &SessionSpec) {
if std::env::var("GROK_PERF_SESSION_SRC").is_ok() {
return;
}
synth::write_rewind_jsonl(path, spec);
}
// ───────────────────────── updates.jsonl stats ─────────────────────────
/// Per-kind statistics for the generated/loaded updates file.
#[derive(Default)]
struct KindStats {
count: BTreeMap<String, u64>,
bytes: BTreeMap<String, u64>,
}
fn file_size_mb(path: &Path) -> f64 {
std::fs::metadata(path).map(|m| m.len()).unwrap_or(0) as f64 / 1e6
}
@ -382,9 +212,9 @@ fn print_kind_breakdown(label: &str, stats: &KindStats) {
async fn phase_breakdown_real_functions() {
let root = TempDir::new().unwrap();
let cwd = TempDir::new().unwrap();
let opts = GenOpts::from_env();
let spec = perf_spec();
let (info, dir) = prepare_session(root.path(), cwd.path(), &opts).await;
let (info, dir) = prepare_session(root.path(), cwd.path(), &spec).await;
let updates_path = dir.join("updates.jsonl");
let rewind_path = dir.join("rewind_points.jsonl");
@ -397,7 +227,7 @@ async fn phase_breakdown_real_functions() {
let adapter = JsonlStorageAdapter::with_root(root.path().to_path_buf());
// Phase A: load_light core (summary + chat_history) — what mvp_agent's
// Phase A: load_light core (summary + chat_history), what mvp_agent's
// `load_light` blocks on before replay.
let t = Instant::now();
let light = adapter
@ -405,12 +235,9 @@ async fn phase_breakdown_real_functions() {
.await
.expect("load_session_without_updates");
let full_load_light = t.elapsed();
// load_light no longer reads rewind_points.jsonl (deferred/lazy), so 0 by
// construction — `PersistedDataLight` has no rewind field.
let light_rewind_in_load = 0usize;
drop(light);
// Lazy rewind path (T2): the deferred cost moved here. The picker only needs
// Lazy rewind path: the deferred cost moved here. The picker only needs
// a cheap metadata scan; an actual rewind triggers the full content load.
// Both read the same file that `load_light` no longer touches.
use xai_grok_workspace::session::file_state::FileStateTracker;
@ -431,8 +258,8 @@ async fn phase_breakdown_real_functions() {
"picker metadata scan must see every rewind point"
);
// Phase A': isolate rewind cost — delete rewind file and re-measure. The
// delta is the rewind-point deserialization (full file-content snapshots).
// Phase A': isolate rewind cost by deleting the rewind file and re-measuring.
// The delta is the rewind-point deserialization (full file-content snapshots).
std::fs::remove_file(&rewind_path).ok();
let t = Instant::now();
let _light2 = adapter
@ -441,12 +268,14 @@ async fn phase_breakdown_real_functions() {
.expect("load_session_without_updates (no rewind)");
let load_light_no_rewind = t.elapsed();
// restore for downstream/manual reruns
generate_or_restore_rewind(&rewind_path, &opts);
generate_or_restore_rewind(&rewind_path, &spec);
let rewind_cost = full_load_light.saturating_sub(load_light_no_rewind);
// Phase B: updates replay parse — production `load_updates_for_replay_at`
// reads the whole file, typed-parses every line, applies rewind filtering.
// Phase B: updates replay parse. The typed `load_updates_for_replay_at`
// reads the whole file, typed-parses every line, and applies rewind
// filtering; production now streams via `stream_replay_updates_at`, so this
// measures the materialize-all parse cost.
let t = Instant::now();
let replayed = load_updates_for_replay_at(info.id.0.as_ref(), root.path())
.expect("load_updates_for_replay_at")
@ -458,7 +287,7 @@ async fn phase_breakdown_real_functions() {
eprintln!("\n[perf] ===== PRE-RENDER LOAD PHASE BREAKDOWN (real production fns) =====");
eprintln!(" rewind_points (on disk) : {num_rewind}");
eprintln!(" rewind_points loaded in load : {light_rewind_in_load} (deferred → lazy)");
eprintln!(" rewind_points loaded in load : 0 (deferred → lazy)");
eprintln!(" updates replayed (acp) : {}", replayed.len());
eprintln!(" ----------------------------------------------------------------");
eprintln!(
@ -495,38 +324,15 @@ async fn phase_breakdown_real_functions() {
assert!(!stats.bytes.is_empty(), "expected a non-empty updates file");
}
/// Re-create the rewind file after the isolation step deletes it (synthetic
/// case). For a real session copy we cannot regenerate; leave it absent.
fn generate_or_restore_rewind(path: &Path, opts: &GenOpts) {
if std::env::var("GROK_PERF_SESSION_SRC").is_ok() {
return;
}
generate_rewind_jsonl(path, opts);
}
// ───────────────────────── TEST 2: true e2e ─────────────────────────
use std::cell::RefCell;
use std::rc::Rc;
use std::sync::Arc;
use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
use xai_acp_lib::{
AcpAgentGatewayReceiver as GatewayReceiver, AcpAgentGatewaySender as GatewaySender,
LineBufferedRead,
};
use xai_grok_shell::agent::config::Config as AgentConfig;
use xai_grok_shell::agent::mvp_agent::MvpAgent;
const DUPLEX_BUFFER_BYTES: usize = 16 * 1024 * 1024;
/// Counts replayed notifications and records first/last receipt timestamps so
/// we can see how long the client streams history before `load` returns.
#[derive(Default)]
struct LoadCounters {
count: u64,
/// `available_commands_update` notifications forwarded during the load. T1
/// skips the (thousands of) historical ones, so this must stay tiny.
/// `available_commands_update` notifications forwarded during the load.
/// History replay skips the (thousands of) historical ones, so this stays tiny.
acu_count: u64,
first_at: Option<Instant>,
last_at: Option<Instant>,
@ -616,7 +422,7 @@ async fn full_session_load_e2e() {
let grok_home = TempDir::new().unwrap();
let cwd = TempDir::new().unwrap();
let opts = GenOpts::from_env();
let spec = perf_spec();
let instr_log = grok_home.path().join("instr.jsonl");
// SAFETY: single-threaded current-thread runtime; set before any agent code
@ -641,7 +447,7 @@ async fn full_session_load_e2e() {
.with(xai_grok_shell::instrumentation::layer::<Registry>())
.try_init();
let (info, dir) = prepare_session(grok_home.path(), cwd.path(), &opts).await;
let (info, dir) = prepare_session(grok_home.path(), cwd.path(), &spec).await;
let updates_path = dir.join("updates.jsonl");
let rewind_path = dir.join("rewind_points.jsonl");
eprintln!(
@ -652,85 +458,33 @@ async fn full_session_load_e2e() {
let stats = updates_kind_breakdown(&updates_path);
print_kind_breakdown("e2e", &stats);
// Zero-data-loss guard (C1): a pure load must never rewrite rewind_points.jsonl
// (T2 reads it lazily, never on the load path). Captured here, asserted after.
// Zero-data-loss guard: a pure load must never rewrite rewind_points.jsonl
// (it is read lazily, never on the load path). Captured here, asserted after.
let rewind_path_guard = rewind_path.clone();
let rewind_fp_before = file_fingerprint(&rewind_path_guard);
let local = tokio::task::LocalSet::new();
local
.run_until(async move {
let agent_config = AgentConfig::default();
let auth_manager = Arc::new(agent_config.create_auth_manager());
let (gw_tx, gw_rx) = tokio::sync::mpsc::unbounded_channel();
let gateway = GatewaySender::new(gw_tx);
let agent =
MvpAgent::new(gateway, &agent_config, auth_manager, None).expect("valid config");
let (c2a_a, c2a_b) = tokio::io::duplex(DUPLEX_BUFFER_BYTES);
let (a2c_a, a2c_b) = tokio::io::duplex(DUPLEX_BUFFER_BYTES);
// Agent side.
let agent_incoming = LineBufferedRead::spawn_local(c2a_b.compat());
let (agent_conn, agent_io) =
acp::AgentSideConnection::new(agent, a2c_a.compat_write(), agent_incoming, |fut| {
tokio::task::spawn_local(fut);
});
tokio::task::spawn_local(
GatewayReceiver::new(gw_rx, agent_conn)
.with_on_meta(xai_file_utils::trace_context::span_from_meta_traceparent)
.run(),
);
tokio::task::spawn_local(agent_io);
// Client side.
let counters = Rc::new(RefCell::new(LoadCounters::default()));
let client = CountingClient {
counters: counters.clone(),
};
let client_incoming = LineBufferedRead::spawn_local(a2c_b.compat());
let (client_conn, client_io) =
acp::ClientSideConnection::new(client, c2a_a.compat_write(), client_incoming, |fut| {
tokio::task::spawn_local(fut);
});
tokio::task::spawn_local(client_io);
use acp::Agent as _;
// initialize + authenticate (api-key, like the pager does).
let init = tokio::time::timeout(
Duration::from_secs(60),
client_conn.initialize(acp::InitializeRequest::new(acp::ProtocolVersion::V1).client_capabilities(acp::ClientCapabilities::new().fs(acp::FileSystemCapabilities::new()).terminal(false)).meta(serde_json::json!({
"startupHints": { "nonInteractive": true, "skipGitStatus": true, "skipProjectLayout": true },
"clientType": "perf-test",
"clientVersion": "0.0-test",
}).as_object().cloned())),
let loaded = load_session_via_agent(
client,
"perf-test",
info.id.clone(),
cwd.path().to_path_buf(),
)
.await
.expect("initialize timed out")
.expect("initialize failed");
.await;
let load_started = loaded.load_started;
let load_elapsed = loaded.load_elapsed;
// Keep the connection alive so the post-load re-advertise still arrives.
let _client_conn = loaded.client_conn;
if let Some(method) = init.auth_methods.iter().find(|m| &*m.id().0 == "xai.api_key") {
let _ = client_conn
.authenticate(acp::AuthenticateRequest::new(method.id().clone()).meta(serde_json::json!({ "headless": true }).as_object().cloned()))
.await;
}
// The measurement: time the full session/load round-trip.
let load_started = Instant::now();
let resp = tokio::time::timeout(
Duration::from_secs(180),
client_conn.load_session(acp::LoadSessionRequest::new(info.id.clone(), cwd.path().to_path_buf())),
)
.await
.expect("session/load timed out (>180s)")
.expect("session/load failed");
let load_elapsed = load_started.elapsed();
let _ = resp;
// Snapshot replay results immediately — BEFORE the post-load
// AdvertiseCommands re-advertise can arrive — so `acu_replayed` is the
// count of ACUs forwarded during history replay (the T1 skip count).
// Snapshot replay results immediately, before the post-load
// AdvertiseCommands re-advertise can arrive, so `acu_replayed` counts
// the ACUs forwarded during history replay (the skip count).
let (replay_count, acu_replayed, ttfn, ttln) = {
let c = counters.borrow();
(
@ -763,8 +517,8 @@ async fn full_session_load_e2e() {
let mut phases = parse_instrumentation_log(&instr_log);
phases.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
// T1 guard: the historical available_commands_update copies (3197 in
// the pathological real session, hundreds in the synthetic one) must
// Replay-skip guard: the historical available_commands_update copies
// (3197 in the pathological real session, hundreds synthetic) must
// NOT be replayed.
let acu_persisted = stats.count.get("available_commands_update").copied().unwrap_or(0);
@ -786,13 +540,13 @@ async fn full_session_load_e2e() {
eprintln!("================================================================\n");
assert!(replay_count > 0, "expected replayed notifications during load");
// C1: the lazy rewind file must be byte-for-byte unchanged by a load.
// The lazy rewind file must be byte-for-byte unchanged by a load.
assert_eq!(
file_fingerprint(&rewind_path_guard),
rewind_fp_before,
"rewind_points.jsonl must be unchanged after a load (zero data loss)"
);
// The thousands of persisted ACUs must be skipped on replay (T1)...
// The thousands of persisted ACUs must be skipped on replay...
assert!(
acu_persisted > 100,
"fixture should have many persisted ACUs to exercise the skip"

46
crates/codegen/xai-grok-shell/tests/test_leader_soak.rs
View file

@ -32,6 +32,7 @@ use xai_grok_shell::leader::{
ClientCapabilities, ClientMode, LeaderClient, LeaderServerControlState, LeaderServerMetadata,
run_leader_server,
};
use xai_grok_test_support::resources::ResourceSnapshot;
const SIMPLEX_BUF: usize = 8 * 1024 * 1024;
@ -42,41 +43,6 @@ fn env_u64(key: &str, default: u64) -> u64 {
.unwrap_or(default)
}
/// Resident set size of THIS process (leader server + agent are in-process).
/// Copied from `xai-codebase-graph/tests/memory_integration.rs`.
fn rss_bytes() -> Option<usize> {
#[cfg(target_os = "linux")]
{
let status = std::fs::read_to_string("/proc/self/status").ok()?;
for line in status.lines() {
if let Some(val) = line.strip_prefix("VmRSS:") {
let kb: usize = val.trim().trim_end_matches(" kB").trim().parse().ok()?;
return Some(kb * 1024);
}
}
None
}
#[cfg(target_os = "macos")]
{
use std::process::Command;
let output = Command::new("ps")
.args(["-o", "rss=", "-p", &std::process::id().to_string()])
.output()
.ok()?;
let kb: usize = String::from_utf8_lossy(&output.stdout)
.trim()
.parse()
.ok()?;
Some(kb * 1024)
}
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
{
None
}
}
/// `leader.response.send_failed` entries written by THIS process.
fn send_failed_count() -> usize {
let Some(bytes) = xai_grok_telemetry::unified_log::snapshot_log() else {
@ -281,7 +247,7 @@ async fn leader_soak_churning_clients_no_leaks_no_zombies() {
)
.await;
let rss_baseline = rss_bytes();
let rss_before = ResourceSnapshot::capture();
let soak_deadline = tokio::time::Instant::now() + Duration::from_secs(soak_secs);
let workdir_str = workdir.path().to_string_lossy().to_string();
let mut cycles: u64 = 0;
@ -379,8 +345,12 @@ async fn leader_soak_churning_clients_no_leaks_no_zombies() {
);
// ── RSS bound ─────────────────────────────────────────────────
if let (Some(before), Some(after)) = (rss_baseline, rss_bytes()) {
let growth_mb = after.saturating_sub(before) as f64 / (1024.0 * 1024.0);
let rss_after = ResourceSnapshot::capture();
let growth = rss_after.growth_from(&rss_before);
if let (Some(before), Some(after), Some(growth_bytes)) =
(rss_before.rss, rss_after.rss, growth.rss)
{
let growth_mb = growth_bytes as f64 / (1024.0 * 1024.0);
eprintln!(
"[soak] rss: {:.1} MB -> {:.1} MB (growth {growth_mb:.1} MB)",
before as f64 / (1024.0 * 1024.0),

183
crates/codegen/xai-grok-shell/tests/test_nonblocking_startup.rs Normal file
View file

@ -0,0 +1,183 @@
//! Non-blocking startup regression tests. `#[ignore]`; requires pre-built binary.
//!
//! These exercise leader startup through the persistent-leader fixture: the
//! leader must bind its socket and become ready without blocking on the remote
//! `/settings` + `/v1/models` fetch, and must self-heal its catalog once the
//! endpoint recovers.
//!
//! ```bash
//! cargo test -p xai-grok-shell --test test_nonblocking_startup -- --ignored
//! ```
#![cfg(unix)]
mod common;
use std::time::{Duration, Instant};
use xai_grok_test_support::leader::LeaderFixture;
use xai_grok_test_support::*;
async fn poll_until(ceiling: Duration, interval: Duration, condition: impl Fn() -> bool) -> bool {
let deadline = Instant::now() + ceiling;
while Instant::now() < deadline {
if condition() {
return true;
}
tokio::time::sleep(interval).await;
}
false
}
/// A hanging proxy must not delay leader readiness: the fixture (which waits for
/// the leader socket) must come up and a session must be created well within the
/// blocking-fetch window.
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn leader_ready_while_proxy_hangs() {
tokio::task::LocalSet::new()
.run_until(async {
let server = MockInferenceServer::start().await.unwrap();
server.set_hang(true);
let workdir = git_workdir();
let sandbox = TestSandbox::new();
let started = Instant::now();
let fixture = LeaderFixture::start(&server, workdir.workspace(), &sandbox)
.await
.expect("leader must become ready while the proxy hangs");
let mut clients = Vec::new();
common::leader::run_with_cleanup(&fixture, &mut clients, |fixture, clients| {
Box::pin(async move {
clients.push(
fixture
.spawn_client(&server, workdir.workspace(), &sandbox)
.await
.expect("spawn leader client"),
);
clients[0].initialize().await;
let _session = clients[0].create_session(workdir.workspace()).await;
let elapsed = started.elapsed();
assert!(
elapsed < Duration::from_secs(25),
"startup took {elapsed:?} with a hanging proxy; readiness appears \
to block on the network fetch\nstderr:\n{}",
clients[0].stderr_text(),
);
})
})
.await;
})
.await;
}
/// The background catalog refresh must re-fetch and push `x.ai/models/update`
/// once a previously-hanging endpoint recovers.
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn catalog_self_heals_after_endpoint_recovers() {
tokio::task::LocalSet::new()
.run_until(async {
let server = MockInferenceServer::start().await.unwrap();
server.set_hang(true);
let workdir = git_workdir();
let sandbox = TestSandbox::new();
let fixture = LeaderFixture::start(&server, workdir.workspace(), &sandbox)
.await
.expect("leader must become ready while the proxy hangs");
let mut clients = Vec::new();
common::leader::run_with_cleanup(&fixture, &mut clients, |fixture, clients| {
Box::pin(async move {
clients.push(
fixture
.spawn_client(&server, workdir.workspace(), &sandbox)
.await
.expect("spawn leader client"),
);
clients[0].initialize().await;
// Recover the endpoint. The background catalog refresh
// (5s-base backoff) re-fetches and pushes `x.ai/models/update`.
server.set_hang(false);
let healed =
poll_until(Duration::from_secs(60), Duration::from_millis(500), || {
clients[0].models_update_count() > 0
})
.await;
assert!(
healed,
"no x.ai/models/update after the endpoint recovered\nstderr:\n{}",
clients[0].stderr_text(),
);
})
})
.await;
})
.await;
}
/// Custom-backend reality: a user points at their own backend that serves
/// `/v1/models` + chat but blocks the cli-chat-proxy `/settings` (404). The
/// leader must boot fast, load its catalog from the served models, and run a
/// real prompt, even though remote settings never arrive.
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn leader_usable_when_settings_blocked_but_models_served() {
tokio::task::LocalSet::new()
.run_until(async {
// Models + chat are served; `set_settings` is never called, so
// `/v1/settings` 404s, mimicking a blocked proxy settings endpoint.
let server = MockInferenceServer::start().await.unwrap();
let workdir = git_workdir();
let sandbox = TestSandbox::new();
let started = Instant::now();
let fixture = LeaderFixture::start(&server, workdir.workspace(), &sandbox)
.await
.expect("leader must become ready with /settings blocked");
let mut clients = Vec::new();
common::leader::run_with_cleanup(&fixture, &mut clients, |fixture, clients| {
Box::pin(async move {
clients.push(
fixture
.spawn_client(&server, workdir.workspace(), &sandbox)
.await
.expect("spawn leader client"),
);
clients[0].initialize().await;
let session = clients[0].create_session(workdir.workspace()).await;
// The custom backend serves chat, so a prompt round-trips
// despite the blocked settings endpoint.
clients[0]
.prompt(&session, "ping")
.await
.expect("prompt against the custom backend");
assert!(
server.has_chat_completion_request() || server.has_responses_request(),
"the prompt must reach the served chat endpoint\nstderr:\n{}",
clients[0].stderr_text(),
);
assert!(
started.elapsed() < Duration::from_secs(25),
"startup/usage blocked on the unreachable /settings\nstderr:\n{}",
clients[0].stderr_text(),
);
assert_eq!(
clients[0].settings_update_count(),
0,
"no settings update should land while /settings is blocked",
);
})
})
.await;
})
.await;
}

68
crates/codegen/xai-grok-shell/tests/test_nonblocking_startup_offline.rs Normal file
View file

@ -0,0 +1,68 @@
//! Leader boots from local data when the endpoint is fully unreachable
//! (connection refused), not merely hanging. `#[ignore]`: needs the built binary.
//!
//! ```bash
//! cargo test -p xai-grok-shell --test test_nonblocking_startup_offline -- --ignored
//! ```
#![cfg(unix)]
mod common;
use std::time::{Duration, Instant};
use xai_grok_test_support::leader::LeaderFixture;
use xai_grok_test_support::*;
/// A loopback URL on a closed port (bind, read addr, drop); refuses instantly.
fn closed_port_base_url() -> String {
let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind ephemeral port");
let addr = listener.local_addr().expect("local_addr");
drop(listener);
format!("http://{addr}/v1")
}
async fn assert_boots_fast(base_url: String, scenario: &'static str) {
let workdir = git_workdir();
let sandbox = TestSandbox::new();
let started = Instant::now();
let fixture = LeaderFixture::start_with_base_url(&base_url, workdir.workspace(), &sandbox)
.await
.unwrap_or_else(|error| {
panic!("[{scenario}] leader never became ready with an unreachable endpoint: {error}")
});
let mut clients = Vec::new();
common::leader::run_with_cleanup(&fixture, &mut clients, |fixture, clients| {
Box::pin(async move {
clients.push(
fixture
.spawn_client_with_base_url(&base_url, workdir.workspace(), &sandbox)
.await
.unwrap_or_else(|error| panic!("[{scenario}] spawn leader client: {error}")),
);
clients[0].initialize().await;
// Catalog resolves offline (built-in/cache), so session creation succeeds.
let _session = clients[0].create_session(workdir.workspace()).await;
let elapsed = started.elapsed();
assert!(
elapsed < Duration::from_secs(25),
"[{scenario}] startup took {elapsed:?} with an unreachable endpoint; \
readiness appears to block on the network fetch\nstderr:\n{}",
clients[0].stderr_text(),
);
})
})
.await;
}
#[tokio::test]
#[ignore] // needs the built binary
async fn leader_ready_with_connection_refused() {
tokio::task::LocalSet::new()
.run_until(async {
assert_boots_fast(closed_port_base_url(), "connection-refused").await;
})
.await;
}

730
crates/codegen/xai-grok-shell/tests/test_session_load_memory.rs Normal file
View file

@ -0,0 +1,730 @@
//! Memory tests for the session-load path: prove the resume peek borrows the
//! transcript instead of copying, and bound peak memory with full cleanup.
//! Resuming a large session once OOM-killed the process under a cgroup cap.
//!
//! Run:
//! cargo test -p xai-grok-shell --features dhat-heap --test test_session_load_memory \
//! session_load_dhat_bounded_and_freed -- --ignored --nocapture
#![cfg(unix)]
#[cfg(feature = "dhat-heap")]
#[global_allocator]
static DHAT_ALLOC: dhat::Alloc = dhat::Alloc;
use xai_grok_shell::session::storage::{JsonlStorageAdapter, StorageAdapter, prepare_replay_lines};
use xai_grok_shell::session::testkit::synth::{self, SessionSpec};
#[cfg(feature = "dhat-heap")]
use std::path::Path;
#[cfg(feature = "dhat-heap")]
use xai_grok_shell::session::info::Info;
use tempfile::TempDir;
const BYTES_PER_MB: f64 = 1024.0 * 1024.0;
const BYTES_PER_MB_U64: u64 = 1024 * 1024;
fn file_len(path: &std::path::Path) -> u64 {
// Fail loud: a silent 0 would collapse the ratio budget instead of
// surfacing a missing or unreadable fixture.
std::fs::metadata(path).expect("stat updates.jsonl").len()
}
fn env_parse<T: std::str::FromStr>(key: &str, default: T) -> T {
let Ok(text) = std::env::var(key) else {
return default;
};
match text.parse() {
Ok(value) => value,
Err(_) => {
eprintln!("[test_session_load] ignoring unparseable {key}={text:?}; using default");
default
}
}
}
fn memory_spec() -> SessionSpec {
SessionSpec::from_env_prefixed("SESSION_LOAD", SessionSpec::default())
}
// Replay keeps one user chunk plus the agent chunks per turn and drops the
// redundant ACUs, mirroring `synth::prepare_session` and `prepare_replay_lines`.
fn expected_replayed_lines(spec: &SessionSpec) -> usize {
spec.turns * (1 + spec.agent_chunks_per_turn)
}
/// Non-ignored zero-copy guard: every replay line must borrow from the
/// transcript, so an owned-copy regression fails here in CI.
#[tokio::test]
async fn prepare_replay_lines_borrows_the_transcript() {
let spec = SessionSpec {
turns: 3,
acu_per_turn: 2,
catalog_commands: 2,
catalog_desc_len: 8,
agent_chunks_per_turn: 2,
agent_chunk_len: 32,
rewind_points: 0,
files_per_rewind: 0,
file_content_len: 0,
};
let root = TempDir::new().unwrap();
let cwd = TempDir::new().unwrap();
let (_info, dir) = synth::prepare_session(root.path(), cwd.path(), &spec).await;
let transcript =
std::fs::read_to_string(dir.join("updates.jsonl")).expect("read updates.jsonl");
let prepared = prepare_replay_lines(&transcript, None);
assert_eq!(
prepared.lines.len(),
expected_replayed_lines(&spec),
"replay line count regressed"
);
let start = transcript.as_ptr() as usize;
let end = start + transcript.len();
for line in &prepared.lines {
let line_start = line.as_ptr() as usize;
assert!(
line_start >= start && line_start + line.len() <= end,
"replay line must borrow from the transcript (zero-copy), not own a copy"
);
}
}
/// Let ready tasks drain and timer-driven cleanup run before reading heap
/// stats, so a drop's frees show up in the next `curr_bytes`/`curr_blocks`.
#[cfg(feature = "dhat-heap")]
async fn quiesce() {
const YIELDS: usize = 50;
const SETTLE: std::time::Duration = std::time::Duration::from_millis(10);
for _ in 0..YIELDS {
tokio::task::yield_now().await;
}
tokio::time::sleep(SETTLE).await;
}
#[cfg(feature = "dhat-heap")]
async fn run_load_cycle(adapter: &JsonlStorageAdapter, info: &Info, updates_path: &Path) -> usize {
let light = adapter
.load_session_without_updates(info)
.await
.expect("load_session_without_updates");
let transcript = std::fs::read_to_string(updates_path).expect("read updates.jsonl");
let prepared = prepare_replay_lines(&transcript, None);
let replayed_lines = prepared.lines.len();
drop(prepared);
drop(transcript);
drop(light);
quiesce().await;
replayed_lines
}
/// Env-derived gates and cycle counts, separated from the measured results. The
/// caller clamps `cycles` to at least one so the per-cycle divisions are safe.
#[cfg(feature = "dhat-heap")]
struct DhatBudget {
warmup: usize,
cycles: usize,
ratio: f64,
abs_budget_mb: u64,
max_bytes_per_cycle: i64,
max_blocks_per_cycle: i64,
}
/// Heap readings captured across the measured window, named so nothing can
/// silently transpose the same-typed counts.
#[cfg(feature = "dhat-heap")]
struct DhatMeasured {
replayed_lines: usize,
expected_lines: usize,
on_disk_bytes: u64,
peak_over_baseline: u64,
net_bytes: i64,
net_blocks: i64,
}
/// The measured window paired with the budget it is judged against; every gate
/// threshold derives from `budget`, so nothing is stored twice.
#[cfg(feature = "dhat-heap")]
struct DhatOutcome<'a> {
budget: &'a DhatBudget,
measured: DhatMeasured,
}
#[cfg(feature = "dhat-heap")]
impl DhatOutcome<'_> {
fn ratio_budget_bytes(&self) -> u64 {
(self.budget.ratio * self.measured.on_disk_bytes as f64) as u64
}
fn abs_budget_bytes(&self) -> u64 {
self.budget.abs_budget_mb * BYTES_PER_MB_U64
}
fn per_cycle_bytes(&self) -> i64 {
self.measured.net_bytes / self.budget.cycles as i64
}
fn per_cycle_blocks(&self) -> i64 {
self.measured.net_blocks / self.budget.cycles as i64
}
fn no_spike(&self) -> bool {
self.measured.peak_over_baseline < self.ratio_budget_bytes()
&& self.measured.peak_over_baseline < self.abs_budget_bytes()
}
fn cleaned_up(&self) -> bool {
self.per_cycle_bytes() < self.budget.max_bytes_per_cycle
&& self.per_cycle_blocks() < self.budget.max_blocks_per_cycle
}
fn pass(&self) -> bool {
self.no_spike()
&& self.cleaned_up()
&& self.measured.replayed_lines == self.measured.expected_lines
}
}
#[cfg(feature = "dhat-heap")]
#[test]
fn dhat_outcome_verdict_arithmetic() {
let budget = DhatBudget {
warmup: 0,
cycles: 4,
ratio: 2.0,
abs_budget_mb: 1,
max_bytes_per_cycle: 100,
max_blocks_per_cycle: 10,
};
let ok = DhatOutcome {
budget: &budget,
measured: DhatMeasured {
replayed_lines: 5,
expected_lines: 5,
on_disk_bytes: 1024,
peak_over_baseline: 1000,
net_bytes: 40,
net_blocks: 4,
},
};
assert_eq!(ok.per_cycle_bytes(), 10);
assert_eq!(ok.per_cycle_blocks(), 1);
assert!(ok.no_spike() && ok.cleaned_up() && ok.pass());
// Peak over the ratio budget (2x the 1024-byte file) trips no_spike.
let ratio_spike = DhatOutcome {
budget: &budget,
measured: DhatMeasured {
replayed_lines: 5,
expected_lines: 5,
on_disk_bytes: 1024,
peak_over_baseline: 4096,
net_bytes: 0,
net_blocks: 0,
},
};
assert!(!ratio_spike.no_spike() && !ratio_spike.pass());
// Per-cycle residual over the gate trips cleaned_up.
let leak = DhatOutcome {
budget: &budget,
measured: DhatMeasured {
replayed_lines: 5,
expected_lines: 5,
on_disk_bytes: 1024,
peak_over_baseline: 1000,
net_bytes: 4 * 200,
net_blocks: 4 * 20,
},
};
assert_eq!(leak.per_cycle_bytes(), 200);
assert!(!leak.cleaned_up() && !leak.pass());
// Clean gates but a mismatched replay count still fails pass.
let miscount = DhatOutcome {
budget: &budget,
measured: DhatMeasured {
replayed_lines: 4,
expected_lines: 5,
on_disk_bytes: 1024,
peak_over_baseline: 1000,
net_bytes: 40,
net_blocks: 4,
},
};
assert!(miscount.no_spike() && miscount.cleaned_up() && !miscount.pass());
// A peak under the ratio budget but over the absolute budget trips no_spike
// via its other arm.
let abs_budget = DhatBudget {
ratio: 1.0,
..budget
};
let abs_spike = DhatOutcome {
budget: &abs_budget,
measured: DhatMeasured {
replayed_lines: 5,
expected_lines: 5,
on_disk_bytes: 2 * 1024 * 1024,
peak_over_baseline: 1024 * 1024 + 1,
net_bytes: 0,
net_blocks: 0,
},
};
assert!(!abs_spike.no_spike() && !abs_spike.pass());
}
#[cfg(feature = "dhat-heap")]
fn report_summary(o: &DhatOutcome<'_>) {
eprintln!(
"SESSION_LOAD_DHAT_SUMMARY {}",
serde_json::json!({
"mode": "dhat-heap",
"cycles": o.budget.cycles,
"warmup": o.budget.warmup,
"replayed_lines": o.measured.replayed_lines,
"expected_lines": o.measured.expected_lines,
"on_disk_updates_bytes": o.measured.on_disk_bytes,
"on_disk_updates_mb": o.measured.on_disk_bytes as f64 / BYTES_PER_MB,
"peak_over_baseline_bytes": o.measured.peak_over_baseline,
"peak_over_baseline_mb": o.measured.peak_over_baseline as f64 / BYTES_PER_MB,
"peak_over_on_disk": o.measured.peak_over_baseline as f64 / o.measured.on_disk_bytes.max(1) as f64,
"ratio_budget": o.budget.ratio,
"ratio_budget_mb": o.ratio_budget_bytes() as f64 / BYTES_PER_MB,
"abs_budget_mb": o.budget.abs_budget_mb,
"no_spike": o.no_spike(),
"per_cycle_residual_bytes": o.per_cycle_bytes(),
"per_cycle_residual_blocks": o.per_cycle_blocks(),
"max_bytes_per_cycle": o.budget.max_bytes_per_cycle,
"max_blocks_per_cycle": o.budget.max_blocks_per_cycle,
"net_window_bytes": o.measured.net_bytes,
"net_window_blocks": o.measured.net_blocks,
"cleaned_up": o.cleaned_up(),
"pass": o.pass(),
})
);
}
#[cfg(feature = "dhat-heap")]
fn assert_bounds(o: &DhatOutcome<'_>) {
assert_eq!(
o.measured.replayed_lines, o.measured.expected_lines,
"replayed line count must equal the non-ACU update count"
);
assert!(
o.no_spike(),
"load peak {:.1} MB over baseline is {:.2}x the {:.1} MB on-disk updates and exceeds a gate \
(RATIO {}x = {:.1} MB, ABSOLUTE {} MB); load memory is super-linear in session size",
o.measured.peak_over_baseline as f64 / BYTES_PER_MB,
o.measured.peak_over_baseline as f64 / o.measured.on_disk_bytes.max(1) as f64,
o.measured.on_disk_bytes as f64 / BYTES_PER_MB,
o.budget.ratio,
o.ratio_budget_bytes() as f64 / BYTES_PER_MB,
o.budget.abs_budget_mb,
);
assert!(
o.cleaned_up(),
"leak: {} bytes/cycle and {} blocks/cycle retained over {} load/drop cycles \
({} net bytes, {} net blocks) exceed the {}-byte / {}-block gate; load does not free \
everything",
o.per_cycle_bytes(),
o.per_cycle_blocks(),
o.budget.cycles,
o.measured.net_bytes,
o.measured.net_blocks,
o.budget.max_bytes_per_cycle,
o.budget.max_blocks_per_cycle,
);
}
#[cfg(feature = "dhat-heap")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[ignore = "memory soak; run with --features dhat-heap --ignored --nocapture"]
async fn session_load_dhat_bounded_and_freed() {
let opts = memory_spec();
let root = TempDir::new().unwrap();
let cwd = TempDir::new().unwrap();
let (info, dir) = synth::prepare_session(root.path(), cwd.path(), &opts).await;
let updates_path = dir.join("updates.jsonl");
let on_disk_bytes = file_len(&updates_path);
let expected_lines = expected_replayed_lines(&opts);
let budget = DhatBudget {
warmup: env_parse("SESSION_LOAD_WARMUP", 3usize),
cycles: env_parse("SESSION_LOAD_CYCLES", 8usize).max(1),
ratio: env_parse("SESSION_LOAD_HEAP_RATIO", 4.0),
abs_budget_mb: env_parse("SESSION_LOAD_MAX_PEAK_HEAP_MB", 512u64),
max_bytes_per_cycle: env_parse("SESSION_LOAD_MAX_BYTES_PER_CYCLE", 1i64 << 20),
max_blocks_per_cycle: env_parse("SESSION_LOAD_MAX_BLOCKS_PER_CYCLE", 128i64),
};
let adapter = JsonlStorageAdapter::with_root(root.path().to_path_buf());
let profiler = dhat::Profiler::builder().testing().build();
for _ in 0..budget.warmup {
let _ = run_load_cycle(&adapter, &info, &updates_path).await;
}
let window_before = dhat::HeapStats::get();
let mut replayed_lines = 0usize;
for _ in 0..budget.cycles {
replayed_lines = run_load_cycle(&adapter, &info, &updates_path).await;
}
let window_after = dhat::HeapStats::get();
drop(profiler);
// `max_bytes` is a running maximum over the profiler's whole life (warmup
// included), so subtracting the post-warmup baseline yields a conservative
// upper bound on the load peak, never an underestimate.
let peak_over_baseline =
(window_after.max_bytes as u64).saturating_sub(window_before.curr_bytes as u64);
// Net change across the measured window; goes negative if a cycle frees more
// than warmup left resident, which still satisfies the leak gate.
let net_bytes = window_after.curr_bytes as i64 - window_before.curr_bytes as i64;
let net_blocks = window_after.curr_blocks as i64 - window_before.curr_blocks as i64;
let outcome = DhatOutcome {
budget: &budget,
measured: DhatMeasured {
replayed_lines,
expected_lines,
on_disk_bytes,
peak_over_baseline,
net_bytes,
net_blocks,
},
};
report_summary(&outcome);
assert_bounds(&outcome);
}
// dhat replaces the global allocator and perturbs RSS, so the RSS-based forms
// only compile without the `dhat-heap` feature.
#[cfg(not(feature = "dhat-heap"))]
mod rss {
use super::*;
use std::cell::RefCell;
use std::path::PathBuf;
use std::rc::Rc;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use agent_client_protocol::{self as acp};
use xai_grok_test_support::resources::ResourceSnapshot;
const SAMPLE_INTERVAL_MS: u64 = 3;
/// Poll RSS from a thread; the load path is synchronous, so this is how its
/// peak gets captured.
fn spawn_rss_sampler(stop: Arc<AtomicBool>) -> std::thread::JoinHandle<usize> {
std::thread::spawn(move || {
let mut peak = ResourceSnapshot::capture_rss().unwrap_or(0);
while !stop.load(Ordering::Relaxed) {
if let Some(r) = ResourceSnapshot::capture_rss() {
peak = peak.max(r);
}
std::thread::sleep(Duration::from_millis(SAMPLE_INTERVAL_MS));
}
if let Some(r) = ResourceSnapshot::capture_rss() {
peak = peak.max(r);
}
peak
})
}
/// Owns the RSS baseline and the background sampler for one measured load.
struct RssSampler {
baseline: Option<usize>,
stop: Arc<AtomicBool>,
handle: std::thread::JoinHandle<usize>,
}
impl RssSampler {
/// Capture the RSS baseline and start a background sampler.
fn start() -> Self {
let baseline = ResourceSnapshot::capture_rss();
let stop = Arc::new(AtomicBool::new(false));
let handle = spawn_rss_sampler(stop.clone());
Self {
baseline,
stop,
handle,
}
}
/// Stop sampling and build the outcome. Takes a final synchronous read
/// first as a best-effort backstop: when the caller still holds the
/// loaded state it pins a load that peaked and freed between the
/// sampler's ticks; otherwise the background sampler is the sole source.
fn finish(self, budget_mb: u64) -> RssOutcome {
let final_rss = ResourceSnapshot::capture_rss().unwrap_or(0);
self.stop.store(true, Ordering::Relaxed);
let peak_rss = self.handle.join().expect("sampler thread").max(final_rss);
RssOutcome {
baseline: self.baseline,
peak_rss,
budget_mb,
}
}
}
struct RssOutcome {
baseline: Option<usize>,
peak_rss: usize,
budget_mb: u64,
}
impl RssOutcome {
fn baseline_bytes(&self) -> usize {
self.baseline.unwrap_or(0)
}
fn measurable(&self) -> bool {
self.baseline.is_some() && self.peak_rss > 0
}
fn peak_growth_bytes(&self) -> usize {
self.peak_rss.saturating_sub(self.baseline_bytes())
}
fn budget_bytes(&self) -> u64 {
self.budget_mb * BYTES_PER_MB_U64
}
fn within_budget(&self) -> bool {
(self.peak_growth_bytes() as u64) < self.budget_bytes()
}
fn pass(&self) -> bool {
!self.measurable() || self.within_budget()
}
}
#[test]
fn rss_outcome_verdict_arithmetic() {
let mb = BYTES_PER_MB_U64 as usize;
let over = RssOutcome {
baseline: Some(mb),
peak_rss: mb + 3 * mb,
budget_mb: 2,
};
assert_eq!(over.peak_growth_bytes(), 3 * mb);
assert!(!over.within_budget());
assert!(!over.pass());
let under = RssOutcome {
baseline: Some(mb),
peak_rss: mb + mb,
budget_mb: 2,
};
assert!(under.within_budget());
assert!(under.pass());
// An unmeasurable baseline passes vacuously.
let unmeasurable = RssOutcome {
baseline: None,
peak_rss: 0,
budget_mb: 1,
};
assert!(!unmeasurable.measurable());
assert!(unmeasurable.pass());
}
fn report_summary(mode: &str, counts: serde_json::Value, on_disk_bytes: u64, o: &RssOutcome) {
let mut summary = serde_json::json!({
"mode": mode,
"on_disk_updates_bytes": on_disk_bytes,
"on_disk_updates_mb": on_disk_bytes as f64 / BYTES_PER_MB,
"baseline_rss_mb": o.baseline_bytes() as f64 / BYTES_PER_MB,
"peak_rss_mb": o.peak_rss as f64 / BYTES_PER_MB,
"peak_rss_growth_mb": o.peak_growth_bytes() as f64 / BYTES_PER_MB,
"budget_mb": o.budget_mb,
"rss_measurable": o.measurable(),
"pass": o.pass(),
});
let obj = summary
.as_object_mut()
.expect("summary literal is a JSON object");
let extra = counts.as_object().expect("counts must be a JSON object");
for (k, v) in extra {
obj.insert(k.clone(), v.clone());
}
eprintln!("SESSION_LOAD_MEMORY_SUMMARY {summary}");
}
fn assert_bounds(label: Option<&str>, on_disk_bytes: u64, o: &RssOutcome) {
if o.measurable() {
let prefix = label.map(|l| format!("{l} ")).unwrap_or_default();
assert!(
o.within_budget(),
"{prefix}peak RSS grew {:.1} MB over baseline while loading a {:.1} MB updates file \
(bound {} MB)",
o.peak_growth_bytes() as f64 / BYTES_PER_MB,
on_disk_bytes as f64 / BYTES_PER_MB,
o.budget_mb,
);
} else {
eprintln!("[soak] RSS measurement unavailable on this platform; bound skipped");
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[ignore = "peak-memory soak; run with --ignored --nocapture"]
async fn session_load_peak_rss_under_budget() {
let opts = memory_spec();
let root = TempDir::new().unwrap();
let cwd = TempDir::new().unwrap();
let (info, dir) = synth::prepare_session(root.path(), cwd.path(), &opts).await;
let updates_path = dir.join("updates.jsonl");
let on_disk_bytes = file_len(&updates_path);
let expected_lines = expected_replayed_lines(&opts);
let budget_mb = env_parse("SESSION_LOAD_MAX_PEAK_MB", 1024u64);
let adapter = JsonlStorageAdapter::with_root(root.path().to_path_buf());
let sampler = RssSampler::start();
let light = adapter
.load_session_without_updates(&info)
.await
.expect("load_session_without_updates");
let transcript = std::fs::read_to_string(&updates_path).expect("read updates.jsonl");
let prepared = prepare_replay_lines(&transcript, None);
let replayed = prepared.lines.len();
let outcome = sampler.finish(budget_mb);
drop(prepared);
drop(transcript);
drop(light);
// Report before asserting so a count regression still emits the summary.
report_summary(
"rss",
serde_json::json!({
"replayed_lines": replayed,
"expected_lines": expected_lines,
}),
on_disk_bytes,
&outcome,
);
assert_eq!(
replayed, expected_lines,
"replayed line count must equal the non-ACU update count"
);
assert_bounds(None, on_disk_bytes, &outcome);
}
struct CountingClient {
count: Rc<RefCell<u64>>,
}
#[async_trait::async_trait(?Send)]
impl acp::Client for CountingClient {
async fn request_permission(
&self,
args: acp::RequestPermissionRequest,
) -> acp::Result<acp::RequestPermissionResponse> {
let outcome = args
.options
.first()
.map(|o| {
acp::RequestPermissionOutcome::Selected(acp::SelectedPermissionOutcome::new(
o.option_id.clone(),
))
})
.unwrap_or(acp::RequestPermissionOutcome::Cancelled);
Ok(acp::RequestPermissionResponse::new(outcome))
}
async fn session_notification(&self, _args: acp::SessionNotification) -> acp::Result<()> {
*self.count.borrow_mut() += 1;
Ok(())
}
}
async fn count_replayed_notifications(session_id: acp::SessionId, cwd: PathBuf) -> u64 {
let local = tokio::task::LocalSet::new();
local
.run_until(async move {
let count = Rc::new(RefCell::new(0u64));
let client = CountingClient {
count: count.clone(),
};
let loaded = xai_grok_shell::session::testkit::e2e::load_session_via_agent(
client, "mem-soak", session_id, cwd,
)
.await;
drop(loaded);
*count.borrow()
})
.await
}
#[tokio::test(flavor = "current_thread")]
#[ignore = "heavy: builds a full MvpAgent and loads a large session; run with --ignored --nocapture"]
async fn session_load_e2e_peak_rss() {
let _ = rustls::crypto::ring::default_provider().install_default();
let server = xai_grok_test_support::MockInferenceServer::start()
.await
.unwrap();
let grok_home = TempDir::new().unwrap();
let cwd = TempDir::new().unwrap();
let opts = memory_spec();
let budget_mb = env_parse("SESSION_LOAD_MAX_PEAK_MB", 1024u64);
// SAFETY: single-threaded current-thread runtime; set before any agent
// code reads these process-globals (same pattern as session_load_perf).
unsafe {
std::env::set_var("GROK_HOME", grok_home.path());
std::env::set_var("GROK_CLI_CHAT_PROXY_BASE_URL", server.url());
std::env::set_var("GROK_XAI_API_BASE_URL", server.url());
std::env::set_var("XAI_API_KEY", "test-key-for-ci");
std::env::set_var("GROK_TELEMETRY_ENABLED", "false");
std::env::set_var("GROK_FEEDBACK_ENABLED", "false");
std::env::set_var("GROK_TRACE_UPLOAD", "false");
}
let (info, dir) = synth::prepare_session(grok_home.path(), cwd.path(), &opts).await;
let on_disk_bytes = file_len(&dir.join("updates.jsonl"));
let sampler = RssSampler::start();
let replay_count =
count_replayed_notifications(info.id.clone(), cwd.path().to_path_buf()).await;
// The agent load already dropped the loaded state, so the peak here comes
// from the background sampler; the final read is only a backstop.
let outcome = sampler.finish(budget_mb);
report_summary(
"rss-e2e",
serde_json::json!({ "replayed_notifications": replay_count }),
on_disk_bytes,
&outcome,
);
// At least one notification per synthesized turn must replay; a near-empty
// replay that still grew memory would otherwise pass silently.
assert!(
replay_count >= opts.turns as u64,
"replayed {replay_count} notifications, expected at least {} (one per turn)",
opts.turns,
);
assert_bounds(Some("e2e"), on_disk_bytes, &outcome);
}
}

4
crates/codegen/xai-grok-shell/tests/test_settings_refresh.rs
View file

@ -129,7 +129,7 @@ async fn test_fetch_settings_blocking_round_trip() {
let result = tokio::task::spawn_blocking({
let url = server.url().to_string();
let auth = auth.clone();
move || xai_grok_shell::remote::fetch_settings_blocking(&url, &auth, None)
move || xai_grok_shell::remote::fetch_settings_blocking(&url, &auth, None).into_option()
})
.await
.unwrap();
@ -146,7 +146,7 @@ async fn test_fetch_settings_blocking_round_trip() {
let result = tokio::task::spawn_blocking({
let url = server.url().to_string();
let auth = auth.clone();
move || xai_grok_shell::remote::fetch_settings_blocking(&url, &auth, None)
move || xai_grok_shell::remote::fetch_settings_blocking(&url, &auth, None).into_option()
})
.await
.unwrap();

796
crates/codegen/xai-grok-shell/tests/test_subagent_soak.rs Normal file
View file

@ -0,0 +1,796 @@
//! Subagent lifecycle soak: churn spawn/run/completion/eviction and assert
//! threads, fds, and heap/RSS reach steady state. A stub `ChildRunner` drives
//! the real coordinator/transport.
//!
//! SUBAGENT_SOAK_CYCLES=20000 cargo test -p xai-grok-shell \
//! [--features dhat-heap] --test test_subagent_soak -- --ignored --nocapture
#![cfg(unix)]
#[cfg(feature = "dhat-heap")]
#[global_allocator]
static DHAT_ALLOC: dhat::Alloc = dhat::Alloc;
use std::sync::Arc;
use std::time::Duration;
use serde::ser::SerializeMap;
use serde::{Serialize, Serializer};
use strum::{EnumCount, IntoEnumIterator};
use tokio_util::sync::CancellationToken;
use xai_grok_test_support::env::env_parse;
use xai_grok_test_support::resources::{ResourceGrowth, ResourceSnapshot};
use xai_grok_tools::implementations::grok_build::task::backend::{ChannelBackend, SubagentBackend};
use xai_grok_tools::implementations::grok_build::task::coordinator::{
ChildCompletion, ChildControl, ChildRunOutput, ChildRunRequest, ChildRunner, CoordinatorConfig,
LocalBoxFuture, MAX_COMPLETED_ENTRIES, StartedChild, SubagentCoordinator, SubagentProgress,
};
use xai_grok_tools::implementations::grok_build::task::types::{
SubagentDescribeOutcome, SubagentOwner, SubagentRegistryCounts, SubagentRequest,
SubagentResult, SubagentValidateTypeOutcome,
};
const PARENT_SESSION_ID: &str = "subagent-soak-parent";
#[derive(Clone, Copy, strum::EnumCount, strum::EnumIter)]
enum Metric {
Rss,
Threads,
Fds,
}
impl Metric {
fn label(self) -> &'static str {
match self {
Metric::Rss => "rss",
Metric::Threads => "threads",
Metric::Fds => "fds",
}
}
/// RSS reports raw bytes, so its key names the unit.
fn summary_key(self) -> &'static str {
match self {
Metric::Rss => "rss_bytes",
Metric::Threads => "threads",
Metric::Fds => "fds",
}
}
fn unit(self) -> Option<&'static str> {
match self {
Metric::Rss => Some("MiB"),
Metric::Threads | Metric::Fds => None,
}
}
fn budget(self, bounds: &Bounds) -> f64 {
match self {
Metric::Rss => bounds.max_rss_growth_mib as f64,
Metric::Threads => bounds.max_thread_growth as f64,
Metric::Fds => bounds.max_fd_growth as f64,
}
}
/// RSS growth samples are bytes; convert to MiB for the budget comparison.
fn growth_in_budget_unit(self, raw: usize) -> f64 {
match self {
Metric::Rss => bytes_to_mib(raw),
Metric::Threads | Metric::Fds => raw as f64,
}
}
}
/// Reads a metric's field from a snapshot or a growth delta so serialization and
/// the gates share one projection instead of repeating it.
trait MetricValue {
fn value_of(&self, metric: Metric) -> Option<usize>;
}
impl MetricValue for ResourceSnapshot {
fn value_of(&self, metric: Metric) -> Option<usize> {
// Destructure so a new resource field is a compile error here, not a
// silently dropped metric.
let ResourceSnapshot { rss, threads, fds } = *self;
match metric {
Metric::Rss => rss,
Metric::Threads => threads,
Metric::Fds => fds,
}
}
}
impl MetricValue for ResourceGrowth {
fn value_of(&self, metric: Metric) -> Option<usize> {
let ResourceGrowth { rss, threads, fds } = *self;
match metric {
Metric::Rss => rss,
Metric::Threads => threads,
Metric::Fds => fds,
}
}
}
fn serialize_metrics<T: MetricValue, S: Serializer>(
value: &T,
serializer: S,
) -> Result<S::Ok, S::Error> {
let mut map = serializer.serialize_map(Some(Metric::COUNT))?;
for metric in Metric::iter() {
map.serialize_entry(metric.summary_key(), &value.value_of(metric))?;
}
map.end()
}
fn bytes_to_mib(bytes: usize) -> f64 {
bytes as f64 / (1024.0 * 1024.0)
}
#[cfg_attr(not(feature = "dhat-heap"), allow(dead_code))]
#[derive(Clone, Copy, Serialize)]
struct HeapSample {
blocks: i64,
bytes: i64,
}
#[derive(Clone, Copy, Serialize)]
struct HeapMetrics {
before: HeapSample,
after: HeapSample,
blocks_per_cycle: f64,
bytes_per_cycle: f64,
}
impl HeapMetrics {
fn new(before: HeapSample, after: HeapSample, cycles: u64) -> Self {
// `SUBAGENT_SOAK_CYCLES=0` would otherwise divide by zero and feed
// NaN/inf into the leak gates.
let cycles = cycles.max(1) as f64;
Self {
before,
after,
blocks_per_cycle: (after.blocks - before.blocks) as f64 / cycles,
bytes_per_cycle: (after.bytes - before.bytes) as f64 / cycles,
}
}
}
#[derive(Serialize)]
struct Bounds {
#[serde(rename = "warmup_cycles")]
warmup: u64,
#[serde(rename = "measured_cycles")]
measure: u64,
max_thread_growth: u64,
max_fd_growth: u64,
max_rss_growth_mib: u64,
max_blocks_per_cycle: f64,
max_bytes_per_cycle: f64,
}
impl Bounds {
fn from_env() -> Self {
Self {
// Default warmup to the completed-entry cap so the ring is saturated
// and the measured window observes steady-state eviction rather than
// one-time cache fill.
warmup: env_parse("SUBAGENT_SOAK_WARMUP", MAX_COMPLETED_ENTRIES as u64),
measure: env_parse("SUBAGENT_SOAK_CYCLES", 512u64),
max_thread_growth: env_parse("SUBAGENT_SOAK_MAX_THREAD_GROWTH", 32u64),
max_fd_growth: env_parse("SUBAGENT_SOAK_MAX_FD_GROWTH", 64u64),
max_rss_growth_mib: env_parse("SUBAGENT_SOAK_MAX_RSS_GROWTH_MIB", 256u64),
max_blocks_per_cycle: env_parse("SUBAGENT_SOAK_MAX_BLOCKS_PER_CYCLE", 2.0f64),
max_bytes_per_cycle: env_parse("SUBAGENT_SOAK_MAX_BYTES_PER_CYCLE", 4096.0f64),
}
}
}
#[derive(Serialize)]
struct Measurement {
#[serde(serialize_with = "serialize_metrics")]
before: ResourceSnapshot,
#[serde(serialize_with = "serialize_metrics")]
after: ResourceSnapshot,
#[serde(serialize_with = "serialize_metrics")]
growth: ResourceGrowth,
#[serde(serialize_with = "serialize_counts")]
counts: SubagentRegistryCounts,
heap: Option<HeapMetrics>,
quiesced: bool,
}
fn serialize_counts<S: Serializer>(
counts: &SubagentRegistryCounts,
serializer: S,
) -> Result<S::Ok, S::Error> {
// Exhaustive destructure so a new count field is a compile error here, not a
// silently dropped summary key.
let SubagentRegistryCounts {
pending,
active,
completed,
} = counts;
let mut map = serializer.serialize_map(Some(3))?;
map.serialize_entry("pending", pending)?;
map.serialize_entry("active", active)?;
map.serialize_entry("completed", completed)?;
map.end()
}
#[derive(Serialize)]
struct Summary<'a> {
#[serde(flatten)]
bounds: &'a Bounds,
#[serde(flatten)]
measurement: &'a Measurement,
}
fn heap_capture() -> Option<HeapSample> {
#[cfg(feature = "dhat-heap")]
{
let stats = dhat::HeapStats::get();
Some(HeapSample {
blocks: stats.curr_blocks as i64,
bytes: stats.curr_bytes as i64,
})
}
#[cfg(not(feature = "dhat-heap"))]
{
None
}
}
async fn quiesce(backend: &ChannelBackend) -> bool {
const MAX_POLLS: usize = 200;
const SLEEP: Duration = Duration::from_millis(5);
for _ in 0..MAX_POLLS {
let counts = backend.registry_counts().await;
if counts.pending == 0 && counts.active == 0 {
return true;
}
tokio::time::sleep(SLEEP).await;
}
let counts = backend.registry_counts().await;
eprintln!(
"[soak] quiesce budget expired with pending={} active={}; snapshot may be noisy",
counts.pending, counts.active
);
false
}
#[derive(Clone)]
struct SoakControl {
cancellation: CancellationToken,
}
impl ChildControl for SoakControl {
type ProgressFuture = std::future::Ready<SubagentProgress>;
fn progress(&self) -> Self::ProgressFuture {
std::future::ready(SubagentProgress::default())
}
fn cancel(&self) {
self.cancellation.cancel();
}
}
struct SoakRunner;
impl ChildRunner for SoakRunner {
type Control = SoakControl;
type CompletionData = ();
type RunFuture = LocalBoxFuture<ChildRunOutput<()>>;
type ValidateFuture = LocalBoxFuture<SubagentValidateTypeOutcome>;
type DescribeFuture = LocalBoxFuture<SubagentDescribeOutcome>;
fn run(&self, run: ChildRunRequest<Self::Control>) -> Self::RunFuture {
Box::pin(async move {
let ChildRunRequest {
request,
cancellation,
reporter,
} = run;
let promoted = reporter
.started(StartedChild {
child_session_id: request.id.clone(),
persona: None,
resumed_from: request.resume_from.clone(),
child_cwd: request.cwd.clone().unwrap_or_default(),
worktree_path: None,
effective_model_id: "soak-model".to_owned(),
definition_background: false,
control: SoakControl {
cancellation: cancellation.clone(),
},
})
.await;
if !promoted || cancellation.is_cancelled() {
return ChildRunOutput {
result: SubagentResult {
success: false,
cancelled: true,
error: Some("cancelled before start".to_owned()),
subagent_id: request.id.clone(),
child_session_id: request.id,
..Default::default()
},
completion_data: (),
snapshot_ref: None,
};
}
ChildRunOutput {
result: SubagentResult {
success: true,
output: Arc::from("soak child output"),
subagent_id: request.id.clone(),
child_session_id: request.id,
tool_calls: 1,
turns: 1,
..Default::default()
},
completion_data: (),
snapshot_ref: None,
}
})
}
fn validate_type(&self, _subagent_type: String, _parent: String) -> Self::ValidateFuture {
Box::pin(std::future::ready(SubagentValidateTypeOutcome::Ok))
}
fn describe_type(
&self,
_subagent_type: String,
_harness_agent_type: Option<String>,
_parent: String,
) -> Self::DescribeFuture {
Box::pin(std::future::ready(SubagentDescribeOutcome::Unavailable))
}
fn on_completed(&self, _completion: ChildCompletion<Self::CompletionData>) {}
}
fn soak_request(id: String, background: bool) -> SubagentRequest {
SubagentRequest {
id,
prompt: "soak work".to_owned(),
description: "soak child".to_owned(),
subagent_type: "explore".to_owned(),
parent_session_id: PARENT_SESSION_ID.to_owned(),
parent_prompt_id: Some("soak-prompt".to_owned()),
resume_from: None,
cwd: None,
runtime_overrides: Default::default(),
run_in_background: background,
surface_completion: true,
await_to_completion: false,
fork_context: false,
owner: SubagentOwner::Task,
cancel_token: CancellationToken::new(),
}
}
async fn run_cycle(backend: &ChannelBackend, i: u64) {
let fg = backend
.spawn(soak_request(format!("fg-{i}"), false))
.await
.expect("foreground spawn round-trips through the coordinator");
assert!(fg.success, "cycle {i}: foreground child must complete");
let bg_id = format!("bg-{i}");
let bg = backend
.spawn(soak_request(bg_id.clone(), true))
.await
.expect("background spawn round-trips through the coordinator");
assert!(bg.success, "cycle {i}: background child must complete");
let blocking = true;
let timeout_ms = Some(5_000);
let snapshot = backend.query(&bg_id, blocking, timeout_ms).await;
assert!(
snapshot.is_some(),
"cycle {i}: completed subagent must be queryable"
);
}
async fn warmup(backend: &ChannelBackend, cycles: u64) -> bool {
for i in 0..cycles {
run_cycle(backend, i).await;
}
quiesce(backend).await
}
async fn measure(backend: &ChannelBackend, bounds: &Bounds, warmup_quiesced: bool) -> Measurement {
let heap_before = heap_capture();
let before = ResourceSnapshot::capture();
// Continue ids past the warmup window so measured cycles use fresh entries
// and keep exercising eviction instead of colliding with warmup ids.
for i in bounds.warmup..(bounds.warmup + bounds.measure) {
run_cycle(backend, i).await;
}
// A warmup that never drained already poisons the `before` baseline, so skip
// the measured-window drain and report the window as not quiesced.
let quiesced = warmup_quiesced && quiesce(backend).await;
let heap_after = heap_capture();
let after = ResourceSnapshot::capture();
let counts = backend.registry_counts().await;
Measurement {
before,
after,
growth: after.growth_from(&before),
counts,
heap: heap_before
.zip(heap_after)
.map(|(before, after)| HeapMetrics::new(before, after, bounds.measure)),
quiesced,
}
}
fn check_bounds(bounds: &Bounds, m: &Measurement) -> Vec<String> {
// Drain first: a non-quiesced window has nonzero counts and noisy growth, so
// report the quiesce failure alone; the gates below only mean anything once
// drained.
if !m.quiesced {
return vec![
"quiesce budget expired before the measured window drained; soak result is unreliable"
.to_owned(),
];
}
let mut failures = Vec::new();
if m.counts.pending != 0 {
failures.push(format!(
"no subagent may remain pending, saw {}",
m.counts.pending
));
}
if m.counts.active != 0 {
failures.push(format!(
"no subagent may remain active, saw {}",
m.counts.active
));
}
if m.counts.completed > MAX_COMPLETED_ENTRIES {
failures.push(format!(
"completed retention must stay bounded by its cap, saw {}",
m.counts.completed
));
}
for metric in Metric::iter() {
let Some(raw) = m.growth.value_of(metric) else {
continue;
};
let growth = metric.growth_in_budget_unit(raw);
let budget = metric.budget(bounds);
if growth > budget {
let unit = metric.unit().map(|u| format!(" {u}")).unwrap_or_default();
failures.push(format!(
"{}: grew {growth:.1}{unit} over the soak (bound {budget:.1}{unit})",
metric.label()
));
}
}
if let Some(h) = m.heap {
let measure = bounds.measure;
if h.blocks_per_cycle > bounds.max_blocks_per_cycle {
failures.push(format!(
"block-count leak: {:.3} blocks/cycle retained ({} over {measure} cycles) \
exceeds the {} gate",
h.blocks_per_cycle,
h.after.blocks - h.before.blocks,
bounds.max_blocks_per_cycle
));
}
if h.bytes_per_cycle > bounds.max_bytes_per_cycle {
failures.push(format!(
"byte leak: {:.1} bytes/cycle retained ({} over {measure} cycles) \
exceeds the {} gate",
h.bytes_per_cycle,
h.after.bytes - h.before.bytes,
bounds.max_bytes_per_cycle
));
}
}
failures
}
fn assert_bounds(bounds: &Bounds, m: &Measurement) {
let failures = check_bounds(bounds, m);
assert!(
failures.is_empty(),
"subagent soak bounds violated:\n - {}",
failures.join("\n - ")
);
}
/// Keep this the only test in the binary that creates a `dhat::Profiler`.
#[tokio::test(flavor = "current_thread")]
#[ignore = "subagent soak; run with --ignored (SUBAGENT_SOAK_CYCLES bounds the measured window)"]
async fn subagent_lifecycle_soak_bounds_threads_fds_and_heap() {
#[cfg(feature = "dhat-heap")]
let _profiler = dhat::Profiler::builder().testing().build();
let bounds = Bounds::from_env();
let local = tokio::task::LocalSet::new();
local
.run_until(async move {
let (command_tx, command_rx) = tokio::sync::mpsc::unbounded_channel();
let config = CoordinatorConfig {
foreground_budget: Duration::from_secs(600),
..CoordinatorConfig::default()
};
tokio::task::spawn_local(
SubagentCoordinator::new(command_rx, SoakRunner, config).run(),
);
let backend = ChannelBackend::new(command_tx);
let warmup_quiesced = warmup(&backend, bounds.warmup).await;
let measurement = measure(&backend, &bounds, warmup_quiesced).await;
let summary = Summary {
bounds: &bounds,
measurement: &measurement,
};
eprintln!(
"SUBAGENT_SOAK_SUMMARY {}",
serde_json::to_string(&summary).expect("summary serializes")
);
assert_bounds(&bounds, &measurement);
})
.await;
}
mod tests {
use super::*;
#[test]
fn value_of_reads_the_matching_slot_of_snapshot_and_growth() {
let snapshot = ResourceSnapshot {
rss: Some(11),
threads: Some(22),
fds: Some(33),
};
assert_eq!(snapshot.value_of(Metric::Rss), Some(11));
assert_eq!(snapshot.value_of(Metric::Threads), Some(22));
assert_eq!(snapshot.value_of(Metric::Fds), Some(33));
let growth = ResourceGrowth {
rss: Some(1),
threads: None,
fds: Some(3),
};
assert_eq!(growth.value_of(Metric::Rss), Some(1));
assert_eq!(growth.value_of(Metric::Threads), None);
assert_eq!(growth.value_of(Metric::Fds), Some(3));
}
#[test]
fn serialize_metrics_keys_match_summary_keys_in_order() {
#[derive(Serialize)]
struct Wrap(#[serde(serialize_with = "serialize_metrics")] ResourceSnapshot);
let snapshot = ResourceSnapshot {
rss: Some(1),
threads: None,
fds: Some(3),
};
let json = serde_json::to_string(&Wrap(snapshot)).expect("snapshot serializes");
assert_eq!(json, r#"{"rss_bytes":1,"threads":null,"fds":3}"#);
}
#[test]
fn bytes_to_mib_divides_by_1024_squared() {
assert_eq!(bytes_to_mib(0), 0.0);
assert_eq!(bytes_to_mib(1024 * 1024), 1.0);
assert_eq!(bytes_to_mib(3 * 1024 * 1024), 3.0);
}
#[test]
fn growth_in_budget_unit_scales_only_rss() {
assert_eq!(Metric::Rss.growth_in_budget_unit(2 * 1024 * 1024), 2.0);
assert_eq!(Metric::Threads.growth_in_budget_unit(7), 7.0);
assert_eq!(Metric::Fds.growth_in_budget_unit(7), 7.0);
}
#[test]
fn budget_reads_per_metric_bound() {
let bounds = Bounds {
warmup: 0,
measure: 0,
max_thread_growth: 3,
max_fd_growth: 5,
max_rss_growth_mib: 7,
max_blocks_per_cycle: 1.0,
max_bytes_per_cycle: 2.0,
};
assert_eq!(Metric::Rss.budget(&bounds), 7.0);
assert_eq!(Metric::Threads.budget(&bounds), 3.0);
assert_eq!(Metric::Fds.budget(&bounds), 5.0);
}
#[test]
fn heap_metrics_clamps_zero_cycles() {
let before = HeapSample {
blocks: 10,
bytes: 100,
};
let after = HeapSample {
blocks: 20,
bytes: 400,
};
let heap = HeapMetrics::new(before, after, 0);
assert!(heap.blocks_per_cycle.is_finite());
assert!(heap.bytes_per_cycle.is_finite());
assert_eq!(heap.blocks_per_cycle, 10.0);
assert_eq!(heap.bytes_per_cycle, 300.0);
}
fn generous_bounds() -> Bounds {
Bounds {
warmup: 0,
measure: 4,
max_thread_growth: 100,
max_fd_growth: 100,
max_rss_growth_mib: 100,
max_blocks_per_cycle: 10.0,
max_bytes_per_cycle: 10_000.0,
}
}
fn drained(growth: ResourceGrowth, heap: Option<HeapMetrics>) -> Measurement {
Measurement {
before: ResourceSnapshot::default(),
after: ResourceSnapshot::default(),
growth,
counts: SubagentRegistryCounts {
pending: 0,
active: 0,
completed: 0,
},
heap,
quiesced: true,
}
}
#[test]
fn check_bounds_passes_a_clean_drained_window() {
let m = drained(ResourceGrowth::default(), None);
assert!(check_bounds(&generous_bounds(), &m).is_empty());
}
#[test]
fn check_bounds_reports_non_quiesce_first_and_alone() {
let mut m = drained(ResourceGrowth::default(), None);
m.quiesced = false;
m.counts.pending = 3;
let failures = check_bounds(&generous_bounds(), &m);
assert_eq!(failures.len(), 1);
assert!(failures[0].contains("quiesce"));
}
#[test]
fn check_bounds_flags_over_budget_growth() {
let growth = ResourceGrowth {
rss: Some(200 * 1024 * 1024),
threads: Some(0),
fds: Some(0),
};
let failures = check_bounds(&generous_bounds(), &drained(growth, None));
assert!(
failures.iter().any(|f| f.starts_with("rss:")),
"{failures:?}"
);
}
#[test]
fn check_bounds_treats_the_budget_as_an_inclusive_max() {
let growth = ResourceGrowth {
rss: Some(100 * 1024 * 1024),
threads: Some(100),
fds: Some(100),
};
assert!(check_bounds(&generous_bounds(), &drained(growth, None)).is_empty());
}
#[test]
fn check_bounds_flags_nonzero_counts_and_heap_leak() {
let mut m = drained(
ResourceGrowth::default(),
Some(HeapMetrics {
before: HeapSample {
blocks: 0,
bytes: 0,
},
after: HeapSample {
blocks: 0,
bytes: 0,
},
blocks_per_cycle: 0.0,
bytes_per_cycle: 1_000_000.0,
}),
);
m.counts.active = 2;
let failures = check_bounds(&generous_bounds(), &m);
assert!(
failures.iter().any(|f| f.contains("active")),
"{failures:?}"
);
assert!(
failures.iter().any(|f| f.contains("byte leak")),
"{failures:?}"
);
}
#[test]
fn check_bounds_flags_pending_while_quiesced() {
let mut m = drained(ResourceGrowth::default(), None);
m.counts.pending = 3;
let failures = check_bounds(&generous_bounds(), &m);
assert!(
failures.iter().any(|f| f.contains("pending")),
"{failures:?}"
);
}
#[test]
fn check_bounds_flags_completed_over_cap() {
let mut m = drained(ResourceGrowth::default(), None);
m.counts.completed = MAX_COMPLETED_ENTRIES + 1;
let failures = check_bounds(&generous_bounds(), &m);
assert!(
failures.iter().any(|f| f.contains("completed retention")),
"{failures:?}"
);
}
#[test]
fn check_bounds_flags_thread_and_fd_over_budget() {
let growth = ResourceGrowth {
rss: Some(0),
threads: Some(200),
fds: Some(200),
};
let failures = check_bounds(&generous_bounds(), &drained(growth, None));
assert!(
failures.iter().any(|f| f.starts_with("threads:")),
"{failures:?}"
);
assert!(
failures.iter().any(|f| f.starts_with("fds:")),
"{failures:?}"
);
}
#[test]
fn check_bounds_flags_block_count_leak() {
let m = drained(
ResourceGrowth::default(),
Some(HeapMetrics {
before: HeapSample {
blocks: 0,
bytes: 0,
},
after: HeapSample {
blocks: 0,
bytes: 0,
},
blocks_per_cycle: 50.0,
bytes_per_cycle: 0.0,
}),
);
let failures = check_bounds(&generous_bounds(), &m);
assert!(
failures.iter().any(|f| f.contains("block-count leak")),
"{failures:?}"
);
}
}

87
crates/codegen/xai-grok-shell/tests/testkit_synth_roundtrip.rs Normal file
View file

@ -0,0 +1,87 @@
//! Non-ignored guard for the testkit's core path: synthesize a small session
//! with [`synth::prepare_session`], then confirm the production replay reader
//! parses every persisted update back with the right per-kind counts.
//!
//! `load_updates_for_replay_at` is the typed reader and keeps every update
//! (only `Xai` updates are dropped); the redundant-ACU skip is a later
//! line-based step in the client replay path, not asserted here.
use agent_client_protocol as acp;
use tempfile::TempDir;
use xai_grok_shell::session::storage::{
JsonlStorageAdapter, StorageAdapter, load_updates_for_replay_at,
};
use xai_grok_shell::session::testkit::synth::{self, SessionSpec};
#[tokio::test]
async fn synth_replay_roundtrip_parses_every_persisted_update() {
let root = TempDir::new().unwrap();
let cwd = TempDir::new().unwrap();
// Distinct per-turn counts so a miscount of any kind is unambiguous.
let spec = SessionSpec {
turns: 3,
acu_per_turn: 2,
catalog_commands: 2,
catalog_desc_len: 8,
agent_chunks_per_turn: 4,
agent_chunk_len: 16,
rewind_points: 0,
files_per_rewind: 0,
file_content_len: 0,
};
let (info, _dir) = synth::prepare_session(root.path(), cwd.path(), &spec).await;
let replayed = load_updates_for_replay_at(info.id.0.as_ref(), root.path())
.expect("load_updates_for_replay_at")
.unwrap_or_default();
let count = |pred: fn(&acp::SessionUpdate) -> bool| replayed.iter().filter(|u| pred(u)).count();
let users = count(|u| matches!(u, acp::SessionUpdate::UserMessageChunk(_)));
let acus = count(|u| matches!(u, acp::SessionUpdate::AvailableCommandsUpdate(_)));
let agents = count(|u| matches!(u, acp::SessionUpdate::AgentMessageChunk(_)));
assert_eq!(users, spec.turns, "one user chunk per turn");
assert_eq!(
acus,
spec.turns * spec.acu_per_turn,
"every ACU is preserved"
);
assert_eq!(
agents,
spec.turns * spec.agent_chunks_per_turn,
"every agent chunk is preserved"
);
assert_eq!(
replayed.len(),
spec.turns * (1 + spec.acu_per_turn + spec.agent_chunks_per_turn),
"no update is dropped or duplicated by the typed replay reader"
);
}
/// The adapter-driven bench generator reaches its byte target and emits updates
/// the production reader parses. Synchronous because `synthesize_to_target_bytes`
/// drives the adapter on its own runtime.
#[test]
fn synthesize_to_target_bytes_reaches_target_and_parses() {
let root = TempDir::new().unwrap();
let target: u64 = 8 * 1024;
let info = synth::synthesize_to_target_bytes(root.path(), target);
let adapter = JsonlStorageAdapter::with_root(root.path().to_path_buf());
let updates_path = adapter.updates_file_path(&info).expect("updates path");
let len = std::fs::metadata(&updates_path)
.expect("stat updates.jsonl")
.len();
assert!(
len >= target,
"updates.jsonl ({len} B) reached the target ({target} B)"
);
let replayed = load_updates_for_replay_at(info.id.0.as_ref(), root.path())
.expect("load_updates_for_replay_at")
.unwrap_or_default();
assert!(!replayed.is_empty(), "the emitted updates parse back");
}