From 2d00a4e6e6f1979b7cc21fc3677e5523002d6cf8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Wed, 22 Jul 2026 11:08:54 -0400 Subject: [PATCH] update --- AGENTS.md | 40 ++ README.md | 150 +++--- .../kigi-shell/src/agent/auth_method.rs | 166 ++++++- crates/codegen/kigi-shell/src/agent/config.rs | 80 +++- crates/codegen/kigi-shell/src/agent/models.rs | 91 ++++ .../kigi-shell/src/agent/models_fetch.rs | 62 ++- .../src/agent/mvp_agent/agent_ops.rs | 80 +++- .../kigi-shell/src/agent/mvp_agent/tests.rs | 21 +- .../tests/api_key_channel_leak_tests.rs | 313 +++++++++++++ .../kigi-shell/src/agent/subagent/mod.rs | 26 +- .../src/agent/subagent/tests/mod.rs | 20 +- .../src/agent/subagent/tests/rest.rs | 50 +- .../kigi-shell/src/auth/oauth_registry.rs | 187 ++++++-- .../kigi-shell/src/session/acp_session.rs | 13 +- .../session/acp_session_impl/prompt_build.rs | 12 +- .../session/acp_session_impl/sampler_turn.rs | 189 +++++--- .../auth_error_no_retry_tests.rs | 26 +- .../session_bearer_leak_platform_tests.rs | 289 ++++++++++++ .../session_bearer_leak_tests.rs | 436 ++++++++++++++++++ 19 files changed, 1955 insertions(+), 296 deletions(-) create mode 100644 crates/codegen/kigi-shell/src/agent/mvp_agent/tests/api_key_channel_leak_tests.rs create mode 100644 crates/codegen/kigi-shell/src/session/acp_session_tests/session_bearer_leak_platform_tests.rs create mode 100644 crates/codegen/kigi-shell/src/session/acp_session_tests/session_bearer_leak_tests.rs diff --git a/AGENTS.md b/AGENTS.md index 432aed0..62b43c6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -259,6 +259,46 @@ edges stay deterministic Rust. The harness appends a terminal 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. +- INFERENCE-AUTH INVARIANT (security): a SESSION bearer may ride a request ONLY + when the endpoint's platform `uses_oauth()` — `kimi-code` (the primary + session) or one of the four subscription-OAuth platforms (their OWN pooled + `AuthManager` via `oauth_registry::manager_for_model`, which is why they keep + a live `bearer_resolver` and mid-session refresh despite non-first-party base + URLs). Every API-key registry platform is refused, because + `manager_for_model` falls through to the PRIMARY manager for a non-OAuth key + and `SamplingClient::post` REPLACES the request's auth header from the + resolver. A model with NO platform (a bare slug / `[model.*]` entry) is + decided by the ENDPOINT — `util::is_effective_coding_endpoint_url` (the + effective `KIGI_CODE_BASE_URL` deployment, loopback dev proxies, and the + compiled production endpoint), never a blanket allow: BYOK is + `has_own_credentials()`, which probes `std::env::var` at call time, so a + `[model.*]` block with an unset `env_key` classifies `NotByok`. Both are + `auth_method::platform_takes_session_credential(platform, base_url)`, the + single predicate enforced at BOTH channels that reach the wire — the + `bearer_resolver` (outer term of `auth_method::session_token_auth_gate`) and + the `api_key` (`MvpAgent::session_token_for_model`, + `oauth_registry::session_key_for_endpoint` for the aux / summary / + subagent-override paths, and `sampler_turn::aux_bearer_resolver` for the + stamped aux configs). +- MODEL→PLATFORM LOOKUP (security): `SamplingConfig::model` is the BARE routing + slug, and duplicate slugs across platforms are BY DESIGN — an API-key platform + and its subscription-OAuth twin list identical ids (`xai`/`xai-grok`, + `anthropic`/`claude-pro-max`, `openai`/`openai-codex`), with the API-key + platform FIRST in `PlatformId::ALL`. The auth layer therefore resolves the + platform from the catalog KEY the picker selected + (`ModelsManager::current_model_id`), via + `agent::models::managed_key_for_slug`/`platform_for_slug`; anything else falls + back to the picker's own `resolve_catalog_key`, and + `config::find_model_by_id`'s slug scan takes the LAST match so the two can + never disagree. Resolving the wrong twin costs the OAuth platform its live + `bearer_resolver` (unrecoverable 401 ~1h in), its Messages adaptation and its + Copilot/Codex identity headers. +- CATALOG VISIBILITY: `platform_wire_model_to_entry` stamps + `supported_in_api = platform != KimiCode`. `ModelInfo::visible_for_auth` + reads only the PRIMARY manager's auth mode, so gating the other OAuth + platforms on it would hide every model from a user who signed in with ONLY a + Claude Pro/Max, ChatGPT, Copilot, or Grok subscription. Only `kimi-code` + rides the primary session, so only it may be gated on it. - 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/README.md b/README.md index 6ce417b..3321595 100644 --- a/README.md +++ b/README.md @@ -8,19 +8,14 @@ autonomous, self-verifying agent loops — planned, parallelized, adversarially verified, and merged back, end to end.

