Publish harness and TUI open-source

initial sync from the monorepo

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

@@ -0,0 +1,118 @@
+//! Credential dependency-inversion seam for outbound HTTP made by the
+//! data-collector. Shell installs `ShellAuthCredentialProvider` wrapping
+//! `AuthManager` + `TokenRefresher`; data-collector code holds an
+//! `Arc<dyn AuthCredentialProvider>`.
+
+use reqwest::RequestBuilder;
+
+use crate::visibility::HttpAuth;
+
+/// Snapshot of the currently effective credentials. Used by callers
+/// that build their own header maps (the OTel OTLP exporter) or that
+/// need the bearer prefix for 401-attribution telemetry.
+#[derive(Clone, Debug, Default)]
+pub struct CredentialSnapshot {
+    /// Bearer token. `None` when no auth is configured (CI / `--api-key` headless).
+    pub token: Option<String>,
+    /// User identifier matching the bearer token's owner. `None` when no auth
+    /// is configured or when the underlying provider has no concept of user
+    /// identity (`StaticAuthCredentialProvider`). Read by the OTel layer to
+    /// populate the `user.id` resource attribute.
+    pub user_id: Option<String>,
+    /// Team identifier from OAuth. `None` for personal accounts or when
+    /// no auth is configured.
+    pub team_id: Option<String>,
+    /// `uuidv5(NAMESPACE_OID, deployment_key)`, set only for deployment-key auth.
+    pub deployment_id: Option<String>,
+    /// `uuidv5(NAMESPACE_OID, api_key)`, set only for `AuthMode::ApiKey`.
+    pub api_key_id: Option<String>,
+    /// Org id from the OIDC `organizationId` claim; `None` for personal / deployment-key auth.
+    pub organization_id: Option<String>,
+}
+
+/// Source of truth for outbound auth on data-collector requests.
+///
+/// Supertrait of `HttpAuth` so a single impl satisfies both this trait
+/// (refresh-aware snapshot + 401 recovery) and the visibility seam
+/// (header construction). Callers add headers via `HttpAuth::apply`.
+#[async_trait::async_trait]
+pub trait AuthCredentialProvider: HttpAuth + Send + Sync + 'static {
+    /// Return the current credential snapshot. Implementations should
+    /// issue a cheap disk re-read (`AuthManager::refresh`) before
+    /// snapshotting so callers see updates from sibling processes
+    /// (`grok-desktop`, `grok login`). The `token` field MUST mirror
+    /// the bearer that `HttpAuth::apply` would send on the wire so
+    /// 401-attribution prefixes match the actual request.
+    fn snapshot(&self) -> CredentialSnapshot;
+
+    /// Attempt to obtain a fresh token. Returns `true` if a different
+    /// token was obtained -- caller should retry the failed request once.
+    /// Returns `false` if no refresher is configured or refresh failed.
+    async fn refresh_after_unauthorized(&self) -> bool;
+
+    /// Whether `X-XAI-Token-Auth` should be sent with the bearer token.
+    /// `false` for deployment keys (bare Bearer), `true` for user/OAuth tokens.
+    /// See `GrokAuthCredentials::apply()` for the wire format contract.
+    fn needs_token_auth_header(&self) -> bool {
+        true
+    }
+
+    /// Whether the provider holds a credential worth a real outbound attempt —
+    /// an unexpired token (in memory or on disk), or a static key. Default
+    /// `true` always attempts.
+    fn has_usable_credential(&self) -> bool {
+        true
+    }
+}
+
+/// Static credential provider. Used by tests and by callers that pass a
+/// raw `&str` token with no `AuthManager` available.
+///
+/// `apply()` delegates to the underlying `HttpAuth::apply()`.
+/// `refresh_after_unauthorized()` always returns `false`.
+///
+/// `bearer` is the wire bearer the inner `HttpAuth` will send in the
+/// `Authorization` header. Stored alongside the inner so `snapshot().token`
+/// returns the same prefix that goes out on the wire (used by
+/// 401-attribution telemetry). `None` when no bearer is configured.
+pub struct StaticAuthCredentialProvider {
+    inner: Box<dyn HttpAuth>,
+    bearer: Option<String>,
+}
+
+impl StaticAuthCredentialProvider {
+    /// Wrap `inner` so callers see it as an `AuthCredentialProvider`. Pass
+    /// the bearer token that `inner.apply()` will send in the `Authorization`
+    /// header so `snapshot().token` reflects the wire bearer truthfully.
+    pub fn new(inner: Box<dyn HttpAuth>, bearer: Option<String>) -> Self {
+        Self { inner, bearer }
+    }
+}
+
+impl std::fmt::Debug for StaticAuthCredentialProvider {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("StaticAuthCredentialProvider")
+            .field("has_bearer", &self.bearer.is_some())
+            .finish()
+    }
+}
+
+impl HttpAuth for StaticAuthCredentialProvider {
+    fn apply(&self, builder: RequestBuilder, base_url: &str) -> RequestBuilder {
+        self.inner.apply(builder, base_url)
+    }
+}
+
+#[async_trait::async_trait]
+impl AuthCredentialProvider for StaticAuthCredentialProvider {
+    fn snapshot(&self) -> CredentialSnapshot {
+        CredentialSnapshot {
+            token: self.bearer.clone(),
+            ..Default::default()
+        }
+    }
+
+    async fn refresh_after_unauthorized(&self) -> bool {
+        false
+    }
+}

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

