Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d6e49bcc7d | ||
|
|
d2358b037c | ||
|
|
6ba24db019 | ||
|
|
10149f50dd | ||
|
|
e53a66d113 | ||
|
|
f6eaafa3da | ||
|
|
13ccab980c | ||
|
|
b9b7e6c989 | ||
|
|
f403c38a94 | ||
|
|
27d009cb6e | ||
|
|
815bd99356 | ||
|
|
5189bc5e28 | ||
|
|
0e4356c34e | ||
|
|
63ff114f71 | ||
|
|
66d5194ee3 | ||
|
|
1bb10ef7f1 | ||
|
|
0e352c9914 | ||
|
|
fc7c2a1b9b | ||
|
|
1e57cd9225 | ||
|
|
c2d1067f0d | ||
|
|
48d89c7830 | ||
|
|
2d00a4e6e6 | ||
|
|
422e241e13 | ||
|
|
8179438278 | ||
|
|
5a9183b08b | ||
|
|
dbce6bf305 | ||
|
|
8a26460251 | ||
|
|
bc4e76db96 | ||
|
|
347311bfe5 | ||
|
|
010f6c3be3 | ||
|
|
c02b4b1ed7 | ||
|
|
8245deb373 | ||
|
|
ebf11057f8 | ||
|
|
193d16f6d5 | ||
|
|
9ee40b13d0 | ||
|
|
f825132983 | ||
|
|
0a147bd8a5 | ||
|
|
e7dbb98207 | ||
|
|
7a2cd8a726 | ||
|
|
e373ff8b29 | ||
|
|
9953a26b8d | ||
|
|
49e6414c29 | ||
|
|
7efb4b07cc | ||
|
|
b86722f508 | ||
|
|
23e94939c0 | ||
|
|
fdf9b956f5 | ||
|
|
83e6935189 | ||
|
|
c5ddaec71e | ||
|
|
99d99fb47a | ||
|
|
c4fe0750e2 | ||
|
|
d6f216facd | ||
|
|
4d2cf7ffd4 | ||
|
|
02cf5deebd | ||
|
|
db05ed0751 | ||
|
|
771de1822d | ||
|
|
bb5cbff62d | ||
|
|
1579558b56 | ||
|
|
4d1e4fdc52 | ||
|
|
4361b03259 | ||
|
|
2c29d3ecd5 | ||
|
|
de534582c1 | ||
|
|
54c9177ce2 |
@@ -1,79 +0,0 @@
|
||||
name: CI
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
pull_request:
|
||||
|
||||
env:
|
||||
CARGO_TERM_COLOR: always
|
||||
|
||||
jobs:
|
||||
gates:
|
||||
name: check / clippy / fmt / deny (${{ matrix.os }})
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
os: [macos-14, ubuntu-24.04]
|
||||
runs-on: ${{ matrix.os }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Install toolchain (rust-toolchain.toml)
|
||||
run: rustup show
|
||||
- name: Install dotslash (protoc launcher)
|
||||
run: cargo install dotslash --locked
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
- name: cargo fmt
|
||||
run: cargo fmt --all --check
|
||||
- name: cargo check
|
||||
run: cargo check --workspace --all-targets --locked
|
||||
- name: cargo clippy
|
||||
run: cargo clippy --workspace --all-targets --locked -- -D warnings
|
||||
- name: Install cargo-deny
|
||||
run: cargo install cargo-deny --locked
|
||||
- name: cargo deny advisories
|
||||
run: cargo deny check advisories
|
||||
|
||||
test:
|
||||
name: test (${{ matrix.os }})
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
os: [macos-14, ubuntu-24.04]
|
||||
runs-on: ${{ matrix.os }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Install toolchain (rust-toolchain.toml)
|
||||
run: rustup show
|
||||
- name: Install dotslash (protoc launcher)
|
||||
run: cargo install dotslash --locked
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
- name: cargo test
|
||||
run: cargo test --workspace --locked
|
||||
|
||||
perf:
|
||||
name: performance budgets (${{ matrix.os }})
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
os: [macos-14, ubuntu-24.04]
|
||||
runs-on: ${{ matrix.os }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Install toolchain (rust-toolchain.toml)
|
||||
run: rustup show
|
||||
- name: Install dotslash (protoc launcher)
|
||||
run: cargo install dotslash --locked
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
- name: Install hyperfine (macOS)
|
||||
if: runner.os == 'macOS'
|
||||
run: brew install hyperfine
|
||||
- name: Install hyperfine (Linux)
|
||||
if: runner.os == 'Linux'
|
||||
run: sudo apt-get update && sudo apt-get install -y hyperfine
|
||||
- name: Build release binaries
|
||||
run: |
|
||||
cargo build --release -p kigi-bin --locked
|
||||
cargo build --release -p kigi-pager-pty-harness --bin pty-scenario --locked
|
||||
- name: Enforce performance budgets
|
||||
run: scripts/bench.sh target/release/kigi
|
||||
@@ -0,0 +1,85 @@
|
||||
name: Warm build cache
|
||||
|
||||
# Release builds run on TAG refs, and GitHub Actions cache isolation only
|
||||
# lets a run restore caches created on its OWN ref or the DEFAULT branch.
|
||||
# No workflow ran on `main`, so every release compiled the whole workspace
|
||||
# cold on all five targets (~50 min wall clock, gated by Windows). This
|
||||
# workflow builds the same `release-dist` profile on `main` so tag builds
|
||||
# restore a warm default-branch cache.
|
||||
#
|
||||
# Triggers: dependency-affecting pushes to main (each release's version-bump
|
||||
# commit warms the cache for the NEXT release), a weekly refresh so the
|
||||
# cache never hits GitHub's 7-day unused-eviction, and manual dispatch.
|
||||
#
|
||||
# The setup steps mirror release.yml's build job (toolchain, target,
|
||||
# dotslash/protoc, rust-cache key) — keep them in lockstep, or the cache
|
||||
# key won't match and releases go back to cold builds.
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: ["main"]
|
||||
paths:
|
||||
- "Cargo.lock"
|
||||
- "Cargo.toml"
|
||||
- "rust-toolchain.toml"
|
||||
- ".github/workflows/warm-cache.yml"
|
||||
schedule:
|
||||
- cron: "17 5 * * 1"
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
env:
|
||||
CARGO_TERM_COLOR: always
|
||||
|
||||
concurrency:
|
||||
group: warm-cache
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
warm:
|
||||
name: warm (${{ matrix.target }})
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
include:
|
||||
- target: aarch64-apple-darwin
|
||||
os: macos-14
|
||||
- target: x86_64-apple-darwin
|
||||
os: macos-14
|
||||
- target: x86_64-unknown-linux-gnu
|
||||
os: ubuntu-24.04
|
||||
- target: aarch64-unknown-linux-gnu
|
||||
os: ubuntu-24.04-arm
|
||||
- target: x86_64-pc-windows-msvc
|
||||
os: windows-2022
|
||||
runs-on: ${{ matrix.os }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install toolchain (rust-toolchain.toml)
|
||||
run: rustup show
|
||||
|
||||
- name: Add build target
|
||||
run: rustup target add ${{ matrix.target }}
|
||||
|
||||
- name: Install dotslash (protoc launcher)
|
||||
run: cargo install dotslash --locked
|
||||
|
||||
- name: Install protoc (Windows PATH fallback)
|
||||
if: runner.os == 'Windows'
|
||||
shell: pwsh
|
||||
run: |
|
||||
$url = "https://github.com/protocolbuffers/protobuf/releases/download/v29.3/protoc-29.3-win64.zip"
|
||||
Invoke-WebRequest -Uri $url -OutFile protoc.zip
|
||||
Expand-Archive protoc.zip -DestinationPath "$env:USERPROFILE\protoc"
|
||||
Add-Content $env:GITHUB_PATH "$env:USERPROFILE\protoc\bin"
|
||||
|
||||
# Same key as release.yml so tag builds restore this cache verbatim.
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
key: ${{ matrix.target }}
|
||||
|
||||
- name: Build kigi (release-dist)
|
||||
run: cargo build --profile release-dist -p kigi-bin --locked --target ${{ matrix.target }}
|
||||
@@ -18,9 +18,12 @@ import) or any `KIMI_*` env var.
|
||||
|
||||
- **Zero egress**: outbound connections are limited to
|
||||
`auth.kimi.com`, `api.kimi.com`, `api.moonshot.cn`, `api.moonshot.ai`,
|
||||
GitHub Releases domains, and user-configured MCP servers. No telemetry,
|
||||
no analytics, ever. `crates/codegen/kigi-env` is the single home of
|
||||
first-party endpoints.
|
||||
GitHub Releases domains, user-configured MCP servers, the endpoints of
|
||||
provider platforms the user has credentialed, and `models.dev` (model
|
||||
metadata refresh — reached ONLY when an enabled platform's `/models` wire
|
||||
lacks metadata, `wire_serves_metadata=false`; Kimi/Moonshot never trigger
|
||||
it; `KIGI_MODELS_DEV_URL=0` disables). No telemetry, no analytics, ever.
|
||||
`crates/codegen/kigi-env` is the single home of first-party endpoints.
|
||||
- **Toolchain**: Rust 1.97.0 (rust-toolchain.toml), edition 2024.
|
||||
- **Gates** (all must stay green):
|
||||
`cargo check --workspace --all-targets`,
|
||||
@@ -29,9 +32,19 @@ import) or any `KIMI_*` env var.
|
||||
- **Observability is local**: `kigi-log` (unified session log, `--debug`
|
||||
firehose, subsystem file logs, opt-in instrumentation) writes under
|
||||
`~/.kigi` only. Its zero-network property is a contract.
|
||||
- **Atomic file replace goes through `util::fs::replace_file`** (tmp+rename
|
||||
commit step; async callers wrap in `spawn_blocking`). Never inline a bare
|
||||
`fs::rename` replace: Windows `MoveFileExW(REPLACE_EXISTING)` fails with a
|
||||
sharing violation while AV/indexer/cloud-sync holds the destination open —
|
||||
the "persists on macOS, silently doesn't on Windows" class (a /model
|
||||
switch that never stuck). Plain rename stays correct only for true moves
|
||||
whose destination doesn't pre-exist (worktree-pool markers, corrupt-file
|
||||
backups). Write failures must at least `warn!` — never `let _ =`.
|
||||
- The root `Cargo.toml` is hand-maintained (upstream's generator is not in
|
||||
this repo). Members sorted; versions inherited from
|
||||
`workspace.package.version` (0.1.0).
|
||||
`workspace.package.version` — the single source of truth for the release
|
||||
version (`kigi_version::VERSION` derives from it; the release workflow
|
||||
gates the `v*` tag against it).
|
||||
|
||||
## Layout
|
||||
|
||||
@@ -49,6 +62,17 @@ import) or any `KIMI_*` env var.
|
||||
- `third_party/` — vendored Mermaid rendering stack (untouched policy).
|
||||
- `bin/protoc` — dotslash launcher used by proto codegen.
|
||||
|
||||
## Storage discipline
|
||||
|
||||
- Tests that touch the filesystem MUST use `tempfile::TempDir` (drop
|
||||
cleans up) — never bare `std::env::temp_dir()` + `create_dir_all`,
|
||||
which leaks directories into the OS temp root forever.
|
||||
- `target/` grows past 150GB across repeated full-workspace builds
|
||||
(incremental is already off); run `cargo clean` when it exceeds
|
||||
~50GB and at milestone boundaries.
|
||||
- Graph node worktrees are removed right after a successful merge-back;
|
||||
only FAILED nodes keep theirs for postmortem.
|
||||
|
||||
## Test seams
|
||||
|
||||
Cross-crate test hooks are behind the `test-support` cargo feature
|
||||
@@ -56,6 +80,332 @@ Cross-crate test hooks are behind the `test-support` cargo feature
|
||||
dependents' `[dev-dependencies]`. Don't expose new test seams as plain
|
||||
`#[cfg(test)]` items across crate boundaries.
|
||||
|
||||
## Graph mode (`/graph`, post-0.1.x — plan.md in the parent dir)
|
||||
|
||||
A deterministic DAG scheduler layered over the goal engine: `/graph
|
||||
<objective>` decomposes the objective into nodes (graph planner subagent
|
||||
→ Agentproof-style static gate in `graph_plan.rs`), then executes each
|
||||
node as one ordinary goal — the agentic loop lives INSIDE the node; the
|
||||
edges stay deterministic Rust. The harness appends a terminal
|
||||
`gn-final` verification node depending on every planner node.
|
||||
|
||||
- Enabled by default (`KIGI_GRAPH=0` is the off-switch; the G0 gray
|
||||
release is over); availability additionally requires the goal harness
|
||||
(`BuiltinGate::Graph`).
|
||||
- Key modules (kigi-shell): `session/graph_tracker.rs` (pure state
|
||||
machine; reuses `GoalStatus`/`GoalPhase`/`GoalPauseReason`),
|
||||
`session/graph_plan.rs` (planner-JSON contract + validation + fnv id
|
||||
canonicalization), `session/graph_planner.rs` (planner runner, reuses
|
||||
the goal planner spawn plumbing),
|
||||
`session/acp_session_impl/graph.rs` (orchestration seam).
|
||||
- Seam points: `handle_prompt` intercepts GraphSet/GraphResume; the
|
||||
in-turn loop's `EndTurn` arm calls `run_graph_round_end()` to advance
|
||||
nodes within the same turn; goal auto-pauses cascade to the graph in
|
||||
`auto_pause_goal_if_active_inner`; node goals are armed with the
|
||||
REMAINING graph budget so `enforce_goal_token_budget` cascades trips.
|
||||
- Persistence: `PersistenceMsg::GraphModeState(Option<..>)` →
|
||||
`<session_dir>/graph/state.json` (`None` tombstones after clear);
|
||||
immutable per-version baselines `graph/graph.baseline.v{N}.json`;
|
||||
per-node goal artifacts archived to `graph/<node_id>/`. Restore
|
||||
demotes `Active`→`UserPaused` and `Running`→`Ready` (re-run is safe:
|
||||
the verifier gates completion).
|
||||
- `/goal` and `/graph` are mutually exclusive while the graph owns the
|
||||
engine; e2e suite: `acp_session_tests/graph/graph_e2e_tests.rs`.
|
||||
- Parallel fan-out (G1): with `KIGI_GRAPH_CONCURRENCY > 1` (default 3,
|
||||
clamp [1,8]) and ≥2 `Ready` nodes, `drive_graph` runs batches via
|
||||
`acp_session_impl/graph_workers.rs` — per node a bounded
|
||||
worker↔verifier subagent loop (`KIGI_GRAPH_NODE_ROUNDS`, default 3;
|
||||
`general-purpose` children; worktree isolation on round 1, resume
|
||||
keeps context+worktree on later rounds; `NODE_RESULT:` /
|
||||
`NODE_VERDICT:` terminal contracts parsed fail-closed), then
|
||||
SEQUENTIAL merge-back via `kigi_workspace::worktree::apply_worktree`
|
||||
(`ApplyMode::Merge`); a conflict fails the node and blocks its
|
||||
dependents while other chains continue. `gn-final` always runs
|
||||
serially on the full goal engine. Concurrency=1 is byte-identical to
|
||||
the serial G0 path. Ceiling: a worker round exceeding the foreground
|
||||
subagent await budget (600s) is cancelled and retried via resume.
|
||||
- G2: `BudgetLimited` is resumable — a budget trip demotes in-flight
|
||||
nodes to `Ready` (resource stop, not a verdict) and
|
||||
`/graph resume --budget <tokens>` re-arms with fresh headroom. The
|
||||
pager shows a graph status chip driven by the `GraphUpdated` wire
|
||||
variant (`extensions/notification.rs`), emitted from the single
|
||||
`persist_graph_state` chokepoint (checkpoint ⇔ badge tick); old
|
||||
pagers degrade via `#[serde(other)] Unknown`. PTY scenarios:
|
||||
`graph_slash_presession{,_disabled}.yaml`.
|
||||
- G3: dynamic replan. `DISCOVERED: <text>` line markers (fence-stripped,
|
||||
placeholder-filtered) from workers/verifiers/the serial node's final
|
||||
text queue as `pending_discoveries`; at each dispatch boundary
|
||||
`maybe_replan_graph` (`acp_session_impl/graph_replan.rs`) runs a
|
||||
replanner subagent producing an APPEND-ONLY appendix
|
||||
(`validate_replan`: existing-id deps allowed, edges onto `gn-final`
|
||||
rejected — they would cycle after the final-gating extension),
|
||||
bumps `plan_version`, freezes `graph.baseline.v{N}.json`, and regates
|
||||
`gn-final` (Ready→Waiting). Bounded by `KIGI_GRAPH_REPLAN_CAP`
|
||||
(default 3, 0 = off); past the cap — and after the final node has
|
||||
achieved — discoveries drain to history only. Replan failure degrades
|
||||
(history + notice); it never pauses a working graph.
|
||||
- G4: the graph follows the repo. Every checkpoint projects to
|
||||
`.kigi/graph.jsonl` at the git root (`session/graph_project.rs`,
|
||||
header line + one node per line, atomic write); single writer via an
|
||||
fs2 flock sidecar; other instances get read-only `/graph status`.
|
||||
Fresh sessions revive via `/graph resume` (load UNDER the lock,
|
||||
from_snapshot demotions apply). All lock-then-mutate sites
|
||||
identity-check the projected `graph_id`; kigi never commits the file.
|
||||
- G5: `/graph show` renders box-drawing DAG art
|
||||
(`session/graph_render.rs`, Sugiyama-lite: longest-path layers, dummy
|
||||
pass-throughs, barycenter ordering, bus lanes). Wider than 120 cols
|
||||
degrades to the status tree.
|
||||
- G6: plan-boundary topology optimizer
|
||||
(`acp_session_impl/graph_optimize.rs`; `KIGI_GRAPH_OPTIMIZER=0`
|
||||
disables). Restricted ops (`remove_dep`/`reorder`/`merge`/`split`)
|
||||
validated by `graph_plan::apply_optimization`: pending-only targets,
|
||||
immutable nodes byte-identical in the result, terminal gate rebuilt,
|
||||
whole-graph acyclicity. Applied passes bump `plan_version` and share
|
||||
the replan cap; `{"ops": []}` is a respected free no-op; failures
|
||||
degrade.
|
||||
|
||||
## Provider registry & API-key auth (post-0.1.3 expansion)
|
||||
|
||||
- The platform registry is compiled-in spec rows in `kigi-models`
|
||||
(`PlatformSpec`; adding a platform = enum variant + `ALL` entry + `spec()`
|
||||
arm + row; registry tests enforce completeness/uniqueness/row shape).
|
||||
- API-key resolution precedence, per platform: platform env var(s) >
|
||||
`auth.json` scope named by the platform id (`moonshot-cn`, …) >
|
||||
legacy `[platforms.<id>]` in config.toml (read-only fallback).
|
||||
- The TUI login picker persists pasted keys to `auth.json` (platform-id
|
||||
scope, `api_key` mode) — never to config.toml. The keyring holds ONLY the
|
||||
OAuth session scope; platform keys are file-only.
|
||||
- Auth method ids over ACP equal the platform ids; interactive picker rows
|
||||
are built generically from advertised methods (`AuthMethodKind::
|
||||
ApiKeyPlatform`), so new registry rows appear in the picker with no TUI
|
||||
changes.
|
||||
- Every OAuth picker row carries ITS OWN method id (`PendingMenuItem::Login
|
||||
{ method_id }` → `Action::LoginWith`); an unknown id fails closed. The
|
||||
id-less `Action::Login` (auto-login, 401 re-auth) resolves the first
|
||||
interactive method. `/login` opens the picker (`Action::OpenLoginPicker`),
|
||||
never a flow directly; mid-session its last row is Cancel, not Quit.
|
||||
- `_meta.connected` on an advertised method is DISPLAY state (green badge on
|
||||
the picker): stamped at `initialize()` from stored credentials
|
||||
(`connected_method_ids` + `stamp_connected_meta`), kept fresh TUI-side
|
||||
after in-session logins (`auth_in_flight_method` → `AuthComplete`). It is
|
||||
never an authorization input.
|
||||
- Refreshable-OAuth providers beyond Kimi Code use a GENERIC path, NOT Kimi's
|
||||
bespoke wire. A `uses_oauth` platform carrying `oauth: Some(&OAuthConfig)`
|
||||
(client id / auth host / start+token paths / `token_host` / `scope` /
|
||||
`scope_key` / optional extra device field / `flow` / `token_body`) drives a
|
||||
scope-keyed `AuthManager::new_oauth_provider` +
|
||||
`refresh::GenericDeviceRefresher` (selected by `build_refresher` via
|
||||
`oauth_config_for_scope_key`; the refresher dispatches the refresh body by
|
||||
`token_body`: form → `auth::oauth_device`, JSON → `auth::oauth_pkce`,
|
||||
`GithubCopilotExchange` → `auth::github_copilot` copilot-token re-mint). Kimi
|
||||
Code keeps `oauth: None` and its bespoke path unchanged. The interactive
|
||||
login is dispatched by `OAuthConfig.flow` (in `run_oauth_provider_flow`):
|
||||
- `OAuthFlow::DeviceCode` → `auth::oauth_device` (RFC-8628 device-code, plain
|
||||
kigi UA, no X-Msh headers). Provider: `xai-grok` (`scope_key oauth/xai`,
|
||||
base `api.x.ai/v1`, form token body, same wire as the API-key `xai` row).
|
||||
- `OAuthFlow::PkceLocalhost { redirect_port, redirect_path }` →
|
||||
`auth::oauth_pkce` (authorization-code + PKCE S256,
|
||||
`127.0.0.1:redirect_port{redirect_path}` loopback with STRICT `state`
|
||||
validation + manual-paste fallback, authorize host possibly ≠ token host).
|
||||
The `token_body` selects the login dialect: `Json` = claude (`state ==
|
||||
verifier`, JSON exchange carrying `state`), `Form` = codex (fresh-random
|
||||
`state`, FORM exchange without `state` via `exchange_code_form`).
|
||||
`OAuthConfig.authorize_extra` appends provider-only authorize params (empty
|
||||
for all but codex). Providers:
|
||||
- `claude-pro-max` (`scope_key oauth/claude-pro-max`, port 53692
|
||||
`/callback`, JSON body, base `api.anthropic.com/v1`, Anthropic Messages +
|
||||
listing wire reached with an OAuth `sk-ant-oat…` Bearer). Its Messages
|
||||
requests take the OAuth adaptation — `anthropic-beta claude-code-…,oauth-…`
|
||||
+ `claude-cli` UA + `x-app cli` + the required "You are Claude Code…"
|
||||
system prefix — gated on `SamplerConfig.anthropic_oauth` (claude-pro-max
|
||||
only), so API-key `anthropic`/`minimax` Messages requests stay
|
||||
byte-identical. Its `/v1/models` listing rides the same Bearer +
|
||||
oauth-beta headers. THINKING REPLAY (`prune_replayed_thinking`,
|
||||
all Messages requests): Anthropic validates every replayed
|
||||
`thinking` block (signature model-bound, non-empty required), so
|
||||
only the final assistant message's signed thinking is replayed and
|
||||
only while its tool loop is open (request ends on the tool
|
||||
results); everything else — unsigned cross-backend history, `tco_*`
|
||||
Responses blobs, stale-model blocks — is stripped, or the request
|
||||
400s "Invalid `signature` in `thinking` block".
|
||||
- `openai-codex` (ChatGPT Plus/Pro, `scope_key oauth/openai-codex`, port
|
||||
1455 `/auth/callback`, FORM body, authorize+token host `auth.openai.com`,
|
||||
client `app_EMoam…`, scope `openid profile email offline_access`, the 3
|
||||
authorize-extra params `id_token_add_organizations`/
|
||||
`codex_cli_simplified_flow`/`originator=codex_cli_rs`). Refresh is a plain
|
||||
`refresh_token` FORM grant (the generic refresher's `Form` path →
|
||||
`auth::oauth_device`). The minted `access_token` is a JWT; login FAILS
|
||||
FAST unless it carries the `["https://api.openai.com/auth"]
|
||||
["chatgpt_account_id"]` claim (`chatgpt_account_id_from_jwt`) — that
|
||||
account id is NOT persisted but re-derived STATELESSLY from the current
|
||||
bearer at every request. INFERENCE reuses the EXISTING Responses wire
|
||||
against base `chatgpt.com/backend-api/codex` (→ `{base}/responses`) with
|
||||
a Codex-gated adaptation (`SamplerConfig.openai_codex` /
|
||||
`PlatformId::sends_codex_responses_headers()`): headers
|
||||
`chatgpt-account-id` (per-request from the JWT), `originator codex_cli_rs`,
|
||||
`OpenAI-Beta responses=experimental`, a codex `User-Agent`; `store:false`
|
||||
is the shared Responses default. BODY adaptation
|
||||
(`adapt_body_for_codex_backend`, same gate): the backend 400s
|
||||
`role:system` input ("System messages are not allowed") — system items
|
||||
are hoisted into the top-level `instructions` field — and stateless
|
||||
reasoning replay requires `include:["reasoning.encrypted_content"]`.
|
||||
API-key `openai` Responses requests carry NONE of this
|
||||
(byte-identical, pinned by a control wire test). `reasoning.effort`
|
||||
carries the thinking level (incl. the codex-only `ultra`). NO
|
||||
websocket, NO base_instructions.
|
||||
CATALOG is HARDCODED (`PlatformId::hardcoded_catalog` →
|
||||
`openai_codex_wire_models`, mapped through the SAME
|
||||
`platform_wire_model_to_entry` output): exactly the 4 `visibility=list` &&
|
||||
`supported_in_api=true` models (`gpt-5.6-sol/terra/luna`, `gpt-5.5`, ctx
|
||||
272000, per-model efforts) — NO live `/models` fetch, NO codex-CLI /
|
||||
`~/.codex` dependency; `gpt-5.3-codex-spark` (api=false) and
|
||||
`gpt-5.4`/`gpt-5.4-mini`/`codex-auto-review` (hidden) are EXCLUDED.
|
||||
- `OAuthFlow::GithubDeviceCopilot` → `auth::github_copilot` (TWO-STAGE).
|
||||
Provider: `github-copilot` (`scope_key oauth/github-copilot`, base
|
||||
`api.individual.githubcopilot.com`, ChatCompletions wire). Stage 1 is an
|
||||
RFC-8628 device flow on `github.com` (client `Iv1.b507a08c87ecfe98`, scope
|
||||
`read:user`) whose errors ride a `200` body (not `4xx`) — it mints the
|
||||
DURABLE github token. Stage 2 (`GET api.github.com/copilot_internal/v2/token`
|
||||
with `copilot_exchange` + editor headers) re-mints the SHORT-LIVED copilot
|
||||
token. Persisted as `KimiAuth.key = copilot token`, `refresh_token = github
|
||||
token`, `expires_at = copilot expiry`; the "refresh" is a copilot-token
|
||||
RE-MINT (GET, not a `refresh_token` grant). Every `/models` listing AND
|
||||
`/chat/completions` request carries the VS Code editor-identity headers
|
||||
(`User-Agent GitHubCopilotChat/…`, `Editor-Version`, `Editor-Plugin-Version`,
|
||||
`Copilot-Integration-Id`; `+X-GitHub-Api-Version` on `/models`, `+X-Initiator
|
||||
user` on inference) — gated on `SamplerConfig.github_copilot` /
|
||||
`PlatformId::sends_copilot_editor_headers()` so every other ChatCompletions
|
||||
provider stays byte-identical. WIRE-COMPAT SCOPE: Kigi is one-wire-per-
|
||||
platform, so the catalog is FILTERED (`parse_github_copilot_listing`) to the
|
||||
openai-completions-served models — keep iff `model_picker_enabled` &&
|
||||
`policy.state != "disabled"` && `tool_calls != false` AND the id is NOT a
|
||||
`claude-(haiku|sonnet|opus)-[45]` (anthropic-messages) or `gpt-5/oswe/mai-`
|
||||
(responses-only) model. Those excluded models need per-model wire routing
|
||||
(deferred, documented debt), NOT included lest they fail at inference.
|
||||
KNOWN LIMITATION: Kigi does NOT port Pi's per-model policy-acceptance step
|
||||
(`POST {base}/models/{id}/policy {state:"enabled"}`). A kept model whose
|
||||
Copilot policy is unconfigured can list yet `403` at inference until the user
|
||||
enables it once in GitHub's UI — a deliberate omission (it mutates account
|
||||
state and is unverifiable without a live Copilot account), not a silent gap.
|
||||
These are INTERACTIVE login rows advertised right after `kimi-code`
|
||||
(`AuthMethodKind::OAuthPlatform`, in `PlatformId::ALL` order: `xai-grok`,
|
||||
`claude-pro-max`, `github-copilot`, `openai-codex`). The catalog fetch resolves each such platform's OWN session
|
||||
token (`resolve_generic_oauth_tokens`, refreshed on expiry) and routes
|
||||
`platform.oauth().is_some()` → `platform.base_url()` (kimi-code alone →
|
||||
`proxy_url()`). Tokens/codes/verifiers are NEVER logged.
|
||||
- A stored subscription-OAuth session IS a catalog fetch source. Every
|
||||
fetch-PLAN decision that cannot afford real token resolution — the startup
|
||||
prefetch arming gate, `on_auth_changed`'s wipe guard
|
||||
(`should_wipe_catalog_on_auth_change`), and cache-origin computation — uses
|
||||
the sync presence probes `models_fetch::stored_oauth_platforms` /
|
||||
`stored_oauth_token_stubs` (auth.json scope scan; names only, no bearers).
|
||||
All three must see the SAME enabled-platform set the real fetch enables, or
|
||||
a claude-pro-max-only user boots onto the bundled Kimi table with an empty
|
||||
picker. Managed catalog entries stamp `meta.provider` (platform display
|
||||
name) for the client's `/model` picker; the picker itself lists the fetched
|
||||
catalog, which is by construction the credentialed providers' models.
|
||||
- INFERENCE-AUTH CHOKEPOINT (security): `auth::credential_authority::CredentialAuthority`
|
||||
is the ONE authority answering "WHICH credential — if any — may ride this
|
||||
request?". It holds the session's effective `EndpointsConfig` and
|
||||
its primary `AuthManager` PRIVATELY and answers only
|
||||
`(platform, base_url)` questions: `credential_class` → `CredentialClass::{Pooled,
|
||||
Primary, None}` (the outer term of `auth_method::session_token_auth_gate`),
|
||||
`manager_for` (the governing manager
|
||||
for refresh / 401 recovery), `credential_for` (the request's `api_key`) and
|
||||
`bearer_resolver_for` (the aux/summary/session resolver). Do NOT re-derive this
|
||||
rule anywhere else — three rounds of leaks came from exactly that.
|
||||
The rule: a subscription-OAuth platform rides ITS OWN pooled `AuthManager` (so
|
||||
it keeps a live `bearer_resolver` and mid-session refresh despite a
|
||||
non-first-party base URL) and ONLY at its own `platform.base_url()` — a
|
||||
`[model."claude-pro-max/x"]` override keeps `info.id` but can point `base_url`
|
||||
anywhere. `kimi-code` and a platform-less model (a bare slug / `[model.*]`
|
||||
entry) ride the PRIMARY session, and ONLY at the session's own effective coding
|
||||
endpoint: `EndpointsConfig::proxy_url()` (which prefers `[endpoints]
|
||||
coding_api_base_url` from **config.toml** — what the managed-config sync writes
|
||||
— over `KIGI_CODE_BASE_URL`), `models_base_url`, loopback, or 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`. Every API-key registry platform rides NOTHING.
|
||||
STRUCTURAL ENFORCEMENT: the authority is the only producer of
|
||||
`SessionCredential`, an opaque type with no production constructor, and every
|
||||
API that stamps a session bearer onto a request (`resolve_credentials`,
|
||||
`resolve_aux_model_sampling_config`, `try_resolve_model_credentials`,
|
||||
`resolve_chat_state_auth_type`) takes
|
||||
`Option<&SessionCredential>` rather than `Option<&str>` — a new call site
|
||||
cannot express the leak. `stamp_session_local_sampler_fields` likewise takes the
|
||||
aux `bearer_resolver` explicitly instead of copying the session's and relying on
|
||||
the caller to re-point it, and `sampler_turn::aux_bearer_resolver_for` is the
|
||||
ONE definition of the aux/summary resolver rule (the session actor and
|
||||
`MvpAgent::build_summary_client` both call it; a private second copy is how the
|
||||
summary client stayed ungated after M3). NEVER HAND-CARRY A CREDENTIAL TO A
|
||||
GUARD (C1): "may **a** session credential ride here" is `true` for
|
||||
a subscription-OAuth platform at its own host — where the credential that may
|
||||
ride is that platform's POOLED token, never the primary. Ask
|
||||
`credential_for(platform, base_url)` and stamp what it returns, so the question
|
||||
and the credential are the same object; where you cannot (a credential the
|
||||
authority does not own), MATCH on `credential_class` rather than reach for a
|
||||
boolean. There is deliberately no second, similarly-named predicate to pick
|
||||
wrongly. The shared `sampling_config.api_key`
|
||||
(the subagent baseline and the unresolved-model fallback) has exactly TWO
|
||||
production writers, both guarded by the authority:
|
||||
`MvpAgent::stamp_session_credential` (the `cached_token` / `kimi.com/oidc`
|
||||
login handlers and the `new_session` / `load_session` seed), which asks
|
||||
`credential_for`; and the `xai.api_key` handler in `acp_agent.rs`, which stamps
|
||||
the house `KIGI_API_KEY` read from the environment — a credential the authority
|
||||
does not own — only when `credential_class` is `Primary`, the class the house
|
||||
key rides and the one every OAuth platform's own host is NOT.
|
||||
- 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, held PER SESSION in
|
||||
`SessionActor::selected_catalog_key` (seeded at spawn by
|
||||
`agent::models::selected_catalog_key_for_spawn`, rewritten by `SetSessionModel`,
|
||||
CLEARED by `OverrideModelName` when the rename makes it stale), via
|
||||
`agent::models::entry_for_slug`/`platform_for_slug`.
|
||||
NEVER `ModelsManager::current_model_id()`: that cell is process-global,
|
||||
last-writer-wins across concurrent sessions, and Leader mode never writes it at
|
||||
all (`agent/handlers/model_switch.rs`). The shared `MvpAgent::sampling_config`
|
||||
is BUILT from that cell — but ONCE, at startup, and never rebuilt, while the
|
||||
cell moves on every non-Leader switch. Its guards therefore read
|
||||
`MvpAgent::sampling_config_platform`, the platform captured WITH the config by
|
||||
the same `ModelsManager::sampling_config()` call, never a fresh lookup against
|
||||
the live cell: once the two drift, re-resolving the config's bare slug falls
|
||||
through to `resolve_catalog_key`'s `.rev()` scan, answers the API-key twin, and
|
||||
a post-expiry `kigi login` silently leaves the EXPIRED bearer in the config
|
||||
that seeds every subagent (H-a). REFUSE RATHER THAN GUESS (H-b):
|
||||
when the per-session key does not name the slug, `platform_for_slug` returns
|
||||
`None` for a slug that collides across platforms rather than trusting
|
||||
`resolve_catalog_key`'s `.rev()` last match — which is the subscription-OAuth
|
||||
twin, so the guess hands an API-key session the pooled bearer that REPLACES its
|
||||
own key on the wire. `None` then routes purely by the ENDPOINT, which for an
|
||||
OAuth host means no credential, no resolver and no adaptation.
|
||||
Anything else (aux models, subagent
|
||||
overrides) 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 — and hands the API-key twin's session a pooled
|
||||
OAuth bearer stamped over the user's own key.
|
||||
- CATALOG VISIBILITY: `platform_wire_model_to_entry` stamps
|
||||
`supported_in_api = platform != KimiCode`. `ModelInfo::visible_for_auth`
|
||||
reads only the PRIMARY manager's auth mode, so gating the other OAuth
|
||||
platforms on it would hide every model from a user who signed in with ONLY a
|
||||
Claude Pro/Max, ChatGPT, Copilot, or Grok subscription. Only `kimi-code`
|
||||
rides the primary session, so only it may be gated on it.
|
||||
- Model metadata (context window, thinking levels) comes from the provider
|
||||
wire when served; metadata-poor listings are enriched from models.dev
|
||||
(`kigi-models/src/enrichment.rs` — bundled raw snapshot regenerated by
|
||||
`scripts/gen_enrichment_snapshot.py`, single Rust transform
|
||||
`parse_api_json` for bundled + runtime refresh; 24h cache
|
||||
`~/.kigi/models_dev_cache.json`). Wire values always win; enrichment
|
||||
never invents model availability. Canonical reasoning efforts:
|
||||
none/minimal/low/medium/high/xhigh/max/ultra (`max` split from `xhigh`
|
||||
2026-07; `ultra` is codex-only, above `max`, surfaced only via a model's
|
||||
server-declared effort menu; Kimi wire spells its top tier `max`, kimi_compat
|
||||
renames).
|
||||
|
||||
## Milestones (PRD §8.3)
|
||||
|
||||
- M0 (done): rename, deletions (voice/telemetry/announcements/marketplace/
|
||||
|
||||
Generated
+64
-62
@@ -5442,7 +5442,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-acp-lib"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"agent-client-protocol",
|
||||
"async-trait",
|
||||
@@ -5456,7 +5456,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-agent"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"chrono",
|
||||
"dirs 6.0.0",
|
||||
@@ -5486,7 +5486,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-agent-lifecycle"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"tokio",
|
||||
@@ -5495,7 +5495,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-auth"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"http 1.4.2",
|
||||
@@ -5508,7 +5508,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-bin"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"clap",
|
||||
@@ -5543,7 +5543,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-chat-state"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"indexmap",
|
||||
"kigi-compaction",
|
||||
@@ -5560,7 +5560,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-codebase-graph"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"ahash",
|
||||
"clap",
|
||||
@@ -5596,7 +5596,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-compaction"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"async-trait",
|
||||
@@ -5609,7 +5609,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-config"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"base64",
|
||||
"blake3",
|
||||
@@ -5632,7 +5632,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-config-types"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"agent-client-protocol",
|
||||
"indexmap",
|
||||
@@ -5646,7 +5646,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-crash-handler"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"backtrace",
|
||||
"libc",
|
||||
@@ -5657,7 +5657,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-env"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"tracing",
|
||||
"url",
|
||||
@@ -5665,7 +5665,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-fast-worktree"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"bytes",
|
||||
@@ -5697,7 +5697,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-file-utils"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"aws-config",
|
||||
@@ -5721,7 +5721,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-fsnotify"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"criterion",
|
||||
"dunce",
|
||||
@@ -5742,7 +5742,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-gix-status"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"gix",
|
||||
"kigi-test-utils",
|
||||
@@ -5752,7 +5752,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-hooks"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"fastrand",
|
||||
"kigi-config",
|
||||
@@ -5771,7 +5771,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-hooks-plugins-types"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"serde_json",
|
||||
@@ -5779,7 +5779,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-http"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"kigi-auth",
|
||||
"kigi-log",
|
||||
@@ -5794,7 +5794,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-hunk-tracker"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"chrono",
|
||||
"dunce",
|
||||
@@ -5815,14 +5815,14 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-interjection-core"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "kigi-log"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"chrono",
|
||||
@@ -5840,7 +5840,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-markdown"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"anstyle",
|
||||
"anstyle-lossy",
|
||||
@@ -5864,14 +5864,14 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-markdown-core"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"pulldown-cmark",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "kigi-mcp"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"agent-client-protocol",
|
||||
"async-trait",
|
||||
@@ -5908,7 +5908,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-memory"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"arc-swap",
|
||||
@@ -5942,7 +5942,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-mermaid"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"fontdb",
|
||||
"image",
|
||||
@@ -5960,16 +5960,17 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-models"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"kigi-env",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "kigi-pager-minimal"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"chrono",
|
||||
"crossterm",
|
||||
@@ -5986,7 +5987,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-pager-pty-harness"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"alacritty_terminal",
|
||||
"anyhow",
|
||||
@@ -6011,7 +6012,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-pager-render"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"agent-client-protocol",
|
||||
"anstyle",
|
||||
@@ -6063,7 +6064,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-paths"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"camino",
|
||||
"serde",
|
||||
@@ -6073,7 +6074,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-prompt-queue"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"serde_json",
|
||||
@@ -6081,7 +6082,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-proto-build"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"pbjson-build",
|
||||
@@ -6092,7 +6093,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-ratatui-inline"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"ansi-width",
|
||||
"anstyle-parse 0.2.7",
|
||||
@@ -6109,7 +6110,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-ratatui-textarea"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"arboard",
|
||||
"chrono",
|
||||
@@ -6130,7 +6131,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-sampler"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"async-openai",
|
||||
"async-stream",
|
||||
@@ -6153,10 +6154,11 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-sampling-types"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"assert_matches",
|
||||
"async-openai",
|
||||
"base64",
|
||||
"indexmap",
|
||||
"kigi-compaction",
|
||||
"kigi-tools",
|
||||
@@ -6169,7 +6171,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-sandbox"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"chrono",
|
||||
@@ -6190,7 +6192,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-secrets"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"regex",
|
||||
"serde_json",
|
||||
@@ -6228,7 +6230,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-shell"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"agent-client-protocol",
|
||||
"anyhow",
|
||||
@@ -6365,7 +6367,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-shell-base"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"chrono",
|
||||
@@ -6390,7 +6392,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-sqlite-journal"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"rusqlite",
|
||||
@@ -6401,7 +6403,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-subagent-resolution"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"kigi-sampling-types",
|
||||
"kigi-tool-types",
|
||||
@@ -6416,7 +6418,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-system-power"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"windows-sys 0.59.0",
|
||||
"zbus",
|
||||
@@ -6424,7 +6426,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-test-support"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"agent-client-protocol",
|
||||
"anyhow",
|
||||
@@ -6446,7 +6448,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-test-utils"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"runfiles",
|
||||
"tracing",
|
||||
@@ -6455,11 +6457,11 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-token-estimation"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
|
||||
[[package]]
|
||||
name = "kigi-tool-protocol"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"kigi-tool-types",
|
||||
"serde",
|
||||
@@ -6470,7 +6472,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-tool-runtime"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"async-trait",
|
||||
@@ -6488,7 +6490,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-tool-types"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"minijinja",
|
||||
"schemars 1.2.1",
|
||||
@@ -6498,7 +6500,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-tools"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"arc-swap",
|
||||
@@ -6575,7 +6577,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-tools-api"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"kigi-proto-build",
|
||||
"kigi-tool-protocol",
|
||||
@@ -6588,11 +6590,11 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-tracing-macros"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
|
||||
[[package]]
|
||||
name = "kigi-tty-utils"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"nix 0.30.1",
|
||||
@@ -6602,7 +6604,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-tui"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"agent-client-protocol",
|
||||
"ansi-to-tui",
|
||||
@@ -6689,7 +6691,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-update"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"dunce",
|
||||
@@ -6718,14 +6720,14 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-version"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"semver",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "kigi-workspace"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"agent-client-protocol",
|
||||
"anyhow",
|
||||
@@ -6804,7 +6806,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "kigi-workspace-types"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"base64",
|
||||
"chrono",
|
||||
@@ -8838,7 +8840,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "ptyctl"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"alacritty_terminal",
|
||||
"anyhow",
|
||||
@@ -8856,7 +8858,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "ptyctl-cli"
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"axum",
|
||||
|
||||
+1
-1
@@ -76,7 +76,7 @@ members = [
|
||||
]
|
||||
|
||||
[workspace.package]
|
||||
version = "0.1.1"
|
||||
version = "0.1.6"
|
||||
edition = "2024"
|
||||
license = "Apache-2.0"
|
||||
|
||||
|
||||
@@ -19,3 +19,8 @@ Modifications relative to the upstream work (Apache License 2.0 §4(b) notice):
|
||||
This product includes code ported from openai/codex and sst/opencode; see
|
||||
crates/codegen/kigi-tools/THIRD_PARTY_NOTICES.md for the license terms and
|
||||
the per-file change notices.
|
||||
|
||||
This product bundles model metadata derived from models.dev
|
||||
(https://github.com/sst/models.dev, MIT License) as
|
||||
crates/codegen/kigi-models/enrichment_snapshot.json, and may refresh it
|
||||
from https://models.dev/api.json at runtime.
|
||||
|
||||
@@ -2,17 +2,23 @@
|
||||
|
||||
<h1>Kigi (<code>kigi</code>) 🌘</h1>
|
||||
|
||||
**Kigi** is an unofficial Kimi Code CLI community build — a terminal-based
|
||||
AI coding agent re-targeted at the Kimi Code subscription API and the
|
||||
Moonshot open platform, built on the Apache-2.0 sources of
|
||||
[xai-org/grok-build](https://github.com/xai-org/grok-build).
|
||||
<h3>🕸️ The world's first CLI with built-in <em>Graph Engineering</em></h3>
|
||||
|
||||
It runs as a full-screen TUI that understands your codebase, edits files,
|
||||
executes shell commands, searches the web, and manages long-running tasks —
|
||||
interactively, headlessly for scripting/CI, or embedded in editors via the
|
||||
Agent Client Protocol (ACP).
|
||||
<p><code>/graph</code> turns one objective into a dependency graph of
|
||||
autonomous, self-verifying agent loops — planned, parallelized,
|
||||
adversarially verified, and merged back, end to end.</p>
|
||||
|
||||
**Kigi** is a coding agent that lives in your terminal. It reads the repo,
|
||||
writes the patch, runs the tests, and keeps going while you do something else.
|
||||
Full-screen, headless in CI with `-p`, or docked in your editor over ACP.
|
||||
|
||||
**Already paying for Claude Pro/Max, ChatGPT Plus/Pro, GitHub Copilot, or
|
||||
Grok? Sign in and use it.** No API key, no second bill. Rather bring your own
|
||||
key? OpenAI, Anthropic, Google, DeepSeek, Groq, Moonshot and
|
||||
[two dozen more](#providers-and-api-keys) are wired in.
|
||||
|
||||
[Installation](#installation) ·
|
||||
[Graph engineering](#graph-engineering) ·
|
||||
[Providers and API keys](#providers-and-api-keys) ·
|
||||
[Building from source](#building-from-source) ·
|
||||
[Coexistence with the official CLI](#coexistence-with-the-official-kimi-cli) ·
|
||||
@@ -28,10 +34,6 @@ Agent Client Protocol (ACP).
|
||||
|
||||
## Installation
|
||||
|
||||
Prebuilt single-file binaries for macOS (arm64/x86_64), Linux (arm64/x86_64),
|
||||
and Windows (x86_64) are published on
|
||||
[GitHub Releases](https://github.com/ZacharyZhang-NY/Kigi-CLI/releases):
|
||||
|
||||
```sh
|
||||
# macOS / Linux
|
||||
curl -fsSL https://raw.githubusercontent.com/ZacharyZhang-NY/Kigi-CLI/main/install.sh | bash
|
||||
@@ -43,61 +45,115 @@ irm https://raw.githubusercontent.com/ZacharyZhang-NY/Kigi-CLI/main/install.ps1
|
||||
```
|
||||
|
||||
```sh
|
||||
kigi --version # kigi 0.1.1 … unofficial Kimi Code CLI community build
|
||||
kigi login # sign in with your Kimi Code subscription (device-code flow)
|
||||
kigi # start the TUI
|
||||
kigi login # pick a provider, sign in
|
||||
kigi # go
|
||||
```
|
||||
|
||||
The installer verifies every download against the release's `SHA256SUMS`,
|
||||
installs into `~/.kigi/bin/kigi` (`%USERPROFILE%\.kigi\bin\kigi.exe` on
|
||||
Windows), and prints the PATH line to add. Later releases arrive through the
|
||||
built-in self-updater (`kigi update`, gated by `KIGI_AUTO_UPDATE`), which
|
||||
pulls from the same GitHub Releases feed.
|
||||
Single file, no runtime. macOS and Linux on arm64/x86_64, Windows on x86_64,
|
||||
checksummed against the release's `SHA256SUMS`. `kigi update` handles upgrades.
|
||||
|
||||
## Graph engineering
|
||||
|
||||
Every other agent runs a loop: think, act, repeat — one thread, one thing at a
|
||||
time. `/graph` runs a dependency graph instead. Work that doesn't block other
|
||||
work happens at the same time, in separate worktrees, and nothing merges until
|
||||
something else has tried to tear it apart.
|
||||
|
||||
```
|
||||
/graph <objective> [--budget <tokens>] # decompose + run fully autonomously
|
||||
/graph status # node tree, budget, current work
|
||||
/graph show # box-drawing DAG view
|
||||
/graph pause | resume [--budget <n>] # halt / continue (budget top-up)
|
||||
/graph clear # abandon the graph
|
||||
```
|
||||
|
||||
One command runs the whole thing, start to finish:
|
||||
|
||||
- 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.
|
||||
|
||||
State lives in `.kigi/graph.jsonl`, next to your code. Close the laptop, come
|
||||
back tomorrow, `/graph resume`. A teammate can pick it up from the same file.
|
||||
|
||||
On by default. `KIGI_GRAPH=0` turns it off; `KIGI_GRAPH_CONCURRENCY` (default
|
||||
3) controls how many nodes run at once.
|
||||
|
||||
## Providers and API keys
|
||||
|
||||
Kigi talks to a fixed three-platform registry:
|
||||
29 platforms ship compiled in: 5 you sign into, 24 you hand a key. Nothing is
|
||||
registered at runtime — if it's not in this list, it's not there.
|
||||
|
||||
| Platform id | Base URL | Auth |
|
||||
| ------------- | -------------------------------- | --------------------------- |
|
||||
| `kimi-code` | `https://api.kimi.com/coding/v1` | Kimi Code subscription OAuth (`kigi login`) |
|
||||
| `moonshot-cn` | `https://api.moonshot.cn/v1` | Moonshot open-platform API key |
|
||||
| `moonshot-ai` | `https://api.moonshot.ai/v1` | Moonshot open-platform 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.
|
||||
|
||||
Moonshot API keys come from the environment or `~/.kigi/config.toml`
|
||||
(environment wins; values are never logged):
|
||||
| Platform id | Provider | Sign-in |
|
||||
| ---------------- | ------------------------- | ------------------------------------------- |
|
||||
| `kimi-code` | Kimi Code (original target)| Subscription OAuth (device code) |
|
||||
| `claude-pro-max` | Claude Pro/Max | Subscription OAuth (browser, PKCE) |
|
||||
| `openai-codex` | ChatGPT Plus/Pro (Codex) | Subscription OAuth (browser, PKCE) |
|
||||
| `github-copilot` | GitHub Copilot | Subscription OAuth (device code) |
|
||||
| `xai-grok` | xAI Grok | Subscription OAuth (device code) |
|
||||
|
||||
You get whatever models your plan actually serves — the list is fetched at
|
||||
sign-in, not hardcoded. (ChatGPT/Codex is the exception: its backend publishes
|
||||
no model endpoint, so those four are compiled in.)
|
||||
|
||||
**API-key providers.** Export the env var, or drop the key in
|
||||
`~/.kigi/config.toml`. Keys are never logged.
|
||||
|
||||
| Provider | Platform id | API key env |
|
||||
| ------------------------- | ---------------------- | ----------------------------------------------- |
|
||||
| Moonshot (moonshot.cn) | `moonshot-cn` | `KIGI_MOONSHOT_CN_API_KEY` (or `KIGI_MOONSHOT_API_KEY`) |
|
||||
| Moonshot (moonshot.ai) | `moonshot-ai` | `KIGI_MOONSHOT_AI_API_KEY` (or `KIGI_MOONSHOT_API_KEY`) |
|
||||
| OpenAI | `openai` | `OPENAI_API_KEY` |
|
||||
| Anthropic | `anthropic` | `ANTHROPIC_API_KEY` |
|
||||
| DeepSeek | `deepseek` | `DEEPSEEK_API_KEY` |
|
||||
| Groq | `groq` | `GROQ_API_KEY` |
|
||||
| Mistral | `mistral` | `MISTRAL_API_KEY` |
|
||||
| Fireworks AI | `fireworks` | `FIREWORKS_API_KEY` |
|
||||
| Google Gemini | `google` | `GEMINI_API_KEY` |
|
||||
| OpenRouter | `openrouter` | `OPENROUTER_API_KEY` |
|
||||
| Together AI | `together` | `TOGETHER_API_KEY` |
|
||||
| Cerebras | `cerebras` | `CEREBRAS_API_KEY` |
|
||||
| NVIDIA NIM | `nvidia` | `NVIDIA_API_KEY` |
|
||||
| Vercel AI Gateway | `vercel-ai-gateway` | `AI_GATEWAY_API_KEY` |
|
||||
| xAI (Grok) | `xai` | `XAI_API_KEY` |
|
||||
| Qwen Token Plan | `qwen-token-plan` | `QWEN_TOKEN_PLAN_API_KEY` |
|
||||
| Qwen Token Plan (China) | `qwen-token-plan-cn` | `QWEN_TOKEN_PLAN_CN_API_KEY` |
|
||||
| Kimi For Coding | `kimi-coding` | `KIMI_API_KEY` |
|
||||
| Z.AI | `zai` | `ZAI_API_KEY` |
|
||||
| Z.AI Coding (China) | `zai-coding-cn` | `ZAI_CODING_CN_API_KEY` |
|
||||
| Xiaomi MiMo | `xiaomi` | `XIAOMI_API_KEY` |
|
||||
| Xiaomi Token Plan (China) | `xiaomi-token-plan-cn` | `XIAOMI_TOKEN_PLAN_CN_API_KEY` |
|
||||
| MiniMax | `minimax` | `MINIMAX_API_KEY` |
|
||||
| MiniMax (China) | `minimax-cn` | `MINIMAX_CN_API_KEY` |
|
||||
|
||||
```sh
|
||||
export KIGI_MOONSHOT_API_KEY=sk-... # applies to both open platforms
|
||||
export KIGI_MOONSHOT_CN_API_KEY=sk-... # platform-scoped, beats the generic name
|
||||
export KIGI_MOONSHOT_AI_API_KEY=sk-...
|
||||
export OPENAI_API_KEY=sk-...
|
||||
export XAI_API_KEY=xai-...
|
||||
```
|
||||
|
||||
```toml
|
||||
# ~/.kigi/config.toml
|
||||
[platforms.moonshot-cn]
|
||||
[platforms.openai]
|
||||
api_key = "sk-..."
|
||||
|
||||
[platforms.moonshot-ai]
|
||||
api_key = "sk-..."
|
||||
[platforms.xai]
|
||||
api_key = "xai-..."
|
||||
```
|
||||
|
||||
On login and on startup Kigi syncs each configured platform's model list
|
||||
from `GET {base}/models` and shows the merged catalog in the model picker
|
||||
(catalog keys are `{platform_id}/{model_id}`). Models that advertise
|
||||
selectable thinking levels (e.g. K3's `low`/`high`/`max`) expose them in
|
||||
`/model` and `/effort`. If the sync fails, the last cached catalog is used;
|
||||
with no cache, a small built-in fallback list applies. Model selection
|
||||
resolves as `--model` CLI flag > `KIGI_DEFAULT_MODEL` > `[models] default`
|
||||
in config.toml > server-delivered list > built-in fallback.
|
||||
Model lists sync on startup. Pick one with `/model`, set its thinking level
|
||||
with `/effort`.
|
||||
|
||||
`KIGI_CODE_BASE_URL` re-points the subscription platform (useful for
|
||||
testing); `KIGI_MOONSHOT_CN_BASE_URL` / `KIGI_MOONSHOT_AI_BASE_URL` are the
|
||||
equivalent dev/test overrides for the open platforms.
|
||||
|
||||
The web `search`/`fetch` tools ride the Kimi Code subscription services and
|
||||
are present only on OAuth sessions — API-key-only sessions run without
|
||||
them, matching the official client.
|
||||
Web `search`/`fetch` need a Kimi Code subscription; API-key sessions run
|
||||
without them, same as the official client.
|
||||
|
||||
## Building from source
|
||||
|
||||
@@ -113,19 +169,15 @@ launcher at `bin/protoc`; install dotslash (`brew install dotslash` or
|
||||
|
||||
## Coexistence with the official Kimi CLI
|
||||
|
||||
Kigi is not affiliated with Moonshot AI or xAI, and it coexists with the
|
||||
official `kimi` CLI on the same machine: independent binary name,
|
||||
independent config directory (`~/.kigi`), independent keyring credentials
|
||||
(service `kigi`), and a `KIGI_*` environment-variable namespace. Nothing
|
||||
the official client installs or stores is ever read at runtime or written.
|
||||
On first launch Kigi offers a **one-time, strictly read-only** import of
|
||||
your existing `~/.kimi` configuration (MCP servers, custom providers,
|
||||
default model) via `kigi import-kimi` — file contents and mtimes under
|
||||
`~/.kimi` are left untouched, verified by tests.
|
||||
Kigi started as an unofficial Kimi Code CLI — a community fork of
|
||||
[xai-org/grok-build](https://github.com/xai-org/grok-build), not affiliated
|
||||
with Moonshot AI or xAI. It keeps its own binary, its own `~/.kigi`, its own
|
||||
keyring entry, and its own `KIGI_*` env vars, and never touches what the
|
||||
official `kimi` CLI installed. `kigi import-kimi` copies your old config over
|
||||
once, read-only.
|
||||
|
||||
Kigi is **zero-telemetry**: the only outbound connections are the
|
||||
inference/auth APIs you configure, GitHub Releases for updates, and MCP
|
||||
servers you add.
|
||||
**Zero telemetry.** It talks to the APIs you configured, GitHub Releases, and
|
||||
your own MCP servers. Nothing else.
|
||||
|
||||
## License
|
||||
|
||||
|
||||
@@ -272,6 +272,7 @@ mod tests {
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: std::num::NonZeroU64::new(128_000).unwrap(),
|
||||
reasoning_effort: None,
|
||||
|
||||
@@ -23,6 +23,7 @@ fn test_config_with_window(context_window: u64) -> SamplingConfig {
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: NonZeroU64::new(context_window)
|
||||
.expect("test context_window must be non-zero"),
|
||||
@@ -913,6 +914,7 @@ async fn update_sampling_config_is_queryable() {
|
||||
temperature: Some(0.5),
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: NonZeroU64::new(200_000).unwrap(),
|
||||
reasoning_effort: None,
|
||||
@@ -1298,6 +1300,7 @@ async fn build_request_uses_sampling_config() {
|
||||
temperature: Some(0.7),
|
||||
top_p: Some(0.9),
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: NonZeroU64::new(128_000).unwrap(),
|
||||
reasoning_effort: None,
|
||||
@@ -3401,6 +3404,7 @@ async fn sampling_config_survives_compaction_replacement() {
|
||||
temperature: Some(0.7),
|
||||
top_p: Some(0.95),
|
||||
api_backend: ApiBackend::Responses,
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: NonZeroU64::new(500_000).unwrap(),
|
||||
reasoning_effort: None,
|
||||
@@ -3481,6 +3485,7 @@ async fn model_metadata_lost_after_compaction_then_recovered_on_next_turn() {
|
||||
temperature: Some(0.7),
|
||||
top_p: Some(0.95),
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: NonZeroU64::new(500_000).unwrap(),
|
||||
reasoning_effort: None,
|
||||
@@ -3569,6 +3574,7 @@ async fn context_window_downgrade_triggers_auto_compact() {
|
||||
temperature: Some(0.7),
|
||||
top_p: Some(0.95),
|
||||
api_backend: ApiBackend::Responses,
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: NonZeroU64::new(500_000).unwrap(),
|
||||
reasoning_effort: None,
|
||||
|
||||
@@ -370,6 +370,7 @@ mod tests {
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: std::num::NonZeroU64::new(128_000).unwrap(),
|
||||
reasoning_effort: None,
|
||||
|
||||
@@ -178,6 +178,7 @@ mod tests {
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: NonZeroU64::new(128_000).unwrap(),
|
||||
reasoning_effort: None,
|
||||
@@ -221,6 +222,7 @@ mod tests {
|
||||
temperature: Some(0.7),
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: NonZeroU64::new(128_000).unwrap(),
|
||||
reasoning_effort: None,
|
||||
|
||||
@@ -9,6 +9,7 @@ description = "Kimi platform registry, /models wire contract, capability derivat
|
||||
kigi-env = { workspace = true }
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
serde_json = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,388 @@
|
||||
//! models.dev metadata enrichment (multi-provider expansion).
|
||||
//!
|
||||
//! Most provider `/models` listings return bare ids — no context window, no
|
||||
//! thinking levels. This module carries per-model metadata keyed by
|
||||
//! models.dev provider id, sourced from the bundled snapshot
|
||||
//! (`enrichment_snapshot.json`, regenerated by
|
||||
//! `scripts/gen_enrichment_snapshot.py`) or a runtime refresh fetched by the
|
||||
//! shell. It NEVER invents model availability: the live listing is the only
|
||||
//! source of which models exist — enrichment only fills metadata gaps, and
|
||||
//! wire-served values always win (see [`enrich_wire_model`]).
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::LazyLock;
|
||||
|
||||
/// One model's enrichment metadata. All fields optional-by-default so the
|
||||
/// snapshot stays minimal.
|
||||
#[derive(Debug, Clone, Default, PartialEq, serde::Deserialize, serde::Serialize)]
|
||||
pub struct EnrichmentModel {
|
||||
/// Max context window in tokens (`limit.context`).
|
||||
#[serde(default, skip_serializing_if = "is_zero")]
|
||||
pub context: u64,
|
||||
/// Max output tokens (`limit.output`); fills a wire-unserved output cap
|
||||
/// (Anthropic 400s when `max_tokens` exceeds the model's limit).
|
||||
#[serde(default, skip_serializing_if = "is_zero")]
|
||||
pub output: u64,
|
||||
/// Model supports reasoning/thinking.
|
||||
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
|
||||
pub reasoning: bool,
|
||||
/// Selectable effort levels (canonical tokens, e.g. ["low","high","max"]).
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub efforts: Vec<String>,
|
||||
/// Accepts image input.
|
||||
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
|
||||
pub image_in: bool,
|
||||
/// Supports tool calling (used later to filter non-agentic listings).
|
||||
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
|
||||
pub tool_call: bool,
|
||||
/// Human display name.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub name: Option<String>,
|
||||
}
|
||||
|
||||
fn is_zero(v: &u64) -> bool {
|
||||
*v == 0
|
||||
}
|
||||
|
||||
/// `models.dev provider id -> model id -> metadata`.
|
||||
pub type EnrichmentCatalog = BTreeMap<String, BTreeMap<String, EnrichmentModel>>;
|
||||
|
||||
/// The bundled snapshot, embedded at compile time — RAW models.dev shape,
|
||||
/// filtered to the providers kigi references (pure-filter script:
|
||||
/// `scripts/gen_enrichment_snapshot.py`). OFFLINE FALLBACK for the runtime
|
||||
/// refresh; both parse through [`parse_api_json`] so there is exactly one
|
||||
/// field interpretation.
|
||||
pub const ENRICHMENT_SNAPSHOT_JSON: &str = include_str!("../enrichment_snapshot.json");
|
||||
|
||||
static BUNDLED: LazyLock<EnrichmentCatalog> = LazyLock::new(|| {
|
||||
// Baked-in JSON — a mismatch here is a developer error, not a runtime
|
||||
// condition (same policy as default_models.json).
|
||||
parse_api_json(ENRICHMENT_SNAPSHOT_JSON, None)
|
||||
.expect("enrichment_snapshot.json: invalid JSON (regenerate via script)")
|
||||
});
|
||||
|
||||
// ── Raw models.dev api.json shape (parse-only) ──────────────────────────────
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RawProvider {
|
||||
#[serde(default)]
|
||||
models: BTreeMap<String, RawModel>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize, Default)]
|
||||
#[serde(default)]
|
||||
struct RawModel {
|
||||
name: Option<String>,
|
||||
reasoning: bool,
|
||||
reasoning_options: Vec<RawReasoningOption>,
|
||||
limit: RawLimit,
|
||||
modalities: RawModalities,
|
||||
tool_call: bool,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize, Default)]
|
||||
#[serde(default)]
|
||||
struct RawReasoningOption {
|
||||
r#type: String,
|
||||
values: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize, Default)]
|
||||
#[serde(default)]
|
||||
struct RawLimit {
|
||||
context: u64,
|
||||
output: u64,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize, Default)]
|
||||
#[serde(default)]
|
||||
struct RawModalities {
|
||||
input: Vec<String>,
|
||||
}
|
||||
|
||||
/// Parse a models.dev `api.json` document (full download or the bundled
|
||||
/// filtered snapshot) into the in-memory catalog. `keep`: restrict to these
|
||||
/// provider ids (runtime refresh filters the full 3MB download to the
|
||||
/// registry's providers); `None` keeps everything present.
|
||||
///
|
||||
/// Blast-radius confinement: providers are filtered on RAW keys first and
|
||||
/// typed-parsed individually — a schema drift in one of the ~70 providers
|
||||
/// kigi never keeps cannot fail the whole refresh, and a malformed KEPT
|
||||
/// provider is warn-skipped (its models fall to defaults) rather than
|
||||
/// killing the others. Only a document that isn't a JSON object errors
|
||||
/// (the caller falls back to cache/bundled — never a silently empty
|
||||
/// catalog). Keys starting with `_` (provenance stamps) are skipped.
|
||||
pub fn parse_api_json(
|
||||
json: &str,
|
||||
keep: Option<&std::collections::BTreeSet<&str>>,
|
||||
) -> Result<EnrichmentCatalog, serde_json::Error> {
|
||||
let raw: BTreeMap<String, serde_json::Value> = serde_json::from_str(json)?;
|
||||
let mut catalog = EnrichmentCatalog::new();
|
||||
for (pid, value) in raw {
|
||||
if pid.starts_with('_') {
|
||||
continue;
|
||||
}
|
||||
if let Some(keep) = keep
|
||||
&& !keep.contains(pid.as_str())
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let provider: RawProvider = match serde_json::from_value(value) {
|
||||
Ok(p) => p,
|
||||
Err(e) => {
|
||||
tracing::warn!(provider = %pid, error = %e,
|
||||
"models.dev provider entry malformed; skipping (models fall to defaults)");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let models = provider
|
||||
.models
|
||||
.into_iter()
|
||||
.map(|(mid, m)| {
|
||||
let efforts = m
|
||||
.reasoning_options
|
||||
.into_iter()
|
||||
.find(|o| o.r#type == "effort" && !o.values.is_empty())
|
||||
.map(|o| o.values)
|
||||
.unwrap_or_default();
|
||||
let meta = EnrichmentModel {
|
||||
context: m.limit.context,
|
||||
output: m.limit.output,
|
||||
reasoning: m.reasoning,
|
||||
efforts,
|
||||
image_in: m.modalities.input.iter().any(|s| s == "image"),
|
||||
tool_call: m.tool_call,
|
||||
name: m.name,
|
||||
};
|
||||
(mid, meta)
|
||||
})
|
||||
.collect();
|
||||
catalog.insert(pid, models);
|
||||
}
|
||||
Ok(catalog)
|
||||
}
|
||||
|
||||
/// The compiled-in enrichment catalog.
|
||||
pub fn bundled_enrichment() -> &'static EnrichmentCatalog {
|
||||
&BUNDLED
|
||||
}
|
||||
|
||||
/// Look up one model's metadata. `None` when the provider or model is
|
||||
/// unknown to the catalog (callers fall through to defaults).
|
||||
pub fn lookup<'a>(
|
||||
catalog: &'a EnrichmentCatalog,
|
||||
models_dev_id: &str,
|
||||
model_id: &str,
|
||||
) -> Option<&'a EnrichmentModel> {
|
||||
catalog.get(models_dev_id)?.get(model_id)
|
||||
}
|
||||
|
||||
/// Fill metadata gaps on a live listing entry. WIRE WINS: a field the
|
||||
/// provider served is never overwritten — enrichment only supplies what the
|
||||
/// wire left absent/zero. Never changes the model id (availability stays
|
||||
/// wire-truth).
|
||||
pub fn enrich_wire_model(wire: &mut crate::WireModel, meta: &EnrichmentModel) {
|
||||
if wire.context_length == 0 && meta.context > 0 {
|
||||
wire.context_length = meta.context;
|
||||
}
|
||||
if wire.max_output_tokens == 0 && meta.output > 0 {
|
||||
wire.max_output_tokens = meta.output;
|
||||
}
|
||||
if meta.reasoning {
|
||||
wire.supports_reasoning = true;
|
||||
}
|
||||
if meta.image_in {
|
||||
wire.supports_image_in = true;
|
||||
}
|
||||
if wire.display_name.is_none() {
|
||||
wire.display_name = meta.name.clone();
|
||||
}
|
||||
if wire.think_efforts.is_none() && !meta.efforts.is_empty() {
|
||||
wire.think_efforts = Some(crate::WireThinkEfforts {
|
||||
support: true,
|
||||
valid_efforts: meta.efforts.clone(),
|
||||
// The provider's implicit default applies when the user doesn't
|
||||
// pick a level; models.dev doesn't record one.
|
||||
default_effort: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// The bundled snapshot parses and covers the providers this expansion
|
||||
/// references; kimi entries cross-check against the live wire values the
|
||||
/// registry already knows (guards against a corrupted regeneration).
|
||||
#[test]
|
||||
fn bundled_snapshot_parses_and_cross_checks_kimi() {
|
||||
let catalog = bundled_enrichment();
|
||||
assert!(
|
||||
catalog.len() >= 25,
|
||||
"snapshot lost providers: {}",
|
||||
catalog.len()
|
||||
);
|
||||
let k3 = lookup(catalog, "kimi-for-coding", "k3").expect("k3 present");
|
||||
assert_eq!(k3.context, 1_048_576, "k3 context must match the live wire");
|
||||
assert_eq!(k3.efforts, ["low", "high", "max"]);
|
||||
assert!(k3.reasoning);
|
||||
let opus = lookup(catalog, "anthropic", "claude-opus-4-8").expect("opus present");
|
||||
assert_eq!(opus.context, 1_000_000);
|
||||
assert_eq!(opus.efforts, ["low", "medium", "high", "xhigh", "max"]);
|
||||
}
|
||||
|
||||
/// Wire-served fields are never overwritten; absent fields are filled.
|
||||
#[test]
|
||||
fn enrich_fills_gaps_and_never_overwrites_wire() {
|
||||
let meta = EnrichmentModel {
|
||||
context: 400_000,
|
||||
output: 64_000,
|
||||
reasoning: true,
|
||||
efforts: vec!["low".into(), "high".into()],
|
||||
image_in: true,
|
||||
name: Some("GPT Test".into()),
|
||||
..Default::default()
|
||||
};
|
||||
// Bare listing entry (OpenAI-style: id only).
|
||||
let mut bare: crate::WireModel =
|
||||
serde_json::from_value(serde_json::json!({ "id": "gpt-test" })).unwrap();
|
||||
enrich_wire_model(&mut bare, &meta);
|
||||
assert_eq!(bare.context_length, 400_000);
|
||||
assert_eq!(bare.max_output_tokens, 64_000, "output cap filled");
|
||||
assert!(bare.supports_reasoning);
|
||||
assert!(bare.supports_image_in);
|
||||
assert_eq!(bare.display_name.as_deref(), Some("GPT Test"));
|
||||
let efforts = bare.think_efforts.expect("efforts filled");
|
||||
assert!(efforts.support);
|
||||
assert_eq!(efforts.valid_efforts, ["low", "high"]);
|
||||
assert_eq!(efforts.default_effort, None);
|
||||
|
||||
// Wire-served entry: nothing may change.
|
||||
let mut served: crate::WireModel = serde_json::from_value(serde_json::json!({
|
||||
"id": "gpt-test",
|
||||
"context_length": 123,
|
||||
"max_output_tokens": 77,
|
||||
"display_name": "Wire Name",
|
||||
"think_efforts": { "support": true, "valid_efforts": ["max"] }
|
||||
}))
|
||||
.unwrap();
|
||||
enrich_wire_model(&mut served, &meta);
|
||||
assert_eq!(served.context_length, 123, "wire context wins");
|
||||
assert_eq!(served.max_output_tokens, 77, "wire output cap wins");
|
||||
assert_eq!(served.display_name.as_deref(), Some("Wire Name"));
|
||||
assert_eq!(
|
||||
served.think_efforts.unwrap().valid_efforts,
|
||||
["max"],
|
||||
"wire efforts win"
|
||||
);
|
||||
}
|
||||
|
||||
/// The runtime refresh parses the FULL api.json (extra fields like cost/
|
||||
/// env/doc present) and filters to the registry's provider ids.
|
||||
#[test]
|
||||
fn parse_api_json_filters_and_tolerates_unknown_fields() {
|
||||
let full = serde_json::json!({
|
||||
"openai": {
|
||||
"id": "openai", "env": ["OPENAI_API_KEY"], "doc": "https://x",
|
||||
"models": {
|
||||
"gpt-test": {
|
||||
"name": "GPT Test",
|
||||
"reasoning": true,
|
||||
"reasoning_options": [
|
||||
{"type": "effort", "values": ["low", "high", "max"]}
|
||||
],
|
||||
"limit": {"context": 400000, "output": 128000},
|
||||
"modalities": {"input": ["text", "image"], "output": ["text"]},
|
||||
"tool_call": true,
|
||||
"cost": {"input": 1.25, "output": 10}
|
||||
}
|
||||
}
|
||||
},
|
||||
"unwanted-provider": { "models": { "m": {} } }
|
||||
})
|
||||
.to_string();
|
||||
let keep: std::collections::BTreeSet<&str> = ["openai"].into();
|
||||
let catalog = parse_api_json(&full, Some(&keep)).unwrap();
|
||||
assert!(!catalog.contains_key("unwanted-provider"));
|
||||
let m = lookup(&catalog, "openai", "gpt-test").unwrap();
|
||||
assert_eq!(m.context, 400_000);
|
||||
assert_eq!(m.output, 128_000);
|
||||
assert_eq!(m.efforts, ["low", "high", "max"]);
|
||||
assert!(m.reasoning && m.image_in && m.tool_call);
|
||||
assert_eq!(m.name.as_deref(), Some("GPT Test"));
|
||||
// Malformed document errors — callers must fall back loudly, never
|
||||
// proceed with a silently empty catalog.
|
||||
assert!(parse_api_json("not json", None).is_err());
|
||||
}
|
||||
|
||||
/// Every registry `models_dev_id` must be covered by the bundled
|
||||
/// snapshot — a registry row added without updating the script's
|
||||
/// TARGETS would silently diverge bundled-vs-refresh behavior.
|
||||
#[test]
|
||||
fn bundled_snapshot_covers_every_registry_models_dev_id() {
|
||||
let catalog = bundled_enrichment();
|
||||
for platform in crate::PlatformId::ALL {
|
||||
if let Some(dev_id) = platform.models_dev_id() {
|
||||
assert!(
|
||||
catalog.contains_key(dev_id),
|
||||
"{}: models_dev_id {dev_id:?} missing from the bundled \
|
||||
snapshot — add it to scripts/gen_enrichment_snapshot.py \
|
||||
TARGETS and regenerate",
|
||||
platform.as_str(),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Field-coverage guard against script/parser drift: for every field
|
||||
/// `RawModel` reads, at least one bundled model must carry a non-default
|
||||
/// value. A regeneration that dropped a MODEL_KEYS entry (e.g.
|
||||
/// `modalities`) would zero that field across the whole snapshot and
|
||||
/// fail here, instead of silently diverging from runtime refreshes.
|
||||
#[test]
|
||||
fn bundled_snapshot_carries_every_parsed_field() {
|
||||
let all: Vec<&EnrichmentModel> = bundled_enrichment()
|
||||
.values()
|
||||
.flat_map(|models| models.values())
|
||||
.collect();
|
||||
assert!(all.iter().any(|m| m.context > 0), "no context anywhere");
|
||||
assert!(all.iter().any(|m| m.output > 0), "no output anywhere");
|
||||
assert!(all.iter().any(|m| m.reasoning), "no reasoning anywhere");
|
||||
assert!(
|
||||
all.iter().any(|m| !m.efforts.is_empty()),
|
||||
"no efforts anywhere"
|
||||
);
|
||||
assert!(all.iter().any(|m| m.image_in), "no image_in anywhere");
|
||||
assert!(all.iter().any(|m| m.tool_call), "no tool_call anywhere");
|
||||
assert!(all.iter().any(|m| m.name.is_some()), "no names anywhere");
|
||||
}
|
||||
|
||||
/// One malformed provider (schema drift on models.dev) must not kill
|
||||
/// the whole refresh — kept siblings still parse; only a non-object
|
||||
/// document errors.
|
||||
#[test]
|
||||
fn malformed_provider_is_skipped_not_fatal() {
|
||||
let doc = serde_json::json!({
|
||||
"good": { "models": { "m": { "limit": {"context": 7} } } },
|
||||
"drifted": { "models": "this is not an object" },
|
||||
"_meta": { "source": "stamp, must be ignored" }
|
||||
})
|
||||
.to_string();
|
||||
let catalog = parse_api_json(&doc, None).expect("document parses");
|
||||
assert_eq!(
|
||||
lookup(&catalog, "good", "m").map(|m| m.context),
|
||||
Some(7),
|
||||
"sibling providers must survive one drifted provider"
|
||||
);
|
||||
assert!(!catalog.contains_key("drifted"));
|
||||
assert!(!catalog.contains_key("_meta"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lookup_misses_are_none() {
|
||||
let catalog = bundled_enrichment();
|
||||
assert!(lookup(catalog, "no-such-provider", "x").is_none());
|
||||
assert!(lookup(catalog, "openai", "no-such-model").is_none());
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -97,7 +97,7 @@ impl ContentController {
|
||||
// pins them to the mock so no PTY test can reach a live endpoint.
|
||||
("KIGI_MOONSHOT_CN_BASE_URL".into(), self.url()),
|
||||
("KIGI_MOONSHOT_AI_BASE_URL".into(), self.url()),
|
||||
("XAI_API_KEY".into(), "test-key-for-ci".into()),
|
||||
("KIGI_API_KEY".into(), "test-key-for-ci".into()),
|
||||
("KIGI_TELEMETRY_ENABLED".into(), "false".into()),
|
||||
("KIGI_FEEDBACK_ENABLED".into(), "false".into()),
|
||||
("KIGI_TRACE_UPLOAD".into(), "false".into()),
|
||||
@@ -287,7 +287,7 @@ mod tests {
|
||||
assert_eq!(get("KIGI_API_BASE_URL"), Some(content.url()));
|
||||
assert_eq!(get("KIGI_MOONSHOT_CN_BASE_URL"), Some(content.url()));
|
||||
assert_eq!(get("KIGI_MOONSHOT_AI_BASE_URL"), Some(content.url()));
|
||||
assert_eq!(get("XAI_API_KEY").as_deref(), Some("test-key-for-ci"));
|
||||
assert_eq!(get("KIGI_API_KEY").as_deref(), Some("test-key-for-ci"));
|
||||
assert_eq!(get("KIGI_TELEMETRY_ENABLED").as_deref(), Some("false"));
|
||||
assert_eq!(get("KIGI_FEEDBACK_ENABLED").as_deref(), Some("false"));
|
||||
assert_eq!(get("KIGI_TRACE_UPLOAD").as_deref(), Some("false"));
|
||||
|
||||
@@ -96,11 +96,12 @@ pub fn seed_fake_oauth(content: &ContentController, user: &str) {
|
||||
.expect("seed fake oauth auth.json");
|
||||
}
|
||||
|
||||
/// [`ContentController::env_for_pager`] minus `XAI_API_KEY`, so the entry
|
||||
/// written by [`seed_fake_oauth`] is the active credential.
|
||||
/// [`ContentController::env_for_pager`] minus the house BYOK key
|
||||
/// (`KIGI_API_KEY`), so the entry written by [`seed_fake_oauth`] is the active
|
||||
/// credential.
|
||||
pub fn oauth_env_for_pager(content: &ContentController) -> Vec<(String, String)> {
|
||||
let mut env = content.env_for_pager();
|
||||
env.retain(|(k, _)| k != "XAI_API_KEY");
|
||||
env.retain(|(k, _)| k != "KIGI_API_KEY");
|
||||
env
|
||||
}
|
||||
|
||||
|
||||
@@ -89,6 +89,10 @@ mod tests {
|
||||
top_p: None,
|
||||
api_backend: ApiBackend::ChatCompletions,
|
||||
auth_scheme: Default::default(),
|
||||
anthropic_oauth: false,
|
||||
github_copilot: false,
|
||||
openai_codex: false,
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: IndexMap::new(),
|
||||
context_window: 8192,
|
||||
force_http1: false,
|
||||
|
||||
@@ -41,6 +41,35 @@ pub use kigi_sampling_types::ApiBackend;
|
||||
const AGENT_PRODUCT: &str = "kigi";
|
||||
const ANTHROPIC_DEFAULT_MAX_TOKENS: u32 = 128_000;
|
||||
|
||||
/// Prepend the required Claude-Code system block to a Messages request's
|
||||
/// `system` field (Claude Pro/Max OAuth path). Anthropic inspects the FIRST
|
||||
/// system block, so the prefix is inserted as a distinct leading `text` block
|
||||
/// while preserving any caller-supplied prompt (string or block form).
|
||||
/// Idempotent: a leading block already equal to the prefix is not re-added.
|
||||
fn prepend_claude_code_system_prefix(system: &mut Option<messages::SystemParam>) {
|
||||
use messages::{SystemParam, TextBlock};
|
||||
let text_block = |text: String| TextBlock {
|
||||
r#type: "text".to_string(),
|
||||
text,
|
||||
cache_control: None,
|
||||
};
|
||||
let mut blocks = match system.take() {
|
||||
None => Vec::new(),
|
||||
Some(SystemParam::Text(text)) => vec![text_block(text)],
|
||||
Some(SystemParam::Blocks(blocks)) => blocks,
|
||||
};
|
||||
let already_present = blocks
|
||||
.first()
|
||||
.is_some_and(|b| b.text == kigi_sampling_types::CLAUDE_CODE_SYSTEM_PREFIX);
|
||||
if !already_present {
|
||||
blocks.insert(
|
||||
0,
|
||||
text_block(kigi_sampling_types::CLAUDE_CODE_SYSTEM_PREFIX.to_string()),
|
||||
);
|
||||
}
|
||||
*system = Some(SystemParam::Blocks(blocks));
|
||||
}
|
||||
|
||||
/// Parse the `Retry-After` response header as delta-seconds.
|
||||
/// Our inference backends only emit integer seconds (never HTTP-date),
|
||||
/// so we only handle that form. HTTP-dates silently return `None` and
|
||||
@@ -70,6 +99,9 @@ fn deserialize_response_event(data: &str) -> Result<rs::ResponseStreamEvent> {
|
||||
Err(first_err) => {
|
||||
// Try sanitizing: parse as Value, strip unknown tools, retry.
|
||||
if let Ok(mut value) = serde_json::from_str::<serde_json::Value>(data) {
|
||||
// A `max` reasoning-effort echo is unrepresentable in the
|
||||
// typed enum; drop it so the event parses.
|
||||
kigi_sampling_types::normalize_effort_echo(&mut value);
|
||||
// Strip tools that async_openai's rs::Tool can't deserialize
|
||||
// (e.g., xAI-specific "x_search"). Instead of maintaining a
|
||||
// hardcoded allowlist, try deserializing each tool entry —
|
||||
@@ -267,8 +299,14 @@ struct ClientDefaults {
|
||||
top_p: Option<f32>,
|
||||
api_backend: ApiBackend,
|
||||
auth_scheme: AuthScheme,
|
||||
chat_compat: kigi_sampling_types::ChatCompat,
|
||||
stream_tool_calls: bool,
|
||||
doom_loop_recovery: Option<kigi_sampling_types::DoomLoopRecoveryPolicy>,
|
||||
/// Claude Pro/Max OAuth Messages adaptation (see [`SamplerConfig`]).
|
||||
anthropic_oauth: bool,
|
||||
/// ChatGPT/Codex Responses adaptation (see [`SamplerConfig`]). Gates the
|
||||
/// per-request `chatgpt-account-id` header derived from the bearer JWT.
|
||||
openai_codex: bool,
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
@@ -368,6 +406,14 @@ impl SamplingClient {
|
||||
)
|
||||
})?;
|
||||
headers.insert(HeaderName::from_static("x-api-key"), header_value);
|
||||
if config.api_backend == kigi_sampling_types::ApiBackend::Messages {
|
||||
// The real Anthropic Messages wire rejects requests
|
||||
// without this; compatible endpoints ignore it.
|
||||
headers.insert(
|
||||
HeaderName::from_static("anthropic-version"),
|
||||
HeaderValue::from_static(kigi_sampling_types::ANTHROPIC_VERSION),
|
||||
);
|
||||
}
|
||||
}
|
||||
AuthScheme::Bearer => {
|
||||
let bearer = format!("Bearer {}", api_key);
|
||||
@@ -386,6 +432,74 @@ impl SamplingClient {
|
||||
}
|
||||
}
|
||||
|
||||
// Claude Pro/Max OAuth identity headers (claude-pro-max only). The
|
||||
// OAuth `sk-ant-oat…` bearer is Claude-Code-scoped, so Anthropic
|
||||
// rejects the Messages request without the oauth beta + claude-cli
|
||||
// identity. Gated on `anthropic_oauth` so API-key anthropic/minimax
|
||||
// requests carry none of this and stay byte-identical. `Accept` is set
|
||||
// per-request (text/event-stream for streams), so it is NOT added here.
|
||||
if config.anthropic_oauth {
|
||||
headers.insert(
|
||||
HeaderName::from_static("anthropic-version"),
|
||||
HeaderValue::from_static(kigi_sampling_types::ANTHROPIC_VERSION),
|
||||
);
|
||||
headers.insert(
|
||||
HeaderName::from_static("anthropic-beta"),
|
||||
HeaderValue::from_static(kigi_sampling_types::ANTHROPIC_OAUTH_BETA),
|
||||
);
|
||||
headers.insert(
|
||||
HeaderName::from_static("x-app"),
|
||||
HeaderValue::from_static("cli"),
|
||||
);
|
||||
headers.insert(
|
||||
HeaderName::from_static("anthropic-dangerous-direct-browser-access"),
|
||||
HeaderValue::from_static("true"),
|
||||
);
|
||||
}
|
||||
|
||||
// GitHub Copilot editor-identity headers (github-copilot only). Copilot's
|
||||
// proxy validates the VS Code editor identity, so the ChatCompletions
|
||||
// request MUST carry it. `User-Agent` is set in the UA block below (it
|
||||
// would otherwise be overwritten); here we add the other three editor
|
||||
// headers plus `X-Initiator: user`. Gated on `github_copilot` so every
|
||||
// other ChatCompletions provider (groq, …) stays byte-identical.
|
||||
if config.github_copilot {
|
||||
headers.insert(
|
||||
HeaderName::from_static("editor-version"),
|
||||
HeaderValue::from_static(kigi_sampling_types::COPILOT_EDITOR_VERSION),
|
||||
);
|
||||
headers.insert(
|
||||
HeaderName::from_static("editor-plugin-version"),
|
||||
HeaderValue::from_static(kigi_sampling_types::COPILOT_EDITOR_PLUGIN_VERSION),
|
||||
);
|
||||
headers.insert(
|
||||
HeaderName::from_static("copilot-integration-id"),
|
||||
HeaderValue::from_static(kigi_sampling_types::COPILOT_INTEGRATION_ID),
|
||||
);
|
||||
headers.insert(
|
||||
HeaderName::from_static("x-initiator"),
|
||||
HeaderValue::from_static(kigi_sampling_types::COPILOT_INITIATOR),
|
||||
);
|
||||
}
|
||||
|
||||
// ChatGPT/Codex Responses identity headers (openai-codex only). The
|
||||
// Codex backend authorizes the OAuth bearer AND validates the Codex
|
||||
// client identity. The static pieces (`originator`, `OpenAI-Beta`) ride
|
||||
// every request; `chatgpt-account-id` is dynamic (derived per-request
|
||||
// from the bearer JWT in `post()`), and `User-Agent` is set in the UA
|
||||
// block below. Gated on `openai_codex` so API-key `openai` Responses
|
||||
// requests stay byte-identical (`store: false` is the shared default).
|
||||
if config.openai_codex {
|
||||
headers.insert(
|
||||
HeaderName::from_static("originator"),
|
||||
HeaderValue::from_static(kigi_sampling_types::CODEX_ORIGINATOR),
|
||||
);
|
||||
headers.insert(
|
||||
HeaderName::from_static("openai-beta"),
|
||||
HeaderValue::from_static(kigi_sampling_types::CODEX_OPENAI_BETA),
|
||||
);
|
||||
}
|
||||
|
||||
// Apply all extra headers verbatim. This is the single
|
||||
// injection point for proxy-auth headers and any other URL- or
|
||||
// environment-specific headers the session decides to set.
|
||||
@@ -403,12 +517,23 @@ impl SamplingClient {
|
||||
// (PRD F3: auth is a plain bearer; kimi-cli sends only User-Agent
|
||||
// plus the OAuth device headers, src/kimi_cli/llm.py:317-323).
|
||||
{
|
||||
let ua_string = match config.origin_client.as_ref() {
|
||||
// Claude Pro/Max OAuth presents the claude-cli identity; GitHub
|
||||
// Copilot presents the VS Code Copilot Chat identity; every other
|
||||
// path keeps the kigi User-Agent.
|
||||
let ua_string = if config.anthropic_oauth {
|
||||
kigi_sampling_types::CLAUDE_CODE_USER_AGENT.to_string()
|
||||
} else if config.github_copilot {
|
||||
kigi_sampling_types::COPILOT_USER_AGENT.to_string()
|
||||
} else if config.openai_codex {
|
||||
kigi_sampling_types::CODEX_USER_AGENT.to_string()
|
||||
} else {
|
||||
match config.origin_client.as_ref() {
|
||||
Some(origin) => user_agent_string_for(origin),
|
||||
None => user_agent_string_for(&OriginClientInfo {
|
||||
product: AGENT_PRODUCT.to_string(),
|
||||
version: Some(agent_version()),
|
||||
}),
|
||||
}
|
||||
};
|
||||
if let Ok(v) = HeaderValue::from_str(&ua_string) {
|
||||
headers.insert(USER_AGENT, v);
|
||||
@@ -445,8 +570,11 @@ impl SamplingClient {
|
||||
top_p: config.top_p,
|
||||
api_backend: config.api_backend,
|
||||
auth_scheme: config.auth_scheme,
|
||||
chat_compat: config.chat_compat,
|
||||
stream_tool_calls: config.stream_tool_calls,
|
||||
doom_loop_recovery: config.doom_loop_recovery,
|
||||
anthropic_oauth: config.anthropic_oauth,
|
||||
openai_codex: config.openai_codex,
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
@@ -486,6 +614,23 @@ impl SamplingClient {
|
||||
}
|
||||
}
|
||||
}
|
||||
// ChatGPT/Codex `chatgpt-account-id` (openai-codex only): derive it
|
||||
// STATELESSLY from the bearer that will actually ride this request (the
|
||||
// resolver-fresh one just set, or the construction-time bearer in
|
||||
// `default_headers`) by decoding its JWT claim. A refreshed token still
|
||||
// carries the claim, so there is no persisted account-id field.
|
||||
// openai-codex-gated → API-key `openai` never gets this header.
|
||||
// SECURITY: the bearer and account id are never logged here.
|
||||
if self.defaults.openai_codex
|
||||
&& let Some(bearer) = headers
|
||||
.get(AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| s.strip_prefix("Bearer "))
|
||||
&& let Some(account_id) = kigi_sampling_types::chatgpt_account_id_from_jwt(bearer)
|
||||
&& let Ok(v) = HeaderValue::from_str(&account_id)
|
||||
{
|
||||
headers.insert(HeaderName::from_static("chatgpt-account-id"), v);
|
||||
}
|
||||
{
|
||||
let auth_prefix = headers
|
||||
.get(AUTHORIZATION)
|
||||
@@ -595,6 +740,9 @@ impl SamplingClient {
|
||||
|| lower.contains("apikey")
|
||||
|| lower.contains("token")
|
||||
|| lower.contains("secret")
|
||||
// `chatgpt-account-id` (Codex) and any other account identifier: a
|
||||
// stable per-user id that must never reach a log.
|
||||
|| lower.contains("account-id")
|
||||
}
|
||||
|
||||
/// Format a single header for error messages, redacting sensitive values.
|
||||
@@ -764,7 +912,10 @@ impl SamplingClient {
|
||||
tracing::error!("Failed to serialize chat/completions request: {}", e);
|
||||
SamplingError::Serialization(e)
|
||||
})?;
|
||||
crate::kimi_compat::adapt_chat_completions_body(&mut request_body);
|
||||
crate::kimi_compat::adapt_chat_completions_body_for(
|
||||
self.defaults.chat_compat,
|
||||
&mut request_body,
|
||||
);
|
||||
|
||||
let http_request = self
|
||||
.post(self.endpoint("chat/completions"))
|
||||
@@ -818,7 +969,10 @@ impl SamplingClient {
|
||||
tracing::error!("Failed to serialize chat/completions request: {}", e);
|
||||
SamplingError::Serialization(e)
|
||||
})?;
|
||||
crate::kimi_compat::adapt_chat_completions_body(&mut request_body);
|
||||
crate::kimi_compat::adapt_chat_completions_body_for(
|
||||
self.defaults.chat_compat,
|
||||
&mut request_body,
|
||||
);
|
||||
|
||||
let http_request = self
|
||||
.post(self.endpoint("chat/completions"))
|
||||
@@ -1028,6 +1182,10 @@ impl SamplingClient {
|
||||
// it in post-serialize. This is the last surviving piece of the
|
||||
// old raw_output machinery.
|
||||
kigi_sampling_types::patch_reasoning_text_types(&mut request_body);
|
||||
kigi_sampling_types::patch_reasoning_effort(&mut request_body, request.reasoning_effort);
|
||||
if self.defaults.openai_codex {
|
||||
kigi_sampling_types::adapt_body_for_codex_backend(&mut request_body);
|
||||
}
|
||||
let http_request = self.post(self.endpoint("responses")).json(&request_body);
|
||||
|
||||
let response = http_request.send().await.map_err(|e| {
|
||||
@@ -1074,7 +1232,16 @@ impl SamplingClient {
|
||||
});
|
||||
}
|
||||
|
||||
let response_obj = serde_json::from_slice::<rs::Response>(&bytes).map_err(|e| {
|
||||
let mut response_value =
|
||||
serde_json::from_slice::<serde_json::Value>(&bytes).map_err(|e| {
|
||||
let raw_body = String::from_utf8_lossy(&bytes);
|
||||
tracing::error!(error = %e, raw_body = %raw_body, "Response body is not JSON");
|
||||
SamplingError::Serialization(e)
|
||||
})?;
|
||||
// A `max` effort echo is unrepresentable in the typed enum — drop it
|
||||
// rather than failing the whole response.
|
||||
kigi_sampling_types::normalize_effort_echo(&mut response_value);
|
||||
let response_obj = serde_json::from_value::<rs::Response>(response_value).map_err(|e| {
|
||||
let raw_body = String::from_utf8_lossy(&bytes);
|
||||
tracing::error!(
|
||||
error = %e,
|
||||
@@ -1156,6 +1323,10 @@ impl SamplingClient {
|
||||
}
|
||||
}
|
||||
kigi_sampling_types::patch_reasoning_text_types(&mut request_body);
|
||||
kigi_sampling_types::patch_reasoning_effort(&mut request_body, request.reasoning_effort);
|
||||
if self.defaults.openai_codex {
|
||||
kigi_sampling_types::adapt_body_for_codex_backend(&mut request_body);
|
||||
}
|
||||
// Fresh per attempt so signals never leak across retries; `None`
|
||||
// (check disabled) sends no header and does no peek work per event.
|
||||
let doom_loop = self
|
||||
@@ -1314,6 +1485,15 @@ impl SamplingClient {
|
||||
|
||||
/// Apply default configuration to a Messages API request.
|
||||
fn apply_message_defaults(&self, request: &mut MessagesRequestWrapper) -> Result<()> {
|
||||
// Claude Pro/Max OAuth adaptation (claude-pro-max only): the OAuth
|
||||
// token is Claude-Code-scoped, so the request MUST lead with the exact
|
||||
// "You are Claude Code…" system block or Anthropic rejects it. Prepend
|
||||
// it as a distinct first system block, preserving any caller prompt.
|
||||
// Gated on `anthropic_oauth` so API-key anthropic/minimax are untouched.
|
||||
if self.defaults.anthropic_oauth {
|
||||
prepend_claude_code_system_prefix(&mut request.inner.system);
|
||||
}
|
||||
|
||||
// Apply model default if not specified
|
||||
if request.inner.model.is_empty() {
|
||||
request.inner.model = self.defaults.model.clone();
|
||||
@@ -1687,6 +1867,7 @@ impl SamplingClient {
|
||||
let responses_request: rs::CreateResponse = (&request).into();
|
||||
|
||||
let mut wrapper = CreateResponseWrapper::new(responses_request);
|
||||
wrapper.reasoning_effort = request.reasoning_effort;
|
||||
wrapper.x_kigi_conv_id = x_kigi_conv_id;
|
||||
wrapper.x_kigi_req_id = x_kigi_req_id;
|
||||
wrapper.x_kigi_session_id = x_kigi_session_id;
|
||||
@@ -1720,6 +1901,7 @@ impl SamplingClient {
|
||||
let responses_request: rs::CreateResponse = (&request).into();
|
||||
|
||||
let mut wrapper = CreateResponseWrapper::new(responses_request);
|
||||
wrapper.reasoning_effort = request.reasoning_effort;
|
||||
wrapper.x_kigi_conv_id = x_kigi_conv_id;
|
||||
wrapper.x_kigi_req_id = x_kigi_req_id;
|
||||
wrapper.x_kigi_session_id = x_kigi_session_id;
|
||||
@@ -1856,6 +2038,10 @@ mod tests {
|
||||
top_p: None,
|
||||
api_backend: ApiBackend::ChatCompletions,
|
||||
auth_scheme: AuthScheme::Bearer,
|
||||
anthropic_oauth: false,
|
||||
github_copilot: false,
|
||||
openai_codex: false,
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: IndexMap::new(),
|
||||
context_window: 8192,
|
||||
force_http1: false,
|
||||
@@ -1874,6 +2060,337 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
/// The real Anthropic Messages wire rejects requests without
|
||||
/// `anthropic-version`; the XApiKey+Messages client must carry it in its
|
||||
/// default headers, and Bearer/ChatCompletions clients must NOT.
|
||||
#[test]
|
||||
fn x_api_key_messages_client_sends_anthropic_version() {
|
||||
let mut config = minimal_config();
|
||||
config.auth_scheme = AuthScheme::XApiKey;
|
||||
config.api_backend = ApiBackend::Messages;
|
||||
let client = SamplingClient::new(config).expect("client builds");
|
||||
assert_eq!(
|
||||
client
|
||||
.default_headers
|
||||
.get("anthropic-version")
|
||||
.and_then(|v| v.to_str().ok()),
|
||||
Some(kigi_sampling_types::ANTHROPIC_VERSION)
|
||||
);
|
||||
assert!(client.default_headers.get("x-api-key").is_some());
|
||||
|
||||
let bearer = SamplingClient::new(minimal_config()).expect("client builds");
|
||||
assert!(
|
||||
bearer.default_headers.get("anthropic-version").is_none(),
|
||||
"non-anthropic clients must not grow the header"
|
||||
);
|
||||
let mut x_api_chat = minimal_config();
|
||||
x_api_chat.auth_scheme = AuthScheme::XApiKey;
|
||||
let x_api_chat = SamplingClient::new(x_api_chat).expect("client builds");
|
||||
assert!(
|
||||
x_api_chat
|
||||
.default_headers
|
||||
.get("anthropic-version")
|
||||
.is_none(),
|
||||
"XApiKey without the Messages backend must not grow the header"
|
||||
);
|
||||
}
|
||||
|
||||
/// Claude Pro/Max OAuth Messages client (`anthropic_oauth = true`, Bearer,
|
||||
/// Messages) carries the full OAuth identity: Bearer auth, anthropic-version,
|
||||
/// the oauth `anthropic-beta`, `x-app: cli`, the claude-cli User-Agent, and
|
||||
/// the direct-browser-access header.
|
||||
#[test]
|
||||
fn anthropic_oauth_messages_client_sends_oauth_identity_headers() {
|
||||
let mut config = minimal_config();
|
||||
config.api_key = Some("sk-ant-oat-secret".to_string());
|
||||
config.auth_scheme = AuthScheme::Bearer;
|
||||
config.api_backend = ApiBackend::Messages;
|
||||
config.anthropic_oauth = true;
|
||||
let client = SamplingClient::new(config).expect("client builds");
|
||||
let h = &client.default_headers;
|
||||
assert_eq!(
|
||||
h.get(AUTHORIZATION).and_then(|v| v.to_str().ok()),
|
||||
Some("Bearer sk-ant-oat-secret"),
|
||||
"OAuth path rides Authorization: Bearer, never x-api-key"
|
||||
);
|
||||
assert!(
|
||||
h.get("x-api-key").is_none(),
|
||||
"OAuth path must not send x-api-key"
|
||||
);
|
||||
assert_eq!(
|
||||
h.get("anthropic-version").and_then(|v| v.to_str().ok()),
|
||||
Some(kigi_sampling_types::ANTHROPIC_VERSION)
|
||||
);
|
||||
assert_eq!(
|
||||
h.get("anthropic-beta").and_then(|v| v.to_str().ok()),
|
||||
Some(kigi_sampling_types::ANTHROPIC_OAUTH_BETA)
|
||||
);
|
||||
assert_eq!(h.get("x-app").and_then(|v| v.to_str().ok()), Some("cli"));
|
||||
assert_eq!(
|
||||
h.get(USER_AGENT).and_then(|v| v.to_str().ok()),
|
||||
Some(kigi_sampling_types::CLAUDE_CODE_USER_AGENT)
|
||||
);
|
||||
assert_eq!(
|
||||
h.get("anthropic-dangerous-direct-browser-access")
|
||||
.and_then(|v| v.to_str().ok()),
|
||||
Some("true")
|
||||
);
|
||||
}
|
||||
|
||||
/// REGRESSION: an API-key Anthropic Messages client (XApiKey, NOT oauth)
|
||||
/// carries NONE of the OAuth identity — no anthropic-beta, no x-app, and
|
||||
/// the kigi User-Agent — so the API-key path stays byte-identical.
|
||||
#[test]
|
||||
fn api_key_anthropic_messages_client_has_no_oauth_identity() {
|
||||
let mut config = minimal_config();
|
||||
config.auth_scheme = AuthScheme::XApiKey;
|
||||
config.api_backend = ApiBackend::Messages;
|
||||
// anthropic_oauth stays false.
|
||||
let client = SamplingClient::new(config).expect("client builds");
|
||||
let h = &client.default_headers;
|
||||
assert!(
|
||||
h.get("anthropic-beta").is_none(),
|
||||
"API-key anthropic must NOT send the oauth beta"
|
||||
);
|
||||
assert!(
|
||||
h.get("x-app").is_none(),
|
||||
"API-key anthropic must NOT send x-app"
|
||||
);
|
||||
assert!(h.get("anthropic-dangerous-direct-browser-access").is_none());
|
||||
assert!(
|
||||
h.get(USER_AGENT)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.is_some_and(|ua| ua.starts_with("kigi/")),
|
||||
"API-key anthropic keeps the kigi User-Agent"
|
||||
);
|
||||
}
|
||||
|
||||
/// GitHub Copilot ChatCompletions client (`github_copilot = true`) carries
|
||||
/// the VS Code Copilot editor-identity headers + `X-Initiator: user` and
|
||||
/// presents the Copilot User-Agent (overriding the kigi UA).
|
||||
#[test]
|
||||
fn github_copilot_client_sends_editor_identity_headers() {
|
||||
let mut config = minimal_config();
|
||||
config.github_copilot = true;
|
||||
let client = SamplingClient::new(config).expect("client builds");
|
||||
let h = &client.default_headers;
|
||||
assert_eq!(
|
||||
h.get(USER_AGENT).and_then(|v| v.to_str().ok()),
|
||||
Some(kigi_sampling_types::COPILOT_USER_AGENT),
|
||||
"Copilot presents the VS Code Copilot User-Agent"
|
||||
);
|
||||
assert_eq!(
|
||||
h.get("editor-version").and_then(|v| v.to_str().ok()),
|
||||
Some(kigi_sampling_types::COPILOT_EDITOR_VERSION)
|
||||
);
|
||||
assert_eq!(
|
||||
h.get("editor-plugin-version").and_then(|v| v.to_str().ok()),
|
||||
Some(kigi_sampling_types::COPILOT_EDITOR_PLUGIN_VERSION)
|
||||
);
|
||||
assert_eq!(
|
||||
h.get("copilot-integration-id")
|
||||
.and_then(|v| v.to_str().ok()),
|
||||
Some(kigi_sampling_types::COPILOT_INTEGRATION_ID)
|
||||
);
|
||||
assert_eq!(
|
||||
h.get("x-initiator").and_then(|v| v.to_str().ok()),
|
||||
Some("user"),
|
||||
"inference carries X-Initiator: user"
|
||||
);
|
||||
}
|
||||
|
||||
/// SECURITY REGRESSION: `chatgpt-account-id` (the Codex per-user account
|
||||
/// identifier) must be redacted by the header loggers. It is not a token, so
|
||||
/// the substring list has to name it explicitly — without this, one debug
|
||||
/// request log would write the user's stable ChatGPT account id to disk.
|
||||
#[test]
|
||||
fn account_id_headers_are_redacted_from_logs() {
|
||||
assert!(
|
||||
SamplingClient::is_sensitive_header("chatgpt-account-id"),
|
||||
"chatgpt-account-id must be treated as sensitive"
|
||||
);
|
||||
assert!(
|
||||
SamplingClient::is_sensitive_header("ChatGPT-Account-Id"),
|
||||
"redaction is case-insensitive"
|
||||
);
|
||||
let rendered = SamplingClient::format_header("chatgpt-account-id", "acct-abc-123");
|
||||
assert!(
|
||||
rendered.contains("[REDACTED]") && !rendered.contains("acct-abc-123"),
|
||||
"the account id value must never appear verbatim, got {rendered:?}"
|
||||
);
|
||||
// Sanity: an ordinary header is still shown (redaction stays targeted).
|
||||
assert!(!SamplingClient::is_sensitive_header("content-type"));
|
||||
}
|
||||
|
||||
/// REGRESSION: a plain ChatCompletions client (github-copilot OFF, standing
|
||||
/// in for groq) carries NONE of the Copilot editor headers and keeps the
|
||||
/// kigi User-Agent — every other ChatCompletions provider stays untouched.
|
||||
#[test]
|
||||
fn plain_chat_completions_client_has_no_copilot_editor_headers() {
|
||||
// github_copilot stays false (as it is for groq and every other
|
||||
// ChatCompletions platform).
|
||||
let client = SamplingClient::new(minimal_config()).expect("client builds");
|
||||
let h = &client.default_headers;
|
||||
assert!(
|
||||
h.get("editor-version").is_none(),
|
||||
"groq must NOT send Editor-Version"
|
||||
);
|
||||
assert!(
|
||||
h.get("editor-plugin-version").is_none(),
|
||||
"groq must NOT send Editor-Plugin-Version"
|
||||
);
|
||||
assert!(
|
||||
h.get("copilot-integration-id").is_none(),
|
||||
"groq must NOT send Copilot-Integration-Id"
|
||||
);
|
||||
assert!(
|
||||
h.get("x-initiator").is_none(),
|
||||
"groq must NOT send X-Initiator"
|
||||
);
|
||||
assert!(
|
||||
h.get(USER_AGENT)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.is_some_and(|ua| ua.starts_with("kigi/")),
|
||||
"groq keeps the kigi User-Agent"
|
||||
);
|
||||
}
|
||||
|
||||
// A JWT whose payload carries `["https://api.openai.com/auth"]
|
||||
// ["chatgpt_account_id"] = "acct-test-42"` (alg=none; signature is cosmetic
|
||||
// — the extractor only base64url-decodes the payload segment).
|
||||
const CODEX_TEST_JWT: &str = "eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0.eyJodHRwczovL2FwaS5vcGVuYWkuY29tL2F1dGgiOnsiY2hhdGdwdF9hY2NvdW50X2lkIjoiYWNjdC10ZXN0LTQyIn0sInN1YiI6InVzZXItMSJ9.sig";
|
||||
|
||||
/// openai-codex Responses client (`openai_codex = true`): the STATIC codex
|
||||
/// identity headers (`originator`, `OpenAI-Beta`, codex `User-Agent`) ride
|
||||
/// `default_headers`, and a built `post()` request derives `chatgpt-account-
|
||||
/// id` from the bearer JWT. The Responses body default keeps `store: false`.
|
||||
#[test]
|
||||
fn openai_codex_client_sends_codex_identity_headers() {
|
||||
let mut config = minimal_config();
|
||||
config.api_backend = ApiBackend::Responses;
|
||||
config.openai_codex = true;
|
||||
config.api_key = Some(CODEX_TEST_JWT.to_string());
|
||||
let client = SamplingClient::new(config).expect("client builds");
|
||||
let h = &client.default_headers;
|
||||
assert_eq!(
|
||||
h.get("originator").and_then(|v| v.to_str().ok()),
|
||||
Some(kigi_sampling_types::CODEX_ORIGINATOR),
|
||||
"codex must send originator: codex_cli_rs"
|
||||
);
|
||||
assert_eq!(
|
||||
h.get("openai-beta").and_then(|v| v.to_str().ok()),
|
||||
Some(kigi_sampling_types::CODEX_OPENAI_BETA),
|
||||
"codex must opt into OpenAI-Beta: responses=experimental"
|
||||
);
|
||||
assert_eq!(
|
||||
h.get(USER_AGENT).and_then(|v| v.to_str().ok()),
|
||||
Some(kigi_sampling_types::CODEX_USER_AGENT),
|
||||
"codex presents the codex User-Agent"
|
||||
);
|
||||
// The account id is derived PER REQUEST from the bearer JWT in post().
|
||||
let req = client
|
||||
.post("https://chatgpt.com/backend-api/codex/responses")
|
||||
.build()
|
||||
.expect("build request");
|
||||
assert_eq!(
|
||||
req.headers()
|
||||
.get("chatgpt-account-id")
|
||||
.and_then(|v| v.to_str().ok()),
|
||||
Some("acct-test-42"),
|
||||
"chatgpt-account-id must be decoded from the bearer JWT claim"
|
||||
);
|
||||
}
|
||||
|
||||
/// REGRESSION: an API-key `openai` Responses client (`openai_codex = false`)
|
||||
/// carries NONE of the Codex identity headers — not the static ones and not
|
||||
/// the per-request `chatgpt-account-id` — so its request stays byte-identical.
|
||||
#[test]
|
||||
fn api_key_openai_responses_client_has_no_codex_headers() {
|
||||
let mut config = minimal_config();
|
||||
config.api_backend = ApiBackend::Responses;
|
||||
// openai_codex stays false (as it is for API-key openai). Even with a
|
||||
// JWT-shaped key, no account-id header is derived.
|
||||
config.api_key = Some(CODEX_TEST_JWT.to_string());
|
||||
let client = SamplingClient::new(config).expect("client builds");
|
||||
let h = &client.default_headers;
|
||||
assert!(
|
||||
h.get("originator").is_none(),
|
||||
"openai must NOT send originator"
|
||||
);
|
||||
assert!(
|
||||
h.get("openai-beta").is_none(),
|
||||
"openai must NOT send the codex OpenAI-Beta"
|
||||
);
|
||||
assert!(
|
||||
h.get(USER_AGENT)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.is_some_and(|ua| ua.starts_with("kigi/")),
|
||||
"API-key openai keeps the kigi User-Agent"
|
||||
);
|
||||
let req = client
|
||||
.post("https://api.openai.com/v1/responses")
|
||||
.build()
|
||||
.expect("build request");
|
||||
assert!(
|
||||
req.headers().get("chatgpt-account-id").is_none(),
|
||||
"API-key openai must NOT send chatgpt-account-id (regression)"
|
||||
);
|
||||
}
|
||||
|
||||
/// The system-prompt prefix is prepended as a distinct leading `text`
|
||||
/// block for each `system` shape (absent / string / blocks), preserving the
|
||||
/// caller's prompt, and is idempotent (not stamped twice).
|
||||
#[test]
|
||||
fn claude_code_system_prefix_prepends_and_is_idempotent() {
|
||||
use messages::{SystemParam, TextBlock};
|
||||
let prefix = kigi_sampling_types::CLAUDE_CODE_SYSTEM_PREFIX;
|
||||
|
||||
// Absent system → a single prefix block.
|
||||
let mut none = None;
|
||||
prepend_claude_code_system_prefix(&mut none);
|
||||
match none {
|
||||
Some(SystemParam::Blocks(b)) => {
|
||||
assert_eq!(b.len(), 1);
|
||||
assert_eq!(b[0].text, prefix);
|
||||
}
|
||||
other => panic!("expected one prefix block, got {other:?}"),
|
||||
}
|
||||
|
||||
// String system → [prefix, original].
|
||||
let mut text = Some(SystemParam::Text("do the thing".into()));
|
||||
prepend_claude_code_system_prefix(&mut text);
|
||||
match text {
|
||||
Some(SystemParam::Blocks(b)) => {
|
||||
assert_eq!(b.len(), 2);
|
||||
assert_eq!(b[0].text, prefix);
|
||||
assert_eq!(b[1].text, "do the thing");
|
||||
}
|
||||
other => panic!("expected two blocks, got {other:?}"),
|
||||
}
|
||||
|
||||
// Idempotent: a leading prefix block is not re-added.
|
||||
let mut already = Some(SystemParam::Blocks(vec![
|
||||
TextBlock {
|
||||
r#type: "text".into(),
|
||||
text: prefix.to_string(),
|
||||
cache_control: None,
|
||||
},
|
||||
TextBlock {
|
||||
r#type: "text".into(),
|
||||
text: "tail".into(),
|
||||
cache_control: None,
|
||||
},
|
||||
]));
|
||||
prepend_claude_code_system_prefix(&mut already);
|
||||
match already {
|
||||
Some(SystemParam::Blocks(b)) => {
|
||||
assert_eq!(b.len(), 2, "prefix must not be duplicated");
|
||||
assert_eq!(b[0].text, prefix);
|
||||
}
|
||||
other => panic!("expected unchanged blocks, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Verify the serialized shape of StreamingChatRequest matches the
|
||||
/// expected wire format: all ChatCompletionRequest fields flattened at
|
||||
/// top level, plus `stream: true` and `stream_options.include_usage: true`.
|
||||
@@ -2017,6 +2534,7 @@ mod tests {
|
||||
api_key: Some("bearer-key-abc123".to_string()),
|
||||
api_backend: ApiBackend::Messages,
|
||||
auth_scheme: AuthScheme::Bearer,
|
||||
chat_compat: Default::default(),
|
||||
..minimal_config()
|
||||
};
|
||||
let client = SamplingClient::new(cfg).expect("client should build");
|
||||
@@ -2183,6 +2701,7 @@ mod tests {
|
||||
api_key: Some("stale-bearer".to_string()),
|
||||
api_backend: ApiBackend::Messages,
|
||||
auth_scheme: AuthScheme::Bearer,
|
||||
chat_compat: Default::default(),
|
||||
bearer_resolver: Some(std::sync::Arc::new(StaticBearerResolver("fresh-bearer"))),
|
||||
..minimal_config()
|
||||
};
|
||||
@@ -2211,6 +2730,7 @@ mod tests {
|
||||
api_key: Some("stale-bearer".to_string()),
|
||||
api_backend: ApiBackend::Responses,
|
||||
auth_scheme: AuthScheme::Bearer,
|
||||
chat_compat: Default::default(),
|
||||
bearer_resolver: Some(std::sync::Arc::new(StaticBearerResolver("fresh-bearer"))),
|
||||
..minimal_config()
|
||||
};
|
||||
|
||||
@@ -33,7 +33,7 @@ pub enum AuthScheme {
|
||||
///
|
||||
/// `SamplerConfig` is the single source of truth for sampler
|
||||
/// configuration. The shell builds it directly (see
|
||||
/// `agent::config::resolve_model_to_sampling_config` and
|
||||
/// `agent::config::sampling_config_for_model` and
|
||||
/// `session::acp_session::SessionActor::reconstruct_full_config`) by
|
||||
/// composing chat-state's `kigi_sampling_types::SamplingConfig`
|
||||
/// with `Credentials` (api key, client version).
|
||||
@@ -55,6 +55,29 @@ pub struct SamplerConfig {
|
||||
pub api_backend: ApiBackend,
|
||||
#[serde(default)]
|
||||
pub auth_scheme: AuthScheme,
|
||||
/// Claude Pro/Max OAuth adaptation (claude-pro-max only). When true the
|
||||
/// Messages request carries the OAuth identity headers (`anthropic-beta`
|
||||
/// oauth, `claude-cli` User-Agent, `x-app: cli`) and its system prompt is
|
||||
/// prefixed with the required "You are Claude Code…" line. Gated so the
|
||||
/// API-key `anthropic` + `minimax` Messages requests stay byte-identical.
|
||||
#[serde(default)]
|
||||
pub anthropic_oauth: bool,
|
||||
/// GitHub Copilot ChatCompletions adaptation (github-copilot only). When
|
||||
/// true the request carries the VS Code Copilot editor-identity headers
|
||||
/// (User-Agent `GitHubCopilotChat/…`, `Editor-Version`,
|
||||
/// `Editor-Plugin-Version`, `Copilot-Integration-Id`) plus `X-Initiator:
|
||||
/// user`. Gated so every other ChatCompletions provider (groq, …) stays
|
||||
/// byte-identical.
|
||||
#[serde(default)]
|
||||
pub github_copilot: bool,
|
||||
/// ChatGPT/Codex Responses adaptation (openai-codex only). When true the
|
||||
/// `/codex/responses` request carries the Codex identity headers
|
||||
/// (`chatgpt-account-id` derived per-request from the bearer JWT,
|
||||
/// `originator: codex_cli_rs`, `OpenAI-Beta: responses=experimental`, a codex
|
||||
/// `User-Agent`). Gated so the API-key `openai` Responses requests stay
|
||||
/// byte-identical (`store: false` is already the shared Responses default).
|
||||
#[serde(default)]
|
||||
pub openai_codex: bool,
|
||||
/// Extra request headers applied verbatim. The sampler never inspects
|
||||
/// the URL to derive headers; callers (the session) inject proxy auth
|
||||
/// and other access headers here before constructing the config.
|
||||
@@ -70,6 +93,11 @@ pub struct SamplerConfig {
|
||||
|
||||
// Reasoning effort
|
||||
pub reasoning_effort: Option<ReasoningEffort>,
|
||||
/// ChatCompletions body-adaptation dialect (per-platform; BYOK/custom
|
||||
/// endpoints default to the historical Kimi behavior; lenient default on
|
||||
/// deserialize so persisted configs from before the field parse).
|
||||
#[serde(default)]
|
||||
pub chat_compat: kigi_sampling_types::ChatCompat,
|
||||
|
||||
/// Client identity for the User-Agent header (`kigi/{version}` plus an
|
||||
/// optional origin product). The old xAI proxy's identity headers
|
||||
@@ -135,9 +163,13 @@ impl Default for SamplerConfig {
|
||||
model: String::new(),
|
||||
max_completion_tokens: None,
|
||||
temperature: None,
|
||||
chat_compat: kigi_sampling_types::ChatCompat::default(),
|
||||
top_p: None,
|
||||
api_backend: ApiBackend::default(),
|
||||
auth_scheme: AuthScheme::default(),
|
||||
anthropic_oauth: false,
|
||||
github_copilot: false,
|
||||
openai_codex: false,
|
||||
extra_headers: IndexMap::new(),
|
||||
context_window: 0,
|
||||
force_http1: false,
|
||||
|
||||
@@ -29,6 +29,89 @@ pub(crate) fn adapt_chat_completions_body(body: &mut Value) {
|
||||
adapt_tool_schemas(body);
|
||||
}
|
||||
|
||||
/// Dialect-dispatched body adaptation. Kimi keeps the full historical
|
||||
/// pipeline (thinking + message hygiene + schema normalization — all built
|
||||
/// for the Kimi wire's strictness); DeepSeek differs ONLY in how thinking
|
||||
/// rides the body; Passthrough providers take OpenAI-style bodies verbatim
|
||||
/// (their `reasoning_effort` scalar is already the wire form).
|
||||
pub(crate) fn adapt_chat_completions_body_for(
|
||||
compat: kigi_sampling_types::ChatCompat,
|
||||
body: &mut Value,
|
||||
) {
|
||||
match compat {
|
||||
kigi_sampling_types::ChatCompat::Kimi => adapt_chat_completions_body(body),
|
||||
kigi_sampling_types::ChatCompat::DeepSeek => {
|
||||
adapt_thinking_deepseek(body);
|
||||
strip_kigi_private_message_fields(body);
|
||||
}
|
||||
kigi_sampling_types::ChatCompat::Passthrough => {
|
||||
strip_kigi_private_message_fields(body);
|
||||
}
|
||||
kigi_sampling_types::ChatCompat::StrictOpenAi => {
|
||||
strip_kigi_private_message_fields(body);
|
||||
strip_stream_options(body);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Mistral's strict Pydantic validator 422-rejects `stream_options`
|
||||
/// (`extra_forbidden` on `stream_options.include_usage`; its request model
|
||||
/// has no such field). kigi injects `stream_options.include_usage` on every
|
||||
/// streaming request for the other providers, so strip the whole object for
|
||||
/// Mistral. Streaming usage falls back to token estimation (as for any
|
||||
/// provider that omits streaming usage).
|
||||
fn strip_stream_options(body: &mut Value) {
|
||||
if let Some(obj) = body.as_object_mut() {
|
||||
obj.remove("stream_options");
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove kigi-internal history artifacts from input messages before they
|
||||
/// reach a non-Kimi wire. `reasoning_content` is Kimi's replayed-thinking
|
||||
/// field (Kimi consumes it; DeepSeek documents it as prefix-mode-only and
|
||||
/// historically 400s on it; other providers don't know it) and `model_id`
|
||||
/// is kigi's private per-message provenance. Kimi's own pipeline handles
|
||||
/// these in `adapt_messages`.
|
||||
fn strip_kigi_private_message_fields(body: &mut Value) {
|
||||
let Some(messages) = body.get_mut("messages").and_then(|m| m.as_array_mut()) else {
|
||||
return;
|
||||
};
|
||||
for message in messages {
|
||||
if let Some(obj) = message.as_object_mut() {
|
||||
obj.remove("reasoning_content");
|
||||
obj.remove("model_id");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// DeepSeek spells the thinking control `thinking:{type, reasoning_effort}`
|
||||
/// (api-docs.deepseek.com create-chat-completion; the server maps
|
||||
/// low/medium→high and xhigh→max itself, so the canonical level passes
|
||||
/// through verbatim). `none` disables thinking; absent leaves the server
|
||||
/// default (enabled).
|
||||
fn adapt_thinking_deepseek(body: &mut Value) {
|
||||
let Some(obj) = body.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
let Some(effort) = obj.remove("reasoning_effort") else {
|
||||
return;
|
||||
};
|
||||
let Some(level) = effort.as_str().map(str::to_owned) else {
|
||||
return;
|
||||
};
|
||||
if level == "none" {
|
||||
obj.insert(
|
||||
"thinking".to_string(),
|
||||
serde_json::json!({ "type": "disabled" }),
|
||||
);
|
||||
} else {
|
||||
obj.insert(
|
||||
"thinking".to_string(),
|
||||
serde_json::json!({ "type": "enabled", "reasoning_effort": level }),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Map the OpenAI-style `reasoning_effort` knob onto Kimi's `thinking`
|
||||
/// request field and drop `reasoning_effort` from the wire.
|
||||
///
|
||||
@@ -280,6 +363,123 @@ mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn deepseek_dialect_spells_thinking_reasoning_effort() {
|
||||
use kigi_sampling_types::ChatCompat;
|
||||
// Official docs: thinking:{type, reasoning_effort}; server maps
|
||||
// low/medium→high, xhigh→max itself — levels pass through verbatim.
|
||||
let mut body = json!({ "model": "deepseek-v4-pro", "reasoning_effort": "high" });
|
||||
adapt_chat_completions_body_for(ChatCompat::DeepSeek, &mut body);
|
||||
assert_eq!(body.get("reasoning_effort"), None);
|
||||
assert_eq!(
|
||||
body["thinking"],
|
||||
json!({ "type": "enabled", "reasoning_effort": "high" })
|
||||
);
|
||||
|
||||
let mut body = json!({ "model": "deepseek-v4-flash", "reasoning_effort": "max" });
|
||||
adapt_chat_completions_body_for(ChatCompat::DeepSeek, &mut body);
|
||||
assert_eq!(
|
||||
body["thinking"],
|
||||
json!({ "type": "enabled", "reasoning_effort": "max" })
|
||||
);
|
||||
|
||||
// none disables; absent leaves the server default (no thinking key).
|
||||
let mut body = json!({ "reasoning_effort": "none" });
|
||||
adapt_chat_completions_body_for(ChatCompat::DeepSeek, &mut body);
|
||||
assert_eq!(body["thinking"], json!({ "type": "disabled" }));
|
||||
let mut body = json!({ "model": "deepseek-chat" });
|
||||
adapt_chat_completions_body_for(ChatCompat::DeepSeek, &mut body);
|
||||
assert_eq!(body.get("thinking"), None);
|
||||
|
||||
// DeepSeek does NOT get kimi's message/tool-schema rewrites (empty
|
||||
// assistant tool-call content survives — DeepSeek's documented
|
||||
// function-calling round-trip uses that shape), but kigi-private
|
||||
// fields are stripped: replayed reasoning_content is prefix-mode-only
|
||||
// on the DeepSeek wire (historically a 400 in input messages).
|
||||
let mut body = json!({
|
||||
"reasoning_effort": "high",
|
||||
"messages": [
|
||||
{ "role": "assistant", "content": "", "tool_calls": [{}],
|
||||
"reasoning_content": "replayed thinking", "model_id": "kigi/x" }
|
||||
]
|
||||
});
|
||||
adapt_chat_completions_body_for(ChatCompat::DeepSeek, &mut body);
|
||||
assert_eq!(body["messages"][0]["content"], json!(""));
|
||||
assert_eq!(body["messages"][0].get("reasoning_content"), None);
|
||||
assert_eq!(body["messages"][0].get("model_id"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strict_openai_dialect_strips_stream_options_and_private_fields() {
|
||||
use kigi_sampling_types::ChatCompat;
|
||||
// Mistral 422s on stream_options (extra_forbidden) and doesn't know
|
||||
// kigi's private message fields; OpenAI-style reasoning_effort stays.
|
||||
let mut body = json!({
|
||||
"model": "mistral-medium-latest",
|
||||
"reasoning_effort": "high",
|
||||
"stream": true,
|
||||
"stream_options": { "include_usage": true },
|
||||
"messages": [
|
||||
{ "role": "assistant", "content": "hi",
|
||||
"reasoning_content": "internal", "model_id": "kigi/x" }
|
||||
]
|
||||
});
|
||||
adapt_chat_completions_body_for(ChatCompat::StrictOpenAi, &mut body);
|
||||
assert_eq!(
|
||||
body.get("stream_options"),
|
||||
None,
|
||||
"stream_options must be stripped"
|
||||
);
|
||||
assert_eq!(body["stream"], json!(true), "stream flag stays");
|
||||
assert_eq!(
|
||||
body["reasoning_effort"],
|
||||
json!("high"),
|
||||
"OpenAI-style effort passes through (Mistral accepts it natively)"
|
||||
);
|
||||
assert_eq!(body["messages"][0].get("reasoning_content"), None);
|
||||
assert_eq!(body["messages"][0].get("model_id"), None);
|
||||
assert_eq!(body["messages"][0]["content"], json!("hi"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn passthrough_dialect_leaves_openai_body_verbatim() {
|
||||
use kigi_sampling_types::ChatCompat;
|
||||
// Verbatim EXCEPT kigi-private history artifacts, which no non-Kimi
|
||||
// wire understands.
|
||||
let mut body = json!({
|
||||
"model": "gpt-oss",
|
||||
"reasoning_effort": "high",
|
||||
"messages": [
|
||||
{ "role": "user", "content": "hi" },
|
||||
{ "role": "assistant", "content": "yo",
|
||||
"reasoning_content": "internal", "model_id": "kigi/x" }
|
||||
]
|
||||
});
|
||||
adapt_chat_completions_body_for(ChatCompat::Passthrough, &mut body);
|
||||
assert_eq!(
|
||||
body,
|
||||
json!({
|
||||
"model": "gpt-oss",
|
||||
"reasoning_effort": "high",
|
||||
"messages": [
|
||||
{ "role": "user", "content": "hi" },
|
||||
{ "role": "assistant", "content": "yo" }
|
||||
]
|
||||
}),
|
||||
"reasoning_effort stays OpenAI-style; private fields are stripped"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kimi_dialect_dispatch_matches_legacy_pipeline() {
|
||||
use kigi_sampling_types::ChatCompat;
|
||||
let mut via_dispatch = json!({ "model": "k3", "reasoning_effort": "max" });
|
||||
adapt_chat_completions_body_for(ChatCompat::Kimi, &mut via_dispatch);
|
||||
let mut via_legacy = json!({ "model": "k3", "reasoning_effort": "max" });
|
||||
adapt_chat_completions_body(&mut via_legacy);
|
||||
assert_eq!(via_dispatch, via_legacy, "Kimi dispatch = legacy pipeline");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reasoning_effort_maps_to_kimi_thinking_field() {
|
||||
// Level rides along as thinking.effort (live wire: 200 with
|
||||
@@ -292,8 +492,9 @@ mod tests {
|
||||
json!({ "type": "enabled", "effort": "high" })
|
||||
);
|
||||
|
||||
// Canonical `xhigh` is spelled `max` on the Kimi wire (the K3
|
||||
// valid_efforts vocabulary is low/high/max).
|
||||
// Legacy canonical `xhigh` (pre-Max configs/sessions) is spelled
|
||||
// `max` on the Kimi wire (the K3 valid_efforts vocabulary is
|
||||
// low/high/max — there is no `xhigh` there).
|
||||
let mut body = json!({ "model": "k3", "reasoning_effort": "xhigh" });
|
||||
adapt_chat_completions_body(&mut body);
|
||||
assert_eq!(
|
||||
@@ -301,6 +502,15 @@ mod tests {
|
||||
json!({ "type": "enabled", "effort": "max" })
|
||||
);
|
||||
|
||||
// Canonical `max` (what the K3 menu token parses to since the
|
||||
// ReasoningEffort::Max split) passes through unchanged.
|
||||
let mut body = json!({ "model": "k3", "reasoning_effort": "max" });
|
||||
adapt_chat_completions_body(&mut body);
|
||||
assert_eq!(
|
||||
body["thinking"],
|
||||
json!({ "type": "enabled", "effort": "max" })
|
||||
);
|
||||
|
||||
// kimi.py:218: "off" (our ReasoningEffort::None) → disabled, and no
|
||||
// effort key (a disabled+effort combination would be contradictory).
|
||||
let mut body = json!({ "reasoning_effort": "none" });
|
||||
|
||||
@@ -517,6 +517,85 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
/// Deserialize a full chunk from a JSON `delta` so it flows through the
|
||||
/// `#[serde(from = "RawChatChunkDelta")]` content-split path (Mistral
|
||||
/// sends array `content`).
|
||||
fn chunk_from_delta(delta: serde_json::Value) -> ChatCompletionChunk {
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"id": "c",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 0,
|
||||
"model": "mistral-medium-latest",
|
||||
"choices": [{ "index": 0, "delta": delta, "finish_reason": null }],
|
||||
}))
|
||||
.expect("mistral chunk deserializes")
|
||||
}
|
||||
|
||||
/// End-to-end: a Mistral reasoning stream (array `content` thinking
|
||||
/// deltas → a transition delta carrying both a thinking and a text chunk
|
||||
/// → plain-string answer deltas) is split by deserialization and consumed
|
||||
/// into the SAME reasoning-sibling + assistant-answer result the
|
||||
/// `reasoning_content` string path produces. Closes the array-path
|
||||
/// integration gap.
|
||||
#[tokio::test]
|
||||
async fn mistral_reasoning_array_stream_splits_into_reasoning_and_answer() {
|
||||
let chunks: Vec<Result<ChatCompletionChunk, SamplingError>> = vec![
|
||||
Ok(chunk_from_delta(serde_json::json!({ "content": [
|
||||
{ "type": "thinking",
|
||||
"thinking": [{ "type": "text", "text": "Let me think. " }] }
|
||||
]}))),
|
||||
Ok(chunk_from_delta(serde_json::json!({ "content": [
|
||||
{ "type": "thinking",
|
||||
"thinking": [{ "type": "text", "text": "It's 22." }] }
|
||||
]}))),
|
||||
// Transition: one array with a closing thinking chunk AND the
|
||||
// first answer text chunk.
|
||||
Ok(chunk_from_delta(serde_json::json!({ "content": [
|
||||
{ "type": "thinking", "thinking": [{ "type": "text", "text": " Done." }] },
|
||||
{ "type": "text", "text": "Answer: " }
|
||||
]}))),
|
||||
// Answer phase: plain-string deltas (no longer arrays).
|
||||
Ok(chunk_from_delta(serde_json::json!({ "content": "22." }))),
|
||||
Ok(final_chunk(FinishReason::Stop)),
|
||||
];
|
||||
let raw = stream::iter(chunks).boxed();
|
||||
let events = collect(stream_chat_completions(
|
||||
raw,
|
||||
None,
|
||||
rid(),
|
||||
Duration::from_secs(60),
|
||||
))
|
||||
.await;
|
||||
|
||||
// Channel tokens: thinking rode the Reasoning channel, answer the Text
|
||||
// channel — never crossed.
|
||||
let mut reasoning = String::new();
|
||||
let mut answer = String::new();
|
||||
for e in &events {
|
||||
if let SamplingEvent::ChannelToken { channel, text, .. } = e {
|
||||
match channel {
|
||||
SamplingChannel::Reasoning => reasoning.push_str(text),
|
||||
SamplingChannel::Text => answer.push_str(text),
|
||||
}
|
||||
}
|
||||
}
|
||||
assert_eq!(reasoning, "Let me think. It's 22. Done.");
|
||||
assert_eq!(answer, "Answer: 22.");
|
||||
|
||||
// The accumulated final response carries the same split.
|
||||
match events.last().unwrap() {
|
||||
SamplingEvent::Completed { response, .. } => {
|
||||
let r = response
|
||||
.reasoning_items()
|
||||
.next()
|
||||
.expect("array thinking became a reasoning sibling");
|
||||
let rs::SummaryPart::SummaryText(t) = &r.summary[0];
|
||||
assert_eq!(t.text, "Let me think. It's 22. Done.");
|
||||
}
|
||||
other => panic!("expected Completed, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_call_stream_emits_deltas_and_assembles_final_call() {
|
||||
// First chunk has id + name + part of arguments.
|
||||
|
||||
@@ -78,6 +78,10 @@ fn test_config(base_url: String, model: &str) -> SamplerConfig {
|
||||
top_p: None,
|
||||
api_backend: ApiBackend::ChatCompletions,
|
||||
auth_scheme: Default::default(),
|
||||
anthropic_oauth: false,
|
||||
github_copilot: false,
|
||||
openai_codex: false,
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: IndexMap::new(),
|
||||
context_window: 128_000,
|
||||
force_http1: false,
|
||||
|
||||
@@ -7,6 +7,7 @@ description = "Pure data types for the xAI sampling / chat-completion API layer"
|
||||
|
||||
[dependencies]
|
||||
async-openai = { workspace = true }
|
||||
base64 = { workspace = true }
|
||||
indexmap = { workspace = true, features = ["serde"] }
|
||||
reqwest = { workspace = true }
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
|
||||
@@ -250,8 +250,13 @@ pub struct AssistantItem {
|
||||
/// `response.reasoning.effort` (Responses API). Stored beside
|
||||
/// `model_id`/`model_fingerprint` so per-response effort survives
|
||||
/// mid-session model/effort switches. `None` for synthetic items and
|
||||
/// backends that don't echo it.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
/// backends that don't echo it. Lenient on read: an unknown token from a
|
||||
/// newer kigi drops to `None` rather than failing the history line.
|
||||
#[serde(
|
||||
default,
|
||||
skip_serializing_if = "Option::is_none",
|
||||
deserialize_with = "crate::types::lenient_reasoning_effort_opt"
|
||||
)]
|
||||
pub reasoning_effort: Option<crate::ReasoningEffort>,
|
||||
}
|
||||
|
||||
@@ -2153,7 +2158,10 @@ impl From<&ConversationRequest> for rs::CreateResponse {
|
||||
prompt_cache_key: None,
|
||||
prompt_cache_retention: None,
|
||||
reasoning: Some(rs::Reasoning {
|
||||
effort: req.reasoning_effort.map(|e| e.to_responses_api()),
|
||||
// Effort is written onto the serialized body by
|
||||
// `patch_reasoning_effort` (the typed enum cannot spell
|
||||
// `max`); the wrapper carries the canonical value.
|
||||
effort: None,
|
||||
summary: Some(rs::ReasoningSummary::Concise),
|
||||
}),
|
||||
safety_identifier: None,
|
||||
@@ -3193,6 +3201,8 @@ pub fn build_messages_request(req: &ConversationRequest) -> crate::messages::Mes
|
||||
flush_assistant(&mut pending_assistant, &mut messages);
|
||||
flush_tool_results(&mut pending_tool_results, &mut messages);
|
||||
|
||||
prune_replayed_thinking(&mut messages);
|
||||
|
||||
// Attach cache_control: {type: "ephemeral"} to last system block
|
||||
if let Some(last) = system_blocks.last_mut() {
|
||||
last.cache_control = Some(CacheControl {
|
||||
@@ -3282,6 +3292,74 @@ pub fn build_messages_request(req: &ConversationRequest) -> crate::messages::Mes
|
||||
}
|
||||
}
|
||||
|
||||
/// Strip replayed `thinking` blocks the Anthropic Messages API would
|
||||
/// reject — keep exactly the one it requires.
|
||||
///
|
||||
/// Anthropic validates EVERY `thinking` block in the request: the
|
||||
/// signature is bound to the emitting model and must be non-empty, so
|
||||
/// history from another backend (`encrypted_content: None` replays as
|
||||
/// `signature: ""`), a Responses-API `tco_*` blob (signature bytes with no
|
||||
/// text), or a block signed by a DIFFERENT model after a mid-session
|
||||
/// `/model` switch 400s the whole request with
|
||||
/// "messages.N.content.0: Invalid `signature` in `thinking` block".
|
||||
///
|
||||
/// The API only NEEDS thinking for the ACTIVE tool-use continuation: the
|
||||
/// final assistant message whose tool_use results follow must carry its
|
||||
/// signed thinking back verbatim. Prior turns' thinking is ignored even
|
||||
/// when valid (Pi/Claude Code replay exactly this way). So: keep the
|
||||
/// final assistant message's thinking when the loop is open and the block
|
||||
/// is genuinely signed (non-empty text AND signature — an open loop can
|
||||
/// never span a model switch, so that signature is always the current
|
||||
/// model's); strip every other thinking block. An assistant message left
|
||||
/// EMPTY by the strip (a thinking-only aborted turn) is removed — the API
|
||||
/// rejects empty content arrays.
|
||||
fn prune_replayed_thinking(messages: &mut Vec<crate::messages::Message>) {
|
||||
use crate::messages::{ContentBlock, MessageContent, MessageRole};
|
||||
let last_assistant = messages
|
||||
.iter()
|
||||
.rposition(|m| matches!(m.role, MessageRole::Assistant));
|
||||
let active_tool_loop = last_assistant.is_some_and(|i| {
|
||||
let has_tool_use = matches!(
|
||||
&messages[i].content,
|
||||
MessageContent::Blocks(blocks)
|
||||
if blocks.iter().any(|b| matches!(b, ContentBlock::ToolUse { .. }))
|
||||
);
|
||||
// The loop is OPEN only while the request ends on the tool results:
|
||||
// a later plain user turn closes it (the results answered, the model
|
||||
// replied — its thinking is history the API ignores or rejects).
|
||||
let continuation = &messages[i + 1..];
|
||||
let ends_on_results = !continuation.is_empty()
|
||||
&& continuation.iter().all(|m| {
|
||||
matches!(
|
||||
&m.content,
|
||||
MessageContent::Blocks(blocks)
|
||||
if blocks.iter().any(|b| matches!(b, ContentBlock::ToolResult { .. }))
|
||||
)
|
||||
});
|
||||
has_tool_use && ends_on_results
|
||||
});
|
||||
let mut index = 0;
|
||||
messages.retain_mut(|m| {
|
||||
let i = index;
|
||||
index += 1;
|
||||
if !matches!(m.role, MessageRole::Assistant) {
|
||||
return true;
|
||||
}
|
||||
let MessageContent::Blocks(blocks) = &mut m.content else {
|
||||
return true;
|
||||
};
|
||||
let keep_thinking = active_tool_loop && Some(i) == last_assistant;
|
||||
blocks.retain(|b| match b {
|
||||
ContentBlock::Thinking {
|
||||
thinking,
|
||||
signature,
|
||||
} => keep_thinking && !thinking.is_empty() && !signature.is_empty(),
|
||||
_ => true,
|
||||
});
|
||||
!blocks.is_empty()
|
||||
});
|
||||
}
|
||||
|
||||
/// Convert a MessagesResponse to a single Assistant `ConversationItem`.
|
||||
///
|
||||
/// Note: Anthropic `Thinking` blocks are dropped here because this `From`
|
||||
@@ -5076,13 +5154,35 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assistant_item_unknown_reasoning_effort_token_degrades_to_none() {
|
||||
// History lines written by a NEWER kigi (grown effort vocabulary)
|
||||
// must not fail this binary's history parse — the lenient
|
||||
// deserializer must be wired on the AssistantItem field itself.
|
||||
let json = serde_json::json!({
|
||||
"role": "assistant",
|
||||
"content": "hi",
|
||||
"reasoning_effort": "hypermax"
|
||||
});
|
||||
let item: AssistantItem = serde_json::from_value(json).expect("history line must survive");
|
||||
assert_eq!(item.reasoning_effort, None);
|
||||
let json = serde_json::json!({
|
||||
"role": "assistant",
|
||||
"content": "hi",
|
||||
"reasoning_effort": "max"
|
||||
});
|
||||
let item: AssistantItem = serde_json::from_value(json).unwrap();
|
||||
assert_eq!(item.reasoning_effort, Some(crate::ReasoningEffort::Max));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_messages_request_wire_format_for_supported_variants() {
|
||||
for (variant, expected) in [
|
||||
(crate::ReasoningEffort::Low, "low"),
|
||||
(crate::ReasoningEffort::Medium, "medium"),
|
||||
(crate::ReasoningEffort::High, "high"),
|
||||
(crate::ReasoningEffort::Xhigh, "max"),
|
||||
(crate::ReasoningEffort::Xhigh, "xhigh"),
|
||||
(crate::ReasoningEffort::Max, "max"),
|
||||
] {
|
||||
let req = messages_test_request(Some(variant));
|
||||
let msgs = build_messages_request(&req);
|
||||
@@ -5131,6 +5231,7 @@ mod tests {
|
||||
(crate::ReasoningEffort::Medium, "medium"),
|
||||
(crate::ReasoningEffort::High, "high"),
|
||||
(crate::ReasoningEffort::Xhigh, "xhigh"),
|
||||
(crate::ReasoningEffort::Max, "max"),
|
||||
] {
|
||||
let req = ConversationRequest::from_items(vec![ConversationItem::user("hi")])
|
||||
.with_model("test");
|
||||
@@ -5163,6 +5264,10 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_responses_request_carries_reasoning_effort_nested() {
|
||||
// Wire contract since the Max split: the typed conversion leaves
|
||||
// reasoning.effort UNSET (async-openai's enum cannot spell `max`);
|
||||
// `patch_reasoning_effort` writes the canonical token onto the
|
||||
// serialized body — every level, `xhigh` and `max` distinct.
|
||||
for (variant, expected) in [
|
||||
(crate::ReasoningEffort::None, "none"),
|
||||
(crate::ReasoningEffort::Minimal, "minimal"),
|
||||
@@ -5170,6 +5275,7 @@ mod tests {
|
||||
(crate::ReasoningEffort::Medium, "medium"),
|
||||
(crate::ReasoningEffort::High, "high"),
|
||||
(crate::ReasoningEffort::Xhigh, "xhigh"),
|
||||
(crate::ReasoningEffort::Max, "max"),
|
||||
] {
|
||||
let req = ConversationRequest {
|
||||
reasoning_effort: Some(variant),
|
||||
@@ -5177,15 +5283,56 @@ mod tests {
|
||||
.with_model("test")
|
||||
};
|
||||
let resp: crate::rs::CreateResponse = (&req).into();
|
||||
let json = serde_json::to_value(&resp).unwrap();
|
||||
let mut json = serde_json::to_value(&resp).unwrap();
|
||||
assert_eq!(
|
||||
json.pointer("/reasoning/effort"),
|
||||
None,
|
||||
"typed conversion must leave effort unset ({variant:?})"
|
||||
);
|
||||
crate::patch_reasoning_effort(&mut json, req.reasoning_effort);
|
||||
assert_eq!(
|
||||
json.pointer("/reasoning/effort").and_then(|v| v.as_str()),
|
||||
Some(expected),
|
||||
"{variant:?} should serialize as reasoning.effort={expected:?}; got: {json:#}",
|
||||
"{variant:?} should be patched as reasoning.effort={expected:?}; got: {json:#}",
|
||||
);
|
||||
assert_eq!(
|
||||
json.pointer("/reasoning/summary").and_then(|v| v.as_str()),
|
||||
Some("concise"),
|
||||
"the patch must not clobber reasoning.summary"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_effort_echo_drops_only_unrepresentable_tokens() {
|
||||
// `max` echo (bare response shape) is dropped so typed parsing
|
||||
// succeeds; known tokens pass through; stream-event envelopes are
|
||||
// handled via /response/reasoning.
|
||||
let mut bare = serde_json::json!({ "reasoning": { "effort": "max", "summary": "c" } });
|
||||
crate::normalize_effort_echo(&mut bare);
|
||||
assert_eq!(bare.pointer("/reasoning/effort"), None);
|
||||
assert_eq!(
|
||||
bare.pointer("/reasoning/summary").and_then(|v| v.as_str()),
|
||||
Some("c")
|
||||
);
|
||||
|
||||
let mut known = serde_json::json!({ "reasoning": { "effort": "high" } });
|
||||
crate::normalize_effort_echo(&mut known);
|
||||
assert_eq!(
|
||||
known.pointer("/reasoning/effort").and_then(|v| v.as_str()),
|
||||
Some("high"),
|
||||
"representable echoes must pass through"
|
||||
);
|
||||
|
||||
let mut event = serde_json::json!({ "response": { "reasoning": { "effort": "max" } } });
|
||||
crate::normalize_effort_echo(&mut event);
|
||||
assert_eq!(event.pointer("/response/reasoning/effort"), None);
|
||||
|
||||
let mut absent = serde_json::json!({ "output": [] });
|
||||
crate::normalize_effort_echo(&mut absent);
|
||||
assert_eq!(absent, serde_json::json!({ "output": [] }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_responses_request_omits_effort_when_unset() {
|
||||
let req =
|
||||
@@ -5255,6 +5402,160 @@ mod tests {
|
||||
/// messages while setting top-level `thinking: null` — the Messages API
|
||||
/// rejects this with a 400. Verify that stripped reasoning produces a
|
||||
/// valid request with no thinking blocks in messages.
|
||||
/// Collect `(message_index, thinking, signature)` for every thinking
|
||||
/// block in a built Messages request.
|
||||
fn thinking_blocks(json: &serde_json::Value) -> Vec<(usize, String, String)> {
|
||||
let mut out = Vec::new();
|
||||
for (i, m) in json["messages"].as_array().unwrap().iter().enumerate() {
|
||||
if let Some(content) = m.get("content").and_then(|c| c.as_array()) {
|
||||
for b in content {
|
||||
if b.get("type").and_then(|t| t.as_str()) == Some("thinking") {
|
||||
out.push((
|
||||
i,
|
||||
b["thinking"].as_str().unwrap_or_default().to_string(),
|
||||
b["signature"].as_str().unwrap_or_default().to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn reasoning(text: &str, encrypted: Option<&str>) -> ConversationItem {
|
||||
ConversationItem::Reasoning(rs::ReasoningItem {
|
||||
id: String::new(),
|
||||
summary: if text.is_empty() {
|
||||
vec![]
|
||||
} else {
|
||||
vec![rs::SummaryPart::SummaryText(rs::SummaryTextContent {
|
||||
text: text.to_string(),
|
||||
})]
|
||||
},
|
||||
content: None,
|
||||
encrypted_content: encrypted.map(str::to_owned),
|
||||
status: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn assistant_text(text: &str) -> ConversationItem {
|
||||
ConversationItem::Assistant(AssistantItem {
|
||||
content: text.into(),
|
||||
tool_calls: vec![],
|
||||
model_id: None,
|
||||
model_fingerprint: None,
|
||||
reasoning_effort: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// Anthropic validates EVERY replayed `thinking` block: an unsigned one
|
||||
/// (history synthesized by another backend replays as `signature: ""`),
|
||||
/// a Responses `tco_*` blob (signature with no text), or a block signed
|
||||
/// by a different model after a mid-session `/model` switch 400s the
|
||||
/// whole request — "messages.N.content.0: Invalid `signature` in
|
||||
/// `thinking` block". The API only NEEDS thinking for the active
|
||||
/// tool-use continuation, so outside one the builder must replay NO
|
||||
/// thinking blocks at all.
|
||||
#[test]
|
||||
fn messages_request_strips_thinking_outside_active_tool_loop() {
|
||||
let req = ConversationRequest::from_items(vec![
|
||||
ConversationItem::system("sys"),
|
||||
ConversationItem::user("q1"),
|
||||
// Cross-backend history: unsigned reasoning (the Windows repro —
|
||||
// session started on another model, then switched to Claude).
|
||||
reasoning("some thinking", None),
|
||||
assistant_text("a1"),
|
||||
ConversationItem::user("q2"),
|
||||
// Responses-API blob: signature-shaped bytes, no text.
|
||||
reasoning("", Some("tco_blob")),
|
||||
assistant_text("a2"),
|
||||
ConversationItem::user("q3"),
|
||||
// Genuinely signed — but its tool loop (none) is closed, so the
|
||||
// API ignores it when valid and 400s it after a model switch.
|
||||
reasoning("signed thinking", Some("sig-real")),
|
||||
assistant_text("a3"),
|
||||
ConversationItem::user("q4"),
|
||||
]);
|
||||
let json = serde_json::to_value(build_messages_request(&req)).unwrap();
|
||||
assert_eq!(
|
||||
thinking_blocks(&json),
|
||||
vec![],
|
||||
"no thinking block may be replayed outside an active tool loop:\n{json:#}"
|
||||
);
|
||||
}
|
||||
|
||||
/// The active tool-use continuation is the one place Anthropic REQUIRES
|
||||
/// the signed thinking block back: the final assistant message issued
|
||||
/// tool_use and its results follow. Exactly that block is kept; a
|
||||
/// prior turn's signed thinking is still stripped.
|
||||
#[test]
|
||||
fn messages_request_keeps_signed_thinking_for_active_tool_loop() {
|
||||
let req = ConversationRequest::from_items(vec![
|
||||
ConversationItem::user("q0"),
|
||||
reasoning("old turn", Some("sig-old")),
|
||||
assistant_text("a0"),
|
||||
ConversationItem::user("q1"),
|
||||
reasoning("current turn", Some("sig-current")),
|
||||
ConversationItem::Assistant(AssistantItem {
|
||||
content: "".into(),
|
||||
tool_calls: vec![ToolCall {
|
||||
id: std::sync::Arc::from("tc1"),
|
||||
name: "read_file".to_string(),
|
||||
arguments: std::sync::Arc::from("{}"),
|
||||
}],
|
||||
model_id: None,
|
||||
model_fingerprint: None,
|
||||
reasoning_effort: None,
|
||||
}),
|
||||
ConversationItem::tool_result("tc1", "file contents"),
|
||||
]);
|
||||
let json = serde_json::to_value(build_messages_request(&req)).unwrap();
|
||||
let blocks = thinking_blocks(&json);
|
||||
assert_eq!(
|
||||
blocks.len(),
|
||||
1,
|
||||
"exactly the active loop's thinking survives:\n{json:#}"
|
||||
);
|
||||
let (msg_idx, thinking, signature) = &blocks[0];
|
||||
assert_eq!(thinking, "current turn");
|
||||
assert_eq!(signature, "sig-current");
|
||||
// It sits at content.0 of the final assistant message.
|
||||
let msg = &json["messages"].as_array().unwrap()[*msg_idx];
|
||||
assert_eq!(msg["role"], "assistant");
|
||||
assert_eq!(msg["content"][0]["type"], "thinking");
|
||||
assert!(
|
||||
msg["content"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.any(|b| b["type"] == "tool_use"),
|
||||
"the kept thinking belongs to the tool_use turn"
|
||||
);
|
||||
}
|
||||
|
||||
/// A thinking-only assistant turn (aborted before any text/tool output)
|
||||
/// must not survive as an EMPTY assistant message after the strip —
|
||||
/// Anthropic rejects empty content arrays.
|
||||
#[test]
|
||||
fn messages_request_drops_assistant_message_emptied_by_thinking_strip() {
|
||||
let req = ConversationRequest::from_items(vec![
|
||||
ConversationItem::user("q"),
|
||||
reasoning("aborted turn thinking", Some("sig")),
|
||||
assistant_text(""),
|
||||
ConversationItem::user("follow-up"),
|
||||
]);
|
||||
let json = serde_json::to_value(build_messages_request(&req)).unwrap();
|
||||
for m in json["messages"].as_array().unwrap() {
|
||||
if let Some(content) = m.get("content").and_then(|c| c.as_array()) {
|
||||
assert!(
|
||||
!content.is_empty(),
|
||||
"no message may ship an empty content array:\n{json:#}"
|
||||
);
|
||||
}
|
||||
}
|
||||
assert_eq!(thinking_blocks(&json), vec![]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_btw_stripped_reasoning_produces_no_thinking_blocks() {
|
||||
// Simulate a conversation where the model responded with thinking.
|
||||
|
||||
@@ -487,7 +487,69 @@ pub enum FinishReason {
|
||||
FunctionCall,
|
||||
}
|
||||
|
||||
/// Wire `content` field on a chat response/delta. OpenAI-compatible providers
|
||||
/// send a plain string, but reasoning providers (Mistral) send an ARRAY of
|
||||
/// typed chunks: `{"type":"text","text":".."}` for the answer and
|
||||
/// `{"type":"thinking","thinking":[{"type":"text","text":".."}],...}` for the
|
||||
/// chain of thought (verified against the mistralai/client-python SDK models).
|
||||
/// The chunk union is OPEN — unknown chunk types are ignored, never fatal.
|
||||
#[derive(Deserialize)]
|
||||
#[serde(untagged)]
|
||||
enum WireContent {
|
||||
Text(String),
|
||||
Chunks(Vec<Value>),
|
||||
}
|
||||
|
||||
/// Read a chunk's `text` string field into `dst` when present (defensive: a
|
||||
/// reference/tool-reference/unknown chunk simply has no `text`).
|
||||
fn push_chunk_text(dst: &mut String, chunk: &Value) {
|
||||
if let Some(t) = chunk.get("text").and_then(Value::as_str) {
|
||||
dst.push_str(t);
|
||||
}
|
||||
}
|
||||
|
||||
fn non_empty(s: String) -> Option<String> {
|
||||
if s.is_empty() { None } else { Some(s) }
|
||||
}
|
||||
|
||||
impl WireContent {
|
||||
/// Split into `(visible answer text, thinking text)`. A plain string is
|
||||
/// the answer with no thinking; an array routes `text` chunks to the
|
||||
/// answer and the nested `text` of `thinking` chunks to the reasoning
|
||||
/// channel. Byte-identical for string content (the only shape non-Mistral
|
||||
/// providers ever send).
|
||||
fn split(self) -> (Option<String>, Option<String>) {
|
||||
match self {
|
||||
// Preserve string content verbatim (incl. empty) — byte-identical
|
||||
// to the pre-change `content: Option<String>` for every provider
|
||||
// that sends a string.
|
||||
WireContent::Text(s) => (Some(s), None),
|
||||
WireContent::Chunks(chunks) => {
|
||||
let mut answer = String::new();
|
||||
let mut thinking = String::new();
|
||||
for c in &chunks {
|
||||
match c.get("type").and_then(Value::as_str) {
|
||||
Some("text") => push_chunk_text(&mut answer, c),
|
||||
Some("thinking") => {
|
||||
// `thinking` is a nested list of chunks.
|
||||
if let Some(inner) = c.get("thinking").and_then(Value::as_array) {
|
||||
for ic in inner {
|
||||
push_chunk_text(&mut thinking, ic);
|
||||
}
|
||||
}
|
||||
}
|
||||
// reference / tool_reference / unknown → not text.
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
(non_empty(answer), non_empty(thinking))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
#[serde(from = "RawChatResponseMessage")]
|
||||
pub struct ChatResponseMessage {
|
||||
pub role: Role,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
@@ -502,6 +564,38 @@ pub struct ChatResponseMessage {
|
||||
pub citations: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
/// Deserialization mirror of [`ChatResponseMessage`] that accepts array
|
||||
/// `content` and folds thinking chunks into `reasoning_content`.
|
||||
#[derive(Deserialize)]
|
||||
struct RawChatResponseMessage {
|
||||
role: Role,
|
||||
#[serde(default)]
|
||||
content: Option<WireContent>,
|
||||
#[serde(default)]
|
||||
reasoning_content: Option<String>,
|
||||
#[serde(default)]
|
||||
tool_calls: Vec<ToolCallResponse>,
|
||||
#[serde(default)]
|
||||
tool_call_id: Option<String>,
|
||||
#[serde(default)]
|
||||
citations: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
impl From<RawChatResponseMessage> for ChatResponseMessage {
|
||||
fn from(r: RawChatResponseMessage) -> Self {
|
||||
let (answer, thinking) = r.content.map(WireContent::split).unwrap_or((None, None));
|
||||
ChatResponseMessage {
|
||||
role: r.role,
|
||||
content: answer,
|
||||
// The wire's own `reasoning_content` wins; else thinking chunks.
|
||||
reasoning_content: r.reasoning_content.or(thinking),
|
||||
tool_calls: r.tool_calls,
|
||||
tool_call_id: r.tool_call_id,
|
||||
citations: r.citations,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct ToolCallResponse {
|
||||
pub id: String,
|
||||
@@ -649,6 +743,7 @@ pub struct ToolCallFunctionDelta {
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone, Default)]
|
||||
#[serde(from = "RawChatChunkDelta")]
|
||||
pub struct ChatChunkDelta {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub role: Option<Role>,
|
||||
@@ -656,16 +751,43 @@ pub struct ChatChunkDelta {
|
||||
pub content: Option<String>,
|
||||
pub reasoning_content: Option<String>,
|
||||
/// Tool call deltas. Handles `null` in JSON as empty vec.
|
||||
#[serde(
|
||||
default,
|
||||
skip_serializing_if = "Vec::is_empty",
|
||||
deserialize_with = "deserialize_null_default"
|
||||
)]
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub tool_calls: Vec<ToolCallDelta>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_call_id: Option<String>,
|
||||
}
|
||||
|
||||
/// Deserialization mirror of [`ChatChunkDelta`] that accepts array `content`
|
||||
/// (Mistral reasoning streams a thinking-chunk array during the thinking
|
||||
/// phase, then plain-string answer deltas) and folds thinking into
|
||||
/// `reasoning_content`.
|
||||
#[derive(Deserialize)]
|
||||
struct RawChatChunkDelta {
|
||||
#[serde(default)]
|
||||
role: Option<Role>,
|
||||
#[serde(default)]
|
||||
content: Option<WireContent>,
|
||||
#[serde(default)]
|
||||
reasoning_content: Option<String>,
|
||||
#[serde(default, deserialize_with = "deserialize_null_default")]
|
||||
tool_calls: Vec<ToolCallDelta>,
|
||||
#[serde(default)]
|
||||
tool_call_id: Option<String>,
|
||||
}
|
||||
|
||||
impl From<RawChatChunkDelta> for ChatChunkDelta {
|
||||
fn from(r: RawChatChunkDelta) -> Self {
|
||||
let (answer, thinking) = r.content.map(WireContent::split).unwrap_or((None, None));
|
||||
ChatChunkDelta {
|
||||
role: r.role,
|
||||
content: answer,
|
||||
reasoning_content: r.reasoning_content.or(thinking),
|
||||
tool_calls: r.tool_calls,
|
||||
tool_call_id: r.tool_call_id,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parameters to control realtime data.
|
||||
#[derive(Serialize, Deserialize, Clone, Debug)]
|
||||
pub struct SearchParameters {
|
||||
@@ -787,22 +909,23 @@ pub enum ReasoningEffort {
|
||||
Medium,
|
||||
High,
|
||||
Xhigh,
|
||||
/// Distinct top tier above `xhigh` (OpenAI Responses and Anthropic
|
||||
/// Messages both accept `xhigh` AND `max` as separate levels in 2026;
|
||||
/// the Kimi wire spells its top tier `max` with no `xhigh`).
|
||||
Max,
|
||||
/// Codex-only top tier above `max` (the ChatGPT Codex backend exposes an
|
||||
/// `ultra` reasoning effort on its flagship models). Reachable ONLY via a
|
||||
/// model's server-declared effort menu (openai-codex); no built-in fallback
|
||||
/// menu offers it, so other providers never emit it.
|
||||
Ultra,
|
||||
}
|
||||
|
||||
impl ReasoningEffort {
|
||||
pub fn to_responses_api(self) -> crate::rs::ReasoningEffort {
|
||||
match self {
|
||||
Self::None => crate::rs::ReasoningEffort::None,
|
||||
Self::Minimal => crate::rs::ReasoningEffort::Minimal,
|
||||
Self::Low => crate::rs::ReasoningEffort::Low,
|
||||
Self::Medium => crate::rs::ReasoningEffort::Medium,
|
||||
Self::High => crate::rs::ReasoningEffort::High,
|
||||
Self::Xhigh => crate::rs::ReasoningEffort::Xhigh,
|
||||
}
|
||||
}
|
||||
|
||||
/// Inverse of [`to_responses_api`](Self::to_responses_api): the effort the
|
||||
/// Responses API echoes back on `response.reasoning.effort`.
|
||||
/// The canonical effort behind a typed Responses-API echo
|
||||
/// (`response.reasoning.effort`). `max` echoes never reach the typed
|
||||
/// enum — [`normalize_effort_echo`] drops them pre-parse (async-openai
|
||||
/// has no such variant); the request direction writes the wire string
|
||||
/// via [`patch_reasoning_effort`].
|
||||
pub fn from_responses_api(effort: crate::rs::ReasoningEffort) -> Self {
|
||||
match effort {
|
||||
crate::rs::ReasoningEffort::None => Self::None,
|
||||
@@ -822,17 +945,25 @@ impl ReasoningEffort {
|
||||
Self::Medium => "medium",
|
||||
Self::High => "high",
|
||||
Self::Xhigh => "xhigh",
|
||||
Self::Max => "max",
|
||||
Self::Ultra => "ultra",
|
||||
}
|
||||
}
|
||||
|
||||
/// Anthropic Messages API `output_config.effort` string; `None` for unsupported variants.
|
||||
/// Anthropic Messages API effort string; `None` for unsupported variants.
|
||||
/// `xhigh` and `max` are distinct levels on the 2026 Messages API (both
|
||||
/// appear in `GET /v1/models` `capabilities.effort`). `ultra` is codex-only
|
||||
/// and never selected on an Anthropic model, but maps to its own string for
|
||||
/// completeness (the Responses path writes effort via `as_str`, not this).
|
||||
pub fn to_messages_api(self) -> Option<&'static str> {
|
||||
match self {
|
||||
Self::None | Self::Minimal => None,
|
||||
Self::Low => Some("low"),
|
||||
Self::Medium => Some("medium"),
|
||||
Self::High => Some("high"),
|
||||
Self::Xhigh => Some("max"),
|
||||
Self::Xhigh => Some("xhigh"),
|
||||
Self::Max => Some("max"),
|
||||
Self::Ultra => Some("ultra"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -853,19 +984,277 @@ impl std::str::FromStr for ReasoningEffort {
|
||||
"low" => Ok(Self::Low),
|
||||
"medium" => Ok(Self::Medium),
|
||||
"high" => Ok(Self::High),
|
||||
"xhigh" | "max" => Ok(Self::Xhigh), // max is a CLI/UX alias of xhigh
|
||||
"xhigh" => Ok(Self::Xhigh),
|
||||
"max" => Ok(Self::Max),
|
||||
"ultra" => Ok(Self::Ultra),
|
||||
_ => Err(format!(
|
||||
"invalid reasoning effort: {s:?} (expected one of: none, minimal, low, medium, high, xhigh, max)"
|
||||
"invalid reasoning effort: {s:?} (expected one of: none, minimal, low, medium, high, xhigh, max, ultra)"
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Canonical wire parse only (`max` → `Xhigh`); remapped menu ids need a model catalog.
|
||||
/// Canonical wire parse; remapped menu ids need a model catalog.
|
||||
pub fn parse_canonical_effort_token(token: &str) -> Option<ReasoningEffort> {
|
||||
token.parse().ok()
|
||||
}
|
||||
|
||||
/// Deserialize an optional effort LENIENTLY for persisted stores (session
|
||||
/// summaries, chat history): a token this binary doesn't know — written by a
|
||||
/// newer kigi after the effort vocabulary grew, exactly what happened when
|
||||
/// `max` split from `xhigh` — degrades to `None` with a warning instead of
|
||||
/// failing the whole record, so listings and resume survive version
|
||||
/// rollback. Non-string values still error (real corruption stays loud), and
|
||||
/// config-TOML parsing stays strict — its layer warn-skips explicitly.
|
||||
pub fn lenient_reasoning_effort_opt<'de, D>(d: D) -> Result<Option<ReasoningEffort>, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
let raw = Option::<String>::deserialize(d)?;
|
||||
Ok(raw.and_then(|s| match s.parse::<ReasoningEffort>() {
|
||||
Ok(effort) => Some(effort),
|
||||
Err(error) => {
|
||||
tracing::warn!(token = %s, %error, "persisted reasoning_effort unknown; dropping");
|
||||
None
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
/// Write the canonical effort onto a serialized Responses request body
|
||||
/// (`body.reasoning.effort`). This is the ONLY place effort reaches the
|
||||
/// Responses wire: the typed `rs::ReasoningEffort` tops out at `xhigh`, so
|
||||
/// `max` must be written post-serialize. A `None` effort leaves the body
|
||||
/// untouched (the provider default applies).
|
||||
pub fn patch_reasoning_effort(body: &mut Value, effort: Option<ReasoningEffort>) {
|
||||
let Some(effort) = effort else { return };
|
||||
let Some(obj) = body.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
let reasoning = obj
|
||||
.entry("reasoning")
|
||||
.or_insert_with(|| Value::Object(serde_json::Map::new()));
|
||||
if let Some(reasoning) = reasoning.as_object_mut() {
|
||||
reasoning.insert(
|
||||
"effort".to_string(),
|
||||
Value::String(effort.as_str().to_string()),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Adapt a serialized Responses request body to the ChatGPT/Codex backend
|
||||
/// contract (`chatgpt.com/backend-api/codex/responses`) — ported from the
|
||||
/// same reference as the identity headers (official Codex CLI + Pi's
|
||||
/// `api/openai-codex-responses.ts`):
|
||||
///
|
||||
/// 1. The backend rejects `role: system` input items outright
|
||||
/// (400 `{"detail":"System messages are not allowed"}`); its system
|
||||
/// channel is the top-level `instructions` field. Hoist every system
|
||||
/// input message there (order preserved, blank-line joined, appended to
|
||||
/// any existing instructions) and remove them from `input`.
|
||||
/// 2. `store` is always `false` on this backend, so reasoning continuity
|
||||
/// is stateless: `include: ["reasoning.encrypted_content"]` is required
|
||||
/// for the response to carry replayable encrypted reasoning.
|
||||
///
|
||||
/// openai-codex-GATED at the call sites — API-key `openai` Responses
|
||||
/// bodies stay byte-identical.
|
||||
pub fn adapt_body_for_codex_backend(body: &mut Value) {
|
||||
// 1. Hoist system messages into `instructions`.
|
||||
let mut hoisted: Vec<String> = Vec::new();
|
||||
if let Some(input) = body.get_mut("input").and_then(|v| v.as_array_mut()) {
|
||||
input.retain(|item| {
|
||||
let is_system = item.get("role").and_then(|r| r.as_str()) == Some("system");
|
||||
if is_system {
|
||||
match item.get("content") {
|
||||
Some(Value::String(s)) => hoisted.push(s.clone()),
|
||||
Some(Value::Array(parts)) => {
|
||||
for p in parts {
|
||||
if let Some(t) = p.get("text").and_then(|t| t.as_str()) {
|
||||
hoisted.push(t.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
!is_system
|
||||
});
|
||||
}
|
||||
if !hoisted.is_empty() {
|
||||
let mut instructions = body
|
||||
.get("instructions")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(str::to_owned)
|
||||
.unwrap_or_default();
|
||||
for part in hoisted {
|
||||
if !instructions.is_empty() {
|
||||
instructions.push_str("\n\n");
|
||||
}
|
||||
instructions.push_str(&part);
|
||||
}
|
||||
body["instructions"] = Value::String(instructions);
|
||||
}
|
||||
|
||||
// 2. Request replayable encrypted reasoning.
|
||||
let include = body
|
||||
.as_object_mut()
|
||||
.map(|obj| obj.entry("include").or_insert_with(|| Value::Array(vec![])));
|
||||
if let Some(Value::Array(entries)) = include {
|
||||
let key = Value::String("reasoning.encrypted_content".to_string());
|
||||
if !entries.contains(&key) {
|
||||
entries.push(key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Neutralize a `reasoning.effort` echo the typed `rs` enum cannot parse
|
||||
/// (`max`): remove it so response deserialization succeeds. The turn's
|
||||
/// canonical effort lives in the session sampling config regardless; only
|
||||
/// the per-item echo metadata is dropped, recorded as not-echoed.
|
||||
/// (Ceiling: async-openai lacks a Max variant; delete this when it grows
|
||||
/// one.) Handles both bare response bodies (`/reasoning/effort`) and
|
||||
/// stream-event envelopes (`/response/reasoning/effort`).
|
||||
pub fn normalize_effort_echo(value: &mut Value) {
|
||||
for path in ["/reasoning", "/response/reasoning"] {
|
||||
if let Some(reasoning) = value.pointer_mut(path).and_then(|v| v.as_object_mut())
|
||||
&& let Some(effort) = reasoning.get("effort").and_then(|v| v.as_str())
|
||||
&& crate::rs::ReasoningEffort::deserialize(serde_json::Value::String(
|
||||
effort.to_string(),
|
||||
))
|
||||
.is_err()
|
||||
{
|
||||
tracing::debug!(
|
||||
effort,
|
||||
"reasoning.effort echo unrepresentable in the typed enum; dropping"
|
||||
);
|
||||
reasoning.remove("effort");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The `anthropic-version` header value kigi speaks on Anthropic-style
|
||||
/// wires (Messages inference and the /v1/models listing).
|
||||
pub const ANTHROPIC_VERSION: &str = "2023-06-01";
|
||||
|
||||
/// `anthropic-beta` value for the Claude-Code OAuth (Pro/Max subscription)
|
||||
/// path — required on BOTH the Messages inference request and the `/v1/models`
|
||||
/// listing when the bearer is an OAuth `sk-ant-oat…` token. API-key Anthropic
|
||||
/// and MiniMax NEVER send this (their requests stay byte-identical).
|
||||
pub const ANTHROPIC_OAUTH_BETA: &str = "claude-code-20250219,oauth-2025-04-20";
|
||||
|
||||
/// User-Agent kigi presents on the Claude-Code OAuth path (mirrors the
|
||||
/// official Claude Code CLI). OAuth-gated: unrelated to the default kigi UA.
|
||||
pub const CLAUDE_CODE_USER_AGENT: &str = "claude-cli/2.1.75";
|
||||
|
||||
/// System-prompt prefix REQUIRED on the Claude-Code OAuth Messages path: the
|
||||
/// OAuth token is Claude-Code-scoped, so Anthropic rejects the request unless
|
||||
/// the system prompt's first block is exactly this line. OAuth-gated.
|
||||
pub const CLAUDE_CODE_SYSTEM_PREFIX: &str =
|
||||
"You are Claude Code, Anthropic's official CLI for Claude.";
|
||||
|
||||
// ── GitHub Copilot editor-identity headers ──────────────────────────────────
|
||||
// The VS Code Copilot Chat client identity. Copilot's proxy authorizes the
|
||||
// short-lived copilot token AND validates these editor headers, so they ride
|
||||
// the `copilot_internal/v2/token` exchange, the `/models` listing, and every
|
||||
// `/chat/completions` inference request. github-copilot-GATED: no other
|
||||
// platform sends them, so their requests stay byte-identical. Values are
|
||||
// non-secret wire constants (ported from Pi `api/github-copilot-headers.ts` +
|
||||
// `auth/oauth/github-copilot.ts`).
|
||||
|
||||
/// `User-Agent` for the Copilot path (overrides the default kigi UA, OAuth-gated).
|
||||
pub const COPILOT_USER_AGENT: &str = "GitHubCopilotChat/0.35.0";
|
||||
/// `Editor-Version` — the host editor Copilot believes it is talking to.
|
||||
pub const COPILOT_EDITOR_VERSION: &str = "vscode/1.107.0";
|
||||
/// `Editor-Plugin-Version` — the Copilot Chat plugin build.
|
||||
pub const COPILOT_EDITOR_PLUGIN_VERSION: &str = "copilot-chat/0.35.0";
|
||||
/// `Copilot-Integration-Id` — the integration the token is scoped to.
|
||||
pub const COPILOT_INTEGRATION_ID: &str = "vscode-chat";
|
||||
/// `X-GitHub-Api-Version` — sent ONLY on the `/models` listing.
|
||||
pub const COPILOT_API_VERSION: &str = "2026-06-01";
|
||||
/// `X-Initiator` value — sent ONLY on inference (`user`, per the spec).
|
||||
pub const COPILOT_INITIATOR: &str = "user";
|
||||
|
||||
// ── ChatGPT/Codex (openai-codex) OAuth-inference headers ─────────────────────
|
||||
// The ChatGPT Codex backend authorizes an OAuth bearer AND validates the Codex
|
||||
// client identity. These ride the `/codex/responses` inference request ONLY.
|
||||
// openai-codex-GATED: no other Responses provider (API-key `openai`) sends them,
|
||||
// so their requests stay byte-identical. Values are non-secret wire constants
|
||||
// (ported from the official Codex CLI + Pi `api/openai-codex-responses.ts`).
|
||||
|
||||
/// `originator` header identifying the Codex CLI client (matches the authorize
|
||||
/// `originator` param).
|
||||
pub const CODEX_ORIGINATOR: &str = "codex_cli_rs";
|
||||
/// `OpenAI-Beta` opt-in the Codex Responses endpoint requires.
|
||||
pub const CODEX_OPENAI_BETA: &str = "responses=experimental";
|
||||
/// `User-Agent` presented on the Codex path (overrides the default kigi UA,
|
||||
/// openai-codex-gated). The Codex backend does not strictly validate the UA
|
||||
/// string (Pi ships its own and it works), so this is a stable best-effort
|
||||
/// identity, not a pinned build.
|
||||
pub const CODEX_USER_AGENT: &str = "codex_cli_rs/0.104.0";
|
||||
/// JWT payload claim namespace carrying the ChatGPT account id.
|
||||
const CODEX_JWT_AUTH_CLAIM: &str = "https://api.openai.com/auth";
|
||||
|
||||
/// Extract the `chatgpt_account_id` from a Codex OAuth access token (a JWT):
|
||||
/// base64url-decode the payload segment and read
|
||||
/// `["https://api.openai.com/auth"]["chatgpt_account_id"]`. Returns `None` when
|
||||
/// the token is not a well-formed JWT or the claim is missing/empty.
|
||||
///
|
||||
/// Used BOTH at login (fail-fast: a token without the claim is useless) and at
|
||||
/// inference (the header is derived STATELESSLY from the current bearer, so a
|
||||
/// refreshed token — which still carries the claim — needs no persisted field).
|
||||
///
|
||||
/// SECURITY: the token, its payload, and the returned account id are NEVER
|
||||
/// logged by this function or its callers.
|
||||
pub fn chatgpt_account_id_from_jwt(token: &str) -> Option<String> {
|
||||
use base64::Engine;
|
||||
// A JWT is exactly three dot-separated segments; anything else is not a
|
||||
// token we can read (fail closed rather than decode a lookalike).
|
||||
let mut segments = token.split('.');
|
||||
let (_header, payload_b64, _signature) = (segments.next()?, segments.next()?, segments.next()?);
|
||||
if segments.next().is_some() {
|
||||
return None;
|
||||
}
|
||||
// JWT payloads are base64url without padding; be tolerant of either.
|
||||
let bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
|
||||
.decode(payload_b64)
|
||||
.or_else(|_| base64::engine::general_purpose::URL_SAFE.decode(payload_b64))
|
||||
.ok()?;
|
||||
let claims: serde_json::Value = serde_json::from_slice(&bytes).ok()?;
|
||||
let account_id = claims
|
||||
.get(CODEX_JWT_AUTH_CLAIM)?
|
||||
.get("chatgpt_account_id")?
|
||||
.as_str()?;
|
||||
(!account_id.is_empty()).then(|| account_id.to_string())
|
||||
}
|
||||
|
||||
/// ChatCompletions request-body adaptation dialect. Providers disagree on
|
||||
/// how thinking rides an OpenAI-compatible body: Kimi wants
|
||||
/// `thinking:{type,effort}`, DeepSeek wants
|
||||
/// `thinking:{type,reasoning_effort}`, most others take the OpenAI-style
|
||||
/// `reasoning_effort` scalar untouched.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ChatCompat {
|
||||
/// Kimi wire (`thinking:{type,effort}`, canonical xhigh spelled max).
|
||||
/// The default: BYOK/custom ChatCompletions endpoints keep the
|
||||
/// historical kigi behavior.
|
||||
#[default]
|
||||
Kimi,
|
||||
/// DeepSeek wire (`thinking:{type,reasoning_effort}`, official docs:
|
||||
/// low/medium map to high and xhigh to max server-side).
|
||||
DeepSeek,
|
||||
/// Leave the body as-is (OpenAI-style `reasoning_effort` passes through).
|
||||
Passthrough,
|
||||
/// Strict OpenAI-compatible validators (Mistral, Cerebras) reject any
|
||||
/// out-of-schema request field with a 4xx (`additionalProperties:false`).
|
||||
/// kigi injects `stream_options.include_usage` on every streaming
|
||||
/// request, which such validators reject, so it is stripped (streaming
|
||||
/// usage falls back to token estimation). `reasoning_effort` passes
|
||||
/// through; private message fields are stripped like Passthrough.
|
||||
/// (Serde alias `mistral` keeps sessions persisted before the rename.)
|
||||
#[serde(alias = "mistral")]
|
||||
StrictOpenAi,
|
||||
}
|
||||
|
||||
pub const REASONING_EFFORT_META_KEY: &str = "reasoningEffort";
|
||||
pub const SUPPORTS_REASONING_EFFORT_META_KEY: &str = "supportsReasoningEffort";
|
||||
|
||||
@@ -1054,6 +1443,11 @@ pub struct SamplingConfig {
|
||||
pub max_completion_tokens: Option<u32>,
|
||||
pub temperature: Option<f32>,
|
||||
pub top_p: Option<f32>,
|
||||
/// ChatCompletions body-adaptation dialect (per-platform; serde-default
|
||||
/// Kimi keeps pre-field sessions and BYOK endpoints on the historical
|
||||
/// behavior).
|
||||
#[serde(default)]
|
||||
pub chat_compat: ChatCompat,
|
||||
/// Which API backend to use for this model
|
||||
#[serde(default)]
|
||||
pub api_backend: ApiBackend,
|
||||
@@ -1081,6 +1475,12 @@ pub struct CreateResponseWrapper {
|
||||
/// The inner Responses API request.
|
||||
pub inner: crate::rs::CreateResponse,
|
||||
|
||||
/// Canonical reasoning effort for this request. The typed
|
||||
/// `inner.reasoning.effort` stays `None` (async-openai's enum cannot
|
||||
/// spell `max`); [`patch_reasoning_effort`] writes this value onto the
|
||||
/// serialized body just before send.
|
||||
pub reasoning_effort: Option<ReasoningEffort>,
|
||||
|
||||
/// Custom header: conversation ID for tracking.
|
||||
pub x_kigi_conv_id: Option<String>,
|
||||
|
||||
@@ -1107,6 +1507,7 @@ impl CreateResponseWrapper {
|
||||
pub fn new(inner: crate::rs::CreateResponse) -> Self {
|
||||
Self {
|
||||
inner,
|
||||
reasoning_effort: None,
|
||||
x_kigi_conv_id: None,
|
||||
x_kigi_req_id: None,
|
||||
x_kigi_session_id: None,
|
||||
@@ -1215,6 +1616,173 @@ mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
/// The Codex backend rejects `role: system` input outright
|
||||
/// (400 `{"detail":"System messages are not allowed"}`) — its system
|
||||
/// channel is the top-level `instructions` field, and stateless
|
||||
/// (`store: false`) reasoning replay needs
|
||||
/// `include: ["reasoning.encrypted_content"]`. The adapter must hoist
|
||||
/// every system item (string or parts content, order preserved),
|
||||
/// append to existing instructions, and leave the rest of the input
|
||||
/// untouched.
|
||||
#[test]
|
||||
fn codex_adapter_hoists_system_messages_and_requests_encrypted_reasoning() {
|
||||
let mut body = json!({
|
||||
"model": "gpt-5.2-codex",
|
||||
"instructions": "base",
|
||||
"input": [
|
||||
{"type": "message", "role": "system", "content": "sys head"},
|
||||
{"type": "message", "role": "user", "content": "hello"},
|
||||
{"type": "message", "role": "system", "content": [
|
||||
{"type": "input_text", "text": "memory reminder"}
|
||||
]},
|
||||
{"type": "message", "role": "assistant", "content": "hi"}
|
||||
]
|
||||
});
|
||||
adapt_body_for_codex_backend(&mut body);
|
||||
|
||||
let input = body["input"].as_array().unwrap();
|
||||
assert_eq!(input.len(), 2, "system items removed from input: {body:#}");
|
||||
assert!(
|
||||
input.iter().all(|i| i["role"] != "system"),
|
||||
"no system role may remain: {body:#}"
|
||||
);
|
||||
assert_eq!(
|
||||
body["instructions"], "base\n\nsys head\n\nmemory reminder",
|
||||
"system content hoisted into instructions, order preserved"
|
||||
);
|
||||
assert_eq!(
|
||||
body["include"],
|
||||
json!(["reasoning.encrypted_content"]),
|
||||
"stateless reasoning replay requires the include"
|
||||
);
|
||||
|
||||
// Idempotent: a second pass changes nothing.
|
||||
let before = body.clone();
|
||||
adapt_body_for_codex_backend(&mut body);
|
||||
assert_eq!(body, before);
|
||||
}
|
||||
|
||||
/// No system items and no prior instructions: input untouched, no
|
||||
/// empty-string instructions invented, include still requested.
|
||||
#[test]
|
||||
fn codex_adapter_without_system_messages_only_adds_include() {
|
||||
let mut body = json!({
|
||||
"model": "gpt-5.2-codex",
|
||||
"input": [{"type": "message", "role": "user", "content": "q"}]
|
||||
});
|
||||
adapt_body_for_codex_backend(&mut body);
|
||||
assert!(body.get("instructions").is_none(), "{body:#}");
|
||||
assert_eq!(body["input"].as_array().unwrap().len(), 1);
|
||||
assert_eq!(body["include"], json!(["reasoning.encrypted_content"]));
|
||||
}
|
||||
|
||||
/// String content (the only shape non-Mistral providers send) stays the
|
||||
/// answer verbatim with no thinking — byte-identical to the pre-change
|
||||
/// deserialization.
|
||||
/// The StrictOpenAi dialect kept the serde alias `mistral`, so sessions
|
||||
/// persisted before the rename still deserialize.
|
||||
#[test]
|
||||
fn chat_compat_mistral_alias_deserializes_to_strict_openai() {
|
||||
let v: ChatCompat = serde_json::from_str("\"mistral\"").unwrap();
|
||||
assert_eq!(v, ChatCompat::StrictOpenAi);
|
||||
// New value round-trips as strict_open_ai.
|
||||
let v: ChatCompat = serde_json::from_str("\"strict_open_ai\"").unwrap();
|
||||
assert_eq!(v, ChatCompat::StrictOpenAi);
|
||||
assert_eq!(
|
||||
serde_json::to_string(&ChatCompat::StrictOpenAi).unwrap(),
|
||||
"\"strict_open_ai\""
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chunk_delta_string_content_unchanged() {
|
||||
let delta: ChatChunkDelta =
|
||||
serde_json::from_value(json!({ "content": "hello world" })).unwrap();
|
||||
assert_eq!(delta.content.as_deref(), Some("hello world"));
|
||||
assert_eq!(delta.reasoning_content, None);
|
||||
// A DeepSeek-style wire reasoning_content still rides its own field.
|
||||
let delta: ChatChunkDelta =
|
||||
serde_json::from_value(json!({ "content": "hi", "reasoning_content": "because" }))
|
||||
.unwrap();
|
||||
assert_eq!(delta.content.as_deref(), Some("hi"));
|
||||
assert_eq!(delta.reasoning_content.as_deref(), Some("because"));
|
||||
// Null / absent content → None.
|
||||
let delta: ChatChunkDelta = serde_json::from_value(json!({})).unwrap();
|
||||
assert_eq!(delta.content, None);
|
||||
}
|
||||
|
||||
/// Mistral reasoning: array content routes `text` chunks to the answer
|
||||
/// and the nested `text` of `thinking` chunks to reasoning_content
|
||||
/// (verified shapes from the mistralai/client-python SDK).
|
||||
#[test]
|
||||
fn chunk_delta_array_content_splits_thinking_and_answer() {
|
||||
// Thinking phase: array with only a thinking chunk (nested list).
|
||||
let delta: ChatChunkDelta = serde_json::from_value(json!({
|
||||
"content": [
|
||||
{ "type": "thinking", "thinking": [{ "type": "text", "text": "step 1" }],
|
||||
"signature": null, "closed": false }
|
||||
]
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(delta.content, None, "thinking phase has no visible answer");
|
||||
assert_eq!(delta.reasoning_content.as_deref(), Some("step 1"));
|
||||
|
||||
// Answer phase: plain string delta.
|
||||
let delta: ChatChunkDelta =
|
||||
serde_json::from_value(json!({ "content": "the answer" })).unwrap();
|
||||
assert_eq!(delta.content.as_deref(), Some("the answer"));
|
||||
assert_eq!(delta.reasoning_content, None);
|
||||
|
||||
// Transition delta: one array carrying both a closing thinking chunk
|
||||
// and the first text chunk → both channels populated.
|
||||
let delta: ChatChunkDelta = serde_json::from_value(json!({
|
||||
"content": [
|
||||
{ "type": "thinking", "thinking": [{ "type": "text", "text": "done" }] },
|
||||
{ "type": "text", "text": "Answer: 22" }
|
||||
]
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(delta.content.as_deref(), Some("Answer: 22"));
|
||||
assert_eq!(delta.reasoning_content.as_deref(), Some("done"));
|
||||
}
|
||||
|
||||
/// The chunk union is OPEN: unknown chunk types and reference chunks are
|
||||
/// ignored, never fatal (SDK's UnknownContentChunk fallback).
|
||||
#[test]
|
||||
fn content_array_tolerates_unknown_chunk_types() {
|
||||
let delta: ChatChunkDelta = serde_json::from_value(json!({
|
||||
"content": [
|
||||
{ "type": "reference", "reference_ids": [1, 2] },
|
||||
{ "type": "some_future_chunk", "payload": { "x": 1 } },
|
||||
{ "type": "text", "text": "visible" }
|
||||
]
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(delta.content.as_deref(), Some("visible"));
|
||||
assert_eq!(delta.reasoning_content, None);
|
||||
}
|
||||
|
||||
/// Non-streaming ChatResponseMessage gets the same array handling.
|
||||
#[test]
|
||||
fn response_message_array_content_splits() {
|
||||
let msg: ChatResponseMessage = serde_json::from_value(json!({
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{ "type": "thinking", "thinking": [{ "type": "text", "text": "reasoning" }] },
|
||||
{ "type": "text", "text": "final" }
|
||||
]
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(msg.content.as_deref(), Some("final"));
|
||||
assert_eq!(msg.reasoning_content.as_deref(), Some("reasoning"));
|
||||
// String content still works.
|
||||
let msg: ChatResponseMessage = serde_json::from_value(json!({
|
||||
"role": "assistant", "content": "plain"
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(msg.content.as_deref(), Some("plain"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reasoning_effort_serde_lowercase_round_trip() {
|
||||
for v in [
|
||||
@@ -1224,6 +1792,7 @@ mod tests {
|
||||
ReasoningEffort::Medium,
|
||||
ReasoningEffort::High,
|
||||
ReasoningEffort::Xhigh,
|
||||
ReasoningEffort::Max,
|
||||
] {
|
||||
let json = serde_json::to_string(&v).unwrap();
|
||||
assert_eq!(json, format!("\"{}\"", v.as_str()), "serialize {v:?}");
|
||||
@@ -1231,31 +1800,112 @@ mod tests {
|
||||
assert_eq!(back, v, "round-trip {v:?}");
|
||||
}
|
||||
assert!(serde_json::from_str::<ReasoningEffort>("\"BOGUS\"").is_err());
|
||||
assert!(serde_json::from_str::<ReasoningEffort>("\"max\"").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reasoning_effort_from_str_accepts_max_as_xhigh() {
|
||||
fn reasoning_effort_from_str_max_and_xhigh_are_distinct() {
|
||||
assert_eq!(
|
||||
"max".parse::<ReasoningEffort>().unwrap(),
|
||||
ReasoningEffort::Xhigh
|
||||
ReasoningEffort::Max
|
||||
);
|
||||
assert_eq!(
|
||||
"MAX".parse::<ReasoningEffort>().unwrap(),
|
||||
ReasoningEffort::Xhigh
|
||||
ReasoningEffort::Max
|
||||
);
|
||||
assert_eq!(
|
||||
"xhigh".parse::<ReasoningEffort>().unwrap(),
|
||||
ReasoningEffort::Xhigh
|
||||
);
|
||||
assert_eq!(ReasoningEffort::Xhigh.as_str(), "xhigh");
|
||||
// Messages API: distinct wire tokens per level (2026 capabilities).
|
||||
assert_eq!(ReasoningEffort::Xhigh.to_messages_api(), Some("xhigh"));
|
||||
assert_eq!(ReasoningEffort::Max.to_messages_api(), Some("max"));
|
||||
}
|
||||
|
||||
/// The codex-only `ultra` tier parses, serializes, and patches onto a
|
||||
/// Responses body as `reasoning.effort = "ultra"` (the crux of surfacing a
|
||||
/// codex model's full thinking menu). It is a DISTINCT level above `max`.
|
||||
#[test]
|
||||
fn reasoning_effort_ultra_is_a_distinct_codex_tier() {
|
||||
assert_eq!(
|
||||
"ultra".parse::<ReasoningEffort>().unwrap(),
|
||||
ReasoningEffort::Ultra
|
||||
);
|
||||
assert_eq!(
|
||||
"ULTRA".parse::<ReasoningEffort>().unwrap(),
|
||||
ReasoningEffort::Ultra
|
||||
);
|
||||
assert_ne!(ReasoningEffort::Ultra, ReasoningEffort::Max);
|
||||
assert_eq!(ReasoningEffort::Ultra.as_str(), "ultra");
|
||||
let json = serde_json::to_string(&ReasoningEffort::Ultra).unwrap();
|
||||
assert_eq!(json, "\"ultra\"");
|
||||
assert_eq!(
|
||||
serde_json::from_str::<ReasoningEffort>("\"ultra\"").unwrap(),
|
||||
ReasoningEffort::Ultra
|
||||
);
|
||||
let mut body = serde_json::json!({ "model": "gpt-5.6-sol" });
|
||||
patch_reasoning_effort(&mut body, Some(ReasoningEffort::Ultra));
|
||||
assert_eq!(body["reasoning"]["effort"], "ultra");
|
||||
}
|
||||
|
||||
/// The account id is decoded STATELESSLY from the bearer JWT payload's
|
||||
/// `["https://api.openai.com/auth"]["chatgpt_account_id"]` claim; a token
|
||||
/// without the claim (or not a JWT) yields `None` (login fails fast on it).
|
||||
#[test]
|
||||
fn chatgpt_account_id_extracted_from_jwt_claim() {
|
||||
use base64::Engine;
|
||||
let make_jwt = |payload: serde_json::Value| -> String {
|
||||
let header =
|
||||
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(br#"{"alg":"none"}"#);
|
||||
let body = base64::engine::general_purpose::URL_SAFE_NO_PAD
|
||||
.encode(serde_json::to_vec(&payload).unwrap());
|
||||
format!("{header}.{body}.sig")
|
||||
};
|
||||
let good = make_jwt(serde_json::json!({
|
||||
"https://api.openai.com/auth": { "chatgpt_account_id": "acct-abc-123" },
|
||||
"sub": "user-1"
|
||||
}));
|
||||
assert_eq!(
|
||||
chatgpt_account_id_from_jwt(&good).as_deref(),
|
||||
Some("acct-abc-123")
|
||||
);
|
||||
// Claim namespace present but no account id → None (fail-fast).
|
||||
let no_account = make_jwt(serde_json::json!({
|
||||
"https://api.openai.com/auth": { "user_id": "u" }
|
||||
}));
|
||||
assert_eq!(chatgpt_account_id_from_jwt(&no_account), None);
|
||||
// Empty account id → None.
|
||||
let empty = make_jwt(serde_json::json!({
|
||||
"https://api.openai.com/auth": { "chatgpt_account_id": "" }
|
||||
}));
|
||||
assert_eq!(chatgpt_account_id_from_jwt(&empty), None);
|
||||
// Not a JWT (no payload segment) → None.
|
||||
assert_eq!(chatgpt_account_id_from_jwt("not-a-jwt"), None);
|
||||
assert_eq!(chatgpt_account_id_from_jwt(""), None);
|
||||
// A JWT is EXACTLY three segments: a lookalike with too few or too many
|
||||
// is rejected outright rather than decoded (fail closed).
|
||||
let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(
|
||||
serde_json::to_vec(&serde_json::json!({
|
||||
"https://api.openai.com/auth": { "chatgpt_account_id": "acct-abc-123" }
|
||||
}))
|
||||
.unwrap(),
|
||||
);
|
||||
assert_eq!(
|
||||
chatgpt_account_id_from_jwt(&format!("hdr.{payload}")),
|
||||
None,
|
||||
"two segments is not a JWT"
|
||||
);
|
||||
assert_eq!(
|
||||
chatgpt_account_id_from_jwt(&format!("hdr.{payload}.sig.extra")),
|
||||
None,
|
||||
"four segments is not a JWT"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_canonical_effort_token_helper() {
|
||||
assert_eq!(
|
||||
parse_canonical_effort_token("max"),
|
||||
Some(ReasoningEffort::Xhigh)
|
||||
Some(ReasoningEffort::Max)
|
||||
);
|
||||
assert_eq!(
|
||||
parse_canonical_effort_token("high"),
|
||||
@@ -1415,7 +2065,8 @@ mod tests {
|
||||
);
|
||||
let bad_type = as_map(serde_json::json!({"reasoningEffort": 3}));
|
||||
assert_eq!(parse_reasoning_effort_meta(Some(&bad_type)), None);
|
||||
let unknown = as_map(serde_json::json!({"reasoningEffort": "ULTRA"}));
|
||||
// `ultra` is now a real codex tier; a genuinely-unknown token still None.
|
||||
let unknown = as_map(serde_json::json!({"reasoningEffort": "MEGA"}));
|
||||
assert_eq!(parse_reasoning_effort_meta(Some(&unknown)), None);
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
//! Filesystem primitives shared across the shell.
|
||||
|
||||
use std::io;
|
||||
use std::path::Path;
|
||||
|
||||
/// Replace `dest` with `tmp` — the commit step of every tmp+rename atomic
|
||||
/// write in the product. This is the ONE place that knows how to make that
|
||||
/// commit stick on Windows; call sites must never inline a bare
|
||||
/// `fs::rename` replace again.
|
||||
///
|
||||
/// Unix `rename(2)` replaces atomically and needs no help. Windows
|
||||
/// `MoveFileExW(REPLACE_EXISTING)` fails with a sharing violation while
|
||||
/// ANOTHER process (antivirus scanner, search indexer, cloud sync) holds
|
||||
/// `dest` open — the classic "persists on macOS, silently doesn't on
|
||||
/// Windows" failure (a model switch that never sticks, a stale models
|
||||
/// cache). On Windows a failed rename therefore deletes the destination
|
||||
/// first (the pattern `auth/storage.rs` shipped first) and retries with two
|
||||
/// short back-offs for scanners that hold the file for a few milliseconds.
|
||||
///
|
||||
/// On final failure the tmp file is removed (no litter) and the error is
|
||||
/// returned — callers decide severity, but MUST at least log it (errors
|
||||
/// never pass silently).
|
||||
pub fn replace_file(tmp: &Path, dest: &Path) -> io::Result<()> {
|
||||
let result = replace_file_inner(tmp, dest);
|
||||
if result.is_err() {
|
||||
let _ = std::fs::remove_file(tmp);
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
#[cfg(not(windows))]
|
||||
fn replace_file_inner(tmp: &Path, dest: &Path) -> io::Result<()> {
|
||||
std::fs::rename(tmp, dest)
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
fn replace_file_inner(tmp: &Path, dest: &Path) -> io::Result<()> {
|
||||
let mut last = match std::fs::rename(tmp, dest) {
|
||||
Ok(()) => return Ok(()),
|
||||
Err(e) => e,
|
||||
};
|
||||
for backoff_ms in [0u64, 10, 50] {
|
||||
if backoff_ms > 0 {
|
||||
std::thread::sleep(std::time::Duration::from_millis(backoff_ms));
|
||||
}
|
||||
// Delete-first: marks an open-with-delete-sharing file for deletion
|
||||
// and clears the way for a plain rename; harmless when absent.
|
||||
let _ = std::fs::remove_file(dest);
|
||||
match std::fs::rename(tmp, dest) {
|
||||
Ok(()) => return Ok(()),
|
||||
Err(e) => last = e,
|
||||
}
|
||||
}
|
||||
Err(last)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// The common contract on every platform: replace over an existing
|
||||
/// destination, create a missing one, and error (cleaning the tmp)
|
||||
/// when the tmp itself is missing.
|
||||
#[test]
|
||||
fn replace_file_commits_and_cleans_up() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let dest = dir.path().join("target.json");
|
||||
let tmp = dir.path().join("target.json.tmp");
|
||||
|
||||
// Create-missing.
|
||||
std::fs::write(&tmp, b"v1").unwrap();
|
||||
replace_file(&tmp, &dest).expect("create");
|
||||
assert_eq!(std::fs::read(&dest).unwrap(), b"v1");
|
||||
assert!(!tmp.exists(), "tmp must be consumed");
|
||||
|
||||
// Replace-existing.
|
||||
std::fs::write(&tmp, b"v2").unwrap();
|
||||
replace_file(&tmp, &dest).expect("replace");
|
||||
assert_eq!(std::fs::read(&dest).unwrap(), b"v2");
|
||||
assert!(!tmp.exists());
|
||||
|
||||
// Missing tmp → error, dest untouched.
|
||||
let err = replace_file(&tmp, &dest).expect_err("missing tmp must fail");
|
||||
assert_eq!(err.kind(), io::ErrorKind::NotFound);
|
||||
assert_eq!(std::fs::read(&dest).unwrap(), b"v2");
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
pub mod event_id;
|
||||
pub mod fs;
|
||||
pub mod kigi_home;
|
||||
pub mod secure_file;
|
||||
pub mod tips;
|
||||
@@ -34,7 +35,17 @@ pub fn random_f64() -> f64 {
|
||||
pub fn probabilistic_sample(rate: f64) -> bool {
|
||||
random_f64() < rate
|
||||
}
|
||||
fn matches_trusted_base_url(candidate: &str, trusted_base: &str) -> bool {
|
||||
/// True when `candidate` is `trusted_base` or a path below it, comparing
|
||||
/// scheme, host and effective port exactly (so suffix attacks such as
|
||||
/// `api.kimi.com.evil.example` never match).
|
||||
///
|
||||
/// Public because the credential chokepoint
|
||||
/// ([`kigi_shell::auth::credential_authority`](../../../kigi_shell/auth/credential_authority/index.html))
|
||||
/// must compare a request URL against the SESSION's effective endpoints
|
||||
/// (`[endpoints] coding_api_base_url` from config.toml, `models_base_url`, a
|
||||
/// platform's own registry host) — none of which the env-var-only predicates
|
||||
/// below can see.
|
||||
pub fn matches_trusted_base_url(candidate: &str, trusted_base: &str) -> bool {
|
||||
let Ok(candidate) = reqwest::Url::parse(candidate) else {
|
||||
return false;
|
||||
};
|
||||
|
||||
@@ -109,21 +109,43 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
t = asset_triple
|
||||
);
|
||||
|
||||
// Transient CDN hiccups (502/503/timeouts) are retried with backoff: a
|
||||
// single flaky response must not kill a ~50-minute release build (it
|
||||
// did — the v0.1.5 tag build failed on one 502). Genuine failures
|
||||
// (404, offline) still error out with the offline-build hint.
|
||||
let bytes: Vec<u8> = {
|
||||
let resp = reqwest::blocking::get(&url).map_err(|e| {
|
||||
format!(
|
||||
"Failed to download ripgrep: {}\nSet KIGI_SHELL_BUNDLE_RG_PATH to a local rg for offline builds.",
|
||||
e
|
||||
)
|
||||
})?;
|
||||
if !resp.status().is_success() {
|
||||
return Err(format!(
|
||||
"HTTP {} downloading ripgrep. Set KIGI_SHELL_BUNDLE_RG_PATH for offline builds.",
|
||||
resp.status()
|
||||
)
|
||||
.into());
|
||||
let mut last_err = String::new();
|
||||
let mut bytes = None;
|
||||
for (attempt, backoff_secs) in [0u64, 2, 8].into_iter().enumerate() {
|
||||
if backoff_secs > 0 {
|
||||
std::thread::sleep(std::time::Duration::from_secs(backoff_secs));
|
||||
}
|
||||
resp.bytes()?.to_vec()
|
||||
match reqwest::blocking::get(&url) {
|
||||
Ok(resp) if resp.status().is_success() => match resp.bytes() {
|
||||
Ok(b) => {
|
||||
bytes = Some(b.to_vec());
|
||||
break;
|
||||
}
|
||||
Err(e) => last_err = format!("reading ripgrep body: {e}"),
|
||||
},
|
||||
Ok(resp) => {
|
||||
let status = resp.status();
|
||||
last_err = format!("HTTP {status} downloading ripgrep");
|
||||
// Only server-side/transient statuses are worth retrying.
|
||||
if !(status.is_server_error() || status.as_u16() == 429) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(e) => last_err = format!("Failed to download ripgrep: {e}"),
|
||||
}
|
||||
println!(
|
||||
"cargo:warning=ripgrep download attempt {} failed: {last_err}",
|
||||
attempt + 1
|
||||
);
|
||||
}
|
||||
bytes.ok_or_else(|| {
|
||||
format!("{last_err}. Set KIGI_SHELL_BUNDLE_RG_PATH for offline builds.")
|
||||
})?
|
||||
};
|
||||
|
||||
let gz = flate2::read::GzDecoder::new(&bytes[..]);
|
||||
|
||||
@@ -170,9 +170,7 @@ fn write_data_file_atomic(
|
||||
let json = serde_json::to_string_pretty(sessions)
|
||||
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
|
||||
fs::write(tmp_path, json.as_bytes())?;
|
||||
fs::rename(tmp_path, data_path).inspect_err(|_| {
|
||||
let _ = fs::remove_file(tmp_path);
|
||||
})
|
||||
crate::util::fs::replace_file(tmp_path, data_path)
|
||||
}
|
||||
|
||||
fn is_pid_alive(pid: u32) -> bool {
|
||||
|
||||
@@ -178,7 +178,13 @@ async fn prefetch_models(agent_config: &AgentConfig) -> Option<IndexMap<String,
|
||||
|
||||
if auth.is_some() || endpoints.has_custom_endpoint() || platform_keys.any() {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
prefetch_models_blocking(&endpoints, auth.as_ref(), fetch_auth, &platform_keys)
|
||||
prefetch_models_blocking(
|
||||
&endpoints,
|
||||
auth.as_ref(),
|
||||
&Default::default(),
|
||||
fetch_auth,
|
||||
&platform_keys,
|
||||
)
|
||||
})
|
||||
.await
|
||||
.ok()
|
||||
@@ -621,6 +627,7 @@ pub async fn run_leader(
|
||||
crate::agent::models::prefetch_models_blocking(
|
||||
&endpoints_for_prefetch,
|
||||
auth_for_prefetch.as_ref(),
|
||||
&Default::default(),
|
||||
fetch_auth_for_prefetch,
|
||||
&platform_keys_for_prefetch,
|
||||
)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -311,7 +311,7 @@ mod tests {
|
||||
w.kind == ModelOverrideWarningKind::DuplicateAlias
|
||||
&& w.field.as_deref() == Some("send_compactions_remaining")
|
||||
}));
|
||||
let resolved = crate::agent::config::resolve_model_list(&cfg, None);
|
||||
let resolved = crate::agent::config::resolve_model_list(&cfg, None, &Default::default());
|
||||
assert!(resolved.contains_key("kigi-4.5"));
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,432 @@
|
||||
//! models.dev enrichment catalog loading (agent-side IO).
|
||||
//!
|
||||
//! Providers whose `/models` wire serves no context/thinking metadata get it
|
||||
//! from models.dev (see `kigi_models::enrichment`). This module owns the IO:
|
||||
//! a 24h-TTL disk cache under `~/.kigi`, a runtime refresh of
|
||||
//! `https://models.dev/api.json` (filtered to registry providers before
|
||||
//! caching), and the bundled-snapshot fallback. NO NETWORK unless some
|
||||
//! enabled platform actually needs enrichment (`wire_serves_metadata` false)
|
||||
//! — with only Kimi/Moonshot configured this module never leaves disk.
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use kigi_models::enrichment::{EnrichmentCatalog, bundled_enrichment, parse_api_json};
|
||||
|
||||
/// Override the refresh URL (e2e mock), or disable refresh entirely with
|
||||
/// `0`/`off` (bundled snapshot + existing cache only).
|
||||
pub(crate) const MODELS_DEV_URL_ENV: &str = "KIGI_MODELS_DEV_URL";
|
||||
const DEFAULT_MODELS_DEV_URL: &str = "https://models.dev/api.json";
|
||||
const CACHE_FILE: &str = "models_dev_cache.json";
|
||||
const CACHE_TTL: Duration = Duration::from_secs(24 * 60 * 60);
|
||||
|
||||
#[derive(serde::Serialize, serde::Deserialize)]
|
||||
struct DiskCache {
|
||||
/// Unix seconds of the successful fetch.
|
||||
fetched_at: u64,
|
||||
/// The kigi version that wrote the cache — a different binary (upgrade
|
||||
/// OR downgrade) refetches rather than trusting old filtering rules.
|
||||
#[serde(default)]
|
||||
kigi_version: String,
|
||||
/// The keep-set the catalog was filtered to. A registry change (new
|
||||
/// provider row, models_dev_id rename) invalidates the cache instead of
|
||||
/// silently serving a catalog missing the new provider for up to 24h.
|
||||
#[serde(default)]
|
||||
keep_set: Vec<String>,
|
||||
/// Already filtered + transformed catalog.
|
||||
catalog: EnrichmentCatalog,
|
||||
}
|
||||
|
||||
fn current_keep_set() -> Vec<String> {
|
||||
registry_models_dev_ids()
|
||||
.into_iter()
|
||||
.map(str::to_owned)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Whether any of `platforms` needs enrichment at all.
|
||||
pub(crate) fn any_platform_needs_enrichment(platforms: &[kigi_models::PlatformId]) -> bool {
|
||||
platforms.iter().any(|p| !p.wire_serves_metadata())
|
||||
}
|
||||
|
||||
/// The registry's models.dev provider ids (the refresh filter).
|
||||
fn registry_models_dev_ids() -> BTreeSet<&'static str> {
|
||||
kigi_models::PlatformId::ALL
|
||||
.into_iter()
|
||||
.filter_map(|p| p.models_dev_id())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Cache dir override (tests re-home the cache away from the real
|
||||
/// `~/.kigi` — same pattern as `KIGI_MODELS_CACHE_DIR`).
|
||||
pub(crate) const MODELS_DEV_CACHE_DIR_ENV: &str = "KIGI_MODELS_DEV_CACHE_DIR";
|
||||
|
||||
fn cache_path() -> std::path::PathBuf {
|
||||
match std::env::var(MODELS_DEV_CACHE_DIR_ENV) {
|
||||
Ok(dir) if !dir.trim().is_empty() => std::path::PathBuf::from(dir).join(CACHE_FILE),
|
||||
_ => crate::util::kigi_home::kigi_home().join(CACHE_FILE),
|
||||
}
|
||||
}
|
||||
|
||||
fn refresh_url() -> Option<String> {
|
||||
match std::env::var(MODELS_DEV_URL_ENV) {
|
||||
Ok(v)
|
||||
if matches!(
|
||||
v.trim().to_ascii_lowercase().as_str(),
|
||||
"0" | "off" | "false"
|
||||
) =>
|
||||
{
|
||||
None
|
||||
}
|
||||
Ok(v) if !v.trim().is_empty() => Some(v.trim().to_string()),
|
||||
_ => Some(DEFAULT_MODELS_DEV_URL.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn now_unix() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Load the enrichment catalog for a fetch pass over `enabled` platforms.
|
||||
///
|
||||
/// Fast path: nothing needs enrichment → empty catalog, zero IO (the merge
|
||||
/// branch is never taken for wire-served platforms, and NOT forcing the
|
||||
/// bundled parse keeps its cost/panic surface off the kimi/moonshot path).
|
||||
/// Otherwise: fresh valid disk cache → use it; else refresh over HTTP
|
||||
/// (filter + transform + best-effort cache write); on refresh failure fall
|
||||
/// back to a STALE cache, then the bundled snapshot — each step logged,
|
||||
/// never silent.
|
||||
pub(crate) fn load_enrichment_catalog(
|
||||
enabled: &[kigi_models::PlatformId],
|
||||
) -> std::borrow::Cow<'static, EnrichmentCatalog> {
|
||||
if !any_platform_needs_enrichment(enabled) {
|
||||
return std::borrow::Cow::Owned(EnrichmentCatalog::new());
|
||||
}
|
||||
load_enrichment_catalog_at(&cache_path())
|
||||
}
|
||||
|
||||
/// Path-injectable core (tests use a tempdir path directly — no env, no
|
||||
/// `kigi_home()` OnceLock interaction).
|
||||
fn load_enrichment_catalog_at(
|
||||
path: &std::path::Path,
|
||||
) -> std::borrow::Cow<'static, EnrichmentCatalog> {
|
||||
use std::borrow::Cow;
|
||||
|
||||
let cached: Option<DiskCache> =
|
||||
std::fs::read_to_string(path)
|
||||
.ok()
|
||||
.and_then(|s| match serde_json::from_str(&s) {
|
||||
Ok(c) => Some(c),
|
||||
Err(e) => {
|
||||
tracing::warn!(path = %path.display(), error = %e,
|
||||
"models.dev cache unreadable; refetching");
|
||||
None
|
||||
}
|
||||
});
|
||||
let now = now_unix();
|
||||
let cache_is_fresh = cached.as_ref().is_some_and(|c| {
|
||||
// A future fetched_at (clock jump backwards, corrupt stamp) is
|
||||
// stale, not fresh-forever; a different binary or keep-set means the
|
||||
// cache was filtered under other rules — refetch instead of serving
|
||||
// a catalog that may miss newly-registered providers for 24h.
|
||||
c.fetched_at <= now
|
||||
&& now - c.fetched_at < CACHE_TTL.as_secs()
|
||||
&& c.kigi_version == kigi_version::VERSION
|
||||
&& c.keep_set == current_keep_set()
|
||||
});
|
||||
if cache_is_fresh {
|
||||
tracing::debug!("models.dev enrichment: fresh disk cache");
|
||||
return Cow::Owned(cached.expect("cache_is_fresh implies Some").catalog);
|
||||
}
|
||||
|
||||
match refresh_url() {
|
||||
Some(url) => match fetch_and_filter(&url) {
|
||||
Ok(catalog) => {
|
||||
let cache = DiskCache {
|
||||
fetched_at: now_unix(),
|
||||
kigi_version: kigi_version::VERSION.to_string(),
|
||||
keep_set: current_keep_set(),
|
||||
catalog,
|
||||
};
|
||||
// Best-effort write: a read-only home must not fail the fetch.
|
||||
match serde_json::to_string(&cache) {
|
||||
Ok(body) => {
|
||||
if let Err(e) = crate::util::config::atomic_write_string(path, &body) {
|
||||
tracing::warn!(path = %path.display(), error = %e,
|
||||
"models.dev cache write failed; continuing in-memory");
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "models.dev cache serialize failed")
|
||||
}
|
||||
}
|
||||
tracing::info!("models.dev enrichment refreshed");
|
||||
Cow::Owned(cache.catalog)
|
||||
}
|
||||
Err(e) => {
|
||||
if let Some(c) = cached {
|
||||
tracing::warn!(error = %e,
|
||||
"models.dev refresh failed; using STALE cache");
|
||||
Cow::Owned(c.catalog)
|
||||
} else {
|
||||
tracing::warn!(error = %e,
|
||||
"models.dev refresh failed; using bundled snapshot");
|
||||
Cow::Borrowed(bundled_enrichment())
|
||||
}
|
||||
}
|
||||
},
|
||||
None => {
|
||||
tracing::info!("models.dev refresh disabled; using cache/bundled");
|
||||
match cached {
|
||||
Some(c) => Cow::Owned(c.catalog),
|
||||
None => Cow::Borrowed(bundled_enrichment()),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn fetch_and_filter(url: &str) -> anyhow::Result<EnrichmentCatalog> {
|
||||
let response = crate::http::shared_blocking_client().get(url).send()?;
|
||||
let status = response.status();
|
||||
anyhow::ensure!(status.is_success(), "GET {url}: HTTP {}", status.as_u16());
|
||||
let body = response.text()?;
|
||||
let keep = registry_models_dev_ids();
|
||||
Ok(parse_api_json(&body, Some(&keep))?)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use kigi_test_support::EnvGuard;
|
||||
use serial_test::serial;
|
||||
|
||||
/// Wire-served-only platform sets never trigger IO — and never force the
|
||||
/// bundled parse (empty owned catalog; the merge branch is gated off).
|
||||
/// Kimi/Moonshot users therefore keep a zero-egress, zero-cache fetch
|
||||
/// path even now that enrichment-needing platforms (OpenAI) exist.
|
||||
#[test]
|
||||
fn wire_served_platforms_get_empty_catalog_without_io() {
|
||||
let wire_served = [
|
||||
kigi_models::PlatformId::KimiCode,
|
||||
kigi_models::PlatformId::MoonshotCn,
|
||||
kigi_models::PlatformId::MoonshotAi,
|
||||
];
|
||||
assert!(!any_platform_needs_enrichment(&wire_served));
|
||||
let catalog = load_enrichment_catalog(&wire_served);
|
||||
assert!(catalog.is_empty());
|
||||
// The full registry now DOES need enrichment (OpenAI is
|
||||
// wire_serves_metadata=false) — the fast path must not hide that.
|
||||
assert!(any_platform_needs_enrichment(&kigi_models::PlatformId::ALL));
|
||||
}
|
||||
|
||||
fn cache_file_in(dir: &tempfile::TempDir) -> std::path::PathBuf {
|
||||
dir.path().join(CACHE_FILE)
|
||||
}
|
||||
|
||||
fn write_cache(path: &std::path::Path, fetched_at: u64, versioned: bool) {
|
||||
let cache = DiskCache {
|
||||
fetched_at,
|
||||
kigi_version: if versioned {
|
||||
kigi_version::VERSION.to_string()
|
||||
} else {
|
||||
"0.0.0-other".to_string()
|
||||
},
|
||||
keep_set: current_keep_set(),
|
||||
catalog: EnrichmentCatalog::from([(
|
||||
"moonshotai".to_string(),
|
||||
std::collections::BTreeMap::from([(
|
||||
"from-cache".to_string(),
|
||||
kigi_models::enrichment::EnrichmentModel {
|
||||
context: 111,
|
||||
..Default::default()
|
||||
},
|
||||
)]),
|
||||
)]),
|
||||
};
|
||||
std::fs::write(path, serde_json::to_string(&cache).unwrap()).unwrap();
|
||||
}
|
||||
|
||||
/// Fresh valid cache short-circuits — no HTTP (no mock server mounted:
|
||||
/// a fetch attempt would fail and fall to bundled, failing the assert).
|
||||
#[test]
|
||||
#[serial]
|
||||
fn fresh_valid_cache_is_served_without_refresh() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = cache_file_in(&dir);
|
||||
write_cache(&path, now_unix(), true);
|
||||
let _url = EnvGuard::set(MODELS_DEV_URL_ENV, "http://127.0.0.1:1/api.json");
|
||||
let catalog = load_enrichment_catalog_at(&path);
|
||||
assert!(
|
||||
kigi_models::enrichment::lookup(&catalog, "moonshotai", "from-cache").is_some(),
|
||||
"fresh cache must be served"
|
||||
);
|
||||
}
|
||||
|
||||
/// Version/keep-set/future-stamp guards: each invalidates a fresh-aged
|
||||
/// cache. Refresh is disabled here, so invalidation falls through to the
|
||||
/// STALE cache (resource degradation, not data loss) — proving both the
|
||||
/// guard firing and the fallback order.
|
||||
#[test]
|
||||
#[serial]
|
||||
fn cache_guards_invalidate_and_fall_back_to_stale() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = cache_file_in(&dir);
|
||||
let _url = EnvGuard::set(MODELS_DEV_URL_ENV, "0");
|
||||
// Wrong binary version → not fresh → (refresh disabled) → stale used.
|
||||
write_cache(&path, now_unix(), false);
|
||||
let catalog = load_enrichment_catalog_at(&path);
|
||||
assert!(
|
||||
kigi_models::enrichment::lookup(&catalog, "moonshotai", "from-cache").is_some(),
|
||||
"stale-fallback must still serve the cached data"
|
||||
);
|
||||
// Future fetched_at → same path (guard fired: debug-log absence is
|
||||
// not observable here; the behavioral pin is refresh-disabled + the
|
||||
// wiremock test below proving a fired guard refetches).
|
||||
write_cache(&path, now_unix() + 10_000, true);
|
||||
let catalog = load_enrichment_catalog_at(&path);
|
||||
assert!(kigi_models::enrichment::lookup(&catalog, "moonshotai", "from-cache").is_some());
|
||||
}
|
||||
|
||||
/// An invalidated cache (wrong version) REFETCHES when refresh is
|
||||
/// enabled: wiremock expect(1) proves the HTTP call happened; the new
|
||||
/// cache file carries the current version + keep-set and the fetched
|
||||
/// content replaces the stale entry.
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
#[serial]
|
||||
async fn invalidated_cache_refetches_and_rewrites() {
|
||||
let _ = tracing_subscriber::fmt()
|
||||
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
|
||||
.with_test_writer()
|
||||
.try_init();
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("GET"))
|
||||
.and(wiremock::matchers::path("/api.json"))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"moonshotai": { "models": { "from-refresh": {
|
||||
"limit": {"context": 222}
|
||||
}}}
|
||||
})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = cache_file_in(&dir);
|
||||
write_cache(&path, now_unix(), false);
|
||||
let _url = EnvGuard::set(MODELS_DEV_URL_ENV, format!("{}/api.json", server.uri()));
|
||||
let path2 = path.clone();
|
||||
let catalog =
|
||||
tokio::task::spawn_blocking(move || load_enrichment_catalog_at(&path2).into_owned())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
kigi_models::enrichment::lookup(&catalog, "moonshotai", "from-refresh").is_some(),
|
||||
"guard-invalidated cache must refetch"
|
||||
);
|
||||
let rewritten: DiskCache =
|
||||
serde_json::from_str(&std::fs::read_to_string(path).unwrap()).unwrap();
|
||||
assert_eq!(rewritten.kigi_version, kigi_version::VERSION);
|
||||
assert_eq!(rewritten.keep_set, current_keep_set());
|
||||
assert!(rewritten.catalog.contains_key("moonshotai"));
|
||||
}
|
||||
|
||||
/// Corrupted cache file → refetch (not a crash, not trust-garbage).
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
#[serial]
|
||||
async fn corrupted_cache_refetches() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("GET"))
|
||||
.and(wiremock::matchers::path("/api.json"))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200)
|
||||
.set_body_json(serde_json::json!({ "moonshotai": { "models": {} } })),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = cache_file_in(&dir);
|
||||
std::fs::write(&path, "not json {").unwrap();
|
||||
let _url = EnvGuard::set(MODELS_DEV_URL_ENV, format!("{}/api.json", server.uri()));
|
||||
let path2 = path.clone();
|
||||
let catalog =
|
||||
tokio::task::spawn_blocking(move || load_enrichment_catalog_at(&path2).into_owned())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(catalog.contains_key("moonshotai"));
|
||||
}
|
||||
|
||||
/// Refresh failure with NO cache → bundled snapshot fallback.
|
||||
#[test]
|
||||
#[serial]
|
||||
fn refresh_failure_without_cache_falls_back_to_bundled() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = cache_file_in(&dir);
|
||||
let _url = EnvGuard::set(MODELS_DEV_URL_ENV, "http://127.0.0.1:1/api.json");
|
||||
let catalog = load_enrichment_catalog_at(&path);
|
||||
assert!(
|
||||
kigi_models::enrichment::lookup(&catalog, "kimi-for-coding", "k3").is_some(),
|
||||
"bundled snapshot must back a total refresh failure"
|
||||
);
|
||||
assert!(!path.exists(), "failed refresh must not write a cache");
|
||||
}
|
||||
|
||||
/// Kill switch through the FULL load path (not just refresh_url): no
|
||||
/// cache + refresh disabled → bundled, no HTTP attempted (an attempt
|
||||
/// against the sentinel URL would be a hang/refusal, not bundled data).
|
||||
#[test]
|
||||
#[serial]
|
||||
fn kill_switch_full_path_serves_bundled() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = cache_file_in(&dir);
|
||||
for token in ["0", "off", "FALSE", " Off "] {
|
||||
let _url = EnvGuard::set(MODELS_DEV_URL_ENV, token);
|
||||
assert!(refresh_url().is_none(), "token {token:?} must disable");
|
||||
let catalog = load_enrichment_catalog_at(&path);
|
||||
assert!(kigi_models::enrichment::lookup(&catalog, "kimi-for-coding", "k3").is_some());
|
||||
}
|
||||
}
|
||||
|
||||
/// Refresh path: mock server → transform runs and the keep-set filters
|
||||
/// to registry providers (`moonshotai` is a real registry models_dev id;
|
||||
/// unknown providers are dropped).
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
#[serial]
|
||||
async fn refresh_transforms_and_filters_to_registry_ids() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("GET"))
|
||||
.and(wiremock::matchers::path("/api.json"))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"moonshotai": { "models": { "kimi-test": {
|
||||
"limit": {"context": 262144},
|
||||
"reasoning": true,
|
||||
"reasoning_options": [
|
||||
{"type": "effort", "values": ["low", "high"]}
|
||||
]
|
||||
}}},
|
||||
"not-in-registry": { "models": { "m": {} } }
|
||||
})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let url = format!("{}/api.json", server.uri());
|
||||
let catalog = tokio::task::spawn_blocking(move || fetch_and_filter(&url))
|
||||
.await
|
||||
.unwrap()
|
||||
.expect("fetch must succeed");
|
||||
assert!(
|
||||
!catalog.contains_key("not-in-registry"),
|
||||
"keep-set must drop providers outside the registry"
|
||||
);
|
||||
let m = kigi_models::enrichment::lookup(&catalog, "moonshotai", "kimi-test")
|
||||
.expect("registry provider survives the filter");
|
||||
assert_eq!(m.context, 262_144);
|
||||
assert_eq!(m.efforts, ["low", "high"]);
|
||||
}
|
||||
}
|
||||
@@ -110,6 +110,29 @@ pub(crate) async fn apply(
|
||||
.models_manager
|
||||
.model_supports_reasoning_effort(model_id.0.as_ref())
|
||||
{
|
||||
// Legacy migration: pre-split sessions persisted canonical
|
||||
// `xhigh` for models whose live menu now spells the top tier
|
||||
// `max` (K3). The wire is identical either way (kimi_compat
|
||||
// renames xhigh→max), but the menu has no xhigh-valued row, so
|
||||
// display/active-row would drift from the model vocabulary and
|
||||
// the stale token would be re-persisted forever. Migrate once.
|
||||
let eff = if eff == kigi_sampling_types::ReasoningEffort::Xhigh
|
||||
&& !agent
|
||||
.models_manager
|
||||
.model_offers_effort(model_id.0.as_ref(), eff)
|
||||
&& agent.models_manager.model_offers_effort(
|
||||
model_id.0.as_ref(),
|
||||
kigi_sampling_types::ReasoningEffort::Max,
|
||||
) {
|
||||
tracing::info!(
|
||||
session_id = % session_id.0,
|
||||
"set_session_model: migrating legacy xhigh override to max \
|
||||
(model menu offers max, not xhigh)"
|
||||
);
|
||||
kigi_sampling_types::ReasoningEffort::Max
|
||||
} else {
|
||||
eff
|
||||
};
|
||||
tracing::info!(
|
||||
session_id = % session_id.0, effort = % eff,
|
||||
"set_session_model: applying reasoning_effort override from meta"
|
||||
@@ -169,9 +192,18 @@ pub(crate) async fn apply(
|
||||
model.map(|e| &e.info),
|
||||
)
|
||||
};
|
||||
// H4: hand the session the catalog KEY the picker actually resolved. The
|
||||
// slug in `model_sampling.model` cannot distinguish `xai/grok-*` from
|
||||
// `xai-grok/grok-*` (duplicate ids across an API-key platform and its
|
||||
// subscription-OAuth twin are by design), and the process-global
|
||||
// `current_model_id()` below is not written at all in Leader mode.
|
||||
let catalog_key =
|
||||
crate::agent::models::resolve_catalog_key(&agent.models_manager.models(), &model_id)
|
||||
.map(|k| k.0.to_string());
|
||||
let (tx, rx) = oneshot::channel();
|
||||
let _ = handle.cmd_tx.send(SessionCommand::SetSessionModel {
|
||||
sampling_config: model_sampling,
|
||||
catalog_key,
|
||||
use_concise,
|
||||
apply_prompt_override,
|
||||
skip_prompt_rewrite: did_rebuild || model_unchanged,
|
||||
|
||||
@@ -4,6 +4,7 @@ pub mod auth_method;
|
||||
pub mod chat_modes;
|
||||
pub mod config;
|
||||
pub mod config_model_override_parse;
|
||||
pub(crate) mod enrichment_fetch;
|
||||
mod ext_parsers;
|
||||
pub(crate) mod feedback_client;
|
||||
pub mod folder_trust;
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -165,7 +165,7 @@ impl acp::Agent for MvpAgent {
|
||||
&crate::util::kigi_home::kigi_home(),
|
||||
)
|
||||
{
|
||||
unsafe { std::env::set_var("XAI_API_KEY", &api_key) };
|
||||
unsafe { std::env::set_var("KIGI_API_KEY", &api_key) };
|
||||
tracing::info!("auth: loaded API key from auth.json (xai::api_key scope)");
|
||||
kigi_log::unified_log::info(
|
||||
"auth: loaded API key from auth.json (xai::api_key scope)",
|
||||
@@ -221,7 +221,24 @@ impl acp::Agent for MvpAgent {
|
||||
has_cached_token,
|
||||
login_label: None,
|
||||
});
|
||||
let auth_methods = built.methods;
|
||||
let mut auth_methods = built.methods;
|
||||
// Connected badges for the client's login picker: probe stored
|
||||
// credentials once (auth.json scopes + resolved platform keys) and
|
||||
// stamp `_meta.connected` on every method that already has one.
|
||||
{
|
||||
let store = crate::auth::read_auth_json(
|
||||
&crate::util::kigi_home::kigi_home().join("auth.json"),
|
||||
)
|
||||
.unwrap_or_default();
|
||||
let keys = crate::agent::models::PlatformApiKeys::resolve_from_effective_config();
|
||||
let connected = auth_method::connected_method_ids(
|
||||
has_cached_token,
|
||||
has_external_api_key,
|
||||
|scope| store.contains_key(scope),
|
||||
|p| keys.key_for(p).is_some(),
|
||||
);
|
||||
auth_method::stamp_connected_meta(&mut auth_methods, &connected);
|
||||
}
|
||||
kigi_log::unified_log::info(
|
||||
"auth: initialize() built auth_methods for ACP response",
|
||||
None,
|
||||
@@ -330,10 +347,28 @@ impl acp::Agent for MvpAgent {
|
||||
);
|
||||
match arguments.method_id.0.as_ref() {
|
||||
auth_method::XAI_API_KEY_METHOD_ID => {
|
||||
// C1: the SECOND writer of the shared `sampling_config.api_key`.
|
||||
// The house `KIGI_API_KEY` is the user's own key for the
|
||||
// session's own endpoint, so the stamp requires the authority to
|
||||
// classify that endpoint `CredentialClass::Primary` — NOT merely
|
||||
// "takes some session credential", which is also true on a
|
||||
// subscription-OAuth platform's host, where this key has no
|
||||
// business (the shared config is the subagent baseline and the
|
||||
// unresolved-model fallback). The key is still persisted to
|
||||
// auth.json either way; only the stamp is guarded.
|
||||
let takes_house_key = self.shared_config_takes_house_key();
|
||||
let mut sampling_config = self.sampling_config.borrow_mut();
|
||||
if sampling_config.api_key.is_none() {
|
||||
if let Ok(api_key) = auth_method::read_xai_api_key_env() {
|
||||
if takes_house_key {
|
||||
sampling_config.api_key = Some(api_key.clone());
|
||||
} else {
|
||||
tracing::debug!(
|
||||
model = sampling_config.model.as_str(),
|
||||
"auth: house api key withheld from the shared sampling config \
|
||||
(its endpoint is not this session's own)"
|
||||
);
|
||||
}
|
||||
if let Err(e) = crate::auth::store_api_key(
|
||||
&crate::util::kigi_home::kigi_home(),
|
||||
&api_key,
|
||||
@@ -357,7 +392,7 @@ impl acp::Agent for MvpAgent {
|
||||
return Err(
|
||||
acp::Error::auth_required()
|
||||
.data(
|
||||
"Set XAI_API_KEY or add api_key/env_key to config.toml.",
|
||||
"Set KIGI_API_KEY or add api_key/env_key to config.toml.",
|
||||
),
|
||||
);
|
||||
}
|
||||
@@ -389,7 +424,7 @@ impl acp::Agent for MvpAgent {
|
||||
),
|
||||
),
|
||||
);
|
||||
let Some(auth) = self.auth_manager.current() else {
|
||||
let Some(_auth) = self.auth_manager.current() else {
|
||||
let message = if self.auth_manager.is_expired() {
|
||||
"Session expired, re-authentication required"
|
||||
} else {
|
||||
@@ -408,9 +443,15 @@ impl acp::Agent for MvpAgent {
|
||||
.await;
|
||||
};
|
||||
self.emit_settings_update_notification();
|
||||
{
|
||||
let mut sampling_config = self.sampling_config.borrow_mut();
|
||||
sampling_config.api_key = Some(auth.key);
|
||||
// H2/C1: route the stamp through the ONE guard, which asks the
|
||||
// authority which credential governs the shared config rather
|
||||
// than being handed this one. That config may already point at
|
||||
// a third-party model or at another provider's subscription
|
||||
// host (it is the subagent baseline and the unresolved-model
|
||||
// fallback), and `auth.key` authorizes only the session's own
|
||||
// coding endpoint. The manager already holds this token, so the
|
||||
// authority reads it back where it belongs.
|
||||
if self.stamp_session_credential(true) {
|
||||
tracing::debug!(
|
||||
"auth: cached_token handler set api_key (SessionToken)"
|
||||
);
|
||||
@@ -484,9 +525,15 @@ impl acp::Agent for MvpAgent {
|
||||
err.message = e.to_string();
|
||||
err
|
||||
})?;
|
||||
{
|
||||
let mut sampling_config = self.sampling_config.borrow_mut();
|
||||
sampling_config.api_key = Some(auth.key.clone());
|
||||
// C1: hot-swap FIRST, then let the authority read the fresh
|
||||
// token back where it belongs. Nothing hand-carries `auth.key`
|
||||
// to the shared config any more — the stamp is whatever
|
||||
// credential governs that config's own model + endpoint, which
|
||||
// for a session whose current model is another provider's
|
||||
// subscription model is that provider's pooled token, and for a
|
||||
// third-party host is nothing at all.
|
||||
self.auth_manager.hot_swap(auth.clone());
|
||||
if self.stamp_session_credential(true) {
|
||||
tracing::debug!(
|
||||
"auth: kimi.com/oidc handler set api_key (SessionToken)"
|
||||
);
|
||||
@@ -496,7 +543,6 @@ impl acp::Agent for MvpAgent {
|
||||
None,
|
||||
);
|
||||
}
|
||||
self.auth_manager.hot_swap(auth.clone());
|
||||
self.emit_settings_update_notification();
|
||||
self.set_auth_method(arguments.method_id.clone());
|
||||
self.models_manager.on_auth_changed().await;
|
||||
@@ -508,15 +554,22 @@ impl acp::Agent for MvpAgent {
|
||||
);
|
||||
Ok(self.auth_response_with_meta())
|
||||
}
|
||||
auth_method::MOONSHOT_CN_METHOD_ID | auth_method::MOONSHOT_AI_METHOD_ID => {
|
||||
let platform = auth_method::moonshot_platform_for_method_id(
|
||||
_ => {
|
||||
if let Some(platform) = auth_method::AuthMethodKind::from_id(
|
||||
&arguments.method_id,
|
||||
)
|
||||
.expect("match arm guarantees a moonshot method id");
|
||||
self.authenticate_moonshot(platform, arguments.method_id.clone())
|
||||
.oauth_platform()
|
||||
{
|
||||
self.authenticate_oauth_platform(platform, arguments).await
|
||||
} else if let Some(platform) =
|
||||
auth_method::platform_for_method_id(&arguments.method_id)
|
||||
{
|
||||
self.authenticate_api_key_platform(
|
||||
platform,
|
||||
arguments.method_id.clone(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
_ => {
|
||||
} else {
|
||||
Err(
|
||||
acp::Error::invalid_params()
|
||||
.data(
|
||||
@@ -526,6 +579,7 @@ impl acp::Agent for MvpAgent {
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
async fn new_session(
|
||||
&self,
|
||||
arguments: acp::NewSessionRequest,
|
||||
@@ -754,6 +808,7 @@ impl acp::Agent for MvpAgent {
|
||||
persisted_signals: None,
|
||||
persisted_plan_mode: None,
|
||||
persisted_goal_mode: None,
|
||||
persisted_graph_mode: None,
|
||||
persisted_announcement_state: None,
|
||||
session_meta: arguments.meta.as_ref(),
|
||||
model_agent_type: model_agent_type.as_deref(),
|
||||
@@ -985,6 +1040,7 @@ impl acp::Agent for MvpAgent {
|
||||
signals: persisted_signals,
|
||||
announcement_state: persisted_announcement_state,
|
||||
goal_mode_state: _persisted_goal_mode,
|
||||
graph_mode_state: _persisted_graph_mode,
|
||||
} = persistence_info;
|
||||
let restored_awaiting_plan_approval = persisted_plan_mode
|
||||
.as_ref()
|
||||
@@ -1207,6 +1263,7 @@ impl acp::Agent for MvpAgent {
|
||||
persisted_signals,
|
||||
persisted_plan_mode,
|
||||
persisted_goal_mode: _persisted_goal_mode,
|
||||
persisted_graph_mode: _persisted_graph_mode,
|
||||
persisted_announcement_state,
|
||||
session_meta: request_meta.as_ref(),
|
||||
model_agent_type: persisted_agent_name.as_deref(),
|
||||
|
||||
@@ -25,20 +25,35 @@ impl MvpAgent {
|
||||
primary: &SamplingConfig,
|
||||
) -> Result<(OaiCompatClient, String), acp::Error> {
|
||||
let slug = self.resolve_session_summary_model();
|
||||
let session_key = self.auth_manager.current_or_expired().map(|a| a.key.clone());
|
||||
let models = self.models_manager.models();
|
||||
// Resolve the aux token by the summary model's OWN platform AND
|
||||
// endpoint: a grok (oauth-platform) summary model draws its pooled grok
|
||||
// token or `None`, and an API-key registry platform draws NOTHING —
|
||||
// NEVER the primary Kimi session token (which `resolve_credentials`
|
||||
// would otherwise stamp onto an api.x.ai / api.deepseek.com request).
|
||||
// The first-party subscription channel still gets the primary
|
||||
// (byte-identical).
|
||||
let authority = self.credential_authority();
|
||||
let session_key = authority.credential_for_slug(&models, None, &slug);
|
||||
let endpoints = self.models_manager.endpoints();
|
||||
let alpha_test_key = self.cfg.borrow().endpoints.alpha_test_key.clone();
|
||||
let config = match crate::agent::config::resolve_aux_model_sampling_config(
|
||||
&slug,
|
||||
&models,
|
||||
&endpoints,
|
||||
session_key.as_deref(),
|
||||
session_key.as_ref(),
|
||||
alpha_test_key,
|
||||
) {
|
||||
Some(mut cfg) => {
|
||||
cfg.attribution_callback = primary.attribution_callback.clone();
|
||||
cfg.bearer_resolver = primary.bearer_resolver.clone();
|
||||
// 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. The chokepoint
|
||||
// resolves the resolver from the SUMMARY model's platform +
|
||||
// endpoint instead; the session's is never even read.
|
||||
cfg.bearer_resolver =
|
||||
self.summary_bearer_resolver(&models, &slug, &cfg.base_url);
|
||||
cfg.max_retries = primary.max_retries;
|
||||
cfg
|
||||
}
|
||||
@@ -52,6 +67,36 @@ impl MvpAgent {
|
||||
let client = OaiCompatClient::new(config).map_err(map_sampling_err_to_acp)?;
|
||||
Ok((client, model))
|
||||
}
|
||||
/// The `bearer_resolver` the SESSION-SUMMARY client may carry — the SAME
|
||||
/// rule the session actor's aux path applies
|
||||
/// ([`crate::session::acp_session::sampler_turn::aux_bearer_resolver_for`]),
|
||||
/// not a second copy of it.
|
||||
///
|
||||
/// M3 completed: the aux path was gated on the session-token gate and this
|
||||
/// one was not, so an api-key / house-key session whose
|
||||
/// `[model.session-summary]` block carries its OWN `env_key` on the
|
||||
/// session's own coding endpoint had that key REPLACED on the wire by the
|
||||
/// primary bearer on every summary request. Named (rather than inlined
|
||||
/// above) so the gate is reachable from a test — the ungated version stayed
|
||||
/// green because the resolver is consumed by `OaiCompatClient::new`.
|
||||
///
|
||||
/// Aux slugs are not the session's selection, so `current_key = None`
|
||||
/// (H-b then refuses a collided slug rather than guessing its OAuth twin).
|
||||
pub(super) fn summary_bearer_resolver(
|
||||
&self,
|
||||
models: &indexmap::IndexMap<String, ModelEntry>,
|
||||
slug: &str,
|
||||
base_url: &str,
|
||||
) -> Option<kigi_sampler::SharedBearerResolver> {
|
||||
let auth_method = self.auth_method_id.load();
|
||||
crate::session::acp_session::sampler_turn::aux_bearer_resolver_for(
|
||||
&self.credential_authority(),
|
||||
auth_method.as_deref(),
|
||||
crate::agent::models::platform_for_slug(models, None, slug),
|
||||
crate::agent::config::resolve_model_auth_facts(slug).byok,
|
||||
base_url,
|
||||
)
|
||||
}
|
||||
/// `true` for session-based ACP auth methods.
|
||||
fn is_session_based_auth(&self) -> bool {
|
||||
self.auth_method_id
|
||||
@@ -310,8 +355,13 @@ impl MvpAgent {
|
||||
pub(crate) fn command_availability(
|
||||
&self,
|
||||
) -> crate::session::slash_commands::CommandAvailability {
|
||||
let goal = self.cfg.borrow().resolve_goal().value;
|
||||
crate::session::slash_commands::CommandAvailability {
|
||||
goal: self.cfg.borrow().resolve_goal().value,
|
||||
goal,
|
||||
// Same convention as /goal: the flag is known at initialize
|
||||
// time, so advertise pre-session; the in-session path
|
||||
// re-checks the live toolset.
|
||||
graph: goal && self.cfg.borrow().resolve_graph().value,
|
||||
..crate::session::slash_commands::CommandAvailability::default()
|
||||
}
|
||||
}
|
||||
@@ -447,18 +497,19 @@ impl MvpAgent {
|
||||
)
|
||||
.await
|
||||
}
|
||||
/// `authenticate(moonshot-cn / moonshot-ai)`: interactive open-platform
|
||||
/// API-key login from the welcome picker.
|
||||
/// `authenticate(<api-key platform id>)`: interactive API-key login from
|
||||
/// the welcome picker for any non-OAuth registry platform.
|
||||
///
|
||||
/// Reloads the platform keys from disk+env (the TUI persists the pasted
|
||||
/// key to `[platforms.<id>]` in config.toml immediately before this call),
|
||||
/// fails with an actionable error when none is configured, validates the
|
||||
/// key against `GET {platform_base}/models`, then marks the session
|
||||
/// authenticated exactly like an external API key: publish the method id
|
||||
/// (NOT session-based — no token refresh), swap the freshly-stamped config
|
||||
/// into the models manager, and trigger the model sync so the catalog
|
||||
/// gains the platform's entries. The key itself is never logged.
|
||||
pub(super) async fn authenticate_moonshot(
|
||||
/// key to auth.json under the platform-id scope immediately before this
|
||||
/// call), fails with an actionable error when none is configured,
|
||||
/// validates the key against `GET {platform_base}/models`, then marks the
|
||||
/// session authenticated exactly like an external API key: publish the
|
||||
/// method id (NOT session-based — no token refresh), swap the
|
||||
/// freshly-stamped config into the models manager, and trigger the model
|
||||
/// sync so the catalog gains the platform's entries. The key itself is
|
||||
/// never logged.
|
||||
pub(super) async fn authenticate_api_key_platform(
|
||||
&self,
|
||||
platform: kigi_models::PlatformId,
|
||||
method_id: acp::AuthMethodId,
|
||||
@@ -475,10 +526,11 @@ impl MvpAgent {
|
||||
Some("platform_key_invalid_or_missing"),
|
||||
);
|
||||
})?;
|
||||
// Swap the on-disk config (now carrying the key) into the models
|
||||
// manager so `apply_platform_credentials` stamps the platform's
|
||||
// catalog entries; a parse failure keeps the last-known-good config
|
||||
// (`on_auth_changed` below still re-resolves keys from disk itself).
|
||||
// Rebuild the catalog from the on-disk config: the rebuild freshly
|
||||
// resolves platform keys (env > auth.json > config), so the key just
|
||||
// persisted to auth.json is stamped onto the platform's entries; a
|
||||
// parse failure keeps the last-known-good config (`on_auth_changed`
|
||||
// below still re-resolves keys from disk itself).
|
||||
match crate::config::load_effective_config()
|
||||
.map_err(|e| e.to_string())
|
||||
.and_then(|raw| crate::agent::config::Config::new_from_toml_cfg(&raw))
|
||||
@@ -507,6 +559,78 @@ impl MvpAgent {
|
||||
.and_then(|v| v.as_object().cloned());
|
||||
Ok(AuthenticateResponse::new().meta(meta))
|
||||
}
|
||||
/// `authenticate(<generic-oauth platform id>)`: interactive device-code
|
||||
/// login for a `uses_oauth` platform carrying an `OAuthConfig` (xai-grok).
|
||||
///
|
||||
/// Uses a per-provider [`AuthManager`] scoped to the platform's `scope_key`
|
||||
/// (NOT the primary Kimi manager) so the minted session is persisted under
|
||||
/// its own `auth.json` scope, then triggers a catalog re-sync so the
|
||||
/// platform's models appear (their bearer is resolved per-provider at fetch
|
||||
/// / sampling time). The tokens are never logged.
|
||||
pub(super) async fn authenticate_oauth_platform(
|
||||
&self,
|
||||
platform: kigi_models::PlatformId,
|
||||
arguments: acp::AuthenticateRequest,
|
||||
) -> Result<AuthenticateResponse, acp::Error> {
|
||||
let method_id = arguments.method_id.clone();
|
||||
let oauth = platform
|
||||
.oauth()
|
||||
.expect("oauth_platform() guarantees a device-code OAuthConfig");
|
||||
let auth_meta = AuthRequestMeta::from_json(arguments.meta.as_ref());
|
||||
tracing::info!(
|
||||
method = method_id.0.as_ref(),
|
||||
headless = auth_meta.headless,
|
||||
reauth = auth_meta.reauth,
|
||||
"auth: generic oauth device login",
|
||||
);
|
||||
// M4: the SAME home the pool reads from
|
||||
// (`oauth_registry::pool_home()`), not `kigi_home()` directly —
|
||||
// identical in production, but a login driven from a lib test would
|
||||
// otherwise write into the developer's real `~/.kigi` while every
|
||||
// inference-time lookup read the disposable test pool home.
|
||||
let kigi_home = crate::auth::oauth_registry::pool_home();
|
||||
let auth_manager =
|
||||
std::sync::Arc::new(crate::auth::AuthManager::new_oauth_provider(&kigi_home, oauth));
|
||||
auth_manager.configure_refresher();
|
||||
|
||||
let flow_result = if !auth_meta.headless {
|
||||
let (url_tx, url_rx) = tokio::sync::oneshot::channel();
|
||||
let (code_tx, code_rx) = tokio::sync::mpsc::channel(1);
|
||||
*self.auth_code_tx.borrow_mut() = Some(code_tx);
|
||||
*self.auth_url_rx.borrow_mut() = Some(url_rx);
|
||||
let result = crate::auth::run_oauth_provider_flow(
|
||||
&auth_manager,
|
||||
oauth,
|
||||
auth_meta.reauth,
|
||||
Some(crate::auth::AuthChannels {
|
||||
url_tx: Some(url_tx),
|
||||
code_rx,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
*self.auth_code_tx.borrow_mut() = None;
|
||||
*self.auth_url_rx.borrow_mut() = None;
|
||||
result
|
||||
} else {
|
||||
crate::auth::run_oauth_provider_flow(&auth_manager, oauth, auth_meta.reauth, None).await
|
||||
};
|
||||
|
||||
let (_auth, _did_auth) = flow_result.map_err(|e| {
|
||||
emit_login_span(false, method_id.0.as_ref(), None, Some("login_flow_failed"));
|
||||
let mut err = acp::Error::auth_required();
|
||||
err.message = e.to_string();
|
||||
err
|
||||
})?;
|
||||
|
||||
// Do NOT stamp this token onto the shared sampling_config: it authorizes
|
||||
// ONLY this platform's models (api.x.ai/v1), not the primary session.
|
||||
// The catalog re-sync below resolves it per-provider.
|
||||
self.set_auth_method(method_id.clone());
|
||||
self.models_manager.on_auth_changed().await;
|
||||
emit_login_span(true, method_id.0.as_ref(), None, None);
|
||||
Ok(self.auth_response_with_meta())
|
||||
}
|
||||
|
||||
pub(crate) fn deployment_key(&self) -> Option<String> {
|
||||
self.cfg.borrow().endpoints.deployment_key.clone()
|
||||
}
|
||||
@@ -586,21 +710,60 @@ impl MvpAgent {
|
||||
);
|
||||
Ok(entry.clone())
|
||||
}
|
||||
/// This agent's credential chokepoint: the EFFECTIVE endpoints (so a
|
||||
/// managed `[endpoints] coding_api_base_url` deployment keeps the session
|
||||
/// bearer — H3) plus the primary manager, which the authority keeps private.
|
||||
pub(crate) fn credential_authority(
|
||||
&self,
|
||||
) -> crate::auth::credential_authority::CredentialAuthority {
|
||||
crate::auth::credential_authority::CredentialAuthority::new(
|
||||
self.models_manager.endpoints(),
|
||||
Some(self.auth_manager.clone()),
|
||||
)
|
||||
}
|
||||
/// The SESSION credential for `model`, resolved by the model's OWN platform
|
||||
/// AND endpoint at the chokepoint — the single guard against the
|
||||
/// api_key-channel token leak (C1).
|
||||
///
|
||||
/// A subscription-OAuth model draws from ITS OWN pooled manager (`None`,
|
||||
/// never the Kimi key, when that provider has no stored session). Every
|
||||
/// API-key registry platform and every `[model.*]` block pointed at a
|
||||
/// third-party host gets `None`: `resolve_credentials` would otherwise take
|
||||
/// the `else if let Some(key) = session_key` arm and set
|
||||
/// `api_key = <Kimi bearer>` with the THIRD-PARTY `base_url`, which
|
||||
/// `SamplingClient` builds into `Authorization: Bearer …` — reachable with
|
||||
/// ZERO configuration, since `default_models.json` bundles `moonshot-cn/*`
|
||||
/// and `moonshot-ai/*` entries a Kimi-subscription user sees on first
|
||||
/// launch / offline. The first-party subscription channel uses the primary,
|
||||
/// and only under a session-based auth method — byte-identical.
|
||||
///
|
||||
/// SECURITY: the resolved token is never logged.
|
||||
fn session_token_for_model(
|
||||
&self,
|
||||
model: &ModelEntry,
|
||||
) -> Option<crate::auth::credential_authority::SessionCredential> {
|
||||
let is_primary_channel = crate::auth::credential_authority::entry_platform(model)
|
||||
.is_none_or(|platform| platform.oauth().is_none());
|
||||
if is_primary_channel && !self.is_session_based_auth() {
|
||||
return None;
|
||||
}
|
||||
self.credential_authority().credential_for_model(model)
|
||||
}
|
||||
pub(crate) fn prepare_sampling_config_for_model(
|
||||
&self,
|
||||
model: &ModelEntry,
|
||||
origin_client: Option<crate::http::OriginClientInfo>,
|
||||
) -> SamplingConfig {
|
||||
let session = if self.is_session_based_auth() {
|
||||
self.auth_manager.current_or_expired()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
// Resolve the session token by the MODEL's platform, not the primary
|
||||
// auth method: an oauth-platform model (xai-grok) uses its OWN
|
||||
// pool-backed token (`None` — never the Kimi key — when the user has not
|
||||
// logged into that provider), closing the api_key-channel leak where the
|
||||
// primary Kimi session token was stamped onto a grok request. A
|
||||
// first-party / Kimi model is unchanged. GUARANTEE: a grok model's
|
||||
// api_key is its own grok token or `None`, never the primary Kimi key.
|
||||
let session = self.session_token_for_model(model);
|
||||
let has_session_key = session.is_some();
|
||||
let mut credentials = resolve_credentials(
|
||||
model,
|
||||
session.as_ref().map(|a| a.key.as_str()),
|
||||
);
|
||||
let mut credentials = resolve_credentials(model, session.as_ref());
|
||||
if !has_session_key && credentials.auth_type == kigi_chat_state::AuthType::ApiKey
|
||||
&& !model.has_own_credentials() && self.is_session_based_auth()
|
||||
{
|
||||
@@ -787,7 +950,8 @@ impl MvpAgent {
|
||||
models_manager: crate::agent::models::ModelsManager,
|
||||
) -> Self {
|
||||
models_manager.set_gateway(gateway.clone());
|
||||
let sampling_config = models_manager.sampling_config();
|
||||
// H-a: the config AND the platform it was built from, from ONE call.
|
||||
let baseline = models_manager.sampling_config();
|
||||
let storage_mode = cfg.storage_mode;
|
||||
let default_yolo_mode = cfg.default_yolo_mode;
|
||||
let default_auto_mode = cfg.default_auto_mode;
|
||||
@@ -842,7 +1006,8 @@ impl MvpAgent {
|
||||
models_manager,
|
||||
cfg: RefCell::new(cfg.clone()),
|
||||
auth_method_id: crate::agent::auth_method::new_shared_auth_method_id(None),
|
||||
sampling_config: RefCell::new(sampling_config),
|
||||
sampling_config: RefCell::new(baseline.config),
|
||||
sampling_config_platform: std::cell::Cell::new(baseline.platform),
|
||||
auth_manager,
|
||||
auth_code_tx: RefCell::new(None),
|
||||
auth_url_rx: RefCell::new(None),
|
||||
@@ -1432,32 +1597,128 @@ impl MvpAgent {
|
||||
);
|
||||
(serde_json::json!({ "options" : config_options }), serde_json::json!(detail))
|
||||
}
|
||||
/// The registry platform the SHARED `sampling_config` routes to: the one
|
||||
/// captured WITH the config when it was built, never a fresh lookup.
|
||||
///
|
||||
/// H-a: `ModelsManager::sampling_config` builds the shared config from the
|
||||
/// catalog entry `current_model_id()` named AT THAT MOMENT, and the config
|
||||
/// is never rebuilt. A guard that re-resolves `config.model` (the BARE
|
||||
/// routing slug) against the LIVE cell therefore answers about a DIFFERENT
|
||||
/// entry the moment the two drift — and they drift on every non-Leader model
|
||||
/// switch (`handlers/model_switch.rs`) and on catalog reselection.
|
||||
/// `entry_for_slug_resolution`'s `entry.info.model == slug` test then fails
|
||||
/// and the resolution falls through to `resolve_catalog_key`'s `.rev()`
|
||||
/// scan; `PlatformId::ALL` lists `kimi-code` first and its API-key twin
|
||||
/// `kimi-coding` 19th, so the LAST match is the twin, which takes no session
|
||||
/// credential. For a Kimi-subscription user who also has `KIMI_API_KEY` set,
|
||||
/// a post-expiry `kigi login` then found no governing manager and — because
|
||||
/// the stamp only overwrites on success — left the EXPIRED bearer in place:
|
||||
/// every unresolved-model fallback and every subagent baseline turn 401'd
|
||||
/// until restart.
|
||||
fn shared_config_platform(&self) -> Option<kigi_models::PlatformId> {
|
||||
self.sampling_config_platform.get()
|
||||
}
|
||||
/// H2/C1: stamp the shared `sampling_config`'s OWN governing session
|
||||
/// credential — whatever the authority says that is for this config's model
|
||||
/// and endpoint.
|
||||
///
|
||||
/// The shared config is not inert: `resolve_sampling_config_for_model`
|
||||
/// returns it verbatim whenever a model id fails to resolve, and
|
||||
/// `SubagentSpawnContext` clones it as every subagent's baseline — so an
|
||||
/// `api_key` stamped here reaches the wire against whatever `base_url` the
|
||||
/// config carries. The login/seed sites used to stamp it unconditionally,
|
||||
/// exactly the mistake `authenticate_oauth_platform` already documents ("Do
|
||||
/// NOT stamp this token onto the shared sampling_config").
|
||||
///
|
||||
/// C1: this function does NOT take a credential. The previous shape took
|
||||
/// `key: String` — always the primary Kimi bearer — and guarded it with a
|
||||
/// predicate asking whether **a** session credential may ride. For a
|
||||
/// subscription-OAuth platform at its own host that is correctly `true`, but
|
||||
/// the credential that may ride there is that platform's POOLED token: a
|
||||
/// Claude Pro/Max user running `kigi login` stamped the Kimi subscription
|
||||
/// bearer onto a config routed at `api.anthropic.com`. Asking
|
||||
/// `credential_for` instead makes the question and the credential the same
|
||||
/// object, so the pairing cannot be wrong — and there is no primary handle
|
||||
/// here to hand-carry (M2).
|
||||
///
|
||||
/// `overwrite = false` keeps the historical "only if missing" seeding
|
||||
/// behaviour; the login handlers pass `true` because a fresh login must
|
||||
/// replace a stale bearer, and they call this AFTER the manager holds the
|
||||
/// new token so it is read back through the authority.
|
||||
///
|
||||
/// SECURITY: the token is never logged.
|
||||
pub(super) fn stamp_session_credential(&self, overwrite: bool) -> bool {
|
||||
let (model, base_url) = {
|
||||
let sampling_config = self.sampling_config.borrow();
|
||||
if !overwrite && sampling_config.api_key.is_some() {
|
||||
return false;
|
||||
}
|
||||
(
|
||||
sampling_config.model.clone(),
|
||||
sampling_config.base_url.clone(),
|
||||
)
|
||||
};
|
||||
let platform = self.shared_config_platform();
|
||||
let Some(credential) = self.credential_authority().credential_for(platform, &base_url)
|
||||
else {
|
||||
tracing::debug!(
|
||||
model = model.as_str(),
|
||||
"auth: no session credential governs the shared sampling config \
|
||||
(its model + endpoint take none, or the governing provider has no session)"
|
||||
);
|
||||
return false;
|
||||
};
|
||||
self.sampling_config.borrow_mut().api_key = Some(credential.expose().to_owned());
|
||||
true
|
||||
}
|
||||
/// Whether the shared `sampling_config` may receive a credential that
|
||||
/// authorizes only the SESSION's own first-party endpoint — the house
|
||||
/// `KIGI_API_KEY` the `xai.api_key` login handler reads from the
|
||||
/// environment, which the authority does not own and so cannot produce.
|
||||
///
|
||||
/// The house key rides the [`CredentialClass::Primary`] channel and no
|
||||
/// other. `Pooled` is deliberately excluded: a subscription-OAuth platform's
|
||||
/// own host DOES take a session credential, but that credential is the
|
||||
/// platform's pooled token, where the house key has no business (C1).
|
||||
pub(super) fn shared_config_takes_house_key(&self) -> bool {
|
||||
let base_url = self.sampling_config.borrow().base_url.clone();
|
||||
matches!(
|
||||
self.credential_authority()
|
||||
.credential_class(self.shared_config_platform(), &base_url),
|
||||
crate::auth::credential_authority::CredentialClass::Primary
|
||||
)
|
||||
}
|
||||
/// Seed the global sampling config with login auth when available.
|
||||
///
|
||||
/// Only sets the `api_key` if missing. Does NOT resolve `base_url` from
|
||||
/// Only sets the `api_key` if missing, and only with the credential that
|
||||
/// governs the config's own model + endpoint (see
|
||||
/// [`Self::stamp_session_credential`]). Does NOT resolve `base_url` from
|
||||
/// `current_model_id` — that's deferred to session creation time to avoid
|
||||
/// cross-client contamination in leader mode (where `current_model_id` is
|
||||
/// shared mutable state).
|
||||
pub(super) fn seed_client_config_auth_if_available(&self) {
|
||||
let mut sampling_config = self.sampling_config.borrow_mut();
|
||||
if sampling_config.api_key.is_none() {
|
||||
if let Some(auth) = self.auth_manager.current_or_expired() {
|
||||
sampling_config.api_key = Some(auth.key);
|
||||
if self.sampling_config.borrow().api_key.is_some() {
|
||||
return;
|
||||
}
|
||||
if self.stamp_session_credential(false) {
|
||||
tracing::debug!("auth: seed_client_config set auth (SessionToken)");
|
||||
kigi_log::unified_log::debug(
|
||||
"auth: seed_client_config set auth (SessionToken)",
|
||||
None,
|
||||
None,
|
||||
);
|
||||
} else if !self
|
||||
return;
|
||||
}
|
||||
// No credential was stamped. Only the total-absence case is worth
|
||||
// warning about: a withheld-by-endpoint seed is the rule working.
|
||||
if self.auth_manager.current_or_expired().is_none()
|
||||
&& !self
|
||||
.models_manager
|
||||
.models()
|
||||
.values()
|
||||
.any(|m| m.has_own_credentials())
|
||||
{
|
||||
tracing::warn!(
|
||||
"No credentials found: no login token and no model api_key/env_key"
|
||||
);
|
||||
tracing::warn!("No credentials found: no login token and no model api_key/env_key");
|
||||
kigi_log::unified_log::warn(
|
||||
"No credentials found: no login token and no model api_key/env_key",
|
||||
None,
|
||||
@@ -1465,7 +1726,6 @@ impl MvpAgent {
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
/// Allocate the next monotonic telemetry turn number for a session.
|
||||
///
|
||||
/// Returns the current turn number and advances the counter. The counter is
|
||||
@@ -1677,6 +1937,7 @@ impl MvpAgent {
|
||||
persisted_signals,
|
||||
persisted_plan_mode,
|
||||
persisted_goal_mode,
|
||||
persisted_graph_mode,
|
||||
persisted_announcement_state,
|
||||
session_meta,
|
||||
model_agent_type,
|
||||
@@ -2111,6 +2372,7 @@ impl MvpAgent {
|
||||
let web_fetch_config = self.prepare_web_fetch_config();
|
||||
let write_file_enabled = self.cfg.borrow().resolve_write_file().value;
|
||||
let goal_enabled = self.cfg.borrow().resolve_goal().value;
|
||||
let graph_enabled = self.cfg.borrow().resolve_graph().value;
|
||||
let subagents_enabled = self.cfg.borrow().subagents_enabled;
|
||||
let ask_user_question_enabled = parse_ask_user_question_from_meta(session_meta)
|
||||
.unwrap_or_else(|| self.cfg.borrow().resolve_ask_user_question().value);
|
||||
@@ -2183,12 +2445,34 @@ impl MvpAgent {
|
||||
}
|
||||
let (mut handle, agent_system_prompt, session_thread) = {
|
||||
let _timer = crate::instrumentation_timer!("session.spawn_actor_call");
|
||||
let session_key = self.auth_manager.current_or_expired().map(|a| a.key);
|
||||
// Classify the credential's auth_type against the bearer this
|
||||
// model's endpoint would ACTUALLY receive, not the raw primary:
|
||||
// an API-key-platform model has no session credential at all, and
|
||||
// reading the primary here reported one.
|
||||
//
|
||||
// H-a: the platform is resolved with the SAME catalog key the
|
||||
// actor about to be spawned seeds itself with
|
||||
// (`selected_catalog_key_for_spawn`, spawn.rs), so this
|
||||
// classification and every per-turn decision that session makes
|
||||
// agree by construction instead of re-resolving the bare slug.
|
||||
let models_for_key = self.models_manager.models();
|
||||
let session_key = self.credential_authority().credential_for(
|
||||
crate::agent::models::platform_for_slug(
|
||||
&models_for_key,
|
||||
crate::agent::models::selected_catalog_key_for_spawn(
|
||||
&models_for_key,
|
||||
&session_model_id,
|
||||
)
|
||||
.as_deref(),
|
||||
sampling_config.model.as_str(),
|
||||
),
|
||||
&sampling_config.base_url,
|
||||
);
|
||||
let credentials = kigi_chat_state::Credentials {
|
||||
api_key: sampling_config.api_key.clone(),
|
||||
auth_type: crate::agent::config::resolve_chat_state_auth_type(
|
||||
sampling_config.model.as_str(),
|
||||
session_key.as_deref(),
|
||||
session_key.as_ref(),
|
||||
self.auth_type(),
|
||||
),
|
||||
alpha_test_key: self.alpha_test_key(),
|
||||
@@ -2314,6 +2598,7 @@ impl MvpAgent {
|
||||
persisted_signals,
|
||||
persisted_plan_mode,
|
||||
persisted_goal_mode,
|
||||
persisted_graph_mode,
|
||||
persisted_announcement_state,
|
||||
self.memory_config.clone(),
|
||||
feedback_flags,
|
||||
@@ -2328,6 +2613,7 @@ impl MvpAgent {
|
||||
app_builder_deployer_config,
|
||||
write_file_enabled,
|
||||
goal_enabled,
|
||||
graph_enabled,
|
||||
subagents_enabled,
|
||||
ask_user_question_enabled,
|
||||
client_hooks,
|
||||
|
||||
@@ -120,6 +120,7 @@ pub(crate) struct SessionSpawnOptions<'a> {
|
||||
pub persisted_signals: Option<crate::session::signals::SessionSignals>,
|
||||
pub persisted_plan_mode: Option<crate::session::plan_mode::PlanModeSnapshot>,
|
||||
pub persisted_goal_mode: Option<crate::session::goal_tracker::GoalOrchestration>,
|
||||
pub persisted_graph_mode: Option<crate::session::graph_tracker::GraphOrchestration>,
|
||||
pub persisted_announcement_state: Option<
|
||||
crate::session::announcement_state::AnnouncementState,
|
||||
>,
|
||||
@@ -257,6 +258,7 @@ pub(crate) fn chat_session_spawn_options<'a>(
|
||||
persisted_signals: None,
|
||||
persisted_plan_mode: None,
|
||||
persisted_goal_mode: None,
|
||||
persisted_graph_mode: None,
|
||||
persisted_announcement_state: None,
|
||||
session_meta,
|
||||
model_agent_type,
|
||||
@@ -499,6 +501,19 @@ pub struct MvpAgent {
|
||||
/// only api_key is written here (same for all clients). Per-session base_url
|
||||
/// is resolved at session creation time in `new_session` / `load_session`.
|
||||
pub(crate) sampling_config: RefCell<SamplingConfig>,
|
||||
/// The registry platform [`Self::sampling_config`] was BUILT from, captured
|
||||
/// in the same `ModelsManager::sampling_config()` call that produced it.
|
||||
///
|
||||
/// H-a: the shared config is built ONCE (`Self::with_models`) and never
|
||||
/// rebuilt, but `ModelsManager::current_model_id()` moves on every
|
||||
/// non-Leader model switch and on catalog reselection. Its guards
|
||||
/// (`stamp_session_credential`, `shared_config_takes_house_key`)
|
||||
/// therefore read THIS cell, never the live one: after a switch the live
|
||||
/// cell names a different entry, and re-resolving the config's bare slug
|
||||
/// against it fell through to `resolve_catalog_key`'s `.rev()` scan —
|
||||
/// answering the API-key twin, which takes no session credential, so a
|
||||
/// successful `kigi login` silently failed to replace the expired bearer.
|
||||
pub(crate) sampling_config_platform: std::cell::Cell<Option<kigi_models::PlatformId>>,
|
||||
pub(crate) auth_manager: Arc<AuthManager>,
|
||||
pub(crate) models_manager: crate::agent::models::ModelsManager,
|
||||
/// Forwards pasted codes from `handle_auth_submit_code` to the auth flow.
|
||||
|
||||
@@ -1158,6 +1158,80 @@ fn build_agent_with_auth(auth: crate::auth::KimiAuth) -> MvpAgent {
|
||||
let cfg = AgentConfig::default();
|
||||
MvpAgent::new(gateway, &cfg, auth_manager, None).expect("valid test config")
|
||||
}
|
||||
/// Regression (token-leak, Facet B via the api_key channel): under a
|
||||
/// session-based (Kimi) primary auth method, `prepare_sampling_config_for_model`
|
||||
/// must resolve the session token by the MODEL's platform — so a grok
|
||||
/// (oauth-platform) model NEVER carries the primary Kimi session key as its
|
||||
/// `api_key`, while a first-party Kimi model still does. Relies on the
|
||||
/// process-global OAuth pool (there is no per-session snapshot).
|
||||
///
|
||||
/// Reverting the fix (resolving the session token from the primary regardless of
|
||||
/// the model's platform) fails the grok assertion below — it would stamp the
|
||||
/// live Kimi key on a request bound for api.x.ai.
|
||||
#[tokio::test]
|
||||
#[serial_test::serial]
|
||||
async fn prepare_sampling_config_never_stamps_kimi_key_on_grok_model() {
|
||||
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::{EndpointsConfig, ModelEntry};
|
||||
use kigi_test_support::EnvGuard;
|
||||
|
||||
const KIMI_KEY: &str = "kimi-session-secret-DO-NOT-LEAK";
|
||||
|
||||
// No ambient BYOK env key: a grok model with no stored oauth session then
|
||||
// resolves to no api_key at all, rather than a global-key fallback that could
|
||||
// mask the leak under test.
|
||||
let _house = EnvGuard::unset(HOUSE_API_KEY_ENV_VAR);
|
||||
let _xai = EnvGuard::unset(XAI_API_KEY_ENV_VAR);
|
||||
let _legacy = EnvGuard::unset(LEGACY_XAI_API_KEY_ENV_VAR);
|
||||
|
||||
// Primary: a live Kimi session token under a session-based auth method.
|
||||
let agent = build_agent_with_auth(crate::auth::KimiAuth {
|
||||
key: KIMI_KEY.to_string(),
|
||||
auth_mode: crate::auth::AuthMode::OAuth,
|
||||
..crate::auth::KimiAuth::test_default()
|
||||
});
|
||||
agent.set_auth_method(acp::AuthMethodId::new(CACHED_TOKEN_AUTH_METHOD_ID));
|
||||
|
||||
let endpoints = EndpointsConfig::default();
|
||||
|
||||
// First-party SUBSCRIPTION model (kimi-code): the primary session key IS its
|
||||
// api_key — the byte-identical primary path, and proof the Kimi token is
|
||||
// live (so it WOULD leak if mis-routed onto a grok request). This assertion
|
||||
// also confirms the session-based primary path is active.
|
||||
//
|
||||
// This used to use `moonshot-cn/kimi-k2-0905-preview` and assert the SAME
|
||||
// thing, which encoded the C1 defect: moonshot-cn is an API-key registry
|
||||
// platform on `api.moonshot.cn`, NOT first-party, so "must carry the primary
|
||||
// session key" was asserting the leak. `api_key_channel_leak_tests` now pins
|
||||
// the opposite for every moonshot entry.
|
||||
let mut kimi_model = ModelEntry::fallback("kimi-for-coding", &endpoints);
|
||||
kimi_model.info.id = Some("kimi-code/kimi-for-coding".to_string());
|
||||
kimi_model.info.base_url = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url.to_string();
|
||||
assert!(!kimi_model.has_own_credentials());
|
||||
let kimi_cfg = agent.prepare_sampling_config_for_model(&kimi_model, None);
|
||||
assert_eq!(
|
||||
kimi_cfg.api_key.as_deref(),
|
||||
Some(KIMI_KEY),
|
||||
"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
|
||||
// pool-backed manager, INDEPENDENT of the Kimi primary — so its api_key can
|
||||
// NEVER be the Kimi session key.
|
||||
let mut grok_model = ModelEntry::fallback("grok-4-latest", &endpoints);
|
||||
grok_model.info.id = Some("xai-grok/grok-4-latest".to_string());
|
||||
assert!(!grok_model.has_own_credentials());
|
||||
let grok_cfg = agent.prepare_sampling_config_for_model(&grok_model, None);
|
||||
assert_ne!(
|
||||
grok_cfg.api_key.as_deref(),
|
||||
Some(KIMI_KEY),
|
||||
"LEAK: a grok model must never carry the primary Kimi session key as api_key"
|
||||
);
|
||||
}
|
||||
/// Regression: boot-time plugin discovery is deferred past ACP
|
||||
/// `initialize`, so the shared plugin registry starts empty.
|
||||
/// `resolve_mcp_servers` reads that snapshot to merge plugin-contributed
|
||||
@@ -1214,6 +1288,10 @@ async fn ensure_plugin_registry_lazily_populates_snapshot() {
|
||||
);
|
||||
}
|
||||
mod subagent_spawn_context_tests;
|
||||
/// LEAK guard for the `api_key` channel (C1/C2), through the real
|
||||
/// `prepare_sampling_config_for_model` resolution path.
|
||||
mod api_key_channel_leak_tests;
|
||||
mod chokepoint_leak_tests;
|
||||
/// No load in flight and no session → the wait returns immediately
|
||||
/// (the caller then surfaces "unknown session id" exactly as before).
|
||||
#[tokio::test]
|
||||
@@ -1904,10 +1982,12 @@ async fn cached_token_fallthrough_prefers_api_key_for_deployment_key() {
|
||||
#[serial_test::serial]
|
||||
async fn cached_token_fallthrough_falls_to_kigi_com_without_credentials() {
|
||||
use crate::agent::auth_method::{
|
||||
KIMI_CODE_METHOD_ID, LEGACY_XAI_API_KEY_ENV_VAR, XAI_API_KEY_ENV_VAR,
|
||||
HOUSE_API_KEY_ENV_VAR, KIMI_CODE_METHOD_ID, LEGACY_XAI_API_KEY_ENV_VAR,
|
||||
XAI_API_KEY_ENV_VAR,
|
||||
};
|
||||
use kigi_test_support::EnvGuard;
|
||||
let _lockdown = EnvGuard::unset("KIGI_DISABLE_API_KEY_AUTH");
|
||||
let _house = EnvGuard::unset(HOUSE_API_KEY_ENV_VAR);
|
||||
let _new = EnvGuard::unset(XAI_API_KEY_ENV_VAR);
|
||||
let _legacy = EnvGuard::unset(LEGACY_XAI_API_KEY_ENV_VAR);
|
||||
let agent = build_minimal_agent_for_tests();
|
||||
|
||||
@@ -0,0 +1,314 @@
|
||||
//! 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;
|
||||
|
||||
pub(super) 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]`.
|
||||
pub(super) 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.
|
||||
pub(super) 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.
|
||||
pub(super) 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: make `CredentialAuthority::governing_manager`'s
|
||||
/// `Some(platform) => None` arm return `self.primary.clone()` and every `api_key`
|
||||
/// below becomes `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: make `CredentialAuthority::is_session_coding_endpoint` return
|
||||
/// `true` unconditionally and `api_key` here becomes `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"
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,589 @@
|
||||
//! LEAK GUARD (the chokepoint's own call sites) — H1, H2 and the H3 regression.
|
||||
//!
|
||||
//! The companion of `api_key_channel_leak_tests`, which drives
|
||||
//! `MvpAgent::prepare_sampling_config_for_model`. These three pin the OTHER
|
||||
//! producers of an outgoing credential: `ModelsManager::sampling_config()` (the
|
||||
//! agent-wide baseline), the shared `MvpAgent::sampling_config` the login/seed
|
||||
//! paths stamp, and the SESSION's effective coding endpoint as configured from
|
||||
//! config.toml rather than the environment.
|
||||
//!
|
||||
//! Every assertion reads what the real resolution path produced; nothing is
|
||||
//! hand-stamped.
|
||||
|
||||
use super::super::*;
|
||||
use super::api_key_channel_leak_tests::{
|
||||
KIMI_TOKEN, kimi_session_agent, platform_entry, without_ambient_byok_env,
|
||||
};
|
||||
use crate::agent::auth_method::CACHED_TOKEN_AUTH_METHOD_ID;
|
||||
use crate::agent::config::{Config as AgentConfig, EndpointsConfig, ModelEntry};
|
||||
use crate::auth::credential_authority::CredentialClass;
|
||||
use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig};
|
||||
use kigi_sampler::BearerResolver;
|
||||
use kigi_test_support::EnvGuard;
|
||||
|
||||
/// Re-seed the shared config the way `MvpAgent::with_models` does: the config
|
||||
/// AND the platform it was BUILT from, from ONE `ModelsManager::sampling_config()`
|
||||
/// against whatever `current_model_id` names right now. A test can therefore
|
||||
/// never set one without the other — which is the whole point of H-a.
|
||||
fn rebuild_shared_config(agent: &MvpAgent) {
|
||||
let baseline = agent.models_manager.sampling_config();
|
||||
agent.sampling_config_platform.set(baseline.platform);
|
||||
*agent.sampling_config.borrow_mut() = baseline.config;
|
||||
}
|
||||
|
||||
/// H1 — `ModelsManager::sampling_config()`, the OTHER `api_key` producer, and a
|
||||
/// byte-for-byte repeat of the round-1 defect: it resolved the session bearer
|
||||
/// itself (`platform.oauth()`, else `auth_manager.current_or_expired()`), so
|
||||
/// every non-OAuth platform got the primary Kimi token.
|
||||
///
|
||||
/// This config is not incidental — it is the `MvpAgent` baseline
|
||||
/// (`Self::with_models`), which `resolve_sampling_config_for_model` returns
|
||||
/// verbatim for an unresolved model id and `SubagentSpawnContext` clones as
|
||||
/// every subagent's baseline, so its `api_key` reaches the wire against its own
|
||||
/// `base_url`. Zero-config repro: a Kimi subscription + a bundled
|
||||
/// `moonshot-cn/*` default.
|
||||
///
|
||||
/// Revert-to-red (L: this edit COMPILES — the previous wording named a
|
||||
/// `Option<String>` argument that the `Option<&SessionCredential>` signature
|
||||
/// rejects, so it could never have been run): in
|
||||
/// `ModelsManager::sampling_config`, ask the authority about the SESSION's
|
||||
/// endpoint instead of the current model's own —
|
||||
/// `.credential_for(None, &config.endpoints.proxy_url())` in place of
|
||||
/// `.credential_for_model(current_model)`. That is the round-1 defect's shape
|
||||
/// (the credential decided by something other than the model's own platform +
|
||||
/// endpoint) and every `assert_ne!` below sees `Some(KIMI_TOKEN)`.
|
||||
#[tokio::test]
|
||||
#[serial_test::serial]
|
||||
async fn models_manager_sampling_config_never_carries_the_kimi_bearer() {
|
||||
let _env = without_ambient_byok_env();
|
||||
let (_dir, agent) = kimi_session_agent();
|
||||
|
||||
for (catalog_key, slug, base_url) in [
|
||||
(
|
||||
"moonshot-cn/kimi-k2-turbo-preview",
|
||||
"kimi-k2-turbo-preview",
|
||||
"https://api.moonshot.cn/v1",
|
||||
),
|
||||
(
|
||||
"deepseek/deepseek-chat",
|
||||
"deepseek-chat",
|
||||
"https://api.deepseek.com/v1",
|
||||
),
|
||||
("openai/gpt-5.2", "gpt-5.2", "https://api.openai.com/v1"),
|
||||
] {
|
||||
agent
|
||||
.models_manager
|
||||
.insert_test_entry(catalog_key, platform_entry(catalog_key, slug, base_url));
|
||||
agent
|
||||
.models_manager
|
||||
.set_current_model_id(acp::ModelId::new(catalog_key));
|
||||
let cfg = agent.models_manager.sampling_config().config;
|
||||
assert_eq!(cfg.base_url, base_url, "{catalog_key}: routed to its own host");
|
||||
assert_ne!(
|
||||
cfg.api_key.as_deref(),
|
||||
Some(KIMI_TOKEN),
|
||||
"LEAK: ModelsManager::sampling_config sent the Kimi bearer to {base_url}"
|
||||
);
|
||||
}
|
||||
|
||||
// …and the first-party subscription channel is unchanged, which is what
|
||||
// proves the Kimi bearer was reachable above.
|
||||
let kimi_key = "kimi-code/kimi-for-coding";
|
||||
agent.models_manager.insert_test_entry(
|
||||
kimi_key,
|
||||
platform_entry(
|
||||
kimi_key,
|
||||
"kimi-for-coding",
|
||||
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
|
||||
),
|
||||
);
|
||||
agent
|
||||
.models_manager
|
||||
.set_current_model_id(acp::ModelId::new(kimi_key));
|
||||
assert_eq!(
|
||||
agent
|
||||
.models_manager
|
||||
.sampling_config()
|
||||
.config
|
||||
.api_key
|
||||
.as_deref(),
|
||||
Some(KIMI_TOKEN),
|
||||
"the kimi-code subscription channel must be byte-identical"
|
||||
);
|
||||
}
|
||||
|
||||
/// H2 — the SHARED `MvpAgent::sampling_config`. `seed_client_config_auth_if_available`
|
||||
/// (from `new_session` / `load_session`) and the `cached_token` / `kimi.com/oidc`
|
||||
/// login handlers all stamped `sampling_config.api_key = Some(<primary bearer>)`
|
||||
/// with NO platform or endpoint guard — while that very config may point at a
|
||||
/// third-party model, and is both the subagent baseline and the
|
||||
/// unresolved-model fallback. `agent_ops`' generic-OAuth login handler already
|
||||
/// documents the correct rule ("Do NOT stamp this token onto the shared
|
||||
/// sampling_config"); the Kimi handlers violated it.
|
||||
///
|
||||
/// Revert-to-red: make `stamp_session_credential` skip the authority entirely —
|
||||
/// `self.sampling_config.borrow_mut().api_key =
|
||||
/// self.auth_manager.current_or_expired().map(|a| a.key); return true;` — and
|
||||
/// the third-party rows below become `Some(KIMI_TOKEN)`.
|
||||
#[tokio::test]
|
||||
#[serial_test::serial]
|
||||
async fn shared_sampling_config_is_never_stamped_off_the_session_endpoint() {
|
||||
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",
|
||||
),
|
||||
(
|
||||
"moonshot-cn/kimi-k2-turbo-preview",
|
||||
"kimi-k2-turbo-preview",
|
||||
"https://api.moonshot.cn/v1",
|
||||
),
|
||||
] {
|
||||
agent
|
||||
.models_manager
|
||||
.insert_test_entry(catalog_key, platform_entry(catalog_key, slug, base_url));
|
||||
agent
|
||||
.models_manager
|
||||
.set_current_model_id(acp::ModelId::new(catalog_key));
|
||||
rebuild_shared_config(&agent);
|
||||
{
|
||||
let mut shared = agent.sampling_config.borrow_mut();
|
||||
assert_eq!(shared.model, slug, "{catalog_key}: built from this entry");
|
||||
assert_eq!(shared.base_url, base_url);
|
||||
shared.api_key = None;
|
||||
}
|
||||
// The `new_session` / `load_session` seed…
|
||||
agent.seed_client_config_auth_if_available();
|
||||
assert_eq!(
|
||||
agent.sampling_config.borrow().api_key, None,
|
||||
"LEAK: seeding stamped the Kimi bearer onto a config routed at {base_url}"
|
||||
);
|
||||
// …and the login handlers, which overwrite rather than seed.
|
||||
assert!(
|
||||
!agent.stamp_session_credential(true),
|
||||
"LEAK: a login handler stamped the Kimi bearer onto a config routed at {base_url}"
|
||||
);
|
||||
assert_eq!(agent.sampling_config.borrow().api_key, None);
|
||||
}
|
||||
|
||||
// Byte-identical on the session's own endpoint: both paths still stamp.
|
||||
let kimi_key = "kimi-code/kimi-for-coding";
|
||||
agent.models_manager.insert_test_entry(
|
||||
kimi_key,
|
||||
platform_entry(
|
||||
kimi_key,
|
||||
"kimi-for-coding",
|
||||
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
|
||||
),
|
||||
);
|
||||
agent
|
||||
.models_manager
|
||||
.set_current_model_id(acp::ModelId::new(kimi_key));
|
||||
rebuild_shared_config(&agent);
|
||||
agent.sampling_config.borrow_mut().api_key = None;
|
||||
agent.seed_client_config_auth_if_available();
|
||||
assert_eq!(
|
||||
agent.sampling_config.borrow().api_key.as_deref(),
|
||||
Some(KIMI_TOKEN),
|
||||
"the subscription endpoint must still be seeded (this is what makes the \
|
||||
assertions above meaningful)"
|
||||
);
|
||||
}
|
||||
|
||||
/// C1 — THE CRITICAL. The login stamp used to pair the right QUESTION (*may
|
||||
/// **a** session credential ride here?*) with the wrong CREDENTIAL (always
|
||||
/// `auth_manager.current_or_expired().key`, the primary Kimi bearer). For a
|
||||
/// subscription-OAuth platform at its OWN host the answer is correctly "yes" —
|
||||
/// the credential that may ride there is that platform's POOLED token
|
||||
/// ([`CredentialClass::Pooled`]) — so a Claude Pro/Max user whose current model is
|
||||
/// `claude-pro-max/*` ran `kigi login` and the Kimi subscription bearer landed
|
||||
/// on a config routed at `api.anthropic.com`. From there it reaches the wire via
|
||||
/// `resolve_sampling_config_for_model`'s verbatim fallback (offline / stale
|
||||
/// catalog) and via `SubagentSpawnContext`'s baseline clone.
|
||||
///
|
||||
/// All FOUR subscription platforms, each at its own registry host, derived from
|
||||
/// the registry so the fixture cannot drift.
|
||||
///
|
||||
/// The pooled managers are empty here (`pool_home()` is a per-process temp path
|
||||
/// under `cfg(test)`), so the correct answer is `None` — and `None` is also what
|
||||
/// proves the primary is not being substituted, because the same agent DOES
|
||||
/// stamp `KIMI_TOKEN` on its own coding endpoint at the end of the test.
|
||||
///
|
||||
/// Revert-to-red (production, compiles): restore the old shape in
|
||||
/// `MvpAgent::stamp_session_credential` —
|
||||
/// ```ignore
|
||||
/// if self.credential_authority().credential_class(platform, &base_url)
|
||||
/// == CredentialClass::None
|
||||
/// {
|
||||
/// return false;
|
||||
/// }
|
||||
/// let Some(auth) = self.auth_manager.current_or_expired() else { return false };
|
||||
/// self.sampling_config.borrow_mut().api_key = Some(auth.key);
|
||||
/// true
|
||||
/// ```
|
||||
/// and every OAuth row below becomes `Some(KIMI_TOKEN)`.
|
||||
#[tokio::test]
|
||||
#[serial_test::serial]
|
||||
async fn oauth_platform_shared_config_never_receives_the_primary_on_login() {
|
||||
let _env = without_ambient_byok_env();
|
||||
let (_dir, agent) = kimi_session_agent();
|
||||
|
||||
for (platform_id, slug) in [
|
||||
("claude-pro-max", "claude-opus-4-8"),
|
||||
("openai-codex", "gpt-5.5-codex"),
|
||||
("github-copilot", "gpt-4.1"),
|
||||
("xai-grok", "grok-4.5"),
|
||||
] {
|
||||
let platform = kigi_models::PlatformId::parse(platform_id).expect("known platform");
|
||||
let base_url = platform.base_url();
|
||||
let catalog_key = format!("{platform_id}/{slug}");
|
||||
agent
|
||||
.models_manager
|
||||
.insert_test_entry(&catalog_key, platform_entry(&catalog_key, slug, &base_url));
|
||||
agent
|
||||
.models_manager
|
||||
.set_current_model_id(acp::ModelId::new(catalog_key.clone()));
|
||||
rebuild_shared_config(&agent);
|
||||
{
|
||||
let mut shared = agent.sampling_config.borrow_mut();
|
||||
assert_eq!(shared.model, slug, "{catalog_key}: built from this entry");
|
||||
assert_eq!(shared.base_url, base_url);
|
||||
shared.api_key = None;
|
||||
}
|
||||
|
||||
// Precondition: this endpoint DOES take a session credential — that is
|
||||
// exactly why guarding a hand-carried primary with "may a session
|
||||
// credential ride?" was the bug. The class names WHICH one: the
|
||||
// platform's POOLED token, never the primary / house key.
|
||||
assert_eq!(
|
||||
agent
|
||||
.credential_authority()
|
||||
.credential_class(Some(platform), &base_url),
|
||||
CredentialClass::Pooled,
|
||||
"{catalog_key}: precondition — its own host takes its POOLED token, \
|
||||
and never the primary / house credential"
|
||||
);
|
||||
|
||||
// `kigi login` (cached_token and kimi.com/oidc both land here) …
|
||||
assert!(
|
||||
!agent.stamp_session_credential(true),
|
||||
"LEAK: a Kimi login stamped a credential onto a config routed at {base_url}"
|
||||
);
|
||||
assert_eq!(
|
||||
agent.sampling_config.borrow().api_key,
|
||||
None,
|
||||
"LEAK: {catalog_key} received a bearer that is not its own pooled token"
|
||||
);
|
||||
// … and the `new_session` / `load_session` seed.
|
||||
agent.seed_client_config_auth_if_available();
|
||||
assert_ne!(
|
||||
agent.sampling_config.borrow().api_key.as_deref(),
|
||||
Some(KIMI_TOKEN),
|
||||
"LEAK: seeding sent the Kimi subscription bearer to {base_url}"
|
||||
);
|
||||
}
|
||||
|
||||
// The first-party channel is untouched — this is what makes every
|
||||
// assertion above meaningful (the Kimi bearer IS live and IS stampable).
|
||||
let kimi_key = "kimi-code/kimi-for-coding";
|
||||
agent.models_manager.insert_test_entry(
|
||||
kimi_key,
|
||||
platform_entry(
|
||||
kimi_key,
|
||||
"kimi-for-coding",
|
||||
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
|
||||
),
|
||||
);
|
||||
agent
|
||||
.models_manager
|
||||
.set_current_model_id(acp::ModelId::new(kimi_key));
|
||||
rebuild_shared_config(&agent);
|
||||
agent.sampling_config.borrow_mut().api_key = None;
|
||||
assert!(agent.stamp_session_credential(true));
|
||||
assert_eq!(
|
||||
agent.sampling_config.borrow().api_key.as_deref(),
|
||||
Some(KIMI_TOKEN),
|
||||
"the Kimi subscription channel must be byte-identical"
|
||||
);
|
||||
}
|
||||
|
||||
/// H-a (AVAILABILITY, HIGH) — the stamp guard must resolve the model the shared
|
||||
/// config was BUILT from, not the one `current_model_id()` names NOW.
|
||||
///
|
||||
/// The shared config is built ONCE (`MvpAgent::with_models`) and never rebuilt,
|
||||
/// while `current_model_id()` moves on every non-Leader model switch
|
||||
/// (`handlers/model_switch.rs`) and on catalog reselection. Once they drift, a
|
||||
/// guard re-resolving the config's BARE slug against the live cell fails
|
||||
/// `entry_for_slug_resolution`'s `entry.info.model == slug` test and falls
|
||||
/// through to `resolve_catalog_key`'s `.rev()` scan.
|
||||
///
|
||||
/// `kimi-code` (subscription, `uses_oauth`) and `kimi-coding` (API-key twin,
|
||||
/// SAME coding host, `uses_oauth: false`) list the same routing slug, and
|
||||
/// `PlatformId::ALL` puts `kimi-code` 1st and `kimi-coding` 19th — so that scan
|
||||
/// answers `kimi-coding`, which takes NO session credential and therefore has NO
|
||||
/// governing manager. Because `stamp_session_credential` only overwrites ON
|
||||
/// SUCCESS, a Kimi-subscription user who also has `KIMI_API_KEY` set kept the
|
||||
/// EXPIRED bearer in the shared config after a successful re-login: every
|
||||
/// unresolved-model fallback and every subagent baseline turn 401'd until
|
||||
/// restart.
|
||||
///
|
||||
/// The switch happens AFTER the config is built — the previous version of this
|
||||
/// test left both on the same model, so it could not catch this.
|
||||
///
|
||||
/// Revert-to-red (production, compiles): make `MvpAgent::shared_config_platform`
|
||||
/// re-resolve from the live cell instead of returning the captured value —
|
||||
/// ```ignore
|
||||
/// fn shared_config_platform(&self) -> Option<kigi_models::PlatformId> {
|
||||
/// let model = self.sampling_config.borrow().model.clone();
|
||||
/// let current_key = self.models_manager.current_model_id();
|
||||
/// crate::agent::models::platform_for_slug(
|
||||
/// &self.models_manager.models(),
|
||||
/// Some(current_key.0.as_ref()),
|
||||
/// &model,
|
||||
/// )
|
||||
/// }
|
||||
/// ```
|
||||
/// and the re-login assertion below sees the stale bearer.
|
||||
#[tokio::test]
|
||||
#[serial_test::serial]
|
||||
async fn relogin_restamps_the_shared_config_after_the_model_switched() {
|
||||
let _env = without_ambient_byok_env();
|
||||
let (_dir, agent) = kimi_session_agent();
|
||||
|
||||
let coding_host = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url;
|
||||
let slug = "kimi-for-coding";
|
||||
// Insertion order mirrors `PlatformId::ALL`: the subscription platform
|
||||
// first, its API-key twin later — so the `.rev()` scan answers the twin.
|
||||
for catalog_key in ["kimi-code/kimi-for-coding", "kimi-coding/kimi-for-coding"] {
|
||||
agent
|
||||
.models_manager
|
||||
.insert_test_entry(catalog_key, platform_entry(catalog_key, slug, coding_host));
|
||||
}
|
||||
assert_eq!(
|
||||
crate::agent::models::platform_for_slug(&agent.models_manager.models(), None, slug),
|
||||
Some(kigi_models::PlatformId::KimiCoding),
|
||||
"precondition: a `None`-keyed slug scan answers the API-key twin, which takes \
|
||||
no session credential"
|
||||
);
|
||||
|
||||
// Startup: the picker selected the SUBSCRIPTION entry, and the shared config
|
||||
// was built from it (config + platform, one call).
|
||||
agent
|
||||
.models_manager
|
||||
.set_current_model_id(acp::ModelId::new("kimi-code/kimi-for-coding"));
|
||||
rebuild_shared_config(&agent);
|
||||
assert_eq!(
|
||||
agent.sampling_config.borrow().model,
|
||||
slug,
|
||||
"precondition: the shared config carries the BARE slug, which collides"
|
||||
);
|
||||
|
||||
// The user switches model: `current_model_id` moves to the API-key twin
|
||||
// while the shared config still represents the subscription entry.
|
||||
agent
|
||||
.models_manager
|
||||
.set_current_model_id(acp::ModelId::new("kimi-coding/kimi-for-coding"));
|
||||
|
||||
// The session expires and `kigi login` re-mints it. The handlers overwrite
|
||||
// (`stamp_session_credential(true)`) AFTER the manager holds the new token.
|
||||
agent.sampling_config.borrow_mut().api_key = Some("expired-bearer".to_string());
|
||||
assert!(
|
||||
agent.stamp_session_credential(true),
|
||||
"a successful re-login must restamp the shared config"
|
||||
);
|
||||
assert_eq!(
|
||||
agent.sampling_config.borrow().api_key.as_deref(),
|
||||
Some(KIMI_TOKEN),
|
||||
"the re-login must REPLACE the expired bearer: resolving the API-key twin \
|
||||
instead finds no governing manager, leaves the stale token in place, and \
|
||||
401s every subagent baseline turn until restart"
|
||||
);
|
||||
|
||||
// And the seed path (`new_session` / `load_session`) agrees.
|
||||
agent.sampling_config.borrow_mut().api_key = None;
|
||||
agent.seed_client_config_auth_if_available();
|
||||
assert_eq!(
|
||||
agent.sampling_config.borrow().api_key.as_deref(),
|
||||
Some(KIMI_TOKEN),
|
||||
);
|
||||
}
|
||||
|
||||
/// H3 (REGRESSION) — a MANAGED deployment configures its coding endpoint with
|
||||
/// `[endpoints] coding_api_base_url` in **config.toml** (what the managed-config
|
||||
/// sync writes), NOT the `KIGI_CODE_BASE_URL` env var. The previous round's
|
||||
/// predicate knew only the env var, so `EndpointsConfig::proxy_url()`'s
|
||||
/// config-key branch was invisible: every model inheriting the managed endpoint
|
||||
/// classified third-party, lost its api_key AND its resolver, and 401'd on every
|
||||
/// turn.
|
||||
///
|
||||
/// NOTE: no env var is set anywhere in this test — that is the point.
|
||||
///
|
||||
/// Revert-to-red: drop the `proxy_url()` arm from
|
||||
/// `CredentialAuthority::is_session_coding_endpoint` and both `assert_eq!`s
|
||||
/// below become `None`.
|
||||
#[tokio::test]
|
||||
#[serial_test::serial]
|
||||
async fn managed_config_toml_coding_endpoint_keeps_the_session_bearer() {
|
||||
let _env = without_ambient_byok_env();
|
||||
let _no_env_override = EnvGuard::unset("KIGI_CODE_BASE_URL");
|
||||
let managed = "https://proxy.acme.example/v1";
|
||||
|
||||
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 cfg = AgentConfig {
|
||||
endpoints: EndpointsConfig {
|
||||
coding_api_base_url: Some(managed.to_string()),
|
||||
..EndpointsConfig::default()
|
||||
},
|
||||
..AgentConfig::default()
|
||||
};
|
||||
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
let agent = MvpAgent::new(GatewaySender::new(tx), &cfg, auth_manager, None)
|
||||
.expect("valid test config");
|
||||
agent.set_auth_method(acp::AuthMethodId::new(CACHED_TOKEN_AUTH_METHOD_ID));
|
||||
assert_eq!(
|
||||
agent.models_manager.endpoints().proxy_url(),
|
||||
managed,
|
||||
"precondition: the session's effective coding endpoint is the config.toml key"
|
||||
);
|
||||
|
||||
// A `[model.*]` entry inheriting the managed endpoint …
|
||||
let mut bare = ModelEntry::fallback("kigi-4.5", &cfg.endpoints);
|
||||
bare.info.id = None;
|
||||
bare.info.base_url = managed.to_string();
|
||||
assert_eq!(
|
||||
agent
|
||||
.prepare_sampling_config_for_model(&bare, None)
|
||||
.api_key
|
||||
.as_deref(),
|
||||
Some(KIMI_TOKEN),
|
||||
"a managed deployment must still receive the session bearer"
|
||||
);
|
||||
|
||||
// … and the kimi-code catalog entry, whose base_url IS `proxy_url()`.
|
||||
let kimi = platform_entry("kimi-code/kimi-for-coding", "kimi-for-coding", managed);
|
||||
assert_eq!(
|
||||
agent
|
||||
.prepare_sampling_config_for_model(&kimi, None)
|
||||
.api_key
|
||||
.as_deref(),
|
||||
Some(KIMI_TOKEN),
|
||||
"the managed kimi-code entry must still receive the session bearer"
|
||||
);
|
||||
|
||||
// The spec's requirement was "rides AND still refreshes": an api_key alone
|
||||
// freezes at login and 401s unrecoverably ~1h in. The managed endpoint must
|
||||
// also keep a LIVE manager — the primary's, so mid-session refresh and 401
|
||||
// recovery run against the credential that actually owns that host.
|
||||
for platform in [None, Some(kigi_models::PlatformId::KimiCode)] {
|
||||
let manager = agent
|
||||
.credential_authority()
|
||||
.manager_for(platform, managed)
|
||||
.unwrap_or_else(|| {
|
||||
panic!("{platform:?}: a managed deployment must keep a live manager")
|
||||
});
|
||||
assert!(
|
||||
std::sync::Arc::ptr_eq(&manager, &agent.auth_manager),
|
||||
"{platform:?}: and it must be the session's OWN primary manager"
|
||||
);
|
||||
let resolver = agent
|
||||
.credential_authority()
|
||||
.bearer_resolver_for(platform, managed)
|
||||
.unwrap_or_else(|| panic!("{platform:?}: … exposed as a live bearer_resolver"));
|
||||
assert_eq!(
|
||||
resolver.current_bearer(),
|
||||
Some(KIMI_TOKEN.to_string()),
|
||||
"{platform:?}: the resolver reads the live primary session bearer"
|
||||
);
|
||||
}
|
||||
|
||||
// The guard still holds: a third-party host under the SAME managed config
|
||||
// gets nothing — no key, and no resolver either.
|
||||
let third_party = platform_entry(
|
||||
"deepseek/deepseek-chat",
|
||||
"deepseek-chat",
|
||||
"https://api.deepseek.com/v1",
|
||||
);
|
||||
assert_eq!(
|
||||
agent
|
||||
.prepare_sampling_config_for_model(&third_party, None)
|
||||
.api_key,
|
||||
None,
|
||||
"LEAK: a managed deployment must not widen the trust set to third parties"
|
||||
);
|
||||
assert!(
|
||||
agent
|
||||
.credential_authority()
|
||||
.bearer_resolver_for(
|
||||
Some(kigi_models::PlatformId::DeepSeek),
|
||||
"https://api.deepseek.com/v1"
|
||||
)
|
||||
.is_none(),
|
||||
"LEAK: nor hand it a live resolver over the primary"
|
||||
);
|
||||
}
|
||||
|
||||
/// M3 (COMPLETION) — the SUMMARY client's `bearer_resolver` must honour the
|
||||
/// session-token gate, exactly as the session actor's aux path does.
|
||||
///
|
||||
/// `build_summary_client` set `cfg.bearer_resolver =
|
||||
/// authority.bearer_resolver_for(platform_for_slug(…), &cfg.base_url)` with NO
|
||||
/// gate while `SessionActor::aux_bearer_resolver` had one. `SamplingClient::post`
|
||||
/// REPLACES the request's auth header from that resolver, so an api-key /
|
||||
/// house-key session whose `[model.session-summary]` block carries its OWN
|
||||
/// `env_key` on the session's own coding endpoint had that key overwritten by
|
||||
/// the primary bearer on every summary request. Both now go through ONE rule,
|
||||
/// `sampler_turn::aux_bearer_resolver_for`.
|
||||
///
|
||||
/// The summary slug is deliberately absent from the catalog and from any
|
||||
/// config: it classifies `NotByok` definitively, so the only variable left is
|
||||
/// the ACP auth method — which is the gate term under test.
|
||||
///
|
||||
/// Revert-to-red (production, compiles): in `MvpAgent::summary_bearer_resolver`,
|
||||
/// return `self.credential_authority().bearer_resolver_for(platform, base_url)`
|
||||
/// directly (the pre-fix shape) and the api-key row below resolves `KIMI_TOKEN`.
|
||||
#[tokio::test]
|
||||
#[serial_test::serial]
|
||||
async fn summary_client_resolver_honours_the_session_gate() {
|
||||
let _env = without_ambient_byok_env();
|
||||
let (_dir, agent) = kimi_session_agent();
|
||||
let coding_host = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url;
|
||||
let slug = "kigi-summary-aux-not-in-any-catalog";
|
||||
let models = agent.models_manager.models();
|
||||
|
||||
// A session-based method on the session's OWN endpoint: byte-identical, the
|
||||
// summary model keeps a LIVE resolver over the primary.
|
||||
agent.set_auth_method(acp::AuthMethodId::new(CACHED_TOKEN_AUTH_METHOD_ID));
|
||||
let resolver = agent
|
||||
.summary_bearer_resolver(&models, slug, coding_host)
|
||||
.expect("the first-party subscription summary channel keeps its resolver");
|
||||
assert_eq!(
|
||||
resolver.current_bearer(),
|
||||
Some(KIMI_TOKEN.to_string()),
|
||||
"…and it reads the live primary session bearer"
|
||||
);
|
||||
|
||||
// An API-KEY session: the gate is inactive, so the summary model's own key
|
||||
// must survive to the wire instead of being replaced by the primary.
|
||||
agent.set_auth_method(acp::AuthMethodId::new(
|
||||
crate::agent::auth_method::XAI_API_KEY_METHOD_ID,
|
||||
));
|
||||
assert!(
|
||||
agent
|
||||
.summary_bearer_resolver(&models, slug, coding_host)
|
||||
.is_none(),
|
||||
"LEAK: an api-key session's summary model had its own key replaced on the \
|
||||
wire by the primary bearer"
|
||||
);
|
||||
}
|
||||
@@ -176,6 +176,7 @@ async fn handle_connection(ws: WebSocket, state: Arc<ServerState>, peer_addr: So
|
||||
prefetch_models_blocking(
|
||||
&agent_config.endpoints,
|
||||
auth.as_ref(),
|
||||
&Default::default(),
|
||||
fetch_auth,
|
||||
&platform_keys,
|
||||
)
|
||||
|
||||
@@ -59,6 +59,8 @@ fn effort_label(effort: ReasoningEffort) -> String {
|
||||
ReasoningEffort::Medium => "Medium",
|
||||
ReasoningEffort::High => "High",
|
||||
ReasoningEffort::Xhigh => "X-High",
|
||||
ReasoningEffort::Max => "Max",
|
||||
ReasoningEffort::Ultra => "Ultra",
|
||||
}
|
||||
.to_string()
|
||||
}
|
||||
|
||||
@@ -1036,6 +1036,7 @@ pub(crate) async fn handle_subagent_request(
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
if verbatim_mirror_fork {
|
||||
None
|
||||
} else if let Some(scope) = agent_memory_scope {
|
||||
@@ -1069,6 +1070,8 @@ pub(crate) async fn handle_subagent_request(
|
||||
ctx.app_builder_deployer_config.clone(),
|
||||
ctx.write_file_enabled,
|
||||
ctx.goal_enabled,
|
||||
// Graph mode is a depth-0 harness; child sessions never drive it.
|
||||
false,
|
||||
true,
|
||||
ctx.ask_user_question_enabled,
|
||||
ctx.client_hooks.clone(),
|
||||
|
||||
@@ -880,6 +880,24 @@ async fn read_parent_sampling_config(
|
||||
let auth_scheme = crate::agent::config::try_resolve_model_credentials(&cfg.model, None)
|
||||
.map(|r| r.auth_scheme)
|
||||
.unwrap_or_default();
|
||||
// Claude Pro/Max OAuth Messages adaptation inherits from the parent
|
||||
// model's platform (claude-pro-max → true); every other platform,
|
||||
// and BYOK, → false, so the API-key paths stay byte-identical.
|
||||
let anthropic_oauth = kigi_models::parse_managed_model_key(ctx.model_id.0.as_ref())
|
||||
.is_some_and(|(platform, _)| {
|
||||
platform.oauth().is_some()
|
||||
&& platform.wire_api() == kigi_models::PlatformWireApi::Messages
|
||||
});
|
||||
// GitHub Copilot editor headers inherit from the parent model's
|
||||
// platform (github-copilot → true); every other platform / BYOK →
|
||||
// false, so the other ChatCompletions paths stay byte-identical.
|
||||
let github_copilot = kigi_models::parse_managed_model_key(ctx.model_id.0.as_ref())
|
||||
.is_some_and(|(platform, _)| platform.sends_copilot_editor_headers());
|
||||
// ChatGPT/Codex headers inherit from the parent model's platform
|
||||
// (openai-codex → true); every other platform / BYOK → false, so the
|
||||
// API-key openai Responses path stays byte-identical.
|
||||
let openai_codex = kigi_models::parse_managed_model_key(ctx.model_id.0.as_ref())
|
||||
.is_some_and(|(platform, _)| platform.sends_codex_responses_headers());
|
||||
let inherited = kigi_sampler::SamplerConfig {
|
||||
api_key: creds.api_key,
|
||||
base_url: cfg.base_url,
|
||||
@@ -889,6 +907,10 @@ async fn read_parent_sampling_config(
|
||||
top_p: cfg.top_p,
|
||||
api_backend: cfg.api_backend,
|
||||
auth_scheme,
|
||||
anthropic_oauth,
|
||||
github_copilot,
|
||||
openai_codex,
|
||||
chat_compat: cfg.chat_compat,
|
||||
extra_headers,
|
||||
context_window: cfg.context_window.get(),
|
||||
reasoning_effort: cfg.reasoning_effort,
|
||||
@@ -981,9 +1003,21 @@ fn resolve_model_override_to_config(
|
||||
} else {
|
||||
acp::ModelId::new(entry.info().model.clone())
|
||||
};
|
||||
let session_key = ctx.auth.as_ref().map(|a| a.key.as_str());
|
||||
// Resolve the child's session token by the OVERRIDE model's OWN platform
|
||||
// AND endpoint, not the parent's primary auth: a grok (oauth-platform)
|
||||
// override draws its pooled grok token or `None`, and an API-key registry
|
||||
// platform / a third-party `[model.*]` host draws NOTHING — NEVER the
|
||||
// primary Kimi session token, which `resolve_credentials` would otherwise
|
||||
// stamp onto the child's api.x.ai / api.moonshot.cn credentials. The
|
||||
// first-party subscription channel still resolves to the primary
|
||||
// (byte-identical).
|
||||
let session_key = crate::auth::credential_authority::CredentialAuthority::new(
|
||||
ctx.models_manager.endpoints(),
|
||||
Some(ctx.auth_manager.clone()),
|
||||
)
|
||||
.credential_for_model(&entry);
|
||||
let has_session_key = session_key.is_some();
|
||||
let mut credentials = resolve_credentials(&entry, session_key);
|
||||
let mut credentials = resolve_credentials(&entry, session_key.as_ref());
|
||||
credentials.auth_type = subagent_auth_type(Some(&entry), &ctx.auth_method_id);
|
||||
let resolved_auth_type = credentials.auth_type;
|
||||
let config = sampling_config_for_model(&entry, credentials, ctx.alpha_test_key.clone());
|
||||
|
||||
@@ -3131,19 +3131,25 @@ fn fresh_tool_model_rejects_unavailable_exact_key_over_visible_slug_collision()
|
||||
"validation must inspect the unavailable exact-key entry selected by execution"
|
||||
);
|
||||
}
|
||||
/// Validation must inspect the SAME slug-collision entry execution selects.
|
||||
/// Both go through `find_model_by_id`, whose slug scan takes the LAST match —
|
||||
/// aligned with the picker's `resolve_catalog_key` so the auth layer and the
|
||||
/// picker can never resolve different platforms for one slug (the H5 collision).
|
||||
/// So a blocked LAST entry must be rejected even though an available earlier one
|
||||
/// shares the slug.
|
||||
#[test]
|
||||
fn fresh_tool_model_rejects_unavailable_first_slug_collision() {
|
||||
fn fresh_tool_model_rejects_unavailable_last_slug_collision() {
|
||||
let mut models = indexmap::IndexMap::new();
|
||||
let mut unavailable_first = test_model_entry("shared-routing-slug");
|
||||
unavailable_first.info.user_selectable = false;
|
||||
models.insert("blocked-first".to_string(), unavailable_first);
|
||||
models.insert("visible-second".to_string(), test_model_entry("shared-routing-slug"));
|
||||
models.insert("visible-first".to_string(), test_model_entry("shared-routing-slug"));
|
||||
let mut unavailable_last = test_model_entry("shared-routing-slug");
|
||||
unavailable_last.info.user_selectable = false;
|
||||
models.insert("blocked-last".to_string(), unavailable_last);
|
||||
assert_eq!(
|
||||
super::handle_request::task_model_override_error(Some("shared-routing-slug"),
|
||||
ModelOverrideProvenance::Tool, false, & models, false,).as_deref(),
|
||||
Some("Unknown Task.model slug 'shared-routing-slug'. Valid model slugs: \
|
||||
visible-second. Omit `model` to inherit the parent model."),
|
||||
"validation must inspect the first routing-slug entry selected by execution"
|
||||
visible-first. Omit `model` to inherit the parent model."),
|
||||
"validation must inspect the last routing-slug entry selected by execution"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
@@ -3272,6 +3278,7 @@ fn test_sampling_config(model_slug: &str) -> kigi_sampling_types::SamplingConfig
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: NonZeroU64::new(256_000).expect("non-zero context window"),
|
||||
reasoning_effort: None,
|
||||
|
||||
@@ -2467,6 +2467,111 @@ async fn resolve_subagent_config_override_unknown_model_falls_through_to_inherit
|
||||
assert_eq!(config.model, "kigi-4.5");
|
||||
assert_eq!(model_id.0.as_ref(), "kigi-4.5");
|
||||
}
|
||||
/// Build an `Arc<AuthManager>` (primary Kimi) holding `key` as its live bearer.
|
||||
/// The `TempDir` is returned so the caller keeps it alive.
|
||||
fn kimi_primary_with_token(key: &str) -> (tempfile::TempDir, std::sync::Arc<crate::auth::AuthManager>) {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let manager = std::sync::Arc::new(crate::auth::AuthManager::new(
|
||||
dir.path(),
|
||||
crate::auth::KimiCodeConfig::default(),
|
||||
));
|
||||
manager.hot_swap(crate::auth::KimiAuth {
|
||||
key: key.to_string(),
|
||||
auth_mode: crate::auth::AuthMode::OAuth,
|
||||
..crate::auth::KimiAuth::test_default()
|
||||
});
|
||||
(dir, manager)
|
||||
}
|
||||
/// LEAK 2 (subagent model-override): a grok (oauth-platform) override with a
|
||||
/// Kimi primary must NEVER receive the primary Kimi session token as its
|
||||
/// `api_key` — it draws grok's own pooled token (or `None`). Revert-to-red: the
|
||||
/// pre-fix code passed `ctx.auth` (Kimi) straight to `resolve_credentials`, so
|
||||
/// `config.api_key == "kimi-secret"` and this assertion fails.
|
||||
#[tokio::test]
|
||||
async fn subagent_override_grok_model_never_leaks_kimi_session_token() {
|
||||
let (_kd, manager) = kimi_primary_with_token("kimi-secret");
|
||||
let mut grok = test_model_entry("grok-4-latest");
|
||||
grok.info.id = Some("xai-grok/grok-4-latest".to_string());
|
||||
grok.info.base_url = "https://api.x.ai/v1".to_string();
|
||||
let mut models = indexmap::IndexMap::new();
|
||||
models.insert("grok".to_string(), grok);
|
||||
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("grok", &ctx).expect("grok override resolves to a config");
|
||||
assert_ne!(
|
||||
config.api_key.as_deref(),
|
||||
Some("kimi-secret"),
|
||||
"a grok override must never receive the primary Kimi session token",
|
||||
);
|
||||
}
|
||||
/// Byte-identical guard: an override on the SESSION's own first-party endpoint
|
||||
/// (the kimi-code subscription channel) still resolves to the primary session
|
||||
/// token — the primary path is unchanged.
|
||||
#[tokio::test]
|
||||
async fn subagent_override_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: make `CredentialAuthority::governing_manager`'s
|
||||
/// `Some(platform) => None` arm return `self.primary.clone()` and `api_key`
|
||||
/// becomes `Some("kimi-secret")`.
|
||||
#[tokio::test]
|
||||
async fn subagent_override_api_key_platform_never_gets_the_primary_token() {
|
||||
let (_kd, manager) = kimi_primary_with_token("kimi-secret");
|
||||
let mut entry = test_model_entry("kimi-k2-0905-preview");
|
||||
entry.info.id = Some("moonshot-cn/kimi-k2".to_string());
|
||||
entry.info.base_url = "https://api.moonshot.cn/v1".to_string();
|
||||
let mut models = indexmap::IndexMap::new();
|
||||
models.insert("k2".to_string(), entry);
|
||||
let mut ctx = ctx_with_toggle(HashMap::new());
|
||||
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("k2", &ctx)
|
||||
.expect("an API-key-platform override still resolves to a config");
|
||||
assert_ne!(
|
||||
config.api_key.as_deref(),
|
||||
Some("kimi-secret"),
|
||||
"LEAK: an API-key-platform override must never receive the primary Kimi session token",
|
||||
);
|
||||
}
|
||||
/// An unresolvable `AgentDefinition.model` pin (model absent from
|
||||
/// `available_models`) falls through to inherit the parent model.
|
||||
#[tokio::test]
|
||||
|
||||
@@ -0,0 +1,568 @@
|
||||
//! THE credential chokepoint.
|
||||
//!
|
||||
//! One authority answers, for every outgoing inference request, the only
|
||||
//! question that matters: **which credential — if any — may ride it?**
|
||||
//! ([`CredentialAuthority::credential_class`]). Before this module the answer was
|
||||
//! re-derived — differently — at every call site (`ModelsManager`, `MvpAgent`,
|
||||
//! `SessionActor`, the aux/summary/subagent paths), and three separate rounds
|
||||
//! of fixes each closed some sites and missed others.
|
||||
//!
|
||||
//! # How omission is structurally prevented
|
||||
//!
|
||||
//! 1. [`SessionCredential`] wraps the bearer and has **no production
|
||||
//! constructor outside this module**. The only function in the crate that
|
||||
//! can build one is [`CredentialAuthority::credential_for`], which *requires*
|
||||
//! `(platform, base_url)` and holds the session's `EndpointsConfig` and
|
||||
//! primary [`AuthManager`] privately.
|
||||
//! 2. Every API that stamps a session credential onto a request —
|
||||
//! `resolve_credentials`, `resolve_aux_model_sampling_config`,
|
||||
//! `try_resolve_model_credentials`,
|
||||
//! `resolve_chat_state_auth_type` — takes `Option<&SessionCredential>`,
|
||||
//! never `Option<&str>`. A new call site therefore *cannot compile* a leak:
|
||||
//! there is no way to produce the value without going through the rule.
|
||||
//! 3. The authority owns the primary manager privately and exposes it only via
|
||||
//! [`CredentialAuthority::manager_for`] /
|
||||
//! [`CredentialAuthority::bearer_resolver_for`], which take the same
|
||||
//! `(platform, base_url)` pair — so the `bearer_resolver` sink is funnelled
|
||||
//! through the identical rule as the `api_key` sink.
|
||||
//! 4. A guard asks [`CredentialAuthority::credential_class`] and MATCHES on the
|
||||
//! answer. There is no second, similarly-named boolean to pick by mistake:
|
||||
//! the round-3 defect (C1) was `takes_session_credential` — *may **a**
|
||||
//! session credential ride?* — paired with a hand-carried PRIMARY bearer,
|
||||
//! and the two predicates that made that pairing expressible are gone.
|
||||
//!
|
||||
//! SECURITY: no token is ever logged, `Debug`-printed or `Display`ed here.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::agent::config::{EndpointsConfig, ModelEntry};
|
||||
use crate::auth::AuthManager;
|
||||
|
||||
/// A session bearer this authority has cleared for one specific request
|
||||
/// endpoint.
|
||||
///
|
||||
/// Opaque by construction: the inner `String` is private, the type is not
|
||||
/// `Debug`/`Clone`-into-`String`, and the only production constructor is
|
||||
/// [`CredentialAuthority::credential_for`]. See the module docs for why that
|
||||
/// matters.
|
||||
pub(crate) struct SessionCredential(String);
|
||||
|
||||
impl SessionCredential {
|
||||
/// The raw bearer. SECURITY: callers stamp this straight onto a request —
|
||||
/// never log it.
|
||||
pub(crate) fn expose(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
|
||||
/// Test-only forgery, so unit tests can exercise the *downstream*
|
||||
/// credential plumbing (`resolve_credentials`' BYOK-vs-session precedence,
|
||||
/// aux config shapes) without standing up an `AuthManager`. Deliberately
|
||||
/// `#[cfg(test)]`: production code has no way to build one.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn for_test(key: &str) -> Self {
|
||||
Self(key.to_owned())
|
||||
}
|
||||
}
|
||||
|
||||
/// WHICH credential — if any — may ride a request routed to a given
|
||||
/// `(platform, base_url)` pair.
|
||||
///
|
||||
/// ONE question with three answers, replacing the two look-alike booleans
|
||||
/// `takes_session_credential` / `takes_primary_credential` (identical
|
||||
/// signatures, near-identical names, opposite answers on a subscription host).
|
||||
/// C1 was caused by asking the first and stamping the credential the second
|
||||
/// describes; with a single classifier a call site must MATCH on the answer, so
|
||||
/// that mistake is no longer expressible.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(crate) enum CredentialClass {
|
||||
/// `platform`'s OWN pooled subscription-OAuth token, at its own registry
|
||||
/// host. NEVER the primary session bearer and never the house key.
|
||||
Pooled,
|
||||
/// The credential that authorizes the SESSION's own coding endpoint:
|
||||
/// the primary (`kimi-code` / platform-less) bearer.
|
||||
///
|
||||
/// Deliberately NOT split into a separate `HouseKey` variant: the house
|
||||
/// `KIGI_API_KEY` is accepted by exactly this endpoint and no other, so it
|
||||
/// rides precisely this class. A fourth variant would re-create the
|
||||
/// two-similar-answers hazard this enum exists to remove.
|
||||
Primary,
|
||||
/// Nothing rides: every API-key registry platform, an OAuth platform
|
||||
/// redirected off its own host, and any endpoint that is not the session's.
|
||||
None,
|
||||
}
|
||||
|
||||
impl CredentialClass {
|
||||
/// Stable label for structured logs. SECURITY: names a channel, never a
|
||||
/// token.
|
||||
pub(crate) fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Pooled => "pooled",
|
||||
Self::Primary => "primary",
|
||||
Self::None => "none",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The single authority over inference-time session credentials.
|
||||
///
|
||||
/// Construct one from the session's EFFECTIVE endpoints plus its primary
|
||||
/// (first-party / Kimi) manager, then ask it about a request. Cheap to build
|
||||
/// (a handful of `Option<String>` clones + an `Arc` clone).
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct CredentialAuthority {
|
||||
/// The session's effective `[endpoints]` — config.toml layered over env.
|
||||
/// H3: `EndpointsConfig::proxy_url()` prefers `[endpoints]
|
||||
/// coding_api_base_url` from **config.toml** and only then falls back to
|
||||
/// `KIGI_CODE_BASE_URL`. A predicate that knows only the env var makes a
|
||||
/// managed/enterprise deployment lose its session bearer entirely (401 on
|
||||
/// every turn), so the endpoints are part of the authority's identity, not
|
||||
/// an afterthought.
|
||||
endpoints: EndpointsConfig,
|
||||
/// The primary session manager. PRIVATE: nothing hands it back, so a path
|
||||
/// holding a `CredentialAuthority` cannot reach `current_or_expired()`
|
||||
/// without naming an endpoint.
|
||||
primary: Option<Arc<AuthManager>>,
|
||||
}
|
||||
|
||||
impl CredentialAuthority {
|
||||
pub(crate) fn new(endpoints: EndpointsConfig, primary: Option<Arc<AuthManager>>) -> Self {
|
||||
Self { endpoints, primary }
|
||||
}
|
||||
|
||||
/// THE rule, stated once.
|
||||
///
|
||||
/// - a subscription-OAuth platform (claude-pro-max, openai-codex,
|
||||
/// github-copilot, xai-grok) rides ITS OWN pooled manager — never the
|
||||
/// primary — and only to its own registry host (L10: a
|
||||
/// `[model."claude-pro-max/x"]` override keeps `info.id` but can point
|
||||
/// `base_url` anywhere, and used to ship the Claude OAuth bearer there);
|
||||
/// - `kimi-code` — the one `uses_oauth` platform with no `OAuthConfig` —
|
||||
/// rides the PRIMARY session, and only at the session's own effective
|
||||
/// coding endpoint;
|
||||
/// - every API-key registry platform (deepseek, openai, anthropic,
|
||||
/// moonshot-*, …) rides NOTHING: its credential is that platform's API
|
||||
/// key, already resolved into the catalog entry;
|
||||
/// - a platform-less model (a bare slug or a `[model.*]` block) is decided
|
||||
/// purely by the ENDPOINT — BYOK detection probes `std::env::var` at call
|
||||
/// time, so an unset/mistyped `env_key` must not turn into "send the
|
||||
/// subscription bearer to `api.openai.com`".
|
||||
pub(crate) fn credential_class(
|
||||
&self,
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
base_url: &str,
|
||||
) -> CredentialClass {
|
||||
match platform {
|
||||
Some(platform) => match platform.oauth() {
|
||||
Some(_) if self.endpoint_is_platform_host(platform, base_url) => {
|
||||
CredentialClass::Pooled
|
||||
}
|
||||
// An OAuth platform pointed at a host that is NOT its own.
|
||||
Some(_) => CredentialClass::None,
|
||||
None if platform.uses_oauth() && self.is_session_coding_endpoint(base_url) => {
|
||||
CredentialClass::Primary
|
||||
}
|
||||
// `kimi-code` off the session's endpoint, and every API-key
|
||||
// registry platform.
|
||||
None => CredentialClass::None,
|
||||
},
|
||||
None if self.is_session_coding_endpoint(base_url) => CredentialClass::Primary,
|
||||
None => CredentialClass::None,
|
||||
}
|
||||
}
|
||||
|
||||
/// The manager behind [`Self::credential_class`]. Derived from the class, so
|
||||
/// the rule is stated exactly once and the two can never disagree.
|
||||
fn governing_manager(
|
||||
&self,
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
base_url: &str,
|
||||
) -> Option<Arc<AuthManager>> {
|
||||
match self.credential_class(platform, base_url) {
|
||||
CredentialClass::Pooled => {
|
||||
platform
|
||||
.and_then(kigi_models::PlatformId::oauth)
|
||||
.map(|oauth| {
|
||||
crate::auth::oauth_registry::global_manager_for(
|
||||
&crate::auth::oauth_registry::pool_home(),
|
||||
oauth,
|
||||
)
|
||||
})
|
||||
}
|
||||
CredentialClass::Primary => self.primary.clone(),
|
||||
CredentialClass::None => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether `base_url` is the SESSION's own coding endpoint: the effective
|
||||
/// `[endpoints] coding_api_base_url` from **config.toml** (what a managed /
|
||||
/// enterprise deployment actually sets — H3), the `models_base_url`
|
||||
/// custom-endpoint mode, the `KIGI_CODE_BASE_URL` env override, a loopback
|
||||
/// dev proxy, or the compiled production endpoint.
|
||||
///
|
||||
/// Deliberately NOT [`crate::util::is_first_party_url`], which is
|
||||
/// production-only and would break every custom deployment.
|
||||
fn is_session_coding_endpoint(&self, base_url: &str) -> bool {
|
||||
if crate::util::is_effective_coding_endpoint_url(base_url) {
|
||||
return true;
|
||||
}
|
||||
if crate::util::matches_trusted_base_url(base_url, &self.endpoints.proxy_url()) {
|
||||
return true;
|
||||
}
|
||||
self.endpoints
|
||||
.models_base_url
|
||||
.as_deref()
|
||||
.is_some_and(|models_base| crate::util::matches_trusted_base_url(base_url, models_base))
|
||||
}
|
||||
|
||||
/// Whether `base_url` is `platform`'s own registry host — the guard that
|
||||
/// keeps a subscription-OAuth bearer from riding a redirected `[model.*]`
|
||||
/// override to a third party (L10).
|
||||
fn endpoint_is_platform_host(&self, platform: kigi_models::PlatformId, base_url: &str) -> bool {
|
||||
crate::util::matches_trusted_base_url(base_url, &platform.base_url())
|
||||
}
|
||||
|
||||
/// The `AuthManager` that governs this request's bearer resolution,
|
||||
/// mid-session refresh and 401 recovery — or `None` when no session
|
||||
/// credential may ride (fail fast; never a silent fallback to the primary).
|
||||
pub(crate) fn manager_for(
|
||||
&self,
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
base_url: &str,
|
||||
) -> Option<Arc<AuthManager>> {
|
||||
self.governing_manager(platform, base_url)
|
||||
}
|
||||
|
||||
/// The session bearer to stamp as this request's `api_key`, or `None`.
|
||||
///
|
||||
/// The ONLY production constructor of [`SessionCredential`].
|
||||
pub(crate) fn credential_for(
|
||||
&self,
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
base_url: &str,
|
||||
) -> Option<SessionCredential> {
|
||||
self.governing_manager(platform, base_url)
|
||||
.and_then(|am| am.current_or_expired())
|
||||
.map(|auth| SessionCredential(auth.key))
|
||||
}
|
||||
|
||||
/// A live sampler `bearer_resolver` over the governing manager, so the
|
||||
/// request keeps mid-session refresh / 401 recovery against the credential
|
||||
/// that actually belongs to its host.
|
||||
pub(crate) fn bearer_resolver_for(
|
||||
&self,
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
base_url: &str,
|
||||
) -> Option<kigi_sampler::SharedBearerResolver> {
|
||||
self.manager_for(platform, base_url)
|
||||
.map(crate::session::acp_session::sampler_turn::auth_manager_bearer_resolver)
|
||||
}
|
||||
|
||||
/// [`Self::credential_for`] for a resolved catalog entry: derives the
|
||||
/// platform and the base URL from the SAME entry, so the two can never be
|
||||
/// mismatched by a call site.
|
||||
pub(crate) fn credential_for_model(&self, entry: &ModelEntry) -> Option<SessionCredential> {
|
||||
let info = entry.info();
|
||||
self.credential_for(entry_platform(entry), &info.base_url)
|
||||
}
|
||||
|
||||
/// [`Self::credential_for`] for the catalog model a routing slug resolves
|
||||
/// to. `current_key` is the SESSION's own selected catalog key (see
|
||||
/// [`crate::agent::models::entry_for_slug`]); pass `None` for aux /
|
||||
/// override slugs, which are not the session's selection.
|
||||
///
|
||||
/// M5: a slug that is NOT in the catalog resolves through the SAME endpoint
|
||||
/// rule against the aux fallback endpoint
|
||||
/// (`EndpointsConfig::resolve_inference_base_url`, which is exactly where
|
||||
/// `resolve_aux_model_sampling_config`'s Tier-2 entry routes) instead of
|
||||
/// being handed the primary unconditionally — the old "first-party by
|
||||
/// construction" justification was false once `models_base_url` could point
|
||||
/// anywhere.
|
||||
///
|
||||
/// M6: the platform and the base URL come from ONE
|
||||
/// [`crate::agent::models::entry_for_slug`] lookup, so they can no longer
|
||||
/// disagree (the aux path used to resolve the platform with `current_key`
|
||||
/// and the credential with a separate `find_model_by_id`).
|
||||
pub(crate) fn credential_for_slug(
|
||||
&self,
|
||||
models: &indexmap::IndexMap<String, ModelEntry>,
|
||||
current_key: Option<&str>,
|
||||
slug: &str,
|
||||
) -> Option<SessionCredential> {
|
||||
match crate::agent::models::entry_for_slug(models, current_key, slug) {
|
||||
Some(entry) => self.credential_for_model(entry),
|
||||
None => self.credential_for(None, &self.endpoints.resolve_inference_base_url()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The registry platform a catalog entry belongs to (`info.id` is the managed
|
||||
/// key `{platform}/{model}`). `None` for a bare / `[model.*]` entry.
|
||||
pub(crate) fn entry_platform(entry: &ModelEntry) -> Option<kigi_models::PlatformId> {
|
||||
entry
|
||||
.info()
|
||||
.id
|
||||
.as_deref()
|
||||
.and_then(kigi_models::parse_managed_model_key)
|
||||
.map(|(platform, _)| platform)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::auth::{AuthMode, KimiAuth, KimiCodeConfig};
|
||||
|
||||
/// A primary holding a fixed in-memory bearer. The `TempDir` is returned so
|
||||
/// the caller keeps it alive; the token is read from memory, so on-disk
|
||||
/// contents are irrelevant.
|
||||
fn primary(key: &str) -> (tempfile::TempDir, Arc<AuthManager>) {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let manager = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
|
||||
manager.hot_swap(KimiAuth {
|
||||
key: key.to_string(),
|
||||
auth_mode: AuthMode::OAuth,
|
||||
..KimiAuth::test_default()
|
||||
});
|
||||
(dir, manager)
|
||||
}
|
||||
|
||||
fn authority(endpoints: EndpointsConfig, primary: Arc<AuthManager>) -> CredentialAuthority {
|
||||
CredentialAuthority::new(endpoints, Some(primary))
|
||||
}
|
||||
|
||||
fn platform(id: &str) -> kigi_models::PlatformId {
|
||||
kigi_models::PlatformId::parse(id).expect("known platform")
|
||||
}
|
||||
|
||||
/// H3 (REGRESSION): the effective coding endpoint is
|
||||
/// `EndpointsConfig::proxy_url()`, which prefers `[endpoints]
|
||||
/// coding_api_base_url` from **config.toml** — the key the managed-config
|
||||
/// sync writes. A predicate that knows only `KIGI_CODE_BASE_URL` classifies
|
||||
/// such a deployment as third-party, withholds the api_key AND the
|
||||
/// resolver, and 401s on every turn.
|
||||
///
|
||||
/// Revert-to-red: drop the `proxy_url()` arm from
|
||||
/// `is_session_coding_endpoint` (leaving only
|
||||
/// `is_effective_coding_endpoint_url`) and every assertion here fails —
|
||||
/// with NO env var set anywhere in the test.
|
||||
#[test]
|
||||
fn config_toml_coding_endpoint_still_rides_the_session_bearer() {
|
||||
let (_d, kimi) = primary("kimi-tok");
|
||||
let managed = "https://proxy.acme.com/v1";
|
||||
let auth = authority(
|
||||
EndpointsConfig {
|
||||
coding_api_base_url: Some(managed.to_string()),
|
||||
..EndpointsConfig::default()
|
||||
},
|
||||
kimi.clone(),
|
||||
);
|
||||
assert_eq!(
|
||||
auth.credential_class(None, managed),
|
||||
CredentialClass::Primary,
|
||||
"a [model.*] entry inheriting the managed coding endpoint takes the session bearer"
|
||||
);
|
||||
assert_eq!(
|
||||
auth.credential_for(None, managed)
|
||||
.map(|c| c.expose().to_owned()),
|
||||
Some("kimi-tok".to_string()),
|
||||
"the managed deployment must still receive the session bearer"
|
||||
);
|
||||
assert!(
|
||||
auth.manager_for(None, managed).is_some(),
|
||||
"and must keep a live manager, or it loses refresh and 401 recovery"
|
||||
);
|
||||
// kimi-code entries route to `proxy_url()` too (models_fetch's
|
||||
// `platform_fetch_base`), so the platform arm must honour it as well.
|
||||
assert_eq!(
|
||||
auth.credential_for(Some(platform("kimi-code")), managed)
|
||||
.map(|c| c.expose().to_owned()),
|
||||
Some("kimi-tok".to_string()),
|
||||
);
|
||||
// A DIFFERENT authority (no managed key configured) must NOT trust it.
|
||||
let default_auth = authority(EndpointsConfig::default(), kimi);
|
||||
assert_eq!(
|
||||
default_auth.credential_class(None, managed),
|
||||
CredentialClass::None,
|
||||
"the managed host is only trusted for the session that configured it"
|
||||
);
|
||||
}
|
||||
|
||||
/// The `models_base_url` custom-endpoint mode is equally invisible to the
|
||||
/// env-var-only predicate.
|
||||
#[test]
|
||||
fn config_toml_models_base_url_still_rides_the_session_bearer() {
|
||||
let (_d, kimi) = primary("kimi-tok");
|
||||
let custom = "https://models.acme.internal/v1";
|
||||
let auth = authority(
|
||||
EndpointsConfig {
|
||||
models_base_url: Some(custom.to_string()),
|
||||
..EndpointsConfig::default()
|
||||
},
|
||||
kimi,
|
||||
);
|
||||
assert_eq!(
|
||||
auth.credential_for(None, custom)
|
||||
.map(|c| c.expose().to_owned()),
|
||||
Some("kimi-tok".to_string()),
|
||||
);
|
||||
}
|
||||
|
||||
/// The compiled production endpoint and loopback proxies are unchanged.
|
||||
#[test]
|
||||
fn production_and_loopback_endpoints_are_unchanged() {
|
||||
let (_d, kimi) = primary("kimi-tok");
|
||||
let auth = authority(EndpointsConfig::default(), kimi);
|
||||
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_eq!(
|
||||
auth.credential_class(None, url),
|
||||
CredentialClass::Primary,
|
||||
"{url}: the session's own endpoint is byte-identical"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// LEAK guard: every API-key registry platform, and any platform-less model
|
||||
/// on a third-party host, gets NO session credential and NO manager.
|
||||
#[test]
|
||||
fn third_party_endpoints_never_receive_the_primary() {
|
||||
let (_d, kimi) = primary("kimi-tok");
|
||||
let auth = authority(EndpointsConfig::default(), kimi);
|
||||
for id in [
|
||||
"deepseek",
|
||||
"openai",
|
||||
"anthropic",
|
||||
"moonshot-cn",
|
||||
"moonshot-ai",
|
||||
] {
|
||||
let p = platform(id);
|
||||
assert!(
|
||||
auth.credential_for(Some(p), &p.base_url()).is_none(),
|
||||
"LEAK: {id} is an API-key platform — no session bearer may ride there"
|
||||
);
|
||||
assert!(auth.manager_for(Some(p), &p.base_url()).is_none());
|
||||
}
|
||||
for url in [
|
||||
"https://api.openai.com/v1",
|
||||
"https://api.deepseek.com/v1",
|
||||
"https://api.moonshot.cn/v1",
|
||||
"",
|
||||
] {
|
||||
assert!(
|
||||
auth.credential_for(None, url).is_none(),
|
||||
"LEAK: {url} is a third-party host"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// L10: an OAuth platform whose `[model.*]` override redirects `base_url`
|
||||
/// to a third-party host keeps `info.id` — and must NOT ship that
|
||||
/// platform's pooled OAuth bearer there.
|
||||
#[tokio::test]
|
||||
async fn oauth_platform_redirected_to_a_third_party_host_gets_nothing() {
|
||||
let (_d, kimi) = primary("kimi-tok");
|
||||
let auth = authority(EndpointsConfig::default(), kimi);
|
||||
for id in [
|
||||
"claude-pro-max",
|
||||
"openai-codex",
|
||||
"github-copilot",
|
||||
"xai-grok",
|
||||
] {
|
||||
let p = platform(id);
|
||||
assert!(
|
||||
auth.manager_for(Some(p), &p.base_url()).is_some(),
|
||||
"{id} keeps its own pooled manager on its own host"
|
||||
);
|
||||
assert!(
|
||||
auth.manager_for(Some(p), "https://third.party/v1")
|
||||
.is_none(),
|
||||
"LEAK: {id} redirected to a third-party host must ship no bearer"
|
||||
);
|
||||
assert_eq!(
|
||||
auth.credential_class(Some(p), "https://third.party/v1"),
|
||||
CredentialClass::None,
|
||||
"LEAK: {id} redirected to a third-party host takes no session credential"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// C1 — a subscription platform's own host DOES take a session credential,
|
||||
/// but it is that platform's POOLED token, never the primary / house key.
|
||||
/// Two look-alike booleans used to encode this, and picking the wrong one is
|
||||
/// the whole defect; one classifier makes the distinction impossible to
|
||||
/// mis-read.
|
||||
#[test]
|
||||
fn a_subscription_host_classifies_pooled_never_primary() {
|
||||
let (_d, kimi) = primary("kimi-tok");
|
||||
let auth = authority(EndpointsConfig::default(), kimi);
|
||||
for id in [
|
||||
"claude-pro-max",
|
||||
"openai-codex",
|
||||
"github-copilot",
|
||||
"xai-grok",
|
||||
] {
|
||||
let p = platform(id);
|
||||
assert_eq!(
|
||||
auth.credential_class(Some(p), &p.base_url()),
|
||||
CredentialClass::Pooled,
|
||||
"LEAK: {id}'s own host takes its POOLED token — never the primary / house key"
|
||||
);
|
||||
}
|
||||
// Every API-key registry platform: nothing at all.
|
||||
for id in ["deepseek", "openai", "anthropic", "moonshot-cn"] {
|
||||
let p = platform(id);
|
||||
assert_eq!(
|
||||
auth.credential_class(Some(p), &p.base_url()),
|
||||
CredentialClass::None
|
||||
);
|
||||
}
|
||||
// The primary channel is unchanged: kimi-code and a platform-less model
|
||||
// on the session's own endpoint, and nothing on a third-party host.
|
||||
for url in [
|
||||
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
|
||||
"http://127.0.0.1:8080/v1",
|
||||
] {
|
||||
assert_eq!(auth.credential_class(None, url), CredentialClass::Primary);
|
||||
assert_eq!(
|
||||
auth.credential_class(Some(platform("kimi-code")), url),
|
||||
CredentialClass::Primary
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
auth.credential_class(None, "https://api.openai.com/v1"),
|
||||
CredentialClass::None
|
||||
);
|
||||
}
|
||||
|
||||
/// The four subscription-OAuth platforms draw from their OWN pooled
|
||||
/// managers — never the primary Kimi one, even under a Kimi session.
|
||||
#[tokio::test]
|
||||
async fn oauth_platforms_never_resolve_the_primary() {
|
||||
let (_d, kimi) = primary("kimi-tok");
|
||||
let auth = authority(EndpointsConfig::default(), kimi.clone());
|
||||
for id in [
|
||||
"claude-pro-max",
|
||||
"openai-codex",
|
||||
"github-copilot",
|
||||
"xai-grok",
|
||||
] {
|
||||
let p = platform(id);
|
||||
let resolved = auth
|
||||
.manager_for(Some(p), &p.base_url())
|
||||
.expect("pooled manager");
|
||||
assert!(
|
||||
!Arc::ptr_eq(&resolved, &kimi),
|
||||
"{id} must NOT resolve the primary Kimi manager"
|
||||
);
|
||||
assert_ne!(
|
||||
auth.credential_for(Some(p), &p.base_url())
|
||||
.map(|c| c.expose().to_owned()),
|
||||
Some("kimi-tok".to_string()),
|
||||
"{id} must never receive the primary Kimi bearer"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -13,15 +13,70 @@
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::auth::kimi_oauth::{
|
||||
DeviceAuthorization, DevicePollResult, poll_device_token, request_device_authorization,
|
||||
};
|
||||
use kigi_models::OAuthConfig;
|
||||
|
||||
use crate::auth::kimi_oauth::{DeviceAuthorization, DevicePollResult};
|
||||
use crate::auth::{AuthChannels, AuthManager, AuthUrlInfo, AuthUrlMode, KimiAuth};
|
||||
|
||||
/// Extra wait added to the poll interval when the server answers `slow_down`
|
||||
/// (OAuth-standard device-flow backpressure).
|
||||
const SLOW_DOWN_INCREMENT_SECS: u64 = 5;
|
||||
|
||||
/// The wire behind a device-code login. The `Kimi` arm calls the bespoke Kimi
|
||||
/// Code wire (X-Msh headers, `/api/oauth/*`) verbatim — byte-identical to the
|
||||
/// pre-generalization path; the `Generic` arm drives a registry
|
||||
/// [`OAuthConfig`] provider (xai-grok) through [`crate::auth::oauth_device`].
|
||||
enum DeviceFlowBackend<'a> {
|
||||
Kimi {
|
||||
host: &'a str,
|
||||
},
|
||||
Generic(&'a OAuthConfig),
|
||||
/// GitHub Copilot two-stage flow: the device authorization is the generic
|
||||
/// one, but the token POLL reads GitHub's 200-body errors, and login
|
||||
/// FINALIZES the durable github token into a copilot session token via the
|
||||
/// Stage-2 exchange (see [`DeviceFlowBackend::finalize`]).
|
||||
GithubCopilot(&'a OAuthConfig),
|
||||
}
|
||||
|
||||
impl DeviceFlowBackend<'_> {
|
||||
async fn request(&self) -> anyhow::Result<DeviceAuthorization> {
|
||||
match self {
|
||||
Self::Kimi { host } => {
|
||||
crate::auth::kimi_oauth::request_device_authorization(host).await
|
||||
}
|
||||
Self::Generic(cfg) | Self::GithubCopilot(cfg) => {
|
||||
crate::auth::oauth_device::request_device_authorization(cfg).await
|
||||
}
|
||||
}
|
||||
}
|
||||
async fn poll(&self, device_code: &str) -> anyhow::Result<DevicePollResult> {
|
||||
match self {
|
||||
Self::Kimi { host } => {
|
||||
crate::auth::kimi_oauth::poll_device_token(host, device_code).await
|
||||
}
|
||||
Self::Generic(cfg) => {
|
||||
crate::auth::oauth_device::poll_device_token(cfg, device_code).await
|
||||
}
|
||||
Self::GithubCopilot(cfg) => {
|
||||
crate::auth::github_copilot::poll_github_device_token(cfg, device_code).await
|
||||
}
|
||||
}
|
||||
}
|
||||
/// Transform the device-grant credential before it is persisted. The Kimi
|
||||
/// and generic flows persist the poll result verbatim; the GitHub Copilot
|
||||
/// flow exchanges the durable github token (in `auth.key`) for the
|
||||
/// short-lived copilot token, persisting BOTH (copilot as `key`, github as
|
||||
/// `refresh_token`).
|
||||
async fn finalize(&self, auth: KimiAuth) -> anyhow::Result<KimiAuth> {
|
||||
match self {
|
||||
Self::Kimi { .. } | Self::Generic(_) => Ok(auth),
|
||||
Self::GithubCopilot(cfg) => {
|
||||
crate::auth::github_copilot::exchange_copilot_token(cfg, &auth.key).await
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Outcome of one full poll loop over a single device authorization.
|
||||
enum PollLoopOutcome {
|
||||
/// Access token issued.
|
||||
@@ -40,11 +95,46 @@ pub async fn run_device_code_login_channels(
|
||||
host: &str,
|
||||
auth_manager: &Arc<AuthManager>,
|
||||
channels: &mut Option<AuthChannels>,
|
||||
) -> anyhow::Result<(KimiAuth, bool)> {
|
||||
run_device_code_login_backend(DeviceFlowBackend::Kimi { host }, auth_manager, channels).await
|
||||
}
|
||||
|
||||
/// Device-code login for a GENERIC [`OAuthConfig`] provider (xai-grok). Same
|
||||
/// TUI/CLI presentation as the Kimi login; only the wire differs.
|
||||
pub async fn run_device_code_login_generic(
|
||||
oauth: &OAuthConfig,
|
||||
auth_manager: &Arc<AuthManager>,
|
||||
channels: &mut Option<AuthChannels>,
|
||||
) -> anyhow::Result<(KimiAuth, bool)> {
|
||||
run_device_code_login_backend(DeviceFlowBackend::Generic(oauth), auth_manager, channels).await
|
||||
}
|
||||
|
||||
/// GitHub Copilot two-stage login (github-copilot): the same device-flow
|
||||
/// presentation as the generic path, but the token poll reads GitHub's 200-body
|
||||
/// errors and the minted github token is finalized into a copilot session token
|
||||
/// before it is persisted (see [`DeviceFlowBackend::finalize`]).
|
||||
pub async fn run_device_code_login_github_copilot(
|
||||
oauth: &OAuthConfig,
|
||||
auth_manager: &Arc<AuthManager>,
|
||||
channels: &mut Option<AuthChannels>,
|
||||
) -> anyhow::Result<(KimiAuth, bool)> {
|
||||
run_device_code_login_backend(
|
||||
DeviceFlowBackend::GithubCopilot(oauth),
|
||||
auth_manager,
|
||||
channels,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn run_device_code_login_backend(
|
||||
backend: DeviceFlowBackend<'_>,
|
||||
auth_manager: &Arc<AuthManager>,
|
||||
channels: &mut Option<AuthChannels>,
|
||||
) -> anyhow::Result<(KimiAuth, bool)> {
|
||||
let interactive_tui = channels.is_some();
|
||||
let mut channels = channels.take();
|
||||
loop {
|
||||
let device_auth = request_device_authorization(host).await?;
|
||||
let device_auth = backend.request().await?;
|
||||
let display_uri = device_auth.verification_uri_complete.clone();
|
||||
|
||||
if interactive_tui {
|
||||
@@ -62,10 +152,14 @@ pub async fn run_device_code_login_channels(
|
||||
prompt_on_stderr(&device_auth).await;
|
||||
}
|
||||
|
||||
match complete_device_code_login(host, &device_auth).await? {
|
||||
match complete_device_code_login(&backend, &device_auth).await? {
|
||||
PollLoopOutcome::Done(auth) => {
|
||||
// Finalize before persisting: the GitHub Copilot flow exchanges
|
||||
// the durable github token for the short-lived copilot token
|
||||
// here; the Kimi / generic flows pass the credential through.
|
||||
let auth = backend.finalize(*auth).await?;
|
||||
let auth = auth_manager
|
||||
.update(*auth)
|
||||
.update(auth)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to save credentials: {e}"))?;
|
||||
return Ok((auth, true));
|
||||
@@ -112,7 +206,7 @@ async fn prompt_on_stderr(device_auth: &DeviceAuthorization) {
|
||||
/// Poll the token endpoint until the user approves, the device code expires
|
||||
/// (→ [`PollLoopOutcome::Restart`]), or the wire fails.
|
||||
async fn complete_device_code_login(
|
||||
host: &str,
|
||||
backend: &DeviceFlowBackend<'_>,
|
||||
device_auth: &DeviceAuthorization,
|
||||
) -> anyhow::Result<PollLoopOutcome> {
|
||||
let mut poll_interval = std::time::Duration::from_secs(device_auth.interval.max(1) as u64);
|
||||
@@ -120,7 +214,7 @@ async fn complete_device_code_login(
|
||||
// Sleep first: an immediate poll on a fresh code only returns
|
||||
// authorization_pending (and risks slow_down).
|
||||
tokio::time::sleep(poll_interval).await;
|
||||
match poll_device_token(host, &device_auth.device_code).await? {
|
||||
match backend.poll(&device_auth.device_code).await? {
|
||||
DevicePollResult::Success(auth) => {
|
||||
tracing::info!("auth: device login authorized");
|
||||
return Ok(PollLoopOutcome::Done(auth));
|
||||
@@ -148,8 +242,8 @@ async fn complete_device_code_login(
|
||||
/// Open `url` in the browser off-thread: `webbrowser::open` is synchronous and
|
||||
/// would stall the single-threaded TUI loop. Returns `true` on success so the
|
||||
/// caller can decide how to notify the user (eprintln on CLI, nothing on TUI
|
||||
/// where the URL is already rendered in the widget).
|
||||
async fn open_browser_detached(url: &str) -> bool {
|
||||
/// where the URL is already rendered in the widget). Shared with the PKCE flow.
|
||||
pub(super) async fn open_browser_detached(url: &str) -> bool {
|
||||
// Unit tests drive the full login flow against mock servers — their
|
||||
// fixture URLs must never reach a real browser.
|
||||
if cfg!(test) {
|
||||
@@ -338,6 +432,100 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
/// GitHub Copilot two-stage login e2e (mock wire): device authorization →
|
||||
/// poll (pending → github token) → Stage-2 copilot-token exchange (Bearer
|
||||
/// github token + editor headers → copilot token + expiry). The persisted
|
||||
/// credential keys the COPILOT token, keeps the GITHUB token as
|
||||
/// `refresh_token`, and carries the copilot expiry.
|
||||
#[tokio::test]
|
||||
async fn github_copilot_two_stage_login_persists_copilot_and_github() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
let cfg = OAuthConfig {
|
||||
auth_host: host,
|
||||
token_host: host,
|
||||
copilot_exchange: Some((host, "/copilot_internal/v2/token")),
|
||||
..kigi_models::COPILOT_OAUTH_CONFIG
|
||||
};
|
||||
// Stage 1a: device authorization (github.com/login/device/code).
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/login/device/code"))
|
||||
.and(body_string_contains("client_id=Iv1.b507a08c87ecfe98"))
|
||||
.and(body_string_contains("scope=read"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"device_code": "gh-dev-1",
|
||||
"user_code": "WDJB-MJHT",
|
||||
"verification_uri": "https://github.com/login/device",
|
||||
"expires_in": 900,
|
||||
"interval": 0,
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
// Stage 1b: token poll — GitHub returns errors AND success in a 200 body.
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/login/oauth/access_token"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_json(serde_json::json!({ "error": "authorization_pending" })),
|
||||
)
|
||||
.up_to_n_times(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/login/oauth/access_token"))
|
||||
.and(body_string_contains("device_code=gh-dev-1"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_json(serde_json::json!({ "access_token": "gho_github_tok" })),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
// Stage 2: copilot-token exchange (Bearer github token + editor headers).
|
||||
let future = (chrono::Utc::now() + chrono::Duration::minutes(30)).timestamp();
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/copilot_internal/v2/token"))
|
||||
.and(wiremock::matchers::header(
|
||||
"Authorization",
|
||||
"Bearer gho_github_tok",
|
||||
))
|
||||
.and(wiremock::matchers::header(
|
||||
"Editor-Plugin-Version",
|
||||
"copilot-chat/0.35.0",
|
||||
))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(
|
||||
serde_json::json!({ "token": "copilot-session-tok", "expires_at": future }),
|
||||
))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mgr = auth_manager(&dir);
|
||||
let mut channels = None;
|
||||
let (auth, is_new) = run_device_code_login_github_copilot(&cfg, &mgr, &mut channels)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(is_new);
|
||||
assert_eq!(
|
||||
auth.key, "copilot-session-tok",
|
||||
"key = copilot session token"
|
||||
);
|
||||
assert_eq!(
|
||||
auth.refresh_token.as_deref(),
|
||||
Some("gho_github_tok"),
|
||||
"the durable github token is persisted as refresh_token"
|
||||
);
|
||||
assert!(
|
||||
auth.expires_at.is_some_and(|e| e > chrono::Utc::now()),
|
||||
"the copilot expiry must be persisted"
|
||||
);
|
||||
assert_eq!(
|
||||
mgr.current_or_expired().map(|a| a.key),
|
||||
Some("copilot-session-tok".into()),
|
||||
"login must land the copilot token in the manager cache"
|
||||
);
|
||||
}
|
||||
|
||||
/// A 5xx from the token endpoint is a hard error (kimi-cli parity).
|
||||
#[tokio::test]
|
||||
async fn server_error_during_poll_fails_login() {
|
||||
|
||||
@@ -67,6 +67,178 @@ pub async fn run_auth_flow(
|
||||
run_auth_flow_inner(auth_manager, kimi_code_config, reauth, false, channels).await
|
||||
}
|
||||
|
||||
/// Login flow for a GENERIC OAuth provider (xai-grok device-code,
|
||||
/// claude-pro-max PKCE-localhost): use a valid cached session unless re-authing,
|
||||
/// otherwise dispatch by `oauth.flow` to the device-code or PKCE-localhost login
|
||||
/// (persisting under the provider's own scope via `auth_manager`). Unlike the
|
||||
/// Kimi flow this does not run the silent-refresh dance — the login's
|
||||
/// `AuthManager::update` persists a fresh token set directly.
|
||||
pub async fn run_oauth_provider_flow(
|
||||
auth_manager: &Arc<AuthManager>,
|
||||
oauth: &'static kigi_models::OAuthConfig,
|
||||
reauth: bool,
|
||||
channels: Option<AuthChannels>,
|
||||
) -> anyhow::Result<(KimiAuth, bool)> {
|
||||
tracing::info!(
|
||||
scope_key = oauth.scope_key,
|
||||
reauth,
|
||||
"auth: starting generic oauth login"
|
||||
);
|
||||
if reauth {
|
||||
auth_manager.clear()?;
|
||||
}
|
||||
if !reauth && let Some(auth) = auth_manager.current() {
|
||||
tracing::info!(
|
||||
scope_key = oauth.scope_key,
|
||||
"auth: using cached oauth session"
|
||||
);
|
||||
return Ok((auth, false));
|
||||
}
|
||||
let mut channels = channels;
|
||||
match oauth.flow {
|
||||
kigi_models::OAuthFlow::DeviceCode => {
|
||||
crate::auth::device_code::run_device_code_login_generic(
|
||||
oauth,
|
||||
auth_manager,
|
||||
&mut channels,
|
||||
)
|
||||
.await
|
||||
}
|
||||
kigi_models::OAuthFlow::PkceLocalhost {
|
||||
redirect_port,
|
||||
redirect_path,
|
||||
} => {
|
||||
run_pkce_localhost_login(
|
||||
oauth,
|
||||
redirect_port,
|
||||
redirect_path,
|
||||
auth_manager,
|
||||
&mut channels,
|
||||
)
|
||||
.await
|
||||
}
|
||||
kigi_models::OAuthFlow::GithubDeviceCopilot => {
|
||||
crate::auth::device_code::run_device_code_login_github_copilot(
|
||||
oauth,
|
||||
auth_manager,
|
||||
&mut channels,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// PKCE-localhost login (claude-pro-max JSON, openai-codex FORM): generate PKCE,
|
||||
/// present the browser authorize URL (TUI channel or stderr), open the browser,
|
||||
/// then await the code from EITHER the `127.0.0.1:{redirect_port}{redirect_path}`
|
||||
/// loopback callback OR a manual paste (headless fallback). Exchange it at the
|
||||
/// token endpoint and persist.
|
||||
///
|
||||
/// The `token_body` selects the wire dialect: `Json` is the claude path
|
||||
/// (`state == verifier`, JSON exchange carrying `state`); `Form` is the codex
|
||||
/// path (fresh-random `state`, FORM exchange WITHOUT `state`, then a FAIL-FAST
|
||||
/// check that the minted JWT carries a `chatgpt_account_id` — a token without it
|
||||
/// is useless for inference, so the login bails rather than persisting it).
|
||||
///
|
||||
/// SECURITY: the verifier / code / tokens / JWT / account id are never logged;
|
||||
/// the loopback binds `127.0.0.1` only and validates `state` strictly.
|
||||
async fn run_pkce_localhost_login(
|
||||
oauth: &'static kigi_models::OAuthConfig,
|
||||
redirect_port: u16,
|
||||
redirect_path: &'static str,
|
||||
auth_manager: &Arc<AuthManager>,
|
||||
channels: &mut Option<AuthChannels>,
|
||||
) -> anyhow::Result<(KimiAuth, bool)> {
|
||||
use crate::auth::oauth_pkce;
|
||||
|
||||
// The token-body encoding also selects the PKCE `state` convention: the JSON
|
||||
// dialect is Claude's/Pi's `state == verifier`, while form-encoded endpoints
|
||||
// (the OAuth norm) get an INDEPENDENT random state. A new form provider
|
||||
// inheriting the standard random state is correct by default.
|
||||
let uses_form_exchange = matches!(oauth.token_body, kigi_models::OAuthTokenBody::Form);
|
||||
let pkce = if uses_form_exchange {
|
||||
oauth_pkce::generate_pkce_random_state()
|
||||
} else {
|
||||
oauth_pkce::generate_pkce()
|
||||
};
|
||||
let redirect = oauth_pkce::redirect_uri(redirect_port, redirect_path);
|
||||
let authorize_url = oauth_pkce::build_authorize_url(oauth, &redirect, &pkce);
|
||||
|
||||
let mut chans = channels.take();
|
||||
if let Some(tx) = chans.as_mut().and_then(|c| c.url_tx.take()) {
|
||||
// TUI: push the URL BEFORE opening the browser (never block the UI on a
|
||||
// slow/headless browser launch).
|
||||
let _ = tx.send(AuthUrlInfo {
|
||||
url: authorize_url.clone(),
|
||||
mode: AuthUrlMode::Device,
|
||||
});
|
||||
crate::auth::device_code::open_browser_detached(&authorize_url).await;
|
||||
} else {
|
||||
eprintln!();
|
||||
eprintln!("To sign in, open this URL in your browser:");
|
||||
eprintln!();
|
||||
eprintln!(" {authorize_url}");
|
||||
eprintln!();
|
||||
if !crate::auth::device_code::open_browser_detached(&authorize_url).await {
|
||||
eprintln!(" (Could not open the browser automatically — open the URL above.)");
|
||||
eprintln!();
|
||||
}
|
||||
eprintln!("Waiting for the sign-in to complete...");
|
||||
}
|
||||
|
||||
let code = await_pkce_code(redirect_port, redirect_path, &pkce, chans.as_mut()).await?;
|
||||
let auth = if uses_form_exchange {
|
||||
oauth_pkce::exchange_code_form(oauth, &code, &pkce, &redirect).await?
|
||||
} else {
|
||||
oauth_pkce::exchange_code(oauth, &code, &pkce, &redirect).await?
|
||||
};
|
||||
// FAIL FAST: an access token that yields no chatgpt_account_id cannot
|
||||
// authorize inference (it becomes the `chatgpt-account-id` header) — bail
|
||||
// rather than persist a dead session. Gated on the EXPLICIT provider fact,
|
||||
// never on the token-body encoding.
|
||||
if oauth.requires_chatgpt_account_id
|
||||
&& kigi_sampling_types::chatgpt_account_id_from_jwt(&auth.key).is_none()
|
||||
{
|
||||
anyhow::bail!(
|
||||
"ChatGPT login did not return a usable account id \
|
||||
(the access token is missing the chatgpt_account_id claim)"
|
||||
);
|
||||
}
|
||||
let auth = auth_manager
|
||||
.update(auth)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to save credentials: {e}"))?;
|
||||
Ok((auth, true))
|
||||
}
|
||||
|
||||
/// Await the authorization code: the loopback callback is primary; when a TUI
|
||||
/// channel is present, a pasted code (redirect URL / `code#state` / bare code)
|
||||
/// is accepted concurrently as a headless fallback. State is validated in both
|
||||
/// arms (strict on the loopback, mismatch-rejecting on the paste).
|
||||
async fn await_pkce_code(
|
||||
redirect_port: u16,
|
||||
redirect_path: &str,
|
||||
pkce: &crate::auth::oauth_pkce::PkceCodes,
|
||||
channels: Option<&mut AuthChannels>,
|
||||
) -> anyhow::Result<String> {
|
||||
use crate::auth::oauth_pkce;
|
||||
match channels {
|
||||
Some(ch) => {
|
||||
tokio::select! {
|
||||
code = oauth_pkce::await_loopback_code(redirect_port, redirect_path, &pkce.state) => code,
|
||||
pasted = ch.code_rx.recv() => {
|
||||
let pasted = pasted
|
||||
.ok_or_else(|| anyhow::anyhow!("auth code channel closed before a code arrived"))?;
|
||||
let params = oauth_pkce::parse_manual_paste(&pasted)?;
|
||||
oauth_pkce::validate_pasted_state(¶ms, &pkce.state)?;
|
||||
Ok(params.code)
|
||||
}
|
||||
}
|
||||
}
|
||||
None => oauth_pkce::await_loopback_code(redirect_port, redirect_path, &pkce.state).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_auth_flow_inner(
|
||||
auth_manager: &Arc<AuthManager>,
|
||||
_kimi_code_config: &KimiCodeConfig,
|
||||
|
||||
@@ -0,0 +1,485 @@
|
||||
//! GitHub Copilot two-stage OAuth wire (github-copilot), driven by a registry
|
||||
//! [`kigi_models::OAuthConfig`] whose `flow` is [`OAuthFlow::GithubDeviceCopilot`].
|
||||
//!
|
||||
//! Stage 1 — RFC-8628 device flow on `auth_host` (github.com). The device
|
||||
//! authorization POST is the generic one ([`super::oauth_device`]); the token
|
||||
//! POLL is Copilot-specific because GitHub returns its device errors in a `200`
|
||||
//! body (`{error: "authorization_pending"|"slow_down"|"expired_token"}`), not a
|
||||
//! `4xx`, and the success payload carries ONLY `access_token` (the DURABLE
|
||||
//! GitHub token — no refresh token, no expiry).
|
||||
//!
|
||||
//! Stage 2 — copilot-token exchange: `GET {copilot_exchange}` bearing the GitHub
|
||||
//! token + the editor headers re-mints the SHORT-LIVED copilot session token
|
||||
//! (`{token, expires_at}`). This runs at login ([`exchange_copilot_token`]) and
|
||||
//! on every "refresh" ([`remint_copilot_token`], dispatched by the generic
|
||||
//! refresher) — the github token is unchanged and re-persisted as the
|
||||
//! `refresh_token`; the copilot token becomes the `key`.
|
||||
//!
|
||||
//! SECURITY: the github token and the copilot token are NEVER logged (only
|
||||
//! non-secret events: poll succeeded, copilot token minted/re-minted).
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use kigi_models::OAuthConfig;
|
||||
use serde::Deserialize;
|
||||
|
||||
use super::kimi_oauth::{DevicePollResult, RefreshError};
|
||||
use super::model::{AuthMode, KimiAuth};
|
||||
|
||||
/// RFC-8628 device grant type (shared with the generic device wire).
|
||||
const DEVICE_GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type:device_code";
|
||||
/// Copilot-exchange retry budget over 5xx / network blips (parity with the
|
||||
/// device/PKCE refresh wires); 401/403 fails fast (the github token is dead).
|
||||
const MAX_EXCHANGE_RETRIES: u32 = 3;
|
||||
const RETRYABLE_EXCHANGE_STATUSES: [u16; 5] = [429, 500, 502, 503, 504];
|
||||
|
||||
/// The four VS Code Copilot editor-identity headers every Copilot request
|
||||
/// carries. Non-secret wire constants owned by `kigi_sampling_types` (the
|
||||
/// single source shared with the sampler's inference gate).
|
||||
fn editor_headers() -> [(&'static str, &'static str); 4] {
|
||||
[
|
||||
("User-Agent", kigi_sampling_types::COPILOT_USER_AGENT),
|
||||
(
|
||||
"Editor-Version",
|
||||
kigi_sampling_types::COPILOT_EDITOR_VERSION,
|
||||
),
|
||||
(
|
||||
"Editor-Plugin-Version",
|
||||
kigi_sampling_types::COPILOT_EDITOR_PLUGIN_VERSION,
|
||||
),
|
||||
(
|
||||
"Copilot-Integration-Id",
|
||||
kigi_sampling_types::COPILOT_INTEGRATION_ID,
|
||||
),
|
||||
]
|
||||
}
|
||||
|
||||
/// GitHub's device-token poll response: EITHER `access_token` (the durable
|
||||
/// GitHub token) OR an `error` (in a `200` body). No refresh token / expiry.
|
||||
#[derive(Deserialize, Default)]
|
||||
struct GithubDeviceTokenResponse {
|
||||
#[serde(default)]
|
||||
access_token: Option<String>,
|
||||
#[serde(default)]
|
||||
error: Option<String>,
|
||||
}
|
||||
|
||||
/// One poll of `POST {auth_host}{token_path}` (github.com/login/oauth/
|
||||
/// access_token) with the device grant. GitHub answers `200` for BOTH success
|
||||
/// and the pending/slow_down/expired errors, so the outcome is read from the
|
||||
/// body, not the status. On success the [`KimiAuth`] carries the GitHub token as
|
||||
/// `key` with NO refresh token / expiry — the caller finalizes it via the
|
||||
/// copilot exchange before persisting.
|
||||
pub(crate) async fn poll_github_device_token(
|
||||
cfg: &OAuthConfig,
|
||||
device_code: &str,
|
||||
) -> anyhow::Result<DevicePollResult> {
|
||||
let url = format!("{}{}", cfg.auth_host.trim_end_matches('/'), cfg.token_path);
|
||||
let resp = crate::http::shared_client()
|
||||
.post(&url)
|
||||
.header("Accept", "application/json")
|
||||
.header("User-Agent", kigi_sampling_types::COPILOT_USER_AGENT)
|
||||
.form(&[
|
||||
("client_id", cfg.client_id),
|
||||
("device_code", device_code),
|
||||
("grant_type", DEVICE_GRANT_TYPE),
|
||||
])
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Token polling request failed: {e}"))?;
|
||||
|
||||
let status = resp.status();
|
||||
if status.is_server_error() {
|
||||
anyhow::bail!("Token polling server error: {status}");
|
||||
}
|
||||
let body = resp.bytes().await?;
|
||||
let parsed: GithubDeviceTokenResponse = serde_json::from_slice(&body).unwrap_or_default();
|
||||
|
||||
if let Some(access) = parsed.access_token.filter(|t| !t.is_empty()) {
|
||||
tracing::info!("auth: github device poll succeeded, github token issued (copilot)");
|
||||
return Ok(DevicePollResult::Success(Box::new(github_token_auth(
|
||||
access,
|
||||
))));
|
||||
}
|
||||
match parsed.error.as_deref() {
|
||||
Some("expired_token") => {
|
||||
tracing::info!("auth: github device code expired; restarting (copilot)");
|
||||
Ok(DevicePollResult::Expired)
|
||||
}
|
||||
Some(error) => {
|
||||
tracing::debug!(error, "auth: github device poll pending (copilot)");
|
||||
Ok(DevicePollResult::Pending {
|
||||
error: error.to_owned(),
|
||||
description: None,
|
||||
})
|
||||
}
|
||||
None => Ok(DevicePollResult::Pending {
|
||||
error: "missing_access_token".to_owned(),
|
||||
description: None,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// A transient [`KimiAuth`] holding ONLY the durable GitHub token (no refresh
|
||||
/// token / expiry) — the intermediate device-flow result, finalized by the
|
||||
/// copilot exchange before it is ever persisted.
|
||||
fn github_token_auth(github_token: String) -> KimiAuth {
|
||||
KimiAuth {
|
||||
key: github_token,
|
||||
auth_mode: AuthMode::OAuth,
|
||||
create_time: Utc::now(),
|
||||
user_id: String::new(),
|
||||
email: None,
|
||||
refresh_token: None,
|
||||
expires_at: None,
|
||||
expires_in: None,
|
||||
scope: None,
|
||||
token_type: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Stage-2 copilot-token exchange response (`GET copilot_internal/v2/token`).
|
||||
/// `endpoints`/`proxy-ep` are ignored: Kigi resolves the base URL from the
|
||||
/// platform registry (the individual-subscription endpoint, or the
|
||||
/// `KIGI_COPILOT_BASE_URL` override).
|
||||
#[derive(Deserialize)]
|
||||
struct CopilotTokenResponse {
|
||||
token: String,
|
||||
/// Unix seconds when the copilot token expires (~30 min out).
|
||||
expires_at: i64,
|
||||
}
|
||||
|
||||
/// Materialize the persisted credential from a copilot-token exchange:
|
||||
/// `key` = the short-lived copilot token, `refresh_token` = the DURABLE github
|
||||
/// token (so every re-mint re-exchanges it), `expires_at` = the copilot expiry.
|
||||
fn copilot_auth(resp: CopilotTokenResponse, github_token: &str) -> anyhow::Result<KimiAuth> {
|
||||
let now = Utc::now();
|
||||
// FAIL-FAST: an uninterpretable expiry means we cannot schedule the re-mint,
|
||||
// so reject it rather than silently falling back to a long default TTL — which
|
||||
// would let the ~30-min copilot token 401 on the wire ~30 min later.
|
||||
let expires_at = DateTime::from_timestamp(resp.expires_at, 0).ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"copilot token has an out-of-range expires_at: {}",
|
||||
resp.expires_at
|
||||
)
|
||||
})?;
|
||||
// The manager's dynamic threshold (`max(300, expires_in × 0.5)`) drives the
|
||||
// proactive re-mint; `expires_in` is the copilot token's remaining life.
|
||||
let expires_in = (expires_at - now).num_seconds();
|
||||
Ok(KimiAuth {
|
||||
key: resp.token,
|
||||
auth_mode: AuthMode::OAuth,
|
||||
create_time: now,
|
||||
user_id: String::new(),
|
||||
email: None,
|
||||
refresh_token: Some(github_token.to_owned()),
|
||||
expires_at: Some(expires_at),
|
||||
expires_in: Some(expires_in),
|
||||
scope: None,
|
||||
token_type: Some("bearer".to_owned()),
|
||||
})
|
||||
}
|
||||
|
||||
/// The `(host, path)` of the copilot-token exchange endpoint, or a fatal error
|
||||
/// when the config lacks it (a non-Copilot config reaching this wire is a bug).
|
||||
fn exchange_url(cfg: &OAuthConfig) -> anyhow::Result<String> {
|
||||
let (host, path) = cfg.copilot_exchange.ok_or_else(|| {
|
||||
anyhow::anyhow!("github-copilot config missing copilot_exchange endpoint")
|
||||
})?;
|
||||
Ok(format!("{}{path}", host.trim_end_matches('/')))
|
||||
}
|
||||
|
||||
/// `GET {copilot_exchange}` bearing the github token + editor headers.
|
||||
async fn send_copilot_exchange(
|
||||
cfg: &OAuthConfig,
|
||||
github_token: &str,
|
||||
) -> anyhow::Result<reqwest::Response> {
|
||||
let url = exchange_url(cfg)?;
|
||||
let mut req = crate::http::shared_client()
|
||||
.get(&url)
|
||||
.header("Accept", "application/json")
|
||||
.header("Authorization", format!("Bearer {github_token}"));
|
||||
for (name, value) in editor_headers() {
|
||||
req = req.header(name, value);
|
||||
}
|
||||
req.send()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("copilot-token exchange request failed: {e}"))
|
||||
}
|
||||
|
||||
/// Exchange the durable GitHub token for a copilot session token (login path).
|
||||
/// FAIL-FAST: a non-2xx response aborts login (never a silent fallback).
|
||||
pub(crate) async fn exchange_copilot_token(
|
||||
cfg: &OAuthConfig,
|
||||
github_token: &str,
|
||||
) -> anyhow::Result<KimiAuth> {
|
||||
tracing::info!(
|
||||
scope_key = cfg.scope_key,
|
||||
"auth: exchanging github token for copilot token"
|
||||
);
|
||||
let resp = send_copilot_exchange(cfg, github_token).await?;
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
tracing::warn!(%status, scope_key = cfg.scope_key, "auth: copilot-token exchange failed");
|
||||
anyhow::bail!("Copilot token exchange failed (HTTP {status}): {body}");
|
||||
}
|
||||
let parsed: CopilotTokenResponse = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("malformed copilot token payload: {e}"))?;
|
||||
tracing::info!(scope_key = cfg.scope_key, "auth: copilot token minted");
|
||||
copilot_auth(parsed, github_token)
|
||||
}
|
||||
|
||||
/// Re-mint the copilot token from the durable GitHub token (refresh path,
|
||||
/// dispatched from the generic refresher). This is NOT a `refresh_token` grant:
|
||||
/// it re-runs the copilot exchange. `refresh_token` is the github token; the
|
||||
/// returned [`KimiAuth`] preserves it. Retries 5xx / network blips; 401/403
|
||||
/// (the github token is revoked) fails fast as [`RefreshError::Unauthorized`].
|
||||
pub(crate) async fn remint_copilot_token(
|
||||
cfg: &OAuthConfig,
|
||||
github_token: &str,
|
||||
) -> Result<KimiAuth, RefreshError> {
|
||||
let mut last_error = String::from("no attempt made");
|
||||
for attempt in 0..MAX_EXCHANGE_RETRIES {
|
||||
if attempt > 0 {
|
||||
let backoff = std::time::Duration::from_secs(1 << (attempt - 1));
|
||||
tracing::warn!(
|
||||
attempt,
|
||||
backoff_secs = backoff.as_secs(),
|
||||
"auth: retrying copilot-token re-mint"
|
||||
);
|
||||
tokio::time::sleep(backoff).await;
|
||||
}
|
||||
let resp = match send_copilot_exchange(cfg, github_token).await {
|
||||
Ok(resp) => resp,
|
||||
Err(e) => {
|
||||
last_error = format!("{e}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let status = resp.status().as_u16();
|
||||
let bytes = resp.bytes().await.unwrap_or_default();
|
||||
if status == 401 || status == 403 {
|
||||
return Err(RefreshError::Unauthorized {
|
||||
status,
|
||||
description: "GitHub token rejected at copilot-token exchange.".to_owned(),
|
||||
});
|
||||
}
|
||||
if status == 200 {
|
||||
return match serde_json::from_slice::<CopilotTokenResponse>(&bytes) {
|
||||
Ok(parsed) => match copilot_auth(parsed, github_token) {
|
||||
Ok(auth) => {
|
||||
tracing::info!(scope_key = cfg.scope_key, "auth: copilot token re-minted");
|
||||
Ok(auth)
|
||||
}
|
||||
Err(e) => Err(RefreshError::Fatal {
|
||||
status,
|
||||
description: format!("{e}"),
|
||||
}),
|
||||
},
|
||||
Err(e) => Err(RefreshError::Fatal {
|
||||
status,
|
||||
description: format!("malformed copilot token payload: {e}"),
|
||||
}),
|
||||
};
|
||||
}
|
||||
let description = format!("copilot-token exchange failed (HTTP {status}).");
|
||||
if RETRYABLE_EXCHANGE_STATUSES.contains(&status) {
|
||||
last_error = description;
|
||||
continue;
|
||||
}
|
||||
return Err(RefreshError::Fatal {
|
||||
status,
|
||||
description,
|
||||
});
|
||||
}
|
||||
Err(RefreshError::Exhausted { last_error })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use chrono::Duration;
|
||||
use kigi_models::COPILOT_OAUTH_CONFIG;
|
||||
use wiremock::matchers::{body_string_contains, header, method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
/// A COPILOT_OAUTH_CONFIG pointed at a mock server for both stages.
|
||||
fn mock_cfg(host: &'static str, exchange: &'static str) -> OAuthConfig {
|
||||
OAuthConfig {
|
||||
auth_host: host,
|
||||
token_host: host,
|
||||
copilot_exchange: Some((exchange, "/copilot_internal/v2/token")),
|
||||
..COPILOT_OAUTH_CONFIG
|
||||
}
|
||||
}
|
||||
|
||||
/// GitHub's device poll returns pending errors in a 200 body — mapped to
|
||||
/// Pending (authorization_pending / slow_down) and Expired (expired_token),
|
||||
/// never mis-read as a token.
|
||||
#[tokio::test]
|
||||
async fn github_device_poll_maps_200_body_errors() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/login/oauth/access_token"))
|
||||
.and(body_string_contains("grant_type=urn"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_json(serde_json::json!({ "error": "authorization_pending" })),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let result = poll_github_device_token(&mock_cfg(host, host), "dev-1")
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
matches!(result, DevicePollResult::Pending { error, .. } if error == "authorization_pending")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn github_device_poll_expired_restarts() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/login/oauth/access_token"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_json(serde_json::json!({ "error": "expired_token" })),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let result = poll_github_device_token(&mock_cfg(host, host), "dev-1")
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(result, DevicePollResult::Expired));
|
||||
}
|
||||
|
||||
/// A successful poll yields the DURABLE github token as `key` with NO
|
||||
/// refresh token / expiry (the copilot exchange finalizes it next).
|
||||
#[tokio::test]
|
||||
async fn github_device_poll_success_is_bare_github_token() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/login/oauth/access_token"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_json(serde_json::json!({ "access_token": "gho_github_tok" })),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let DevicePollResult::Success(auth) = poll_github_device_token(&mock_cfg(host, host), "d")
|
||||
.await
|
||||
.unwrap()
|
||||
else {
|
||||
panic!("expected success");
|
||||
};
|
||||
assert_eq!(auth.key, "gho_github_tok");
|
||||
assert_eq!(
|
||||
auth.refresh_token, None,
|
||||
"github token is not a refresh grant"
|
||||
);
|
||||
assert_eq!(auth.expires_at, None, "the github token is long-lived");
|
||||
}
|
||||
|
||||
/// The Stage-2 exchange rides the github Bearer + editor headers and maps
|
||||
/// `{token, expires_at}` onto `key=copilot`, `refresh_token=github`, with a
|
||||
/// future `expires_at`.
|
||||
#[tokio::test]
|
||||
async fn copilot_exchange_maps_token_and_persists_github_as_refresh() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
let future = (Utc::now() + Duration::minutes(30)).timestamp();
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/copilot_internal/v2/token"))
|
||||
.and(header("Authorization", "Bearer gho_github_tok"))
|
||||
.and(header("Editor-Version", "vscode/1.107.0"))
|
||||
.and(header("Copilot-Integration-Id", "vscode-chat"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(
|
||||
serde_json::json!({ "token": "tid=abc;copilot-tok", "expires_at": future }),
|
||||
))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let auth = exchange_copilot_token(&mock_cfg(host, host), "gho_github_tok")
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(auth.key, "tid=abc;copilot-tok", "key = copilot token");
|
||||
assert_eq!(
|
||||
auth.refresh_token.as_deref(),
|
||||
Some("gho_github_tok"),
|
||||
"the durable github token is persisted as refresh_token"
|
||||
);
|
||||
assert!(
|
||||
auth.expires_at.is_some_and(|e| e > Utc::now()),
|
||||
"copilot expiry must be in the future"
|
||||
);
|
||||
}
|
||||
|
||||
/// The copilot re-mint (refresh) re-exchanges the github token for a NEW
|
||||
/// copilot token, keeping the github token as refresh_token.
|
||||
#[tokio::test]
|
||||
async fn copilot_remint_returns_new_copilot_token() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
let future = (Utc::now() + Duration::minutes(30)).timestamp();
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/copilot_internal/v2/token"))
|
||||
.and(header("Authorization", "Bearer gho_github_tok"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(
|
||||
serde_json::json!({ "token": "copilot-tok-2", "expires_at": future }),
|
||||
))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let auth = remint_copilot_token(&mock_cfg(host, host), "gho_github_tok")
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(auth.key, "copilot-tok-2");
|
||||
assert_eq!(auth.refresh_token.as_deref(), Some("gho_github_tok"));
|
||||
}
|
||||
|
||||
/// A 401 at the exchange (github token revoked) fails fast as Unauthorized
|
||||
/// (drives the manager's permanent-failure / re-login path).
|
||||
#[tokio::test]
|
||||
async fn copilot_remint_401_is_unauthorized() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/copilot_internal/v2/token"))
|
||||
.respond_with(ResponseTemplate::new(401))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let err = remint_copilot_token(&mock_cfg(host, host), "dead-github-tok")
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
matches!(err, RefreshError::Unauthorized { status: 401, .. }),
|
||||
"got {err:?}"
|
||||
);
|
||||
}
|
||||
|
||||
/// FAIL-FAST: an out-of-range `expires_at` is rejected rather than silently
|
||||
/// degrading to a long default TTL (which would 401 on the wire ~30 min in).
|
||||
#[tokio::test]
|
||||
async fn copilot_exchange_rejects_out_of_range_expiry() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/copilot_internal/v2/token"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(
|
||||
serde_json::json!({ "token": "copilot-tok", "expires_at": i64::MAX }),
|
||||
))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let err = exchange_copilot_token(&mock_cfg(host, host), "gho_github_tok")
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
format!("{err}").contains("out-of-range expires_at"),
|
||||
"got {err}"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -130,8 +130,9 @@ fn with_device_headers(
|
||||
}
|
||||
|
||||
/// Defend against control characters / non-https redirects from a
|
||||
/// compromised or mis-configured OAuth host.
|
||||
fn validate_verification_uri(uri: &str) -> anyhow::Result<()> {
|
||||
/// compromised or mis-configured OAuth host. Shared with the generic
|
||||
/// device-code wire ([`super::oauth_device`]).
|
||||
pub(crate) fn validate_verification_uri(uri: &str) -> anyhow::Result<()> {
|
||||
if uri.chars().any(|c| c.is_ascii_control()) {
|
||||
anyhow::bail!("Server returned invalid verification URI");
|
||||
}
|
||||
|
||||
@@ -12,6 +12,7 @@ use tokio_util::sync::CancellationToken;
|
||||
|
||||
#[path = "manager/lock.rs"]
|
||||
mod lock;
|
||||
pub(crate) use lock::try_lock_auth_file_nonblocking;
|
||||
#[path = "manager/sleep_gate.rs"]
|
||||
mod sleep_gate;
|
||||
|
||||
@@ -384,6 +385,59 @@ impl AuthManager {
|
||||
)
|
||||
}
|
||||
|
||||
/// Build a manager for a GENERIC device-code OAuth provider (xai-grok),
|
||||
/// scoped to `oauth.scope_key`. Unlike [`Self::new`] this path is
|
||||
/// file-store only (no keyring — that is gated to the default Kimi install)
|
||||
/// and ignores the Kimi-specific `KIGI_AUTH` inline-credential env; it
|
||||
/// otherwise shares the same multi-scope `auth.json` (honoring
|
||||
/// `KIGI_AUTH_PATH`). The refresher is selected from the scope by
|
||||
/// [`super::refresh::build_refresher`].
|
||||
pub(crate) fn new_oauth_provider(
|
||||
kigi_home: &Path,
|
||||
oauth: &'static kigi_models::OAuthConfig,
|
||||
) -> Self {
|
||||
let scope = oauth.scope_key.to_owned();
|
||||
let path = std::env::var("KIGI_AUTH_PATH")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|_| kigi_home.join("auth.json"));
|
||||
|
||||
let (auth, disk_state) = match read_auth_json(&path) {
|
||||
Ok(map) => {
|
||||
let found = lookup_auth(&map, &scope);
|
||||
let state = if found.is_some() {
|
||||
DiskAuthState::Ok
|
||||
} else {
|
||||
DiskAuthState::EntryMissing
|
||||
};
|
||||
(found, state)
|
||||
}
|
||||
Err(e) => {
|
||||
let state = if e.kind() == std::io::ErrorKind::NotFound {
|
||||
DiskAuthState::FileMissing
|
||||
} else {
|
||||
DiskAuthState::Unreadable
|
||||
};
|
||||
(None, state)
|
||||
}
|
||||
};
|
||||
kigi_log::unified_log::info(
|
||||
"AuthManager::new_oauth_provider",
|
||||
None,
|
||||
Some(serde_json::json!({
|
||||
"scope": &scope,
|
||||
"found": auth.is_some(),
|
||||
"is_expired": auth.as_ref().map(is_expired),
|
||||
})),
|
||||
);
|
||||
Self::assemble(
|
||||
auth,
|
||||
path,
|
||||
scope,
|
||||
KimiCodeConfig::default(),
|
||||
Some(disk_state),
|
||||
)
|
||||
}
|
||||
|
||||
/// Single field-assembly point for [`Self::new`]'s two construction paths
|
||||
/// (inline `KIGI_AUTH` vs. on-disk `auth.json`), which differ only in the
|
||||
/// threaded fields. One literal means a newly added field can't be silently
|
||||
@@ -788,6 +842,13 @@ impl AuthManager {
|
||||
&self.kimi_code_config
|
||||
}
|
||||
|
||||
/// The auth.json / keyring scope key this manager persists under
|
||||
/// (`oauth/kimi-code` for Kimi, `oauth/xai` for xai-grok, …). Drives the
|
||||
/// refresher selection in [`super::refresh::build_refresher`].
|
||||
pub(crate) fn scope(&self) -> &str {
|
||||
&self.scope
|
||||
}
|
||||
|
||||
/// Handle notified after every successful token refresh.
|
||||
///
|
||||
/// Used by [`ModelsManager`] to trigger model catalog recovery
|
||||
|
||||
@@ -1,13 +1,18 @@
|
||||
pub(crate) mod attribution;
|
||||
mod config;
|
||||
pub(crate) mod credential_authority;
|
||||
pub mod credential_provider;
|
||||
pub(crate) mod device;
|
||||
pub mod device_code;
|
||||
pub mod error;
|
||||
mod flow;
|
||||
pub(crate) mod github_copilot;
|
||||
pub(crate) mod kimi_oauth;
|
||||
pub(crate) mod manager;
|
||||
mod model;
|
||||
pub(crate) mod oauth_device;
|
||||
pub(crate) mod oauth_pkce;
|
||||
pub(crate) mod oauth_registry;
|
||||
pub(crate) mod recovery;
|
||||
pub(crate) mod refresh;
|
||||
mod storage;
|
||||
@@ -17,7 +22,8 @@ pub(crate) use flow::try_ensure_session_noninteractive;
|
||||
pub use flow::{
|
||||
AuthChannels, AuthUrlInfo, AuthUrlMode, LogoutResult, ensure_authenticated,
|
||||
ensure_authenticated_or_noninteractive, perform_logout, run_auth_flow,
|
||||
run_auth_flow_with_stderr_bridge, run_cli_login, run_cli_logout, try_ensure_fresh_auth,
|
||||
run_auth_flow_with_stderr_bridge, run_cli_login, run_cli_logout, run_oauth_provider_flow,
|
||||
try_ensure_fresh_auth,
|
||||
};
|
||||
mod meta;
|
||||
pub use device::device_headers;
|
||||
@@ -27,5 +33,6 @@ pub use meta::AuthMeta;
|
||||
pub use model::{AuthMode, KimiAuth, lookup_auth};
|
||||
pub(crate) use model::{TOKEN_TTL, is_expired, token_suffix};
|
||||
pub use storage::{
|
||||
clear_api_key, read_api_key, read_auth_json, read_token_by_scope, store_api_key,
|
||||
clear_api_key, read_api_key, read_auth_json, read_platform_api_key, read_token_by_scope,
|
||||
store_api_key, store_platform_api_key,
|
||||
};
|
||||
|
||||
@@ -0,0 +1,423 @@
|
||||
//! Generic RFC-8628 device-code OAuth wire, driven by a registry
|
||||
//! [`kigi_models::OAuthConfig`] (xai-grok today; Copilot/Claude later).
|
||||
//!
|
||||
//! Three `application/x-www-form-urlencoded` POSTs against `{auth_host}`:
|
||||
//!
|
||||
//! - `POST {device_path}` — form `client_id` + `scope` + the optional
|
||||
//! `extra_device_field` (e.g. `referrer=kigi`)
|
||||
//! - `POST {token_path}` (poll) — form `client_id` + `device_code` +
|
||||
//! `grant_type=urn:ietf:params:oauth:grant-type:device_code`
|
||||
//! - `POST {token_path}` (refresh) — form `client_id` +
|
||||
//! `grant_type=refresh_token` + `refresh_token`, with the same exponential
|
||||
//! backoff / status handling as the Kimi wire.
|
||||
//!
|
||||
//! Unlike [`super::kimi_oauth`] this sends NO X-Msh device headers — just the
|
||||
//! shared kigi `User-Agent` and `Accept: application/json`. Access/refresh
|
||||
//! tokens are NEVER logged (only non-secret events: requested, poll succeeded,
|
||||
//! refreshed).
|
||||
|
||||
use kigi_models::OAuthConfig;
|
||||
use serde::Deserialize;
|
||||
|
||||
use super::kimi_oauth::{
|
||||
DeviceAuthorization, DevicePollResult, RefreshError, TokenResponse, validate_verification_uri,
|
||||
};
|
||||
|
||||
const DEVICE_GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type:device_code";
|
||||
const REFRESH_GRANT_TYPE: &str = "refresh_token";
|
||||
/// Refresh retry budget over the retryable statuses / network blips.
|
||||
const MAX_REFRESH_RETRIES: u32 = 3;
|
||||
/// HTTP statuses worth retrying a refresh for (kimi-cli parity).
|
||||
const RETRYABLE_REFRESH_STATUSES: [u16; 5] = [429, 500, 502, 503, 504];
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct DeviceAuthorizationResponse {
|
||||
user_code: String,
|
||||
device_code: String,
|
||||
#[serde(default)]
|
||||
verification_uri: Option<String>,
|
||||
/// Optional here (the Kimi wire requires it): Pi's xAI response may omit
|
||||
/// `verification_uri_complete` and carry only `verification_uri`.
|
||||
#[serde(default)]
|
||||
verification_uri_complete: Option<String>,
|
||||
#[serde(default)]
|
||||
expires_in: Option<i64>,
|
||||
#[serde(default)]
|
||||
interval: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Default)]
|
||||
struct OAuthErrorBody {
|
||||
#[serde(default)]
|
||||
error: Option<String>,
|
||||
#[serde(default)]
|
||||
error_description: Option<String>,
|
||||
}
|
||||
|
||||
fn oauth_url(host: &str, path: &str) -> String {
|
||||
format!("{}{path}", host.trim_end_matches('/'))
|
||||
}
|
||||
|
||||
/// The device-authorization form fields: `client_id`, `scope`, and the
|
||||
/// optional non-standard `extra_device_field`.
|
||||
fn device_form(cfg: &OAuthConfig) -> Vec<(&'static str, &'static str)> {
|
||||
let mut form = vec![("client_id", cfg.client_id), ("scope", cfg.scope)];
|
||||
if let Some((name, value)) = cfg.extra_device_field {
|
||||
form.push((name, value));
|
||||
}
|
||||
form
|
||||
}
|
||||
|
||||
/// `POST {auth_host}{device_path}` — start a device login.
|
||||
pub(crate) async fn request_device_authorization(
|
||||
cfg: &OAuthConfig,
|
||||
) -> anyhow::Result<DeviceAuthorization> {
|
||||
let url = oauth_url(cfg.auth_host, cfg.device_path);
|
||||
tracing::info!(url = %url, "auth: requesting device authorization (generic oauth)");
|
||||
let resp = crate::http::shared_client()
|
||||
.post(&url)
|
||||
.header("Accept", "application/json")
|
||||
.form(&device_form(cfg))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
tracing::warn!(%status, "auth: device authorization failed (generic oauth)");
|
||||
anyhow::bail!("Device authorization failed (HTTP {status}): {body}");
|
||||
}
|
||||
let parsed: DeviceAuthorizationResponse = resp.json().await?;
|
||||
|
||||
if !parsed
|
||||
.user_code
|
||||
.chars()
|
||||
.all(|c| c.is_ascii_alphanumeric() || c == '-')
|
||||
{
|
||||
anyhow::bail!("Server returned invalid user_code format (expected [A-Z0-9-])");
|
||||
}
|
||||
// Pi forces the displayed URI to https; we require a valid https (or
|
||||
// localhost) verification target, preferring the pre-filled complete form.
|
||||
let verification_uri_complete = parsed
|
||||
.verification_uri_complete
|
||||
.clone()
|
||||
.or_else(|| parsed.verification_uri.clone())
|
||||
.ok_or_else(|| anyhow::anyhow!("Server returned no verification URI"))?;
|
||||
validate_verification_uri(&verification_uri_complete)?;
|
||||
if let Some(ref uri) = parsed.verification_uri {
|
||||
validate_verification_uri(uri)?;
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
user_code = %parsed.user_code,
|
||||
interval = parsed.interval.unwrap_or(5),
|
||||
expires_in = ?parsed.expires_in,
|
||||
"auth: device authorization issued (generic oauth)"
|
||||
);
|
||||
Ok(DeviceAuthorization {
|
||||
user_code: parsed.user_code,
|
||||
device_code: parsed.device_code,
|
||||
verification_uri: parsed.verification_uri.filter(|u| !u.is_empty()),
|
||||
verification_uri_complete,
|
||||
expires_in: parsed.expires_in.filter(|&e| e > 0),
|
||||
interval: parsed.interval.unwrap_or(5),
|
||||
})
|
||||
}
|
||||
|
||||
/// One poll of `POST {auth_host}{token_path}` with the device grant.
|
||||
pub(crate) async fn poll_device_token(
|
||||
cfg: &OAuthConfig,
|
||||
device_code: &str,
|
||||
) -> anyhow::Result<DevicePollResult> {
|
||||
let url = oauth_url(cfg.auth_host, cfg.token_path);
|
||||
let resp = crate::http::shared_client()
|
||||
.post(&url)
|
||||
.header("Accept", "application/json")
|
||||
.form(&[
|
||||
("client_id", cfg.client_id),
|
||||
("device_code", device_code),
|
||||
("grant_type", DEVICE_GRANT_TYPE),
|
||||
])
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Token polling request failed: {e}"))?;
|
||||
|
||||
let status = resp.status();
|
||||
if status.is_server_error() {
|
||||
anyhow::bail!("Token polling server error: {status}");
|
||||
}
|
||||
let body = resp.bytes().await?;
|
||||
if status.is_success() {
|
||||
if let Ok(tokens) = serde_json::from_slice::<TokenResponse>(&body) {
|
||||
tracing::info!("auth: device poll succeeded, access token issued (generic oauth)");
|
||||
return Ok(DevicePollResult::Success(Box::new(tokens.into_auth())));
|
||||
}
|
||||
tracing::warn!(
|
||||
"auth: device poll returned 200 without access_token; continuing (generic oauth)"
|
||||
);
|
||||
return Ok(DevicePollResult::Pending {
|
||||
error: "missing_access_token".to_owned(),
|
||||
description: None,
|
||||
});
|
||||
}
|
||||
let err: OAuthErrorBody = serde_json::from_slice(&body).unwrap_or_default();
|
||||
let error = err.error.unwrap_or_else(|| "unknown_error".to_owned());
|
||||
if error == "expired_token" {
|
||||
tracing::info!(
|
||||
"auth: device code expired; restarting device authorization (generic oauth)"
|
||||
);
|
||||
return Ok(DevicePollResult::Expired);
|
||||
}
|
||||
tracing::debug!(error = %error, "auth: device poll pending (generic oauth)");
|
||||
Ok(DevicePollResult::Pending {
|
||||
error,
|
||||
description: err.error_description,
|
||||
})
|
||||
}
|
||||
|
||||
/// `POST {auth_host}{token_path}` with `grant_type=refresh_token`. Retries the
|
||||
/// retryable statuses / network errors with exponential backoff; 401/403
|
||||
/// returns immediately as [`RefreshError::Unauthorized`].
|
||||
pub(crate) async fn refresh_token(
|
||||
cfg: &OAuthConfig,
|
||||
refresh_token: &str,
|
||||
) -> Result<super::model::KimiAuth, RefreshError> {
|
||||
let url = oauth_url(cfg.auth_host, cfg.token_path);
|
||||
let mut last_error = String::from("no attempt made");
|
||||
for attempt in 0..MAX_REFRESH_RETRIES {
|
||||
if attempt > 0 {
|
||||
let backoff = std::time::Duration::from_secs(1 << (attempt - 1));
|
||||
tracing::warn!(
|
||||
attempt,
|
||||
backoff_secs = backoff.as_secs(),
|
||||
last_error = %last_error,
|
||||
"auth: retrying token refresh (generic oauth)"
|
||||
);
|
||||
tokio::time::sleep(backoff).await;
|
||||
}
|
||||
tracing::info!(attempt, "auth: token refresh attempt (generic oauth)");
|
||||
let send_result = crate::http::shared_client()
|
||||
.post(&url)
|
||||
.header("Accept", "application/json")
|
||||
.form(&[
|
||||
("client_id", cfg.client_id),
|
||||
("grant_type", REFRESH_GRANT_TYPE),
|
||||
("refresh_token", refresh_token),
|
||||
])
|
||||
.send()
|
||||
.await;
|
||||
|
||||
let resp = match send_result {
|
||||
Ok(resp) => resp,
|
||||
Err(e) => {
|
||||
last_error = format!("network error: {e}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let status = resp.status().as_u16();
|
||||
let body = resp.bytes().await.unwrap_or_default();
|
||||
if status == 401 || status == 403 {
|
||||
let err: OAuthErrorBody = serde_json::from_slice(&body).unwrap_or_default();
|
||||
return Err(RefreshError::Unauthorized {
|
||||
status,
|
||||
description: err
|
||||
.error_description
|
||||
.unwrap_or_else(|| "Token refresh unauthorized.".to_owned()),
|
||||
});
|
||||
}
|
||||
if status == 200 {
|
||||
return match serde_json::from_slice::<TokenResponse>(&body) {
|
||||
Ok(tokens) => Ok(tokens.into_auth()),
|
||||
Err(e) => Err(RefreshError::Fatal {
|
||||
status,
|
||||
description: format!("malformed token payload: {e}"),
|
||||
}),
|
||||
};
|
||||
}
|
||||
let err: OAuthErrorBody = serde_json::from_slice(&body).unwrap_or_default();
|
||||
let description = err
|
||||
.error_description
|
||||
.unwrap_or_else(|| format!("Token refresh failed (HTTP {status})."));
|
||||
if RETRYABLE_REFRESH_STATUSES.contains(&status) {
|
||||
last_error = description;
|
||||
continue;
|
||||
}
|
||||
return Err(RefreshError::Fatal {
|
||||
status,
|
||||
description,
|
||||
});
|
||||
}
|
||||
Err(RefreshError::Exhausted { last_error })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use kigi_models::XAI_OAUTH_CONFIG;
|
||||
use wiremock::matchers::{body_string_contains, method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
/// An OAuthConfig pointed at a mock server (copies XAI's client_id/scope/
|
||||
/// paths but overrides the host).
|
||||
fn mock_cfg(host: &'static str) -> OAuthConfig {
|
||||
OAuthConfig {
|
||||
auth_host: host,
|
||||
..XAI_OAUTH_CONFIG
|
||||
}
|
||||
}
|
||||
|
||||
fn token_json(access: &str, refresh: &str) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"access_token": access,
|
||||
"refresh_token": refresh,
|
||||
"expires_in": 3600,
|
||||
"scope": "grok-cli:access",
|
||||
"token_type": "bearer",
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn device_authorization_sends_client_scope_and_referrer() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth2/device/code"))
|
||||
.and(body_string_contains(
|
||||
"client_id=b1a00492-073a-47ea-816f-4c329264a828",
|
||||
))
|
||||
.and(body_string_contains("scope=openid"))
|
||||
.and(body_string_contains("referrer=kigi"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"user_code": "GROK-1234",
|
||||
"device_code": "dev-xai-1",
|
||||
"verification_uri": "https://x.ai/device",
|
||||
"verification_uri_complete": "https://x.ai/device?user_code=GROK-1234",
|
||||
"expires_in": 900,
|
||||
"interval": 5,
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let auth = request_device_authorization(&mock_cfg(host)).await.unwrap();
|
||||
assert_eq!(auth.user_code, "GROK-1234");
|
||||
assert_eq!(auth.device_code, "dev-xai-1");
|
||||
assert_eq!(
|
||||
auth.verification_uri_complete,
|
||||
"https://x.ai/device?user_code=GROK-1234"
|
||||
);
|
||||
assert_eq!(auth.expires_in, Some(900));
|
||||
}
|
||||
|
||||
/// A response with only `verification_uri` (no `_complete`) still yields a
|
||||
/// valid display URI.
|
||||
#[tokio::test]
|
||||
async fn device_authorization_falls_back_to_verification_uri() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth2/device/code"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"user_code": "GROK-9",
|
||||
"device_code": "d",
|
||||
"verification_uri": "https://x.ai/device",
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let auth = request_device_authorization(&mock_cfg(host)).await.unwrap();
|
||||
assert_eq!(auth.verification_uri_complete, "https://x.ai/device");
|
||||
assert_eq!(auth.interval, 5, "default interval");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn poll_success_builds_auth_with_expiry() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth2/token"))
|
||||
.and(body_string_contains("grant_type=urn"))
|
||||
.and(body_string_contains("device_code=dev-xai-1"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_json(token_json("grok-at", "grok-rt")),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let result = poll_device_token(&mock_cfg(host), "dev-xai-1")
|
||||
.await
|
||||
.unwrap();
|
||||
let DevicePollResult::Success(auth) = result else {
|
||||
panic!("expected success, got {result:?}");
|
||||
};
|
||||
assert_eq!(auth.key, "grok-at");
|
||||
assert_eq!(auth.refresh_token.as_deref(), Some("grok-rt"));
|
||||
assert_eq!(auth.expires_in, Some(3600));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn poll_maps_authorization_pending_to_pending() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth2/token"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(400)
|
||||
.set_body_json(serde_json::json!({ "error": "authorization_pending" })),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let result = poll_device_token(&mock_cfg(host), "dev-xai-1")
|
||||
.await
|
||||
.unwrap();
|
||||
match result {
|
||||
DevicePollResult::Pending { error, .. } => assert_eq!(error, "authorization_pending"),
|
||||
other => panic!("expected pending, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_success_round_trip() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth2/token"))
|
||||
.and(body_string_contains("grant_type=refresh_token"))
|
||||
.and(body_string_contains("refresh_token=grok-rt-old"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_json(token_json("grok-at-new", "grok-rt-new")),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let auth = refresh_token(&mock_cfg(host), "grok-rt-old").await.unwrap();
|
||||
assert_eq!(auth.key, "grok-at-new");
|
||||
assert_eq!(auth.refresh_token.as_deref(), Some("grok-rt-new"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_401_maps_to_unauthorized() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth2/token"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(401)
|
||||
.set_body_json(serde_json::json!({ "error_description": "refresh revoked" })),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let err = refresh_token(&mock_cfg(host), "grok-rt-dead")
|
||||
.await
|
||||
.unwrap_err();
|
||||
match err {
|
||||
RefreshError::Unauthorized {
|
||||
status,
|
||||
description,
|
||||
} => {
|
||||
assert_eq!(status, 401);
|
||||
assert_eq!(description, "refresh revoked");
|
||||
}
|
||||
other => panic!("expected Unauthorized, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,904 @@
|
||||
//! Generic authorization-code + PKCE (S256) OAuth wire with a `127.0.0.1`
|
||||
//! loopback callback, driven by a registry [`kigi_models::OAuthConfig`] whose
|
||||
//! `flow` is [`OAuthFlow::PkceLocalhost`] (claude-pro-max JSON + openai-codex
|
||||
//! FORM).
|
||||
//!
|
||||
//! Shape (Pi `earendil-works/pi` `auth/oauth/{anthropic,openai-codex}.ts`):
|
||||
//! - `verifier = base64url(32 random bytes)`; `challenge = base64url(SHA-256(
|
||||
//! verifier))`. `state` is `verifier` for claude ([`generate_pkce`]) or a
|
||||
//! fresh-random value for codex ([`generate_pkce_random_state`]).
|
||||
//! - Browser opens `{auth_host}{device_path}?client_id&response_type=code&
|
||||
//! scope&redirect_uri&state&code_challenge&code_challenge_method=S256` plus any
|
||||
//! `authorize_extra` params (codex only).
|
||||
//! - The code returns to a loopback listener on `127.0.0.1:{redirect_port}`
|
||||
//! answering ONLY `{redirect_path}` (claude `/callback`, codex
|
||||
//! `/auth/callback`), with STRICT `state` validation (a mismatch is rejected —
|
||||
//! CSRF guard). A manual paste (redirect URL / `code#state` / bare code) is
|
||||
//! accepted as a headless fallback.
|
||||
//! - Code → token exchange POSTs `{token_host}{token_path}` as JSON
|
||||
//! ([`exchange_code`], claude) or FORM ([`exchange_code_form`], codex; NO
|
||||
//! `state` field). Refresh: claude JSON here ([`refresh_token`], rotating);
|
||||
//! codex takes the generic device refresher's FORM path.
|
||||
//!
|
||||
//! SECURITY: the verifier, authorization code, access token, and refresh token
|
||||
//! are NEVER logged (only non-secret events: authorize URL requested, callback
|
||||
//! received, token issued, token refreshed).
|
||||
|
||||
use anyhow::Context;
|
||||
use base64::Engine;
|
||||
use kigi_models::{OAuthConfig, OAuthTokenBody};
|
||||
use serde::Deserialize;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use super::kimi_oauth::{RefreshError, TokenResponse};
|
||||
use super::model::KimiAuth;
|
||||
|
||||
const CODE_GRANT_TYPE: &str = "authorization_code";
|
||||
const REFRESH_GRANT_TYPE: &str = "refresh_token";
|
||||
/// Refresh retry budget over the retryable statuses / network blips.
|
||||
const MAX_REFRESH_RETRIES: u32 = 3;
|
||||
/// HTTP statuses worth retrying a refresh for (parity with the device wire).
|
||||
const RETRYABLE_REFRESH_STATUSES: [u16; 5] = [429, 500, 502, 503, 504];
|
||||
|
||||
/// PKCE secrets for one login attempt. Two `state` conventions ship:
|
||||
/// [`generate_pkce`] sets `state == verifier` (Pi's/Claude's convention), while
|
||||
/// [`generate_pkce_random_state`] mints an INDEPENDENT random state (the OAuth
|
||||
/// standard, used by ChatGPT/Codex). Either way the state is validated on the
|
||||
/// callback as the CSRF guard.
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct PkceCodes {
|
||||
/// `code_verifier` — the 43-char base64url secret, sent at token exchange.
|
||||
pub verifier: String,
|
||||
/// `code_challenge = base64url(SHA-256(verifier))`, sent at authorize.
|
||||
pub challenge: String,
|
||||
/// `state` — the verifier itself, or an independent random value depending
|
||||
/// on the provider's dialect; validated on the callback (CSRF guard).
|
||||
pub state: String,
|
||||
}
|
||||
|
||||
/// Generate PKCE S256 codes: `verifier = base64url(32 random bytes)`,
|
||||
/// `challenge = base64url(SHA-256(verifier))`, `state = verifier`.
|
||||
pub(crate) fn generate_pkce() -> PkceCodes {
|
||||
use rand::RngCore;
|
||||
let mut raw = [0u8; 32];
|
||||
rand::rng().fill_bytes(&mut raw);
|
||||
let verifier = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(raw);
|
||||
let digest = Sha256::digest(verifier.as_bytes());
|
||||
let challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(digest);
|
||||
PkceCodes {
|
||||
state: verifier.clone(),
|
||||
verifier,
|
||||
challenge,
|
||||
}
|
||||
}
|
||||
|
||||
/// Like [`generate_pkce`] but with an INDEPENDENT fresh-random `state` (16
|
||||
/// random bytes) instead of `state == verifier`. The ChatGPT/Codex flow uses a
|
||||
/// distinct state (the verifier never doubles as the CSRF token there), so the
|
||||
/// verifier stays out of the state carried on the loopback callback.
|
||||
pub(crate) fn generate_pkce_random_state() -> PkceCodes {
|
||||
use rand::RngCore;
|
||||
let mut raw = [0u8; 16];
|
||||
rand::rng().fill_bytes(&mut raw);
|
||||
let state = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(raw);
|
||||
PkceCodes {
|
||||
state,
|
||||
..generate_pkce()
|
||||
}
|
||||
}
|
||||
|
||||
/// The loopback redirect URI for a PKCE-localhost provider (claude `/callback`,
|
||||
/// codex `/auth/callback`).
|
||||
pub(crate) fn redirect_uri(redirect_port: u16, redirect_path: &str) -> String {
|
||||
format!("http://localhost:{redirect_port}{redirect_path}")
|
||||
}
|
||||
|
||||
/// Build the browser authorize URL:
|
||||
/// `{auth_host}{device_path}?client_id&response_type=code&scope&redirect_uri&
|
||||
/// state&code_challenge&code_challenge_method=S256` plus any config
|
||||
/// `authorize_extra` params (empty for every config but codex, so their URLs
|
||||
/// stay byte-identical).
|
||||
pub(crate) fn build_authorize_url(
|
||||
cfg: &OAuthConfig,
|
||||
redirect_uri: &str,
|
||||
pkce: &PkceCodes,
|
||||
) -> String {
|
||||
let base = format!("{}{}", cfg.auth_host.trim_end_matches('/'), cfg.device_path);
|
||||
let mut serializer = url::form_urlencoded::Serializer::new(String::new());
|
||||
serializer
|
||||
.append_pair("client_id", cfg.client_id)
|
||||
.append_pair("response_type", "code")
|
||||
.append_pair("scope", cfg.scope)
|
||||
.append_pair("redirect_uri", redirect_uri)
|
||||
.append_pair("state", &pkce.state)
|
||||
.append_pair("code_challenge", &pkce.challenge)
|
||||
.append_pair("code_challenge_method", "S256");
|
||||
for (key, value) in cfg.authorize_extra {
|
||||
serializer.append_pair(key, value);
|
||||
}
|
||||
format!("{base}?{}", serializer.finish())
|
||||
}
|
||||
|
||||
/// `code` + `state` extracted from a callback (loopback query OR manual paste).
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct CallbackParams {
|
||||
pub code: String,
|
||||
pub state: Option<String>,
|
||||
}
|
||||
|
||||
/// Parse `code`/`state` from the raw query string of a `/callback?…` request
|
||||
/// (e.g. `code=abc&state=xyz`). An `error=` param surfaces as an `Err`.
|
||||
pub(crate) fn parse_callback_query(query: &str) -> anyhow::Result<CallbackParams> {
|
||||
let mut code = None;
|
||||
let mut state = None;
|
||||
let mut error = None;
|
||||
for (k, v) in url::form_urlencoded::parse(query.as_bytes()) {
|
||||
match k.as_ref() {
|
||||
"code" => code = Some(v.into_owned()),
|
||||
"state" => state = Some(v.into_owned()),
|
||||
"error" => error = Some(v.into_owned()),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
if let Some(error) = error {
|
||||
anyhow::bail!("Authorization server returned an error: {error}");
|
||||
}
|
||||
let code = code.context("callback missing authorization code")?;
|
||||
if code.is_empty() {
|
||||
anyhow::bail!("callback authorization code was empty");
|
||||
}
|
||||
Ok(CallbackParams { code, state })
|
||||
}
|
||||
|
||||
/// Parse a MANUAL paste (headless fallback). Accepts, in order:
|
||||
/// - a full redirect URL (`http://localhost:…/callback?code=…&state=…`),
|
||||
/// - a `code#state` pair (Anthropic's console shows this form),
|
||||
/// - a bare `code` (state then unknown → `None`, caller validation applies).
|
||||
pub(crate) fn parse_manual_paste(input: &str) -> anyhow::Result<CallbackParams> {
|
||||
let trimmed = input.trim();
|
||||
if trimmed.is_empty() {
|
||||
anyhow::bail!("empty paste");
|
||||
}
|
||||
// Full redirect URL.
|
||||
if trimmed.starts_with("http://") || trimmed.starts_with("https://") {
|
||||
let url = url::Url::parse(trimmed).context("pasted value is not a valid URL")?;
|
||||
return parse_callback_query(url.query().unwrap_or_default());
|
||||
}
|
||||
// `code#state`.
|
||||
if let Some((code, state)) = trimmed.split_once('#') {
|
||||
if code.is_empty() {
|
||||
anyhow::bail!("pasted code was empty");
|
||||
}
|
||||
return Ok(CallbackParams {
|
||||
code: code.to_owned(),
|
||||
state: (!state.is_empty()).then(|| state.to_owned()),
|
||||
});
|
||||
}
|
||||
// Bare code.
|
||||
Ok(CallbackParams {
|
||||
code: trimmed.to_owned(),
|
||||
state: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// STRICT state validation (CSRF guard): the callback `state` MUST be present
|
||||
/// AND equal to the expected value. A mismatch (or absence, when a paste has no
|
||||
/// state) is rejected — the flow NEVER proceeds on an unverified callback.
|
||||
pub(crate) fn validate_state(params: &CallbackParams, expected_state: &str) -> anyhow::Result<()> {
|
||||
match params.state.as_deref() {
|
||||
Some(state) if state == expected_state => Ok(()),
|
||||
Some(_) => anyhow::bail!("OAuth state mismatch — rejecting callback (CSRF guard)"),
|
||||
None => anyhow::bail!("OAuth callback carried no state — rejecting (CSRF guard)"),
|
||||
}
|
||||
}
|
||||
|
||||
/// State validation for a MANUAL paste (headless fallback): a present state
|
||||
/// MUST match (mismatch rejected — CSRF guard), but an ABSENT state is allowed
|
||||
/// — a bare-code paste is user-initiated (not a network-reachable callback), so
|
||||
/// there is no state to check. The loopback path uses the stricter
|
||||
/// [`validate_state`] (an absent state there IS rejected).
|
||||
pub(crate) fn validate_pasted_state(
|
||||
params: &CallbackParams,
|
||||
expected_state: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
match params.state.as_deref() {
|
||||
Some(state) if state == expected_state => Ok(()),
|
||||
Some(_) => anyhow::bail!("OAuth state mismatch — rejecting pasted code (CSRF guard)"),
|
||||
None => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
/// The token-endpoint URL (`{token_host}{token_path}`).
|
||||
fn token_url(cfg: &OAuthConfig) -> String {
|
||||
format!("{}{}", cfg.token_host.trim_end_matches('/'), cfg.token_path)
|
||||
}
|
||||
|
||||
/// POST the token endpoint with a JSON body, honoring `cfg.token_body`. Claude
|
||||
/// is JSON; a `Form`-bodied config would be handled by the device wire, so the
|
||||
/// PKCE path asserts JSON (never silently mis-encodes).
|
||||
async fn post_token_json(
|
||||
cfg: &OAuthConfig,
|
||||
body: serde_json::Value,
|
||||
) -> reqwest::Result<reqwest::Response> {
|
||||
debug_assert!(
|
||||
matches!(cfg.token_body, OAuthTokenBody::Json),
|
||||
"PKCE token exchange expects a JSON token body"
|
||||
);
|
||||
crate::http::shared_client()
|
||||
.post(token_url(cfg))
|
||||
.header("Accept", "application/json")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
}
|
||||
|
||||
/// Exchange an authorization `code` for a token set (JSON body):
|
||||
/// `{grant_type:"authorization_code", code, state, client_id, redirect_uri,
|
||||
/// code_verifier}`. Returns the materialized [`KimiAuth`].
|
||||
pub(crate) async fn exchange_code(
|
||||
cfg: &OAuthConfig,
|
||||
code: &str,
|
||||
pkce: &PkceCodes,
|
||||
redirect_uri: &str,
|
||||
) -> anyhow::Result<KimiAuth> {
|
||||
let body = serde_json::json!({
|
||||
"grant_type": CODE_GRANT_TYPE,
|
||||
"code": code,
|
||||
"state": pkce.state,
|
||||
"client_id": cfg.client_id,
|
||||
"redirect_uri": redirect_uri,
|
||||
"code_verifier": pkce.verifier,
|
||||
});
|
||||
tracing::info!(
|
||||
scope_key = cfg.scope_key,
|
||||
"auth: exchanging code for token (pkce)"
|
||||
);
|
||||
let resp = post_token_json(cfg, body)
|
||||
.await
|
||||
.context("token exchange request failed")?;
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
tracing::warn!(%status, scope_key = cfg.scope_key, "auth: code exchange failed (pkce)");
|
||||
anyhow::bail!("Token exchange failed (HTTP {status}): {body}");
|
||||
}
|
||||
let tokens: TokenResponse = resp.json().await.context("malformed token payload")?;
|
||||
tracing::info!(
|
||||
scope_key = cfg.scope_key,
|
||||
"auth: pkce code exchange succeeded"
|
||||
);
|
||||
Ok(tokens.into_auth())
|
||||
}
|
||||
|
||||
/// Exchange an authorization `code` for a token set with a FORM body (codex):
|
||||
/// `{grant_type=authorization_code, client_id, code, code_verifier,
|
||||
/// redirect_uri}`. Unlike [`exchange_code`], the `state` is NOT sent in the
|
||||
/// token body (the ChatGPT/Codex token endpoint does not expect it). Asserts a
|
||||
/// `Form` config so a JSON provider can never silently mis-encode.
|
||||
pub(crate) async fn exchange_code_form(
|
||||
cfg: &OAuthConfig,
|
||||
code: &str,
|
||||
pkce: &PkceCodes,
|
||||
redirect_uri: &str,
|
||||
) -> anyhow::Result<KimiAuth> {
|
||||
debug_assert!(
|
||||
matches!(cfg.token_body, OAuthTokenBody::Form),
|
||||
"PKCE form exchange expects a Form token body"
|
||||
);
|
||||
tracing::info!(
|
||||
scope_key = cfg.scope_key,
|
||||
"auth: exchanging code for token (pkce form)"
|
||||
);
|
||||
let resp = crate::http::shared_client()
|
||||
.post(token_url(cfg))
|
||||
.header("Accept", "application/json")
|
||||
.form(&[
|
||||
("grant_type", CODE_GRANT_TYPE),
|
||||
("client_id", cfg.client_id),
|
||||
("code", code),
|
||||
("code_verifier", pkce.verifier.as_str()),
|
||||
("redirect_uri", redirect_uri),
|
||||
])
|
||||
.send()
|
||||
.await
|
||||
.context("token exchange request failed")?;
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
tracing::warn!(%status, scope_key = cfg.scope_key, "auth: code exchange failed (pkce form)");
|
||||
anyhow::bail!("Token exchange failed (HTTP {status}): {body}");
|
||||
}
|
||||
let tokens: TokenResponse = resp.json().await.context("malformed token payload")?;
|
||||
tracing::info!(
|
||||
scope_key = cfg.scope_key,
|
||||
"auth: pkce form code exchange succeeded"
|
||||
);
|
||||
Ok(tokens.into_auth())
|
||||
}
|
||||
|
||||
/// `POST {token_host}{token_path}` with `grant_type=refresh_token` (JSON body).
|
||||
/// Claude ROTATES the refresh token, so the caller MUST persist the returned
|
||||
/// one. Retries the retryable statuses / network errors with exponential
|
||||
/// backoff; 401/403 returns immediately as [`RefreshError::Unauthorized`].
|
||||
pub(crate) async fn refresh_token(
|
||||
cfg: &OAuthConfig,
|
||||
refresh_token: &str,
|
||||
) -> Result<KimiAuth, RefreshError> {
|
||||
let mut last_error = String::from("no attempt made");
|
||||
for attempt in 0..MAX_REFRESH_RETRIES {
|
||||
if attempt > 0 {
|
||||
let backoff = std::time::Duration::from_secs(1 << (attempt - 1));
|
||||
tracing::warn!(
|
||||
attempt,
|
||||
backoff_secs = backoff.as_secs(),
|
||||
last_error = %last_error,
|
||||
"auth: retrying token refresh (pkce)"
|
||||
);
|
||||
tokio::time::sleep(backoff).await;
|
||||
}
|
||||
tracing::info!(
|
||||
attempt,
|
||||
scope_key = cfg.scope_key,
|
||||
"auth: token refresh attempt (pkce)"
|
||||
);
|
||||
let body = serde_json::json!({
|
||||
"grant_type": REFRESH_GRANT_TYPE,
|
||||
"client_id": cfg.client_id,
|
||||
"refresh_token": refresh_token,
|
||||
});
|
||||
let resp = match post_token_json(cfg, body).await {
|
||||
Ok(resp) => resp,
|
||||
Err(e) => {
|
||||
last_error = format!("network error: {e}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let status = resp.status().as_u16();
|
||||
let bytes = resp.bytes().await.unwrap_or_default();
|
||||
if status == 401 || status == 403 {
|
||||
let err: OAuthErrorBody = serde_json::from_slice(&bytes).unwrap_or_default();
|
||||
return Err(RefreshError::Unauthorized {
|
||||
status,
|
||||
description: err
|
||||
.error_description
|
||||
.unwrap_or_else(|| "Token refresh unauthorized.".to_owned()),
|
||||
});
|
||||
}
|
||||
if status == 200 {
|
||||
return match serde_json::from_slice::<TokenResponse>(&bytes) {
|
||||
Ok(tokens) => Ok(tokens.into_auth()),
|
||||
Err(e) => Err(RefreshError::Fatal {
|
||||
status,
|
||||
description: format!("malformed token payload: {e}"),
|
||||
}),
|
||||
};
|
||||
}
|
||||
let err: OAuthErrorBody = serde_json::from_slice(&bytes).unwrap_or_default();
|
||||
let description = err
|
||||
.error_description
|
||||
.unwrap_or_else(|| format!("Token refresh failed (HTTP {status})."));
|
||||
if RETRYABLE_REFRESH_STATUSES.contains(&status) {
|
||||
last_error = description;
|
||||
continue;
|
||||
}
|
||||
return Err(RefreshError::Fatal {
|
||||
status,
|
||||
description,
|
||||
});
|
||||
}
|
||||
Err(RefreshError::Exhausted { last_error })
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Default)]
|
||||
struct OAuthErrorBody {
|
||||
#[serde(default)]
|
||||
error_description: Option<String>,
|
||||
}
|
||||
|
||||
/// Bind a loopback HTTP listener on `127.0.0.1:{redirect_port}` and wait for a
|
||||
/// single `GET {redirect_path}?code=…&state=…`, validating `state` STRICTLY
|
||||
/// against `expected_state` (mismatch → rejected). Returns the authorization
|
||||
/// code.
|
||||
///
|
||||
/// The listener answers ONLY `redirect_path` (claude `/callback`, codex
|
||||
/// `/auth/callback`); any other path gets 404. It binds `127.0.0.1` (never
|
||||
/// `0.0.0.0`), so no non-loopback host can reach it.
|
||||
pub(crate) async fn await_loopback_code(
|
||||
redirect_port: u16,
|
||||
redirect_path: &str,
|
||||
expected_state: &str,
|
||||
) -> anyhow::Result<String> {
|
||||
let listener = tokio::net::TcpListener::bind(("127.0.0.1", redirect_port))
|
||||
.await
|
||||
.with_context(|| format!("could not bind loopback 127.0.0.1:{redirect_port}"))?;
|
||||
tracing::info!(port = redirect_port, "auth: pkce loopback listener bound");
|
||||
loop {
|
||||
let (stream, _peer) = listener.accept().await.context("loopback accept failed")?;
|
||||
match handle_loopback_conn(stream, redirect_path, expected_state).await {
|
||||
LoopbackOutcome::Code(code) => return Ok(code),
|
||||
LoopbackOutcome::Rejected(err) => return Err(err),
|
||||
// Not the callback GET (favicon, health probe): keep listening.
|
||||
LoopbackOutcome::Ignore => continue,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum LoopbackOutcome {
|
||||
Code(String),
|
||||
Rejected(anyhow::Error),
|
||||
Ignore,
|
||||
}
|
||||
|
||||
/// Read the request line of one loopback connection, answer with a small HTML
|
||||
/// page, and classify the outcome. STRICT: a `/callback` with a bad/missing
|
||||
/// state is [`LoopbackOutcome::Rejected`] (the browser sees an error page).
|
||||
async fn handle_loopback_conn(
|
||||
mut stream: tokio::net::TcpStream,
|
||||
redirect_path: &str,
|
||||
expected_state: &str,
|
||||
) -> LoopbackOutcome {
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
|
||||
// Read only enough for the request line — a GET has no body.
|
||||
let mut buf = [0u8; 8192];
|
||||
let n = match stream.read(&mut buf).await {
|
||||
Ok(0) => return LoopbackOutcome::Ignore,
|
||||
Ok(n) => n,
|
||||
Err(_) => return LoopbackOutcome::Ignore,
|
||||
};
|
||||
let head = String::from_utf8_lossy(&buf[..n]);
|
||||
let Some(request_line) = head.lines().next() else {
|
||||
return LoopbackOutcome::Ignore;
|
||||
};
|
||||
// `GET {redirect_path}?code=…&state=… HTTP/1.1`
|
||||
let mut parts = request_line.split_whitespace();
|
||||
let (Some(method), Some(target)) = (parts.next(), parts.next()) else {
|
||||
return LoopbackOutcome::Ignore;
|
||||
};
|
||||
if method != "GET" {
|
||||
let _ = write_http(&mut stream, 405, "Method Not Allowed").await;
|
||||
return LoopbackOutcome::Ignore;
|
||||
}
|
||||
let (path, query) = target.split_once('?').unwrap_or((target, ""));
|
||||
if path != redirect_path {
|
||||
let _ = write_http(&mut stream, 404, "Not Found").await;
|
||||
return LoopbackOutcome::Ignore;
|
||||
}
|
||||
|
||||
let result = parse_callback_query(query)
|
||||
.and_then(|params| validate_state(¶ms, expected_state).map(|()| params.code));
|
||||
match result {
|
||||
Ok(code) => {
|
||||
let _ = write_http(
|
||||
&mut stream,
|
||||
200,
|
||||
"Signed in. You can close this window and return to kigi.",
|
||||
)
|
||||
.await;
|
||||
let _ = stream.flush().await;
|
||||
LoopbackOutcome::Code(code)
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = write_http(&mut stream, 400, "Login failed — return to kigi and retry.").await;
|
||||
let _ = stream.flush().await;
|
||||
LoopbackOutcome::Rejected(e)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Write a minimal HTTP/1.1 response with an HTML body.
|
||||
async fn write_http(
|
||||
stream: &mut tokio::net::TcpStream,
|
||||
status: u16,
|
||||
message: &str,
|
||||
) -> std::io::Result<()> {
|
||||
use tokio::io::AsyncWriteExt;
|
||||
let reason = match status {
|
||||
200 => "OK",
|
||||
400 => "Bad Request",
|
||||
404 => "Not Found",
|
||||
405 => "Method Not Allowed",
|
||||
_ => "Error",
|
||||
};
|
||||
let body = format!("<!doctype html><meta charset=utf-8><p>{message}</p>");
|
||||
let response = format!(
|
||||
"HTTP/1.1 {status} {reason}\r\nContent-Type: text/html; charset=utf-8\r\n\
|
||||
Content-Length: {}\r\nConnection: close\r\n\r\n{body}",
|
||||
body.len()
|
||||
);
|
||||
stream.write_all(response.as_bytes()).await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use kigi_models::CLAUDE_OAUTH_CONFIG;
|
||||
use wiremock::matchers::{body_string_contains, method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
/// A config pointed at a mock token host (copies Claude's client_id/scope/
|
||||
/// paths but overrides the token host).
|
||||
fn mock_cfg(token_host: &'static str) -> OAuthConfig {
|
||||
OAuthConfig {
|
||||
token_host,
|
||||
..CLAUDE_OAUTH_CONFIG
|
||||
}
|
||||
}
|
||||
|
||||
/// PKCE codes: verifier/challenge are non-empty base64url (no padding), the
|
||||
/// challenge is the base64url SHA-256 of the verifier, and state == verifier.
|
||||
#[test]
|
||||
fn generate_pkce_produces_valid_s256_codes() {
|
||||
let pkce = generate_pkce();
|
||||
assert_eq!(pkce.state, pkce.verifier, "state must equal the verifier");
|
||||
assert!(!pkce.verifier.is_empty() && !pkce.challenge.is_empty());
|
||||
for s in [&pkce.verifier, &pkce.challenge] {
|
||||
assert!(
|
||||
s.chars()
|
||||
.all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_'),
|
||||
"base64url (no pad) only: {s}"
|
||||
);
|
||||
assert!(!s.contains('='), "no padding: {s}");
|
||||
}
|
||||
// challenge == base64url(SHA-256(verifier)).
|
||||
let expect = base64::engine::general_purpose::URL_SAFE_NO_PAD
|
||||
.encode(Sha256::digest(pkce.verifier.as_bytes()));
|
||||
assert_eq!(pkce.challenge, expect);
|
||||
// Fresh entropy each call.
|
||||
assert_ne!(pkce.verifier, generate_pkce().verifier);
|
||||
}
|
||||
|
||||
/// The authorize URL carries the fixed params + the PKCE state and S256
|
||||
/// challenge, and targets `claude.ai/oauth/authorize`.
|
||||
#[test]
|
||||
fn authorize_url_has_state_and_s256_challenge() {
|
||||
let pkce = generate_pkce();
|
||||
let redirect = redirect_uri(53692, "/callback");
|
||||
let url = build_authorize_url(&CLAUDE_OAUTH_CONFIG, &redirect, &pkce);
|
||||
let parsed = url::Url::parse(&url).expect("valid URL");
|
||||
assert_eq!(parsed.host_str(), Some("claude.ai"));
|
||||
assert_eq!(parsed.path(), "/oauth/authorize");
|
||||
let q: std::collections::HashMap<_, _> = parsed.query_pairs().into_owned().collect();
|
||||
assert_eq!(q.get("response_type").map(String::as_str), Some("code"));
|
||||
assert_eq!(
|
||||
q.get("code_challenge_method").map(String::as_str),
|
||||
Some("S256")
|
||||
);
|
||||
assert_eq!(
|
||||
q.get("state").map(String::as_str),
|
||||
Some(pkce.state.as_str())
|
||||
);
|
||||
assert_eq!(
|
||||
q.get("code_challenge").map(String::as_str),
|
||||
Some(pkce.challenge.as_str())
|
||||
);
|
||||
assert_eq!(
|
||||
q.get("client_id").map(String::as_str),
|
||||
Some(CLAUDE_OAUTH_CONFIG.client_id)
|
||||
);
|
||||
assert_eq!(
|
||||
q.get("redirect_uri").map(String::as_str),
|
||||
Some(redirect.as_str())
|
||||
);
|
||||
// The verifier itself must NEVER appear in the browser URL.
|
||||
assert!(
|
||||
!url.contains("code_verifier"),
|
||||
"the verifier must not ride the authorize URL"
|
||||
);
|
||||
}
|
||||
|
||||
/// STRICT state validation: an exact match passes; a mismatch or an absent
|
||||
/// state is REJECTED (CSRF guard — the flow must never proceed).
|
||||
#[test]
|
||||
fn state_validation_is_strict() {
|
||||
let ok = CallbackParams {
|
||||
code: "c".into(),
|
||||
state: Some("expected".into()),
|
||||
};
|
||||
assert!(validate_state(&ok, "expected").is_ok());
|
||||
let mismatch = CallbackParams {
|
||||
code: "c".into(),
|
||||
state: Some("attacker".into()),
|
||||
};
|
||||
assert!(
|
||||
validate_state(&mismatch, "expected").is_err(),
|
||||
"a state mismatch MUST be rejected"
|
||||
);
|
||||
let missing = CallbackParams {
|
||||
code: "c".into(),
|
||||
state: None,
|
||||
};
|
||||
assert!(
|
||||
validate_state(&missing, "expected").is_err(),
|
||||
"an absent state MUST be rejected"
|
||||
);
|
||||
}
|
||||
|
||||
/// A loopback `/callback` with the WRONG state is rejected end-to-end (the
|
||||
/// listener returns an error, never a code) — the CSRF guard on the wire.
|
||||
#[tokio::test]
|
||||
async fn loopback_rejects_state_mismatch() {
|
||||
// Ephemeral port: bind, learn the port, then drive a client at it.
|
||||
let probe = tokio::net::TcpListener::bind(("127.0.0.1", 0))
|
||||
.await
|
||||
.unwrap();
|
||||
let port = probe.local_addr().unwrap().port();
|
||||
drop(probe);
|
||||
|
||||
let server =
|
||||
tokio::spawn(
|
||||
async move { await_loopback_code(port, "/callback", "the-real-state").await },
|
||||
);
|
||||
// Give the listener a moment to bind.
|
||||
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
||||
// Attacker callback: valid code, WRONG state.
|
||||
let _ = reqwest::get(format!(
|
||||
"http://127.0.0.1:{port}/callback?code=stolen&state=wrong-state"
|
||||
))
|
||||
.await;
|
||||
let outcome = server.await.unwrap();
|
||||
let err = outcome.expect_err("a state mismatch must be rejected, never yield a code");
|
||||
assert!(err.to_string().contains("state mismatch"), "{err}");
|
||||
}
|
||||
|
||||
/// A loopback `/callback` with the MATCHING state yields the code.
|
||||
#[tokio::test]
|
||||
async fn loopback_returns_code_on_valid_state() {
|
||||
let probe = tokio::net::TcpListener::bind(("127.0.0.1", 0))
|
||||
.await
|
||||
.unwrap();
|
||||
let port = probe.local_addr().unwrap().port();
|
||||
drop(probe);
|
||||
|
||||
let server =
|
||||
tokio::spawn(async move { await_loopback_code(port, "/callback", "good-state").await });
|
||||
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
||||
let _ = reqwest::get(format!(
|
||||
"http://127.0.0.1:{port}/callback?code=auth-code-123&state=good-state"
|
||||
))
|
||||
.await;
|
||||
let code = server.await.unwrap().expect("valid state yields the code");
|
||||
assert_eq!(code, "auth-code-123");
|
||||
}
|
||||
|
||||
/// Manual-paste parsing: full redirect URL, `code#state`, and bare code.
|
||||
#[test]
|
||||
fn manual_paste_parses_all_three_forms() {
|
||||
let from_url =
|
||||
parse_manual_paste("http://localhost:53692/callback?code=abc123&state=st-9").unwrap();
|
||||
assert_eq!(from_url.code, "abc123");
|
||||
assert_eq!(from_url.state.as_deref(), Some("st-9"));
|
||||
|
||||
let from_hash = parse_manual_paste("abc123#st-9").unwrap();
|
||||
assert_eq!(from_hash.code, "abc123");
|
||||
assert_eq!(from_hash.state.as_deref(), Some("st-9"));
|
||||
|
||||
let bare = parse_manual_paste(" abc123 ").unwrap();
|
||||
assert_eq!(bare.code, "abc123");
|
||||
assert_eq!(bare.state, None);
|
||||
|
||||
assert!(parse_manual_paste("").is_err());
|
||||
// A pasted redirect that carries an error param surfaces the error.
|
||||
assert!(parse_manual_paste("http://localhost/callback?error=access_denied").is_err());
|
||||
}
|
||||
|
||||
/// Code → token exchange: JSON body carries the grant + verifier, response
|
||||
/// materializes a `KimiAuth` with the rotating refresh token.
|
||||
#[tokio::test]
|
||||
async fn exchange_code_posts_json_and_returns_auth() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/oauth/token"))
|
||||
.and(body_string_contains(
|
||||
"\"grant_type\":\"authorization_code\"",
|
||||
))
|
||||
.and(body_string_contains("\"code\":\"auth-code-xyz\""))
|
||||
.and(body_string_contains("\"code_verifier\""))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"access_token": "sk-ant-oat-new",
|
||||
"refresh_token": "sk-ant-ort-new",
|
||||
"expires_in": 3600,
|
||||
"token_type": "bearer",
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let cfg = mock_cfg(host);
|
||||
let pkce = generate_pkce();
|
||||
let auth = exchange_code(
|
||||
&cfg,
|
||||
"auth-code-xyz",
|
||||
&pkce,
|
||||
&redirect_uri(53692, "/callback"),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(auth.key, "sk-ant-oat-new");
|
||||
assert_eq!(auth.refresh_token.as_deref(), Some("sk-ant-ort-new"));
|
||||
assert_eq!(auth.expires_in, Some(3600));
|
||||
}
|
||||
|
||||
/// Refresh rotates the refresh token (JSON body, refresh grant).
|
||||
#[tokio::test]
|
||||
async fn refresh_rotates_refresh_token() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/oauth/token"))
|
||||
.and(body_string_contains("\"grant_type\":\"refresh_token\""))
|
||||
.and(body_string_contains("\"refresh_token\":\"ort-old\""))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"access_token": "oat-fresh",
|
||||
"refresh_token": "ort-rotated",
|
||||
"expires_in": 3600,
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let auth = refresh_token(&mock_cfg(host), "ort-old").await.unwrap();
|
||||
assert_eq!(auth.key, "oat-fresh");
|
||||
assert_eq!(
|
||||
auth.refresh_token.as_deref(),
|
||||
Some("ort-rotated"),
|
||||
"the rotated refresh token must be adopted"
|
||||
);
|
||||
}
|
||||
|
||||
/// A 401 on refresh maps to Unauthorized (drives the permanent-failure path).
|
||||
#[tokio::test]
|
||||
async fn refresh_401_maps_to_unauthorized() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/oauth/token"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(401)
|
||||
.set_body_json(serde_json::json!({ "error_description": "refresh revoked" })),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let err = refresh_token(&mock_cfg(host), "ort-dead")
|
||||
.await
|
||||
.unwrap_err();
|
||||
match err {
|
||||
RefreshError::Unauthorized {
|
||||
status,
|
||||
description,
|
||||
} => {
|
||||
assert_eq!(status, 401);
|
||||
assert_eq!(description, "refresh revoked");
|
||||
}
|
||||
other => panic!("expected Unauthorized, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
// ── ChatGPT/Codex PKCE (openai-codex) ────────────────────────────────────
|
||||
|
||||
/// Codex PKCE uses an INDEPENDENT fresh-random state (NOT `state ==
|
||||
/// verifier`) so the verifier never rides the callback.
|
||||
#[test]
|
||||
fn codex_pkce_state_is_independent_of_the_verifier() {
|
||||
let pkce = generate_pkce_random_state();
|
||||
assert_ne!(
|
||||
pkce.state, pkce.verifier,
|
||||
"codex state must be fresh-random, not the verifier"
|
||||
);
|
||||
assert!(!pkce.state.is_empty() && !pkce.verifier.is_empty());
|
||||
// Challenge is still the S256 of the verifier.
|
||||
let expect = base64::engine::general_purpose::URL_SAFE_NO_PAD
|
||||
.encode(Sha256::digest(pkce.verifier.as_bytes()));
|
||||
assert_eq!(pkce.challenge, expect);
|
||||
assert_ne!(
|
||||
generate_pkce_random_state().state,
|
||||
pkce.state,
|
||||
"fresh state"
|
||||
);
|
||||
}
|
||||
|
||||
/// The codex authorize URL carries the PKCE state + S256 challenge AND the
|
||||
/// three codex-only extra params, and targets `auth.openai.com/oauth/
|
||||
/// authorize` with the `/auth/callback` redirect. The verifier never rides it.
|
||||
#[test]
|
||||
fn codex_authorize_url_has_state_challenge_and_three_extra_params() {
|
||||
use kigi_models::CODEX_OAUTH_CONFIG;
|
||||
let pkce = generate_pkce_random_state();
|
||||
let redirect = redirect_uri(1455, "/auth/callback");
|
||||
assert_eq!(redirect, "http://localhost:1455/auth/callback");
|
||||
let url = build_authorize_url(&CODEX_OAUTH_CONFIG, &redirect, &pkce);
|
||||
let parsed = url::Url::parse(&url).expect("valid URL");
|
||||
assert_eq!(parsed.host_str(), Some("auth.openai.com"));
|
||||
assert_eq!(parsed.path(), "/oauth/authorize");
|
||||
let q: std::collections::HashMap<_, _> = parsed.query_pairs().into_owned().collect();
|
||||
assert_eq!(
|
||||
q.get("state").map(String::as_str),
|
||||
Some(pkce.state.as_str())
|
||||
);
|
||||
assert_eq!(
|
||||
q.get("code_challenge").map(String::as_str),
|
||||
Some(pkce.challenge.as_str())
|
||||
);
|
||||
assert_eq!(
|
||||
q.get("code_challenge_method").map(String::as_str),
|
||||
Some("S256")
|
||||
);
|
||||
// The three codex-only extra params.
|
||||
assert_eq!(
|
||||
q.get("id_token_add_organizations").map(String::as_str),
|
||||
Some("true")
|
||||
);
|
||||
assert_eq!(
|
||||
q.get("codex_cli_simplified_flow").map(String::as_str),
|
||||
Some("true")
|
||||
);
|
||||
assert_eq!(
|
||||
q.get("originator").map(String::as_str),
|
||||
Some("codex_cli_rs")
|
||||
);
|
||||
assert!(
|
||||
!url.contains("code_verifier"),
|
||||
"the verifier must not ride the authorize URL"
|
||||
);
|
||||
}
|
||||
|
||||
/// The claude authorize URL is UNCHANGED (no extra params) — its empty
|
||||
/// `authorize_extra` keeps it byte-identical.
|
||||
#[test]
|
||||
fn claude_authorize_url_carries_no_extra_params() {
|
||||
let pkce = generate_pkce();
|
||||
let url = build_authorize_url(
|
||||
&CLAUDE_OAUTH_CONFIG,
|
||||
&redirect_uri(53692, "/callback"),
|
||||
&pkce,
|
||||
);
|
||||
assert!(!url.contains("id_token_add_organizations"));
|
||||
assert!(!url.contains("codex_cli_simplified_flow"));
|
||||
assert!(!url.contains("originator"));
|
||||
}
|
||||
|
||||
fn codex_mock_cfg(token_host: &'static str) -> OAuthConfig {
|
||||
OAuthConfig {
|
||||
token_host,
|
||||
..kigi_models::CODEX_OAUTH_CONFIG
|
||||
}
|
||||
}
|
||||
|
||||
/// Codex code→token exchange posts a FORM body carrying the grant + code +
|
||||
/// verifier + redirect_uri, and NOTABLY NO `state` field (the codex token
|
||||
/// endpoint does not expect it). Response materializes a `KimiAuth`.
|
||||
#[tokio::test]
|
||||
async fn codex_exchange_code_posts_form_without_state() {
|
||||
use wiremock::matchers::{body_string_contains, header, method, path};
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth/token"))
|
||||
.and(header("content-type", "application/x-www-form-urlencoded"))
|
||||
.and(body_string_contains("grant_type=authorization_code"))
|
||||
.and(body_string_contains("code=codex-auth-code"))
|
||||
.and(body_string_contains("code_verifier="))
|
||||
.and(body_string_contains("redirect_uri="))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"access_token": "codex-access-jwt",
|
||||
"refresh_token": "codex-refresh",
|
||||
"expires_in": 3600,
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let cfg = codex_mock_cfg(host);
|
||||
let pkce = generate_pkce_random_state();
|
||||
let auth = exchange_code_form(
|
||||
&cfg,
|
||||
"codex-auth-code",
|
||||
&pkce,
|
||||
"http://localhost:1455/auth/callback",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(auth.key, "codex-access-jwt");
|
||||
assert_eq!(auth.refresh_token.as_deref(), Some("codex-refresh"));
|
||||
assert_eq!(auth.expires_in, Some(3600));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,193 @@
|
||||
//! Process-global per-provider OAuth `AuthManager` pool for INFERENCE-time auth.
|
||||
//!
|
||||
//! A session binds its primary (Kimi / first-party) [`AuthManager`] for the
|
||||
//! subscription path, but a `uses_oauth` platform that carries an
|
||||
//! [`kigi_models::OAuthConfig`] (xai-grok today) needs its OWN scope-keyed
|
||||
//! manager for every per-turn decision — bearer resolution, proactive /
|
||||
//! on-expiry refresh, and 401 recovery. Reusing the Kimi manager for a grok
|
||||
//! turn would transmit the Kimi subscription bearer to `api.x.ai` (a
|
||||
//! cross-provider leak, guaranteed 401) and, without proactive refresh, would
|
||||
//! 401 every turn once the ~1h grok token expired until a process restart.
|
||||
//!
|
||||
//! The pool is the SINGLE SOURCE OF TRUTH: one long-lived `AuthManager` per
|
||||
//! generic-oauth scope, each wired with the SAME lifecycle as the primary Kimi
|
||||
//! manager (`configure_refresher()` + `start_proactive_refresh()`) so the
|
||||
//! on-disk token stays fresh and a 401 recovers via the provider's own manager.
|
||||
//! Managers are built ON DEMAND from the on-disk token ([`global_manager_for`]),
|
||||
//! so a login landing AFTER a session spawned self-heals — no frozen per-session
|
||||
//! snapshot.
|
||||
//!
|
||||
//! ROUTING LIVES ELSEWHERE. This module is only the pool; the decision of which
|
||||
//! credential governs a request belongs to the single chokepoint,
|
||||
//! [`crate::auth::credential_authority::CredentialAuthority`]. Keeping the two
|
||||
//! apart is deliberate: three rounds of leaks came from routing rules being
|
||||
//! re-derived per call site.
|
||||
//!
|
||||
//! SECURITY: access/refresh tokens and resolved bearers are NEVER logged here.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
use parking_lot::Mutex;
|
||||
|
||||
use crate::auth::AuthManager;
|
||||
|
||||
/// Process-wide pool of live per-scope OAuth managers.
|
||||
///
|
||||
/// Auth is process-global (one user), so a single manager per scope is correct
|
||||
/// and lets the proactive-refresh task start exactly once per scope no matter
|
||||
/// how many sessions spawn. Keyed by the OAuth `scope_key` (`oauth/xai`, …).
|
||||
fn oauth_manager_pool() -> &'static Mutex<HashMap<&'static str, Arc<AuthManager>>> {
|
||||
static POOL: OnceLock<Mutex<HashMap<&'static str, Arc<AuthManager>>>> = OnceLock::new();
|
||||
POOL.get_or_init(|| Mutex::new(HashMap::new()))
|
||||
}
|
||||
|
||||
/// The kigi home EVERY OAuth-provider construction in this crate resolves from
|
||||
/// — the pool, the catalog fetch's per-platform token resolution, the
|
||||
/// aux/summary token routing and the session's inference manager — so they can
|
||||
/// never read different homes.
|
||||
///
|
||||
/// Production: [`crate::util::kigi_home::kigi_home`]. LIB TESTS: a
|
||||
/// per-process path under the system temp dir that is deliberately **never
|
||||
/// created**. 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.
|
||||
///
|
||||
/// M4 — THE LIMIT, STATED: `cfg(test)` is set only for THIS crate's `--lib`
|
||||
/// tests. `crates/codegen/kigi-shell/tests/*.rs` link the library built WITHOUT
|
||||
/// it, so for an integration test this resolves the real home unless that test
|
||||
/// binary itself isolates one, which it must do through the two overrides the
|
||||
/// auth stack already honours and BEFORE anything resolves `kigi_home()`:
|
||||
/// `KIGI_SHARE_DIR` (read by `kigi_home()`, a `OnceLock`) or `KIGI_AUTH_PATH`
|
||||
/// (read by [`AuthManager::new_oauth_provider`], which pins the token file
|
||||
/// outright and so overrides this home entirely). 12 of the 28 integration
|
||||
/// binaries under `crates/codegen/kigi-shell/tests/` set `KIGI_SHARE_DIR`; the
|
||||
/// other 16 never reach an OAuth-platform inference path today, which is a
|
||||
/// property of those tests, not a guarantee of this function. No
|
||||
/// production-readable env override is added here on purpose: a knob that
|
||||
/// redirects where OAuth tokens are read from is not worth a test convenience.
|
||||
///
|
||||
/// M8: this used to be a `static OnceLock<TempDir>`. Statics are never dropped,
|
||||
/// so that leaked one temp directory per test binary — against the project's
|
||||
/// "tests are TempDir self-cleaning" discipline. Nothing is created here
|
||||
/// instead, and nothing in the lib-test suite creates it: a manager reads a
|
||||
/// missing `auth.json` as "no session", and the only two paths that WRITE one
|
||||
/// are a successful token refresh (which needs a stored refresh token that by
|
||||
/// construction does not exist here) and a completed device login
|
||||
/// ([`crate::agent::mvp_agent::MvpAgent::authenticate_oauth_platform`], which
|
||||
/// M4 repointed at this same home). Both require the network, so no lib test
|
||||
/// performs either. That is an observation about the suite, not an invariant of
|
||||
/// this function — [`tests::test_pool_home_is_disposable_and_never_the_real_home`]
|
||||
/// asserts the directory does not exist and is the tripwire if one ever does
|
||||
/// (the path is per-PROCESS under the system temp dir, so the blast radius of a
|
||||
/// future login-driving test is one disposable directory, never `~/.kigi`).
|
||||
pub(crate) fn pool_home() -> std::path::PathBuf {
|
||||
#[cfg(test)]
|
||||
{
|
||||
std::env::temp_dir().join(format!("kigi-oauth-pool-test-{}", std::process::id()))
|
||||
}
|
||||
#[cfg(not(test))]
|
||||
crate::util::kigi_home::kigi_home()
|
||||
}
|
||||
|
||||
/// Get-or-create the process-global manager for `oauth`, wiring the same
|
||||
/// refresher + proactive-refresh lifecycle as the primary Kimi manager the
|
||||
/// FIRST time a scope is seen. The manager reads the on-disk token at
|
||||
/// construction (thereafter kept fresh by the proactive-refresh loop), so a
|
||||
/// grok login that lands after this scope was first built is adopted on the
|
||||
/// manager's own refresh tick — no session ever needs re-spawning.
|
||||
///
|
||||
/// MUST be called from within a Tokio runtime (the proactive-refresh loop
|
||||
/// spawns a task, mirroring the primary).
|
||||
pub(crate) fn global_manager_for(
|
||||
kigi_home: &Path,
|
||||
oauth: &'static kigi_models::OAuthConfig,
|
||||
) -> Arc<AuthManager> {
|
||||
let mut pool = oauth_manager_pool().lock();
|
||||
if let Some(existing) = pool.get(oauth.scope_key) {
|
||||
return existing.clone();
|
||||
}
|
||||
let manager = Arc::new(AuthManager::new_oauth_provider(kigi_home, oauth));
|
||||
manager.configure_refresher();
|
||||
// Never-cancelled token = process-lifetime, matching the api-server /
|
||||
// per-session eager-refresh sites that pass a fresh token.
|
||||
manager.start_proactive_refresh(tokio_util::sync::CancellationToken::new());
|
||||
pool.insert(oauth.scope_key, manager.clone());
|
||||
manager
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn oauth_for(id: &str) -> &'static kigi_models::OAuthConfig {
|
||||
kigi_models::PlatformId::parse(id)
|
||||
.expect("known platform")
|
||||
.oauth()
|
||||
.expect("subscription-OAuth platform carries an OAuthConfig")
|
||||
}
|
||||
|
||||
/// Each subscription-OAuth platform gets its OWN process-global pooled
|
||||
/// manager, and no two share one. (Which credential governs a REQUEST is
|
||||
/// not decided here — see
|
||||
/// [`crate::auth::credential_authority::CredentialAuthority`].)
|
||||
#[tokio::test]
|
||||
async fn every_oauth_scope_gets_its_own_pooled_manager() {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let ids = [
|
||||
"xai-grok",
|
||||
"claude-pro-max",
|
||||
"github-copilot",
|
||||
"openai-codex",
|
||||
];
|
||||
let managers: Vec<_> = ids
|
||||
.iter()
|
||||
.map(|id| global_manager_for(home.path(), oauth_for(id)))
|
||||
.collect();
|
||||
for (i, a) in managers.iter().enumerate() {
|
||||
assert!(
|
||||
Arc::ptr_eq(a, &global_manager_for(home.path(), oauth_for(ids[i]))),
|
||||
"{}: the pool must return the SAME manager for a scope",
|
||||
ids[i]
|
||||
);
|
||||
for (j, b) in managers.iter().enumerate() {
|
||||
if i != j {
|
||||
assert!(
|
||||
!Arc::ptr_eq(a, b),
|
||||
"{} and {} must not share a pooled manager",
|
||||
ids[i],
|
||||
ids[j]
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// M8: the test pool home is a per-process path that is never created, so a
|
||||
/// test binary leaves nothing behind (and never resolves the developer's
|
||||
/// real `~/.kigi`, whose stored OAuth tokens the pool would otherwise read
|
||||
/// and proactively refresh over the network).
|
||||
///
|
||||
/// This is also the tripwire for the cleanup claim in [`pool_home`]: a
|
||||
/// completed device login through `authenticate_oauth_platform` WOULD create
|
||||
/// this directory, so if a lib test ever drives one, this assertion fires
|
||||
/// and the cleanup has to be added rather than silently regressing the
|
||||
/// "tests are TempDir self-cleaning" discipline.
|
||||
#[test]
|
||||
fn test_pool_home_is_disposable_and_never_the_real_home() {
|
||||
let home = pool_home();
|
||||
assert!(
|
||||
home.starts_with(std::env::temp_dir()),
|
||||
"the test pool home must live under the system temp dir, got {home:?}"
|
||||
);
|
||||
assert!(
|
||||
!home.exists(),
|
||||
"the test pool home must not be created — nothing to clean up"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,314 @@
|
||||
//! Generic OAuth token refresher for any [`kigi_models::OAuthConfig`] provider,
|
||||
//! driven through the [`TokenRefresher`] seam. The wire call is selected by the
|
||||
//! config's `token_body`: form-encoded `grant_type=refresh_token` (xai-grok,
|
||||
//! openai-codex), JSON (claude-pro-max), or the GitHub Copilot copilot-token
|
||||
//! RE-MINT. Providers with `requires_chatgpt_account_id` additionally fail fast
|
||||
//! when the refreshed token drops the claim.
|
||||
//!
|
||||
//! Structurally identical to [`super::kimi_refresher::KimiRefresher`] — same
|
||||
//! sibling-adoption + post-401 grace — but the wire call goes through
|
||||
//! [`crate::auth::oauth_device`] (no X-Msh headers) instead of the Kimi wire.
|
||||
//! Access/refresh tokens are NEVER logged.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use kigi_models::{OAuthConfig, OAuthTokenBody};
|
||||
|
||||
use crate::auth::error::RefreshTokenFailedReason;
|
||||
use crate::auth::kimi_oauth::RefreshError;
|
||||
use crate::auth::manager::RefreshReason;
|
||||
use crate::auth::{github_copilot, oauth_device, oauth_pkce};
|
||||
|
||||
use super::{AuthSnapshot, RefreshOutcome, TokenRefresher};
|
||||
|
||||
/// Grace period after a 401/403 before concluding the refresh token is dead:
|
||||
/// a concurrent instance may still be persisting its rotated token.
|
||||
const POST_UNAUTHORIZED_GRACE: std::time::Duration = std::time::Duration::from_secs(1);
|
||||
|
||||
pub(crate) struct GenericDeviceRefresher {
|
||||
auth: Arc<dyn AuthSnapshot>,
|
||||
cfg: &'static OAuthConfig,
|
||||
}
|
||||
|
||||
impl GenericDeviceRefresher {
|
||||
pub(crate) fn new(auth: Arc<dyn AuthSnapshot>, cfg: &'static OAuthConfig) -> Self {
|
||||
Self { auth, cfg }
|
||||
}
|
||||
|
||||
/// Post-401 sibling check: wait a beat, re-read the persisted credential,
|
||||
/// and adopt it when its refresh token differs from the rejected one.
|
||||
async fn adopt_rotation_after_unauthorized(&self, tried_rt: &str) -> Option<RefreshOutcome> {
|
||||
tokio::time::sleep(POST_UNAUTHORIZED_GRACE).await;
|
||||
let latest = self.auth.read_disk_auth()?;
|
||||
let latest_rt = latest.refresh_token.as_deref()?;
|
||||
if latest_rt == tried_rt {
|
||||
return None;
|
||||
}
|
||||
kigi_log::unified_log::info(
|
||||
"auth.refresh.adopted_rotation_after_401",
|
||||
None,
|
||||
Some(serde_json::json!({
|
||||
"scope_key": self.cfg.scope_key,
|
||||
"adopted_rt_prefix": crate::auth::token_suffix(latest_rt),
|
||||
"rejected_rt_prefix": crate::auth::token_suffix(tried_rt),
|
||||
})),
|
||||
);
|
||||
Some(RefreshOutcome::success(latest))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl TokenRefresher for GenericDeviceRefresher {
|
||||
async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome {
|
||||
tracing::info!(
|
||||
?reason,
|
||||
scope_key = self.cfg.scope_key,
|
||||
"auth: generic refresh attempt"
|
||||
);
|
||||
|
||||
let disk_auth = self.auth.read_disk_auth();
|
||||
|
||||
// Sibling short-circuit: a valid persisted token whose key differs from
|
||||
// in-memory means another process refreshed already — adopt directly.
|
||||
if let Some(ref d) = disk_auth
|
||||
&& !crate::auth::is_expired(d)
|
||||
&& self.auth.current().map(|a| a.key).as_deref() != Some(&d.key)
|
||||
{
|
||||
kigi_log::unified_log::info(
|
||||
"auth.refresh.adopted_sibling_token",
|
||||
None,
|
||||
Some(serde_json::json!({
|
||||
"scope_key": self.cfg.scope_key,
|
||||
"disk_key_prefix": crate::auth::token_suffix(&d.key),
|
||||
})),
|
||||
);
|
||||
return RefreshOutcome::success(d.clone());
|
||||
}
|
||||
|
||||
let Some(auth) = super::resolve_refresh_credential(self.auth.as_ref(), disk_auth, reason)
|
||||
else {
|
||||
tracing::warn!(
|
||||
?reason,
|
||||
"auth: no credential available for refresh (generic)"
|
||||
);
|
||||
return RefreshOutcome::transient("no token with refresh_token available");
|
||||
};
|
||||
let Some(refresh_token) = auth.refresh_token.clone() else {
|
||||
tracing::warn!(
|
||||
?reason,
|
||||
"auth: resolved credential has no refresh token (generic)"
|
||||
);
|
||||
return RefreshOutcome::transient("credential has no refresh token");
|
||||
};
|
||||
|
||||
tracing::info!(
|
||||
rt_prefix = crate::auth::token_suffix(&refresh_token),
|
||||
expires_at = ?auth.expires_at,
|
||||
"auth: sending refresh_token grant (generic oauth)"
|
||||
);
|
||||
|
||||
// Refresh over the provider's token-body encoding: xai's endpoint is
|
||||
// form-encoded (device wire); Claude's is JSON (PKCE wire); GitHub
|
||||
// Copilot's "refresh" is a copilot-token RE-MINT — a `GET
|
||||
// copilot_internal/v2/token` bearing the durable github token (the
|
||||
// `refresh_token` field here), NOT a refresh_token grant. All three
|
||||
// return the same `Result<KimiAuth, RefreshError>`.
|
||||
let wire_result = match self.cfg.token_body {
|
||||
OAuthTokenBody::Form => oauth_device::refresh_token(self.cfg, &refresh_token).await,
|
||||
OAuthTokenBody::Json => oauth_pkce::refresh_token(self.cfg, &refresh_token).await,
|
||||
OAuthTokenBody::GithubCopilotExchange => {
|
||||
github_copilot::remint_copilot_token(self.cfg, &refresh_token).await
|
||||
}
|
||||
};
|
||||
match wire_result {
|
||||
Ok(new_auth)
|
||||
if self.cfg.requires_chatgpt_account_id
|
||||
&& kigi_sampling_types::chatgpt_account_id_from_jwt(&new_auth.key)
|
||||
.is_none() =>
|
||||
{
|
||||
// FAIL FAST (never silently): the refreshed token carries no
|
||||
// `chatgpt_account_id`, so every inference request would go out
|
||||
// WITHOUT the required `chatgpt-account-id` header and draw an
|
||||
// opaque backend 4xx. Surface it as a permanent failure so the
|
||||
// user is told to re-login.
|
||||
tracing::warn!(
|
||||
scope_key = self.cfg.scope_key,
|
||||
"auth: refreshed token is missing the chatgpt_account_id claim"
|
||||
);
|
||||
RefreshOutcome::permanent(RefreshTokenFailedReason::Other, Some(refresh_token))
|
||||
}
|
||||
Ok(new_auth) => {
|
||||
kigi_log::unified_log::info(
|
||||
"auth.refresh.token_rotated",
|
||||
None,
|
||||
Some(serde_json::json!({
|
||||
"scope_key": self.cfg.scope_key,
|
||||
"new_key_prefix": crate::auth::token_suffix(&new_auth.key),
|
||||
"expires_at": new_auth.expires_at.map(|e| e.to_rfc3339()),
|
||||
})),
|
||||
);
|
||||
RefreshOutcome::success(new_auth)
|
||||
}
|
||||
Err(RefreshError::Unauthorized {
|
||||
status,
|
||||
description,
|
||||
}) => {
|
||||
tracing::warn!(status, %description, "auth: refresh token rejected (generic)");
|
||||
if let Some(adopted) = self.adopt_rotation_after_unauthorized(&refresh_token).await
|
||||
{
|
||||
return adopted;
|
||||
}
|
||||
kigi_log::unified_log::warn(
|
||||
"auth.refresh.unauthorized",
|
||||
None,
|
||||
Some(serde_json::json!({
|
||||
"scope_key": self.cfg.scope_key,
|
||||
"status": status,
|
||||
"description": description,
|
||||
"rt_prefix": crate::auth::token_suffix(&refresh_token),
|
||||
})),
|
||||
);
|
||||
RefreshOutcome::permanent(
|
||||
RefreshTokenFailedReason::RefreshTokenRejected,
|
||||
Some(refresh_token),
|
||||
)
|
||||
}
|
||||
Err(
|
||||
e @ (RefreshError::Exhausted { .. }
|
||||
| RefreshError::Fatal { .. }
|
||||
| RefreshError::Local(_)),
|
||||
) => {
|
||||
tracing::warn!(error = %e, "auth: refresh attempt failed (transient, generic)");
|
||||
kigi_log::unified_log::warn(
|
||||
"auth.refresh.transient_wire_failure",
|
||||
None,
|
||||
Some(serde_json::json!({
|
||||
"scope_key": self.cfg.scope_key,
|
||||
"error": format!("{e}"),
|
||||
})),
|
||||
);
|
||||
RefreshOutcome::transient(format!("token refresh failed: {e}"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::auth::model::KimiAuth;
|
||||
use chrono::{Duration, Utc};
|
||||
use kigi_models::XAI_OAUTH_CONFIG;
|
||||
use parking_lot::Mutex;
|
||||
use wiremock::matchers::{body_string_contains, method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
struct FakeSnapshot {
|
||||
current: Mutex<Option<KimiAuth>>,
|
||||
disk: Mutex<Option<KimiAuth>>,
|
||||
}
|
||||
impl FakeSnapshot {
|
||||
fn new(current: Option<KimiAuth>, disk: Option<KimiAuth>) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
current: Mutex::new(current),
|
||||
disk: Mutex::new(disk),
|
||||
})
|
||||
}
|
||||
}
|
||||
impl AuthSnapshot for FakeSnapshot {
|
||||
fn current(&self) -> Option<KimiAuth> {
|
||||
self.current
|
||||
.lock()
|
||||
.clone()
|
||||
.filter(|a| !crate::auth::is_expired(a))
|
||||
}
|
||||
fn expired_auth(&self) -> Option<KimiAuth> {
|
||||
self.current.lock().clone().filter(crate::auth::is_expired)
|
||||
}
|
||||
fn read_disk_auth(&self) -> Option<KimiAuth> {
|
||||
self.disk.lock().clone()
|
||||
}
|
||||
fn is_expired(&self) -> bool {
|
||||
self.current
|
||||
.lock()
|
||||
.as_ref()
|
||||
.is_some_and(crate::auth::is_expired)
|
||||
}
|
||||
}
|
||||
|
||||
fn expired_session(key: &str, rt: &str) -> KimiAuth {
|
||||
KimiAuth {
|
||||
key: key.into(),
|
||||
refresh_token: Some(rt.into()),
|
||||
expires_at: Some(Utc::now() - Duration::hours(1)),
|
||||
expires_in: Some(3600),
|
||||
..KimiAuth::test_default()
|
||||
}
|
||||
}
|
||||
|
||||
fn mock_cfg(host: &'static str) -> OAuthConfig {
|
||||
OAuthConfig {
|
||||
auth_host: host,
|
||||
..XAI_OAUTH_CONFIG
|
||||
}
|
||||
}
|
||||
|
||||
/// A successful refresh rotates the token via the generic wire.
|
||||
#[tokio::test]
|
||||
async fn refresh_success_returns_rotated_token() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth2/token"))
|
||||
.and(body_string_contains("refresh_token=grok-rt-old"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"access_token": "grok-at-new",
|
||||
"refresh_token": "grok-rt-new",
|
||||
"expires_in": 3600,
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let stale = expired_session("grok-at-old", "grok-rt-old");
|
||||
let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale));
|
||||
let cfg: &'static OAuthConfig = Box::leak(Box::new(mock_cfg(host)));
|
||||
let refresher = GenericDeviceRefresher::new(snap, cfg);
|
||||
let outcome = refresher.refresh(RefreshReason::PreRequest).await;
|
||||
let RefreshOutcome::Success(new_auth) = outcome else {
|
||||
panic!("expected success, got {outcome:?}");
|
||||
};
|
||||
assert_eq!(new_auth.key, "grok-at-new");
|
||||
assert_eq!(new_auth.refresh_token.as_deref(), Some("grok-rt-new"));
|
||||
}
|
||||
|
||||
/// A 401 on refresh (with no sibling rotation) tombstones the rejected
|
||||
/// refresh token as a permanent failure.
|
||||
#[tokio::test]
|
||||
async fn unauthorized_is_permanent_failure() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth2/token"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(401)
|
||||
.set_body_json(serde_json::json!({ "error_description": "revoked" })),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let stale = expired_session("grok-at-old", "grok-rt-dead");
|
||||
let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale));
|
||||
let cfg: &'static OAuthConfig = Box::leak(Box::new(mock_cfg(host)));
|
||||
let refresher = GenericDeviceRefresher::new(snap, cfg);
|
||||
let outcome = refresher.refresh(RefreshReason::PreRequest).await;
|
||||
let RefreshOutcome::PermanentFailure {
|
||||
error,
|
||||
rejected_refresh_token,
|
||||
} = outcome
|
||||
else {
|
||||
panic!("expected permanent failure, got {outcome:?}");
|
||||
};
|
||||
assert_eq!(error.reason, RefreshTokenFailedReason::RefreshTokenRejected);
|
||||
assert_eq!(rejected_refresh_token.as_deref(), Some("grok-rt-dead"));
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
mod generic_refresher;
|
||||
mod kimi_refresher;
|
||||
|
||||
use std::sync::Arc;
|
||||
@@ -6,6 +7,7 @@ use crate::auth::manager::AuthManager;
|
||||
pub(crate) use crate::auth::manager::RefreshReason;
|
||||
use crate::auth::model::KimiAuth;
|
||||
|
||||
pub(crate) use generic_refresher::GenericDeviceRefresher;
|
||||
pub(crate) use kimi_refresher::KimiRefresher;
|
||||
|
||||
/// Read-only view of `AuthManager` for refreshers. Enforces the
|
||||
@@ -120,8 +122,20 @@ pub(crate) trait TokenRefresher: Send + Sync {
|
||||
async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome;
|
||||
}
|
||||
|
||||
/// Build the production refresher against `kigi_env::oauth_host()`.
|
||||
/// Build the production refresher for this manager's scope. A scope that maps
|
||||
/// to a generic device-code [`kigi_models::OAuthConfig`] (xai-grok) gets the
|
||||
/// [`GenericDeviceRefresher`]; every other scope — Kimi Code, whose registry
|
||||
/// `oauth` field is `None` by design — gets the bespoke [`KimiRefresher`]
|
||||
/// against `kigi_env::oauth_host()`.
|
||||
pub(crate) fn build_refresher(auth_manager: Arc<AuthManager>) -> Arc<dyn TokenRefresher> {
|
||||
match kigi_models::oauth_config_for_scope_key(auth_manager.scope()) {
|
||||
Some(cfg) => {
|
||||
let snapshot: Arc<dyn AuthSnapshot> = auth_manager;
|
||||
Arc::new(GenericDeviceRefresher::new(snapshot, cfg))
|
||||
}
|
||||
None => {
|
||||
let snapshot: Arc<dyn AuthSnapshot> = auth_manager;
|
||||
Arc::new(KimiRefresher::new(snapshot, kigi_env::oauth_host()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -96,9 +96,9 @@ pub(crate) fn disable_mock_keyring_for_test() {
|
||||
fn keyring_entry() -> Result<&'static keyring::Entry, keyring::Error> {
|
||||
static ENTRY: std::sync::OnceLock<Result<keyring::Entry, keyring::Error>> =
|
||||
std::sync::OnceLock::new();
|
||||
match ENTRY.get_or_init(|| {
|
||||
keyring::Entry::new(KEYRING_SERVICE, crate::auth::config::KIMI_CODE_OAUTH_SCOPE)
|
||||
}) {
|
||||
match ENTRY
|
||||
.get_or_init(|| keyring::Entry::new(KEYRING_SERVICE, crate::auth::KIMI_CODE_OAUTH_SCOPE))
|
||||
{
|
||||
Ok(entry) => Ok(entry),
|
||||
// `keyring::Error` is not `Clone`; surface a stable equivalent.
|
||||
Err(e) => {
|
||||
@@ -430,17 +430,12 @@ fn write_store_to(path: &Path, auth_store: &AuthStore) -> std::io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Atomic write: tmp + rename. Unix `rename(2)` replaces atomically;
|
||||
/// Windows `rename` requires removing the target first.
|
||||
/// Atomic write: tmp + Windows-safe replace (see `util::fs::replace_file`,
|
||||
/// which this site's inline delete-first pattern graduated into).
|
||||
fn write_auth_json_atomic(auth_file: &Path, auth_store: &AuthStore) -> std::io::Result<()> {
|
||||
let tmp = auth_file.with_extension(format!("json.{}.tmp", std::process::id()));
|
||||
write_store_to(&tmp, auth_store)?;
|
||||
#[cfg(windows)]
|
||||
{
|
||||
let _ = std::fs::remove_file(auth_file);
|
||||
}
|
||||
std::fs::rename(&tmp, auth_file)?;
|
||||
Ok(())
|
||||
crate::util::fs::replace_file(&tmp, auth_file)
|
||||
}
|
||||
|
||||
/// Non-atomic fallback: truncate and rewrite `auth.json` in place.
|
||||
@@ -536,6 +531,71 @@ pub fn store_api_key(kigi_home: &Path, api_key: &str) -> std::io::Result<()> {
|
||||
write_auth_json(&path, &map)
|
||||
}
|
||||
|
||||
/// Read an API-key platform's key from auth.json. The scope is the platform
|
||||
/// id itself (`anthropic`, `moonshot-cn`, …) — the stable per-provider
|
||||
/// auth.json key contract. `None` when absent or unreadable.
|
||||
pub fn read_platform_api_key(
|
||||
kigi_home: &Path,
|
||||
platform: kigi_models::PlatformId,
|
||||
) -> Option<String> {
|
||||
let path = kigi_home.join("auth.json");
|
||||
let map = read_auth_json(&path).ok()?;
|
||||
map.get(platform.as_str()).map(|a| a.key.clone())
|
||||
}
|
||||
|
||||
/// Store an API-key platform's key in auth.json under its platform-id scope.
|
||||
/// Same corrupt-recovery + atomic-write path as [`store_api_key`]; all other
|
||||
/// scopes (OAuth session, other platforms) are preserved.
|
||||
///
|
||||
/// SECURITY: the key must never be logged; errors carry only IO context.
|
||||
pub fn store_platform_api_key(
|
||||
kigi_home: &Path,
|
||||
platform: kigi_models::PlatformId,
|
||||
api_key: &str,
|
||||
) -> std::io::Result<()> {
|
||||
if platform.uses_oauth() {
|
||||
// Real error, not debug_assert: an OAuth scope written here would
|
||||
// shadow the session entry in release builds too.
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
format!(
|
||||
"{} authenticates via OAuth and takes no API key",
|
||||
platform.as_str()
|
||||
),
|
||||
));
|
||||
}
|
||||
let path = kigi_home.join("auth.json");
|
||||
// Serialize with the manager's cross-process auth.json writers (token
|
||||
// refresh holds the same flock): an unlocked read-modify-write here
|
||||
// could write back a pre-refresh map and revert a rotated refresh
|
||||
// token (token-family revocation → forced re-login). Bounded retry;
|
||||
// a sustained holder fails loudly rather than racing.
|
||||
let mut lock = None;
|
||||
for _ in 0..20 {
|
||||
lock = super::manager::try_lock_auth_file_nonblocking(&path);
|
||||
if lock.is_some() {
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(25));
|
||||
}
|
||||
let Some(_lock) = lock else {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::WouldBlock,
|
||||
"auth.json is locked by another kigi process; try again",
|
||||
));
|
||||
};
|
||||
let mut map = read_auth_json_or_empty_recovering_corrupt(&path)?;
|
||||
map.insert(
|
||||
platform.as_str().to_owned(),
|
||||
KimiAuth {
|
||||
key: api_key.to_owned(),
|
||||
auth_mode: AuthMode::ApiKey,
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
write_auth_json(&path, &map)
|
||||
}
|
||||
|
||||
/// Remove the `kigi::api_key` scope from auth.json.
|
||||
pub fn clear_api_key(kigi_home: &Path) -> std::io::Result<()> {
|
||||
let path = kigi_home.join("auth.json");
|
||||
@@ -550,6 +610,80 @@ pub fn clear_api_key(kigi_home: &Path) -> std::io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod platform_key_tests {
|
||||
use super::*;
|
||||
|
||||
/// Store/read round-trip under the platform-id scope; the OAuth session
|
||||
/// scope and other platform scopes in the same file are preserved.
|
||||
#[test]
|
||||
fn platform_key_round_trip_preserves_other_scopes() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let home = dir.path();
|
||||
// Pre-existing OAuth session entry must survive platform-key writes.
|
||||
let path = home.join("auth.json");
|
||||
let mut map = AuthStore::new();
|
||||
map.insert(
|
||||
crate::auth::KIMI_CODE_OAUTH_SCOPE.to_owned(),
|
||||
KimiAuth {
|
||||
key: "oauth-token".to_owned(),
|
||||
auth_mode: AuthMode::OAuth,
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
write_auth_json(&path, &map).unwrap();
|
||||
|
||||
store_platform_api_key(home, kigi_models::PlatformId::MoonshotCn, "sk-cn").unwrap();
|
||||
store_platform_api_key(home, kigi_models::PlatformId::MoonshotAi, "sk-ai").unwrap();
|
||||
|
||||
assert_eq!(
|
||||
read_platform_api_key(home, kigi_models::PlatformId::MoonshotCn).as_deref(),
|
||||
Some("sk-cn")
|
||||
);
|
||||
assert_eq!(
|
||||
read_platform_api_key(home, kigi_models::PlatformId::MoonshotAi).as_deref(),
|
||||
Some("sk-ai")
|
||||
);
|
||||
let stored = read_auth_json(&path).unwrap();
|
||||
assert_eq!(
|
||||
stored
|
||||
.get(crate::auth::KIMI_CODE_OAUTH_SCOPE)
|
||||
.map(|a| a.key.as_str()),
|
||||
Some("oauth-token"),
|
||||
"platform-key writes must not clobber the OAuth session scope"
|
||||
);
|
||||
assert_eq!(
|
||||
stored.get("moonshot-cn").map(|a| a.auth_mode.clone()),
|
||||
Some(AuthMode::ApiKey),
|
||||
"platform keys are stored as api_key mode under the platform id"
|
||||
);
|
||||
}
|
||||
|
||||
/// The OAuth platform takes no API key — a real error in release builds.
|
||||
#[test]
|
||||
fn storing_key_for_oauth_platform_is_invalid_input() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let err = store_platform_api_key(dir.path(), kigi_models::PlatformId::KimiCode, "sk-x")
|
||||
.expect_err("oauth platform must reject api keys");
|
||||
assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
|
||||
assert!(
|
||||
!dir.path().join("auth.json").exists(),
|
||||
"rejected write must not create auth.json"
|
||||
);
|
||||
}
|
||||
|
||||
/// Missing file reads as None (not an error) — resolution treats absent
|
||||
/// auth.json as "no stored key".
|
||||
#[test]
|
||||
fn reading_platform_key_without_auth_json_is_none() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
assert_eq!(
|
||||
read_platform_api_key(dir.path(), kigi_models::PlatformId::MoonshotCn),
|
||||
None
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod write_fallback_tests {
|
||||
use super::*;
|
||||
|
||||
@@ -670,10 +670,7 @@ fn write_import_marker(config_path: &Path) -> anyhow::Result<()> {
|
||||
let _ = std::fs::remove_file(&tmp);
|
||||
return Err(e.into());
|
||||
}
|
||||
if let Err(e) = std::fs::rename(&tmp, config_path) {
|
||||
let _ = std::fs::remove_file(&tmp);
|
||||
return Err(e.into());
|
||||
}
|
||||
crate::util::fs::replace_file(&tmp, config_path)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -854,7 +851,7 @@ fn apply_items_to_config(config_path: &Path, items: &[ImportableItem]) -> anyhow
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
std::fs::write(&tmp, &toml_str)?;
|
||||
std::fs::rename(&tmp, config_path)?;
|
||||
crate::util::fs::replace_file(&tmp, config_path)?;
|
||||
info!(
|
||||
path = %config_path.display(),
|
||||
count,
|
||||
@@ -1204,7 +1201,7 @@ fn apply_hooks_to_dir(hooks_dir: &Path, items: &[ImportableItem]) -> anyhow::Res
|
||||
let json_str = serde_json::to_string_pretty(&root)?;
|
||||
let tmp = target.with_extension("json.tmp");
|
||||
std::fs::write(&tmp, &json_str)?;
|
||||
std::fs::rename(&tmp, &target)?;
|
||||
crate::util::fs::replace_file(&tmp, &target)?;
|
||||
info!(
|
||||
path = %target.display(),
|
||||
count,
|
||||
|
||||
@@ -90,7 +90,7 @@ pub fn save_import_state(state: &ImportState) -> std::io::Result<()> {
|
||||
// `claude_import_state.json.tmp` (the last extension is replaced).
|
||||
let tmp = path.with_extension("json.tmp");
|
||||
std::fs::write(&tmp, &json)?;
|
||||
std::fs::rename(&tmp, &path)?;
|
||||
crate::util::fs::replace_file(&tmp, &path)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@@ -35,7 +35,11 @@ impl AuthStatus {
|
||||
.unwrap_or(&origin);
|
||||
return Self::LoggedIn(host.to_owned());
|
||||
}
|
||||
let models = crate::agent::config::resolve_model_list(agent_config, None);
|
||||
let models = crate::agent::config::resolve_model_list(
|
||||
agent_config,
|
||||
None,
|
||||
&crate::agent::models::PlatformApiKeys::resolve(&agent_config.platforms),
|
||||
);
|
||||
if crate::agent::auth_method::should_advertise_xai_api_key(models.values())
|
||||
&& let Some(name) = models
|
||||
.iter()
|
||||
@@ -88,7 +92,9 @@ pub async fn list_models(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::agent::auth_method::{LEGACY_XAI_API_KEY_ENV_VAR, XAI_API_KEY_ENV_VAR};
|
||||
use crate::agent::auth_method::{
|
||||
HOUSE_API_KEY_ENV_VAR, LEGACY_XAI_API_KEY_ENV_VAR, XAI_API_KEY_ENV_VAR,
|
||||
};
|
||||
use crate::agent::config::Config;
|
||||
use crate::auth::{AuthMode, KimiAuth};
|
||||
use kigi_test_support::EnvGuard;
|
||||
@@ -98,10 +104,11 @@ mod tests {
|
||||
///
|
||||
/// Uses `KIGI_AUTH_PATH` (not `KIGI_SHARE_DIR`) so a OnceLock-cached real home
|
||||
/// with `auth.json` cannot leak into these tests.
|
||||
fn isolate_auth_sources() -> (tempfile::TempDir, [EnvGuard; 7]) {
|
||||
fn isolate_auth_sources() -> (tempfile::TempDir, [EnvGuard; 8]) {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let auth_path = dir.path().join("no-auth.json");
|
||||
let guards = [
|
||||
EnvGuard::unset(HOUSE_API_KEY_ENV_VAR),
|
||||
EnvGuard::unset(XAI_API_KEY_ENV_VAR),
|
||||
EnvGuard::unset(LEGACY_XAI_API_KEY_ENV_VAR),
|
||||
EnvGuard::unset("KIGI_AUTH"),
|
||||
|
||||
@@ -56,18 +56,18 @@ fn handle_set_api_key(args: &acp::ExtRequest) -> ExtResult {
|
||||
crate::auth::clear_api_key(&kigi_home)
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
// SAFETY: ext_method is single-threaded per agent
|
||||
unsafe { std::env::remove_var("XAI_API_KEY") };
|
||||
unsafe { std::env::remove_var("KIGI_API_KEY") };
|
||||
} else {
|
||||
crate::auth::store_api_key(&kigi_home, k)
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
// SAFETY: ext_method is single-threaded per agent
|
||||
unsafe { std::env::set_var("XAI_API_KEY", k) };
|
||||
unsafe { std::env::set_var("KIGI_API_KEY", k) };
|
||||
}
|
||||
} else {
|
||||
crate::auth::clear_api_key(&kigi_home)
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
// SAFETY: ext_method is single-threaded per agent
|
||||
unsafe { std::env::remove_var("XAI_API_KEY") };
|
||||
unsafe { std::env::remove_var("KIGI_API_KEY") };
|
||||
}
|
||||
ExtMethodResult::success(serde_json::json!({ "ok": true }))
|
||||
.to_ext_response()
|
||||
|
||||
@@ -897,6 +897,36 @@ pub enum SessionUpdate {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
usage: Option<PromptUsage>,
|
||||
},
|
||||
/// Graph mode (`/graph`) progress for the pager's status chip.
|
||||
/// Wire tag `graph_updated`; `status: "cleared"` tells the pager to
|
||||
/// drop its graph state (same sentinel convention as `GoalUpdated`).
|
||||
/// Old pagers degrade to [`Self::Unknown`] silently.
|
||||
GraphUpdated {
|
||||
graph_id: String,
|
||||
objective: String,
|
||||
/// Goal-status vocabulary (`active`, paused family,
|
||||
/// `budget_limited`, `complete`) plus `cleared`.
|
||||
status: String,
|
||||
/// `idle` | `planning` | `executing`.
|
||||
phase: String,
|
||||
plan_version: u32,
|
||||
total_nodes: u32,
|
||||
achieved_nodes: u32,
|
||||
failed_nodes: u32,
|
||||
running_nodes: u32,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
current_node: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
current_node_title: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
token_budget: Option<i64>,
|
||||
#[serde(default)]
|
||||
tokens_spent: i64,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
last_event: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pause_message: Option<String>,
|
||||
},
|
||||
/// Catch-all for unrecognized session update types.
|
||||
/// Allows forward/backward compatibility when variants are added or removed.
|
||||
/// All fields from the unrecognized variant are discarded during deserialization.
|
||||
@@ -2293,3 +2323,88 @@ mod tests {
|
||||
assert!(serde_json::from_str::<SessionUpdate>(missing_stop_reason).is_err());
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod graph_updated_wire_tests {
|
||||
use super::*;
|
||||
|
||||
fn full_graph_updated() -> SessionUpdate {
|
||||
SessionUpdate::GraphUpdated {
|
||||
graph_id: "g-1".into(),
|
||||
objective: "ship it".into(),
|
||||
status: "active".into(),
|
||||
phase: "executing".into(),
|
||||
plan_version: 2,
|
||||
total_nodes: 5,
|
||||
achieved_nodes: 2,
|
||||
failed_nodes: 1,
|
||||
running_nodes: 1,
|
||||
current_node: Some("gn-abc".into()),
|
||||
current_node_title: Some("Node C".into()),
|
||||
token_budget: Some(10_000),
|
||||
tokens_spent: 4_200,
|
||||
last_event: Some("node_achieved".into()),
|
||||
pause_message: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn graph_updated_round_trips_with_snake_case_tag() {
|
||||
let update = full_graph_updated();
|
||||
let json = serde_json::to_value(&update).unwrap();
|
||||
assert_eq!(json["sessionUpdate"], "graph_updated");
|
||||
assert_eq!(json["achieved_nodes"], 2);
|
||||
assert_eq!(json["current_node_title"], "Node C");
|
||||
// Omitted optionals must not serialize at all.
|
||||
assert!(json.get("pause_message").is_none());
|
||||
let back: SessionUpdate = serde_json::from_value(json).unwrap();
|
||||
assert_eq!(back, update);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn graph_updated_minimal_payload_fills_defaults() {
|
||||
// Only the required fields on the wire: every optional absent,
|
||||
// `tokens_spent` relies on #[serde(default)].
|
||||
let json = serde_json::json!({
|
||||
"sessionUpdate": "graph_updated",
|
||||
"graph_id": "g-2",
|
||||
"objective": "o",
|
||||
"status": "cleared",
|
||||
"phase": "idle",
|
||||
"plan_version": 0,
|
||||
"total_nodes": 0,
|
||||
"achieved_nodes": 0,
|
||||
"failed_nodes": 0,
|
||||
"running_nodes": 0,
|
||||
});
|
||||
let update: SessionUpdate = serde_json::from_value(json).unwrap();
|
||||
match update {
|
||||
SessionUpdate::GraphUpdated {
|
||||
status,
|
||||
tokens_spent,
|
||||
current_node,
|
||||
token_budget,
|
||||
pause_message,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(status, "cleared");
|
||||
assert_eq!(tokens_spent, 0, "#[serde(default)] must backfill");
|
||||
assert!(current_node.is_none());
|
||||
assert!(token_budget.is_none());
|
||||
assert!(pause_message.is_none());
|
||||
}
|
||||
other => panic!("expected GraphUpdated, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// An OLD pager (this enum before the variant existed) must degrade
|
||||
/// a graph_updated payload to `Unknown` — pinned by feeding an
|
||||
/// unknown-tag payload through today's enum, which uses the same
|
||||
/// #[serde(other)] mechanism.
|
||||
#[test]
|
||||
fn unknown_tags_still_degrade_gracefully() {
|
||||
let json = serde_json::json!({ "sessionUpdate": "graph_updated_v99", "x": 1 });
|
||||
let update: SessionUpdate = serde_json::from_value(json).unwrap();
|
||||
assert_eq!(update, SessionUpdate::Unknown);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -160,18 +160,22 @@ struct KimiProviderToml {
|
||||
}
|
||||
|
||||
/// Built-in kigi platform a kimi provider duplicates, if any: provider type
|
||||
/// `kimi` is the Kimi Code subscription channel; the two Moonshot open
|
||||
/// platforms are recognized by their fixed production hosts (the same hosts
|
||||
/// `kigi_models::PlatformId::base_url` compiles in).
|
||||
/// `kimi` is the Kimi Code subscription channel; API-key platforms are
|
||||
/// recognized by their production hosts (the same hosts
|
||||
/// `kigi_models::PlatformId::base_url` compiles in — moonshot, openai, and
|
||||
/// every future registry row automatically).
|
||||
fn builtin_platform(provider: &KimiProviderToml) -> Option<PlatformId> {
|
||||
if provider.provider_type == "kimi" {
|
||||
return Some(PlatformId::KimiCode);
|
||||
}
|
||||
match url_host(&provider.base_url) {
|
||||
Some("api.moonshot.cn") => Some(PlatformId::MoonshotCn),
|
||||
Some("api.moonshot.ai") => Some(PlatformId::MoonshotAi),
|
||||
_ => None,
|
||||
let host = url_host(&provider.base_url)?;
|
||||
PlatformId::ALL.into_iter().find(|platform| {
|
||||
if platform.uses_oauth() {
|
||||
return false;
|
||||
}
|
||||
let base = platform.base_url();
|
||||
url_host(&base) == Some(host)
|
||||
})
|
||||
}
|
||||
|
||||
/// Host component of an http(s) URL. `None` for other schemes.
|
||||
@@ -524,10 +528,7 @@ pub fn apply_at(plan: &KimiImportPlan, kigi_home: &Path) -> anyhow::Result<KimiA
|
||||
let _ = std::fs::remove_file(&tmp);
|
||||
return Err(e.into());
|
||||
}
|
||||
if let Err(e) = std::fs::rename(&tmp, &config_path) {
|
||||
let _ = std::fs::remove_file(&tmp);
|
||||
return Err(e.into());
|
||||
}
|
||||
crate::util::fs::replace_file(&tmp, &config_path)?;
|
||||
info!(
|
||||
path = %config_path.display(),
|
||||
added = applied.total_added(),
|
||||
@@ -607,7 +608,7 @@ model = "kimi-for-coding"
|
||||
max_context_size = 262144
|
||||
|
||||
[models.my-openai]
|
||||
provider = "openrouter"
|
||||
provider = "customllm"
|
||||
model = "gpt-x"
|
||||
max_context_size = 128000
|
||||
|
||||
@@ -616,9 +617,9 @@ type = "kimi"
|
||||
base_url = "https://api.kimi.com/coding/v1"
|
||||
api_key = "sk-kimi-secret"
|
||||
|
||||
[providers.openrouter]
|
||||
[providers.customllm]
|
||||
type = "openai_legacy"
|
||||
base_url = "https://openrouter.ai/api/v1"
|
||||
base_url = "https://llm.example.test/v1"
|
||||
api_key = "sk-or-secret"
|
||||
"#,
|
||||
)
|
||||
@@ -707,7 +708,7 @@ api_key = "sk-or-secret"
|
||||
let m = &plan.custom_models[0];
|
||||
assert_eq!(m.alias, "my-openai");
|
||||
assert_eq!(m.model, "gpt-x");
|
||||
assert_eq!(m.base_url, "https://openrouter.ai/api/v1");
|
||||
assert_eq!(m.base_url, "https://llm.example.test/v1");
|
||||
assert_eq!(m.api_key.as_deref(), Some("sk-or-secret"));
|
||||
assert_eq!(m.context_window, Some(128_000));
|
||||
|
||||
@@ -846,7 +847,7 @@ api_key = "sk-ms"
|
||||
assert_eq!(model["model"].as_str().unwrap(), "gpt-x");
|
||||
assert_eq!(
|
||||
model["base_url"].as_str().unwrap(),
|
||||
"https://openrouter.ai/api/v1"
|
||||
"https://llm.example.test/v1"
|
||||
);
|
||||
assert_eq!(model["api_key"].as_str().unwrap(), "sk-or-secret");
|
||||
assert_eq!(model["context_window"].as_integer().unwrap(), 128_000);
|
||||
|
||||
@@ -76,7 +76,7 @@ pub fn map_sampling_err_to_acp(err: SamplingError) -> acp::Error {
|
||||
&& crate::agent::auth_method::has_xai_api_key_env()
|
||||
{
|
||||
format!(
|
||||
"{message}\n\nYou have an API key set (XAI_API_KEY). \
|
||||
"{message}\n\nYou have an API key set (KIGI_API_KEY). \
|
||||
Your cached OAuth session is being used instead. \
|
||||
To use your API key, run `kigi logout` or type /logout in the TUI."
|
||||
)
|
||||
@@ -441,24 +441,32 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
/// Helper: run a closure with XAI_API_KEY temporarily set (or cleared).
|
||||
/// Cleans up even if the closure panics.
|
||||
/// Helper: run a closure with the house BYOK key (KIGI_API_KEY) temporarily
|
||||
/// set (or cleared). Clears every house-key env (KIGI_API_KEY plus the
|
||||
/// back-compat XAI_API_KEY / KIGI_CODE_XAI_API_KEY) so the "no key" case is
|
||||
/// hermetic. Cleans up even if the closure panics.
|
||||
fn with_api_key_env<F: FnOnce()>(key: Option<&str>, f: F) {
|
||||
let prev_house = std::env::var("KIGI_API_KEY").ok();
|
||||
let prev = std::env::var("XAI_API_KEY").ok();
|
||||
let prev_legacy = std::env::var("KIGI_CODE_XAI_API_KEY").ok();
|
||||
// SAFETY: serial_test ensures no concurrent env mutation.
|
||||
unsafe {
|
||||
std::env::remove_var("KIGI_API_KEY");
|
||||
std::env::remove_var("XAI_API_KEY");
|
||||
std::env::remove_var("KIGI_CODE_XAI_API_KEY");
|
||||
if let Some(k) = key {
|
||||
std::env::set_var("XAI_API_KEY", k);
|
||||
std::env::set_var("KIGI_API_KEY", k);
|
||||
}
|
||||
}
|
||||
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
|
||||
// Restore original state.
|
||||
unsafe {
|
||||
std::env::remove_var("KIGI_API_KEY");
|
||||
std::env::remove_var("XAI_API_KEY");
|
||||
std::env::remove_var("KIGI_CODE_XAI_API_KEY");
|
||||
if let Some(v) = prev_house {
|
||||
std::env::set_var("KIGI_API_KEY", v);
|
||||
}
|
||||
if let Some(v) = prev {
|
||||
std::env::set_var("XAI_API_KEY", v);
|
||||
}
|
||||
|
||||
@@ -93,6 +93,14 @@ pub(crate) use types::*;
|
||||
pub use types::{TodoGateDecision, TodoGateReason};
|
||||
#[path = "acp_session_impl/goal.rs"]
|
||||
mod goal;
|
||||
#[path = "acp_session_impl/graph.rs"]
|
||||
mod graph;
|
||||
#[path = "acp_session_impl/graph_optimize.rs"]
|
||||
mod graph_optimize;
|
||||
#[path = "acp_session_impl/graph_replan.rs"]
|
||||
mod graph_replan;
|
||||
#[path = "acp_session_impl/graph_workers.rs"]
|
||||
mod graph_workers;
|
||||
#[path = "acp_session_impl/interjection.rs"]
|
||||
mod interjection;
|
||||
#[path = "acp_session_impl/tool_calls.rs"]
|
||||
@@ -127,7 +135,7 @@ use prompt_build::*;
|
||||
mod session_mode;
|
||||
use session_mode::*;
|
||||
#[path = "acp_session_impl/sampler_turn.rs"]
|
||||
mod sampler_turn;
|
||||
pub(crate) mod sampler_turn;
|
||||
use sampler_turn::*;
|
||||
#[path = "acp_session_impl/tool_dispatch.rs"]
|
||||
mod tool_dispatch;
|
||||
@@ -409,6 +417,21 @@ pub(crate) struct SessionActor {
|
||||
/// [`SessionActor::model_auth_facts`].
|
||||
pub(crate) model_auth_facts:
|
||||
std::cell::RefCell<Option<(String, crate::agent::config::ModelAuthFacts)>>,
|
||||
/// The catalog KEY this session's model was selected by (`{platform}/{model}`
|
||||
/// for a registry model), owned PER SESSION.
|
||||
///
|
||||
/// H4: `SamplingConfig::model` is the bare routing slug, and duplicate slugs
|
||||
/// across an API-key platform and its subscription-OAuth twin
|
||||
/// (`xai`/`xai-grok`, `anthropic`/`claude-pro-max`, `openai`/`openai-codex`)
|
||||
/// are BY DESIGN, so the slug alone cannot name the platform. This used to
|
||||
/// be read from `ModelsManager::current_model_id()` — a single
|
||||
/// PROCESS-GLOBAL cell that Leader mode never writes
|
||||
/// (`agent/handlers/model_switch.rs`) and that is last-writer-wins across
|
||||
/// concurrent sessions, so both collision directions resolved the wrong
|
||||
/// platform: the subscription session lost its resolver (unrecoverable 401
|
||||
/// ~1h in) and the API-key session got the pooled OAuth bearer stamped over
|
||||
/// its own `sk-…` key. Written at spawn and on every `SetSessionModel`.
|
||||
pub(crate) selected_catalog_key: std::cell::RefCell<Option<String>>,
|
||||
/// 401-attribution callback. Joined with the bearer the
|
||||
/// sampler sends on the wire to emit an `auth 401 attribution`
|
||||
/// event at each of the six `OaiCompatClient` 401 arms in
|
||||
@@ -589,6 +612,34 @@ pub(crate) struct SessionActor {
|
||||
/// Goal mode orchestration tracker. Session-scoped state for the
|
||||
/// Design-Execute-Verify loop. Modeled after `plan_mode` above.
|
||||
pub(crate) goal_tracker: Arc<parking_lot::Mutex<crate::session::goal_tracker::GoalTracker>>,
|
||||
/// Whether graph mode (`/graph`) is enabled for this session (feature
|
||||
/// flag `KIGI_GRAPH`). Availability additionally requires the goal
|
||||
/// harness — graph nodes execute as goals.
|
||||
pub(crate) graph_enabled: bool,
|
||||
/// Graph mode orchestration tracker: the deterministic DAG scheduler
|
||||
/// layered over the goal engine. Modeled after `goal_tracker` above;
|
||||
/// all graph state logic lives in `graph_tracker.rs`.
|
||||
pub(crate) graph_tracker: Arc<parking_lot::Mutex<crate::session::graph_tracker::GraphTracker>>,
|
||||
/// Max graph nodes running concurrently (1 = serial G0 behavior).
|
||||
/// Cached at actor construction from `resolve_graph_concurrency`.
|
||||
pub(crate) graph_concurrency: u32,
|
||||
/// Max worker↔verifier rounds per parallel graph node before the
|
||||
/// node fails. Cached at actor construction.
|
||||
pub(crate) graph_node_rounds: u32,
|
||||
/// Max replan passes per graph (0 = replanning off). Cached at
|
||||
/// actor construction.
|
||||
pub(crate) graph_replan_cap: u32,
|
||||
/// Topology optimizer switch (plan-boundary passes; shares the
|
||||
/// replan cap). Cached at actor construction.
|
||||
pub(crate) graph_optimizer_enabled: bool,
|
||||
/// `.kigi` dir at the git root, when the session cwd is in a git
|
||||
/// repo — home of the project-level shared graph projection.
|
||||
pub(crate) graph_project_dir: Option<std::path::PathBuf>,
|
||||
/// Held single-writer lock on the project graph. `Some` while this
|
||||
/// session owns the graph (created or resumed it); dropped on
|
||||
/// `/graph clear`.
|
||||
pub(crate) graph_project_lock:
|
||||
std::cell::RefCell<Option<crate::session::graph_project::ProjectGraphLock>>,
|
||||
/// `task_id`s of background tasks (and monitors) that originated during
|
||||
/// the goal turn — either spawned by the goal model itself or reparented
|
||||
/// from a harness verifier/planner subagent on its exit. Their late
|
||||
@@ -985,6 +1036,9 @@ impl SessionActor {
|
||||
hooks: self.hook_registry.borrow().is_some(),
|
||||
plugins: self.plugin_registry.borrow().is_some(),
|
||||
goal,
|
||||
// Graph rides the goal harness: nodes execute as goals, so
|
||||
// `/graph` is only real when `/goal` is.
|
||||
graph: self.graph_enabled && goal,
|
||||
}
|
||||
}
|
||||
/// Names of every tool registered with the session's tool bridge.
|
||||
@@ -1093,7 +1147,7 @@ fn persist_chat_history_jsonl_sync(session_info: &SessionInfo, conversation: &[C
|
||||
buf.push(b'\n');
|
||||
}
|
||||
std::fs::File::create(&tmp_path)?.write_all(&buf)?;
|
||||
std::fs::rename(&tmp_path, &final_path)?;
|
||||
crate::util::fs::replace_file(&tmp_path, &final_path)?;
|
||||
Ok(())
|
||||
})();
|
||||
if let Err(e) = result {
|
||||
@@ -1217,6 +1271,17 @@ mod rewind_synthetic_turn_tests;
|
||||
#[cfg(test)]
|
||||
#[path = "acp_session_tests/rewrite_zero_turn_prefix_tests.rs"]
|
||||
mod rewrite_zero_turn_prefix_tests;
|
||||
/// The same guard for the model→platform lookup (the dual-credential slug
|
||||
/// collision) and for the stamped aux/summary configs.
|
||||
#[cfg(test)]
|
||||
#[path = "acp_session_tests/session_bearer_leak_platform_tests.rs"]
|
||||
mod session_bearer_leak_platform_tests;
|
||||
/// LEAK guard: the primary Kimi subscription bearer must never ride a request
|
||||
/// to an API-key registry platform's host, while the subscription-OAuth
|
||||
/// platforms keep a live resolver from their OWN pooled manager.
|
||||
#[cfg(test)]
|
||||
#[path = "acp_session_tests/session_bearer_leak_tests.rs"]
|
||||
mod session_bearer_leak_tests;
|
||||
/// Pins the `SubagentFinished` usage-fold attribution gate.
|
||||
#[cfg(test)]
|
||||
#[path = "acp_session_tests/subagent_usage_fold_tests.rs"]
|
||||
@@ -1423,6 +1488,9 @@ mod goal_strategist_e2e_tests;
|
||||
#[path = "acp_session_tests/goal/goal_summarizer_e2e_tests.rs"]
|
||||
mod goal_summarizer_e2e_tests;
|
||||
#[cfg(test)]
|
||||
#[path = "acp_session_tests/graph/graph_e2e_tests.rs"]
|
||||
mod graph_e2e_tests;
|
||||
#[cfg(test)]
|
||||
#[path = "acp_session_tests/idle_resume_tests.rs"]
|
||||
mod idle_resume_tests;
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -1904,6 +1904,29 @@ impl SessionActor {
|
||||
tokens_used,
|
||||
finished_marginal,
|
||||
);
|
||||
// Graph cascade: node goals are armed with the REMAINING graph
|
||||
// budget, so a node-level budget trip is the graph-level trip.
|
||||
if self.graph_harness_enabled() && self.graph_tracker.lock().is_active() {
|
||||
tracing::warn!("graph: node goal budget trip cascades to graph BudgetLimited");
|
||||
{
|
||||
let mut tracker = self.graph_tracker.lock();
|
||||
// Budget integrity: charge the tripped node's partial burn
|
||||
// BEFORE budget_limit clears current_node — otherwise the
|
||||
// top-up arithmetic runs on an under-counted ledger.
|
||||
if let Some(node_id) = tracker.current_node_id().map(str::to_owned) {
|
||||
tracker.charge_node_tokens(&node_id, tokens_used);
|
||||
}
|
||||
tracker.budget_limit();
|
||||
}
|
||||
self.persist_graph_state();
|
||||
self.send_slash_command_output(&format!(
|
||||
"Graph token budget reached ({tokens_used} tokens this node) — graph \
|
||||
stopped. Top up with /graph resume --budget <tokens>, or /graph clear \
|
||||
to abandon."
|
||||
))
|
||||
.await;
|
||||
return true;
|
||||
}
|
||||
self.send_slash_command_output(&format!(
|
||||
"Goal token budget reached ({tokens_used} of {budget} tokens) — goal \
|
||||
stopped. Use /goal clear, then /goal <objective> to start a new one."
|
||||
@@ -2265,18 +2288,21 @@ impl SessionActor {
|
||||
message: Option<String>,
|
||||
) -> bool {
|
||||
let current_tokens = self.chat_state_handle.get_total_tokens().await as i64;
|
||||
{
|
||||
let graph_message = message.clone();
|
||||
let goal_paused = {
|
||||
let mut tracker = self.goal_tracker.lock();
|
||||
if tracker.status() != Some(crate::session::goal_tracker::GoalStatus::Active) {
|
||||
return false;
|
||||
}
|
||||
// The early-return above guarantees `Active`, so the pause
|
||||
// transition always succeeds here.
|
||||
if tracker.status() == Some(crate::session::goal_tracker::GoalStatus::Active) {
|
||||
// Active is guaranteed here, so the transition succeeds.
|
||||
match message {
|
||||
Some(msg) => tracker.pause_with_message(reason, msg),
|
||||
None => tracker.pause(reason),
|
||||
};
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
};
|
||||
if goal_paused {
|
||||
self.clear_pending_classifier_completions();
|
||||
let (tokens_used, finished_marginal) = self.goal_tokens(current_tokens);
|
||||
let notify = self.goal_notify_sender();
|
||||
@@ -2288,7 +2314,31 @@ impl SessionActor {
|
||||
self.emit_event(crate::session::events::Event::GoalAutoPaused {
|
||||
reason: reason.into(),
|
||||
});
|
||||
true
|
||||
}
|
||||
// Graph cascade chokepoint: every goal auto-pause path funnels
|
||||
// through here. It runs REGARDLESS of whether a goal pause
|
||||
// applied — a cancel can land while the graph is Active with no
|
||||
// node goal in the engine (during graph planning, or on the node
|
||||
// boundary between engine reset and goal creation), and the
|
||||
// graph must still lose its self-driving status. `pause` is
|
||||
// Active-only, so double cascades are idempotent.
|
||||
if self.graph_harness_enabled() && self.graph_tracker.lock().is_active() {
|
||||
let node = self
|
||||
.graph_tracker
|
||||
.lock()
|
||||
.current_node_id()
|
||||
.map(str::to_owned);
|
||||
tracing::info!(reason = ?reason, node = ?node, goal_paused, "graph: cascading pause to graph");
|
||||
let detail = match (&node, &graph_message) {
|
||||
(Some(n), Some(msg)) => format!("Node {n} paused: {msg}"),
|
||||
(Some(n), None) => format!("Node {n} paused ({reason:?})"),
|
||||
(None, Some(msg)) => format!("Paused with no node goal in flight: {msg}"),
|
||||
(None, None) => format!("Paused with no node goal in flight ({reason:?})"),
|
||||
};
|
||||
self.graph_tracker.lock().pause_with_message(reason, detail);
|
||||
self.persist_graph_state();
|
||||
}
|
||||
goal_paused
|
||||
}
|
||||
|
||||
/// Match the last assistant message text (via
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,171 @@
|
||||
//! Topology optimizer (G6): a plan-boundary review pass that may issue
|
||||
//! a RESTRICTED set of graph edits — remove false deps (restoring
|
||||
//! parallelism), reorder pending priority, merge tiny nodes, split
|
||||
//! oversized ones — over Waiting/Ready nodes only.
|
||||
//!
|
||||
//! The optimizer changes GRAPH DATA only; the executor stays pure
|
||||
//! deterministic Rust. It fires ① right after initial planning and
|
||||
//! ② at each replan boundary (piggybacked), never mid-execution.
|
||||
//! Applied (non-empty) passes bump `plan_version`, freeze a baseline,
|
||||
//! and consume a slot of the SHARED replan cap; an explicit `[]` is a
|
||||
//! respected no-op consuming nothing. Failure degrades — the current
|
||||
//! graph keeps running. `KIGI_GRAPH_OPTIMIZER=0` disables entirely.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::super::goal_planner::{ChannelSpawner, GoalPlannerSpawner};
|
||||
use super::super::graph_plan;
|
||||
use super::super::graph_planner::{ArtifactPassSpec, run_graph_artifact_pass};
|
||||
use super::SessionActor;
|
||||
|
||||
const OPTIMIZER_PROMPT_TEMPLATE: &str = include_str!("../templates/graph_optimizer_prompt.md");
|
||||
|
||||
impl SessionActor {
|
||||
/// One optimizer pass at a plan boundary. No-op when disabled, when
|
||||
/// the shared cap is exhausted, or when the graph is not Active.
|
||||
pub(super) async fn maybe_optimize_graph(&self) {
|
||||
if !self.graph_optimizer_enabled {
|
||||
return;
|
||||
}
|
||||
let (replan_runs, current_graph, history_text, graph_file, next_version) = {
|
||||
let tracker = self.graph_tracker.lock();
|
||||
let Some(state) = tracker.snapshot() else {
|
||||
return;
|
||||
};
|
||||
if state.status != crate::session::goal_tracker::GoalStatus::Active {
|
||||
return;
|
||||
}
|
||||
let compact: Vec<serde_json::Value> = state
|
||||
.nodes
|
||||
.iter()
|
||||
.map(|n| {
|
||||
serde_json::json!({
|
||||
"id": n.id,
|
||||
"title": n.title,
|
||||
"spec": n.spec,
|
||||
"status": format!("{:?}", n.status),
|
||||
"deps": n.deps.iter().map(|d| d.on.clone()).collect::<Vec<_>>(),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
let history_text = state
|
||||
.history
|
||||
.iter()
|
||||
.rev()
|
||||
.take(12)
|
||||
.map(|h| {
|
||||
format!(
|
||||
"- {:?} {} {}",
|
||||
h.event,
|
||||
h.node_id.as_deref().unwrap_or("-"),
|
||||
h.detail.as_deref().unwrap_or("")
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
(
|
||||
state.replan_runs,
|
||||
serde_json::to_string(&compact).unwrap_or_default(),
|
||||
history_text,
|
||||
tracker
|
||||
.artifacts_dir()
|
||||
.join(format!("optimize.v{}.json", state.plan_version + 1)),
|
||||
state.plan_version + 1,
|
||||
)
|
||||
};
|
||||
if self.graph_replan_cap == 0 || replan_runs >= self.graph_replan_cap {
|
||||
tracing::info!(
|
||||
replan_runs,
|
||||
cap = self.graph_replan_cap,
|
||||
"graph optimizer: shared cap exhausted; skipping pass"
|
||||
);
|
||||
return;
|
||||
}
|
||||
let Some(event_tx) = self.tool_context.subagent_event_tx.clone() else {
|
||||
return;
|
||||
};
|
||||
let objective = self
|
||||
.graph_tracker
|
||||
.lock()
|
||||
.objective()
|
||||
.map(str::to_owned)
|
||||
.unwrap_or_default();
|
||||
let parent_prompt_id = self
|
||||
.current_prompt_id
|
||||
.lock()
|
||||
.expect("current_prompt_id mutex poisoned")
|
||||
.clone();
|
||||
let spawner: Arc<dyn GoalPlannerSpawner> = Arc::new(ChannelSpawner {
|
||||
event_tx,
|
||||
parent_session_id: self.session_id_string(),
|
||||
parent_prompt_id,
|
||||
cwd: Some(self.tool_context.cwd.as_str().to_owned()),
|
||||
role_override: Default::default(),
|
||||
events: Some(self.events.writer()),
|
||||
});
|
||||
let tool_names = self.resolve_inherit_role_tool_names().await;
|
||||
let sections = format!(
|
||||
"OBJECTIVE:\n{objective}\n\nCURRENT GRAPH:\n{current_graph}\n\n\
|
||||
EXECUTION HISTORY:\n{history_text}\n"
|
||||
);
|
||||
tracing::info!(next_version, "graph optimizer: firing");
|
||||
let json = match run_graph_artifact_pass(
|
||||
spawner,
|
||||
ArtifactPassSpec {
|
||||
template: OPTIMIZER_PROMPT_TEMPLATE,
|
||||
sections: §ions,
|
||||
graph_file: &graph_file,
|
||||
tool_names: &tool_names,
|
||||
role: "graph optimizer",
|
||||
},
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(json) => json,
|
||||
Err(reason) => {
|
||||
// Degrade: an enhancement pass never blocks the graph.
|
||||
tracing::warn!(%reason, "graph optimizer: pass failed; keeping current plan");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let existing = self
|
||||
.graph_tracker
|
||||
.lock()
|
||||
.snapshot()
|
||||
.map(|s| s.nodes.clone())
|
||||
.unwrap_or_default();
|
||||
match graph_plan::apply_optimization(&existing, &json) {
|
||||
Ok(None) => {
|
||||
tracing::info!("graph optimizer: no ops (already good)");
|
||||
}
|
||||
Ok(Some(optimized)) => {
|
||||
let n_before = existing.len();
|
||||
let n_after = optimized.len();
|
||||
self.graph_tracker.lock().install_optimized_nodes(optimized);
|
||||
let all_nodes = self
|
||||
.graph_tracker
|
||||
.lock()
|
||||
.snapshot()
|
||||
.map(|s| s.nodes.clone())
|
||||
.unwrap_or_default();
|
||||
if let Err(err) = self.write_graph_baseline(&all_nodes).await {
|
||||
tracing::warn!(%err, "graph optimizer: baseline write failed (audit gap only)");
|
||||
}
|
||||
self.persist_graph_state();
|
||||
tracing::info!(
|
||||
next_version,
|
||||
n_before,
|
||||
n_after,
|
||||
"graph optimizer: plan optimized"
|
||||
);
|
||||
self.send_slash_command_output(&format!(
|
||||
"Graph optimized (v{next_version}): {n_before} → {n_after} node(s)."
|
||||
))
|
||||
.await;
|
||||
}
|
||||
Err(err) => {
|
||||
tracing::warn!(%err, "graph optimizer: ops rejected; keeping current plan");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,236 @@
|
||||
//! Dynamic replan (G3): fold `DISCOVERED:` items surfaced during node
|
||||
//! execution into the graph at dispatch boundaries.
|
||||
//!
|
||||
//! SGH version discipline: the running plan is immutable inside a
|
||||
//! version — a replan appends new nodes (never edits existing ones),
|
||||
//! bumps `plan_version`, and freezes a new immutable baseline. The pass
|
||||
//! is BOUNDED by `KIGI_GRAPH_REPLAN_CAP` (default 3, 0 = off): past the
|
||||
//! cap, discoveries drain to history only, so the graph always
|
||||
//! converges. Replan failure DEGRADES (discoveries kept in history, the
|
||||
//! graph keeps running) — unlike initial planning, a working graph is
|
||||
//! never paused because an enhancement pass failed.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::super::goal_planner::{ChannelSpawner, GoalPlannerSpawner};
|
||||
use super::super::graph_planner::{
|
||||
GRAPH_REPLANNER_SUBAGENT_DESCRIPTION, GraphPlannerOutcome, GraphReplannerInputs,
|
||||
run_graph_replanner,
|
||||
};
|
||||
use super::SessionActor;
|
||||
|
||||
impl SessionActor {
|
||||
/// Replan boundary, called at the top of every `drive_graph`
|
||||
/// iteration. No-op without pending discoveries.
|
||||
pub(super) async fn maybe_replan_graph(&self) {
|
||||
let (pending, replan_runs) = {
|
||||
let tracker = self.graph_tracker.lock();
|
||||
let Some(state) = tracker.snapshot() else {
|
||||
return;
|
||||
};
|
||||
(state.pending_discoveries.clone(), state.replan_runs)
|
||||
};
|
||||
if pending.is_empty() {
|
||||
return;
|
||||
}
|
||||
// Budget first: a budget-dead graph must not spend two replanner
|
||||
// runs right before the dispatch loop trips BudgetLimited. The
|
||||
// discoveries STAY QUEUED (persisted) — a later
|
||||
// `/graph resume --budget` top-up re-enters and replans with
|
||||
// budget actually available.
|
||||
if self.graph_tracker.lock().remaining_budget() == Some(0) {
|
||||
tracing::info!("graph replan: budget exhausted; keeping discoveries queued");
|
||||
return;
|
||||
}
|
||||
let final_achieved = self
|
||||
.graph_tracker
|
||||
.lock()
|
||||
.node(super::super::graph_tracker::FINAL_NODE_ID)
|
||||
.is_some_and(|n| n.status == super::super::graph_tracker::NodeStatus::Achieved);
|
||||
if final_achieved {
|
||||
// The whole-objective gate already passed; late discoveries
|
||||
// (typically from the final verification itself) are
|
||||
// advisory — appending nodes now would ship work the
|
||||
// terminal gate never re-verified.
|
||||
let n = self.graph_tracker.lock().drain_discoveries_to_history();
|
||||
self.persist_graph_state();
|
||||
tracing::info!(
|
||||
drained = n,
|
||||
"graph replan: final already achieved; history only"
|
||||
);
|
||||
return;
|
||||
}
|
||||
if self.graph_replan_cap == 0 {
|
||||
// Feature off: quiet drain (history keeps the audit trail).
|
||||
let n = self.graph_tracker.lock().drain_discoveries_to_history();
|
||||
self.persist_graph_state();
|
||||
tracing::info!(drained = n, "graph replan: disabled (cap 0); history only");
|
||||
return;
|
||||
}
|
||||
if replan_runs >= self.graph_replan_cap {
|
||||
let n = self.graph_tracker.lock().drain_discoveries_to_history();
|
||||
self.persist_graph_state();
|
||||
tracing::warn!(
|
||||
drained = n,
|
||||
replan_runs,
|
||||
cap = self.graph_replan_cap,
|
||||
"graph replan: cap exhausted; discoveries recorded in history only"
|
||||
);
|
||||
self.send_slash_command_output(&format!(
|
||||
"Graph replan cap reached ({replan_runs}/{}); {n} discover{} recorded in \
|
||||
history only — the graph will converge on the current plan.",
|
||||
self.graph_replan_cap,
|
||||
if n == 1 { "y" } else { "ies" },
|
||||
))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(event_tx) = self.tool_context.subagent_event_tx.clone() else {
|
||||
tracing::warn!("graph replan: no subagent coordinator; keeping discoveries queued");
|
||||
return;
|
||||
};
|
||||
let (existing, objective, current_graph, discoveries_text, graph_file, next_version) = {
|
||||
let tracker = self.graph_tracker.lock();
|
||||
let Some(state) = tracker.snapshot() else {
|
||||
return;
|
||||
};
|
||||
let compact: Vec<serde_json::Value> = state
|
||||
.nodes
|
||||
.iter()
|
||||
.map(|n| {
|
||||
serde_json::json!({
|
||||
"id": n.id,
|
||||
"title": n.title,
|
||||
"status": format!("{:?}", n.status),
|
||||
"deps": n.deps.iter().map(|d| d.on.clone()).collect::<Vec<_>>(),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
let discoveries_text = pending
|
||||
.iter()
|
||||
.map(|d| format!("- (from {}) {}", d.from_node, d.description))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
(
|
||||
state.nodes.clone(),
|
||||
state.objective.clone(),
|
||||
serde_json::to_string(&compact).unwrap_or_default(),
|
||||
discoveries_text,
|
||||
tracker
|
||||
.artifacts_dir()
|
||||
.join(format!("replan.v{}.json", state.plan_version + 1)),
|
||||
state.plan_version + 1,
|
||||
)
|
||||
};
|
||||
let parent_prompt_id = self
|
||||
.current_prompt_id
|
||||
.lock()
|
||||
.expect("current_prompt_id mutex poisoned")
|
||||
.clone();
|
||||
let spawner: Arc<dyn GoalPlannerSpawner> = Arc::new(ChannelSpawner {
|
||||
event_tx,
|
||||
parent_session_id: self.session_id_string(),
|
||||
parent_prompt_id,
|
||||
cwd: Some(self.tool_context.cwd.as_str().to_owned()),
|
||||
role_override: Default::default(),
|
||||
events: Some(self.events.writer()),
|
||||
});
|
||||
let tool_names = self.resolve_inherit_role_tool_names().await;
|
||||
|
||||
let mut feedback = String::new();
|
||||
for attempt in 1..=2u32 {
|
||||
tracing::info!(
|
||||
attempt,
|
||||
next_version,
|
||||
role = GRAPH_REPLANNER_SUBAGENT_DESCRIPTION,
|
||||
pending = pending.len(),
|
||||
"graph replan: firing"
|
||||
);
|
||||
match run_graph_replanner(
|
||||
spawner.clone(),
|
||||
&existing,
|
||||
GraphReplannerInputs {
|
||||
objective: &objective,
|
||||
current_graph: ¤t_graph,
|
||||
discoveries: &discoveries_text,
|
||||
feedback: &feedback,
|
||||
graph_file: &graph_file,
|
||||
tool_names: &tool_names,
|
||||
inherit_tool_names: &tool_names,
|
||||
},
|
||||
)
|
||||
.await
|
||||
{
|
||||
GraphPlannerOutcome::Planned(appendix) if appendix.is_empty() => {
|
||||
// Escape hatch: everything already covered. The pass
|
||||
// still counts against the cap.
|
||||
tracing::info!("graph replan: empty appendix (already covered)");
|
||||
{
|
||||
let mut tracker = self.graph_tracker.lock();
|
||||
tracker.drain_discoveries_to_history();
|
||||
if let Some(state) = tracker.snapshot_mut() {
|
||||
state.replan_runs += 1;
|
||||
}
|
||||
}
|
||||
self.persist_graph_state();
|
||||
return;
|
||||
}
|
||||
GraphPlannerOutcome::Planned(appendix) => {
|
||||
let added = appendix.len();
|
||||
self.graph_tracker.lock().append_replan_nodes(appendix);
|
||||
// Freeze the new version's immutable baseline (full
|
||||
// node set post-append; create_new keeps v{N-1}
|
||||
// byte-identical forever).
|
||||
let all_nodes = self
|
||||
.graph_tracker
|
||||
.lock()
|
||||
.snapshot()
|
||||
.map(|s| s.nodes.clone())
|
||||
.unwrap_or_default();
|
||||
if let Err(err) = self.write_graph_baseline(&all_nodes).await {
|
||||
tracing::warn!(%err, "graph replan: baseline write failed (audit gap only)");
|
||||
}
|
||||
self.persist_graph_state();
|
||||
tracing::info!(added, next_version, "graph replan: appendix installed");
|
||||
self.send_slash_command_output(&format!(
|
||||
"Graph replanned (v{next_version}): {added} node(s) added from \
|
||||
discovered work."
|
||||
))
|
||||
.await;
|
||||
// Plan-boundary optimizer pass ② (piggybacked on the
|
||||
// replan version boundary).
|
||||
self.maybe_optimize_graph().await;
|
||||
return;
|
||||
}
|
||||
GraphPlannerOutcome::Invalid { reason } if attempt == 1 => {
|
||||
tracing::warn!(%reason, "graph replan: invalid appendix; retrying with feedback");
|
||||
feedback = format!(
|
||||
"Your previous replan JSON failed validation:\n{reason}\n\
|
||||
Rewrite the file fixing exactly this."
|
||||
);
|
||||
}
|
||||
GraphPlannerOutcome::Invalid { reason }
|
||||
| GraphPlannerOutcome::FailClosed { reason } => {
|
||||
// Degrade, never pause a working graph for a failed
|
||||
// enhancement pass. The run still counts.
|
||||
tracing::warn!(%reason, "graph replan: failed; draining discoveries to history");
|
||||
{
|
||||
let mut tracker = self.graph_tracker.lock();
|
||||
tracker.drain_discoveries_to_history();
|
||||
if let Some(state) = tracker.snapshot_mut() {
|
||||
state.replan_runs += 1;
|
||||
}
|
||||
}
|
||||
self.persist_graph_state();
|
||||
self.send_slash_command_output(&format!(
|
||||
"Graph replan failed ({reason}); discovered work recorded in \
|
||||
history only."
|
||||
))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,979 @@
|
||||
//! Parallel graph-node execution: worker/verifier subagent pairs.
|
||||
//!
|
||||
//! In parallel mode (`KIGI_GRAPH_CONCURRENCY > 1` with ≥2 `Ready`
|
||||
//! nodes) a node does NOT run on the session goal engine — it runs as a
|
||||
//! harness-internal `general-purpose` subagent (the implementer toolset)
|
||||
//! in its OWN git worktree, adversarially checked by a read-only
|
||||
//! verifier subagent, with a bounded worker↔verifier round loop
|
||||
//! (`graph_node_rounds`). Achieved nodes merge back into the main tree
|
||||
//! SEQUENTIALLY via `kigi_workspace`'s 3-way `apply_worktree`; a merge
|
||||
//! conflict fails the node (its dependents block; other chains
|
||||
//! continue). The terminal `gn-final` node always runs serially on the
|
||||
//! full goal engine because it depends on every other node.
|
||||
//!
|
||||
//! Known ceiling: a worker round that outlives the foreground subagent
|
||||
//! await budget (default 600s) is cancelled and counted as a failed
|
||||
//! round with an explicit gap; the next round resumes the same child
|
||||
//! session. Fetching results from auto-backgrounded children would need
|
||||
//! completed-store plumbing — deferred until real usage demands it.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use kigi_tools::implementations::kigi::task::types::{
|
||||
SubagentEvent, SubagentRequest, SubagentRuntimeOverrides,
|
||||
};
|
||||
|
||||
use super::SessionActor;
|
||||
|
||||
const WORKER_PROMPT_TEMPLATE: &str = include_str!("../templates/graph_node_worker_prompt.md");
|
||||
const VERIFIER_PROMPT_TEMPLATE: &str = include_str!("../templates/graph_node_verifier_prompt.md");
|
||||
|
||||
// Terminal-contract parsing
|
||||
|
||||
/// The worker's parsed claim, from the trailing `NODE_RESULT:` line.
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub(crate) enum WorkerClaim {
|
||||
Done { summary: String },
|
||||
Blocked { reason: String },
|
||||
Unparseable,
|
||||
}
|
||||
|
||||
/// Drop ``` fenced code blocks so a QUOTED marker (the templates
|
||||
/// themselves contain fenced `NODE_RESULT:`/`NODE_VERDICT:` examples a
|
||||
/// child may echo) can never be parsed as the real terminal line.
|
||||
fn strip_fenced_blocks(output: &str) -> String {
|
||||
let mut kept = String::with_capacity(output.len());
|
||||
let mut in_fence = false;
|
||||
for line in output.lines() {
|
||||
if line.trim_start().starts_with("```") {
|
||||
in_fence = !in_fence;
|
||||
continue;
|
||||
}
|
||||
if !in_fence {
|
||||
kept.push_str(line);
|
||||
kept.push('\n');
|
||||
}
|
||||
}
|
||||
kept
|
||||
}
|
||||
|
||||
/// The LAST line that STARTS with `marker` (line-anchored — a
|
||||
/// mid-sentence mention never matches), plus everything after it.
|
||||
fn last_marker_line(output: &str, marker: &str) -> Option<(String, String)> {
|
||||
let lines: Vec<&str> = output.lines().collect();
|
||||
let idx = lines
|
||||
.iter()
|
||||
.rposition(|l| l.trim_start().trim_start_matches('`').starts_with(marker))?;
|
||||
let value = lines[idx]
|
||||
.trim_start()
|
||||
.trim_start_matches('`')
|
||||
.trim_start_matches(marker)
|
||||
.trim()
|
||||
.trim_matches('`')
|
||||
.to_owned();
|
||||
let tail = lines[idx + 1..].join("\n").trim().to_owned();
|
||||
Some((value, tail))
|
||||
}
|
||||
|
||||
/// Parse the last line-anchored `NODE_RESULT:` marker outside fenced
|
||||
/// blocks; text after the marker line is the summary/reason.
|
||||
/// Fail-closed: no marker ⇒ `Unparseable`.
|
||||
pub(crate) fn parse_worker_claim(output: &str) -> WorkerClaim {
|
||||
let stripped = strip_fenced_blocks(output);
|
||||
let Some((value, tail)) = last_marker_line(&stripped, "NODE_RESULT:") else {
|
||||
return WorkerClaim::Unparseable;
|
||||
};
|
||||
match value.as_str() {
|
||||
"done" => WorkerClaim::Done { summary: tail },
|
||||
"blocked" => WorkerClaim::Blocked {
|
||||
reason: if tail.is_empty() {
|
||||
"no reason given".to_owned()
|
||||
} else {
|
||||
tail
|
||||
},
|
||||
},
|
||||
_ => WorkerClaim::Unparseable,
|
||||
}
|
||||
}
|
||||
|
||||
/// The verifier's parsed verdict, from the trailing `NODE_VERDICT:` line.
|
||||
/// Fail-closed: anything unparseable is `NotAchieved` with that fact as
|
||||
/// the gap.
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub(crate) enum NodeVerdict {
|
||||
Achieved,
|
||||
NotAchieved { gaps: Vec<String> },
|
||||
}
|
||||
|
||||
pub(crate) fn parse_node_verdict(output: &str) -> NodeVerdict {
|
||||
let stripped = strip_fenced_blocks(output);
|
||||
let Some((value, tail)) = last_marker_line(&stripped, "NODE_VERDICT:") else {
|
||||
return NodeVerdict::NotAchieved {
|
||||
gaps: vec!["verifier response lacked a NODE_VERDICT line".to_owned()],
|
||||
};
|
||||
};
|
||||
match value.as_str() {
|
||||
"achieved" => NodeVerdict::Achieved,
|
||||
"not_achieved" => {
|
||||
let gaps: Vec<String> = tail
|
||||
.lines()
|
||||
.map(str::trim)
|
||||
.filter(|l| !l.is_empty() && *l != "GAPS:")
|
||||
.map(|l| l.trim_start_matches('-').trim().to_owned())
|
||||
.filter(|l| !l.is_empty())
|
||||
.collect();
|
||||
NodeVerdict::NotAchieved {
|
||||
gaps: if gaps.is_empty() {
|
||||
vec!["verifier rejected without naming gaps".to_owned()]
|
||||
} else {
|
||||
gaps
|
||||
},
|
||||
}
|
||||
}
|
||||
other => NodeVerdict::NotAchieved {
|
||||
gaps: vec![format!("unrecognized verdict token {other:?}")],
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Line-anchored `DISCOVERED:` items outside fenced blocks. The
|
||||
/// placeholder filter (`<`) drops template echoes ("<one-line
|
||||
/// description …>") a child may parrot back.
|
||||
pub(crate) fn parse_discovered_lines(output: &str) -> Vec<String> {
|
||||
strip_fenced_blocks(output)
|
||||
.lines()
|
||||
.filter_map(|l| {
|
||||
l.trim_start()
|
||||
.trim_start_matches('`')
|
||||
.strip_prefix("DISCOVERED:")
|
||||
})
|
||||
.map(str::trim)
|
||||
// Placeholder-echo defense: drop only the templates' literal
|
||||
// "<one-line description …>" shape, not every '<' (legit Rust
|
||||
// discoveries mention generics like Vec<String>).
|
||||
.filter(|d| !d.is_empty() && !d.starts_with('<'))
|
||||
.map(str::to_owned)
|
||||
.collect()
|
||||
}
|
||||
|
||||
// Spawner seam (mockable in tests)
|
||||
|
||||
pub(crate) struct WorkerSpawnSpec {
|
||||
pub prompt: String,
|
||||
pub description: String,
|
||||
/// Explicit child cwd (verifiers run in the worker's worktree).
|
||||
pub cwd: Option<String>,
|
||||
/// Mint an isolated worktree for the child (first worker round).
|
||||
pub isolation_worktree: bool,
|
||||
/// Resume a prior child session (later worker rounds keep context
|
||||
/// AND the worktree).
|
||||
pub resume_from: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct WorkerSpawnOutcome {
|
||||
pub success: bool,
|
||||
pub cancelled: bool,
|
||||
pub backgrounded: bool,
|
||||
pub output: String,
|
||||
pub error: Option<String>,
|
||||
pub child_session_id: String,
|
||||
pub tokens_used: u64,
|
||||
pub worktree_path: Option<String>,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
pub(crate) trait GraphWorkerSpawner: Send + Sync {
|
||||
/// Spawn one child and await its terminal result. `Err` = transport
|
||||
/// failure (coordinator gone).
|
||||
async fn spawn(&self, id: &str, spec: WorkerSpawnSpec) -> Result<WorkerSpawnOutcome, String>;
|
||||
/// Best-effort cancel of a still-running child (budget overrun).
|
||||
async fn cancel(&self, subagent_id: &str);
|
||||
}
|
||||
|
||||
/// Production spawner: raw harness-internal `SubagentEvent::Spawn`,
|
||||
/// exactly the goal-classifier wire (`surface_completion: false`, no
|
||||
/// fork), plus worktree isolation / cwd override for node work.
|
||||
pub(crate) struct GraphWorkerChannelSpawner {
|
||||
pub event_tx: tokio::sync::mpsc::UnboundedSender<SubagentEvent>,
|
||||
pub parent_session_id: String,
|
||||
pub parent_prompt_id: Option<String>,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl GraphWorkerSpawner for GraphWorkerChannelSpawner {
|
||||
async fn spawn(&self, id: &str, spec: WorkerSpawnSpec) -> Result<WorkerSpawnOutcome, String> {
|
||||
let (result_tx, result_rx) = tokio::sync::oneshot::channel();
|
||||
let request = SubagentRequest {
|
||||
id: id.to_string(),
|
||||
prompt: spec.prompt,
|
||||
description: spec.description,
|
||||
// The implementer toolset (full read/edit/bash inventory);
|
||||
// verifier read-only-ness is prompt-enforced, same as the
|
||||
// goal skeptic panel.
|
||||
subagent_type: "general-purpose".to_string(),
|
||||
parent_session_id: self.parent_session_id.clone(),
|
||||
parent_prompt_id: self.parent_prompt_id.clone(),
|
||||
resume_from: spec.resume_from,
|
||||
cwd: spec.cwd,
|
||||
runtime_overrides: SubagentRuntimeOverrides {
|
||||
isolation: spec
|
||||
.isolation_worktree
|
||||
.then_some(kigi_tool_types::SubagentIsolationMode::Worktree),
|
||||
..Default::default()
|
||||
},
|
||||
run_in_background: false,
|
||||
// Harness-internal: never surfaces to the model's idle reminder.
|
||||
surface_completion: false,
|
||||
fork_context: false,
|
||||
result_tx,
|
||||
};
|
||||
if self
|
||||
.event_tx
|
||||
.send(SubagentEvent::Spawn(Box::new(request)))
|
||||
.is_err()
|
||||
{
|
||||
return Err("subagent coordinator channel closed".to_owned());
|
||||
}
|
||||
let result = result_rx
|
||||
.await
|
||||
.map_err(|_| "subagent result channel dropped".to_owned())?;
|
||||
Ok(WorkerSpawnOutcome {
|
||||
success: result.success,
|
||||
cancelled: result.cancelled,
|
||||
backgrounded: result.backgrounded,
|
||||
output: result.output.to_string(),
|
||||
error: result.error.clone(),
|
||||
child_session_id: result.child_session_id.clone(),
|
||||
tokens_used: result.tokens_used,
|
||||
worktree_path: result.worktree_path.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn cancel(&self, subagent_id: &str) {
|
||||
use kigi_tools::implementations::kigi::task::types::{
|
||||
SubagentCancelRequest, SubagentCancelTarget,
|
||||
};
|
||||
let (respond_to, ack) = tokio::sync::oneshot::channel();
|
||||
let _ = self
|
||||
.event_tx
|
||||
.send(SubagentEvent::Cancel(SubagentCancelRequest {
|
||||
target: SubagentCancelTarget::SubagentId(subagent_id.to_string()),
|
||||
respond_to,
|
||||
}));
|
||||
let _ = ack.await;
|
||||
}
|
||||
}
|
||||
|
||||
// Per-node bounded closed loop
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct NodeRunReport {
|
||||
pub node_id: String,
|
||||
pub achieved: bool,
|
||||
/// Worker summary on success; failure reason otherwise.
|
||||
pub detail: String,
|
||||
pub rounds: u32,
|
||||
pub tokens_used: i64,
|
||||
pub worktree_path: Option<String>,
|
||||
/// Last worker child session id (audit link, stored on the node).
|
||||
pub worker_session_id: Option<String>,
|
||||
/// `DISCOVERED:` items surfaced by the workers/verifiers (deduped).
|
||||
pub discoveries: Vec<String>,
|
||||
}
|
||||
|
||||
fn worker_prompt(node_objective: &str, gaps: &[String]) -> String {
|
||||
let mut p = String::with_capacity(WORKER_PROMPT_TEMPLATE.len() + node_objective.len() + 256);
|
||||
p.push_str(WORKER_PROMPT_TEMPLATE);
|
||||
p.push_str("\n\nNODE OBJECTIVE:\n");
|
||||
p.push_str(node_objective);
|
||||
if !gaps.is_empty() {
|
||||
p.push_str("\n\nGAPS (from the previous verification round — close exactly these):\n");
|
||||
for gap in gaps {
|
||||
p.push_str("- ");
|
||||
p.push_str(gap);
|
||||
p.push('\n');
|
||||
}
|
||||
}
|
||||
p
|
||||
}
|
||||
|
||||
fn verifier_prompt(node_objective: &str, worker_summary: &str) -> String {
|
||||
// Neutralize terminal-contract tokens in the worker-controlled
|
||||
// summary so a lazy/adversarial claim cannot smuggle marker lines
|
||||
// into the verifier's context.
|
||||
let safe_summary = worker_summary
|
||||
.replace("NODE_VERDICT", "NODE-VERDICT")
|
||||
.replace("NODE_RESULT", "NODE-RESULT")
|
||||
.replace("DISCOVERED", "DISCOVERED-");
|
||||
format!(
|
||||
"{VERIFIER_PROMPT_TEMPLATE}\n\nNODE OBJECTIVE (the contract to judge):\n{node_objective}\n\n\
|
||||
IMPLEMENTER'S CLAIM (audit it, do not trust it):\n{safe_summary}\n"
|
||||
)
|
||||
}
|
||||
|
||||
/// Drive one node through bounded worker↔verifier rounds. Never panics;
|
||||
/// every failure path returns a `NodeRunReport` with a precise reason.
|
||||
pub(crate) async fn run_node_to_verdict(
|
||||
spawner: &Arc<dyn GraphWorkerSpawner>,
|
||||
node_id: &str,
|
||||
node_objective: &str,
|
||||
rounds_cap: u32,
|
||||
) -> NodeRunReport {
|
||||
let mut tokens: i64 = 0;
|
||||
let mut discoveries: Vec<String> = Vec::new();
|
||||
let mut gaps: Vec<String> = Vec::new();
|
||||
let mut resume_from: Option<String> = None;
|
||||
let mut worktree_path: Option<String> = None;
|
||||
let mut last_gaps_summary = String::new();
|
||||
|
||||
for round in 1..=rounds_cap {
|
||||
let spawn_id = format!("graph-{node_id}-w{round}-{}", uuid::Uuid::now_v7());
|
||||
let spec_was_isolated = resume_from.is_none();
|
||||
let spec = WorkerSpawnSpec {
|
||||
prompt: worker_prompt(node_objective, &gaps),
|
||||
description: format!("graph node worker ({node_id})"),
|
||||
cwd: None,
|
||||
// Fresh worktree only on the first round; resumes reuse it.
|
||||
isolation_worktree: spec_was_isolated,
|
||||
resume_from: resume_from.clone(),
|
||||
};
|
||||
tracing::info!(%node_id, round, resumed = resume_from.is_some(), "graph worker: round start");
|
||||
let outcome = match spawner.spawn(&spawn_id, spec).await {
|
||||
Ok(o) => o,
|
||||
Err(err) => {
|
||||
return NodeRunReport {
|
||||
node_id: node_id.to_owned(),
|
||||
achieved: false,
|
||||
detail: format!("worker transport failure: {err}"),
|
||||
rounds: round,
|
||||
tokens_used: tokens,
|
||||
worktree_path,
|
||||
worker_session_id: resume_from,
|
||||
discoveries: discoveries.clone(),
|
||||
};
|
||||
}
|
||||
};
|
||||
let round_requested_isolation = spec_was_isolated;
|
||||
tokens = tokens.saturating_add(outcome.tokens_used as i64);
|
||||
if outcome.worktree_path.is_some() {
|
||||
worktree_path = outcome.worktree_path.clone();
|
||||
}
|
||||
// Adopt-guard: an in-band spawn failure carries an EMPTY child id;
|
||||
// adopting it would make the next round a fresh UNISOLATED spawn
|
||||
// in the shared tree while verify/merge still target the stale
|
||||
// worktree. Keep the last valid id (or None ⇒ re-mint isolation).
|
||||
if !outcome.child_session_id.is_empty() {
|
||||
resume_from = Some(outcome.child_session_id.clone());
|
||||
}
|
||||
for d in parse_discovered_lines(&outcome.output) {
|
||||
if !discoveries.contains(&d) {
|
||||
discoveries.push(d);
|
||||
}
|
||||
}
|
||||
|
||||
if outcome.cancelled {
|
||||
return NodeRunReport {
|
||||
node_id: node_id.to_owned(),
|
||||
achieved: false,
|
||||
detail: "worker cancelled".to_owned(),
|
||||
rounds: round,
|
||||
tokens_used: tokens,
|
||||
worktree_path,
|
||||
worker_session_id: resume_from,
|
||||
discoveries,
|
||||
};
|
||||
}
|
||||
if outcome.backgrounded {
|
||||
// Ceiling (see module doc): cancel the runaway child and
|
||||
// burn the round; the resume keeps its context.
|
||||
tracing::warn!(%node_id, round, "graph worker: exceeded foreground await budget; cancelling round");
|
||||
// Cancel by the SPAWN REQUEST id — the coordinator's cancel
|
||||
// maps are keyed by it, not by the child session id.
|
||||
spawner.cancel(&spawn_id).await;
|
||||
gaps = vec![
|
||||
"the previous round exceeded the foreground time budget and was cancelled; \
|
||||
split the remaining work into smaller, faster steps"
|
||||
.to_owned(),
|
||||
];
|
||||
last_gaps_summary = gaps.join("; ");
|
||||
continue;
|
||||
}
|
||||
if !outcome.success {
|
||||
let err = outcome.error.unwrap_or_else(|| "unknown error".to_owned());
|
||||
tracing::warn!(%node_id, round, %err, "graph worker: round failed");
|
||||
gaps = vec![format!("the previous round failed with an error: {err}")];
|
||||
last_gaps_summary = gaps.join("; ");
|
||||
continue;
|
||||
}
|
||||
|
||||
// Isolation guard: a SUCCESSFUL first round that came back with
|
||||
// no worktree means isolation silently degraded (non-git dir,
|
||||
// worktree creation failure, or the snapshot-disposal flag
|
||||
// deleted it before we saw it). Running parallel writers in the
|
||||
// shared tree — or merging a disposed tree — is never OK.
|
||||
if round_requested_isolation && outcome.worktree_path.is_none() {
|
||||
return NodeRunReport {
|
||||
node_id: node_id.to_owned(),
|
||||
achieved: false,
|
||||
detail: "worktree isolation unavailable for this node (non-git directory, \
|
||||
worktree creation failure, or KIGI_SUBAGENT_WORKTREE_SNAPSHOT \
|
||||
disposal); parallel execution requires isolation"
|
||||
.to_owned(),
|
||||
rounds: round,
|
||||
tokens_used: tokens,
|
||||
worktree_path: None,
|
||||
worker_session_id: resume_from,
|
||||
discoveries: discoveries.clone(),
|
||||
};
|
||||
}
|
||||
|
||||
match parse_worker_claim(&outcome.output) {
|
||||
WorkerClaim::Blocked { reason } => {
|
||||
return NodeRunReport {
|
||||
node_id: node_id.to_owned(),
|
||||
achieved: false,
|
||||
detail: format!("worker reported blocked: {reason}"),
|
||||
rounds: round,
|
||||
tokens_used: tokens,
|
||||
worktree_path,
|
||||
worker_session_id: resume_from,
|
||||
discoveries: discoveries.clone(),
|
||||
};
|
||||
}
|
||||
WorkerClaim::Unparseable => {
|
||||
gaps = vec![
|
||||
"the previous round's final message lacked the required NODE_RESULT line"
|
||||
.to_owned(),
|
||||
];
|
||||
last_gaps_summary = gaps.join("; ");
|
||||
continue;
|
||||
}
|
||||
WorkerClaim::Done { summary } => {
|
||||
let verify_id = format!("graph-{node_id}-v{round}-{}", uuid::Uuid::now_v7());
|
||||
let verify_spec = WorkerSpawnSpec {
|
||||
prompt: verifier_prompt(node_objective, &summary),
|
||||
description: format!("graph node verifier ({node_id})"),
|
||||
// The verifier inspects the worker's worktree.
|
||||
cwd: worktree_path.clone(),
|
||||
isolation_worktree: false,
|
||||
resume_from: None,
|
||||
};
|
||||
let verdict = match spawner.spawn(&verify_id, verify_spec).await {
|
||||
Ok(v) => {
|
||||
tokens = tokens.saturating_add(v.tokens_used as i64);
|
||||
for d in parse_discovered_lines(&v.output) {
|
||||
if !discoveries.contains(&d) {
|
||||
discoveries.push(d);
|
||||
}
|
||||
}
|
||||
if v.success {
|
||||
parse_node_verdict(&v.output)
|
||||
} else {
|
||||
// Fail CLOSED: an unverified claim never passes.
|
||||
NodeVerdict::NotAchieved {
|
||||
gaps: vec![format!(
|
||||
"verifier run failed ({}); the claim is unverified",
|
||||
v.error.unwrap_or_else(|| "unknown error".to_owned())
|
||||
)],
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => NodeVerdict::NotAchieved {
|
||||
gaps: vec![format!("verifier transport failure: {err}")],
|
||||
},
|
||||
};
|
||||
match verdict {
|
||||
NodeVerdict::Achieved => {
|
||||
tracing::info!(%node_id, round, tokens, "graph worker: node verified achieved");
|
||||
return NodeRunReport {
|
||||
node_id: node_id.to_owned(),
|
||||
achieved: true,
|
||||
detail: summary,
|
||||
rounds: round,
|
||||
tokens_used: tokens,
|
||||
worktree_path,
|
||||
worker_session_id: resume_from,
|
||||
discoveries: discoveries.clone(),
|
||||
};
|
||||
}
|
||||
NodeVerdict::NotAchieved { gaps: new_gaps } => {
|
||||
tracing::info!(%node_id, round, gap_count = new_gaps.len(), "graph worker: verifier rejected round");
|
||||
last_gaps_summary = new_gaps.join("; ");
|
||||
gaps = new_gaps;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
NodeRunReport {
|
||||
node_id: node_id.to_owned(),
|
||||
achieved: false,
|
||||
detail: format!(
|
||||
"verification rejected after {rounds_cap} rounds; last gaps: {last_gaps_summary}"
|
||||
),
|
||||
rounds: rounds_cap,
|
||||
tokens_used: tokens,
|
||||
worktree_path,
|
||||
worker_session_id: resume_from,
|
||||
discoveries,
|
||||
}
|
||||
}
|
||||
|
||||
/// `git rev-parse HEAD` of `dir`, `None` outside a git repo.
|
||||
async fn git_head(dir: &std::path::Path) -> Option<String> {
|
||||
let output = tokio::process::Command::new("git")
|
||||
.arg("-C")
|
||||
.arg(dir)
|
||||
.args(["rev-parse", "HEAD"])
|
||||
.output()
|
||||
.await
|
||||
.ok()?;
|
||||
output
|
||||
.status
|
||||
.success()
|
||||
.then(|| String::from_utf8_lossy(&output.stdout).trim().to_owned())
|
||||
}
|
||||
|
||||
// SessionActor integration
|
||||
|
||||
impl SessionActor {
|
||||
/// Production worker spawner wired to this session's coordinator.
|
||||
fn graph_worker_spawner(&self) -> Option<Arc<dyn GraphWorkerSpawner>> {
|
||||
let event_tx = self.tool_context.subagent_event_tx.clone()?;
|
||||
let parent_prompt_id = self
|
||||
.current_prompt_id
|
||||
.lock()
|
||||
.expect("current_prompt_id mutex poisoned")
|
||||
.clone();
|
||||
Some(Arc::new(GraphWorkerChannelSpawner {
|
||||
event_tx,
|
||||
parent_session_id: self.session_id_string(),
|
||||
parent_prompt_id,
|
||||
}))
|
||||
}
|
||||
|
||||
/// Run one parallel batch of `Ready` nodes to their verdicts, then
|
||||
/// merge achieved worktrees back SEQUENTIALLY in batch order. All
|
||||
/// tracker mutations + persistence happen here; the caller re-reads
|
||||
/// the tracker afterwards.
|
||||
pub(super) async fn run_graph_parallel_batch(&self, node_ids: Vec<String>) {
|
||||
let Some(spawner) = self.graph_worker_spawner() else {
|
||||
tracing::error!("graph batch: no subagent coordinator; pausing graph");
|
||||
self.graph_tracker.lock().pause_with_message(
|
||||
crate::session::goal_tracker::GoalPauseReason::Infra,
|
||||
"No subagent coordinator available for parallel execution".to_owned(),
|
||||
);
|
||||
self.persist_graph_state();
|
||||
return;
|
||||
};
|
||||
// Compose objectives + mark Running under one lock pass.
|
||||
let mut jobs: Vec<(String, String)> = Vec::with_capacity(node_ids.len());
|
||||
{
|
||||
let mut tracker = self.graph_tracker.lock();
|
||||
let Some(snapshot) = tracker.snapshot() else {
|
||||
return;
|
||||
};
|
||||
let total = snapshot.nodes.len();
|
||||
let objective = snapshot.objective.clone();
|
||||
for id in &node_ids {
|
||||
if let Some(pos) = snapshot.nodes.iter().position(|n| n.id == *id) {
|
||||
let node = &snapshot.nodes[pos];
|
||||
jobs.push((
|
||||
id.clone(),
|
||||
super::graph::node_goal_objective(&objective, node, pos + 1, total),
|
||||
));
|
||||
}
|
||||
}
|
||||
for (id, _) in &jobs {
|
||||
tracker.mark_node_running(id, String::new());
|
||||
}
|
||||
// `current_node` means "the node on the serial goal engine";
|
||||
// batch nodes are tracked by their own Running status.
|
||||
if let Some(s) = tracker.snapshot_mut() {
|
||||
s.current_node = None;
|
||||
}
|
||||
}
|
||||
self.persist_graph_state();
|
||||
// Merge-base integrity: if the main repo HEAD moves during the
|
||||
// batch (external commit), apply_worktree would diff against the
|
||||
// wrong base and silently reverse-apply those commits. Capture
|
||||
// HEAD now; every merge re-checks it.
|
||||
let head_at_fanout = git_head(self.tool_context.cwd.as_path()).await;
|
||||
let rounds_cap = self.graph_node_rounds;
|
||||
tracing::info!(batch = jobs.len(), rounds_cap, "graph batch: fan-out");
|
||||
|
||||
let reports = futures::future::join_all(jobs.iter().map(|(id, objective)| {
|
||||
let spawner = spawner.clone();
|
||||
async move { run_node_to_verdict(&spawner, id, objective, rounds_cap).await }
|
||||
}))
|
||||
.await;
|
||||
|
||||
// Sequential merge + tracker resolution in batch order.
|
||||
let mut achieved = 0usize;
|
||||
let mut failed = 0usize;
|
||||
for report in reports {
|
||||
// Stamp the worker session id for audit (goal_id slot).
|
||||
if let Some(worker_id) = &report.worker_session_id
|
||||
&& let Some(node) = self
|
||||
.graph_tracker
|
||||
.lock()
|
||||
.snapshot_mut()
|
||||
.and_then(|s| s.nodes.iter_mut().find(|n| n.id == report.node_id))
|
||||
{
|
||||
node.goal_id = Some(worker_id.clone());
|
||||
}
|
||||
if !report.discoveries.is_empty() {
|
||||
// A failed node's discoveries are still real work.
|
||||
let ds: Vec<crate::session::graph_tracker::Discovery> = report
|
||||
.discoveries
|
||||
.iter()
|
||||
.map(|d| crate::session::graph_tracker::Discovery {
|
||||
from_node: report.node_id.clone(),
|
||||
description: d.clone(),
|
||||
})
|
||||
.collect();
|
||||
self.graph_tracker.lock().queue_discoveries(ds);
|
||||
}
|
||||
if !report.achieved {
|
||||
failed += 1;
|
||||
{
|
||||
let mut tracker = self.graph_tracker.lock();
|
||||
// Budget integrity: a failed node's tokens were still
|
||||
// spent — charge them before failing the node.
|
||||
tracker.charge_node_tokens(&report.node_id, report.tokens_used);
|
||||
tracker.mark_node_failed(&report.node_id, report.detail.clone());
|
||||
}
|
||||
self.persist_graph_state();
|
||||
continue;
|
||||
}
|
||||
match self
|
||||
.merge_node_worktree(
|
||||
&report.node_id,
|
||||
report.worktree_path.as_deref(),
|
||||
head_at_fanout.as_deref(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(()) => {
|
||||
achieved += 1;
|
||||
self.graph_tracker.lock().mark_node_achieved(
|
||||
&report.node_id,
|
||||
report.rounds,
|
||||
report.tokens_used,
|
||||
);
|
||||
}
|
||||
Err(detail) => {
|
||||
failed += 1;
|
||||
self.graph_tracker
|
||||
.lock()
|
||||
.mark_node_failed(&report.node_id, detail);
|
||||
}
|
||||
}
|
||||
self.persist_graph_state();
|
||||
}
|
||||
tracing::info!(achieved, failed, "graph batch: settled");
|
||||
self.send_slash_command_output(&format!(
|
||||
"Graph batch settled: {achieved} node(s) achieved, {failed} failed."
|
||||
))
|
||||
.await;
|
||||
}
|
||||
|
||||
/// Merge one achieved node's worktree back into the main tree with
|
||||
/// the 3-way apply. `None` worktree (isolation soft-fallback) means
|
||||
/// the worker already wrote in the shared tree — nothing to merge.
|
||||
pub(super) async fn merge_node_worktree(
|
||||
&self,
|
||||
node_id: &str,
|
||||
worktree_path: Option<&str>,
|
||||
expected_main_head: Option<&str>,
|
||||
) -> Result<(), String> {
|
||||
let Some(worktree_path) = worktree_path else {
|
||||
tracing::warn!(
|
||||
%node_id,
|
||||
"graph merge: worker ran without worktree isolation (soft fallback); nothing to merge"
|
||||
);
|
||||
return Ok(());
|
||||
};
|
||||
use kigi_workspace::worktree::{
|
||||
ApplyMode, ApplyWorktreeRequest, ApplyWorktreeResponse, apply_worktree,
|
||||
};
|
||||
// apply_worktree diffs against the main repo HEAD AT APPLY TIME;
|
||||
// if HEAD moved since fan-out, that diff would silently
|
||||
// reverse-apply the external commits. Fail the node loudly.
|
||||
if let Some(expected) = expected_main_head {
|
||||
let current = git_head(self.tool_context.cwd.as_path()).await;
|
||||
if current.as_deref() != Some(expected) {
|
||||
return Err(format!(
|
||||
"main repository HEAD moved during the batch (was {expected}, now {}); \
|
||||
merge aborted for safety — /graph resume re-runs the node",
|
||||
current.as_deref().unwrap_or("unknown")
|
||||
));
|
||||
}
|
||||
}
|
||||
let request = ApplyWorktreeRequest {
|
||||
session_id: self.session_id_string(),
|
||||
worktree_path: worktree_path.to_owned(),
|
||||
mode: ApplyMode::Merge,
|
||||
};
|
||||
match apply_worktree(&request).await {
|
||||
Ok(ApplyWorktreeResponse::Success { files, .. }) => {
|
||||
tracing::info!(%node_id, files = files.len(), "graph merge: applied");
|
||||
// Storage discipline: the changes now live in the main
|
||||
// tree, so the worktree is dead weight — remove it.
|
||||
// Best-effort (a failed removal only leaks disk, never
|
||||
// progress) but always logged. Failed nodes KEEP their
|
||||
// worktree for postmortem.
|
||||
if let Err(err) = kigi_workspace::worktree::remove_subagent_worktree(
|
||||
std::path::Path::new(worktree_path),
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(%node_id, %err, "graph merge: worktree cleanup failed");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Ok(ApplyWorktreeResponse::Conflicts { conflicts, .. }) => {
|
||||
let names: Vec<String> = conflicts.iter().map(|c| c.path.clone()).collect();
|
||||
tracing::warn!(%node_id, ?names, "graph merge: conflicts; failing node");
|
||||
Err(format!("merge conflict in: {}", names.join(", ")))
|
||||
}
|
||||
Err(err) => {
|
||||
tracing::warn!(%node_id, %err, "graph merge: apply failed");
|
||||
Err(format!("worktree apply failed: {err}"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn worker_claim_parses_done_blocked_and_garbage() {
|
||||
assert_eq!(
|
||||
parse_worker_claim("work...\nNODE_RESULT: done\nBuilt X; tests pass."),
|
||||
WorkerClaim::Done {
|
||||
summary: "Built X; tests pass.".to_owned()
|
||||
}
|
||||
);
|
||||
assert_eq!(
|
||||
parse_worker_claim("NODE_RESULT: blocked\nno compiler available"),
|
||||
WorkerClaim::Blocked {
|
||||
reason: "no compiler available".to_owned()
|
||||
}
|
||||
);
|
||||
assert_eq!(
|
||||
parse_worker_claim("all done, promise!"),
|
||||
WorkerClaim::Unparseable
|
||||
);
|
||||
// Last marker wins (a quoted earlier marker cannot spoof).
|
||||
assert_eq!(
|
||||
parse_worker_claim(
|
||||
"NODE_RESULT: done\nold\n...more work...\nNODE_RESULT: blocked\nreal"
|
||||
),
|
||||
WorkerClaim::Blocked {
|
||||
reason: "real".to_owned()
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn verdict_parses_achieved_gaps_and_fails_closed() {
|
||||
assert_eq!(
|
||||
parse_node_verdict("checked\nNODE_VERDICT: achieved"),
|
||||
NodeVerdict::Achieved
|
||||
);
|
||||
assert_eq!(
|
||||
parse_node_verdict(
|
||||
"NODE_VERDICT: not_achieved\nGAPS:\n- test suite not run\n- claim B unverified"
|
||||
),
|
||||
NodeVerdict::NotAchieved {
|
||||
gaps: vec![
|
||||
"test suite not run".to_owned(),
|
||||
"claim B unverified".to_owned()
|
||||
]
|
||||
}
|
||||
);
|
||||
assert!(matches!(
|
||||
parse_node_verdict("looks good to me"),
|
||||
NodeVerdict::NotAchieved { gaps } if gaps[0].contains("lacked a NODE_VERDICT")
|
||||
));
|
||||
assert!(matches!(
|
||||
parse_node_verdict("NODE_VERDICT: maybe"),
|
||||
NodeVerdict::NotAchieved { gaps } if gaps[0].contains("unrecognized verdict")
|
||||
));
|
||||
assert!(matches!(
|
||||
parse_node_verdict("NODE_VERDICT: not_achieved"),
|
||||
NodeVerdict::NotAchieved { gaps } if gaps[0].contains("without naming gaps")
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fenced_template_echo_cannot_spoof_markers() {
|
||||
// A worker echoing the template's fenced examples must stay
|
||||
// unparseable; only its own line-anchored terminal marker counts.
|
||||
let echoed = "Here is my plan:\n```\nNODE_RESULT: done\n```\nstill working...";
|
||||
assert_eq!(parse_worker_claim(echoed), WorkerClaim::Unparseable);
|
||||
let real = "```\nNODE_RESULT: blocked\n```\n...work...\nNODE_RESULT: done\nall built";
|
||||
assert_eq!(
|
||||
parse_worker_claim(real),
|
||||
WorkerClaim::Done {
|
||||
summary: "all built".to_owned()
|
||||
}
|
||||
);
|
||||
// Mid-sentence mention is not a marker (line-anchored scan).
|
||||
assert_eq!(
|
||||
parse_worker_claim("I will print NODE_RESULT: done when finished"),
|
||||
WorkerClaim::Unparseable
|
||||
);
|
||||
// Same discipline for the verifier.
|
||||
assert!(matches!(
|
||||
parse_node_verdict("quoting:\n```\nNODE_VERDICT: achieved\n```\nhmm"),
|
||||
NodeVerdict::NotAchieved { .. }
|
||||
));
|
||||
}
|
||||
|
||||
struct MockSpawner {
|
||||
replies: std::sync::Mutex<std::collections::VecDeque<WorkerSpawnOutcome>>,
|
||||
specs: std::sync::Mutex<Vec<(String, Option<String>, bool)>>,
|
||||
cancels: std::sync::Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl GraphWorkerSpawner for MockSpawner {
|
||||
async fn spawn(
|
||||
&self,
|
||||
id: &str,
|
||||
spec: WorkerSpawnSpec,
|
||||
) -> Result<WorkerSpawnOutcome, String> {
|
||||
self.specs.lock().unwrap().push((
|
||||
id.to_owned(),
|
||||
spec.resume_from.clone(),
|
||||
spec.isolation_worktree,
|
||||
));
|
||||
Ok(self
|
||||
.replies
|
||||
.lock()
|
||||
.unwrap()
|
||||
.pop_front()
|
||||
.expect("unexpected extra spawn"))
|
||||
}
|
||||
async fn cancel(&self, subagent_id: &str) {
|
||||
self.cancels.lock().unwrap().push(subagent_id.to_owned());
|
||||
}
|
||||
}
|
||||
|
||||
fn outcome(output: &str) -> WorkerSpawnOutcome {
|
||||
WorkerSpawnOutcome {
|
||||
success: true,
|
||||
cancelled: false,
|
||||
backgrounded: false,
|
||||
output: output.to_owned(),
|
||||
error: None,
|
||||
child_session_id: "child-1".to_owned(),
|
||||
tokens_used: 5,
|
||||
worktree_path: Some("/wt".to_owned()),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn backgrounded_round_cancels_by_spawn_id_and_resumes_next_round() {
|
||||
let mut bg = outcome("");
|
||||
bg.backgrounded = true;
|
||||
let replies = std::collections::VecDeque::from(vec![
|
||||
bg, // round 1: budget overrun
|
||||
outcome("NODE_RESULT: done\nfinished"), // round 2: worker done
|
||||
outcome("NODE_VERDICT: achieved"), // round 2: verifier
|
||||
]);
|
||||
// Keep a concrete handle for assertions; hand the trait object in.
|
||||
let mock = Arc::new(MockSpawner {
|
||||
replies: std::sync::Mutex::new(replies),
|
||||
specs: std::sync::Mutex::new(Vec::new()),
|
||||
cancels: std::sync::Mutex::new(Vec::new()),
|
||||
});
|
||||
let spawner: Arc<dyn GraphWorkerSpawner> = mock.clone();
|
||||
let report = run_node_to_verdict(&spawner, "gn-x", "do x", 3).await;
|
||||
assert!(report.achieved, "{}", report.detail);
|
||||
assert_eq!(report.rounds, 2, "backgrounded round burned, retry won");
|
||||
assert_eq!(report.tokens_used, 15, "all three spawns charged");
|
||||
|
||||
let cancels = mock.cancels.lock().unwrap().clone();
|
||||
assert_eq!(cancels.len(), 1, "runaway child cancelled once");
|
||||
assert!(
|
||||
cancels[0].starts_with("graph-gn-x-w1-"),
|
||||
"cancel must target the SPAWN REQUEST id (coordinator map key), got {}",
|
||||
cancels[0]
|
||||
);
|
||||
let specs = mock.specs.lock().unwrap().clone();
|
||||
assert_eq!(specs.len(), 3);
|
||||
assert!(
|
||||
specs[0].1.is_none() && specs[0].2,
|
||||
"round 1: fresh + isolated"
|
||||
);
|
||||
assert_eq!(
|
||||
specs[1].1.as_deref(),
|
||||
Some("child-1"),
|
||||
"round 2 resumes the backgrounded child's session"
|
||||
);
|
||||
assert!(!specs[1].2, "resume never re-mints isolation");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_child_id_is_never_adopted_as_resume_target() {
|
||||
// In-band spawn failure: success=false, child_session_id="".
|
||||
let failed = WorkerSpawnOutcome {
|
||||
success: false,
|
||||
cancelled: false,
|
||||
backgrounded: false,
|
||||
output: String::new(),
|
||||
error: Some("boom".to_owned()),
|
||||
child_session_id: String::new(),
|
||||
tokens_used: 0,
|
||||
worktree_path: None,
|
||||
};
|
||||
let replies = std::collections::VecDeque::from(vec![
|
||||
failed, // round 1: in-band failure
|
||||
outcome("NODE_RESULT: done\nfinished"), // round 2: worker (fresh, isolated)
|
||||
outcome("NODE_VERDICT: achieved"), // round 2: verifier
|
||||
]);
|
||||
let mock = Arc::new(MockSpawner {
|
||||
replies: std::sync::Mutex::new(replies),
|
||||
specs: std::sync::Mutex::new(Vec::new()),
|
||||
cancels: std::sync::Mutex::new(Vec::new()),
|
||||
});
|
||||
let spawner: Arc<dyn GraphWorkerSpawner> = mock.clone();
|
||||
let report = run_node_to_verdict(&spawner, "gn-y", "do y", 3).await;
|
||||
assert!(report.achieved, "{}", report.detail);
|
||||
let specs = mock.specs.lock().unwrap().clone();
|
||||
assert!(
|
||||
specs[1].1.is_none(),
|
||||
"an empty child id must NOT be adopted; retry is a fresh spawn"
|
||||
);
|
||||
assert!(
|
||||
specs[1].2,
|
||||
"fresh retry re-mints worktree isolation (no unisolated escape)"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn successful_isolated_round_without_worktree_fails_the_node() {
|
||||
let mut no_wt = outcome("NODE_RESULT: done\nfinished");
|
||||
no_wt.worktree_path = None;
|
||||
let replies = std::collections::VecDeque::from(vec![no_wt]);
|
||||
let mock = Arc::new(MockSpawner {
|
||||
replies: std::sync::Mutex::new(replies),
|
||||
specs: std::sync::Mutex::new(Vec::new()),
|
||||
cancels: std::sync::Mutex::new(Vec::new()),
|
||||
});
|
||||
let spawner: Arc<dyn GraphWorkerSpawner> = mock.clone();
|
||||
let report = run_node_to_verdict(&spawner, "gn-z", "do z", 3).await;
|
||||
assert!(!report.achieved);
|
||||
assert!(
|
||||
report.detail.contains("isolation unavailable"),
|
||||
"{}",
|
||||
report.detail
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -5,12 +5,18 @@ impl SessionActor {
|
||||
pub(super) async fn handle_set_session_model(
|
||||
&self,
|
||||
sampling_config: kigi_sampler::SamplerConfig,
|
||||
catalog_key: Option<String>,
|
||||
use_concise: bool,
|
||||
apply_prompt_override: bool,
|
||||
skip_prompt_rewrite: bool,
|
||||
auto_compact_threshold_percent: u8,
|
||||
) -> Result<acp::ModelId, acp::Error> {
|
||||
let model_id = acp::ModelId::new(sampling_config.model.clone());
|
||||
// H4: record the picker's catalog KEY as this SESSION's own selection.
|
||||
// `sampling_config.model` is the ambiguous bare slug; the key is what
|
||||
// disambiguates an API-key platform from its subscription-OAuth twin,
|
||||
// and it must never come from the process-global `current_model_id()`.
|
||||
*self.selected_catalog_key.borrow_mut() = catalog_key;
|
||||
let new_context_window = self.compaction.context_window_override.unwrap_or_else(|| {
|
||||
std::num::NonZeroU64::new(sampling_config.context_window).unwrap_or_else(|| {
|
||||
std::num::NonZeroU64::new(DEFAULT_CONTEXT_WINDOW)
|
||||
@@ -53,22 +59,24 @@ impl SessionActor {
|
||||
temperature: sampling_config.temperature,
|
||||
top_p: sampling_config.top_p,
|
||||
api_backend: sampling_config.api_backend.clone(),
|
||||
chat_compat: sampling_config.chat_compat,
|
||||
extra_headers: sampling_config.extra_headers.clone(),
|
||||
context_window: new_context_window,
|
||||
reasoning_effort: sampling_config.reasoning_effort,
|
||||
stream_tool_calls: Some(sampling_config.stream_tool_calls),
|
||||
});
|
||||
let existing = self.chat_state_handle.get_credentials().await;
|
||||
let session_key = self
|
||||
.auth_manager
|
||||
.as_ref()
|
||||
.and_then(|am| am.current_or_expired().map(|a| a.key));
|
||||
// Read the session bearer from the switched-to model's OWN manager: a
|
||||
// grok model reads the xai-grok token (used only to classify the
|
||||
// credential's auth_type here), never the Kimi one. Kimi / non-oauth
|
||||
// models resolve to the primary — byte-identical.
|
||||
let session_key = self.session_credential_for_model(&sampling_config.model);
|
||||
self.chat_state_handle
|
||||
.update_credentials(kigi_chat_state::Credentials {
|
||||
api_key: sampling_config.api_key.clone(),
|
||||
auth_type: crate::agent::config::resolve_chat_state_auth_type(
|
||||
sampling_config.model.as_str(),
|
||||
session_key.as_deref(),
|
||||
session_key.as_ref(),
|
||||
existing.auth_type,
|
||||
),
|
||||
alpha_test_key: existing.alpha_test_key,
|
||||
|
||||
@@ -624,10 +624,20 @@ impl SessionActor {
|
||||
let resolved_describe = self
|
||||
.resolve_aux_sampler_config(&self.image_description_model)
|
||||
.await;
|
||||
// LEAK 1b: the aux bearer_resolver is decided at the chokepoint from the
|
||||
// IMAGE-DESCRIBE model's own platform + endpoint and passed in
|
||||
// explicitly, so an aux model on another provider can never inherit the
|
||||
// session (Kimi) resolver and have its own key overwritten on the aux
|
||||
// host. The `None` fallback yields the SESSION config verbatim, whose
|
||||
// own resolver must stay as-is.
|
||||
let describe_resolver = resolved_describe
|
||||
.as_ref()
|
||||
.map(|cfg| self.aux_bearer_resolver(&self.image_description_model, &cfg.base_url));
|
||||
let (describe_model, sampler_config) =
|
||||
crate::agent::config::finalize_image_describe_sampler_config(
|
||||
resolved_describe,
|
||||
&active_session_config,
|
||||
describe_resolver.flatten(),
|
||||
Some(self.max_retries),
|
||||
);
|
||||
let client = kigi_sampler::SamplingClient::new(sampler_config).map_err(|e| {
|
||||
|
||||
@@ -300,10 +300,10 @@ pub(super) async fn run_session(
|
||||
SessionActor::maybe_start_running_task(session.clone(), completion_tx
|
||||
.clone()). await; } SessionCommand::SessionMode { session_mode, responds_to }
|
||||
=> { session.handle_session_mode(session_mode). await; let _ = responds_to
|
||||
.send(()); } SessionCommand::SetSessionModel { sampling_config, use_concise,
|
||||
apply_prompt_override, skip_prompt_rewrite, auto_compact_threshold_percent,
|
||||
responds_to } => { let updated_model_id = session
|
||||
.handle_set_session_model(sampling_config, use_concise,
|
||||
.send(()); } SessionCommand::SetSessionModel { sampling_config, catalog_key,
|
||||
use_concise, apply_prompt_override, skip_prompt_rewrite,
|
||||
auto_compact_threshold_percent, responds_to } => { let updated_model_id =
|
||||
session.handle_set_session_model(sampling_config, catalog_key, use_concise,
|
||||
apply_prompt_override, skip_prompt_rewrite, auto_compact_threshold_percent).
|
||||
await; let _ = responds_to.send(updated_model_id); }
|
||||
SessionCommand::RebuildAgentForDefinition { definition, responds_to } => {
|
||||
@@ -319,11 +319,24 @@ pub(super) async fn run_session(
|
||||
.signals_handle().set_primary_model(& model_name); cfg.model = model_name
|
||||
.clone(); cfg.extra_headers.extend(extra_headers); if let Some(cw) =
|
||||
context_window && session.compaction.context_window_override.is_none() { cfg
|
||||
.context_window = cw; } session.chat_state_handle
|
||||
.context_window = cw; } let override_base_url = cfg
|
||||
.base_url.clone(); session.chat_state_handle
|
||||
.update_sampling_config(cfg); let existing = session.chat_state_handle
|
||||
.get_credentials(). await; if let Some(r) = crate
|
||||
::agent::config::try_resolve_model_credentials(model_name.as_str(), existing
|
||||
.api_key.as_deref()) { session.chat_state_handle
|
||||
.get_credentials(). await;
|
||||
// H-c: the rename makes the session's own selected catalog
|
||||
// key stale unless it still names this model; a stale key is
|
||||
// exactly what the model→platform rule must not trust.
|
||||
session.retain_selected_catalog_key_for(& model_name);
|
||||
// The override model routes to the SAME endpoint the session
|
||||
// already had; ask the chokepoint whether that endpoint takes
|
||||
// a session credential rather than re-offering the key
|
||||
// already in chat state. The platform comes from the session's
|
||||
// OWN lookup so this and every later turn agree.
|
||||
let override_session_key = session.credential_authority()
|
||||
.credential_for(session.model_platform(model_name.as_str()), &
|
||||
override_base_url); if let Some(r) = crate
|
||||
::agent::config::try_resolve_model_credentials(model_name.as_str(),
|
||||
override_session_key.as_ref()) { session.chat_state_handle
|
||||
.update_credentials(kigi_chat_state::Credentials { api_key : r.api_key,
|
||||
auth_type : r.auth_type, alpha_test_key : existing.alpha_test_key, }); } session.model_auth_facts
|
||||
.replace(None); } } SessionCommand::GetCurrentModel { responds_to } => { let
|
||||
|
||||
@@ -29,37 +29,105 @@ pub(super) fn is_auth_tool_error(err: &kigi_tool_runtime::ToolError) -> bool {
|
||||
/// Gate inputs bundled with the composed decision so the 401-recovery log can
|
||||
/// report the components.
|
||||
#[derive(Clone, Copy)]
|
||||
struct SessionTokenAuthGate {
|
||||
pub(crate) struct SessionTokenAuthGate {
|
||||
is_session_based: bool,
|
||||
model_byok: crate::agent::auth_method::ModelByok,
|
||||
/// Whether the request targets a first-party host. Lets an `Unknown`
|
||||
/// BYOK status still refresh against the first-party cli-chat-proxy hosts without
|
||||
/// risking a session-token leak to a third-party BYOK endpoint.
|
||||
endpoint_is_first_party: bool,
|
||||
/// WHICH credential this model's platform/endpoint pair accepts, per the
|
||||
/// single credential chokepoint
|
||||
/// ([`crate::auth::credential_authority::CredentialAuthority::credential_class`]).
|
||||
/// `None` for every API-key registry platform, which keeps the primary Kimi
|
||||
/// bearer off `api.deepseek.com` / `api.openai.com` / … .
|
||||
credential_class: crate::auth::credential_authority::CredentialClass,
|
||||
}
|
||||
impl SessionTokenAuthGate {
|
||||
/// Single place `is_session_based` / `endpoint_is_first_party` are derived,
|
||||
/// so all call sites assemble the gate identically.
|
||||
fn new(
|
||||
/// 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::model_platform`]) that
|
||||
/// [`SessionActor::auth_manager_for_endpoint`] uses, so the gate's verdict
|
||||
/// and the manager actually wrapped as the bearer resolver can never
|
||||
/// disagree. `authority` is that same chokepoint, so the gate cannot answer
|
||||
/// the endpoint question differently from the manager routing.
|
||||
pub(crate) fn new(
|
||||
auth_method_id: Option<&acp::AuthMethodId>,
|
||||
model_byok: crate::agent::auth_method::ModelByok,
|
||||
base_url: &str,
|
||||
model_platform: Option<kigi_models::PlatformId>,
|
||||
authority: &crate::auth::credential_authority::CredentialAuthority,
|
||||
) -> Self {
|
||||
Self {
|
||||
// L13: a model whose OWN credential is a pooled subscription-OAuth
|
||||
// session is session-based BY ITSELF, whatever the primary ACP
|
||||
// method is. A user logged in with an API-KEY platform (e.g.
|
||||
// `deepseek`) who selects a `claude-pro-max/*` model still gets that
|
||||
// platform's pooled bearer as the request's `api_key` — without this
|
||||
// term the gate would be inactive, so the config would carry NO
|
||||
// resolver: the token freezes at selection time and the session dies
|
||||
// with an unrecoverable 401 once it expires (~1h). The outer
|
||||
// `credential_class` conjunct keeps this confined to that
|
||||
// platform's own host.
|
||||
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_platform.is_some_and(|p| p.oauth().is_some()),
|
||||
model_byok,
|
||||
endpoint_is_first_party: crate::util::is_first_party_url(base_url),
|
||||
credential_class: authority.credential_class(model_platform, base_url),
|
||||
}
|
||||
}
|
||||
fn active(self) -> bool {
|
||||
pub(crate) fn active(self) -> bool {
|
||||
crate::agent::auth_method::session_token_auth_gate(
|
||||
self.is_session_based,
|
||||
self.model_byok,
|
||||
self.endpoint_is_first_party,
|
||||
self.credential_class,
|
||||
)
|
||||
}
|
||||
}
|
||||
/// THE aux / summary `bearer_resolver` rule, stated ONCE.
|
||||
///
|
||||
/// `SamplingClient::post` REPLACES the request's auth header from the resolver,
|
||||
/// so an aux model on a different provider would have its own correctly-resolved
|
||||
/// key overwritten by the session bearer ON THE AUX HOST. An OAuth aux model
|
||||
/// gets a live resolver over ITS OWN pooled manager (keeping mid-session
|
||||
/// refresh); a first-party aux model gets the primary's, but ONLY when the
|
||||
/// session-token gate is active; everything else gets `None`, so the aux model's
|
||||
/// own key survives to the wire.
|
||||
///
|
||||
/// M3 — the FIRST-PARTY case honours the gate, which is what the old "copy
|
||||
/// `active_session_config.bearer_resolver`" shape did implicitly: that field is
|
||||
/// `None` whenever the gate is inactive. Without it, a BYOK / api-key session
|
||||
/// with a `[model.*]` aux entry carrying its OWN `env_key` on the session's own
|
||||
/// coding endpoint has that key REPLACED on the wire by the primary bearer on
|
||||
/// every image-describe / auto-mode-classifier / summary request. A
|
||||
/// subscription-OAuth aux model is deliberately NOT gated this way: its pooled
|
||||
/// token IS its credential, and withholding the resolver only costs it
|
||||
/// mid-session refresh (L13).
|
||||
///
|
||||
/// Shared by [`SessionActor::aux_bearer_resolver`] and
|
||||
/// `MvpAgent::summary_bearer_resolver`: the summary client is built by the
|
||||
/// AGENT, not the session actor, and its own private copy of this rule is
|
||||
/// exactly how it stayed ungated after M3 closed the session-actor side.
|
||||
pub(crate) fn aux_bearer_resolver_for(
|
||||
authority: &crate::auth::credential_authority::CredentialAuthority,
|
||||
auth_method_id: Option<&acp::AuthMethodId>,
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
model_byok: crate::agent::auth_method::ModelByok,
|
||||
base_url: &str,
|
||||
) -> Option<kigi_sampler::SharedBearerResolver> {
|
||||
let is_primary_channel = platform.is_none_or(|p| p.oauth().is_none());
|
||||
if is_primary_channel
|
||||
&& !SessionTokenAuthGate::new(auth_method_id, model_byok, base_url, platform, authority)
|
||||
.active()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
authority.bearer_resolver_for(platform, base_url)
|
||||
}
|
||||
/// Run a tool call; on an auth-shaped failure, attempt recovery via
|
||||
/// `AuthManager` and one retry. When `shared_recovery` is `Some`, concurrent
|
||||
/// 401s in the same batch deduplicate via `OnceCell::get_or_init`.
|
||||
@@ -103,6 +171,31 @@ where
|
||||
result
|
||||
}
|
||||
}
|
||||
/// Wraps an [`AuthManager`](crate::auth::AuthManager) as a sampler
|
||||
/// [`BearerResolver`](kigi_sampler::BearerResolver), resolving the live
|
||||
/// (current-or-expired) bearer at request time. Shared by
|
||||
/// [`SessionActor::reconstruct_full_config`] (the session model) and the
|
||||
/// aux-model bearer routing
|
||||
/// ([`CredentialAuthority::bearer_resolver_for`](crate::auth::credential_authority::CredentialAuthority::bearer_resolver_for))
|
||||
/// so both wrap ONE definition. SECURITY: the bearer is resolved per request
|
||||
/// and never logged.
|
||||
pub(crate) struct AuthManagerBearerResolver(pub(crate) std::sync::Arc<crate::auth::AuthManager>);
|
||||
impl std::fmt::Debug for AuthManagerBearerResolver {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("AuthManagerBearerResolver").finish()
|
||||
}
|
||||
}
|
||||
impl kigi_sampler::BearerResolver for AuthManagerBearerResolver {
|
||||
fn current_bearer(&self) -> Option<String> {
|
||||
self.0.current_or_expired().map(|a| a.key)
|
||||
}
|
||||
}
|
||||
/// Wrap `am` as a shared sampler bearer resolver.
|
||||
pub(crate) fn auth_manager_bearer_resolver(
|
||||
am: std::sync::Arc<crate::auth::AuthManager>,
|
||||
) -> kigi_sampler::SharedBearerResolver {
|
||||
std::sync::Arc::new(AuthManagerBearerResolver(am))
|
||||
}
|
||||
impl SessionActor {
|
||||
pub(super) async fn prepare_tool_definitions_timed(&self) -> (Vec<ToolDefinition>, u64) {
|
||||
let mcp_wait_start = std::time::Instant::now();
|
||||
@@ -146,8 +239,8 @@ impl SessionActor {
|
||||
let plan_active = self.plan_mode.lock().is_active();
|
||||
filter_cursor_tools_by_plan_mode(defs, plan_active)
|
||||
}
|
||||
/// Memoized per-model [`ModelAuthFacts`](crate::agent::config::ModelAuthFacts),
|
||||
/// keyed by `model_id`.
|
||||
/// Memoized per-model [`ModelAuthFacts`](crate::agent::config::ModelAuthFacts)
|
||||
/// for the SESSION's own model, keyed by `model_id`.
|
||||
///
|
||||
/// A fresh `Unknown` (config currently unparseable) falls back to the last
|
||||
/// definite value for the same `model_id` rather than demoting a live session
|
||||
@@ -156,6 +249,32 @@ impl SessionActor {
|
||||
/// `model_id`, keying on `model_id` alone is insufficient — each
|
||||
/// model/credential chokepoint must clear this memo (`replace(None)`).
|
||||
pub(super) fn model_auth_facts(&self, model_id: &str) -> crate::agent::config::ModelAuthFacts {
|
||||
self.resolve_auth_facts(model_id, true)
|
||||
}
|
||||
/// [`Self::model_auth_facts`] for a model that is NOT the session's own — an
|
||||
/// AUX / summary / image-describe slug.
|
||||
///
|
||||
/// Identical resolution, but it NEVER WRITES the slot. The memo is a SINGLE
|
||||
/// slot: when the aux path shared it, one classifier or image-describe call
|
||||
/// evicted the session model's entry, and (a) the next
|
||||
/// [`Self::reconstruct_full_config`] paid another `load_effective_config()`
|
||||
/// + `resolve_model_list()` — the per-turn disk read M7/M9 removed — while
|
||||
/// (b) a transient `Unknown` for the SESSION model then had no same-`model_id`
|
||||
/// definite value to fall back to, so it degraded to `endpoint_is_first_party`
|
||||
/// — `false` for every subscription-OAuth host, costing the session its
|
||||
/// `bearer_resolver` and 401ing unrecoverably ~1h in (the failure L13
|
||||
/// prevents). Reading a matching entry is still allowed: it can only hit when
|
||||
/// the slot already names this same slug.
|
||||
fn aux_model_auth_facts(&self, model_id: &str) -> crate::agent::config::ModelAuthFacts {
|
||||
self.resolve_auth_facts(model_id, false)
|
||||
}
|
||||
/// Shared body of [`Self::model_auth_facts`] / [`Self::aux_model_auth_facts`].
|
||||
/// `memoize` is the ONLY difference, so the two can never resolve differently.
|
||||
fn resolve_auth_facts(
|
||||
&self,
|
||||
model_id: &str,
|
||||
memoize: bool,
|
||||
) -> crate::agent::config::ModelAuthFacts {
|
||||
use crate::agent::auth_method::ModelByok;
|
||||
if let Some((cached_id, facts)) = self.model_auth_facts.borrow().as_ref()
|
||||
&& cached_id == model_id
|
||||
@@ -172,17 +291,187 @@ impl SessionActor {
|
||||
}
|
||||
return fresh;
|
||||
}
|
||||
if memoize {
|
||||
*self.model_auth_facts.borrow_mut() = Some((model_id.to_string(), fresh));
|
||||
}
|
||||
fresh
|
||||
}
|
||||
/// Gate inputs for `model_id` routed to `base_url`. See
|
||||
/// Gate inputs for the SESSION model `model_id` routed to `base_url`. See
|
||||
/// [`crate::agent::auth_method::session_token_auth_gate`] for the rationale
|
||||
/// (`base_url` keeps an `Unknown` BYOK status refreshable only
|
||||
/// against first-party xAI hosts).
|
||||
fn auth_gate(&self, model_id: &str, base_url: &str) -> SessionTokenAuthGate {
|
||||
let byok = self.model_auth_facts(model_id).byok;
|
||||
let auth_method = self.auth_method_id.load();
|
||||
SessionTokenAuthGate::new(auth_method.as_deref(), byok, base_url)
|
||||
SessionTokenAuthGate::new(
|
||||
auth_method.as_deref(),
|
||||
byok,
|
||||
base_url,
|
||||
self.model_platform(model_id),
|
||||
&self.credential_authority(),
|
||||
)
|
||||
}
|
||||
/// This session's credential chokepoint: its EFFECTIVE endpoints (so a
|
||||
/// managed `[endpoints] coding_api_base_url` deployment keeps the session
|
||||
/// bearer — H3) plus its primary manager, which the authority keeps
|
||||
/// private. Every inference-auth question this actor asks goes through it.
|
||||
pub(crate) fn credential_authority(
|
||||
&self,
|
||||
) -> crate::auth::credential_authority::CredentialAuthority {
|
||||
crate::auth::credential_authority::CredentialAuthority::new(
|
||||
self.models_manager.endpoints(),
|
||||
self.auth_manager.clone(),
|
||||
)
|
||||
}
|
||||
/// The [`AuthManager`](crate::auth::AuthManager) that governs INFERENCE auth
|
||||
/// for the routing slug `model` against the endpoint the request will
|
||||
/// ACTUALLY be sent to, from the ONE chokepoint
|
||||
/// ([`crate::auth::credential_authority::CredentialAuthority`]).
|
||||
///
|
||||
/// A subscription-OAuth platform routes to ITS OWN scope-keyed pooled
|
||||
/// manager; `kimi-code` and a platform-less model on the session's own
|
||||
/// coding endpoint route to the primary; every API-key registry platform,
|
||||
/// and any endpoint that is neither, routes to `None` — fail fast, never a
|
||||
/// silent fallback to the primary. `None` also for a BYOK / test session
|
||||
/// with no primary.
|
||||
///
|
||||
/// Callers pass the LIVE sampling config's `base_url` so the manager, the
|
||||
/// gate and the wire can never be resolved against three different endpoints
|
||||
/// (an `OverrideModelName` session keeps its original `base_url` under a
|
||||
/// routing name absent from the catalog). A `model`-only sibling that
|
||||
/// re-derived the endpoint from the CATALOG instead used to exist beside
|
||||
/// this; it had zero callers and was deleted rather than left as a second,
|
||||
/// unexercised way to answer the same question (`lib.rs`'s
|
||||
/// `#![allow(dead_code)]` means such a helper raises no warning).
|
||||
pub(super) fn auth_manager_for_endpoint(
|
||||
&self,
|
||||
model: &str,
|
||||
base_url: &str,
|
||||
) -> Option<std::sync::Arc<crate::auth::AuthManager>> {
|
||||
self.credential_authority()
|
||||
.manager_for(self.model_platform(model), base_url)
|
||||
}
|
||||
/// The SESSION credential (if any) that may ride a request for the routing
|
||||
/// slug `model`. The only producer is the chokepoint.
|
||||
pub(super) fn session_credential_for_model(
|
||||
&self,
|
||||
model: &str,
|
||||
) -> Option<crate::auth::credential_authority::SessionCredential> {
|
||||
self.credential_authority()
|
||||
.credential_for(self.model_platform(model), &self.model_base_url(model))
|
||||
}
|
||||
/// The base URL a routing slug actually resolves to in the live catalog.
|
||||
/// Falls back to the session's own inference endpoint for an unlisted slug,
|
||||
/// which is exactly where `resolve_aux_model_sampling_config`'s Tier-2
|
||||
/// fallback entry routes — so the endpoint the rule is applied to is always
|
||||
/// the endpoint the request is sent to.
|
||||
fn model_base_url(&self, model: &str) -> String {
|
||||
let models = self.models_manager.models();
|
||||
match crate::agent::config::find_model_by_id(&models, model) {
|
||||
Some(entry) => entry.info().base_url.clone(),
|
||||
None => self.models_manager.endpoints().resolve_inference_base_url(),
|
||||
}
|
||||
}
|
||||
/// This SESSION's own selected catalog key (H4) — never the process-global
|
||||
/// `ModelsManager::current_model_id()`, which Leader mode never writes and
|
||||
/// which is last-writer-wins across concurrent sessions.
|
||||
pub(super) fn selected_catalog_key(&self) -> Option<String> {
|
||||
self.selected_catalog_key.borrow().clone()
|
||||
}
|
||||
/// Keep the session's own selected catalog key consistent with an
|
||||
/// `OverrideModelName` rename: KEEP it when it still names `model_name`
|
||||
/// (same entry, new routing name), otherwise CLEAR it.
|
||||
///
|
||||
/// H-c: `OverrideModelName` is the one command that rewrites
|
||||
/// `SamplingConfig::model` without going through `SetSessionModel`, so it
|
||||
/// used to leave the field naming a model the session is no longer on.
|
||||
/// Clearing rather than re-resolving is deliberate: re-resolving would put
|
||||
/// `resolve_catalog_key`'s `.rev()` guess INTO the field the whole rule
|
||||
/// treats as the session's deliberate selection, and a cleared field
|
||||
/// refuses a collided slug instead of guessing its OAuth twin (H-b).
|
||||
pub(super) fn retain_selected_catalog_key_for(&self, model_name: &str) {
|
||||
let models = self.models_manager.models();
|
||||
let still_names_it = self.selected_catalog_key().is_some_and(|key| {
|
||||
key == model_name
|
||||
|| models
|
||||
.get(key.as_str())
|
||||
.is_some_and(|entry| entry.info.model == model_name)
|
||||
});
|
||||
if !still_names_it {
|
||||
*self.selected_catalog_key.borrow_mut() = None;
|
||||
}
|
||||
}
|
||||
/// The registry platform the routing slug `model` belongs to, from the SAME
|
||||
/// lookup [`Self::auth_manager_for_endpoint`] routes on. `None` for a bare /
|
||||
/// `[model.*]` / unlisted model.
|
||||
pub(super) fn model_platform(&self, model: &str) -> Option<kigi_models::PlatformId> {
|
||||
let models = self.models_manager.models();
|
||||
crate::agent::models::platform_for_slug(
|
||||
&models,
|
||||
self.selected_catalog_key().as_deref(),
|
||||
model,
|
||||
)
|
||||
}
|
||||
/// Whether `model` routes to the Claude Pro/Max OAuth-Messages platform
|
||||
/// (claude-pro-max) — the gate for the sampler's OAuth Messages adaptation
|
||||
/// (identity headers + "You are Claude Code" system prefix). A generic-OAuth
|
||||
/// platform speaking the Messages wire; every other model (incl. xai-grok,
|
||||
/// which is ChatCompletions) returns `false`, keeping the API-key Anthropic
|
||||
/// / MiniMax Messages requests byte-identical.
|
||||
fn model_is_anthropic_oauth(&self, model: &str) -> bool {
|
||||
self.model_platform(model).is_some_and(|platform| {
|
||||
platform.oauth().is_some()
|
||||
&& platform.wire_api() == kigi_models::PlatformWireApi::Messages
|
||||
})
|
||||
}
|
||||
/// Whether `model` routes to the GitHub Copilot ChatCompletions platform
|
||||
/// (github-copilot) — the gate for the sampler's editor-identity headers +
|
||||
/// `X-Initiator`. Every other model returns `false`, keeping the other
|
||||
/// ChatCompletions providers byte-identical.
|
||||
fn model_is_github_copilot(&self, model: &str) -> bool {
|
||||
self.model_platform(model)
|
||||
.is_some_and(kigi_models::PlatformId::sends_copilot_editor_headers)
|
||||
}
|
||||
/// Whether `model` routes to the ChatGPT/Codex Responses platform
|
||||
/// (openai-codex) — the gate for the sampler's Codex identity headers
|
||||
/// (`chatgpt-account-id` + originator + OpenAI-Beta). Every other model
|
||||
/// returns `false`, keeping the API-key `openai` Responses request
|
||||
/// byte-identical.
|
||||
fn model_is_openai_codex(&self, model: &str) -> bool {
|
||||
self.model_platform(model)
|
||||
.is_some_and(kigi_models::PlatformId::sends_codex_responses_headers)
|
||||
}
|
||||
/// The `bearer_resolver` an AUX / summary `SamplerConfig` may carry — the
|
||||
/// shared [`aux_bearer_resolver_for`] rule applied to the AUX model's own
|
||||
/// platform + endpoint.
|
||||
///
|
||||
/// The aux config never inherits the session resolver: it is passed this
|
||||
/// value explicitly (see
|
||||
/// [`crate::agent::config::stamp_session_local_sampler_fields`]), so
|
||||
/// "forgot to re-point" is not expressible.
|
||||
///
|
||||
/// The BYOK status comes from [`Self::aux_model_auth_facts`], which does NOT
|
||||
/// write the session model's single-slot memo.
|
||||
pub(super) fn aux_bearer_resolver(
|
||||
&self,
|
||||
slug: &str,
|
||||
base_url: &str,
|
||||
) -> Option<kigi_sampler::SharedBearerResolver> {
|
||||
let auth_method = self.auth_method_id.load();
|
||||
// An aux slug is NOT the session's selection, so it must not be resolved
|
||||
// against `selected_catalog_key` — the same rule the aux `api_key` obeys
|
||||
// (`credential_for_slug(.., None, ..)`). Keying an aux model on the
|
||||
// SESSION's selection let a colliding same-vendor slug resolve the OAuth
|
||||
// twin, whose pooled resolver would then overwrite the user's own key on
|
||||
// the aux request.
|
||||
let models = self.models_manager.models();
|
||||
aux_bearer_resolver_for(
|
||||
&self.credential_authority(),
|
||||
auth_method.as_deref(),
|
||||
crate::agent::models::platform_for_slug(&models, None, slug),
|
||||
self.aux_model_auth_facts(slug).byok,
|
||||
base_url,
|
||||
)
|
||||
}
|
||||
/// Emit a unified-log breadcrumb whenever the session-token refresh gate is
|
||||
/// evaluated with an **`Unknown`** per-model BYOK status on a session-based
|
||||
@@ -202,8 +491,8 @@ impl SessionActor {
|
||||
let ctx = serde_json::json!(
|
||||
{ "site" : site, "model_byok" : gate.model_byok.as_str(), "is_session_based"
|
||||
: gate.is_session_based, "endpoint_is_first_party" : gate
|
||||
.endpoint_is_first_party, "refresh_active" : refresh_active, "base_url" :
|
||||
base_url, }
|
||||
.endpoint_is_first_party, "credential_class" : gate.credential_class
|
||||
.as_str(), "refresh_active" : refresh_active, "base_url" : base_url, }
|
||||
);
|
||||
let sid = Some(self.session_info.id.0.as_ref());
|
||||
if refresh_active {
|
||||
@@ -237,18 +526,6 @@ impl SessionActor {
|
||||
}
|
||||
}
|
||||
}
|
||||
#[allow(clippy::items_after_statements)]
|
||||
struct AuthManagerBearerResolver(std::sync::Arc<crate::auth::AuthManager>);
|
||||
impl std::fmt::Debug for AuthManagerBearerResolver {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("AuthManagerBearerResolver").finish()
|
||||
}
|
||||
}
|
||||
impl kigi_sampler::BearerResolver for AuthManagerBearerResolver {
|
||||
fn current_bearer(&self) -> Option<String> {
|
||||
self.0.current_or_expired().map(|a| a.key)
|
||||
}
|
||||
}
|
||||
let cfg = self
|
||||
.chat_state_handle
|
||||
.get_sampling_config()
|
||||
@@ -260,6 +537,7 @@ impl SessionActor {
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: std::num::NonZeroU64::new(256_000).unwrap(),
|
||||
reasoning_effort: None,
|
||||
@@ -268,10 +546,32 @@ impl SessionActor {
|
||||
let creds = self.chat_state_handle.get_credentials().await;
|
||||
let model_facts = self.model_auth_facts(cfg.model.as_str());
|
||||
let auth_method = self.auth_method_id.load();
|
||||
let gate =
|
||||
SessionTokenAuthGate::new(auth_method.as_deref(), model_facts.byok, &cfg.base_url);
|
||||
let gate = SessionTokenAuthGate::new(
|
||||
auth_method.as_deref(),
|
||||
model_facts.byok,
|
||||
&cfg.base_url,
|
||||
self.model_platform(cfg.model.as_str()),
|
||||
&self.credential_authority(),
|
||||
);
|
||||
let use_bearer_resolver = gate.active();
|
||||
self.log_auth_gate_unknown("reconstruct_full_config", gate, &cfg.base_url);
|
||||
// Resolve the bearer from the ACTIVE model's OWN manager: a grok model
|
||||
// wraps the xai-grok manager, never the Kimi one (captured before
|
||||
// `cfg.model` is moved into the struct below). `None` when the gate is
|
||||
// inactive or the oauth provider has no manager (fail-fast, no Kimi
|
||||
// fallback).
|
||||
let inference_auth_manager = if use_bearer_resolver {
|
||||
self.auth_manager_for_endpoint(&cfg.model, &cfg.base_url)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
// Claude Pro/Max OAuth Messages adaptation for THIS turn's model
|
||||
// (captured before `cfg.model` is moved into the struct below).
|
||||
let anthropic_oauth = self.model_is_anthropic_oauth(&cfg.model);
|
||||
// GitHub Copilot editor-identity headers for THIS turn's model.
|
||||
let github_copilot = self.model_is_github_copilot(&cfg.model);
|
||||
// ChatGPT/Codex identity headers for THIS turn's model.
|
||||
let openai_codex = self.model_is_openai_codex(&cfg.model);
|
||||
let auth_scheme = model_facts.auth_scheme;
|
||||
let mut extra_headers = cfg.extra_headers;
|
||||
crate::agent::config::inject_url_derived_headers(
|
||||
@@ -312,6 +612,10 @@ impl SessionActor {
|
||||
top_p: cfg.top_p,
|
||||
api_backend: cfg.api_backend,
|
||||
auth_scheme,
|
||||
anthropic_oauth,
|
||||
github_copilot,
|
||||
openai_codex,
|
||||
chat_compat: cfg.chat_compat,
|
||||
extra_headers,
|
||||
context_window: cfg.context_window.get(),
|
||||
reasoning_effort: cfg.reasoning_effort,
|
||||
@@ -321,15 +625,7 @@ impl SessionActor {
|
||||
idle_timeout_secs: None,
|
||||
origin_client: self.origin_client.clone(),
|
||||
attribution_callback: self.attribution_callback.clone(),
|
||||
bearer_resolver: if use_bearer_resolver {
|
||||
self.auth_manager
|
||||
.as_ref()
|
||||
.map(|am| -> kigi_sampler::SharedBearerResolver {
|
||||
std::sync::Arc::new(AuthManagerBearerResolver(am.clone()))
|
||||
})
|
||||
} else {
|
||||
None
|
||||
},
|
||||
bearer_resolver: inference_auth_manager.map(auth_manager_bearer_resolver),
|
||||
supports_backend_search: self.supports_backend_search.get(),
|
||||
compactions_remaining: self.compactions_remaining.get(),
|
||||
compaction_at_tokens: self.compaction_at_tokens.get(),
|
||||
@@ -459,17 +755,27 @@ impl SessionActor {
|
||||
slug: &str,
|
||||
) -> Option<kigi_sampler::SamplerConfig> {
|
||||
let creds = self.chat_state_handle.get_credentials().await;
|
||||
let session_key = self
|
||||
.auth_manager
|
||||
.as_ref()
|
||||
.and_then(|am| am.current_or_expired().map(|a| a.key.clone()));
|
||||
let models = self.models_manager.models();
|
||||
// Resolve the aux token by the aux model's OWN platform AND endpoint: a
|
||||
// grok (oauth-platform) aux model draws its pooled grok token or `None`,
|
||||
// and an API-key registry platform draws NOTHING — NEVER the primary
|
||||
// Kimi session token (which `resolve_credentials` would otherwise stamp
|
||||
// onto an api.x.ai / api.deepseek.com request). The first-party
|
||||
// subscription channel still gets the primary (byte-identical).
|
||||
// M6: ONE lookup. The platform AND the base URL the rule is applied to
|
||||
// both come from `credential_for_slug`'s single resolution of `slug`
|
||||
// against this catalog, so the platform and the endpoint can no longer
|
||||
// disagree (they were previously resolved by two different lookups).
|
||||
// Aux slugs are not the session's selection, so no `current_key`.
|
||||
let session_key = self
|
||||
.credential_authority()
|
||||
.credential_for_slug(&models, None, slug);
|
||||
let endpoints = self.models_manager.endpoints();
|
||||
crate::agent::config::resolve_aux_model_sampling_config(
|
||||
slug,
|
||||
&models,
|
||||
&endpoints,
|
||||
session_key.as_deref(),
|
||||
session_key.as_ref(),
|
||||
creds.alpha_test_key.clone(),
|
||||
)
|
||||
}
|
||||
@@ -484,9 +790,15 @@ impl SessionActor {
|
||||
) -> Option<(kigi_sampler::SamplingClient, String)> {
|
||||
let active_session_config = self.reconstruct_full_config().await;
|
||||
let mut cfg = self.resolve_aux_sampler_config(slug).await?;
|
||||
// LEAK 1b: the aux classifier must NOT inherit the SESSION model's
|
||||
// (Kimi) bearer_resolver — the resolver is decided by the AUX model's
|
||||
// own platform + endpoint at the chokepoint and passed in explicitly,
|
||||
// so there is no "copy then remember to re-point" step to forget.
|
||||
let aux_resolver = self.aux_bearer_resolver(slug, &cfg.base_url);
|
||||
crate::agent::config::stamp_session_local_sampler_fields(
|
||||
&mut cfg,
|
||||
&active_session_config,
|
||||
aux_resolver,
|
||||
Some(self.max_retries),
|
||||
);
|
||||
let model = cfg.model.clone();
|
||||
@@ -634,6 +946,7 @@ impl SessionActor {
|
||||
session_id = % self.session_info.id.0, is_session_based = gate
|
||||
.is_session_based, model_byok = gate.model_byok.as_str(),
|
||||
endpoint_is_first_party = gate.endpoint_is_first_party,
|
||||
credential_class = gate.credential_class.as_str(),
|
||||
"auth recovery: sampler 401 not refreshable (api-key auth) — surfacing 401",
|
||||
);
|
||||
kigi_log::unified_log::warn(
|
||||
@@ -643,7 +956,8 @@ impl SessionActor {
|
||||
{ "kind" : error.kind.as_str(), "status_code" : error
|
||||
.status_code, "is_session_based" : gate.is_session_based,
|
||||
"model_byok" : gate.model_byok.as_str(),
|
||||
"endpoint_is_first_party" : gate.endpoint_is_first_party, }
|
||||
"endpoint_is_first_party" : gate.endpoint_is_first_party,
|
||||
"credential_class" : gate.credential_class.as_str(), }
|
||||
)),
|
||||
);
|
||||
}
|
||||
@@ -659,7 +973,17 @@ impl SessionActor {
|
||||
)),
|
||||
);
|
||||
}
|
||||
if auth_recovery_eligible && let Some(ref am) = self.auth_manager {
|
||||
// Recover via the ACTIVE model's OWN manager: a grok 401 recovers the
|
||||
// xai-grok session via the xai-grok manager, never the Kimi one. For a
|
||||
// Kimi / non-oauth model this resolves to the primary — byte-identical.
|
||||
if auth_recovery_eligible {
|
||||
let (recovery_model, recovery_base_url) = self
|
||||
.chat_state_handle
|
||||
.get_sampling_config()
|
||||
.await
|
||||
.map(|c| (c.model, c.base_url))
|
||||
.unwrap_or_default();
|
||||
if let Some(am) = self.auth_manager_for_endpoint(&recovery_model, &recovery_base_url) {
|
||||
if am.try_recover_unauthorized().await {
|
||||
tracing::info!(
|
||||
session_id = % self.session_info.id.0,
|
||||
@@ -683,6 +1007,7 @@ impl SessionActor {
|
||||
None,
|
||||
);
|
||||
}
|
||||
}
|
||||
if matches!(error.kind, SamplingErrorKind::IdleTimeout) {
|
||||
self.signals_handle().record_idle_timeout();
|
||||
}
|
||||
@@ -844,14 +1169,17 @@ impl SessionActor {
|
||||
}
|
||||
/// Proactively refresh the auth token if near expiry.
|
||||
pub(super) async fn refresh_token_if_expired(&self) {
|
||||
if let Some(ref am) = self.auth_manager {
|
||||
let creds = self.chat_state_handle.get_credentials().await;
|
||||
let (model_id, base_url) = self
|
||||
.chat_state_handle
|
||||
.get_sampling_config()
|
||||
.await
|
||||
.map(|c| (c.model, c.base_url))
|
||||
.unwrap_or_default();
|
||||
// Refresh the ACTIVE model's OWN manager: a grok model refreshes the
|
||||
// xai-grok token via the xai-grok manager, never the Kimi one. For a
|
||||
// Kimi / non-oauth model this resolves to the primary — byte-identical.
|
||||
if let Some(am) = self.auth_manager_for_endpoint(&model_id, &base_url) {
|
||||
let creds = self.chat_state_handle.get_credentials().await;
|
||||
if self.auth_gate(&model_id, &base_url).active()
|
||||
&& let Ok(key) = am.get_valid_token().await
|
||||
{
|
||||
@@ -882,6 +1210,25 @@ impl SessionActor {
|
||||
.map(|c| c.model)
|
||||
.unwrap_or_default();
|
||||
let Some(ref key) = current_key else { return };
|
||||
// M7/M9: a registry-platform model's key normally comes from that
|
||||
// platform's credential resolved into its catalog entry, so 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.
|
||||
//
|
||||
// But a `[model."deepseek/deepseek-chat"]` override DOES keep the base
|
||||
// entry's `info.id` (`ConfigModelOverride::apply`), so "has a platform"
|
||||
// does NOT imply "has no `[model.*]` block" — skipping on the platform
|
||||
// alone would freeze an on-disk key rotation for the whole session.
|
||||
// Skip only when the catalog entry carries no own credential at all,
|
||||
// which is exactly the "key came from the platform, not from config"
|
||||
// case the disk read cannot improve on.
|
||||
if self.model_platform(¤t_model_id).is_some()
|
||||
&& !self.model_has_own_credential(¤t_model_id)
|
||||
{
|
||||
return;
|
||||
}
|
||||
let Some(new_key) = self.reload_api_key_from_config(¤t_model_id) else {
|
||||
return;
|
||||
};
|
||||
@@ -896,6 +1243,15 @@ impl SessionActor {
|
||||
creds.api_key = Some(new_key);
|
||||
self.chat_state_handle.update_credentials(creds);
|
||||
}
|
||||
/// Whether the live catalog entry for `slug` carries its own credential —
|
||||
/// an `api_key`/`env_key` from a `[model.*]` block, which a config edit can
|
||||
/// rotate mid-session. A platform entry whose key came from the platform
|
||||
/// credential has none.
|
||||
fn model_has_own_credential(&self, slug: &str) -> bool {
|
||||
let models = self.models_manager.models();
|
||||
crate::agent::config::find_model_by_id(&models, slug)
|
||||
.is_some_and(crate::agent::config::ModelEntry::has_own_credentials)
|
||||
}
|
||||
fn reload_api_key_from_config(&self, current_model_id: &str) -> Option<String> {
|
||||
let raw_config = crate::config::load_effective_config()
|
||||
.map_err(|e| tracing::warn!(error = % e, "Failed to reload config"))
|
||||
@@ -973,3 +1329,47 @@ impl SessionActor {
|
||||
.push_assistant_response(assistant_item);
|
||||
}
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod bearer_resolver_tests {
|
||||
use super::AuthManagerBearerResolver;
|
||||
use kigi_sampler::BearerResolver;
|
||||
|
||||
/// LEAK 1b: the shared `AuthManagerBearerResolver` resolves the LIVE bearer
|
||||
/// of the manager it wraps. So an aux bearer_resolver built over grok's OWN
|
||||
/// (oauth) pooled manager yields grok's token (or `None`) — NEVER the Kimi
|
||||
/// session token that a Kimi-manager resolver would. The
|
||||
/// [`CredentialAuthority::bearer_resolver_for`](crate::auth::credential_authority::CredentialAuthority::bearer_resolver_for)
|
||||
/// wraps exactly this grok manager for a grok aux model.
|
||||
#[tokio::test]
|
||||
async fn resolver_resolves_the_wrapped_manager_never_kimi() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let kimi = std::sync::Arc::new(crate::auth::AuthManager::new(
|
||||
dir.path(),
|
||||
crate::auth::KimiCodeConfig::default(),
|
||||
));
|
||||
kimi.hot_swap(crate::auth::KimiAuth {
|
||||
key: "kimi-tok".to_string(),
|
||||
auth_mode: crate::auth::AuthMode::OAuth,
|
||||
..crate::auth::KimiAuth::test_default()
|
||||
});
|
||||
// The Kimi-manager resolver yields the Kimi bearer.
|
||||
assert_eq!(
|
||||
AuthManagerBearerResolver(kimi.clone()).current_bearer(),
|
||||
Some("kimi-tok".to_string()),
|
||||
);
|
||||
// The grok (oauth) pooled manager is distinct — its resolver never
|
||||
// yields the Kimi bearer (grok's own token, or None).
|
||||
let oauth = kigi_models::PlatformId::XaiGrok
|
||||
.oauth()
|
||||
.expect("xai-grok carries an OAuthConfig");
|
||||
let grok = crate::auth::oauth_registry::global_manager_for(
|
||||
&crate::auth::oauth_registry::pool_home(),
|
||||
oauth,
|
||||
);
|
||||
assert_ne!(
|
||||
AuthManagerBearerResolver(grok).current_bearer(),
|
||||
Some("kimi-tok".to_string()),
|
||||
"a grok aux bearer_resolver must never resolve the Kimi session token",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -747,16 +747,24 @@ impl SessionActor {
|
||||
ok_end_turn(0, None)
|
||||
}
|
||||
BuiltinAction::GoalPause => {
|
||||
if self.graph_owns_goal_engine() {
|
||||
self.send_slash_command_output(
|
||||
"A graph owns the goal engine. Use /graph pause instead.",
|
||||
)
|
||||
.await;
|
||||
return ok_end_turn(0, None);
|
||||
}
|
||||
let current_tokens = self.chat_state_handle.get_total_tokens().await as i64;
|
||||
use crate::session::goal_tracker::{GoalPauseReason, GoalStatus};
|
||||
let (msg, changed) = {
|
||||
let mut tracker = self.goal_tracker.lock();
|
||||
match tracker.status() {
|
||||
Some(GoalStatus::Active) => {
|
||||
debug_assert!(
|
||||
tracker.pause(GoalPauseReason::User),
|
||||
"Active goal must pause"
|
||||
);
|
||||
// Side effect OUTSIDE the assert: debug_assert!
|
||||
// strips its condition in release builds, which
|
||||
// would silently skip the pause itself.
|
||||
let paused = tracker.pause(GoalPauseReason::User);
|
||||
debug_assert!(paused, "Active goal must pause");
|
||||
("Goal paused. Use /goal resume to continue.", true)
|
||||
}
|
||||
Some(
|
||||
@@ -790,27 +798,149 @@ impl SessionActor {
|
||||
unreachable!("GoalResume is intercepted in handle_prompt")
|
||||
}
|
||||
BuiltinAction::GoalClear => {
|
||||
self.goal_tracker.lock().clear();
|
||||
// `/goal clear` is a deliberate user reset — drop both
|
||||
// streaks so stale counters from the previous goal
|
||||
// can't leak into the next one.
|
||||
self.goal_continuation_streak
|
||||
.store(0, std::sync::atomic::Ordering::Relaxed);
|
||||
self.goal_blocked_streak
|
||||
.store(0, std::sync::atomic::Ordering::Relaxed);
|
||||
// Drop goal-turn-origin task ids so a future goal's drain
|
||||
// doesn't suppress the next goal's (or post-goal) tasks.
|
||||
self.goal_turn_task_ids.lock().clear();
|
||||
// Clear per-subagent token records so stale entries
|
||||
// from the previous goal don't leak into the next.
|
||||
self.subagent_token_records.lock().clear();
|
||||
self.clear_pending_classifier_completions();
|
||||
// Emit a cleared notification so the pager drops goal state.
|
||||
let update = crate::session::goal_orchestrator::build_goal_cleared();
|
||||
self.send_xai_notification(update).await;
|
||||
if self.graph_owns_goal_engine() {
|
||||
self.send_slash_command_output(
|
||||
"A graph owns the goal engine. Use /graph clear instead.",
|
||||
)
|
||||
.await;
|
||||
return ok_end_turn(0, None);
|
||||
}
|
||||
// `/goal clear` is a deliberate user reset — the shared
|
||||
// helper drops the tracker, both streaks, goal-turn task
|
||||
// ids, per-subagent token records, pending classifier
|
||||
// claims, and notifies the pager. Shared with the graph
|
||||
// node boundary and `/graph clear`.
|
||||
self.reset_goal_engine_state().await;
|
||||
self.send_slash_command_output("Goal cleared.").await;
|
||||
ok_end_turn(0, None)
|
||||
}
|
||||
BuiltinAction::GraphStatus => {
|
||||
let msg = self.graph_status_message().await;
|
||||
self.send_slash_command_output(&msg).await;
|
||||
ok_end_turn(0, None)
|
||||
}
|
||||
BuiltinAction::GraphShow => {
|
||||
// Box-drawing DAG; wider than the budget (or no graph)
|
||||
// degrades to the indented status tree — wrapped box art
|
||||
// is worse than no art.
|
||||
const SHOW_WIDTH_BUDGET: usize = 120;
|
||||
let rendered =
|
||||
self.graph_tracker.lock().snapshot().and_then(|s| {
|
||||
crate::session::graph_render::render_dag(s, SHOW_WIDTH_BUDGET)
|
||||
});
|
||||
let msg = match rendered {
|
||||
Some(dag) => dag,
|
||||
None => self.graph_status_message().await,
|
||||
};
|
||||
self.send_slash_command_output(&msg).await;
|
||||
ok_end_turn(0, None)
|
||||
}
|
||||
BuiltinAction::GraphPause => {
|
||||
use crate::session::goal_tracker::{GoalPauseReason, GoalStatus};
|
||||
let (msg, changed) = {
|
||||
let mut tracker = self.graph_tracker.lock();
|
||||
match tracker.status() {
|
||||
Some(GoalStatus::Active) => {
|
||||
// Side effect OUTSIDE the assert: debug_assert!
|
||||
// strips its condition in release builds, which
|
||||
// would silently skip the pause itself.
|
||||
let paused = tracker.pause(GoalPauseReason::User);
|
||||
debug_assert!(paused, "Active graph must pause");
|
||||
("Graph paused. Use /graph resume to continue.", true)
|
||||
}
|
||||
Some(s) if s.is_paused() => ("Graph is already paused.", false),
|
||||
Some(GoalStatus::Complete) => ("Graph is already complete.", false),
|
||||
Some(GoalStatus::BudgetLimited) => ("Graph is budget-limited.", false),
|
||||
Some(_) | None => ("No graph is currently set.", false),
|
||||
}
|
||||
};
|
||||
if changed {
|
||||
// Pause the running node's goal too so the in-turn
|
||||
// loop stops at the next round boundary.
|
||||
self.auto_pause_goal_if_active(GoalPauseReason::User).await;
|
||||
self.persist_graph_state();
|
||||
}
|
||||
self.send_slash_command_output(msg).await;
|
||||
ok_end_turn(0, None)
|
||||
}
|
||||
BuiltinAction::GraphClear => {
|
||||
let had_graph = self.graph_tracker.lock().snapshot().is_some();
|
||||
// Clear the goal engine ONLY when the graph owns it
|
||||
// (Active/paused ⇒ the engine's goal is a node goal). A
|
||||
// terminal graph (Complete/BudgetLimited) may coexist
|
||||
// with an unrelated standalone /goal the user started
|
||||
// afterwards — that goal must survive /graph clear.
|
||||
if self.graph_owns_goal_engine() {
|
||||
self.reset_goal_engine_state().await;
|
||||
}
|
||||
// Projection teardown only when there is something of
|
||||
// OURS to un-project: with no session graph, deleting
|
||||
// .kigi/graph.jsonl would destroy another session's
|
||||
// revivable graph while replying "No graph is set".
|
||||
let session_graph_id = self
|
||||
.graph_tracker
|
||||
.lock()
|
||||
.snapshot()
|
||||
.map(|s| s.graph_id.clone());
|
||||
self.graph_tracker.lock().clear();
|
||||
if had_graph {
|
||||
// Take the writer lock if we don't hold it (e.g. a
|
||||
// session-snapshot-restored graph cleared before any
|
||||
// resume). Busy = another instance owns the project
|
||||
// graph; local-only clear is then correct.
|
||||
match self.acquire_project_graph_writer() {
|
||||
Ok(true) => {
|
||||
// Identity check: only remove a projection
|
||||
// that belongs to the graph being cleared.
|
||||
let foreign = match (self.projected_graph_id(), &session_graph_id) {
|
||||
(Some(projected), Some(ours)) => projected != *ours,
|
||||
_ => false,
|
||||
};
|
||||
if foreign {
|
||||
tracing::info!(
|
||||
"graph clear: projection belongs to a different \
|
||||
graph; leaving .kigi/graph.jsonl in place"
|
||||
);
|
||||
self.graph_project_lock.borrow_mut().take();
|
||||
}
|
||||
}
|
||||
Ok(false) => {
|
||||
tracing::info!(
|
||||
"graph clear: another instance holds the project \
|
||||
graph; local session state cleared only"
|
||||
);
|
||||
}
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
%err,
|
||||
"graph clear: project lock acquisition failed; \
|
||||
.kigi/graph.jsonl may survive as stale"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
// persist runs BEFORE the lock drops so the projection
|
||||
// removal (when we hold writer rights on OUR graph)
|
||||
// executes; without the lock it is a session-only clear.
|
||||
self.persist_graph_state();
|
||||
self.graph_project_lock.borrow_mut().take();
|
||||
self.send_slash_command_output(if had_graph {
|
||||
"Graph cleared."
|
||||
} else {
|
||||
"No graph is currently set."
|
||||
})
|
||||
.await;
|
||||
ok_end_turn(0, None)
|
||||
}
|
||||
// GraphSet / GraphResume are intercepted in handle_prompt
|
||||
// (like GoalSet / GoalResume) so a successful setup/resume
|
||||
// flows through to model inference.
|
||||
BuiltinAction::GraphSet { .. } => {
|
||||
unreachable!("GraphSet is intercepted in handle_prompt")
|
||||
}
|
||||
BuiltinAction::GraphResume { .. } => {
|
||||
unreachable!("GraphResume is intercepted in handle_prompt")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -136,6 +136,7 @@ pub(crate) async fn spawn_session_actor(
|
||||
persisted_signals: Option<crate::session::signals::SessionSignals>,
|
||||
persisted_plan_mode: Option<crate::session::plan_mode::PlanModeSnapshot>,
|
||||
persisted_goal_mode: Option<crate::session::goal_tracker::GoalOrchestration>,
|
||||
persisted_graph_mode: Option<crate::session::graph_tracker::GraphOrchestration>,
|
||||
persisted_announcement_state: Option<crate::session::announcement_state::AnnouncementState>,
|
||||
memory_config: Option<crate::config::MemoryConfig>,
|
||||
feedback_flags: crate::session::feedback_manager::FeedbackFlags,
|
||||
@@ -150,6 +151,7 @@ pub(crate) async fn spawn_session_actor(
|
||||
app_builder_deployer_config: kigi_tools::implementations::kigi::deploy_app::AppBuilderDeployerConfig,
|
||||
write_file_enabled: bool,
|
||||
goal_enabled: bool,
|
||||
graph_enabled: bool,
|
||||
subagents_enabled: bool,
|
||||
ask_user_question_enabled: bool,
|
||||
client_hooks: crate::extensions::hooks::ClientHooks,
|
||||
@@ -344,6 +346,7 @@ pub(crate) async fn spawn_session_actor(
|
||||
temperature: sampling_config.temperature,
|
||||
top_p: sampling_config.top_p,
|
||||
api_backend: sampling_config.api_backend.clone(),
|
||||
chat_compat: sampling_config.chat_compat,
|
||||
extra_headers: sampling_config.extra_headers.clone(),
|
||||
context_window: context_window_override.unwrap_or(baseline_context_window),
|
||||
reasoning_effort: sampling_config.reasoning_effort,
|
||||
@@ -429,6 +432,17 @@ pub(crate) async fn spawn_session_actor(
|
||||
};
|
||||
Arc::new(parking_lot::Mutex::new(tracker))
|
||||
};
|
||||
let graph_project_dir =
|
||||
crate::session::graph_project::project_graph_dir(tool_context.cwd.as_path());
|
||||
let graph_tracker = {
|
||||
let session_dir = crate::session::persistence::session_dir(&session_info);
|
||||
let tracker = if let Some(snapshot) = persisted_graph_mode {
|
||||
crate::session::graph_tracker::GraphTracker::from_snapshot(session_dir, snapshot)
|
||||
} else {
|
||||
crate::session::graph_tracker::GraphTracker::new(session_dir)
|
||||
};
|
||||
Arc::new(parking_lot::Mutex::new(tracker))
|
||||
};
|
||||
let current_prompt_mode = Arc::new(parking_lot::Mutex::new(PromptMode::Agent));
|
||||
let turn_prompt_mode = Arc::new(parking_lot::Mutex::new(PromptMode::Agent));
|
||||
let task_output_tool_name = Arc::new(std::sync::OnceLock::new());
|
||||
@@ -966,10 +980,21 @@ pub(crate) async fn spawn_session_actor(
|
||||
}
|
||||
};
|
||||
let doom_loop_recovery = effective_config.resolve_doom_loop_recovery();
|
||||
let session_model_id_for_actor = session_model_id.clone();
|
||||
let session = Arc::new_cyclic(|weak: &std::sync::Weak<SessionActor>| SessionActor {
|
||||
session_info: session_info.clone(),
|
||||
auth_method_id,
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
// H4: seed the session's OWN selected catalog key from the model it was
|
||||
// spawned with, resolved through the picker's lookup. Never the
|
||||
// process-global `current_model_id()`. H-c: the rule lives in
|
||||
// `selected_catalog_key_for_spawn` so it is covered by a test.
|
||||
selected_catalog_key: std::cell::RefCell::new(
|
||||
crate::agent::models::selected_catalog_key_for_spawn(
|
||||
&models_manager.models(),
|
||||
&session_model_id_for_actor,
|
||||
),
|
||||
),
|
||||
attribution_callback,
|
||||
auth_manager,
|
||||
state,
|
||||
@@ -1086,6 +1111,14 @@ pub(crate) async fn spawn_session_actor(
|
||||
goal_harness_enabled: std::sync::atomic::AtomicBool::new(false),
|
||||
goal_harness_availability_reconciled: std::sync::atomic::AtomicBool::new(false),
|
||||
goal_tracker,
|
||||
graph_enabled,
|
||||
graph_tracker,
|
||||
graph_concurrency: effective_config.resolve_graph_concurrency(),
|
||||
graph_node_rounds: effective_config.resolve_graph_node_rounds(),
|
||||
graph_replan_cap: effective_config.resolve_graph_replan_cap(),
|
||||
graph_optimizer_enabled: effective_config.resolve_graph_optimizer_enabled(),
|
||||
graph_project_dir,
|
||||
graph_project_lock: std::cell::RefCell::new(None),
|
||||
goal_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
@@ -1229,6 +1262,18 @@ pub(crate) async fn spawn_session_actor(
|
||||
),
|
||||
)
|
||||
.await;
|
||||
// A restored graph was demoted (Active→UserPaused, Running→Ready) IN
|
||||
// MEMORY after the updates-log replay, whose last GraphUpdated still
|
||||
// shows the pre-shutdown Active state. Re-emit truth once so a
|
||||
// reattached pager never renders a stale self-driving chip — and
|
||||
// best-effort reclaim project writership so the shared file gets the
|
||||
// demoted truth too (Busy = another instance owns it; skip quietly).
|
||||
if session.graph_tracker.lock().snapshot().is_some() {
|
||||
if let Some(msg) = session.claim_project_graph_for_resume() {
|
||||
tracing::info!(%msg, "graph restore: project writership not reclaimed");
|
||||
}
|
||||
session.persist_graph_state();
|
||||
}
|
||||
if let Some(ref display_cwd) = prompt_display_cwd {
|
||||
session
|
||||
.agent
|
||||
@@ -1511,6 +1556,7 @@ pub(crate) async fn spawn_session_on_thread(
|
||||
persisted_signals: Option<crate::session::signals::SessionSignals>,
|
||||
persisted_plan_mode: Option<crate::session::plan_mode::PlanModeSnapshot>,
|
||||
persisted_goal_mode: Option<crate::session::goal_tracker::GoalOrchestration>,
|
||||
persisted_graph_mode: Option<crate::session::graph_tracker::GraphOrchestration>,
|
||||
persisted_announcement_state: Option<crate::session::announcement_state::AnnouncementState>,
|
||||
memory_config: Option<crate::config::MemoryConfig>,
|
||||
feedback_flags: crate::session::feedback_manager::FeedbackFlags,
|
||||
@@ -1525,6 +1571,7 @@ pub(crate) async fn spawn_session_on_thread(
|
||||
app_builder_deployer_config: kigi_tools::implementations::kigi::deploy_app::AppBuilderDeployerConfig,
|
||||
write_file_enabled: bool,
|
||||
goal_enabled: bool,
|
||||
graph_enabled: bool,
|
||||
subagents_enabled: bool,
|
||||
ask_user_question_enabled: bool,
|
||||
client_hooks: crate::extensions::hooks::ClientHooks,
|
||||
@@ -1653,6 +1700,7 @@ pub(crate) async fn spawn_session_on_thread(
|
||||
persisted_signals,
|
||||
persisted_plan_mode,
|
||||
persisted_goal_mode,
|
||||
persisted_graph_mode,
|
||||
persisted_announcement_state,
|
||||
memory_config,
|
||||
feedback_flags,
|
||||
@@ -1667,6 +1715,7 @@ pub(crate) async fn spawn_session_on_thread(
|
||||
app_builder_deployer_config,
|
||||
write_file_enabled,
|
||||
goal_enabled,
|
||||
graph_enabled,
|
||||
subagents_enabled,
|
||||
ask_user_question_enabled,
|
||||
client_hooks,
|
||||
|
||||
@@ -478,6 +478,14 @@ impl SessionActor {
|
||||
if let Some(running_task) = running_task {
|
||||
running_task.abort();
|
||||
}
|
||||
// Re-sweep subagents AFTER the abort: the first sweep raced the
|
||||
// still-live turn future, which may have spawned NEW harness
|
||||
// children (graph batch workers/verifiers) between the sweep and
|
||||
// the abort. Coordinator-channel FIFO guarantees this second
|
||||
// Cancel lands after any Spawn the turn issued before it died.
|
||||
if cancel_subagents && let Some(prompt_id) = cancelled_prompt_id.as_deref() {
|
||||
self.cancel_subagents_for_prompt_id(prompt_id);
|
||||
}
|
||||
// The aborted turn's `BlockingWaitGuard`s drop asynchronously (they
|
||||
// live in tool futures owned by the drainer task / subagent spawn
|
||||
// task). Until they do, `queue_input` would read a stale depth > 0 and
|
||||
|
||||
@@ -305,10 +305,28 @@ impl SessionActor {
|
||||
objective,
|
||||
token_budget,
|
||||
} => {
|
||||
// The graph owns the goal engine while set: a
|
||||
// manual /goal would corrupt the running node.
|
||||
if self.graph_owns_goal_engine() {
|
||||
self.send_slash_command_output(
|
||||
"A graph owns the goal engine. Use /graph status, or /graph \
|
||||
clear before /goal.",
|
||||
)
|
||||
.await;
|
||||
return ok_end_turn(0, None);
|
||||
}
|
||||
let reminder = self.setup_goal(&objective, token_budget).await;
|
||||
vec![text_block(reminder), text_block(objective)]
|
||||
}
|
||||
BuiltinAction::GoalResume => match self.resume_goal().await {
|
||||
BuiltinAction::GoalResume => {
|
||||
if self.graph_owns_goal_engine() {
|
||||
self.send_slash_command_output(
|
||||
"A graph owns the goal engine. Use /graph resume instead.",
|
||||
)
|
||||
.await;
|
||||
return ok_end_turn(0, None);
|
||||
}
|
||||
match self.resume_goal().await {
|
||||
GoalResumeOutcome::Inference { reminder, user_msg } => {
|
||||
self.send_slash_command_output(&user_msg).await;
|
||||
vec![text_block(reminder)]
|
||||
@@ -317,7 +335,33 @@ impl SessionActor {
|
||||
self.send_slash_command_output(&msg).await;
|
||||
return ok_end_turn(0, None);
|
||||
}
|
||||
}
|
||||
}
|
||||
BuiltinAction::GraphSet {
|
||||
objective,
|
||||
token_budget,
|
||||
} => match self.setup_graph(&objective, token_budget).await {
|
||||
super::graph::GraphSetupOutcome::Inference { reminder, user_msg } => {
|
||||
self.send_slash_command_output(&user_msg).await;
|
||||
vec![text_block(reminder), text_block(objective)]
|
||||
}
|
||||
super::graph::GraphSetupOutcome::Message(msg) => {
|
||||
self.send_slash_command_output(&msg).await;
|
||||
return ok_end_turn(0, None);
|
||||
}
|
||||
},
|
||||
BuiltinAction::GraphResume { extra_budget } => {
|
||||
match self.resume_graph(extra_budget).await {
|
||||
super::graph::GraphSetupOutcome::Inference { reminder, user_msg } => {
|
||||
self.send_slash_command_output(&user_msg).await;
|
||||
vec![text_block(reminder)]
|
||||
}
|
||||
super::graph::GraphSetupOutcome::Message(msg) => {
|
||||
self.send_slash_command_output(&msg).await;
|
||||
return ok_end_turn(0, None);
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => return self.execute_builtin_slash_command(action).await,
|
||||
}
|
||||
}
|
||||
@@ -688,13 +732,35 @@ impl SessionActor {
|
||||
self.goal_tracker.lock().status(),
|
||||
);
|
||||
if !goal_active {
|
||||
break round;
|
||||
// The node goal may have resolved MID-round without a
|
||||
// graph cascade (e.g. a classifier-disabled completion
|
||||
// applied by the mid-turn drainer). Consult the graph
|
||||
// seam before ending the turn so the graph advances or
|
||||
// settles loudly instead of stranding Active forever.
|
||||
match self.run_graph_round_end().await {
|
||||
Some(node_reminder) => {
|
||||
self.inject_goal_continuation_message(node_reminder).await;
|
||||
continue;
|
||||
}
|
||||
None => break round,
|
||||
}
|
||||
}
|
||||
match self.run_goal_round_end().await {
|
||||
GoalRoundDecision::Continue(directive) => {
|
||||
self.inject_goal_continuation_message(directive).await;
|
||||
}
|
||||
GoalRoundDecision::EndTurn => break round,
|
||||
GoalRoundDecision::EndTurn => {
|
||||
// Graph seam: when the node goal resolved, the
|
||||
// graph may advance to the next node inside the
|
||||
// SAME turn (multi-loop closed loop). None ends
|
||||
// the turn for real (graph done/paused/absent).
|
||||
match self.run_graph_round_end().await {
|
||||
Some(node_reminder) => {
|
||||
self.inject_goal_continuation_message(node_reminder).await;
|
||||
}
|
||||
None => break round,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
+36
-11
@@ -485,23 +485,44 @@ async fn no_legacy_hint_for_oidc_auth() {
|
||||
#[test]
|
||||
fn session_token_auth_gate_truth_table() {
|
||||
use crate::agent::auth_method::{ModelByok, session_token_auth_gate as gate};
|
||||
use crate::auth::credential_authority::CredentialClass;
|
||||
// Non-session methods never refresh, regardless of BYOK status or endpoint.
|
||||
// `Pooled` (an OAuth platform's own pool) and `Primary` (kimi-code, or a
|
||||
// bare / [model.*] model on the session's own endpoint) behave identically
|
||||
// here: each names a credential that IS refreshable on that host.
|
||||
for fp in [false, true] {
|
||||
assert!(!gate(false, ModelByok::NotByok, fp));
|
||||
assert!(!gate(false, ModelByok::Byok, fp));
|
||||
assert!(!gate(false, ModelByok::Unknown, fp));
|
||||
// Session method: a definite classification ignores the endpoint —
|
||||
// NotByok always refreshes (only ever routes to the session endpoint),
|
||||
// a genuine per-model Byok never does.
|
||||
assert!(gate(true, ModelByok::NotByok, fp));
|
||||
assert!(!gate(true, ModelByok::Byok, fp));
|
||||
for class in [CredentialClass::Pooled, CredentialClass::Primary] {
|
||||
assert!(!gate(false, ModelByok::NotByok, fp, class));
|
||||
assert!(!gate(false, ModelByok::Byok, fp, class));
|
||||
assert!(!gate(false, ModelByok::Unknown, fp, class));
|
||||
// Session method on an endpoint that DOES take a session
|
||||
// credential: a definite classification ignores the endpoint —
|
||||
// NotByok refreshes, a genuine per-model Byok never does.
|
||||
assert!(gate(true, ModelByok::NotByok, fp, class));
|
||||
assert!(!gate(true, ModelByok::Byok, fp, class));
|
||||
}
|
||||
// …and an API-key registry platform endpoint (`CredentialClass::None`)
|
||||
// is refused on every arm, first-party flag included: the leak guard.
|
||||
assert!(!gate(true, ModelByok::NotByok, fp, CredentialClass::None));
|
||||
assert!(!gate(true, ModelByok::Byok, fp, CredentialClass::None));
|
||||
assert!(!gate(true, ModelByok::Unknown, fp, CredentialClass::None));
|
||||
}
|
||||
// Session method + Unknown BYOK: refresh only against a first-party xAI
|
||||
// host, so a transiently-unclassifiable config can't demote a live session
|
||||
// (the stale-token 401 regression) yet the session token never leaks to a
|
||||
// third-party BYOK endpoint. This arm was unconditionally `false` pre-fix.
|
||||
assert!(gate(true, ModelByok::Unknown, true));
|
||||
assert!(!gate(true, ModelByok::Unknown, false));
|
||||
assert!(gate(
|
||||
true,
|
||||
ModelByok::Unknown,
|
||||
true,
|
||||
CredentialClass::Primary
|
||||
));
|
||||
assert!(!gate(
|
||||
true,
|
||||
ModelByok::Unknown,
|
||||
false,
|
||||
CredentialClass::Primary
|
||||
));
|
||||
}
|
||||
|
||||
/// Pre-fix, the gate read `auth_type` and skipped recovery here, 401'ing every
|
||||
@@ -847,7 +868,11 @@ async fn set_session_model_invalidates_byok_memo_for_same_model_id() {
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: crate::sampling::ApiBackend::ChatCompletions,
|
||||
chat_compat: Default::default(),
|
||||
auth_scheme: Default::default(),
|
||||
anthropic_oauth: false,
|
||||
github_copilot: false,
|
||||
openai_codex: false,
|
||||
extra_headers: Default::default(),
|
||||
context_window: 256_000,
|
||||
force_http1: false,
|
||||
@@ -865,7 +890,7 @@ async fn set_session_model_invalidates_byok_memo_for_same_model_id() {
|
||||
header_injector: None,
|
||||
};
|
||||
let _ = actor
|
||||
.handle_set_session_model(cfg, false, false, true, 85)
|
||||
.handle_set_session_model(cfg, None, false, false, true, 85)
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
|
||||
@@ -44,7 +44,11 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
auth_scheme: Default::default(),
|
||||
anthropic_oauth: false,
|
||||
github_copilot: false,
|
||||
openai_codex: false,
|
||||
extra_headers: Default::default(),
|
||||
context_window: 100_000,
|
||||
force_http1: false,
|
||||
@@ -87,6 +91,7 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: std::num::NonZeroU64::new(100_000).unwrap(),
|
||||
reasoning_effort: None,
|
||||
@@ -104,6 +109,7 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
|
||||
session_info,
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state: TokioMutex::new(State {
|
||||
@@ -212,6 +218,16 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
|
||||
"/tmp/test-session",
|
||||
)),
|
||||
)),
|
||||
graph_enabled: false,
|
||||
graph_tracker: Arc::new(parking_lot::Mutex::new(
|
||||
crate::session::graph_tracker::GraphTracker::new(std::env::temp_dir()),
|
||||
)),
|
||||
graph_concurrency: 1,
|
||||
graph_node_rounds: 3,
|
||||
graph_replan_cap: 3,
|
||||
graph_optimizer_enabled: false,
|
||||
graph_project_dir: None,
|
||||
graph_project_lock: std::cell::RefCell::new(None),
|
||||
goal_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
@@ -328,7 +344,11 @@ async fn first_turn_memory_injection_persists_to_chat_history() {
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
auth_scheme: Default::default(),
|
||||
anthropic_oauth: false,
|
||||
github_copilot: false,
|
||||
openai_codex: false,
|
||||
context_window: 100_000,
|
||||
force_http1: false,
|
||||
max_retries: None,
|
||||
@@ -371,6 +391,7 @@ async fn first_turn_memory_injection_persists_to_chat_history() {
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: std::num::NonZeroU64::new(100_000).unwrap(),
|
||||
reasoning_effort: None,
|
||||
@@ -456,7 +477,11 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
auth_scheme: Default::default(),
|
||||
anthropic_oauth: false,
|
||||
github_copilot: false,
|
||||
openai_codex: false,
|
||||
context_window: 100_000,
|
||||
force_http1: false,
|
||||
max_retries: None,
|
||||
@@ -503,6 +528,7 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: std::num::NonZeroU64::new(100_000).unwrap(),
|
||||
reasoning_effort: None,
|
||||
@@ -537,6 +563,7 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
|
||||
session_info: session_info.clone(),
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state: TokioMutex::new(State {
|
||||
@@ -648,6 +675,16 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
|
||||
"/tmp/test-session",
|
||||
)),
|
||||
)),
|
||||
graph_enabled: false,
|
||||
graph_tracker: Arc::new(parking_lot::Mutex::new(
|
||||
crate::session::graph_tracker::GraphTracker::new(std::env::temp_dir()),
|
||||
)),
|
||||
graph_concurrency: 1,
|
||||
graph_node_rounds: 3,
|
||||
graph_replan_cap: 3,
|
||||
graph_optimizer_enabled: false,
|
||||
graph_project_dir: None,
|
||||
graph_project_lock: std::cell::RefCell::new(None),
|
||||
goal_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
@@ -790,6 +827,7 @@ async fn cancel_running_task_teardown_clears_running_and_pending_work() {
|
||||
},
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state,
|
||||
@@ -893,6 +931,16 @@ async fn cancel_running_task_teardown_clears_running_and_pending_work() {
|
||||
"/tmp/test-session",
|
||||
)),
|
||||
)),
|
||||
graph_enabled: false,
|
||||
graph_tracker: Arc::new(parking_lot::Mutex::new(
|
||||
crate::session::graph_tracker::GraphTracker::new(std::env::temp_dir()),
|
||||
)),
|
||||
graph_concurrency: 1,
|
||||
graph_node_rounds: 3,
|
||||
graph_replan_cap: 3,
|
||||
graph_optimizer_enabled: false,
|
||||
graph_project_dir: None,
|
||||
graph_project_lock: std::cell::RefCell::new(None),
|
||||
goal_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
@@ -1702,7 +1750,11 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() {
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: kigi_sampler::ApiBackend::Responses,
|
||||
chat_compat: Default::default(),
|
||||
auth_scheme: Default::default(),
|
||||
anthropic_oauth: false,
|
||||
github_copilot: false,
|
||||
openai_codex: false,
|
||||
extra_headers: Default::default(),
|
||||
context_window: 100_000,
|
||||
force_http1: false,
|
||||
@@ -1768,6 +1820,7 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() {
|
||||
},
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state,
|
||||
@@ -1871,6 +1924,16 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() {
|
||||
"/tmp/test-session",
|
||||
)),
|
||||
)),
|
||||
graph_enabled: false,
|
||||
graph_tracker: Arc::new(parking_lot::Mutex::new(
|
||||
crate::session::graph_tracker::GraphTracker::new(std::env::temp_dir()),
|
||||
)),
|
||||
graph_concurrency: 1,
|
||||
graph_node_rounds: 3,
|
||||
graph_replan_cap: 3,
|
||||
graph_optimizer_enabled: false,
|
||||
graph_project_dir: None,
|
||||
graph_project_lock: std::cell::RefCell::new(None),
|
||||
goal_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -104,6 +104,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
|
||||
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,
|
||||
@@ -127,6 +128,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
|
||||
attribution_callback: None,
|
||||
auth_method_id: test_auth_method_id("cached_token"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
auth_manager: {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mgr = std::sync::Arc::new(crate::auth::AuthManager::new(
|
||||
@@ -241,6 +243,16 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
|
||||
"/tmp/test-session",
|
||||
)),
|
||||
)),
|
||||
graph_enabled: false,
|
||||
graph_tracker: Arc::new(parking_lot::Mutex::new(
|
||||
crate::session::graph_tracker::GraphTracker::new(std::env::temp_dir()),
|
||||
)),
|
||||
graph_concurrency: 1,
|
||||
graph_node_rounds: 3,
|
||||
graph_replan_cap: 3,
|
||||
graph_optimizer_enabled: false,
|
||||
graph_project_dir: None,
|
||||
graph_project_lock: std::cell::RefCell::new(None),
|
||||
goal_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
|
||||
+36
@@ -52,6 +52,7 @@ async fn create_test_actor(
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: std::num::NonZeroU64::new(context_window)
|
||||
.expect("test context_window must be non-zero"),
|
||||
@@ -71,6 +72,7 @@ async fn create_test_actor(
|
||||
rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(),
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state,
|
||||
@@ -173,6 +175,16 @@ async fn create_test_actor(
|
||||
"/tmp/test-session",
|
||||
)),
|
||||
)),
|
||||
graph_enabled: false,
|
||||
graph_tracker: Arc::new(parking_lot::Mutex::new(
|
||||
crate::session::graph_tracker::GraphTracker::new(std::env::temp_dir()),
|
||||
)),
|
||||
graph_concurrency: 1,
|
||||
graph_node_rounds: 3,
|
||||
graph_replan_cap: 3,
|
||||
graph_optimizer_enabled: false,
|
||||
graph_project_dir: None,
|
||||
graph_project_lock: std::cell::RefCell::new(None),
|
||||
goal_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
@@ -476,6 +488,7 @@ async fn create_test_actor_with_memory(
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: std::num::NonZeroU64::new(context_window)
|
||||
.expect("test context_window must be non-zero"),
|
||||
@@ -499,6 +512,7 @@ async fn create_test_actor_with_memory(
|
||||
rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(),
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state,
|
||||
@@ -611,6 +625,16 @@ async fn create_test_actor_with_memory(
|
||||
"/tmp/test-session",
|
||||
)),
|
||||
)),
|
||||
graph_enabled: false,
|
||||
graph_tracker: Arc::new(parking_lot::Mutex::new(
|
||||
crate::session::graph_tracker::GraphTracker::new(std::env::temp_dir()),
|
||||
)),
|
||||
graph_concurrency: 1,
|
||||
graph_node_rounds: 3,
|
||||
graph_replan_cap: 3,
|
||||
graph_optimizer_enabled: false,
|
||||
graph_project_dir: None,
|
||||
graph_project_lock: std::cell::RefCell::new(None),
|
||||
goal_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
@@ -1220,6 +1244,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
|
||||
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,
|
||||
@@ -1243,6 +1268,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
|
||||
rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(),
|
||||
auth_method_id: test_auth_method_id("cached_token"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
auth_manager: {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mgr = std::sync::Arc::new(crate::auth::AuthManager::new(
|
||||
@@ -1360,6 +1386,16 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
|
||||
"/tmp/test-session",
|
||||
)),
|
||||
)),
|
||||
graph_enabled: false,
|
||||
graph_tracker: Arc::new(parking_lot::Mutex::new(
|
||||
crate::session::graph_tracker::GraphTracker::new(std::env::temp_dir()),
|
||||
)),
|
||||
graph_concurrency: 1,
|
||||
graph_node_rounds: 3,
|
||||
graph_replan_cap: 3,
|
||||
graph_optimizer_enabled: false,
|
||||
graph_project_dir: None,
|
||||
graph_project_lock: std::cell::RefCell::new(None),
|
||||
goal_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
|
||||
@@ -104,6 +104,7 @@ async fn create_test_actor_with_memory(
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: std::num::NonZeroU64::new(context_window)
|
||||
.expect("test context_window must be non-zero"),
|
||||
@@ -126,6 +127,7 @@ async fn create_test_actor_with_memory(
|
||||
},
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state,
|
||||
@@ -235,6 +237,16 @@ async fn create_test_actor_with_memory(
|
||||
"/tmp/test-session",
|
||||
)),
|
||||
)),
|
||||
graph_enabled: false,
|
||||
graph_tracker: Arc::new(parking_lot::Mutex::new(
|
||||
crate::session::graph_tracker::GraphTracker::new(std::env::temp_dir()),
|
||||
)),
|
||||
graph_concurrency: 1,
|
||||
graph_node_rounds: 3,
|
||||
graph_replan_cap: 3,
|
||||
graph_optimizer_enabled: false,
|
||||
graph_project_dir: None,
|
||||
graph_project_lock: std::cell::RefCell::new(None),
|
||||
goal_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
|
||||
+11
@@ -78,6 +78,7 @@ pub(super) async fn make_replay_send_update_fixture() -> ReplaySendUpdateFixture
|
||||
},
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state,
|
||||
@@ -181,6 +182,16 @@ pub(super) async fn make_replay_send_update_fixture() -> ReplaySendUpdateFixture
|
||||
"/tmp/test-session",
|
||||
)),
|
||||
)),
|
||||
graph_enabled: false,
|
||||
graph_tracker: Arc::new(parking_lot::Mutex::new(
|
||||
crate::session::graph_tracker::GraphTracker::new(std::env::temp_dir()),
|
||||
)),
|
||||
graph_concurrency: 1,
|
||||
graph_node_rounds: 3,
|
||||
graph_replan_cap: 3,
|
||||
graph_optimizer_enabled: false,
|
||||
graph_project_dir: None,
|
||||
graph_project_lock: std::cell::RefCell::new(None),
|
||||
goal_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
|
||||
+990
@@ -0,0 +1,990 @@
|
||||
//! 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 kigi_test_support::EnvGuard;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// The host BOTH halves of the `anthropic` / `claude-pro-max` collision route
|
||||
/// to, derived from the registry (as the sibling at
|
||||
/// [`oauth_platform_models_keep_a_live_resolver_from_their_own_pool`] does) so
|
||||
/// the fixture cannot drift, with the twin agreement asserted rather than
|
||||
/// assumed — the collision is only a collision because both platforms serve the
|
||||
/// same host.
|
||||
fn anthropic_collision_host() -> String {
|
||||
let oauth_host = kigi_models::PlatformId::ClaudeProMax.base_url();
|
||||
assert_eq!(
|
||||
kigi_models::PlatformId::Anthropic.base_url(),
|
||||
oauth_host,
|
||||
"the API-key platform and its subscription-OAuth twin must serve the same host, \
|
||||
or this fixture is not testing the dual-credential collision"
|
||||
);
|
||||
oauth_host
|
||||
}
|
||||
|
||||
/// Ambient BYOK env unset. `resolve_model_auth_facts` probes `std::env::var` at
|
||||
/// call time, so a developer (or CI) holding `ANTHROPIC_API_KEY` flips the
|
||||
/// fixture to `Byok` and switches off the session-token gate for a reason that
|
||||
/// has nothing to do with the platform lookup under test. Every holder must be
|
||||
/// `#[serial]`.
|
||||
fn anthropic_collision_env_guard() -> [EnvGuard; 2] {
|
||||
[
|
||||
EnvGuard::unset("ANTHROPIC_API_KEY"),
|
||||
EnvGuard::unset("KIGI_CODE_BASE_URL"),
|
||||
]
|
||||
}
|
||||
|
||||
#[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 {
|
||||
// The base URL is the platform's OWN registry host, exactly as
|
||||
// `models_fetch::platform_fetch_base` builds every fetched entry —
|
||||
// derived here rather than hard-coded so the fixture cannot drift
|
||||
// from the registry (L10 compares against precisely this).
|
||||
for (catalog_key, slug) in [
|
||||
("claude-pro-max/claude-opus-4-8", "claude-opus-4-8"),
|
||||
("github-copilot/gpt-4.1", "gpt-4.1"),
|
||||
("xai-grok/grok-4-latest", "grok-4-latest"),
|
||||
("openai-codex/gpt-5.5", "gpt-5.5"),
|
||||
] {
|
||||
let base_url = kigi_models::parse_managed_model_key(catalog_key)
|
||||
.expect("managed key")
|
||||
.0
|
||||
.base_url();
|
||||
let base_url = base_url.as_str();
|
||||
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 = actor
|
||||
.credential_authority()
|
||||
.manager_for(
|
||||
kigi_models::parse_managed_model_key(catalog_key).map(|(p, _)| p),
|
||||
base_url,
|
||||
)
|
||||
.expect("an OAuth platform on its own host 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
|
||||
/// `CredentialAuthority::bearer_resolver_for`). `SamplingClient::post` REPLACES
|
||||
/// the request's auth header from the resolver, so an API-key-platform aux model
|
||||
/// would have its own key overwritten by the Kimi bearer ON THE AUX HOST.
|
||||
///
|
||||
/// Revert-to-red: make `CredentialAuthority::governing_manager`'s
|
||||
/// `Some(platform) => None` arm return `self.primary.clone()` and the deepseek /
|
||||
/// openai / moonshot rows resolve `KIMI_TOKEN`.
|
||||
#[test]
|
||||
fn aux_bearer_resolver_clears_the_session_resolver_off_a_third_party_aux_host() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let primary = Arc::new(crate::auth::AuthManager::new(
|
||||
dir.path(),
|
||||
crate::auth::KimiCodeConfig::default(),
|
||||
));
|
||||
primary.hot_swap(crate::auth::KimiAuth {
|
||||
key: KIMI_TOKEN.to_string(),
|
||||
auth_mode: crate::auth::AuthMode::OAuth,
|
||||
..crate::auth::KimiAuth::test_default()
|
||||
});
|
||||
let authority = crate::auth::credential_authority::CredentialAuthority::new(
|
||||
crate::agent::config::EndpointsConfig::default(),
|
||||
Some(primary),
|
||||
);
|
||||
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!(
|
||||
authority
|
||||
.bearer_resolver_for(platform(key), base_url)
|
||||
.is_none(),
|
||||
"LEAK: an aux model on {base_url} must not inherit the session bearer resolver"
|
||||
);
|
||||
}
|
||||
assert!(
|
||||
authority
|
||||
.bearer_resolver_for(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 = authority
|
||||
.bearer_resolver_for(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 on ITS OWN host gets a LIVE resolver over
|
||||
// its own pool (empty here), never the Kimi primary. L10: the same model
|
||||
// redirected to a third-party host gets NOTHING.
|
||||
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 p = platform(key).expect("managed key");
|
||||
let resolved = authority
|
||||
.bearer_resolver_for(Some(p), &p.base_url())
|
||||
.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"
|
||||
);
|
||||
assert!(
|
||||
authority
|
||||
.bearer_resolver_for(Some(p), "https://example.invalid/v1")
|
||||
.is_none(),
|
||||
"LEAK ({key}): an OAuth aux model redirected off its own host gets nothing"
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// H4 — TWO CONCURRENT SESSIONS on a colliding slug. The model→platform lookup
|
||||
/// used to key on `ModelsManager::current_model_id()`, a single PROCESS-GLOBAL
|
||||
/// `RwLock<acp::ModelId>` written by whichever session switched last. With one
|
||||
/// session on `xai-grok/grok-4.5` and another on `xai/grok-4.5` — same routing
|
||||
/// slug, by design — the loser resolved the OTHER session's platform:
|
||||
/// the subscription session lost its live resolver (unrecoverable 401 ~1h in)
|
||||
/// and the API-key session got the pooled OAuth bearer stamped over its own
|
||||
/// `sk-…` key, which the provider rejects.
|
||||
///
|
||||
/// Here the global cell is deliberately set to the API-key twin for BOTH
|
||||
/// sessions (last writer wins, and it was the API-key one). Each session must
|
||||
/// still resolve ITS OWN selection.
|
||||
///
|
||||
/// Revert-to-red: replace `self.selected_catalog_key()` in
|
||||
/// `SessionActor::model_platform` with
|
||||
/// `Some(self.models_manager.current_model_id().0.as_ref())`. Under THIS
|
||||
/// fixture both of the subscription session's assertions fail. (L: the fixture
|
||||
/// is what makes that true — `managed_entry` carries `api_key: None` and
|
||||
/// `actor_with_catalog` pins `NotByok`, which is exactly what a FETCHED registry
|
||||
/// entry resolves to. A user who additionally sets `ANTHROPIC_API_KEY` /
|
||||
/// `[model.*] env_key` classifies `Byok`, the gate is inactive for that reason
|
||||
/// alone, and only the `anthropic_oauth` assertion would still catch the
|
||||
/// mis-resolution — hence the env guard below.)
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
#[serial_test::serial]
|
||||
async fn concurrent_sessions_on_a_colliding_slug_each_resolve_their_own_platform() {
|
||||
let _env = anthropic_collision_env_guard();
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
let api_key_twin = "anthropic/claude-opus-4-8";
|
||||
let oauth_twin = "claude-pro-max/claude-opus-4-8";
|
||||
let slug = "claude-opus-4-8";
|
||||
let host = &anthropic_collision_host();
|
||||
let catalog = || {
|
||||
vec![
|
||||
// API-key platform FIRST, exactly as `PlatformId::ALL` orders them.
|
||||
managed_entry(api_key_twin, slug, host),
|
||||
managed_entry(oauth_twin, slug, host),
|
||||
]
|
||||
};
|
||||
|
||||
let (_d1, subscription, _r1) =
|
||||
actor_with_catalog(catalog(), oauth_twin, "unused").await;
|
||||
let (_d2, api_key, _r2) =
|
||||
actor_with_catalog(catalog(), api_key_twin, "sk-ant-user").await;
|
||||
// The other session switched last: the process-global cell now names
|
||||
// the API-key twin for BOTH.
|
||||
for actor in [&subscription, &api_key] {
|
||||
actor
|
||||
.models_manager
|
||||
.set_current_model_id(acp::ModelId::new(api_key_twin.to_string()));
|
||||
}
|
||||
|
||||
let sub_cfg = subscription.reconstruct_full_config().await;
|
||||
let resolver = sub_cfg.bearer_resolver.as_ref().expect(
|
||||
"the subscription session must keep its own live bearer_resolver even when \
|
||||
another session switched the process-global model last",
|
||||
);
|
||||
assert_ne!(
|
||||
resolver.current_bearer(),
|
||||
Some(KIMI_TOKEN.to_string()),
|
||||
"it must read the claude-pro-max pool, never the Kimi primary"
|
||||
);
|
||||
assert!(
|
||||
sub_cfg.anthropic_oauth,
|
||||
"the Claude OAuth Messages adaptation must follow the SUBSCRIPTION session"
|
||||
);
|
||||
|
||||
let api_cfg = api_key.reconstruct_full_config().await;
|
||||
assert!(
|
||||
api_cfg.bearer_resolver.is_none(),
|
||||
"LEAK: the API-key session must get no session bearer_resolver"
|
||||
);
|
||||
assert!(
|
||||
!api_cfg.anthropic_oauth,
|
||||
"the API-key session must not get the OAuth Messages adaptation"
|
||||
);
|
||||
assert_eq!(
|
||||
api_cfg.api_key.as_deref(),
|
||||
Some("sk-ant-user"),
|
||||
"the API-key session keeps its own provider key"
|
||||
);
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// H4 in LEADER mode, where `agent/handlers/model_switch.rs` skips
|
||||
/// `set_current_model_id` ENTIRELY, so the process-global cell is frozen at the
|
||||
/// startup default for the whole process lifetime. `platform_for_slug` then
|
||||
/// fell through to the `.rev()` scan, which returns the LAST match — the OAuth
|
||||
/// twin — so a Leader-mode session on the API-KEY twin was handed the pooled
|
||||
/// OAuth bearer plus the Messages adaptation, and Anthropic rejects both.
|
||||
///
|
||||
/// Revert-to-red: same edit as above; with the global cell naming the startup
|
||||
/// default (not in this catalog) the `.rev()` fallback resolves
|
||||
/// `claude-pro-max/*` and both assertions fail.
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
#[serial_test::serial]
|
||||
async fn leader_mode_session_resolves_its_own_platform_without_the_global_cell() {
|
||||
let _env = anthropic_collision_env_guard();
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
let slug = "claude-opus-4-8";
|
||||
let host = &anthropic_collision_host();
|
||||
let (_dir, actor, _rx) = actor_with_catalog(
|
||||
vec![
|
||||
managed_entry("anthropic/claude-opus-4-8", slug, host),
|
||||
managed_entry("claude-pro-max/claude-opus-4-8", slug, host),
|
||||
],
|
||||
"anthropic/claude-opus-4-8",
|
||||
"sk-ant-user",
|
||||
)
|
||||
.await;
|
||||
// Leader mode never writes the global cell: it still names the
|
||||
// startup default, which is not in this catalog at all.
|
||||
assert!(
|
||||
!actor
|
||||
.models_manager
|
||||
.models()
|
||||
.contains_key(actor.models_manager.current_model_id().0.as_ref()),
|
||||
"precondition: the process-global model id is stale (Leader mode)"
|
||||
);
|
||||
|
||||
let cfg = actor.reconstruct_full_config().await;
|
||||
assert!(
|
||||
cfg.bearer_resolver.is_none(),
|
||||
"LEAK: a Leader-mode API-key session must get no session bearer_resolver"
|
||||
);
|
||||
assert!(
|
||||
!cfg.anthropic_oauth,
|
||||
"a Leader-mode API-key session must not get the OAuth Messages adaptation"
|
||||
);
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// L13 — a user whose ACP auth method is an API-KEY registry platform (e.g.
|
||||
/// `deepseek`) can still SELECT a subscription-OAuth model, and the chokepoint
|
||||
/// hands it that platform's pooled bearer as the request's `api_key`. The gate
|
||||
/// keys on the primary method, which is not session-based, so the config used
|
||||
/// to carry NO `bearer_resolver`: the pooled token froze at selection time and
|
||||
/// the session died with an unrecoverable 401 once it expired (~1h).
|
||||
///
|
||||
/// The model's own credential now makes the gate session-based, confined to
|
||||
/// that platform's own host by the gate's `credential_class` conjunct.
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn oauth_model_under_an_api_key_auth_method_keeps_its_pooled_resolver() {
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
let base_url = kigi_models::PlatformId::ClaudeProMax.base_url();
|
||||
let (_dir, actor, _rx) = actor_with_catalog(
|
||||
vec![managed_entry(
|
||||
"claude-pro-max/claude-opus-4-8",
|
||||
"claude-opus-4-8",
|
||||
&base_url,
|
||||
)],
|
||||
"claude-pro-max/claude-opus-4-8",
|
||||
"unused",
|
||||
)
|
||||
.await;
|
||||
// The PRIMARY ACP method is an API-key registry platform login.
|
||||
actor
|
||||
.auth_method_id
|
||||
.store(Some(std::sync::Arc::new(acp::AuthMethodId::new(
|
||||
"deepseek",
|
||||
))));
|
||||
|
||||
let cfg = actor.reconstruct_full_config().await;
|
||||
let resolver = cfg.bearer_resolver.as_ref().expect(
|
||||
"a subscription-OAuth model keeps a live resolver whatever the primary \
|
||||
ACP auth method is, or it cannot refresh mid-session",
|
||||
);
|
||||
assert_ne!(
|
||||
resolver.current_bearer(),
|
||||
Some(KIMI_TOKEN.to_string()),
|
||||
"and it reads the claude-pro-max pool, never the primary"
|
||||
);
|
||||
|
||||
// An API-key-platform model under the same method stays resolver-free.
|
||||
let (_d2, deepseek, _r2) = actor_with_catalog(
|
||||
vec![managed_entry(
|
||||
"deepseek/deepseek-chat",
|
||||
"deepseek-chat",
|
||||
"https://api.deepseek.com/v1",
|
||||
)],
|
||||
"deepseek/deepseek-chat",
|
||||
"sk-deepseek",
|
||||
)
|
||||
.await;
|
||||
deepseek
|
||||
.auth_method_id
|
||||
.store(Some(std::sync::Arc::new(acp::AuthMethodId::new(
|
||||
"deepseek",
|
||||
))));
|
||||
assert!(
|
||||
deepseek
|
||||
.reconstruct_full_config()
|
||||
.await
|
||||
.bearer_resolver
|
||||
.is_none(),
|
||||
"LEAK: an API-key-platform model must never get a session resolver"
|
||||
);
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// H-b — a `None` or STALE per-session catalog key must REFUSE, not degrade to
|
||||
/// the subscription-OAuth twin.
|
||||
///
|
||||
/// `model_platform` falls through to `resolve_catalog_key`'s `.rev()` scan when
|
||||
/// the session's own key does not name the slug, and that scan returns the LAST
|
||||
/// match — the OAuth twin, because `PlatformId::ALL` orders every API-key
|
||||
/// platform first. Combined with the L13 disjunct (a model whose own credential
|
||||
/// is a pooled OAuth session is session-based BY ITSELF), an API-KEY session on
|
||||
/// `anthropic/claude-opus-4-8` with no per-session key got
|
||||
/// `is_session_based = true`, `credential_class = Pooled` (same
|
||||
/// host) and, at `NotByok`, an ACTIVE gate — so `manager_for` handed it the
|
||||
/// Claude POOLED manager, whose `bearer_resolver` REPLACES the user's own
|
||||
/// `sk-ant-…` on the wire, plus the OAuth Messages adaptation. Anthropic rejects
|
||||
/// both. This is exactly what H4 prevents, reached through the `None` path.
|
||||
///
|
||||
/// A key that is absent or names a different model is not evidence for either
|
||||
/// twin: resolve to NO platform, which the chokepoint then decides purely by the
|
||||
/// ENDPOINT (the OAuth host is not this session's coding endpoint ⇒ nothing
|
||||
/// rides).
|
||||
///
|
||||
/// Revert-to-red (production, compiles): delete the
|
||||
/// `platform.oauth().is_some() && !disambiguated && slug_collides_across_platforms(..)`
|
||||
/// refusal from `crate::agent::models::platform_for_slug` and every
|
||||
/// `bearer_resolver` / `anthropic_oauth` assertion below fails.
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
#[serial_test::serial]
|
||||
async fn a_missing_or_stale_session_key_refuses_instead_of_guessing_the_oauth_twin() {
|
||||
let _env = anthropic_collision_env_guard();
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
let api_key_twin = "anthropic/claude-opus-4-8";
|
||||
let slug = "claude-opus-4-8";
|
||||
let host = anthropic_collision_host();
|
||||
let catalog = || {
|
||||
vec![
|
||||
// API-key platform FIRST, exactly as `PlatformId::ALL` orders them.
|
||||
managed_entry(api_key_twin, slug, &host),
|
||||
managed_entry("claude-pro-max/claude-opus-4-8", slug, &host),
|
||||
]
|
||||
};
|
||||
|
||||
for (case, stale_key) in [
|
||||
// No key at all: a session spawned on a model that left the
|
||||
// catalog, or one an older build never seeded.
|
||||
("absent", None),
|
||||
// Stale: an `OverrideModelName` rename, or a key naming a model
|
||||
// this session is no longer on.
|
||||
("stale", Some("claude-pro-max/some-other-model".to_string())),
|
||||
] {
|
||||
let (_dir, actor, _rx) =
|
||||
actor_with_catalog(catalog(), api_key_twin, "sk-ant-user").await;
|
||||
*actor.selected_catalog_key.borrow_mut() = stale_key;
|
||||
|
||||
let cfg = actor.reconstruct_full_config().await;
|
||||
assert!(
|
||||
cfg.bearer_resolver.is_none(),
|
||||
"{case}: LEAK — an unresolvable selection must get NO session bearer \
|
||||
resolver; the pooled OAuth bearer would REPLACE the user's own key"
|
||||
);
|
||||
assert!(
|
||||
!cfg.anthropic_oauth,
|
||||
"{case}: nor the Claude OAuth Messages adaptation"
|
||||
);
|
||||
assert_eq!(
|
||||
cfg.api_key.as_deref(),
|
||||
Some("sk-ant-user"),
|
||||
"{case}: the user's own provider key must survive untouched"
|
||||
);
|
||||
assert!(
|
||||
actor
|
||||
.credential_authority()
|
||||
.manager_for(
|
||||
crate::agent::models::platform_for_slug(
|
||||
&actor.models_manager.models(),
|
||||
actor.selected_catalog_key().as_deref(),
|
||||
slug,
|
||||
),
|
||||
&host,
|
||||
)
|
||||
.is_none(),
|
||||
"{case}: and no manager either — refuse, never guess"
|
||||
);
|
||||
}
|
||||
|
||||
// …while a session that DID select the OAuth twin still gets its
|
||||
// pooled resolver: the refusal is about the guess, not the platform.
|
||||
let (_dir, selected, _rx) =
|
||||
actor_with_catalog(catalog(), "claude-pro-max/claude-opus-4-8", "unused").await;
|
||||
let cfg = selected.reconstruct_full_config().await;
|
||||
assert!(
|
||||
cfg.bearer_resolver.is_some() && cfg.anthropic_oauth,
|
||||
"a DELIBERATE subscription selection keeps its pooled resolver and \
|
||||
adaptation (this is what makes the refusals above meaningful)"
|
||||
);
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// H-c — coverage for the FIRST of the two production writers of
|
||||
/// `selected_catalog_key`: the spawn seed
|
||||
/// (`crate::agent::models::selected_catalog_key_for_spawn`, called from
|
||||
/// `spawn.rs`). Every other test in this module sets the field by hand, so a
|
||||
/// wrong seed was silent.
|
||||
///
|
||||
/// Both spawn shapes are covered: a FRESH session, spawned on the catalog key
|
||||
/// the picker resolved, and a RESUME/LOAD, which spawns with the RAW persisted
|
||||
/// `summary.current_model_id` — a BARE routing slug after any `SetSessionModel`,
|
||||
/// since `handle_set_session_model` persists `sampling_config.model`. This seed
|
||||
/// is where that slug becomes a key. The assertion is end-to-end: the seeded key
|
||||
/// is fed to the very function the auth layer keys on.
|
||||
#[test]
|
||||
fn spawn_seeds_the_session_key_the_auth_layer_keys_on() {
|
||||
let slug = "claude-opus-4-8";
|
||||
let host = anthropic_collision_host();
|
||||
let api_key_twin = "anthropic/claude-opus-4-8";
|
||||
let oauth_twin = "claude-pro-max/claude-opus-4-8";
|
||||
let models: indexmap::IndexMap<String, crate::agent::config::ModelEntry> = [
|
||||
managed_entry(api_key_twin, slug, &host),
|
||||
managed_entry(oauth_twin, slug, &host),
|
||||
]
|
||||
.into_iter()
|
||||
.collect();
|
||||
|
||||
// FRESH: spawned on the catalog key. Idempotent, and it disambiguates.
|
||||
for selected in [api_key_twin, oauth_twin] {
|
||||
let seeded = crate::agent::models::selected_catalog_key_for_spawn(
|
||||
&models,
|
||||
&acp::ModelId::new(selected.to_string()),
|
||||
);
|
||||
assert_eq!(
|
||||
seeded.as_deref(),
|
||||
Some(selected),
|
||||
"a fresh session must record the catalog key it was spawned with"
|
||||
);
|
||||
assert_eq!(
|
||||
crate::agent::models::platform_for_slug(&models, seeded.as_deref(), slug),
|
||||
kigi_models::parse_managed_model_key(selected).map(|(p, _)| p),
|
||||
"…and that key must resolve THIS session's own platform for the bare slug"
|
||||
);
|
||||
}
|
||||
|
||||
// RESUME/LOAD: `acp_agent::load_session` spawns with the RAW persisted id,
|
||||
// which after any model switch is the bare routing slug. THIS seed is what
|
||||
// turns it into a key — the picker's `.rev()` answer, which is the resume
|
||||
// default for a collided slug.
|
||||
let resumed = crate::agent::models::selected_catalog_key_for_spawn(
|
||||
&models,
|
||||
&acp::ModelId::new(slug.to_string()),
|
||||
);
|
||||
assert_eq!(
|
||||
resumed.as_deref(),
|
||||
Some(oauth_twin),
|
||||
"a bare persisted slug resolves through the picker's own lookup"
|
||||
);
|
||||
|
||||
// A model that is no longer in the catalog seeds NOTHING, which (H-b) then
|
||||
// refuses rather than guessing a twin.
|
||||
assert_eq!(
|
||||
crate::agent::models::selected_catalog_key_for_spawn(
|
||||
&models,
|
||||
&acp::ModelId::new("gone/model".to_string()),
|
||||
),
|
||||
None,
|
||||
"a model that left the catalog must not seed a key"
|
||||
);
|
||||
}
|
||||
|
||||
/// H-c — coverage for the SECOND production writer: `SetSessionModel`
|
||||
/// (`handle_set_session_model`), the picker's own path. The existing test
|
||||
/// through this handler passes `None`, so a handler that dropped the key on the
|
||||
/// floor stayed green.
|
||||
///
|
||||
/// End-to-end: after the switch the session's per-turn config must carry the
|
||||
/// SELECTED twin's pooled resolver and adaptation, even though the bare slug in
|
||||
/// the config is ambiguous.
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
#[serial_test::serial]
|
||||
async fn set_session_model_records_the_key_the_next_turn_resolves_on() {
|
||||
let _env = anthropic_collision_env_guard();
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
let slug = "claude-opus-4-8";
|
||||
let host = anthropic_collision_host();
|
||||
let api_key_twin = "anthropic/claude-opus-4-8";
|
||||
let oauth_twin = "claude-pro-max/claude-opus-4-8";
|
||||
let (_dir, actor, _rx) = actor_with_catalog(
|
||||
vec![
|
||||
managed_entry(api_key_twin, slug, &host),
|
||||
managed_entry(oauth_twin, slug, &host),
|
||||
],
|
||||
api_key_twin,
|
||||
"sk-ant-user",
|
||||
)
|
||||
.await;
|
||||
|
||||
// Switch to the SUBSCRIPTION twin, exactly as
|
||||
// `agent/handlers/model_switch.rs` does: the ambiguous slug in the
|
||||
// sampler config plus the catalog KEY the picker resolved.
|
||||
let models = actor.models_manager.models();
|
||||
let entry = models.get(oauth_twin).expect("catalog entry");
|
||||
let sampler = crate::agent::config::sampling_config_for_model(
|
||||
entry,
|
||||
crate::agent::config::resolve_credentials(entry, None),
|
||||
None,
|
||||
);
|
||||
actor
|
||||
.handle_set_session_model(
|
||||
sampler,
|
||||
Some(oauth_twin.to_string()),
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
85,
|
||||
)
|
||||
.await
|
||||
.expect("model switch");
|
||||
|
||||
assert_eq!(
|
||||
actor.selected_catalog_key().as_deref(),
|
||||
Some(oauth_twin),
|
||||
"SetSessionModel must record the picker's catalog key"
|
||||
);
|
||||
let cfg = actor.reconstruct_full_config().await;
|
||||
assert!(
|
||||
cfg.bearer_resolver.is_some(),
|
||||
"the switched-to subscription model must keep a live pooled resolver"
|
||||
);
|
||||
assert!(
|
||||
cfg.anthropic_oauth,
|
||||
"…and the Claude OAuth Messages adaptation"
|
||||
);
|
||||
|
||||
// And back: switching to the API-key twin must UNDO both.
|
||||
let entry = models.get(api_key_twin).expect("catalog entry");
|
||||
let sampler = crate::agent::config::sampling_config_for_model(
|
||||
entry,
|
||||
crate::agent::config::resolve_credentials(entry, None),
|
||||
None,
|
||||
);
|
||||
actor
|
||||
.handle_set_session_model(
|
||||
sampler,
|
||||
Some(api_key_twin.to_string()),
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
85,
|
||||
)
|
||||
.await
|
||||
.expect("model switch");
|
||||
let cfg = actor.reconstruct_full_config().await;
|
||||
assert!(
|
||||
cfg.bearer_resolver.is_none() && !cfg.anthropic_oauth,
|
||||
"LEAK: switching back to the API-key twin must drop the pooled resolver \
|
||||
and the OAuth adaptation"
|
||||
);
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// H-c — `OverrideModelName` is the one command that rewrites
|
||||
/// `SamplingConfig::model` WITHOUT going through `SetSessionModel`, so it used
|
||||
/// to leave `selected_catalog_key` naming a model the session is no longer on.
|
||||
/// It must keep the field consistent: KEEP when the key still names the new
|
||||
/// routing name, CLEAR otherwise — never re-resolve, which would put the
|
||||
/// `.rev()` guess into the field the rule treats as a deliberate selection.
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
#[serial_test::serial]
|
||||
async fn override_model_name_keeps_the_session_key_consistent() {
|
||||
let _env = anthropic_collision_env_guard();
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
let slug = "claude-opus-4-8";
|
||||
let host = anthropic_collision_host();
|
||||
let oauth_twin = "claude-pro-max/claude-opus-4-8";
|
||||
let (_dir, actor, _rx) = actor_with_catalog(
|
||||
vec![
|
||||
managed_entry("anthropic/claude-opus-4-8", slug, &host),
|
||||
managed_entry(oauth_twin, slug, &host),
|
||||
],
|
||||
oauth_twin,
|
||||
"unused",
|
||||
)
|
||||
.await;
|
||||
|
||||
// A rename to the SAME model's routing slug (or to its catalog key)
|
||||
// keeps the selection.
|
||||
for same in [slug, oauth_twin] {
|
||||
actor.retain_selected_catalog_key_for(same);
|
||||
assert_eq!(
|
||||
actor.selected_catalog_key().as_deref(),
|
||||
Some(oauth_twin),
|
||||
"{same}: still names the selected entry — keep it"
|
||||
);
|
||||
}
|
||||
|
||||
// A rename to a DIFFERENT name makes the key stale: clear it, so the
|
||||
// collided slug refuses (H-b) instead of resolving the old model.
|
||||
actor.retain_selected_catalog_key_for("some-harness-model-name");
|
||||
assert_eq!(
|
||||
actor.selected_catalog_key(),
|
||||
None,
|
||||
"a stale key must be cleared, not carried into the next turn's \
|
||||
platform lookup"
|
||||
);
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// M3 — the FIRST-PARTY aux case must honour the session gate, which is what
|
||||
/// the old shape did implicitly.
|
||||
///
|
||||
/// `stamp_session_local_sampler_fields` used to copy
|
||||
/// `active_session_config.bearer_resolver`, and that field is `None` whenever
|
||||
/// the gate is inactive. Re-pointing the aux resolver at the chokepoint (the
|
||||
/// LEAK 1b fix) made it `Some(primary)` for the session's own coding endpoint
|
||||
/// REGARDLESS of the gate — so a BYOK / api-key session with a `[model.*]` aux
|
||||
/// entry carrying its own key on that endpoint had that key REPLACED by the
|
||||
/// primary bearer on every image-describe / auto-mode-classifier / summary
|
||||
/// request (`SamplingClient::post` overrides the auth header from the resolver).
|
||||
///
|
||||
/// Revert-to-red (production, compiles): delete the
|
||||
/// `if is_primary_channel && !SessionTokenAuthGate::new(…).active()` early
|
||||
/// return from `sampler_turn::aux_bearer_resolver_for` and the first two rows
|
||||
/// below resolve `KIMI_TOKEN`.
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn first_party_aux_resolver_honours_the_session_gate() {
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
let coding_host = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url;
|
||||
let aux_slug = "kigi-aux";
|
||||
let mut info = crate::agent::config::ModelInfo::fallback(aux_slug);
|
||||
info.id = None; // a `[model.kigi-aux]` block, not a registry entry
|
||||
info.base_url = coding_host.to_string();
|
||||
let aux_entry = crate::agent::config::ModelEntry {
|
||||
info,
|
||||
api_key: None,
|
||||
env_key: None,
|
||||
api_base_url: None,
|
||||
};
|
||||
|
||||
// (case, ACP auth method, the aux model's own BYOK status, expected)
|
||||
for (case, auth_method, byok, expect_resolver) in [
|
||||
(
|
||||
"an API-key session: the aux model's own key must survive",
|
||||
"deepseek",
|
||||
crate::agent::auth_method::ModelByok::NotByok,
|
||||
false,
|
||||
),
|
||||
(
|
||||
"a BYOK aux entry under a session method: its env_key wins",
|
||||
"cached_token",
|
||||
crate::agent::auth_method::ModelByok::Byok,
|
||||
false,
|
||||
),
|
||||
(
|
||||
"the first-party subscription aux channel: byte-identical",
|
||||
"cached_token",
|
||||
crate::agent::auth_method::ModelByok::NotByok,
|
||||
true,
|
||||
),
|
||||
] {
|
||||
let (_dir, actor, _rx) = actor_with_catalog(
|
||||
vec![(aux_slug.to_string(), aux_entry.clone())],
|
||||
aux_slug,
|
||||
"",
|
||||
)
|
||||
.await;
|
||||
actor
|
||||
.auth_method_id
|
||||
.store(Some(Arc::new(acp::AuthMethodId::new(auth_method))));
|
||||
actor.model_auth_facts.replace(Some((
|
||||
aux_slug.to_string(),
|
||||
crate::agent::config::ModelAuthFacts {
|
||||
byok,
|
||||
auth_scheme: Default::default(),
|
||||
},
|
||||
)));
|
||||
|
||||
let resolved = actor.aux_bearer_resolver(aux_slug, coding_host);
|
||||
assert_eq!(
|
||||
resolved.is_some(),
|
||||
expect_resolver,
|
||||
"{case}: aux resolver presence on the session's own endpoint"
|
||||
);
|
||||
if let Some(resolver) = resolved {
|
||||
assert_eq!(
|
||||
resolver.current_bearer(),
|
||||
Some(KIMI_TOKEN.to_string()),
|
||||
"{case}: and when it IS kept it is the primary's, live"
|
||||
);
|
||||
}
|
||||
}
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// M-aux (REGRESSION this remediation introduced) — an aux call must NOT evict
|
||||
/// the SESSION model's memoized auth facts.
|
||||
///
|
||||
/// `SessionActor::model_auth_facts` is a SINGLE slot. When `aux_bearer_resolver`
|
||||
/// began asking it about the AUX slug, a definite result overwrote the session
|
||||
/// model's entry, and:
|
||||
/// (a) the next `reconstruct_full_config` re-paid `load_effective_config()` +
|
||||
/// `resolve_model_list()` — the per-turn disk read M7/M9 removed — on top
|
||||
/// of the one the aux call itself paid; and
|
||||
/// (b) the memo's documented purpose (a transient `Unknown` falling back to
|
||||
/// the last DEFINITE value FOR THE SAME model_id) was defeated: with the
|
||||
/// aux slug in the slot, the session model's `Unknown` degrades to
|
||||
/// `endpoint_is_first_party`, which is `false` for every
|
||||
/// subscription-OAuth host — the session loses its `bearer_resolver` and
|
||||
/// 401s unrecoverably ~1h in, the failure L13 exists to prevent.
|
||||
///
|
||||
/// Round 3's deleted `repoint_aux_bearer_resolver` never touched the memo.
|
||||
///
|
||||
/// Revert-to-red (production, compiles): make `SessionActor::aux_bearer_resolver`
|
||||
/// call `self.model_auth_facts(slug)` instead of `self.aux_model_auth_facts(slug)`
|
||||
/// — the slot then names the aux slug and both assertions below fail.
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
#[serial_test::serial]
|
||||
async fn an_aux_call_does_not_evict_the_session_models_auth_facts() {
|
||||
let _env = anthropic_collision_env_guard();
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
let session_slug = "claude-opus-4-8";
|
||||
let host = anthropic_collision_host();
|
||||
let oauth_twin = "claude-pro-max/claude-opus-4-8";
|
||||
let (_dir, actor, _rx) = actor_with_catalog(
|
||||
vec![managed_entry(oauth_twin, session_slug, &host)],
|
||||
oauth_twin,
|
||||
"unused",
|
||||
)
|
||||
.await;
|
||||
|
||||
// The session model's DEFINITE facts, as a turn would have memoized
|
||||
// them.
|
||||
actor.model_auth_facts.replace(Some((
|
||||
session_slug.to_string(),
|
||||
crate::agent::config::ModelAuthFacts {
|
||||
byok: crate::agent::auth_method::ModelByok::NotByok,
|
||||
auth_scheme: Default::default(),
|
||||
},
|
||||
)));
|
||||
|
||||
// An aux turn: the auto-mode classifier / image-describe slug, which
|
||||
// is NOT the session's model.
|
||||
let _ = actor.aux_bearer_resolver("kigi-aux-classifier", &host);
|
||||
|
||||
let memo = actor.model_auth_facts.borrow();
|
||||
let (cached_id, facts) = memo
|
||||
.as_ref()
|
||||
.expect("the session model's memo must survive an aux call");
|
||||
assert_eq!(
|
||||
cached_id, session_slug,
|
||||
"an aux call evicted the SESSION model's memo: the next turn re-reads \
|
||||
config from disk, and a transient Unknown loses its definite fallback"
|
||||
);
|
||||
assert_eq!(facts.byok, crate::agent::auth_method::ModelByok::NotByok);
|
||||
})
|
||||
.await;
|
||||
}
|
||||
@@ -0,0 +1,437 @@
|
||||
//! 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, the manager lookup 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/M8): 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 per-process temp path that is
|
||||
//! never created, so every pooled manager is empty — exactly what the
|
||||
//! assertions need (a live resolver that is provably NOT the Kimi one) — and
|
||||
//! the binary leaves nothing behind.
|
||||
|
||||
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");
|
||||
// H4: the SESSION owns its selection. The process-global
|
||||
// `ModelsManager::current_model_id()` is deliberately left UNSET (it still
|
||||
// names the startup default, which is not in this catalog) — exactly what
|
||||
// Leader mode produces, since `agent/handlers/model_switch.rs` never calls
|
||||
// `set_current_model_id` there, and what a second concurrent session on a
|
||||
// colliding slug produces (last writer wins). Every assertion below
|
||||
// therefore rides the per-session key, not the global cell.
|
||||
*actor.selected_catalog_key.borrow_mut() = Some(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 the `credential_class` conjunct 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 config.toml `[endpoints] coding_api_base_url`
|
||||
/// deployment, a `KIGI_CODE_BASE_URL` override, or a local dev proxy) it must
|
||||
/// keep one — that is why the predicate is not `is_first_party_url`.
|
||||
///
|
||||
/// Revert-to-red: make `CredentialAuthority::is_session_coding_endpoint` return
|
||||
/// `true` unconditionally and a Kimi resolver lands 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;
|
||||
}
|
||||
|
||||
/// The Kimi / first-party subscription channel must be BYTE-IDENTICAL: the
|
||||
/// session model keeps its live bearer_resolver AND the pre-flight refresh
|
||||
/// still heals a stale buffered key. This is also what proves the Kimi bearer is
|
||||
/// live in every LEAK assertion in this module — it WOULD leak if the guard
|
||||
/// were missing.
|
||||
///
|
||||
/// (L12: the slug-collision commentary that used to sit here belongs to the
|
||||
/// collision tests in `session_bearer_leak_platform_tests`, which is where its
|
||||
/// revert-to-red actually reproduces; on this first-party test it never could.)
|
||||
#[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;
|
||||
}
|
||||
@@ -168,6 +168,7 @@ pub(crate) async fn create_test_actor_ex(
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: std::num::NonZeroU64::new(context_window)
|
||||
.expect("test context_window must be non-zero"),
|
||||
@@ -187,6 +188,7 @@ pub(crate) async fn create_test_actor_ex(
|
||||
},
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state,
|
||||
@@ -286,6 +288,16 @@ pub(crate) async fn create_test_actor_ex(
|
||||
"/tmp/test-session",
|
||||
)),
|
||||
)),
|
||||
graph_enabled: false,
|
||||
graph_tracker: Arc::new(parking_lot::Mutex::new(
|
||||
crate::session::graph_tracker::GraphTracker::new(std::env::temp_dir()),
|
||||
)),
|
||||
graph_concurrency: 1,
|
||||
graph_node_rounds: 3,
|
||||
graph_replan_cap: 3,
|
||||
graph_optimizer_enabled: false,
|
||||
graph_project_dir: None,
|
||||
graph_project_lock: std::cell::RefCell::new(None),
|
||||
goal_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
|
||||
@@ -145,6 +145,13 @@ pub enum SessionCommand {
|
||||
},
|
||||
SetSessionModel {
|
||||
sampling_config: kigi_sampler::SamplerConfig,
|
||||
/// The catalog KEY the picker resolved (`{platform}/{model}` for a
|
||||
/// registry model), which `sampling_config.model` — the bare routing
|
||||
/// slug — cannot express when an API-key platform and its
|
||||
/// subscription-OAuth twin list the same id. The session stores it as
|
||||
/// its OWN selection instead of reading the process-global
|
||||
/// `ModelsManager::current_model_id()` (H4).
|
||||
catalog_key: Option<String>,
|
||||
use_concise: bool,
|
||||
/// When `false`, skip the system prompt rewrite (concise/default swap).
|
||||
/// Set to `false` for forked sessions so mid-session model switches
|
||||
|
||||
@@ -2149,6 +2149,7 @@ mod inline_auto_compact_flow_tests {
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: std::num::NonZeroU64::new(context_window)
|
||||
.expect("test context_window must be non-zero"),
|
||||
@@ -2167,6 +2168,7 @@ mod inline_auto_compact_flow_tests {
|
||||
},
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state,
|
||||
@@ -2268,6 +2270,16 @@ mod inline_auto_compact_flow_tests {
|
||||
"/tmp/test-session",
|
||||
)),
|
||||
)),
|
||||
graph_enabled: false,
|
||||
graph_tracker: Arc::new(parking_lot::Mutex::new(
|
||||
crate::session::graph_tracker::GraphTracker::new(std::env::temp_dir()),
|
||||
)),
|
||||
graph_concurrency: 1,
|
||||
graph_node_rounds: 3,
|
||||
graph_replan_cap: 3,
|
||||
graph_optimizer_enabled: false,
|
||||
graph_project_dir: None,
|
||||
graph_project_lock: std::cell::RefCell::new(None),
|
||||
goal_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
||||
|
||||
@@ -594,10 +594,10 @@ async fn write_patch_file_atomic(path: &Path, body: &str) -> std::io::Result<()>
|
||||
.unwrap_or("goal-classifier.patch");
|
||||
let tmp = dir.join(format!(".{file_name}.{}.tmp", uuid::Uuid::now_v7()));
|
||||
tokio::fs::write(&tmp, body).await?;
|
||||
if let Err(err) = tokio::fs::rename(&tmp, path).await {
|
||||
let _ = tokio::fs::remove_file(&tmp).await;
|
||||
return Err(err);
|
||||
}
|
||||
let dest = path.to_path_buf();
|
||||
tokio::task::spawn_blocking(move || crate::util::fs::replace_file(&tmp, &dest))
|
||||
.await
|
||||
.map_err(std::io::Error::other)??;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@@ -154,7 +154,9 @@ pub enum GoalPauseReason {
|
||||
}
|
||||
|
||||
impl GoalPauseReason {
|
||||
fn to_status(self) -> GoalStatus {
|
||||
/// Also used by the graph tracker (`graph_tracker.rs`), which reuses
|
||||
/// the goal status vocabulary for graph-level pauses.
|
||||
pub(crate) fn to_status(self) -> GoalStatus {
|
||||
match self {
|
||||
Self::User => GoalStatus::UserPaused,
|
||||
Self::BackOff => GoalStatus::BackOffPaused,
|
||||
@@ -166,7 +168,8 @@ impl GoalPauseReason {
|
||||
|
||||
/// Short, stable label stashed in the `GoalPaused` history entry's
|
||||
/// `detail` so the pager's Recent History distinguishes pause causes.
|
||||
fn history_detail(self) -> &'static str {
|
||||
/// Shared with the graph tracker's `GraphPaused` entries.
|
||||
pub(crate) fn history_detail(self) -> &'static str {
|
||||
match self {
|
||||
Self::User => "user",
|
||||
Self::BackOff => "back_off",
|
||||
@@ -892,7 +895,8 @@ impl GoalTracker {
|
||||
};
|
||||
let dest = goal_dir.join(name);
|
||||
let _ = std::fs::create_dir_all(&goal_dir);
|
||||
if std::fs::rename(&src, &dest).is_ok() || copy_no_follow(&src, &dest).is_ok() {
|
||||
if crate::util::fs::replace_file(&src, &dest).is_ok() || copy_no_follow(&src, &dest).is_ok()
|
||||
{
|
||||
append_skeptic_reports(&scratch_root, &dest);
|
||||
o.last_classifier_details_path = Some(dest.to_string_lossy().into_owned());
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user