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

Synced from monorepo

Browse source 
Changes:
- Persist submitter identity for /feedback
- Let custom models use rotating tokens from named auth providers
- Template stale tool/param name literals in server-native descriptions
- Minimal mode commits thinking in full, lookups as one-liners
- Tighten durable append internals
- Nudge model to end turn on no-op bash commands
- Per-fetch signing nonce in the managed-config envelope, with a server-side replay probe
- Include working tree in startup status
This commit is contained in:
grokkybara[bot] 2026-07-20 18:06:59 +01:00
commit a881e6703f
140 changed files with 6746 additions and 2377 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -651,6 +651,25 @@ pub struct RuntimeResolutionContext<'a> {
/// CLI `--storage-mode` override. `None` = defer to env/remote/default.
pub storage_mode: Option<&'a str>,
}
/// First-party credential env vars scrubbed from a BYOK auth-provider helper's
/// environment so it can't inherit the keys Grok uses for its own first-party
/// requests. Keep in sync with every first-party credential env read across the
/// crate: `auth::manager` (`GROK_AUTH`/`GROK_AUTH_PATH`), `auth_method`
/// (`XAI_API_KEY`/legacy), and the credential-bearing `env_string(...)` reads in
/// `EndpointsConfig::default`. The `provider_helper_env_scrubs_first_party_credentials`
/// test pins this against an independent audited literal, so any change here must
/// be mirrored (and re-audited) there.
pub(crate) const FIRST_PARTY_CREDENTIAL_ENV_VARS: &[&str] = &[
crate::agent::auth_method::XAI_API_KEY_ENV_VAR,
crate::agent::auth_method::LEGACY_XAI_API_KEY_ENV_VAR,
"GROK_AUTH",
"GROK_AUTH_PATH",
"GROK_DEPLOYMENT_KEY",
"GROK_EXTRA_AUTH_KEY",
"GROK_TRACE_UPLOAD_CREDENTIALS_FILE",
"OTEL_EXPORTER_OTLP_HEADERS",
"GROK_INTERNAL_OTLP_HEADERS",
];
/// Read an env var as a trimmed string. Returns `None` if unset or empty/whitespace-only.
pub(crate) fn env_string(name: &str) -> Option<String> {
let value = std::env::var(name).ok()?;
@ -1280,10 +1299,15 @@ pub struct Config {
/// `[model.*]` overrides from config.toml. Resolve via `resolve_model_list()`.
#[serde(skip)]
pub config_models: IndexMap<String, ConfigModelOverride>,
/// Warnings from `[model.*]` parsing; surfaced by `grok inspect`.
/// Warnings from `[model.*]` and `[auth_provider.*]` parsing; surfaced by
/// `grok inspect`.
#[serde(skip)]
pub model_override_warnings: Vec<super::config_model_override_parse::ModelOverrideWarning>,
pub config_warnings: Vec<super::config_model_override_parse::ConfigWarning>,
pub grok_com_config: GrokComConfig,
/// `[auth_provider.<name>]` tables, populated by
/// [`parse_auth_providers`] from trusted config layers only.
#[serde(skip)]
pub auth_providers: IndexMap<String, crate::auth::AuthProviderConfig>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub shortcuts: Option<toml::Value>,
/// Written by the client via `config_toml_edit`; absorbed so it isn't
@ -1708,8 +1732,9 @@ impl Default for Config {
doom_loop_recovery: crate::util::config::DoomLoopRecoverySettings::default(),
auto_mode: AutoModeConfig::default(),
config_models: IndexMap::new(),
model_override_warnings: Vec::new(),
config_warnings: Vec::new(),
grok_com_config: GrokComConfig::default(),
auth_providers: IndexMap::new(),
shortcuts: None,
hints: None,
ui: UiConfig::default(),
@ -1792,6 +1817,101 @@ impl Default for Config {
cfg
}
}
/// Parse `[auth_provider.<name>]` tables leniently: a malformed entry warns
/// (surfaced by `grok inspect`) and is skipped, so it fails closed for the
/// models referencing it instead of failing the whole config.
fn parse_auth_providers(
raw_config: &toml::Value,
) -> (
IndexMap<String, crate::auth::AuthProviderConfig>,
Vec<super::config_model_override_parse::ConfigWarning>,
) {
use super::config_model_override_parse::{ConfigWarning, ConfigWarningKind};
let mut providers = IndexMap::new();
let mut warnings = Vec::new();
let Some(section) = raw_config.get("auth_provider") else {
return (providers, warnings);
};
let Some(table) = section.as_table() else {
warnings.push(ConfigWarning::auth_provider_section(
ConfigWarningKind::NotATable,
format!(
"`auth_provider` must be a table of [auth_provider.<name>] entries, got {}; \
all auth providers ignored",
section.type_str()
),
));
return (providers, warnings);
};
for (name, value) in table {
let mut unknown = Vec::new();
match serde_ignored::deserialize::<_, _, crate::auth::AuthProviderConfig>(
value.clone(),
|path| unknown.push(path.to_string()),
) {
Ok(provider) => {
for key in unknown {
warnings.push(ConfigWarning::auth_provider(
name,
Some(key.as_str()),
ConfigWarningKind::UnknownField,
"unrecognized key; field ignored".to_owned(),
));
}
if !provider.is_usable() {
warnings.push(ConfigWarning::auth_provider(
name,
Some("command"),
ConfigWarningKind::InvalidValue,
"missing or empty command; referencing models resolve \
with no credential"
.to_owned(),
));
}
let skew = crate::auth::PROVIDER_TOKEN_EXPIRY_SKEW_SECS;
if provider.token_ttl_secs.is_some_and(|ttl| ttl <= skew) {
warnings.push(ConfigWarning::auth_provider(
name,
Some("token_ttl_secs"),
ConfigWarningKind::InvalidValue,
format!(
"at or below the {skew}s refresh margin; the command will \
run before every turn"
),
));
}
if let Some(timeout) = provider.timeout_secs
&& !(1..=crate::auth::PROVIDER_TIMEOUT_CEILING_SECS).contains(&timeout)
{
let ceiling = crate::auth::PROVIDER_TIMEOUT_CEILING_SECS;
warnings.push(ConfigWarning::auth_provider(
name,
Some("timeout_secs"),
ConfigWarningKind::InvalidValue,
if timeout == 0 {
"below the 1 second minimum; clamped to 1".to_owned()
} else {
format!("above the {ceiling}s maximum; clamped to {ceiling}")
},
));
}
providers.insert(name.clone(), provider);
}
Err(error) => {
warnings.push(ConfigWarning::auth_provider(
name,
None,
ConfigWarningKind::InvalidValue,
format!(
"failed to parse ({error}); provider skipped, referencing models \
resolve with no credential"
),
));
}
}
}
(providers, warnings)
}
impl Config {
/// Reject invalid glob patterns in the model-filter lists at config load, so
/// a typo fails loudly instead of silently changing availability.
@ -1847,9 +1967,9 @@ impl Config {
let raw_config = &Self::expand_auth_alias(raw_config);
let super::config_model_override_parse::ParsedModelOverrides {
models: config_models,
warnings: model_override_warnings,
warnings: config_warnings,
} = super::config_model_override_parse::parse_model_overrides(raw_config);
super::config_model_override_parse::log_model_override_warnings(&model_override_warnings);
let (auth_providers, auth_provider_warnings) = parse_auth_providers(raw_config);
let mut base = toml::Value::try_from(Self::default()).map_err(|e| e.to_string())?;
if let toml::Value::Table(ref mut t) = base {
t.remove("model");
@ -1857,6 +1977,7 @@ impl Config {
let mut raw_without_model_sections = raw_config.clone();
if let toml::Value::Table(ref mut t) = raw_without_model_sections {
t.remove("model");
t.remove("auth_provider");
}
crate::config::deep_merge_toml(&mut base, &raw_without_model_sections);
let (mut config, user_unused) =
@ -1868,7 +1989,33 @@ impl Config {
);
}
config.config_models = config_models;
config.model_override_warnings = model_override_warnings;
config.config_warnings = config_warnings;
config.auth_providers = auth_providers;
config.config_warnings.extend(auth_provider_warnings);
let declared_provider_names: std::collections::HashSet<&str> = raw_config
.get("auth_provider")
.and_then(toml::Value::as_table)
.map(|t| t.keys().map(String::as_str).collect())
.unwrap_or_default();
for (model_key, model) in &config.config_models {
if let Some(ref name) = model.auth_provider
&& !config.auth_providers.contains_key(name)
&& !declared_provider_names.contains(name.as_str())
{
config.config_warnings.push(
super::config_model_override_parse::ConfigWarning::model(
model_key,
Some("auth_provider"),
super::config_model_override_parse::ConfigWarningKind::InvalidValue,
format!(
"references [auth_provider.{name}], which is not defined; \
the model resolves with no provider credential"
),
),
);
}
}
super::config_model_override_parse::log_config_warnings(&config.config_warnings);
if config.grok_com_config.oidc.is_none() {
config.grok_com_config.oidc = OidcAuthConfig::from_env();
}
@ -3196,11 +3343,24 @@ pub fn resolve_model_list(
let entry = model_override.apply(key, base, &cfg.endpoints);
tracing::debug!(
model_key = % key, base_url = % entry.info.base_url, has_api_key = entry
.api_key.is_some(), env_key = ? entry.env_key, had_base,
.api_key.is_some(), env_key = ? entry.env_key, auth_provider = entry
.auth_provider.as_ref().map(| p | p.name.as_str()), had_base,
"config model override applied"
);
resolved.insert(key.clone(), entry);
}
for (key, entry) in resolved.iter_mut() {
if let Some(ref mut provider) = entry.auth_provider {
let config = cfg.auth_providers.get(&provider.name);
if config.is_none() {
tracing::debug!(
model_key = % key, provider = % provider.name,
"model references an undefined [auth_provider.*] table"
);
}
provider.attach_trusted_config(config);
}
}
{
let default_cw = DEFAULT_CONTEXT_WINDOW;
let donors: std::collections::HashMap<String, (std::num::NonZeroU64, ApiBackend)> =
@ -3582,6 +3742,10 @@ pub struct ConfigModelOverride {
pub api_key: Option<String>,
/// Env var name(s) for the provider key — string or array in config.toml.
pub env_key: Option<EnvKeys>,
/// Name of a `[auth_provider.<name>]` credential helper that mints
/// this model's bearer token. Static `api_key` / `env_key` win when both
/// are set.
pub auth_provider: Option<String>,
pub api_base_url: Option<String>,
pub max_completion_tokens: Option<u32>,
pub temperature: Option<f32>,
@ -3708,10 +3872,15 @@ impl ConfigModelOverride {
if self.env_key.is_some() {
entry.env_key.clone_from(&self.env_key);
}
if let Some(ref name) = self.auth_provider {
entry.auth_provider = Some(crate::auth::AuthProviderRef::unresolved(name.clone()));
}
if self.api_base_url.is_some() {
entry.api_base_url.clone_from(&self.api_base_url);
}
if self.supported_in_api.is_none() && (self.api_key.is_some() || self.env_key.is_some()) {
if self.supported_in_api.is_none()
&& (self.api_key.is_some() || self.env_key.is_some() || self.auth_provider.is_some())
{
entry.info.supported_in_api = true;
}
entry
@ -3897,6 +4066,11 @@ pub struct ModelEntry {
pub info: ModelInfo,
pub api_key: Option<String>,
pub env_key: Option<EnvKeys>,
/// Named credential helper (`[model.<id>] auth_provider = "<name>"`),
/// resolved against `[auth_provider.<name>]` by `resolve_model_list`.
/// Config-file models only: the built-in catalog never carries one.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub auth_provider: Option<crate::auth::AuthProviderRef>,
/// When set, `base_url` is used for session auth, `api_base_url` for API-key auth.
pub api_base_url: Option<String>,
}
@ -3909,6 +4083,7 @@ impl ModelEntry {
info,
api_key: None,
env_key: None,
auth_provider: None,
api_base_url: None,
}
}
@ -3920,19 +4095,31 @@ impl ModelEntry {
info: ModelInfo::from_config(entry),
api_key: entry.api_key.clone(),
env_key: entry.env_key.clone(),
auth_provider: None,
api_base_url: entry.api_base_url.clone(),
}
}
/// Non-empty `api_key`, else first non-empty resolved `env_key`.
/// `None` → fall through to session / global key.
/// `None` → fall through to session / global key. Static only: never
/// consults auth-provider tokens.
pub(crate) fn own_credential(&self) -> Option<String> {
first_own_credential(self.api_key.as_deref(), self.env_key.as_ref())
}
/// `true` when the model has a non-empty `api_key` or an `env_key` that
/// resolves to a non-empty value.
/// Probes `std::env::var` at call time — result is not stable across env changes.
/// The provider governing this model's bearer: `None` when a static
/// `api_key`/`env_key` resolves. The turn paths consult this, so a
/// shadowed provider never runs.
pub(crate) fn effective_auth_provider(&self) -> Option<&crate::auth::AuthProviderRef> {
if self.own_credential().is_some() {
return None;
}
self.auth_provider.as_ref()
}
/// `true` when the model has a non-empty `api_key`, an `env_key` that
/// resolves to a non-empty value, or a named auth provider.
/// Probes `std::env::var` at call time: result is not stable across env
/// changes. Never executes a provider command.
pub fn has_own_credentials(&self) -> bool {
self.own_credential().is_some()
self.own_credential().is_some() || self.auth_provider.is_some()
}
}
impl std::ops::Deref for ModelEntry {
@ -4306,10 +4493,8 @@ pub(crate) fn first_own_credential(
.map(str::to_owned)
.or_else(|| env_key.and_then(EnvKeys::resolve_value))
}
/// Resolve credentials for a model.
/// Priority: model api_key/env_key > session token > XAI_API_KEY.
///
/// When `env_key` lists multiple names, the first set non-empty value is used.
/// Priority: model api_key/env_key > cached auth-provider token > session
/// token > XAI_API_KEY.
pub fn resolve_credentials(model: &ModelEntry, session_key: Option<&str>) -> ResolvedCredentials {
let info = model.info();
let (api_key, base_url, auth_type) = if let Some(key) = model.own_credential() {
@ -4318,6 +4503,13 @@ pub fn resolve_credentials(model: &ModelEntry, session_key: Option<&str>) -> Res
info.base_url.clone(),
xai_chat_state::AuthType::ApiKey,
)
} else if let Some(provider) = model.auth_provider.as_ref() {
debug_assert!(model.effective_auth_provider().is_some());
(
provider.cached_token(),
info.base_url.clone(),
xai_chat_state::AuthType::ApiKey,
)
} else if let Some(key) = session_key {
(
Some(key.to_owned()),
@ -4425,23 +4617,37 @@ pub struct ModelAuthFacts {
pub byok: ModelByok,
pub auth_scheme: AuthScheme,
}
/// Resolve `model_id` to its auth facts from one effective-config load.
/// Load/parse failure → `byok = Unknown`; model absent from the catalog →
/// `NotByok`. An empty `model_id` (no sampling config yet) → `Unknown`, not
/// `NotByok`, so the gate isn't activated for an unidentified model.
pub fn resolve_model_auth_facts(model_id: &str) -> ModelAuthFacts {
/// Resolve `model_id` to its auth facts and auth-provider reference from one
/// effective-config load; both ride the same memo (see
/// `SessionActor::model_auth_memo`). Load/parse failure → `byok = Unknown`;
/// model absent from the catalog → `NotByok`. An empty `model_id` (no sampling
/// config yet) → `Unknown`, not `NotByok`, so the gate isn't activated for an
/// unidentified model.
pub fn resolve_model_auth_facts_and_provider(
model_id: &str,
) -> (ModelAuthFacts, Option<crate::auth::AuthProviderRef>) {
if model_id.is_empty() {
return ModelAuthFacts {
byok: ModelByok::Unknown,
auth_scheme: AuthScheme::default(),
};
return (
ModelAuthFacts {
byok: ModelByok::Unknown,
auth_scheme: AuthScheme::default(),
},
None,
);
}
with_resolved_model(model_id, |lookup| ModelAuthFacts {
byok: byok_from_lookup(&lookup),
auth_scheme: match lookup {
ModelLookup::Loaded(Some(e)) => e.info().auth_scheme,
_ => AuthScheme::default(),
},
with_resolved_model(model_id, |lookup| {
let facts = ModelAuthFacts {
byok: byok_from_lookup(&lookup),
auth_scheme: match lookup {
ModelLookup::Loaded(Some(e)) => e.info().auth_scheme,
_ => AuthScheme::default(),
},
};
let provider = match lookup {
ModelLookup::Loaded(Some(e)) => e.effective_auth_provider().cloned(),
_ => None,
};
(facts, provider)
})
}
fn byok_from_lookup(lookup: &ModelLookup) -> ModelByok {
@ -4502,6 +4708,13 @@ pub fn resolve_aux_model_sampling_config(
if sampler.api_key.is_some() {
return Some(sampler);
}
if entry.effective_auth_provider().is_some() {
tracing::warn!(
model = % model_id,
"aux model uses an auth provider with no cached token; the caller falls back to its session default"
);
return None;
}
}
let xai_bearer = session_key
.map(|s| s.to_owned())
@ -4545,6 +4758,7 @@ pub fn resolve_aux_model_sampling_config(
},
api_key: Some(bearer),
env_key: None,
auth_provider: None,
api_base_url: None,
};
let credentials = resolve_credentials_enforced(&entry, session_key, disable_api_key_auth);
@ -4564,18 +4778,14 @@ pub fn resolve_aux_model_sampling_config(
);
None
}
/// Finalize image-describe model + sampler config for user attachments.
/// Shared so the aux resolve happy path and the
/// `None` fallback cannot diverge between those entry points.
///
/// On aux resolve `Some`, stamp session-local fields (client id, attribution, bearer,
/// retries) onto the helper config. On `None`, fall back to the active session model and
/// full config (not forcing `image_description_model` onto the agent endpoint, which 404s
/// on BYOK / non-proxy routes for internal slugs like `grok-build`).
/// Stamp the session-local fields (client id, attribution, bearer resolver,
/// retries) from the active session onto a routed aux `SamplerConfig` so a
/// helper model keeps the session's auth/attribution. Shared by image-describe
/// and the auto-mode classifier so the two can't drift.
///
/// The resolver gate is host-based, stricter than `session_token_auth_gate`:
/// a session-token deployment on a custom `models_base_url` loses aux-sampler
/// refresh, rather than risk the session bearer on a third-party endpoint.
pub fn stamp_session_local_sampler_fields(
cfg: &mut SamplerConfig,
active_session_config: &SamplerConfig,
@ -4584,9 +4794,19 @@ pub fn stamp_session_local_sampler_fields(
) {
cfg.client_identifier = client_identifier;
cfg.attribution_callback = active_session_config.attribution_callback.clone();
cfg.bearer_resolver = active_session_config.bearer_resolver.clone();
if crate::util::is_xai_api_bearer_url(&cfg.base_url) {
cfg.bearer_resolver = active_session_config.bearer_resolver.clone();
}
cfg.max_retries = max_retries;
}
/// Finalize image-describe model + sampler config for user attachments.
/// Shared so the aux resolve happy path and the `None` fallback cannot
/// diverge between those entry points.
///
/// On aux resolve `Some`, stamp session-local fields onto the helper config.
/// On `None`, fall back to the active session model and full config (not
/// forcing `image_description_model` onto the agent endpoint, which 404s on
/// BYOK / non-proxy routes for internal slugs like `grok-build`).
pub fn finalize_image_describe_sampler_config(
resolved_aux: Option<SamplerConfig>,
active_session_config: &SamplerConfig,
@ -4768,6 +4988,7 @@ fn resolve_hidden_default_web_search_sampling_config(
},
api_key: None,
env_key: None,
auth_provider: None,
api_base_url: None,
};
let credentials = resolve_credentials_enforced(&entry, session_key, disable_api_key_auth);
@ -4791,6 +5012,13 @@ pub fn resolve_web_search_sampling_config(
) -> Option<SamplerConfig> {
let resolved = if let Some(entry) = find_model_by_id(models, model_id).cloned() {
let credentials = resolve_credentials_enforced(&entry, session_key, disable_api_key_auth);
if credentials.api_key.is_none() && entry.effective_auth_provider().is_some() {
tracing::warn!(
web_search_model = % model_id,
"web search model uses an auth provider with no cached token; disabling web search"
);
return None;
}
Some(sampling_config_for_model(
&entry,
credentials,
@ -5330,6 +5558,256 @@ reasoning_effort = "low"
assert_eq!(resolved.base_url, "https://vendor.example/v1");
assert_eq!(resolved.api_key.as_deref(), Some("vendor-key"));
}
/// Cold cache falls back to the session model, never the xAI proxy;
/// warm cache serves the provider token at the provider endpoint.
#[tokio::test]
async fn aux_model_with_auth_provider_never_reroutes() {
let endpoints = EndpointsConfig::default();
let provider = crate::auth::AuthProviderRef::new(
"aux-provider-test".into(),
crate::auth::AuthProviderConfig {
command: "printf aux-token".into(),
args: None,
token_ttl_secs: Some(3600),
timeout_secs: None,
},
);
let mut entry = test_model_entry("m", "https://litellm.example/v1", None, None, None);
entry.auth_provider = Some(provider.clone());
let mut catalog = IndexMap::new();
catalog.insert("proxied-aux".to_string(), entry);
assert!(
resolve_aux_model_sampling_config(
"proxied-aux",
&catalog,
&endpoints,
Some("session-jwt"),
false,
None,
None,
)
.is_none(),
"cold provider cache must not reroute the aux model through the xAI proxy"
);
let _ = provider.ensure_fresh_token(None).await;
let resolved = resolve_aux_model_sampling_config(
"proxied-aux",
&catalog,
&endpoints,
Some("session-jwt"),
false,
None,
None,
)
.expect("warm cache resolves");
assert_eq!(resolved.base_url, "https://litellm.example/v1");
assert_eq!(resolved.api_key.as_deref(), Some("aux-token"));
}
/// The session bearer resolver must never be stamped onto a third-party
/// sampler: the sampler substitutes the resolver's bearer at request
/// time.
#[test]
fn session_resolver_is_not_stamped_onto_third_party_samplers() {
#[derive(Debug)]
struct SessionResolver;
impl xai_grok_sampler::BearerResolver for SessionResolver {
fn current_bearer(&self) -> Option<String> {
Some("session-jwt".into())
}
}
let session_cfg = SamplerConfig {
bearer_resolver: Some(std::sync::Arc::new(SessionResolver)),
..SamplerConfig::default()
};
let mut third_party = SamplerConfig {
base_url: "https://litellm.corp.example/v1".into(),
..SamplerConfig::default()
};
stamp_session_local_sampler_fields(&mut third_party, &session_cfg, None, None);
assert!(
third_party.bearer_resolver.is_none(),
"a third-party endpoint must keep its resolved credential"
);
let mut first_party = SamplerConfig {
base_url: EndpointsConfig::default().resolve_inference_base_url(),
..SamplerConfig::default()
};
stamp_session_local_sampler_fields(&mut first_party, &session_cfg, None, None);
assert!(
first_party.bearer_resolver.is_some(),
"first-party aux samplers keep the session refresh behavior"
);
}
/// A cold cache disables web search rather than sending an
/// unauthenticated request.
#[tokio::test]
async fn web_search_with_auth_provider_requires_warm_cache() {
let endpoints = EndpointsConfig::default();
let provider = crate::auth::AuthProviderRef::new(
"web-search-provider-test".into(),
crate::auth::AuthProviderConfig {
command: "printf ws-token".into(),
args: None,
token_ttl_secs: Some(3600),
timeout_secs: None,
},
);
let mut entry = test_model_entry("m", "https://litellm.example/v1", None, None, None);
entry.auth_provider = Some(provider.clone());
let mut catalog = IndexMap::new();
catalog.insert("proxied-search".to_string(), entry);
assert!(
resolve_web_search_sampling_config(
"proxied-search",
&catalog,
Some("session-jwt"),
false,
None,
None,
&endpoints,
)
.is_none(),
"a cold provider cache must disable web search, not send an unauthenticated request"
);
let _ = provider.ensure_fresh_token(None).await;
let resolved = resolve_web_search_sampling_config(
"proxied-search",
&catalog,
Some("session-jwt"),
false,
None,
None,
&endpoints,
)
.expect("warm cache resolves");
assert_eq!(resolved.api_key.as_deref(), Some("ws-token"));
}
/// The lenient parser warns per problem and never fails the whole
/// config.
#[test]
fn auth_provider_parse_warnings_are_lenient_and_specific() {
use super::super::config_model_override_parse::{ConfigWarningKind, WarningTarget};
let raw_config: toml::Value = toml::from_str(
r#"
[auth_provider.good]
command = "printf ok"
[auth_provider.bad-type]
command = "printf x"
token_ttl_secs = "not-a-number"
[auth_provider.typo]
command = "printf y"
timeout_seconds = 5
[auth_provider.commandless]
token_ttl_secs = 60
[auth_provider.short-ttl]
command = "printf x"
token_ttl_secs = 60
[auth_provider.zero-timeout]
command = "printf x"
timeout_secs = 0
[auth_provider.slow]
command = "printf x"
timeout_secs = 601
[model.orphaned]
model = "m"
base_url = "https://x.example/v1"
context_window = 200000
auth_provider = "does-not-exist"
"#,
)
.unwrap();
let cfg =
Config::new_from_toml_cfg(&raw_config).expect("one bad table must not fail the config");
assert!(cfg.auth_providers.contains_key("good"));
assert!(
!cfg.auth_providers.contains_key("bad-type"),
"malformed entry is skipped (fails closed)"
);
let has_provider = |name: &str, field: Option<&str>, kind: ConfigWarningKind| {
cfg.config_warnings.iter().any(|w| {
w.kind == kind
&& matches!(
& w.target, WarningTarget::AuthProvider { name : n, field : f
}
if n == name && f.as_deref() == field
)
})
};
assert!(has_provider(
"bad-type",
None,
ConfigWarningKind::InvalidValue
));
assert!(has_provider(
"typo",
Some("timeout_seconds"),
ConfigWarningKind::UnknownField
));
assert!(has_provider(
"commandless",
Some("command"),
ConfigWarningKind::InvalidValue
));
assert!(has_provider(
"short-ttl",
Some("token_ttl_secs"),
ConfigWarningKind::InvalidValue
));
assert!(has_provider(
"zero-timeout",
Some("timeout_secs"),
ConfigWarningKind::InvalidValue
));
assert!(has_provider(
"slow",
Some("timeout_secs"),
ConfigWarningKind::InvalidValue
));
let provider_reason = |name: &str| {
cfg.config_warnings
.iter()
.find(|w| {
matches!(
& w.target, WarningTarget::AuthProvider { name : n, field : f }
if n == name && f.as_deref() == Some("timeout_secs")
)
})
.map(|w| w.reason.as_str())
.unwrap_or_default()
.to_owned()
};
assert!(provider_reason("zero-timeout").contains("clamped to 1"));
assert!(provider_reason("slow").contains("clamped to 600"));
assert!(
cfg.config_warnings.iter().any(|w| {
w.kind == ConfigWarningKind::InvalidValue
&& matches!(& w.target, WarningTarget::Model
{ field, .. }
if field.as_deref() == Some("auth_provider"))
}),
"undefined reference warns at parse time: {:?}",
cfg.config_warnings
);
let raw_config: toml::Value = toml::from_str(r#"auth_provider = "oops""#).unwrap();
let cfg = Config::new_from_toml_cfg(&raw_config)
.expect("a non-table auth_provider must not fail the config");
assert!(cfg.auth_providers.is_empty());
assert!(
cfg.config_warnings.iter().any(|w| {
matches!(w.target, WarningTarget::AuthProviderSection)
&& w.kind == ConfigWarningKind::NotATable
}),
"non-table section warns: {:?}",
cfg.config_warnings
);
}
#[test]
fn web_search_disable_api_key_auth_swaps_first_party_key_for_session() {
let endpoints = EndpointsConfig::default();
@ -5379,6 +5857,199 @@ reasoning_effort = "low"
assert_eq!(model.info.base_url, "https://api.example.com/v1");
assert_eq!(model.api_key, Some("sk-test-key-12345".to_string()));
}
#[test]
fn parses_auth_provider_tables_and_model_reference() {
let raw_config: toml::Value = toml::from_str(
r#"
[auth_provider.litellm]
command = "/usr/local/bin/litellm-token"
args = ["--scope", "corp"]
token_ttl_secs = 3600
timeout_secs = 10
[model.proxied-claude]
model = "claude-sonnet-4-5"
base_url = "https://litellm.corp.example/v1"
context_window = 200000
auth_provider = "litellm"
"#,
)
.unwrap();
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
assert_eq!(
cfg.auth_providers.get("litellm"),
Some(&crate::auth::AuthProviderConfig {
command: "/usr/local/bin/litellm-token".into(),
args: Some(vec!["--scope".into(), "corp".into()]),
token_ttl_secs: Some(3600),
timeout_secs: Some(10),
})
);
let resolved = resolve_model_list(&cfg, None);
let model = resolved.get("proxied-claude").expect("model should exist");
let provider = model
.auth_provider
.as_ref()
.expect("model should reference the provider");
assert_eq!(provider.name, "litellm");
assert_eq!(provider.config.command, "/usr/local/bin/litellm-token");
assert_eq!(provider.config.token_ttl_secs, Some(3600));
assert!(
model.has_own_credentials(),
"provider-backed models classify as BYOK (session token must not leak)"
);
assert!(
model.info.supported_in_api,
"declaring an auth provider implies supported_in_api"
);
}
/// A static key shadows a fully defined provider through the real
/// `resolve_model_list` + `attach_trusted_config` pipeline (not a
/// hand-built ref): the static key wins even with the provider cache warm.
#[tokio::test]
async fn static_key_shadows_defined_provider_through_pipeline() {
let raw_config: toml::Value = toml::from_str(
r#"
[auth_provider.understudy]
command = "printf provider-token"
token_ttl_secs = 3600
[model.dual-auth]
model = "m"
base_url = "https://switchboard.example/v1"
context_window = 200000
api_key = "sk-house-key"
auth_provider = "understudy"
"#,
)
.unwrap();
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
let resolved = resolve_model_list(&cfg, None);
let model = resolved.get("dual-auth").expect("model should exist");
assert_eq!(
model.effective_auth_provider().map(|p| p.name.as_str()),
None,
"a static key shadows the provider after real resolution"
);
let provider = model.auth_provider.as_ref().unwrap().clone();
let _ = provider.ensure_fresh_token(None).await;
let creds = resolve_credentials(model, Some("session-jwt"));
assert_eq!(creds.api_key.as_deref(), Some("sk-house-key"));
assert_eq!(creds.auth_type, xai_chat_state::AuthType::ApiKey);
assert_eq!(creds.base_url, "https://switchboard.example/v1");
}
#[test]
fn undefined_auth_provider_fails_closed() {
let raw_config: toml::Value = toml::from_str(
r#"
[model.orphan]
model = "m"
base_url = "https://third-party.example/v1"
context_window = 200000
auth_provider = "nope"
"#,
)
.unwrap();
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
let resolved = resolve_model_list(&cfg, None);
let model = resolved.get("orphan").expect("model should exist");
let provider = model.auth_provider.as_ref().unwrap();
assert_eq!(provider.name, "nope");
assert!(
provider.config.command.is_empty(),
"undefined provider keeps an empty command"
);
assert!(model.has_own_credentials());
let creds = resolve_credentials(model, Some("session-jwt"));
assert_eq!(creds.api_key, None);
}
#[tokio::test]
async fn resolve_credentials_serves_cached_provider_token() {
use xai_chat_state::AuthType;
let mut model = test_model_entry("m", "https://litellm.example/v1", None, None, None);
let provider = crate::auth::AuthProviderRef::new(
"resolve-creds-test".into(),
crate::auth::AuthProviderConfig {
command: "printf provider-minted-token".into(),
args: None,
token_ttl_secs: Some(3600),
timeout_secs: None,
},
);
model.auth_provider = Some(provider.clone());
let creds = resolve_credentials(&model, Some("session-jwt"));
assert_eq!(creds.api_key, None, "cold cache must not run the command");
let _ = provider.ensure_fresh_token(None).await;
let creds = resolve_credentials(&model, Some("session-jwt"));
assert_eq!(creds.api_key.as_deref(), Some("provider-minted-token"));
assert_eq!(creds.auth_type, AuthType::ApiKey);
assert_eq!(creds.base_url, "https://litellm.example/v1");
}
/// A set `env_key` shadows even a warm provider cache at resolve time, so
/// the static credential wins on the wire and the provider never governs.
#[tokio::test]
async fn set_env_key_shadows_warm_provider_at_resolve_time() {
use xai_grok_test_support::EnvGuard;
let var = "GROK_TEST_ENVKEY_SHADOW";
let _guard = EnvGuard::set(var, "env-token");
let mut model = test_model_entry("m", "https://litellm.example/v1", None, Some(var), None);
let provider = crate::auth::AuthProviderRef::new(
"env-shadow-test".into(),
crate::auth::AuthProviderConfig {
command: "printf provider-token".into(),
args: None,
token_ttl_secs: Some(3600),
timeout_secs: None,
},
);
model.auth_provider = Some(provider.clone());
let _ = provider.ensure_fresh_token(None).await;
assert_eq!(
model.effective_auth_provider().map(|p| p.name.as_str()),
None,
"a resolvable env_key shadows the provider"
);
let creds = resolve_credentials(&model, Some("session-jwt"));
assert_eq!(
creds.api_key.as_deref(),
Some("env-token"),
"a set env_key must win over a warm provider cache"
);
}
/// A catalog deserialized from bytes cannot smuggle a runnable command.
#[test]
fn prefetched_entry_provider_config_comes_from_trusted_tables_only() {
let mut entry = test_model_entry("m", "https://cache.example/v1", None, None, None);
let smuggled: crate::auth::AuthProviderRef = serde_json::from_str(
r#"{"name": "cache-smuggle-test", "config": {"command": "evil"}}"#,
)
.unwrap();
entry.auth_provider = Some(smuggled);
let mut prefetched = IndexMap::new();
prefetched.insert("cached-model".to_string(), entry);
let cfg = Config::default();
let resolved = resolve_model_list(&cfg, Some(prefetched.clone()));
let provider = resolved["cached-model"].auth_provider.as_ref().unwrap();
assert_eq!(
resolve_credentials(&resolved["cached-model"], Some("session-jwt")).api_key,
None,
"an unusable provider fails closed"
);
assert_eq!(provider.config, crate::auth::AuthProviderConfig::default());
let mut cfg = Config::default();
cfg.auth_providers.insert(
"cache-smuggle-test".to_string(),
crate::auth::AuthProviderConfig {
command: "printf local".to_string(),
args: None,
token_ttl_secs: None,
timeout_secs: None,
},
);
let resolved = resolve_model_list(&cfg, Some(prefetched));
let provider = resolved["cached-model"].auth_provider.as_ref().unwrap();
assert_eq!(provider.config.command, "printf local");
}
fn test_model_entry(
model: &str,
base_url: &str,
@ -5421,6 +6092,7 @@ reasoning_effort = "low"
},
api_key: api_key.map(|s| s.to_string()),
env_key: env_key.map(EnvKeys::single),
auth_provider: None,
api_base_url: api_base_url.map(|s| s.to_string()),
}
}
@ -5957,7 +6629,10 @@ reasoning_effort = "low"
}
#[test]
fn resolve_model_auth_facts_empty_model_id_is_unknown() {
assert_eq!(resolve_model_auth_facts("").byok, ModelByok::Unknown);
assert_eq!(
resolve_model_auth_facts_and_provider("").0.byok,
ModelByok::Unknown
);
}
#[test]
fn user_override_adds_api_key_to_default_model() {
@ -10608,6 +11283,7 @@ default = "grok-4.5"
},
api_key: None,
env_key: None,
auth_provider: None,
api_base_url: None,
}
}

397
crates/codegen/xai-grok-shell/src/agent/config_model_override_parse.rs
View file

@ -1,5 +1,8 @@
//! Resilient parsing for `[model.<id>]` TOML overrides.
//!
//! It also defines [`ConfigWarning`] and [`WarningTarget`], the shared warning
//! vocabulary; the `[auth_provider.*]` parser in `config.rs` emits them too.
//!
//! A model entry must survive a bad field: warn and skip the field, never
//! drop the model (managed configs must not lose catalog entries).
//!
@ -9,7 +12,7 @@
//! fail to parse on their own are pruned (one warning each) and the table is
//! parsed again. Non-table values are dropped with a warning.
//!
//! Warnings are retained on `Config::model_override_warnings` and surfaced by
//! Warnings are retained on `Config::config_warnings` and surfaced by
//! `grok inspect`.
use indexmap::IndexMap;
@ -17,10 +20,10 @@ use serde::Serialize;
use super::config::ConfigModelOverride;
/// Category for a [`ModelOverrideWarning`].
/// Category for a [`ConfigWarning`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize)]
#[serde(rename_all = "kebab-case")]
pub enum ModelOverrideWarningKind {
pub enum ConfigWarningKind {
/// Field name not recognized; field ignored.
UnknownField,
/// Value failed to parse; field skipped.
@ -29,29 +32,127 @@ pub enum ModelOverrideWarningKind {
DuplicateAlias,
/// Entry value is not a TOML table; entry dropped.
NotATable,
/// Fields are individually valid but conflict (e.g. `auth_provider`
/// shadowed by `api_key`/`env_key`); all fields kept, one is inert.
ConflictingFields,
/// Entry failed to parse even after skipping invalid fields; the model
/// keeps an empty override.
UnparseableEntry,
}
/// One skipped field or dropped entry from `[model.*]` parsing.
/// What a [`ConfigWarning`] is about. Serialize-only: `grok inspect --json`
/// emits it, nothing deserializes it back.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)]
#[serde(tag = "target", rename_all = "camelCase")]
pub enum WarningTarget {
/// The `[model]` section as a whole (e.g. not a table).
ModelSection,
/// A `[model.<key>]` entry; `field` names a key when the warning is
/// field-specific.
Model {
key: String,
#[serde(skip_serializing_if = "Option::is_none")]
field: Option<String>,
},
/// The `[auth_provider]` section as a whole.
AuthProviderSection,
/// An `[auth_provider.<name>]` table; `field` names a key when the
/// warning is field-specific.
AuthProvider {
name: String,
#[serde(skip_serializing_if = "Option::is_none")]
field: Option<String>,
},
}
impl WarningTarget {
/// The config path, e.g. `model."grok-4.5"` or `auth_provider."litellm"`.
pub(crate) fn label(&self) -> String {
match self {
Self::ModelSection => "model".to_owned(),
Self::Model { key, .. } => format!("model.\"{key}\""),
Self::AuthProviderSection => "auth_provider".to_owned(),
Self::AuthProvider { name, .. } => format!("auth_provider.\"{name}\""),
}
}
pub(crate) fn field(&self) -> Option<&str> {
match self {
Self::Model { field, .. } | Self::AuthProvider { field, .. } => field.as_deref(),
Self::ModelSection | Self::AuthProviderSection => None,
}
}
}
/// One skipped field or dropped entry from config parsing.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ModelOverrideWarning {
/// `None` when the warning is about the `[model]` section itself.
#[serde(skip_serializing_if = "Option::is_none")]
pub model_key: Option<String>,
/// `None` for warnings about the entry as a whole.
#[serde(skip_serializing_if = "Option::is_none")]
pub field: Option<String>,
pub kind: ModelOverrideWarningKind,
pub struct ConfigWarning {
#[serde(flatten)]
pub target: WarningTarget,
pub kind: ConfigWarningKind,
pub reason: String,
}
/// Result of [`parse_model_overrides`].
impl ConfigWarning {
pub(crate) fn model(
key: &str,
field: Option<&str>,
kind: ConfigWarningKind,
reason: String,
) -> Self {
let target = WarningTarget::Model {
key: key.to_owned(),
field: field.map(str::to_owned),
};
Self {
target,
kind,
reason,
}
}
pub(crate) fn model_section(kind: ConfigWarningKind, reason: String) -> Self {
Self {
target: WarningTarget::ModelSection,
kind,
reason,
}
}
pub(crate) fn auth_provider(
name: &str,
field: Option<&str>,
kind: ConfigWarningKind,
reason: String,
) -> Self {
let target = WarningTarget::AuthProvider {
name: name.to_owned(),
field: field.map(str::to_owned),
};
Self {
target,
kind,
reason,
}
}
pub(crate) fn auth_provider_section(kind: ConfigWarningKind, reason: String) -> Self {
Self {
target: WarningTarget::AuthProviderSection,
kind,
reason,
}
}
pub(crate) fn field(&self) -> Option<&str> {
self.target.field()
}
}
pub(crate) struct ParsedModelOverrides {
pub models: IndexMap<String, ConfigModelOverride>,
pub warnings: Vec<ModelOverrideWarning>,
pub warnings: Vec<ConfigWarning>,
}
/// Parses every `[model.<id>]` entry in `raw_config`, returning the overrides
@ -63,28 +164,26 @@ pub(crate) fn parse_model_overrides(raw_config: &toml::Value) -> ParsedModelOver
return ParsedModelOverrides { models, warnings };
};
let Some(table) = section.as_table() else {
warnings.push(ModelOverrideWarning {
model_key: None,
field: None,
kind: ModelOverrideWarningKind::NotATable,
reason: format!(
warnings.push(ConfigWarning::model_section(
ConfigWarningKind::NotATable,
format!(
"`model` must be a table of [model.<id>] entries, got {}; all model overrides ignored",
section.type_str()
),
});
));
return ParsedModelOverrides { models, warnings };
};
for (model_key, value) in table {
let Some(entry_table) = value.as_table() else {
warnings.push(ModelOverrideWarning {
model_key: Some(model_key.clone()),
field: None,
kind: ModelOverrideWarningKind::NotATable,
reason: format!(
warnings.push(ConfigWarning::model(
model_key,
None,
ConfigWarningKind::NotATable,
format!(
"expected a table like [model.\"{model_key}\"], got {}; entry dropped",
value.type_str()
),
});
));
continue;
};
let (entry, entry_warnings) = parse_model_override_table(model_key, entry_table.clone());
@ -96,7 +195,7 @@ pub(crate) fn parse_model_overrides(raw_config: &toml::Value) -> ParsedModelOver
/// Logs the warnings when they differ from the previous parse, so a
/// persistently broken config logs once per process instead of once per parse.
pub(crate) fn log_model_override_warnings(warnings: &[ModelOverrideWarning]) {
pub(crate) fn log_config_warnings(warnings: &[ConfigWarning]) {
use std::hash::{Hash as _, Hasher as _};
use std::sync::atomic::{AtomicU64, Ordering};
@ -115,8 +214,8 @@ pub(crate) fn log_model_override_warnings(warnings: &[ModelOverrideWarning]) {
for warning in warnings {
tracing::warn!(
model = warning.model_key.as_deref().unwrap_or("(section)"),
field = warning.field.as_deref().unwrap_or("(entry)"),
path = %warning.target.label(),
field = warning.field().unwrap_or("(entry)"),
kind = ?warning.kind,
reason = %warning.reason,
"model_override: skipped invalid config"
@ -133,13 +232,13 @@ pub(crate) fn log_model_override_warnings(warnings: &[ModelOverrideWarning]) {
fn parse_model_override_table(
model_key: &str,
mut table: toml::map::Map<String, toml::Value>,
) -> (ConfigModelOverride, Vec<ModelOverrideWarning>) {
) -> (ConfigModelOverride, Vec<ConfigWarning>) {
let mut warnings = Vec::new();
dedupe_aliases(model_key, &mut table, &mut warnings);
// Unknown-field warnings come from whichever parse produces the returned
// entry, so both paths report them identically.
match deserialize_with_unknown_fields(table.clone()) {
let (entry, mut warnings) = match deserialize_with_unknown_fields(table.clone()) {
Ok((entry, unknown)) => {
warnings.extend(unknown_field_warnings(model_key, unknown));
(entry, warnings)
@ -155,19 +254,57 @@ fn parse_model_override_table(
// Reachable only when fields conflict jointly, e.g. an
// alias pair missing from `ALIASES`. Keep the model
// rather than dropping it.
warnings.push(ModelOverrideWarning {
model_key: Some(model_key.to_owned()),
field: None,
kind: ModelOverrideWarningKind::UnparseableEntry,
reason: format!(
warnings.push(ConfigWarning::model(
model_key,
None,
ConfigWarningKind::UnparseableEntry,
format!(
"failed to parse after skipping invalid fields ({error}); using empty override"
),
});
));
(ConfigModelOverride::default(), warnings)
}
}
}
};
if entry.auth_provider.is_some() {
// A non-empty `api_key` always shadows; an `env_key` only shadows when
// its variable resolves at runtime, which parse time can't know. Warn
// accordingly so the message matches what actually happens.
let has_static_api_key = entry
.api_key
.as_deref()
.map(str::trim)
.is_some_and(|k| !k.is_empty());
if has_static_api_key {
warnings.push(ConfigWarning::model(
model_key,
Some("auth_provider"),
ConfigWarningKind::ConflictingFields,
"auth_provider is shadowed by api_key on this model; the static \
key always takes precedence, so the provider never runs"
.to_owned(),
));
} else if entry
.env_key
.as_ref()
.and_then(crate::agent::config::EnvKeys::primary)
.is_some()
{
warnings.push(ConfigWarning::model(
model_key,
Some("auth_provider"),
ConfigWarningKind::ConflictingFields,
"auth_provider may be shadowed by env_key on this model; env_key \
takes precedence when its variable resolves to a value, \
otherwise the provider runs"
.to_owned(),
));
}
}
(entry, warnings)
}
/// `(canonical, legacy)` key pairs that serde rejects as duplicate fields
@ -181,7 +318,7 @@ const ALIASES: &[(&str, &str)] = &[("compactions_remaining", "send_compactions_r
fn dedupe_aliases(
model_key: &str,
table: &mut toml::map::Map<String, toml::Value>,
warnings: &mut Vec<ModelOverrideWarning>,
warnings: &mut Vec<ConfigWarning>,
) {
for &(canonical, legacy) in ALIASES {
if !(table.contains_key(canonical) && table.contains_key(legacy)) {
@ -190,21 +327,21 @@ fn dedupe_aliases(
match field_parse_error(canonical, &table[canonical]) {
None => {
table.remove(legacy);
warnings.push(ModelOverrideWarning {
model_key: Some(model_key.to_owned()),
field: Some(legacy.to_owned()),
kind: ModelOverrideWarningKind::DuplicateAlias,
reason: format!("legacy alias of {canonical}; skipped in favor of {canonical}"),
});
warnings.push(ConfigWarning::model(
model_key,
Some(legacy),
ConfigWarningKind::DuplicateAlias,
format!("legacy alias of {canonical}; skipped in favor of {canonical}"),
));
}
Some(error) => {
table.remove(canonical);
warnings.push(ModelOverrideWarning {
model_key: Some(model_key.to_owned()),
field: Some(canonical.to_owned()),
kind: ModelOverrideWarningKind::InvalidValue,
reason: format!("{error}; skipped in favor of {legacy}"),
});
warnings.push(ConfigWarning::model(
model_key,
Some(canonical),
ConfigWarningKind::InvalidValue,
format!("{error}; skipped in favor of {legacy}"),
));
}
}
}
@ -222,14 +359,16 @@ fn deserialize_with_unknown_fields(
Ok((entry, unknown))
}
fn unknown_field_warnings(model_key: &str, unknown: Vec<String>) -> Vec<ModelOverrideWarning> {
fn unknown_field_warnings(model_key: &str, unknown: Vec<String>) -> Vec<ConfigWarning> {
unknown
.into_iter()
.map(|field| ModelOverrideWarning {
model_key: Some(model_key.to_owned()),
field: Some(field),
kind: ModelOverrideWarningKind::UnknownField,
reason: "unknown field".to_owned(),
.map(|field| {
ConfigWarning::model(
model_key,
Some(field.as_str()),
ConfigWarningKind::UnknownField,
"unknown field".to_owned(),
)
})
.collect()
}
@ -239,17 +378,17 @@ fn unknown_field_warnings(model_key: &str, unknown: Vec<String>) -> Vec<ModelOve
fn prune_invalid_fields(
model_key: &str,
table: &mut toml::map::Map<String, toml::Value>,
warnings: &mut Vec<ModelOverrideWarning>,
warnings: &mut Vec<ConfigWarning>,
) {
table.retain(|field, value| match field_parse_error(field, value) {
None => true,
Some(error) => {
warnings.push(ModelOverrideWarning {
model_key: Some(model_key.to_owned()),
field: Some(field.to_owned()),
kind: ModelOverrideWarningKind::InvalidValue,
reason: error.to_string(),
});
warnings.push(ConfigWarning::model(
model_key,
Some(field),
ConfigWarningKind::InvalidValue,
error.to_string(),
));
false
}
});
@ -277,12 +416,7 @@ mod tests {
crate::agent::config::Config::new_from_toml_cfg(&raw).expect("config should parse")
}
fn parse_raw(
toml_str: &str,
) -> (
IndexMap<String, ConfigModelOverride>,
Vec<ModelOverrideWarning>,
) {
fn parse_raw(toml_str: &str) -> (IndexMap<String, ConfigModelOverride>, Vec<ConfigWarning>) {
let raw: toml::Value = toml::from_str(toml_str).unwrap();
let ParsedModelOverrides { models, warnings } = parse_model_overrides(&raw);
(models, warnings)
@ -307,9 +441,9 @@ mod tests {
model.compactions_remaining,
Some(CompactionsRemaining::Fixed(1))
);
assert!(cfg.model_override_warnings.iter().any(|w| {
w.kind == ModelOverrideWarningKind::DuplicateAlias
&& w.field.as_deref() == Some("send_compactions_remaining")
assert!(cfg.config_warnings.iter().any(|w| {
w.kind == ConfigWarningKind::DuplicateAlias
&& w.field() == Some("send_compactions_remaining")
}));
let resolved = crate::agent::config::resolve_model_list(&cfg, None);
assert!(resolved.contains_key("grok-4.5"));
@ -329,7 +463,7 @@ mod tests {
model.compactions_remaining,
Some(CompactionsRemaining::Fixed(2))
);
assert!(cfg.model_override_warnings.is_empty());
assert!(cfg.config_warnings.is_empty());
}
#[test]
@ -348,9 +482,8 @@ mod tests {
.expect("grok-4.5 must remain in catalog");
assert_eq!(model.model.as_deref(), Some("grok-4.5"));
assert!(model.reasoning_effort.is_none());
assert!(cfg.model_override_warnings.iter().any(|w| {
w.kind == ModelOverrideWarningKind::InvalidValue
&& w.field.as_deref() == Some("reasoning_effort")
assert!(cfg.config_warnings.iter().any(|w| {
w.kind == ConfigWarningKind::InvalidValue && w.field() == Some("reasoning_effort")
}));
}
@ -372,12 +505,12 @@ mod tests {
);
assert_eq!(
warnings,
vec![ModelOverrideWarning {
model_key: Some("grok-4.5".to_owned()),
field: Some("future_field".to_owned()),
kind: ModelOverrideWarningKind::UnknownField,
reason: "unknown field".to_owned(),
}]
vec![ConfigWarning::model(
"grok-4.5",
Some("future_field"),
ConfigWarningKind::UnknownField,
"unknown field".to_owned(),
)]
);
}
@ -389,7 +522,7 @@ mod tests {
let (_, warnings) = parse_raw(toml_str);
warnings
.into_iter()
.filter(|w| w.kind == ModelOverrideWarningKind::UnknownField)
.filter(|w| w.kind == ConfigWarningKind::UnknownField)
.collect::<Vec<_>>()
};
let fast = unknown_of(
@ -407,7 +540,7 @@ mod tests {
);
assert_eq!(fast, slow);
assert_eq!(fast.len(), 1);
assert_eq!(fast[0].field.as_deref(), Some("temprature"));
assert_eq!(fast[0].field(), Some("temprature"));
}
#[test]
@ -428,8 +561,7 @@ mod tests {
);
assert!(entry.temperature.is_none());
assert!(warnings.iter().any(|w| {
w.kind == ModelOverrideWarningKind::InvalidValue
&& w.field.as_deref() == Some("temperature")
w.kind == ConfigWarningKind::InvalidValue && w.field() == Some("temperature")
}));
// All fields invalid: the model stays, with an empty override.
@ -447,7 +579,7 @@ mod tests {
assert!(
warnings
.iter()
.all(|w| w.kind == ModelOverrideWarningKind::InvalidValue)
.all(|w| w.kind == ConfigWarningKind::InvalidValue)
);
}
@ -466,8 +598,8 @@ mod tests {
Some(CompactionsRemaining::Fixed(2))
);
assert_eq!(warnings.len(), 1);
assert_eq!(warnings[0].kind, ModelOverrideWarningKind::InvalidValue);
assert_eq!(warnings[0].field.as_deref(), Some("compactions_remaining"));
assert_eq!(warnings[0].kind, ConfigWarningKind::InvalidValue);
assert_eq!(warnings[0].field(), Some("compactions_remaining"));
}
#[test]
@ -475,9 +607,8 @@ mod tests {
let (models, warnings) = parse_raw(r#"model = "grok-4""#);
assert!(models.is_empty());
assert_eq!(warnings.len(), 1);
assert_eq!(warnings[0].kind, ModelOverrideWarningKind::NotATable);
assert_eq!(warnings[0].model_key, None);
assert_eq!(warnings[0].field, None);
assert_eq!(warnings[0].kind, ConfigWarningKind::NotATable);
assert!(matches!(warnings[0].target, WarningTarget::ModelSection));
}
#[test]
@ -490,9 +621,12 @@ mod tests {
);
assert!(models.is_empty(), "a scalar cannot define a model");
assert_eq!(warnings.len(), 1);
assert_eq!(warnings[0].kind, ModelOverrideWarningKind::NotATable);
assert_eq!(warnings[0].model_key.as_deref(), Some("oops"));
assert_eq!(warnings[0].field, None);
assert_eq!(warnings[0].kind, ConfigWarningKind::NotATable);
assert!(matches!(
&warnings[0].target,
WarningTarget::Model { key, field: None }
if key == "oops"
));
}
/// Exhaustive literal (no `..`): a new struct field is a compile error
@ -505,6 +639,7 @@ mod tests {
description: Some("desc".into()),
api_key: Some("key".into()),
env_key: Some(crate::agent::config::EnvKeys::single("ENV_KEY")),
auth_provider: Some("corp-gateway".into()),
api_base_url: Some("https://api.example.com".into()),
max_completion_tokens: Some(1024),
temperature: Some(0.5),
@ -541,10 +676,7 @@ mod tests {
fn parse_single_entry(
entry: toml::map::Map<String, toml::Value>,
) -> (
IndexMap<String, ConfigModelOverride>,
Vec<ModelOverrideWarning>,
) {
) -> (IndexMap<String, ConfigModelOverride>, Vec<ConfigWarning>) {
let mut model_table = toml::map::Map::new();
model_table.insert("m".to_owned(), toml::Value::Table(entry));
let mut root = toml::map::Map::new();
@ -555,14 +687,73 @@ mod tests {
}
#[test]
fn fully_populated_override_round_trips_without_warnings() {
fn fully_populated_override_round_trips_with_only_the_shadowing_warning() {
let serialized = toml::Value::try_from(fully_populated_override()).unwrap();
let (models, warnings) = parse_single_entry(serialized.as_table().unwrap().clone());
assert_eq!(warnings, Vec::new(), "no field may be skipped or unknown");
// The exhaustive literal deliberately sets `api_key`, `env_key`, AND
// `auth_provider`: the one legal-but-warned combination. Any other
// warning (skipped/unknown field) still fails the guard.
let unexpected: Vec<_> = warnings
.iter()
.filter(|w| w.kind != ConfigWarningKind::ConflictingFields)
.collect();
assert_eq!(unexpected, Vec::<&ConfigWarning>::new());
assert_eq!(warnings.len(), 1);
let reparsed = toml::Value::try_from(models.get("m").unwrap()).unwrap();
assert_eq!(reparsed, serialized, "round-trip must be lossless");
}
/// `auth_provider` alongside `api_key`/`env_key` warns (static keys
/// win in `resolve_credentials`, so the provider never runs) but keeps
/// both fields.
#[test]
fn auth_provider_shadowed_by_static_key_warns() {
let mut entry = toml::map::Map::new();
entry.insert("api_key".to_owned(), toml::Value::String("sk-x".into()));
entry.insert(
"auth_provider".to_owned(),
toml::Value::String("corp".into()),
);
let (models, warnings) = parse_single_entry(entry);
assert_eq!(warnings.len(), 1);
assert_eq!(warnings[0].kind, ConfigWarningKind::ConflictingFields);
assert_eq!(warnings[0].field(), Some("auth_provider"));
let parsed = models.get("m").unwrap();
assert_eq!(parsed.api_key.as_deref(), Some("sk-x"));
assert_eq!(parsed.auth_provider.as_deref(), Some("corp"));
// Provider alone: no warning.
let mut entry = toml::map::Map::new();
entry.insert(
"auth_provider".to_owned(),
toml::Value::String("corp".into()),
);
let (_, warnings) = parse_single_entry(entry);
assert_eq!(warnings, Vec::new());
// env_key is only a conditional shadow: warn, but as "may be shadowed".
let mut entry = toml::map::Map::new();
entry.insert("env_key".to_owned(), toml::Value::String("MY_KEY".into()));
entry.insert(
"auth_provider".to_owned(),
toml::Value::String("corp".into()),
);
let (_, warnings) = parse_single_entry(entry);
assert_eq!(warnings.len(), 1);
assert_eq!(warnings[0].kind, ConfigWarningKind::ConflictingFields);
assert!(warnings[0].reason.contains("may be shadowed"));
// An empty api_key does not shadow, so it must not warn.
let mut entry = toml::map::Map::new();
entry.insert("api_key".to_owned(), toml::Value::String(" ".into()));
entry.insert(
"auth_provider".to_owned(),
toml::Value::String("corp".into()),
);
let (_, warnings) = parse_single_entry(entry);
assert_eq!(warnings, Vec::new());
}
/// Drift guard: every `#[serde(alias)]` on [`ConfigModelOverride`] must
/// have a matching `ALIASES` pair, and vice versa. An unregistered alias
/// would send both-keys configs to the empty-override fallback.
@ -632,8 +823,8 @@ mod tests {
"canonical value must be retained"
);
assert_eq!(warnings.len(), 1);
assert_eq!(warnings[0].kind, ModelOverrideWarningKind::DuplicateAlias);
assert_eq!(warnings[0].field.as_deref(), Some(legacy));
assert_eq!(warnings[0].kind, ConfigWarningKind::DuplicateAlias);
assert_eq!(warnings[0].field(), Some(legacy));
}
}
}

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

@ -1388,6 +1388,7 @@ fn build_prefetched_map(
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);
@ -2015,6 +2016,7 @@ mod tests {
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;
@ -2027,6 +2029,7 @@ mod tests {
info: config::ModelInfo::fallback("plain-model"),
api_key: None,
env_key: None,
auth_provider: None,
api_base_url: None,
},
);
@ -2037,6 +2040,7 @@ mod tests {
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;
@ -2207,6 +2211,7 @@ mod tests {
info: config::ModelInfo::fallback("test-model"),
api_key: None,
env_key: None,
auth_provider: None,
api_base_url: None,
},
);
@ -2261,6 +2266,7 @@ mod tests {
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;
@ -2283,6 +2289,7 @@ mod tests {
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);
@ -2310,6 +2317,7 @@ mod tests {
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;
@ -2328,6 +2336,7 @@ mod tests {
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;
@ -2434,6 +2443,7 @@ mod tests {
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;
@ -2443,6 +2453,7 @@ mod tests {
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);
@ -2485,6 +2496,7 @@ mod tests {
info: config::ModelInfo::fallback(model_id),
api_key: None,
env_key: None,
auth_provider: None,
api_base_url: None,
}
}
@ -3268,6 +3280,7 @@ mod tests {
info: config::ModelInfo::fallback("static-one"),
api_key: None,
env_key: None,
auth_provider: None,
api_base_url: None,
},
);
@ -3295,6 +3308,7 @@ mod tests {
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;
@ -3304,6 +3318,7 @@ mod tests {
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);

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

@ -58,10 +58,12 @@ impl MvpAgent {
client_version,
) {
Some(mut cfg) => {
cfg.client_identifier = primary.client_identifier.clone();
cfg.attribution_callback = primary.attribution_callback.clone();
cfg.bearer_resolver = primary.bearer_resolver.clone();
cfg.max_retries = primary.max_retries;
crate::agent::config::stamp_session_local_sampler_fields(
&mut cfg,
primary,
primary.client_identifier.clone(),
primary.max_retries,
);
cfg
}
None => {

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

@ -2114,6 +2114,7 @@ fn find_model_by_id_prefers_key_then_falls_back_to_slug() {
},
api_key: None,
env_key: None,
auth_provider: None,
api_base_url: None,
};
let mut models = indexmap::IndexMap::new();

37
crates/codegen/xai-grok-shell/src/agent/relay.rs
View file

@ -509,24 +509,25 @@ where
let mut keepalive = tokio::time::interval(Duration::from_secs(KEEPALIVE_INTERVAL_SECS));
loop {
tokio::select! {
_ = cancel_write.cancelled() => break, msg_opt = from_agent_rx.recv() =>
{ match msg_opt { Some(msg) => { if
tracing::enabled!(tracing::Level::DEBUG) { if let Ok(json_val) =
serde_json::from_str::< serde_json::Value > (& msg) { let method =
json_val.get("method").and_then(| m | m.as_str()); let line_to_print =
match method { Some("session/update") => { let params = json_val
.get("params").unwrap_or(& serde_json::Value::Null);
format!("acp_outbound::session/update::{params}") } Some(m) =>
format!("acp_outbound::{m}"), None => "acp_outbound::response"
.to_string(), }; debug!("{line_to_print}"); } else {
debug!("acp_outbound::response"); } } if ! msg.is_empty() && let Err(e) =
ws_outbound.send(Message::Text(Utf8Bytes::from(msg))). await {
warn!(error = ? e, "failed to send to WS"); break; } } None => {
info!("Agent outbound channel closed"); break; } } } _ = keepalive.tick()
=> { tprintln!("ws::keep_alive_tick"); if let Err(e) = ws_outbound
.send(Message::Ping(Vec::new().into())). await {
tprintln!("ws::keep_alive::error::{:?}", & e); break; } }
}
_ = cancel_write.cancelled() => break, msg_opt = from_agent_rx.recv() =>
{ match msg_opt { Some(msg) => { if
tracing::enabled!(tracing::Level::DEBUG) { if let Ok(json_val) =
serde_json::from_str::< serde_json::Value > (& msg) { let method =
json_val.get("method").and_then(| m | m.as_str()); let line_to_print =
match method { Some("session/update") => { let params = json_val
.get("params").unwrap_or(& serde_json::Value::Null);
format!("acp_outbound::session/update::{params}") } Some(m) =>
format!("acp_outbound::{m}"), None => "acp_outbound::response"
.to_string(), }; debug!("{line_to_print}"); } else {
debug!("acp_outbound::response"); } }
if ! msg.is_empty() && let Err(e) =
ws_outbound.send(Message::Text(Utf8Bytes::from(msg))). await {
warn!(error = ? e, "failed to send to WS"); break; } } None => {
info!("Agent outbound channel closed"); break; } } } _ = keepalive.tick()
=> { tprintln!("ws::keep_alive_tick"); if let Err(e) = ws_outbound
.send(Message::Ping(Vec::new().into())). await {
tprintln!("ws::keep_alive::error::{:?}", & e); break; } }
}
}
anyhow::Ok(())
};

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

@ -1115,7 +1115,8 @@ async fn cancel_with_outcome_returns_variant_for_active_finished_unknown() {
);
assert!(
matches!(coordinator.cancel_with_outcome("sub-done"),
SubagentCancelOutcome::AlreadyFinished { status } if status == "completed")
SubagentCancelOutcome::AlreadyFinished { status }
if status == "completed")
);
assert!(
matches!(coordinator.cancel_with_outcome("nonexistent"),
@ -1850,7 +1851,8 @@ fn resume_vs_fork_helper_shapes_differ() {
assert!(
! matches!(resumed.conversation.get(1), Some(ConversationItem::User(u)) if u
.content.iter().any(| p | matches!(p,
xai_grok_sampling_types::conversation::ContentPart::Text { text } if text
xai_grok_sampling_types::conversation::ContentPart::Text { text }
if text
.contains("<background_context>"))))
);
}
@ -1912,7 +1914,8 @@ fn verbatim_fork_keeps_items_byte_for_byte_when_small() {
.any(|i| {
matches!(
i, ConversationItem::User(u) if u.content.iter().any(| p |
matches!(p, ContentPart::Text { text } if text.contains(needle)))
matches!(p, ContentPart::Text { text }
if text.contains(needle)))
)
})
};
@ -1954,7 +1957,8 @@ fn verbatim_fork_falls_back_to_summary_on_incomplete_tail() {
assert_eq!(ctx.prefix_len, Some(2));
assert!(
ctx.conversation.iter().any(| i | { matches!(i, ConversationItem::User(u) if u
.content.iter().any(| p | matches!(p, ContentPart::Text { text } if text
.content.iter().any(| p | matches!(p, ContentPart::Text { text }
if text
.contains("<background_context>")))) }),
"summarized fallback must produce a background_context blob"
);
@ -1995,7 +1999,8 @@ fn verbatim_fork_falls_back_to_summary_when_oversize() {
.any(|i| {
matches!(
i, ConversationItem::User(u) if u.content.iter().any(| p | matches!(p,
ContentPart::Text { text } if text.contains("<background_context>")))
ContentPart::Text { text }
if text.contains("<background_context>")))
)
});
assert!(has_blob, "oversize fallback must produce a background_context blob");
@ -3305,6 +3310,7 @@ fn test_model_entry(model_id: &str) -> crate::agent::config::ModelEntry {
},
api_key: None,
env_key: None,
auth_provider: None,
api_base_url: None,
}
}

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

@ -274,7 +274,8 @@ fn compaction_preserves_inherited_prefix() {
.any(|p| {
matches!(
p, xai_grok_sampling_types::conversation::ContentPart::Text {
text } if text.contains("<background_context>")
text }
if text.contains("<background_context>")
)
})
} else {
@ -2870,6 +2871,51 @@ async fn resolve_subagent_agent_definition_unknown_model_falls_through_to_inheri
assert_eq!(config.model, "grok-4.5");
assert_eq!(model_id.0.as_ref(), "grok-4.5");
}
/// Spawn-time credentials are cache-only: a cold spawn has no key,
/// never the parent session key.
#[tokio::test]
async fn subagent_override_provider_model_spawns_cache_only_credentials() {
use xai_grok_agent::config::ModelOverride;
let dir = tempfile::tempdir().unwrap();
let provider = crate::auth::test_counting_provider(
"test-subagent-spawn",
dir.path(),
);
let mut entry = test_model_entry("proxied-model");
entry.info.base_url = "https://gateway.example/v1".to_string();
entry.auth_provider = Some(provider.clone());
let mut models = indexmap::IndexMap::new();
models.insert("proxied".to_string(), entry);
let mut ctx = ctx_with_toggle(HashMap::new());
ctx.sampling_config.model = "grok-4.5".to_string();
ctx.model_id = acp::ModelId::new("grok-4.5");
ctx.available_models = models;
ctx.auth = Some(crate::auth::GrokAuth {
key: "parent-session-jwt".to_string(),
..Default::default()
});
ctx.subagent_model_overrides.insert("explore".to_string(), "proxied".to_string());
let (config, model_id) = resolve_subagent_sampling_config(
"explore",
&ModelOverride::Inherit,
&ctx,
)
.await;
assert_eq!(model_id.0.as_ref(), "proxied");
assert_eq!(
config.api_key, None,
"a cold cache spawns with no key, never the parent session key"
);
provider.ensure_fresh_token(None).await.rotated().unwrap();
let (config, _) = resolve_subagent_sampling_config(
"explore",
&ModelOverride::Inherit,
&ctx,
)
.await;
assert_eq!(config.api_key.as_deref(), Some("tok-1"));
assert_eq!(config.base_url, "https://gateway.example/v1");
}
#[test]
fn key_prefix_truncates_to_8_chars() {
let key = Some("eyJ0eXAiOiJhbGciOiJSUzI1NiJ9".to_string());

603
crates/codegen/xai-grok-shell/src/auth/auth_provider.rs Normal file
View file

@ -0,0 +1,603 @@
//! Model auth providers (`[auth_provider.<name>]`).
//!
//! A model opts in with `auth_provider = "<name>"`; the named table declares a
//! command that prints a fresh bearer token, which this module mints, caches,
//! and rotates for that model's requests.
//!
//! The minted token stays in memory only ([`AUTH_PROVIDER_SLOTS`] and chat
//! state, never `auth.json`); the command is a credential helper that owns its
//! own durable storage and OAuth2 refresh. See "Where model auth providers fit
//! (and don't)"
//! in `docs/internal/AUTH.md`.
//!
//! This is distinct from the `AuthCredentialProvider` HTTP consumers in
//! [`crate::auth::credential_provider`].
use super::token_output::{expiry_after_seconds, parse_token_output};
/// One named `[auth_provider.<name>]` table, honored only from the trusted
/// config layers (`parse_auth_providers`). A new field here needs a
/// `parse_auth_providers` warning decision.
#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Deserialize)]
#[serde(default)]
pub struct AuthProviderConfig {
/// Command that prints a bearer token on stdout, bare or as JSON
/// `{access_token, expires_in}`. Without `args` it runs via `sh -c`.
pub command: String,
/// Arguments for `command`. When present (even empty), the command runs
/// directly with no shell; `command` is a program name on `PATH`, or a path.
pub args: Option<Vec<String>>,
/// Fallback token lifetime in seconds, used when the command's output
/// carries no `expires_in`. Takes precedence over a JWT `exp` claim.
pub token_ttl_secs: Option<u64>,
/// Maximum seconds to wait for the command (default 30, clamped to 1..=600).
/// A turn waits up to this long on a mint, so keep helpers fast and
/// non-interactive.
pub timeout_secs: Option<u64>,
}
impl AuthProviderConfig {
pub(crate) fn is_usable(&self) -> bool {
!self.command.trim().is_empty()
}
}
/// A model's reference to a named auth provider, built by `resolve_model_list`.
#[derive(Clone, serde::Serialize, serde::Deserialize)]
#[serde(from = "AuthProviderRefData", into = "AuthProviderRefData")]
pub struct AuthProviderRef {
pub(crate) name: String,
pub(crate) config: AuthProviderConfig,
slot: ProviderSlot,
/// `true` once the trusted table is attached. A ref revived from bytes is
/// `false` and never mints or reads until [`AuthProviderRef::attach_trusted_config`]
/// joins the shared slot for its name.
resolved: bool,
}
/// Serialized form: the name only, so persisted bytes never carry a command.
#[derive(serde::Serialize, serde::Deserialize)]
struct AuthProviderRefData {
name: String,
}
impl From<AuthProviderRefData> for AuthProviderRef {
fn from(data: AuthProviderRefData) -> Self {
AuthProviderRef::unresolved(data.name)
}
}
impl From<AuthProviderRef> for AuthProviderRefData {
fn from(provider: AuthProviderRef) -> Self {
Self {
name: provider.name,
}
}
}
impl AuthProviderRef {
/// Production uses `unresolved` + `attach_trusted_config`.
#[cfg(test)]
pub(crate) fn new(name: String, config: AuthProviderConfig) -> Self {
let slot = provider_slot(&name);
Self {
name,
config,
slot,
resolved: true,
}
}
/// The in-memory form of a ref revived from bytes;
/// [`AuthProviderRef::attach_trusted_config`] resolves it.
pub(crate) fn unresolved(name: String) -> Self {
Self {
name,
config: AuthProviderConfig::default(),
slot: ProviderSlot::default(),
resolved: false,
}
}
/// Re-attach the trusted config for this name at model resolution
/// (`None` = the table was removed, leaving an unusable config). The ref
/// becomes authoritative, joins the shared slot for its name, and may mint.
pub(crate) fn attach_trusted_config(&mut self, config: Option<&AuthProviderConfig>) {
self.config = config.cloned().unwrap_or_default();
self.slot = provider_slot(&self.name);
self.resolved = true;
}
}
/// Ignores the slot; a deserialized ref compares unequal until resolution
/// re-attaches its config.
impl PartialEq for AuthProviderRef {
fn eq(&self, other: &Self) -> bool {
self.name == other.name && self.config == other.config
}
}
impl Eq for AuthProviderRef {}
impl std::fmt::Debug for AuthProviderRef {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AuthProviderRef")
.field("name", &self.name)
.field("config", &self.config)
.field("resolved", &self.resolved)
.finish_non_exhaustive()
}
}
struct MintedProviderToken {
token: String,
/// Handed back to the command on the next run; never sent on the wire.
refresh_token: Option<String>,
/// Drives the 401 fresh-mint guard.
minted_at: std::time::Instant,
expires_at: Option<chrono::DateTime<chrono::Utc>>,
/// The table version that minted the token; a different version reads as
/// stale (see [`token_identity`]), so edits re-mint.
minted_with: AuthProviderConfig,
}
/// The async lock is held across the command run, single-flighting mints
/// per provider name (shared across sessions). This dedupes concurrent
/// successes; a persistently failing helper is retried per waiter, each bounded
/// by the timeout clamp.
type ProviderSlot = std::sync::Arc<tokio::sync::Mutex<Option<MintedProviderToken>>>;
/// Shared token slots, one per resolved provider name. Bounded by the configured
/// provider names (only `attach_trusted_config` and test `new` insert), so no
/// eviction.
static AUTH_PROVIDER_SLOTS: std::sync::OnceLock<
std::sync::Mutex<std::collections::HashMap<String, ProviderSlot>>,
> = std::sync::OnceLock::new();
fn provider_slot(name: &str) -> ProviderSlot {
let map = AUTH_PROVIDER_SLOTS.get_or_init(Default::default);
let mut map = map.lock().unwrap_or_else(|e| e.into_inner());
map.entry(name.to_owned()).or_default().clone()
}
/// Pre-refresh margin: re-mint when the token expires within this window.
pub(crate) const PROVIDER_TOKEN_EXPIRY_SKEW_SECS: u64 = 60;
const PROVIDER_TOKEN_EXPIRY_SKEW: chrono::Duration =
chrono::Duration::seconds(PROVIDER_TOKEN_EXPIRY_SKEW_SECS as i64);
/// 401 fresh-mint guard: a token minted this recently is never re-minted on
/// rejection. Same idea as the guard in `unauthorized_recovery`, with a shorter
/// window because a provider mint is local and cheap.
const PROVIDER_TOKEN_FRESH_MINT_GUARD: std::time::Duration = std::time::Duration::from_secs(30);
const DEFAULT_PROVIDER_TIMEOUT_SECS: u64 = 30;
/// The effective mint timeout is clamped to `[1, this]`. A configured value
/// outside the range is honored up to the bound and draws a parse warning,
/// since a turn waits on the mint.
pub(crate) const PROVIDER_TIMEOUT_CEILING_SECS: u64 = 600;
/// Caps on the helper's captured output so a runaway command can't exhaust
/// memory before the timeout fires. A bearer (even a large JWT) is far under
/// the stdout cap; stderr only ever appears truncated in the failure log.
const PROVIDER_STDOUT_CAP_BYTES: u64 = 1 << 20; // 1 MiB
const PROVIDER_STDERR_CAP_BYTES: u64 = 64 << 10; // 64 KiB
/// The table fields that shape the minted token; a cached token minted under a
/// different set reads as stale, so a config edit re-mints. Destructured so a
/// new `AuthProviderConfig` field is a compile error until it is classified as
/// token-shaping (add it here) or an execution knob like `timeout_secs`
/// (editing it never invalidates).
fn token_identity(config: &AuthProviderConfig) -> (&str, Option<&[String]>, Option<u64>) {
let AuthProviderConfig {
command,
args,
token_ttl_secs,
timeout_secs: _,
} = config;
(command, args.as_deref(), *token_ttl_secs)
}
fn minted_token_is_stale(minted: &MintedProviderToken, config: &AuthProviderConfig) -> bool {
token_identity(&minted.minted_with) != token_identity(config)
|| minted
.expires_at
.is_some_and(|at| chrono::Utc::now() + PROVIDER_TOKEN_EXPIRY_SKEW >= at)
}
/// Log the missing-command warning once per provider, then at debug, so a
/// misconfigured model doesn't warn on every turn.
fn warn_empty_command(name: &str) {
static WARNED: std::sync::OnceLock<std::sync::Mutex<std::collections::HashSet<String>>> =
std::sync::OnceLock::new();
let first = WARNED
.get_or_init(Default::default)
.lock()
.unwrap_or_else(|e| e.into_inner())
.insert(name.to_owned());
const MSG: &str = "auth provider has no usable command: the [auth_provider.*] table is \
missing from the trusted config layers, or its `command` is empty";
if first {
tracing::warn!(provider = %name, "{MSG}");
} else {
tracing::debug!(provider = %name, "{MSG}");
}
}
/// Read up to `keep` bytes into `buf`, then drain and discard any remainder so
/// the child never blocks on a full pipe. Memory stays bounded by `keep`.
async fn read_capped<R>(reader: R, keep: u64, buf: &mut Vec<u8>) -> std::io::Result<()>
where
R: tokio::io::AsyncRead + Unpin,
{
use tokio::io::AsyncReadExt;
let mut limited = reader.take(keep);
limited.read_to_end(buf).await?;
tokio::io::copy(&mut limited.into_inner(), &mut tokio::io::sink()).await?;
Ok(())
}
/// Remove every first-party credential from the helper's environment. BYOK
/// isolates these keys on the wire, so the helper (the agent puts them in its
/// own env at startup) must not inherit them.
fn scrub_first_party_credentials(cmd: &mut tokio::process::Command) {
for var in crate::agent::config::FIRST_PARTY_CREDENTIAL_ENV_VARS {
cmd.env_remove(var);
}
}
/// Spawn `cmd`, capture stdout/stderr with a byte cap (reading both
/// concurrently so a full pipe on one can't deadlock the other; a runaway helper
/// is drained to a sink past the cap so it can't wedge the wait), and bound the
/// whole run by `timeout`. Exceeding the stdout cap is an error.
///
/// On timeout the child's entire process group is killed. The helper is a group
/// leader (`detach_command`'s `setsid`), so a compound `sh -c` helper's
/// grandchildren -- and the `GROK_AUTH_PROVIDER_*` credentials in their env --
/// do not outlive the reported timeout; `kill_on_drop` alone would reap only the
/// direct child.
async fn run_capped(
cmd: &mut tokio::process::Command,
timeout: std::time::Duration,
) -> anyhow::Result<std::process::Output> {
let mut child = cmd
.spawn()
.map_err(|e| anyhow::anyhow!("command failed to start: {e}"))?;
// Enroll the child's process group so the timeout path can tear down the
// whole tree. Best-effort: if enrollment fails, `kill_on_drop` still reaps
// the direct child.
let mut group = xai_grok_tools::util::ProcessGroup::new()
.map_err(|e| anyhow::anyhow!("process group setup failed: {e}"))?;
if let Err(e) = group.attach(&child) {
tracing::debug!(error = %e, "auth provider: could not enroll helper process group");
}
let stdout = child.stdout.take().expect("stdout is piped");
let stderr = child.stderr.take().expect("stderr is piped");
let mut out_buf = Vec::new();
let mut err_buf = Vec::new();
// One extra stdout byte so an over-cap write is detectable, not truncated.
// The stderr read is advisory (it only feeds the failure log), so only
// stdout governs the mint.
let capture = async {
let (out_res, err_res) = tokio::join!(
read_capped(stdout, PROVIDER_STDOUT_CAP_BYTES + 1, &mut out_buf),
read_capped(stderr, PROVIDER_STDERR_CAP_BYTES, &mut err_buf),
);
if let Err(e) = err_res {
tracing::debug!(error = %e, "auth provider: stderr capture failed (advisory)");
}
out_res.map_err(|e| anyhow::anyhow!("reading command stdout: {e}"))?;
child
.wait()
.await
.map_err(|e| anyhow::anyhow!("waiting on command: {e}"))
};
let status = match tokio::time::timeout(timeout, capture).await {
Ok(res) => res?,
Err(_elapsed) => {
let _ = group.kill();
anyhow::bail!("command timed out after {}s", timeout.as_secs());
}
};
if out_buf.len() as u64 > PROVIDER_STDOUT_CAP_BYTES {
anyhow::bail!("command wrote more than {PROVIDER_STDOUT_CAP_BYTES} bytes to stdout");
}
Ok(std::process::Output {
status,
stdout: out_buf,
stderr: err_buf,
})
}
async fn mint_provider_token(
provider: &AuthProviderRef,
mark_expired: bool,
previous: Option<&MintedProviderToken>,
) -> anyhow::Result<MintedProviderToken> {
use std::process::Stdio;
let name = &provider.name;
let config = &provider.config;
// Clamp to [1, ceiling]: the slot lock is held across the run, so an
// unbounded timeout would let one hung helper stall every turn sharing this
// provider name. The ceiling is a hard bound, not just a parse warning.
let timeout_secs = config
.timeout_secs
.unwrap_or(DEFAULT_PROVIDER_TIMEOUT_SECS)
.clamp(1, PROVIDER_TIMEOUT_CEILING_SECS);
tracing::info!(
provider = %name,
mark_expired,
timeout_secs,
"auth provider: running helper command"
);
let mut cmd = match config.args {
Some(ref args) => {
// Direct exec: the program name is a PATH lookup, so trim stray
// whitespace that would otherwise fail to resolve.
let mut cmd = tokio::process::Command::new(config.command.trim());
cmd.args(args);
cmd
}
None => {
let mut cmd = tokio::process::Command::new("sh");
cmd.args(["-c", &config.command]);
cmd
}
};
cmd.stdin(Stdio::null())
.stdout(Stdio::piped())
// Capture stderr for the failure log; inheriting corrupts the TUI.
.stderr(Stdio::piped())
// Reaps the direct child if the future is dropped; `run_capped`
// additionally kills the whole process group on timeout.
.kill_on_drop(true);
if mark_expired {
cmd.env("GROK_AUTH_EXPIRED", "1");
}
// Git-credential-helper handback: give the command the last stored
// credential so it can refresh instead of re-authenticating.
if let Some(prev) = previous {
cmd.env("GROK_AUTH_PROVIDER_ACCESS_TOKEN", &prev.token);
if let Some(refresh) = &prev.refresh_token {
cmd.env("GROK_AUTH_PROVIDER_REFRESH_TOKEN", refresh);
}
if let Some(expires_at) = prev.expires_at {
cmd.env("GROK_AUTH_PROVIDER_EXPIRES_AT", expires_at.to_rfc3339());
}
}
xai_grok_tools::util::detach_command(&mut cmd);
cmd.envs(xai_grok_tools::util::pager_env());
// Scrub last so nothing above can reintroduce a first-party credential.
scrub_first_party_credentials(&mut cmd);
let output = run_capped(&mut cmd, std::time::Duration::from_secs(timeout_secs)).await?;
let parsed = match parse_token_output(&output) {
Ok(parsed) => parsed,
Err(e) => {
let stderr = String::from_utf8_lossy(&output.stderr);
anyhow::bail!(
"{e} (stderr: {})",
crate::util::truncate(stderr.trim(), 300)
);
}
};
let expires_at = parsed
.expires_at
.or_else(|| config.token_ttl_secs.and_then(expiry_after_seconds))
.or_else(|| crate::auth::parse_jwt_expiration(&parsed.access_token));
tracing::info!(
provider = %name,
mark_expired,
expires_at = ?expires_at,
"auth provider minted token"
);
Ok(MintedProviderToken {
token: parsed.access_token,
refresh_token: parsed.refresh_token,
minted_at: std::time::Instant::now(),
expires_at,
minted_with: config.clone(),
})
}
#[derive(Debug, PartialEq, Eq)]
#[must_use = "a rotated token must be written to chat-state, or the wire keeps the stale key"]
pub(crate) enum ProviderRefreshOutcome {
/// `current_key` is already the fresh cached token; nothing to write.
Unchanged,
/// A token that should replace `current_key` on the wire.
Rotated(String),
/// The provider is unusable (unresolved or removed); already warned.
Unusable,
/// The mint ran and failed (logged).
MintFailed,
}
impl ProviderRefreshOutcome {
pub(crate) fn rotated(self) -> Option<String> {
match self {
Self::Rotated(token) => Some(token),
Self::Unchanged | Self::Unusable | Self::MintFailed => None,
}
}
}
impl AuthProviderRef {
/// The slot, locked for a mutating operation. A removed provider drops
/// its cached token and yields `None`, failing closed. An unresolved ref
/// (revived from bytes) fails closed without touching the shared slot.
async fn locked_slot(
&self,
) -> Option<tokio::sync::OwnedMutexGuard<Option<MintedProviderToken>>> {
if !self.resolved {
return None;
}
let mut slot = self.slot.clone().lock_owned().await;
if !self.config.is_usable() {
if slot.take().is_some() {
tracing::warn!(
provider = %self.name,
"auth provider removed from config: dropping its cached token"
);
}
warn_empty_command(&self.name);
return None;
}
Some(slot)
}
/// Cache-only read for sync resolution: never runs the command, blocks, or
/// mutates. `None` for an unresolved ref, a cold or stale cache, or a mint
/// in progress; minting happens pre-turn via [`AuthProviderRef::ensure_fresh_token`].
pub(crate) fn cached_token(&self) -> Option<String> {
if !self.resolved {
return None;
}
if !self.config.is_usable() {
warn_empty_command(&self.name);
return None;
}
// A mint in progress holds the lock; treat it as a miss rather than
// block the sync path.
let Ok(guard) = self.slot.try_lock() else {
tracing::debug!(provider = %self.name, "cache read skipped: mint in progress");
return None;
};
guard
.as_ref()
.filter(|m| !minted_token_is_stale(m, &self.config))
.map(|m| m.token.clone())
}
/// The token that should replace `current_key` on the wire: serves the
/// fresh cached token when chat-state lags behind a rotation, mints when
/// the cache is cold or stale. Mints or rotates a bearer; unrelated to an
/// OAuth refresh token.
pub(crate) async fn ensure_fresh_token(
&self,
current_key: Option<&str>,
) -> ProviderRefreshOutcome {
let Some(mut slot) = self.locked_slot().await else {
return ProviderRefreshOutcome::Unusable;
};
if let Some(ref minted) = *slot
&& !minted_token_is_stale(minted, &self.config)
{
return if current_key == Some(minted.token.as_str()) {
ProviderRefreshOutcome::Unchanged
} else {
ProviderRefreshOutcome::Rotated(minted.token.clone())
};
}
let mark_expired = slot.is_some();
let minted = match mint_provider_token(self, mark_expired, slot.as_ref()).await {
Ok(minted) => minted,
Err(e) => {
tracing::warn!(
provider = %self.name,
error = %e,
"auth provider pre-turn mint failed"
);
return ProviderRefreshOutcome::MintFailed;
}
};
let token = minted.token.clone();
*slot = Some(minted);
ProviderRefreshOutcome::Rotated(token)
}
/// The replacement for a server-rejected `rejected_key` (chat-state's
/// current key): a fresher cached token is adopted without a re-run,
/// otherwise the command runs once. `None` for a token minted moments ago
/// under the current table (the fresh-mint guard, which an edited table
/// bypasses).
pub(crate) async fn recover_rejected_token(&self, rejected_key: &str) -> Option<String> {
let mut slot = self.locked_slot().await?;
if let Some(ref minted) = *slot {
if minted.token != rejected_key && !minted_token_is_stale(minted, &self.config) {
return Some(minted.token.clone());
}
if minted.token == rejected_key
&& token_identity(&minted.minted_with) == token_identity(&self.config)
&& minted.minted_at.elapsed() < PROVIDER_TOKEN_FRESH_MINT_GUARD
{
tracing::warn!(
provider = %self.name,
"auth provider token rejected moments after mint: not \
re-running (fresh-mint guard); surfacing the 401"
);
return None;
}
}
tracing::info!(provider = %self.name, "auth provider token rejected: re-minting");
let minted = match mint_provider_token(self, true, slot.as_ref()).await {
Ok(minted) => minted,
Err(e) => {
tracing::warn!(
provider = %self.name,
error = %e,
"auth provider 401 re-mint failed"
);
// The server rejected the cached token and the re-mint failed;
// mark it stale so it is not re-served next turn (fail closed).
// The entry stays so its refresh token still feeds the next
// handback attempt.
if let Some(minted) = slot.as_mut() {
minted.expires_at = Some(chrono::Utc::now());
}
return None;
}
};
let token = minted.token.clone();
*slot = Some(minted);
Some(token)
}
}
/// Backdate a provider's mint time past the fresh-mint guard.
#[cfg(test)]
pub(crate) fn test_backdate_provider_mint(name: &str, age: std::time::Duration) {
let slot = provider_slot(name);
let mut slot = slot
.try_lock()
.expect("no mint in flight during test mutation");
if let Some(ref mut minted) = *slot {
minted.minted_at = std::time::Instant::now()
.checked_sub(age)
.expect("backdate before the process epoch");
}
}
/// A counting provider that prints "tok-1", "tok-2", ... on successive runs.
#[cfg(test)]
pub(crate) fn test_counting_provider(name: &str, dir: &std::path::Path) -> AuthProviderRef {
let counter = dir.join("count");
AuthProviderRef::new(
name.to_owned(),
AuthProviderConfig {
command: format!(
"echo run >> {c}; printf 'tok-%s' \"$(wc -l < {c} | tr -d ' ')\"",
c = counter.display()
),
args: None,
token_ttl_secs: Some(3600),
timeout_secs: None,
},
)
}
#[cfg(test)]
fn test_expire_provider_token(name: &str) {
let slot = provider_slot(name);
let mut slot = slot
.try_lock()
.expect("no mint in flight during test mutation");
if let Some(ref mut minted) = *slot {
minted.expires_at = Some(chrono::Utc::now() - chrono::Duration::seconds(1));
}
}
#[cfg(test)]
#[path = "auth_provider_tests.rs"]
mod tests;

763
crates/codegen/xai-grok-shell/src/auth/auth_provider_tests.rs Normal file
View file

@ -0,0 +1,763 @@
// Slot names are process-global, so every test uses a unique name (no #[serial]
// needed). No test mutates the process env: the scrub test sets its leak values
// on the child command instead.
use super::test_counting_provider as counting_provider;
use super::*;
#[tokio::test]
async fn provider_token_is_cached_while_fresh() {
let dir = tempfile::tempdir().unwrap();
let provider = counting_provider("test-cache", dir.path());
assert_eq!(
provider.cached_token(),
None,
"cache-only read must miss on a cold cache without running the command"
);
let first = provider.ensure_fresh_token(None).await.rotated().unwrap();
let second = provider.ensure_fresh_token(None).await.rotated().unwrap();
assert_eq!(first, "tok-1");
assert_eq!(second, "tok-1", "fresh token must be served from cache");
assert_eq!(
provider.cached_token().as_deref(),
Some("tok-1"),
"sync cache-only read must serve the warm cache"
);
}
#[tokio::test]
async fn provider_token_reminted_when_expired() {
let dir = tempfile::tempdir().unwrap();
let provider = counting_provider("test-expiry", dir.path());
assert_eq!(
provider.ensure_fresh_token(None).await.rotated().unwrap(),
"tok-1"
);
test_expire_provider_token("test-expiry");
assert_eq!(
provider.cached_token(),
None,
"cache-only read must not serve a stale token"
);
assert_eq!(
provider.ensure_fresh_token(None).await.rotated().unwrap(),
"tok-2",
"expired token must be re-minted"
);
}
#[tokio::test]
async fn provider_pre_turn_refresh_semantics() {
let dir = tempfile::tempdir().unwrap();
let provider = counting_provider("test-stale", dir.path());
let token = provider.ensure_fresh_token(None).await.rotated().unwrap();
assert_eq!(
provider.ensure_fresh_token(Some(&token)).await,
ProviderRefreshOutcome::Unchanged,
"fresh matching token must not be re-minted pre-turn"
);
assert_eq!(
provider
.ensure_fresh_token(Some("lagging-chat-state-key"))
.await
.rotated()
.as_deref(),
Some("tok-1"),
"chat-state lagging behind a rotation adopts the fresh cached token"
);
test_expire_provider_token("test-stale");
assert_eq!(
provider
.ensure_fresh_token(Some(&token))
.await
.rotated()
.as_deref(),
Some("tok-2"),
"stale token must be re-minted pre-turn"
);
}
#[tokio::test]
async fn provider_401_recovery_has_fresh_mint_guard() {
let dir = tempfile::tempdir().unwrap();
let provider = counting_provider("test-401", dir.path());
let token = provider.ensure_fresh_token(None).await.rotated().unwrap();
assert_eq!(
provider.recover_rejected_token(&token).await,
None,
"a token minted moments ago must not be re-minted on 401 (loop guard)"
);
test_backdate_provider_mint("test-401", std::time::Duration::from_secs(60));
assert_eq!(
provider.recover_rejected_token(&token).await.as_deref(),
Some("tok-2"),
"an aged rejected token is re-minted once"
);
assert_eq!(
provider.recover_rejected_token(&token).await.as_deref(),
Some("tok-2"),
"a rejection of the already-replaced key adopts the fresh token without a re-run"
);
}
/// Regression: a warm cache must not outlive the provider's config.
#[tokio::test]
async fn provider_removed_from_config_drops_cached_token() {
let dir = tempfile::tempdir().unwrap();
let provider = counting_provider("test-removed", dir.path());
let token = provider.ensure_fresh_token(None).await.rotated().unwrap();
let removed = AuthProviderRef::new("test-removed".to_owned(), AuthProviderConfig::default());
assert_eq!(
removed.cached_token(),
None,
"empty command must fail closed even with a warm cache"
);
assert_eq!(
removed.ensure_fresh_token(Some(&token)).await,
ProviderRefreshOutcome::Unusable
);
let restored = counting_provider("test-removed", dir.path());
assert_eq!(
restored
.ensure_fresh_token(Some(&token))
.await
.rotated()
.as_deref(),
Some("tok-2"),
"the removed provider's token must not survive in the slot"
);
}
#[tokio::test]
async fn provider_config_edit_invalidates_cached_token() {
let dir = tempfile::tempdir().unwrap();
let old = counting_provider("test-freshen", dir.path());
assert_eq!(
old.ensure_fresh_token(None).await.rotated().unwrap(),
"tok-1"
);
let edited = AuthProviderRef::new(
"test-freshen".to_owned(),
AuthProviderConfig {
command: "printf edited-token".to_owned(),
args: None,
token_ttl_secs: Some(3600),
timeout_secs: None,
},
);
assert_eq!(
edited.cached_token(),
None,
"the unexpired old token must not be served under the edited table"
);
assert_eq!(
edited
.ensure_fresh_token(Some("tok-1"))
.await
.rotated()
.as_deref(),
Some("edited-token"),
"refresh must run the edited command without waiting for expiry"
);
}
/// The fresh-mint guard applies per table version.
#[tokio::test]
async fn provider_401_recovery_reminted_under_edited_config() {
let dir = tempfile::tempdir().unwrap();
let old = counting_provider("test-401-edited", dir.path());
let token = old.ensure_fresh_token(None).await.rotated().unwrap();
let edited = AuthProviderRef::new(
"test-401-edited".to_owned(),
AuthProviderConfig {
command: "printf new-config-token".to_owned(),
args: None,
token_ttl_secs: Some(3600),
timeout_secs: None,
},
);
assert_eq!(
edited.recover_rejected_token(&token).await.as_deref(),
Some("new-config-token"),
"recovery must run the edited command, not adopt the old-table token"
);
}
/// Editing only `timeout_secs` keeps the token; it is not part of
/// `token_identity`.
#[tokio::test]
async fn provider_timeout_edit_does_not_invalidate_token() {
let dir = tempfile::tempdir().unwrap();
let provider = counting_provider("test-timeout-edit", dir.path());
provider.ensure_fresh_token(None).await.rotated().unwrap();
let retimed = AuthProviderRef::new(
"test-timeout-edit".to_owned(),
AuthProviderConfig {
command: provider.config.command.clone(),
args: None,
token_ttl_secs: Some(3600),
timeout_secs: Some(5),
},
);
assert_eq!(
retimed.cached_token().as_deref(),
Some("tok-1"),
"a timeout-only edit must not invalidate the cached token"
);
}
#[tokio::test]
async fn attach_trusted_config_lets_a_revived_ref_mint() {
let dir = tempfile::tempdir().unwrap();
let template = counting_provider("test-attach", dir.path());
let mut revived: AuthProviderRef = serde_json::from_str(r#"{"name": "test-attach"}"#).unwrap();
assert_eq!(
revived.ensure_fresh_token(None).await,
ProviderRefreshOutcome::Unusable
);
revived.attach_trusted_config(Some(&template.config));
assert_eq!(
revived.ensure_fresh_token(None).await.rotated().as_deref(),
Some("tok-1"),
"a re-attached ref must be able to mint"
);
}
/// A ref revived from bytes never mutates the shared slot: a mutating
/// call fails closed and leaves a resolved ref's token intact.
#[tokio::test]
async fn deserialized_ref_never_drops_the_shared_token() {
let dir = tempfile::tempdir().unwrap();
let resolved = counting_provider("test-unresolved", dir.path());
resolved.ensure_fresh_token(None).await.rotated().unwrap();
let revived: AuthProviderRef = serde_json::from_str(r#"{"name": "test-unresolved"}"#).unwrap();
assert_eq!(
revived.ensure_fresh_token(None).await,
ProviderRefreshOutcome::Unusable
);
assert_eq!(revived.recover_rejected_token("tok-1").await, None);
assert_eq!(
resolved.cached_token().as_deref(),
Some("tok-1"),
"the resolved ref's token must survive a mutating call on the stub"
);
}
/// A ref serializes to its name only: the revived ref carries no command
/// and fails closed until re-attached, while the shared slot still serves
/// resolved refs of the same name.
#[tokio::test]
async fn provider_ref_serializes_name_only_and_drops_config() {
let dir = tempfile::tempdir().unwrap();
let provider = counting_provider("test-serde", dir.path());
provider.ensure_fresh_token(None).await.rotated().unwrap();
let bytes = serde_json::to_string(&provider).unwrap();
assert!(bytes.contains("test-serde"));
assert!(
!bytes.contains("tok-%s") && !bytes.contains("command"),
"the serialized form must carry the name only: {bytes}"
);
let revived: AuthProviderRef = serde_json::from_str(&bytes).unwrap();
assert_eq!(revived.name, "test-serde");
assert_eq!(
revived.config,
AuthProviderConfig::default(),
"a serialized command must not survive deserialization"
);
assert_eq!(
revived.cached_token(),
None,
"an unresolved ref fails closed"
);
let same_name = counting_provider("test-serde", dir.path());
assert_eq!(
same_name.cached_token().as_deref(),
Some("tok-1"),
"the shared slot still serves refs constructed with the real config"
);
}
#[tokio::test]
async fn provider_refresh_sets_expired_env() {
let provider = AuthProviderRef::new(
"test-expired-env".to_owned(),
AuthProviderConfig {
command: "printf 'tok-%s' \"${GROK_AUTH_EXPIRED:-0}\"".to_owned(),
args: None,
token_ttl_secs: Some(3600),
timeout_secs: None,
},
);
assert_eq!(
provider.ensure_fresh_token(None).await.rotated().as_deref(),
Some("tok-0"),
"first mint runs without GROK_AUTH_EXPIRED"
);
test_expire_provider_token("test-expired-env");
assert_eq!(
provider.ensure_fresh_token(None).await.rotated().as_deref(),
Some("tok-1"),
"re-mints run with GROK_AUTH_EXPIRED=1"
);
}
#[tokio::test]
async fn provider_concurrent_mints_single_flight() {
let dir = tempfile::tempdir().unwrap();
let counter = dir.path().join("count");
let provider = AuthProviderRef::new(
"test-single-flight".to_owned(),
AuthProviderConfig {
command: format!(
"sleep 0.3; echo run >> {c}; printf 'tok-%s' \"$(wc -l < {c} | tr -d ' ')\"",
c = counter.display()
),
args: None,
token_ttl_secs: Some(3600),
timeout_secs: None,
},
);
let (a, b) = tokio::join!(
provider.ensure_fresh_token(None),
provider.ensure_fresh_token(None)
);
assert_eq!(a.rotated().as_deref(), Some("tok-1"));
assert_eq!(
b.rotated().as_deref(),
Some("tok-1"),
"second caller adopts, never re-runs"
);
let runs = std::fs::read_to_string(&counter).unwrap().lines().count();
assert_eq!(runs, 1, "the command must run exactly once");
}
/// Proven by staleness: an expiry inside the 60s skew re-mints, a
/// distant one serves from cache.
#[tokio::test]
async fn provider_expiry_source_precedence() {
fn short_jwt() -> String {
// exp within the skew window: stale immediately if consumed.
jwt_with_exp(chrono::Utc::now().timestamp() + 30)
}
fn long_jwt() -> String {
jwt_with_exp(chrono::Utc::now().timestamp() + 7200)
}
fn jwt_with_exp(exp: i64) -> String {
jsonwebtoken::encode(
&jsonwebtoken::Header::default(),
&serde_json::json!({ "exp": exp }),
&jsonwebtoken::EncodingKey::from_secret(b"test"),
)
.unwrap()
}
async fn mints_after_first(
name: &str,
command: String,
token_ttl_secs: Option<u64>,
counter: &std::path::Path,
) -> usize {
let provider = AuthProviderRef::new(
name.to_owned(),
AuthProviderConfig {
command,
args: None,
token_ttl_secs,
timeout_secs: None,
},
);
let first = provider
.ensure_fresh_token(None)
.await
.rotated()
.expect("first mint");
let _ = provider.ensure_fresh_token(Some(&first)).await;
std::fs::read_to_string(counter).unwrap().lines().count()
}
let dir = tempfile::tempdir().unwrap();
// expires_in=10 (stale) wins over token_ttl_secs=3600 (fresh): re-mints.
let c1 = dir.path().join("c1");
let cmd1 = format!(
"echo run >> {}; printf '{{\"access_token\":\"t1\",\"expires_in\":10}}'",
c1.display()
);
assert_eq!(
mints_after_first("test-exp-expires-in", cmd1, Some(3600), &c1).await,
2,
"expires_in must win over token_ttl_secs"
);
// token_ttl_secs=1 (stale) wins over a 2h JWT exp (fresh): re-mints.
let c2 = dir.path().join("c2");
let cmd2 = format!("echo run >> {}; printf '{}'", c2.display(), long_jwt());
assert_eq!(
mints_after_first("test-exp-ttl", cmd2, Some(1), &c2).await,
2,
"token_ttl_secs must win over the JWT exp claim"
);
// JWT exp alone: a near-expiry claim (inside the skew) re-mints,
// proving the claim is consumed when nothing else is configured.
let c3 = dir.path().join("c3");
let cmd3 = format!("echo run >> {}; printf '{}'", c3.display(), short_jwt());
assert_eq!(
mints_after_first("test-exp-jwt", cmd3, None, &c3).await,
2,
"the JWT exp claim must apply when expires_in and token_ttl_secs are absent"
);
}
#[tokio::test]
async fn provider_unusable_expiry_still_mints() {
let provider = AuthProviderRef::new(
"test-overflow".to_owned(),
AuthProviderConfig {
command: format!(
"printf '{{\"access_token\":\"t\",\"expires_in\":{}}}'",
u64::MAX
),
args: None,
token_ttl_secs: Some(u64::MAX),
timeout_secs: None,
},
);
assert_eq!(
provider.ensure_fresh_token(None).await.rotated().as_deref(),
Some("t"),
"an unusable expiry still mints; the token just has no expiry"
);
assert_eq!(
provider.ensure_fresh_token(Some("t")).await,
ProviderRefreshOutcome::Unchanged,
"no expiry source: never proactively re-minted"
);
}
#[tokio::test]
async fn provider_args_run_without_a_shell() {
let provider = AuthProviderRef::new(
"test-args".to_owned(),
AuthProviderConfig {
command: "printf".to_owned(),
// Shell metacharacters stay literal under direct exec.
args: Some(vec!["tok-$HOME;42".to_owned()]),
token_ttl_secs: Some(3600),
timeout_secs: None,
},
);
assert_eq!(
provider.ensure_fresh_token(None).await.rotated().as_deref(),
Some("tok-$HOME;42"),
);
}
#[tokio::test]
async fn provider_command_times_out() {
let provider = AuthProviderRef::new(
"test-timeout".to_owned(),
AuthProviderConfig {
command: "sleep 20; printf never".to_owned(),
args: None,
token_ttl_secs: None,
timeout_secs: Some(1),
},
);
let start = std::time::Instant::now();
assert_eq!(
provider.ensure_fresh_token(None).await,
ProviderRefreshOutcome::MintFailed
);
assert!(
start.elapsed().as_secs() < 5,
"1s timeout_secs must bound the mint (took {}s)",
start.elapsed().as_secs()
);
}
#[tokio::test]
async fn provider_zero_timeout_clamps_to_one_second() {
// `timeout_secs = 0` clamps up to the 1s floor, so an instant helper mints
// rather than failing immediately.
let fast = AuthProviderRef::new(
"test-zero-timeout-fast".to_owned(),
AuthProviderConfig {
command: "printf tok".to_owned(),
args: None,
token_ttl_secs: Some(3600),
timeout_secs: Some(0),
},
);
assert_eq!(
fast.ensure_fresh_token(None).await.rotated().as_deref(),
Some("tok")
);
// ...and clamps down from the 30s default: a helper that runs past 1s times
// out, proving the effective bound is the clamp, not the default.
let slow = AuthProviderRef::new(
"test-zero-timeout-slow".to_owned(),
AuthProviderConfig {
command: "sleep 5; printf tok".to_owned(),
args: None,
token_ttl_secs: Some(3600),
timeout_secs: Some(0),
},
);
assert!(
matches!(
slow.ensure_fresh_token(None).await,
ProviderRefreshOutcome::MintFailed
),
"a >1s helper under timeout_secs=0 must time out at the 1s clamp"
);
}
/// The distinct mint-failure modes (timeout, spawn failure, ran-but-no-token)
/// surface distinct, greppable error messages so operators can triage them.
#[tokio::test]
async fn mint_error_messages_distinguish_failure_modes() {
let timed_out = AuthProviderRef::new(
"test-classify-timeout".to_owned(),
AuthProviderConfig {
command: "sleep 20".to_owned(),
args: None,
token_ttl_secs: None,
timeout_secs: Some(1),
},
);
let err = mint_provider_token(&timed_out, false, None)
.await
.err()
.expect("timeout must fail the mint");
assert!(err.to_string().contains("timed out"), "got: {err}");
let missing = AuthProviderRef::new(
"test-classify-spawn".to_owned(),
AuthProviderConfig {
command: "/nonexistent/provider-binary".to_owned(),
args: Some(vec![]),
token_ttl_secs: None,
timeout_secs: Some(5),
},
);
let err = mint_provider_token(&missing, false, None)
.await
.err()
.expect("spawn failure must fail the mint");
assert!(err.to_string().contains("failed to start"), "got: {err}");
let empty_output = AuthProviderRef::new(
"test-classify-permanent".to_owned(),
AuthProviderConfig {
command: "printf ''".to_owned(),
args: None,
token_ttl_secs: None,
timeout_secs: Some(5),
},
);
let err = mint_provider_token(&empty_output, false, None)
.await
.err()
.expect("empty output must fail the mint");
assert!(err.to_string().contains("no output"), "got: {err}");
}
/// On an in-session re-mint, the prior credential is handed back to the command
/// via `GROK_AUTH_PROVIDER_*`, so a refresh-grant command can refresh instead of
/// re-authenticating. Nothing is written to disk.
#[tokio::test]
async fn re_mint_hands_the_prior_token_back_to_the_command() {
let provider = AuthProviderRef::new(
"test-handback".to_owned(),
AuthProviderConfig {
command: "printf 'seen-%s' \"${GROK_AUTH_PROVIDER_ACCESS_TOKEN:-none}\"".to_owned(),
args: None,
token_ttl_secs: Some(3600),
timeout_secs: None,
},
);
let first = provider.ensure_fresh_token(None).await.rotated().unwrap();
assert_eq!(first, "seen-none", "the first mint has no prior credential");
test_expire_provider_token("test-handback");
assert_eq!(
provider
.ensure_fresh_token(Some(&first))
.await
.rotated()
.as_deref(),
Some("seen-seen-none"),
"the re-mint must receive the prior access token via env"
);
}
/// A 401 whose re-mint fails invalidates the rejected token, so it is not
/// re-served next turn (fail closed) even while still locally unexpired.
#[tokio::test]
async fn failed_401_remint_invalidates_the_cached_token() {
let dir = tempfile::tempdir().unwrap();
let counter = dir.path().join("count");
// Mints tok-1 on the first run, then exits non-zero on every later run.
let provider = AuthProviderRef::new(
"test-401-invalidate".to_owned(),
AuthProviderConfig {
command: format!(
"echo run >> {c}; n=$(wc -l < {c} | tr -d ' '); \
[ \"$n\" = 1 ] && printf 'tok-1' || exit 1",
c = counter.display()
),
args: None,
token_ttl_secs: Some(3600),
timeout_secs: None,
},
);
let token = provider.ensure_fresh_token(None).await.rotated().unwrap();
assert_eq!(token, "tok-1");
// Age past the fresh-mint guard so recovery attempts a re-mint.
test_backdate_provider_mint("test-401-invalidate", PROVIDER_TOKEN_FRESH_MINT_GUARD * 2);
assert_eq!(
provider.recover_rejected_token(&token).await,
None,
"a failed re-mint surfaces the 401"
);
assert_eq!(
provider.cached_token(),
None,
"a rejected token whose re-mint failed must not be re-served"
);
}
/// A pre-turn re-mint that fails over a now-stale cached token leaves nothing
/// servable: the stale token is never handed to the wire (mirror of the 401
/// path, for the pre-turn path).
#[tokio::test]
async fn failed_pre_turn_mint_does_not_serve_the_stale_token() {
let dir = tempfile::tempdir().unwrap();
let counter = dir.path().join("count");
let provider = AuthProviderRef::new(
"test-pre-turn-stale".to_owned(),
AuthProviderConfig {
command: format!(
"echo run >> {c}; n=$(wc -l < {c} | tr -d ' '); \
[ \"$n\" = 1 ] && printf 'tok-1' || exit 1",
c = counter.display()
),
args: None,
token_ttl_secs: Some(3600),
timeout_secs: None,
},
);
let token = provider.ensure_fresh_token(None).await.rotated().unwrap();
assert_eq!(token, "tok-1");
// Make the cached token stale so the next pre-turn call re-mints (and fails).
test_expire_provider_token("test-pre-turn-stale");
assert!(matches!(
provider.ensure_fresh_token(Some(token.as_str())).await,
ProviderRefreshOutcome::MintFailed
));
assert_eq!(
provider.cached_token(),
None,
"a stale token whose pre-turn re-mint failed must not be served"
);
}
/// A helper that writes past the stdout cap fails closed (permanent), so a
/// runaway command can't exhaust memory or put a huge token on the wire.
#[tokio::test]
async fn provider_output_over_cap_fails_closed() {
let over = PROVIDER_STDOUT_CAP_BYTES + 4096;
let provider = AuthProviderRef::new(
"test-stdout-cap".to_owned(),
AuthProviderConfig {
command: format!("head -c {over} /dev/zero"),
args: None,
token_ttl_secs: None,
timeout_secs: Some(5),
},
);
let err = mint_provider_token(&provider, false, None)
.await
.err()
.expect("over-cap output must fail the mint");
assert!(
err.to_string().contains("more than"),
"an over-cap write must be reported as such, got: {err}"
);
assert_eq!(
provider.ensure_fresh_token(None).await,
ProviderRefreshOutcome::MintFailed
);
}
/// Every first-party credential env var is scrubbed from the helper, so a BYOK
/// helper never inherits the keys BYOK isolates on the wire.
///
/// The test drives its set/echo from an independent audited `EXPECTED` list, not
/// from the scrub const, so it is not tautological: dropping an entry from
/// `FIRST_PARTY_CREDENTIAL_ENV_VARS` alone leaves that var set on the command and
/// trips the assert below, and removing one from both requires deliberately
/// editing this audited list.
///
/// The leak values are set on the child command, not the process env, so the
/// test is hermetic: it needs no `#[serial]` and cannot race a sibling test that
/// reads a first-party credential (e.g. the `auth::manager` session tests).
#[tokio::test]
async fn provider_helper_env_scrubs_first_party_credentials() {
// The credentials a BYOK helper must never inherit. Editing this list is the
// audit checkpoint: it must equal the production scrub const.
const EXPECTED: &[&str] = &[
"XAI_API_KEY",
"GROK_CODE_XAI_API_KEY",
"GROK_AUTH",
"GROK_AUTH_PATH",
"GROK_DEPLOYMENT_KEY",
"GROK_EXTRA_AUTH_KEY",
"GROK_TRACE_UPLOAD_CREDENTIALS_FILE",
"OTEL_EXPORTER_OTLP_HEADERS",
"GROK_INTERNAL_OTLP_HEADERS",
];
assert_eq!(
crate::agent::config::FIRST_PARTY_CREDENTIAL_ENV_VARS,
EXPECTED,
"the scrub list changed: re-audit that every entry is a first-party \
credential a BYOK helper must not inherit, then update EXPECTED"
);
// Echo each expected var back; the scrub must leave every one empty. A
// scrub-const entry that EXPECTED still lists but production stopped removing
// stays at its leak value and surfaces here.
let echo = EXPECTED
.iter()
.map(|v| format!("${{{v}-}}"))
.collect::<Vec<_>>()
.join("");
let mut cmd = tokio::process::Command::new("sh");
cmd.args(["-c", &format!("printf 'tok[%s]' \"{echo}\"")]);
for var in EXPECTED {
cmd.env(var, "first-party-leak");
}
super::scrub_first_party_credentials(&mut cmd);
let output = cmd.output().await.expect("helper spawns");
assert_eq!(
String::from_utf8_lossy(&output.stdout),
"tok[]",
"no first-party credential may survive into the helper env"
);
}

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

@ -1,60 +1,11 @@
use crate::auth::token_output::parse_token_output;
use crate::auth::{AuthMode, GrokAuth};
#[derive(serde::Deserialize)]
pub(crate) struct ExternalAuthOutput {
pub access_token: String,
#[serde(default)]
pub refresh_token: Option<String>,
#[serde(default)]
pub expires_in: Option<u64>,
/// Token issuer. An xAI issuer marks the credential as first-party;
/// see [`GrokAuth::is_xai_auth`].
#[serde(default)]
pub issuer: Option<String>,
}
/// Parse process output (stdout) into a `GrokAuth`. Accepts bare token or JSON.
/// Parse stdout into a session-credential `GrokAuth`.
pub(crate) fn parse_output(output: &std::process::Output) -> anyhow::Result<GrokAuth> {
if !output.status.success() {
anyhow::bail!("exited with {}", output.status);
}
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_owned();
if stdout.is_empty() {
anyhow::bail!("produced no output on stdout");
}
let (token, refresh_token, expires_at, issuer) =
if let Ok(parsed) = serde_json::from_str::<ExternalAuthOutput>(&stdout) {
tracing::debug!(
has_refresh_token = parsed.refresh_token.is_some(),
expires_in = ?parsed.expires_in,
issuer = ?parsed.issuer,
"auth: parsed external provider output as JSON"
);
let expires_at = parsed
.expires_in
.map(|secs| chrono::Utc::now() + chrono::Duration::seconds(secs as i64));
let issuer = parsed
.issuer
.map(|i| i.trim().to_owned())
.filter(|i| !i.is_empty());
(
parsed.access_token,
parsed.refresh_token,
expires_at,
issuer,
)
} else {
tracing::debug!(
stdout_len = stdout.len(),
"auth: treating output as bare token"
);
(stdout, None, None, None)
};
let parsed = parse_token_output(output)?;
Ok(GrokAuth {
key: token,
key: parsed.access_token,
auth_mode: AuthMode::External,
create_time: chrono::Utc::now(),
user_id: String::new(),
@ -74,20 +25,25 @@ pub(crate) fn parse_output(output: &std::process::Output) -> anyhow::Result<Grok
team_blocked_reasons: vec![],
coding_data_retention_opt_out: crate::auth::default_coding_data_retention_opt_out(),
has_grok_code_access: None,
refresh_token,
expires_at,
oidc_issuer: issuer,
refresh_token: parsed.refresh_token,
expires_at: parsed.expires_at,
oidc_issuer: parsed.issuer,
oidc_client_id: None,
})
}
/// Sync version for mid-session refresh. 5s timeout for refresh, 60s for initial.
pub(crate) fn run_external_auth_sync(command: &str, is_refresh: bool) -> Option<GrokAuth> {
let timeout_secs = if is_refresh { 5 } else { 60 };
run_auth_command(command, timeout_secs, is_refresh)
}
/// Runs `command` via `sh -c`; `mark_expired` sets `GROK_AUTH_EXPIRED=1` so the
/// helper can distinguish re-mints from first runs.
fn run_auth_command(command: &str, timeout_secs: u64, mark_expired: bool) -> Option<GrokAuth> {
use std::process::{Command, Stdio};
let timeout_secs = if is_refresh { 5 } else { 60 };
tracing::info!(cmd = %command, is_refresh, timeout_secs, "auth: running external auth provider (sync)");
tracing::info!(cmd = %command, mark_expired, timeout_secs, "auth: running external auth provider (sync)");
let mut cmd = Command::new("sh");
cmd.args(["-c", command])
@ -95,7 +51,7 @@ pub(crate) fn run_external_auth_sync(command: &str, is_refresh: bool) -> Option<
.stdout(Stdio::piped())
// Pipe stderr — inherit would corrupt the TUI alternate screen.
.stderr(Stdio::piped());
if is_refresh {
if mark_expired {
cmd.env("GROK_AUTH_EXPIRED", "1");
}
xai_grok_tools::util::detach_std_command(&mut cmd);
@ -224,14 +180,13 @@ mod tests {
}
#[test]
fn parse_output_malformed_json_falls_back_to_bare() {
fn parse_output_json_shaped_but_invalid_is_err() {
let output = std::process::Output {
status: std::process::Command::new("true").status().unwrap(),
stdout: b"{not valid json}".to_vec(),
stderr: vec![],
};
let auth = parse_output(&output).unwrap();
assert_eq!(auth.key, "{not valid json}");
assert!(parse_output(&output).is_err());
}
#[test]

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

@ -1,4 +1,5 @@
pub(crate) mod attribution;
mod auth_provider;
mod config;
pub mod credential_provider;
#[path = "devbox_login_stub.rs"]
@ -15,7 +16,14 @@ pub(crate) mod recovery;
pub(crate) mod refresh;
pub(crate) mod single_flight;
mod storage;
mod token_output;
pub(crate) mod token_type;
pub use auth_provider::{AuthProviderConfig, AuthProviderRef};
pub(crate) use auth_provider::{
PROVIDER_TIMEOUT_CEILING_SECS, PROVIDER_TOKEN_EXPIRY_SKEW_SECS, ProviderRefreshOutcome,
};
#[cfg(test)]
pub(crate) use auth_provider::{test_backdate_provider_mint, test_counting_provider};
pub(crate) use config::LEGACY_AUTH_SCOPE;
pub use config::{
ForceLoginTeam, GrokComConfig, OAuth2ProviderConfig, OidcAuthConfig, PreferredAuthMethod,

155
crates/codegen/xai-grok-shell/src/auth/token_output.rs Normal file
View file

@ -0,0 +1,155 @@
//! Shared parser for an auth command's stdout.
//!
//! Both auth paths run a command that prints a bearer token and parse it here:
//! the session external-auth path ([`super::external_auth`]) and the per-model
//! provider mint ([`super::auth_provider`]).
#[derive(serde::Deserialize)]
pub(crate) struct ExternalAuthOutput {
pub access_token: String,
#[serde(default)]
pub refresh_token: Option<String>,
#[serde(default)]
pub expires_in: Option<u64>,
/// An xAI issuer marks the credential as first-party
/// (see [`crate::auth::GrokAuth::is_xai_auth`]).
#[serde(default)]
pub issuer: Option<String>,
}
/// A bearer must be a single line: reject control characters (including an
/// interior newline) so a malformed token can never be smuggled onto an HTTP
/// header, rather than relying on the HTTP layer to reject it later.
fn reject_control_chars(token: &str) -> anyhow::Result<()> {
if token.contains(char::is_control) {
anyhow::bail!("token contains control characters");
}
Ok(())
}
/// `now + secs`, or `None` on overflow.
pub(crate) fn expiry_after_seconds(secs: u64) -> Option<chrono::DateTime<chrono::Utc>> {
let secs = i64::try_from(secs).ok()?;
chrono::Utc::now().checked_add_signed(chrono::Duration::try_seconds(secs)?)
}
pub(crate) struct ParsedTokenOutput {
pub access_token: String,
pub refresh_token: Option<String>,
pub expires_at: Option<chrono::DateTime<chrono::Utc>>,
pub issuer: Option<String>,
}
/// Accepts a bare token or JSON `{access_token, expires_in, issuer, ...}`. A
/// non-zero exit, non-UTF-8 or empty stdout, an empty `access_token`, or
/// JSON-object output that is not a valid token payload are all errors, so a
/// malformed mint fails closed rather than putting garbage on the wire.
pub(crate) fn parse_token_output(
output: &std::process::Output,
) -> anyhow::Result<ParsedTokenOutput> {
if !output.status.success() {
anyhow::bail!("exited with {}", output.status);
}
let stdout = std::str::from_utf8(&output.stdout)
.map_err(|_| anyhow::anyhow!("produced non-UTF-8 output on stdout"))?
.trim();
if stdout.is_empty() {
anyhow::bail!("produced no output on stdout");
}
// Output that starts with `{` is meant to be a token payload: require it to
// parse and carry a non-empty access_token. Anything else is a bare token
// (JWTs and opaque tokens never start with `{`), so an error object like
// `{"error":"expired"}` can never be mistaken for a bearer.
if stdout.starts_with('{') {
let parsed: ExternalAuthOutput = serde_json::from_str(stdout)
.map_err(|e| anyhow::anyhow!("produced JSON that is not a token payload: {e}"))?;
let access_token = parsed.access_token.trim().to_owned();
if access_token.is_empty() {
anyhow::bail!("produced JSON with an empty access_token");
}
reject_control_chars(&access_token)?;
tracing::debug!(
has_refresh_token = parsed.refresh_token.is_some(),
expires_in = ?parsed.expires_in,
issuer = ?parsed.issuer,
"auth: parsed external provider output as JSON"
);
return Ok(ParsedTokenOutput {
access_token,
refresh_token: parsed.refresh_token,
expires_at: parsed.expires_in.and_then(expiry_after_seconds),
issuer: parsed
.issuer
.map(|i| i.trim().to_owned())
.filter(|i| !i.is_empty()),
});
}
reject_control_chars(stdout)?;
tracing::debug!(
stdout_len = stdout.len(),
"auth: treating output as bare token"
);
Ok(ParsedTokenOutput {
access_token: stdout.to_owned(),
refresh_token: None,
expires_at: None,
issuer: None,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn expiry_after_seconds_returns_none_on_overflow() {
assert_eq!(expiry_after_seconds(u64::MAX), None);
assert_eq!(expiry_after_seconds(u64::try_from(i64::MAX).unwrap()), None);
assert!(expiry_after_seconds(3600).is_some());
}
/// The provider path reads `refresh_token`, which the bare-token fallback
/// cannot carry; only JSON output does.
#[test]
fn parse_token_output_reads_refresh_token_from_json_only() {
let ok = |stdout: &str| std::process::Output {
status: std::process::Command::new("true").status().unwrap(),
stdout: stdout.as_bytes().to_vec(),
stderr: vec![],
};
let parsed =
parse_token_output(&ok(r#"{"access_token":"a","refresh_token":"r"}"#)).unwrap();
assert_eq!(parsed.access_token, "a");
assert_eq!(parsed.refresh_token.as_deref(), Some("r"));
assert_eq!(parse_token_output(&ok("bare")).unwrap().refresh_token, None);
}
/// JSON-shaped output must be a valid, non-empty token payload; a botched or
/// error payload fails closed instead of going on the wire as a bearer.
#[test]
fn parse_token_output_rejects_invalid_json_payloads() {
let ok = |stdout: &str| std::process::Output {
status: std::process::Command::new("true").status().unwrap(),
stdout: stdout.as_bytes().to_vec(),
stderr: vec![],
};
assert!(parse_token_output(&ok(r#"{"access_token":""}"#)).is_err());
assert!(parse_token_output(&ok(r#"{"access_token":" "}"#)).is_err());
assert!(parse_token_output(&ok(r#"{"error":"expired"}"#)).is_err());
assert!(parse_token_output(&ok("{not valid json}")).is_err());
// A JSON payload's access_token is trimmed of surrounding whitespace.
let parsed = parse_token_output(&ok("{\"access_token\":\" tok \"}")).unwrap();
assert_eq!(parsed.access_token, "tok");
// An interior control character is rejected on both paths, so a
// malformed token can never reach an HTTP header.
assert!(parse_token_output(&ok("{\"access_token\":\"tok\\ninjected\"}")).is_err());
assert!(parse_token_output(&ok("tok\ninjected")).is_err());
}
}

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

@ -2191,7 +2191,10 @@ extra_rule_dirs = ["/c/rules"]
let leaked: Vec<&RequirementSource> = r
.sources
.iter()
.filter(|s| matches!(s, RequirementSource::Settings { path } if path == &tempdir_claude))
.filter(|s| {
matches!(s, RequirementSource::Settings { path }
if path == &tempdir_claude)
})
.collect();
assert!(
leaked.is_empty(),

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

@ -766,7 +766,8 @@ mod tests {
.expect("first event should dispatch within 2s")
.expect("channel open");
assert!(
matches!(update, ConfigUpdate::ProjectMcpServersChanged { cwd: ref c } if *c == cwd),
matches!(update, ConfigUpdate::ProjectMcpServersChanged { cwd: ref c }
if *c == cwd),
"first project event must dispatch"
);
@ -790,7 +791,8 @@ mod tests {
.expect("changed content should dispatch within 2s")
.expect("channel open");
assert!(
matches!(update, ConfigUpdate::ProjectMcpServersChanged { cwd: ref c } if *c == cwd),
matches!(update, ConfigUpdate::ProjectMcpServersChanged { cwd: ref c }
if *c == cwd),
"changed project config must dispatch"
);

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

@ -2691,6 +2691,29 @@ fn config_layers_user_overrides_managed() {
Some(crate ::agent::config::TelemetryMode::Enabled), cfg.features.telemetry
);
}
/// A provider in a trusted disk layer resolves through the real
/// `ConfigLayers` → `effective_config_disk_only` → parse seam that the
/// direct-TOML parse tests bypass. (`ConfigLayers` has no project slot, so
/// a repo `.grok/config.toml` structurally cannot supply one.)
#[test]
fn auth_provider_honored_only_from_trusted_disk_layers() {
let layers = ConfigLayers {
managed: toml::from_str(
"[auth_provider.corp]\ncommand = \"/usr/local/bin/corp-token\"\n",
)
.unwrap(),
..Default::default()
};
let cfg = crate::agent::config::Config::new_from_toml_cfg(
&layers.effective_config_disk_only(),
)
.unwrap();
assert_eq!(
cfg.auth_providers.get("corp").map(| c | c.command.as_str()),
Some("/usr/local/bin/corp-token"),
"a provider in a trusted disk layer is honored"
);
}
/// REGRESSION: the real enterprise two-file merge —
/// `managed_config.toml` (proxy + BYO model host) layered with
/// `requirements.toml` (deployment key + S3 trace upload) via the actual

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

@ -1447,9 +1447,10 @@ mod official_source_tests {
assert_eq!(sources.len(), 1);
assert_eq!(sources[0].name, "my-plugins");
assert!(matches!(
&sources[0].kind,
xai_grok_plugin_marketplace::SourceKind::Local { path } if path == &dir
));
&sources[0].kind,
xai_grok_plugin_marketplace::SourceKind::Local { path }
if path == &dir
));
// The path must not be mangled into a git URL.
let raw = std::fs::read_to_string(&config_path).unwrap();
assert!(!raw.contains("git ="), "{raw}");

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

@ -73,10 +73,9 @@ pub struct InspectReport {
pub lsp_servers: Vec<LspServerEntry>,
pub config_sources: ConfigSources,
pub external_compat: ExternalCompatReport,
/// Warnings from `[model.*]` parsing.
/// Warnings from `[model.*]` and `[auth_provider.*]` parsing.
#[serde(skip_serializing_if = "Vec::is_empty")]
pub model_override_warnings:
Vec<crate::agent::config_model_override_parse::ModelOverrideWarning>,
pub config_warnings: Vec<crate::agent::config_model_override_parse::ConfigWarning>,
}
#[derive(Debug, Serialize)]
@ -380,9 +379,9 @@ async fn build_report(cwd: &Path) -> InspectReport {
}
let lsp = list_lsp_servers(cwd, &discovered_plugins);
let configs = list_config_sources(cwd);
let model_override_warnings = parsed_config
let config_warnings = parsed_config
.as_ref()
.map(|c| c.model_override_warnings.clone())
.map(|c| c.config_warnings.clone())
.unwrap_or_default();
InspectReport {
@ -405,7 +404,7 @@ async fn build_report(cwd: &Path) -> InspectReport {
lsp_servers: lsp,
config_sources: configs,
external_compat,
model_override_warnings,
config_warnings,
}
}
@ -1242,35 +1241,27 @@ fn disabled_compat_tags(
}
}
/// Renders the "Model Overrides" section of the human report; empty when
/// there are no warnings.
fn render_model_override_warnings(
warnings: &[crate::agent::config_model_override_parse::ModelOverrideWarning],
/// Renders the "Config Warnings" section of the human report; empty when
/// there are no warnings. Covers `[model.*]` overrides and the
/// `[auth_provider.*]` tables, which share the same warning channel.
fn render_config_warnings(
warnings: &[crate::agent::config_model_override_parse::ConfigWarning],
) -> String {
use std::fmt::Write as _;
if warnings.is_empty() {
return String::new();
}
let mut out = String::from("\n Model Overrides\n");
let _ = writeln!(
out,
" {TREE} {} warning(s) (models with invalid fields kept in catalog)",
warnings.len()
);
let mut out = String::from("\n Config Warnings\n");
let _ = writeln!(out, " {TREE} {} warning(s)", warnings.len());
for w in warnings {
let target = match w.model_key.as_deref() {
Some(key) => format!("[model.\"{key}\"]"),
None => "[model]".to_owned(),
};
match w.field.as_deref() {
Some(field) => {
let _ = writeln!(out, " {TREE} {target} {field} — {}", w.reason);
}
None => {
let _ = writeln!(out, " {TREE} {target} — {}", w.reason);
}
}
let field = w.field().map(|f| format!(" {f}")).unwrap_or_default();
let _ = writeln!(
out,
" {TREE} [{}]{field} — {}",
w.target.label(),
w.reason
);
}
out
}
@ -1545,10 +1536,7 @@ fn print_human(r: &InspectReport) {
println!(" {TREE} Project: (none)");
}
print!(
"{}",
render_model_override_warnings(&r.model_override_warnings)
);
print!("{}", render_config_warnings(&r.config_warnings));
print!("{}", render_harness_compatibility(&r.external_compat));
}
@ -1845,7 +1833,7 @@ mod tests {
/// Model-override warnings flow from an effective config through `Config`
/// to the human renderer and the JSON report.
#[test]
fn model_override_warnings_inspect_smoke() {
fn config_warnings_inspect_smoke() {
let effective: toml::Value = toml::from_str(
r#"
[model."grok-4.5"]
@ -1858,23 +1846,23 @@ mod tests {
)
.unwrap();
let cfg = crate::agent::config::Config::new_from_toml_cfg(&effective).unwrap();
let warnings = cfg.model_override_warnings;
let warnings = cfg.config_warnings;
assert!(
warnings
.iter()
.any(|w| w.field.as_deref() == Some("send_compactions_remaining")),
.any(|w| w.field() == Some("send_compactions_remaining")),
"duplicate alias should warn: {warnings:?}"
);
assert!(
warnings
.iter()
.any(|w| w.field.as_deref() == Some("reasoning_effort")),
.any(|w| w.field() == Some("reasoning_effort")),
"invalid enum should warn: {warnings:?}"
);
assert!(cfg.config_models.contains_key("grok-4.5"));
let human = render_model_override_warnings(&warnings);
assert!(human.contains("Model Overrides"), "{human}");
let human = render_config_warnings(&warnings);
assert!(human.contains("Config Warnings"), "{human}");
assert!(
human.contains("[model.\"grok-4.5\"] send_compactions_remaining"),
"{human}"
@ -1883,7 +1871,33 @@ mod tests {
human.contains("[model.\"grok-4.5\"] reasoning_effort"),
"{human}"
);
assert_eq!(render_model_override_warnings(&[]), "");
// Auth-provider warnings render under their own table syntax.
let provider_warning =
crate::agent::config_model_override_parse::ConfigWarning::auth_provider(
"litellm",
Some("command"),
crate::agent::config_model_override_parse::ConfigWarningKind::InvalidValue,
"missing or empty command".to_owned(),
);
let human = render_config_warnings(&[provider_warning]);
assert!(
human.contains("[auth_provider.\"litellm\"] command"),
"{human}"
);
// A dotted provider name renders whole; the field splits off the
// right.
let dotted = crate::agent::config_model_override_parse::ConfigWarning::auth_provider(
"corp.gateway",
Some("token_ttl_secs"),
crate::agent::config_model_override_parse::ConfigWarningKind::InvalidValue,
"at or below the refresh margin".to_owned(),
);
let human = render_config_warnings(&[dotted]);
assert!(
human.contains("[auth_provider.\"corp.gateway\"] token_ttl_secs"),
"{human}"
);
assert_eq!(render_config_warnings(&[]), "");
let json = serde_json::to_value(&warnings).unwrap();
let alias_warning = json
@ -1892,7 +1906,8 @@ mod tests {
.iter()
.find(|w| w["field"] == "send_compactions_remaining")
.expect("alias warning present in JSON");
assert_eq!(alias_warning["modelKey"], "grok-4.5");
assert_eq!(alias_warning["target"], "model");
assert_eq!(alias_warning["key"], "grok-4.5");
assert_eq!(alias_warning["kind"], "duplicate-alias");
assert!(
alias_warning["reason"]

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

@ -858,13 +858,14 @@ mod tests {
match start_result {
Ok(started) => {
assert!(matches!(
started,
ControlPayload::CpuProfileStarted {
svg_path,
frequency_hz: 200,
..
} if svg_path == output_path
));
started,
ControlPayload::CpuProfileStarted {
svg_path,
frequency_hz: 200,
..
}
if svg_path == output_path
));
let status = client
.send_control(ControlCommand::CpuProfileStatus)
@ -872,15 +873,16 @@ mod tests {
.unwrap()
.unwrap();
assert!(matches!(
status,
ControlPayload::CpuProfileStatus {
active: true,
stopping: false,
svg_path: Some(path),
frequency_hz: Some(200),
..
} if path == output_path
));
status,
ControlPayload::CpuProfileStatus {
active: true,
stopping: false,
svg_path: Some(path),
frequency_hz: Some(200),
..
}
if path == output_path
));
let stopped = client
.send_control(ControlCommand::StopCpuProfile)
@ -888,9 +890,10 @@ mod tests {
.unwrap()
.unwrap();
assert!(matches!(
stopped,
ControlPayload::CpuProfileStopped { svg_path, .. } if svg_path == output_path
));
stopped,
ControlPayload::CpuProfileStopped { svg_path, .. }
if svg_path == output_path
));
assert!(output_path.exists());
}
Err(error) => {
@ -987,15 +990,16 @@ mod tests {
.unwrap()
.unwrap();
assert!(matches!(
status,
ControlPayload::CpuProfileStatus {
active: false,
stopping: true,
svg_path: Some(path),
frequency_hz: Some(200),
..
} if path == output_path
));
status,
ControlPayload::CpuProfileStatus {
active: false,
stopping: true,
svg_path: Some(path),
frequency_hz: Some(200),
..
}
if path == output_path
));
let leader_info = client_b
.send_control(ControlCommand::GetLeaderInfo)
@ -1033,9 +1037,10 @@ mod tests {
let stopped = stop_task.await.unwrap().unwrap().unwrap();
assert!(matches!(
stopped,
ControlPayload::CpuProfileStopped { svg_path, .. } if svg_path == output_path
));
stopped,
ControlPayload::CpuProfileStopped { svg_path, .. }
if svg_path == output_path
));
assert_eq!(
stop_calls.lock().unwrap().as_slice(),
std::slice::from_ref(&output_path)

86
crates/codegen/xai-grok-shell/src/leader/protocol.rs
View file

@ -445,15 +445,16 @@ mod tests {
let received: ClientMessage = read_message(&mut server).await.unwrap();
assert!(matches!(
received,
ClientMessage::Control {
request_id,
command: ControlCommand::StartCpuProfile {
output: Some(output),
frequency_hz: Some(250),
},
} if request_id == "req-1" && output == "/tmp/profile.folded"
));
received,
ClientMessage::Control {
request_id,
command: ControlCommand::StartCpuProfile {
output: Some(output),
frequency_hz: Some(250),
},
}
if request_id == "req-1" && output == "/tmp/profile.folded"
));
}
#[tokio::test]
@ -539,21 +540,22 @@ mod tests {
let json = serde_json::to_string(&msg).unwrap();
let decoded: ServerMessage = serde_json::from_str(&json).unwrap();
assert!(matches!(
decoded,
ServerMessage::Registered {
client_id: 7,
ready: true,
leader_protocol_version: Some(LEADER_PROTOCOL_VERSION),
leader_binary_version: Some(_),
leader_capabilities: Some(LeaderCapabilities {
control_v1: true,
runtime_cpu_profile: true,
profile_formats,
workspace_exposure: true,
relaunch_v1: true,
}),
} if profile_formats == vec![ProfileArtifactFormat::Svg]
));
decoded,
ServerMessage::Registered {
client_id: 7,
ready: true,
leader_protocol_version: Some(LEADER_PROTOCOL_VERSION),
leader_binary_version: Some(_),
leader_capabilities: Some(LeaderCapabilities {
control_v1: true,
runtime_cpu_profile: true,
profile_formats,
workspace_exposure: true,
relaunch_v1: true,
}),
}
if profile_formats == vec![ProfileArtifactFormat::Svg]
));
}
#[test]
@ -637,14 +639,15 @@ mod tests {
let received: ClientMessage = read_message(&mut server).await.unwrap();
assert!(matches!(
received,
ClientMessage::Control {
request_id,
command: ControlCommand::WorkspaceStart { hub_url: Some(url), cwd },
} if request_id == "ws-1"
&& url == "wss://hub.example/v1/tools"
&& cwd == "/home/u/proj"
));
received,
ClientMessage::Control {
request_id,
command: ControlCommand::WorkspaceStart { hub_url: Some(url), cwd },
}
if request_id == "ws-1"
&& url == "wss://hub.example/v1/tools"
&& cwd == "/home/u/proj"
));
}
#[test]
@ -669,15 +672,16 @@ mod tests {
let json = r#"{"type":"workspace_status","state":"none","uptime_ms":0,"active_tool_calls":0,"pid":1}"#;
let decoded: ControlPayload = serde_json::from_str(json).unwrap();
assert!(matches!(
decoded,
ControlPayload::WorkspaceStatus {
state,
hub_url: None,
cwd: None,
sessions,
..
} if state == "none" && sessions.is_empty()
));
decoded,
ControlPayload::WorkspaceStatus {
state,
hub_url: None,
cwd: None,
sessions,
..
}
if state == "none" && sessions.is_empty()
));
}
#[test]

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

@ -3065,7 +3065,8 @@ mod tests {
assert!(
matches!(response, ServerMessage::ControlResult { request_id, result :
Ok(ControlPayload::CpuProfileStatus { active : false, stopping : false,
started_at : None, svg_path : None, frequency_hz : None, }), } if request_id
started_at : None, svg_path : None, frequency_hz : None, }), }
if request_id
== "status-1")
);
assert!(

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

@ -232,10 +232,11 @@ async fn fetch_managed_config(
token: &str,
source: ManagedConfigSource,
max_attempts: u32,
echo_principal: Option<&str>,
) -> Result<ManagedConfigResponse, ManagedConfigError> {
crate::http::send_with_retry_escaping_pool(
move |client: reqwest::Client| async move {
fetch_managed_config_once(&client, url, token, source).await
fetch_managed_config_once(&client, url, token, source, echo_principal).await
},
max_attempts,
|e: &ManagedConfigError| e.is_retryable(),
@ -324,14 +325,25 @@ async fn fetch_managed_config_once(
url: &str,
token: &str,
source: ManagedConfigSource,
echo_principal: Option<&str>,
) -> Result<ManagedConfigResponse, ManagedConfigError> {
let resp = match client
let mut request = client
.get(url)
.header("Authorization", format!("Bearer {}", token))
.timeout(std::time::Duration::from_secs(15))
.send()
.await
.timeout(std::time::Duration::from_secs(15));
// Replay-probe echo (telemetry only). Skip on invalid HeaderValue so a
// corrupt sidecar never bricks the fetch (echo is fail-open).
if let Some(nonce) = xai_grok_config::signed_policy::stored_envelope_nonce(
&crate::util::grok_home::grok_home(),
echo_principal,
) && let Ok(value) = reqwest::header::HeaderValue::from_str(&nonce)
{
request = request.header(
xai_grok_config::signed_policy::MANAGED_CONFIG_NONCE_ECHO_HEADER,
value,
);
}
let resp = match request.send().await {
Ok(r) if r.status().is_success() => r,
Ok(r) => {
let status = r.status().as_u16();
@ -544,7 +556,11 @@ async fn fetch_for_principal(
if let Some(dk) = resolve_deployment_key() {
let source = ManagedConfigSource::DeploymentKey;
match fetch_managed_config(&url, &dk, source, max_attempts).await {
// Echo binds to the deployment this key last synced (marker-bound; None
// on first sync or after a key rotation — then there is nothing to echo).
let echo_principal = crate::config::managed_deployment_id(&deployment_key_fingerprint(&dk));
match fetch_managed_config(&url, &dk, source, max_attempts, echo_principal.as_deref()).await
{
// A rejected dk (stale env/config) must not starve a valid team
// sign-in: fall through. Network/5xx do NOT — same unreachable
// server, double the latency for nothing.
@ -569,6 +585,7 @@ async fn fetch_for_principal(
&auth.key,
ManagedConfigSource::TeamOauth,
max_attempts,
auth.team_id.as_deref(),
)
.await?;
return Ok(FetchedConfig::Team {

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

@ -410,6 +410,7 @@ fn served_principal_prefers_deployment_id() {
requirements: None,
fail_closed: false,
expires_at: 0,
nonce: String::new(),
key_id: "v1".into(),
};
assert_eq!(

14
crates/codegen/xai-grok-shell/src/remote/sync.rs
View file

@ -43,6 +43,20 @@ pub struct RemoteSync {
}
impl RemoteSync {
#[cfg(test)]
pub(crate) fn test_observer() -> (Self, mpsc::UnboundedReceiver<acp::SessionNotification>) {
let (tx, mut rx) = mpsc::unbounded_channel();
let (observed_tx, observed_rx) = mpsc::unbounded_channel();
tokio::spawn(async move {
while let Some(message) = rx.recv().await {
if let SyncMsg::Queue(notification) = message {
let _ = observed_tx.send(*notification);
}
}
});
(Self { tx }, observed_rx)
}
/// Metadata is included on every flush to keep the backend session row current.
pub(crate) fn new(
session_id: String,

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

@ -560,6 +560,14 @@ impl PreparedToolCall {
#[cfg(test)]
pub(crate) use crate::session::streaming_capture::STREAMING_CAPTURE_MAX_BYTES;
pub(crate) use crate::session::streaming_capture::StreamingTurnCapture;
/// One memoized model's auth state, keyed by model id; see
/// [`SessionActor::model_auth_memo`] for the invalidation contract.
#[derive(Clone)]
pub(crate) struct ModelAuthMemo {
pub(crate) model_id: String,
pub(crate) facts: crate::agent::config::ModelAuthFacts,
pub(crate) provider: Option<crate::auth::AuthProviderRef>,
}
/// Phase 3: Post-flight handling after dispatch (inline in execute_tool_calls for now).
pub(crate) struct SessionActor {
pub(crate) session_info: SessionInfo,
@ -569,10 +577,17 @@ pub(crate) struct SessionActor {
/// fresh, isolated handle seeded once at spawn (frozen for their lifetime).
/// `None` until the agent has selected a method.
pub(crate) auth_method_id: crate::agent::auth_method::SharedAuthMethodId,
/// Memoized per-model auth facts, keyed by model id — see
/// [`SessionActor::model_auth_facts`].
pub(crate) model_auth_facts:
std::cell::RefCell<Option<(String, crate::agent::config::ModelAuthFacts)>>,
/// Memoized per-model auth state, read through
/// [`SessionActor::model_auth_facts`] and
/// [`SessionActor::model_auth_provider`].
///
/// A fresh `Unknown` (config currently unparseable) falls back to the
/// last definite value for the same model rather than demoting a live
/// session to non-refreshable api-key mode. Because a config edit can
/// turn the selected model into a per-model BYOK model without changing
/// its id, keying on the id alone is insufficient: each model/credential
/// chokepoint must clear this memo (`replace(None)`).
pub(crate) model_auth_memo: std::cell::RefCell<Option<ModelAuthMemo>>,
/// 401-attribution callback. Joined with the bearer the
/// sampler sends on the wire to emit an `auth 401 attribution`
/// event at each of the six `OaiCompatClient` 401 arms in

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

@ -74,7 +74,7 @@ impl SessionActor {
alpha_test_key: existing.alpha_test_key,
client_version: sampling_config.client_version.clone(),
});
self.model_auth_facts.replace(None);
self.invalidate_model_auth_memo();
self.signals_handle()
.record_model_usage(&sampling_config.model);
if apply_prompt_override && !skip_prompt_rewrite {

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

File diff suppressed because it is too large Load diff

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

@ -150,34 +150,129 @@ impl SessionActor {
let plan_active = self.plan_mode.lock().is_active();
filter_cursor_tools_by_plan_mode(defs, plan_active)
}
/// Memoized per-model [`ModelAuthFacts`](crate::agent::config::ModelAuthFacts),
/// keyed by `model_id`.
///
/// A fresh `Unknown` (config currently unparseable) falls back to the last
/// definite value for the same `model_id` rather than demoting a live session
/// to non-refreshable api-key mode. Because a config edit can turn the
/// currently-selected model into a per-model BYOK model without changing
/// `model_id`, keying on `model_id` alone is insufficient — each
/// model/credential chokepoint must clear this memo (`replace(None)`).
pub(super) fn model_auth_facts(&self, model_id: &str) -> crate::agent::config::ModelAuthFacts {
self.model_auth_state(model_id).0
}
pub(super) fn model_auth_provider(
&self,
model_id: &str,
) -> Option<crate::auth::AuthProviderRef> {
self.model_auth_state(model_id).1
}
/// Drop the memoized per-model auth state; see [`Self::model_auth_memo`]
/// for why each model/credential chokepoint must call this.
pub(crate) fn invalidate_model_auth_memo(&self) {
self.model_auth_memo.replace(None);
}
/// Reads and populates [`Self::model_auth_memo`]; a fresh `Unknown`
/// falls back to the last definite entry (see the field's contract).
fn model_auth_state(
&self,
model_id: &str,
) -> (
crate::agent::config::ModelAuthFacts,
Option<crate::auth::AuthProviderRef>,
) {
use crate::agent::auth_method::ModelByok;
if let Some((cached_id, facts)) = self.model_auth_facts.borrow().as_ref()
&& cached_id == model_id
&& facts.byok != ModelByok::Unknown
use crate::session::acp_session::ModelAuthMemo;
if let Some(memo) = self.model_auth_memo.borrow().as_ref()
&& memo.model_id == model_id
&& memo.facts.byok != ModelByok::Unknown
{
return *facts;
return (memo.facts, memo.provider.clone());
}
let fresh = crate::agent::config::resolve_model_auth_facts(model_id);
let (fresh, provider) =
crate::agent::config::resolve_model_auth_facts_and_provider(model_id);
if fresh.byok == ModelByok::Unknown {
if let Some((cached_id, facts)) = self.model_auth_facts.borrow().as_ref()
&& cached_id == model_id
if let Some(memo) = self.model_auth_memo.borrow().as_ref()
&& memo.model_id == model_id
{
return *facts;
return (memo.facts, memo.provider.clone());
}
return fresh;
return (fresh, provider);
}
*self.model_auth_facts.borrow_mut() = Some((model_id.to_string(), fresh));
fresh
*self.model_auth_memo.borrow_mut() = Some(ModelAuthMemo {
model_id: model_id.to_string(),
facts: fresh,
provider: provider.clone(),
});
(fresh, provider)
}
/// The single writer of a provider mint/rotation into chat-state credentials.
async fn set_chat_api_key(&self, new_key: String) {
let mut creds = self.chat_state_handle.get_credentials().await;
creds.api_key = Some(new_key);
self.chat_state_handle.update_credentials(creds);
}
/// Pre-turn arm for a provider-backed model: mint on a cold cache,
/// re-mint near expiry, and adopt a rotation chat-state missed. No-op
/// when `current_key` is already the fresh cached token.
async fn refresh_provider_token_pre_turn(
&self,
provider: &crate::auth::AuthProviderRef,
current_key: Option<&str>,
model_id: &str,
) {
match provider.ensure_fresh_token(current_key).await {
crate::auth::ProviderRefreshOutcome::Rotated(new_key) => {
tracing::info!(
model = % model_id, provider = % provider.name, cold = current_key
.is_none(), "auth provider token rotated pre-turn"
);
self.set_chat_api_key(new_key).await;
}
crate::auth::ProviderRefreshOutcome::Unchanged => {}
crate::auth::ProviderRefreshOutcome::MintFailed => {
tracing::warn!(
session_id = % self.session_info.id.0, provider = % provider.name,
model = % model_id, "auth provider pre-turn refresh failed"
);
xai_grok_telemetry::unified_log::warn(
"auth provider pre-turn refresh failed",
Some(self.session_info.id.0.as_ref()),
Some(serde_json::json!(
{ "provider" : provider.name, "model" : model_id, "cold" :
current_key.is_none(), }
)),
);
}
crate::auth::ProviderRefreshOutcome::Unusable => {}
}
}
/// 401 arm for a provider-backed model: re-run the helper once and
/// resubmit. A missing key means the cold mint failed and the request
/// went out unauthenticated, so mint instead. Returns `false` when the
/// fresh-mint guard blocked the re-run or the helper failed; the 401
/// then surfaces as a terminal error.
async fn try_provider_401_recovery(&self, provider: &crate::auth::AuthProviderRef) -> bool {
let rejected_key = self.chat_state_handle.get_credentials().await.api_key;
let recovered = match rejected_key {
Some(ref rejected_key) => provider.recover_rejected_token(rejected_key).await,
None => provider.ensure_fresh_token(None).await.rotated(),
};
let Some(new_key) = recovered else {
tracing::warn!(
session_id = % self.session_info.id.0, provider = % provider.name,
"auth recovery: sampler 401, provider re-mint declined or failed"
);
xai_grok_telemetry::unified_log::warn(
"auth recovery: sampler 401, provider re-mint declined or failed",
Some(self.session_info.id.0.as_ref()),
Some(serde_json::json!({ "provider" : provider.name })),
);
return false;
};
tracing::info!(
session_id = % self.session_info.id.0, provider = % provider.name,
"auth recovery: sampler 401, auth provider re-mint, retrying"
);
xai_grok_telemetry::unified_log::info(
"auth recovery: sampler 401, auth provider re-mint, retrying",
Some(self.session_info.id.0.as_ref()),
None,
);
self.set_chat_api_key(new_key).await;
true
}
/// Gate inputs for `model_id` routed to `base_url`. See
/// [`crate::agent::auth_method::session_token_auth_gate`] for the rationale
@ -642,17 +737,23 @@ impl SessionActor {
.data(detailed_message);
return Err(acp_err);
}
let (failed_model_id, failed_base_url) = self
.chat_state_handle
.get_sampling_config()
.await
.map(|c| (c.model, c.base_url))
.unwrap_or_default();
let auth_provider =
if matches!(error.kind, SamplingErrorKind::Auth) || error.status_code == Some(401) {
self.model_auth_provider(&failed_model_id)
} else {
None
};
let auth_recovery_eligible = matches!(error.kind, SamplingErrorKind::Auth) && {
let (model_id, base_url) = self
.chat_state_handle
.get_sampling_config()
.await
.map(|c| (c.model, c.base_url))
.unwrap_or_default();
let gate = self.auth_gate(&model_id, &base_url);
let gate = self.auth_gate(&failed_model_id, &failed_base_url);
let eligible = gate.active();
self.log_auth_gate_unknown("handle_sampling_failure", gate, &base_url);
if !eligible {
self.log_auth_gate_unknown("handle_sampling_failure", gate, &failed_base_url);
if !eligible && auth_provider.is_none() {
tracing::warn!(
session_id = % self.session_info.id.0, is_session_based = gate
.is_session_based, model_byok = gate.model_byok.as_str(),
@ -672,7 +773,14 @@ impl SessionActor {
}
eligible
};
if !matches!(error.kind, SamplingErrorKind::Auth) && error.status_code == Some(401) {
debug_assert!(
!(auth_recovery_eligible && auth_provider.is_some()),
"a provider-backed model must not be session-recovery-eligible"
);
if !matches!(error.kind, SamplingErrorKind::Auth)
&& error.status_code == Some(401)
&& auth_provider.is_none()
{
xai_grok_telemetry::unified_log::warn(
"auth recovery: sampler 401 not eligible (non-auth error kind)",
Some(self.session_info.id.0.as_ref()),
@ -735,6 +843,12 @@ impl SessionActor {
None,
);
}
if let Some(ref provider) = auth_provider
&& self.try_provider_401_recovery(provider).await
{
self.prepare_sampler_for_turn().await;
return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit);
}
if matches!(error.kind, SamplingErrorKind::IdleTimeout) {
self.signals_handle().record_idle_timeout();
}
@ -807,6 +921,14 @@ impl SessionActor {
let mut msg = format!("{detailed_message}\n");
msg.push_str(&format!("\n Model: {current_model}"));
msg.push_str(&format!("\n Auth: {auth_mode_str}"));
if let Some(ref provider) = auth_provider {
msg.push_str(
&format!(
"\n Provider: [auth_provider.{}] (check the provider command and the debug log)",
provider.name
),
);
}
msg.push_str(&format!("\n Version: {client_version}"));
if available.is_empty() {
msg.push_str("\n Available: (none)");
@ -975,6 +1097,15 @@ impl SessionActor {
.await
.map(|c| c.model)
.unwrap_or_default();
if let Some(provider) = self.model_auth_provider(&current_model_id) {
self.refresh_provider_token_pre_turn(
&provider,
current_key.as_deref(),
&current_model_id,
)
.await;
return;
}
let Some(ref key) = current_key else { return };
if !is_jwt_expired_or_near(key, REFRESH_THRESHOLD) {
if let Some(exp) = parse_jwt_expiration(key) {

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

@ -1151,7 +1151,7 @@ pub(crate) async fn spawn_session_actor(
let session = Arc::new_cyclic(|weak: &std::sync::Weak<SessionActor>| SessionActor {
session_info: session_info.clone(),
auth_method_id,
model_auth_facts: std::cell::RefCell::new(None),
model_auth_memo: std::cell::RefCell::new(None),
attribution_callback,
auth_manager,
state,

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

@ -487,7 +487,8 @@ mod stop_gate_snapshot_tests {
]);
assert!(
matches!(&results[0], HookRunResult::Success { hook_name, .. } if hook_name == "gate"),
matches!(&results[0], HookRunResult::Success { hook_name, .. }
if hook_name == "gate"),
"a discarded decision must read as success, got {:?}",
results[0]
);

5
crates/codegen/xai-grok-shell/src/session/acp_session_impl/types.rs
View file

@ -19,8 +19,9 @@ pub(crate) enum SamplerFailureRecovery {
/// Compaction ran. The turn loop should rebuild the request from
/// the compacted conversation and resubmit.
CompactAndResubmit,
/// Auth 401 recovery succeeded (devbox re-mint or OIDC refresh).
/// The turn loop should resubmit once with the fresh token.
/// Auth 401 recovery succeeded (devbox re-mint, OIDC refresh, or auth
/// provider re-mint). The turn loop should resubmit once with the
/// fresh token.
RefreshAuthAndResubmit,
}

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

@ -243,11 +243,9 @@ async fn sampler_401_with_api_key_auth_skips_refresh_and_surfaces_error() {
.await;
}
/// Per-turn pre-flight refresh dispatches on `AuthManager`'s
/// `TokenType`, not `creds.auth_type`. Pins that a stale
/// When `creds.auth_type` is `ApiKey` (BYOK model), the pre-flight
/// refresh must NOT fire — the model's own API key must not be
/// overwritten by the session JWT.
/// Per-turn pre-flight refresh must not fire when `creds.auth_type` is
/// `ApiKey` (a BYOK model): the model's own API key must not be overwritten
/// by the session JWT.
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial(attribution_emit_count)]
async fn pre_flight_refresh_skips_api_key_auth_type() {
@ -659,12 +657,8 @@ async fn no_legacy_hint_for_oidc_auth() {
.await;
}
// Regression: a live OIDC session whose `creds.auth_type` has
// transiently collapsed to `ApiKey` (session-token cache miss + `XAI_API_KEY`)
// must still drive the live bearer resolver, be eligible for 401 retry, and get
// its stale `api_key` healed — the gate keys off the stable `auth_method_id`,
// not the collapsible `auth_type`.
// Regression group: a live session whose `auth_type` transiently reads `ApiKey`
// must still recover, because the gate keys off the stable `auth_method_id`.
#[test]
fn session_token_auth_gate_truth_table() {
use crate::agent::auth_method::{ModelByok, session_token_auth_gate as gate};
@ -904,13 +898,13 @@ async fn session_born_on_api_key_recovers_after_oidc_login_without_restart() {
.await;
}
// Per-model BYOK memo (`SessionActor::model_auth_facts`): a definite cached
// Per-model BYOK memo (`SessionActor::model_auth_memo`): a definite cached
// status is served without recomputing, and the memo keys on `model_id`.
/// The cache-hit branch is what lets a later config parse failure (`Unknown`)
/// fall back to the last-known-good status.
#[tokio::test(flavor = "current_thread")]
async fn model_auth_facts_memo_serves_cached_status_and_keys_on_model() {
async fn model_auth_memo_serves_cached_status_and_keys_on_model() {
use crate::agent::auth_method::ModelByok;
use crate::agent::config::ModelAuthFacts;
let local = tokio::task::LocalSet::new();
@ -924,13 +918,16 @@ async fn model_auth_facts_memo_serves_cached_status_and_keys_on_model() {
)
.await;
actor.model_auth_facts.replace(Some((
"model-a".to_string(),
ModelAuthFacts {
byok: ModelByok::Byok,
auth_scheme: Default::default(),
},
)));
actor
.model_auth_memo
.replace(Some(crate::session::acp_session::ModelAuthMemo {
model_id: "model-a".to_string(),
facts: ModelAuthFacts {
byok: ModelByok::Byok,
auth_scheme: Default::default(),
},
provider: None,
}));
// Cache hit: served without consulting config.
assert_eq!(actor.model_auth_facts("model-a").byok, ModelByok::Byok);
@ -965,13 +962,16 @@ async fn reconstruct_full_config_no_bearer_resolver_for_byok_model_on_session_me
.await
.map(|c| c.model)
.unwrap_or_default();
actor.model_auth_facts.replace(Some((
model,
ModelAuthFacts {
byok: ModelByok::Byok,
auth_scheme: Default::default(),
},
)));
actor
.model_auth_memo
.replace(Some(crate::session::acp_session::ModelAuthMemo {
model_id: model,
facts: ModelAuthFacts {
byok: ModelByok::Byok,
auth_scheme: Default::default(),
},
provider: None,
}));
let cfg = actor.reconstruct_full_config().await;
@ -1010,13 +1010,16 @@ async fn set_session_model_invalidates_byok_memo_for_same_model_id() {
.map(|c| c.model)
.unwrap_or_default();
actor.model_auth_facts.replace(Some((
model.clone(),
ModelAuthFacts {
byok: ModelByok::NotByok,
auth_scheme: Default::default(),
},
)));
actor
.model_auth_memo
.replace(Some(crate::session::acp_session::ModelAuthMemo {
model_id: model.clone(),
facts: ModelAuthFacts {
byok: ModelByok::NotByok,
auth_scheme: Default::default(),
},
provider: None,
}));
// Switch to the same model_id, now a per-model BYOK model on a
// third-party endpoint.
@ -1054,10 +1057,330 @@ async fn set_session_model_invalidates_byok_memo_for_same_model_id() {
.await;
assert!(
actor.model_auth_facts.borrow().is_none(),
actor.model_auth_memo.borrow().is_none(),
"a model switch must invalidate the per-model BYOK memo so the next \
reconstruct recomputes under the current config"
);
})
.await;
}
use crate::auth::test_counting_provider as counting_provider;
/// Seed the per-model memo so `model_auth_provider` resolves without a
/// config load.
async fn seed_provider_memo(actor: &Arc<SessionActor>, provider: crate::auth::AuthProviderRef) {
let model = actor
.chat_state_handle
.get_sampling_config()
.await
.map(|c| c.model)
.unwrap_or_default();
actor
.model_auth_memo
.replace(Some(crate::session::acp_session::ModelAuthMemo {
model_id: model,
facts: crate::agent::config::ModelAuthFacts {
byok: crate::agent::auth_method::ModelByok::Byok,
auth_scheme: Default::default(),
},
provider: Some(provider),
}));
}
/// Regression: switching from a provider-backed model to a first-party model
/// must drop the minted provider token from the chat credentials, so it can
/// never ride a later request to `api.x.ai`. Mirrors the forward direction in
/// `set_session_model_invalidates_byok_memo_for_same_model_id`.
#[tokio::test(flavor = "current_thread")]
async fn switch_to_first_party_model_drops_minted_provider_token() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let dir = tempfile::tempdir().unwrap();
let provider = counting_provider("hall-pass", dir.path());
let token = provider.ensure_fresh_token(None).await.rotated().unwrap();
assert_eq!(token, "tok-1");
let (actor, _rx) =
make_actor_with_auth_and_credentials(None, xai_chat_state::AuthType::ApiKey, token)
.await;
seed_provider_memo(&actor, provider).await;
let model = actor
.chat_state_handle
.get_sampling_config()
.await
.map(|c| c.model)
.unwrap_or_default();
let cfg = xai_grok_sampler::SamplerConfig {
api_key: Some("session-jwt".to_string()),
base_url: "https://api.x.ai/v1".to_string(),
model,
max_completion_tokens: None,
temperature: None,
top_p: None,
api_backend: crate::sampling::ApiBackend::ChatCompletions,
auth_scheme: Default::default(),
extra_headers: Default::default(),
context_window: 256_000,
client_version: None,
force_http1: false,
max_retries: None,
stream_tool_calls: false,
idle_timeout_secs: None,
client_identifier: None,
reasoning_effort: None,
deployment_id: None,
user_id: None,
origin_client: None,
attribution_callback: None,
bearer_resolver: None,
supports_backend_search: false,
compactions_remaining: None,
compaction_at_tokens: None,
doom_loop_recovery: None,
header_injector: None,
};
let _ = actor
.handle_set_session_model(cfg, false, false, true, 85)
.await;
let creds = actor.chat_state_handle.get_credentials().await;
assert_eq!(
creds.api_key.as_deref(),
Some("session-jwt"),
"switching to a first-party model must install the session credential, \
not the minted provider token"
);
})
.await;
}
/// Arm 4c: a 401 on a provider-backed model re-mints once and resubmits.
#[tokio::test(flavor = "current_thread")]
async fn sampler_401_on_provider_model_remints_and_resubmits() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let dir = tempfile::tempdir().unwrap();
let provider = counting_provider("test-4c-recover", dir.path());
let token = provider.ensure_fresh_token(None).await.rotated().unwrap();
assert_eq!(token, "tok-1");
let (actor, _rx) =
make_actor_with_auth_and_credentials(None, xai_chat_state::AuthType::ApiKey, token)
.await;
seed_provider_memo(&actor, provider).await;
crate::auth::test_backdate_provider_mint(
"test-4c-recover",
std::time::Duration::from_secs(60),
);
let result = actor.handle_sampling_failure(auth_error()).await;
assert!(
matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)),
"provider 401 must re-mint and resubmit"
);
let creds = actor.chat_state_handle.get_credentials().await;
assert_eq!(
creds.api_key.as_deref(),
Some("tok-2"),
"chat-state credentials must carry the re-minted token"
);
})
.await;
}
/// Arm 4c also fires for a bare 401 that did not classify as `Auth`-kind.
#[tokio::test(flavor = "current_thread")]
async fn sampler_non_auth_kind_401_on_provider_model_still_recovers() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let dir = tempfile::tempdir().unwrap();
let provider = counting_provider("test-4c-non-auth-kind", dir.path());
let token = provider.ensure_fresh_token(None).await.rotated().unwrap();
let (actor, _rx) =
make_actor_with_auth_and_credentials(None, xai_chat_state::AuthType::ApiKey, token)
.await;
seed_provider_memo(&actor, provider).await;
crate::auth::test_backdate_provider_mint(
"test-4c-non-auth-kind",
std::time::Duration::from_secs(60),
);
let mut error = auth_error();
error.kind = xai_grok_sampler::SamplingErrorKind::Api;
let result = actor.handle_sampling_failure(error).await;
assert!(
matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)),
"a non-Auth-kind 401 on a provider model must still recover via 4c"
);
let creds = actor.chat_state_handle.get_credentials().await;
assert_eq!(creds.api_key.as_deref(), Some("tok-2"));
})
.await;
}
/// A 401 on a request that went out with no key mints instead of
/// recovering.
#[tokio::test(flavor = "current_thread")]
async fn sampler_401_with_no_key_on_provider_model_mints_and_resubmits() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let dir = tempfile::tempdir().unwrap();
let provider = counting_provider("test-4c-no-key", dir.path());
let (actor, _rx) = make_actor_with_auth_and_credentials(
None,
xai_chat_state::AuthType::ApiKey,
"placeholder".to_string(),
)
.await;
let mut creds = actor.chat_state_handle.get_credentials().await;
creds.api_key = None;
actor.chat_state_handle.update_credentials(creds);
seed_provider_memo(&actor, provider).await;
let result = actor.handle_sampling_failure(auth_error()).await;
assert!(
matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)),
"an unauthenticated 401 on a provider model must mint and resubmit"
);
let creds = actor.chat_state_handle.get_credentials().await;
assert_eq!(creds.api_key.as_deref(), Some("tok-1"));
})
.await;
}
/// A provider model's 401 goes through the provider, never the session
/// refresher (4a/4b vs 4c exclusivity). The actor uses a session-based method,
/// so the gate would be active for a non-BYOK model; the BYOK memo is what
/// shadows it, which is the invariant under test.
#[tokio::test(flavor = "current_thread")]
async fn sampler_401_on_provider_model_never_refreshes_session() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let dir = tempfile::tempdir().unwrap();
let provider = counting_provider("test-4c-exclusive", dir.path());
let token = provider.ensure_fresh_token(None).await.rotated().unwrap();
let called = Arc::new(AtomicBool::new(false));
let refresher: Arc<dyn crate::auth::refresh::TokenRefresher> =
Arc::new(AlwaysSucceedRefresher {
called: called.clone(),
});
let (_dir, am) = auth_manager_with_refresher(refresher);
let (actor, _rx) = make_actor_with_method_and_credentials(
Some(am),
"cached_token",
xai_chat_state::AuthType::SessionToken,
token,
)
.await;
seed_provider_memo(&actor, provider).await;
crate::auth::test_backdate_provider_mint(
"test-4c-exclusive",
std::time::Duration::from_secs(60),
);
let result = actor.handle_sampling_failure(auth_error()).await;
assert!(
matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)),
"the provider arm must recover"
);
assert!(
!called.load(Ordering::SeqCst),
"session refresh must never fire for a provider-backed model"
);
let creds = actor.chat_state_handle.get_credentials().await;
assert_eq!(creds.api_key.as_deref(), Some("tok-2"));
})
.await;
}
/// The pre-turn mirror of the exclusivity test: a cold cache mints the
/// provider token into chat-state, and the session refresher never fires. The
/// actor uses a session-based method, so the gate would be active for a
/// non-BYOK model; the BYOK memo is what keeps the refresher silent.
#[tokio::test(flavor = "current_thread")]
async fn pre_turn_on_provider_model_never_installs_session_token() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let dir = tempfile::tempdir().unwrap();
let provider = counting_provider("test-preturn-exclusive", dir.path());
let called = Arc::new(AtomicBool::new(false));
let refresher: Arc<dyn crate::auth::refresh::TokenRefresher> =
Arc::new(AlwaysSucceedRefresher {
called: called.clone(),
});
let (_dir, am) = auth_manager_with_refresher(refresher);
let (actor, _rx) = make_actor_with_method_and_credentials(
Some(am),
"cached_token",
xai_chat_state::AuthType::SessionToken,
"placeholder".to_string(),
)
.await;
// Cold cache: no key on the wire yet.
let mut creds = actor.chat_state_handle.get_credentials().await;
creds.api_key = None;
actor.chat_state_handle.update_credentials(creds);
seed_provider_memo(&actor, provider).await;
actor.refresh_token_if_expired().await;
let creds = actor.chat_state_handle.get_credentials().await;
assert_eq!(
creds.api_key.as_deref(),
Some("tok-1"),
"the cold pre-turn hook must mint the provider token"
);
assert!(
!called.load(Ordering::SeqCst),
"the session refresher must never fire for a provider-backed model"
);
})
.await;
}
/// A token rejected moments after mint surfaces the 401 (fresh-mint
/// guard).
#[tokio::test(flavor = "current_thread")]
async fn sampler_401_on_fresh_provider_token_surfaces_error() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let dir = tempfile::tempdir().unwrap();
let provider = counting_provider("test-4c-guard", dir.path());
let token = provider.ensure_fresh_token(None).await.rotated().unwrap();
let (actor, _rx) = make_actor_with_auth_and_credentials(
None,
xai_chat_state::AuthType::ApiKey,
token.clone(),
)
.await;
seed_provider_memo(&actor, provider).await;
let result = actor.handle_sampling_failure(auth_error()).await;
assert!(
result.is_err(),
"a fresh-minted rejected token must surface the 401, not loop"
);
let creds = actor.chat_state_handle.get_credentials().await;
assert_eq!(
creds.api_key.as_deref(),
Some(token.as_str()),
"credentials must be unchanged when the guard blocks the re-mint"
);
})
.await;
}

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

@ -109,7 +109,7 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
let actor = Arc::new(SessionActor {
session_info,
auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None),
model_auth_memo: std::cell::RefCell::new(None),
attribution_callback: None,
auth_manager: None,
state: TokioMutex::new(State {
@ -561,7 +561,7 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
let actor = Arc::new(SessionActor {
session_info: session_info.clone(),
auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None),
model_auth_memo: std::cell::RefCell::new(None),
attribution_callback: None,
auth_manager: None,
state: TokioMutex::new(State {
@ -833,7 +833,7 @@ async fn cancel_running_task_teardown_clears_running_and_pending_work() {
cwd: cwd.as_str().to_string(),
},
auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None),
model_auth_memo: std::cell::RefCell::new(None),
attribution_callback: None,
auth_manager: None,
state,
@ -2065,7 +2065,7 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() {
cwd: cwd.as_str().to_string(),
},
auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None),
model_auth_memo: std::cell::RefCell::new(None),
attribution_callback: None,
auth_manager: None,
state,
@ -2325,10 +2325,11 @@ async fn skill_reminder_deferred_while_turn_running_flushed_when_idle() {
.iter()
.filter(|item| {
matches!(
item, ConversationItem::User(u) if u.content.iter().any(| p |
matches!(p, xai_grok_sampling_types::ContentPart::Text { text } if
text.contains("pdf-tools")))
)
item, ConversationItem::User(u) if u.content.iter().any(| p |
matches!(p, xai_grok_sampling_types::ContentPart::Text { text }
if
text.contains("pdf-tools")))
)
})
.count()
}

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

@ -2714,6 +2714,7 @@ fn catalog_with(
info,
api_key: None,
env_key: None,
auth_provider: None,
api_base_url: None,
},
);

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

@ -127,7 +127,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
},
attribution_callback: None,
auth_method_id: test_auth_method_id("cached_token"),
model_auth_facts: std::cell::RefCell::new(None),
model_auth_memo: std::cell::RefCell::new(None),
auth_manager: {
let dir = tempfile::tempdir().unwrap();
let mgr = std::sync::Arc::new(crate::auth::AuthManager::new(

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

@ -70,7 +70,7 @@ async fn create_test_actor(
},
rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(),
auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None),
model_auth_memo: std::cell::RefCell::new(None),
attribution_callback: None,
auth_manager: None,
state,
@ -503,7 +503,7 @@ async fn create_test_actor_with_memory(
},
rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(),
auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None),
model_auth_memo: std::cell::RefCell::new(None),
attribution_callback: None,
auth_manager: None,
state,
@ -1255,7 +1255,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
},
rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(),
auth_method_id: test_auth_method_id("cached_token"),
model_auth_facts: std::cell::RefCell::new(None),
model_auth_memo: std::cell::RefCell::new(None),
auth_manager: {
let dir = tempfile::tempdir().unwrap();
let mgr = std::sync::Arc::new(crate::auth::AuthManager::new(

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

@ -38,6 +38,7 @@ fn detector_entry(
info,
api_key: None,
env_key: None,
auth_provider: None,
api_base_url: None,
}
}

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

@ -123,7 +123,7 @@ async fn create_test_actor_with_memory(
cwd: cwd.as_str().to_string(),
},
auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None),
model_auth_memo: std::cell::RefCell::new(None),
attribution_callback: None,
auth_manager: None,
state,

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

@ -180,6 +180,7 @@ async fn build_session_info_sources_show_model_fingerprint_from_catalog() {
info: ModelInfo::fallback("test"),
api_key: None,
env_key: None,
auth_provider: None,
api_base_url: None,
};
entry.info.show_model_fingerprint = false;

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

@ -77,7 +77,7 @@ pub(super) async fn make_replay_send_update_fixture() -> ReplaySendUpdateFixture
cwd: cwd.as_str().to_string(),
},
auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None),
model_auth_memo: std::cell::RefCell::new(None),
attribution_callback: None,
auth_manager: None,
state,

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

@ -202,7 +202,7 @@ pub(crate) async fn create_test_actor_ex(
cwd: cwd.as_str().to_string(),
},
auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None),
model_auth_memo: std::cell::RefCell::new(None),
attribution_callback: None,
auth_manager: None,
state,

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

@ -2202,7 +2202,7 @@ mod inline_auto_compact_flow_tests {
cwd: cwd.as_str().to_string(),
},
auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None),
model_auth_memo: std::cell::RefCell::new(None),
attribution_callback: None,
auth_manager: None,
state,

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

@ -138,34 +138,9 @@ mod feedback_tests {
} else {
Some("could be better".into())
},
feedback_categories: vec![],
message_id: None,
model_id: Some("grok-3-fast".into()),
resolved_model_id: Some("grok-4.5".into()),
model_fingerprint: None,
context_type: None,
feature_name: None,
tool_name: None,
experiment_id: None,
comparison_id: None,
preferred_model_id: None,
preference_strength: None,
preference_reasons: vec![],
request_id: None,
client_version: None,
shell_version: None,
extension_host: None,
metadata: None,
last_user_message: None,
last_assistant_message: None,
tool_outcomes: vec![],
session_cwd: None,
compaction_count: None,
context_window_usage: None,
context_tokens_used: None,
context_window_tokens: None,
terminal_info: None,
unified_log_url: None,
..Default::default()
}
}
@ -308,7 +283,8 @@ pub enum PersistenceMsg {
Update(SessionUpdate),
AppendUpdateDurablyAndAck {
update: SessionUpdate,
respond_to: tokio::sync::oneshot::Sender<io::Result<()>>,
respond_to:
tokio::sync::oneshot::Sender<Result<(), crate::session::storage::AppendUpdateError>>,
},
ContentChunk(PersistenceContentChunk),
Chat(ConversationItem),
@ -1369,25 +1345,88 @@ mod generated_title_tests {
pub struct PersistenceHandle {
pub tx: mpsc::UnboundedSender<PersistenceMsg>,
/// Explicit flag set only by [`Self::noop`]. Do not treat a closed sender
/// alone as noop — a real persistence actor may exit and drop its receiver.
noop: bool,
}
#[derive(Debug)]
pub enum DurableAppendError {
NotCommitted(io::Error),
Committed(io::Error),
AcknowledgementLost(io::Error),
}
impl std::fmt::Display for DurableAppendError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::NotCommitted(error)
| Self::Committed(error)
| Self::AcknowledgementLost(error) => error.fmt(formatter),
}
}
}
impl std::error::Error for DurableAppendError {}
impl From<crate::session::storage::AppendUpdateError> for DurableAppendError {
fn from(error: crate::session::storage::AppendUpdateError) -> Self {
use crate::session::storage::AppendUpdateError;
match error {
AppendUpdateError::NotCommitted(error) => Self::NotCommitted(error),
AppendUpdateError::Committed(error) => Self::Committed(error),
}
}
}
impl PersistenceHandle {
/// Create a no-op persistence handle that silently discards all messages.
///
/// Used for subagent child sessions that don't need disk persistence
/// (their results are captured by the parent via the oneshot channel).
pub fn noop() -> Self {
let (tx, _rx) = mpsc::unbounded_channel();
Self { tx, noop: true }
}
/// `true` only for handles created via [`Self::noop`].
pub fn is_noop(&self) -> bool {
self.noop
}
/// Append after older buffered updates and wait for the durable barrier.
///
/// [`DurableAppendError::NotCommitted`] is safe to retry; [`DurableAppendError::Committed`]
/// means the replay line landed; [`DurableAppendError::AcknowledgementLost`] has unknown status.
/// No-op handles return `Unsupported`.
pub async fn append_update_durably(
&self,
update: SessionUpdate,
) -> Result<(), DurableAppendError> {
if self.noop {
return Err(DurableAppendError::NotCommitted(io::Error::new(
io::ErrorKind::Unsupported,
"durable session update append is unsupported by a no-op persistence handle",
)));
}
let (respond_to, response) = tokio::sync::oneshot::channel();
self.tx
.send(PersistenceMsg::AppendUpdateDurablyAndAck { update, respond_to })
.map_err(|_| {
DurableAppendError::NotCommitted(io::Error::new(
io::ErrorKind::BrokenPipe,
"session persistence actor stopped before durable append dispatch",
))
})?;
response
.await
.map_err(|_| {
DurableAppendError::AcknowledgementLost(io::Error::new(
io::ErrorKind::BrokenPipe,
"session persistence actor stopped before durable append acknowledgement",
))
})?
.map_err(DurableAppendError::from)
}
}
enum PendingAppendOutcome {
CommittedOk(acp::SessionNotification),
CommittedErr(acp::SessionNotification, io::Error),
NotCommittedErr(acp::SessionNotification, io::Error),
}
struct SessionPersistence {
@ -1517,46 +1556,69 @@ impl SessionPersistence {
}
fn finish_pending_append(
pending: &mut Option<acp::SessionNotification>,
notification: acp::SessionNotification,
result: Result<(), crate::session::storage::AppendUpdateError>,
) -> Result<acp::SessionNotification, io::Error> {
) -> PendingAppendOutcome {
match result {
Ok(()) => Ok(notification),
Ok(()) => PendingAppendOutcome::CommittedOk(notification),
Err(crate::session::storage::AppendUpdateError::NotCommitted(error)) => {
*pending = Some(notification);
Err(error)
PendingAppendOutcome::NotCommittedErr(notification, error)
}
Err(crate::session::storage::AppendUpdateError::Committed(error)) => {
PendingAppendOutcome::CommittedErr(notification, error)
}
Err(crate::session::storage::AppendUpdateError::Committed(error)) => Err(error),
}
}
async fn drain_pending(&mut self) -> io::Result<()> {
/// Restore uncommitted failures; sync committed records before returning errors.
async fn drain_pending(&mut self) -> Result<(), crate::session::storage::AppendUpdateError> {
if let Some(notification) = self.pending_notification.take() {
let result = self
.write_update(&SessionUpdate::Acp(Box::new(notification.clone())))
.await;
match Self::finish_pending_append(
&mut self.pending_notification,
notification.clone(),
result,
) {
Ok(notification) => self.queue_acp_sync(notification),
Err(error) => {
if self.pending_notification.is_none() {
self.queue_acp_sync(notification);
}
return Err(error);
match Self::finish_pending_append(notification, result) {
PendingAppendOutcome::CommittedOk(notification) => {
self.queue_acp_sync(notification);
}
PendingAppendOutcome::CommittedErr(notification, error) => {
self.queue_acp_sync(notification);
return Err(crate::session::storage::AppendUpdateError::Committed(error));
}
PendingAppendOutcome::NotCommittedErr(notification, error) => {
self.pending_notification = Some(notification);
return Err(crate::session::storage::AppendUpdateError::NotCommitted(
error,
));
}
}
}
Ok(())
}
async fn handle_durable_append(
&mut self,
update: SessionUpdate,
) -> Result<(), crate::session::storage::AppendUpdateError> {
self.drain_pending().await?;
let result = self
.storage
.append_update_durable_commit_aware(&self.info, &update)
.await;
match (&update, &result) {
(SessionUpdate::Acp(notification), Ok(()))
| (
SessionUpdate::Acp(notification),
Err(crate::session::storage::AppendUpdateError::Committed(_)),
) => self.queue_acp_sync((**notification).clone()),
_ => {}
}
result
}
/// Flush any pending merged ACP notification to disk and remote sync.
async fn flush_pending(&mut self) {
if let Err(error) = self.drain_pending().await {
tracing::warn!(?error, "failed to write pending update");
tracing::warn!(%error, "failed to write pending update");
}
if let Some(sync) = &self.remote_sync {
sync.flush();
@ -1624,17 +1686,7 @@ impl SessionPersistence {
}
}
PersistenceMsg::AppendUpdateDurablyAndAck { update, respond_to } => {
let result = async {
self.drain_pending().await?;
self.storage
.append_update_durable(&self.info, &update)
.await?;
if let SessionUpdate::Acp(notification) = update {
self.queue_acp_sync(*notification);
}
Ok(())
}
.await;
let result = self.handle_durable_append(update).await;
let _ = respond_to.send(result);
}
PersistenceMsg::Chat(chat_msg) => {

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

@ -1,4 +1,5 @@
use super::*;
use crate::session::storage::jsonl::AppendDurability;
struct ActorGuard {
handle: PersistenceHandle,
@ -13,6 +14,14 @@ impl ActorGuard {
}
fn test_actor(info: Info, storage: Arc<dyn StorageAdapter>) -> ActorGuard {
test_actor_with_remote_sync(info, storage, None)
}
fn test_actor_with_remote_sync(
info: Info,
storage: Arc<dyn StorageAdapter>,
remote_sync: Option<RemoteSync>,
) -> ActorGuard {
let (tx, rx) = mpsc::unbounded_channel();
let summary_tx = tx.clone();
let sampling_client = OaiCompatClient::new(xai_grok_sampler::SamplerConfig::default()).unwrap();
@ -22,7 +31,7 @@ fn test_actor(info: Info, storage: Arc<dyn StorageAdapter>) -> ActorGuard {
storage,
pending_notification: None,
rx,
remote_sync: None,
remote_sync,
relay_sync: None,
summary: crate::session::summary::SummaryGenerator::new(
crate::session::summary::SummaryConfig {
@ -55,50 +64,210 @@ fn neutral_update(info: &Info, text: &str) -> SessionUpdate {
SessionUpdate::Acp(Box::new(notification(info, text)))
}
#[test]
fn committed_error_does_not_restore_pending_notification() {
let notification = notification(
&Info {
id: acp::SessionId::new("committed-update"),
cwd: "/test".into(),
},
"committed",
);
let mut pending = None;
let result = SessionPersistence::finish_pending_append(
&mut pending,
notification,
Err(crate::session::storage::AppendUpdateError::Committed(
io::Error::other("summary patch failed"),
)),
);
assert_eq!(result.unwrap_err().to_string(), "summary patch failed");
assert!(pending.is_none());
fn break_summary_writes(dir: &std::path::Path) {
let summary = dir.join("summary.json");
std::fs::remove_file(&summary).unwrap();
std::fs::create_dir(summary).unwrap();
}
async fn recv_observed(
observed: &mut tokio::sync::mpsc::UnboundedReceiver<acp::SessionNotification>,
) -> acp::SessionNotification {
tokio::time::timeout(std::time::Duration::from_secs(1), observed.recv())
.await
.expect("remote sync timed out")
.expect("remote sync observer closed")
}
#[test]
fn uncommitted_error_restores_pending_notification() {
let notification = notification(
&Info {
id: acp::SessionId::new("uncommitted-update"),
cwd: "/test".into(),
},
"pending",
);
let mut pending = None;
let result = SessionPersistence::finish_pending_append(
&mut pending,
notification,
Err(crate::session::storage::AppendUpdateError::NotCommitted(
io::Error::other("append failed"),
)),
);
assert!(result.is_err());
assert!(pending.is_some());
fn committed_error_returns_sync_disposition() {
let info = Info {
id: acp::SessionId::new("committed-update"),
cwd: "/test".into(),
};
let notification = notification(&info, "committed");
let PendingAppendOutcome::CommittedErr(sync_notification, error) =
SessionPersistence::finish_pending_append(
notification,
Err(crate::session::storage::AppendUpdateError::Committed(
io::Error::other("summary patch failed"),
)),
)
else {
panic!("expected committed failure");
};
assert_eq!(sync_notification.session_id, info.id);
assert_eq!(error.to_string(), "summary patch failed");
}
#[test]
fn uncommitted_error_returns_restore_disposition() {
let info = Info {
id: acp::SessionId::new("uncommitted-update"),
cwd: "/test".into(),
};
let notification = notification(&info, "pending");
let PendingAppendOutcome::NotCommittedErr(pending_notification, error) =
SessionPersistence::finish_pending_append(
notification,
Err(crate::session::storage::AppendUpdateError::NotCommitted(
io::Error::other("append failed"),
)),
)
else {
panic!("expected uncommitted failure");
};
assert_eq!(pending_notification.session_id, info.id);
assert_eq!(error.to_string(), "append failed");
}
#[tokio::test]
async fn durable_ack_drains_pending_update_in_fifo_order() {
async fn noop_handle_rejects_durable_append() {
let info = Info {
id: acp::SessionId::new("noop-durable-update"),
cwd: "/test".into(),
};
assert!(matches!(
PersistenceHandle::noop()
.append_update_durably(neutral_update(&info, "durable"))
.await,
Err(DurableAppendError::NotCommitted(error))
if error.kind() == io::ErrorKind::Unsupported
));
}
#[tokio::test]
async fn pending_drain_disposition_controls_remote_sync() {
let info = Info {
id: acp::SessionId::new("pending-remote-sync"),
cwd: "/test".into(),
};
let storage = JsonlStorageAdapter::with_update_append_probe("/unused".into(), |_| {
Err(io::Error::other("append failed"))
});
let (remote_sync, mut observed) = RemoteSync::test_observer();
let actor = test_actor_with_remote_sync(info.clone(), Arc::new(storage), Some(remote_sync));
actor
.handle
.tx
.send(PersistenceMsg::Update(neutral_update(&info, "pending")))
.unwrap();
assert!(matches!(
actor
.handle
.append_update_durably(neutral_update(&info, "durable"))
.await,
Err(DurableAppendError::NotCommitted(_))
));
assert!(observed.try_recv().is_err());
actor.stop().await;
let dir = tempfile::tempdir().unwrap();
let attempts = Arc::new(std::sync::Mutex::new(Vec::new()));
let observed_attempts = attempts.clone();
let storage = Arc::new(JsonlStorageAdapter::with_update_append_probe(
dir.path().to_path_buf(),
move |durability| {
observed_attempts.lock().unwrap().push(durability);
Ok(())
},
));
storage
.init_session(&info, default_model_id())
.await
.unwrap();
let (remote_sync, mut observed) = RemoteSync::test_observer();
let actor = test_actor_with_remote_sync(info.clone(), storage, Some(remote_sync));
actor
.handle
.tx
.send(PersistenceMsg::Update(neutral_update(&info, "pending")))
.unwrap();
break_summary_writes(dir.path());
assert!(matches!(
actor
.handle
.append_update_durably(neutral_update(&info, "durable"))
.await,
Err(DurableAppendError::Committed(_))
));
let synced = recv_observed(&mut observed).await;
assert_eq!(synced.session_id, info.id);
assert!(matches!(
attempts.lock().unwrap().as_slice(),
[AppendDurability::Buffered]
));
actor.stop().await;
}
#[tokio::test]
async fn durable_append_committed_failure_is_synced() {
let dir = tempfile::tempdir().unwrap();
let info = Info {
id: acp::SessionId::new("durable-remote-sync"),
cwd: "/test".into(),
};
let storage = Arc::new(JsonlStorageAdapter::with_explicit_session_dir(
dir.path().to_path_buf(),
));
storage
.init_session(&info, default_model_id())
.await
.unwrap();
break_summary_writes(dir.path());
let (remote_sync, mut observed) = RemoteSync::test_observer();
let actor = test_actor_with_remote_sync(info.clone(), storage, Some(remote_sync));
assert!(matches!(
actor
.handle
.append_update_durably(neutral_update(&info, "durable"))
.await,
Err(DurableAppendError::Committed(_))
));
let synced = recv_observed(&mut observed).await;
assert_eq!(synced.session_id, info.id);
actor.stop().await;
}
#[tokio::test]
async fn failed_pending_drain_retains_record_and_skips_durable_update() {
let info = Info {
id: acp::SessionId::new("durable-drain-failure"),
cwd: "/test".into(),
};
let attempts = Arc::new(std::sync::Mutex::new(Vec::new()));
let observed = attempts.clone();
let storage =
JsonlStorageAdapter::with_update_append_probe("/unused".into(), move |durability| {
observed.lock().unwrap().push(durability);
Err(io::Error::other("pending append failed"))
});
let actor = test_actor(info.clone(), Arc::new(storage));
actor
.handle
.tx
.send(PersistenceMsg::Update(neutral_update(&info, "pending")))
.unwrap();
for _ in 0..2 {
assert_eq!(
actor
.handle
.append_update_durably(neutral_update(&info, "durable"))
.await
.unwrap_err()
.to_string(),
"pending append failed"
);
}
assert!(matches!(
attempts.lock().unwrap().as_slice(),
[AppendDurability::Buffered, AppendDurability::Buffered]
));
actor.stop().await;
}
#[tokio::test]
async fn durable_append_drains_pending_update_in_fifo_order() {
let dir = tempfile::tempdir().unwrap();
let info = Info {
id: acp::SessionId::new("durable-update"),
@ -117,16 +286,11 @@ async fn durable_ack_drains_pending_update_in_fifo_order() {
.tx
.send(PersistenceMsg::Update(neutral_update(&info, "before")))
.unwrap();
let (respond_to, response) = tokio::sync::oneshot::channel();
actor
.handle
.tx
.send(PersistenceMsg::AppendUpdateDurablyAndAck {
update: neutral_update(&info, "durable"),
respond_to,
})
.append_update_durably(neutral_update(&info, "durable"))
.await
.unwrap();
response.await.unwrap().unwrap();
let summary = storage.load_summary(&info).await.unwrap();
assert_eq!(summary.num_messages, 2);

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

@ -44,7 +44,10 @@ async fn ordinary_and_durable_appends_keep_every_physical_line_parseable() {
let durable = tokio::spawn(async move {
for index in 0..N {
durable
.append_update_durable(&info_b, &update(&info_b, format!("durable-{index}")))
.append_update_durable_commit_aware(
&info_b,
&update(&info_b, format!("durable-{index}")),
)
.await
.unwrap();
}
@ -71,69 +74,45 @@ async fn append_commit_is_reported_when_bookkeeping_fails() {
.init_session(&info, default_model_id())
.await
.unwrap();
let result = adapter
.append_update_with_bookkeeping(&info, &update(&info, "committed".into()), async {
Err(io::Error::other("summary patch failed"))
})
.await;
let summary = dir.path().join("summary.json");
std::fs::remove_file(&summary).unwrap();
std::fs::create_dir(&summary).unwrap();
assert!(matches!(
result,
adapter
.append_update_durable_commit_aware(&info, &update(&info, "committed".into()))
.await,
Err(crate::session::storage::AppendUpdateError::Committed(_))
));
let bytes = std::fs::read(dir.path().join("updates.jsonl")).unwrap();
let parsed = bytes
.split(|byte| *byte == b'\n')
.filter(|line| !line.is_empty())
.map(serde_json::from_slice::<SessionUpdateEnvelope>)
.collect::<Result<Vec<_>, _>>()
.unwrap();
assert_eq!(parsed.len(), 1);
}
#[test]
fn lock_serializes_tail_heal_and_complete_record() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("updates.jsonl");
std::fs::write(&path, b"torn").unwrap();
JsonlStorageAdapter::append_jsonl_line_sync(
&path,
b"{\"record\":1}\n".to_vec(),
AppendDurability::Buffered,
)
.unwrap();
assert_eq!(
std::fs::read_to_string(path).unwrap(),
"torn\n{\"record\":1}\n"
std::fs::read_to_string(dir.path().join("updates.jsonl"))
.unwrap()
.lines()
.count(),
1
);
}
#[test]
fn directory_barrier_failure_is_retried_even_after_file_exists() {
use std::sync::atomic::{AtomicUsize, Ordering};
static ATTEMPTS: AtomicUsize = AtomicUsize::new(0);
static TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
let _guard = TEST_LOCK.lock().unwrap();
fn sync_file(file: &std::fs::File) -> io::Result<()> {
file.sync_all()
}
fn flaky_parent(_path: &Path) -> io::Result<()> {
if ATTEMPTS.fetch_add(1, Ordering::SeqCst) == 0 {
let mut attempts = 0;
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("updates.jsonl");
let mut flaky_parent = || {
attempts += 1;
if attempts == 1 {
Err(io::Error::other("directory barrier failed"))
} else {
Ok(())
}
}
ATTEMPTS.store(0, Ordering::SeqCst);
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("updates.jsonl");
};
assert!(
JsonlStorageAdapter::append_jsonl_line_sync_with(
&path,
b"{\"record\":1}\n".to_vec(),
AppendDurability::Durable,
sync_file,
flaky_parent,
std::fs::File::sync_all,
&mut flaky_parent,
)
.is_err()
);
@ -141,33 +120,9 @@ fn directory_barrier_failure_is_retried_even_after_file_exists() {
&path,
b"{\"record\":1}\n".to_vec(),
AppendDurability::Durable,
sync_file,
flaky_parent,
std::fs::File::sync_all,
&mut flaky_parent,
)
.unwrap();
assert_eq!(ATTEMPTS.load(Ordering::SeqCst), 2);
}
#[test]
fn file_barrier_error_propagates() {
fn fail(_file: &std::fs::File) -> io::Result<()> {
Err(io::Error::other("file barrier failed"))
}
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("updates.jsonl");
let error = JsonlStorageAdapter::append_jsonl_line_sync_with(
&path,
b"{\"record\":1}\n".to_vec(),
AppendDurability::Durable,
fail,
|_| Ok(()),
)
.unwrap_err();
assert_eq!(error.to_string(), "file barrier failed");
}
#[cfg(target_os = "macos")]
#[test]
fn darwin_fullfsync_seam_reports_invalid_descriptor() {
assert!(JsonlStorageAdapter::fullfsync_raw(-1).is_err());
assert_eq!(attempts, 2);
}

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

@ -15,28 +15,25 @@ use std::io::{self, Read, Seek, Write};
use std::os::fd::AsRawFd;
use std::path::{Path, PathBuf};
use xai_grok_workspace::session::file_state::RewindPoint;
/// How the adapter resolves the session directory on disk.
///
/// - `FromRoot` (default): computes `{root}/sessions/{urlencoded(cwd)}/{session_id}/`
/// - `Explicit`: uses a caller-provided directory directly, ignoring `Info` fields.
/// Used for subagent child sessions whose files live under the parent's session dir.
#[derive(Clone)]
enum SessionDirMode {
/// Existing behavior: root + sessions/{cwd}/{id}/
FromRoot(PathBuf),
/// New: use this directory directly (for subagent children).
Explicit(PathBuf),
}
pub(super) enum AppendDurability {
#[derive(Clone, Copy)]
pub(crate) enum AppendDurability {
Buffered,
Durable,
}
/// JSONL-based storage adapter (legacy format)
/// Stores sessions in {root}/sessions/{url_encoded_cwd}/{session_id}/
/// JSONL storage under `{root}/sessions/{url_encoded_cwd}/{session_id}/`.
#[derive(Clone)]
pub struct JsonlStorageAdapter {
dir_mode: SessionDirMode,
#[cfg(test)]
update_append_probe: Option<std::sync::Arc<AppendProbe>>,
}
#[cfg(test)]
type AppendProbe = dyn Fn(AppendDurability) -> io::Result<()> + Send + Sync;
impl Default for JsonlStorageAdapter {
fn default() -> Self {
Self::new()
@ -46,11 +43,15 @@ impl JsonlStorageAdapter {
pub fn new() -> Self {
Self {
dir_mode: SessionDirMode::FromRoot(crate::util::grok_home::grok_home()),
#[cfg(test)]
update_append_probe: None,
}
}
pub fn with_root(root_dir: PathBuf) -> Self {
Self {
dir_mode: SessionDirMode::FromRoot(root_dir),
#[cfg(test)]
update_append_probe: None,
}
}
/// Create an adapter that writes directly to `session_dir`, bypassing
@ -61,6 +62,18 @@ impl JsonlStorageAdapter {
pub fn with_explicit_session_dir(session_dir: PathBuf) -> Self {
Self {
dir_mode: SessionDirMode::Explicit(session_dir),
#[cfg(test)]
update_append_probe: None,
}
}
#[cfg(test)]
pub(crate) fn with_update_append_probe(
session_dir: PathBuf,
append_probe: impl Fn(AppendDurability) -> io::Result<()> + Send + Sync + 'static,
) -> Self {
Self {
dir_mode: SessionDirMode::Explicit(session_dir),
update_append_probe: Some(std::sync::Arc::new(append_probe)),
}
}
/// Load chat history from a specific directory.
@ -247,8 +260,19 @@ impl JsonlStorageAdapter {
line.push(b'\n');
self.append_jsonl_line(path, line).await
}
/// Append one newline-terminated JSONL record to `path`, healing a torn
/// tail first.
async fn append_jsonl_line(&self, path: PathBuf, line: Vec<u8>) -> io::Result<()> {
Self::append_jsonl_line_blocking(path, line, AppendDurability::Buffered).await
}
async fn append_jsonl_line_blocking(
path: PathBuf,
line: Vec<u8>,
durability: AppendDurability,
) -> io::Result<()> {
tokio::task::spawn_blocking(move || Self::append_jsonl_line_sync(&path, line, durability))
.await
.map_err(io::Error::other)?
}
/// Append one JSONL record, healing a torn tail before writing.
///
/// Appends are not crash-atomic: a process kill / `ENOSPC` mid-`write_all`
/// (e.g. the auto-update leader relaunch aborting a persistence actor
@ -263,47 +287,24 @@ impl JsonlStorageAdapter {
/// the torn record is terminated as its own (single) corrupt line. This
/// bounds the damage of any torn write to exactly one record, which the
/// lenient readers (e.g. [`Self::read_chat_history_sync`]) then skip.
async fn append_jsonl_line(&self, path: PathBuf, line: Vec<u8>) -> io::Result<()> {
Self::append_jsonl_line_locked(path, line, AppendDurability::Buffered).await
}
async fn append_jsonl_line_locked(
path: PathBuf,
line: Vec<u8>,
durability: AppendDurability,
) -> io::Result<()> {
tokio::task::spawn_blocking(move || Self::append_jsonl_line_sync(&path, line, durability))
.await
.map_err(io::Error::other)?
}
fn append_jsonl_line_sync(
path: &Path,
line: Vec<u8>,
durability: AppendDurability,
) -> io::Result<()> {
Self::append_jsonl_line_sync_with(
path,
line,
durability,
Self::sync_file_durable,
Self::sync_parent_directory,
)
Self::append_jsonl_line_sync_with(path, line, durability, Self::sync_file_durable, || {
Self::sync_parent_directory(path)
})
}
fn append_jsonl_line_sync_with(
path: &Path,
mut line: Vec<u8>,
durability: AppendDurability,
sync_file: fn(&std::fs::File) -> io::Result<()>,
sync_parent: fn(&Path) -> io::Result<()>,
mut sync_file: impl FnMut(&std::fs::File) -> io::Result<()>,
mut sync_parent: impl FnMut() -> io::Result<()>,
) -> io::Result<()> {
debug_assert!(line.ends_with(b"\n"), "JSONL record must end with \\n");
let lock_path = path.with_extension("jsonl.lock");
let lock = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(lock_path)?;
lock.lock_exclusive()?;
let lock = Self::lock_append(path)?;
let result = (|| {
let mut file = OpenOptions::new()
.read(true)
@ -317,7 +318,8 @@ impl JsonlStorageAdapter {
file.read_exact(&mut last)?;
if last[0] != b'\n' {
tracing::warn!(
path = % path.display(), "terminating torn jsonl tail"
path = % path.display(),
"jsonl file has a torn trailing line (previous append crashed mid-write?); terminating it before appending"
);
line.insert(0, b'\n');
}
@ -327,7 +329,7 @@ impl JsonlStorageAdapter {
if matches!(durability, AppendDurability::Durable) {
sync_file(&file)?;
drop(file);
sync_parent(path)?;
sync_parent()?;
} else {
drop(file);
}
@ -336,6 +338,18 @@ impl JsonlStorageAdapter {
let _ = lock.unlock();
result
}
/// Lock tail healing, append, and barriers through `<target>.jsonl.lock`.
/// Full-file [`Self::write_jsonl`] atomic-rename rewrites bypass this append-only lock.
fn lock_append(path: &Path) -> io::Result<std::fs::File> {
let lock = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(path.with_extension("jsonl.lock"))?;
lock.lock_exclusive()?;
Ok(lock)
}
#[cfg(target_os = "macos")]
fn sync_file_durable(file: &std::fs::File) -> io::Result<()> {
file.sync_all()?;
@ -413,28 +427,36 @@ impl JsonlStorageAdapter {
update: &super::SessionUpdate,
durability: AppendDurability,
) -> io::Result<()> {
#[cfg(test)]
if let Some(append_probe) = &self.update_append_probe {
append_probe(durability)?;
}
let envelope = SessionUpdateEnvelope::from_update(update)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
let mut line = serde_json::to_vec(&envelope)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
line.push(b'\n');
Self::append_jsonl_line_locked(path, line, durability).await
Self::append_jsonl_line_blocking(path, line, durability).await
}
pub(crate) async fn append_update_with_bookkeeping<F>(
async fn append_update_with_bookkeeping(
&self,
info: &Info,
update: &super::SessionUpdate,
bookkeeping: F,
) -> Result<(), super::AppendUpdateError>
where
F: std::future::Future<Output = io::Result<()>>,
{
self.append_update_to_file(self.updates_file(info), update, AppendDurability::Buffered)
durability: AppendDurability,
) -> Result<(), super::AppendUpdateError> {
self.append_update_to_file(self.updates_file(info), update, durability)
.await
.map_err(super::AppendUpdateError::NotCommitted)?;
bookkeeping
.await
.map_err(super::AppendUpdateError::Committed)
self.apply_summary_patch(
info,
super::summary_write::SummaryPatch {
record_activity: true,
messages: Some(super::summary_write::CounterOp::Increment(1)),
..Default::default()
},
)
.await
.map_err(super::AppendUpdateError::Committed)
}
/// Read session updates from an updates.jsonl file, handling both envelope and legacy formats.
///
@ -1075,36 +1097,16 @@ impl StorageAdapter for JsonlStorageAdapter {
info: &Info,
update: &super::SessionUpdate,
) -> Result<(), super::AppendUpdateError> {
self.append_update_with_bookkeeping(
info,
update,
self.apply_summary_patch(
info,
super::summary_write::SummaryPatch {
record_activity: true,
messages: Some(super::summary_write::CounterOp::Increment(1)),
..Default::default()
},
),
)
.await
self.append_update_with_bookkeeping(info, update, AppendDurability::Buffered)
.await
}
async fn append_update_durable(
async fn append_update_durable_commit_aware(
&self,
info: &Info,
update: &super::SessionUpdate,
) -> io::Result<()> {
self.append_update_to_file(self.updates_file(info), update, AppendDurability::Durable)
.await?;
self.apply_summary_patch(
info,
super::summary_write::SummaryPatch {
record_activity: true,
messages: Some(super::summary_write::CounterOp::Increment(1)),
..Default::default()
},
)
.await
) -> Result<(), super::AppendUpdateError> {
self.append_update_with_bookkeeping(info, update, AppendDurability::Durable)
.await
}
async fn append_chat_message(&self, info: &Info, message: &ConversationItem) -> io::Result<()> {
self.append_jsonl(self.chat_file(info), message).await?;

60
crates/codegen/xai-grok-shell/src/session/storage/jsonl/tests.rs
View file

@ -1138,41 +1138,14 @@ async fn test_append_feedback_creates_file_and_persists() {
dismissed: false,
submission: Some(FeedbackSubmission {
session_id: "test-session-123".into(),
user_id: None,
client_type: ClientType::Tui,
feedback_type: FeedbackType::Rating,
turn_number: Some(3),
rating_type: Some(RatingType::Thumbs),
rating_value: Some(1),
feedback_text: None,
feedback_categories: vec![],
message_id: None,
model_id: Some("grok-3-fast".into()),
resolved_model_id: Some("grok-4.5".into()),
model_fingerprint: None,
context_type: None,
feature_name: None,
tool_name: None,
experiment_id: None,
comparison_id: None,
preferred_model_id: None,
preference_strength: None,
preference_reasons: vec![],
request_id: None,
client_version: None,
shell_version: None,
extension_host: None,
metadata: None,
last_user_message: None,
last_assistant_message: None,
tool_outcomes: vec![],
session_cwd: None,
compaction_count: None,
context_window_usage: None,
context_tokens_used: None,
context_window_tokens: None,
terminal_info: None,
unified_log_url: None,
..Default::default()
}),
});
adapter.append_feedback(&info, &user_entry).await.unwrap();
@ -1838,9 +1811,7 @@ fn write_test_summary(
}
#[test]
fn scan_session_dirs_returns_empty_for_explicit_mode() {
let adapter = JsonlStorageAdapter {
dir_mode: SessionDirMode::Explicit(PathBuf::from("/fake")),
};
let adapter = JsonlStorageAdapter::with_explicit_session_dir(PathBuf::from("/fake"));
assert!(adapter.scan_session_dirs(None).is_empty());
}
#[test]
@ -2080,7 +2051,8 @@ fn strip_invalid_images_corrupt_base64_stripped() {
if let ConversationItem::User(u) = &items[0] {
assert_eq!(u.content.len(), 2);
assert!(
matches!(& u.content[1], ContentPart::Text { text } if text
matches!(& u.content[1], ContentPart::Text { text }
if text
.contains("invalid data"))
);
} else {
@ -2117,7 +2089,8 @@ fn strip_invalid_images_http_url_untouched() {
assert_eq!(strip_invalid_images(& mut items), 0);
assert!(
matches!(& items[0], ConversationItem::User(u) if matches!(& u.content[0],
ContentPart::Image { url : u } if u.as_ref() == "https://example.com/photo.jpg"))
ContentPart::Image { url : u }
if u.as_ref() == "https://example.com/photo.jpg"))
);
}
#[test]
@ -2146,19 +2119,23 @@ fn strip_invalid_images_mixed_valid_and_invalid() {
if let ConversationItem::User(u) = &items[0] {
assert_eq!(u.content.len(), 4);
assert!(
matches!(& u.content[0], ContentPart::Text { text } if text.as_ref() ==
matches!(& u.content[0], ContentPart::Text { text }
if text.as_ref() ==
"check these")
);
assert!(
matches!(& u.content[1], ContentPart::Image { url } if url.as_ref() ==
matches!(& u.content[1], ContentPart::Image { url }
if url.as_ref() ==
valid_url.as_str())
);
assert!(
matches!(& u.content[2], ContentPart::Text { text } if text
matches!(& u.content[2], ContentPart::Text { text }
if text
.contains("invalid data"))
);
assert!(
matches!(& u.content[3], ContentPart::Image { url } if url.as_ref() ==
matches!(& u.content[3], ContentPart::Image { url }
if url.as_ref() ==
"https://example.com/img.png")
);
} else {
@ -2197,7 +2174,8 @@ fn strip_invalid_images_heals_tool_result_images() {
};
assert_eq!(t.images.len(), 1, "only the invalid image is removed");
assert!(
matches!(& t.images[0], ContentPart::Image { url } if url.as_ref() == good_url
matches!(& t.images[0], ContentPart::Image { url }
if url.as_ref() == good_url
.as_str())
);
}
@ -2243,7 +2221,8 @@ fn strip_invalid_images_truncated_jpeg_stripped() {
assert_eq!(strip_invalid_images(& mut items), 1);
assert!(
matches!(& items[0], ConversationItem::User(u) if matches!(& u.content[1],
ContentPart::Text { text } if text.contains("invalid data")))
ContentPart::Text { text }
if text.contains("invalid data")))
);
}
#[test]
@ -2547,7 +2526,8 @@ fn read_chat_history_quarantines_original_on_image_strip() {
let (_, chat_path, items) = load_raw_chat(&temp_dir, raw.as_bytes());
assert!(
matches!(& items[0], ConversationItem::User(u) if matches!(& u.content[0],
ContentPart::Text { text } if text.contains("invalid data")))
ContentPart::Text { text }
if text.contains("invalid data")))
);
let quarantine = chat_path.with_extension("jsonl.corrupt");
assert_eq!(

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

@ -944,15 +944,16 @@ pub trait StorageAdapter: Send + Sync {
.map_err(AppendUpdateError::NotCommitted)
}
/// Append one update with the ordinary bookkeeping and a durable log barrier.
///
/// Adapters without this capability return `Unsupported`; callers must tolerate a duplicate
/// record when retrying an error that occurred after the append reached storage.
async fn append_update_durable(&self, _info: &Info, _update: &SessionUpdate) -> io::Result<()> {
Err(io::Error::new(
/// Append one update durably, preserving whether the replay record committed before failure.
async fn append_update_durable_commit_aware(
&self,
_info: &Info,
_update: &SessionUpdate,
) -> Result<(), AppendUpdateError> {
Err(AppendUpdateError::NotCommitted(io::Error::new(
io::ErrorKind::Unsupported,
"durable session update append is unsupported",
))
)))
}
/// Append a chat message and increment counter

36
crates/codegen/xai-grok-shell/src/session/user_message.rs
View file

@ -104,26 +104,49 @@ pub async fn compute_vcs_status_block(
working_directory: &Path,
vcs_kind: VcsKind,
) -> Option<String> {
use xai_grok_workspace::file_system::{git_status, jj_status};
use xai_grok_workspace::file_system::{git_status_short, jj_status};
if matches!(vcs_kind, VcsKind::None) {
return None;
}
let _timer = crate::instrumentation_timer!("session.user_prefix.vcs_status");
let timeout = std::time::Duration::from_secs(2);
let mut timer = crate::instrumentation_timer!("session.user_prefix.vcs_status");
timer.with_field("vcs", if vcs_kind.is_jj() { "jj" } else { "git" });
timer.with_field(
"status_mode",
if vcs_kind.is_jj() {
"jj"
} else {
"short_untracked_normal"
},
);
timer.with_field("timeout_ms", 5_000_u64);
let timeout = std::time::Duration::from_secs(5);
let result = if vcs_kind.is_jj() {
tokio::time::timeout(timeout, jj_status(working_directory)).await
} else {
tokio::time::timeout(timeout, git_status(working_directory)).await
tokio::time::timeout(timeout, git_status_short(working_directory)).await
};
match result {
Ok(Ok(status)) => Some(format_vcs_status_block(&status, vcs_kind)),
Ok(Ok(status)) => {
timer.with_field("outcome", "success");
timer.with_field("output_bytes", status.len() as u64);
let status = if vcs_kind.is_jj() {
Some(status)
} else {
xai_grok_agent::prompt::user_message::normalize_git_status(&status)
};
status.map(|status| format_vcs_status_block(&status, vcs_kind))
}
Ok(Err(e)) => {
timer.with_field("outcome", "error");
timer.with_field("output_bytes", 0_u64);
tracing::warn!("user prefix VCS status failed: {e}");
None
}
Err(_) => {
tracing::warn!(vcs = ?vcs_kind, "user prefix VCS status timed out after 2s");
timer.with_field("outcome", "timeout");
timer.with_field("output_bytes", 0_u64);
tracing::warn!(vcs = ?vcs_kind, "user prefix VCS status timed out after 5s");
None
}
}
@ -166,7 +189,6 @@ mod tests {
use super::*;
use xai_grok_workspace::file_system::FsError;
/// Verify that construct_user_message completes within the 2s git_status
/// timeout even when pointed at a non-existent directory (git commands
/// fail instantly → no timeout path exercised, but validates the happy
/// path doesn't regress).

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

File diff suppressed because it is too large Load diff

22
crates/codegen/xai-grok-shell/src/upload/trace.rs
View file

@ -476,7 +476,6 @@ pub(crate) fn mime_type_to_extension(mime_type: &str) -> &str {
_ => "bin",
}
}
/// Path format: {session_id}/turn_{N}/full_prompt.txt
pub(crate) async fn upload_full_prompt_txt(ctx: &PromptTraceContext, _full_prompt: &str) {
super::manifest::skip_artifact(
&ctx.artifact_tracker,
@ -977,7 +976,6 @@ pub(crate) async fn upload_permission_events(
)
.await;
}
/// Path format: {session_id}/turn_{N}/turn_messages.json
pub(crate) async fn upload_turn_messages(
ctx: &PromptTraceContext,
_capture: xai_chat_state::TurnCapture,
@ -1009,7 +1007,7 @@ pub(crate) struct SessionStateBuildError {
/// zero-byte payload the viewer treats as "no history" (harness pairs always
/// carry ≥1 message, so this is only a safety floor).
pub(crate) fn build_chat_history_session_state(
_messages: &[xai_grok_sampling_types::conversation::ConversationItem],
messages: &[xai_grok_sampling_types::conversation::ConversationItem],
) -> Result<Vec<u8>, SessionStateBuildError> {
use flate2::Compression;
use flate2::write::GzEncoder;
@ -1019,7 +1017,10 @@ pub(crate) fn build_chat_history_session_state(
error: error.into(),
}
}
let jsonl = Vec::new();
let jsonl = {
let _ = messages;
Vec::new()
};
let mut archive_data = Vec::new();
{
let encoder = GzEncoder::new(&mut archive_data, Compression::default());
@ -2418,19 +2419,6 @@ mod tests {
out
}
#[test]
fn chat_history_session_state_omits_conversation_items() {
use xai_grok_sampling_types::conversation::ConversationItem;
let messages = vec![
ConversationItem::user("verify whether the change compiles"),
ConversationItem::assistant("PASS: the change compiles and tests pass"),
];
let archive = build_chat_history_session_state(&messages).unwrap();
let entries = read_tar_gz_entries(&archive);
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].0, "chat_history.jsonl");
assert!(entries[0].1.is_empty());
}
#[test]
fn chat_history_session_state_empty_messages_yields_valid_empty_archive() {
let archive = build_chat_history_session_state(&[]).unwrap();
let entries = read_tar_gz_entries(&archive);