diff --git a/SOURCE_REV b/SOURCE_REV index a92b503..e31e81e 100644 --- a/SOURCE_REV +++ b/SOURCE_REV @@ -1 +1 @@ -ba69d70c2f7d70a130a323b2becdf137af784c7f +c5c4ce03436b4bb2cec43d3feaa27dee0109bf37 diff --git a/crates/codegen/xai-chat-state/src/actor/request_builder.rs b/crates/codegen/xai-chat-state/src/actor/request_builder.rs index 4a708aa..3c0dc02 100644 --- a/crates/codegen/xai-chat-state/src/actor/request_builder.rs +++ b/crates/codegen/xai-chat-state/src/actor/request_builder.rs @@ -598,12 +598,13 @@ mod tests { fn has_placeholder(item: &ConversationItem) -> bool { matches!( - item, - ConversationItem::User(u) if u.content.iter().any(|p| matches!( - p, - ContentPart::Text { text } if text.as_ref() == IMAGE_COMPACT_PLACEHOLDER - )) - ) + item, + ConversationItem::User(u) if u.content.iter().any(|p| matches!( + p, + ContentPart::Text { text } + if text.as_ref() == IMAGE_COMPACT_PLACEHOLDER + )) + ) } // Images are sized ~100 KB so the ~235 B placeholder that replaces an diff --git a/crates/codegen/xai-grok-agent/src/config.rs b/crates/codegen/xai-grok-agent/src/config.rs index bcf87ce..fe0dfad 100644 --- a/crates/codegen/xai-grok-agent/src/config.rs +++ b/crates/codegen/xai-grok-agent/src/config.rs @@ -2020,10 +2020,12 @@ description: Minimal agent assert_eq!(v, McpServerRef::Named("slack".to_string())); let v: McpServerRef = serde_json::from_value(serde_json::json!({ "s" : { "type" : "stdio" } })).unwrap(); - assert!(matches!(v, McpServerRef::Inline { ref name, .. } if name == "s")); + assert!(matches!(v, McpServerRef::Inline { ref name, .. } +if name == "s")); let v: McpServerRef = serde_json::from_value(serde_json::json!({ "name" : "s", "type" : "stdio" })).unwrap(); - assert!(matches!(v, McpServerRef::Inline { ref name, .. } if name == "s")); + assert!(matches!(v, McpServerRef::Inline { ref name, .. } +if name == "s")); assert!( serde_json::from_value::(serde_json::json!({ "type" : "stdio" })) diff --git a/crates/codegen/xai-grok-agent/src/prompt/user_message.rs b/crates/codegen/xai-grok-agent/src/prompt/user_message.rs index a12c284..e67d2bc 100644 --- a/crates/codegen/xai-grok-agent/src/prompt/user_message.rs +++ b/crates/codegen/xai-grok-agent/src/prompt/user_message.rs @@ -35,7 +35,7 @@ pub const GIT_STATUS_CHARACTER_LIMIT: usize = 10_000; /// and no empty code fence is emitted), otherwise the status capped at /// [`GIT_STATUS_CHARACTER_LIMIT`] -- snapped back to the last newline -- with /// the `... (git status truncated)` marker appended. -fn normalize_git_status(status: &str) -> Option { +pub fn normalize_git_status(status: &str) -> Option { let status = status.trim(); if status.is_empty() { return None; diff --git a/crates/codegen/xai-grok-config/src/config_override.rs b/crates/codegen/xai-grok-config/src/config_override.rs index bf33971..cf5f808 100644 --- a/crates/codegen/xai-grok-config/src/config_override.rs +++ b/crates/codegen/xai-grok-config/src/config_override.rs @@ -75,14 +75,12 @@ pub fn patch_touches_any(patch: &toml::Table, paths: &[PatchPath]) -> bool { paths.iter().any(|p| patch_touches_path(patch, p)) } -/// Keys stripped from every applied patch so an override can't re-introduce a -/// nested `version_overrides`/`campaigns` array (recursive re-injection). This -/// const owns the recursive-injection keys for every override kind; [`apply_patches`] -/// takes the strip list as a parameter so the strip step itself stays key-agnostic. -pub const PATCH_STRIP_KEYS: &[&str] = &["version_overrides", "campaigns"]; +/// Keys stripped from every applied patch: an override cannot re-inject nested +/// `version_overrides`/`campaigns` or define `[auth_provider.*]` command tables. +pub const PATCH_STRIP_KEYS: &[&str] = &["version_overrides", "campaigns", "auth_provider"]; /// Deep-merge each patch in iteration order (later wins on a leaf), stripping -/// `strip_keys` from every patch first. +/// `strip_keys` (top level) first. pub fn apply_patches( config: &mut toml::Value, patches: impl IntoIterator, @@ -133,10 +131,28 @@ mod tests { let mut p = toml::Table::new(); p.insert("version_overrides".into(), toml::Value::Array(vec![])); p.insert("campaigns".into(), toml::Value::Array(vec![])); + p.insert( + "auth_provider".into(), + toml::Value::Table(toml::Table::new()), + ); p.insert("keep".into(), toml::Value::Boolean(true)); apply_patches(&mut cfg2, std::iter::once(p), PATCH_STRIP_KEYS); assert!(cfg2.get("version_overrides").is_none()); assert!(cfg2.get("campaigns").is_none()); + assert!(cfg2.get("auth_provider").is_none()); assert_eq!(cfg2["keep"].as_bool(), Some(true)); + + // Top-level strip only: a model may still reference a local provider by name. + let mut cfg3 = toml::Value::Table(toml::Table::new()); + let p = table( + "[auth_provider.injected]\ncommand = \"evil\"\n\ + [model.x]\nauth_provider = \"local-name\"\n", + ); + apply_patches(&mut cfg3, std::iter::once(p), PATCH_STRIP_KEYS); + assert!(cfg3.get("auth_provider").is_none()); + assert_eq!( + cfg3["model"]["x"]["auth_provider"].as_str(), + Some("local-name") + ); } } diff --git a/crates/codegen/xai-grok-config/src/signed_policy.rs b/crates/codegen/xai-grok-config/src/signed_policy.rs index b7e84b6..c82630d 100644 --- a/crates/codegen/xai-grok-config/src/signed_policy.rs +++ b/crates/codegen/xai-grok-config/src/signed_policy.rs @@ -9,8 +9,8 @@ //! marker stays the (best-effort) authority. use base64::Engine; pub use prod_mc_cli_chat_proxy_types::{ - MANAGED_IDENTITY_TYP, MANAGED_POLICY_TYP, ManagedIdentityClaim, SignatureEnvelope, - SignedPayload, now_unix, + MANAGED_CONFIG_NONCE_ECHO_HEADER, MANAGED_IDENTITY_TYP, MANAGED_POLICY_TYP, + ManagedIdentityClaim, SignatureEnvelope, SignedPayload, is_server_nonce_shape, now_unix, }; /// Compiled-in trusted Ed25519 public keys, `(key_id, raw 32 bytes)`; more than one /// entry only during a rotation. Empty ships dark (see [`verification_active`]). @@ -327,6 +327,27 @@ fn write_envelope_at(path: &std::path::Path, sidecar: &SignatureEnvelope) -> std .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; crate::fs_atomic::write_atomically(path, &json, Some(0o600)) } +/// Persisted envelope nonce for [`MANAGED_CONFIG_NONCE_ECHO_HEADER`] (unverified; +/// telemetry only, never a trust input). Both guards fail open by skipping the +/// echo: only the server mint shape (header-safe, so a corrupt sidecar can't brick +/// the fetch), and only a payload issued to `fetch_principal`. A leftover sidecar +/// from a prior identity must not read as a cross-tenant replay upstream. +pub fn stored_envelope_nonce( + home: &std::path::Path, + fetch_principal: Option<&str>, +) -> Option { + let fetch_principal = fetch_principal?; + let SidecarRead::Present(sidecar) = read_sidecar(home) else { + return None; + }; + let payload: SignedPayload = serde_json::from_str(&sidecar.signed_payload).ok()?; + let issued_to = payload + .deployment_id + .as_deref() + .or(payload.team_id.as_deref()); + (issued_to == Some(fetch_principal) && is_server_nonce_shape(&payload.nonce)) + .then_some(payload.nonce) +} /// Whether an authentic claim IMPOSES fail-closed enforcement: verified, bound to /// the KNOWN `expected_principal`, in-date vs the caller-clamped `now_unix`, and /// `fail_closed`. Anything else imposes nothing: permissive (must not override a diff --git a/crates/codegen/xai-grok-config/src/signed_policy/tests.rs b/crates/codegen/xai-grok-config/src/signed_policy/tests.rs index f903efe..d3a893c 100644 --- a/crates/codegen/xai-grok-config/src/signed_policy/tests.rs +++ b/crates/codegen/xai-grok-config/src/signed_policy/tests.rs @@ -33,6 +33,7 @@ fn payload() -> SignedPayload { requirements: Some("[features]\nweb_fetch = false\n".into()), fail_closed: false, expires_at: 4_000_000_000, + nonce: String::new(), key_id: "v1".into(), } } @@ -888,6 +889,7 @@ fn unknown_signed_key_id_is_rejected() { let home = dir.path(); let (kp, pubkey) = test_keypair(); let p = SignedPayload { + nonce: String::new(), key_id: "v9".into(), fail_closed: true, ..payload() @@ -930,6 +932,7 @@ fn rotation_selects_the_trusted_key_by_signed_key_id() { let v1 = sign(&kp1, &payload()); let v2_payload = SignedPayload { + nonce: String::new(), key_id: "v2".into(), ..payload() }; diff --git a/crates/codegen/xai-grok-hooks/src/dispatcher.rs b/crates/codegen/xai-grok-hooks/src/dispatcher.rs index d396b6c..fb14219 100644 --- a/crates/codegen/xai-grok-hooks/src/dispatcher.rs +++ b/crates/codegen/xai-grok-hooks/src/dispatcher.rs @@ -738,7 +738,8 @@ mod tests { ); assert_eq!(result.results.len(), 1); assert!( - matches!(&result.results[0], HookRunResult::Failed { hook_name, .. } if hook_name == "crasher"), + matches!(&result.results[0], HookRunResult::Failed { hook_name, .. } +if hook_name == "crasher"), "the failure must still appear in run_results for UI scrollback, got {:?}", result.results ); diff --git a/crates/codegen/xai-grok-markdown/src/mermaid.rs b/crates/codegen/xai-grok-markdown/src/mermaid.rs index 61d17f8..84bc3fc 100644 --- a/crates/codegen/xai-grok-markdown/src/mermaid.rs +++ b/crates/codegen/xai-grok-markdown/src/mermaid.rs @@ -3807,11 +3807,14 @@ mod tests { ) .unwrap(); assert!(s.items.iter().any(|it| matches!(it, - SeqItem::Message { text: Some(t), .. } if t.contains("call ") && !t.contains("<")))); + SeqItem::Message { text: Some(t), .. } +if t.contains("call ") && !t.contains("<")))); assert!(s.items.iter().any(|it| matches!(it, - SeqItem::Note { text, .. } if text.contains("memo ") && !text.contains("<")))); + SeqItem::Note { text, .. } +if text.contains("memo ") && !text.contains("<")))); assert!(s.items.iter().any(|it| matches!(it, - SeqItem::Divider { text } if text.contains("c ") && !text.contains("<")))); + SeqItem::Divider { text } +if text.contains("c ") && !text.contains("<")))); // Class members and ER attributes have no clean quoted form (splitter // fragments unquoted `;`; ER drops quoted text as a comment), so exercise diff --git a/crates/codegen/xai-grok-mcp/src/servers.rs b/crates/codegen/xai-grok-mcp/src/servers.rs index c0ebe1f..d4e1d46 100644 --- a/crates/codegen/xai-grok-mcp/src/servers.rs +++ b/crates/codegen/xai-grok-mcp/src/servers.rs @@ -170,7 +170,8 @@ impl InitProgress { /// True iff every per-server handshake has settled and `finish_init` /// has fired. Pairs with [`Self::is_in_progress`]. pub fn is_complete(&self) -> bool { - matches!(self, Self::Finished { handshaking } if handshaking.is_empty()) + matches!(self, Self::Finished { handshaking } +if handshaking.is_empty()) } /// True iff any init work is outstanding — either we are pre- diff --git a/crates/codegen/xai-grok-memory/src/dream.rs b/crates/codegen/xai-grok-memory/src/dream.rs index a10d110..b5920f6 100644 --- a/crates/codegen/xai-grok-memory/src/dream.rs +++ b/crates/codegen/xai-grok-memory/src/dream.rs @@ -867,7 +867,8 @@ mod tests { let result = execute_dream(&lock, &storage, response, 5, 300, &sdir, &[]); assert!( - matches!(result.status, DreamStatus::Completed { chars_written } if chars_written == response.chars().count()) + matches!(result.status, DreamStatus::Completed { chars_written } +if chars_written == response.chars().count()) ); assert_eq!(result.sessions_eligible, 5); assert_eq!(result.cleaned_stems.len(), 0); diff --git a/crates/codegen/xai-grok-pager-minimal/src/commit.rs b/crates/codegen/xai-grok-pager-minimal/src/commit.rs index e6af70f..421ba78 100644 --- a/crates/codegen/xai-grok-pager-minimal/src/commit.rs +++ b/crates/codegen/xai-grok-pager-minimal/src/commit.rs @@ -115,21 +115,18 @@ pub fn is_committable(entry: &ScrollbackEntry, turn_running: bool, is_last: bool } /// The display mode a block should be committed in (minimal mode, print-once). -/// -/// Independent of the interactive `default_display_mode` / `finished_display_mode` -/// because committed scrollback can't be re-folded later (it is static terminal -/// text). The per-type fidelity policy (design decision K9) lives here in one -/// place: messages full, reasoning collapsed-but-expandable, tool output -/// truncated, diffs always full. pub fn minimal_commit_display_mode(block: &RenderBlock) -> DisplayMode { match block { - // Diffs are the key artifact of an edit — always full. RenderBlock::ToolCall(ToolCallBlock::Edit(_)) => DisplayMode::Expanded, - // Other tool calls: truncated (first/last N + hidden-line count). + RenderBlock::ToolCall( + tc @ (ToolCallBlock::Search(_) + | ToolCallBlock::Read(_) + | ToolCallBlock::ListDir(_) + | ToolCallBlock::MemorySearch(_) + | ToolCallBlock::IntegrationSearch(_)), + ) if tc.is_success() => DisplayMode::Collapsed, RenderBlock::ToolCall(_) => DisplayMode::Truncated, - // Reasoning: collapsed marker ("Thought for Xs"); expandable via Ctrl+E. - RenderBlock::Thinking(_) => DisplayMode::Collapsed, - // Messages, system/session events, etc.: full. + RenderBlock::Thinking(_) => DisplayMode::Expanded, _ => DisplayMode::Expanded, } } @@ -1342,7 +1339,7 @@ mod tests { fn commit_display_mode_policy() { assert_eq!( minimal_commit_display_mode(&RenderBlock::thinking("reasoning")), - DisplayMode::Collapsed + DisplayMode::Expanded ); assert_eq!( minimal_commit_display_mode(&RenderBlock::edit("file.rs", None)), @@ -1357,4 +1354,42 @@ mod tests { DisplayMode::Expanded ); } + + #[test] + fn commit_display_mode_lookups_collapse_on_success_only() { + use xai_grok_pager::scrollback::blocks::{ + ListDirToolCallBlock, ReadToolCallBlock, SearchToolCallBlock, + }; + + assert_eq!( + minimal_commit_display_mode(&RenderBlock::search("pat", 3, vec![])), + DisplayMode::Collapsed + ); + assert_eq!( + minimal_commit_display_mode(&RenderBlock::read("src/lib.rs", None)), + DisplayMode::Collapsed + ); + assert_eq!( + minimal_commit_display_mode(&RenderBlock::list_dir_with_output("src", "a.rs\nb.rs")), + DisplayMode::Collapsed + ); + + for failed in [ + RenderBlock::ToolCall(ToolCallBlock::Search( + SearchToolCallBlock::new("pat").with_error("regex parse error"), + )), + RenderBlock::ToolCall(ToolCallBlock::Read( + ReadToolCallBlock::new("gone.rs").with_error("file not found"), + )), + RenderBlock::ToolCall(ToolCallBlock::ListDir( + ListDirToolCallBlock::new("gone/").with_error("no such directory"), + )), + ] { + assert_eq!( + minimal_commit_display_mode(&failed), + DisplayMode::Truncated, + "failed lookup must stay truncated: {failed:?}" + ); + } + } } diff --git a/crates/codegen/xai-grok-pager-pty-harness/src/scroll_matrix/runner.rs b/crates/codegen/xai-grok-pager-pty-harness/src/scroll_matrix/runner.rs index 41cd971..3f5c130 100644 --- a/crates/codegen/xai-grok-pager-pty-harness/src/scroll_matrix/runner.rs +++ b/crates/codegen/xai-grok-pager-pty-harness/src/scroll_matrix/runner.rs @@ -397,21 +397,18 @@ mod tests { assert!(check_quiet(0).is_pass()); assert!(check_quiet(QUIET_MAX_FRAMES).is_pass()); let result = check_quiet(QUIET_MAX_FRAMES + 1); - assert!( - matches!(result, InvariantResult::Violated { ref detail } if detail.contains("churn")) - ); + assert!(matches!(result, InvariantResult::Violated { ref detail } +if detail.contains("churn"))); } #[test] fn screen_rejects_streaming_sessions_and_marker_loss() { let streaming = check_screen(SessionKind::Streaming, 100, Some(100), &[]); - assert!( - matches!(streaming, InvariantResult::Violated { ref detail } if detail.contains("streaming")) - ); + assert!(matches!(streaming, InvariantResult::Violated { ref detail } +if detail.contains("streaming"))); let lost = check_screen(SessionKind::BottomPinned, 100, None, &[]); - assert!( - matches!(lost, InvariantResult::Violated { ref detail } if detail.contains("no marker")) - ); + assert!(matches!(lost, InvariantResult::Violated { ref detail } +if detail.contains("no marker"))); // Empty capture ⇒ no movement expected; a matching marker passes. assert!(check_screen(SessionKind::BottomPinned, 100, Some(100), &[]).is_pass()); let moved = check_screen(SessionKind::BottomPinned, 100, Some(97), &[]); diff --git a/crates/codegen/xai-grok-pager/src/app/acp_handler/mcp.rs b/crates/codegen/xai-grok-pager/src/app/acp_handler/mcp.rs index 29b11a4..14a3c3c 100644 --- a/crates/codegen/xai-grok-pager/src/app/acp_handler/mcp.rs +++ b/crates/codegen/xai-grok-pager/src/app/acp_handler/mcp.rs @@ -157,9 +157,10 @@ pub(super) fn handle_mcp_tools_changed(notif: &acp::ExtNotification, app: &mut A pub(super) fn agent_has_pending_mcps_fetch(app: &AppView, agent_id: AgentId) -> bool { app.pending_effects.iter().any(|e| { matches!( - e, - Effect::FetchMcpsList { agent_id: a, .. } if *a == agent_id - ) + e, + Effect::FetchMcpsList { agent_id: a, .. } + if *a == agent_id + ) }) } diff --git a/crates/codegen/xai-grok-pager/src/app/acp_handler/tests/announcements.rs b/crates/codegen/xai-grok-pager/src/app/acp_handler/tests/announcements.rs index da6bbbc..6bb5e5d 100644 --- a/crates/codegen/xai-grok-pager/src/app/acp_handler/tests/announcements.rs +++ b/crates/codegen/xai-grok-pager/src/app/acp_handler/tests/announcements.rs @@ -94,7 +94,8 @@ assert!( app.pending_effects.iter().any(|e| matches!( e, - Effect::PersistAnnouncementsHidden { hidden_ids } if hidden_ids == &expected + Effect::PersistAnnouncementsHidden { hidden_ids } +if hidden_ids == &expected )), "prune must persist the shrunken set, got {:?}", app.pending_effects diff --git a/crates/codegen/xai-grok-pager/src/app/agent_view/queue.rs b/crates/codegen/xai-grok-pager/src/app/agent_view/queue.rs index a8cdafd..d2be29f 100644 --- a/crates/codegen/xai-grok-pager/src/app/agent_view/queue.rs +++ b/crates/codegen/xai-grok-pager/src/app/agent_view/queue.rs @@ -25,9 +25,11 @@ impl AgentView { // BEFORE the removal so a potential auto-hide pane switch can't hit // the editing lock (see queue_edit.rs ordering invariant). if matches!( - self.prompt_mode, - PromptMode::EditingQueued { id: editing_id, server_id: None, .. } if editing_id == id - ) { + self.prompt_mode, + PromptMode::EditingQueued { id: editing_id, server_id: None, .. } + if editing_id == id + ) + { self.exit_editing_mode(); } self.queue.select_after_delete(id); diff --git a/crates/codegen/xai-grok-pager/src/app/app_view.rs b/crates/codegen/xai-grok-pager/src/app/app_view.rs index 8d29c5f..f10b209 100644 --- a/crates/codegen/xai-grok-pager/src/app/app_view.rs +++ b/crates/codegen/xai-grok-pager/src/app/app_view.rs @@ -343,8 +343,9 @@ impl VoiceState { /// it). `/voice` and toggle-style starts leave this false. pub(crate) fn hold(&self) -> bool { matches!( - self, Self::ColdStart { hold, .. } | Self::Recording { hold, .. } if * hold - ) + self, Self::ColdStart { hold, .. } | Self::Recording { hold, .. } + if * hold + ) } } /// Entry in the session picker list on the welcome screen. diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/queue.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/queue.rs index 8b37cbe..f6eb9a3 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/queue.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/queue.rs @@ -965,7 +965,8 @@ mod tests { // Turn ends → should drain "second" (front, not being edited) + FetchBilling. let effects = dispatch(end_turn(), &mut app); assert_eq!(effects.len(), 2); - assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "second")); + assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } +if text == "second")); assert!(matches!( &effects[1], Effect::FetchBilling { silent: true, .. } @@ -991,7 +992,8 @@ mod tests { // DrainQueue should pop and send. let effects = dispatch(Action::DrainQueue, &mut app); assert_eq!(effects.len(), 1); - assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "queued")); + assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } +if text == "queued")); assert_eq!(app.agents[&id].session.queue_len(), 0); } @@ -1960,7 +1962,8 @@ mod tests { let effects = dispatch(Action::DrainQueue, &mut app); assert_eq!(effects.len(), 1); assert!( - matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "p3-edited"), + matches!(&effects[0], Effect::SendPrompt { text, .. } +if text == "p3-edited"), "should send the edited prompt, got: {:?}", effects[0] ); diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/session/lifecycle.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/session/lifecycle.rs index 8f75b09..3a382b6 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/session/lifecycle.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/session/lifecycle.rs @@ -790,10 +790,6 @@ pub(in crate::app::dispatch) fn skip_picker_and_create_session( model_id: None, preferred_session_id, chat_kind, - - - - }] } pub(in crate::app::dispatch) fn handle_session_created( diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/billing.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/billing.rs index 358b126..09e7e85 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/billing.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/billing.rs @@ -492,7 +492,8 @@ fn show_usage_returns_fetch_billing_effect() { // together and renders a single summary. assert_eq!(effects.len(), 1, "got: {effects:?}"); assert!( - matches!(&effects[0], Effect::FetchBilling { agent_id, silent } if *agent_id == AgentId(0) && !*silent), + matches!(&effects[0], Effect::FetchBilling { agent_id, silent } +if *agent_id == AgentId(0) && !*silent), "effect should be a non-silent FetchBilling, got: {effects:?}" ); } @@ -859,7 +860,8 @@ fn free_usage_failure_opens_paywall_modal() { // 1. Real send. let effects = dispatch(Action::SendPrompt("draw me a cat".into()), &mut app); assert!( - matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "draw me a cat"), + matches!(&effects[0], Effect::SendPrompt { text, .. } +if text == "draw me a cat"), "send must dispatch: {effects:?}" ); let prompt_id = app.agents[&id].session.current_prompt_id.clone(); @@ -1051,7 +1053,8 @@ fn unknown_non_restricted_command_still_passes_through() { assert_eq!(effects.len(), 1); assert!( - matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "/frobnicate arg"), + matches!(&effects[0], Effect::SendPrompt { text, .. } +if text == "/frobnicate arg"), "unknown command must still pass through: {effects:?}" ); assert!( diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/cta_e2e.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/cta_e2e.rs index b8aee70..906d534 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/cta_e2e.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/cta_e2e.rs @@ -218,9 +218,10 @@ fn plugin_cta_catalog_load_recomputes_match_for_typed_draft() { &mut app, ); assert!(matches!( - &app.agents[&id].plugin_cta.phase, - CtaPhase::Matched { name, .. } if name == "zzctaplugin" - )); + &app.agents[&id].plugin_cta.phase, + CtaPhase::Matched { name, .. } + if name == "zzctaplugin" + )); } #[test] diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/dashboard.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/dashboard.rs index 3ee1e26..a78b234 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/dashboard.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/dashboard.rs @@ -568,7 +568,8 @@ async fn dashboard_change_location_valid_updates_cwd_and_closes_modal() { assert!( effects .iter() - .any(|e| matches!(e, Effect::SetWorkingDir { path } if path == &target)) + .any(|e| matches!(e, Effect::SetWorkingDir { path } +if path == &target)) ); assert!( app.dashboard.as_ref().unwrap().location_picker.is_none(), @@ -2187,9 +2188,10 @@ fn dashboard_dispatch_applies_pending_model_and_plan() { let new_id = *app.agents.keys().next().unwrap(); // CreateSession carries the staged model id. assert!(effects.iter().any(|e| matches!( - e, - Effect::CreateSession { model_id: Some(m), .. } if *m == model_id - ))); + e, + Effect::CreateSession { model_id: Some(m), .. } + if *m == model_id + ))); let agent = &app.agents[&new_id]; assert_eq!( agent.session.deferred_model_switch, @@ -2230,9 +2232,10 @@ fn dashboard_new_agent_button_applies_pending_model_and_plan() { let new_id = *app.agents.keys().next().unwrap(); // CreateSession carries the staged model id. assert!(effects.iter().any(|e| matches!( - e, - Effect::CreateSession { model_id: Some(m), .. } if *m == model_id - ))); + e, + Effect::CreateSession { model_id: Some(m), .. } + if *m == model_id + ))); let agent = &app.agents[&new_id]; assert_eq!( agent.session.deferred_model_switch, @@ -2277,7 +2280,8 @@ fn dashboard_deferred_plan_mode_applied_on_session_created() { assert!( effects .iter() - .any(|e| matches!(e, Effect::SetSessionMode { session_id: s, .. } if *s == session_id)), + .any(|e| matches!(e, Effect::SetSessionMode { session_id: s, .. } +if *s == session_id)), "SessionCreated must emit SetSessionMode for the deferred plan mode" ); } @@ -5077,7 +5081,8 @@ fn dashboard_peek_reply_to_idle_agent_sends() { /* attach */ false, ); assert_eq!(effects.len(), 1); - assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "please continue")); + assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } +if text == "please continue")); assert!(app.agents[&AgentId(0)].session.state.is_turn_running()); assert_eq!(app.agents[&AgentId(0)].session.queue_len(), 0); // Reply draft cleared after sending. diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/modes.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/modes.rs index b883f42..d358ca5 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/modes.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/modes.rs @@ -161,7 +161,8 @@ fn slash_plan_no_args_not_in_plan_enters_plan_mode() { // Should emit SetSessionMode to enter plan mode. assert_eq!(effects.len(), 1); assert!( - matches!(&effects[0], Effect::SetSessionMode { mode_id, .. } if &*mode_id.0 == "plan"), + matches!(&effects[0], Effect::SetSessionMode { mode_id, .. } +if &*mode_id.0 == "plan"), "expected SetSessionMode(plan), got: {effects:?}" ); // Optimistic pending state should be set. @@ -1932,9 +1933,10 @@ fn cycle_always_approve_with_nudge_jumps_to_plan() { ); assert!( effects.iter().any(|e| matches!( - e, - Effect::SetSessionMode { mode_id, .. } if &*mode_id.0 == "plan" - )), + e, + Effect::SetSessionMode { mode_id, .. } + if &*mode_id.0 == "plan" + )), "expected SetSessionMode(plan), got {effects:?}" ); assert!( @@ -1988,9 +1990,10 @@ fn cycle_auto_with_nudge_jumps_to_plan() { ); assert!( effects.iter().any(|e| matches!( - e, - Effect::SetSessionMode { mode_id, .. } if &*mode_id.0 == "plan" - )), + e, + Effect::SetSessionMode { mode_id, .. } + if &*mode_id.0 == "plan" + )), "expected SetSessionMode(plan), got {effects:?}" ); assert!( @@ -2312,7 +2315,8 @@ fn set_plan_mode_idempotency_uses_pending_over_active() { "OFF from EFFECTIVE-ON must emit Effect::SetSessionMode (not idempotent)" ); assert!( - matches!(&effects[0], Effect::SetSessionMode { mode_id, .. } if &*mode_id.0 == "default"), + matches!(&effects[0], Effect::SetSessionMode { mode_id, .. } +if &*mode_id.0 == "default"), "OFF transition must emit SetSessionMode(default): {effects:?}" ); let agent = app.agents.get(&AgentId(0)).unwrap(); diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/prompt.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/prompt.rs index 28bdc37..28c958c 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/prompt.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/prompt.rs @@ -534,7 +534,8 @@ fn send_prompt_produces_effect_and_clears_input() { // Prompt is enqueued and immediately drained (agent was idle). assert_eq!(effects.len(), 1); - assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "hello")); + assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } +if text == "hello")); assert!(app.agents[&id].prompt.text().is_empty()); assert!(app.agents[&id].session.state.is_turn_running()); assert_eq!(app.agents[&id].scrollback.len(), 1); @@ -754,7 +755,8 @@ fn chip_submit_while_enqueued_clears_follow_up_chips() { assert!( !effects .iter() - .any(|e| matches!(e, Effect::SendPrompt { text, .. } if text == "Summarize")), + .any(|e| matches!(e, Effect::SendPrompt { text, .. } +if text == "Summarize")), "chip must be enqueued, not immediate-sent, got {effects:?}" ); // The chips are cleared on the enqueue path too (the bug fix). @@ -979,7 +981,8 @@ fn multiple_queued_prompts_drain_one_per_turn() { // Turn end → drain "b" + FetchBilling. let effects = dispatch(end_turn(), &mut app); - assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "b")); + assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } +if text == "b")); assert!(matches!( &effects[1], Effect::FetchBilling { silent: true, .. } @@ -988,7 +991,8 @@ fn multiple_queued_prompts_drain_one_per_turn() { // Turn end → drain "c" + FetchBilling. let effects = dispatch(end_turn(), &mut app); - assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "c")); + assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } +if text == "c")); assert!(matches!( &effects[1], Effect::FetchBilling { silent: true, .. } @@ -1825,7 +1829,8 @@ fn cancel_with_queued_prompt_drains_on_completion() { ); assert_eq!(effects.len(), 2); - assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "queued")); + assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } +if text == "queued")); assert!(matches!( &effects[1], Effect::FetchBilling { silent: true, .. } @@ -1907,7 +1912,8 @@ fn cancel_with_multiple_queued_prompts_drains_only_front_prompt() { ); assert_eq!(effects.len(), 2); - assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "queued-1")); + assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } +if text == "queued-1")); assert!(matches!( &effects[1], Effect::FetchBilling { silent: true, .. } @@ -2172,7 +2178,8 @@ fn slash_unknown_command_passthrough_enqueues_prompt() { let effects = dispatch(Action::SendPrompt("/unknown-cmd arg1".into()), &mut app); // Unknown slash command → PassThrough → enqueue as prompt. assert_eq!(effects.len(), 1); - assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "/unknown-cmd arg1")); + assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } +if text == "/unknown-cmd arg1")); assert!(app.agents[&id].prompt.text().is_empty()); } @@ -2184,7 +2191,8 @@ fn non_slash_prompt_still_works() { let effects = dispatch(Action::SendPrompt("hello world".into()), &mut app); assert_eq!(effects.len(), 1); - assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "hello world")); + assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } +if text == "hello world")); assert!(app.agents[&id].prompt.text().is_empty()); } diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/rewind.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/rewind.rs index 9830764..e7159f8 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/rewind.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/rewind.rs @@ -67,7 +67,8 @@ fn rewind_then_resubmit_drains_immediately_and_discards_orphan() { // User edits and re-submits without waiting. let effects = dispatch(Action::SendPrompt("second".into()), &mut app); assert_eq!(effects.len(), 1); - assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "second")); + assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } +if text == "second")); assert!(app.agents[&id].session.state.is_turn_running()); let second_pid = app.agents[&id].session.current_prompt_id.clone(); assert!(second_pid.is_some()); @@ -379,9 +380,10 @@ fn inline_edit_conversation_only_success_resubmits_and_closes_editor() { ); assert!( - effects.iter().any( - |e| matches!(e, Effect::SendPrompt { text, .. } if text == "fix the bug properly") - ), + effects + .iter() + .any(|e| matches!(e, Effect::SendPrompt { text, .. } +if text == "fix the bug properly")), "edited prompt must be sent, got {effects:?}" ); let agent = &app.agents[&id]; @@ -467,9 +469,10 @@ fn inline_edit_all_mode_previews_confirms_and_resubmits() { &mut app, ); assert!( - effects.iter().any( - |e| matches!(e, Effect::SendPrompt { text, .. } if text == "fix the bug properly") - ), + effects + .iter() + .any(|e| matches!(e, Effect::SendPrompt { text, .. } +if text == "fix the bug properly")), "got {effects:?}" ); assert!( @@ -722,9 +725,10 @@ fn inline_edit_resubmit_sends_slash_text_literally() { ); assert!( - effects.iter().any( - |e| matches!(e, Effect::SendPrompt { text, .. } if text == "/etc/hosts is wrong, fix it") - ), + effects + .iter() + .any(|e| matches!(e, Effect::SendPrompt { text, .. } +if text == "/etc/hosts is wrong, fix it")), "slash-lookalike edit must be sent as a prompt, got {effects:?}" ); } diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/router.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/router.rs index 17e8635..37efa04 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/router.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/router.rs @@ -505,7 +505,8 @@ fn dispatch_send_prompt_announcements_via_registry() { effects .iter() .any(|e| matches!(e, Effect::PersistAnnouncementsHidden { - hidden_ids } if hidden_ids.contains("crit-a"))), + hidden_ids } +if hidden_ids.contains("crit-a"))), "expected persist effect carrying the hidden id, got {effects:?}" ); assert!(app.hidden_announcement_ids.contains("crit-a")); @@ -569,7 +570,8 @@ fn announcements_show_clears_visible_critical_ids_only() { effects .iter() .any(|e| matches!(e, Effect::PersistAnnouncementsHidden { - hidden_ids } if ! hidden_ids.contains("outage-a"))), + hidden_ids } +if ! hidden_ids.contains("outage-a"))), "expected persist effect without the un-hidden id, got {effects:?}" ); assert_eq!(shown_banner_id(&app).as_deref(), Some("outage-a")); @@ -661,7 +663,8 @@ fn announcements_show_clears_hidden_promo_ids() { effects .iter() .any(|e| matches!(e, Effect::PersistAnnouncementsHidden { - hidden_ids } if ! hidden_ids.contains("promo-a"))), + hidden_ids } +if ! hidden_ids.contains("promo-a"))), "expected persist effect without the un-hidden promo id, got {effects:?}" ); assert_eq!(shown_banner_id(&app).as_deref(), Some("promo-a")); @@ -685,7 +688,8 @@ fn switch_model_dispatch_produces_effect_and_sets_pending() { ); assert_eq!(effects.len(), 1); assert!( - matches!(& effects[0], Effect::SwitchModel { model_id : mid, .. } if mid == & + matches!(& effects[0], Effect::SwitchModel { model_id : mid, .. } +if mid == & model_id) ); assert!(app.agents[&id].session.model_switch_pending); @@ -706,7 +710,8 @@ fn switch_model_allowed_when_agent_chat_kind() { ); assert_eq!(effects.len(), 1); assert!( - matches!(& effects[0], Effect::SwitchModel { model_id : mid, .. } if mid == & + matches!(& effects[0], Effect::SwitchModel { model_id : mid, .. } +if mid == & model_id) ); assert!(app.agents[&id].session.model_switch_pending); @@ -726,7 +731,8 @@ fn switch_model_allowed_when_app_chat_mode() { ); assert_eq!(effects.len(), 1); assert!( - matches!(& effects[0], Effect::SwitchModel { model_id : mid, .. } if mid == & + matches!(& effects[0], Effect::SwitchModel { model_id : mid, .. } +if mid == & model_id) ); assert!(app.agents[&id].session.model_switch_pending); @@ -946,7 +952,8 @@ fn acp_bootstrap_command_executes_as_passthrough() { let effects = dispatch(Action::SendPrompt("/flush".into()), &mut app); assert_eq!(effects.len(), 1); assert!( - matches!(& effects[0], Effect::SendPrompt { text, .. } if text == "/flush"), + matches!(& effects[0], Effect::SendPrompt { text, .. } +if text == "/flush"), "ACP command should passthrough, got: {effects:?}" ); } @@ -1081,7 +1088,8 @@ fn acp_command_with_args_passthrough_includes_args() { let effects = dispatch(Action::SendPrompt("/search find bugs".into()), &mut app); assert_eq!(effects.len(), 1); assert!( - matches!(& effects[0], Effect::SendPrompt { text, .. } if text == + matches!(& effects[0], Effect::SendPrompt { text, .. } +if text == "/search find bugs"), "ACP passthrough should preserve args, got: {effects:?}" ); @@ -1324,7 +1332,8 @@ fn view_catalog_entry_emits_fetch_effect() { ); assert_eq!(effects.len(), 1); assert!( - matches!(& effects[0], Effect::FetchCatalogEntry { kind, name } if kind == + matches!(& effects[0], Effect::FetchCatalogEntry { kind, name } +if kind == "persona" && name == "researcher") ); } diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/fork.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/fork.rs index 591c1e1..fd2ee2a 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/fork.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/fork.rs @@ -96,9 +96,10 @@ fn worktree_forked_with_restore_shows_summary_in_scrollback() { // Should emit LoadSession. assert_eq!(effects.len(), 1); assert!(matches!( - &effects[0], - Effect::LoadSession { session_id, .. } if session_id == "forked-sess-2" - )); + &effects[0], + Effect::LoadSession { session_id, .. } + if session_id == "forked-sess-2" + )); // Scrollback should contain the restore summary. let has_restore_msg = app.agents[&id] .scrollback diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/lifecycle.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/lifecycle.rs index 72be456..ecdcc0a 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/lifecycle.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/lifecycle.rs @@ -137,7 +137,8 @@ fn session_created_sets_session_id() { ); assert_eq!(effects.len(), 7); assert!( - matches!(& effects[0], Effect::FetchPromptHistory { session_id, .. } if + matches!(& effects[0], Effect::FetchPromptHistory { session_id, .. } +if session_id == "new-session-123") ); assert!(matches!(&effects[1], Effect::FetchSessionAgentName { .. })); @@ -491,7 +492,8 @@ fn worktree_session_created_drains_queued_prompts() { assert!( effects .iter() - .any(|e| matches!(e, Effect::SendPrompt { text, .. } if text == + .any(|e| matches!(e, Effect::SendPrompt { text, .. } +if text == "hello")) ); assert!( @@ -525,7 +527,8 @@ fn session_created_drains_queued_prompts() { assert!( effects .iter() - .any(|e| matches!(e, Effect::SendPrompt { text, .. } if text == + .any(|e| matches!(e, Effect::SendPrompt { text, .. } +if text == "queued msg")) ); assert!( @@ -825,7 +828,8 @@ fn deferred_model_switch_applied_on_session_created() { effects .iter() .any(|e| matches!(e, Effect::SwitchModel { agent_id : a_id, - session_id : s_id, model_id : m_id, .. } if * a_id == id && * s_id == session_id + session_id : s_id, model_id : m_id, .. } +if * a_id == id && * s_id == session_id && * m_id == model_id)) ); } @@ -864,7 +868,8 @@ fn deferred_model_switch_applied_on_worktree_session_created() { effects .iter() .any(|e| matches!(e, Effect::SwitchModel { agent_id : a_id, - session_id : s_id, model_id : m_id, .. } if * a_id == id && * s_id == session_id + session_id : s_id, model_id : m_id, .. } +if * a_id == id && * s_id == session_id && * m_id == model_id)) ); } @@ -1164,7 +1169,8 @@ fn deferred_worktree_ref_replays_through_gate() { effects .iter() .any(|e| matches!(e, Effect::CreateWorktreeSession { git_ref : - Some(r), .. } if r == "feature-branch")), + Some(r), .. } +if r == "feature-branch")), "the deferred --worktree replays with its git ref", ); assert!( @@ -1225,7 +1231,8 @@ fn gated_worktree_without_load_id_preserves_stashed_resume() { effects .iter() .any(|e| matches!(e, Effect::CreateWorktreeSession { - load_session_id : Some(id), .. } if id == "resume-me")), + load_session_id : Some(id), .. } +if id == "resume-me")), "the deferred worktree replays with the preserved resume id", ); assert!(app.deferred_startup.session.is_none()); @@ -1297,7 +1304,8 @@ fn gated_worktree_with_none_companions_preserves_stashed_label_and_ref() { effects .iter() .any(|e| matches!(e, Effect::CreateWorktreeSession { - load_session_id : Some(id), label : Some(l), git_ref : Some(r), .. } if id == + load_session_id : Some(id), label : Some(l), git_ref : Some(r), .. } +if id == "mysess" && l == "mylabel" && r == "featbranch")), "the deferred worktree replays with the preserved id, label, and ref", ); @@ -1411,7 +1419,8 @@ fn auth_complete_retries_stashed_prompt_after_mid_session_login() { assert!( effects .iter() - .any(|e| matches!(e, Effect::SendPrompt { text, .. } if text == + .any(|e| matches!(e, Effect::SendPrompt { text, .. } +if text == "retry me")), "the stashed prompt must be auto-resubmitted, got: {effects:?}" ); @@ -1604,13 +1613,15 @@ async fn project_selected_creates_session_and_sends_prompt() { assert!( effects .iter() - .any(|e| matches!(e, Effect::SetWorkingDir { path } if path == & + .any(|e| matches!(e, Effect::SetWorkingDir { path } +if path == & selected)) ); assert!( effects .iter() - .any(|e| matches!(e, Effect::CreateSession { cwd, .. } if cwd == + .any(|e| matches!(e, Effect::CreateSession { cwd, .. } +if cwd == & selected)) ); assert_eq!(app.agents[&id].session.queue_len(), 1); @@ -1837,7 +1848,8 @@ fn set_plan_mode_on_from_off_emits_set_session_mode() { ); assert_eq!(effects.len(), 1); assert!( - matches!(& effects[0], Effect::SetSessionMode { mode_id, .. } if &* mode_id.0 == + matches!(& effects[0], Effect::SetSessionMode { mode_id, .. } +if &* mode_id.0 == "plan"), "expected SetSessionMode(plan), got: {effects:?}" ); diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/load.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/load.rs index cdcf1ad..a96eabd 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/load.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/load.rs @@ -629,7 +629,8 @@ fn resume_known_session_id_loads_not_creates() { assert!( effects .iter() - .any(|e| matches!(e, Effect::LoadSession { session_id, .. } if + .any(|e| matches!(e, Effect::LoadSession { session_id, .. } +if session_id == "resume-known-id")), "expected LoadSession, got {effects:?}" ); @@ -964,7 +965,8 @@ fn resume_unknown_session_still_creates_new_agent() { effects .iter() .any(|e| matches!(e, Effect::LoadSession { agent_id, session_id, - .. } if * agent_id == new_id && session_id == "sess-never-open")) + .. } +if * agent_id == new_id && session_id == "sess-never-open")) ); } /// Stale `attached_agent` (not equal to visible agent) must not re-arm overlay. @@ -1027,7 +1029,8 @@ fn resume_conversation_does_not_focus_build_id_collision() { effects .iter() .any(|e| matches!(e, Effect::LoadSession { session_id, chat_kind - : true, .. } if session_id == "shared-id")) + : true, .. } +if session_id == "shared-id")) ); assert!(!app.agents[&agent_0].chat_kind); } @@ -1165,7 +1168,8 @@ fn resume_after_load_failed_reissues_load() { assert!( effects .iter() - .any(|e| matches!(e, Effect::LoadSession { agent_id, .. } if * + .any(|e| matches!(e, Effect::LoadSession { agent_id, .. } +if * agent_id == agent_0)) ); assert!(app.agents[&agent_0].loading_placeholder_id.is_some()); @@ -1188,7 +1192,8 @@ fn resume_after_load_failed_reissues_load() { effects .iter() .any(|e| matches!(e, Effect::LoadSession { agent_id, session_id, - .. } if * agent_id != agent_0 && session_id == "fail-then-retry")), + .. } +if * agent_id != agent_0 && session_id == "fail-then-retry")), "retry after failure must emit LoadSession for a new agent, got {effects:?}" ); assert_eq!(app.agents.len(), count_before + 1); diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/settings.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/settings.rs index 66811b0..125f72b 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/settings.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/settings.rs @@ -295,7 +295,8 @@ fn slash_model_valid_dispatches_set_default_model_with_switch_and_persist() { effects[0], ); assert!( - matches!(& effects[1], Effect::SwitchModel { model_id : mid, .. } if mid == & + matches!(& effects[1], Effect::SwitchModel { model_id : mid, .. } +if mid == & model_id), "second effect must be SwitchModel(), got {:?}", effects[1], @@ -993,7 +994,8 @@ fn clear_default_model_persists_but_keeps_live_current() { ); assert!( matches!(& effects[0], Effect::PersistSetting { key : "default_model", value : - crate ::settings::SettingValue::String(s), .. } if s.is_empty()), + crate ::settings::SettingValue::String(s), .. } +if s.is_empty()), "expected PersistSetting(default_model, ''), got {:?}", effects[0], ); @@ -1028,9 +1030,13 @@ fn set_default_model_resolves_known_name() { assert_eq!(effects.len(), 2); assert!( matches!(& effects[0], Effect::PersistSetting { key : "default_model", value : - crate ::settings::SettingValue::String(s), .. } if s == "grok-4.5") + crate ::settings::SettingValue::String(s), .. } +if s == "grok-4.5") + ); + assert!( + matches!(& effects[1], Effect::SwitchModel { model_id : mid, .. } +if mid == & id) ); - assert!(matches!(& effects[1], Effect::SwitchModel { model_id : mid, .. } if mid == & id)); assert_eq!(app.agents[&agent_id].session.models.current, Some(id)); } /// Re-dispatching the same model diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/task_result.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/task_result.rs index 3348163..5f4fd20 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/task_result.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/task_result.rs @@ -660,9 +660,10 @@ fn switch_model_complete_success_updates_model_and_pushes_message() { // PersistPreferredModel effect emitted. assert_eq!(effects.len(), 1); assert!(matches!( - &effects[0], - Effect::PersistPreferredModel { model_id: mid, .. } if *mid == model_id.clone() - )); + &effects[0], + Effect::PersistPreferredModel { model_id: mid, .. } + if *mid == model_id.clone() + )); } #[test] diff --git a/crates/codegen/xai-grok-pager/src/app/effects/tests.rs b/crates/codegen/xai-grok-pager/src/app/effects/tests.rs index 299b777..4a7484b 100644 --- a/crates/codegen/xai-grok-pager/src/app/effects/tests.rs +++ b/crates/codegen/xai-grok-pager/src/app/effects/tests.rs @@ -249,7 +249,8 @@ fn parse_subagent_kill_outcome_reads_typed_outcome() { ); assert!( matches!(parse_subagent_kill_outcome(r#"{"result":{"subagentId":"sa-1","cancelled":false,"outcome":{"kind":"already_finished","status":"completed"}}}"#), - SubagentKillOutcome::NothingLive { status : Some(s) } if s == "completed") + SubagentKillOutcome::NothingLive { status : Some(s) } +if s == "completed") ); assert!( matches!(parse_subagent_kill_outcome(r#"{"result":{"subagentId":"sa-1","cancelled":false,"outcome":{"kind":"not_found"}}}"#), @@ -302,7 +303,8 @@ fn parse_subagent_kill_outcome_round_trips_agent_serialization() { .unwrap(); assert!( matches!(parse_subagent_kill_outcome(& wire), SubagentKillOutcome::NothingLive { - status : Some(s) } if s == "failed") + status : Some(s) } +if s == "failed") ); } /// A top-level payload (no `result` envelope), error envelopes, and @@ -1200,7 +1202,8 @@ async fn check_marketplace_updates_dispatches_update_and_skips_failed_notificati MarketplaceAction::Update { source_url_or_path, plugin_relative_path, - } if source_url_or_path == "https://example.com/plugins.git" + } +if source_url_or_path == "https://example.com/plugins.git" && plugin_relative_path == "plugins/test-plugin" => { saw_update_for_task.store(true, Ordering::SeqCst); } diff --git a/crates/codegen/xai-grok-pager/src/app/leader_cluster/scenarios.rs b/crates/codegen/xai-grok-pager/src/app/leader_cluster/scenarios.rs index 8cf8724..756bc47 100644 --- a/crates/codegen/xai-grok-pager/src/app/leader_cluster/scenarios.rs +++ b/crates/codegen/xai-grok-pager/src/app/leader_cluster/scenarios.rs @@ -258,9 +258,11 @@ fn leader_kill_reconnect_reloads_without_duplicating_history() { let deadline = tokio::time::Instant::now() + Duration::from_secs(30); loop { if matches!( - *status_rx.borrow_and_update(), - ConnectionStatus::Connected { generation } if generation >= 1 - ) { + *status_rx.borrow_and_update(), + ConnectionStatus::Connected { generation } + if generation >= 1 + ) + { break; } assert!( diff --git a/crates/codegen/xai-grok-pager/src/headless.rs b/crates/codegen/xai-grok-pager/src/headless.rs index 5b8a655..bdfaf0b 100644 --- a/crates/codegen/xai-grok-pager/src/headless.rs +++ b/crates/codegen/xai-grok-pager/src/headless.rs @@ -2086,9 +2086,10 @@ mod tests { }), ); assert!(matches!( - handle_ext_notification(¬if, OutputFormat::Plain), - ExtEvent::TaskBackgrounded { task_id, is_monitor: false } if task_id == "task-abc" - )); + handle_ext_notification(¬if, OutputFormat::Plain), + ExtEvent::TaskBackgrounded { task_id, is_monitor: false } + if task_id == "task-abc" + )); } #[test] @@ -2102,9 +2103,10 @@ mod tests { }), ); assert!(matches!( - handle_ext_notification(¬if, OutputFormat::Plain), - ExtEvent::TaskBackgrounded { task_id, is_monitor: true } if task_id == "mon-1" - )); + handle_ext_notification(¬if, OutputFormat::Plain), + ExtEvent::TaskBackgrounded { task_id, is_monitor: true } + if task_id == "mon-1" + )); } #[test] @@ -2122,9 +2124,10 @@ mod tests { }), ); assert!(matches!( - handle_ext_notification(¬if, OutputFormat::Plain), - ExtEvent::TaskCompleted { task_id } if task_id == "task-abc" - )); + handle_ext_notification(¬if, OutputFormat::Plain), + ExtEvent::TaskCompleted { task_id } + if task_id == "task-abc" + )); } #[test] @@ -2141,9 +2144,10 @@ mod tests { }), ); assert!(matches!( - handle_ext_notification(&spawned, OutputFormat::Plain), - ExtEvent::SubagentSpawned { subagent_id } if subagent_id == "sub-1" - )); + handle_ext_notification(&spawned, OutputFormat::Plain), + ExtEvent::SubagentSpawned { subagent_id } + if subagent_id == "sub-1" + )); let finished = make_ext_notif( "x.ai/session_notification", serde_json::json!({ @@ -2157,9 +2161,10 @@ mod tests { }), ); assert!(matches!( - handle_ext_notification(&finished, OutputFormat::Plain), - ExtEvent::SubagentFinished { subagent_id } if subagent_id == "sub-1" - )); + handle_ext_notification(&finished, OutputFormat::Plain), + ExtEvent::SubagentFinished { subagent_id } + if subagent_id == "sub-1" + )); } #[test] diff --git a/crates/codegen/xai-grok-pager/src/plugin_cmd.rs b/crates/codegen/xai-grok-pager/src/plugin_cmd.rs index 7212d23..afac4c8 100644 --- a/crates/codegen/xai-grok-pager/src/plugin_cmd.rs +++ b/crates/codegen/xai-grok-pager/src/plugin_cmd.rs @@ -888,9 +888,10 @@ fn marketplace_add( if u.trim_end_matches(".git") == normalized) }) } - MarketplaceAddInput::LocalPath(path) => sources - .iter() - .any(|s| matches!(&s.kind, SourceKind::Local { path: p } if p == path)), + MarketplaceAddInput::LocalPath(path) => sources.iter().any(|s| { + matches!(&s.kind, SourceKind::Local { path: p } +if p == path) + }), }; if already_configured { bail!("Marketplace source already configured: {identity}"); diff --git a/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/execute.rs b/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/execute.rs index aa056ae..e38ec2a 100644 --- a/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/execute.rs +++ b/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/execute.rs @@ -579,10 +579,13 @@ impl ExecuteToolCallBlock { .with_joiner(joiner.clone()), ); } - // Ellipsis (non-selectable, breaks range continuity) + let hidden = total - threshold; lines.push( - BlockLine::separator(Line::from(Span::styled("\u{2026}", theme.muted()))) - .with_panel_background(theme.bg_dark), + BlockLine::separator(Line::from(Span::styled( + format!("\u{2026} +{hidden} lines"), + theme.muted(), + ))) + .with_panel_background(theme.bg_dark), ); // Last M lines: range base + 1 (distinct from first chunk) for (wrapped_line, joiner) in diff --git a/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/list_dir.rs b/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/list_dir.rs index 23e303a..9dd8303 100644 --- a/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/list_dir.rs +++ b/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/list_dir.rs @@ -97,9 +97,6 @@ impl ListDirToolCallBlock { self.output = output.into(); } - /// Render collapsed line: `List path`. - /// - /// When `width` is provided, the path is fish-shortened to fit. fn collapsed_line(&self, theme: &Theme, muted: bool, width: Option) -> Line<'static> { let text_style = if muted { theme.muted() @@ -114,15 +111,32 @@ impl ListDirToolCallBlock { }; let prefix = "List "; + let entry_count = self.output.lines().filter(|l| !l.trim().is_empty()).count(); + let suffix = if self.error.is_none() && entry_count > 0 { + let s = if entry_count == 1 { "y" } else { "ies" }; + format!(" ({entry_count} entr{s})") + } else { + String::new() + }; + let suffix_fits = width.is_none_or(|w| prefix.len() + suffix.len() < w); + let effective_suffix = if suffix_fits { suffix.as_str() } else { "" }; + let path_budget = width - .map(|w| w.saturating_sub(prefix.len())) + .map(|w| { + w.saturating_sub(prefix.len()) + .saturating_sub(effective_suffix.len()) + }) .unwrap_or(usize::MAX); let path = crate::render::tool_paths::shorten_path(&self.path, path_budget); - Line::from(vec![ + let mut spans = vec![ Span::styled(prefix, bold_style), Span::styled(path, path_style), - ]) + ]; + if !effective_suffix.is_empty() { + spans.push(Span::styled(effective_suffix.to_string(), theme.muted())); + } + Line::from(spans) } /// Header line with only the path span selectable (exclude "List " prefix). @@ -224,3 +238,49 @@ impl BlockContent for ListDirToolCallBlock { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::scrollback::types::BlockContext; + + fn ctx() -> BlockContext { + BlockContext { + width: 80, + mode: DisplayMode::Collapsed, + is_running: false, + raw: false, + max_lines: None, + appearance: Default::default(), + is_selected: false, + cwd: None, + } + } + + fn header_text(block: &ListDirToolCallBlock) -> String { + block.output(&ctx()).lines[0] + .content + .spans + .iter() + .map(|s| s.content.as_ref()) + .collect() + } + + #[test] + fn collapsed_header_shows_entry_count() { + let block = ListDirToolCallBlock::new("src").with_output("a.rs\nb.rs\nsub/\n"); + assert_eq!(header_text(&block), "List src (3 entries)"); + + let single = ListDirToolCallBlock::new("src").with_output("lonely.rs\n"); + assert_eq!(header_text(&single), "List src (1 entry)"); + } + + #[test] + fn collapsed_header_omits_count_when_empty_or_failed() { + let empty = ListDirToolCallBlock::new("src"); + assert_eq!(header_text(&empty), "List src"); + + let failed = ListDirToolCallBlock::new("gone").with_error("no such directory"); + assert_eq!(header_text(&failed), "List gone"); + } +} diff --git a/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/mod.rs b/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/mod.rs index cdd22ef..7539ca5 100644 --- a/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/mod.rs +++ b/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/mod.rs @@ -342,6 +342,25 @@ impl ToolCallBlock { } } + /// Whether the tool call finished without an error. + pub fn is_success(&self) -> bool { + match self { + ToolCallBlock::Execute(b) => b.is_success(), + ToolCallBlock::Read(b) => b.is_success(), + ToolCallBlock::Edit(b) => b.is_success(), + ToolCallBlock::Search(b) => b.is_success(), + ToolCallBlock::ListDir(b) => b.is_success(), + ToolCallBlock::WebFetch(b) => b.is_success(), + ToolCallBlock::WebSearch(b) => b.is_success(), + ToolCallBlock::IntegrationSearch(b) => b.is_success(), + ToolCallBlock::UseTool(b) => b.is_success(), + ToolCallBlock::MemorySearch(b) => b.is_success(), + ToolCallBlock::Skill(b) => b.is_success(), + ToolCallBlock::Other(b) => b.is_success(), + ToolCallBlock::Lifecycle(_) => true, + } + } + /// Set `started_at` on the inner variant block. /// /// Unlike `transfer_timing_from`, this works across variant boundaries diff --git a/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/use_tool.rs b/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/use_tool.rs index 7b71d55..e78d2bd 100644 --- a/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/use_tool.rs +++ b/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/use_tool.rs @@ -11,8 +11,8 @@ use crate::scrollback::types::{ }; use crate::theme::Theme; -/// Max lines of output shown inline before truncation. const MAX_INLINE_LINES: usize = 10; +const TRUNCATED_INLINE_LINES: usize = 3; /// Use tool call — dispatching to an MCP integration tool. #[derive(Debug, Clone)] @@ -175,7 +175,11 @@ impl BlockContent for UseToolCallBlock { } } - // Output preview + let max_inline = if ctx.mode == DisplayMode::Truncated { + TRUNCATED_INLINE_LINES + } else { + MAX_INLINE_LINES + }; if let Some(ref output) = self.output { lines.push(Line::from("").into()); lines @@ -185,8 +189,8 @@ impl BlockContent for UseToolCallBlock { let content_lines: Vec<&str> = output.lines().collect(); for (i, line) in content_lines.iter().enumerate() { - if i >= MAX_INLINE_LINES { - let remaining = content_lines.len() - MAX_INLINE_LINES; + if i >= max_inline { + let remaining = content_lines.len() - max_inline; lines.push( BlockLine::from(Line::from(Span::styled( format!( @@ -284,3 +288,58 @@ impl BlockContent for UseToolCallBlock { Some(Text::from(vec![self.header_line(&theme, false, None)])) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::scrollback::types::BlockContext; + + fn ctx(mode: DisplayMode) -> BlockContext { + BlockContext { + width: 80, + mode, + is_running: false, + raw: false, + max_lines: None, + appearance: Default::default(), + is_selected: false, + cwd: None, + } + } + + fn rendered_text(block: &UseToolCallBlock, mode: DisplayMode) -> String { + block + .output(&ctx(mode)) + .lines + .iter() + .map(|l| { + l.content + .spans + .iter() + .map(|s| s.content.as_ref()) + .collect::() + }) + .collect::>() + .join("\n") + } + + #[test] + fn truncated_caps_inline_output_tighter_than_expanded() { + let mut block = UseToolCallBlock::new("linear__list_issues"); + let content: Vec = (1..=12).map(|i| format!("l{i:02} row")).collect(); + block.output = Some(content.join("\n")); + + let truncated = rendered_text(&block, DisplayMode::Truncated); + assert!(truncated.contains("l03"), "truncated:\n{truncated}"); + assert!(!truncated.contains("l04"), "truncated:\n{truncated}"); + assert!( + truncated.contains("(9 more lines"), + "truncated:\n{truncated}" + ); + + let expanded = rendered_text(&block, DisplayMode::Expanded); + assert!(expanded.contains("l10"), "expanded:\n{expanded}"); + assert!(!expanded.contains("l11"), "expanded:\n{expanded}"); + assert!(expanded.contains("(2 more lines"), "expanded:\n{expanded}"); + } +} diff --git a/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/web_fetch.rs b/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/web_fetch.rs index e47d7e0..d90765a 100644 --- a/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/web_fetch.rs +++ b/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/web_fetch.rs @@ -11,8 +11,8 @@ use crate::scrollback::types::{ }; use crate::theme::Theme; -/// Max lines of content shown inline before truncation. const MAX_INLINE_LINES: usize = 10; +const TRUNCATED_INLINE_LINES: usize = 3; /// Web fetch tool call — fetching a URL and returning markdown content. #[derive(Debug, Clone)] @@ -238,9 +238,11 @@ impl BlockContent for WebFetchToolCallBlock { lines.push(BlockLine::separator(meta)); } - // Content preview with bg_dark background, capped at - // MAX_INLINE_LINES. Full content is available via the - // fullscreen viewer (Enter/o). + let max_inline = if ctx.mode == DisplayMode::Truncated { + TRUNCATED_INLINE_LINES + } else { + MAX_INLINE_LINES + }; if let Some(ref output) = self.output { lines.push(Line::from("").into()); @@ -252,12 +254,12 @@ impl BlockContent for WebFetchToolCallBlock { let total_lines = output.lines().count(); for (i, line) in output.lines().enumerate() { - if i >= MAX_INLINE_LINES { + if i >= max_inline { lines.push( BlockLine::from(Line::from(Span::styled( format!( "{indent}... ({} more lines, press Enter to view)", - total_lines - MAX_INLINE_LINES + total_lines - max_inline ), theme.dim(), ))) @@ -347,3 +349,58 @@ impl BlockContent for WebFetchToolCallBlock { Some(Text::from(vec![self.header_line(&theme, false, None)])) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::scrollback::types::BlockContext; + + fn ctx(mode: DisplayMode) -> BlockContext { + BlockContext { + width: 80, + mode, + is_running: false, + raw: false, + max_lines: None, + appearance: Default::default(), + is_selected: false, + cwd: None, + } + } + + fn rendered_text(block: &WebFetchToolCallBlock, mode: DisplayMode) -> String { + block + .output(&ctx(mode)) + .lines + .iter() + .map(|l| { + l.content + .spans + .iter() + .map(|s| s.content.as_ref()) + .collect::() + }) + .collect::>() + .join("\n") + } + + #[test] + fn truncated_caps_inline_content_tighter_than_expanded() { + let content: Vec = (1..=12).map(|i| format!("l{i:02} body")).collect(); + let block = + WebFetchToolCallBlock::new("https://example.com").with_output(content.join("\n")); + + let truncated = rendered_text(&block, DisplayMode::Truncated); + assert!(truncated.contains("l03"), "truncated:\n{truncated}"); + assert!(!truncated.contains("l04"), "truncated:\n{truncated}"); + assert!( + truncated.contains("(9 more lines"), + "truncated:\n{truncated}" + ); + + let expanded = rendered_text(&block, DisplayMode::Expanded); + assert!(expanded.contains("l10"), "expanded:\n{expanded}"); + assert!(!expanded.contains("l11"), "expanded:\n{expanded}"); + assert!(expanded.contains("(2 more lines"), "expanded:\n{expanded}"); + } +} diff --git a/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/web_search.rs b/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/web_search.rs index 3c7498e..4afb807 100644 --- a/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/web_search.rs +++ b/crates/codegen/xai-grok-pager/src/scrollback/blocks/tool/web_search.rs @@ -13,8 +13,8 @@ use crate::scrollback::types::{ }; use crate::theme::Theme; -/// Max lines of content shown inline before truncation. const MAX_INLINE_LINES: usize = 10; +const TRUNCATED_INLINE_LINES: usize = 3; /// Max number of domain names shown in the sources summary line. const MAX_INLINE_SOURCES: usize = 3; @@ -258,9 +258,11 @@ impl BlockContent for WebSearchToolCallBlock { }) .collect(); - // Content preview with bg_dark background, capped at - // MAX_INLINE_LINES. Full content is available via the - // fullscreen viewer (Enter/o). + let max_inline = if ctx.mode == DisplayMode::Truncated { + TRUNCATED_INLINE_LINES + } else { + MAX_INLINE_LINES + }; if let Some(ref content) = self.content { lines.push(BlockLine::separator(Line::from(""))); @@ -272,8 +274,8 @@ impl BlockContent for WebSearchToolCallBlock { let content_lines: Vec<&str> = content.lines().collect(); for (i, line) in content_lines.iter().enumerate() { - if i >= MAX_INLINE_LINES { - let remaining = content_lines.len() - MAX_INLINE_LINES; + if i >= max_inline { + let remaining = content_lines.len() - max_inline; lines.push( BlockLine::from(Line::from(Span::styled( format!( @@ -379,3 +381,58 @@ impl BlockContent for WebSearchToolCallBlock { Some(Text::from(vec![self.header_line(&theme, false, None)])) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::scrollback::types::BlockContext; + + fn ctx(mode: DisplayMode) -> BlockContext { + BlockContext { + width: 80, + mode, + is_running: false, + raw: false, + max_lines: None, + appearance: Default::default(), + is_selected: false, + cwd: None, + } + } + + fn rendered_text(block: &WebSearchToolCallBlock, mode: DisplayMode) -> String { + block + .output(&ctx(mode)) + .lines + .iter() + .map(|l| { + l.content + .spans + .iter() + .map(|s| s.content.as_ref()) + .collect::() + }) + .collect::>() + .join("\n") + } + + #[test] + fn truncated_caps_inline_content_tighter_than_expanded() { + let mut block = WebSearchToolCallBlock::new("rust async traits"); + let content: Vec = (1..=12).map(|i| format!("l{i:02} result")).collect(); + block.content = Some(content.join("\n")); + + let truncated = rendered_text(&block, DisplayMode::Truncated); + assert!(truncated.contains("l03"), "truncated:\n{truncated}"); + assert!(!truncated.contains("l04"), "truncated:\n{truncated}"); + assert!( + truncated.contains("(9 more lines"), + "truncated:\n{truncated}" + ); + + let expanded = rendered_text(&block, DisplayMode::Expanded); + assert!(expanded.contains("l10"), "expanded:\n{expanded}"); + assert!(!expanded.contains("l11"), "expanded:\n{expanded}"); + assert!(expanded.contains("(2 more lines"), "expanded:\n{expanded}"); + } +} diff --git a/crates/codegen/xai-grok-pager/src/slash/commands/mod.rs b/crates/codegen/xai-grok-pager/src/slash/commands/mod.rs index 355bb6d..a1c67f7 100644 --- a/crates/codegen/xai-grok-pager/src/slash/commands/mod.rs +++ b/crates/codegen/xai-grok-pager/src/slash/commands/mod.rs @@ -79,7 +79,6 @@ pub fn builtin_commands() -> Vec> { Arc::new(docs::DocsCommand), Arc::new(home::HomeCommand), Arc::new(new::NewCommand), - Arc::new(fork::ForkCommand), Arc::new(compact::CompactCommand), Arc::new(copy::CopyCommand), @@ -117,7 +116,6 @@ pub fn builtin_commands() -> Vec> { Arc::new(mcps::McpsCommand), Arc::new(btw::BtwCommand), Arc::new(recap::RecapCommand), - Arc::new(terminal_setup::TerminalSetupCommand), Arc::new(voice::VoiceCommand), Arc::new(loop_cmd::LoopCommand), diff --git a/crates/codegen/xai-grok-pager/src/views/dashboard/render.rs b/crates/codegen/xai-grok-pager/src/views/dashboard/render.rs index ab9e9db..77cb2af 100644 --- a/crates/codegen/xai-grok-pager/src/views/dashboard/render.rs +++ b/crates/codegen/xai-grok-pager/src/views/dashboard/render.rs @@ -6040,9 +6040,10 @@ mod tests { ); assert!( lines.iter().any(|l| matches!( - l, - DashboardLine::Header { state, count } if *state == RowState::Working && *count == 2 - )), + l, + DashboardLine::Header { state, count } + if *state == RowState::Working && *count == 2 + )), "collapsed Working header must still render with its true count", ); let working_rows = lines @@ -6080,7 +6081,8 @@ mod tests { assert!( lines .iter() - .any(|l| matches!(l, DashboardLine::PinnedHeader { count } if *count == 1)), + .any(|l| matches!(l, DashboardLine::PinnedHeader { count } +if *count == 1)), "collapsed Pinned header must still render", ); // The pinned row is hidden; the (non-pinned) Working row remains. @@ -6142,9 +6144,10 @@ mod tests { // Header still shows the TRUE total, not the visible count. assert!( lines.iter().any(|l| matches!( - l, - DashboardLine::Header { state, count } if *state == RowState::Idle && *count == total as usize - )), + l, + DashboardLine::Header { state, count } + if *state == RowState::Idle && *count == total as usize + )), "Idle header keeps the true total count", ); } diff --git a/crates/codegen/xai-grok-pager/src/views/memory_modal.rs b/crates/codegen/xai-grok-pager/src/views/memory_modal.rs index 98c4b4a..9bc2229 100644 --- a/crates/codegen/xai-grok-pager/src/views/memory_modal.rs +++ b/crates/codegen/xai-grok-pager/src/views/memory_modal.rs @@ -595,7 +595,8 @@ fn render_file_list(buf: &mut Buffer, area: Rect, state: &mut MemoryModalState, ); if is_selected - && matches!(state.mode, MemoryModalMode::ConfirmingDelete { idx } if idx == filt_idx) + && matches!(state.mode, MemoryModalMode::ConfirmingDelete { idx } +if idx == filt_idx) { let hint = " [x to confirm]"; let hint_w = hint.len() as u16; diff --git a/crates/codegen/xai-grok-pager/src/views/settings_modal/state.rs b/crates/codegen/xai-grok-pager/src/views/settings_modal/state.rs index 4a0e302..0f7356a 100644 --- a/crates/codegen/xai-grok-pager/src/views/settings_modal/state.rs +++ b/crates/codegen/xai-grok-pager/src/views/settings_modal/state.rs @@ -275,21 +275,20 @@ impl SettingsModalState { self.invalidate_filter(); if let Some(key) = subpane_key { - let still_visible = self - .rows - .iter() - .any(|r| matches!(r, RowEntry::Setting { key: k, .. } if *k == key)); + let still_visible = self.rows.iter().any(|r| { + matches!(r, RowEntry::Setting { key: k, .. } +if *k == key) + }); if !still_visible { self.transition_to_browse(); } } if let Some(key) = prev_key { - if let Some(idx) = self - .rows - .iter() - .position(|r| matches!(r, RowEntry::Setting { key: k, .. } if *k == key)) - { + if let Some(idx) = self.rows.iter().position(|r| { + matches!(r, RowEntry::Setting { key: k, .. } +if *k == key) + }) { self.selected = idx; } else { self.selected = self diff --git a/crates/codegen/xai-grok-pager/src/views/settings_modal/tests.rs b/crates/codegen/xai-grok-pager/src/views/settings_modal/tests.rs index 9d1cb3f..40541c5 100644 --- a/crates/codegen/xai-grok-pager/src/views/settings_modal/tests.rs +++ b/crates/codegen/xai-grok-pager/src/views/settings_modal/tests.rs @@ -36,13 +36,17 @@ fn contextual_hints_group_sub_sheet_flow() { let group_idx = s .rows .iter() - .position(|r| matches!(r, RowEntry::Setting { key, .. } if *key == "contextual_hints")) + .position(|r| { + matches!(r, RowEntry::Setting { key, .. } +if *key == "contextual_hints") + }) .expect("group row present"); assert!( !s.rows.iter().any(|r| matches!( - r, - RowEntry::Setting { key, .. } if key.starts_with("contextual_hints.") - )), + r, + RowEntry::Setting { key, .. } + if key.starts_with("contextual_hints.") + )), "child rows must be hidden from the top-level list", ); @@ -4118,7 +4122,10 @@ fn advance_next_recovers_when_selection_is_hidden() { let compact_idx = s .rows .iter() - .position(|r| matches!(r, RowEntry::Setting { key, .. } if *key == "compact_mode")) + .position(|r| { + matches!(r, RowEntry::Setting { key, .. } +if *key == "compact_mode") + }) .unwrap(); s.selected = compact_idx; // Advance: lands on the first visible setting (show_timestamps). @@ -4127,7 +4134,10 @@ fn advance_next_recovers_when_selection_is_hidden() { let show_ts_idx = s .rows .iter() - .position(|r| matches!(r, RowEntry::Setting { key, .. } if *key == "show_timestamps")) + .position(|r| { + matches!(r, RowEntry::Setting { key, .. } +if *key == "show_timestamps") + }) .unwrap(); assert_eq!(s.selected, show_ts_idx); } @@ -4150,7 +4160,10 @@ fn advance_prev_recovers_when_selection_is_hidden() { let compact_idx = s .rows .iter() - .position(|r| matches!(r, RowEntry::Setting { key, .. } if *key == "compact_mode")) + .position(|r| { + matches!(r, RowEntry::Setting { key, .. } +if *key == "compact_mode") + }) .unwrap(); s.selected = compact_idx; let moved = s.advance_prev(); @@ -4158,7 +4171,10 @@ fn advance_prev_recovers_when_selection_is_hidden() { let simple_idx = s .rows .iter() - .position(|r| matches!(r, RowEntry::Setting { key, .. } if *key == "simple_mode")) + .position(|r| { + matches!(r, RowEntry::Setting { key, .. } +if *key == "simple_mode") + }) .unwrap(); assert_eq!(s.selected, simple_idx); } @@ -4258,10 +4274,10 @@ fn section_headers_have_blank_line_above_except_first() { for cat in SettingCategory::ALL { // Skip categories the default registry doesn't populate // (e.g. Session — no settings registered). - let has_setting = s - .rows - .iter() - .any(|r| matches!(r, RowEntry::Header { category } if category == cat)); + let has_setting = s.rows.iter().any(|r| { + matches!(r, RowEntry::Header { category } +if category == cat) + }); if !has_setting { continue; } @@ -4607,7 +4623,10 @@ fn two_line_row_hit_rect_spans_both_lines() { let row_idx = s .rows .iter() - .position(|r| matches!(r, RowEntry::Setting { key, .. } if *key == "coding_data_sharing")) + .position(|r| { + matches!(r, RowEntry::Setting { key, .. } +if *key == "coding_data_sharing") + }) .expect("coding_data_sharing must be registered"); // Render at a narrow width so coding_data_sharing forces a // two-line layout. @@ -4672,7 +4691,10 @@ fn two_line_row_with_expansion_renders_three_segments() { let row_idx = s .rows .iter() - .position(|r| matches!(r, RowEntry::Setting { key, .. } if *key == "coding_data_sharing")) + .position(|r| { + matches!(r, RowEntry::Setting { key, .. } +if *key == "coding_data_sharing") + }) .expect("coding_data_sharing must be registered"); s.selected = row_idx; s.expanded_keys.insert("coding_data_sharing"); @@ -4730,7 +4752,10 @@ fn group_row_renders_expanded_description() { let row_idx = s .rows .iter() - .position(|r| matches!(r, RowEntry::Setting { key, .. } if *key == "contextual_hints")) + .position(|r| { + matches!(r, RowEntry::Setting { key, .. } +if *key == "contextual_hints") + }) .expect("contextual_hints group must be registered"); s.selected = row_idx; s.expanded_keys.insert("contextual_hints"); @@ -5769,7 +5794,10 @@ fn enter_picker_for(key: &'static str) -> SettingsModalState { let row_idx = s .rows .iter() - .position(|r| matches!(r, RowEntry::Setting { key: k, .. } if *k == key)) + .position(|r| { + matches!(r, RowEntry::Setting { key: k, .. } +if *k == key) + }) .unwrap_or_else(|| panic!("no row for key `{key}` in default registry")); assert!(s.select_at(row_idx), "select_at({row_idx})"); assert!( diff --git a/crates/codegen/xai-grok-pager/src/views/shortcuts_help.rs b/crates/codegen/xai-grok-pager/src/views/shortcuts_help.rs index 1332193..9b41b4a 100644 --- a/crates/codegen/xai-grok-pager/src/views/shortcuts_help.rs +++ b/crates/codegen/xai-grok-pager/src/views/shortcuts_help.rs @@ -1681,9 +1681,10 @@ mod tests { let entries = build_entries(&all_contexts(), ®istry, true); let has_row = entries.iter().any(|e| { matches!( - e, - ShortcutsHelpEntry::Hint { item, .. } if item.label == "mouse reporting" - ) + e, + ShortcutsHelpEntry::Hint { item, .. } + if item.label == "mouse reporting" + ) }); assert!( !has_row, @@ -1776,21 +1777,24 @@ mod tests { let has_todos = entries.iter().any(|e| { matches!( - e, - ShortcutsHelpEntry::Hint { item, .. } if item.label == "todos" - ) + e, + ShortcutsHelpEntry::Hint { item, .. } + if item.label == "todos" + ) }); let has_sessions = entries.iter().any(|e| { matches!( - e, - ShortcutsHelpEntry::Hint { item, .. } if item.label == "sessions" - ) + e, + ShortcutsHelpEntry::Hint { item, .. } + if item.label == "sessions" + ) }); let has_queue = entries.iter().any(|e| { matches!( - e, - ShortcutsHelpEntry::Hint { item, .. } if item.label == "queue" - ) + e, + ShortcutsHelpEntry::Hint { item, .. } + if item.label == "queue" + ) }); assert!(has_todos, "should include toggle todos"); assert!(has_sessions, "should include open sessions"); @@ -1835,13 +1839,14 @@ mod tests { .iter() .find(|e| { matches!( - e, - ShortcutsHelpEntry::Hint { - item, - action_id: None, - .. - } if item.label == "paste" - ) + e, + ShortcutsHelpEntry::Hint { + item, + action_id: None, + .. + } + if item.label == "paste" + ) }) .expect("cheatsheet should list paste"); let ShortcutsHelpEntry::Hint { @@ -1922,9 +1927,10 @@ mod tests { let nav_dimmed = entries.iter().any(|e| { matches!( - e, - ShortcutsHelpEntry::Hint { item, dimmed: true, .. } if item.label == "nav" - ) + e, + ShortcutsHelpEntry::Hint { item, dimmed: true, .. } + if item.label == "nav" + ) }); assert!( nav_dimmed, @@ -1933,17 +1939,19 @@ mod tests { let quit_bright = entries.iter().any(|e| { matches!( - e, - ShortcutsHelpEntry::Hint { item, dimmed: false, .. } if item.label == "quit" - ) + e, + ShortcutsHelpEntry::Hint { item, dimmed: false, .. } + if item.label == "quit" + ) }); assert!(quit_bright, "quit should not be dimmed (When::Always)"); let cancel_bright = entries.iter().any(|e| { matches!( - e, - ShortcutsHelpEntry::Hint { item, dimmed: false, .. } if item.label == "cancel" - ) + e, + ShortcutsHelpEntry::Hint { item, dimmed: false, .. } + if item.label == "cancel" + ) }); assert!( cancel_bright, @@ -1959,9 +1967,10 @@ mod tests { let send_dimmed = entries.iter().any(|e| { matches!( - e, - ShortcutsHelpEntry::Hint { item, dimmed: true, .. } if item.label == "send" - ) + e, + ShortcutsHelpEntry::Hint { item, dimmed: true, .. } + if item.label == "send" + ) }); assert!( send_dimmed, @@ -1970,9 +1979,10 @@ mod tests { let nav_dimmed = entries.iter().any(|e| { matches!( - e, - ShortcutsHelpEntry::Hint { item, dimmed: true, .. } if item.label == "nav" - ) + e, + ShortcutsHelpEntry::Hint { item, dimmed: true, .. } + if item.label == "nav" + ) }); assert!( nav_dimmed, @@ -2638,14 +2648,15 @@ mod tests { .iter() .position(|e| { matches!( - e, - ShortcutsHelpEntry::Hint { - item, - action_id: None, - long_help: Some(_), - .. - } if item.label == "paste" - ) + e, + ShortcutsHelpEntry::Hint { + item, + action_id: None, + long_help: Some(_), + .. + } + if item.label == "paste" + ) }) .expect("paste pseudo-row with long_help"); assert_eq!( @@ -3024,9 +3035,10 @@ mod tests { for label in ["top", "btm", "copy", "copy cmd"] { let present = entries.iter().any(|e| { matches!( - e, - ShortcutsHelpEntry::Hint { item, .. } if item.label == label - ) + e, + ShortcutsHelpEntry::Hint { item, .. } + if item.label == label + ) }); assert!( !present, @@ -3315,14 +3327,15 @@ mod tests { .iter() .position(|e| { matches!( - e, - ShortcutsHelpEntry::Hint { - item, - action_id: None, - long_help: Some(_), - .. - } if item.label == "paste" - ) + e, + ShortcutsHelpEntry::Hint { + item, + action_id: None, + long_help: Some(_), + .. + } + if item.label == "paste" + ) }) .expect("paste pseudo-row with long_help"); let key_id = ExpandKey::Pseudo("paste"); diff --git a/crates/codegen/xai-grok-pager/src/views/tasks_pane.rs b/crates/codegen/xai-grok-pager/src/views/tasks_pane.rs index 875addc..5a5bcf7 100644 --- a/crates/codegen/xai-grok-pager/src/views/tasks_pane.rs +++ b/crates/codegen/xai-grok-pager/src/views/tasks_pane.rs @@ -992,11 +992,10 @@ impl TasksPane { }; if changed { self.rebuild_entries(); - if let Some(header) = self - .entries - .iter() - .find(|e| matches!(e, TaskEntry::Header { group: g, .. } if *g == group)) - { + if let Some(header) = self.entries.iter().find(|e| { + matches!(e, TaskEntry::Header { group: g, .. } +if *g == group) + }) { let id = header.stable_id(); self.list_state.select_by_id(id); } diff --git a/crates/codegen/xai-grok-pager/src/views/welcome/mod.rs b/crates/codegen/xai-grok-pager/src/views/welcome/mod.rs index 8e9577e..f3bbf23 100644 --- a/crates/codegen/xai-grok-pager/src/views/welcome/mod.rs +++ b/crates/codegen/xai-grok-pager/src/views/welcome/mod.rs @@ -2949,10 +2949,12 @@ mod tests { // Verify headers assert!( - matches!(&result[0], crate::views::picker::PickerEntry::Header { label } if label == &"fw-1") + matches!(&result[0], crate::views::picker::PickerEntry::Header { label } +if label == &"fw-1") ); assert!( - matches!(&result[2], crate::views::picker::PickerEntry::Header { label } if label == &"xai") + matches!(&result[2], crate::views::picker::PickerEntry::Header { label } +if label == &"xai") ); } @@ -2981,11 +2983,13 @@ mod tests { Some("zzz"), ); assert!( - matches!(&result[0], crate::views::picker::PickerEntry::Header { label } if label == &"zzz"), + matches!(&result[0], crate::views::picker::PickerEntry::Header { label } +if label == &"zzz"), "current repo group pinned first" ); assert!( - matches!(&result[2], crate::views::picker::PickerEntry::Header { label } if label == &"aaa"), + matches!(&result[2], crate::views::picker::PickerEntry::Header { label } +if label == &"aaa"), "remaining group follows alphabetically" ); } diff --git a/crates/codegen/xai-grok-pager/tests/pty_e2e/bash_full_output_double_click_fold_pty.rs b/crates/codegen/xai-grok-pager/tests/pty_e2e/bash_full_output_double_click_fold_pty.rs index d419093..efc416e 100644 --- a/crates/codegen/xai-grok-pager/tests/pty_e2e/bash_full_output_double_click_fold_pty.rs +++ b/crates/codegen/xai-grok-pager/tests/pty_e2e/bash_full_output_double_click_fold_pty.rs @@ -42,15 +42,17 @@ async fn bash_full_output_double_click_fold_pty() { harness .wait_for_text("L12", Duration::from_secs(30)) .expect("bash output tail"); + // Live tail can show L06–L12 while L01 is still clipped; wait for + // expand-on-finish before asserting the head is present. harness - .wait_for_text("L06", Duration::from_secs(10)) + .wait_for_text("L01", Duration::from_secs(15)) .unwrap_or_else(|_| { panic!( - "finished ! command must show its full output (middle lines); got:\n{}", + "finished ! command must not truncate output (L01 missing)\nscreen:\n{}", harness.screen_contents() ) }); - for line in ["L01", "L03", "L09"] { + for line in ["L03", "L06", "L09"] { assert!( harness.contains_text(line), "finished ! command must not truncate output ({line} missing)\nscreen:\n{}", diff --git a/crates/codegen/xai-grok-pager/tests/pty_e2e/minimal/minimal_transcript_expands_collapsed_thinking.rs b/crates/codegen/xai-grok-pager/tests/pty_e2e/minimal/minimal_commits_thinking_body_to_scrollback.rs similarity index 60% rename from crates/codegen/xai-grok-pager/tests/pty_e2e/minimal/minimal_transcript_expands_collapsed_thinking.rs rename to crates/codegen/xai-grok-pager/tests/pty_e2e/minimal/minimal_commits_thinking_body_to_scrollback.rs index c105c9a..58c273b 100644 --- a/crates/codegen/xai-grok-pager/tests/pty_e2e/minimal/minimal_transcript_expands_collapsed_thinking.rs +++ b/crates/codegen/xai-grok-pager/tests/pty_e2e/minimal/minimal_commits_thinking_body_to_scrollback.rs @@ -6,20 +6,11 @@ use crate::common::*; /// so screen assertions can tell the two apart. const REASONING_SENTINEL: &str = "REASONINGSENTINEL"; -/// Dogfood bug: "I don't see thoughts in the transcript". With thinking -/// enabled (`[ui] show_thinking_blocks` — the default, set -/// explicitly here so the test doesn't depend on the rollout default), -/// minimal commits reasoning as a **collapsed** `Thought for Xs` header -/// (print-once display policy) — the body is intentionally not in the live -/// scrollback. The advertised full-fidelity `/transcript` view must therefore -/// render the thinking body **expanded**, or the reasoning is unreachable. -/// -/// Flow: stream a reasoning+text turn → the answer commits, the reasoning -/// collapses to its header (body nowhere on screen) → `/transcript` with -/// `PAGER=cat` dumps the full view → the reasoning body appears. +/// `[ui] show_thinking_blocks` is set explicitly (not left to the default) +/// so the test doesn't depend on the rollout default. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[ignore] -async fn minimal_transcript_expands_collapsed_thinking() { +async fn minimal_commits_thinking_body_to_scrollback() { // The model must run on the Responses backend — reasoning summary deltas // are a Responses-API stream shape (the scripted events below). let content = ContentController::start_with_models(vec![ @@ -54,10 +45,7 @@ async fn minimal_transcript_expands_collapsed_thinking() { ) .expect("write config"); - // Minimal env + PAGER=cat (non-interactive dump, same as - // `minimal_transcript_opens_in_pager`). - let mut env = content.env_for_pager(); - env.push(("PAGER".to_string(), "cat".to_string())); + let env = content.env_for_pager(); let env_refs: Vec<(&str, &str)> = env.iter().map(|(k, v)| (k.as_str(), v.as_str())).collect(); let binary = pager_binary().expect("resolve pager binary"); let mut harness = PtyHarness::new(&binary, DEFAULT_ROWS, DEFAULT_COLS, MINIMAL_ARGS, &env_refs) @@ -73,34 +61,17 @@ async fn minimal_transcript_expands_collapsed_thinking() { .wait_for_full_text(MOCK_RESPONSE_SENTINEL, Duration::from_secs(30)) .expect("turn committed"); - // The reasoning committed as its collapsed header: the body is NOT in the - // live view (that's the print-once display policy, not a bug)… harness .wait_for_full_text("Thought for", Duration::from_secs(10)) - .expect("collapsed thinking header committed"); - assert!( - !harness.full_text().contains(REASONING_SENTINEL), - "reasoning body must be collapsed in the live view\nfull:\n{}", - harness.full_text() - ); - - // …so the transcript is the only way to read it. cat dumps the full view. - inject_keys_paced(&mut harness, b"/transcript"); - harness.inject_keys(b"\r").expect("submit /transcript"); - + .expect("thinking header committed"); harness - .wait_for_full_text(REASONING_SENTINEL, Duration::from_secs(15)) + .wait_for_full_text(REASONING_SENTINEL, Duration::from_secs(10)) .unwrap_or_else(|e| { panic!( - "transcript must expand the collapsed thinking body: {e}\nfull:\n{}", + "reasoning body must be committed to scrollback: {e}\nfull:\n{}", harness.full_text() ) }); - - // And the inline TUI survives the suspend/restore round trip. - harness - .wait_for_text(MINIMAL_IDLE_SENTINEL, Duration::from_secs(10)) - .expect("inline TUI restored after the pager exited"); assert!( !harness.contains_text("panicked"), "pager panicked\nscreen:\n{}", diff --git a/crates/codegen/xai-grok-pager/tests/pty_e2e/minimal/minimal_lookup_commits_one_line_summary.rs b/crates/codegen/xai-grok-pager/tests/pty_e2e/minimal/minimal_lookup_commits_one_line_summary.rs new file mode 100644 index 0000000..1911275 --- /dev/null +++ b/crates/codegen/xai-grok-pager/tests/pty_e2e/minimal/minimal_lookup_commits_one_line_summary.rs @@ -0,0 +1,74 @@ +// Per-test-case module for the `pty_e2e` integration test crate. +#[allow(unused_imports)] +use crate::common::*; + +const BODY_SENTINEL: &str = "READBODYONLYSENTINEL"; +const DONE_SENTINEL: &str = "LOOKUP_TURN_DONE"; + +/// Uses `read_file` rather than `grep`: the grep tool shells out to `rg`, +/// which is absent from the Bazel remote-exec sandbox (only xai-grok-tools' +/// own test targets ship `@ripgrep_hermetic//:rg`), and a failed spawn +/// degrades to a zero-match result that vacuously passes the absence assert. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[ignore] +async fn minimal_lookup_commits_one_line_summary() { + let content = ContentController::start().await.expect("start content"); + + let fixture = content.home().join("haystack.txt"); + std::fs::write(&fixture, format!("{BODY_SENTINEL} body line\n")).expect("write fixture"); + + enqueue_tool_turn( + &content, + "call_read", + "read_file", + json!({ "target_file": fixture.to_string_lossy() }).to_string(), + ); + content.set_response(DONE_SENTINEL); + + let mut harness = spawn_minimal_in_dir( + &content, + DEFAULT_ROWS, + DEFAULT_COLS, + &["--yolo", "--trust"], + content.home(), + ); + wait_minimal_ready(&mut harness); + + harness + .inject_keys(format!("{PROMPT}\r").as_bytes()) + .expect("submit prompt"); + harness + .wait_for_full_text(DONE_SENTINEL, Duration::from_secs(60)) + .expect("tool turn settles"); + harness + .wait_for_text(MINIMAL_IDLE_SENTINEL, Duration::from_secs(20)) + .expect("return to idle"); + + harness + .wait_for_full_text("haystack.txt", Duration::from_secs(10)) + .expect("read header committed"); + assert!( + !harness.full_text().contains(BODY_SENTINEL), + "successful read must commit as a one-line header, without file \ + content\nfull:\n{}", + harness.full_text() + ); + + harness.inject_keys(b"\x05").expect("ctrl+e expand"); + harness + .wait_for_full_text(BODY_SENTINEL, Duration::from_secs(10)) + .unwrap_or_else(|e| { + panic!( + "Ctrl+E must re-print the read with its file content: {e}\nfull:\n{}", + harness.full_text() + ) + }); + + assert!( + !harness.contains_text("panicked"), + "pager panicked\nscreen:\n{}", + harness.screen_contents() + ); + + quit_minimal(&mut harness); +} diff --git a/crates/codegen/xai-grok-pager/tests/pty_e2e/minimal/mod.rs b/crates/codegen/xai-grok-pager/tests/pty_e2e/minimal/mod.rs index ee9f12d..f0c1e35 100644 --- a/crates/codegen/xai-grok-pager/tests/pty_e2e/minimal/mod.rs +++ b/crates/codegen/xai-grok-pager/tests/pty_e2e/minimal/mod.rs @@ -9,6 +9,7 @@ mod minimal_cli_screen_mode_does_not_persist; mod minimal_commits_response_to_scrollback; +mod minimal_commits_thinking_body_to_scrollback; mod minimal_committed_content_survives_overlay_grow; mod minimal_continue_reprints_transcript; mod minimal_ctrl_c_arms_and_quits; @@ -16,6 +17,7 @@ mod minimal_double_esc_committed_queued_prompt_single_render; mod minimal_esc_mid_turn_is_swallowed; mod minimal_flush_left_no_hpad; mod minimal_help_opens_command_palette; +mod minimal_lookup_commits_one_line_summary; mod minimal_new_session_keeps_history_and_resets; mod minimal_queue_indicator_shows_while_running; mod minimal_resize_preserves_committed_scrollback; @@ -25,6 +27,5 @@ mod minimal_short_response_stays_on_screen; mod minimal_slash_dropdown_dismisses_with_esc; mod minimal_slash_switches_from_fullscreen; mod minimal_slash_switches_to_fullscreen; -mod minimal_transcript_expands_collapsed_thinking; mod minimal_transcript_opens_in_pager; mod minimal_transcript_pager_restore_no_artifacts; diff --git a/crates/codegen/xai-grok-pager/tests/settings_e2e.rs b/crates/codegen/xai-grok-pager/tests/settings_e2e.rs index 7895549..938f2e0 100644 --- a/crates/codegen/xai-grok-pager/tests/settings_e2e.rs +++ b/crates/codegen/xai-grok-pager/tests/settings_e2e.rs @@ -148,7 +148,10 @@ fn row_idx_for(state: &SettingsModalState, target: &str) -> usize { state .rows .iter() - .position(|r| matches!(r, RowEntry::Setting { key, .. } if *key == target)) + .position(|r| { + matches!(r, RowEntry::Setting { key, .. } +if *key == target) + }) .unwrap_or_else(|| panic!("setting `{target}` not present in modal rows")) } @@ -2146,10 +2149,10 @@ fn d_key_emits_open_reset_confirm_for_every_setting() { // without key-release reporting, which tests run without). Skip settings // with no visible row; their reset path is covered by the dispatch // round-trip tests. - let has_row = s - .rows - .iter() - .any(|r| matches!(r, RowEntry::Setting { key, .. } if *key == meta.key)); + let has_row = s.rows.iter().any(|r| { + matches!(r, RowEntry::Setting { key, .. } +if *key == meta.key) + }); if !has_row { continue; } @@ -3031,7 +3034,8 @@ fn pr6_permission_mode_picker_enter_dispatches_set_permission_mode_commit() { let _ = handle_settings_key(&mut s, &press(KeyCode::Enter)); assert!( - matches!(s.mode(), SettingsModalMode::PickingEnum { key, .. } if key == "permission_mode"), + matches!(s.mode(), SettingsModalMode::PickingEnum { key, .. } +if key == "permission_mode"), "Enter on permission_mode row must open the picker, got {:?}", s.mode(), ); @@ -3312,7 +3316,8 @@ fn pr11_picker_commit_for_default_dispatches_set_permission_mode_default() { navigate_to(&mut s, "permission_mode"); let _ = handle_settings_key(&mut s, &press(KeyCode::Enter)); assert!( - matches!(s.mode(), SettingsModalMode::PickingEnum { key, .. } if key == "permission_mode"), + matches!(s.mode(), SettingsModalMode::PickingEnum { key, .. } +if key == "permission_mode"), "Enter on permission_mode row must open the picker, got {:?}", s.mode(), ); @@ -3357,7 +3362,8 @@ fn pr11_picker_commit_for_ask_dispatches_set_permission_mode_ask() { navigate_to(&mut s, "permission_mode"); let _ = handle_settings_key(&mut s, &press(KeyCode::Enter)); assert!( - matches!(s.mode(), SettingsModalMode::PickingEnum { key, .. } if key == "permission_mode"), + matches!(s.mode(), SettingsModalMode::PickingEnum { key, .. } +if key == "permission_mode"), "Enter on permission_mode row must open the picker, got {:?}", s.mode(), ); @@ -4221,7 +4227,8 @@ fn pr14_default_model_picker_commits_resolved_model_id() { "Enter on DynamicEnum row must transition to PickingEnum, got {outcome:?}" ); assert!( - matches!(s.mode(), SettingsModalMode::PickingEnum { key, .. } if key == "default_model"), + matches!(s.mode(), SettingsModalMode::PickingEnum { key, .. } +if key == "default_model"), "Enter must transition to PickingEnum for default_model" ); @@ -4336,7 +4343,8 @@ fn pr14_mouse_click_on_dynamic_enum_row_opens_picker() { "second click on DynamicEnum row must open picker, got {outcome:?}", ); assert!( - matches!(s.mode(), SettingsModalMode::PickingEnum { key, .. } if key == "default_model"), + matches!(s.mode(), SettingsModalMode::PickingEnum { key, .. } +if key == "default_model"), "second click on DynamicEnum row must transition to PickingEnum, got {:?}", s.mode(), ); @@ -4380,7 +4388,8 @@ fn pr8_mouse_click_on_int_row_opens_editor() { "second click on Int row must be Changed, got {outcome:?}", ); assert!( - matches!(s.mode(), SettingsModalMode::EditingValue { key, .. } if key == "max_thoughts_width"), + matches!(s.mode(), SettingsModalMode::EditingValue { key, .. } +if key == "max_thoughts_width"), "second click on Int row must transition to EditingValue, got {:?}", s.mode(), ); @@ -6880,7 +6889,8 @@ fn scroll_speed_mouse_click_opens_editor() { "second click on focused Int row must enter the editor, got {outcome:?}" ); assert!( - matches!(s.mode(), SettingsModalMode::EditingValue { key, .. } if key == "scroll_speed"), + matches!(s.mode(), SettingsModalMode::EditingValue { key, .. } +if key == "scroll_speed"), "mode must be EditingValue(scroll_speed) after Enter-equivalent click, got {:?}", s.mode(), ); @@ -7070,7 +7080,8 @@ fn scroll_lines_mouse_click_opens_editor() { "second click on focused Int row must enter the editor, got {outcome:?}" ); assert!( - matches!(s.mode(), SettingsModalMode::EditingValue { key, .. } if key == "scroll_lines"), + matches!(s.mode(), SettingsModalMode::EditingValue { key, .. } +if key == "scroll_lines"), "mode must be EditingValue(scroll_lines), got {:?}", s.mode(), ); diff --git a/crates/codegen/xai-grok-plugin-marketplace/src/config.rs b/crates/codegen/xai-grok-plugin-marketplace/src/config.rs index 35a1831..d0013f8 100644 --- a/crates/codegen/xai-grok-plugin-marketplace/src/config.rs +++ b/crates/codegen/xai-grok-plugin-marketplace/src/config.rs @@ -281,9 +281,8 @@ mod tests { let sources = load_sources(&config); assert_eq!(sources.len(), 1); assert_eq!(sources[0].name, "Local Dev"); - assert!( - matches!(&sources[0].kind, SourceKind::Local { path } if path == &PathBuf::from("/home/user/plugins")) - ); + assert!(matches!(&sources[0].kind, SourceKind::Local { path } +if path == &PathBuf::from("/home/user/plugins"))); } #[test] @@ -300,9 +299,8 @@ mod tests { let sources = load_sources(&config); assert_eq!(sources.len(), 1); assert_eq!(sources[0].name, "xAI Official"); - assert!( - matches!(&sources[0].kind, SourceKind::Git { url, branch } if url.contains("xai-org") && branch.as_deref() == Some("main")) - ); + assert!(matches!(&sources[0].kind, SourceKind::Git { url, branch } +if url.contains("xai-org") && branch.as_deref() == Some("main"))); } #[test] @@ -397,9 +395,8 @@ mod tests { extract_marketplace_entries(marketplaces, &mut seen, &mut sources); assert_eq!(sources.len(), 1); assert_eq!(sources[0].name, "my-marketplace"); - assert!( - matches!(&sources[0].kind, SourceKind::Git { url, .. } if url == "https://github.com/anthropics/claude-plugins-official.git") - ); + assert!(matches!(&sources[0].kind, SourceKind::Git { url, .. } +if url == "https://github.com/anthropics/claude-plugins-official.git")); } #[test] @@ -420,9 +417,8 @@ mod tests { let mut sources = Vec::new(); extract_marketplace_entries(marketplaces, &mut seen, &mut sources); assert_eq!(sources.len(), 1); - assert!( - matches!(&sources[0].kind, SourceKind::Git { url, .. } if url == "git@github.com:org/repo.git") - ); + assert!(matches!(&sources[0].kind, SourceKind::Git { url, .. } +if url == "git@github.com:org/repo.git")); } #[test] diff --git a/crates/codegen/xai-grok-sampling-types/src/conversation.rs b/crates/codegen/xai-grok-sampling-types/src/conversation.rs index ad485d4..072c094 100644 --- a/crates/codegen/xai-grok-sampling-types/src/conversation.rs +++ b/crates/codegen/xai-grok-sampling-types/src/conversation.rs @@ -3742,18 +3742,20 @@ mod tests { }; assert_eq!(u.content.len(), 2); assert_matches!( - &u.content[1], - ContentPart::Image { url } if url.as_ref() == "https://example.com/image.png" - ); + &u.content[1], + ContentPart::Image { url } + if url.as_ref() == "https://example.com/image.png" + ); // Convert to chat request and verify let chat_msg = conversation_item_to_chat_message(user); let blocks = chat_msg.content.blocks(); assert_eq!(blocks.len(), 2); assert_matches!( - &blocks[1], - ChatContentBlock::ImageUrl { image_url } if image_url.url == "https://example.com/image.png" - ); + &blocks[1], + ChatContentBlock::ImageUrl { image_url } + if image_url.url == "https://example.com/image.png" + ); } #[test] @@ -4936,7 +4938,8 @@ mod tests { let chat_msg = conversation_item_to_chat_message(user); let blocks = chat_msg.content.blocks(); assert_eq!(blocks.len(), 4); - assert_matches!(&blocks[0], ChatContentBlock::Text { text } if text == "Compare these images:"); + assert_matches!(&blocks[0], ChatContentBlock::Text { text } +if text == "Compare these images:"); assert_matches!(&blocks[1], ChatContentBlock::ImageUrl { .. }); assert_matches!(&blocks[2], ChatContentBlock::ImageUrl { .. }); assert_matches!(&blocks[3], ChatContentBlock::ImageUrl { .. }); @@ -7647,11 +7650,11 @@ mod tests { ); }; assert_eq!(blocks.len(), 2); + assert!(matches!(&blocks[0], ChatContentBlock::Text { text } +if text == "Read image file: photo.png")); assert!( - matches!(&blocks[0], ChatContentBlock::Text { text } if text == "Read image file: photo.png") - ); - assert!( - matches!(&blocks[1], ChatContentBlock::ImageUrl { image_url } if image_url.url == "data:image/png;base64,iVBOR") + matches!(&blocks[1], ChatContentBlock::ImageUrl { image_url } +if image_url.url == "data:image/png;base64,iVBOR") ); } @@ -7717,10 +7720,12 @@ mod tests { }; assert_eq!(inner.len(), 2); assert!( - matches!(&inner[0], crate::messages::ContentBlock::Text { text, .. } if text == "Read image file: photo.png") + matches!(&inner[0], crate::messages::ContentBlock::Text { text, .. } +if text == "Read image file: photo.png") ); assert!( - matches!(&inner[1], crate::messages::ContentBlock::Image { source: crate::messages::ImageSource::Base64 { media_type, data } } if media_type == "image/png" && data == "iVBOR") + matches!(&inner[1], crate::messages::ContentBlock::Image { source: crate::messages::ImageSource::Base64 { media_type, data } } +if media_type == "image/png" && data == "iVBOR") ); } @@ -7770,7 +7775,8 @@ mod tests { if let ConversationItem::ToolResult(t) = &back { assert_eq!(t.images.len(), 1); - assert!(matches!(&t.images[0], ContentPart::Image { url } if url.contains("iVBOR"))); + assert!(matches!(&t.images[0], ContentPart::Image { url } +if url.contains("iVBOR"))); } else { panic!("Expected ToolResult"); } diff --git a/crates/codegen/xai-grok-sampling-types/src/error.rs b/crates/codegen/xai-grok-sampling-types/src/error.rs index 1407754..792f512 100644 --- a/crates/codegen/xai-grok-sampling-types/src/error.rs +++ b/crates/codegen/xai-grok-sampling-types/src/error.rs @@ -214,13 +214,14 @@ impl SamplingError { /// a new session. pub fn is_encrypted_content_error(&self) -> bool { matches!( - self, - SamplingError::Api { - status: StatusCode::BAD_REQUEST, - message, - .. - } if message.contains("encrypted_content") - ) + self, + SamplingError::Api { + status: StatusCode::BAD_REQUEST, + message, + .. + } + if message.contains("encrypted_content") + ) } /// The API rejected the request because an inline image could not be @@ -228,13 +229,14 @@ impl SamplingError { /// Exact-case match — consistent with `is_encrypted_content_error`. pub fn is_image_processing_error(&self) -> bool { matches!( - self, - SamplingError::Api { - status, - message, - .. - } if matches!(status.as_u16(), 400 | 500) && message.contains("Could not process image") - ) + self, + SamplingError::Api { + status, + message, + .. + } + if matches!(status.as_u16(), 400 | 500) && message.contains("Could not process image") + ) } pub fn is_retryable(&self) -> bool { diff --git a/crates/codegen/xai-grok-shell-base/src/cpu_profile.rs b/crates/codegen/xai-grok-shell-base/src/cpu_profile.rs index 60eda74..f9727c6 100644 --- a/crates/codegen/xai-grok-shell-base/src/cpu_profile.rs +++ b/crates/codegen/xai-grok-shell-base/src/cpu_profile.rs @@ -1003,12 +1003,13 @@ mod tests { let stop_handle = manager.take_stop_handle().unwrap(); assert!(matches!( - manager.status(), - CpuProfileStatus::Stopping { - svg_path: status_path, - .. - } if status_path == svg_path - )); + manager.status(), + CpuProfileStatus::Stopping { + svg_path: status_path, + .. + } + if status_path == svg_path + )); let err = manager .start_with_engine_for_test( @@ -1136,13 +1137,14 @@ mod tests { let _stop_handle = manager.take_stop_handle().unwrap(); assert!(matches!( - manager.status(), - CpuProfileStatus::Stopping { - svg_path: status_path, - frequency_hz: DEFAULT_FREQUENCY_HZ, - .. - } if status_path == svg_path - )); + manager.status(), + CpuProfileStatus::Stopping { + svg_path: status_path, + frequency_hz: DEFAULT_FREQUENCY_HZ, + .. + } + if status_path == svg_path + )); } #[test] diff --git a/crates/codegen/xai-grok-shell/README.md b/crates/codegen/xai-grok-shell/README.md index dc17c5b..493e5ac 100644 --- a/crates/codegen/xai-grok-shell/README.md +++ b/crates/codegen/xai-grok-shell/README.md @@ -334,6 +334,44 @@ Common log messages: | `auth: external auth provider timed out (likely needs interactive auth), killing` | Binary didn't exit before the timeout (60s initial, 5s mid-session refresh) and was killed | | `auth: failed to start external auth provider` | The command couldn't be spawned (e.g. binary not found) | +### Per-Model Auth Providers + +`auth_provider_command` above replaces Grok's *session* auth: it mints the token sent to xAI's backend. If you instead want xAI models on normal xAI login while **other models** route through a gateway (LiteLLM, corporate proxy) whose bearer tokens rotate, use a named auth provider — the rotating-token analogue of a per-model `api_key`/`env_key`. + +```toml +# ~/.grok/config.toml +[auth_provider.litellm] +command = "/usr/local/bin/litellm-token" # run via `sh -c` +token_ttl_secs = 3600 # optional: see below +timeout_secs = 10 # optional: command timeout (default 30) + +[model.proxied-claude] +model = "claude-sonnet-4-5" +base_url = "https://litellm.corp.example/v1" +context_window = 200000 +auth_provider = "litellm" +``` + +**Contract** (same stdout contract as `auth_provider_command`; the `issuer` field is accepted but unused here, and `refresh_token`, when present, is handed back to the command on refresh): + +- Without `args`, the command runs via POSIX `sh -c`, so it can be a binary path, a script, or a pipeline. With `args = ["..."]`, the command runs directly with those arguments and no shell: `command` is a program name resolved via `PATH`, or a path. Use `args` to avoid shell quoting, and on Windows, where there is no `sh`. +- stdout: a bare token, or JSON `{"access_token": "...", "expires_in": 3600}`. +- stderr: logged when the command fails; exit 0 = success. +- `GROK_AUTH_EXPIRED=1` is set whenever Grok re-mints over a token still cached in memory, whether from near-expiry rotation or a rejection. The first mint on a cold cache runs without it. + +**Token lifecycle:** + +- Tokens are cached in memory per provider and shared by every model referencing the provider; nothing is written to disk. The command is a credential helper: it owns durable storage and OAuth2 refresh (keychain, its own dotdir, etc.), exactly like `gcloud auth print-access-token` or a git credential helper. On an in-session re-mint the last credential is handed back via `GROK_AUTH_PROVIDER_ACCESS_TOKEN` (and, when present, `GROK_AUTH_PROVIDER_REFRESH_TOKEN` / `GROK_AUTH_PROVIDER_EXPIRES_AT`), so a refresh-grant command can refresh instead of re-authenticating. The command must be non-interactive and fast; do any interactive login out of band, and Grok re-runs the command on restart to re-mint. +- Grok runs the command before a chat turn when the token is missing or within about a minute of expiring, and once more after the server rejects a token. A token rejected within 30 seconds of being fetched is not refetched again, so a broken helper surfaces one clear error instead of looping. +- Token lifetime comes from `expires_in` in the command's JSON output, else `token_ttl_secs`, else the token's own JWT expiry claim. With none of these, tokens are only replaced after the server rejects one. +- Commands run with a `timeout_secs` bound (default 30, clamped to 1..=600) and are killed on timeout. A turn waits on the run, so keep helpers fast and non-interactive. +- Active sessions pick up edits or removal of a provider table at the next model switch or new session. Once picked up, an edit invalidates the cached token, so the edited command runs at the next use; removal drops the cached token. +- Helper models (web search, session summary, image description) read the shared cache and never run the command; point them at providers your chat model keeps warm. Subagents refresh tokens the same way their parent session does. + +**Interaction with other credentials:** a literal `api_key`/`env_key` on the model wins over its `auth_provider`. Provider-backed models are BYOK: your xAI session token is never sent to their endpoints, and a failing provider command fails the request rather than falling back to the session token. + +**Security:** provider commands execute code, so they are honored only from trusted config layers (`~/.grok/config.toml`, managed config, requirements). A project's `.grok/config.toml` can never define one. Whatever layer sets a model's `base_url` decides where that model's minted token is sent, and `base_url` (unlike the provider table) is not stripped from remote or campaign patches, the same as for a static `env_key`. Keep provider tables and the model `base_url` in layers you trust. The command inherits Grok's environment (so it sees `PATH`, `HOME`, and any other secrets there), but Grok's own first-party credentials (`XAI_API_KEY`, `GROK_DEPLOYMENT_KEY`, and related keys) are removed so a BYOK helper never receives them; write helpers that read only what they need, and prefer the `GROK_AUTH_PROVIDER_*` handback for the prior credential. + ### Using auth.json for API Access If you've authenticated with `grok login`, you can use the stored credentials to call the CLI chat proxy directly via curl. The proxy requires specific headers that mirror what the grok CLI sends internally: @@ -1707,13 +1745,14 @@ name = "Display Name" # Shown in model picker description = "Model description" # Optional description api_key = "sk-..." # API key for this provider (optional) env_key = "OPENAI_API_KEY" # Env var(s) holding the API key (string or array; first set wins) +auth_provider = "corp-gateway" # Named credential helper for rotating tokens (optional) temperature = 0.7 # Sampling temperature (0.0-2.0) top_p = 0.95 # Nucleus sampling parameter max_completion_tokens = 8192 # Max tokens per response context_window = 256000 # Total context window in tokens (for auto-compact) ``` -**Credential resolution order:** `api_key` → `env_key` → `XAI_API_KEY`. If neither `api_key` nor `env_key` is set, Grok falls back to the global `XAI_API_KEY` environment variable. +**Credential resolution order:** `api_key` → `env_key` → cached `auth_provider` token (terminal: a cache miss resolves to no credential, never the session token) → session token → `XAI_API_KEY`. See [Per-Model Auth Providers](#per-model-auth-providers). The `context_window` parameter is used to calculate when auto-compact should trigger. If not specified, Grok falls back to built-in defaults for known models. diff --git a/crates/codegen/xai-grok-shell/src/agent/config.rs b/crates/codegen/xai-grok-shell/src/agent/config.rs index d076f39..7a9e9cd 100644 --- a/crates/codegen/xai-grok-shell/src/agent/config.rs +++ b/crates/codegen/xai-grok-shell/src/agent/config.rs @@ -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 { 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, - /// 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, + pub config_warnings: Vec, pub grok_com_config: GrokComConfig, + /// `[auth_provider.]` tables, populated by + /// [`parse_auth_providers`] from trusted config layers only. + #[serde(skip)] + pub auth_providers: IndexMap, #[serde(default, skip_serializing_if = "Option::is_none")] pub shortcuts: Option, /// 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.]` 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, + Vec, +) { + 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.] 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 = @@ -3582,6 +3742,10 @@ pub struct ConfigModelOverride { pub api_key: Option, /// Env var name(s) for the provider key — string or array in config.toml. pub env_key: Option, + /// Name of a `[auth_provider.]` credential helper that mints + /// this model's bearer token. Static `api_key` / `env_key` win when both + /// are set. + pub auth_provider: Option, pub api_base_url: Option, pub max_completion_tokens: Option, pub temperature: Option, @@ -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, pub env_key: Option, + /// Named credential helper (`[model.] auth_provider = ""`), + /// resolved against `[auth_provider.]` 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, /// When set, `base_url` is used for session auth, `api_base_url` for API-key auth. pub api_base_url: Option, } @@ -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 { 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) { 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, 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 { 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 { + 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, } } diff --git a/crates/codegen/xai-grok-shell/src/agent/config_model_override_parse.rs b/crates/codegen/xai-grok-shell/src/agent/config_model_override_parse.rs index a3d64b3..596601f 100644 --- a/crates/codegen/xai-grok-shell/src/agent/config_model_override_parse.rs +++ b/crates/codegen/xai-grok-shell/src/agent/config_model_override_parse.rs @@ -1,5 +1,8 @@ //! Resilient parsing for `[model.]` 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.]` entry; `field` names a key when the warning is + /// field-specific. + Model { + key: String, + #[serde(skip_serializing_if = "Option::is_none")] + field: Option, + }, + /// The `[auth_provider]` section as a whole. + AuthProviderSection, + /// An `[auth_provider.]` table; `field` names a key when the + /// warning is field-specific. + AuthProvider { + name: String, + #[serde(skip_serializing_if = "Option::is_none")] + field: Option, + }, +} + +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, - /// `None` for warnings about the entry as a whole. - #[serde(skip_serializing_if = "Option::is_none")] - pub field: Option, - 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, - pub warnings: Vec, + pub warnings: Vec, } /// Parses every `[model.]` 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.] 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, -) -> (ConfigModelOverride, Vec) { +) -> (ConfigModelOverride, Vec) { 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, - warnings: &mut Vec, + warnings: &mut Vec, ) { 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) -> Vec { +fn unknown_field_warnings(model_key: &str, unknown: Vec) -> Vec { 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) -> Vec, - warnings: &mut Vec, + warnings: &mut Vec, ) { 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, - Vec, - ) { + fn parse_raw(toml_str: &str) -> (IndexMap, Vec) { 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::>() }; 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, - ) -> ( - IndexMap, - Vec, - ) { + ) -> (IndexMap, Vec) { 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)); } } } diff --git a/crates/codegen/xai-grok-shell/src/agent/models.rs b/crates/codegen/xai-grok-shell/src/agent/models.rs index a1c9b68..01e262a 100644 --- a/crates/codegen/xai-grok-shell/src/agent/models.rs +++ b/crates/codegen/xai-grok-shell/src/agent/models.rs @@ -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); diff --git a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/agent_ops.rs b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/agent_ops.rs index 7f506cb..15982fc 100644 --- a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/agent_ops.rs +++ b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/agent_ops.rs @@ -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 => { diff --git a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests.rs b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests.rs index 91324ed..fdc1900 100644 --- a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests.rs +++ b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests.rs @@ -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(); diff --git a/crates/codegen/xai-grok-shell/src/agent/relay.rs b/crates/codegen/xai-grok-shell/src/agent/relay.rs index af6fa0b..0001bf8 100644 --- a/crates/codegen/xai-grok-shell/src/agent/relay.rs +++ b/crates/codegen/xai-grok-shell/src/agent/relay.rs @@ -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(()) }; diff --git a/crates/codegen/xai-grok-shell/src/agent/subagent/tests/mod.rs b/crates/codegen/xai-grok-shell/src/agent/subagent/tests/mod.rs index f21c5a3..118ab26 100644 --- a/crates/codegen/xai-grok-shell/src/agent/subagent/tests/mod.rs +++ b/crates/codegen/xai-grok-shell/src/agent/subagent/tests/mod.rs @@ -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("")))) ); } @@ -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("")))) }), "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(""))) + ContentPart::Text { text } +if text.contains(""))) ) }); 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, } } diff --git a/crates/codegen/xai-grok-shell/src/agent/subagent/tests/rest.rs b/crates/codegen/xai-grok-shell/src/agent/subagent/tests/rest.rs index 993110a..b48f79f 100644 --- a/crates/codegen/xai-grok-shell/src/agent/subagent/tests/rest.rs +++ b/crates/codegen/xai-grok-shell/src/agent/subagent/tests/rest.rs @@ -274,7 +274,8 @@ fn compaction_preserves_inherited_prefix() { .any(|p| { matches!( p, xai_grok_sampling_types::conversation::ContentPart::Text { - text } if text.contains("") + text } +if text.contains("") ) }) } 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()); diff --git a/crates/codegen/xai-grok-shell/src/auth/auth_provider.rs b/crates/codegen/xai-grok-shell/src/auth/auth_provider.rs new file mode 100644 index 0000000..7599c76 --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/auth/auth_provider.rs @@ -0,0 +1,603 @@ +//! Model auth providers (`[auth_provider.]`). +//! +//! A model opts in with `auth_provider = ""`; 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.]` 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>, + /// 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, + /// 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, +} + +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 for AuthProviderRef { + fn from(data: AuthProviderRefData) -> Self { + AuthProviderRef::unresolved(data.name) + } +} + +impl From 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, + /// Drives the 401 fresh-mint guard. + minted_at: std::time::Instant, + expires_at: Option>, + /// 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>>; + +/// 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::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) { + 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::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(reader: R, keep: u64, buf: &mut Vec) -> 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 { + 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 { + 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 { + 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>> { + 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 { + 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 { + 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; diff --git a/crates/codegen/xai-grok-shell/src/auth/auth_provider_tests.rs b/crates/codegen/xai-grok-shell/src/auth/auth_provider_tests.rs new file mode 100644 index 0000000..cd260a1 --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/auth/auth_provider_tests.rs @@ -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, + 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::>() + .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" + ); +} diff --git a/crates/codegen/xai-grok-shell/src/auth/external_auth.rs b/crates/codegen/xai-grok-shell/src/auth/external_auth.rs index f7d1704..1077b43 100644 --- a/crates/codegen/xai-grok-shell/src/auth/external_auth.rs +++ b/crates/codegen/xai-grok-shell/src/auth/external_auth.rs @@ -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, - #[serde(default)] - pub expires_in: Option, - /// Token issuer. An xAI issuer marks the credential as first-party; - /// see [`GrokAuth::is_xai_auth`]. - #[serde(default)] - pub issuer: Option, -} - -/// 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 { - 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::(&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 Option { + 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 { 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] diff --git a/crates/codegen/xai-grok-shell/src/auth/mod.rs b/crates/codegen/xai-grok-shell/src/auth/mod.rs index 16be304..4a595b1 100644 --- a/crates/codegen/xai-grok-shell/src/auth/mod.rs +++ b/crates/codegen/xai-grok-shell/src/auth/mod.rs @@ -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, diff --git a/crates/codegen/xai-grok-shell/src/auth/token_output.rs b/crates/codegen/xai-grok-shell/src/auth/token_output.rs new file mode 100644 index 0000000..1570e0e --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/auth/token_output.rs @@ -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, + #[serde(default)] + pub expires_in: Option, + /// An xAI issuer marks the credential as first-party + /// (see [`crate::auth::GrokAuth::is_xai_auth`]). + #[serde(default)] + pub issuer: Option, +} + +/// 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> { + 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, + pub expires_at: Option>, + pub issuer: Option, +} + +/// 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 { + 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()); + } +} diff --git a/crates/codegen/xai-grok-shell/src/claude_import.rs b/crates/codegen/xai-grok-shell/src/claude_import.rs index 477bebc..fc76a7d 100644 --- a/crates/codegen/xai-grok-shell/src/claude_import.rs +++ b/crates/codegen/xai-grok-shell/src/claude_import.rs @@ -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(), diff --git a/crates/codegen/xai-grok-shell/src/config/reloader.rs b/crates/codegen/xai-grok-shell/src/config/reloader.rs index d64d568..9d288b3 100644 --- a/crates/codegen/xai-grok-shell/src/config/reloader.rs +++ b/crates/codegen/xai-grok-shell/src/config/reloader.rs @@ -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" ); diff --git a/crates/codegen/xai-grok-shell/src/config/tests.rs b/crates/codegen/xai-grok-shell/src/config/tests.rs index 2da2071..f23e9a2 100644 --- a/crates/codegen/xai-grok-shell/src/config/tests.rs +++ b/crates/codegen/xai-grok-shell/src/config/tests.rs @@ -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 diff --git a/crates/codegen/xai-grok-shell/src/extensions/marketplace.rs b/crates/codegen/xai-grok-shell/src/extensions/marketplace.rs index 9c1fed0..40a2e0f 100644 --- a/crates/codegen/xai-grok-shell/src/extensions/marketplace.rs +++ b/crates/codegen/xai-grok-shell/src/extensions/marketplace.rs @@ -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}"); diff --git a/crates/codegen/xai-grok-shell/src/inspect/mod.rs b/crates/codegen/xai-grok-shell/src/inspect/mod.rs index 899fcb7..c2807ae 100644 --- a/crates/codegen/xai-grok-shell/src/inspect/mod.rs +++ b/crates/codegen/xai-grok-shell/src/inspect/mod.rs @@ -73,10 +73,9 @@ pub struct InspectReport { pub lsp_servers: Vec, 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, + pub config_warnings: Vec, } #[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"] diff --git a/crates/codegen/xai-grok-shell/src/leader/client.rs b/crates/codegen/xai-grok-shell/src/leader/client.rs index cc709b4..07762e2 100644 --- a/crates/codegen/xai-grok-shell/src/leader/client.rs +++ b/crates/codegen/xai-grok-shell/src/leader/client.rs @@ -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) diff --git a/crates/codegen/xai-grok-shell/src/leader/protocol.rs b/crates/codegen/xai-grok-shell/src/leader/protocol.rs index b67ba44..0b08239 100644 --- a/crates/codegen/xai-grok-shell/src/leader/protocol.rs +++ b/crates/codegen/xai-grok-shell/src/leader/protocol.rs @@ -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] diff --git a/crates/codegen/xai-grok-shell/src/leader/server.rs b/crates/codegen/xai-grok-shell/src/leader/server.rs index d23d9cc..8ff4001 100644 --- a/crates/codegen/xai-grok-shell/src/leader/server.rs +++ b/crates/codegen/xai-grok-shell/src/leader/server.rs @@ -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!( diff --git a/crates/codegen/xai-grok-shell/src/managed_config.rs b/crates/codegen/xai-grok-shell/src/managed_config.rs index 82a9485..7120e43 100644 --- a/crates/codegen/xai-grok-shell/src/managed_config.rs +++ b/crates/codegen/xai-grok-shell/src/managed_config.rs @@ -232,10 +232,11 @@ async fn fetch_managed_config( token: &str, source: ManagedConfigSource, max_attempts: u32, + echo_principal: Option<&str>, ) -> Result { 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 { - 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 { diff --git a/crates/codegen/xai-grok-shell/src/managed_config/tests.rs b/crates/codegen/xai-grok-shell/src/managed_config/tests.rs index 3d39868..24a7e29 100644 --- a/crates/codegen/xai-grok-shell/src/managed_config/tests.rs +++ b/crates/codegen/xai-grok-shell/src/managed_config/tests.rs @@ -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!( diff --git a/crates/codegen/xai-grok-shell/src/remote/sync.rs b/crates/codegen/xai-grok-shell/src/remote/sync.rs index 7c67f93..df1d5b6 100644 --- a/crates/codegen/xai-grok-shell/src/remote/sync.rs +++ b/crates/codegen/xai-grok-shell/src/remote/sync.rs @@ -43,6 +43,20 @@ pub struct RemoteSync { } impl RemoteSync { + #[cfg(test)] + pub(crate) fn test_observer() -> (Self, mpsc::UnboundedReceiver) { + 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, diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session.rs b/crates/codegen/xai-grok-shell/src/session/acp_session.rs index 9d63bca..f261b58 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session.rs @@ -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, +} /// 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>, + /// 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>, /// 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 diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/model_switch.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/model_switch.rs index 2338367..b5b4ef0 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/model_switch.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/model_switch.rs @@ -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 { diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/run_loop.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/run_loop.rs index f9484b9..523bd9a 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/run_loop.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/run_loop.rs @@ -204,665 +204,674 @@ pub(super) async fn run_session( tokio::pin!(dream_check_sleep); loop { tokio::select! { - biased; _ = & mut idle_flush_sleep, if session.idle_flush_timeout.is_some() - && session.memory.is_enabled() && ! session.memory.is_flushing - .load(std::sync::atomic::Ordering::Relaxed) => { let current_len = session - .chat_state_handle.get_conversation_len(). await; let last_len = session - .last_idle_flush_conversation_len.load(std::sync::atomic::Ordering::Relaxed); - if current_len > last_len { tracing::info!(target : - xai_grok_telemetry::memory_log::TARGET, - "MEMORY_IDLE_FLUSH: timer fired (conversation {last_len} → {current_len})"); - session.last_idle_flush_conversation_len.store(current_len, - std::sync::atomic::Ordering::Relaxed); tokio::task::spawn_local({ let session - = session.clone(); async move { if ! session.run_memory_flush("interval", - None). await { tracing::info!(target : - xai_grok_telemetry::memory_log::TARGET, - "MEMORY_IDLE_FLUSH: skipped — another flush already in progress"); } } }); - } else { tracing::debug!(target : xai_grok_telemetry::memory_log::TARGET, - "MEMORY_IDLE_FLUSH: skipped, no new messages since last flush (len={current_len})"); - } if let Some(timeout) = session.idle_flush_timeout { idle_flush_sleep - .as_mut().reset(tokio::time::Instant::now() + timeout); } } _ = & mut - dream_check_sleep, if session.dream_check_timeout.is_some() && session.memory - .is_enabled() => { tracing::debug!(target : - xai_grok_telemetry::memory_log::TARGET, "MEMORY_DREAM_CHECK: timer fired"); - tokio::task::spawn_local({ let session = session.clone(); async move { - session.maybe_run_dream(). await; } }); if let Some(timeout) = session - .dream_check_timeout { dream_check_sleep.as_mut() - .reset(tokio::time::Instant::now() + timeout); } } changed = model_switch_rx - .changed() => { if changed.is_ok() { let new_gen = * model_switch_rx - .borrow_and_update(); session.handle_model_switch_for_laziness(new_gen). - await; } } event = chat_state_event_rx.recv() => { match event { - Some(xai_chat_state::ChatStateEvent::ConversationReset { new_len }) => { - session.last_idle_flush_conversation_len.store(new_len, - std::sync::atomic::Ordering::Relaxed); session.memory.context_injected - .store(false, std::sync::atomic::Ordering::Relaxed); } - Some(xai_chat_state::ChatStateEvent::ImageBudget { body_bytes, trigger_bytes, - reclaim_target_bytes, inline_images, needs_image_compaction, evicted, - body_bytes_after, }) => { - xai_grok_telemetry::unified_log::info("shell.image_budget", Some(session - .session_info.id.0.as_ref()), Some(serde_json::json!({ "body_bytes" : - body_bytes, "body_bytes_after" : body_bytes_after, "trigger_bytes" : - trigger_bytes, "reclaim_target_bytes" : reclaim_target_bytes, "inline_images" - : inline_images, "images_remaining" : inline_images.saturating_sub(evicted), - "needs_image_compaction" : needs_image_compaction, "evicted" : evicted, - })),); } Some(xai_chat_state::ChatStateEvent::PromptIndexChanged { .. }) | - Some(xai_chat_state::ChatStateEvent::TokensUpdated { .. }) => {} None => {} } - } maybe_event = event_rx.recv() => { if let Some(event) = maybe_event { match - event { SessionEvent::Notification(notification) => { let out = replay_buffer - .consume_chunk(notification); match out { None => {} Some((first, second)) => - { session.emit_buffered(first). await; if let Some(second) = second { session - .emit_buffered(second). await; } } } } SessionEvent::FlushReplay { respond_to - } => { if let Some(notification) = replay_buffer.flush() { session - .emit_buffered(notification). await; } if let Some(tx) = respond_to { let _ = - tx.send(()); } } } } } maybe_completion = completion_rx.recv() => { let - Some((prompt_id, result)) = maybe_completion else { if let Some(cancel) = & - session.sync_loop_cancel { cancel.cancel(); } cleanup_session_scratch(& - session); return; }; if let Some(notification) = replay_buffer.flush() { - session.emit_buffered(notification). await; } let (turn_succeeded, - infra_pause_message) = SessionActor::post_turn_goal_degradation_plan(& - result); session.handle_completion(prompt_id, result). await; session - .drain_monitor_buffer_to_pending(). await; if let Some(message) = - infra_pause_message { session.apply_infra_pause_after_turn_err(message). - await; } session.handle_turn_end(turn_succeeded). await; if session - .flush_stranded_interjections(). await { - tracing::info!("Flushed stranded interjection(s) into prompt turns"); } - SessionActor::maybe_start_running_task(session.clone(), completion_tx - .clone()). await; SessionActor::maybe_drain_notifications(session.clone(), - completion_tx.clone()). await; session.emit_session_idle_if_idle(). await; { - let s = session.clone(); tokio::task::spawn_local(async move { s - .maybe_fire_laziness_check(). await; }); } } maybe_cmd = cmd_rx.recv() => { - let Some(cmd) = maybe_cmd else { let envelope = session - .fire_hook(xai_grok_hooks::event::HookEventName::SessionEnd, None, - xai_grok_hooks::event::HookPayload::SessionEnd { reason : "channel_closed" - .to_string(), turn_count : None, tool_call_count : None, },); if let - Some(registry) = session.hook_registry.borrow().clone() { let ctx = session - .hook_run_ctx(); let results = - xai_grok_hooks::dispatcher::dispatch_non_blocking(& registry, - xai_grok_hooks::event::HookEventName::SessionEnd, & envelope, & ctx,). await; - session.send_hook_execution("session_end", None, None, & results). await; } - session.dispatch_session_end_stop("channel_closed"). await; let mut - session_end_result = "disabled"; let mut total_chunks_at_end = 0usize; if ! - session.startup_hints.is_subagent { if let Some(storage) = session.memory - .storage() { let conversation = session.chat_state_handle.get_conversation(). - await; let result = crate ::session::memory::hooks::on_session_end(& storage, - & conversation, & session.session_info.id.0, session.memory.save_on_end,); - session_end_result = match & result { crate - ::session::memory::hooks::SessionEndResult::Written(_) => "written", crate - ::session::memory::hooks::SessionEndResult::Skipped => "skipped", crate - ::session::memory::hooks::SessionEndResult::Failed(_) => "failed", }; - total_chunks_at_end = storage.total_chunk_count(); let telem = session.memory - .telemetry_snapshot(); tracing::info!(target : - xai_grok_telemetry::memory_log::TARGET, result = ? result, tool_searches = - telem.tool_search_count, injection_searches = telem.injection_count, - recovery_searches = telem.compaction_recovery_count, - "MEMORY_SESSION_END: channel closed, session summary saved"); if let crate - ::session::memory::hooks::SessionEndResult::Written(ref path_str) = result { - session.reindex_and_embed(std::path::Path::new(path_str), "session"). await; - session.send_xai_notification(XaiSessionUpdate::MemorySessionSaved { path : - path_str.clone(), }). await; } } } else { tracing::debug!(target : - xai_grok_telemetry::memory_log::TARGET, - "MEMORY_SUBAGENT_SKIP: skipping on_session_end for subagent session"); } - session.maybe_run_dream(). await; let telem = session.memory - .telemetry_snapshot(); session.emit_memory_session_summary(& telem, - total_chunks_at_end, session_end_result); if let Some(notification) = - replay_buffer.flush() { session.emit_buffered(notification). await; } { let - model_id = session.current_model_id(). await; if let Some(signals) = session - .signals_handle().snapshot(). await { - xai_grok_telemetry::session_ctx::log_event(xai_grok_telemetry::events::SessionEnded - { duration_secs : session.session_start.elapsed().as_secs(), turn_count : - signals.turn_count as u64, tool_call_count : signals.tool_call_count as u64, - compaction_count : signals.compaction_count as u64, model_id, },); } } if let - Some(cancel) = & session.sync_loop_cancel { cancel.cancel(); } session - .feedback_manager.shutdown(session.upload_queue.get()). await; if ! session - .startup_hints.is_subagent { session.persist_background_task_manifest(). - await; } cleanup_session_scratch(& session); return; }; match cmd { - SessionCommand::Initialize { system_prompt } => { session - .initialize(system_prompt). await; let s = session.clone(); let handle = - tokio::task::spawn_local(async move { s.build_prefix_background(). await }); - session.deferred_prefix.arm(handle); } SessionCommand::ReplaceSystemPrompt { - system_prompt } => { session.handle_replace_system_prompt(system_prompt). - await; } SessionCommand::RestorePlanApproval => { let s = session.clone(); - let completion_tx = completion_tx.clone(); tokio::task::spawn_local(async - move { s.resume_plan_approval(completion_tx). await; }); } - SessionCommand::Prompt { prompt_id, prompt_blocks, prompt_mode, - artifact_upload_ctx, client_identifier, screen_mode, verbatim, traceparent, - json_schema, send_now, admission, respond_to, persist_ack, parsed_prompt_tx } - => { let origin = super::PromptOrigin::from_prompt_id(& prompt_id); let - (actor_admitted, task_wake_fallback) = match admission { Some(admission) => { - let fallback = session.admit_task_completion_wake(& origin, admission). - await; (fallback.is_some(), fallback) } None => (true, None), }; if ! - actor_admitted { SessionActor::respond_removed_prompt(respond_to); continue; - } session.ensure_prefix_ready(). await; if ! origin.is_synthetic() { if let - Some(gate) = & session.tool_context.task_wake_suppressed { gate.set(false); } - let mut state = session.state.lock(). await; state.notifications_suppressed = - false; xai_grok_telemetry::unified_log::info("shell.task_wake.gate_cleared", - Some(session.session_info.id.0.as_ref()), Some(serde_json::json!({ "reason" : - "user_intake" })),); session.user_input_generation.fetch_add(1, - std::sync::atomic::Ordering::AcqRel); } if origin.is_synthetic() { let state - = session.state.lock(). await; let has_running = state.running_task - .is_some(); let queue_depth = state.pending_inputs.len(); drop(state); - tracing::info!(prompt_id = % prompt_id, has_running_task = has_running, - queue_depth = queue_depth, - "auto-wake: session actor received synthetic prompt"); } if let Some(ref tp) - = traceparent { let meta = serde_json::json!({ "traceparent" : tp }); - xai_file_utils::trace_context::link_current_span_to_meta(& meta); } let - (trace_gcs_config, artifact_tracker) = match artifact_upload_ctx { Some(tu) - => (Some(tu.gcs_config), Some(tu.artifact_tracker)), None => (None, None), }; - let cancel_for_send_now = session.queue_input(prompt_blocks, prompt_id, - prompt_mode, trace_gcs_config, artifact_tracker, client_identifier, - screen_mode, verbatim, json_schema, send_now, task_wake_fallback, respond_to, - persist_ack, parsed_prompt_tx). await; if cancel_for_send_now { session - .cancel_turn_for_send_now(& mut replay_buffer). await; } - SessionActor::maybe_start_running_task(session.clone(), completion_tx - .clone()). await; } SessionCommand::SessionMode { session_mode, responds_to } - => { session.handle_session_mode(session_mode). await; let _ = responds_to - .send(()); } SessionCommand::SetSessionModel { sampling_config, use_concise, - apply_prompt_override, skip_prompt_rewrite, auto_compact_threshold_percent, - responds_to } => { let updated_model_id = session - .handle_set_session_model(sampling_config, use_concise, - apply_prompt_override, skip_prompt_rewrite, auto_compact_threshold_percent). - await; let _ = responds_to.send(updated_model_id); } - SessionCommand::RebuildAgentForDefinition { definition, responds_to } => { - let outcome = session.handle_rebuild_agent_for_definition(definition). await; - let _ = responds_to.send(outcome); } SessionCommand::OverrideModelName { - model_name, extra_headers, context_window } => { if let Some(mut cfg) = - session.chat_state_handle.get_sampling_config(). await { - tracing::info!(target : SESSION_LOG, session_id = % session.session_info.id, - old_model = % cfg.model, new_model = % model_name, extra_header_count = - extra_headers.len(), old_context_window = cfg.context_window.get(), - new_context_window = ? context_window.map(| cw | cw.get()), - "OVERRIDE_MODEL: changing model name in sampling config"); session - .signals_handle().set_primary_model(& model_name); cfg.model = model_name - .clone(); cfg.extra_headers.extend(extra_headers); if let Some(cw) = - context_window && session.compaction.context_window_override.is_none() { cfg - .context_window = cw; } session.chat_state_handle - .update_sampling_config(cfg); let existing = session.chat_state_handle - .get_credentials(). await; if let Some(r) = crate - ::agent::config::try_resolve_model_credentials(model_name.as_str(), existing - .api_key.as_deref()) { session.chat_state_handle - .update_credentials(xai_chat_state::Credentials { api_key : r.api_key, - auth_type : r.auth_type, alpha_test_key : existing.alpha_test_key, - client_version : existing.client_version, }); } session.model_auth_facts - .replace(None); } } SessionCommand::GetCurrentModel { responds_to } => { let - model = session.chat_state_handle.get_sampling_config(). await .map(| c | c - .model).unwrap_or_default(); let _ = responds_to.send(model); } - SessionCommand::GetCurrentPromptMode { responds_to } => { let mode = * - session.current_prompt_mode.lock(); let _ = responds_to.send(mode); } - SessionCommand::GetModelMetadata { responds_to } => { let id = session - .chat_state_handle.get_last_model_metadata(). await; let _ = responds_to - .send(id); } SessionCommand::GetSessionInfo { responds_to } => { let info = - session.build_session_info(). await; let _ = responds_to.send(info); } - SessionCommand::BackgroundForegroundCommand { tool_call_id, respond_to } => { - let result = session.agent.borrow().tool_bridge() - .background_foreground_command(& tool_call_id). await; let _ = respond_to - .send(result); } SessionCommand::KillBackgroundTask { task_id, respond_to } - => { let result = session.agent.borrow().tool_bridge().kill_background_task(& - task_id). await .map_err(| e | e.to_string()); let _ = respond_to - .send(result); } SessionCommand::DeleteScheduledTask { task_id, respond_to } - => { let result = session.agent.borrow().tool_bridge() - .delete_scheduled_task(& task_id). await .map_err(| e | e.to_string()); let _ - = respond_to.send(result); } SessionCommand::ListTasks { respond_to } => { - let result = session.agent.borrow().tool_bridge().list_tasks(). await; let _ - = respond_to.send(result); } SessionCommand::GetHooksList { respond_to } => { - use crate ::extensions::hooks::hook_spec_to_info; let hooks = match &* - session.hook_registry.borrow() { Some(registry) => registry.all_hooks() - .iter().map(| spec | hook_spec_to_info(spec)).collect(), None => Vec::new(), - }; let project_trusted = crate - ::agent::folder_trust::project_scope_allowed(std::path::Path::new(& session - .session_info.cwd),); let _ = respond_to - .send(xai_hooks_plugins_types::HooksListResponse { hooks, project_trusted, - load_errors : session.hook_load_errors.borrow().clone(), }); } - SessionCommand::HooksAction { action, respond_to } => { let outcome = session - .handle_hooks_action(action). await; let _ = respond_to.send(outcome); } - SessionCommand::NotifyPluginUpdates { updates } => { session - .send_xai_notification(XaiSessionUpdate::PluginUpdatesInstalled { updates },) - . await; } SessionCommand::PluginsAction { action, respond_to } => { let - outcome = session.handle_plugins_action(action). await; let _ = respond_to - .send(outcome); } SessionCommand::PluginsList { respond_to } => { let _ = - respond_to.send(session.plugin_registry.borrow().clone()); } - SessionCommand::DispatchNotificationHook { notification_type, message, title, - level, } => { session.dispatch_notification_hook(& notification_type, - message, title, level,). await; } SessionCommand::DropMonitorNotifications { - task_id } => { { let mut state = session.state.lock(). await; state - .pending_notifications.retain(| n | { ! matches!(& n.source, - NotificationSource::MonitorEvent { task_id : tid } if tid == & task_id) }); } - if let Some(buffer) = & session.tool_context.monitor_event_buffer { let - dropped = buffer.drain_matching(| e | e.task_id == task_id); if ! dropped - .is_empty() { tracing::debug!(task_id = % task_id, dropped = dropped.len(), - "dropped buffered monitor events after TaskCompleted auto-wake"); } } } - SessionCommand::InjectNotification { prompt_id, prompt_blocks, priority, - source } => { let is_turn_active = session.tool_context.is_turn_active - .as_ref().map(| f | f.load(std::sync::atomic::Ordering::Relaxed)) - .unwrap_or(false); if is_turn_active && priority == - NotificationPriority::Next { if let Some(buffer) = & session.tool_context - .monitor_event_buffer { let non_text_count = prompt_blocks.iter().filter(| b - | ! matches!(b, acp::ContentBlock::Text(_))).count(); if non_text_count > 0 { - tracing::debug!(non_text_count, - "Non-text content blocks dropped in mid-turn monitor event routing"); } let - event_text = prompt_blocks.iter().filter_map(| b | { if let - acp::ContentBlock::Text(t) = b { Some(t.text.clone()) } else { None } }) - .collect::< Vec < _ >> ().join("\n"); let task_id = source.task_id() - .to_owned(); const MAX_BUFFER_EVENTS : usize = 50; buffer - .push_capped(xai_grok_tools::implementations::grok_build::task::types::MonitorEventNotification - { task_id : task_id.clone(), event_text, owner_session_id : Some(session - .session_info.id.0.to_string(),), }, MAX_BUFFER_EVENTS,); - tracing::debug!(task_id = % task_id, - "Routed monitor event to mid-turn buffer"); } } else { { let mut state = - session.state.lock(). await; SessionActor::push_pending_notification(& mut - state, PendingNotification { prompt_id, prompt_blocks, priority, source, },); - } SessionActor::maybe_drain_notifications(session.clone(), completion_tx - .clone()). await; } } SessionCommand::RecordGoalTurnTaskIds { task_ids } => { - session.record_reparented_goal_turn_task_ids(task_ids); } - SessionCommand::RemoveQueuedPrompt { id, expected_version, owner } => { - session.handle_remove_queued_prompt(& id, expected_version, owner.as_deref()) - . await; } SessionCommand::ReorderQueue { ordered_ids } => { session - .handle_reorder_queue(& ordered_ids). await; } SessionCommand::ClearQueue { - owner } => { session.handle_clear_queue(owner.as_deref()). await; } - SessionCommand::EditQueuedPrompt { id, new_text, editor } => { session - .handle_edit_queued_prompt(& id, new_text, editor.as_deref()). await; } - SessionCommand::InterjectQueuedPrompt { id, expected_version, owner, new_text - } => { let cancel_for_send_now = session.handle_interject_queued_prompt(& id, - expected_version, owner.as_deref(), new_text.as_deref()). await; if - cancel_for_send_now { session.cancel_turn_for_send_now(& mut replay_buffer). - await; } SessionActor::maybe_start_running_task(session.clone(), - completion_tx.clone()). await; } SessionCommand::Cancel { cancel_subagents, - kill_background_tasks, rewind_if_pristine, trigger, } => { if let - Some(notification) = replay_buffer.flush() { session - .emit_buffered(notification). await; } session.pending_interjections.clear(); - let suppress_task_wakes = trigger.as_deref() == Some("ctrl_c"); session - .cancel_running_task(cancel_subagents, kill_background_tasks, - rewind_if_pristine, trigger,). await; session.auto_pause_goal_if_active(crate - ::session::goal_tracker::GoalPauseReason::User,). await; - SessionActor::maybe_start_running_task(session.clone(), completion_tx - .clone()). await; if ! suppress_task_wakes { - SessionActor::maybe_drain_notifications(session.clone(), completion_tx - .clone(),). await; } } SessionCommand::CompactSession { user_context, - respond_to } => { let s = session.clone(); tokio::task::spawn_local(async - move { let compact_session = s.run_compact(user_context). await; let _ = - respond_to.send(compact_session); }); } SessionCommand::ReloadPlugins { - registry } => { if ! session.startup_hints.is_subagent { let registry = - session.preserve_session_plugin_dirs(registry); session - .apply_plugin_registry_snapshot(registry). await; } } - SessionCommand::ReloadHooks => { if ! session.startup_hints.is_subagent { let - _ = session.reload_hooks_impl(). await; } } - SessionCommand::RefreshSkillBaseline => { let s = session.clone(); - tokio::task::spawn_local(async move { let cwd = s.tool_context.cwd.as_path() - .to_string_lossy(); let skills_config = crate ::util::config::load_config(). - await .skills; let pr = s.plugin_registry.borrow().clone(); let new_skills = - xai_grok_agent::prompt::skills::list_skills_with_plugins(Some(& cwd), & - skills_config, pr.as_deref(), s.rebuild_spec.compat,). await; - tracing::info!(skills = new_skills.len(), - "refreshed skill baseline after bundle sync"); let bridge = s.agent.borrow() - .tool_bridge().clone(); bridge.update_skill_baseline(new_skills). await; if - let Some(effects) = bridge.apply_pending_skill_update(). await { s - .apply_skill_update_effects(effects). await; } }); } - SessionCommand::FlushMemory { respond_to } => { let s = session.clone(); - tokio::task::spawn_local(async move { if s.memory.is_enabled() { let - did_flush = s.run_memory_flush("user_requested", None). await; let _ = - respond_to.send(Ok(did_flush)); } else { let _ = respond_to - .send(Err(acp::Error::invalid_request() - .data("memory is not enabled for this session".to_string()))); } }); } - SessionCommand::SetYoloMode { enabled } => { let was = session.permissions - .is_yolo_mode(); tracing::info!("Session received SetYoloMode: {}", enabled); - session.permissions.set_yolo_mode(enabled); let actual = session.permissions - .is_yolo_mode(); if let Some(enabled) = yolo_toggle_report(was, actual) { - session.emit_event(crate ::session::events::Event::YoloToggled { enabled }); - } } SessionCommand::SetAutoMode { enabled } => { let enabled = enabled && - crate ::util::config::auto_permission_mode_enabled_from_disk(); - tracing::info!("Session received SetAutoMode: {}", enabled); session - .permissions.set_auto_mode(enabled); if enabled { session - .wire_permission_auto_llm_classifier(). await; } else { session.permissions - .set_llm_side_query_wired(false); } } SessionCommand::ResetPermissionState => - { session.permissions.reset_state(); tracing::info!(session_id = % session - .session_info.id, "Permission state reset via notification"); } - SessionCommand::Rewind { request, respond_to } => { let s = session.clone(); - tokio::task::spawn_local(async move { let result = s.handle_rewind(request). - await; let _ = respond_to.send(result); }); } SessionCommand::RepairHistory { - dry_run, respond_to } => { let s = session.clone(); - tokio::task::spawn_local(async move { let result = s - .handle_repair_history(dry_run). await; let _ = respond_to.send(result); }); - } SessionCommand::GetRewindPoints { respond_to } => { let response = session - .get_rewind_points(). await; let _ = respond_to.send(response); } - SessionCommand::GetRewindFileCounts { respond_to } => { let _ = respond_to - .send(session.rewind_file_counts(). await); } - SessionCommand::ReconcileRewindTracker { target_prompt_index } => { session - .merge_rewind_tracker_from(target_prompt_index). await; } - SessionCommand::XaiSessionNotification { notification } => { session - .handle_xai_session_notification(notification). await; } - SessionCommand::RecordSubagentUsage { by_model, parent_prompt_id, incomplete, - respond_to, } => { use super::updates::SubagentUsageApply; match session - .record_subagent_usage(& by_model, parent_prompt_id.as_deref(), incomplete,). - await { Ok(SubagentUsageApply::AttributedToPrompt) => { let _ = respond_to - .send(()); } Ok(SubagentUsageApply::SessionOnly) => { let _ = session - .mark_subagent_usage_not_applied(parent_prompt_id.as_deref(),). await; let _ - = respond_to.send(()); } Err(()) => {} } } - SessionCommand::MarkSubagentUsageNotApplied { parent_prompt_id, respond_to, } - => { if session.mark_apply_miss_incomplete(parent_prompt_id.as_deref()). - await { let _ = respond_to.send(()); } } - SessionCommand::ErrorPathUsageFallback { prompt_id, respond_to, } => { let - pid = prompt_id.or_else(|| { session.current_prompt_id.lock().ok().and_then(| - g | g.clone()) }); let usage = match pid.as_deref() { Some(id) => session - .error_path_usage_fallback(id). await, None => { match session - .chat_state_handle.try_get_prompt_usage(). await { Ok(ledger) => { crate - ::extensions::notification::PromptUsage::for_error_path(ledger.as_ref(), - false,) } Err(()) => { crate - ::extensions::notification::PromptUsage::for_error_path(None, true,) } } } }; - let _ = respond_to.send(usage); } SessionCommand::SetNextTraceTurn { - next_trace_turn, request_id, } => { let _ = session.notifications - .persistence_tx.send(PersistenceMsg::NextTraceTurn { next_trace_turn, - request_id, }); } SessionCommand::CopyFile { respond_to } => { if let - Some(notification) = replay_buffer.flush() { session - .emit_buffered(notification). await; } let _ = session.notifications - .persistence_tx.send(PersistenceMsg::CopyFile { one_shot : respond_to }); } - SessionCommand::IsBusy { respond_to } => { let busy = { let state = session - .state.lock(). await; state_is_busy(& state) }; let _ = respond_to - .send(busy); } SessionCommand::FlushComplete { respond_to } => { if let - Some(notification) = replay_buffer.flush() { session - .emit_buffered(notification). await; } let _ = session.notifications - .persistence_tx.send(PersistenceMsg::FlushAndAck { respond_to }); } - SessionCommand::UpdateMcpServers { mcp_servers, respond_to } => { if session - .startup_hints.is_subagent { tracing::debug!(session_id = % session - .session_info.id.0, "Skipping UpdateMcpServers for subagent session",); let _ - = respond_to.send(Ok(())); continue; } - tracing::info!("Updating MCP servers for session '{}' ({} servers)", session - .session_info.id.0, mcp_servers.len()); session.reseed_mcp_output_cap(). - await; let (diff, dispatch_event_tx) = { let mut mcp_state = session - .mcp_state.lock(). await; let diff = mcp_state - .update_configs_diff(mcp_servers); let tx = mcp_state.client_event_tx(); - (diff, tx) }; let Some(diff) = diff else { - tracing::debug!("MCP configs unchanged for session '{}', skipping re-initialization", - session.session_info.id.0); let _ = respond_to.send(Ok(())); continue; }; if - (! diff.added.is_empty() || ! diff.removed.is_empty()) && let Some(tx) = & - dispatch_event_tx { let _ = tx - .send(xai_grok_mcp::servers::McpClientEvent::ConfigDiff { added : diff.added - .clone(), removed : diff.removed.clone(), },); } for name in & diff.removed { - let prefix = format!("{}{}", name, crate - ::session::mcp_servers::MCP_TOOL_NAME_DELIMITER); let removed_count = session - .agent.borrow().tool_bridge().unregister_tools_by_prefix(& prefix); - tracing::info!(server = name.as_str(), tools_removed = removed_count, - "Unregistered tools for removed MCP server"); } let session_for_mcp = session - .clone(); tokio::task::spawn_local(async move { session_for_mcp - .ensure_mcp_tools_initialized(). await; let _ = respond_to.send(Ok(())); }); - } SessionCommand::ToggleMcpServer { server_name, enabled, server_config, - respond_to } => { session.events - .emit(xai_file_utils::events::Event::McpServerToggled { server_name : - server_name.clone(), enabled, }); let mut mcp_state = session.mcp_state - .lock(). await; let mut configs = mcp_state.configs.clone(); if enabled { if - let Some(config) = server_config { configs.retain(| c | { crate - ::session::mcp_servers::mcp_server_name(c) != server_name }); configs - .push(config); } else { let already_present = configs.iter().any(| c | { - crate ::session::mcp_servers::mcp_server_name(c) == server_name }); if - already_present { drop(mcp_state); let _ = respond_to.send(Ok(())); continue; - } drop(mcp_state); let _ = respond_to.send(Err(acp::Error::invalid_params() - .data(format!("server '{}' not found in config", server_name)))); continue; } - } else { configs.retain(| c | crate - ::session::mcp_servers::mcp_server_name(c) != server_name); } let diff = - mcp_state.update_configs_diff(configs); let dispatch_event_tx = mcp_state - .client_event_tx(); drop(mcp_state); let Some(diff) = diff else { let _ = - respond_to.send(Ok(())); continue; }; if (! diff.added.is_empty() || ! diff - .removed.is_empty()) && let Some(tx) = & dispatch_event_tx { let _ = tx - .send(xai_grok_mcp::servers::McpClientEvent::ConfigDiff { added : diff.added - .clone(), removed : diff.removed.clone(), },); } for name in & diff.removed { - let prefix = format!("{}{}", name, crate - ::session::mcp_servers::MCP_TOOL_NAME_DELIMITER); let removed_count = session - .agent.borrow().tool_bridge().unregister_tools_by_prefix(& prefix); - tracing::info!(server = name.as_str(), tools_removed = removed_count, - "Unregistered tools for toggled MCP server"); } let session_for_mcp = session - .clone(); let sname = server_name.clone(); tokio::task::spawn_local(async - move { session_for_mcp.ensure_mcp_tools_initialized(). await; if let Err(e) = - crate ::util::config::save_mcp_server_enabled(& sname, enabled,). await { - tracing::warn!(server = sname.as_str(), error = % e, - "Failed to persist server enabled state to config"); } let _ = respond_to - .send(Ok(())); }); } SessionCommand::ToggleMcpTool { server_name, tool_name, - enabled, is_managed_gateway, respond_to } => { if is_managed_gateway { let - mut disabled_tools = crate - ::util::config::get_all_mcp_disabled_tools(std::path::Path::new(& session - .session_info.cwd)); if tool_name.is_empty() { let set = disabled_tools - .entry(crate ::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY - .to_string()).or_default(); if enabled { set.remove(& server_name); } else { - set.insert(server_name.clone()); } if set.is_empty() { disabled_tools - .remove(crate ::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY); } } - else if enabled { if let Some(set) = disabled_tools.get_mut(& server_name) { - set.remove(& tool_name); if set.is_empty() { disabled_tools.remove(& - server_name); } } } else { disabled_tools.entry(server_name.clone()) - .or_default().insert(tool_name.clone()); } session - .refresh_mcp_snapshot_and_schedule_reminder_with_disabled(& disabled_tools,). - await; session.refresh_goal_harness_enabled(). await; let disabled_vec : Vec - < String > = if tool_name.is_empty() { disabled_tools.get(crate - ::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY).map(| s | s.iter() - .cloned().collect()).unwrap_or_default() } else { disabled_tools.get(& - server_name).map(| s | s.iter().cloned().collect()).unwrap_or_default() }; - let notifications = session.notifications.gateway.clone(); let session_id = - session.session_info.id.0.clone(); let server_for_persist = if tool_name - .is_empty() { crate ::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY - .to_string() } else { server_name.clone() }; tokio::task::spawn_local(async - move { if let Err(e) = crate ::util::config::save_mcp_disabled_tools(& - server_for_persist, & disabled_vec,). await { tracing::warn!(server = - server_for_persist.as_str(), error = % e, - "Failed to persist disabled_tools to config"); } let payload = crate - ::extensions::mcp::McpToolsChanged { session_id : session_id.to_string(), - server_name : String::new(), tools : Vec::new(), }; if let Ok(params) = - serde_json::value::to_raw_value(& payload) { notifications - .forward_fire_and_forget(acp::ExtNotification::new("x.ai/mcp/tools_changed", - params.into())); } let _ = respond_to.send(Ok(())); }); continue; } let - qualified = format!("{}{}{}", server_name, crate - ::session::mcp_servers::MCP_TOOL_NAME_DELIMITER, tool_name,); let mut - mcp_state = session.mcp_state.lock(). await; if enabled { if let Some(set) = - mcp_state.disabled_tools.get_mut(& server_name) { set.remove(& tool_name); if - set.is_empty() { mcp_state.disabled_tools.remove(& server_name); } } if let - Some(reg) = mcp_state.disabled_tool_registrations.remove(& qualified) && reg - .model_visible { let bridge = session.agent.borrow().tool_bridge().clone(); - if let Err(e) = bridge.register_mcp_tools(reg.name, reg.tool, Some(reg - .input_schema)). await { tracing::warn!(tool = qualified.as_str(), error = % - e, "Failed to re-register toggled MCP tool"); } } } else { let bridge = - session.agent.borrow().tool_bridge().clone(); let tool_def = bridge - .tool_definitions(). await .into_iter().find(| d | d.function.name == - qualified); if let Some(def) = tool_def { let meta = mcp_state.mcp_tool_meta - .get(& qualified).cloned(); let schema = def.function.parameters.clone(); let - mcp_tool = crate ::session::mcp_servers::McpTool::new(tool_name.clone(), def - .function.description.clone().unwrap_or_default(), server_name.clone(), - session.mcp_state.clone(), schema, meta,); if let Some(reg) = mcp_tool - .into_registration() { mcp_state.disabled_tool_registrations.insert(qualified - .clone(), reg); } } bridge.unregister_tool_by_name(& qualified); mcp_state - .disabled_tools.entry(server_name.clone()).or_default().insert(tool_name - .clone()); } let disabled_vec : Vec < String > = mcp_state.disabled_tools - .get(& server_name).map(| s | s.iter().cloned().collect()) - .unwrap_or_default(); drop(mcp_state); session - .refresh_mcp_snapshot_and_schedule_reminder(). await; session - .refresh_goal_harness_enabled(). await; let notifications = session - .notifications.gateway.clone(); let session_id = session.session_info.id.0 - .clone(); let server_for_persist = server_name.clone(); - tokio::task::spawn_local(async move { if let Err(e) = crate - ::util::config::save_mcp_disabled_tools(& server_for_persist, & - disabled_vec,). await { tracing::warn!(server = server_for_persist.as_str(), - error = % e, "Failed to persist disabled_tools to config"); } let payload = - crate ::extensions::mcp::McpToolsChanged { session_id : session_id - .to_string(), server_name : String::new(), tools : Vec::new(), }; if let - Ok(params) = serde_json::value::to_raw_value(& payload) { notifications - .forward_fire_and_forget(acp::ExtNotification::new(crate - ::extensions::mcp::mcp_methods::TOOLS_CHANGED, params.into())); } let _ = - respond_to.send(Ok(())); }); } SessionCommand::SnapshotMcpPool { respond_to } - => { let mcp_state = session.mcp_state.lock(). await; let pool = if mcp_state - .owned_clients.is_empty() && mcp_state.shared_clients.is_empty() { None } - else { Some(crate ::session::mcp_servers::SharedMcpPool::from_state(& - mcp_state)) }; let _ = respond_to.send(pool); } - SessionCommand::SnapshotClientHooks { respond_to } => { let _ = respond_to - .send(session.client_hooks.borrow().clone()); } - SessionCommand::SnapshotToolDefinitions { respond_to } => { let defs = - session.prepare_tool_definitions_inner(). await; let specs = session - .turn_base_tool_specs(& defs); let _ = respond_to.send(specs); } - SessionCommand::SetClientHooks { hooks } => { * session.client_hooks - .borrow_mut() = hooks; } SessionCommand::GetMcpStatus { respond_to } => { let - mcp_state = session.mcp_state.clone(); let tool_bridge = session.agent - .borrow().tool_bridge().clone(); let writer = session.events.writer(); - tokio::task::spawn_local(async move { let snapshot = crate - ::extensions::mcp::build_mcp_status(& mcp_state, & tool_bridge, Some(& - writer),). await; let _ = respond_to.send(snapshot); }); } - SessionCommand::CallMcpTool { server_name, server_url, tool_name, arguments, - respond_to } => { let mcp_state = session.mcp_state.clone(); - tokio::task::spawn_local(async move { let result = crate - ::extensions::mcp::call_mcp_tool(& mcp_state, & server_name, server_url - .as_deref(), & tool_name, arguments,). await; let _ = respond_to - .send(result); }); } SessionCommand::ReadMcpResource { server_name, uri, - respond_to } => { let mcp_state = session.mcp_state.clone(); - tokio::task::spawn_local(async move { let result = crate - ::extensions::mcp::read_mcp_resource(& mcp_state, & server_name, & uri,). - await; let _ = respond_to.send(result); }); } SessionCommand::McpAuthStatus { - respond_to } => { let mcp_state = session.mcp_state.clone(); - tokio::task::spawn_local(async move { let state = mcp_state.lock(). await; - let entries : Vec < _ > = state.auth_required.iter().map(| name | { crate - ::extensions::mcp::McpAuthStatusEntry { server_name : name.clone(), status : - "needs_auth", } }).collect(); let _ = respond_to.send(entries); }); } - SessionCommand::McpAuthTrigger { server_name, respond_to } => { let s = - session.clone(); tokio::task::spawn_local(async move { let result = s - .handle_mcp_auth_trigger(& server_name). await; let _ = respond_to - .send(result); }); } SessionCommand::GetManagedGatewayDisabledTools { - respond_to } => { let disabled_tools = crate - ::util::config::get_all_mcp_disabled_tools(std::path::Path::new(& session - .session_info.cwd),); let _ = respond_to.send(disabled_tools); } - SessionCommand::RetryAuthRequiredServers { respond_to } => { let s = session - .clone(); tokio::task::spawn_local(async move { s - .retry_auth_required_servers(). await; let _ = respond_to.send(()); }); } - SessionCommand::RefreshMcpSearchIndex => { session - .refresh_mcp_snapshot_and_schedule_reminder(). await; } - SessionCommand::TriggerTestFeedback { tier, mode, respond_to } => { let s = - session.clone(); tokio::task::spawn_local(async move { let request = s - .feedback_manager.force_feedback_request(tier, mode). await; let notification - = crate ::extensions::notification::FeedbackRequestNotification::from(request - .clone()); s.send_feedback_notification(request). await; let resp = - ExtMethodResult::success(notification).to_ext_response(); let _ = respond_to - .send(resp); }); } SessionCommand::PersistFeedback(entry) => { let _ = - session.notifications.persistence_tx.send(PersistenceMsg::Feedback(* entry)); - } SessionCommand::AdvertiseCommands => { session - .send_available_commands_update(). await; } SessionCommand::ReloadSkills => { - let s = session.clone(); tokio::task::spawn_local(async move { s - .reload_skills_from_disk(). await; }); } - SessionCommand::DispatchSessionStartHook { source } => { let envelope = - session.fire_hook(xai_grok_hooks::event::HookEventName::SessionStart, None, - xai_grok_hooks::event::HookPayload::SessionStart { source, model_id : None, - agent_type : None, },); if let Some(registry) = session.hook_registry - .borrow().clone() { let ctx = session.hook_run_ctx(); let results = - xai_grok_hooks::dispatcher::dispatch_non_blocking(& registry, - xai_grok_hooks::event::HookEventName::SessionStart, & envelope, & ctx,). - await; session.send_hook_execution("session_start", None, None, & results). - await; } } SessionCommand::GetFeedbackContext { turn_number, responds_to } => - { let s = session.clone(); tokio::task::spawn_local(async move { use - prod_mc_cli_chat_proxy_types::feedback_types::FeedbackToolOutcome; let - turn_idx = turn_number.and_then(| n | usize::try_from(n).ok()); let - (last_user_message, last_assistant_message) = match turn_idx { Some(n) => { - let conv = s.chat_state_handle.get_conversation(). await; - turn_texts_for_feedback(& conv, n) } None => { tokio::join!(s - .chat_state_handle.get_last_user_query_text(), s.chat_state_handle - .get_last_assistant_text(),) } }; let sh = s.signals_handle(); let (signals, - tool_outcomes) = tokio::join!(sh.snapshot(), sh.last_turn_tool_outcomes(),); - let signals = signals.unwrap_or_default(); let ctx = FeedbackContext { - last_user_message, last_assistant_message, tool_outcomes : tool_outcomes - .into_iter().map(| o | FeedbackToolOutcome { tool_name : o.tool_name, calls : - o.successes + o.failures, failures : o.failures, }).collect(), - compaction_count : signals.compaction_count as i64, context_window_usage : - signals.context_window_usage, context_tokens_used : signals - .context_tokens_used, context_window_tokens : signals.context_window_tokens, - session_cwd : s.tool_context.cwd.as_path().to_string_lossy().to_string(), }; - let _ = responds_to.send(ctx); }); } SessionCommand::GetActiveAgent { - responds_to } => { let agent_type = session.active_agent_type.lock().clone(); - let _ = responds_to.send(agent_type); } SessionCommand::SideQuestion { - question, respond_to } => { let s = session.clone(); - tokio::task::spawn_local(async move { let result = s.handle_side_question(& - question). await; let _ = respond_to.send(result); }); } - SessionCommand::Recap { auto } => { let s = session.clone(); - tokio::task::spawn_local(async move { s.handle_recap(auto). await; }); } - SessionCommand::AISuggest { prefix, cwd, model_override, respond_to } => { - let s = session.clone(); tokio::task::spawn_local(async move { let result = s - .handle_ai_suggest(& prefix, & cwd, model_override.as_deref()). await; let _ - = respond_to.send(result); }); } SessionCommand::SuggestPrompt { - model_override, respond_to } => { let s = session.clone(); - tokio::task::spawn_local(async move { let result = s - .handle_suggest_prompt(model_override.as_deref()). await; let _ = respond_to - .send(result); }); } SessionCommand::RewriteMemoryNote { raw_text, - context_summary, respond_to } => { let s = session.clone(); - tokio::task::spawn_local(async move { let result = s - .handle_rewrite_memory_note(& raw_text, & context_summary). await; let _ = - respond_to.send(result); }); } SessionCommand::Interject { text, id, images } - => { session.broadcast_interjection(& text, id.as_deref()); session.events - .emit(crate ::session::events::Event::Interjected { source : crate - ::session::events::InterjectionSource::Direct, image_count : images.len() as - u32, redirect_kind : crate ::session::events::RedirectKind::Interjection, }); - let turn_running = session.current_prompt_id.lock().ok().and_then(| g | g - .clone()).is_some(); if turn_running { session.pending_interjections - .push(PendingInterjection { text, attachments : images, }); - tracing::info!("Queued mid-turn interjection"); } else { session - .queue_interjection_fallback_prompt(text, images, true). await; - SessionActor::maybe_start_running_task(session.clone(), completion_tx - .clone(),). await; } } SessionCommand::GoalSummaryTurn { prompt_text } => { - let prompt_id = format!("goal-summary-{}", uuid::Uuid::now_v7()); let - prompt_blocks = - vec![acp::ContentBlock::Text(acp::TextContent::new(prompt_text))]; let - (respond_to, _) = tokio::sync::oneshot::channel(); { let mut state = session - .state.lock(). await; state.pending_inputs.push_back(InputItem { prompt_id, - prompt_blocks, prompt_mode : crate ::session::plan_mode::PromptMode::Agent, - trace_gcs_config : None, artifact_tracker : None, client_identifier : None, - screen_mode : None, verbatim : true, json_schema : None, origin : - super::PromptOrigin::GoalSummary, task_wake_fallback : None, respond_to, - persist_ack : None, parsed_prompt_tx : None, queue_meta : None, send_now : - false, }); } SessionActor::maybe_start_running_task(session.clone(), - completion_tx.clone()). await; } SessionCommand::TakeTurnMessages { - respond_to } => { let result = session.chat_state_handle.take_turn_messages() - . await; let _ = respond_to.send(result); } - SessionCommand::TakeHarnessTraceTurns { respond_to } => { let result = - session.chat_state_handle.take_harness_trace_turns(). await; let _ = - respond_to.send(result); } SessionCommand::TakeStreamingCapture { prompt_id, - respond_to } => { let taken = { let mut cap = session.streaming_turn_capture - .lock(); if cap.prompt_id.as_deref() == Some(prompt_id.as_str()) { - Some(std::mem::take(& mut * cap)) } else { if ! cap.is_empty() { - tracing::warn!(requested_prompt_id = % prompt_id, slot_prompt_id = ? cap - .prompt_id, - "streaming_capture race: live slot belongs to a different prompt; \ - dropping streaming_partial.json for the requested turn",); - } None } }; let result = taken.and_then(| mut cap | { cap - .finalize_for_upload(); (! cap.is_empty()).then_some(cap) }); let _ = - respond_to.send(result); } SessionCommand::PersistGitHead { commit, branch } - => { let _ = session.notifications.persistence_tx - .send(PersistenceMsg::GitHead { commit, branch },); } SessionCommand::Shutdown => { if let Some(notification) = replay_buffer - .flush() { session.emit_buffered(notification). await; } session - .drop_pending_synthetic_items(). await; let envelope = session - .fire_hook(xai_grok_hooks::event::HookEventName::SessionEnd, None, - xai_grok_hooks::event::HookPayload::SessionEnd { reason : "shutdown" - .to_string(), turn_count : None, tool_call_count : None, },); if let - Some(registry) = session.hook_registry.borrow().clone() { let ctx = session - .hook_run_ctx(); let results = - xai_grok_hooks::dispatcher::dispatch_non_blocking(& registry, - xai_grok_hooks::event::HookEventName::SessionEnd, & envelope, & ctx,). await; - session.send_hook_execution("session_end", None, None, & results). await; } - session.dispatch_session_end_stop("shutdown"). await; let mut - session_end_result = "disabled"; let mut total_chunks_at_end = 0usize; if ! - session.startup_hints.is_subagent { if let Some(storage) = session.memory - .storage() { let conversation = session.chat_state_handle.get_conversation(). - await; let result = crate ::session::memory::hooks::on_session_end(& storage, - & conversation, & session.session_info.id.0, session.memory.save_on_end,); - session_end_result = match & result { crate - ::session::memory::hooks::SessionEndResult::Written(_) => "written", crate - ::session::memory::hooks::SessionEndResult::Skipped => "skipped", crate - ::session::memory::hooks::SessionEndResult::Failed(_) => "failed", }; - total_chunks_at_end = storage.total_chunk_count(); let telem = session.memory - .telemetry_snapshot(); tracing::info!(target : - xai_grok_telemetry::memory_log::TARGET, result = ? result, tool_searches = - telem.tool_search_count, injection_searches = telem.injection_count, - recovery_searches = telem.compaction_recovery_count, - "MEMORY_SESSION_END: session summary saved"); if let crate - ::session::memory::hooks::SessionEndResult::Written(ref path_str) = result { - session.reindex_and_embed(std::path::Path::new(path_str), "session"). await; - session.send_xai_notification(XaiSessionUpdate::MemorySessionSaved { path : - path_str.clone(), }). await; } } } else { tracing::debug!(target : - xai_grok_telemetry::memory_log::TARGET, - "MEMORY_SUBAGENT_SKIP: skipping on_session_end for subagent session"); } - session.maybe_run_dream(). await; let telem = session.memory - .telemetry_snapshot(); session.emit_memory_session_summary(& telem, - total_chunks_at_end, session_end_result); if let Some(cancel) = & session - .sync_loop_cancel { cancel.cancel(); } session.feedback_manager - .shutdown(session.upload_queue.get()). await; if ! session.startup_hints - .is_subagent { session.persist_background_task_manifest(). await; } - cleanup_session_scratch(& session); return; } } } - } + biased; _ = & mut idle_flush_sleep, if session.idle_flush_timeout.is_some() + && session.memory.is_enabled() && ! session.memory.is_flushing + .load(std::sync::atomic::Ordering::Relaxed) => { let current_len = session + .chat_state_handle.get_conversation_len(). await; let last_len = session + .last_idle_flush_conversation_len.load(std::sync::atomic::Ordering::Relaxed); + if current_len > last_len { tracing::info!(target : + xai_grok_telemetry::memory_log::TARGET, + "MEMORY_IDLE_FLUSH: timer fired (conversation {last_len} → {current_len})"); + session.last_idle_flush_conversation_len.store(current_len, + std::sync::atomic::Ordering::Relaxed); tokio::task::spawn_local({ let session + = session.clone(); async move { if ! session.run_memory_flush("interval", + None). await { tracing::info!(target : + xai_grok_telemetry::memory_log::TARGET, + "MEMORY_IDLE_FLUSH: skipped — another flush already in progress"); } } }); + } else { tracing::debug!(target : xai_grok_telemetry::memory_log::TARGET, + "MEMORY_IDLE_FLUSH: skipped, no new messages since last flush (len={current_len})"); + } + if let Some(timeout) = session.idle_flush_timeout { idle_flush_sleep + .as_mut().reset(tokio::time::Instant::now() + timeout); } } _ = & mut + dream_check_sleep, if session.dream_check_timeout.is_some() && session.memory + .is_enabled() => { tracing::debug!(target : + xai_grok_telemetry::memory_log::TARGET, "MEMORY_DREAM_CHECK: timer fired"); + tokio::task::spawn_local({ let session = session.clone(); async move { + session.maybe_run_dream(). await; } }); if let Some(timeout) = session + .dream_check_timeout { dream_check_sleep.as_mut() + .reset(tokio::time::Instant::now() + timeout); } } changed = model_switch_rx + .changed() => { if changed.is_ok() { let new_gen = * model_switch_rx + .borrow_and_update(); session.handle_model_switch_for_laziness(new_gen). + await; } } event = chat_state_event_rx.recv() => { match event { + Some(xai_chat_state::ChatStateEvent::ConversationReset { new_len }) => { + session.last_idle_flush_conversation_len.store(new_len, + std::sync::atomic::Ordering::Relaxed); session.memory.context_injected + .store(false, std::sync::atomic::Ordering::Relaxed); } + Some(xai_chat_state::ChatStateEvent::ImageBudget { body_bytes, trigger_bytes, + reclaim_target_bytes, inline_images, needs_image_compaction, evicted, + body_bytes_after, }) => { + xai_grok_telemetry::unified_log::info("shell.image_budget", Some(session + .session_info.id.0.as_ref()), Some(serde_json::json!({ "body_bytes" : + body_bytes, "body_bytes_after" : body_bytes_after, "trigger_bytes" : + trigger_bytes, "reclaim_target_bytes" : reclaim_target_bytes, "inline_images" + : inline_images, "images_remaining" : inline_images.saturating_sub(evicted), + "needs_image_compaction" : needs_image_compaction, "evicted" : evicted, + })),); } Some(xai_chat_state::ChatStateEvent::PromptIndexChanged { .. }) | + Some(xai_chat_state::ChatStateEvent::TokensUpdated { .. }) => {} None => {} } + } maybe_event = event_rx.recv() => { if let Some(event) = maybe_event { match + event { SessionEvent::Notification(notification) => { let out = replay_buffer + .consume_chunk(notification); match out { None => {} Some((first, second)) => + { session.emit_buffered(first). await; if let Some(second) = second { session + .emit_buffered(second). await; } } } } SessionEvent::FlushReplay { respond_to + } => { if let Some(notification) = replay_buffer.flush() { session + .emit_buffered(notification). await; } + if let Some(tx) = respond_to { let _ = + tx.send(()); } } } } } maybe_completion = completion_rx.recv() => { let + Some((prompt_id, result)) = maybe_completion else { if let Some(cancel) = & + session.sync_loop_cancel { cancel.cancel(); } cleanup_session_scratch(& + session); return; }; if let Some(notification) = replay_buffer.flush() { + session.emit_buffered(notification). await; } let (turn_succeeded, + infra_pause_message) = SessionActor::post_turn_goal_degradation_plan(& + result); session.handle_completion(prompt_id, result). await; session + .drain_monitor_buffer_to_pending(). await; if let Some(message) = + infra_pause_message { session.apply_infra_pause_after_turn_err(message). + await; } session.handle_turn_end(turn_succeeded). await; if session + .flush_stranded_interjections(). await { + tracing::info!("Flushed stranded interjection(s) into prompt turns"); } + SessionActor::maybe_start_running_task(session.clone(), completion_tx + .clone()). await; SessionActor::maybe_drain_notifications(session.clone(), + completion_tx.clone()). await; session.emit_session_idle_if_idle(). await; { + let s = session.clone(); tokio::task::spawn_local(async move { s + .maybe_fire_laziness_check(). await; }); } } maybe_cmd = cmd_rx.recv() => { + let Some(cmd) = maybe_cmd else { let envelope = session + .fire_hook(xai_grok_hooks::event::HookEventName::SessionEnd, None, + xai_grok_hooks::event::HookPayload::SessionEnd { reason : "channel_closed" + .to_string(), turn_count : None, tool_call_count : None, },); if let + Some(registry) = session.hook_registry.borrow().clone() { let ctx = session + .hook_run_ctx(); let results = + xai_grok_hooks::dispatcher::dispatch_non_blocking(& registry, + xai_grok_hooks::event::HookEventName::SessionEnd, & envelope, & ctx,). await; + session.send_hook_execution("session_end", None, None, & results). await; } + session.dispatch_session_end_stop("channel_closed"). await; let mut + session_end_result = "disabled"; let mut total_chunks_at_end = 0usize; if ! + session.startup_hints.is_subagent { if let Some(storage) = session.memory + .storage() { let conversation = session.chat_state_handle.get_conversation(). + await; let result = crate ::session::memory::hooks::on_session_end(& storage, + & conversation, & session.session_info.id.0, session.memory.save_on_end,); + session_end_result = match & result { crate + ::session::memory::hooks::SessionEndResult::Written(_) => "written", crate + ::session::memory::hooks::SessionEndResult::Skipped => "skipped", crate + ::session::memory::hooks::SessionEndResult::Failed(_) => "failed", }; + total_chunks_at_end = storage.total_chunk_count(); let telem = session.memory + .telemetry_snapshot(); tracing::info!(target : + xai_grok_telemetry::memory_log::TARGET, result = ? result, tool_searches = + telem.tool_search_count, injection_searches = telem.injection_count, + recovery_searches = telem.compaction_recovery_count, + "MEMORY_SESSION_END: channel closed, session summary saved"); if let crate + ::session::memory::hooks::SessionEndResult::Written(ref path_str) = result { + session.reindex_and_embed(std::path::Path::new(path_str), "session"). await; + session.send_xai_notification(XaiSessionUpdate::MemorySessionSaved { path : + path_str.clone(), }). await; } } } else { tracing::debug!(target : + xai_grok_telemetry::memory_log::TARGET, + "MEMORY_SUBAGENT_SKIP: skipping on_session_end for subagent session"); } + session.maybe_run_dream(). await; let telem = session.memory + .telemetry_snapshot(); session.emit_memory_session_summary(& telem, + total_chunks_at_end, session_end_result); if let Some(notification) = + replay_buffer.flush() { session.emit_buffered(notification). await; } + { let + model_id = session.current_model_id(). await; if let Some(signals) = session + .signals_handle().snapshot(). await { + xai_grok_telemetry::session_ctx::log_event(xai_grok_telemetry::events::SessionEnded + { duration_secs : session.session_start.elapsed().as_secs(), turn_count : + signals.turn_count as u64, tool_call_count : signals.tool_call_count as u64, + compaction_count : signals.compaction_count as u64, model_id, },); } } + if let + Some(cancel) = & session.sync_loop_cancel { cancel.cancel(); } session + .feedback_manager.shutdown(session.upload_queue.get()). await; if ! session + .startup_hints.is_subagent { session.persist_background_task_manifest(). + await; } cleanup_session_scratch(& session); return; }; match cmd { + SessionCommand::Initialize { system_prompt } => { session + .initialize(system_prompt). await; let s = session.clone(); let handle = + tokio::task::spawn_local(async move { s.build_prefix_background(). await }); + session.deferred_prefix.arm(handle); } SessionCommand::ReplaceSystemPrompt { + system_prompt } => { session.handle_replace_system_prompt(system_prompt). + await; } SessionCommand::RestorePlanApproval => { let s = session.clone(); + let completion_tx = completion_tx.clone(); tokio::task::spawn_local(async + move { s.resume_plan_approval(completion_tx). await; }); } + SessionCommand::Prompt { prompt_id, prompt_blocks, prompt_mode, + artifact_upload_ctx, client_identifier, screen_mode, verbatim, traceparent, + json_schema, send_now, admission, respond_to, persist_ack, parsed_prompt_tx } + => { let origin = super::PromptOrigin::from_prompt_id(& prompt_id); let + (actor_admitted, task_wake_fallback) = match admission { Some(admission) => { + let fallback = session.admit_task_completion_wake(& origin, admission). + await; (fallback.is_some(), fallback) } None => (true, None), }; if ! + actor_admitted { SessionActor::respond_removed_prompt(respond_to); continue; + } session.ensure_prefix_ready(). await; if ! origin.is_synthetic() { if let + Some(gate) = & session.tool_context.task_wake_suppressed { gate.set(false); } + let mut state = session.state.lock(). await; state.notifications_suppressed = + false; xai_grok_telemetry::unified_log::info("shell.task_wake.gate_cleared", + Some(session.session_info.id.0.as_ref()), Some(serde_json::json!({ "reason" : + "user_intake" })),); session.user_input_generation.fetch_add(1, + std::sync::atomic::Ordering::AcqRel); } + if origin.is_synthetic() { let state + = session.state.lock(). await; let has_running = state.running_task + .is_some(); let queue_depth = state.pending_inputs.len(); drop(state); + tracing::info!(prompt_id = % prompt_id, has_running_task = has_running, + queue_depth = queue_depth, + "auto-wake: session actor received synthetic prompt"); } + if let Some(ref tp) + = traceparent { let meta = serde_json::json!({ "traceparent" : tp }); + xai_file_utils::trace_context::link_current_span_to_meta(& meta); } let + (trace_gcs_config, artifact_tracker) = match artifact_upload_ctx { Some(tu) + => (Some(tu.gcs_config), Some(tu.artifact_tracker)), None => (None, None), }; + let cancel_for_send_now = session.queue_input(prompt_blocks, prompt_id, + prompt_mode, trace_gcs_config, artifact_tracker, client_identifier, + screen_mode, verbatim, json_schema, send_now, task_wake_fallback, respond_to, + persist_ack, parsed_prompt_tx). await; if cancel_for_send_now { session + .cancel_turn_for_send_now(& mut replay_buffer). await; } + SessionActor::maybe_start_running_task(session.clone(), completion_tx + .clone()). await; } SessionCommand::SessionMode { session_mode, responds_to } + => { session.handle_session_mode(session_mode). await; let _ = responds_to + .send(()); } SessionCommand::SetSessionModel { sampling_config, use_concise, + apply_prompt_override, skip_prompt_rewrite, auto_compact_threshold_percent, + responds_to } => { let updated_model_id = session + .handle_set_session_model(sampling_config, use_concise, + apply_prompt_override, skip_prompt_rewrite, auto_compact_threshold_percent). + await; let _ = responds_to.send(updated_model_id); } + SessionCommand::RebuildAgentForDefinition { definition, responds_to } => { + let outcome = session.handle_rebuild_agent_for_definition(definition). await; + let _ = responds_to.send(outcome); } SessionCommand::OverrideModelName { + model_name, extra_headers, context_window } => { if let Some(mut cfg) = + session.chat_state_handle.get_sampling_config(). await { + tracing::info!(target : SESSION_LOG, session_id = % session.session_info.id, + old_model = % cfg.model, new_model = % model_name, extra_header_count = + extra_headers.len(), old_context_window = cfg.context_window.get(), + new_context_window = ? context_window.map(| cw | cw.get()), + "OVERRIDE_MODEL: changing model name in sampling config"); session + .signals_handle().set_primary_model(& model_name); cfg.model = model_name + .clone(); cfg.extra_headers.extend(extra_headers); if let Some(cw) = + context_window && session.compaction.context_window_override.is_none() { cfg + .context_window = cw; } session.chat_state_handle + .update_sampling_config(cfg); let existing = session.chat_state_handle + .get_credentials(). await; if let Some(r) = crate + ::agent::config::try_resolve_model_credentials(model_name.as_str(), existing + .api_key.as_deref()) { session.chat_state_handle + .update_credentials(xai_chat_state::Credentials { api_key : r.api_key, + auth_type : r.auth_type, alpha_test_key : existing.alpha_test_key, + client_version : existing.client_version, }); } session + .invalidate_model_auth_memo(); } } SessionCommand::GetCurrentModel { + responds_to } => { let model = session.chat_state_handle + .get_sampling_config(). await .map(| c | c.model).unwrap_or_default(); let _ + = responds_to.send(model); } SessionCommand::GetCurrentPromptMode { + responds_to } => { let mode = * session.current_prompt_mode.lock(); let _ = + responds_to.send(mode); } SessionCommand::GetModelMetadata { responds_to } => + { let id = session.chat_state_handle.get_last_model_metadata(). await; let _ + = responds_to.send(id); } SessionCommand::GetSessionInfo { responds_to } => { + let info = session.build_session_info(). await; let _ = responds_to + .send(info); } SessionCommand::BackgroundForegroundCommand { tool_call_id, + respond_to } => { let result = session.agent.borrow().tool_bridge() + .background_foreground_command(& tool_call_id). await; let _ = respond_to + .send(result); } SessionCommand::KillBackgroundTask { task_id, respond_to } + => { let result = session.agent.borrow().tool_bridge().kill_background_task(& + task_id). await .map_err(| e | e.to_string()); let _ = respond_to + .send(result); } SessionCommand::DeleteScheduledTask { task_id, respond_to } + => { let result = session.agent.borrow().tool_bridge() + .delete_scheduled_task(& task_id). await .map_err(| e | e.to_string()); let _ + = respond_to.send(result); } SessionCommand::ListTasks { respond_to } => { + let result = session.agent.borrow().tool_bridge().list_tasks(). await; let _ + = respond_to.send(result); } SessionCommand::GetHooksList { respond_to } => { + use crate ::extensions::hooks::hook_spec_to_info; let hooks = match &* + session.hook_registry.borrow() { Some(registry) => registry.all_hooks() + .iter().map(| spec | hook_spec_to_info(spec)).collect(), None => Vec::new(), + }; let project_trusted = crate + ::agent::folder_trust::project_scope_allowed(std::path::Path::new(& session + .session_info.cwd),); let _ = respond_to + .send(xai_hooks_plugins_types::HooksListResponse { hooks, project_trusted, + load_errors : session.hook_load_errors.borrow().clone(), }); } + SessionCommand::HooksAction { action, respond_to } => { let outcome = session + .handle_hooks_action(action). await; let _ = respond_to.send(outcome); } + SessionCommand::NotifyPluginUpdates { updates } => { session + .send_xai_notification(XaiSessionUpdate::PluginUpdatesInstalled { updates },) + . await; } SessionCommand::PluginsAction { action, respond_to } => { let + outcome = session.handle_plugins_action(action). await; let _ = respond_to + .send(outcome); } SessionCommand::PluginsList { respond_to } => { let _ = + respond_to.send(session.plugin_registry.borrow().clone()); } + SessionCommand::DispatchNotificationHook { notification_type, message, title, + level, } => { session.dispatch_notification_hook(& notification_type, + message, title, level,). await; } SessionCommand::DropMonitorNotifications { + task_id } => { { let mut state = session.state.lock(). await; state + .pending_notifications.retain(| n | { ! matches!(& n.source, + NotificationSource::MonitorEvent { task_id : tid } + if tid == & task_id) }); } + if let Some(buffer) = & session.tool_context.monitor_event_buffer { let + dropped = buffer.drain_matching(| e | e.task_id == task_id); if ! dropped + .is_empty() { tracing::debug!(task_id = % task_id, dropped = dropped.len(), + "dropped buffered monitor events after TaskCompleted auto-wake"); } } } + SessionCommand::InjectNotification { prompt_id, prompt_blocks, priority, + source } => { let is_turn_active = session.tool_context.is_turn_active + .as_ref().map(| f | f.load(std::sync::atomic::Ordering::Relaxed)) + .unwrap_or(false); if is_turn_active && priority == + NotificationPriority::Next { if let Some(buffer) = & session.tool_context + .monitor_event_buffer { let non_text_count = prompt_blocks.iter().filter(| b + | ! matches!(b, acp::ContentBlock::Text(_))).count(); if non_text_count > 0 { + tracing::debug!(non_text_count, + "Non-text content blocks dropped in mid-turn monitor event routing"); } let + event_text = prompt_blocks.iter().filter_map(| b | { if let + acp::ContentBlock::Text(t) = b { Some(t.text.clone()) } else { None } }) + .collect::< Vec < _ >> ().join("\n"); let task_id = source.task_id() + .to_owned(); const MAX_BUFFER_EVENTS : usize = 50; buffer + .push_capped(xai_grok_tools::implementations::grok_build::task::types::MonitorEventNotification + { task_id : task_id.clone(), event_text, owner_session_id : Some(session + .session_info.id.0.to_string(),), }, MAX_BUFFER_EVENTS,); + tracing::debug!(task_id = % task_id, + "Routed monitor event to mid-turn buffer"); } } else { { let mut state = + session.state.lock(). await; SessionActor::push_pending_notification(& mut + state, PendingNotification { prompt_id, prompt_blocks, priority, source, },); + } SessionActor::maybe_drain_notifications(session.clone(), completion_tx + .clone()). await; } } SessionCommand::RecordGoalTurnTaskIds { task_ids } => { + session.record_reparented_goal_turn_task_ids(task_ids); } + SessionCommand::RemoveQueuedPrompt { id, expected_version, owner } => { + session.handle_remove_queued_prompt(& id, expected_version, owner.as_deref()) + . await; } SessionCommand::ReorderQueue { ordered_ids } => { session + .handle_reorder_queue(& ordered_ids). await; } SessionCommand::ClearQueue { + owner } => { session.handle_clear_queue(owner.as_deref()). await; } + SessionCommand::EditQueuedPrompt { id, new_text, editor } => { session + .handle_edit_queued_prompt(& id, new_text, editor.as_deref()). await; } + SessionCommand::InterjectQueuedPrompt { id, expected_version, owner, new_text + } => { let cancel_for_send_now = session.handle_interject_queued_prompt(& id, + expected_version, owner.as_deref(), new_text.as_deref()). await; if + cancel_for_send_now { session.cancel_turn_for_send_now(& mut replay_buffer). + await; } SessionActor::maybe_start_running_task(session.clone(), + completion_tx.clone()). await; } SessionCommand::Cancel { cancel_subagents, + kill_background_tasks, rewind_if_pristine, trigger, } => { if let + Some(notification) = replay_buffer.flush() { session + .emit_buffered(notification). await; } session.pending_interjections.clear(); + let suppress_task_wakes = trigger.as_deref() == Some("ctrl_c"); session + .cancel_running_task(cancel_subagents, kill_background_tasks, + rewind_if_pristine, trigger,). await; session.auto_pause_goal_if_active(crate + ::session::goal_tracker::GoalPauseReason::User,). await; + SessionActor::maybe_start_running_task(session.clone(), completion_tx + .clone()). await; if ! suppress_task_wakes { + SessionActor::maybe_drain_notifications(session.clone(), completion_tx + .clone(),). await; } } SessionCommand::CompactSession { user_context, + respond_to } => { let s = session.clone(); tokio::task::spawn_local(async + move { let compact_session = s.run_compact(user_context). await; let _ = + respond_to.send(compact_session); }); } SessionCommand::ReloadPlugins { + registry } => { if ! session.startup_hints.is_subagent { let registry = + session.preserve_session_plugin_dirs(registry); session + .apply_plugin_registry_snapshot(registry). await; } } + SessionCommand::ReloadHooks => { if ! session.startup_hints.is_subagent { let + _ = session.reload_hooks_impl(). await; } } + SessionCommand::RefreshSkillBaseline => { let s = session.clone(); + tokio::task::spawn_local(async move { let cwd = s.tool_context.cwd.as_path() + .to_string_lossy(); let skills_config = crate ::util::config::load_config(). + await .skills; let pr = s.plugin_registry.borrow().clone(); let new_skills = + xai_grok_agent::prompt::skills::list_skills_with_plugins(Some(& cwd), & + skills_config, pr.as_deref(), s.rebuild_spec.compat,). await; + tracing::info!(skills = new_skills.len(), + "refreshed skill baseline after bundle sync"); let bridge = s.agent.borrow() + .tool_bridge().clone(); bridge.update_skill_baseline(new_skills). await; if + let Some(effects) = bridge.apply_pending_skill_update(). await { s + .apply_skill_update_effects(effects). await; } }); } + SessionCommand::FlushMemory { respond_to } => { let s = session.clone(); + tokio::task::spawn_local(async move { if s.memory.is_enabled() { let + did_flush = s.run_memory_flush("user_requested", None). await; let _ = + respond_to.send(Ok(did_flush)); } else { let _ = respond_to + .send(Err(acp::Error::invalid_request() + .data("memory is not enabled for this session".to_string()))); } }); } + SessionCommand::SetYoloMode { enabled } => { let was = session.permissions + .is_yolo_mode(); tracing::info!("Session received SetYoloMode: {}", enabled); + session.permissions.set_yolo_mode(enabled); let actual = session.permissions + .is_yolo_mode(); if let Some(enabled) = yolo_toggle_report(was, actual) { + session.emit_event(crate ::session::events::Event::YoloToggled { enabled }); + } } SessionCommand::SetAutoMode { enabled } => { let enabled = enabled && + crate ::util::config::auto_permission_mode_enabled_from_disk(); + tracing::info!("Session received SetAutoMode: {}", enabled); session + .permissions.set_auto_mode(enabled); if enabled { session + .wire_permission_auto_llm_classifier(). await; } else { session.permissions + .set_llm_side_query_wired(false); } } SessionCommand::ResetPermissionState => + { session.permissions.reset_state(); tracing::info!(session_id = % session + .session_info.id, "Permission state reset via notification"); } + SessionCommand::Rewind { request, respond_to } => { let s = session.clone(); + tokio::task::spawn_local(async move { let result = s.handle_rewind(request). + await; let _ = respond_to.send(result); }); } SessionCommand::RepairHistory { + dry_run, respond_to } => { let s = session.clone(); + tokio::task::spawn_local(async move { let result = s + .handle_repair_history(dry_run). await; let _ = respond_to.send(result); }); + } SessionCommand::GetRewindPoints { respond_to } => { let response = session + .get_rewind_points(). await; let _ = respond_to.send(response); } + SessionCommand::GetRewindFileCounts { respond_to } => { let _ = respond_to + .send(session.rewind_file_counts(). await); } + SessionCommand::ReconcileRewindTracker { target_prompt_index } => { session + .merge_rewind_tracker_from(target_prompt_index). await; } + SessionCommand::XaiSessionNotification { notification } => { session + .handle_xai_session_notification(notification). await; } + SessionCommand::RecordSubagentUsage { by_model, parent_prompt_id, incomplete, + respond_to, } => { use super::updates::SubagentUsageApply; match session + .record_subagent_usage(& by_model, parent_prompt_id.as_deref(), incomplete,). + await { Ok(SubagentUsageApply::AttributedToPrompt) => { let _ = respond_to + .send(()); } Ok(SubagentUsageApply::SessionOnly) => { let _ = session + .mark_subagent_usage_not_applied(parent_prompt_id.as_deref(),). await; let _ + = respond_to.send(()); } Err(()) => {} } } + SessionCommand::MarkSubagentUsageNotApplied { parent_prompt_id, respond_to, } + => { if session.mark_apply_miss_incomplete(parent_prompt_id.as_deref()). + await { let _ = respond_to.send(()); } } + SessionCommand::ErrorPathUsageFallback { prompt_id, respond_to, } => { let + pid = prompt_id.or_else(|| { session.current_prompt_id.lock().ok().and_then(| + g | g.clone()) }); let usage = match pid.as_deref() { Some(id) => session + .error_path_usage_fallback(id). await, None => { match session + .chat_state_handle.try_get_prompt_usage(). await { Ok(ledger) => { crate + ::extensions::notification::PromptUsage::for_error_path(ledger.as_ref(), + false,) } Err(()) => { crate + ::extensions::notification::PromptUsage::for_error_path(None, true,) } } } }; + let _ = respond_to.send(usage); } SessionCommand::SetNextTraceTurn { + next_trace_turn, request_id, } => { let _ = session.notifications + .persistence_tx.send(PersistenceMsg::NextTraceTurn { next_trace_turn, + request_id, }); } SessionCommand::CopyFile { respond_to } => { if let + Some(notification) = replay_buffer.flush() { session + .emit_buffered(notification). await; } let _ = session.notifications + .persistence_tx.send(PersistenceMsg::CopyFile { one_shot : respond_to }); } + SessionCommand::IsBusy { respond_to } => { let busy = { let state = session + .state.lock(). await; state_is_busy(& state) }; let _ = respond_to + .send(busy); } SessionCommand::FlushComplete { respond_to } => { if let + Some(notification) = replay_buffer.flush() { session + .emit_buffered(notification). await; } let _ = session.notifications + .persistence_tx.send(PersistenceMsg::FlushAndAck { respond_to }); } + SessionCommand::UpdateMcpServers { mcp_servers, respond_to } => { if session + .startup_hints.is_subagent { tracing::debug!(session_id = % session + .session_info.id.0, "Skipping UpdateMcpServers for subagent session",); let _ + = respond_to.send(Ok(())); continue; } + tracing::info!("Updating MCP servers for session '{}' ({} servers)", session + .session_info.id.0, mcp_servers.len()); session.reseed_mcp_output_cap(). + await; let (diff, dispatch_event_tx) = { let mut mcp_state = session + .mcp_state.lock(). await; let diff = mcp_state + .update_configs_diff(mcp_servers); let tx = mcp_state.client_event_tx(); + (diff, tx) }; let Some(diff) = diff else { + tracing::debug!("MCP configs unchanged for session '{}', skipping re-initialization", + session.session_info.id.0); let _ = respond_to.send(Ok(())); continue; }; if + (! diff.added.is_empty() || ! diff.removed.is_empty()) && let Some(tx) = & + dispatch_event_tx { let _ = tx + .send(xai_grok_mcp::servers::McpClientEvent::ConfigDiff { added : diff.added + .clone(), removed : diff.removed.clone(), },); } for name in & diff.removed { + let prefix = format!("{}{}", name, crate + ::session::mcp_servers::MCP_TOOL_NAME_DELIMITER); let removed_count = session + .agent.borrow().tool_bridge().unregister_tools_by_prefix(& prefix); + tracing::info!(server = name.as_str(), tools_removed = removed_count, + "Unregistered tools for removed MCP server"); } let session_for_mcp = session + .clone(); tokio::task::spawn_local(async move { session_for_mcp + .ensure_mcp_tools_initialized(). await; let _ = respond_to.send(Ok(())); }); + } SessionCommand::ToggleMcpServer { server_name, enabled, server_config, + respond_to } => { session.events + .emit(xai_file_utils::events::Event::McpServerToggled { server_name : + server_name.clone(), enabled, }); let mut mcp_state = session.mcp_state + .lock(). await; let mut configs = mcp_state.configs.clone(); if enabled { if + let Some(config) = server_config { configs.retain(| c | { crate + ::session::mcp_servers::mcp_server_name(c) != server_name }); configs + .push(config); } else { let already_present = configs.iter().any(| c | { + crate ::session::mcp_servers::mcp_server_name(c) == server_name }); if + already_present { drop(mcp_state); let _ = respond_to.send(Ok(())); continue; + } drop(mcp_state); let _ = respond_to.send(Err(acp::Error::invalid_params() + .data(format!("server '{}' not found in config", server_name)))); continue; } + } else { configs.retain(| c | crate + ::session::mcp_servers::mcp_server_name(c) != server_name); } let diff = + mcp_state.update_configs_diff(configs); let dispatch_event_tx = mcp_state + .client_event_tx(); drop(mcp_state); let Some(diff) = diff else { let _ = + respond_to.send(Ok(())); continue; }; if (! diff.added.is_empty() || ! diff + .removed.is_empty()) && let Some(tx) = & dispatch_event_tx { let _ = tx + .send(xai_grok_mcp::servers::McpClientEvent::ConfigDiff { added : diff.added + .clone(), removed : diff.removed.clone(), },); } for name in & diff.removed { + let prefix = format!("{}{}", name, crate + ::session::mcp_servers::MCP_TOOL_NAME_DELIMITER); let removed_count = session + .agent.borrow().tool_bridge().unregister_tools_by_prefix(& prefix); + tracing::info!(server = name.as_str(), tools_removed = removed_count, + "Unregistered tools for toggled MCP server"); } let session_for_mcp = session + .clone(); let sname = server_name.clone(); tokio::task::spawn_local(async + move { session_for_mcp.ensure_mcp_tools_initialized(). await; if let Err(e) = + crate ::util::config::save_mcp_server_enabled(& sname, enabled,). await { + tracing::warn!(server = sname.as_str(), error = % e, + "Failed to persist server enabled state to config"); } let _ = respond_to + .send(Ok(())); }); } SessionCommand::ToggleMcpTool { server_name, tool_name, + enabled, is_managed_gateway, respond_to } => { if is_managed_gateway { let + mut disabled_tools = crate + ::util::config::get_all_mcp_disabled_tools(std::path::Path::new(& session + .session_info.cwd)); if tool_name.is_empty() { let set = disabled_tools + .entry(crate ::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY + .to_string()).or_default(); if enabled { set.remove(& server_name); } else { + set.insert(server_name.clone()); } + if set.is_empty() { disabled_tools + .remove(crate ::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY); } } + else if enabled { if let Some(set) = disabled_tools.get_mut(& server_name) { + set.remove(& tool_name); if set.is_empty() { disabled_tools.remove(& + server_name); } } } else { disabled_tools.entry(server_name.clone()) + .or_default().insert(tool_name.clone()); } session + .refresh_mcp_snapshot_and_schedule_reminder_with_disabled(& disabled_tools,). + await; session.refresh_goal_harness_enabled(). await; let disabled_vec : Vec + < String > = if tool_name.is_empty() { disabled_tools.get(crate + ::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY).map(| s | s.iter() + .cloned().collect()).unwrap_or_default() } else { disabled_tools.get(& + server_name).map(| s | s.iter().cloned().collect()).unwrap_or_default() }; + let notifications = session.notifications.gateway.clone(); let session_id = + session.session_info.id.0.clone(); let server_for_persist = if tool_name + .is_empty() { crate ::util::config::MANAGED_GATEWAY_DISABLED_CONNECTORS_KEY + .to_string() } else { server_name.clone() }; tokio::task::spawn_local(async + move { if let Err(e) = crate ::util::config::save_mcp_disabled_tools(& + server_for_persist, & disabled_vec,). await { tracing::warn!(server = + server_for_persist.as_str(), error = % e, + "Failed to persist disabled_tools to config"); } let payload = crate + ::extensions::mcp::McpToolsChanged { session_id : session_id.to_string(), + server_name : String::new(), tools : Vec::new(), }; if let Ok(params) = + serde_json::value::to_raw_value(& payload) { notifications + .forward_fire_and_forget(acp::ExtNotification::new("x.ai/mcp/tools_changed", + params.into())); } let _ = respond_to.send(Ok(())); }); continue; } let + qualified = format!("{}{}{}", server_name, crate + ::session::mcp_servers::MCP_TOOL_NAME_DELIMITER, tool_name,); let mut + mcp_state = session.mcp_state.lock(). await; if enabled { if let Some(set) = + mcp_state.disabled_tools.get_mut(& server_name) { set.remove(& tool_name); if + set.is_empty() { mcp_state.disabled_tools.remove(& server_name); } } + if let + Some(reg) = mcp_state.disabled_tool_registrations.remove(& qualified) && reg + .model_visible { let bridge = session.agent.borrow().tool_bridge().clone(); + if let Err(e) = bridge.register_mcp_tools(reg.name, reg.tool, Some(reg + .input_schema)). await { tracing::warn!(tool = qualified.as_str(), error = % + e, "Failed to re-register toggled MCP tool"); } } } else { let bridge = + session.agent.borrow().tool_bridge().clone(); let tool_def = bridge + .tool_definitions(). await .into_iter().find(| d | d.function.name == + qualified); if let Some(def) = tool_def { let meta = mcp_state.mcp_tool_meta + .get(& qualified).cloned(); let schema = def.function.parameters.clone(); let + mcp_tool = crate ::session::mcp_servers::McpTool::new(tool_name.clone(), def + .function.description.clone().unwrap_or_default(), server_name.clone(), + session.mcp_state.clone(), schema, meta,); if let Some(reg) = mcp_tool + .into_registration() { mcp_state.disabled_tool_registrations.insert(qualified + .clone(), reg); } } bridge.unregister_tool_by_name(& qualified); mcp_state + .disabled_tools.entry(server_name.clone()).or_default().insert(tool_name + .clone()); } let disabled_vec : Vec < String > = mcp_state.disabled_tools + .get(& server_name).map(| s | s.iter().cloned().collect()) + .unwrap_or_default(); drop(mcp_state); session + .refresh_mcp_snapshot_and_schedule_reminder(). await; session + .refresh_goal_harness_enabled(). await; let notifications = session + .notifications.gateway.clone(); let session_id = session.session_info.id.0 + .clone(); let server_for_persist = server_name.clone(); + tokio::task::spawn_local(async move { if let Err(e) = crate + ::util::config::save_mcp_disabled_tools(& server_for_persist, & + disabled_vec,). await { tracing::warn!(server = server_for_persist.as_str(), + error = % e, "Failed to persist disabled_tools to config"); } let payload = + crate ::extensions::mcp::McpToolsChanged { session_id : session_id + .to_string(), server_name : String::new(), tools : Vec::new(), }; if let + Ok(params) = serde_json::value::to_raw_value(& payload) { notifications + .forward_fire_and_forget(acp::ExtNotification::new(crate + ::extensions::mcp::mcp_methods::TOOLS_CHANGED, params.into())); } let _ = + respond_to.send(Ok(())); }); } SessionCommand::SnapshotMcpPool { respond_to } + => { let mcp_state = session.mcp_state.lock(). await; let pool = if mcp_state + .owned_clients.is_empty() && mcp_state.shared_clients.is_empty() { None } + else { Some(crate ::session::mcp_servers::SharedMcpPool::from_state(& + mcp_state)) }; let _ = respond_to.send(pool); } + SessionCommand::SnapshotClientHooks { respond_to } => { let _ = respond_to + .send(session.client_hooks.borrow().clone()); } + SessionCommand::SnapshotToolDefinitions { respond_to } => { let defs = + session.prepare_tool_definitions_inner(). await; let specs = session + .turn_base_tool_specs(& defs); let _ = respond_to.send(specs); } + SessionCommand::SetClientHooks { hooks } => { * session.client_hooks + .borrow_mut() = hooks; } SessionCommand::GetMcpStatus { respond_to } => { let + mcp_state = session.mcp_state.clone(); let tool_bridge = session.agent + .borrow().tool_bridge().clone(); let writer = session.events.writer(); + tokio::task::spawn_local(async move { let snapshot = crate + ::extensions::mcp::build_mcp_status(& mcp_state, & tool_bridge, Some(& + writer),). await; let _ = respond_to.send(snapshot); }); } + SessionCommand::CallMcpTool { server_name, server_url, tool_name, arguments, + respond_to } => { let mcp_state = session.mcp_state.clone(); + tokio::task::spawn_local(async move { let result = crate + ::extensions::mcp::call_mcp_tool(& mcp_state, & server_name, server_url + .as_deref(), & tool_name, arguments,). await; let _ = respond_to + .send(result); }); } SessionCommand::ReadMcpResource { server_name, uri, + respond_to } => { let mcp_state = session.mcp_state.clone(); + tokio::task::spawn_local(async move { let result = crate + ::extensions::mcp::read_mcp_resource(& mcp_state, & server_name, & uri,). + await; let _ = respond_to.send(result); }); } SessionCommand::McpAuthStatus { + respond_to } => { let mcp_state = session.mcp_state.clone(); + tokio::task::spawn_local(async move { let state = mcp_state.lock(). await; + let entries : Vec < _ > = state.auth_required.iter().map(| name | { crate + ::extensions::mcp::McpAuthStatusEntry { server_name : name.clone(), status : + "needs_auth", } }).collect(); let _ = respond_to.send(entries); }); } + SessionCommand::McpAuthTrigger { server_name, respond_to } => { let s = + session.clone(); tokio::task::spawn_local(async move { let result = s + .handle_mcp_auth_trigger(& server_name). await; let _ = respond_to + .send(result); }); } SessionCommand::GetManagedGatewayDisabledTools { + respond_to } => { let disabled_tools = crate + ::util::config::get_all_mcp_disabled_tools(std::path::Path::new(& session + .session_info.cwd),); let _ = respond_to.send(disabled_tools); } + SessionCommand::RetryAuthRequiredServers { respond_to } => { let s = session + .clone(); tokio::task::spawn_local(async move { s + .retry_auth_required_servers(). await; let _ = respond_to.send(()); }); } + SessionCommand::RefreshMcpSearchIndex => { session + .refresh_mcp_snapshot_and_schedule_reminder(). await; } + SessionCommand::TriggerTestFeedback { tier, mode, respond_to } => { let s = + session.clone(); tokio::task::spawn_local(async move { let request = s + .feedback_manager.force_feedback_request(tier, mode). await; let notification + = crate ::extensions::notification::FeedbackRequestNotification::from(request + .clone()); s.send_feedback_notification(request). await; let resp = + ExtMethodResult::success(notification).to_ext_response(); let _ = respond_to + .send(resp); }); } SessionCommand::PersistFeedback(entry) => { let _ = + session.notifications.persistence_tx.send(PersistenceMsg::Feedback(* entry)); + } SessionCommand::AdvertiseCommands => { session + .send_available_commands_update(). await; } SessionCommand::ReloadSkills => { + let s = session.clone(); tokio::task::spawn_local(async move { s + .reload_skills_from_disk(). await; }); } + SessionCommand::DispatchSessionStartHook { source } => { let envelope = + session.fire_hook(xai_grok_hooks::event::HookEventName::SessionStart, None, + xai_grok_hooks::event::HookPayload::SessionStart { source, model_id : None, + agent_type : None, },); if let Some(registry) = session.hook_registry + .borrow().clone() { let ctx = session.hook_run_ctx(); let results = + xai_grok_hooks::dispatcher::dispatch_non_blocking(& registry, + xai_grok_hooks::event::HookEventName::SessionStart, & envelope, & ctx,). + await; session.send_hook_execution("session_start", None, None, & results). + await; } } SessionCommand::GetFeedbackContext { turn_number, responds_to } => + { let s = session.clone(); tokio::task::spawn_local(async move { use + prod_mc_cli_chat_proxy_types::feedback_types::FeedbackToolOutcome; let + turn_idx = turn_number.and_then(| n | usize::try_from(n).ok()); let + (last_user_message, last_assistant_message) = match turn_idx { Some(n) => { + let conv = s.chat_state_handle.get_conversation(). await; + turn_texts_for_feedback(& conv, n) } None => { tokio::join!(s + .chat_state_handle.get_last_user_query_text(), s.chat_state_handle + .get_last_assistant_text(),) } }; let sh = s.signals_handle(); let (signals, + tool_outcomes) = tokio::join!(sh.snapshot(), sh.last_turn_tool_outcomes(),); + let signals = signals.unwrap_or_default(); let ctx = FeedbackContext { + last_user_message, last_assistant_message, tool_outcomes : tool_outcomes + .into_iter().map(| o | FeedbackToolOutcome { tool_name : o.tool_name, calls : + o.successes + o.failures, failures : o.failures, }).collect(), + compaction_count : signals.compaction_count as i64, context_window_usage : + signals.context_window_usage, context_tokens_used : signals + .context_tokens_used, context_window_tokens : signals.context_window_tokens, + session_cwd : s.tool_context.cwd.as_path().to_string_lossy().to_string(), }; + let _ = responds_to.send(ctx); }); } SessionCommand::GetActiveAgent { + responds_to } => { let agent_type = session.active_agent_type.lock().clone(); + let _ = responds_to.send(agent_type); } SessionCommand::SideQuestion { + question, respond_to } => { let s = session.clone(); + tokio::task::spawn_local(async move { let result = s.handle_side_question(& + question). await; let _ = respond_to.send(result); }); } + SessionCommand::Recap { auto } => { let s = session.clone(); + tokio::task::spawn_local(async move { s.handle_recap(auto). await; }); } + SessionCommand::AISuggest { prefix, cwd, model_override, respond_to } => { + let s = session.clone(); tokio::task::spawn_local(async move { let result = s + .handle_ai_suggest(& prefix, & cwd, model_override.as_deref()). await; let _ + = respond_to.send(result); }); } SessionCommand::SuggestPrompt { + model_override, respond_to } => { let s = session.clone(); + tokio::task::spawn_local(async move { let result = s + .handle_suggest_prompt(model_override.as_deref()). await; let _ = respond_to + .send(result); }); } SessionCommand::RewriteMemoryNote { raw_text, + context_summary, respond_to } => { let s = session.clone(); + tokio::task::spawn_local(async move { let result = s + .handle_rewrite_memory_note(& raw_text, & context_summary). await; let _ = + respond_to.send(result); }); } SessionCommand::Interject { text, id, images } + => { session.broadcast_interjection(& text, id.as_deref()); session.events + .emit(crate ::session::events::Event::Interjected { source : crate + ::session::events::InterjectionSource::Direct, image_count : images.len() as + u32, redirect_kind : crate ::session::events::RedirectKind::Interjection, }); + let turn_running = session.current_prompt_id.lock().ok().and_then(| g | g + .clone()).is_some(); if turn_running { session.pending_interjections + .push(PendingInterjection { text, attachments : images, }); + tracing::info!("Queued mid-turn interjection"); } else { session + .queue_interjection_fallback_prompt(text, images, true). await; + SessionActor::maybe_start_running_task(session.clone(), completion_tx + .clone(),). await; } } SessionCommand::GoalSummaryTurn { prompt_text } => { + let prompt_id = format!("goal-summary-{}", uuid::Uuid::now_v7()); let + prompt_blocks = + vec![acp::ContentBlock::Text(acp::TextContent::new(prompt_text))]; let + (respond_to, _) = tokio::sync::oneshot::channel(); { let mut state = session + .state.lock(). await; state.pending_inputs.push_back(InputItem { prompt_id, + prompt_blocks, prompt_mode : crate ::session::plan_mode::PromptMode::Agent, + trace_gcs_config : None, artifact_tracker : None, client_identifier : None, + screen_mode : None, verbatim : true, json_schema : None, origin : + super::PromptOrigin::GoalSummary, task_wake_fallback : None, respond_to, + persist_ack : None, parsed_prompt_tx : None, queue_meta : None, send_now : + false, }); } SessionActor::maybe_start_running_task(session.clone(), + completion_tx.clone()). await; } SessionCommand::TakeTurnMessages { + respond_to } => { let result = session.chat_state_handle.take_turn_messages() + . await; let _ = respond_to.send(result); } + SessionCommand::TakeHarnessTraceTurns { respond_to } => { let result = + session.chat_state_handle.take_harness_trace_turns(). await; let _ = + respond_to.send(result); } SessionCommand::TakeStreamingCapture { prompt_id, + respond_to } => { let taken = { let mut cap = session.streaming_turn_capture + .lock(); if cap.prompt_id.as_deref() == Some(prompt_id.as_str()) { + Some(std::mem::take(& mut * cap)) } else { if ! cap.is_empty() { + tracing::warn!(requested_prompt_id = % prompt_id, slot_prompt_id = ? cap + .prompt_id, + "streaming_capture race: live slot belongs to a different prompt; \ + dropping streaming_partial.json for the requested turn",); + } None } }; let result = taken.and_then(| mut cap | { cap + .finalize_for_upload(); (! cap.is_empty()).then_some(cap) }); let _ = + respond_to.send(result); } SessionCommand::PersistGitHead { commit, branch } + => { let _ = session.notifications.persistence_tx + .send(PersistenceMsg::GitHead { commit, branch },); } SessionCommand::Shutdown => { if let Some(notification) = replay_buffer + .flush() { session.emit_buffered(notification). await; } session + .drop_pending_synthetic_items(). await; let envelope = session + .fire_hook(xai_grok_hooks::event::HookEventName::SessionEnd, None, + xai_grok_hooks::event::HookPayload::SessionEnd { reason : "shutdown" + .to_string(), turn_count : None, tool_call_count : None, },); if let + Some(registry) = session.hook_registry.borrow().clone() { let ctx = session + .hook_run_ctx(); let results = + xai_grok_hooks::dispatcher::dispatch_non_blocking(& registry, + xai_grok_hooks::event::HookEventName::SessionEnd, & envelope, & ctx,). await; + session.send_hook_execution("session_end", None, None, & results). await; } + session.dispatch_session_end_stop("shutdown"). await; let mut + session_end_result = "disabled"; let mut total_chunks_at_end = 0usize; if ! + session.startup_hints.is_subagent { if let Some(storage) = session.memory + .storage() { let conversation = session.chat_state_handle.get_conversation(). + await; let result = crate ::session::memory::hooks::on_session_end(& storage, + & conversation, & session.session_info.id.0, session.memory.save_on_end,); + session_end_result = match & result { crate + ::session::memory::hooks::SessionEndResult::Written(_) => "written", crate + ::session::memory::hooks::SessionEndResult::Skipped => "skipped", crate + ::session::memory::hooks::SessionEndResult::Failed(_) => "failed", }; + total_chunks_at_end = storage.total_chunk_count(); let telem = session.memory + .telemetry_snapshot(); tracing::info!(target : + xai_grok_telemetry::memory_log::TARGET, result = ? result, tool_searches = + telem.tool_search_count, injection_searches = telem.injection_count, + recovery_searches = telem.compaction_recovery_count, + "MEMORY_SESSION_END: session summary saved"); if let crate + ::session::memory::hooks::SessionEndResult::Written(ref path_str) = result { + session.reindex_and_embed(std::path::Path::new(path_str), "session"). await; + session.send_xai_notification(XaiSessionUpdate::MemorySessionSaved { path : + path_str.clone(), }). await; } } } else { tracing::debug!(target : + xai_grok_telemetry::memory_log::TARGET, + "MEMORY_SUBAGENT_SKIP: skipping on_session_end for subagent session"); } + session.maybe_run_dream(). await; let telem = session.memory + .telemetry_snapshot(); session.emit_memory_session_summary(& telem, + total_chunks_at_end, session_end_result); if let Some(cancel) = & session + .sync_loop_cancel { cancel.cancel(); } session.feedback_manager + .shutdown(session.upload_queue.get()). await; if ! session.startup_hints + .is_subagent { session.persist_background_task_manifest(). await; } + cleanup_session_scratch(& session); return; } } } + } } } /// Extract the user query text and assistant response text for the diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/sampler_turn.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/sampler_turn.rs index 710dfae..6ff5b57 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/sampler_turn.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/sampler_turn.rs @@ -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 { + 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, + ) { 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(¤t_model_id) { + self.refresh_provider_token_pre_turn( + &provider, + current_key.as_deref(), + ¤t_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) { diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/spawn.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/spawn.rs index c12365b..68cbb43 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/spawn.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/spawn.rs @@ -1151,7 +1151,7 @@ pub(crate) async fn spawn_session_actor( let session = Arc::new_cyclic(|weak: &std::sync::Weak| 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, diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/stop_gate.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/stop_gate.rs index c15b23c..eb2de0d 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/stop_gate.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/stop_gate.rs @@ -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] ); diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/types.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/types.rs index 50702e8..a61385c 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/types.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/types.rs @@ -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, } diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs index ebab4b9..0875209 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs @@ -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, 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 = + 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 = + 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; +} diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/cancel_running_task_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/cancel_running_task_tests.rs index 296e396..b19a293 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/cancel_running_task_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/cancel_running_task_tests.rs @@ -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() } diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/goal/goal_classifier_e2e_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/goal/goal_classifier_e2e_tests.rs index 6494f52..1827f01 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/goal/goal_classifier_e2e_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/goal/goal_classifier_e2e_tests.rs @@ -2714,6 +2714,7 @@ fn catalog_with( info, api_key: None, env_key: None, + auth_provider: None, api_base_url: None, }, ); diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/idle_resume_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/idle_resume_tests.rs index cebe291..b1539b7 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/idle_resume_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/idle_resume_tests.rs @@ -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( diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs index d0250f6..e2b4e40 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs @@ -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( diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/laziness/laziness_integration_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/laziness/laziness_integration_tests.rs index 35c0e4c..e3092b0 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/laziness/laziness_integration_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/laziness/laziness_integration_tests.rs @@ -38,6 +38,7 @@ fn detector_entry( info, api_key: None, env_key: None, + auth_provider: None, api_base_url: None, } } diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/memory_config_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/memory_config_tests.rs index 6be2fee..a17231c 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/memory_config_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/memory_config_tests.rs @@ -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, diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/record_response_token_usage_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/record_response_token_usage_tests.rs index b35d95c..1ec5621 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/record_response_token_usage_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/record_response_token_usage_tests.rs @@ -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; diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs index a15dd0d..13491dc 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs @@ -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, diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/support.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/support.rs index ecb1cdc..96adb56 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/support.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/support.rs @@ -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, diff --git a/crates/codegen/xai-grok-shell/src/session/compaction.rs b/crates/codegen/xai-grok-shell/src/session/compaction.rs index e1300fb..b868e29 100644 --- a/crates/codegen/xai-grok-shell/src/session/compaction.rs +++ b/crates/codegen/xai-grok-shell/src/session/compaction.rs @@ -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, diff --git a/crates/codegen/xai-grok-shell/src/session/persistence.rs b/crates/codegen/xai-grok-shell/src/session/persistence.rs index 1220f70..5a86ed6 100644 --- a/crates/codegen/xai-grok-shell/src/session/persistence.rs +++ b/crates/codegen/xai-grok-shell/src/session/persistence.rs @@ -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>, + respond_to: + tokio::sync::oneshot::Sender>, }, ContentChunk(PersistenceContentChunk), Chat(ConversationItem), @@ -1369,25 +1345,88 @@ mod generated_title_tests { pub struct PersistenceHandle { pub tx: mpsc::UnboundedSender, - /// 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 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, notification: acp::SessionNotification, result: Result<(), crate::session::storage::AppendUpdateError>, - ) -> Result { + ) -> 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) => { diff --git a/crates/codegen/xai-grok-shell/src/session/persistence_tests.rs b/crates/codegen/xai-grok-shell/src/session/persistence_tests.rs index a6d15a8..573a7cf 100644 --- a/crates/codegen/xai-grok-shell/src/session/persistence_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/persistence_tests.rs @@ -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) -> ActorGuard { + test_actor_with_remote_sync(info, storage, None) +} + +fn test_actor_with_remote_sync( + info: Info, + storage: Arc, + remote_sync: Option, +) -> 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) -> 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 { + 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); diff --git a/crates/codegen/xai-grok-shell/src/session/storage/jsonl/durable_tests.rs b/crates/codegen/xai-grok-shell/src/session/storage/jsonl/durable_tests.rs index afb6f8a..21d14a4 100644 --- a/crates/codegen/xai-grok-shell/src/session/storage/jsonl/durable_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/storage/jsonl/durable_tests.rs @@ -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::) - .collect::, _>>() - .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); } diff --git a/crates/codegen/xai-grok-shell/src/session/storage/jsonl/mod.rs b/crates/codegen/xai-grok-shell/src/session/storage/jsonl/mod.rs index 10cb10a..624249f 100644 --- a/crates/codegen/xai-grok-shell/src/session/storage/jsonl/mod.rs +++ b/crates/codegen/xai-grok-shell/src/session/storage/jsonl/mod.rs @@ -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>, } +#[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) -> io::Result<()> { + Self::append_jsonl_line_blocking(path, line, AppendDurability::Buffered).await + } + async fn append_jsonl_line_blocking( + path: PathBuf, + line: Vec, + 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) -> io::Result<()> { - Self::append_jsonl_line_locked(path, line, AppendDurability::Buffered).await - } - async fn append_jsonl_line_locked( - path: PathBuf, - line: Vec, - 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, 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, 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 `.jsonl.lock`. + /// Full-file [`Self::write_jsonl`] atomic-rename rewrites bypass this append-only lock. + fn lock_append(path: &Path) -> io::Result { + 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( + async fn append_update_with_bookkeeping( &self, info: &Info, update: &super::SessionUpdate, - bookkeeping: F, - ) -> Result<(), super::AppendUpdateError> - where - F: std::future::Future>, - { - 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?; diff --git a/crates/codegen/xai-grok-shell/src/session/storage/jsonl/tests.rs b/crates/codegen/xai-grok-shell/src/session/storage/jsonl/tests.rs index 2e33ee2..f743b38 100644 --- a/crates/codegen/xai-grok-shell/src/session/storage/jsonl/tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/storage/jsonl/tests.rs @@ -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!( diff --git a/crates/codegen/xai-grok-shell/src/session/storage/mod.rs b/crates/codegen/xai-grok-shell/src/session/storage/mod.rs index 05ed6ae..2c9c4a8 100644 --- a/crates/codegen/xai-grok-shell/src/session/storage/mod.rs +++ b/crates/codegen/xai-grok-shell/src/session/storage/mod.rs @@ -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 diff --git a/crates/codegen/xai-grok-shell/src/session/user_message.rs b/crates/codegen/xai-grok-shell/src/session/user_message.rs index 563c22f..819853c 100644 --- a/crates/codegen/xai-grok-shell/src/session/user_message.rs +++ b/crates/codegen/xai-grok-shell/src/session/user_message.rs @@ -104,26 +104,49 @@ pub async fn compute_vcs_status_block( working_directory: &Path, vcs_kind: VcsKind, ) -> Option { - 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). diff --git a/crates/codegen/xai-grok-shell/src/tools/notification_bridge.rs b/crates/codegen/xai-grok-shell/src/tools/notification_bridge.rs index 4bc23cf..86264e3 100644 --- a/crates/codegen/xai-grok-shell/src/tools/notification_bridge.rs +++ b/crates/codegen/xai-grok-shell/src/tools/notification_bridge.rs @@ -1,24 +1,19 @@ //! Notification bridge: translates `xai-grok-tools` `ToolNotification` events //! into `xai-grok-shell`'s native systems (ACP gateway, hunk tracker, file state tracker). - +use crate::session::commands::SessionCommand; +use crate::session::commands::{NotificationPriority, NotificationSource}; +use crate::session::persistence::PersistenceMsg; +use agent_client_protocol::{self as acp, Client as _}; use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; - -use agent_client_protocol::{self as acp, Client as _}; use tokio::sync::{Mutex as TokioMutex, mpsc}; use xai_acp_lib::AcpAgentGatewaySender as GatewaySender; use xai_grok_tools::notification::types::{ToolNotification, ToolNotificationHandle}; use xai_grok_tools::types::output::{BashOutput, ToolOutput}; -use xai_hunk_tracker::HunkTrackerHandle; - -use crate::session::commands::SessionCommand; -use crate::session::commands::{NotificationPriority, NotificationSource}; -use crate::session::persistence::PersistenceMsg; use xai_grok_workspace::session::file_state::FileStateTracker; - +use xai_hunk_tracker::HunkTrackerHandle; const TASK_WAKE_ADMISSION_TIMEOUT: std::time::Duration = std::time::Duration::from_millis(250); - /// Configuration for the notification bridge. pub struct NotificationBridgeConfig { /// ACP gateway for sending streaming updates to TUI @@ -97,7 +92,6 @@ pub struct NotificationBridgeConfig { /// `SessionActor::set_goal_loop_active_resource` for the rationale. pub goal_loop_active: std::sync::Arc, } - /// Snapshot a shared `OnceLock` tool-name slot as a borrowed `&str`. /// Returns `None` if the slot is still unset (toolset not yet finalized) /// or if the resolved value is `None` (no such tool registered in this @@ -105,32 +99,24 @@ pub struct NotificationBridgeConfig { pub(crate) fn resolved_tool_name(slot: &std::sync::OnceLock>) -> Option<&str> { slot.get().and_then(|v| v.as_deref()) } - /// Stamp a bridge-emitted notification's meta before it forks into /// persistence + broadcast — see `util::event_id::ensure_event_id_meta`. fn stamp_event_id(config: &NotificationBridgeConfig, meta: &mut Option) { crate::util::event_id::ensure_event_id_meta(&config.session_id.0, meta); } - /// Create a `ToolNotificationHandle` and spawn a bridge task that /// translates notifications into shell-native systems. pub fn spawn_notification_bridge(config: NotificationBridgeConfig) -> ToolNotificationHandle { let (handle, mut rx) = ToolNotificationHandle::channel(); - tokio::task::spawn_local(async move { - // Per-tool-call byte offset for incremental delta computation. - // Only used when `config.incremental_bash_output` is true. let mut offsets: HashMap = HashMap::new(); - while let Some(notification) = rx.recv().await { handle_notification(&config, notification, &mut offsets).await; } tracing::debug!("Notification bridge task exiting (sender dropped)"); }); - handle } - /// Emit a `CurrentModeUpdate` for the given [`SessionMode`] — persisted to /// `updates.jsonl` so session replay re-applies the mode, and forwarded to /// the gateway so the pager updates live. @@ -145,14 +131,11 @@ async fn emit_current_mode_update( )), ); stamp_event_id(config, &mut notification.meta); - let _ = config.persistence_tx.send(PersistenceMsg::Update( crate::session::storage::SessionUpdate::Acp(Box::new(notification.clone())), )); - config.gateway.forward_fire_and_forget(notification); } - /// Handle a single notification by forwarding it to the appropriate shell system. async fn handle_notification( config: &NotificationBridgeConfig, @@ -161,25 +144,19 @@ async fn handle_notification( ) { match notification { ToolNotification::BashOutputChunk(chunk) => { - // Compute output and output_delta based on incremental mode. let (output, output_delta) = if config.incremental_bash_output { let prev_offset = offsets.get(&chunk.base.tool_call_id).copied().unwrap_or(0); let full = &chunk.base.output; let delta = if prev_offset <= full.len() { full[prev_offset..].to_vec() } else { - // Buffer shrank (e.g. terminal clear / reset). - // Send the full buffer and reset offset. full.clone() }; offsets.insert(chunk.base.tool_call_id.clone(), full.len()); - // In incremental mode: output is empty, delta carries the bytes. (Vec::new(), Some(delta)) } else { (chunk.base.output.clone(), None) }; - - // Build a ToolOutput::Bash from the chunk for the TUI to parse let bash_output = ToolOutput::Bash(BashOutput { output_for_prompt: BashOutput::make_output_for_prompt(&String::from_utf8_lossy( &chunk.base.output, @@ -197,8 +174,6 @@ async fn handle_notification( output_delta, was_bare_echo: false, }); - - // Send ACP ToolCallUpdate with InProgress status for TUI streaming let update = acp::SessionUpdate::ToolCallUpdate(acp::ToolCallUpdate::new( acp::ToolCallId::new(chunk.base.tool_call_id.clone()), acp::ToolCallUpdateFields::new() @@ -212,12 +187,9 @@ async fn handle_notification( )); let mut notification = acp::SessionNotification::new(config.session_id.clone(), update); stamp_event_id(config, &mut notification.meta); - // Always persist — even when the gateway gate is closed, so bash - // output survives replay when the client later calls loadSession. let _ = config.persistence_tx.send(PersistenceMsg::Update( crate::session::storage::SessionUpdate::Acp(Box::new(notification.clone())), )); - // Only forward to the client if the gateway gate is open. if config .gateway_enabled .load(std::sync::atomic::Ordering::Relaxed) @@ -225,44 +197,31 @@ async fn handle_notification( let _ = config.gateway.session_notification(notification).await; } } - ToolNotification::BashExecutionComplete(complete) => { - // Clean up offset tracking for this tool call. offsets.remove(&complete.base.tool_call_id); tracing::debug!( - tool_call_id = %complete.base.tool_call_id, - exit_code = ?complete.exit_code, - "Bash execution complete notification received" + tool_call_id = % complete.base.tool_call_id, exit_code = ? complete + .exit_code, "Bash execution complete notification received" ); } - ToolNotification::BashExecutionTimeout(timeout) => { tracing::debug!( - tool_call_id = %timeout.base.tool_call_id, - elapsed = ?timeout.elapsed, + tool_call_id = % timeout.base.tool_call_id, elapsed = ? timeout.elapsed, "Bash execution timeout notification received" ); } - ToolNotification::BashExecutionFailed(failed) => { tracing::warn!( - tool_call_id = %failed.tool_call_id, - error = %failed.error, + tool_call_id = % failed.tool_call_id, error = % failed.error, "Bash execution failed notification received" ); } - ToolNotification::BashExecutionBackgrounded(bg) => { tracing::debug!( - tool_call_id = %bg.base.tool_call_id, - task_id = %bg.task_id, - command = %bg.base.command, - output_file = %bg.output_file.display(), + tool_call_id = % bg.base.tool_call_id, task_id = % bg.task_id, command = + % bg.base.command, output_file = % bg.output_file.display(), "Bash execution backgrounded notification received — forwarding to TUI" ); - - // Forward as x.ai/task_backgrounded ExtNotification so the TUI can - // correlate tool_call_id with task_id and populate the tasks panel. let mut notification = crate::extensions::notification::SessionNotification { session_id: config.session_id.clone(), update: crate::extensions::notification::SessionUpdate::TaskBackgrounded { @@ -281,12 +240,9 @@ async fn handle_notification( stamp_event_id(config, &mut meta_map); notification.meta = meta_map.map(serde_json::Value::Object); } - - // Persist so task correlation survives reconnect/replay. let _ = config.persistence_tx.send(PersistenceMsg::Update( crate::session::storage::SessionUpdate::Xai(Box::new(notification.clone())), )); - let params = serde_json::to_value(¬ification) .and_then(|v| serde_json::value::to_raw_value(&v)) .ok(); @@ -296,7 +252,6 @@ async fn handle_notification( config.gateway.forward_fire_and_forget(ext_notification); } } - ToolNotification::FileWritten(written) => { let prompt_index = *config.prompt_index.lock().await; config.hunk_tracker_handle.record_agent_write( @@ -305,7 +260,6 @@ async fn handle_notification( prompt_index, written.previous_content.clone(), ); - if written.previous_content.is_some() || written.is_new_file { config .file_state_tracker @@ -317,14 +271,11 @@ async fn handle_notification( ) .await; } - tracing::debug!( - path = %written.absolute_path.display(), - is_new_file = written.is_new_file, - "FileWritten notification forwarded to hunk tracker" + path = % written.absolute_path.display(), is_new_file = written + .is_new_file, "FileWritten notification forwarded to hunk tracker" ); } - ToolNotification::TaskCompleted(task_snapshot) => { let is_monitor = task_snapshot.kind == xai_grok_tools::computer::types::TaskKind::Monitor; @@ -332,24 +283,15 @@ async fn handle_notification( let goal_loop_active = config .goal_loop_active .load(std::sync::atomic::Ordering::Relaxed); - - // Natural monitor exit uses the same immediate wake path as bash; - // x.ai/task_completed still drives the pager UI in every branch. let mut will_wake = false; if task_snapshot.block_waited || task_snapshot.explicitly_killed { - // The blocking wait or kill result already reports completion. } else if goal_loop_active { - // Goal loop active: suppress the wake (synthetic prompt + the - // idle-gated fallback); surfaces 2/3 drain it. See - // `set_goal_loop_active_resource`. tracing::info!( - task_id = %task_id, - is_monitor, + task_id = % task_id, is_monitor, "auto-wake: suppressed completion (goal loop active)" ); } else if config.auto_wake_enabled { config.task_completion_reservations.reserve(task_id.clone()); - let tool_name = resolved_tool_name(&config.task_output_tool_name); let read_name = resolved_tool_name(&config.read_tool_name); let body = if is_monitor { @@ -367,7 +309,6 @@ async fn handle_notification( let message = xai_grok_tools::reminders::wrap_reminder(&body); let prompt_id = format!("task-completed-{task_id}"); let prompt_blocks = vec![acp::ContentBlock::Text(acp::TextContent::new(message))]; - let synthetic_trace_tx = config .synthetic_trace_tx .lock() @@ -376,9 +317,7 @@ async fn handle_notification( let (respond_to, completion_rx) = tokio::sync::oneshot::channel(); let (admission_tx, admission_rx) = tokio::sync::oneshot::channel(); tracing::info!( - task_id = %task_id, - prompt_id = %prompt_id, - is_monitor, + task_id = % task_id, prompt_id = % prompt_id, is_monitor, "auto-wake: requesting synthetic prompt admission for completed background task" ); let enqueued = config @@ -437,15 +376,12 @@ async fn handle_notification( xai_grok_telemetry::unified_log::info( "shell.task_wake.bridge_admission", Some(config.session_id.0.as_ref()), - Some(serde_json::json!({ - "task_id": &task_id, - "monitor": is_monitor, - "enqueued": enqueued, - "admitted": admitted, - "gate": config.task_wake_suppressed.get(), - })), + Some(serde_json::json!( + { "task_id" : & task_id, "monitor" : is_monitor, "enqueued" : + enqueued, "admitted" : admitted, "gate" : config + .task_wake_suppressed.get(), } + )), ); - if will_wake { if is_monitor { let _ = @@ -466,7 +402,7 @@ async fn handle_notification( .is_ok(); if copy_requested { tracing::info!( - task_id = %task_id, + task_id = % task_id, "auto-wake: sending synthetic turn trace request" ); let _ = trace_tx.send(crate::upload::turn::SyntheticTurnTraceRequest { @@ -477,19 +413,18 @@ async fn handle_notification( }); } else { tracing::debug!( - task_id = %task_id, + task_id = % task_id, "auto-wake: session snapshot request failed, skipping trace request" ); } } else { tracing::debug!( - task_id = %task_id, + task_id = % task_id, "auto-wake: no synthetic trace consumer, skipping trace request" ); } } } else { - // Auto-wake disabled — fall back to idle-gated notification drain. let tool_name = resolved_tool_name(&config.task_output_tool_name); let read_name = resolved_tool_name(&config.read_tool_name); let message = if is_monitor { @@ -528,8 +463,6 @@ async fn handle_notification( source, }); } - - // When a task is complete send notifications to the client so it can act on it let mut notification = crate::extensions::notification::SessionNotification { session_id: config.session_id.clone(), update: crate::extensions::notification::SessionUpdate::TaskCompleted { @@ -543,12 +476,9 @@ async fn handle_notification( stamp_event_id(config, &mut meta_map); notification.meta = meta_map.map(serde_json::Value::Object); } - - // Persist so task completion history survives reconnect/replay. let _ = config.persistence_tx.send(PersistenceMsg::Update( crate::session::storage::SessionUpdate::Xai(Box::new(notification.clone())), )); - let params = serde_json::to_value(¬ification) .and_then(|v| serde_json::value::to_raw_value(&v)) .ok(); @@ -557,7 +487,6 @@ async fn handle_notification( acp::ExtNotification::new("x.ai/task_completed", params.into()); config.gateway.forward_fire_and_forget(notification); } - let _ = config .session_cmd_tx .send(SessionCommand::DispatchNotificationHook { @@ -567,33 +496,23 @@ async fn handle_notification( level: Some("info".into()), }); } - ToolNotification::PlanModeEntered(entered) => { let activated = config.plan_mode.lock().activate_from_tool(); if activated { *config.current_prompt_mode.lock() = crate::session::plan_mode::PromptMode::Plan; *config.turn_prompt_mode.lock() = crate::session::plan_mode::PromptMode::Plan; - let snapshot = config.plan_mode.lock().snapshot(); let _ = config .persistence_tx .send(PersistenceMsg::PlanModeState(snapshot)); - - // Notify the frontend immediately so the plan-mode chip appears in the UI - // (currentModeId = 'plan'). Without this the agent can silently enter plan - // mode via the EnterPlanMode tool and the UI would never update. emit_current_mode_update(config, xai_grok_tools::types::SessionMode::Plan).await; } tracing::info!( - tool_call_id = %entered.tool_call_id, - activated, + tool_call_id = % entered.tool_call_id, activated, "Plan mode entered via EnterPlanMode tool" ); } - ToolNotification::PlanModeExited(exited) => { - // v1: auto-approve. A full implementation would present an - // approval UI with reject/feedback options. let deactivated = { let mut tracker = config.plan_mode.lock(); let deactivated = tracker.deactivate_approved(); @@ -609,71 +528,53 @@ async fn handle_notification( if deactivated { *config.current_prompt_mode.lock() = crate::session::plan_mode::PromptMode::Agent; *config.turn_prompt_mode.lock() = crate::session::plan_mode::PromptMode::Agent; - let snapshot = config.plan_mode.lock().snapshot(); let _ = config .persistence_tx .send(PersistenceMsg::PlanModeState(snapshot)); - - // Mirror the entry path: emit a `CurrentModeUpdate("default")` - // so the pager flips out of plan mode without having to - // string-match tool titles. Persist + forward so the next - // session replay also sees the exit. emit_current_mode_update(config, xai_grok_tools::types::SessionMode::Default).await; } tracing::info!( - tool_call_id = %exited.tool_call_id, - deactivated, - has_plan = exited.plan_content.is_some(), - "Plan mode exited via ExitPlanMode tool" + tool_call_id = % exited.tool_call_id, deactivated, has_plan = exited + .plan_content.is_some(), "Plan mode exited via ExitPlanMode tool" ); } - ToolNotification::UserQuestionAsked(asked) => { - tracing::info!( - tool_call_id = %asked.tool_call_id, - "User question asked" - ); + tracing::info!(tool_call_id = % asked.tool_call_id, "User question asked"); } - ToolNotification::LspServerStarting(s) => { - tracing::debug!(server = %s.server_name, command = %s.command, "LSP server starting"); + tracing::debug!( + server = % s.server_name, command = % s.command, "LSP server starting" + ); } ToolNotification::LspServerReady(s) => { - tracing::info!(server = %s.server_name, "LSP server ready"); + tracing::info!(server = % s.server_name, "LSP server ready"); } ToolNotification::LspServerCrashed(s) => { - tracing::warn!(server = %s.server_name, "LSP server crashed"); + tracing::warn!(server = % s.server_name, "LSP server crashed"); } ToolNotification::LspServerRetrying(s) => { tracing::warn!( - server = %s.server_name, - attempt = s.attempt, - max_restarts = s.max_restarts, - backoff_ms = s.backoff_ms, - "LSP server retrying" + server = % s.server_name, attempt = s.attempt, max_restarts = s + .max_restarts, backoff_ms = s.backoff_ms, "LSP server retrying" ); } ToolNotification::LspServerFailed(s) => { - tracing::error!(server = %s.server_name, error = %s.error, "LSP server failed"); + tracing::error!( + server = % s.server_name, error = % s.error, "LSP server failed" + ); } - ToolNotification::ScheduledTaskFired(fired) => { tracing::info!( - task_id = %fired.task_id, - schedule = %fired.human_schedule, - subagent_id = fired.subagent_id.as_deref().unwrap_or(""), - "Scheduled task fired" + task_id = % fired.task_id, schedule = % fired.human_schedule, subagent_id + = fired.subagent_id.as_deref().unwrap_or(""), "Scheduled task fired" ); - if fired.subagent_id.is_none() { - let inject_payload = serde_json::json!({ - "sessionId": config.session_id, - "taskId": &fired.task_id, - "prompt": &fired.prompt, - "humanSchedule": &fired.human_schedule, - "nextFireAt": &fired.next_fire_at, - }); + let inject_payload = serde_json::json!( + { "sessionId" : config.session_id, "taskId" : & fired.task_id, + "prompt" : & fired.prompt, "humanSchedule" : & fired.human_schedule, + "nextFireAt" : & fired.next_fire_at, } + ); if let Ok(params) = serde_json::value::to_raw_value(&inject_payload) { config .gateway @@ -683,7 +584,6 @@ async fn handle_notification( )); } } - let fired_notif = crate::extensions::notification::SessionNotification { session_id: config.session_id.clone(), update: crate::extensions::notification::SessionUpdate::ScheduledTaskFired { @@ -706,34 +606,22 @@ async fn handle_notification( )); } } - ToolNotification::MonitorEvent(event) => { - // Cross-session guard: in leader mode many sessions share one agent - // process, so drop events whose owner isn't this bridge's session - // (else session A's monitor injects reminders into session B). - // `None` owners (legacy backends) pass through. let my_session = config.session_id.0.as_ref(); if let Some(owner) = event.owner_session_id.as_deref() && owner != my_session { - // WARN (not debug) to surface the leader-mode mis-route in logs. tracing::warn!( - task_id = %event.task_id, - description = %event.description, - monitor_owner = %owner, - bridge_session = %my_session, + task_id = % event.task_id, description = % event.description, + monitor_owner = % owner, bridge_session = % my_session, "Dropped cross-session monitor event: owner does not match this bridge's session" ); return; } - tracing::debug!( - task_id = %event.task_id, - description = %event.description, + task_id = % event.task_id, description = % event.description, "Monitor event received, injecting into session" ); - - // Forward to pager -- raw text for the bg task stdout buffer. let notification = crate::extensions::notification::SessionNotification { session_id: config.session_id.clone(), update: crate::extensions::notification::SessionUpdate::MonitorEvent { @@ -754,19 +642,13 @@ async fn handle_notification( params.into(), )); } - - // If this monitor already auto-woke via TaskCompleted, do not inject - // model-facing notifications (avoids a second NotificationDrain turn - // with the same ended signal). Pager UI still got the event above. if config.task_completion_reservations.contains(&event.task_id) { tracing::debug!( - task_id = %event.task_id, + task_id = % event.task_id, "skipping model inject for monitor event: task already auto-woke via TaskCompleted" ); return; } - - // Inject the event into the notification queue for idle-gated drain. let prompt_id = format!("monitor-{}-{}", event.task_id, uuid::Uuid::now_v7()); let prompt_blocks = vec![acp::ContentBlock::Text(acp::TextContent::new( event.event_text, @@ -782,10 +664,8 @@ async fn handle_notification( }, }); } - ToolNotification::ScheduledTaskRemoved(removed) => { - tracing::info!(task_id = %removed.task_id, "Scheduled task removed"); - + tracing::info!(task_id = % removed.task_id, "Scheduled task removed"); let mut notification = crate::extensions::notification::SessionNotification { session_id: config.session_id.clone(), update: crate::extensions::notification::SessionUpdate::ScheduledTaskDeleted { @@ -798,8 +678,6 @@ async fn handle_notification( stamp_event_id(config, &mut meta_map); notification.meta = meta_map.map(serde_json::Value::Object); } - // Persist the deletion too, so replay on resume nets out a removed - // loop instead of resurrecting it from a persisted `created` line. let _ = config.persistence_tx.send(PersistenceMsg::Update( crate::session::storage::SessionUpdate::Xai(Box::new(notification.clone())), )); @@ -814,10 +692,8 @@ async fn handle_notification( )); } } - ToolNotification::ScheduledTaskCreated(created) => { - tracing::info!(task_id = %created.task_id, "Scheduled task created"); - + tracing::info!(task_id = % created.task_id, "Scheduled task created"); let mut notification = crate::extensions::notification::SessionNotification { session_id: config.session_id.clone(), update: crate::extensions::notification::SessionUpdate::ScheduledTaskCreated { @@ -833,14 +709,6 @@ async fn handle_notification( stamp_event_id(config, &mut meta_map); notification.meta = meta_map.map(serde_json::Value::Object); } - // Persist so the loop survives reconnect/replay, mirroring - // TaskBackgrounded/TaskCompleted. Without this, a second terminal - // that resumes the session restores monitors and subagents (whose - // notifications are persisted) but NOT loops, so the "watching" cue - // and Tasks pane undercount until the loop next fires. Create/delete - // are infrequent (bounded by the number of live loops); the - // recurring `_fired` notification is deliberately NOT persisted to - // avoid unbounded log growth. let _ = config.persistence_tx.send(PersistenceMsg::Update( crate::session::storage::SessionUpdate::Xai(Box::new(notification.clone())), )); @@ -857,13 +725,11 @@ async fn handle_notification( } } } - #[cfg(test)] mod tests { use super::*; use xai_grok_tools::computer::types::TaskKind; use xai_grok_tools::types::TaskSnapshot; - /// Drive the admission handshake inline so receiver assertions observe the /// bridge's command order without racing a detached proxy task. async fn handle_notification_with_admission( @@ -875,10 +741,10 @@ mod tests { ) { let notification = handle_notification(config, notification, offsets); tokio::pin!(notification); - let mut command = tokio::select! { - _ = &mut notification => panic!("notification completed before requesting admission"), - command = cmd_rx.recv() => command.expect("expected task-wake prompt"), + _ = & mut notification => + panic!("notification completed before requesting admission"), command = + cmd_rx.recv() => command.expect("expected task-wake prompt"), }; let SessionCommand::Prompt { admission, .. } = &mut command else { panic!("expected task-wake prompt"); @@ -895,7 +761,6 @@ mod tests { .expect("test command receiver must remain open"); notification.await; } - fn make_test_config() -> ( NotificationBridgeConfig, mpsc::UnboundedReceiver, @@ -903,7 +768,6 @@ mod tests { let (config, _gateway_rx, _persistence_rx, session_cmd_rx) = make_test_config_full(); (config, session_cmd_rx) } - #[allow(clippy::type_complexity)] fn make_test_config_full() -> ( NotificationBridgeConfig, @@ -913,7 +777,6 @@ mod tests { ) { make_test_config_full_raw() } - #[allow(clippy::type_complexity)] fn make_test_config_full_raw() -> ( NotificationBridgeConfig, @@ -960,7 +823,6 @@ mod tests { }; (config, gateway_rx, persistence_rx, session_cmd_rx) } - fn make_task_snapshot(task_id: &str, kind: TaskKind) -> TaskSnapshot { TaskSnapshot { task_id: task_id.into(), @@ -981,7 +843,6 @@ mod tests { owner_session_id: None, } } - #[tokio::test] async fn bash_task_completed_injects_bash_task_completed_source() { let (config, mut cmd_rx) = make_test_config(); @@ -992,10 +853,8 @@ mod tests { let snapshot = make_task_snapshot("bg-123", TaskKind::Bash); let notification = ToolNotification::TaskCompleted(snapshot); let mut offsets = HashMap::new(); - handle_notification_with_admission(&config, notification, &mut offsets, &mut cmd_rx, true) .await; - let command = cmd_rx.try_recv().expect("expected Prompt"); match command { SessionCommand::Prompt { @@ -1017,7 +876,6 @@ mod tests { } _ => panic!("expected Prompt"), } - let cmd3 = cmd_rx .try_recv() .expect("expected DispatchNotificationHook for task_complete"); @@ -1036,7 +894,6 @@ mod tests { _ => panic!("expected DispatchNotificationHook"), } } - /// Gap 1: while a goal loop is active, a completed background bash task /// must NOT fire the synthetic auto-wake prompt — an async "task completed" /// wake mid-goal derails a weak model. It must also NOT be marked @@ -1054,35 +911,31 @@ mod tests { .store(true, std::sync::atomic::Ordering::Relaxed); let snapshot = make_task_snapshot("bg-goal", TaskKind::Bash); let mut offsets = HashMap::new(); - handle_notification( &config, ToolNotification::TaskCompleted(snapshot), &mut offsets, ) .await; - - // No synthetic prompt / CopyFile / InjectNotification while the goal - // loop drives the turn — only the Notification hook dispatch. match cmd_rx .try_recv() .expect("expected DispatchNotificationHook for task_complete") { SessionCommand::DispatchNotificationHook { notification_type, .. - } => assert_eq!(notification_type, "task_complete"), + } => { + assert_eq!(notification_type, "task_complete") + } _ => panic!("unexpected session command"), } assert!( cmd_rx.try_recv().is_err(), "goal-loop-active bash completion must not inject auto-wake commands" ); - // Not marked reserved: surface 2 must be free to drain it. assert!( config.task_completion_reservations.snapshot().is_empty(), "goal-loop-active completion must not be marked reserved" ); - // The pager UI notification must still be emitted. let mut found_ext = false; while let Ok(msg) = gateway_rx.try_recv() { if let xai_acp_lib::AcpClientMessage::ExtNotification(args) = msg @@ -1096,7 +949,6 @@ mod tests { "x.ai/task_completed ExtNotification must still be sent for UI" ); } - /// Gap 1 (preserve non-goal behavior): with the goal loop inactive — the /// default for a normal session — a completed bash task DOES fire the /// synthetic auto-wake prompt AND is marked reserved so surface @@ -1110,7 +962,6 @@ mod tests { .expect("slot is fresh in this test fixture"); let snapshot = make_task_snapshot("bg-normal", TaskKind::Bash); let mut offsets = HashMap::new(); - handle_notification_with_admission( &config, ToolNotification::TaskCompleted(snapshot), @@ -1119,7 +970,6 @@ mod tests { true, ) .await; - assert!(matches!( cmd_rx.try_recv(), Ok(SessionCommand::Prompt { .. }) @@ -1133,7 +983,6 @@ mod tests { vec!["bg-normal".to_string()], ); } - fn task_completed_will_wake( gateway_rx: &mut mpsc::UnboundedReceiver, ) -> Option { @@ -1147,7 +996,6 @@ mod tests { } None } - /// The completion notification carries the wake verdict — the pager keys /// its between-turns status line on it (skip when a wake response /// follows, emit when nothing else will mark the moment). @@ -1189,7 +1037,6 @@ mod tests { trace_rx.try_recv().is_ok(), "accepted admission must request a synthetic-turn trace" ); - let (config, mut gateway_rx, mut persistence_rx, mut cmd_rx) = make_test_config_full(); let (trace_tx, mut trace_rx) = mpsc::unbounded_channel(); *config @@ -1243,7 +1090,6 @@ mod tests { "declined admission must still persist x.ai/task_completed" ); } - #[tokio::test(start_paused = true)] async fn stalled_admission_is_bounded_and_task_completion_still_emits() { let (config, mut gateway_rx, mut persistence_rx, mut cmd_rx) = make_test_config_full_raw(); @@ -1258,16 +1104,15 @@ mod tests { &mut offsets, ); tokio::pin!(notification); - tokio::select! { - _ = &mut notification => panic!("admission should still be waiting"), - command = cmd_rx.recv() => assert!(matches!(command, Some(SessionCommand::Prompt { .. }))), + _ = & mut notification => panic!("admission should still be waiting"), + command = cmd_rx.recv() => assert!(matches!(command, + Some(SessionCommand::Prompt { .. }))), } tokio::time::advance(TASK_WAKE_ADMISSION_TIMEOUT + std::time::Duration::from_millis(1)) .await; tokio::task::yield_now().await; notification.await; - assert_eq!(task_completed_will_wake(&mut gateway_rx), Some(false)); assert!( config.task_completion_reservations.contains("bg-stalled"), @@ -1287,7 +1132,6 @@ mod tests { } assert!(persisted_completion); } - #[tokio::test(start_paused = true)] async fn timed_out_monitor_admission_queues_one_fallback_and_late_actor_drops_prompt() { let (config, mut gateway_rx, _persistence_rx, mut cmd_rx) = make_test_config_full_raw(); @@ -1303,14 +1147,13 @@ mod tests { ); tokio::pin!(notification); let prompt = tokio::select! { - _ = &mut notification => panic!("admission should still be waiting"), + _ = & mut notification => panic!("admission should still be waiting"), command = cmd_rx.recv() => command.expect("prompt command"), }; tokio::time::advance(TASK_WAKE_ADMISSION_TIMEOUT + std::time::Duration::from_millis(1)) .await; tokio::task::yield_now().await; notification.await; - let SessionCommand::Prompt { admission: Some(admission), respond_to, @@ -1319,10 +1162,11 @@ mod tests { else { panic!("expected task wake prompt"); }; - assert!(matches!( - admission.fallback.source, - NotificationSource::MonitorCompleted { ref task_id } if task_id == "mon-timeout" - )); + assert!( + matches!(admission.fallback.source, NotificationSource::MonitorCompleted { + ref task_id } +if task_id == "mon-timeout") + ); assert!(admission.respond_to.send(true).is_err()); let _ = respond_to.send(Ok(crate::session::commands::PromptTurnOk { stop_reason: acp::StopReason::Cancelled, @@ -1332,7 +1176,6 @@ mod tests { structured_output: None, usage: None, })); - assert!(matches!( cmd_rx.try_recv(), Ok(SessionCommand::DispatchNotificationHook { .. }) @@ -1344,7 +1187,6 @@ mod tests { "the late actor fallback retains the reservation until user delivery" ); } - #[tokio::test] async fn task_completed_stamps_will_wake_false_when_session_channel_closed() { let (config, mut gateway_rx, _persistence_rx, cmd_rx) = make_test_config_full_raw(); @@ -1372,7 +1214,6 @@ mod tests { config.task_completion_reservations.release("bg-dead"); assert!(!config.task_completion_reservations.contains("bg-dead")); } - /// Gap 1 (adjacent branch): the goal-loop arm sits BEFORE the /// `auto_wake_enabled == false` `InjectNotification` fallback, so an /// auto-wake-DISABLED completion mid-goal must also be suppressed — it must @@ -1387,22 +1228,21 @@ mod tests { .store(true, std::sync::atomic::Ordering::Relaxed); let snapshot = make_task_snapshot("bg-disabled-goal", TaskKind::Bash); let mut offsets = HashMap::new(); - handle_notification( &config, ToolNotification::TaskCompleted(snapshot), &mut offsets, ) .await; - - // No InjectNotification during the goal loop, even with auto-wake disabled. match cmd_rx .try_recv() .expect("expected DispatchNotificationHook for task_complete") { SessionCommand::DispatchNotificationHook { notification_type, .. - } => assert_eq!(notification_type, "task_complete"), + } => { + assert_eq!(notification_type, "task_complete") + } _ => panic!("unexpected session command"), } assert!( @@ -1411,7 +1251,6 @@ mod tests { ); assert!(config.task_completion_reservations.snapshot().is_empty()); } - /// Natural monitor exit (including exit code 0) must immediate-auto-wake /// the same way bash does — not only via the idle-gated MonitorEvent path. /// Also drops queued MonitorEvents so a second NotificationDrain turn is @@ -1428,7 +1267,6 @@ mod tests { snapshot.command = "tail -f deploy.log".into(); snapshot.exit_code = Some(0); let mut offsets = HashMap::new(); - handle_notification_with_admission( &config, ToolNotification::TaskCompleted(snapshot), @@ -1437,7 +1275,6 @@ mod tests { true, ) .await; - let cmd = cmd_rx.try_recv().expect("expected Prompt auto-wake"); match cmd { SessionCommand::Prompt { @@ -1485,7 +1322,6 @@ mod tests { vec!["mon-456".to_string()], ); } - #[tokio::test] async fn declined_quiet_monitor_wake_queues_canonical_deferred_completion() { let (config, _gateway_rx, mut persistence_rx, mut cmd_rx) = make_test_config_full(); @@ -1494,7 +1330,6 @@ mod tests { .set(Some("get_command_or_subagent_output".to_string())) .expect("slot is fresh in this test fixture"); let mut offsets = HashMap::new(); - handle_notification_with_admission( &config, ToolNotification::TaskCompleted(make_task_snapshot("mon-declined", TaskKind::Monitor)), @@ -1503,7 +1338,6 @@ mod tests { false, ) .await; - assert!(matches!( cmd_rx.try_recv(), Ok(SessionCommand::Prompt { .. }) @@ -1531,7 +1365,6 @@ mod tests { "the actor owns reservation release after queuing the deferred fallback" ); } - /// After TaskCompleted auto-wake reserves the task, late pipeline /// MonitorEvents must not inject another model-facing notification. #[tokio::test] @@ -1541,7 +1374,6 @@ mod tests { .task_completion_reservations .reserve("mon-done".into()); let mut offsets = HashMap::new(); - handle_notification( &config, ToolNotification::MonitorEvent(xai_grok_tools::notification::types::MonitorEvent { @@ -1554,14 +1386,11 @@ mod tests { &mut offsets, ) .await; - - // No InjectNotification — only the TaskCompleted wake should talk to the model. assert!( cmd_rx.try_recv().is_err(), "post-auto-wake MonitorEvent must not InjectNotification" ); } - /// Explicit kill of a monitor still skips auto-wake — the model already /// got the kill_task tool result. #[tokio::test] @@ -1570,21 +1399,21 @@ mod tests { let mut snapshot = make_task_snapshot("mon-killed", TaskKind::Monitor); snapshot.explicitly_killed = true; let mut offsets = HashMap::new(); - handle_notification( &config, ToolNotification::TaskCompleted(snapshot), &mut offsets, ) .await; - match cmd_rx .try_recv() .expect("expected DispatchNotificationHook for task_complete") { SessionCommand::DispatchNotificationHook { notification_type, .. - } => assert_eq!(notification_type, "task_complete"), + } => { + assert_eq!(notification_type, "task_complete") + } _ => panic!("unexpected session command"), } assert!( @@ -1593,7 +1422,6 @@ mod tests { ); assert!(config.task_completion_reservations.snapshot().is_empty()); } - /// Goal-loop suppression applies to monitor completions too. #[tokio::test] async fn monitor_task_completed_suppressed_during_goal_loop() { @@ -1603,21 +1431,21 @@ mod tests { .store(true, std::sync::atomic::Ordering::Relaxed); let snapshot = make_task_snapshot("mon-goal", TaskKind::Monitor); let mut offsets = HashMap::new(); - handle_notification( &config, ToolNotification::TaskCompleted(snapshot), &mut offsets, ) .await; - match cmd_rx .try_recv() .expect("expected DispatchNotificationHook for task_complete") { SessionCommand::DispatchNotificationHook { notification_type, .. - } => assert_eq!(notification_type, "task_complete"), + } => { + assert_eq!(notification_type, "task_complete") + } _ => panic!("unexpected session command"), } assert!( @@ -1626,12 +1454,8 @@ mod tests { ); assert!(config.task_completion_reservations.snapshot().is_empty()); } - #[tokio::test] async fn scheduled_task_created_is_persisted() { - // A `/loop` create must be persisted (like TaskBackgrounded) so a - // second terminal that resumes the session restores the loop from - // replay — otherwise it stays invisible until the loop next fires. let (config, _gateway_rx, mut persistence_rx, _cmd_rx) = make_test_config_full(); let notification = ToolNotification::ScheduledTaskCreated( xai_grok_tools::notification::types::ScheduledTaskCreated { @@ -1642,9 +1466,7 @@ mod tests { }, ); let mut offsets = HashMap::new(); - handle_notification(&config, notification, &mut offsets).await; - let msg = persistence_rx .try_recv() .expect("scheduled_task_created must be persisted"); @@ -1667,7 +1489,6 @@ mod tests { _ => panic!("expected PersistenceMsg::Update(Xai(ScheduledTaskCreated))"), } } - /// Persisted⇒stamped contract at the bridge's highest-frequency emitter: /// the persisted bash-output line carries an `eventId`, and the live /// broadcast carries the SAME id (the meta is minted before the @@ -1689,9 +1510,7 @@ mod tests { }, ); let mut offsets = HashMap::new(); - handle_notification(&config, notification, &mut offsets).await; - let persisted_id = match persistence_rx.try_recv().expect("chunk must be persisted") { PersistenceMsg::Update(crate::session::storage::SessionUpdate::Acp(notif)) => notif .meta @@ -1702,7 +1521,6 @@ mod tests { .to_string(), other => panic!("expected PersistenceMsg::Update(Acp(..)), got {other:?}"), }; - let broadcast_id = match gateway_rx.try_recv().expect("chunk must be broadcast") { xai_acp_lib::AcpClientMessage::SessionNotification(args) => args .request @@ -1716,11 +1534,8 @@ mod tests { }; assert_eq!(persisted_id, broadcast_id); } - #[tokio::test] async fn scheduled_task_removed_is_persisted() { - // The deletion must also persist so replay nets out a removed loop - // instead of resurrecting it from the persisted `created` line. let (config, _gateway_rx, mut persistence_rx, _cmd_rx) = make_test_config_full(); let notification = ToolNotification::ScheduledTaskRemoved( xai_grok_tools::notification::types::ScheduledTaskRemoved { @@ -1728,9 +1543,7 @@ mod tests { }, ); let mut offsets = HashMap::new(); - handle_notification(&config, notification, &mut offsets).await; - let msg = persistence_rx .try_recv() .expect("scheduled_task_removed must be persisted"); @@ -1748,7 +1561,6 @@ mod tests { _ => panic!("expected PersistenceMsg::Update(Xai(ScheduledTaskDeleted))"), } } - fn xai_persisted_event_id( notif: &crate::extensions::notification::SessionNotification, ) -> Option { @@ -1759,7 +1571,6 @@ mod tests { .and_then(|v| v.as_str()) .map(str::to_string) } - /// Per-site stamp pins for the bridge emitters not covered by the /// representative chokepoint tests: deleting any one `stamp_event_id` /// call must fail a test (an id-less persisted line silently disables @@ -1784,9 +1595,7 @@ mod tests { }, ); let mut offsets = HashMap::new(); - handle_notification(&config, notification, &mut offsets).await; - match persistence_rx.try_recv().expect("must persist") { PersistenceMsg::Update(crate::session::storage::SessionUpdate::Xai(notif)) => { assert!(xai_persisted_event_id(¬if).is_some()); @@ -1794,21 +1603,17 @@ mod tests { _ => panic!("expected Xai update"), } } - #[tokio::test] async fn task_completed_persisted_line_is_stamped() { let (config, _gateway_rx, mut persistence_rx, _cmd_rx) = make_test_config_full(); - // Monitor kind: persists without the bash auto-wake side effects. let snapshot = make_task_snapshot("mon-1", TaskKind::Monitor); let mut offsets = HashMap::new(); - handle_notification( &config, ToolNotification::TaskCompleted(snapshot), &mut offsets, ) .await; - match persistence_rx.try_recv().expect("must persist") { PersistenceMsg::Update(crate::session::storage::SessionUpdate::Xai(notif)) => { assert!(xai_persisted_event_id(¬if).is_some()); @@ -1816,13 +1621,10 @@ mod tests { _ => panic!("expected Xai update"), } } - #[tokio::test] async fn current_mode_update_persisted_line_is_stamped() { let (config, _gateway_rx, mut persistence_rx, _cmd_rx) = make_test_config_full(); - emit_current_mode_update(&config, xai_grok_tools::types::SessionMode::Plan).await; - match persistence_rx.try_recv().expect("must persist") { PersistenceMsg::Update(crate::session::storage::SessionUpdate::Acp(notif)) => { assert!(matches!( @@ -1842,13 +1644,8 @@ mod tests { _ => panic!("expected Acp update"), } } - #[tokio::test] async fn scheduled_task_fired_is_not_persisted() { - // `_fired` recurs on every interval; persisting it would grow the - // updates log without bound. Loops are restored from create/delete, so - // the fire stays gateway-only (the pager self-heals the entry on a live - // fire if needed). let (config, _gateway_rx, mut persistence_rx, _cmd_rx) = make_test_config_full(); let notification = ToolNotification::ScheduledTaskFired( xai_grok_tools::notification::types::ScheduledTaskFired { @@ -1860,15 +1657,12 @@ mod tests { }, ); let mut offsets = HashMap::new(); - handle_notification(&config, notification, &mut offsets).await; - assert!( persistence_rx.try_recv().is_err(), "scheduled_task_fired must NOT be persisted (recurring \u{2192} unbounded log growth)" ); } - fn make_monitor_event_notification(task_id: &str, owner: Option<&str>) -> ToolNotification { ToolNotification::MonitorEvent(xai_grok_tools::notification::types::MonitorEvent { task_id: task_id.into(), @@ -1878,21 +1672,12 @@ mod tests { owner_session_id: owner.map(str::to_string), }) } - #[tokio::test] async fn cross_session_monitor_event_is_dropped() { - // The bridge belongs to "test-session"; the event is owned by a - // different session. In leader mode (one agent process, many sessions) - // this is the cross-session leak: without the owner guard the foreign - // monitor would inject a `` reminder into this session's - // conversation. Assert it is fully dropped — no conversation injection - // and no pager forward. let (config, mut gateway_rx, _persistence_rx, mut cmd_rx) = make_test_config_full(); let notification = make_monitor_event_notification("mon-foreign", Some("other-session")); let mut offsets = HashMap::new(); - handle_notification(&config, notification, &mut offsets).await; - assert!( cmd_rx.try_recv().is_err(), "cross-session monitor event must not be injected into this session" @@ -1907,38 +1692,31 @@ mod tests { } } } - #[tokio::test] async fn same_session_monitor_event_is_injected() { - // Owner matches the bridge's own session id ("test-session") -> deliver. let (config, mut cmd_rx) = make_test_config(); let notification = make_monitor_event_notification("mon-own", Some("test-session")); let mut offsets = HashMap::new(); - handle_notification(&config, notification, &mut offsets).await; - match cmd_rx .try_recv() .expect("own-session monitor event must be injected") { SessionCommand::InjectNotification { source, .. } => match source { - NotificationSource::MonitorEvent { task_id } => assert_eq!(task_id, "mon-own"), + NotificationSource::MonitorEvent { task_id } => { + assert_eq!(task_id, "mon-own") + } _ => panic!("expected MonitorEvent notification source"), }, _ => panic!("expected InjectNotification"), } } - #[tokio::test] async fn legacy_monitor_event_without_owner_is_injected() { - // Legacy / non-grok-build backends record no owner; such events must - // pass through unchanged for backwards compatibility. let (config, mut cmd_rx) = make_test_config(); let notification = make_monitor_event_notification("mon-legacy", None); let mut offsets = HashMap::new(); - handle_notification(&config, notification, &mut offsets).await; - assert!( matches!( cmd_rx @@ -1952,7 +1730,6 @@ mod tests { "legacy monitor event should be injected as a MonitorEvent notification" ); } - #[tokio::test] async fn block_waited_task_skips_auto_wake_prompt() { let (config, mut gateway_rx, _persistence_rx, mut cmd_rx) = make_test_config_full(); @@ -1960,26 +1737,22 @@ mod tests { snapshot.block_waited = true; let notification = ToolNotification::TaskCompleted(snapshot); let mut offsets = HashMap::new(); - handle_notification(&config, notification, &mut offsets).await; - - // block_waited tasks must NOT inject a synthetic prompt — the - // blocking caller already received the result directly. match cmd_rx .try_recv() .expect("expected DispatchNotificationHook for task_complete") { SessionCommand::DispatchNotificationHook { notification_type, .. - } => assert_eq!(notification_type, "task_complete"), + } => { + assert_eq!(notification_type, "task_complete") + } _ => panic!("unexpected session command"), } assert!( cmd_rx.try_recv().is_err(), "block_waited completion should not send Prompt or InjectNotification" ); - - // The x.ai/task_completed ExtNotification for UI updates must still be sent. let mut found_ext = false; while let Ok(msg) = gateway_rx.try_recv() { if let xai_acp_lib::AcpClientMessage::ExtNotification(args) = msg @@ -1993,7 +1766,6 @@ mod tests { "x.ai/task_completed ExtNotification must still be sent for UI" ); } - #[tokio::test] async fn explicitly_killed_task_skips_auto_wake_prompt() { let (config, mut gateway_rx, _persistence_rx, mut cmd_rx) = make_test_config_full(); @@ -2001,26 +1773,22 @@ mod tests { snapshot.explicitly_killed = true; let notification = ToolNotification::TaskCompleted(snapshot); let mut offsets = HashMap::new(); - handle_notification(&config, notification, &mut offsets).await; - - // explicitly_killed tasks must NOT inject a synthetic prompt — the - // model already received the KillTaskResult from the kill tool. match cmd_rx .try_recv() .expect("expected DispatchNotificationHook for task_complete") { SessionCommand::DispatchNotificationHook { notification_type, .. - } => assert_eq!(notification_type, "task_complete"), + } => { + assert_eq!(notification_type, "task_complete") + } _ => panic!("unexpected session command"), } assert!( cmd_rx.try_recv().is_err(), "explicitly_killed completion should not send Prompt or InjectNotification" ); - - // The x.ai/task_completed ExtNotification for UI updates must still be sent. let mut found_ext = false; while let Ok(msg) = gateway_rx.try_recv() { if let xai_acp_lib::AcpClientMessage::ExtNotification(args) = msg @@ -2034,7 +1802,6 @@ mod tests { "x.ai/task_completed ExtNotification must still be sent for UI" ); } - #[tokio::test] async fn bash_task_completed_falls_back_when_auto_wake_disabled() { let (mut config, mut cmd_rx) = make_test_config(); @@ -2046,10 +1813,7 @@ mod tests { let snapshot = make_task_snapshot("bg-disabled", TaskKind::Bash); let notification = ToolNotification::TaskCompleted(snapshot); let mut offsets = HashMap::new(); - handle_notification(&config, notification, &mut offsets).await; - - // With auto-wake disabled, should use InjectNotification (not Prompt). let cmd = cmd_rx.try_recv().expect("expected InjectNotification"); match cmd { SessionCommand::InjectNotification { @@ -2061,10 +1825,11 @@ mod tests { } => { assert!(prompt_id.starts_with("bash-completed-")); assert_eq!(priority, NotificationPriority::Later); - assert!(matches!( - source, - NotificationSource::BashTaskCompleted { ref task_id } if task_id == "bg-disabled" - )); + assert!( + matches!(source, NotificationSource::BashTaskCompleted { ref task_id + } +if task_id == "bg-disabled") + ); let text = match &prompt_blocks[0] { acp::ContentBlock::Text(t) => &t.text, _ => panic!("expected text block"), @@ -2075,7 +1840,6 @@ mod tests { } _ => panic!("expected InjectNotification"), } - let hook_cmd = cmd_rx .try_recv() .expect("expected DispatchNotificationHook for task_complete"); @@ -2094,17 +1858,14 @@ mod tests { _ => panic!("expected DispatchNotificationHook"), } } - #[tokio::test] async fn bash_completion_uses_single_task_id_clone() { let (config, mut cmd_rx) = make_test_config(); let snapshot = make_task_snapshot("unique-id-789", TaskKind::Bash); let notification = ToolNotification::TaskCompleted(snapshot); let mut offsets = HashMap::new(); - handle_notification_with_admission(&config, notification, &mut offsets, &mut cmd_rx, true) .await; - let cmd = cmd_rx.try_recv().unwrap(); if let SessionCommand::Prompt { prompt_id, .. } = cmd { assert_eq!(prompt_id, "task-completed-unique-id-789"); @@ -2112,41 +1873,32 @@ mod tests { panic!("expected Prompt"); } } - fn extract_current_mode_id(notification: &acp::SessionNotification) -> Option<&str> { match ¬ification.update { acp::SessionUpdate::CurrentModeUpdate(cmu) => Some(cmu.current_mode_id.0.as_ref()), _ => None, } } - /// Regression: `PlanModeExited` must emit `CurrentModeUpdate("default")` /// onto both the gateway and the persistence stream. Without this, /// agent-driven plan approvals leave the TUI stuck in plan mode. #[tokio::test] async fn plan_mode_exited_emits_current_mode_update_default() { let (config, mut gateway_rx, mut persistence_rx, _cmd_rx) = make_test_config_full(); - - // Pre-condition: agent path requires plan mode to be Active first - // so `deactivate_approved` actually flips state and triggers the emit. { let mut tracker = config.plan_mode.lock(); assert!(tracker.activate_from_tool()); } *config.current_prompt_mode.lock() = crate::session::plan_mode::PromptMode::Plan; *config.turn_prompt_mode.lock() = crate::session::plan_mode::PromptMode::Plan; - let notification = ToolNotification::PlanModeExited(xai_grok_tools::notification::types::PlanModeExited { tool_call_id: "tc-exit-1".into(), plan_content: Some("- step 1".into()), plan_file_path: "/tmp/test-session/plan.md".into(), }); - let mut offsets = HashMap::new(); handle_notification(&config, notification, &mut offsets).await; - - // Gateway: one CurrentModeUpdate("default"). let mut gateway_modes = Vec::new(); while let Ok(msg) = gateway_rx.try_recv() { if let xai_acp_lib::AcpClientMessage::SessionNotification(args) = msg @@ -2160,8 +1912,6 @@ mod tests { vec!["default".to_string()], "PlanModeExited should emit exactly one CurrentModeUpdate(default) to the gateway" ); - - // Persistence: same notification persisted so replay re-applies the exit. let mut persisted_modes = Vec::new(); while let Ok(msg) = persistence_rx.try_recv() { if let PersistenceMsg::Update(crate::session::storage::SessionUpdate::Acp(notif)) = msg @@ -2175,14 +1925,11 @@ mod tests { vec!["default".to_string()], "PlanModeExited should persist exactly one CurrentModeUpdate(default)" ); - - // Session-level prompt mode was reset. assert!(matches!( *config.current_prompt_mode.lock(), crate::session::plan_mode::PromptMode::Agent )); } - /// Default (grok) polarity: the exit_plan_mode tool result is the model's /// only exit signal, so an approved `PlanModeExited` must NOT arm the /// deferred exit reminder — in memory or in the persisted snapshot. @@ -2190,22 +1937,18 @@ mod tests { #[tokio::test] async fn plan_mode_exited_does_not_arm_exit_reminder_by_default() { let (config, _gateway_rx, mut persistence_rx, _cmd_rx) = make_test_config_full(); - { let mut tracker = config.plan_mode.lock(); assert!(tracker.activate_from_tool()); } - let notification = ToolNotification::PlanModeExited(xai_grok_tools::notification::types::PlanModeExited { tool_call_id: "tc-exit-grok".into(), plan_content: Some("- step 1".into()), plan_file_path: "/tmp/test-session/plan.md".into(), }); - let mut offsets = HashMap::new(); handle_notification(&config, notification, &mut offsets).await; - assert!( !config.plan_mode.lock().has_pending_exit_reminder(), "approved exit must not arm the deferred exit reminder" @@ -2224,7 +1967,6 @@ mod tests { "persisted plan-mode snapshot must not carry the exit reminder" ); } - /// Gated counterpart: when `queue_exit_reminder_on_approved_exit` is /// set, an approved `PlanModeExited` must arm the next-turn exit /// reminder and persist it. @@ -2234,22 +1976,18 @@ mod tests { config .queue_exit_reminder_on_approved_exit .store(true, std::sync::atomic::Ordering::Relaxed); - { let mut tracker = config.plan_mode.lock(); assert!(tracker.activate_from_tool()); } - let notification = ToolNotification::PlanModeExited(xai_grok_tools::notification::types::PlanModeExited { tool_call_id: "tc-exit-gated".into(), plan_content: Some("- step 1".into()), plan_file_path: "/tmp/test-session/plan.md".into(), }); - let mut offsets = HashMap::new(); handle_notification(&config, notification, &mut offsets).await; - assert!( config.plan_mode.lock().has_pending_exit_reminder(), "gated approved exit must arm the next-turn exit reminder" @@ -2268,22 +2006,18 @@ mod tests { "persisted plan-mode snapshot must carry the armed exit reminder" ); } - /// Symmetric to the exit test: `PlanModeEntered` emits /// `CurrentModeUpdate("plan")`. #[tokio::test] async fn plan_mode_entered_emits_current_mode_update_plan() { let (config, mut gateway_rx, mut persistence_rx, _cmd_rx) = make_test_config_full(); - let notification = ToolNotification::PlanModeEntered( xai_grok_tools::notification::types::PlanModeEntered { tool_call_id: "tc-enter-1".into(), }, ); - let mut offsets = HashMap::new(); handle_notification(&config, notification, &mut offsets).await; - let mut gateway_modes = Vec::new(); while let Ok(msg) = gateway_rx.try_recv() { if let xai_acp_lib::AcpClientMessage::SessionNotification(args) = msg @@ -2293,7 +2027,6 @@ mod tests { } } assert_eq!(gateway_modes, vec!["plan".to_string()]); - let mut persisted_modes = Vec::new(); while let Ok(msg) = persistence_rx.try_recv() { if let PersistenceMsg::Update(crate::session::storage::SessionUpdate::Acp(notif)) = msg @@ -2304,7 +2037,6 @@ mod tests { } assert_eq!(persisted_modes, vec!["plan".to_string()]); } - /// Build a completed-bash `TaskSnapshot` whose `output` is large enough /// to trip the inline-completion truncation cap, with a concrete /// `output_file` path so the disk-pointer footer is exercised end-to-end. @@ -2328,7 +2060,6 @@ mod tests { owner_session_id: None, } } - /// Extract the auto-wake prompt text emitted on the session command channel. fn auto_wake_prompt_text(cmd_rx: &mut mpsc::UnboundedReceiver) -> String { let cmd = cmd_rx.try_recv().expect("expected Prompt"); @@ -2340,7 +2071,6 @@ mod tests { _ => panic!("expected Prompt"), } } - /// Extract the InjectNotification prompt text emitted on the session /// command channel (auto-wake-disabled fallback path). fn inject_notification_prompt_text( @@ -2355,7 +2085,6 @@ mod tests { _ => panic!("expected InjectNotification"), } } - /// Bash completion with a large output and no polling tool (compat-harness /// toolset) renders the truncation marker AND the disk-pointer footer /// pointing the model at `output_file` via the resolved Read tool name. @@ -2365,7 +2094,6 @@ mod tests { #[tokio::test] async fn bash_completion_renders_disk_pointer_footer_in_both_branches() { let output_file = PathBuf::from("/tmp/bg-disk-pointer.log"); - let (config_auto, mut cmd_rx_auto) = make_test_config(); config_auto .read_tool_name @@ -2398,7 +2126,6 @@ mod tests { prompt.contains("bg-disk-1"), "auto-wake: prompt must reference task id" ); - let (mut config_no_wake, mut cmd_rx_no_wake) = make_test_config(); config_no_wake.auto_wake_enabled = false; config_no_wake diff --git a/crates/codegen/xai-grok-shell/src/upload/trace.rs b/crates/codegen/xai-grok-shell/src/upload/trace.rs index 8e4d87b..b374c4f 100644 --- a/crates/codegen/xai-grok-shell/src/upload/trace.rs +++ b/crates/codegen/xai-grok-shell/src/upload/trace.rs @@ -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, 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); diff --git a/crates/codegen/xai-grok-shell/tests/test_auth_provider_e2e.rs b/crates/codegen/xai-grok-shell/tests/test_auth_provider_e2e.rs new file mode 100644 index 0000000..3ad0e02 --- /dev/null +++ b/crates/codegen/xai-grok-shell/tests/test_auth_provider_e2e.rs @@ -0,0 +1,316 @@ +//! End-to-end test for `[auth_provider.]` per-model credential helpers. +//! +//! Runs the built grok binary headless against the mock inference server with +//! a config-defined BYOK model whose bearer comes from a mock auth binary (a +//! script the test writes to disk). The harness's `XAI_API_KEY` stands in for +//! the session-tier credential. +//! +//! `#[ignore]` (needs a built binary). The CI lifecycle lanes run it against +//! the release artifact via `GROK_BINARY`; run locally (auto-builds the pager): +//! ```bash +//! cargo test -p xai-grok-shell --test test_auth_provider_e2e -- --ignored +//! ``` +//! +//! Unix-only: the mock helpers are `sh` scripts run via `sh -c`. +#![cfg(unix)] + +use xai_grok_test_support::*; + +#[tokio::test] +#[ignore] +async fn provider_backed_model_sends_minted_token_on_the_wire() { + let server = MockInferenceServer::start() + .await + .expect("start mock server"); + let workdir = git_workdir(); + let home = tempfile::TempDir::new().unwrap(); + + let grok_home = home.path().join(".grok"); + std::fs::create_dir_all(&grok_home).expect("create .grok home"); + + let counter = grok_home.join("mint-count"); + let helper = grok_home.join("mock-auth.sh"); + std::fs::write( + &helper, + format!( + "#!/bin/sh\necho run >> {}\nprintf 'gateway-tok-1'\n", + counter.display() + ), + ) + .expect("write mock auth binary"); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(&helper, std::fs::Permissions::from_mode(0o755)).unwrap(); + } + + std::fs::write( + grok_home.join("config.toml"), + format!( + r#"[auth_provider.gateway] +command = "{helper}" +token_ttl_secs = 3600 + +[model.proxied-gateway] +model = "mock-gateway-model" +base_url = "{base}" +context_window = 200000 +auth_provider = "gateway" +"#, + helper = helper.display(), + // Already ends in `/v1`. + base = server.url(), + ), + ) + .expect("write config.toml"); + + let mut cmd = tokio::process::Command::new(grok_binary()); + cmd.args([ + "-p", + "say hi", + "--yolo", + "--model", + "proxied-gateway", + "--max-turns", + "1", + "--output-format", + "json", + ]) + .arg("--cwd") + .arg(workdir.path()) + .current_dir(workdir.path()) + .stdin(std::process::Stdio::null()) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .kill_on_drop(true); + xai_grok_test_support::env::test_env_cmd_tokio(&mut cmd, &server.url(), home.path()); + // Don't attach to a developer's ambient leader; spawn fresh against the mock. + cmd.env_remove("GROK_LEADER_SOCKET"); + + let result = run_headless_with_cmd(cmd).await; + assert_headless_success(&result, "auth provider e2e", Some(&server)); + + let runs = std::fs::read_to_string(&counter) + .expect("helper must have run") + .lines() + .count(); + assert_eq!(runs, 1, "one turn mints exactly once"); + + let requests = server.requests(); + // The mock server plays both the first-party and the provider role on + // one host, so the leak assertions are scoped to the inference path. + let chat = requests + .iter() + .find(|e| e.method == "POST" && e.path.contains("chat/completions")) + .unwrap_or_else(|| { + panic!( + "no POST /v1/chat/completions request logged; requests:\n{}", + server.request_log_summary() + ) + }); + assert_eq!( + chat.authorization.as_deref(), + Some("Bearer gateway-tok-1"), + "the inference request must carry the minted provider token; requests:\n{}", + server.request_log_summary() + ); + // `test-key-for-ci` is the harness's XAI_API_KEY; it stands in for the + // session credential, which resolves below the provider arm. + assert!( + !requests.iter().any(|e| { + e.path.contains("chat/completions") + && e.authorization + .as_deref() + .is_some_and(|a| a.contains("test-key-for-ci")) + }), + "the session credential must never reach the provider-backed endpoint; requests:\n{}", + server.request_log_summary() + ); +} + +/// A model that references an undefined `[auth_provider.]` must fail +/// closed end to end: the request goes out without a bearer (which the mock +/// rejects), and the session credential is never substituted onto the wire. +#[tokio::test] +#[ignore] +async fn undefined_provider_fails_closed_and_never_leaks_session_key() { + // Reject any bearer: nothing legitimate can satisfy this, since the model + // references a provider that is never defined. + let server = MockInferenceServer::start_with_required_auth( + vec![MockModelEntry::new("mock-gateway-model")], + "never-issued-token", + ) + .await + .expect("start mock server"); + let workdir = git_workdir(); + let home = tempfile::TempDir::new().unwrap(); + + let grok_home = home.path().join(".grok"); + std::fs::create_dir_all(&grok_home).expect("create .grok home"); + + // Model references `gateway`, but no `[auth_provider.gateway]` table exists. + std::fs::write( + grok_home.join("config.toml"), + format!( + r#"[model.proxied-gateway] +model = "mock-gateway-model" +base_url = "{base}" +context_window = 200000 +auth_provider = "gateway" +"#, + base = server.url(), + ), + ) + .expect("write config.toml"); + + let mut cmd = tokio::process::Command::new(grok_binary()); + cmd.args([ + "-p", + "say hi", + "--yolo", + "--model", + "proxied-gateway", + "--max-turns", + "1", + "--output-format", + "json", + ]) + .arg("--cwd") + .arg(workdir.path()) + .current_dir(workdir.path()) + .stdin(std::process::Stdio::null()) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .kill_on_drop(true); + xai_grok_test_support::env::test_env_cmd_tokio(&mut cmd, &server.url(), home.path()); + cmd.env_remove("GROK_LEADER_SOCKET"); + + // The turn is expected to fail (the mock 401s the unauthenticated request); + // we assert on the wire, not the exit code. + let _ = run_headless_with_cmd(cmd).await; + + let requests = server.requests(); + // Non-vacuity: the model was actually exercised. + assert!( + requests + .iter() + .any(|e| e.method == "POST" && e.path.contains("chat/completions")), + "no chat request attempted; requests:\n{}", + server.request_log_summary() + ); + assert!( + !requests.iter().any(|e| { + e.path.contains("chat/completions") + && e.authorization + .as_deref() + .is_some_and(|a| a.contains("test-key-for-ci")) + }), + "an undefined provider must fail closed, never sending the session \ + credential; requests:\n{}", + server.request_log_summary() + ); +} + +/// The documented `args` + JSON-output shape, end to end: a provider with +/// `args = [...]` runs the helper directly (no shell) and parses a JSON +/// `{access_token, expires_in}` payload, and the minted token reaches the wire. +#[tokio::test] +#[ignore] +async fn provider_with_args_and_json_output_sends_minted_token() { + let server = MockInferenceServer::start() + .await + .expect("start mock server"); + let workdir = git_workdir(); + let home = tempfile::TempDir::new().unwrap(); + + let grok_home = home.path().join(".grok"); + std::fs::create_dir_all(&grok_home).expect("create .grok home"); + + // The helper records the args it was invoked with (proving direct exec, no + // shell) and prints a JSON token payload. + let seen_args = grok_home.join("seen-args"); + let helper = grok_home.join("mock-auth-json.sh"); + std::fs::write( + &helper, + format!( + "#!/bin/sh\nprintf '%s' \"$*\" > {}\n\ + printf '{{\"access_token\":\"gateway-tok-json\",\"expires_in\":3600}}'\n", + seen_args.display() + ), + ) + .expect("write mock auth binary"); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(&helper, std::fs::Permissions::from_mode(0o755)).unwrap(); + } + + std::fs::write( + grok_home.join("config.toml"), + format!( + r#"[auth_provider.gateway] +command = "{helper}" +args = ["--profile", "corp"] +token_ttl_secs = 3600 +timeout_secs = 10 + +[model.proxied-gateway] +model = "mock-gateway-model" +base_url = "{base}" +context_window = 200000 +auth_provider = "gateway" +"#, + helper = helper.display(), + base = server.url(), + ), + ) + .expect("write config.toml"); + + let mut cmd = tokio::process::Command::new(grok_binary()); + cmd.args([ + "-p", + "say hi", + "--yolo", + "--model", + "proxied-gateway", + "--max-turns", + "1", + "--output-format", + "json", + ]) + .arg("--cwd") + .arg(workdir.path()) + .current_dir(workdir.path()) + .stdin(std::process::Stdio::null()) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .kill_on_drop(true); + xai_grok_test_support::env::test_env_cmd_tokio(&mut cmd, &server.url(), home.path()); + cmd.env_remove("GROK_LEADER_SOCKET"); + + let result = run_headless_with_cmd(cmd).await; + assert_headless_success(&result, "auth provider args/json e2e", Some(&server)); + + let args = std::fs::read_to_string(&seen_args).expect("helper must have run"); + assert_eq!( + args, "--profile corp", + "args must be passed directly to the helper with no shell" + ); + + let requests = server.requests(); + let chat = requests + .iter() + .find(|e| e.method == "POST" && e.path.contains("chat/completions")) + .unwrap_or_else(|| { + panic!( + "no POST /v1/chat/completions request logged; requests:\n{}", + server.request_log_summary() + ) + }); + assert_eq!( + chat.authorization.as_deref(), + Some("Bearer gateway-tok-json"), + "the JSON access_token must reach the wire; requests:\n{}", + server.request_log_summary() + ); +} diff --git a/crates/codegen/xai-grok-shell/tests/test_leader_stdio_integration.rs b/crates/codegen/xai-grok-shell/tests/test_leader_stdio_integration.rs index 7537b08..6140db6 100644 --- a/crates/codegen/xai-grok-shell/tests/test_leader_stdio_integration.rs +++ b/crates/codegen/xai-grok-shell/tests/test_leader_stdio_integration.rs @@ -581,9 +581,10 @@ async fn test_runtime_profile_start_status_stop_across_clients() { return; // pprof can't start in sandbox — skip }; assert!(matches!( - started, - ControlPayload::CpuProfileStarted { svg_path, .. } if svg_path == output_path - )); + started, + ControlPayload::CpuProfileStarted { svg_path, .. } + if svg_path == output_path + )); let status = client_b .send_control(ControlCommand::CpuProfileStatus) @@ -595,15 +596,16 @@ async fn test_runtime_profile_start_status_stop_across_clients() { "registration should stay consistent with status behavior" ); 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_b .send_control(ControlCommand::StopCpuProfile) @@ -611,9 +613,10 @@ async fn test_runtime_profile_start_status_stop_across_clients() { .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()); } else { let error = client_a @@ -726,9 +729,10 @@ async fn test_runtime_profile_creates_missing_parent_directory_end_to_end() { return; // pprof can't start in sandbox — skip }; assert!(matches!( - started, - ControlPayload::CpuProfileStarted { svg_path, .. } if svg_path == nested_output - )); + started, + ControlPayload::CpuProfileStarted { svg_path, .. } + if svg_path == nested_output + )); let stopped = client .send_control(ControlCommand::StopCpuProfile) @@ -736,9 +740,10 @@ async fn test_runtime_profile_creates_missing_parent_directory_end_to_end() { .unwrap() .unwrap(); assert!(matches!( - stopped, - ControlPayload::CpuProfileStopped { svg_path, .. } if svg_path == nested_output - )); + stopped, + ControlPayload::CpuProfileStopped { svg_path, .. } + if svg_path == nested_output + )); assert!(nested_output.exists()); } else { let error = client diff --git a/crates/codegen/xai-grok-test-support/src/env.rs b/crates/codegen/xai-grok-test-support/src/env.rs index d2113f8..c15f2a9 100644 --- a/crates/codegen/xai-grok-test-support/src/env.rs +++ b/crates/codegen/xai-grok-test-support/src/env.rs @@ -77,7 +77,13 @@ fn ensure_local_grok_binary(binary: &Path) { let cargo = std::env::var("CARGO").unwrap_or_else(|_| "cargo".to_string()); let output = Command::new(&cargo) .current_dir(workspace_root()) - .args(["build", "-p", "xai-grok-pager", "--bin", "xai-grok-pager"]) + .args([ + "build", + "-p", + "xai-grok-pager-bin", + "--bin", + "xai-grok-pager", + ]) .output() .unwrap_or_else(|e| panic!("failed to spawn {cargo} to build xai-grok-pager: {e}")); diff --git a/crates/codegen/xai-grok-tools-api/src/config_validation.rs b/crates/codegen/xai-grok-tools-api/src/config_validation.rs index a230437..9824b1b 100644 --- a/crates/codegen/xai-grok-tools-api/src/config_validation.rs +++ b/crates/codegen/xai-grok-tools-api/src/config_validation.rs @@ -170,9 +170,10 @@ mod tests { fn invalid_json_is_a_parse_error() { let err = parse_params_json(0, "t", Some("{not json")).unwrap_err(); assert!(matches!( - err.kind, - ToolConfigEntryErrorKind::ParamsJsonParse { raw, .. } if raw == "{not json" - )); + err.kind, + ToolConfigEntryErrorKind::ParamsJsonParse { raw, .. } + if raw == "{not json" + )); } #[test] @@ -206,9 +207,10 @@ mod tests { assert_eq!(err.field_path(), "tools[2].name_override"); assert!( matches!( - &err.kind, - ToolConfigEntryErrorKind::NameOverrideInvalid { name: n, .. } if n == name - ), + &err.kind, + ToolConfigEntryErrorKind::NameOverrideInvalid { name: n, .. } + if n == name + ), "name={name:?} kind={:?}", err.kind ); diff --git a/crates/codegen/xai-grok-tools/src/implementations/codex/apply_patch/tool.rs b/crates/codegen/xai-grok-tools/src/implementations/codex/apply_patch/tool.rs index d22fb30..4c9f236 100644 --- a/crates/codegen/xai-grok-tools/src/implementations/codex/apply_patch/tool.rs +++ b/crates/codegen/xai-grok-tools/src/implementations/codex/apply_patch/tool.rs @@ -23,7 +23,7 @@ use super::{apply, errors::ApplyPatchError}; // ─── Description ───────────────────────────────────────────────────── /// Tool description derived from the codex `apply_patch_tool_instructions.md`. -const DESCRIPTION: &str = r#"Use the `apply_patch` tool to edit files. +const DESCRIPTION: &str = r#"Use this tool to edit files. Your patch language is a stripped‑down, file‑oriented diff format designed to be easy to parse and safe to apply. You can think of it as a high‑level envelope: *** Begin Patch diff --git a/crates/codegen/xai-grok-tools/src/implementations/grok_build/bash/mod.rs b/crates/codegen/xai-grok-tools/src/implementations/grok_build/bash/mod.rs index b4d0791..279f930 100644 --- a/crates/codegen/xai-grok-tools/src/implementations/grok_build/bash/mod.rs +++ b/crates/codegen/xai-grok-tools/src/implementations/grok_build/bash/mod.rs @@ -277,9 +277,11 @@ pub struct BashToolInput { pub description: String, /// Set to true for long-running commands that should run in the background (e.g., dev servers, long builds). - /// Returns a task_id immediately while the command keeps running in the background; you are notified on completion, so do not poll or sleep-wait for it. + /// Returns a task id immediately while the command keeps running in the background; you are notified on completion, so do not poll or sleep-wait for it. + // "task id" stays plain English: the kill/get-output input params are + // renameable, so naming a literal key here goes stale after randomization. #[schemars( - description = "Set to true for long-running commands that should run in the background (e.g., dev servers, long builds). Returns a task_id immediately while the command keeps running in the background; you are notified on completion, so do not poll or sleep-wait for it." + description = "Set to true for long-running commands that should run in the background (e.g., dev servers, long builds). Returns a task id immediately while the command keeps running in the background; you are notified on completion, so do not poll or sleep-wait for it." )] #[serde( default, @@ -403,6 +405,43 @@ fn annotations(bash: &BashOutput) -> String { s } +const NOOP_END_TURN_REMINDER: &str = "\n\ + You appear to be running empty commands to stay active while waiting for background work. \ + End your turn — you will be woken automatically when there is something to do.\n\ + "; + +fn is_noop_command(command: &str) -> bool { + let trimmed = command.trim(); + trimmed.is_empty() || trimmed == "true" || trimmed == ":" || is_pure_status_print(trimmed) +} + +fn is_pure_status_print(trimmed: &str) -> bool { + if !(matches!(trimmed, "echo" | "printf") + || trimmed.starts_with("echo ") + || trimmed.starts_with("printf ")) + { + return false; + } + let mut in_single = false; + let mut in_double = false; + let mut chars = trimmed.chars(); + while let Some(c) = chars.next() { + match c { + '\\' if !in_single => { + chars.next(); + } + '\'' if !in_double => in_single = !in_single, + '"' if !in_single => in_double = !in_double, + '$' | '`' if !in_single => return false, + ';' | '&' | '|' | '<' | '>' | '(' | ')' | '\n' if !in_single && !in_double => { + return false; + } + _ => {} + } + } + true +} + /// Build the full DEFAULT prompt text from a `BashOutput`. /// /// - Normal: `exit: N [annotations]\n` @@ -438,7 +477,12 @@ pub(crate) fn format_default_prompt(bash: &BashOutput) -> String { Some(reason) => format!("exit: killed ({}){}", reason, annotations(bash)), None => format!("exit: {}{}", bash.exit_code, annotations(bash)), }; - format!("{}\n{}", header, output_str) + let prompt = format!("{}\n{}", header, output_str); + if bash.signal.is_none() && is_noop_command(&bash.command) { + format!("{}\n\n{}", prompt.trim_end(), NOOP_END_TURN_REMINDER) + } else { + prompt + } } } @@ -1336,9 +1380,33 @@ impl BashTool { Some(default_ms.min(budget_ms).max(1)) } + /// Background retrieval hint naming the get-output tool and its task-ids + /// param. Kind-wide resolution is correct here: this names *another* + /// tool's param (the get-output tool), not this bash tool's own schema + /// key — do not switch it to invoking-tool param names. + async fn background_retrieval_hint( + resources: &SharedResources, + task_id: &str, + ) -> Result { + let res = resources.lock().await; + let renderer = res.require::()?; + let get_task_name = renderer + .render("${{ tools.by_kind.background_task_action }}") + .unwrap_or_else(|_| "get_command_or_subagent_output".to_string()); + let task_ids_param = renderer + .param_for_kind(ToolKind::BackgroundTaskAction, "task_ids") + .unwrap_or("task_ids"); + Ok(format!( + "Use {get_task_name} tool with {task_ids_param}=[\"{task_id}\"] to retrieve the output." + )) + } + + /// Model-facing input schema. `timeout_param_name` is the client-facing + /// timeout field (canonical or alias) — must match the remapped key. fn exported_input_schema( input_schema: &serde_json::Value, params: &BashParams, + timeout_param_name: &str, ) -> serde_json::Value { let background_enabled = Self::background_enabled(params); let auto_bg = Self::auto_background_on_timeout_enabled(params); @@ -1353,19 +1421,21 @@ impl BashTool { let max_ms = Self::effective_max_timeout_ms(params); let default_ms = Self::effective_default_timeout_ms(params); // `max_timeout_secs` is a foreground-only ceiling; background - // `timeout: 0` is always unbounded, so this note is + // `{name}: 0` is always unbounded, so this note is // unconditional. - let bg_zero = "`timeout: 0` in background mode disables the wrapper timeout entirely; the task runs until it exits or is killed via the kill task tool."; + let bg_zero = format!( + "`{timeout_param_name}: 0` in background mode disables the wrapper timeout entirely; the task runs until it exits or is killed via the kill task tool." + ); // Keep main-style auto-bg wording (no FG-budget ms advertised). // Follow-up: surface effective_auto_bg_wait_ms / FG budget here // once we deliberately change model-facing copy. let desc = if auto_bg { format!( - "Optional timeout in milliseconds (max {max_ms}). Default: {default_ms}. If not specified, commands exceeding the default timeout will be automatically backgrounded. {bg_zero}" + "Optional {timeout_param_name} in milliseconds (max {max_ms}). Default: {default_ms}. If not specified, commands exceeding the default timeout will be automatically backgrounded. {bg_zero}" ) } else { format!( - "Optional timeout in milliseconds (max {max_ms}). Default: {default_ms}. {bg_zero}" + "Optional {timeout_param_name} in milliseconds (max {max_ms}). Default: {default_ms}. {bg_zero}" ) }; timeout_prop.insert("description".to_string(), serde_json::json!(desc)); @@ -1425,10 +1495,10 @@ impl BashTool { r#"Run a ${%- if is_windows %} shell command${%- else %} bash command${%- endif %} and return its output. Usage notes: - - You can specify an optional ${{ params.execute.timeout }} in milliseconds (up to ${{ max_timeout_ms | default(300000) }}ms). ${%- if auto_background_on_timeout %} If not specified, commands exceeding the default timeout will be automatically backgrounded instead of killed. You will receive a task_id to check output later.${%- else %} If not specified, commands will timeout after ${{ default_timeout_ms | default(120000) }}ms.${%- endif %} + - You can specify an optional ${{ params.execute.timeout }} in milliseconds (up to ${{ max_timeout_ms | default(300000) }}ms). ${%- if auto_background_on_timeout %} If not specified, commands exceeding the default timeout will be automatically backgrounded instead of killed. You will receive a task id to check output later.${%- else %} If not specified, commands will timeout after ${{ default_timeout_ms | default(120000) }}ms.${%- endif %} - Timeout enforcement: when the timeout fires, the wrapper${%- if is_windows %} terminates the child's Job Object, killing every descendant process immediately (no graceful-termination grace period).${%- else %} kills the child process group (SIGTERM, escalated to SIGKILL after a ~1s grace period). Descendants that did not detach via `setsid` / `nohup` will also be killed.${%- endif %} `${{ params.execute.timeout }}: 0` in `${%- if params is defined and params.execute is defined and params.execute.is_background %}${{ params.execute.is_background }}${%- else %}background${%- endif %}: true` mode disables the wrapper timeout entirely; the child's lifetime is owned by the model via ${{ tools.by_kind.kill_task_action }}. - If the output exceeds {max_output_bytes} characters, output will be truncated before being returned to you. - - You can use the ${{ params.execute.is_background }} parameter to run the command in the background (e.g., dev servers, long builds): it returns a task_id immediately and keeps running in the background. You are notified on completion, so do not poll or sleep-wait for it.${%- if has_unix_utilities %} You do not need to use '&' at the end of the command when using this parameter.${%- endif %} + - You can use the ${{ params.execute.is_background }} parameter to run the command in the background (e.g., dev servers, long builds): it returns a task id immediately and keeps running in the background. You are notified on completion, so do not poll or sleep-wait for it.${%- if has_unix_utilities %} You do not need to use '&' at the end of the command when using this parameter.${%- endif %} ${%- if shell_uses_semicolon %} - '&&' is not supported in this shell; chain sequential commands with ';'. ${%- endif %} @@ -1545,7 +1615,16 @@ impl crate::types::tool_metadata::ToolMetadata for BashTool { let params: BashParams = serde_json::from_value(effective_params.clone()).unwrap_or_default(); let description = Self::rendered_description(description_override, renderer, ¶ms); - let exported_schema = Self::exported_input_schema(input_schema, ¶ms); + // Only this tool's param_map renames schema property keys — do not + // fall back to kind-wide renderer aliases (another Execute tool's + // override could advertise e.g. max_wait while this schema still + // exposes timeout). + let timeout_param_name = param_map + .get("timeout") + .map(String::as_str) + .unwrap_or("timeout"); + let exported_schema = + Self::exported_input_schema(input_schema, ¶ms, timeout_param_name); let remapped_schema = if param_map.is_empty() { exported_schema } else { @@ -1987,13 +2066,7 @@ impl xai_tool_runtime::Tool for BashTool { description: Some(input.description.clone()), }); - // Build the retrieval hint with the resolved tool name; the task id - // is passed as a single-element `task_ids` array (the only arg). - let __res = resources.lock().await; - let renderer = __res.require::()?; - let get_task_name = renderer - .render("${{ tools.by_kind.background_task_action }}") - .unwrap_or_else(|_| "get_command_or_subagent_output".to_string()); + let retrieval_hint = Self::background_retrieval_hint(&resources, &task_id).await?; Ok(BashToolOutput::Background(BackgroundTaskStarted { task_id: task_id.clone(), @@ -2002,10 +2075,7 @@ impl xai_tool_runtime::Tool for BashTool { status: "running".to_string(), command: input.command, summary: format!("Background task {} started", task_id), - retrieval_hint: format!( - "Use {} tool with task_ids=[\"{}\"] to retrieve the output.", - get_task_name, task_id - ), + retrieval_hint, pre_formatted: None, pid: bg_pid, })) @@ -2090,13 +2160,8 @@ impl xai_tool_runtime::Tool for BashTool { description: Some(input.description.clone()), }); - // Build the retrieval hint with the resolved tool name; the task - // id is passed as a single-element `task_ids` array (the only arg). - let __res = resources.lock().await; - let renderer = __res.require::()?; - let get_task_name = renderer - .render("${{ tools.by_kind.background_task_action }}") - .unwrap_or_else(|_| "get_command_or_subagent_output".to_string()); + let retrieval_hint = + Self::background_retrieval_hint(&resources, tool_call_id.as_str()).await?; let summary = if auto_backgrounded { format!( @@ -2118,11 +2183,7 @@ impl xai_tool_runtime::Tool for BashTool { status: "running".to_string(), command: input.command, summary, - retrieval_hint: format!( - "Use {} tool with task_ids=[\"{}\"] to retrieve the output.", - get_task_name, - tool_call_id.as_str() - ), + retrieval_hint, pre_formatted: None, // Real PID from the foreground spawn surfaced via // `TerminalRunResult::pid`. Adapters rely on this @@ -3260,12 +3321,15 @@ mod tests { #[tokio::test] async fn tool_name_mapping_in_background_hint() { let mut resources = make_resources(MockTerminal::background_ok("t1")); - // Custom model-facing tool name. The task id is always passed via the - // canonical single-element `task_ids` array (the param name is not - // overridable in the hint, matching the subagent-started footers). + // Custom model-facing tool AND param names — the hint must track both + // (a hardcoded `task_ids` goes stale after randomization renames). resources.insert(TemplateRenderer::new( [(ToolKind::BackgroundTaskAction, "GetOutput".to_string())].into(), - HashMap::new(), + [( + ToolKind::BackgroundTaskAction, + HashMap::from([("task_ids".to_string(), "jobs".to_string())]), + )] + .into(), )); let tool = BashTool; @@ -3285,8 +3349,8 @@ mod tests { bg.retrieval_hint ); assert!( - bg.retrieval_hint.contains("task_ids=[\"t1\"]"), - "Hint should pass the task id via a single-element task_ids array: {}", + bg.retrieval_hint.contains("jobs=[\"t1\"]"), + "Hint should pass the task id via the renamed task_ids param: {}", bg.retrieval_hint ); } @@ -3303,7 +3367,7 @@ mod tests { output: output.as_bytes().to_vec(), output_for_prompt: BashOutput::make_output_for_prompt(output), exit_code, - command: "echo test".to_string(), + command: "cat test".to_string(), truncated: false, signal: None, timed_out: false, @@ -3453,6 +3517,68 @@ mod tests { ); } + fn bash_output_with_command(command: &str, output: &str) -> BashOutput { + BashOutput { + output: output.as_bytes().to_vec(), + output_for_prompt: BashOutput::make_output_for_prompt(output), + exit_code: 0, + command: command.to_string(), + truncated: false, + signal: None, + timed_out: false, + description: None, + current_dir: "/tmp".to_string(), + output_file: String::new(), + total_bytes: output.len(), + output_delta: None, + was_bare_echo: false, + } + } + + #[test] + fn default_prompt_noop_command_appends_end_turn_reminder() { + for cmd in [ + "true", + ":", + "", + " ", + "\t\n", + "echo ok", + "echo \"Healthy.\"", + "echo \"s14=198; s11 full. Healthy.\"", + "printf hi", + "printf 'done\\n'", + ] { + let prompt = format_default_prompt(&bash_output_with_command(cmd, "")); + assert!( + prompt.contains(NOOP_END_TURN_REMINDER), + "no-op command {cmd:?} should append the end-turn reminder, got: {prompt:?}" + ); + } + } + + #[test] + fn default_prompt_normal_command_has_no_end_turn_reminder() { + for cmd in [ + "true && echo hi", + "run-true", + "grep : file", + "cat file", + "echo $VAR", + "echo x > f", + "echo a | cat", + "echo $(date)", + "echo hi; ls", + "printf '%s' \"$x\"", + ] { + let prompt = format_default_prompt(&bash_output_with_command(cmd, "hi\n")); + assert!( + !prompt.contains(""), + "normal command {cmd:?} must not append the end-turn reminder, got: {prompt:?}" + ); + } + } + // ─── contains_background_operator unit tests ─── mod background_operator_tests { @@ -4002,7 +4128,7 @@ mod tests { } fn timeout_desc(params: &BashParams) -> String { - let schema = BashTool::exported_input_schema(&base_schema(), params); + let schema = BashTool::exported_input_schema(&base_schema(), params, "timeout"); schema["properties"]["timeout"]["description"] .as_str() .expect("timeout description") @@ -4110,7 +4236,7 @@ mod tests { desc.contains("Default: 30000") || desc.contains("30000"), "default must track config: {desc}" ); - let schema = BashTool::exported_input_schema(&base_schema(), ¶ms); + let schema = BashTool::exported_input_schema(&base_schema(), ¶ms, "timeout"); assert_eq!( schema["properties"]["timeout"]["maximum"].as_u64(), Some(60_000) @@ -4133,6 +4259,78 @@ mod tests { ); } + /// Property description must track rename after `remap_schema_properties` + /// (regression: stale `` `timeout: 0` `` under `properties.`). + #[test] + fn schema_property_description_tracks_renamed_timeout() { + let param_map = + std::collections::HashMap::from([("timeout".to_string(), "max_wait".to_string())]); + let exported = + BashTool::exported_input_schema(&base_schema(), &BashParams::default(), "max_wait"); + let remapped = crate::util::remap::remap_schema_properties(&exported, ¶m_map); + let desc = remapped["properties"]["max_wait"]["description"] + .as_str() + .expect("max_wait description"); + assert!( + desc.contains("Optional max_wait in milliseconds") + && desc.contains("`max_wait: 0`"), + "renamed timeout must appear in property description:\n{desc}" + ); + assert!( + !desc.contains("`timeout: 0`") + && !desc.contains("Optional timeout in milliseconds"), + "canonical timeout must not remain in property description:\n{desc}" + ); + } + + /// Kind-wide renderer aliases must not rewrite this tool's property + /// description when this tool's own param_map did not rename timeout — + /// schema keys only follow param_map. + #[test] + fn schema_property_description_ignores_kind_wide_timeout_alias() { + use crate::types::tool_metadata::ToolMetadata; + + let renderer = TemplateRenderer::new( + HashMap::from([(ToolKind::Execute, "run_terminal_cmd".to_string())]), + HashMap::from([( + ToolKind::Execute, + // Another Execute tool (or identity-seed collision) renamed + // timeout kind-wide; this bash tool's param_map is empty. + HashMap::from([("timeout".to_string(), "max_wait".to_string())]), + )]), + ); + let def = ToolMetadata::versioned_definition( + &BashTool, + None, + "run_terminal_cmd", + None, + &renderer, + &HashMap::new(), + &base_schema(), + &serde_json::json!({}), + ); + let props = def + .function + .parameters + .get("properties") + .expect("properties"); + assert!( + props.get("timeout").is_some() && props.get("max_wait").is_none(), + "empty param_map must keep schema key timeout, got: {props}" + ); + let desc = props["timeout"]["description"] + .as_str() + .expect("timeout description"); + assert!( + desc.contains("Optional timeout in milliseconds") && desc.contains("`timeout: 0`"), + "property description must match schema key, not kind-wide alias:\n{desc}" + ); + assert!( + !desc.contains("max_wait"), + "kind-wide alias must not leak into property description:\n{desc}" + ); + } + #[test] fn tool_description_timeout_numbers_track_config() { let params = BashParams { diff --git a/crates/codegen/xai-grok-tools/src/implementations/grok_build/read_file/mod.rs b/crates/codegen/xai-grok-tools/src/implementations/grok_build/read_file/mod.rs index 81c6196..b6db3c1 100644 --- a/crates/codegen/xai-grok-tools/src/implementations/grok_build/read_file/mod.rs +++ b/crates/codegen/xai-grok-tools/src/implementations/grok_build/read_file/mod.rs @@ -299,6 +299,7 @@ pub(crate) async fn run_read_file( contract_version: Option<&str>, resources: SharedResources, streamable_out: Option<&mut bool>, + invoking_param_names: &crate::types::resources::InvokingToolParamNames, ) -> Result { let (cwd, display_cwd, fs, hints_enabled); { @@ -473,11 +474,13 @@ pub(crate) async fn run_read_file( .render("${{ tools.by_kind.execute }}") .map_err(|e| xai_tool_runtime::ToolError::invalid_arguments(e.to_string()))?; } + let offset_param = invoking_param_names.resolve("offset"); + let limit_param = invoking_param_names.resolve("limit"); let single_content_line = extracted.raw_output.lines().count() <= 1; let single_line_hint = if single_content_line && !execute_name.is_empty() { format!( "\nNote: the requested read is a single very long line, so \ - line-based offset/limit cannot narrow it further. Use the \ + line-based {offset_param}/{limit_param} cannot narrow it further. Use the \ '{execute_name}' tool to extract the parts you need (e.g. \ `jq`, `python3`, or `cut -c`)." ) @@ -493,18 +496,16 @@ pub(crate) async fn run_read_file( .limit .map_or_else(|| "to end".to_string(), |v| v.to_string()); format!( - "The requested line range (offset={}, limit={}) contains {} tokens, \ - which exceeds the maximum allowed tokens ({} tokens).\n\ - Try a smaller `limit`, a different starting `offset`, \ - or use the '{}' tool to search for specific content.{}", - off, lim, token_count, MAX_NUM_TOKENS, grep_name, single_line_hint + "The requested line range ({offset_param}={off}, {limit_param}={lim}) contains {token_count} tokens, \ + which exceeds the maximum allowed tokens ({MAX_NUM_TOKENS} tokens).\n\ + Try a smaller `{limit_param}`, a different starting `{offset_param}`, \ + or use the '{grep_name}' tool to search for specific content.{single_line_hint}" ) } else { format!( - "File content ({} tokens) exceeds maximum allowed tokens ({} tokens).\n\ - Please use offset and limit parameters to read a shorter range, \ - or use the '{}' to search for specific content.{}", - token_count, MAX_NUM_TOKENS, grep_name, single_line_hint + "File content ({token_count} tokens) exceeds maximum allowed tokens ({MAX_NUM_TOKENS} tokens).\n\ + Please use {offset_param} and {limit_param} parameters to read a shorter range, \ + or use the '{grep_name}' to search for specific content.{single_line_hint}" ) }; return Ok(ReadFileOutput::FileTooLarge(msg)); @@ -601,21 +602,22 @@ impl xai_tool_runtime::Tool for ReadFileTool { }); }; Box::pin(async_stream::stream! { - match ReadFileTool::read_with_streamability(& ctx, input). await { - Ok((output, streamable)) => { if streamable && let - ReadFileOutput::FileContent(fc) = & output && ! fc.content.is_empty() { - let content = fc.content.as_bytes(); let mut last_total : u64 = 0; let - mut window_start = 0usize; while window_start < content.len() { let mut - window_end = (window_start + STREAM_DELTA_TARGET_BYTES).min(content - .len()); while window_end > window_start && ! fc.content - .is_char_boundary(window_end) { window_end -= 1; } if let Some(p) = - xai_tool_runtime::stream_chunk(spec, & content[..window_end], window_end - as u64, & mut last_total, false,) { yield - xai_tool_runtime::ToolStreamItem::Progress(p); } window_start = - window_end; } } yield - xai_tool_runtime::ToolStreamItem::Terminal(Ok(output)); } Err(e) => yield - xai_tool_runtime::ToolStreamItem::Terminal(Err(e)), } - }) + match ReadFileTool::read_with_streamability(& ctx, input). await { + Ok((output, streamable)) => { if streamable && let + ReadFileOutput::FileContent(fc) = & output && ! fc.content.is_empty() { + let content = fc.content.as_bytes(); let mut last_total : u64 = 0; let + mut window_start = 0usize; while window_start < content.len() { let mut + window_end = (window_start + STREAM_DELTA_TARGET_BYTES).min(content + .len()); while window_end > window_start && ! fc.content + .is_char_boundary(window_end) { window_end -= 1; } + if let Some(p) = + xai_tool_runtime::stream_chunk(spec, & content[..window_end], window_end + as u64, & mut last_total, false,) { yield + xai_tool_runtime::ToolStreamItem::Progress(p); } window_start = + window_end; } } yield + xai_tool_runtime::ToolStreamItem::Terminal(Ok(output)); } Err(e) => yield + xai_tool_runtime::ToolStreamItem::Terminal(Err(e)), } + }) } #[tracing::instrument(name = "tool.read_file", skip_all, fields(path = %input.path))] async fn run( @@ -643,12 +645,14 @@ impl ReadFileTool { .map(|c| c.0.clone()); let bv = crate::types::tool_metadata::behavior_version(ctx); let mut streamable_text = false; + let invoking = crate::types::tool_metadata::invoking_param_names(ctx); let output = run_read_file( input, cwd_override.clone(), bv.as_deref(), resources.clone(), Some(&mut streamable_text), + &invoking, ) .await?; Ok((output, streamable_text)) @@ -1007,6 +1011,63 @@ mod tests { other => panic!("Expected FileTooLarge, got {:?}", other), } } + /// Regression: FileTooLarge must name *this* tool's schema keys, not + /// whatever a sibling Read tool last wrote into the kind-wide param map. + #[tokio::test] + async fn token_limit_error_uses_invoking_tool_param_names_not_kind_wide() { + let tmp = TempDir::new().unwrap(); + let line = "x".repeat(200); + let big_content = std::iter::repeat_n(line.as_str(), 1100) + .collect::>() + .join("\n"); + std::fs::write(tmp.path().join("big.txt"), &big_content).unwrap(); + let tool = ReadFileTool; + let mut resources = test_resources(tmp.path()); + resources.insert(TemplateRenderer::new( + [(ToolKind::Search, "Grep".to_string())].into(), + [( + ToolKind::Read, + [ + ("offset".to_string(), "poisoned_offset".to_string()), + ("limit".to_string(), "poisoned_limit".to_string()), + ] + .into(), + )] + .into(), + )); + let input = ReadFileInput { + path: "big.txt".to_string(), + offset: Some(1), + limit: Some(800), + pages: None, + format: None, + }; + let mut ctx = test_ctx(resources.into_shared()); + ctx.extensions + .insert(crate::types::resources::InvokingToolParamNames( + [ + ("offset".to_string(), "start_line".to_string()), + ("limit".to_string(), "max_lines".to_string()), + ] + .into(), + )); + let result = xai_tool_runtime::Tool::run(&tool, ctx, input) + .await + .unwrap(); + match result { + ReadFileOutput::FileTooLarge(msg) => { + assert!( + msg.contains("start_line=1") && msg.contains("max_lines=800"), + "expected invoking-tool names, got: {msg}" + ); + assert!( + !msg.contains("poisoned_offset") && !msg.contains("poisoned_limit"), + "must not use kind-wide sibling renames: {msg}" + ); + } + other => panic!("Expected FileTooLarge, got {:?}", other), + } + } #[test] fn test_extract_file_content_lines_basic() { let extracted = extract_file_content_lines("1\n2\r\n3\n", None, None, 4); diff --git a/crates/codegen/xai-grok-tools/src/implementations/grok_build_concise/read_file.rs b/crates/codegen/xai-grok-tools/src/implementations/grok_build_concise/read_file.rs index 8613bbb..c719aee 100644 --- a/crates/codegen/xai-grok-tools/src/implementations/grok_build_concise/read_file.rs +++ b/crates/codegen/xai-grok-tools/src/implementations/grok_build_concise/read_file.rs @@ -6,7 +6,7 @@ const DESCRIPTION_CONCISE: &str = r#"Reads a file from the computer's filesystem It is okay to read a file that does not exist; an error will be returned. Usage: -- You can optionally specify a line offset and limit (especially handy for long files). +- You can optionally specify ${{ params.read.offset }} and ${{ params.read.limit }} (especially handy for long files). - Lines in the output are numbered starting at 1, using following format: LINE_NUMBER→LINE_CONTENT. - You have the capability to call multiple tools in a single response. It is always better to speculatively read multiple files as a batch that are potentially useful."#; use crate::types::output::ReadFileOutput; @@ -84,7 +84,8 @@ impl xai_tool_runtime::Tool for ReadFileConciseTool { .map(|c| c.0.clone()); // `None`: the concise tool does not stream, so it needs no // text-path streamability signal (see `run_read_file`). - let result = run_read_file(input, cwd_override, None, resources, None).await?; + let invoking = crate::types::tool_metadata::invoking_param_names(&ctx); + let result = run_read_file(input, cwd_override, None, resources, None, &invoking).await?; match result { ReadFileOutput::FileContent(mut fc) => { @@ -133,6 +134,33 @@ mod tests { ); } + #[test] + fn description_template_tracks_renamed_offset_limit() { + use crate::types::template_renderer::TemplateRenderer; + use crate::types::tool_metadata::ToolMetadata; + use std::collections::HashMap; + + let tools = HashMap::from([(ToolKind::Read, "read_file".to_string())]); + let params = HashMap::from([( + ToolKind::Read, + HashMap::from([ + ("offset".to_string(), "start_line".to_string()), + ("limit".to_string(), "num_lines".to_string()), + ]), + )]); + let rendered = TemplateRenderer::new(tools, params) + .render(ToolMetadata::description_template(&ReadFileConciseTool)) + .unwrap(); + assert!( + rendered.contains("start_line and num_lines"), + "renamed offset/limit must appear:\n{rendered}" + ); + assert!( + !rendered.contains("a line offset and limit"), + "canonical offset/limit must not remain after rename:\n{rendered}" + ); + } + #[tokio::test] async fn concise_mode_uses_concise_content() { let tmp = TempDir::new().unwrap(); diff --git a/crates/codegen/xai-grok-tools/src/implementations/grok_build_hashline/edit/mod.rs b/crates/codegen/xai-grok-tools/src/implementations/grok_build_hashline/edit/mod.rs index 64c1615..0957ef3 100644 --- a/crates/codegen/xai-grok-tools/src/implementations/grok_build_hashline/edit/mod.rs +++ b/crates/codegen/xai-grok-tools/src/implementations/grok_build_hashline/edit/mod.rs @@ -47,7 +47,7 @@ Operations (use the "op" field): "write" — Replace entire file content (no anchors needed). { "op": "write", "content": "full file content here" } -Batch edits: pass multiple operations in "edits". They are validated against the +Batch edits: pass multiple operations in "${{ params.edit.edits }}". They are validated against the pre-edit snapshot and applied atomically bottom-up — if any anchor fails validation, ALL edits in the batch are rejected (none are applied). Overlapping ranges are also rejected. @@ -461,6 +461,35 @@ mod tests { resources } + #[test] + fn description_template_tracks_renamed_edits() { + use crate::types::template_renderer::TemplateRenderer; + use crate::types::tool::ToolKind; + use crate::types::tool_metadata::ToolMetadata; + use std::collections::HashMap; + + let tools = HashMap::from([ + (ToolKind::Edit, "hashline_edit".to_string()), + (ToolKind::Read, "hashline_read".to_string()), + (ToolKind::Search, "hashline_grep".to_string()), + ]); + let params = HashMap::from([( + ToolKind::Edit, + HashMap::from([("edits".to_string(), "changes".to_string())]), + )]); + let rendered = TemplateRenderer::new(tools, params) + .render(ToolMetadata::description_template(&HashlineEditTool)) + .unwrap(); + assert!( + rendered.contains("pass multiple operations in \"changes\""), + "renamed edits param must appear:\n{rendered}" + ); + assert!( + !rendered.contains("in \"edits\""), + "canonical edits must not remain after rename:\n{rendered}" + ); + } + fn anchors_for(content: &str) -> Vec { use crate::implementations::grok_build_hashline::anchor::split_lines; use crate::implementations::grok_build_hashline::edit::apply::anchor_suffix; diff --git a/crates/codegen/xai-grok-tools/src/implementations/grok_build_hashline/read_file.rs b/crates/codegen/xai-grok-tools/src/implementations/grok_build_hashline/read_file.rs index 860a900..8dbde12 100644 --- a/crates/codegen/xai-grok-tools/src/implementations/grok_build_hashline/read_file.rs +++ b/crates/codegen/xai-grok-tools/src/implementations/grok_build_hashline/read_file.rs @@ -86,8 +86,8 @@ use the fresh anchors returned by ${{ tools.by_kind.edit }} or re-read the file. Usage: - The ${{ params.read.target_file }} parameter must be an absolute path, not a relative path - By default reads up to {max_lines_read} lines from the beginning -- Optionally specify offset and limit for large files -- Can read images (PNG, JPG, etc.) and PDF files (each page rendered as an image; use `pages` parameter for PDFs with more than 10 pages, max 20 per call) +- Optionally specify ${{ params.read.offset }} and ${{ params.read.limit }} for large files +- Can read images (PNG, JPG, etc.) and PDF files (each page rendered as an image; use ${{ params.read.pages }} for PDFs with more than 10 pages, max 20 per call) - You can call multiple tools in a single response - If you read a file that exists but has empty contents you will receive a system reminder warning in place of file contents."#; @@ -188,7 +188,16 @@ impl xai_tool_runtime::Tool for HashlineReadTool { .map(|c| c.0.clone()); // `None`: the hashline tool does not stream, so it needs no // text-path streamability signal (see `run_read_file`). - let result = run_read_file(input, cwd_override, None, resources.clone(), None).await?; + let invoking = crate::types::tool_metadata::invoking_param_names(&ctx); + let result = run_read_file( + input, + cwd_override, + None, + resources.clone(), + None, + &invoking, + ) + .await?; match result { ReadFileOutput::FileContent(mut fc) => { @@ -374,6 +383,38 @@ mod tests { assert!(ToolMetadata::description_template(&hashline).contains("tools.by_kind.edit")); } + #[test] + fn description_template_tracks_renamed_offset_limit() { + use crate::types::template_renderer::TemplateRenderer; + use crate::types::tool::ToolKind; + use crate::types::tool_metadata::ToolMetadata; + use std::collections::HashMap; + + let tools = HashMap::from([ + (ToolKind::Read, "hashline_read".to_string()), + (ToolKind::Edit, "hashline_edit".to_string()), + ]); + let params = HashMap::from([( + ToolKind::Read, + HashMap::from([ + ("target_file".to_string(), "target_file".to_string()), + ("offset".to_string(), "start_line".to_string()), + ("limit".to_string(), "max_lines".to_string()), + ]), + )]); + let rendered = TemplateRenderer::new(tools, params) + .render(ToolMetadata::description_template(&HashlineReadTool)) + .unwrap(); + assert!( + rendered.contains("start_line and max_lines for large files"), + "renamed offset/limit must appear:\n{rendered}" + ); + assert!( + !rendered.contains("offset and limit for large files"), + "canonical offset/limit must not remain after rename:\n{rendered}" + ); + } + #[tokio::test] async fn read_basic_file() { let tmp = TempDir::new().unwrap(); diff --git a/crates/codegen/xai-grok-tools/src/implementations/lsp/types.rs b/crates/codegen/xai-grok-tools/src/implementations/lsp/types.rs index 53332be..1902b30 100644 --- a/crates/codegen/xai-grok-tools/src/implementations/lsp/types.rs +++ b/crates/codegen/xai-grok-tools/src/implementations/lsp/types.rs @@ -16,7 +16,8 @@ pub enum LspConfig { impl LspConfig { pub fn is_enabled(&self) -> bool { - matches!(self, Self::Enabled { servers, .. } if !servers.is_empty()) + matches!(self, Self::Enabled { servers, .. } +if !servers.is_empty()) } } diff --git a/crates/codegen/xai-grok-tools/src/implementations/opencode/bash/mod.rs b/crates/codegen/xai-grok-tools/src/implementations/opencode/bash/mod.rs index 3ba86db..0d077c1 100644 --- a/crates/codegen/xai-grok-tools/src/implementations/opencode/bash/mod.rs +++ b/crates/codegen/xai-grok-tools/src/implementations/opencode/bash/mod.rs @@ -64,7 +64,7 @@ Before executing the command, please follow these steps: - Capture the output of the command. Usage notes: - - The command argument is required. + - The ${{ params.execute.command }} argument is required. - You can specify an optional ${{ params.execute.timeout }} in milliseconds. If not specified, commands will use the default timeout. - It is very helpful if you write a clear, concise description of what this command does in 5-10 words. - If the output exceeds {max_output_bytes} characters, output will be truncated before being returned to you. @@ -111,7 +111,7 @@ Git Safety Protocol: - CRITICAL: If you already pushed to remote, NEVER amend unless user explicitly requests it (requires force push) - NEVER commit changes unless the user explicitly asks you to. It is VERY IMPORTANT to only commit when explicitly asked, otherwise the user will feel that you are being too proactive. -1. You can call multiple tools in a single response. When multiple independent pieces of information are requested and all commands are likely to succeed, run multiple tool calls in parallel for optimal performance. run the following bash commands in parallel, each using the Bash tool: +1. You can call multiple tools in a single response. When multiple independent pieces of information are requested and all commands are likely to succeed, run multiple tool calls in parallel for optimal performance. run the following bash commands in parallel, each using this tool: - Run a git status command to see all untracked files. - Run a git diff command to see both staged and unstaged changes that will be committed. - Run a git log command to see recent commit messages, so that you can follow this repository's commit message style. @@ -134,11 +134,11 @@ Important notes: - If there are no changes to commit (i.e., no untracked files and no modifications), do not create an empty commit # Creating pull requests -Use the gh command via the Bash tool for ALL GitHub-related tasks including working with issues, pull requests, checks, and releases. If given a GitHub URL use the gh command to get the information needed. +Use the gh command via this tool for ALL GitHub-related tasks including working with issues, pull requests, checks, and releases. If given a GitHub URL use the gh command to get the information needed. IMPORTANT: When the user asks you to create a pull request, follow these steps carefully: -1. You can call multiple tools in a single response. When multiple independent pieces of information are requested and all commands are likely to succeed, run multiple tool calls in parallel for optimal performance. run the following bash commands in parallel using the Bash tool, in order to understand the current state of the branch since it diverged from the main branch: +1. You can call multiple tools in a single response. When multiple independent pieces of information are requested and all commands are likely to succeed, run multiple tool calls in parallel for optimal performance. run the following bash commands in parallel using this tool, in order to understand the current state of the branch since it diverged from the main branch: - Run a git status command to see all untracked files - Run a git diff command to see both staged and unstaged changes that will be committed - Check if the current branch tracks a remote branch and is up to date with the remote, so you know if you need to push to the remote @@ -597,6 +597,35 @@ mod tests { ); } + #[test] + fn description_template_tracks_renamed_command() { + use crate::types::template_renderer::TemplateRenderer; + use crate::types::tool::ToolKind; + use crate::types::tool_metadata::ToolMetadata; + use std::collections::HashMap; + + let tools = HashMap::from([(ToolKind::Execute, "run_command".to_string())]); + let params = HashMap::from([( + ToolKind::Execute, + HashMap::from([ + ("command".to_string(), "script".to_string()), + ("timeout".to_string(), "timeout".to_string()), + ]), + )]); + let rendered = TemplateRenderer::new(tools, params) + .render(ToolMetadata::description_template(&BashTool)) + .unwrap(); + assert!( + rendered.contains("The script argument is required."), + "renamed command must appear:\n{rendered}" + ); + assert!( + !rendered.contains("The command argument is required.") + && !rendered.contains("Bash tool"), + "stale command/tool-name literals must not remain:\n{rendered}" + ); + } + fn make_input(command: &str) -> BashInput { BashInput { command: command.to_string(), diff --git a/crates/codegen/xai-grok-tools/src/implementations/opencode/read/mod.rs b/crates/codegen/xai-grok-tools/src/implementations/opencode/read/mod.rs index 0fc36cc..2041254 100644 --- a/crates/codegen/xai-grok-tools/src/implementations/opencode/read/mod.rs +++ b/crates/codegen/xai-grok-tools/src/implementations/opencode/read/mod.rs @@ -40,13 +40,13 @@ Assume this tool is able to read all files on the machine. If the User provides Usage: - The ${{ params.read.filePath }} parameter must be an absolute path, not a relative path - By default, it reads up to {max_lines_read} lines starting from the beginning of the file -- You can optionally specify a line offset and limit (especially handy for long files), but it's recommended to read the whole file by not providing these parameters +- You can optionally specify ${{ params.read.offset }} and ${{ params.read.limit }} (especially handy for long files), but it's recommended to read the whole file by not providing these parameters - Any lines longer than {max_chars_per_line} characters will be truncated - Results are returned using cat -n format, with line numbers starting at 1. The format is: LINE_NUMBER→LINE_CONTENT, where LINE_NUMBER is right-aligned and padded with spaces - This tool can read images (eg PNG, JPG, etc). When reading an image file the contents are presented visually as this tool uses multimodal LLMs. - This tool can read PDF files (.pdf). PDFs are processed page by page, extracting both text and visual content for analysis. - This tool can read Jupyter notebooks (.ipynb files) and returns all cells with their outputs, combining code, text, and visualizations. -- This tool can only read files, not directories. To read a directory, use an ls command via the Bash tool. +- This tool can only read files, not directories.${%- if tools.by_kind.execute %} To read a directory, use an ls command via the ${{ tools.by_kind.execute }} tool.${%- endif %} - You can call multiple tools in a single response. It is always better to speculatively read multiple potentially useful files in parallel. - You will regularly be asked to read screenshots. If the user provides a path to a screenshot, ALWAYS use this tool to view the file at the path. This tool will work with all temporary file paths. - If you read a file that exists but has empty contents you will receive a system reminder warning in place of file contents."#; @@ -147,16 +147,20 @@ impl xai_tool_runtime::Tool for ReadTool { ctx: xai_tool_runtime::ToolCallContext, input: ReadInput, ) -> Result { - use crate::types::tool_metadata::{resolve_cwd, shared_resources}; + use crate::types::tool_metadata::{invoking_param_names, resolve_cwd, shared_resources}; let resources = shared_resources(&ctx)?; + // Client-facing `offset` name for runtime "read beyond…" hints; a + // rename must not tell the model to pass a key its schema lacks. + let invoking = invoking_param_names(&ctx); + let offset_param = invoking.resolve("offset"); // ── Validate offset ───────────────────────────────────────── if let Some(offset) = input.offset && offset < 1 { - return Ok(ReadFileOutput::FileReadError( - "offset must be >= 1".to_string(), - )); + return Ok(ReadFileOutput::FileReadError(format!( + "{offset_param} must be >= 1" + ))); } // ── Resolve path (single lock acquisition) ───────────────── @@ -188,7 +192,7 @@ impl xai_tool_runtime::Tool for ReadTool { // BRANCH A: DIRECTORY // ═══════════════════════════════════════════════════════════ if metadata.is_dir() { - return Ok(read_directory(&path, input.offset, input.limit).await); + return Ok(read_directory(&path, input.offset, input.limit, offset_param).await); } // ═══════════════════════════════════════════════════════════ @@ -265,8 +269,7 @@ impl xai_tool_runtime::Tool for ReadTool { // Validate offset against file size. if total_lines > 0 && start >= total_lines { return Ok(ReadFileOutput::FileReadError(format!( - "Offset {} is out of range for this file ({} lines)", - offset, total_lines, + "{offset_param} {offset} is out of range for this file ({total_lines} lines)" ))); } @@ -329,16 +332,14 @@ impl xai_tool_runtime::Tool for ReadTool { let footer = if truncated_by_bytes { format!( - "\n\n(Output capped at 50 KB. Showing lines {}-{}. Use offset={} to continue.)", - offset, last_read_line, next_offset, + "\n\n(Output capped at 50 KB. Showing lines {offset}-{last_read_line}. Use {offset_param}={next_offset} to continue.)" ) } else if has_more_lines { format!( - "\n\n(Showing lines {}-{} of {}. Use offset={} to continue.)", - offset, last_read_line, total_lines, next_offset, + "\n\n(Showing lines {offset}-{last_read_line} of {total_lines}. Use {offset_param}={next_offset} to continue.)" ) } else { - format!("\n\n(End of file - total {} lines)", total_lines) + format!("\n\n(End of file - total {total_lines} lines)") }; let formatted = format!( @@ -370,6 +371,8 @@ async fn read_directory( path: &std::path::Path, offset: Option, limit: Option, + // Client-facing `offset` param name for the "read beyond…" hint. + offset_param: &str, ) -> ReadFileOutput { let mut entries = Vec::new(); @@ -423,11 +426,9 @@ async fn read_directory( let truncated = (start + shown) < total; let entries_footer = if truncated { + let beyond = offset_val + shown; format!( - "\n(Showing {} of {} entries. Use 'offset' parameter to read beyond entry {})", - shown, - total, - offset_val + shown, + "\n(Showing {shown} of {total} entries. Use the {offset_param} parameter to read beyond entry {beyond})" ) } else { format!("\n({} entries)", total) @@ -513,6 +514,115 @@ mod tests { resources } + #[test] + fn description_template_tracks_renamed_offset_limit_and_execute() { + use crate::types::template_renderer::TemplateRenderer; + use crate::types::tool::ToolKind; + use crate::types::tool_metadata::ToolMetadata; + use std::collections::HashMap; + + let tools = HashMap::from([ + (ToolKind::Read, "read".to_string()), + (ToolKind::Execute, "run_command".to_string()), + ]); + let params = HashMap::from([( + ToolKind::Read, + HashMap::from([ + ("filePath".to_string(), "filePath".to_string()), + ("offset".to_string(), "start_line".to_string()), + ("limit".to_string(), "max_lines".to_string()), + ]), + )]); + let rendered = TemplateRenderer::new(tools, params) + .render(ToolMetadata::description_template(&ReadTool)) + .unwrap(); + assert!( + rendered.contains("start_line and max_lines"), + "renamed offset/limit must appear:\n{rendered}" + ); + assert!( + rendered.contains("via the run_command tool"), + "resolved execute tool name must appear:\n{rendered}" + ); + assert!( + !rendered.contains("a line offset and limit") && !rendered.contains("Bash tool"), + "stale offset/limit/Bash-tool literals must not remain:\n{rendered}" + ); + } + + /// Runtime "read beyond…" footer must name this tool's client-facing + /// offset param, not the canonical `offset`, after a rename. + #[tokio::test] + async fn runtime_footer_tracks_renamed_offset() { + let tmp = TempDir::new().unwrap(); + let canonical_tmp = dunce::canonicalize(tmp.path()).unwrap(); + let file_path = canonical_tmp.join("big.txt"); + let content = (1..=100) + .map(|i| format!("line {i}")) + .collect::>() + .join("\n"); + std::fs::write(&file_path, &content).unwrap(); + + let resources = test_resources(&canonical_tmp); + let mut ctx = test_ctx(resources.into_shared()); + ctx.extensions + .insert(crate::types::resources::InvokingToolParamNames( + [("offset".to_string(), "start_line".to_string())].into(), + )); + + let input = ReadInput { + file_path: file_path.to_string_lossy().to_string(), + offset: None, + limit: Some(5), + }; + let result = xai_tool_runtime::Tool::run(&ReadTool, ctx, input) + .await + .unwrap(); + match result { + ReadFileOutput::FileContent(fc) => { + assert!( + fc.content.contains("Use start_line=6 to continue") + && !fc.content.contains("Use offset="), + "footer must use renamed offset param: {}", + fc.content + ); + } + other => panic!("Expected FileContent, got {other:?}"), + } + } + + /// The invalid-offset validation error must name this tool's client-facing + /// offset param, not the canonical `offset`, after a rename. (Fires before + /// path resolution, so no file is needed.) + #[tokio::test] + async fn validation_error_tracks_renamed_offset() { + let tmp = TempDir::new().unwrap(); + let resources = test_resources(tmp.path()); + let mut ctx = test_ctx(resources.into_shared()); + ctx.extensions + .insert(crate::types::resources::InvokingToolParamNames( + [("offset".to_string(), "start_line".to_string())].into(), + )); + + let input = ReadInput { + file_path: "whatever.txt".to_string(), + offset: Some(0), + limit: None, + }; + let result = xai_tool_runtime::Tool::run(&ReadTool, ctx, input) + .await + .unwrap(); + match result { + ReadFileOutput::FileReadError(msg) => { + assert!( + msg.contains("start_line must be >= 1") && !msg.contains("offset must"), + "validation error must use renamed offset param: {msg}" + ); + } + other => panic!("Expected FileReadError, got {other:?}"), + } + } + #[tokio::test] async fn read_text_file_basic() { let tmp = TempDir::new().unwrap(); diff --git a/crates/codegen/xai-grok-tools/src/registry/proto_convert.rs b/crates/codegen/xai-grok-tools/src/registry/proto_convert.rs index 31b5aaa..cb83bbc 100644 --- a/crates/codegen/xai-grok-tools/src/registry/proto_convert.rs +++ b/crates/codegen/xai-grok-tools/src/registry/proto_convert.rs @@ -140,9 +140,10 @@ mod tests { assert_eq!(err.tool_id, "GrokBuild:bash"); assert_eq!(err.field_path(), "tools[3].params_json"); assert!(matches!( - &err.kind, - ToolConfigEntryErrorKind::ParamsJsonParse { raw, .. } if raw == "{not json" - )); + &err.kind, + ToolConfigEntryErrorKind::ParamsJsonParse { raw, .. } + if raw == "{not json" + )); } #[test] @@ -180,9 +181,10 @@ mod tests { assert_eq!(err.field_path(), "tools[2].name_override"); assert!( matches!( - &err.kind, - ToolConfigEntryErrorKind::NameOverrideInvalid { name: n, .. } if n == name - ), + &err.kind, + ToolConfigEntryErrorKind::NameOverrideInvalid { name: n, .. } + if n == name + ), "name={name:?} kind={:?}", err.kind ); diff --git a/crates/codegen/xai-grok-tools/src/registry/types.rs b/crates/codegen/xai-grok-tools/src/registry/types.rs index c1cd579..a4912f9 100644 --- a/crates/codegen/xai-grok-tools/src/registry/types.rs +++ b/crates/codegen/xai-grok-tools/src/registry/types.rs @@ -1414,6 +1414,9 @@ impl FinalizedToolset { let mut ctx = xai_tool_runtime::ToolCallContext::new(parent_ctx.call_id.clone()); ctx.extensions.insert(self.resources.clone()); ctx.extensions.insert_arc(Arc::clone(&self.renderer)); + ctx.extensions.insert( + crate::types::resources::InvokingToolParamNames::from_reverse_params(&reverse_params), + ); if let Some(cwd) = parent_ctx.extensions.get::() { ctx.extensions.insert((*cwd).clone()); } @@ -1542,6 +1545,9 @@ impl FinalizedToolset { let mut ctx = xai_tool_runtime::ToolCallContext::new(rt_call_id); ctx.extensions.insert(self.resources.clone()); ctx.extensions.insert_arc(Arc::clone(&self.renderer)); + ctx.extensions.insert( + crate::types::resources::InvokingToolParamNames::from_reverse_params(&reverse_params), + ); if let Some(cwd) = cwd_override { ctx.extensions.insert(xai_tool_runtime::Cwd(cwd)); } diff --git a/crates/codegen/xai-grok-tools/src/types/resources.rs b/crates/codegen/xai-grok-tools/src/types/resources.rs index 42108b2..ed7c060 100644 --- a/crates/codegen/xai-grok-tools/src/types/resources.rs +++ b/crates/codegen/xai-grok-tools/src/types/resources.rs @@ -712,6 +712,35 @@ impl ParamNameMapping { .unwrap_or(canonical) } } +/// Canonical → client-facing param names for the tool currently executing. +/// +/// Stamped onto [`xai_tool_runtime::ToolCallContext::extensions`] by +/// `prepare_dispatch` / `call_raw` from that tool's own +/// `params_name_overrides`. Prefer this over kind-wide +/// [`crate::types::template_renderer::TemplateRenderer::param_for_kind`] when +/// naming params in that tool's own errors — multiple tools can share a +/// `ToolKind` with different renames, and the kind map is first/last-wins. +#[derive(Debug, Clone, Default)] +pub struct InvokingToolParamNames(pub HashMap); +impl InvokingToolParamNames { + /// Build from a client→canonical reverse map (the dispatch remap direction). + pub fn from_reverse_params(reverse_params: &HashMap) -> Self { + Self( + reverse_params + .iter() + .map(|(client, canonical)| (canonical.clone(), client.clone())) + .collect(), + ) + } + /// Resolve a canonical parameter name for the invoking tool. + /// Falls back to the canonical name if not in the map. + pub fn resolve<'a>(&'a self, canonical: &'a str) -> &'a str { + self.0 + .get(canonical) + .map(String::as_str) + .unwrap_or(canonical) + } +} /// Map of `ToolKind` → client-facing tool name. /// /// Built at finalize time from the enabled tools and client name overrides. @@ -1203,6 +1232,17 @@ mod tests { assert_eq!(mapping.resolve("other_tool", "old_string"), "old_string"); } #[test] + fn invoking_tool_param_names_from_reverse_and_resolve() { + let reverse = HashMap::from([ + ("start_line".to_string(), "offset".to_string()), + ("max_lines".to_string(), "limit".to_string()), + ]); + let names = InvokingToolParamNames::from_reverse_params(&reverse); + assert_eq!(names.resolve("offset"), "start_line"); + assert_eq!(names.resolve("limit"), "max_lines"); + assert_eq!(names.resolve("path"), "path"); + } + #[test] fn params_deref() { let p = Params(EditConfig { skip_read_before_edit: true, diff --git a/crates/codegen/xai-grok-tools/src/types/tool_metadata.rs b/crates/codegen/xai-grok-tools/src/types/tool_metadata.rs index d87f8b4..54b18a5 100644 --- a/crates/codegen/xai-grok-tools/src/types/tool_metadata.rs +++ b/crates/codegen/xai-grok-tools/src/types/tool_metadata.rs @@ -180,3 +180,19 @@ pub fn behavior_version(ctx: &xai_tool_runtime::ToolCallContext) -> Option() .map(|v| v.0.clone()) } + +/// This tool's own canonical→client param-name map, stamped on the dispatch +/// context by `prepare_dispatch` / `call_raw`. Returns an empty (identity) +/// map when absent — e.g. unit tests that call `Tool::run` directly — so +/// callers resolve to canonical names. Prefer this over kind-wide +/// [`crate::types::template_renderer::TemplateRenderer::param_for_kind`] when +/// naming *this* tool's own params (a sibling tool sharing the `ToolKind` +/// can rename the same field differently). +pub fn invoking_param_names( + ctx: &xai_tool_runtime::ToolCallContext, +) -> crate::types::resources::InvokingToolParamNames { + ctx.extensions + .get::() + .map(|arc| (*arc).clone()) + .unwrap_or_default() +} diff --git a/crates/codegen/xai-grok-workspace/src/file_system/git_status.rs b/crates/codegen/xai-grok-workspace/src/file_system/git_status.rs index f1ad39c..5e6f1c1 100644 --- a/crates/codegen/xai-grok-workspace/src/file_system/git_status.rs +++ b/crates/codegen/xai-grok-workspace/src/file_system/git_status.rs @@ -64,24 +64,22 @@ fn collapse_status_spaces(s: &str) -> String { out } -/// Short git status for the templated user message. -/// -/// Runs `git status --short --branch` and returns its output with consecutive -/// spaces collapsed via [`collapse_status_spaces`]: a leading `## ` -/// line followed by the file change list (or just `## ` on a clean -/// tree). This matches the body embedded in the `` block -/// byte-for-byte. pub async fn git_status_short(working_directory: impl Into) -> Result { let working_directory = working_directory.into(); tokio::task::spawn_blocking(move || { let output = xai_tty_utils::git_command() - .args(["status", "--short", "--branch"]) + .args(["status", "--short", "--branch", "--untracked-files=normal"]) .current_dir(&working_directory) .stdout(std::process::Stdio::piped()) .stderr(std::process::Stdio::null()) .output() - .map_err(|e| FsError::Other(format!("git status --short --branch failed: {}", e)))?; + .map_err(|e| { + FsError::Other(format!( + "git status --short --branch --untracked-files=normal failed: {}", + e + )) + })?; if !output.status.success() { return Err(FsError::Other(format!( diff --git a/crates/codegen/xai-grok-workspace/src/permission/resolution.rs b/crates/codegen/xai-grok-workspace/src/permission/resolution.rs index 50f08c4..72d78c6 100644 --- a/crates/codegen/xai-grok-workspace/src/permission/resolution.rs +++ b/crates/codegen/xai-grok-workspace/src/permission/resolution.rs @@ -2480,7 +2480,8 @@ mod tests { let (entries, logs) = parse_mcp_entries_capturing_logs(&json, "deniedMcpServers"); assert_eq!(entries.len(), 1); assert!( - matches!(&entries[0], AllowedMcpServer::Name { name } if name == "internal-only"), + matches!(&entries[0], AllowedMcpServer::Name { name } +if name == "internal-only"), "expected a Name entry, got {entries:?}" ); assert!( diff --git a/crates/codegen/xai-grok-workspace/src/permission/types.rs b/crates/codegen/xai-grok-workspace/src/permission/types.rs index ddd6e21..107a748 100644 --- a/crates/codegen/xai-grok-workspace/src/permission/types.rs +++ b/crates/codegen/xai-grok-workspace/src/permission/types.rs @@ -534,7 +534,8 @@ mod tests { }); let access = AccessKind::from(&input); assert!( - matches!(access, AccessKind::MCPTool { ref name, ref input } if name == + matches!(access, AccessKind::MCPTool { ref name, ref input } +if name == "linear__save_issue" && input["title"] == "test"), "UseTool should produce AccessKind::MCPTool carrying the inner tool name and args, got {access:?}" ); diff --git a/crates/codegen/xai-grok-workspace/src/upload/mod.rs b/crates/codegen/xai-grok-workspace/src/upload/mod.rs index 1973998..51dfd38 100644 --- a/crates/codegen/xai-grok-workspace/src/upload/mod.rs +++ b/crates/codegen/xai-grok-workspace/src/upload/mod.rs @@ -392,7 +392,8 @@ mod tests { let cfg = source.resolve(); assert_eq!(cfg.bucket_url.as_deref(), Some("gs://placeholder")); assert!( - matches!(& cfg.upload_method, UploadMethod::Proxy { proxy_base_url, .. } if + matches!(& cfg.upload_method, UploadMethod::Proxy { proxy_base_url, .. } +if proxy_base_url == "https://proxy.example/v1"), "resolve() must carry the proxy upload method + base url" ); diff --git a/crates/codegen/xai-hunk-tracker/src/actor/file_utils.rs b/crates/codegen/xai-hunk-tracker/src/actor/file_utils.rs index eac502a..d70ed44 100644 --- a/crates/codegen/xai-hunk-tracker/src/actor/file_utils.rs +++ b/crates/codegen/xai-hunk-tracker/src/actor/file_utils.rs @@ -197,9 +197,8 @@ mod tests { // Create content larger than MAX_TRACKED_TEXT_BYTES let large = vec![b'a'; MAX_TRACKED_TEXT_BYTES + 1]; let state = classify_bytes(&large); - assert!( - matches!(state, FileContentState::TooLarge { byte_len } if byte_len == MAX_TRACKED_TEXT_BYTES + 1) - ); + assert!(matches!(state, FileContentState::TooLarge { byte_len } +if byte_len == MAX_TRACKED_TEXT_BYTES + 1)); } #[test] @@ -232,9 +231,8 @@ mod tests { let large = vec![b'a'; MAX_TRACKED_TEXT_BYTES + 1]; std::fs::write(&path, &large).unwrap(); let state = read_file_bounded(&path).await; - assert!( - matches!(state, FileContentState::TooLarge { byte_len } if byte_len == MAX_TRACKED_TEXT_BYTES + 1) - ); + assert!(matches!(state, FileContentState::TooLarge { byte_len } +if byte_len == MAX_TRACKED_TEXT_BYTES + 1)); } // === LFS pointer tests === @@ -276,9 +274,8 @@ mod tests { let pointer = b"version https://git-lfs.github.com/spec/v1\noid sha256:abc123\nsize 12345\n"; let state = classify_bytes(pointer); - assert!( - matches!(state, FileContentState::LfsPointer { byte_len } if byte_len == pointer.len()) - ); + assert!(matches!(state, FileContentState::LfsPointer { byte_len } +if byte_len == pointer.len())); } #[test] @@ -287,7 +284,8 @@ mod tests { .to_string(); let len = pointer.len(); let state = classify_string(pointer); - assert!(matches!(state, FileContentState::LfsPointer { byte_len } if byte_len == len)); + assert!(matches!(state, FileContentState::LfsPointer { byte_len } +if byte_len == len)); } #[tokio::test] @@ -298,9 +296,8 @@ mod tests { b"version https://git-lfs.github.com/spec/v1\noid sha256:abc123\nsize 12345\n"; std::fs::write(&path, pointer).unwrap(); let state = read_file_bounded(&path).await; - assert!( - matches!(state, FileContentState::LfsPointer { byte_len } if byte_len == pointer.len()) - ); + assert!(matches!(state, FileContentState::LfsPointer { byte_len } +if byte_len == pointer.len())); } #[tokio::test] diff --git a/crates/common/xai-computer-hub-mcp-adapter/src/bridge.rs b/crates/common/xai-computer-hub-mcp-adapter/src/bridge.rs index 9047cb8..827231c 100644 --- a/crates/common/xai-computer-hub-mcp-adapter/src/bridge.rs +++ b/crates/common/xai-computer-hub-mcp-adapter/src/bridge.rs @@ -560,12 +560,10 @@ mod tests { match output { ToolOutputWire::Mcp { blocks } => { assert_eq!(blocks.len(), 2); - assert!( - matches!(&blocks[0], McpBlock::Text { text } if text == "result text") - ); - assert!( - matches!(&blocks[1], McpBlock::Image { mime_type, .. } if mime_type == "image/png") - ); + assert!(matches!(&blocks[0], McpBlock::Text { text } +if text == "result text")); + assert!(matches!(&blocks[1], McpBlock::Image { mime_type, .. } +if mime_type == "image/png")); } other => panic!("expected Mcp blocks, got {other:?}"), } diff --git a/crates/common/xai-computer-hub-sdk/src/notification.rs b/crates/common/xai-computer-hub-sdk/src/notification.rs index 028bb5b..88c6bbd 100644 --- a/crates/common/xai-computer-hub-sdk/src/notification.rs +++ b/crates/common/xai-computer-hub-sdk/src/notification.rs @@ -233,7 +233,8 @@ mod tests { }); let notif = HubNotification::parse(&value).expect("should parse as Unknown, not None"); assert!( - matches!(notif, HubNotification::Unknown { ref method, .. } if method == "tool.notification"), + matches!(notif, HubNotification::Unknown { ref method, .. } +if method == "tool.notification"), "tool.notification without envelope session_id should fall back to Unknown, got {notif:?}" ); } @@ -273,7 +274,8 @@ mod tests { }); let notif = HubNotification::parse(&value).expect("should parse as Unknown, not None"); assert!( - matches!(notif, HubNotification::Unknown { ref method, .. } if method == "tools_changed"), + matches!(notif, HubNotification::Unknown { ref method, .. } +if method == "tools_changed"), "malformed tools_changed should fall back to Unknown, got {notif:?}" ); } @@ -355,7 +357,8 @@ mod tests { }); let notif = HubNotification::parse(&value).expect("should parse as Unknown, not None"); assert!( - matches!(notif, HubNotification::Unknown { ref method, .. } if method == "tool.notification"), + matches!(notif, HubNotification::Unknown { ref method, .. } +if method == "tool.notification"), "malformed tool.notification should fall back to Unknown, got {notif:?}" ); } diff --git a/crates/common/xai-tool-protocol/tests/identifier_validation.rs b/crates/common/xai-tool-protocol/tests/identifier_validation.rs index 4cb930e..561cb73 100644 --- a/crates/common/xai-tool-protocol/tests/identifier_validation.rs +++ b/crates/common/xai-tool-protocol/tests/identifier_validation.rs @@ -53,7 +53,8 @@ fn tool_id_rejects_disallowed_characters() { ] { let err = ToolId::new(bad).unwrap_err(); assert!( - matches!(err, IdError::InvalidFormat { ref value } if value == bad), + matches!(err, IdError::InvalidFormat { ref value } +if value == bad), "expected InvalidFormat for {bad:?}, got {err:?}" ); } @@ -73,7 +74,8 @@ fn tool_id_rejects_empty_segments_around_separator() { for bad in [":foo", "foo:", ":"] { let err = ToolId::new(bad).unwrap_err(); assert!( - matches!(err, IdError::InvalidFormat { ref value } if value == bad), + matches!(err, IdError::InvalidFormat { ref value } +if value == bad), "expected InvalidFormat for {bad:?}, got {err:?}" ); } @@ -103,7 +105,8 @@ fn server_id_rejects_reserved_auto_prefix() { for bad in ["auto:my-server", "auto:", "auto:tool:read_file"] { let err = ServerId::new(bad).unwrap_err(); assert!( - matches!(err, IdError::ReservedPrefix { ref value } if value == bad), + matches!(err, IdError::ReservedPrefix { ref value } +if value == bad), "expected ReservedPrefix for {bad:?}, got {err:?}" ); } diff --git a/crates/common/xai-tool-runtime/src/render.rs b/crates/common/xai-tool-runtime/src/render.rs index 73e7d0d..93adb43 100644 --- a/crates/common/xai-tool-runtime/src/render.rs +++ b/crates/common/xai-tool-runtime/src/render.rs @@ -534,7 +534,8 @@ mod tests { ]); let blocks = extract_content_blocks(&v); assert_eq!(blocks.len(), 2); - assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "a")); + assert!(matches!(&blocks[0], ContentBlock::Text { text } +if text == "a")); assert!(matches!(&blocks[1], ContentBlock::Image { .. })); } @@ -652,7 +653,8 @@ mod tests { assert_eq!(blocks.len(), 2); // First block is the remainder text (field order in JSON objects // is not guaranteed, so just check it's Text and non-empty). - assert!(matches!(&blocks[0], ContentBlock::Text { text } if text.contains("summary"))); + assert!(matches!(&blocks[0], ContentBlock::Text { text } +if text.contains("summary"))); assert!(matches!(&blocks[1], ContentBlock::Image { .. })); } @@ -668,7 +670,8 @@ mod tests { let blocks = extract_content_blocks(&v); // results field → 2 blocks extracted, metadata → remainder text. assert_eq!(blocks.len(), 3); - assert!(matches!(&blocks[0], ContentBlock::Text { text } if text.contains("metadata"))); + assert!(matches!(&blocks[0], ContentBlock::Text { text } +if text.contains("metadata"))); assert_eq!(blocks[1], ContentBlock::Text { text: "a".into() }); assert_eq!(blocks[2], ContentBlock::Text { text: "b".into() }); } diff --git a/crates/common/xai-tool-runtime/tests/error_conversion.rs b/crates/common/xai-tool-runtime/tests/error_conversion.rs index 2fa6f41..2e0b872 100644 --- a/crates/common/xai-tool-runtime/tests/error_conversion.rs +++ b/crates/common/xai-tool-runtime/tests/error_conversion.rs @@ -45,16 +45,16 @@ fn invalid_arguments_round_trips_message_and_details() { fn not_found_maps_to_tool_not_found() { let err = ToolError::not_found(tid("missing"), "tool 'missing' not registered"); let wire: ToolErrorWire = err.into(); - assert!(matches!(wire, ToolErrorWire::ToolNotFound { tool_id } if tool_id == tid("missing"))); + assert!(matches!(wire, ToolErrorWire::ToolNotFound { tool_id } +if tool_id == tid("missing"))); } #[test] fn permission_denied_round_trips_reason() { let err = ToolError::permission_denied("not authorised for write"); let wire: ToolErrorWire = err.into(); - assert!( - matches!(wire, ToolErrorWire::PermissionDenied { reason } if reason == "not authorised for write") - ); + assert!(matches!(wire, ToolErrorWire::PermissionDenied { reason } +if reason == "not authorised for write")); } #[test] @@ -93,7 +93,8 @@ fn timeout_with_details() { fn cancelled_round_trips_tool_id() { let err = ToolError::cancelled(tid("paused"), "user cancelled"); let wire: ToolErrorWire = err.into(); - assert!(matches!(wire, ToolErrorWire::Cancelled { tool_id } if tool_id == tid("paused"))); + assert!(matches!(wire, ToolErrorWire::Cancelled { tool_id } +if tool_id == tid("paused"))); } #[test] diff --git a/prod/mc/cli-chat-proxy-types/src/deployment_config_types.rs b/prod/mc/cli-chat-proxy-types/src/deployment_config_types.rs index 1e7c298..e295aaf 100644 --- a/prod/mc/cli-chat-proxy-types/src/deployment_config_types.rs +++ b/prod/mc/cli-chat-proxy-types/src/deployment_config_types.rs @@ -4,16 +4,28 @@ use serde::{Deserialize, Serialize}; -/// The payload format version the server currently signs. Bump when the payload -/// gains semantics (e.g. an anti-replay counter or a key-fingerprint binding) so -/// verifiers can distinguish generations; `0` means a pre-versioned payload. -pub const SIGNED_PAYLOAD_VERSION: u32 = 1; +/// The payload format version the server currently signs. Informational for +/// now: no verifier gates on it (every version verifies the same); enforce a +/// minimum only once the fleet has rotated past older generations. +/// `0` = pre-versioned, `1` = first versioned payload, `2` = per-fetch `nonce`. +pub const SIGNED_PAYLOAD_VERSION: u32 = 2; /// Domain-separation tags inside the signed bytes: both message types share one /// signing key, so each verifier requires its own tag (no cross-substitution). pub const MANAGED_POLICY_TYP: &str = "grok.managed_policy.v1"; pub const MANAGED_IDENTITY_TYP: &str = "grok.managed_identity.v1"; +/// Client echoes its persisted envelope `nonce` on this header for the server probe. +pub const MANAGED_CONFIG_NONCE_ECHO_HEADER: &str = "x-grok-managed-config-nonce"; + +/// Shape of a server-minted nonce (16 random bytes as hex): what `fresh_nonce` +/// produces and the only shape the client echoes (hex is HTTP-header-safe). +/// Shared so a mint change cannot silently disable the echo: the proxy pins +/// its mint against this, the client gates its echo on it. +pub fn is_server_nonce_shape(nonce: &str) -> bool { + nonce.len() == 32 && nonce.bytes().all(|b| b.is_ascii_hexdigit()) +} + /// The exact bytes the server signs: the served policy, the principal it is /// bound to, and an expiry. Serialized once on the server and shipped verbatim /// as `signed_payload`, so the client verifies the received bytes directly @@ -42,6 +54,10 @@ pub struct SignedPayload { pub fail_closed: bool, /// Unix seconds after which the signature is no longer trusted. pub expires_at: u64, + /// Per-response nonce in the signed bytes (echoed on [`MANAGED_CONFIG_NONCE_ECHO_HEADER`]). + /// `default` empty keeps pre-nonce sidecars verifiable. + #[serde(default)] + pub nonce: String, /// Identifies the signing key, so a rotation can be distinguished. pub key_id: String, } @@ -144,6 +160,7 @@ mod tests { requirements: None, fail_closed: false, expires_at: 4_000_000_000, + nonce: "9f86d081884c7d6594a85abf0f0cf96b".into(), key_id: "v1".into(), }; let json = serde_json::to_string(&versioned).unwrap(); @@ -159,6 +176,12 @@ mod tests { legacy.typ, "", "an untagged payload parses (verifiers reject it)" ); + assert_eq!( + legacy.nonce, "", + "pre-nonce payloads default to an empty nonce" + ); + assert!(is_server_nonce_shape("0123456789abcdef0123456789abcdef")); + assert!(!is_server_nonce_shape("short")); } /// The claim round-trips; `fail_closed` is additive (absent → permissive). diff --git a/prod/mc/cli-chat-proxy-types/src/feedback_types.rs b/prod/mc/cli-chat-proxy-types/src/feedback_types.rs index 281533c..a57011f 100644 --- a/prod/mc/cli-chat-proxy-types/src/feedback_types.rs +++ b/prod/mc/cli-chat-proxy-types/src/feedback_types.rs @@ -442,6 +442,13 @@ pub struct FeedbackSubmission { /// Backend URL linking this feedback to its server-side session log. #[serde(default, skip_serializing_if = "Option::is_none")] pub unified_log_url: Option, + + /// Client-reported display name; unverified, never used for authorization. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub author_name: Option, + /// Client-reported email; unverified, never used for authorization. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub author_email: Option, } impl FeedbackSubmission { @@ -460,8 +467,8 @@ impl FeedbackSubmission { s } - /// Remove session context and model metadata, preserving only the - /// user's rating/text and essential identifiers (session_id, client_type). + /// Remove session context and model metadata. The author identity is + /// preserved so the author stays reachable after the strip. pub fn strip_metadata(&mut self) { self.model_id = None; self.resolved_model_id = None; diff --git a/prod/mc/cli-chat-proxy-types/src/metadata_types.rs b/prod/mc/cli-chat-proxy-types/src/metadata_types.rs index eaecf32..83b3531 100644 --- a/prod/mc/cli-chat-proxy-types/src/metadata_types.rs +++ b/prod/mc/cli-chat-proxy-types/src/metadata_types.rs @@ -36,7 +36,8 @@ use serde::{Deserialize, Serialize}; /// has no configured effort. /// v1.23: Removed `prompt`, `full_prompt`, and `truncated_prompt_local_path` /// from metadata.json (prompt content is no longer uploaded in metadata). -pub const GCS_SCHEMA_VERSION: &str = "v1.23"; +/// v1.24: Prompt metadata updates. +pub const GCS_SCHEMA_VERSION: &str = "v1.24"; /// OS-level sandbox state for a trace turn (local `xai-grok-sandbox`, not cloud sandbox). #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct LocalSandboxTelemetry {