Synced from monorepo

Synced from monorepo

Changes:
- Temporarily disable session share link creation in the TUI
- Do not approve plan on empty Enter from the revise prompt
- Expose chat product Skills via ACP available_commands_update
- Return immediately from a blocking wait on an already-completed ACP task
- Split headless pager module for clearer structure
- Stop git worktree prune from removing user registrations on resume
- Use compaction sampler tokenizer for item token counts
- Opt-in extra root CAs via GROK_EXTRA_CA_BUNDLE
- Cancel all session subagents when the user stops
- Let the session persistence actor exit when its session ends
- Make fullscreen terminal resize much cheaper on long sessions
- Report honestly from kill_task when an ACP task does not exist
- Hide /usage for external-auth deployments
- Forward the history-load trailer’s computer_reason to the client
- Remove ineffective no-op tool reminder
- Declare slash-command screen-mode support in one place
- Keep settings enum picker on the committed value until Enter
- Reap a PTY’s full process tree
- Stream tool calls from headless mode over ACP
- Bridge gateway task lifecycle to ACP for chat session background tasks
- Don’t warn about truncated history on a suppressed replay
- Fit full-replace summarizer input and recover on context-length errors
- Stop dropping agents over an unrecognized frontmatter color
- Add /undo as a slash alias for /rewind
- Harden sleep/wake token-refresh paths against forced re-login
- Add session/list ACP method
- Give each sampling backend its own conversion module
- Treat an unenrolled child process as a lint error
- Suppress the cancelled marker on send-now wake turns
- Stop tearing down Roslyn on every edit, and read C# diagnostics

Source-Revision: 2a28b4a86cfc4a4c133c35b7fc2a6a9964387c39

crates/codegen/xai-grok-auth/src/lib.rs

@@ -10,5 +10,5 @@ pub mod visibility;
 
 pub use auth_provider::{AuthCredentialProvider, CredentialSnapshot, StaticAuthCredentialProvider};
 #[cfg(feature = "middleware")]
-pub use retry_middleware::AuthRetryMiddleware;
+pub use retry_middleware::{AuthRetryMiddleware, StampedBearerSuffix, execute_with_stamp};
 pub use visibility::HttpAuth;

crates/codegen/xai-grok-auth/src/retry_middleware.rs

@@ -8,6 +8,47 @@ use reqwest_middleware::{Error, Middleware, Next};
 
 use crate::AuthCredentialProvider;
 
+/// Tail fragment (last [`STAMPED_BEARER_SUFFIX_LEN`] chars) of the bearer
+/// this middleware stamped, recorded into the request's `http::Extensions`
+/// at stamp time. 401-attribution sites read it back via
+/// [`execute_with_stamp`] instead of re-resolving at record time, which
+/// races with the refresh the 401 itself triggers. Absent ⇒ nothing was
+/// stamped; a retry overwrites it, so it always describes the attempt whose
+/// response the caller holds. Only the tail crosses this boundary — JWT
+/// heads are a shared constant, and the tail is safe for sinks to log.
+#[derive(Clone, Debug)]
+pub struct StampedBearerSuffix(pub String);
+
+/// Length of [`StampedBearerSuffix`]. Matches `token_suffix` in
+/// xai-grok-shell (the comparison site for 401 attribution).
+const STAMPED_BEARER_SUFFIX_LEN: usize = 12;
+
+/// Last [`STAMPED_BEARER_SUFFIX_LEN`] chars, counting chars from the end
+/// so a non-ASCII credential cannot cause a byte-boundary panic.
+fn bearer_suffix(token: &str) -> &str {
+    match token
+        .char_indices()
+        .rev()
+        .nth(STAMPED_BEARER_SUFFIX_LEN - 1)
+    {
+        Some((i, _)) => &token[i..],
+        None => token,
+    }
+}
+
+/// Execute `req` on a middleware-wrapped client and return the response
+/// together with the [`StampedBearerSuffix`] the auth middleware recorded
+/// (if it stamped anything). The one blessed way for 401-attribution
+/// call sites to learn what was actually sent on the wire.
+pub async fn execute_with_stamp(
+    client: &reqwest_middleware::ClientWithMiddleware,
+    req: Request,
+) -> reqwest_middleware::Result<(Response, Option<StampedBearerSuffix>)> {
+    let mut ext = http::Extensions::new();
+    let resp = client.execute_with_extensions(req, &mut ext).await?;
+    Ok((resp, ext.get::<StampedBearerSuffix>().cloned()))
+}
+
 pub struct AuthRetryMiddleware {
     credentials: Arc<dyn AuthCredentialProvider>,
     max_retries: u32,
@@ -22,11 +63,12 @@ impl AuthRetryMiddleware {
     }
 }
 
