ZacharyZhang-NY/Kigi-CLI
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

update

Browse Source
This commit is contained in:
ZacharyZhang-NY
2026-07-22 11:08:54 -04:00
parent 422e241e13
commit 2d00a4e6e6
19 changed files with 1955 additions and 296 deletions
Download Patch File Download Diff File Expand all files Collapse all files
AGENTS.md
+40
View File
@@ -259,6 +259,46 @@ edges stay deterministic Rust. The harness appends a terminal
token (`resolve_generic_oauth_tokens`, refreshed on expiry) and routes token (`resolve_generic_oauth_tokens`, refreshed on expiry) and routes
`platform.oauth().is_some()` → `platform.base_url()` (kimi-code alone → `platform.oauth().is_some()` → `platform.base_url()` (kimi-code alone →
`proxy_url()`). Tokens/codes/verifiers are NEVER logged. `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 - Model metadata (context window, thinking levels) comes from the provider
wire when served; metadata-poor listings are enriched from models.dev wire when served; metadata-poor listings are enriched from models.dev
(`kigi-models/src/enrichment.rs` — bundled raw snapshot regenerated by (`kigi-models/src/enrichment.rs` — bundled raw snapshot regenerated by
README.md
+58 -92
View File
@@ -8,19 +8,14 @@
autonomous, self-verifying agent loops — planned, parallelized, autonomous, self-verifying agent loops — planned, parallelized,
adversarially verified, and merged back, end to end.</p> adversarially verified, and merged back, end to end.</p>
**Kigi** started as an unofficial Kimi Code CLI community build, a **Kigi** is a coding agent that lives in your terminal. It reads the repo,
terminal-based AI coding agent re-targeted at the Kimi Code subscription API, writes the patch, runs the tests, and keeps going while you do something else.
built on the Apache-2.0 sources of Full-screen, headless in CI with `-p`, or docked in your editor over ACP.
[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).
It runs as a full-screen TUI that understands your codebase, edits files, **Already paying for Claude Pro/Max, ChatGPT Plus/Pro, GitHub Copilot, or
executes shell commands, searches the web, and manages long-running tasks, Grok? Sign in and use it.** No API key, no second bill. Rather bring your own
interactively, headlessly for scripting/CI, or embedded in editors via the key? OpenAI, Anthropic, Google, DeepSeek, Groq, Moonshot and
Agent Client Protocol (ACP). [two dozen more](#providers-and-api-keys) are wired in.
[Installation](#installation) · [Installation](#installation) ·
[Graph engineering](#graph-engineering) · [Graph engineering](#graph-engineering) ·
@@ -39,10 +34,6 @@ Agent Client Protocol (ACP).
## Installation ## 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 ```sh
# macOS / Linux # macOS / Linux
curl -fsSL https://raw.githubusercontent.com/ZacharyZhang-NY/Kigi-CLI/main/install.sh | bash 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 ```sh
kigi --version # kigi 0.1.1 … unofficial Kimi Code CLI community build kigi login # pick a provider, sign in
kigi login # sign in with your Kimi Code subscription (device-code flow) kigi # go
kigi # start the TUI
``` ```
The installer verifies every download against the release's `SHA256SUMS`, Single file, no runtime. macOS and Linux on arm64/x86_64, Windows on x86_64,
installs into `~/.kigi/bin/kigi` (`%USERPROFILE%\.kigi\bin\kigi.exe` on checksummed against the release's `SHA256SUMS`. `kigi update` handles upgrades.
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.
## Graph engineering ## Graph engineering
Kigi is the first CLI to ship *graph engineering* as a first-class command: Every other agent runs a loop: think, act, repeat — one thread, one thing at a
where a loop drives one agent, a graph is the programmable organization time. `/graph` runs a dependency graph instead. Work that doesn't block other
connecting many. work happens at the same time, in separate worktrees, and nothing merges until
something else has tried to tear it apart.
``` ```
/graph <objective> [--budget <tokens>] # decompose + run fully autonomously /graph <objective> [--budget <tokens>] # decompose + run fully autonomously
@@ -81,49 +67,46 @@ connecting many.
/graph clear # abandon the graph /graph clear # abandon the graph
``` ```
One `/graph <objective>` runs the whole closed loop: a planner subagent One command runs the whole thing, start to finish:
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.
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 State lives in `.kigi/graph.jsonl`, next to your code. Close the laptop, come
# macOS / Linux back tomorrow, `/graph resume`. A teammate can pick it up from the same file.
echo 'export KIGI_GRAPH=0' >> ~/.zshrc # or ~/.bashrc / ~/.bash_profile
```
```powershell On by default. `KIGI_GRAPH=0` turns it off; `KIGI_GRAPH_CONCURRENCY` (default
# Windows PowerShell 3) controls how many nodes run at once.
[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).
## Providers and API keys ## Providers and API keys
Kigi ships a fixed registry of 25 platforms: the Kimi Code subscription plus 29 platforms ship compiled in: 5 you sign into, 24 you hand a key. Nothing is
24 API-key providers. There is no dynamic provider registration; each is a registered at runtime — if it's not in this list, it's not there.
compiled-in spec.
**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 | | Platform id | Provider | Sign-in |
| ----------- | -------------------------------- | ------------------------------------------- | | ---------------- | ------------------------- | ------------------------------------------- |
| `kimi-code` | `https://api.kimi.com/coding/v1` | Kimi Code subscription OAuth (`kigi login`) | | `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 You get whatever models your plan actually serves — the list is fetched at
`~/.kigi/config.toml` under `[platforms.<id>]`. The environment wins, a sign-in, not hardcoded. (ChatGPT/Codex is the exception: its backend publishes
platform-scoped name beats a generic one, and keys are never logged. 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 | | Provider | Platform id | API key env |
| ------------------------- | ---------------------- | ----------------------------------------------- | | ------------------------- | ---------------------- | ----------------------------------------------- |
@@ -166,24 +149,11 @@ api_key = "sk-..."
api_key = "xai-..." api_key = "xai-..."
``` ```
On login and on startup Kigi syncs each configured platform's model list Model lists sync on startup. Pick one with `/model`, set its thinking level
from `GET {base}/models` and shows the merged catalog in the model picker with `/effort`.
(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.
Each platform's base URL can be re-pointed for dev/test with Web `search`/`fetch` need a Kimi Code subscription; API-key sessions run
`KIGI_<PLATFORM>_BASE_URL` (e.g. `KIGI_CODE_BASE_URL`, without them, same as the official client.
`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.
## Building from source ## Building from source
@@ -199,19 +169,15 @@ launcher at `bin/protoc`; install dotslash (`brew install dotslash` or
## Coexistence with the official Kimi CLI ## Coexistence with the official Kimi CLI
Kigi is not affiliated with Moonshot AI or xAI, and it coexists with the Kigi started as an unofficial Kimi Code CLI — a community fork of
official `kimi` CLI on the same machine: independent binary name, [xai-org/grok-build](https://github.com/xai-org/grok-build), not affiliated
independent config directory (`~/.kigi`), independent keyring credentials with Moonshot AI or xAI. It keeps its own binary, its own `~/.kigi`, its own
(service `kigi`), and a `KIGI_*` environment-variable namespace. Nothing keyring entry, and its own `KIGI_*` env vars, and never touches what the
the official client installs or stores is ever read at runtime or written. official `kimi` CLI installed. `kigi import-kimi` copies your old config over
On first launch Kigi offers a **one-time, strictly read-only** import of once, read-only.
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 is **zero-telemetry**: the only outbound connections are the **Zero telemetry.** It talks to the APIs you configured, GitHub Releases, and
inference/auth APIs you configure, GitHub Releases for updates, and MCP your own MCP servers. Nothing else.
servers you add.
## License ## License
crates/codegen/kigi-shell/src/agent/auth_method.rs
+158 -8
View File
@@ -302,14 +302,26 @@ impl ModelByok {
/// buffered token on every turn and 401s with `bad-credentials` until restart. /// buffered token on every turn and 401s with `bad-credentials` until restart.
/// It refreshes when `endpoint_is_first_party` — the request targets the /// It refreshes when `endpoint_is_first_party` — the request targets the
/// first-party API, where sending the session token cannot leak to a /// first-party API, where sending the session token cannot leak to a
/// third-party BYOK endpoint. A definite `NotByok` always refreshes (it only /// third-party BYOK endpoint. A definite `NotByok` refreshes (within the
/// ever routes to the session endpoint); a definite `Byok` never does. /// 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( pub fn session_token_auth_gate(
is_session_based_method: bool, is_session_based_method: bool,
model_byok: ModelByok, model_byok: ModelByok,
endpoint_is_first_party: bool, endpoint_is_first_party: bool,
endpoint_takes_session_credential: bool,
) -> bool { ) -> bool {
is_session_based_method is_session_based_method
&& endpoint_takes_session_credential
&& match model_byok { && match model_byok {
ModelByok::NotByok => true, ModelByok::NotByok => true,
ModelByok::Byok => false, 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<kigi_models::PlatformId>,
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 = pub const AUTH_ERROR_SESSION_EXPIRED: &str =
"Session expired. Run `kigi login` to re-authenticate."; "Session expired. Run `kigi login` to re-authenticate.";
@@ -683,15 +731,117 @@ mod tests {
#[test] #[test]
fn session_token_auth_gate_matrix() { fn session_token_auth_gate_matrix() {
// Session method + NotByok → refresh. // Session method + NotByok → refresh (endpoint takes the session cred).
assert!(session_token_auth_gate(true, ModelByok::NotByok, false)); assert!(session_token_auth_gate(
true,
ModelByok::NotByok,
false,
true
));
// Session method + Byok → never. // 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. // Session method + Unknown → only on first-party endpoints.
assert!(session_token_auth_gate(true, ModelByok::Unknown, true)); assert!(session_token_auth_gate(
assert!(!session_token_auth_gate(true, ModelByok::Unknown, false)); true,
ModelByok::Unknown,
true,
true
));
assert!(!session_token_auth_gate(
true,
ModelByok::Unknown,
false,
true
));
// Non-session method → never. // 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). /// RAII guard restoring an env var on drop (panic-safe).
crates/codegen/kigi-shell/src/agent/config.rs
+58 -22
View File
@@ -2779,13 +2779,20 @@ pub fn default_model_entries(endpoints: &EndpointsConfig) -> IndexMap<String, Mo
} }
/// Resolve a model against the available model map. /// Resolve a model against the available model map.
/// Checks the map key (id) first, then falls back to a slug scan. /// Checks the map key (id) first, then falls back to a slug scan.
///
/// The slug scan takes the LAST match, exactly like the picker's
/// [`crate::agent::models::resolve_catalog_key`]. Duplicate slugs across
/// platforms are by design (`moonshot-cn` and `moonshot-ai` list the same ids;
/// so do every API-key platform and its subscription-OAuth twin), and a
/// first-match scan here made the auth layer resolve a DIFFERENT entry than the
/// one the picker selected — the H5 slug collision. One direction, one answer.
pub fn find_model_by_id<'a>( pub fn find_model_by_id<'a>(
models: &'a IndexMap<String, ModelEntry>, models: &'a IndexMap<String, ModelEntry>,
model_id: &str, model_id: &str,
) -> Option<&'a ModelEntry> { ) -> Option<&'a ModelEntry> {
models models
.get(model_id) .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: /// Whether the EFFECTIVE Auto-mode classifier model supports reasoning effort:
/// the model actually routed to (`aux_model` when the aux sampler resolved) else /// the model actually routed to (`aux_model` when the aux sampler resolved) else
@@ -4643,14 +4650,22 @@ reasoning_effort = "low"
}); });
(dir, manager) (dir, manager)
} }
/// LEAK 1a (aux/summary model): the aux `session_key` is resolved by the aux /// LEAK 1a (aux/summary model `api_key` channel): the aux `session_key` is
/// model's OWN platform (as `build_summary_client` / `resolve_aux_sampler_config` /// resolved by the aux model's OWN platform AND endpoint (as
/// now do). A grok (oauth-platform) aux model's resolved sampler `api_key` is /// `build_summary_client` / `resolve_aux_sampler_config` now do), so it is
/// therefore grok's own token or `None` — NEVER the primary Kimi key — while a /// never the primary Kimi key on a host that does not own it.
/// first-party / non-oauth aux model still gets the primary (byte-identical). ///
/// - 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] #[tokio::test]
async fn aux_model_session_key_is_platform_scoped_never_leaking_kimi() { 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 (_kd, kimi) = kimi_primary("kimi-tok");
let endpoints = EndpointsConfig::default(); let endpoints = EndpointsConfig::default();
// A grok aux catalog entry (managed id → oauth platform), no own key. // 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()); grok.info.id = Some("xai-grok/grok-4-latest".to_string());
let mut grok_catalog = IndexMap::new(); let mut grok_catalog = IndexMap::new();
grok_catalog.insert("grok".to_string(), grok); grok_catalog.insert("grok".to_string(), grok);
let grok_key = crate::auth::oauth_registry::session_key_for_model( let grok_key = crate::auth::oauth_registry::session_key_for_catalog_model(
home.path(), &grok_catalog,
"xai-grok/grok-4-latest", "grok",
Some(&kimi), Some(&kimi),
); );
let grok_cfg = resolve_aux_model_sampling_config( let grok_cfg = resolve_aux_model_sampling_config(
@@ -4675,10 +4690,11 @@ reasoning_effort = "low"
Some("kimi-tok"), Some("kimi-tok"),
"a grok aux model must never receive the primary Kimi session token", "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( let mut k2 = test_model_entry(
"kimi-k2-0905-preview", "kimi-k2-0905-preview",
"https://vendor/v1", "https://api.moonshot.cn/v1",
None, None,
None, None,
None, None,
@@ -4686,23 +4702,43 @@ reasoning_effort = "low"
k2.info.id = Some("moonshot-cn/kimi-k2".to_string()); k2.info.id = Some("moonshot-cn/kimi-k2".to_string());
let mut k2_catalog = IndexMap::new(); let mut k2_catalog = IndexMap::new();
k2_catalog.insert("k2".to_string(), k2); k2_catalog.insert("k2".to_string(), k2);
let k2_key = crate::auth::oauth_registry::session_key_for_model( assert_eq!(
home.path(), crate::auth::oauth_registry::session_key_for_catalog_model(
"moonshot-cn/kimi-k2", &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), Some(&kimi),
); );
let k2_cfg = resolve_aux_model_sampling_config( let kimi_cfg = resolve_aux_model_sampling_config(
"k2", "kimi-for-coding",
&k2_catalog, &kimi_catalog,
&endpoints, &endpoints,
k2_key.as_deref(), kimi_key.as_deref(),
None, 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!( assert_eq!(
k2_cfg.api_key.as_deref(), kimi_cfg.api_key.as_deref(),
Some("kimi-tok"), 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] #[test]
crates/codegen/kigi-shell/src/agent/models.rs
+91
View File
@@ -1917,6 +1917,51 @@ pub(crate) fn resolve_catalog_key(
.map(|(key, _)| acp::ModelId::new(key.clone())) .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<String, ModelEntry>,
current_key: Option<&str>,
slug: &str,
) -> Option<String> {
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<String, ModelEntry>,
current_key: Option<&str>,
slug: &str,
) -> Option<kigi_models::PlatformId> {
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** /// Catalog key for a persisted session model id, restricted to **selectable**
/// entries. A selectable exact-key match wins (as in [`resolve_catalog_key`]); /// entries. A selectable exact-key match wins (as in [`resolve_catalog_key`]);
/// otherwise the last selectable entry whose routing slug matches `id`, so a /// otherwise the last selectable entry whose routing slug matches `id`, so a
@@ -3668,6 +3713,52 @@ mod tests {
assert!(!info.visible_for_auth(false)); 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) ── // ── duplicate model slug re-keying (A/B experiment "auto" alias) ──
fn make_entry_config(model: &str, name: Option<&str>) -> config::ModelEntryConfig { fn make_entry_config(model: &str, name: Option<&str>) -> config::ModelEntryConfig {
crates/codegen/kigi-shell/src/agent/models_fetch.rs
+51 -11
View File
@@ -645,9 +645,15 @@ pub(crate) fn platform_wire_model_to_entry(
inference_idle_timeout_secs: None, inference_idle_timeout_secs: None,
max_retries: None, max_retries: None,
hidden: false, hidden: false,
// Subscription models require the OAuth session; open-platform // `supported_in_api: false` hides a model unless the PRIMARY session is
// models are usable by API-key users. // an OAuth session (`ModelInfo::visible_for_auth`). Only `kimi-code`
supported_in_api: !platform.uses_oauth(), // 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, supports_backend_search: false,
compactions_remaining: None, compactions_remaining: None,
compaction_at_tokens: None, compaction_at_tokens: None,
@@ -898,6 +904,16 @@ fn get_string_map(
mod tests { mod tests {
use super::*; 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 + /// OpenAI-cycle e2e (mock wire): a polluted bare-id `/models` listing +
/// a models.dev refresh produce a catalog with ONLY chat models, enriched /// a models.dev refresh produce a catalog with ONLY chat models, enriched
/// context windows / efforts, and the Responses backend — the full /// context windows / efforts, and the Responses backend — the full
@@ -1277,8 +1293,14 @@ mod tests {
"enrichment fills the context window from models.dev anthropic" "enrichment fills the context window from models.dev anthropic"
); );
assert!( assert!(
!opus.supported_in_api, opus.supported_in_api,
"subscription (uses_oauth) models require the OAuth session" "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" "context window comes from models.dev github-copilot enrichment"
); );
assert!( assert!(
!entry.supported_in_api, entry.supported_in_api,
"subscription (uses_oauth) models require the OAuth session" "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.context_window.get(), 272_000);
assert_eq!(sol.name.as_deref(), Some("GPT-5.6-Sol")); assert_eq!(sol.name.as_deref(), Some("GPT-5.6-Sol"));
assert!( assert!(
!sol.supported_in_api, sol.supported_in_api,
"subscription (uses_oauth) models require the OAuth session" "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!(sol.supports_reasoning_effort);
assert_eq!( assert_eq!(
@@ -3563,8 +3597,14 @@ mod tests {
"an OAuth channel carries no api-key env" "an OAuth channel carries no api-key env"
); );
assert!( assert!(
!entry.supported_in_api, entry.supported_in_api,
"subscription models require the OAuth session (not the public 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). // Passthrough dialect (identical to the API-key xai wire).
let model_entry = crate::agent::config::ModelEntry::from_config_entry(entry); let model_entry = crate::agent::config::ModelEntry::from_config_entry(entry);
crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs
+56 -24
View File
@@ -25,17 +25,19 @@ impl MvpAgent {
primary: &SamplingConfig, primary: &SamplingConfig,
) -> Result<(OaiCompatClient, String), acp::Error> { ) -> Result<(OaiCompatClient, String), acp::Error> {
let slug = self.resolve_session_summary_model(); let slug = self.resolve_session_summary_model();
// Resolve the aux token by the summary model's OWN platform: a grok let models = self.models_manager.models();
// (oauth-platform) summary model draws its pooled grok token or `None` // Resolve the aux token by the summary model's OWN platform AND
// — NEVER the primary Kimi session token (which `resolve_credentials` // endpoint: a grok (oauth-platform) summary model draws its pooled grok
// would otherwise stamp onto an api.x.ai request). A first-party / // token or `None`, and an API-key registry platform draws NOTHING —
// non-oauth summary model still gets the primary (byte-identical). // NEVER the primary Kimi session token (which `resolve_credentials`
let session_key = crate::auth::oauth_registry::session_key_for_model( // would otherwise stamp onto an api.x.ai / api.deepseek.com request).
&crate::util::kigi_home::kigi_home(), // 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, &slug,
Some(&self.auth_manager), Some(&self.auth_manager),
); );
let models = self.models_manager.models();
let endpoints = self.models_manager.endpoints(); let endpoints = self.models_manager.endpoints();
let alpha_test_key = self.cfg.borrow().endpoints.alpha_test_key.clone(); let alpha_test_key = self.cfg.borrow().endpoints.alpha_test_key.clone();
let config = match crate::agent::config::resolve_aux_model_sampling_config( let config = match crate::agent::config::resolve_aux_model_sampling_config(
@@ -47,7 +49,17 @@ impl MvpAgent {
) { ) {
Some(mut cfg) => { Some(mut cfg) => {
cfg.attribution_callback = primary.attribution_callback.clone(); 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.max_retries = primary.max_retries;
cfg cfg
} }
@@ -669,31 +681,51 @@ impl MvpAgent {
); );
Ok(entry.clone()) Ok(entry.clone())
} }
/// Resolve the SESSION token for `model` by the model's OWN platform — the /// Resolve the SESSION token for `model` by the model's OWN platform AND
/// single guard against the api_key-channel token leak. /// endpoint — the single guard against the api_key-channel token leak.
/// ///
/// An oauth-platform model (xai-grok) draws its session token from ITS OWN /// - an oauth-platform model (xai-grok, claude-pro-max, github-copilot,
/// process-global pool manager (build-on-demand from the on-disk grok token, /// openai-codex) draws its session token from ITS OWN process-global pool
/// proactively refreshed), INDEPENDENT of the primary `auth_method`; when /// manager (built on demand from the on-disk token, proactively
/// that provider has no stored session the token is `None` — NEVER the /// refreshed), INDEPENDENT of the primary `auth_method`; when that
/// primary Kimi key. Every other model (first-party / Kimi) uses the primary /// provider has no stored session the token is `None` — NEVER the primary
/// session manager, and only under a session-based auth method — /// Kimi key;
/// byte-identical to the pre-fix path. SECURITY: the resolved token is never /// - an endpoint that does not take a session credential at all
/// logged. /// ([`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 = <Kimi bearer>` 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<crate::auth::KimiAuth> { fn session_token_for_model(&self, model: &ModelEntry) -> Option<crate::auth::KimiAuth> {
if let Some(oauth) = model let info = model.info();
.info() let platform = info
.id .id
.as_deref() .as_deref()
.and_then(kigi_models::parse_managed_model_key) .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( return crate::auth::oauth_registry::global_manager_for(
&crate::util::kigi_home::kigi_home(), &crate::auth::oauth_registry::pool_home(),
oauth, oauth,
) )
.current_or_expired(); .current_or_expired();
} }
if !crate::agent::auth_method::platform_takes_session_credential(
platform,
&info.base_url,
) {
return None;
}
if self.is_session_based_auth() { if self.is_session_based_auth() {
self.auth_manager.current_or_expired() self.auth_manager.current_or_expired()
} else { } else {
crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs
+16 -5
View File
@@ -1197,18 +1197,26 @@ async fn prepare_sampling_config_never_stamps_kimi_key_on_grok_model() {
let endpoints = EndpointsConfig::default(); let endpoints = EndpointsConfig::default();
// First-party Kimi model (non-oauth platform): the primary session key IS // First-party SUBSCRIPTION model (kimi-code): the primary session key IS its
// its api_key — the byte-identical primary path, and proof the Kimi token is // 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 // live (so it WOULD leak if mis-routed onto a grok request). This assertion
// also confirms the session-based primary path is active. // 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()); assert!(!kimi_model.has_own_credentials());
let kimi_cfg = agent.prepare_sampling_config_for_model(&kimi_model, None); let kimi_cfg = agent.prepare_sampling_config_for_model(&kimi_model, None);
assert_eq!( assert_eq!(
kimi_cfg.api_key.as_deref(), kimi_cfg.api_key.as_deref(),
Some(KIMI_KEY), 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 // 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; 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 /// No load in flight and no session → the wait returns immediately
/// (the caller then surfaces "unknown session id" exactly as before). /// (the caller then surfaces "unknown session id" exactly as before).
#[tokio::test] #[tokio::test]
crates/codegen/kigi-shell/src/agent/mvp_agent/tests/api_key_channel_leak_tests.rs
+313
View File
@@ -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 = <Kimi token>`
//! with the THIRD-PARTY `base_url`, and `SamplingClient` builds
//! `Authorization: Bearer <api_key>` 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 <Kimi OAuth token>` 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"
);
}
crates/codegen/kigi-shell/src/agent/subagent/mod.rs
+16 -10
View File
@@ -1003,16 +1003,22 @@ fn resolve_model_override_to_config(
} else { } else {
acp::ModelId::new(entry.info().model.clone()) acp::ModelId::new(entry.info().model.clone())
}; };
// Resolve the child's session token by the OVERRIDE model's OWN platform, // 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 // AND endpoint, not the parent's primary auth: a grok (oauth-platform)
// pooled grok token or `None` — NEVER the primary Kimi session token (which // override draws its pooled grok token or `None`, and an API-key registry
// `resolve_credentials` would otherwise stamp onto the child's api.x.ai // platform / a third-party `[model.*]` host draws NOTHING — NEVER the
// credentials, leaking it in the logout-mid-session edge). A first-party / // primary Kimi session token, which `resolve_credentials` would otherwise
// non-oauth override still resolves to the primary (byte-identical). // stamp onto the child's api.x.ai / api.moonshot.cn credentials. The
let managed_key = entry.info().id.as_deref().unwrap_or(model_id); // first-party subscription channel still resolves to the primary
let session_key = crate::auth::oauth_registry::session_key_for_model( // (byte-identical).
&crate::util::kigi_home::kigi_home(), let session_key = crate::auth::oauth_registry::session_key_for_endpoint(
managed_key, entry
.info()
.id
.as_deref()
.and_then(kigi_models::parse_managed_model_key)
.map(|(platform, _)| platform),
&entry.info().base_url,
Some(&ctx.auth_manager), Some(&ctx.auth_manager),
); );
let has_session_key = session_key.is_some(); let has_session_key = session_key.is_some();
crates/codegen/kigi-shell/src/agent/subagent/tests/mod.rs
+13 -7
View File
@@ -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 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] #[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 models = indexmap::IndexMap::new();
let mut unavailable_first = test_model_entry("shared-routing-slug"); models.insert("visible-first".to_string(), test_model_entry("shared-routing-slug"));
unavailable_first.info.user_selectable = false; let mut unavailable_last = test_model_entry("shared-routing-slug");
models.insert("blocked-first".to_string(), unavailable_first); unavailable_last.info.user_selectable = false;
models.insert("visible-second".to_string(), test_model_entry("shared-routing-slug")); models.insert("blocked-last".to_string(), unavailable_last);
assert_eq!( assert_eq!(
super::handle_request::task_model_override_error(Some("shared-routing-slug"), super::handle_request::task_model_override_error(Some("shared-routing-slug"),
ModelOverrideProvenance::Tool, false, & models, false,).as_deref(), ModelOverrideProvenance::Tool, false, & models, false,).as_deref(),
Some("Unknown Task.model slug 'shared-routing-slug'. Valid model slugs: \ Some("Unknown Task.model slug 'shared-routing-slug'. Valid model slugs: \
visible-second. Omit `model` to inherit the parent model."), visible-first. Omit `model` to inherit the parent model."),
"validation must inspect the first routing-slug entry selected by execution" "validation must inspect the last routing-slug entry selected by execution"
); );
} }
#[test] #[test]
crates/codegen/kigi-shell/src/agent/subagent/tests/rest.rs
+42 -8
View File
@@ -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", "a grok override must never receive the primary Kimi session token",
); );
} }
/// Byte-identical guard: a non-oauth override with a Kimi primary still resolves /// Byte-identical guard: an override on the SESSION's own first-party endpoint
/// to the primary session token — passes both before and after the fix (the /// (the kimi-code subscription channel) still resolves to the primary session
/// non-oauth path is unchanged). /// token — the primary path is unchanged.
#[tokio::test] #[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 (_kd, manager) = kimi_primary_with_token("kimi-secret");
let mut entry = test_model_entry("kimi-k2-0905-preview"); let mut entry = test_model_entry("kimi-k2-0905-preview");
entry.info.id = Some("moonshot-cn/kimi-k2".to_string()); 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(); let mut models = indexmap::IndexMap::new();
models.insert("k2".to_string(), entry); models.insert("k2".to_string(), entry);
let mut ctx = ctx_with_toggle(HashMap::new()); 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() ..crate::auth::KimiAuth::test_default()
}); });
ctx.auth_manager = manager; ctx.auth_manager = manager;
let (config, _model_id) = let (config, _model_id) = resolve_model_override_to_config("k2", &ctx)
resolve_model_override_to_config("k2", &ctx).expect("non-oauth override resolves to a config"); .expect("an API-key-platform override still resolves to a config");
assert_eq!( assert_ne!(
config.api_key.as_deref(), config.api_key.as_deref(),
Some("kimi-secret"), 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 /// An unresolvable `AgentDefinition.model` pin (model absent from
crates/codegen/kigi-shell/src/auth/oauth_registry.rs
+153 -34
View File
@@ -13,11 +13,10 @@
//! generic-oauth scope, each wired with the SAME lifecycle as the primary Kimi //! generic-oauth scope, each wired with the SAME lifecycle as the primary Kimi
//! manager (`configure_refresher()` + `start_proactive_refresh()`) so the //! manager (`configure_refresher()` + `start_proactive_refresh()`) so the
//! on-disk token stays fresh and a 401 recovers via the provider's own manager. //! 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 //! Managers are built ON DEMAND from the on-disk token ([`global_manager_for`]),
//! on-disk token via [`global_manager_for`], so a login that lands AFTER a //! so a login landing AFTER a session spawned self-heals — no frozen per-session
//! session spawned self-heals — there is no frozen per-session snapshot to go //! snapshot. [`manager_for_model`] routes a managed catalog key to the pool
//! stale. [`manager_for_model`] routes a managed catalog key to the pool (oauth //! (oauth platform) or to the session's primary (everything else).
//! platform) or to the session's primary (everything else).
//! //!
//! SECURITY: access/refresh tokens and resolved bearers are NEVER logged here. //! SECURITY: access/refresh tokens and resolved bearers are NEVER logged here.
@@ -39,6 +38,31 @@ fn oauth_manager_pool() -> &'static Mutex<HashMap<&'static str, Arc<AuthManager>
POOL.get_or_init(|| Mutex::new(HashMap::new())) 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<tempfile::TempDir> = 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 /// Get-or-create the process-global manager for `oauth`, wiring the same
/// refresher + proactive-refresh lifecycle as the primary Kimi manager the /// refresher + proactive-refresh lifecycle as the primary Kimi manager the
/// FIRST time a scope is seen. The manager reads the on-disk token at /// 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() primary.cloned()
} }
/// The SESSION token (the raw bearer/key string) that governs INFERENCE auth /// The SESSION token (the raw bearer/key string) that may ride an INFERENCE
/// for `managed_key`, resolved by the model's OWN platform. Thin wrapper over /// request routed to `platform` at `base_url`. Used by the aux-model, summary
/// [`manager_for_model`] used by the aux-model and subagent-override wire paths /// and subagent-override wire paths, where the result is stamped straight into
/// so a `{platform}/{model}` key never receives the primary token of a /// [`crate::agent::config::resolve_credentials`] as the request's `api_key`.
/// DIFFERENT provider.
/// ///
/// A generic device-code OAuth platform (xai-grok) draws its token from ITS OWN /// - a generic device-code OAuth platform (xai-grok, claude-pro-max,
/// pooled manager; when that provider has no stored session the result is /// github-copilot, openai-codex) draws from ITS OWN pooled manager; when that
/// `None` — NEVER the primary Kimi key. Every other key routes to `primary` and /// provider has no stored session the result is `None` — never `primary`;
/// yields the primary's current-or-expired token, byte-identical to reading it /// - `kimi-code`, and a platform-less model whose endpoint IS the session's own
/// directly. SECURITY: the resolved token is never logged. /// coding endpoint (incl. a `KIGI_CODE_BASE_URL` deployment or a loopback
pub(crate) fn session_key_for_model( /// proxy), yield the primary's current-or-expired token — byte-identical to
kigi_home: &Path, /// reading it directly;
managed_key: &str, /// - 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<kigi_models::PlatformId>,
base_url: &str,
primary: Option<&Arc<AuthManager>>, primary: Option<&Arc<AuthManager>>,
) -> Option<String> { ) -> Option<String> {
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()) .and_then(|am| am.current_or_expired())
.map(|a| a.key) .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<String, crate::agent::config::ModelEntry>,
slug: &str,
primary: Option<&Arc<AuthManager>>,
) -> Option<String> {
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)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
@@ -159,6 +226,20 @@ mod tests {
.expect("openai-codex carries an OAuthConfig") .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<AuthManager>>,
) -> Option<String> {
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/<model>` turn resolves to the process-global pooled /// An `openai-codex/<model>` turn resolves to the process-global pooled
/// openai-codex manager (its OWN `oauth/openai-codex` scope), NEVER the /// openai-codex manager (its OWN `oauth/openai-codex` scope), NEVER the
/// primary Kimi manager — the same leak-safe routing as the other OAuth /// 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" "openai-codex and claude-pro-max must not share a pooled manager"
); );
assert_ne!( 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()), Some("kimi-tok".to_string()),
"an openai-codex model must never receive the primary Kimi token" "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 // 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). // Kimi bearer — it draws from the copilot pool (its own token, or None).
assert_ne!( 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()), Some("kimi-tok".to_string()),
"a github-copilot model must never receive the primary Kimi token" "a github-copilot model must never receive the primary Kimi token"
); );
@@ -256,9 +337,8 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn session_key_for_claude_pro_max_is_never_the_kimi_primary() { async fn session_key_for_claude_pro_max_is_never_the_kimi_primary() {
let (_kd, kimi) = primary_with_token("kimi-tok"); let (_kd, kimi) = primary_with_token("kimi-tok");
let home = tempfile::tempdir().unwrap();
assert_ne!( 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()), Some("kimi-tok".to_string()),
"a claude-pro-max model must never receive the primary Kimi session token" "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 /// `session_key_for_endpoint`: the endpoints that genuinely ride the
/// token exactly as reading it directly would — byte-identical to the /// PRIMARY session — `kimi-code` (the subscription channel) and a
/// pre-fix aux/override wire path (no runtime / pool touched). /// 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] #[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 (_kd, kimi) = primary_with_token("kimi-tok");
let home = tempfile::tempdir().unwrap(); assert_eq!(
for key in ["moonshot-cn/kimi-k2", "kimi-k2-0905-preview"] { 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!( assert_eq!(
session_key_for_model(home.path(), key, Some(&kimi)), session_key_for_endpoint(None, url, Some(&kimi)),
Some("kimi-tok".to_string()), 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 /// 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 /// 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 /// pooled xai manager (its own token, or `None`). This is the exact source
@@ -365,15 +485,14 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn session_key_for_grok_is_never_the_kimi_primary() { async fn session_key_for_grok_is_never_the_kimi_primary() {
let (_kd, kimi) = primary_with_token("kimi-tok"); let (_kd, kimi) = primary_with_token("kimi-tok");
let home = tempfile::tempdir().unwrap();
assert_ne!( 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()), Some("kimi-tok".to_string()),
"a grok aux/override model must never receive the primary Kimi session token" "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. // Even with `None` primary the routing is unchanged: grok → pool, never a panic.
assert_ne!( 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()), Some("kimi-tok".to_string()),
); );
} }
crates/codegen/kigi-shell/src/session/acp_session.rs
+12 -1
View File
@@ -135,7 +135,7 @@ use prompt_build::*;
mod session_mode; mod session_mode;
use session_mode::*; use session_mode::*;
#[path = "acp_session_impl/sampler_turn.rs"] #[path = "acp_session_impl/sampler_turn.rs"]
mod sampler_turn; pub(crate) mod sampler_turn;
use sampler_turn::*; use sampler_turn::*;
#[path = "acp_session_impl/tool_dispatch.rs"] #[path = "acp_session_impl/tool_dispatch.rs"]
mod tool_dispatch; mod tool_dispatch;
@@ -1256,6 +1256,17 @@ mod rewind_synthetic_turn_tests;
#[cfg(test)] #[cfg(test)]
#[path = "acp_session_tests/rewrite_zero_turn_prefix_tests.rs"] #[path = "acp_session_tests/rewrite_zero_turn_prefix_tests.rs"]
mod rewrite_zero_turn_prefix_tests; 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. /// Pins the `SubagentFinished` usage-fold attribution gate.
#[cfg(test)] #[cfg(test)]
#[path = "acp_session_tests/subagent_usage_fold_tests.rs"] #[path = "acp_session_tests/subagent_usage_fold_tests.rs"]
crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs
+5 -7
View File
@@ -634,14 +634,12 @@ impl SessionActor {
&active_session_config, &active_session_config,
Some(self.max_retries), Some(self.max_retries),
); );
// A grok (oauth-platform) image-describe model must not inherit the // An image-describe model on another provider must not inherit the
// session (Kimi) bearer_resolver stamped by `finalize_*`; re-point it at // 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 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 { if aux_resolved {
self.repoint_aux_bearer_resolver_for_oauth( self.repoint_aux_bearer_resolver(&mut sampler_config, &self.image_description_model);
&mut sampler_config,
&self.image_description_model,
);
} }
let client = kigi_sampler::SamplingClient::new(sampler_config).map_err(|e| { let client = kigi_sampler::SamplingClient::new(sampler_config).map_err(|e| {
acp::Error::internal_error().data(format!( acp::Error::internal_error().data(format!(
crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs
+132 -57
View File
@@ -36,20 +36,37 @@ struct SessionTokenAuthGate {
/// BYOK status still refresh against the first-party cli-chat-proxy hosts without /// BYOK status still refresh against the first-party cli-chat-proxy hosts without
/// risking a session-token leak to a third-party BYOK endpoint. /// risking a session-token leak to a third-party BYOK endpoint.
endpoint_is_first_party: bool, 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 { impl SessionTokenAuthGate {
/// Single place `is_session_based` / `endpoint_is_first_party` are derived, /// 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( fn new(
auth_method_id: Option<&acp::AuthMethodId>, auth_method_id: Option<&acp::AuthMethodId>,
model_byok: crate::agent::auth_method::ModelByok, model_byok: crate::agent::auth_method::ModelByok,
base_url: &str, base_url: &str,
model_platform: Option<kigi_models::PlatformId>,
) -> Self { ) -> Self {
Self { Self {
is_session_based: auth_method_id is_session_based: auth_method_id
.is_some_and(crate::agent::auth_method::is_session_based_method), .is_some_and(crate::agent::auth_method::is_session_based_method),
model_byok, model_byok,
endpoint_is_first_party: crate::util::is_first_party_url(base_url), 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 { fn active(self) -> bool {
@@ -57,6 +74,7 @@ impl SessionTokenAuthGate {
self.is_session_based, self.is_session_based,
self.model_byok, self.model_byok,
self.endpoint_is_first_party, self.endpoint_is_first_party,
self.endpoint_takes_session_credential,
) )
} }
} }
@@ -127,6 +145,43 @@ fn auth_manager_bearer_resolver(
) -> kigi_sampler::SharedBearerResolver { ) -> kigi_sampler::SharedBearerResolver {
std::sync::Arc::new(AuthManagerBearerResolver(am)) 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<kigi_sampler::SharedBearerResolver>,
platform: Option<kigi_models::PlatformId>,
base_url: &str,
) -> Option<kigi_sampler::SharedBearerResolver> {
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 { impl SessionActor {
pub(super) async fn prepare_tool_definitions_timed(&self) -> (Vec<ToolDefinition>, u64) { pub(super) async fn prepare_tool_definitions_timed(&self) -> (Vec<ToolDefinition>, u64) {
let mcp_wait_start = std::time::Instant::now(); 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 { fn auth_gate(&self, model_id: &str, base_url: &str) -> SessionTokenAuthGate {
let byok = self.model_auth_facts(model_id).byok; let byok = self.model_auth_facts(model_id).byok;
let auth_method = self.auth_method_id.load(); 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 /// The [`AuthManager`](crate::auth::AuthManager) that governs INFERENCE auth
/// for the model whose routing slug is `model` (the sampling config's /// for the model whose routing slug is `model` (the sampling config's
@@ -238,26 +298,30 @@ impl SessionActor {
// and resolves to the primary. // and resolves to the primary.
let managed_key = self.managed_key_for_model(model); let managed_key = self.managed_key_for_model(model);
crate::auth::oauth_registry::manager_for_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), managed_key.as_deref().unwrap_or(model),
self.auth_manager.as_ref(), self.auth_manager.as_ref(),
) )
} }
/// Recover the managed catalog key (`{platform}/{model}`) for a routing slug /// Recover the managed catalog key (`{platform}/{model}`) for a routing slug
/// from the live catalog. `None` for a bare / config / unlisted model. /// 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<String> { fn managed_key_for_model(&self, model: &str) -> Option<String> {
let models = self.models_manager.models(); 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 /// The registry platform the routing slug `model` belongs to, from the SAME
/// OAuth platform (xai-grok). The gate for re-pointing an aux model's /// lookup [`Self::auth_manager_for_model`] routes on. `None` for a bare /
/// bearer_resolver away from the session (Kimi) resolver — a first-party / /// `[model.*]` / unlisted model.
/// non-oauth model returns `false` and keeps the stamped session resolver. fn model_platform(&self, model: &str) -> Option<kigi_models::PlatformId> {
fn model_is_oauth_platform(&self, model: &str) -> bool { let models = self.models_manager.models();
let managed_key = self.managed_key_for_model(model); let current = self.models_manager.current_model_id();
kigi_models::parse_managed_model_key(managed_key.as_deref().unwrap_or(model)) crate::agent::models::platform_for_slug(&models, Some(current.0.as_ref()), model)
.and_then(|(platform, _)| platform.oauth())
.is_some()
} }
/// Whether `model` routes to the Claude Pro/Max OAuth-Messages platform /// Whether `model` routes to the Claude Pro/Max OAuth-Messages platform
/// (claude-pro-max) — the gate for the sampler's OAuth Messages adaptation /// (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 /// which is ChatCompletions) returns `false`, keeping the API-key Anthropic
/// / MiniMax Messages requests byte-identical. /// / MiniMax Messages requests byte-identical.
fn model_is_anthropic_oauth(&self, model: &str) -> bool { fn model_is_anthropic_oauth(&self, model: &str) -> bool {
let managed_key = self.managed_key_for_model(model); self.model_platform(model).is_some_and(|platform| {
kigi_models::parse_managed_model_key(managed_key.as_deref().unwrap_or(model)).is_some_and( platform.oauth().is_some()
|(platform, _)| { && platform.wire_api() == kigi_models::PlatformWireApi::Messages
platform.oauth().is_some() })
&& platform.wire_api() == kigi_models::PlatformWireApi::Messages
},
)
} }
/// Whether `model` routes to the GitHub Copilot ChatCompletions platform /// Whether `model` routes to the GitHub Copilot ChatCompletions platform
/// (github-copilot) — the gate for the sampler's editor-identity headers + /// (github-copilot) — the gate for the sampler's editor-identity headers +
/// `X-Initiator`. Every other model returns `false`, keeping the other /// `X-Initiator`. Every other model returns `false`, keeping the other
/// ChatCompletions providers byte-identical. /// ChatCompletions providers byte-identical.
fn model_is_github_copilot(&self, model: &str) -> bool { fn model_is_github_copilot(&self, model: &str) -> bool {
let managed_key = self.managed_key_for_model(model); self.model_platform(model)
kigi_models::parse_managed_model_key(managed_key.as_deref().unwrap_or(model)) .is_some_and(kigi_models::PlatformId::sends_copilot_editor_headers)
.is_some_and(|(platform, _)| platform.sends_copilot_editor_headers())
} }
/// Whether `model` routes to the ChatGPT/Codex Responses platform /// Whether `model` routes to the ChatGPT/Codex Responses platform
/// (openai-codex) — the gate for the sampler's Codex identity headers /// (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 /// returns `false`, keeping the API-key `openai` Responses request
/// byte-identical. /// byte-identical.
fn model_is_openai_codex(&self, model: &str) -> bool { fn model_is_openai_codex(&self, model: &str) -> bool {
let managed_key = self.managed_key_for_model(model); self.model_platform(model)
kigi_models::parse_managed_model_key(managed_key.as_deref().unwrap_or(model)) .is_some_and(kigi_models::PlatformId::sends_codex_responses_headers)
.is_some_and(|(platform, _)| platform.sends_codex_responses_headers())
} }
/// LEAK guard for the stamped aux paths (auto-mode classifier, image /// LEAK guard for the stamped aux paths (auto-mode classifier, image
/// describe). After [`crate::agent::config::stamp_session_local_sampler_fields`] /// describe). Applies [`aux_bearer_resolver`] to the config
/// has copied the SESSION model's `bearer_resolver` onto an aux /// [`crate::agent::config::stamp_session_local_sampler_fields`] just stamped
/// `SamplerConfig`, re-point it at the AUX model's OWN platform manager when /// the SESSION model's `bearer_resolver` onto.
/// the aux model is an oauth platform (xai-grok) — so a grok aux model never pub(super) fn repoint_aux_bearer_resolver(
/// 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(
&self, &self,
cfg: &mut kigi_sampler::SamplerConfig, cfg: &mut kigi_sampler::SamplerConfig,
slug: &str, slug: &str,
) { ) {
if self.model_is_oauth_platform(slug) cfg.bearer_resolver = aux_bearer_resolver(
&& let Some(manager) = self.auth_manager_for_model(slug) cfg.bearer_resolver.take(),
{ self.model_platform(slug),
cfg.bearer_resolver = Some(auth_manager_bearer_resolver(manager)); &cfg.base_url,
} );
} }
/// Emit a unified-log breadcrumb whenever the session-token refresh gate is /// Emit a unified-log breadcrumb whenever the session-token refresh gate is
/// evaluated with an **`Unknown`** per-model BYOK status on a session-based /// evaluated with an **`Unknown`** per-model BYOK status on a session-based
@@ -330,8 +385,9 @@ impl SessionActor {
let ctx = serde_json::json!( let ctx = serde_json::json!(
{ "site" : site, "model_byok" : gate.model_byok.as_str(), "is_session_based" { "site" : site, "model_byok" : gate.model_byok.as_str(), "is_session_based"
: gate.is_session_based, "endpoint_is_first_party" : gate : gate.is_session_based, "endpoint_is_first_party" : gate
.endpoint_is_first_party, "refresh_active" : refresh_active, "base_url" : .endpoint_is_first_party, "endpoint_takes_session_credential" : gate
base_url, } .endpoint_takes_session_credential, "refresh_active" : refresh_active,
"base_url" : base_url, }
); );
let sid = Some(self.session_info.id.0.as_ref()); let sid = Some(self.session_info.id.0.as_ref());
if refresh_active { if refresh_active {
@@ -385,8 +441,12 @@ impl SessionActor {
let creds = self.chat_state_handle.get_credentials().await; let creds = self.chat_state_handle.get_credentials().await;
let model_facts = self.model_auth_facts(cfg.model.as_str()); let model_facts = self.model_auth_facts(cfg.model.as_str());
let auth_method = self.auth_method_id.load(); let auth_method = self.auth_method_id.load();
let gate = let gate = SessionTokenAuthGate::new(
SessionTokenAuthGate::new(auth_method.as_deref(), model_facts.byok, &cfg.base_url); auth_method.as_deref(),
model_facts.byok,
&cfg.base_url,
self.model_platform(cfg.model.as_str()),
);
let use_bearer_resolver = gate.active(); let use_bearer_resolver = gate.active();
self.log_auth_gate_unknown("reconstruct_full_config", gate, &cfg.base_url); 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 // Resolve the bearer from the ACTIVE model's OWN manager: a grok model
@@ -589,17 +649,18 @@ impl SessionActor {
slug: &str, slug: &str,
) -> Option<kigi_sampler::SamplerConfig> { ) -> Option<kigi_sampler::SamplerConfig> {
let creds = self.chat_state_handle.get_credentials().await; let creds = self.chat_state_handle.get_credentials().await;
// Resolve the aux token by the aux model's OWN platform: a grok let models = self.models_manager.models();
// (oauth-platform) aux model draws its pooled grok token or `None` — // Resolve the aux token by the aux model's OWN platform AND endpoint: a
// NEVER the primary Kimi session token (which `resolve_credentials` // grok (oauth-platform) aux model draws its pooled grok token or `None`,
// would otherwise stamp onto an api.x.ai request). A first-party / // and an API-key registry platform draws NOTHING — NEVER the primary
// non-oauth aux model still gets the primary (byte-identical). // Kimi session token (which `resolve_credentials` would otherwise stamp
let session_key = crate::auth::oauth_registry::session_key_for_model( // onto an api.x.ai / api.deepseek.com request). The first-party
&crate::util::kigi_home::kigi_home(), // subscription channel still gets the primary (byte-identical).
let session_key = crate::auth::oauth_registry::session_key_for_catalog_model(
&models,
slug, slug,
self.auth_manager.as_ref(), self.auth_manager.as_ref(),
); );
let models = self.models_manager.models();
let endpoints = self.models_manager.endpoints(); let endpoints = self.models_manager.endpoints();
crate::agent::config::resolve_aux_model_sampling_config( crate::agent::config::resolve_aux_model_sampling_config(
slug, slug,
@@ -625,10 +686,11 @@ impl SessionActor {
&active_session_config, &active_session_config,
Some(self.max_retries), Some(self.max_retries),
); );
// LEAK 1b: a grok aux classifier must not inherit the SESSION model's // LEAK 1b: the aux classifier must not inherit the SESSION model's
// (Kimi) bearer_resolver stamped above; re-point it at grok's own // (Kimi) bearer_resolver stamped above — re-point it at an OAuth aux
// manager (its pooled token, or None). No-op for a non-oauth aux. // model's own manager, or clear it for an API-key-platform / third-party
self.repoint_aux_bearer_resolver_for_oauth(&mut cfg, slug); // aux endpoint. No-op for the first-party subscription channel.
self.repoint_aux_bearer_resolver(&mut cfg, slug);
let model = cfg.model.clone(); let model = cfg.model.clone();
let client = kigi_sampler::SamplingClient::new(cfg) let client = kigi_sampler::SamplingClient::new(cfg)
.map_err(|e| { .map_err(|e| {
@@ -774,6 +836,8 @@ impl SessionActor {
session_id = % self.session_info.id.0, is_session_based = gate session_id = % self.session_info.id.0, is_session_based = gate
.is_session_based, model_byok = gate.model_byok.as_str(), .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,
"auth recovery: sampler 401 not refreshable (api-key auth) — surfacing 401", "auth recovery: sampler 401 not refreshable (api-key auth) — surfacing 401",
); );
kigi_log::unified_log::warn( kigi_log::unified_log::warn(
@@ -783,7 +847,9 @@ impl SessionActor {
{ "kind" : error.kind.as_str(), "status_code" : error { "kind" : error.kind.as_str(), "status_code" : error
.status_code, "is_session_based" : gate.is_session_based, .status_code, "is_session_based" : gate.is_session_based,
"model_byok" : gate.model_byok.as_str(), "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) .map(|c| c.model)
.unwrap_or_default(); .unwrap_or_default();
let Some(ref key) = current_key else { return }; 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(&current_model_id).is_some() {
return;
}
let Some(new_key) = self.reload_api_key_from_config(&current_model_id) else { let Some(new_key) = self.reload_api_key_from_config(&current_model_id) else {
return; return;
}; };
@@ -1161,7 +1236,7 @@ mod bearer_resolver_tests {
.oauth() .oauth()
.expect("xai-grok carries an OAuthConfig"); .expect("xai-grok carries an OAuthConfig");
let grok = crate::auth::oauth_registry::global_manager_for( let grok = crate::auth::oauth_registry::global_manager_for(
&crate::util::kigi_home::kigi_home(), &crate::auth::oauth_registry::pool_home(),
oauth, oauth,
); );
assert_ne!( assert_ne!(
crates/codegen/kigi-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs
+16 -10
View File
@@ -487,21 +487,27 @@ fn session_token_auth_gate_truth_table() {
use crate::agent::auth_method::{ModelByok, session_token_auth_gate as gate}; use crate::agent::auth_method::{ModelByok, session_token_auth_gate as gate};
// Non-session methods never refresh, regardless of BYOK status or endpoint. // Non-session methods never refresh, regardless of BYOK status or endpoint.
for fp in [false, true] { for fp in [false, true] {
assert!(!gate(false, ModelByok::NotByok, fp)); assert!(!gate(false, ModelByok::NotByok, fp, true));
assert!(!gate(false, ModelByok::Byok, fp)); assert!(!gate(false, ModelByok::Byok, fp, true));
assert!(!gate(false, ModelByok::Unknown, fp)); assert!(!gate(false, ModelByok::Unknown, fp, true));
// Session method: a definite classification ignores the endpoint — // Session method on an endpoint that DOES take the session credential
// NotByok always refreshes (only ever routes to the session endpoint), // (kimi-code, an OAuth platform's own pool, or a bare / [model.*]
// a genuine per-model Byok never does. // model): a definite classification ignores the endpoint — NotByok
assert!(gate(true, ModelByok::NotByok, fp)); // refreshes, a genuine per-model Byok never does.
assert!(!gate(true, ModelByok::Byok, fp)); 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 // Session method + Unknown BYOK: refresh only against a first-party xAI
// host, so a transiently-unclassifiable config can't demote a live session // 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 // (the stale-token 401 regression) yet the session token never leaks to a
// third-party BYOK endpoint. This arm was unconditionally `false` pre-fix. // third-party BYOK endpoint. This arm was unconditionally `false` pre-fix.
assert!(gate(true, ModelByok::Unknown, true)); assert!(gate(true, ModelByok::Unknown, true, true));
assert!(!gate(true, ModelByok::Unknown, false)); assert!(!gate(true, ModelByok::Unknown, false, true));
} }
/// Pre-fix, the gate read `auth_type` and skipped recovery here, 401'ing every /// Pre-fix, the gate read `auth_type` and skipped recovery here, 401'ing every
crates/codegen/kigi-shell/src/session/acp_session_tests/session_bearer_leak_platform_tests.rs
+289
View File
@@ -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<String> {
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"
);
}
});
}
crates/codegen/kigi-shell/src/session/acp_session_tests/session_bearer_leak_tests.rs
+436
View File
@@ -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<AuthManager>) {
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<SessionActor>,
mpsc::UnboundedReceiver<PersistenceMsg>,
) {
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<SessionActor>,
mpsc::UnboundedReceiver<PersistenceMsg>,
) {
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 <KIMI_TOKEN>` 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;
}