Skip to content

Commit 7eca1fb

Browse files
pthurlowpthurlow
authored
Add sandbox JWT support
1 parent 6095dfa commit 7eca1fb

6 files changed

Lines changed: 714 additions & 95 deletions

File tree

‎src/api.rs‎

Lines changed: 55 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -47,24 +47,61 @@ impl ApiClient {
4747
}
4848
};
4949

50-
let api_key_fallback = profile_config
51-
.api_key
52-
.as_deref()
53-
.filter(|k| !k.is_empty() && *k != "PLACEHOLDER");
54-
55-
// Pre-flight: return the cached JWT if valid, refresh it if
56-
// close to expiry, or mint a new one from the API key. The
57-
// returned string is a JWT — that's what we send on the wire.
58-
let access_token = match crate::jwt::ensure_access_token(&profile_config, api_key_fallback)
59-
{
60-
Ok(t) => t,
61-
Err(e) => {
62-
eprintln!("{}", format!("error: {e}").red());
63-
eprintln!(
64-
"Run {} to log in, or pass --api-key.",
65-
"hotdata auth".cyan()
66-
);
67-
std::process::exit(1);
50+
// Auth source precedence:
51+
//
52+
// 1. `HOTDATA_SANDBOX_TOKEN` env var — a `sandbox run` child
53+
// is executing with the parent's credentials scrubbed.
54+
// Refresh in-memory via `HOTDATA_SANDBOX_REFRESH_TOKEN` if
55+
// the JWT is close to expiry; never write to disk (the
56+
// child's FS may not be writable).
57+
// 2. `~/.hotdata/sandbox_session.json` — the user ran
58+
// `hotdata sandbox set <id>` (or `sandbox new` / `sandbox
59+
// run` in the parent shell). The sandbox JWT is the active
60+
// bearer for *every* command until `sandbox set` (with no
61+
// id) clears the file.
62+
// 3. `~/.hotdata/session.json` + optional api_key fallback —
63+
// normal user-scoped CLI session.
64+
let api_url = profile_config.api_url.to_string();
65+
let access_token = if std::env::var("HOTDATA_SANDBOX_TOKEN").is_ok() {
66+
match crate::sandbox_session::refresh_from_env(&api_url) {
67+
Some(t) => t,
68+
None => {
69+
eprintln!("{}", "error: HOTDATA_SANDBOX_TOKEN is empty".red());
70+
std::process::exit(1);
71+
}
72+
}
73+
} else if crate::sandbox_session::load().is_some() {
74+
match crate::sandbox_session::ensure_access_token(&api_url) {
75+
Some(t) => t,
76+
None => {
77+
eprintln!("{}", "error: sandbox session expired".red());
78+
eprintln!(
79+
"Run {} to clear it, or {} to re-mint.",
80+
"hotdata sandbox set".cyan(),
81+
"hotdata sandbox set <id>".cyan(),
82+
);
83+
std::process::exit(1);
84+
}
85+
}
86+
} else {
87+
let api_key_fallback = profile_config
88+
.api_key
89+
.as_deref()
90+
.filter(|k| !k.is_empty() && *k != "PLACEHOLDER");
91+
92+
// Pre-flight: return the cached JWT if valid, refresh it if
93+
// close to expiry, or mint a new one from the API key. The
94+
// returned string is a JWT — that's what we send on the wire.
95+
match crate::jwt::ensure_access_token(&profile_config, api_key_fallback) {
96+
Ok(t) => t,
97+
Err(e) => {
98+
eprintln!("{}", format!("error: {e}").red());
99+
eprintln!(
100+
"Run {} to log in, or pass --api-key.",
101+
"hotdata auth".cyan()
102+
);
103+
std::process::exit(1);
104+
}
68105
}
69106
};
70107

‎src/auth.rs‎

Lines changed: 73 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -26,21 +26,37 @@ pub enum AuthStatus {
2626
}
2727

2828
pub fn check_status(profile_config: &config::ProfileConfig) -> AuthStatus {
29-
let api_key_fallback = profile_config
30-
.api_key
31-
.as_deref()
32-
.filter(|k| !k.is_empty() && *k != "PLACEHOLDER");
33-
34-
// PKCE-origin sessions don't write an api_key, so absence of a key
35-
// alone isn't "not configured" — only true if there's also no
36-
// cached JWT session to validate.
37-
if api_key_fallback.is_none() && crate::jwt::load_session().is_none() {
38-
return AuthStatus::NotConfigured;
39-
}
29+
// Same precedence as `ApiClient::new`:
30+
// 1. `sandbox run` child via env var
31+
// 2. on-disk sandbox session (sandbox set <id>)
32+
// 3. user-scoped CLI session / api_key fallback
33+
let api_url = profile_config.api_url.to_string();
34+
let access_token = if let Some((sandbox_jwt, _)) =
35+
crate::sandbox_session::sandbox_token_in_use()
36+
{
37+
sandbox_jwt
38+
} else if crate::sandbox_session::load().is_some() {
39+
match crate::sandbox_session::ensure_access_token(&api_url) {
40+
Some(t) => t,
41+
None => return AuthStatus::Invalid(401),
42+
}
43+
} else {
44+
let api_key_fallback = profile_config
45+
.api_key
46+
.as_deref()
47+
.filter(|k| !k.is_empty() && *k != "PLACEHOLDER");
48+
49+
// PKCE-origin sessions don't write an api_key, so absence of a key
50+
// alone isn't "not configured" — only true if there's also no
51+
// cached JWT session to validate.
52+
if api_key_fallback.is_none() && crate::jwt::load_session().is_none() {
53+
return AuthStatus::NotConfigured;
54+
}
4055

41-
let access_token = match crate::jwt::ensure_access_token(profile_config, api_key_fallback) {
42-
Ok(t) => t,
43-
Err(_) => return AuthStatus::Invalid(401),
56+
match crate::jwt::ensure_access_token(profile_config, api_key_fallback) {
57+
Ok(t) => t,
58+
Err(_) => return AuthStatus::Invalid(401),
59+
}
4460
};
4561

4662
let url = format!("{}/workspaces", profile_config.api_url);
@@ -64,26 +80,50 @@ pub fn status(profile: &str) {
6480
}
6581
};
6682

67-
// The credential the CLI is *about to use*. Note: even when an
68-
// override is set, the wire credential is still a JWT (minted on
69-
// demand from the override) — but we report the user-visible source.
70-
let method_label = match profile_config.api_key_source {
71-
ApiKeySource::Flag => "API Key flag",
72-
ApiKeySource::Env => "API Key env",
73-
ApiKeySource::Config => "CLI Session",
83+
// The credential the CLI is *about to use*. Precedence matches
84+
// `ApiClient::new`: env-var sandbox token (sandbox run child) >
85+
// on-disk sandbox session (sandbox set <id>) > user CLI session.
86+
let env_sandbox = crate::sandbox_session::sandbox_token_in_use();
87+
let disk_sandbox = if env_sandbox.is_none() {
88+
crate::sandbox_session::load()
89+
} else {
90+
None
7491
};
75-
76-
// For Flag/Env we mask the api_key the user supplied. For the
77-
// CLI session path we mask the refresh_token — it's stable across
78-
// commands (unlike the 5-min access_token), so the tail stays
79-
// recognizable between runs.
80-
let credential_tail = match profile_config.api_key_source {
81-
ApiKeySource::Flag | ApiKeySource::Env => profile_config
82-
.api_key
83-
.as_deref()
84-
.map(crate::util::mask_credential),
85-
ApiKeySource::Config => crate::jwt::load_session()
86-
.map(|s| crate::util::mask_credential(&s.refresh_token)),
92+
let (method_label, credential_tail) = if let Some((token, sandbox_id)) = &env_sandbox {
93+
let label = match sandbox_id {
94+
Some(id) => format!("Sandbox {id}"),
95+
None => "Sandbox Session".to_string(),
96+
};
97+
(label, Some(crate::util::mask_credential(token)))
98+
} else if let Some(s) = &disk_sandbox {
99+
// Use the refresh token for the displayed tail — it's stable
100+
// across refreshes (the access token rotates every 3 days), so
101+
// the tail stays recognizable between runs.
102+
let label = if s.sandbox_id.is_empty() {
103+
"Sandbox Session".to_string()
104+
} else {
105+
format!("Sandbox {}", s.sandbox_id)
106+
};
107+
(label, Some(crate::util::mask_credential(&s.refresh_token)))
108+
} else {
109+
let label = match profile_config.api_key_source {
110+
ApiKeySource::Flag => "API Key flag",
111+
ApiKeySource::Env => "API Key env",
112+
ApiKeySource::Config => "CLI Session",
113+
};
114+
// For Flag/Env we mask the api_key the user supplied. For
115+
// the CLI session path we mask the refresh_token — it's
116+
// stable across commands (unlike the 5-min access_token),
117+
// so the tail stays recognizable between runs.
118+
let tail = match profile_config.api_key_source {
119+
ApiKeySource::Flag | ApiKeySource::Env => profile_config
120+
.api_key
121+
.as_deref()
122+
.map(crate::util::mask_credential),
123+
ApiKeySource::Config => crate::jwt::load_session()
124+
.map(|s| crate::util::mask_credential(&s.refresh_token)),
125+
};
126+
(label.to_string(), tail)
87127
};
88128
let method_suffix = match credential_tail {
89129
Some(tail) => format!(" - {method_label} [{tail}]"),

‎src/main.rs‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ mod queries;
1414
mod query;
1515
mod results;
1616
mod sandbox;
17+
mod sandbox_session;
1718
mod skill;
1819
mod table;
1920
mod tables;
@@ -81,7 +82,49 @@ fn resolve_workspace(provided: Option<String>) -> String {
8182
}
8283
}
8384

85+
// libc::atexit (no extra crate needed — the symbol is linked by default).
86+
// Callbacks registered here fire even when subcommands call
87+
// `std::process::exit`, which Rust's `Drop` would otherwise miss.
88+
unsafe extern "C" {
89+
fn atexit(callback: extern "C" fn()) -> i32;
90+
}
91+
92+
/// Runs once at process exit. Prints a sandbox footer on stderr when
93+
/// the CLI is running under an on-disk sandbox session (i.e. the user
94+
/// ran `hotdata sandbox set <id>` to enter it from this shell). Stays
95+
/// silent when the sandbox comes from `HOTDATA_SANDBOX_TOKEN` in the
96+
/// environment: that means we're inside a `sandbox run` child, and
97+
/// the parent already announced the sandbox once at spawn time.
98+
/// Stderr keeps stdout clean for callers parsing JSON/YAML output.
99+
extern "C" fn print_sandbox_footer() {
100+
use crossterm::style::Stylize;
101+
102+
// Inside a `sandbox run` child — parent printed the banner already.
103+
if sandbox_session::sandbox_token_in_use().is_some() {
104+
return;
105+
}
106+
let Some(session) = sandbox_session::load() else {
107+
return;
108+
};
109+
if session.sandbox_id.is_empty() {
110+
return;
111+
}
112+
eprintln!(
113+
"{}",
114+
format!(
115+
"current sandbox: {} use 'hotdata sandbox set' to change",
116+
session.sandbox_id
117+
)
118+
.dark_grey(),
119+
);
120+
}
121+
84122
fn main() {
123+
// Register before `Cli::parse`, since `--help` / `--version` exit
124+
// from inside the parser. Safety: `atexit` is async-signal-safe;
125+
// the callback only reads env vars / files and writes to stderr.
126+
unsafe { atexit(print_sandbox_footer) };
127+
85128
dotenvy::dotenv().ok();
86129
let cli = Cli::parse();
87130

0 commit comments

Comments
 (0)