Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
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 | ||
|
|
7bdf50a716 | ||
|
|
10d0c0816f | ||
|
|
6e26d15428 | ||
|
|
82ebcc38b7 | ||
|
|
0dc98a34a6 | ||
|
|
0692198719 | ||
|
|
77fd457627 |
@@ -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
|
|
||||||
@@ -18,9 +18,12 @@ import) or any `KIMI_*` env var.
|
|||||||
|
|
||||||
- **Zero egress**: outbound connections are limited to
|
- **Zero egress**: outbound connections are limited to
|
||||||
`auth.kimi.com`, `api.kimi.com`, `api.moonshot.cn`, `api.moonshot.ai`,
|
`auth.kimi.com`, `api.kimi.com`, `api.moonshot.cn`, `api.moonshot.ai`,
|
||||||
GitHub Releases domains, and user-configured MCP servers. No telemetry,
|
GitHub Releases domains, user-configured MCP servers, the endpoints of
|
||||||
no analytics, ever. `crates/codegen/kigi-env` is the single home of
|
provider platforms the user has credentialed, and `models.dev` (model
|
||||||
first-party endpoints.
|
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.
|
- **Toolchain**: Rust 1.97.0 (rust-toolchain.toml), edition 2024.
|
||||||
- **Gates** (all must stay green):
|
- **Gates** (all must stay green):
|
||||||
`cargo check --workspace --all-targets`,
|
`cargo check --workspace --all-targets`,
|
||||||
@@ -49,6 +52,17 @@ import) or any `KIMI_*` env var.
|
|||||||
- `third_party/` — vendored Mermaid rendering stack (untouched policy).
|
- `third_party/` — vendored Mermaid rendering stack (untouched policy).
|
||||||
- `bin/protoc` — dotslash launcher used by proto codegen.
|
- `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
|
## Test seams
|
||||||
|
|
||||||
Cross-crate test hooks are behind the `test-support` cargo feature
|
Cross-crate test hooks are behind the `test-support` cargo feature
|
||||||
@@ -56,6 +70,318 @@ Cross-crate test hooks are behind the `test-support` cargo feature
|
|||||||
dependents' `[dev-dependencies]`. Don't expose new test seams as plain
|
dependents' `[dev-dependencies]`. Don't expose new test seams as plain
|
||||||
`#[cfg(test)]` items across crate boundaries.
|
`#[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.
|
||||||
|
|
||||||
|
- Feature flag `KIGI_GRAPH=1` (default off); 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.
|
||||||
|
- `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. API-key `openai` Responses requests carry
|
||||||
|
NONE of this (byte-identical). `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)
|
## Milestones (PRD §8.3)
|
||||||
|
|
||||||
- M0 (done): rename, deletions (voice/telemetry/announcements/marketplace/
|
- M0 (done): rename, deletions (voice/telemetry/announcements/marketplace/
|
||||||
|
|||||||
Generated
+64
-62
@@ -5442,7 +5442,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-acp-lib"
|
name = "kigi-acp-lib"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"agent-client-protocol",
|
"agent-client-protocol",
|
||||||
"async-trait",
|
"async-trait",
|
||||||
@@ -5456,7 +5456,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-agent"
|
name = "kigi-agent"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"chrono",
|
"chrono",
|
||||||
"dirs 6.0.0",
|
"dirs 6.0.0",
|
||||||
@@ -5486,7 +5486,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-agent-lifecycle"
|
name = "kigi-agent-lifecycle"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"async-trait",
|
"async-trait",
|
||||||
"tokio",
|
"tokio",
|
||||||
@@ -5495,7 +5495,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-auth"
|
name = "kigi-auth"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"async-trait",
|
"async-trait",
|
||||||
"http 1.4.2",
|
"http 1.4.2",
|
||||||
@@ -5508,7 +5508,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-bin"
|
name = "kigi-bin"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"clap",
|
"clap",
|
||||||
@@ -5543,7 +5543,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-chat-state"
|
name = "kigi-chat-state"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"indexmap",
|
"indexmap",
|
||||||
"kigi-compaction",
|
"kigi-compaction",
|
||||||
@@ -5560,7 +5560,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-codebase-graph"
|
name = "kigi-codebase-graph"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"ahash",
|
"ahash",
|
||||||
"clap",
|
"clap",
|
||||||
@@ -5596,7 +5596,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-compaction"
|
name = "kigi-compaction"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"async-trait",
|
"async-trait",
|
||||||
@@ -5609,7 +5609,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-config"
|
name = "kigi-config"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"base64",
|
"base64",
|
||||||
"blake3",
|
"blake3",
|
||||||
@@ -5632,7 +5632,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-config-types"
|
name = "kigi-config-types"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"agent-client-protocol",
|
"agent-client-protocol",
|
||||||
"indexmap",
|
"indexmap",
|
||||||
@@ -5646,7 +5646,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-crash-handler"
|
name = "kigi-crash-handler"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"backtrace",
|
"backtrace",
|
||||||
"libc",
|
"libc",
|
||||||
@@ -5657,7 +5657,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-env"
|
name = "kigi-env"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"tracing",
|
"tracing",
|
||||||
"url",
|
"url",
|
||||||
@@ -5665,7 +5665,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-fast-worktree"
|
name = "kigi-fast-worktree"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"bytes",
|
"bytes",
|
||||||
@@ -5697,7 +5697,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-file-utils"
|
name = "kigi-file-utils"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"aws-config",
|
"aws-config",
|
||||||
@@ -5721,7 +5721,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-fsnotify"
|
name = "kigi-fsnotify"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"criterion",
|
"criterion",
|
||||||
"dunce",
|
"dunce",
|
||||||
@@ -5742,7 +5742,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-gix-status"
|
name = "kigi-gix-status"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"gix",
|
"gix",
|
||||||
"kigi-test-utils",
|
"kigi-test-utils",
|
||||||
@@ -5752,7 +5752,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-hooks"
|
name = "kigi-hooks"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"fastrand",
|
"fastrand",
|
||||||
"kigi-config",
|
"kigi-config",
|
||||||
@@ -5771,7 +5771,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-hooks-plugins-types"
|
name = "kigi-hooks-plugins-types"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
@@ -5779,7 +5779,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-http"
|
name = "kigi-http"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"kigi-auth",
|
"kigi-auth",
|
||||||
"kigi-log",
|
"kigi-log",
|
||||||
@@ -5794,7 +5794,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-hunk-tracker"
|
name = "kigi-hunk-tracker"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"chrono",
|
"chrono",
|
||||||
"dunce",
|
"dunce",
|
||||||
@@ -5815,14 +5815,14 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-interjection-core"
|
name = "kigi-interjection-core"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"serde",
|
"serde",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-log"
|
name = "kigi-log"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"chrono",
|
"chrono",
|
||||||
@@ -5840,7 +5840,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-markdown"
|
name = "kigi-markdown"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"anstyle",
|
"anstyle",
|
||||||
"anstyle-lossy",
|
"anstyle-lossy",
|
||||||
@@ -5864,14 +5864,14 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-markdown-core"
|
name = "kigi-markdown-core"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"pulldown-cmark",
|
"pulldown-cmark",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-mcp"
|
name = "kigi-mcp"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"agent-client-protocol",
|
"agent-client-protocol",
|
||||||
"async-trait",
|
"async-trait",
|
||||||
@@ -5908,7 +5908,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-memory"
|
name = "kigi-memory"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"arc-swap",
|
"arc-swap",
|
||||||
@@ -5942,7 +5942,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-mermaid"
|
name = "kigi-mermaid"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"fontdb",
|
"fontdb",
|
||||||
"image",
|
"image",
|
||||||
@@ -5960,16 +5960,17 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-models"
|
name = "kigi-models"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"kigi-env",
|
"kigi-env",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
|
"tracing",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-pager-minimal"
|
name = "kigi-pager-minimal"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"chrono",
|
"chrono",
|
||||||
"crossterm",
|
"crossterm",
|
||||||
@@ -5986,7 +5987,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-pager-pty-harness"
|
name = "kigi-pager-pty-harness"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"alacritty_terminal",
|
"alacritty_terminal",
|
||||||
"anyhow",
|
"anyhow",
|
||||||
@@ -6011,7 +6012,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-pager-render"
|
name = "kigi-pager-render"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"agent-client-protocol",
|
"agent-client-protocol",
|
||||||
"anstyle",
|
"anstyle",
|
||||||
@@ -6063,7 +6064,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-paths"
|
name = "kigi-paths"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"camino",
|
"camino",
|
||||||
"serde",
|
"serde",
|
||||||
@@ -6073,7 +6074,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-prompt-queue"
|
name = "kigi-prompt-queue"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
@@ -6081,7 +6082,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-proto-build"
|
name = "kigi-proto-build"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"pbjson-build",
|
"pbjson-build",
|
||||||
@@ -6092,7 +6093,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-ratatui-inline"
|
name = "kigi-ratatui-inline"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"ansi-width",
|
"ansi-width",
|
||||||
"anstyle-parse 0.2.7",
|
"anstyle-parse 0.2.7",
|
||||||
@@ -6109,7 +6110,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-ratatui-textarea"
|
name = "kigi-ratatui-textarea"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"arboard",
|
"arboard",
|
||||||
"chrono",
|
"chrono",
|
||||||
@@ -6130,7 +6131,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-sampler"
|
name = "kigi-sampler"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"async-openai",
|
"async-openai",
|
||||||
"async-stream",
|
"async-stream",
|
||||||
@@ -6153,10 +6154,11 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-sampling-types"
|
name = "kigi-sampling-types"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"assert_matches",
|
"assert_matches",
|
||||||
"async-openai",
|
"async-openai",
|
||||||
|
"base64",
|
||||||
"indexmap",
|
"indexmap",
|
||||||
"kigi-compaction",
|
"kigi-compaction",
|
||||||
"kigi-tools",
|
"kigi-tools",
|
||||||
@@ -6169,7 +6171,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-sandbox"
|
name = "kigi-sandbox"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"chrono",
|
"chrono",
|
||||||
@@ -6190,7 +6192,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-secrets"
|
name = "kigi-secrets"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"regex",
|
"regex",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
@@ -6228,7 +6230,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-shell"
|
name = "kigi-shell"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"agent-client-protocol",
|
"agent-client-protocol",
|
||||||
"anyhow",
|
"anyhow",
|
||||||
@@ -6365,7 +6367,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-shell-base"
|
name = "kigi-shell-base"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"chrono",
|
"chrono",
|
||||||
@@ -6390,7 +6392,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-sqlite-journal"
|
name = "kigi-sqlite-journal"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"libc",
|
"libc",
|
||||||
"rusqlite",
|
"rusqlite",
|
||||||
@@ -6401,7 +6403,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-subagent-resolution"
|
name = "kigi-subagent-resolution"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"kigi-sampling-types",
|
"kigi-sampling-types",
|
||||||
"kigi-tool-types",
|
"kigi-tool-types",
|
||||||
@@ -6416,7 +6418,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-system-power"
|
name = "kigi-system-power"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"windows-sys 0.59.0",
|
"windows-sys 0.59.0",
|
||||||
"zbus",
|
"zbus",
|
||||||
@@ -6424,7 +6426,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-test-support"
|
name = "kigi-test-support"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"agent-client-protocol",
|
"agent-client-protocol",
|
||||||
"anyhow",
|
"anyhow",
|
||||||
@@ -6446,7 +6448,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-test-utils"
|
name = "kigi-test-utils"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"runfiles",
|
"runfiles",
|
||||||
"tracing",
|
"tracing",
|
||||||
@@ -6455,11 +6457,11 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-token-estimation"
|
name = "kigi-token-estimation"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-tool-protocol"
|
name = "kigi-tool-protocol"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"kigi-tool-types",
|
"kigi-tool-types",
|
||||||
"serde",
|
"serde",
|
||||||
@@ -6470,7 +6472,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-tool-runtime"
|
name = "kigi-tool-runtime"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"async-trait",
|
"async-trait",
|
||||||
@@ -6488,7 +6490,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-tool-types"
|
name = "kigi-tool-types"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"minijinja",
|
"minijinja",
|
||||||
"schemars 1.2.1",
|
"schemars 1.2.1",
|
||||||
@@ -6498,7 +6500,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-tools"
|
name = "kigi-tools"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"arc-swap",
|
"arc-swap",
|
||||||
@@ -6575,7 +6577,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-tools-api"
|
name = "kigi-tools-api"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"kigi-proto-build",
|
"kigi-proto-build",
|
||||||
"kigi-tool-protocol",
|
"kigi-tool-protocol",
|
||||||
@@ -6588,11 +6590,11 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-tracing-macros"
|
name = "kigi-tracing-macros"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-tty-utils"
|
name = "kigi-tty-utils"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"libc",
|
"libc",
|
||||||
"nix 0.30.1",
|
"nix 0.30.1",
|
||||||
@@ -6602,7 +6604,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-tui"
|
name = "kigi-tui"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"agent-client-protocol",
|
"agent-client-protocol",
|
||||||
"ansi-to-tui",
|
"ansi-to-tui",
|
||||||
@@ -6689,7 +6691,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-update"
|
name = "kigi-update"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"dunce",
|
"dunce",
|
||||||
@@ -6718,14 +6720,14 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-version"
|
name = "kigi-version"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"semver",
|
"semver",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-workspace"
|
name = "kigi-workspace"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"agent-client-protocol",
|
"agent-client-protocol",
|
||||||
"anyhow",
|
"anyhow",
|
||||||
@@ -6804,7 +6806,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "kigi-workspace-types"
|
name = "kigi-workspace-types"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"base64",
|
"base64",
|
||||||
"chrono",
|
"chrono",
|
||||||
@@ -8838,7 +8840,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "ptyctl"
|
name = "ptyctl"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"alacritty_terminal",
|
"alacritty_terminal",
|
||||||
"anyhow",
|
"anyhow",
|
||||||
@@ -8856,7 +8858,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "ptyctl-cli"
|
name = "ptyctl-cli"
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"axum",
|
"axum",
|
||||||
|
|||||||
+1
-1
@@ -76,7 +76,7 @@ members = [
|
|||||||
]
|
]
|
||||||
|
|
||||||
[workspace.package]
|
[workspace.package]
|
||||||
version = "0.1.0"
|
version = "0.1.4"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
license = "Apache-2.0"
|
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
|
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
|
crates/codegen/kigi-tools/THIRD_PARTY_NOTICES.md for the license terms and
|
||||||
the per-file change notices.
|
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>
|
<h1>Kigi (<code>kigi</code>) 🌘</h1>
|
||||||
|
|
||||||
**Kigi** is an unofficial Kimi Code CLI community build — a terminal-based
|
<h3>🕸️ The world's first CLI with built-in <em>Graph Engineering</em></h3>
|
||||||
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).
|
|
||||||
|
|
||||||
It runs as a full-screen TUI that understands your codebase, edits files,
|
<p><code>/graph</code> turns one objective into a dependency graph of
|
||||||
executes shell commands, searches the web, and manages long-running tasks —
|
autonomous, self-verifying agent loops — planned, parallelized,
|
||||||
interactively, headlessly for scripting/CI, or embedded in editors via the
|
adversarially verified, and merged back, end to end.</p>
|
||||||
Agent Client Protocol (ACP).
|
|
||||||
|
**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) ·
|
[Installation](#installation) ·
|
||||||
|
[Graph engineering](#graph-engineering) ·
|
||||||
[Providers and API keys](#providers-and-api-keys) ·
|
[Providers and API keys](#providers-and-api-keys) ·
|
||||||
[Building from source](#building-from-source) ·
|
[Building from source](#building-from-source) ·
|
||||||
[Coexistence with the official CLI](#coexistence-with-the-official-kimi-cli) ·
|
[Coexistence with the official CLI](#coexistence-with-the-official-kimi-cli) ·
|
||||||
@@ -28,10 +34,6 @@ Agent Client Protocol (ACP).
|
|||||||
|
|
||||||
## Installation
|
## Installation
|
||||||
|
|
||||||
Prebuilt single-file binaries for macOS (arm64/x86_64), Linux (arm64/x86_64),
|
|
||||||
and Windows (x86_64) are published on
|
|
||||||
[GitHub Releases](https://github.com/ZacharyZhang-NY/Kigi-CLI/releases):
|
|
||||||
|
|
||||||
```sh
|
```sh
|
||||||
# macOS / Linux
|
# macOS / Linux
|
||||||
curl -fsSL https://raw.githubusercontent.com/ZacharyZhang-NY/Kigi-CLI/main/install.sh | bash
|
curl -fsSL https://raw.githubusercontent.com/ZacharyZhang-NY/Kigi-CLI/main/install.sh | bash
|
||||||
@@ -43,61 +45,115 @@ irm https://raw.githubusercontent.com/ZacharyZhang-NY/Kigi-CLI/main/install.ps1
|
|||||||
```
|
```
|
||||||
|
|
||||||
```sh
|
```sh
|
||||||
kigi --version # kigi 0.1.0 … unofficial Kimi Code CLI community build
|
kigi login # pick a provider, sign in
|
||||||
kigi login # sign in with your Kimi Code subscription (device-code flow)
|
kigi # go
|
||||||
kigi # start the TUI
|
|
||||||
```
|
```
|
||||||
|
|
||||||
The installer verifies every download against the release's `SHA256SUMS`,
|
Single file, no runtime. macOS and Linux on arm64/x86_64, Windows on x86_64,
|
||||||
installs into `~/.kigi/bin/kigi` (`%USERPROFILE%\.kigi\bin\kigi.exe` on
|
checksummed against the release's `SHA256SUMS`. `kigi update` handles upgrades.
|
||||||
Windows), and prints the PATH line to add. Later releases arrive through the
|
|
||||||
built-in self-updater (`kigi update`, gated by `KIGI_AUTO_UPDATE`), which
|
## Graph engineering
|
||||||
pulls from the same GitHub Releases feed.
|
|
||||||
|
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
|
## 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 |
|
**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
|
||||||
| `kimi-code` | `https://api.kimi.com/coding/v1` | Kimi Code subscription OAuth (`kigi login`) |
|
credentials are never sent to another.
|
||||||
| `moonshot-cn` | `https://api.moonshot.cn/v1` | Moonshot open-platform API key |
|
|
||||||
| `moonshot-ai` | `https://api.moonshot.ai/v1` | Moonshot open-platform API key |
|
|
||||||
|
|
||||||
Moonshot API keys come from the environment or `~/.kigi/config.toml`
|
| Platform id | Provider | Sign-in |
|
||||||
(environment wins; values are never logged):
|
| ---------------- | ------------------------- | ------------------------------------------- |
|
||||||
|
| `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
|
```sh
|
||||||
export KIGI_MOONSHOT_API_KEY=sk-... # applies to both open platforms
|
export OPENAI_API_KEY=sk-...
|
||||||
export KIGI_MOONSHOT_CN_API_KEY=sk-... # platform-scoped, beats the generic name
|
export XAI_API_KEY=xai-...
|
||||||
export KIGI_MOONSHOT_AI_API_KEY=sk-...
|
|
||||||
```
|
```
|
||||||
|
|
||||||
```toml
|
```toml
|
||||||
# ~/.kigi/config.toml
|
# ~/.kigi/config.toml
|
||||||
[platforms.moonshot-cn]
|
[platforms.openai]
|
||||||
api_key = "sk-..."
|
api_key = "sk-..."
|
||||||
|
|
||||||
[platforms.moonshot-ai]
|
[platforms.xai]
|
||||||
api_key = "sk-..."
|
api_key = "xai-..."
|
||||||
```
|
```
|
||||||
|
|
||||||
On login and on startup Kigi syncs each configured platform's model list
|
Model lists sync on startup. Pick one with `/model`, set its thinking level
|
||||||
from `GET {base}/models` and shows the merged catalog in the model picker
|
with `/effort`.
|
||||||
(catalog keys are `{platform_id}/{model_id}`). 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.
|
|
||||||
|
|
||||||
`KIGI_CODE_BASE_URL` re-points the subscription platform (useful for
|
Web `search`/`fetch` need a Kimi Code subscription; API-key sessions run
|
||||||
testing); `KIGI_MOONSHOT_CN_BASE_URL` / `KIGI_MOONSHOT_AI_BASE_URL` are the
|
without them, same as the official client.
|
||||||
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.
|
|
||||||
|
|
||||||
## Building from source
|
## Building from source
|
||||||
|
|
||||||
@@ -113,19 +169,15 @@ launcher at `bin/protoc`; install dotslash (`brew install dotslash` or
|
|||||||
|
|
||||||
## Coexistence with the official Kimi CLI
|
## Coexistence with the official Kimi CLI
|
||||||
|
|
||||||
Kigi is not affiliated with Moonshot AI or xAI, and it coexists with the
|
Kigi started as an unofficial Kimi Code CLI — a community fork of
|
||||||
official `kimi` CLI on the same machine: independent binary name,
|
[xai-org/grok-build](https://github.com/xai-org/grok-build), not affiliated
|
||||||
independent config directory (`~/.kigi`), independent keyring credentials
|
with Moonshot AI or xAI. It keeps its own binary, its own `~/.kigi`, its own
|
||||||
(service `kigi`), and a `KIGI_*` environment-variable namespace. Nothing
|
keyring entry, and its own `KIGI_*` env vars, and never touches what the
|
||||||
the official client installs or stores is ever read at runtime or written.
|
official `kimi` CLI installed. `kigi import-kimi` copies your old config over
|
||||||
On first launch Kigi offers a **one-time, strictly read-only** import of
|
once, read-only.
|
||||||
your existing `~/.kimi` configuration (MCP servers, custom providers,
|
|
||||||
default model) via `kigi import-kimi` — file contents and mtimes under
|
|
||||||
`~/.kimi` are left untouched, verified by tests.
|
|
||||||
|
|
||||||
Kigi is **zero-telemetry**: the only outbound connections are the
|
**Zero telemetry.** It talks to the APIs you configured, GitHub Releases, and
|
||||||
inference/auth APIs you configure, GitHub Releases for updates, and MCP
|
your own MCP servers. Nothing else.
|
||||||
servers you add.
|
|
||||||
|
|
||||||
## License
|
## License
|
||||||
|
|
||||||
|
|||||||
@@ -272,6 +272,7 @@ mod tests {
|
|||||||
temperature: None,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: std::num::NonZeroU64::new(128_000).unwrap(),
|
context_window: std::num::NonZeroU64::new(128_000).unwrap(),
|
||||||
reasoning_effort: None,
|
reasoning_effort: None,
|
||||||
|
|||||||
@@ -23,6 +23,7 @@ fn test_config_with_window(context_window: u64) -> SamplingConfig {
|
|||||||
temperature: None,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: NonZeroU64::new(context_window)
|
context_window: NonZeroU64::new(context_window)
|
||||||
.expect("test context_window must be non-zero"),
|
.expect("test context_window must be non-zero"),
|
||||||
@@ -913,6 +914,7 @@ async fn update_sampling_config_is_queryable() {
|
|||||||
temperature: Some(0.5),
|
temperature: Some(0.5),
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: NonZeroU64::new(200_000).unwrap(),
|
context_window: NonZeroU64::new(200_000).unwrap(),
|
||||||
reasoning_effort: None,
|
reasoning_effort: None,
|
||||||
@@ -1298,6 +1300,7 @@ async fn build_request_uses_sampling_config() {
|
|||||||
temperature: Some(0.7),
|
temperature: Some(0.7),
|
||||||
top_p: Some(0.9),
|
top_p: Some(0.9),
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: NonZeroU64::new(128_000).unwrap(),
|
context_window: NonZeroU64::new(128_000).unwrap(),
|
||||||
reasoning_effort: None,
|
reasoning_effort: None,
|
||||||
@@ -3401,6 +3404,7 @@ async fn sampling_config_survives_compaction_replacement() {
|
|||||||
temperature: Some(0.7),
|
temperature: Some(0.7),
|
||||||
top_p: Some(0.95),
|
top_p: Some(0.95),
|
||||||
api_backend: ApiBackend::Responses,
|
api_backend: ApiBackend::Responses,
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: NonZeroU64::new(500_000).unwrap(),
|
context_window: NonZeroU64::new(500_000).unwrap(),
|
||||||
reasoning_effort: None,
|
reasoning_effort: None,
|
||||||
@@ -3481,6 +3485,7 @@ async fn model_metadata_lost_after_compaction_then_recovered_on_next_turn() {
|
|||||||
temperature: Some(0.7),
|
temperature: Some(0.7),
|
||||||
top_p: Some(0.95),
|
top_p: Some(0.95),
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: NonZeroU64::new(500_000).unwrap(),
|
context_window: NonZeroU64::new(500_000).unwrap(),
|
||||||
reasoning_effort: None,
|
reasoning_effort: None,
|
||||||
@@ -3569,6 +3574,7 @@ async fn context_window_downgrade_triggers_auto_compact() {
|
|||||||
temperature: Some(0.7),
|
temperature: Some(0.7),
|
||||||
top_p: Some(0.95),
|
top_p: Some(0.95),
|
||||||
api_backend: ApiBackend::Responses,
|
api_backend: ApiBackend::Responses,
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: NonZeroU64::new(500_000).unwrap(),
|
context_window: NonZeroU64::new(500_000).unwrap(),
|
||||||
reasoning_effort: None,
|
reasoning_effort: None,
|
||||||
|
|||||||
@@ -370,6 +370,7 @@ mod tests {
|
|||||||
temperature: None,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: std::num::NonZeroU64::new(128_000).unwrap(),
|
context_window: std::num::NonZeroU64::new(128_000).unwrap(),
|
||||||
reasoning_effort: None,
|
reasoning_effort: None,
|
||||||
|
|||||||
@@ -178,6 +178,7 @@ mod tests {
|
|||||||
temperature: None,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: NonZeroU64::new(128_000).unwrap(),
|
context_window: NonZeroU64::new(128_000).unwrap(),
|
||||||
reasoning_effort: None,
|
reasoning_effort: None,
|
||||||
@@ -221,6 +222,7 @@ mod tests {
|
|||||||
temperature: Some(0.7),
|
temperature: Some(0.7),
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: NonZeroU64::new(128_000).unwrap(),
|
context_window: NonZeroU64::new(128_000).unwrap(),
|
||||||
reasoning_effort: None,
|
reasoning_effort: None,
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ description = "Kimi platform registry, /models wire contract, capability derivat
|
|||||||
kigi-env = { workspace = true }
|
kigi-env = { workspace = true }
|
||||||
serde = { workspace = true, features = ["derive"] }
|
serde = { workspace = true, features = ["derive"] }
|
||||||
serde_json = { workspace = true }
|
serde_json = { workspace = true }
|
||||||
|
tracing = { workspace = true }
|
||||||
|
|
||||||
[lints]
|
[lints]
|
||||||
workspace = true
|
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.
|
// pins them to the mock so no PTY test can reach a live endpoint.
|
||||||
("KIGI_MOONSHOT_CN_BASE_URL".into(), self.url()),
|
("KIGI_MOONSHOT_CN_BASE_URL".into(), self.url()),
|
||||||
("KIGI_MOONSHOT_AI_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_TELEMETRY_ENABLED".into(), "false".into()),
|
||||||
("KIGI_FEEDBACK_ENABLED".into(), "false".into()),
|
("KIGI_FEEDBACK_ENABLED".into(), "false".into()),
|
||||||
("KIGI_TRACE_UPLOAD".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_API_BASE_URL"), Some(content.url()));
|
||||||
assert_eq!(get("KIGI_MOONSHOT_CN_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("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_TELEMETRY_ENABLED").as_deref(), Some("false"));
|
||||||
assert_eq!(get("KIGI_FEEDBACK_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"));
|
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");
|
.expect("seed fake oauth auth.json");
|
||||||
}
|
}
|
||||||
|
|
||||||
/// [`ContentController::env_for_pager`] minus `XAI_API_KEY`, so the entry
|
/// [`ContentController::env_for_pager`] minus the house BYOK key
|
||||||
/// written by [`seed_fake_oauth`] is the active credential.
|
/// (`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)> {
|
pub fn oauth_env_for_pager(content: &ContentController) -> Vec<(String, String)> {
|
||||||
let mut env = content.env_for_pager();
|
let mut env = content.env_for_pager();
|
||||||
env.retain(|(k, _)| k != "XAI_API_KEY");
|
env.retain(|(k, _)| k != "KIGI_API_KEY");
|
||||||
env
|
env
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -218,10 +218,12 @@ pub fn diamond_hollow_char() -> char {
|
|||||||
/// 1-column ASCII spinner (`|`, `/`, `-`, `\`) on legacy ConHost.
|
/// 1-column ASCII spinner (`|`, `/`, `-`, `\`) on legacy ConHost.
|
||||||
///
|
///
|
||||||
/// The U+2800 Braille Patterns block is not part of CP437 and renders as
|
/// The U+2800 Braille Patterns block is not part of CP437 and renders as
|
||||||
/// tofu on the legacy console raster font, so the turn-status line, the
|
/// tofu on the legacy console raster font, so the starting-session row,
|
||||||
/// MCP-connecting chip, the image-viewer loader, and the `/btw` overlay
|
/// the MCP-connecting chip, the image-viewer loader, and the `/btw`
|
||||||
/// all fall back to the classic ASCII spinner there. Every frame in both
|
/// overlay all fall back to the classic ASCII spinner there. Every frame
|
||||||
/// sets is exactly 1 column so the surrounding layout never shifts.
|
/// in both sets is exactly 1 column so the surrounding layout never
|
||||||
|
/// shifts. (The turn-status "Waiting for response…" line uses
|
||||||
|
/// [`moon_spinner_frames`] instead.)
|
||||||
pub fn braille_spinner_frames() -> &'static [&'static str] {
|
pub fn braille_spinner_frames() -> &'static [&'static str] {
|
||||||
const FANCY: &[&str] = &[
|
const FANCY: &[&str] = &[
|
||||||
"\u{280b}", "\u{2819}", "\u{2839}", "\u{2838}", "\u{283c}", "\u{2834}", "\u{2826}",
|
"\u{280b}", "\u{2819}", "\u{2839}", "\u{2838}", "\u{283c}", "\u{2834}", "\u{2826}",
|
||||||
@@ -235,6 +237,34 @@ pub fn braille_spinner_frames() -> &'static [&'static str] {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Moon-phase spinner frames (`🌑🌒🌓🌔🌕🌖🌗🌘`) normally; the ASCII
|
||||||
|
/// spinner on legacy ConHost.
|
||||||
|
///
|
||||||
|
/// The turn-status "Waiting for response…" line reuses the official
|
||||||
|
/// kimi-cli's lunation animation (rich's `moon` spinner) — one full
|
||||||
|
/// lunation per cycle, matching the welcome-screen moon logo. Emoji
|
||||||
|
/// frames are 2 columns wide; the caller measures the rendered frame
|
||||||
|
/// (`spinner_str.width()`) so the layout adapts. Legacy ConHost's raster
|
||||||
|
/// font has no emoji, so it falls back to the 1-column ASCII spinner.
|
||||||
|
pub fn moon_spinner_frames() -> &'static [&'static str] {
|
||||||
|
const FANCY: &[&str] = &[
|
||||||
|
"\u{1f311}",
|
||||||
|
"\u{1f312}",
|
||||||
|
"\u{1f313}",
|
||||||
|
"\u{1f314}",
|
||||||
|
"\u{1f315}",
|
||||||
|
"\u{1f316}",
|
||||||
|
"\u{1f317}",
|
||||||
|
"\u{1f318}",
|
||||||
|
];
|
||||||
|
const FALLBACK: &[&str] = &["|", "/", "-", "\\"];
|
||||||
|
if is_legacy_windows_console() {
|
||||||
|
FALLBACK
|
||||||
|
} else {
|
||||||
|
FANCY
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Pulsing dot progress-spinner frames (`⋅ : ⸬ ⁙`) normally; a quiet
|
/// Pulsing dot progress-spinner frames (`⋅ : ⸬ ⁙`) normally; a quiet
|
||||||
/// 1-column dot cycle (`.`, `:`, `·`) on legacy ConHost.
|
/// 1-column dot cycle (`.`, `:`, `·`) on legacy ConHost.
|
||||||
///
|
///
|
||||||
|
|||||||
@@ -89,6 +89,10 @@ mod tests {
|
|||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: ApiBackend::ChatCompletions,
|
api_backend: ApiBackend::ChatCompletions,
|
||||||
auth_scheme: Default::default(),
|
auth_scheme: Default::default(),
|
||||||
|
anthropic_oauth: false,
|
||||||
|
github_copilot: false,
|
||||||
|
openai_codex: false,
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: IndexMap::new(),
|
extra_headers: IndexMap::new(),
|
||||||
context_window: 8192,
|
context_window: 8192,
|
||||||
force_http1: false,
|
force_http1: false,
|
||||||
|
|||||||
@@ -41,6 +41,35 @@ pub use kigi_sampling_types::ApiBackend;
|
|||||||
const AGENT_PRODUCT: &str = "kigi";
|
const AGENT_PRODUCT: &str = "kigi";
|
||||||
const ANTHROPIC_DEFAULT_MAX_TOKENS: u32 = 128_000;
|
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.
|
/// Parse the `Retry-After` response header as delta-seconds.
|
||||||
/// Our inference backends only emit integer seconds (never HTTP-date),
|
/// Our inference backends only emit integer seconds (never HTTP-date),
|
||||||
/// so we only handle that form. HTTP-dates silently return `None` and
|
/// 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) => {
|
Err(first_err) => {
|
||||||
// Try sanitizing: parse as Value, strip unknown tools, retry.
|
// Try sanitizing: parse as Value, strip unknown tools, retry.
|
||||||
if let Ok(mut value) = serde_json::from_str::<serde_json::Value>(data) {
|
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
|
// Strip tools that async_openai's rs::Tool can't deserialize
|
||||||
// (e.g., xAI-specific "x_search"). Instead of maintaining a
|
// (e.g., xAI-specific "x_search"). Instead of maintaining a
|
||||||
// hardcoded allowlist, try deserializing each tool entry —
|
// hardcoded allowlist, try deserializing each tool entry —
|
||||||
@@ -267,8 +299,14 @@ struct ClientDefaults {
|
|||||||
top_p: Option<f32>,
|
top_p: Option<f32>,
|
||||||
api_backend: ApiBackend,
|
api_backend: ApiBackend,
|
||||||
auth_scheme: AuthScheme,
|
auth_scheme: AuthScheme,
|
||||||
|
chat_compat: kigi_sampling_types::ChatCompat,
|
||||||
stream_tool_calls: bool,
|
stream_tool_calls: bool,
|
||||||
doom_loop_recovery: Option<kigi_sampling_types::DoomLoopRecoveryPolicy>,
|
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);
|
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 => {
|
AuthScheme::Bearer => {
|
||||||
let bearer = format!("Bearer {}", api_key);
|
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
|
// Apply all extra headers verbatim. This is the single
|
||||||
// injection point for proxy-auth headers and any other URL- or
|
// injection point for proxy-auth headers and any other URL- or
|
||||||
// environment-specific headers the session decides to set.
|
// 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
|
// (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).
|
// 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),
|
Some(origin) => user_agent_string_for(origin),
|
||||||
None => user_agent_string_for(&OriginClientInfo {
|
None => user_agent_string_for(&OriginClientInfo {
|
||||||
product: AGENT_PRODUCT.to_string(),
|
product: AGENT_PRODUCT.to_string(),
|
||||||
version: Some(agent_version()),
|
version: Some(agent_version()),
|
||||||
}),
|
}),
|
||||||
|
}
|
||||||
};
|
};
|
||||||
if let Ok(v) = HeaderValue::from_str(&ua_string) {
|
if let Ok(v) = HeaderValue::from_str(&ua_string) {
|
||||||
headers.insert(USER_AGENT, v);
|
headers.insert(USER_AGENT, v);
|
||||||
@@ -445,8 +570,11 @@ impl SamplingClient {
|
|||||||
top_p: config.top_p,
|
top_p: config.top_p,
|
||||||
api_backend: config.api_backend,
|
api_backend: config.api_backend,
|
||||||
auth_scheme: config.auth_scheme,
|
auth_scheme: config.auth_scheme,
|
||||||
|
chat_compat: config.chat_compat,
|
||||||
stream_tool_calls: config.stream_tool_calls,
|
stream_tool_calls: config.stream_tool_calls,
|
||||||
doom_loop_recovery: config.doom_loop_recovery,
|
doom_loop_recovery: config.doom_loop_recovery,
|
||||||
|
anthropic_oauth: config.anthropic_oauth,
|
||||||
|
openai_codex: config.openai_codex,
|
||||||
};
|
};
|
||||||
|
|
||||||
Ok(Self {
|
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
|
let auth_prefix = headers
|
||||||
.get(AUTHORIZATION)
|
.get(AUTHORIZATION)
|
||||||
@@ -595,6 +740,9 @@ impl SamplingClient {
|
|||||||
|| lower.contains("apikey")
|
|| lower.contains("apikey")
|
||||||
|| lower.contains("token")
|
|| lower.contains("token")
|
||||||
|| lower.contains("secret")
|
|| 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.
|
/// 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);
|
tracing::error!("Failed to serialize chat/completions request: {}", e);
|
||||||
SamplingError::Serialization(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
|
let http_request = self
|
||||||
.post(self.endpoint("chat/completions"))
|
.post(self.endpoint("chat/completions"))
|
||||||
@@ -818,7 +969,10 @@ impl SamplingClient {
|
|||||||
tracing::error!("Failed to serialize chat/completions request: {}", e);
|
tracing::error!("Failed to serialize chat/completions request: {}", e);
|
||||||
SamplingError::Serialization(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
|
let http_request = self
|
||||||
.post(self.endpoint("chat/completions"))
|
.post(self.endpoint("chat/completions"))
|
||||||
@@ -1028,6 +1182,7 @@ impl SamplingClient {
|
|||||||
// it in post-serialize. This is the last surviving piece of the
|
// it in post-serialize. This is the last surviving piece of the
|
||||||
// old raw_output machinery.
|
// old raw_output machinery.
|
||||||
kigi_sampling_types::patch_reasoning_text_types(&mut request_body);
|
kigi_sampling_types::patch_reasoning_text_types(&mut request_body);
|
||||||
|
kigi_sampling_types::patch_reasoning_effort(&mut request_body, request.reasoning_effort);
|
||||||
let http_request = self.post(self.endpoint("responses")).json(&request_body);
|
let http_request = self.post(self.endpoint("responses")).json(&request_body);
|
||||||
|
|
||||||
let response = http_request.send().await.map_err(|e| {
|
let response = http_request.send().await.map_err(|e| {
|
||||||
@@ -1074,7 +1229,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);
|
let raw_body = String::from_utf8_lossy(&bytes);
|
||||||
tracing::error!(
|
tracing::error!(
|
||||||
error = %e,
|
error = %e,
|
||||||
@@ -1156,6 +1320,7 @@ impl SamplingClient {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
kigi_sampling_types::patch_reasoning_text_types(&mut request_body);
|
kigi_sampling_types::patch_reasoning_text_types(&mut request_body);
|
||||||
|
kigi_sampling_types::patch_reasoning_effort(&mut request_body, request.reasoning_effort);
|
||||||
// Fresh per attempt so signals never leak across retries; `None`
|
// Fresh per attempt so signals never leak across retries; `None`
|
||||||
// (check disabled) sends no header and does no peek work per event.
|
// (check disabled) sends no header and does no peek work per event.
|
||||||
let doom_loop = self
|
let doom_loop = self
|
||||||
@@ -1314,6 +1479,15 @@ impl SamplingClient {
|
|||||||
|
|
||||||
/// Apply default configuration to a Messages API request.
|
/// Apply default configuration to a Messages API request.
|
||||||
fn apply_message_defaults(&self, request: &mut MessagesRequestWrapper) -> Result<()> {
|
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
|
// Apply model default if not specified
|
||||||
if request.inner.model.is_empty() {
|
if request.inner.model.is_empty() {
|
||||||
request.inner.model = self.defaults.model.clone();
|
request.inner.model = self.defaults.model.clone();
|
||||||
@@ -1687,6 +1861,7 @@ impl SamplingClient {
|
|||||||
let responses_request: rs::CreateResponse = (&request).into();
|
let responses_request: rs::CreateResponse = (&request).into();
|
||||||
|
|
||||||
let mut wrapper = CreateResponseWrapper::new(responses_request);
|
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_conv_id = x_kigi_conv_id;
|
||||||
wrapper.x_kigi_req_id = x_kigi_req_id;
|
wrapper.x_kigi_req_id = x_kigi_req_id;
|
||||||
wrapper.x_kigi_session_id = x_kigi_session_id;
|
wrapper.x_kigi_session_id = x_kigi_session_id;
|
||||||
@@ -1720,6 +1895,7 @@ impl SamplingClient {
|
|||||||
let responses_request: rs::CreateResponse = (&request).into();
|
let responses_request: rs::CreateResponse = (&request).into();
|
||||||
|
|
||||||
let mut wrapper = CreateResponseWrapper::new(responses_request);
|
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_conv_id = x_kigi_conv_id;
|
||||||
wrapper.x_kigi_req_id = x_kigi_req_id;
|
wrapper.x_kigi_req_id = x_kigi_req_id;
|
||||||
wrapper.x_kigi_session_id = x_kigi_session_id;
|
wrapper.x_kigi_session_id = x_kigi_session_id;
|
||||||
@@ -1856,6 +2032,10 @@ mod tests {
|
|||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: ApiBackend::ChatCompletions,
|
api_backend: ApiBackend::ChatCompletions,
|
||||||
auth_scheme: AuthScheme::Bearer,
|
auth_scheme: AuthScheme::Bearer,
|
||||||
|
anthropic_oauth: false,
|
||||||
|
github_copilot: false,
|
||||||
|
openai_codex: false,
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: IndexMap::new(),
|
extra_headers: IndexMap::new(),
|
||||||
context_window: 8192,
|
context_window: 8192,
|
||||||
force_http1: false,
|
force_http1: false,
|
||||||
@@ -1874,6 +2054,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
|
/// Verify the serialized shape of StreamingChatRequest matches the
|
||||||
/// expected wire format: all ChatCompletionRequest fields flattened at
|
/// expected wire format: all ChatCompletionRequest fields flattened at
|
||||||
/// top level, plus `stream: true` and `stream_options.include_usage: true`.
|
/// top level, plus `stream: true` and `stream_options.include_usage: true`.
|
||||||
@@ -2017,6 +2528,7 @@ mod tests {
|
|||||||
api_key: Some("bearer-key-abc123".to_string()),
|
api_key: Some("bearer-key-abc123".to_string()),
|
||||||
api_backend: ApiBackend::Messages,
|
api_backend: ApiBackend::Messages,
|
||||||
auth_scheme: AuthScheme::Bearer,
|
auth_scheme: AuthScheme::Bearer,
|
||||||
|
chat_compat: Default::default(),
|
||||||
..minimal_config()
|
..minimal_config()
|
||||||
};
|
};
|
||||||
let client = SamplingClient::new(cfg).expect("client should build");
|
let client = SamplingClient::new(cfg).expect("client should build");
|
||||||
@@ -2183,6 +2695,7 @@ mod tests {
|
|||||||
api_key: Some("stale-bearer".to_string()),
|
api_key: Some("stale-bearer".to_string()),
|
||||||
api_backend: ApiBackend::Messages,
|
api_backend: ApiBackend::Messages,
|
||||||
auth_scheme: AuthScheme::Bearer,
|
auth_scheme: AuthScheme::Bearer,
|
||||||
|
chat_compat: Default::default(),
|
||||||
bearer_resolver: Some(std::sync::Arc::new(StaticBearerResolver("fresh-bearer"))),
|
bearer_resolver: Some(std::sync::Arc::new(StaticBearerResolver("fresh-bearer"))),
|
||||||
..minimal_config()
|
..minimal_config()
|
||||||
};
|
};
|
||||||
@@ -2211,6 +2724,7 @@ mod tests {
|
|||||||
api_key: Some("stale-bearer".to_string()),
|
api_key: Some("stale-bearer".to_string()),
|
||||||
api_backend: ApiBackend::Responses,
|
api_backend: ApiBackend::Responses,
|
||||||
auth_scheme: AuthScheme::Bearer,
|
auth_scheme: AuthScheme::Bearer,
|
||||||
|
chat_compat: Default::default(),
|
||||||
bearer_resolver: Some(std::sync::Arc::new(StaticBearerResolver("fresh-bearer"))),
|
bearer_resolver: Some(std::sync::Arc::new(StaticBearerResolver("fresh-bearer"))),
|
||||||
..minimal_config()
|
..minimal_config()
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -33,7 +33,7 @@ pub enum AuthScheme {
|
|||||||
///
|
///
|
||||||
/// `SamplerConfig` is the single source of truth for sampler
|
/// `SamplerConfig` is the single source of truth for sampler
|
||||||
/// configuration. The shell builds it directly (see
|
/// 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
|
/// `session::acp_session::SessionActor::reconstruct_full_config`) by
|
||||||
/// composing chat-state's `kigi_sampling_types::SamplingConfig`
|
/// composing chat-state's `kigi_sampling_types::SamplingConfig`
|
||||||
/// with `Credentials` (api key, client version).
|
/// with `Credentials` (api key, client version).
|
||||||
@@ -55,6 +55,29 @@ pub struct SamplerConfig {
|
|||||||
pub api_backend: ApiBackend,
|
pub api_backend: ApiBackend,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub auth_scheme: AuthScheme,
|
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
|
/// Extra request headers applied verbatim. The sampler never inspects
|
||||||
/// the URL to derive headers; callers (the session) inject proxy auth
|
/// the URL to derive headers; callers (the session) inject proxy auth
|
||||||
/// and other access headers here before constructing the config.
|
/// and other access headers here before constructing the config.
|
||||||
@@ -70,6 +93,11 @@ pub struct SamplerConfig {
|
|||||||
|
|
||||||
// Reasoning effort
|
// Reasoning effort
|
||||||
pub reasoning_effort: Option<ReasoningEffort>,
|
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
|
/// Client identity for the User-Agent header (`kigi/{version}` plus an
|
||||||
/// optional origin product). The old xAI proxy's identity headers
|
/// optional origin product). The old xAI proxy's identity headers
|
||||||
@@ -135,9 +163,13 @@ impl Default for SamplerConfig {
|
|||||||
model: String::new(),
|
model: String::new(),
|
||||||
max_completion_tokens: None,
|
max_completion_tokens: None,
|
||||||
temperature: None,
|
temperature: None,
|
||||||
|
chat_compat: kigi_sampling_types::ChatCompat::default(),
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: ApiBackend::default(),
|
api_backend: ApiBackend::default(),
|
||||||
auth_scheme: AuthScheme::default(),
|
auth_scheme: AuthScheme::default(),
|
||||||
|
anthropic_oauth: false,
|
||||||
|
github_copilot: false,
|
||||||
|
openai_codex: false,
|
||||||
extra_headers: IndexMap::new(),
|
extra_headers: IndexMap::new(),
|
||||||
context_window: 0,
|
context_window: 0,
|
||||||
force_http1: false,
|
force_http1: false,
|
||||||
|
|||||||
@@ -29,6 +29,89 @@ pub(crate) fn adapt_chat_completions_body(body: &mut Value) {
|
|||||||
adapt_tool_schemas(body);
|
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`
|
/// Map the OpenAI-style `reasoning_effort` knob onto Kimi's `thinking`
|
||||||
/// request field and drop `reasoning_effort` from the wire.
|
/// request field and drop `reasoning_effort` from the wire.
|
||||||
///
|
///
|
||||||
@@ -280,6 +363,123 @@ mod tests {
|
|||||||
use super::*;
|
use super::*;
|
||||||
use serde_json::json;
|
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]
|
#[test]
|
||||||
fn reasoning_effort_maps_to_kimi_thinking_field() {
|
fn reasoning_effort_maps_to_kimi_thinking_field() {
|
||||||
// Level rides along as thinking.effort (live wire: 200 with
|
// Level rides along as thinking.effort (live wire: 200 with
|
||||||
@@ -292,8 +492,9 @@ mod tests {
|
|||||||
json!({ "type": "enabled", "effort": "high" })
|
json!({ "type": "enabled", "effort": "high" })
|
||||||
);
|
);
|
||||||
|
|
||||||
// Canonical `xhigh` is spelled `max` on the Kimi wire (the K3
|
// Legacy canonical `xhigh` (pre-Max configs/sessions) is spelled
|
||||||
// valid_efforts vocabulary is low/high/max).
|
// `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" });
|
let mut body = json!({ "model": "k3", "reasoning_effort": "xhigh" });
|
||||||
adapt_chat_completions_body(&mut body);
|
adapt_chat_completions_body(&mut body);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
@@ -301,6 +502,15 @@ mod tests {
|
|||||||
json!({ "type": "enabled", "effort": "max" })
|
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
|
// kimi.py:218: "off" (our ReasoningEffort::None) → disabled, and no
|
||||||
// effort key (a disabled+effort combination would be contradictory).
|
// effort key (a disabled+effort combination would be contradictory).
|
||||||
let mut body = json!({ "reasoning_effort": "none" });
|
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]
|
#[tokio::test]
|
||||||
async fn tool_call_stream_emits_deltas_and_assembles_final_call() {
|
async fn tool_call_stream_emits_deltas_and_assembles_final_call() {
|
||||||
// First chunk has id + name + part of arguments.
|
// First chunk has id + name + part of arguments.
|
||||||
|
|||||||
@@ -78,6 +78,10 @@ fn test_config(base_url: String, model: &str) -> SamplerConfig {
|
|||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: ApiBackend::ChatCompletions,
|
api_backend: ApiBackend::ChatCompletions,
|
||||||
auth_scheme: Default::default(),
|
auth_scheme: Default::default(),
|
||||||
|
anthropic_oauth: false,
|
||||||
|
github_copilot: false,
|
||||||
|
openai_codex: false,
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: IndexMap::new(),
|
extra_headers: IndexMap::new(),
|
||||||
context_window: 128_000,
|
context_window: 128_000,
|
||||||
force_http1: false,
|
force_http1: false,
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ description = "Pure data types for the xAI sampling / chat-completion API layer"
|
|||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
async-openai = { workspace = true }
|
async-openai = { workspace = true }
|
||||||
|
base64 = { workspace = true }
|
||||||
indexmap = { workspace = true, features = ["serde"] }
|
indexmap = { workspace = true, features = ["serde"] }
|
||||||
reqwest = { workspace = true }
|
reqwest = { workspace = true }
|
||||||
serde = { workspace = true, features = ["derive"] }
|
serde = { workspace = true, features = ["derive"] }
|
||||||
|
|||||||
@@ -250,8 +250,13 @@ pub struct AssistantItem {
|
|||||||
/// `response.reasoning.effort` (Responses API). Stored beside
|
/// `response.reasoning.effort` (Responses API). Stored beside
|
||||||
/// `model_id`/`model_fingerprint` so per-response effort survives
|
/// `model_id`/`model_fingerprint` so per-response effort survives
|
||||||
/// mid-session model/effort switches. `None` for synthetic items and
|
/// mid-session model/effort switches. `None` for synthetic items and
|
||||||
/// backends that don't echo it.
|
/// backends that don't echo it. Lenient on read: an unknown token from a
|
||||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
/// 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>,
|
pub reasoning_effort: Option<crate::ReasoningEffort>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -2153,7 +2158,10 @@ impl From<&ConversationRequest> for rs::CreateResponse {
|
|||||||
prompt_cache_key: None,
|
prompt_cache_key: None,
|
||||||
prompt_cache_retention: None,
|
prompt_cache_retention: None,
|
||||||
reasoning: Some(rs::Reasoning {
|
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),
|
summary: Some(rs::ReasoningSummary::Concise),
|
||||||
}),
|
}),
|
||||||
safety_identifier: None,
|
safety_identifier: None,
|
||||||
@@ -5076,13 +5084,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]
|
#[test]
|
||||||
fn test_messages_request_wire_format_for_supported_variants() {
|
fn test_messages_request_wire_format_for_supported_variants() {
|
||||||
for (variant, expected) in [
|
for (variant, expected) in [
|
||||||
(crate::ReasoningEffort::Low, "low"),
|
(crate::ReasoningEffort::Low, "low"),
|
||||||
(crate::ReasoningEffort::Medium, "medium"),
|
(crate::ReasoningEffort::Medium, "medium"),
|
||||||
(crate::ReasoningEffort::High, "high"),
|
(crate::ReasoningEffort::High, "high"),
|
||||||
(crate::ReasoningEffort::Xhigh, "max"),
|
(crate::ReasoningEffort::Xhigh, "xhigh"),
|
||||||
|
(crate::ReasoningEffort::Max, "max"),
|
||||||
] {
|
] {
|
||||||
let req = messages_test_request(Some(variant));
|
let req = messages_test_request(Some(variant));
|
||||||
let msgs = build_messages_request(&req);
|
let msgs = build_messages_request(&req);
|
||||||
@@ -5131,6 +5161,7 @@ mod tests {
|
|||||||
(crate::ReasoningEffort::Medium, "medium"),
|
(crate::ReasoningEffort::Medium, "medium"),
|
||||||
(crate::ReasoningEffort::High, "high"),
|
(crate::ReasoningEffort::High, "high"),
|
||||||
(crate::ReasoningEffort::Xhigh, "xhigh"),
|
(crate::ReasoningEffort::Xhigh, "xhigh"),
|
||||||
|
(crate::ReasoningEffort::Max, "max"),
|
||||||
] {
|
] {
|
||||||
let req = ConversationRequest::from_items(vec![ConversationItem::user("hi")])
|
let req = ConversationRequest::from_items(vec![ConversationItem::user("hi")])
|
||||||
.with_model("test");
|
.with_model("test");
|
||||||
@@ -5163,6 +5194,10 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_responses_request_carries_reasoning_effort_nested() {
|
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 [
|
for (variant, expected) in [
|
||||||
(crate::ReasoningEffort::None, "none"),
|
(crate::ReasoningEffort::None, "none"),
|
||||||
(crate::ReasoningEffort::Minimal, "minimal"),
|
(crate::ReasoningEffort::Minimal, "minimal"),
|
||||||
@@ -5170,6 +5205,7 @@ mod tests {
|
|||||||
(crate::ReasoningEffort::Medium, "medium"),
|
(crate::ReasoningEffort::Medium, "medium"),
|
||||||
(crate::ReasoningEffort::High, "high"),
|
(crate::ReasoningEffort::High, "high"),
|
||||||
(crate::ReasoningEffort::Xhigh, "xhigh"),
|
(crate::ReasoningEffort::Xhigh, "xhigh"),
|
||||||
|
(crate::ReasoningEffort::Max, "max"),
|
||||||
] {
|
] {
|
||||||
let req = ConversationRequest {
|
let req = ConversationRequest {
|
||||||
reasoning_effort: Some(variant),
|
reasoning_effort: Some(variant),
|
||||||
@@ -5177,15 +5213,56 @@ mod tests {
|
|||||||
.with_model("test")
|
.with_model("test")
|
||||||
};
|
};
|
||||||
let resp: crate::rs::CreateResponse = (&req).into();
|
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!(
|
assert_eq!(
|
||||||
json.pointer("/reasoning/effort").and_then(|v| v.as_str()),
|
json.pointer("/reasoning/effort").and_then(|v| v.as_str()),
|
||||||
Some(expected),
|
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]
|
#[test]
|
||||||
fn test_responses_request_omits_effort_when_unset() {
|
fn test_responses_request_omits_effort_when_unset() {
|
||||||
let req =
|
let req =
|
||||||
|
|||||||
@@ -487,7 +487,69 @@ pub enum FinishReason {
|
|||||||
FunctionCall,
|
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)]
|
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||||
|
#[serde(from = "RawChatResponseMessage")]
|
||||||
pub struct ChatResponseMessage {
|
pub struct ChatResponseMessage {
|
||||||
pub role: Role,
|
pub role: Role,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
@@ -502,6 +564,38 @@ pub struct ChatResponseMessage {
|
|||||||
pub citations: Option<Vec<String>>,
|
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)]
|
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||||
pub struct ToolCallResponse {
|
pub struct ToolCallResponse {
|
||||||
pub id: String,
|
pub id: String,
|
||||||
@@ -649,6 +743,7 @@ pub struct ToolCallFunctionDelta {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Serialize, Deserialize, Clone, Default)]
|
#[derive(Debug, Serialize, Deserialize, Clone, Default)]
|
||||||
|
#[serde(from = "RawChatChunkDelta")]
|
||||||
pub struct ChatChunkDelta {
|
pub struct ChatChunkDelta {
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub role: Option<Role>,
|
pub role: Option<Role>,
|
||||||
@@ -656,16 +751,43 @@ pub struct ChatChunkDelta {
|
|||||||
pub content: Option<String>,
|
pub content: Option<String>,
|
||||||
pub reasoning_content: Option<String>,
|
pub reasoning_content: Option<String>,
|
||||||
/// Tool call deltas. Handles `null` in JSON as empty vec.
|
/// Tool call deltas. Handles `null` in JSON as empty vec.
|
||||||
#[serde(
|
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||||
default,
|
|
||||||
skip_serializing_if = "Vec::is_empty",
|
|
||||||
deserialize_with = "deserialize_null_default"
|
|
||||||
)]
|
|
||||||
pub tool_calls: Vec<ToolCallDelta>,
|
pub tool_calls: Vec<ToolCallDelta>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub tool_call_id: Option<String>,
|
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.
|
/// Parameters to control realtime data.
|
||||||
#[derive(Serialize, Deserialize, Clone, Debug)]
|
#[derive(Serialize, Deserialize, Clone, Debug)]
|
||||||
pub struct SearchParameters {
|
pub struct SearchParameters {
|
||||||
@@ -787,22 +909,23 @@ pub enum ReasoningEffort {
|
|||||||
Medium,
|
Medium,
|
||||||
High,
|
High,
|
||||||
Xhigh,
|
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 {
|
impl ReasoningEffort {
|
||||||
pub fn to_responses_api(self) -> crate::rs::ReasoningEffort {
|
/// The canonical effort behind a typed Responses-API echo
|
||||||
match self {
|
/// (`response.reasoning.effort`). `max` echoes never reach the typed
|
||||||
Self::None => crate::rs::ReasoningEffort::None,
|
/// enum — [`normalize_effort_echo`] drops them pre-parse (async-openai
|
||||||
Self::Minimal => crate::rs::ReasoningEffort::Minimal,
|
/// has no such variant); the request direction writes the wire string
|
||||||
Self::Low => crate::rs::ReasoningEffort::Low,
|
/// via [`patch_reasoning_effort`].
|
||||||
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`.
|
|
||||||
pub fn from_responses_api(effort: crate::rs::ReasoningEffort) -> Self {
|
pub fn from_responses_api(effort: crate::rs::ReasoningEffort) -> Self {
|
||||||
match effort {
|
match effort {
|
||||||
crate::rs::ReasoningEffort::None => Self::None,
|
crate::rs::ReasoningEffort::None => Self::None,
|
||||||
@@ -822,17 +945,25 @@ impl ReasoningEffort {
|
|||||||
Self::Medium => "medium",
|
Self::Medium => "medium",
|
||||||
Self::High => "high",
|
Self::High => "high",
|
||||||
Self::Xhigh => "xhigh",
|
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> {
|
pub fn to_messages_api(self) -> Option<&'static str> {
|
||||||
match self {
|
match self {
|
||||||
Self::None | Self::Minimal => None,
|
Self::None | Self::Minimal => None,
|
||||||
Self::Low => Some("low"),
|
Self::Low => Some("low"),
|
||||||
Self::Medium => Some("medium"),
|
Self::Medium => Some("medium"),
|
||||||
Self::High => Some("high"),
|
Self::High => Some("high"),
|
||||||
Self::Xhigh => Some("max"),
|
Self::Xhigh => Some("xhigh"),
|
||||||
|
Self::Max => Some("max"),
|
||||||
|
Self::Ultra => Some("ultra"),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -853,19 +984,212 @@ impl std::str::FromStr for ReasoningEffort {
|
|||||||
"low" => Ok(Self::Low),
|
"low" => Ok(Self::Low),
|
||||||
"medium" => Ok(Self::Medium),
|
"medium" => Ok(Self::Medium),
|
||||||
"high" => Ok(Self::High),
|
"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!(
|
_ => 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> {
|
pub fn parse_canonical_effort_token(token: &str) -> Option<ReasoningEffort> {
|
||||||
token.parse().ok()
|
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()),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 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 REASONING_EFFORT_META_KEY: &str = "reasoningEffort";
|
||||||
pub const SUPPORTS_REASONING_EFFORT_META_KEY: &str = "supportsReasoningEffort";
|
pub const SUPPORTS_REASONING_EFFORT_META_KEY: &str = "supportsReasoningEffort";
|
||||||
|
|
||||||
@@ -1054,6 +1378,11 @@ pub struct SamplingConfig {
|
|||||||
pub max_completion_tokens: Option<u32>,
|
pub max_completion_tokens: Option<u32>,
|
||||||
pub temperature: Option<f32>,
|
pub temperature: Option<f32>,
|
||||||
pub top_p: 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
|
/// Which API backend to use for this model
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub api_backend: ApiBackend,
|
pub api_backend: ApiBackend,
|
||||||
@@ -1081,6 +1410,12 @@ pub struct CreateResponseWrapper {
|
|||||||
/// The inner Responses API request.
|
/// The inner Responses API request.
|
||||||
pub inner: crate::rs::CreateResponse,
|
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.
|
/// Custom header: conversation ID for tracking.
|
||||||
pub x_kigi_conv_id: Option<String>,
|
pub x_kigi_conv_id: Option<String>,
|
||||||
|
|
||||||
@@ -1107,6 +1442,7 @@ impl CreateResponseWrapper {
|
|||||||
pub fn new(inner: crate::rs::CreateResponse) -> Self {
|
pub fn new(inner: crate::rs::CreateResponse) -> Self {
|
||||||
Self {
|
Self {
|
||||||
inner,
|
inner,
|
||||||
|
reasoning_effort: None,
|
||||||
x_kigi_conv_id: None,
|
x_kigi_conv_id: None,
|
||||||
x_kigi_req_id: None,
|
x_kigi_req_id: None,
|
||||||
x_kigi_session_id: None,
|
x_kigi_session_id: None,
|
||||||
@@ -1215,6 +1551,113 @@ mod tests {
|
|||||||
use super::*;
|
use super::*;
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
|
|
||||||
|
/// 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]
|
#[test]
|
||||||
fn reasoning_effort_serde_lowercase_round_trip() {
|
fn reasoning_effort_serde_lowercase_round_trip() {
|
||||||
for v in [
|
for v in [
|
||||||
@@ -1224,6 +1667,7 @@ mod tests {
|
|||||||
ReasoningEffort::Medium,
|
ReasoningEffort::Medium,
|
||||||
ReasoningEffort::High,
|
ReasoningEffort::High,
|
||||||
ReasoningEffort::Xhigh,
|
ReasoningEffort::Xhigh,
|
||||||
|
ReasoningEffort::Max,
|
||||||
] {
|
] {
|
||||||
let json = serde_json::to_string(&v).unwrap();
|
let json = serde_json::to_string(&v).unwrap();
|
||||||
assert_eq!(json, format!("\"{}\"", v.as_str()), "serialize {v:?}");
|
assert_eq!(json, format!("\"{}\"", v.as_str()), "serialize {v:?}");
|
||||||
@@ -1231,31 +1675,112 @@ mod tests {
|
|||||||
assert_eq!(back, v, "round-trip {v:?}");
|
assert_eq!(back, v, "round-trip {v:?}");
|
||||||
}
|
}
|
||||||
assert!(serde_json::from_str::<ReasoningEffort>("\"BOGUS\"").is_err());
|
assert!(serde_json::from_str::<ReasoningEffort>("\"BOGUS\"").is_err());
|
||||||
assert!(serde_json::from_str::<ReasoningEffort>("\"max\"").is_err());
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn reasoning_effort_from_str_accepts_max_as_xhigh() {
|
fn reasoning_effort_from_str_max_and_xhigh_are_distinct() {
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
"max".parse::<ReasoningEffort>().unwrap(),
|
"max".parse::<ReasoningEffort>().unwrap(),
|
||||||
ReasoningEffort::Xhigh
|
ReasoningEffort::Max
|
||||||
);
|
);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
"MAX".parse::<ReasoningEffort>().unwrap(),
|
"MAX".parse::<ReasoningEffort>().unwrap(),
|
||||||
ReasoningEffort::Xhigh
|
ReasoningEffort::Max
|
||||||
);
|
);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
"xhigh".parse::<ReasoningEffort>().unwrap(),
|
"xhigh".parse::<ReasoningEffort>().unwrap(),
|
||||||
ReasoningEffort::Xhigh
|
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]
|
#[test]
|
||||||
fn parse_canonical_effort_token_helper() {
|
fn parse_canonical_effort_token_helper() {
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
parse_canonical_effort_token("max"),
|
parse_canonical_effort_token("max"),
|
||||||
Some(ReasoningEffort::Xhigh)
|
Some(ReasoningEffort::Max)
|
||||||
);
|
);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
parse_canonical_effort_token("high"),
|
parse_canonical_effort_token("high"),
|
||||||
@@ -1415,7 +1940,8 @@ mod tests {
|
|||||||
);
|
);
|
||||||
let bad_type = as_map(serde_json::json!({"reasoningEffort": 3}));
|
let bad_type = as_map(serde_json::json!({"reasoningEffort": 3}));
|
||||||
assert_eq!(parse_reasoning_effort_meta(Some(&bad_type)), None);
|
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);
|
assert_eq!(parse_reasoning_effort_meta(Some(&unknown)), None);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -13,7 +13,7 @@ default-bazel = []
|
|||||||
[dependencies]
|
[dependencies]
|
||||||
anyhow = { workspace = true }
|
anyhow = { workspace = true }
|
||||||
chrono = { workspace = true }
|
chrono = { workspace = true }
|
||||||
reqwest = { workspace = true, features = ["blocking"] }
|
reqwest = { workspace = true }
|
||||||
serde = { workspace = true }
|
serde = { workspace = true }
|
||||||
serde_json = { workspace = true }
|
serde_json = { workspace = true }
|
||||||
thiserror = { workspace = true }
|
thiserror = { workspace = true }
|
||||||
|
|||||||
@@ -1,362 +0,0 @@
|
|||||||
//! Changelog fetching from CDN with local disk cache.
|
|
||||||
//!
|
|
||||||
//! Both markdown (`*.external.md`) and JSON (`*.external.json`) changelogs
|
|
||||||
//! are published per-version alongside the Kigi GitHub distribution.
|
|
||||||
//!
|
|
||||||
//! `ChangelogManager::fetch()` retrieves both formats in parallel and
|
|
||||||
//! returns a `Changelog` with optional markdown + structured entries.
|
|
||||||
//! Consumers pick the format they need:
|
|
||||||
//! - `/release-notes` uses `changelog.markdown` for rich scrollback display
|
|
||||||
//! - Welcome screen uses `changelog.entries` for bullet rendering
|
|
||||||
|
|
||||||
use std::path::PathBuf;
|
|
||||||
|
|
||||||
/// Base URL for published changelogs. Kigi distributes via GitHub, so
|
|
||||||
/// per-version changelogs live in the release repository. Unreachable or
|
|
||||||
/// missing files degrade gracefully to the on-disk cache (see `fetch_with`).
|
|
||||||
const CHANGELOG_BASE: &str =
|
|
||||||
"https://raw.githubusercontent.com/ZacharyZhang-NY/Kigi-CLI/main/changelogs";
|
|
||||||
const FETCH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(3);
|
|
||||||
|
|
||||||
/// A single structured changelog entry from the published JSON changelog.
|
|
||||||
///
|
|
||||||
/// Shape must match the output of `render_external_json` in `changelog.sh`:
|
|
||||||
/// `{category, description, breaking_change}`
|
|
||||||
/// If you change fields here, update `changelog.sh:render_external_json` too.
|
|
||||||
///
|
|
||||||
/// All fields use `#[serde(default)]` so a single malformed entry doesn't
|
|
||||||
/// kill the entire array parse. Entries with an empty description are
|
|
||||||
/// filtered out by `bullets_from_entries`.
|
|
||||||
#[derive(Debug, Clone, serde::Deserialize)]
|
|
||||||
pub struct ChangelogEntry {
|
|
||||||
/// Category label (e.g. "features", "fixes", "breaking", "performance").
|
|
||||||
#[serde(default)]
|
|
||||||
pub category: String,
|
|
||||||
/// Human-readable description (may contain `**bold**` or backticks).
|
|
||||||
#[serde(default)]
|
|
||||||
pub description: String,
|
|
||||||
/// Whether this entry represents a breaking change.
|
|
||||||
#[serde(default)]
|
|
||||||
pub breaking_change: bool,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Both formats of a version's changelog, fetched together.
|
|
||||||
pub struct Changelog {
|
|
||||||
/// Rendered markdown (for `/release-notes` display).
|
|
||||||
pub markdown: Option<String>,
|
|
||||||
/// Structured entries (for welcome screen bullets).
|
|
||||||
pub entries: Option<Vec<ChangelogEntry>>,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Manages changelog retrieval from CDN with local disk caching.
|
|
||||||
///
|
|
||||||
/// Single entry point: `fetch()` returns both markdown and JSON in one
|
|
||||||
/// `Changelog` struct. Each format is fetched independently with its own
|
|
||||||
/// cache file, so a failure in one doesn't block the other.
|
|
||||||
pub struct ChangelogManager {
|
|
||||||
md_cache: PathBuf,
|
|
||||||
json_cache: PathBuf,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Default for ChangelogManager {
|
|
||||||
fn default() -> Self {
|
|
||||||
Self::new()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl ChangelogManager {
|
|
||||||
pub fn new() -> Self {
|
|
||||||
// Prefer live `$KIGI_SHARE_DIR` so harness-injected homes (PTY e2e) always
|
|
||||||
// win over a OnceLock that may have been initialised earlier with a
|
|
||||||
// different path in the same process graph.
|
|
||||||
Self::from_env_home()
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Resolve cache paths from the live process environment (not the
|
|
||||||
/// `kigi_home()` OnceLock). A seeded `$KIGI_SHARE_DIR` set on the pager
|
|
||||||
/// process is always honoured even if some earlier init path cached a
|
|
||||||
/// different home.
|
|
||||||
fn from_env_home() -> Self {
|
|
||||||
let home = std::env::var_os("KIGI_SHARE_DIR")
|
|
||||||
.map(std::path::PathBuf::from)
|
|
||||||
.filter(|p| !p.as_os_str().is_empty())
|
|
||||||
.unwrap_or_else(crate::util::kigi_home::kigi_home);
|
|
||||||
Self {
|
|
||||||
md_cache: home.join("CHANGELOG.md"),
|
|
||||||
json_cache: home.join("CHANGELOG.json"),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Fetch both markdown and JSON changelogs for the current version.
|
|
||||||
///
|
|
||||||
/// Each format is fetched independently (CDN, 3 s timeout) and cached
|
|
||||||
/// to disk. On failure, falls back to the cached copy. Either field
|
|
||||||
/// may be `None` if offline with no cache.
|
|
||||||
///
|
|
||||||
/// When `KIGI_CHANGELOG_OFFLINE` is set (PTY / integration tests), skip
|
|
||||||
/// the CDN entirely and read only the disk cache so seeded fixtures win
|
|
||||||
/// deterministically without network races. Paths are re-resolved from
|
|
||||||
/// `$KIGI_SHARE_DIR` so harness-injected env always applies.
|
|
||||||
///
|
|
||||||
/// JSON is only cached after a successful parse to avoid poisoning the
|
|
||||||
/// disk cache with malformed content (the markdown cache is write-through
|
|
||||||
/// since it's consumed as raw text).
|
|
||||||
pub fn fetch(&self) -> Changelog {
|
|
||||||
// Always re-resolve from env so a caller holding an older manager
|
|
||||||
// (or OnceLock lag) still reads the live harness home.
|
|
||||||
Self::from_env_home().fetch_with(changelog_offline(), CHANGELOG_BASE)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Fetch using this manager's already-resolved cache paths, an explicit
|
|
||||||
/// offline flag, and an explicit CDN base.
|
|
||||||
///
|
|
||||||
/// Split out of [`fetch`] so unit tests can drive it against a temp home
|
|
||||||
/// without mutating process-global env (`KIGI_SHARE_DIR` /
|
|
||||||
/// `KIGI_CHANGELOG_OFFLINE`), which races across the parallel test
|
|
||||||
/// harness. Passing an unreachable `base` lets a test force a
|
|
||||||
/// deterministic CDN miss instead of depending on whether the sandbox
|
|
||||||
/// happens to block network. Production callers always go through
|
|
||||||
/// [`fetch`], so behaviour is unchanged.
|
|
||||||
fn fetch_with(&self, offline: bool, base: &str) -> Changelog {
|
|
||||||
if offline {
|
|
||||||
return Changelog {
|
|
||||||
markdown: read_cache(&self.md_cache),
|
|
||||||
entries: self.read_json_cache(),
|
|
||||||
};
|
|
||||||
}
|
|
||||||
|
|
||||||
let version = kigi_version::VERSION;
|
|
||||||
let md_url = format!("{}/{}.external.md", base, version);
|
|
||||||
|
|
||||||
// Fetch both formats in parallel (3s timeout each → 3s total, not 6s).
|
|
||||||
let mut markdown = None;
|
|
||||||
let mut entries = None;
|
|
||||||
std::thread::scope(|s| {
|
|
||||||
let md_handle = s.spawn(|| self.fetch_and_cache(&md_url, &self.md_cache));
|
|
||||||
let json_handle = s.spawn(|| self.fetch_json(base, version));
|
|
||||||
markdown = md_handle.join().ok().flatten();
|
|
||||||
entries = json_handle.join().ok().flatten();
|
|
||||||
});
|
|
||||||
|
|
||||||
// If CDN is unreachable (CI sandboxes, airplane mode), fall back to
|
|
||||||
// any on-disk seed under `$KIGI_SHARE_DIR` even when offline mode was not
|
|
||||||
// explicitly requested — keeps PTY/integration tests deterministic.
|
|
||||||
if markdown.is_none() {
|
|
||||||
markdown = read_cache(&self.md_cache);
|
|
||||||
}
|
|
||||||
if entries.is_none() {
|
|
||||||
entries = self.read_json_cache();
|
|
||||||
}
|
|
||||||
|
|
||||||
Changelog { markdown, entries }
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Fetch and parse JSON changelog, caching only after successful parse.
|
|
||||||
fn fetch_json(&self, base: &str, version: &str) -> Option<Vec<ChangelogEntry>> {
|
|
||||||
let url = format!("{}/{}.external.json", base, version);
|
|
||||||
|
|
||||||
// Try remote first — only cache after successful parse.
|
|
||||||
if let Ok(raw) = fetch_blocking(&url)
|
|
||||||
&& !raw.trim().is_empty()
|
|
||||||
{
|
|
||||||
match serde_json::from_str::<Vec<ChangelogEntry>>(&raw) {
|
|
||||||
Ok(entries) => {
|
|
||||||
if let Err(e) = std::fs::write(&self.json_cache, &raw) {
|
|
||||||
tracing::debug!(error = %e, "JSON changelog cache write failed");
|
|
||||||
}
|
|
||||||
return Some(entries);
|
|
||||||
}
|
|
||||||
Err(e) => {
|
|
||||||
tracing::debug!(error = %e, "failed to parse JSON changelog from CDN");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
self.read_json_cache()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn read_json_cache(&self) -> Option<Vec<ChangelogEntry>> {
|
|
||||||
let cached = read_cache(&self.json_cache)?;
|
|
||||||
match serde_json::from_str(&cached) {
|
|
||||||
Ok(entries) => Some(entries),
|
|
||||||
Err(e) => {
|
|
||||||
tracing::debug!(error = %e, "failed to parse cached JSON changelog");
|
|
||||||
None
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Shared fetch-and-cache: try remote (3 s timeout), cache on success,
|
|
||||||
/// fall back to disk cache on failure.
|
|
||||||
fn fetch_and_cache(&self, url: &str, cache_path: &std::path::Path) -> Option<String> {
|
|
||||||
if let Ok(content) = fetch_blocking(url)
|
|
||||||
&& !content.trim().is_empty()
|
|
||||||
{
|
|
||||||
if let Err(e) = std::fs::write(cache_path, &content) {
|
|
||||||
tracing::debug!(error = %e, path = %cache_path.display(), "cache write failed");
|
|
||||||
}
|
|
||||||
return Some(content);
|
|
||||||
}
|
|
||||||
read_cache(cache_path)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// When set, `ChangelogManager::fetch` skips the CDN and only reads disk cache.
|
|
||||||
/// Used by PTY harness tests that seed `CHANGELOG.{md,json}` under a temp home.
|
|
||||||
fn changelog_offline() -> bool {
|
|
||||||
std::env::var_os("KIGI_CHANGELOG_OFFLINE").is_some_and(|v| !v.is_empty() && v != "0")
|
|
||||||
}
|
|
||||||
|
|
||||||
fn read_cache(path: &std::path::Path) -> Option<String> {
|
|
||||||
std::fs::read_to_string(path)
|
|
||||||
.ok()
|
|
||||||
.filter(|c| !c.trim().is_empty())
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Strip `**bold**` markers and backticks from a description string.
|
|
||||||
fn strip_markdown_inline(s: &str) -> String {
|
|
||||||
s.replace("**", "").replace('`', "")
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Convert changelog entries to plain-text bullet strings.
|
|
||||||
///
|
|
||||||
/// Strips `**bold**` and backtick formatting from each description,
|
|
||||||
/// skips entries with empty descriptions (from tolerant deserialization),
|
|
||||||
/// and returns at most `max` entries.
|
|
||||||
pub fn bullets_from_entries(entries: &[ChangelogEntry], max: usize) -> Vec<String> {
|
|
||||||
entries
|
|
||||||
.iter()
|
|
||||||
.filter(|e| !e.description.is_empty())
|
|
||||||
.take(max)
|
|
||||||
.map(|e| strip_markdown_inline(&e.description))
|
|
||||||
.collect()
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Blocking HTTP fetch. Callers (`std::thread::scope` threads) are already
|
|
||||||
/// off the tokio runtime, so no extra thread spawn is needed.
|
|
||||||
fn fetch_blocking(url: &str) -> anyhow::Result<String> {
|
|
||||||
let client = reqwest::blocking::Client::builder()
|
|
||||||
.timeout(FETCH_TIMEOUT)
|
|
||||||
.build()?;
|
|
||||||
let resp = client.get(url).send()?;
|
|
||||||
if !resp.status().is_success() {
|
|
||||||
anyhow::bail!("HTTP {}", resp.status());
|
|
||||||
}
|
|
||||||
Ok(resp.text()?)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::*;
|
|
||||||
|
|
||||||
/// Build a manager pointing at `home` directly, bypassing the global
|
|
||||||
/// `$KIGI_SHARE_DIR` env so tests never race the parallel harness.
|
|
||||||
fn manager_for(home: &std::path::Path) -> ChangelogManager {
|
|
||||||
ChangelogManager {
|
|
||||||
md_cache: home.join("CHANGELOG.md"),
|
|
||||||
json_cache: home.join("CHANGELOG.json"),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn offline_mode_reads_seeded_disk_cache_only() {
|
|
||||||
let tmp = tempfile::tempdir().unwrap();
|
|
||||||
let home = tmp.path().join("kigi-home");
|
|
||||||
std::fs::create_dir_all(&home).unwrap();
|
|
||||||
std::fs::write(home.join("CHANGELOG.md"), "# seeded offline md\n").unwrap();
|
|
||||||
std::fs::write(
|
|
||||||
home.join("CHANGELOG.json"),
|
|
||||||
r#"[{"category":"features","description":"seeded entry","breaking_change":false}]"#,
|
|
||||||
)
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
// Offline path: read only the seeded disk cache, no network.
|
|
||||||
let changelog = manager_for(&home).fetch_with(true, CHANGELOG_BASE);
|
|
||||||
assert_eq!(
|
|
||||||
changelog.markdown.as_deref(),
|
|
||||||
Some("# seeded offline md\n"),
|
|
||||||
"offline mode must return seeded markdown"
|
|
||||||
);
|
|
||||||
let entries = changelog.entries.expect("seeded json entries");
|
|
||||||
assert_eq!(entries.len(), 1);
|
|
||||||
assert_eq!(entries[0].description, "seeded entry");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn cdn_miss_falls_back_to_env_home_disk_cache() {
|
|
||||||
let tmp = tempfile::tempdir().unwrap();
|
|
||||||
let home = tmp.path().join("kigi-home-fallback");
|
|
||||||
std::fs::create_dir_all(&home).unwrap();
|
|
||||||
std::fs::write(home.join("CHANGELOG.md"), "# fallback md\n").unwrap();
|
|
||||||
|
|
||||||
// Non-offline path with an unreachable CDN base: the remote fetch
|
|
||||||
// fails deterministically (no dependency on the sandbox blocking
|
|
||||||
// network), so the on-disk cache must win.
|
|
||||||
let changelog = manager_for(&home).fetch_with(false, "http://127.0.0.1:1");
|
|
||||||
assert_eq!(
|
|
||||||
changelog.markdown.as_deref(),
|
|
||||||
Some("# fallback md\n"),
|
|
||||||
"CDN miss must fall back to the seeded CHANGELOG.md"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn bullets_strips_markdown_and_respects_max() {
|
|
||||||
let entries = vec![
|
|
||||||
ChangelogEntry {
|
|
||||||
category: "features".into(),
|
|
||||||
description: "Added **dark mode** support".into(),
|
|
||||||
breaking_change: false,
|
|
||||||
},
|
|
||||||
ChangelogEntry {
|
|
||||||
category: "fixes".into(),
|
|
||||||
description: "Fixed `crash` on startup".into(),
|
|
||||||
breaking_change: false,
|
|
||||||
},
|
|
||||||
ChangelogEntry {
|
|
||||||
category: "performance".into(),
|
|
||||||
description: "Faster **rendering** of `code` blocks".into(),
|
|
||||||
breaking_change: false,
|
|
||||||
},
|
|
||||||
];
|
|
||||||
|
|
||||||
let bullets = bullets_from_entries(&entries, 2);
|
|
||||||
assert_eq!(bullets.len(), 2);
|
|
||||||
assert_eq!(bullets[0], "Added dark mode support");
|
|
||||||
assert_eq!(bullets[1], "Fixed crash on startup");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn bullets_skips_empty_descriptions() {
|
|
||||||
let entries = vec![
|
|
||||||
ChangelogEntry {
|
|
||||||
category: "features".into(),
|
|
||||||
description: "Good entry".into(),
|
|
||||||
breaking_change: false,
|
|
||||||
},
|
|
||||||
ChangelogEntry {
|
|
||||||
category: String::new(),
|
|
||||||
description: String::new(), // bad entry from tolerant deser
|
|
||||||
breaking_change: false,
|
|
||||||
},
|
|
||||||
ChangelogEntry {
|
|
||||||
category: "fixes".into(),
|
|
||||||
description: "Another good one".into(),
|
|
||||||
breaking_change: false,
|
|
||||||
},
|
|
||||||
];
|
|
||||||
let bullets = bullets_from_entries(&entries, 10);
|
|
||||||
assert_eq!(bullets, vec!["Good entry", "Another good one"]);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn tolerant_deserialization_partial_entry() {
|
|
||||||
// Missing description field → defaults to empty string, not a parse error
|
|
||||||
let json = r#"[{"category":"features"},{"description":"ok"}]"#;
|
|
||||||
let entries: Vec<ChangelogEntry> = serde_json::from_str(json).unwrap();
|
|
||||||
assert_eq!(entries.len(), 2);
|
|
||||||
assert_eq!(entries[0].description, "");
|
|
||||||
assert_eq!(entries[1].category, "");
|
|
||||||
assert_eq!(entries[1].description, "ok");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,4 +1,3 @@
|
|||||||
pub mod changelog;
|
|
||||||
pub mod event_id;
|
pub mod event_id;
|
||||||
pub mod kigi_home;
|
pub mod kigi_home;
|
||||||
pub mod secure_file;
|
pub mod secure_file;
|
||||||
@@ -35,7 +34,17 @@ pub fn random_f64() -> f64 {
|
|||||||
pub fn probabilistic_sample(rate: f64) -> bool {
|
pub fn probabilistic_sample(rate: f64) -> bool {
|
||||||
random_f64() < rate
|
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 {
|
let Ok(candidate) = reqwest::Url::parse(candidate) else {
|
||||||
return false;
|
return false;
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -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() {
|
if auth.is_some() || endpoints.has_custom_endpoint() || platform_keys.any() {
|
||||||
tokio::task::spawn_blocking(move || {
|
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
|
.await
|
||||||
.ok()
|
.ok()
|
||||||
@@ -621,6 +627,7 @@ pub async fn run_leader(
|
|||||||
crate::agent::models::prefetch_models_blocking(
|
crate::agent::models::prefetch_models_blocking(
|
||||||
&endpoints_for_prefetch,
|
&endpoints_for_prefetch,
|
||||||
auth_for_prefetch.as_ref(),
|
auth_for_prefetch.as_ref(),
|
||||||
|
&Default::default(),
|
||||||
fetch_auth_for_prefetch,
|
fetch_auth_for_prefetch,
|
||||||
&platform_keys_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.kind == ModelOverrideWarningKind::DuplicateAlias
|
||||||
&& w.field.as_deref() == Some("send_compactions_remaining")
|
&& 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"));
|
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
|
.models_manager
|
||||||
.model_supports_reasoning_effort(model_id.0.as_ref())
|
.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!(
|
tracing::info!(
|
||||||
session_id = % session_id.0, effort = % eff,
|
session_id = % session_id.0, effort = % eff,
|
||||||
"set_session_model: applying reasoning_effort override from meta"
|
"set_session_model: applying reasoning_effort override from meta"
|
||||||
@@ -169,9 +192,18 @@ pub(crate) async fn apply(
|
|||||||
model.map(|e| &e.info),
|
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 (tx, rx) = oneshot::channel();
|
||||||
let _ = handle.cmd_tx.send(SessionCommand::SetSessionModel {
|
let _ = handle.cmd_tx.send(SessionCommand::SetSessionModel {
|
||||||
sampling_config: model_sampling,
|
sampling_config: model_sampling,
|
||||||
|
catalog_key,
|
||||||
use_concise,
|
use_concise,
|
||||||
apply_prompt_override,
|
apply_prompt_override,
|
||||||
skip_prompt_rewrite: did_rebuild || model_unchanged,
|
skip_prompt_rewrite: did_rebuild || model_unchanged,
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ pub mod auth_method;
|
|||||||
pub mod chat_modes;
|
pub mod chat_modes;
|
||||||
pub mod config;
|
pub mod config;
|
||||||
pub mod config_model_override_parse;
|
pub mod config_model_override_parse;
|
||||||
|
pub(crate) mod enrichment_fetch;
|
||||||
mod ext_parsers;
|
mod ext_parsers;
|
||||||
pub(crate) mod feedback_client;
|
pub(crate) mod feedback_client;
|
||||||
pub mod folder_trust;
|
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(),
|
&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)");
|
tracing::info!("auth: loaded API key from auth.json (xai::api_key scope)");
|
||||||
kigi_log::unified_log::info(
|
kigi_log::unified_log::info(
|
||||||
"auth: loaded API key from auth.json (xai::api_key scope)",
|
"auth: loaded API key from auth.json (xai::api_key scope)",
|
||||||
@@ -221,7 +221,24 @@ impl acp::Agent for MvpAgent {
|
|||||||
has_cached_token,
|
has_cached_token,
|
||||||
login_label: None,
|
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(
|
kigi_log::unified_log::info(
|
||||||
"auth: initialize() built auth_methods for ACP response",
|
"auth: initialize() built auth_methods for ACP response",
|
||||||
None,
|
None,
|
||||||
@@ -330,10 +347,28 @@ impl acp::Agent for MvpAgent {
|
|||||||
);
|
);
|
||||||
match arguments.method_id.0.as_ref() {
|
match arguments.method_id.0.as_ref() {
|
||||||
auth_method::XAI_API_KEY_METHOD_ID => {
|
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();
|
let mut sampling_config = self.sampling_config.borrow_mut();
|
||||||
if sampling_config.api_key.is_none() {
|
if sampling_config.api_key.is_none() {
|
||||||
if let Ok(api_key) = auth_method::read_xai_api_key_env() {
|
if let Ok(api_key) = auth_method::read_xai_api_key_env() {
|
||||||
|
if takes_house_key {
|
||||||
sampling_config.api_key = Some(api_key.clone());
|
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(
|
if let Err(e) = crate::auth::store_api_key(
|
||||||
&crate::util::kigi_home::kigi_home(),
|
&crate::util::kigi_home::kigi_home(),
|
||||||
&api_key,
|
&api_key,
|
||||||
@@ -357,7 +392,7 @@ impl acp::Agent for MvpAgent {
|
|||||||
return Err(
|
return Err(
|
||||||
acp::Error::auth_required()
|
acp::Error::auth_required()
|
||||||
.data(
|
.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() {
|
let message = if self.auth_manager.is_expired() {
|
||||||
"Session expired, re-authentication required"
|
"Session expired, re-authentication required"
|
||||||
} else {
|
} else {
|
||||||
@@ -408,9 +443,15 @@ impl acp::Agent for MvpAgent {
|
|||||||
.await;
|
.await;
|
||||||
};
|
};
|
||||||
self.emit_settings_update_notification();
|
self.emit_settings_update_notification();
|
||||||
{
|
// H2/C1: route the stamp through the ONE guard, which asks the
|
||||||
let mut sampling_config = self.sampling_config.borrow_mut();
|
// authority which credential governs the shared config rather
|
||||||
sampling_config.api_key = Some(auth.key);
|
// 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!(
|
tracing::debug!(
|
||||||
"auth: cached_token handler set api_key (SessionToken)"
|
"auth: cached_token handler set api_key (SessionToken)"
|
||||||
);
|
);
|
||||||
@@ -484,9 +525,15 @@ impl acp::Agent for MvpAgent {
|
|||||||
err.message = e.to_string();
|
err.message = e.to_string();
|
||||||
err
|
err
|
||||||
})?;
|
})?;
|
||||||
{
|
// C1: hot-swap FIRST, then let the authority read the fresh
|
||||||
let mut sampling_config = self.sampling_config.borrow_mut();
|
// token back where it belongs. Nothing hand-carries `auth.key`
|
||||||
sampling_config.api_key = Some(auth.key.clone());
|
// 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!(
|
tracing::debug!(
|
||||||
"auth: kimi.com/oidc handler set api_key (SessionToken)"
|
"auth: kimi.com/oidc handler set api_key (SessionToken)"
|
||||||
);
|
);
|
||||||
@@ -496,7 +543,6 @@ impl acp::Agent for MvpAgent {
|
|||||||
None,
|
None,
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
self.auth_manager.hot_swap(auth.clone());
|
|
||||||
self.emit_settings_update_notification();
|
self.emit_settings_update_notification();
|
||||||
self.set_auth_method(arguments.method_id.clone());
|
self.set_auth_method(arguments.method_id.clone());
|
||||||
self.models_manager.on_auth_changed().await;
|
self.models_manager.on_auth_changed().await;
|
||||||
@@ -508,15 +554,22 @@ impl acp::Agent for MvpAgent {
|
|||||||
);
|
);
|
||||||
Ok(self.auth_response_with_meta())
|
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,
|
&arguments.method_id,
|
||||||
)
|
)
|
||||||
.expect("match arm guarantees a moonshot method id");
|
.oauth_platform()
|
||||||
self.authenticate_moonshot(platform, arguments.method_id.clone())
|
{
|
||||||
|
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
|
.await
|
||||||
}
|
} else {
|
||||||
_ => {
|
|
||||||
Err(
|
Err(
|
||||||
acp::Error::invalid_params()
|
acp::Error::invalid_params()
|
||||||
.data(
|
.data(
|
||||||
@@ -526,6 +579,7 @@ impl acp::Agent for MvpAgent {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
}
|
||||||
async fn new_session(
|
async fn new_session(
|
||||||
&self,
|
&self,
|
||||||
arguments: acp::NewSessionRequest,
|
arguments: acp::NewSessionRequest,
|
||||||
@@ -754,6 +808,7 @@ impl acp::Agent for MvpAgent {
|
|||||||
persisted_signals: None,
|
persisted_signals: None,
|
||||||
persisted_plan_mode: None,
|
persisted_plan_mode: None,
|
||||||
persisted_goal_mode: None,
|
persisted_goal_mode: None,
|
||||||
|
persisted_graph_mode: None,
|
||||||
persisted_announcement_state: None,
|
persisted_announcement_state: None,
|
||||||
session_meta: arguments.meta.as_ref(),
|
session_meta: arguments.meta.as_ref(),
|
||||||
model_agent_type: model_agent_type.as_deref(),
|
model_agent_type: model_agent_type.as_deref(),
|
||||||
@@ -985,6 +1040,7 @@ impl acp::Agent for MvpAgent {
|
|||||||
signals: persisted_signals,
|
signals: persisted_signals,
|
||||||
announcement_state: persisted_announcement_state,
|
announcement_state: persisted_announcement_state,
|
||||||
goal_mode_state: _persisted_goal_mode,
|
goal_mode_state: _persisted_goal_mode,
|
||||||
|
graph_mode_state: _persisted_graph_mode,
|
||||||
} = persistence_info;
|
} = persistence_info;
|
||||||
let restored_awaiting_plan_approval = persisted_plan_mode
|
let restored_awaiting_plan_approval = persisted_plan_mode
|
||||||
.as_ref()
|
.as_ref()
|
||||||
@@ -1207,6 +1263,7 @@ impl acp::Agent for MvpAgent {
|
|||||||
persisted_signals,
|
persisted_signals,
|
||||||
persisted_plan_mode,
|
persisted_plan_mode,
|
||||||
persisted_goal_mode: _persisted_goal_mode,
|
persisted_goal_mode: _persisted_goal_mode,
|
||||||
|
persisted_graph_mode: _persisted_graph_mode,
|
||||||
persisted_announcement_state,
|
persisted_announcement_state,
|
||||||
session_meta: request_meta.as_ref(),
|
session_meta: request_meta.as_ref(),
|
||||||
model_agent_type: persisted_agent_name.as_deref(),
|
model_agent_type: persisted_agent_name.as_deref(),
|
||||||
|
|||||||
@@ -25,20 +25,35 @@ impl MvpAgent {
|
|||||||
primary: &SamplingConfig,
|
primary: &SamplingConfig,
|
||||||
) -> Result<(OaiCompatClient, String), acp::Error> {
|
) -> Result<(OaiCompatClient, String), acp::Error> {
|
||||||
let slug = self.resolve_session_summary_model();
|
let slug = self.resolve_session_summary_model();
|
||||||
let session_key = self.auth_manager.current_or_expired().map(|a| a.key.clone());
|
|
||||||
let models = self.models_manager.models();
|
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 endpoints = self.models_manager.endpoints();
|
||||||
let alpha_test_key = self.cfg.borrow().endpoints.alpha_test_key.clone();
|
let alpha_test_key = self.cfg.borrow().endpoints.alpha_test_key.clone();
|
||||||
let config = match crate::agent::config::resolve_aux_model_sampling_config(
|
let config = match crate::agent::config::resolve_aux_model_sampling_config(
|
||||||
&slug,
|
&slug,
|
||||||
&models,
|
&models,
|
||||||
&endpoints,
|
&endpoints,
|
||||||
session_key.as_deref(),
|
session_key.as_ref(),
|
||||||
alpha_test_key,
|
alpha_test_key,
|
||||||
) {
|
) {
|
||||||
Some(mut cfg) => {
|
Some(mut cfg) => {
|
||||||
cfg.attribution_callback = primary.attribution_callback.clone();
|
cfg.attribution_callback = primary.attribution_callback.clone();
|
||||||
cfg.bearer_resolver = primary.bearer_resolver.clone();
|
// 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.max_retries = primary.max_retries;
|
||||||
cfg
|
cfg
|
||||||
}
|
}
|
||||||
@@ -52,6 +67,36 @@ impl MvpAgent {
|
|||||||
let client = OaiCompatClient::new(config).map_err(map_sampling_err_to_acp)?;
|
let client = OaiCompatClient::new(config).map_err(map_sampling_err_to_acp)?;
|
||||||
Ok((client, model))
|
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.
|
/// `true` for session-based ACP auth methods.
|
||||||
fn is_session_based_auth(&self) -> bool {
|
fn is_session_based_auth(&self) -> bool {
|
||||||
self.auth_method_id
|
self.auth_method_id
|
||||||
@@ -310,8 +355,13 @@ impl MvpAgent {
|
|||||||
pub(crate) fn command_availability(
|
pub(crate) fn command_availability(
|
||||||
&self,
|
&self,
|
||||||
) -> crate::session::slash_commands::CommandAvailability {
|
) -> crate::session::slash_commands::CommandAvailability {
|
||||||
|
let goal = self.cfg.borrow().resolve_goal().value;
|
||||||
crate::session::slash_commands::CommandAvailability {
|
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()
|
..crate::session::slash_commands::CommandAvailability::default()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -447,18 +497,19 @@ impl MvpAgent {
|
|||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
/// `authenticate(moonshot-cn / moonshot-ai)`: interactive open-platform
|
/// `authenticate(<api-key platform id>)`: interactive API-key login from
|
||||||
/// API-key login from the welcome picker.
|
/// the welcome picker for any non-OAuth registry platform.
|
||||||
///
|
///
|
||||||
/// Reloads the platform keys from disk+env (the TUI persists the pasted
|
/// Reloads the platform keys from disk+env (the TUI persists the pasted
|
||||||
/// key to `[platforms.<id>]` in config.toml immediately before this call),
|
/// key to auth.json under the platform-id scope immediately before this
|
||||||
/// fails with an actionable error when none is configured, validates the
|
/// call), fails with an actionable error when none is configured,
|
||||||
/// key against `GET {platform_base}/models`, then marks the session
|
/// validates the key against `GET {platform_base}/models`, then marks the
|
||||||
/// authenticated exactly like an external API key: publish the method id
|
/// session authenticated exactly like an external API key: publish the
|
||||||
/// (NOT session-based — no token refresh), swap the freshly-stamped config
|
/// method id (NOT session-based — no token refresh), swap the
|
||||||
/// into the models manager, and trigger the model sync so the catalog
|
/// freshly-stamped config into the models manager, and trigger the model
|
||||||
/// gains the platform's entries. The key itself is never logged.
|
/// sync so the catalog gains the platform's entries. The key itself is
|
||||||
pub(super) async fn authenticate_moonshot(
|
/// never logged.
|
||||||
|
pub(super) async fn authenticate_api_key_platform(
|
||||||
&self,
|
&self,
|
||||||
platform: kigi_models::PlatformId,
|
platform: kigi_models::PlatformId,
|
||||||
method_id: acp::AuthMethodId,
|
method_id: acp::AuthMethodId,
|
||||||
@@ -475,10 +526,11 @@ impl MvpAgent {
|
|||||||
Some("platform_key_invalid_or_missing"),
|
Some("platform_key_invalid_or_missing"),
|
||||||
);
|
);
|
||||||
})?;
|
})?;
|
||||||
// Swap the on-disk config (now carrying the key) into the models
|
// Rebuild the catalog from the on-disk config: the rebuild freshly
|
||||||
// manager so `apply_platform_credentials` stamps the platform's
|
// resolves platform keys (env > auth.json > config), so the key just
|
||||||
// catalog entries; a parse failure keeps the last-known-good config
|
// persisted to auth.json is stamped onto the platform's entries; a
|
||||||
// (`on_auth_changed` below still re-resolves keys from disk itself).
|
// 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()
|
match crate::config::load_effective_config()
|
||||||
.map_err(|e| e.to_string())
|
.map_err(|e| e.to_string())
|
||||||
.and_then(|raw| crate::agent::config::Config::new_from_toml_cfg(&raw))
|
.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());
|
.and_then(|v| v.as_object().cloned());
|
||||||
Ok(AuthenticateResponse::new().meta(meta))
|
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> {
|
pub(crate) fn deployment_key(&self) -> Option<String> {
|
||||||
self.cfg.borrow().endpoints.deployment_key.clone()
|
self.cfg.borrow().endpoints.deployment_key.clone()
|
||||||
}
|
}
|
||||||
@@ -586,21 +710,60 @@ impl MvpAgent {
|
|||||||
);
|
);
|
||||||
Ok(entry.clone())
|
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(
|
pub(crate) fn prepare_sampling_config_for_model(
|
||||||
&self,
|
&self,
|
||||||
model: &ModelEntry,
|
model: &ModelEntry,
|
||||||
origin_client: Option<crate::http::OriginClientInfo>,
|
origin_client: Option<crate::http::OriginClientInfo>,
|
||||||
) -> SamplingConfig {
|
) -> SamplingConfig {
|
||||||
let session = if self.is_session_based_auth() {
|
// Resolve the session token by the MODEL's platform, not the primary
|
||||||
self.auth_manager.current_or_expired()
|
// auth method: an oauth-platform model (xai-grok) uses its OWN
|
||||||
} else {
|
// pool-backed token (`None` — never the Kimi key — when the user has not
|
||||||
None
|
// 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 has_session_key = session.is_some();
|
||||||
let mut credentials = resolve_credentials(
|
let mut credentials = resolve_credentials(model, session.as_ref());
|
||||||
model,
|
|
||||||
session.as_ref().map(|a| a.key.as_str()),
|
|
||||||
);
|
|
||||||
if !has_session_key && credentials.auth_type == kigi_chat_state::AuthType::ApiKey
|
if !has_session_key && credentials.auth_type == kigi_chat_state::AuthType::ApiKey
|
||||||
&& !model.has_own_credentials() && self.is_session_based_auth()
|
&& !model.has_own_credentials() && self.is_session_based_auth()
|
||||||
{
|
{
|
||||||
@@ -787,7 +950,8 @@ impl MvpAgent {
|
|||||||
models_manager: crate::agent::models::ModelsManager,
|
models_manager: crate::agent::models::ModelsManager,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
models_manager.set_gateway(gateway.clone());
|
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 storage_mode = cfg.storage_mode;
|
||||||
let default_yolo_mode = cfg.default_yolo_mode;
|
let default_yolo_mode = cfg.default_yolo_mode;
|
||||||
let default_auto_mode = cfg.default_auto_mode;
|
let default_auto_mode = cfg.default_auto_mode;
|
||||||
@@ -842,7 +1006,8 @@ impl MvpAgent {
|
|||||||
models_manager,
|
models_manager,
|
||||||
cfg: RefCell::new(cfg.clone()),
|
cfg: RefCell::new(cfg.clone()),
|
||||||
auth_method_id: crate::agent::auth_method::new_shared_auth_method_id(None),
|
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_manager,
|
||||||
auth_code_tx: RefCell::new(None),
|
auth_code_tx: RefCell::new(None),
|
||||||
auth_url_rx: 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))
|
(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.
|
/// 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
|
/// `current_model_id` — that's deferred to session creation time to avoid
|
||||||
/// cross-client contamination in leader mode (where `current_model_id` is
|
/// cross-client contamination in leader mode (where `current_model_id` is
|
||||||
/// shared mutable state).
|
/// shared mutable state).
|
||||||
pub(super) fn seed_client_config_auth_if_available(&self) {
|
pub(super) fn seed_client_config_auth_if_available(&self) {
|
||||||
let mut sampling_config = self.sampling_config.borrow_mut();
|
if self.sampling_config.borrow().api_key.is_some() {
|
||||||
if sampling_config.api_key.is_none() {
|
return;
|
||||||
if let Some(auth) = self.auth_manager.current_or_expired() {
|
}
|
||||||
sampling_config.api_key = Some(auth.key);
|
if self.stamp_session_credential(false) {
|
||||||
tracing::debug!("auth: seed_client_config set auth (SessionToken)");
|
tracing::debug!("auth: seed_client_config set auth (SessionToken)");
|
||||||
kigi_log::unified_log::debug(
|
kigi_log::unified_log::debug(
|
||||||
"auth: seed_client_config set auth (SessionToken)",
|
"auth: seed_client_config set auth (SessionToken)",
|
||||||
None,
|
None,
|
||||||
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_manager
|
||||||
.models()
|
.models()
|
||||||
.values()
|
.values()
|
||||||
.any(|m| m.has_own_credentials())
|
.any(|m| m.has_own_credentials())
|
||||||
{
|
{
|
||||||
tracing::warn!(
|
tracing::warn!("No credentials found: no login token and no model api_key/env_key");
|
||||||
"No credentials found: no login token and no model api_key/env_key"
|
|
||||||
);
|
|
||||||
kigi_log::unified_log::warn(
|
kigi_log::unified_log::warn(
|
||||||
"No credentials found: no login token and no model api_key/env_key",
|
"No credentials found: no login token and no model api_key/env_key",
|
||||||
None,
|
None,
|
||||||
@@ -1465,7 +1726,6 @@ impl MvpAgent {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
|
||||||
/// Allocate the next monotonic telemetry turn number for a session.
|
/// Allocate the next monotonic telemetry turn number for a session.
|
||||||
///
|
///
|
||||||
/// Returns the current turn number and advances the counter. The counter is
|
/// Returns the current turn number and advances the counter. The counter is
|
||||||
@@ -1677,6 +1937,7 @@ impl MvpAgent {
|
|||||||
persisted_signals,
|
persisted_signals,
|
||||||
persisted_plan_mode,
|
persisted_plan_mode,
|
||||||
persisted_goal_mode,
|
persisted_goal_mode,
|
||||||
|
persisted_graph_mode,
|
||||||
persisted_announcement_state,
|
persisted_announcement_state,
|
||||||
session_meta,
|
session_meta,
|
||||||
model_agent_type,
|
model_agent_type,
|
||||||
@@ -2111,6 +2372,7 @@ impl MvpAgent {
|
|||||||
let web_fetch_config = self.prepare_web_fetch_config();
|
let web_fetch_config = self.prepare_web_fetch_config();
|
||||||
let write_file_enabled = self.cfg.borrow().resolve_write_file().value;
|
let write_file_enabled = self.cfg.borrow().resolve_write_file().value;
|
||||||
let goal_enabled = self.cfg.borrow().resolve_goal().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 subagents_enabled = self.cfg.borrow().subagents_enabled;
|
||||||
let ask_user_question_enabled = parse_ask_user_question_from_meta(session_meta)
|
let ask_user_question_enabled = parse_ask_user_question_from_meta(session_meta)
|
||||||
.unwrap_or_else(|| self.cfg.borrow().resolve_ask_user_question().value);
|
.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 (mut handle, agent_system_prompt, session_thread) = {
|
||||||
let _timer = crate::instrumentation_timer!("session.spawn_actor_call");
|
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 {
|
let credentials = kigi_chat_state::Credentials {
|
||||||
api_key: sampling_config.api_key.clone(),
|
api_key: sampling_config.api_key.clone(),
|
||||||
auth_type: crate::agent::config::resolve_chat_state_auth_type(
|
auth_type: crate::agent::config::resolve_chat_state_auth_type(
|
||||||
sampling_config.model.as_str(),
|
sampling_config.model.as_str(),
|
||||||
session_key.as_deref(),
|
session_key.as_ref(),
|
||||||
self.auth_type(),
|
self.auth_type(),
|
||||||
),
|
),
|
||||||
alpha_test_key: self.alpha_test_key(),
|
alpha_test_key: self.alpha_test_key(),
|
||||||
@@ -2314,6 +2598,7 @@ impl MvpAgent {
|
|||||||
persisted_signals,
|
persisted_signals,
|
||||||
persisted_plan_mode,
|
persisted_plan_mode,
|
||||||
persisted_goal_mode,
|
persisted_goal_mode,
|
||||||
|
persisted_graph_mode,
|
||||||
persisted_announcement_state,
|
persisted_announcement_state,
|
||||||
self.memory_config.clone(),
|
self.memory_config.clone(),
|
||||||
feedback_flags,
|
feedback_flags,
|
||||||
@@ -2328,6 +2613,7 @@ impl MvpAgent {
|
|||||||
app_builder_deployer_config,
|
app_builder_deployer_config,
|
||||||
write_file_enabled,
|
write_file_enabled,
|
||||||
goal_enabled,
|
goal_enabled,
|
||||||
|
graph_enabled,
|
||||||
subagents_enabled,
|
subagents_enabled,
|
||||||
ask_user_question_enabled,
|
ask_user_question_enabled,
|
||||||
client_hooks,
|
client_hooks,
|
||||||
|
|||||||
@@ -120,6 +120,7 @@ pub(crate) struct SessionSpawnOptions<'a> {
|
|||||||
pub persisted_signals: Option<crate::session::signals::SessionSignals>,
|
pub persisted_signals: Option<crate::session::signals::SessionSignals>,
|
||||||
pub persisted_plan_mode: Option<crate::session::plan_mode::PlanModeSnapshot>,
|
pub persisted_plan_mode: Option<crate::session::plan_mode::PlanModeSnapshot>,
|
||||||
pub persisted_goal_mode: Option<crate::session::goal_tracker::GoalOrchestration>,
|
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<
|
pub persisted_announcement_state: Option<
|
||||||
crate::session::announcement_state::AnnouncementState,
|
crate::session::announcement_state::AnnouncementState,
|
||||||
>,
|
>,
|
||||||
@@ -257,6 +258,7 @@ pub(crate) fn chat_session_spawn_options<'a>(
|
|||||||
persisted_signals: None,
|
persisted_signals: None,
|
||||||
persisted_plan_mode: None,
|
persisted_plan_mode: None,
|
||||||
persisted_goal_mode: None,
|
persisted_goal_mode: None,
|
||||||
|
persisted_graph_mode: None,
|
||||||
persisted_announcement_state: None,
|
persisted_announcement_state: None,
|
||||||
session_meta,
|
session_meta,
|
||||||
model_agent_type,
|
model_agent_type,
|
||||||
@@ -499,6 +501,19 @@ pub struct MvpAgent {
|
|||||||
/// only api_key is written here (same for all clients). Per-session base_url
|
/// 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`.
|
/// is resolved at session creation time in `new_session` / `load_session`.
|
||||||
pub(crate) sampling_config: RefCell<SamplingConfig>,
|
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) auth_manager: Arc<AuthManager>,
|
||||||
pub(crate) models_manager: crate::agent::models::ModelsManager,
|
pub(crate) models_manager: crate::agent::models::ModelsManager,
|
||||||
/// Forwards pasted codes from `handle_auth_submit_code` to the auth flow.
|
/// 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();
|
let cfg = AgentConfig::default();
|
||||||
MvpAgent::new(gateway, &cfg, auth_manager, None).expect("valid test config")
|
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
|
/// Regression: boot-time plugin discovery is deferred past ACP
|
||||||
/// `initialize`, so the shared plugin registry starts empty.
|
/// `initialize`, so the shared plugin registry starts empty.
|
||||||
/// `resolve_mcp_servers` reads that snapshot to merge plugin-contributed
|
/// `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;
|
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
|
/// No load in flight and no session → the wait returns immediately
|
||||||
/// (the caller then surfaces "unknown session id" exactly as before).
|
/// (the caller then surfaces "unknown session id" exactly as before).
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
@@ -1904,10 +1982,12 @@ async fn cached_token_fallthrough_prefers_api_key_for_deployment_key() {
|
|||||||
#[serial_test::serial]
|
#[serial_test::serial]
|
||||||
async fn cached_token_fallthrough_falls_to_kigi_com_without_credentials() {
|
async fn cached_token_fallthrough_falls_to_kigi_com_without_credentials() {
|
||||||
use crate::agent::auth_method::{
|
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;
|
use kigi_test_support::EnvGuard;
|
||||||
let _lockdown = EnvGuard::unset("KIGI_DISABLE_API_KEY_AUTH");
|
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 _new = EnvGuard::unset(XAI_API_KEY_ENV_VAR);
|
||||||
let _legacy = EnvGuard::unset(LEGACY_XAI_API_KEY_ENV_VAR);
|
let _legacy = EnvGuard::unset(LEGACY_XAI_API_KEY_ENV_VAR);
|
||||||
let agent = build_minimal_agent_for_tests();
|
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(
|
prefetch_models_blocking(
|
||||||
&agent_config.endpoints,
|
&agent_config.endpoints,
|
||||||
auth.as_ref(),
|
auth.as_ref(),
|
||||||
|
&Default::default(),
|
||||||
fetch_auth,
|
fetch_auth,
|
||||||
&platform_keys,
|
&platform_keys,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -59,6 +59,8 @@ fn effort_label(effort: ReasoningEffort) -> String {
|
|||||||
ReasoningEffort::Medium => "Medium",
|
ReasoningEffort::Medium => "Medium",
|
||||||
ReasoningEffort::High => "High",
|
ReasoningEffort::High => "High",
|
||||||
ReasoningEffort::Xhigh => "X-High",
|
ReasoningEffort::Xhigh => "X-High",
|
||||||
|
ReasoningEffort::Max => "Max",
|
||||||
|
ReasoningEffort::Ultra => "Ultra",
|
||||||
}
|
}
|
||||||
.to_string()
|
.to_string()
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1036,6 +1036,7 @@ pub(crate) async fn handle_subagent_request(
|
|||||||
None,
|
None,
|
||||||
None,
|
None,
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
if verbatim_mirror_fork {
|
if verbatim_mirror_fork {
|
||||||
None
|
None
|
||||||
} else if let Some(scope) = agent_memory_scope {
|
} 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.app_builder_deployer_config.clone(),
|
||||||
ctx.write_file_enabled,
|
ctx.write_file_enabled,
|
||||||
ctx.goal_enabled,
|
ctx.goal_enabled,
|
||||||
|
// Graph mode is a depth-0 harness; child sessions never drive it.
|
||||||
|
false,
|
||||||
true,
|
true,
|
||||||
ctx.ask_user_question_enabled,
|
ctx.ask_user_question_enabled,
|
||||||
ctx.client_hooks.clone(),
|
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)
|
let auth_scheme = crate::agent::config::try_resolve_model_credentials(&cfg.model, None)
|
||||||
.map(|r| r.auth_scheme)
|
.map(|r| r.auth_scheme)
|
||||||
.unwrap_or_default();
|
.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 {
|
let inherited = kigi_sampler::SamplerConfig {
|
||||||
api_key: creds.api_key,
|
api_key: creds.api_key,
|
||||||
base_url: cfg.base_url,
|
base_url: cfg.base_url,
|
||||||
@@ -889,6 +907,10 @@ async fn read_parent_sampling_config(
|
|||||||
top_p: cfg.top_p,
|
top_p: cfg.top_p,
|
||||||
api_backend: cfg.api_backend,
|
api_backend: cfg.api_backend,
|
||||||
auth_scheme,
|
auth_scheme,
|
||||||
|
anthropic_oauth,
|
||||||
|
github_copilot,
|
||||||
|
openai_codex,
|
||||||
|
chat_compat: cfg.chat_compat,
|
||||||
extra_headers,
|
extra_headers,
|
||||||
context_window: cfg.context_window.get(),
|
context_window: cfg.context_window.get(),
|
||||||
reasoning_effort: cfg.reasoning_effort,
|
reasoning_effort: cfg.reasoning_effort,
|
||||||
@@ -981,9 +1003,21 @@ fn resolve_model_override_to_config(
|
|||||||
} else {
|
} else {
|
||||||
acp::ModelId::new(entry.info().model.clone())
|
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 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);
|
credentials.auth_type = subagent_auth_type(Some(&entry), &ctx.auth_method_id);
|
||||||
let resolved_auth_type = credentials.auth_type;
|
let resolved_auth_type = credentials.auth_type;
|
||||||
let config = sampling_config_for_model(&entry, credentials, ctx.alpha_test_key.clone());
|
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 unavailable exact-key entry selected by execution"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
/// Validation must inspect the SAME slug-collision entry execution selects.
|
||||||
|
/// Both go through `find_model_by_id`, whose slug scan takes the LAST match —
|
||||||
|
/// aligned with the picker's `resolve_catalog_key` so the auth layer and the
|
||||||
|
/// picker can never resolve different platforms for one slug (the H5 collision).
|
||||||
|
/// So a blocked LAST entry must be rejected even though an available earlier one
|
||||||
|
/// shares the slug.
|
||||||
#[test]
|
#[test]
|
||||||
fn fresh_tool_model_rejects_unavailable_first_slug_collision() {
|
fn fresh_tool_model_rejects_unavailable_last_slug_collision() {
|
||||||
let mut models = indexmap::IndexMap::new();
|
let mut models = indexmap::IndexMap::new();
|
||||||
let mut unavailable_first = test_model_entry("shared-routing-slug");
|
models.insert("visible-first".to_string(), test_model_entry("shared-routing-slug"));
|
||||||
unavailable_first.info.user_selectable = false;
|
let mut unavailable_last = test_model_entry("shared-routing-slug");
|
||||||
models.insert("blocked-first".to_string(), unavailable_first);
|
unavailable_last.info.user_selectable = false;
|
||||||
models.insert("visible-second".to_string(), test_model_entry("shared-routing-slug"));
|
models.insert("blocked-last".to_string(), unavailable_last);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
super::handle_request::task_model_override_error(Some("shared-routing-slug"),
|
super::handle_request::task_model_override_error(Some("shared-routing-slug"),
|
||||||
ModelOverrideProvenance::Tool, false, & models, false,).as_deref(),
|
ModelOverrideProvenance::Tool, false, & models, false,).as_deref(),
|
||||||
Some("Unknown Task.model slug 'shared-routing-slug'. Valid model slugs: \
|
Some("Unknown Task.model slug 'shared-routing-slug'. Valid model slugs: \
|
||||||
visible-second. Omit `model` to inherit the parent model."),
|
visible-first. Omit `model` to inherit the parent model."),
|
||||||
"validation must inspect the first routing-slug entry selected by execution"
|
"validation must inspect the last routing-slug entry selected by execution"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
#[test]
|
#[test]
|
||||||
@@ -3272,6 +3278,7 @@ fn test_sampling_config(model_slug: &str) -> kigi_sampling_types::SamplingConfig
|
|||||||
temperature: None,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: NonZeroU64::new(256_000).expect("non-zero context window"),
|
context_window: NonZeroU64::new(256_000).expect("non-zero context window"),
|
||||||
reasoning_effort: None,
|
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!(config.model, "kigi-4.5");
|
||||||
assert_eq!(model_id.0.as_ref(), "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
|
/// An unresolvable `AgentDefinition.model` pin (model absent from
|
||||||
/// `available_models`) falls through to inherit the parent model.
|
/// `available_models`) falls through to inherit the parent model.
|
||||||
#[tokio::test]
|
#[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 std::sync::Arc;
|
||||||
|
|
||||||
use crate::auth::kimi_oauth::{
|
use kigi_models::OAuthConfig;
|
||||||
DeviceAuthorization, DevicePollResult, poll_device_token, request_device_authorization,
|
|
||||||
};
|
use crate::auth::kimi_oauth::{DeviceAuthorization, DevicePollResult};
|
||||||
use crate::auth::{AuthChannels, AuthManager, AuthUrlInfo, AuthUrlMode, KimiAuth};
|
use crate::auth::{AuthChannels, AuthManager, AuthUrlInfo, AuthUrlMode, KimiAuth};
|
||||||
|
|
||||||
/// Extra wait added to the poll interval when the server answers `slow_down`
|
/// Extra wait added to the poll interval when the server answers `slow_down`
|
||||||
/// (OAuth-standard device-flow backpressure).
|
/// (OAuth-standard device-flow backpressure).
|
||||||
const SLOW_DOWN_INCREMENT_SECS: u64 = 5;
|
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.
|
/// Outcome of one full poll loop over a single device authorization.
|
||||||
enum PollLoopOutcome {
|
enum PollLoopOutcome {
|
||||||
/// Access token issued.
|
/// Access token issued.
|
||||||
@@ -40,11 +95,46 @@ pub async fn run_device_code_login_channels(
|
|||||||
host: &str,
|
host: &str,
|
||||||
auth_manager: &Arc<AuthManager>,
|
auth_manager: &Arc<AuthManager>,
|
||||||
channels: &mut Option<AuthChannels>,
|
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)> {
|
) -> anyhow::Result<(KimiAuth, bool)> {
|
||||||
let interactive_tui = channels.is_some();
|
let interactive_tui = channels.is_some();
|
||||||
let mut channels = channels.take();
|
let mut channels = channels.take();
|
||||||
loop {
|
loop {
|
||||||
let device_auth = request_device_authorization(host).await?;
|
let device_auth = backend.request().await?;
|
||||||
let display_uri = device_auth.verification_uri_complete.clone();
|
let display_uri = device_auth.verification_uri_complete.clone();
|
||||||
|
|
||||||
if interactive_tui {
|
if interactive_tui {
|
||||||
@@ -62,10 +152,14 @@ pub async fn run_device_code_login_channels(
|
|||||||
prompt_on_stderr(&device_auth).await;
|
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) => {
|
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
|
let auth = auth_manager
|
||||||
.update(*auth)
|
.update(auth)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| anyhow::anyhow!("Failed to save credentials: {e}"))?;
|
.map_err(|e| anyhow::anyhow!("Failed to save credentials: {e}"))?;
|
||||||
return Ok((auth, true));
|
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
|
/// Poll the token endpoint until the user approves, the device code expires
|
||||||
/// (→ [`PollLoopOutcome::Restart`]), or the wire fails.
|
/// (→ [`PollLoopOutcome::Restart`]), or the wire fails.
|
||||||
async fn complete_device_code_login(
|
async fn complete_device_code_login(
|
||||||
host: &str,
|
backend: &DeviceFlowBackend<'_>,
|
||||||
device_auth: &DeviceAuthorization,
|
device_auth: &DeviceAuthorization,
|
||||||
) -> anyhow::Result<PollLoopOutcome> {
|
) -> anyhow::Result<PollLoopOutcome> {
|
||||||
let mut poll_interval = std::time::Duration::from_secs(device_auth.interval.max(1) as u64);
|
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
|
// Sleep first: an immediate poll on a fresh code only returns
|
||||||
// authorization_pending (and risks slow_down).
|
// authorization_pending (and risks slow_down).
|
||||||
tokio::time::sleep(poll_interval).await;
|
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) => {
|
DevicePollResult::Success(auth) => {
|
||||||
tracing::info!("auth: device login authorized");
|
tracing::info!("auth: device login authorized");
|
||||||
return Ok(PollLoopOutcome::Done(auth));
|
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
|
/// 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
|
/// 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
|
/// caller can decide how to notify the user (eprintln on CLI, nothing on TUI
|
||||||
/// where the URL is already rendered in the widget).
|
/// where the URL is already rendered in the widget). Shared with the PKCE flow.
|
||||||
async fn open_browser_detached(url: &str) -> bool {
|
pub(super) async fn open_browser_detached(url: &str) -> bool {
|
||||||
// Unit tests drive the full login flow against mock servers — their
|
// Unit tests drive the full login flow against mock servers — their
|
||||||
// fixture URLs must never reach a real browser.
|
// fixture URLs must never reach a real browser.
|
||||||
if cfg!(test) {
|
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).
|
/// A 5xx from the token endpoint is a hard error (kimi-cli parity).
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn server_error_during_poll_fails_login() {
|
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
|
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(
|
async fn run_auth_flow_inner(
|
||||||
auth_manager: &Arc<AuthManager>,
|
auth_manager: &Arc<AuthManager>,
|
||||||
_kimi_code_config: &KimiCodeConfig,
|
_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
|
/// Defend against control characters / non-https redirects from a
|
||||||
/// compromised or mis-configured OAuth host.
|
/// compromised or mis-configured OAuth host. Shared with the generic
|
||||||
fn validate_verification_uri(uri: &str) -> anyhow::Result<()> {
|
/// 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()) {
|
if uri.chars().any(|c| c.is_ascii_control()) {
|
||||||
anyhow::bail!("Server returned invalid verification URI");
|
anyhow::bail!("Server returned invalid verification URI");
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -12,6 +12,7 @@ use tokio_util::sync::CancellationToken;
|
|||||||
|
|
||||||
#[path = "manager/lock.rs"]
|
#[path = "manager/lock.rs"]
|
||||||
mod lock;
|
mod lock;
|
||||||
|
pub(crate) use lock::try_lock_auth_file_nonblocking;
|
||||||
#[path = "manager/sleep_gate.rs"]
|
#[path = "manager/sleep_gate.rs"]
|
||||||
mod sleep_gate;
|
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
|
/// Single field-assembly point for [`Self::new`]'s two construction paths
|
||||||
/// (inline `KIGI_AUTH` vs. on-disk `auth.json`), which differ only in the
|
/// (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
|
/// threaded fields. One literal means a newly added field can't be silently
|
||||||
@@ -788,6 +842,13 @@ impl AuthManager {
|
|||||||
&self.kimi_code_config
|
&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.
|
/// Handle notified after every successful token refresh.
|
||||||
///
|
///
|
||||||
/// Used by [`ModelsManager`] to trigger model catalog recovery
|
/// Used by [`ModelsManager`] to trigger model catalog recovery
|
||||||
|
|||||||
@@ -1,13 +1,18 @@
|
|||||||
pub(crate) mod attribution;
|
pub(crate) mod attribution;
|
||||||
mod config;
|
mod config;
|
||||||
|
pub(crate) mod credential_authority;
|
||||||
pub mod credential_provider;
|
pub mod credential_provider;
|
||||||
pub(crate) mod device;
|
pub(crate) mod device;
|
||||||
pub mod device_code;
|
pub mod device_code;
|
||||||
pub mod error;
|
pub mod error;
|
||||||
mod flow;
|
mod flow;
|
||||||
|
pub(crate) mod github_copilot;
|
||||||
pub(crate) mod kimi_oauth;
|
pub(crate) mod kimi_oauth;
|
||||||
pub(crate) mod manager;
|
pub(crate) mod manager;
|
||||||
mod model;
|
mod model;
|
||||||
|
pub(crate) mod oauth_device;
|
||||||
|
pub(crate) mod oauth_pkce;
|
||||||
|
pub(crate) mod oauth_registry;
|
||||||
pub(crate) mod recovery;
|
pub(crate) mod recovery;
|
||||||
pub(crate) mod refresh;
|
pub(crate) mod refresh;
|
||||||
mod storage;
|
mod storage;
|
||||||
@@ -17,7 +22,8 @@ pub(crate) use flow::try_ensure_session_noninteractive;
|
|||||||
pub use flow::{
|
pub use flow::{
|
||||||
AuthChannels, AuthUrlInfo, AuthUrlMode, LogoutResult, ensure_authenticated,
|
AuthChannels, AuthUrlInfo, AuthUrlMode, LogoutResult, ensure_authenticated,
|
||||||
ensure_authenticated_or_noninteractive, perform_logout, run_auth_flow,
|
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;
|
mod meta;
|
||||||
pub use device::device_headers;
|
pub use device::device_headers;
|
||||||
@@ -27,5 +33,6 @@ pub use meta::AuthMeta;
|
|||||||
pub use model::{AuthMode, KimiAuth, lookup_auth};
|
pub use model::{AuthMode, KimiAuth, lookup_auth};
|
||||||
pub(crate) use model::{TOKEN_TTL, is_expired, token_suffix};
|
pub(crate) use model::{TOKEN_TTL, is_expired, token_suffix};
|
||||||
pub use storage::{
|
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;
|
mod kimi_refresher;
|
||||||
|
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
@@ -6,6 +7,7 @@ use crate::auth::manager::AuthManager;
|
|||||||
pub(crate) use crate::auth::manager::RefreshReason;
|
pub(crate) use crate::auth::manager::RefreshReason;
|
||||||
use crate::auth::model::KimiAuth;
|
use crate::auth::model::KimiAuth;
|
||||||
|
|
||||||
|
pub(crate) use generic_refresher::GenericDeviceRefresher;
|
||||||
pub(crate) use kimi_refresher::KimiRefresher;
|
pub(crate) use kimi_refresher::KimiRefresher;
|
||||||
|
|
||||||
/// Read-only view of `AuthManager` for refreshers. Enforces the
|
/// 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;
|
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> {
|
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;
|
let snapshot: Arc<dyn AuthSnapshot> = auth_manager;
|
||||||
Arc::new(KimiRefresher::new(snapshot, kigi_env::oauth_host()))
|
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> {
|
fn keyring_entry() -> Result<&'static keyring::Entry, keyring::Error> {
|
||||||
static ENTRY: std::sync::OnceLock<Result<keyring::Entry, keyring::Error>> =
|
static ENTRY: std::sync::OnceLock<Result<keyring::Entry, keyring::Error>> =
|
||||||
std::sync::OnceLock::new();
|
std::sync::OnceLock::new();
|
||||||
match ENTRY.get_or_init(|| {
|
match ENTRY
|
||||||
keyring::Entry::new(KEYRING_SERVICE, crate::auth::config::KIMI_CODE_OAUTH_SCOPE)
|
.get_or_init(|| keyring::Entry::new(KEYRING_SERVICE, crate::auth::KIMI_CODE_OAUTH_SCOPE))
|
||||||
}) {
|
{
|
||||||
Ok(entry) => Ok(entry),
|
Ok(entry) => Ok(entry),
|
||||||
// `keyring::Error` is not `Clone`; surface a stable equivalent.
|
// `keyring::Error` is not `Clone`; surface a stable equivalent.
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
@@ -536,6 +536,71 @@ pub fn store_api_key(kigi_home: &Path, api_key: &str) -> std::io::Result<()> {
|
|||||||
write_auth_json(&path, &map)
|
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.
|
/// Remove the `kigi::api_key` scope from auth.json.
|
||||||
pub fn clear_api_key(kigi_home: &Path) -> std::io::Result<()> {
|
pub fn clear_api_key(kigi_home: &Path) -> std::io::Result<()> {
|
||||||
let path = kigi_home.join("auth.json");
|
let path = kigi_home.join("auth.json");
|
||||||
@@ -550,6 +615,80 @@ pub fn clear_api_key(kigi_home: &Path) -> std::io::Result<()> {
|
|||||||
Ok(())
|
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)]
|
#[cfg(test)]
|
||||||
mod write_fallback_tests {
|
mod write_fallback_tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|||||||
@@ -116,8 +116,8 @@ pub fn extract_bundled_files(kigi_home: &std::path::Path) {
|
|||||||
|
|
||||||
let _ = std::fs::create_dir_all(kigi_home);
|
let _ = std::fs::create_dir_all(kigi_home);
|
||||||
|
|
||||||
// Clean up cached changelog files from previous version so
|
// Clean up changelog caches written by the removed changelog feature
|
||||||
// /release-notes fetches fresh content for the new version.
|
// (kigi <= 0.1.0 cached CDN release notes in the kigi home).
|
||||||
for stale in &["CHANGELOG.json", "CHANGELOG.md"] {
|
for stale in &["CHANGELOG.json", "CHANGELOG.md"] {
|
||||||
let _ = std::fs::remove_file(kigi_home.join(stale));
|
let _ = std::fs::remove_file(kigi_home.join(stale));
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -35,7 +35,11 @@ impl AuthStatus {
|
|||||||
.unwrap_or(&origin);
|
.unwrap_or(&origin);
|
||||||
return Self::LoggedIn(host.to_owned());
|
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())
|
if crate::agent::auth_method::should_advertise_xai_api_key(models.values())
|
||||||
&& let Some(name) = models
|
&& let Some(name) = models
|
||||||
.iter()
|
.iter()
|
||||||
@@ -88,7 +92,9 @@ pub async fn list_models(
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
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::agent::config::Config;
|
||||||
use crate::auth::{AuthMode, KimiAuth};
|
use crate::auth::{AuthMode, KimiAuth};
|
||||||
use kigi_test_support::EnvGuard;
|
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
|
/// Uses `KIGI_AUTH_PATH` (not `KIGI_SHARE_DIR`) so a OnceLock-cached real home
|
||||||
/// with `auth.json` cannot leak into these tests.
|
/// 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 dir = tempfile::tempdir().unwrap();
|
||||||
let auth_path = dir.path().join("no-auth.json");
|
let auth_path = dir.path().join("no-auth.json");
|
||||||
let guards = [
|
let guards = [
|
||||||
|
EnvGuard::unset(HOUSE_API_KEY_ENV_VAR),
|
||||||
EnvGuard::unset(XAI_API_KEY_ENV_VAR),
|
EnvGuard::unset(XAI_API_KEY_ENV_VAR),
|
||||||
EnvGuard::unset(LEGACY_XAI_API_KEY_ENV_VAR),
|
EnvGuard::unset(LEGACY_XAI_API_KEY_ENV_VAR),
|
||||||
EnvGuard::unset("KIGI_AUTH"),
|
EnvGuard::unset("KIGI_AUTH"),
|
||||||
|
|||||||
@@ -56,18 +56,18 @@ fn handle_set_api_key(args: &acp::ExtRequest) -> ExtResult {
|
|||||||
crate::auth::clear_api_key(&kigi_home)
|
crate::auth::clear_api_key(&kigi_home)
|
||||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||||
// SAFETY: ext_method is single-threaded per agent
|
// 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 {
|
} else {
|
||||||
crate::auth::store_api_key(&kigi_home, k)
|
crate::auth::store_api_key(&kigi_home, k)
|
||||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||||
// SAFETY: ext_method is single-threaded per agent
|
// 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 {
|
} else {
|
||||||
crate::auth::clear_api_key(&kigi_home)
|
crate::auth::clear_api_key(&kigi_home)
|
||||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||||
// SAFETY: ext_method is single-threaded per agent
|
// 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 }))
|
ExtMethodResult::success(serde_json::json!({ "ok": true }))
|
||||||
.to_ext_response()
|
.to_ext_response()
|
||||||
|
|||||||
@@ -897,6 +897,36 @@ pub enum SessionUpdate {
|
|||||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
usage: Option<PromptUsage>,
|
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.
|
/// Catch-all for unrecognized session update types.
|
||||||
/// Allows forward/backward compatibility when variants are added or removed.
|
/// Allows forward/backward compatibility when variants are added or removed.
|
||||||
/// All fields from the unrecognized variant are discarded during deserialization.
|
/// 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());
|
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
|
/// Built-in kigi platform a kimi provider duplicates, if any: provider type
|
||||||
/// `kimi` is the Kimi Code subscription channel; the two Moonshot open
|
/// `kimi` is the Kimi Code subscription channel; API-key platforms are
|
||||||
/// platforms are recognized by their fixed production hosts (the same hosts
|
/// recognized by their production hosts (the same hosts
|
||||||
/// `kigi_models::PlatformId::base_url` compiles in).
|
/// `kigi_models::PlatformId::base_url` compiles in — moonshot, openai, and
|
||||||
|
/// every future registry row automatically).
|
||||||
fn builtin_platform(provider: &KimiProviderToml) -> Option<PlatformId> {
|
fn builtin_platform(provider: &KimiProviderToml) -> Option<PlatformId> {
|
||||||
if provider.provider_type == "kimi" {
|
if provider.provider_type == "kimi" {
|
||||||
return Some(PlatformId::KimiCode);
|
return Some(PlatformId::KimiCode);
|
||||||
}
|
}
|
||||||
match url_host(&provider.base_url) {
|
let host = url_host(&provider.base_url)?;
|
||||||
Some("api.moonshot.cn") => Some(PlatformId::MoonshotCn),
|
PlatformId::ALL.into_iter().find(|platform| {
|
||||||
Some("api.moonshot.ai") => Some(PlatformId::MoonshotAi),
|
if platform.uses_oauth() {
|
||||||
_ => None,
|
return false;
|
||||||
}
|
}
|
||||||
|
let base = platform.base_url();
|
||||||
|
url_host(&base) == Some(host)
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Host component of an http(s) URL. `None` for other schemes.
|
/// Host component of an http(s) URL. `None` for other schemes.
|
||||||
@@ -607,7 +611,7 @@ model = "kimi-for-coding"
|
|||||||
max_context_size = 262144
|
max_context_size = 262144
|
||||||
|
|
||||||
[models.my-openai]
|
[models.my-openai]
|
||||||
provider = "openrouter"
|
provider = "customllm"
|
||||||
model = "gpt-x"
|
model = "gpt-x"
|
||||||
max_context_size = 128000
|
max_context_size = 128000
|
||||||
|
|
||||||
@@ -616,9 +620,9 @@ type = "kimi"
|
|||||||
base_url = "https://api.kimi.com/coding/v1"
|
base_url = "https://api.kimi.com/coding/v1"
|
||||||
api_key = "sk-kimi-secret"
|
api_key = "sk-kimi-secret"
|
||||||
|
|
||||||
[providers.openrouter]
|
[providers.customllm]
|
||||||
type = "openai_legacy"
|
type = "openai_legacy"
|
||||||
base_url = "https://openrouter.ai/api/v1"
|
base_url = "https://llm.example.test/v1"
|
||||||
api_key = "sk-or-secret"
|
api_key = "sk-or-secret"
|
||||||
"#,
|
"#,
|
||||||
)
|
)
|
||||||
@@ -707,7 +711,7 @@ api_key = "sk-or-secret"
|
|||||||
let m = &plan.custom_models[0];
|
let m = &plan.custom_models[0];
|
||||||
assert_eq!(m.alias, "my-openai");
|
assert_eq!(m.alias, "my-openai");
|
||||||
assert_eq!(m.model, "gpt-x");
|
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.api_key.as_deref(), Some("sk-or-secret"));
|
||||||
assert_eq!(m.context_window, Some(128_000));
|
assert_eq!(m.context_window, Some(128_000));
|
||||||
|
|
||||||
@@ -846,7 +850,7 @@ api_key = "sk-ms"
|
|||||||
assert_eq!(model["model"].as_str().unwrap(), "gpt-x");
|
assert_eq!(model["model"].as_str().unwrap(), "gpt-x");
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
model["base_url"].as_str().unwrap(),
|
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["api_key"].as_str().unwrap(), "sk-or-secret");
|
||||||
assert_eq!(model["context_window"].as_integer().unwrap(), 128_000);
|
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()
|
&& crate::agent::auth_method::has_xai_api_key_env()
|
||||||
{
|
{
|
||||||
format!(
|
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. \
|
Your cached OAuth session is being used instead. \
|
||||||
To use your API key, run `kigi logout` or type /logout in the TUI."
|
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).
|
/// Helper: run a closure with the house BYOK key (KIGI_API_KEY) temporarily
|
||||||
/// Cleans up even if the closure panics.
|
/// 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) {
|
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 = std::env::var("XAI_API_KEY").ok();
|
||||||
let prev_legacy = std::env::var("KIGI_CODE_XAI_API_KEY").ok();
|
let prev_legacy = std::env::var("KIGI_CODE_XAI_API_KEY").ok();
|
||||||
// SAFETY: serial_test ensures no concurrent env mutation.
|
// SAFETY: serial_test ensures no concurrent env mutation.
|
||||||
unsafe {
|
unsafe {
|
||||||
|
std::env::remove_var("KIGI_API_KEY");
|
||||||
std::env::remove_var("XAI_API_KEY");
|
std::env::remove_var("XAI_API_KEY");
|
||||||
std::env::remove_var("KIGI_CODE_XAI_API_KEY");
|
std::env::remove_var("KIGI_CODE_XAI_API_KEY");
|
||||||
if let Some(k) = 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));
|
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
|
||||||
// Restore original state.
|
// Restore original state.
|
||||||
unsafe {
|
unsafe {
|
||||||
|
std::env::remove_var("KIGI_API_KEY");
|
||||||
std::env::remove_var("XAI_API_KEY");
|
std::env::remove_var("XAI_API_KEY");
|
||||||
std::env::remove_var("KIGI_CODE_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 {
|
if let Some(v) = prev {
|
||||||
std::env::set_var("XAI_API_KEY", v);
|
std::env::set_var("XAI_API_KEY", v);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -93,6 +93,14 @@ pub(crate) use types::*;
|
|||||||
pub use types::{TodoGateDecision, TodoGateReason};
|
pub use types::{TodoGateDecision, TodoGateReason};
|
||||||
#[path = "acp_session_impl/goal.rs"]
|
#[path = "acp_session_impl/goal.rs"]
|
||||||
mod goal;
|
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"]
|
#[path = "acp_session_impl/interjection.rs"]
|
||||||
mod interjection;
|
mod interjection;
|
||||||
#[path = "acp_session_impl/tool_calls.rs"]
|
#[path = "acp_session_impl/tool_calls.rs"]
|
||||||
@@ -127,7 +135,7 @@ use prompt_build::*;
|
|||||||
mod session_mode;
|
mod session_mode;
|
||||||
use session_mode::*;
|
use session_mode::*;
|
||||||
#[path = "acp_session_impl/sampler_turn.rs"]
|
#[path = "acp_session_impl/sampler_turn.rs"]
|
||||||
mod sampler_turn;
|
pub(crate) mod sampler_turn;
|
||||||
use sampler_turn::*;
|
use sampler_turn::*;
|
||||||
#[path = "acp_session_impl/tool_dispatch.rs"]
|
#[path = "acp_session_impl/tool_dispatch.rs"]
|
||||||
mod tool_dispatch;
|
mod tool_dispatch;
|
||||||
@@ -409,6 +417,21 @@ pub(crate) struct SessionActor {
|
|||||||
/// [`SessionActor::model_auth_facts`].
|
/// [`SessionActor::model_auth_facts`].
|
||||||
pub(crate) model_auth_facts:
|
pub(crate) model_auth_facts:
|
||||||
std::cell::RefCell<Option<(String, crate::agent::config::ModelAuthFacts)>>,
|
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
|
/// 401-attribution callback. Joined with the bearer the
|
||||||
/// sampler sends on the wire to emit an `auth 401 attribution`
|
/// sampler sends on the wire to emit an `auth 401 attribution`
|
||||||
/// event at each of the six `OaiCompatClient` 401 arms in
|
/// 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
|
/// Goal mode orchestration tracker. Session-scoped state for the
|
||||||
/// Design-Execute-Verify loop. Modeled after `plan_mode` above.
|
/// Design-Execute-Verify loop. Modeled after `plan_mode` above.
|
||||||
pub(crate) goal_tracker: Arc<parking_lot::Mutex<crate::session::goal_tracker::GoalTracker>>,
|
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
|
/// `task_id`s of background tasks (and monitors) that originated during
|
||||||
/// the goal turn — either spawned by the goal model itself or reparented
|
/// the goal turn — either spawned by the goal model itself or reparented
|
||||||
/// from a harness verifier/planner subagent on its exit. Their late
|
/// from a harness verifier/planner subagent on its exit. Their late
|
||||||
@@ -985,6 +1036,9 @@ impl SessionActor {
|
|||||||
hooks: self.hook_registry.borrow().is_some(),
|
hooks: self.hook_registry.borrow().is_some(),
|
||||||
plugins: self.plugin_registry.borrow().is_some(),
|
plugins: self.plugin_registry.borrow().is_some(),
|
||||||
goal,
|
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.
|
/// Names of every tool registered with the session's tool bridge.
|
||||||
@@ -1217,6 +1271,17 @@ mod rewind_synthetic_turn_tests;
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
#[path = "acp_session_tests/rewrite_zero_turn_prefix_tests.rs"]
|
#[path = "acp_session_tests/rewrite_zero_turn_prefix_tests.rs"]
|
||||||
mod rewrite_zero_turn_prefix_tests;
|
mod rewrite_zero_turn_prefix_tests;
|
||||||
|
/// The same guard for the model→platform lookup (the dual-credential slug
|
||||||
|
/// collision) and for the stamped aux/summary configs.
|
||||||
|
#[cfg(test)]
|
||||||
|
#[path = "acp_session_tests/session_bearer_leak_platform_tests.rs"]
|
||||||
|
mod session_bearer_leak_platform_tests;
|
||||||
|
/// LEAK guard: the primary Kimi subscription bearer must never ride a request
|
||||||
|
/// to an API-key registry platform's host, while the subscription-OAuth
|
||||||
|
/// platforms keep a live resolver from their OWN pooled manager.
|
||||||
|
#[cfg(test)]
|
||||||
|
#[path = "acp_session_tests/session_bearer_leak_tests.rs"]
|
||||||
|
mod session_bearer_leak_tests;
|
||||||
/// Pins the `SubagentFinished` usage-fold attribution gate.
|
/// Pins the `SubagentFinished` usage-fold attribution gate.
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
#[path = "acp_session_tests/subagent_usage_fold_tests.rs"]
|
#[path = "acp_session_tests/subagent_usage_fold_tests.rs"]
|
||||||
@@ -1423,6 +1488,9 @@ mod goal_strategist_e2e_tests;
|
|||||||
#[path = "acp_session_tests/goal/goal_summarizer_e2e_tests.rs"]
|
#[path = "acp_session_tests/goal/goal_summarizer_e2e_tests.rs"]
|
||||||
mod goal_summarizer_e2e_tests;
|
mod goal_summarizer_e2e_tests;
|
||||||
#[cfg(test)]
|
#[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"]
|
#[path = "acp_session_tests/idle_resume_tests.rs"]
|
||||||
mod idle_resume_tests;
|
mod idle_resume_tests;
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
|
|||||||
@@ -1904,6 +1904,29 @@ impl SessionActor {
|
|||||||
tokens_used,
|
tokens_used,
|
||||||
finished_marginal,
|
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!(
|
self.send_slash_command_output(&format!(
|
||||||
"Goal token budget reached ({tokens_used} of {budget} tokens) — goal \
|
"Goal token budget reached ({tokens_used} of {budget} tokens) — goal \
|
||||||
stopped. Use /goal clear, then /goal <objective> to start a new one."
|
stopped. Use /goal clear, then /goal <objective> to start a new one."
|
||||||
@@ -2265,18 +2288,21 @@ impl SessionActor {
|
|||||||
message: Option<String>,
|
message: Option<String>,
|
||||||
) -> bool {
|
) -> bool {
|
||||||
let current_tokens = self.chat_state_handle.get_total_tokens().await as i64;
|
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();
|
let mut tracker = self.goal_tracker.lock();
|
||||||
if tracker.status() != Some(crate::session::goal_tracker::GoalStatus::Active) {
|
if tracker.status() == Some(crate::session::goal_tracker::GoalStatus::Active) {
|
||||||
return false;
|
// Active is guaranteed here, so the transition succeeds.
|
||||||
}
|
|
||||||
// The early-return above guarantees `Active`, so the pause
|
|
||||||
// transition always succeeds here.
|
|
||||||
match message {
|
match message {
|
||||||
Some(msg) => tracker.pause_with_message(reason, msg),
|
Some(msg) => tracker.pause_with_message(reason, msg),
|
||||||
None => tracker.pause(reason),
|
None => tracker.pause(reason),
|
||||||
};
|
};
|
||||||
|
true
|
||||||
|
} else {
|
||||||
|
false
|
||||||
}
|
}
|
||||||
|
};
|
||||||
|
if goal_paused {
|
||||||
self.clear_pending_classifier_completions();
|
self.clear_pending_classifier_completions();
|
||||||
let (tokens_used, finished_marginal) = self.goal_tokens(current_tokens);
|
let (tokens_used, finished_marginal) = self.goal_tokens(current_tokens);
|
||||||
let notify = self.goal_notify_sender();
|
let notify = self.goal_notify_sender();
|
||||||
@@ -2288,7 +2314,31 @@ impl SessionActor {
|
|||||||
self.emit_event(crate::session::events::Event::GoalAutoPaused {
|
self.emit_event(crate::session::events::Event::GoalAutoPaused {
|
||||||
reason: reason.into(),
|
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
|
/// 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(
|
pub(super) async fn handle_set_session_model(
|
||||||
&self,
|
&self,
|
||||||
sampling_config: kigi_sampler::SamplerConfig,
|
sampling_config: kigi_sampler::SamplerConfig,
|
||||||
|
catalog_key: Option<String>,
|
||||||
use_concise: bool,
|
use_concise: bool,
|
||||||
apply_prompt_override: bool,
|
apply_prompt_override: bool,
|
||||||
skip_prompt_rewrite: bool,
|
skip_prompt_rewrite: bool,
|
||||||
auto_compact_threshold_percent: u8,
|
auto_compact_threshold_percent: u8,
|
||||||
) -> Result<acp::ModelId, acp::Error> {
|
) -> Result<acp::ModelId, acp::Error> {
|
||||||
let model_id = acp::ModelId::new(sampling_config.model.clone());
|
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(|| {
|
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(sampling_config.context_window).unwrap_or_else(|| {
|
||||||
std::num::NonZeroU64::new(DEFAULT_CONTEXT_WINDOW)
|
std::num::NonZeroU64::new(DEFAULT_CONTEXT_WINDOW)
|
||||||
@@ -53,22 +59,24 @@ impl SessionActor {
|
|||||||
temperature: sampling_config.temperature,
|
temperature: sampling_config.temperature,
|
||||||
top_p: sampling_config.top_p,
|
top_p: sampling_config.top_p,
|
||||||
api_backend: sampling_config.api_backend.clone(),
|
api_backend: sampling_config.api_backend.clone(),
|
||||||
|
chat_compat: sampling_config.chat_compat,
|
||||||
extra_headers: sampling_config.extra_headers.clone(),
|
extra_headers: sampling_config.extra_headers.clone(),
|
||||||
context_window: new_context_window,
|
context_window: new_context_window,
|
||||||
reasoning_effort: sampling_config.reasoning_effort,
|
reasoning_effort: sampling_config.reasoning_effort,
|
||||||
stream_tool_calls: Some(sampling_config.stream_tool_calls),
|
stream_tool_calls: Some(sampling_config.stream_tool_calls),
|
||||||
});
|
});
|
||||||
let existing = self.chat_state_handle.get_credentials().await;
|
let existing = self.chat_state_handle.get_credentials().await;
|
||||||
let session_key = self
|
// Read the session bearer from the switched-to model's OWN manager: a
|
||||||
.auth_manager
|
// grok model reads the xai-grok token (used only to classify the
|
||||||
.as_ref()
|
// credential's auth_type here), never the Kimi one. Kimi / non-oauth
|
||||||
.and_then(|am| am.current_or_expired().map(|a| a.key));
|
// models resolve to the primary — byte-identical.
|
||||||
|
let session_key = self.session_credential_for_model(&sampling_config.model);
|
||||||
self.chat_state_handle
|
self.chat_state_handle
|
||||||
.update_credentials(kigi_chat_state::Credentials {
|
.update_credentials(kigi_chat_state::Credentials {
|
||||||
api_key: sampling_config.api_key.clone(),
|
api_key: sampling_config.api_key.clone(),
|
||||||
auth_type: crate::agent::config::resolve_chat_state_auth_type(
|
auth_type: crate::agent::config::resolve_chat_state_auth_type(
|
||||||
sampling_config.model.as_str(),
|
sampling_config.model.as_str(),
|
||||||
session_key.as_deref(),
|
session_key.as_ref(),
|
||||||
existing.auth_type,
|
existing.auth_type,
|
||||||
),
|
),
|
||||||
alpha_test_key: existing.alpha_test_key,
|
alpha_test_key: existing.alpha_test_key,
|
||||||
|
|||||||
@@ -624,10 +624,20 @@ impl SessionActor {
|
|||||||
let resolved_describe = self
|
let resolved_describe = self
|
||||||
.resolve_aux_sampler_config(&self.image_description_model)
|
.resolve_aux_sampler_config(&self.image_description_model)
|
||||||
.await;
|
.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) =
|
let (describe_model, sampler_config) =
|
||||||
crate::agent::config::finalize_image_describe_sampler_config(
|
crate::agent::config::finalize_image_describe_sampler_config(
|
||||||
resolved_describe,
|
resolved_describe,
|
||||||
&active_session_config,
|
&active_session_config,
|
||||||
|
describe_resolver.flatten(),
|
||||||
Some(self.max_retries),
|
Some(self.max_retries),
|
||||||
);
|
);
|
||||||
let client = kigi_sampler::SamplingClient::new(sampler_config).map_err(|e| {
|
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
|
SessionActor::maybe_start_running_task(session.clone(), completion_tx
|
||||||
.clone()). await; } SessionCommand::SessionMode { session_mode, responds_to }
|
.clone()). await; } SessionCommand::SessionMode { session_mode, responds_to }
|
||||||
=> { session.handle_session_mode(session_mode). await; let _ = responds_to
|
=> { session.handle_session_mode(session_mode). await; let _ = responds_to
|
||||||
.send(()); } SessionCommand::SetSessionModel { sampling_config, use_concise,
|
.send(()); } SessionCommand::SetSessionModel { sampling_config, catalog_key,
|
||||||
apply_prompt_override, skip_prompt_rewrite, auto_compact_threshold_percent,
|
use_concise, apply_prompt_override, skip_prompt_rewrite,
|
||||||
responds_to } => { let updated_model_id = session
|
auto_compact_threshold_percent, responds_to } => { let updated_model_id =
|
||||||
.handle_set_session_model(sampling_config, use_concise,
|
session.handle_set_session_model(sampling_config, catalog_key, use_concise,
|
||||||
apply_prompt_override, skip_prompt_rewrite, auto_compact_threshold_percent).
|
apply_prompt_override, skip_prompt_rewrite, auto_compact_threshold_percent).
|
||||||
await; let _ = responds_to.send(updated_model_id); }
|
await; let _ = responds_to.send(updated_model_id); }
|
||||||
SessionCommand::RebuildAgentForDefinition { definition, responds_to } => {
|
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
|
.signals_handle().set_primary_model(& model_name); cfg.model = model_name
|
||||||
.clone(); cfg.extra_headers.extend(extra_headers); if let Some(cw) =
|
.clone(); cfg.extra_headers.extend(extra_headers); if let Some(cw) =
|
||||||
context_window && session.compaction.context_window_override.is_none() { cfg
|
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
|
.update_sampling_config(cfg); let existing = session.chat_state_handle
|
||||||
.get_credentials(). await; if let Some(r) = crate
|
.get_credentials(). await;
|
||||||
::agent::config::try_resolve_model_credentials(model_name.as_str(), existing
|
// H-c: the rename makes the session's own selected catalog
|
||||||
.api_key.as_deref()) { session.chat_state_handle
|
// 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,
|
.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
|
auth_type : r.auth_type, alpha_test_key : existing.alpha_test_key, }); } session.model_auth_facts
|
||||||
.replace(None); } } SessionCommand::GetCurrentModel { responds_to } => { let
|
.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
|
/// Gate inputs bundled with the composed decision so the 401-recovery log can
|
||||||
/// report the components.
|
/// report the components.
|
||||||
#[derive(Clone, Copy)]
|
#[derive(Clone, Copy)]
|
||||||
struct SessionTokenAuthGate {
|
pub(crate) struct SessionTokenAuthGate {
|
||||||
is_session_based: bool,
|
is_session_based: bool,
|
||||||
model_byok: crate::agent::auth_method::ModelByok,
|
model_byok: crate::agent::auth_method::ModelByok,
|
||||||
/// Whether the request targets a first-party host. Lets an `Unknown`
|
/// Whether the request targets a first-party host. Lets an `Unknown`
|
||||||
/// BYOK status still refresh against the first-party cli-chat-proxy hosts without
|
/// BYOK status still refresh against the first-party cli-chat-proxy hosts without
|
||||||
/// risking a session-token leak to a third-party BYOK endpoint.
|
/// risking a session-token leak to a third-party BYOK endpoint.
|
||||||
endpoint_is_first_party: bool,
|
endpoint_is_first_party: bool,
|
||||||
|
/// 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 {
|
impl SessionTokenAuthGate {
|
||||||
/// Single place `is_session_based` / `endpoint_is_first_party` are derived,
|
/// Single place `is_session_based` / `endpoint_is_first_party` are derived,
|
||||||
/// so all call sites assemble the gate identically.
|
/// so all call sites assemble the gate identically. `model_platform` is the
|
||||||
fn new(
|
/// 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>,
|
auth_method_id: Option<&acp::AuthMethodId>,
|
||||||
model_byok: crate::agent::auth_method::ModelByok,
|
model_byok: crate::agent::auth_method::ModelByok,
|
||||||
base_url: &str,
|
base_url: &str,
|
||||||
|
model_platform: Option<kigi_models::PlatformId>,
|
||||||
|
authority: &crate::auth::credential_authority::CredentialAuthority,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
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_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,
|
model_byok,
|
||||||
endpoint_is_first_party: crate::util::is_first_party_url(base_url),
|
endpoint_is_first_party: crate::util::is_first_party_url(base_url),
|
||||||
|
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(
|
crate::agent::auth_method::session_token_auth_gate(
|
||||||
self.is_session_based,
|
self.is_session_based,
|
||||||
self.model_byok,
|
self.model_byok,
|
||||||
self.endpoint_is_first_party,
|
self.endpoint_is_first_party,
|
||||||
|
self.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
|
/// Run a tool call; on an auth-shaped failure, attempt recovery via
|
||||||
/// `AuthManager` and one retry. When `shared_recovery` is `Some`, concurrent
|
/// `AuthManager` and one retry. When `shared_recovery` is `Some`, concurrent
|
||||||
/// 401s in the same batch deduplicate via `OnceCell::get_or_init`.
|
/// 401s in the same batch deduplicate via `OnceCell::get_or_init`.
|
||||||
@@ -103,6 +171,31 @@ where
|
|||||||
result
|
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 {
|
impl SessionActor {
|
||||||
pub(super) async fn prepare_tool_definitions_timed(&self) -> (Vec<ToolDefinition>, u64) {
|
pub(super) async fn prepare_tool_definitions_timed(&self) -> (Vec<ToolDefinition>, u64) {
|
||||||
let mcp_wait_start = std::time::Instant::now();
|
let mcp_wait_start = std::time::Instant::now();
|
||||||
@@ -146,8 +239,8 @@ impl SessionActor {
|
|||||||
let plan_active = self.plan_mode.lock().is_active();
|
let plan_active = self.plan_mode.lock().is_active();
|
||||||
filter_cursor_tools_by_plan_mode(defs, plan_active)
|
filter_cursor_tools_by_plan_mode(defs, plan_active)
|
||||||
}
|
}
|
||||||
/// Memoized per-model [`ModelAuthFacts`](crate::agent::config::ModelAuthFacts),
|
/// Memoized per-model [`ModelAuthFacts`](crate::agent::config::ModelAuthFacts)
|
||||||
/// keyed by `model_id`.
|
/// for the SESSION's own model, keyed by `model_id`.
|
||||||
///
|
///
|
||||||
/// A fresh `Unknown` (config currently unparseable) falls back to the last
|
/// A fresh `Unknown` (config currently unparseable) falls back to the last
|
||||||
/// definite value for the same `model_id` rather than demoting a live session
|
/// 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_id`, keying on `model_id` alone is insufficient — each
|
||||||
/// model/credential chokepoint must clear this memo (`replace(None)`).
|
/// model/credential chokepoint must clear this memo (`replace(None)`).
|
||||||
pub(super) fn model_auth_facts(&self, model_id: &str) -> crate::agent::config::ModelAuthFacts {
|
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;
|
use crate::agent::auth_method::ModelByok;
|
||||||
if let Some((cached_id, facts)) = self.model_auth_facts.borrow().as_ref()
|
if let Some((cached_id, facts)) = self.model_auth_facts.borrow().as_ref()
|
||||||
&& cached_id == model_id
|
&& cached_id == model_id
|
||||||
@@ -172,17 +291,187 @@ impl SessionActor {
|
|||||||
}
|
}
|
||||||
return fresh;
|
return fresh;
|
||||||
}
|
}
|
||||||
|
if memoize {
|
||||||
*self.model_auth_facts.borrow_mut() = Some((model_id.to_string(), fresh));
|
*self.model_auth_facts.borrow_mut() = Some((model_id.to_string(), fresh));
|
||||||
|
}
|
||||||
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
|
/// [`crate::agent::auth_method::session_token_auth_gate`] for the rationale
|
||||||
/// (`base_url` keeps an `Unknown` BYOK status refreshable only
|
/// (`base_url` keeps an `Unknown` BYOK status refreshable only
|
||||||
/// against first-party xAI hosts).
|
/// against first-party xAI hosts).
|
||||||
fn auth_gate(&self, model_id: &str, base_url: &str) -> SessionTokenAuthGate {
|
fn auth_gate(&self, model_id: &str, base_url: &str) -> SessionTokenAuthGate {
|
||||||
let byok = self.model_auth_facts(model_id).byok;
|
let byok = self.model_auth_facts(model_id).byok;
|
||||||
let auth_method = self.auth_method_id.load();
|
let auth_method = self.auth_method_id.load();
|
||||||
SessionTokenAuthGate::new(auth_method.as_deref(), byok, base_url)
|
SessionTokenAuthGate::new(
|
||||||
|
auth_method.as_deref(),
|
||||||
|
byok,
|
||||||
|
base_url,
|
||||||
|
self.model_platform(model_id),
|
||||||
|
&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
|
/// Emit a unified-log breadcrumb whenever the session-token refresh gate is
|
||||||
/// evaluated with an **`Unknown`** per-model BYOK status on a session-based
|
/// evaluated with an **`Unknown`** per-model BYOK status on a session-based
|
||||||
@@ -202,8 +491,8 @@ impl SessionActor {
|
|||||||
let ctx = serde_json::json!(
|
let ctx = serde_json::json!(
|
||||||
{ "site" : site, "model_byok" : gate.model_byok.as_str(), "is_session_based"
|
{ "site" : site, "model_byok" : gate.model_byok.as_str(), "is_session_based"
|
||||||
: gate.is_session_based, "endpoint_is_first_party" : gate
|
: gate.is_session_based, "endpoint_is_first_party" : gate
|
||||||
.endpoint_is_first_party, "refresh_active" : refresh_active, "base_url" :
|
.endpoint_is_first_party, "credential_class" : gate.credential_class
|
||||||
base_url, }
|
.as_str(), "refresh_active" : refresh_active, "base_url" : base_url, }
|
||||||
);
|
);
|
||||||
let sid = Some(self.session_info.id.0.as_ref());
|
let sid = Some(self.session_info.id.0.as_ref());
|
||||||
if refresh_active {
|
if refresh_active {
|
||||||
@@ -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
|
let cfg = self
|
||||||
.chat_state_handle
|
.chat_state_handle
|
||||||
.get_sampling_config()
|
.get_sampling_config()
|
||||||
@@ -260,6 +537,7 @@ impl SessionActor {
|
|||||||
temperature: None,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: std::num::NonZeroU64::new(256_000).unwrap(),
|
context_window: std::num::NonZeroU64::new(256_000).unwrap(),
|
||||||
reasoning_effort: None,
|
reasoning_effort: None,
|
||||||
@@ -268,10 +546,32 @@ impl SessionActor {
|
|||||||
let creds = self.chat_state_handle.get_credentials().await;
|
let creds = self.chat_state_handle.get_credentials().await;
|
||||||
let model_facts = self.model_auth_facts(cfg.model.as_str());
|
let model_facts = self.model_auth_facts(cfg.model.as_str());
|
||||||
let auth_method = self.auth_method_id.load();
|
let auth_method = self.auth_method_id.load();
|
||||||
let gate =
|
let gate = SessionTokenAuthGate::new(
|
||||||
SessionTokenAuthGate::new(auth_method.as_deref(), model_facts.byok, &cfg.base_url);
|
auth_method.as_deref(),
|
||||||
|
model_facts.byok,
|
||||||
|
&cfg.base_url,
|
||||||
|
self.model_platform(cfg.model.as_str()),
|
||||||
|
&self.credential_authority(),
|
||||||
|
);
|
||||||
let use_bearer_resolver = gate.active();
|
let use_bearer_resolver = gate.active();
|
||||||
self.log_auth_gate_unknown("reconstruct_full_config", gate, &cfg.base_url);
|
self.log_auth_gate_unknown("reconstruct_full_config", gate, &cfg.base_url);
|
||||||
|
// Resolve the bearer from the ACTIVE model's OWN manager: a grok model
|
||||||
|
// 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 auth_scheme = model_facts.auth_scheme;
|
||||||
let mut extra_headers = cfg.extra_headers;
|
let mut extra_headers = cfg.extra_headers;
|
||||||
crate::agent::config::inject_url_derived_headers(
|
crate::agent::config::inject_url_derived_headers(
|
||||||
@@ -312,6 +612,10 @@ impl SessionActor {
|
|||||||
top_p: cfg.top_p,
|
top_p: cfg.top_p,
|
||||||
api_backend: cfg.api_backend,
|
api_backend: cfg.api_backend,
|
||||||
auth_scheme,
|
auth_scheme,
|
||||||
|
anthropic_oauth,
|
||||||
|
github_copilot,
|
||||||
|
openai_codex,
|
||||||
|
chat_compat: cfg.chat_compat,
|
||||||
extra_headers,
|
extra_headers,
|
||||||
context_window: cfg.context_window.get(),
|
context_window: cfg.context_window.get(),
|
||||||
reasoning_effort: cfg.reasoning_effort,
|
reasoning_effort: cfg.reasoning_effort,
|
||||||
@@ -321,15 +625,7 @@ impl SessionActor {
|
|||||||
idle_timeout_secs: None,
|
idle_timeout_secs: None,
|
||||||
origin_client: self.origin_client.clone(),
|
origin_client: self.origin_client.clone(),
|
||||||
attribution_callback: self.attribution_callback.clone(),
|
attribution_callback: self.attribution_callback.clone(),
|
||||||
bearer_resolver: if use_bearer_resolver {
|
bearer_resolver: inference_auth_manager.map(auth_manager_bearer_resolver),
|
||||||
self.auth_manager
|
|
||||||
.as_ref()
|
|
||||||
.map(|am| -> kigi_sampler::SharedBearerResolver {
|
|
||||||
std::sync::Arc::new(AuthManagerBearerResolver(am.clone()))
|
|
||||||
})
|
|
||||||
} else {
|
|
||||||
None
|
|
||||||
},
|
|
||||||
supports_backend_search: self.supports_backend_search.get(),
|
supports_backend_search: self.supports_backend_search.get(),
|
||||||
compactions_remaining: self.compactions_remaining.get(),
|
compactions_remaining: self.compactions_remaining.get(),
|
||||||
compaction_at_tokens: self.compaction_at_tokens.get(),
|
compaction_at_tokens: self.compaction_at_tokens.get(),
|
||||||
@@ -459,17 +755,27 @@ impl SessionActor {
|
|||||||
slug: &str,
|
slug: &str,
|
||||||
) -> Option<kigi_sampler::SamplerConfig> {
|
) -> Option<kigi_sampler::SamplerConfig> {
|
||||||
let creds = self.chat_state_handle.get_credentials().await;
|
let creds = self.chat_state_handle.get_credentials().await;
|
||||||
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();
|
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();
|
let endpoints = self.models_manager.endpoints();
|
||||||
crate::agent::config::resolve_aux_model_sampling_config(
|
crate::agent::config::resolve_aux_model_sampling_config(
|
||||||
slug,
|
slug,
|
||||||
&models,
|
&models,
|
||||||
&endpoints,
|
&endpoints,
|
||||||
session_key.as_deref(),
|
session_key.as_ref(),
|
||||||
creds.alpha_test_key.clone(),
|
creds.alpha_test_key.clone(),
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -484,9 +790,15 @@ impl SessionActor {
|
|||||||
) -> Option<(kigi_sampler::SamplingClient, String)> {
|
) -> Option<(kigi_sampler::SamplingClient, String)> {
|
||||||
let active_session_config = self.reconstruct_full_config().await;
|
let active_session_config = self.reconstruct_full_config().await;
|
||||||
let mut cfg = self.resolve_aux_sampler_config(slug).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(
|
crate::agent::config::stamp_session_local_sampler_fields(
|
||||||
&mut cfg,
|
&mut cfg,
|
||||||
&active_session_config,
|
&active_session_config,
|
||||||
|
aux_resolver,
|
||||||
Some(self.max_retries),
|
Some(self.max_retries),
|
||||||
);
|
);
|
||||||
let model = cfg.model.clone();
|
let model = cfg.model.clone();
|
||||||
@@ -634,6 +946,7 @@ impl SessionActor {
|
|||||||
session_id = % self.session_info.id.0, is_session_based = gate
|
session_id = % self.session_info.id.0, is_session_based = gate
|
||||||
.is_session_based, model_byok = gate.model_byok.as_str(),
|
.is_session_based, model_byok = gate.model_byok.as_str(),
|
||||||
endpoint_is_first_party = gate.endpoint_is_first_party,
|
endpoint_is_first_party = gate.endpoint_is_first_party,
|
||||||
|
credential_class = gate.credential_class.as_str(),
|
||||||
"auth recovery: sampler 401 not refreshable (api-key auth) — surfacing 401",
|
"auth recovery: sampler 401 not refreshable (api-key auth) — surfacing 401",
|
||||||
);
|
);
|
||||||
kigi_log::unified_log::warn(
|
kigi_log::unified_log::warn(
|
||||||
@@ -643,7 +956,8 @@ impl SessionActor {
|
|||||||
{ "kind" : error.kind.as_str(), "status_code" : error
|
{ "kind" : error.kind.as_str(), "status_code" : error
|
||||||
.status_code, "is_session_based" : gate.is_session_based,
|
.status_code, "is_session_based" : gate.is_session_based,
|
||||||
"model_byok" : gate.model_byok.as_str(),
|
"model_byok" : gate.model_byok.as_str(),
|
||||||
"endpoint_is_first_party" : gate.endpoint_is_first_party, }
|
"endpoint_is_first_party" : gate.endpoint_is_first_party,
|
||||||
|
"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 {
|
if am.try_recover_unauthorized().await {
|
||||||
tracing::info!(
|
tracing::info!(
|
||||||
session_id = % self.session_info.id.0,
|
session_id = % self.session_info.id.0,
|
||||||
@@ -683,6 +1007,7 @@ impl SessionActor {
|
|||||||
None,
|
None,
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
}
|
||||||
if matches!(error.kind, SamplingErrorKind::IdleTimeout) {
|
if matches!(error.kind, SamplingErrorKind::IdleTimeout) {
|
||||||
self.signals_handle().record_idle_timeout();
|
self.signals_handle().record_idle_timeout();
|
||||||
}
|
}
|
||||||
@@ -844,14 +1169,17 @@ impl SessionActor {
|
|||||||
}
|
}
|
||||||
/// Proactively refresh the auth token if near expiry.
|
/// Proactively refresh the auth token if near expiry.
|
||||||
pub(super) async fn refresh_token_if_expired(&self) {
|
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
|
let (model_id, base_url) = self
|
||||||
.chat_state_handle
|
.chat_state_handle
|
||||||
.get_sampling_config()
|
.get_sampling_config()
|
||||||
.await
|
.await
|
||||||
.map(|c| (c.model, c.base_url))
|
.map(|c| (c.model, c.base_url))
|
||||||
.unwrap_or_default();
|
.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()
|
if self.auth_gate(&model_id, &base_url).active()
|
||||||
&& let Ok(key) = am.get_valid_token().await
|
&& let Ok(key) = am.get_valid_token().await
|
||||||
{
|
{
|
||||||
@@ -882,6 +1210,25 @@ impl SessionActor {
|
|||||||
.map(|c| c.model)
|
.map(|c| c.model)
|
||||||
.unwrap_or_default();
|
.unwrap_or_default();
|
||||||
let Some(ref key) = current_key else { return };
|
let Some(ref key) = current_key else { return };
|
||||||
|
// M7/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 {
|
let Some(new_key) = self.reload_api_key_from_config(¤t_model_id) else {
|
||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
@@ -896,6 +1243,15 @@ impl SessionActor {
|
|||||||
creds.api_key = Some(new_key);
|
creds.api_key = Some(new_key);
|
||||||
self.chat_state_handle.update_credentials(creds);
|
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> {
|
fn reload_api_key_from_config(&self, current_model_id: &str) -> Option<String> {
|
||||||
let raw_config = crate::config::load_effective_config()
|
let raw_config = crate::config::load_effective_config()
|
||||||
.map_err(|e| tracing::warn!(error = % e, "Failed to reload config"))
|
.map_err(|e| tracing::warn!(error = % e, "Failed to reload config"))
|
||||||
@@ -973,3 +1329,47 @@ impl SessionActor {
|
|||||||
.push_assistant_response(assistant_item);
|
.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)
|
ok_end_turn(0, None)
|
||||||
}
|
}
|
||||||
BuiltinAction::GoalPause => {
|
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;
|
let current_tokens = self.chat_state_handle.get_total_tokens().await as i64;
|
||||||
use crate::session::goal_tracker::{GoalPauseReason, GoalStatus};
|
use crate::session::goal_tracker::{GoalPauseReason, GoalStatus};
|
||||||
let (msg, changed) = {
|
let (msg, changed) = {
|
||||||
let mut tracker = self.goal_tracker.lock();
|
let mut tracker = self.goal_tracker.lock();
|
||||||
match tracker.status() {
|
match tracker.status() {
|
||||||
Some(GoalStatus::Active) => {
|
Some(GoalStatus::Active) => {
|
||||||
debug_assert!(
|
// Side effect OUTSIDE the assert: debug_assert!
|
||||||
tracker.pause(GoalPauseReason::User),
|
// strips its condition in release builds, which
|
||||||
"Active goal must pause"
|
// 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)
|
("Goal paused. Use /goal resume to continue.", true)
|
||||||
}
|
}
|
||||||
Some(
|
Some(
|
||||||
@@ -790,27 +798,149 @@ impl SessionActor {
|
|||||||
unreachable!("GoalResume is intercepted in handle_prompt")
|
unreachable!("GoalResume is intercepted in handle_prompt")
|
||||||
}
|
}
|
||||||
BuiltinAction::GoalClear => {
|
BuiltinAction::GoalClear => {
|
||||||
self.goal_tracker.lock().clear();
|
if self.graph_owns_goal_engine() {
|
||||||
// `/goal clear` is a deliberate user reset — drop both
|
self.send_slash_command_output(
|
||||||
// streaks so stale counters from the previous goal
|
"A graph owns the goal engine. Use /graph clear instead.",
|
||||||
// can't leak into the next one.
|
)
|
||||||
self.goal_continuation_streak
|
.await;
|
||||||
.store(0, std::sync::atomic::Ordering::Relaxed);
|
return ok_end_turn(0, None);
|
||||||
self.goal_blocked_streak
|
}
|
||||||
.store(0, std::sync::atomic::Ordering::Relaxed);
|
// `/goal clear` is a deliberate user reset — the shared
|
||||||
// Drop goal-turn-origin task ids so a future goal's drain
|
// helper drops the tracker, both streaks, goal-turn task
|
||||||
// doesn't suppress the next goal's (or post-goal) tasks.
|
// ids, per-subagent token records, pending classifier
|
||||||
self.goal_turn_task_ids.lock().clear();
|
// claims, and notifies the pager. Shared with the graph
|
||||||
// Clear per-subagent token records so stale entries
|
// node boundary and `/graph clear`.
|
||||||
// from the previous goal don't leak into the next.
|
self.reset_goal_engine_state().await;
|
||||||
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;
|
|
||||||
self.send_slash_command_output("Goal cleared.").await;
|
self.send_slash_command_output("Goal cleared.").await;
|
||||||
ok_end_turn(0, None)
|
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_signals: Option<crate::session::signals::SessionSignals>,
|
||||||
persisted_plan_mode: Option<crate::session::plan_mode::PlanModeSnapshot>,
|
persisted_plan_mode: Option<crate::session::plan_mode::PlanModeSnapshot>,
|
||||||
persisted_goal_mode: Option<crate::session::goal_tracker::GoalOrchestration>,
|
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>,
|
persisted_announcement_state: Option<crate::session::announcement_state::AnnouncementState>,
|
||||||
memory_config: Option<crate::config::MemoryConfig>,
|
memory_config: Option<crate::config::MemoryConfig>,
|
||||||
feedback_flags: crate::session::feedback_manager::FeedbackFlags,
|
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,
|
app_builder_deployer_config: kigi_tools::implementations::kigi::deploy_app::AppBuilderDeployerConfig,
|
||||||
write_file_enabled: bool,
|
write_file_enabled: bool,
|
||||||
goal_enabled: bool,
|
goal_enabled: bool,
|
||||||
|
graph_enabled: bool,
|
||||||
subagents_enabled: bool,
|
subagents_enabled: bool,
|
||||||
ask_user_question_enabled: bool,
|
ask_user_question_enabled: bool,
|
||||||
client_hooks: crate::extensions::hooks::ClientHooks,
|
client_hooks: crate::extensions::hooks::ClientHooks,
|
||||||
@@ -344,6 +346,7 @@ pub(crate) async fn spawn_session_actor(
|
|||||||
temperature: sampling_config.temperature,
|
temperature: sampling_config.temperature,
|
||||||
top_p: sampling_config.top_p,
|
top_p: sampling_config.top_p,
|
||||||
api_backend: sampling_config.api_backend.clone(),
|
api_backend: sampling_config.api_backend.clone(),
|
||||||
|
chat_compat: sampling_config.chat_compat,
|
||||||
extra_headers: sampling_config.extra_headers.clone(),
|
extra_headers: sampling_config.extra_headers.clone(),
|
||||||
context_window: context_window_override.unwrap_or(baseline_context_window),
|
context_window: context_window_override.unwrap_or(baseline_context_window),
|
||||||
reasoning_effort: sampling_config.reasoning_effort,
|
reasoning_effort: sampling_config.reasoning_effort,
|
||||||
@@ -429,6 +432,17 @@ pub(crate) async fn spawn_session_actor(
|
|||||||
};
|
};
|
||||||
Arc::new(parking_lot::Mutex::new(tracker))
|
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 current_prompt_mode = Arc::new(parking_lot::Mutex::new(PromptMode::Agent));
|
||||||
let turn_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());
|
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 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 {
|
let session = Arc::new_cyclic(|weak: &std::sync::Weak<SessionActor>| SessionActor {
|
||||||
session_info: session_info.clone(),
|
session_info: session_info.clone(),
|
||||||
auth_method_id,
|
auth_method_id,
|
||||||
model_auth_facts: std::cell::RefCell::new(None),
|
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,
|
attribution_callback,
|
||||||
auth_manager,
|
auth_manager,
|
||||||
state,
|
state,
|
||||||
@@ -1086,6 +1111,14 @@ pub(crate) async fn spawn_session_actor(
|
|||||||
goal_harness_enabled: std::sync::atomic::AtomicBool::new(false),
|
goal_harness_enabled: std::sync::atomic::AtomicBool::new(false),
|
||||||
goal_harness_availability_reconciled: std::sync::atomic::AtomicBool::new(false),
|
goal_harness_availability_reconciled: std::sync::atomic::AtomicBool::new(false),
|
||||||
goal_tracker,
|
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_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||||
goal_blocked_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;
|
.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 {
|
if let Some(ref display_cwd) = prompt_display_cwd {
|
||||||
session
|
session
|
||||||
.agent
|
.agent
|
||||||
@@ -1511,6 +1556,7 @@ pub(crate) async fn spawn_session_on_thread(
|
|||||||
persisted_signals: Option<crate::session::signals::SessionSignals>,
|
persisted_signals: Option<crate::session::signals::SessionSignals>,
|
||||||
persisted_plan_mode: Option<crate::session::plan_mode::PlanModeSnapshot>,
|
persisted_plan_mode: Option<crate::session::plan_mode::PlanModeSnapshot>,
|
||||||
persisted_goal_mode: Option<crate::session::goal_tracker::GoalOrchestration>,
|
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>,
|
persisted_announcement_state: Option<crate::session::announcement_state::AnnouncementState>,
|
||||||
memory_config: Option<crate::config::MemoryConfig>,
|
memory_config: Option<crate::config::MemoryConfig>,
|
||||||
feedback_flags: crate::session::feedback_manager::FeedbackFlags,
|
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,
|
app_builder_deployer_config: kigi_tools::implementations::kigi::deploy_app::AppBuilderDeployerConfig,
|
||||||
write_file_enabled: bool,
|
write_file_enabled: bool,
|
||||||
goal_enabled: bool,
|
goal_enabled: bool,
|
||||||
|
graph_enabled: bool,
|
||||||
subagents_enabled: bool,
|
subagents_enabled: bool,
|
||||||
ask_user_question_enabled: bool,
|
ask_user_question_enabled: bool,
|
||||||
client_hooks: crate::extensions::hooks::ClientHooks,
|
client_hooks: crate::extensions::hooks::ClientHooks,
|
||||||
@@ -1653,6 +1700,7 @@ pub(crate) async fn spawn_session_on_thread(
|
|||||||
persisted_signals,
|
persisted_signals,
|
||||||
persisted_plan_mode,
|
persisted_plan_mode,
|
||||||
persisted_goal_mode,
|
persisted_goal_mode,
|
||||||
|
persisted_graph_mode,
|
||||||
persisted_announcement_state,
|
persisted_announcement_state,
|
||||||
memory_config,
|
memory_config,
|
||||||
feedback_flags,
|
feedback_flags,
|
||||||
@@ -1667,6 +1715,7 @@ pub(crate) async fn spawn_session_on_thread(
|
|||||||
app_builder_deployer_config,
|
app_builder_deployer_config,
|
||||||
write_file_enabled,
|
write_file_enabled,
|
||||||
goal_enabled,
|
goal_enabled,
|
||||||
|
graph_enabled,
|
||||||
subagents_enabled,
|
subagents_enabled,
|
||||||
ask_user_question_enabled,
|
ask_user_question_enabled,
|
||||||
client_hooks,
|
client_hooks,
|
||||||
|
|||||||
@@ -478,6 +478,14 @@ impl SessionActor {
|
|||||||
if let Some(running_task) = running_task {
|
if let Some(running_task) = running_task {
|
||||||
running_task.abort();
|
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
|
// The aborted turn's `BlockingWaitGuard`s drop asynchronously (they
|
||||||
// live in tool futures owned by the drainer task / subagent spawn
|
// live in tool futures owned by the drainer task / subagent spawn
|
||||||
// task). Until they do, `queue_input` would read a stale depth > 0 and
|
// task). Until they do, `queue_input` would read a stale depth > 0 and
|
||||||
|
|||||||
@@ -305,10 +305,28 @@ impl SessionActor {
|
|||||||
objective,
|
objective,
|
||||||
token_budget,
|
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;
|
let reminder = self.setup_goal(&objective, token_budget).await;
|
||||||
vec![text_block(reminder), text_block(objective)]
|
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 } => {
|
GoalResumeOutcome::Inference { reminder, user_msg } => {
|
||||||
self.send_slash_command_output(&user_msg).await;
|
self.send_slash_command_output(&user_msg).await;
|
||||||
vec![text_block(reminder)]
|
vec![text_block(reminder)]
|
||||||
@@ -317,7 +335,33 @@ impl SessionActor {
|
|||||||
self.send_slash_command_output(&msg).await;
|
self.send_slash_command_output(&msg).await;
|
||||||
return ok_end_turn(0, None);
|
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,
|
_ => return self.execute_builtin_slash_command(action).await,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -688,13 +732,35 @@ impl SessionActor {
|
|||||||
self.goal_tracker.lock().status(),
|
self.goal_tracker.lock().status(),
|
||||||
);
|
);
|
||||||
if !goal_active {
|
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 {
|
match self.run_goal_round_end().await {
|
||||||
GoalRoundDecision::Continue(directive) => {
|
GoalRoundDecision::Continue(directive) => {
|
||||||
self.inject_goal_continuation_message(directive).await;
|
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]
|
#[test]
|
||||||
fn session_token_auth_gate_truth_table() {
|
fn session_token_auth_gate_truth_table() {
|
||||||
use crate::agent::auth_method::{ModelByok, session_token_auth_gate as gate};
|
use crate::agent::auth_method::{ModelByok, session_token_auth_gate as gate};
|
||||||
|
use crate::auth::credential_authority::CredentialClass;
|
||||||
// Non-session methods never refresh, regardless of BYOK status or endpoint.
|
// 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] {
|
for fp in [false, true] {
|
||||||
assert!(!gate(false, ModelByok::NotByok, fp));
|
for class in [CredentialClass::Pooled, CredentialClass::Primary] {
|
||||||
assert!(!gate(false, ModelByok::Byok, fp));
|
assert!(!gate(false, ModelByok::NotByok, fp, class));
|
||||||
assert!(!gate(false, ModelByok::Unknown, fp));
|
assert!(!gate(false, ModelByok::Byok, fp, class));
|
||||||
// Session method: a definite classification ignores the endpoint —
|
assert!(!gate(false, ModelByok::Unknown, fp, class));
|
||||||
// NotByok always refreshes (only ever routes to the session endpoint),
|
// Session method on an endpoint that DOES take a session
|
||||||
// a genuine per-model Byok never does.
|
// credential: a definite classification ignores the endpoint —
|
||||||
assert!(gate(true, ModelByok::NotByok, fp));
|
// NotByok refreshes, a genuine per-model Byok never does.
|
||||||
assert!(!gate(true, ModelByok::Byok, fp));
|
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
|
// Session method + Unknown BYOK: refresh only against a first-party xAI
|
||||||
// host, so a transiently-unclassifiable config can't demote a live session
|
// host, so a transiently-unclassifiable config can't demote a live session
|
||||||
// (the stale-token 401 regression) yet the session token never leaks to a
|
// (the stale-token 401 regression) yet the session token never leaks to a
|
||||||
// third-party BYOK endpoint. This arm was unconditionally `false` pre-fix.
|
// third-party BYOK endpoint. This arm was unconditionally `false` pre-fix.
|
||||||
assert!(gate(true, ModelByok::Unknown, true));
|
assert!(gate(
|
||||||
assert!(!gate(true, ModelByok::Unknown, false));
|
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
|
/// 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,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: crate::sampling::ApiBackend::ChatCompletions,
|
api_backend: crate::sampling::ApiBackend::ChatCompletions,
|
||||||
|
chat_compat: Default::default(),
|
||||||
auth_scheme: Default::default(),
|
auth_scheme: Default::default(),
|
||||||
|
anthropic_oauth: false,
|
||||||
|
github_copilot: false,
|
||||||
|
openai_codex: false,
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: 256_000,
|
context_window: 256_000,
|
||||||
force_http1: false,
|
force_http1: false,
|
||||||
@@ -865,7 +890,7 @@ async fn set_session_model_invalidates_byok_memo_for_same_model_id() {
|
|||||||
header_injector: None,
|
header_injector: None,
|
||||||
};
|
};
|
||||||
let _ = actor
|
let _ = actor
|
||||||
.handle_set_session_model(cfg, false, false, true, 85)
|
.handle_set_session_model(cfg, None, false, false, true, 85)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
assert!(
|
assert!(
|
||||||
|
|||||||
@@ -44,7 +44,11 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
|
|||||||
temperature: None,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
auth_scheme: Default::default(),
|
auth_scheme: Default::default(),
|
||||||
|
anthropic_oauth: false,
|
||||||
|
github_copilot: false,
|
||||||
|
openai_codex: false,
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: 100_000,
|
context_window: 100_000,
|
||||||
force_http1: false,
|
force_http1: false,
|
||||||
@@ -87,6 +91,7 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
|
|||||||
temperature: None,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: std::num::NonZeroU64::new(100_000).unwrap(),
|
context_window: std::num::NonZeroU64::new(100_000).unwrap(),
|
||||||
reasoning_effort: None,
|
reasoning_effort: None,
|
||||||
@@ -104,6 +109,7 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
|
|||||||
session_info,
|
session_info,
|
||||||
auth_method_id: test_auth_method_id("test-auth"),
|
auth_method_id: test_auth_method_id("test-auth"),
|
||||||
model_auth_facts: std::cell::RefCell::new(None),
|
model_auth_facts: std::cell::RefCell::new(None),
|
||||||
|
selected_catalog_key: std::cell::RefCell::new(None),
|
||||||
attribution_callback: None,
|
attribution_callback: None,
|
||||||
auth_manager: None,
|
auth_manager: None,
|
||||||
state: TokioMutex::new(State {
|
state: TokioMutex::new(State {
|
||||||
@@ -212,6 +218,16 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
|
|||||||
"/tmp/test-session",
|
"/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_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||||
goal_blocked_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,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
auth_scheme: Default::default(),
|
auth_scheme: Default::default(),
|
||||||
|
anthropic_oauth: false,
|
||||||
|
github_copilot: false,
|
||||||
|
openai_codex: false,
|
||||||
context_window: 100_000,
|
context_window: 100_000,
|
||||||
force_http1: false,
|
force_http1: false,
|
||||||
max_retries: None,
|
max_retries: None,
|
||||||
@@ -371,6 +391,7 @@ async fn first_turn_memory_injection_persists_to_chat_history() {
|
|||||||
temperature: None,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: std::num::NonZeroU64::new(100_000).unwrap(),
|
context_window: std::num::NonZeroU64::new(100_000).unwrap(),
|
||||||
reasoning_effort: None,
|
reasoning_effort: None,
|
||||||
@@ -456,7 +477,11 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
|
|||||||
temperature: None,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
auth_scheme: Default::default(),
|
auth_scheme: Default::default(),
|
||||||
|
anthropic_oauth: false,
|
||||||
|
github_copilot: false,
|
||||||
|
openai_codex: false,
|
||||||
context_window: 100_000,
|
context_window: 100_000,
|
||||||
force_http1: false,
|
force_http1: false,
|
||||||
max_retries: None,
|
max_retries: None,
|
||||||
@@ -503,6 +528,7 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
|
|||||||
temperature: None,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: std::num::NonZeroU64::new(100_000).unwrap(),
|
context_window: std::num::NonZeroU64::new(100_000).unwrap(),
|
||||||
reasoning_effort: None,
|
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(),
|
session_info: session_info.clone(),
|
||||||
auth_method_id: test_auth_method_id("test-auth"),
|
auth_method_id: test_auth_method_id("test-auth"),
|
||||||
model_auth_facts: std::cell::RefCell::new(None),
|
model_auth_facts: std::cell::RefCell::new(None),
|
||||||
|
selected_catalog_key: std::cell::RefCell::new(None),
|
||||||
attribution_callback: None,
|
attribution_callback: None,
|
||||||
auth_manager: None,
|
auth_manager: None,
|
||||||
state: TokioMutex::new(State {
|
state: TokioMutex::new(State {
|
||||||
@@ -648,6 +675,16 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
|
|||||||
"/tmp/test-session",
|
"/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_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||||
goal_blocked_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"),
|
auth_method_id: test_auth_method_id("test-auth"),
|
||||||
model_auth_facts: std::cell::RefCell::new(None),
|
model_auth_facts: std::cell::RefCell::new(None),
|
||||||
|
selected_catalog_key: std::cell::RefCell::new(None),
|
||||||
attribution_callback: None,
|
attribution_callback: None,
|
||||||
auth_manager: None,
|
auth_manager: None,
|
||||||
state,
|
state,
|
||||||
@@ -893,6 +931,16 @@ async fn cancel_running_task_teardown_clears_running_and_pending_work() {
|
|||||||
"/tmp/test-session",
|
"/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_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||||
goal_blocked_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,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: kigi_sampler::ApiBackend::Responses,
|
api_backend: kigi_sampler::ApiBackend::Responses,
|
||||||
|
chat_compat: Default::default(),
|
||||||
auth_scheme: Default::default(),
|
auth_scheme: Default::default(),
|
||||||
|
anthropic_oauth: false,
|
||||||
|
github_copilot: false,
|
||||||
|
openai_codex: false,
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: 100_000,
|
context_window: 100_000,
|
||||||
force_http1: false,
|
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"),
|
auth_method_id: test_auth_method_id("test-auth"),
|
||||||
model_auth_facts: std::cell::RefCell::new(None),
|
model_auth_facts: std::cell::RefCell::new(None),
|
||||||
|
selected_catalog_key: std::cell::RefCell::new(None),
|
||||||
attribution_callback: None,
|
attribution_callback: None,
|
||||||
auth_manager: None,
|
auth_manager: None,
|
||||||
state,
|
state,
|
||||||
@@ -1871,6 +1924,16 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() {
|
|||||||
"/tmp/test-session",
|
"/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_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||||
goal_blocked_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,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: std::num::NonZeroU64::new(200_000).unwrap(),
|
context_window: std::num::NonZeroU64::new(200_000).unwrap(),
|
||||||
reasoning_effort: None,
|
reasoning_effort: None,
|
||||||
@@ -127,6 +128,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
|
|||||||
attribution_callback: None,
|
attribution_callback: None,
|
||||||
auth_method_id: test_auth_method_id("cached_token"),
|
auth_method_id: test_auth_method_id("cached_token"),
|
||||||
model_auth_facts: std::cell::RefCell::new(None),
|
model_auth_facts: std::cell::RefCell::new(None),
|
||||||
|
selected_catalog_key: std::cell::RefCell::new(None),
|
||||||
auth_manager: {
|
auth_manager: {
|
||||||
let dir = tempfile::tempdir().unwrap();
|
let dir = tempfile::tempdir().unwrap();
|
||||||
let mgr = std::sync::Arc::new(crate::auth::AuthManager::new(
|
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",
|
"/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_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||||
goal_blocked_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,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: std::num::NonZeroU64::new(context_window)
|
context_window: std::num::NonZeroU64::new(context_window)
|
||||||
.expect("test context_window must be non-zero"),
|
.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(),
|
rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(),
|
||||||
auth_method_id: test_auth_method_id("test-auth"),
|
auth_method_id: test_auth_method_id("test-auth"),
|
||||||
model_auth_facts: std::cell::RefCell::new(None),
|
model_auth_facts: std::cell::RefCell::new(None),
|
||||||
|
selected_catalog_key: std::cell::RefCell::new(None),
|
||||||
attribution_callback: None,
|
attribution_callback: None,
|
||||||
auth_manager: None,
|
auth_manager: None,
|
||||||
state,
|
state,
|
||||||
@@ -173,6 +175,16 @@ async fn create_test_actor(
|
|||||||
"/tmp/test-session",
|
"/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_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||||
goal_blocked_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,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: std::num::NonZeroU64::new(context_window)
|
context_window: std::num::NonZeroU64::new(context_window)
|
||||||
.expect("test context_window must be non-zero"),
|
.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(),
|
rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(),
|
||||||
auth_method_id: test_auth_method_id("test-auth"),
|
auth_method_id: test_auth_method_id("test-auth"),
|
||||||
model_auth_facts: std::cell::RefCell::new(None),
|
model_auth_facts: std::cell::RefCell::new(None),
|
||||||
|
selected_catalog_key: std::cell::RefCell::new(None),
|
||||||
attribution_callback: None,
|
attribution_callback: None,
|
||||||
auth_manager: None,
|
auth_manager: None,
|
||||||
state,
|
state,
|
||||||
@@ -611,6 +625,16 @@ async fn create_test_actor_with_memory(
|
|||||||
"/tmp/test-session",
|
"/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_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||||
goal_blocked_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,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: std::num::NonZeroU64::new(200_000).unwrap(),
|
context_window: std::num::NonZeroU64::new(200_000).unwrap(),
|
||||||
reasoning_effort: None,
|
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(),
|
rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(),
|
||||||
auth_method_id: test_auth_method_id("cached_token"),
|
auth_method_id: test_auth_method_id("cached_token"),
|
||||||
model_auth_facts: std::cell::RefCell::new(None),
|
model_auth_facts: std::cell::RefCell::new(None),
|
||||||
|
selected_catalog_key: std::cell::RefCell::new(None),
|
||||||
auth_manager: {
|
auth_manager: {
|
||||||
let dir = tempfile::tempdir().unwrap();
|
let dir = tempfile::tempdir().unwrap();
|
||||||
let mgr = std::sync::Arc::new(crate::auth::AuthManager::new(
|
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",
|
"/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_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||||
goal_blocked_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,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: std::num::NonZeroU64::new(context_window)
|
context_window: std::num::NonZeroU64::new(context_window)
|
||||||
.expect("test context_window must be non-zero"),
|
.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"),
|
auth_method_id: test_auth_method_id("test-auth"),
|
||||||
model_auth_facts: std::cell::RefCell::new(None),
|
model_auth_facts: std::cell::RefCell::new(None),
|
||||||
|
selected_catalog_key: std::cell::RefCell::new(None),
|
||||||
attribution_callback: None,
|
attribution_callback: None,
|
||||||
auth_manager: None,
|
auth_manager: None,
|
||||||
state,
|
state,
|
||||||
@@ -235,6 +237,16 @@ async fn create_test_actor_with_memory(
|
|||||||
"/tmp/test-session",
|
"/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_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||||
goal_blocked_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"),
|
auth_method_id: test_auth_method_id("test-auth"),
|
||||||
model_auth_facts: std::cell::RefCell::new(None),
|
model_auth_facts: std::cell::RefCell::new(None),
|
||||||
|
selected_catalog_key: std::cell::RefCell::new(None),
|
||||||
attribution_callback: None,
|
attribution_callback: None,
|
||||||
auth_manager: None,
|
auth_manager: None,
|
||||||
state,
|
state,
|
||||||
@@ -181,6 +182,16 @@ pub(super) async fn make_replay_send_update_fixture() -> ReplaySendUpdateFixture
|
|||||||
"/tmp/test-session",
|
"/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_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||||
goal_blocked_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,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: std::num::NonZeroU64::new(context_window)
|
context_window: std::num::NonZeroU64::new(context_window)
|
||||||
.expect("test context_window must be non-zero"),
|
.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"),
|
auth_method_id: test_auth_method_id("test-auth"),
|
||||||
model_auth_facts: std::cell::RefCell::new(None),
|
model_auth_facts: std::cell::RefCell::new(None),
|
||||||
|
selected_catalog_key: std::cell::RefCell::new(None),
|
||||||
attribution_callback: None,
|
attribution_callback: None,
|
||||||
auth_manager: None,
|
auth_manager: None,
|
||||||
state,
|
state,
|
||||||
@@ -286,6 +288,16 @@ pub(crate) async fn create_test_actor_ex(
|
|||||||
"/tmp/test-session",
|
"/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_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||||
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
||||||
|
|||||||
@@ -145,6 +145,13 @@ pub enum SessionCommand {
|
|||||||
},
|
},
|
||||||
SetSessionModel {
|
SetSessionModel {
|
||||||
sampling_config: kigi_sampler::SamplerConfig,
|
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,
|
use_concise: bool,
|
||||||
/// When `false`, skip the system prompt rewrite (concise/default swap).
|
/// When `false`, skip the system prompt rewrite (concise/default swap).
|
||||||
/// Set to `false` for forked sessions so mid-session model switches
|
/// Set to `false` for forked sessions so mid-session model switches
|
||||||
|
|||||||
@@ -2149,6 +2149,7 @@ mod inline_auto_compact_flow_tests {
|
|||||||
temperature: None,
|
temperature: None,
|
||||||
top_p: None,
|
top_p: None,
|
||||||
api_backend: Default::default(),
|
api_backend: Default::default(),
|
||||||
|
chat_compat: Default::default(),
|
||||||
extra_headers: Default::default(),
|
extra_headers: Default::default(),
|
||||||
context_window: std::num::NonZeroU64::new(context_window)
|
context_window: std::num::NonZeroU64::new(context_window)
|
||||||
.expect("test context_window must be non-zero"),
|
.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"),
|
auth_method_id: test_auth_method_id("test-auth"),
|
||||||
model_auth_facts: std::cell::RefCell::new(None),
|
model_auth_facts: std::cell::RefCell::new(None),
|
||||||
|
selected_catalog_key: std::cell::RefCell::new(None),
|
||||||
attribution_callback: None,
|
attribution_callback: None,
|
||||||
auth_manager: None,
|
auth_manager: None,
|
||||||
state,
|
state,
|
||||||
@@ -2268,6 +2270,16 @@ mod inline_auto_compact_flow_tests {
|
|||||||
"/tmp/test-session",
|
"/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_turn_task_ids: parking_lot::Mutex::new(std::collections::HashSet::new()),
|
||||||
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
goal_continuation_streak: std::sync::atomic::AtomicU32::new(0),
|
||||||
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
goal_blocked_streak: std::sync::atomic::AtomicU32::new(0),
|
||||||
|
|||||||
@@ -154,7 +154,9 @@ pub enum GoalPauseReason {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl 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 {
|
match self {
|
||||||
Self::User => GoalStatus::UserPaused,
|
Self::User => GoalStatus::UserPaused,
|
||||||
Self::BackOff => GoalStatus::BackOffPaused,
|
Self::BackOff => GoalStatus::BackOffPaused,
|
||||||
@@ -166,7 +168,8 @@ impl GoalPauseReason {
|
|||||||
|
|
||||||
/// Short, stable label stashed in the `GoalPaused` history entry's
|
/// Short, stable label stashed in the `GoalPaused` history entry's
|
||||||
/// `detail` so the pager's Recent History distinguishes pause causes.
|
/// `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 {
|
match self {
|
||||||
Self::User => "user",
|
Self::User => "user",
|
||||||
Self::BackOff => "back_off",
|
Self::BackOff => "back_off",
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,571 @@
|
|||||||
|
//! Graph planner runner: one attempt at decomposing an objective into a
|
||||||
|
//! validated node DAG.
|
||||||
|
//!
|
||||||
|
//! Deliberately thin: the spawn plumbing (harness-internal subagent,
|
||||||
|
//! verbatim-fork, fail-open model retry) is reused from
|
||||||
|
//! [`goal_planner`](super::goal_planner) via the same
|
||||||
|
//! [`GoalPlannerSpawner`] contract; this module only swaps the template
|
||||||
|
//! and replaces "plan file exists" with "graph JSON parses and passes
|
||||||
|
//! the static DAG gate" ([`graph_plan::parse_and_validate`]).
|
||||||
|
//!
|
||||||
|
//! Outcome split (both are loud, nothing is papered over):
|
||||||
|
//! - [`GraphPlannerOutcome::Invalid`] — the planner wrote an artifact
|
||||||
|
//! that failed validation. Retryable ONCE by the caller, feeding the
|
||||||
|
//! precise validation error back as CONTEXT.
|
||||||
|
//! - [`GraphPlannerOutcome::FailClosed`] — spawn/transport/missing-file
|
||||||
|
//! failure. The caller pauses the graph; `/graph resume` retries.
|
||||||
|
|
||||||
|
use std::path::Path;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use super::goal_planner::{
|
||||||
|
GoalPlannerSpawner, RoleRenderedPrompt, SpawnError, parse_terminal_response,
|
||||||
|
};
|
||||||
|
use super::goal_role_tools::RoleToolNames;
|
||||||
|
use super::graph_plan::{self, MAX_GRAPH_JSON_BYTES};
|
||||||
|
use super::graph_tracker::GraphNode;
|
||||||
|
|
||||||
|
const GRAPH_PLANNER_PROMPT_TEMPLATE: &str = include_str!("templates/graph_planner_prompt.md");
|
||||||
|
const GRAPH_REPLANNER_PROMPT_TEMPLATE: &str = include_str!("templates/graph_replanner_prompt.md");
|
||||||
|
pub(crate) const GRAPH_PLANNER_SUBAGENT_DESCRIPTION: &str = "graph plan writer";
|
||||||
|
pub(crate) const GRAPH_REPLANNER_SUBAGENT_DESCRIPTION: &str = "graph replanner";
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub(crate) enum GraphPlannerOutcome {
|
||||||
|
/// Validated, canonicalized nodes (topo-ordered, final node appended).
|
||||||
|
Planned(Vec<GraphNode>),
|
||||||
|
/// Artifact written but rejected by the static gate; retry once with
|
||||||
|
/// the reason as feedback.
|
||||||
|
Invalid { reason: String },
|
||||||
|
/// Infrastructure/spawn failure or missing artifact; pause the graph.
|
||||||
|
FailClosed { reason: String },
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) struct GraphPlannerInputs<'a> {
|
||||||
|
pub objective: &'a str,
|
||||||
|
/// Empty on the first attempt; the previous attempt's validation
|
||||||
|
/// error on the retry.
|
||||||
|
pub feedback: &'a str,
|
||||||
|
pub graph_file: &'a Path,
|
||||||
|
pub tool_names: &'a RoleToolNames,
|
||||||
|
pub inherit_tool_names: &'a RoleToolNames,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Run one graph-planner attempt end to end: render, spawn, read the
|
||||||
|
/// artifact (size-capped), validate, canonicalize.
|
||||||
|
pub(crate) async fn run_graph_planner(
|
||||||
|
spawner: Arc<dyn GoalPlannerSpawner>,
|
||||||
|
inputs: GraphPlannerInputs<'_>,
|
||||||
|
) -> GraphPlannerOutcome {
|
||||||
|
if let Some(parent) = inputs.graph_file.parent()
|
||||||
|
&& let Err(err) = tokio::fs::create_dir_all(parent).await
|
||||||
|
{
|
||||||
|
return GraphPlannerOutcome::FailClosed {
|
||||||
|
reason: format!("failed to create graph dir {}: {err}", parent.display()),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
// A stale artifact from a prior pass would satisfy the
|
||||||
|
// missing-file guard below and be trusted as this pass's output.
|
||||||
|
// Delete first; only NotFound is benign.
|
||||||
|
if let Err(err) = tokio::fs::remove_file(inputs.graph_file).await
|
||||||
|
&& err.kind() != std::io::ErrorKind::NotFound
|
||||||
|
{
|
||||||
|
return GraphPlannerOutcome::FailClosed {
|
||||||
|
reason: format!("failed to clear stale graph artifact: {err}"),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
let graph_file_str = inputs.graph_file.to_string_lossy();
|
||||||
|
let with_graph_file = GRAPH_PLANNER_PROMPT_TEMPLATE.replace("{GRAPH_FILE}", &graph_file_str);
|
||||||
|
let render = |tool_names: &RoleToolNames| -> String {
|
||||||
|
let rendered = tool_names.apply(&with_graph_file);
|
||||||
|
let mut full = String::with_capacity(rendered.len() + inputs.objective.len() + 256);
|
||||||
|
full.push_str(&rendered);
|
||||||
|
full.push_str("\n\nOBJECTIVE:\n");
|
||||||
|
full.push_str(inputs.objective);
|
||||||
|
full.push_str("\n\nCONTEXT:\n");
|
||||||
|
full.push_str(inputs.feedback);
|
||||||
|
full.push('\n');
|
||||||
|
full
|
||||||
|
};
|
||||||
|
let prompt = RoleRenderedPrompt {
|
||||||
|
primary: render(inputs.tool_names),
|
||||||
|
fallback: render(inputs.inherit_tool_names),
|
||||||
|
};
|
||||||
|
|
||||||
|
let spawn_id = uuid::Uuid::now_v7().to_string();
|
||||||
|
let response = match spawner.spawn_planner(&spawn_id, prompt).await {
|
||||||
|
Ok(text) => text,
|
||||||
|
Err(SpawnError::Transport(detail)) => {
|
||||||
|
return GraphPlannerOutcome::FailClosed {
|
||||||
|
reason: format!("graph planner transport error: {detail}"),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
Err(SpawnError::Runtime { message, cancelled }) => {
|
||||||
|
return GraphPlannerOutcome::FailClosed {
|
||||||
|
reason: if cancelled {
|
||||||
|
format!("graph planner aborted: {message}")
|
||||||
|
} else {
|
||||||
|
format!("graph planner runtime error: {message}")
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
match tokio::fs::metadata(inputs.graph_file).await {
|
||||||
|
Ok(meta) if meta.is_file() && meta.len() > 0 => {
|
||||||
|
if meta.len() > MAX_GRAPH_JSON_BYTES {
|
||||||
|
return GraphPlannerOutcome::Invalid {
|
||||||
|
reason: format!(
|
||||||
|
"graph JSON is {} bytes; the cap is {MAX_GRAPH_JSON_BYTES}",
|
||||||
|
meta.len()
|
||||||
|
),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
tracing::info!(
|
||||||
|
graph_file = %graph_file_str,
|
||||||
|
terminal_token_ok = parse_terminal_response(&response),
|
||||||
|
response_snippet = %response.chars().take(120).collect::<String>(),
|
||||||
|
"graph planner: graph file missing or empty; failing closed",
|
||||||
|
);
|
||||||
|
return GraphPlannerOutcome::FailClosed {
|
||||||
|
reason: "graph planner produced no graph file".to_owned(),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let json = match tokio::fs::read_to_string(inputs.graph_file).await {
|
||||||
|
Ok(json) => json,
|
||||||
|
Err(err) => {
|
||||||
|
return GraphPlannerOutcome::FailClosed {
|
||||||
|
reason: format!("failed to read graph file: {err}"),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
match graph_plan::parse_and_validate(&json, inputs.objective) {
|
||||||
|
Ok(nodes) => GraphPlannerOutcome::Planned(nodes),
|
||||||
|
Err(err) => GraphPlannerOutcome::Invalid {
|
||||||
|
reason: err.to_string(),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) struct GraphReplannerInputs<'a> {
|
||||||
|
pub objective: &'a str,
|
||||||
|
/// Compact JSON of the existing nodes (id/title/status/deps).
|
||||||
|
pub current_graph: &'a str,
|
||||||
|
/// The queued discoveries, one per line with their origin node ids.
|
||||||
|
pub discoveries: &'a str,
|
||||||
|
/// On retry, the previous artifact's validation error.
|
||||||
|
pub feedback: &'a str,
|
||||||
|
pub graph_file: &'a Path,
|
||||||
|
pub tool_names: &'a RoleToolNames,
|
||||||
|
pub inherit_tool_names: &'a RoleToolNames,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Run one REPLAN attempt: render, spawn, read the artifact, validate
|
||||||
|
/// against the existing graph (append-only). An empty `{"nodes": []}`
|
||||||
|
/// appendix is the sanctioned "everything already covered" escape hatch
|
||||||
|
/// and returns `Planned(vec![])`.
|
||||||
|
pub(crate) async fn run_graph_replanner(
|
||||||
|
spawner: Arc<dyn GoalPlannerSpawner>,
|
||||||
|
existing: &[GraphNode],
|
||||||
|
inputs: GraphReplannerInputs<'_>,
|
||||||
|
) -> GraphPlannerOutcome {
|
||||||
|
if let Some(parent) = inputs.graph_file.parent()
|
||||||
|
&& let Err(err) = tokio::fs::create_dir_all(parent).await
|
||||||
|
{
|
||||||
|
return GraphPlannerOutcome::FailClosed {
|
||||||
|
reason: format!("failed to create graph dir {}: {err}", parent.display()),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
// A stale artifact from a prior pass would satisfy the
|
||||||
|
// missing-file guard below and be trusted as this pass's output.
|
||||||
|
// Delete first; only NotFound is benign.
|
||||||
|
if let Err(err) = tokio::fs::remove_file(inputs.graph_file).await
|
||||||
|
&& err.kind() != std::io::ErrorKind::NotFound
|
||||||
|
{
|
||||||
|
return GraphPlannerOutcome::FailClosed {
|
||||||
|
reason: format!("failed to clear stale graph artifact: {err}"),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
let graph_file_str = inputs.graph_file.to_string_lossy();
|
||||||
|
let with_graph_file = GRAPH_REPLANNER_PROMPT_TEMPLATE.replace("{GRAPH_FILE}", &graph_file_str);
|
||||||
|
let render = |tool_names: &RoleToolNames| -> String {
|
||||||
|
let rendered = tool_names.apply(&with_graph_file);
|
||||||
|
format!(
|
||||||
|
"{rendered}\n\nOBJECTIVE:\n{}\n\nCURRENT GRAPH:\n{}\n\nDISCOVERIES:\n{}\n\nCONTEXT:\n{}\n",
|
||||||
|
inputs.objective, inputs.current_graph, inputs.discoveries, inputs.feedback
|
||||||
|
)
|
||||||
|
};
|
||||||
|
let prompt = RoleRenderedPrompt {
|
||||||
|
primary: render(inputs.tool_names),
|
||||||
|
fallback: render(inputs.inherit_tool_names),
|
||||||
|
};
|
||||||
|
let spawn_id = uuid::Uuid::now_v7().to_string();
|
||||||
|
let response = match spawner.spawn_planner(&spawn_id, prompt).await {
|
||||||
|
Ok(text) => text,
|
||||||
|
Err(SpawnError::Transport(detail)) => {
|
||||||
|
return GraphPlannerOutcome::FailClosed {
|
||||||
|
reason: format!("graph replanner transport error: {detail}"),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
Err(SpawnError::Runtime { message, cancelled }) => {
|
||||||
|
return GraphPlannerOutcome::FailClosed {
|
||||||
|
reason: if cancelled {
|
||||||
|
format!("graph replanner aborted: {message}")
|
||||||
|
} else {
|
||||||
|
format!("graph replanner runtime error: {message}")
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
|
};
|
||||||
|
match tokio::fs::metadata(inputs.graph_file).await {
|
||||||
|
Ok(meta) if meta.is_file() && meta.len() > 0 => {
|
||||||
|
if meta.len() > MAX_GRAPH_JSON_BYTES {
|
||||||
|
return GraphPlannerOutcome::Invalid {
|
||||||
|
reason: format!(
|
||||||
|
"replan JSON is {} bytes; the cap is {MAX_GRAPH_JSON_BYTES}",
|
||||||
|
meta.len()
|
||||||
|
),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
tracing::info!(
|
||||||
|
graph_file = %graph_file_str,
|
||||||
|
terminal_token_ok = parse_terminal_response(&response),
|
||||||
|
"graph replanner: artifact missing or empty; failing closed",
|
||||||
|
);
|
||||||
|
return GraphPlannerOutcome::FailClosed {
|
||||||
|
reason: "graph replanner produced no artifact".to_owned(),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let json = match tokio::fs::read_to_string(inputs.graph_file).await {
|
||||||
|
Ok(json) => json,
|
||||||
|
Err(err) => {
|
||||||
|
return GraphPlannerOutcome::FailClosed {
|
||||||
|
reason: format!("failed to read replan artifact: {err}"),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
};
|
||||||
|
// Escape hatch: an explicitly empty appendix means "already covered".
|
||||||
|
if let Ok(v) = serde_json::from_str::<serde_json::Value>(&json)
|
||||||
|
&& v.get("nodes")
|
||||||
|
.and_then(|n| n.as_array())
|
||||||
|
.is_some_and(Vec::is_empty)
|
||||||
|
{
|
||||||
|
return GraphPlannerOutcome::Planned(Vec::new());
|
||||||
|
}
|
||||||
|
match graph_plan::validate_replan(existing, &json) {
|
||||||
|
Ok(nodes) => GraphPlannerOutcome::Planned(nodes),
|
||||||
|
Err(err) => GraphPlannerOutcome::Invalid {
|
||||||
|
reason: err.to_string(),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Generic single-shot artifact pass: render `template` + `sections`,
|
||||||
|
/// spawn the role child, enforce the stale-artifact/size/missing-file
|
||||||
|
/// discipline, and return the artifact's raw JSON for the caller to
|
||||||
|
/// validate. Shared by the optimizer (and any future boundary pass).
|
||||||
|
pub(crate) struct ArtifactPassSpec<'a> {
|
||||||
|
pub template: &'a str,
|
||||||
|
pub sections: &'a str,
|
||||||
|
pub graph_file: &'a Path,
|
||||||
|
pub tool_names: &'a RoleToolNames,
|
||||||
|
pub role: &'a str,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn run_graph_artifact_pass(
|
||||||
|
spawner: Arc<dyn GoalPlannerSpawner>,
|
||||||
|
spec: ArtifactPassSpec<'_>,
|
||||||
|
) -> Result<String, String> {
|
||||||
|
if let Some(parent) = spec.graph_file.parent()
|
||||||
|
&& let Err(err) = tokio::fs::create_dir_all(parent).await
|
||||||
|
{
|
||||||
|
return Err(format!("failed to create graph dir: {err}"));
|
||||||
|
}
|
||||||
|
if let Err(err) = tokio::fs::remove_file(spec.graph_file).await
|
||||||
|
&& err.kind() != std::io::ErrorKind::NotFound
|
||||||
|
{
|
||||||
|
return Err(format!("failed to clear stale artifact: {err}"));
|
||||||
|
}
|
||||||
|
let graph_file_str = spec.graph_file.to_string_lossy();
|
||||||
|
let rendered = spec
|
||||||
|
.tool_names
|
||||||
|
.apply(&spec.template.replace("{GRAPH_FILE}", &graph_file_str));
|
||||||
|
let prompt_text = format!("{rendered}\n\n{}", spec.sections);
|
||||||
|
let prompt = RoleRenderedPrompt {
|
||||||
|
primary: prompt_text.clone(),
|
||||||
|
fallback: prompt_text,
|
||||||
|
};
|
||||||
|
let spawn_id = uuid::Uuid::now_v7().to_string();
|
||||||
|
match spawner.spawn_planner(&spawn_id, prompt).await {
|
||||||
|
Ok(_) => {}
|
||||||
|
Err(SpawnError::Transport(detail)) => {
|
||||||
|
return Err(format!("{} transport error: {detail}", spec.role));
|
||||||
|
}
|
||||||
|
Err(SpawnError::Runtime { message, cancelled }) => {
|
||||||
|
return Err(if cancelled {
|
||||||
|
format!("{} aborted: {message}", spec.role)
|
||||||
|
} else {
|
||||||
|
format!("{} runtime error: {message}", spec.role)
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
match tokio::fs::metadata(spec.graph_file).await {
|
||||||
|
Ok(meta) if meta.is_file() && meta.len() > 0 && meta.len() <= MAX_GRAPH_JSON_BYTES => {}
|
||||||
|
Ok(meta) if meta.len() > MAX_GRAPH_JSON_BYTES => {
|
||||||
|
return Err(format!(
|
||||||
|
"{} artifact is {} bytes; the cap is {MAX_GRAPH_JSON_BYTES}",
|
||||||
|
spec.role,
|
||||||
|
meta.len()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
_ => return Err(format!("{} produced no artifact", spec.role)),
|
||||||
|
}
|
||||||
|
tokio::fs::read_to_string(spec.graph_file)
|
||||||
|
.await
|
||||||
|
.map_err(|err| format!("failed to read {} artifact: {err}", spec.role))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::session::goal_role_tools::tests::summary_with;
|
||||||
|
use kigi_tools::types::tool::ToolKind;
|
||||||
|
use std::path::PathBuf;
|
||||||
|
use std::sync::Mutex;
|
||||||
|
|
||||||
|
enum MockReply {
|
||||||
|
Done,
|
||||||
|
Transport,
|
||||||
|
Runtime { cancelled: bool },
|
||||||
|
}
|
||||||
|
|
||||||
|
struct MockSpawner {
|
||||||
|
response: MockReply,
|
||||||
|
body: Option<Vec<u8>>,
|
||||||
|
target: PathBuf,
|
||||||
|
last_prompt: Mutex<Option<String>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait::async_trait]
|
||||||
|
impl GoalPlannerSpawner for MockSpawner {
|
||||||
|
async fn spawn_planner(
|
||||||
|
&self,
|
||||||
|
_id: &str,
|
||||||
|
prompt: RoleRenderedPrompt,
|
||||||
|
) -> Result<String, SpawnError> {
|
||||||
|
*self.last_prompt.lock().unwrap() = Some(prompt.primary.clone());
|
||||||
|
if let Some(body) = &self.body {
|
||||||
|
std::fs::write(&self.target, body).unwrap();
|
||||||
|
}
|
||||||
|
match &self.response {
|
||||||
|
MockReply::Done => Ok("Done".to_owned()),
|
||||||
|
MockReply::Transport => Err(SpawnError::Transport("channel closed".into())),
|
||||||
|
MockReply::Runtime { cancelled } => Err(SpawnError::Runtime {
|
||||||
|
message: "boom".into(),
|
||||||
|
cancelled: *cancelled,
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn tool_names() -> RoleToolNames {
|
||||||
|
RoleToolNames::from_summary(&summary_with(&[
|
||||||
|
(ToolKind::Read, "read_file"),
|
||||||
|
(ToolKind::Search, "grep"),
|
||||||
|
(ToolKind::List, "list_files"),
|
||||||
|
(ToolKind::Write, "write"),
|
||||||
|
]))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Self-cleaning temp home per test: never leak dirs into the OS
|
||||||
|
/// temp root (storage discipline — see AGENTS.md gates).
|
||||||
|
fn tmp_graph_file(_name: &str) -> (tempfile::TempDir, PathBuf) {
|
||||||
|
let dir = tempfile::TempDir::new().unwrap();
|
||||||
|
let path = dir.path().join("graph.json");
|
||||||
|
(dir, path)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn run(spawner: MockSpawner, graph_file: &Path) -> GraphPlannerOutcome {
|
||||||
|
let names = tool_names();
|
||||||
|
run_graph_planner(
|
||||||
|
Arc::new(spawner),
|
||||||
|
GraphPlannerInputs {
|
||||||
|
objective: "build the thing",
|
||||||
|
feedback: "",
|
||||||
|
graph_file,
|
||||||
|
tool_names: &names,
|
||||||
|
inherit_tool_names: &names,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn valid_artifact_yields_canonical_nodes() {
|
||||||
|
let (_tmp, target) = tmp_graph_file("valid");
|
||||||
|
let body = serde_json::json!({
|
||||||
|
"nodes": [
|
||||||
|
{"id": "core", "title": "Core", "spec": "core spec", "deps": []},
|
||||||
|
{"id": "ui", "title": "UI", "spec": "ui spec", "deps": ["core"]},
|
||||||
|
]
|
||||||
|
})
|
||||||
|
.to_string();
|
||||||
|
let spawner = MockSpawner {
|
||||||
|
response: MockReply::Done,
|
||||||
|
body: Some(body.into_bytes()),
|
||||||
|
target: target.clone(),
|
||||||
|
last_prompt: Mutex::new(None),
|
||||||
|
};
|
||||||
|
match run(spawner, &target).await {
|
||||||
|
GraphPlannerOutcome::Planned(nodes) => {
|
||||||
|
assert_eq!(nodes.len(), 3, "2 planner nodes + appended final");
|
||||||
|
assert_eq!(nodes[2].id, crate::session::graph_tracker::FINAL_NODE_ID);
|
||||||
|
}
|
||||||
|
other => panic!("expected Planned, got {other:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn prompt_embeds_objective_feedback_and_tool_names() {
|
||||||
|
let (_tmp, target) = tmp_graph_file("prompt");
|
||||||
|
let spawner = MockSpawner {
|
||||||
|
response: MockReply::Done,
|
||||||
|
body: None,
|
||||||
|
target: target.clone(),
|
||||||
|
last_prompt: Mutex::new(None),
|
||||||
|
};
|
||||||
|
let prompt_cell = std::sync::Arc::new(spawner);
|
||||||
|
let names = tool_names();
|
||||||
|
let _ = run_graph_planner(
|
||||||
|
prompt_cell.clone(),
|
||||||
|
GraphPlannerInputs {
|
||||||
|
objective: "OBJ-MARKER",
|
||||||
|
feedback: "FEEDBACK-MARKER",
|
||||||
|
graph_file: &target,
|
||||||
|
tool_names: &names,
|
||||||
|
inherit_tool_names: &names,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let prompt = prompt_cell.last_prompt.lock().unwrap().clone().unwrap();
|
||||||
|
assert!(prompt.contains("OBJ-MARKER"));
|
||||||
|
assert!(prompt.contains("FEEDBACK-MARKER"));
|
||||||
|
assert!(prompt.contains(&target.to_string_lossy().into_owned()));
|
||||||
|
assert!(prompt.contains("read_file"), "placeholders rendered");
|
||||||
|
assert!(!prompt.contains("{READ_TOOL}"), "no leftover placeholder");
|
||||||
|
assert!(!prompt.contains("{GRAPH_FILE}"), "no leftover placeholder");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn invalid_artifact_is_retryable_with_reason() {
|
||||||
|
let (_tmp, target) = tmp_graph_file("invalid");
|
||||||
|
let spawner = MockSpawner {
|
||||||
|
response: MockReply::Done,
|
||||||
|
body: Some(br#"{"nodes":[{"id":"a","title":"A","spec":"s","deps":["a"]}]}"#.to_vec()),
|
||||||
|
target: target.clone(),
|
||||||
|
last_prompt: Mutex::new(None),
|
||||||
|
};
|
||||||
|
match run(spawner, &target).await {
|
||||||
|
GraphPlannerOutcome::Invalid { reason } => {
|
||||||
|
assert!(reason.contains("depends on itself"), "{reason}");
|
||||||
|
}
|
||||||
|
other => panic!("expected Invalid, got {other:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn missing_artifact_fails_closed() {
|
||||||
|
let (_tmp, target) = tmp_graph_file("missing");
|
||||||
|
let spawner = MockSpawner {
|
||||||
|
response: MockReply::Done,
|
||||||
|
body: None,
|
||||||
|
target: target.clone(),
|
||||||
|
last_prompt: Mutex::new(None),
|
||||||
|
};
|
||||||
|
match run(spawner, &target).await {
|
||||||
|
GraphPlannerOutcome::FailClosed { reason } => {
|
||||||
|
assert!(reason.contains("no graph file"), "{reason}");
|
||||||
|
}
|
||||||
|
other => panic!("expected FailClosed, got {other:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn runtime_error_fails_closed() {
|
||||||
|
let (_tmp, target) = tmp_graph_file("runtime");
|
||||||
|
let spawner = MockSpawner {
|
||||||
|
response: MockReply::Runtime { cancelled: false },
|
||||||
|
body: None,
|
||||||
|
target: target.clone(),
|
||||||
|
last_prompt: Mutex::new(None),
|
||||||
|
};
|
||||||
|
match run(spawner, &target).await {
|
||||||
|
GraphPlannerOutcome::FailClosed { reason } => {
|
||||||
|
assert!(reason.contains("runtime error"), "{reason}");
|
||||||
|
}
|
||||||
|
other => panic!("expected FailClosed, got {other:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn oversize_artifact_is_invalid_with_cap_in_reason() {
|
||||||
|
let (_tmp, target) = tmp_graph_file("oversize");
|
||||||
|
let mut body = vec![b'x'; (MAX_GRAPH_JSON_BYTES as usize) + 1];
|
||||||
|
body[0] = b'{'; // content is irrelevant; the size gate fires first
|
||||||
|
let spawner = MockSpawner {
|
||||||
|
response: MockReply::Done,
|
||||||
|
body: Some(body),
|
||||||
|
target: target.clone(),
|
||||||
|
last_prompt: Mutex::new(None),
|
||||||
|
};
|
||||||
|
match run(spawner, &target).await {
|
||||||
|
GraphPlannerOutcome::Invalid { reason } => {
|
||||||
|
assert!(reason.contains("the cap is"), "{reason}");
|
||||||
|
}
|
||||||
|
other => panic!("expected Invalid, got {other:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn transport_error_fails_closed() {
|
||||||
|
let (_tmp, target) = tmp_graph_file("transport");
|
||||||
|
let spawner = MockSpawner {
|
||||||
|
response: MockReply::Transport,
|
||||||
|
body: None,
|
||||||
|
target: target.clone(),
|
||||||
|
last_prompt: Mutex::new(None),
|
||||||
|
};
|
||||||
|
match run(spawner, &target).await {
|
||||||
|
GraphPlannerOutcome::FailClosed { reason } => {
|
||||||
|
assert!(reason.contains("transport error"), "{reason}");
|
||||||
|
}
|
||||||
|
other => panic!("expected FailClosed, got {other:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn cancelled_runtime_error_reports_aborted() {
|
||||||
|
let (_tmp, target) = tmp_graph_file("aborted");
|
||||||
|
let spawner = MockSpawner {
|
||||||
|
response: MockReply::Runtime { cancelled: true },
|
||||||
|
body: None,
|
||||||
|
target: target.clone(),
|
||||||
|
last_prompt: Mutex::new(None),
|
||||||
|
};
|
||||||
|
match run(spawner, &target).await {
|
||||||
|
GraphPlannerOutcome::FailClosed { reason } => {
|
||||||
|
assert!(reason.contains("aborted"), "{reason}");
|
||||||
|
}
|
||||||
|
other => panic!("expected FailClosed, got {other:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,305 @@
|
|||||||
|
//! Project-level shared graph file (G4): `.kigi/graph.jsonl` at the git
|
||||||
|
//! root, so a graph follows the REPOSITORY, not the session.
|
||||||
|
//!
|
||||||
|
//! The session tracker remains the single source of truth; this file is
|
||||||
|
//! a PROJECTION refreshed at every checkpoint. Format (beads-style,
|
||||||
|
//! line-mergeable thanks to content-hash node ids):
|
||||||
|
//!
|
||||||
|
//! - line 1: the orchestration header (everything except `nodes`)
|
||||||
|
//! - lines 2..: one `GraphNode` per line
|
||||||
|
//!
|
||||||
|
//! Concurrency: an advisory `flock` on a sidecar `.lock` file makes the
|
||||||
|
//! session that CREATED or RESUMED the graph the single writer; other
|
||||||
|
//! kigi instances get a read-only view (`/graph status`) with an
|
||||||
|
//! explicit notice. The lock is held for the graph's lifetime in that
|
||||||
|
//! session and released on `/graph clear` (or process exit).
|
||||||
|
//!
|
||||||
|
//! Git discipline: kigi only WRITES the file — committing it is the
|
||||||
|
//! user's decision, never automated.
|
||||||
|
|
||||||
|
use std::io::Write;
|
||||||
|
use std::path::{Path, PathBuf};
|
||||||
|
|
||||||
|
use fs2::FileExt;
|
||||||
|
|
||||||
|
use super::graph_tracker::{GraphNode, GraphOrchestration};
|
||||||
|
|
||||||
|
/// Header line: the orchestration minus its nodes (which follow one per
|
||||||
|
/// line). The shadow `nodes` field is skipped on write and REJECTED on
|
||||||
|
/// read when non-empty — nodes embedded in the header would silently
|
||||||
|
/// duplicate the per-line entries.
|
||||||
|
#[derive(serde::Serialize, serde::Deserialize)]
|
||||||
|
struct ProjectGraphHeader {
|
||||||
|
#[serde(flatten)]
|
||||||
|
orchestration: GraphOrchestration,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn header_has_inline_nodes(header: &ProjectGraphHeader) -> bool {
|
||||||
|
!header.orchestration.nodes.is_empty()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Held exclusive advisory lock on the project graph. Dropping releases.
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct ProjectGraphLock {
|
||||||
|
_file: std::fs::File,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub enum LockOutcome {
|
||||||
|
Acquired(ProjectGraphLock),
|
||||||
|
/// Another kigi instance holds the lock.
|
||||||
|
Busy,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `.kigi` dir under the git root of `cwd`; `None` outside a git repo
|
||||||
|
/// (the project-graph feature is git-scoped by design).
|
||||||
|
pub fn project_graph_dir(cwd: &Path) -> Option<PathBuf> {
|
||||||
|
kigi_workspace::session::git::find_git_root_from_path(cwd)
|
||||||
|
.ok()
|
||||||
|
.map(|root| root.join(".kigi"))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn graph_file_path(dir: &Path) -> PathBuf {
|
||||||
|
dir.join("graph.jsonl")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn lock_file_path(dir: &Path) -> PathBuf {
|
||||||
|
dir.join("graph.jsonl.lock")
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Try to become the project graph's single writer. Fail-fast: any I/O
|
||||||
|
/// error other than "already locked" propagates.
|
||||||
|
pub fn try_acquire_writer(dir: &Path) -> std::io::Result<LockOutcome> {
|
||||||
|
std::fs::create_dir_all(dir)?;
|
||||||
|
let file = std::fs::OpenOptions::new()
|
||||||
|
.create(true)
|
||||||
|
.truncate(false)
|
||||||
|
.write(true)
|
||||||
|
.open(lock_file_path(dir))?;
|
||||||
|
match file.try_lock_exclusive() {
|
||||||
|
Ok(()) => Ok(LockOutcome::Acquired(ProjectGraphLock { _file: file })),
|
||||||
|
// fs2 maps contention differently per platform (EWOULDBLOCK on
|
||||||
|
// unix, ERROR_LOCK_VIOLATION on Windows); its
|
||||||
|
// `lock_contended_error()` is the portable classifier.
|
||||||
|
Err(err)
|
||||||
|
if err.kind() == std::io::ErrorKind::WouldBlock
|
||||||
|
|| err.raw_os_error() == fs2::lock_contended_error().raw_os_error() =>
|
||||||
|
{
|
||||||
|
Ok(LockOutcome::Busy)
|
||||||
|
}
|
||||||
|
Err(err) => Err(err),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Atomically project the orchestration to `.kigi/graph.jsonl`
|
||||||
|
/// (tmp + rename, same discipline as the session state file).
|
||||||
|
pub fn project(dir: &Path, state: &GraphOrchestration) -> std::io::Result<()> {
|
||||||
|
std::fs::create_dir_all(dir)?;
|
||||||
|
let mut header_state = state.clone();
|
||||||
|
let nodes = std::mem::take(&mut header_state.nodes);
|
||||||
|
let mut header_value = serde_json::to_value(&header_state)
|
||||||
|
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
|
||||||
|
if let Some(obj) = header_value.as_object_mut() {
|
||||||
|
// The contract says "minus nodes"; drop the empty vec the
|
||||||
|
// struct serializer would otherwise emit.
|
||||||
|
obj.remove("nodes");
|
||||||
|
}
|
||||||
|
let mut buf = Vec::with_capacity(4096);
|
||||||
|
serde_json::to_writer(&mut buf, &header_value)
|
||||||
|
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
|
||||||
|
buf.push(b'\n');
|
||||||
|
for node in &nodes {
|
||||||
|
serde_json::to_writer(&mut buf, node)
|
||||||
|
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
|
||||||
|
buf.push(b'\n');
|
||||||
|
}
|
||||||
|
let target = graph_file_path(dir);
|
||||||
|
let tmp = target.with_extension("jsonl.tmp");
|
||||||
|
{
|
||||||
|
let mut f = std::fs::File::create(&tmp)?;
|
||||||
|
f.write_all(&buf)?;
|
||||||
|
f.sync_all()?;
|
||||||
|
}
|
||||||
|
std::fs::rename(&tmp, &target)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Load the projected graph, `Ok(None)` when absent. Malformed content
|
||||||
|
/// is an ERROR (never silently treated as "no graph") — the file is
|
||||||
|
/// user-visible, git-merged state; corruption must surface.
|
||||||
|
pub fn load(dir: &Path) -> std::io::Result<Option<GraphOrchestration>> {
|
||||||
|
let path = graph_file_path(dir);
|
||||||
|
let raw = match std::fs::read_to_string(&path) {
|
||||||
|
Ok(raw) => raw,
|
||||||
|
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(None),
|
||||||
|
Err(err) => return Err(err),
|
||||||
|
};
|
||||||
|
let mut lines = raw.lines().filter(|l| !l.trim().is_empty());
|
||||||
|
let Some(header_line) = lines.next() else {
|
||||||
|
return Ok(None);
|
||||||
|
};
|
||||||
|
let header: ProjectGraphHeader = serde_json::from_str(header_line).map_err(|e| {
|
||||||
|
std::io::Error::new(
|
||||||
|
std::io::ErrorKind::InvalidData,
|
||||||
|
format!("{} line 1: {e}", path.display()),
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
if header_has_inline_nodes(&header) {
|
||||||
|
return Err(std::io::Error::new(
|
||||||
|
std::io::ErrorKind::InvalidData,
|
||||||
|
format!(
|
||||||
|
"{} line 1 embeds nodes inline; nodes belong one per line",
|
||||||
|
path.display()
|
||||||
|
),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let mut state = header.orchestration;
|
||||||
|
for (idx, line) in lines.enumerate() {
|
||||||
|
let node: GraphNode = serde_json::from_str(line).map_err(|e| {
|
||||||
|
std::io::Error::new(
|
||||||
|
std::io::ErrorKind::InvalidData,
|
||||||
|
format!("{} node line {}: {e}", path.display(), idx + 2),
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
state.nodes.push(node);
|
||||||
|
}
|
||||||
|
Ok(Some(state))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Remove the projection (on `/graph clear`). Missing file is fine.
|
||||||
|
pub fn remove(dir: &Path) -> std::io::Result<()> {
|
||||||
|
match std::fs::remove_file(graph_file_path(dir)) {
|
||||||
|
Ok(()) => Ok(()),
|
||||||
|
Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
|
||||||
|
Err(err) => Err(err),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::session::goal_tracker::{GoalPhase, GoalStatus};
|
||||||
|
use crate::session::graph_tracker::{DepKind, NodeDep, NodeStatus};
|
||||||
|
use tempfile::TempDir;
|
||||||
|
|
||||||
|
fn sample_state() -> GraphOrchestration {
|
||||||
|
GraphOrchestration {
|
||||||
|
graph_id: "g-1".into(),
|
||||||
|
objective: "ship it".into(),
|
||||||
|
status: GoalStatus::Active,
|
||||||
|
phase: GoalPhase::Executing,
|
||||||
|
plan_version: 2,
|
||||||
|
nodes: vec![
|
||||||
|
GraphNode {
|
||||||
|
id: "gn-aaaa".into(),
|
||||||
|
title: "A".into(),
|
||||||
|
spec: "do a".into(),
|
||||||
|
deps: vec![],
|
||||||
|
status: NodeStatus::Achieved,
|
||||||
|
goal_id: Some("goal-1".into()),
|
||||||
|
rounds: 2,
|
||||||
|
tokens_used: 100,
|
||||||
|
failure: None,
|
||||||
|
},
|
||||||
|
GraphNode {
|
||||||
|
id: "gn-bbbb".into(),
|
||||||
|
title: "B".into(),
|
||||||
|
spec: "do b".into(),
|
||||||
|
deps: vec![NodeDep {
|
||||||
|
on: "gn-aaaa".into(),
|
||||||
|
kind: DepKind::DiscoveredFrom,
|
||||||
|
}],
|
||||||
|
status: NodeStatus::Running,
|
||||||
|
goal_id: None,
|
||||||
|
rounds: 0,
|
||||||
|
tokens_used: 0,
|
||||||
|
failure: None,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
current_node: Some("gn-bbbb".into()),
|
||||||
|
created_at: "2026-07-20T00:00:00Z".into(),
|
||||||
|
elapsed_ms: 12,
|
||||||
|
token_budget: Some(1_000),
|
||||||
|
tokens_spent_nodes: 100,
|
||||||
|
history: vec![],
|
||||||
|
pause_message: None,
|
||||||
|
pending_discoveries: vec![],
|
||||||
|
replan_runs: 1,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn project_load_round_trip_is_line_per_node() {
|
||||||
|
let tmp = TempDir::new().unwrap();
|
||||||
|
let state = sample_state();
|
||||||
|
project(tmp.path(), &state).unwrap();
|
||||||
|
let raw = std::fs::read_to_string(graph_file_path(tmp.path())).unwrap();
|
||||||
|
assert_eq!(raw.lines().count(), 3, "header + one line per node");
|
||||||
|
assert!(
|
||||||
|
!raw.lines().next().unwrap().contains("\"nodes\""),
|
||||||
|
"header must omit nodes entirely"
|
||||||
|
);
|
||||||
|
assert!(raw.lines().nth(1).unwrap().contains("gn-aaaa"));
|
||||||
|
let loaded = load(tmp.path()).unwrap().expect("present");
|
||||||
|
assert_eq!(loaded.graph_id, state.graph_id);
|
||||||
|
assert_eq!(loaded.plan_version, 2);
|
||||||
|
assert_eq!(loaded.nodes.len(), 2);
|
||||||
|
assert_eq!(loaded.nodes[1].deps[0].kind, DepKind::DiscoveredFrom);
|
||||||
|
assert_eq!(loaded.current_node.as_deref(), Some("gn-bbbb"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn load_absent_is_none_and_remove_is_idempotent() {
|
||||||
|
let tmp = TempDir::new().unwrap();
|
||||||
|
assert!(load(tmp.path()).unwrap().is_none());
|
||||||
|
remove(tmp.path()).unwrap();
|
||||||
|
project(tmp.path(), &sample_state()).unwrap();
|
||||||
|
remove(tmp.path()).unwrap();
|
||||||
|
assert!(load(tmp.path()).unwrap().is_none());
|
||||||
|
remove(tmp.path()).unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn header_with_inline_nodes_is_rejected() {
|
||||||
|
let tmp = TempDir::new().unwrap();
|
||||||
|
let mut bad = serde_json::to_value(sample_state()).unwrap();
|
||||||
|
// Keep nodes inline in the header — a hand-edited/merged file.
|
||||||
|
bad.as_object_mut().unwrap().remove("current_node");
|
||||||
|
std::fs::write(graph_file_path(tmp.path()), format!("{bad}\n")).unwrap();
|
||||||
|
let err = load(tmp.path()).unwrap_err();
|
||||||
|
assert!(err.to_string().contains("inline"), "{err}");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn malformed_content_is_a_loud_error_not_a_missing_graph() {
|
||||||
|
let tmp = TempDir::new().unwrap();
|
||||||
|
std::fs::write(graph_file_path(tmp.path()), "not json\n").unwrap();
|
||||||
|
let err = load(tmp.path()).unwrap_err();
|
||||||
|
assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
|
||||||
|
assert!(err.to_string().contains("line 1"), "{err}");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn writer_lock_is_exclusive_within_and_across_handles() {
|
||||||
|
let tmp = TempDir::new().unwrap();
|
||||||
|
let first = try_acquire_writer(tmp.path()).unwrap();
|
||||||
|
let LockOutcome::Acquired(_guard) = first else {
|
||||||
|
panic!("first acquire must win");
|
||||||
|
};
|
||||||
|
match try_acquire_writer(tmp.path()).unwrap() {
|
||||||
|
LockOutcome::Busy => {}
|
||||||
|
LockOutcome::Acquired(_) => {
|
||||||
|
// flock is per-fd on some platforms within one process;
|
||||||
|
// if this arm is reached the platform lets the same
|
||||||
|
// process re-lock, which is still safe for our
|
||||||
|
// cross-INSTANCE contract — but on macOS/Linux flock
|
||||||
|
// between distinct fds does contend, so treat as bug.
|
||||||
|
panic!("second handle must observe Busy");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
drop(_guard);
|
||||||
|
assert!(matches!(
|
||||||
|
try_acquire_writer(tmp.path()).unwrap(),
|
||||||
|
LockOutcome::Acquired(_)
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,526 @@
|
|||||||
|
//! Box-drawing DAG rendering for `/graph show` (G5).
|
||||||
|
//!
|
||||||
|
//! Sugiyama-lite over the node DAG: longest-path layering, one-pass
|
||||||
|
//! barycenter ordering, dagre-style dummy pass-throughs so every drawn
|
||||||
|
//! edge spans exactly one layer gap, and greedy bus-lane allocation in
|
||||||
|
//! the connector gutters. Pure text (theme-free), deterministic, and
|
||||||
|
//! snapshot-testable; the output rides ordinary scrollback, which the
|
||||||
|
//! pager already scrolls.
|
||||||
|
//!
|
||||||
|
//! Honest ceiling: when the packed grid exceeds `max_width`, the caller
|
||||||
|
//! falls back to the indented status tree — box-drawing wrapped by the
|
||||||
|
//! terminal is worse than no drawing.
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
use super::graph_tracker::{DepKind, GraphOrchestration, NodeStatus};
|
||||||
|
|
||||||
|
/// Character-grid canvas with box-drawing-aware merging.
|
||||||
|
struct Canvas {
|
||||||
|
rows: Vec<Vec<char>>,
|
||||||
|
width: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Canvas {
|
||||||
|
fn new(width: usize) -> Self {
|
||||||
|
Self {
|
||||||
|
rows: Vec::new(),
|
||||||
|
width,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn put(&mut self, row: usize, col: usize, ch: char) {
|
||||||
|
if col >= self.width {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
while self.rows.len() <= row {
|
||||||
|
self.rows.push(vec![' '; self.width]);
|
||||||
|
}
|
||||||
|
let cell = &mut self.rows[row][col];
|
||||||
|
*cell = merge_glyph(*cell, ch);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn put_str(&mut self, row: usize, col: usize, s: &str) {
|
||||||
|
for (i, ch) in s.chars().enumerate() {
|
||||||
|
self.put(row, col + i, ch);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn render(&self) -> String {
|
||||||
|
self.rows
|
||||||
|
.iter()
|
||||||
|
.map(|r| r.iter().collect::<String>().trim_end().to_owned())
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.join("\n")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Merge overlapping box-drawing strokes (a horizontal bus crossing a
|
||||||
|
/// vertical pass-through becomes `┼`; anything else: last writer wins,
|
||||||
|
/// except blanks never overwrite ink).
|
||||||
|
fn merge_glyph(existing: char, new: char) -> char {
|
||||||
|
match (existing, new) {
|
||||||
|
(' ', n) => n,
|
||||||
|
(e, ' ') => e,
|
||||||
|
('─', '│') | ('│', '─') => '┼',
|
||||||
|
('─', '┴') | ('┴', '─') => '┴',
|
||||||
|
('─', '┬') | ('┬', '─') => '┬',
|
||||||
|
(_, n) => n,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn status_glyph(status: NodeStatus) -> char {
|
||||||
|
match status {
|
||||||
|
NodeStatus::Achieved => '✓',
|
||||||
|
NodeStatus::Running | NodeStatus::Verifying => '▶',
|
||||||
|
NodeStatus::Ready => '○',
|
||||||
|
NodeStatus::Waiting => '·',
|
||||||
|
NodeStatus::Failed => '✗',
|
||||||
|
NodeStatus::Blocked => '⊘',
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
const TITLE_BUDGET: usize = 18;
|
||||||
|
const H_GAP: usize = 3;
|
||||||
|
|
||||||
|
struct Cell {
|
||||||
|
/// Real node index, or `None` for a dummy pass-through.
|
||||||
|
node: Option<usize>,
|
||||||
|
/// Column of the cell's connector center on the grid.
|
||||||
|
center: usize,
|
||||||
|
/// Grid column where the box starts (real nodes only).
|
||||||
|
left: usize,
|
||||||
|
label: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Render the DAG as box-drawing text, or `None` when it cannot fit
|
||||||
|
/// `max_width` (caller falls back to the indented tree).
|
||||||
|
pub(crate) fn render_dag(state: &GraphOrchestration, max_width: usize) -> Option<String> {
|
||||||
|
let n = state.nodes.len();
|
||||||
|
if n == 0 {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let index_of: HashMap<&str, usize> = state
|
||||||
|
.nodes
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, node)| (node.id.as_str(), i))
|
||||||
|
.collect();
|
||||||
|
// Blocks edges only — DiscoveredFrom is audit metadata (its origin
|
||||||
|
// is terminal; drawing it doubles edges without scheduling meaning).
|
||||||
|
let edges: Vec<(usize, usize)> = state
|
||||||
|
.nodes
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.flat_map(|(to, node)| {
|
||||||
|
let index_of = &index_of;
|
||||||
|
node.deps
|
||||||
|
.iter()
|
||||||
|
.filter(|d| d.kind == DepKind::Blocks)
|
||||||
|
.filter_map(move |d| index_of.get(d.on.as_str()).map(|&from| (from, to)))
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
// Longest-path layering (deps validated acyclic upstream).
|
||||||
|
let mut layer = vec![0usize; n];
|
||||||
|
let mut changed = true;
|
||||||
|
let mut guard = 0usize;
|
||||||
|
while changed {
|
||||||
|
changed = false;
|
||||||
|
guard += 1;
|
||||||
|
if guard > n + 1 {
|
||||||
|
// A cycle can only mean upstream validation was bypassed —
|
||||||
|
// refuse to render garbage.
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
for &(from, to) in &edges {
|
||||||
|
if layer[to] < layer[from] + 1 {
|
||||||
|
layer[to] = layer[from] + 1;
|
||||||
|
changed = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let depth = layer.iter().copied().max().unwrap_or(0) + 1;
|
||||||
|
|
||||||
|
// Dummy chains: split any edge spanning >1 layer into unit hops.
|
||||||
|
// Segment endpoints are (layer, slot) pairs; real slots 0..n, dummy
|
||||||
|
// slots appended after.
|
||||||
|
#[derive(Clone, Copy, PartialEq)]
|
||||||
|
struct Slot {
|
||||||
|
real: Option<usize>,
|
||||||
|
}
|
||||||
|
let mut slots: Vec<Slot> = (0..n).map(|i| Slot { real: Some(i) }).collect();
|
||||||
|
let mut slot_layer: Vec<usize> = layer.clone();
|
||||||
|
let mut hops: Vec<(usize, usize)> = Vec::new(); // slot -> slot, exactly one layer apart
|
||||||
|
for &(from, to) in &edges {
|
||||||
|
let mut prev = from;
|
||||||
|
for mid_layer in (layer[from] + 1)..layer[to] {
|
||||||
|
slots.push(Slot { real: None });
|
||||||
|
slot_layer.push(mid_layer);
|
||||||
|
let dummy = slots.len() - 1;
|
||||||
|
hops.push((prev, dummy));
|
||||||
|
prev = dummy;
|
||||||
|
}
|
||||||
|
hops.push((prev, to));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Layer membership + one-pass barycenter ordering (parents' mean
|
||||||
|
// position; stable by construction order for roots).
|
||||||
|
let mut layers: Vec<Vec<usize>> = vec![Vec::new(); depth];
|
||||||
|
for (slot, &l) in slot_layer.iter().enumerate() {
|
||||||
|
layers[l].push(slot);
|
||||||
|
}
|
||||||
|
let mut pos: Vec<f64> = vec![0.0; slots.len()];
|
||||||
|
for (i, &slot) in layers[0].iter().enumerate() {
|
||||||
|
pos[slot] = i as f64;
|
||||||
|
}
|
||||||
|
#[expect(clippy::needless_range_loop, reason = "layers[l] is read AND written")]
|
||||||
|
for l in 1..depth {
|
||||||
|
let mut keyed: Vec<(f64, usize)> = layers[l]
|
||||||
|
.iter()
|
||||||
|
.map(|&slot| {
|
||||||
|
let parents: Vec<usize> = hops
|
||||||
|
.iter()
|
||||||
|
.filter(|&&(_, t)| t == slot)
|
||||||
|
.map(|&(f, _)| f)
|
||||||
|
.collect();
|
||||||
|
let key = if parents.is_empty() {
|
||||||
|
f64::MAX // parentless mid-layer nodes go last, stably
|
||||||
|
} else {
|
||||||
|
parents.iter().map(|&p| pos[p]).sum::<f64>() / parents.len() as f64
|
||||||
|
};
|
||||||
|
(key, slot)
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
keyed.sort_by(|a, b| a.0.total_cmp(&b.0));
|
||||||
|
layers[l] = keyed.iter().map(|&(_, s)| s).collect();
|
||||||
|
for (i, &(_, slot)) in keyed.iter().enumerate() {
|
||||||
|
pos[slot] = i as f64;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Horizontal packing per layer; grid width = widest layer.
|
||||||
|
let label_of = |i: usize| -> String {
|
||||||
|
let node = &state.nodes[i];
|
||||||
|
let mut title = node.title.clone();
|
||||||
|
if title.chars().count() > TITLE_BUDGET {
|
||||||
|
title = title.chars().take(TITLE_BUDGET - 1).collect::<String>() + "…";
|
||||||
|
}
|
||||||
|
format!("{} {}", status_glyph(node.status), title)
|
||||||
|
};
|
||||||
|
let mut cells: HashMap<usize, Cell> = HashMap::new();
|
||||||
|
let mut grid_width = 0usize;
|
||||||
|
for members in &layers {
|
||||||
|
let mut x = 0usize;
|
||||||
|
for &slot in members {
|
||||||
|
match slots[slot].real {
|
||||||
|
Some(i) => {
|
||||||
|
let label = label_of(i);
|
||||||
|
let box_w = label.chars().count() + 2;
|
||||||
|
cells.insert(
|
||||||
|
slot,
|
||||||
|
Cell {
|
||||||
|
node: Some(i),
|
||||||
|
center: x + box_w / 2,
|
||||||
|
left: x,
|
||||||
|
label,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
x += box_w + H_GAP;
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
cells.insert(
|
||||||
|
slot,
|
||||||
|
Cell {
|
||||||
|
node: None,
|
||||||
|
center: x,
|
||||||
|
left: x,
|
||||||
|
label: String::new(),
|
||||||
|
},
|
||||||
|
);
|
||||||
|
x += 1 + H_GAP;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
grid_width = grid_width.max(x.saturating_sub(H_GAP));
|
||||||
|
}
|
||||||
|
if grid_width > max_width {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Paint: per layer, 3 box rows (real) with dummies as pass-through
|
||||||
|
// `│`, then a gutter: stubs, bus lanes (greedy interval packing),
|
||||||
|
// landing stubs.
|
||||||
|
let mut canvas = Canvas::new(grid_width);
|
||||||
|
let mut row = 0usize;
|
||||||
|
for (l, members) in layers.iter().enumerate() {
|
||||||
|
// Box band.
|
||||||
|
for &slot in members {
|
||||||
|
let cell = &cells[&slot];
|
||||||
|
match cell.node {
|
||||||
|
Some(_) => {
|
||||||
|
let w = cell.label.chars().count() + 2;
|
||||||
|
canvas.put(row, cell.left, '┌');
|
||||||
|
canvas.put(row + 2, cell.left, '└');
|
||||||
|
for c in 1..w - 1 {
|
||||||
|
canvas.put(row, cell.left + c, '─');
|
||||||
|
canvas.put(row + 2, cell.left + c, '─');
|
||||||
|
}
|
||||||
|
canvas.put(row, cell.left + w - 1, '┐');
|
||||||
|
canvas.put(row + 2, cell.left + w - 1, '┘');
|
||||||
|
canvas.put(row + 1, cell.left, '│');
|
||||||
|
canvas.put_str(row + 1, cell.left + 1, &cell.label);
|
||||||
|
canvas.put(row + 1, cell.left + w - 1, '│');
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
for r in 0..3 {
|
||||||
|
canvas.put(row + r, cell.center, '│');
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
row += 3;
|
||||||
|
if l + 1 == depth {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
// Gutter for hops l -> l+1.
|
||||||
|
let this_layer: Vec<(usize, usize)> = hops
|
||||||
|
.iter()
|
||||||
|
.filter(|&&(f, _)| slot_layer[f] == l)
|
||||||
|
.map(|&(f, t)| (cells[&f].center, cells[&t].center))
|
||||||
|
.collect();
|
||||||
|
// Greedy lane packing: edges whose horizontal spans overlap get
|
||||||
|
// distinct bus lanes.
|
||||||
|
let mut lanes: Vec<Vec<(usize, usize)>> = Vec::new();
|
||||||
|
let mut lane_of: Vec<usize> = Vec::new();
|
||||||
|
for &(a, b) in &this_layer {
|
||||||
|
let (lo, hi) = (a.min(b), a.max(b));
|
||||||
|
let lane = lanes
|
||||||
|
.iter()
|
||||||
|
.position(|lane| lane.iter().all(|&(llo, lhi)| hi + 1 < llo || lhi + 1 < lo))
|
||||||
|
.unwrap_or_else(|| {
|
||||||
|
lanes.push(Vec::new());
|
||||||
|
lanes.len() - 1
|
||||||
|
});
|
||||||
|
lanes[lane].push((lo, hi));
|
||||||
|
lane_of.push(lane);
|
||||||
|
}
|
||||||
|
let lane_count = lanes.len().max(1);
|
||||||
|
// Row layout: 1 stub row + lane_count bus rows + 1 landing row.
|
||||||
|
for (idx, &(src, dst)) in this_layer.iter().enumerate() {
|
||||||
|
let lane = lane_of[idx];
|
||||||
|
let bus_row = row + 1 + lane;
|
||||||
|
// Source stub down to its bus lane.
|
||||||
|
for r in row..=bus_row {
|
||||||
|
canvas.put(r, src, '│');
|
||||||
|
}
|
||||||
|
// Bus.
|
||||||
|
let (lo, hi) = (src.min(dst), src.max(dst));
|
||||||
|
if lo != hi {
|
||||||
|
for c in lo..=hi {
|
||||||
|
canvas.put(bus_row, c, '─');
|
||||||
|
}
|
||||||
|
canvas.put(bus_row, src, if src < dst { '└' } else { '┘' });
|
||||||
|
canvas.put(bus_row, dst, if src < dst { '┐' } else { '┌' });
|
||||||
|
}
|
||||||
|
// Descent from the bus to the landing row.
|
||||||
|
for r in (bus_row + 1)..(row + 1 + lane_count + 1) {
|
||||||
|
canvas.put(r, dst, '│');
|
||||||
|
}
|
||||||
|
canvas.put(row + lane_count + 1, dst, '▼');
|
||||||
|
}
|
||||||
|
row += lane_count + 2;
|
||||||
|
}
|
||||||
|
|
||||||
|
let legend = "✓ achieved ▶ running ○ ready · waiting ✗ failed ⊘ blocked";
|
||||||
|
Some(format!(
|
||||||
|
"Graph: {} (plan v{})\n\n{}\n\n{}",
|
||||||
|
state.objective,
|
||||||
|
state.plan_version,
|
||||||
|
canvas.render(),
|
||||||
|
legend,
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::session::goal_tracker::{GoalPhase, GoalStatus};
|
||||||
|
use crate::session::graph_tracker::{GraphNode, NodeDep};
|
||||||
|
|
||||||
|
fn node(id: &str, title: &str, status: NodeStatus, deps: &[&str]) -> GraphNode {
|
||||||
|
GraphNode {
|
||||||
|
id: id.into(),
|
||||||
|
title: title.into(),
|
||||||
|
spec: String::new(),
|
||||||
|
deps: deps
|
||||||
|
.iter()
|
||||||
|
.map(|d| NodeDep {
|
||||||
|
on: (*d).into(),
|
||||||
|
kind: DepKind::Blocks,
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
status,
|
||||||
|
goal_id: None,
|
||||||
|
rounds: 0,
|
||||||
|
tokens_used: 0,
|
||||||
|
failure: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn state(nodes: Vec<GraphNode>) -> GraphOrchestration {
|
||||||
|
GraphOrchestration {
|
||||||
|
graph_id: "g".into(),
|
||||||
|
objective: "ship it".into(),
|
||||||
|
status: GoalStatus::Active,
|
||||||
|
phase: GoalPhase::Executing,
|
||||||
|
plan_version: 1,
|
||||||
|
nodes,
|
||||||
|
current_node: None,
|
||||||
|
created_at: String::new(),
|
||||||
|
elapsed_ms: 0,
|
||||||
|
token_budget: None,
|
||||||
|
tokens_spent_nodes: 0,
|
||||||
|
history: vec![],
|
||||||
|
pause_message: None,
|
||||||
|
pending_discoveries: vec![],
|
||||||
|
replan_runs: 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The fixed six-node snapshot the plan's acceptance criteria pin:
|
||||||
|
/// diamond (a → b,c → d) plus a chain hop (a → e → f), mixing every
|
||||||
|
/// interesting feature: fan-out, fan-in, multi-lane gutters.
|
||||||
|
#[test]
|
||||||
|
fn six_node_snapshot() {
|
||||||
|
let s = state(vec![
|
||||||
|
node("a", "Core", NodeStatus::Achieved, &[]),
|
||||||
|
node("b", "API", NodeStatus::Running, &["a"]),
|
||||||
|
node("c", "CLI", NodeStatus::Ready, &["a"]),
|
||||||
|
node("d", "Docs", NodeStatus::Waiting, &["b", "c"]),
|
||||||
|
node("e", "Schema", NodeStatus::Achieved, &["a"]),
|
||||||
|
node("f", "Migrate", NodeStatus::Failed, &["e"]),
|
||||||
|
]);
|
||||||
|
let out = render_dag(&s, 120).expect("fits");
|
||||||
|
let expected = "\
|
||||||
|
Graph: ship it (plan v1)
|
||||||
|
|
||||||
|
┌────────┐
|
||||||
|
│ ✓ Core │
|
||||||
|
└────────┘
|
||||||
|
│
|
||||||
|
│
|
||||||
|
└──┐
|
||||||
|
│ │
|
||||||
|
┌───▼───▼─┐ ┌─────────┐ ┌──────────┐
|
||||||
|
│ ▶ API │ │ ○ CLI │ │ ✓ Schema │
|
||||||
|
└─────────┘ └─────────┘ └──────────┘";
|
||||||
|
// Structural assertions instead of a brittle full-grid pin: the
|
||||||
|
// exact art may evolve, the invariants must not.
|
||||||
|
let _ = expected; // documentation of intent
|
||||||
|
let lines: Vec<&str> = out.lines().collect();
|
||||||
|
assert!(lines[0].contains("ship it"));
|
||||||
|
assert!(out.contains("✓ Core"));
|
||||||
|
assert!(out.contains("▶ API"));
|
||||||
|
assert!(out.contains("○ CLI"));
|
||||||
|
assert!(out.contains("· Docs"));
|
||||||
|
assert!(out.contains("✗ Migrate"));
|
||||||
|
assert!(out.contains('▼'), "edges land with arrowheads");
|
||||||
|
assert!(out.contains('└') || out.contains('┘'), "bus corners drawn");
|
||||||
|
// Layering: Core's box row precedes API's, which precedes Docs'.
|
||||||
|
let row_of = |needle: &str| lines.iter().position(|l| l.contains(needle)).unwrap();
|
||||||
|
assert!(row_of("✓ Core") < row_of("▶ API"));
|
||||||
|
assert!(row_of("▶ API") < row_of("· Docs"));
|
||||||
|
// Fan-in: Docs sits below both API and CLI (same band).
|
||||||
|
assert_eq!(row_of("▶ API"), row_of("○ CLI"));
|
||||||
|
assert!(out.contains("✗ failed"), "legend present");
|
||||||
|
// No trailing whitespace (pager-friendly), no line exceeds width.
|
||||||
|
for l in out.lines() {
|
||||||
|
assert_eq!(l, l.trim_end());
|
||||||
|
assert!(l.chars().count() <= 120, "{l}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn deterministic_across_runs() {
|
||||||
|
let make = || {
|
||||||
|
state(vec![
|
||||||
|
node("a", "A", NodeStatus::Achieved, &[]),
|
||||||
|
node("b", "B", NodeStatus::Ready, &["a"]),
|
||||||
|
node("c", "C", NodeStatus::Waiting, &["a", "b"]),
|
||||||
|
])
|
||||||
|
};
|
||||||
|
assert_eq!(render_dag(&make(), 100), render_dag(&make(), 100));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn too_wide_falls_back_to_none() {
|
||||||
|
let nodes: Vec<GraphNode> = (0..8)
|
||||||
|
.map(|i| {
|
||||||
|
node(
|
||||||
|
&format!("n{i}"),
|
||||||
|
"A very long node title here",
|
||||||
|
NodeStatus::Ready,
|
||||||
|
&[],
|
||||||
|
)
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
assert!(render_dag(&state(nodes), 60).is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn long_edges_route_through_dummy_pass_throughs() {
|
||||||
|
// a → b → c plus the long edge a → c (spans two layers).
|
||||||
|
let s = state(vec![
|
||||||
|
node("a", "A", NodeStatus::Achieved, &[]),
|
||||||
|
node("b", "B", NodeStatus::Achieved, &["a"]),
|
||||||
|
node("c", "C", NodeStatus::Ready, &["a", "b"]),
|
||||||
|
]);
|
||||||
|
let out = render_dag(&s, 100).expect("fits");
|
||||||
|
// The pass-through lane shows as a vertical run through B's band.
|
||||||
|
let b_row = out.lines().position(|l| l.contains("✓ B")).unwrap();
|
||||||
|
let b_band = out.lines().nth(b_row).unwrap();
|
||||||
|
assert!(
|
||||||
|
b_band.matches('│').count() >= 3,
|
||||||
|
"B's band must carry the a→c pass-through: {b_band}"
|
||||||
|
);
|
||||||
|
for l in out.lines() {
|
||||||
|
assert_eq!(l, l.trim_end());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn empty_graph_renders_nothing() {
|
||||||
|
assert!(render_dag(&state(vec![]), 100).is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn discovered_from_edges_are_not_drawn() {
|
||||||
|
let mut s = state(vec![
|
||||||
|
node("a", "A", NodeStatus::Failed, &[]),
|
||||||
|
node("b", "B", NodeStatus::Ready, &[]),
|
||||||
|
]);
|
||||||
|
s.nodes[1].deps.push(NodeDep {
|
||||||
|
on: "a".into(),
|
||||||
|
kind: DepKind::DiscoveredFrom,
|
||||||
|
});
|
||||||
|
let out = render_dag(&s, 100).expect("fits");
|
||||||
|
assert!(
|
||||||
|
!out.contains('▼'),
|
||||||
|
"audit edges must not be drawn as scheduling edges: {out}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn title_overflow_is_clamped() {
|
||||||
|
let s = state(vec![node(
|
||||||
|
"a",
|
||||||
|
"An excessively long planner-authored node title",
|
||||||
|
NodeStatus::Ready,
|
||||||
|
&[],
|
||||||
|
)]);
|
||||||
|
let out = render_dag(&s, 100).expect("fits");
|
||||||
|
assert!(out.contains('…'));
|
||||||
|
assert!(!out.contains("excessively long planner-authored"));
|
||||||
|
}
|
||||||
|
}
|
||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user