-fn apply_auth_header(req: &mut Request, token: &str) {
+fn apply_auth_header(req: &mut Request, token: &str, extensions: &mut http::Extensions) {
     match HeaderValue::from_str(&format!("Bearer {token}")) {
         Ok(val) => {
             req.headers_mut()
                 .insert(reqwest::header::AUTHORIZATION, val);
+            extensions.insert(StampedBearerSuffix(bearer_suffix(token).to_string()));
         }
         Err(e) => {
             tracing::warn!(error = %e, "auth retry: failed to build Authorization header");
@@ -43,7 +85,7 @@ impl Middleware for AuthRetryMiddleware {
         next: Next<'_>,
     ) -> Result<Response, Error> {
         if let Some(ref token) = self.credentials.snapshot().token {
-            apply_auth_header(&mut req, token);
+            apply_auth_header(&mut req, token, extensions);
         }
 
         let backup = req.try_clone();
@@ -67,7 +109,7 @@ impl Middleware for AuthRetryMiddleware {
             let Some(mut retry) = backup.try_clone() else {
                 break;
             };
-            apply_auth_header(&mut retry, token);
+            apply_auth_header(&mut retry, token, extensions);
             last_resp = next.clone().run(retry, extensions).await?;
             if last_resp.status() != StatusCode::UNAUTHORIZED {
                 return Ok(last_resp);
@@ -251,6 +293,73 @@ mod tests {
         mock.assert_async().await;
     }
 
+    /// The stamp must describe the bearer of the attempt whose response
+    /// the caller holds: after a 401 → refresh → retry, that is the
+    /// FRESH token, not the stale one stamped on the first attempt.
+    #[tokio::test]
+    async fn execute_with_stamp_reports_last_stamped_bearer() {
+        let mut server = mockito::Server::new_async().await;
+        let m401 = server
+            .mock("GET", "/api")
+            .match_header("authorization", "Bearer stale-token")
+            .with_status(401)
+            .create_async()
+            .await;
+        let m200 = server
+            .mock("GET", "/api")
+            .match_header("authorization", "Bearer fresh-token")
+            .with_status(200)
+            .create_async()
+            .await;
+
+        let p = Arc::new(SimulatedAuthManager::simulated(
+            "stale-token",
+            "fresh-token",
+        ));
+        let client = build_client(p, 1).await;
+
+        let req = client.get(format!("{}/api", server.url())).build().unwrap();
+        let (resp, stamp) = execute_with_stamp(&client, req).await.unwrap();
+        assert_eq!(resp.status(), 200);
+        // ≤ 12 chars → the suffix is the whole token.
+        assert_eq!(stamp.expect("bearer was stamped").0, "fresh-token");
+        m401.assert_async().await;
+        m200.assert_async().await;
+    }
+
+    /// No credential ⇒ no stamp: attribution must see "nothing was sent",
+    /// not an empty string or a stale record.
+    #[tokio::test]
+    async fn execute_with_stamp_is_none_when_nothing_stamped() {
+        let mut server = mockito::Server::new_async().await;
+        let m = server
+            .mock("GET", "/")
+            .with_status(401)
+            .create_async()
+            .await;
+
+        let p = Arc::new(MockProvider::new(None, false));
+        let client = build_client(p, 0).await;
+
+        let req = client.get(server.url()).build().unwrap();
+        let (resp, stamp) = execute_with_stamp(&client, req).await.unwrap();
+        assert_eq!(resp.status(), 401);
+        assert!(stamp.is_none(), "no credential must mean no stamp");
+        m.assert_async().await;
+    }
+
+    #[test]
+    fn bearer_suffix_takes_char_safe_tail() {
+        assert_eq!(
+            bearer_suffix("eyJ0eXAiOiJh.head.tail-distinct"),
+            "ail-distinct"
+        );
+        assert_eq!(bearer_suffix("short"), "short");
+        assert_eq!(bearer_suffix(""), "");
+        // 13 multi-byte chars: a byte-index cut would land mid-char.
+        assert_eq!(bearer_suffix("ééééééééééééé"), "éééééééééééé");
+    }
+
     #[tokio::test]
     async fn test_max_retries_bounds_attempts() {
         let mut server = mockito::Server::new_async().await;