forge_gateway/
binding.rs

1//! Session registration: the binding of `(org, project, session)` to
2//! an account, and the env set a child is stamped with.
3//!
4//! The child never holds a real credential. Its `ANTHROPIC_BASE_URL`
5//! points at this listener with the three routing segments in the
6//! path, and the credential it does hold is a dummy delivered through
7//! the SAME variable the account's real credential uses - moving a
8//! credential between variables changes which entitlements reach the
9//! backend, not only which header carries it.
10
11use std::collections::HashMap;
12
13use forge_primitives::account::Provider;
14use parking_lot::Mutex;
15
16use crate::account::AccountKey;
17
18/// What the CLI sends when it has no real credential. Authorises
19/// nothing upstream; the gateway replaces it before forwarding.
20pub const DUMMY_CREDENTIAL: &str = "forge-gateway-unused";
21
22pub(crate) const OAUTH_VARIABLE: &str = "CLAUDE_CODE_OAUTH_TOKEN";
23pub(crate) const AUTH_TOKEN_VARIABLE: &str = "ANTHROPIC_AUTH_TOKEN";
24const API_KEY_VARIABLE: &str = "ANTHROPIC_API_KEY";
25const ALT_BASE_URL_VARIABLE: &str = "CLAUDE_CODE_API_BASE_URL";
26const BASE_URL_VARIABLE: &str = "ANTHROPIC_BASE_URL";
27
28/// One spawn's registration: who owns the session and where it sits.
29#[derive(Debug, Clone, PartialEq, Eq)]
30pub struct Registration {
31    pub org: String,
32    pub project: String,
33    pub session: String,
34    pub account: AccountKey,
35    pub provider: Provider,
36}
37
38/// The env a child is stamped with, in insertion order.
39#[derive(Clone, PartialEq, Eq)]
40pub struct EnvSet(pub Vec<(String, String)>);
41
42/// Hand-written because the set carries live values for the keys the
43/// gateway does not own; only the key names may print.
44impl std::fmt::Debug for EnvSet {
45    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46        let pairs: Vec<(&str, &str)> =
47            self.0.iter().map(|(key, _)| (key.as_str(), "[redacted]")).collect();
48        f.debug_tuple("EnvSet").field(&pairs).finish()
49    }
50}
51
52impl EnvSet {
53    pub fn into_map(self) -> HashMap<String, String> {
54        self.0.into_iter().collect()
55    }
56
57    pub fn get(&self, key: &str) -> Option<&str> {
58        self.0.iter().find(|(k, _)| k == key).map(|(_, v)| v.as_str())
59    }
60}
61
62/// The session bindings, keyed by the three routing segments.
63#[derive(Default)]
64pub struct Bindings {
65    by_session: Mutex<HashMap<(String, String, String), AccountKey>>,
66}
67
68impl Bindings {
69    /// Register `registration` and derive the env set its child is
70    /// stamped with, from the account's real env. Re-registering the
71    /// same three segments replaces the prior binding, so a respawned
72    /// CLI answers to this generation of the session, not the last.
73    pub fn register(
74        &self,
75        registration: &Registration,
76        listener_base: &str,
77        account_env: &HashMap<String, String>,
78    ) -> EnvSet {
79        let segments =
80            (registration.org.clone(), registration.project.clone(), registration.session.clone());
81        self.by_session.lock().insert(segments, registration.account.clone());
82
83        let mut env: Vec<(String, String)> = account_env
84            .iter()
85            .map(|(k, v)| (k.clone(), v.clone()))
86            .filter(|(k, _)| {
87                k != BASE_URL_VARIABLE
88                    && k != ALT_BASE_URL_VARIABLE
89                    && k != OAUTH_VARIABLE
90                    && k != AUTH_TOKEN_VARIABLE
91                    && k != API_KEY_VARIABLE
92            })
93            .collect();
94        for (key, value) in stamped_gateway_keys(registration, listener_base) {
95            if let Some(slot) = env.iter_mut().find(|(k, _)| *k == key) {
96                slot.1 = value;
97            } else {
98                env.push((key, value));
99            }
100        }
101        EnvSet(env)
102    }
103
104    /// Re-register `registration` and return ONLY the four
105    /// gateway-owned keys for a respawn's launch-settings overrides:
106    /// the base URL naming this listener, the alt base-url slot, the
107    /// dummy in the account's credential variable, and the forced-empty
108    /// API key. The account and project env already live in the bridge
109    /// from the original spawn, so carrying the whole env here would
110    /// let a key declared in both layers revert to its account value.
111    ///
112    /// `registration.session` is the segment the respawned child runs
113    /// under; `replaced_session` is the one the child it replaces ran
114    /// under, dropped under the same lock as the insert, so no lookup
115    /// lands between the two. Without that drop a session keeps one
116    /// binding for every id it has ever held.
117    pub fn respawn_env_overrides(
118        &self,
119        registration: &Registration,
120        replaced_session: &str,
121        listener_base: &str,
122    ) -> HashMap<String, String> {
123        {
124            let mut map = self.by_session.lock();
125            map.remove(&(
126                registration.org.clone(),
127                registration.project.clone(),
128                replaced_session.to_owned(),
129            ));
130            map.insert(
131                (
132                    registration.org.clone(),
133                    registration.project.clone(),
134                    registration.session.clone(),
135                ),
136                registration.account.clone(),
137            );
138        }
139        stamped_gateway_keys(registration, listener_base).into_iter().collect()
140    }
141
142    /// The account bound to the three routing segments, if any.
143    pub fn binding_for(&self, org: &str, project: &str, session: &str) -> Option<AccountKey> {
144        self.by_session
145            .lock()
146            .get(&(org.to_owned(), project.to_owned(), session.to_owned()))
147            .cloned()
148    }
149
150    /// Remove and return the binding for the three routing segments,
151    /// read and removed under one acquisition. A caller that acts on
152    /// what it read - cooling the account it has proved exhausted - must
153    /// not race a re-bind of the same triple between the two, or it
154    /// cools one account and drops another's binding.
155    pub fn take_binding(&self, org: &str, project: &str, session: &str) -> Option<AccountKey> {
156        self.by_session.lock().remove(&(org.to_owned(), project.to_owned(), session.to_owned()))
157    }
158
159    /// Bind a selected account to the three routing segments.
160    pub fn bind(&self, org: &str, project: &str, session: &str, account: AccountKey) {
161        self.by_session
162            .lock()
163            .insert((org.to_owned(), project.to_owned(), session.to_owned()), account);
164    }
165
166    /// Drop the binding, so the session's next request selects again.
167    pub fn unbind(&self, org: &str, project: &str, session: &str) {
168        self.by_session.lock().remove(&(org.to_owned(), project.to_owned(), session.to_owned()));
169    }
170
171    /// Drop every binding onto `account`, so those sessions select
172    /// again. The usage probe knows only the account it proved
173    /// exhausted.
174    pub fn unbind_for_account(&self, account: &AccountKey) {
175        self.by_session.lock().retain(|_, bound| bound != account);
176    }
177}
178
179/// The four gateway-owned keys, stamped fresh for `registration`:
180/// every other env key composes per the documented precedence, these
181/// four are the gateway's alone.
182fn stamped_gateway_keys(registration: &Registration, listener_base: &str) -> Vec<(String, String)> {
183    let base_url = format!(
184        "{}/{}/{}/{}",
185        listener_base, registration.org, registration.project, registration.session
186    );
187    let credential_variable = credential_variable_for(registration.provider);
188    vec![
189        (BASE_URL_VARIABLE.to_owned(), base_url.clone()),
190        (ALT_BASE_URL_VARIABLE.to_owned(), base_url),
191        (credential_variable.to_owned(), DUMMY_CREDENTIAL.to_owned()),
192        (API_KEY_VARIABLE.to_owned(), String::new()),
193    ]
194}
195
196/// The variable the account's real credential lives in, which is the
197/// variable its dummy has to travel in too. The one home for the rule:
198/// the config load maps the flat token onto it, the stamp puts the
199/// dummy in it, and the forward leg reads the real credential back out
200/// of it.
201pub fn credential_variable_for(provider: Provider) -> &'static str {
202    if provider.uses_base_url() { AUTH_TOKEN_VARIABLE } else { OAUTH_VARIABLE }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    fn anthropic_env() -> HashMap<String, String> {
210        HashMap::from([
211            ("CLAUDE_CODE_OAUTH_TOKEN".to_owned(), "real-oauth-token".to_owned()),
212            ("CLAUDE_CODE_MAX_CONTEXT_TOKENS".to_owned(), "1000000".to_owned()),
213        ])
214    }
215
216    fn base_url_env() -> HashMap<String, String> {
217        HashMap::from([
218            ("ANTHROPIC_BASE_URL".to_owned(), "https://openrouter.ai/api".to_owned()),
219            ("ANTHROPIC_AUTH_TOKEN".to_owned(), "real-openrouter-key".to_owned()),
220            ("ANTHROPIC_API_KEY".to_owned(), String::new()),
221        ])
222    }
223
224    fn registration(provider: Provider) -> Registration {
225        Registration {
226            org: "Busytools".to_owned(),
227            project: "forge".to_owned(),
228            session: "session-1".to_owned(),
229            account: AccountKey("OpenRouter".to_owned()),
230            provider,
231        }
232    }
233
234    fn env_contains_secret(env: &EnvSet, secret: &str) -> bool {
235        env.0.iter().any(|(_, v)| v.contains(secret))
236    }
237
238    #[test]
239    fn the_base_url_is_always_the_listener_with_the_three_segments() {
240        let bindings = Bindings::default();
241        let env = bindings.register(
242            &registration(Provider::Openrouter),
243            "http://127.0.0.1:8787",
244            &base_url_env(),
245        );
246        assert_eq!(
247            env.get("ANTHROPIC_BASE_URL"),
248            Some("http://127.0.0.1:8787/Busytools/forge/session-1"),
249            "the three routing segments ride the stamped base URL",
250        );
251    }
252
253    #[test]
254    fn a_base_url_account_gets_its_dummy_in_the_auth_token_variable() {
255        let bindings = Bindings::default();
256        let env = bindings.register(
257            &registration(Provider::Openrouter),
258            "http://127.0.0.1:8787",
259            &base_url_env(),
260        );
261        assert_eq!(
262            env.get("ANTHROPIC_AUTH_TOKEN"),
263            Some(DUMMY_CREDENTIAL),
264            "the dummy travels in the same variable the real credential uses",
265        );
266        assert!(
267            !env_contains_secret(&env, "real-openrouter-key"),
268            "the real credential never reaches the child",
269        );
270    }
271
272    #[test]
273    fn an_anthropic_account_gets_its_dummy_in_the_oauth_variable() {
274        let bindings = Bindings::default();
275        let env = bindings.register(
276            &registration(Provider::Anthropic),
277            "http://127.0.0.1:8787",
278            &anthropic_env(),
279        );
280        assert_eq!(
281            env.get("CLAUDE_CODE_OAUTH_TOKEN"),
282            Some(DUMMY_CREDENTIAL),
283            "the OAuth variable carries the dummy, not the real setup token",
284        );
285        assert!(
286            !env_contains_secret(&env, "real-oauth-token"),
287            "the real token never reaches the child"
288        );
289    }
290
291    #[test]
292    fn an_anthropic_account_never_passes_a_stray_auth_token_to_the_child() {
293        let bindings = Bindings::default();
294        let mut env = anthropic_env();
295        env.insert("ANTHROPIC_AUTH_TOKEN".to_owned(), "stray-base-url-key".to_owned());
296        let stamped =
297            bindings.register(&registration(Provider::Anthropic), "http://127.0.0.1:8787", &env);
298        assert_eq!(
299            stamped.get("ANTHROPIC_AUTH_TOKEN"),
300            None,
301            "the credential variable the account does not use is removed, not passed through",
302        );
303        assert!(
304            !env_contains_secret(&stamped, "stray-base-url-key"),
305            "a foreign credential never reaches the child",
306        );
307        assert_eq!(
308            stamped.get("CLAUDE_CODE_OAUTH_TOKEN"),
309            Some(DUMMY_CREDENTIAL),
310            "the account's own variable still carries the dummy",
311        );
312    }
313
314    #[test]
315    fn the_api_key_is_forced_empty() {
316        let bindings = Bindings::default();
317        let env = bindings.register(
318            &registration(Provider::Openrouter),
319            "http://127.0.0.1:8787",
320            &base_url_env(),
321        );
322        assert_eq!(
323            env.get("ANTHROPIC_API_KEY"),
324            Some(""),
325            "an empty x-api-key source is the same as unset"
326        );
327    }
328
329    #[test]
330    fn the_alt_base_url_is_never_left_pointing_elsewhere() {
331        let bindings = Bindings::default();
332        let mut env_with_alt = base_url_env();
333        env_with_alt
334            .insert("CLAUDE_CODE_API_BASE_URL".to_owned(), "http://localhost:18765".to_owned());
335        let env = bindings.register(
336            &registration(Provider::Openrouter),
337            "http://127.0.0.1:8787",
338            &env_with_alt,
339        );
340        assert_eq!(
341            env.get("CLAUDE_CODE_API_BASE_URL"),
342            Some("http://127.0.0.1:8787/Busytools/forge/session-1"),
343            "the lower-precedence variable must not bypass the redirect",
344        );
345    }
346
347    #[test]
348    fn unrelated_env_passes_through_verbatim() {
349        let bindings = Bindings::default();
350        let env = bindings.register(
351            &registration(Provider::Anthropic),
352            "http://127.0.0.1:8787",
353            &anthropic_env(),
354        );
355        assert_eq!(
356            env.get("CLAUDE_CODE_MAX_CONTEXT_TOKENS"),
357            Some("1000000"),
358            "an unrelated env key rides through untouched",
359        );
360    }
361
362    #[test]
363    fn reregistration_replaces_the_prior_binding() {
364        let bindings = Bindings::default();
365        let mut first = registration(Provider::Openrouter);
366        first.account = AccountKey("OldAccount".to_owned());
367        bindings.register(&first, "http://127.0.0.1:8787", &base_url_env());
368        bindings.register(
369            &registration(Provider::Openrouter),
370            "http://127.0.0.1:8787",
371            &base_url_env(),
372        );
373        assert_eq!(
374            bindings.binding_for("Busytools", "forge", "session-1"),
375            Some(AccountKey("OpenRouter".to_owned())),
376            "a respawn's registration replaces the previous generation's binding",
377        );
378    }
379
380    #[test]
381    fn a_respawn_binding_names_the_segment_the_child_runs_under() {
382        let bindings = Bindings::default();
383        bindings.register(
384            &registration(Provider::Anthropic),
385            "http://127.0.0.1:8787",
386            &anthropic_env(),
387        );
388        let mut respawned = registration(Provider::Anthropic);
389        respawned.session = "session-2".to_owned();
390
391        let overrides =
392            bindings.respawn_env_overrides(&respawned, "session-1", "http://127.0.0.1:8787");
393
394        assert_eq!(
395            overrides.get("ANTHROPIC_BASE_URL").map(String::as_str),
396            Some("http://127.0.0.1:8787/Busytools/forge/session-2"),
397            "the respawned child's base URL names the id it runs under",
398        );
399        assert_eq!(
400            bindings.binding_for("Busytools", "forge", "session-2"),
401            Some(AccountKey("OpenRouter".to_owned())),
402            "the binding moves to the new segment",
403        );
404        assert_eq!(
405            bindings.binding_for("Busytools", "forge", "session-1"),
406            None,
407            "the id the session left behind keeps no binding",
408        );
409    }
410
411    #[test]
412    fn an_unregistered_session_has_no_binding() {
413        let bindings = Bindings::default();
414        assert_eq!(bindings.binding_for("Busytools", "forge", "ghost"), None);
415    }
416
417    /// The set carries live values for the keys the gateway does not
418    /// own, so only the key names may print.
419    #[test]
420    fn a_value_never_reaches_the_env_set_debug_output() {
421        let set = EnvSet(vec![
422            ("MY_API_KEY".to_owned(), "env-set-supersecret".to_owned()),
423            ("ANTHROPIC_BASE_URL".to_owned(), "http://127.0.0.1:8787/x".to_owned()),
424        ]);
425        let rendered = format!("{set:?}");
426        assert!(!rendered.contains("env-set-supersecret"), "the stamped value leaked: {rendered}");
427        assert!(rendered.contains("MY_API_KEY"), "the key name stays readable: {rendered}");
428    }
429}
forge_gateway/binding.rs

forge_gateway/
binding.rs
