Add sandbox JWT support

src/api.rs

@@ -47,24 +47,61 @@ impl ApiClient {
             }
         };
 
-        let api_key_fallback = profile_config
-            .api_key
-            .as_deref()
-            .filter(|k| !k.is_empty() && *k != "PLACEHOLDER");
-
-        // Pre-flight: return the cached JWT if valid, refresh it if
-        // close to expiry, or mint a new one from the API key. The
-        // returned string is a JWT — that's what we send on the wire.
-        let access_token = match crate::jwt::ensure_access_token(&profile_config, api_key_fallback)
-        {
-            Ok(t) => t,
-            Err(e) => {
-                eprintln!("{}", format!("error: {e}").red());
-                eprintln!(
-                    "Run {} to log in, or pass --api-key.",
-                    "hotdata auth".cyan()
-                );
-                std::process::exit(1);
+        // Auth source precedence:
+        //
+        // 1. `HOTDATA_SANDBOX_TOKEN` env var — a `sandbox run` child
+        //    is executing with the parent's credentials scrubbed.
+        //    Refresh in-memory via `HOTDATA_SANDBOX_REFRESH_TOKEN` if
+        //    the JWT is close to expiry; never write to disk (the
+        //    child's FS may not be writable).
+        // 2. `~/.hotdata/sandbox_session.json` — the user ran
+        //    `hotdata sandbox set <id>` (or `sandbox new` / `sandbox
+        //    run` in the parent shell). The sandbox JWT is the active
+        //    bearer for *every* command until `sandbox set` (with no
+        //    id) clears the file.
+        // 3. `~/.hotdata/session.json` + optional api_key fallback —
+        //    normal user-scoped CLI session.
+        let api_url = profile_config.api_url.to_string();
+        let access_token = if std::env::var("HOTDATA_SANDBOX_TOKEN").is_ok() {
+            match crate::sandbox_session::refresh_from_env(&api_url) {
+                Some(t) => t,
+                None => {
+                    eprintln!("{}", "error: HOTDATA_SANDBOX_TOKEN is empty".red());
+                    std::process::exit(1);
+                }
+            }
+        } else if crate::sandbox_session::load().is_some() {
+            match crate::sandbox_session::ensure_access_token(&api_url) {
+                Some(t) => t,
+                None => {
+                    eprintln!("{}", "error: sandbox session expired".red());
+                    eprintln!(
+                        "Run {} to clear it, or {} to re-mint.",
+                        "hotdata sandbox set".cyan(),
+                        "hotdata sandbox set <id>".cyan(),
+                    );
+                    std::process::exit(1);
+                }
+            }
+        } else {
+            let api_key_fallback = profile_config
+                .api_key
+                .as_deref()
+                .filter(|k| !k.is_empty() && *k != "PLACEHOLDER");
+
+            // Pre-flight: return the cached JWT if valid, refresh it if
+            // close to expiry, or mint a new one from the API key. The
+            // returned string is a JWT — that's what we send on the wire.
+            match crate::jwt::ensure_access_token(&profile_config, api_key_fallback) {
+                Ok(t) => t,
+                Err(e) => {
+                    eprintln!("{}", format!("error: {e}").red());
+                    eprintln!(
+                        "Run {} to log in, or pass --api-key.",
+                        "hotdata auth".cyan()
+                    );
+                    std::process::exit(1);
+                }
             }
         };
 

src/auth.rs

@@ -26,21 +26,37 @@ pub enum AuthStatus {
 }
 
 pub fn check_status(profile_config: &config::ProfileConfig) -> AuthStatus {
-    let api_key_fallback = profile_config
-        .api_key
-        .as_deref()
-        .filter(|k| !k.is_empty() && *k != "PLACEHOLDER");
-
-    // PKCE-origin sessions don't write an api_key, so absence of a key
-    // alone isn't "not configured" — only true if there's also no
-    // cached JWT session to validate.
-    if api_key_fallback.is_none() && crate::jwt::load_session().is_none() {
-        return AuthStatus::NotConfigured;
-    }
+    // Same precedence as `ApiClient::new`:
+    // 1. `sandbox run` child via env var
+    // 2. on-disk sandbox session (sandbox set <id>)
+    // 3. user-scoped CLI session / api_key fallback
+    let api_url = profile_config.api_url.to_string();
+    let access_token = if let Some((sandbox_jwt, _)) =
+        crate::sandbox_session::sandbox_token_in_use()
+    {
+        sandbox_jwt
+    } else if crate::sandbox_session::load().is_some() {
+        match crate::sandbox_session::ensure_access_token(&api_url) {
+            Some(t) => t,
+            None => return AuthStatus::Invalid(401),
+        }
+    } else {
+        let api_key_fallback = profile_config
+            .api_key
+            .as_deref()
+            .filter(|k| !k.is_empty() && *k != "PLACEHOLDER");
+
+        // PKCE-origin sessions don't write an api_key, so absence of a key
+        // alone isn't "not configured" — only true if there's also no
+        // cached JWT session to validate.
+        if api_key_fallback.is_none() && crate::jwt::load_session().is_none() {
+            return AuthStatus::NotConfigured;
+        }
 
-    let access_token = match crate::jwt::ensure_access_token(profile_config, api_key_fallback) {
-        Ok(t) => t,
-        Err(_) => return AuthStatus::Invalid(401),
+        match crate::jwt::ensure_access_token(profile_config, api_key_fallback) {
+            Ok(t) => t,
+            Err(_) => return AuthStatus::Invalid(401),
+        }
     };
 
     let url = format!("{}/workspaces", profile_config.api_url);
@@ -64,26 +80,50 @@ pub fn status(profile: &str) {
         }
     };
 
-    // The credential the CLI is *about to use*. Note: even when an
-    // override is set, the wire credential is still a JWT (minted on
-    // demand from the override) — but we report the user-visible source.
-    let method_label = match profile_config.api_key_source {
-        ApiKeySource::Flag => "API Key flag",
-        ApiKeySource::Env => "API Key env",
-        ApiKeySource::Config => "CLI Session",
+    // The credential the CLI is *about to use*. Precedence matches
+    // `ApiClient::new`: env-var sandbox token (sandbox run child) >
+    // on-disk sandbox session (sandbox set <id>) > user CLI session.
+    let env_sandbox = crate::sandbox_session::sandbox_token_in_use();
+    let disk_sandbox = if env_sandbox.is_none() {
+        crate::sandbox_session::load()
+    } else {
+        None
     };
-
-    // For Flag/Env we mask the api_key the user supplied. For the
-    // CLI session path we mask the refresh_token — it's stable across
-    // commands (unlike the 5-min access_token), so the tail stays
-    // recognizable between runs.
-    let credential_tail = match profile_config.api_key_source {
-        ApiKeySource::Flag | ApiKeySource::Env => profile_config
-            .api_key
-            .as_deref()
-            .map(crate::util::mask_credential),
-        ApiKeySource::Config => crate::jwt::load_session()
-            .map(|s| crate::util::mask_credential(&s.refresh_token)),
+    let (method_label, credential_tail) = if let Some((token, sandbox_id)) = &env_sandbox {
+        let label = match sandbox_id {
+            Some(id) => format!("Sandbox {id}"),
+            None => "Sandbox Session".to_string(),
+        };
+        (label, Some(crate::util::mask_credential(token)))
+    } else if let Some(s) = &disk_sandbox {
+        // Use the refresh token for the displayed tail — it's stable
+        // across refreshes (the access token rotates every 3 days), so
+        // the tail stays recognizable between runs.
+        let label = if s.sandbox_id.is_empty() {
+            "Sandbox Session".to_string()
+        } else {
+            format!("Sandbox {}", s.sandbox_id)
+        };
+        (label, Some(crate::util::mask_credential(&s.refresh_token)))
+    } else {
+        let label = match profile_config.api_key_source {
+            ApiKeySource::Flag => "API Key flag",
+            ApiKeySource::Env => "API Key env",
+            ApiKeySource::Config => "CLI Session",
+        };
+        // For Flag/Env we mask the api_key the user supplied. For
+        // the CLI session path we mask the refresh_token — it's
+        // stable across commands (unlike the 5-min access_token),
+        // so the tail stays recognizable between runs.
+        let tail = match profile_config.api_key_source {
+            ApiKeySource::Flag | ApiKeySource::Env => profile_config
+                .api_key
+                .as_deref()
+                .map(crate::util::mask_credential),
+            ApiKeySource::Config => crate::jwt::load_session()
+                .map(|s| crate::util::mask_credential(&s.refresh_token)),
+        };
+        (label.to_string(), tail)
     };
     let method_suffix = match credential_tail {
         Some(tail) => format!(" - {method_label} [{tail}]"),

src/main.rs

@@ -14,6 +14,7 @@ mod queries;
 mod query;
 mod results;
 mod sandbox;
+mod sandbox_session;
 mod skill;
 mod table;
 mod tables;
@@ -81,7 +82,49 @@ fn resolve_workspace(provided: Option<String>) -> String {
     }
 }
 
+// libc::atexit (no extra crate needed — the symbol is linked by default).
+// Callbacks registered here fire even when subcommands call
+// `std::process::exit`, which Rust's `Drop` would otherwise miss.
+unsafe extern "C" {
+    fn atexit(callback: extern "C" fn()) -> i32;
+}
+
+/// Runs once at process exit. Prints a sandbox footer on stderr when
+/// the CLI is running under an on-disk sandbox session (i.e. the user
+/// ran `hotdata sandbox set <id>` to enter it from this shell). Stays
+/// silent when the sandbox comes from `HOTDATA_SANDBOX_TOKEN` in the
+/// environment: that means we're inside a `sandbox run` child, and
+/// the parent already announced the sandbox once at spawn time.
+/// Stderr keeps stdout clean for callers parsing JSON/YAML output.
+extern "C" fn print_sandbox_footer() {
+    use crossterm::style::Stylize;
+
+    // Inside a `sandbox run` child — parent printed the banner already.
+    if sandbox_session::sandbox_token_in_use().is_some() {
+        return;
+    }
+    let Some(session) = sandbox_session::load() else {
+        return;
+    };
+    if session.sandbox_id.is_empty() {
+        return;
+    }
+    eprintln!(
+        "{}",
+        format!(
+            "current sandbox: {} use 'hotdata sandbox set' to change",
+            session.sandbox_id
+        )
+        .dark_grey(),
+    );
+}
+
 fn main() {
+    // Register before `Cli::parse`, since `--help` / `--version` exit
+    // from inside the parser. Safety: `atexit` is async-signal-safe;
+    // the callback only reads env vars / files and writes to stderr.
+    unsafe { atexit(print_sandbox_footer) };
+
     dotenvy::dotenv().ok();
     let cli = Cli::parse();