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

fix(models): subscription-OAuth sessions are a fetch source for the catalog

Browse Source 
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.
This commit is contained in:
ZacharyZhang-NY
2026-07-22 17:16:38 -04:00
parent 0e352c9914
commit 1bb10ef7f1
2 changed files with 242 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/codegen/kigi-shell/src/agent/models.rs
+135 -19
View File
@@ -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<KimiAuth>) -> Option<PrefetchEnv>
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<KimiAuth>) -> Option<PrefetchEnv>
/// 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<KimiAuth>,
endpoints: config::EndpointsConfig,
platform_keys: PlatformApiKeys,
has_stored_oauth: bool,
remote_fetch_enabled: bool,
) -> Option<PrefetchEnv> {
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,
crates/codegen/kigi-shell/src/agent/models_fetch.rs
+107
View File
@@ -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<kigi_models::PlatformId> {
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/<provider>` 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"
);
}
}