@@ -0,0 +1,14 @@
+//! Auth dependency-inversion seam shared between `xai-file-utils`
+//! (the holder) and `xai-grok-shell` (the implementer). Keeps shell types
+//! out of data-collector's import graph while still letting refresh-aware
+//! token resolution drive HTTP requests.
+
+pub mod auth_provider;
+#[cfg(feature = "middleware")]
+pub mod retry_middleware;
+pub mod visibility;
+
+pub use auth_provider::{AuthCredentialProvider, CredentialSnapshot, StaticAuthCredentialProvider};
+#[cfg(feature = "middleware")]
+pub use retry_middleware::AuthRetryMiddleware;
+pub use visibility::HttpAuth;

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

@@ -0,0 +1,272 @@
+//! `reqwest-middleware` layer: stamps auth headers and retries on 401.
+//! Gated behind the `middleware` cargo feature.
+
+use std::sync::Arc;
+
+use reqwest::{Request, Response, StatusCode, header::HeaderValue};
+use reqwest_middleware::{Error, Middleware, Next};
+
+use crate::AuthCredentialProvider;
+
+pub struct AuthRetryMiddleware {
+    credentials: Arc<dyn AuthCredentialProvider>,
+    max_retries: u32,
+}
+
+impl AuthRetryMiddleware {
+    pub fn new(credentials: Arc<dyn AuthCredentialProvider>, max_retries: u32) -> Self {
+        Self {
+            credentials,
+            max_retries,
+        }
+    }
+}
+
+fn apply_auth_header(req: &mut Request, token: &str) {
+    match HeaderValue::from_str(&format!("Bearer {token}")) {
+        Ok(val) => {
+            req.headers_mut()
+                .insert(reqwest::header::AUTHORIZATION, val);
+        }
+        Err(e) => {
+            tracing::warn!(error = %e, "auth retry: failed to build Authorization header");
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl Middleware for AuthRetryMiddleware {
+    async fn handle(
+        &self,
+        mut req: Request,
+        extensions: &mut http::Extensions,
+        next: Next<'_>,
+    ) -> Result<Response, Error> {
+        if let Some(ref token) = self.credentials.snapshot().token {
+            apply_auth_header(&mut req, token);
+        }
+
+        let backup = req.try_clone();
+        let resp = next.clone().run(req, extensions).await?;
+
+        if resp.status() != StatusCode::UNAUTHORIZED || self.max_retries == 0 {
+            return Ok(resp);
+        }
+        let Some(backup) = backup else {
+            return Ok(resp);
+        };
+
+        let mut last_resp = resp;
+        for _ in 0..self.max_retries {
+            if !self.credentials.refresh_after_unauthorized().await {
+                break;
+            }
+            let Some(ref token) = self.credentials.snapshot().token else {
+                break;
+            };
+            let Some(mut retry) = backup.try_clone() else {
+                break;
+            };
+            apply_auth_header(&mut retry, token);
+            last_resp = next.clone().run(retry, extensions).await?;
+            if last_resp.status() != StatusCode::UNAUTHORIZED {
+                return Ok(last_resp);
+            }
+        }
+
+        Ok(last_resp)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{CredentialSnapshot, HttpAuth};
+    use reqwest_middleware::ClientBuilder;
+    use std::sync::Mutex;
+
+    struct MockProvider {
+        token: Mutex<Option<String>>,
+        refresh_result: bool,
+        refresh_count: Mutex<u32>,
+    }
+
+    impl MockProvider {
+        fn new(token: Option<&str>, refresh_result: bool) -> Self {
+            Self {
+                token: Mutex::new(token.map(|s| s.to_owned())),
+                refresh_result,
+                refresh_count: Mutex::new(0),
+            }
+        }
+
+        fn refresh_count(&self) -> u32 {
+            *self.refresh_count.lock().unwrap()
+        }
+    }
+
+    impl HttpAuth for MockProvider {
+        fn apply(&self, b: reqwest::RequestBuilder, _: &str) -> reqwest::RequestBuilder {
+            b
+        }
+    }
+
+    #[async_trait::async_trait]
+    impl AuthCredentialProvider for MockProvider {
+        fn snapshot(&self) -> CredentialSnapshot {
+            CredentialSnapshot {
+                token: self.token.lock().unwrap().clone(),
+                ..Default::default()
+            }
+        }
+
+        async fn refresh_after_unauthorized(&self) -> bool {
+            *self.refresh_count.lock().unwrap() += 1;
+            self.refresh_result
+        }
+    }
+
+    async fn build_client(
+        provider: Arc<dyn AuthCredentialProvider>,
+        max_retries: u32,
+    ) -> reqwest_middleware::ClientWithMiddleware {
+        ClientBuilder::new(reqwest::Client::new())
+            .with(AuthRetryMiddleware::new(provider, max_retries))
+            .build()
+    }
+
+    #[tokio::test]
+    async fn test_401_no_refresh_returns_401() {
+        let mut server = mockito::Server::new_async().await;
+        let m = server
+            .mock("GET", "/")
+            .with_status(401)
+            .expect(1)
+            .create_async()
+            .await;
+
+        let p = Arc::new(MockProvider::new(Some("tok"), false));
+        let client = build_client(p.clone(), 1).await;
+
+        let resp = client.get(server.url()).send().await.unwrap();
+        assert_eq!(resp.status(), 401);
+        assert_eq!(p.refresh_count(), 1);
+        m.assert_async().await;
+    }
+
+    /// Simulates a real auth manager: starts with stale token, refresh swaps to fresh.
+    struct SimulatedAuthManager {
+        token: Mutex<Option<String>>,
+        fresh_token: String,
+        refresh_count: Mutex<u32>,
+    }
+
+    impl SimulatedAuthManager {
+        fn simulated(stale: &str, fresh: &str) -> Self {
+            Self {
+                token: Mutex::new(Some(stale.to_owned())),
+                fresh_token: fresh.to_owned(),
+                refresh_count: Mutex::new(0),
+            }
+        }
+    }
+
+    impl HttpAuth for SimulatedAuthManager {
+        fn apply(&self, b: reqwest::RequestBuilder, _: &str) -> reqwest::RequestBuilder {
+            b
+        }
+    }
+
+    #[async_trait::async_trait]
+    impl AuthCredentialProvider for SimulatedAuthManager {
+        fn snapshot(&self) -> CredentialSnapshot {
+            CredentialSnapshot {
+                token: self.token.lock().unwrap().clone(),
+                ..Default::default()
+            }
+        }
+
+        async fn refresh_after_unauthorized(&self) -> bool {
+            *self.refresh_count.lock().unwrap() += 1;
+            *self.token.lock().unwrap() = Some(self.fresh_token.clone());
+            true
+        }
+    }
+
+    #[tokio::test]
+    async fn test_e2e_stale_token_refreshed_and_retried() {
+        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)
+            .with_body(r#"{"ok":true}"#)
+            .create_async()
+            .await;
+
+        let p = Arc::new(SimulatedAuthManager::simulated(
+            "stale-token",
+            "fresh-token",
+        ));
+        let client = build_client(p.clone(), 1).await;
+
+        let resp = client
+            .get(format!("{}/api", server.url()))
+            .send()
+            .await
+            .unwrap();
+        assert_eq!(resp.status(), 200);
+        assert_eq!(*p.refresh_count.lock().unwrap(), 1);
+        m401.assert_async().await;
+        m200.assert_async().await;
+    }
+
+    #[tokio::test]
+    async fn test_e2e_auth_header_stamped_automatically() {
+        let mut server = mockito::Server::new_async().await;
+        let mock = server
+            .mock("GET", "/api")
+            .match_header("authorization", "Bearer my-token")
+            .with_status(200)
+            .create_async()
+            .await;
+
+        let p = Arc::new(MockProvider::new(Some("my-token"), false));
+        let client = build_client(p.clone(), 1).await;
+
+        let resp = client
+            .get(format!("{}/api", server.url()))
+            .send()
+            .await
+            .unwrap();
+        assert_eq!(resp.status(), 200);
+        assert_eq!(p.refresh_count(), 0);
+        mock.assert_async().await;
+    }
+
+    #[tokio::test]
+    async fn test_max_retries_bounds_attempts() {
+        let mut server = mockito::Server::new_async().await;
+        let m = server
+            .mock("GET", "/")
+            .with_status(401)
+            .expect(4)
+            .create_async()
+            .await;
+
+        let p = Arc::new(MockProvider::new(Some("tok"), true));
+        let client = build_client(p.clone(), 3).await;
+
+        let resp = client.get(server.url()).send().await.unwrap();
+        assert_eq!(resp.status(), 401);
+        assert_eq!(p.refresh_count(), 3);
+        m.assert_async().await;
+    }
+}

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

@@ -0,0 +1,7 @@
+/// Apply auth headers to outbound visibility requests.
+/// Implemented by `xai-grok-shell::util::grok_auth_credentials::GrokAuthCredentials`
+/// to keep credential construction owned by shell while letting data-collector
+/// build the request without reaching back into shell types.
+pub trait HttpAuth: Send + Sync {
+    fn apply(&self, builder: reqwest::RequestBuilder, base_url: &str) -> reqwest::RequestBuilder;
+}