From 1b8ced3f243a0beebf1226808a1ee111ef55e1d6 Mon Sep 17 00:00:00 2001 From: Peter Jausovec Date: Mon, 19 Oct 2026 13:40:25 -0700 Subject: [PATCH 1/6] support ChatGPT subscription routing through the gateway Signed-off-by: Peter Jausovec --- README.md | 134 +- crates/agent/src/api.rs | 46 +- crates/agent/src/daemon.rs | 13 +- crates/agent/src/llm_proxy.rs | 2627 +---------------- crates/agent/src/llm_proxy/cache.rs | 306 ++ crates/agent/src/llm_proxy/routes.rs | 251 ++ crates/agent/src/llm_proxy/tests.rs | 2146 ++++++++++++++ .../src/provider/codex/integration_tests.rs | 105 + crates/agent/src/provider/codex/mod.rs | 37 +- crates/agent/src/provider/codex/reconcile.rs | 297 +- crates/agent/tests/common/gateway.rs | 22 +- crates/controller/src/service.rs | 26 +- crates/core/src/config.rs | 124 +- examples/codex-subscription/README.md | 80 + examples/codex-subscription/agentgateway.yaml | 36 + examples/codex-subscription/compose.yaml | 19 + examples/codex-subscription/config.yaml | 20 + examples/codex-subscription/dex.yaml | 24 + frontend/controller/package.json | 3 +- .../src/views/ConfigurationView.stories.tsx | 27 + .../src/views/ConfigurationView.tsx | 11 +- frontend/pnpm-lock.yaml | 93 +- schema/daemon-config.json | 13 +- schema/daemon-config.md | 3 +- 24 files changed, 3724 insertions(+), 2739 deletions(-) create mode 100644 crates/agent/src/llm_proxy/cache.rs create mode 100644 crates/agent/src/llm_proxy/routes.rs create mode 100644 crates/agent/src/llm_proxy/tests.rs create mode 100644 examples/codex-subscription/README.md create mode 100644 examples/codex-subscription/agentgateway.yaml create mode 100644 examples/codex-subscription/compose.yaml create mode 100644 examples/codex-subscription/config.yaml create mode 100644 examples/codex-subscription/dex.yaml diff --git a/README.md b/README.md index 3bf0e89e..c72e562e 100644 --- a/README.md +++ b/README.md @@ -96,9 +96,12 @@ for prerequisites, test credentials, and a walkthrough of the local services. ## Local LLM proxy -For clients without an API-key helper, enable the daemon's optional loopback proxy -in the local configuration file (startup settings are local only, not delivered by -the controller). The proxy is available in `--user` mode only: it hands out the +Use the daemon's optional loopback proxy for clients without an API-key helper, +or to preserve a client's provider subscription while authenticating separately +to the gateway. Codex ChatGPT subscriptions and VS Code Copilot passthrough use +this same proxy. Enable it in the local configuration file (startup settings are +local only, not delivered by the controller). The proxy is available in `--user` +mode only: it hands out the current user's gateway credential, so a system daemon rejects `daemon.llmProxy.listen`. If the address cannot be bound at startup, the daemon keeps running without the proxy and reports `llmProxy.bound: false` with the @@ -119,14 +122,19 @@ Then start the daemon as usual: agentdesktop daemon --user --config .yaml ``` -Point the client's model endpoint at `http://127.0.0.1:4000/v1/chat/completions`, +For a manually configured client, point its model endpoint at +`http://127.0.0.1:4000/v1/chat/completions`, `/v1/messages`, or `/v1/responses`, according to the protocol it uses, and give it the pairing header (see **Pairing** below). The proxy forwards request bodies (buffered, see below) to `llmGateway.proxyUrl` (or `llmGateway.url` when unset), streams the response back, and adds the credential from the same authentication flow used by the daemon's API-key helpers. Login and refresh stay with the -existing authentication implementation. No client API key is needed. Incoming -authorization/API-key headers are replaced, and upstream HTTP errors and +existing authentication implementation. In gateway-only mode, no client API key +is needed and incoming authorization/API-key headers are replaced. In passthrough +mode, the client's provider token is moved to `x-llm-token` while gateway identity +occupies `Authorization`; the gateway validates that identity before restoring +the provider bearer upstream. The managed integrations below generate their +client endpoints and pairing settings automatically. Upstream HTTP errors and redirects are passed back without following redirects. The incoming path and query are appended to the gateway base URL: with @@ -137,25 +145,23 @@ Avoid repeating `/v1` in both URLs. **Without the proxy.** When the daemon runs without its loopback proxy (`daemon.llmProxy.listen` unset, or the address could not be bound), `llmGateway.whenProxyUnavailable` decides what happens to the files of the -programs that use the proxy (`programs.copilot`, `programs.vscode`). With -`failClosed`, the default, the daemon leaves their managed entries as they -are: the Copilot CLI and VS Code stay pointed at the loopback port, and those -requests fail instead of reaching GitHub past the gateway. This needs entries -to be in place: on a first apply without the proxy, or after the entries were -removed by hand, nothing points at the loopback port, so VS Code on GitHub's -models talks to GitHub directly and the Copilot CLI and VS Code on own models -have no gateway models; the status says so. If the address could -not be bound because another process holds it, that process receives those -requests, with VS Code's Copilot tokens and the pairing value: use the default -on devices with a single user, as for the proxy in general. With -`failOpen`, the daemon removes the entries until the proxy is back: VS Code on -GitHub's models then talks to GitHub directly, and the Copilot CLI and VS Code -on own models lose the gateway's models. Either way the programs report +programs that use the proxy (`programs.copilot`, `programs.vscode`, and Codex +with `auth: subscription`). With `failClosed`, the default, existing managed +entries stay pointed at the loopback proxy and requests fail until it returns. +Before the first successful apply, or if those entries were removed, there is +no managed routing to retain: the tools keep their own settings. With `failOpen`, +the daemon removes its managed entries, including Codex's managed configuration +file, until the proxy returns. VS Code using GitHub's models and Codex using its +remaining native configuration may then connect directly to their provider; +gateway-only custom models are no longer available. If another process occupies +the loopback port, it can receive requests containing provider tokens and the +pairing value; reserve the port and use this setup on single-user devices. +Either way the programs report `inactive` with the reason, and the next apply with the proxy writes the current entries. A stopped daemon changes no file, so the tools fail until it runs again, whatever the policy. The policy covers only the files the daemon -writes: whether developers can reach GitHub's models without the gateway is -set by the organisation's Copilot policy and egress rules. Set the key only +writes: whether developers can reach their providers without the gateway is +set by the organisation's provider policies and egress rules. Set the key only after the controller and the daemons run a version that knows it; older ones reject a configuration that carries it. @@ -164,6 +170,7 @@ llmGateway: url: https://gateway.example.com whenProxyUnavailable: failOpen # default failClosed ``` + The proxy does not discover models or rewrite model IDs; configure those in the client. VS Code's built-in **Custom Endpoint** provider can use this endpoint without the experimental extension in `vscode/`. The provider reads @@ -196,7 +203,7 @@ the header is added. Controller JWT authentication uses client ID `vscode` by default on the prefix-less route; choose another ID with `daemon.llmProxy.clientId`. Every client -ID the proxy presents (`vscode` or your choice, plus `copilot-cli` and +ID the proxy presents (`vscode` or your choice, plus `codex`, `copilot-cli`, and `vscode-copilot` for the routes below) must be in the controller's `allowedClientIds`. With `authentication.type: oidc` the client ID has no effect: the OIDC token is the same for every route. The listen address must be loopback. @@ -204,7 +211,7 @@ OIDC token is the same for every route. The listen address must be loopback. **Routes.** Besides the prefix-less route (the manual shape used above, with the credential mode from `llmGateway.githubOAuth` and the upstream from `llmGateway.proxyUrl`), the proxy serves fixed path prefixes for the managed -Copilot programs, one per program route, so the client ID is decided by the file +programs, one per program route, so the client ID is decided by the file the daemon writes and never by request headers. The prefix is stripped before forwarding. A path whose first segment looks like a route name without matching one exactly is refused with `404 route_unknown` rather than served as the @@ -215,6 +222,7 @@ paths on the prefix-less route: | Prefix | Client ID | Sends upstream | Upstream | | --- | --- | --- | --- | | `/copilot-cli/` | `copilot-cli` | gateway identity only | `llmGateway.url` | +| `/codex/` | `codex` | gateway identity in `Authorization`, Codex's ChatGPT token in `x-llm-token`; HTTP Responses requests | `llmGateway.proxyUrl` (required) | | `/vscode-copilot/` | `vscode-copilot` | gateway identity only | `llmGateway.url` | | `/vscode-copilot-passthrough/` | `vscode-copilot` | gateway identity in `Authorization`, the client's own bearer token in `x-llm-token` | `llmGateway.proxyUrl` (required) | | `/vscode-copilot-capi//` | `vscode-copilot` | as the pass-through route; a credential-less `GET /_ping` is forwarded without `x-llm-token`; the only route that tunnels WebSocket upgrades | `llmGateway.proxyUrl` (required) | @@ -229,7 +237,10 @@ manual setup copies it from there. The state directory in `--user` mode is home directory from `HOME` or `USERPROFILE`, on every platform), unless `daemon.stateDir` says otherwise. The Copilot CLI's `providers.json` carries it as `"headers": {"x-agentdesktop-pairing": ""}` -on the provider entry. It keeps other local users on a shared host +on the provider entry, and Codex's generated provider carries it in `http_headers`. +The daemon +writes that Codex file owner-only and redacts its pairing-bearing diffs. Pairing +keeps other local users on a shared host and browser pages off the proxy, provided the client files that carry it are owner-only too. It does not restrict the current user's own processes, which can obtain a credential from the daemon directly. If the value cannot be created the @@ -310,6 +321,61 @@ VS Code on GitHub's models, with what to check at each step (drift, conflicts, removal) and the settings the GitHub organisation owner controls ([`examples/copilot/github-side.md`](examples/copilot/github-side.md)). +### Managed program: Codex ChatGPT subscriptions + +Set `programs.codex.auth: subscription` to keep Codex's own ChatGPT login while +the local proxy authenticates separately to the gateway: + +```yaml +daemon: + user: true + llmProxy: + listen: 127.0.0.1:4012 + +llmGateway: + url: https://gateway.example.com + proxyUrl: https://gateway.example.com/codex-subscription + authentication: + type: oidc + issuer: https://login.example.com + clientId: agentdesktop + +programs: + codex: + auth: subscription +``` + +`daemon.llmProxy.listen` enables the shared proxy locally; `proxyUrl` selects the +gateway's ChatGPT passthrough route. On controller-managed devices, distribute +`llmGateway` and `programs` through controller policy, but keep `daemon.*` local. +For controller JWT authentication, authorize client ID `codex` in +`allowedClientIds`. Upgrade both controller and daemon before using this option. + +Agentdesktop generates the Codex provider pointed at `/codex`, inserts the pairing +header, and enables ChatGPT authentication without a conflicting gateway +credential command. No manual Codex provider, `daemon.llmProxy.clientId`, or +`githubOAuth` configuration is needed. WebSockets are disabled for this HTTP-only +route. Keep Codex signed in with `codex login`; Codex owns the ChatGPT credential +and its refresh, while Agentdesktop obtains and refreshes gateway identity. + +The gateway must validate organization identity from `Authorization`, restore the +ChatGPT bearer from `x-llm-token`, and remove that transport header before +forwarding. Preserve Codex's account-identification headers and map +`/codex-subscription/responses` to +`https://chatgpt.com/backend-api/codex/responses`, not `api.openai.com/v1`. + +This requires a per-user daemon, not the Windows system service. Agentdesktop +manages `~/.codex/config.toml` in user mode; a pre-existing personal file without +its ownership header is reported as a conflict rather than overwritten. Migrate +needed settings into `programs.codex.managedConfig` and resolve that conflict +without removing Codex's separate login file. A `daemon.codex.config` override +must name an ordinary configuration layer Codex reads, not `requirements.toml`. + +See the [complete Codex subscription example](examples/codex-subscription/README.md) +for the gateway, sign-in, and production boundaries. The real Codex 0.153.4 CLI is +tested with synthetic credentials and a mock gateway, not a paid OpenAI account. +Validate live subscription access, token refresh, and provider terms before rollout. + ### Managed program: GitHub Copilot CLI `programs.copilot` makes the daemon write the Copilot CLI's BYOK provider @@ -766,10 +832,12 @@ usage: Then send agentgateway's access log to the fleet API with that token. `root` is the CA that signed the controller's TLS certificate. The added fields carry -the user, device, and tool from the controller JWT. Billing is derived from the -composite credential in the `Authorization` header, so subscription usage is -recognized for both controller-JWT and OIDC gateways. Usage without the added -fields is shown as unknown. +the user, device, and tool from the controller JWT. The billing expression below +recognizes Claude's composite subscription credential with either controller-JWT +or OIDC gateway authentication. Codex uses separate headers instead: classify its +dedicated passthrough route as subscription traffic, or use the controller JWT's +`billing` claim when available. Usage without the added identity fields is shown +as unknown. ```yaml frontendPolicies: @@ -797,7 +865,8 @@ frontendPolicies: Requests to models without a catalog price count toward tokens but not cost, and the UI says how many there were. - **Subscription (API rates)** is what requests billed to a Claude - subscription (`auth: subscription`) would have cost at API rates. It is + or Codex ChatGPT subscription (`auth: subscription`) would have cost at API + rates, when the gateway exports priced token usage for those requests. It is shown on its own and never added to cost. The controller marks the credentials it issues for such requests with a `billing: subscription` claim; daemons older than the controller's usage support don't ask for the @@ -805,7 +874,10 @@ frontendPolicies: - **Not metered:** VS Code with `copilotChat: githubModels` reaches GitHub through the gateway's pass-through route, which is not an LLM route, so its requests leave no usage records. The desktop app says so when the program is - configured that way. + configured that way. Codex's subscription example also uses raw HTTP + passthrough; forwarding and a subscription-billing claim alone do not extract + token counts or prices. The gateway must export those fields for such usage + to appear in the totals. - Hourly totals are kept for 90 days and daily totals for 13 months. agentgateway does not retry an export, so requests made while the controller is unreachable are not counted. Usage stays after a device is deleted. diff --git a/crates/agent/src/api.rs b/crates/agent/src/api.rs index e2f4fd07..72d31ea2 100644 --- a/crates/agent/src/api.rs +++ b/crates/agent/src/api.rs @@ -11,7 +11,7 @@ use tokio::sync::{mpsc, oneshot, watch}; use agentdesktop_core::{ config::{ - ControllerConnectionConfig, DaemonConfig, LlmGatewayAuthentication, ProgramAuthentication, + ControllerConnectionConfig, DaemonConfig, LlmGatewayAuthentication, SubscriptionAuth, VsCodeCopilotChat, valid_client_id, }, model::{ @@ -443,6 +443,25 @@ pub(crate) async fn gateway_credential( effective: &DaemonConfig, client_id: &str, include_subscription: bool, +) -> Result { + resolve_gateway_credential(state, effective, client_id, include_subscription, false).await +} + +pub(crate) async fn gateway_proxy_credential( + state: &AppState, + effective: &DaemonConfig, + client_id: &str, + subscription: bool, +) -> Result { + resolve_gateway_credential(state, effective, client_id, false, subscription).await +} + +async fn resolve_gateway_credential( + state: &AppState, + effective: &DaemonConfig, + client_id: &str, + include_subscription: bool, + passthrough_subscription: bool, ) -> Result { let gateway = effective.llm_gateway.as_ref().ok_or_else(|| { ( @@ -464,9 +483,14 @@ pub(crate) async fn gateway_credential( } // Local transport permissions authenticate the user, not the calling // process. The client ID selects an allowed policy within that boundary. - remote::llm_gateway_credential(controller, &state.state_dir, client_id, false) - .await - .map(|credential| (credential, false)) + remote::llm_gateway_credential( + controller, + &state.state_dir, + client_id, + passthrough_subscription, + ) + .await + .map(|credential| (credential, false)) } Some(LlmGatewayAuthentication::Oidc { issuer, @@ -669,19 +693,7 @@ async fn llm_gateway_login( } fn program_uses_subscription(config: &DaemonConfig, client_id: &str) -> bool { - match client_id { - "claude-code" => config - .programs - .claude_code - .as_ref() - .is_some_and(|program| program.auth == Some(ProgramAuthentication::Subscription)), - "claude-desktop" => config - .programs - .claude_desktop - .as_ref() - .is_some_and(|program| program.auth == Some(ProgramAuthentication::Subscription)), - _ => false, - } + config.programs.subscription_auth(client_id) == Some(SubscriptionAuth::CredentialHelper) } #[cfg(test)] diff --git a/crates/agent/src/daemon.rs b/crates/agent/src/daemon.rs index c60a9469..d624b6f2 100644 --- a/crates/agent/src/daemon.rs +++ b/crates/agent/src/daemon.rs @@ -951,14 +951,11 @@ fn start_gateway_authentication( }; let authentication = gateway.authentication.clone(); let github = gateway.github_oauth.clone().filter(|_| proxy_enabled); - let subscription = config.programs.claude_code.as_ref().is_some_and(|program| { - program.auth == Some(agentdesktop_core::config::ProgramAuthentication::Subscription) - }) || config - .programs - .claude_desktop - .as_ref() - .is_some_and(|program| { - program.auth == Some(agentdesktop_core::config::ProgramAuthentication::Subscription) + let subscription = ["claude-code", "claude-desktop"] + .into_iter() + .any(|client_id| { + config.programs.subscription_auth(client_id) + == Some(agentdesktop_core::config::SubscriptionAuth::CredentialHelper) }); if authentication.is_none() && !subscription { return; diff --git a/crates/agent/src/llm_proxy.rs b/crates/agent/src/llm_proxy.rs index 393224a6..95bf79d7 100644 --- a/crates/agent/src/llm_proxy.rs +++ b/crates/agent/src/llm_proxy.rs @@ -1,3 +1,15 @@ +mod cache; +mod routes; + +pub(crate) use routes::{CAPI_ROUTE, CODEX_ROUTE, PAIRING_HEADER}; +use routes::{ + RouteCredential, Upstream, bearer_or_header, constant_time_eq, host_is_loopback, match_route, + path_escapes_base, resembles_route, split_path_pairing, upstream_base, +}; + +pub(crate) use cache::CredentialCache; +use cache::{Deadline, now_unix}; + use std::{convert::Infallible, net::SocketAddr, path::Path, sync::Arc, time::Duration}; use anyhow::Context; @@ -17,7 +29,7 @@ use hyper_util::{ }; use tokio::{net::TcpListener, task::JoinSet}; -use agentdesktop_core::config::{GitHubTokenSource, LlmGatewayConfig}; +use agentdesktop_core::config::GitHubTokenSource; use crate::api::{self, AppState}; @@ -42,308 +54,9 @@ const RESPONSE_HEADERS_TIMEOUT: Duration = Duration::from_secs(600); /// front of it. Applies to controller-issued credentials only: an OIDC /// credential may need a browser login. const CREDENTIAL_FETCH_TIMEOUT: Duration = Duration::from_secs(20); -/// A cached controller credential is not used within this margin of its expiry. -const CREDENTIAL_EXPIRY_MARGIN: Duration = Duration::from_secs(30); -/// Longest a controller credential is served from the cache, counted from its -/// fetch. This bounds how long a revoked device keeps using the proxy: -/// revocation is enforced at issuance, the gateway validates tokens statelessly. -const CREDENTIAL_CACHE_TTL: Duration = Duration::from_secs(60); - -/// Listener-local cache of controller-issued gateway credentials, one per client -/// id. A controller fetch is a fresh mTLS connection and round trip; agent-mode -/// sessions issue bursts of requests, so a short cache keeps them off the -/// controller. Fetches are single-flight per key: concurrent misses wait for -/// the one fetch in progress and then read its result from the cache. The -/// entries lock is never held across a fetch, so a slow controller delays only -/// the requests waiting for that key, not cache hits on other routes. -/// OIDC credentials never go through here; the secret store is their cache. -#[derive(Default)] -pub(crate) struct CredentialCache { - entries: std::sync::Mutex>, - flights: std::sync::Mutex>>>, -} - -struct CachedCredential { - credential: String, - /// The enrolled device the credential was issued to. A logout or a - /// re-enrollment changes it, and the entry is then ignored. - device_id: String, - /// Monotonic deadline (wall-clock expiry at fetch time, minus the margin, - /// capped by CREDENTIAL_CACHE_TTL) and the wall-clock deadline itself. Both - /// must be in the future: the monotonic clock does not advance across a - /// suspend, the wall clock can jump; neither alone is trusted. - valid_until: std::time::Instant, - valid_until_unix: u64, - /// The credential's own expiry (the controller JWT `exp`), on both clocks: - /// a tunnel opened with the credential runs until then, not until the - /// cache stops reusing it. - expires: std::time::Instant, - expires_unix: u64, -} - -/// A point in time on both clocks (see `CachedCredential`): a tunnel closes -/// when either says the credential has expired. -#[derive(Clone, Copy, Debug)] -struct Deadline { - monotonic: std::time::Instant, - unix: u64, -} - -impl Deadline { - fn from_unix(expires_unix: u64) -> Self { - Self { - monotonic: std::time::Instant::now() - + Duration::from_secs(expires_unix.saturating_sub(now_unix())), - unix: expires_unix, - } - } -} - -fn now_unix() -> u64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_secs()) - .unwrap_or(0) -} - -impl CredentialCache { - /// The lock a fetch for the key holds; one per key, kept for the - /// listener's lifetime (keys are client id, gateway and audience). - fn flight(&self, key: &str) -> Arc> { - self.flights - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .entry(key.to_owned()) - .or_default() - .clone() - } - - /// The cached credential for the key, or one from `fetch`, which is cached - /// at once. Single-flight: a caller that finds a fetch in progress waits - /// for it and reads its result. Returns the credential, its own expiry, and - /// whether it came from the cache. - async fn get_or_fetch( - &self, - key: &str, - device_id: &str, - fetch: F, - ) -> Result<(String, Deadline, bool), E> - where - F: FnOnce() -> Fut, - Fut: - std::future::Future>, - { - let flight = self.flight(key); - let _flight = flight.lock().await; - if let Some((credential, deadline)) = self.lookup(key, device_id) { - return Ok((credential, deadline, true)); - } - let fetched = fetch().await?; - self.insert(key, device_id, &fetched); - let deadline = Deadline::from_unix(fetched.expires_at_unix_seconds); - Ok((fetched.credential, deadline, false)) - } - - /// A cached credential for the key and device, if one is still valid. - #[cfg(all(test, target_os = "linux"))] - fn get(&self, key: &str, device_id: &str) -> Option { - self.lookup(key, device_id) - .map(|(credential, _)| credential) - } - - /// A cached credential together with its own expiry, read in one go so a - /// tunnel opened with it always gets the deadline of the credential it - /// was actually opened with. - fn lookup(&self, key: &str, device_id: &str) -> Option<(String, Deadline)> { - let entries = self - .entries - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - entries - .get(key) - .filter(|entry| { - entry.device_id == device_id - && std::time::Instant::now() < entry.valid_until - && now_unix() < entry.valid_until_unix - }) - .map(|entry| { - ( - entry.credential.clone(), - Deadline { - monotonic: entry.expires, - unix: entry.expires_unix, - }, - ) - }) - } - - /// Remember a freshly fetched credential until the earlier of its expiry - /// margin and the cache TTL. - fn insert( - &self, - key: &str, - device_id: &str, - fetched: &agentdesktop_core::model::LlmGatewayCredential, - ) { - let now = now_unix(); - let remaining = Duration::from_secs(fetched.expires_at_unix_seconds.saturating_sub(now)); - let lifetime = remaining - .saturating_sub(CREDENTIAL_EXPIRY_MARGIN) - .min(CREDENTIAL_CACHE_TTL); - let mut entries = self - .entries - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - if lifetime > Duration::ZERO { - entries.insert( - key.to_owned(), - CachedCredential { - credential: fetched.credential.clone(), - device_id: device_id.to_owned(), - valid_until: std::time::Instant::now() + lifetime, - valid_until_unix: now + lifetime.as_secs(), - expires: std::time::Instant::now() + remaining, - expires_unix: fetched.expires_at_unix_seconds, - }, - ); - } else { - tracing::debug!( - key, - "controller credential already within its expiry margin; not cached" - ); - entries.remove(key); - } - } - - /// Drop the entry for the key if it still holds the given credential. A - /// concurrent request that already replaced it with a fresh one is left alone. - fn invalidate_if(&self, key: &str, credential: &str) { - let mut entries = self - .entries - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - if entries - .get(key) - .is_some_and(|entry| entry.credential == credential) - { - entries.remove(key); - } - } -} - -/// Header that carries the per-device pairing value on every route. -pub(crate) const PAIRING_HEADER: &str = "x-agentdesktop-pairing"; /// File under the daemon state directory that holds the pairing value. const PAIRING_FILE: &str = "llm-proxy-pairing"; -/// How a route obtains the credential that goes upstream in `x-llm-token`. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(crate) enum RouteCredential { - /// Only the gateway identity is sent; any client credential is dropped. - Gateway, - /// The client's own bearer token is moved to `x-llm-token` and the gateway - /// identity takes `Authorization` (the VS Code Copilot pass-through shape). - Passthrough, -} - -/// Which configured gateway base URL a route forwards to. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(crate) enum Upstream { - /// `llmGateway.url`, shared with the credential-helper programs. - Url, - /// `llmGateway.proxyUrl` when set, else `llmGateway.url`. - ProxyUrl, -} - -/// A fixed path prefix the proxy serves for one managed program. -/// -/// The prefix is what a reconciler writes into that program's client file, so -/// the client id is decided by the file the daemon owns, never by request -/// headers. The prefix is stripped before forwarding. -#[derive(Clone, Copy, Debug)] -pub(crate) struct ProxyRoute { - pub prefix: &'static str, - pub client_id: &'static str, - pub credential: RouteCredential, - pub upstream: Upstream, - /// The pairing is the first path segment after the prefix (the header is - /// ignored): for a client that cannot add a header to its requests. - pub pairing_in_path: bool, - /// WebSocket upgrades are tunnelled instead of refused. - pub tunnels_upgrades: bool, - /// A credential-less `GET /_ping` is forwarded without `x-llm-token`. - pub credential_less_ping: bool, -} - -/// VS Code Copilot Chat's CAPI pass-through route (the target of -/// `github.copilot.advanced.debug.overrideCapiUrl`): the pairing travels as -/// the first path segment after the prefix instead of in `PAIRING_HEADER`, -/// because VS Code sends these requests itself and has no setting that adds a -/// header to them. Also the only route a WebSocket upgrade is tunnelled on -/// (the `GET /responses` Auto/agent conversation). -pub(crate) const CAPI_ROUTE: &str = "/vscode-copilot-capi"; - -/// Routes for the Copilot programs. A prefix matches only at a `/` boundary, so -/// `/vscode-copilot-passthrough/...` never matches `/vscode-copilot`. -pub(crate) const ROUTES: &[ProxyRoute] = &[ - ProxyRoute { - prefix: "/vscode-copilot-passthrough", - client_id: "vscode-copilot", - credential: RouteCredential::Passthrough, - upstream: Upstream::ProxyUrl, - pairing_in_path: false, - tunnels_upgrades: false, - credential_less_ping: false, - }, - ProxyRoute { - prefix: CAPI_ROUTE, - client_id: "vscode-copilot", - credential: RouteCredential::Passthrough, - upstream: Upstream::ProxyUrl, - pairing_in_path: true, - tunnels_upgrades: true, - credential_less_ping: true, - }, - ProxyRoute { - prefix: "/vscode-copilot", - client_id: "vscode-copilot", - credential: RouteCredential::Gateway, - upstream: Upstream::Url, - pairing_in_path: false, - tunnels_upgrades: false, - credential_less_ping: false, - }, - ProxyRoute { - prefix: "/copilot-cli", - client_id: "copilot-cli", - credential: RouteCredential::Gateway, - upstream: Upstream::Url, - pairing_in_path: false, - tunnels_upgrades: false, - credential_less_ping: false, - }, -]; - -/// Splits `rest` (the path remainder after `CAPI_ROUTE`'s prefix, always -/// starting with `/`) into the first non-empty path segment (the pairing -/// value, still percent-encoded as sent) and the remainder starting with `/` -/// (or `/` itself when nothing follows), so the pairing segment never reaches -/// `match_route`'s caller, the forwarded upstream path, logs, or plan -/// reports. `None` when the first segment is missing or empty (bare prefix, -/// `/`, or `//...`). -/// -/// Called from `forward` for `CAPI_ROUTE` only, in place of the `PAIRING_HEADER` -/// check; the segment is compared in constant time and stripped before the -/// request is forwarded. -pub(crate) fn split_path_pairing(rest: &str) -> Option<(&str, &str)> { - let without_slash = rest.strip_prefix('/')?; - let (segment, remainder) = match without_slash.find('/') { - Some(index) => (&without_slash[..index], &without_slash[index..]), - None => (without_slash, "/"), - }; - (!segment.is_empty()).then_some((segment, remainder)) -} - /// Whether the request asks for a WebSocket upgrade: `Connection` lists the /// `upgrade` token and `Upgrade` lists `websocket`, both case-insensitive. fn is_websocket_upgrade(headers: &HeaderMap) -> bool { @@ -650,131 +363,6 @@ impl ProxyError { } } -/// Only loopback hosts: a page on another origin that resolves a name to -/// 127.0.0.1 (DNS rebinding) sends that name in `Host`. Accepts any loopback -/// IP (the listen validation allows all of 127.0.0.0/8 and ::1), bracketed or -/// bare IPv6, `localhost` in any case with an optional trailing dot, each with -/// or without a port. A request without `Host` (HTTP/1.0) is refused. -fn host_is_loopback(headers: &HeaderMap) -> bool { - let Some(host) = headers.get(HOST).and_then(|value| value.to_str().ok()) else { - return false; - }; - let host = host.trim(); - let name = if let Some(rest) = host.strip_prefix('[') { - match rest.split_once(']') { - Some((name, port)) if port.is_empty() || port.starts_with(':') => name, - _ => return false, - } - } else if host.matches(':').count() > 1 { - // Bare IPv6 without brackets cannot carry a port. - host - } else { - host.rsplit_once(':').map_or(host, |(name, _)| name) - }; - if let Ok(ip) = name.parse::() { - return ip.is_loopback(); - } - name.trim_end_matches('.').eq_ignore_ascii_case("localhost") -} - -/// Refuse `..` segments so a route cannot escape the upstream base it pins. -/// The check runs on a percent-decoded copy (the forwarded path stays as sent), -/// so `%2e%2e` and encoded separators are caught too; a malformed escape is -/// refused outright. -fn path_escapes_base(path: &str) -> bool { - let Some(decoded) = percent_decode(path) else { - return true; - }; - decoded - .split(|byte| *byte == b'/' || *byte == b'\\') - .any(|segment| segment == b"..") -} - -/// Decodes `%XX` escapes to bytes. Escapes that are not two hex digits are -/// malformed and yield `None`; the bytes are not required to be UTF-8, since -/// the check only looks for separators and dots. -fn percent_decode(input: &str) -> Option> { - let bytes = input.as_bytes(); - let mut out = Vec::with_capacity(bytes.len()); - let mut i = 0; - while i < bytes.len() { - if bytes[i] == b'%' { - let hex = bytes.get(i + 1..i + 3)?; - if !hex.iter().all(u8::is_ascii_hexdigit) { - return None; - } - let value = u8::from_str_radix(std::str::from_utf8(hex).ok()?, 16).ok()?; - out.push(value); - i += 3; - } else { - out.push(bytes[i]); - i += 1; - } - } - Some(out) -} - -/// Whether the first path segment looks like one of the managed route names -/// without matching one exactly (wrong case, extra suffix, doubled slash). -/// A gateway sub-path whose first segment starts with a route name (for -/// example `/vscode-copilot-proxy/...`) is refused on the prefix-less route -/// as well; the check is on the literal path, so an encoded near miss is not -/// caught and reaches the gateway on the prefix-less route. -fn resembles_route(path: &str) -> bool { - let first = path - .trim_start_matches('/') - .split('/') - .next() - .unwrap_or("") - .to_ascii_lowercase(); - !first.is_empty() - && ROUTES.iter().any(|route| { - let name = &route.prefix[1..]; - first == name || first.starts_with(name) || name.starts_with(&first) && first.len() >= 8 - }) -} - -fn constant_time_eq(a: &[u8], b: &[u8]) -> bool { - if a.len() != b.len() { - return false; - } - a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0 -} - -/// The route for a request path, and the path with the prefix removed. -fn match_route(path: &str) -> Option<(&'static ProxyRoute, &str)> { - ROUTES.iter().find_map(|route| { - let rest = path.strip_prefix(route.prefix)?; - (rest.is_empty() || rest.starts_with('/')) - .then_some((route, if rest.is_empty() { "/" } else { rest })) - }) -} - -fn bearer_or_header(headers: &HeaderMap, header: &str) -> Option { - headers - .get(header) - .or_else(|| headers.get(AUTHORIZATION)) - .and_then(|value| value.to_str().ok()) - .map(|value| strip_bearer(value).trim().to_owned()) - .filter(|value| !value.is_empty()) -} - -/// The token of a `Bearer` credential; authentication schemes are -/// case-insensitive (RFC 9110, section 11.1). -fn strip_bearer(value: &str) -> &str { - match value.split_once(' ') { - Some((scheme, token)) if scheme.eq_ignore_ascii_case("bearer") => token, - _ => value, - } -} - -fn upstream_base(gateway: &LlmGatewayConfig, upstream: Upstream) -> &url::Url { - match upstream { - Upstream::Url => &gateway.url, - Upstream::ProxyUrl => gateway.proxy_url.as_ref().unwrap_or(&gateway.url), - } -} - async fn forward( mut request: Request, client: &ProxyClient, @@ -1051,6 +639,10 @@ async fn forward( }, }; let client_token_sent = upstream_token.is_some(); + let subscription = client_token_sent + && matched_route.is_some_and(|route| route.prefix == CODEX_ROUTE) + && effective.programs.subscription_auth(&client_id) + == Some(agentdesktop_core::config::SubscriptionAuth::ClientPassthrough); if let Some(token) = upstream_token { let mut value = HeaderValue::from_str(&token).map_err(|_| { ProxyError::new( @@ -1099,10 +691,11 @@ async fn forward( .ok() .flatten() .map(|device_id| { - ( - format!("{client_id}\u{0}{}\u{0}{audience}", gateway.url), - device_id, - ) + let mut key = format!("{client_id}\u{0}{}\u{0}{audience}", gateway.url); + if subscription { + key.push_str("\0subscription"); + } + (key, device_id) }), _ => None, }; @@ -1144,8 +737,14 @@ async fn forward( let cache = cache.clone(); let cache_key = cache_key.clone(); tokio::spawn(async move { - let fetch = - || api::gateway_credential(&state, &effective, &client_id, false); + let fetch = || { + api::gateway_proxy_credential( + &state, + &effective, + &client_id, + subscription, + ) + }; match &cache_key { Some((key, device_id)) => { cache.get_or_fetch(key, device_id, fetch).await @@ -1313,2168 +912,4 @@ fn strip_hop_headers(headers: &mut HeaderMap) { } #[cfg(all(test, target_os = "linux"))] -mod tests { - use super::*; - use crate::{enrollment::EnrollmentState, secret_store::SecretStore}; - use agentdesktop_core::{ - config::parse_daemon, - model::{DaemonInfo, DaemonScope, Discovery}, - }; - use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; - use http_body_util::StreamBody; - use hyper::body::Frame; - use sha2::{Digest, Sha256}; - use std::sync::Arc; - use tokio::sync::{mpsc, oneshot, watch}; - use tokio_stream::wrappers::ReceiverStream; - - // Covers the real socket path, credential replacement/rotation, opaque bytes, - // header forwarding, and delivery before the upstream response is complete. - #[test] - fn bearer_scheme_is_case_insensitive() { - for value in ["Bearer tok", "bearer tok", "BEARER tok", "tok"] { - let mut headers = HeaderMap::new(); - headers.insert(AUTHORIZATION, HeaderValue::from_str(value).unwrap()); - assert_eq!( - bearer_or_header(&headers, "x-llm-token").as_deref(), - Some("tok"), - "{value}" - ); - } - let mut headers = HeaderMap::new(); - headers.insert("x-llm-token", HeaderValue::from_static("bearer gho_x")); - assert_eq!( - bearer_or_header(&headers, "x-llm-token").as_deref(), - Some("gho_x") - ); - } - - #[tokio::test] - async fn streams_opaque_requests_with_current_credentials() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let (release_tx, release_rx) = oneshot::channel(); - let release = Arc::new(tokio::sync::Mutex::new(Some(release_rx))); - let upstream_task = tokio::spawn(async move { - let (socket, _) = upstream.accept().await.unwrap(); - hyper::server::conn::http1::Builder::new().serve_connection( - TokioIo::new(socket), service_fn(move |request: Request| { - let release = release.clone(); - async move { - assert_eq!(request.uri(), "/gateway/v1/messages?value=%2F"); - assert_eq!(request.headers()[HOST], address.to_string()); - assert!(!request.headers().contains_key("x-api-key")); - assert!(!request.headers().contains_key("x-hop")); - assert!(!request.headers().contains_key(PAIRING_HEADER)); - assert_eq!(request.headers()["x-llm-token"], "ghu_test"); - assert_eq!(request.headers()["anthropic-version"], "2023-06-01"); - let gate = release.lock().await.take(); - assert_eq!(request.headers()[AUTHORIZATION], - if gate.is_some() { "Bearer first" } else { "Bearer second" }); - assert_eq!(request.into_body().collect().await.unwrap().to_bytes(), - Bytes::from_static(b"\xffnot-json\x00")); - let (tx, rx) = mpsc::channel::, Infallible>>(2); - tokio::spawn(async move { - tx.send(Ok(Frame::data(Bytes::from_static(b"data: first\n\n")))).await.unwrap(); - if let Some(gate) = gate { gate.await.unwrap(); } - let _ = tx.send(Ok(Frame::data(Bytes::from_static(b"data: last\n\n")))).await; - }); - Ok::<_, Infallible>(Response::builder() - .header("content-type", "text/event-stream") - .header("connection", "x-upstream-hop") - .header("x-upstream-hop", "remove") - .body(StreamBody::new(ReceiverStream::new(rx))).unwrap()) - } - }) - ).await.unwrap(); - }); - let dir = tempfile::tempdir().unwrap(); - let store = SecretStore::new(dir.path()).unwrap(); - let account = URL_SAFE_NO_PAD.encode(Sha256::digest(b"https://issuer.example/\0proxy-test")); - let save_token = |token: &str| store.set("dev.agentdesktop.gateway-oidc", &account, - &serde_json::json!({"accessToken": token, "refreshToken": null, - "expiresAtUnixSeconds": 4_000_000_000u64, - "tokenEndpoint": "https://issuer.example/token"}).to_string()).unwrap(); - save_token("first"); - store.set("dev.agentdesktop.github-oauth", "github-test", &serde_json::json!({ - "accessToken": "ghu_test", "expiresAt": 4_000_000_000u64, - "refreshToken": null, "refreshExpiresAt": 0, - }).to_string()).unwrap(); - let config = parse_daemon(&format!("llmGateway:\n url: http://{address}/gateway\n githubOAuth:\n clientId: github-test\n authentication:\n type: oidc\n issuer: https://issuer.example/\n clientId: proxy-test\n")).unwrap(); - let (_, discovery) = watch::channel(Arc::new(Discovery { agents: vec![], model_runtimes: vec![] })); - let state = AppState { - controller_status: None, - config, - daemon_info: DaemonInfo { version: "test".into(), scope: DaemonScope::User, - config_path: String::new(), state_directory: String::new(), - inventory_interval: Duration::from_secs(60), controller: None, llm_proxy: None, - usage_receiver: None }, - discovery, enrollment: EnrollmentState::new(false), state_dir: dir.path().to_owned(), - oidc_callback_listen: None, telemetry: None, logout: None, usage: None, - programs: None, - }; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = Client::builder(TokioExecutor::new()).build_http(); - let mut release_tx = Some(release_tx); - for _ in 0..2 { - let request = Request::builder().method(Method::POST) - .uri(format!("http://{proxy_address}/v1/messages?value=%2F")) - .header(AUTHORIZATION, "Bearer must-not-forward") - .header("x-api-key", "must-not-forward") - .header("x-llm-token", "must-not-forward") - .header(PAIRING_HEADER, PAIRING) - .header(CONNECTION, "x-hop").header("x-hop", "remove") - .header("anthropic-version", "2023-06-01") - .body(Full::new(Bytes::from_static(b"\xffnot-json\x00"))).unwrap(); - let response = client.request(request).await.unwrap(); - assert_eq!(response.status(), StatusCode::OK); - assert_eq!(response.headers()["content-type"], "text/event-stream"); - assert!(!response.headers().contains_key("x-upstream-hop")); - let mut body = response.into_body(); - let first = body.frame().await.unwrap().unwrap().into_data().unwrap(); - if let Some(release) = release_tx.take() { - assert_eq!(first, "data: first\n\n"); - release.send(()).unwrap(); - } - let rest = body.collect().await.unwrap().to_bytes(); - assert_eq!([first.as_ref(), rest.as_ref()].concat(), b"data: first\n\ndata: last\n\n"); - save_token("second"); - } - proxy.abort(); - let _ = proxy.await; - drop(client); - upstream_task.abort(); - }).await.unwrap(); - } - - // source: request. The client's own credential is what reaches the gateway - // in x-llm-token, and the daemon contributes only the gateway identity. - // This is the VS Code Copilot shape: Copilot has already exchanged its - // GitHub session for a provider token and sends it in Authorization. - #[tokio::test] - async fn forwards_the_client_credential_when_source_is_request() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let upstream_task = tokio::spawn(async move { - let (socket, _) = upstream.accept().await.unwrap(); - hyper::server::conn::http1::Builder::new() - .serve_connection( - TokioIo::new(socket), - service_fn(move |request: Request| async move { - // The Bearer prefix is stripped here and re-added by the - // gateway, so the token travels bare. - assert_eq!(request.headers()["x-llm-token"], "tid=from-client"); - assert_eq!(request.headers()[AUTHORIZATION], "Bearer identity"); - Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static(b"ok")))) - }), - ) - .await - .unwrap(); - }); - let (dir, state) = request_source_state(address).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let request = Request::builder() - .method(Method::POST) - .uri(format!("http://{proxy_address}/chat/completions")) - .header(PAIRING_HEADER, PAIRING) - .header(AUTHORIZATION, "Bearer tid=from-client") - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - let response = client.request(request).await.unwrap(); - assert_eq!(response.status(), StatusCode::OK); - proxy.abort(); - let _ = proxy.await; - drop(client); - upstream_task.abort(); - drop(dir); - }) - .await - .unwrap(); - } - - // Without a client credential the gateway would swap in an empty - // Authorization and the provider would fail in a way that looks like a - // gateway fault, so the proxy refuses instead of forwarding. - #[tokio::test] - async fn rejects_a_request_with_no_client_credential() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let (dir, state) = request_source_state(address).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let request = Request::builder() - .method(Method::POST) - .uri(format!("http://{proxy_address}/chat/completions")) - .header(PAIRING_HEADER, PAIRING) - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - let response = client.request(request).await.unwrap(); - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); - // The default route is paired too: without the header it is refused - // before any credential work. - let request = Request::builder() - .method(Method::POST) - .uri(format!("http://{proxy_address}/chat/completions")) - .header(AUTHORIZATION, "Bearer tid=from-client") - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - assert_eq!( - error_code(client.request(request).await.unwrap()).await, - (StatusCode::FORBIDDEN, "pairing_invalid".to_owned()) - ); - proxy.abort(); - let _ = proxy.await; - drop(client); - drop(dir); - }) - .await - .unwrap(); - } - - const PAIRING: &str = "pairing-secret-value-0123456789abcdef"; - - fn default_config() -> ProxyConfig { - ProxyConfig { - default_client_id: "vscode".into(), - pairing: Arc::from(PAIRING), - } - } - - async fn error_code(response: Response) -> (StatusCode, String) { - let status = response.status(); - assert_eq!(response.headers()[CONTENT_TYPE], "application/json"); - let body = response.into_body().collect().await.unwrap().to_bytes(); - let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); - (status, json["error"]["code"].as_str().unwrap().to_owned()) - } - - #[tokio::test] - async fn refuses_non_loopback_hosts_and_options_with_json_errors() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let (dir, state) = request_source_state(upstream.local_addr().unwrap()).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let rebound = Request::builder() - .method(Method::POST) - .uri(format!( - "http://{proxy_address}/copilot-cli/v1/chat/completions" - )) - .header(HOST, "attacker.example") - .header(PAIRING_HEADER, "pairing-secret-value-0123456789abcdef") - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - assert_eq!( - error_code(client.request(rebound).await.unwrap()).await, - (StatusCode::FORBIDDEN, "host_not_allowed".to_owned()) - ); - let options = Request::builder() - .method(Method::OPTIONS) - .uri(format!( - "http://{proxy_address}/copilot-cli/v1/chat/completions" - )) - .body(Full::new(Bytes::new())) - .unwrap(); - assert_eq!( - error_code(client.request(options).await.unwrap()).await, - ( - StatusCode::METHOD_NOT_ALLOWED, - "method_not_allowed".to_owned() - ) - ); - proxy.abort(); - let _ = proxy.await; - drop(client); - drop(dir); - }) - .await - .unwrap(); - } - - // A prefixed route needs the pairing value, strips its prefix, sends only the - // gateway identity (no client credential leaks through), and is refused with a - // JSON error without the pairing. - #[tokio::test] - async fn prefixed_gateway_route_requires_pairing_and_strips_the_prefix() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let upstream_task = tokio::spawn(async move { - let (socket, _) = upstream.accept().await.unwrap(); - hyper::server::conn::http1::Builder::new() - .serve_connection( - TokioIo::new(socket), - service_fn(move |request: Request| async move { - assert_eq!(request.uri(), "/v1/chat/completions?stream=true"); - assert_eq!(request.headers()[AUTHORIZATION], "Bearer identity"); - assert!(!request.headers().contains_key("x-llm-token")); - assert!(!request.headers().contains_key(PAIRING_HEADER)); - Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static(b"ok")))) - }), - ) - .await - .unwrap(); - }); - let (dir, state) = request_source_state(address).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state.clone(), default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let unpaired = Request::builder() - .method(Method::POST) - .uri(format!( - "http://{proxy_address}/copilot-cli/v1/chat/completions" - )) - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - assert_eq!( - error_code(client.request(unpaired).await.unwrap()).await, - (StatusCode::FORBIDDEN, "pairing_invalid".to_owned()) - ); - // The pairing is only accepted in its own header, never as the bearer. - let as_bearer = Request::builder() - .method(Method::POST) - .uri(format!( - "http://{proxy_address}/copilot-cli/v1/chat/completions" - )) - .header(AUTHORIZATION, format!("Bearer {PAIRING}")) - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - assert_eq!( - error_code(client.request(as_bearer).await.unwrap()).await, - (StatusCode::FORBIDDEN, "pairing_invalid".to_owned()) - ); - let wrong = Request::builder() - .method(Method::POST) - .uri(format!( - "http://{proxy_address}/copilot-cli/v1/chat/completions" - )) - .header(PAIRING_HEADER, "pairing-secret-value-0123456789abcdeX") - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - assert_eq!( - error_code(client.request(wrong).await.unwrap()).await, - (StatusCode::FORBIDDEN, "pairing_invalid".to_owned()) - ); - let paired = Request::builder() - .method(Method::POST) - .uri(format!( - "http://{proxy_address}/copilot-cli/v1/chat/completions?stream=true" - )) - .header(PAIRING_HEADER, PAIRING) - .header(AUTHORIZATION, "Bearer a-real-key-by-mistake") - .header("x-llm-token", "must-not-forward") - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - let response = client.request(paired).await.unwrap(); - assert_eq!(response.status(), StatusCode::OK); - proxy.abort(); - let _ = proxy.await; - // A '..' segment cannot escape the upstream base the route pins. - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let request = Request::builder() - .method(Method::POST) - .uri(format!("http://{proxy_address}/copilot-cli/../admin")) - .header(PAIRING_HEADER, PAIRING) - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - assert_eq!( - error_code(client.request(request).await.unwrap()).await, - (StatusCode::BAD_REQUEST, "path_invalid".to_owned()) - ); - proxy.abort(); - let _ = proxy.await; - drop(client); - upstream_task.abort(); - drop(dir); - }) - .await - .unwrap(); - } - - // The pass-through route moves the client's token to x-llm-token, keeps the - // gateway identity in Authorization, forwards to proxyUrl, and refuses a - // request without a client token. The pairing must come in its own header - // here, because Authorization carries the client token. - #[tokio::test] - async fn passthrough_route_moves_the_client_token_and_uses_proxy_url() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let upstream_task = tokio::spawn(async move { - let (socket, _) = upstream.accept().await.unwrap(); - hyper::server::conn::http1::Builder::new() - .serve_connection( - TokioIo::new(socket), - service_fn(move |request: Request| async move { - assert_eq!(request.uri(), "/copilot-proxy/chat/completions"); - assert_eq!(request.headers()["x-llm-token"], "tid=from-client"); - assert_eq!(request.headers()[AUTHORIZATION], "Bearer identity"); - assert!(!request.headers().contains_key(PAIRING_HEADER)); - Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static(b"ok")))) - }), - ) - .await - .unwrap(); - }); - let (dir, mut state) = request_source_state(address).await; - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - // Without proxyUrl the pass-through route is unavailable rather than - // sending x-llm-token to the plain LLM route. - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state.clone(), default_config())); - let request = Request::builder() - .method(Method::POST) - .uri(format!( - "http://{proxy_address}/vscode-copilot-passthrough/chat/completions" - )) - .header(PAIRING_HEADER, PAIRING) - .header(AUTHORIZATION, "Bearer tid=from-client") - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - assert_eq!( - error_code(client.request(request).await.unwrap()).await, - ( - StatusCode::SERVICE_UNAVAILABLE, - "agentdesktop_unavailable".to_owned() - ) - ); - proxy.abort(); - let _ = proxy.await; - state.config.llm_gateway.as_mut().unwrap().proxy_url = - Some(format!("http://{address}/copilot-proxy").parse().unwrap()); - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let no_token = Request::builder() - .method(Method::POST) - .uri(format!( - "http://{proxy_address}/vscode-copilot-passthrough/chat/completions" - )) - .header(PAIRING_HEADER, "pairing-secret-value-0123456789abcdef") - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - assert_eq!( - error_code(client.request(no_token).await.unwrap()).await, - ( - StatusCode::UNAUTHORIZED, - "client_credential_missing".to_owned() - ) - ); - // The pairing in Authorization is not a client token: refused, never forwarded. - let pairing_as_token = Request::builder() - .method(Method::POST) - .uri(format!( - "http://{proxy_address}/vscode-copilot-passthrough/chat/completions" - )) - .header(PAIRING_HEADER, PAIRING) - .header(AUTHORIZATION, format!("Bearer {PAIRING}")) - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - assert_eq!( - error_code(client.request(pairing_as_token).await.unwrap()).await, - ( - StatusCode::UNAUTHORIZED, - "client_credential_missing".to_owned() - ) - ); - let request = Request::builder() - .method(Method::POST) - .uri(format!( - "http://{proxy_address}/vscode-copilot-passthrough/chat/completions" - )) - .header(PAIRING_HEADER, "pairing-secret-value-0123456789abcdef") - .header(AUTHORIZATION, "Bearer tid=from-client") - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - let response = client.request(request).await.unwrap(); - assert_eq!(response.status(), StatusCode::OK); - proxy.abort(); - let _ = proxy.await; - drop(client); - upstream_task.abort(); - drop(dir); - }) - .await - .unwrap(); - } - - #[test] - fn route_matching_prefers_the_longer_prefix_and_strips_it() { - let (route, rest) = match_route("/vscode-copilot-passthrough/chat/completions").unwrap(); - assert_eq!(route.client_id, "vscode-copilot"); - assert_eq!(route.credential, RouteCredential::Passthrough); - assert_eq!(rest, "/chat/completions"); - let (route, rest) = match_route("/vscode-copilot/v1/chat/completions").unwrap(); - assert_eq!(route.credential, RouteCredential::Gateway); - assert_eq!(rest, "/v1/chat/completions"); - assert_eq!(match_route("/copilot-cli").unwrap().1, "/"); - assert!(match_route("/copilot-client/v1").is_none()); - assert!(match_route("/v1/chat/completions").is_none()); - } - - #[test] - fn host_check_accepts_loopback_forms_only() { - let mut headers = HeaderMap::new(); - let accepted = [ - "127.0.0.1:18095", - "localhost:18095", - "[::1]:18095", - "localhost", - "127.0.0.1", - "127.0.0.2:1", - "Localhost:18095", - "LOCALHOST", - "localhost.:18095", - "[::1]", - "::1", - ]; - for ok in accepted { - headers.insert(HOST, HeaderValue::from_static(ok)); - assert!(host_is_loopback(&headers), "{ok}"); - } - let refused = [ - "attacker.example", - "127.0.0.1.attacker.example:18095", - "10.0.0.5:18095", - "", - "[::1", - "[fe80::1]:18095", - "localhost.attacker.example", - "127.0.0.1:18095:1", - ]; - for bad in refused { - headers.insert(HOST, HeaderValue::from_static(bad)); - assert!(!host_is_loopback(&headers), "{bad}"); - } - headers.remove(HOST); - assert!(!host_is_loopback(&headers)); - } - - #[test] - fn pairing_is_created_once_and_reused() { - let dir = tempfile::tempdir().unwrap(); - let first = load_or_create_pairing(dir.path()).unwrap(); - let second = load_or_create_pairing(dir.path()).unwrap(); - assert_eq!(first, second); - assert!(first.len() >= 43); - let metadata = std::fs::metadata(dir.path().join(PAIRING_FILE)).unwrap(); - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - assert_eq!(metadata.permissions().mode() & 0o777, 0o600); - } - let _ = metadata; - } - - #[test] - fn pairing_file_is_regenerated_only_when_unusable() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join(PAIRING_FILE); - // Too short: replaced. - std::fs::write(&path, b"short").unwrap(); - let regenerated = load_or_create_pairing(dir.path()).unwrap(); - assert!(regenerated.len() >= 43); - assert_eq!(std::fs::read_to_string(&path).unwrap(), *regenerated); - // Not text: replaced. - std::fs::write(&path, [0xff, 0xfe, 0x00, 0x01]).unwrap(); - let again = load_or_create_pairing(dir.path()).unwrap(); - assert_ne!(again, regenerated); - assert_eq!(std::fs::read_to_string(&path).unwrap(), *again); - // Any other read error is an error, and the value on disk is untouched. - let other = tempfile::tempdir().unwrap(); - std::fs::create_dir(other.path().join(PAIRING_FILE)).unwrap(); - let error = load_or_create_pairing(other.path()).unwrap_err(); - assert!(format!("{error:#}").contains("read "), "{error:#}"); - assert!(other.path().join(PAIRING_FILE).is_dir()); - // An existing value that cannot be read (permissions) is kept as is. - #[cfg(unix)] - if !nix_is_root() { - use std::os::unix::fs::PermissionsExt; - std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o000)).unwrap(); - let error = load_or_create_pairing(dir.path()).unwrap_err(); - assert!(format!("{error:#}").contains("read "), "{error:#}"); - std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600)).unwrap(); - assert_eq!(*load_or_create_pairing(dir.path()).unwrap(), *again); - } - // A state directory that cannot be created is an error, not a pairing. - let blocked = tempfile::tempdir().unwrap(); - let file = blocked.path().join("state"); - std::fs::write(&file, b"x").unwrap(); - assert!(load_or_create_pairing(&file.join("nested")).is_err()); - } - - #[cfg(unix)] - fn nix_is_root() -> bool { - std::fs::metadata("/proc/self") - .map(|m| { - use std::os::unix::fs::MetadataExt; - m.uid() == 0 - }) - .unwrap_or(false) - } - - #[test] - fn route_table_maps_each_prefix_to_its_identity_mode_and_upstream() { - let expected = [ - ( - "/copilot-cli/v1/x", - "copilot-cli", - RouteCredential::Gateway, - Upstream::Url, - ), - ( - "/vscode-copilot/v1/x", - "vscode-copilot", - RouteCredential::Gateway, - Upstream::Url, - ), - ( - "/vscode-copilot-passthrough/chat/completions", - "vscode-copilot", - RouteCredential::Passthrough, - Upstream::ProxyUrl, - ), - ( - "/vscode-copilot-capi/PAIRING/responses", - "vscode-copilot", - RouteCredential::Passthrough, - Upstream::ProxyUrl, - ), - ]; - for (path, client_id, credential, upstream) in expected { - let (route, _) = match_route(path).unwrap_or_else(|| panic!("{path}")); - assert_eq!(route.client_id, client_id, "{path}"); - assert_eq!(route.credential, credential, "{path}"); - assert_eq!(route.upstream, upstream, "{path}"); - } - assert_eq!(ROUTES.len(), 4); - assert_eq!(CAPI_ROUTE, "/vscode-copilot-capi"); - } - - #[test] - fn path_guard_catches_encoded_traversal_and_near_miss_routes() { - for escaping in [ - "/copilot-cli/../admin", - "/copilot-cli/%2e%2e/admin", - "/copilot-cli/%2E%2E/admin", - "/copilot-cli/%2e%2e%2fadmin", - "/copilot-cli/..%5cadmin", - "/copilot-cli/%zz", - "/copilot-cli/%+f", - "/copilot-cli/%2", - ] { - assert!(path_escapes_base(escaping), "{escaping}"); - } - for fine in [ - "/copilot-cli/v1/chat/completions", - "/v1/messages", - "/copilot-cli/a.b/c%20d", - "/copilot-cli/caf%e9", - "/copilot-cli/%2e%2ex", - ] { - assert!(!path_escapes_base(fine), "{fine}"); - } - for near in [ - "/copilot-cli2/v1", - "/Copilot-cli/v1", - "//copilot-cli/v1", - "/vscode-copilotx/v1", - ] { - assert!( - match_route(near).is_none() && resembles_route(near), - "{near}" - ); - } - for other in ["/v1/chat/completions", "/chat/completions", "/models"] { - assert!(!resembles_route(other), "{other}"); - } - } - - // The pairing is checked before any configuration is read: an unpaired - // caller gets 403 even when the applied configuration is unreadable, and a - // paired caller gets a generic message without file detail. - #[tokio::test] - async fn pairing_is_checked_before_configuration_is_read() { - tokio::time::timeout(Duration::from_secs(15), async { - let dir = tempfile::tempdir().unwrap(); - std::fs::write( - dir.path().join("remote-config.yaml"), - "llmGateway: [not: an: object", - ) - .unwrap(); - let config = parse_daemon("controller:\n address: https://127.0.0.1:1\n").unwrap(); - let (_, discovery) = watch::channel(Arc::new(Discovery { - agents: vec![], - model_runtimes: vec![], - })); - let state = AppState { - controller_status: None, - config, - daemon_info: DaemonInfo { - version: "test".into(), - scope: DaemonScope::User, - config_path: String::new(), - state_directory: String::new(), - inventory_interval: Duration::from_secs(60), - controller: None, - llm_proxy: None, - usage_receiver: None, - }, - discovery, - enrollment: EnrollmentState::new(true), - state_dir: dir.path().to_owned(), - oidc_callback_listen: None, - telemetry: None, - logout: None, - usage: None, - programs: None, - }; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let unpaired = Request::builder() - .method(Method::POST) - .uri(format!( - "http://{proxy_address}/copilot-cli/v1/chat/completions" - )) - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - assert_eq!( - error_code(client.request(unpaired).await.unwrap()).await, - (StatusCode::FORBIDDEN, "pairing_invalid".to_owned()) - ); - // The path checks are pure string checks and run before the - // configuration is read: a paired caller with a bad path gets the - // path error, not the configuration error. - for (bad_path, code) in [ - ("/copilot-cli/../admin", "path_invalid"), - ("/Copilot-cli/v1/chat/completions", "route_unknown"), - ] { - let request = Request::builder() - .method(Method::POST) - .uri(format!("http://{proxy_address}{bad_path}")) - .header(PAIRING_HEADER, PAIRING) - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - let (_, got) = error_code(client.request(request).await.unwrap()).await; - assert_eq!(got, code, "{bad_path}"); - } - let paired = Request::builder() - .method(Method::POST) - .uri(format!( - "http://{proxy_address}/copilot-cli/v1/chat/completions" - )) - .header(PAIRING_HEADER, PAIRING) - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - let response = client.request(paired).await.unwrap(); - assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); - let body = response.into_body().collect().await.unwrap().to_bytes(); - let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); - let message = json["error"]["message"].as_str().unwrap(); - assert!( - !message.contains("remote-config") - && !message.contains(dir.path().to_str().unwrap()), - "{message}" - ); - proxy.abort(); - let _ = proxy.await; - drop(client); - drop(dir); - }) - .await - .unwrap(); - } - - // The retry applies only to a credential served from the cache. OIDC - // credentials are never cached, so a 401 from the gateway is returned as is - // after a single attempt, with the buffered body delivered once. The cached - // path is covered by `cached_credential_is_used_and_dropped_on_401` below. - #[tokio::test] - async fn uncached_credentials_are_not_retried_on_401() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let attempts = Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let seen = attempts.clone(); - let upstream_task = tokio::spawn(async move { - loop { - let (socket, _) = upstream.accept().await.unwrap(); - let seen = seen.clone(); - tokio::spawn(async move { - let _ = hyper::server::conn::http1::Builder::new() - .serve_connection( - TokioIo::new(socket), - service_fn(move |request: Request| { - let seen = seen.clone(); - async move { - let n = - seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - assert_eq!( - request.headers()[AUTHORIZATION], - "Bearer identity" - ); - let body = - request.into_body().collect().await.unwrap().to_bytes(); - assert_eq!(body, "{\"n\":1}"); - // One attempt per client request: the first - // is answered 401, the second 200. - let status = if n == 1 { - StatusCode::OK - } else { - StatusCode::UNAUTHORIZED - }; - Ok::<_, Infallible>( - Response::builder() - .status(status) - .body(Full::new(Bytes::from_static(b"upstream"))) - .unwrap(), - ) - } - }), - ) - .await; - }); - } - }); - let (dir, state) = request_source_state(address).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let request = || { - Request::builder() - .method(Method::POST) - .uri(format!("http://{proxy_address}/chat/completions")) - .header(PAIRING_HEADER, PAIRING) - .header(AUTHORIZATION, "Bearer tid=from-client") - .body(Full::new(Bytes::from_static(b"{\"n\":1}"))) - .unwrap() - }; - let response = client.request(request()).await.unwrap(); - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); - assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); - let response = client.request(request()).await.unwrap(); - assert_eq!(response.status(), StatusCode::OK); - assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 2); - proxy.abort(); - let _ = proxy.await; - drop(client); - upstream_task.abort(); - drop(dir); - }) - .await - .unwrap(); - } - - // A pre-seeded cache stands in for a controller: the proxy sends the cached - // credential; when the gateway accepts it the entry stays; when the gateway - // rejects it the entry is dropped and the request is retried, which here - // needs a controller that does not exist, so the client gets - // `agentdesktop_credential` and the gateway saw the body exactly once. - #[tokio::test] - async fn cached_credential_is_used_and_dropped_on_401() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - type Seen = Arc)>>>; - let seen: Seen = Arc::default(); - let record = seen.clone(); - let upstream_task = tokio::spawn(async move { - loop { - let (socket, _) = upstream.accept().await.unwrap(); - let record = record.clone(); - tokio::spawn(async move { - let _ = hyper::server::conn::http1::Builder::new() - .serve_connection( - TokioIo::new(socket), - service_fn(move |request: Request| { - let record = record.clone(); - async move { - let authorization = request - .headers() - .get(AUTHORIZATION) - .and_then(|v| v.to_str().ok()) - .unwrap_or("") - .to_owned(); - let body = request - .into_body() - .collect() - .await - .unwrap() - .to_bytes() - .to_vec(); - let status = if authorization == "Bearer good" { - StatusCode::OK - } else { - StatusCode::UNAUTHORIZED - }; - record.lock().unwrap().push((authorization, body)); - Ok::<_, Infallible>( - Response::builder() - .status(status) - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(), - ) - } - }), - ) - .await; - }); - } - }); - let (dir, state) = controller_jwt_state(address, false).await; - let cache = Arc::new(CredentialCache::default()); - let key = format!("copilot-cli\u{0}http://{address}/\u{0}agentgateway"); - let issued = |value: &str| agentdesktop_core::model::LlmGatewayCredential { - credential: value.to_owned(), - expires_at_unix_seconds: now_unix() + 600, - }; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve_with_cache( - listener, - state, - default_config(), - cache.clone(), - )); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let request = |body: &'static [u8]| { - Request::builder() - .method(Method::POST) - .uri(format!( - "http://{proxy_address}/copilot-cli/v1/chat/completions" - )) - .header(PAIRING_HEADER, PAIRING) - .body(Full::new(Bytes::from_static(body))) - .unwrap() - }; - // Accepted: served from the cache, entry kept. - cache.insert(&key, "device-test", &issued("good")); - let response = client.request(request(b"{\"n\":1}")).await.unwrap(); - assert_eq!(response.status(), StatusCode::OK); - assert_eq!(cache.get(&key, "device-test").as_deref(), Some("good")); - // Rejected: entry dropped, one retry that needs a controller. - cache.insert(&key, "device-test", &issued("stale")); - let (status, code) = - error_code(client.request(request(b"{\"n\":2}")).await.unwrap()).await; - assert_eq!(code, "agentdesktop_credential", "{status}"); - assert_eq!(cache.get(&key, "device-test"), None); - let seen = seen.lock().unwrap().clone(); - assert_eq!( - seen, - vec![ - ("Bearer good".to_owned(), b"{\"n\":1}".to_vec()), - ("Bearer stale".to_owned(), b"{\"n\":2}".to_vec()), - ], - "each body reaches the gateway exactly once per attempt" - ); - proxy.abort(); - let _ = proxy.await; - drop(client); - upstream_task.abort(); - drop(dir); - }) - .await - .unwrap(); - } - - // With a client token forwarded (prefix-less route, githubOAuth source - // request) a 401 is returned as is: no retry, and the cached gateway - // credential is kept. - #[tokio::test] - async fn client_token_401_is_not_retried_and_keeps_the_cached_credential() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let attempts = Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let seen = attempts.clone(); - let upstream_task = tokio::spawn(async move { - loop { - let (socket, _) = upstream.accept().await.unwrap(); - let seen = seen.clone(); - tokio::spawn(async move { - let _ = hyper::server::conn::http1::Builder::new() - .serve_connection( - TokioIo::new(socket), - service_fn(move |request: Request| { - let seen = seen.clone(); - async move { - seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - assert_eq!( - request.headers()[AUTHORIZATION], - "Bearer cached" - ); - assert_eq!(request.headers()["x-llm-token"], "tid=bad"); - Ok::<_, Infallible>( - Response::builder() - .status(StatusCode::UNAUTHORIZED) - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(), - ) - } - }), - ) - .await; - }); - } - }); - let (dir, state) = controller_jwt_state(address, true).await; - let cache = Arc::new(CredentialCache::default()); - let key = format!("vscode\u{0}http://{address}/\u{0}agentgateway"); - cache.insert( - &key, - "device-test", - &agentdesktop_core::model::LlmGatewayCredential { - credential: "cached".to_owned(), - expires_at_unix_seconds: now_unix() + 600, - }, - ); - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve_with_cache( - listener, - state, - default_config(), - cache.clone(), - )); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let request = Request::builder() - .method(Method::POST) - .uri(format!("http://{proxy_address}/chat/completions")) - .header(PAIRING_HEADER, PAIRING) - .header(AUTHORIZATION, "Bearer tid=bad") - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - let response = client.request(request).await.unwrap(); - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); - assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); - assert_eq!(cache.get(&key, "device-test").as_deref(), Some("cached")); - proxy.abort(); - let _ = proxy.await; - drop(client); - upstream_task.abort(); - drop(dir); - }) - .await - .unwrap(); - } - - // Controller-JWT state with an identity.json but no controller configured: - // the device id is readable for the cache key, a fetch fails at once - // ("requires a controller"), so no timeout is involved. - async fn controller_jwt_state( - address: std::net::SocketAddr, - client_token: bool, - ) -> (tempfile::TempDir, AppState) { - let dir = tempfile::tempdir().unwrap(); - std::fs::write( - dir.path().join("identity.json"), - serde_json::json!({ - "deviceId": "device-test", - "clientCertificatePem": "", - "clientCertificateExpiresAtUnixSeconds": 4_000_000_000u64, - "oauthTokenEndpoint": "https://controller.example/token", - "oauthClientId": "device", - }) - .to_string(), - ) - .unwrap(); - let github = if client_token { - " githubOAuth:\n source: request\n" - } else { - "" - }; - let config = parse_daemon(&format!( - "llmGateway:\n url: http://{address}\n{github} authentication:\n type: controllerJwt\n audience: agentgateway\n allowedClientIds: [copilot-cli, vscode]\n" - )).unwrap(); - let (_, discovery) = watch::channel(Arc::new(Discovery { - agents: vec![], - model_runtimes: vec![], - })); - let state = AppState { - controller_status: None, - config, - daemon_info: DaemonInfo { - version: "test".into(), - scope: DaemonScope::User, - config_path: String::new(), - state_directory: String::new(), - inventory_interval: Duration::from_secs(60), - controller: None, - llm_proxy: None, - usage_receiver: None, - }, - discovery, - enrollment: EnrollmentState::new(false), - state_dir: dir.path().to_owned(), - oidc_callback_listen: None, - telemetry: None, - logout: None, - usage: None, - programs: None, - }; - (dir, state) - } - - #[tokio::test] - async fn refuses_a_body_over_the_cap_with_a_json_error() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let (dir, state) = request_source_state(upstream.local_addr().unwrap()).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let request = Request::builder() - .method(Method::POST) - .uri(format!("http://{proxy_address}/chat/completions")) - .header(PAIRING_HEADER, PAIRING) - .header(AUTHORIZATION, "Bearer tid=from-client") - .body(Full::new(Bytes::from(vec![b'x'; MAX_REQUEST_BODY + 1]))) - .unwrap(); - assert_eq!( - error_code(client.request(request).await.unwrap()).await, - (StatusCode::PAYLOAD_TOO_LARGE, "body_too_large".to_owned()) - ); - proxy.abort(); - let _ = proxy.await; - drop(client); - drop(dir); - }) - .await - .unwrap(); - } - - #[tokio::test] - async fn concurrent_misses_share_one_fetch_per_key() { - let far = now_unix() + 3600; - let cache = Arc::new(CredentialCache::default()); - let fetches = Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let mut tasks = tokio::task::JoinSet::new(); - for key in ["a", "a", "a", "a", "a", "b", "b"] { - let cache = cache.clone(); - let fetches = fetches.clone(); - tasks.spawn(async move { - cache - .get_or_fetch(key, "device-1", || async move { - let n = fetches.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - tokio::time::sleep(Duration::from_millis(50)).await; - Ok::<_, ()>(agentdesktop_core::model::LlmGatewayCredential { - credential: format!("{key}{n}"), - expires_at_unix_seconds: far, - }) - }) - .await - .map(|(credential, _, _)| (key, credential)) - }); - } - let mut results = Vec::new(); - while let Some(result) = tasks.join_next().await { - results.push(result.unwrap().unwrap()); - } - // One fetch per key; every caller of a key got that fetch's credential. - assert_eq!(fetches.load(std::sync::atomic::Ordering::SeqCst), 2); - for key in ["a", "b"] { - let mut credentials: Vec<_> = results - .iter() - .filter(|(k, _)| *k == key) - .map(|(_, c)| c) - .collect(); - credentials.dedup(); - assert_eq!(credentials.len(), 1, "{key}: {results:?}"); - } - // Cached at fetch time, before any gateway answer. - assert!(cache.get("a", "device-1").is_some()); - } - - #[test] - fn credential_cache_deadline_is_the_credential_expiry_not_the_reuse_window() { - let cache = CredentialCache::default(); - let expires = now_unix() + 3600; - cache.insert( - "k", - "device-1", - &agentdesktop_core::model::LlmGatewayCredential { - credential: "jwt".to_owned(), - expires_at_unix_seconds: expires, - }, - ); - let (credential, deadline) = cache.lookup("k", "device-1").expect("cached"); - assert_eq!(credential, "jwt"); - assert_eq!(deadline.unix, expires); - let remaining = deadline.monotonic - std::time::Instant::now(); - assert!( - remaining > CREDENTIAL_CACHE_TTL && remaining <= Duration::from_secs(3600), - "tunnel deadline must follow the credential's expiry, got {remaining:?}" - ); - // The reuse window is still the short one. - let entry = cache.entries.lock().unwrap(); - assert!(entry["k"].valid_until <= std::time::Instant::now() + CREDENTIAL_CACHE_TTL); - } - - #[test] - fn credential_cache_honours_device_ttl_and_invalidation() { - let cache = CredentialCache::default(); - let far = now_unix() + 600; - let issued = |value: &str, expires: u64| agentdesktop_core::model::LlmGatewayCredential { - credential: value.to_owned(), - expires_at_unix_seconds: expires, - }; - cache.insert("k", "device-1", &issued("one", far)); - assert_eq!(cache.get("k", "device-1").as_deref(), Some("one")); - assert_eq!( - cache.get("k", "device-2"), - None, - "another device never sees it" - ); - assert_eq!(cache.get("other", "device-1"), None); - cache.invalidate_if("k", "someone-else"); - assert_eq!( - cache.get("k", "device-1").as_deref(), - Some("one"), - "a different credential leaves the entry alone" - ); - cache.invalidate_if("k", "one"); - assert_eq!(cache.get("k", "device-1"), None); - // Already within the expiry margin: not cached at all. - cache.insert("m", "device-1", &issued("m1", now_unix() + 10)); - assert_eq!(cache.get("m", "device-1"), None); - // Lifetime is capped by the TTL, not the token's own expiry. - cache.insert("t", "device-1", &issued("t1", far)); - let entries = cache.entries.lock().unwrap(); - let entry = &entries["t"]; - assert!(entry.valid_until <= std::time::Instant::now() + CREDENTIAL_CACHE_TTL); - assert!(entry.valid_until_unix <= now_unix() + CREDENTIAL_CACHE_TTL.as_secs()); - } - - async fn request_source_state(address: std::net::SocketAddr) -> (tempfile::TempDir, AppState) { - let dir = tempfile::tempdir().unwrap(); - let store = SecretStore::new(dir.path()).unwrap(); - let account = - URL_SAFE_NO_PAD.encode(Sha256::digest(b"https://issuer.example/\0proxy-test")); - store - .set( - "dev.agentdesktop.gateway-oidc", - &account, - &serde_json::json!({"accessToken": "identity", "refreshToken": null, - "expiresAtUnixSeconds": 4_000_000_000u64, - "tokenEndpoint": "https://issuer.example/token"}) - .to_string(), - ) - .unwrap(); - let config = parse_daemon(&format!( - "llmGateway:\n url: http://{address}\n githubOAuth:\n source: request\n authentication:\n type: oidc\n issuer: https://issuer.example/\n clientId: proxy-test\n" - )).unwrap(); - let (_, discovery) = watch::channel(Arc::new(Discovery { - agents: vec![], - model_runtimes: vec![], - })); - let state = AppState { - controller_status: None, - config, - daemon_info: DaemonInfo { - version: "test".into(), - scope: DaemonScope::User, - config_path: String::new(), - state_directory: String::new(), - inventory_interval: Duration::from_secs(60), - controller: None, - llm_proxy: None, - usage_receiver: None, - }, - discovery, - enrollment: EnrollmentState::new(false), - state_dir: dir.path().to_owned(), - oidc_callback_listen: None, - telemetry: None, - logout: None, - usage: None, - programs: None, - }; - (dir, state) - } - - // --- vscode-copilot-capi: path pairing, /_ping and the WebSocket tunnel. - - use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::TcpStream, - }; - - /// `request_source_state` (OIDC identity "identity", `githubOAuth: source: - /// request`) with `llmGateway.proxyUrl` pointed at `upstream`, matching - /// the CAPI route's Passthrough/ProxyUrl shape used by most tests below. - async fn capi_state(upstream: std::net::SocketAddr) -> (tempfile::TempDir, AppState) { - let (dir, mut state) = request_source_state(upstream).await; - state.config.llm_gateway.as_mut().unwrap().proxy_url = - Some(format!("http://{upstream}").parse().unwrap()); - (dir, state) - } - - /// Controller-JWT state (cacheable credential, unlike OIDC) with - /// `llmGateway.proxyUrl` pointed at `upstream`, for the credential-expiry - /// tunnel test, which needs a credential it can seed into the cache with a - /// controlled deadline. - async fn capi_controller_jwt_state( - upstream: std::net::SocketAddr, - ) -> (tempfile::TempDir, AppState) { - let dir = tempfile::tempdir().unwrap(); - std::fs::write( - dir.path().join("identity.json"), - serde_json::json!({ - "deviceId": "device-test", - "clientCertificatePem": "", - "clientCertificateExpiresAtUnixSeconds": 4_000_000_000u64, - "oauthTokenEndpoint": "https://controller.example/token", - "oauthClientId": "device", - }) - .to_string(), - ) - .unwrap(); - let config = parse_daemon(&format!( - "llmGateway:\n url: http://{upstream}\n proxyUrl: http://{upstream}\n authentication:\n type: controllerJwt\n audience: agentgateway\n allowedClientIds: [vscode-copilot]\n" - )) - .unwrap(); - let (_, discovery) = watch::channel(Arc::new(Discovery { - agents: vec![], - model_runtimes: vec![], - })); - let state = AppState { - controller_status: None, - config, - daemon_info: DaemonInfo { - version: "test".into(), - scope: DaemonScope::User, - config_path: String::new(), - state_directory: String::new(), - inventory_interval: Duration::from_secs(60), - controller: None, - llm_proxy: None, - usage_receiver: None, - }, - discovery, - enrollment: EnrollmentState::new(false), - state_dir: dir.path().to_owned(), - oidc_callback_listen: None, - telemetry: None, - logout: None, - usage: None, - programs: None, - }; - (dir, state) - } - - /// Whether no connection reached `listener` within a short window: proof - /// that a refused request was never forwarded upstream. - async fn no_connection_arrives(listener: &TcpListener) -> bool { - tokio::time::timeout(Duration::from_millis(200), listener.accept()) - .await - .is_err() - } - - /// Reads a raw HTTP/1.1 request head (request line + headers, byte at a - /// time up to the blank line) from a hand-rolled test upstream: good - /// enough for a WebSocket upgrade request, which has no body. Header - /// names are folded to lowercase for case-insensitive lookup. - async fn read_http_head( - stream: &mut TcpStream, - ) -> (String, String, std::collections::HashMap) { - let mut buffer = Vec::new(); - let mut byte = [0u8; 1]; - while !buffer.ends_with(b"\r\n\r\n") { - stream.read_exact(&mut byte).await.unwrap(); - buffer.push(byte[0]); - } - let text = String::from_utf8(buffer).unwrap(); - let mut lines = text.split("\r\n"); - let mut parts = lines.next().unwrap().split_whitespace(); - let method = parts.next().unwrap().to_owned(); - let path = parts.next().unwrap().to_owned(); - let mut headers = std::collections::HashMap::new(); - for line in lines { - if let Some((name, value)) = line.split_once(':') { - headers.insert(name.trim().to_ascii_lowercase(), value.trim().to_owned()); - } - } - (method, path, headers) - } - - /// RFC 6455's fixed GUID: concatenated with the client's - /// `Sec-WebSocket-Key` and SHA-1'd (then base64'd) to produce - /// `Sec-WebSocket-Accept`. - const WS_GUID: &str = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"; - - fn ws_accept(key: &str) -> String { - use base64::engine::general_purpose::STANDARD; - use sha1::{Digest, Sha1}; - STANDARD.encode(Sha1::digest(format!("{key}{WS_GUID}").as_bytes())) - } - - /// Encodes a single, unfragmented text frame (opcode `0x1`). Client - /// frames must be masked (`mask: true`, a fixed key: the test payloads - /// are not attacker-controlled), server frames must not be. - fn ws_frame(payload: &[u8], mask: bool) -> Vec { - let mut frame = vec![0x81u8]; - let mask_bit = if mask { 0x80 } else { 0x00 }; - let len = payload.len(); - if len < 126 { - frame.push(mask_bit | len as u8); - } else { - frame.push(mask_bit | 126); - frame.extend_from_slice(&(len as u16).to_be_bytes()); - } - if mask { - let key = [0x11u8, 0x22, 0x33, 0x44]; - frame.extend_from_slice(&key); - frame.extend(payload.iter().enumerate().map(|(i, b)| b ^ key[i % 4])); - } else { - frame.extend_from_slice(payload); - } - frame - } - - /// Reads exactly one, unfragmented, non-huge frame and returns its opcode - /// and unmasked payload. - async fn ws_read_frame(stream: &mut TcpStream) -> (u8, Vec) { - let mut header = [0u8; 2]; - stream.read_exact(&mut header).await.unwrap(); - let opcode = header[0] & 0x0f; - let masked = header[1] & 0x80 != 0; - let mut len = (header[1] & 0x7f) as usize; - if len == 126 { - let mut extended = [0u8; 2]; - stream.read_exact(&mut extended).await.unwrap(); - len = u16::from_be_bytes(extended) as usize; - } - let mask_key = if masked { - let mut key = [0u8; 4]; - stream.read_exact(&mut key).await.unwrap(); - Some(key) - } else { - None - }; - let mut payload = vec![0u8; len]; - stream.read_exact(&mut payload).await.unwrap(); - if let Some(key) = mask_key { - for (index, byte) in payload.iter_mut().enumerate() { - *byte ^= key[index % 4]; - } - } - (opcode, payload) - } - - /// A fixed, valid `Sec-WebSocket-Key` (the RFC 6455 example key): its - /// `Sec-WebSocket-Accept` is computed by `ws_accept`, never hand-checked - /// against a literal. - const SEC_WEBSOCKET_KEY: &str = "dGhlIHNhbXBsZSBub25jZQ=="; - - fn upgrade_request(uri: String) -> Request> { - Request::builder() - .method(Method::GET) - .uri(uri) - .header(CONNECTION, "Upgrade") - .header("upgrade", "websocket") - .header("sec-websocket-version", "13") - .header("sec-websocket-key", SEC_WEBSOCKET_KEY) - .body(Full::new(Bytes::new())) - .unwrap() - } - - #[tokio::test] - async fn capi_route_accepts_and_strips_path_pairing_never_forwarding_it() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let upstream_task = tokio::spawn(async move { - let (socket, _) = upstream.accept().await.unwrap(); - hyper::server::conn::http1::Builder::new() - .serve_connection( - TokioIo::new(socket), - service_fn(move |request: Request| async move { - // The pairing segment must never reach the upstream path. - assert_eq!(request.uri(), "/chat/completions"); - assert!(!request.headers().contains_key(PAIRING_HEADER)); - Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static(b"ok")))) - }), - ) - .await - .unwrap(); - }); - let (dir, state) = capi_state(address).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let request = Request::builder() - .method(Method::POST) - .uri(format!( - "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/chat/completions" - )) - .header(AUTHORIZATION, "Bearer tid=from-client") - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - let response = client.request(request).await.unwrap(); - assert_eq!( - response.status(), - StatusCode::OK, - "path pairing must be accepted on {CAPI_ROUTE}" - ); - proxy.abort(); - let _ = proxy.await; - drop(client); - upstream_task.abort(); - drop(dir); - }) - .await - .unwrap(); - } - - #[tokio::test] - async fn capi_route_rejects_missing_empty_or_wrong_pairing_segment_even_with_a_valid_header() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - // Dropped rather than accepted from: today none of these four - // paths are rejected before forwarding (path-pairing extraction - // does not exist yet), so leaving the listener up would have - // every case hang until the gateway timeout instead of failing - // fast with a clean (if wrong) status. - drop(upstream); - let (dir, state) = capi_state(address).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - for path in [ - CAPI_ROUTE.to_owned(), // no segment at all - format!("{CAPI_ROUTE}/"), // no segment at all - format!("{CAPI_ROUTE}//chat/completions"), // empty segment - format!("{CAPI_ROUTE}/wrong-pairing/chat/completions"), // wrong segment - ] { - let request = Request::builder() - .method(Method::POST) - .uri(format!("http://{proxy_address}{path}")) - // A valid header must not substitute for the path segment - // on this route: the header is ignored here. - .header(PAIRING_HEADER, PAIRING) - .header(AUTHORIZATION, "Bearer tid=from-client") - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - assert_eq!( - error_code(client.request(request).await.unwrap()).await, - (StatusCode::FORBIDDEN, "pairing_invalid".to_owned()), - "{path}" - ); - } - proxy.abort(); - let _ = proxy.await; - drop(client); - drop(dir); - }) - .await - .unwrap(); - } - - #[tokio::test] - async fn capi_route_rejects_a_double_dot_segment_after_the_pairing() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let (dir, state) = capi_state(address).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let request = Request::builder() - .method(Method::POST) - .uri(format!( - "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/../admin" - )) - .header(AUTHORIZATION, "Bearer tid=from-client") - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - assert_eq!( - error_code(client.request(request).await.unwrap()).await, - (StatusCode::BAD_REQUEST, "path_invalid".to_owned()) - ); - proxy.abort(); - let _ = proxy.await; - drop(client); - drop(dir); - }) - .await - .unwrap(); - } - - #[tokio::test] - async fn capi_route_moves_the_client_token_to_x_llm_token() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let upstream_task = tokio::spawn(async move { - let (socket, _) = upstream.accept().await.unwrap(); - hyper::server::conn::http1::Builder::new() - .serve_connection( - TokioIo::new(socket), - service_fn(move |request: Request| async move { - assert_eq!(request.uri(), "/chat/completions"); - assert_eq!(request.headers()["x-llm-token"], "tid=from-client"); - assert_eq!(request.headers()[AUTHORIZATION], "Bearer identity"); - Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static(b"ok")))) - }), - ) - .await - .unwrap(); - }); - let (dir, state) = capi_state(address).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let request = Request::builder() - .method(Method::POST) - .uri(format!( - "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/chat/completions" - )) - .header(AUTHORIZATION, "Bearer tid=from-client") - .body(Full::new(Bytes::from_static(b"{}"))) - .unwrap(); - let response = client.request(request).await.unwrap(); - assert_eq!(response.status(), StatusCode::OK); - proxy.abort(); - let _ = proxy.await; - drop(client); - upstream_task.abort(); - drop(dir); - }) - .await - .unwrap(); - } - - #[tokio::test] - async fn capi_route_ping_without_credential_is_forwarded_without_x_llm_token() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let upstream_task = tokio::spawn(async move { - let (socket, _) = upstream.accept().await.unwrap(); - hyper::server::conn::http1::Builder::new() - .serve_connection( - TokioIo::new(socket), - service_fn(move |request: Request| async move { - assert_eq!(request.uri(), "/_ping"); - assert!(!request.headers().contains_key("x-llm-token")); - Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static(b"ok")))) - }), - ) - .await - .unwrap(); - }); - let (dir, state) = capi_state(address).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let request = Request::builder() - .method(Method::GET) - .uri(format!( - "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/_ping" - )) - .body(Full::new(Bytes::new())) - .unwrap(); - let response = client.request(request).await.unwrap(); - assert_eq!(response.status(), StatusCode::OK); - proxy.abort(); - let _ = proxy.await; - drop(client); - upstream_task.abort(); - drop(dir); - }) - .await - .unwrap(); - } - - #[tokio::test] - async fn capi_route_ping_post_without_credential_is_401() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let (dir, state) = capi_state(address).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let request = Request::builder() - .method(Method::POST) - .uri(format!( - "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/_ping" - )) - .body(Full::new(Bytes::new())) - .unwrap(); - assert_eq!( - error_code(client.request(request).await.unwrap()).await, - ( - StatusCode::UNAUTHORIZED, - "client_credential_missing".to_owned() - ) - ); - proxy.abort(); - let _ = proxy.await; - drop(client); - drop(dir); - }) - .await - .unwrap(); - } - - #[tokio::test] - async fn capi_route_ping_without_pairing_is_403() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let (dir, state) = capi_state(address).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let request = Request::builder() - .method(Method::GET) - .uri(format!("http://{proxy_address}{CAPI_ROUTE}/wrong/_ping")) - .body(Full::new(Bytes::new())) - .unwrap(); - assert_eq!( - error_code(client.request(request).await.unwrap()).await, - (StatusCode::FORBIDDEN, "pairing_invalid".to_owned()) - ); - proxy.abort(); - let _ = proxy.await; - drop(client); - drop(dir); - }) - .await - .unwrap(); - } - - #[tokio::test] - async fn capi_route_upgrade_without_pairing_is_403_and_not_forwarded() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let (dir, state) = capi_state(address).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let request = upgrade_request(format!( - "http://{proxy_address}{CAPI_ROUTE}/wrong/responses" - )); - assert_eq!( - error_code(client.request(request).await.unwrap()).await, - (StatusCode::FORBIDDEN, "pairing_invalid".to_owned()) - ); - assert!( - no_connection_arrives(&upstream).await, - "an upgrade refused for its pairing must not be forwarded upstream" - ); - proxy.abort(); - let _ = proxy.await; - drop(client); - drop(dir); - }) - .await - .unwrap(); - } - - #[tokio::test] - async fn upgrade_on_another_route_is_400_but_other_upgrade_tokens_still_forward() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let upstream_task = tokio::spawn(async move { - loop { - let (socket, _) = match upstream.accept().await { - Ok(accepted) => accepted, - Err(_) => break, - }; - tokio::spawn(async move { - let _ = hyper::server::conn::http1::Builder::new() - .serve_connection( - TokioIo::new(socket), - service_fn(|_request: Request| async move { - Ok::<_, Infallible>(Response::new(Full::new( - Bytes::from_static(b"ok"), - ))) - }), - ) - .await; - }); - } - }); - let (dir, state) = request_source_state(address).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - // Only /vscode-copilot-capi tunnels an upgrade; on any other - // route Upgrade: websocket is refused outright. - let websocket_upgrade = Request::builder() - .method(Method::GET) - .uri(format!("http://{proxy_address}/copilot-cli/v1/models")) - .header(PAIRING_HEADER, PAIRING) - .header(AUTHORIZATION, "Bearer tid=from-client") - .header(CONNECTION, "upgrade") - .header("upgrade", "websocket") - .body(Full::new(Bytes::new())) - .unwrap(); - // Checked as a plain status first (not via the shared - // `error_code` helper, which assumes a JSON error body): today - // this request is not refused at all and is forwarded as an - // ordinary request, so the response has no content-type header - // for `error_code` to read. - let response = client.request(websocket_upgrade).await.unwrap(); - assert_eq!( - response.status(), - StatusCode::BAD_REQUEST, - "an Upgrade: websocket request on a non-CAPI route must be refused, not forwarded" - ); - let body = response.into_body().collect().await.unwrap().to_bytes(); - let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); - assert_eq!(json["error"]["code"], "upgrade_not_supported"); - // Only the websocket token is refused; another Upgrade token on - // the same route is stripped (as a hop-by-hop header) and - // forwarded as an ordinary request, as before this feature. - let other_upgrade = Request::builder() - .method(Method::GET) - .uri(format!("http://{proxy_address}/copilot-cli/v1/models")) - .header(PAIRING_HEADER, PAIRING) - .header(AUTHORIZATION, "Bearer tid=from-client") - .header(CONNECTION, "upgrade") - .header("upgrade", "h2c") - .body(Full::new(Bytes::new())) - .unwrap(); - let response = client.request(other_upgrade).await.unwrap(); - assert_eq!( - response.status(), - StatusCode::OK, - "a non-websocket Upgrade token must still be forwarded" - ); - proxy.abort(); - let _ = proxy.await; - drop(client); - upstream_task.abort(); - drop(dir); - }) - .await - .unwrap(); - } - - #[tokio::test] - async fn capi_route_non_101_upstream_response_passes_through() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let upstream_task = tokio::spawn(async move { - let (socket, _) = upstream.accept().await.unwrap(); - hyper::server::conn::http1::Builder::new() - .serve_connection( - TokioIo::new(socket), - service_fn(|_request: Request| async move { - Ok::<_, Infallible>( - Response::builder() - .status(StatusCode::BAD_GATEWAY) - .body(Full::new(Bytes::from_static(b"nope"))) - .unwrap(), - ) - }), - ) - .await - .unwrap(); - }); - let (dir, state) = capi_state(address).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let mut request = upgrade_request(format!( - "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/responses" - )); - request - .headers_mut() - .insert(AUTHORIZATION, "Bearer tid=from-client".parse().unwrap()); - let response = client.request(request).await.unwrap(); - assert_eq!(response.status(), StatusCode::BAD_GATEWAY); - let body = response.into_body().collect().await.unwrap().to_bytes(); - assert_eq!(body, Bytes::from_static(b"nope")); - proxy.abort(); - let _ = proxy.await; - drop(client); - upstream_task.abort(); - drop(dir); - }) - .await - .unwrap(); - } - - #[tokio::test] - async fn capi_route_websocket_tunnel_relays_the_101_and_one_frame_each_way_with_gateway_credentials() - { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let upstream_task = tokio::spawn(async move { - let (mut socket, _) = upstream.accept().await.unwrap(); - let (method, path, headers) = read_http_head(&mut socket).await; - assert_eq!(method, "GET"); - assert_eq!(path, "/responses"); - assert_eq!( - headers.get("authorization").map(String::as_str), - Some("Bearer identity"), - "the gateway identity must reach the upstream on the handshake" - ); - assert_eq!( - headers.get("x-llm-token").map(String::as_str), - Some("tid=from-client"), - "the client's own token must reach the upstream on the handshake" - ); - let accept = - ws_accept(headers.get("sec-websocket-key").map(String::as_str).unwrap_or("")); - socket - .write_all( - format!( - "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {accept}\r\n\r\n" - ) - .as_bytes(), - ) - .await - .unwrap(); - let (opcode, payload) = ws_read_frame(&mut socket).await; - assert_eq!(opcode, 0x1); - assert_eq!(payload, b"hello from vs code"); - socket - .write_all(&ws_frame(b"hello from github", false)) - .await - .unwrap(); - }); - let (dir, state) = capi_state(address).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let mut request = upgrade_request(format!( - "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/responses" - )); - request - .headers_mut() - .insert(AUTHORIZATION, "Bearer tid=from-client".parse().unwrap()); - let response = client.request(request).await.unwrap(); - assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS); - assert_eq!(response.headers()["upgrade"], "websocket"); - assert!( - response.headers()["connection"] - .to_str() - .unwrap() - .eq_ignore_ascii_case("upgrade") - ); - assert_eq!( - response.headers()["sec-websocket-accept"], - ws_accept(SEC_WEBSOCKET_KEY) - ); - let upgraded = hyper::upgrade::on(response).await.unwrap(); - let mut io = TokioIo::new(upgraded); - io.write_all(&ws_frame(b"hello from vs code", true)) - .await - .unwrap(); - let mut header = [0u8; 2]; - io.read_exact(&mut header).await.unwrap(); - assert_eq!(header[0] & 0x0f, 0x1, "expected a text frame back"); - let len = (header[1] & 0x7f) as usize; - assert!(header[1] & 0x80 == 0, "a server frame must not be masked"); - let mut payload = vec![0u8; len]; - io.read_exact(&mut payload).await.unwrap(); - assert_eq!(payload, b"hello from github"); - proxy.abort(); - let _ = proxy.await; - drop(client); - upstream_task.abort(); - drop(dir); - }) - .await - .unwrap(); - } - - #[tokio::test] - async fn dropping_the_proxy_closes_an_open_capi_tunnel() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let upstream_task = tokio::spawn(async move { - let (mut socket, _) = upstream.accept().await.unwrap(); - let (_, _, headers) = read_http_head(&mut socket).await; - let accept = - ws_accept(headers.get("sec-websocket-key").map(String::as_str).unwrap_or("")); - socket - .write_all( - format!( - "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {accept}\r\n\r\n" - ) - .as_bytes(), - ) - .await - .unwrap(); - // Held open until the client side sees it close. - let mut buffer = [0u8; 1]; - let _ = socket.read(&mut buffer).await; - }); - let (dir, state) = capi_state(address).await; - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve(listener, state, default_config())); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let mut request = upgrade_request(format!( - "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/responses" - )); - request - .headers_mut() - .insert(AUTHORIZATION, "Bearer tid=from-client".parse().unwrap()); - let response = client.request(request).await.unwrap(); - assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS); - let upgraded = hyper::upgrade::on(response).await.unwrap(); - let mut io = TokioIo::new(upgraded); - proxy.abort(); - let _ = proxy.await; - // The tunnel task is owned by the connection's JoinSet: dropping - // the server must close it, and the client sees EOF. - let mut buffer = [0u8; 1]; - let read = tokio::time::timeout(Duration::from_secs(5), io.read(&mut buffer)) - .await - .expect("dropping the proxy must close the open tunnel promptly") - .unwrap(); - assert_eq!(read, 0, "expected EOF once the tunnel closes"); - upstream_task.abort(); - drop(client); - drop(dir); - }) - .await - .unwrap(); - } - - #[tokio::test] - async fn capi_tunnel_closes_when_the_cached_credential_expires() { - tokio::time::timeout(Duration::from_secs(15), async { - let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = upstream.local_addr().unwrap(); - let upstream_task = tokio::spawn(async move { - let (mut socket, _) = upstream.accept().await.unwrap(); - let (_, _, headers) = read_http_head(&mut socket).await; - let accept = - ws_accept(headers.get("sec-websocket-key").map(String::as_str).unwrap_or("")); - socket - .write_all( - format!( - "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {accept}\r\n\r\n" - ) - .as_bytes(), - ) - .await - .unwrap(); - let mut buffer = [0u8; 1]; - let _ = socket.read(&mut buffer).await; - }); - let (dir, state) = capi_controller_jwt_state(address).await; - let cache = Arc::new(CredentialCache::default()); - let key = format!("vscode-copilot\u{0}http://{address}/\u{0}agentgateway"); - // Inserted directly (bypassing `insert`'s expiry-margin floor, - // which would otherwise refuse to cache anything this close to - // expiry). The cache reuses the entry for a minute, so the - // request always gets it however slowly it arrives; the - // credential itself expires in 300 ms, and that is the deadline - // the tunnel must close at. - cache.entries.lock().unwrap().insert( - key, - CachedCredential { - credential: "short-lived".to_owned(), - device_id: "device-test".to_owned(), - valid_until: std::time::Instant::now() + Duration::from_secs(60), - valid_until_unix: now_unix() + 60, - expires: std::time::Instant::now() + Duration::from_millis(300), - expires_unix: now_unix() + 60, - }, - ); - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let proxy_address = listener.local_addr().unwrap(); - let proxy = tokio::spawn(serve_with_cache( - listener, - state, - default_config(), - cache, - )); - let client: Client> = - Client::builder(TokioExecutor::new()).build_http(); - let mut request = upgrade_request(format!( - "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/responses" - )); - request - .headers_mut() - .insert(AUTHORIZATION, "Bearer tid=from-client".parse().unwrap()); - let response = client.request(request).await.unwrap(); - assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS); - let upgraded = hyper::upgrade::on(response).await.unwrap(); - let mut io = TokioIo::new(upgraded); - // Once the cached credential's deadline passes, the tunnel closes - // on its own, with the proxy still running. - let mut buffer = [0u8; 1]; - let read = tokio::time::timeout(Duration::from_secs(5), io.read(&mut buffer)) - .await - .expect("the tunnel must close once its credential expires") - .unwrap(); - assert_eq!(read, 0, "expected EOF once the tunnel closes"); - proxy.abort(); - let _ = proxy.await; - upstream_task.abort(); - drop(client); - drop(dir); - }) - .await - .unwrap(); - } -} +mod tests; diff --git a/crates/agent/src/llm_proxy/cache.rs b/crates/agent/src/llm_proxy/cache.rs new file mode 100644 index 00000000..ef1afee9 --- /dev/null +++ b/crates/agent/src/llm_proxy/cache.rs @@ -0,0 +1,306 @@ +use std::{sync::Arc, time::Duration}; + +/// A cached controller credential is not used within this margin of its expiry. +const CREDENTIAL_EXPIRY_MARGIN: Duration = Duration::from_secs(30); +/// Longest a controller credential is served from the cache, counted from its +/// fetch. This bounds how long a revoked device keeps using the proxy: +/// revocation is enforced at issuance, the gateway validates tokens statelessly. +const CREDENTIAL_CACHE_TTL: Duration = Duration::from_secs(60); + +/// Listener-local cache of controller-issued gateway credentials, one per client +/// id. A controller fetch is a fresh mTLS connection and round trip; agent-mode +/// sessions issue bursts of requests, so a short cache keeps them off the +/// controller. Fetches are single-flight per key: concurrent misses wait for +/// the one fetch in progress and then read its result from the cache. The +/// entries lock is never held across a fetch, so a slow controller delays only +/// the requests waiting for that key, not cache hits on other routes. +/// OIDC credentials never go through here; the secret store is their cache. +#[derive(Default)] +pub(crate) struct CredentialCache { + entries: std::sync::Mutex>, + flights: std::sync::Mutex>>>, +} + +struct CachedCredential { + credential: String, + /// The enrolled device the credential was issued to. A logout or a + /// re-enrollment changes it, and the entry is then ignored. + device_id: String, + /// Monotonic deadline (wall-clock expiry at fetch time, minus the margin, + /// capped by CREDENTIAL_CACHE_TTL) and the wall-clock deadline itself. Both + /// must be in the future: the monotonic clock does not advance across a + /// suspend, the wall clock can jump; neither alone is trusted. + valid_until: std::time::Instant, + valid_until_unix: u64, + /// The credential's own expiry (the controller JWT `exp`), on both clocks: + /// a tunnel opened with the credential runs until then, not until the + /// cache stops reusing it. + expires: std::time::Instant, + expires_unix: u64, +} + +/// A point in time on both clocks (see `CachedCredential`): a tunnel closes +/// when either says the credential has expired. +#[derive(Clone, Copy, Debug)] +pub(super) struct Deadline { + pub(super) monotonic: std::time::Instant, + pub(super) unix: u64, +} + +impl Deadline { + pub(super) fn from_unix(expires_unix: u64) -> Self { + Self { + monotonic: std::time::Instant::now() + + Duration::from_secs(expires_unix.saturating_sub(now_unix())), + unix: expires_unix, + } + } +} + +pub(super) fn now_unix() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) +} + +impl CredentialCache { + /// The lock a fetch for the key holds; one per key, kept for the + /// listener's lifetime (keys are client id, gateway and audience). + fn flight(&self, key: &str) -> Arc> { + self.flights + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .entry(key.to_owned()) + .or_default() + .clone() + } + + /// The cached credential for the key, or one from `fetch`, which is cached + /// at once. Single-flight: a caller that finds a fetch in progress waits + /// for it and reads its result. Returns the credential, its own expiry, and + /// whether it came from the cache. + pub(super) async fn get_or_fetch( + &self, + key: &str, + device_id: &str, + fetch: F, + ) -> Result<(String, Deadline, bool), E> + where + F: FnOnce() -> Fut, + Fut: + std::future::Future>, + { + let flight = self.flight(key); + let _flight = flight.lock().await; + if let Some((credential, deadline)) = self.lookup(key, device_id) { + return Ok((credential, deadline, true)); + } + let fetched = fetch().await?; + self.insert(key, device_id, &fetched); + let deadline = Deadline::from_unix(fetched.expires_at_unix_seconds); + Ok((fetched.credential, deadline, false)) + } + + /// A cached credential for the key and device, if one is still valid. + #[cfg(all(test, target_os = "linux"))] + pub(super) fn get(&self, key: &str, device_id: &str) -> Option { + self.lookup(key, device_id) + .map(|(credential, _)| credential) + } + + /// A cached credential together with its own expiry, read in one go so a + /// tunnel opened with it always gets the deadline of the credential it + /// was actually opened with. + pub(super) fn lookup(&self, key: &str, device_id: &str) -> Option<(String, Deadline)> { + let entries = self + .entries + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + entries + .get(key) + .filter(|entry| { + entry.device_id == device_id + && std::time::Instant::now() < entry.valid_until + && now_unix() < entry.valid_until_unix + }) + .map(|entry| { + ( + entry.credential.clone(), + Deadline { + monotonic: entry.expires, + unix: entry.expires_unix, + }, + ) + }) + } + + /// Remember a freshly fetched credential until the earlier of its expiry + /// margin and the cache TTL. + pub(super) fn insert( + &self, + key: &str, + device_id: &str, + fetched: &agentdesktop_core::model::LlmGatewayCredential, + ) { + let now = now_unix(); + let remaining = Duration::from_secs(fetched.expires_at_unix_seconds.saturating_sub(now)); + let lifetime = remaining + .saturating_sub(CREDENTIAL_EXPIRY_MARGIN) + .min(CREDENTIAL_CACHE_TTL); + let mut entries = self + .entries + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if lifetime > Duration::ZERO { + entries.insert( + key.to_owned(), + CachedCredential { + credential: fetched.credential.clone(), + device_id: device_id.to_owned(), + valid_until: std::time::Instant::now() + lifetime, + valid_until_unix: now + lifetime.as_secs(), + expires: std::time::Instant::now() + remaining, + expires_unix: fetched.expires_at_unix_seconds, + }, + ); + } else { + tracing::debug!( + key, + "controller credential already within its expiry margin; not cached" + ); + entries.remove(key); + } + } + + /// Drop the entry for the key if it still holds the given credential. A + /// concurrent request that already replaced it with a fresh one is left alone. + pub(super) fn invalidate_if(&self, key: &str, credential: &str) { + let mut entries = self + .entries + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if entries + .get(key) + .is_some_and(|entry| entry.credential == credential) + { + entries.remove(key); + } + } +} + +#[cfg(all(test, target_os = "linux"))] +impl CredentialCache { + pub(super) fn expire_at_for_test(&self, key: &str, deadline: Deadline) { + let mut entries = self.entries.lock().unwrap(); + let entry = entries.get_mut(key).expect("cached test credential"); + entry.expires = deadline.monotonic; + entry.expires_unix = deadline.unix; + } +} + +#[cfg(all(test, target_os = "linux"))] +mod tests { + use super::*; + #[tokio::test] + async fn concurrent_misses_share_one_fetch_per_key() { + let far = now_unix() + 3600; + let cache = Arc::new(CredentialCache::default()); + let fetches = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let mut tasks = tokio::task::JoinSet::new(); + for key in ["a", "a", "a", "a", "a", "b", "b"] { + let cache = cache.clone(); + let fetches = fetches.clone(); + tasks.spawn(async move { + cache + .get_or_fetch(key, "device-1", || async move { + let n = fetches.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(50)).await; + Ok::<_, ()>(agentdesktop_core::model::LlmGatewayCredential { + credential: format!("{key}{n}"), + expires_at_unix_seconds: far, + }) + }) + .await + .map(|(credential, _, _)| (key, credential)) + }); + } + let mut results = Vec::new(); + while let Some(result) = tasks.join_next().await { + results.push(result.unwrap().unwrap()); + } + // One fetch per key; every caller of a key got that fetch's credential. + assert_eq!(fetches.load(std::sync::atomic::Ordering::SeqCst), 2); + for key in ["a", "b"] { + let mut credentials: Vec<_> = results + .iter() + .filter(|(k, _)| *k == key) + .map(|(_, c)| c) + .collect(); + credentials.dedup(); + assert_eq!(credentials.len(), 1, "{key}: {results:?}"); + } + // Cached at fetch time, before any gateway answer. + assert!(cache.get("a", "device-1").is_some()); + } + + #[test] + fn credential_cache_deadline_is_the_credential_expiry_not_the_reuse_window() { + let cache = CredentialCache::default(); + let expires = now_unix() + 3600; + cache.insert( + "k", + "device-1", + &agentdesktop_core::model::LlmGatewayCredential { + credential: "jwt".to_owned(), + expires_at_unix_seconds: expires, + }, + ); + let (credential, deadline) = cache.lookup("k", "device-1").expect("cached"); + assert_eq!(credential, "jwt"); + assert_eq!(deadline.unix, expires); + let remaining = deadline.monotonic - std::time::Instant::now(); + assert!( + remaining > CREDENTIAL_CACHE_TTL && remaining <= Duration::from_secs(3600), + "tunnel deadline must follow the credential's expiry, got {remaining:?}" + ); + // The reuse window is still the short one. + let entry = cache.entries.lock().unwrap(); + assert!(entry["k"].valid_until <= std::time::Instant::now() + CREDENTIAL_CACHE_TTL); + } + + #[test] + fn credential_cache_honours_device_ttl_and_invalidation() { + let cache = CredentialCache::default(); + let far = now_unix() + 600; + let issued = |value: &str, expires: u64| agentdesktop_core::model::LlmGatewayCredential { + credential: value.to_owned(), + expires_at_unix_seconds: expires, + }; + cache.insert("k", "device-1", &issued("one", far)); + assert_eq!(cache.get("k", "device-1").as_deref(), Some("one")); + assert_eq!( + cache.get("k", "device-2"), + None, + "another device never sees it" + ); + assert_eq!(cache.get("other", "device-1"), None); + cache.invalidate_if("k", "someone-else"); + assert_eq!( + cache.get("k", "device-1").as_deref(), + Some("one"), + "a different credential leaves the entry alone" + ); + cache.invalidate_if("k", "one"); + assert_eq!(cache.get("k", "device-1"), None); + // Already within the expiry margin: not cached at all. + cache.insert("m", "device-1", &issued("m1", now_unix() + 10)); + assert_eq!(cache.get("m", "device-1"), None); + // Lifetime is capped by the TTL, not the token's own expiry. + cache.insert("t", "device-1", &issued("t1", far)); + let entries = cache.entries.lock().unwrap(); + let entry = &entries["t"]; + assert!(entry.valid_until <= std::time::Instant::now() + CREDENTIAL_CACHE_TTL); + assert!(entry.valid_until_unix <= now_unix() + CREDENTIAL_CACHE_TTL.as_secs()); + } +} diff --git a/crates/agent/src/llm_proxy/routes.rs b/crates/agent/src/llm_proxy/routes.rs new file mode 100644 index 00000000..aeea3acd --- /dev/null +++ b/crates/agent/src/llm_proxy/routes.rs @@ -0,0 +1,251 @@ +use agentdesktop_core::config::LlmGatewayConfig; +use hyper::{ + HeaderMap, + header::{AUTHORIZATION, HOST}, +}; + +/// Header that carries the per-device pairing value on every route. +pub(crate) const PAIRING_HEADER: &str = "x-agentdesktop-pairing"; + +/// How a route obtains the credential that goes upstream in `x-llm-token`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum RouteCredential { + /// Only the gateway identity is sent; any client credential is dropped. + Gateway, + /// The client's own bearer token is moved to `x-llm-token` and the gateway + /// identity takes `Authorization` (the VS Code Copilot pass-through shape). + Passthrough, +} + +/// Which configured gateway base URL a route forwards to. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Upstream { + /// `llmGateway.url`, shared with the credential-helper programs. + Url, + /// `llmGateway.proxyUrl` when set, else `llmGateway.url`. + ProxyUrl, +} + +/// A fixed path prefix the proxy serves for one managed program. +/// +/// The prefix is what a reconciler writes into that program's client file, so +/// the client id is decided by the file the daemon owns, never by request +/// headers. The prefix is stripped before forwarding. +#[derive(Clone, Copy, Debug)] +pub(crate) struct ProxyRoute { + pub prefix: &'static str, + pub client_id: &'static str, + pub credential: RouteCredential, + pub upstream: Upstream, + /// The pairing is the first path segment after the prefix (the header is + /// ignored): for a client that cannot add a header to its requests. + pub pairing_in_path: bool, + /// WebSocket upgrades are tunnelled instead of refused. + pub tunnels_upgrades: bool, + /// A credential-less `GET /_ping` is forwarded without `x-llm-token`. + pub credential_less_ping: bool, +} + +/// VS Code Copilot Chat's CAPI pass-through route (the target of +/// `github.copilot.advanced.debug.overrideCapiUrl`): the pairing travels as +/// the first path segment after the prefix instead of in `PAIRING_HEADER`, +/// because VS Code sends these requests itself and has no setting that adds a +/// header to them. Also the only route a WebSocket upgrade is tunnelled on +/// (the `GET /responses` Auto/agent conversation). +pub(crate) const CAPI_ROUTE: &str = "/vscode-copilot-capi"; +pub(crate) const CODEX_ROUTE: &str = "/codex"; + +/// Routes for the managed programs. A prefix matches only at a `/` boundary, so +/// `/vscode-copilot-passthrough/...` never matches `/vscode-copilot`. +pub(crate) const ROUTES: &[ProxyRoute] = &[ + ProxyRoute { + prefix: CODEX_ROUTE, + client_id: "codex", + credential: RouteCredential::Passthrough, + upstream: Upstream::ProxyUrl, + pairing_in_path: false, + tunnels_upgrades: false, + credential_less_ping: false, + }, + ProxyRoute { + prefix: "/vscode-copilot-passthrough", + client_id: "vscode-copilot", + credential: RouteCredential::Passthrough, + upstream: Upstream::ProxyUrl, + pairing_in_path: false, + tunnels_upgrades: false, + credential_less_ping: false, + }, + ProxyRoute { + prefix: CAPI_ROUTE, + client_id: "vscode-copilot", + credential: RouteCredential::Passthrough, + upstream: Upstream::ProxyUrl, + pairing_in_path: true, + tunnels_upgrades: true, + credential_less_ping: true, + }, + ProxyRoute { + prefix: "/vscode-copilot", + client_id: "vscode-copilot", + credential: RouteCredential::Gateway, + upstream: Upstream::Url, + pairing_in_path: false, + tunnels_upgrades: false, + credential_less_ping: false, + }, + ProxyRoute { + prefix: "/copilot-cli", + client_id: "copilot-cli", + credential: RouteCredential::Gateway, + upstream: Upstream::Url, + pairing_in_path: false, + tunnels_upgrades: false, + credential_less_ping: false, + }, +]; + +/// Splits `rest` (the path remainder after `CAPI_ROUTE`'s prefix, always +/// starting with `/`) into the first non-empty path segment (the pairing +/// value, still percent-encoded as sent) and the remainder starting with `/` +/// (or `/` itself when nothing follows), so the pairing segment never reaches +/// `match_route`'s caller, the forwarded upstream path, logs, or plan +/// reports. `None` when the first segment is missing or empty (bare prefix, +/// `/`, or `//...`). +/// +/// Called from `forward` for `CAPI_ROUTE` only, in place of the `PAIRING_HEADER` +/// check; the segment is compared in constant time and stripped before the +/// request is forwarded. +pub(crate) fn split_path_pairing(rest: &str) -> Option<(&str, &str)> { + let without_slash = rest.strip_prefix('/')?; + let (segment, remainder) = match without_slash.find('/') { + Some(index) => (&without_slash[..index], &without_slash[index..]), + None => (without_slash, "/"), + }; + (!segment.is_empty()).then_some((segment, remainder)) +} + +/// Only loopback hosts: a page on another origin that resolves a name to +/// 127.0.0.1 (DNS rebinding) sends that name in `Host`. Accepts any loopback +/// IP (the listen validation allows all of 127.0.0.0/8 and ::1), bracketed or +/// bare IPv6, `localhost` in any case with an optional trailing dot, each with +/// or without a port. A request without `Host` (HTTP/1.0) is refused. +pub(super) fn host_is_loopback(headers: &HeaderMap) -> bool { + let Some(host) = headers.get(HOST).and_then(|value| value.to_str().ok()) else { + return false; + }; + let host = host.trim(); + let name = if let Some(rest) = host.strip_prefix('[') { + match rest.split_once(']') { + Some((name, port)) if port.is_empty() || port.starts_with(':') => name, + _ => return false, + } + } else if host.matches(':').count() > 1 { + // Bare IPv6 without brackets cannot carry a port. + host + } else { + host.rsplit_once(':').map_or(host, |(name, _)| name) + }; + if let Ok(ip) = name.parse::() { + return ip.is_loopback(); + } + name.trim_end_matches('.').eq_ignore_ascii_case("localhost") +} + +/// Refuse `..` segments so a route cannot escape the upstream base it pins. +/// The check runs on a percent-decoded copy (the forwarded path stays as sent), +/// so `%2e%2e` and encoded separators are caught too; a malformed escape is +/// refused outright. +pub(super) fn path_escapes_base(path: &str) -> bool { + let Some(decoded) = percent_decode(path) else { + return true; + }; + decoded + .split(|byte| *byte == b'/' || *byte == b'\\') + .any(|segment| segment == b"..") +} + +/// Decodes `%XX` escapes to bytes. Escapes that are not two hex digits are +/// malformed and yield `None`; the bytes are not required to be UTF-8, since +/// the check only looks for separators and dots. +fn percent_decode(input: &str) -> Option> { + let bytes = input.as_bytes(); + let mut out = Vec::with_capacity(bytes.len()); + let mut i = 0; + while i < bytes.len() { + if bytes[i] == b'%' { + let hex = bytes.get(i + 1..i + 3)?; + if !hex.iter().all(u8::is_ascii_hexdigit) { + return None; + } + let value = u8::from_str_radix(std::str::from_utf8(hex).ok()?, 16).ok()?; + out.push(value); + i += 3; + } else { + out.push(bytes[i]); + i += 1; + } + } + Some(out) +} + +/// Whether the first path segment looks like one of the managed route names +/// without matching one exactly (wrong case, extra suffix, doubled slash). +/// A gateway sub-path whose first segment starts with a route name (for +/// example `/vscode-copilot-proxy/...`) is refused on the prefix-less route +/// as well; the check is on the literal path, so an encoded near miss is not +/// caught and reaches the gateway on the prefix-less route. +pub(super) fn resembles_route(path: &str) -> bool { + let first = path + .trim_start_matches('/') + .split('/') + .next() + .unwrap_or("") + .to_ascii_lowercase(); + !first.is_empty() + && ROUTES.iter().any(|route| { + let name = &route.prefix[1..]; + first == name || first.starts_with(name) || name.starts_with(&first) && first.len() >= 8 + }) +} + +pub(super) fn constant_time_eq(a: &[u8], b: &[u8]) -> bool { + if a.len() != b.len() { + return false; + } + a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0 +} + +/// The route for a request path, and the path with the prefix removed. +pub(super) fn match_route(path: &str) -> Option<(&'static ProxyRoute, &str)> { + ROUTES.iter().find_map(|route| { + let rest = path.strip_prefix(route.prefix)?; + (rest.is_empty() || rest.starts_with('/')) + .then_some((route, if rest.is_empty() { "/" } else { rest })) + }) +} + +pub(super) fn bearer_or_header(headers: &HeaderMap, header: &str) -> Option { + headers + .get(header) + .or_else(|| headers.get(AUTHORIZATION)) + .and_then(|value| value.to_str().ok()) + .map(|value| strip_bearer(value).trim().to_owned()) + .filter(|value| !value.is_empty()) +} + +/// The token of a `Bearer` credential; authentication schemes are +/// case-insensitive (RFC 9110, section 11.1). +fn strip_bearer(value: &str) -> &str { + match value.split_once(' ') { + Some((scheme, token)) if scheme.eq_ignore_ascii_case("bearer") => token, + _ => value, + } +} + +pub(super) fn upstream_base(gateway: &LlmGatewayConfig, upstream: Upstream) -> &url::Url { + match upstream { + Upstream::Url => &gateway.url, + Upstream::ProxyUrl => gateway.proxy_url.as_ref().unwrap_or(&gateway.url), + } +} diff --git a/crates/agent/src/llm_proxy/tests.rs b/crates/agent/src/llm_proxy/tests.rs new file mode 100644 index 00000000..318e43e6 --- /dev/null +++ b/crates/agent/src/llm_proxy/tests.rs @@ -0,0 +1,2146 @@ +use super::routes::ROUTES; +use super::*; +use crate::{enrollment::EnrollmentState, secret_store::SecretStore}; +use agentdesktop_core::{ + config::parse_daemon, + model::{DaemonInfo, DaemonScope, Discovery}, +}; +use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; +use http_body_util::StreamBody; +use hyper::body::Frame; +use sha2::{Digest, Sha256}; +use std::sync::Arc; +use tokio::sync::{mpsc, oneshot, watch}; +use tokio_stream::wrappers::ReceiverStream; + +// Covers the real socket path, credential replacement/rotation, opaque bytes, +// header forwarding, and delivery before the upstream response is complete. +#[test] +fn bearer_scheme_is_case_insensitive() { + for value in ["Bearer tok", "bearer tok", "BEARER tok", "tok"] { + let mut headers = HeaderMap::new(); + headers.insert(AUTHORIZATION, HeaderValue::from_str(value).unwrap()); + assert_eq!( + bearer_or_header(&headers, "x-llm-token").as_deref(), + Some("tok"), + "{value}" + ); + } + let mut headers = HeaderMap::new(); + headers.insert("x-llm-token", HeaderValue::from_static("bearer gho_x")); + assert_eq!( + bearer_or_header(&headers, "x-llm-token").as_deref(), + Some("gho_x") + ); +} + +#[tokio::test] +async fn streams_opaque_requests_with_current_credentials() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let (release_tx, release_rx) = oneshot::channel(); + let release = Arc::new(tokio::sync::Mutex::new(Some(release_rx))); + let upstream_task = tokio::spawn(async move { + let (socket, _) = upstream.accept().await.unwrap(); + hyper::server::conn::http1::Builder::new().serve_connection( + TokioIo::new(socket), service_fn(move |request: Request| { + let release = release.clone(); + async move { + assert_eq!(request.uri(), "/gateway/v1/messages?value=%2F"); + assert_eq!(request.headers()[HOST], address.to_string()); + assert!(!request.headers().contains_key("x-api-key")); + assert!(!request.headers().contains_key("x-hop")); + assert!(!request.headers().contains_key(PAIRING_HEADER)); + assert_eq!(request.headers()["x-llm-token"], "ghu_test"); + assert_eq!(request.headers()["anthropic-version"], "2023-06-01"); + let gate = release.lock().await.take(); + assert_eq!(request.headers()[AUTHORIZATION], + if gate.is_some() { "Bearer first" } else { "Bearer second" }); + assert_eq!(request.into_body().collect().await.unwrap().to_bytes(), + Bytes::from_static(b"\xffnot-json\x00")); + let (tx, rx) = mpsc::channel::, Infallible>>(2); + tokio::spawn(async move { + tx.send(Ok(Frame::data(Bytes::from_static(b"data: first\n\n")))).await.unwrap(); + if let Some(gate) = gate { gate.await.unwrap(); } + let _ = tx.send(Ok(Frame::data(Bytes::from_static(b"data: last\n\n")))).await; + }); + Ok::<_, Infallible>(Response::builder() + .header("content-type", "text/event-stream") + .header("connection", "x-upstream-hop") + .header("x-upstream-hop", "remove") + .body(StreamBody::new(ReceiverStream::new(rx))).unwrap()) + } + }) + ).await.unwrap(); + }); + let dir = tempfile::tempdir().unwrap(); + let store = SecretStore::new(dir.path()).unwrap(); + let account = URL_SAFE_NO_PAD.encode(Sha256::digest(b"https://issuer.example/\0proxy-test")); + let save_token = |token: &str| store.set("dev.agentdesktop.gateway-oidc", &account, + &serde_json::json!({"accessToken": token, "refreshToken": null, + "expiresAtUnixSeconds": 4_000_000_000u64, + "tokenEndpoint": "https://issuer.example/token"}).to_string()).unwrap(); + save_token("first"); + store.set("dev.agentdesktop.github-oauth", "github-test", &serde_json::json!({ + "accessToken": "ghu_test", "expiresAt": 4_000_000_000u64, + "refreshToken": null, "refreshExpiresAt": 0, + }).to_string()).unwrap(); + let config = parse_daemon(&format!("llmGateway:\n url: http://{address}/gateway\n githubOAuth:\n clientId: github-test\n authentication:\n type: oidc\n issuer: https://issuer.example/\n clientId: proxy-test\n")).unwrap(); + let (_, discovery) = watch::channel(Arc::new(Discovery { agents: vec![], model_runtimes: vec![] })); + let state = AppState { + controller_status: None, + config, + daemon_info: DaemonInfo { version: "test".into(), scope: DaemonScope::User, + config_path: String::new(), state_directory: String::new(), + inventory_interval: Duration::from_secs(60), controller: None, llm_proxy: None, + usage_receiver: None }, + discovery, enrollment: EnrollmentState::new(false), state_dir: dir.path().to_owned(), + oidc_callback_listen: None, telemetry: None, logout: None, usage: None, + programs: None, + }; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = Client::builder(TokioExecutor::new()).build_http(); + let mut release_tx = Some(release_tx); + for _ in 0..2 { + let request = Request::builder().method(Method::POST) + .uri(format!("http://{proxy_address}/v1/messages?value=%2F")) + .header(AUTHORIZATION, "Bearer must-not-forward") + .header("x-api-key", "must-not-forward") + .header("x-llm-token", "must-not-forward") + .header(PAIRING_HEADER, PAIRING) + .header(CONNECTION, "x-hop").header("x-hop", "remove") + .header("anthropic-version", "2023-06-01") + .body(Full::new(Bytes::from_static(b"\xffnot-json\x00"))).unwrap(); + let response = client.request(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response.headers()["content-type"], "text/event-stream"); + assert!(!response.headers().contains_key("x-upstream-hop")); + let mut body = response.into_body(); + let first = body.frame().await.unwrap().unwrap().into_data().unwrap(); + if let Some(release) = release_tx.take() { + assert_eq!(first, "data: first\n\n"); + release.send(()).unwrap(); + } + let rest = body.collect().await.unwrap().to_bytes(); + assert_eq!([first.as_ref(), rest.as_ref()].concat(), b"data: first\n\ndata: last\n\n"); + save_token("second"); + } + proxy.abort(); + let _ = proxy.await; + drop(client); + upstream_task.abort(); + }).await.unwrap(); +} + +// source: request. The client's own credential is what reaches the gateway +// in x-llm-token, and the daemon contributes only the gateway identity. +// This is the VS Code Copilot shape: Copilot has already exchanged its +// GitHub session for a provider token and sends it in Authorization. +#[tokio::test] +async fn forwards_the_client_credential_when_source_is_request() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let upstream_task = tokio::spawn(async move { + let (socket, _) = upstream.accept().await.unwrap(); + hyper::server::conn::http1::Builder::new() + .serve_connection( + TokioIo::new(socket), + service_fn(move |request: Request| async move { + // The Bearer prefix is stripped here and re-added by the + // gateway, so the token travels bare. + assert_eq!(request.headers()["x-llm-token"], "tid=from-client"); + assert_eq!(request.headers()[AUTHORIZATION], "Bearer identity"); + Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static(b"ok")))) + }), + ) + .await + .unwrap(); + }); + let (dir, state) = request_source_state(address).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let request = Request::builder() + .method(Method::POST) + .uri(format!("http://{proxy_address}/chat/completions")) + .header(PAIRING_HEADER, PAIRING) + .header(AUTHORIZATION, "Bearer tid=from-client") + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + let response = client.request(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + proxy.abort(); + let _ = proxy.await; + drop(client); + upstream_task.abort(); + drop(dir); + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn codex_subscription_separates_credentials_and_preserves_responses_request() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let upstream_task = tokio::spawn(async move { + let (socket, _) = upstream.accept().await.unwrap(); + hyper::server::conn::http1::Builder::new() + .serve_connection( + TokioIo::new(socket), + service_fn(move |request: Request| async move { + assert_eq!(request.uri(), "/backend-api/codex/responses?test=1"); + assert_eq!(request.headers()[AUTHORIZATION], "Bearer identity"); + assert_eq!( + request.headers()["x-llm-token"], + "chatgpt-subscription-token" + ); + assert_eq!(request.headers()["chatgpt-account-id"], "test-account"); + assert!(!request.headers().contains_key(PAIRING_HEADER)); + assert_eq!( + request.into_body().collect().await.unwrap().to_bytes(), + "{\"model\":\"gpt-5.4\"}" + ); + Ok::<_, Infallible>( + Response::builder() + .header(CONTENT_TYPE, "text/event-stream") + .body(Full::new(Bytes::from_static( + b"data: subscription-response\n\n", + ))) + .unwrap(), + ) + }), + ) + .await + .unwrap(); + }); + let (dir, mut state) = request_source_state(address).await; + let gateway = state.config.llm_gateway.as_mut().unwrap(); + gateway.github_oauth = None; + gateway.proxy_url = Some( + format!("http://{address}/backend-api/codex") + .parse() + .unwrap(), + ); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + for (pairing, credential, expected) in [ + ( + "wrong-pairing", + Some("chatgpt-subscription-token"), + StatusCode::FORBIDDEN, + ), + (PAIRING, None, StatusCode::UNAUTHORIZED), + (PAIRING, Some(PAIRING), StatusCode::UNAUTHORIZED), + (PAIRING, Some("chatgpt-subscription-token"), StatusCode::OK), + ] { + let mut request = Request::builder() + .method(Method::POST) + .uri(format!("http://{proxy_address}/codex/responses?test=1")) + .header(PAIRING_HEADER, pairing) + .header("chatgpt-account-id", "test-account"); + if let Some(credential) = credential { + request = request.header(AUTHORIZATION, format!("Bearer {credential}")); + } + let response = client + .request( + request + .body(Full::new(Bytes::from_static(b"{\"model\":\"gpt-5.4\"}"))) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), expected); + if expected == StatusCode::OK { + assert_eq!(response.headers()[CONTENT_TYPE], "text/event-stream"); + assert_eq!( + response.into_body().collect().await.unwrap().to_bytes(), + "data: subscription-response\n\n" + ); + } + } + proxy.abort(); + let _ = proxy.await; + upstream_task.abort(); + drop(dir); + }) + .await + .unwrap(); +} + +// Without a client credential the gateway would swap in an empty +// Authorization and the provider would fail in a way that looks like a +// gateway fault, so the proxy refuses instead of forwarding. +#[tokio::test] +async fn rejects_a_request_with_no_client_credential() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let (dir, state) = request_source_state(address).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let request = Request::builder() + .method(Method::POST) + .uri(format!("http://{proxy_address}/chat/completions")) + .header(PAIRING_HEADER, PAIRING) + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + let response = client.request(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + // The default route is paired too: without the header it is refused + // before any credential work. + let request = Request::builder() + .method(Method::POST) + .uri(format!("http://{proxy_address}/chat/completions")) + .header(AUTHORIZATION, "Bearer tid=from-client") + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + assert_eq!( + error_code(client.request(request).await.unwrap()).await, + (StatusCode::FORBIDDEN, "pairing_invalid".to_owned()) + ); + proxy.abort(); + let _ = proxy.await; + drop(client); + drop(dir); + }) + .await + .unwrap(); +} + +const PAIRING: &str = "pairing-secret-value-0123456789abcdef"; + +fn default_config() -> ProxyConfig { + ProxyConfig { + default_client_id: "vscode".into(), + pairing: Arc::from(PAIRING), + } +} + +async fn error_code(response: Response) -> (StatusCode, String) { + let status = response.status(); + assert_eq!(response.headers()[CONTENT_TYPE], "application/json"); + let body = response.into_body().collect().await.unwrap().to_bytes(); + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + (status, json["error"]["code"].as_str().unwrap().to_owned()) +} + +#[tokio::test] +async fn refuses_non_loopback_hosts_and_options_with_json_errors() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let (dir, state) = request_source_state(upstream.local_addr().unwrap()).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let rebound = Request::builder() + .method(Method::POST) + .uri(format!( + "http://{proxy_address}/copilot-cli/v1/chat/completions" + )) + .header(HOST, "attacker.example") + .header(PAIRING_HEADER, "pairing-secret-value-0123456789abcdef") + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + assert_eq!( + error_code(client.request(rebound).await.unwrap()).await, + (StatusCode::FORBIDDEN, "host_not_allowed".to_owned()) + ); + let options = Request::builder() + .method(Method::OPTIONS) + .uri(format!( + "http://{proxy_address}/copilot-cli/v1/chat/completions" + )) + .body(Full::new(Bytes::new())) + .unwrap(); + assert_eq!( + error_code(client.request(options).await.unwrap()).await, + ( + StatusCode::METHOD_NOT_ALLOWED, + "method_not_allowed".to_owned() + ) + ); + proxy.abort(); + let _ = proxy.await; + drop(client); + drop(dir); + }) + .await + .unwrap(); +} + +// A prefixed route needs the pairing value, strips its prefix, sends only the +// gateway identity (no client credential leaks through), and is refused with a +// JSON error without the pairing. +#[tokio::test] +async fn prefixed_gateway_route_requires_pairing_and_strips_the_prefix() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let upstream_task = tokio::spawn(async move { + let (socket, _) = upstream.accept().await.unwrap(); + hyper::server::conn::http1::Builder::new() + .serve_connection( + TokioIo::new(socket), + service_fn(move |request: Request| async move { + assert_eq!(request.uri(), "/v1/chat/completions?stream=true"); + assert_eq!(request.headers()[AUTHORIZATION], "Bearer identity"); + assert!(!request.headers().contains_key("x-llm-token")); + assert!(!request.headers().contains_key(PAIRING_HEADER)); + Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static(b"ok")))) + }), + ) + .await + .unwrap(); + }); + let (dir, state) = request_source_state(address).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state.clone(), default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let unpaired = Request::builder() + .method(Method::POST) + .uri(format!( + "http://{proxy_address}/copilot-cli/v1/chat/completions" + )) + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + assert_eq!( + error_code(client.request(unpaired).await.unwrap()).await, + (StatusCode::FORBIDDEN, "pairing_invalid".to_owned()) + ); + // The pairing is only accepted in its own header, never as the bearer. + let as_bearer = Request::builder() + .method(Method::POST) + .uri(format!( + "http://{proxy_address}/copilot-cli/v1/chat/completions" + )) + .header(AUTHORIZATION, format!("Bearer {PAIRING}")) + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + assert_eq!( + error_code(client.request(as_bearer).await.unwrap()).await, + (StatusCode::FORBIDDEN, "pairing_invalid".to_owned()) + ); + let wrong = Request::builder() + .method(Method::POST) + .uri(format!( + "http://{proxy_address}/copilot-cli/v1/chat/completions" + )) + .header(PAIRING_HEADER, "pairing-secret-value-0123456789abcdeX") + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + assert_eq!( + error_code(client.request(wrong).await.unwrap()).await, + (StatusCode::FORBIDDEN, "pairing_invalid".to_owned()) + ); + let paired = Request::builder() + .method(Method::POST) + .uri(format!( + "http://{proxy_address}/copilot-cli/v1/chat/completions?stream=true" + )) + .header(PAIRING_HEADER, PAIRING) + .header(AUTHORIZATION, "Bearer a-real-key-by-mistake") + .header("x-llm-token", "must-not-forward") + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + let response = client.request(paired).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + proxy.abort(); + let _ = proxy.await; + // A '..' segment cannot escape the upstream base the route pins. + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let request = Request::builder() + .method(Method::POST) + .uri(format!("http://{proxy_address}/copilot-cli/../admin")) + .header(PAIRING_HEADER, PAIRING) + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + assert_eq!( + error_code(client.request(request).await.unwrap()).await, + (StatusCode::BAD_REQUEST, "path_invalid".to_owned()) + ); + proxy.abort(); + let _ = proxy.await; + drop(client); + upstream_task.abort(); + drop(dir); + }) + .await + .unwrap(); +} + +// The pass-through route moves the client's token to x-llm-token, keeps the +// gateway identity in Authorization, forwards to proxyUrl, and refuses a +// request without a client token. The pairing must come in its own header +// here, because Authorization carries the client token. +#[tokio::test] +async fn passthrough_route_moves_the_client_token_and_uses_proxy_url() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let upstream_task = tokio::spawn(async move { + let (socket, _) = upstream.accept().await.unwrap(); + hyper::server::conn::http1::Builder::new() + .serve_connection( + TokioIo::new(socket), + service_fn(move |request: Request| async move { + assert_eq!(request.uri(), "/copilot-proxy/chat/completions"); + assert_eq!(request.headers()["x-llm-token"], "tid=from-client"); + assert_eq!(request.headers()[AUTHORIZATION], "Bearer identity"); + assert!(!request.headers().contains_key(PAIRING_HEADER)); + Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static(b"ok")))) + }), + ) + .await + .unwrap(); + }); + let (dir, mut state) = request_source_state(address).await; + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + // Without proxyUrl the pass-through route is unavailable rather than + // sending x-llm-token to the plain LLM route. + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state.clone(), default_config())); + let request = Request::builder() + .method(Method::POST) + .uri(format!( + "http://{proxy_address}/vscode-copilot-passthrough/chat/completions" + )) + .header(PAIRING_HEADER, PAIRING) + .header(AUTHORIZATION, "Bearer tid=from-client") + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + assert_eq!( + error_code(client.request(request).await.unwrap()).await, + ( + StatusCode::SERVICE_UNAVAILABLE, + "agentdesktop_unavailable".to_owned() + ) + ); + proxy.abort(); + let _ = proxy.await; + state.config.llm_gateway.as_mut().unwrap().proxy_url = + Some(format!("http://{address}/copilot-proxy").parse().unwrap()); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let no_token = Request::builder() + .method(Method::POST) + .uri(format!( + "http://{proxy_address}/vscode-copilot-passthrough/chat/completions" + )) + .header(PAIRING_HEADER, "pairing-secret-value-0123456789abcdef") + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + assert_eq!( + error_code(client.request(no_token).await.unwrap()).await, + ( + StatusCode::UNAUTHORIZED, + "client_credential_missing".to_owned() + ) + ); + // The pairing in Authorization is not a client token: refused, never forwarded. + let pairing_as_token = Request::builder() + .method(Method::POST) + .uri(format!( + "http://{proxy_address}/vscode-copilot-passthrough/chat/completions" + )) + .header(PAIRING_HEADER, PAIRING) + .header(AUTHORIZATION, format!("Bearer {PAIRING}")) + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + assert_eq!( + error_code(client.request(pairing_as_token).await.unwrap()).await, + ( + StatusCode::UNAUTHORIZED, + "client_credential_missing".to_owned() + ) + ); + let request = Request::builder() + .method(Method::POST) + .uri(format!( + "http://{proxy_address}/vscode-copilot-passthrough/chat/completions" + )) + .header(PAIRING_HEADER, "pairing-secret-value-0123456789abcdef") + .header(AUTHORIZATION, "Bearer tid=from-client") + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + let response = client.request(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + proxy.abort(); + let _ = proxy.await; + drop(client); + upstream_task.abort(); + drop(dir); + }) + .await + .unwrap(); +} + +#[test] +fn route_matching_prefers_the_longer_prefix_and_strips_it() { + let (route, rest) = match_route("/vscode-copilot-passthrough/chat/completions").unwrap(); + assert_eq!(route.client_id, "vscode-copilot"); + assert_eq!(route.credential, RouteCredential::Passthrough); + assert_eq!(rest, "/chat/completions"); + let (route, rest) = match_route("/vscode-copilot/v1/chat/completions").unwrap(); + assert_eq!(route.credential, RouteCredential::Gateway); + assert_eq!(rest, "/v1/chat/completions"); + assert_eq!(match_route("/copilot-cli").unwrap().1, "/"); + assert!(match_route("/copilot-client/v1").is_none()); + assert!(match_route("/v1/chat/completions").is_none()); +} + +#[test] +fn host_check_accepts_loopback_forms_only() { + let mut headers = HeaderMap::new(); + let accepted = [ + "127.0.0.1:18095", + "localhost:18095", + "[::1]:18095", + "localhost", + "127.0.0.1", + "127.0.0.2:1", + "Localhost:18095", + "LOCALHOST", + "localhost.:18095", + "[::1]", + "::1", + ]; + for ok in accepted { + headers.insert(HOST, HeaderValue::from_static(ok)); + assert!(host_is_loopback(&headers), "{ok}"); + } + let refused = [ + "attacker.example", + "127.0.0.1.attacker.example:18095", + "10.0.0.5:18095", + "", + "[::1", + "[fe80::1]:18095", + "localhost.attacker.example", + "127.0.0.1:18095:1", + ]; + for bad in refused { + headers.insert(HOST, HeaderValue::from_static(bad)); + assert!(!host_is_loopback(&headers), "{bad}"); + } + headers.remove(HOST); + assert!(!host_is_loopback(&headers)); +} + +#[test] +fn pairing_is_created_once_and_reused() { + let dir = tempfile::tempdir().unwrap(); + let first = load_or_create_pairing(dir.path()).unwrap(); + let second = load_or_create_pairing(dir.path()).unwrap(); + assert_eq!(first, second); + assert!(first.len() >= 43); + let metadata = std::fs::metadata(dir.path().join(PAIRING_FILE)).unwrap(); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + assert_eq!(metadata.permissions().mode() & 0o777, 0o600); + } + let _ = metadata; +} + +#[test] +fn pairing_file_is_regenerated_only_when_unusable() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join(PAIRING_FILE); + // Too short: replaced. + std::fs::write(&path, b"short").unwrap(); + let regenerated = load_or_create_pairing(dir.path()).unwrap(); + assert!(regenerated.len() >= 43); + assert_eq!(std::fs::read_to_string(&path).unwrap(), *regenerated); + // Not text: replaced. + std::fs::write(&path, [0xff, 0xfe, 0x00, 0x01]).unwrap(); + let again = load_or_create_pairing(dir.path()).unwrap(); + assert_ne!(again, regenerated); + assert_eq!(std::fs::read_to_string(&path).unwrap(), *again); + // Any other read error is an error, and the value on disk is untouched. + let other = tempfile::tempdir().unwrap(); + std::fs::create_dir(other.path().join(PAIRING_FILE)).unwrap(); + let error = load_or_create_pairing(other.path()).unwrap_err(); + assert!(format!("{error:#}").contains("read "), "{error:#}"); + assert!(other.path().join(PAIRING_FILE).is_dir()); + // An existing value that cannot be read (permissions) is kept as is. + #[cfg(unix)] + if !nix_is_root() { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o000)).unwrap(); + let error = load_or_create_pairing(dir.path()).unwrap_err(); + assert!(format!("{error:#}").contains("read "), "{error:#}"); + std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600)).unwrap(); + assert_eq!(*load_or_create_pairing(dir.path()).unwrap(), *again); + } + // A state directory that cannot be created is an error, not a pairing. + let blocked = tempfile::tempdir().unwrap(); + let file = blocked.path().join("state"); + std::fs::write(&file, b"x").unwrap(); + assert!(load_or_create_pairing(&file.join("nested")).is_err()); +} + +#[cfg(unix)] +fn nix_is_root() -> bool { + std::fs::metadata("/proc/self") + .map(|m| { + use std::os::unix::fs::MetadataExt; + m.uid() == 0 + }) + .unwrap_or(false) +} + +#[test] +fn route_table_maps_each_prefix_to_its_identity_mode_and_upstream() { + let expected = [ + ( + "/codex/responses", + "codex", + RouteCredential::Passthrough, + Upstream::ProxyUrl, + ), + ( + "/copilot-cli/v1/x", + "copilot-cli", + RouteCredential::Gateway, + Upstream::Url, + ), + ( + "/vscode-copilot/v1/x", + "vscode-copilot", + RouteCredential::Gateway, + Upstream::Url, + ), + ( + "/vscode-copilot-passthrough/chat/completions", + "vscode-copilot", + RouteCredential::Passthrough, + Upstream::ProxyUrl, + ), + ( + "/vscode-copilot-capi/PAIRING/responses", + "vscode-copilot", + RouteCredential::Passthrough, + Upstream::ProxyUrl, + ), + ]; + for (path, client_id, credential, upstream) in expected { + let (route, _) = match_route(path).unwrap_or_else(|| panic!("{path}")); + assert_eq!(route.client_id, client_id, "{path}"); + assert_eq!(route.credential, credential, "{path}"); + assert_eq!(route.upstream, upstream, "{path}"); + } + assert_eq!(ROUTES.len(), 5); + assert_eq!(CAPI_ROUTE, "/vscode-copilot-capi"); +} + +#[test] +fn path_guard_catches_encoded_traversal_and_near_miss_routes() { + for escaping in [ + "/copilot-cli/../admin", + "/copilot-cli/%2e%2e/admin", + "/copilot-cli/%2E%2E/admin", + "/copilot-cli/%2e%2e%2fadmin", + "/copilot-cli/..%5cadmin", + "/copilot-cli/%zz", + "/copilot-cli/%+f", + "/copilot-cli/%2", + ] { + assert!(path_escapes_base(escaping), "{escaping}"); + } + for fine in [ + "/copilot-cli/v1/chat/completions", + "/v1/messages", + "/copilot-cli/a.b/c%20d", + "/copilot-cli/caf%e9", + "/copilot-cli/%2e%2ex", + ] { + assert!(!path_escapes_base(fine), "{fine}"); + } + for near in [ + "/copilot-cli2/v1", + "/Copilot-cli/v1", + "//copilot-cli/v1", + "/vscode-copilotx/v1", + ] { + assert!( + match_route(near).is_none() && resembles_route(near), + "{near}" + ); + } + for other in ["/v1/chat/completions", "/chat/completions", "/models"] { + assert!(!resembles_route(other), "{other}"); + } +} + +// The pairing is checked before any configuration is read: an unpaired +// caller gets 403 even when the applied configuration is unreadable, and a +// paired caller gets a generic message without file detail. +#[tokio::test] +async fn pairing_is_checked_before_configuration_is_read() { + tokio::time::timeout(Duration::from_secs(15), async { + let dir = tempfile::tempdir().unwrap(); + std::fs::write( + dir.path().join("remote-config.yaml"), + "llmGateway: [not: an: object", + ) + .unwrap(); + let config = parse_daemon("controller:\n address: https://127.0.0.1:1\n").unwrap(); + let (_, discovery) = watch::channel(Arc::new(Discovery { + agents: vec![], + model_runtimes: vec![], + })); + let state = AppState { + controller_status: None, + config, + daemon_info: DaemonInfo { + version: "test".into(), + scope: DaemonScope::User, + config_path: String::new(), + state_directory: String::new(), + inventory_interval: Duration::from_secs(60), + controller: None, + llm_proxy: None, + usage_receiver: None, + }, + discovery, + enrollment: EnrollmentState::new(true), + state_dir: dir.path().to_owned(), + oidc_callback_listen: None, + telemetry: None, + logout: None, + usage: None, + programs: None, + }; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let unpaired = Request::builder() + .method(Method::POST) + .uri(format!( + "http://{proxy_address}/copilot-cli/v1/chat/completions" + )) + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + assert_eq!( + error_code(client.request(unpaired).await.unwrap()).await, + (StatusCode::FORBIDDEN, "pairing_invalid".to_owned()) + ); + // The path checks are pure string checks and run before the + // configuration is read: a paired caller with a bad path gets the + // path error, not the configuration error. + for (bad_path, code) in [ + ("/copilot-cli/../admin", "path_invalid"), + ("/Copilot-cli/v1/chat/completions", "route_unknown"), + ] { + let request = Request::builder() + .method(Method::POST) + .uri(format!("http://{proxy_address}{bad_path}")) + .header(PAIRING_HEADER, PAIRING) + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + let (_, got) = error_code(client.request(request).await.unwrap()).await; + assert_eq!(got, code, "{bad_path}"); + } + let paired = Request::builder() + .method(Method::POST) + .uri(format!( + "http://{proxy_address}/copilot-cli/v1/chat/completions" + )) + .header(PAIRING_HEADER, PAIRING) + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + let response = client.request(paired).await.unwrap(); + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); + let body = response.into_body().collect().await.unwrap().to_bytes(); + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + let message = json["error"]["message"].as_str().unwrap(); + assert!( + !message.contains("remote-config") && !message.contains(dir.path().to_str().unwrap()), + "{message}" + ); + proxy.abort(); + let _ = proxy.await; + drop(client); + drop(dir); + }) + .await + .unwrap(); +} + +// The retry applies only to a credential served from the cache. OIDC +// credentials are never cached, so a 401 from the gateway is returned as is +// after a single attempt, with the buffered body delivered once. The cached +// path is covered by `cached_credential_is_used_and_dropped_on_401` below. +#[tokio::test] +async fn uncached_credentials_are_not_retried_on_401() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let attempts = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let seen = attempts.clone(); + let upstream_task = tokio::spawn(async move { + loop { + let (socket, _) = upstream.accept().await.unwrap(); + let seen = seen.clone(); + tokio::spawn(async move { + let _ = hyper::server::conn::http1::Builder::new() + .serve_connection( + TokioIo::new(socket), + service_fn(move |request: Request| { + let seen = seen.clone(); + async move { + let n = seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + assert_eq!(request.headers()[AUTHORIZATION], "Bearer identity"); + let body = + request.into_body().collect().await.unwrap().to_bytes(); + assert_eq!(body, "{\"n\":1}"); + // One attempt per client request: the first + // is answered 401, the second 200. + let status = if n == 1 { + StatusCode::OK + } else { + StatusCode::UNAUTHORIZED + }; + Ok::<_, Infallible>( + Response::builder() + .status(status) + .body(Full::new(Bytes::from_static(b"upstream"))) + .unwrap(), + ) + } + }), + ) + .await; + }); + } + }); + let (dir, state) = request_source_state(address).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let request = || { + Request::builder() + .method(Method::POST) + .uri(format!("http://{proxy_address}/chat/completions")) + .header(PAIRING_HEADER, PAIRING) + .header(AUTHORIZATION, "Bearer tid=from-client") + .body(Full::new(Bytes::from_static(b"{\"n\":1}"))) + .unwrap() + }; + let response = client.request(request()).await.unwrap(); + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); + let response = client.request(request()).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 2); + proxy.abort(); + let _ = proxy.await; + drop(client); + upstream_task.abort(); + drop(dir); + }) + .await + .unwrap(); +} + +// A pre-seeded cache stands in for a controller: the proxy sends the cached +// credential; when the gateway accepts it the entry stays; when the gateway +// rejects it the entry is dropped and the request is retried, which here +// needs a controller that does not exist, so the client gets +// `agentdesktop_credential` and the gateway saw the body exactly once. +#[tokio::test] +async fn cached_credential_is_used_and_dropped_on_401() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + type Seen = Arc)>>>; + let seen: Seen = Arc::default(); + let record = seen.clone(); + let upstream_task = tokio::spawn(async move { + loop { + let (socket, _) = upstream.accept().await.unwrap(); + let record = record.clone(); + tokio::spawn(async move { + let _ = hyper::server::conn::http1::Builder::new() + .serve_connection( + TokioIo::new(socket), + service_fn(move |request: Request| { + let record = record.clone(); + async move { + let authorization = request + .headers() + .get(AUTHORIZATION) + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_owned(); + let body = request + .into_body() + .collect() + .await + .unwrap() + .to_bytes() + .to_vec(); + let status = if authorization == "Bearer good" { + StatusCode::OK + } else { + StatusCode::UNAUTHORIZED + }; + record.lock().unwrap().push((authorization, body)); + Ok::<_, Infallible>( + Response::builder() + .status(status) + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(), + ) + } + }), + ) + .await; + }); + } + }); + let (dir, state) = controller_jwt_state(address, false).await; + let cache = Arc::new(CredentialCache::default()); + let key = format!("copilot-cli\u{0}http://{address}/\u{0}agentgateway"); + let issued = |value: &str| agentdesktop_core::model::LlmGatewayCredential { + credential: value.to_owned(), + expires_at_unix_seconds: now_unix() + 600, + }; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve_with_cache( + listener, + state, + default_config(), + cache.clone(), + )); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let request = |body: &'static [u8]| { + Request::builder() + .method(Method::POST) + .uri(format!( + "http://{proxy_address}/copilot-cli/v1/chat/completions" + )) + .header(PAIRING_HEADER, PAIRING) + .body(Full::new(Bytes::from_static(body))) + .unwrap() + }; + // Accepted: served from the cache, entry kept. + cache.insert(&key, "device-test", &issued("good")); + let response = client.request(request(b"{\"n\":1}")).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(cache.get(&key, "device-test").as_deref(), Some("good")); + // Rejected: entry dropped, one retry that needs a controller. + cache.insert(&key, "device-test", &issued("stale")); + let (status, code) = error_code(client.request(request(b"{\"n\":2}")).await.unwrap()).await; + assert_eq!(code, "agentdesktop_credential", "{status}"); + assert_eq!(cache.get(&key, "device-test"), None); + let seen = seen.lock().unwrap().clone(); + assert_eq!( + seen, + vec![ + ("Bearer good".to_owned(), b"{\"n\":1}".to_vec()), + ("Bearer stale".to_owned(), b"{\"n\":2}".to_vec()), + ], + "each body reaches the gateway exactly once per attempt" + ); + proxy.abort(); + let _ = proxy.await; + drop(client); + upstream_task.abort(); + drop(dir); + }) + .await + .unwrap(); +} + +// With a client token forwarded (prefix-less route, githubOAuth source +// request) a 401 is returned as is: no retry, and the cached gateway +// credential is kept. +#[tokio::test] +async fn client_token_401_is_not_retried_and_keeps_the_cached_credential() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let attempts = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let seen = attempts.clone(); + let upstream_task = tokio::spawn(async move { + loop { + let (socket, _) = upstream.accept().await.unwrap(); + let seen = seen.clone(); + tokio::spawn(async move { + let _ = hyper::server::conn::http1::Builder::new() + .serve_connection( + TokioIo::new(socket), + service_fn(move |request: Request| { + let seen = seen.clone(); + async move { + seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + assert_eq!(request.headers()[AUTHORIZATION], "Bearer cached"); + assert_eq!(request.headers()["x-llm-token"], "tid=bad"); + Ok::<_, Infallible>( + Response::builder() + .status(StatusCode::UNAUTHORIZED) + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(), + ) + } + }), + ) + .await; + }); + } + }); + let (dir, state) = controller_jwt_state(address, true).await; + let cache = Arc::new(CredentialCache::default()); + let key = format!("vscode\u{0}http://{address}/\u{0}agentgateway"); + cache.insert( + &key, + "device-test", + &agentdesktop_core::model::LlmGatewayCredential { + credential: "cached".to_owned(), + expires_at_unix_seconds: now_unix() + 600, + }, + ); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve_with_cache( + listener, + state, + default_config(), + cache.clone(), + )); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let request = Request::builder() + .method(Method::POST) + .uri(format!("http://{proxy_address}/chat/completions")) + .header(PAIRING_HEADER, PAIRING) + .header(AUTHORIZATION, "Bearer tid=bad") + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + let response = client.request(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); + assert_eq!(cache.get(&key, "device-test").as_deref(), Some("cached")); + proxy.abort(); + let _ = proxy.await; + drop(client); + upstream_task.abort(); + drop(dir); + }) + .await + .unwrap(); +} + +// Controller-JWT state with an identity.json but no controller configured: +// the device id is readable for the cache key, a fetch fails at once +// ("requires a controller"), so no timeout is involved. +async fn controller_jwt_state( + address: std::net::SocketAddr, + client_token: bool, +) -> (tempfile::TempDir, AppState) { + let dir = tempfile::tempdir().unwrap(); + std::fs::write( + dir.path().join("identity.json"), + serde_json::json!({ + "deviceId": "device-test", + "clientCertificatePem": "", + "clientCertificateExpiresAtUnixSeconds": 4_000_000_000u64, + "oauthTokenEndpoint": "https://controller.example/token", + "oauthClientId": "device", + }) + .to_string(), + ) + .unwrap(); + let github = if client_token { + " githubOAuth:\n source: request\n" + } else { + "" + }; + let config = parse_daemon(&format!( + "llmGateway:\n url: http://{address}\n{github} authentication:\n type: controllerJwt\n audience: agentgateway\n allowedClientIds: [copilot-cli, vscode]\n" + )).unwrap(); + let (_, discovery) = watch::channel(Arc::new(Discovery { + agents: vec![], + model_runtimes: vec![], + })); + let state = AppState { + controller_status: None, + config, + daemon_info: DaemonInfo { + version: "test".into(), + scope: DaemonScope::User, + config_path: String::new(), + state_directory: String::new(), + inventory_interval: Duration::from_secs(60), + controller: None, + llm_proxy: None, + usage_receiver: None, + }, + discovery, + enrollment: EnrollmentState::new(false), + state_dir: dir.path().to_owned(), + oidc_callback_listen: None, + telemetry: None, + logout: None, + usage: None, + programs: None, + }; + (dir, state) +} + +#[tokio::test] +async fn refuses_a_body_over_the_cap_with_a_json_error() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let (dir, state) = request_source_state(upstream.local_addr().unwrap()).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let request = Request::builder() + .method(Method::POST) + .uri(format!("http://{proxy_address}/chat/completions")) + .header(PAIRING_HEADER, PAIRING) + .header(AUTHORIZATION, "Bearer tid=from-client") + .body(Full::new(Bytes::from(vec![b'x'; MAX_REQUEST_BODY + 1]))) + .unwrap(); + assert_eq!( + error_code(client.request(request).await.unwrap()).await, + (StatusCode::PAYLOAD_TOO_LARGE, "body_too_large".to_owned()) + ); + proxy.abort(); + let _ = proxy.await; + drop(client); + drop(dir); + }) + .await + .unwrap(); +} + +async fn request_source_state(address: std::net::SocketAddr) -> (tempfile::TempDir, AppState) { + let dir = tempfile::tempdir().unwrap(); + let store = SecretStore::new(dir.path()).unwrap(); + let account = URL_SAFE_NO_PAD.encode(Sha256::digest(b"https://issuer.example/\0proxy-test")); + store + .set( + "dev.agentdesktop.gateway-oidc", + &account, + &serde_json::json!({"accessToken": "identity", "refreshToken": null, + "expiresAtUnixSeconds": 4_000_000_000u64, + "tokenEndpoint": "https://issuer.example/token"}) + .to_string(), + ) + .unwrap(); + let config = parse_daemon(&format!( + "llmGateway:\n url: http://{address}\n githubOAuth:\n source: request\n authentication:\n type: oidc\n issuer: https://issuer.example/\n clientId: proxy-test\n" + )).unwrap(); + let (_, discovery) = watch::channel(Arc::new(Discovery { + agents: vec![], + model_runtimes: vec![], + })); + let state = AppState { + controller_status: None, + config, + daemon_info: DaemonInfo { + version: "test".into(), + scope: DaemonScope::User, + config_path: String::new(), + state_directory: String::new(), + inventory_interval: Duration::from_secs(60), + controller: None, + llm_proxy: None, + usage_receiver: None, + }, + discovery, + enrollment: EnrollmentState::new(false), + state_dir: dir.path().to_owned(), + oidc_callback_listen: None, + telemetry: None, + logout: None, + usage: None, + programs: None, + }; + (dir, state) +} + +// --- vscode-copilot-capi: path pairing, /_ping and the WebSocket tunnel. + +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpStream, +}; + +/// `request_source_state` (OIDC identity "identity", `githubOAuth: source: +/// request`) with `llmGateway.proxyUrl` pointed at `upstream`, matching +/// the CAPI route's Passthrough/ProxyUrl shape used by most tests below. +async fn capi_state(upstream: std::net::SocketAddr) -> (tempfile::TempDir, AppState) { + let (dir, mut state) = request_source_state(upstream).await; + state.config.llm_gateway.as_mut().unwrap().proxy_url = + Some(format!("http://{upstream}").parse().unwrap()); + (dir, state) +} + +/// Controller-JWT state (cacheable credential, unlike OIDC) with +/// `llmGateway.proxyUrl` pointed at `upstream`, for the credential-expiry +/// tunnel test, which needs a credential it can seed into the cache with a +/// controlled deadline. +async fn capi_controller_jwt_state( + upstream: std::net::SocketAddr, +) -> (tempfile::TempDir, AppState) { + let dir = tempfile::tempdir().unwrap(); + std::fs::write( + dir.path().join("identity.json"), + serde_json::json!({ + "deviceId": "device-test", + "clientCertificatePem": "", + "clientCertificateExpiresAtUnixSeconds": 4_000_000_000u64, + "oauthTokenEndpoint": "https://controller.example/token", + "oauthClientId": "device", + }) + .to_string(), + ) + .unwrap(); + let config = parse_daemon(&format!( + "llmGateway:\n url: http://{upstream}\n proxyUrl: http://{upstream}\n authentication:\n type: controllerJwt\n audience: agentgateway\n allowedClientIds: [vscode-copilot]\n" + )) + .unwrap(); + let (_, discovery) = watch::channel(Arc::new(Discovery { + agents: vec![], + model_runtimes: vec![], + })); + let state = AppState { + controller_status: None, + config, + daemon_info: DaemonInfo { + version: "test".into(), + scope: DaemonScope::User, + config_path: String::new(), + state_directory: String::new(), + inventory_interval: Duration::from_secs(60), + controller: None, + llm_proxy: None, + usage_receiver: None, + }, + discovery, + enrollment: EnrollmentState::new(false), + state_dir: dir.path().to_owned(), + oidc_callback_listen: None, + telemetry: None, + logout: None, + usage: None, + programs: None, + }; + (dir, state) +} + +/// Whether no connection reached `listener` within a short window: proof +/// that a refused request was never forwarded upstream. +async fn no_connection_arrives(listener: &TcpListener) -> bool { + tokio::time::timeout(Duration::from_millis(200), listener.accept()) + .await + .is_err() +} + +/// Reads a raw HTTP/1.1 request head (request line + headers, byte at a +/// time up to the blank line) from a hand-rolled test upstream: good +/// enough for a WebSocket upgrade request, which has no body. Header +/// names are folded to lowercase for case-insensitive lookup. +async fn read_http_head( + stream: &mut TcpStream, +) -> (String, String, std::collections::HashMap) { + let mut buffer = Vec::new(); + let mut byte = [0u8; 1]; + while !buffer.ends_with(b"\r\n\r\n") { + stream.read_exact(&mut byte).await.unwrap(); + buffer.push(byte[0]); + } + let text = String::from_utf8(buffer).unwrap(); + let mut lines = text.split("\r\n"); + let mut parts = lines.next().unwrap().split_whitespace(); + let method = parts.next().unwrap().to_owned(); + let path = parts.next().unwrap().to_owned(); + let mut headers = std::collections::HashMap::new(); + for line in lines { + if let Some((name, value)) = line.split_once(':') { + headers.insert(name.trim().to_ascii_lowercase(), value.trim().to_owned()); + } + } + (method, path, headers) +} + +/// RFC 6455's fixed GUID: concatenated with the client's +/// `Sec-WebSocket-Key` and SHA-1'd (then base64'd) to produce +/// `Sec-WebSocket-Accept`. +const WS_GUID: &str = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"; + +fn ws_accept(key: &str) -> String { + use base64::engine::general_purpose::STANDARD; + use sha1::{Digest, Sha1}; + STANDARD.encode(Sha1::digest(format!("{key}{WS_GUID}").as_bytes())) +} + +/// Encodes a single, unfragmented text frame (opcode `0x1`). Client +/// frames must be masked (`mask: true`, a fixed key: the test payloads +/// are not attacker-controlled), server frames must not be. +fn ws_frame(payload: &[u8], mask: bool) -> Vec { + let mut frame = vec![0x81u8]; + let mask_bit = if mask { 0x80 } else { 0x00 }; + let len = payload.len(); + if len < 126 { + frame.push(mask_bit | len as u8); + } else { + frame.push(mask_bit | 126); + frame.extend_from_slice(&(len as u16).to_be_bytes()); + } + if mask { + let key = [0x11u8, 0x22, 0x33, 0x44]; + frame.extend_from_slice(&key); + frame.extend(payload.iter().enumerate().map(|(i, b)| b ^ key[i % 4])); + } else { + frame.extend_from_slice(payload); + } + frame +} + +/// Reads exactly one, unfragmented, non-huge frame and returns its opcode +/// and unmasked payload. +async fn ws_read_frame(stream: &mut TcpStream) -> (u8, Vec) { + let mut header = [0u8; 2]; + stream.read_exact(&mut header).await.unwrap(); + let opcode = header[0] & 0x0f; + let masked = header[1] & 0x80 != 0; + let mut len = (header[1] & 0x7f) as usize; + if len == 126 { + let mut extended = [0u8; 2]; + stream.read_exact(&mut extended).await.unwrap(); + len = u16::from_be_bytes(extended) as usize; + } + let mask_key = if masked { + let mut key = [0u8; 4]; + stream.read_exact(&mut key).await.unwrap(); + Some(key) + } else { + None + }; + let mut payload = vec![0u8; len]; + stream.read_exact(&mut payload).await.unwrap(); + if let Some(key) = mask_key { + for (index, byte) in payload.iter_mut().enumerate() { + *byte ^= key[index % 4]; + } + } + (opcode, payload) +} + +/// A fixed, valid `Sec-WebSocket-Key` (the RFC 6455 example key): its +/// `Sec-WebSocket-Accept` is computed by `ws_accept`, never hand-checked +/// against a literal. +const SEC_WEBSOCKET_KEY: &str = "dGhlIHNhbXBsZSBub25jZQ=="; + +fn upgrade_request(uri: String) -> Request> { + Request::builder() + .method(Method::GET) + .uri(uri) + .header(CONNECTION, "Upgrade") + .header("upgrade", "websocket") + .header("sec-websocket-version", "13") + .header("sec-websocket-key", SEC_WEBSOCKET_KEY) + .body(Full::new(Bytes::new())) + .unwrap() +} + +#[tokio::test] +async fn capi_route_accepts_and_strips_path_pairing_never_forwarding_it() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let upstream_task = tokio::spawn(async move { + let (socket, _) = upstream.accept().await.unwrap(); + hyper::server::conn::http1::Builder::new() + .serve_connection( + TokioIo::new(socket), + service_fn(move |request: Request| async move { + // The pairing segment must never reach the upstream path. + assert_eq!(request.uri(), "/chat/completions"); + assert!(!request.headers().contains_key(PAIRING_HEADER)); + Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static(b"ok")))) + }), + ) + .await + .unwrap(); + }); + let (dir, state) = capi_state(address).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let request = Request::builder() + .method(Method::POST) + .uri(format!( + "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/chat/completions" + )) + .header(AUTHORIZATION, "Bearer tid=from-client") + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + let response = client.request(request).await.unwrap(); + assert_eq!( + response.status(), + StatusCode::OK, + "path pairing must be accepted on {CAPI_ROUTE}" + ); + proxy.abort(); + let _ = proxy.await; + drop(client); + upstream_task.abort(); + drop(dir); + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn capi_route_rejects_missing_empty_or_wrong_pairing_segment_even_with_a_valid_header() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + // Dropped rather than accepted from: today none of these four + // paths are rejected before forwarding (path-pairing extraction + // does not exist yet), so leaving the listener up would have + // every case hang until the gateway timeout instead of failing + // fast with a clean (if wrong) status. + drop(upstream); + let (dir, state) = capi_state(address).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + for path in [ + CAPI_ROUTE.to_owned(), // no segment at all + format!("{CAPI_ROUTE}/"), // no segment at all + format!("{CAPI_ROUTE}//chat/completions"), // empty segment + format!("{CAPI_ROUTE}/wrong-pairing/chat/completions"), // wrong segment + ] { + let request = Request::builder() + .method(Method::POST) + .uri(format!("http://{proxy_address}{path}")) + // A valid header must not substitute for the path segment + // on this route: the header is ignored here. + .header(PAIRING_HEADER, PAIRING) + .header(AUTHORIZATION, "Bearer tid=from-client") + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + assert_eq!( + error_code(client.request(request).await.unwrap()).await, + (StatusCode::FORBIDDEN, "pairing_invalid".to_owned()), + "{path}" + ); + } + proxy.abort(); + let _ = proxy.await; + drop(client); + drop(dir); + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn capi_route_rejects_a_double_dot_segment_after_the_pairing() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let (dir, state) = capi_state(address).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let request = Request::builder() + .method(Method::POST) + .uri(format!( + "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/../admin" + )) + .header(AUTHORIZATION, "Bearer tid=from-client") + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + assert_eq!( + error_code(client.request(request).await.unwrap()).await, + (StatusCode::BAD_REQUEST, "path_invalid".to_owned()) + ); + proxy.abort(); + let _ = proxy.await; + drop(client); + drop(dir); + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn capi_route_moves_the_client_token_to_x_llm_token() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let upstream_task = tokio::spawn(async move { + let (socket, _) = upstream.accept().await.unwrap(); + hyper::server::conn::http1::Builder::new() + .serve_connection( + TokioIo::new(socket), + service_fn(move |request: Request| async move { + assert_eq!(request.uri(), "/chat/completions"); + assert_eq!(request.headers()["x-llm-token"], "tid=from-client"); + assert_eq!(request.headers()[AUTHORIZATION], "Bearer identity"); + Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static(b"ok")))) + }), + ) + .await + .unwrap(); + }); + let (dir, state) = capi_state(address).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let request = Request::builder() + .method(Method::POST) + .uri(format!( + "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/chat/completions" + )) + .header(AUTHORIZATION, "Bearer tid=from-client") + .body(Full::new(Bytes::from_static(b"{}"))) + .unwrap(); + let response = client.request(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + proxy.abort(); + let _ = proxy.await; + drop(client); + upstream_task.abort(); + drop(dir); + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn capi_route_ping_without_credential_is_forwarded_without_x_llm_token() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let upstream_task = tokio::spawn(async move { + let (socket, _) = upstream.accept().await.unwrap(); + hyper::server::conn::http1::Builder::new() + .serve_connection( + TokioIo::new(socket), + service_fn(move |request: Request| async move { + assert_eq!(request.uri(), "/_ping"); + assert!(!request.headers().contains_key("x-llm-token")); + Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static(b"ok")))) + }), + ) + .await + .unwrap(); + }); + let (dir, state) = capi_state(address).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let request = Request::builder() + .method(Method::GET) + .uri(format!( + "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/_ping" + )) + .body(Full::new(Bytes::new())) + .unwrap(); + let response = client.request(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + proxy.abort(); + let _ = proxy.await; + drop(client); + upstream_task.abort(); + drop(dir); + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn capi_route_ping_post_without_credential_is_401() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let (dir, state) = capi_state(address).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let request = Request::builder() + .method(Method::POST) + .uri(format!( + "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/_ping" + )) + .body(Full::new(Bytes::new())) + .unwrap(); + assert_eq!( + error_code(client.request(request).await.unwrap()).await, + ( + StatusCode::UNAUTHORIZED, + "client_credential_missing".to_owned() + ) + ); + proxy.abort(); + let _ = proxy.await; + drop(client); + drop(dir); + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn capi_route_ping_without_pairing_is_403() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let (dir, state) = capi_state(address).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let request = Request::builder() + .method(Method::GET) + .uri(format!("http://{proxy_address}{CAPI_ROUTE}/wrong/_ping")) + .body(Full::new(Bytes::new())) + .unwrap(); + assert_eq!( + error_code(client.request(request).await.unwrap()).await, + (StatusCode::FORBIDDEN, "pairing_invalid".to_owned()) + ); + proxy.abort(); + let _ = proxy.await; + drop(client); + drop(dir); + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn capi_route_upgrade_without_pairing_is_403_and_not_forwarded() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let (dir, state) = capi_state(address).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let request = upgrade_request(format!( + "http://{proxy_address}{CAPI_ROUTE}/wrong/responses" + )); + assert_eq!( + error_code(client.request(request).await.unwrap()).await, + (StatusCode::FORBIDDEN, "pairing_invalid".to_owned()) + ); + assert!( + no_connection_arrives(&upstream).await, + "an upgrade refused for its pairing must not be forwarded upstream" + ); + proxy.abort(); + let _ = proxy.await; + drop(client); + drop(dir); + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn upgrade_on_another_route_is_400_but_other_upgrade_tokens_still_forward() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let upstream_task = tokio::spawn(async move { + loop { + let (socket, _) = match upstream.accept().await { + Ok(accepted) => accepted, + Err(_) => break, + }; + tokio::spawn(async move { + let _ = hyper::server::conn::http1::Builder::new() + .serve_connection( + TokioIo::new(socket), + service_fn(|_request: Request| async move { + Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static( + b"ok", + )))) + }), + ) + .await; + }); + } + }); + let (dir, state) = request_source_state(address).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + // Only /vscode-copilot-capi tunnels an upgrade; on any other + // route Upgrade: websocket is refused outright. + let websocket_upgrade = Request::builder() + .method(Method::GET) + .uri(format!("http://{proxy_address}/copilot-cli/v1/models")) + .header(PAIRING_HEADER, PAIRING) + .header(AUTHORIZATION, "Bearer tid=from-client") + .header(CONNECTION, "upgrade") + .header("upgrade", "websocket") + .body(Full::new(Bytes::new())) + .unwrap(); + // Checked as a plain status first (not via the shared + // `error_code` helper, which assumes a JSON error body): today + // this request is not refused at all and is forwarded as an + // ordinary request, so the response has no content-type header + // for `error_code` to read. + let response = client.request(websocket_upgrade).await.unwrap(); + assert_eq!( + response.status(), + StatusCode::BAD_REQUEST, + "an Upgrade: websocket request on a non-CAPI route must be refused, not forwarded" + ); + let body = response.into_body().collect().await.unwrap().to_bytes(); + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(json["error"]["code"], "upgrade_not_supported"); + // Only the websocket token is refused; another Upgrade token on + // the same route is stripped (as a hop-by-hop header) and + // forwarded as an ordinary request, as before this feature. + let other_upgrade = Request::builder() + .method(Method::GET) + .uri(format!("http://{proxy_address}/copilot-cli/v1/models")) + .header(PAIRING_HEADER, PAIRING) + .header(AUTHORIZATION, "Bearer tid=from-client") + .header(CONNECTION, "upgrade") + .header("upgrade", "h2c") + .body(Full::new(Bytes::new())) + .unwrap(); + let response = client.request(other_upgrade).await.unwrap(); + assert_eq!( + response.status(), + StatusCode::OK, + "a non-websocket Upgrade token must still be forwarded" + ); + proxy.abort(); + let _ = proxy.await; + drop(client); + upstream_task.abort(); + drop(dir); + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn capi_route_non_101_upstream_response_passes_through() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let upstream_task = tokio::spawn(async move { + let (socket, _) = upstream.accept().await.unwrap(); + hyper::server::conn::http1::Builder::new() + .serve_connection( + TokioIo::new(socket), + service_fn(|_request: Request| async move { + Ok::<_, Infallible>( + Response::builder() + .status(StatusCode::BAD_GATEWAY) + .body(Full::new(Bytes::from_static(b"nope"))) + .unwrap(), + ) + }), + ) + .await + .unwrap(); + }); + let (dir, state) = capi_state(address).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let mut request = upgrade_request(format!( + "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/responses" + )); + request + .headers_mut() + .insert(AUTHORIZATION, "Bearer tid=from-client".parse().unwrap()); + let response = client.request(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + let body = response.into_body().collect().await.unwrap().to_bytes(); + assert_eq!(body, Bytes::from_static(b"nope")); + proxy.abort(); + let _ = proxy.await; + drop(client); + upstream_task.abort(); + drop(dir); + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn capi_route_websocket_tunnel_relays_the_101_and_one_frame_each_way_with_gateway_credentials() + { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let upstream_task = tokio::spawn(async move { + let (mut socket, _) = upstream.accept().await.unwrap(); + let (method, path, headers) = read_http_head(&mut socket).await; + assert_eq!(method, "GET"); + assert_eq!(path, "/responses"); + assert_eq!( + headers.get("authorization").map(String::as_str), + Some("Bearer identity"), + "the gateway identity must reach the upstream on the handshake" + ); + assert_eq!( + headers.get("x-llm-token").map(String::as_str), + Some("tid=from-client"), + "the client's own token must reach the upstream on the handshake" + ); + let accept = + ws_accept(headers.get("sec-websocket-key").map(String::as_str).unwrap_or("")); + socket + .write_all( + format!( + "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {accept}\r\n\r\n" + ) + .as_bytes(), + ) + .await + .unwrap(); + let (opcode, payload) = ws_read_frame(&mut socket).await; + assert_eq!(opcode, 0x1); + assert_eq!(payload, b"hello from vs code"); + socket + .write_all(&ws_frame(b"hello from github", false)) + .await + .unwrap(); + }); + let (dir, state) = capi_state(address).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let mut request = upgrade_request(format!( + "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/responses" + )); + request + .headers_mut() + .insert(AUTHORIZATION, "Bearer tid=from-client".parse().unwrap()); + let response = client.request(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS); + assert_eq!(response.headers()["upgrade"], "websocket"); + assert!( + response.headers()["connection"] + .to_str() + .unwrap() + .eq_ignore_ascii_case("upgrade") + ); + assert_eq!( + response.headers()["sec-websocket-accept"], + ws_accept(SEC_WEBSOCKET_KEY) + ); + let upgraded = hyper::upgrade::on(response).await.unwrap(); + let mut io = TokioIo::new(upgraded); + io.write_all(&ws_frame(b"hello from vs code", true)) + .await + .unwrap(); + let mut header = [0u8; 2]; + io.read_exact(&mut header).await.unwrap(); + assert_eq!(header[0] & 0x0f, 0x1, "expected a text frame back"); + let len = (header[1] & 0x7f) as usize; + assert!(header[1] & 0x80 == 0, "a server frame must not be masked"); + let mut payload = vec![0u8; len]; + io.read_exact(&mut payload).await.unwrap(); + assert_eq!(payload, b"hello from github"); + proxy.abort(); + let _ = proxy.await; + drop(client); + upstream_task.abort(); + drop(dir); + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn dropping_the_proxy_closes_an_open_capi_tunnel() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let upstream_task = tokio::spawn(async move { + let (mut socket, _) = upstream.accept().await.unwrap(); + let (_, _, headers) = read_http_head(&mut socket).await; + let accept = + ws_accept(headers.get("sec-websocket-key").map(String::as_str).unwrap_or("")); + socket + .write_all( + format!( + "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {accept}\r\n\r\n" + ) + .as_bytes(), + ) + .await + .unwrap(); + // Held open until the client side sees it close. + let mut buffer = [0u8; 1]; + let _ = socket.read(&mut buffer).await; + }); + let (dir, state) = capi_state(address).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve(listener, state, default_config())); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let mut request = upgrade_request(format!( + "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/responses" + )); + request + .headers_mut() + .insert(AUTHORIZATION, "Bearer tid=from-client".parse().unwrap()); + let response = client.request(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS); + let upgraded = hyper::upgrade::on(response).await.unwrap(); + let mut io = TokioIo::new(upgraded); + proxy.abort(); + let _ = proxy.await; + // The tunnel task is owned by the connection's JoinSet: dropping + // the server must close it, and the client sees EOF. + let mut buffer = [0u8; 1]; + let read = tokio::time::timeout(Duration::from_secs(5), io.read(&mut buffer)) + .await + .expect("dropping the proxy must close the open tunnel promptly") + .unwrap(); + assert_eq!(read, 0, "expected EOF once the tunnel closes"); + upstream_task.abort(); + drop(client); + drop(dir); + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn capi_tunnel_closes_when_the_cached_credential_expires() { + tokio::time::timeout(Duration::from_secs(15), async { + let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = upstream.local_addr().unwrap(); + let upstream_task = tokio::spawn(async move { + let (mut socket, _) = upstream.accept().await.unwrap(); + let (_, _, headers) = read_http_head(&mut socket).await; + let accept = + ws_accept(headers.get("sec-websocket-key").map(String::as_str).unwrap_or("")); + socket + .write_all( + format!( + "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {accept}\r\n\r\n" + ) + .as_bytes(), + ) + .await + .unwrap(); + let mut buffer = [0u8; 1]; + let _ = socket.read(&mut buffer).await; + }); + let (dir, state) = capi_controller_jwt_state(address).await; + let cache = Arc::new(CredentialCache::default()); + let key = format!("vscode-copilot\u{0}http://{address}/\u{0}agentgateway"); + cache.insert( + &key, + "device-test", + &agentdesktop_core::model::LlmGatewayCredential { + credential: "short-lived".to_owned(), + expires_at_unix_seconds: now_unix() + 600, + }, + ); + cache.expire_at_for_test(&key, Deadline { + monotonic: std::time::Instant::now() + Duration::from_millis(300), + unix: now_unix() + 60, + }); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let proxy = tokio::spawn(serve_with_cache( + listener, + state, + default_config(), + cache, + )); + let client: Client> = + Client::builder(TokioExecutor::new()).build_http(); + let mut request = upgrade_request(format!( + "http://{proxy_address}{CAPI_ROUTE}/{PAIRING}/responses" + )); + request + .headers_mut() + .insert(AUTHORIZATION, "Bearer tid=from-client".parse().unwrap()); + let response = client.request(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS); + let upgraded = hyper::upgrade::on(response).await.unwrap(); + let mut io = TokioIo::new(upgraded); + // Once the cached credential's deadline passes, the tunnel closes + // on its own, with the proxy still running. + let mut buffer = [0u8; 1]; + let read = tokio::time::timeout(Duration::from_secs(5), io.read(&mut buffer)) + .await + .expect("the tunnel must close once its credential expires") + .unwrap(); + assert_eq!(read, 0, "expected EOF once the tunnel closes"); + proxy.abort(); + let _ = proxy.await; + upstream_task.abort(); + drop(client); + drop(dir); + }) + .await + .unwrap(); +} diff --git a/crates/agent/src/provider/codex/integration_tests.rs b/crates/agent/src/provider/codex/integration_tests.rs index a1b67239..a93769c6 100644 --- a/crates/agent/src/provider/codex/integration_tests.rs +++ b/crates/agent/src/provider/codex/integration_tests.rs @@ -1,7 +1,9 @@ use std::time::{Duration, Instant}; use anyhow::{Context, ensure}; +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; use tracing::info; use crate::common::{Container, Gateway, TOKEN}; @@ -13,6 +15,109 @@ const COMPANY_SETTINGS: &str = "/etc/codex/config.toml"; const VERSION: &str = "0.153.4"; const PROMPT: &str = "Return the provider integration response."; +#[tokio::test] +async fn managed_chatgpt_subscription_lifecycle() -> anyhow::Result<()> { + Container::run( + "codex", + "crates/agent/src/provider/codex/testdata/Dockerfile", + async |container| { + let gateway = Gateway::start().await?; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let proxy_address = listener.local_addr()?; + drop(listener); + let state_dir = "/home/tester/.local/state/codex-subscription"; + let socket = format!("{state_dir}/agentdesktop.sock"); + let issuer = "https://codex-fixture.example/"; + let client_id = "codex-fixture"; + let account = URL_SAFE_NO_PAD.encode(Sha256::digest(format!("{issuer}\0{client_id}"))); + let entry = Sha256::digest(format!("dev.agentdesktop.gateway-oidc\0{account}")); + let entry = entry.iter().map(|byte| format!("{byte:02x}")).collect::(); + container.write( + &format!("{state_dir}/secrets/{entry}"), + &json!({ + "accessToken": TOKEN, + "refreshToken": null, + "expiresAtUnixSeconds": 4_000_000_000u64, + "tokenEndpoint": "https://codex-fixture.example/token", + }).to_string(), + ).await?; + let claims = URL_SAFE_NO_PAD.encode(json!({ + "email": "fixture@example.com", + "exp": 4_000_000_000u64, + "https://api.openai.com/auth": { + "chatgpt_account_id": "fixture-account", + "chatgpt_user_id": "fixture-user", + "chatgpt_plan_type": "plus", + }, + }).to_string()); + container.write( + "/home/tester/.codex/auth.json", + &json!({ + "auth_mode": "chatgpt", + "tokens": { + "id_token": format!("e30.{claims}.fixture"), + "access_token": "codex-subscription-fixture", + "refresh_token": "fixture-refresh-token", + "account_id": "fixture-account", + }, + "last_refresh": "2090-01-01T00:00:00Z", + }).to_string(), + ).await?; + container.exec(&["chown", "-R", "tester:tester", state_dir, "/home/tester/.codex"]).await?; + container.exec(&["chmod", "700", state_dir, &format!("{state_dir}/secrets")]).await?; + container.write(CONFIG, &json!({ + "daemon": { + "user": true, + "stateDir": state_dir, + "llmProxy": { "listen": proxy_address.to_string() }, + }, + "llmGateway": { + "url": gateway.url, + "proxyUrl": format!("{}backend-api/codex", gateway.url), + "authentication": { "type": "oidc", "issuer": issuer, "clientId": client_id }, + }, + "programs": { "codex": { + "auth": "subscription", + "managedConfig": { + "model": "gpt-5.4", + "cli_auth_credentials_store": "file", + "model_providers": { "agentdesktop": { "request_max_retries": 0, "stream_max_retries": 0 } }, + }, + } }, + }).to_string()).await?; + container.start_process("subscription-daemon", &[ + "runuser", "--preserve-environment", "-u", "tester", "--", + "agentdesktop", "daemon", "--user", "--config", CONFIG, + ]).await?; + container.wait_ready(&["agentdesktop", "--socket", &socket, "status"]).await?; + let managed = container.read(USER_SETTINGS).await?; + let parsed: toml::Value = toml::from_str(&managed)?; + let provider = &parsed["model_providers"]["agentdesktop"]; + ensure!(provider["requires_openai_auth"].as_bool() == Some(true), "ChatGPT authentication missing"); + ensure!(provider.get("auth").is_none(), "gateway credential command conflicts with ChatGPT auth"); + let output = container.exec_as("tester", &[ + "env", "HTTP_PROXY=http://127.0.0.1:9", "HTTPS_PROXY=http://127.0.0.1:9", "NO_PROXY=127.0.0.1,localhost", + "codex", "exec", "--skip-git-repo-check", "--ephemeral", "--json", PROMPT, + ], Duration::from_secs(90)).await?; + ensure!(output.contains("provider-integration-ok"), "native Codex subscription request failed: {output}"); + ensure!(gateway.requests().iter().any(|request| { + request["path"] == "/backend-api/codex/responses" + && request["authorization"] == format!("Bearer {TOKEN}") + && request["llmToken"] == "codex-subscription-fixture" + && request["chatgptAccount"] == "fixture-account" + && request["pairing"].is_null() + && request["body"]["input"].to_string().contains(PROMPT) + }), "Codex subscription headers did not reach the mock gateway correctly"); + container.stop_process("subscription-daemon").await?; + container.write(CONFIG, &json!({"daemon": {"user": true, "stateDir": state_dir}, "programs": {}}).to_string()).await?; + container.exec_as("tester", &["agentdesktop", "daemon", "--user", "--config", CONFIG, "--once"], Duration::from_secs(15)).await?; + container.exec(&["test", "!", "-e", USER_SETTINGS]).await?; + container.exec(&["test", "-f", "/home/tester/.codex/auth.json"]).await?; + Ok(()) + }, + ).await +} + #[tokio::test] async fn managed_gateway_lifecycle() -> anyhow::Result<()> { Container::run( diff --git a/crates/agent/src/provider/codex/mod.rs b/crates/agent/src/provider/codex/mod.rs index 0fceb0c6..4e7c7529 100644 --- a/crates/agent/src/provider/codex/mod.rs +++ b/crates/agent/src/provider/codex/mod.rs @@ -1,6 +1,6 @@ use std::path::PathBuf; -use agentdesktop_core::config::DaemonConfig; +use agentdesktop_core::config::{DaemonConfig, ProgramAuthentication}; use agentdesktop_core::model::Discovery; use super::{Provider, ReconcileContext}; @@ -41,6 +41,17 @@ impl Provider for Codex { } fn plan(&self, ctx: &ReconcileContext, config: &DaemonConfig) -> anyhow::Result { + if !ctx.merge_user_settings + && config + .programs + .codex + .as_ref() + .is_some_and(|program| program.auth == Some(ProgramAuthentication::Subscription)) + { + anyhow::bail!( + "programs.codex.auth subscription requires --user (or daemon.user: true)" + ); + } let configured = config.programs.codex.as_ref().map(|provider| { let gateway = config .llm_gateway @@ -54,6 +65,9 @@ impl Provider for Codex { &ctx.credential_helper, &ctx.socket, config.sandbox.as_ref(), + ctx.llm_proxy + .as_ref() + .map(|proxy| (proxy.address, &*proxy.pairing)), configured, &plan, )?; @@ -65,3 +79,24 @@ impl Provider for Codex { pub fn default_codex_managed_config_path() -> PathBuf { PathBuf::from("/etc/codex/managed_config.toml") } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn subscription_requires_a_user_mode_daemon() { + let config = agentdesktop_core::config::parse_daemon( + "llmGateway:\n url: https://gateway.example.com\n proxyUrl: https://gateway.example.com/codex\n authentication:\n type: controllerJwt\n audience: gateway\n allowedClientIds: [codex]\nprograms:\n codex:\n auth: subscription\n", + ) + .unwrap(); + let context = ReconcileContext { + merge_user_settings: false, + credential_helper: PathBuf::from("agentdesktop"), + socket: PathBuf::from("socket"), + llm_proxy: None, + }; + let error = Codex::default().plan(&context, &config).err().unwrap(); + assert!(error.to_string().contains("requires --user")); + } +} diff --git a/crates/agent/src/provider/codex/reconcile.rs b/crates/agent/src/provider/codex/reconcile.rs index 5f2c6956..150f4a42 100644 --- a/crates/agent/src/provider/codex/reconcile.rs +++ b/crates/agent/src/provider/codex/reconcile.rs @@ -1,9 +1,9 @@ use super::Codex; -use std::path::Path; +use std::{net::SocketAddr, path::Path}; use agentdesktop_core::config::{ - CodexConfig, LlmGatewayAuthentication, LlmGatewayConfig, SandboxConfig, + CodexConfig, LlmGatewayAuthentication, LlmGatewayConfig, ProgramAuthentication, SandboxConfig, }; use anyhow::Context; use serde_json::{Value, json}; @@ -23,6 +23,7 @@ pub(super) fn plan( credential_helper: &Path, socket: &Path, sandbox: Option<&SandboxConfig>, + proxy: Option<(SocketAddr, &str)>, config: Option<(&CodexConfig, Option<&LlmGatewayConfig>)>, plan: &ReconcilePlan, ) -> anyhow::Result<()> { @@ -30,7 +31,25 @@ pub(super) fn plan( return remove(path, plan); }; - let settings = managed_config(config, gateway, credential_helper, socket, sandbox)?; + let subscription = config.auth == Some(ProgramAuthentication::Subscription); + if subscription && proxy.is_none() { + let gateway = gateway.context("Codex subscription requires an LLM gateway")?; + if crate::reconcile::fail_closed_without_proxy(gateway, plan, || { + plan.read(path).is_ok_and(|contents| { + contents.starts_with(MANAGED_HEADER.as_bytes()) && carries_pairing(&contents) + }) + }) { + return Ok(()); + } + plan.inactive(crate::reconcile::FAIL_OPEN_REASON); + return remove(path, plan); + } + let mode = if subscription { + 0o600 + } else { + MANAGED_FILE_MODE + }; + let settings = managed_config(config, gateway, credential_helper, socket, sandbox, proxy)?; let mut contents = MANAGED_HEADER.as_bytes().to_vec(); contents.extend_from_slice( toml::to_string_pretty(&settings) @@ -52,10 +71,7 @@ pub(super) fn plan( }; let action = match existing.as_deref() { // Identical bytes at a mode looser than planned: rewritten (update). - Some(existing) - if existing == contents - && !crate::reconcile::grants_beyond(path, MANAGED_FILE_MODE) => - { + Some(existing) if existing == contents && !crate::reconcile::grants_beyond(path, mode) => { debug!( program = Codex::ID, action = "unchanged", @@ -67,43 +83,64 @@ pub(super) fn plan( } Some(existing) if existing.starts_with(MANAGED_HEADER.as_bytes()) => "update", Some(existing) => { - plan.record_diff( - Codex::DISPLAY_NAME, - "configuration", - "conflict", - path, - Some(existing), - Some(&contents), - ); + record_configuration(plan, path, "conflict", Some(existing), Some(&contents)); return Ok(()); } None => "create", }; - plan.write_file(path, &contents, MANAGED_FILE_MODE)?; + plan.write_file(path, &contents, mode)?; debug!( program = Codex::ID, action, path = %path.display(), "planned managed configuration" ); - plan.record_diff( - Codex::DISPLAY_NAME, - "configuration", - action, - path, - existing.as_deref(), - Some(&contents), - ); + record_configuration(plan, path, action, existing.as_deref(), Some(&contents)); Ok(()) } +fn carries_pairing(contents: &[u8]) -> bool { + contents + .windows(crate::llm_proxy::PAIRING_HEADER.len()) + .any(|part| part == crate::llm_proxy::PAIRING_HEADER.as_bytes()) +} + +fn record_configuration( + plan: &ReconcilePlan, + path: &Path, + action: &str, + before: Option<&[u8]>, + after: Option<&[u8]>, +) { + if before.is_some_and(carries_pairing) || after.is_some_and(carries_pairing) { + plan.record_diff( + Codex::DISPLAY_NAME, + "configuration", + action, + path, + None, + None, + ); + } else { + plan.record_diff( + Codex::DISPLAY_NAME, + "configuration", + action, + path, + before, + after, + ); + } +} + fn managed_config( config: &CodexConfig, gateway: Option<&LlmGatewayConfig>, credential_helper: &Path, socket: &Path, sandbox: Option<&SandboxConfig>, + proxy: Option<(SocketAddr, &str)>, ) -> anyhow::Result { let mut settings = serde_json::to_value(&config.managed_config) .context("serialize Codex pass-through managed configuration")?; @@ -152,7 +189,31 @@ fn managed_config( "base_url": responses_base_url(gateway), "wire_api": "responses", }); - if gateway + if config.auth == Some(ProgramAuthentication::Subscription) { + let (listen, pairing) = proxy.context("Codex subscription requires the local LLM proxy")?; + provider["base_url"] = json!(format!("http://{listen}{}", crate::llm_proxy::CODEX_ROUTE)); + provider["requires_openai_auth"] = json!(true); + provider["supports_websockets"] = json!(false); + provider["http_headers"] = json!({ crate::llm_proxy::PAIRING_HEADER: pairing }); + settings["forced_login_method"] = json!("chatgpt"); + if let Some(existing) = settings + .get_mut("model_providers") + .and_then(|providers| providers.get_mut(provider_name)) + .and_then(Value::as_object_mut) + { + for field in [ + "auth", + "env_key", + "experimental_bearer_token", + "aws", + "gateway_oauth", + "env_http_headers", + "http_headers", + ] { + existing.remove(field); + } + } + } else if gateway .authentication .as_ref() .is_some_and(LlmGatewayAuthentication::uses_credential_helper) @@ -200,7 +261,7 @@ fn remove(path: &Path, plan: &ReconcilePlan) -> anyhow::Result<()> { path = %path.display(), "planned managed configuration" ); - plan.record(Codex::DISPLAY_NAME, "configuration", "remove", path); + record_configuration(plan, path, "remove", Some(&contents), None); Ok(()) } Ok(_) => { @@ -236,6 +297,186 @@ mod tests { use super::managed_config; + fn subscription_config() -> agentdesktop_core::config::DaemonConfig { + parse_daemon( + r#" +llmGateway: + url: https://gateway.example.com/llm + proxyUrl: https://gateway.example.com/chatgpt + authentication: + type: controllerJwt + audience: agentgateway + allowedClientIds: [codex] +programs: + codex: + auth: subscription + managedConfig: + model: gpt-5.4 + model_providers: + agentdesktop: + auth: {command: conflicting-helper} + env_key: CONFLICTING_KEY + experimental_bearer_token: conflicting-token + http_headers: {Authorization: "Bearer conflicting-token"} + env_http_headers: {Authorization: CONFLICTING_KEY} +"#, + ) + .unwrap() + } + + #[test] + fn subscription_provider_keeps_chatgpt_auth_and_pairs_with_the_proxy() { + let config = subscription_config(); + let settings = managed_config( + config.programs.codex.as_ref().unwrap(), + config.llm_gateway.as_ref(), + Path::new("agentdesktop"), + Path::new("agentdesktop.sock"), + None, + Some(("127.0.0.1:4010".parse().unwrap(), "pairing-secret")), + ) + .unwrap(); + assert_eq!(settings["model"], "gpt-5.4"); + assert_eq!(settings["forced_login_method"], "chatgpt"); + let provider = &settings["model_providers"]["agentdesktop"]; + assert_eq!(provider["base_url"], "http://127.0.0.1:4010/codex"); + assert_eq!(provider["requires_openai_auth"], true); + assert_eq!(provider["supports_websockets"], false); + assert_eq!(provider["wire_api"], "responses"); + assert_eq!( + provider["http_headers"], + serde_json::json!({ + crate::llm_proxy::PAIRING_HEADER: "pairing-secret" + }) + ); + for field in [ + "auth", + "env_key", + "experimental_bearer_token", + "env_http_headers", + ] { + assert!(provider.get(field).is_none(), "unexpected {field}"); + } + let contents = toml::to_string_pretty(&settings).unwrap(); + assert!(!contents.contains("conflicting")); + let parsed: toml::Value = toml::from_str(&contents).unwrap(); + assert_eq!(parsed["model_provider"].as_str(), Some("agentdesktop")); + } + + #[test] + fn subscription_lifecycle_redacts_pairing_and_obeys_proxy_unavailable_policy() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("config.toml"); + let mut config = subscription_config(); + let run_plan = |config: &agentdesktop_core::config::DaemonConfig, proxy| { + let plan = crate::reconcile::ReconcilePlan::default(); + super::plan( + &path, + Path::new("agentdesktop"), + Path::new("agentdesktop.sock"), + None, + proxy, + Some(( + config.programs.codex.as_ref().unwrap(), + config.llm_gateway.as_ref(), + )), + &plan, + ) + .unwrap(); + plan + }; + let absent = run_plan(&config, None); + assert_eq!( + absent.inactive_reason().as_deref(), + Some(crate::reconcile::FAIL_CLOSED_NOT_IN_PLACE_REASON) + ); + absent.apply().unwrap(); + assert!(!path.exists()); + for pairing in ["first-pairing-secret", "second-pairing-secret"] { + let plan = run_plan(&config, Some(("127.0.0.1:4010".parse().unwrap(), pairing))); + assert!(!plan.render().contains("pairing-secret")); + plan.apply().unwrap(); + assert!(std::fs::read_to_string(&path).unwrap().contains(pairing)); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + assert_eq!( + std::fs::metadata(&path).unwrap().permissions().mode() & 0o777, + 0o600 + ); + std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o644)).unwrap(); + } + } + let contents = std::fs::read(&path).unwrap(); + let closed = run_plan(&config, None); + assert_eq!( + closed.inactive_reason().as_deref(), + Some(crate::reconcile::FAIL_CLOSED_REASON) + ); + closed.apply().unwrap(); + assert_eq!(std::fs::read(&path).unwrap(), contents); + config.llm_gateway.as_mut().unwrap().when_proxy_unavailable = + agentdesktop_core::config::ProxyUnavailable::FailOpen; + let open = run_plan(&config, None); + assert_eq!( + open.inactive_reason().as_deref(), + Some(crate::reconcile::FAIL_OPEN_REASON) + ); + assert!(!open.render().contains("pairing-secret")); + open.apply().unwrap(); + assert!(!path.exists()); + } + + #[test] + fn subscription_configuration_is_redacted_on_conflict_and_when_disabled() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("config.toml"); + let config = subscription_config(); + std::fs::write(&path, "model = 'personal'\n").unwrap(); + let plan = crate::reconcile::ReconcilePlan::default(); + super::plan( + &path, + Path::new("agentdesktop"), + Path::new("agentdesktop.sock"), + None, + Some(("127.0.0.1:4010".parse().unwrap(), "pairing-secret")), + Some(( + config.programs.codex.as_ref().unwrap(), + config.llm_gateway.as_ref(), + )), + &plan, + ) + .unwrap(); + assert!(plan.has_conflicts()); + assert!(!plan.render().contains("pairing-secret")); + assert_eq!( + std::fs::read_to_string(&path).unwrap(), + "model = 'personal'\n" + ); + std::fs::write( + &path, + format!( + "{}x-agentdesktop-pairing = 'pairing-secret'\n", + super::MANAGED_HEADER + ), + ) + .unwrap(); + let plan = crate::reconcile::ReconcilePlan::default(); + super::plan( + &path, + Path::new("agentdesktop"), + Path::new("socket"), + None, + None, + None, + &plan, + ) + .unwrap(); + assert!(!plan.render().contains("pairing-secret")); + plan.apply().unwrap(); + assert!(!path.exists()); + } + #[test] fn pass_through_settings_are_merged_with_managed_gateway_values() { let config = parse_daemon( @@ -268,6 +509,7 @@ programs: Path::new("/usr/local/bin/agentdesktop"), Path::new("/run/agentdesktop/agentdesktop.sock"), None, + None, ) .expect("merged settings"); @@ -314,6 +556,7 @@ programs: Path::new("agentdesktop"), Path::new("agentdesktop.sock"), config.sandbox.as_ref(), + None, ) .expect("sandbox settings"); @@ -365,6 +608,7 @@ programs: Path::new("agentdesktop"), Path::new("agentdesktop.sock"), None, + None, Some((config.programs.codex.as_ref().unwrap(), None)), plan, ) @@ -387,6 +631,7 @@ programs: Path::new("agentdesktop"), Path::new("agentdesktop.sock"), None, + None, Some((config.programs.codex.as_ref().unwrap(), None)), plan, ) diff --git a/crates/agent/tests/common/gateway.rs b/crates/agent/tests/common/gateway.rs index d4d15708..021cdd47 100644 --- a/crates/agent/tests/common/gateway.rs +++ b/crates/agent/tests/common/gateway.rs @@ -74,19 +74,37 @@ async fn handle( let api_key = headers.get("x-api-key").and_then(|v| v.to_str().ok()); let llm_token = headers.get("x-llm-token").and_then(|v| v.to_str().ok()); tracing::debug!(%method, path, "Gateway fixture request"); - requests.lock().unwrap().push(json!({"path": path, "method": method.as_str(), "authorization": authorization, "apiKey": api_key, "llmToken": llm_token, "body": body})); + let chatgpt_account = headers + .get("chatgpt-account-id") + .and_then(|value| value.to_str().ok()); + let pairing = headers + .get("x-agentdesktop-pairing") + .and_then(|value| value.to_str().ok()); + requests.lock().unwrap().push(json!({"path": path, "method": method.as_str(), "authorization": authorization, "apiKey": api_key, "llmToken": llm_token, "chatgptAccount": chatgpt_account, "pairing": pairing, "body": body})); if path == "/health" { return Json(json!({"ready": true})).into_response(); } if !matches!( path, - "/v1/messages" | "/v1/messages/count_tokens" | "/v1/responses" + "/v1/messages" + | "/v1/messages/count_tokens" + | "/v1/responses" + | "/backend-api/codex/responses" ) { return StatusCode::NOT_FOUND.into_response(); } if authorization != Some(format!("Bearer {TOKEN}").as_str()) && api_key != Some(TOKEN) { return (StatusCode::UNAUTHORIZED, Json(json!({"type": "error", "error": {"type": "authentication_error", "message": "Expected the test API key"}}))).into_response(); } + if path == "/backend-api/codex/responses" { + if llm_token != Some("codex-subscription-fixture") + || chatgpt_account != Some("fixture-account") + || pairing.is_some() + { + return (StatusCode::UNAUTHORIZED, "invalid subscription passthrough").into_response(); + } + return responses(&body); + } if path == "/v1/responses" { return responses(&body); } diff --git a/crates/controller/src/service.rs b/crates/controller/src/service.rs index 1073540f..6fe3570b 100644 --- a/crates/controller/src/service.rs +++ b/crates/controller/src/service.rs @@ -15,7 +15,7 @@ use tracing::{debug, error, info, warn}; use uuid::Uuid; use x509_parser::{extensions::GeneralName, parse_x509_certificate}; -use agentdesktop_core::config::{LlmGatewayAuthentication, ProgramAuthentication, ProgramsConfig}; +use agentdesktop_core::config::{LlmGatewayAuthentication, ProgramsConfig}; use agentdesktop_usage::UsageError; use crate::{ @@ -62,17 +62,7 @@ impl FleetAgentService { } fn subscription_is_configured_for_client(programs: &ProgramsConfig, client_id: &str) -> bool { - match client_id { - "claude-code" => programs - .claude_code - .as_ref() - .is_some_and(|program| program.auth == Some(ProgramAuthentication::Subscription)), - "claude-desktop" => programs - .claude_desktop - .as_ref() - .is_some_and(|program| program.auth == Some(ProgramAuthentication::Subscription)), - _ => false, - } + programs.subscription_auth(client_id).is_some() } #[tonic::async_trait] @@ -658,6 +648,18 @@ mod tests { &config.programs, "codex" )); + let codex = agentdesktop_core::config::parse_daemon( + "llmGateway:\n url: https://gateway.example.com\n proxyUrl: https://gateway.example.com/codex\n authentication:\n type: controllerJwt\n audience: gateway\n allowedClientIds: [codex]\nprograms:\n codex:\n auth: subscription\n", + ) + .unwrap(); + assert!(subscription_is_configured_for_client( + &codex.programs, + "codex" + )); + assert!(!subscription_is_configured_for_client( + &codex.programs, + "claude-code" + )); } #[tokio::test] diff --git a/crates/core/src/config.rs b/crates/core/src/config.rs index 2ac17f4b..43d787df 100644 --- a/crates/core/src/config.rs +++ b/crates/core/src/config.rs @@ -277,7 +277,7 @@ pub struct LlmGatewayConfig { )] pub github_oauth: Option, /// What the programs behind the local LLM proxy (`programs.copilot`, - /// `programs.vscode`) do while the daemon runs without that proxy: + /// `programs.vscode`, or Codex with `auth: subscription`) do while the daemon runs without that proxy: /// `failClosed` (the default) leaves their managed entries in place, so /// the tools stay pointed at the loopback port and fail instead /// of falling back to GitHub past the gateway (with no entries in place @@ -707,6 +707,25 @@ pub struct ProgramsConfig { } impl ProgramsConfig { + pub fn subscription_auth(&self, client_id: &str) -> Option { + let (auth, mechanism) = match client_id { + "claude-code" => ( + self.claude_code.as_ref()?.auth.as_ref(), + SubscriptionAuth::CredentialHelper, + ), + "claude-desktop" => ( + self.claude_desktop.as_ref()?.auth.as_ref(), + SubscriptionAuth::CredentialHelper, + ), + "codex" => ( + self.codex.as_ref()?.auth.as_ref(), + SubscriptionAuth::ClientPassthrough, + ), + _ => return None, + }; + (auth == Some(&ProgramAuthentication::Subscription)).then_some(mechanism) + } + fn is_empty(&self) -> bool { self.claude_code.is_none() && self.claude_desktop.is_none() @@ -718,6 +737,12 @@ impl ProgramsConfig { } } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum SubscriptionAuth { + CredentialHelper, + ClientPassthrough, +} + /// Settings reconciled into Claude Code's managed configuration. /// /// Arbitrary keys are written directly to Agentdesktop's managed-settings @@ -768,6 +793,10 @@ pub struct CodexConfig { /// Whether this program uses the top-level LLM gateway. #[serde(default = "default_true", skip_serializing_if = "is_true")] pub use_llm_gateway: bool, + /// Keep Codex's ChatGPT subscription login and route it through the local proxy. + /// Requires user mode, gateway authentication, and `llmGateway.proxyUrl`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub auth: Option, /// Arbitrary values written to Codex's organization-managed TOML configuration. /// /// Use Codex's native snake_case configuration keys. TOML has no null value, @@ -1502,6 +1531,25 @@ fn validate_daemon( } } + if let Some(codex) = &programs.codex + && codex.auth == Some(ProgramAuthentication::Subscription) + { + if !codex.use_llm_gateway || llm_gateway.is_none() { + anyhow::bail!( + "programs.codex.auth subscription requires that program to use an LLM gateway" + ); + } + let gateway = llm_gateway.expect("gateway checked above"); + if gateway.authentication.is_none() { + anyhow::bail!( + "programs.codex.auth subscription requires oidc or controllerJwt gateway authentication" + ); + } + if gateway.proxy_url.is_none() { + anyhow::bail!("programs.codex.auth subscription requires llmGateway.proxyUrl"); + } + } + if let Some(open_code) = &programs.open_code && llm_gateway.is_some() && open_code.use_llm_gateway @@ -1833,6 +1881,10 @@ usage: "../../../examples/claude-subscription/config.yaml" )) .expect("Claude subscription user configuration example"); + super::parse_local_daemon(include_str!( + "../../../examples/codex-subscription/config.yaml" + )) + .expect("Codex subscription user configuration example"); let standalone = super::parse_local_daemon(include_str!("../../../examples/standalone/config.yaml")) .expect("standalone daemon configuration example"); @@ -1931,6 +1983,76 @@ programs: ); } + #[test] + fn subscription_auth_distinguishes_helper_credentials_from_client_passthrough() { + let config = parse_daemon( + "llmGateway:\n url: https://gateway.example.com\n proxyUrl: https://gateway.example.com/subscription\n authentication:\n type: controllerJwt\n audience: gateway\n allowedClientIds: [claude-code, claude-desktop, codex]\nprograms:\n claudeCode:\n auth: subscription\n claudeDesktop: {}\n codex:\n auth: subscription\n", + ) + .unwrap(); + assert_eq!( + config.programs.subscription_auth("claude-code"), + Some(super::SubscriptionAuth::CredentialHelper) + ); + assert_eq!( + config.programs.subscription_auth("codex"), + Some(super::SubscriptionAuth::ClientPassthrough) + ); + for client_id in ["claude-desktop", "copilot-cli", "opencode", "unknown"] { + assert_eq!(config.programs.subscription_auth(client_id), None); + } + assert_eq!( + super::ProgramsConfig::default().subscription_auth("codex"), + None + ); + } + + #[test] + fn codex_subscription_requires_an_authenticated_proxy_gateway() { + let valid = r#" +llmGateway: + url: https://gateway.example.com + proxyUrl: https://gateway.example.com/codex + authentication: + type: controllerJwt + audience: agentgateway + allowedClientIds: [codex] +programs: + codex: + auth: subscription +"#; + let config = parse_daemon(valid).expect("valid Codex subscription configuration"); + assert_eq!( + config.programs.codex.unwrap().auth, + Some(super::ProgramAuthentication::Subscription) + ); + for (yaml, message) in [ + ( + "programs:\n codex:\n auth: subscription\n".to_owned(), + "use an LLM gateway", + ), + ( + valid.replace(" auth: subscription", " auth: subscription\n useLlmGateway: false"), + "use an LLM gateway", + ), + ( + valid.replace(" proxyUrl: https://gateway.example.com/codex\n", ""), + "proxyUrl", + ), + ( + valid.replace( + " authentication:\n type: controllerJwt\n audience: agentgateway\n allowedClientIds: [codex]\n", + "", + ), + "gateway authentication", + ), + ] { + assert!( + parse_daemon(&yaml).unwrap_err().to_string().contains(message), + "expected {message} for {yaml}" + ); + } + } + #[test] fn claude_subscription_composes_with_oidc() { let daemon = parse_daemon( diff --git a/examples/codex-subscription/README.md b/examples/codex-subscription/README.md new file mode 100644 index 00000000..4762eecb --- /dev/null +++ b/examples/codex-subscription/README.md @@ -0,0 +1,80 @@ +# Codex subscription through Agentdesktop + +Use Codex's existing ChatGPT login for inference while separately authenticating +to the organization's gateway. Agentdesktop manages the Codex provider, local +proxy route, and pairing header; Codex owns ChatGPT sign-in and token refresh. + +This requires a user-mode daemon. The Windows system service cannot run the +per-user proxy. Update both the controller and daemon before distributing +`programs.codex.auth: subscription` to a fleet. + +## Run + +From the repository root, start the example gateway and identity provider: + +```sh +docker compose -f examples/codex-subscription/compose.yaml up -d +agentdesktop daemon --user --config examples/codex-subscription/config.yaml +``` + +Complete organization sign-in in the browser. The local Dex account is +`admin@example.com` with password `password`. Codex must also be signed in to +ChatGPT using its normal `codex login` flow. Run Codex in another terminal after +the daemon applies the configuration. + +Agentdesktop manages `~/.codex/config.toml` in user mode. A pre-existing file +without Agentdesktop's ownership header is reported as a conflict and is not +overwritten. Migrate the settings you need into `programs.codex.managedConfig` +and resolve the reported conflict before testing; do not delete Codex's separate +login file. A `daemon.codex.config` override must point to a configuration layer +the installed Codex version actually reads, not `requirements.toml`. + +The reconciler generates a provider with `requires_openai_auth = true`, forces +ChatGPT login, disables WebSockets for this HTTP-only route, and writes the +pairing header into an owner-only configuration file. It does not write a +gateway credential command or copy ChatGPT tokens into the file. + +## Credential flow + +```text +Codex -> local /codex proxy -> gateway -> ChatGPT Codex backend +``` + +The local proxy puts gateway identity in `Authorization` and Codex's bearer in +`x-llm-token`. The gateway strictly validates organization identity first, then +restores the ChatGPT bearer and strips `x-llm-token`. The gateway forwards the +account-identification headers and maps `/codex-subscription/responses` to +`https://chatgpt.com/backend-api/codex/responses`, not the API-key endpoint at +`api.openai.com/v1`. + +For controller-issued gateway credentials, authorize client ID `codex` and set +the same subscription configuration in the controller's effective policy. The +controller then marks these requests as subscription-billed. OIDC-only deployments +must classify subscription traffic using the dedicated gateway route when +collecting usage. + +The proxy must be enabled locally with `daemon.llmProxy.listen`; `proxyUrl` and +the program configuration belong in the controller-distributed policy on enrolled +devices. `llmGateway.githubOAuth` is not required for the managed `/codex` route. +When the proxy is unavailable, the existing `whenProxyUnavailable` policy applies: +`failClosed` retains existing managed settings; `failOpen` removes the managed +file. Before the first successful apply there is no managed configuration to keep. + +## Validation and Production + +The integration test runs the real Codex 0.153.4 CLI with synthetic ChatGPT/OIDC +credentials against a mock gateway and verifies streamed responses, separated +credentials, account headers, and cleanup. It does not validate a paid ChatGPT +account against OpenAI or certify provider terms. Validate both login-refresh +flows and actual subscription access before rollout. + +This local example intentionally uses loopback HTTP and a demonstration Dex +account. Use HTTPS and your production identity provider outside local testing. +Protect pairing and credential files with user-only permissions or Windows ACLs, +keep tokens out of logs, and review provider terms and credential-handling policy. + +Stop the example services with: + +```sh +docker compose -f examples/codex-subscription/compose.yaml down +``` diff --git a/examples/codex-subscription/agentgateway.yaml b/examples/codex-subscription/agentgateway.yaml new file mode 100644 index 00000000..f57fe118 --- /dev/null +++ b/examples/codex-subscription/agentgateway.yaml @@ -0,0 +1,36 @@ +gateways: + gateway: + port: 4002 + +routes: +- name: codex-subscription + gateways: gateway + matches: + - path: + pathPrefix: /codex-subscription + policies: + jwtAuth: + mode: strict + issuer: http://127.0.0.1:5558/dex + audiences: [agentdesktop-local] + jwks: + url: http://dex:5558/dex/keys + urlRewrite: + authority: + full: chatgpt.com + path: + prefix: /backend-api/codex + timeout: + requestTimeout: 120s + backends: + - host: chatgpt.com:443 + policies: + backendTLS: {} + backendAuth: + passthrough: {} + transformations: + request: + set: + authorization: '"Bearer " + request.headers["x-llm-token"]' + remove: + - x-llm-token diff --git a/examples/codex-subscription/compose.yaml b/examples/codex-subscription/compose.yaml new file mode 100644 index 00000000..0e434afd --- /dev/null +++ b/examples/codex-subscription/compose.yaml @@ -0,0 +1,19 @@ +services: + dex: + image: ghcr.io/dexidp/dex:v2.45.1 + command: dex serve /etc/dex/config.yaml + ports: + - 127.0.0.1:5558:5558 + volumes: + - ./dex.yaml:/etc/dex/config.yaml:ro + + agentgateway: + image: cr.agentgateway.dev/agentgateway:v1.4.1 + command: -f /etc/agentgateway/config.yaml + depends_on: + - dex + restart: on-failure + ports: + - 127.0.0.1:4002:4002 + volumes: + - ./agentgateway.yaml:/etc/agentgateway/config.yaml:ro diff --git a/examples/codex-subscription/config.yaml b/examples/codex-subscription/config.yaml new file mode 100644 index 00000000..d012e7ed --- /dev/null +++ b/examples/codex-subscription/config.yaml @@ -0,0 +1,20 @@ +daemon: + user: true + llmProxy: + listen: 127.0.0.1:4012 + +llmGateway: + url: http://127.0.0.1:4002 + proxyUrl: http://127.0.0.1:4002/codex-subscription + authentication: + type: oidc + issuer: http://127.0.0.1:5558/dex + clientId: agentdesktop-local + scopes: [openid, email, offline_access] + allowInsecure: true + +programs: + codex: + auth: subscription + managedConfig: + model: gpt-5.4 diff --git a/examples/codex-subscription/dex.yaml b/examples/codex-subscription/dex.yaml new file mode 100644 index 00000000..a3fcb108 --- /dev/null +++ b/examples/codex-subscription/dex.yaml @@ -0,0 +1,24 @@ +issuer: http://127.0.0.1:5558/dex + +storage: + type: memory + +web: + http: 0.0.0.0:5558 + +oauth2: + skipApprovalScreen: true + +staticClients: +- id: agentdesktop-local + name: Agentdesktop Codex subscription example + public: true + redirectURIs: + - http://127.0.0.1:51327/callback + +enablePasswordDB: true +staticPasswords: +- email: admin@example.com + username: admin + userID: local-admin + hash: $2y$10$cA5xff5AaaAoxB8VZjwXUeXZIlH8V0sLf/JkIBXHip6gnQsyKJiHW diff --git a/frontend/controller/package.json b/frontend/controller/package.json index 1fece4bc..071214be 100644 --- a/frontend/controller/package.json +++ b/frontend/controller/package.json @@ -17,7 +17,8 @@ "@agentdesktop/ui": "workspace:*", "lucide-react": "^1.55.0", "react": "^19.3.0", - "react-dom": "^19.3.0" + "react-dom": "^19.3.0", + "yaml": "^2.9.1" }, "devDependencies": { "@biomejs/biome": "^2.5.15", diff --git a/frontend/controller/src/views/ConfigurationView.stories.tsx b/frontend/controller/src/views/ConfigurationView.stories.tsx index 91c05611..6ccefee6 100644 --- a/frontend/controller/src/views/ConfigurationView.stories.tsx +++ b/frontend/controller/src/views/ConfigurationView.stories.tsx @@ -71,6 +71,33 @@ export const CopilotFailOpen: Story = { }, }; +export const CodexSubscription: Story = { + args: { + initialConfig: { + llmGateway: { + url: "https://gateway.example.com", + proxyUrl: "https://gateway.example.com/chatgpt", + whenProxyUnavailable: "failOpen", + authentication: { type: "controllerJwt", audience: "agentgateway" }, + }, + programs: { codex: { auth: "subscription" } }, + }, + }, + play: async ({ canvas, canvasElement }) => { + await expect( + canvas.getByRole("checkbox", { + name: /Fail open without the local proxy/, + }), + ).toBeChecked(); + const output = canvasElement.querySelector(".output-card code"); + await expect(output?.textContent).toContain('auth: "subscription"'); + await expect(output).toHaveTextContent("whenProxyUnavailable: failOpen"); + await expect(output).toHaveTextContent( + "https://gateway.example.com/chatgpt", + ); + }, +}; + export const ActiveProgramSelection: Story = { args: { initialConfig: programSelectionDaemonConfig }, play: async ({ canvas, canvasElement }) => { diff --git a/frontend/controller/src/views/ConfigurationView.tsx b/frontend/controller/src/views/ConfigurationView.tsx index 77742929..0a0bc153 100644 --- a/frontend/controller/src/views/ConfigurationView.tsx +++ b/frontend/controller/src/views/ConfigurationView.tsx @@ -8,6 +8,7 @@ import { Trash2, } from "lucide-react"; import { useEffect, useRef, useState } from "react"; +import { parse } from "yaml"; import type { AgentDraft, AgentKind, DaemonConfigDocument } from "../types"; @@ -48,7 +49,15 @@ export function ConfigurationView({ return [definition?.label ?? agent.kind]; }); const sandboxUnavailable = incompatibleSandboxAgentNames.length > 0; - const usesLocalProxy = agents.some((agent) => proxyAgents.has(agent.kind)); + const usesLocalProxy = agents.some((agent) => { + if (proxyAgents.has(agent.kind)) return true; + if (agent.kind !== "codex" || !agent.useGateway) return false; + try { + return parse(agent.settings)?.auth === "subscription"; + } catch { + return false; + } + }); const yaml = daemonConfigYaml({ gateway, gatewayUrl, diff --git a/frontend/pnpm-lock.yaml b/frontend/pnpm-lock.yaml index c27a17ab..d5a9c13b 100644 --- a/frontend/pnpm-lock.yaml +++ b/frontend/pnpm-lock.yaml @@ -22,6 +22,9 @@ importers: react-dom: specifier: ^19.3.0 version: 19.3.0(react@19.3.0) + yaml: + specifier: ^2.9.1 + version: 2.9.1 devDependencies: '@biomejs/biome': specifier: ^2.5.15 @@ -31,10 +34,10 @@ importers: version: 10.6.1(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0)) '@storybook/addon-vitest': specifier: ^10.6.1 - version: 10.6.1(@vitest/browser-playwright@5.0.3)(@vitest/browser@5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2))(vitest@5.0.3))(@vitest/runner@4.1.11)(react@19.3.0)(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(vitest@5.0.3) + version: 10.6.1(@vitest/browser-playwright@5.0.3)(@vitest/browser@5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1))(vitest@5.0.3))(@vitest/runner@4.1.11)(react@19.3.0)(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(vitest@5.0.3) '@storybook/react-vite': specifier: ^10.6.1 - version: 10.6.1(@types/react-dom@19.3.0(@types/react@19.3.0))(@types/react@19.3.0)(react-dom@19.3.0(react@19.3.0))(react@19.3.0)(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(typescript@7.0.2)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)) + version: 10.6.1(@types/react-dom@19.3.0(@types/react@19.3.0))(@types/react@19.3.0)(react-dom@19.3.0(react@19.3.0))(react@19.3.0)(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(typescript@7.0.2)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1)) '@types/node': specifier: ^26.6.5 version: 26.6.5 @@ -46,10 +49,10 @@ importers: version: 19.3.0(@types/react@19.3.0) '@vitejs/plugin-react': specifier: ^6.1.2 - version: 6.1.2(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)) + version: 6.1.2(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1)) '@vitest/browser-playwright': specifier: ^5.0.3 - version: 5.0.3(playwright@1.64.0)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2))(vitest@5.0.3) + version: 5.0.3(playwright@1.64.0)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1))(vitest@5.0.3) playwright: specifier: ^1.64.0 version: 1.64.0 @@ -61,10 +64,10 @@ importers: version: 7.0.2 vite: specifier: ^8.3.4 - version: 8.3.4(@types/node@26.6.5)(esbuild@0.28.2) + version: 8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1) vitest: specifier: ^5.0.3 - version: 5.0.3(@types/node@26.6.5)(@vitest/browser-playwright@5.0.3)(@vitest/ui@5.0.3)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)) + version: 5.0.3(@types/node@26.6.5)(@vitest/browser-playwright@5.0.3)(@vitest/ui@5.0.3)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1)) desktop: dependencies: @@ -92,10 +95,10 @@ importers: version: 10.6.1(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0)) '@storybook/addon-vitest': specifier: ^10.6.1 - version: 10.6.1(@vitest/browser-playwright@5.0.3)(@vitest/browser@5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2))(vitest@5.0.3))(@vitest/runner@4.1.11)(react@19.3.0)(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(vitest@5.0.3) + version: 10.6.1(@vitest/browser-playwright@5.0.3)(@vitest/browser@5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1))(vitest@5.0.3))(@vitest/runner@4.1.11)(react@19.3.0)(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(vitest@5.0.3) '@storybook/react-vite': specifier: ^10.6.1 - version: 10.6.1(@types/react-dom@19.3.0(@types/react@19.3.0))(@types/react@19.3.0)(react-dom@19.3.0(react@19.3.0))(react@19.3.0)(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(typescript@7.0.2)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)) + version: 10.6.1(@types/react-dom@19.3.0(@types/react@19.3.0))(@types/react@19.3.0)(react-dom@19.3.0(react@19.3.0))(react@19.3.0)(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(typescript@7.0.2)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1)) '@tauri-apps/cli': specifier: ^2.12.2 version: 2.12.2 @@ -110,10 +113,10 @@ importers: version: 19.3.0(@types/react@19.3.0) '@vitejs/plugin-react': specifier: ^6.1.2 - version: 6.1.2(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)) + version: 6.1.2(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1)) '@vitest/browser-playwright': specifier: ^5.0.3 - version: 5.0.3(playwright@1.64.0)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2))(vitest@5.0.3) + version: 5.0.3(playwright@1.64.0)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1))(vitest@5.0.3) playwright: specifier: ^1.64.0 version: 1.64.0 @@ -125,10 +128,10 @@ importers: version: 7.0.2 vite: specifier: ^8.3.4 - version: 8.3.4(@types/node@26.6.5)(esbuild@0.28.2) + version: 8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1) vitest: specifier: ^5.0.3 - version: 5.0.3(@types/node@26.6.5)(@vitest/browser-playwright@5.0.3)(@vitest/ui@5.0.3)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)) + version: 5.0.3(@types/node@26.6.5)(@vitest/browser-playwright@5.0.3)(@vitest/ui@5.0.3)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1)) ui: dependencies: @@ -2024,6 +2027,11 @@ packages: yallist@3.1.1: resolution: {integrity: sha512-a4UGQaWPH59mOXUYnAG2ewncQS4i4F43Tv3JoAM+s2VDAmS9NsK8GpDMLrCHPksFT7h3K6TOoUNn2pb7RoXx4g==} + yaml@2.9.1: + resolution: {integrity: sha512-3NxN8+78OdzbT7C/WjGsyfPAtJaN3FNDsWxv7Y7mcDsT/oOmgW8BpyQQFFBnvZE3j9Y2Sdz1ULFLezL7Eb2yFw==} + engines: {node: '>= 14.6'} + hasBin: true + snapshots: '@adobe/css-tools@4.5.0': {} @@ -2277,11 +2285,11 @@ snapshots: '@esbuild/win32-x64@0.28.2': optional: true - '@joshwooding/vite-plugin-react-docgen-typescript@0.7.0(typescript@7.0.2)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2))': + '@joshwooding/vite-plugin-react-docgen-typescript@0.7.0(typescript@7.0.2)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1))': dependencies: glob: 13.0.6 react-docgen-typescript: 2.4.0(typescript@7.0.2) - vite: 8.3.4(@types/node@26.6.5)(esbuild@0.28.2) + vite: 8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1) optionalDependencies: typescript: 7.0.2 @@ -2559,24 +2567,24 @@ snapshots: axe-core: 4.13.0 storybook: 10.6.1(@types/react@19.3.0)(react@19.3.0) - '@storybook/addon-vitest@10.6.1(@vitest/browser-playwright@5.0.3)(@vitest/browser@5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2))(vitest@5.0.3))(@vitest/runner@4.1.11)(react@19.3.0)(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(vitest@5.0.3)': + '@storybook/addon-vitest@10.6.1(@vitest/browser-playwright@5.0.3)(@vitest/browser@5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1))(vitest@5.0.3))(@vitest/runner@4.1.11)(react@19.3.0)(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(vitest@5.0.3)': dependencies: '@storybook/global': 5.0.0 '@storybook/icons': 2.1.0(react@19.3.0) storybook: 10.6.1(@types/react@19.3.0)(react@19.3.0) optionalDependencies: - '@vitest/browser': 5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2))(vitest@5.0.3) - '@vitest/browser-playwright': 5.0.3(playwright@1.64.0)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2))(vitest@5.0.3) + '@vitest/browser': 5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1))(vitest@5.0.3) + '@vitest/browser-playwright': 5.0.3(playwright@1.64.0)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1))(vitest@5.0.3) '@vitest/runner': 4.1.11 - vitest: 5.0.3(@types/node@26.6.5)(@vitest/browser-playwright@5.0.3)(@vitest/ui@5.0.3)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)) + vitest: 5.0.3(@types/node@26.6.5)(@vitest/browser-playwright@5.0.3)(@vitest/ui@5.0.3)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1)) transitivePeerDependencies: - react - '@storybook/builder-vite@10.6.1(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2))': + '@storybook/builder-vite@10.6.1(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1))': dependencies: storybook: 10.6.1(@types/react@19.3.0)(react@19.3.0) ts-dedent: 2.3.0 - vite: 8.3.4(@types/node@26.6.5)(esbuild@0.28.2) + vite: 8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1) '@storybook/global@5.0.0': {} @@ -2593,11 +2601,11 @@ snapshots: '@types/react': 19.3.0 '@types/react-dom': 19.3.0(@types/react@19.3.0) - '@storybook/react-vite@10.6.1(@types/react-dom@19.3.0(@types/react@19.3.0))(@types/react@19.3.0)(react-dom@19.3.0(react@19.3.0))(react@19.3.0)(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(typescript@7.0.2)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2))': + '@storybook/react-vite@10.6.1(@types/react-dom@19.3.0(@types/react@19.3.0))(@types/react@19.3.0)(react-dom@19.3.0(react@19.3.0))(react@19.3.0)(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(typescript@7.0.2)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1))': dependencies: - '@joshwooding/vite-plugin-react-docgen-typescript': 0.7.0(typescript@7.0.2)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)) + '@joshwooding/vite-plugin-react-docgen-typescript': 0.7.0(typescript@7.0.2)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1)) '@rollup/pluginutils': 5.4.0 - '@storybook/builder-vite': 10.6.1(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)) + '@storybook/builder-vite': 10.6.1(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1)) '@storybook/react': 10.6.1(@types/react-dom@19.3.0(@types/react@19.3.0))(@types/react@19.3.0)(react-dom@19.3.0(react@19.3.0))(react@19.3.0)(storybook@10.6.1(@types/react@19.3.0)(react@19.3.0))(typescript@7.0.2) empathic: 2.0.1 magic-string: 1.2.3 @@ -2607,7 +2615,7 @@ snapshots: resolve: 1.22.12 storybook: 10.6.1(@types/react@19.3.0)(react@19.3.0) tsconfig-paths: 4.2.0 - vite: 8.3.4(@types/node@26.6.5)(esbuild@0.28.2) + vite: 8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1) optionalDependencies: typescript: 7.0.2 transitivePeerDependencies: @@ -2818,35 +2826,35 @@ snapshots: '@typescript/typescript-win32-x64@7.0.2': optional: true - '@vitejs/plugin-react@6.1.2(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2))': + '@vitejs/plugin-react@6.1.2(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1))': dependencies: '@rolldown/pluginutils': 1.0.1 - vite: 8.3.4(@types/node@26.6.5)(esbuild@0.28.2) + vite: 8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1) - '@vitest/browser-playwright@5.0.3(playwright@1.64.0)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2))(vitest@5.0.3)': + '@vitest/browser-playwright@5.0.3(playwright@1.64.0)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1))(vitest@5.0.3)': dependencies: - '@vitest/browser': 5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2))(vitest@5.0.3) - '@vitest/mocker': 5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)) + '@vitest/browser': 5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1))(vitest@5.0.3) + '@vitest/mocker': 5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1)) playwright: 1.64.0 tinyrainbow: 3.1.1 - vitest: 5.0.3(@types/node@26.6.5)(@vitest/browser-playwright@5.0.3)(@vitest/ui@5.0.3)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)) + vitest: 5.0.3(@types/node@26.6.5)(@vitest/browser-playwright@5.0.3)(@vitest/ui@5.0.3)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1)) transitivePeerDependencies: - bufferutil - msw - utf-8-validate - vite - '@vitest/browser@5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2))(vitest@5.0.3)': + '@vitest/browser@5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1))(vitest@5.0.3)': dependencies: '@blazediff/core': 1.10.0 - '@vitest/mocker': 5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)) + '@vitest/mocker': 5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1)) '@vitest/ui': 5.0.3(vitest@5.0.3) '@vitest/utils': 5.0.3 magic-string: 1.2.3 pngjs: 7.0.0 sirv: 3.0.2 tinyrainbow: 3.1.1 - vitest: 5.0.3(@types/node@26.6.5)(@vitest/browser-playwright@5.0.3)(@vitest/ui@5.0.3)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)) + vitest: 5.0.3(@types/node@26.6.5)(@vitest/browser-playwright@5.0.3)(@vitest/ui@5.0.3)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1)) ws: 8.21.3 transitivePeerDependencies: - bufferutil @@ -2862,14 +2870,14 @@ snapshots: chai: 5.3.3 tinyrainbow: 2.0.0 - '@vitest/mocker@5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2))': + '@vitest/mocker@5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1))': dependencies: '@jridgewell/trace-mapping': 0.3.31 '@vitest/spy': 5.0.3 estree-walker: 3.0.3 magic-string: 1.2.3 optionalDependencies: - vite: 8.3.4(@types/node@26.6.5)(esbuild@0.28.2) + vite: 8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1) '@vitest/pretty-format@3.2.4': dependencies: @@ -2904,7 +2912,7 @@ snapshots: pathe: 2.0.3 sirv: 3.0.2 tinyrainbow: 3.1.1 - vitest: 5.0.3(@types/node@26.6.5)(@vitest/browser-playwright@5.0.3)(@vitest/ui@5.0.3)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)) + vitest: 5.0.3(@types/node@26.6.5)(@vitest/browser-playwright@5.0.3)(@vitest/ui@5.0.3)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1)) '@vitest/utils@3.2.4': dependencies: @@ -3485,7 +3493,7 @@ snapshots: dependencies: react: 19.3.0 - vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2): + vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1): dependencies: lightningcss: 1.33.0 picomatch: 4.0.7 @@ -3496,11 +3504,12 @@ snapshots: '@types/node': 26.6.5 esbuild: 0.28.2 fsevents: 2.3.3 + yaml: 2.9.1 - vitest@5.0.3(@types/node@26.6.5)(@vitest/browser-playwright@5.0.3)(@vitest/ui@5.0.3)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)): + vitest@5.0.3(@types/node@26.6.5)(@vitest/browser-playwright@5.0.3)(@vitest/ui@5.0.3)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1)): dependencies: '@types/chai': 5.2.3 - '@vitest/mocker': 5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)) + '@vitest/mocker': 5.0.3(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1)) chai: 6.2.2 es-module-lexer: 2.3.2 expect-type: 1.4.0 @@ -3511,11 +3520,11 @@ snapshots: tinybench: 6.1.4 tinyexec: 1.3.0 tinyglobby: 0.2.17 - vite: 8.3.4(@types/node@26.6.5)(esbuild@0.28.2) + vite: 8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1) why-is-node-running: 3.2.1 optionalDependencies: '@types/node': 26.6.5 - '@vitest/browser-playwright': 5.0.3(playwright@1.64.0)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2))(vitest@5.0.3) + '@vitest/browser-playwright': 5.0.3(playwright@1.64.0)(vite@8.3.4(@types/node@26.6.5)(esbuild@0.28.2)(yaml@2.9.1))(vitest@5.0.3) '@vitest/ui': 5.0.3(vitest@5.0.3) transitivePeerDependencies: - msw @@ -3529,3 +3538,5 @@ snapshots: is-wsl: 3.1.1 yallist@3.1.1: {} + + yaml@2.9.1: {} diff --git a/schema/daemon-config.json b/schema/daemon-config.json index 91d8f061..342b7c1f 100644 --- a/schema/daemon-config.json +++ b/schema/daemon-config.json @@ -137,6 +137,17 @@ "description": "Settings reconciled into Codex's organization-managed configuration.\n\nValues under `managedConfig` are written to Codex's `managed_config.toml`.\nWhen generated LLM-gateway settings overlap with those values,\nAgentdesktop's generated values take precedence.", "type": "object", "properties": { + "auth": { + "description": "Keep Codex's ChatGPT subscription login and route it through the local proxy.\nRequires user mode, gateway authentication, and `llmGateway.proxyUrl`.", + "anyOf": [ + { + "$ref": "#/$defs/ProgramAuthentication" + }, + { + "type": "null" + } + ] + }, "managedConfig": { "description": "Arbitrary values written to Codex's organization-managed TOML configuration.\n\nUse Codex's native snake_case configuration keys. TOML has no null value,\nso null values cannot be reconciled.", "type": "object", @@ -508,7 +519,7 @@ "type": "string" }, "whenProxyUnavailable": { - "description": "What the programs behind the local LLM proxy (`programs.copilot`,\n`programs.vscode`) do while the daemon runs without that proxy:\n`failClosed` (the default) leaves their managed entries in place, so\nthe tools stay pointed at the loopback port and fail instead\nof falling back to GitHub past the gateway (with no entries in place\nyet, the tools keep their own settings); `failOpen` removes them\nuntil the proxy is back (VS Code on `githubModels` then talks to GitHub\ndirectly, the Copilot CLI and VS Code on `ownModels` lose the gateway's\nmodels).", + "description": "What the programs behind the local LLM proxy (`programs.copilot`,\n`programs.vscode`, or Codex with `auth: subscription`) do while the daemon runs without that proxy:\n`failClosed` (the default) leaves their managed entries in place, so\nthe tools stay pointed at the loopback port and fail instead\nof falling back to GitHub past the gateway (with no entries in place\nyet, the tools keep their own settings); `failOpen` removes them\nuntil the proxy is back (VS Code on `githubModels` then talks to GitHub\ndirectly, the Copilot CLI and VS Code on `ownModels` lose the gateway's\nmodels).", "$ref": "#/$defs/ProxyUnavailable" } }, diff --git a/schema/daemon-config.md b/schema/daemon-config.md index 957ce2ed..5e6434a8 100644 --- a/schema/daemon-config.md +++ b/schema/daemon-config.md @@ -54,7 +54,7 @@ |`llmGateway.githubOAuth.source`|enum|Where the GitHub credential comes from.
Possible values: `deviceFlow`, `request`.| |`llmGateway.proxyUrl`|string|Base URL the local LLM proxy forwards to, when it differs from `url`.

`url` is shared by every program that sets `useLlmGateway`, so it can
only carry one path prefix. A gateway that puts each provider behind its
own prefix therefore cannot serve both a program and the proxy from one
value. Setting this leaves `url` to the programs and gives the proxy its
own target. The same rules as for `url` apply.| |`llmGateway.url`|string|Base HTTP or HTTPS URL of the LLM gateway.

The URL must include a host and cannot include credentials, a query, or a fragment.| -|`llmGateway.whenProxyUnavailable`|enum|What the programs behind the local LLM proxy (`programs.copilot`,
`programs.vscode`) do while the daemon runs without that proxy:
`failClosed` (the default) leaves their managed entries in place, so
the tools stay pointed at the loopback port and fail instead
of falling back to GitHub past the gateway (with no entries in place
yet, the tools keep their own settings); `failOpen` removes them
until the proxy is back (VS Code on `githubModels` then talks to GitHub
directly, the Copilot CLI and VS Code on `ownModels` lose the gateway's
models).
Possible values: `failClosed`, `failOpen`.| +|`llmGateway.whenProxyUnavailable`|enum|What the programs behind the local LLM proxy (`programs.copilot`,
`programs.vscode`, or Codex with `auth: subscription`) do while the daemon runs without that proxy:
`failClosed` (the default) leaves their managed entries in place, so
the tools stay pointed at the loopback port and fail instead
of falling back to GitHub past the gateway (with no entries in place
yet, the tools keep their own settings); `failOpen` removes them
until the proxy is back (VS Code on `githubModels` then talks to GitHub
directly, the Copilot CLI and VS Code on `ownModels` lose the gateway's
models).
Possible values: `failClosed`, `failOpen`.| |`programSelection`|object|Whether the device user may choose which configured programs are managed.

A daemon that predates this setting rejects a configuration that
contains it, so upgrade devices before setting it in
controller-delivered policy.| |`programSelection.allowUserChoice`|boolean|Let the device user opt out of all configured programs, or choose the
ones to manage, with `agentdesktop programs set`.
Agentdesktop removes the settings it wrote for a program the user opts
out of, and reports the program as opted out.

Defaults to `false`: every configured program is managed and a choice
stored on the device is ignored until this is enabled again.| |`programs`|object|Per-program settings reconciled on this device.| @@ -65,6 +65,7 @@ |`programs.claudeDesktop.auth`|enum|Upstream authentication used by this agent.
Possible values: `subscription`.| |`programs.claudeDesktop.useLlmGateway`|boolean|Whether this program uses the top-level LLM gateway.| |`programs.codex`|object|Codex managed configuration.| +|`programs.codex.auth`|enum|Keep Codex's ChatGPT subscription login and route it through the local proxy.
Requires user mode, gateway authentication, and `llmGateway.proxyUrl`.
Possible values: `subscription`.| |`programs.codex.managedConfig`|object|Arbitrary values written to Codex's organization-managed TOML configuration.

Use Codex's native snake_case configuration keys. TOML has no null value,
so null values cannot be reconciled.| |`programs.codex.useLlmGateway`|boolean|Whether this program uses the top-level LLM gateway.| |`programs.copilot`|object|GitHub Copilot CLI managed configuration.| From 25a3913b8fcb80a19d7f049a45bf65ea08cb0374 Mon Sep 17 00:00:00 2001 From: Peter Jausovec Date: Mon, 19 Oct 2026 16:52:17 -0700 Subject: [PATCH 2/6] pr feedback, implement codex the same way as claude Signed-off-by: Peter Jausovec --- README.md | 121 +++--- crates/agent/src/api.rs | 13 +- crates/agent/src/daemon.rs | 23 +- crates/agent/src/lib.rs | 4 +- crates/agent/src/oidc.rs | 159 +++++++- .../src/provider/codex/integration_tests.rs | 45 +-- crates/agent/src/provider/codex/mod.rs | 18 +- crates/agent/src/provider/codex/reconcile.rs | 83 ++-- crates/agent/src/subscription.rs | 85 +++- .../anthropic.rs} | 140 +++---- crates/agent/src/subscription_oauth/mod.rs | 313 +++++++++++++++ crates/agent/src/subscription_oauth/openai.rs | 375 ++++++++++++++++++ crates/agent/tests/common/gateway.rs | 20 +- crates/core/src/config.rs | 48 +-- examples/codex-subscription/README.md | 88 ++-- examples/codex-subscription/agentgateway.yaml | 54 ++- examples/codex-subscription/compose.yaml | 2 +- examples/codex-subscription/config.yaml | 5 +- frontend/controller/package.json | 3 +- .../src/views/ConfigurationView.stories.tsx | 10 +- .../src/views/ConfigurationView.tsx | 11 +- frontend/pnpm-lock.yaml | 6 +- schema/daemon-config.json | 4 +- schema/daemon-config.md | 4 +- 24 files changed, 1211 insertions(+), 423 deletions(-) rename crates/agent/src/{anthropic_oauth.rs => subscription_oauth/anthropic.rs} (54%) create mode 100644 crates/agent/src/subscription_oauth/mod.rs create mode 100644 crates/agent/src/subscription_oauth/openai.rs diff --git a/README.md b/README.md index c72e562e..9d9136ba 100644 --- a/README.md +++ b/README.md @@ -94,12 +94,47 @@ The desktop and fleet interfaces also expose the MCP server and skill inventory. See the [standalone quickstart](https://agentdesktop.dev/docs/getting-started/standalone/) for prerequisites, test credentials, and a walkthrough of the local services. +## Provider subscriptions + +Claude Code, Claude Desktop, and Codex can connect provider subscriptions through +Agentdesktop's OAuth flow using `auth: subscription`. Agentdesktop owns token +storage and refresh and supplies a composite credential through each client's +credential helper. These integrations connect directly to the gateway and do not +require the local LLM proxy. + +For Codex: + +```yaml +llmGateway: + url: https://gateway.example.com/v1 + authentication: + type: oidc + issuer: https://login.example.com + clientId: agentdesktop +programs: + codex: + auth: subscription +``` + +Complete organization sign-in and the optional ChatGPT step in Agentdesktop's +browser flow. No separate `codex login` is needed. The helper returns +`agentdesktop:::`; the gateway +validates the identity before restoring the ChatGPT bearer and account header. +Claude retains its existing three-field credential without the account suffix. + +Use Agentgateway's `llm` configuration with the OpenAI provider's base URL set to +`https://chatgpt.com/backend-api/codex`, not `api.openai.com/v1`. This keeps model +routing, LLM policies, and token-usage processing active. See the +[Codex subscription example](examples/codex-subscription/README.md) +for the complete gateway and local sign-in setup. Native Codex 0.153.4 is tested +with synthetic credentials and a mock gateway, not a paid OpenAI account. + ## Local LLM proxy Use the daemon's optional loopback proxy for clients without an API-key helper, or to preserve a client's provider subscription while authenticating separately -to the gateway. Codex ChatGPT subscriptions and VS Code Copilot passthrough use -this same proxy. Enable it in the local configuration file (startup settings are +to the gateway, such as VS Code Copilot passthrough. Claude and Codex subscription +helpers use the direct flow above instead. Enable the proxy in the local configuration file (startup settings are local only, not delivered by the controller). The proxy is available in `--user` mode only: it hands out the current user's gateway credential, so a system daemon rejects @@ -145,14 +180,12 @@ Avoid repeating `/v1` in both URLs. **Without the proxy.** When the daemon runs without its loopback proxy (`daemon.llmProxy.listen` unset, or the address could not be bound), `llmGateway.whenProxyUnavailable` decides what happens to the files of the -programs that use the proxy (`programs.copilot`, `programs.vscode`, and Codex -with `auth: subscription`). With `failClosed`, the default, existing managed +programs that use the proxy (`programs.copilot`, `programs.vscode`). With `failClosed`, the default, existing managed entries stay pointed at the loopback proxy and requests fail until it returns. Before the first successful apply, or if those entries were removed, there is no managed routing to retain: the tools keep their own settings. With `failOpen`, -the daemon removes its managed entries, including Codex's managed configuration -file, until the proxy returns. VS Code using GitHub's models and Codex using its -remaining native configuration may then connect directly to their provider; +the daemon removes its proxy-managed entries until the proxy returns. VS Code using GitHub's models +may then connect directly to its provider; gateway-only custom models are no longer available. If another process occupies the loopback port, it can receive requests containing provider tokens and the pairing value; reserve the port and use this setup on single-user devices. @@ -222,7 +255,7 @@ paths on the prefix-less route: | Prefix | Client ID | Sends upstream | Upstream | | --- | --- | --- | --- | | `/copilot-cli/` | `copilot-cli` | gateway identity only | `llmGateway.url` | -| `/codex/` | `codex` | gateway identity in `Authorization`, Codex's ChatGPT token in `x-llm-token`; HTTP Responses requests | `llmGateway.proxyUrl` (required) | +| `/codex/` | `codex` | manual/legacy client-token passthrough; not used by `auth: subscription` | `llmGateway.proxyUrl` (required) | | `/vscode-copilot/` | `vscode-copilot` | gateway identity only | `llmGateway.url` | | `/vscode-copilot-passthrough/` | `vscode-copilot` | gateway identity in `Authorization`, the client's own bearer token in `x-llm-token` | `llmGateway.proxyUrl` (required) | | `/vscode-copilot-capi//` | `vscode-copilot` | as the pass-through route; a credential-less `GET /_ping` is forwarded without `x-llm-token`; the only route that tunnels WebSocket upgrades | `llmGateway.proxyUrl` (required) | @@ -237,9 +270,7 @@ manual setup copies it from there. The state directory in `--user` mode is home directory from `HOME` or `USERPROFILE`, on every platform), unless `daemon.stateDir` says otherwise. The Copilot CLI's `providers.json` carries it as `"headers": {"x-agentdesktop-pairing": ""}` -on the provider entry, and Codex's generated provider carries it in `http_headers`. -The daemon -writes that Codex file owner-only and redacts its pairing-bearing diffs. Pairing +on the provider entry. Pairing keeps other local users on a shared host and browser pages off the proxy, provided the client files that carry it are owner-only too. It does not restrict the current user's own processes, which can @@ -321,61 +352,6 @@ VS Code on GitHub's models, with what to check at each step (drift, conflicts, removal) and the settings the GitHub organisation owner controls ([`examples/copilot/github-side.md`](examples/copilot/github-side.md)). -### Managed program: Codex ChatGPT subscriptions - -Set `programs.codex.auth: subscription` to keep Codex's own ChatGPT login while -the local proxy authenticates separately to the gateway: - -```yaml -daemon: - user: true - llmProxy: - listen: 127.0.0.1:4012 - -llmGateway: - url: https://gateway.example.com - proxyUrl: https://gateway.example.com/codex-subscription - authentication: - type: oidc - issuer: https://login.example.com - clientId: agentdesktop - -programs: - codex: - auth: subscription -``` - -`daemon.llmProxy.listen` enables the shared proxy locally; `proxyUrl` selects the -gateway's ChatGPT passthrough route. On controller-managed devices, distribute -`llmGateway` and `programs` through controller policy, but keep `daemon.*` local. -For controller JWT authentication, authorize client ID `codex` in -`allowedClientIds`. Upgrade both controller and daemon before using this option. - -Agentdesktop generates the Codex provider pointed at `/codex`, inserts the pairing -header, and enables ChatGPT authentication without a conflicting gateway -credential command. No manual Codex provider, `daemon.llmProxy.clientId`, or -`githubOAuth` configuration is needed. WebSockets are disabled for this HTTP-only -route. Keep Codex signed in with `codex login`; Codex owns the ChatGPT credential -and its refresh, while Agentdesktop obtains and refreshes gateway identity. - -The gateway must validate organization identity from `Authorization`, restore the -ChatGPT bearer from `x-llm-token`, and remove that transport header before -forwarding. Preserve Codex's account-identification headers and map -`/codex-subscription/responses` to -`https://chatgpt.com/backend-api/codex/responses`, not `api.openai.com/v1`. - -This requires a per-user daemon, not the Windows system service. Agentdesktop -manages `~/.codex/config.toml` in user mode; a pre-existing personal file without -its ownership header is reported as a conflict rather than overwritten. Migrate -needed settings into `programs.codex.managedConfig` and resolve that conflict -without removing Codex's separate login file. A `daemon.codex.config` override -must name an ordinary configuration layer Codex reads, not `requirements.toml`. - -See the [complete Codex subscription example](examples/codex-subscription/README.md) -for the gateway, sign-in, and production boundaries. The real Codex 0.153.4 CLI is -tested with synthetic credentials and a mock gateway, not a paid OpenAI account. -Validate live subscription access, token refresh, and provider terms before rollout. - ### Managed program: GitHub Copilot CLI `programs.copilot` makes the daemon write the Copilot CLI's BYOK provider @@ -833,9 +809,8 @@ usage: Then send agentgateway's access log to the fleet API with that token. `root` is the CA that signed the controller's TLS certificate. The added fields carry the user, device, and tool from the controller JWT. The billing expression below -recognizes Claude's composite subscription credential with either controller-JWT -or OIDC gateway authentication. Codex uses separate headers instead: classify its -dedicated passthrough route as subscription traffic, or use the controller JWT's +recognizes Claude and Codex composite subscription credentials with either controller-JWT +or OIDC gateway authentication. You can also use the controller JWT's `billing` claim when available. Usage without the added identity fields is shown as unknown. @@ -874,10 +849,10 @@ frontendPolicies: - **Not metered:** VS Code with `copilotChat: githubModels` reaches GitHub through the gateway's pass-through route, which is not an LLM route, so its requests leave no usage records. The desktop app says so when the program is - configured that way. Codex's subscription example also uses raw HTTP - passthrough; forwarding and a subscription-billing claim alone do not extract - token counts or prices. The gateway must export those fields for such usage - to appear in the totals. + configured that way. Codex's subscription example instead uses the LLM + pipeline, which extracts token usage; configure the gateway's usage exporter + to include it in Agentdesktop totals. The dedicated example listener marks + its access logs as subscription-billed rather than organization API-key cost. - Hourly totals are kept for 90 days and daily totals for 13 months. agentgateway does not retry an export, so requests made while the controller is unreachable are not counted. Usage stays after a device is deleted. diff --git a/crates/agent/src/api.rs b/crates/agent/src/api.rs index 72d31ea2..ab2db468 100644 --- a/crates/agent/src/api.rs +++ b/crates/agent/src/api.rs @@ -551,6 +551,7 @@ async fn resolve_gateway_credential( &state.state_dir, state.oidc_callback_listen, continue_in_browser, + client_id, ) .await } else { @@ -568,10 +569,14 @@ async fn controller_subscription_credential( client_id: &str, ) -> Result { let bad_gateway = |error: anyhow::Error| (StatusCode::BAD_GATEWAY, format!("{error:#}")); - let subscription = - subscription::credential(&state.state_dir, state.oidc_callback_listen, false) - .await - .map_err(bad_gateway)?; + let subscription = subscription::credential( + &state.state_dir, + state.oidc_callback_listen, + false, + client_id, + ) + .await + .map_err(bad_gateway)?; let identity = remote::llm_gateway_credential( controller, &state.state_dir, diff --git a/crates/agent/src/daemon.rs b/crates/agent/src/daemon.rs index d624b6f2..8b22fcb8 100644 --- a/crates/agent/src/daemon.rs +++ b/crates/agent/src/daemon.rs @@ -951,12 +951,18 @@ fn start_gateway_authentication( }; let authentication = gateway.authentication.clone(); let github = gateway.github_oauth.clone().filter(|_| proxy_enabled); - let subscription = ["claude-code", "claude-desktop"] + let subscription_clients: Vec<_> = ["claude-code", "claude-desktop", "codex"] .into_iter() - .any(|client_id| { + .filter(|client_id| { config.programs.subscription_auth(client_id) == Some(agentdesktop_core::config::SubscriptionAuth::CredentialHelper) - }); + }) + .filter(|client_id| { + *client_id != "claude-desktop" + || config.programs.subscription_auth("claude-code").is_none() + }) + .collect(); + let subscription = !subscription_clients.is_empty(); if authentication.is_none() && !subscription { return; } @@ -1013,15 +1019,16 @@ fn start_gateway_authentication( continue_in_browser = acquired.interactive && subscription; tracing::info!(%issuer, "LLM gateway OIDC authentication ready"); } - if subscription { - tracing::info!("starting Anthropic subscription authentication"); - crate::anthropic_oauth::credential( + for (index, client_id) in subscription_clients.into_iter().enumerate() { + tracing::info!(client_id, "starting provider subscription authentication"); + crate::subscription::credential( &state_dir, callback_listen, - !continue_in_browser, + continue_in_browser && index == 0, + client_id, ) .await?; - tracing::info!("Anthropic subscription authentication ready"); + tracing::info!(client_id, "provider subscription authentication ready"); } if let Some(github) = github { // source: request has no credential of its own to acquire; the diff --git a/crates/agent/src/lib.rs b/crates/agent/src/lib.rs index 1bf0f6de..ed16e58b 100644 --- a/crates/agent/src/lib.rs +++ b/crates/agent/src/lib.rs @@ -1,4 +1,4 @@ -pub mod anthropic_oauth; +pub use subscription_oauth::anthropic as anthropic_oauth; pub mod api; pub mod cli; pub mod daemon; @@ -9,6 +9,7 @@ pub mod identity; mod inventory_watch; mod llm_proxy; pub mod oidc; +use subscription_oauth::openai as openai_oauth; mod program_choice; pub mod provider; pub mod reconcile; @@ -16,6 +17,7 @@ pub mod remote; pub mod secret_store; pub mod secure_fs; pub mod subscription; +mod subscription_oauth; mod tick; #[cfg(windows)] diff --git a/crates/agent/src/oidc.rs b/crates/agent/src/oidc.rs index bac61be7..b969b94b 100644 --- a/crates/agent/src/oidc.rs +++ b/crates/agent/src/oidc.rs @@ -648,8 +648,54 @@ pub(crate) async fn wait_for_authorization_code_with_page( callback_listen: Option, page: AuthorizationPage, open_browser: bool, +) -> anyhow::Result> { + wait_for_authorization_code_with_callback( + authorization_url, + redirect_uri, + expected_state, + callback_listen, + page, + open_browser, + None, + ) + .await +} + +pub(crate) async fn wait_for_subscription_authorization_code( + authorization_url: &str, + provider_callback: &Url, + expected_state: String, + callback_listen: Option, + open_browser: bool, +) -> anyhow::Result> { + let prompt_uri = + Url::parse("http://localhost:51327/callback").expect("valid subscription prompt URI"); + wait_for_authorization_code_with_callback( + authorization_url, + &prompt_uri, + expected_state, + callback_listen, + AuthorizationPage::Subscription, + open_browser, + Some(provider_callback), + ) + .await +} + +async fn wait_for_authorization_code_with_callback( + authorization_url: &str, + redirect_uri: &Url, + expected_state: String, + callback_listen: Option, + page: AuthorizationPage, + open_browser: bool, + provider_callback: Option<&Url>, ) -> anyhow::Result> { let listener = bind_callback(redirect_uri, callback_listen).await?; + let provider_listener = match provider_callback { + Some(callback) => Some(bind_callback(callback, None).await?), + None => None, + }; let (result_sender, result_receiver) = oneshot::channel(); let state = CallbackState { expected_state, @@ -658,13 +704,32 @@ pub(crate) async fn wait_for_authorization_code_with_page( result: Arc::new(Mutex::new(Some(result_sender))), }; let callback_path = redirect_uri.path().to_owned(); - let app = Router::new() + let mut app = Router::new() .route("/", get(authorization_prompt)) .route("/continue", get(continue_authorization)) .route("/skip", get(skip_authorization)) .route("/flow-ready", get(flow_ready)) - .route(&callback_path, get(callback)) - .with_state(state); + .route(&callback_path, get(callback)); + if let Some(provider_callback) = provider_callback + && provider_callback.path() != callback_path + { + app = app.route(provider_callback.path(), get(callback)); + } + let app = app.with_state(state); + let provider_server = provider_listener.map(|listener| { + let app = app.clone(); + let (shutdown, stopped) = oneshot::channel::<()>(); + ( + shutdown, + tokio::spawn(async move { + axum::serve(listener, app) + .with_graceful_shutdown(async move { + let _ = stopped.await; + }) + .await + }), + ) + }); let (shutdown_sender, shutdown_receiver) = oneshot::channel::<()>(); let server = tokio::spawn(async move { axum::serve(listener, app) @@ -693,13 +758,22 @@ pub(crate) async fn wait_for_authorization_code_with_page( tracing::info!(%prompt_url, "continuing in existing Agentdesktop authorization page"); } - let authorization_code = tokio::time::timeout(CALLBACK_TIMEOUT, result_receiver) - .await - .context("timed out waiting for OIDC callback")? - .context("OIDC callback server stopped")??; + let authorization_code = async { + tokio::time::timeout(CALLBACK_TIMEOUT, result_receiver) + .await + .context("timed out waiting for OIDC callback")? + .context("OIDC callback server stopped")? + } + .await; let _ = shutdown_sender.send(()); + if let Some((shutdown, provider_server)) = provider_server { + let _ = shutdown.send(()); + provider_server + .await + .context("join provider OAuth callback server")??; + } server.await.context("join OIDC callback server")??; - Ok(authorization_code) + authorization_code } async fn authorization_prompt(State(state): State) -> Html { @@ -991,6 +1065,75 @@ mod tests { )); } + #[tokio::test] + async fn subscription_prompt_accepts_a_separate_registered_provider_callback() { + tokio::time::timeout(std::time::Duration::from_secs(10), async { + let prompt_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let prompt_address = prompt_listener.local_addr().unwrap(); + let provider_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let provider_address = provider_listener.local_addr().unwrap(); + drop((prompt_listener, provider_listener)); + let prompt = url::Url::parse(&format!("http://{prompt_address}/callback")).unwrap(); + let mut provider = + url::Url::parse(&format!("http://{provider_address}/auth/callback")).unwrap(); + let callback = provider.clone(); + let task = tokio::spawn(async move { + super::wait_for_authorization_code_with_callback( + "https://provider.example/authorize", + &prompt, + "expected-state".to_owned(), + None, + super::AuthorizationPage::Subscription, + false, + Some(&callback), + ) + .await + }); + let client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .unwrap(); + loop { + if let Ok(response) = client + .get(format!("http://{prompt_address}/flow-ready")) + .send() + .await + && response.status().is_success() + { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + } + let response = client + .get(format!("http://{prompt_address}/continue")) + .send() + .await + .unwrap(); + assert_eq!( + response.headers()["location"], + "https://provider.example/authorize" + ); + provider + .query_pairs_mut() + .extend_pairs([("state", "expected-state"), ("code", "provider-code")]); + let response = client.get(provider).send().await.unwrap(); + assert!(response.status().is_success()); + assert!( + response + .text() + .await + .unwrap() + .contains("Agentdesktop connected") + ); + assert_eq!( + task.await.unwrap().unwrap().as_deref(), + Some("provider-code") + ); + }) + .await + .unwrap(); + } + #[test] fn identity_page_only_mentions_subscription_when_configured() { let without_subscription = AuthorizationPage::Identity { diff --git a/crates/agent/src/provider/codex/integration_tests.rs b/crates/agent/src/provider/codex/integration_tests.rs index a93769c6..0aadeec9 100644 --- a/crates/agent/src/provider/codex/integration_tests.rs +++ b/crates/agent/src/provider/codex/integration_tests.rs @@ -22,9 +22,6 @@ async fn managed_chatgpt_subscription_lifecycle() -> anyhow::Result<()> { "crates/agent/src/provider/codex/testdata/Dockerfile", async |container| { let gateway = Gateway::start().await?; - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; - let proxy_address = listener.local_addr()?; - drop(listener); let state_dir = "/home/tester/.local/state/codex-subscription"; let socket = format!("{state_dir}/agentdesktop.sock"); let issuer = "https://codex-fixture.example/"; @@ -41,39 +38,27 @@ async fn managed_chatgpt_subscription_lifecycle() -> anyhow::Result<()> { "tokenEndpoint": "https://codex-fixture.example/token", }).to_string(), ).await?; - let claims = URL_SAFE_NO_PAD.encode(json!({ - "email": "fixture@example.com", - "exp": 4_000_000_000u64, - "https://api.openai.com/auth": { - "chatgpt_account_id": "fixture-account", - "chatgpt_user_id": "fixture-user", - "chatgpt_plan_type": "plus", - }, - }).to_string()); + let entry = Sha256::digest(b"dev.agentdesktop.openai-subscription\0codex"); + let entry = entry.iter().map(|byte| format!("{byte:02x}")).collect::(); container.write( - "/home/tester/.codex/auth.json", + &format!("{state_dir}/secrets/{entry}"), &json!({ - "auth_mode": "chatgpt", - "tokens": { - "id_token": format!("e30.{claims}.fixture"), - "access_token": "codex-subscription-fixture", - "refresh_token": "fixture-refresh-token", - "account_id": "fixture-account", - }, - "last_refresh": "2090-01-01T00:00:00Z", + "accessToken": "codex-subscription-fixture", + "refreshToken": "fixture-refresh-token", + "accountId": "fixture-account", + "expiresAtUnixSeconds": 4_000_000_000u64, }).to_string(), ).await?; + container.exec(&["mkdir", "-p", "/home/tester/.codex"]).await?; container.exec(&["chown", "-R", "tester:tester", state_dir, "/home/tester/.codex"]).await?; container.exec(&["chmod", "700", state_dir, &format!("{state_dir}/secrets")]).await?; container.write(CONFIG, &json!({ "daemon": { "user": true, "stateDir": state_dir, - "llmProxy": { "listen": proxy_address.to_string() }, }, "llmGateway": { - "url": gateway.url, - "proxyUrl": format!("{}backend-api/codex", gateway.url), + "url": format!("{}backend-api/codex", gateway.url), "authentication": { "type": "oidc", "issuer": issuer, "clientId": client_id }, }, "programs": { "codex": { @@ -93,8 +78,9 @@ async fn managed_chatgpt_subscription_lifecycle() -> anyhow::Result<()> { let managed = container.read(USER_SETTINGS).await?; let parsed: toml::Value = toml::from_str(&managed)?; let provider = &parsed["model_providers"]["agentdesktop"]; - ensure!(provider["requires_openai_auth"].as_bool() == Some(true), "ChatGPT authentication missing"); - ensure!(provider.get("auth").is_none(), "gateway credential command conflicts with ChatGPT auth"); + ensure!(provider["requires_openai_auth"].as_bool() == Some(false), "Codex must use Agentdesktop OAuth instead of native ChatGPT login"); + ensure!(provider["auth"]["command"].as_str() == Some("/usr/local/bin/agentdesktop"), "Agentdesktop subscription helper missing"); + container.exec(&["test", "!", "-e", "/home/tester/.codex/auth.json"]).await?; let output = container.exec_as("tester", &[ "env", "HTTP_PROXY=http://127.0.0.1:9", "HTTPS_PROXY=http://127.0.0.1:9", "NO_PROXY=127.0.0.1,localhost", "codex", "exec", "--skip-git-repo-check", "--ephemeral", "--json", PROMPT, @@ -102,9 +88,8 @@ async fn managed_chatgpt_subscription_lifecycle() -> anyhow::Result<()> { ensure!(output.contains("provider-integration-ok"), "native Codex subscription request failed: {output}"); ensure!(gateway.requests().iter().any(|request| { request["path"] == "/backend-api/codex/responses" - && request["authorization"] == format!("Bearer {TOKEN}") - && request["llmToken"] == "codex-subscription-fixture" - && request["chatgptAccount"] == "fixture-account" + && request["authorization"] == format!("Bearer agentdesktop:codex-subscription-fixture:{TOKEN}:fixture-account") + && request["llmToken"].is_null() && request["pairing"].is_null() && request["body"]["input"].to_string().contains(PROMPT) }), "Codex subscription headers did not reach the mock gateway correctly"); @@ -112,7 +97,7 @@ async fn managed_chatgpt_subscription_lifecycle() -> anyhow::Result<()> { container.write(CONFIG, &json!({"daemon": {"user": true, "stateDir": state_dir}, "programs": {}}).to_string()).await?; container.exec_as("tester", &["agentdesktop", "daemon", "--user", "--config", CONFIG, "--once"], Duration::from_secs(15)).await?; container.exec(&["test", "!", "-e", USER_SETTINGS]).await?; - container.exec(&["test", "-f", "/home/tester/.codex/auth.json"]).await?; + container.exec(&["test", "!", "-e", "/home/tester/.codex/auth.json"]).await?; Ok(()) }, ).await diff --git a/crates/agent/src/provider/codex/mod.rs b/crates/agent/src/provider/codex/mod.rs index 4e7c7529..cdd94ac4 100644 --- a/crates/agent/src/provider/codex/mod.rs +++ b/crates/agent/src/provider/codex/mod.rs @@ -1,6 +1,6 @@ use std::path::PathBuf; -use agentdesktop_core::config::{DaemonConfig, ProgramAuthentication}; +use agentdesktop_core::config::DaemonConfig; use agentdesktop_core::model::Discovery; use super::{Provider, ReconcileContext}; @@ -41,17 +41,6 @@ impl Provider for Codex { } fn plan(&self, ctx: &ReconcileContext, config: &DaemonConfig) -> anyhow::Result { - if !ctx.merge_user_settings - && config - .programs - .codex - .as_ref() - .is_some_and(|program| program.auth == Some(ProgramAuthentication::Subscription)) - { - anyhow::bail!( - "programs.codex.auth subscription requires --user (or daemon.user: true)" - ); - } let configured = config.programs.codex.as_ref().map(|provider| { let gateway = config .llm_gateway @@ -85,7 +74,7 @@ mod tests { use super::*; #[test] - fn subscription_requires_a_user_mode_daemon() { + fn subscription_can_use_a_system_mode_credential_helper() { let config = agentdesktop_core::config::parse_daemon( "llmGateway:\n url: https://gateway.example.com\n proxyUrl: https://gateway.example.com/codex\n authentication:\n type: controllerJwt\n audience: gateway\n allowedClientIds: [codex]\nprograms:\n codex:\n auth: subscription\n", ) @@ -96,7 +85,6 @@ mod tests { socket: PathBuf::from("socket"), llm_proxy: None, }; - let error = Codex::default().plan(&context, &config).err().unwrap(); - assert!(error.to_string().contains("requires --user")); + assert!(Codex::default().plan(&context, &config).is_ok()); } } diff --git a/crates/agent/src/provider/codex/reconcile.rs b/crates/agent/src/provider/codex/reconcile.rs index 150f4a42..bdd879c0 100644 --- a/crates/agent/src/provider/codex/reconcile.rs +++ b/crates/agent/src/provider/codex/reconcile.rs @@ -23,7 +23,7 @@ pub(super) fn plan( credential_helper: &Path, socket: &Path, sandbox: Option<&SandboxConfig>, - proxy: Option<(SocketAddr, &str)>, + _proxy: Option<(SocketAddr, &str)>, config: Option<(&CodexConfig, Option<&LlmGatewayConfig>)>, plan: &ReconcilePlan, ) -> anyhow::Result<()> { @@ -32,24 +32,12 @@ pub(super) fn plan( }; let subscription = config.auth == Some(ProgramAuthentication::Subscription); - if subscription && proxy.is_none() { - let gateway = gateway.context("Codex subscription requires an LLM gateway")?; - if crate::reconcile::fail_closed_without_proxy(gateway, plan, || { - plan.read(path).is_ok_and(|contents| { - contents.starts_with(MANAGED_HEADER.as_bytes()) && carries_pairing(&contents) - }) - }) { - return Ok(()); - } - plan.inactive(crate::reconcile::FAIL_OPEN_REASON); - return remove(path, plan); - } let mode = if subscription { 0o600 } else { MANAGED_FILE_MODE }; - let settings = managed_config(config, gateway, credential_helper, socket, sandbox, proxy)?; + let settings = managed_config(config, gateway, credential_helper, socket, sandbox, None)?; let mut contents = MANAGED_HEADER.as_bytes().to_vec(); contents.extend_from_slice( toml::to_string_pretty(&settings) @@ -140,7 +128,7 @@ fn managed_config( credential_helper: &Path, socket: &Path, sandbox: Option<&SandboxConfig>, - proxy: Option<(SocketAddr, &str)>, + _proxy: Option<(SocketAddr, &str)>, ) -> anyhow::Result { let mut settings = serde_json::to_value(&config.managed_config) .context("serialize Codex pass-through managed configuration")?; @@ -190,12 +178,13 @@ fn managed_config( "wire_api": "responses", }); if config.auth == Some(ProgramAuthentication::Subscription) { - let (listen, pairing) = proxy.context("Codex subscription requires the local LLM proxy")?; - provider["base_url"] = json!(format!("http://{listen}{}", crate::llm_proxy::CODEX_ROUTE)); - provider["requires_openai_auth"] = json!(true); + provider["base_url"] = json!(gateway.url.as_str().trim_end_matches('/')); + provider["requires_openai_auth"] = json!(false); provider["supports_websockets"] = json!(false); - provider["http_headers"] = json!({ crate::llm_proxy::PAIRING_HEADER: pairing }); - settings["forced_login_method"] = json!("chatgpt"); + settings + .as_object_mut() + .unwrap() + .remove("forced_login_method"); if let Some(existing) = settings .get_mut("model_providers") .and_then(|providers| providers.get_mut(provider_name)) @@ -213,15 +202,17 @@ fn managed_config( existing.remove(field); } } - } else if gateway + } + if gateway .authentication .as_ref() .is_some_and(LlmGatewayAuthentication::uses_credential_helper) { - let timeout_ms = if matches!( - gateway.authentication, - Some(LlmGatewayAuthentication::Oidc { .. }) - ) { + let timeout_ms = if config.auth == Some(ProgramAuthentication::Subscription) + || matches!( + gateway.authentication, + Some(LlmGatewayAuthentication::Oidc { .. }) + ) { 600_000 } else { 5_000 @@ -325,7 +316,7 @@ programs: } #[test] - fn subscription_provider_keeps_chatgpt_auth_and_pairs_with_the_proxy() { + fn subscription_provider_uses_the_agentdesktop_oauth_helper_directly() { let config = subscription_config(); let settings = managed_config( config.programs.codex.as_ref().unwrap(), @@ -337,20 +328,17 @@ programs: ) .unwrap(); assert_eq!(settings["model"], "gpt-5.4"); - assert_eq!(settings["forced_login_method"], "chatgpt"); + assert!(settings.get("forced_login_method").is_none()); let provider = &settings["model_providers"]["agentdesktop"]; - assert_eq!(provider["base_url"], "http://127.0.0.1:4010/codex"); - assert_eq!(provider["requires_openai_auth"], true); + assert_eq!(provider["base_url"], "https://gateway.example.com/llm"); + assert_eq!(provider["requires_openai_auth"], false); assert_eq!(provider["supports_websockets"], false); assert_eq!(provider["wire_api"], "responses"); - assert_eq!( - provider["http_headers"], - serde_json::json!({ - crate::llm_proxy::PAIRING_HEADER: "pairing-secret" - }) - ); + assert_eq!(provider["auth"]["command"], "agentdesktop"); + assert_eq!(provider["auth"]["args"][4], "codex"); + assert_eq!(provider["auth"]["timeout_ms"], 600_000); for field in [ - "auth", + "http_headers", "env_key", "experimental_bearer_token", "env_http_headers", @@ -364,7 +352,7 @@ programs: } #[test] - fn subscription_lifecycle_redacts_pairing_and_obeys_proxy_unavailable_policy() { + fn subscription_lifecycle_is_independent_of_proxy_availability() { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("config.toml"); let mut config = subscription_config(); @@ -386,17 +374,14 @@ programs: plan }; let absent = run_plan(&config, None); - assert_eq!( - absent.inactive_reason().as_deref(), - Some(crate::reconcile::FAIL_CLOSED_NOT_IN_PLACE_REASON) - ); + assert!(absent.inactive_reason().is_none()); absent.apply().unwrap(); - assert!(!path.exists()); + assert!(path.exists()); for pairing in ["first-pairing-secret", "second-pairing-secret"] { let plan = run_plan(&config, Some(("127.0.0.1:4010".parse().unwrap(), pairing))); assert!(!plan.render().contains("pairing-secret")); plan.apply().unwrap(); - assert!(std::fs::read_to_string(&path).unwrap().contains(pairing)); + assert!(!std::fs::read_to_string(&path).unwrap().contains(pairing)); #[cfg(unix)] { use std::os::unix::fs::PermissionsExt; @@ -409,22 +394,16 @@ programs: } let contents = std::fs::read(&path).unwrap(); let closed = run_plan(&config, None); - assert_eq!( - closed.inactive_reason().as_deref(), - Some(crate::reconcile::FAIL_CLOSED_REASON) - ); + assert!(closed.inactive_reason().is_none()); closed.apply().unwrap(); assert_eq!(std::fs::read(&path).unwrap(), contents); config.llm_gateway.as_mut().unwrap().when_proxy_unavailable = agentdesktop_core::config::ProxyUnavailable::FailOpen; let open = run_plan(&config, None); - assert_eq!( - open.inactive_reason().as_deref(), - Some(crate::reconcile::FAIL_OPEN_REASON) - ); + assert!(open.inactive_reason().is_none()); assert!(!open.render().contains("pairing-secret")); open.apply().unwrap(); - assert!(!path.exists()); + assert!(path.exists()); } #[test] diff --git a/crates/agent/src/subscription.rs b/crates/agent/src/subscription.rs index 24c2a14c..fc2907b8 100644 --- a/crates/agent/src/subscription.rs +++ b/crates/agent/src/subscription.rs @@ -2,15 +2,22 @@ use std::{net::SocketAddr, path::Path}; use agentdesktop_core::model::LlmGatewayCredential; -use crate::anthropic_oauth; +use crate::{anthropic_oauth, openai_oauth}; + +pub struct SubscriptionCredential { + pub token: LlmGatewayCredential, + pub account_id: Option, +} pub async fn compose( identity: LlmGatewayCredential, state_dir: &Path, callback_listen: Option, continue_in_browser: bool, + client_id: &str, ) -> anyhow::Result { - let subscription = credential(state_dir, callback_listen, continue_in_browser).await?; + let subscription = + credential(state_dir, callback_listen, continue_in_browser, client_id).await?; Ok(combine(subscription, identity)) } @@ -19,23 +26,50 @@ pub async fn credential( state_dir: &Path, callback_listen: Option, continue_in_browser: bool, -) -> anyhow::Result> { - anthropic_oauth::credential(state_dir, callback_listen, !continue_in_browser).await + client_id: &str, +) -> anyhow::Result> { + match client_id { + "claude-code" | "claude-desktop" => { + Ok( + anthropic_oauth::credential(state_dir, callback_listen, !continue_in_browser) + .await? + .map(|token| SubscriptionCredential { + token, + account_id: None, + }), + ) + } + "codex" => Ok( + openai_oauth::credential(state_dir, callback_listen, !continue_in_browser) + .await? + .map(|(token, account_id)| SubscriptionCredential { + token, + account_id: Some(account_id), + }), + ), + _ => anyhow::bail!("subscription authentication is not supported for client {client_id}"), + } } pub fn combine( - subscription: Option, + subscription: Option, identity: LlmGatewayCredential, ) -> LlmGatewayCredential { let Some(subscription) = subscription else { return identity; }; + let mut credential = format!( + "agentdesktop:{}:{}", + subscription.token.credential, identity.credential + ); + if let Some(account_id) = subscription.account_id { + credential.push(':'); + credential.push_str(&account_id); + } LlmGatewayCredential { - credential: format!( - "agentdesktop:{}:{}", - subscription.credential, identity.credential - ), + credential, expires_at_unix_seconds: subscription + .token .expires_at_unix_seconds .min(identity.expires_at_unix_seconds), } @@ -45,14 +79,17 @@ pub fn combine( mod tests { use agentdesktop_core::model::LlmGatewayCredential; - use super::combine; + use super::{SubscriptionCredential, combine}; #[test] fn credential_contains_subscription_then_identity_and_uses_earliest_expiry() { let credential = combine( - Some(LlmGatewayCredential { - credential: "subscription".to_owned(), - expires_at_unix_seconds: 200, + Some(SubscriptionCredential { + token: LlmGatewayCredential { + credential: "subscription".to_owned(), + expires_at_unix_seconds: 200, + }, + account_id: None, }), LlmGatewayCredential { credential: "identity".to_owned(), @@ -63,6 +100,28 @@ mod tests { assert_eq!(credential.expires_at_unix_seconds, 100); } + #[test] + fn chatgpt_credential_keeps_identity_third_and_appends_account() { + let credential = combine( + Some(SubscriptionCredential { + token: LlmGatewayCredential { + credential: "chatgpt-token".to_owned(), + expires_at_unix_seconds: 100, + }, + account_id: Some("account-1".to_owned()), + }), + LlmGatewayCredential { + credential: "identity".to_owned(), + expires_at_unix_seconds: 200, + }, + ); + assert_eq!( + credential.credential, + "agentdesktop:chatgpt-token:identity:account-1" + ); + assert_eq!(credential.expires_at_unix_seconds, 100); + } + #[test] fn skipped_subscription_uses_identity_credential() { let credential = combine( diff --git a/crates/agent/src/anthropic_oauth.rs b/crates/agent/src/subscription_oauth/anthropic.rs similarity index 54% rename from crates/agent/src/anthropic_oauth.rs rename to crates/agent/src/subscription_oauth/anthropic.rs index 6a59946d..6c78ac9c 100644 --- a/crates/agent/src/anthropic_oauth.rs +++ b/crates/agent/src/subscription_oauth/anthropic.rs @@ -1,17 +1,12 @@ -use std::{ - net::SocketAddr, - path::Path, - sync::atomic::{AtomicBool, Ordering}, -}; - -use agentdesktop_core::http::ClientExt; -use agentdesktop_core::model::LlmGatewayCredential; +use std::{net::SocketAddr, path::Path}; + +use agentdesktop_core::{http::ClientExt, model::LlmGatewayCredential}; use anyhow::{Context, bail}; use serde::{Deserialize, Serialize}; -use tokio::sync::Mutex; use url::Url; -use crate::{oidc, secret_store::SecretStore}; +use super::{LoginState, OAuthProvider}; +use crate::subscription::SubscriptionCredential; // These are the production Claude subscription OAuth settings used by Claude Code // 2.1.237. Anthropic does not currently publish them as a stable third-party API. @@ -21,12 +16,9 @@ const CLIENT_ID: &str = "9d1c250a-e61b-44d9-88ed-5944d1962f5e"; const REDIRECT_URI: &str = "http://localhost:51327/callback"; const SCOPE: &str = "user:inference"; const REQUESTED_LIFETIME_SECONDS: u64 = 365 * 24 * 60 * 60; -const EXPIRY_SKEW_SECONDS: u64 = 60; -const SECRET_SERVICE: &str = "dev.agentdesktop.anthropic-subscription"; -const SECRET_ACCOUNT: &str = "claude"; -static LOGIN: Mutex<()> = Mutex::const_new(()); -static DECLINED_FOR_PROCESS: AtomicBool = AtomicBool::new(false); +static LOGIN: LoginState = LoginState::new(); +struct Anthropic; #[derive(Deserialize, Serialize)] #[serde(rename_all = "camelCase")] @@ -58,61 +50,54 @@ pub async fn credential( callback_listen: Option, open_browser: bool, ) -> anyhow::Result> { - let state_dir = state_dir.to_owned(); - tokio::spawn(async move { credential_inner(&state_dir, callback_listen, open_browser).await }) - .await - .context("join Anthropic subscription OAuth task")? + Ok( + super::credential(&Anthropic, &LOGIN, state_dir, callback_listen, open_browser) + .await? + .map(|credential| credential.token), + ) } -async fn credential_inner( - state_dir: &Path, - callback_listen: Option, - open_browser: bool, -) -> anyhow::Result> { - let _login = LOGIN.lock().await; - let store = SecretStore::new(state_dir)?; - let redirect_uri = Url::parse(REDIRECT_URI).expect("static Anthropic redirect URI is valid"); - if let Some(stored) = load(&store)? - && stored.expires_at_unix_seconds > now().saturating_add(EXPIRY_SKEW_SECONDS) - { - if !open_browser { - oidc::continue_subscription_page(&redirect_uri, callback_listen, true).await?; - } - return Ok(Some(as_credential(&stored))); +#[async_trait::async_trait] +impl OAuthProvider for Anthropic { + type Token = StoredToken; + const NAME: &'static str = "Anthropic"; + const SERVICE: &'static str = "dev.agentdesktop.anthropic-subscription"; + const ACCOUNT: &'static str = "claude"; + const REDIRECT_URI: &'static str = REDIRECT_URI; + + fn authorization_url(&self, redirect: &Url, state: &str, challenge: &str) -> Url { + authorization_url(redirect, state, challenge) + } + + fn expires_at(&self, token: &StoredToken) -> u64 { + token.expires_at_unix_seconds } - if DECLINED_FOR_PROCESS.load(Ordering::Relaxed) { - if !open_browser { - oidc::continue_subscription_page(&redirect_uri, callback_listen, false).await?; + + fn credential_from_token(&self, token: StoredToken) -> SubscriptionCredential { + SubscriptionCredential { + token: LlmGatewayCredential { + credential: token.access_token, + expires_at_unix_seconds: token.expires_at_unix_seconds, + }, + account_id: None, } - return Ok(None); } - let state = oidc::random_secret(); - let (verifier, challenge) = oidc::pkce(); - let authorization_url = authorization_url(&redirect_uri, &state, &challenge); - let authorization_code = oidc::wait_for_authorization_code_with_page( - authorization_url.as_str(), - &redirect_uri, - state.clone(), - callback_listen, - oidc::AuthorizationPage::Subscription, - open_browser, - ) - .await?; - let Some(authorization_code) = authorization_code else { - DECLINED_FOR_PROCESS.store(true, Ordering::Relaxed); - return Ok(None); - }; - let response = exchange_authorization_code(&authorization_code, &verifier, &state).await?; - if !response.token_type.eq_ignore_ascii_case("bearer") { - bail!("Anthropic OAuth endpoint returned unsupported token type"); + async fn exchange( + &self, + code: &str, + verifier: &str, + state: &str, + ) -> anyhow::Result { + let response = exchange_authorization_code(code, verifier, state).await?; + if !response.token_type.eq_ignore_ascii_case("bearer") { + bail!("Anthropic OAuth endpoint returned unsupported token type"); + } + Ok(StoredToken { + access_token: response.access_token, + expires_at_unix_seconds: super::now().saturating_add(response.expires_in), + }) } - let stored = StoredToken { - access_token: response.access_token, - expires_at_unix_seconds: now().saturating_add(response.expires_in), - }; - save(&store, &stored)?; - Ok(Some(as_credential(&stored))) } fn authorization_url(redirect_uri: &Url, state: &str, challenge: &str) -> Url { @@ -153,35 +138,6 @@ async fn exchange_authorization_code( .context("exchange Anthropic subscription authorization code") } -fn load(store: &SecretStore) -> anyhow::Result> { - store - .get_optional(SECRET_SERVICE, SECRET_ACCOUNT)? - .map(|value| serde_json::from_str(&value).context("decode stored Anthropic OAuth token")) - .transpose() -} - -fn save(store: &SecretStore, token: &StoredToken) -> anyhow::Result<()> { - store.set( - SECRET_SERVICE, - SECRET_ACCOUNT, - &serde_json::to_string(token).context("encode Anthropic OAuth token")?, - ) -} - -fn as_credential(token: &StoredToken) -> LlmGatewayCredential { - LlmGatewayCredential { - credential: token.access_token.clone(), - expires_at_unix_seconds: token.expires_at_unix_seconds, - } -} - -fn now() -> u64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs() -} - #[cfg(test)] mod tests { use std::collections::HashMap; diff --git a/crates/agent/src/subscription_oauth/mod.rs b/crates/agent/src/subscription_oauth/mod.rs new file mode 100644 index 00000000..58128346 --- /dev/null +++ b/crates/agent/src/subscription_oauth/mod.rs @@ -0,0 +1,313 @@ +use std::{ + net::SocketAddr, + path::Path, + sync::atomic::{AtomicBool, Ordering}, +}; + +use anyhow::Context; +use serde::{Serialize, de::DeserializeOwned}; +use tokio::sync::Mutex; +use url::Url; + +use crate::{oidc, secret_store::SecretStore, subscription::SubscriptionCredential}; + +pub mod anthropic; +pub(crate) mod openai; + +const PROMPT_URI: &str = "http://localhost:51327/callback"; +const EXPIRY_SKEW_SECONDS: u64 = 60; + +struct LoginState { + lock: Mutex<()>, + declined: AtomicBool, +} + +impl LoginState { + const fn new() -> Self { + Self { + lock: Mutex::const_new(()), + declined: AtomicBool::new(false), + } + } +} + +#[async_trait::async_trait] +trait OAuthProvider: Send + Sync { + type Token: Serialize + DeserializeOwned + Send + Sync; + + const NAME: &'static str; + const SERVICE: &'static str; + const ACCOUNT: &'static str; + const REDIRECT_URI: &'static str; + + fn authorization_url(&self, redirect: &Url, state: &str, challenge: &str) -> Url; + fn expires_at(&self, token: &Self::Token) -> u64; + fn credential_from_token(&self, token: Self::Token) -> SubscriptionCredential; + async fn exchange( + &self, + code: &str, + verifier: &str, + state: &str, + ) -> anyhow::Result; + async fn refresh(&self, _token: &Self::Token) -> anyhow::Result> { + Ok(None) + } +} + +async fn credential( + provider: &'static P, + login: &'static LoginState, + state_dir: &Path, + callback_listen: Option, + open_browser: bool, +) -> anyhow::Result> { + let state_dir = state_dir.to_owned(); + tokio::spawn(async move { + credential_inner(provider, login, &state_dir, callback_listen, open_browser).await + }) + .await + .with_context(|| format!("join {} subscription OAuth task", P::NAME))? +} + +async fn credential_inner( + provider: &P, + login: &LoginState, + state_dir: &Path, + callback_listen: Option, + open_browser: bool, +) -> anyhow::Result> { + let _login = login.lock.lock().await; + let store = SecretStore::new(state_dir)?; + let prompt = Url::parse(PROMPT_URI).expect("valid subscription prompt URI"); + if let Some(stored) = load::

(&store)? { + let ready = if provider.expires_at(&stored) > now().saturating_add(EXPIRY_SKEW_SECONDS) { + Some(stored) + } else if let Some(refreshed) = provider.refresh(&stored).await? { + save::

(&store, &refreshed)?; + Some(refreshed) + } else { + None + }; + if let Some(token) = ready { + if !open_browser { + oidc::continue_subscription_page(&prompt, callback_listen, true).await?; + } + return Ok(Some(provider.credential_from_token(token))); + } + } + if login.declined.load(Ordering::Relaxed) { + if !open_browser { + oidc::continue_subscription_page(&prompt, callback_listen, false).await?; + } + return Ok(None); + } + let redirect = + Url::parse(P::REDIRECT_URI).context("parse subscription provider callback URI")?; + let state = oidc::random_secret(); + let (verifier, challenge) = oidc::pkce(); + let authorization_url = provider.authorization_url(&redirect, &state, &challenge); + let code = if redirect == prompt { + oidc::wait_for_authorization_code_with_page( + authorization_url.as_str(), + &redirect, + state.clone(), + callback_listen, + oidc::AuthorizationPage::Subscription, + open_browser, + ) + .await? + } else { + oidc::wait_for_subscription_authorization_code( + authorization_url.as_str(), + &redirect, + state.clone(), + callback_listen, + open_browser, + ) + .await? + }; + let Some(code) = code else { + login.declined.store(true, Ordering::Relaxed); + return Ok(None); + }; + let token = provider.exchange(&code, &verifier, &state).await?; + save::

(&store, &token)?; + Ok(Some(provider.credential_from_token(token))) +} + +fn load(store: &SecretStore) -> anyhow::Result> { + store + .get_optional(P::SERVICE, P::ACCOUNT)? + .map(|value| { + serde_json::from_str(&value) + .with_context(|| format!("decode stored {} subscription", P::NAME)) + }) + .transpose() +} + +fn save(store: &SecretStore, token: &P::Token) -> anyhow::Result<()> { + store.set( + P::SERVICE, + P::ACCOUNT, + &serde_json::to_string(token) + .with_context(|| format!("encode {} subscription", P::NAME))?, + ) +} + +fn now() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +#[cfg(all(test, target_os = "linux"))] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use agentdesktop_core::model::LlmGatewayCredential; + use serde::{Deserialize, Serialize}; + + use super::*; + + #[derive(Deserialize, Serialize)] + #[serde(rename_all = "camelCase")] + struct Token { + access_token: String, + expires_at_unix_seconds: u64, + } + + #[derive(Default)] + struct TestProvider { + refreshes: AtomicUsize, + fail_refresh: bool, + } + + #[async_trait::async_trait] + impl OAuthProvider for TestProvider { + type Token = Token; + const NAME: &'static str = "test-provider"; + const SERVICE: &'static str = "dev.agentdesktop.test-subscription"; + const ACCOUNT: &'static str = "test"; + const REDIRECT_URI: &'static str = PROMPT_URI; + + fn authorization_url(&self, _: &Url, _: &str, _: &str) -> Url { + panic!("test must not open provider authorization") + } + fn expires_at(&self, token: &Token) -> u64 { + token.expires_at_unix_seconds + } + fn credential_from_token(&self, token: Token) -> SubscriptionCredential { + SubscriptionCredential { + token: LlmGatewayCredential { + credential: token.access_token, + expires_at_unix_seconds: token.expires_at_unix_seconds, + }, + account_id: None, + } + } + async fn exchange(&self, _: &str, _: &str, _: &str) -> anyhow::Result { + panic!("test must not exchange an authorization code") + } + async fn refresh(&self, _: &Token) -> anyhow::Result> { + self.refreshes.fetch_add(1, Ordering::SeqCst); + if self.fail_refresh { + anyhow::bail!("transient refresh failure"); + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + Ok(Some(Token { + access_token: "refreshed".to_owned(), + expires_at_unix_seconds: now() + 3600, + })) + } + } + + #[tokio::test] + async fn cached_credentials_are_reused_and_concurrent_expiry_refreshes_once() { + let directory = tempfile::tempdir().unwrap(); + let store = SecretStore::new(directory.path()).unwrap(); + let provider = TestProvider::default(); + let login = LoginState::new(); + save::( + &store, + &Token { + access_token: "cached".to_owned(), + expires_at_unix_seconds: now() + 3600, + }, + ) + .unwrap(); + let cached = credential_inner(&provider, &login, directory.path(), None, true) + .await + .unwrap() + .unwrap(); + assert_eq!(cached.token.credential, "cached"); + assert_eq!(provider.refreshes.load(Ordering::SeqCst), 0); + save::( + &store, + &Token { + access_token: "expired".to_owned(), + expires_at_unix_seconds: 1, + }, + ) + .unwrap(); + let (first, second) = tokio::join!( + credential_inner(&provider, &login, directory.path(), None, true), + credential_inner(&provider, &login, directory.path(), None, true), + ); + assert_eq!(first.unwrap().unwrap().token.credential, "refreshed"); + assert_eq!(second.unwrap().unwrap().token.credential, "refreshed"); + assert_eq!(provider.refreshes.load(Ordering::SeqCst), 1); + assert_eq!( + load::(&store).unwrap().unwrap().access_token, + "refreshed" + ); + } + + #[tokio::test] + async fn failed_refresh_preserves_the_stored_credential() { + let directory = tempfile::tempdir().unwrap(); + let store = SecretStore::new(directory.path()).unwrap(); + let provider = TestProvider { + fail_refresh: true, + ..TestProvider::default() + }; + save::( + &store, + &Token { + access_token: "expired".to_owned(), + expires_at_unix_seconds: 1, + }, + ) + .unwrap(); + let before = store + .get(TestProvider::SERVICE, TestProvider::ACCOUNT) + .unwrap(); + assert!( + credential_inner(&provider, &LoginState::new(), directory.path(), None, true) + .await + .is_err() + ); + assert_eq!( + store + .get(TestProvider::SERVICE, TestProvider::ACCOUNT) + .unwrap(), + before + ); + } + + #[tokio::test] + async fn declined_sign_in_is_remembered_without_opening_authorization() { + let directory = tempfile::tempdir().unwrap(); + let login = LoginState::new(); + login.declined.store(true, Ordering::Relaxed); + let provider = TestProvider::default(); + assert!( + credential_inner(&provider, &login, directory.path(), None, true) + .await + .unwrap() + .is_none() + ); + assert_eq!(provider.refreshes.load(Ordering::SeqCst), 0); + assert!(!LoginState::new().declined.load(Ordering::Relaxed)); + } +} diff --git a/crates/agent/src/subscription_oauth/openai.rs b/crates/agent/src/subscription_oauth/openai.rs new file mode 100644 index 00000000..ef943c0b --- /dev/null +++ b/crates/agent/src/subscription_oauth/openai.rs @@ -0,0 +1,375 @@ +use std::{net::SocketAddr, path::Path, time::Duration}; + +use agentdesktop_core::model::LlmGatewayCredential; +use anyhow::{Context, bail}; +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use url::Url; + +use super::{EXPIRY_SKEW_SECONDS, LoginState, OAuthProvider}; +use crate::subscription::SubscriptionCredential; + +const AUTHORIZATION_ENDPOINT: &str = "https://auth.openai.com/oauth/authorize"; +const TOKEN_ENDPOINT: &str = "https://auth.openai.com/oauth/token"; +const CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann"; +const REDIRECT_URI: &str = "http://localhost:1455/auth/callback"; + +static LOGIN: LoginState = LoginState::new(); +struct OpenAi; + +#[derive(Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +struct StoredToken { + access_token: String, + refresh_token: String, + account_id: String, + expires_at_unix_seconds: u64, +} + +#[derive(Deserialize)] +struct TokenResponse { + access_token: String, + refresh_token: Option, + id_token: Option, + expires_in: Option, + token_type: Option, +} + +#[derive(Debug)] +struct ReauthorizationRequired; + +impl std::fmt::Display for ReauthorizationRequired { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("ChatGPT subscription requires sign-in again") + } +} + +impl std::error::Error for ReauthorizationRequired {} + +pub async fn credential( + state_dir: &Path, + callback_listen: Option, + open_browser: bool, +) -> anyhow::Result> { + Ok( + super::credential(&OpenAi, &LOGIN, state_dir, callback_listen, open_browser) + .await? + .map(|credential| { + ( + credential.token, + credential + .account_id + .expect("OpenAI adapter supplies account ID"), + ) + }), + ) +} + +#[async_trait::async_trait] +impl OAuthProvider for OpenAi { + type Token = StoredToken; + const NAME: &'static str = "ChatGPT"; + const SERVICE: &'static str = "dev.agentdesktop.openai-subscription"; + const ACCOUNT: &'static str = "codex"; + const REDIRECT_URI: &'static str = REDIRECT_URI; + + fn authorization_url(&self, redirect: &Url, state: &str, challenge: &str) -> Url { + authorization_url(redirect, state, challenge) + } + + fn expires_at(&self, token: &StoredToken) -> u64 { + token.expires_at_unix_seconds + } + + fn credential_from_token(&self, token: StoredToken) -> SubscriptionCredential { + SubscriptionCredential { + token: LlmGatewayCredential { + credential: token.access_token, + expires_at_unix_seconds: token.expires_at_unix_seconds, + }, + account_id: Some(token.account_id), + } + } + + async fn exchange( + &self, + code: &str, + verifier: &str, + _state: &str, + ) -> anyhow::Result { + let form = [ + ("grant_type", "authorization_code"), + ("client_id", CLIENT_ID), + ("code", code), + ("code_verifier", verifier), + ("redirect_uri", REDIRECT_URI), + ]; + let response = request_token(TOKEN_ENDPOINT, &form, false).await?; + stored_token(response, None, super::now()) + } + + async fn refresh(&self, token: &StoredToken) -> anyhow::Result> { + match request_token(TOKEN_ENDPOINT, &refresh_form(&token.refresh_token), true).await { + Ok(response) => Ok(Some(stored_token(response, Some(token), super::now())?)), + Err(error) if error.is::() => Ok(None), + Err(error) => Err(error), + } + } +} + +fn authorization_url(redirect_uri: &Url, state: &str, challenge: &str) -> Url { + let mut endpoint = Url::parse(AUTHORIZATION_ENDPOINT).expect("valid ChatGPT authorization URL"); + endpoint.query_pairs_mut().extend_pairs([ + ("response_type", "code"), + ("client_id", CLIENT_ID), + ("redirect_uri", redirect_uri.as_str()), + ("scope", "openid profile email offline_access"), + ("code_challenge", challenge), + ("code_challenge_method", "S256"), + ("id_token_add_organizations", "true"), + ("codex_cli_simplified_flow", "true"), + ("state", state), + ("originator", "codex_cli_rs"), + ]); + endpoint +} + +fn refresh_form(refresh_token: &str) -> [(&str, &str); 3] { + [ + ("grant_type", "refresh_token"), + ("client_id", CLIENT_ID), + ("refresh_token", refresh_token), + ] +} + +async fn request_token( + endpoint: &str, + form: &[(&str, &str)], + refreshing: bool, +) -> anyhow::Result { + let response = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .connect_timeout(Duration::from_secs(10)) + .timeout(Duration::from_secs(30)) + .build()? + .post(endpoint) + .form(form) + .send() + .await + .context("request ChatGPT subscription OAuth token")?; + let status = response.status(); + if !status.is_success() { + let invalid_grant = response.json::().await.ok().is_some_and(|body| { + matches!( + body["error"].as_str(), + Some("invalid_grant" | "invalid_refresh_token") + ) + }); + if refreshing && (status == reqwest::StatusCode::UNAUTHORIZED || invalid_grant) { + return Err(ReauthorizationRequired.into()); + } + bail!("ChatGPT OAuth token endpoint rejected the request ({status})"); + } + response + .json() + .await + .context("decode ChatGPT OAuth token response") +} + +fn claims(token: &str) -> anyhow::Result { + let payload = token + .split('.') + .nth(1) + .context("ChatGPT token has no claims")?; + let payload = URL_SAFE_NO_PAD + .decode(payload) + .context("decode ChatGPT token claims")?; + serde_json::from_slice(&payload).context("parse ChatGPT token claims") +} + +fn stored_token( + response: TokenResponse, + previous: Option<&StoredToken>, + now: u64, +) -> anyhow::Result { + if response + .token_type + .as_deref() + .is_some_and(|kind| !kind.eq_ignore_ascii_case("bearer")) + { + bail!("ChatGPT OAuth returned an unsupported token type"); + } + if response.access_token.is_empty() + || response + .access_token + .bytes() + .any(|byte| byte <= b' ' || byte == b':') + { + bail!("ChatGPT OAuth returned an invalid access token"); + } + let token_claims = response.id_token.as_deref().map(claims).transpose()?; + let account_id = token_claims + .as_ref() + .and_then(|value| value["https://api.openai.com/auth"]["chatgpt_account_id"].as_str()) + .or_else(|| previous.map(|token| token.account_id.as_str())) + .context("ChatGPT OAuth token has no account ID")?; + if account_id.is_empty() + || !account_id + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || b"._-".contains(&byte)) + { + bail!("ChatGPT OAuth returned an invalid account ID"); + } + if previous.is_some_and(|token| token.account_id != account_id) { + bail!("ChatGPT token refresh changed the account; sign in again"); + } + let expires_at_unix_seconds = response + .expires_in + .map(|lifetime| now.saturating_add(lifetime)) + .or_else(|| claims(&response.access_token).ok()?.get("exp")?.as_u64()) + .context("ChatGPT OAuth token has no expiry")?; + if expires_at_unix_seconds <= now.saturating_add(EXPIRY_SKEW_SECONDS) { + bail!("ChatGPT OAuth returned an expired or short-lived access token"); + } + let refresh_token = response + .refresh_token + .or_else(|| previous.map(|token| token.refresh_token.clone())) + .filter(|token| !token.is_empty()) + .context("ChatGPT OAuth returned no refresh token")?; + Ok(StoredToken { + access_token: response.access_token, + refresh_token, + account_id: account_id.to_owned(), + expires_at_unix_seconds, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn response(account_id: &str) -> TokenResponse { + let payload = URL_SAFE_NO_PAD.encode( + serde_json::json!({ + "https://api.openai.com/auth": {"chatgpt_account_id": account_id} + }) + .to_string(), + ); + TokenResponse { + access_token: "access-token".to_owned(), + refresh_token: Some("refresh-token".to_owned()), + id_token: Some(format!("e30.{payload}.signature")), + expires_in: Some(3600), + token_type: Some("Bearer".to_owned()), + } + } + + #[test] + fn authorization_uses_registered_callback_and_pkce() { + let url = authorization_url(&Url::parse(REDIRECT_URI).unwrap(), "state", "challenge"); + let parameters: std::collections::HashMap<_, _> = url.query_pairs().collect(); + assert_eq!(parameters["client_id"], CLIENT_ID); + assert_eq!( + parameters["redirect_uri"], + "http://localhost:1455/auth/callback" + ); + assert_eq!(parameters["code_challenge_method"], "S256"); + assert_eq!(parameters["state"], "state"); + } + + #[test] + fn token_refresh_retains_account_and_rotates_refresh_token() { + let previous = stored_token(response("account-1"), None, 100).unwrap(); + let mut refreshed = response("account-1"); + refreshed.id_token = None; + refreshed.refresh_token = Some("rotated-refresh".to_owned()); + let stored = stored_token(refreshed, Some(&previous), 200).unwrap(); + assert_eq!(stored.account_id, "account-1"); + assert_eq!(stored.refresh_token, "rotated-refresh"); + assert_eq!(stored.expires_at_unix_seconds, 3800); + assert!(stored_token(response("account-2"), Some(&previous), 200).is_err()); + } + + #[test] + fn invalid_account_and_short_lived_tokens_are_rejected() { + assert!(stored_token(response("account:injection"), None, 100).is_err()); + let mut short = response("account-1"); + short.expires_in = Some(1); + assert!(stored_token(short, None, 100).is_err()); + } + + async fn token_server( + status: reqwest::StatusCode, + body: Value, + ) -> (String, tokio::task::JoinHandle<()>) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}/token", listener.local_addr().unwrap()); + let app = axum::Router::new().route( + "/token", + axum::routing::post( + move |axum::Form(form): axum::Form>| { + let body = body.clone(); + async move { + assert_eq!(form["grant_type"], "refresh_token"); + assert_eq!(form["client_id"], CLIENT_ID); + assert_eq!(form["refresh_token"], "original-refresh"); + (status, axum::Json(body)) + } + }, + ), + ); + let task = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + (endpoint, task) + } + + #[tokio::test] + async fn refresh_uses_the_registered_client_and_preserves_omitted_rotation() { + let (endpoint, server) = token_server( + reqwest::StatusCode::OK, + serde_json::json!({ + "access_token": "refreshed-access", "expires_in": 3600, "token_type": "Bearer" + }), + ) + .await; + let response = request_token(&endpoint, &refresh_form("original-refresh"), true) + .await + .unwrap(); + let mut previous = stored_token(self::response("account-1"), None, 100).unwrap(); + previous.refresh_token = "original-refresh".to_owned(); + let stored = stored_token(response, Some(&previous), 200).unwrap(); + assert_eq!(stored.refresh_token, "original-refresh"); + assert_eq!(stored.access_token, "refreshed-access"); + assert_eq!(stored.account_id, "account-1"); + server.abort(); + } + + #[tokio::test] + async fn rejected_refresh_reauthenticates_but_transient_errors_do_not_leak_tokens() { + for (status, code, reauthorize) in [ + (reqwest::StatusCode::BAD_REQUEST, "invalid_grant", true), + ( + reqwest::StatusCode::SERVICE_UNAVAILABLE, + "temporary_failure", + false, + ), + ] { + let (endpoint, server) = token_server( + status, + serde_json::json!({ + "error": code, "error_description": "secret-refresh-token-must-not-leak" + }), + ) + .await; + let error = request_token(&endpoint, &refresh_form("original-refresh"), true) + .await + .err() + .unwrap(); + assert_eq!(error.is::(), reauthorize); + assert!(!format!("{error:#}").contains("secret-refresh-token")); + server.abort(); + } + } +} diff --git a/crates/agent/tests/common/gateway.rs b/crates/agent/tests/common/gateway.rs index 021cdd47..44611f42 100644 --- a/crates/agent/tests/common/gateway.rs +++ b/crates/agent/tests/common/gateway.rs @@ -93,18 +93,26 @@ async fn handle( ) { return StatusCode::NOT_FOUND.into_response(); } - if authorization != Some(format!("Bearer {TOKEN}").as_str()) && api_key != Some(TOKEN) { - return (StatusCode::UNAUTHORIZED, Json(json!({"type": "error", "error": {"type": "authentication_error", "message": "Expected the test API key"}}))).into_response(); - } if path == "/backend-api/codex/responses" { - if llm_token != Some("codex-subscription-fixture") - || chatgpt_account != Some("fixture-account") + if authorization + != Some( + format!("Bearer agentdesktop:codex-subscription-fixture:{TOKEN}:fixture-account") + .as_str(), + ) + || llm_token.is_some() || pairing.is_some() { - return (StatusCode::UNAUTHORIZED, "invalid subscription passthrough").into_response(); + return ( + StatusCode::UNAUTHORIZED, + "invalid composite ChatGPT credential", + ) + .into_response(); } return responses(&body); } + if authorization != Some(format!("Bearer {TOKEN}").as_str()) && api_key != Some(TOKEN) { + return (StatusCode::UNAUTHORIZED, Json(json!({"type": "error", "error": {"type": "authentication_error", "message": "Expected the test API key"}}))).into_response(); + } if path == "/v1/responses" { return responses(&body); } diff --git a/crates/core/src/config.rs b/crates/core/src/config.rs index 43d787df..f58fae26 100644 --- a/crates/core/src/config.rs +++ b/crates/core/src/config.rs @@ -277,7 +277,7 @@ pub struct LlmGatewayConfig { )] pub github_oauth: Option, /// What the programs behind the local LLM proxy (`programs.copilot`, - /// `programs.vscode`, or Codex with `auth: subscription`) do while the daemon runs without that proxy: + /// `programs.vscode`) do while the daemon runs without that proxy: /// `failClosed` (the default) leaves their managed entries in place, so /// the tools stay pointed at the loopback port and fail instead /// of falling back to GitHub past the gateway (with no entries in place @@ -719,7 +719,7 @@ impl ProgramsConfig { ), "codex" => ( self.codex.as_ref()?.auth.as_ref(), - SubscriptionAuth::ClientPassthrough, + SubscriptionAuth::CredentialHelper, ), _ => return None, }; @@ -793,8 +793,7 @@ pub struct CodexConfig { /// Whether this program uses the top-level LLM gateway. #[serde(default = "default_true", skip_serializing_if = "is_true")] pub use_llm_gateway: bool, - /// Keep Codex's ChatGPT subscription login and route it through the local proxy. - /// Requires user mode, gateway authentication, and `llmGateway.proxyUrl`. + /// Connect a ChatGPT subscription through Agentdesktop's OAuth credential helper. #[serde(default, skip_serializing_if = "Option::is_none")] pub auth: Option, /// Arbitrary values written to Codex's organization-managed TOML configuration. @@ -1502,6 +1501,16 @@ fn validate_daemon( .as_ref() .is_some_and(|program| program.use_llm_gateway), ), + ( + "codex", + programs.codex.as_ref().is_some_and(|program| { + program.auth == Some(ProgramAuthentication::Subscription) + }), + programs + .codex + .as_ref() + .is_some_and(|program| program.use_llm_gateway), + ), ] { if subscription && (!uses_gateway || llm_gateway.is_none()) { @@ -1531,25 +1540,6 @@ fn validate_daemon( } } - if let Some(codex) = &programs.codex - && codex.auth == Some(ProgramAuthentication::Subscription) - { - if !codex.use_llm_gateway || llm_gateway.is_none() { - anyhow::bail!( - "programs.codex.auth subscription requires that program to use an LLM gateway" - ); - } - let gateway = llm_gateway.expect("gateway checked above"); - if gateway.authentication.is_none() { - anyhow::bail!( - "programs.codex.auth subscription requires oidc or controllerJwt gateway authentication" - ); - } - if gateway.proxy_url.is_none() { - anyhow::bail!("programs.codex.auth subscription requires llmGateway.proxyUrl"); - } - } - if let Some(open_code) = &programs.open_code && llm_gateway.is_some() && open_code.use_llm_gateway @@ -1984,7 +1974,7 @@ programs: } #[test] - fn subscription_auth_distinguishes_helper_credentials_from_client_passthrough() { + fn subscription_auth_uses_helpers_for_claude_and_codex() { let config = parse_daemon( "llmGateway:\n url: https://gateway.example.com\n proxyUrl: https://gateway.example.com/subscription\n authentication:\n type: controllerJwt\n audience: gateway\n allowedClientIds: [claude-code, claude-desktop, codex]\nprograms:\n claudeCode:\n auth: subscription\n claudeDesktop: {}\n codex:\n auth: subscription\n", ) @@ -1995,7 +1985,7 @@ programs: ); assert_eq!( config.programs.subscription_auth("codex"), - Some(super::SubscriptionAuth::ClientPassthrough) + Some(super::SubscriptionAuth::CredentialHelper) ); for client_id in ["claude-desktop", "copilot-cli", "opencode", "unknown"] { assert_eq!(config.programs.subscription_auth(client_id), None); @@ -2007,7 +1997,7 @@ programs: } #[test] - fn codex_subscription_requires_an_authenticated_proxy_gateway() { + fn codex_subscription_requires_an_authenticated_gateway_without_a_proxy() { let valid = r#" llmGateway: url: https://gateway.example.com @@ -2021,6 +2011,8 @@ programs: auth: subscription "#; let config = parse_daemon(valid).expect("valid Codex subscription configuration"); + parse_daemon(&valid.replace(" proxyUrl: https://gateway.example.com/codex\n", "")) + .expect("Codex helper does not require proxyUrl"); assert_eq!( config.programs.codex.unwrap().auth, Some(super::ProgramAuthentication::Subscription) @@ -2034,10 +2026,6 @@ programs: valid.replace(" auth: subscription", " auth: subscription\n useLlmGateway: false"), "use an LLM gateway", ), - ( - valid.replace(" proxyUrl: https://gateway.example.com/codex\n", ""), - "proxyUrl", - ), ( valid.replace( " authentication:\n type: controllerJwt\n audience: agentgateway\n allowedClientIds: [codex]\n", diff --git a/examples/codex-subscription/README.md b/examples/codex-subscription/README.md index 4762eecb..30f7bca0 100644 --- a/examples/codex-subscription/README.md +++ b/examples/codex-subscription/README.md @@ -1,11 +1,12 @@ # Codex subscription through Agentdesktop -Use Codex's existing ChatGPT login for inference while separately authenticating -to the organization's gateway. Agentdesktop manages the Codex provider, local -proxy route, and pairing header; Codex owns ChatGPT sign-in and token refresh. +Connect ChatGPT through Agentdesktop's subscription OAuth flow, just like Claude. +Agentdesktop owns subscription sign-in, secure token storage, refresh, and the +Codex credential helper. Codex connects directly to the gateway; no local LLM +proxy, pairing header, or separate `codex login` is needed. -This requires a user-mode daemon. The Windows system service cannot run the -per-user proxy. Update both the controller and daemon before distributing +The example runs as the current user. System mode uses the same credential-helper +contract, with credentials belonging to the daemon's account. Update both controller and daemon before distributing `programs.codex.auth: subscription` to a fleet. ## Run @@ -18,9 +19,10 @@ agentdesktop daemon --user --config examples/codex-subscription/config.yaml ``` Complete organization sign-in in the browser. The local Dex account is -`admin@example.com` with password `password`. Codex must also be signed in to -ChatGPT using its normal `codex login` flow. Run Codex in another terminal after -the daemon applies the configuration. +`admin@example.com` with password `password`. Continue to the optional subscription +step and sign in to ChatGPT. The Agentdesktop page uses port 51327 and OpenAI's +registered callback uses port 1455; keep both available during sign-in. Run Codex +in another terminal after the daemon applies the configuration. Agentdesktop manages `~/.codex/config.toml` in user mode. A pre-existing file without Agentdesktop's ownership header is reported as a conflict and is not @@ -29,48 +31,70 @@ and resolve the reported conflict before testing; do not delete Codex's separate login file. A `daemon.codex.config` override must point to a configuration layer the installed Codex version actually reads, not `requirements.toml`. -The reconciler generates a provider with `requires_openai_auth = true`, forces -ChatGPT login, disables WebSockets for this HTTP-only route, and writes the -pairing header into an owner-only configuration file. It does not write a -gateway credential command or copy ChatGPT tokens into the file. +The reconciler generates a Responses provider with `requires_openai_auth = false` +and an Agentdesktop credential command. It never writes subscription tokens to +Codex configuration or reads Codex's own login file. ## Credential flow ```text -Codex -> local /codex proxy -> gateway -> ChatGPT Codex backend +Codex -> gateway -> ChatGPT Codex backend ``` -The local proxy puts gateway identity in `Authorization` and Codex's bearer in -`x-llm-token`. The gateway strictly validates organization identity first, then -restores the ChatGPT bearer and strips `x-llm-token`. The gateway forwards the -account-identification headers and maps `/codex-subscription/responses` to -`https://chatgpt.com/backend-api/codex/responses`, not the API-key endpoint at -`api.openai.com/v1`. +The helper supplies a composite bearer in `Authorization`: + +```text +agentdesktop::: +``` + +The first three fields match Claude's format; the fourth carries the ChatGPT +account ID. The gateway strictly validates the identity in the third field first, +then restores the provider bearer and `ChatGPT-Account-Id` together with one +header-replacement expression, preserving other request headers. + +The example uses Agentgateway 1.6.0's **`llm` configuration**, not a raw HTTP +route or opaque passthrough. Codex sends `/v1/responses` to the LLM listener. +The OpenAI provider's `params.baseUrl` points at +`https://chatgpt.com/backend-api/codex`, so the upstream request reaches +`/backend-api/codex/responses`, not the API-key endpoint at `api.openai.com/v1`. +Normal model routing, LLM policies, and token-usage processing remain active. +Add rate limits or guardrails through the usual LLM configuration as needed. For controller-issued gateway credentials, authorize client ID `codex` and set the same subscription configuration in the controller's effective policy. The -controller then marks these requests as subscription-billed. OIDC-only deployments -must classify subscription traffic using the dedicated gateway route when -collecting usage. - -The proxy must be enabled locally with `daemon.llmProxy.listen`; `proxyUrl` and -the program configuration belong in the controller-distributed policy on enrolled -devices. `llmGateway.githubOAuth` is not required for the managed `/codex` route. -When the proxy is unavailable, the existing `whenProxyUnavailable` policy applies: -`failClosed` retains existing managed settings; `failOpen` removes the managed -file. Before the first successful apply there is no managed configuration to keep. +controller then marks these requests as subscription-billed. The dedicated LLM +listener also labels its access logs with `agentdesktop.billing: subscription`, +including for OIDC-only deployments. Keep that label when adding an OTLP usage +exporter; exporting to Agentdesktop still requires the usage configuration +described in the main README. + +On enrolled devices, put `llmGateway` and `programs` in controller-distributed +policy. Set `llmGateway.url` to the LLM listener's `/v1` base URL. `proxyUrl`, `githubOAuth`, +`daemon.llmProxy.listen`, and `whenProxyUnavailable` are not used for this mode. +If subscription sign-in is skipped, the helper supplies identity alone; this +example's strict LLM authentication policy rejects it rather than using a paid fallback. ## Validation and Production The integration test runs the real Codex 0.153.4 CLI with synthetic ChatGPT/OIDC -credentials against a mock gateway and verifies streamed responses, separated -credentials, account headers, and cleanup. It does not validate a paid ChatGPT +credentials in Agentdesktop's store against a mock gateway, with no Codex login +file and no proxy. It verifies direct helper authentication, the account metadata, +streamed responses, and cleanup. It does not validate a paid ChatGPT account against OpenAI or certify provider terms. Validate both login-refresh flows and actual subscription access before rollout. +The example was also checked against a real Agentgateway 1.6.0 LLM pipeline with +a local mock upstream: model mapping, Responses-specific fields, streaming, +provider/account headers, malformed-credential rejection, and Prometheus token +metrics were verified. A fixture reporting 10 input and 5 output tokens produced +those counts in the gateway's LLM metrics. This does not establish live OpenAI +subscription compatibility or provider terms. + This local example intentionally uses loopback HTTP and a demonstration Dex account. Use HTTPS and your production identity provider outside local testing. -Protect pairing and credential files with user-only permissions or Windows ACLs, +This adapter follows Codex 0.153.4's public OAuth client registration and protocol; +they are provider-controlled and are not a guaranteed third-party API. +Protect credential files with user-only permissions or Windows ACLs, keep tokens out of logs, and review provider terms and credential-handling policy. Stop the example services with: diff --git a/examples/codex-subscription/agentgateway.yaml b/examples/codex-subscription/agentgateway.yaml index f57fe118..175046e5 100644 --- a/examples/codex-subscription/agentgateway.yaml +++ b/examples/codex-subscription/agentgateway.yaml @@ -1,36 +1,34 @@ -gateways: - gateway: - port: 4002 +# yaml-language-server: $schema=https://raw.githubusercontent.com/agentgateway/agentgateway/v1.6.0/schema/config.json -routes: -- name: codex-subscription - gateways: gateway - matches: - - path: - pathPrefix: /codex-subscription +frontendPolicies: + accessLog: + add: + agentdesktop.billing: '"subscription"' + +llm: + port: 4002 policies: jwtAuth: mode: strict + location: + expression: | + request.headers["authorization"].with(credential, + credential.matches("^Bearer agentdesktop:[^: ]+:[^: ]+:[^: ]+$") + ? credential.split(":")[2] : "") issuer: http://127.0.0.1:5558/dex audiences: [agentdesktop-local] jwks: url: http://dex:5558/dex/keys - urlRewrite: - authority: - full: chatgpt.com - path: - prefix: /backend-api/codex - timeout: - requestTimeout: 120s - backends: - - host: chatgpt.com:443 - policies: - backendTLS: {} - backendAuth: - passthrough: {} - transformations: - request: - set: - authorization: '"Bearer " + request.headers["x-llm-token"]' - remove: - - x-llm-token + transformations: + request: + replace: | + request.headers["authorization"].split(" ")[1].split(":").with(parts, + request.headers.merge({ + "authorization": "Bearer " + parts[1], + "chatgpt-account-id": parts[3] + })) + models: + - name: "*" + provider: openai + params: + baseUrl: https://chatgpt.com/backend-api/codex diff --git a/examples/codex-subscription/compose.yaml b/examples/codex-subscription/compose.yaml index 0e434afd..93a53088 100644 --- a/examples/codex-subscription/compose.yaml +++ b/examples/codex-subscription/compose.yaml @@ -8,7 +8,7 @@ services: - ./dex.yaml:/etc/dex/config.yaml:ro agentgateway: - image: cr.agentgateway.dev/agentgateway:v1.4.1 + image: cr.agentgateway.dev/agentgateway:v1.6.0 command: -f /etc/agentgateway/config.yaml depends_on: - dex diff --git a/examples/codex-subscription/config.yaml b/examples/codex-subscription/config.yaml index d012e7ed..402c8019 100644 --- a/examples/codex-subscription/config.yaml +++ b/examples/codex-subscription/config.yaml @@ -1,11 +1,8 @@ daemon: user: true - llmProxy: - listen: 127.0.0.1:4012 llmGateway: - url: http://127.0.0.1:4002 - proxyUrl: http://127.0.0.1:4002/codex-subscription + url: http://127.0.0.1:4002/v1 authentication: type: oidc issuer: http://127.0.0.1:5558/dex diff --git a/frontend/controller/package.json b/frontend/controller/package.json index 071214be..1fece4bc 100644 --- a/frontend/controller/package.json +++ b/frontend/controller/package.json @@ -17,8 +17,7 @@ "@agentdesktop/ui": "workspace:*", "lucide-react": "^1.55.0", "react": "^19.3.0", - "react-dom": "^19.3.0", - "yaml": "^2.9.1" + "react-dom": "^19.3.0" }, "devDependencies": { "@biomejs/biome": "^2.5.15", diff --git a/frontend/controller/src/views/ConfigurationView.stories.tsx b/frontend/controller/src/views/ConfigurationView.stories.tsx index 6ccefee6..d5f42f5b 100644 --- a/frontend/controller/src/views/ConfigurationView.stories.tsx +++ b/frontend/controller/src/views/ConfigurationView.stories.tsx @@ -75,9 +75,7 @@ export const CodexSubscription: Story = { args: { initialConfig: { llmGateway: { - url: "https://gateway.example.com", - proxyUrl: "https://gateway.example.com/chatgpt", - whenProxyUnavailable: "failOpen", + url: "https://gateway.example.com/chatgpt", authentication: { type: "controllerJwt", audience: "agentgateway" }, }, programs: { codex: { auth: "subscription" } }, @@ -85,13 +83,13 @@ export const CodexSubscription: Story = { }, play: async ({ canvas, canvasElement }) => { await expect( - canvas.getByRole("checkbox", { + canvas.queryByRole("checkbox", { name: /Fail open without the local proxy/, }), - ).toBeChecked(); + ).toBeNull(); const output = canvasElement.querySelector(".output-card code"); await expect(output?.textContent).toContain('auth: "subscription"'); - await expect(output).toHaveTextContent("whenProxyUnavailable: failOpen"); + await expect(output).not.toHaveTextContent("whenProxyUnavailable"); await expect(output).toHaveTextContent( "https://gateway.example.com/chatgpt", ); diff --git a/frontend/controller/src/views/ConfigurationView.tsx b/frontend/controller/src/views/ConfigurationView.tsx index 0a0bc153..77742929 100644 --- a/frontend/controller/src/views/ConfigurationView.tsx +++ b/frontend/controller/src/views/ConfigurationView.tsx @@ -8,7 +8,6 @@ import { Trash2, } from "lucide-react"; import { useEffect, useRef, useState } from "react"; -import { parse } from "yaml"; import type { AgentDraft, AgentKind, DaemonConfigDocument } from "../types"; @@ -49,15 +48,7 @@ export function ConfigurationView({ return [definition?.label ?? agent.kind]; }); const sandboxUnavailable = incompatibleSandboxAgentNames.length > 0; - const usesLocalProxy = agents.some((agent) => { - if (proxyAgents.has(agent.kind)) return true; - if (agent.kind !== "codex" || !agent.useGateway) return false; - try { - return parse(agent.settings)?.auth === "subscription"; - } catch { - return false; - } - }); + const usesLocalProxy = agents.some((agent) => proxyAgents.has(agent.kind)); const yaml = daemonConfigYaml({ gateway, gatewayUrl, diff --git a/frontend/pnpm-lock.yaml b/frontend/pnpm-lock.yaml index d5a9c13b..a85b742e 100644 --- a/frontend/pnpm-lock.yaml +++ b/frontend/pnpm-lock.yaml @@ -22,9 +22,6 @@ importers: react-dom: specifier: ^19.3.0 version: 19.3.0(react@19.3.0) - yaml: - specifier: ^2.9.1 - version: 2.9.1 devDependencies: '@biomejs/biome': specifier: ^2.5.15 @@ -3539,4 +3536,5 @@ snapshots: yallist@3.1.1: {} - yaml@2.9.1: {} + yaml@2.9.1: + optional: true diff --git a/schema/daemon-config.json b/schema/daemon-config.json index 342b7c1f..ce3aab74 100644 --- a/schema/daemon-config.json +++ b/schema/daemon-config.json @@ -138,7 +138,7 @@ "type": "object", "properties": { "auth": { - "description": "Keep Codex's ChatGPT subscription login and route it through the local proxy.\nRequires user mode, gateway authentication, and `llmGateway.proxyUrl`.", + "description": "Connect a ChatGPT subscription through agentdesktop's OAuth credential helper.", "anyOf": [ { "$ref": "#/$defs/ProgramAuthentication" @@ -519,7 +519,7 @@ "type": "string" }, "whenProxyUnavailable": { - "description": "What the programs behind the local LLM proxy (`programs.copilot`,\n`programs.vscode`, or Codex with `auth: subscription`) do while the daemon runs without that proxy:\n`failClosed` (the default) leaves their managed entries in place, so\nthe tools stay pointed at the loopback port and fail instead\nof falling back to GitHub past the gateway (with no entries in place\nyet, the tools keep their own settings); `failOpen` removes them\nuntil the proxy is back (VS Code on `githubModels` then talks to GitHub\ndirectly, the Copilot CLI and VS Code on `ownModels` lose the gateway's\nmodels).", + "description": "What the programs behind the local LLM proxy (`programs.copilot`,\n`programs.vscode`) do while the daemon runs without that proxy:\n`failClosed` (the default) leaves their managed entries in place, so\nthe tools stay pointed at the loopback port and fail instead\nof falling back to GitHub past the gateway (with no entries in place\nyet, the tools keep their own settings); `failOpen` removes them\nuntil the proxy is back (VS Code on `githubModels` then talks to GitHub\ndirectly, the Copilot CLI and VS Code on `ownModels` lose the gateway's\nmodels).", "$ref": "#/$defs/ProxyUnavailable" } }, diff --git a/schema/daemon-config.md b/schema/daemon-config.md index 5e6434a8..f16cb82d 100644 --- a/schema/daemon-config.md +++ b/schema/daemon-config.md @@ -54,7 +54,7 @@ |`llmGateway.githubOAuth.source`|enum|Where the GitHub credential comes from.
Possible values: `deviceFlow`, `request`.| |`llmGateway.proxyUrl`|string|Base URL the local LLM proxy forwards to, when it differs from `url`.

`url` is shared by every program that sets `useLlmGateway`, so it can
only carry one path prefix. A gateway that puts each provider behind its
own prefix therefore cannot serve both a program and the proxy from one
value. Setting this leaves `url` to the programs and gives the proxy its
own target. The same rules as for `url` apply.| |`llmGateway.url`|string|Base HTTP or HTTPS URL of the LLM gateway.

The URL must include a host and cannot include credentials, a query, or a fragment.| -|`llmGateway.whenProxyUnavailable`|enum|What the programs behind the local LLM proxy (`programs.copilot`,
`programs.vscode`, or Codex with `auth: subscription`) do while the daemon runs without that proxy:
`failClosed` (the default) leaves their managed entries in place, so
the tools stay pointed at the loopback port and fail instead
of falling back to GitHub past the gateway (with no entries in place
yet, the tools keep their own settings); `failOpen` removes them
until the proxy is back (VS Code on `githubModels` then talks to GitHub
directly, the Copilot CLI and VS Code on `ownModels` lose the gateway's
models).
Possible values: `failClosed`, `failOpen`.| +|`llmGateway.whenProxyUnavailable`|enum|What the programs behind the local LLM proxy (`programs.copilot`,
`programs.vscode`) do while the daemon runs without that proxy:
`failClosed` (the default) leaves their managed entries in place, so
the tools stay pointed at the loopback port and fail instead
of falling back to GitHub past the gateway (with no entries in place
yet, the tools keep their own settings); `failOpen` removes them
until the proxy is back (VS Code on `githubModels` then talks to GitHub
directly, the Copilot CLI and VS Code on `ownModels` lose the gateway's
models).
Possible values: `failClosed`, `failOpen`.| |`programSelection`|object|Whether the device user may choose which configured programs are managed.

A daemon that predates this setting rejects a configuration that
contains it, so upgrade devices before setting it in
controller-delivered policy.| |`programSelection.allowUserChoice`|boolean|Let the device user opt out of all configured programs, or choose the
ones to manage, with `agentdesktop programs set`.
Agentdesktop removes the settings it wrote for a program the user opts
out of, and reports the program as opted out.

Defaults to `false`: every configured program is managed and a choice
stored on the device is ignored until this is enabled again.| |`programs`|object|Per-program settings reconciled on this device.| @@ -65,7 +65,7 @@ |`programs.claudeDesktop.auth`|enum|Upstream authentication used by this agent.
Possible values: `subscription`.| |`programs.claudeDesktop.useLlmGateway`|boolean|Whether this program uses the top-level LLM gateway.| |`programs.codex`|object|Codex managed configuration.| -|`programs.codex.auth`|enum|Keep Codex's ChatGPT subscription login and route it through the local proxy.
Requires user mode, gateway authentication, and `llmGateway.proxyUrl`.
Possible values: `subscription`.| +|`programs.codex.auth`|enum|Connect a ChatGPT subscription through Agentdesktop's OAuth credential helper.
Possible values: `subscription`.| |`programs.codex.managedConfig`|object|Arbitrary values written to Codex's organization-managed TOML configuration.

Use Codex's native snake_case configuration keys. TOML has no null value,
so null values cannot be reconciled.| |`programs.codex.useLlmGateway`|boolean|Whether this program uses the top-level LLM gateway.| |`programs.copilot`|object|GitHub Copilot CLI managed configuration.| From fcd819cd80bab9da03461235ae75cf58f46bccb5 Mon Sep 17 00:00:00 2001 From: Peter Jausovec Date: Mon, 19 Oct 2026 17:50:33 -0700 Subject: [PATCH 3/6] schema fix Signed-off-by: Peter Jausovec --- schema/daemon-config.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/schema/daemon-config.json b/schema/daemon-config.json index ce3aab74..084035cb 100644 --- a/schema/daemon-config.json +++ b/schema/daemon-config.json @@ -138,7 +138,7 @@ "type": "object", "properties": { "auth": { - "description": "Connect a ChatGPT subscription through agentdesktop's OAuth credential helper.", + "description": "Connect a ChatGPT subscription through Agentdesktop's OAuth credential helper.", "anyOf": [ { "$ref": "#/$defs/ProgramAuthentication" From 899df9bcc20885fe32be6ff969a1884876e171c0 Mon Sep 17 00:00:00 2001 From: Peter Jausovec Date: Tue, 20 Oct 2026 12:08:23 -0700 Subject: [PATCH 4/6] pr feedback Signed-off-by: Peter Jausovec --- crates/agent/src/oidc.rs | 65 +++++++++++++++++--- crates/agent/src/provider/claude_code/mod.rs | 1 + crates/agent/src/provider/codex/mod.rs | 2 + crates/agent/src/provider/codex/reconcile.rs | 56 ++++++++++++++--- crates/agent/src/provider/copilot/tests.rs | 1 + crates/agent/src/provider/mod.rs | 2 + crates/agent/src/provider/vscode/tests.rs | 2 + crates/agent/src/reconcile/mod.rs | 1 + crates/agent/src/subscription_oauth/mod.rs | 45 ++++++++------ 9 files changed, 141 insertions(+), 34 deletions(-) diff --git a/crates/agent/src/oidc.rs b/crates/agent/src/oidc.rs index b969b94b..b6efd245 100644 --- a/crates/agent/src/oidc.rs +++ b/crates/agent/src/oidc.rs @@ -828,15 +828,15 @@ pub(crate) async fn continue_subscription_page( .route("/flow-ready", get(continued_flow_ready)) .with_state(state); tokio::spawn(async move { - let result = axum::serve(listener, app) + axum::serve(listener, app) .with_graceful_shutdown(async move { let _ = tokio::time::timeout(CALLBACK_TIMEOUT, viewed_receiver).await; }) - .await; - if let Err(error) = result { - tracing::warn!(%error, "serve continued Agentdesktop authorization page"); - } - }); + .await + }) + .await + .context("join continued Agentdesktop authorization page")? + .context("serve continued Agentdesktop authorization page")?; Ok(()) } @@ -950,7 +950,7 @@ mod tests { use super::{ AuthorizationPage, DevicePoll, TokenResponse, apply_refreshed_tokens, - classify_device_token_error, poll_device_token, + classify_device_token_error, continue_subscription_page, poll_device_token, }; use crate::identity::{Identity, OAuthCredentials}; @@ -1134,6 +1134,57 @@ mod tests { .unwrap(); } + #[tokio::test] + async fn continued_subscription_page_waits_for_view_and_then_stops() { + tokio::time::timeout(Duration::from_secs(10), async { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + drop(listener); + let redirect = url::Url::parse("http://localhost:51327/callback").unwrap(); + let server = tokio::spawn(async move { + continue_subscription_page(&redirect, Some(address), true).await + }); + let client = reqwest::Client::new(); + loop { + if let Ok(response) = client + .get(format!("http://{address}/flow-ready")) + .send() + .await + && response.status().is_success() + { + break; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + assert!( + !server.is_finished(), + "listener must remain owned until page visit" + ); + let response = client + .get(format!("http://{address}/")) + .send() + .await + .unwrap(); + assert!( + response + .text() + .await + .unwrap() + .contains("Agentdesktop connected") + ); + server.await.unwrap().unwrap(); + assert!( + client + .get(format!("http://{address}/")) + .send() + .await + .is_err() + ); + }) + .await + .unwrap(); + } + #[test] fn identity_page_only_mentions_subscription_when_configured() { let without_subscription = AuthorizationPage::Identity { diff --git a/crates/agent/src/provider/claude_code/mod.rs b/crates/agent/src/provider/claude_code/mod.rs index 0c016487..93536bcc 100644 --- a/crates/agent/src/provider/claude_code/mod.rs +++ b/crates/agent/src/provider/claude_code/mod.rs @@ -140,6 +140,7 @@ mod tests { let context = ReconcileContext { llm_proxy: None, merge_user_settings: false, + user_mode: false, credential_helper: PathBuf::from(r"C:\Program Files\Agent Desktop\agentdesktop.exe"), socket: PathBuf::from(r"\\.\pipe\agentdesktop"), }; diff --git a/crates/agent/src/provider/codex/mod.rs b/crates/agent/src/provider/codex/mod.rs index cdd94ac4..0c9d29ad 100644 --- a/crates/agent/src/provider/codex/mod.rs +++ b/crates/agent/src/provider/codex/mod.rs @@ -54,6 +54,7 @@ impl Provider for Codex { &ctx.credential_helper, &ctx.socket, config.sandbox.as_ref(), + ctx.user_mode, ctx.llm_proxy .as_ref() .map(|proxy| (proxy.address, &*proxy.pairing)), @@ -81,6 +82,7 @@ mod tests { .unwrap(); let context = ReconcileContext { merge_user_settings: false, + user_mode: false, credential_helper: PathBuf::from("agentdesktop"), socket: PathBuf::from("socket"), llm_proxy: None, diff --git a/crates/agent/src/provider/codex/reconcile.rs b/crates/agent/src/provider/codex/reconcile.rs index bdd879c0..f4a0ccd1 100644 --- a/crates/agent/src/provider/codex/reconcile.rs +++ b/crates/agent/src/provider/codex/reconcile.rs @@ -23,6 +23,7 @@ pub(super) fn plan( credential_helper: &Path, socket: &Path, sandbox: Option<&SandboxConfig>, + user_mode: bool, _proxy: Option<(SocketAddr, &str)>, config: Option<(&CodexConfig, Option<&LlmGatewayConfig>)>, plan: &ReconcilePlan, @@ -32,7 +33,7 @@ pub(super) fn plan( }; let subscription = config.auth == Some(ProgramAuthentication::Subscription); - let mode = if subscription { + let mode = if subscription && user_mode { 0o600 } else { MANAGED_FILE_MODE @@ -178,7 +179,6 @@ fn managed_config( "wire_api": "responses", }); if config.auth == Some(ProgramAuthentication::Subscription) { - provider["base_url"] = json!(gateway.url.as_str().trim_end_matches('/')); provider["requires_openai_auth"] = json!(false); provider["supports_websockets"] = json!(false); settings @@ -330,7 +330,7 @@ programs: assert_eq!(settings["model"], "gpt-5.4"); assert!(settings.get("forced_login_method").is_none()); let provider = &settings["model_providers"]["agentdesktop"]; - assert_eq!(provider["base_url"], "https://gateway.example.com/llm"); + assert_eq!(provider["base_url"], "https://gateway.example.com/llm/v1"); assert_eq!(provider["requires_openai_auth"], false); assert_eq!(provider["supports_websockets"], false); assert_eq!(provider["wire_api"], "responses"); @@ -356,13 +356,14 @@ programs: let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("config.toml"); let mut config = subscription_config(); - let run_plan = |config: &agentdesktop_core::config::DaemonConfig, proxy| { + let run_plan = |config: &agentdesktop_core::config::DaemonConfig, proxy, user_mode| { let plan = crate::reconcile::ReconcilePlan::default(); super::plan( &path, Path::new("agentdesktop"), Path::new("agentdesktop.sock"), None, + user_mode, proxy, Some(( config.programs.codex.as_ref().unwrap(), @@ -373,12 +374,16 @@ programs: .unwrap(); plan }; - let absent = run_plan(&config, None); + let absent = run_plan(&config, None, true); assert!(absent.inactive_reason().is_none()); absent.apply().unwrap(); assert!(path.exists()); for pairing in ["first-pairing-secret", "second-pairing-secret"] { - let plan = run_plan(&config, Some(("127.0.0.1:4010".parse().unwrap(), pairing))); + let plan = run_plan( + &config, + Some(("127.0.0.1:4010".parse().unwrap(), pairing)), + true, + ); assert!(!plan.render().contains("pairing-secret")); plan.apply().unwrap(); assert!(!std::fs::read_to_string(&path).unwrap().contains(pairing)); @@ -393,19 +398,50 @@ programs: } } let contents = std::fs::read(&path).unwrap(); - let closed = run_plan(&config, None); + let closed = run_plan(&config, None, true); assert!(closed.inactive_reason().is_none()); closed.apply().unwrap(); assert_eq!(std::fs::read(&path).unwrap(), contents); config.llm_gateway.as_mut().unwrap().when_proxy_unavailable = agentdesktop_core::config::ProxyUnavailable::FailOpen; - let open = run_plan(&config, None); + let open = run_plan(&config, None, true); assert!(open.inactive_reason().is_none()); assert!(!open.render().contains("pairing-secret")); open.apply().unwrap(); assert!(path.exists()); } + #[cfg(unix)] + #[test] + fn system_mode_subscription_configuration_remains_readable_by_client_users() { + use std::os::unix::fs::PermissionsExt; + + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("managed_config.toml"); + let config = subscription_config(); + let plan = crate::reconcile::ReconcilePlan::default(); + super::plan( + &path, + Path::new("agentdesktop"), + Path::new("agentdesktop.sock"), + None, + false, + None, + Some(( + config.programs.codex.as_ref().unwrap(), + config.llm_gateway.as_ref(), + )), + &plan, + ) + .unwrap(); + plan.apply().unwrap(); + + assert_eq!( + std::fs::metadata(&path).unwrap().permissions().mode() & 0o777, + 0o644 + ); + } + #[test] fn subscription_configuration_is_redacted_on_conflict_and_when_disabled() { let dir = tempfile::tempdir().unwrap(); @@ -418,6 +454,7 @@ programs: Path::new("agentdesktop"), Path::new("agentdesktop.sock"), None, + false, Some(("127.0.0.1:4010".parse().unwrap(), "pairing-secret")), Some(( config.programs.codex.as_ref().unwrap(), @@ -446,6 +483,7 @@ programs: Path::new("agentdesktop"), Path::new("socket"), None, + false, None, None, &plan, @@ -587,6 +625,7 @@ programs: Path::new("agentdesktop"), Path::new("agentdesktop.sock"), None, + false, None, Some((config.programs.codex.as_ref().unwrap(), None)), plan, @@ -610,6 +649,7 @@ programs: Path::new("agentdesktop"), Path::new("agentdesktop.sock"), None, + false, None, Some((config.programs.codex.as_ref().unwrap(), None)), plan, diff --git a/crates/agent/src/provider/copilot/tests.rs b/crates/agent/src/provider/copilot/tests.rs index 4de39f1c..10270711 100644 --- a/crates/agent/src/provider/copilot/tests.rs +++ b/crates/agent/src/provider/copilot/tests.rs @@ -766,6 +766,7 @@ fn system_mode_rejects_a_configured_program_before_writing_anything() { .unwrap(); let context = |merge_user_settings: bool| ReconcileContext { merge_user_settings, + user_mode: merge_user_settings, credential_helper: dir.path().join("agentdesktop"), socket: dir.path().join("agentdesktop.sock"), llm_proxy: None, diff --git a/crates/agent/src/provider/mod.rs b/crates/agent/src/provider/mod.rs index b22310d3..d54d648d 100644 --- a/crates/agent/src/provider/mod.rs +++ b/crates/agent/src/provider/mod.rs @@ -46,6 +46,8 @@ pub trait Provider: Send + Sync { #[derive(Clone)] pub struct ReconcileContext { pub merge_user_settings: bool, + /// Whether the daemon is running in user mode rather than managing system files. + pub user_mode: bool, pub credential_helper: PathBuf, pub socket: PathBuf, /// The bound loopback LLM proxy, for programs whose client files point at diff --git a/crates/agent/src/provider/vscode/tests.rs b/crates/agent/src/provider/vscode/tests.rs index 91d37bbe..0aee5def 100644 --- a/crates/agent/src/provider/vscode/tests.rs +++ b/crates/agent/src/provider/vscode/tests.rs @@ -658,6 +658,7 @@ fn system_mode_rejects_a_configured_program_before_writing_anything() { .unwrap(); let context = |merge_user_settings: bool| ReconcileContext { merge_user_settings, + user_mode: merge_user_settings, credential_helper: dir.path().join("agentdesktop"), socket: dir.path().join("agentdesktop.sock"), llm_proxy: None, @@ -709,6 +710,7 @@ fn github_models_through_the_provider_removes_the_chat_models_entry_and_writes_s let pairing: std::sync::Arc = "PAIRING-FIXTURE".into(); let context = ReconcileContext { merge_user_settings: true, + user_mode: true, credential_helper: dir.path().join("agentdesktop"), socket: dir.path().join("agentdesktop.sock"), llm_proxy: Some(crate::llm_proxy::LlmProxyContext { diff --git a/crates/agent/src/reconcile/mod.rs b/crates/agent/src/reconcile/mod.rs index f4449fbd..622c248e 100644 --- a/crates/agent/src/reconcile/mod.rs +++ b/crates/agent/src/reconcile/mod.rs @@ -70,6 +70,7 @@ impl Reconciler { Self { context: ReconcileContext { merge_user_settings, + user_mode: merge_user_settings, credential_helper, socket, llm_proxy: None, diff --git a/crates/agent/src/subscription_oauth/mod.rs b/crates/agent/src/subscription_oauth/mod.rs index 58128346..2570ee69 100644 --- a/crates/agent/src/subscription_oauth/mod.rs +++ b/crates/agent/src/subscription_oauth/mod.rs @@ -17,6 +17,8 @@ pub(crate) mod openai; const PROMPT_URI: &str = "http://localhost:51327/callback"; const EXPIRY_SKEW_SECONDS: u64 = 60; +static LISTENER_LOCK: Mutex<()> = Mutex::const_new(()); + struct LoginState { lock: Mutex<()>, declined: AtomicBool, @@ -90,6 +92,7 @@ async fn credential_inner( }; if let Some(token) = ready { if !open_browser { + let _listener = LISTENER_LOCK.lock().await; oidc::continue_subscription_page(&prompt, callback_listen, true).await?; } return Ok(Some(provider.credential_from_token(token))); @@ -97,6 +100,7 @@ async fn credential_inner( } if login.declined.load(Ordering::Relaxed) { if !open_browser { + let _listener = LISTENER_LOCK.lock().await; oidc::continue_subscription_page(&prompt, callback_listen, false).await?; } return Ok(None); @@ -106,25 +110,28 @@ async fn credential_inner( let state = oidc::random_secret(); let (verifier, challenge) = oidc::pkce(); let authorization_url = provider.authorization_url(&redirect, &state, &challenge); - let code = if redirect == prompt { - oidc::wait_for_authorization_code_with_page( - authorization_url.as_str(), - &redirect, - state.clone(), - callback_listen, - oidc::AuthorizationPage::Subscription, - open_browser, - ) - .await? - } else { - oidc::wait_for_subscription_authorization_code( - authorization_url.as_str(), - &redirect, - state.clone(), - callback_listen, - open_browser, - ) - .await? + let code = { + let _listener = LISTENER_LOCK.lock().await; + if redirect == prompt { + oidc::wait_for_authorization_code_with_page( + authorization_url.as_str(), + &redirect, + state.clone(), + callback_listen, + oidc::AuthorizationPage::Subscription, + open_browser, + ) + .await? + } else { + oidc::wait_for_subscription_authorization_code( + authorization_url.as_str(), + &redirect, + state.clone(), + callback_listen, + open_browser, + ) + .await? + } }; let Some(code) = code else { login.declined.store(true, Ordering::Relaxed); From d324484659c3dc1185bd442122f107ff43061e70 Mon Sep 17 00:00:00 2001 From: Peter Jausovec Date: Wed, 21 Oct 2026 09:11:29 -0700 Subject: [PATCH 5/6] fix ci failure Signed-off-by: Peter Jausovec --- crates/agent/src/provider/codex/reconcile.rs | 36 ++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/crates/agent/src/provider/codex/reconcile.rs b/crates/agent/src/provider/codex/reconcile.rs index f4a0ccd1..95899e7a 100644 --- a/crates/agent/src/provider/codex/reconcile.rs +++ b/crates/agent/src/provider/codex/reconcile.rs @@ -18,6 +18,15 @@ const MANAGED_FILE_MODE: u32 = 0o644; const MANAGED_HEADER: &str = "# Managed by Agentdesktop. Manual changes will be replaced.\n"; +fn codex_responses_base_url(gateway: &LlmGatewayConfig, subscription: bool) -> String { + let path = gateway.url.path().trim_end_matches('/'); + if subscription && path.ends_with("/backend-api/codex") { + gateway.url.to_string().trim_end_matches('/').to_owned() + } else { + responses_base_url(gateway) + } +} + pub(super) fn plan( path: &Path, credential_helper: &Path, @@ -172,13 +181,14 @@ fn managed_config( return Ok(settings); }; + let subscription = config.auth == Some(ProgramAuthentication::Subscription); let provider_name = "agentdesktop"; let mut provider = json!({ "name": "Agentdesktop", - "base_url": responses_base_url(gateway), + "base_url": codex_responses_base_url(gateway, subscription), "wire_api": "responses", }); - if config.auth == Some(ProgramAuthentication::Subscription) { + if subscription { provider["requires_openai_auth"] = json!(false); provider["supports_websockets"] = json!(false); settings @@ -351,6 +361,28 @@ programs: assert_eq!(parsed["model_provider"].as_str(), Some("agentdesktop")); } + #[test] + fn subscription_provider_preserves_the_chatgpt_codex_endpoint() { + let mut config = subscription_config(); + config.llm_gateway.as_mut().unwrap().url = "https://gateway.example.com/backend-api/codex" + .parse() + .unwrap(); + let settings = managed_config( + config.programs.codex.as_ref().unwrap(), + config.llm_gateway.as_ref(), + Path::new("agentdesktop"), + Path::new("agentdesktop.sock"), + None, + None, + ) + .unwrap(); + + assert_eq!( + settings["model_providers"]["agentdesktop"]["base_url"], + "https://gateway.example.com/backend-api/codex" + ); + } + #[test] fn subscription_lifecycle_is_independent_of_proxy_availability() { let dir = tempfile::tempdir().unwrap(); From 426b819eb361b4a76c7711dd5fec159b5d74a379 Mon Sep 17 00:00:00 2001 From: Peter Jausovec Date: Wed, 21 Oct 2026 09:24:51 -0700 Subject: [PATCH 6/6] fix linter issues Signed-off-by: Peter Jausovec --- crates/agent/src/provider/codex/mod.rs | 3 -- crates/agent/src/provider/codex/reconcile.rs | 39 +++++++++----------- 2 files changed, 17 insertions(+), 25 deletions(-) diff --git a/crates/agent/src/provider/codex/mod.rs b/crates/agent/src/provider/codex/mod.rs index 0c9d29ad..a80c3100 100644 --- a/crates/agent/src/provider/codex/mod.rs +++ b/crates/agent/src/provider/codex/mod.rs @@ -55,9 +55,6 @@ impl Provider for Codex { &ctx.socket, config.sandbox.as_ref(), ctx.user_mode, - ctx.llm_proxy - .as_ref() - .map(|proxy| (proxy.address, &*proxy.pairing)), configured, &plan, )?; diff --git a/crates/agent/src/provider/codex/reconcile.rs b/crates/agent/src/provider/codex/reconcile.rs index 95899e7a..50c185fb 100644 --- a/crates/agent/src/provider/codex/reconcile.rs +++ b/crates/agent/src/provider/codex/reconcile.rs @@ -33,7 +33,6 @@ pub(super) fn plan( socket: &Path, sandbox: Option<&SandboxConfig>, user_mode: bool, - _proxy: Option<(SocketAddr, &str)>, config: Option<(&CodexConfig, Option<&LlmGatewayConfig>)>, plan: &ReconcilePlan, ) -> anyhow::Result<()> { @@ -389,22 +388,20 @@ programs: let path = dir.path().join("config.toml"); let mut config = subscription_config(); let run_plan = |config: &agentdesktop_core::config::DaemonConfig, proxy, user_mode| { - let plan = crate::reconcile::ReconcilePlan::default(); - super::plan( - &path, - Path::new("agentdesktop"), - Path::new("agentdesktop.sock"), - None, + use crate::provider::{Provider, ReconcileContext}; + + let context = ReconcileContext { + merge_user_settings: false, user_mode, - proxy, - Some(( - config.programs.codex.as_ref().unwrap(), - config.llm_gateway.as_ref(), - )), - &plan, - ) - .unwrap(); - plan + credential_helper: Path::new("agentdesktop").to_path_buf(), + socket: Path::new("agentdesktop.sock").to_path_buf(), + llm_proxy: proxy, + }; + super::Codex { + managed_config_path: path.clone(), + } + .plan(&context, config) + .unwrap() }; let absent = run_plan(&config, None, true); assert!(absent.inactive_reason().is_none()); @@ -413,7 +410,10 @@ programs: for pairing in ["first-pairing-secret", "second-pairing-secret"] { let plan = run_plan( &config, - Some(("127.0.0.1:4010".parse().unwrap(), pairing)), + Some(crate::llm_proxy::LlmProxyContext { + address: "127.0.0.1:4010".parse().unwrap(), + pairing: pairing.into(), + }), true, ); assert!(!plan.render().contains("pairing-secret")); @@ -458,7 +458,6 @@ programs: Path::new("agentdesktop.sock"), None, false, - None, Some(( config.programs.codex.as_ref().unwrap(), config.llm_gateway.as_ref(), @@ -487,7 +486,6 @@ programs: Path::new("agentdesktop.sock"), None, false, - Some(("127.0.0.1:4010".parse().unwrap(), "pairing-secret")), Some(( config.programs.codex.as_ref().unwrap(), config.llm_gateway.as_ref(), @@ -517,7 +515,6 @@ programs: None, false, None, - None, &plan, ) .unwrap(); @@ -658,7 +655,6 @@ programs: Path::new("agentdesktop.sock"), None, false, - None, Some((config.programs.codex.as_ref().unwrap(), None)), plan, ) @@ -682,7 +678,6 @@ programs: Path::new("agentdesktop.sock"), None, false, - None, Some((config.programs.codex.as_ref().unwrap(), None)), plan, )