diff --git a/SOURCE_REV b/SOURCE_REV index 93d1715..c2214f4 100644 --- a/SOURCE_REV +++ b/SOURCE_REV @@ -1 +1 @@ -d02693a856a54f1030695b36b91d276e96b30b23 +91d8cf309110a3b879c1b8198f7525aed545dfb4 diff --git a/crates/codegen/xai-grok-config-types/src/lib.rs b/crates/codegen/xai-grok-config-types/src/lib.rs index a7fda0a..63408a7 100644 --- a/crates/codegen/xai-grok-config-types/src/lib.rs +++ b/crates/codegen/xai-grok-config-types/src/lib.rs @@ -697,8 +697,9 @@ pub struct RemoteSettings { pub managed_mcps_enabled: Option, #[serde(default)] pub managed_mcp_gateway_tools_enabled: Option, - /// Fleet kill switch for the **external OTEL** stream (customer - /// collectors). Restrictive-only by construction: there is deliberately + /// Remote-policy disable lever for the **external OTEL** stream (customer + /// collectors); feeds `ExternalOtelRemotePolicy.force_disable`. + /// Restrictive-only by construction: there is deliberately /// no `external_otel_enabled` remote field — remote settings are fetched /// per-run and never persisted, so a remote "enable" could never reach /// init; org-wide enable ships via managed config instead. Applied diff --git a/crates/codegen/xai-grok-http/src/lib.rs b/crates/codegen/xai-grok-http/src/lib.rs index 4a6e8c1..9552383 100644 --- a/crates/codegen/xai-grok-http/src/lib.rs +++ b/crates/codegen/xai-grok-http/src/lib.rs @@ -5,11 +5,11 @@ //! provides four clients for non-sampling traffic (the first three //! public and cached, the last crate-internal and built on demand): //! -//! - `shared_client` -- a `OnceLock`-cached async client for general +//! - `shared_client`: a `OnceLock`-cached async client for general //! use (telemetry, feedback, settings, etc.). -//! - `shared_upload_client` -- a `OnceLock`-cached client for GCS +//! - `shared_upload_client`: a `OnceLock`-cached client for GCS //! uploads with aggressive connection pool eviction. -//! - `shared_blocking_client` -- a blocking client for the early +//! - `shared_startup_blocking_client`: a blocking client for the early //! model prefetch (runs before the async runtime is available). //! - `fresh_http1_client` -- a crate-internal, on-demand, pool-less //! HTTP/1.1 client used by `send_with_retry_escaping_pool` for the @@ -30,6 +30,42 @@ use std::sync::OnceLock; use xai_grok_workspace::permission::ClientType; +/// Per-attempt ceiling for a startup `/settings` or `/v1/models` fetch; raising +/// it delays how soon the background refresh gives up and retries. +pub const STARTUP_FETCH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); +/// Cap on non-interactive boot auth (token refresh or cold-start mint); a mint +/// that exceeds it leaves the leader session-less and is retried off the +/// readiness path. +pub const STARTUP_AUTH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); +/// Ceiling on a single startup token-refresh round trip, kept separate from +/// `STARTUP_FETCH_TIMEOUT` so the two tune independently; on timeout the caller +/// proceeds with cached or no credentials and re-auths later. +pub const STARTUP_AUTH_REFRESH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); +/// Outer bound on a single settings-reapply task, which drives up to +/// `SETTINGS_FETCH_MAX_ATTEMPTS` fetches. +pub const SETTINGS_REAPPLY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); +/// Attempt budget for the background settings fetch; bounds proxy load while +/// still covering a brief blip. +pub const SETTINGS_FETCH_MAX_ATTEMPTS: u32 = 3; +// A `401` self-heal may add one more bounded fetch beyond this cap; that fetch +// is cut off fail-closed and retried later, so the cap only needs to cover the +// common path. +const _: () = assert!( + SETTINGS_REAPPLY_TIMEOUT.as_millis() + > STARTUP_FETCH_TIMEOUT.as_millis() * (1 + SETTINGS_FETCH_MAX_ATTEMPTS as u128), + "SETTINGS_REAPPLY_TIMEOUT must exceed STARTUP_FETCH_TIMEOUT * (1 + MAX_ATTEMPTS)" +); + +/// Lower bound for a client's leader-connect timeout: a slow-but-valid boot +/// (bounded startup auth plus the rest of leader startup and the connect +/// handshake) must never be aborted. The pager bounds its connect by this value, +/// reached via the shell's `http` re-export. +pub const MIN_CLIENT_CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(120); +const _: () = assert!( + MIN_CLIENT_CONNECT_TIMEOUT.as_millis() >= 2 * STARTUP_AUTH_TIMEOUT.as_millis(), + "MIN_CLIENT_CONNECT_TIMEOUT must stay >= 2x STARTUP_AUTH_TIMEOUT" +); + /// Startup span timer, local to this crate. /// /// Replaces `xai_grok_shell::instrumentation_timer!`, which cannot be referenced @@ -487,7 +523,8 @@ where Err(last_err.expect("send_with_retry_escaping_pool ran at least one attempt")) } -/// Returns a shared [`reqwest::blocking::Client`], creating it on first call. +/// Shared blocking client for startup fetches. Carries `STARTUP_FETCH_TIMEOUT` +/// as the connect+read ceiling; do not reuse for long-lived requests. /// /// This avoids redundant TLS certificate loading for blocking HTTP calls /// (e.g., model prefetching during startup). The blocking client is separate @@ -501,14 +538,14 @@ where /// (~60-100s; 30s is a conservative default) closes it. The HTTP/2 keepalive-ping /// setters that `shared_client()` uses are NOT exposed on reqwest's blocking /// `ClientBuilder` (0.12), so only the idle/TCP-eviction half applies here. -pub fn shared_blocking_client() -> reqwest::blocking::Client { +pub fn shared_startup_blocking_client() -> reqwest::blocking::Client { static BLOCKING_CLIENT: OnceLock = OnceLock::new(); BLOCKING_CLIENT .get_or_init(|| { let _timer = startup_timer!("startup.http_blocking_client_build"); reqwest::blocking::Client::builder() - .connect_timeout(std::time::Duration::from_secs(30)) - .timeout(std::time::Duration::from_secs(30)) + .connect_timeout(STARTUP_FETCH_TIMEOUT) + .timeout(STARTUP_FETCH_TIMEOUT) .user_agent(process_user_agent_string()) .pool_idle_timeout(std::time::Duration::from_secs(30)) .tcp_keepalive(std::time::Duration::from_secs(30)) diff --git a/crates/codegen/xai-grok-pager-bin/src/main.rs b/crates/codegen/xai-grok-pager-bin/src/main.rs index f52d448..2053e71 100644 --- a/crates/codegen/xai-grok-pager-bin/src/main.rs +++ b/crates/codegen/xai-grok-pager-bin/src/main.rs @@ -1026,6 +1026,12 @@ async fn run_agent_command( } } }); + if matches!( + agent_args.mode, + Some(AgentCmd::Leader(_) | AgentCmd::Stdio | AgentCmd::Headless(_) | AgentCmd::Serve(_)) + ) { + xai_grok_shell::agent::app::suppress_otel(); + } init_tracing_simple("agent"); let _otel_guard = xai_grok_telemetry::otel_layer::otel_guard(); xai_grok_telemetry::instrumentation::install_panic_hook(); diff --git a/crates/codegen/xai-grok-pager-pty-harness/src/pty.rs b/crates/codegen/xai-grok-pager-pty-harness/src/pty.rs index e227d7d..c8c063f 100644 --- a/crates/codegen/xai-grok-pager-pty-harness/src/pty.rs +++ b/crates/codegen/xai-grok-pager-pty-harness/src/pty.rs @@ -574,6 +574,7 @@ const HOST_TERMINAL_ENV_VARS: &[&str] = &[ "CMUX_SOCKET_PATH", "CMUX_PANEL_ID", "CMUX_BUNDLE_ID", + "HERDR_ENV", // Embedded editor markers (embedded_editor_from_env). "NVIM", "NVIM_LISTEN_ADDRESS", diff --git a/crates/codegen/xai-grok-pager-render/src/gboom/mod.rs b/crates/codegen/xai-grok-pager-render/src/gboom/mod.rs index 42915d2..02b1068 100644 --- a/crates/codegen/xai-grok-pager-render/src/gboom/mod.rs +++ b/crates/codegen/xai-grok-pager-render/src/gboom/mod.rs @@ -1,6 +1,9 @@ //! `/gboom` easter egg: a tiny single-level raycaster shooter rendered in //! the terminal via the kitty graphics protocol. //! +//! Not production code — this is for fun and is not maintained to the +//! standards in `crates/codegen/AGENTS.md`. +//! //! Typing `/gboom` (and nothing else) opens a modal overlay — the same //! surface the imagine-video player uses — and streams PNG frames via //! per-frame kitty `a=T` retransmission at the ~30 fps animation tick. The diff --git a/crates/codegen/xai-grok-pager-render/src/terminal/mod.rs b/crates/codegen/xai-grok-pager-render/src/terminal/mod.rs index 9a72beb..79a1e47 100644 --- a/crates/codegen/xai-grok-pager-render/src/terminal/mod.rs +++ b/crates/codegen/xai-grok-pager-render/src/terminal/mod.rs @@ -15,6 +15,7 @@ pub mod image; pub mod keyboard; pub mod overlay; pub(crate) mod probe; +pub mod term_version; pub mod tmux_probe; pub mod xtversion; @@ -27,6 +28,7 @@ pub use keyboard::{ KeyboardCapabilities, ModifierDelivery, ModifierFate, keyboard_capabilities, keyboard_capabilities_for_host, }; +pub use term_version::{TermVersion, TermVersionSource}; #[cfg(test)] mod test; @@ -195,6 +197,15 @@ pub enum MultiplexerKind { /// cmux (Ghostty-backed macOS terminal multiplexer). #[strum(to_string = "cmux")] Cmux, + /// herdr, a libghostty-backed agent multiplexer + /// ([ogulcancelik/herdr](https://github.com/ogulcancelik/herdr)). + /// + /// Its embedded emulator answers CSI queries itself, so it counts as + /// CSI-intercepting. The accepted cost: a herdr pane typically has no + /// other version signal (brand `Unknown`, no `TERM_PROGRAM_VERSION`), but + /// the XTVERSION reply describes herdr's engine rather than the host. + #[strum(to_string = "herdr")] + Herdr, /// No recognized multiplexer detected (does not rule out unknown ones). #[default] #[strum(to_string = "None detected")] @@ -203,9 +214,10 @@ pub enum MultiplexerKind { impl MultiplexerKind { /// Whether this multiplexer intercepts CSI queries (e.g. XTVERSION) - /// instead of passing them through to the outer terminal. + /// instead of passing them through to the outer terminal. See + /// [`Self::Herdr`] for the version signal herdr gives up by being here. pub fn intercepts_csi_queries(self) -> bool { - matches!(self, Self::Tmux | Self::Screen | Self::Zellij) + matches!(self, Self::Tmux | Self::Screen | Self::Zellij | Self::Herdr) } } @@ -275,10 +287,16 @@ pub struct TerminalContext { /// Value of tmux's `extended-keys` global option (`"on"`, `"off"`, /// `"always"`); populated only when `multiplexer == Tmux`. pub tmux_extended_keys: Option, - /// The `TERM_PROGRAM_VERSION` environment variable (e.g. `"3.5.6"` for - /// iTerm2, `"1.1.3"` for Ghostty). Used for version-gating features - /// that require a minimum terminal version. + /// The `TERM_PROGRAM_VERSION` environment variable, falling back to + /// `LC_TERMINAL_VERSION` (e.g. `"3.5.6"` for iTerm2, `"1.1.3"` for + /// Ghostty). Used for version-gating features that require a minimum + /// terminal version. Raw and **ungated**: inside tmux this is tmux's own + /// version. For a brand-corroborated value use [`Self::env_term_version`]. pub term_program_version: Option, + /// The brand-corroborated counterpart to `term_program_version` (see + /// [`term_version`]). Resolved once in [`build_terminal_context_from_env`], + /// so it does not re-derive if `brand` or `vte_version` change afterwards. + pub env_term_version: Option, } impl TerminalContext { @@ -553,12 +571,25 @@ impl TerminalContext { } } + /// The best available terminal version and the source that reported it. + /// + /// Only the environment arm exists today. Final precedence is + /// `da2 > xtversion > env` — probe arms insert **above** it, since a live + /// self-report cannot be inherited or go stale. + pub fn term_version(&self) -> (String, TermVersionSource) { + match &self.env_term_version { + Some(v) => (v.version.clone(), v.source), + None => (String::new(), TermVersionSource::None), + } + } + /// Extract a flat snapshot of terminal details for telemetry. pub fn telemetry_snapshot(&self) -> xai_grok_telemetry::events::TerminalTelemetry { let os = crate::host::HostOs::current(); let server = crate::host::DisplayServer::current(); let kb = self.keyboard_capabilities(); let route = crate::clipboard::clipboard_route(); + let (term_version, term_version_source) = self.term_version(); xai_grok_telemetry::events::TerminalTelemetry { brand: self.brand.to_string(), multiplexer: self.multiplexer.to_string(), @@ -572,6 +603,8 @@ impl TerminalContext { enter_modifier_fate: kb.enter_modifier.to_string(), tmux_version: self.tmux_version_or_na().to_owned(), xtversion: xtversion::detected().unwrap_or("").to_owned(), + term_version, + term_version_source: term_version_source.to_string(), hyperlink_osc8: self.hyperlink_capabilities().osc8.to_string(), hyperlink_skip_reason: self.hyperlink_skip_reason().unwrap_or("none").to_owned(), clipboard_route: route.to_string(), @@ -879,8 +912,11 @@ fn infer_byobu_backend_from_mux_markers(env: &HashMap) -> Option /// 1. Explicit `BYOBU_BACKEND` beats generic `TMUX`/`STY` clues. /// 2. `TMUX` beats `ZELLIJ` (tmux can nest inside Zellij but not vice-versa). /// 3. `STY` (GNU screen) is only chosen when neither `TMUX` nor `ZELLIJ` is set. -/// 4. cmux markers classify only when no tmux/zellij/screen (or explicit -/// Byobu backend) won — so a real mux nested inside cmux still wins. +/// 4. herdr and cmux markers classify only when no tmux/zellij/screen (or +/// explicit Byobu backend) won — a real mux nested inside either still +/// wins. Between the two, `HERDR_ENV` beats the `CMUX_*` markers: the +/// shape they appear in together is herdr running in a cmux panel, where +/// the `CMUX_*` values are inherited rather than fresh. /// /// This ensures one deterministic classification even when multiple markers /// are present (e.g., an inherited `ZELLIJ` var inside a tmux pane). @@ -897,7 +933,9 @@ pub fn detect_multiplexer_from_env(env: &HashMap) -> Multiplexer } // Standard multiplexer markers, tmux > Zellij > screen. - // Nested real multiplexers inside cmux must win over cmux itself. + // Nested real multiplexers inside herdr/cmux must win over the host mux. + // A herdr daemon first started from tmux freezes that TMUX into every pane: + // classified tmux here, so OSC 52 gets a tmux DCS wrap herdr renders as text. if env_get(env, "TMUX").is_some() { return MultiplexerKind::Tmux; } @@ -907,6 +945,12 @@ pub fn detect_multiplexer_from_env(env: &HashMap) -> Multiplexer if env_get(env, "STY").is_some() { return MultiplexerKind::Screen; } + // herdr sets HERDR_ENV=1 in every pane, overrides TERM to xterm-256color + // and never sets TERM_PROGRAM, so this marker is its only documented, + // stable signal. + if env_get(env, "HERDR_ENV").is_some() { + return MultiplexerKind::Herdr; + } // cmux sets non-empty CMUX_SOCKET_PATH / CMUX_PANEL_ID / CMUX_BUNDLE_ID; // CMUX_SOCKET may be present but empty — env_get filters empties. if env_get(env, "CMUX_SOCKET_PATH").is_some() @@ -952,6 +996,8 @@ pub fn build_terminal_context_from_env(env: &HashMap) -> Termina let term_program_version = env_get(env, "TERM_PROGRAM_VERSION") .or_else(|| env_get(env, "LC_TERMINAL_VERSION")) .map(|s| s.to_owned()); + // Resolved here: the brand each version var must corroborate is in hand. + let env_term_version = term_version::detect_env_term_version(env, brand); TerminalContext { brand, @@ -967,6 +1013,7 @@ pub fn build_terminal_context_from_env(env: &HashMap) -> Termina vte_version, tmux_extended_keys: None, term_program_version, + env_term_version, } } diff --git a/crates/codegen/xai-grok-pager-render/src/terminal/term_version.rs b/crates/codegen/xai-grok-pager-render/src/terminal/term_version.rs new file mode 100644 index 0000000..ef33962 --- /dev/null +++ b/crates/codegen/xai-grok-pager-render/src/terminal/term_version.rs @@ -0,0 +1,312 @@ +//! Terminal version capture from environment variables. +//! +//! **A version variable is trusted only when the environment corroborates the +//! brand it belongs to.** These variables cross process, SSH and multiplexer +//! boundaries, so an uncorroborated version is as likely to describe another +//! program as the terminal drawing our output. For a variable that is itself +//! the brand marker (`WEZTERM_VERSION`, `VTE_VERSION`), corroboration means no +//! stronger marker outranked it in `detect_terminal_brand_from_env`. +//! +//! Never read: `ZELLIJ_VERSION`, which would make an Alacritty pane inside +//! Zellij report Zellij's number as its own, and `KONSOLE_VERSION`, which has +//! no `TerminalName::Konsole` to attach to. + +use std::collections::HashMap; + +use super::{TerminalName, terminal_name_from_term_program}; + +/// Which environment variable produced a [`TermVersion`]. +/// +/// The rendered labels are stable telemetry values — do not rename them. +#[derive(Clone, Copy, Debug, Eq, PartialEq, strum::Display)] +#[strum(serialize_all = "snake_case")] +pub enum TermVersionSource { + None, + /// `TERM_PROGRAM_VERSION`, or its SSH-surviving `LC_TERMINAL_VERSION` + /// mirror (iTerm2 only). + TermProgram, + #[strum(serialize = "wezterm")] + WezTerm, + Vte, +} + +/// A terminal version together with the source that reported it. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct TermVersion { + /// Raw, exactly as the source reported it (trimmed). Shapes vary by + /// terminal: `"3.5.6"`, `"20240203-110809-5046fc22"`, `"7402"`. + pub version: String, + pub source: TermVersionSource, +} + +impl TermVersion { + fn new(version: &str, source: TermVersionSource) -> Self { + Self { + version: version.to_owned(), + source, + } + } +} + +/// Whether the brand `TERM_PROGRAM` names vouches for `env_brand`'s version. +/// +/// Identity, widened for VS Code forks: they export `TERM_PROGRAM=vscode` from +/// the same host process that writes their brand marker and draws our output, +/// so the version is that host's and `brand` records whose numbering it is. +/// One-directional, so a leaked marker cannot borrow another brand's version; +/// Zed is excluded as it is not an xterm.js host. +fn corroborates(named: TerminalName, env_brand: TerminalName) -> bool { + named == env_brand + || (named == TerminalName::VsCode + && matches!( + env_brand, + TerminalName::VsCode | TerminalName::Cursor | TerminalName::Windsurf + )) +} + +/// Look up an env value, trimmed; `env_get` alone would pass whitespace. +fn env_trimmed<'a>(env: &'a HashMap, key: &str) -> Option<&'a str> { + let value = super::env_get(env, key)?.trim(); + (!value.is_empty()).then_some(value) +} + +/// Resolve the terminal version from the environment, taking the first +/// variable corroborated by `env_brand`. +/// +/// `env_brand` is the *pre-refinement* brand, before +/// `refine_unknown_brand_for_host` may rewrite `TerminalContext::brand`: +/// the native-Windows `Unknown -> WindowsTerminal` guess must not license a +/// version attribution. +pub(super) fn detect_env_term_version( + env: &HashMap, + env_brand: TerminalName, +) -> Option { + // tmux >= 3.2 exports TERM_PROGRAM=tmux and its own TERM_PROGRAM_VERSION, + // which ungated would land on whichever brand marker survived inside the + // tmux server environment. + let named_brand = env_trimmed(env, "TERM_PROGRAM").and_then(terminal_name_from_term_program); + if let Some(version) = env_trimmed(env, "TERM_PROGRAM_VERSION") + && named_brand.is_some_and(|named| corroborates(named, env_brand)) + { + return Some(TermVersion::new(version, TermVersionSource::TermProgram)); + } + + // LC_TERMINAL_VERSION survives SSH where TERM_PROGRAM_VERSION does not, + // but only iTerm2 sets the pair. + if let Some(version) = env_trimmed(env, "LC_TERMINAL_VERSION") + && env_trimmed(env, "LC_TERMINAL").is_some_and(|v| v.eq_ignore_ascii_case("iterm2")) + && env_brand == TerminalName::Iterm2 + { + return Some(TermVersion::new(version, TermVersionSource::TermProgram)); + } + + if let Some(version) = env_trimmed(env, "WEZTERM_VERSION") + && env_brand == TerminalName::WezTerm + { + return Some(TermVersion::new(version, TermVersionSource::WezTerm)); + } + + // Brand-only: `TerminalContext::is_vte_based()` also accepts a present + // `vte_version`, which is the candidate here — routing the gate through it + // would make it vacuous. + if let Some(version) = env_trimmed(env, "VTE_VERSION") + && env_brand.is_vte_based() + { + return Some(TermVersion::new(version, TermVersionSource::Vte)); + } + + None +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::terminal::{build_terminal_context_from_env, env_from}; + + fn resolved(pairs: &[(&str, &str)]) -> (String, TermVersionSource) { + build_terminal_context_from_env(&env_from(pairs)).term_version() + } + + /// Unlike the others, this reads ambient host state and warms process-wide + /// caches; the asserted fields still come only from the injected map. + fn snapshot(pairs: &[(&str, &str)]) -> xai_grok_telemetry::events::TerminalTelemetry { + build_terminal_context_from_env(&env_from(pairs)).telemetry_snapshot() + } + + #[test] + fn source_labels_are_pinned() { + assert_eq!(TermVersionSource::None.to_string(), "none"); + assert_eq!(TermVersionSource::TermProgram.to_string(), "term_program"); + assert_eq!(TermVersionSource::WezTerm.to_string(), "wezterm"); + assert_eq!(TermVersionSource::Vte.to_string(), "vte"); + } + + #[test] + fn term_program_version_wins_when_term_program_names_the_brand() { + let (version, source) = resolved(&[ + ("TERM_PROGRAM", "iTerm.app"), + ("TERM_PROGRAM_VERSION", "3.5.6"), + ("LC_TERMINAL", "iTerm2"), + ("LC_TERMINAL_VERSION", "3.4.0"), + ]); + assert_eq!(version, "3.5.6"); + assert_eq!(source, TermVersionSource::TermProgram); + } + + #[test] + fn term_program_version_wins_over_wezterm_version() { + let (version, source) = resolved(&[ + ("TERM_PROGRAM", "WezTerm"), + ("TERM_PROGRAM_VERSION", "20240203-110809-5046fc22"), + ("WEZTERM_VERSION", "20230712-072601-f4abf8fd"), + ]); + assert_eq!(version, "20240203-110809-5046fc22"); + assert_eq!(source, TermVersionSource::TermProgram); + } + + #[test] + fn vscode_and_its_forks_keep_the_vscode_host_version() { + let (version, source) = resolved(&[ + ("VSCODE_GIT_ASKPASS_MAIN", "/home/u/.vscode-server/askpass"), + ("TERM_PROGRAM", "vscode"), + ("TERM_PROGRAM_VERSION", "1.99.3"), + ]); + assert_eq!(version, "1.99.3"); + assert_eq!(source, TermVersionSource::TermProgram); + + let fork = build_terminal_context_from_env(&env_from(&[ + ("CURSOR_TRACE_ID", "abc123"), + ("TERM_PROGRAM", "vscode"), + ("TERM_PROGRAM_VERSION", "1.99.3"), + ])); + assert_eq!(fork.brand, TerminalName::Cursor); + let (version, source) = fork.term_version(); + assert_eq!(version, "1.99.3"); + assert_eq!(source, TermVersionSource::TermProgram); + } + + #[test] + fn the_vscode_widening_is_one_directional() { + // TERM_PROGRAM names Zed while the brand chain resolves Cursor — + // neither vouches for the other. + let (version, source) = resolved(&[ + ("CURSOR_TRACE_ID", "abc123"), + ("TERM_PROGRAM", "zed"), + ("TERM_PROGRAM_VERSION", "0.180.0"), + ]); + assert_eq!(version, ""); + assert_eq!(source, TermVersionSource::None); + } + + #[test] + fn lc_terminal_version_wins_when_term_program_version_is_absent() { + let (version, source) = + resolved(&[("LC_TERMINAL", "iTerm2"), ("LC_TERMINAL_VERSION", "3.5.6")]); + assert_eq!(version, "3.5.6"); + assert_eq!(source, TermVersionSource::TermProgram); + } + + #[test] + fn lc_terminal_version_ignored_without_lc_terminal() { + let (version, source) = resolved(&[ + ("ITERM_SESSION_ID", "w0t0p0:1234"), + ("LC_TERMINAL_VERSION", "3.5.6"), + ]); + assert_eq!(version, ""); + assert_eq!(source, TermVersionSource::None); + } + + #[test] + fn wezterm_version_wins_over_vte_version() { + let (version, source) = resolved(&[ + ("WEZTERM_VERSION", "20240203-110809-5046fc22"), + ("VTE_VERSION", "7402"), + ]); + assert_eq!(version, "20240203-110809-5046fc22"); + assert_eq!(source, TermVersionSource::WezTerm); + } + + #[test] + fn vte_version_wins_for_a_vte_brand() { + let (version, source) = resolved(&[("VTE_VERSION", "7402")]); + assert_eq!(version, "7402"); + assert_eq!(source, TermVersionSource::Vte); + } + + #[test] + fn wezterm_version_ignored_for_another_brand() { + // An inherited WEZTERM_VERSION is not the Ghostty session's version. + let (version, source) = resolved(&[ + ("TERM_PROGRAM", "Ghostty"), + ("WEZTERM_VERSION", "20240203-110809-5046fc22"), + ]); + assert_eq!(version, ""); + assert_eq!(source, TermVersionSource::None); + } + + #[test] + fn vte_version_ignored_for_a_non_vte_brand() { + let (version, source) = resolved(&[("TERM", "alacritty"), ("VTE_VERSION", "7402")]); + assert_eq!(version, ""); + assert_eq!(source, TermVersionSource::None); + } + + #[test] + fn tmux_term_program_version_is_not_the_terminal_version() { + // tmux >= 3.2 exports TERM_PROGRAM=tmux plus its own version, and + // iTerm2's releases are also 3.5.x — an ungated value would be + // indistinguishable from a real one. + let (version, source) = resolved(&[ + ("TMUX", "/tmp/tmux-501/default,12345,0"), + ("TERM_PROGRAM", "tmux"), + ("TERM_PROGRAM_VERSION", "3.5"), + ("ITERM_SESSION_ID", "w0t0p0:1234"), + ]); + assert_eq!(version, ""); + assert_eq!(source, TermVersionSource::None); + } + + #[test] + fn iterm2_in_tmux_falls_through_to_lc_terminal_version() { + let (version, source) = resolved(&[ + ("TMUX", "/tmp/tmux-501/default,12345,0"), + ("TERM_PROGRAM", "tmux"), + ("TERM_PROGRAM_VERSION", "3.5"), + ("LC_TERMINAL", "iTerm2"), + ("LC_TERMINAL_VERSION", "3.5.6"), + ]); + assert_eq!(version, "3.5.6"); + assert_eq!(source, TermVersionSource::TermProgram); + } + + #[test] + fn blank_version_is_absent() { + let (version, source) = resolved(&[ + ("TERM_PROGRAM", "Ghostty"), + ("TERM_PROGRAM_VERSION", " \t "), + ]); + assert_eq!(version, ""); + assert_eq!(source, TermVersionSource::None); + } + + #[test] + fn surrounding_whitespace_is_trimmed() { + let (version, source) = resolved(&[ + ("TERM_PROGRAM", "Ghostty"), + ("TERM_PROGRAM_VERSION", " 1.1.3\n"), + ]); + assert_eq!(version, "1.1.3"); + assert_eq!(source, TermVersionSource::TermProgram); + } + + #[test] + fn telemetry_snapshot_carries_version_and_source() { + let populated = snapshot(&[("VTE_VERSION", "7402")]); + assert_eq!(populated.term_version, "7402"); + assert_eq!(populated.term_version_source, "vte"); + + let empty = snapshot(&[("TERM", "xterm-256color")]); + assert_eq!(empty.term_version, ""); + assert_eq!(empty.term_version_source, "none"); + } +} diff --git a/crates/codegen/xai-grok-pager-render/src/terminal/test.rs b/crates/codegen/xai-grok-pager-render/src/terminal/test.rs index 2eeb37c..4494c44 100644 --- a/crates/codegen/xai-grok-pager-render/src/terminal/test.rs +++ b/crates/codegen/xai-grok-pager-render/src/terminal/test.rs @@ -416,6 +416,34 @@ fn mux_tmux_nested_inside_cmux_wins() { assert_eq!(detect_multiplexer_from_env(&env), MultiplexerKind::Tmux); } +#[test] +fn mux_herdr_from_herdr_env() { + let env = env_from(&[("HERDR_ENV", "1")]); + assert_eq!(detect_multiplexer_from_env(&env), MultiplexerKind::Herdr); +} + +#[test] +fn mux_tmux_nested_inside_herdr_wins() { + // tmux started inside a herdr pane, or a stale TMUX frozen into the pane by + // herdr's daemon — indistinguishable from the env, and tmux wins either way. + let env = env_from(&[ + ("TMUX", "/tmp/tmux-501/default,12345,0"), + ("HERDR_ENV", "1"), + ]); + assert_eq!(detect_multiplexer_from_env(&env), MultiplexerKind::Tmux); +} + +#[test] +fn mux_herdr_nested_inside_cmux_wins() { + // herdr in a cmux panel inherits the CMUX_* markers; the inner layer wins. + let env = env_from(&[ + ("CMUX_SOCKET_PATH", "/tmp/cmux.sock"), + ("CMUX_PANEL_ID", "1"), + ("HERDR_ENV", "1"), + ]); + assert_eq!(detect_multiplexer_from_env(&env), MultiplexerKind::Herdr); +} + // -- ambiguous marker precedence ------------------------------------------ #[test] @@ -1122,6 +1150,21 @@ fn mux_zellij_not_from_version_only() { ); } +#[test] +fn zellij_version_is_never_the_terminal_version() { + // The multiplexer's version must never be attributed to the emulator. + let env = env_from(&[ + ("TERM", "alacritty"), + ("ZELLIJ", "0"), + ("ZELLIJ_SESSION_NAME", "main"), + ("ZELLIJ_VERSION", "0.43.1"), + ]); + let ctx = build_terminal_context_from_env(&env); + assert_eq!(ctx.brand, TerminalName::Alacritty); + assert_eq!(ctx.multiplexer, MultiplexerKind::Zellij); + assert_eq!(ctx.term_version(), (String::new(), TermVersionSource::None)); +} + // -- Byobu inference edge cases ------------------------------------------- #[test] @@ -1692,6 +1735,28 @@ fn shift_enter_available_unknown_with_multiplexer() { assert!(!ctx.shift_enter_unavailable()); } +#[test] +fn herdr_over_ssh_pane_does_not_skip_kitty_keyboard() { + // A real herdr pane reached over SSH: no TERM_PROGRAM, so the brand stays + // Unknown. HERDR_ENV is what keeps this out of the unknown-no-multiplexer + // skip, which would otherwise drop Shift+Enter (herdr speaks KKP). + let env = env_from(&[ + ("HERDR_ENV", "1"), + ("HERDR_PANE_ID", "3"), + ("TERM", "xterm-256color"), + ("COLORTERM", "truecolor"), + ("SSH_CONNECTION", "10.0.0.1 52000 10.0.0.2 22"), + ]); + let ctx = build_terminal_context_from_env(&env); + assert!(ctx.is_ssh); + assert_eq!(ctx.brand, TerminalName::Unknown); + // shift_enter_unavailable() reads env_brand, not brand. + assert_eq!(ctx.env_brand, TerminalName::Unknown); + assert_eq!(ctx.multiplexer, MultiplexerKind::Herdr); + assert_eq!(ctx.kitty_skip_reason(), None); + assert!(!ctx.shift_enter_unavailable()); +} + #[test] fn ctrl_dot_unreliable_on_vte() { let ctx = TerminalContext { diff --git a/crates/codegen/xai-grok-pager-render/src/terminal/xtversion.rs b/crates/codegen/xai-grok-pager-render/src/terminal/xtversion.rs index f347b92..2c7fbce 100644 --- a/crates/codegen/xai-grok-pager-render/src/terminal/xtversion.rs +++ b/crates/codegen/xai-grok-pager-render/src/terminal/xtversion.rs @@ -193,6 +193,10 @@ mod tests { !gate_allows_probe(&ctx(brand, MultiplexerKind::Tmux)), "{brand:?} under tmux should be skipped" ); + assert!( + !gate_allows_probe(&ctx(brand, MultiplexerKind::Herdr)), + "{brand:?} under herdr should be skipped" + ); } // JediTerm renders the query as garbage and must never be probed. assert!(!gate_allows_probe(&ctx( diff --git a/crates/codegen/xai-grok-pager/docs/user-guide/19-plan-mode.md b/crates/codegen/xai-grok-pager/docs/user-guide/19-plan-mode.md index 2869480..71a415b 100644 --- a/crates/codegen/xai-grok-pager/docs/user-guide/19-plan-mode.md +++ b/crates/codegen/xai-grok-pager/docs/user-guide/19-plan-mode.md @@ -77,6 +77,7 @@ Scroll the plan with the arrow keys or `j`/`k`. The action bar shows these short | `a` | Approve the plan and start building. With pending comments, this reads `approve w/ comments` and sends them alongside the approval. | | `s` | Request changes. Focus moves to the prompt so you can type revision notes; press `Enter` to send them. | | `c` | Comment on the selected line or line range. | +| `y` | Copy the full plan to the clipboard. | | `q` | Quit plan -- abandon the plan without approving and turn plan mode off. | Press `Tab` to move focus between the plan preview and the prompt. diff --git a/crates/codegen/xai-grok-pager/src/acp/mod.rs b/crates/codegen/xai-grok-pager/src/acp/mod.rs index 04999cf..cf43974 100644 --- a/crates/codegen/xai-grok-pager/src/acp/mod.rs +++ b/crates/codegen/xai-grok-pager/src/acp/mod.rs @@ -212,7 +212,7 @@ pub async fn connect(cancel: &CancellationToken, flags: ConnectFlags) -> Result< startup_auth_metadata(&auth_methods); let (needs_login, login_label, login_method_id, auth_start_mode, auth_meta) = - eager_auth_or_login_fallback( + bounded_eager_auth( &tx, &auth_methods, default_auth_method_id.as_ref(), @@ -325,7 +325,7 @@ pub async fn connect_via_leader( startup_auth_metadata(&auth_methods); let (needs_login, login_label, login_method_id, auth_start_mode, auth_meta) = - eager_auth_or_login_fallback( + bounded_eager_auth( &tx, &auth_methods, default_auth_method_id.as_ref(), @@ -710,6 +710,49 @@ async fn eager_auth_or_login_fallback( } } +/// [`eager_auth_or_login_fallback`] bounded by `STARTUP_AUTH_REFRESH_TIMEOUT`, +/// so a hung agent cannot gate the first draw. On timeout the inputs pass +/// through unchanged and the agent finishes authentication in the background. +async fn bounded_eager_auth( + tx: &AcpAgentTx, + auth_methods: &[acp::AuthMethod], + default_auth_method_id: Option<&acp::AuthMethodId>, + needs_login: bool, + login_label: Option, + login_method_id: Option, + auth_start_mode: AuthStartMode, +) -> ( + bool, + Option, + Option, + AuthStartMode, + Option, +) { + match tokio::time::timeout( + xai_grok_shell::http::STARTUP_AUTH_REFRESH_TIMEOUT, + eager_auth_or_login_fallback( + tx, + auth_methods, + default_auth_method_id, + needs_login, + login_label.clone(), + login_method_id.clone(), + auth_start_mode, + ), + ) + .await + { + Ok(resolved) => resolved, + Err(_) => ( + needs_login, + login_label, + login_method_id, + auth_start_mode, + None, + ), + } +} + /// Authenticate with the agent using the agent's chosen default method. /// /// Prefer `defaultAuthMethodId` from initialize meta when present and listed. diff --git a/crates/codegen/xai-grok-pager/src/acp/spawn.rs b/crates/codegen/xai-grok-pager/src/acp/spawn.rs index ead8e86..9fba58b 100644 --- a/crates/codegen/xai-grok-pager/src/acp/spawn.rs +++ b/crates/codegen/xai-grok-pager/src/acp/spawn.rs @@ -188,8 +188,13 @@ pub async fn spawn_grok_shell( // re-login). No-op where the OS listener is unavailable. auth_manager.start_system_power_listener(); + // Both embedded-agent paths (`--no-leader` and leader fallback) converge + // here, so the agent's external-OTEL gate is applied exactly once, before boot. + xai_grok_shell::agent::app::apply_otel_config(&auth_manager, &agent_config.grok_com_config); + // Best-effort refresh of managed policy before bootstrap reads it (repairs a wrong-identity/missing - // cache). Never errors — the OS-protected system/MDM layers still apply. + // cache). Never errors — the OS-protected system/MDM layers still apply, and every network step + // inside is bounded (SESSION_START_AUTH_DEADLINE / SyncBudget::SessionStart). xai_grok_shell::managed_config::ensure_managed_policy_present(&auth_manager).await; // Run the full bootstrap sequence: config resolution, process-level @@ -200,6 +205,9 @@ pub async fn spawn_grok_shell( models_manager .list_models(RefreshStrategy::OnlineIfUncached) .await; + // Self-heal a cold-cache/failed boot fetch once the backend recovers, + // matching the leader and stdio paths. + models_manager.spawn_background_refresh(); let agent_cancel = cancel.child_token(); let (acp_client, acp_agent) = acp_channels(); diff --git a/crates/codegen/xai-grok-pager/src/app/acp_handler/settings.rs b/crates/codegen/xai-grok-pager/src/app/acp_handler/settings.rs index 2c424bb..7a4cd9d 100644 --- a/crates/codegen/xai-grok-pager/src/app/acp_handler/settings.rs +++ b/crates/codegen/xai-grok-pager/src/app/acp_handler/settings.rs @@ -58,6 +58,20 @@ pub(super) fn handle_settings_update(notif: &acp::ExtNotification, app: &mut App return false; }; + // Reseed this process's remote-campaign cache. In leader mode no in-process + // agent seeds the TUI process, and the bounded startup prefetch can miss — + // without this reseed a remote campaign stays invisible to + // `resolve_dismissable_campaigns`, so a `/model` pick never records its + // dismissal and the leader re-nudges every new session. Idempotent in + // embedded mode, where the in-process agent seeds the same cache. + if let Some(campaigns) = update.campaigns.clone() { + let rs = xai_grok_shell::util::config::RemoteSettings { + campaigns, + ..Default::default() + }; + xai_grok_shell::util::config::set_remote_campaigns_from_settings(Some(&rs)); + } + if let Some(v) = update.auto_permission_mode_enabled { // Keep the pager's auto-permission-mode gate live with the remote settings // remote tier (the leader caches it agent-side; the pager process needs @@ -523,6 +537,11 @@ pub(super) struct PagerSettingsUpdate { // remote_settings also emits gen-ordered `x.ai/announcements/update` // (emit_announcements_if_changed), and a gen-less apply on this path could // clobber a newer push. Single ingest path: handle_announcements_update. + /// Remote campaigns snapshot. `Some` whenever the shell has settings + /// (empty = campaigns withdrawn); `None`/omitted (settings-less push, + /// older shell) must leave this process's campaign cache untouched. + #[serde(default)] + campaigns: Option>, #[serde(default)] gate_message: Option, #[serde(default)] diff --git a/crates/codegen/xai-grok-pager/src/app/agent_view/plan.rs b/crates/codegen/xai-grok-pager/src/app/agent_view/plan.rs index fcb66f9..b3072d1 100644 --- a/crates/codegen/xai-grok-pager/src/app/agent_view/plan.rs +++ b/crates/codegen/xai-grok-pager/src/app/agent_view/plan.rs @@ -79,7 +79,7 @@ impl AgentView { /// the shell-read file body), then falls back to the on-disk plan file. /// Request body first keeps file-backed previews working when the path /// resolution fails or the file disappears between intercept and open. - fn plan_body_for_preview(&self) -> Option { + pub(super) fn plan_body_for_preview(&self) -> Option { if let Some(content) = self .plan_approval_view .as_ref() diff --git a/crates/codegen/xai-grok-pager/src/app/agent_view/render.rs b/crates/codegen/xai-grok-pager/src/app/agent_view/render.rs index 100c0ec..20576f1 100644 --- a/crates/codegen/xai-grok-pager/src/app/agent_view/render.rs +++ b/crates/codegen/xai-grok-pager/src/app/agent_view/render.rs @@ -128,7 +128,7 @@ impl AgentView { ] } } - PlanApprovalFocus::Preview => vec![], + PlanApprovalFocus::Preview => vec![HintItem::new(key!('y'), "copy plan")], } } /// Returns the *exact* hints the bottom shortcuts bar would render right now. @@ -208,6 +208,7 @@ impl AgentView { } else { let mut h = vec![ HintItem::new(key!('c'), "comment"), + HintItem::new(key!('y'), "copy plan"), HintItem::new(key!('f', CONTROL), "fullscreen"), ]; if !self.plan_comments.is_empty() { @@ -3197,6 +3198,7 @@ impl AgentView { } else { let mut h = vec![ HintItem::new(key!('c'), "comment"), + HintItem::new(key!('y'), "copy plan"), HintItem::new(key!('f', CONTROL), "fullscreen"), ]; if !self.plan_comments.is_empty() { @@ -3289,6 +3291,7 @@ impl AgentView { .with_pending(pending_hint) .render(layout.shortcuts, buf); } + let line_viewer_toast = self.active_toast_message().map(|s| s.to_string()); let is_plan_viewer = self.is_plan_viewer(); let has_plan_comments = !self.plan_comments.is_empty(); let casual_commenting = self.is_casual_commenting(); @@ -3335,6 +3338,26 @@ impl AgentView { &theme, effective_comment_count, ); + let toast_area = viewer + .last_popup_area + .or(viewer.last_modal_area) + .unwrap_or(overlay_area); + if let Some(ref msg) = line_viewer_toast + && toast_area.height > 0 + && let Some(toast_text) = fit_toast_text(msg, toast_area.width.saturating_sub(1)) + { + let w = toast_text.chars().count() as u16; + let tx = toast_area.right().saturating_sub(w + 1); + let ty = toast_area.bottom().saturating_sub(1); + for (i, ch) in toast_text.chars().enumerate() { + if let Some(cell) = buf.cell_mut((tx + i as u16, ty)) { + cell.set_char(ch); + cell.fg = theme.accent_user; + cell.bg = theme.bg_base; + cell.modifier = ratatui::prelude::Modifier::BOLD; + } + } + } let in_plan_approval = self.plan_approval_view.is_some(); let on_comment = in_plan_approval && viewer @@ -3360,11 +3383,15 @@ impl AgentView { } else { h.push(HintItem::new(key!('a'), "approve")); } + h.push(HintItem::new(key!('y'), "copy plan")); h.push(HintItem::new(key!('q'), "quit plan")); h.push(HintItem::new(key!(Tab), "prompt")); h } else if in_plan_approval { - let mut h = vec![HintItem::new(key!('c'), "comment")]; + let mut h = vec![ + HintItem::new(key!('c'), "comment"), + HintItem::new(key!('y'), "copy plan"), + ]; if approval_has_comments { h.push(HintItem::new(key!('s'), "send")); } else { @@ -3390,9 +3417,13 @@ impl AgentView { vec![ HintItem::new(key!(Enter), "edit"), HintItem::new(key!('x'), "delete"), + HintItem::new(key!('y'), "copy plan"), ] } else { - vec![HintItem::new(key!('c'), "comment")] + vec![ + HintItem::new(key!('c'), "comment"), + HintItem::new(key!('y'), "copy plan"), + ] }; if has_plan_comments { h.push(HintItem::new(key!('s'), "send")); diff --git a/crates/codegen/xai-grok-pager/src/app/agent_view/viewer.rs b/crates/codegen/xai-grok-pager/src/app/agent_view/viewer.rs index 855d1c3..859d905 100644 --- a/crates/codegen/xai-grok-pager/src/app/agent_view/viewer.rs +++ b/crates/codegen/xai-grok-pager/src/app/agent_view/viewer.rs @@ -54,6 +54,19 @@ impl AgentView { } } + pub(super) fn copy_plan_full(&mut self) -> InputOutcome { + let text = self + .line_viewer + .as_ref() + .and_then(|v| v.markdown_content_for_feedback()) + .filter(|s| !s.is_empty()) + .or_else(|| self.plan_body_for_preview()); + if let Some(text) = text { + self.copy_to_clipboard(&text); + } + InputOutcome::Changed + } + /// Handle a key event while the line viewer is open. pub(super) fn handle_line_viewer_key(&mut self, key: &KeyEvent) -> InputOutcome { let in_plan_approval = self.plan_approval_view.is_some(); @@ -185,8 +198,10 @@ impl AgentView { self.confirm_line_viewer(false); return InputOutcome::Changed; } - // y: copy selected line(s) to system clipboard. if key!('y').matches(key) { + if self.is_plan_viewer() { + return self.copy_plan_full(); + } if let Some(ref viewer) = self.line_viewer { let text = if viewer.list_state.visual_mode { if let Some(ref range) = viewer.list_state.multi_range() { @@ -219,8 +234,10 @@ impl AgentView { } return InputOutcome::Changed; } - // Y: copy filename to clipboard. if key!('Y').matches(key) { + if self.is_plan_viewer() { + return InputOutcome::Changed; + } if let Some(ref viewer) = self.line_viewer { let name = viewer .title_override @@ -392,6 +409,7 @@ impl AgentView { let abandon_area = viewer.plan_ref().and_then(|p| p.abandon_button_area); let approve_area = viewer.plan_ref().and_then(|p| p.approve_button_area); let comment_btn_area = viewer.plan_ref().and_then(|p| p.comment_button_area); + let copy_btn_area = viewer.plan_ref().and_then(|p| p.copy_button_area); // Cached `is_plan_viewer()` so we don't need to call self while // the line_viewer is mutably borrowed below. let is_plan_preview = @@ -440,6 +458,9 @@ impl AgentView { // patterns just above. return InputOutcome::Changed; } + if copy_btn_area.is_some_and(|a| a.contains((mouse.column, mouse.row).into())) { + return self.copy_plan_full(); + } if send_area.is_some_and(|a| a.contains((mouse.column, mouse.row).into())) { if self.plan_approval_view.is_some() { if let Some(ref mut pav) = self.plan_approval_view { @@ -528,6 +549,13 @@ impl AgentView { viewer.plan_mut().comment_hovered = comment_btn_hover; changed = true; } + let copy_btn_hover = + copy_btn_area.is_some_and(|a| a.contains((mouse.column, mouse.row).into())); + let prev_copy_btn = viewer.plan_ref().is_some_and(|p| p.copy_hovered); + if copy_btn_hover != prev_copy_btn { + viewer.plan_mut().copy_hovered = copy_btn_hover; + changed = true; + } if self.plan_approval_view.is_some() && let Some(area) = popup_area && area.contains((mouse.column, mouse.row).into()) diff --git a/crates/codegen/xai-grok-pager/src/app/display_refresh_startup.rs b/crates/codegen/xai-grok-pager/src/app/display_refresh_startup.rs index 0b89885..9e4de14 100644 --- a/crates/codegen/xai-grok-pager/src/app/display_refresh_startup.rs +++ b/crates/codegen/xai-grok-pager/src/app/display_refresh_startup.rs @@ -123,6 +123,8 @@ fn spawn_terminal_and_display_refresh_telemetry(tel: StartupTel) { terminal.tmux_version = %t.tmux_version, terminal.term_var = %t.term_var, terminal.xtversion = %t.xtversion, + terminal.term_version = %t.term_version, + terminal.term_version_source = %t.term_version_source, ) .entered(); tracing::info!("terminal environment detected"); diff --git a/crates/codegen/xai-grok-pager/src/app/effects/mod.rs b/crates/codegen/xai-grok-pager/src/app/effects/mod.rs index d9f63f6..33ce532 100644 --- a/crates/codegen/xai-grok-pager/src/app/effects/mod.rs +++ b/crates/codegen/xai-grok-pager/src/app/effects/mod.rs @@ -4219,10 +4219,11 @@ pub(crate) fn execute( .to_owned() }); xai_grok_shell::remote::fetch_settings_blocking( - &proxy_base, - &auth, - None, - ) + &proxy_base, + &auth, + None, + ) + .into_option() }) .await .ok() diff --git a/crates/codegen/xai-grok-pager/src/app/mod.rs b/crates/codegen/xai-grok-pager/src/app/mod.rs index 853a6d9..59c88f1 100644 --- a/crates/codegen/xai-grok-pager/src/app/mod.rs +++ b/crates/codegen/xai-grok-pager/src/app/mod.rs @@ -458,6 +458,23 @@ fn resolve_hunk_tracker_mode( .find(|s| !s.is_empty()) .map(str::to_owned) } +/// Run a connect future bounded by cancellation and `timeout`, so a hung leader +/// or embedded spawn cannot strand the user on a blank screen. +async fn bounded_connect( + cancel: &CancellationToken, + timeout: std::time::Duration, + target: &str, + connect: impl std::future::Future>, +) -> anyhow::Result { + tokio::select! { + biased; + () = cancel.cancelled() => Err(anyhow::anyhow!("startup cancelled before {target} connected")), + r = connect => r, + () = tokio::time::sleep(timeout) => { + Err(anyhow::anyhow!("timed out after {}s connecting to {target}", timeout.as_secs())) + } + } +} /// Main entry point: connect to agent, init terminal, run event loop, restore. /// /// If a session ID is provided via `--resume` / `--load` / `--continue`, the @@ -487,9 +504,16 @@ pub async fn run( xai_grok_shell::auth::GrokComConfig::default() } }; - let refreshed_auth = xai_grok_shell::auth::try_ensure_fresh_auth(&grok_com_config).await; - let early_prefetch = - xai_grok_shell::agent::models::start_early_prefetch_with_auth(refreshed_auth); + let refreshed_auth = tokio::time::timeout( + xai_grok_shell::http::STARTUP_AUTH_REFRESH_TIMEOUT, + xai_grok_shell::auth::try_ensure_fresh_auth(&grok_com_config), + ) + .await + .unwrap_or(None); + let early_prefetch = match refreshed_auth { + Some(auth) => xai_grok_shell::agent::models::start_early_prefetch_with_auth(Some(auth)), + None => xai_grok_shell::agent::models::start_early_prefetch(Some(grok_com_config.clone())), + }; xai_grok_shell::agent::mvp_agent::warm_async_http_client(); tokio::task::spawn_blocking(|| {}); if let Ok(cwd) = std::env::current_dir() { @@ -641,23 +665,6 @@ pub async fn run( default_yolo_mode: launch_yolo.yolo, default_auto_mode: launch_auto && !launch_yolo.yolo, }; - let mut connection = if use_leader { - let conn = crate::acp::connect_via_leader(&cancel, connect_flags, &raw_config).await?; - tracing::info!( - elapsed_ms = startup_start.elapsed().as_millis() as u64, - "Connected via leader" - ); - conn - } else { - let conn = crate::acp::connect(&cancel, connect_flags).await?; - tracing::info!( - elapsed_ms = startup_start.elapsed().as_millis() as u64, - "Connected directly (non-leader)" - ); - conn - }; - let agent_guard = - crate::acp::spawn::AgentShutdownGuard::new(cancel.clone(), connection.agent_thread.take()); let mut config_watcher = crate::appearance::ConfigWatcher::start().await?; let alt_screen_config_mode = config_watcher.current().alt_screen; let term_ctx = crate::terminal::terminal_context(); @@ -728,6 +735,53 @@ pub async fn run( if let Some(ref t) = session_title { set_terminal_title(t); } + const CONNECT_UI_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); + let fallback_flags = use_leader.then(|| connect_flags.clone()); + let primary_target = if use_leader { + "the grok leader" + } else { + "the embedded agent" + }; + let connect_result = bounded_connect(&cancel, CONNECT_UI_TIMEOUT, primary_target, async { + if use_leader { + crate::acp::connect_via_leader(&cancel, connect_flags, &raw_config).await + } else { + crate::acp::connect(&cancel, connect_flags).await + } + }) + .await; + let (connect_result, embedded_fallback) = match connect_result { + Err(e) if use_leader && !cancel.is_cancelled() => { + tracing::warn!(error = %e, "leader connect failed; falling back to embedded agent"); + let flags = fallback_flags.expect("set on the use_leader path"); + let fallback = + bounded_connect(&cancel, CONNECT_UI_TIMEOUT, "the embedded agent", async { + crate::acp::connect(&cancel, flags).await + }) + .await; + (fallback, true) + } + other => (other, false), + }; + let mut connection = match connect_result { + Ok(conn) => { + tracing::info!( + elapsed_ms = startup_start.elapsed().as_millis() as u64, + use_leader = use_leader && !embedded_fallback, + embedded_fallback, + "Connected" + ); + conn + } + Err(e) => { + crate::unified_log::flush_blocking().await; + let _ = restore_terminal(terminal, writer_thread, screen_mode); + cancel.cancel(); + return Err(e); + } + }; + let agent_guard = + crate::acp::spawn::AgentShutdownGuard::new(cancel.clone(), connection.agent_thread.take()); let effective_args = PagerArgs { resume_session: None, load_session: None, @@ -1486,6 +1540,31 @@ mod tests { let toml_str = format!("[cli]\nuse_leader = {enabled}"); toml::from_str(&toml_str).unwrap() } + #[tokio::test] + async fn bounded_connect_times_out_when_the_target_stalls() { + let cancel = CancellationToken::new(); + let r = bounded_connect( + &cancel, + std::time::Duration::from_millis(20), + "the test target", + std::future::pending::>(), + ) + .await; + assert!(r.is_err_and(|e| e.to_string().contains("timed out"))); + } + #[tokio::test] + async fn bounded_connect_returns_err_on_cancel() { + let cancel = CancellationToken::new(); + cancel.cancel(); + let r = bounded_connect( + &cancel, + std::time::Duration::from_secs(60), + "the test target", + std::future::pending::>(), + ) + .await; + assert!(r.is_err_and(|e| e.to_string().contains("cancelled"))); + } #[test] fn terminal_title_strips_control_characters() { assert_eq!( diff --git a/crates/codegen/xai-grok-pager/src/app/subagent.rs b/crates/codegen/xai-grok-pager/src/app/subagent.rs index a5d0f99..8c118e0 100644 --- a/crates/codegen/xai-grok-pager/src/app/subagent.rs +++ b/crates/codegen/xai-grok-pager/src/app/subagent.rs @@ -6,6 +6,7 @@ use serde::Deserialize; use std::sync::Arc; use std::time::Instant; +use xai_grok_shell::session::storage::{ReplayEmission, stream_replay_updates_at}; /// Enriched subagent tracking info. /// /// Keyed by `child_session_id` in `AgentView::subagent_sessions`. @@ -107,15 +108,24 @@ struct SubagentMetaSlice { #[serde(default)] worktree_path: Option, } +/// Grok home for the replay path. In production this is just `grok_home()`; the +/// whole test override below is `#[cfg(test)]`, so no thread-local or dead +/// always-false branch ships in release. +#[cfg(not(test))] +fn effective_grok_home() -> std::path::PathBuf { + xai_grok_shell::util::grok_home::grok_home() +} +#[cfg(test)] thread_local! { static REPLAY_GROK_HOME: std::cell::RefCell> = const { std::cell::RefCell::new(None) }; } -/// Override grok home for disk-replay unit tests (thread-local; production never sets this). +/// Override grok home for disk-replay unit tests (thread-local). #[cfg(test)] pub(crate) fn set_replay_grok_home_for_tests(home: Option) { REPLAY_GROK_HOME.with(|h| *h.borrow_mut() = home); } +#[cfg(test)] fn effective_grok_home() -> std::path::PathBuf { if let Some(home) = REPLAY_GROK_HOME.with(|h| h.borrow().clone()) { return home; @@ -161,39 +171,31 @@ fn enrich_from_meta_with_home( info.child_cwd = meta.child_cwd.map(Arc::from); info.worktree_path = meta.worktree_path.map(Arc::from); } -/// Best-effort replay of inherited conversation for a child subagent. -/// -/// Reads `updates.jsonl` from the child session directory via -/// [`load_updates_for_replay`], then feeds ACP updates through the child's -/// tracker with replay semantics. No-ops when the child session or file is -/// missing (typical for a live spawn before the shell has persisted updates). +/// Best-effort replay of a child's inherited conversation, streamed one typed +/// update at a time so a large inherited transcript is not materialized as a +/// full `Vec` of typed structs (peak stays near the file size rather than +/// several multiples of it). No-ops when the child session or file is missing. pub(crate) fn replay_inherited_updates( child_view: &mut crate::app::agent_view::AgentView, child_session_id: &str, ) { let home = effective_grok_home(); - let updates = match xai_grok_shell::session::storage::load_updates_for_replay_at( - child_session_id, - &home, - ) { - Ok(Some(u)) => u, - Ok(None) => return, - Err(e) => { - tracing::debug!(session_id = %child_session_id, error = %e, "failed to load updates for replay"); - return; - } - }; let replay_meta = crate::acp::meta::NotificationMeta { is_replay: true, ..Default::default() }; - let replayed_any = !updates.is_empty(); - for update in updates { + let outcome = match stream_replay_updates_at(child_session_id, &home, |update| { child_view .session .handle_update(update, &replay_meta, &mut child_view.scrollback); - } - if replayed_any { + }) { + Ok(outcome) => outcome, + Err(e) => { + tracing::warn!(session_id = %child_session_id, error = %e, "failed to read updates for replay"); + return; + } + }; + if outcome == ReplayEmission::Emitted { crate::memory_release::release_retained_memory_with("subagent-replay"); } } diff --git a/crates/codegen/xai-grok-pager/src/diagnostics/fix_tests.rs b/crates/codegen/xai-grok-pager/src/diagnostics/fix_tests.rs index dd2e980..f804bd2 100644 --- a/crates/codegen/xai-grok-pager/src/diagnostics/fix_tests.rs +++ b/crates/codegen/xai-grok-pager/src/diagnostics/fix_tests.rs @@ -73,6 +73,7 @@ fn terminal() -> TerminalContext { vte_version: None, tmux_extended_keys: None, term_program_version: None, + env_term_version: None, } } diff --git a/crates/codegen/xai-grok-pager/src/doctor_cmd/json.rs b/crates/codegen/xai-grok-pager/src/doctor_cmd/json.rs index f436330..fecbea2 100644 --- a/crates/codegen/xai-grok-pager/src/doctor_cmd/json.rs +++ b/crates/codegen/xai-grok-pager/src/doctor_cmd/json.rs @@ -379,6 +379,7 @@ pub(super) fn multiplexer(kind: MultiplexerKind) -> &'static str { MultiplexerKind::Screen => "screen", MultiplexerKind::Zellij => "zellij", MultiplexerKind::Cmux => "cmux", + MultiplexerKind::Herdr => "herdr", MultiplexerKind::Undetected => "undetected", } } diff --git a/crates/codegen/xai-grok-pager/src/doctor_cmd/tests.rs b/crates/codegen/xai-grok-pager/src/doctor_cmd/tests.rs index 84e638d..838bcc5 100644 --- a/crates/codegen/xai-grok-pager/src/doctor_cmd/tests.rs +++ b/crates/codegen/xai-grok-pager/src/doctor_cmd/tests.rs @@ -864,10 +864,11 @@ fn stable_mapping_tables_are_complete() { MultiplexerKind::Screen, MultiplexerKind::Zellij, MultiplexerKind::Cmux, + MultiplexerKind::Herdr, MultiplexerKind::Undetected, ] .map(multiplexer), - ["tmux", "screen", "zellij", "cmux", "undetected"] + ["tmux", "screen", "zellij", "cmux", "herdr", "undetected"] ); assert_eq!( [ diff --git a/crates/codegen/xai-grok-pager/src/input/mouse.rs b/crates/codegen/xai-grok-pager/src/input/mouse.rs index 8d06c97..d8bcdc6 100644 --- a/crates/codegen/xai-grok-pager/src/input/mouse.rs +++ b/crates/codegen/xai-grok-pager/src/input/mouse.rs @@ -277,12 +277,15 @@ impl ScrollConfigOverrides { } /// Multiplexers that re-encode mouse into their own SGR stream (tmux with -/// `mouse on`, screen, zellij all re-emit per pane). Cmux is a Ghostty-backed -/// passthrough and keeps the outer brand's stream. +/// `mouse on`, screen, zellij, herdr all re-emit per pane). Cmux is a +/// Ghostty-backed passthrough and keeps the outer brand's stream. fn multiplexer_reencodes_mouse(multiplexer: MultiplexerKind) -> bool { matches!( multiplexer, - MultiplexerKind::Tmux | MultiplexerKind::Screen | MultiplexerKind::Zellij + MultiplexerKind::Tmux + | MultiplexerKind::Screen + | MultiplexerKind::Zellij + | MultiplexerKind::Herdr ) } @@ -311,8 +314,8 @@ impl ScrollConfig { } /// Derive scroll normalization defaults from detected terminal metadata. - /// tmux/screen/zellij re-encode mouse into their own SGR stream, so the - /// outer brand's events-per-tick/pacing calibration describes the wrong + /// tmux/screen/zellij/herdr re-encode mouse into their own SGR stream, so + /// the outer brand's events-per-tick/pacing calibration describes the wrong /// producer — trusting an outer ept=3 profile under tmux under-counts 3x /// per notch when the multiplexer re-chunks to one event. Under those /// multiplexers the brand table is replaced by a conservative ept=1 diff --git a/crates/codegen/xai-grok-pager/src/input/mouse/tests.rs b/crates/codegen/xai-grok-pager/src/input/mouse/tests.rs index 4a67ab0..9a2e707 100644 --- a/crates/codegen/xai-grok-pager/src/input/mouse/tests.rs +++ b/crates/codegen/xai-grok-pager/src/input/mouse/tests.rs @@ -1147,7 +1147,7 @@ fn ghostty_duplicate_reports_do_not_feed_accel_banding() { #[test] fn multiplexed_sessions_use_conservative_profile_regardless_of_brand() { - // tmux/screen/zellij re-encode mouse into their own SGR stream, so + // tmux/screen/zellij/herdr re-encode mouse into their own SGR stream, so // the outer brand's ept/pacing calibration is wrong under them: the // conservative ept=1 shape applies no matter the brand. Cmux is a // passthrough and Undetected means no multiplexer — both keep the @@ -1172,6 +1172,7 @@ fn multiplexed_sessions_use_conservative_profile_regardless_of_brand() { MultiplexerKind::Tmux, MultiplexerKind::Screen, MultiplexerKind::Zellij, + MultiplexerKind::Herdr, ] { for brand in brands { let cfg = ScrollConfig::from_terminal_context(brand, mux, Default::default()); diff --git a/crates/codegen/xai-grok-pager/src/views/file_search/line_viewer.rs b/crates/codegen/xai-grok-pager/src/views/file_search/line_viewer.rs index 35ad2f9..61094cd 100644 --- a/crates/codegen/xai-grok-pager/src/views/file_search/line_viewer.rs +++ b/crates/codegen/xai-grok-pager/src/views/file_search/line_viewer.rs @@ -553,6 +553,8 @@ pub struct PlanViewerExtras { pub comment_hovered: bool, pub abandon_button_area: Option, pub abandon_hovered: bool, + pub copy_button_area: Option, + pub copy_hovered: bool, pub last_click_at: Option, pub gutter_drag_start: Option, pub gutter_drag_end: Option, @@ -822,8 +824,7 @@ impl LineViewerState { } /// Whether the plan modal should render the action-button footer. - /// True for both modes: plan-approval (q/c/s|a) and casual - /// (c/s — quit via the close-X button instead of a footer button). + /// True for plan-approval and casual plan preview (not plain file preview). pub fn show_footer(&self) -> bool { self.plan .as_ref() @@ -1477,10 +1478,7 @@ pub fn render_line_viewer( // Buttons use the same `key bold + label dim` treatment as // `render_modal_shortcuts`, sit in a single row separated by // ` | `, centered within the modal frame. - // - // - Plan-approval: q quit | c comment | s send / a approve - // - Casual preview: c comment | s send (no `q` — - // the close-X button handles closing in casual mode) + // Casual preview omits `q` (close via the X button). if viewer.show_footer() && inner.height >= 2 { let div_y = inner.y + inner.height - 2; let div_style = Style::default().fg(theme.gray_dim).bg(theme.bg_base); @@ -1492,11 +1490,15 @@ pub fn render_line_viewer( let abandon_hovered = viewer.plan_ref().is_some_and(|p| p.abandon_hovered); let comment_hovered = viewer.plan_ref().is_some_and(|p| p.comment_hovered); let approve_hovered = viewer.plan_ref().is_some_and(|p| p.approve_hovered); + let copy_hovered = viewer.plan_ref().is_some_and(|p| p.copy_hovered); let is_approval = viewer.feedback_active(); let comment_spans = build_shortcut_button('c', "comment", comment_hovered, theme); let comment_w: u16 = comment_spans.iter().map(|s| s.width() as u16).sum(); + let copy_spans = build_shortcut_button('y', "copy plan", copy_hovered, theme); + let copy_w: u16 = copy_spans.iter().map(|s| s.width() as u16).sum(); + // In approval mode, always show `a approve`. When there are // pending review comments, also show `s revise` (request changes). // In approval mode, show `a approve` (or `a approve w/ comments` @@ -1555,18 +1557,20 @@ pub fn render_line_viewer( let sep_w: u16 = 5; // separator is fixed-width ASCII; matches modal_window.rs:565 let sep_style = Style::default().fg(theme.gray_dim).bg(theme.bg_base); - // Total width: [action] + (sep + revise)? + sep + comment[badge?] + (sep + quit)? - let mut total_w: u16 = 0; + let mut base_w: u16 = 0; if action_w > 0 { - total_w = total_w.saturating_add(action_w).saturating_add(sep_w); + base_w = base_w.saturating_add(action_w).saturating_add(sep_w); } if revise_w > 0 { - total_w = total_w.saturating_add(revise_w).saturating_add(sep_w); + base_w = base_w.saturating_add(revise_w).saturating_add(sep_w); } - total_w = total_w.saturating_add(comment_w).saturating_add(badge_w); + base_w = base_w.saturating_add(comment_w).saturating_add(badge_w); if let Some((_, w)) = &quit_spans { - total_w = total_w.saturating_add(sep_w).saturating_add(*w); + base_w = base_w.saturating_add(sep_w).saturating_add(*w); } + let with_copy_w = base_w.saturating_add(sep_w).saturating_add(copy_w); + let show_copy = with_copy_w <= inner.width; + let total_w = if show_copy { with_copy_w } else { base_w }; if total_w <= inner.width { let mut x = inner.x + (inner.width - total_w) / 2; @@ -1620,6 +1624,20 @@ pub fn render_line_viewer( x += badge_w; } + if show_copy { + buf.set_string(x, bottom_y, separator, sep_style); + x += sep_w; + let copy_x = x; + for span in ©_spans { + let w = span.width() as u16; + buf.set_span(x, bottom_y, span, w); + x += w; + } + viewer.plan_mut().copy_button_area = Some(Rect::new(copy_x, bottom_y, copy_w, 1)); + } else { + viewer.plan_mut().copy_button_area = None; + } + // Quit button — approval mode only. if let Some((spans, w)) = quit_spans { buf.set_string(x, bottom_y, separator, sep_style); @@ -1640,6 +1658,7 @@ pub fn render_line_viewer( let plan = viewer.plan_mut(); plan.approve_button_area = None; plan.comment_button_area = None; + plan.copy_button_area = None; plan.abandon_button_area = None; } } @@ -1683,6 +1702,20 @@ mod tests { ); } + #[test] + fn plan_preview_exposes_full_raw_markdown_for_copy() { + let body = "# Plan\n\n- Do the thing\n- Then ship"; + let mut viewer = LineViewerState::open_markdown_content("plan.md", body.to_owned(), None) + .expect("markdown content should open"); + viewer.kind = LineViewerKind::PlanPreview; + viewer.prepare_layout(80, 20); + + assert_eq!( + viewer.markdown_content_for_feedback().as_deref(), + Some(body) + ); + } + fn line_text(line: &Line<'_>) -> String { line.spans .iter() diff --git a/crates/codegen/xai-grok-pager/src/views/prompt_widget/mod.rs b/crates/codegen/xai-grok-pager/src/views/prompt_widget/mod.rs index 3ac8bb5..3c96319 100644 --- a/crates/codegen/xai-grok-pager/src/views/prompt_widget/mod.rs +++ b/crates/codegen/xai-grok-pager/src/views/prompt_widget/mod.rs @@ -1826,6 +1826,8 @@ impl PromptWidget { terminal.multiplexer = %evt.terminal.multiplexer, terminal.is_ssh = evt.terminal.is_ssh, terminal.term_var = %evt.terminal.term_var, + terminal.term_version = %evt.terminal.term_version, + terminal.term_version_source = %evt.terminal.term_version_source, key.code = %evt.key_code, key.modifiers = %evt.key_modifiers, key.kind = %evt.key_kind, diff --git a/crates/codegen/xai-grok-pager/tests/pty_e2e/embedded_mode_boots_without_hanging_on_blocked_backend.rs b/crates/codegen/xai-grok-pager/tests/pty_e2e/embedded_mode_boots_without_hanging_on_blocked_backend.rs new file mode 100644 index 0000000..8559550 --- /dev/null +++ b/crates/codegen/xai-grok-pager/tests/pty_e2e/embedded_mode_boots_without_hanging_on_blocked_backend.rs @@ -0,0 +1,64 @@ +// Per-test-case module for the `pty_e2e` integration test crate. +#[allow(unused_imports)] +use super::common::*; + +/// 1a. **Embedded mode (`--no-leader`) boots without hanging on a blocked backend.** +/// +/// Enterprise deployments set `[cli] use_leader = false` and point at their own +/// backend, often with the grok.com proxy blocked. A TCP listener that accepts +/// but never replies stands in for that endpoint, so every startup HTTP call +/// stalls until the client's own bounded timeout fires. The welcome screen must +/// render anyway; a hang here means some boot path went unbounded. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[ignore] +async fn embedded_mode_boots_without_hanging_on_blocked_backend() { + // Accept connections but never respond, holding the streams open. + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + let addr = listener.local_addr().expect("addr"); + std::thread::spawn(move || { + let mut held = Vec::new(); + for stream in listener.incoming() { + match stream { + Ok(s) => held.push(s), + Err(_) => break, + } + } + }); + let base = format!("http://{addr}/v1"); + + let home = tempfile::tempdir().expect("home"); + let grok_home = home.path().join(".grok"); + std::fs::create_dir_all(&grok_home).unwrap(); + let env = [ + ("HOME", home.path().to_str().unwrap()), + ("GROK_HOME", grok_home.to_str().unwrap()), + ("XAI_API_KEY", "test-key-for-ci"), + ("GROK_CLI_CHAT_PROXY_BASE_URL", base.as_str()), + ("GROK_XAI_API_BASE_URL", base.as_str()), + ("GROK_TELEMETRY_ENABLED", "false"), + ("GROK_FEEDBACK_ENABLED", "false"), + ("GROK_TRACE_UPLOAD", "false"), + ]; + + let binary = pager_binary().expect("resolve pager binary"); + let mut harness = PtyHarness::new_inherited_env( + &binary, + DEFAULT_ROWS, + DEFAULT_COLS, + &["--no-leader"], + &env, + None, + ) + .expect("spawn pager"); + + // 30s exceeds the sum of the bounded startup fetches (auth + settings + + // catalog, ~5s each), so a timeout here means an unbounded wait, not a + // slow one. + harness + .wait_for_text(WELCOME_SCREEN_SENTINEL, Duration::from_secs(30)) + .expect("embedded welcome must render despite the blocked backend (boot went unbounded)"); + let screen = harness.screen_contents(); + assert!(!screen.contains("panicked"), "panic on screen:\n{screen}"); + + harness.quit().expect("clean quit"); +} diff --git a/crates/codegen/xai-grok-pager/tests/pty_e2e_smoke.rs b/crates/codegen/xai-grok-pager/tests/pty_e2e_smoke.rs index 7c5c18e..e5be7d4 100644 --- a/crates/codegen/xai-grok-pager/tests/pty_e2e_smoke.rs +++ b/crates/codegen/xai-grok-pager/tests/pty_e2e_smoke.rs @@ -17,6 +17,8 @@ mod auto_compact_top_row; mod basename_path_demo_pty; #[path = "pty_e2e/doubled_lines_out_of_band_repro.rs"] mod doubled_lines_out_of_band_repro; +#[path = "pty_e2e/embedded_mode_boots_without_hanging_on_blocked_backend.rs"] +mod embedded_mode_boots_without_hanging_on_blocked_backend; #[path = "pty_e2e/initial_prompt_positional_auto_submits.rs"] mod initial_prompt_positional_auto_submits; #[path = "pty_e2e/input_echoes_at_idle_prompt.rs"] diff --git a/crates/codegen/xai-grok-shell/Cargo.toml b/crates/codegen/xai-grok-shell/Cargo.toml index 57fba20..a10700d 100644 --- a/crates/codegen/xai-grok-shell/Cargo.toml +++ b/crates/codegen/xai-grok-shell/Cargo.toml @@ -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 diff --git a/crates/codegen/xai-grok-shell/benches/fork_copy.rs b/crates/codegen/xai-grok-shell/benches/fork_copy.rs index e0cfe9b..eb4b273 100644 --- a/crates/codegen/xai-grok-shell/benches/fork_copy.rs +++ b/crates/codegen/xai-grok-shell/benches/fork_copy.rs @@ -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 { - 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") diff --git a/crates/codegen/xai-grok-shell/src/agent/app.rs b/crates/codegen/xai-grok-shell/src/agent/app.rs index 613bce4..24e4ee7 100644 --- a/crates/codegen/xai-grok-shell/src/agent/app.rs +++ b/crates/codegen/xai-grok-shell/src/agent/app.rs @@ -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> { - 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, memory_config: Option, ) -> 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 = 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 = crate::auth::try_noninteractive_auth_no_mint(ctx).await; // ── Phase 6b: Legacy devbox auth migration ───────────────────────────── let auth: Option = migrate_devbox_auth_if_legacy(auth, &agent_config).await; - let auth_for_prefetch: Option = 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 = - 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) { + 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, + legacy: Option, + } + 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, 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>>> = + 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 diff --git a/crates/codegen/xai-grok-shell/src/agent/config.rs b/crates/codegen/xai-grok-shell/src/agent/config.rs index 35f0cc7..ff33fda 100644 --- a/crates/codegen/xai-grok-shell/src/agent/config.rs +++ b/crates/codegen/xai-grok-shell/src/agent/config.rs @@ -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.` from Claude-compat `managed_settings.json`. `Some(true)` /// indicates a force-off signal from a Mac-MDM-style admin policy. diff --git a/crates/codegen/xai-grok-shell/src/agent/init.rs b/crates/codegen/xai-grok-shell/src/agent/init.rs index 0eca53f..799c685 100644 --- a/crates/codegen/xai-grok-shell/src/agent/init.rs +++ b/crates/codegen/xai-grok-shell/src/agent/init.rs @@ -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); } diff --git a/crates/codegen/xai-grok-shell/src/agent/mod.rs b/crates/codegen/xai-grok-shell/src/agent/mod.rs index 8ae4d41..27ccd83 100644 --- a/crates/codegen/xai-grok-shell/src/agent/mod.rs +++ b/crates/codegen/xai-grok-shell/src/agent/mod.rs @@ -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; diff --git a/crates/codegen/xai-grok-shell/src/agent/models.rs b/crates/codegen/xai-grok-shell/src/agent/models.rs index d3413a2..d5ad1dc 100644 --- a/crates/codegen/xai-grok-shell/src/agent/models.rs +++ b/crates/codegen/xai-grok-shell/src/agent/models.rs @@ -1,5 +1,7 @@ //! Model fetching, resolution, and management. +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; @@ -29,10 +31,6 @@ pub(crate) enum ModelFetchAuth { impl ModelFetchAuth { /// custom_endpoint > session > deployment > API key. - /// - /// A `deployment_key` outranks an ambient `XAI_API_KEY` so a stray env key - /// can't redirect model fetching from the deployment's entitlement-gated - /// proxy to a raw `/v1/models` endpoint that lists the full model registry. pub(crate) fn resolve(endpoints: &config::EndpointsConfig, has_cached_session: bool) -> Self { if endpoints.has_custom_endpoint() { Self::CustomEndpoint @@ -58,7 +56,7 @@ impl ModelFetchAuth { #[derive(serde::Serialize, serde::Deserialize, PartialEq, Eq, Clone, Debug)] #[serde(rename_all = "snake_case")] -enum CacheAuthMethod { +pub(crate) enum CacheAuthMethod { Session, ApiKey, Deployment, @@ -95,57 +93,57 @@ pub(crate) fn task_model_error_for_catalog( } /// Thread-safe model manager. -/// -/// Owns the auth manager, config, and gateway needed to refresh models. -/// Uses `parking_lot::RwLock` for short clone-and-release access. #[derive(Clone)] pub struct ModelsManager { inner: Arc, } +/// Catalog fields written together under one lock, so readers never see a torn mix. +#[derive(Default)] +struct CatalogState { + prefetched: Option>, + models: IndexMap, + etag: Option, + /// Gates whether the apply path reselects the default (first real catalog) + has_fetched_real_catalog: bool, + /// `allowed_models` matched nothing; the prompt path blocks instead. + allowlist_excludes_all: bool, +} + struct Inner { - prefetched: RwLock>>, - models: RwLock>, + catalog: RwLock, current_model_id: RwLock, current_reasoning_effort: RwLock>, - etag: RwLock>, - /// Set once a real catalog has been fetched; gates whether - /// `apply_refresh_result` calls `reselect_default_model` (first - /// time) or `reselect_current_model_if_missing` (subsequent). - /// Reset in `clear()` for identity changes. - has_fetched_real_catalog: RwLock, // ── Owned context for self-contained refresh ──────────────── auth_manager: Arc, cfg: RwLock, fetch_auth: RwLock, gateway: RwLock>, cache: ModelsCacheManager, + endpoint: Arc, /// Guard to prevent overlapping retry loops. retry_in_flight: AtomicBool, - /// `allowed_models` matched nothing in the fetched catalog; the prompt path - /// blocks rather than run on the bundled default. Set in `apply_refresh_result`. - allowlist_excludes_all: AtomicBool, - /// Layer-3 LazinessDetector model-switch signal. Carries a - /// monotonically-increasing generation counter (`u64`) that is - /// bumped whenever the current model id actually changes via - /// [`Self::set_current_model_id`]. - /// - /// Two consumer patterns: - /// 1. `subscribe_model_switch().changed().await` — used by the - /// `SessionActor` main loop to react to a switch (e.g. zero - /// the per-session nudge counter). Critically, `watch::Receiver` - /// only resolves `.changed()` on changes that happen **after** - /// subscription — there is no stored-permit hazard akin to - /// `tokio::sync::Notify::notify_one()`. - /// 2. `model_switch_generation()` — cheap snapshot read used by - /// `maybe_fire_laziness_check`'s polling loop to detect a - /// switch that occurred during the idle wait or sampler call. - /// - /// `watch::Sender` natively fans out to every subscriber, so this - /// replaces the previous `RwLock>>` listener - /// registry — no manual fan-out, no listener-leak risk, no - /// `unregister` API to maintain. + /// Single-flight for the etag-triggered background refresh (`spawn_fetch`). + refresh_in_flight: AtomicBool, + /// Model-switch signal: a generation counter bumped when the current model id changes. model_switch_watch: tokio::sync::watch::Sender, + /// Set once the user explicitly picks a model (`/model`); guards the + /// first-catalog reselect from clobbering that choice. + user_selected_model: AtomicBool, +} + +/// Clears an in-flight flag on drop so a panicking task can't wedge future refreshes. +struct RetryInFlightGuard(Arc); +impl Drop for RetryInFlightGuard { + fn drop(&mut self) { + self.0.retry_in_flight.store(false, Ordering::Release); + } +} +struct RefreshInFlightGuard(Arc); +impl Drop for RefreshInFlightGuard { + fn drop(&mut self) { + self.0.refresh_in_flight.store(false, Ordering::Release); + } } impl Default for ModelsManager { @@ -162,6 +160,76 @@ impl Default for ModelsManager { } } +/// Builder for [`ModelsManager`]; transport and disk cache default to production (tests override them). +pub(crate) struct ModelsManagerBuilder { + prefetched: Option>, + models: IndexMap, + current_model_id: acp::ModelId, + auth_manager: Arc, + cfg: config::Config, + endpoint: Arc, + cache: ModelsCacheManager, +} + +impl ModelsManagerBuilder { + pub(crate) fn new( + prefetched: Option>, + models: IndexMap, + current_model_id: acp::ModelId, + auth_manager: Arc, + cfg: config::Config, + ) -> Self { + Self { + prefetched, + models, + current_model_id, + auth_manager, + cfg, + endpoint: Arc::new(HttpModelsEndpoint), + cache: ModelsCacheManager::new(), + } + } + + #[cfg(test)] + pub(crate) fn endpoint(mut self, endpoint: Arc) -> Self { + self.endpoint = endpoint; + self + } + + #[cfg(test)] + pub(crate) fn cache(mut self, cache: ModelsCacheManager) -> Self { + self.cache = cache; + self + } + + pub(crate) fn build(self) -> ModelsManager { + let has_session = self.auth_manager.current_or_expired().is_some(); + let fetch_auth = ModelFetchAuth::resolve(&self.cfg.endpoints, has_session); + let current_reasoning_effort = self.cfg.models.default_reasoning_effort; + ModelsManager { + inner: Arc::new(Inner { + catalog: RwLock::new(CatalogState { + prefetched: self.prefetched, + models: self.models, + ..Default::default() + }), + current_model_id: RwLock::new(self.current_model_id), + current_reasoning_effort: RwLock::new(current_reasoning_effort), + auth_manager: self.auth_manager, + cfg: RwLock::new(self.cfg), + fetch_auth: RwLock::new(fetch_auth), + gateway: RwLock::new(None), + cache: self.cache, + endpoint: self.endpoint, + retry_in_flight: AtomicBool::new(false), + refresh_in_flight: AtomicBool::new(false), + model_switch_watch: tokio::sync::watch::channel(0u64).0, + user_selected_model: AtomicBool::new(false), + }), + } + } +} + impl ModelsManager { pub(crate) fn new( prefetched: Option>, @@ -170,51 +238,20 @@ impl ModelsManager { auth_manager: Arc, cfg: config::Config, ) -> Self { - let has_session = auth_manager.current_or_expired().is_some(); - let fetch_auth = ModelFetchAuth::resolve(&cfg.endpoints, has_session); - let current_reasoning_effort = cfg.models.default_reasoning_effort; - Self { - inner: Arc::new(Inner { - prefetched: RwLock::new(prefetched), - models: RwLock::new(models), - current_model_id: RwLock::new(current_model_id), - current_reasoning_effort: RwLock::new(current_reasoning_effort), - etag: RwLock::new(None), - has_fetched_real_catalog: RwLock::new(false), - auth_manager, - cfg: RwLock::new(cfg), - fetch_auth: RwLock::new(fetch_auth), - gateway: RwLock::new(None), - cache: ModelsCacheManager::new(), - retry_in_flight: AtomicBool::new(false), - allowlist_excludes_all: AtomicBool::new(false), - model_switch_watch: tokio::sync::watch::channel(0u64).0, - }), - } + ModelsManagerBuilder::new(prefetched, models, current_model_id, auth_manager, cfg).build() } /// Subscribe to model-switch events. Returns a `watch::Receiver` - /// carrying the monotonic generation counter. `.changed()` only - /// resolves on switches that occur **after** subscription, so - /// there is no stored-permit hazard (the bug that motivated - /// replacing the previous `Arc` design). pub fn subscribe_model_switch(&self) -> tokio::sync::watch::Receiver { self.inner.model_switch_watch.subscribe() } - /// Cheap snapshot of the current model-switch generation. Used by - /// `maybe_fire_laziness_check`'s polling loop to detect a switch - /// that occurred during the idle wait or sampler call without - /// having to allocate a fresh `Receiver` per fire. + /// Cheap snapshot of the current model-switch generation, for the laziness-check poll loop. pub fn model_switch_generation(&self) -> u64 { *self.inner.model_switch_watch.borrow() } /// Build from a resolved config. Falls back to bundled default if no models available. - /// - /// When `prefetched_models` is `None`, the disk cache is consulted so that - /// server-side models are available for default-model resolution even when - /// the caller didn't do an explicit prefetch. pub fn from_config( cfg: &config::Config, prefetched_models: Option>, @@ -237,8 +274,6 @@ impl ModelsManager { let has_prefetched = prefetched_models.is_some(); let catalog = resolve_model_catalog(cfg, prefetched_models.clone()); - // Validate only against a real catalog; a bundled-only first run defers - // to the async fetch (`apply_refresh_result`). if has_prefetched { validate_selectable(cfg, &catalog)?; } @@ -262,7 +297,7 @@ impl ModelsManager { cfg.clone(), ); if has_prefetched { - *mgr.inner.has_fetched_real_catalog.write() = true; + mgr.inner.catalog.write().has_fetched_real_catalog = true; } Ok(mgr) } @@ -272,19 +307,14 @@ impl ModelsManager { } /// Swap config, rebuild catalog, and reselect the model. - /// - /// Calls `reselect_default_model` when the preferred model changed - /// (and is `Some`); otherwise `reselect_current_model_if_missing`. pub fn apply_config(&self, new_config: config::Config) { - // Reject an invalid reload instead of mutating live state: bad globs or - // (once a real catalog exists) an allowlist that excludes everything. if let Err(e) = new_config.validate_model_filters() { tracing::error!(error = %e, "ignoring config reload: invalid model filters"); return; } - let prefetched = self.inner.prefetched.read().clone(); + let prefetched = self.inner.catalog.read().prefetched.clone(); let new_catalog = resolve_model_catalog(&new_config, prefetched); - let has_real_catalog = *self.inner.has_fetched_real_catalog.read(); + let has_real_catalog = self.inner.catalog.read().has_fetched_real_catalog; if has_real_catalog && let Err(e) = validate_selectable(&new_config, &new_catalog) { tracing::error!(error = %e, "ignoring config reload: allowed_models excludes all models"); return; @@ -302,22 +332,15 @@ impl ModelsManager { *self.inner.fetch_auth.write() = ModelFetchAuth::resolve(&new_config.endpoints, has_session); *self.inner.cfg.write() = new_config.clone(); - // Recompute the prompt-block flag so a corrective reload unblocks. - if has_real_catalog { - let excludes_all = allowlist_matches_nothing(&new_config, &new_catalog); - self.inner - .allowlist_excludes_all - .store(excludes_all, Ordering::Relaxed); + { + let mut cat = self.inner.catalog.write(); + if has_real_catalog { + cat.allowlist_excludes_all = allowlist_matches_nothing(&new_config, &new_catalog); + } + cat.models = new_catalog; } - *self.inner.models.write() = new_catalog; - // A preferred-model flip caused only by a campaign overlay appearing or - // disappearing must not yank an in-flight session whose current model is - // still usable — the campaign applies to /new sessions only. let preferred_changed = new_preferred != old_preferred && new_preferred.is_some(); - // Recognize an appearing OR withdrawing campaign from the - // `default_is_campaign_driven` flag on each config (no disk I/O); correct - // even when the user has no base default (where a value compare would miss). let mut campaign_defaults = std::collections::HashSet::new(); if new_config.models.default_is_campaign_driven && let Some(d) = &new_preferred @@ -330,7 +353,8 @@ impl ModelsManager { let campaign_only_flip = is_campaign_only_flip(&old_preferred, &new_preferred, &campaign_defaults); let current_still_ok = { - let models = self.inner.models.read(); + let cat = self.inner.catalog.read(); + let models = &cat.models; let cur = self.inner.current_model_id.read(); models .get(cur.0.as_ref()) @@ -342,19 +366,20 @@ impl ModelsManager { self.reselect_current_model_if_missing(&new_config); } - // Push the new catalog to connected clients (`x.ai/models/update`). - // Without this, a long-running agent (leader mode) correctly swaps - // its in-memory catalog on a config.toml `[model.*]`/`[models]` edit, - // but already-connected clients keep rendering the stale model list - // until they reconnect. No-op when no gateway is attached (tests, - // pre-init). + self.notify_models_updated(); + } + + /// [`Self::apply_config`] plus an unconditional default re-resolve, for remote-settings arrival while no session exists. + pub fn apply_config_reselecting_default(&self, new_config: config::Config) { + self.apply_config(new_config.clone()); + self.reselect_default_model(&new_config); self.notify_models_updated(); } // ── Accessors ─────────────────────────────────────────────────── pub fn models(&self) -> IndexMap { - self.inner.models.read().clone() + self.inner.catalog.read().models.clone() } pub fn endpoints(&self) -> config::EndpointsConfig { @@ -370,11 +395,10 @@ impl ModelsManager { } /// ACP-visible (non-hidden) projection of the catalog. - /// The catalog coming from `resolve_model_catalog` already has - /// allowed_models + disabled_models + hidden_models applied. pub fn available(&self) -> IndexMap { let snapshot = { - let models = self.inner.models.read(); + let cat = self.inner.catalog.read(); + let models = &cat.models; models.clone() }; @@ -388,8 +412,9 @@ impl ModelsManager { pub(crate) fn task_model_error(&self, requested: &str) -> Option { let is_session_auth = self.is_session_auth(); - let models = self.inner.models.read(); - task_model_error_for_catalog(requested, &models, is_session_auth) + let cat = self.inner.catalog.read(); + let models = &cat.models; + task_model_error_for_catalog(requested, models, is_session_auth) } pub fn current_model_id(&self) -> acp::ModelId { @@ -397,11 +422,13 @@ impl ModelsManager { } pub fn set_current_model_id(&self, id: acp::ModelId) { - // Only bump the model-switch generation on a real change. - // The pager's `/model` handler can call this with the - // already-active id during re-resolution; bumping the counter - // in that case would needlessly cancel a healthy in-flight - // classifier call and zero the per-session nudge counter. + self.inner + .user_selected_model + .store(true, Ordering::Relaxed); + self.set_current_model_id_internal(id); + } + + fn set_current_model_id_internal(&self, id: acp::ModelId) { let changed = { let mut cur = self.inner.current_model_id.write(); let changed = *cur != id; @@ -415,25 +442,21 @@ impl ModelsManager { } } - /// Look up the per-model Layer-3 LazinessDetector config for the - /// model identified by `model_id`. Returns the default (disabled) - /// config when the id isn't in the catalog — same fallback - /// semantics as the `auto_compact_threshold_percent` lookup. + /// Per-model Layer-3 LazinessDetector config for `model_id` (disabled default when absent). pub fn laziness_detector_for(&self, model_id: &str) -> config::LazinessDetectorPerModelConfig { self.inner - .models + .catalog .read() + .models .get(model_id) .map(|e| e.info().laziness_detector.clone()) .unwrap_or_default() } /// Test-only catalog poke: inserts a `ModelEntry` keyed by `id`, - /// allowing integration tests to enable Layer-3 features per - /// model without spinning up the full config-merge pipeline. #[cfg(test)] pub(crate) fn insert_test_entry(&self, id: impl Into, entry: ModelEntry) { - self.inner.models.write().insert(id.into(), entry); + self.inner.catalog.write().models.insert(id.into(), entry); } pub fn current_reasoning_effort(&self) -> Option { @@ -447,33 +470,29 @@ impl ModelsManager { /// Whether the given model supports reasoning effort according to the catalog. pub fn model_supports_reasoning_effort(&self, model_id: &str) -> bool { self.inner - .models + .catalog .read() + .models .get(model_id) .map(|e| e.info().supports_reasoning_effort) .unwrap_or(false) } - /// The catalog default reasoning effort for `model_id`, if the catalog - /// pins one. Used as the final fallback when neither the session handle - /// nor the global config sets an explicit effort, so surfaced config stays - /// consistent with the effort sampling actually uses. pub fn model_default_reasoning_effort(&self, model_id: &str) -> Option { self.inner - .models + .catalog .read() + .models .get(model_id) .and_then(|e| e.info().reasoning_effort) } /// The raw catalog `reasoning_efforts` list for `model_id` with no fallback, - /// empty when the catalog pins none (caller falls back to the built-in - /// session modes). Distinct from the pager's gate-first, fallback-applied - /// `ModelState::reasoning_effort_options`. pub fn model_reasoning_efforts(&self, model_id: &str) -> Vec { self.inner - .models + .catalog .read() + .models .get(model_id) .map(|e| e.info().reasoning_efforts.clone()) .unwrap_or_default() @@ -481,8 +500,9 @@ impl ModelsManager { pub fn model_supports_backend_search(&self, model_id: &str) -> bool { self.inner - .models + .catalog .read() + .models .get(model_id) .map(|e| e.info().supports_backend_search) .unwrap_or(false) @@ -493,8 +513,9 @@ impl ModelsManager { model_id: &str, ) -> Option { self.inner - .models + .catalog .read() + .models .get(model_id) .and_then(|e| e.info().compactions_remaining) } @@ -504,69 +525,56 @@ impl ModelsManager { model_id: &str, ) -> Option { self.inner - .models + .catalog .read() + .models .get(model_id) .and_then(|e| e.info().compaction_at_tokens) } /// Catalog opt-in to display the served-checkpoint fingerprint for this model. - /// - /// `model_id` may be a routing slug (`config.model`, e.g. `grok-4.5`) - /// OR a catalog key; the catalog map is keyed by the config key, which can - /// differ from the slug for custom/enterprise ids (e.g. key `enterprise-grok-build` - /// → slug `grok-4.5`). Resolve to the catalog key first so a slug - /// caller still finds the opted-in entry. pub fn model_show_model_fingerprint(&self, model_id: &str) -> bool { - let models = self.inner.models.read(); - resolve_catalog_key(&models, &acp::ModelId::new(model_id)) + let cat = self.inner.catalog.read(); + let models = &cat.models; + resolve_catalog_key(models, &acp::ModelId::new(model_id)) .and_then(|key| models.get(key.0.as_ref())) .map(|e| e.info().show_model_fingerprint) .unwrap_or(false) } /// Resolved next-prompt-suggestion model pin from the live config - /// (`env > [models] prompt_suggestion > remote settings`); tracks config - /// hot-reloads via [`Self::apply_config`]. Consumed catalog-guarded by - /// `handle_suggest_prompt`. pub fn prompt_suggest_model_pin(&self) -> crate::config::PromptSuggestModelPin { self.inner.cfg.read().prompt_suggest_model_pin.clone() } /// Whether `model_id` resolves in the current catalog — as a config key - /// or a routing slug (see [`resolve_catalog_key`]). Deliberately checks - /// the full catalog rather than the user-selectable projection: auxiliary - /// background calls need a *sampleable* model, and hidden or - /// non-selectable entries are still sampleable. pub fn model_in_catalog(&self, model_id: &str) -> bool { - let models = self.inner.models.read(); - resolve_catalog_key(&models, &acp::ModelId::new(model_id)).is_some() + let cat = self.inner.catalog.read(); + let models = &cat.models; + resolve_catalog_key(models, &acp::ModelId::new(model_id)).is_some() } #[cfg(test)] fn prefetched(&self) -> Option> { - self.inner.prefetched.read().clone() + self.inner.catalog.read().prefetched.clone() } #[cfg(test)] fn has_fetched_real_catalog(&self) -> bool { - *self.inner.has_fetched_real_catalog.read() + self.inner.catalog.read().has_fetched_real_catalog } // ── Mutations ─────────────────────────────────────────────────── fn rebuild(&self, cfg: &config::Config, prefetched: Option>) { - *self.inner.models.write() = resolve_model_catalog(cfg, prefetched); + self.inner.catalog.write().models = resolve_model_catalog(cfg, prefetched); } /// Refresh models when the etag changes. - /// - /// Writes etag optimistically before spawning the fetch to coalesce - /// concurrent callers seeing the same new etag. pub async fn refresh_if_new_etag(&self, etag: String) { let same_etag = { - let current = self.inner.etag.read(); - current.as_deref() == Some(etag.as_str()) + let cat = self.inner.catalog.read(); + cat.etag.as_deref() == Some(etag.as_str()) }; if same_etag { let fetch_auth = *self.inner.fetch_auth.read(); @@ -576,19 +584,11 @@ impl ModelsManager { .await; return; } - *self.inner.etag.write() = Some(etag.clone()); tracing::info!(etag = %etag, "models etag changed, refreshing"); - self.do_refresh(Some(etag), RefreshStrategy::Online); + self.spawn_fetch(Some(etag)); } /// Auth identity changed: invalidate disk cache and refresh the catalog. - /// - /// Safe on OIDC token recovery after idle: we never drop a successfully-fetched - /// catalog on transient failure. Only fall back to the bundled default when - /// we have never had a real catalog (`!has_fetched_real_catalog`), or via - /// the genuine no-auth path (`clear()`). - /// - /// Respects the auth snapshot / hot-swap discipline. pub async fn on_auth_changed(&self) { let config = self.inner.cfg.read().clone(); crate::agent::init::update_telemetry_config(&config, &self.inner.auth_manager); @@ -603,13 +603,14 @@ impl ModelsManager { return; } - // Never eagerly drop prefetched on auth recovery. Only fall back to - // bundled defaults when we have never had a real catalog. Resolved once - // so the fetch and the failure-vs-disabled classification below agree. let remote_fetch_enabled = crate::util::config::resolve_remote_fetch_enabled(); self.fetch_and_apply_inner(remote_fetch_enabled).await; - if !*self.inner.has_fetched_real_catalog.read() && self.inner.prefetched.read().is_none() { + let needs_bundled_fallback = { + let cat = self.inner.catalog.read(); + !cat.has_fetched_real_catalog && cat.prefetched.is_none() + }; + if needs_bundled_fallback { if remote_fetch_enabled { xai_grok_telemetry::unified_log::warn( "model catalog: falling back to bundled defaults only", @@ -620,16 +621,11 @@ impl ModelsManager { })), ); } else { - // Deliberate no-fetch state, not a failure: no warn-class log. tracing::debug!("model catalog: bundled defaults in use (remote_fetch disabled)"); } - self.rebuild(&config, None); // first-time only: no fetched catalog, use bundled defaults + self.rebuild(&config, None); self.reselect_current_model_if_missing(&config); - // Schedule background retries so we recover once the network is - // back (e.g. after sleep/resume when the first fetch races DNS). - // With remote_fetch disabled a retry can never succeed, so none is - // scheduled. if remote_fetch_enabled { self.spawn_catalog_retry(); } @@ -638,7 +634,6 @@ impl ModelsManager { self.notify_models_updated(); } - /// Notify clients about the current model catalog. fn notify_models_updated(&self) { let available = self.available(); let current = self.current_model_id(); @@ -663,37 +658,12 @@ impl ModelsManager { } } - /// Hot-reload the catalog from `~/.grok/models_cache.json` after an - /// external write (detected by the config file watcher). - /// - /// A long-running leader otherwise only refreshes its catalog from its - /// *own* fetch paths (startup prefetch, auth change, response-header etag). - /// When another grok process sharing `~/.grok` (a `--no-leader` run, a - /// newer client, grok-desktop) fetches a fresher `/v1/models` catalog and - /// persists it, this picks it up without a network round-trip. - /// - /// Guards, in order: - /// 1. `load_fresh` — rejects stale (TTL), version-mismatched, - /// auth-method-mismatched, or origin-mismatched cache files (another - /// process running with different credentials or pointed at a - /// different backend must not poison this catalog). - /// 2. Content dedup — the leader itself rewrites the cache file - /// (`persist` after fetch, `renew_ttl` on same-etag responses), and the - /// watcher has no self-write suppression. If the cached models match - /// the in-memory prefetched catalog this is a no-op (the etag is still - /// adopted so `refresh_if_new_etag` doesn't refetch needlessly). - /// - /// On a real change: swaps the prefetched catalog, rebuilds, re-resolves - /// the configured default when this is the first real catalog (otherwise - /// reselects the current model if it disappeared), and notifies clients. + /// Hot-reload the catalog from `~/.grok/models_cache.json` after an external write (config-watcher detected). pub fn reload_from_disk_cache(&self) { self.reload_from_cache_manager(&self.inner.cache); } /// Core of [`Self::reload_from_disk_cache`], parameterized over the cache - /// manager so tests can point it at a temp file (the production - /// `ModelsCacheManager` path is fixed to `grok_home()`, a process-wide - /// `OnceLock`). fn reload_from_cache_manager(&self, cache: &ModelsCacheManager) { let fetch_auth = *self.inner.fetch_auth.read(); let Some(cached) = cache.load_fresh(&fetch_auth.cache_auth_method(), &self.cache_origin()) @@ -702,20 +672,15 @@ impl ModelsManager { return; }; - // Self-write / no-change dedup by content. `ModelEntry` doesn't impl - // `PartialEq` (nested config types), so compare the serialized form — - // catalogs are small (tens of entries) and writes are debounced. let same_content = { - let prefetched = self.inner.prefetched.read(); - prefetched.as_ref().is_some_and(|current| { + let cat = self.inner.catalog.read(); + cat.prefetched.as_ref().is_some_and(|current| { serde_json::to_string(current).ok() == serde_json::to_string(&cached.models).ok() }) }; if same_content { - // Adopt the (possibly newer) etag without a rebuild so the next - // response-header comparison in `refresh_if_new_etag` is accurate. if cached.etag.is_some() { - *self.inner.etag.write() = cached.etag; + self.inner.catalog.write().etag = cached.etag; } tracing::debug!("models cache changed on disk but catalog is identical; skipping"); return; @@ -723,37 +688,7 @@ impl ModelsManager { let cfg = self.inner.cfg.read().clone(); let count = cached.models.len(); - // Capture whether this is the first real catalog (mirrors - // `apply_refresh_result`): if the leader bootstrapped on bundled - // defaults, the configured default must be re-resolved against the - // real catalog rather than left on a placeholder. - let first_real_catalog = { - let mut flag = self.inner.has_fetched_real_catalog.write(); - let was_first = !*flag; - *flag = true; - was_first - }; - *self.inner.prefetched.write() = Some(cached.models.clone()); - self.rebuild(&cfg, Some(cached.models)); - *self.inner.etag.write() = cached.etag; - if first_real_catalog { - self.reselect_default_model(&cfg); - } else { - self.reselect_current_model_if_missing(&cfg); - } - - // Recompute the prompt-block flag (mirrors `apply_refresh_result`) so - // a corrective external cache write unlatches a previously latched - // "allowlist excludes everything" state instead of keeping prompts - // blocked against a stale catalog. - let excludes_all = allowlist_matches_nothing(&cfg, &self.inner.models.read()); - self.inner - .allowlist_excludes_all - .store(excludes_all, Ordering::Relaxed); - if excludes_all { - tracing::error!("allowed_models excludes all fetched models; prompts will be blocked"); - } - + self.apply_catalog(&cfg, cached.models, cached.etag); tracing::info!(count, "model catalog hot-reloaded from disk cache"); xai_grok_telemetry::unified_log::info( "model catalog: reloaded from external disk-cache write", @@ -764,18 +699,17 @@ impl ModelsManager { } /// Retry model catalog fetch in the background with exponential backoff. - /// - /// Spawned when `on_auth_changed` falls back to bundled defaults. Uses the - /// crate-standard `execute_with_backoff` (5 attempts, 5s base, 60s cap) and - /// notifies clients on success so the UI recovers after sleep/resume without - /// requiring a manual restart. fn spawn_catalog_retry(&self) { - // Deliberate no-fetch state: a retry loop can never succeed, so don't - // start one (defensive re-check; the spawn site already gates). + self.spawn_catalog_retry_with_backoff(crate::tools::retry::BackoffConfig::new( + 5, 5_000, 60_000, + )); + } + + /// [`Self::spawn_catalog_retry`] with an injectable backoff (fast in tests). + fn spawn_catalog_retry_with_backoff(&self, backoff: crate::tools::retry::BackoffConfig) { if !crate::util::config::resolve_remote_fetch_enabled() { return; } - // Prevent overlapping retry loops. if self .inner .retry_in_flight @@ -788,21 +722,19 @@ impl ModelsManager { let mgr = self.clone(); tokio::task::spawn(async move { - let backoff = crate::tools::retry::BackoffConfig::new(5, 5_000, 60_000); - + let _retry_guard = RetryInFlightGuard(mgr.inner.clone()); let result = crate::tools::retry::execute_with_backoff( &backoff, || { let mgr = mgr.clone(); async move { - // Bail out early if another code path already loaded a real catalog. - if *mgr.inner.has_fetched_real_catalog.read() { + if mgr.inner.catalog.read().has_fetched_real_catalog { return Ok(()); } mgr.fetch_and_apply().await; - if *mgr.inner.has_fetched_real_catalog.read() { + if mgr.inner.catalog.read().has_fetched_real_catalog { Ok(()) } else { Err("model catalog fetch returned no models") @@ -841,22 +773,24 @@ impl ModelsManager { ); } } - - mgr.inner.retry_in_flight.store(false, Ordering::Release); }); } + /// One-shot background catalog refresh after readiness; no-op when a fresh disk cache already loaded a real catalog. + pub fn spawn_background_refresh(&self) { + if self.inner.catalog.read().has_fetched_real_catalog { + tracing::debug!( + "skipping startup background model refresh: fresh cache already loaded" + ); + return; + } + self.spawn_catalog_retry(); + } + /// Refresh the model catalog on every auth token refresh. - /// - /// Listens for [`AuthManager::refresh_notifier`] signals directly, - /// bypassing the FSEvents file watcher which can silently stop - /// delivering events on macOS after resume from sleep. On each - /// notification the catalog is re-fetched from the server; if the - /// fetch succeeds and the catalog changed, clients are notified - /// via `x.ai/models/update`. pub fn start_auth_refresh_watcher(&self, notify: Arc) { let mgr = self.clone(); - let had_catalog_at_start = *self.inner.has_fetched_real_catalog.read(); + let had_catalog_at_start = self.inner.catalog.read().has_fetched_real_catalog; xai_grok_telemetry::unified_log::info( "model catalog: auth refresh watcher started", None, @@ -868,15 +802,13 @@ impl ModelsManager { tokio::spawn(async move { loop { notify.notified().await; - // Deliberate no-fetch state: skip the refresh entirely so the - // failure-classifying logs below keep meaning "actually failed". if !crate::util::config::resolve_remote_fetch_enabled() { tracing::debug!( "model catalog: auth refresh watcher skipped (remote_fetch disabled)" ); continue; } - let had_catalog = *mgr.inner.has_fetched_real_catalog.read(); + let had_catalog = mgr.inner.catalog.read().has_fetched_real_catalog; let old_count = mgr.available().len(); xai_grok_telemetry::unified_log::info( "model catalog: auth refresh watcher triggered", @@ -887,7 +819,7 @@ impl ModelsManager { })), ); mgr.fetch_and_apply().await; - let has_catalog = *mgr.inner.has_fetched_real_catalog.read(); + let has_catalog = mgr.inner.catalog.read().has_fetched_real_catalog; let new_count = mgr.available().len(); if has_catalog { if !had_catalog || new_count != old_count { @@ -917,12 +849,11 @@ impl ModelsManager { /// Wipe in-memory state so a previous identity's catalog doesn't leak. fn clear(&self) { - *self.inner.prefetched.write() = None; - *self.inner.models.write() = IndexMap::new(); - *self.inner.etag.write() = None; - *self.inner.has_fetched_real_catalog.write() = false; + *self.inner.catalog.write() = CatalogState::default(); + // A new identity starts fresh: drop the prior user's pick so its + // first catalog reselects that identity's default. self.inner - .allowlist_excludes_all + .user_selected_model .store(false, Ordering::Relaxed); } @@ -963,7 +894,6 @@ impl ModelsManager { } /// Disk-cache origin key for this manager's current endpoints/auth shape - /// (see [`ModelsCache::origin`]). fn cache_origin(&self) -> String { let endpoints = self.inner.cfg.read().endpoints.clone(); let fetch_auth = *self.inner.fetch_auth.read(); @@ -980,28 +910,81 @@ impl ModelsManager { return false; }; let cfg = self.inner.cfg.read().clone(); - *self.inner.has_fetched_real_catalog.write() = true; - *self.inner.prefetched.write() = Some(cached.models.clone()); - self.rebuild(&cfg, Some(cached.models)); - *self.inner.etag.write() = cached.etag; + self.apply_catalog(&cfg, cached.models, cached.etag); true } + /// A catalog-fetch session refresh bounded by `STARTUP_AUTH_REFRESH_TIMEOUT`. + /// A hung IdP on a cold cache degrades to a session-less fetch (the + /// bundled/cache catalog stays and the next refresh retries) instead of + /// stalling boot, mirroring the readiness path's no-mint auth bound. + async fn bounded_startup_auth(auth_manager: &Arc) -> Option { + Self::bounded_auth_refresh(async { auth_manager.auth().await.ok() }).await + } + + /// Bounds an auth-refresh future to `STARTUP_AUTH_REFRESH_TIMEOUT`, yielding + /// `None` on timeout. Split out so the timeout contract is unit-testable + /// without a live IdP. + async fn bounded_auth_refresh(fut: F) -> Option + where + F: std::future::Future>, + { + match tokio::time::timeout(crate::http::STARTUP_AUTH_REFRESH_TIMEOUT, fut).await { + Ok(auth) => auth, + Err(_) => { + tracing::warn!( + timeout_secs = crate::http::STARTUP_AUTH_REFRESH_TIMEOUT.as_secs(), + "model catalog: auth refresh timed out; fetching without a fresh session" + ); + None + } + } + } + fn spawn_fetch(&self, new_etag: Option) { - // Degrade to Offline: keep serving the current (cache/static) catalog. - if !crate::util::config::resolve_remote_fetch_enabled() { + self.spawn_fetch_inner( + new_etag, + crate::util::config::resolve_remote_fetch_enabled(), + ); + } + + /// `remote_fetch_enabled` is a parameter so tests can drive the gate without touching on-disk config. + fn spawn_fetch_inner(&self, new_etag: Option, remote_fetch_enabled: bool) { + if !remote_fetch_enabled { tracing::info!("model catalog refresh skipped: remote_fetch disabled"); return; } + if self + .inner + .refresh_in_flight + .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed) + .is_err() + { + tracing::debug!("model catalog refresh already in flight, skipping"); + return; + } let cfg = self.inner.cfg.read().clone(); let endpoints = cfg.endpoints.clone(); let fetch_auth = *self.inner.fetch_auth.read(); let auth_manager = self.inner.auth_manager.clone(); + let endpoint = self.inner.endpoint.clone(); let mgr = self.clone(); tokio::task::spawn(async move { - let auth = auth_manager.auth().await.ok(); - let new_prefetched = fetch_models_async(endpoints, auth, fetch_auth).await; + let _refresh_guard = RefreshInFlightGuard(mgr.inner.clone()); + let auth = Self::bounded_startup_auth(&auth_manager).await; + let new_prefetched = match tokio::time::timeout( + crate::http::STARTUP_FETCH_TIMEOUT, + endpoint.fetch_models(endpoints, auth, fetch_auth), + ) + .await + { + Ok(models) => models, + Err(_) => { + tracing::warn!("etag-triggered model refresh timed out"); + None + } + }; if !mgr.apply_refresh_result(&cfg, new_prefetched, new_etag) { return; } @@ -1010,27 +993,6 @@ impl ModelsManager { }); } - /// Fetch models, rebuild state, and notify clients. - fn do_refresh(&self, new_etag: Option, strategy: RefreshStrategy) { - match strategy { - RefreshStrategy::Offline => { - if self.try_load_cache() { - tracing::info!("models manager refreshed from cache (offline)"); - } - } - RefreshStrategy::OnlineIfUncached => { - if self.try_load_cache() { - tracing::info!("models manager refreshed from cache (online_if_uncached)"); - return; - } - self.spawn_fetch(new_etag); - } - RefreshStrategy::Online => { - self.spawn_fetch(new_etag); - } - } - } - /// Resolve the model list: tries cache first, then fetches from the network. pub async fn list_models(&self, strategy: RefreshStrategy) { match strategy { @@ -1055,14 +1017,12 @@ impl ModelsManager { } /// `remote_fetch_enabled` is a parameter so tests can drive the gate - /// without touching on-disk config layers. async fn fetch_and_apply_inner(&self, remote_fetch_enabled: bool) { - // Degrade to Offline: keep serving the current (cache/static) catalog. if !remote_fetch_enabled { tracing::info!("model catalog refresh skipped: remote_fetch disabled"); return; } - let auth = self.inner.auth_manager.auth().await.ok(); + let auth = Self::bounded_startup_auth(&self.inner.auth_manager).await; let has_auth = auth.is_some(); let fetch_auth = *self.inner.fetch_auth.read(); let cfg = self.inner.cfg.read().clone(); @@ -1074,7 +1034,22 @@ impl ModelsManager { "fetch_auth": format!("{fetch_auth:?}"), })), ); - let new_prefetched = fetch_models_async(cfg.endpoints.clone(), auth, fetch_auth).await; + let endpoint = self.inner.endpoint.clone(); + let new_prefetched = match tokio::time::timeout( + crate::http::STARTUP_FETCH_TIMEOUT, + endpoint.fetch_models(cfg.endpoints.clone(), auth, fetch_auth), + ) + .await + { + Ok(res) => res, + Err(_elapsed) => { + tracing::warn!( + timeout_secs = crate::http::STARTUP_FETCH_TIMEOUT.as_secs(), + "model catalog fetch timed out" + ); + None + } + }; let success = self.apply_refresh_result(&cfg, new_prefetched, None); if success { xai_grok_telemetry::unified_log::info( @@ -1087,6 +1062,37 @@ impl ModelsManager { } } + /// Publish a resolved catalog under one atomic write, then reselect the model (default on first real catalog, else keep current if present). + fn apply_catalog( + &self, + cfg: &config::Config, + models: IndexMap, + new_etag: Option, + ) { + let (first_real_catalog, excludes_all) = { + let mut cat = self.inner.catalog.write(); + let first_real_catalog = !cat.has_fetched_real_catalog; + cat.has_fetched_real_catalog = true; + cat.prefetched = Some(models); + cat.models = resolve_model_catalog(cfg, cat.prefetched.clone()); + cat.etag = new_etag; + cat.allowlist_excludes_all = allowlist_matches_nothing(cfg, &cat.models); + (first_real_catalog, cat.allowlist_excludes_all) + }; + if excludes_all { + tracing::error!("allowed_models excludes all fetched models; prompts will be blocked"); + } + + // Respect an explicit pre-catalog `/model` pick: auto-select the + // default on the first catalog only when the user hasn't chosen. + // Either way a now-invalid selection is replaced. + if first_real_catalog && !self.inner.user_selected_model.load(Ordering::Relaxed) { + self.reselect_default_model(cfg); + } else { + self.reselect_current_model_if_missing(cfg); + } + } + fn apply_refresh_result( &self, config: &config::Config, @@ -1099,50 +1105,25 @@ impl ModelsManager { "model catalog refresh failed", None, Some(serde_json::json!({ - "had_real_catalog": *self.inner.has_fetched_real_catalog.read(), + "had_real_catalog": self.inner.catalog.read().has_fetched_real_catalog, })), ); return false; }; - - let first_real_catalog = { - let mut flag = self.inner.has_fetched_real_catalog.write(); - let was_first = !*flag; - *flag = true; - was_first - }; - *self.inner.prefetched.write() = Some(new_prefetched.clone()); - self.rebuild(config, Some(new_prefetched)); - *self.inner.etag.write() = new_etag; - - // Can't exit a running app; flag it so the prompt path blocks instead. - let excludes_all = allowlist_matches_nothing(config, &self.inner.models.read()); - self.inner - .allowlist_excludes_all - .store(excludes_all, Ordering::Relaxed); - if excludes_all { - tracing::error!("allowed_models excludes all fetched models; prompts will be blocked"); - } - - if first_real_catalog { - self.reselect_default_model(config); - } else { - self.reselect_current_model_if_missing(config); - } + self.apply_catalog(config, new_prefetched, new_etag); true } pub fn allowlist_excludes_all(&self) -> bool { - self.inner.allowlist_excludes_all.load(Ordering::Relaxed) + self.inner.catalog.read().allowlist_excludes_all } /// Re-pick the default if `current_model_id` is gone from the catalog *or* - /// is no longer `user_selectable` (e.g. a config reload narrowed - /// `allowed_models`), so UI and sampling don't disagree on the active model. fn reselect_current_model_if_missing(&self, config: &config::Config) { let current = self.inner.current_model_id.read().clone(); let needs_reselection = { - let models = self.inner.models.read(); + let cat = self.inner.catalog.read(); + let models = &cat.models; match models.get(current.0.as_ref()) { None => true, Some(entry) => !entry.info.user_selectable, @@ -1152,25 +1133,24 @@ impl ModelsManager { return; } let (key, _, source) = { - let models = self.inner.models.read(); - resolve_default_model(config, &models, self.is_session_auth()) + let cat = self.inner.catalog.read(); + let models = &cat.models; + resolve_default_model(config, models, self.is_session_auth()) }; let new_id = acp::ModelId::new(Arc::from(key)); tracing::info!( old = %current.0, new = %new_id.0, source = %source, "current model not in new catalog, reselecting default" ); - *self.inner.current_model_id.write() = new_id; + self.set_current_model_id_internal(new_id); } /// Re-resolve the default model against the current catalog. - /// - /// Called on first catalog fetch and when `apply_config` detects a - /// preferred-model change. fn reselect_default_model(&self, config: &config::Config) { let (key, _, source) = { - let models = self.inner.models.read(); - resolve_default_model(config, &models, self.is_session_auth()) + let cat = self.inner.catalog.read(); + let models = &cat.models; + resolve_default_model(config, models, self.is_session_auth()) }; let new_id = acp::ModelId::new(Arc::from(key)); let current = self.inner.current_model_id.read().clone(); @@ -1179,7 +1159,7 @@ impl ModelsManager { old = %current.0, new = %new_id.0, source = %source, "re-resolved default model after catalog populated" ); - *self.inner.current_model_id.write() = new_id; + self.set_current_model_id_internal(new_id); } } } @@ -1197,2432 +1177,19 @@ pub enum RefreshStrategy { OnlineIfUncached, } -// ── Disk cache ────────────────────────────────────────────────────────────── +mod cache; +mod endpoint; +mod fetch; +mod resolution; -const MODELS_CACHE_FILE: &str = "models_cache.json"; -const CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(300); - -#[derive(serde::Serialize, serde::Deserialize)] -struct ModelsCache { - fetched_at: DateTime, - #[serde(default, skip_serializing_if = "Option::is_none")] - grok_version: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - auth_method: Option, - /// Models-list URL this catalog was fetched from - /// ([`crate::remote::models_list_url`]). Compared on load so a cache - /// written against one backend is a miss for another: entries embed - /// absolute `base_url`s, so adopting a foreign-origin cache silently - /// re-points inference (the windows lifecycle e2e failed exactly this - /// way — test 1's mock-server catalog, cached in the shared profile, - /// sent test 2's prompts to a dead port). `None` (legacy files) never - /// matches. - #[serde(default, skip_serializing_if = "Option::is_none")] - origin: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - etag: Option, - models: IndexMap, -} - -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 - } -} - -struct CacheResult { - models: IndexMap, - etag: Option, -} - -struct ModelsCacheManager { - path: std::path::PathBuf, - ttl: std::time::Duration, -} - -impl ModelsCacheManager { - fn new() -> Self { - Self { - path: crate::util::grok_home::grok_home().join(MODELS_CACHE_FILE), - ttl: CACHE_TTL, - } - } - - /// Sync; used by `prefetch_models_blocking`. Will be removed once startup - /// prefetch is async. - fn load_fresh( - &self, - expected_auth: &CacheAuthMethod, - expected_origin: &str, - ) -> Option { - 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, - }) - } - - /// Sync; see `load_fresh` note. - fn persist( - &self, - models: &IndexMap, - 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); - } - - 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::(&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"); - } - - /// Sync; see `load_fresh` note. - 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"), - } - } - - /// Sync; see `load_fresh` note. - fn atomic_write(&self, cache: &ModelsCache) { - if let Some(parent) = self.path.parent() { - let _ = std::fs::create_dir_all(parent); - } - let tmp = self.path.with_extension("json.tmp"); - if let Ok(json) = serde_json::to_vec_pretty(cache) - && std::fs::write(&tmp, &json).is_ok() - { - let _ = std::fs::rename(&tmp, &self.path); - } - } - - async fn atomic_write_async(&self, cache: &ModelsCache) { - if let Some(parent) = self.path.parent() { - let _ = tokio::fs::create_dir_all(parent).await; - } - let tmp = self.path.with_extension("json.tmp"); - let Ok(json) = serde_json::to_vec_pretty(cache) else { - return; - }; - if tokio::fs::write(&tmp, &json).await.is_ok() { - let _ = tokio::fs::rename(&tmp, &self.path).await; - } - } -} - -// ── Fetch ─────────────────────────────────────────────────────────────────── - -/// Build the prefetched model map from a flat list of entries. -/// -/// Each entry is keyed by its `id` field (falling back to the `model` slug -/// when `id` is absent). This lets A/B experiments that share the same -/// routing slug (e.g. "Auto" and "Grok Build" both route to `grok-build`) -/// coexist in the catalog without collision. -fn build_prefetched_map( - models: Vec, - api_base_url_override: Option, -) -> IndexMap { - let mut map: IndexMap = 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> { - 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 -/// prefetch thread and the leader's startup phase so the settings gate lives -/// once. The remote_fetch knob is resolved a single time so the two fetch -/// decisions cannot disagree mid-startup. -pub(crate) fn prefetch_models_and_settings_blocking( - endpoints: &config::EndpointsConfig, - auth: Option<&GrokAuth>, - fetch_auth: ModelFetchAuth, -) -> ( - Option>, - Option, -) { - let remote_fetch_enabled = crate::util::config::resolve_remote_fetch_enabled(); - let models = prefetch_models_blocking_gated(endpoints, auth, fetch_auth, remote_fetch_enabled); - // Settings need a grok.com session; skip for BYOK. - 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(), - ) - } - _ => 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> { - let cache_auth = fetch_auth.cache_auth_method(); - // Same URL the fetch below will hit — the cache is only valid for it. - 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); - } - - // Every catalog fetch in the product funnels through here, so this single - // gate also covers callers that don't go through the prefetch-env check - // (leader, headless, stdio, server). Cache above is local and stays usable. - 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); - - // NOTE: inheriting context_window / agent_type / api_backend - // from hardcoded defaults is handled centrally in - // `resolve_model_list` (config.rs), not here. Don't re-add it. - - 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>, - pub settings: Option, -} - -/// Handle for a startup prefetch thread. -pub type EarlyPrefetchHandle = std::thread::JoinHandle; - -struct PrefetchEnv { - auth: Option, - endpoints: config::EndpointsConfig, - model_fetch_auth: ModelFetchAuth, -} - -fn resolve_prefetch_env_with_auth(auth: Option) -> Option { - let _timer = crate::instrumentation_timer!("startup.early_prefetch_launch"); - // Config-aware (not env-only) so the prefetch can't leak the bearer to api.x.ai. - let mut endpoints = config::EndpointsConfig::from_effective_config(); - - if endpoints.deployment_key.is_none() { - endpoints.deployment_key = crate::managed_config::resolve_deployment_key(); - } - - resolve_prefetch_env_from_parts( - auth, - endpoints, - crate::util::config::resolve_remote_fetch_enabled(), - ) -} - -/// Decision core of [`resolve_prefetch_env_with_auth`], split from the config -/// loading so the gate is unit-testable. -/// -/// `remote_fetch_enabled = false` wins over every credential shape AND over -/// `has_custom_endpoint()` (which otherwise forces the prefetch to run): the -/// explicit off switch must hold even when a stray login, `XAI_API_KEY`, or -/// `deployment_key` would re-arm the prefetch — and with it the `/v1/settings` -/// fetch and the deployment-config sync on the prefetch thread. -fn resolve_prefetch_env_from_parts( - auth: Option, - endpoints: config::EndpointsConfig, - remote_fetch_enabled: bool, -) -> Option { - 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, - }) -} - -fn resolve_prefetch_env(grok_com_config: Option) -> Option { - let grok_home = crate::util::grok_home::grok_home(); - let auth_manager = AuthManager::new(&grok_home, grok_com_config.unwrap_or_default()); - let auth = auth_manager.current(); - resolve_prefetch_env_with_auth(auth) -} - -/// Start model + settings prefetch on a background thread using pre-resolved auth. -/// -/// When the caller has already obtained valid credentials (e.g. via -/// `try_ensure_fresh_auth`), pass them here to avoid re-reading stale cached -/// credentials from disk. -pub fn start_early_prefetch_with_auth(auth: Option) -> Option { - let env = resolve_prefetch_env_with_auth(auth)?; - Some(spawn_prefetch_thread(env, true)) -} - -/// Start model + settings prefetch on a background thread. -/// -/// Convenience wrapper that reads cached auth from disk. Prefer -/// `start_early_prefetch_with_auth` when you have pre-resolved credentials. -/// Also runs a best-effort managed-config sync when the cache is stale. -pub fn start_early_prefetch(grok_com_config: Option) -> Option { - let env = resolve_prefetch_env(grok_com_config)?; - Some(spawn_prefetch_thread(env, true)) -} - -/// Prefetch models + remote settings only — **no** managed-config sync. -/// -/// Used before the managed-policy gate so a kill-switch can apply on cold start -/// without healing a tampered on-disk policy before the fail-closed gate runs. -pub fn start_early_prefetch_settings_only( - grok_com_config: Option, -) -> Option { - let env = resolve_prefetch_env(grok_com_config)?; - Some(spawn_prefetch_thread(env, false)) -} - -fn spawn_prefetch_thread(env: PrefetchEnv, sync_managed: bool) -> 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, - ); - if sync_managed - && (env.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() - && let Ok(rt) = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - { - crate::managed_config::clear_orphan(); - let _ = rt.block_on(crate::managed_config::sync()); - } - - EarlyPrefetchResult { models, settings } - }) -} - -/// Map a model id (catalog key or routing slug) to its catalog key. -/// -/// Sessions persist the routing slug (`[model.X].model`, e.g. `grok-4.5`); -/// the catalog and `/model` picker use config keys (e.g. `enterprise-grok-build`). -/// Last slug match wins so user overrides beat defaults (matches `MvpAgent::resolve_model_id`). -pub(crate) fn resolve_catalog_key( - models: &IndexMap, - id: &acp::ModelId, -) -> Option { - 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** -/// entries. A selectable exact-key match wins (as in [`resolve_catalog_key`]); -/// otherwise the last selectable entry whose routing slug matches `id`, so a -/// non-selectable exact-key entry never shadows a selectable slug match. -pub(crate) fn selectable_catalog_key_for_persisted( - models: &IndexMap, - available: &IndexMap, - id: &acp::ModelId, -) -> Option { - 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 -/// is an active campaign default, i.e. the change is attributable to a campaign -/// overlay appearing/disappearing rather than a user/CLI/env edit. -fn is_campaign_only_flip( - old_preferred: &Option, - new_preferred: &Option, - campaign_defaults: &std::collections::HashSet, -) -> 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 -/// back to the bundled default when the catalog is empty or the preferred -/// model isn't present. -pub(crate) fn resolve_default_model( - cfg: &config::Config, - catalog: &IndexMap, - is_session_auth: bool, -) -> (String, ModelEntry, config::ConfigSource) { - let visible: IndexMap = 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()); - } - // Pre-catalog/degenerate only: nothing selectable. Set the bundled - // default's flag from `allowed_models` so no reader treats it as allowed. - 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" - ); - } - // A campaign default missing from the catalog falls back to the - // pre-campaign default before the first-visible fallback. Gated - // on the missing pref actually being the campaign-driven config - // value — a CLI/env pref that misses the catalog is not a - // campaign problem and must not detour through campaign state. - 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, - is_session_auth: bool, -) -> IndexMap { - let visible: IndexMap = 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`. Patterns (globset syntax: `*`, `?`, `[...]`) are matched -/// against either the catalog key or the 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 for config to reject. - pub(crate) fn compile(patterns: Option<&Vec>) -> Result, Vec> { - 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` -/// (remove), `allowed_models` (mark `user_selectable`), `hidden_models` (mark -/// `hidden`). Special/internal models (web_search, subagents, …) resolve via -/// `find_model_by_id`/`models()` and ignore `user_selectable`, so they need no -/// exemption. Globs are validated at load (`Config::validate_model_filters`); -/// the arms here fail closed if one slips through. -pub fn resolve_model_catalog( - cfg: &config::Config, - prefetched: Option>, -) -> IndexMap { - let mut catalog: IndexMap = 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"); - } - } - - // None/empty allowlist = allow all. - 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; - } - } - } - - // Persisted default first; CLI override below wins when set. - // Only apply if the model supports reasoning effort. - 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); - } - - // Skip non-reasoning models so we don't send the field to providers that reject it. - // Also skip models whose effort menu does not include the override (e.g. `--effort none` - // must not stamp `none` onto grok-4.5, which only offers low/medium/high). - 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. -/// -/// Uses the server `reasoning_efforts` menu when present; otherwise the -/// built-in low/medium/high/xhigh set (same as the pager legacy menu — no -/// `none`/`minimal`). -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. -/// (An excluded *default* does not count — that is recoverable by reselection.) -pub(crate) fn allowlist_matches_nothing( - cfg: &config::Config, - catalog: &IndexMap, -) -> 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 that -/// excludes an explicitly configured default (`default`/`-m`). Run only against a -/// real catalog (cache/prefetch/fetched), not the bundled bootstrap set. -pub(crate) fn validate_selectable( - cfg: &config::Config, - catalog: &IndexMap, -) -> 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(()) -} - -/// Async wrapper around `prefetch_models_blocking`. -pub(crate) async fn fetch_models_async( - endpoints: config::EndpointsConfig, - auth: Option, - fetch_auth: ModelFetchAuth, -) -> Option> { - tokio::task::spawn_blocking(move || { - prefetch_models_blocking(&endpoints, auth.as_ref(), fetch_auth) - }) - .await - .unwrap_or(None) -} +pub(crate) use cache::*; +pub(crate) use endpoint::*; +pub(crate) use fetch::*; +pub use fetch::{ + EarlyPrefetchHandle, EarlyPrefetchResult, start_early_prefetch, + start_early_prefetch_settings_only, start_early_prefetch_with_auth, +}; +pub(crate) use resolution::*; #[cfg(test)] -mod tests { - use super::*; - - fn test_manager() -> ModelsManager { - let _ = tracing_subscriber::fmt() - .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) - .with_test_writer() - .try_init(); - // Use a temp dir so AuthManager finds no credentials — ensures - // refresh_async bails at the auth check without needing a tokio runtime. - let tmp = std::env::temp_dir().join("grok-test-models-manager"); - let auth_manager = Arc::new(AuthManager::new(&tmp, GrokComConfig::default())); - ModelsManager::new( - None, - IndexMap::new(), - acp::ModelId::new("default"), - auth_manager, - config::Config::default(), - ) - } - - fn config_from_toml(toml: &str) -> config::Config { - config::Config::new_from_toml_cfg(&toml::from_str(toml).unwrap()).unwrap() - } - - #[test] - fn model_show_model_fingerprint_reads_catalog_flag() { - let mgr = test_manager(); - - // Entry with the catalog flag set → accessor returns true. - let mut flagged = ModelEntry { - info: config::ModelInfo::fallback("fp-model"), - api_key: None, - env_key: None, - auth_provider: None, - api_base_url: None, - }; - flagged.info.show_model_fingerprint = true; - mgr.insert_test_entry("fp-model", flagged); - - // Entry without the flag → defaults false. - mgr.insert_test_entry( - "plain-model", - ModelEntry { - info: config::ModelInfo::fallback("plain-model"), - api_key: None, - env_key: None, - auth_provider: None, - api_base_url: None, - }, - ); - - // Catalog KEY differs from the routing SLUG (custom/enterprise id): the - // map is keyed "enterprise-key" but the model slug is "enterprise-slug". - let mut custom = ModelEntry { - info: config::ModelInfo::fallback("enterprise-slug"), - api_key: None, - env_key: None, - auth_provider: None, - api_base_url: None, - }; - custom.info.show_model_fingerprint = true; - mgr.insert_test_entry("enterprise-key", custom); - - assert!(mgr.model_show_model_fingerprint("fp-model")); - assert!(!mgr.model_show_model_fingerprint("plain-model")); - // Unknown model id → false (no catalog entry). - assert!(!mgr.model_show_model_fingerprint("missing-model")); - // Lookup by the routing SLUG must resolve to the differing catalog KEY — - // a direct `.get(slug)` would miss this entry and wrongly return false. - assert!( - mgr.model_show_model_fingerprint("enterprise-slug"), - "slug lookup must resolve to the catalog key and read the flag", - ); - // Lookup by the catalog KEY itself still works (exact-match path). - assert!(mgr.model_show_model_fingerprint("enterprise-key")); - } - - /// The active model must be selectable, not the first entry of the - /// un-allowlisted catalog. - #[test] - fn default_model_honors_allowlist_when_no_default_set() { - let cfg = config_from_toml( - r#" - [models] - allowed_models = ["keep-*"] - [model.zzz-first] - model = "zzz-first" - base_url = "https://api.x.ai/v1" - context_window = 256000 - [model.keep-one] - model = "keep-one" - base_url = "https://api.x.ai/v1" - context_window = 256000 - "#, - ); - let catalog = resolve_model_catalog(&cfg, None); - let (_key, entry, _src) = resolve_default_model(&cfg, &catalog, true); - assert!( - entry.info.user_selectable, - "picked non-selectable {}", - entry.model - ); - } - - #[test] - fn validate_selectable_rejects_bad_allowlists() { - // Excluded explicit default → error names the default. - let excluded = config_from_toml( - r#" - [models] - default = "grok-3" - allowed_models = ["grok-4*"] - [model.grok-3] - model = "grok-3" - base_url = "https://api.x.ai/v1" - context_window = 256000 - [model.grok-4] - model = "grok-4" - base_url = "https://api.x.ai/v1" - context_window = 256000 - "#, - ); - let catalog = resolve_model_catalog(&excluded, None); - assert!( - validate_selectable(&excluded, &catalog) - .unwrap_err() - .contains("grok-3") - ); - - // Matches nothing → error. - let zero = config_from_toml( - r#" - [models] - allowed_models = ["nomatch-*"] - [model.grok-4] - model = "grok-4" - base_url = "https://api.x.ai/v1" - context_window = 256000 - "#, - ); - let catalog = resolve_model_catalog(&zero, None); - assert!(validate_selectable(&zero, &catalog).is_err()); - } - - #[tokio::test] - async fn refresh_if_new_etag_skips_when_same() { - let mgr = test_manager(); - // Set initial etag - *mgr.inner.etag.write() = Some("\"abc123\"".to_string()); - - // Same etag — should be a no-op (etag stays the same) - mgr.refresh_if_new_etag("\"abc123\"".to_string()).await; - assert_eq!( - mgr.inner.etag.read().as_deref(), - Some("\"abc123\""), - "etag should remain unchanged when same" - ); - } - - #[tokio::test] - async fn set_current_model_id_change_fires_watch_to_all_subscribers() { - // Two subscribers (simulating two SessionActors sharing one - // ModelsManager catalog) both observe the change. Fast-path - // "same id" must NOT bump the generation. - let mgr = test_manager(); - let mut rx_a = mgr.subscribe_model_switch(); - let mut rx_b = mgr.subscribe_model_switch(); - let initial_a = *rx_a.borrow_and_update(); - let initial_b = *rx_b.borrow_and_update(); - assert_eq!(initial_a, initial_b); - - // Same id is the fast path — no bump. - mgr.set_current_model_id(acp::ModelId::new("default")); - // Force-yield so any spurious wakeup would have a chance to - // surface. `try_recv` on a watch channel: use a timeout-zero - // race; if `.changed()` resolves within 25ms we have a bug. - let same_id_ticked = - tokio::time::timeout(std::time::Duration::from_millis(25), rx_a.changed()) - .await - .is_ok(); - assert!( - !same_id_ticked, - "set_current_model_id(same id) must NOT bump the watch generation", - ); - - // Real switch: both subscribers see the change. - mgr.set_current_model_id(acp::ModelId::new("grok-4")); - tokio::time::timeout(std::time::Duration::from_millis(100), rx_a.changed()) - .await - .expect("rx_a saw the switch") - .expect("watch channel still open"); - tokio::time::timeout(std::time::Duration::from_millis(100), rx_b.changed()) - .await - .expect("rx_b saw the switch") - .expect("watch channel still open"); - assert_ne!(*rx_a.borrow(), initial_a); - assert_eq!(*rx_a.borrow(), *rx_b.borrow()); - assert!(mgr.model_switch_generation() > initial_a); - } - - #[tokio::test] - async fn model_switch_generation_snapshot_reflects_current_state() { - let mgr = test_manager(); - let start = mgr.model_switch_generation(); - mgr.set_current_model_id(acp::ModelId::new("grok-4")); - assert_eq!(mgr.model_switch_generation(), start + 1); - // Idempotent: same id → no bump. - mgr.set_current_model_id(acp::ModelId::new("grok-4")); - assert_eq!(mgr.model_switch_generation(), start + 1); - // Another real change: another bump. - mgr.set_current_model_id(acp::ModelId::new("grok-3")); - assert_eq!(mgr.model_switch_generation(), start + 2); - } - - #[test] - fn rebuild_updates_models_and_available() { - let mgr = test_manager(); - assert!(mgr.models().is_empty()); - assert!(mgr.available().is_empty()); - - let cfg = config::Config::default(); - let mut prefetched = IndexMap::new(); - prefetched.insert( - "test-model".to_string(), - ModelEntry { - info: config::ModelInfo::fallback("test-model"), - api_key: None, - env_key: None, - auth_provider: None, - api_base_url: None, - }, - ); - - mgr.rebuild(&cfg, Some(prefetched)); - - assert!( - !mgr.models().is_empty(), - "models should be populated after rebuild" - ); - } - - #[test] - fn current_reasoning_effort_round_trip() { - let mgr = test_manager(); - assert_eq!(mgr.current_reasoning_effort(), None); - - mgr.set_current_reasoning_effort(Some(ReasoningEffort::High)); - assert_eq!(mgr.current_reasoning_effort(), Some(ReasoningEffort::High)); - - mgr.set_current_reasoning_effort(None); - assert_eq!(mgr.current_reasoning_effort(), None); - } - - #[test] - fn current_reasoning_effort_seeded_from_config() { - let tmp = std::env::temp_dir().join("grok-test-models-manager-seed"); - let auth_manager = Arc::new(AuthManager::new(&tmp, GrokComConfig::default())); - let mut cfg = config::Config::default(); - cfg.models.default_reasoning_effort = Some(ReasoningEffort::Xhigh); - let mgr = ModelsManager::new( - None, - IndexMap::new(), - acp::ModelId::new("default"), - auth_manager, - cfg, - ); - assert_eq!(mgr.current_reasoning_effort(), Some(ReasoningEffort::Xhigh),); - } - - #[test] - fn default_reasoning_effort_only_stamps_supporting_model() { - use indexmap::IndexMap; - - // Model that supports reasoning effort — effort should be applied. - let mut cfg = config::Config::default(); - cfg.models.default = Some("reasoning-model".to_string()); - cfg.models.default_reasoning_effort = Some(ReasoningEffort::High); - - let mut prefetched = IndexMap::new(); - let mut reasoning_entry = ModelEntry { - info: config::ModelInfo::fallback("reasoning-model"), - api_key: None, - env_key: None, - auth_provider: None, - api_base_url: None, - }; - reasoning_entry.info.supports_reasoning_effort = true; - prefetched.insert("reasoning-model".to_string(), reasoning_entry); - - let catalog = resolve_model_catalog(&cfg, Some(prefetched)); - assert_eq!( - catalog["reasoning-model"].info.reasoning_effort, - Some(ReasoningEffort::High), - "reasoning-supporting default model should be stamped", - ); - - // Model that does NOT support reasoning effort — effort must NOT be applied. - let mut cfg = config::Config::default(); - cfg.models.default = Some("plain-model".to_string()); - cfg.models.default_reasoning_effort = Some(ReasoningEffort::High); - - let mut prefetched = IndexMap::new(); - let plain_entry = ModelEntry { - info: config::ModelInfo::fallback("plain-model"), - api_key: None, - env_key: None, - auth_provider: None, - api_base_url: None, - }; - prefetched.insert("plain-model".to_string(), plain_entry); - - let catalog = resolve_model_catalog(&cfg, Some(prefetched)); - assert_eq!( - catalog["plain-model"].info.reasoning_effort, None, - "non-reasoning default model must NOT be stamped with persisted effort", - ); - } - - #[test] - fn reasoning_effort_override_skips_models_that_do_not_offer_level() { - use indexmap::IndexMap; - use xai_grok_sampling_types::ReasoningEffortOption; - - let cfg = config::Config { - reasoning_effort_override: Some(ReasoningEffort::None), - ..Default::default() - }; - - let mut prefetched = IndexMap::new(); - // 4.5-style: supports effort, menu is high only (no none). - let mut no_none = ModelEntry { - info: config::ModelInfo::fallback("grok-4.5"), - api_key: None, - env_key: None, - auth_provider: None, - api_base_url: None, - }; - no_none.info.supports_reasoning_effort = true; - no_none.info.reasoning_efforts = vec![ReasoningEffortOption { - id: "high".into(), - value: ReasoningEffort::High, - label: "High".into(), - description: None, - default: true, - }]; - no_none.info.reasoning_effort = Some(ReasoningEffort::High); - prefetched.insert("grok-4.5".to_string(), no_none); - - // Model that explicitly offers none. - let mut with_none = ModelEntry { - info: config::ModelInfo::fallback("legacy-none"), - api_key: None, - env_key: None, - auth_provider: None, - api_base_url: None, - }; - with_none.info.supports_reasoning_effort = true; - with_none.info.reasoning_efforts = vec![ReasoningEffortOption { - id: "none".into(), - value: ReasoningEffort::None, - label: "None".into(), - description: None, - default: true, - }]; - prefetched.insert("legacy-none".to_string(), with_none); - - let catalog = resolve_model_catalog(&cfg, Some(prefetched)); - assert_eq!( - catalog["grok-4.5"].info.reasoning_effort, - Some(ReasoningEffort::High), - "--effort none must not stamp onto models that do not offer none" - ); - assert_eq!( - catalog["legacy-none"].info.reasoning_effort, - Some(ReasoningEffort::None), - "models that list none should still accept the override" - ); - } - - #[test] - fn config_menu_only_model_derives_support_and_default() { - // The config-TOML path: a model configured with ONLY `reasoning_efforts` - // (no `supports_reasoning_effort`, no scalar `reasoning_effort`) must read - // as supported with the marked-default option's value on the internal - // gates that BugBot flagged (support gate + wire default). - let mut cfg = config::Config::default(); - cfg.config_models.insert( - "menu-only".to_string(), - config::ConfigModelOverride { - reasoning_efforts: vec![ - ReasoningEffortOption { - id: "balanced".to_string(), - value: ReasoningEffort::Medium, - label: "Balanced".to_string(), - description: None, - default: false, - }, - ReasoningEffortOption { - id: "deep".to_string(), - value: ReasoningEffort::Xhigh, - label: "Deep".to_string(), - description: None, - default: true, - }, - ], - ..Default::default() - }, - ); - // A sibling with no menu must stay underived (empty-list path unchanged). - cfg.config_models - .insert("plain".to_string(), config::ConfigModelOverride::default()); - - let catalog = resolve_model_catalog(&cfg, None); - let info = &catalog["menu-only"].info; - assert!( - info.supports_reasoning_effort, - "menu-only model must derive support" - ); - assert_eq!( - info.reasoning_effort, - Some(ReasoningEffort::Xhigh), - "derived default = marked-default option value" - ); - assert!(!catalog["plain"].info.supports_reasoning_effort); - assert_eq!(catalog["plain"].info.reasoning_effort, None); - - // The internal getters read those derived fields. - let tmp = std::env::temp_dir().join("grok-test-models-manager-menu-only"); - let auth_manager = Arc::new(AuthManager::new(&tmp, GrokComConfig::default())); - let mgr = ModelsManager::new( - None, - catalog, - acp::ModelId::new("menu-only"), - auth_manager, - cfg, - ); - assert!(mgr.model_supports_reasoning_effort("menu-only")); - assert_eq!( - mgr.model_default_reasoning_effort("menu-only"), - Some(ReasoningEffort::Xhigh) - ); - assert_eq!(mgr.model_reasoning_efforts("menu-only").len(), 2); - assert!(!mgr.model_supports_reasoning_effort("plain")); - assert_eq!(mgr.model_default_reasoning_effort("plain"), None); - } - - #[test] - fn cli_reasoning_effort_override_only_stamps_supporting_models() { - use indexmap::IndexMap; - - let cfg = config::Config { - reasoning_effort_override: Some(ReasoningEffort::High), - ..config::Config::default() - }; - - let mut prefetched = IndexMap::new(); - let mut reasoning_entry = ModelEntry { - info: config::ModelInfo::fallback("reasoning-model"), - api_key: None, - env_key: None, - auth_provider: None, - api_base_url: None, - }; - reasoning_entry.info.supports_reasoning_effort = true; - prefetched.insert("reasoning-model".to_string(), reasoning_entry); - - let plain_entry = ModelEntry { - info: config::ModelInfo::fallback("plain-model"), - api_key: None, - env_key: None, - auth_provider: None, - api_base_url: None, - }; - prefetched.insert("plain-model".to_string(), plain_entry); - - let catalog = resolve_model_catalog(&cfg, Some(prefetched)); - assert_eq!( - catalog["reasoning-model"].info.reasoning_effort, - Some(ReasoningEffort::High), - "reasoning-supporting model should be stamped", - ); - assert_eq!( - catalog["plain-model"].info.reasoning_effort, None, - "non-reasoning model must NOT be stamped", - ); - } - - #[test] - fn apply_refresh_result_only_updates_etag_on_success() { - let mgr = test_manager(); - let cfg = config::Config::default(); - *mgr.inner.etag.write() = Some("\"old\"".to_string()); - - assert!( - !mgr.apply_refresh_result(&cfg, None, Some("\"new\"".to_string())), - "failed refresh should report no update" - ); - assert_eq!( - mgr.inner.etag.read().as_deref(), - Some("\"old\""), - "etag should remain unchanged when refresh fails" - ); - assert!( - mgr.prefetched().is_none(), - "prefetched models should stay unchanged" - ); - } - - fn make_model_entry(model_id: &str) -> ModelEntry { - ModelEntry { - info: config::ModelInfo::fallback(model_id), - api_key: None, - env_key: None, - auth_provider: None, - api_base_url: None, - } - } - - fn make_prefetched(ids: &[&str]) -> IndexMap { - ids.iter() - .map(|id| (id.to_string(), make_model_entry(id))) - .collect() - } - - // ── auth-change refresh: has_fetched_real_catalog flag ───────────── - - #[test] - fn first_apply_refresh_reselects_default_model() { - let mgr = test_manager(); - let mut cfg = config::Config::default(); - cfg.models.default = Some("grok-3".to_string()); - - assert!(!mgr.has_fetched_real_catalog()); - - let prefetched = make_prefetched(&["grok-3", "grok-4"]); - mgr.apply_refresh_result(&cfg, Some(prefetched), None); - - assert!(mgr.has_fetched_real_catalog()); - assert_eq!(mgr.current_model_id().0.as_ref(), "grok-3"); - } - - #[test] - fn subsequent_apply_refresh_preserves_user_model() { - let mgr = test_manager(); - let mut cfg = config::Config::default(); - cfg.models.default = Some("grok-3".to_string()); - - let prefetched = make_prefetched(&["grok-3", "grok-4"]); - mgr.apply_refresh_result(&cfg, Some(prefetched), None); - mgr.set_current_model_id(acp::ModelId::new("grok-4")); - - // Simulate on_auth_changed clearing prefetched + etag. - *mgr.inner.prefetched.write() = None; - *mgr.inner.etag.write() = None; - - let prefetched = make_prefetched(&["grok-3", "grok-4"]); - mgr.apply_refresh_result(&cfg, Some(prefetched), None); - - assert_eq!( - mgr.current_model_id().0.as_ref(), - "grok-4", - "user's model selection must survive auth-change refresh" - ); - } - - #[test] - fn subsequent_refresh_reselects_when_model_removed() { - let mgr = test_manager(); - let mut cfg = config::Config::default(); - cfg.models.default = Some("grok-3".to_string()); - - let prefetched = make_prefetched(&["grok-3", "grok-4"]); - mgr.apply_refresh_result(&cfg, Some(prefetched), None); - mgr.set_current_model_id(acp::ModelId::new("grok-4")); - - // Second refresh with grok-4 removed. - let prefetched = make_prefetched(&["grok-3", "grok-4.5"]); - mgr.apply_refresh_result(&cfg, Some(prefetched), None); - - assert_eq!( - mgr.current_model_id().0.as_ref(), - "grok-3", - "should fall back to config default when current is removed" - ); - } - - #[test] - fn failed_refresh_does_not_set_has_fetched_real_catalog() { - let mgr = test_manager(); - let cfg = config::Config::default(); - - mgr.apply_refresh_result(&cfg, None, None); - - assert!( - !mgr.has_fetched_real_catalog(), - "failed refresh must not flip has_fetched_real_catalog" - ); - } - - // ── apply_config: honor changed preferred model from config ──────── - - #[test] - fn apply_config_honors_new_preferred_model() { - let mgr = test_manager(); - let mut cfg = config::Config::default(); - cfg.models.default = Some("grok-3".to_string()); - - let prefetched = make_prefetched(&["grok-3", "grok-4"]); - mgr.apply_refresh_result(&cfg, Some(prefetched), None); - mgr.set_current_model_id(acp::ModelId::new("grok-4")); - - // Simulate stale inner cfg (no default) from a racing auth refresh. - let mut stale_cfg = config::Config::default(); - stale_cfg.models.default = None; - *mgr.inner.cfg.write() = stale_cfg; - - let mut new_cfg = config::Config::default(); - new_cfg.models.default = Some("grok-3".to_string()); - mgr.apply_config(new_cfg); - - assert_eq!( - mgr.current_model_id().0.as_ref(), - "grok-3", - "apply_config must honor updated preferred model from config" - ); - } - - #[test] - fn apply_config_preserves_current_when_preferred_unchanged() { - let mgr = test_manager(); - let cfg = config::Config::default(); - - let prefetched = make_prefetched(&["grok-3", "grok-4"]); - mgr.apply_refresh_result(&cfg, Some(prefetched), None); - - mgr.set_current_model_id(acp::ModelId::new("grok-4")); - - // Unrelated config change — preferred model unchanged. - let new_cfg = config::Config::default(); - mgr.apply_config(new_cfg); - - assert_eq!( - mgr.current_model_id().0.as_ref(), - "grok-4", - "apply_config must not reset model when preferred hasn't changed" - ); - } - - #[test] - fn apply_config_falls_back_when_preferred_not_in_catalog() { - let mgr = test_manager(); - let mut cfg = config::Config::default(); - cfg.models.default = Some("grok-3".to_string()); - - let prefetched = make_prefetched(&["grok-3", "grok-4"]); - mgr.apply_refresh_result(&cfg, Some(prefetched), None); - - mgr.set_current_model_id(acp::ModelId::new("grok-4")); - - // Preferred model not in catalog — falls back to first entry. - let mut new_cfg = config::Config::default(); - new_cfg.models.default = Some("grok-nonexistent".to_string()); - mgr.apply_config(new_cfg); - - let current = mgr.current_model_id(); - let first_available = mgr.available().keys().next().unwrap().clone(); - assert_eq!( - current.0.as_ref(), - first_available.0.as_ref(), - "should fall back to first visible model when preferred not in catalog" - ); - } - - #[test] - fn apply_config_both_none_preferred_preserves_current() { - let mgr = test_manager(); - let cfg = config::Config::default(); - let prefetched = make_prefetched(&["grok-3", "grok-4"]); - mgr.apply_refresh_result(&cfg, Some(prefetched), None); - mgr.set_current_model_id(acp::ModelId::new("grok-4")); - let new_cfg = config::Config::default(); - mgr.apply_config(new_cfg); - - assert_eq!( - mgr.current_model_id().0.as_ref(), - "grok-4", - "both-None preferred must preserve user's runtime model" - ); - } - - #[test] - fn apply_config_old_some_new_none_preserves_current() { - let mgr = test_manager(); - let mut cfg = config::Config::default(); - cfg.models.default = Some("grok-3".to_string()); - - let prefetched = make_prefetched(&["grok-3", "grok-4"]); - mgr.apply_refresh_result(&cfg, Some(prefetched), None); - assert_eq!(mgr.current_model_id().0.as_ref(), "grok-3"); - - mgr.set_current_model_id(acp::ModelId::new("grok-4")); - - // [models] default removed — is_some() guard prevents reset. - let new_cfg = config::Config::default(); - mgr.apply_config(new_cfg); - - assert_eq!( - mgr.current_model_id().0.as_ref(), - "grok-4", - "old=Some new=None must not reset model (is_some guard)" - ); - } - - // ── end-to-end: auth refresh + config reload compose correctly ─── - - #[test] - fn auth_refresh_then_config_reload_preserves_user_model() { - let mgr = test_manager(); - let mut cfg = config::Config::default(); - cfg.models.default = Some("grok-3".to_string()); - - // Initial fetch. - let prefetched = make_prefetched(&["grok-3", "grok-4"]); - mgr.apply_refresh_result(&cfg, Some(prefetched), None); - - // User runs /model grok-4. - mgr.set_current_model_id(acp::ModelId::new("grok-4")); - - // Auth refresh races — clears prefetched/etag. - *mgr.inner.prefetched.write() = None; - *mgr.inner.etag.write() = None; - - // Second fetch must preserve user's model. - let prefetched = make_prefetched(&["grok-3", "grok-4"]); - mgr.apply_refresh_result(&cfg, Some(prefetched), None); - assert_eq!(mgr.current_model_id().0.as_ref(), "grok-4"); - - // Config reload with persisted preference. - let mut new_cfg = config::Config::default(); - new_cfg.models.default = Some("grok-4".to_string()); - mgr.apply_config(new_cfg); - assert_eq!(mgr.current_model_id().0.as_ref(), "grok-4"); - } - - // ── disk-cache hot-reload (external models_cache.json writes) ──── - - fn test_cache_manager(dir: &std::path::Path) -> ModelsCacheManager { - ModelsCacheManager { - path: dir.join(MODELS_CACHE_FILE), - ttl: CACHE_TTL, - } - } - - /// An external process persisting a fresh catalog must be picked up: - /// catalog swapped, etag adopted, real-catalog flag set. - #[test] - fn reload_from_disk_cache_applies_external_catalog() { - let mgr = test_manager(); - let tmp = tempfile::TempDir::new().unwrap(); - let cache = test_cache_manager(tmp.path()); - - let auth_method = mgr.inner.fetch_auth.read().cache_auth_method(); - cache.persist( - &make_prefetched(&["grok-4.5", "grok-4.3"]), - Some("etag-ext"), - auth_method, - &mgr.cache_origin(), - ); - - mgr.reload_from_cache_manager(&cache); - - assert!(mgr.has_fetched_real_catalog()); - assert!(mgr.models().contains_key("grok-4.5")); - assert!(mgr.models().contains_key("grok-4.3")); - assert_eq!(mgr.inner.etag.read().as_deref(), Some("etag-ext")); - } - - /// A latched "allowlist excludes everything" prompt block must clear when - /// an external cache write delivers a catalog the allowlist matches — - /// `reload_from_cache_manager` recomputes `allowlist_excludes_all` after - /// the rebuild, like `apply_refresh_result` does. - #[test] - fn reload_from_disk_cache_recomputes_allowlist_excludes_all() { - let mgr = test_manager(); - let cfg = config_from_toml("[models]\nallowed_models = [\"keep-*\"]"); - - // Latch the flag: neither the fetched model nor the bundled defaults - // merged by `resolve_model_catalog` match `keep-*`. - mgr.apply_refresh_result(&cfg, Some(make_prefetched(&["other-1"])), None); - assert!( - mgr.allowlist_excludes_all(), - "setup: allowlist should exclude the entire catalog" - ); - // `apply_refresh_result` borrows the config without storing it, while - // `reload_from_cache_manager` reads `inner.cfg` — install it there. - *mgr.inner.cfg.write() = cfg.clone(); - - // External process persists a catalog containing an allowed model. - let tmp = tempfile::TempDir::new().unwrap(); - let cache = test_cache_manager(tmp.path()); - let auth_method = mgr.inner.fetch_auth.read().cache_auth_method(); - cache.persist( - &make_prefetched(&["keep-1"]), - Some("etag-keep"), - auth_method, - &mgr.cache_origin(), - ); - - mgr.reload_from_cache_manager(&cache); - - assert!(mgr.models().contains_key("keep-1")); - assert!( - !mgr.allowlist_excludes_all(), - "corrective external cache write must unlatch the prompt block" - ); - } - - /// When the *first* real catalog arrives via an external cache write (the - /// leader never completed its own fetch), the configured `[models]` - /// default must be resolved — mirroring `apply_refresh_result`'s - /// first-catalog branch — instead of staying on the bundled placeholder. - #[test] - fn reload_from_disk_cache_resolves_default_on_first_catalog() { - let mgr = test_manager(); - assert!(!mgr.has_fetched_real_catalog()); - let cfg = config_from_toml("[models]\ndefault = \"keep-1\""); - // `reload_from_cache_manager` reads the manager's stored config. - *mgr.inner.cfg.write() = cfg.clone(); - - let tmp = tempfile::TempDir::new().unwrap(); - let cache = test_cache_manager(tmp.path()); - let auth_method = mgr.inner.fetch_auth.read().cache_auth_method(); - cache.persist( - &make_prefetched(&["keep-1", "other-1"]), - Some("etag-first"), - auth_method, - &mgr.cache_origin(), - ); - - mgr.reload_from_cache_manager(&cache); - - assert!(mgr.has_fetched_real_catalog()); - assert_eq!( - mgr.current_model_id().0.as_ref(), - "keep-1", - "first real catalog must resolve the configured default" - ); - } - - /// A cache write whose catalog matches the in-memory prefetched map (the - /// leader's own `persist`/`renew_ttl` self-writes, or a same-content fetch - /// by another process) must be a no-op apart from adopting the etag — no - /// rebuild, no model reselection. - #[test] - fn reload_from_disk_cache_skips_identical_catalog_and_adopts_etag() { - let mgr = test_manager(); - let cfg = config::Config::default(); - let prefetched = make_prefetched(&["grok-3", "grok-4"]); - mgr.apply_refresh_result(&cfg, Some(prefetched.clone()), Some("etag-a".into())); - mgr.set_current_model_id(acp::ModelId::new("grok-4")); - - let tmp = tempfile::TempDir::new().unwrap(); - let cache = test_cache_manager(tmp.path()); - let auth_method = mgr.inner.fetch_auth.read().cache_auth_method(); - cache.persist( - &prefetched, - Some("etag-b"), - auth_method, - &mgr.cache_origin(), - ); - - mgr.reload_from_cache_manager(&cache); - - assert_eq!( - mgr.current_model_id().0.as_ref(), - "grok-4", - "identical catalog must not disturb the user's model" - ); - assert_eq!( - mgr.inner.etag.read().as_deref(), - Some("etag-b"), - "etag should be adopted so refresh_if_new_etag stays accurate" - ); - } - - /// A cache file older than the TTL is rejected by `load_fresh` — the - /// watcher event arrives within the debounce window of the write, so a - /// stale file means the write was not a fresh fetch. - #[test] - fn reload_from_disk_cache_ignores_stale_cache() { - let mgr = test_manager(); - let tmp = tempfile::TempDir::new().unwrap(); - let cache = test_cache_manager(tmp.path()); - let auth_method = mgr.inner.fetch_auth.read().cache_auth_method(); - let stale = ModelsCache { - fetched_at: Utc::now() - ChronoDuration::seconds(3600), - grok_version: Some(xai_grok_version::VERSION.to_string()), - auth_method: Some(auth_method), - origin: Some(mgr.cache_origin()), - etag: Some("etag-stale".into()), - models: make_prefetched(&["grok-stale"]), - }; - cache.atomic_write(&stale); - - mgr.reload_from_cache_manager(&cache); - - assert!(!mgr.models().contains_key("grok-stale")); - assert!(mgr.inner.etag.read().is_none()); - } - - /// A cache persisted by a process running with different credentials - /// (e.g. an API-key `--no-leader` run next to a session-auth leader) - /// must not poison this manager's catalog. - #[test] - fn reload_from_disk_cache_ignores_auth_method_mismatch() { - let mgr = test_manager(); - let tmp = tempfile::TempDir::new().unwrap(); - let cache = test_cache_manager(tmp.path()); - let current = mgr.inner.fetch_auth.read().cache_auth_method(); - let other = if current == CacheAuthMethod::Session { - CacheAuthMethod::ApiKey - } else { - CacheAuthMethod::Session - }; - cache.persist( - &make_prefetched(&["grok-other-auth"]), - Some("etag-x"), - other, - &mgr.cache_origin(), - ); - - mgr.reload_from_cache_manager(&cache); - - assert!(!mgr.models().contains_key("grok-other-auth")); - } - - /// A cache persisted by a process pointed at a *different backend* (env - /// override, another deployment, a test's mock server) must not poison - /// this manager's catalog: cached entries embed absolute `base_url`s from - /// their origin, so adopting them silently re-points inference. This is - /// the windows-x86_64 lifecycle e2e failure mode — the shared-profile - /// cache from test 1's mock sent test 2's prompts to a dead port. - #[test] - fn reload_from_disk_cache_ignores_origin_mismatch() { - let mgr = test_manager(); - let tmp = tempfile::TempDir::new().unwrap(); - let cache = test_cache_manager(tmp.path()); - let auth_method = mgr.inner.fetch_auth.read().cache_auth_method(); - cache.persist( - &make_prefetched(&["grok-other-origin"]), - Some("etag-y"), - auth_method, - "http://127.0.0.1:49953/v1/models", - ); - - mgr.reload_from_cache_manager(&cache); - - assert!(!mgr.models().contains_key("grok-other-origin")); - assert!(mgr.inner.etag.read().is_none()); - } - - /// A legacy cache file written before the `origin` field existed must be - /// treated as a miss (`None` origin never matches) — its entries could - /// have come from anywhere. - #[test] - fn reload_from_disk_cache_ignores_legacy_cache_without_origin() { - let mgr = test_manager(); - let tmp = tempfile::TempDir::new().unwrap(); - let cache = test_cache_manager(tmp.path()); - let auth_method = mgr.inner.fetch_auth.read().cache_auth_method(); - let legacy = ModelsCache { - fetched_at: Utc::now(), - grok_version: Some(xai_grok_version::VERSION.to_string()), - auth_method: Some(auth_method), - origin: None, - etag: Some("etag-legacy".into()), - models: make_prefetched(&["grok-legacy"]), - }; - cache.atomic_write(&legacy); - - mgr.reload_from_cache_manager(&cache); - - assert!(!mgr.models().contains_key("grok-legacy")); - } - - // ── clear() resets has_fetched_real_catalog ────────────────────── - - #[test] - fn clear_resets_has_fetched_real_catalog() { - let mgr = test_manager(); - let mut cfg = config::Config::default(); - cfg.models.default = Some("grok-3".to_string()); - - let prefetched = make_prefetched(&["grok-3", "grok-4"]); - mgr.apply_refresh_result(&cfg, Some(prefetched), None); - assert!(mgr.has_fetched_real_catalog()); - - mgr.clear(); - assert!(!mgr.has_fetched_real_catalog()); - - // New identity fetch — resolves default via reselect_default_model. - let prefetched = make_prefetched(&["grok-4.5", "grok-4.3"]); - mgr.apply_refresh_result(&cfg, Some(prefetched), None); - let first_available = mgr.available().keys().next().unwrap().clone(); - assert_eq!( - mgr.current_model_id().0.as_ref(), - first_available.0.as_ref() - ); - } - - /// A flip is "campaign-only" iff the preferred changed and either side is an - /// active campaign default. - #[test] - fn is_campaign_only_flip_detects_campaign_driven_changes() { - let camp: std::collections::HashSet = ["beta".into()].into_iter().collect(); - // New side is the campaign default (campaign appearing) → campaign-only. - assert!(is_campaign_only_flip( - &Some("alpha".into()), - &Some("beta".into()), - &camp - )); - // Old side was the campaign default (campaign withdrawing) → campaign-only. - assert!(is_campaign_only_flip( - &Some("beta".into()), - &Some("alpha".into()), - &camp - )); - // Neither side a campaign default → ordinary user/CLI/env flip. - assert!(!is_campaign_only_flip( - &Some("alpha".into()), - &Some("gamma".into()), - &camp - )); - // No change, cleared default, or empty campaign set → never campaign-only. - assert!(!is_campaign_only_flip( - &Some("beta".into()), - &Some("beta".into()), - &camp - )); - assert!(!is_campaign_only_flip(&Some("beta".into()), &None, &camp)); - assert!(!is_campaign_only_flip( - &Some("alpha".into()), - &Some("beta".into()), - &std::collections::HashSet::new() - )); - } - - /// A campaign-only flip must NOT reselect a live session whose current model - /// is still selectable; a non-campaign flip must. "Campaign-driven" is marked - /// by `default_is_campaign_driven` on the incoming config. - #[test] - fn campaign_only_flip_does_not_reselect_live_session() { - let mgr = test_manager(); - let mut cfg = config::Config::default(); - cfg.models.default = Some("alpha".to_string()); - mgr.apply_refresh_result(&cfg, Some(make_prefetched(&["alpha", "beta"])), None); - *mgr.inner.cfg.write() = cfg.clone(); // old_preferred = "alpha" - assert_eq!(mgr.current_model_id().0.as_ref(), "alpha"); - - let mut new_cfg = config::Config::default(); - new_cfg.models.default = Some("beta".to_string()); - new_cfg.models.default_is_campaign_driven = true; // campaign overriding - mgr.apply_config(new_cfg); - assert_eq!( - mgr.current_model_id().0.as_ref(), - "alpha", - "campaign-only flip must not yank a still-selectable live session" - ); - - // Control: same flip with no campaign (no pre_campaign_default) → reselect. - let mgr2 = test_manager(); - let mut cfg2 = config::Config::default(); - cfg2.models.default = Some("alpha".to_string()); - mgr2.apply_refresh_result(&cfg2, Some(make_prefetched(&["alpha", "beta"])), None); - *mgr2.inner.cfg.write() = cfg2.clone(); - let mut new_cfg2 = config::Config::default(); - new_cfg2.models.default = Some("beta".to_string()); - mgr2.apply_config(new_cfg2); - assert_eq!( - mgr2.current_model_id().0.as_ref(), - "beta", - "a non-campaign preferred change must reselect" - ); - } - - /// A campaign default missing from the catalog falls back to - /// `pre_campaign_default`, then to the first visible model — and only when - /// the missing pref is actually the campaign-driven config value. - #[test] - fn unavailable_campaign_default_falls_back_to_config_default() { - let catalog = make_prefetched(&["real-model", "other-model"]); - - let mut cfg = config::Config::default(); - cfg.models.default = Some("missing-model".to_string()); - cfg.models.default_is_campaign_driven = true; - cfg.models.pre_campaign_default = Some("real-model".to_string()); - let (key, _, _) = resolve_default_model(&cfg, &catalog, true); - assert_eq!( - key, "real-model", - "must fall back to the pre-campaign default" - ); - - // Control: pre-campaign default also absent → first visible model. - let mut cfg2 = config::Config::default(); - cfg2.models.default = Some("missing-model".to_string()); - cfg2.models.default_is_campaign_driven = true; - cfg2.models.pre_campaign_default = Some("also-missing".to_string()); - let (key2, _, _) = resolve_default_model(&cfg2, &catalog, true); - assert_eq!(&key2, catalog.keys().next().unwrap()); - - // Control: not campaign-driven (e.g. stale recovery value alongside a - // user-set default) → the campaign detour must NOT fire; a missing - // config pref falls to the first visible model. - let mut cfg3 = config::Config::default(); - cfg3.models.default = Some("missing-model".to_string()); - cfg3.models.pre_campaign_default = Some("real-model".to_string()); - let (key3, _, _) = resolve_default_model(&cfg3, &catalog, true); - assert_eq!( - &key3, - catalog.keys().next().unwrap(), - "non-campaign catalog miss must not recover via campaign state" - ); - - // Control: CLI override misses the catalog while campaign state is set - // → CLI is not a campaign problem; no campaign detour. - let mut cfg4 = config::Config { - default_model_override: Some("missing-cli-model".to_string()), - ..Default::default() - }; - cfg4.models.default = Some("campaign-model".to_string()); - cfg4.models.default_is_campaign_driven = true; - cfg4.models.pre_campaign_default = Some("real-model".to_string()); - let (key4, _, _) = resolve_default_model(&cfg4, &catalog, true); - assert_eq!( - &key4, - catalog.keys().next().unwrap(), - "a CLI pref miss must not detour through pre_campaign_default" - ); - } - - // ── ModelFetchAuth::resolve priority tests ────────────────────── - - use serial_test::serial; - use xai_grok_test_support::EnvGuard; - - #[test] - #[serial] - fn resolve_custom_endpoint_always_wins() { - let _key = EnvGuard::set("XAI_API_KEY", "test-key"); - let endpoints = config::EndpointsConfig { - models_base_url: Some("https://custom.example.com".to_owned()), - ..config::EndpointsConfig::default() - }; - assert_eq!( - ModelFetchAuth::resolve(&endpoints, true), - ModelFetchAuth::CustomEndpoint, - ); - assert_eq!( - ModelFetchAuth::resolve(&endpoints, false), - ModelFetchAuth::CustomEndpoint, - ); - } - - #[test] - #[serial] - fn resolve_cached_session_wins_over_api_key() { - let _key = EnvGuard::set("XAI_API_KEY", "test-key"); - let endpoints = config::EndpointsConfig::default(); - assert_eq!( - ModelFetchAuth::resolve(&endpoints, true), - ModelFetchAuth::Session, - "cached session should take priority over API key", - ); - } - - #[test] - #[serial] - fn resolve_api_key_used_when_no_session() { - let _key = EnvGuard::set("XAI_API_KEY", "test-key"); - let endpoints = config::EndpointsConfig::default(); - assert_eq!( - ModelFetchAuth::resolve(&endpoints, false), - ModelFetchAuth::ApiKey, - "API key should be used when no cached session exists", - ); - } - - #[test] - #[serial] - fn resolve_falls_back_to_session_when_nothing_set() { - let _unset = EnvGuard::unset("XAI_API_KEY"); - let _unset_legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY"); - let endpoints = config::EndpointsConfig::default(); - assert_eq!( - ModelFetchAuth::resolve(&endpoints, false), - ModelFetchAuth::Session, - "should fall back to Session when nothing else is configured", - ); - } - - #[test] - #[serial] - fn resolve_deployment_key_when_no_session_or_api_key() { - let _unset = EnvGuard::unset("XAI_API_KEY"); - let _unset_legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY"); - let endpoints = config::EndpointsConfig { - deployment_key: Some("deploy-key".to_owned()), - ..config::EndpointsConfig::default() - }; - assert_eq!( - ModelFetchAuth::resolve(&endpoints, false), - ModelFetchAuth::Deployment, - ); - } - - /// `deployment_key` outranks a stray `XAI_API_KEY`, but session wins over both. - #[test] - #[serial] - fn resolve_deployment_key_outranks_ambient_api_key() { - let _key = EnvGuard::set("XAI_API_KEY", "stray-env-key"); - let endpoints = config::EndpointsConfig { - deployment_key: Some("deploy-key".to_owned()), - ..config::EndpointsConfig::default() - }; - assert_eq!( - ModelFetchAuth::resolve(&endpoints, false), - ModelFetchAuth::Deployment, - "managed deployment_key should outrank an ambient XAI_API_KEY", - ); - assert_eq!( - ModelFetchAuth::resolve(&endpoints, true), - ModelFetchAuth::Session, - "an active session should still win over a managed deployment", - ); - } - - // ── remote_fetch gate: resolve_prefetch_env_from_parts ─────────── - - /// remote_fetch=false must return `None` against every re-arming shape at - /// once — session auth, ambient `XAI_API_KEY`, `deployment_key`, AND a - /// custom models endpoint (which normally forces the prefetch to run). - #[test] - #[serial] - fn prefetch_env_none_when_remote_fetch_disabled_despite_credentials() { - let _key = EnvGuard::set("XAI_API_KEY", "stray-env-key"); - let endpoints = config::EndpointsConfig { - deployment_key: Some("deploy-key".to_owned()), - models_base_url: Some("https://custom.example.com".to_owned()), - ..config::EndpointsConfig::default() - }; - assert!( - resolve_prefetch_env_from_parts( - Some(GrokAuth::test_default()), - endpoints.clone(), - false, - ) - .is_none(), - "session auth must not re-arm the prefetch when remote_fetch is off", - ); - assert!( - resolve_prefetch_env_from_parts(None, endpoints, false).is_none(), - "API key / deployment key / custom endpoint must not re-arm it either", - ); - } - - /// Inverse sanity: with remote_fetch enabled the same credential shapes DO - /// arm the prefetch, and the credential-less default still doesn't. - #[test] - #[serial] - fn prefetch_env_resolves_when_remote_fetch_enabled() { - let _unset = EnvGuard::unset("XAI_API_KEY"); - let _unset_legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY"); - let endpoints = config::EndpointsConfig { - deployment_key: Some("deploy-key".to_owned()), - ..config::EndpointsConfig::default() - }; - assert!(resolve_prefetch_env_from_parts(None, endpoints, true).is_some()); - assert!( - resolve_prefetch_env_from_parts(None, config::EndpointsConfig::default(), true) - .is_none(), - "no credentials and no custom endpoint must stay a no-prefetch launch", - ); - } - - /// remote_fetch=false: an online catalog refresh is a no-op — nothing is - /// fetched, no real-catalog flag is set, and the static catalog keeps - /// resolving. Covers `list_models`/`do_refresh` online strategies too, - /// which funnel into `fetch_and_apply`/`spawn_fetch`. - #[tokio::test] - async fn fetch_and_apply_degrades_offline_when_remote_fetch_disabled() { - let mgr = test_manager(); - mgr.insert_test_entry( - "static-one", - ModelEntry { - info: config::ModelInfo::fallback("static-one"), - api_key: None, - env_key: None, - auth_provider: None, - api_base_url: None, - }, - ); - - mgr.fetch_and_apply_inner(false).await; - - assert!( - !mgr.has_fetched_real_catalog(), - "no catalog fetch may be recorded when remote_fetch is disabled", - ); - assert!( - mgr.models().contains_key("static-one"), - "the static catalog must keep resolving", - ); - } - - // ── supported_in_api tests ────────────────────────────────────── - - #[test] - fn default_model_skips_oauth_only_for_api_key_users() { - let cfg = config::Config::default(); - let mut catalog = IndexMap::new(); - - let mut oauth_only = ModelEntry { - info: config::ModelInfo::fallback("oauth-only"), - api_key: None, - env_key: None, - auth_provider: None, - api_base_url: None, - }; - oauth_only.info.supported_in_api = false; - catalog.insert("oauth-only".to_string(), oauth_only); - - let public = ModelEntry { - info: config::ModelInfo::fallback("public-model"), - api_key: None, - env_key: None, - auth_provider: None, - api_base_url: None, - }; - catalog.insert("public-model".to_string(), public); - - // API-key user: default should NOT be the oauth-only model - let (key, _, _) = resolve_default_model(&cfg, &catalog, false); - assert_ne!( - key, "oauth-only", - "API-key default must not be an OAuth-only model" - ); - assert_eq!(key, "public-model"); - - // OAuth user: oauth-only is valid as default (it's first in the map) - let (key, _, _) = resolve_default_model(&cfg, &catalog, true); - assert!( - key == "oauth-only" || key == "public-model", - "OAuth user should be able to use either model as default" - ); - } - - #[test] - fn visible_for_auth_logic() { - let mut info = config::ModelInfo::fallback("test"); - - // Default: visible to everyone - assert!(info.visible_for_auth(true)); - assert!(info.visible_for_auth(false)); - - // hidden = true: invisible to everyone - info.hidden = true; - assert!(!info.visible_for_auth(true)); - assert!(!info.visible_for_auth(false)); - - // hidden = false, supported_in_api = false: visible to session only - info.hidden = false; - info.supported_in_api = false; - assert!(info.visible_for_auth(true)); - assert!(!info.visible_for_auth(false)); - } - - // ── duplicate model slug re-keying (A/B experiment "auto" alias) ── - - fn make_entry_config(model: &str, name: Option<&str>) -> config::ModelEntryConfig { - make_entry_config_with_id(None, model, name) - } - - fn make_entry_config_with_id( - id: Option<&str>, - model: &str, - name: Option<&str>, - ) -> config::ModelEntryConfig { - config::ModelEntryConfig { - id: id.map(|s| s.to_owned()), - model: model.to_owned(), - base_url: "https://test.api/v1".to_owned(), - name: name.map(|n| n.to_owned()), - description: None, - max_completion_tokens: None, - temperature: None, - top_p: None, - api_key: None, - env_key: None, - api_backend: Default::default(), - context_window: std::num::NonZeroU64::new(200_000).unwrap(), - auto_compact_threshold_percent: None, - system_prompt_label: None, - extra_headers: IndexMap::new(), - api_base_url: None, - use_concise: false, - agent_type: config::default_agent_type(), - inference_idle_timeout_secs: None, - max_retries: None, - hidden: false, - supported_in_api: true, - auth_scheme: None, - reasoning_effort: None, - supports_reasoning_effort: false, - reasoning_efforts: Vec::new(), - supports_backend_search: false, - compactions_remaining: None, - compaction_at_tokens: None, - show_model_fingerprint: false, - stream_tool_calls: None, - laziness_detector: config::LazinessDetectorPerModelConfig::default(), - } - } - - /// Experiment: two entries share the same routing slug but have distinct ids. - /// Both survive, keyed by their respective ids. - #[test] - fn build_prefetched_map_distinct_ids_same_slug() { - let entries = vec![ - make_entry_config_with_id(Some("auto"), "grok-build", Some("Auto")), - make_entry_config_with_id(Some("grok-build"), "grok-build", Some("Grok Build")), - make_entry_config_with_id( - Some("grok-composer-2.5-fast"), - "grok-composer-2.5-fast", - Some("Grok Fast"), - ), - ]; - let map = build_prefetched_map(entries, None); - - assert_eq!(map.len(), 3, "all three entries should survive"); - assert!(map.contains_key("auto")); - assert!(map.contains_key("grok-build")); - assert!(map.contains_key("grok-composer-2.5-fast")); - assert_eq!( - map["auto"].info.model, "grok-build", - "auto entry should still route to grok-build" - ); - assert_eq!(map["grok-build"].info.model, "grok-build"); - } - - /// No id field — falls back to model slug as key. - #[test] - fn build_prefetched_map_no_id_falls_back_to_slug() { - let entries = vec![ - make_entry_config("model-a", Some("Model A")), - make_entry_config("model-b", Some("Model B")), - ]; - let map = build_prefetched_map(entries, None); - - assert_eq!(map.len(), 2); - assert!(map.contains_key("model-a")); - assert!(map.contains_key("model-b")); - } - - /// Duplicate ids — second overwrites first (same as duplicate slugs before). - #[test] - fn build_prefetched_map_duplicate_id_overwrites() { - let entries = vec![ - make_entry_config_with_id(Some("grok-build"), "grok-build", Some("First")), - make_entry_config_with_id(Some("grok-build"), "grok-build", Some("Second")), - ]; - let map = build_prefetched_map(entries, None); - - assert_eq!(map.len(), 1, "duplicate id: second overwrites first"); - assert_eq!(map["grok-build"].info.name.as_deref(), Some("Second")); - } - - /// Regression: resolve_default_model must match by id before scanning - /// by model slug, otherwise entries sharing a slug resolve to whichever - /// appears first in the catalog. - #[test] - fn resolve_default_model_prefers_id_over_model_slug() { - let mut catalog: IndexMap = IndexMap::new(); - catalog.insert( - "auto-grok-build".to_string(), - make_model_entry("grok-build"), - ); - catalog.insert("grok-build".to_string(), make_model_entry("grok-build")); - - let mut cfg = config::Config::default(); - cfg.models.default = Some("grok-build".to_string()); - - let (key, _, _) = resolve_default_model(&cfg, &catalog, true); - assert_eq!(key, "grok-build", "must match id, not first slug hit"); - } - - /// No id field — falls back to slug as key. - #[test] - fn build_prefetched_map_none_id_falls_back_to_slug() { - let entries = vec![make_entry_config_with_id( - None, - "grok-build", - Some("Grok Build"), - )]; - let map = build_prefetched_map(entries, None); - - assert_eq!(map.len(), 1); - assert!(map.contains_key("grok-build")); - } - - // ── persisted model id → catalog key (session resume) ───────────── - - #[test] - fn resolve_catalog_key_maps_routing_slug_to_config_key() { - let mut models = IndexMap::new(); - models.insert( - "enterprise-grok-build".to_string(), - make_model_entry("grok-4.5"), - ); - models.insert("grok-4.3".to_string(), make_model_entry("grok-4.3")); - - let persisted = acp::ModelId::new("grok-4.5"); - let key = resolve_catalog_key(&models, &persisted).expect("slug must resolve"); - assert_eq!(key.0.as_ref(), "enterprise-grok-build"); - } - - #[test] - fn resolve_catalog_key_prefers_exact_key_match() { - let mut models = IndexMap::new(); - models.insert("grok-4.5".to_string(), make_model_entry("grok-4.5")); - - let persisted = acp::ModelId::new("grok-4.5"); - let key = resolve_catalog_key(&models, &persisted).expect("exact key must resolve"); - assert_eq!(key.0.as_ref(), "grok-4.5"); - } - - #[test] - fn resolve_catalog_key_last_slug_match_wins() { - let mut models = IndexMap::new(); - models.insert( - "default-grok-build".to_string(), - make_model_entry("grok-4.5"), - ); - models.insert("user-grok-build".to_string(), make_model_entry("grok-4.5")); - - let persisted = acp::ModelId::new("grok-4.5"); - let key = resolve_catalog_key(&models, &persisted).expect("slug must resolve"); - assert_eq!(key.0.as_ref(), "user-grok-build"); - } - - #[test] - fn selectable_catalog_key_for_persisted_none_when_resolved_not_available() { - let mut models = IndexMap::new(); - models.insert( - "enterprise-grok-build".to_string(), - make_model_entry("grok-4.5"), - ); - - let available: IndexMap<_, _> = IndexMap::new(); - let persisted = acp::ModelId::new("grok-4.5"); - assert!(selectable_catalog_key_for_persisted(&models, &available, &persisted).is_none()); - } - - #[test] - fn selectable_prefers_available_identity_over_non_selectable_exact_key() { - let mut models = IndexMap::new(); - models.insert("grok-build".to_string(), make_model_entry("grok-build")); - models.insert( - "enterprise-grok-build".to_string(), - make_model_entry("grok-build"), - ); - models.insert("grok-4.3".to_string(), make_model_entry("grok-4.3")); - - let available = test_available_keys(&["enterprise-grok-build", "grok-4.3"]); - - let persisted = acp::ModelId::new("grok-build"); - assert_eq!( - resolve_catalog_key(&models, &persisted) - .expect("exact key exists") - .0 - .as_ref(), - "grok-build" - ); - let key = selectable_catalog_key_for_persisted(&models, &available, &persisted) - .expect("must resolve to selectable section"); - assert_eq!(key.0.as_ref(), "enterprise-grok-build"); - } - - #[test] - fn selectable_matches_routing_slug_when_no_exact_key() { - let mut models = IndexMap::new(); - models.insert( - "enterprise-grok-build".to_string(), - make_model_entry("grok-build"), - ); - models.insert("grok-4.3".to_string(), make_model_entry("grok-4.3")); - - let available = test_available_keys(&["enterprise-grok-build", "grok-4.3"]); - - let persisted = acp::ModelId::new("grok-build"); - let key = selectable_catalog_key_for_persisted(&models, &available, &persisted) - .expect("slug must resolve to selectable key"); - assert_eq!(key.0.as_ref(), "enterprise-grok-build"); - } - - /// A persisted *selectable* catalog key binds to itself even when a later - /// selectable section's routing slug equals that key (exact key wins). - #[test] - fn selectable_prefers_exact_key_over_later_slug_match() { - let mut models = IndexMap::new(); - models.insert("grok-build".to_string(), make_model_entry("grok-4.5")); - models.insert("other".to_string(), make_model_entry("grok-build")); - - let available = test_available_keys(&["grok-build", "other"]); - - let persisted = acp::ModelId::new("grok-build"); - let key = selectable_catalog_key_for_persisted(&models, &available, &persisted) - .expect("exact selectable key must win"); - assert_eq!(key.0.as_ref(), "grok-build"); - } - - fn test_available_keys(keys: &[&str]) -> IndexMap { - keys.iter() - .map(|k| { - let id = acp::ModelId::new(*k); - (id.clone(), acp::ModelInfo::new(id, (*k).to_string())) - }) - .collect() - } -} +mod tests; diff --git a/crates/codegen/xai-grok-shell/src/agent/models/cache.rs b/crates/codegen/xai-grok-shell/src/agent/models/cache.rs new file mode 100644 index 0000000..829fa2b --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/agent/models/cache.rs @@ -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, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) grok_version: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) auth_method: Option, + /// 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, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) etag: Option, + pub(crate) models: IndexMap, +} + +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, + pub(crate) etag: Option, +} + +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 { + 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, + 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::(&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; + } + } +} diff --git a/crates/codegen/xai-grok-shell/src/agent/models/endpoint.rs b/crates/codegen/xai-grok-shell/src/agent/models/endpoint.rs new file mode 100644 index 0000000..5638cfb --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/agent/models/endpoint.rs @@ -0,0 +1,41 @@ +use super::*; + +/// Boxed future returned by [`ModelsEndpoint::fetch_models`]. +pub(crate) type ModelsFetchFuture = + Pin>> + Send>>; + +/// Injectable `/v1/models` transport; tests inject a fake. +pub(crate) trait ModelsEndpoint: Send + Sync { + fn fetch_models( + &self, + endpoints: config::EndpointsConfig, + auth: Option, + 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, + 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, + fetch_auth: ModelFetchAuth, +) -> Option> { + tokio::task::spawn_blocking(move || { + prefetch_models_blocking(&endpoints, auth.as_ref(), fetch_auth) + }) + .await + .unwrap_or(None) +} diff --git a/crates/codegen/xai-grok-shell/src/agent/models/fetch.rs b/crates/codegen/xai-grok-shell/src/agent/models/fetch.rs new file mode 100644 index 0000000..f433236 --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/agent/models/fetch.rs @@ -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, + api_base_url_override: Option, +) -> IndexMap { + let mut map: IndexMap = 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> { + 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>, + Option, +) { + 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> { + 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>, + pub settings: Option, +} + +/// Handle for a startup prefetch thread. +pub type EarlyPrefetchHandle = std::thread::JoinHandle; + +pub(crate) struct PrefetchEnv { + pub(crate) auth: Option, + 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, + endpoints: config::EndpointsConfig, + remote_fetch_enabled: bool, +) -> Option { + 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) -> Option { + 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, + sync_managed: bool, +) -> Option { + 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) -> Option { + 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, +) -> Option { + start_early_prefetch_impl(grok_com_config, false) +} + +fn start_early_prefetch_impl( + grok_com_config: Option, + sync_managed: bool, +) -> Option { + 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 + }); + }); +} diff --git a/crates/codegen/xai-grok-shell/src/agent/models/resolution.rs b/crates/codegen/xai-grok-shell/src/agent/models/resolution.rs new file mode 100644 index 0000000..15872e7 --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/agent/models/resolution.rs @@ -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, + id: &acp::ModelId, +) -> Option { + 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, + available: &IndexMap, + id: &acp::ModelId, +) -> Option { + 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, + new_preferred: &Option, + campaign_defaults: &std::collections::HashSet, +) -> 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, + is_session_auth: bool, +) -> (String, ModelEntry, config::ConfigSource) { + let visible: IndexMap = 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, + is_session_auth: bool, +) -> IndexMap { + let visible: IndexMap = 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>) -> Result, Vec> { + 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 { + let mut catalog: IndexMap = 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, +) -> 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, +) -> 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(()) +} diff --git a/crates/codegen/xai-grok-shell/src/agent/models/tests.rs b/crates/codegen/xai-grok-shell/src/agent/models/tests.rs new file mode 100644 index 0000000..485786e --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/agent/models/tests.rs @@ -0,0 +1,1952 @@ +use super::*; + +fn test_manager() -> ModelsManager { + let _ = tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .with_test_writer() + .try_init(); + let tmp = std::env::temp_dir().join("grok-test-models-manager"); + let auth_manager = Arc::new(AuthManager::new(&tmp, GrokComConfig::default())); + ModelsManagerBuilder::new( + None, + IndexMap::new(), + acp::ModelId::new("default"), + auth_manager, + config::Config::default(), + ) + .cache(test_cache_manager(&tmp)) + .build() +} + +#[tokio::test] +async fn catalog_retry_recovers_after_endpoint_returns() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + struct RecoveringEndpoint { + calls: Arc, + catalog: IndexMap, + } + impl ModelsEndpoint for RecoveringEndpoint { + fn fetch_models( + &self, + _endpoints: config::EndpointsConfig, + _auth: Option, + _fetch_auth: ModelFetchAuth, + ) -> ModelsFetchFuture { + let n = self.calls.fetch_add(1, Ordering::SeqCst); + let out = if n == 0 { + None + } else { + Some(self.catalog.clone()) + }; + Box::pin(async move { out }) + } + } + + let calls = Arc::new(AtomicUsize::new(0)); + let tmp = std::env::temp_dir().join("grok-test-catalog-retry"); + let auth_manager = Arc::new(AuthManager::new(&tmp, GrokComConfig::default())); + let mgr = ModelsManagerBuilder::new( + None, + IndexMap::new(), + acp::ModelId::new("default"), + auth_manager, + config::Config::default(), + ) + .endpoint(Arc::new(RecoveringEndpoint { + calls: calls.clone(), + catalog: make_prefetched(&["grok-4"]), + })) + .build(); + assert!(!mgr.has_fetched_real_catalog()); + + mgr.spawn_catalog_retry_with_backoff(crate::tools::retry::BackoffConfig::new(5, 1, 10)); + + let mut recovered = false; + for _ in 0..200 { + if mgr.has_fetched_real_catalog() { + recovered = true; + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + assert!( + recovered, + "catalog retry did not recover after the endpoint returned" + ); + assert!(mgr.models().contains_key("grok-4")); + assert!( + calls.load(Ordering::SeqCst) >= 2, + "expected a failed attempt then a success", + ); +} + +#[tokio::test] +async fn offline_strategy_serves_cache_without_fetching() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + struct CountingEndpoint { + calls: Arc, + } + impl ModelsEndpoint for CountingEndpoint { + fn fetch_models( + &self, + _endpoints: config::EndpointsConfig, + _auth: Option, + _fetch_auth: ModelFetchAuth, + ) -> ModelsFetchFuture { + self.calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { None }) + } + } + + let calls = Arc::new(AtomicUsize::new(0)); + let tmp = tempfile::TempDir::new().unwrap(); + let auth_manager = Arc::new(AuthManager::new(tmp.path(), GrokComConfig::default())); + let mgr = ModelsManagerBuilder::new( + None, + IndexMap::new(), + acp::ModelId::new("default"), + auth_manager, + config_from_toml("[models]\ndefault = \"grok-4.5\""), + ) + .endpoint(Arc::new(CountingEndpoint { + calls: calls.clone(), + })) + .cache(test_cache_manager(tmp.path())) + .build(); + + let seeder = test_cache_manager(tmp.path()); + let auth_method = mgr.inner.fetch_auth.read().cache_auth_method(); + seeder.persist( + &make_prefetched(&["grok-4.5"]), + Some("etag-x"), + auth_method, + &mgr.cache_origin(), + ); + + mgr.list_models(RefreshStrategy::Offline).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 0, + "Offline must serve the disk cache, never hit the transport", + ); + assert!(mgr.models().contains_key("grok-4.5")); + assert!(mgr.has_fetched_real_catalog()); + assert_eq!( + mgr.current_model_id().0.as_ref(), + "grok-4.5", + "first real catalog from the disk cache must resolve the configured default", + ); +} + +#[tokio::test] +async fn auth_refresh_watcher_refetches_on_notify() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + struct NotifyEndpoint { + calls: Arc, + catalog: IndexMap, + } + impl ModelsEndpoint for NotifyEndpoint { + fn fetch_models( + &self, + _endpoints: config::EndpointsConfig, + _auth: Option, + _fetch_auth: ModelFetchAuth, + ) -> ModelsFetchFuture { + self.calls.fetch_add(1, Ordering::SeqCst); + let catalog = self.catalog.clone(); + Box::pin(async move { Some(catalog) }) + } + } + + let calls = Arc::new(AtomicUsize::new(0)); + let tmp = std::env::temp_dir().join("grok-test-auth-refresh-watcher"); + let auth_manager = Arc::new(AuthManager::new(&tmp, GrokComConfig::default())); + let mgr = ModelsManagerBuilder::new( + None, + IndexMap::new(), + acp::ModelId::new("default"), + auth_manager, + config::Config::default(), + ) + .endpoint(Arc::new(NotifyEndpoint { + calls: calls.clone(), + catalog: make_prefetched(&["grok-4"]), + })) + .build(); + assert!(!mgr.has_fetched_real_catalog()); + + let notify = Arc::new(tokio::sync::Notify::new()); + mgr.start_auth_refresh_watcher(notify.clone()); + notify.notify_one(); + + let mut updated = false; + for _ in 0..200 { + if mgr.has_fetched_real_catalog() { + updated = true; + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + assert!(updated, "watcher did not re-fetch the catalog on notify"); + assert!(mgr.models().contains_key("grok-4")); + assert!(calls.load(Ordering::SeqCst) >= 1); +} + +#[tokio::test(start_paused = true)] +async fn hanging_fetch_does_not_block_refresh() { + struct HangingEndpoint; + impl ModelsEndpoint for HangingEndpoint { + fn fetch_models( + &self, + _endpoints: config::EndpointsConfig, + _auth: Option, + _fetch_auth: ModelFetchAuth, + ) -> ModelsFetchFuture { + Box::pin(std::future::pending()) + } + } + + let tmp = std::env::temp_dir().join("grok-test-hanging-fetch"); + let auth_manager = Arc::new(AuthManager::new(&tmp, GrokComConfig::default())); + let mgr = ModelsManagerBuilder::new( + None, + IndexMap::new(), + acp::ModelId::new("default"), + auth_manager, + config::Config::default(), + ) + .endpoint(Arc::new(HangingEndpoint)) + .build(); + + tokio::time::timeout( + crate::http::STARTUP_FETCH_TIMEOUT * 10, + mgr.fetch_and_apply_inner(true), + ) + .await + .expect("fetch_and_apply_inner must return despite a hanging endpoint"); + + assert!( + !mgr.has_fetched_real_catalog(), + "a timed-out fetch must not mark a real catalog", + ); +} + +#[tokio::test(start_paused = true)] +async fn slow_fetch_within_timeout_still_applies() { + // "Slow but succeeds": a fetch that returns just under STARTUP_FETCH_TIMEOUT + // must still be applied, not degraded to offline. + struct SlowEndpoint { + catalog: IndexMap, + delay: std::time::Duration, + } + impl ModelsEndpoint for SlowEndpoint { + fn fetch_models( + &self, + _endpoints: config::EndpointsConfig, + _auth: Option, + _fetch_auth: ModelFetchAuth, + ) -> ModelsFetchFuture { + let catalog = self.catalog.clone(); + let delay = self.delay; + Box::pin(async move { + tokio::time::sleep(delay).await; + Some(catalog) + }) + } + } + + let tmp = tempfile::TempDir::new().unwrap(); + let auth_manager = Arc::new(AuthManager::new(tmp.path(), GrokComConfig::default())); + let mgr = ModelsManagerBuilder::new( + None, + IndexMap::new(), + acp::ModelId::new("default"), + auth_manager, + config::Config::default(), + ) + .endpoint(Arc::new(SlowEndpoint { + catalog: make_prefetched(&["grok-4"]), + delay: crate::http::STARTUP_FETCH_TIMEOUT / 2, + })) + .build(); + + mgr.fetch_and_apply_inner(true).await; + assert!( + mgr.has_fetched_real_catalog(), + "a fetch within the timeout must apply, not degrade", + ); + assert!(mgr.models().contains_key("grok-4")); +} + +#[tokio::test(start_paused = true)] +async fn etag_refresh_is_bounded_and_single_flighted() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + struct CountingHangEndpoint { + calls: Arc, + } + impl ModelsEndpoint for CountingHangEndpoint { + fn fetch_models( + &self, + _endpoints: config::EndpointsConfig, + _auth: Option, + _fetch_auth: ModelFetchAuth, + ) -> ModelsFetchFuture { + self.calls.fetch_add(1, Ordering::SeqCst); + Box::pin(std::future::pending()) + } + } + + let calls = Arc::new(AtomicUsize::new(0)); + let tmp = tempfile::TempDir::new().unwrap(); + let auth_manager = Arc::new(AuthManager::new(tmp.path(), GrokComConfig::default())); + let mgr = ModelsManagerBuilder::new( + None, + IndexMap::new(), + acp::ModelId::new("default"), + auth_manager, + config::Config::default(), + ) + .endpoint(Arc::new(CountingHangEndpoint { + calls: calls.clone(), + })) + .build(); + + // First etag change spawns a bounded fetch; let the task register in-flight. + mgr.spawn_fetch_inner(Some("etag-1".into()), true); + tokio::task::yield_now().await; + // Single-flight: a second spawn while one is in flight must not fetch again. + mgr.spawn_fetch_inner(Some("etag-2".into()), true); + tokio::task::yield_now().await; + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "single-flight: only one etag fetch in flight at a time", + ); + + // Advance past the bound so the hung fetch is abandoned and the guard clears. + tokio::time::sleep(crate::http::STARTUP_FETCH_TIMEOUT * 2).await; + tokio::task::yield_now().await; + + // Guard released → a later etag change fetches again. + mgr.spawn_fetch_inner(Some("etag-3".into()), true); + tokio::task::yield_now().await; + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "after the timeout cleared the in-flight guard, a new etag fetch proceeds", + ); + + // remote_fetch disabled is a no-op: no additional fetch. + mgr.spawn_fetch_inner(Some("etag-4".into()), false); + tokio::task::yield_now().await; + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "disabled gate must not fetch" + ); +} + +fn config_from_toml(toml: &str) -> config::Config { + config::Config::new_from_toml_cfg(&toml::from_str(toml).unwrap()).unwrap() +} + +#[test] +fn model_show_model_fingerprint_reads_catalog_flag() { + let mgr = test_manager(); + + let mut flagged = ModelEntry { + info: config::ModelInfo::fallback("fp-model"), + api_key: None, + env_key: None, + auth_provider: None, + api_base_url: None, + }; + flagged.info.show_model_fingerprint = true; + mgr.insert_test_entry("fp-model", flagged); + + mgr.insert_test_entry( + "plain-model", + ModelEntry { + info: config::ModelInfo::fallback("plain-model"), + api_key: None, + env_key: None, + auth_provider: None, + api_base_url: None, + }, + ); + + let mut custom = ModelEntry { + info: config::ModelInfo::fallback("enterprise-slug"), + api_key: None, + env_key: None, + auth_provider: None, + api_base_url: None, + }; + custom.info.show_model_fingerprint = true; + mgr.insert_test_entry("enterprise-key", custom); + + assert!(mgr.model_show_model_fingerprint("fp-model")); + assert!(!mgr.model_show_model_fingerprint("plain-model")); + assert!(!mgr.model_show_model_fingerprint("missing-model")); + assert!( + mgr.model_show_model_fingerprint("enterprise-slug"), + "slug lookup must resolve to the catalog key and read the flag", + ); + assert!(mgr.model_show_model_fingerprint("enterprise-key")); +} + +#[test] +fn default_model_honors_allowlist_when_no_default_set() { + let cfg = config_from_toml( + r#" + [models] + allowed_models = ["keep-*"] + [model.zzz-first] + model = "zzz-first" + base_url = "https://api.x.ai/v1" + context_window = 256000 + [model.keep-one] + model = "keep-one" + base_url = "https://api.x.ai/v1" + context_window = 256000 + "#, + ); + let catalog = resolve_model_catalog(&cfg, None); + let (_key, entry, _src) = resolve_default_model(&cfg, &catalog, true); + assert!( + entry.info.user_selectable, + "picked non-selectable {}", + entry.model + ); +} + +#[test] +fn validate_selectable_rejects_bad_allowlists() { + let excluded = config_from_toml( + r#" + [models] + default = "grok-3" + allowed_models = ["grok-4*"] + [model.grok-3] + model = "grok-3" + base_url = "https://api.x.ai/v1" + context_window = 256000 + [model.grok-4] + model = "grok-4" + base_url = "https://api.x.ai/v1" + context_window = 256000 + "#, + ); + let catalog = resolve_model_catalog(&excluded, None); + assert!( + validate_selectable(&excluded, &catalog) + .unwrap_err() + .contains("grok-3") + ); + + let zero = config_from_toml( + r#" + [models] + allowed_models = ["nomatch-*"] + [model.grok-4] + model = "grok-4" + base_url = "https://api.x.ai/v1" + context_window = 256000 + "#, + ); + let catalog = resolve_model_catalog(&zero, None); + assert!(validate_selectable(&zero, &catalog).is_err()); +} + +#[tokio::test] +async fn refresh_if_new_etag_skips_when_same() { + let mgr = test_manager(); + mgr.inner.catalog.write().etag = Some("\"abc123\"".to_string()); + + mgr.refresh_if_new_etag("\"abc123\"".to_string()).await; + assert_eq!( + mgr.inner.catalog.read().etag.as_deref(), + Some("\"abc123\""), + "etag should remain unchanged when same" + ); +} + +#[tokio::test] +async fn set_current_model_id_change_fires_watch_to_all_subscribers() { + let mgr = test_manager(); + let mut rx_a = mgr.subscribe_model_switch(); + let mut rx_b = mgr.subscribe_model_switch(); + let initial_a = *rx_a.borrow_and_update(); + let initial_b = *rx_b.borrow_and_update(); + assert_eq!(initial_a, initial_b); + + mgr.set_current_model_id(acp::ModelId::new("default")); + let same_id_ticked = tokio::time::timeout(std::time::Duration::from_millis(25), rx_a.changed()) + .await + .is_ok(); + assert!( + !same_id_ticked, + "set_current_model_id(same id) must NOT bump the watch generation", + ); + + mgr.set_current_model_id(acp::ModelId::new("grok-4")); + tokio::time::timeout(std::time::Duration::from_millis(100), rx_a.changed()) + .await + .expect("rx_a saw the switch") + .expect("watch channel still open"); + tokio::time::timeout(std::time::Duration::from_millis(100), rx_b.changed()) + .await + .expect("rx_b saw the switch") + .expect("watch channel still open"); + assert_ne!(*rx_a.borrow(), initial_a); + assert_eq!(*rx_a.borrow(), *rx_b.borrow()); + assert!(mgr.model_switch_generation() > initial_a); +} + +#[tokio::test] +async fn model_switch_generation_snapshot_reflects_current_state() { + let mgr = test_manager(); + let start = mgr.model_switch_generation(); + mgr.set_current_model_id(acp::ModelId::new("grok-4")); + assert_eq!(mgr.model_switch_generation(), start + 1); + mgr.set_current_model_id(acp::ModelId::new("grok-4")); + assert_eq!(mgr.model_switch_generation(), start + 1); + mgr.set_current_model_id(acp::ModelId::new("grok-3")); + assert_eq!(mgr.model_switch_generation(), start + 2); +} + +#[test] +fn first_catalog_reselect_bumps_model_switch_watch() { + let mgr = test_manager(); + let start = mgr.model_switch_generation(); + let cfg = config_from_toml("[models]\ndefault = \"grok-4.5\""); + mgr.apply_refresh_result(&cfg, Some(make_prefetched(&["grok-4.5", "grok-4"])), None); + assert_eq!(mgr.current_model_id().0.as_ref(), "grok-4.5"); + assert!( + mgr.model_switch_generation() > start, + "background reselection must fire the model-switch watch", + ); +} + +#[test] +fn reselect_missing_current_model_bumps_watch() { + let mgr = test_manager(); + let cfg = config::Config::default(); + mgr.apply_refresh_result(&cfg, Some(make_prefetched(&["grok-4", "grok-3"])), None); + mgr.set_current_model_id(acp::ModelId::new("grok-4")); + let start = mgr.model_switch_generation(); + // A later catalog drops the current model → reselect_current_model_if_missing. + mgr.apply_refresh_result(&cfg, Some(make_prefetched(&["grok-3"])), None); + assert_ne!(mgr.current_model_id().0.as_ref(), "grok-4"); + assert!( + mgr.model_switch_generation() > start, + "reselecting away from a removed current model must fire the watch", + ); +} + +#[test] +fn rebuild_updates_models_and_available() { + let mgr = test_manager(); + assert!(mgr.models().is_empty()); + assert!(mgr.available().is_empty()); + + let cfg = config::Config::default(); + let mut prefetched = IndexMap::new(); + prefetched.insert( + "test-model".to_string(), + ModelEntry { + info: config::ModelInfo::fallback("test-model"), + api_key: None, + env_key: None, + auth_provider: None, + api_base_url: None, + }, + ); + + mgr.rebuild(&cfg, Some(prefetched)); + + assert!( + !mgr.models().is_empty(), + "models should be populated after rebuild" + ); +} + +#[test] +fn current_reasoning_effort_round_trip() { + let mgr = test_manager(); + assert_eq!(mgr.current_reasoning_effort(), None); + + mgr.set_current_reasoning_effort(Some(ReasoningEffort::High)); + assert_eq!(mgr.current_reasoning_effort(), Some(ReasoningEffort::High)); + + mgr.set_current_reasoning_effort(None); + assert_eq!(mgr.current_reasoning_effort(), None); +} + +#[test] +fn current_reasoning_effort_seeded_from_config() { + let tmp = std::env::temp_dir().join("grok-test-models-manager-seed"); + let auth_manager = Arc::new(AuthManager::new(&tmp, GrokComConfig::default())); + let mut cfg = config::Config::default(); + cfg.models.default_reasoning_effort = Some(ReasoningEffort::Xhigh); + let mgr = ModelsManager::new( + None, + IndexMap::new(), + acp::ModelId::new("default"), + auth_manager, + cfg, + ); + assert_eq!(mgr.current_reasoning_effort(), Some(ReasoningEffort::Xhigh),); +} + +#[test] +fn default_reasoning_effort_only_stamps_supporting_model() { + use indexmap::IndexMap; + + let mut cfg = config::Config::default(); + cfg.models.default = Some("reasoning-model".to_string()); + cfg.models.default_reasoning_effort = Some(ReasoningEffort::High); + + let mut prefetched = IndexMap::new(); + let mut reasoning_entry = ModelEntry { + info: config::ModelInfo::fallback("reasoning-model"), + api_key: None, + env_key: None, + auth_provider: None, + api_base_url: None, + }; + reasoning_entry.info.supports_reasoning_effort = true; + prefetched.insert("reasoning-model".to_string(), reasoning_entry); + + let catalog = resolve_model_catalog(&cfg, Some(prefetched)); + assert_eq!( + catalog["reasoning-model"].info.reasoning_effort, + Some(ReasoningEffort::High), + "reasoning-supporting default model should be stamped", + ); + + let mut cfg = config::Config::default(); + cfg.models.default = Some("plain-model".to_string()); + cfg.models.default_reasoning_effort = Some(ReasoningEffort::High); + + let mut prefetched = IndexMap::new(); + let plain_entry = ModelEntry { + info: config::ModelInfo::fallback("plain-model"), + api_key: None, + env_key: None, + auth_provider: None, + api_base_url: None, + }; + prefetched.insert("plain-model".to_string(), plain_entry); + + let catalog = resolve_model_catalog(&cfg, Some(prefetched)); + assert_eq!( + catalog["plain-model"].info.reasoning_effort, None, + "non-reasoning default model must NOT be stamped with persisted effort", + ); +} + +#[test] +fn reasoning_effort_override_skips_models_that_do_not_offer_level() { + use indexmap::IndexMap; + use xai_grok_sampling_types::ReasoningEffortOption; + + let cfg = config::Config { + reasoning_effort_override: Some(ReasoningEffort::None), + ..Default::default() + }; + + let mut prefetched = IndexMap::new(); + let mut no_none = ModelEntry { + info: config::ModelInfo::fallback("grok-4.5"), + api_key: None, + env_key: None, + auth_provider: None, + api_base_url: None, + }; + no_none.info.supports_reasoning_effort = true; + no_none.info.reasoning_efforts = vec![ReasoningEffortOption { + id: "high".into(), + value: ReasoningEffort::High, + label: "High".into(), + description: None, + default: true, + }]; + no_none.info.reasoning_effort = Some(ReasoningEffort::High); + prefetched.insert("grok-4.5".to_string(), no_none); + + let mut with_none = ModelEntry { + info: config::ModelInfo::fallback("legacy-none"), + api_key: None, + env_key: None, + auth_provider: None, + api_base_url: None, + }; + with_none.info.supports_reasoning_effort = true; + with_none.info.reasoning_efforts = vec![ReasoningEffortOption { + id: "none".into(), + value: ReasoningEffort::None, + label: "None".into(), + description: None, + default: true, + }]; + prefetched.insert("legacy-none".to_string(), with_none); + + let catalog = resolve_model_catalog(&cfg, Some(prefetched)); + assert_eq!( + catalog["grok-4.5"].info.reasoning_effort, + Some(ReasoningEffort::High), + "--effort none must not stamp onto models that do not offer none" + ); + assert_eq!( + catalog["legacy-none"].info.reasoning_effort, + Some(ReasoningEffort::None), + "models that list none should still accept the override" + ); +} + +#[test] +fn config_menu_only_model_derives_support_and_default() { + let mut cfg = config::Config::default(); + cfg.config_models.insert( + "menu-only".to_string(), + config::ConfigModelOverride { + reasoning_efforts: vec![ + ReasoningEffortOption { + id: "balanced".to_string(), + value: ReasoningEffort::Medium, + label: "Balanced".to_string(), + description: None, + default: false, + }, + ReasoningEffortOption { + id: "deep".to_string(), + value: ReasoningEffort::Xhigh, + label: "Deep".to_string(), + description: None, + default: true, + }, + ], + ..Default::default() + }, + ); + cfg.config_models + .insert("plain".to_string(), config::ConfigModelOverride::default()); + + let catalog = resolve_model_catalog(&cfg, None); + let info = &catalog["menu-only"].info; + assert!( + info.supports_reasoning_effort, + "menu-only model must derive support" + ); + assert_eq!( + info.reasoning_effort, + Some(ReasoningEffort::Xhigh), + "derived default = marked-default option value" + ); + assert!(!catalog["plain"].info.supports_reasoning_effort); + assert_eq!(catalog["plain"].info.reasoning_effort, None); + + let tmp = std::env::temp_dir().join("grok-test-models-manager-menu-only"); + let auth_manager = Arc::new(AuthManager::new(&tmp, GrokComConfig::default())); + let mgr = ModelsManager::new( + None, + catalog, + acp::ModelId::new("menu-only"), + auth_manager, + cfg, + ); + assert!(mgr.model_supports_reasoning_effort("menu-only")); + assert_eq!( + mgr.model_default_reasoning_effort("menu-only"), + Some(ReasoningEffort::Xhigh) + ); + assert_eq!(mgr.model_reasoning_efforts("menu-only").len(), 2); + assert!(!mgr.model_supports_reasoning_effort("plain")); + assert_eq!(mgr.model_default_reasoning_effort("plain"), None); +} + +#[test] +fn cli_reasoning_effort_override_only_stamps_supporting_models() { + use indexmap::IndexMap; + + let cfg = config::Config { + reasoning_effort_override: Some(ReasoningEffort::High), + ..config::Config::default() + }; + + let mut prefetched = IndexMap::new(); + let mut reasoning_entry = ModelEntry { + info: config::ModelInfo::fallback("reasoning-model"), + api_key: None, + env_key: None, + auth_provider: None, + api_base_url: None, + }; + reasoning_entry.info.supports_reasoning_effort = true; + prefetched.insert("reasoning-model".to_string(), reasoning_entry); + + let plain_entry = ModelEntry { + info: config::ModelInfo::fallback("plain-model"), + api_key: None, + env_key: None, + auth_provider: None, + api_base_url: None, + }; + prefetched.insert("plain-model".to_string(), plain_entry); + + let catalog = resolve_model_catalog(&cfg, Some(prefetched)); + assert_eq!( + catalog["reasoning-model"].info.reasoning_effort, + Some(ReasoningEffort::High), + "reasoning-supporting model should be stamped", + ); + assert_eq!( + catalog["plain-model"].info.reasoning_effort, None, + "non-reasoning model must NOT be stamped", + ); +} + +#[test] +fn apply_refresh_result_only_updates_etag_on_success() { + let mgr = test_manager(); + let cfg = config::Config::default(); + mgr.inner.catalog.write().etag = Some("\"old\"".to_string()); + + assert!( + !mgr.apply_refresh_result(&cfg, None, Some("\"new\"".to_string())), + "failed refresh should report no update" + ); + assert_eq!( + mgr.inner.catalog.read().etag.as_deref(), + Some("\"old\""), + "etag should remain unchanged when refresh fails" + ); + assert!( + mgr.prefetched().is_none(), + "prefetched models should stay unchanged" + ); +} + +fn make_model_entry(model_id: &str) -> ModelEntry { + ModelEntry { + info: config::ModelInfo::fallback(model_id), + api_key: None, + env_key: None, + auth_provider: None, + api_base_url: None, + } +} + +fn make_prefetched(ids: &[&str]) -> IndexMap { + ids.iter() + .map(|id| (id.to_string(), make_model_entry(id))) + .collect() +} + +// ── startup background refresh ───────────────────────────────────── + +#[test] +fn spawn_background_refresh_is_noop_when_real_catalog_present() { + let mgr = test_manager(); + mgr.inner.catalog.write().has_fetched_real_catalog = true; + mgr.spawn_background_refresh(); // must not panic (no tokio::spawn taken) + assert!(mgr.has_fetched_real_catalog()); +} + +#[test] +fn from_config_without_prefetch_produces_usable_catalog() { + let tmp = tempfile::TempDir::new().unwrap(); + let auth_manager = Arc::new(AuthManager::new(tmp.path(), GrokComConfig::default())); + let cfg = config::Config::default(); + + let mgr = ModelsManager::from_config(&cfg, None, auth_manager).unwrap(); + + let cat = mgr.inner.catalog.read(); + let catalog = &cat.models; + assert!( + !catalog.is_empty(), + "zero-network boot must produce at least one model in the internal catalog" + ); + let default = mgr.current_model_id(); + assert!( + catalog.contains_key(default.0.as_ref()), + "default model {:?} not in internal catalog: {:?}", + default, + catalog.keys().collect::>() + ); + drop(cat); + assert!( + !mgr.has_fetched_real_catalog(), + "cold-cache boot must not claim a real catalog" + ); +} + +// ── auth-change refresh: has_fetched_real_catalog flag ───────────── + +#[test] +fn first_apply_refresh_reselects_default_model() { + let mgr = test_manager(); + let mut cfg = config::Config::default(); + cfg.models.default = Some("grok-3".to_string()); + + assert!(!mgr.has_fetched_real_catalog()); + + let prefetched = make_prefetched(&["grok-3", "grok-4"]); + mgr.apply_refresh_result(&cfg, Some(prefetched), None); + + assert!(mgr.has_fetched_real_catalog()); + assert_eq!(mgr.current_model_id().0.as_ref(), "grok-3"); +} + +#[test] +fn subsequent_apply_refresh_preserves_user_model() { + let mgr = test_manager(); + let mut cfg = config::Config::default(); + cfg.models.default = Some("grok-3".to_string()); + + let prefetched = make_prefetched(&["grok-3", "grok-4"]); + mgr.apply_refresh_result(&cfg, Some(prefetched), None); + mgr.set_current_model_id(acp::ModelId::new("grok-4")); + + mgr.inner.catalog.write().prefetched = None; + mgr.inner.catalog.write().etag = None; + + let prefetched = make_prefetched(&["grok-3", "grok-4"]); + mgr.apply_refresh_result(&cfg, Some(prefetched), None); + + assert_eq!( + mgr.current_model_id().0.as_ref(), + "grok-4", + "user's model selection must survive auth-change refresh" + ); +} + +#[test] +fn subsequent_refresh_reselects_when_model_removed() { + let mgr = test_manager(); + let mut cfg = config::Config::default(); + cfg.models.default = Some("grok-3".to_string()); + + let prefetched = make_prefetched(&["grok-3", "grok-4"]); + mgr.apply_refresh_result(&cfg, Some(prefetched), None); + mgr.set_current_model_id(acp::ModelId::new("grok-4")); + + let prefetched = make_prefetched(&["grok-3", "grok-4.5"]); + mgr.apply_refresh_result(&cfg, Some(prefetched), None); + + assert_eq!( + mgr.current_model_id().0.as_ref(), + "grok-3", + "should fall back to config default when current is removed" + ); +} + +#[test] +fn failed_refresh_does_not_set_has_fetched_real_catalog() { + let mgr = test_manager(); + let cfg = config::Config::default(); + + mgr.apply_refresh_result(&cfg, None, None); + + assert!( + !mgr.has_fetched_real_catalog(), + "failed refresh must not flip has_fetched_real_catalog" + ); +} + +// ── apply_config: honor changed preferred model from config ──────── + +#[test] +fn apply_config_honors_new_preferred_model() { + let mgr = test_manager(); + let mut cfg = config::Config::default(); + cfg.models.default = Some("grok-3".to_string()); + + let prefetched = make_prefetched(&["grok-3", "grok-4"]); + mgr.apply_refresh_result(&cfg, Some(prefetched), None); + mgr.set_current_model_id(acp::ModelId::new("grok-4")); + + let mut stale_cfg = config::Config::default(); + stale_cfg.models.default = None; + *mgr.inner.cfg.write() = stale_cfg; + + let mut new_cfg = config::Config::default(); + new_cfg.models.default = Some("grok-3".to_string()); + mgr.apply_config(new_cfg); + + assert_eq!( + mgr.current_model_id().0.as_ref(), + "grok-3", + "apply_config must honor updated preferred model from config" + ); +} + +#[test] +fn apply_config_preserves_current_when_preferred_unchanged() { + let mgr = test_manager(); + let cfg = config::Config::default(); + + let prefetched = make_prefetched(&["grok-3", "grok-4"]); + mgr.apply_refresh_result(&cfg, Some(prefetched), None); + + mgr.set_current_model_id(acp::ModelId::new("grok-4")); + + let new_cfg = config::Config::default(); + mgr.apply_config(new_cfg); + + assert_eq!( + mgr.current_model_id().0.as_ref(), + "grok-4", + "apply_config must not reset model when preferred hasn't changed" + ); +} + +#[test] +fn apply_config_falls_back_when_preferred_not_in_catalog() { + let mgr = test_manager(); + let mut cfg = config::Config::default(); + cfg.models.default = Some("grok-3".to_string()); + + let prefetched = make_prefetched(&["grok-3", "grok-4"]); + mgr.apply_refresh_result(&cfg, Some(prefetched), None); + + mgr.set_current_model_id(acp::ModelId::new("grok-4")); + + let mut new_cfg = config::Config::default(); + new_cfg.models.default = Some("grok-nonexistent".to_string()); + mgr.apply_config(new_cfg); + + let current = mgr.current_model_id(); + let first_available = mgr.available().keys().next().unwrap().clone(); + assert_eq!( + current.0.as_ref(), + first_available.0.as_ref(), + "should fall back to first visible model when preferred not in catalog" + ); +} + +#[test] +fn apply_config_both_none_preferred_preserves_current() { + let mgr = test_manager(); + let cfg = config::Config::default(); + let prefetched = make_prefetched(&["grok-3", "grok-4"]); + mgr.apply_refresh_result(&cfg, Some(prefetched), None); + mgr.set_current_model_id(acp::ModelId::new("grok-4")); + let new_cfg = config::Config::default(); + mgr.apply_config(new_cfg); + + assert_eq!( + mgr.current_model_id().0.as_ref(), + "grok-4", + "both-None preferred must preserve user's runtime model" + ); +} + +#[test] +fn apply_config_old_some_new_none_preserves_current() { + let mgr = test_manager(); + let mut cfg = config::Config::default(); + cfg.models.default = Some("grok-3".to_string()); + + let prefetched = make_prefetched(&["grok-3", "grok-4"]); + mgr.apply_refresh_result(&cfg, Some(prefetched), None); + assert_eq!(mgr.current_model_id().0.as_ref(), "grok-3"); + + mgr.set_current_model_id(acp::ModelId::new("grok-4")); + + let new_cfg = config::Config::default(); + mgr.apply_config(new_cfg); + + assert_eq!( + mgr.current_model_id().0.as_ref(), + "grok-4", + "old=Some new=None must not reset model (is_some guard)" + ); +} + +// ── end-to-end: auth refresh + config reload compose correctly ─── + +#[test] +fn auth_refresh_then_config_reload_preserves_user_model() { + let mgr = test_manager(); + let mut cfg = config::Config::default(); + cfg.models.default = Some("grok-3".to_string()); + + let prefetched = make_prefetched(&["grok-3", "grok-4"]); + mgr.apply_refresh_result(&cfg, Some(prefetched), None); + + mgr.set_current_model_id(acp::ModelId::new("grok-4")); + + mgr.inner.catalog.write().prefetched = None; + mgr.inner.catalog.write().etag = None; + + let prefetched = make_prefetched(&["grok-3", "grok-4"]); + mgr.apply_refresh_result(&cfg, Some(prefetched), None); + assert_eq!(mgr.current_model_id().0.as_ref(), "grok-4"); + + let mut new_cfg = config::Config::default(); + new_cfg.models.default = Some("grok-4".to_string()); + mgr.apply_config(new_cfg); + assert_eq!(mgr.current_model_id().0.as_ref(), "grok-4"); +} + +// ── disk-cache hot-reload (external models_cache.json writes) ──── + +fn test_cache_manager(dir: &std::path::Path) -> ModelsCacheManager { + ModelsCacheManager { + path: dir.join(MODELS_CACHE_FILE), + ttl: CACHE_TTL, + } +} + +#[test] +fn reload_from_disk_cache_applies_external_catalog() { + let mgr = test_manager(); + let tmp = tempfile::TempDir::new().unwrap(); + let cache = test_cache_manager(tmp.path()); + + let auth_method = mgr.inner.fetch_auth.read().cache_auth_method(); + cache.persist( + &make_prefetched(&["grok-4.5", "grok-4.3"]), + Some("etag-ext"), + auth_method, + &mgr.cache_origin(), + ); + + mgr.reload_from_cache_manager(&cache); + + assert!(mgr.has_fetched_real_catalog()); + assert!(mgr.models().contains_key("grok-4.5")); + assert!(mgr.models().contains_key("grok-4.3")); + assert_eq!(mgr.inner.catalog.read().etag.as_deref(), Some("etag-ext")); +} + +#[test] +fn reload_from_disk_cache_recomputes_allowlist_excludes_all() { + let mgr = test_manager(); + let cfg = config_from_toml("[models]\nallowed_models = [\"keep-*\"]"); + + mgr.apply_refresh_result(&cfg, Some(make_prefetched(&["other-1"])), None); + assert!( + mgr.allowlist_excludes_all(), + "setup: allowlist should exclude the entire catalog" + ); + *mgr.inner.cfg.write() = cfg.clone(); + + let tmp = tempfile::TempDir::new().unwrap(); + let cache = test_cache_manager(tmp.path()); + let auth_method = mgr.inner.fetch_auth.read().cache_auth_method(); + cache.persist( + &make_prefetched(&["keep-1"]), + Some("etag-keep"), + auth_method, + &mgr.cache_origin(), + ); + + mgr.reload_from_cache_manager(&cache); + + assert!(mgr.models().contains_key("keep-1")); + assert!( + !mgr.allowlist_excludes_all(), + "corrective external cache write must unlatch the prompt block" + ); +} + +#[test] +fn reload_from_disk_cache_resolves_default_on_first_catalog() { + let mgr = test_manager(); + assert!(!mgr.has_fetched_real_catalog()); + let cfg = config_from_toml("[models]\ndefault = \"keep-1\""); + *mgr.inner.cfg.write() = cfg.clone(); + + let tmp = tempfile::TempDir::new().unwrap(); + let cache = test_cache_manager(tmp.path()); + let auth_method = mgr.inner.fetch_auth.read().cache_auth_method(); + cache.persist( + &make_prefetched(&["keep-1", "other-1"]), + Some("etag-first"), + auth_method, + &mgr.cache_origin(), + ); + + mgr.reload_from_cache_manager(&cache); + + assert!(mgr.has_fetched_real_catalog()); + assert_eq!( + mgr.current_model_id().0.as_ref(), + "keep-1", + "first real catalog must resolve the configured default" + ); +} + +#[test] +fn reload_from_disk_cache_skips_identical_catalog_and_adopts_etag() { + let mgr = test_manager(); + let cfg = config::Config::default(); + let prefetched = make_prefetched(&["grok-3", "grok-4"]); + mgr.apply_refresh_result(&cfg, Some(prefetched.clone()), Some("etag-a".into())); + mgr.set_current_model_id(acp::ModelId::new("grok-4")); + + let tmp = tempfile::TempDir::new().unwrap(); + let cache = test_cache_manager(tmp.path()); + let auth_method = mgr.inner.fetch_auth.read().cache_auth_method(); + cache.persist( + &prefetched, + Some("etag-b"), + auth_method, + &mgr.cache_origin(), + ); + + mgr.reload_from_cache_manager(&cache); + + assert_eq!( + mgr.current_model_id().0.as_ref(), + "grok-4", + "identical catalog must not disturb the user's model" + ); + assert_eq!( + mgr.inner.catalog.read().etag.as_deref(), + Some("etag-b"), + "etag should be adopted so refresh_if_new_etag stays accurate" + ); +} + +#[test] +fn reload_from_disk_cache_ignores_stale_cache() { + let mgr = test_manager(); + let tmp = tempfile::TempDir::new().unwrap(); + let cache = test_cache_manager(tmp.path()); + let auth_method = mgr.inner.fetch_auth.read().cache_auth_method(); + let stale = ModelsCache { + fetched_at: Utc::now() - ChronoDuration::seconds(3600), + grok_version: Some(xai_grok_version::VERSION.to_string()), + auth_method: Some(auth_method), + origin: Some(mgr.cache_origin()), + etag: Some("etag-stale".into()), + models: make_prefetched(&["grok-stale"]), + }; + cache.atomic_write(&stale); + + mgr.reload_from_cache_manager(&cache); + + assert!(!mgr.models().contains_key("grok-stale")); + assert!(mgr.inner.catalog.read().etag.is_none()); +} + +#[test] +fn reload_from_disk_cache_ignores_auth_method_mismatch() { + let mgr = test_manager(); + let tmp = tempfile::TempDir::new().unwrap(); + let cache = test_cache_manager(tmp.path()); + let current = mgr.inner.fetch_auth.read().cache_auth_method(); + let other = if current == CacheAuthMethod::Session { + CacheAuthMethod::ApiKey + } else { + CacheAuthMethod::Session + }; + cache.persist( + &make_prefetched(&["grok-other-auth"]), + Some("etag-x"), + other, + &mgr.cache_origin(), + ); + + mgr.reload_from_cache_manager(&cache); + + assert!(!mgr.models().contains_key("grok-other-auth")); +} + +#[test] +fn reload_from_disk_cache_ignores_origin_mismatch() { + let mgr = test_manager(); + let tmp = tempfile::TempDir::new().unwrap(); + let cache = test_cache_manager(tmp.path()); + let auth_method = mgr.inner.fetch_auth.read().cache_auth_method(); + cache.persist( + &make_prefetched(&["grok-other-origin"]), + Some("etag-y"), + auth_method, + "http://127.0.0.1:49953/v1/models", + ); + + mgr.reload_from_cache_manager(&cache); + + assert!(!mgr.models().contains_key("grok-other-origin")); + assert!(mgr.inner.catalog.read().etag.is_none()); +} + +#[test] +fn reload_from_disk_cache_ignores_legacy_cache_without_origin() { + let mgr = test_manager(); + let tmp = tempfile::TempDir::new().unwrap(); + let cache = test_cache_manager(tmp.path()); + let auth_method = mgr.inner.fetch_auth.read().cache_auth_method(); + let legacy = ModelsCache { + fetched_at: Utc::now(), + grok_version: Some(xai_grok_version::VERSION.to_string()), + auth_method: Some(auth_method), + origin: None, + etag: Some("etag-legacy".into()), + models: make_prefetched(&["grok-legacy"]), + }; + cache.atomic_write(&legacy); + + mgr.reload_from_cache_manager(&cache); + + assert!(!mgr.models().contains_key("grok-legacy")); +} + +// ── clear() resets has_fetched_real_catalog ────────────────────── + +#[test] +fn clear_resets_has_fetched_real_catalog() { + let mgr = test_manager(); + let mut cfg = config::Config::default(); + cfg.models.default = Some("grok-3".to_string()); + + let prefetched = make_prefetched(&["grok-3", "grok-4"]); + mgr.apply_refresh_result(&cfg, Some(prefetched), None); + assert!(mgr.has_fetched_real_catalog()); + + mgr.clear(); + assert!(!mgr.has_fetched_real_catalog()); + + let prefetched = make_prefetched(&["grok-4.5", "grok-4.3"]); + mgr.apply_refresh_result(&cfg, Some(prefetched), None); + let first_available = mgr.available().keys().next().unwrap().clone(); + assert_eq!( + mgr.current_model_id().0.as_ref(), + first_available.0.as_ref() + ); +} + +#[test] +fn is_campaign_only_flip_detects_campaign_driven_changes() { + let camp: std::collections::HashSet = ["beta".into()].into_iter().collect(); + assert!(is_campaign_only_flip( + &Some("alpha".into()), + &Some("beta".into()), + &camp + )); + assert!(is_campaign_only_flip( + &Some("beta".into()), + &Some("alpha".into()), + &camp + )); + assert!(!is_campaign_only_flip( + &Some("alpha".into()), + &Some("gamma".into()), + &camp + )); + assert!(!is_campaign_only_flip( + &Some("beta".into()), + &Some("beta".into()), + &camp + )); + assert!(!is_campaign_only_flip(&Some("beta".into()), &None, &camp)); + assert!(!is_campaign_only_flip( + &Some("alpha".into()), + &Some("beta".into()), + &std::collections::HashSet::new() + )); +} + +#[test] +fn campaign_only_flip_does_not_reselect_live_session() { + let mgr = test_manager(); + let mut cfg = config::Config::default(); + cfg.models.default = Some("alpha".to_string()); + mgr.apply_refresh_result(&cfg, Some(make_prefetched(&["alpha", "beta"])), None); + *mgr.inner.cfg.write() = cfg.clone(); // old_preferred = "alpha" + assert_eq!(mgr.current_model_id().0.as_ref(), "alpha"); + + let mut new_cfg = config::Config::default(); + new_cfg.models.default = Some("beta".to_string()); + new_cfg.models.default_is_campaign_driven = true; // campaign overriding + mgr.apply_config(new_cfg); + assert_eq!( + mgr.current_model_id().0.as_ref(), + "alpha", + "campaign-only flip must not yank a still-selectable live session" + ); + + let mgr2 = test_manager(); + let mut cfg2 = config::Config::default(); + cfg2.models.default = Some("alpha".to_string()); + mgr2.apply_refresh_result(&cfg2, Some(make_prefetched(&["alpha", "beta"])), None); + *mgr2.inner.cfg.write() = cfg2.clone(); + let mut new_cfg2 = config::Config::default(); + new_cfg2.models.default = Some("beta".to_string()); + mgr2.apply_config(new_cfg2); + assert_eq!( + mgr2.current_model_id().0.as_ref(), + "beta", + "a non-campaign preferred change must reselect" + ); +} + +#[test] +fn unavailable_campaign_default_falls_back_to_config_default() { + let catalog = make_prefetched(&["real-model", "other-model"]); + + let mut cfg = config::Config::default(); + cfg.models.default = Some("missing-model".to_string()); + cfg.models.default_is_campaign_driven = true; + cfg.models.pre_campaign_default = Some("real-model".to_string()); + let (key, _, _) = resolve_default_model(&cfg, &catalog, true); + assert_eq!( + key, "real-model", + "must fall back to the pre-campaign default" + ); + + let mut cfg2 = config::Config::default(); + cfg2.models.default = Some("missing-model".to_string()); + cfg2.models.default_is_campaign_driven = true; + cfg2.models.pre_campaign_default = Some("also-missing".to_string()); + let (key2, _, _) = resolve_default_model(&cfg2, &catalog, true); + assert_eq!(&key2, catalog.keys().next().unwrap()); + + let mut cfg3 = config::Config::default(); + cfg3.models.default = Some("missing-model".to_string()); + cfg3.models.pre_campaign_default = Some("real-model".to_string()); + let (key3, _, _) = resolve_default_model(&cfg3, &catalog, true); + assert_eq!( + &key3, + catalog.keys().next().unwrap(), + "non-campaign catalog miss must not recover via campaign state" + ); + + let mut cfg4 = config::Config { + default_model_override: Some("missing-cli-model".to_string()), + ..Default::default() + }; + cfg4.models.default = Some("campaign-model".to_string()); + cfg4.models.default_is_campaign_driven = true; + cfg4.models.pre_campaign_default = Some("real-model".to_string()); + let (key4, _, _) = resolve_default_model(&cfg4, &catalog, true); + assert_eq!( + &key4, + catalog.keys().next().unwrap(), + "a CLI pref miss must not detour through pre_campaign_default" + ); +} + +// ── ModelFetchAuth::resolve priority tests ────────────────────── + +use serial_test::serial; +use xai_grok_test_support::EnvGuard; + +#[test] +#[serial] +fn resolve_custom_endpoint_always_wins() { + let _key = EnvGuard::set("XAI_API_KEY", "test-key"); + let endpoints = config::EndpointsConfig { + models_base_url: Some("https://custom.example.com".to_owned()), + ..config::EndpointsConfig::default() + }; + assert_eq!( + ModelFetchAuth::resolve(&endpoints, true), + ModelFetchAuth::CustomEndpoint, + ); + assert_eq!( + ModelFetchAuth::resolve(&endpoints, false), + ModelFetchAuth::CustomEndpoint, + ); +} + +#[test] +#[serial] +fn resolve_cached_session_wins_over_api_key() { + let _key = EnvGuard::set("XAI_API_KEY", "test-key"); + let endpoints = config::EndpointsConfig::default(); + assert_eq!( + ModelFetchAuth::resolve(&endpoints, true), + ModelFetchAuth::Session, + "cached session should take priority over API key", + ); +} + +#[test] +#[serial] +fn resolve_api_key_used_when_no_session() { + let _key = EnvGuard::set("XAI_API_KEY", "test-key"); + let endpoints = config::EndpointsConfig::default(); + assert_eq!( + ModelFetchAuth::resolve(&endpoints, false), + ModelFetchAuth::ApiKey, + "API key should be used when no cached session exists", + ); +} + +#[test] +#[serial] +fn resolve_falls_back_to_session_when_nothing_set() { + let _unset = EnvGuard::unset("XAI_API_KEY"); + let _unset_legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY"); + let endpoints = config::EndpointsConfig::default(); + assert_eq!( + ModelFetchAuth::resolve(&endpoints, false), + ModelFetchAuth::Session, + "should fall back to Session when nothing else is configured", + ); +} + +#[test] +#[serial] +fn resolve_deployment_key_when_no_session_or_api_key() { + let _unset = EnvGuard::unset("XAI_API_KEY"); + let _unset_legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY"); + let endpoints = config::EndpointsConfig { + deployment_key: Some("deploy-key".to_owned()), + ..config::EndpointsConfig::default() + }; + assert_eq!( + ModelFetchAuth::resolve(&endpoints, false), + ModelFetchAuth::Deployment, + ); +} + +#[test] +#[serial] +fn resolve_deployment_key_outranks_ambient_api_key() { + let _key = EnvGuard::set("XAI_API_KEY", "stray-env-key"); + let endpoints = config::EndpointsConfig { + deployment_key: Some("deploy-key".to_owned()), + ..config::EndpointsConfig::default() + }; + assert_eq!( + ModelFetchAuth::resolve(&endpoints, false), + ModelFetchAuth::Deployment, + "managed deployment_key should outrank an ambient XAI_API_KEY", + ); + assert_eq!( + ModelFetchAuth::resolve(&endpoints, true), + ModelFetchAuth::Session, + "an active session should still win over a managed deployment", + ); +} + +// ── remote_fetch gate: resolve_prefetch_env_from_parts ─────────── + +#[test] +#[serial] +fn prefetch_env_none_when_remote_fetch_disabled_despite_credentials() { + let _key = EnvGuard::set("XAI_API_KEY", "stray-env-key"); + let endpoints = config::EndpointsConfig { + deployment_key: Some("deploy-key".to_owned()), + models_base_url: Some("https://custom.example.com".to_owned()), + ..config::EndpointsConfig::default() + }; + assert!( + resolve_prefetch_env_from_parts(Some(GrokAuth::test_default()), endpoints.clone(), false,) + .is_none(), + "session auth must not re-arm the prefetch when remote_fetch is off", + ); + assert!( + resolve_prefetch_env_from_parts(None, endpoints, false).is_none(), + "API key / deployment key / custom endpoint must not re-arm it either", + ); +} + +#[test] +#[serial] +fn prefetch_env_resolves_when_remote_fetch_enabled() { + let _unset = EnvGuard::unset("XAI_API_KEY"); + let _unset_legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY"); + let endpoints = config::EndpointsConfig { + deployment_key: Some("deploy-key".to_owned()), + ..config::EndpointsConfig::default() + }; + assert!(resolve_prefetch_env_from_parts(None, endpoints, true).is_some()); + assert!( + resolve_prefetch_env_from_parts(None, config::EndpointsConfig::default(), true).is_none(), + "no credentials and no custom endpoint must stay a no-prefetch launch", + ); +} + +#[tokio::test] +async fn fetch_and_apply_degrades_offline_when_remote_fetch_disabled() { + let mgr = test_manager(); + mgr.insert_test_entry( + "static-one", + ModelEntry { + info: config::ModelInfo::fallback("static-one"), + api_key: None, + env_key: None, + auth_provider: None, + api_base_url: None, + }, + ); + + mgr.fetch_and_apply_inner(false).await; + + assert!( + !mgr.has_fetched_real_catalog(), + "no catalog fetch may be recorded when remote_fetch is disabled", + ); + assert!( + mgr.models().contains_key("static-one"), + "the static catalog must keep resolving", + ); +} + +// ── supported_in_api tests ────────────────────────────────────── + +#[test] +fn default_model_skips_oauth_only_for_api_key_users() { + let cfg = config::Config::default(); + let mut catalog = IndexMap::new(); + + let mut oauth_only = ModelEntry { + info: config::ModelInfo::fallback("oauth-only"), + api_key: None, + env_key: None, + auth_provider: None, + api_base_url: None, + }; + oauth_only.info.supported_in_api = false; + catalog.insert("oauth-only".to_string(), oauth_only); + + let public = ModelEntry { + info: config::ModelInfo::fallback("public-model"), + api_key: None, + env_key: None, + auth_provider: None, + api_base_url: None, + }; + catalog.insert("public-model".to_string(), public); + + let (key, _, _) = resolve_default_model(&cfg, &catalog, false); + assert_ne!( + key, "oauth-only", + "API-key default must not be an OAuth-only model" + ); + assert_eq!(key, "public-model"); + + let (key, _, _) = resolve_default_model(&cfg, &catalog, true); + assert!( + key == "oauth-only" || key == "public-model", + "OAuth user should be able to use either model as default" + ); +} + +#[test] +fn visible_for_auth_logic() { + let mut info = config::ModelInfo::fallback("test"); + + assert!(info.visible_for_auth(true)); + assert!(info.visible_for_auth(false)); + + info.hidden = true; + assert!(!info.visible_for_auth(true)); + assert!(!info.visible_for_auth(false)); + + info.hidden = false; + info.supported_in_api = false; + assert!(info.visible_for_auth(true)); + assert!(!info.visible_for_auth(false)); +} + +// ── duplicate model slug re-keying (A/B experiment "auto" alias) ── + +fn make_entry_config(model: &str, name: Option<&str>) -> config::ModelEntryConfig { + make_entry_config_with_id(None, model, name) +} + +fn make_entry_config_with_id( + id: Option<&str>, + model: &str, + name: Option<&str>, +) -> config::ModelEntryConfig { + config::ModelEntryConfig { + id: id.map(|s| s.to_owned()), + model: model.to_owned(), + base_url: "https://test.api/v1".to_owned(), + name: name.map(|n| n.to_owned()), + description: None, + max_completion_tokens: None, + temperature: None, + top_p: None, + api_key: None, + env_key: None, + api_backend: Default::default(), + context_window: std::num::NonZeroU64::new(200_000).unwrap(), + auto_compact_threshold_percent: None, + system_prompt_label: None, + extra_headers: IndexMap::new(), + api_base_url: None, + use_concise: false, + agent_type: config::default_agent_type(), + inference_idle_timeout_secs: None, + max_retries: None, + hidden: false, + supported_in_api: true, + auth_scheme: None, + reasoning_effort: None, + supports_reasoning_effort: false, + reasoning_efforts: Vec::new(), + supports_backend_search: false, + compactions_remaining: None, + compaction_at_tokens: None, + show_model_fingerprint: false, + stream_tool_calls: None, + laziness_detector: config::LazinessDetectorPerModelConfig::default(), + } +} + +#[test] +fn build_prefetched_map_distinct_ids_same_slug() { + let entries = vec![ + make_entry_config_with_id(Some("auto"), "grok-build", Some("Auto")), + make_entry_config_with_id(Some("grok-build"), "grok-build", Some("Grok Build")), + make_entry_config_with_id( + Some("experimental-fast"), + "experimental-fast", + Some("Grok Fast"), + ), + ]; + let map = build_prefetched_map(entries, None); + + assert_eq!(map.len(), 3, "all three entries should survive"); + assert!(map.contains_key("auto")); + assert!(map.contains_key("grok-build")); + assert!(map.contains_key("experimental-fast")); + assert_eq!( + map["auto"].info.model, "grok-build", + "auto entry should still route to grok-build" + ); + assert_eq!(map["grok-build"].info.model, "grok-build"); +} + +#[test] +fn build_prefetched_map_no_id_falls_back_to_slug() { + let entries = vec![ + make_entry_config("model-a", Some("Model A")), + make_entry_config("model-b", Some("Model B")), + ]; + let map = build_prefetched_map(entries, None); + + assert_eq!(map.len(), 2); + assert!(map.contains_key("model-a")); + assert!(map.contains_key("model-b")); +} + +#[test] +fn build_prefetched_map_duplicate_id_overwrites() { + let entries = vec![ + make_entry_config_with_id(Some("grok-build"), "grok-build", Some("First")), + make_entry_config_with_id(Some("grok-build"), "grok-build", Some("Second")), + ]; + let map = build_prefetched_map(entries, None); + + assert_eq!(map.len(), 1, "duplicate id: second overwrites first"); + assert_eq!(map["grok-build"].info.name.as_deref(), Some("Second")); +} + +#[test] +fn resolve_default_model_prefers_id_over_model_slug() { + let mut catalog: IndexMap = IndexMap::new(); + catalog.insert( + "auto-grok-build".to_string(), + make_model_entry("grok-build"), + ); + catalog.insert("grok-build".to_string(), make_model_entry("grok-build")); + + let mut cfg = config::Config::default(); + cfg.models.default = Some("grok-build".to_string()); + + let (key, _, _) = resolve_default_model(&cfg, &catalog, true); + assert_eq!(key, "grok-build", "must match id, not first slug hit"); +} + +#[test] +fn build_prefetched_map_none_id_falls_back_to_slug() { + let entries = vec![make_entry_config_with_id( + None, + "grok-build", + Some("Grok Build"), + )]; + let map = build_prefetched_map(entries, None); + + assert_eq!(map.len(), 1); + assert!(map.contains_key("grok-build")); +} + +// ── persisted model id → catalog key (session resume) ───────────── + +#[test] +fn resolve_catalog_key_maps_routing_slug_to_config_key() { + let mut models = IndexMap::new(); + models.insert( + "enterprise-grok-build".to_string(), + make_model_entry("grok-4.5"), + ); + models.insert("grok-4.3".to_string(), make_model_entry("grok-4.3")); + + let persisted = acp::ModelId::new("grok-4.5"); + let key = resolve_catalog_key(&models, &persisted).expect("slug must resolve"); + assert_eq!(key.0.as_ref(), "enterprise-grok-build"); +} + +#[test] +fn resolve_catalog_key_prefers_exact_key_match() { + let mut models = IndexMap::new(); + models.insert("grok-4.5".to_string(), make_model_entry("grok-4.5")); + + let persisted = acp::ModelId::new("grok-4.5"); + let key = resolve_catalog_key(&models, &persisted).expect("exact key must resolve"); + assert_eq!(key.0.as_ref(), "grok-4.5"); +} + +#[test] +fn resolve_catalog_key_last_slug_match_wins() { + let mut models = IndexMap::new(); + models.insert( + "default-grok-build".to_string(), + make_model_entry("grok-4.5"), + ); + models.insert("user-grok-build".to_string(), make_model_entry("grok-4.5")); + + let persisted = acp::ModelId::new("grok-4.5"); + let key = resolve_catalog_key(&models, &persisted).expect("slug must resolve"); + assert_eq!(key.0.as_ref(), "user-grok-build"); +} + +#[test] +fn selectable_catalog_key_for_persisted_none_when_resolved_not_available() { + let mut models = IndexMap::new(); + models.insert( + "enterprise-grok-build".to_string(), + make_model_entry("grok-4.5"), + ); + + let available: IndexMap<_, _> = IndexMap::new(); + let persisted = acp::ModelId::new("grok-4.5"); + assert!(selectable_catalog_key_for_persisted(&models, &available, &persisted).is_none()); +} + +#[test] +fn selectable_prefers_available_identity_over_non_selectable_exact_key() { + let mut models = IndexMap::new(); + models.insert("grok-build".to_string(), make_model_entry("grok-build")); + models.insert( + "enterprise-grok-build".to_string(), + make_model_entry("grok-build"), + ); + models.insert("grok-4.3".to_string(), make_model_entry("grok-4.3")); + + let available = test_available_keys(&["enterprise-grok-build", "grok-4.3"]); + + let persisted = acp::ModelId::new("grok-build"); + assert_eq!( + resolve_catalog_key(&models, &persisted) + .expect("exact key exists") + .0 + .as_ref(), + "grok-build" + ); + let key = selectable_catalog_key_for_persisted(&models, &available, &persisted) + .expect("must resolve to selectable section"); + assert_eq!(key.0.as_ref(), "enterprise-grok-build"); +} + +#[test] +fn selectable_matches_routing_slug_when_no_exact_key() { + let mut models = IndexMap::new(); + models.insert( + "enterprise-grok-build".to_string(), + make_model_entry("grok-build"), + ); + models.insert("grok-4.3".to_string(), make_model_entry("grok-4.3")); + + let available = test_available_keys(&["enterprise-grok-build", "grok-4.3"]); + + let persisted = acp::ModelId::new("grok-build"); + let key = selectable_catalog_key_for_persisted(&models, &available, &persisted) + .expect("slug must resolve to selectable key"); + assert_eq!(key.0.as_ref(), "enterprise-grok-build"); +} + +#[test] +fn selectable_prefers_exact_key_over_later_slug_match() { + let mut models = IndexMap::new(); + models.insert("grok-build".to_string(), make_model_entry("grok-4.5")); + models.insert("other".to_string(), make_model_entry("grok-build")); + + let available = test_available_keys(&["grok-build", "other"]); + + let persisted = acp::ModelId::new("grok-build"); + let key = selectable_catalog_key_for_persisted(&models, &available, &persisted) + .expect("exact selectable key must win"); + assert_eq!(key.0.as_ref(), "grok-build"); +} + +fn test_available_keys(keys: &[&str]) -> IndexMap { + keys.iter() + .map(|k| { + let id = acp::ModelId::new(*k); + (id.clone(), acp::ModelInfo::new(id, (*k).to_string())) + }) + .collect() +} + +#[tokio::test(start_paused = true)] +async fn bounded_auth_refresh_times_out_to_none() { + // A hung IdP (never-ready auth future) must degrade to None within the + // bound so a cold-cache boot fetch can't stall on it. + let started = tokio::time::Instant::now(); + let result = + ModelsManager::bounded_auth_refresh(std::future::pending::>()).await; + assert!(result.is_none(), "a hung auth refresh must yield None"); + assert!( + started.elapsed() >= crate::http::STARTUP_AUTH_REFRESH_TIMEOUT, + "must wait the full bound before giving up", + ); +} + +#[tokio::test] +async fn bounded_auth_refresh_passes_through_ready_value() { + let result = + ModelsManager::bounded_auth_refresh(async { Some(GrokAuth::test_default()) }).await; + assert!( + result.is_some(), + "a ready session must pass through unchanged" + ); +} + +#[tokio::test] +async fn explicit_model_pick_survives_first_real_catalog() { + // Non-blocking boot lets the user pick a model before the first real + // catalog lands; that pick must not be clobbered by default reselection. + let mgr = test_manager(); + let cfg = config_from_toml("[models]\ndefault = \"grok-4.5\""); + mgr.set_current_model_id(acp::ModelId::new("grok-4")); + mgr.apply_refresh_result(&cfg, Some(make_prefetched(&["grok-4.5", "grok-4"])), None); + assert_eq!( + mgr.current_model_id().0.as_ref(), + "grok-4", + "an explicit /model pick must survive the first real catalog", + ); +} + +#[tokio::test] +async fn identity_switch_clears_user_pick_latch() { + // After an identity change (`clear()`), the new identity's first catalog must + // reselect its own default rather than inherit the prior user's pick. + let mgr = test_manager(); + let cfg = config_from_toml("[models]\ndefault = \"grok-4.5\""); + mgr.set_current_model_id(acp::ModelId::new("grok-4")); + mgr.clear(); + mgr.apply_refresh_result(&cfg, Some(make_prefetched(&["grok-4.5", "grok-4"])), None); + assert_eq!( + mgr.current_model_id().0.as_ref(), + "grok-4.5", + "a new identity's first catalog must reselect the default after clear()", + ); +} diff --git a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/acp_agent.rs b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/acp_agent.rs index 3de532e..8f04884 100644 --- a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/acp_agent.rs +++ b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/acp_agent.rs @@ -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 = 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 = 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, diff --git a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/agent_ops.rs b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/agent_ops.rs index 976393e..27cebba 100644 --- a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/agent_ops.rs +++ b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/agent_ops.rs @@ -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 { 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 { + 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>, + task: impl std::future::Future + '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>); + 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 { diff --git a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/mod.rs b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/mod.rs index 1c7e1fb..6754170 100644 --- a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/mod.rs +++ b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/mod.rs @@ -502,6 +502,13 @@ struct SettingsUpdateNotification { tips: Option>, slash_command_tags: Option>, announcements: Option>, + /// 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>, gate_message: Option, gate_url: Option, gate_label: Option, @@ -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, - /// 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>, + /// 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, + /// 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, + /// 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>, + /// 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>, /// 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, + /// Test-only: counts `spawn_settings_reapply` tasks spawned past the + /// in-flight guard. + #[cfg(test)] + settings_reapply_spawn_count: std::cell::Cell, + /// Test-only: counts `spawn_post_auth_settings` tasks spawned past its + /// own guard. + #[cfg(test)] + post_auth_settings_spawn_count: std::cell::Cell, } -/// 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()), diff --git a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests.rs b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests.rs index 55b8630..e4ce229 100644 --- a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests.rs +++ b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests.rs @@ -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, +) { + 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, +) -> 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] diff --git a/crates/codegen/xai-grok-shell/src/agent/otel_gate.rs b/crates/codegen/xai-grok-shell/src/agent/otel_gate.rs new file mode 100644 index 0000000..c71370c --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/agent/otel_gate.rs @@ -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>, +} + +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 { + 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)" + ); + } +} diff --git a/crates/codegen/xai-grok-shell/src/agent/server.rs b/crates/codegen/xai-grok-shell/src/agent/server.rs index 8b0f01d..8ae0a65 100644 --- a/crates/codegen/xai-grok-shell/src/agent/server.rs +++ b/crates/codegen/xai-grok-shell/src/agent/server.rs @@ -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), diff --git a/crates/codegen/xai-grok-shell/src/agent/subscription_check.rs b/crates/codegen/xai-grok-shell/src/agent/subscription_check.rs index 2d6a48d..b66a00a 100644 --- a/crates/codegen/xai-grok-shell/src/agent/subscription_check.rs +++ b/crates/codegen/xai-grok-shell/src/agent/subscription_check.rs @@ -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, } /// 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, @@ -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 { diff --git a/crates/codegen/xai-grok-shell/src/auth/flow.rs b/crates/codegen/xai-grok-shell/src/auth/flow.rs index 0ac73cb..38c75cb 100644 --- a/crates/codegen/xai-grok-shell/src/auth/flow.rs +++ b/crates/codegen/xai-grok-shell/src/auth/flow.rs @@ -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 { - 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 { + 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) -> Option { 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 Option { - 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) -> Option { + 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 { } /// 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, - grok_com_config: &GrokComConfig, ) -> Option { + 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(); + } } diff --git a/crates/codegen/xai-grok-shell/src/auth/manager.rs b/crates/codegen/xai-grok-shell/src/auth/manager.rs index ca6e4bb..cd6d782 100644 --- a/crates/codegen/xai-grok-shell/src/auth/manager.rs +++ b/crates/codegen/xai-grok-shell/src/auth/manager.rs @@ -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 diff --git a/crates/codegen/xai-grok-shell/src/auth/manager_tests.rs b/crates/codegen/xai-grok-shell/src/auth/manager_tests.rs index 489d3da..840b429 100644 --- a/crates/codegen/xai-grok-shell/src/auth/manager_tests.rs +++ b/crates/codegen/xai-grok-shell/src/auth/manager_tests.rs @@ -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..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)), diff --git a/crates/codegen/xai-grok-shell/src/auth/mod.rs b/crates/codegen/xai-grok-shell/src/auth/mod.rs index 4a595b1..ab30ae7 100644 --- a/crates/codegen/xai-grok-shell/src/auth/mod.rs +++ b/crates/codegen/xai-grok-shell/src/auth/mod.rs @@ -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, diff --git a/crates/codegen/xai-grok-shell/src/auth/storage.rs b/crates/codegen/xai-grok-shell/src/auth/storage.rs index a232267..335165d 100644 --- a/crates/codegen/xai-grok-shell/src/auth/storage.rs +++ b/crates/codegen/xai-grok-shell/src/auth/storage.rs @@ -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> = 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. diff --git a/crates/codegen/xai-grok-shell/src/config/mod.rs b/crates/codegen/xai-grok-shell/src/config/mod.rs index 49da5f9..098cc78 100644 --- a/crates/codegen/xai-grok-shell/src/config/mod.rs +++ b/crates/codegen/xai-grok-shell/src/config/mod.rs @@ -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) diff --git a/crates/codegen/xai-grok-shell/src/config/reloader.rs b/crates/codegen/xai-grok-shell/src/config/reloader.rs index d64d568..de09cdd 100644 --- a/crates/codegen/xai-grok-shell/src/config/reloader.rs +++ b/crates/codegen/xai-grok-shell/src/config/reloader.rs @@ -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)?; diff --git a/crates/codegen/xai-grok-shell/src/config/tests.rs b/crates/codegen/xai-grok-shell/src/config/tests.rs index de38f0c..40d015d 100644 --- a/crates/codegen/xai-grok-shell/src/config/tests.rs +++ b/crates/codegen/xai-grok-shell/src/config/tests.rs @@ -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 + ); +} diff --git a/crates/codegen/xai-grok-shell/src/leader/client.rs b/crates/codegen/xai-grok-shell/src/leader/client.rs index cc709b4..e74717b 100644 --- a/crates/codegen/xai-grok-shell/src/leader/client.rs +++ b/crates/codegen/xai-grok-shell/src/leader/client.rs @@ -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. /// diff --git a/crates/codegen/xai-grok-shell/src/leader/server.rs b/crates/codegen/xai-grok-shell/src/leader/server.rs index b7a181c..3c80a92 100644 --- a/crates/codegen/xai-grok-shell/src/leader/server.rs +++ b/crates/codegen/xai-grok-shell/src/leader/server.rs @@ -849,7 +849,7 @@ fn make_leader_starting_error(json: &serde_json::Value) -> Option { "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, /// Atomic flag: `true` while the agent has pending (in-flight) requests pub agent_busy: Arc, - /// 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, /// Set the shutdown reason before cancelling so clients receive the correct `ShuttingDown` /// reason. The default value is [`ShutdownReason::Manual`]; send diff --git a/crates/codegen/xai-grok-shell/src/remote/client.rs b/crates/codegen/xai-grok-shell/src/remote/client.rs index dff7a7d..4394973 100644 --- a/crates/codegen/xai-grok-shell/src/remote/client.rs +++ b/crates/codegen/xai-grok-shell/src/remote/client.rs @@ -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), + /// 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 { + 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 { - 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 { - 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, diff --git a/crates/codegen/xai-grok-shell/src/remote/mod.rs b/crates/codegen/xai-grok-shell/src/remote/mod.rs index e97cda4..0894ab2 100644 --- a/crates/codegen/xai-grok-shell/src/remote/mod.rs +++ b/crates/codegen/xai-grok-shell/src/remote/mod.rs @@ -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}; diff --git a/crates/codegen/xai-grok-shell/src/session/mod.rs b/crates/codegen/xai-grok-shell/src/session/mod.rs index 2ec442a..79b9796 100644 --- a/crates/codegen/xai-grok-shell/src/session/mod.rs +++ b/crates/codegen/xai-grok-shell/src/session/mod.rs @@ -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; diff --git a/crates/codegen/xai-grok-shell/src/session/persistence.rs b/crates/codegen/xai-grok-shell/src/session/persistence.rs index 0049f5a..cf34f85 100644 --- a/crates/codegen/xai-grok-shell/src/session/persistence.rs +++ b/crates/codegen/xai-grok-shell/src/session/persistence.rs @@ -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, + }, 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, rx: mpsc::UnboundedReceiver, remote_sync: Option, + /// 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, @@ -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) { + 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, + 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, diff --git a/crates/codegen/xai-grok-shell/src/session/persistence_tests.rs b/crates/codegen/xai-grok-shell/src/session/persistence_tests.rs index 573a7cf..4be1436 100644 --- a/crates/codegen/xai-grok-shell/src/session/persistence_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/persistence_tests.rs @@ -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(); diff --git a/crates/codegen/xai-grok-shell/src/session/storage/mod.rs b/crates/codegen/xai-grok-shell/src/session/storage/mod.rs index 8988807..6445399 100644 --- a/crates/codegen/xai-grok-shell/src/session/storage/mod.rs +++ b/crates/codegen/xai-grok-shell/src/session/storage/mod.rs @@ -1308,81 +1308,119 @@ pub(crate) struct RawChunkMetaPeek { pub host_turn: Option, } -/// 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 }, + /// 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(items: Vec, classify: impl Fn(&T) -> RewindStep) -> Vec { + let mut result: Vec = Vec::with_capacity(items.len()); let mut prompt_starts: Vec = Vec::new(); let mut tracker = UserRunTurnTracker::new(); - for line in &lines { - let (raw_params, is_xai) = if let Ok(env) = serde_json::from_str::>(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::>(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::>(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::>(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) -> Vec { let has_rewinds = updates.iter().any(|u| { matches!( @@ -1396,39 +1434,7 @@ pub fn filter_rewind_updates(updates: Vec) -> Vec if !has_rewinds { return updates; } - - let mut result: Vec = Vec::with_capacity(updates.len()); - let mut prompt_starts: Vec = 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 `` and `` 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 `` / `` 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>> { @@ -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>> { + 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 { + 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> { + let mut acp_updates: Vec = 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> { 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>> { - 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( + session_id: &str, + grok_home: &std::path::Path, + f: F, +) -> std::io::Result { + let Some(updates_path) = resolve_replay_updates_path(session_id, grok_home)? else { + return Ok(ReplayEmission::Empty); }; - - let all: Vec = iter.filter_map(|r| r.ok()).collect(); - let filtered = filter_rewind_updates(all); - - let acp_updates: Vec = 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( + updates_path: &std::path::Path, + mut f: F, +) -> std::io::Result { + // 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, - pub total_live: usize, - /// Replayed spawns with no matching finish (a rewind can drop the finish) — + pub(crate) max_event_seq: Option, + 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 = filter_rewind_lines(lines.clone()) + .iter() + .map(|l| ser(&SessionUpdateEnvelope::from_str(l).unwrap())) + .collect(); + let typed: Vec = lines + .iter() + .map(|l| SessionUpdateEnvelope::from_str(l).unwrap()) + .collect(); + let via_updates: Vec = 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 = raw + .lines() + .map(|l| SessionUpdateEnvelope::from_str(l).unwrap()) + .collect(); + let reference: Vec = 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::>(), + reference.iter().map(ser).collect::>(), + ); + } + // ── prepare_replay_lines tests ─────────────────────────────────────────── /// Envelope with _meta at the params level (where the real agent puts it). diff --git a/crates/codegen/xai-grok-shell/src/session/testkit/e2e.rs b/crates/codegen/xai-grok-shell/src/session/testkit/e2e.rs new file mode 100644 index 0000000..ec20b71 --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/session/testkit/e2e.rs @@ -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( + 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, + } +} diff --git a/crates/codegen/xai-grok-shell/src/session/testkit/mod.rs b/crates/codegen/xai-grok-shell/src/session/testkit/mod.rs new file mode 100644 index 0000000..1e3638f --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/session/testkit/mod.rs @@ -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; diff --git a/crates/codegen/xai-grok-shell/src/session/testkit/synth/bench.rs b/crates/codegen/xai-grok-shell/src/session/testkit/synth/bench.rs new file mode 100644 index 0000000..88b69a4 --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/session/testkit/synth/bench.rs @@ -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 { + 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 +} diff --git a/crates/codegen/xai-grok-shell/src/session/testkit/synth/mod.rs b/crates/codegen/xai-grok-shell/src/session/testkit/synth/mod.rs new file mode 100644 index 0000000..f945eef --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/session/testkit/synth/mod.rs @@ -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}; diff --git a/crates/codegen/xai-grok-shell/src/session/testkit/synth/replay.rs b/crates/codegen/xai-grok-shell/src/session/testkit/synth/replay.rs new file mode 100644 index 0000000..ab958db --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/session/testkit/synth/replay.rs @@ -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(key: &str, found: Option, 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 `_*` env overrides on top of `defaults`, scaling `turns` and + /// `rewind_points` by `_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) -> 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 = (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 `/sessions//` 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" + ); + } +} diff --git a/crates/codegen/xai-grok-shell/src/util/config/campaigns.rs b/crates/codegen/xai-grok-shell/src/util/config/campaigns.rs index 652a347..84479c2 100644 --- a/crates/codegen/xai-grok-shell/src/util/config/campaigns.rs +++ b/crates/codegen/xai-grok-shell/src/util/config/campaigns.rs @@ -237,6 +237,59 @@ pub fn load_effective_config_disk_only() -> std::io::Result { 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, +} + +/// 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 { + 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, +) -> Option { + 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 { 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 = ["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] diff --git a/crates/codegen/xai-grok-shell/src/util/config/mod.rs b/crates/codegen/xai-grok-shell/src/util/config/mod.rs index 9f027d4..d54d400 100644 --- a/crates/codegen/xai-grok-shell/src/util/config/mod.rs +++ b/crates/codegen/xai-grok-shell/src/util/config/mod.rs @@ -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::*; diff --git a/crates/codegen/xai-grok-shell/tests/session_fork_replay_memory.rs b/crates/codegen/xai-grok-shell/tests/session_fork_replay_memory.rs new file mode 100644 index 0000000..c7d3c7e --- /dev/null +++ b/crates/codegen/xai-grok-shell/tests/session_fork_replay_memory.rs @@ -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 { + let iter = UpdatesIterator::open(updates_path) + .expect("open updates") + .expect("updates file exists"); + let all: Vec = 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 = 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 = 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" + ); +} diff --git a/crates/codegen/xai-grok-shell/tests/session_load_perf.rs b/crates/codegen/xai-grok-shell/tests/session_load_perf.rs index 3b426c9..bbcef09 100644 --- a/crates/codegen/xai-grok-shell/tests/session_load_perf.rs +++ b/crates/codegen/xai-grok-shell/tests/session_load_perf.rs @@ -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 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=`. //! -//! 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 = (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":}`. -/// -/// 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, - bytes: BTreeMap, -} - -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 `/sessions//` without depending on the (internal) -/// cwd encoder: scan the one level of cwd dirs for a child named ``. -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, + bytes: BTreeMap, +} + 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, last_at: Option, @@ -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::()) .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" diff --git a/crates/codegen/xai-grok-shell/tests/test_leader_soak.rs b/crates/codegen/xai-grok-shell/tests/test_leader_soak.rs index db0d2a1..0916afe 100644 --- a/crates/codegen/xai-grok-shell/tests/test_leader_soak.rs +++ b/crates/codegen/xai-grok-shell/tests/test_leader_soak.rs @@ -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 { - #[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), diff --git a/crates/codegen/xai-grok-shell/tests/test_nonblocking_startup.rs b/crates/codegen/xai-grok-shell/tests/test_nonblocking_startup.rs new file mode 100644 index 0000000..0c30724 --- /dev/null +++ b/crates/codegen/xai-grok-shell/tests/test_nonblocking_startup.rs @@ -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; +} diff --git a/crates/codegen/xai-grok-shell/tests/test_nonblocking_startup_offline.rs b/crates/codegen/xai-grok-shell/tests/test_nonblocking_startup_offline.rs new file mode 100644 index 0000000..f3186d6 --- /dev/null +++ b/crates/codegen/xai-grok-shell/tests/test_nonblocking_startup_offline.rs @@ -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; +} diff --git a/crates/codegen/xai-grok-shell/tests/test_session_load_memory.rs b/crates/codegen/xai-grok-shell/tests/test_session_load_memory.rs new file mode 100644 index 0000000..163d28b --- /dev/null +++ b/crates/codegen/xai-grok-shell/tests/test_session_load_memory.rs @@ -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(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) -> std::thread::JoinHandle { + 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, + stop: Arc, + handle: std::thread::JoinHandle, + } + + 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, + 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>, + } + + #[async_trait::async_trait(?Send)] + impl acp::Client for CountingClient { + async fn request_permission( + &self, + args: acp::RequestPermissionRequest, + ) -> acp::Result { + 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); + } +} diff --git a/crates/codegen/xai-grok-shell/tests/test_settings_refresh.rs b/crates/codegen/xai-grok-shell/tests/test_settings_refresh.rs index 2fd25c3..27100bf 100644 --- a/crates/codegen/xai-grok-shell/tests/test_settings_refresh.rs +++ b/crates/codegen/xai-grok-shell/tests/test_settings_refresh.rs @@ -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(); diff --git a/crates/codegen/xai-grok-shell/tests/test_subagent_soak.rs b/crates/codegen/xai-grok-shell/tests/test_subagent_soak.rs new file mode 100644 index 0000000..57d5523 --- /dev/null +++ b/crates/codegen/xai-grok-shell/tests/test_subagent_soak.rs @@ -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; +} + +impl MetricValue for ResourceSnapshot { + fn value_of(&self, metric: Metric) -> Option { + // 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 { + let ResourceGrowth { rss, threads, fds } = *self; + match metric { + Metric::Rss => rss, + Metric::Threads => threads, + Metric::Fds => fds, + } + } +} + +fn serialize_metrics( + value: &T, + serializer: S, +) -> Result { + 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, + quiesced: bool, +} + +fn serialize_counts( + counts: &SubagentRegistryCounts, + serializer: S, +) -> Result { + // 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 { + #[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; + + 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>; + type ValidateFuture = LocalBoxFuture; + type DescribeFuture = LocalBoxFuture; + + fn run(&self, run: ChildRunRequest) -> 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, + _parent: String, + ) -> Self::DescribeFuture { + Box::pin(std::future::ready(SubagentDescribeOutcome::Unavailable)) + } + + fn on_completed(&self, _completion: ChildCompletion) {} +} + +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 { + // 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) -> 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:?}" + ); + } +} diff --git a/crates/codegen/xai-grok-shell/tests/testkit_synth_roundtrip.rs b/crates/codegen/xai-grok-shell/tests/testkit_synth_roundtrip.rs new file mode 100644 index 0000000..0528e72 --- /dev/null +++ b/crates/codegen/xai-grok-shell/tests/testkit_synth_roundtrip.rs @@ -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"); +} diff --git a/crates/codegen/xai-grok-telemetry/src/events.rs b/crates/codegen/xai-grok-telemetry/src/events.rs index 886b30a..c6e138a 100644 --- a/crates/codegen/xai-grok-telemetry/src/events.rs +++ b/crates/codegen/xai-grok-telemetry/src/events.rs @@ -1218,6 +1218,10 @@ pub struct TerminalTelemetry { pub term_var: String, pub tmux_version: String, pub xtversion: String, + /// Raw, as its source reported it — shapes vary (`"3.5.6"`, + /// `"20240203-110809-5046fc22"`, `"7402"`). Empty when unknown. + pub term_version: String, + pub term_version_source: String, pub host_os: String, pub display_server: String, pub modifier_cmd_fate: String, @@ -1847,6 +1851,8 @@ mod tests { term_var: "xterm-256color".into(), tmux_version: "".into(), xtversion: "".into(), + term_version: "".into(), + term_version_source: "none".into(), host_os: "linux".into(), display_server: "unknown".into(), modifier_cmd_fate: "unknown".into(), diff --git a/crates/codegen/xai-grok-telemetry/src/external/mod.rs b/crates/codegen/xai-grok-telemetry/src/external/mod.rs index c34cc47..de29f30 100644 --- a/crates/codegen/xai-grok-telemetry/src/external/mod.rs +++ b/crates/codegen/xai-grok-telemetry/src/external/mod.rs @@ -72,7 +72,7 @@ impl IdentityAttrs { /// reach init). #[derive(Debug, Clone, Copy, Default)] pub struct ExternalOtelRemotePolicy { - /// Fleet kill switch: flush, then drop subsequent emissions in-process. + /// Remote-policy force-disable: flush, then drop subsequent emissions in-process. pub force_disable: bool, /// Force the content gates off regardless of local env/config. pub lock_content_gates: bool, @@ -219,17 +219,55 @@ fn active_handle() -> Option> { handle().filter(|ext| ext.active.load(Ordering::Relaxed)) } +/// Fail-closed OTEL gate. Defaults open; the leader closes it before init and +/// re-opens it when settings arrive (or immediately for a pure env-API-key +/// leader, which has no remote policy to fetch). +/// +/// On the leader, opening is the synchronizing event: `OtelGate::apply_and_open` +/// applies the remote force-disable (`active = false`) and then opens here, so +/// an emitter whose `Acquire` read observes the `Release` open also observes +/// `active = false`; the emit-path `active` load can therefore stay `Relaxed`. +/// Closing is fail-safe and stays `Relaxed`. The follower path force-disables +/// without re-opening and relies on eventual visibility, acceptable because the +/// policy is tighten-only. +static SETTINGS_RESOLVED: AtomicBool = AtomicBool::new(true); + +/// Close the gate (leader preinit + account switch). +pub fn suppress_external_otel_until_settings() { + SETTINGS_RESOLVED.store(false, Ordering::Relaxed); +} + +/// Open the gate. `Release` publishes the force-disable applied just before it. +pub fn mark_external_otel_settings_resolved() { + if !SETTINGS_RESOLVED.swap(true, Ordering::Release) { + tracing::debug!("external otel: settings resolved, emission gate opened"); + } +} + +/// Read the gate. `Acquire` pairs with the `Release` open. +#[inline] +pub fn is_settings_gate_open() -> bool { + SETTINGS_RESOLVED.load(Ordering::Acquire) +} + /// Cheap check used by the fan-out hook and the split-sink call sites: -/// registry present AND the runtime emission gate set. A stale `true` read -/// only costs a wasted mapping, never an export ([`emit`] re-checks). +/// registry present AND the runtime emission gate set AND the settings gate +/// open. A stale `true` read only costs a wasted mapping, never an export +/// ([`emit`] re-checks). pub fn is_active() -> bool { - matches!(EXTERNAL.get(), Some(Some(ext)) if ext.active.load(Ordering::Relaxed)) + is_settings_gate_open() + && matches!(EXTERNAL.get(), Some(Some(ext)) if ext.active.load(Ordering::Relaxed)) } /// Map and emit one typed telemetry event. No-op unless the stream is active /// and the event has an `external = …` mapping. Synchronous and cheap (the /// batch processor queues; nothing blocks on I/O). pub fn emit(data: &T) { + // Fail-closed: suppress until the leader confirms the remote policy; open by + // default for everyone else. + if !is_settings_gate_open() { + return; + } let Some(ext) = active_handle() else { return; }; @@ -267,8 +305,8 @@ pub(crate) fn set_identity_on(ext: &ExternalTelemetry, attrs: IdentityAttrs) { } /// Apply remote policy when `RemoteSettings` arrive (post-auth, alongside -/// [`set_identity`]). **TIGHTEN-ONLY**: may clear `active` (fleet kill switch -/// — flushes, then drops subsequent emissions) and may force content gates +/// [`set_identity`]). **TIGHTEN-ONLY**: may clear `active` (remote-policy +/// force-disable, flushes then drops subsequent emissions) and may force content gates /// off; it can never enable a stream that env/config left off, and never /// loosens gates mid-run. pub fn apply_remote_policy(policy: ExternalOtelRemotePolicy) { diff --git a/crates/codegen/xai-grok-telemetry/src/external/tests.rs b/crates/codegen/xai-grok-telemetry/src/external/tests.rs index 7c6be62..72bc40e 100644 --- a/crates/codegen/xai-grok-telemetry/src/external/tests.rs +++ b/crates/codegen/xai-grok-telemetry/src/external/tests.rs @@ -960,6 +960,44 @@ async fn remote_gate_lock_forces_gates_off_and_never_on() { ); } +/// All tests that mutate `SETTINGS_RESOLVED` must hold this lock. +static GATE_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + +#[test] +fn settings_gate_suppresses_until_resolved() { + let _serial = GATE_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + + struct RestoreGate; + impl Drop for RestoreGate { + fn drop(&mut self) { + super::mark_external_otel_settings_resolved(); + } + } + let _restore = RestoreGate; + + // Baseline: default open. + super::mark_external_otel_settings_resolved(); + assert!(super::is_settings_gate_open(), "gate defaults open"); + + // Leader closes it at the start of its auth/network phase. + super::suppress_external_otel_until_settings(); + assert!( + !super::is_settings_gate_open(), + "gate must be closed until settings resolve" + ); + assert!( + !super::is_active(), + "is_active must be false while the settings gate is closed" + ); + + // Settings response arrives (policy evaluated) → reopen. + super::mark_external_otel_settings_resolved(); + assert!( + super::is_settings_gate_open(), + "gate must reopen after settings are resolved" + ); +} + // ───────────────────────────────────────────────────────────────────────────── // Metric increment derivation // ───────────────────────────────────────────────────────────────────────────── diff --git a/crates/codegen/xai-grok-telemetry/src/otel_layer/redact.rs b/crates/codegen/xai-grok-telemetry/src/otel_layer/redact.rs index 40668ff..19c89c7 100644 --- a/crates/codegen/xai-grok-telemetry/src/otel_layer/redact.rs +++ b/crates/codegen/xai-grok-telemetry/src/otel_layer/redact.rs @@ -137,6 +137,8 @@ pub(super) static ALLOWED_STRING_KEYS: &[&str] = &[ "terminal.multiplexer", "terminal.tmux_version", "terminal.term_var", + "terminal.term_version", + "terminal.term_version_source", "skip_reason", "auto_cadence_reason", ]; @@ -491,6 +493,8 @@ mod tests { "terminal.multiplexer", "terminal.tmux_version", "terminal.term_var", + "terminal.term_version", + "terminal.term_version_source", "skip_reason", "auto_cadence_reason", ]; diff --git a/crates/codegen/xai-grok-test-support/src/env.rs b/crates/codegen/xai-grok-test-support/src/env.rs index 5538841..ad2ffc0 100644 --- a/crates/codegen/xai-grok-test-support/src/env.rs +++ b/crates/codegen/xai-grok-test-support/src/env.rs @@ -6,6 +6,21 @@ use std::process::Command; use crate::sandbox::TestSandbox; +/// Parse env var `key` into `T`, falling back to `default` when it is unset or +/// present-but-unparseable (warning in the latter case). +pub fn env_parse(key: &str, default: T) -> T { + let Ok(raw) = std::env::var(key) else { + return default; + }; + match raw.parse() { + Ok(value) => value, + Err(_) => { + eprintln!("[test-support] ignoring unparseable {key}={raw:?}; using default"); + default + } + } +} + /// RAII guard for a single environment variable in `#[serial]` tests: snapshots /// the prior value on construction, applies the change, then restores the prior /// value (or unsets it) on drop — even if an assertion panics. Restoring rather @@ -110,7 +125,9 @@ pub fn grok_binary() -> PathBuf { if let Ok(path) = std::env::var("GROK_BINARY") { let p = PathBuf::from(path); assert!(p.exists(), "GROK_BINARY does not exist: {}", p.display()); - return p; + // Bazel's GROK_BINARY is runfiles-relative; the harness spawns the child + // with a different cwd, so absolutize against the (runfiles-root) cwd now. + return std::path::absolute(&p).unwrap_or(p); } if let Ok(path) = std::env::var("CARGO_BIN_EXE_xai-grok-pager") { diff --git a/crates/codegen/xai-grok-test-support/src/leader.rs b/crates/codegen/xai-grok-test-support/src/leader.rs index 898b0b4..89ed82a 100644 --- a/crates/codegen/xai-grok-test-support/src/leader.rs +++ b/crates/codegen/xai-grok-test-support/src/leader.rs @@ -56,6 +56,8 @@ pub struct Capture { chunks: std::sync::Mutex>, notification_count: AtomicU32, reconnected_count: AtomicU32, + models_update_count: AtomicU32, + settings_update_count: AtomicU32, } struct LeaderAcpClient { @@ -96,10 +98,23 @@ impl acp::Client for LeaderAcpClient { } async fn ext_notification(&self, args: acp::ExtNotification) -> acp::Result<()> { - if &*args.method == "x.ai/leader_reconnected" { - self.capture - .reconnected_count - .fetch_add(1, Ordering::SeqCst); + match &*args.method { + "x.ai/leader_reconnected" => { + self.capture + .reconnected_count + .fetch_add(1, Ordering::SeqCst); + } + "x.ai/models/update" => { + self.capture + .models_update_count + .fetch_add(1, Ordering::SeqCst); + } + "x.ai/settings/update" => { + self.capture + .settings_update_count + .fetch_add(1, Ordering::SeqCst); + } + _ => {} } Ok(()) } @@ -180,12 +195,46 @@ impl LeaderFixture { Self::start_with_binary_timeout(binary, server, cwd, sandbox, Duration::from_secs(30)).await } + /// Start a leader pointed at an arbitrary base URL; offline-startup tests + /// aim it at an unreachable endpoint to prove it boots from local data. + pub async fn start_with_base_url( + base_url: &str, + cwd: &Path, + sandbox: &TestSandbox, + ) -> io::Result { + Self::start_with_binary_base_url_timeout( + &grok_binary(), + base_url, + cwd, + sandbox, + Duration::from_secs(30), + ) + .await + } + async fn start_with_binary_timeout( binary: &Path, server: &MockInferenceServer, cwd: &Path, sandbox: &TestSandbox, readiness_timeout: Duration, + ) -> io::Result { + Self::start_with_binary_base_url_timeout( + binary, + &server.url(), + cwd, + sandbox, + readiness_timeout, + ) + .await + } + + async fn start_with_binary_base_url_timeout( + binary: &Path, + base_url: &str, + cwd: &Path, + sandbox: &TestSandbox, + readiness_timeout: Duration, ) -> io::Result { let socket = sandbox.grok_home().join("leader.sock"); let lock = sandbox.grok_home().join("leader.lock"); @@ -202,11 +251,11 @@ impl LeaderFixture { .stdout(std::process::Stdio::null()); sandbox.apply_to_std_command(&mut cmd); cmd.envs(xai_tty_utils::pager_env()) - .env("GROK_CLI_CHAT_PROXY_BASE_URL", server.url()) - .env("GROK_XAI_API_BASE_URL", server.url()) - .env("GROK_MODELS_BASE_URL", server.url()) - .env("GROK_FEEDBACK_BASE_URL", server.url()) - .env("GROK_TRACE_UPLOAD_URL", server.url()) + .env("GROK_CLI_CHAT_PROXY_BASE_URL", base_url) + .env("GROK_XAI_API_BASE_URL", base_url) + .env("GROK_MODELS_BASE_URL", base_url) + .env("GROK_FEEDBACK_BASE_URL", base_url) + .env("GROK_TRACE_UPLOAD_URL", base_url) .env("XAI_API_KEY", "test-key-for-ci") .env("GROK_LEADER_SOCKET", &socket) .env("RUST_LOG", "xai_grok_shell=debug"); @@ -262,6 +311,18 @@ impl LeaderFixture { server: &MockInferenceServer, cwd: &Path, sandbox: &TestSandbox, + ) -> io::Result { + self.spawn_client_with_base_url(&server.url(), cwd, sandbox) + .await + } + + /// Spawn a relay client whose own endpoints point at an arbitrary base URL; + /// pairs with [`Self::start_with_base_url`] for a fully offline stack. + pub async fn spawn_client_with_base_url( + &self, + base_url: &str, + cwd: &Path, + sandbox: &TestSandbox, ) -> io::Result { let binary = self .inner @@ -269,7 +330,7 @@ impl LeaderFixture { .unwrap_or_else(|error| error.into_inner()) .binary .clone(); - self.spawn_client_with_binary(&binary, server, cwd, sandbox) + self.spawn_client_with_binary_base_url(&binary, base_url, cwd, sandbox) .await } @@ -279,6 +340,17 @@ impl LeaderFixture { server: &MockInferenceServer, cwd: &Path, sandbox: &TestSandbox, + ) -> io::Result { + self.spawn_client_with_binary_base_url(binary, &server.url(), cwd, sandbox) + .await + } + + async fn spawn_client_with_binary_base_url( + &self, + binary: &Path, + base_url: &str, + cwd: &Path, + sandbox: &TestSandbox, ) -> io::Result { let socket = self .inner @@ -289,7 +361,7 @@ impl LeaderFixture { let registration = FixtureClientRegistration::new(&self.inner); LeaderStdioClient::spawn_with_binary_and_socket( binary, - server, + base_url, cwd, sandbox, socket, @@ -550,7 +622,7 @@ fn wait_std_child_bounded( impl LeaderStdioClient { async fn spawn_with_binary_and_socket( binary: &Path, - server: &MockInferenceServer, + base_url: &str, cwd: &Path, sandbox: &TestSandbox, leader_socket: PathBuf, @@ -565,11 +637,11 @@ impl LeaderStdioClient { .label("grok leader stdio client") .stdin(TestStdin::Piped) .stdout(TestOutput::Piped) - .env("GROK_CLI_CHAT_PROXY_BASE_URL", server.url()) - .env("GROK_XAI_API_BASE_URL", server.url()) - .env("GROK_MODELS_BASE_URL", server.url()) - .env("GROK_FEEDBACK_BASE_URL", server.url()) - .env("GROK_TRACE_UPLOAD_URL", server.url()) + .env("GROK_CLI_CHAT_PROXY_BASE_URL", base_url) + .env("GROK_XAI_API_BASE_URL", base_url) + .env("GROK_MODELS_BASE_URL", base_url) + .env("GROK_FEEDBACK_BASE_URL", base_url) + .env("GROK_TRACE_UPLOAD_URL", base_url) .env("XAI_API_KEY", "test-key-for-ci") .env("GROK_LEADER_SOCKET", leader_socket) .env("RUST_LOG", "xai_grok_shell=debug"), @@ -753,6 +825,16 @@ impl LeaderStdioClient { pub fn notification_count(&self) -> u32 { self.capture.notification_count.load(Ordering::SeqCst) } + + /// Count of `x.ai/models/update` notifications received (catalog self-heal). + pub fn models_update_count(&self) -> u32 { + self.capture.models_update_count.load(Ordering::SeqCst) + } + + /// Count of `x.ai/settings/update` notifications received (settings self-heal). + pub fn settings_update_count(&self) -> u32 { + self.capture.settings_update_count.load(Ordering::SeqCst) + } } pub fn leader_lock_path(home: &Path) -> PathBuf { diff --git a/crates/codegen/xai-grok-test-support/src/lib.rs b/crates/codegen/xai-grok-test-support/src/lib.rs index 2eb4b19..fa14bc6 100644 --- a/crates/codegen/xai-grok-test-support/src/lib.rs +++ b/crates/codegen/xai-grok-test-support/src/lib.rs @@ -21,6 +21,7 @@ //! - [`grok_binary`] — Resolve the grok binary path (GROK_BINARY env or cargo_bin) //! - [`spawn_counting_server`] — Connection-counting HTTP/1.1 server for wire/pooling tests //! - [`uds_proxy::UdsProxy`] — Frame-aware fault-injection proxy for leader IPC sockets (unix) +//! - [`ResourceSnapshot`] — RSS/threads/fds sampling for soak tests /// Multiply a harness timeout by `GROK_TEST_TIMEOUT_SCALE` (positive integer, /// default 1). CI lanes on shared runner pools raise it so pool load slows /// tests instead of failing them (see the Grok Build merge CI workflow). @@ -41,6 +42,7 @@ mod inference_override; pub mod leader; pub mod mock_server; pub mod process; +pub mod resources; pub mod sandbox; pub mod scripted; pub mod sse; @@ -67,4 +69,5 @@ pub use process::{ TestOutput, TestOutputSnapshot, TestProcess, TestProcessConfig, TestProcessState, TestProcessStderr, TestProcessStdout, TestProcessTermination, TestProcessTree, TestStdin, }; +pub use resources::{ResourceGrowth, ResourceSnapshot}; pub use sandbox::{TestSandbox, TestSandboxBuilder}; diff --git a/crates/codegen/xai-grok-test-support/src/mock_server.rs b/crates/codegen/xai-grok-test-support/src/mock_server.rs index fd1e5d3..50c9294 100644 --- a/crates/codegen/xai-grok-test-support/src/mock_server.rs +++ b/crates/codegen/xai-grok-test-support/src/mock_server.rs @@ -267,6 +267,9 @@ pub struct MockInferenceServer { chunk_delay: Arc>>, /// Mock `/v1/storage` 401 gate + accepted-upload record. storage: Arc, + /// When set, `/v1/models` and `/v1/settings` hang forever (never + /// respond); see [`Self::set_hang`]. + hang: Arc, /// See [`Self::set_user_subscription_tier`]. user_tier: Arc>>, } @@ -306,6 +309,7 @@ impl MockInferenceServer { let messages_stop_reason = Arc::new(std::sync::RwLock::new("end_turn".to_string())); let chunk_delay = Arc::new(std::sync::RwLock::new(None::)); let storage = Arc::new(StorageState::default()); + let hang = Arc::new(std::sync::atomic::AtomicBool::new(false)); let user_tier = Arc::new(std::sync::RwLock::new(None::)); let app = Self::build_router( log.clone(), @@ -317,6 +321,7 @@ impl MockInferenceServer { messages_stop_reason.clone(), chunk_delay.clone(), storage.clone(), + hang.clone(), user_tier.clone(), ); @@ -356,6 +361,7 @@ impl MockInferenceServer { messages_stop_reason, chunk_delay, storage, + hang, user_tier, }) } @@ -426,6 +432,12 @@ impl MockInferenceServer { self.set_settings(json!({ "allow_access": true })); } + /// Make `/v1/models` and `/v1/settings` hang forever, standing in for a + /// black-holed backend in non-blocking-startup tests. + pub fn set_hang(&self, hang: bool) { + self.hang.store(hang, std::sync::atomic::Ordering::Release); + } + /// Set the `subscriptionTier` served by `GET /v1/user`. `None` /// (default) omits the field, which the shell treats as "no qualifying /// subscription" (free tier). @@ -663,8 +675,11 @@ impl MockInferenceServer { messages_stop_reason: Arc>, chunk_delay: Arc>>, storage: Arc, + hang: Arc, user_tier: Arc>>, ) -> Router { + let hang_models = hang.clone(); + let hang_settings = hang; let log_cc = log.clone(); let log_rs = log.clone(); let log_msg = log.clone(); @@ -920,8 +935,12 @@ impl MockInferenceServer { move || { let log = log.clone(); let models = models.clone(); + let hang = hang_models.clone(); async move { log.record("GET", "/v1/models", None, None, Vec::new()); + if hang.load(std::sync::atomic::Ordering::Acquire) { + tokio::time::sleep(Duration::from_secs(3600)).await; + } let models_json = models.read().unwrap().clone(); Json(json!({ "object": "list", @@ -935,12 +954,18 @@ impl MockInferenceServer { "/v1/settings", get({ let log = log.clone(); + let settings = settings.clone(); + let hang = hang_settings.clone(); move || { let log = log.clone(); let settings = settings.clone(); let overrides = overrides_settings.clone(); + let hang = hang.clone(); async move { log.record("GET", "/v1/settings", None, None, Vec::new()); + if hang.load(std::sync::atomic::Ordering::Acquire) { + tokio::time::sleep(Duration::from_secs(3600)).await; + } // Scripted one-shots take precedence (FIFO), so a // test can serve a transient payload (e.g. one // stale gated snapshot) and fall back to the diff --git a/crates/codegen/xai-grok-test-support/src/resources.rs b/crates/codegen/xai-grok-test-support/src/resources.rs new file mode 100644 index 0000000..22cd3b6 --- /dev/null +++ b/crates/codegen/xai-grok-test-support/src/resources.rs @@ -0,0 +1,134 @@ +//! Generic OS resource snapshots for soak tests. No shell types: `rss_bytes` +//! reads `/proc` (Linux) or shells out to `ps` (macOS); the task/fd counters +//! are Linux-only and return `None` elsewhere. + +/// RSS (bytes), live threads, and open fds sampled together. `None` marks a +/// metric the platform can't report. +#[derive(Clone, Copy, Debug, Default)] +pub struct ResourceSnapshot { + pub rss: Option, + pub threads: Option, + pub fds: Option, +} + +/// Saturating per-field growth of one [`ResourceSnapshot`] over an earlier +/// baseline. A distinct type from a snapshot so a delta can't be mistaken for +/// an absolute sample. `None` marks a field either side couldn't report. +#[derive(Clone, Copy, Debug, Default)] +pub struct ResourceGrowth { + pub rss: Option, + pub threads: Option, + pub fds: Option, +} + +impl ResourceSnapshot { + pub fn capture() -> Self { + Self { + rss: rss_bytes(), + threads: thread_count(), + fds: fd_count(), + } + } + + /// RSS only, skipping the thread and fd probes. For hot sampling loops that + /// use just `rss`: on Linux this avoids the per-tick `/proc/self/{task,fd}` + /// directory scans. The RSS read itself still shells out to `ps` on macOS. + pub fn capture_rss() -> Option { + rss_bytes() + } + + /// Growth of `self` (after) over `baseline` (before); see [`ResourceGrowth`]. + pub fn growth_from(&self, baseline: &ResourceSnapshot) -> ResourceGrowth { + let delta = |after: Option, before: Option| { + before.zip(after).map(|(b, a)| a.saturating_sub(b)) + }; + ResourceGrowth { + rss: delta(self.rss, baseline.rss), + threads: delta(self.threads, baseline.threads), + fds: delta(self.fds, baseline.fds), + } + } +} + +fn rss_bytes() -> Option { + #[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 + } +} + +fn thread_count() -> Option { + #[cfg(target_os = "linux")] + { + Some(std::fs::read_dir("/proc/self/task").ok()?.count()) + } + #[cfg(not(target_os = "linux"))] + { + None + } +} + +/// The read's own transient fd closes with the iterator, so before and after +/// samples stay symmetric. +fn fd_count() -> Option { + #[cfg(target_os = "linux")] + { + Some(std::fs::read_dir("/proc/self/fd").ok()?.count()) + } + #[cfg(not(target_os = "linux"))] + { + None + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn growth_from_saturates_and_propagates_none() { + let before = ResourceSnapshot { + rss: Some(100), + threads: Some(5), + fds: None, + }; + let after = ResourceSnapshot { + rss: Some(30), + threads: Some(9), + fds: Some(3), + }; + let growth = after.growth_from(&before); + assert_eq!(growth.rss, Some(0), "a shrink saturates to zero"); + assert_eq!(growth.threads, Some(4), "growth is the delta"); + assert_eq!( + growth.fds, None, + "a missing baseline sample propagates None" + ); + } +} diff --git a/crates/codegen/xai-grok-tools/src/implementations/grok_build/task/coordinator.rs b/crates/codegen/xai-grok-tools/src/implementations/grok_build/task/coordinator.rs index c818b27..f82ee66 100644 --- a/crates/codegen/xai-grok-tools/src/implementations/grok_build/task/coordinator.rs +++ b/crates/codegen/xai-grok-tools/src/implementations/grok_build/task/coordinator.rs @@ -20,9 +20,9 @@ use tokio::sync::{mpsc, oneshot}; use super::coordinator_state::{ ActiveChild, BlockingWaiter, BufferedCompletion, ChildRecord, CompletedChild, InternalEvent, - ListRequest, MAX_COMPLETED_ENTRIES, PendingChild, ProgressFuture, ProgressTarget, ReplyFuture, - TaggedFuture, active_summary, background_at_deadline, background_if_caller_gone, - completed_snapshot, completion_summary, sleep_until, workflow_outstanding, + ListRequest, PendingChild, ProgressFuture, ProgressTarget, ReplyFuture, TaggedFuture, + active_summary, background_at_deadline, background_if_caller_gone, completed_snapshot, + completion_summary, sleep_until, workflow_outstanding, }; use super::types::{ SpawnedSubagentRef, SubagentCancelOutcome, SubagentCancelTarget, SubagentDescribeOutcome, @@ -32,8 +32,8 @@ use super::types::{ pub use super::coordinator_state::{ ChildCompletion, ChildControl, ChildReporter, ChildRunOutput, ChildRunRequest, ChildRunner, - CompletionDisposition, CoordinatorConfig, LocalBoxFuture, SendBoxFuture, StartedChild, - SubagentProgress, + CompletionDisposition, CoordinatorConfig, LocalBoxFuture, MAX_COMPLETED_ENTRIES, SendBoxFuture, + StartedChild, SubagentProgress, }; /// Channel-owned subagent lifecycle actor. diff --git a/crates/codegen/xai-grok-tools/src/implementations/grok_build/task/coordinator_state.rs b/crates/codegen/xai-grok-tools/src/implementations/grok_build/task/coordinator_state.rs index 7d02fe7..d7a83a5 100644 --- a/crates/codegen/xai-grok-tools/src/implementations/grok_build/task/coordinator_state.rs +++ b/crates/codegen/xai-grok-tools/src/implementations/grok_build/task/coordinator_state.rs @@ -13,7 +13,9 @@ use super::types::{ SubagentSnapshotStatus, SubagentValidateTypeOutcome, }; -pub(super) const MAX_COMPLETED_ENTRIES: usize = 1024; +/// Cap on retained completed-subagent entries before the oldest are evicted. +/// Public so the subagent soak test can assert the coordinator stays bounded. +pub const MAX_COMPLETED_ENTRIES: usize = 1024; pub(super) const OUTPUT_UNAVAILABLE_PLACEHOLDER: &str = "[subagent output no longer available]"; pub type LocalBoxFuture = Pin + 'static>>;