From 1bb10ef7f106a28e4181b56f5cb0549599e03d7d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Wed, 22 Jul 2026 17:16:38 -0400 Subject: [PATCH] fix(models): subscription-OAuth sessions are a fetch source for the catalog MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause of 'connected Claude, still shows Kimi / unknown model': every fetch-plan decision consulted only the primary (Kimi) session and API keys — never the stored subscription-OAuth sessions: - on_auth_changed's wipe guard: a claude-pro-max-only login satisfied 'no session, no keys' → catalog wiped, early return BEFORE the fetch and BEFORE notify_models_updated. The TUI kept an empty picker and the prompt bar rendered 'unknown'. Guard decision extracted into the pure should_wipe_catalog_on_auth_change (matrix-tested); a stored OAuth session now vetoes the wipe, so the fetch runs, the first real catalog reselects the default model (first entry = the connected provider's), and kigi/models/update reaches the client. - Startup prefetch: the arming gate ignored stored OAuth sessions and the prefetch thread passed an empty token map — a claude-only user booted onto the bundled Kimi table until a later refresh. The gate now takes has_stored_oauth and the thread resolves each stored session's bearer (refresh-on-expiry) via a current-thread runtime. - Cache origins: from_config's startup cache load and cache_origin() computed the fetch-plan origin with an empty token map, so a claude-inclusive cached catalog never matched at startup. Both now use presence-only stubs (stored_oauth_token_stubs — names only, no bearers) proven equal to the real-token origin by test. New probes in models_fetch: stored_oauth_platforms / stored_oauth_token_ stubs (sync auth.json scope scan; no AuthManager, no secrets). Verified: kigi-shell 5260 tests green, clippy clean. --- crates/codegen/kigi-shell/src/agent/models.rs | 154 +++++++++++++++--- .../kigi-shell/src/agent/models_fetch.rs | 107 ++++++++++++ 2 files changed, 242 insertions(+), 19 deletions(-) diff --git a/crates/codegen/kigi-shell/src/agent/models.rs b/crates/codegen/kigi-shell/src/agent/models.rs index 3e81994..e4c0810 100644 --- a/crates/codegen/kigi-shell/src/agent/models.rs +++ b/crates/codegen/kigi-shell/src/agent/models.rs @@ -334,7 +334,13 @@ impl ModelsManager { &cfg.endpoints, fetch_auth, has_session, - &Default::default(), + // Presence-only stubs: the origin encodes enabled + // platform NAMES, and it must match the fetch path's + // (which enables stored subscription-OAuth platforms) + // or their cached catalog never loads at startup. + &crate::agent::models_fetch::stored_oauth_token_stubs( + &crate::auth::oauth_registry::pool_home(), + ), &platform_keys, ), ) @@ -740,12 +746,15 @@ impl ModelsManager { self.inner.cache.invalidate(); let fetch_auth = ModelFetchAuth::resolve(&config.endpoints); *self.inner.fetch_auth.write() = fetch_auth; - // With no session, no open-platform key, and no custom endpoint there - // is nothing to fetch from: wipe the previous identity's catalog. - if self.inner.auth_manager.current_or_expired().is_none() - && fetch_auth == ModelFetchAuth::Platforms - && !PlatformApiKeys::resolve(&config.platforms).any() - { + if should_wipe_catalog_on_auth_change( + self.inner.auth_manager.current_or_expired().is_some(), + fetch_auth, + PlatformApiKeys::resolve(&config.platforms).any(), + !crate::agent::models_fetch::stored_oauth_platforms( + &crate::auth::oauth_registry::pool_home(), + ) + .is_empty(), + ) { self.clear(); return; } @@ -1153,14 +1162,17 @@ impl ModelsManager { let has_oauth = self.inner.auth_manager.current_or_expired().is_some(); let platform_keys = PlatformApiKeys::resolve(&platforms); // The origin key encodes only enabled-platform NAMES + URLs (never - // tokens). Generic-oauth presence is reflected by the post-login - // `on_auth_changed` re-fetch; an empty map here keeps this sync path - // cheap (no per-provider AuthManager construction on the hot path). + // tokens). Stored subscription-OAuth platforms join via presence-only + // stubs (cheap auth.json scan — no per-provider AuthManager on this + // sync path) so this origin matches the fetch path's, which enables + // those platforms with real bearers. crate::agent::models_fetch::models_fetch_origin( &endpoints, fetch_auth, has_oauth, - &Default::default(), + &crate::agent::models_fetch::stored_oauth_token_stubs( + &crate::auth::oauth_registry::pool_home(), + ), &platform_keys, ) } @@ -1837,6 +1849,10 @@ fn resolve_prefetch_env_with_auth(auth: Option) -> Option auth, endpoints, PlatformApiKeys::resolve_from_effective_config(), + !crate::agent::models_fetch::stored_oauth_platforms( + &crate::auth::oauth_registry::pool_home(), + ) + .is_empty(), crate::util::config::resolve_remote_fetch_enabled(), ) } @@ -1850,12 +1866,38 @@ fn resolve_prefetch_env_with_auth(auth: Option) -> Option /// or a `deployment_key` would re-arm the prefetch — and with it the /// deployment-config sync on the prefetch thread. /// +/// Decision core of [`ModelsManager::on_auth_changed`]'s wipe guard, split +/// from the disk probes so it is unit-testable: `true` = no fetch source +/// exists, wipe the previous identity's catalog. +/// +/// A stored subscription-OAuth session IS a fetch source (its platform +/// passes `enabled_platforms`), so it vetoes the wipe even with no primary +/// (Kimi) session and no API key. Regression: ignoring it meant a +/// claude-pro-max-only login wiped the catalog and returned before the +/// fetch — the session stayed on the bundled Kimi table with an empty +/// picker ("unknown" model). +fn should_wipe_catalog_on_auth_change( + has_primary_session: bool, + fetch_auth: ModelFetchAuth, + has_platform_keys: bool, + has_stored_oauth: bool, +) -> bool { + !has_primary_session + && fetch_auth == ModelFetchAuth::Platforms + && !has_platform_keys + && !has_stored_oauth +} + /// PRD F2 acceptance: a moonshot API key alone (no subscription login) must -/// arm the prefetch so the catalog syncs on startup. +/// arm the prefetch so the catalog syncs on startup. Likewise a stored +/// subscription-OAuth session alone (`has_stored_oauth`, e.g. a +/// claude-pro-max login with no Kimi session): its models are the user's +/// ONLY models, so the prefetch must run for them. fn resolve_prefetch_env_from_parts( auth: Option, endpoints: config::EndpointsConfig, platform_keys: PlatformApiKeys, + has_stored_oauth: bool, remote_fetch_enabled: bool, ) -> Option { if !remote_fetch_enabled { @@ -1865,7 +1907,11 @@ fn resolve_prefetch_env_from_parts( let model_fetch_auth = ModelFetchAuth::resolve(&endpoints); - if auth.is_none() && !endpoints.has_custom_endpoint() && !platform_keys.any() { + if auth.is_none() + && !endpoints.has_custom_endpoint() + && !platform_keys.any() + && !has_stored_oauth + { return None; } @@ -1910,13 +1956,25 @@ fn spawn_prefetch_thread(env: PrefetchEnv) -> EarlyPrefetchHandle { let mut timer = crate::instrumentation_timer!("startup.early_prefetch"); let proxy_endpoint = env.endpoints.proxy_url(); timer.with_field("endpoint", proxy_endpoint.as_str()); + // Resolve each stored subscription-OAuth session's bearer (refreshed + // on expiry) so those platforms are part of the STARTUP fetch plan — + // otherwise a claude-pro-max-only user boots onto the bundled Kimi + // table until some later async refresh happens to run, and the cache + // origin (which encodes enabled platforms) never matches the + // async-path's claude-inclusive origin. + let oauth_tokens = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .map(|rt| { + rt.block_on(crate::agent::models_fetch::resolve_generic_oauth_tokens( + &crate::auth::oauth_registry::pool_home(), + )) + }) + .unwrap_or_default(); let models = prefetch_models_blocking( &env.endpoints, env.auth.as_ref(), - // Startup prefetch does not resolve generic-oauth (xai-grok) - // tokens; those platforms join on the first async catalog refresh - // (post-login `on_auth_changed` / periodic `spawn_fetch`). - &Default::default(), + &oauth_tokens, env.model_fetch_auth, &env.platform_keys, ); @@ -3730,14 +3788,21 @@ mod tests { Some(KimiAuth::test_default()), endpoints.clone(), keys(Some("sk-cn"), None), + true, false, ) .is_none(), "session auth must not re-arm the prefetch when remote_fetch is off", ); assert!( - resolve_prefetch_env_from_parts(None, endpoints, keys(Some("sk-cn"), None), false) - .is_none(), + resolve_prefetch_env_from_parts( + None, + endpoints, + keys(Some("sk-cn"), None), + false, + false + ) + .is_none(), "platform key / custom endpoint must not re-arm it either", ); } @@ -3751,6 +3816,7 @@ mod tests { None, config::EndpointsConfig::default(), keys(None, Some("sk-ai")), + false, true, ); assert!( @@ -3762,6 +3828,7 @@ mod tests { None, config::EndpointsConfig::default(), PlatformApiKeys::default(), + false, true, ) .is_none(), @@ -3769,6 +3836,55 @@ mod tests { ); } + /// A stored subscription-OAuth session ALONE (claude-pro-max login, no + /// Kimi session, no API key) must arm the startup prefetch — its models + /// are the user's only models. Regression: the gate ignored stored OAuth + /// sessions, so such a user booted onto the bundled Kimi table. + #[test] + fn prefetch_env_resolves_with_stored_oauth_session_alone() { + assert!( + resolve_prefetch_env_from_parts( + None, + config::EndpointsConfig::default(), + PlatformApiKeys::default(), + true, + true, + ) + .is_some(), + "a stored subscription-OAuth session alone must arm the model sync", + ); + } + + /// The `on_auth_changed` wipe guard: only the genuinely credential-less + /// shape wipes. A stored subscription-OAuth session vetoes the wipe — + /// the regression that left a claude-pro-max-only login with an empty + /// catalog and an "unknown" model. + #[test] + fn wipe_guard_spares_stored_oauth_sessions() { + use ModelFetchAuth::{CustomEndpoint, Platforms}; + // No credential of any kind → wipe. + assert!(should_wipe_catalog_on_auth_change( + false, Platforms, false, false + )); + // A stored subscription-OAuth session alone → NO wipe (fetch runs). + assert!(!should_wipe_catalog_on_auth_change( + false, Platforms, false, true + )); + // Primary session / platform key / custom endpoint each veto too. + assert!(!should_wipe_catalog_on_auth_change( + true, Platforms, false, false + )); + assert!(!should_wipe_catalog_on_auth_change( + false, Platforms, true, false + )); + assert!(!should_wipe_catalog_on_auth_change( + false, + CustomEndpoint, + false, + false + )); + } + /// remote_fetch=false: an online catalog refresh is a no-op — nothing is /// fetched, no real-catalog flag is set, and the static catalog keeps /// resolving. Covers `list_models`/`do_refresh` online strategies too, diff --git a/crates/codegen/kigi-shell/src/agent/models_fetch.rs b/crates/codegen/kigi-shell/src/agent/models_fetch.rs index 8b53585..7926112 100644 --- a/crates/codegen/kigi-shell/src/agent/models_fetch.rs +++ b/crates/codegen/kigi-shell/src/agent/models_fetch.rs @@ -66,6 +66,38 @@ pub(crate) async fn resolve_generic_oauth_tokens( } tokens } + +/// The generic device-code OAuth platforms with a STORED session scope in +/// auth.json — the cheap sync companion to [`resolve_generic_oauth_tokens`] +/// (no `AuthManager`, no refresh, no secrets read beyond scope presence). +/// +/// Every place that reasons about the FETCH PLAN without resolving bearers +/// (the prefetch arming gate, `on_auth_changed`'s wipe guard, cache-origin +/// computation) must consult this, or a subscription-OAuth-only user (e.g. +/// claude-pro-max with no Kimi session and no API key) is treated as +/// credential-less: catalog wiped/never fetched, session stuck on the +/// bundled Kimi table with an empty picker. +pub(crate) fn stored_oauth_platforms(kigi_home: &std::path::Path) -> Vec { + let Ok(store) = crate::auth::read_auth_json(&kigi_home.join("auth.json")) else { + return Vec::new(); + }; + kigi_models::PlatformId::ALL + .into_iter() + .filter(|p| p.oauth().is_some_and(|o| store.contains_key(o.scope_key))) + .collect() +} + +/// Presence-only stand-in for [`OAuthSessionTokens`] in fetch-plan/origin +/// computations. Values are EMPTY strings: cache origins encode enabled +/// platform NAMES and URLs only ([`models_fetch_origin`]), and this map must +/// never reach a request builder — real bearers come from +/// [`resolve_generic_oauth_tokens`] on the fetch path itself. +pub(crate) fn stored_oauth_token_stubs(kigi_home: &std::path::Path) -> OAuthSessionTokens { + stored_oauth_platforms(kigi_home) + .into_iter() + .map(|p| (p, String::new())) + .collect() +} /// The models-fetch origin key for this endpoints/auth shape. Used as the /// models disk-cache origin: cached entries embed absolute `base_url`s from /// the backend(s) that served them, so a catalog fetched against one fetch @@ -3726,4 +3758,79 @@ mod tests { "https://models.acme.com/v1/models" ); } + + /// `stored_oauth_platforms` maps auth.json `oauth/` scopes to + /// their platforms — presence only, no AuthManager, no refresh. Missing + /// or empty auth.json → empty (self-cleaning TempDir). + #[test] + fn stored_oauth_platforms_reads_scopes_from_auth_json() { + let home = tempfile::tempdir().expect("tempdir"); + assert!( + stored_oauth_platforms(home.path()).is_empty(), + "no auth.json → no stored oauth platforms" + ); + + let mut store = std::collections::BTreeMap::new(); + store.insert( + "oauth/claude-pro-max".to_string(), + crate::auth::KimiAuth::test_default(), + ); + // A platform API-key scope must NOT count as a stored OAuth session. + store.insert( + "deepseek".to_string(), + crate::auth::KimiAuth::test_default(), + ); + std::fs::write( + home.path().join("auth.json"), + serde_json::to_string(&store).expect("serialize store"), + ) + .expect("write auth.json"); + + assert_eq!( + stored_oauth_platforms(home.path()), + vec![kigi_models::PlatformId::ClaudeProMax], + ); + + // The stub map carries the same platform set with EMPTY values. + let stubs = stored_oauth_token_stubs(home.path()); + assert_eq!(stubs.len(), 1); + assert_eq!(stubs[&kigi_models::PlatformId::ClaudeProMax], ""); + } + + /// The origin computed from presence-only stubs equals the origin the + /// fetch path computes from REAL resolved tokens — the whole point of + /// the stubs (a claude-inclusive cached catalog must load at startup). + #[test] + fn stub_origin_matches_real_token_origin() { + use crate::agent::config::EndpointsConfig; + use crate::agent::models::{ModelFetchAuth, PlatformApiKeys}; + let cfg = EndpointsConfig::default(); + let mut real = OAuthSessionTokens::new(); + real.insert( + kigi_models::PlatformId::ClaudeProMax, + "live-bearer".to_string(), + ); + let mut stubs = OAuthSessionTokens::new(); + stubs.insert(kigi_models::PlatformId::ClaudeProMax, String::new()); + let with_real = models_fetch_origin( + &cfg, + ModelFetchAuth::Platforms, + false, + &real, + &PlatformApiKeys::default(), + ); + let with_stubs = models_fetch_origin( + &cfg, + ModelFetchAuth::Platforms, + false, + &stubs, + &PlatformApiKeys::default(), + ); + assert_eq!(with_real, with_stubs); + assert!(with_real.contains("claude-pro-max=")); + assert!( + !with_real.contains("live-bearer"), + "origin never embeds tokens" + ); + } }