diff --git a/README.md b/README.md index 3bf0e89e..9d9136ba 100644 --- a/README.md +++ b/README.md @@ -94,11 +94,49 @@ 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 -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, 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 `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 +157,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 +180,21 @@ 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`). 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 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. +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 +203,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 +236,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 +244,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 +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` | 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) | @@ -229,7 +270,8 @@ 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. 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 @@ -766,10 +808,11 @@ 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 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. ```yaml frontendPolicies: @@ -797,7 +840,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 +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. + 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 e2f4fd07..ab2db468 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, @@ -527,6 +551,7 @@ pub(crate) async fn gateway_credential( &state.state_dir, state.oidc_callback_listen, continue_in_browser, + client_id, ) .await } else { @@ -544,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, @@ -669,19 +698,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..8b22fcb8 100644 --- a/crates/agent/src/daemon.rs +++ b/crates/agent/src/daemon.rs @@ -951,15 +951,18 @@ 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_clients: Vec<_> = ["claude-code", "claude-desktop", "codex"] + .into_iter() + .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; } @@ -1016,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/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/oidc.rs b/crates/agent/src/oidc.rs index bac61be7..b6efd245 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 { @@ -754,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(()) } @@ -876,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}; @@ -991,6 +1065,126 @@ 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(); + } + + #[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/integration_tests.rs b/crates/agent/src/provider/codex/integration_tests.rs index a1b67239..0aadeec9 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,94 @@ 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 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 entry = Sha256::digest(b"dev.agentdesktop.openai-subscription\0codex"); + let entry = entry.iter().map(|byte| format!("{byte:02x}")).collect::(); + container.write( + &format!("{state_dir}/secrets/{entry}"), + &json!({ + "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, + }, + "llmGateway": { + "url": 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(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, + ], 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 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"); + 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", "!", "-e", "/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..a80c3100 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, configured, &plan, )?; @@ -65,3 +66,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_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", + ) + .unwrap(); + let context = ReconcileContext { + merge_user_settings: false, + user_mode: false, + credential_helper: PathBuf::from("agentdesktop"), + socket: PathBuf::from("socket"), + llm_proxy: None, + }; + 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 5f2c6956..50c185fb 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}; @@ -18,11 +18,21 @@ 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, socket: &Path, sandbox: Option<&SandboxConfig>, + user_mode: bool, config: Option<(&CodexConfig, Option<&LlmGatewayConfig>)>, plan: &ReconcilePlan, ) -> anyhow::Result<()> { @@ -30,7 +40,13 @@ 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); + let mode = if subscription && user_mode { + 0o600 + } else { + MANAGED_FILE_MODE + }; + 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) @@ -52,10 +68,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 +80,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")?; @@ -146,21 +180,48 @@ 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 subscription { + provider["requires_openai_auth"] = json!(false); + provider["supports_websockets"] = json!(false); + 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)) + .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); + } + } + } 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 @@ -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,232 @@ 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_uses_the_agentdesktop_oauth_helper_directly() { + 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!(settings.get("forced_login_method").is_none()); + let provider = &settings["model_providers"]["agentdesktop"]; + 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"); + assert_eq!(provider["auth"]["command"], "agentdesktop"); + assert_eq!(provider["auth"]["args"][4], "codex"); + assert_eq!(provider["auth"]["timeout_ms"], 600_000); + for field in [ + "http_headers", + "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_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(); + let path = dir.path().join("config.toml"); + let mut config = subscription_config(); + let run_plan = |config: &agentdesktop_core::config::DaemonConfig, proxy, user_mode| { + use crate::provider::{Provider, ReconcileContext}; + + let context = ReconcileContext { + merge_user_settings: false, + user_mode, + 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()); + absent.apply().unwrap(); + assert!(path.exists()); + for pairing in ["first-pairing-secret", "second-pairing-secret"] { + let plan = run_plan( + &config, + Some(crate::llm_proxy::LlmProxyContext { + address: "127.0.0.1:4010".parse().unwrap(), + pairing: pairing.into(), + }), + true, + ); + 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, 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, 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, + 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(); + 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, + false, + 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, + false, + 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 +555,7 @@ programs: Path::new("/usr/local/bin/agentdesktop"), Path::new("/run/agentdesktop/agentdesktop.sock"), None, + None, ) .expect("merged settings"); @@ -314,6 +602,7 @@ programs: Path::new("agentdesktop"), Path::new("agentdesktop.sock"), config.sandbox.as_ref(), + None, ) .expect("sandbox settings"); @@ -365,6 +654,7 @@ programs: Path::new("agentdesktop"), Path::new("agentdesktop.sock"), None, + false, Some((config.programs.codex.as_ref().unwrap(), None)), plan, ) @@ -387,6 +677,7 @@ programs: Path::new("agentdesktop"), Path::new("agentdesktop.sock"), None, + false, 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.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..2570ee69 --- /dev/null +++ b/crates/agent/src/subscription_oauth/mod.rs @@ -0,0 +1,320 @@ +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; + +static LISTENER_LOCK: Mutex<()> = Mutex::const_new(()); + +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 { + let _listener = LISTENER_LOCK.lock().await; + 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 { + let _listener = LISTENER_LOCK.lock().await; + 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 = { + 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); + 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 d4d15708..44611f42 100644 --- a/crates/agent/tests/common/gateway.rs +++ b/crates/agent/tests/common/gateway.rs @@ -74,16 +74,42 @@ 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 path == "/backend-api/codex/responses" { + 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 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(); } 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..f58fae26 100644 --- a/crates/core/src/config.rs +++ b/crates/core/src/config.rs @@ -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::CredentialHelper, + ), + _ => 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,9 @@ 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, + /// 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. /// /// Use Codex's native snake_case configuration keys. TOML has no null value, @@ -1473,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()) { @@ -1833,6 +1871,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 +1973,74 @@ programs: ); } + #[test] + 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", + ) + .unwrap(); + assert_eq!( + config.programs.subscription_auth("claude-code"), + Some(super::SubscriptionAuth::CredentialHelper) + ); + assert_eq!( + config.programs.subscription_auth("codex"), + Some(super::SubscriptionAuth::CredentialHelper) + ); + 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_gateway_without_a_proxy() { + 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"); + 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) + ); + 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( + " 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..30f7bca0 --- /dev/null +++ b/examples/codex-subscription/README.md @@ -0,0 +1,104 @@ +# Codex subscription through Agentdesktop + +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. + +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 + +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`. 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 +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 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 -> gateway -> ChatGPT Codex backend +``` + +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. 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 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. +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: + +```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..175046e5 --- /dev/null +++ b/examples/codex-subscription/agentgateway.yaml @@ -0,0 +1,34 @@ +# yaml-language-server: $schema=https://raw.githubusercontent.com/agentgateway/agentgateway/v1.6.0/schema/config.json + +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 + 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 new file mode 100644 index 00000000..93a53088 --- /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.6.0 + 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..402c8019 --- /dev/null +++ b/examples/codex-subscription/config.yaml @@ -0,0 +1,17 @@ +daemon: + user: true + +llmGateway: + url: http://127.0.0.1:4002/v1 + 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/src/views/ConfigurationView.stories.tsx b/frontend/controller/src/views/ConfigurationView.stories.tsx index 91c05611..d5f42f5b 100644 --- a/frontend/controller/src/views/ConfigurationView.stories.tsx +++ b/frontend/controller/src/views/ConfigurationView.stories.tsx @@ -71,6 +71,31 @@ export const CopilotFailOpen: Story = { }, }; +export const CodexSubscription: Story = { + args: { + initialConfig: { + llmGateway: { + url: "https://gateway.example.com/chatgpt", + authentication: { type: "controllerJwt", audience: "agentgateway" }, + }, + programs: { codex: { auth: "subscription" } }, + }, + }, + play: async ({ canvas, canvasElement }) => { + await expect( + canvas.queryByRole("checkbox", { + name: /Fail open without the local proxy/, + }), + ).toBeNull(); + const output = canvasElement.querySelector(".output-card code"); + await expect(output?.textContent).toContain('auth: "subscription"'); + await expect(output).not.toHaveTextContent("whenProxyUnavailable"); + await expect(output).toHaveTextContent( + "https://gateway.example.com/chatgpt", + ); + }, +}; + export const ActiveProgramSelection: Story = { args: { initialConfig: programSelectionDaemonConfig }, play: async ({ canvas, canvasElement }) => { diff --git a/schema/daemon-config.json b/schema/daemon-config.json index 91d8f061..084035cb 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": "Connect a ChatGPT subscription through Agentdesktop's OAuth credential helper.", + "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", diff --git a/schema/daemon-config.md b/schema/daemon-config.md index 957ce2ed..f16cb82d 100644 --- a/schema/daemon-config.md +++ b/schema/daemon-config.md @@ -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|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.|