diff --git a/AGENTS.md b/AGENTS.md index bb3b8db..5d2d3cd 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -168,21 +168,36 @@ edges stay deterministic Rust. The harness appends a terminal are built generically from advertised methods (`AuthMethodKind:: ApiKeyPlatform`), so new registry rows appear in the picker with no TUI changes. -- Refreshable-OAuth providers beyond Kimi Code use a GENERIC device-code - (RFC-8628) path, NOT Kimi's bespoke wire. A `uses_oauth` platform carrying - `oauth: Some(&OAuthConfig)` (client id / auth host / device+token paths / - scope / `scope_key` / optional extra device field) drives - `auth::oauth_device` (plain kigi UA, no X-Msh headers) + a scope-keyed - `AuthManager::new_oauth_provider` + `refresh::GenericDeviceRefresher` - (selected by `build_refresher` via `oauth_config_for_scope_key`). Kimi Code - keeps `oauth: None` and its bespoke path unchanged. First such provider: - `xai-grok` (`scope_key oauth/xai`, base `api.x.ai/v1`, same wire as the - API-key `xai` row) — an INTERACTIVE login row advertised right after - `kimi-code` (`AuthMethodKind::OAuthPlatform`). Its `authenticate` arm runs - the generic device flow under its own scope; the catalog fetch resolves each - such platform's OWN session token (`resolve_generic_oauth_tokens`, refreshed - on expiry) and routes `platform.oauth().is_some()` → `platform.base_url()` - (kimi-code alone → `proxy_url()`). Tokens are NEVER logged. +- Refreshable-OAuth providers beyond Kimi Code use a GENERIC path, NOT Kimi's + bespoke wire. A `uses_oauth` platform carrying `oauth: Some(&OAuthConfig)` + (client id / auth host / start+token paths / `token_host` / `scope` / + `scope_key` / optional extra device field / `flow` / `token_body`) drives a + scope-keyed `AuthManager::new_oauth_provider` + + `refresh::GenericDeviceRefresher` (selected by `build_refresher` via + `oauth_config_for_scope_key`; the refresher dispatches the refresh body by + `token_body`: form → `auth::oauth_device`, JSON → `auth::oauth_pkce`). Kimi + Code keeps `oauth: None` and its bespoke path unchanged. The interactive + login is dispatched by `OAuthConfig.flow` (in `run_oauth_provider_flow`): + - `OAuthFlow::DeviceCode` → `auth::oauth_device` (RFC-8628 device-code, plain + kigi UA, no X-Msh headers). Provider: `xai-grok` (`scope_key oauth/xai`, + base `api.x.ai/v1`, form token body, same wire as the API-key `xai` row). + - `OAuthFlow::PkceLocalhost { redirect_port }` → `auth::oauth_pkce` + (authorization-code + PKCE S256, `127.0.0.1:redirect_port/callback` loopback + with STRICT `state` validation + manual-paste fallback, JSON token body, + authorize host ≠ token host). Provider: `claude-pro-max` (`scope_key + oauth/claude-pro-max`, base `api.anthropic.com/v1`, Anthropic Messages + + listing wire reached with an OAuth `sk-ant-oat…` Bearer). Its Messages + requests take the OAuth adaptation — `anthropic-beta claude-code-…,oauth-…` + + `claude-cli` UA + `x-app cli` + the required "You are Claude Code…" system + prefix — gated on `SamplerConfig.anthropic_oauth` (claude-pro-max only), so + API-key `anthropic`/`minimax` Messages requests stay byte-identical. Its + `/v1/models` listing rides the same Bearer + oauth-beta headers. + Both are INTERACTIVE login rows advertised right after `kimi-code` + (`AuthMethodKind::OAuthPlatform`, in `PlatformId::ALL` order: `xai-grok` then + `claude-pro-max`). The catalog fetch resolves each such platform's OWN session + token (`resolve_generic_oauth_tokens`, refreshed on expiry) and routes + `platform.oauth().is_some()` → `platform.base_url()` (kimi-code alone → + `proxy_url()`). Tokens/codes/verifiers are NEVER logged. - Model metadata (context window, thinking levels) comes from the provider wire when served; metadata-poor listings are enriched from models.dev (`kigi-models/src/enrichment.rs` — bundled raw snapshot regenerated by diff --git a/crates/codegen/kigi-models/src/lib.rs b/crates/codegen/kigi-models/src/lib.rs index b583b9a..a1a7344 100644 --- a/crates/codegen/kigi-models/src/lib.rs +++ b/crates/codegen/kigi-models/src/lib.rs @@ -96,8 +96,35 @@ enum BaseUrlSource { }, } -/// Generic RFC-8628 device-code OAuth configuration carried by a `uses_oauth` -/// platform whose login is the GENERIC device-code path (xai-grok today). +/// The interactive login mechanism a `uses_oauth` [`OAuthConfig`] provider +/// drives. `DeviceCode` is the RFC-8628 device flow (xai-grok); `PkceLocalhost` +/// is the authorization-code + PKCE (S256) flow with a loopback callback +/// (Claude Pro/Max). Kimi Code carries neither (its bespoke flow lives in +/// kigi-shell with `oauth: None`). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OAuthFlow { + /// RFC-8628 device-code: POST `device_path`, poll `token_path`. + DeviceCode, + /// Authorization-code + PKCE (S256): browser hits `auth_host`+`device_path` + /// (the authorize endpoint); the code returns to a `127.0.0.1:redirect_port` + /// loopback listener, then is exchanged at `token_host`+`token_path`. + PkceLocalhost { redirect_port: u16 }, +} + +/// Body encoding a provider's token endpoint expects for the code-exchange and +/// refresh POSTs. xAI's `/oauth2/token` is form-encoded; Claude's +/// `/v1/oauth/token` is JSON. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OAuthTokenBody { + /// `application/x-www-form-urlencoded` (xai-grok device wire). + Form, + /// `application/json` (Claude PKCE wire). + Json, +} + +/// Generic OAuth configuration carried by a `uses_oauth` platform whose login +/// is the GENERIC device-code path (xai-grok) or the PKCE-localhost path +/// (claude-pro-max). /// /// Kimi Code keeps its bespoke device flow (client id, `/api/oauth/*` paths, /// X-Msh device headers, `kigi_env::oauth_host()`); its `oauth` field stays @@ -105,15 +132,22 @@ enum BaseUrlSource { /// tokens they mint are NEVER stored in this struct and never logged. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct OAuthConfig { - /// OAuth client id sent on every device-authorization / token call. + /// OAuth client id sent on every authorize / token call. pub client_id: &'static str, - /// Authorization-server origin (no trailing slash), e.g. `https://auth.x.ai`. + /// Authorization-server origin (no trailing slash), e.g. `https://auth.x.ai` + /// (device) or `https://claude.ai` (PKCE authorize host). pub auth_host: &'static str, - /// Device-authorization path (POST), relative to `auth_host`. + /// Start-endpoint path (POST for `DeviceCode` device-authorization; the + /// browser authorize path for `PkceLocalhost`), relative to `auth_host`. pub device_path: &'static str, - /// Token path (POST) — used for BOTH the device grant and refresh. + /// Token-endpoint origin (no trailing slash). Equals `auth_host` for the + /// device wire; for Claude the token host (`https://platform.claude.com`) + /// differs from the authorize host (`https://claude.ai`). + pub token_host: &'static str, + /// Token path (POST) — used for BOTH the initial grant and refresh, + /// relative to `token_host`. pub token_path: &'static str, - /// OAuth scope string requested at device authorization. + /// OAuth scope string requested at authorization. pub scope: &'static str, /// auth.json map key + keyring entry name for this provider's persisted /// session (e.g. `oauth/xai`). @@ -121,6 +155,10 @@ pub struct OAuthConfig { /// A non-standard extra form field sent ONLY on the device-authorization /// request (e.g. `("referrer", "kigi")`). `None` = no extra field. pub extra_device_field: Option<(&'static str, &'static str)>, + /// Interactive login mechanism (device-code vs PKCE-localhost). + pub flow: OAuthFlow, + /// Body encoding the token endpoint expects (form vs JSON). + pub token_body: OAuthTokenBody, } /// xAI / Grok subscription device-code OAuth (ported from Pi @@ -129,10 +167,36 @@ pub const XAI_OAUTH_CONFIG: OAuthConfig = OAuthConfig { client_id: "b1a00492-073a-47ea-816f-4c329264a828", auth_host: "https://auth.x.ai", device_path: "/oauth2/device/code", + token_host: "https://auth.x.ai", token_path: "/oauth2/token", scope: "openid profile email offline_access grok-cli:access api:access", scope_key: "oauth/xai", extra_device_field: Some(("referrer", "kigi")), + flow: OAuthFlow::DeviceCode, + token_body: OAuthTokenBody::Form, +}; + +/// Base-URL override for the Claude Pro/Max OAuth channel (dev/test escape +/// hatch). Production defaults to `https://api.anthropic.com/v1`. +pub const CLAUDE_OAUTH_BASE_URL_ENV: &str = "KIGI_CLAUDE_OAUTH_BASE_URL"; + +/// Claude Pro/Max subscription OAuth (authorization-code + PKCE S256, loopback +/// callback). Authoritative constants from Pi `earendil-works/pi` +/// `auth/oauth/anthropic.ts`: authorize host `https://claude.ai`, token host +/// `https://platform.claude.com` (JSON body), public Claude Code client id. +pub const CLAUDE_OAUTH_CONFIG: OAuthConfig = OAuthConfig { + client_id: "9d1c250a-e61b-44d9-88ed-5944d1962f5e", + auth_host: "https://claude.ai", + device_path: "/oauth/authorize", + token_host: "https://platform.claude.com", + token_path: "/v1/oauth/token", + scope: "org:create_api_key user:profile user:inference user:sessions:claude_code user:mcp_servers user:file_upload", + scope_key: "oauth/claude-pro-max", + extra_device_field: None, + flow: OAuthFlow::PkceLocalhost { + redirect_port: 53692, + }, + token_body: OAuthTokenBody::Json, }; /// The generic device-code OAuth config for a platform, or `None` for API-key @@ -1001,6 +1065,42 @@ const MINIMAX_CN_SPEC: PlatformSpec = PlatformSpec { restrict_to_enriched: false, }; +/// Claude Pro/Max subscription via PKCE-localhost OAuth. Reaches +/// `api.anthropic.com` with an OAuth `sk-ant-oat…` bearer (NOT an API key) — +/// same Anthropic Messages + listing wire as the API-key `anthropic` row, plus +/// the OAuth identity headers + "You are Claude Code" system prefix (gated on +/// this platform's OAuth path in the sampler/fetch, so the API-key rows stay +/// byte-identical). +const CLAUDE_PRO_MAX_SPEC: PlatformSpec = PlatformSpec { + id: "claude-pro-max", + display_name: "Claude Pro/Max", + // WITH /v1 so listing → /v1/models and inference → /v1/messages, matching + // ANTHROPIC_SPEC's base handling. + base_url: BaseUrlSource::EnvOr { + env: CLAUDE_OAUTH_BASE_URL_ENV, + default: "https://api.anthropic.com/v1", + }, + uses_oauth: true, + oauth: Some(&CLAUDE_OAUTH_CONFIG), + allowed_model_prefixes: None, + // OAuth channel: no API key envs (the PKCE session is the bearer). + api_key_envs: &[], + vendor: "Anthropic", + console_host: None, + login_label: Some("Claude Pro/Max (subscription)"), + models_dev_id: Some("anthropic"), + wire_serves_metadata: false, + wire_api: PlatformWireApi::Messages, + listing: ListingDialect::Anthropic, + // Passthrough is ignored for the Messages backend. + chat_compat: PlatformChatCompat::Passthrough, + // OAuth uses Authorization: Bearer, NOT x-api-key. + key_header: PlatformKeyHeader::Bearer, + restrict_to_enriched: false, + key_validation_path: None, + strip_listing_id_prefix: None, +}; + /// The platform registry. Platforms are compiled-in spec rows; there is no /// dynamic provider registration (PRD F2). #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] @@ -1057,12 +1157,15 @@ pub enum PlatformId { MinimaxCn, /// xAI Grok subscription via device-code OAuth (same wire as `Xai`). XaiGrok, + /// Claude Pro/Max subscription via PKCE-localhost OAuth (Anthropic Messages + /// wire reached with an OAuth bearer instead of an API key). + ClaudeProMax, } impl PlatformId { /// All platforms, in catalog precedence order: the subscription channel /// first so "default model = first list item" favors it when present. - pub const ALL: [PlatformId; 26] = [ + pub const ALL: [PlatformId; 27] = [ Self::KimiCode, Self::MoonshotCn, Self::MoonshotAi, @@ -1089,6 +1192,7 @@ impl PlatformId { Self::Minimax, Self::MinimaxCn, Self::XaiGrok, + Self::ClaudeProMax, ]; /// The registry row backing this platform (single source of per-platform @@ -1121,6 +1225,7 @@ impl PlatformId { Self::Minimax => &MINIMAX_SPEC, Self::MinimaxCn => &MINIMAX_CN_SPEC, Self::XaiGrok => &XAI_GROK_SPEC, + Self::ClaudeProMax => &CLAUDE_PRO_MAX_SPEC, } } @@ -1899,7 +2004,7 @@ mod tests { let g = PlatformId::XaiGrok; assert_eq!(g.as_str(), "xai-grok"); assert!(g.uses_oauth()); - // The only two uses_oauth platforms; only xai-grok carries a config. + // Kimi Code is uses_oauth yet carries no generic config (bespoke flow). assert!(PlatformId::KimiCode.uses_oauth()); assert_eq!(PlatformId::KimiCode.oauth(), None); let cfg = g @@ -1933,6 +2038,65 @@ mod tests { assert_eq!(g.base_url(), "https://mock.grok/v1"); } + /// claude-pro-max is the first PKCE-localhost OAuth platform: it carries a + /// PKCE `OAuthConfig` (authorize host ≠ token host, JSON token body), reuses + /// the Anthropic Messages + listing wire with a Bearer key header (OAuth, + /// NOT x-api-key), enriches from models.dev "anthropic", and keys its models + /// under `claude-pro-max/`. + #[test] + fn claude_pro_max_is_a_pkce_oauth_platform() { + let c = PlatformId::ClaudeProMax; + assert_eq!(c.as_str(), "claude-pro-max"); + assert!(c.uses_oauth()); + let cfg = c + .oauth() + .expect("claude-pro-max carries a PKCE OAuthConfig"); + assert_eq!(cfg, &CLAUDE_OAUTH_CONFIG); + assert_eq!(cfg.client_id, "9d1c250a-e61b-44d9-88ed-5944d1962f5e"); + // Authorize host ≠ token host (the distinguishing PKCE trait). + assert_eq!(cfg.auth_host, "https://claude.ai"); + assert_eq!(cfg.device_path, "/oauth/authorize"); + assert_eq!(cfg.token_host, "https://platform.claude.com"); + assert_eq!(cfg.token_path, "/v1/oauth/token"); + assert_eq!( + cfg.scope, + "org:create_api_key user:profile user:inference \ + user:sessions:claude_code user:mcp_servers user:file_upload" + ); + assert_eq!(cfg.scope_key, "oauth/claude-pro-max"); + assert_eq!(cfg.extra_device_field, None); + assert_eq!( + cfg.flow, + OAuthFlow::PkceLocalhost { + redirect_port: 53692 + } + ); + assert_eq!(cfg.token_body, OAuthTokenBody::Json); + // xai stays the device-code / form contract — unaffected. + assert_eq!(XAI_OAUTH_CONFIG.flow, OAuthFlow::DeviceCode); + assert_eq!(XAI_OAUTH_CONFIG.token_body, OAuthTokenBody::Form); + assert_eq!(XAI_OAUTH_CONFIG.token_host, "https://auth.x.ai"); + // Scope-key lookup resolves the config (drives the generic refresher). + assert_eq!( + oauth_config_for_scope_key("oauth/claude-pro-max"), + Some(&CLAUDE_OAUTH_CONFIG) + ); + // Anthropic Messages + listing wire, reached with a Bearer OAuth token. + assert_eq!(c.models_dev_id(), Some("anthropic")); + assert!(!c.restrict_to_enriched()); + assert_eq!(c.key_header(), PlatformKeyHeader::Bearer); + assert_eq!(c.wire_api(), PlatformWireApi::Messages); + assert_eq!(c.listing(), ListingDialect::Anthropic); + assert_eq!(c.api_key_env_names(), &[] as &[&str]); + assert_eq!( + c.managed_model_key("claude-opus-4-8"), + "claude-pro-max/claude-opus-4-8" + ); + let _guard = + kigi_env::EnvVarGuard::set(CLAUDE_OAUTH_BASE_URL_ENV, "https://mock.claude/v1"); + assert_eq!(c.base_url(), "https://mock.claude/v1"); + } + /// A variant missing from `ALL` compiles fine (`ALL`'s length is a plain /// literal) but is silently unparseable and excluded from model sync. /// The exhaustive match below fails compilation when a variant is added, @@ -1967,9 +2131,10 @@ mod tests { PlatformId::Minimax => 23, PlatformId::MinimaxCn => 24, PlatformId::XaiGrok => 25, + PlatformId::ClaudeProMax => 26, } } - const VARIANT_COUNT: usize = 26; // update together with `ordinal` + const VARIANT_COUNT: usize = 27; // update together with `ordinal` let mut seen: Vec = PlatformId::ALL.iter().map(|&p| ordinal(p)).collect(); seen.sort_unstable(); seen.dedup(); diff --git a/crates/codegen/kigi-sampler/src/actor/state.rs b/crates/codegen/kigi-sampler/src/actor/state.rs index aacba13..57021c9 100644 --- a/crates/codegen/kigi-sampler/src/actor/state.rs +++ b/crates/codegen/kigi-sampler/src/actor/state.rs @@ -89,6 +89,7 @@ mod tests { top_p: None, api_backend: ApiBackend::ChatCompletions, auth_scheme: Default::default(), + anthropic_oauth: false, chat_compat: Default::default(), extra_headers: IndexMap::new(), context_window: 8192, diff --git a/crates/codegen/kigi-sampler/src/client.rs b/crates/codegen/kigi-sampler/src/client.rs index e477b16..f050e8c 100644 --- a/crates/codegen/kigi-sampler/src/client.rs +++ b/crates/codegen/kigi-sampler/src/client.rs @@ -41,6 +41,35 @@ pub use kigi_sampling_types::ApiBackend; const AGENT_PRODUCT: &str = "kigi"; const ANTHROPIC_DEFAULT_MAX_TOKENS: u32 = 128_000; +/// Prepend the required Claude-Code system block to a Messages request's +/// `system` field (Claude Pro/Max OAuth path). Anthropic inspects the FIRST +/// system block, so the prefix is inserted as a distinct leading `text` block +/// while preserving any caller-supplied prompt (string or block form). +/// Idempotent: a leading block already equal to the prefix is not re-added. +fn prepend_claude_code_system_prefix(system: &mut Option) { + use messages::{SystemParam, TextBlock}; + let text_block = |text: String| TextBlock { + r#type: "text".to_string(), + text, + cache_control: None, + }; + let mut blocks = match system.take() { + None => Vec::new(), + Some(SystemParam::Text(text)) => vec![text_block(text)], + Some(SystemParam::Blocks(blocks)) => blocks, + }; + let already_present = blocks + .first() + .is_some_and(|b| b.text == kigi_sampling_types::CLAUDE_CODE_SYSTEM_PREFIX); + if !already_present { + blocks.insert( + 0, + text_block(kigi_sampling_types::CLAUDE_CODE_SYSTEM_PREFIX.to_string()), + ); + } + *system = Some(SystemParam::Blocks(blocks)); +} + /// Parse the `Retry-After` response header as delta-seconds. /// Our inference backends only emit integer seconds (never HTTP-date), /// so we only handle that form. HTTP-dates silently return `None` and @@ -273,6 +302,8 @@ struct ClientDefaults { chat_compat: kigi_sampling_types::ChatCompat, stream_tool_calls: bool, doom_loop_recovery: Option, + /// Claude Pro/Max OAuth Messages adaptation (see [`SamplerConfig`]). + anthropic_oauth: bool, } // ============================================================================= @@ -398,6 +429,31 @@ impl SamplingClient { } } + // Claude Pro/Max OAuth identity headers (claude-pro-max only). The + // OAuth `sk-ant-oat…` bearer is Claude-Code-scoped, so Anthropic + // rejects the Messages request without the oauth beta + claude-cli + // identity. Gated on `anthropic_oauth` so API-key anthropic/minimax + // requests carry none of this and stay byte-identical. `Accept` is set + // per-request (text/event-stream for streams), so it is NOT added here. + if config.anthropic_oauth { + headers.insert( + HeaderName::from_static("anthropic-version"), + HeaderValue::from_static(kigi_sampling_types::ANTHROPIC_VERSION), + ); + headers.insert( + HeaderName::from_static("anthropic-beta"), + HeaderValue::from_static(kigi_sampling_types::ANTHROPIC_OAUTH_BETA), + ); + headers.insert( + HeaderName::from_static("x-app"), + HeaderValue::from_static("cli"), + ); + headers.insert( + HeaderName::from_static("anthropic-dangerous-direct-browser-access"), + HeaderValue::from_static("true"), + ); + } + // Apply all extra headers verbatim. This is the single // injection point for proxy-auth headers and any other URL- or // environment-specific headers the session decides to set. @@ -415,12 +471,18 @@ impl SamplingClient { // (PRD F3: auth is a plain bearer; kimi-cli sends only User-Agent // plus the OAuth device headers, src/kimi_cli/llm.py:317-323). { - let ua_string = match config.origin_client.as_ref() { - Some(origin) => user_agent_string_for(origin), - None => user_agent_string_for(&OriginClientInfo { - product: AGENT_PRODUCT.to_string(), - version: Some(agent_version()), - }), + // Claude Pro/Max OAuth path presents the claude-cli identity; + // every other path keeps the kigi User-Agent. + let ua_string = if config.anthropic_oauth { + kigi_sampling_types::CLAUDE_CODE_USER_AGENT.to_string() + } else { + match config.origin_client.as_ref() { + Some(origin) => user_agent_string_for(origin), + None => user_agent_string_for(&OriginClientInfo { + product: AGENT_PRODUCT.to_string(), + version: Some(agent_version()), + }), + } }; if let Ok(v) = HeaderValue::from_str(&ua_string) { headers.insert(USER_AGENT, v); @@ -460,6 +522,7 @@ impl SamplingClient { chat_compat: config.chat_compat, stream_tool_calls: config.stream_tool_calls, doom_loop_recovery: config.doom_loop_recovery, + anthropic_oauth: config.anthropic_oauth, }; Ok(Self { @@ -1344,6 +1407,15 @@ impl SamplingClient { /// Apply default configuration to a Messages API request. fn apply_message_defaults(&self, request: &mut MessagesRequestWrapper) -> Result<()> { + // Claude Pro/Max OAuth adaptation (claude-pro-max only): the OAuth + // token is Claude-Code-scoped, so the request MUST lead with the exact + // "You are Claude Code…" system block or Anthropic rejects it. Prepend + // it as a distinct first system block, preserving any caller prompt. + // Gated on `anthropic_oauth` so API-key anthropic/minimax are untouched. + if self.defaults.anthropic_oauth { + prepend_claude_code_system_prefix(&mut request.inner.system); + } + // Apply model default if not specified if request.inner.model.is_empty() { request.inner.model = self.defaults.model.clone(); @@ -1888,6 +1960,7 @@ mod tests { top_p: None, api_backend: ApiBackend::ChatCompletions, auth_scheme: AuthScheme::Bearer, + anthropic_oauth: false, chat_compat: Default::default(), extra_headers: IndexMap::new(), context_window: 8192, @@ -1942,6 +2015,130 @@ mod tests { ); } + /// Claude Pro/Max OAuth Messages client (`anthropic_oauth = true`, Bearer, + /// Messages) carries the full OAuth identity: Bearer auth, anthropic-version, + /// the oauth `anthropic-beta`, `x-app: cli`, the claude-cli User-Agent, and + /// the direct-browser-access header. + #[test] + fn anthropic_oauth_messages_client_sends_oauth_identity_headers() { + let mut config = minimal_config(); + config.api_key = Some("sk-ant-oat-secret".to_string()); + config.auth_scheme = AuthScheme::Bearer; + config.api_backend = ApiBackend::Messages; + config.anthropic_oauth = true; + let client = SamplingClient::new(config).expect("client builds"); + let h = &client.default_headers; + assert_eq!( + h.get(AUTHORIZATION).and_then(|v| v.to_str().ok()), + Some("Bearer sk-ant-oat-secret"), + "OAuth path rides Authorization: Bearer, never x-api-key" + ); + assert!( + h.get("x-api-key").is_none(), + "OAuth path must not send x-api-key" + ); + assert_eq!( + h.get("anthropic-version").and_then(|v| v.to_str().ok()), + Some(kigi_sampling_types::ANTHROPIC_VERSION) + ); + assert_eq!( + h.get("anthropic-beta").and_then(|v| v.to_str().ok()), + Some(kigi_sampling_types::ANTHROPIC_OAUTH_BETA) + ); + assert_eq!(h.get("x-app").and_then(|v| v.to_str().ok()), Some("cli")); + assert_eq!( + h.get(USER_AGENT).and_then(|v| v.to_str().ok()), + Some(kigi_sampling_types::CLAUDE_CODE_USER_AGENT) + ); + assert_eq!( + h.get("anthropic-dangerous-direct-browser-access") + .and_then(|v| v.to_str().ok()), + Some("true") + ); + } + + /// REGRESSION: an API-key Anthropic Messages client (XApiKey, NOT oauth) + /// carries NONE of the OAuth identity — no anthropic-beta, no x-app, and + /// the kigi User-Agent — so the API-key path stays byte-identical. + #[test] + fn api_key_anthropic_messages_client_has_no_oauth_identity() { + let mut config = minimal_config(); + config.auth_scheme = AuthScheme::XApiKey; + config.api_backend = ApiBackend::Messages; + // anthropic_oauth stays false. + let client = SamplingClient::new(config).expect("client builds"); + let h = &client.default_headers; + assert!( + h.get("anthropic-beta").is_none(), + "API-key anthropic must NOT send the oauth beta" + ); + assert!( + h.get("x-app").is_none(), + "API-key anthropic must NOT send x-app" + ); + assert!(h.get("anthropic-dangerous-direct-browser-access").is_none()); + assert!( + h.get(USER_AGENT) + .and_then(|v| v.to_str().ok()) + .is_some_and(|ua| ua.starts_with("kigi/")), + "API-key anthropic keeps the kigi User-Agent" + ); + } + + /// The system-prompt prefix is prepended as a distinct leading `text` + /// block for each `system` shape (absent / string / blocks), preserving the + /// caller's prompt, and is idempotent (not stamped twice). + #[test] + fn claude_code_system_prefix_prepends_and_is_idempotent() { + use messages::{SystemParam, TextBlock}; + let prefix = kigi_sampling_types::CLAUDE_CODE_SYSTEM_PREFIX; + + // Absent system → a single prefix block. + let mut none = None; + prepend_claude_code_system_prefix(&mut none); + match none { + Some(SystemParam::Blocks(b)) => { + assert_eq!(b.len(), 1); + assert_eq!(b[0].text, prefix); + } + other => panic!("expected one prefix block, got {other:?}"), + } + + // String system → [prefix, original]. + let mut text = Some(SystemParam::Text("do the thing".into())); + prepend_claude_code_system_prefix(&mut text); + match text { + Some(SystemParam::Blocks(b)) => { + assert_eq!(b.len(), 2); + assert_eq!(b[0].text, prefix); + assert_eq!(b[1].text, "do the thing"); + } + other => panic!("expected two blocks, got {other:?}"), + } + + // Idempotent: a leading prefix block is not re-added. + let mut already = Some(SystemParam::Blocks(vec![ + TextBlock { + r#type: "text".into(), + text: prefix.to_string(), + cache_control: None, + }, + TextBlock { + r#type: "text".into(), + text: "tail".into(), + cache_control: None, + }, + ])); + prepend_claude_code_system_prefix(&mut already); + match already { + Some(SystemParam::Blocks(b)) => { + assert_eq!(b.len(), 2, "prefix must not be duplicated"); + assert_eq!(b[0].text, prefix); + } + other => panic!("expected unchanged blocks, got {other:?}"), + } + } + /// Verify the serialized shape of StreamingChatRequest matches the /// expected wire format: all ChatCompletionRequest fields flattened at /// top level, plus `stream: true` and `stream_options.include_usage: true`. diff --git a/crates/codegen/kigi-sampler/src/config.rs b/crates/codegen/kigi-sampler/src/config.rs index b7175ae..af0dd5c 100644 --- a/crates/codegen/kigi-sampler/src/config.rs +++ b/crates/codegen/kigi-sampler/src/config.rs @@ -55,6 +55,13 @@ pub struct SamplerConfig { pub api_backend: ApiBackend, #[serde(default)] pub auth_scheme: AuthScheme, + /// Claude Pro/Max OAuth adaptation (claude-pro-max only). When true the + /// Messages request carries the OAuth identity headers (`anthropic-beta` + /// oauth, `claude-cli` User-Agent, `x-app: cli`) and its system prompt is + /// prefixed with the required "You are Claude Code…" line. Gated so the + /// API-key `anthropic` + `minimax` Messages requests stay byte-identical. + #[serde(default)] + pub anthropic_oauth: bool, /// Extra request headers applied verbatim. The sampler never inspects /// the URL to derive headers; callers (the session) inject proxy auth /// and other access headers here before constructing the config. @@ -144,6 +151,7 @@ impl Default for SamplerConfig { top_p: None, api_backend: ApiBackend::default(), auth_scheme: AuthScheme::default(), + anthropic_oauth: false, extra_headers: IndexMap::new(), context_window: 0, force_http1: false, diff --git a/crates/codegen/kigi-sampler/tests/test_actor.rs b/crates/codegen/kigi-sampler/tests/test_actor.rs index 319ba43..477a004 100644 --- a/crates/codegen/kigi-sampler/tests/test_actor.rs +++ b/crates/codegen/kigi-sampler/tests/test_actor.rs @@ -78,6 +78,7 @@ fn test_config(base_url: String, model: &str) -> SamplerConfig { top_p: None, api_backend: ApiBackend::ChatCompletions, auth_scheme: Default::default(), + anthropic_oauth: false, chat_compat: Default::default(), extra_headers: IndexMap::new(), context_window: 128_000, diff --git a/crates/codegen/kigi-sampling-types/src/types.rs b/crates/codegen/kigi-sampling-types/src/types.rs index 7aaba35..c656591 100644 --- a/crates/codegen/kigi-sampling-types/src/types.rs +++ b/crates/codegen/kigi-sampling-types/src/types.rs @@ -1060,6 +1060,22 @@ pub fn normalize_effort_echo(value: &mut Value) { /// wires (Messages inference and the /v1/models listing). pub const ANTHROPIC_VERSION: &str = "2023-06-01"; +/// `anthropic-beta` value for the Claude-Code OAuth (Pro/Max subscription) +/// path — required on BOTH the Messages inference request and the `/v1/models` +/// listing when the bearer is an OAuth `sk-ant-oat…` token. API-key Anthropic +/// and MiniMax NEVER send this (their requests stay byte-identical). +pub const ANTHROPIC_OAUTH_BETA: &str = "claude-code-20250219,oauth-2025-04-20"; + +/// User-Agent kigi presents on the Claude-Code OAuth path (mirrors the +/// official Claude Code CLI). OAuth-gated: unrelated to the default kigi UA. +pub const CLAUDE_CODE_USER_AGENT: &str = "claude-cli/2.1.75"; + +/// System-prompt prefix REQUIRED on the Claude-Code OAuth Messages path: the +/// OAuth token is Claude-Code-scoped, so Anthropic rejects the request unless +/// the system prompt's first block is exactly this line. OAuth-gated. +pub const CLAUDE_CODE_SYSTEM_PREFIX: &str = + "You are Claude Code, Anthropic's official CLI for Claude."; + /// ChatCompletions request-body adaptation dialect. Providers disagree on /// how thinking rides an OpenAI-compatible body: Kimi wants /// `thinking:{type,effort}`, DeepSeek wants diff --git a/crates/codegen/kigi-shell/src/agent/auth_method.rs b/crates/codegen/kigi-shell/src/agent/auth_method.rs index c8ec282..f082a72 100644 --- a/crates/codegen/kigi-shell/src/agent/auth_method.rs +++ b/crates/codegen/kigi-shell/src/agent/auth_method.rs @@ -595,6 +595,11 @@ mod tests { ); assert_eq!( ids[kimi_pos + 2], + "claude-pro-max", + "claude-pro-max is the next interactive OAuth login, after xai-grok" + ); + assert_eq!( + ids[kimi_pos + 3], MOONSHOT_CN_METHOD_ID, "the api-key rows follow the generic oauth logins" ); @@ -628,6 +633,29 @@ mod tests { assert_eq!(kind.auth_error_message(), AUTH_ERROR_SESSION_EXPIRED); } + /// claude-pro-max classifies as an interactive OAuth login too (the + /// authenticate handler dispatches it to the PKCE-localhost flow by the + /// config's `flow`): session-based, needs a browser, never api-key, and + /// `oauth_platform()` returns ClaudeProMax. + #[test] + fn claude_pro_max_is_an_interactive_oauth_login() { + let id = acp::AuthMethodId::new("claude-pro-max"); + let kind = AuthMethodKind::from_id(&id); + assert_eq!( + kind, + AuthMethodKind::OAuthPlatform(kigi_models::PlatformId::ClaudeProMax) + ); + assert!(kind.needs_interactive_login()); + assert!(kind.is_session_based()); + assert!(!kind.is_api_key()); + assert_eq!( + kind.oauth_platform(), + Some(kigi_models::PlatformId::ClaudeProMax) + ); + // Never an API-key picker target (keeps it out of the paste-box path). + assert_eq!(platform_for_method_id(&id), None); + } + /// The OAuth platform id must never resolve as an API-key platform /// method — `platform_for_method_id`'s `uses_oauth` filter is what keeps /// the generic `authenticate` arm from hijacking the device login. @@ -721,6 +749,7 @@ mod tests { XAI_API_KEY_METHOD_ID, KIMI_CODE_METHOD_ID, "xai-grok", + "claude-pro-max", MOONSHOT_CN_METHOD_ID, MOONSHOT_AI_METHOD_ID, "openai", @@ -770,6 +799,7 @@ mod tests { CACHED_TOKEN_AUTH_METHOD_ID, KIMI_CODE_METHOD_ID, "xai-grok", + "claude-pro-max", MOONSHOT_CN_METHOD_ID, MOONSHOT_AI_METHOD_ID, "openai", @@ -812,6 +842,7 @@ mod tests { CACHED_TOKEN_AUTH_METHOD_ID, KIMI_CODE_METHOD_ID, "xai-grok", + "claude-pro-max", MOONSHOT_CN_METHOD_ID, MOONSHOT_AI_METHOD_ID, "openai", @@ -857,6 +888,7 @@ mod tests { vec![ KIMI_CODE_METHOD_ID, "xai-grok", + "claude-pro-max", MOONSHOT_CN_METHOD_ID, MOONSHOT_AI_METHOD_ID, "openai", diff --git a/crates/codegen/kigi-shell/src/agent/config.rs b/crates/codegen/kigi-shell/src/agent/config.rs index 2ef7fdb..c652245 100644 --- a/crates/codegen/kigi-shell/src/agent/config.rs +++ b/crates/codegen/kigi-shell/src/agent/config.rs @@ -4079,6 +4079,18 @@ pub fn sampling_config_for_model( } }) .unwrap_or_default(); + // Claude Pro/Max OAuth Messages adaptation: a managed key whose platform is + // a generic-OAuth Messages provider (claude-pro-max) drives the OAuth + // identity headers + "You are Claude Code" system prefix in the sampler. + // Gated here so API-key anthropic/minimax (oauth None) stay byte-identical. + let anthropic_oauth = info + .id + .as_deref() + .and_then(kigi_models::parse_managed_model_key) + .is_some_and(|(platform, _)| { + platform.oauth().is_some() + && platform.wire_api() == kigi_models::PlatformWireApi::Messages + }); SamplerConfig { api_key: credentials.api_key, model: model_name, @@ -4088,6 +4100,7 @@ pub fn sampling_config_for_model( top_p, api_backend, auth_scheme: credentials.auth_scheme, + anthropic_oauth, chat_compat, extra_headers, context_window: info.context_window.get(), diff --git a/crates/codegen/kigi-shell/src/agent/models_fetch.rs b/crates/codegen/kigi-shell/src/agent/models_fetch.rs index acdd9de..dacbcfa 100644 --- a/crates/codegen/kigi-shell/src/agent/models_fetch.rs +++ b/crates/codegen/kigi-shell/src/agent/models_fetch.rs @@ -345,9 +345,22 @@ fn fetch_one_platform_models( }; tracing::info!(platform = platform.as_str(), url = %url, "fetching platform models"); let request = match platform.key_header() { - kigi_models::PlatformKeyHeader::Bearer => client - .get(&url) - .header("Authorization", format!("Bearer {}", bearer)), + kigi_models::PlatformKeyHeader::Bearer => { + let mut req = client + .get(&url) + .header("Authorization", format!("Bearer {}", bearer)); + // An Anthropic listing reached with a Bearer key is the OAuth + // channel (claude-pro-max): the /v1/models endpoint requires + // anthropic-version, and the OAuth bearer needs the oauth beta. + // The OpenAI-listing Bearer platforms (xai-grok, api-key OpenAI + // rows) add neither, so their requests stay byte-identical. + if platform.listing() == kigi_models::ListingDialect::Anthropic { + req = req + .header("anthropic-version", kigi_sampling_types::ANTHROPIC_VERSION) + .header("anthropic-beta", kigi_sampling_types::ANTHROPIC_OAUTH_BETA); + } + req + } kigi_models::PlatformKeyHeader::XApiKey => client .get(&url) .header("x-api-key", bearer) @@ -1106,6 +1119,118 @@ mod tests { ); } + /// Claude Pro/Max OAuth fetch e2e (mock wire): `GET /v1/models?limit=1000` + /// gated on the OAuth `Authorization: Bearer` + the oauth `anthropic-beta` + /// (the OAuth listing contract) → anthropic listing → enriched from + /// models.dev "anthropic" → keyed `claude-pro-max/` on the Messages + /// backend. The token is drawn from the claude-pro-max OAuth-session map, + /// NOT an x-api-key. + #[tokio::test(flavor = "multi_thread")] + #[serial_test::serial] + async fn claude_pro_max_oauth_listing_is_bearer_gated_and_keyed() { + let platform_server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/models")) + .and(wiremock::matchers::query_param("limit", "1000")) + // OAuth listing: Bearer token + the oauth beta + anthropic-version. + // NO x-api-key header (that is the API-key `anthropic` path). + .and(wiremock::matchers::header( + "Authorization", + "Bearer sk-ant-oat-session", + )) + // The oauth beta is comma-joined; wiremock's exact `header` matcher + // splits on commas, so assert both tokens via the multi-valued form. + .and(wiremock::matchers::headers( + "anthropic-beta", + vec!["claude-code-20250219", "oauth-2025-04-20"], + )) + .and(wiremock::matchers::header( + "anthropic-version", + "2023-06-01", + )) + .respond_with(wiremock::ResponseTemplate::new(200).set_body_json( + serde_json::json!({ "data": [ + { + "id": "claude-opus-4-8", + "display_name": "Claude Opus 4.8", + "type": "model" + } + ], "has_more": false }), + )) + .expect(1) + .mount(&platform_server) + .await; + let modelsdev_server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/api.json")) + .respond_with(wiremock::ResponseTemplate::new(200).set_body_json( + serde_json::json!({ "anthropic": { "models": { + "claude-opus-4-8": { + "limit": {"context": 1000000, "output": 128000}, + "tool_call": true + } + }}}), + )) + .expect(1) + .mount(&modelsdev_server) + .await; + let cache_dir = tempfile::tempdir().unwrap(); + let _base = kigi_test_support::EnvGuard::set( + kigi_models::CLAUDE_OAUTH_BASE_URL_ENV, + platform_server.uri(), + ); + let _mdev = kigi_test_support::EnvGuard::set( + crate::agent::enrichment_fetch::MODELS_DEV_URL_ENV, + format!("{}/api.json", modelsdev_server.uri()), + ); + let _mdev_cache = kigi_test_support::EnvGuard::set( + crate::agent::enrichment_fetch::MODELS_DEV_CACHE_DIR_ENV, + cache_dir.path(), + ); + + let endpoints = crate::agent::config::EndpointsConfig::default(); + // The listing bearer comes from the claude-pro-max OAuth-session map, + // never a Kimi session (auth=None) or an API key (keys empty). + let mut oauth_tokens = OAuthSessionTokens::new(); + oauth_tokens.insert( + kigi_models::PlatformId::ClaudeProMax, + "sk-ant-oat-session".to_string(), + ); + let keys = crate::agent::models::PlatformApiKeys::default(); + let result = tokio::task::spawn_blocking(move || { + fetch_platform_models_blocking(&endpoints, None, &oauth_tokens, &keys) + }) + .await + .unwrap() + .expect("claude-pro-max oauth fetch must succeed"); + + assert_eq!(result.models.len(), 1); + let opus = &result.models[0]; + assert_eq!( + opus.id.as_deref(), + Some("claude-pro-max/claude-opus-4-8"), + "the entry must key under the claude-pro-max platform" + ); + assert_eq!( + opus.api_backend, + crate::sampling::ApiBackend::Messages, + "claude-pro-max speaks the Messages wire" + ); + assert_eq!( + opus.auth_scheme, None, + "OAuth Bearer entries carry no XApiKey auth scheme" + ); + assert_eq!( + opus.context_window.get(), + 1_000_000, + "enrichment fills the context window from models.dev anthropic" + ); + assert!( + !opus.supported_in_api, + "subscription (uses_oauth) models require the OAuth session" + ); + } + /// DeepSeek-cycle e2e: bare OpenAI-shape listing + enrichment efforts /// (high/max) produce ChatCompletions entries whose sampler config /// speaks the DeepSeek thinking dialect. diff --git a/crates/codegen/kigi-shell/src/agent/subagent/mod.rs b/crates/codegen/kigi-shell/src/agent/subagent/mod.rs index db5fa31..ecfd3f0 100644 --- a/crates/codegen/kigi-shell/src/agent/subagent/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/subagent/mod.rs @@ -880,6 +880,14 @@ async fn read_parent_sampling_config( let auth_scheme = crate::agent::config::try_resolve_model_credentials(&cfg.model, None) .map(|r| r.auth_scheme) .unwrap_or_default(); + // Claude Pro/Max OAuth Messages adaptation inherits from the parent + // model's platform (claude-pro-max → true); every other platform, + // and BYOK, → false, so the API-key paths stay byte-identical. + let anthropic_oauth = kigi_models::parse_managed_model_key(ctx.model_id.0.as_ref()) + .is_some_and(|(platform, _)| { + platform.oauth().is_some() + && platform.wire_api() == kigi_models::PlatformWireApi::Messages + }); let inherited = kigi_sampler::SamplerConfig { api_key: creds.api_key, base_url: cfg.base_url, @@ -889,6 +897,7 @@ async fn read_parent_sampling_config( top_p: cfg.top_p, api_backend: cfg.api_backend, auth_scheme, + anthropic_oauth, chat_compat: cfg.chat_compat, extra_headers, context_window: cfg.context_window.get(), diff --git a/crates/codegen/kigi-shell/src/auth/device_code.rs b/crates/codegen/kigi-shell/src/auth/device_code.rs index e093163..1003331 100644 --- a/crates/codegen/kigi-shell/src/auth/device_code.rs +++ b/crates/codegen/kigi-shell/src/auth/device_code.rs @@ -198,8 +198,8 @@ async fn complete_device_code_login( /// Open `url` in the browser off-thread: `webbrowser::open` is synchronous and /// would stall the single-threaded TUI loop. Returns `true` on success so the /// caller can decide how to notify the user (eprintln on CLI, nothing on TUI -/// where the URL is already rendered in the widget). -async fn open_browser_detached(url: &str) -> bool { +/// where the URL is already rendered in the widget). Shared with the PKCE flow. +pub(super) async fn open_browser_detached(url: &str) -> bool { // Unit tests drive the full login flow against mock servers — their // fixture URLs must never reach a real browser. if cfg!(test) { diff --git a/crates/codegen/kigi-shell/src/auth/flow.rs b/crates/codegen/kigi-shell/src/auth/flow.rs index e6bb001..e0977c5 100644 --- a/crates/codegen/kigi-shell/src/auth/flow.rs +++ b/crates/codegen/kigi-shell/src/auth/flow.rs @@ -67,10 +67,11 @@ pub async fn run_auth_flow( run_auth_flow_inner(auth_manager, kimi_code_config, reauth, false, channels).await } -/// Login flow for a GENERIC device-code OAuth provider (xai-grok): use a valid -/// cached session unless re-authing, otherwise run the generic device flow +/// Login flow for a GENERIC OAuth provider (xai-grok device-code, +/// claude-pro-max PKCE-localhost): use a valid cached session unless re-authing, +/// otherwise dispatch by `oauth.flow` to the device-code or PKCE-localhost login /// (persisting under the provider's own scope via `auth_manager`). Unlike the -/// Kimi flow this does not run the silent-refresh dance — the device flow's +/// Kimi flow this does not run the silent-refresh dance — the login's /// `AuthManager::update` persists a fresh token set directly. pub async fn run_oauth_provider_flow( auth_manager: &Arc, @@ -94,8 +95,96 @@ pub async fn run_oauth_provider_flow( return Ok((auth, false)); } let mut channels = channels; - crate::auth::device_code::run_device_code_login_generic(oauth, auth_manager, &mut channels) + match oauth.flow { + kigi_models::OAuthFlow::DeviceCode => { + crate::auth::device_code::run_device_code_login_generic( + oauth, + auth_manager, + &mut channels, + ) + .await + } + kigi_models::OAuthFlow::PkceLocalhost { redirect_port } => { + run_pkce_localhost_login(oauth, redirect_port, auth_manager, &mut channels).await + } + } +} + +/// PKCE-localhost login (claude-pro-max): generate PKCE, present the browser +/// authorize URL (TUI channel or stderr), open the browser, then await the code +/// from EITHER the `127.0.0.1:{redirect_port}` loopback callback OR a manual +/// paste (headless fallback). Exchange it at the token endpoint and persist. +/// +/// SECURITY: the verifier / code / tokens are never logged; the loopback binds +/// `127.0.0.1` only and validates `state` strictly. +async fn run_pkce_localhost_login( + oauth: &'static kigi_models::OAuthConfig, + redirect_port: u16, + auth_manager: &Arc, + channels: &mut Option, +) -> anyhow::Result<(KimiAuth, bool)> { + use crate::auth::oauth_pkce; + + let pkce = oauth_pkce::generate_pkce(); + let redirect = oauth_pkce::redirect_uri(redirect_port); + let authorize_url = oauth_pkce::build_authorize_url(oauth, &redirect, &pkce); + + let mut chans = channels.take(); + if let Some(tx) = chans.as_mut().and_then(|c| c.url_tx.take()) { + // TUI: push the URL BEFORE opening the browser (never block the UI on a + // slow/headless browser launch). + let _ = tx.send(AuthUrlInfo { + url: authorize_url.clone(), + mode: AuthUrlMode::Device, + }); + crate::auth::device_code::open_browser_detached(&authorize_url).await; + } else { + eprintln!(); + eprintln!("To sign in to Claude Pro/Max, open this URL in your browser:"); + eprintln!(); + eprintln!(" {authorize_url}"); + eprintln!(); + if !crate::auth::device_code::open_browser_detached(&authorize_url).await { + eprintln!(" (Could not open the browser automatically — open the URL above.)"); + eprintln!(); + } + eprintln!("Waiting for the sign-in to complete..."); + } + + let code = await_pkce_code(redirect_port, &pkce, chans.as_mut()).await?; + let auth = oauth_pkce::exchange_code(oauth, &code, &pkce, &redirect).await?; + let auth = auth_manager + .update(auth) .await + .map_err(|e| anyhow::anyhow!("Failed to save credentials: {e}"))?; + Ok((auth, true)) +} + +/// Await the authorization code: the loopback callback is primary; when a TUI +/// channel is present, a pasted code (redirect URL / `code#state` / bare code) +/// is accepted concurrently as a headless fallback. State is validated in both +/// arms (strict on the loopback, mismatch-rejecting on the paste). +async fn await_pkce_code( + redirect_port: u16, + pkce: &crate::auth::oauth_pkce::PkceCodes, + channels: Option<&mut AuthChannels>, +) -> anyhow::Result { + use crate::auth::oauth_pkce; + match channels { + Some(ch) => { + tokio::select! { + code = oauth_pkce::await_loopback_code(redirect_port, &pkce.state) => code, + pasted = ch.code_rx.recv() => { + let pasted = pasted + .ok_or_else(|| anyhow::anyhow!("auth code channel closed before a code arrived"))?; + let params = oauth_pkce::parse_manual_paste(&pasted)?; + oauth_pkce::validate_pasted_state(¶ms, &pkce.state)?; + Ok(params.code) + } + } + } + None => oauth_pkce::await_loopback_code(redirect_port, &pkce.state).await, + } } async fn run_auth_flow_inner( diff --git a/crates/codegen/kigi-shell/src/auth/mod.rs b/crates/codegen/kigi-shell/src/auth/mod.rs index 4068984..1f3b4fe 100644 --- a/crates/codegen/kigi-shell/src/auth/mod.rs +++ b/crates/codegen/kigi-shell/src/auth/mod.rs @@ -9,6 +9,7 @@ pub(crate) mod kimi_oauth; pub(crate) mod manager; mod model; pub(crate) mod oauth_device; +pub(crate) mod oauth_pkce; pub(crate) mod oauth_registry; pub(crate) mod recovery; pub(crate) mod refresh; diff --git a/crates/codegen/kigi-shell/src/auth/oauth_pkce.rs b/crates/codegen/kigi-shell/src/auth/oauth_pkce.rs new file mode 100644 index 0000000..880a6e9 --- /dev/null +++ b/crates/codegen/kigi-shell/src/auth/oauth_pkce.rs @@ -0,0 +1,687 @@ +//! Generic authorization-code + PKCE (S256) OAuth wire with a `127.0.0.1` +//! loopback callback, driven by a registry [`kigi_models::OAuthConfig`] whose +//! `flow` is [`OAuthFlow::PkceLocalhost`] (claude-pro-max today). +//! +//! Shape (Pi `earendil-works/pi` `auth/oauth/anthropic.ts`): +//! - `verifier = base64url(32 random bytes)`; `challenge = base64url(SHA-256( +//! verifier))`; `state = verifier`. +//! - Browser opens `{auth_host}{device_path}?client_id&response_type=code& +//! scope&redirect_uri&state&code_challenge&code_challenge_method=S256`. +//! - The code returns to a loopback listener on `127.0.0.1:{redirect_port}` +//! answering ONLY `/callback`, with STRICT `state` validation (a mismatch is +//! rejected — CSRF guard). A manual paste (redirect URL / `code#state` / bare +//! code) is accepted as a headless fallback. +//! - Code → token exchange and refresh POST `{token_host}{token_path}` as JSON +//! (per `token_body`); the refresh token ROTATES. +//! +//! SECURITY: the verifier, authorization code, access token, and refresh token +//! are NEVER logged (only non-secret events: authorize URL requested, callback +//! received, token issued, token refreshed). + +use anyhow::Context; +use base64::Engine; +use kigi_models::{OAuthConfig, OAuthTokenBody}; +use serde::Deserialize; +use sha2::{Digest, Sha256}; + +use super::kimi_oauth::{RefreshError, TokenResponse}; +use super::model::KimiAuth; + +const CODE_GRANT_TYPE: &str = "authorization_code"; +const REFRESH_GRANT_TYPE: &str = "refresh_token"; +/// Refresh retry budget over the retryable statuses / network blips. +const MAX_REFRESH_RETRIES: u32 = 3; +/// HTTP statuses worth retrying a refresh for (parity with the device wire). +const RETRYABLE_REFRESH_STATUSES: [u16; 5] = [429, 500, 502, 503, 504]; + +/// PKCE secrets for one login attempt. `state == verifier` (Pi's convention: +/// the state is the verifier, so a returned state binds the callback to this +/// attempt AND doubles as the CSRF token). +#[derive(Debug, Clone)] +pub(crate) struct PkceCodes { + /// `code_verifier` — the 43-char base64url secret, sent at token exchange. + pub verifier: String, + /// `code_challenge = base64url(SHA-256(verifier))`, sent at authorize. + pub challenge: String, + /// `state` — equal to `verifier`; validated on the callback (CSRF guard). + pub state: String, +} + +/// Generate PKCE S256 codes: `verifier = base64url(32 random bytes)`, +/// `challenge = base64url(SHA-256(verifier))`, `state = verifier`. +pub(crate) fn generate_pkce() -> PkceCodes { + use rand::RngCore; + let mut raw = [0u8; 32]; + rand::rng().fill_bytes(&mut raw); + let verifier = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(raw); + let digest = Sha256::digest(verifier.as_bytes()); + let challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(digest); + PkceCodes { + state: verifier.clone(), + verifier, + challenge, + } +} + +/// The loopback redirect URI for a PKCE-localhost provider. +pub(crate) fn redirect_uri(redirect_port: u16) -> String { + format!("http://localhost:{redirect_port}/callback") +} + +/// Build the browser authorize URL: +/// `{auth_host}{device_path}?client_id&response_type=code&scope&redirect_uri& +/// state&code_challenge&code_challenge_method=S256`. +pub(crate) fn build_authorize_url( + cfg: &OAuthConfig, + redirect_uri: &str, + pkce: &PkceCodes, +) -> String { + let base = format!("{}{}", cfg.auth_host.trim_end_matches('/'), cfg.device_path); + let query = url::form_urlencoded::Serializer::new(String::new()) + .append_pair("client_id", cfg.client_id) + .append_pair("response_type", "code") + .append_pair("scope", cfg.scope) + .append_pair("redirect_uri", redirect_uri) + .append_pair("state", &pkce.state) + .append_pair("code_challenge", &pkce.challenge) + .append_pair("code_challenge_method", "S256") + .finish(); + format!("{base}?{query}") +} + +/// `code` + `state` extracted from a callback (loopback query OR manual paste). +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct CallbackParams { + pub code: String, + pub state: Option, +} + +/// Parse `code`/`state` from the raw query string of a `/callback?…` request +/// (e.g. `code=abc&state=xyz`). An `error=` param surfaces as an `Err`. +pub(crate) fn parse_callback_query(query: &str) -> anyhow::Result { + let mut code = None; + let mut state = None; + let mut error = None; + for (k, v) in url::form_urlencoded::parse(query.as_bytes()) { + match k.as_ref() { + "code" => code = Some(v.into_owned()), + "state" => state = Some(v.into_owned()), + "error" => error = Some(v.into_owned()), + _ => {} + } + } + if let Some(error) = error { + anyhow::bail!("Authorization server returned an error: {error}"); + } + let code = code.context("callback missing authorization code")?; + if code.is_empty() { + anyhow::bail!("callback authorization code was empty"); + } + Ok(CallbackParams { code, state }) +} + +/// Parse a MANUAL paste (headless fallback). Accepts, in order: +/// - a full redirect URL (`http://localhost:…/callback?code=…&state=…`), +/// - a `code#state` pair (Anthropic's console shows this form), +/// - a bare `code` (state then unknown → `None`, caller validation applies). +pub(crate) fn parse_manual_paste(input: &str) -> anyhow::Result { + let trimmed = input.trim(); + if trimmed.is_empty() { + anyhow::bail!("empty paste"); + } + // Full redirect URL. + if trimmed.starts_with("http://") || trimmed.starts_with("https://") { + let url = url::Url::parse(trimmed).context("pasted value is not a valid URL")?; + return parse_callback_query(url.query().unwrap_or_default()); + } + // `code#state`. + if let Some((code, state)) = trimmed.split_once('#') { + if code.is_empty() { + anyhow::bail!("pasted code was empty"); + } + return Ok(CallbackParams { + code: code.to_owned(), + state: (!state.is_empty()).then(|| state.to_owned()), + }); + } + // Bare code. + Ok(CallbackParams { + code: trimmed.to_owned(), + state: None, + }) +} + +/// STRICT state validation (CSRF guard): the callback `state` MUST be present +/// AND equal to the expected value. A mismatch (or absence, when a paste has no +/// state) is rejected — the flow NEVER proceeds on an unverified callback. +pub(crate) fn validate_state(params: &CallbackParams, expected_state: &str) -> anyhow::Result<()> { + match params.state.as_deref() { + Some(state) if state == expected_state => Ok(()), + Some(_) => anyhow::bail!("OAuth state mismatch — rejecting callback (CSRF guard)"), + None => anyhow::bail!("OAuth callback carried no state — rejecting (CSRF guard)"), + } +} + +/// State validation for a MANUAL paste (headless fallback): a present state +/// MUST match (mismatch rejected — CSRF guard), but an ABSENT state is allowed +/// — a bare-code paste is user-initiated (not a network-reachable callback), so +/// there is no state to check. The loopback path uses the stricter +/// [`validate_state`] (an absent state there IS rejected). +pub(crate) fn validate_pasted_state( + params: &CallbackParams, + expected_state: &str, +) -> anyhow::Result<()> { + match params.state.as_deref() { + Some(state) if state == expected_state => Ok(()), + Some(_) => anyhow::bail!("OAuth state mismatch — rejecting pasted code (CSRF guard)"), + None => Ok(()), + } +} + +/// The token-endpoint URL (`{token_host}{token_path}`). +fn token_url(cfg: &OAuthConfig) -> String { + format!("{}{}", cfg.token_host.trim_end_matches('/'), cfg.token_path) +} + +/// POST the token endpoint with a JSON body, honoring `cfg.token_body`. Claude +/// is JSON; a `Form`-bodied config would be handled by the device wire, so the +/// PKCE path asserts JSON (never silently mis-encodes). +async fn post_token_json( + cfg: &OAuthConfig, + body: serde_json::Value, +) -> reqwest::Result { + debug_assert!( + matches!(cfg.token_body, OAuthTokenBody::Json), + "PKCE token exchange expects a JSON token body" + ); + crate::http::shared_client() + .post(token_url(cfg)) + .header("Accept", "application/json") + .json(&body) + .send() + .await +} + +/// Exchange an authorization `code` for a token set (JSON body): +/// `{grant_type:"authorization_code", code, state, client_id, redirect_uri, +/// code_verifier}`. Returns the materialized [`KimiAuth`]. +pub(crate) async fn exchange_code( + cfg: &OAuthConfig, + code: &str, + pkce: &PkceCodes, + redirect_uri: &str, +) -> anyhow::Result { + let body = serde_json::json!({ + "grant_type": CODE_GRANT_TYPE, + "code": code, + "state": pkce.state, + "client_id": cfg.client_id, + "redirect_uri": redirect_uri, + "code_verifier": pkce.verifier, + }); + tracing::info!( + scope_key = cfg.scope_key, + "auth: exchanging code for token (pkce)" + ); + let resp = post_token_json(cfg, body) + .await + .context("token exchange request failed")?; + let status = resp.status(); + if !status.is_success() { + let body = resp.text().await.unwrap_or_default(); + tracing::warn!(%status, scope_key = cfg.scope_key, "auth: code exchange failed (pkce)"); + anyhow::bail!("Token exchange failed (HTTP {status}): {body}"); + } + let tokens: TokenResponse = resp.json().await.context("malformed token payload")?; + tracing::info!( + scope_key = cfg.scope_key, + "auth: pkce code exchange succeeded" + ); + Ok(tokens.into_auth()) +} + +/// `POST {token_host}{token_path}` with `grant_type=refresh_token` (JSON body). +/// Claude ROTATES the refresh token, so the caller MUST persist the returned +/// one. Retries the retryable statuses / network errors with exponential +/// backoff; 401/403 returns immediately as [`RefreshError::Unauthorized`]. +pub(crate) async fn refresh_token( + cfg: &OAuthConfig, + refresh_token: &str, +) -> Result { + let mut last_error = String::from("no attempt made"); + for attempt in 0..MAX_REFRESH_RETRIES { + if attempt > 0 { + let backoff = std::time::Duration::from_secs(1 << (attempt - 1)); + tracing::warn!( + attempt, + backoff_secs = backoff.as_secs(), + last_error = %last_error, + "auth: retrying token refresh (pkce)" + ); + tokio::time::sleep(backoff).await; + } + tracing::info!( + attempt, + scope_key = cfg.scope_key, + "auth: token refresh attempt (pkce)" + ); + let body = serde_json::json!({ + "grant_type": REFRESH_GRANT_TYPE, + "client_id": cfg.client_id, + "refresh_token": refresh_token, + }); + let resp = match post_token_json(cfg, body).await { + Ok(resp) => resp, + Err(e) => { + last_error = format!("network error: {e}"); + continue; + } + }; + let status = resp.status().as_u16(); + let bytes = resp.bytes().await.unwrap_or_default(); + if status == 401 || status == 403 { + let err: OAuthErrorBody = serde_json::from_slice(&bytes).unwrap_or_default(); + return Err(RefreshError::Unauthorized { + status, + description: err + .error_description + .unwrap_or_else(|| "Token refresh unauthorized.".to_owned()), + }); + } + if status == 200 { + return match serde_json::from_slice::(&bytes) { + Ok(tokens) => Ok(tokens.into_auth()), + Err(e) => Err(RefreshError::Fatal { + status, + description: format!("malformed token payload: {e}"), + }), + }; + } + let err: OAuthErrorBody = serde_json::from_slice(&bytes).unwrap_or_default(); + let description = err + .error_description + .unwrap_or_else(|| format!("Token refresh failed (HTTP {status}).")); + if RETRYABLE_REFRESH_STATUSES.contains(&status) { + last_error = description; + continue; + } + return Err(RefreshError::Fatal { + status, + description, + }); + } + Err(RefreshError::Exhausted { last_error }) +} + +#[derive(Deserialize, Default)] +struct OAuthErrorBody { + #[serde(default)] + error_description: Option, +} + +/// Bind a loopback HTTP listener on `127.0.0.1:{redirect_port}` and wait for a +/// single `GET /callback?code=…&state=…`, validating `state` STRICTLY against +/// `expected_state` (mismatch → rejected). Returns the authorization code. +/// +/// The listener answers ONLY `/callback`; any other path gets 404. It binds +/// `127.0.0.1` (never `0.0.0.0`), so no non-loopback host can reach it. +pub(crate) async fn await_loopback_code( + redirect_port: u16, + expected_state: &str, +) -> anyhow::Result { + let listener = tokio::net::TcpListener::bind(("127.0.0.1", redirect_port)) + .await + .with_context(|| format!("could not bind loopback 127.0.0.1:{redirect_port}"))?; + tracing::info!(port = redirect_port, "auth: pkce loopback listener bound"); + loop { + let (stream, _peer) = listener.accept().await.context("loopback accept failed")?; + match handle_loopback_conn(stream, expected_state).await { + LoopbackOutcome::Code(code) => return Ok(code), + LoopbackOutcome::Rejected(err) => return Err(err), + // Not the /callback GET (favicon, health probe): keep listening. + LoopbackOutcome::Ignore => continue, + } + } +} + +enum LoopbackOutcome { + Code(String), + Rejected(anyhow::Error), + Ignore, +} + +/// Read the request line of one loopback connection, answer with a small HTML +/// page, and classify the outcome. STRICT: a `/callback` with a bad/missing +/// state is [`LoopbackOutcome::Rejected`] (the browser sees an error page). +async fn handle_loopback_conn( + mut stream: tokio::net::TcpStream, + expected_state: &str, +) -> LoopbackOutcome { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + // Read only enough for the request line — a GET has no body. + let mut buf = [0u8; 8192]; + let n = match stream.read(&mut buf).await { + Ok(0) => return LoopbackOutcome::Ignore, + Ok(n) => n, + Err(_) => return LoopbackOutcome::Ignore, + }; + let head = String::from_utf8_lossy(&buf[..n]); + let Some(request_line) = head.lines().next() else { + return LoopbackOutcome::Ignore; + }; + // `GET /callback?code=…&state=… HTTP/1.1` + let mut parts = request_line.split_whitespace(); + let (Some(method), Some(target)) = (parts.next(), parts.next()) else { + return LoopbackOutcome::Ignore; + }; + if method != "GET" { + let _ = write_http(&mut stream, 405, "Method Not Allowed").await; + return LoopbackOutcome::Ignore; + } + let (path, query) = target.split_once('?').unwrap_or((target, "")); + if path != "/callback" { + let _ = write_http(&mut stream, 404, "Not Found").await; + return LoopbackOutcome::Ignore; + } + + let result = parse_callback_query(query) + .and_then(|params| validate_state(¶ms, expected_state).map(|()| params.code)); + match result { + Ok(code) => { + let _ = write_http( + &mut stream, + 200, + "Signed in to Claude Pro/Max. You can close this window and return to kigi.", + ) + .await; + let _ = stream.flush().await; + LoopbackOutcome::Code(code) + } + Err(e) => { + let _ = write_http(&mut stream, 400, "Login failed — return to kigi and retry.").await; + let _ = stream.flush().await; + LoopbackOutcome::Rejected(e) + } + } +} + +/// Write a minimal HTTP/1.1 response with an HTML body. +async fn write_http( + stream: &mut tokio::net::TcpStream, + status: u16, + message: &str, +) -> std::io::Result<()> { + use tokio::io::AsyncWriteExt; + let reason = match status { + 200 => "OK", + 400 => "Bad Request", + 404 => "Not Found", + 405 => "Method Not Allowed", + _ => "Error", + }; + let body = format!("