-**Kigi** started as an unofficial Kimi Code CLI community build, a -terminal-based AI coding agent re-targeted at the Kimi Code subscription API, -built on the Apache-2.0 sources of -[xai-org/grok-build](https://github.com/xai-org/grok-build). It first shipped -wired to Kimi Code and the Moonshot open platform. By request, it now works -with 25 providers: OpenAI, Anthropic, Google, xAI, Groq, Cerebras, OpenRouter, -MiniMax, Z.AI, Qwen, Xiaomi, and more. The full list is under -[Providers and API keys](#providers-and-api-keys). +**Kigi** is a coding agent that lives in your terminal. It reads the repo, +writes the patch, runs the tests, and keeps going while you do something else. +Full-screen, headless in CI with `-p`, or docked in your editor over ACP. -It runs as a full-screen TUI that understands your codebase, edits files, -executes shell commands, searches the web, and manages long-running tasks, -interactively, headlessly for scripting/CI, or embedded in editors via the -Agent Client Protocol (ACP). +**Already paying for Claude Pro/Max, ChatGPT Plus/Pro, GitHub Copilot, or +Grok? Sign in and use it.** No API key, no second bill. Rather bring your own +key? OpenAI, Anthropic, Google, DeepSeek, Groq, Moonshot and +[two dozen more](#providers-and-api-keys) are wired in. [Installation](#installation) · [Graph engineering](#graph-engineering) · @@ -39,10 +34,6 @@ Agent Client Protocol (ACP). ## Installation -Prebuilt single-file binaries for macOS (arm64/x86_64), Linux (arm64/x86_64), -and Windows (x86_64) are published on -[GitHub Releases](https://github.com/ZacharyZhang-NY/Kigi-CLI/releases): - ```sh # macOS / Linux curl -fsSL https://raw.githubusercontent.com/ZacharyZhang-NY/Kigi-CLI/main/install.sh | bash @@ -54,24 +45,19 @@ irm https://raw.githubusercontent.com/ZacharyZhang-NY/Kigi-CLI/main/install.ps1 ``` ```sh -kigi --version # kigi 0.1.1 … unofficial Kimi Code CLI community build -kigi login # sign in with your Kimi Code subscription (device-code flow) -kigi # start the TUI +kigi login # pick a provider, sign in +kigi # go ``` -The installer verifies every download against the release's `SHA256SUMS`, -installs into `~/.kigi/bin/kigi` (`%USERPROFILE%\.kigi\bin\kigi.exe` on -Windows), persists the PATH line for you, and **enables graph engineering -by default** (`KIGI_GRAPH=1`; see [Graph engineering](#graph-engineering) -to disable). Later releases arrive through the -built-in self-updater (`kigi update`, gated by `KIGI_AUTO_UPDATE`), which -pulls from the same GitHub Releases feed. +Single file, no runtime. macOS and Linux on arm64/x86_64, Windows on x86_64, +checksummed against the release's `SHA256SUMS`. `kigi update` handles upgrades. ## Graph engineering -Kigi is the first CLI to ship *graph engineering* as a first-class command: -where a loop drives one agent, a graph is the programmable organization -connecting many. +Every other agent runs a loop: think, act, repeat — one thread, one thing at a +time. `/graph` runs a dependency graph instead. Work that doesn't block other +work happens at the same time, in separate worktrees, and nothing merges until +something else has tried to tear it apart. ``` /graph [--budget ] # decompose + run fully autonomously @@ -81,49 +67,46 @@ connecting many. /graph clear # abandon the graph ``` -One `/graph ` runs the whole closed loop: a planner subagent -decomposes the objective into a validated dependency DAG; independent -nodes fan out as parallel workers in isolated git worktrees, each gated -by an adversarial verifier and merged back three-way; out-of-scope -discoveries (`DISCOVERED:`) replan the graph append-only; a topology -optimizer prunes false dependencies at plan boundaries; and a terminal -verification node re-checks the *whole* objective before the graph -completes. State follows your repo in `.kigi/graph.jsonl`, so a fresh -session — or a teammate — can `/graph resume` where you left off. +One command runs the whole thing, start to finish: -The installer enables it by default. To disable: +- A planner breaks your objective into a dependency DAG, then validates it. +- Independent nodes fan out as parallel workers, each in its own git worktree. +- Every node has to get past an adversarial verifier before it merges back. +- Find something out of scope? Say `DISCOVERED:` and the graph replans — + append-only, so nothing already agreed on gets rewritten. +- Between passes, a topology optimizer drops dependencies that were never real. +- A final node re-checks the *whole* objective before the graph is allowed to + call itself done. -```sh -# macOS / Linux -echo 'export KIGI_GRAPH=0' >> ~/.zshrc # or ~/.bashrc / ~/.bash_profile -``` +State lives in `.kigi/graph.jsonl`, next to your code. Close the laptop, come +back tomorrow, `/graph resume`. A teammate can pick it up from the same file. -```powershell -# Windows PowerShell -[Environment]::SetEnvironmentVariable('KIGI_GRAPH','0','User') -``` - -(One-off instead: `KIGI_GRAPH=0 kigi`.) Tuning knobs: -`KIGI_GRAPH_CONCURRENCY` (parallel nodes, default 3), -`KIGI_GRAPH_NODE_ROUNDS` (worker↔verifier rounds per node, default 3), -`KIGI_GRAPH_REPLAN_CAP` (replan passes, default 3), -`KIGI_GRAPH_OPTIMIZER=0` (disable the optimizer pass). +On by default. `KIGI_GRAPH=0` turns it off; `KIGI_GRAPH_CONCURRENCY` (default +3) controls how many nodes run at once. ## Providers and API keys -Kigi ships a fixed registry of 25 platforms: the Kimi Code subscription plus -24 API-key providers. There is no dynamic provider registration; each is a -compiled-in spec. +29 platforms ship compiled in: 5 you sign into, 24 you hand a key. Nothing is +registered at runtime — if it's not in this list, it's not there. -**Kimi Code** (the original target) uses subscription OAuth, not an API key: +**Sign in with a subscription you already pay for.** Run `kigi login` and pick. +Each provider's token is stored under its own key, and one provider's +credentials are never sent to another. -| Platform id | Base URL | Auth | -| ----------- | -------------------------------- | ------------------------------------------- | -| `kimi-code` | `https://api.kimi.com/coding/v1` | Kimi Code subscription OAuth (`kigi login`) | +| Platform id | Provider | Sign-in | +| ---------------- | ------------------------- | ------------------------------------------- | +| `kimi-code` | Kimi Code (original target)| Subscription OAuth (device code) | +| `claude-pro-max` | Claude Pro/Max | Subscription OAuth (browser, PKCE) | +| `openai-codex` | ChatGPT Plus/Pro (Codex) | Subscription OAuth (browser, PKCE) | +| `github-copilot` | GitHub Copilot | Subscription OAuth (device code) | +| `xai-grok` | xAI Grok | Subscription OAuth (device code) | -**API-key providers.** Set the provider's env var, or put the key in -`~/.kigi/config.toml` under `[platforms.]`. The environment wins, a -platform-scoped name beats a generic one, and keys are never logged. +You get whatever models your plan actually serves — the list is fetched at +sign-in, not hardcoded. (ChatGPT/Codex is the exception: its backend publishes +no model endpoint, so those four are compiled in.) + +**API-key providers.** Export the env var, or drop the key in +`~/.kigi/config.toml`. Keys are never logged. | Provider | Platform id | API key env | | ------------------------- | ---------------------- | ----------------------------------------------- | @@ -166,24 +149,11 @@ api_key = "sk-..." api_key = "xai-..." ``` -On login and on startup Kigi syncs each configured platform's model list -from `GET {base}/models` and shows the merged catalog in the model picker -(catalog keys are `{platform_id}/{model_id}`). Model metadata (context -window, thinking levels) comes from the live listing when the provider -serves it, otherwise from a bundled models.dev snapshot. Models that -advertise selectable thinking levels (e.g. K3's `low`/`high`/`max`) expose -them in `/model` and `/effort`. If the sync fails, the last cached catalog is -used; with no cache, a small built-in fallback list applies. Model selection -resolves as `--model` CLI flag > `KIGI_DEFAULT_MODEL` > `[models] default` -in config.toml > server-delivered list > built-in fallback. +Model lists sync on startup. Pick one with `/model`, set its thinking level +with `/effort`. -Each platform's base URL can be re-pointed for dev/test with -`KIGI__BASE_URL` (e.g. `KIGI_CODE_BASE_URL`, -`KIGI_MOONSHOT_CN_BASE_URL`, `KIGI_OPENAI_BASE_URL`). - -The web `search`/`fetch` tools ride the Kimi Code subscription services and -are present only on OAuth sessions. API-key-only sessions run without them, -matching the official client. +Web `search`/`fetch` need a Kimi Code subscription; API-key sessions run +without them, same as the official client. ## Building from source @@ -199,19 +169,15 @@ launcher at `bin/protoc`; install dotslash (`brew install dotslash` or ## Coexistence with the official Kimi CLI -Kigi is not affiliated with Moonshot AI or xAI, and it coexists with the -official `kimi` CLI on the same machine: independent binary name, -independent config directory (`~/.kigi`), independent keyring credentials -(service `kigi`), and a `KIGI_*` environment-variable namespace. Nothing -the official client installs or stores is ever read at runtime or written. -On first launch Kigi offers a **one-time, strictly read-only** import of -your existing `~/.kimi` configuration (MCP servers, custom providers, -default model) via `kigi import-kimi` — file contents and mtimes under -`~/.kimi` are left untouched, verified by tests. +Kigi started as an unofficial Kimi Code CLI — a community fork of +[xai-org/grok-build](https://github.com/xai-org/grok-build), not affiliated +with Moonshot AI or xAI. It keeps its own binary, its own `~/.kigi`, its own +keyring entry, and its own `KIGI_*` env vars, and never touches what the +official `kimi` CLI installed. `kigi import-kimi` copies your old config over +once, read-only. -Kigi is **zero-telemetry**: the only outbound connections are the -inference/auth APIs you configure, GitHub Releases for updates, and MCP -servers you add. +**Zero telemetry.** It talks to the APIs you configured, GitHub Releases, and +your own MCP servers. Nothing else. ## License diff --git a/crates/codegen/kigi-shell/src/agent/auth_method.rs b/crates/codegen/kigi-shell/src/agent/auth_method.rs index 3b45dba..6b2b371 100644 --- a/crates/codegen/kigi-shell/src/agent/auth_method.rs +++ b/crates/codegen/kigi-shell/src/agent/auth_method.rs @@ -302,14 +302,26 @@ impl ModelByok { /// buffered token on every turn and 401s with `bad-credentials` until restart. /// It refreshes when `endpoint_is_first_party` — the request targets the /// first-party API, where sending the session token cannot leak to a -/// third-party BYOK endpoint. A definite `NotByok` always refreshes (it only -/// ever routes to the session endpoint); a definite `Byok` never does. +/// third-party BYOK endpoint. A definite `NotByok` refreshes (within the +/// session's own providers); a definite `Byok` never does. +/// +/// `endpoint_takes_session_credential` ([`platform_takes_session_credential`]) +/// is the outer, non-negotiable guard and the reason the `NotByok` arm is safe: +/// it is `false` for every API-key registry platform (deepseek, openai, +/// anthropic, moonshot-*, …), whose models classify `NotByok` (they carry no +/// `[model.*]` key) yet route to a THIRD-PARTY inference host while +/// `oauth_registry::manager_for_model` falls through to the primary (Kimi) +/// manager. Without this term the mainstream "session auth + API-key-platform +/// model" configuration stamps the user's Kimi subscription bearer on every +/// request to that host. pub fn session_token_auth_gate( is_session_based_method: bool, model_byok: ModelByok, endpoint_is_first_party: bool, + endpoint_takes_session_credential: bool, ) -> bool { is_session_based_method + && endpoint_takes_session_credential && match model_byok { ModelByok::NotByok => true, ModelByok::Byok => false, @@ -317,6 +329,42 @@ pub fn session_token_auth_gate( } } +/// Whether a session bearer may EVER be stamped on a request routed to +/// `platform` at `base_url` — i.e. whether the credential +/// `oauth_registry::manager_for_model` resolves for such a model belongs to the +/// host that receives it. +/// +/// - a `uses_oauth` platform: `kimi-code` rides the PRIMARY session; each of the +/// four subscription-OAuth platforms (claude-pro-max, openai-codex, +/// github-copilot, xai-grok) rides its OWN pooled `AuthManager`. In both cases +/// the resolved bearer belongs to the host being called, and mid-session +/// refresh/401-recovery must stay live — so this is `true` even though those +/// four have non-first-party base URLs. +/// - every API-key registry platform: its credential is that platform's API key +/// (resolved into the catalog entry), never a session bearer, and +/// `manager_for_model` has no platform manager to route to. `false`. +/// - `None` — a bare slug or a `[model.*]` config entry, which carries no +/// platform at all — is decided by the ENDPOINT, never blanket-allowed: BYOK +/// is `has_own_credentials()`, which probes `std::env::var` at call time, so a +/// `[model.gpt-4o]` block with `base_url = "https://api.openai.com/v1"` and an +/// unset / mistyped `env_key` classifies `NotByok` and would otherwise hand +/// the Kimi subscription bearer to `api.openai.com`. Only the SESSION's own +/// coding endpoint qualifies: [`crate::util::is_effective_coding_endpoint_url`] +/// = the *effective* `KIGI_CODE_BASE_URL` deployment (so custom Kimi +/// deployments keep the session bearer) plus loopback (local dev proxies and +/// test mocks) plus the compiled production endpoint. Deliberately NOT +/// `is_first_party_url`, which is production-only and would break every +/// `KIGI_CODE_BASE_URL` deployment. +pub fn platform_takes_session_credential( + platform: Option, + base_url: &str, +) -> bool { + match platform { + Some(platform) => platform.uses_oauth(), + None => crate::util::is_effective_coding_endpoint_url(base_url), + } +} + pub const AUTH_ERROR_SESSION_EXPIRED: &str = "Session expired. Run `kigi login` to re-authenticate."; @@ -683,15 +731,117 @@ mod tests { #[test] fn session_token_auth_gate_matrix() { - // Session method + NotByok → refresh. - assert!(session_token_auth_gate(true, ModelByok::NotByok, false)); + // Session method + NotByok → refresh (endpoint takes the session cred). + assert!(session_token_auth_gate( + true, + ModelByok::NotByok, + false, + true + )); // Session method + Byok → never. - assert!(!session_token_auth_gate(true, ModelByok::Byok, true)); + assert!(!session_token_auth_gate(true, ModelByok::Byok, true, true)); // Session method + Unknown → only on first-party endpoints. - assert!(session_token_auth_gate(true, ModelByok::Unknown, true)); - assert!(!session_token_auth_gate(true, ModelByok::Unknown, false)); + assert!(session_token_auth_gate( + true, + ModelByok::Unknown, + true, + true + )); + assert!(!session_token_auth_gate( + true, + ModelByok::Unknown, + false, + true + )); // Non-session method → never. - assert!(!session_token_auth_gate(false, ModelByok::NotByok, true)); + assert!(!session_token_auth_gate( + false, + ModelByok::NotByok, + true, + true + )); + // An endpoint that does not take the session credential (every API-key + // registry platform) is refused on EVERY arm — this is the outer guard + // that keeps the Kimi subscription bearer off third-party hosts. + for byok in [ModelByok::NotByok, ModelByok::Byok, ModelByok::Unknown] { + for first_party in [false, true] { + assert!( + !session_token_auth_gate(true, byok, first_party, false), + "byok={byok:?} first_party={first_party}: an API-key-platform \ + endpoint must never receive a session bearer" + ); + } + } + } + + /// The platform → "may a session bearer ride here?" classification. + /// `kimi-code` rides the PRIMARY session; the four subscription-OAuth + /// platforms ride their OWN pooled managers (so they keep a live + /// bearer_resolver despite non-first-party base URLs); every API-key + /// registry platform is refused, whatever the endpoint. + #[test] + fn platform_takes_session_credential_matrix() { + let first_party = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url; + for id in [ + "kimi-code", + "claude-pro-max", + "openai-codex", + "github-copilot", + "xai-grok", + ] { + let platform = kigi_models::PlatformId::parse(id).expect("known platform"); + for url in [first_party, "https://api.anthropic.com/v1"] { + assert!( + platform_takes_session_credential(Some(platform), url), + "{id} rides a session credential (primary or its own pool)" + ); + } + } + for platform in kigi_models::PlatformId::ALL { + if platform.uses_oauth() { + continue; + } + for url in [first_party, "https://api.deepseek.com/v1"] { + assert!( + !platform_takes_session_credential(Some(platform), url), + "{} is an API-key platform — no session bearer may ride to it", + platform.as_str() + ); + } + } + } + + /// C2 regression: a platform-less model (a bare slug or a `[model.*]` entry) + /// is decided by the ENDPOINT, never blanket-allowed. A `[model.gpt-4o]` + /// block whose `env_key` is unset classifies `NotByok`, so before this the + /// `None` arm handed the Kimi subscription bearer to `api.openai.com`. + /// Custom `KIGI_CODE_BASE_URL` deployments and local dev proxies must still + /// keep it (that is why the predicate is not `is_first_party_url`). + #[test] + fn platform_less_model_takes_the_session_credential_only_on_its_own_endpoint() { + for url in [ + kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url, + "http://127.0.0.1:8080/v1", + "http://localhost:3000/v1", + "http://[::1]:9000/v1", + ] { + assert!( + platform_takes_session_credential(None, url), + "{url} is the session's own endpoint (or a local proxy) — unchanged" + ); + } + for url in [ + "https://api.openai.com/v1", + "https://api.deepseek.com/v1", + "https://api.anthropic.com/v1", + "https://api.moonshot.cn/v1", + "", + ] { + assert!( + !platform_takes_session_credential(None, url), + "LEAK: {url} is a third-party host — no session bearer may ride there" + ); + } } /// RAII guard restoring an env var on drop (panic-safe). diff --git a/crates/codegen/kigi-shell/src/agent/config.rs b/crates/codegen/kigi-shell/src/agent/config.rs index c5683d9..1ae9fd9 100644 --- a/crates/codegen/kigi-shell/src/agent/config.rs +++ b/crates/codegen/kigi-shell/src/agent/config.rs @@ -2779,13 +2779,20 @@ pub fn default_model_entries(endpoints: &EndpointsConfig) -> IndexMap( models: &'a IndexMap, model_id: &str, ) -> Option<&'a ModelEntry> { models .get(model_id) - .or_else(|| models.values().find(|m| m.model == model_id)) + .or_else(|| models.values().rev().find(|m| m.model == model_id)) } /// Whether the EFFECTIVE Auto-mode classifier model supports reasoning effort: /// the model actually routed to (`aux_model` when the aux sampler resolved) else @@ -4643,14 +4650,22 @@ reasoning_effort = "low" }); (dir, manager) } - /// LEAK 1a (aux/summary model): the aux `session_key` is resolved by the aux - /// model's OWN platform (as `build_summary_client` / `resolve_aux_sampler_config` - /// now do). A grok (oauth-platform) aux model's resolved sampler `api_key` is - /// therefore grok's own token or `None` — NEVER the primary Kimi key — while a - /// first-party / non-oauth aux model still gets the primary (byte-identical). + /// LEAK 1a (aux/summary model `api_key` channel): the aux `session_key` is + /// resolved by the aux model's OWN platform AND endpoint (as + /// `build_summary_client` / `resolve_aux_sampler_config` now do), so it is + /// never the primary Kimi key on a host that does not own it. + /// + /// - a grok (oauth-platform) aux model → grok's own pooled token or `None`; + /// - a `moonshot-cn` (API-key platform) aux model → NO session key at all + /// (pre-fix it received `kimi-tok`, which `resolve_credentials` then + /// stamped as the `api_key` on an `api.moonshot.cn` request); + /// - a `kimi-code` aux model → the primary, byte-identical. + /// + /// Revert-to-red: dropping the `platform_takes_session_credential` term + /// from `session_key_for_endpoint` makes the moonshot assertion see + /// `Some("kimi-tok")`. #[tokio::test] async fn aux_model_session_key_is_platform_scoped_never_leaking_kimi() { - let home = tempfile::tempdir().unwrap(); let (_kd, kimi) = kimi_primary("kimi-tok"); let endpoints = EndpointsConfig::default(); // A grok aux catalog entry (managed id → oauth platform), no own key. @@ -4658,9 +4673,9 @@ reasoning_effort = "low" grok.info.id = Some("xai-grok/grok-4-latest".to_string()); let mut grok_catalog = IndexMap::new(); grok_catalog.insert("grok".to_string(), grok); - let grok_key = crate::auth::oauth_registry::session_key_for_model( - home.path(), - "xai-grok/grok-4-latest", + let grok_key = crate::auth::oauth_registry::session_key_for_catalog_model( + &grok_catalog, + "grok", Some(&kimi), ); let grok_cfg = resolve_aux_model_sampling_config( @@ -4675,10 +4690,11 @@ reasoning_effort = "low" Some("kimi-tok"), "a grok aux model must never receive the primary Kimi session token", ); - // A non-oauth aux catalog entry still resolves to the primary token. + // An API-key registry platform gets NO session key — the pre-fix + // behaviour handed it the primary Kimi token on api.moonshot.cn. let mut k2 = test_model_entry( "kimi-k2-0905-preview", - "https://vendor/v1", + "https://api.moonshot.cn/v1", None, None, None, @@ -4686,23 +4702,43 @@ reasoning_effort = "low" k2.info.id = Some("moonshot-cn/kimi-k2".to_string()); let mut k2_catalog = IndexMap::new(); k2_catalog.insert("k2".to_string(), k2); - let k2_key = crate::auth::oauth_registry::session_key_for_model( - home.path(), - "moonshot-cn/kimi-k2", + assert_eq!( + crate::auth::oauth_registry::session_key_for_catalog_model( + &k2_catalog, + "k2", + Some(&kimi), + ), + None, + "LEAK: an API-key-platform aux model must receive no session token", + ); + // The first-party subscription channel is byte-identical. + let mut kimi_code = test_model_entry( + "kimi-for-coding", + kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url, + None, + None, + None, + ); + kimi_code.info.id = Some("kimi-code/kimi-for-coding".to_string()); + let mut kimi_catalog = IndexMap::new(); + kimi_catalog.insert("kimi-for-coding".to_string(), kimi_code); + let kimi_key = crate::auth::oauth_registry::session_key_for_catalog_model( + &kimi_catalog, + "kimi-for-coding", Some(&kimi), ); - let k2_cfg = resolve_aux_model_sampling_config( - "k2", - &k2_catalog, + let kimi_cfg = resolve_aux_model_sampling_config( + "kimi-for-coding", + &kimi_catalog, &endpoints, - k2_key.as_deref(), + kimi_key.as_deref(), None, ) - .expect("non-oauth aux resolves via the primary session token"); + .expect("a kimi-code aux model resolves via the primary session token"); assert_eq!( - k2_cfg.api_key.as_deref(), + kimi_cfg.api_key.as_deref(), Some("kimi-tok"), - "a non-oauth aux model must still receive the primary session token", + "the first-party subscription aux path must be byte-identical", ); } #[test] diff --git a/crates/codegen/kigi-shell/src/agent/models.rs b/crates/codegen/kigi-shell/src/agent/models.rs index 8c06ebb..2943bdd 100644 --- a/crates/codegen/kigi-shell/src/agent/models.rs +++ b/crates/codegen/kigi-shell/src/agent/models.rs @@ -1917,6 +1917,51 @@ pub(crate) fn resolve_catalog_key( .map(|(key, _)| acp::ModelId::new(key.clone())) } +/// The managed catalog key (`{platform}/{model}`) a routing slug belongs to. +/// +/// H5: `SamplingConfig::model` is the BARE routing slug, never the catalog key, +/// and duplicate slugs across platforms are BY DESIGN — the registry guarantees +/// an API-key platform and its subscription-OAuth twin list the SAME ids +/// (`xai`/`xai-grok`, `anthropic`/`claude-pro-max`, `openai`/`openai-codex`), +/// and `PlatformId::ALL` orders every API-key platform FIRST. A slug scan +/// therefore resolves the WRONG platform for a user holding both credentials: +/// the OAuth twin loses its live `bearer_resolver` (no mid-session refresh → an +/// unrecoverable 401 ~1h in), its Messages adaptation and its Copilot/Codex +/// identity headers. +/// +/// `current_key` — [`ModelsManager::current_model_id`], the catalog key the +/// picker actually selected — is therefore authoritative whenever it names this +/// slug. Anything else (aux models, subagent overrides, unlisted slugs) falls +/// back to the picker's OWN lookup, [`resolve_catalog_key`], so the auth layer +/// and the picker can never resolve different entries. +pub(crate) fn managed_key_for_slug( + models: &IndexMap, + current_key: Option<&str>, + slug: &str, +) -> Option { + if let Some(entry) = current_key.and_then(|key| models.get(key)) + && (entry.info.model == slug || current_key == Some(slug)) + { + return entry.info.id.clone(); + } + let key = resolve_catalog_key(models, &acp::ModelId::new(slug.to_string()))?; + models.get(key.0.as_ref())?.info.id.clone() +} + +/// The registry platform a routing slug resolves to, via [`managed_key_for_slug`]. +/// `None` for a bare / `[model.*]` / unlisted model. Single definition shared by +/// the session actor's inference-auth chokepoints and the aux/summary paths, so +/// the gate, the manager and the wire adaptations can never disagree. +pub(crate) fn platform_for_slug( + models: &IndexMap, + current_key: Option<&str>, + slug: &str, +) -> Option { + let key = managed_key_for_slug(models, current_key, slug); + kigi_models::parse_managed_model_key(key.as_deref().unwrap_or(slug)) + .map(|(platform, _)| platform) +} + /// Catalog key for a persisted session model id, restricted to **selectable** /// entries. A selectable exact-key match wins (as in [`resolve_catalog_key`]); /// otherwise the last selectable entry whose routing slug matches `id`, so a @@ -3668,6 +3713,52 @@ mod tests { assert!(!info.visible_for_auth(false)); } + /// SHIP-BLOCKER regression: a user who signed in with ONLY a Claude Pro/Max + /// subscription has no PRIMARY (Kimi) session, so `is_session_auth()` is + /// false. Stamping `supported_in_api = !uses_oauth()` therefore hid every + /// one of their models — `available()` returned an empty picker and the + /// whole subscription-OAuth feature was dead for its target user. The + /// claude-pro-max entry must be visible with no primary session at all. + #[test] + fn claude_pro_max_only_user_sees_their_models_in_the_picker() { + let wire: kigi_models::WireModel = + serde_json::from_value(serde_json::json!({ "id": "claude-opus-4-8" })) + .expect("wire model fixture"); + let entry_config = crate::agent::models_fetch::platform_wire_model_to_entry( + kigi_models::PlatformId::ClaudeProMax, + wire, + "https://api.anthropic.com/v1", + ); + let entry = ModelEntry::from_config_entry(&entry_config); + let key = entry + .info + .id + .clone() + .expect("platform entries carry a managed catalog key"); + assert_eq!(key, "claude-pro-max/claude-opus-4-8"); + let mut catalog = IndexMap::new(); + catalog.insert(key.clone(), entry); + + // Empty home ⇒ the primary AuthManager holds no credential at all. + let home = tempfile::tempdir().expect("tempdir"); + let mgr = ModelsManager::new( + None, + catalog, + acp::ModelId::new(Arc::from(key.clone())), + Arc::new(AuthManager::new(home.path(), KimiCodeConfig::default())), + config::Config::default(), + ); + assert!( + !mgr.is_session_auth(), + "a claude-pro-max-only user has no PRIMARY (Kimi) OAuth session" + ); + assert!( + mgr.available() + .contains_key(&acp::ModelId::new(Arc::from(key.clone()))), + "the claude-pro-max model must reach the picker without a primary session" + ); + } + // ── duplicate model slug re-keying (A/B experiment "auto" alias) ── fn make_entry_config(model: &str, name: Option<&str>) -> config::ModelEntryConfig { diff --git a/crates/codegen/kigi-shell/src/agent/models_fetch.rs b/crates/codegen/kigi-shell/src/agent/models_fetch.rs index 16861ee..8b53585 100644 --- a/crates/codegen/kigi-shell/src/agent/models_fetch.rs +++ b/crates/codegen/kigi-shell/src/agent/models_fetch.rs @@ -645,9 +645,15 @@ pub(crate) fn platform_wire_model_to_entry( inference_idle_timeout_secs: None, max_retries: None, hidden: false, - // Subscription models require the OAuth session; open-platform - // models are usable by API-key users. - supported_in_api: !platform.uses_oauth(), + // `supported_in_api: false` hides a model unless the PRIMARY session is + // an OAuth session (`ModelInfo::visible_for_auth`). Only `kimi-code` + // rides that primary session, so only it may be gated on it. Every + // other OAuth platform (claude-pro-max, openai-codex, github-copilot, + // xai-grok) carries its OWN pooled credential, and its models only + // enter the catalog once THAT provider is signed in — gating them on + // the Kimi session would hide every model from a user who signed in + // with only a Claude/ChatGPT/Copilot/Grok subscription. + supported_in_api: platform != kigi_models::PlatformId::KimiCode, supports_backend_search: false, compactions_remaining: None, compaction_at_tokens: None, @@ -898,6 +904,16 @@ fn get_string_map( mod tests { use super::*; + /// Whether a freshly-fetched platform entry shows up in the picker for a + /// user whose PRIMARY session is NOT an OAuth session (`is_session_auth == + /// false`): an API-key user, or — the case that made this a ship-blocker — + /// someone who signed in with ONLY a Claude Pro/Max, ChatGPT, Copilot, or + /// Grok subscription. `ModelInfo::visible_for_auth` is the picker's real + /// predicate (`agent/models.rs` → `available()`). + fn visible_to_non_primary_session_user(entry: &crate::agent::config::ModelEntryConfig) -> bool { + crate::agent::config::ModelEntry::from_config_entry(entry).visible_for_auth(false) + } + /// OpenAI-cycle e2e (mock wire): a polluted bare-id `/models` listing + /// a models.dev refresh produce a catalog with ONLY chat models, enriched /// context windows / efforts, and the Responses backend — the full @@ -1277,8 +1293,14 @@ mod tests { "enrichment fills the context window from models.dev anthropic" ); assert!( - !opus.supported_in_api, - "subscription (uses_oauth) models require the OAuth session" + opus.supported_in_api, + "claude-pro-max carries its OWN pooled credential — it must NOT be \ + gated on the primary (Kimi) session" + ); + assert!( + visible_to_non_primary_session_user(opus), + "a Claude-Pro/Max-only user has no primary OAuth session; their \ + models must still appear in the picker" ); } @@ -1409,8 +1431,14 @@ mod tests { "context window comes from models.dev github-copilot enrichment" ); assert!( - !entry.supported_in_api, - "subscription (uses_oauth) models require the OAuth session" + entry.supported_in_api, + "github-copilot carries its OWN pooled credential — it must NOT be \ + gated on the primary (Kimi) session" + ); + assert!( + visible_to_non_primary_session_user(entry), + "a Copilot-only user has no primary OAuth session; their models \ + must still appear in the picker" ); } @@ -1484,8 +1512,14 @@ mod tests { assert_eq!(sol.context_window.get(), 272_000); assert_eq!(sol.name.as_deref(), Some("GPT-5.6-Sol")); assert!( - !sol.supported_in_api, - "subscription (uses_oauth) models require the OAuth session" + sol.supported_in_api, + "openai-codex carries its OWN pooled credential — it must NOT be \ + gated on the primary (Kimi) session" + ); + assert!( + visible_to_non_primary_session_user(sol), + "a ChatGPT/Codex-only user has no primary OAuth session; their \ + models must still appear in the picker" ); assert!(sol.supports_reasoning_effort); assert_eq!( @@ -3563,8 +3597,14 @@ mod tests { "an OAuth channel carries no api-key env" ); assert!( - !entry.supported_in_api, - "subscription models require the OAuth session (not the public API)" + entry.supported_in_api, + "xai-grok carries its OWN pooled credential — it must NOT be gated \ + on the primary (Kimi) session" + ); + assert!( + visible_to_non_primary_session_user(entry), + "a Grok-only user has no primary OAuth session; their models must \ + still appear in the picker" ); // Passthrough dialect (identical to the API-key xai wire). let model_entry = crate::agent::config::ModelEntry::from_config_entry(entry); diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs index e683f0a..2b4f069 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs @@ -25,17 +25,19 @@ impl MvpAgent { primary: &SamplingConfig, ) -> Result<(OaiCompatClient, String), acp::Error> { let slug = self.resolve_session_summary_model(); - // Resolve the aux token by the summary model's OWN platform: a grok - // (oauth-platform) summary model draws its pooled grok token or `None` - // — NEVER the primary Kimi session token (which `resolve_credentials` - // would otherwise stamp onto an api.x.ai request). A first-party / - // non-oauth summary model still gets the primary (byte-identical). - let session_key = crate::auth::oauth_registry::session_key_for_model( - &crate::util::kigi_home::kigi_home(), + let models = self.models_manager.models(); + // Resolve the aux token by the summary model's OWN platform AND + // endpoint: a grok (oauth-platform) summary model draws its pooled grok + // token or `None`, and an API-key registry platform draws NOTHING — + // NEVER the primary Kimi session token (which `resolve_credentials` + // would otherwise stamp onto an api.x.ai / api.deepseek.com request). + // The first-party subscription channel still gets the primary + // (byte-identical). + let session_key = crate::auth::oauth_registry::session_key_for_catalog_model( + &models, &slug, Some(&self.auth_manager), ); - let models = self.models_manager.models(); let endpoints = self.models_manager.endpoints(); let alpha_test_key = self.cfg.borrow().endpoints.alpha_test_key.clone(); let config = match crate::agent::config::resolve_aux_model_sampling_config( @@ -47,7 +49,17 @@ impl MvpAgent { ) { Some(mut cfg) => { cfg.attribution_callback = primary.attribution_callback.clone(); - cfg.bearer_resolver = primary.bearer_resolver.clone(); + // H4: the SESSION model's bearer_resolver must not ride to a + // summary model on a different provider — + // `SamplingClient::post` REPLACES the request's auth header + // from it, overwriting the summary model's own resolved key on + // ITS host. Route through the one aux-resolver decision. + cfg.bearer_resolver = + crate::session::acp_session::sampler_turn::aux_bearer_resolver( + primary.bearer_resolver.clone(), + crate::agent::models::platform_for_slug(&models, None, &slug), + &cfg.base_url, + ); cfg.max_retries = primary.max_retries; cfg } @@ -669,31 +681,51 @@ impl MvpAgent { ); Ok(entry.clone()) } - /// Resolve the SESSION token for `model` by the model's OWN platform — the - /// single guard against the api_key-channel token leak. + /// Resolve the SESSION token for `model` by the model's OWN platform AND + /// endpoint — the single guard against the api_key-channel token leak. /// - /// An oauth-platform model (xai-grok) draws its session token from ITS OWN - /// process-global pool manager (build-on-demand from the on-disk grok token, - /// proactively refreshed), INDEPENDENT of the primary `auth_method`; when - /// that provider has no stored session the token is `None` — NEVER the - /// primary Kimi key. Every other model (first-party / Kimi) uses the primary - /// session manager, and only under a session-based auth method — - /// byte-identical to the pre-fix path. SECURITY: the resolved token is never - /// logged. + /// - an oauth-platform model (xai-grok, claude-pro-max, github-copilot, + /// openai-codex) draws its session token from ITS OWN process-global pool + /// manager (built on demand from the on-disk token, proactively + /// refreshed), INDEPENDENT of the primary `auth_method`; when that + /// provider has no stored session the token is `None` — NEVER the primary + /// Kimi key; + /// - an endpoint that does not take a session credential at all + /// ([`crate::agent::auth_method::platform_takes_session_credential`] — + /// every API-key registry platform, and any `[model.*]` block pointed at a + /// third-party host) gets `None`. This is C1: `resolve_credentials` takes + /// the `else if let Some(key) = session_key` arm and sets + /// `api_key = ` with the THIRD-PARTY `base_url`, which + /// `SamplingClient` then builds into `Authorization: Bearer …`. It is + /// reachable with ZERO configuration: `default_models.json` bundles + /// `moonshot-cn/*` + `moonshot-ai/*` entries that a Kimi-subscription user + /// sees on first launch / offline; + /// - the first-party subscription channel (kimi-code, a `KIGI_CODE_BASE_URL` + /// deployment, a loopback proxy) uses the primary session manager, and + /// only under a session-based auth method — byte-identical to the pre-fix + /// path. + /// + /// SECURITY: the resolved token is never logged. fn session_token_for_model(&self, model: &ModelEntry) -> Option { - if let Some(oauth) = model - .info() + let info = model.info(); + let platform = info .id .as_deref() .and_then(kigi_models::parse_managed_model_key) - .and_then(|(platform, _)| platform.oauth()) - { + .map(|(platform, _)| platform); + if let Some(oauth) = platform.and_then(kigi_models::PlatformId::oauth) { return crate::auth::oauth_registry::global_manager_for( - &crate::util::kigi_home::kigi_home(), + &crate::auth::oauth_registry::pool_home(), oauth, ) .current_or_expired(); } + if !crate::agent::auth_method::platform_takes_session_credential( + platform, + &info.base_url, + ) { + return None; + } if self.is_session_based_auth() { self.auth_manager.current_or_expired() } else { diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs index 42fa85e..e8f59ae 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs @@ -1197,18 +1197,26 @@ async fn prepare_sampling_config_never_stamps_kimi_key_on_grok_model() { let endpoints = EndpointsConfig::default(); - // First-party Kimi model (non-oauth platform): the primary session key IS - // its api_key — the byte-identical primary path, and proof the Kimi token is + // First-party SUBSCRIPTION model (kimi-code): the primary session key IS its + // api_key — the byte-identical primary path, and proof the Kimi token is // live (so it WOULD leak if mis-routed onto a grok request). This assertion // also confirms the session-based primary path is active. - let mut kimi_model = ModelEntry::fallback("kimi-k2-0905-preview", &endpoints); - kimi_model.info.id = Some("moonshot-cn/kimi-k2-0905-preview".to_string()); + // + // This used to use `moonshot-cn/kimi-k2-0905-preview` and assert the SAME + // thing, which encoded the C1 defect: moonshot-cn is an API-key registry + // platform on `api.moonshot.cn`, NOT first-party, so "must carry the primary + // session key" was asserting the leak. `api_key_channel_leak_tests` now pins + // the opposite for every moonshot entry. + let mut kimi_model = ModelEntry::fallback("kimi-for-coding", &endpoints); + kimi_model.info.id = Some("kimi-code/kimi-for-coding".to_string()); + kimi_model.info.base_url = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url.to_string(); assert!(!kimi_model.has_own_credentials()); let kimi_cfg = agent.prepare_sampling_config_for_model(&kimi_model, None); assert_eq!( kimi_cfg.api_key.as_deref(), Some(KIMI_KEY), - "a first-party Kimi model must carry the primary session key (primary path unchanged)" + "the first-party subscription model must carry the primary session key \ + (primary path unchanged)" ); // xai-grok model (oauth platform): the session token resolves from its OWN @@ -1280,6 +1288,9 @@ async fn ensure_plugin_registry_lazily_populates_snapshot() { ); } mod subagent_spawn_context_tests; +/// LEAK guard for the `api_key` channel (C1/C2), through the real +/// `prepare_sampling_config_for_model` resolution path. +mod api_key_channel_leak_tests; /// No load in flight and no session → the wait returns immediately /// (the caller then surfaces "unknown session id" exactly as before). #[tokio::test] diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests/api_key_channel_leak_tests.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests/api_key_channel_leak_tests.rs new file mode 100644 index 0000000..485b9b1 --- /dev/null +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests/api_key_channel_leak_tests.rs @@ -0,0 +1,313 @@ +//! LEAK GUARD (`api_key` channel) — C1/C2, driven through the REAL resolution +//! path, `MvpAgent::prepare_sampling_config_for_model`. +//! +//! This is the channel the `bearer_resolver` guard does NOT close, and the one +//! the first round of leak tests assumed away by hand-stamping a provider key +//! into chat state. The chain: `session_token_for_model` used to fall through to +//! `self.auth_manager.current_or_expired()` (the primary Kimi bearer) for every +//! non-OAuth model, `resolve_credentials` then took its +//! `else if let Some(key) = session_key` arm and set `api_key = ` +//! with the THIRD-PARTY `base_url`, and `SamplingClient` builds +//! `Authorization: Bearer ` straight into `default_headers` — which +//! `post()` only overrides when a resolver exists, so `bearer_resolver: None` +//! does not save it. +//! +//! Nothing here stamps a credential by hand: every assertion reads what the +//! resolution path actually produced. + +use super::super::*; +use crate::agent::auth_method::{ + CACHED_TOKEN_AUTH_METHOD_ID, HOUSE_API_KEY_ENV_VAR, LEGACY_XAI_API_KEY_ENV_VAR, + XAI_API_KEY_ENV_VAR, +}; +use crate::agent::config::{Config as AgentConfig, EndpointsConfig, EnvKeys, ModelEntry}; +use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig}; +use kigi_test_support::EnvGuard; + +const KIMI_TOKEN: &str = "kimi-subscription-token-DO-NOT-LEAK"; + +/// Ambient BYOK env vars unset, so a model with no resolvable credential ends up +/// with `api_key == None` rather than a global-key fallback that could mask the +/// leak under test. Every test holding these must be `#[serial]`. +fn without_ambient_byok_env() -> [EnvGuard; 3] { + [ + EnvGuard::unset(HOUSE_API_KEY_ENV_VAR), + EnvGuard::unset(XAI_API_KEY_ENV_VAR), + EnvGuard::unset(LEGACY_XAI_API_KEY_ENV_VAR), + ] +} + +/// An `MvpAgent` on a session-based (`cached_token`) ACP method holding a live +/// Kimi subscription bearer — the mainstream configuration in which the leak +/// fires. `(tempdir, agent)`; the tempdir is the auth store and is returned so +/// the caller keeps it alive. +fn kimi_session_agent() -> (tempfile::TempDir, MvpAgent) { + let dir = tempfile::tempdir().expect("tempdir"); + let auth_manager = std::sync::Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + auth_manager.hot_swap(KimiAuth { + key: KIMI_TOKEN.to_string(), + auth_mode: AuthMode::OAuth, + refresh_token: Some("rt".into()), + expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), + ..KimiAuth::test_default() + }); + let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); + let agent = MvpAgent::new( + GatewaySender::new(tx), + &AgentConfig::default(), + auth_manager, + None, + ) + .expect("valid test config"); + agent.set_auth_method(acp::AuthMethodId::new(CACHED_TOKEN_AUTH_METHOD_ID)); + (dir, agent) +} + +/// A catalog entry as `resolve_model_list` builds one for a fetched registry +/// model: managed catalog key, platform base URL, no credential of its own. +fn platform_entry(catalog_key: &str, slug: &str, base_url: &str) -> ModelEntry { + let mut entry = ModelEntry::fallback(slug, &EndpointsConfig::default()); + entry.info.id = Some(catalog_key.to_string()); + entry.info.base_url = base_url.to_string(); + entry +} + +/// C1, the ZERO-CONFIGURATION repro. `default_models.json` bundles +/// `moonshot-cn/*` and `moonshot-ai/*` entries with `api_key: None`, and +/// `resolve_model_list` keeps the bundled defaults whenever no catalog fetch has +/// succeeded — so on first launch / offline a Kimi-subscription user sees them +/// in the picker with no configuration whatsoever. Selecting one used to send +/// `Authorization: Bearer ` to `api.moonshot.cn`, which is NOT +/// first-party. +/// +/// Revert-to-red: dropping the `platform_takes_session_credential` guard from +/// `session_token_for_model` makes every `api_key` below `Some(KIMI_TOKEN)`. +#[tokio::test] +#[serial_test::serial] +async fn bundled_default_moonshot_models_never_carry_the_kimi_bearer() { + let _env = without_ambient_byok_env(); + let (_dir, agent) = kimi_session_agent(); + + let bundled = crate::agent::config::default_model_entries(&EndpointsConfig::default()); + let moonshot: Vec<_> = bundled + .iter() + .filter(|(key, _)| key.starts_with("moonshot-cn/") || key.starts_with("moonshot-ai/")) + .collect(); + assert_eq!( + moonshot.len(), + 4, + "default_models.json still bundles the four moonshot open-platform entries" + ); + + for (key, entry) in moonshot { + assert!( + !entry.has_own_credentials(), + "{key}: the bundled entry carries no credential of its own" + ); + assert!( + !crate::util::is_effective_coding_endpoint_url(&entry.info().base_url), + "{key}: routes to a third-party host ({})", + entry.info().base_url + ); + let cfg = agent.prepare_sampling_config_for_model(entry, None); + assert_ne!( + cfg.api_key.as_deref(), + Some(KIMI_TOKEN), + "LEAK: selecting the bundled {key} sent the Kimi subscription bearer to {}", + entry.info().base_url + ); + assert_eq!( + cfg.base_url, + entry.info().base_url, + "{key}: still routes to its own host (the fix must not reroute traffic)" + ); + } +} + +/// C1 across the API-key registry platform shapes a fetched catalog produces. +#[tokio::test] +#[serial_test::serial] +async fn api_key_platform_models_never_carry_the_kimi_bearer_as_api_key() { + let _env = without_ambient_byok_env(); + let (_dir, agent) = kimi_session_agent(); + + for (catalog_key, slug, base_url) in [ + ("deepseek/deepseek-chat", "deepseek-chat", "https://api.deepseek.com/v1"), + ("openai/gpt-5.2", "gpt-5.2", "https://api.openai.com/v1"), + ("anthropic/claude-opus-4-8", "claude-opus-4-8", "https://api.anthropic.com/v1"), + ("groq/llama-4", "llama-4", "https://api.groq.com/openai/v1"), + ("xai/grok-4.5", "grok-4.5", "https://api.x.ai/v1"), + ] { + let entry = platform_entry(catalog_key, slug, base_url); + let cfg = agent.prepare_sampling_config_for_model(&entry, None); + assert_ne!( + cfg.api_key.as_deref(), + Some(KIMI_TOKEN), + "LEAK: {catalog_key} carried the primary Kimi session bearer to {base_url}" + ); + } +} + +/// C2, the `[model.*]` repro. A `[model.gpt-4o]` block has `info.id == None`, so +/// it has no platform at all — which used to be a blanket allow. BYOK is +/// `has_own_credentials()`, which probes `std::env::var` AT CALL TIME, so an +/// unset (or mistyped) `env_key` classifies the model NotByok and the Kimi +/// bearer went to `api.openai.com` on BOTH channels. +/// +/// Revert-to-red: making the `None` arm of `platform_takes_session_credential` +/// return `true` again makes `api_key` here `Some(KIMI_TOKEN)`. +#[tokio::test] +#[serial_test::serial] +async fn config_model_with_an_unset_env_key_never_carries_the_kimi_bearer() { + let _env = without_ambient_byok_env(); + let _typo = EnvGuard::unset("OPENAI_API_KEY_TYPO"); + let (_dir, agent) = kimi_session_agent(); + + let mut entry = ModelEntry::fallback("gpt-4o", &EndpointsConfig::default()); + entry.info.id = None; // a `[model.gpt-4o]` config block + entry.info.base_url = "https://api.openai.com/v1".to_string(); + entry.env_key = Some(EnvKeys::single("OPENAI_API_KEY_TYPO")); + assert!( + !entry.has_own_credentials(), + "the env var is unset, so this classifies NotByok — the precondition of the defect" + ); + + let cfg = agent.prepare_sampling_config_for_model(&entry, None); + assert_ne!( + cfg.api_key.as_deref(), + Some(KIMI_TOKEN), + "LEAK: a [model.*] block with an unset env_key sent the Kimi bearer to api.openai.com" + ); + assert_eq!(cfg.api_key, None, "no credential resolves — fail fast"); +} + +/// The first-party subscription channel must stay BYTE-IDENTICAL: `kimi-code/*` +/// (and a `[model.*]` block on the session's own coding endpoint, including a +/// `KIGI_CODE_BASE_URL` deployment / a local dev proxy) still carries the +/// primary session key. This assertion is also what proves the Kimi token is +/// live in the tests above — it WOULD leak if the guard were missing. +#[tokio::test] +#[serial_test::serial] +async fn the_first_party_subscription_channel_still_carries_the_session_key() { + let _env = without_ambient_byok_env(); + let (_dir, agent) = kimi_session_agent(); + + let kimi = platform_entry( + "kimi-code/kimi-for-coding", + "kimi-for-coding", + kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url, + ); + assert_eq!( + agent + .prepare_sampling_config_for_model(&kimi, None) + .api_key + .as_deref(), + Some(KIMI_TOKEN), + "the kimi-code subscription channel must be unchanged" + ); + + for base_url in [ + kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url, + "http://127.0.0.1:4141/v1", + "http://localhost:8080/v1", + ] { + let mut bare = ModelEntry::fallback("kigi-4.5", &EndpointsConfig::default()); + bare.info.id = None; + bare.info.base_url = base_url.to_string(); + assert_eq!( + agent + .prepare_sampling_config_for_model(&bare, None) + .api_key + .as_deref(), + Some(KIMI_TOKEN), + "{base_url}: a custom deployment / local proxy keeps the session key" + ); + } +} + +/// A subscription-OAuth model draws its `api_key` from ITS OWN pooled manager, +/// never the Kimi primary — and never falls back to it when that provider has no +/// stored session (the pool home is an empty TempDir under `cfg(test)`). +#[tokio::test] +#[serial_test::serial] +async fn oauth_platform_models_never_carry_the_kimi_bearer_as_api_key() { + let _env = without_ambient_byok_env(); + let (_dir, agent) = kimi_session_agent(); + + for (catalog_key, slug, base_url) in [ + ("xai-grok/grok-4-latest", "grok-4-latest", "https://api.x.ai/v1"), + ( + "claude-pro-max/claude-opus-4-8", + "claude-opus-4-8", + "https://api.anthropic.com/v1", + ), + ("github-copilot/gpt-4.1", "gpt-4.1", "https://api.githubcopilot.com"), + ( + "openai-codex/gpt-5.5", + "gpt-5.5", + "https://chatgpt.com/backend-api/codex", + ), + ] { + let entry = platform_entry(catalog_key, slug, base_url); + let cfg = agent.prepare_sampling_config_for_model(&entry, None); + assert_ne!( + cfg.api_key.as_deref(), + Some(KIMI_TOKEN), + "LEAK: {catalog_key} carried the primary Kimi session bearer to {base_url}" + ); + } +} + +/// H5 at the api_key channel: `resolve_model_id` (the picker's own lookup) must +/// hand `prepare_sampling_config_for_model` the entry the user SELECTED, even +/// when an API-key platform and its subscription-OAuth twin list the same +/// routing slug in `PlatformId::ALL` order. Selecting the OAuth twin by catalog +/// key must not resolve the API-key twin — and vice versa. +#[tokio::test] +#[serial_test::serial] +async fn dual_credential_slug_collision_resolves_the_selected_catalog_key() { + let _env = without_ambient_byok_env(); + let (_dir, agent) = kimi_session_agent(); + + // API-key platform FIRST, exactly as `PlatformId::ALL` orders them. + for key in ["xai/grok-4.5", "xai-grok/grok-4.5"] { + agent.models_manager.insert_test_entry( + key, + platform_entry(key, "grok-4.5", "https://api.x.ai/v1"), + ); + } + + for key in ["xai/grok-4.5", "xai-grok/grok-4.5"] { + let resolved = agent + .resolve_model_id(&acp::ModelId::new(key)) + .expect("both twins resolve"); + assert_eq!( + resolved.info().id.as_deref(), + Some(key), + "selecting {key} must resolve THAT catalog entry, not its slug twin" + ); + } + + // And the bare slug resolves the same entry the picker's `resolve_catalog_key` + // does — one direction, one answer (the auth layer used to first-match). + let by_slug = agent + .resolve_model_id(&acp::ModelId::new("grok-4.5")) + .expect("the bare slug resolves"); + let models = agent.models_manager.models(); + let picker_key = crate::agent::models::resolve_catalog_key( + &models, + &acp::ModelId::new("grok-4.5"), + ) + .expect("the picker resolves the bare slug"); + assert_eq!( + by_slug.info().id.as_deref(), + Some(picker_key.0.as_ref()), + "the auth layer and the picker must resolve the SAME entry for one slug" + ); + assert_eq!( + crate::agent::config::find_model_by_id(&models, "grok-4.5") + .and_then(|e| e.info().id.as_deref()), + Some(picker_key.0.as_ref()), + "find_model_by_id must agree with resolve_catalog_key by construction" + ); +} diff --git a/crates/codegen/kigi-shell/src/agent/subagent/mod.rs b/crates/codegen/kigi-shell/src/agent/subagent/mod.rs index 32929ba..0382711 100644 --- a/crates/codegen/kigi-shell/src/agent/subagent/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/subagent/mod.rs @@ -1003,16 +1003,22 @@ fn resolve_model_override_to_config( } else { acp::ModelId::new(entry.info().model.clone()) }; - // Resolve the child's session token by the OVERRIDE model's OWN platform, - // not the parent's primary auth: a grok (oauth-platform) override draws its - // pooled grok token or `None` — NEVER the primary Kimi session token (which - // `resolve_credentials` would otherwise stamp onto the child's api.x.ai - // credentials, leaking it in the logout-mid-session edge). A first-party / - // non-oauth override still resolves to the primary (byte-identical). - let managed_key = entry.info().id.as_deref().unwrap_or(model_id); - let session_key = crate::auth::oauth_registry::session_key_for_model( - &crate::util::kigi_home::kigi_home(), - managed_key, + // Resolve the child's session token by the OVERRIDE model's OWN platform + // AND endpoint, not the parent's primary auth: a grok (oauth-platform) + // override draws its pooled grok token or `None`, and an API-key registry + // platform / a third-party `[model.*]` host draws NOTHING — NEVER the + // primary Kimi session token, which `resolve_credentials` would otherwise + // stamp onto the child's api.x.ai / api.moonshot.cn credentials. The + // first-party subscription channel still resolves to the primary + // (byte-identical). + let session_key = crate::auth::oauth_registry::session_key_for_endpoint( + entry + .info() + .id + .as_deref() + .and_then(kigi_models::parse_managed_model_key) + .map(|(platform, _)| platform), + &entry.info().base_url, Some(&ctx.auth_manager), ); let has_session_key = session_key.is_some(); diff --git a/crates/codegen/kigi-shell/src/agent/subagent/tests/mod.rs b/crates/codegen/kigi-shell/src/agent/subagent/tests/mod.rs index e569945..dc77ca6 100644 --- a/crates/codegen/kigi-shell/src/agent/subagent/tests/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/subagent/tests/mod.rs @@ -3131,19 +3131,25 @@ fn fresh_tool_model_rejects_unavailable_exact_key_over_visible_slug_collision() "validation must inspect the unavailable exact-key entry selected by execution" ); } +/// Validation must inspect the SAME slug-collision entry execution selects. +/// Both go through `find_model_by_id`, whose slug scan takes the LAST match — +/// aligned with the picker's `resolve_catalog_key` so the auth layer and the +/// picker can never resolve different platforms for one slug (the H5 collision). +/// So a blocked LAST entry must be rejected even though an available earlier one +/// shares the slug. #[test] -fn fresh_tool_model_rejects_unavailable_first_slug_collision() { +fn fresh_tool_model_rejects_unavailable_last_slug_collision() { let mut models = indexmap::IndexMap::new(); - let mut unavailable_first = test_model_entry("shared-routing-slug"); - unavailable_first.info.user_selectable = false; - models.insert("blocked-first".to_string(), unavailable_first); - models.insert("visible-second".to_string(), test_model_entry("shared-routing-slug")); + models.insert("visible-first".to_string(), test_model_entry("shared-routing-slug")); + let mut unavailable_last = test_model_entry("shared-routing-slug"); + unavailable_last.info.user_selectable = false; + models.insert("blocked-last".to_string(), unavailable_last); assert_eq!( super::handle_request::task_model_override_error(Some("shared-routing-slug"), ModelOverrideProvenance::Tool, false, & models, false,).as_deref(), Some("Unknown Task.model slug 'shared-routing-slug'. Valid model slugs: \ - visible-second. Omit `model` to inherit the parent model."), - "validation must inspect the first routing-slug entry selected by execution" + visible-first. Omit `model` to inherit the parent model."), + "validation must inspect the last routing-slug entry selected by execution" ); } #[test] diff --git a/crates/codegen/kigi-shell/src/agent/subagent/tests/rest.rs b/crates/codegen/kigi-shell/src/agent/subagent/tests/rest.rs index 1ff4a61..37f4969 100644 --- a/crates/codegen/kigi-shell/src/agent/subagent/tests/rest.rs +++ b/crates/codegen/kigi-shell/src/agent/subagent/tests/rest.rs @@ -2511,14 +2511,48 @@ async fn subagent_override_grok_model_never_leaks_kimi_session_token() { "a grok override must never receive the primary Kimi session token", ); } -/// Byte-identical guard: a non-oauth override with a Kimi primary still resolves -/// to the primary session token — passes both before and after the fix (the -/// non-oauth path is unchanged). +/// Byte-identical guard: an override on the SESSION's own first-party endpoint +/// (the kimi-code subscription channel) still resolves to the primary session +/// token — the primary path is unchanged. #[tokio::test] -async fn subagent_override_non_oauth_model_still_gets_primary_token() { +async fn subagent_override_first_party_model_still_gets_primary_token() { + let (_kd, manager) = kimi_primary_with_token("kimi-secret"); + let mut entry = test_model_entry("kimi-for-coding"); + entry.info.id = Some("kimi-code/kimi-for-coding".to_string()); + entry.info.base_url = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url.to_string(); + let mut models = indexmap::IndexMap::new(); + models.insert("kfc".to_string(), entry); + let mut ctx = ctx_with_toggle(HashMap::new()); + ctx.available_models = models; + ctx.auth = Some(crate::auth::KimiAuth { + key: "kimi-secret".to_string(), + auth_mode: crate::auth::AuthMode::OAuth, + ..crate::auth::KimiAuth::test_default() + }); + ctx.auth_manager = manager; + let (config, _model_id) = resolve_model_override_to_config("kfc", &ctx) + .expect("first-party override resolves to a config"); + assert_eq!( + config.api_key.as_deref(), + Some("kimi-secret"), + "a first-party override must still receive the primary session token", + ); +} +/// LEAK guard (C1, subagent-override `api_key` channel): an API-key registry +/// platform override must NOT receive the parent's primary Kimi session token — +/// `resolve_credentials` would stamp it as the child's `api_key` on +/// `api.moonshot.cn`. This test previously asserted the opposite +/// (`subagent_override_non_oauth_model_still_gets_primary_token`), which encoded +/// the defect. +/// +/// Revert-to-red: dropping the `platform_takes_session_credential` term from +/// `oauth_registry::session_key_for_endpoint` makes `api_key` `Some("kimi-secret")`. +#[tokio::test] +async fn subagent_override_api_key_platform_never_gets_the_primary_token() { let (_kd, manager) = kimi_primary_with_token("kimi-secret"); let mut entry = test_model_entry("kimi-k2-0905-preview"); entry.info.id = Some("moonshot-cn/kimi-k2".to_string()); + entry.info.base_url = "https://api.moonshot.cn/v1".to_string(); let mut models = indexmap::IndexMap::new(); models.insert("k2".to_string(), entry); let mut ctx = ctx_with_toggle(HashMap::new()); @@ -2529,12 +2563,12 @@ async fn subagent_override_non_oauth_model_still_gets_primary_token() { ..crate::auth::KimiAuth::test_default() }); ctx.auth_manager = manager; - let (config, _model_id) = - resolve_model_override_to_config("k2", &ctx).expect("non-oauth override resolves to a config"); - assert_eq!( + let (config, _model_id) = resolve_model_override_to_config("k2", &ctx) + .expect("an API-key-platform override still resolves to a config"); + assert_ne!( config.api_key.as_deref(), Some("kimi-secret"), - "a non-oauth override must still receive the primary session token", + "LEAK: an API-key-platform override must never receive the primary Kimi session token", ); } /// An unresolvable `AgentDefinition.model` pin (model absent from diff --git a/crates/codegen/kigi-shell/src/auth/oauth_registry.rs b/crates/codegen/kigi-shell/src/auth/oauth_registry.rs index 0edaa05..b38c128 100644 --- a/crates/codegen/kigi-shell/src/auth/oauth_registry.rs +++ b/crates/codegen/kigi-shell/src/auth/oauth_registry.rs @@ -13,11 +13,10 @@ //! generic-oauth scope, each wired with the SAME lifecycle as the primary Kimi //! manager (`configure_refresher()` + `start_proactive_refresh()`) so the //! on-disk token stays fresh and a 401 recovers via the provider's own manager. -//! Managers are built ON DEMAND: the first grok turn (or model switch) reads the -//! on-disk token via [`global_manager_for`], so a login that lands AFTER a -//! session spawned self-heals — there is no frozen per-session snapshot to go -//! stale. [`manager_for_model`] routes a managed catalog key to the pool (oauth -//! platform) or to the session's primary (everything else). +//! Managers are built ON DEMAND from the on-disk token ([`global_manager_for`]), +//! so a login landing AFTER a session spawned self-heals — no frozen per-session +//! snapshot. [`manager_for_model`] routes a managed catalog key to the pool +//! (oauth platform) or to the session's primary (everything else). //! //! SECURITY: access/refresh tokens and resolved bearers are NEVER logged here. @@ -39,6 +38,31 @@ fn oauth_manager_pool() -> &'static Mutex POOL.get_or_init(|| Mutex::new(HashMap::new())) } +/// The kigi home every OAuth-pool call site resolves from. Single definition so +/// the pool, the aux/summary token routing and the session's inference manager +/// can never read different homes. +/// +/// Production: [`crate::util::kigi_home::kigi_home`]. LIB TESTS: a +/// process-lifetime `TempDir`, unconditionally — the pool is process-global and +/// every manager it builds starts a never-cancelled proactive-refresh loop, so +/// a unit test resolving the real `~/.kigi` would read the developer's stored +/// OAuth tokens and, 60 s later, fire REAL refresh requests against them. +/// Deliberately not a per-test opt-in that can be forgotten: `kigi_home()` is +/// itself a `OnceLock` an earlier test has usually already resolved to the real +/// home, so setting `KIGI_SHARE_DIR` in a test cannot pin it after the fact. +pub(crate) fn pool_home() -> std::path::PathBuf { + #[cfg(test)] + { + static TEST_HOME: OnceLock = OnceLock::new(); + TEST_HOME + .get_or_init(|| tempfile::tempdir().expect("tempdir for the test OAuth pool")) + .path() + .to_path_buf() + } + #[cfg(not(test))] + crate::util::kigi_home::kigi_home() +} + /// Get-or-create the process-global manager for `oauth`, wiring the same /// refresher + proactive-refresh lifecycle as the primary Kimi manager the /// FIRST time a scope is seen. The manager reads the on-disk token at @@ -93,27 +117,70 @@ pub(crate) fn manager_for_model( primary.cloned() } -/// The SESSION token (the raw bearer/key string) that governs INFERENCE auth -/// for `managed_key`, resolved by the model's OWN platform. Thin wrapper over -/// [`manager_for_model`] used by the aux-model and subagent-override wire paths -/// so a `{platform}/{model}` key never receives the primary token of a -/// DIFFERENT provider. +/// The SESSION token (the raw bearer/key string) that may ride an INFERENCE +/// request routed to `platform` at `base_url`. Used by the aux-model, summary +/// and subagent-override wire paths, where the result is stamped straight into +/// [`crate::agent::config::resolve_credentials`] as the request's `api_key`. /// -/// A generic device-code OAuth platform (xai-grok) draws its token from ITS OWN -/// pooled manager; when that provider has no stored session the result is -/// `None` — NEVER the primary Kimi key. Every other key routes to `primary` and -/// yields the primary's current-or-expired token, byte-identical to reading it -/// directly. SECURITY: the resolved token is never logged. -pub(crate) fn session_key_for_model( - kigi_home: &Path, - managed_key: &str, +/// - a generic device-code OAuth platform (xai-grok, claude-pro-max, +/// github-copilot, openai-codex) draws from ITS OWN pooled manager; when that +/// provider has no stored session the result is `None` — never `primary`; +/// - `kimi-code`, and a platform-less model whose endpoint IS the session's own +/// coding endpoint (incl. a `KIGI_CODE_BASE_URL` deployment or a loopback +/// proxy), yield the primary's current-or-expired token — byte-identical to +/// reading it directly; +/// - every API-key registry platform, and every `[model.*]` block pointed at a +/// third-party host, yields `None`. Handing them `primary` put the user's +/// Kimi subscription bearer on `api.deepseek.com` / `api.moonshot.cn` / … +/// as the request's `api_key`. +/// +/// SECURITY: the resolved token is never logged. +pub(crate) fn session_key_for_endpoint( + platform: Option, + base_url: &str, primary: Option<&Arc>, ) -> Option { - manager_for_model(kigi_home, managed_key, primary) + if let Some(oauth) = platform.and_then(kigi_models::PlatformId::oauth) { + return global_manager_for(&pool_home(), oauth) + .current_or_expired() + .map(|a| a.key); + } + if !crate::agent::auth_method::platform_takes_session_credential(platform, base_url) { + return None; + } + primary .and_then(|am| am.current_or_expired()) .map(|a| a.key) } +/// [`session_key_for_endpoint`] for the catalog model whose routing slug (or +/// catalog key) is `slug`. +/// +/// A slug absent from the catalog keeps the pre-registry behaviour: the aux +/// resolver's Tier-2 fallback builds its entry against +/// `EndpointsConfig::resolve_inference_base_url` (first-party), so the primary +/// still governs. +pub(crate) fn session_key_for_catalog_model( + models: &indexmap::IndexMap, + slug: &str, + primary: Option<&Arc>, +) -> Option { + let Some(entry) = crate::agent::config::find_model_by_id(models, slug) else { + return primary + .and_then(|am| am.current_or_expired()) + .map(|a| a.key); + }; + let info = entry.info(); + session_key_for_endpoint( + info.id + .as_deref() + .and_then(kigi_models::parse_managed_model_key) + .map(|(platform, _)| platform), + &info.base_url, + primary, + ) +} + #[cfg(test)] mod tests { use super::*; @@ -159,6 +226,20 @@ mod tests { .expect("openai-codex carries an OAuthConfig") } + /// `session_key_for_endpoint` for a managed catalog key, resolving the + /// platform and its base URL from the registry exactly as the catalog entry + /// would. + fn session_key_for_key( + managed_key: &str, + primary: Option<&Arc>, + ) -> Option { + let platform = kigi_models::parse_managed_model_key(managed_key).map(|(p, _)| p); + let base_url = platform + .map(kigi_models::PlatformId::base_url) + .unwrap_or_default(); + session_key_for_endpoint(platform, &base_url, primary) + } + /// An `openai-codex/` turn resolves to the process-global pooled /// openai-codex manager (its OWN `oauth/openai-codex` scope), NEVER the /// primary Kimi manager — the same leak-safe routing as the other OAuth @@ -187,7 +268,7 @@ mod tests { "openai-codex and claude-pro-max must not share a pooled manager" ); assert_ne!( - session_key_for_model(home.path(), "openai-codex/gpt-5.5", Some(&kimi)), + session_key_for_key("openai-codex/gpt-5.5", Some(&kimi)), Some("kimi-tok".to_string()), "an openai-codex model must never receive the primary Kimi token" ); @@ -218,7 +299,7 @@ mod tests { // Fail-fast: even with a Kimi primary, a copilot turn never yields the // Kimi bearer — it draws from the copilot pool (its own token, or None). assert_ne!( - session_key_for_model(home.path(), "github-copilot/gpt-4.1", Some(&kimi)), + session_key_for_key("github-copilot/gpt-4.1", Some(&kimi)), Some("kimi-tok".to_string()), "a github-copilot model must never receive the primary Kimi token" ); @@ -256,9 +337,8 @@ mod tests { #[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)), + session_key_for_key("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" ); @@ -342,22 +422,62 @@ mod tests { ); } - /// `session_key_for_model`: a non-oauth / bare key yields the primary Kimi - /// token exactly as reading it directly would — byte-identical to the - /// pre-fix aux/override wire path (no runtime / pool touched). + /// `session_key_for_endpoint`: the endpoints that genuinely ride the + /// PRIMARY session — `kimi-code` (the subscription channel) and a + /// platform-less model routed at the session's own coding endpoint (a + /// `KIGI_CODE_BASE_URL` deployment or a loopback dev proxy) — yield the + /// primary token exactly as reading it directly would. No runtime / pool + /// touched. #[test] - fn session_key_for_non_oauth_is_the_primary_token() { + fn session_key_for_the_sessions_own_endpoint_is_the_primary_token() { let (_kd, kimi) = primary_with_token("kimi-tok"); - let home = tempfile::tempdir().unwrap(); - for key in ["moonshot-cn/kimi-k2", "kimi-k2-0905-preview"] { + assert_eq!( + session_key_for_key("kimi-code/kimi-for-coding", Some(&kimi)), + Some("kimi-tok".to_string()), + "kimi-code rides the primary session, unchanged" + ); + for url in [ + kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url, + "http://127.0.0.1:4000/v1", + ] { assert_eq!( - session_key_for_model(home.path(), key, Some(&kimi)), + session_key_for_endpoint(None, url, Some(&kimi)), Some("kimi-tok".to_string()), - "{key} (non-oauth) must yield the primary token unchanged" + "{url}: a platform-less model on the session's own endpoint is unchanged" ); } } + /// LEAK guard (aux / summary / subagent-override `api_key` channel): an + /// API-key registry platform, and a `[model.*]` block pointed at a + /// third-party host, must yield NO session token. Handing them the primary + /// stamped the user's Kimi subscription bearer onto `api.moonshot.cn` / + /// `api.deepseek.com` as the request's `api_key` — the channel the + /// `bearer_resolver` guard alone does not close. + /// + /// Revert-to-red: dropping the `platform_takes_session_credential` term + /// from `session_key_for_endpoint` returns `Some("kimi-tok")` here. + #[test] + fn session_key_for_a_third_party_endpoint_is_never_the_primary_token() { + let (_kd, kimi) = primary_with_token("kimi-tok"); + for key in [ + "moonshot-cn/kimi-k2", + "deepseek/deepseek-chat", + "openai/gpt-5", + ] { + assert_eq!( + session_key_for_key(key, Some(&kimi)), + None, + "LEAK: {key} is an API-key platform — no session token may ride there" + ); + } + assert_eq!( + session_key_for_endpoint(None, "https://api.openai.com/v1", Some(&kimi)), + None, + "LEAK: a [model.*] block on a third-party host gets no session token" + ); + } + /// LEAK guard (aux-model + subagent-override token routing): a grok key with /// a Kimi primary NEVER yields the primary Kimi token — it draws from the /// pooled xai manager (its own token, or `None`). This is the exact source @@ -365,15 +485,14 @@ mod tests { #[tokio::test] async fn session_key_for_grok_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(), "xai-grok/grok-4-latest", Some(&kimi)), + session_key_for_key("xai-grok/grok-4-latest", Some(&kimi)), Some("kimi-tok".to_string()), "a grok aux/override model must never receive the primary Kimi session token" ); // Even with `None` primary the routing is unchanged: grok → pool, never a panic. assert_ne!( - session_key_for_model(home.path(), "xai-grok/grok-4-fast", None), + session_key_for_key("xai-grok/grok-4-fast", None), Some("kimi-tok".to_string()), ); } diff --git a/crates/codegen/kigi-shell/src/session/acp_session.rs b/crates/codegen/kigi-shell/src/session/acp_session.rs index 23e13dc..d6cf089 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session.rs @@ -135,7 +135,7 @@ use prompt_build::*; mod session_mode; use session_mode::*; #[path = "acp_session_impl/sampler_turn.rs"] -mod sampler_turn; +pub(crate) mod sampler_turn; use sampler_turn::*; #[path = "acp_session_impl/tool_dispatch.rs"] mod tool_dispatch; @@ -1256,6 +1256,17 @@ mod rewind_synthetic_turn_tests; #[cfg(test)] #[path = "acp_session_tests/rewrite_zero_turn_prefix_tests.rs"] mod rewrite_zero_turn_prefix_tests; +/// The same guard for the model→platform lookup (the dual-credential slug +/// collision) and for the stamped aux/summary configs. +#[cfg(test)] +#[path = "acp_session_tests/session_bearer_leak_platform_tests.rs"] +mod session_bearer_leak_platform_tests; +/// LEAK guard: the primary Kimi subscription bearer must never ride a request +/// to an API-key registry platform's host, while the subscription-OAuth +/// platforms keep a live resolver from their OWN pooled manager. +#[cfg(test)] +#[path = "acp_session_tests/session_bearer_leak_tests.rs"] +mod session_bearer_leak_tests; /// Pins the `SubagentFinished` usage-fold attribution gate. #[cfg(test)] #[path = "acp_session_tests/subagent_usage_fold_tests.rs"] diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs index b7351b8..bb47ba4 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs @@ -634,14 +634,12 @@ impl SessionActor { &active_session_config, Some(self.max_retries), ); - // A grok (oauth-platform) image-describe model must not inherit the - // session (Kimi) bearer_resolver stamped by `finalize_*`; re-point it at - // grok's own manager. No-op for a first-party / non-oauth model. + // An image-describe model on another provider must not inherit the + // session (Kimi) bearer_resolver stamped by `finalize_*`: re-point it at + // an OAuth model's own manager, or clear it for an API-key-platform / + // third-party endpoint. No-op for the first-party subscription channel. if aux_resolved { - self.repoint_aux_bearer_resolver_for_oauth( - &mut sampler_config, - &self.image_description_model, - ); + self.repoint_aux_bearer_resolver(&mut sampler_config, &self.image_description_model); } let client = kigi_sampler::SamplingClient::new(sampler_config).map_err(|e| { acp::Error::internal_error().data(format!( 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 1d67e33..c127b7a 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 @@ -36,20 +36,37 @@ struct SessionTokenAuthGate { /// BYOK status still refresh against the first-party cli-chat-proxy hosts without /// risking a session-token leak to a third-party BYOK endpoint. endpoint_is_first_party: bool, + /// Whether this model's platform is one whose endpoint accepts a session + /// bearer at all (see + /// [`crate::agent::auth_method::platform_takes_session_credential`]). False + /// for every API-key registry platform, which keeps the primary Kimi bearer + /// off `api.deepseek.com` / `api.openai.com` / … . + endpoint_takes_session_credential: bool, } impl SessionTokenAuthGate { /// Single place `is_session_based` / `endpoint_is_first_party` are derived, - /// so all call sites assemble the gate identically. + /// so all call sites assemble the gate identically. `model_platform` is the + /// registry platform the model routes to (`None` for a bare / `[model.*]` + /// entry) — it MUST be derived from the same lookup + /// ([`SessionActor::managed_key_for_model`]) that + /// [`SessionActor::auth_manager_for_model`] uses, so the gate's verdict and + /// the manager actually wrapped as the bearer resolver can never disagree. fn new( auth_method_id: Option<&acp::AuthMethodId>, model_byok: crate::agent::auth_method::ModelByok, base_url: &str, + model_platform: Option, ) -> Self { Self { is_session_based: auth_method_id .is_some_and(crate::agent::auth_method::is_session_based_method), model_byok, endpoint_is_first_party: crate::util::is_first_party_url(base_url), + endpoint_takes_session_credential: + crate::agent::auth_method::platform_takes_session_credential( + model_platform, + base_url, + ), } } fn active(self) -> bool { @@ -57,6 +74,7 @@ impl SessionTokenAuthGate { self.is_session_based, self.model_byok, self.endpoint_is_first_party, + self.endpoint_takes_session_credential, ) } } @@ -127,6 +145,43 @@ fn auth_manager_bearer_resolver( ) -> kigi_sampler::SharedBearerResolver { std::sync::Arc::new(AuthManagerBearerResolver(am)) } +/// The `bearer_resolver` an AUX / summary `SamplerConfig` may carry, given the +/// SESSION model's resolver (`stamped`) and the AUX model's own platform + +/// endpoint. ONE decision shared by image-describe, the auto-mode classifier +/// (via [`SessionActor::repoint_aux_bearer_resolver`]) and +/// `MvpAgent::build_summary_client`. +/// +/// [`crate::agent::config::stamp_session_local_sampler_fields`] copies the +/// session resolver onto every aux config, and `SamplingClient::post` REPLACES +/// the request's auth header from it — so an aux model on a DIFFERENT provider +/// would have its own correctly-resolved key overwritten by the session bearer +/// on the AUX host (H3/H4). Resolve by the aux model instead: +/// - an OAuth platform → a live resolver over ITS OWN pooled manager (so a grok +/// / claude-pro-max / copilot / codex aux model keeps mid-session refresh); +/// - `kimi-code`, or a platform-less model on the session's own coding endpoint +/// → the stamped session resolver, byte-identical; +/// - every API-key registry platform, and any `[model.*]` block pointed at a +/// third-party host → `None`, so the aux model's own key survives to the wire. +/// +/// SECURITY: no token is logged. +pub(crate) fn aux_bearer_resolver( + stamped: Option, + platform: Option, + base_url: &str, +) -> Option { + if let Some(oauth) = platform.and_then(kigi_models::PlatformId::oauth) { + return Some(auth_manager_bearer_resolver( + crate::auth::oauth_registry::global_manager_for( + &crate::auth::oauth_registry::pool_home(), + oauth, + ), + )); + } + if crate::agent::auth_method::platform_takes_session_credential(platform, base_url) { + return stamped; + } + None +} impl SessionActor { pub(super) async fn prepare_tool_definitions_timed(&self) -> (Vec, u64) { let mcp_wait_start = std::time::Instant::now(); @@ -206,7 +261,12 @@ impl SessionActor { fn auth_gate(&self, model_id: &str, base_url: &str) -> SessionTokenAuthGate { let byok = self.model_auth_facts(model_id).byok; let auth_method = self.auth_method_id.load(); - SessionTokenAuthGate::new(auth_method.as_deref(), byok, base_url) + SessionTokenAuthGate::new( + auth_method.as_deref(), + byok, + base_url, + self.model_platform(model_id), + ) } /// The [`AuthManager`](crate::auth::AuthManager) that governs INFERENCE auth /// for the model whose routing slug is `model` (the sampling config's @@ -238,26 +298,30 @@ impl SessionActor { // and resolves to the primary. let managed_key = self.managed_key_for_model(model); crate::auth::oauth_registry::manager_for_model( - &crate::util::kigi_home::kigi_home(), + &crate::auth::oauth_registry::pool_home(), managed_key.as_deref().unwrap_or(model), self.auth_manager.as_ref(), ) } /// Recover the managed catalog key (`{platform}/{model}`) for a routing slug /// from the live catalog. `None` for a bare / config / unlisted model. + /// + /// H5: the catalog KEY the picker selected + /// ([`crate::agent::models::ModelsManager::current_model_id`]) is + /// authoritative — `model` is the ambiguous bare slug. See + /// [`crate::agent::models::managed_key_for_slug`]. fn managed_key_for_model(&self, model: &str) -> Option { let models = self.models_manager.models(); - crate::agent::config::find_model_by_id(&models, model).and_then(|e| e.info().id.clone()) + let current = self.models_manager.current_model_id(); + crate::agent::models::managed_key_for_slug(&models, Some(current.0.as_ref()), model) } - /// Whether the aux/session model `model` routes to a generic device-code - /// OAuth platform (xai-grok). The gate for re-pointing an aux model's - /// bearer_resolver away from the session (Kimi) resolver — a first-party / - /// non-oauth model returns `false` and keeps the stamped session resolver. - fn model_is_oauth_platform(&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)) - .and_then(|(platform, _)| platform.oauth()) - .is_some() + /// The registry platform the routing slug `model` belongs to, from the SAME + /// lookup [`Self::auth_manager_for_model`] routes on. `None` for a bare / + /// `[model.*]` / unlisted model. + fn model_platform(&self, model: &str) -> Option { + let models = self.models_manager.models(); + let current = self.models_manager.current_model_id(); + crate::agent::models::platform_for_slug(&models, Some(current.0.as_ref()), model) } /// Whether `model` routes to the Claude Pro/Max OAuth-Messages platform /// (claude-pro-max) — the gate for the sampler's OAuth Messages adaptation @@ -266,22 +330,18 @@ impl SessionActor { /// 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 - }, - ) + self.model_platform(model).is_some_and(|platform| { + platform.oauth().is_some() + && platform.wire_api() == kigi_models::PlatformWireApi::Messages + }) } /// Whether `model` routes to the GitHub Copilot ChatCompletions platform /// (github-copilot) — the gate for the sampler's editor-identity headers + /// `X-Initiator`. Every other model returns `false`, keeping the other /// ChatCompletions providers byte-identical. fn model_is_github_copilot(&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.sends_copilot_editor_headers()) + self.model_platform(model) + .is_some_and(kigi_models::PlatformId::sends_copilot_editor_headers) } /// Whether `model` routes to the ChatGPT/Codex Responses platform /// (openai-codex) — the gate for the sampler's Codex identity headers @@ -289,28 +349,23 @@ impl SessionActor { /// returns `false`, keeping the API-key `openai` Responses request /// byte-identical. fn model_is_openai_codex(&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.sends_codex_responses_headers()) + self.model_platform(model) + .is_some_and(kigi_models::PlatformId::sends_codex_responses_headers) } /// 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 - /// `SamplerConfig`, re-point it at the AUX model's OWN platform manager when - /// the aux model is an oauth platform (xai-grok) — so a grok aux model never - /// inherits the live Kimi session bearer (→ api.x.ai). No-op for a - /// first-party / non-oauth aux model (keeps the stamped session resolver → - /// byte-identical). SECURITY: no token is logged. - pub(super) fn repoint_aux_bearer_resolver_for_oauth( + /// describe). Applies [`aux_bearer_resolver`] to the config + /// [`crate::agent::config::stamp_session_local_sampler_fields`] just stamped + /// the SESSION model's `bearer_resolver` onto. + pub(super) fn repoint_aux_bearer_resolver( &self, cfg: &mut kigi_sampler::SamplerConfig, slug: &str, ) { - if self.model_is_oauth_platform(slug) - && let Some(manager) = self.auth_manager_for_model(slug) - { - cfg.bearer_resolver = Some(auth_manager_bearer_resolver(manager)); - } + cfg.bearer_resolver = aux_bearer_resolver( + cfg.bearer_resolver.take(), + self.model_platform(slug), + &cfg.base_url, + ); } /// Emit a unified-log breadcrumb whenever the session-token refresh gate is /// evaluated with an **`Unknown`** per-model BYOK status on a session-based @@ -330,8 +385,9 @@ impl SessionActor { let ctx = serde_json::json!( { "site" : site, "model_byok" : gate.model_byok.as_str(), "is_session_based" : gate.is_session_based, "endpoint_is_first_party" : gate - .endpoint_is_first_party, "refresh_active" : refresh_active, "base_url" : - base_url, } + .endpoint_is_first_party, "endpoint_takes_session_credential" : gate + .endpoint_takes_session_credential, "refresh_active" : refresh_active, + "base_url" : base_url, } ); let sid = Some(self.session_info.id.0.as_ref()); if refresh_active { @@ -385,8 +441,12 @@ impl SessionActor { let creds = self.chat_state_handle.get_credentials().await; let model_facts = self.model_auth_facts(cfg.model.as_str()); let auth_method = self.auth_method_id.load(); - let gate = - SessionTokenAuthGate::new(auth_method.as_deref(), model_facts.byok, &cfg.base_url); + let gate = SessionTokenAuthGate::new( + auth_method.as_deref(), + model_facts.byok, + &cfg.base_url, + self.model_platform(cfg.model.as_str()), + ); let use_bearer_resolver = gate.active(); self.log_auth_gate_unknown("reconstruct_full_config", gate, &cfg.base_url); // Resolve the bearer from the ACTIVE model's OWN manager: a grok model @@ -589,17 +649,18 @@ impl SessionActor { slug: &str, ) -> Option { let creds = self.chat_state_handle.get_credentials().await; - // Resolve the aux token by the aux model's OWN platform: a grok - // (oauth-platform) aux model draws its pooled grok token or `None` — - // NEVER the primary Kimi session token (which `resolve_credentials` - // would otherwise stamp onto an api.x.ai request). A first-party / - // non-oauth aux model still gets the primary (byte-identical). - let session_key = crate::auth::oauth_registry::session_key_for_model( - &crate::util::kigi_home::kigi_home(), + let models = self.models_manager.models(); + // Resolve the aux token by the aux model's OWN platform AND endpoint: a + // grok (oauth-platform) aux model draws its pooled grok token or `None`, + // and an API-key registry platform draws NOTHING — NEVER the primary + // Kimi session token (which `resolve_credentials` would otherwise stamp + // onto an api.x.ai / api.deepseek.com request). The first-party + // subscription channel still gets the primary (byte-identical). + let session_key = crate::auth::oauth_registry::session_key_for_catalog_model( + &models, slug, self.auth_manager.as_ref(), ); - let models = self.models_manager.models(); let endpoints = self.models_manager.endpoints(); crate::agent::config::resolve_aux_model_sampling_config( slug, @@ -625,10 +686,11 @@ impl SessionActor { &active_session_config, Some(self.max_retries), ); - // LEAK 1b: a grok aux classifier must not inherit the SESSION model's - // (Kimi) bearer_resolver stamped above; re-point it at grok's own - // manager (its pooled token, or None). No-op for a non-oauth aux. - self.repoint_aux_bearer_resolver_for_oauth(&mut cfg, slug); + // LEAK 1b: the aux classifier must not inherit the SESSION model's + // (Kimi) bearer_resolver stamped above — re-point it at an OAuth aux + // model's own manager, or clear it for an API-key-platform / third-party + // aux endpoint. No-op for the first-party subscription channel. + self.repoint_aux_bearer_resolver(&mut cfg, slug); let model = cfg.model.clone(); let client = kigi_sampler::SamplingClient::new(cfg) .map_err(|e| { @@ -774,6 +836,8 @@ impl SessionActor { session_id = % self.session_info.id.0, is_session_based = gate .is_session_based, model_byok = gate.model_byok.as_str(), endpoint_is_first_party = gate.endpoint_is_first_party, + endpoint_takes_session_credential = gate + .endpoint_takes_session_credential, "auth recovery: sampler 401 not refreshable (api-key auth) — surfacing 401", ); kigi_log::unified_log::warn( @@ -783,7 +847,9 @@ impl SessionActor { { "kind" : error.kind.as_str(), "status_code" : error .status_code, "is_session_based" : gate.is_session_based, "model_byok" : gate.model_byok.as_str(), - "endpoint_is_first_party" : gate.endpoint_is_first_party, } + "endpoint_is_first_party" : gate.endpoint_is_first_party, + "endpoint_takes_session_credential" : gate + .endpoint_takes_session_credential, } )), ); } @@ -1036,6 +1102,15 @@ impl SessionActor { .map(|c| c.model) .unwrap_or_default(); let Some(ref key) = current_key else { return }; + // M7: a registry-platform model's key comes from that platform's + // credential resolved into its catalog entry — it is NEVER a + // `[model.*]` block. With the session gate now inactive for every + // API-key platform, those turns all fell through to here and paid a + // `load_effective_config()` disk read PER TURN, then logged a + // permanently false "Model not found in config.toml [model.*]" warning. + if self.model_platform(¤t_model_id).is_some() { + return; + } let Some(new_key) = self.reload_api_key_from_config(¤t_model_id) else { return; }; @@ -1161,7 +1236,7 @@ mod bearer_resolver_tests { .oauth() .expect("xai-grok carries an OAuthConfig"); let grok = crate::auth::oauth_registry::global_manager_for( - &crate::util::kigi_home::kigi_home(), + &crate::auth::oauth_registry::pool_home(), oauth, ); assert_ne!( 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 6e974ab..e91961d 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 @@ -487,21 +487,27 @@ fn session_token_auth_gate_truth_table() { use crate::agent::auth_method::{ModelByok, session_token_auth_gate as gate}; // Non-session methods never refresh, regardless of BYOK status or endpoint. for fp in [false, true] { - assert!(!gate(false, ModelByok::NotByok, fp)); - assert!(!gate(false, ModelByok::Byok, fp)); - assert!(!gate(false, ModelByok::Unknown, fp)); - // Session method: a definite classification ignores the endpoint — - // NotByok always refreshes (only ever routes to the session endpoint), - // a genuine per-model Byok never does. - assert!(gate(true, ModelByok::NotByok, fp)); - assert!(!gate(true, ModelByok::Byok, fp)); + assert!(!gate(false, ModelByok::NotByok, fp, true)); + assert!(!gate(false, ModelByok::Byok, fp, true)); + assert!(!gate(false, ModelByok::Unknown, fp, true)); + // Session method on an endpoint that DOES take the session credential + // (kimi-code, an OAuth platform's own pool, or a bare / [model.*] + // model): a definite classification ignores the endpoint — NotByok + // refreshes, a genuine per-model Byok never does. + assert!(gate(true, ModelByok::NotByok, fp, true)); + assert!(!gate(true, ModelByok::Byok, fp, true)); + // …and an API-key registry platform endpoint is refused on every arm, + // first-party flag included: that is the leak guard. + assert!(!gate(true, ModelByok::NotByok, fp, false)); + assert!(!gate(true, ModelByok::Byok, fp, false)); + assert!(!gate(true, ModelByok::Unknown, fp, false)); } // Session method + Unknown BYOK: refresh only against a first-party xAI // host, so a transiently-unclassifiable config can't demote a live session // (the stale-token 401 regression) yet the session token never leaks to a // third-party BYOK endpoint. This arm was unconditionally `false` pre-fix. - assert!(gate(true, ModelByok::Unknown, true)); - assert!(!gate(true, ModelByok::Unknown, false)); + assert!(gate(true, ModelByok::Unknown, true, true)); + assert!(!gate(true, ModelByok::Unknown, false, true)); } /// Pre-fix, the gate read `auth_type` and skipped recovery here, 401'ing every diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/session_bearer_leak_platform_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/session_bearer_leak_platform_tests.rs new file mode 100644 index 0000000..17fa575 --- /dev/null +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/session_bearer_leak_platform_tests.rs @@ -0,0 +1,289 @@ +//! LEAK GUARD, part 2: the model→platform lookup (H5) and the AUX resolver +//! decision (H3/H4). Shares the fixtures in +//! [`super::session_bearer_leak_tests`]; see that module's header for the chain +//! and the storage-discipline contract. + +use super::session_bearer_leak_tests::{ + KIMI_TOKEN, actor_on_managed_model, actor_with_catalog, managed_entry, +}; +use super::*; +use kigi_sampler::BearerResolver; +use std::sync::Arc; + +#[tokio::test(flavor = "current_thread")] +async fn dual_credential_slug_collision_resolves_the_selected_oauth_platform() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + // (api-key twin, oauth twin, shared slug, host) + for (api_key_twin, oauth_twin, slug, base_url) in [ + ( + "xai/grok-4.5", + "xai-grok/grok-4.5", + "grok-4.5", + "https://api.x.ai/v1", + ), + ( + "anthropic/claude-opus-4-8", + "claude-pro-max/claude-opus-4-8", + "claude-opus-4-8", + "https://api.anthropic.com/v1", + ), + ( + "openai/gpt-5.5-codex", + "openai-codex/gpt-5.5-codex", + "gpt-5.5-codex", + "https://chatgpt.com/backend-api/codex", + ), + ] { + // API-key platform FIRST, exactly as `PlatformId::ALL` orders them. + let catalog = vec![ + managed_entry(api_key_twin, slug, base_url), + managed_entry(oauth_twin, slug, base_url), + ]; + let (_dir, actor, _rx) = actor_with_catalog(catalog, oauth_twin, "unused").await; + let cfg = actor.reconstruct_full_config().await; + + let resolver = cfg.bearer_resolver.as_ref().unwrap_or_else(|| { + panic!( + "{oauth_twin}: selecting the OAuth twin must keep a LIVE bearer_resolver \ + (mid-session refresh); resolving {api_key_twin} instead drops it" + ) + }); + assert_ne!( + resolver.current_bearer(), + Some(KIMI_TOKEN.to_string()), + "{oauth_twin}: the resolver must read its OWN pool, never the Kimi primary" + ); + + let platform = kigi_models::parse_managed_model_key(oauth_twin) + .expect("managed key") + .0; + assert_eq!( + cfg.anthropic_oauth, + platform.wire_api() == kigi_models::PlatformWireApi::Messages, + "{oauth_twin}: the Claude OAuth Messages adaptation must follow the \ + SELECTED platform" + ); + assert_eq!( + cfg.openai_codex, + platform.sends_codex_responses_headers(), + "{oauth_twin}: the Codex identity headers must follow the SELECTED platform" + ); + assert_eq!( + cfg.github_copilot, + platform.sends_copilot_editor_headers(), + "{oauth_twin}: the Copilot editor headers must follow the SELECTED platform" + ); + } + }) + .await; +} + +/// The other half of H5: selecting the API-KEY twin of a colliding slug must +/// still resolve the API-key platform — no bearer_resolver, no adaptations. The +/// unified lookup must not simply prefer OAuth. +#[tokio::test(flavor = "current_thread")] +async fn dual_credential_slug_collision_resolves_the_selected_api_key_platform() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let catalog = vec![ + managed_entry( + "anthropic/claude-opus-4-8", + "claude-opus-4-8", + "https://api.anthropic.com/v1", + ), + managed_entry( + "claude-pro-max/claude-opus-4-8", + "claude-opus-4-8", + "https://api.anthropic.com/v1", + ), + ]; + let (_dir, actor, _rx) = + actor_with_catalog(catalog, "anthropic/claude-opus-4-8", "sk-ant-byok").await; + let cfg = actor.reconstruct_full_config().await; + assert!( + cfg.bearer_resolver.is_none(), + "selecting the API-key twin must get NO session bearer resolver" + ); + assert!( + !cfg.anthropic_oauth, + "the API-key Anthropic Messages request must stay byte-identical" + ); + assert_eq!(cfg.api_key.as_deref(), Some("sk-ant-byok")); + }) + .await; +} + +/// MANDATORY counterpart: the subscription-OAuth platforms have NON-first-party +/// base URLs, so the fix must not disable their resolver. Each must still get a +/// LIVE `bearer_resolver` — and it must read THAT platform's own pooled +/// `AuthManager`, never the Kimi primary. The pooled managers are empty here (a +/// TempDir pool home), which is what makes `current_bearer() == None` a proof +/// that the Kimi bearer cannot be what they resolve. +#[tokio::test(flavor = "current_thread")] +async fn oauth_platform_models_keep_a_live_resolver_from_their_own_pool() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + for (catalog_key, slug, base_url) in [ + ( + "claude-pro-max/claude-opus-4-8", + "claude-opus-4-8", + "https://api.anthropic.com/v1", + ), + ( + "github-copilot/gpt-4.1", + "gpt-4.1", + "https://api.githubcopilot.com", + ), + ( + "xai-grok/grok-4-latest", + "grok-4-latest", + "https://api.x.ai/v1", + ), + ( + "openai-codex/gpt-5.5", + "gpt-5.5", + "https://chatgpt.com/backend-api/codex", + ), + ] { + let (_dir, actor, _rx) = + actor_on_managed_model(catalog_key, slug, base_url, "unused").await; + let cfg = actor.reconstruct_full_config().await; + let resolver = cfg.bearer_resolver.as_ref().unwrap_or_else(|| { + panic!("{catalog_key}: must keep a live bearer_resolver for refresh") + }); + assert_ne!( + resolver.current_bearer(), + Some(KIMI_TOKEN.to_string()), + "{catalog_key}: the Kimi bearer must never be what it resolves" + ); + + // The resolver is LIVE over that platform's pooled manager: a + // token rotated inside the pool is observed by the + // already-built resolver (this is what mid-session refresh + // does). The pool is read here, never mutated. + let pooled = crate::auth::oauth_registry::manager_for_model( + &crate::auth::oauth_registry::pool_home(), + catalog_key, + actor.auth_manager.as_ref(), + ) + .expect("an OAuth platform always resolves a manager"); + assert!( + !Arc::ptr_eq( + &pooled, + actor.auth_manager.as_ref().expect("primary is present") + ), + "{catalog_key}: must route to its OWN pooled manager, not the Kimi primary" + ); + assert_eq!( + resolver.current_bearer(), + pooled.current_or_expired().map(|a| a.key), + "{catalog_key}: the resolver must read THIS platform's pooled manager" + ); + } + }) + .await; +} + +/// H3/H4 — the stamped AUX paths (image-describe, the auto-mode classifier and +/// the session-summary client all funnel through `aux_bearer_resolver`). +/// `stamp_session_local_sampler_fields` copies the SESSION model's resolver onto +/// every aux config and `SamplingClient::post` REPLACES the request's auth +/// header from it, so an API-key-platform aux model would have its own key +/// overwritten by the Kimi bearer ON THE AUX HOST. +/// +/// Revert-to-red: returning `stamped` unconditionally (the pre-fix +/// "re-point only when the aux model is OAuth" shape) makes the deepseek / +/// openai / `[model.*]`-on-openai.com rows resolve `KIMI_TOKEN`. +#[test] +fn aux_bearer_resolver_clears_the_session_resolver_off_a_third_party_aux_host() { + #[derive(Debug)] + struct Fixed(&'static str); + impl BearerResolver for Fixed { + fn current_bearer(&self) -> Option { + Some(self.0.to_string()) + } + } + let stamped: kigi_sampler::SharedBearerResolver = Arc::new(Fixed(KIMI_TOKEN)); + let platform = + |key: &str| kigi_models::parse_managed_model_key(key).map(|(platform, _)| platform); + + // Cleared: every API-key registry platform, and a `[model.*]` aux model + // pointed at a third-party host. + for (key, base_url) in [ + ("deepseek/deepseek-chat", "https://api.deepseek.com/v1"), + ("openai/gpt-5-mini", "https://api.openai.com/v1"), + ( + "moonshot-cn/kimi-k2-turbo-preview", + "https://api.moonshot.cn/v1", + ), + ] { + assert!( + crate::session::acp_session::sampler_turn::aux_bearer_resolver( + Some(stamped.clone()), + platform(key), + base_url, + ) + .is_none(), + "LEAK: an aux model on {base_url} must not inherit the session bearer resolver" + ); + } + assert!( + crate::session::acp_session::sampler_turn::aux_bearer_resolver( + Some(stamped.clone()), + None, + "https://api.openai.com/v1", + ) + .is_none(), + "LEAK: a [model.*] aux model on a third-party host must not inherit it either" + ); + + // Kept (byte-identical): the first-party subscription channel and a + // platform-less aux model on the session's own endpoint. + for (key, base_url) in [ + ( + Some("kimi-code/kimi-for-coding"), + kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url, + ), + (None, kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url), + (None, "http://127.0.0.1:4141/v1"), + ] { + let resolved = crate::session::acp_session::sampler_turn::aux_bearer_resolver( + Some(stamped.clone()), + key.and_then(platform), + base_url, + ) + .unwrap_or_else(|| panic!("{key:?} @ {base_url} must keep the session resolver")); + assert_eq!(resolved.current_bearer(), Some(KIMI_TOKEN.to_string())); + } + + // Re-pointed: an OAuth aux model gets a LIVE resolver over its OWN pool + // (empty here), never the stamped Kimi one. + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("runtime for the pooled manager's refresh task"); + rt.block_on(async { + for key in [ + "xai-grok/grok-4-latest", + "claude-pro-max/claude-opus-4-8", + "github-copilot/gpt-4.1", + "openai-codex/gpt-5.5", + ] { + let resolved = crate::session::acp_session::sampler_turn::aux_bearer_resolver( + Some(stamped.clone()), + platform(key), + "https://example.invalid/v1", + ) + .unwrap_or_else(|| panic!("{key} must keep a live resolver from its own pool")); + assert_ne!( + resolved.current_bearer(), + Some(KIMI_TOKEN.to_string()), + "{key}: the aux resolver must never resolve the Kimi session bearer" + ); + } + }); +} diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/session_bearer_leak_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/session_bearer_leak_tests.rs new file mode 100644 index 0000000..1f0565d --- /dev/null +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/session_bearer_leak_tests.rs @@ -0,0 +1,436 @@ +//! LEAK GUARD (bearer_resolver channel): the primary (Kimi) subscription bearer +//! must never be stamped on a request to a host that does not own it. +//! +//! Chain the guard closes: a session-based ACP method (`cached_token` / +//! `kimi-code` / any OAuth platform) + a selected API-key-platform model +//! classifies `ModelByok::NotByok` (the model carries no `[model.*]` key), the +//! pre-fix `session_token_auth_gate` returned `true` unconditionally on that +//! arm, `auth_manager_for_model` fell through to the primary Kimi manager for a +//! non-OAuth platform, and `SamplingClient::post` then REPLACED the correctly +//! resolved provider key with the Kimi bearer on the wire. +//! +//! The `api_key` half of the same defect (the config never even gets the +//! provider key, because `resolve_credentials` stamps the session token) is +//! pinned in `agent/mvp_agent/tests/api_key_channel_leak_tests.rs`, which drives +//! the real `prepare_sampling_config_for_model` resolution path. These tests +//! deliberately do NOT hand-stamp a provider key except where the assertion is +//! about the resolver overwriting one that already resolved correctly. +//! +//! The counterpart contract these tests also pin: the four subscription-OAuth +//! platforms have non-first-party base URLs but MUST keep a live +//! `bearer_resolver` drawn from their OWN pooled `AuthManager`, or they lose +//! mid-session token refresh. +//! +//! STORAGE DISCIPLINE (H6): nothing here touches the developer's real `~/.kigi` +//! and nothing hot-swaps the process-global OAuth pool. Under `cfg(test)` +//! `oauth_registry::pool_home()` is a process-lifetime `TempDir`, so every +//! pooled manager is empty — which is exactly what the assertions need (a live +//! resolver that is provably NOT the Kimi one). + +use super::support::*; +use super::*; +use crate::agent::auth_method::ModelByok; +use crate::agent::config::{ModelAuthFacts, ModelEntry, ModelInfo}; +use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig}; +use kigi_sampler::BearerResolver; +use std::sync::Arc; +use tokio::sync::mpsc; + +/// The primary session bearer. Any occurrence of this string in an outgoing +/// request to a third-party host is the defect. +pub(super) const KIMI_TOKEN: &str = "kimi-subscription-token-DO-NOT-LEAK"; + +/// `(tempdir, manager)` standing in for the session's primary Kimi +/// `AuthManager`, holding a live (unexpired) OAuth session bearer. +fn kimi_primary() -> (tempfile::TempDir, Arc) { + let dir = tempfile::tempdir().expect("tempdir"); + let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + am.hot_swap(KimiAuth { + key: KIMI_TOKEN.to_string(), + auth_mode: AuthMode::OAuth, + refresh_token: Some("rt".into()), + expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), + ..KimiAuth::test_default() + }); + (dir, am) +} + +/// One catalog entry: catalog key `catalog_key`, routing slug `slug`, routed at +/// `base_url`, carrying no credential of its own (the shape every fetched +/// registry model has). +pub(super) fn managed_entry(catalog_key: &str, slug: &str, base_url: &str) -> (String, ModelEntry) { + let mut info = ModelInfo::fallback(slug); + info.id = Some(catalog_key.to_string()); + info.base_url = base_url.to_string(); + ( + catalog_key.to_string(), + ModelEntry { + info, + api_key: None, + env_key: None, + api_base_url: None, + }, + ) +} + +/// A `SessionActor` on a session-based ACP method with a live Kimi primary, +/// whose live catalog holds `catalog` and whose SELECTED model is the catalog +/// key `selected` (the picker's own notion of "current"). `wire_key` is the +/// already-correctly-resolved provider credential sitting in chat state. +/// +/// The per-model BYOK memo is pinned to `NotByok` on purpose: that is what a +/// fetched registry model actually resolves to (`resolve_model_auth_facts` only +/// ever sees `default_models.json` + `[model.*]`), and pinning it keeps the test +/// independent of the developer's on-disk `~/.kigi/config.toml`. +pub(super) async fn actor_with_catalog( + catalog: Vec<(String, ModelEntry)>, + selected: &str, + wire_key: &str, +) -> ( + tempfile::TempDir, + Arc, + mpsc::UnboundedReceiver, +) { + let (dir, am) = kimi_primary(); + let (gateway_tx, _gateway_rx) = mpsc::unbounded_channel(); + let (persistence_tx, persistence_rx) = mpsc::unbounded_channel(); + let mut actor = create_test_actor(50_000, 200_000, 85, gateway_tx, persistence_tx).await; + actor.auth_manager = Some(am); + actor.auth_method_id = test_auth_method_id("cached_token"); + + let mut selected_entry = None; + for (key, entry) in catalog { + if key == selected { + selected_entry = Some(entry.clone()); + } + actor.models_manager.insert_test_entry(key, entry); + } + let selected_entry = selected_entry.expect("the selected key must be in the catalog"); + actor + .models_manager + .set_current_model_id(acp::ModelId::new(selected.to_string())); + + let slug = selected_entry.info().model.clone(); + actor + .chat_state_handle + .update_sampling_config(kigi_sampling_types::SamplingConfig { + base_url: selected_entry.info().base_url.clone(), + model: slug.clone(), + max_completion_tokens: None, + temperature: None, + top_p: None, + api_backend: Default::default(), + chat_compat: Default::default(), + extra_headers: Default::default(), + context_window: std::num::NonZeroU64::new(200_000).unwrap(), + reasoning_effort: None, + stream_tool_calls: None, + }); + actor + .chat_state_handle + .update_credentials(kigi_chat_state::Credentials { + api_key: Some(wire_key.to_string()), + auth_type: kigi_chat_state::AuthType::SessionToken, + ..Default::default() + }); + actor.model_auth_facts.replace(Some(( + slug, + ModelAuthFacts { + byok: ModelByok::NotByok, + auth_scheme: Default::default(), + }, + ))); + (dir, Arc::new(actor), persistence_rx) +} + +/// Single-entry convenience over [`actor_with_catalog`]. +pub(super) async fn actor_on_managed_model( + catalog_key: &str, + slug: &str, + base_url: &str, + wire_key: &str, +) -> ( + tempfile::TempDir, + Arc, + mpsc::UnboundedReceiver, +) { + actor_with_catalog( + vec![managed_entry(catalog_key, slug, base_url)], + catalog_key, + wire_key, + ) + .await +} + +/// THE leak test, at the wire. A `deepseek/deepseek-chat` turn on a session +/// (`cached_token`) method with a live Kimi primary must send DeepSeek's own key +/// — the Kimi subscription bearer must not appear anywhere in the request. +/// +/// Revert-to-red: dropping `endpoint_takes_session_credential` from +/// `session_token_auth_gate` puts `Bearer ` on this request. +#[tokio::test(flavor = "multi_thread")] +async fn deepseek_turn_under_a_kimi_session_sends_no_kimi_bearer_on_the_wire() { + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("POST")) + .and(wiremock::matchers::path("/chat/completions")) + .respond_with( + wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "cmpl-1", + "object": "chat.completion", + "created": 0, + "model": "deepseek-chat", + "choices": [{ + "index": 0, + "message": { "role": "assistant", "content": "ok" }, + "finish_reason": "stop" + }] + })), + ) + .mount(&server) + .await; + let uri = server.uri(); + + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (_dir, actor, _rx) = actor_on_managed_model( + "deepseek/deepseek-chat", + "deepseek-chat", + &uri, + "sk-deepseek-provider-key", + ) + .await; + + let cfg = actor.reconstruct_full_config().await; + assert!( + cfg.bearer_resolver.is_none(), + "an API-key platform model must get NO session bearer resolver" + ); + let client = + kigi_sampler::SamplingClient::new(cfg).expect("sampling client must construct"); + let _ = client + .chat_completion(kigi_sampling_types::ChatCompletionRequest::new( + "deepseek-chat", + vec![kigi_sampling_types::ChatRequestMessage::user("hi")], + )) + .await; + }) + .await; + + let requests = server + .received_requests() + .await + .expect("wiremock records requests"); + assert_eq!(requests.len(), 1, "exactly one inference request was sent"); + let auth = requests[0] + .headers + .get("authorization") + .and_then(|v| v.to_str().ok()) + .expect("the request must carry an Authorization header") + .to_string(); + assert!( + !auth.contains(KIMI_TOKEN), + "the Kimi subscription bearer must never reach a third-party inference host" + ); + assert_eq!( + auth, "Bearer sk-deepseek-provider-key", + "the correctly-resolved provider key must survive to the wire" + ); +} + +/// The same guard for every other API-key registry platform shape: OpenAI +/// (Responses), Anthropic (x-api-key/Messages), Groq, Together and Z.AI CN — all +/// classify `NotByok`, all route to a non-first-party host, none may receive a +/// session bearer resolver. +#[tokio::test(flavor = "current_thread")] +async fn api_key_platform_models_get_no_session_bearer_resolver() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + for (catalog_key, slug, base_url) in [ + ("openai/gpt-5.2", "gpt-5.2", "https://api.openai.com/v1"), + ( + "anthropic/claude-opus-4-8", + "claude-opus-4-8", + "https://api.anthropic.com/v1", + ), + ("groq/llama-4", "llama-4", "https://api.groq.com/openai/v1"), + ("together/qwen-3", "qwen-3", "https://api.together.xyz/v1"), + ( + "zai-coding-cn/glm-5", + "glm-5", + "https://open.bigmodel.cn/api/paas/v4", + ), + ] { + let (_dir, actor, _rx) = + actor_on_managed_model(catalog_key, slug, base_url, "sk-provider-key").await; + let cfg = actor.reconstruct_full_config().await; + assert!( + cfg.bearer_resolver.is_none(), + "{catalog_key}: an API-key platform must get no session bearer resolver" + ); + assert_eq!( + cfg.api_key.as_deref(), + Some("sk-provider-key"), + "{catalog_key}: the provider key must stay on the config" + ); + } + }) + .await; +} + +/// C2 at the resolver channel: a `[model.*]` entry has NO platform +/// (`info.id == None`), which used to be a blanket allow. Pointed at a +/// third-party host it must get no session resolver; pointed at the session's +/// own coding endpoint (a `KIGI_CODE_BASE_URL` deployment or a local dev proxy) +/// it must keep one — that is why the predicate is not `is_first_party_url`. +/// +/// Revert-to-red: making the `None` arm of `platform_takes_session_credential` +/// return `true` again puts a Kimi resolver on the openai.com config. +#[tokio::test(flavor = "current_thread")] +async fn config_model_entry_takes_a_session_resolver_only_on_its_own_endpoint() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + for base_url in ["https://api.openai.com/v1", "https://api.deepseek.com/v1"] { + let mut info = ModelInfo::fallback("gpt-4o"); + info.id = None; // a `[model.gpt-4o]` block + info.base_url = base_url.to_string(); + let entry = ModelEntry { + info, + api_key: None, + // An env_key that is NOT set: `has_own_credentials()` probes + // `std::env::var` at call time, so this classifies NotByok. + env_key: None, + api_base_url: None, + }; + let (_dir, actor, _rx) = + actor_with_catalog(vec![("gpt-4o".to_string(), entry)], "gpt-4o", "").await; + let cfg = actor.reconstruct_full_config().await; + assert!( + cfg.bearer_resolver.is_none(), + "LEAK: a [model.*] block at {base_url} must get no session bearer resolver" + ); + } + + for base_url in [ + kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url, + "http://127.0.0.1:4141/v1", + ] { + let mut info = ModelInfo::fallback("kigi-4.5"); + info.id = None; + info.base_url = base_url.to_string(); + let entry = ModelEntry { + info, + api_key: None, + env_key: None, + api_base_url: None, + }; + let (_dir, actor, _rx) = + actor_with_catalog(vec![("kigi-4.5".to_string(), entry)], "kigi-4.5", "").await; + let resolver = actor + .reconstruct_full_config() + .await + .bearer_resolver + .expect("the session's own endpoint keeps the session resolver"); + assert_eq!( + resolver.current_bearer(), + Some(KIMI_TOKEN.to_string()), + "{base_url}: a custom deployment / dev proxy is unchanged" + ); + } + }) + .await; +} + +/// H5 — the slug collision. A user holding BOTH an xAI API key and a Grok +/// subscription has `xai/grok-4.5` AND `xai-grok/grok-4.5` in one catalog, in +/// `PlatformId::ALL` order (`Xai`(15) before `XaiGrok`(25)) and with the SAME +/// routing slug. `cfg.model` is that bare slug, so the auth layer used to +/// first-match the API-key entry: no bearer_resolver, no live refresh (the +/// session dies ~1h in with an unrecoverable 401), and — for the Anthropic and +/// Codex twins — the OAuth Messages adaptation and the Codex identity headers +/// silently dropped. +/// +/// The catalog KEY the picker selected is now authoritative. +/// +/// Revert-to-red: resolving the platform from `find_model_by_id(models, slug)` +/// instead of `current_model_id` resolves `xai/grok-4.5` / +/// `anthropic/claude-opus-4-8` / `openai/gpt-5.5-codex` and every assertion +/// below fails. +#[tokio::test(flavor = "current_thread")] +async fn kimi_first_party_model_still_rides_the_primary_session_bearer() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (_dir, actor, _rx) = actor_on_managed_model( + "kimi-code/kimi-for-coding", + "kimi-for-coding", + kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url, + "stale-buffered-token", + ) + .await; + + let cfg = actor.reconstruct_full_config().await; + let resolver = cfg + .bearer_resolver + .as_ref() + .expect("the subscription model must keep the live session resolver"); + assert_eq!( + resolver.current_bearer(), + Some(KIMI_TOKEN.to_string()), + "the first-party model resolves the primary session bearer" + ); + + actor.refresh_token_if_expired().await; + assert_eq!( + actor + .chat_state_handle + .get_credentials() + .await + .api_key + .as_deref(), + Some(KIMI_TOKEN), + "the first-party pre-flight refresh must still heal the stale key" + ); + }) + .await; +} + +/// The persistence half of the defect: `refresh_token_if_expired` used to write +/// the Kimi session token into `chat_state` `creds.api_key` for ANY +/// session-method turn, from where it propagated to subagents and aux configs. +/// A deepseek turn must leave the provider key untouched. +/// +/// M7 rides along: a registry-platform model must not fall into +/// `reload_api_key_from_config` at all (a `load_effective_config()` disk read +/// per turn plus a permanently false "not found in config.toml" warning), so +/// the key is left exactly as resolved. +#[tokio::test(flavor = "current_thread")] +async fn preflight_refresh_never_writes_the_kimi_token_into_a_platform_credential() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (_dir, actor, _rx) = actor_on_managed_model( + "deepseek/deepseek-chat", + "deepseek-chat", + "https://api.deepseek.com/v1", + "sk-deepseek-provider-key", + ) + .await; + + actor.refresh_token_if_expired().await; + + assert_eq!( + actor + .chat_state_handle + .get_credentials() + .await + .api_key + .as_deref(), + Some("sk-deepseek-provider-key"), + "the Kimi session token must never overwrite a platform credential" + ); + }) + .await; +}