refactor(auth): centralize inference-credential routing in CredentialAuthority
One authority answers 'which credential may ride this request': credential_class / manager_for / credential_for / bearer_resolver_for, keyed by (platform, base_url). SessionCredential is an opaque type with no production constructor, so a new call site cannot re-introduce the session-bearer leak. Platform-scoped tests extended across all bearer channels (session, aux, summary, subagent override). Verified: cargo check --workspace --all-targets clean; kigi-shell and kigi-tui suites green (6611+ tests).
This commit is contained in:
@@ -417,6 +417,21 @@ pub(crate) struct SessionActor {
|
||||
/// [`SessionActor::model_auth_facts`].
|
||||
pub(crate) model_auth_facts:
|
||||
std::cell::RefCell<Option<(String, crate::agent::config::ModelAuthFacts)>>,
|
||||
/// The catalog KEY this session's model was selected by (`{platform}/{model}`
|
||||
/// for a registry model), owned PER SESSION.
|
||||
///
|
||||
/// H4: `SamplingConfig::model` is the bare routing slug, and duplicate slugs
|
||||
/// across an API-key platform and its subscription-OAuth twin
|
||||
/// (`xai`/`xai-grok`, `anthropic`/`claude-pro-max`, `openai`/`openai-codex`)
|
||||
/// are BY DESIGN, so the slug alone cannot name the platform. This used to
|
||||
/// be read from `ModelsManager::current_model_id()` — a single
|
||||
/// PROCESS-GLOBAL cell that Leader mode never writes
|
||||
/// (`agent/handlers/model_switch.rs`) and that is last-writer-wins across
|
||||
/// concurrent sessions, so both collision directions resolved the wrong
|
||||
/// platform: the subscription session lost its resolver (unrecoverable 401
|
||||
/// ~1h in) and the API-key session got the pooled OAuth bearer stamped over
|
||||
/// its own `sk-…` key. Written at spawn and on every `SetSessionModel`.
|
||||
pub(crate) selected_catalog_key: std::cell::RefCell<Option<String>>,
|
||||
/// 401-attribution callback. Joined with the bearer the
|
||||
/// sampler sends on the wire to emit an `auth 401 attribution`
|
||||
/// event at each of the six `OaiCompatClient` 401 arms in
|
||||
|
||||
@@ -5,12 +5,18 @@ impl SessionActor {
|
||||
pub(super) async fn handle_set_session_model(
|
||||
&self,
|
||||
sampling_config: kigi_sampler::SamplerConfig,
|
||||
catalog_key: Option<String>,
|
||||
use_concise: bool,
|
||||
apply_prompt_override: bool,
|
||||
skip_prompt_rewrite: bool,
|
||||
auto_compact_threshold_percent: u8,
|
||||
) -> Result<acp::ModelId, acp::Error> {
|
||||
let model_id = acp::ModelId::new(sampling_config.model.clone());
|
||||
// H4: record the picker's catalog KEY as this SESSION's own selection.
|
||||
// `sampling_config.model` is the ambiguous bare slug; the key is what
|
||||
// disambiguates an API-key platform from its subscription-OAuth twin,
|
||||
// and it must never come from the process-global `current_model_id()`.
|
||||
*self.selected_catalog_key.borrow_mut() = catalog_key;
|
||||
let new_context_window = self.compaction.context_window_override.unwrap_or_else(|| {
|
||||
std::num::NonZeroU64::new(sampling_config.context_window).unwrap_or_else(|| {
|
||||
std::num::NonZeroU64::new(DEFAULT_CONTEXT_WINDOW)
|
||||
@@ -64,15 +70,13 @@ impl SessionActor {
|
||||
// grok model reads the xai-grok token (used only to classify the
|
||||
// credential's auth_type here), never the Kimi one. Kimi / non-oauth
|
||||
// models resolve to the primary — byte-identical.
|
||||
let session_key = self
|
||||
.auth_manager_for_model(&sampling_config.model)
|
||||
.and_then(|am| am.current_or_expired().map(|a| a.key));
|
||||
let session_key = self.session_credential_for_model(&sampling_config.model);
|
||||
self.chat_state_handle
|
||||
.update_credentials(kigi_chat_state::Credentials {
|
||||
api_key: sampling_config.api_key.clone(),
|
||||
auth_type: crate::agent::config::resolve_chat_state_auth_type(
|
||||
sampling_config.model.as_str(),
|
||||
session_key.as_deref(),
|
||||
session_key.as_ref(),
|
||||
existing.auth_type,
|
||||
),
|
||||
alpha_test_key: existing.alpha_test_key,
|
||||
|
||||
@@ -624,23 +624,22 @@ impl SessionActor {
|
||||
let resolved_describe = self
|
||||
.resolve_aux_sampler_config(&self.image_description_model)
|
||||
.await;
|
||||
// LEAK 1b: only re-point the aux bearer_resolver when the aux model
|
||||
// actually resolved (Some) — the `None` fallback yields the SESSION
|
||||
// config, whose Kimi resolver must stay as-is.
|
||||
let aux_resolved = resolved_describe.is_some();
|
||||
let (describe_model, mut sampler_config) =
|
||||
// LEAK 1b: the aux bearer_resolver is decided at the chokepoint from the
|
||||
// IMAGE-DESCRIBE model's own platform + endpoint and passed in
|
||||
// explicitly, so an aux model on another provider can never inherit the
|
||||
// session (Kimi) resolver and have its own key overwritten on the aux
|
||||
// host. The `None` fallback yields the SESSION config verbatim, whose
|
||||
// own resolver must stay as-is.
|
||||
let describe_resolver = resolved_describe
|
||||
.as_ref()
|
||||
.map(|cfg| self.aux_bearer_resolver(&self.image_description_model, &cfg.base_url));
|
||||
let (describe_model, sampler_config) =
|
||||
crate::agent::config::finalize_image_describe_sampler_config(
|
||||
resolved_describe,
|
||||
&active_session_config,
|
||||
describe_resolver.flatten(),
|
||||
Some(self.max_retries),
|
||||
);
|
||||
// An image-describe model on another provider must not inherit the
|
||||
// session (Kimi) bearer_resolver stamped by `finalize_*`: re-point it at
|
||||
// an OAuth model's own manager, or clear it for an API-key-platform /
|
||||
// third-party endpoint. No-op for the first-party subscription channel.
|
||||
if aux_resolved {
|
||||
self.repoint_aux_bearer_resolver(&mut sampler_config, &self.image_description_model);
|
||||
}
|
||||
let client = kigi_sampler::SamplingClient::new(sampler_config).map_err(|e| {
|
||||
acp::Error::internal_error().data(format!(
|
||||
"failed to build image-describe sampling client: {e}"
|
||||
|
||||
@@ -300,10 +300,10 @@ pub(super) async fn run_session(
|
||||
SessionActor::maybe_start_running_task(session.clone(), completion_tx
|
||||
.clone()). await; } SessionCommand::SessionMode { session_mode, responds_to }
|
||||
=> { session.handle_session_mode(session_mode). await; let _ = responds_to
|
||||
.send(()); } SessionCommand::SetSessionModel { sampling_config, use_concise,
|
||||
apply_prompt_override, skip_prompt_rewrite, auto_compact_threshold_percent,
|
||||
responds_to } => { let updated_model_id = session
|
||||
.handle_set_session_model(sampling_config, use_concise,
|
||||
.send(()); } SessionCommand::SetSessionModel { sampling_config, catalog_key,
|
||||
use_concise, apply_prompt_override, skip_prompt_rewrite,
|
||||
auto_compact_threshold_percent, responds_to } => { let updated_model_id =
|
||||
session.handle_set_session_model(sampling_config, catalog_key, use_concise,
|
||||
apply_prompt_override, skip_prompt_rewrite, auto_compact_threshold_percent).
|
||||
await; let _ = responds_to.send(updated_model_id); }
|
||||
SessionCommand::RebuildAgentForDefinition { definition, responds_to } => {
|
||||
@@ -319,11 +319,24 @@ pub(super) async fn run_session(
|
||||
.signals_handle().set_primary_model(& model_name); cfg.model = model_name
|
||||
.clone(); cfg.extra_headers.extend(extra_headers); if let Some(cw) =
|
||||
context_window && session.compaction.context_window_override.is_none() { cfg
|
||||
.context_window = cw; } session.chat_state_handle
|
||||
.context_window = cw; } let override_base_url = cfg
|
||||
.base_url.clone(); session.chat_state_handle
|
||||
.update_sampling_config(cfg); let existing = session.chat_state_handle
|
||||
.get_credentials(). await; if let Some(r) = crate
|
||||
::agent::config::try_resolve_model_credentials(model_name.as_str(), existing
|
||||
.api_key.as_deref()) { session.chat_state_handle
|
||||
.get_credentials(). await;
|
||||
// H-c: the rename makes the session's own selected catalog
|
||||
// key stale unless it still names this model; a stale key is
|
||||
// exactly what the model→platform rule must not trust.
|
||||
session.retain_selected_catalog_key_for(& model_name);
|
||||
// The override model routes to the SAME endpoint the session
|
||||
// already had; ask the chokepoint whether that endpoint takes
|
||||
// a session credential rather than re-offering the key
|
||||
// already in chat state. The platform comes from the session's
|
||||
// OWN lookup so this and every later turn agree.
|
||||
let override_session_key = session.credential_authority()
|
||||
.credential_for(session.model_platform(model_name.as_str()), &
|
||||
override_base_url); if let Some(r) = crate
|
||||
::agent::config::try_resolve_model_credentials(model_name.as_str(),
|
||||
override_session_key.as_ref()) { session.chat_state_handle
|
||||
.update_credentials(kigi_chat_state::Credentials { api_key : r.api_key,
|
||||
auth_type : r.auth_type, alpha_test_key : existing.alpha_test_key, }); } session.model_auth_facts
|
||||
.replace(None); } } SessionCommand::GetCurrentModel { responds_to } => { let
|
||||
|
||||
@@ -29,55 +29,105 @@ pub(super) fn is_auth_tool_error(err: &kigi_tool_runtime::ToolError) -> bool {
|
||||
/// Gate inputs bundled with the composed decision so the 401-recovery log can
|
||||
/// report the components.
|
||||
#[derive(Clone, Copy)]
|
||||
struct SessionTokenAuthGate {
|
||||
pub(crate) struct SessionTokenAuthGate {
|
||||
is_session_based: bool,
|
||||
model_byok: crate::agent::auth_method::ModelByok,
|
||||
/// Whether the request targets a first-party host. Lets an `Unknown`
|
||||
/// BYOK status still refresh against the first-party cli-chat-proxy hosts without
|
||||
/// risking a session-token leak to a third-party BYOK endpoint.
|
||||
endpoint_is_first_party: bool,
|
||||
/// Whether this model's platform is one whose endpoint accepts a session
|
||||
/// bearer at all (see
|
||||
/// [`crate::agent::auth_method::platform_takes_session_credential`]). False
|
||||
/// for every API-key registry platform, which keeps the primary Kimi bearer
|
||||
/// off `api.deepseek.com` / `api.openai.com` / … .
|
||||
endpoint_takes_session_credential: bool,
|
||||
/// WHICH credential this model's platform/endpoint pair accepts, per the
|
||||
/// single credential chokepoint
|
||||
/// ([`crate::auth::credential_authority::CredentialAuthority::credential_class`]).
|
||||
/// `None` for every API-key registry platform, which keeps the primary Kimi
|
||||
/// bearer off `api.deepseek.com` / `api.openai.com` / … .
|
||||
credential_class: crate::auth::credential_authority::CredentialClass,
|
||||
}
|
||||
impl SessionTokenAuthGate {
|
||||
/// Single place `is_session_based` / `endpoint_is_first_party` are derived,
|
||||
/// so all call sites assemble the gate identically. `model_platform` is the
|
||||
/// registry platform the model routes to (`None` for a bare / `[model.*]`
|
||||
/// entry) — it MUST be derived from the same lookup
|
||||
/// ([`SessionActor::managed_key_for_model`]) that
|
||||
/// [`SessionActor::auth_manager_for_model`] uses, so the gate's verdict and
|
||||
/// the manager actually wrapped as the bearer resolver can never disagree.
|
||||
fn new(
|
||||
/// ([`SessionActor::model_platform`]) that
|
||||
/// [`SessionActor::auth_manager_for_endpoint`] uses, so the gate's verdict
|
||||
/// and the manager actually wrapped as the bearer resolver can never
|
||||
/// disagree. `authority` is that same chokepoint, so the gate cannot answer
|
||||
/// the endpoint question differently from the manager routing.
|
||||
pub(crate) fn new(
|
||||
auth_method_id: Option<&acp::AuthMethodId>,
|
||||
model_byok: crate::agent::auth_method::ModelByok,
|
||||
base_url: &str,
|
||||
model_platform: Option<kigi_models::PlatformId>,
|
||||
authority: &crate::auth::credential_authority::CredentialAuthority,
|
||||
) -> Self {
|
||||
Self {
|
||||
// L13: a model whose OWN credential is a pooled subscription-OAuth
|
||||
// session is session-based BY ITSELF, whatever the primary ACP
|
||||
// method is. A user logged in with an API-KEY platform (e.g.
|
||||
// `deepseek`) who selects a `claude-pro-max/*` model still gets that
|
||||
// platform's pooled bearer as the request's `api_key` — without this
|
||||
// term the gate would be inactive, so the config would carry NO
|
||||
// resolver: the token freezes at selection time and the session dies
|
||||
// with an unrecoverable 401 once it expires (~1h). The outer
|
||||
// `credential_class` conjunct keeps this confined to that
|
||||
// platform's own host.
|
||||
is_session_based: auth_method_id
|
||||
.is_some_and(crate::agent::auth_method::is_session_based_method),
|
||||
.is_some_and(crate::agent::auth_method::is_session_based_method)
|
||||
|| model_platform.is_some_and(|p| p.oauth().is_some()),
|
||||
model_byok,
|
||||
endpoint_is_first_party: crate::util::is_first_party_url(base_url),
|
||||
endpoint_takes_session_credential:
|
||||
crate::agent::auth_method::platform_takes_session_credential(
|
||||
model_platform,
|
||||
base_url,
|
||||
),
|
||||
credential_class: authority.credential_class(model_platform, base_url),
|
||||
}
|
||||
}
|
||||
fn active(self) -> bool {
|
||||
pub(crate) fn active(self) -> bool {
|
||||
crate::agent::auth_method::session_token_auth_gate(
|
||||
self.is_session_based,
|
||||
self.model_byok,
|
||||
self.endpoint_is_first_party,
|
||||
self.endpoint_takes_session_credential,
|
||||
self.credential_class,
|
||||
)
|
||||
}
|
||||
}
|
||||
/// THE aux / summary `bearer_resolver` rule, stated ONCE.
|
||||
///
|
||||
/// `SamplingClient::post` REPLACES the request's auth header from the resolver,
|
||||
/// so an aux model on a different provider would have its own correctly-resolved
|
||||
/// key overwritten by the session bearer ON THE AUX HOST. An OAuth aux model
|
||||
/// gets a live resolver over ITS OWN pooled manager (keeping mid-session
|
||||
/// refresh); a first-party aux model gets the primary's, but ONLY when the
|
||||
/// session-token gate is active; everything else gets `None`, so the aux model's
|
||||
/// own key survives to the wire.
|
||||
///
|
||||
/// M3 — the FIRST-PARTY case honours the gate, which is what the old "copy
|
||||
/// `active_session_config.bearer_resolver`" shape did implicitly: that field is
|
||||
/// `None` whenever the gate is inactive. Without it, a BYOK / api-key session
|
||||
/// with a `[model.*]` aux entry carrying its OWN `env_key` on the session's own
|
||||
/// coding endpoint has that key REPLACED on the wire by the primary bearer on
|
||||
/// every image-describe / auto-mode-classifier / summary request. A
|
||||
/// subscription-OAuth aux model is deliberately NOT gated this way: its pooled
|
||||
/// token IS its credential, and withholding the resolver only costs it
|
||||
/// mid-session refresh (L13).
|
||||
///
|
||||
/// Shared by [`SessionActor::aux_bearer_resolver`] and
|
||||
/// `MvpAgent::summary_bearer_resolver`: the summary client is built by the
|
||||
/// AGENT, not the session actor, and its own private copy of this rule is
|
||||
/// exactly how it stayed ungated after M3 closed the session-actor side.
|
||||
pub(crate) fn aux_bearer_resolver_for(
|
||||
authority: &crate::auth::credential_authority::CredentialAuthority,
|
||||
auth_method_id: Option<&acp::AuthMethodId>,
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
model_byok: crate::agent::auth_method::ModelByok,
|
||||
base_url: &str,
|
||||
) -> Option<kigi_sampler::SharedBearerResolver> {
|
||||
let is_primary_channel = platform.is_none_or(|p| p.oauth().is_none());
|
||||
if is_primary_channel
|
||||
&& !SessionTokenAuthGate::new(auth_method_id, model_byok, base_url, platform, authority)
|
||||
.active()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
authority.bearer_resolver_for(platform, base_url)
|
||||
}
|
||||
/// Run a tool call; on an auth-shaped failure, attempt recovery via
|
||||
/// `AuthManager` and one retry. When `shared_recovery` is `Some`, concurrent
|
||||
/// 401s in the same batch deduplicate via `OnceCell::get_or_init`.
|
||||
@@ -125,7 +175,8 @@ where
|
||||
/// [`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 repoint ([`SessionActor::repoint_aux_bearer_resolver_for_oauth`])
|
||||
/// 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>);
|
||||
@@ -140,48 +191,11 @@ impl kigi_sampler::BearerResolver for AuthManagerBearerResolver {
|
||||
}
|
||||
}
|
||||
/// Wrap `am` as a shared sampler bearer resolver.
|
||||
fn auth_manager_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))
|
||||
}
|
||||
/// The `bearer_resolver` an AUX / summary `SamplerConfig` may carry, given the
|
||||
/// SESSION model's resolver (`stamped`) and the AUX model's own platform +
|
||||
/// endpoint. ONE decision shared by image-describe, the auto-mode classifier
|
||||
/// (via [`SessionActor::repoint_aux_bearer_resolver`]) and
|
||||
/// `MvpAgent::build_summary_client`.
|
||||
///
|
||||
/// [`crate::agent::config::stamp_session_local_sampler_fields`] copies the
|
||||
/// session resolver onto every aux config, and `SamplingClient::post` REPLACES
|
||||
/// the request's auth header from it — so an aux model on a DIFFERENT provider
|
||||
/// would have its own correctly-resolved key overwritten by the session bearer
|
||||
/// on the AUX host (H3/H4). Resolve by the aux model instead:
|
||||
/// - an OAuth platform → a live resolver over ITS OWN pooled manager (so a grok
|
||||
/// / claude-pro-max / copilot / codex aux model keeps mid-session refresh);
|
||||
/// - `kimi-code`, or a platform-less model on the session's own coding endpoint
|
||||
/// → the stamped session resolver, byte-identical;
|
||||
/// - every API-key registry platform, and any `[model.*]` block pointed at a
|
||||
/// third-party host → `None`, so the aux model's own key survives to the wire.
|
||||
///
|
||||
/// SECURITY: no token is logged.
|
||||
pub(crate) fn aux_bearer_resolver(
|
||||
stamped: Option<kigi_sampler::SharedBearerResolver>,
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
base_url: &str,
|
||||
) -> Option<kigi_sampler::SharedBearerResolver> {
|
||||
if let Some(oauth) = platform.and_then(kigi_models::PlatformId::oauth) {
|
||||
return Some(auth_manager_bearer_resolver(
|
||||
crate::auth::oauth_registry::global_manager_for(
|
||||
&crate::auth::oauth_registry::pool_home(),
|
||||
oauth,
|
||||
),
|
||||
));
|
||||
}
|
||||
if crate::agent::auth_method::platform_takes_session_credential(platform, base_url) {
|
||||
return stamped;
|
||||
}
|
||||
None
|
||||
}
|
||||
impl SessionActor {
|
||||
pub(super) async fn prepare_tool_definitions_timed(&self) -> (Vec<ToolDefinition>, u64) {
|
||||
let mcp_wait_start = std::time::Instant::now();
|
||||
@@ -225,8 +239,8 @@ impl SessionActor {
|
||||
let plan_active = self.plan_mode.lock().is_active();
|
||||
filter_cursor_tools_by_plan_mode(defs, plan_active)
|
||||
}
|
||||
/// Memoized per-model [`ModelAuthFacts`](crate::agent::config::ModelAuthFacts),
|
||||
/// keyed by `model_id`.
|
||||
/// Memoized per-model [`ModelAuthFacts`](crate::agent::config::ModelAuthFacts)
|
||||
/// for the SESSION's own model, keyed by `model_id`.
|
||||
///
|
||||
/// A fresh `Unknown` (config currently unparseable) falls back to the last
|
||||
/// definite value for the same `model_id` rather than demoting a live session
|
||||
@@ -235,6 +249,32 @@ impl SessionActor {
|
||||
/// `model_id`, keying on `model_id` alone is insufficient — each
|
||||
/// model/credential chokepoint must clear this memo (`replace(None)`).
|
||||
pub(super) fn model_auth_facts(&self, model_id: &str) -> crate::agent::config::ModelAuthFacts {
|
||||
self.resolve_auth_facts(model_id, true)
|
||||
}
|
||||
/// [`Self::model_auth_facts`] for a model that is NOT the session's own — an
|
||||
/// AUX / summary / image-describe slug.
|
||||
///
|
||||
/// Identical resolution, but it NEVER WRITES the slot. The memo is a SINGLE
|
||||
/// slot: when the aux path shared it, one classifier or image-describe call
|
||||
/// evicted the session model's entry, and (a) the next
|
||||
/// [`Self::reconstruct_full_config`] paid another `load_effective_config()`
|
||||
/// + `resolve_model_list()` — the per-turn disk read M7/M9 removed — while
|
||||
/// (b) a transient `Unknown` for the SESSION model then had no same-`model_id`
|
||||
/// definite value to fall back to, so it degraded to `endpoint_is_first_party`
|
||||
/// — `false` for every subscription-OAuth host, costing the session its
|
||||
/// `bearer_resolver` and 401ing unrecoverably ~1h in (the failure L13
|
||||
/// prevents). Reading a matching entry is still allowed: it can only hit when
|
||||
/// the slot already names this same slug.
|
||||
fn aux_model_auth_facts(&self, model_id: &str) -> crate::agent::config::ModelAuthFacts {
|
||||
self.resolve_auth_facts(model_id, false)
|
||||
}
|
||||
/// Shared body of [`Self::model_auth_facts`] / [`Self::aux_model_auth_facts`].
|
||||
/// `memoize` is the ONLY difference, so the two can never resolve differently.
|
||||
fn resolve_auth_facts(
|
||||
&self,
|
||||
model_id: &str,
|
||||
memoize: bool,
|
||||
) -> crate::agent::config::ModelAuthFacts {
|
||||
use crate::agent::auth_method::ModelByok;
|
||||
if let Some((cached_id, facts)) = self.model_auth_facts.borrow().as_ref()
|
||||
&& cached_id == model_id
|
||||
@@ -251,10 +291,12 @@ impl SessionActor {
|
||||
}
|
||||
return fresh;
|
||||
}
|
||||
*self.model_auth_facts.borrow_mut() = Some((model_id.to_string(), fresh));
|
||||
if memoize {
|
||||
*self.model_auth_facts.borrow_mut() = Some((model_id.to_string(), fresh));
|
||||
}
|
||||
fresh
|
||||
}
|
||||
/// Gate inputs for `model_id` routed to `base_url`. See
|
||||
/// Gate inputs for the SESSION model `model_id` routed to `base_url`. See
|
||||
/// [`crate::agent::auth_method::session_token_auth_gate`] for the rationale
|
||||
/// (`base_url` keeps an `Unknown` BYOK status refreshable only
|
||||
/// against first-party xAI hosts).
|
||||
@@ -266,62 +308,109 @@ impl SessionActor {
|
||||
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 model whose routing slug is `model` (the sampling config's
|
||||
/// `model`). A generic device-code OAuth platform (xai-grok) routes to its
|
||||
/// OWN scope-keyed manager from the process-global OAuth pool
|
||||
/// ([`crate::auth::oauth_registry::manager_for_model`], built on demand from
|
||||
/// the on-disk token); every other model routes to the primary Kimi
|
||||
/// `auth_manager`.
|
||||
/// for the routing slug `model` against the endpoint the request will
|
||||
/// ACTUALLY be sent to, from the ONE chokepoint
|
||||
/// ([`crate::auth::credential_authority::CredentialAuthority`]).
|
||||
///
|
||||
/// This is the single chokepoint that keeps a grok turn from ever sending
|
||||
/// the Kimi bearer (Facet B) and gives it its own proactive-refresh + 401
|
||||
/// recovery manager (Facet A). The pool is the single source of truth, so a
|
||||
/// grok login that lands AFTER this session spawned is resolved on the next
|
||||
/// grok turn (no frozen per-session snapshot). Cheap `Arc` clone. Kimi /
|
||||
/// first-party path is byte-identical: a non-oauth model always resolves to
|
||||
/// the primary.
|
||||
/// 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.
|
||||
///
|
||||
/// `None` when there is no governing manager: a BYOK / test session with no
|
||||
/// primary. A grok model always resolves to its pooled manager (never the
|
||||
/// Kimi primary); when the user has not logged into grok that manager simply
|
||||
/// holds no token, so no Kimi bearer can leak.
|
||||
pub(super) fn auth_manager_for_model(
|
||||
/// 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>> {
|
||||
// `model` is the bare routing slug; recover the managed catalog key
|
||||
// (`{platform}/{model}`) so the platform — and thus its OAuth scope — is
|
||||
// unambiguous. A bare / config / unlisted model yields no managed key
|
||||
// and resolves to the primary.
|
||||
let managed_key = self.managed_key_for_model(model);
|
||||
crate::auth::oauth_registry::manager_for_model(
|
||||
&crate::auth::oauth_registry::pool_home(),
|
||||
managed_key.as_deref().unwrap_or(model),
|
||||
self.auth_manager.as_ref(),
|
||||
)
|
||||
self.credential_authority()
|
||||
.manager_for(self.model_platform(model), base_url)
|
||||
}
|
||||
/// Recover the managed catalog key (`{platform}/{model}`) for a routing slug
|
||||
/// from the live catalog. `None` for a bare / config / unlisted model.
|
||||
///
|
||||
/// H5: the catalog KEY the picker selected
|
||||
/// ([`crate::agent::models::ModelsManager::current_model_id`]) is
|
||||
/// authoritative — `model` is the ambiguous bare slug. See
|
||||
/// [`crate::agent::models::managed_key_for_slug`].
|
||||
fn managed_key_for_model(&self, model: &str) -> Option<String> {
|
||||
/// 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();
|
||||
let current = self.models_manager.current_model_id();
|
||||
crate::agent::models::managed_key_for_slug(&models, Some(current.0.as_ref()), model)
|
||||
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_model`] routes on. `None` for a bare /
|
||||
/// lookup [`Self::auth_manager_for_endpoint`] routes on. `None` for a bare /
|
||||
/// `[model.*]` / unlisted model.
|
||||
fn model_platform(&self, model: &str) -> Option<kigi_models::PlatformId> {
|
||||
pub(super) fn model_platform(&self, model: &str) -> Option<kigi_models::PlatformId> {
|
||||
let models = self.models_manager.models();
|
||||
let current = self.models_manager.current_model_id();
|
||||
crate::agent::models::platform_for_slug(&models, Some(current.0.as_ref()), model)
|
||||
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
|
||||
@@ -352,20 +441,37 @@ impl SessionActor {
|
||||
self.model_platform(model)
|
||||
.is_some_and(kigi_models::PlatformId::sends_codex_responses_headers)
|
||||
}
|
||||
/// LEAK guard for the stamped aux paths (auto-mode classifier, image
|
||||
/// describe). Applies [`aux_bearer_resolver`] to the config
|
||||
/// [`crate::agent::config::stamp_session_local_sampler_fields`] just stamped
|
||||
/// the SESSION model's `bearer_resolver` onto.
|
||||
pub(super) fn repoint_aux_bearer_resolver(
|
||||
/// 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,
|
||||
cfg: &mut kigi_sampler::SamplerConfig,
|
||||
slug: &str,
|
||||
) {
|
||||
cfg.bearer_resolver = aux_bearer_resolver(
|
||||
cfg.bearer_resolver.take(),
|
||||
self.model_platform(slug),
|
||||
&cfg.base_url,
|
||||
);
|
||||
base_url: &str,
|
||||
) -> Option<kigi_sampler::SharedBearerResolver> {
|
||||
let auth_method = self.auth_method_id.load();
|
||||
// An aux slug is NOT the session's selection, so it must not be resolved
|
||||
// against `selected_catalog_key` — the same rule the aux `api_key` obeys
|
||||
// (`credential_for_slug(.., None, ..)`). Keying an aux model on the
|
||||
// SESSION's selection let a colliding same-vendor slug resolve the OAuth
|
||||
// twin, whose pooled resolver would then overwrite the user's own key on
|
||||
// the aux request.
|
||||
let models = self.models_manager.models();
|
||||
aux_bearer_resolver_for(
|
||||
&self.credential_authority(),
|
||||
auth_method.as_deref(),
|
||||
crate::agent::models::platform_for_slug(&models, None, slug),
|
||||
self.aux_model_auth_facts(slug).byok,
|
||||
base_url,
|
||||
)
|
||||
}
|
||||
/// Emit a unified-log breadcrumb whenever the session-token refresh gate is
|
||||
/// evaluated with an **`Unknown`** per-model BYOK status on a session-based
|
||||
@@ -385,9 +491,8 @@ impl SessionActor {
|
||||
let ctx = serde_json::json!(
|
||||
{ "site" : site, "model_byok" : gate.model_byok.as_str(), "is_session_based"
|
||||
: gate.is_session_based, "endpoint_is_first_party" : gate
|
||||
.endpoint_is_first_party, "endpoint_takes_session_credential" : gate
|
||||
.endpoint_takes_session_credential, "refresh_active" : refresh_active,
|
||||
"base_url" : base_url, }
|
||||
.endpoint_is_first_party, "credential_class" : gate.credential_class
|
||||
.as_str(), "refresh_active" : refresh_active, "base_url" : base_url, }
|
||||
);
|
||||
let sid = Some(self.session_info.id.0.as_ref());
|
||||
if refresh_active {
|
||||
@@ -446,6 +551,7 @@ impl SessionActor {
|
||||
model_facts.byok,
|
||||
&cfg.base_url,
|
||||
self.model_platform(cfg.model.as_str()),
|
||||
&self.credential_authority(),
|
||||
);
|
||||
let use_bearer_resolver = gate.active();
|
||||
self.log_auth_gate_unknown("reconstruct_full_config", gate, &cfg.base_url);
|
||||
@@ -455,7 +561,7 @@ impl SessionActor {
|
||||
// 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_model(&cfg.model)
|
||||
self.auth_manager_for_endpoint(&cfg.model, &cfg.base_url)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
@@ -656,17 +762,20 @@ impl SessionActor {
|
||||
// Kimi session token (which `resolve_credentials` would otherwise stamp
|
||||
// onto an api.x.ai / api.deepseek.com request). The first-party
|
||||
// subscription channel still gets the primary (byte-identical).
|
||||
let session_key = crate::auth::oauth_registry::session_key_for_catalog_model(
|
||||
&models,
|
||||
slug,
|
||||
self.auth_manager.as_ref(),
|
||||
);
|
||||
// M6: ONE lookup. The platform AND the base URL the rule is applied to
|
||||
// both come from `credential_for_slug`'s single resolution of `slug`
|
||||
// against this catalog, so the platform and the endpoint can no longer
|
||||
// disagree (they were previously resolved by two different lookups).
|
||||
// Aux slugs are not the session's selection, so no `current_key`.
|
||||
let session_key = self
|
||||
.credential_authority()
|
||||
.credential_for_slug(&models, None, slug);
|
||||
let endpoints = self.models_manager.endpoints();
|
||||
crate::agent::config::resolve_aux_model_sampling_config(
|
||||
slug,
|
||||
&models,
|
||||
&endpoints,
|
||||
session_key.as_deref(),
|
||||
session_key.as_ref(),
|
||||
creds.alpha_test_key.clone(),
|
||||
)
|
||||
}
|
||||
@@ -681,16 +790,17 @@ impl SessionActor {
|
||||
) -> Option<(kigi_sampler::SamplingClient, String)> {
|
||||
let active_session_config = self.reconstruct_full_config().await;
|
||||
let mut cfg = self.resolve_aux_sampler_config(slug).await?;
|
||||
// LEAK 1b: the aux classifier must NOT inherit the SESSION model's
|
||||
// (Kimi) bearer_resolver — the resolver is decided by the AUX model's
|
||||
// own platform + endpoint at the chokepoint and passed in explicitly,
|
||||
// so there is no "copy then remember to re-point" step to forget.
|
||||
let aux_resolver = self.aux_bearer_resolver(slug, &cfg.base_url);
|
||||
crate::agent::config::stamp_session_local_sampler_fields(
|
||||
&mut cfg,
|
||||
&active_session_config,
|
||||
aux_resolver,
|
||||
Some(self.max_retries),
|
||||
);
|
||||
// LEAK 1b: the aux classifier must not inherit the SESSION model's
|
||||
// (Kimi) bearer_resolver stamped above — re-point it at an OAuth aux
|
||||
// model's own manager, or clear it for an API-key-platform / third-party
|
||||
// aux endpoint. No-op for the first-party subscription channel.
|
||||
self.repoint_aux_bearer_resolver(&mut cfg, slug);
|
||||
let model = cfg.model.clone();
|
||||
let client = kigi_sampler::SamplingClient::new(cfg)
|
||||
.map_err(|e| {
|
||||
@@ -836,8 +946,7 @@ impl SessionActor {
|
||||
session_id = % self.session_info.id.0, is_session_based = gate
|
||||
.is_session_based, model_byok = gate.model_byok.as_str(),
|
||||
endpoint_is_first_party = gate.endpoint_is_first_party,
|
||||
endpoint_takes_session_credential = gate
|
||||
.endpoint_takes_session_credential,
|
||||
credential_class = gate.credential_class.as_str(),
|
||||
"auth recovery: sampler 401 not refreshable (api-key auth) — surfacing 401",
|
||||
);
|
||||
kigi_log::unified_log::warn(
|
||||
@@ -848,8 +957,7 @@ impl SessionActor {
|
||||
.status_code, "is_session_based" : gate.is_session_based,
|
||||
"model_byok" : gate.model_byok.as_str(),
|
||||
"endpoint_is_first_party" : gate.endpoint_is_first_party,
|
||||
"endpoint_takes_session_credential" : gate
|
||||
.endpoint_takes_session_credential, }
|
||||
"credential_class" : gate.credential_class.as_str(), }
|
||||
)),
|
||||
);
|
||||
}
|
||||
@@ -869,13 +977,13 @@ impl SessionActor {
|
||||
// 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 = self
|
||||
let (recovery_model, recovery_base_url) = self
|
||||
.chat_state_handle
|
||||
.get_sampling_config()
|
||||
.await
|
||||
.map(|c| c.model)
|
||||
.map(|c| (c.model, c.base_url))
|
||||
.unwrap_or_default();
|
||||
if let Some(am) = self.auth_manager_for_model(&recovery_model) {
|
||||
if let Some(am) = self.auth_manager_for_endpoint(&recovery_model, &recovery_base_url) {
|
||||
if am.try_recover_unauthorized().await {
|
||||
tracing::info!(
|
||||
session_id = % self.session_info.id.0,
|
||||
@@ -1070,7 +1178,7 @@ impl SessionActor {
|
||||
// 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_model(&model_id) {
|
||||
if let Some(am) = self.auth_manager_for_endpoint(&model_id, &base_url) {
|
||||
let creds = self.chat_state_handle.get_credentials().await;
|
||||
if self.auth_gate(&model_id, &base_url).active()
|
||||
&& let Ok(key) = am.get_valid_token().await
|
||||
@@ -1102,13 +1210,23 @@ impl SessionActor {
|
||||
.map(|c| c.model)
|
||||
.unwrap_or_default();
|
||||
let Some(ref key) = current_key else { return };
|
||||
// M7: a registry-platform model's key comes from that platform's
|
||||
// credential resolved into its catalog entry — it is NEVER a
|
||||
// `[model.*]` block. With the session gate now inactive for every
|
||||
// API-key platform, those turns all fell through to here and paid a
|
||||
// `load_effective_config()` disk read PER TURN, then logged a
|
||||
// permanently false "Model not found in config.toml [model.*]" warning.
|
||||
if self.model_platform(¤t_model_id).is_some() {
|
||||
// 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 {
|
||||
@@ -1125,6 +1243,15 @@ impl SessionActor {
|
||||
creds.api_key = Some(new_key);
|
||||
self.chat_state_handle.update_credentials(creds);
|
||||
}
|
||||
/// Whether the live catalog entry for `slug` carries its own credential —
|
||||
/// an `api_key`/`env_key` from a `[model.*]` block, which a config edit can
|
||||
/// rotate mid-session. A platform entry whose key came from the platform
|
||||
/// credential has none.
|
||||
fn model_has_own_credential(&self, slug: &str) -> bool {
|
||||
let models = self.models_manager.models();
|
||||
crate::agent::config::find_model_by_id(&models, slug)
|
||||
.is_some_and(crate::agent::config::ModelEntry::has_own_credentials)
|
||||
}
|
||||
fn reload_api_key_from_config(&self, current_model_id: &str) -> Option<String> {
|
||||
let raw_config = crate::config::load_effective_config()
|
||||
.map_err(|e| tracing::warn!(error = % e, "Failed to reload config"))
|
||||
@@ -1211,8 +1338,8 @@ mod bearer_resolver_tests {
|
||||
/// 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
|
||||
/// `repoint_aux_bearer_resolver_for_oauth` fix wraps exactly this grok
|
||||
/// manager for a grok aux model.
|
||||
/// [`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();
|
||||
|
||||
@@ -980,10 +980,21 @@ pub(crate) async fn spawn_session_actor(
|
||||
}
|
||||
};
|
||||
let doom_loop_recovery = effective_config.resolve_doom_loop_recovery();
|
||||
let session_model_id_for_actor = session_model_id.clone();
|
||||
let session = Arc::new_cyclic(|weak: &std::sync::Weak<SessionActor>| SessionActor {
|
||||
session_info: session_info.clone(),
|
||||
auth_method_id,
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
// H4: seed the session's OWN selected catalog key from the model it was
|
||||
// spawned with, resolved through the picker's lookup. Never the
|
||||
// process-global `current_model_id()`. H-c: the rule lives in
|
||||
// `selected_catalog_key_for_spawn` so it is covered by a test.
|
||||
selected_catalog_key: std::cell::RefCell::new(
|
||||
crate::agent::models::selected_catalog_key_for_spawn(
|
||||
&models_manager.models(),
|
||||
&session_model_id_for_actor,
|
||||
),
|
||||
),
|
||||
attribution_callback,
|
||||
auth_manager,
|
||||
state,
|
||||
|
||||
+32
-17
@@ -485,29 +485,44 @@ async fn no_legacy_hint_for_oidc_auth() {
|
||||
#[test]
|
||||
fn session_token_auth_gate_truth_table() {
|
||||
use crate::agent::auth_method::{ModelByok, session_token_auth_gate as gate};
|
||||
use crate::auth::credential_authority::CredentialClass;
|
||||
// Non-session methods never refresh, regardless of BYOK status or endpoint.
|
||||
// `Pooled` (an OAuth platform's own pool) and `Primary` (kimi-code, or a
|
||||
// bare / [model.*] model on the session's own endpoint) behave identically
|
||||
// here: each names a credential that IS refreshable on that host.
|
||||
for fp in [false, true] {
|
||||
assert!(!gate(false, ModelByok::NotByok, fp, true));
|
||||
assert!(!gate(false, ModelByok::Byok, fp, true));
|
||||
assert!(!gate(false, ModelByok::Unknown, fp, true));
|
||||
// Session method on an endpoint that DOES take the session credential
|
||||
// (kimi-code, an OAuth platform's own pool, or a bare / [model.*]
|
||||
// model): a definite classification ignores the endpoint — NotByok
|
||||
// refreshes, a genuine per-model Byok never does.
|
||||
assert!(gate(true, ModelByok::NotByok, fp, true));
|
||||
assert!(!gate(true, ModelByok::Byok, fp, true));
|
||||
// …and an API-key registry platform endpoint is refused on every arm,
|
||||
// first-party flag included: that is the leak guard.
|
||||
assert!(!gate(true, ModelByok::NotByok, fp, false));
|
||||
assert!(!gate(true, ModelByok::Byok, fp, false));
|
||||
assert!(!gate(true, ModelByok::Unknown, fp, false));
|
||||
for class in [CredentialClass::Pooled, CredentialClass::Primary] {
|
||||
assert!(!gate(false, ModelByok::NotByok, fp, class));
|
||||
assert!(!gate(false, ModelByok::Byok, fp, class));
|
||||
assert!(!gate(false, ModelByok::Unknown, fp, class));
|
||||
// Session method on an endpoint that DOES take a session
|
||||
// credential: a definite classification ignores the endpoint —
|
||||
// NotByok refreshes, a genuine per-model Byok never does.
|
||||
assert!(gate(true, ModelByok::NotByok, fp, class));
|
||||
assert!(!gate(true, ModelByok::Byok, fp, class));
|
||||
}
|
||||
// …and an API-key registry platform endpoint (`CredentialClass::None`)
|
||||
// is refused on every arm, first-party flag included: the leak guard.
|
||||
assert!(!gate(true, ModelByok::NotByok, fp, CredentialClass::None));
|
||||
assert!(!gate(true, ModelByok::Byok, fp, CredentialClass::None));
|
||||
assert!(!gate(true, ModelByok::Unknown, fp, CredentialClass::None));
|
||||
}
|
||||
// Session method + Unknown BYOK: refresh only against a first-party xAI
|
||||
// host, so a transiently-unclassifiable config can't demote a live session
|
||||
// (the stale-token 401 regression) yet the session token never leaks to a
|
||||
// third-party BYOK endpoint. This arm was unconditionally `false` pre-fix.
|
||||
assert!(gate(true, ModelByok::Unknown, true, true));
|
||||
assert!(!gate(true, ModelByok::Unknown, false, true));
|
||||
assert!(gate(
|
||||
true,
|
||||
ModelByok::Unknown,
|
||||
true,
|
||||
CredentialClass::Primary
|
||||
));
|
||||
assert!(!gate(
|
||||
true,
|
||||
ModelByok::Unknown,
|
||||
false,
|
||||
CredentialClass::Primary
|
||||
));
|
||||
}
|
||||
|
||||
/// Pre-fix, the gate read `auth_type` and skipped recovery here, 401'ing every
|
||||
@@ -875,7 +890,7 @@ async fn set_session_model_invalidates_byok_memo_for_same_model_id() {
|
||||
header_injector: None,
|
||||
};
|
||||
let _ = actor
|
||||
.handle_set_session_model(cfg, false, false, true, 85)
|
||||
.handle_set_session_model(cfg, None, false, false, true, 85)
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
|
||||
@@ -109,6 +109,7 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
|
||||
session_info,
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state: TokioMutex::new(State {
|
||||
@@ -562,6 +563,7 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
|
||||
session_info: session_info.clone(),
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state: TokioMutex::new(State {
|
||||
@@ -825,6 +827,7 @@ async fn cancel_running_task_teardown_clears_running_and_pending_work() {
|
||||
},
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state,
|
||||
@@ -1817,6 +1820,7 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() {
|
||||
},
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state,
|
||||
|
||||
@@ -128,6 +128,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
|
||||
attribution_callback: None,
|
||||
auth_method_id: test_auth_method_id("cached_token"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
auth_manager: {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mgr = std::sync::Arc::new(crate::auth::AuthManager::new(
|
||||
|
||||
+3
@@ -72,6 +72,7 @@ async fn create_test_actor(
|
||||
rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(),
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state,
|
||||
@@ -511,6 +512,7 @@ async fn create_test_actor_with_memory(
|
||||
rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(),
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state,
|
||||
@@ -1266,6 +1268,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
|
||||
rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(),
|
||||
auth_method_id: test_auth_method_id("cached_token"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
auth_manager: {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mgr = std::sync::Arc::new(crate::auth::AuthManager::new(
|
||||
|
||||
@@ -127,6 +127,7 @@ async fn create_test_actor_with_memory(
|
||||
},
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state,
|
||||
|
||||
+1
@@ -78,6 +78,7 @@ pub(super) async fn make_replay_send_update_fixture() -> ReplaySendUpdateFixture
|
||||
},
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state,
|
||||
|
||||
+770
-69
@@ -8,8 +8,38 @@ use super::session_bearer_leak_tests::{
|
||||
};
|
||||
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();
|
||||
@@ -127,28 +157,21 @@ async fn oauth_platform_models_keep_a_live_resolver_from_their_own_pool() {
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
for (catalog_key, slug, base_url) in [
|
||||
(
|
||||
"claude-pro-max/claude-opus-4-8",
|
||||
"claude-opus-4-8",
|
||||
"https://api.anthropic.com/v1",
|
||||
),
|
||||
(
|
||||
"github-copilot/gpt-4.1",
|
||||
"gpt-4.1",
|
||||
"https://api.githubcopilot.com",
|
||||
),
|
||||
(
|
||||
"xai-grok/grok-4-latest",
|
||||
"grok-4-latest",
|
||||
"https://api.x.ai/v1",
|
||||
),
|
||||
(
|
||||
"openai-codex/gpt-5.5",
|
||||
"gpt-5.5",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
),
|
||||
// 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;
|
||||
@@ -165,12 +188,13 @@ async fn oauth_platform_models_keep_a_live_resolver_from_their_own_pool() {
|
||||
// token rotated inside the pool is observed by the
|
||||
// already-built resolver (this is what mid-session refresh
|
||||
// does). The pool is read here, never mutated.
|
||||
let pooled = crate::auth::oauth_registry::manager_for_model(
|
||||
&crate::auth::oauth_registry::pool_home(),
|
||||
catalog_key,
|
||||
actor.auth_manager.as_ref(),
|
||||
)
|
||||
.expect("an OAuth platform always resolves a manager");
|
||||
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,
|
||||
@@ -189,25 +213,30 @@ async fn oauth_platform_models_keep_a_live_resolver_from_their_own_pool() {
|
||||
}
|
||||
|
||||
/// H3/H4 — the stamped AUX paths (image-describe, the auto-mode classifier and
|
||||
/// the session-summary client all funnel through `aux_bearer_resolver`).
|
||||
/// `stamp_session_local_sampler_fields` copies the SESSION model's resolver onto
|
||||
/// every aux config and `SamplingClient::post` REPLACES the request's auth
|
||||
/// header from it, so an API-key-platform aux model would have its own key
|
||||
/// overwritten by the Kimi bearer ON THE AUX HOST.
|
||||
/// 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: returning `stamped` unconditionally (the pre-fix
|
||||
/// "re-point only when the aux model is OAuth" shape) makes the deepseek /
|
||||
/// openai / `[model.*]`-on-openai.com rows resolve `KIMI_TOKEN`.
|
||||
/// 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() {
|
||||
#[derive(Debug)]
|
||||
struct Fixed(&'static str);
|
||||
impl BearerResolver for Fixed {
|
||||
fn current_bearer(&self) -> Option<String> {
|
||||
Some(self.0.to_string())
|
||||
}
|
||||
}
|
||||
let stamped: kigi_sampler::SharedBearerResolver = Arc::new(Fixed(KIMI_TOKEN));
|
||||
let 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);
|
||||
|
||||
@@ -222,22 +251,16 @@ fn aux_bearer_resolver_clears_the_session_resolver_off_a_third_party_aux_host()
|
||||
),
|
||||
] {
|
||||
assert!(
|
||||
crate::session::acp_session::sampler_turn::aux_bearer_resolver(
|
||||
Some(stamped.clone()),
|
||||
platform(key),
|
||||
base_url,
|
||||
)
|
||||
.is_none(),
|
||||
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!(
|
||||
crate::session::acp_session::sampler_turn::aux_bearer_resolver(
|
||||
Some(stamped.clone()),
|
||||
None,
|
||||
"https://api.openai.com/v1",
|
||||
)
|
||||
.is_none(),
|
||||
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"
|
||||
);
|
||||
|
||||
@@ -251,17 +274,15 @@ fn aux_bearer_resolver_clears_the_session_resolver_off_a_third_party_aux_host()
|
||||
(None, kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url),
|
||||
(None, "http://127.0.0.1:4141/v1"),
|
||||
] {
|
||||
let resolved = crate::session::acp_session::sampler_turn::aux_bearer_resolver(
|
||||
Some(stamped.clone()),
|
||||
key.and_then(platform),
|
||||
base_url,
|
||||
)
|
||||
.unwrap_or_else(|| panic!("{key:?} @ {base_url} must keep the session resolver"));
|
||||
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 gets a LIVE resolver over its OWN pool
|
||||
// (empty here), never the stamped Kimi one.
|
||||
// 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()
|
||||
@@ -273,17 +294,697 @@ fn aux_bearer_resolver_clears_the_session_resolver_off_a_third_party_aux_host()
|
||||
"github-copilot/gpt-4.1",
|
||||
"openai-codex/gpt-5.5",
|
||||
] {
|
||||
let resolved = crate::session::acp_session::sampler_turn::aux_bearer_resolver(
|
||||
Some(stamped.clone()),
|
||||
platform(key),
|
||||
"https://example.invalid/v1",
|
||||
)
|
||||
.unwrap_or_else(|| panic!("{key} must keep a live resolver from its own pool"));
|
||||
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;
|
||||
}
|
||||
|
||||
+29
-28
@@ -5,7 +5,7 @@
|
||||
//! `kimi-code` / any OAuth platform) + a selected API-key-platform model
|
||||
//! classifies `ModelByok::NotByok` (the model carries no `[model.*]` key), the
|
||||
//! pre-fix `session_token_auth_gate` returned `true` unconditionally on that
|
||||
//! arm, `auth_manager_for_model` fell through to the primary Kimi manager for a
|
||||
//! 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.
|
||||
//!
|
||||
@@ -21,11 +21,12 @@
|
||||
//! `bearer_resolver` drawn from their OWN pooled `AuthManager`, or they lose
|
||||
//! mid-session token refresh.
|
||||
//!
|
||||
//! STORAGE DISCIPLINE (H6): nothing here touches the developer's real `~/.kigi`
|
||||
//! and nothing hot-swaps the process-global OAuth pool. Under `cfg(test)`
|
||||
//! `oauth_registry::pool_home()` is a process-lifetime `TempDir`, so every
|
||||
//! pooled manager is empty — which is exactly what the assertions need (a live
|
||||
//! resolver that is provably NOT the Kimi one).
|
||||
//! 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::*;
|
||||
@@ -106,9 +107,14 @@ pub(super) async fn actor_with_catalog(
|
||||
actor.models_manager.insert_test_entry(key, entry);
|
||||
}
|
||||
let selected_entry = selected_entry.expect("the selected key must be in the catalog");
|
||||
actor
|
||||
.models_manager
|
||||
.set_current_model_id(acp::ModelId::new(selected.to_string()));
|
||||
// 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
|
||||
@@ -166,7 +172,7 @@ pub(super) async fn actor_on_managed_model(
|
||||
/// (`cached_token`) method with a live Kimi primary must send DeepSeek's own key
|
||||
/// — the Kimi subscription bearer must not appear anywhere in the request.
|
||||
///
|
||||
/// Revert-to-red: dropping `endpoint_takes_session_credential` from
|
||||
/// 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() {
|
||||
@@ -282,11 +288,12 @@ async fn api_key_platform_models_get_no_session_bearer_resolver() {
|
||||
/// C2 at the resolver channel: a `[model.*]` entry has NO platform
|
||||
/// (`info.id == None`), which used to be a blanket allow. Pointed at a
|
||||
/// third-party host it must get no session resolver; pointed at the session's
|
||||
/// own coding endpoint (a `KIGI_CODE_BASE_URL` deployment or a local dev proxy)
|
||||
/// it must keep one — that is why the predicate is not `is_first_party_url`.
|
||||
/// 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: making the `None` arm of `platform_takes_session_credential`
|
||||
/// return `true` again puts a Kimi resolver on the openai.com config.
|
||||
/// 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();
|
||||
@@ -343,21 +350,15 @@ async fn config_model_entry_takes_a_session_resolver_only_on_its_own_endpoint()
|
||||
.await;
|
||||
}
|
||||
|
||||
/// H5 — the slug collision. A user holding BOTH an xAI API key and a Grok
|
||||
/// subscription has `xai/grok-4.5` AND `xai-grok/grok-4.5` in one catalog, in
|
||||
/// `PlatformId::ALL` order (`Xai`(15) before `XaiGrok`(25)) and with the SAME
|
||||
/// routing slug. `cfg.model` is that bare slug, so the auth layer used to
|
||||
/// first-match the API-key entry: no bearer_resolver, no live refresh (the
|
||||
/// session dies ~1h in with an unrecoverable 401), and — for the Anthropic and
|
||||
/// Codex twins — the OAuth Messages adaptation and the Codex identity headers
|
||||
/// silently dropped.
|
||||
/// The 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.
|
||||
///
|
||||
/// The catalog KEY the picker selected is now authoritative.
|
||||
///
|
||||
/// Revert-to-red: resolving the platform from `find_model_by_id(models, slug)`
|
||||
/// instead of `current_model_id` resolves `xai/grok-4.5` /
|
||||
/// `anthropic/claude-opus-4-8` / `openai/gpt-5.5-codex` and every assertion
|
||||
/// below fails.
|
||||
/// (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();
|
||||
|
||||
@@ -188,6 +188,7 @@ pub(crate) async fn create_test_actor_ex(
|
||||
},
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state,
|
||||
|
||||
@@ -145,6 +145,13 @@ pub enum SessionCommand {
|
||||
},
|
||||
SetSessionModel {
|
||||
sampling_config: kigi_sampler::SamplerConfig,
|
||||
/// The catalog KEY the picker resolved (`{platform}/{model}` for a
|
||||
/// registry model), which `sampling_config.model` — the bare routing
|
||||
/// slug — cannot express when an API-key platform and its
|
||||
/// subscription-OAuth twin list the same id. The session stores it as
|
||||
/// its OWN selection instead of reading the process-global
|
||||
/// `ModelsManager::current_model_id()` (H4).
|
||||
catalog_key: Option<String>,
|
||||
use_concise: bool,
|
||||
/// When `false`, skip the system prompt rewrite (concise/default swap).
|
||||
/// Set to `false` for forked sessions so mid-session model switches
|
||||
|
||||
@@ -2168,6 +2168,7 @@ mod inline_auto_compact_flow_tests {
|
||||
},
|
||||
auth_method_id: test_auth_method_id("test-auth"),
|
||||
model_auth_facts: std::cell::RefCell::new(None),
|
||||
selected_catalog_key: std::cell::RefCell::new(None),
|
||||
attribution_callback: None,
|
||||
auth_manager: None,
|
||||
state,
|
||||
|
||||
Reference in New Issue
Block a user