{message}

"); + let response = format!( + "HTTP/1.1 {status} {reason}\r\nContent-Type: text/html; charset=utf-8\r\n\ + Content-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + stream.write_all(response.as_bytes()).await +} + +#[cfg(test)] +mod tests { + use super::*; + use kigi_models::CLAUDE_OAUTH_CONFIG; + use wiremock::matchers::{body_string_contains, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + /// A config pointed at a mock token host (copies Claude's client_id/scope/ + /// paths but overrides the token host). + fn mock_cfg(token_host: &'static str) -> OAuthConfig { + OAuthConfig { + token_host, + ..CLAUDE_OAUTH_CONFIG + } + } + + /// PKCE codes: verifier/challenge are non-empty base64url (no padding), the + /// challenge is the base64url SHA-256 of the verifier, and state == verifier. + #[test] + fn generate_pkce_produces_valid_s256_codes() { + let pkce = generate_pkce(); + assert_eq!(pkce.state, pkce.verifier, "state must equal the verifier"); + assert!(!pkce.verifier.is_empty() && !pkce.challenge.is_empty()); + for s in [&pkce.verifier, &pkce.challenge] { + assert!( + s.chars() + .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_'), + "base64url (no pad) only: {s}" + ); + assert!(!s.contains('='), "no padding: {s}"); + } + // challenge == base64url(SHA-256(verifier)). + let expect = base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(Sha256::digest(pkce.verifier.as_bytes())); + assert_eq!(pkce.challenge, expect); + // Fresh entropy each call. + assert_ne!(pkce.verifier, generate_pkce().verifier); + } + + /// The authorize URL carries the fixed params + the PKCE state and S256 + /// challenge, and targets `claude.ai/oauth/authorize`. + #[test] + fn authorize_url_has_state_and_s256_challenge() { + let pkce = generate_pkce(); + let redirect = redirect_uri(53692); + let url = build_authorize_url(&CLAUDE_OAUTH_CONFIG, &redirect, &pkce); + let parsed = url::Url::parse(&url).expect("valid URL"); + assert_eq!(parsed.host_str(), Some("claude.ai")); + assert_eq!(parsed.path(), "/oauth/authorize"); + let q: std::collections::HashMap<_, _> = parsed.query_pairs().into_owned().collect(); + assert_eq!(q.get("response_type").map(String::as_str), Some("code")); + assert_eq!( + q.get("code_challenge_method").map(String::as_str), + Some("S256") + ); + assert_eq!( + q.get("state").map(String::as_str), + Some(pkce.state.as_str()) + ); + assert_eq!( + q.get("code_challenge").map(String::as_str), + Some(pkce.challenge.as_str()) + ); + assert_eq!( + q.get("client_id").map(String::as_str), + Some(CLAUDE_OAUTH_CONFIG.client_id) + ); + assert_eq!( + q.get("redirect_uri").map(String::as_str), + Some(redirect.as_str()) + ); + // The verifier itself must NEVER appear in the browser URL. + assert!( + !url.contains("code_verifier"), + "the verifier must not ride the authorize URL" + ); + } + + /// STRICT state validation: an exact match passes; a mismatch or an absent + /// state is REJECTED (CSRF guard — the flow must never proceed). + #[test] + fn state_validation_is_strict() { + let ok = CallbackParams { + code: "c".into(), + state: Some("expected".into()), + }; + assert!(validate_state(&ok, "expected").is_ok()); + let mismatch = CallbackParams { + code: "c".into(), + state: Some("attacker".into()), + }; + assert!( + validate_state(&mismatch, "expected").is_err(), + "a state mismatch MUST be rejected" + ); + let missing = CallbackParams { + code: "c".into(), + state: None, + }; + assert!( + validate_state(&missing, "expected").is_err(), + "an absent state MUST be rejected" + ); + } + + /// A loopback `/callback` with the WRONG state is rejected end-to-end (the + /// listener returns an error, never a code) — the CSRF guard on the wire. + #[tokio::test] + async fn loopback_rejects_state_mismatch() { + // Ephemeral port: bind, learn the port, then drive a client at it. + let probe = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .unwrap(); + let port = probe.local_addr().unwrap().port(); + drop(probe); + + let server = tokio::spawn(async move { await_loopback_code(port, "the-real-state").await }); + // Give the listener a moment to bind. + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + // Attacker callback: valid code, WRONG state. + let _ = reqwest::get(format!( + "http://127.0.0.1:{port}/callback?code=stolen&state=wrong-state" + )) + .await; + let outcome = server.await.unwrap(); + let err = outcome.expect_err("a state mismatch must be rejected, never yield a code"); + assert!(err.to_string().contains("state mismatch"), "{err}"); + } + + /// A loopback `/callback` with the MATCHING state yields the code. + #[tokio::test] + async fn loopback_returns_code_on_valid_state() { + let probe = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .unwrap(); + let port = probe.local_addr().unwrap().port(); + drop(probe); + + let server = tokio::spawn(async move { await_loopback_code(port, "good-state").await }); + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + let _ = reqwest::get(format!( + "http://127.0.0.1:{port}/callback?code=auth-code-123&state=good-state" + )) + .await; + let code = server.await.unwrap().expect("valid state yields the code"); + assert_eq!(code, "auth-code-123"); + } + + /// Manual-paste parsing: full redirect URL, `code#state`, and bare code. + #[test] + fn manual_paste_parses_all_three_forms() { + let from_url = + parse_manual_paste("http://localhost:53692/callback?code=abc123&state=st-9").unwrap(); + assert_eq!(from_url.code, "abc123"); + assert_eq!(from_url.state.as_deref(), Some("st-9")); + + let from_hash = parse_manual_paste("abc123#st-9").unwrap(); + assert_eq!(from_hash.code, "abc123"); + assert_eq!(from_hash.state.as_deref(), Some("st-9")); + + let bare = parse_manual_paste(" abc123 ").unwrap(); + assert_eq!(bare.code, "abc123"); + assert_eq!(bare.state, None); + + assert!(parse_manual_paste("").is_err()); + // A pasted redirect that carries an error param surfaces the error. + assert!(parse_manual_paste("http://localhost/callback?error=access_denied").is_err()); + } + + /// Code → token exchange: JSON body carries the grant + verifier, response + /// materializes a `KimiAuth` with the rotating refresh token. + #[tokio::test] + async fn exchange_code_posts_json_and_returns_auth() { + let server = MockServer::start().await; + let host: &'static str = Box::leak(server.uri().into_boxed_str()); + Mock::given(method("POST")) + .and(path("/v1/oauth/token")) + .and(body_string_contains( + "\"grant_type\":\"authorization_code\"", + )) + .and(body_string_contains("\"code\":\"auth-code-xyz\"")) + .and(body_string_contains("\"code_verifier\"")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "sk-ant-oat-new", + "refresh_token": "sk-ant-ort-new", + "expires_in": 3600, + "token_type": "bearer", + }))) + .expect(1) + .mount(&server) + .await; + let cfg = mock_cfg(host); + let pkce = generate_pkce(); + let auth = exchange_code(&cfg, "auth-code-xyz", &pkce, &redirect_uri(53692)) + .await + .unwrap(); + assert_eq!(auth.key, "sk-ant-oat-new"); + assert_eq!(auth.refresh_token.as_deref(), Some("sk-ant-ort-new")); + assert_eq!(auth.expires_in, Some(3600)); + } + + /// Refresh rotates the refresh token (JSON body, refresh grant). + #[tokio::test] + async fn refresh_rotates_refresh_token() { + let server = MockServer::start().await; + let host: &'static str = Box::leak(server.uri().into_boxed_str()); + Mock::given(method("POST")) + .and(path("/v1/oauth/token")) + .and(body_string_contains("\"grant_type\":\"refresh_token\"")) + .and(body_string_contains("\"refresh_token\":\"ort-old\"")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "oat-fresh", + "refresh_token": "ort-rotated", + "expires_in": 3600, + }))) + .expect(1) + .mount(&server) + .await; + let auth = refresh_token(&mock_cfg(host), "ort-old").await.unwrap(); + assert_eq!(auth.key, "oat-fresh"); + assert_eq!( + auth.refresh_token.as_deref(), + Some("ort-rotated"), + "the rotated refresh token must be adopted" + ); + } + + /// A 401 on refresh maps to Unauthorized (drives the permanent-failure path). + #[tokio::test] + async fn refresh_401_maps_to_unauthorized() { + let server = MockServer::start().await; + let host: &'static str = Box::leak(server.uri().into_boxed_str()); + Mock::given(method("POST")) + .and(path("/v1/oauth/token")) + .respond_with( + ResponseTemplate::new(401) + .set_body_json(serde_json::json!({ "error_description": "refresh revoked" })), + ) + .expect(1) + .mount(&server) + .await; + let err = refresh_token(&mock_cfg(host), "ort-dead") + .await + .unwrap_err(); + match err { + RefreshError::Unauthorized { + status, + description, + } => { + assert_eq!(status, 401); + assert_eq!(description, "refresh revoked"); + } + other => panic!("expected Unauthorized, got {other:?}"), + } + } +} diff --git a/crates/codegen/kigi-shell/src/auth/oauth_registry.rs b/crates/codegen/kigi-shell/src/auth/oauth_registry.rs index aa7b8bc..0ca3af6 100644 --- a/crates/codegen/kigi-shell/src/auth/oauth_registry.rs +++ b/crates/codegen/kigi-shell/src/auth/oauth_registry.rs @@ -141,6 +141,52 @@ mod tests { .expect("xai-grok carries an OAuthConfig") } + fn claude_oauth() -> &'static kigi_models::OAuthConfig { + kigi_models::PlatformId::ClaudeProMax + .oauth() + .expect("claude-pro-max carries an OAuthConfig") + } + + /// A `claude-pro-max/` turn resolves to the process-global pooled + /// claude-pro-max manager (its OWN `oauth/claude-pro-max` scope), NEVER the + /// primary Kimi manager — the same leak-safe routing as xai-grok, and a + /// DISTINCT pool entry from the xai manager. + #[tokio::test] + async fn claude_pro_max_model_resolves_to_its_own_manager_not_kimi() { + let (_kd, kimi) = primary_with_token("kimi-tok"); + let home = tempfile::tempdir().unwrap(); + let resolved = + manager_for_model(home.path(), "claude-pro-max/claude-opus-4-8", Some(&kimi)) + .expect("claude-pro-max model resolves to its pooled manager"); + assert!( + !Arc::ptr_eq(&resolved, &kimi), + "claude-pro-max must NOT resolve to the Kimi manager" + ); + assert!( + Arc::ptr_eq(&resolved, &global_manager_for(home.path(), claude_oauth())), + "claude-pro-max must resolve to its OWN process-global pooled manager" + ); + // And it is a DIFFERENT manager than xai-grok's pooled one. + assert!( + !Arc::ptr_eq(&resolved, &global_manager_for(home.path(), xai_oauth())), + "claude-pro-max and xai-grok must not share a pooled manager" + ); + } + + /// Fail-fast (no Kimi fallback): a claude-pro-max key with a Kimi primary + /// never yields the Kimi session token — it draws from the claude pool (its + /// own token, or `None`), so the Kimi bearer can never reach api.anthropic. + #[tokio::test] + async fn session_key_for_claude_pro_max_is_never_the_kimi_primary() { + let (_kd, kimi) = primary_with_token("kimi-tok"); + let home = tempfile::tempdir().unwrap(); + assert_ne!( + session_key_for_model(home.path(), "claude-pro-max/claude-opus-4-8", Some(&kimi)), + Some("kimi-tok".to_string()), + "a claude-pro-max model must never receive the primary Kimi session token" + ); + } + /// A non-OAuth managed key (moonshot-cn/…) and an unprefixed bare id both /// route to the primary Kimi manager — the Kimi / first-party path is /// untouched and never consults the pool (no runtime needed). diff --git a/crates/codegen/kigi-shell/src/auth/refresh/generic_refresher.rs b/crates/codegen/kigi-shell/src/auth/refresh/generic_refresher.rs index eff9be0..cc41249 100644 --- a/crates/codegen/kigi-shell/src/auth/refresh/generic_refresher.rs +++ b/crates/codegen/kigi-shell/src/auth/refresh/generic_refresher.rs @@ -9,12 +9,12 @@ use std::sync::Arc; -use kigi_models::OAuthConfig; +use kigi_models::{OAuthConfig, OAuthTokenBody}; use crate::auth::error::RefreshTokenFailedReason; use crate::auth::kimi_oauth::RefreshError; use crate::auth::manager::RefreshReason; -use crate::auth::oauth_device::{self}; +use crate::auth::{oauth_device, oauth_pkce}; use super::{AuthSnapshot, RefreshOutcome, TokenRefresher}; @@ -104,7 +104,14 @@ impl TokenRefresher for GenericDeviceRefresher { "auth: sending refresh_token grant (generic oauth)" ); - match oauth_device::refresh_token(self.cfg, &refresh_token).await { + // Refresh over the provider's token-body encoding: xai's endpoint is + // form-encoded (device wire); Claude's is JSON (PKCE wire). Both return + // the same `Result`. + let wire_result = match self.cfg.token_body { + OAuthTokenBody::Form => oauth_device::refresh_token(self.cfg, &refresh_token).await, + OAuthTokenBody::Json => oauth_pkce::refresh_token(self.cfg, &refresh_token).await, + }; + match wire_result { Ok(new_auth) => { kigi_log::unified_log::info( "auth.refresh.token_rotated", diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs index 60cc8ff..233d038 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs @@ -259,6 +259,21 @@ impl SessionActor { .and_then(|(platform, _)| platform.oauth()) .is_some() } + /// Whether `model` routes to the Claude Pro/Max OAuth-Messages platform + /// (claude-pro-max) — the gate for the sampler's OAuth Messages adaptation + /// (identity headers + "You are Claude Code" system prefix). A generic-OAuth + /// platform speaking the Messages wire; every other model (incl. xai-grok, + /// which is ChatCompletions) returns `false`, keeping the API-key Anthropic + /// / MiniMax Messages requests byte-identical. + fn model_is_anthropic_oauth(&self, model: &str) -> bool { + let managed_key = self.managed_key_for_model(model); + kigi_models::parse_managed_model_key(managed_key.as_deref().unwrap_or(model)).is_some_and( + |(platform, _)| { + platform.oauth().is_some() + && platform.wire_api() == kigi_models::PlatformWireApi::Messages + }, + ) + } /// LEAK guard for the stamped aux paths (auto-mode classifier, image /// describe). After [`crate::agent::config::stamp_session_local_sampler_fields`] /// has copied the SESSION model's `bearer_resolver` onto an aux @@ -365,6 +380,9 @@ impl SessionActor { } else { None }; + // Claude Pro/Max OAuth Messages adaptation for THIS turn's model + // (captured before `cfg.model` is moved into the struct below). + let anthropic_oauth = self.model_is_anthropic_oauth(&cfg.model); let auth_scheme = model_facts.auth_scheme; let mut extra_headers = cfg.extra_headers; crate::agent::config::inject_url_derived_headers( @@ -405,6 +423,7 @@ impl SessionActor { top_p: cfg.top_p, api_backend: cfg.api_backend, auth_scheme, + anthropic_oauth, chat_compat: cfg.chat_compat, extra_headers, context_window: cfg.context_window.get(), diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs index 0ce2311..da7a0a9 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs @@ -849,6 +849,7 @@ async fn set_session_model_invalidates_byok_memo_for_same_model_id() { api_backend: crate::sampling::ApiBackend::ChatCompletions, chat_compat: Default::default(), auth_scheme: Default::default(), + anthropic_oauth: false, extra_headers: Default::default(), context_window: 256_000, force_http1: false, diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs index e9155e7..ad60b01 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs @@ -46,6 +46,7 @@ async fn persist_ack_waits_for_disk_flush_before_success() { api_backend: Default::default(), chat_compat: Default::default(), auth_scheme: Default::default(), + anthropic_oauth: false, extra_headers: Default::default(), context_window: 100_000, force_http1: false, @@ -342,6 +343,7 @@ async fn first_turn_memory_injection_persists_to_chat_history() { api_backend: Default::default(), chat_compat: Default::default(), auth_scheme: Default::default(), + anthropic_oauth: false, context_window: 100_000, force_http1: false, max_retries: None, @@ -472,6 +474,7 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history() api_backend: Default::default(), chat_compat: Default::default(), auth_scheme: Default::default(), + anthropic_oauth: false, context_window: 100_000, force_http1: false, max_retries: None, @@ -1740,6 +1743,7 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() { api_backend: kigi_sampler::ApiBackend::Responses, chat_compat: Default::default(), auth_scheme: Default::default(), + anthropic_oauth: false, extra_headers: Default::default(), context_window: 100_000, force_http1: false, diff --git a/crates/codegen/kigi-shell/src/session/helpers/session_compact.rs b/crates/codegen/kigi-shell/src/session/helpers/session_compact.rs index 0cf7dbf..58d2727 100644 --- a/crates/codegen/kigi-shell/src/session/helpers/session_compact.rs +++ b/crates/codegen/kigi-shell/src/session/helpers/session_compact.rs @@ -1590,6 +1590,7 @@ mod reasoning_compaction_regression_tests { api_backend: ApiBackend::ChatCompletions, chat_compat: Default::default(), auth_scheme: Default::default(), + anthropic_oauth: false, extra_headers: Default::default(), context_window: 256_000, force_http1: false, diff --git a/crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs b/crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs index 0786b8b..3e10b0c 100644 --- a/crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs +++ b/crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs @@ -46,6 +46,7 @@ pub(crate) fn ctx_with_toggle(toggle: HashMap) -> SubagentSpawnCon api_backend: Default::default(), chat_compat: Default::default(), auth_scheme: Default::default(), + anthropic_oauth: false, extra_headers: Default::default(), context_window: 256_000, force_http1: false, diff --git a/crates/codegen/kigi-shell/tests/common/mod.rs b/crates/codegen/kigi-shell/tests/common/mod.rs index c2f7217..1cfe592 100644 --- a/crates/codegen/kigi-shell/tests/common/mod.rs +++ b/crates/codegen/kigi-shell/tests/common/mod.rs @@ -38,6 +38,7 @@ pub fn test_sampler_config( top_p: None, api_backend, auth_scheme: Default::default(), + anthropic_oauth: false, chat_compat: Default::default(), extra_headers: extra_headers .iter() diff --git a/crates/codegen/kigi-tui/src/app/app_view.rs b/crates/codegen/kigi-tui/src/app/app_view.rs index 6ae6233..dde44fd 100644 --- a/crates/codegen/kigi-tui/src/app/app_view.rs +++ b/crates/codegen/kigi-tui/src/app/app_view.rs @@ -6894,7 +6894,7 @@ pub(crate) mod tests { #[test] fn pending_menu_items_lists_interactive_methods_plus_quit() { let items = pending_menu_items(&fresh_user_auth_methods(), None); - assert_eq!(items.len(), 27, "26 login rows + Quit, got {items:?}"); + assert_eq!(items.len(), 28, "27 login rows + Quit, got {items:?}"); assert!( matches!(&items[0], PendingMenuItem::Login { label } if label == "Kimi Code (OAuth)"), "row 0 must be the OAuth login, got {:?}", @@ -6905,22 +6905,27 @@ pub(crate) mod tests { "row 1 must be the xai-grok OAuth login (interactive, after kimi-code), got {:?}", items[1] ); + assert!( + matches!(&items[2], PendingMenuItem::Login { label } if label == "Claude Pro/Max (subscription) (OAuth)"), + "row 2 must be the claude-pro-max OAuth login (after xai-grok), got {:?}", + items[2] + ); assert_eq!( - items[2], + items[3], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::MoonshotCn), label: "Moonshot Open Platform (API key \u{b7} moonshot.cn)".into(), } ); assert_eq!( - items[3], + items[4], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::MoonshotAi), label: "Moonshot Open Platform (API key \u{b7} moonshot.ai)".into(), } ); assert_eq!( - items[4], + items[5], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::OpenAi), label: "OpenAI (API key)".into(), @@ -6928,153 +6933,153 @@ pub(crate) mod tests { "new registry rows must appear in the picker with zero TUI changes" ); assert_eq!( - items[5], + items[6], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Anthropic), label: "Anthropic (API key)".into(), } ); assert_eq!( - items[6], + items[7], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::DeepSeek), label: "DeepSeek (API key)".into(), } ); assert_eq!( - items[7], + items[8], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Groq), label: "Groq (API key)".into(), } ); assert_eq!( - items[8], + items[9], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Mistral), label: "Mistral (API key)".into(), } ); assert_eq!( - items[9], + items[10], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Fireworks), label: "Fireworks AI (API key)".into(), } ); assert_eq!( - items[10], + items[11], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Google), label: "Google Gemini (API key)".into(), } ); assert_eq!( - items[11], + items[12], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::OpenRouter), label: "OpenRouter (API key)".into(), } ); assert_eq!( - items[12], + items[13], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Together), label: "Together AI (API key)".into(), } ); assert_eq!( - items[13], + items[14], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Cerebras), label: "Cerebras (API key)".into(), } ); assert_eq!( - items[14], + items[15], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Nvidia), label: "NVIDIA NIM (API key)".into(), } ); assert_eq!( - items[15], + items[16], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Vercel), label: "Vercel AI Gateway (API key)".into(), } ); assert_eq!( - items[16], + items[17], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Xai), label: "xAI (Grok) (API key)".into(), } ); assert_eq!( - items[17], + items[18], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::QwenTokenPlan), label: "Qwen Token Plan (API key)".into(), } ); assert_eq!( - items[18], + items[19], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::QwenTokenPlanCn), label: "Qwen Token Plan China (API key)".into(), } ); assert_eq!( - items[19], + items[20], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::KimiCoding), label: "Kimi For Coding (API key)".into(), } ); assert_eq!( - items[20], + items[21], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Zai), label: "Z.AI (API key)".into(), } ); assert_eq!( - items[21], + items[22], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::ZaiCodingCn), label: "Z.AI Coding China (API key)".into(), } ); assert_eq!( - items[22], + items[23], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Xiaomi), label: "Xiaomi MiMo (API key)".into(), } ); assert_eq!( - items[23], + items[24], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::XiaomiTokenPlanCn), label: "Xiaomi Token Plan China (API key)".into(), } ); assert_eq!( - items[24], + items[25], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Minimax), label: "MiniMax (API key)".into(), } ); assert_eq!( - items[25], + items[26], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::MinimaxCn), label: "MiniMax China (API key)".into(), } ); - assert_eq!(items[26], PendingMenuItem::Quit); + assert_eq!(items[27], PendingMenuItem::Quit); // The non-interactive methods must never appear as rows. let byok = kigi_shell::agent::auth_method::build_auth_methods( kigi_shell::agent::auth_method::AuthMethodsBuildInputs { @@ -7085,8 +7090,8 @@ pub(crate) mod tests { ); assert_eq!( pending_menu_items(&byok.methods, None).len(), - 27, - "xai.api_key / cached_token must not add rows (26 login rows + Quit)" + 28, + "xai.api_key / cached_token must not add rows (27 login rows + Quit)" ); } /// Startup lands on the picker only when there is a real choice: the @@ -7108,7 +7113,9 @@ pub(crate) mod tests { app.auth_state = AuthState::Pending { error: None }; app.welcome_prompt_focused = false; // Interactive OAuth logins come first: row 0 (kimi-code), row 1 - // (xai-grok); the first API-key row (moonshot-cn) is now row 2. + // (xai-grok), row 2 (claude-pro-max); the first API-key row + // (moonshot-cn) is now row 3, so four Downs land on it. + app.handle_input(&key_event(KeyCode::Down, KeyModifiers::NONE)); app.handle_input(&key_event(KeyCode::Down, KeyModifiers::NONE)); app.handle_input(&key_event(KeyCode::Down, KeyModifiers::NONE)); app.handle_input(&key_event(KeyCode::Down, KeyModifiers::NONE)); @@ -7120,7 +7127,7 @@ pub(crate) mod tests { kigi_shell::models::PlatformId::MoonshotCn ))) ), - "Enter on row 1 must open moonshot-cn key entry, got {outcome:?}" + "Enter on the moonshot-cn row must open its key entry, got {outcome:?}" ); // 'l' is muscle-memory for the first (OAuth) row regardless of the // arrow selection. diff --git a/crates/codegen/kigi-tui/src/views/welcome/mod.rs b/crates/codegen/kigi-tui/src/views/welcome/mod.rs index a7aaeec..45c4233 100644 --- a/crates/codegen/kigi-tui/src/views/welcome/mod.rs +++ b/crates/codegen/kigi-tui/src/views/welcome/mod.rs @@ -2106,6 +2106,14 @@ mod tests { // lists ~20 platforms); this asserts content coverage, not fit. let text = render_done_text_h(¶ms, 72); assert!(text.contains("Kimi Code (OAuth)"), "{text}"); + assert!( + text.contains("xAI Grok (subscription) (OAuth)"), + "the xai-grok interactive OAuth login row must render: {text}" + ); + assert!( + text.contains("Claude Pro/Max (subscription) (OAuth)"), + "the claude-pro-max interactive OAuth login row must render: {text}" + ); assert!( text.contains("Moonshot Open Platform (API key \u{b7} moonshot.cn)"), "{text}"