From 021b82443d7a28c5c5462062e76cba09c44686e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Fri, 17 Jul 2026 07:37:29 -0400 Subject: [PATCH] M1/F1: Kimi Code OAuth device-code flow Replace the xAI OAuth stack with the Kimi device authorization grant: - kimi_oauth.rs wire layer (device_authorization + token poll + refresh against kigi_env::oauth_host(); client_id per PRD; retryable statuses 429/5xx with backoff; expired_token restarts authorization) - X-Msh-Device-{Name,Model,Id} headers; device_id minted uuid4-hex at ~/.kigi/device_id (0600) - Storage: system keyring service `kigi`, entry `oauth/kimi-code` (macOS/Windows native backends), atomic-file fallback under ~/.kigi; official client's keyring/~/.kimi never touched - Refresh manager: 60s tick, threshold max(300, expires_in*0.5), 401-tombstone keyed by rejected refresh token with 300s cooldown and rotation auto-clear, cross-process lock with sibling-adoption triple-check, sleep/wake forced refresh - Deleted xAI machinery: enterprise OIDC (PKCE/JWKS/teams), devbox login, external auth provider, JWT tier gating + subscription paywall stack, X-XAI-Token-Auth marker headers, ZDR gates, /user enrichment - kigi login / TUI /login both drive the device flow; login-host display now derives from kigi_env::oauth_host() - 264 auth unit/wiremock tests; live contract probe of auth.kimi.com/api/oauth/device_authorization matches the wire shapes Gates: check/clippy --all-targets clean, fmt, deny ok, kigi-shell lib 5131 tests green. --- Cargo.lock | 53 +- Cargo.toml | 1 + crates/codegen/kigi-auth/src/auth_provider.rs | 12 - crates/codegen/kigi-auth/src/visibility.rs | 2 +- crates/codegen/kigi-bin/src/main.rs | 11 +- crates/codegen/kigi-memory/src/embedding.rs | 1 - .../kigi-pager-pty-harness/src/flows.rs | 2 +- crates/codegen/kigi-sampler/src/client.rs | 6 +- crates/codegen/kigi-sampler/src/config.rs | 3 +- crates/codegen/kigi-shell/Cargo.toml | 7 + crates/codegen/kigi-shell/src/agent/app.rs | 83 +- .../kigi-shell/src/agent/auth_method.rs | 935 +--- crates/codegen/kigi-shell/src/agent/config.rs | 401 +- .../kigi-shell/src/agent/feedback_client.rs | 58 +- crates/codegen/kigi-shell/src/agent/init.rs | 6 +- crates/codegen/kigi-shell/src/agent/mod.rs | 1 - crates/codegen/kigi-shell/src/agent/models.rs | 38 +- .../src/agent/mvp_agent/acp_agent.rs | 287 +- .../src/agent/mvp_agent/agent_ops.rs | 159 +- .../kigi-shell/src/agent/mvp_agent/mod.rs | 489 +- .../kigi-shell/src/agent/mvp_agent/tests.rs | 411 +- .../src/agent/session_registry_client.rs | 12 +- .../kigi-shell/src/agent/subagent/mod.rs | 2 +- .../src/agent/subscription_check.rs | 191 - .../kigi-shell/src/auth/attribution.rs | 20 +- crates/codegen/kigi-shell/src/auth/config.rs | 440 +- .../src/auth/credential_provider.rs | 54 +- .../kigi-shell/src/auth/devbox_login_stub.rs | 37 - crates/codegen/kigi-shell/src/auth/device.rs | 297 ++ .../kigi-shell/src/auth/device_code.rs | 1094 ++--- crates/codegen/kigi-shell/src/auth/error.rs | 65 +- .../kigi-shell/src/auth/external_auth.rs | 283 -- crates/codegen/kigi-shell/src/auth/flow.rs | 1758 +------ crates/codegen/kigi-shell/src/auth/jwt.rs | 45 - .../codegen/kigi-shell/src/auth/kimi_oauth.rs | 626 +++ crates/codegen/kigi-shell/src/auth/manager.rs | 1093 ++--- .../kigi-shell/src/auth/manager/enrichment.rs | 256 - .../kigi-shell/src/auth/manager_tests.rs | 4265 ++--------------- crates/codegen/kigi-shell/src/auth/meta.rs | 22 +- crates/codegen/kigi-shell/src/auth/mod.rs | 30 +- crates/codegen/kigi-shell/src/auth/model.rs | 535 +-- .../codegen/kigi-shell/src/auth/oidc/login.rs | 701 --- .../codegen/kigi-shell/src/auth/oidc/mod.rs | 14 - .../kigi-shell/src/auth/oidc/protocol.rs | 1283 ----- .../kigi-shell/src/auth/oidc/refresh.rs | 249 - .../kigi-shell/src/auth/oidc/test_helpers.rs | 130 - .../codegen/kigi-shell/src/auth/recovery.rs | 434 +- .../refresh/auth_backend_contract_tests.rs | 351 -- .../src/auth/refresh/external_refresher.rs | 176 - .../src/auth/refresh/kimi_refresher.rs | 356 ++ .../kigi-shell/src/auth/refresh/mod.rs | 156 +- .../src/auth/refresh/oidc_refresher.rs | 260 - .../src/auth/refresh/oidc_refresher_tests.rs | 1311 ----- crates/codegen/kigi-shell/src/auth/storage.rs | 268 +- .../codegen/kigi-shell/src/auth/token_type.rs | 56 +- crates/codegen/kigi-shell/src/cli_models.rs | 75 +- .../codegen/kigi-shell/src/config/reloader.rs | 14 +- .../codegen/kigi-shell/src/extensions/auth.rs | 73 +- .../kigi-shell/src/extensions/auth_gate.rs | 8 +- .../kigi-shell/src/extensions/billing.rs | 8 - .../kigi-shell/src/extensions/bundle.rs | 28 +- .../codegen/kigi-shell/src/extensions/mod.rs | 1 - .../kigi-shell/src/extensions/privacy.rs | 91 - .../src/extensions/session_admin.rs | 18 +- .../kigi-shell/src/extensions/share.rs | 26 +- crates/codegen/kigi-shell/src/inspect/mod.rs | 62 - .../codegen/kigi-shell/src/leader/server.rs | 22 +- crates/codegen/kigi-shell/src/lib.rs | 1 - .../codegen/kigi-shell/src/managed_config.rs | 200 +- crates/codegen/kigi-shell/src/mcp_doctor.rs | 8 +- crates/codegen/kigi-shell/src/remote/agent.rs | 3 +- .../src/remote/chat_models_client.rs | 4 - .../codegen/kigi-shell/src/remote/client.rs | 53 +- .../src/remote/conversations_client.rs | 12 +- .../kigi-shell/src/remote/pull_smoke_test.rs | 10 +- .../src/remote/workspaces_client.rs | 6 +- .../session/acp_session_impl/sampler_turn.rs | 87 +- .../auth_error_no_retry_tests.rs | 158 +- .../goal/goal_classifier_e2e_tests.rs | 2 +- .../acp_session_tests/idle_resume_tests.rs | 8 +- .../inline_auto_compact_flow_tests.rs | 8 +- .../media_gen_auth_retry_tests.rs | 22 +- .../reactive_managed_reauth_e2e_tests.rs | 6 +- .../src/session/feedback_manager.rs | 20 +- .../kigi-shell/src/session/persistence.rs | 7 +- .../src/session/unified_list/mod.rs | 12 +- .../src/test_support/lsp_runtime.rs | 2 +- crates/codegen/kigi-shell/src/tier.rs | 58 - .../kigi-shell/src/trace_classifier/mod.rs | 24 +- ...redentials.rs => kigi_auth_credentials.rs} | 45 +- crates/codegen/kigi-shell/src/util/mod.rs | 2 +- .../kigi-shell/tests/signed_managed_config.rs | 145 +- .../tests/signed_managed_config/common.rs | 54 +- .../tests/signed_managed_config_extended.rs | 29 +- .../kigi-shell/tests/team_managed_config.rs | 1846 ------- .../kigi-shell/tests/test_settings_refresh.rs | 2 +- crates/codegen/kigi-tui/src/acp/mod.rs | 6 +- crates/codegen/kigi-tui/src/acp/spawn.rs | 4 +- .../src/app/acp_handler/tests/settings.rs | 3 +- crates/codegen/kigi-tui/src/app/app_view.rs | 184 +- crates/codegen/kigi-tui/src/app/cli.rs | 25 +- .../codegen/kigi-tui/src/app/dispatch/auth.rs | 12 +- .../kigi-tui/src/app/dispatch/billing.rs | 15 +- .../kigi-tui/src/app/dispatch/tests/auth.rs | 4 +- .../src/app/dispatch/tests/billing.rs | 44 - .../src/app/dispatch/tests/task_result.rs | 77 - .../codegen/kigi-tui/src/app/effects/mod.rs | 2 +- crates/codegen/kigi-tui/src/app/event_loop.rs | 5 +- crates/codegen/kigi-tui/src/app/mod.rs | 8 +- .../kigi-tui/src/app/session_startup.rs | 8 +- crates/codegen/kigi-tui/src/sessions_cmd.rs | 6 +- .../subscription_watch_and_gate_verify_pty.rs | 439 -- .../kigi-tui/tests/pty_e2e_config_ui.rs | 2 - crates/codegen/kigi-update/src/version.rs | 2 +- crates/codegen/kigi-workspace/src/hub_auth.rs | 2 +- .../kigi-workspace/src/session/tool_config.rs | 7 - deny.toml | 1 + 117 files changed, 4052 insertions(+), 19900 deletions(-) delete mode 100644 crates/codegen/kigi-shell/src/agent/subscription_check.rs delete mode 100644 crates/codegen/kigi-shell/src/auth/devbox_login_stub.rs create mode 100644 crates/codegen/kigi-shell/src/auth/device.rs delete mode 100644 crates/codegen/kigi-shell/src/auth/external_auth.rs delete mode 100644 crates/codegen/kigi-shell/src/auth/jwt.rs create mode 100644 crates/codegen/kigi-shell/src/auth/kimi_oauth.rs delete mode 100644 crates/codegen/kigi-shell/src/auth/manager/enrichment.rs delete mode 100644 crates/codegen/kigi-shell/src/auth/oidc/login.rs delete mode 100644 crates/codegen/kigi-shell/src/auth/oidc/mod.rs delete mode 100644 crates/codegen/kigi-shell/src/auth/oidc/protocol.rs delete mode 100644 crates/codegen/kigi-shell/src/auth/oidc/refresh.rs delete mode 100644 crates/codegen/kigi-shell/src/auth/oidc/test_helpers.rs delete mode 100644 crates/codegen/kigi-shell/src/auth/refresh/auth_backend_contract_tests.rs delete mode 100644 crates/codegen/kigi-shell/src/auth/refresh/external_refresher.rs create mode 100644 crates/codegen/kigi-shell/src/auth/refresh/kimi_refresher.rs delete mode 100644 crates/codegen/kigi-shell/src/auth/refresh/oidc_refresher.rs delete mode 100644 crates/codegen/kigi-shell/src/auth/refresh/oidc_refresher_tests.rs delete mode 100644 crates/codegen/kigi-shell/src/extensions/privacy.rs delete mode 100644 crates/codegen/kigi-shell/src/tier.rs rename crates/codegen/kigi-shell/src/util/{grok_auth_credentials.rs => kigi_auth_credentials.rs} (82%) delete mode 100644 crates/codegen/kigi-shell/tests/team_managed_config.rs delete mode 100644 crates/codegen/kigi-tui/tests/pty_e2e/subscription_watch_and_gate_verify_pty.rs diff --git a/Cargo.lock b/Cargo.lock index d0c702b..de0f7e2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1960,6 +1960,16 @@ dependencies = [ "unicode-segmentation", ] +[[package]] +name = "core-foundation" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "core-foundation" version = "0.10.1" @@ -5509,6 +5519,20 @@ dependencies = [ "thiserror 2.0.18", ] +[[package]] +name = "keyring" +version = "3.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eebcc3aff044e5944a8fbaf69eb277d11986064cba30c468730e8b9909fb551c" +dependencies = [ + "byteorder", + "log", + "security-framework 2.11.1", + "security-framework 3.7.0", + "windows-sys 0.60.2", + "zeroize", +] + [[package]] name = "kigi-acp-lib" version = "0.1.0" @@ -5757,7 +5781,7 @@ version = "0.1.0" dependencies = [ "base64", "blake3", - "core-foundation", + "core-foundation 0.10.1", "dunce", "kigi-tty-utils", "kigi-version", @@ -6409,6 +6433,7 @@ dependencies = [ "jsonschema", "jsonwebtoken", "kanal", + "keyring", "kigi-acp-lib", "kigi-agent", "kigi-agent-lifecycle", @@ -6508,6 +6533,7 @@ dependencies = [ "walkdir", "webbrowser", "windows 0.61.3", + "wiremock", "zstd", ] @@ -6814,7 +6840,7 @@ dependencies = [ "chrono", "clap", "clap_complete", - "core-foundation", + "core-foundation 0.10.1", "criterion", "crossterm", "derive_more 2.1.1", @@ -9977,7 +10003,7 @@ dependencies = [ "openssl-probe", "rustls-pki-types", "schannel", - "security-framework", + "security-framework 3.7.0", ] [[package]] @@ -9996,7 +10022,7 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "26d1e2536ce4f35f4846aa13bff16bd0ff40157cdb14cc056c7b14ba41233ba0" dependencies = [ - "core-foundation", + "core-foundation 0.10.1", "core-foundation-sys", "jni", "log", @@ -10005,7 +10031,7 @@ dependencies = [ "rustls-native-certs", "rustls-platform-verifier-android", "rustls-webpki", - "security-framework", + "security-framework 3.7.0", "security-framework-sys", "webpki-root-certs", "windows-sys 0.61.2", @@ -10178,6 +10204,19 @@ dependencies = [ "zeroize", ] +[[package]] +name = "security-framework" +version = "2.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "897b2245f0b511c87893af39b033e5ca9cce68824c4d7e7630b5a1d339658d02" +dependencies = [ + "bitflags 2.13.1", + "core-foundation 0.9.4", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + [[package]] name = "security-framework" version = "3.7.0" @@ -10185,7 +10224,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" dependencies = [ "bitflags 2.13.1", - "core-foundation", + "core-foundation 0.10.1", "core-foundation-sys", "libc", "security-framework-sys", @@ -12788,7 +12827,7 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fc95580916af1e68ff6a7be07446fc5db73ebf71cf092de939bbf5f7e189f72" dependencies = [ - "core-foundation", + "core-foundation 0.10.1", "jni", "log", "ndk-context", diff --git a/Cargo.toml b/Cargo.toml index 69be7ce..4d7c4dc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -264,6 +264,7 @@ wait-timeout = "0.2" walkdir = "2" webbrowser = { version = "1.0.4" } which = "8" +keyring = { version = "3", default-features = false } windows = { version = "0.61", features = ["Win32_Security", "Win32_Security_Authorization", "Win32_Foundation", "Win32_System_Threading", "Win32_System_JobObjects", "Win32_System_Console", "Win32_System_Pipes"] } wiremock = "0.6" wl-clipboard-rs = "0.9" diff --git a/crates/codegen/kigi-auth/src/auth_provider.rs b/crates/codegen/kigi-auth/src/auth_provider.rs index a051b8f..392d857 100644 --- a/crates/codegen/kigi-auth/src/auth_provider.rs +++ b/crates/codegen/kigi-auth/src/auth_provider.rs @@ -19,15 +19,10 @@ pub struct CredentialSnapshot { /// identity (`StaticAuthCredentialProvider`). Read by the OTel layer to /// populate the `user.id` resource attribute. pub user_id: Option, - /// Team identifier from OAuth. `None` for personal accounts or when - /// no auth is configured. - pub team_id: Option, /// `uuidv5(NAMESPACE_OID, deployment_key)`, set only for deployment-key auth. pub deployment_id: Option, /// `uuidv5(NAMESPACE_OID, api_key)`, set only for `AuthMode::ApiKey`. pub api_key_id: Option, - /// Org id from the OIDC `organizationId` claim; `None` for personal / deployment-key auth. - pub organization_id: Option, } /// Source of truth for outbound auth on data-collector requests. @@ -50,13 +45,6 @@ pub trait AuthCredentialProvider: HttpAuth + Send + Sync + 'static { /// Returns `false` if no refresher is configured or refresh failed. async fn refresh_after_unauthorized(&self) -> bool; - /// Whether `X-XAI-Token-Auth` should be sent with the bearer token. - /// `false` for deployment keys (bare Bearer), `true` for user/OAuth tokens. - /// See `GrokAuthCredentials::apply()` for the wire format contract. - fn needs_token_auth_header(&self) -> bool { - true - } - /// Whether the provider holds a credential worth a real outbound attempt — /// an unexpired token (in memory or on disk), or a static key. Default /// `true` always attempts. diff --git a/crates/codegen/kigi-auth/src/visibility.rs b/crates/codegen/kigi-auth/src/visibility.rs index 1dbea40..65c4e14 100644 --- a/crates/codegen/kigi-auth/src/visibility.rs +++ b/crates/codegen/kigi-auth/src/visibility.rs @@ -1,5 +1,5 @@ /// Apply auth headers to outbound visibility requests. -/// Implemented by `kigi-shell::util::grok_auth_credentials::GrokAuthCredentials` +/// Implemented by `kigi-shell::util::kigi_auth_credentials::KigiAuthCredentials` /// to keep credential construction owned by shell while letting data-collector /// build the request without reaching back into shell types. pub trait HttpAuth: Send + Sync { diff --git a/crates/codegen/kigi-bin/src/main.rs b/crates/codegen/kigi-bin/src/main.rs index 0d2cf69..7d97726 100644 --- a/crates/codegen/kigi-bin/src/main.rs +++ b/crates/codegen/kigi-bin/src/main.rs @@ -492,7 +492,7 @@ async fn workspace_start( ); } ensure_authenticated( - &agent_config.grok_com_config, + &agent_config.kimi_code_config, false, Some("No cached credentials found. Run `grok login` first."), ) @@ -1682,18 +1682,13 @@ async fn async_main() -> Result<()> { ) .await; } - Command::Login { - legacy: _, - oauth, - device_auth, - devbox, - } => { + Command::Login => { init_tracing_simple("cli"); let config = kigi_shell::config::load_effective_config_disk_only() .map_err(|e| anyhow::anyhow!("Failed to load config: {e}"))?; let config = AgentConfig::new_from_toml_cfg(&config) .map_err(|e| anyhow::anyhow!("Failed to create agent config: {e}"))?; - kigi_shell::auth::run_cli_login(&config, oauth, device_auth, devbox).await?; + kigi_shell::auth::run_cli_login(&config).await?; println!(); kigi_shell::instrumentation::finalize_and_exit(0); } diff --git a/crates/codegen/kigi-memory/src/embedding.rs b/crates/codegen/kigi-memory/src/embedding.rs index cc1e883..8066d47 100644 --- a/crates/codegen/kigi-memory/src/embedding.rs +++ b/crates/codegen/kigi-memory/src/embedding.rs @@ -139,7 +139,6 @@ impl EmbeddingProvider for ApiEmbeddingProvider { let request = kigi_http::shared_client() .post(format!("{}/embeddings", self.api_base)) .json(&body_json) - .header("X-XAI-Token-Auth", "xai-grok-cli") .header("x-grok-client-version", kigi_version::VERSION); let req = match request.build() { diff --git a/crates/codegen/kigi-pager-pty-harness/src/flows.rs b/crates/codegen/kigi-pager-pty-harness/src/flows.rs index 82de6e4..b5adafb 100644 --- a/crates/codegen/kigi-pager-pty-harness/src/flows.rs +++ b/crates/codegen/kigi-pager-pty-harness/src/flows.rs @@ -82,7 +82,7 @@ pub fn seed_fake_oauth(content: &ContentController, user: &str) { r#"{{ "https://auth.x.ai::b1a00492-073a-47ea-816f-4c329264a828": {{ "key": "pty-test-oauth-token", - "auth_mode": "oidc", + "auth_mode": "oauth", "create_time": "2026-01-01T00:00:00Z", "user_id": "{user}", "email": "{user}@test.invalid", diff --git a/crates/codegen/kigi-sampler/src/client.rs b/crates/codegen/kigi-sampler/src/client.rs index fd63f09..a2338af 100644 --- a/crates/codegen/kigi-sampler/src/client.rs +++ b/crates/codegen/kigi-sampler/src/client.rs @@ -2191,8 +2191,10 @@ mod tests { let mut cfg = minimal_config(); cfg.extra_headers .insert("x-test-header".to_string(), "test-value".to_string()); - cfg.extra_headers - .insert("x-XAI-token-auth".to_string(), "xai-grok-cli".to_string()); + cfg.extra_headers.insert( + "x-custom-auth-marker".to_string(), + "marker-value".to_string(), + ); let _client = SamplingClient::new(cfg).expect("client with extra headers should construct"); } diff --git a/crates/codegen/kigi-sampler/src/config.rs b/crates/codegen/kigi-sampler/src/config.rs index 4b6d6c8..29e0089 100644 --- a/crates/codegen/kigi-sampler/src/config.rs +++ b/crates/codegen/kigi-sampler/src/config.rs @@ -38,8 +38,7 @@ pub enum AuthScheme { /// composing chat-state's `kigi_sampling_types::SamplingConfig` /// with `Credentials` (api key, client version). /// -/// URL-derived request headers (e.g. `X-XAI-Token-Auth` for the -/// cli-chat-proxy) are +/// URL-derived request headers are /// folded into [`Self::extra_headers`] by /// `agent::config::inject_url_derived_headers` before the /// `SamplerConfig` is handed to the actor. Auth is selected separately diff --git a/crates/codegen/kigi-shell/Cargo.toml b/crates/codegen/kigi-shell/Cargo.toml index c461bad..9673fa4 100644 --- a/crates/codegen/kigi-shell/Cargo.toml +++ b/crates/codegen/kigi-shell/Cargo.toml @@ -172,12 +172,19 @@ kigi-tool-types = { workspace = true } libc = { workspace = true } nix = { workspace = true } +[target.'cfg(target_os = "macos")'.dependencies] +# System-keyring credential storage for the Kimi Code OAuth session (PRD F1). +keyring = { workspace = true, features = ["apple-native"] } + [target.'cfg(windows)'.dependencies] siphasher = { workspace = true } windows = { workspace = true } +# System-keyring credential storage for the Kimi Code OAuth session (PRD F1). +keyring = { workspace = true, features = ["windows-native"] } [dev-dependencies] criterion = { workspace = true } +wiremock = { workspace = true } filetime = { workspace = true } tempfile = { workspace = true } kigi-memory = { workspace = true, features = [] } diff --git a/crates/codegen/kigi-shell/src/agent/app.rs b/crates/codegen/kigi-shell/src/agent/app.rs index 591ebc8..9bd80b4 100644 --- a/crates/codegen/kigi-shell/src/agent/app.rs +++ b/crates/codegen/kigi-shell/src/agent/app.rs @@ -20,7 +20,7 @@ use crate::agent::config::{Config as AgentConfig, ModelEntry}; use crate::agent::init::{bootstrap, exit_on_config_error}; use crate::agent::models::{ModelFetchAuth, prefetch_models_blocking}; use crate::agent::mvp_agent::MvpAgent; -use crate::auth::{AuthManager, AuthMode, GrokAuth, run_auth_flow}; +use crate::auth::{AuthManager, AuthMode, KimiAuth, run_auth_flow}; use crate::util::kigi_home; use dirs; @@ -399,76 +399,6 @@ pub async fn run_stdio_agent( result } -async fn migrate_devbox_auth_if_legacy( - auth: Option, - agent_config: &AgentConfig, -) -> Option { - let auth = auth?; - if !crate::auth::devbox_login::is_devbox_environment() || auth.auth_mode != AuthMode::WebLogin { - return Some(auth); - } - - info!("Devbox legacy auth detected, attempting migration to OIDC"); - kigi_log::unified_log::info( - "devbox legacy auth migration: starting", - None, - Some(serde_json::json!({ - "user_id": auth.user_id, - "auth_mode": format!("{:?}", auth.auth_mode), - })), - ); - - // save + remove_scope are two non-atomic writes to auth.json (no lock). Safe - // at startup: no concurrent writer yet, and `lookup_auth` prefers the primary - // scope if a reader sees the intermediate state. - let migration_auth_manager = agent_config.create_auth_manager(); - - let new_auth = match crate::auth::devbox_login::mint_devbox_auth(&migration_auth_manager).await - { - Ok(new_auth) => new_auth, - Err(e) => { - tracing::warn!(error = ?e, "devbox legacy auth migration: devbox login helper call failed, continuing with legacy auth"); - kigi_log::unified_log::error( - "devbox legacy auth migration: mint failed", - None, - Some(serde_json::json!({ "error": e.to_string() })), - ); - return Some(auth); - } - }; - match migration_auth_manager - .save_without_enrichment(new_auth) - .await - { - Ok(saved_auth) => { - if let Err(e) = migration_auth_manager.remove_scope(crate::auth::LEGACY_AUTH_SCOPE) { - tracing::warn!(error = ?e, "Failed to remove legacy auth scope entry (non-fatal)"); - } - kigi_log::unified_log::info( - "devbox legacy auth migration: succeeded", - None, - Some(serde_json::json!({ - "user_id": saved_auth.user_id, - "has_refresh_token": saved_auth.refresh_token.is_some(), - "expires_at": saved_auth.expires_at.map(|e| e.to_rfc3339()), - "auth_mode": format!("{:?}", saved_auth.auth_mode), - })), - ); - info!(user_id = %saved_auth.user_id, "Devbox legacy auth migrated to OIDC successfully"); - Some(saved_auth) - } - Err(e) => { - tracing::warn!(error = ?e, "devbox legacy auth migration: failed to save new auth, continuing with legacy"); - kigi_log::unified_log::error( - "devbox legacy auth migration: save failed", - None, - Some(serde_json::json!({ "error": e.to_string() })), - ); - Some(auth) - } - } -} - /// Run the agent in leader mode, accepting IPC connections from multiple clients. /// /// Startup sequence: @@ -681,14 +611,11 @@ pub async fn run_leader( // The IPC server is already accepting connections. Clients that send ACP // messages during this window receive a `leader_starting` error and can retry. - let ctx = &agent_config.grok_com_config; + let ctx = &agent_config.kimi_code_config; // Never interactive: a detached leader has no TTY (forcing OAuth here hung BYOK). - let auth: Option = crate::auth::try_ensure_session_noninteractive(ctx).await; + let auth: Option = crate::auth::try_ensure_session_noninteractive(ctx).await; - // ── Phase 6b: Legacy devbox auth migration ───────────────────────────── - let auth: Option = migrate_devbox_auth_if_legacy(auth, &agent_config).await; - - let auth_for_prefetch: Option = auth.clone(); + let auth_for_prefetch: Option = auth.clone(); let endpoints_for_prefetch = agent_config.endpoints.clone(); let fetch_auth_for_prefetch = ModelFetchAuth::resolve(&endpoints_for_prefetch, auth.is_some()); // The shared pair helper owns the remote_fetch gate for both halves, so a @@ -894,7 +821,7 @@ pub async fn run_leader( if let Some(home) = dirs::home_dir() { watch_paths.push(home.join(".claude.json")); } - let auth_scope = agent_config.grok_com_config.auth_scope(); + let auth_scope = agent_config.kimi_code_config.auth_scope(); // Gated on user_kigi_home() so a cwd-relative .kigi/auth.json is never // read as the user auth store when no home resolves. let initial_auth_key_hash = kigi_config::user_kigi_home() diff --git a/crates/codegen/kigi-shell/src/agent/auth_method.rs b/crates/codegen/kigi-shell/src/agent/auth_method.rs index b9928d0..a09d4d3 100644 --- a/crates/codegen/kigi-shell/src/agent/auth_method.rs +++ b/crates/codegen/kigi-shell/src/agent/auth_method.rs @@ -1,7 +1,6 @@ use agent_client_protocol as acp; use crate::agent::config::ModelEntry; -use crate::auth::PreferredAuthMethod; /// Shared, live handle to the agent's current ACP auth method id. /// @@ -55,18 +54,10 @@ pub fn has_xai_api_key_env() -> bool { /// Probes `std::env` at call time and consults each `ModelEntry` for a /// resolvable api_key/env_key -- both inputs can change between calls, so the /// result is not cached. -/// -/// `disable_api_key_auth` (`[grok_com_config] disable_api_key_auth` / -/// `KIGI_DISABLE_API_KEY_AUTH`) is the admin kill switch: when true the -/// method is never advertised, regardless of available credentials, so -/// `XAI_API_KEY` can't bypass a deployment's forced IdP login. -pub fn should_advertise_xai_api_key<'a, I>(disable_api_key_auth: bool, models: I) -> bool +pub fn should_advertise_xai_api_key<'a, I>(models: I) -> bool where I: IntoIterator, { - if disable_api_key_auth { - return false; - } has_xai_api_key_env() || models.into_iter().any(ModelEntry::has_own_credentials) } @@ -78,25 +69,13 @@ where /// unit-tested without any of that machinery. pub struct AuthMethodsBuildInputs<'a> { /// True if `xai.api_key` should be advertised AT ALL. Caller computes via - /// [`should_advertise_xai_api_key`]. When `preferred_method` is `Oidc`, - /// this is ignored (API key is never advertised under that pin). + /// [`should_advertise_xai_api_key`]. pub has_external_api_key: bool, /// True if a cached session token is available (either present at startup /// or recovered via silent refresh). pub has_cached_token: bool, - /// True if enterprise OIDC is configured. Mutually exclusive with the - /// default `grok.com` method. - pub has_enterprise_oidc: bool, - /// Required when `has_enterprise_oidc` is true; ignored otherwise. - pub enterprise_oidc_issuer: Option<&'a str>, - /// Optional display label for the login method (`grok.com` or `oidc`). + /// Optional display label for the interactive login method. pub login_label: Option<&'a str>, - /// True if `grok_com_config.auth_provider_command` is configured (sets - /// `meta.external_provider = true` on the `grok.com` method). - pub has_auth_provider_command: bool, - /// Config pin (`[auth] preferred_method`). `None` keeps multi-method - /// fallthrough; `Some` is fail-closed (only that method family). - pub preferred_method: Option, } /// Output of [`build_auth_methods`]. @@ -105,123 +84,36 @@ pub struct BuiltAuthMethods { /// `startup_auth_metadata()` reads `methods.first()` to decide whether /// interactive login is needed. pub methods: Vec, - /// The default `auth_method_id` to install on the agent. When unpinned, - /// `cached_token` wins over `xai.api_key` when both are present. When - /// pinned, only the preferred method may appear; `None` means unavailable - /// (fail auth — no cross-method fallthrough). + /// The default `auth_method_id` to install on the agent. `cached_token` + /// wins over `xai.api_key` when both are present; `None` means an + /// interactive login is required. pub default_auth_method_id: Option, } /// Build the `auth_methods` list and default `auth_method_id` from /// pre-computed inputs. /// -/// REGRESSION GUARD: when unpinned and -/// `has_external_api_key` is true, the **first** entry MUST be `xai.api_key`. -/// A prior change deferred it to the END for per-model credentials, which made -/// the pager send per-model-key users to the login screen. Unit tests lock this. +/// REGRESSION GUARD: when `has_external_api_key` is true, the **first** entry +/// MUST be `xai.api_key`. A prior change deferred it to the END for per-model +/// credentials, which made the pager send per-model-key users to the login +/// screen. Unit tests lock this. /// -/// Unpinned ordering (when each method is enabled): +/// Ordering (when each method is enabled): /// 1. `xai.api_key` (if `has_external_api_key`) /// 2. `cached_token` (if `has_cached_token`) -/// 3. exactly one of: -/// - `oidc` (if `has_enterprise_oidc`) -/// - `grok.com` (otherwise) +/// 3. `grok.com` (the Kimi Code device login) /// -/// Unpinned `default_auth_method_id`: +/// `default_auth_method_id`: /// - `cached_token` if `has_cached_token` /// - `xai.api_key` else if `has_external_api_key` /// - `None` otherwise -/// -/// Pinned (`preferred_method`): -/// - `ApiKey`: only `xai.api_key` if available; else empty list + `None` (fail). -/// - `Oidc`: `cached_token` (if any) + interactive login; never `xai.api_key`. -/// Default is `cached_token` when present, else `None` (interactive). pub fn build_auth_methods(inputs: AuthMethodsBuildInputs<'_>) -> BuiltAuthMethods { let AuthMethodsBuildInputs { has_external_api_key, has_cached_token, - has_enterprise_oidc, - enterprise_oidc_issuer, login_label, - has_auth_provider_command, - preferred_method, } = inputs; - match preferred_method { - Some(PreferredAuthMethod::ApiKey) => build_pinned_api_key(has_external_api_key), - Some(PreferredAuthMethod::Oidc) => build_pinned_oidc( - has_cached_token, - has_enterprise_oidc, - enterprise_oidc_issuer, - login_label, - has_auth_provider_command, - ), - None => build_unpinned( - has_external_api_key, - has_cached_token, - has_enterprise_oidc, - enterprise_oidc_issuer, - login_label, - has_auth_provider_command, - ), - } -} - -fn build_pinned_api_key(has_external_api_key: bool) -> BuiltAuthMethods { - if !has_external_api_key { - kigi_log::unified_log::warn( - "auth: preferred_method=api_key but no API key credentials available", - None, - None, - ); - return BuiltAuthMethods { - methods: Vec::new(), - default_auth_method_id: None, - }; - } - BuiltAuthMethods { - methods: vec![xai_api_key_auth_method()], - default_auth_method_id: Some(acp::AuthMethodId::new(XAI_API_KEY_METHOD_ID)), - } -} - -fn build_pinned_oidc( - has_cached_token: bool, - has_enterprise_oidc: bool, - enterprise_oidc_issuer: Option<&str>, - login_label: Option<&str>, - has_auth_provider_command: bool, -) -> BuiltAuthMethods { - let mut methods: Vec = Vec::new(); - let mut default_auth_method_id: Option = None; - - if has_cached_token { - methods.push(cached_token_auth_method()); - default_auth_method_id = Some(acp::AuthMethodId::new(CACHED_TOKEN_AUTH_METHOD_ID)); - } - - push_interactive_login( - &mut methods, - has_enterprise_oidc, - enterprise_oidc_issuer, - login_label, - has_auth_provider_command, - ); - - BuiltAuthMethods { - methods, - default_auth_method_id, - } -} - -fn build_unpinned( - has_external_api_key: bool, - has_cached_token: bool, - has_enterprise_oidc: bool, - enterprise_oidc_issuer: Option<&str>, - login_label: Option<&str>, - has_auth_provider_command: bool, -) -> BuiltAuthMethods { let mut methods: Vec = Vec::new(); let mut default_auth_method_id: Option = None; @@ -233,7 +125,7 @@ fn build_unpinned( if has_cached_token { methods.push(cached_token_auth_method()); // cached_token wins over xai.api_key for default_auth_method_id so - // is_session_based_auth() returns true and OIDC refresh stays alive. + // is_session_based_auth() returns true and OAuth refresh stays alive. let overrode_api_key = default_auth_method_id.is_some(); default_auth_method_id = Some(acp::AuthMethodId::new(CACHED_TOKEN_AUTH_METHOD_ID)); if overrode_api_key { @@ -248,13 +140,7 @@ fn build_unpinned( } } - push_interactive_login( - &mut methods, - has_enterprise_oidc, - enterprise_oidc_issuer, - login_label, - has_auth_provider_command, - ); + methods.push(kimi_code_auth_method(login_label)); BuiltAuthMethods { methods, @@ -262,35 +148,12 @@ fn build_unpinned( } } -fn push_interactive_login( - methods: &mut Vec, - has_enterprise_oidc: bool, - enterprise_oidc_issuer: Option<&str>, - login_label: Option<&str>, - has_auth_provider_command: bool, -) { - if has_enterprise_oidc { - // Caller invariant: `enterprise_oidc_issuer` MUST be `Some(...)` when - // `has_enterprise_oidc` is true. Production callers derive both from - // the same `cfg.grok_com_config.oidc` Option, so the inconsistent - // `(true, None)` combination is a programmer error -- panic loudly - // (matches the original `cfg.grok_com_config.oidc.as_ref().unwrap()` - // call in `MvpAgent::initialize()` before this refactor). - let issuer = enterprise_oidc_issuer - .expect("enterprise_oidc_issuer is required when has_enterprise_oidc is true"); - methods.push(oidc_auth_method(issuer, login_label)); - } else { - methods.push(grok_com_auth_method(login_label, has_auth_provider_command)); - } -} - /// ACP session auth method. Use `is_session_based_method` for classification. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum AuthMethodKind { XaiApiKey, CachedToken, GrokCom, - Oidc, Unknown, } @@ -300,7 +163,6 @@ impl AuthMethodKind { XAI_API_KEY_METHOD_ID => Self::XaiApiKey, CACHED_TOKEN_AUTH_METHOD_ID => Self::CachedToken, KIGI_COM_METHOD_ID => Self::GrokCom, - OIDC_METHOD_ID => Self::Oidc, _ => Self::Unknown, } } @@ -310,14 +172,14 @@ impl AuthMethodKind { matches!(self, Self::XaiApiKey) } - /// `true` for session-based methods (cached_token, grok.com, oidc). + /// `true` for session-based methods (cached_token, interactive login). pub fn is_session_based(self) -> bool { - matches!(self, Self::CachedToken | Self::GrokCom | Self::Oidc) + matches!(self, Self::CachedToken | Self::GrokCom) } - /// Requires user interaction (browser, OIDC redirect, or external auth command). + /// Requires user interaction (device-code login in the browser). pub fn needs_interactive_login(self) -> bool { - matches!(self, Self::GrokCom | Self::Oidc) + matches!(self, Self::GrokCom) } pub fn auth_error_message(self) -> &'static str { @@ -329,7 +191,7 @@ impl AuthMethodKind { } } -/// `true` for session-based ACP methods (cached_token, grok.com, oidc). +/// `true` for session-based ACP methods (cached_token, interactive login). pub fn is_session_based_method(method_id: &acp::AuthMethodId) -> bool { AuthMethodKind::from_id(method_id).is_session_based() } @@ -360,20 +222,18 @@ impl ModelByok { /// /// Gates on stable inputs, not `Credentials.auth_type`: that field collapses /// to `ApiKey` when the session-token cache is momentarily empty and -/// `XAI_API_KEY` is set, which demoted live OIDC sessions to non-refreshable +/// `XAI_API_KEY` is set, which demoted live sessions to non-refreshable /// api-key mode and 401'd every prompt until restart. `model_byok` still /// excludes genuine per-model BYOK, whose keys are not refreshable. /// /// `Unknown` (BYOK status indeterminate — config currently unparseable, no /// sampling config yet, or the per-model memo was cleared) must **not** demote /// a live session to non-refreshable api-key mode: that re-sends the stale -/// buffered token on every turn and 401s with `bad-credentials` until restart -/// (the stale-token regression this gate addresses; fall back rather than -/// demote on `Unknown`). It refreshes when `endpoint_is_first_party` — the -/// request targets a first-party host (cli-chat-proxy / first-party API), -/// where sending the session token cannot leak to a third-party BYOK -/// endpoint. A definite `NotByok` always refreshes (it only ever routes to -/// the session endpoint); a definite `Byok` never does. +/// buffered token on every turn and 401s with `bad-credentials` until restart. +/// It refreshes when `endpoint_is_first_party` — the request targets the +/// first-party API, where sending the session token cannot leak to a +/// third-party BYOK endpoint. A definite `NotByok` always refreshes (it only +/// ever routes to the session endpoint); a definite `Byok` never does. pub fn session_token_auth_gate( is_session_based_method: bool, model_byok: ModelByok, @@ -388,40 +248,21 @@ pub fn session_token_auth_gate( } pub const AUTH_ERROR_SESSION_EXPIRED: &str = - "Session expired. Run `grok login` to re-authenticate."; + "Session expired. Run `kigi login` to re-authenticate."; -pub const AUTH_ERROR_API_KEY: &str = "Authentication failed. Run `grok login`, set XAI_API_KEY, or add api_key to ~/.kigi/config.toml."; +pub const AUTH_ERROR_API_KEY: &str = "Authentication failed. Run `kigi login`, set XAI_API_KEY, or add api_key to ~/.kigi/config.toml."; -/// Next ACP method id when `cached_token` cannot proceed (missing / expired / -/// legacy WebLogin), or `None` when fallthrough is forbidden. -/// -/// Unpinned: prefer non-interactive `xai.api_key` when advertiseable, else -/// interactive `grok.com`. -/// -/// Pinned `oidc`: **no** fallthrough to api_key — return `None` so the caller -/// fails auth. Pinned `api_key` should not reach this path (cached_token is -/// not advertised). -pub fn method_id_after_cached_token_unavailable( - has_external_api_key: bool, - preferred_method: Option, -) -> Option<&'static str> { - match preferred_method { - Some(PreferredAuthMethod::Oidc) | Some(PreferredAuthMethod::ApiKey) => None, - None => Some(if has_external_api_key { - XAI_API_KEY_METHOD_ID - } else { - KIGI_COM_METHOD_ID - }), +/// Next ACP method id when `cached_token` cannot proceed (missing / expired): +/// prefer non-interactive `xai.api_key` when advertiseable, else the +/// interactive device login. +pub fn method_id_after_cached_token_unavailable(has_external_api_key: bool) -> &'static str { + if has_external_api_key { + XAI_API_KEY_METHOD_ID + } else { + KIGI_COM_METHOD_ID } } -/// Error when `preferred_method=api_key` but no key/BYOK credentials exist. -pub const PREFERRED_API_KEY_UNAVAILABLE: &str = "preferred_method=api_key but no API key is configured (set XAI_API_KEY or model api_key/env_key in config.toml)."; - -/// Error when `preferred_method=oidc` but the session path cannot proceed. -pub const PREFERRED_OIDC_UNAVAILABLE: &str = - "preferred_method=oidc but no session is available. Run `grok login` to authenticate."; - pub const XAI_API_KEY_METHOD_ID: &str = "xai.api_key"; pub fn xai_api_key_auth_method() -> acp::AuthMethod { acp::AuthMethod::Agent( @@ -442,39 +283,20 @@ pub fn cached_token_auth_method() -> acp::AuthMethod { acp::AuthMethodId::new(CACHED_TOKEN_AUTH_METHOD_ID), "cached_token".to_string(), ) - .description(Some("Cached token from ~/.kigi/auth.json".to_string())), + .description(Some("Cached Kimi Code session".to_string())), ) } +/// Interactive login method id. The literal `"grok.com"` is kept for ACP +/// wire-compat with the in-repo pager (renaming is a cross-crate wire change +/// deferred to the command-surface milestone). pub const KIGI_COM_METHOD_ID: &str = "grok.com"; -/// xAI OAuth2/OIDC auth. Method id `"grok.com"` kept for ACP wire-compat. -pub fn grok_com_auth_method( - label: Option<&str>, - has_auth_provider_command: bool, -) -> acp::AuthMethod { - let name = label.unwrap_or("Grok"); - let meta = if has_auth_provider_command { - let mut m = acp::Meta::new(); - m.insert("external_provider".to_owned(), serde_json::json!(true)); - Some(m) - } else { - None - }; +/// The Kimi Code device-code login. +pub fn kimi_code_auth_method(label: Option<&str>) -> acp::AuthMethod { + let name = label.unwrap_or("Kimi Code"); acp::AuthMethod::Agent( acp::AuthMethodAgent::new(acp::AuthMethodId::new(KIGI_COM_METHOD_ID), name.to_string()) - .description(Some(format!("Sign in with {name}"))) - .meta(meta), - ) -} - -pub const OIDC_METHOD_ID: &str = "oidc"; -pub fn oidc_auth_method(issuer: &str, label: Option<&str>) -> acp::AuthMethod { - let name = label - .map(|l| l.to_string()) - .unwrap_or_else(|| format!("Single sign-on ({})", issuer)); - acp::AuthMethod::Agent( - acp::AuthMethodAgent::new(acp::AuthMethodId::new(OIDC_METHOD_ID), name.clone()) .description(Some(format!("Sign in with {name}"))), ) } @@ -487,87 +309,97 @@ mod tests { use serial_test::serial; /// When API-key credentials are advertiseable, fall through from a dead - /// `cached_token` to non-interactive `xai.api_key` (not browser OAuth). - /// Covers the both-advertised case (`has_cached_token` true at initialize - /// but session later missing/expired/legacy): advertise order still puts - /// `xai.api_key` first, while `default_auth_method_id` prefers session; - /// after session fails, this helper must still pick `xai.api_key`. + /// `cached_token` to non-interactive `xai.api_key` (not the browser). #[test] fn after_cached_token_unavailable_prefers_api_key_when_advertiseable() { assert_eq!( - method_id_after_cached_token_unavailable(true, None), - Some(XAI_API_KEY_METHOD_ID), + method_id_after_cached_token_unavailable(true), + XAI_API_KEY_METHOD_ID, ); } - /// No advertiseable API-key credentials → interactive `grok.com`. + /// No advertiseable API-key credentials → interactive device login. #[test] - fn after_cached_token_unavailable_falls_to_grok_com_without_api_key() { + fn after_cached_token_unavailable_falls_to_interactive_login() { assert_eq!( - method_id_after_cached_token_unavailable(false, None), - Some(KIGI_COM_METHOD_ID), - ); - } - - /// Pinned methods never fall through across the api_key ↔ oidc boundary. - #[test] - fn after_cached_token_unavailable_fails_closed_when_pinned() { - assert_eq!( - method_id_after_cached_token_unavailable(true, Some(PreferredAuthMethod::Oidc)), - None, - ); - assert_eq!( - method_id_after_cached_token_unavailable(true, Some(PreferredAuthMethod::ApiKey)), - None, + method_id_after_cached_token_unavailable(false), + KIGI_COM_METHOD_ID, ); } /// Classifier matrix for all auth method variants. #[test] fn auth_method_kind_classifier_matrix() { - let session_methods = [ - CACHED_TOKEN_AUTH_METHOD_ID, - KIGI_COM_METHOD_ID, - OIDC_METHOD_ID, - ]; - for method_id in session_methods { - let id = acp::AuthMethodId::new(method_id); - let kind = AuthMethodKind::from_id(&id); - assert!( - kind.is_session_based(), - "{method_id}: kind must be session-based" - ); - assert!( - is_session_based_method(&id), - "{method_id}: wrapper must agree" - ); + let session_methods = [CACHED_TOKEN_AUTH_METHOD_ID, KIGI_COM_METHOD_ID]; + for id in session_methods { + let kind = AuthMethodKind::from_id(&acp::AuthMethodId::new(id)); + assert!(kind.is_session_based(), "{id} must be session-based"); + assert!(!kind.is_api_key(), "{id} must not be api-key"); } - let api_id = acp::AuthMethodId::new(XAI_API_KEY_METHOD_ID); - let api_kind = AuthMethodKind::from_id(&api_id); - assert!(!api_kind.is_session_based()); - assert!(api_kind.is_api_key()); - assert!(!is_session_based_method(&api_id)); - assert!(!is_session_based_method(&acp::AuthMethodId::new( - "unknown-method" - ))); + let api = AuthMethodKind::from_id(&acp::AuthMethodId::new(XAI_API_KEY_METHOD_ID)); + assert!(api.is_api_key()); + assert!(!api.is_session_based()); + assert!(!api.needs_interactive_login()); + let unknown = AuthMethodKind::from_id(&acp::AuthMethodId::new("who-knows")); + assert_eq!(unknown, AuthMethodKind::Unknown); + assert!(!unknown.is_session_based()); + // Only the interactive login needs a browser. + assert!( + AuthMethodKind::from_id(&acp::AuthMethodId::new(KIGI_COM_METHOD_ID)) + .needs_interactive_login() + ); + assert!( + !AuthMethodKind::from_id(&acp::AuthMethodId::new(CACHED_TOKEN_AUTH_METHOD_ID)) + .needs_interactive_login() + ); } - use kigi_test_support::EnvGuard; + #[test] + fn session_token_auth_gate_matrix() { + // Session method + NotByok → refresh. + assert!(session_token_auth_gate(true, ModelByok::NotByok, false)); + // Session method + Byok → never. + assert!(!session_token_auth_gate(true, ModelByok::Byok, true)); + // Session method + Unknown → only on first-party endpoints. + assert!(session_token_auth_gate(true, ModelByok::Unknown, true)); + assert!(!session_token_auth_gate(true, ModelByok::Unknown, false)); + // Non-session method → never. + assert!(!session_token_auth_gate(false, ModelByok::NotByok, true)); + } - // ── Helpers ───────────────────────────────────────────────────────── + /// RAII guard restoring an env var on drop (panic-safe). + struct EnvGuard { + key: &'static str, + prev: Option, + } + impl EnvGuard { + fn set(key: &'static str, value: &str) -> Self { + let prev = std::env::var(key).ok(); + unsafe { std::env::set_var(key, value) }; + Self { key, prev } + } + fn unset(key: &'static str) -> Self { + let prev = std::env::var(key).ok(); + unsafe { std::env::remove_var(key) }; + Self { key, prev } + } + } + impl Drop for EnvGuard { + fn drop(&mut self) { + unsafe { + match self.prev.take() { + Some(v) => std::env::set_var(self.key, v), + None => std::env::remove_var(self.key), + } + } + } + } - /// Default inputs to `build_auth_methods` representing a session-only user - /// with no API key anywhere. Tests override only the fields they care - /// about. fn default_inputs() -> AuthMethodsBuildInputs<'static> { AuthMethodsBuildInputs { has_external_api_key: false, has_cached_token: false, - has_enterprise_oidc: false, - enterprise_oidc_issuer: None, login_label: None, - has_auth_provider_command: false, - preferred_method: None, } } @@ -586,520 +418,105 @@ mod tests { methods.first().map(|m| AuthMethodKind::from_id(m.id())) } - // build_auth_methods regression: pin production call-site ordering. - // Reordering so `xai.api_key` is after login methods must fail the tests below. - - /// BYOK with only per-model `env_key` must list `xai.api_key` first. + /// BYOK: `xai.api_key` must be `auth_methods.first()`; deferred-to-last + /// ordering sends per-model-key users to the login screen. #[test] - fn enterprise_byok_first_method_is_xai_api_key() { - let inputs = AuthMethodsBuildInputs { - has_external_api_key: true, // enterprise user with resolved per-model env_key - has_cached_token: false, + fn byok_first_method_is_xai_api_key() { + let built = build_auth_methods(AuthMethodsBuildInputs { + has_external_api_key: true, ..default_inputs() - }; - let built = build_auth_methods(inputs); - + }); assert_eq!( - first_kind(&built.methods), - Some(AuthMethodKind::XaiApiKey), - "BYOK enterprise-style: auth_methods.first() MUST be xai.api_key \ - (deferred-to-last ordering sends users to the login screen)", + method_ids(&built), + vec![XAI_API_KEY_METHOD_ID, KIGI_COM_METHOD_ID] ); - assert_eq!( - built - .default_auth_method_id - .as_ref() - .map(|id| id.0.as_ref()), - Some(XAI_API_KEY_METHOD_ID), - ); - // Cross-check with the pager-side predicate: the first method must - // not require interactive login, which is the exact condition the - // pager's `startup_auth_metadata()` uses. + assert_eq!(default_id(&built), Some(XAI_API_KEY_METHOD_ID)); assert!( !AuthMethodKind::from_id(built.methods[0].id()).needs_interactive_login(), - "first method MUST NOT need interactive login when xai.api_key is available", + "auth_methods.first() must not need interactive login" ); } - /// BYOK + cached session token: xai.api_key stays first in the methods - /// list (skips login screen), but `default_auth_method_id` is - /// `cached_token` (keeps OIDC refresh alive). + /// API key + cached session: `xai.api_key` stays first in the advertised + /// list, but the session wins the default (refresh stays alive). #[test] fn byok_with_cached_token_keeps_xai_api_key_first() { - let inputs = AuthMethodsBuildInputs { + let built = build_auth_methods(AuthMethodsBuildInputs { has_external_api_key: true, has_cached_token: true, ..default_inputs() - }; - let built = build_auth_methods(inputs); - + }); assert_eq!( - first_kind(&built.methods), - Some(AuthMethodKind::XaiApiKey), - "xai.api_key MUST precede cached_token in advertised order", - ); - // Sanity: cached_token still appears, just second. - assert!( - built - .methods - .iter() - .any(|m| AuthMethodKind::from_id(m.id()) == AuthMethodKind::CachedToken), - "cached_token must still be advertised when present", - ); - // cached_token wins for default_auth_method_id (keeps OIDC refresh alive). - assert_eq!( - built - .default_auth_method_id - .as_ref() - .map(|id| id.0.as_ref()), - Some(CACHED_TOKEN_AUTH_METHOD_ID), + method_ids(&built), + vec![ + XAI_API_KEY_METHOD_ID, + CACHED_TOKEN_AUTH_METHOD_ID, + KIGI_COM_METHOD_ID + ] ); + assert_eq!(default_id(&built), Some(CACHED_TOKEN_AUTH_METHOD_ID)); } - /// Session-only user (no API key anywhere): cached_token first, then - /// `grok.com` — `auth_methods.first()` does NOT need interactive login, - /// so this user also skips the login screen at startup. + /// Session-only user: cached_token first, interactive login as fallback. #[test] fn session_only_user_first_method_is_cached_token() { - let inputs = AuthMethodsBuildInputs { - has_external_api_key: false, + let built = build_auth_methods(AuthMethodsBuildInputs { has_cached_token: true, ..default_inputs() - }; - let built = build_auth_methods(inputs); - - assert_eq!( - first_kind(&built.methods), - Some(AuthMethodKind::CachedToken) - ); - assert_eq!( - built - .default_auth_method_id - .as_ref() - .map(|id| id.0.as_ref()), - Some(CACHED_TOKEN_AUTH_METHOD_ID), - ); - } - - /// Brand-new user (no API key, no cached token): only `grok.com` is - /// advertised, and the pager will (correctly) show the login screen. - /// `default_auth_method_id` is None so the pager falls back to the - /// advertised login method. - #[test] - fn fresh_user_only_advertises_grok_com_and_requires_login() { - let built = build_auth_methods(default_inputs()); - - assert_eq!(first_kind(&built.methods), Some(AuthMethodKind::GrokCom)); - assert!(built.default_auth_method_id.is_none()); - assert_eq!(built.methods.len(), 1); - } - - /// Enterprise OIDC replaces `grok.com` (mutually exclusive). xai.api_key, - /// when present, still leads. - #[test] - fn enterprise_oidc_replaces_grok_com_but_xai_api_key_still_first() { - let inputs = AuthMethodsBuildInputs { - has_external_api_key: true, - has_cached_token: false, - has_enterprise_oidc: true, - enterprise_oidc_issuer: Some("https://sso.example.com"), - ..default_inputs() - }; - let built = build_auth_methods(inputs); - - assert_eq!(first_kind(&built.methods), Some(AuthMethodKind::XaiApiKey)); - assert!( - built - .methods - .iter() - .any(|m| AuthMethodKind::from_id(m.id()) == AuthMethodKind::Oidc), - "oidc must be advertised when has_enterprise_oidc", - ); - assert!( - !built - .methods - .iter() - .any(|m| AuthMethodKind::from_id(m.id()) == AuthMethodKind::GrokCom), - "grok.com and oidc are mutually exclusive", - ); - } - - /// `has_auth_provider_command` is plumbed through to the `grok.com` method - /// as `meta.external_provider = true`. Pinning this here so the pager's - /// `AuthStartMode::Command` path keeps working. - #[test] - fn auth_provider_command_sets_external_provider_meta() { - let inputs = AuthMethodsBuildInputs { - has_auth_provider_command: true, - login_label: Some("Acme Corp"), - ..default_inputs() - }; - let built = build_auth_methods(inputs); - - let grok = built - .methods - .iter() - .find(|m| AuthMethodKind::from_id(m.id()) == AuthMethodKind::GrokCom) - .expect("grok.com must be advertised"); - assert_eq!(grok.name(), "Acme Corp"); - let meta = grok.meta().expect("meta should be set"); - assert_eq!( - meta.get("external_provider").and_then(|v| v.as_bool()), - Some(true), - ); - } - - // ── End-to-end: enterprise TOML -> resolved models -> build_auth_methods ─ - - /// END-TO-END REGRESSION TEST: parses the literal enterprise-style - /// `~/.kigi/config.toml` skeleton from the bug report, walks it through - /// the same predicate (`should_advertise_xai_api_key`) and the same - /// list-builder (`build_auth_methods`) that `MvpAgent::initialize()` uses - /// in production, and asserts that `auth_methods.first()` is `xai.api_key` - /// (which causes the pager to skip the login screen). - /// - /// This is the test that *would have caught* that regression -- if you mentally - /// re-introduce that bug (push xai.api_key LAST when has_external_api_key - /// && !global env var), this test fails because `first_kind` is no longer - /// `XaiApiKey`. - #[test] - #[serial] - fn enterprise_byok_config_does_not_require_login() { - const TEST_ENV_VAR: &str = "TEST_ENTERPRISE_REGRESSION_AUTH_TOKEN"; - - // Make sure no global key is masking the per-model path we're trying - // to exercise. Held until end-of-scope so we restore on panic too. - let _global = EnvGuard::unset(XAI_API_KEY_ENV_VAR); - - let dm = crate::models::default_model(); - let toml: toml::Value = toml::from_str(&format!( - r#" - [model."{dm}"] - model = "{dm}" - base_url = "https://inference.example.com/v1" - context_window = 200000 - env_key = "{TEST_ENV_VAR}" - "#, - )) - .unwrap(); - let cfg = Config::new_from_toml_cfg(&toml).expect("config should parse"); - let models = resolve_model_list(&cfg, None); - let model = models.get(dm).expect("enterprise-style model should exist"); - assert_eq!( - model.env_key.as_ref().map(|k| k.names()), - Some(vec![TEST_ENV_VAR]) - ); - - // Without the env var present, has_own_credentials() returns false, - // the predicate returns false, and the builder advertises only the - // login method. Confirms the predicate isn't trivially true. - { - let _unset = EnvGuard::unset(TEST_ENV_VAR); - let has_external_api_key = should_advertise_xai_api_key(false, models.values()); - assert!(!has_external_api_key); - let built = build_auth_methods(AuthMethodsBuildInputs { - has_external_api_key, - ..default_inputs() - }); - assert_ne!( - first_kind(&built.methods), - Some(AuthMethodKind::XaiApiKey), - "without env_key resolved, xai.api_key must NOT be advertised first", - ); - } - - // With the env var present (the actual enterprise scenario), the predicate - // returns true and the builder MUST put `xai.api_key` first so the - // pager's `startup_auth_metadata()` returns `needs_login = false`. - { - let _set = EnvGuard::set(TEST_ENV_VAR, "enterprise-secret-token"); - let has_external_api_key = should_advertise_xai_api_key(false, models.values()); - assert!(has_external_api_key); - let built = build_auth_methods(AuthMethodsBuildInputs { - has_external_api_key, - // Realistic enterprise user: no cached session token, default - // grok.com login (no enterprise OIDC). - has_cached_token: false, - ..default_inputs() - }); - assert_eq!( - first_kind(&built.methods), - Some(AuthMethodKind::XaiApiKey), - "BYOK: xai.api_key must be auth_methods.first(); deferred-to-last \ - ordering sends enterprise users to the login screen", - ); - assert!( - !AuthMethodKind::from_id(built.methods[0].id()).needs_interactive_login(), - "auth_methods.first() MUST NOT need interactive login -- this \ - is the exact predicate the pager's startup_auth_metadata() \ - uses to decide whether to show the login screen", - ); - } - } - - /// `XAI_API_KEY` alone (no per-model creds) also triggers - /// advertising `xai.api_key` as the first method. Historical "external - /// key" path; covered here so the predicate keeps treating env-var-only - /// users the same as per-model users. - #[test] - #[serial] - fn global_external_api_key_advertises_xai_api_key_first() { - let _set = EnvGuard::set(XAI_API_KEY_ENV_VAR, "xai-external-key"); - let cfg = Config::default(); - let models = resolve_model_list(&cfg, None); - let has_external_api_key = should_advertise_xai_api_key(false, models.values()); - assert!(has_external_api_key); - let built = build_auth_methods(AuthMethodsBuildInputs { - has_external_api_key, - ..default_inputs() - }); - assert_eq!(first_kind(&built.methods), Some(AuthMethodKind::XaiApiKey)); - } - - /// Admin kill switch (`disable_api_key_auth`): the predicate must return - /// false even when credentials are available everywhere (global env var - /// AND per-model env_key), so the builder never advertises `xai.api_key` - /// and the pager sends the user to the deployment's login method instead. - #[test] - #[serial] - fn disable_api_key_auth_suppresses_xai_api_key_method() { - let _set = EnvGuard::set(XAI_API_KEY_ENV_VAR, "xai-external-key"); - let cfg = Config::default(); - let models = resolve_model_list(&cfg, None); - - // Flag off: today's behavior (advertised first). - assert!(should_advertise_xai_api_key(false, models.values())); - - // Flag on: never advertised, regardless of credentials. - let has_external_api_key = should_advertise_xai_api_key(true, models.values()); - assert!(!has_external_api_key); - let built = build_auth_methods(AuthMethodsBuildInputs { - has_external_api_key, - ..default_inputs() - }); - assert!( - !built - .methods - .iter() - .any(|m| AuthMethodKind::from_id(m.id()) == AuthMethodKind::XaiApiKey), - "xai.api_key must not be advertised when disable_api_key_auth is set", - ); - assert_eq!( - first_kind(&built.methods), - Some(AuthMethodKind::GrokCom), - "with api-key auth disabled and no cached token, the login method \ - must lead so the pager requires interactive login", - ); - assert!(built.default_auth_method_id.is_none()); - } - - /// Legacy `KIGI_CODE_XAI_API_KEY` env var is accepted as a fallback - /// when `XAI_API_KEY` is not set, ensuring existing deployments keep working. - #[test] - #[serial] - fn legacy_env_var_fallback_advertises_xai_api_key() { - let _unset_new = EnvGuard::unset(XAI_API_KEY_ENV_VAR); - let _set_legacy = EnvGuard::set(LEGACY_XAI_API_KEY_ENV_VAR, "xai-legacy-key"); - assert!(has_xai_api_key_env()); - assert_eq!(read_xai_api_key_env().unwrap(), "xai-legacy-key"); - - let cfg = Config::default(); - let models = resolve_model_list(&cfg, None); - let has_external_api_key = should_advertise_xai_api_key(false, models.values()); - assert!(has_external_api_key); - } - - /// When both `XAI_API_KEY` and `KIGI_CODE_XAI_API_KEY` are set, - /// the new name takes precedence. - #[test] - #[serial] - fn new_env_var_takes_precedence_over_legacy() { - let _new = EnvGuard::set(XAI_API_KEY_ENV_VAR, "new-key"); - let _legacy = EnvGuard::set(LEGACY_XAI_API_KEY_ENV_VAR, "old-key"); - assert_eq!(read_xai_api_key_env().unwrap(), "new-key"); - } - - // -- grok login --legacy regression coverage ------------------------ - // - // `grok login --legacy` produces a GrokAuth with `auth_mode: WebLogin`, - // `oidc_issuer: None`, and no `expires_at` (30-day hardcoded TTL). - // When this token is present via the `KIGI_AUTH` env var (or via legacy - // scope fallback in auth.json), `AuthManager::new` returns it from - // `current()`, feeding `has_cached_token = true` into `build_auth_methods`. - // This puts `cached_token` first so `startup_auth_metadata()` returns - // `needs_login = false` -- legacy users get frictionless auth, no login - // screen. - // - // This test pins the env-var path (highest priority in AuthManager) end- - // to-end. A regression in KIGI_AUTH JSON parsing or in auth method - // ordering would send legacy-token users to the login screen. - - /// END-TO-END REGRESSION TEST: a legacy auth token (WebLogin, no - /// expires_at) present in the `KIGI_AUTH` env var, with no other auth - /// available, MUST be loaded by `AuthManager` and cause `build_auth_methods` - /// to advertise `cached_token` first. The pager therefore skips the login - /// screen (frictionless legacy auth). This behavior works; the test - /// prevents regressions. - #[test] - #[serial] - fn grok_login_legacy_token_does_not_require_login() { - use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig}; - - // Ensure clean slate for "no other auth available". - let _g1 = EnvGuard::unset("KIGI_AUTH_PATH"); - let _g2 = EnvGuard::unset(XAI_API_KEY_ENV_VAR); - - // Construct a legacy-style token exactly as `grok login --legacy` - // produces: WebLogin mode, no OIDC fields, no refresh_token, no - // expires_at (is_expired falls back to 30-day age check). - let legacy_token = GrokAuth { - key: "legacy-relay-token".into(), - auth_mode: AuthMode::WebLogin, - create_time: chrono::Utc::now(), - user_id: "legacy-user".into(), - email: Some("legacy@example.com".into()), - oidc_issuer: None, - oidc_client_id: None, - refresh_token: None, - expires_at: None, - ..GrokAuth::test_default() - }; - - // Provide it via KIGI_AUTH env var (highest priority code path in - // AuthManager::new). This is the "legacy auth token exists in the env" - // case with no other auth. - let legacy_json = serde_json::to_string(&legacy_token).expect("serialize legacy token"); - let _g = EnvGuard::set("KIGI_AUTH", &legacy_json); - - // AuthManager picks it up from the env var directly (no file needed). - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let mgr = AuthManager::new(dir.path(), cfg); - let current = mgr.current(); - assert!( - current.is_some(), - "legacy token in KIGI_AUTH env MUST be loaded directly -- if this fails, \ - users with legacy auth in env would be sent to the login screen", - ); - assert_eq!( - current.as_ref().unwrap().key, - "legacy-relay-token", - "loaded token must match the one injected via env", - ); - - // derive has_cached_token exactly as initialize() does. - let has_cached_token = mgr.current().is_some(); - assert!(has_cached_token); - - // With only this legacy token (no xai api key), first method must be - // cached_token so pager skips login screen. - let built = build_auth_methods(AuthMethodsBuildInputs { - has_external_api_key: false, - has_cached_token, - ..default_inputs() - }); - - assert_eq!( - first_kind(&built.methods), - Some(AuthMethodKind::CachedToken), - "legacy token in env: cached_token MUST be auth_methods.first() \ - (pager startup_auth_metadata returns needs_login=false)", - ); - assert!( - !AuthMethodKind::from_id(built.methods[0].id()).needs_interactive_login(), - "auth_methods.first() MUST NOT need interactive login when legacy token \ - is in env -- prevents login screen regression", - ); - assert_eq!( - built - .default_auth_method_id - .as_ref() - .map(|id| id.0.as_ref()), - Some(CACHED_TOKEN_AUTH_METHOD_ID), - ); - } - - /// Negative case for the legacy flow: when auth.json does NOT contain a - /// legacy-scope entry, AuthManager::current() is None, - /// has_cached_token is false, and build_auth_methods advertises only - /// the login method. This pins the predicate's "no" answer so the test - /// above isn't trivially passing. - #[test] - #[serial] - fn no_legacy_token_means_no_cached_token_advertised() { - use crate::auth::{AuthManager, GrokComConfig}; - - let _g1 = EnvGuard::unset("KIGI_AUTH"); - let _g2 = EnvGuard::unset("KIGI_AUTH_PATH"); - - let dir = tempfile::tempdir().unwrap(); - // No auth.json in the tempdir. - let cfg = GrokComConfig::default(); - let mgr = AuthManager::new(dir.path(), cfg); - assert!(mgr.current().is_none()); - - let built = build_auth_methods(AuthMethodsBuildInputs { - has_external_api_key: false, - has_cached_token: mgr.current().is_some(), - ..default_inputs() - }); - assert_eq!( - first_kind(&built.methods), - Some(AuthMethodKind::GrokCom), - "no cached token AND no api key: pager must show login (grok.com first)", - ); - } - - // ── preferred_method pin (fail-closed) ────────────────────────────── - - #[test] - fn pin_api_key_with_key_only_advertises_api_key() { - let built = build_auth_methods(AuthMethodsBuildInputs { - has_external_api_key: true, - has_cached_token: true, - preferred_method: Some(PreferredAuthMethod::ApiKey), - ..default_inputs() - }); - assert_eq!(method_ids(&built), vec![XAI_API_KEY_METHOD_ID]); - assert_eq!(default_id(&built), Some(XAI_API_KEY_METHOD_ID)); - } - - #[test] - fn pin_api_key_without_key_fails_closed_even_with_session() { - let built = build_auth_methods(AuthMethodsBuildInputs { - has_external_api_key: false, - has_cached_token: true, - preferred_method: Some(PreferredAuthMethod::ApiKey), - ..default_inputs() - }); - assert!(built.methods.is_empty()); - assert!(built.default_auth_method_id.is_none()); - } - - #[test] - fn pin_oidc_with_session_hides_api_key() { - let built = build_auth_methods(AuthMethodsBuildInputs { - has_external_api_key: true, - has_cached_token: true, - preferred_method: Some(PreferredAuthMethod::Oidc), - ..default_inputs() }); assert_eq!( method_ids(&built), vec![CACHED_TOKEN_AUTH_METHOD_ID, KIGI_COM_METHOD_ID] ); assert_eq!(default_id(&built), Some(CACHED_TOKEN_AUTH_METHOD_ID)); + assert_eq!( + first_kind(&built.methods), + Some(AuthMethodKind::CachedToken) + ); } + /// Fresh user: only the interactive login is advertised; no default + /// method (login required). #[test] - fn pin_oidc_without_session_is_interactive_only() { + fn fresh_user_only_advertises_interactive_login() { + let built = build_auth_methods(default_inputs()); + assert_eq!(method_ids(&built), vec![KIGI_COM_METHOD_ID]); + assert_eq!(default_id(&built), None); + } + + /// `XAI_API_KEY` alone (no per-model creds) triggers advertising + /// `xai.api_key` as the first method. + #[test] + #[serial] + fn global_external_api_key_advertises_xai_api_key_first() { + let _set = EnvGuard::set(XAI_API_KEY_ENV_VAR, "xai-external-key"); + let cfg = Config::default(); + let models = resolve_model_list(&cfg, None); + let has_external_api_key = should_advertise_xai_api_key(models.values()); + assert!(has_external_api_key); let built = build_auth_methods(AuthMethodsBuildInputs { - has_external_api_key: true, - has_cached_token: false, - preferred_method: Some(PreferredAuthMethod::Oidc), + has_external_api_key, ..default_inputs() }); - assert_eq!(method_ids(&built), vec![KIGI_COM_METHOD_ID]); - assert!(built.default_auth_method_id.is_none()); + assert_eq!(first_kind(&built.methods), Some(AuthMethodKind::XaiApiKey)); + } + + /// Legacy env var fallback keeps working. + #[test] + #[serial] + fn legacy_env_var_fallback_advertises_xai_api_key() { + let _unset = EnvGuard::unset(XAI_API_KEY_ENV_VAR); + let _set = EnvGuard::set(LEGACY_XAI_API_KEY_ENV_VAR, "legacy-key"); + assert!(has_xai_api_key_env()); + assert_eq!(read_xai_api_key_env().unwrap(), "legacy-key"); + } + + /// The new env var takes precedence over the legacy one. + #[test] + #[serial] + fn new_env_var_takes_precedence_over_legacy() { + let _new = EnvGuard::set(XAI_API_KEY_ENV_VAR, "new-key"); + let _legacy = EnvGuard::set(LEGACY_XAI_API_KEY_ENV_VAR, "legacy-key"); + assert_eq!(read_xai_api_key_env().unwrap(), "new-key"); } } diff --git a/crates/codegen/kigi-shell/src/agent/config.rs b/crates/codegen/kigi-shell/src/agent/config.rs index 1497c5c..cad117b 100644 --- a/crates/codegen/kigi-shell/src/agent/config.rs +++ b/crates/codegen/kigi-shell/src/agent/config.rs @@ -1,5 +1,5 @@ use crate::agent::auth_method::ModelByok; -use crate::auth::{AuthManager, GrokComConfig, OidcAuthConfig}; +use crate::auth::{AuthManager, KimiCodeConfig}; use crate::remote::DEFAULT_CONTEXT_WINDOW; use crate::{config::StorageMode, sampling::ApiBackend, tools::config::ShellToolsetConfig}; use agent_client_protocol as acp; @@ -1108,7 +1108,7 @@ pub struct Config { /// Warnings from `[model.*]` parsing; surfaced by `grok inspect`. #[serde(skip)] pub model_override_warnings: Vec, - pub grok_com_config: GrokComConfig, + pub kimi_code_config: KimiCodeConfig, #[serde(default, skip_serializing_if = "Option::is_none")] pub shortcuts: Option, /// Written by the client via `config_toml_edit`; absorbed so it isn't @@ -1178,9 +1178,9 @@ pub struct Config { #[serde(default, skip_serializing)] pub managed_mcps: crate::config::ManagedMcpsConfig, /// `[auth]` alias — consumed by `expand_auth_alias` before serde. - /// Typed as `GrokComConfig` (same schema) so sub-field typos are caught. + /// Typed as `KimiCodeConfig` (same schema) so sub-field typos are caught. #[serde(default, skip_serializing)] - pub auth: Option, + pub auth: Option, /// `[desktop]` section — owned by grok-desktop (Electron app), opaque to the CLI agent. #[serde(default, skip_serializing)] pub desktop: Option, @@ -1525,7 +1525,7 @@ impl Default for Config { auto_mode: AutoModeConfig::default(), config_models: IndexMap::new(), model_override_warnings: Vec::new(), - grok_com_config: GrokComConfig::default(), + kimi_code_config: KimiCodeConfig::default(), shortcuts: None, hints: None, ui: UiConfig::default(), @@ -1620,13 +1620,12 @@ impl Config { } Ok(()) } - /// Build an `AuthManager` with the configured proxy URL applied. + /// Build an `AuthManager` for this configuration. pub fn create_auth_manager(&self) -> AuthManager { AuthManager::new( &crate::util::kigi_home::kigi_home(), - self.grok_com_config.clone(), + self.kimi_code_config.clone(), ) - .with_proxy_base_url(&self.endpoints.proxy_url()) } /// Deserialize the merged `base` document, also returning the ignored key /// paths whose top-level key appears in `user_config`. Paths outside it @@ -1679,12 +1678,6 @@ impl Config { } config.config_models = config_models; config.model_override_warnings = model_override_warnings; - if config.grok_com_config.oidc.is_none() { - config.grok_com_config.oidc = OidcAuthConfig::from_env(); - } - if config.grok_com_config.oidc.is_none() && config.grok_com_config.oauth2.is_none() { - config.grok_com_config.oauth2 = crate::auth::OAuth2ProviderConfig::from_env(); - } if config.client_version.is_none() { config.client_version = Self::default().client_version; } @@ -1818,16 +1811,16 @@ impl Config { self.resolve_runtime_fields(&ctx); crate::util::config::set_remote_campaigns_from_settings(self.remote_settings.as_ref()); } - /// If the TOML contains `[auth]`, copy its contents under `[grok_com_config]`. - /// `[grok_com_config]` takes precedence if both are present (explicit wins). + /// If the TOML contains `[auth]`, copy its contents under `[kimi_code_config]`. + /// `[kimi_code_config]` takes precedence if both are present (explicit wins). /// - /// This lets customers write the shorter `[auth.oidc]` instead of `[grok_com_config.oidc]`. + /// This lets customers write the shorter `[auth.oidc]` instead of `[kimi_code_config.oidc]`. fn expand_auth_alias(raw_config: &toml::Value) -> toml::Value { let mut config = raw_config.clone(); if let toml::Value::Table(ref mut table) = config && let Some(auth) = table.remove("auth") { - if let Some(gcc) = table.get_mut("grok_com_config") { + if let Some(gcc) = table.get_mut("kimi_code_config") { if let (toml::Value::Table(gcc_table), toml::Value::Table(auth_table)) = (gcc, &auth) { @@ -1836,7 +1829,7 @@ impl Config { } } } else { - table.insert("grok_com_config".to_owned(), auth); + table.insert("kimi_code_config".to_owned(), auth); } } config @@ -2320,18 +2313,6 @@ impl Config { .default(true) .resolve() } - /// Resolve whether to use grok's default OAuth2 (xAI auth.x.ai). - /// - /// Enterprise OIDC (`oidc` in config.toml) always wins — this only gates - /// the default xAI OAuth2 fallback when no enterprise OIDC is configured. - /// - /// Priority: `--oauth` > KIGI_OAUTH_ENABLED env > default (true = OAuth). - pub fn resolve_grok_oauth(&self, cli_oidc: Option) -> Resolved { - BoolFlag::env("KIGI_OAUTH_ENABLED") - .cli(cli_oidc) - .default(true) - .resolve() - } /// Resolve whether to spawn the per-`Ready`-client transport /// liveness pollers and the session-actor `StatusDispatcher`. /// @@ -3918,42 +3899,6 @@ pub fn resolve_credentials(model: &ModelEntry, session_key: Option<&str>) -> Res auth_scheme, } } -/// `disable_api_key_auth` at the credential seam: swap a first-party xAI API -/// key for the IdP session (absent => request fails => forces login). BYOK -/// (non-xAI `base_url`) is untouched; no-op when the switch is off. -pub fn enforce_disable_api_key_auth( - creds: &mut ResolvedCredentials, - disable_api_key_auth: bool, - session_key: Option<&str>, -) { - if disable_api_key_auth - && creds.auth_type == kigi_chat_state::AuthType::ApiKey - && crate::util::is_first_party_xai_url(&creds.base_url) - { - creds.auth_type = kigi_chat_state::AuthType::SessionToken; - creds.api_key = session_key.map(str::to_owned); - kigi_log::unified_log::debug( - "auth: kill switch blocked a first-party API key at the credential seam", - None, - Some(serde_json::json!( - { "replaced_with_session" : session_key.is_some(), "base_url" : creds - .base_url, } - )), - ); - } -} -/// Resolve credentials for an auxiliary sampling path (web search, image -/// description) with the first-party API-key kill switch applied, so these -/// paths honor `disable_api_key_auth` exactly like the main chat path. -fn resolve_credentials_enforced( - entry: &ModelEntry, - session_key: Option<&str>, - disable_api_key_auth: bool, -) -> ResolvedCredentials { - let mut credentials = resolve_credentials(entry, session_key); - enforce_disable_api_key_auth(&mut credentials, disable_api_key_auth, session_key); - credentials -} /// Try to resolve credentials for a model by loading the effective config. /// Returns `None` (with a warning) if config loading, parsing, or model @@ -3971,12 +3916,7 @@ pub fn try_resolve_model_credentials( .ok()?; let models = resolve_model_list(&cfg, None); let entry = find_model_by_id(&models, model_id)?; - let mut credentials = resolve_credentials(entry, session_key); - enforce_disable_api_key_auth( - &mut credentials, - cfg.grok_com_config.api_key_auth_disabled(), - session_key, - ); + let credentials = resolve_credentials(entry, session_key); Some(credentials) } /// Per-model auth facts (BYOK status + auth scheme) from one effective-config @@ -4045,13 +3985,12 @@ pub fn resolve_aux_model_sampling_config( models: &IndexMap, endpoints: &EndpointsConfig, session_key: Option<&str>, - disable_api_key_auth: bool, alpha_test_key: Option, client_version: Option, ) -> Option { let catalog_entry = find_model_by_id(models, model_id).cloned(); if let Some(entry) = &catalog_entry { - let credentials = resolve_credentials_enforced(entry, session_key, disable_api_key_auth); + let credentials = resolve_credentials(entry, session_key); let sampler = sampling_config_for_model( entry, credentials, @@ -4108,7 +4047,7 @@ pub fn resolve_aux_model_sampling_config( env_key: None, api_base_url: None, }; - let credentials = resolve_credentials_enforced(&entry, session_key, disable_api_key_auth); + let credentials = resolve_credentials(&entry, session_key); let sampler = sampling_config_for_model( &entry, credentials, @@ -4240,11 +4179,7 @@ pub fn sampling_config_for_model( /// URL-derived header logic at the shell boundary so callers downstream see a /// single homogenous header bag. /// -/// * cli-chat-proxy bases get `X-XAI-Token-Auth` and -/// `x-authenticateresponse` headers (mirrors the inline match in the legacy -/// `sampling::Client::new` on `is_cli_chat_proxy_url`). -/// * With the optional non-production feature, matching first-party hosts may -/// get an extra access header from the corresponding key argument. +/// * First-party bases get the client-mode header. /// /// Existing entries are never overwritten so callers can pre-set a value. pub fn inject_url_derived_headers( @@ -4253,12 +4188,6 @@ pub fn inject_url_derived_headers( base_url: &str, ) { if crate::util::is_cli_chat_proxy_url(base_url) { - headers - .entry("X-XAI-Token-Auth".to_string()) - .or_insert_with(|| "xai-grok-cli".to_string()); - headers - .entry("x-authenticateresponse".to_string()) - .or_insert_with(|| "authenticate-response".to_string()); headers .entry(crate::http::CLIENT_MODE_HEADER.to_string()) .or_insert_with(|| crate::http::process_client_mode().to_string()); @@ -4289,7 +4218,6 @@ pub fn resolve_model_to_sampling_config( fn resolve_hidden_default_web_search_sampling_config( model_id: &str, session_key: Option<&str>, - disable_api_key_auth: bool, alpha_test_key: Option, client_version: Option, endpoints: &EndpointsConfig, @@ -4331,7 +4259,7 @@ fn resolve_hidden_default_web_search_sampling_config( env_key: None, api_base_url: None, }; - let credentials = resolve_credentials_enforced(&entry, session_key, disable_api_key_auth); + let credentials = resolve_credentials(&entry, session_key); sampling_config_for_model( &entry, credentials, @@ -4345,13 +4273,12 @@ pub fn resolve_web_search_sampling_config( model_id: &str, models: &IndexMap, session_key: Option<&str>, - disable_api_key_auth: bool, alpha_test_key: Option, client_version: Option, endpoints: &EndpointsConfig, ) -> Option { let resolved = if let Some(entry) = find_model_by_id(models, model_id).cloned() { - let credentials = resolve_credentials_enforced(&entry, session_key, disable_api_key_auth); + let credentials = resolve_credentials(&entry, session_key); Some(sampling_config_for_model( &entry, credentials, @@ -4364,7 +4291,6 @@ pub fn resolve_web_search_sampling_config( Some(resolve_hidden_default_web_search_sampling_config( model_id, session_key, - disable_api_key_auth, alpha_test_key, client_version, endpoints, @@ -4643,28 +4569,21 @@ reasoning_effort = "low" } } #[test] - fn inject_url_derived_headers_adds_proxy_headers_for_cli_chat_proxy_url() { + fn inject_url_derived_headers_adds_client_mode_for_first_party_url() { let mut headers = IndexMap::new(); inject_url_derived_headers( &mut headers, None, kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url, ); - assert_eq!( - headers.get("X-XAI-Token-Auth").map(String::as_str), - Some("xai-grok-cli") - ); - assert_eq!( - headers.get("x-authenticateresponse").map(String::as_str), - Some("authenticate-response") - ); + assert!(headers.get(crate::http::CLIENT_MODE_HEADER).is_some()); + assert!(headers.get("X-XAI-Token-Auth").is_none()); } #[test] - fn inject_url_derived_headers_skips_proxy_headers_for_external_url() { + fn inject_url_derived_headers_skips_headers_for_external_url() { let mut headers = IndexMap::new(); - inject_url_derived_headers(&mut headers, None, "https://api.x.ai/v1"); - assert!(headers.get("X-XAI-Token-Auth").is_none()); - assert!(headers.get("x-authenticateresponse").is_none()); + inject_url_derived_headers(&mut headers, None, "https://api.example.com/v1"); + assert!(headers.get(crate::http::CLIENT_MODE_HEADER).is_none()); } #[test] fn inject_url_derived_headers_preserves_caller_extra_headers() { @@ -4679,24 +4598,6 @@ reasoning_effort = "low" headers.get("x-custom-byok").map(String::as_str), Some("value") ); - assert_eq!( - headers.get("X-XAI-Token-Auth").map(String::as_str), - Some("xai-grok-cli") - ); - } - #[test] - fn inject_url_derived_headers_does_not_overwrite_existing_entries() { - let mut headers = IndexMap::new(); - headers.insert("X-XAI-Token-Auth".to_string(), "caller-set".to_string()); - inject_url_derived_headers( - &mut headers, - None, - kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url, - ); - assert_eq!( - headers.get("X-XAI-Token-Auth").map(String::as_str), - Some("caller-set"), - ); } #[test] fn parses_toolset_overrides() { @@ -4834,7 +4735,6 @@ reasoning_effort = "low" crate::models::default_web_search_model(), &IndexMap::new(), Some("session-token"), - false, None, None, &endpoints, @@ -4891,51 +4791,14 @@ reasoning_effort = "low" None, ), ); - let resolved = resolve_aux_model_sampling_config( - "grok-build", - &catalog, - &endpoints, - None, - false, - None, - None, - ) - .expect("override entry has an API key, so resolution succeeds"); + let resolved = + resolve_aux_model_sampling_config("grok-build", &catalog, &endpoints, None, None, None) + .expect("override entry has an API key, so resolution succeeds"); assert_eq!(resolved.model, "v9m-rl-learnability-tp8"); assert_eq!(resolved.base_url, "https://vendor.example/v1"); assert_eq!(resolved.api_key.as_deref(), Some("vendor-key")); } #[test] - fn web_search_disable_api_key_auth_swaps_first_party_key_for_session() { - let endpoints = EndpointsConfig::default(); - let mut models = IndexMap::new(); - models.insert( - "ws-model".to_string(), - test_model_entry( - "ws-model", - "https://api.x.ai/v1", - Some("first-party-key"), - None, - None, - ), - ); - let resolved = resolve_web_search_sampling_config( - "ws-model", - &models, - Some("session-token"), - true, - None, - None, - &endpoints, - ) - .expect("web search model should resolve"); - assert_eq!( - resolved.api_key.as_deref(), - Some("session-token"), - "first-party API key must be swapped for the session token when disabled" - ); - } - #[test] fn parses_model_api_key() { let raw_config: toml::Value = toml::from_str( r#" @@ -5351,12 +5214,9 @@ reasoning_effort = "low" config.base_url, kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url ); - assert_eq!( - config - .extra_headers - .get("X-XAI-Token-Auth") - .map(String::as_str), - Some("xai-grok-cli") + assert!( + config.extra_headers.get("X-XAI-Token-Auth").is_none(), + "the xAI token-auth marker header must be gone" ); } /// Regression: without a session key, `resolve_credentials` falls through @@ -5376,75 +5236,6 @@ reasoning_effort = "low" auth_scheme: Default::default(), } } - /// `disable_api_key_auth` kill switch (Claude `forceLoginMethod` parity). - #[test] - fn enforce_disable_api_key_auth_blocks_first_party_only() { - use kigi_chat_state::AuthType; - let mut creds = api_key_creds("https://api.x.ai/v1"); - enforce_disable_api_key_auth(&mut creds, false, Some("session-jwt")); - assert_eq!(creds.auth_type, AuthType::ApiKey); - assert_eq!(creds.api_key.as_deref(), Some("xai-secret")); - let mut creds = api_key_creds("https://api.x.ai/v1"); - enforce_disable_api_key_auth(&mut creds, true, Some("session-jwt")); - assert_eq!(creds.auth_type, AuthType::SessionToken); - assert_eq!(creds.api_key.as_deref(), Some("session-jwt")); - let mut creds = api_key_creds("https://api.x.ai/v1"); - enforce_disable_api_key_auth(&mut creds, true, None); - assert_eq!(creds.auth_type, AuthType::SessionToken); - assert_eq!(creds.api_key, None); - let mut creds = api_key_creds("https://api.example.com/v1"); - enforce_disable_api_key_auth(&mut creds, true, Some("session-jwt")); - assert_eq!(creds.auth_type, AuthType::ApiKey); - assert_eq!(creds.api_key.as_deref(), Some("xai-secret")); - let mut creds = ResolvedCredentials { - auth_type: AuthType::SessionToken, - ..api_key_creds("https://api.x.ai/v1") - }; - enforce_disable_api_key_auth(&mut creds, true, Some("session-jwt")); - assert_eq!(creds.auth_type, AuthType::SessionToken); - } - /// Regression for the OVERRIDE_MODEL kill-switch bypass: a first-party model - /// with its own api_key resolves to `ApiKey` (priority 1, beating the - /// session), and the kill switch — now applied inside - /// `try_resolve_model_credentials` — swaps it for the session token. BYOK - /// (non-x.ai) own keys are preserved. (`try_resolve_model_credentials` - /// loads global config, so this exercises its resolve + enforce core.) - #[test] - fn try_resolve_model_credentials_swaps_first_party_own_key_under_kill_switch() { - use kigi_chat_state::AuthType; - let entry = test_model_entry( - "m", - "https://api.x.ai/v1", - Some("xai-model-key"), - None, - None, - ); - let mut creds = resolve_credentials(&entry, Some("session-jwt")); - assert_eq!( - creds.auth_type, - AuthType::ApiKey, - "own key wins over session" - ); - assert_eq!(creds.api_key.as_deref(), Some("xai-model-key")); - enforce_disable_api_key_auth(&mut creds, true, Some("session-jwt")); - assert_eq!( - creds.auth_type, - AuthType::SessionToken, - "swapped under switch" - ); - assert_eq!(creds.api_key.as_deref(), Some("session-jwt")); - let byok = test_model_entry( - "b", - "https://api.example.com/v1", - Some("sk-byok"), - None, - None, - ); - let mut byok_creds = resolve_credentials(&byok, Some("session-jwt")); - enforce_disable_api_key_auth(&mut byok_creds, true, Some("session-jwt")); - assert_eq!(byok_creds.auth_type, AuthType::ApiKey); - assert_eq!(byok_creds.api_key.as_deref(), Some("sk-byok")); - } #[test] fn x_api_key_auth_scheme_flows_from_config_to_sampler() { let mut model = test_model_entry( @@ -6631,124 +6422,16 @@ reasoning_effort = "low" let info = ModelInfo::from_config(&entry); assert_eq!(info.inference_idle_timeout_secs, Some(120)); } + /// The `[auth]` alias and the explicit `[kimi_code_config]` table both + /// deserialize (the auth block currently carries no per-deployment + /// options; the alias machinery is retained for future knobs). #[test] - fn auth_alias_maps_to_grok_com_config() { - let raw: toml::Value = toml::from_str( - r#" - [auth.oidc] - issuer = "https://example.okta.com" - client_id = "test-id" - "#, - ) - .unwrap(); - let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse"); - let oidc = cfg.grok_com_config.oidc.expect("oidc should be set"); - assert_eq!(oidc.issuer, "https://example.okta.com"); - assert_eq!(oidc.client_id, "test-id"); - } - #[test] - fn grok_com_config_still_works() { - let raw: toml::Value = toml::from_str( - r#" - [grok_com_config.oidc] - issuer = "https://example.okta.com" - client_id = "test-id" - "#, - ) - .unwrap(); - let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse"); - let oidc = cfg.grok_com_config.oidc.expect("oidc should be set"); - assert_eq!(oidc.issuer, "https://example.okta.com"); - } - /// `disable_api_key_auth` plumbs through the `[auth]` alias, and absent - /// means None (opt-in knob, zero impact by default). - #[test] - fn disable_api_key_auth_parses_from_auth_alias() { - let absent = Config::new_from_toml_cfg(&toml::from_str("").unwrap()).unwrap(); - assert_eq!(absent.grok_com_config.disable_api_key_auth, None); - let raw: toml::Value = toml::from_str( - r#" - [auth] - disable_api_key_auth = true - "#, - ) - .unwrap(); - let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse"); - assert_eq!(cfg.grok_com_config.disable_api_key_auth, Some(true)); - } - /// `force_login_team_uuid` parses a string (pin), array (any-of), or `[]` - /// (fail closed); absent => None. - #[test] - fn force_login_team_uuid_parses_string_and_array() { - use crate::auth::ForceLoginTeam; - let absent = Config::new_from_toml_cfg(&toml::from_str("").unwrap()).unwrap(); - assert_eq!(absent.grok_com_config.force_login_team_uuid, None); - let raw: toml::Value = toml::from_str( - r#" - [auth] - force_login_team_uuid = "team-abc" - "#, - ) - .unwrap(); - let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse"); - assert_eq!( - cfg.grok_com_config.force_login_team_uuid, - Some(ForceLoginTeam::Single("team-abc".into())), - ); - let raw: toml::Value = toml::from_str( - r#" - [grok_com_config] - force_login_team_uuid = ["team-a", "team-b"] - "#, - ) - .unwrap(); - let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse"); - assert_eq!( - cfg.grok_com_config.force_login_team_uuid, - Some(ForceLoginTeam::AnyOf(vec![ - "team-a".into(), - "team-b".into() - ])), - ); - let raw: toml::Value = toml::from_str( - r#" - [auth] - force_login_team_uuid = [] - "#, - ) - .unwrap(); - let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse"); - assert_eq!( - cfg.grok_com_config.force_login_team_uuid, - Some(ForceLoginTeam::AnyOf(vec![])), - ); - } - /// Pinning a team via `force_login_team_uuid` implies API-key auth is - /// disabled even without an explicit `disable_api_key_auth` (team - /// membership can't be verified from a bare API key, so it needs IdP login). - #[test] - fn force_login_team_uuid_implies_api_key_auth_disabled() { - use crate::auth::{ForceLoginTeam, GrokComConfig}; - let base = GrokComConfig { - disable_api_key_auth: None, - force_login_team_uuid: None, - ..GrokComConfig::default() - }; - assert!(!base.api_key_auth_disabled()); - assert!( - GrokComConfig { - disable_api_key_auth: Some(true), - ..base.clone() - } - .api_key_auth_disabled() - ); - assert!( - GrokComConfig { - force_login_team_uuid: Some(ForceLoginTeam::Single("team-x".into())), - ..base - } - .api_key_auth_disabled() - ); + fn auth_alias_and_kimi_code_config_tables_parse() { + for body in ["[auth]\n", "[kimi_code_config]\n"] { + let raw: toml::Value = toml::from_str(body).unwrap(); + let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse"); + assert_eq!(cfg.kimi_code_config.auth_scope(), "oauth/kimi-code"); + } } fn resolve_models_from_toml( toml_str: &str, @@ -8662,11 +8345,7 @@ agent_type = "cursor" persistent_shell = true [shortcuts] ctrl_k = "search" - [grok_com_config] - token_header = "test" - [auth.oidc] - issuer = "https://sso.corp.com" - client_id = "abc123" + [kimi_code_config] [storage] cleanup_ttl_days = 7 [permission] diff --git a/crates/codegen/kigi-shell/src/agent/feedback_client.rs b/crates/codegen/kigi-shell/src/agent/feedback_client.rs index 62b637e..9f9a41c 100644 --- a/crates/codegen/kigi-shell/src/agent/feedback_client.rs +++ b/crates/codegen/kigi-shell/src/agent/feedback_client.rs @@ -320,14 +320,14 @@ pub struct FeedbackClient { http: reqwest::Client, client: reqwest_middleware::ClientWithMiddleware, base_url: String, - credentials: crate::util::grok_auth_credentials::GrokAuthCredentials, + credentials: crate::util::kigi_auth_credentials::KigiAuthCredentials, session_id: Option, } impl FeedbackClient { pub fn new(base_url: impl Into, user_token: Option) -> Self { let http = crate::http::shared_client(); - let credentials = crate::util::grok_auth_credentials::GrokAuthCredentials::new(user_token); + let credentials = crate::util::kigi_auth_credentials::KigiAuthCredentials::new(user_token); let client = Self::build_middleware_client(&http, &credentials); Self { http, @@ -361,7 +361,7 @@ impl FeedbackClient { base_url: impl Into, user_token: Option, ) -> Self { - let credentials = crate::util::grok_auth_credentials::GrokAuthCredentials::new(user_token); + let credentials = crate::util::kigi_auth_credentials::KigiAuthCredentials::new(user_token); let client = Self::build_middleware_client(&http, &credentials); Self { http, @@ -398,7 +398,7 @@ impl FeedbackClient { fn build_middleware_client( http: &reqwest::Client, - credentials: &crate::util::grok_auth_credentials::GrokAuthCredentials, + credentials: &crate::util::kigi_auth_credentials::KigiAuthCredentials, ) -> reqwest_middleware::ClientWithMiddleware { let provider = Self::make_auth_provider(credentials); // max_retries=0: the middleware stamps the auth header but does NOT @@ -415,7 +415,7 @@ impl FeedbackClient { } fn make_auth_provider( - credentials: &crate::util::grok_auth_credentials::GrokAuthCredentials, + credentials: &crate::util::kigi_auth_credentials::KigiAuthCredentials, ) -> Arc { if let Some(am) = credentials.auth_manager() { Arc::new( @@ -1118,7 +1118,7 @@ mod forbidden_tests { #[cfg(test)] mod auth_refresh_tests { use super::*; - use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig}; + use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig}; use axum::{Router, routing::get}; use chrono::{Duration, Utc}; use std::net::SocketAddr; @@ -1156,14 +1156,14 @@ mod auth_refresh_tests { let (addr, _server) = start_server(router).await; let dir = tempfile::tempdir().unwrap(); - let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - am.hot_swap(GrokAuth { + let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + am.hot_swap(KimiAuth { key: "fresh-from-auth-manager".into(), auth_mode: AuthMode::ApiKey, create_time: Utc::now(), user_id: "user-42".into(), expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }); let client = FeedbackClient::new( @@ -1208,14 +1208,14 @@ mod auth_refresh_tests { let (addr, _server) = start_server(router).await; let dir = tempfile::tempdir().unwrap(); - let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - let fresh = GrokAuth { + let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + let fresh = KimiAuth { key: "fresh-from-auth-manager".into(), auth_mode: AuthMode::ApiKey, create_time: Utc::now(), user_id: "user-42".into(), expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }; am.hot_swap(fresh); @@ -1247,16 +1247,14 @@ mod auth_refresh_tests { _reason: crate::auth::refresh::RefreshReason, ) -> crate::auth::refresh::RefreshOutcome { self.calls.fetch_add(1, Ordering::SeqCst); - crate::auth::refresh::RefreshOutcome::Success(Box::new(GrokAuth { + crate::auth::refresh::RefreshOutcome::Success(Box::new(KimiAuth { key: "fresh-from-refresher".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, create_time: Utc::now(), user_id: "user-42".into(), refresh_token: Some("rt-fresh".into()), expires_at: Some(Utc::now() + Duration::hours(1)), - oidc_issuer: Some("https://issuer.example".into()), - oidc_client_id: Some("test-client".into()), - ..GrokAuth::test_default() + ..KimiAuth::test_default() })) } } @@ -1266,34 +1264,30 @@ mod auth_refresh_tests { #[tokio::test] async fn try_refresh_credentials_picks_up_disk_rotation_without_hitting_idp() { let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); + let cfg = KimiCodeConfig::default(); let scope = cfg.auth_scope(); let am = Arc::new(AuthManager::new(dir.path(), cfg)); // In-memory: stale token (the one the server rejected). - am.hot_swap(GrokAuth { + am.hot_swap(KimiAuth { key: "stale-rejected".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, create_time: Utc::now() - Duration::hours(2), user_id: "user-42".into(), refresh_token: Some("rt-stale".into()), expires_at: Some(Utc::now() + Duration::hours(1)), - oidc_issuer: Some("https://issuer.example".into()), - oidc_client_id: Some("test-client".into()), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }); // Disk: a sibling already rotated to a fresh token. - let disk_auth = GrokAuth { + let disk_auth = KimiAuth { key: "fresh-from-sibling-on-disk".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, create_time: Utc::now(), user_id: "user-42".into(), refresh_token: Some("rt-fresh".into()), expires_at: Some(Utc::now() + Duration::hours(1)), - oidc_issuer: Some("https://issuer.example".into()), - oidc_client_id: Some("test-client".into()), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }; let mut store = std::collections::BTreeMap::new(); store.insert(scope, disk_auth); @@ -1329,15 +1323,15 @@ mod auth_refresh_tests { #[tokio::test] async fn try_refresh_credentials_returns_false_on_terminal_failure() { let dir = tempfile::tempdir().unwrap(); - let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); + let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); // LegacySession: no refresh_token, no recovery possible. - am.hot_swap(GrokAuth { + am.hot_swap(KimiAuth { key: "legacy-rejected".into(), - auth_mode: AuthMode::WebLogin, + auth_mode: AuthMode::OAuth, create_time: Utc::now() - Duration::days(60), user_id: "user-42".into(), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }); let client = FeedbackClient::new("http://example/v1", Some("legacy-rejected".into())) diff --git a/crates/codegen/kigi-shell/src/agent/init.rs b/crates/codegen/kigi-shell/src/agent/init.rs index 2b22e48..e96a1c8 100644 --- a/crates/codegen/kigi-shell/src/agent/init.rs +++ b/crates/codegen/kigi-shell/src/agent/init.rs @@ -79,7 +79,7 @@ fn resolve_config(cfg: &AgentConfig, auth_manager: &AuthManager) -> AgentConfig // thread the result into `cfg.remote_settings` skip this entirely. if cfg.remote_settings.is_none() && let Some(handle) = - crate::agent::models::start_early_prefetch(Some(cfg.grok_com_config.clone())) + crate::agent::models::start_early_prefetch(Some(cfg.kimi_code_config.clone())) { match handle.join() { Ok(result) => { @@ -103,9 +103,9 @@ fn resolve_config(cfg: &AgentConfig, auth_manager: &AuthManager) -> AgentConfig { cfg.storage_mode = StorageMode::resolve(None, cfg.remote_settings.as_ref()); } - // Writeback talks to the code backend; requires grok.com auth. + // Writeback talks to the code backend; requires a Kimi Code session. if cfg.storage_mode == StorageMode::Writeback - && !auth_manager.current().is_some_and(|a| a.is_xai_auth()) + && !auth_manager.current().is_some_and(|a| a.is_session_auth()) { tracing::info!("Writeback is disabled: requires auth with grok.com"); cfg.storage_mode = StorageMode::Local; diff --git a/crates/codegen/kigi-shell/src/agent/mod.rs b/crates/codegen/kigi-shell/src/agent/mod.rs index 65839f8..f3ee54a 100644 --- a/crates/codegen/kigi-shell/src/agent/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/mod.rs @@ -18,7 +18,6 @@ pub mod server; pub mod session_config; pub mod session_registry_client; pub(crate) mod subagent; -pub(crate) mod subscription_check; pub(crate) mod update_chunk_merge; pub use mvp_agent::MvpAgent; diff --git a/crates/codegen/kigi-shell/src/agent/models.rs b/crates/codegen/kigi-shell/src/agent/models.rs index ce4790e..c9da5be 100644 --- a/crates/codegen/kigi-shell/src/agent/models.rs +++ b/crates/codegen/kigi-shell/src/agent/models.rs @@ -10,7 +10,7 @@ use chrono::{DateTime, Duration as ChronoDuration, Utc}; use indexmap::IndexMap; use crate::agent::config::{self, ModelEntry, resolve_credentials, sampling_config_for_model}; -use crate::auth::{AuthManager, GrokAuth, GrokComConfig}; +use crate::auth::{AuthManager, KimiAuth, KimiCodeConfig}; use crate::remote::{FetchModelsResult, fetch_models_blocking}; use crate::sampling::SamplerConfig as SamplingConfig; use globset::{Glob, GlobSet, GlobSetBuilder}; @@ -151,7 +151,7 @@ struct Inner { impl Default for ModelsManager { fn default() -> Self { let kigi_home = crate::util::kigi_home::kigi_home(); - let auth_manager = Arc::new(AuthManager::new(&kigi_home, GrokComConfig::default())); + let auth_manager = Arc::new(AuthManager::new(&kigi_home, KimiCodeConfig::default())); Self::new( None, IndexMap::new(), @@ -1397,7 +1397,7 @@ fn build_prefetched_map( /// Fetch remote models. Checks disk cache first; persists after fetch. pub(crate) fn prefetch_models_blocking( endpoints: &config::EndpointsConfig, - auth: Option<&GrokAuth>, + auth: Option<&KimiAuth>, fetch_auth: ModelFetchAuth, ) -> Option> { prefetch_models_blocking_gated( @@ -1414,7 +1414,7 @@ pub(crate) fn prefetch_models_blocking( /// decisions cannot disagree mid-startup. pub(crate) fn prefetch_models_and_settings_blocking( endpoints: &config::EndpointsConfig, - auth: Option<&GrokAuth>, + auth: Option<&KimiAuth>, fetch_auth: ModelFetchAuth, ) -> ( Option>, @@ -1441,7 +1441,7 @@ pub(crate) fn prefetch_models_and_settings_blocking( /// knob once for both halves. fn prefetch_models_blocking_gated( endpoints: &config::EndpointsConfig, - auth: Option<&GrokAuth>, + auth: Option<&KimiAuth>, fetch_auth: ModelFetchAuth, remote_fetch_enabled: bool, ) -> Option> { @@ -1499,12 +1499,12 @@ pub struct EarlyPrefetchResult { pub type EarlyPrefetchHandle = std::thread::JoinHandle; struct PrefetchEnv { - auth: Option, + auth: Option, endpoints: config::EndpointsConfig, model_fetch_auth: ModelFetchAuth, } -fn resolve_prefetch_env_with_auth(auth: Option) -> Option { +fn resolve_prefetch_env_with_auth(auth: Option) -> Option { let _timer = crate::instrumentation_timer!("startup.early_prefetch_launch"); // Config-aware (not env-only) so the prefetch can't leak the bearer to api.x.ai. let mut endpoints = config::EndpointsConfig::from_effective_config(); @@ -1529,7 +1529,7 @@ fn resolve_prefetch_env_with_auth(auth: Option) -> Option /// `deployment_key` would re-arm the prefetch — and with it the `/v1/settings` /// fetch and the deployment-config sync on the prefetch thread. fn resolve_prefetch_env_from_parts( - auth: Option, + auth: Option, endpoints: config::EndpointsConfig, remote_fetch_enabled: bool, ) -> Option { @@ -1554,9 +1554,9 @@ fn resolve_prefetch_env_from_parts( }) } -fn resolve_prefetch_env(grok_com_config: Option) -> Option { +fn resolve_prefetch_env(kimi_code_config: Option) -> Option { let kigi_home = crate::util::kigi_home::kigi_home(); - let auth_manager = AuthManager::new(&kigi_home, grok_com_config.unwrap_or_default()); + let auth_manager = AuthManager::new(&kigi_home, kimi_code_config.unwrap_or_default()); let auth = auth_manager.current(); resolve_prefetch_env_with_auth(auth) } @@ -1566,7 +1566,7 @@ fn resolve_prefetch_env(grok_com_config: Option) -> Option) -> Option { +pub fn start_early_prefetch_with_auth(auth: Option) -> Option { let env = resolve_prefetch_env_with_auth(auth)?; Some(spawn_prefetch_thread(env)) } @@ -1575,8 +1575,10 @@ pub fn start_early_prefetch_with_auth(auth: Option) -> Option) -> Option { - let env = resolve_prefetch_env(grok_com_config)?; +pub fn start_early_prefetch( + kimi_code_config: Option, +) -> Option { + let env = resolve_prefetch_env(kimi_code_config)?; Some(spawn_prefetch_thread(env)) } @@ -1969,7 +1971,7 @@ pub(crate) fn validate_selectable( /// Async wrapper around `prefetch_models_blocking`. pub(crate) async fn fetch_models_async( endpoints: config::EndpointsConfig, - auth: Option, + auth: Option, fetch_auth: ModelFetchAuth, ) -> Option> { tokio::task::spawn_blocking(move || { @@ -1991,7 +1993,7 @@ mod tests { // Use a temp dir so AuthManager finds no credentials — ensures // refresh_async bails at the auth check without needing a tokio runtime. let tmp = std::env::temp_dir().join("grok-test-models-manager"); - let auth_manager = Arc::new(AuthManager::new(&tmp, GrokComConfig::default())); + let auth_manager = Arc::new(AuthManager::new(&tmp, KimiCodeConfig::default())); ModelsManager::new( None, IndexMap::new(), @@ -2233,7 +2235,7 @@ mod tests { #[test] fn current_reasoning_effort_seeded_from_config() { let tmp = std::env::temp_dir().join("grok-test-models-manager-seed"); - let auth_manager = Arc::new(AuthManager::new(&tmp, GrokComConfig::default())); + let auth_manager = Arc::new(AuthManager::new(&tmp, KimiCodeConfig::default())); let mut cfg = config::Config::default(); cfg.models.default_reasoning_effort = Some(ReasoningEffort::Xhigh); let mgr = ModelsManager::new( @@ -2401,7 +2403,7 @@ mod tests { // The internal getters read those derived fields. let tmp = std::env::temp_dir().join("grok-test-models-manager-menu-only"); - let auth_manager = Arc::new(AuthManager::new(&tmp, GrokComConfig::default())); + let auth_manager = Arc::new(AuthManager::new(&tmp, KimiCodeConfig::default())); let mgr = ModelsManager::new( None, catalog, @@ -3222,7 +3224,7 @@ mod tests { }; assert!( resolve_prefetch_env_from_parts( - Some(GrokAuth::test_default()), + Some(KimiAuth::test_default()), endpoints.clone(), false, ) diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs index ade55ae..a3d2404 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs @@ -48,31 +48,6 @@ impl acp::Agent for MvpAgent { ); }); kigi_workspace::trust::migrate_legacy_hook_trust(); - if let Some(auth) = self.auth_manager.current() { - let user_id = auth.user_id.trim(); - let needs_user_info = user_id.is_empty() - || user_id.eq_ignore_ascii_case("unknown"); - kigi_log::unified_log::info( - "auth init user_info check", - None, - Some( - serde_json::json!( - { "user_id" : user_id, "needs_user_info" : needs_user_info, - "key_prefix" : crate ::auth::token_suffix(& auth.key), - "rt_prefix" : auth.refresh_token.as_deref().map(crate - ::auth::token_suffix), } - ), - ), - ); - if needs_user_info && let Err(e) = self.auth_manager.update(auth).await { - tracing::warn!( - "Failed to refresh user info from proxy during new_session: {}", e - ); - } - } - if !self.tier_allowed.get() && let Some(auth) = self.auth_manager.current() { - self.enforce_grok_code_access(&auth).await; - } self.maybe_sync_bundle_in_background(false); let mut client_type = arguments .meta @@ -186,8 +161,7 @@ impl acp::Agent for MvpAgent { ), ), ); - if !self.cfg.borrow().grok_com_config.api_key_auth_disabled() - && auth_method::read_xai_api_key_env().is_err() + if auth_method::read_xai_api_key_env().is_err() && let Some(api_key) = crate::auth::read_api_key( &crate::util::kigi_home::kigi_home(), ) @@ -200,33 +174,8 @@ impl acp::Agent for MvpAgent { None, ); } - let disable_api_key_auth = self - .cfg - .borrow() - .grok_com_config - .api_key_auth_disabled(); - { - let cfg = self.cfg.borrow(); - let gc = &cfg.grok_com_config; - if disable_api_key_auth || gc.force_login_team_uuid.is_some() { - kigi_log::unified_log::info( - "auth: enterprise login policy active", - None, - Some( - serde_json::json!( - { "force_login_team_uuid" : gc.force_login_team_uuid.as_ref() - .map(| t | format!("{t:?}")), "disable_api_key_auth_knob" : - gc.disable_api_key_auth, "api_key_auth_disabled" : - disable_api_key_auth, } - ), - ), - ); - } - } - let has_external_api_key = auth_method::should_advertise_xai_api_key( - disable_api_key_auth, - self.models_manager.models().values(), - ); + let has_external_api_key = + auth_method::should_advertise_xai_api_key(self.models_manager.models().values()); let init_has_current = self.auth_manager.current().is_some(); let init_is_expired = self.auth_manager.is_expired(); kigi_log::unified_log::info( @@ -267,58 +216,11 @@ impl acp::Agent for MvpAgent { ); } } - let ( - login_label, - has_auth_provider, - has_enterprise_oidc, - enterprise_oidc_issuer, - ) = { - let cfg = self.cfg.borrow(); - let issuer = cfg.grok_com_config.oidc.as_ref().map(|o| o.issuer.clone()); - ( - cfg.grok_com_config.auth_provider_label.clone(), - cfg.grok_com_config.auth_provider_command.is_some(), - cfg.grok_com_config.oidc.is_some(), - issuer, - ) - }; - if has_enterprise_oidc { - let issuer = enterprise_oidc_issuer - .as_deref() - .expect( - "enterprise_oidc_issuer must be Some when has_enterprise_oidc is true", - ); - tracing::info!( - issuer = % issuer, "auth: advertising enterprise OIDC auth method", - ); - kigi_log::unified_log::info( - "auth: advertising enterprise OIDC auth method", - None, - Some(serde_json::json!({ "issuer" : issuer })), - ); - } else { - tracing::info!( - label = ? login_label, has_auth_provider, - "auth: advertising grok.com auth method", - ); - } - let preferred_method = self.cfg.borrow().grok_com_config.preferred_method; - let has_external_api_key = match preferred_method { - Some(crate::auth::PreferredAuthMethod::Oidc) => false, - _ => has_external_api_key, - }; - let has_cached_token = match preferred_method { - Some(crate::auth::PreferredAuthMethod::ApiKey) => false, - _ => has_cached_token, - }; + tracing::info!("auth: advertising Kimi Code device login auth method"); let built = auth_method::build_auth_methods(auth_method::AuthMethodsBuildInputs { has_external_api_key, has_cached_token, - has_enterprise_oidc, - enterprise_oidc_issuer: enterprise_oidc_issuer.as_deref(), - login_label: login_label.as_deref(), - has_auth_provider_command: has_auth_provider, - preferred_method, + login_label: None, }); let auth_methods = built.methods; kigi_log::unified_log::info( @@ -329,9 +231,8 @@ impl acp::Agent for MvpAgent { { "kigi_home" : crate ::util::kigi_home::kigi_home().display() .to_string(), "HOME" : std::env::var("HOME").unwrap_or_else(| _ | "(unset)".into()), "has_external_api_key" : has_external_api_key, - "disable_api_key_auth" : disable_api_key_auth, "has_cached_token" : - has_cached_token, "has_enterprise_oidc" : has_enterprise_oidc, - "init_has_current" : init_has_current, "init_is_expired" : + "has_cached_token" : + has_cached_token, "init_has_current" : init_has_current, "init_is_expired" : init_is_expired, "auth_mode" : self.auth_manager.current().map(| a | format!("{:?}", a.auth_mode)), "methods" : auth_methods.iter().map(| m | m.id().0.as_ref()).collect::< Vec < _ >> (), @@ -438,39 +339,8 @@ impl acp::Agent for MvpAgent { None, Some(serde_json::json!({ "method" : arguments.method_id.0.as_ref() })), ); - if let Some(preferred) = self.cfg.borrow().grok_com_config.preferred_method { - let kind = auth_method::AuthMethodKind::from_id(&arguments.method_id); - let allowed = match preferred { - crate::auth::PreferredAuthMethod::ApiKey => kind.is_api_key(), - crate::auth::PreferredAuthMethod::Oidc => kind.is_session_based(), - }; - if !allowed { - let msg = match preferred { - crate::auth::PreferredAuthMethod::ApiKey => { - auth_method::PREFERRED_API_KEY_UNAVAILABLE - } - crate::auth::PreferredAuthMethod::Oidc => { - "preferred_method=oidc; API-key auth is not allowed." - } - }; - emit_login_span( - false, - arguments.method_id.0.as_ref(), - None, - Some("preferred_method_mismatch"), - ); - return Err(acp::Error::auth_required().data(msg)); - } - } match arguments.method_id.0.as_ref() { auth_method::XAI_API_KEY_METHOD_ID => { - if self.cfg.borrow().grok_com_config.api_key_auth_disabled() { - emit_login_span(false, "api_key", None, Some("disabled_by_admin")); - return Err( - acp::Error::auth_required() - .data("API-key auth is disabled by your administrator."), - ); - } let mut sampling_config = self.sampling_config.borrow_mut(); if sampling_config.api_key.is_none() { if let Ok(api_key) = auth_method::read_xai_api_key_env() { @@ -516,82 +386,23 @@ impl acp::Agent for MvpAgent { return self .authenticate( acp::AuthenticateRequest::new( - acp::AuthMethodId::new(auth_method::OIDC_METHOD_ID), + acp::AuthMethodId::new(auth_method::KIGI_COM_METHOD_ID), ) .meta(arguments.meta), ) .await; } - let current_auth = self.auth_manager.current(); - let has_current = current_auth.is_some(); + let has_current = self.auth_manager.current().is_some(); let is_expired = self.auth_manager.is_expired(); - let is_devbox = crate::auth::devbox_login::is_devbox_environment(); - let is_legacy = current_auth - .as_ref() - .is_some_and(|a| a.auth_mode == crate::auth::AuthMode::WebLogin); kigi_log::unified_log::info( "auth cached_token check", None, Some( serde_json::json!( - { "has_current" : has_current, "is_expired" : is_expired, - "is_devbox" : is_devbox, "is_legacy" : is_legacy, } + { "has_current" : has_current, "is_expired" : is_expired, } ), ), ); - let pin_blocks_oidc_mint = matches!( - self.cfg.borrow().grok_com_config.preferred_method, Some(crate - ::auth::PreferredAuthMethod::ApiKey) - ); - if is_devbox && is_legacy && !pin_blocks_oidc_mint { - kigi_log::unified_log::info( - "auth cached_token: devbox legacy migration starting", - None, - None, - ); - match crate::auth::devbox_login::mint_devbox_auth(&self.auth_manager) - .await - { - Ok(new_auth) => { - match self - .auth_manager - .save_without_enrichment(new_auth) - .await - { - Ok(_) => { - if let Err(e) = self - .auth_manager - .remove_scope(crate::auth::LEGACY_AUTH_SCOPE) - { - tracing::warn!( - error = ? e, - "auth: failed to remove legacy scope (non-fatal)" - ); - } - kigi_log::unified_log::info( - "auth cached_token: devbox legacy migration succeeded", - None, - None, - ); - } - Err(e) => { - kigi_log::unified_log::warn( - "auth cached_token: devbox migration save failed", - None, - Some(serde_json::json!({ "error" : e.to_string() })), - ); - } - } - } - Err(e) => { - kigi_log::unified_log::warn( - "auth cached_token: devbox mint failed, will reject legacy token", - None, - Some(serde_json::json!({ "error" : format!("{e}") })), - ); - } - } - } let Some(auth) = self.auth_manager.current() else { let message = if self.auth_manager.is_expired() { "Session expired, re-authentication required" @@ -610,34 +421,8 @@ impl acp::Agent for MvpAgent { .authenticate_after_cached_token_unavailable(arguments) .await; }; - if auth.auth_mode == crate::auth::AuthMode::WebLogin { - tracing::info!("auth: rejecting legacy WebLogin token"); - kigi_log::unified_log::warn( - "auth cached_token legacy rejected", - None, - Some( - serde_json::json!( - { "auth_mode" : format!("{:?}", auth.auth_mode) } - ), - ), - ); - self.auth_manager.clear_in_memory(); - if let Err(e) = self - .auth_manager - .remove_scope(crate::auth::LEGACY_AUTH_SCOPE) - { - tracing::warn!( - error = ? e, - "auth: failed to remove legacy scope during WebLogin rejection (non-fatal)" - ); - } - return self - .authenticate_after_cached_token_unavailable(arguments) - .await; - } self.refresh_remote_settings(&auth).await; self.emit_settings_update_notification(); - self.enforce_grok_code_access(&auth).await; self.maybe_sync_bundle_in_background(false); { let mut sampling_config = self.sampling_config.borrow_mut(); @@ -660,13 +445,12 @@ impl acp::Agent for MvpAgent { self.maybe_fetch_post_auth_settings().await; Ok(self.auth_response_with_meta()) } - auth_method::KIGI_COM_METHOD_ID | auth_method::OIDC_METHOD_ID => { - let grok_ctx = self.auth_manager.grok_com_config(); + auth_method::KIGI_COM_METHOD_ID => { + let kimi_ctx = self.auth_manager.kimi_code_config().clone(); let auth_meta = AuthRequestMeta::from_json(arguments.meta.as_ref()); tracing::info!( method = arguments.method_id.0.as_ref(), headless = auth_meta - .headless, reauth = auth_meta.reauth, use_oauth = auth_meta - .use_oauth, "auth: inline auth flow", + .headless, reauth = auth_meta.reauth, "auth: inline auth flow", ); kigi_log::unified_log::info( "auth: inline auth flow", @@ -674,31 +458,13 @@ impl acp::Agent for MvpAgent { Some( serde_json::json!( { "method" : arguments.method_id.0.as_ref(), "headless" : - auth_meta.headless, "reauth" : auth_meta.reauth, "use_oauth" - : auth_meta.use_oauth, } + auth_meta.headless, "reauth" : auth_meta.reauth, } ), ), ); if auth_meta.reauth { let _ = self.auth_manager.clear(); } - let cli_oauth = auth_meta.use_oauth.then_some(true); - let use_oidc = self.cfg.borrow().resolve_grok_oauth(cli_oauth); - tracing::debug!( - resolved = use_oidc.value, source = ? use_oidc.source, - "auth: method resolved" - ); - kigi_log::unified_log::debug( - "auth: method resolved", - None, - Some( - serde_json::json!( - { "use_oidc" : use_oidc.value, "source" : format!("{:?}", - use_oidc.source), } - ), - ), - ); - let login_override = auth_meta.login_override(); let (auth, _did_auth) = if !auth_meta.headless { let (url_tx, url_rx) = tokio::sync::oneshot::channel(); let (code_tx, code_rx) = tokio::sync::mpsc::channel(1); @@ -706,14 +472,13 @@ impl acp::Agent for MvpAgent { *self.auth_url_rx.borrow_mut() = Some(url_rx); let result = crate::auth::run_auth_flow_with_stderr_bridge( &self.auth_manager, - grok_ctx, + &kimi_ctx, crate::auth::AuthChannels { url_tx: Some(url_tx), code_rx, }, auth_meta.reauth, auth_meta.force_interactive, - login_override, ) .await; *self.auth_code_tx.borrow_mut() = None; @@ -722,12 +487,9 @@ impl acp::Agent for MvpAgent { } else { crate::auth::run_auth_flow( &self.auth_manager, - grok_ctx, + &kimi_ctx, auth_meta.reauth, None, - None, - None, - login_override, ) .await } @@ -757,11 +519,7 @@ impl acp::Agent for MvpAgent { self.auth_manager.hot_swap(auth.clone()); self.refresh_remote_settings(&auth).await; self.emit_settings_update_notification(); - self.enforce_grok_code_access(&auth).await; self.maybe_sync_bundle_in_background(false); - tokio::task::spawn_local( - crate::managed_config::post_login_sync(Some(auth.clone())), - ); self.set_auth_method(arguments.method_id.clone()); self.models_manager.on_auth_changed().await; if crate::agent::chat_modes::process_chat_mode_enabled() { @@ -959,10 +717,7 @@ impl acp::Agent for MvpAgent { .session_registry_client() .map(|client| crate::session::persistence::RegistryGeneratedTitleSync { client, - suppress_for_zdr: self - .auth_manager - .current_or_expired() - .is_some_and(|a| a.is_zdr_team()), + suppress_for_zdr: false, }); crate::session::persistence::new( &session_info, @@ -1239,10 +994,7 @@ impl acp::Agent for MvpAgent { .session_registry_client() .map(|client| crate::session::persistence::RegistryGeneratedTitleSync { client, - suppress_for_zdr: self - .auth_manager - .current_or_expired() - .is_some_and(|a| a.is_zdr_team()), + suppress_for_zdr: false, }); let (persistence_info, persistence) = crate::session::persistence::load_light( &session_info, @@ -2576,9 +2328,6 @@ impl acp::Agent for MvpAgent { crate::extensions::billing::handle(self, &args).await } "x.ai/share_session" => crate::extensions::share::handle(self, &args).await, - "x.ai/privacy/setCodingDataRetention" => { - crate::extensions::privacy::handle(self, &args).await - } "x.ai/rollout/survey" => { crate::extensions::rollout::handle(self, &args).await } diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs index c7dd4af..431ea17 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs @@ -28,10 +28,9 @@ impl MvpAgent { let session_key = self.auth_manager.current_or_expired().map(|a| a.key.clone()); let models = self.models_manager.models(); let endpoints = self.models_manager.endpoints(); - let (disable_api_key_auth, alpha_test_key, client_version) = { + let (alpha_test_key, client_version) = { let cfg = self.cfg.borrow(); ( - cfg.grok_com_config.api_key_auth_disabled(), cfg.endpoints.alpha_test_key.clone(), cfg.client_version.clone(), ) @@ -41,7 +40,6 @@ impl MvpAgent { &models, &endpoints, session_key.as_deref(), - disable_api_key_auth, alpha_test_key, client_version, ) { @@ -64,7 +62,7 @@ impl MvpAgent { } fn has_proxy_credentials(&self) -> bool { self.cfg.borrow().endpoints.deployment_key.is_some() - || self.auth_manager.current_or_expired().is_some_and(|a| a.is_xai_auth()) + || self.auth_manager.current_or_expired().is_some_and(|a| a.is_session_auth()) } /// `true` for session-based ACP auth methods. fn is_session_based_auth(&self) -> bool { @@ -79,7 +77,7 @@ impl MvpAgent { self.auth_method_id.store(Some(std::sync::Arc::new(id))); } /// Return auth for sync config construction. - pub(super) fn current_or_buffered_auth(&self) -> Option { + pub(super) fn current_or_buffered_auth(&self) -> Option { self.auth_manager .current() .or_else(|| { @@ -101,7 +99,7 @@ impl MvpAgent { fn has_managed_mcp_auth(&self) -> bool { self.auth_manager .current_or_expired() - .is_some_and(|a| a.is_managed_mcp_eligible()) + .is_some_and(|a| a.is_session_auth()) } /// Requires feature flag AND xAI authentication (OIDC or legacy WebLogin). pub(super) fn can_fetch_managed_mcps(&self) -> bool { @@ -195,7 +193,7 @@ impl MvpAgent { .or_else(|| auth_manager.current_or_expired().map(|a| a.key)); if !auth_manager .current_or_expired() - .is_some_and(|a| a.is_managed_mcp_eligible()) + .is_some_and(|a| a.is_session_auth()) { cache.lock().await.disable_gateway_tools(); for tx in session_txs { @@ -401,7 +399,7 @@ impl MvpAgent { let user_token = self .auth_manager .current_or_expired() - .filter(|a| a.is_xai_auth()) + .filter(|a| a.is_session_auth()) .map(|a| a.key.clone()); let cfg = self.cfg.borrow(); let base_url = cfg.endpoints.resolve_feedback_base_url(); @@ -434,7 +432,7 @@ impl MvpAgent { return None; } let auth = self.auth_manager.current_or_expired()?; - if !auth.is_xai_auth() { + if !auth.is_session_auth() { return None; } let key = auth.key.clone(); @@ -498,12 +496,6 @@ impl MvpAgent { ..crate::session::slash_commands::CommandAvailability::default() } } - /// `true` when data collection should be suppressed (team ZDR or - /// coding-data-retention opt-out). Delegates to - /// [`AuthManager::is_data_collection_disabled`]. - pub(crate) fn is_data_collection_disabled(&self) -> bool { - self.auth_manager.is_data_collection_disabled() - } /// Current client type as set by the most recent `initialize()` call. pub(crate) fn client_type(&self) -> ClientType { *self.client_type.borrow() @@ -513,8 +505,8 @@ impl MvpAgent { pub(crate) fn session_turn_number(&self, sid: &acp::SessionId) -> Option { self.session_turn_numbers.borrow().get(sid).copied() } - /// Return the current GrokAuth credentials, if authenticated and not expired. - pub(crate) fn current_auth(&self) -> Option { + /// Return the current KimiAuth credentials, if authenticated and not expired. + pub(crate) fn current_auth(&self) -> Option { self.auth_manager.current() } /// Shared plugin registry handle used by extensions for snapshot/reload. @@ -613,55 +605,21 @@ impl MvpAgent { } } /// When `cached_token` cannot proceed, prefer non-interactive `xai.api_key` - /// iff `should_advertise_xai_api_key`; otherwise `grok.com`. Returns `None` - /// when `preferred_method` is pinned (fail-closed — no cross-method fallthrough). - pub(super) fn cached_token_fallthrough_method_id( - &self, - ) -> Option { - let preferred = self.cfg.borrow().grok_com_config.preferred_method; + /// iff `should_advertise_xai_api_key`; otherwise the interactive device + /// login. + pub(super) fn cached_token_fallthrough_method_id(&self) -> acp::AuthMethodId { let id = auth_method::method_id_after_cached_token_unavailable( - auth_method::should_advertise_xai_api_key( - self.cfg.borrow().grok_com_config.api_key_auth_disabled(), - self.models_manager.models().values(), - ), - preferred, - )?; - Some(acp::AuthMethodId::new(id)) + auth_method::should_advertise_xai_api_key(self.models_manager.models().values()), + ); + acp::AuthMethodId::new(id) } - /// Shared exit for missing/expired/legacy `cached_token`: fall through with - /// `use_oauth` only when the target is interactive `grok.com`. When - /// `preferred_method` is pinned, fail instead of falling through. + /// Shared exit for missing/expired `cached_token`. pub(super) async fn authenticate_after_cached_token_unavailable( &self, arguments: acp::AuthenticateRequest, ) -> Result { - let Some(method_id) = self.cached_token_fallthrough_method_id() else { - let preferred = self.cfg.borrow().grok_com_config.preferred_method; - let msg = match preferred { - Some(crate::auth::PreferredAuthMethod::ApiKey) => { - auth_method::PREFERRED_API_KEY_UNAVAILABLE - } - _ => auth_method::PREFERRED_OIDC_UNAVAILABLE, - }; - tracing::info!( - % msg, "cached_token unavailable; preferred_method forbids fallthrough" - ); - kigi_log::unified_log::warn( - "auth cached_token fallthrough blocked by preferred_method", - None, - Some( - serde_json::json!( - { "preferred_method" : preferred.map(| p | format!("{p:?}")), } - ), - ), - ); - return Err(acp::Error::auth_required().data(msg)); - }; - let meta = if method_id.0.as_ref() == auth_method::KIGI_COM_METHOD_ID { - serde_json::json!({ "use_oauth" : true }).as_object().cloned() - } else { - arguments.meta - }; + let method_id = self.cached_token_fallthrough_method_id(); + let meta = arguments.meta; tracing::info!(fallback = % method_id.0, "cached_token fallthrough"); kigi_log::unified_log::warn( "auth cached_token fallthrough", @@ -693,7 +651,7 @@ impl MvpAgent { /// Agent-level fields materialised at startup (`worktree_type`, /// `restore_code`) are NOT re-resolved here; that requires a /// broader refactor of the init path. - pub(super) async fn refresh_remote_settings(&self, auth: &crate::auth::GrokAuth) { + pub(super) async fn refresh_remote_settings(&self, auth: &crate::auth::KimiAuth) { if !crate::util::config::resolve_remote_fetch_enabled() { tracing::debug!("post-auth settings refresh skipped: remote_fetch disabled"); return; @@ -722,7 +680,7 @@ impl MvpAgent { /// In-flight sessions are unaffected — they snapshot config at creation. pub(super) async fn refresh_settings_and_reapply( &self, - auth: &crate::auth::GrokAuth, + auth: &crate::auth::KimiAuth, ) { self.refresh_remote_settings(auth).await; let cwd = std::env::current_dir().ok(); @@ -746,7 +704,7 @@ impl MvpAgent { /// Callers own their miss logging. pub(super) async fn fetch_remote_settings( &self, - auth: crate::auth::GrokAuth, + auth: crate::auth::KimiAuth, ) -> Option { if !crate::util::config::resolve_remote_fetch_enabled() { tracing::debug!("settings fetch skipped: remote_fetch disabled"); @@ -828,29 +786,16 @@ impl MvpAgent { model: &ModelEntry, origin_client: Option, ) -> SamplingConfig { - let preferred = self.cfg.borrow().grok_com_config.preferred_method; - let session = match preferred { - Some(crate::auth::PreferredAuthMethod::ApiKey) => None, - _ if self.is_session_based_auth() => self.auth_manager.current_or_expired(), - _ => None, + let session = if self.is_session_based_auth() { + self.auth_manager.current_or_expired() + } else { + None }; let has_session_key = session.is_some(); let mut credentials = resolve_credentials( model, session.as_ref().map(|a| a.key.as_str()), ); - if matches!(preferred, Some(crate ::auth::PreferredAuthMethod::Oidc)) - && !model.has_own_credentials() - && credentials.auth_type == kigi_chat_state::AuthType::ApiKey - { - credentials.api_key = None; - credentials.auth_type = kigi_chat_state::AuthType::SessionToken; - } - crate::agent::config::enforce_disable_api_key_auth( - &mut credentials, - self.cfg.borrow().grok_com_config.api_key_auth_disabled(), - session.as_ref().map(|a| a.key.as_str()), - ); if !has_session_key && credentials.auth_type == kigi_chat_state::AuthType::ApiKey && !model.has_own_credentials() && self.is_session_based_auth() { @@ -893,7 +838,7 @@ impl MvpAgent { let user_id = self .auth_manager .current_or_expired() - .filter(|a| a.is_xai_auth()) + .filter(|a| a.is_session_auth()) .map(|a| a.user_id); let mut config = crate::agent::config::sampling_config_for_model( model, @@ -945,39 +890,6 @@ impl MvpAgent { ); (id.clone(), new_config) } - /// Whether the current session is a personal grok.com account on a gated - /// tier (free / X Basic). The Imagine tools stay advertised to the model but - /// are flagged tier-restricted so they short-circuit at call time with the - /// SuperGrok upsell prose (see `ImageGenConfig`/`VideoGenConfig`'s - /// `tier_restricted`). - /// - /// Fails **open** (returns `false`) whenever we can't positively confirm a - /// restricted personal tier — no auth yet, BYOK / API-key sessions, team - /// accounts, and an unknown/absent tier all pass. The server - /// authoritatively zero-limits Imagine for free & X Basic (429), so this - /// client gate is a UX optimization (a clean in-chat upsell instead of a - /// doomed request), never the security boundary — under-restricting is safe, - /// over-restricting would wrongly disable a paid feature. - /// - /// Mirrors the pager's cosmetic slash-command gate - /// ([`crate::tier::is_restricted_tier_name`]); the only difference is the - /// absent-tier policy (the pager hides on `None`, we fail open on `None`). - fn is_tier_restricted_capability(&self) -> bool { - let Some(auth) = self.auth_manager.current() else { - return false; - }; - if !auth.is_xai_auth() || auth.team_id.is_some() { - return false; - } - let tier = self - .cfg - .borrow() - .remote_settings - .as_ref() - .and_then(|rs| rs.subscription_tier_display.clone()) - .or_else(|| jwt_tier_claim(&auth.key)); - tier.as_deref().is_some_and(crate::tier::is_restricted_tier_name) - } /// Build image generation config. /// /// Both BYOK and session (OAuth) users go direct to `xai_api_base_url`. @@ -992,7 +904,6 @@ impl MvpAgent { let Some(ref api_key) = sampling_config.api_key else { return ImageGenConfig::Disabled; }; - let tier_restricted = self.is_tier_restricted_capability(); let cfg = self.cfg.borrow(); let base_url = cfg.endpoints.xai_api_base_url.clone(); let version = cfg @@ -1015,7 +926,7 @@ impl MvpAgent { image_gen_enabled: cfg.resolve_image_gen().value, image_edit_enabled: cfg.resolve_image_edit().value, model_override: cfg.resolve_image_gen_model_override(), - tier_restricted, + tier_restricted: false, } } /// Build deploy-service config. The tool talks directly to the deployer service. @@ -1033,7 +944,6 @@ impl MvpAgent { let Some(api_key) = self.sampling_config.borrow().api_key.clone() else { return VideoGenConfig::Disabled; }; - let tier_restricted = self.is_tier_restricted_capability(); let cfg = self.cfg.borrow(); let zdr_video_output_s3 = cfg .disable_zdr_incompatible_tools @@ -1063,7 +973,7 @@ impl MvpAgent { base_url, extra_headers: headers, zdr_video_output_s3: zdr_video_output_s3.map(Box::new), - tier_restricted, + tier_restricted: false, } } pub(super) fn prepare_web_search_sampling_config(&self) -> Option { @@ -1076,7 +986,6 @@ impl MvpAgent { &model_id, &models, session.as_ref().map(|a| a.key.as_str()), - self.cfg.borrow().grok_com_config.api_key_auth_disabled(), alpha_test_key.clone(), client_version, &self.cfg.borrow().endpoints, @@ -1224,7 +1133,6 @@ impl MvpAgent { interactive_trust_prompted: Rc::new( RefCell::new(std::collections::HashSet::new()), ), - tier_allowed: std::cell::Cell::new(true), storage_mode, default_yolo_mode, default_auto_mode, @@ -1252,9 +1160,6 @@ impl MvpAgent { subagent_coordinator: RefCell::new(subagent_coordinator), monitor_event_buffer: kigi_tools::implementations::grok_build::task::types::MonitorEventBuffer::default(), bundle_sync_in_flight: Arc::new(std::sync::atomic::AtomicBool::new(false)), - post_unblock_jwt_retry_in_flight: Arc::new( - std::sync::atomic::AtomicBool::new(false), - ), workspace_ops: RefCell::new(None), require_gateway_sessions: Rc::new( RefCell::new(std::collections::HashSet::new()), @@ -1268,11 +1173,7 @@ impl MvpAgent { #[cfg(test)] supervisor_spawn_count: std::cell::Cell::new(0), }; - instance - .auth_manager - .configure_refresher( - instance.cfg.borrow().grok_com_config.auth_provider_command.clone(), - ); + instance.auth_manager.configure_refresher(); instance } /// Handle `x.ai/internal/evict_sessions` — the leader server tells us a @@ -2189,7 +2090,7 @@ impl MvpAgent { } None => (kigi_hunk_tracker::HunkTrackerHandle::noop(), None), }; - let has_xai_auth = self.auth_manager.current().is_some_and(|a| a.is_xai_auth()); + let has_xai_auth = self.auth_manager.current().is_some_and(|a| a.is_session_auth()); let loc_tracking_enabled = hunk_tracking_enabled && has_xai_auth && (self .cfg diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs index a5aeb76..e2c7b01 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs @@ -98,75 +98,6 @@ pub(crate) fn reject_direct_hub_cloud_meta( } Ok(()) } -/// Marks a notification's meta field with `isReplay: true` for replayed session updates. -/// If `persist_data` is provided, it will be included in the meta under `x.ai/persist`. -/// Extract the numeric `tier` claim from a JWT access token (no signature -/// verification). Maps the `prod_auth.SubscriptionTier` proto enum values -/// to display-style strings that `normalize_tier` in the telemetry crate -/// will canonicalize for Mixpanel. -pub(crate) fn jwt_tier_claim(jwt: &str) -> Option { - use base64::Engine; - let payload_b64 = jwt.split('.').nth(1)?; - let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD - .decode(payload_b64) - .ok()?; - let claims: serde_json::Value = serde_json::from_slice(&payload).ok()?; - let tier = claims.get("tier")?.as_u64()?; - Some( - match tier { - 1 => "supergrok", - 2 => "x_basic", - 3 => "x_premium", - 4 => "x_premium_plus", - 5 => "supergrok_heavy", - 6 => "supergrok_lite", - 0 => "free", - _ => return Some(tier.to_string()), - } - .to_string(), - ) -} -/// Resolve Mixpanel / AuthMeta `subscription_tier`. -/// -/// Precedence: -/// 1. CCP `/settings` `subscription_tier_display` (when present and non-empty) -/// 2. [`AuthMode::ApiKey`] → `"api_key"` (never free) -/// 3. JWT `tier` claim via [`jwt_tier_claim`] (OAuth free → `"free"`) -pub(crate) fn resolve_subscription_tier_for_telemetry( - display: Option, - auth: Option<&crate::auth::GrokAuth>, -) -> Option { - if let Some(t) = display.filter(|s| !s.trim().is_empty()) { - return Some(t); - } - let auth = auth?; - if auth.auth_mode == crate::auth::AuthMode::ApiKey { - return Some("api_key".into()); - } - jwt_tier_claim(&auth.key) -} -/// Whether a JWT `tier` claim (from [`jwt_tier_claim`]) reflects the live -/// `/user?include=subscription` tier string (from the subscription API / QUALIFYING_TIERS). -/// -/// Post-unblock catalog refresh must not treat *any* present claim as enough: -/// an older paid claim (e.g. `x_basic`) can remain on the access token while -/// `/user` already reports a newly qualifying tier (e.g. `SuperGrokPro`). In -/// that case `/v1/models` would still be targeted at the stale level (the -/// "stale JWT tier skips retry" bug). -pub(crate) fn jwt_claim_matches_user_subscription_tier( - jwt_claim: &str, - user_subscription_tier: &str, -) -> bool { - match user_subscription_tier { - "GrokPro" => jwt_claim == "supergrok", - "XBasic" => jwt_claim == "x_basic", - "XPremium" => jwt_claim == "x_premium", - "XPremiumPlus" => jwt_claim == "x_premium_plus", - "SuperGrokPro" => jwt_claim == "supergrok_heavy", - "SuperGrokLite" => jwt_claim == "supergrok_lite", - _ => false, - } -} fn parse_session_computer_sessions(_meta: Option<&acp::Meta>) -> Option> { None } @@ -617,11 +548,6 @@ pub struct MvpAgent { /// into the detached prompt task; cleared for a workspace on GUI untrust /// (`execute_hooks_action`) so a later re-open can re-prompt. interactive_trust_prompted: Rc>>, - /// Whether the user's subscription tier is in the remote settings `allowed_tiers` - /// list. Set by `enforce_grok_code_access`; defaults to `true` (API-key and - /// external-auth users bypass the check). When `false`, the pager shows a - /// gate CTA instead of the prompt. - tier_allowed: std::cell::Cell, /// Storage mode - determines whether to sync to backend (writeback) or local only storage_mode: StorageMode, /// Default YOLO mode - when true, sessions start with auto-approve enabled. @@ -760,17 +686,6 @@ pub struct MvpAgent { /// on completion without re-borrowing `&self`. `Send` is required /// because the inner `sync_bundle_to_root` now uses `spawn_blocking`. bundle_sync_in_flight: Arc, - /// Single-flight guard for [`spawn_post_unblock_jwt_and_catalog_retry`]. - /// - /// After free→paid unblock the JWT may still lack a `tier` claim for - /// several seconds. Overlapping `CheckSubscription` RPCs (watch debounce, - /// paywall ticks, concurrent in-flight checks) would each otherwise spawn - /// another five-attempt `refresh_chain` backoff loop — multiplying IdP - /// traffic and redundant catalog work. - /// - /// Cleared by [`PostUnblockJwtRetryInFlightGuard`] on task exit (including - /// panic/abort), not only on the normal post-backoff path. - post_unblock_jwt_retry_in_flight: Arc, /// Local workspace ops, built lazily via [`Self::ensure_local_workspace_ops`]. /// The agent never opens Computer Hub as a harness/client; remote cloud /// sandboxes are gateway-owned (`gateway_bridge` / `computer_sessions`). @@ -1013,26 +928,14 @@ struct AuthRequestMeta { headless: bool, #[serde(default)] reauth: bool, - /// `--oauth`: force loopback. The only transport override sent over ACP - /// (loopback is the default; device is opt-in via env/config). - #[serde(default)] - use_oauth: bool, - /// When true, skip cached tokens and force the interactive browser login - /// flow. Used by the `/login` slash command for mid-session re-auth. - /// Unlike `reauth`, this does NOT clear existing credentials — if the - /// user abandons the browser flow, the current session continues. + /// When true, skip cached tokens and force the interactive login flow. + /// Used by the `/login` slash command for mid-session re-auth. Unlike + /// `reauth`, this does NOT clear existing credentials — if the user + /// abandons the device flow, the current session continues. #[serde(default)] force_interactive: bool, } impl AuthRequestMeta { - /// `--oauth` → force loopback; otherwise default (loopback). - fn login_override(&self) -> crate::auth::LoginTransportOverride { - if self.use_oauth { - crate::auth::LoginTransportOverride::ForceLoopback - } else { - crate::auth::LoginTransportOverride::None - } - } fn from_json(meta: Option<&acp::Meta>) -> Self { meta.cloned() .and_then(|value| { @@ -1654,239 +1557,21 @@ impl MvpAgent { } result } - /// Check whether the user has access via remote settings `allow_access`. - /// - /// Non-xAI auth (API keys, enterprise) always passes. For xAI OAuth2 - /// users, reads `allow_access` from remote settings. Defaults to - /// `false` (blocked) when remote settings are unavailable. - pub(super) async fn enforce_grok_code_access(&self, auth: &crate::auth::GrokAuth) { - if !auth.is_xai_auth() { - self.tier_allowed.set(true); - return; - } - let allow = settings_allow_access(self.cfg.borrow().remote_settings.as_ref()); - self.tier_allowed.set(allow); - if !allow { - tracing::info!( - "auth: user blocked by allow_access (remote settings grok_build_access_gate)" - ); - self.retry_subscription_check().await; - } - } - /// Single-shot subscription check called by the pager's "Check - /// subscription" button (`x.ai/auth/check_subscription`). The pager - /// calls this every 5s while the paywall is shown, acting as the poller. - /// - /// Queries `/user?include=subscription` for the live tier from the - /// subscription API. If a qualifying tier is found, does a best-effort - /// JWT refresh and settings re-fetch, lifts the gate, then — when the - /// access token's `tier` claim **matches** that live tier - /// ([`jwt_claim_matches_user_subscription_tier`]; bare `refresh_chain` - /// Ok or any older paid claim is not enough) — fire-and-forgets an - /// explicit model catalog refresh (`ModelsManager::on_auth_changed`) so - /// tier-targeted models appear without restart. - /// Catalog refresh is not awaited so gate lift / auth meta are not - /// blocked on `/v1/models`. Without a matching claim, defers to - /// `spawn_post_unblock_jwt_and_catalog_retry`. - pub(crate) async fn retry_subscription_check(&self) { - let (proxy_base_url, alpha_test_key) = { - let cfg = self.cfg.borrow(); - (cfg.endpoints.proxy_url(), cfg.endpoints.alpha_test_key.clone()) - }; - let user_id = self - .auth_manager - .current() - .map(|a| a.user_id.clone()) - .unwrap_or_default(); - let result = super::subscription_check::single_check( - self.auth_manager.clone(), - &proxy_base_url, - alpha_test_key.as_deref(), - &user_id, - ) - .await; - if let Some(unblocked) = result { - tracing::info!( - new_tier = % unblocked.new_tier, "subscription detected, lifting gate" - ); - kigi_log::unified_log::info( - "paywall_check_gate_lifting", - None, - Some( - serde_json::json!( - { "user_id" : user_id, "new_tier" : unblocked.new_tier, } - ), - ), - ); - if let Some(settings) = unblocked.settings { - { - let mut cfg = self.cfg.borrow_mut(); - cfg.remote_settings = Some(settings); - crate::agent::config::apply_remote_settings_side_effects( - cfg.remote_settings.as_ref(), - ); - } - } - if crate::util::config::resolve_remote_fetch_enabled() - && !settings_allow_access(self.cfg.borrow().remote_settings.as_ref()) - { - tracing::info!( - new_tier = % unblocked.new_tier, - "subscription detected but allow_access still false, keeping gate" - ); - kigi_log::unified_log::warn( - "paywall_check_gate_kept_allow_access_false", - None, - Some( - serde_json::json!( - { "user_id" : user_id, "new_tier" : unblocked.new_tier, } - ), - ), - ); - return; - } - self.tier_allowed.set(true); - let refresh_ok = match self - .auth_manager - .refresh_chain( - crate::auth::token_type::TokenType::OidcSession, - crate::auth::manager::RefreshReason::ServerRejected, - ) - .await - { - Ok(_) => { - tracing::info!("post-unblock: JWT refresh_chain succeeded"); - kigi_log::unified_log::info( - "paywall_check_jwt_refreshed", - None, - Some(serde_json::json!({ "user_id" : user_id })), - ); - true - } - Err(e) => { - tracing::warn!( - error = % e, - "post-unblock: JWT refresh failed, user may need to re-login on next restart" - ); - kigi_log::unified_log::warn( - "paywall_check_error", - None, - Some( - serde_json::json!( - { "user_id" : user_id, "kind" : - "post_unblock_refresh_failed", "detail" : e.to_string(), } - ), - ), - ); - false - } - }; - let jwt_claim = self - .auth_manager - .current_or_expired() - .and_then(|auth| jwt_tier_claim(&auth.key)); - let jwt_matches_new_tier = jwt_claim - .as_ref() - .is_some_and(|claim| jwt_claim_matches_user_subscription_tier( - claim, - &unblocked.new_tier, - )); - if jwt_matches_new_tier { - let models_manager = self.models_manager.clone(); - let user_id_log = user_id.clone(); - let new_tier = unblocked.new_tier.clone(); - let jwt_claim_log = jwt_claim.clone(); - tokio::task::spawn(async move { - kigi_log::unified_log::info( - "model catalog: post_subscription_unblock refresh", - None, - Some( - serde_json::json!( - { "user_id" : user_id_log, "new_tier" : new_tier, - "refresh_ok" : refresh_ok, "jwt_claim" : jwt_claim_log, - "jwt_matches_new_tier" : true, } - ), - ), - ); - models_manager.on_auth_changed().await; - }); - } else { - tracing::warn!( - refresh_ok, jwt_claim = ? jwt_claim, new_tier = % unblocked.new_tier, - "post-unblock: JWT tier claim missing or stale vs live tier; deferring model catalog refresh with retry" - ); - kigi_log::unified_log::warn( - "model catalog: post_subscription_unblock deferred (jwt tier missing or stale)", - None, - Some( - serde_json::json!( - { "user_id" : user_id, "new_tier" : unblocked.new_tier, - "refresh_ok" : refresh_ok, "jwt_claim" : jwt_claim, } - ), - ), - ); - spawn_post_unblock_jwt_and_catalog_retry( - self.auth_manager.clone(), - self.models_manager.clone(), - self.post_unblock_jwt_retry_in_flight.clone(), - user_id.clone(), - unblocked.new_tier.clone(), - ); - } - } else { - kigi_log::unified_log::info( - "paywall_check_no_subscription", - None, - Some(serde_json::json!({ "user_id" : user_id, })), - ); - } - } pub(crate) fn auth_response_with_meta(&self) -> AuthenticateResponse { - let (show_resolved_model, gate, subscription_tier) = { + let show_resolved_model = { let cfg = self.cfg.borrow(); - let rs = cfg.remote_settings.as_ref(); - let gate = rs - .and_then(|s| s.gate_message.as_ref()) - .filter(|m| !m.is_empty()) - .map(|message| crate::auth::GateInfo { - message: message.clone(), - url: rs.and_then(|s| s.gate_url.clone()), - label: rs.and_then(|s| s.gate_label.clone()), - }); - let subscription_tier = rs.and_then(|s| s.subscription_tier_display.clone()); - (rs.and_then(|s| s.show_resolved_model), gate, subscription_tier) + cfg.remote_settings + .as_ref() + .and_then(|s| s.show_resolved_model) }; - let subscription_tier = resolve_subscription_tier_for_telemetry( - subscription_tier, - self.auth_manager.current_or_expired().as_ref(), - ); let meta = self .auth_manager .current() .map(|auth| { - let gate = if !self.tier_allowed.get() && gate.is_none() { - let message = "A subscription is required.".to_string(); - Some(crate::auth::GateInfo { - message, - url: Some( - "https://grok.com/supergrok?referrer=grok-build".to_string(), - ), - label: Some("Subscribe".to_string()), - }) - } else { - gate - }; let auth_meta = crate::auth::AuthMeta { email: auth.email.clone(), auth_mode: Some(format!("{:?}", auth.auth_mode)), - team_id: auth.team_id.clone(), - team_name: auth.team_name.clone(), - is_zdr: auth.is_zdr_team(), - team_role: auth.team_role.clone(), - coding_data_retention_opt_out: auth.coding_data_retention_opt_out, show_resolved_model, - gate, - subscription_tier, }; serde_json::to_value(auth_meta) .ok() @@ -1904,7 +1589,7 @@ impl MvpAgent { let Some(auth) = self.auth_manager.current() else { return; }; - let is_xai_auth = auth.is_xai_auth(); + let is_session_auth = auth.is_session_auth(); let Some(settings) = self.fetch_remote_settings(auth).await else { return; }; @@ -1922,7 +1607,7 @@ impl MvpAgent { None, cfg.remote_settings.as_ref(), ); - if cfg.storage_mode == StorageMode::Writeback && !is_xai_auth { + if cfg.storage_mode == StorageMode::Writeback && !is_session_auth { cfg.storage_mode = StorageMode::Local; } } @@ -2107,160 +1792,6 @@ impl MvpAgent { }); } } -/// Clears [`MvpAgent::post_unblock_jwt_retry_in_flight`] on scope exit — -/// success, exhaustion, cancel/abort, or panic — so the single-flight flag -/// cannot wedge `true` for the rest of the process. -struct PostUnblockJwtRetryInFlightGuard { - flag: Arc, -} -impl Drop for PostUnblockJwtRetryInFlightGuard { - fn drop(&mut self) { - self.flag.store(false, std::sync::atomic::Ordering::Release); - } -} -/// Background retry when post-unblock JWT lacks a tier claim that matches -/// the live `/user` tier. Re-attempts `refresh_chain` and only treats an -/// attempt as success when [`jwt_claim_matches_user_subscription_tier`] -/// holds (bare refresh Ok, free token, or a *stale older* paid claim are -/// all misses). Then refreshes the model catalog. -/// -/// Gate lift already happened; this only recovers the tier-targeted catalog. -/// -/// Single-flight: concurrent unblocks (overlapping `CheckSubscription` -/// RPCs while the JWT is still free/stale-targeted) share one backoff loop -/// via `in_flight`. A second spawn while a loop is running is a no-op. -/// The flag is released by [`PostUnblockJwtRetryInFlightGuard`] (Drop), not -/// only on the happy path after `execute_with_backoff`. -fn spawn_post_unblock_jwt_and_catalog_retry( - auth_manager: std::sync::Arc, - models_manager: crate::agent::models::ModelsManager, - in_flight: Arc, - user_id: String, - new_tier: String, -) { - use std::sync::atomic::Ordering; - if in_flight - .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed) - .is_err() - { - tracing::debug!( - "post-unblock JWT/catalog retry already in flight, skipping duplicate spawn" - ); - kigi_log::unified_log::info( - "model catalog: post_subscription_unblock jwt retry skipped (already in flight)", - None, - Some(serde_json::json!({ "user_id" : user_id, "new_tier" : new_tier, })), - ); - return; - } - tokio::task::spawn(async move { - let _in_flight_guard = PostUnblockJwtRetryInFlightGuard { - flag: in_flight, - }; - let backoff = crate::tools::retry::BackoffConfig::new(5, 2_000, 30_000); - let result = crate::tools::retry::execute_with_backoff( - &backoff, - || { - let auth_manager = auth_manager.clone(); - let new_tier = new_tier.clone(); - async move { - let refresh_result = auth_manager - .refresh_chain( - crate::auth::token_type::TokenType::OidcSession, - crate::auth::manager::RefreshReason::ServerRejected, - ) - .await; - let jwt_claim = auth_manager - .current_or_expired() - .and_then(|auth| jwt_tier_claim(&auth.key)); - let matches = jwt_claim - .as_ref() - .is_some_and(|claim| jwt_claim_matches_user_subscription_tier( - claim, - &new_tier, - )); - if matches { - Ok(()) - } else { - let detail = match (&refresh_result, &jwt_claim) { - (Ok(_), None) => "refresh_ok but no tier claim".to_string(), - (Ok(_), Some(c)) => { - format!( - "refresh_ok but stale tier claim={c} (want {new_tier})" - ) - } - (Err(e), Some(c)) => { - format!( - "refresh_err={e}; stale tier claim={c} (want {new_tier})" - ) - } - (Err(e), None) => e.to_string(), - }; - Err(format!("jwt tier not current: {detail}")) - } - } - }, - |attempt, max_retries, delay| { - let user_id = user_id.clone(); - let new_tier = new_tier.clone(); - async move { - kigi_log::unified_log::warn( - "model catalog: post_subscription_unblock jwt retry scheduled", - None, - Some( - serde_json::json!( - { "user_id" : user_id, "new_tier" : new_tier, "attempt" : - attempt, "max_retries" : max_retries, "delay_ms" : delay - .as_millis() as u64, } - ), - ), - ); - } - }, - ) - .await; - match result { - Ok(()) => { - kigi_log::unified_log::info( - "model catalog: post_subscription_unblock refresh (after jwt retry)", - None, - Some( - serde_json::json!( - { "user_id" : user_id, "new_tier" : new_tier, } - ), - ), - ); - models_manager.on_auth_changed().await; - } - Err(e) => { - kigi_log::unified_log::warn( - "model catalog: post_subscription_unblock jwt retry exhausted", - None, - Some( - serde_json::json!( - { "user_id" : user_id, "new_tier" : new_tier, "error" : e - .to_string(), } - ), - ), - ); - } - } - }); -} -/// Resolve `allow_access` from remote settings. -/// -/// Returns `true` only when remote settings explicitly set `allow_access: true`. -/// Defaults to `false` (blocked) when settings are `None` or the field is -/// absent — matching the `grok_build_access_gate` flag's server-side default. -/// -/// Used by both `enforce_grok_code_access` (initial login gate) and -/// `retry_subscription_check` (poller gate lift) to keep the decision in -/// one place. -pub(crate) fn settings_allow_access( - rs: Option<&crate::util::config::RemoteSettings>, -) -> bool { - rs.and_then(|s| s.allow_access).unwrap_or(false) -} /// Parse `_meta.agentProfile` as a JSON object or string name. /// Returns `None` if absent or invalid. pub(crate) fn parse_agent_profile_from_meta( diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs index b82b8f6..0ac7e34 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs @@ -1,159 +1,4 @@ use super::*; -/// Build an unsigned JWT with a `tier` claim (header.payload.sig base64url). -fn jwt_with_tier(tier: u64) -> String { - use base64::Engine; - let enc = base64::engine::general_purpose::URL_SAFE_NO_PAD; - let header = enc.encode(br#"{"alg":"none"}"#); - let payload = enc.encode(format!(r#"{{"tier":{tier}}}"#).as_bytes()); - format!("{header}.{payload}.sig") -} -#[test] -fn jwt_tier_claim_maps_free_and_paid() { - assert_eq!(jwt_tier_claim(&jwt_with_tier(0)).as_deref(), Some("free")); - assert_eq!( - jwt_tier_claim(&jwt_with_tier(1)).as_deref(), - Some("supergrok") - ); - assert_eq!( - jwt_tier_claim(&jwt_with_tier(2)).as_deref(), - Some("x_basic") - ); - assert_eq!( - jwt_tier_claim(&jwt_with_tier(3)).as_deref(), - Some("x_premium") - ); - assert_eq!( - jwt_tier_claim(&jwt_with_tier(4)).as_deref(), - Some("x_premium_plus") - ); - assert_eq!( - jwt_tier_claim(&jwt_with_tier(5)).as_deref(), - Some("supergrok_heavy") - ); - assert_eq!( - jwt_tier_claim(&jwt_with_tier(6)).as_deref(), - Some("supergrok_lite") - ); - assert_eq!(jwt_tier_claim(&jwt_with_tier(99)).as_deref(), Some("99")); -} -fn auth_with_mode(mode: crate::auth::AuthMode, key: &str) -> crate::auth::GrokAuth { - crate::auth::GrokAuth { - key: key.into(), - auth_mode: mode, - create_time: chrono::Utc::now(), - user_id: "u".into(), - email: None, - first_name: None, - last_name: None, - profile_image_asset_id: None, - principal_type: None, - principal_id: None, - team_id: None, - team_name: None, - team_role: None, - organization_id: None, - organization_name: None, - organization_role: None, - user_blocked_reason: None, - team_blocked_reasons: vec![], - coding_data_retention_opt_out: false, - has_grok_code_access: None, - refresh_token: None, - expires_at: None, - oidc_issuer: None, - oidc_client_id: None, - } -} -#[test] -fn resolve_subscription_tier_prefers_display_then_api_key_then_jwt() { - assert_eq!( - resolve_subscription_tier_for_telemetry(Some("Free".into()), None).as_deref(), - Some("Free") - ); - let api = auth_with_mode(crate::auth::AuthMode::ApiKey, "xai-not-a-jwt"); - assert_eq!( - resolve_subscription_tier_for_telemetry(Some(" ".into()), Some(&api)).as_deref(), - Some("api_key") - ); - assert_eq!( - resolve_subscription_tier_for_telemetry(None, Some(&api)).as_deref(), - Some("api_key") - ); - let oauth = auth_with_mode(crate::auth::AuthMode::Oidc, &jwt_with_tier(0)); - assert_eq!( - resolve_subscription_tier_for_telemetry(None, Some(&oauth)).as_deref(), - Some("free") - ); - assert_ne!( - resolve_subscription_tier_for_telemetry(None, Some(&api)).as_deref(), - Some("free") - ); -} -/// JWT claim ↔ `/user` tier mapping used to gate post-unblock catalog refresh -/// (a stale older paid claim must not skip retry). -#[test] -fn jwt_claim_matches_user_subscription_tier_known_pairs() { - let cases = [ - ("supergrok", "GrokPro"), - ("x_basic", "XBasic"), - ("x_premium", "XPremium"), - ("x_premium_plus", "XPremiumPlus"), - ("supergrok_heavy", "SuperGrokPro"), - ("supergrok_lite", "SuperGrokLite"), - ]; - for (claim, user_tier) in cases { - assert!( - jwt_claim_matches_user_subscription_tier(claim, user_tier), - "{claim} should match {user_tier}" - ); - } -} -#[test] -fn jwt_claim_matches_user_subscription_tier_rejects_stale_and_unknown() { - assert!(!jwt_claim_matches_user_subscription_tier( - "x_basic", - "SuperGrokPro" - )); - assert!(!jwt_claim_matches_user_subscription_tier( - "supergrok", - "SuperGrokPro" - )); - assert!(!jwt_claim_matches_user_subscription_tier("free", "GrokPro")); - assert!(!jwt_claim_matches_user_subscription_tier("", "XPremium")); - assert!(!jwt_claim_matches_user_subscription_tier( - "supergrok_heavy", - "EnterpriseMystery" - )); -} -/// Single-flight flag must clear on Drop even if the retry task panics / -/// aborts mid-backoff (guards against the flag stuck true forever). -#[test] -fn post_unblock_jwt_retry_in_flight_guard_clears_on_drop() { - use std::sync::Arc; - use std::sync::atomic::{AtomicBool, Ordering}; - let flag = Arc::new(AtomicBool::new(true)); - { - let _guard = PostUnblockJwtRetryInFlightGuard { flag: flag.clone() }; - assert!(flag.load(Ordering::Acquire)); - } - assert!( - !flag.load(Ordering::Acquire), - "Drop must release post_unblock_jwt_retry_in_flight" - ); - let flag = Arc::new(AtomicBool::new(true)); - let flag_for_catch = flag.clone(); - let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - let _guard = PostUnblockJwtRetryInFlightGuard { - flag: flag_for_catch, - }; - panic!("simulate retry task panic"); - })); - assert!(result.is_err()); - assert!( - !flag.load(Ordering::Acquire), - "Drop must release flag on panic unwind" - ); -} mod hunk_tracking_mode { use super::super::{plan_hunk_tracking, resolve_hunk_tracking_mode}; use kigi_hunk_tracker::TrackingMode; @@ -354,42 +199,6 @@ fn trace_turn_to_i32_saturates_at_max() { let result = i32::try_from(boundary).unwrap_or(i32::MAX); assert_eq!(result, i32::MAX); } -/// When remote settings are absent (`None`), default to blocked. -#[test] -fn settings_allow_access_none_settings_is_blocked() { - assert!(!settings_allow_access(None)); -} -/// When `allow_access` is `Some(true)`, user is allowed. -#[test] -fn settings_allow_access_true_is_allowed() { - let rs = crate::util::config::RemoteSettings { - allow_access: Some(true), - ..Default::default() - }; - assert!(settings_allow_access(Some(&rs))); -} -/// When `allow_access` is `Some(false)` (remote settings default / rule -/// disabled), user stays blocked — even if they hold a qualifying -/// subscription. This is the regression guard for the bug where -/// `retry_subscription_check` unconditionally lifted the gate. -#[test] -fn settings_allow_access_false_is_blocked() { - let rs = crate::util::config::RemoteSettings { - allow_access: Some(false), - ..Default::default() - }; - assert!(!settings_allow_access(Some(&rs))); -} -/// When `/settings` returned successfully but the field is absent -/// (`None`), default to blocked (conservative). -#[test] -fn settings_allow_access_field_absent_is_blocked() { - let rs = crate::util::config::RemoteSettings { - allow_access: None, - ..Default::default() - }; - assert!(!settings_allow_access(Some(&rs))); -} /// After allocating a turn number, `session_turn_numbers` holds the next /// value (current + 1). This is the value that must be persisted via /// `SetNextTraceTurn` so the counter survives restarts. @@ -1333,10 +1142,10 @@ async fn ext_method_routes_auth_cleared_and_refreshes_resident_sessions() { let local = tokio::task::LocalSet::new(); local .run_until(async { - let agent = build_agent_with_auth(crate::auth::GrokAuth { + let agent = build_agent_with_auth(crate::auth::KimiAuth { key: "eligible".into(), - auth_mode: crate::auth::AuthMode::WebLogin, - ..crate::auth::GrokAuth::test_default() + auth_mode: crate::auth::AuthMode::OAuth, + ..crate::auth::KimiAuth::test_default() }); use acp::Agent as _; agent.managed_mcp_cache.lock().await.enable_gateway_tools(); @@ -1363,22 +1172,22 @@ async fn ext_method_routes_auth_cleared_and_refreshes_resident_sessions() { /// Build a minimal MvpAgent suitable for testing extension methods. fn build_minimal_agent_for_tests() -> MvpAgent { use crate::agent::config::Config as AgentConfig; - use crate::auth::{AuthManager, GrokComConfig}; + use crate::auth::{AuthManager, KimiCodeConfig}; let temp_dir = tempfile::tempdir().unwrap(); let auth_manager = - std::sync::Arc::new(AuthManager::new(temp_dir.path(), GrokComConfig::default())); + std::sync::Arc::new(AuthManager::new(temp_dir.path(), KimiCodeConfig::default())); let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); let gateway = GatewaySender::new(tx); let cfg = AgentConfig::default(); MvpAgent::new(gateway, &cfg, auth_manager, None).expect("valid test config") } /// Build a minimal MvpAgent with pre-loaded auth for gate tests. -fn build_agent_with_auth(auth: crate::auth::GrokAuth) -> MvpAgent { +fn build_agent_with_auth(auth: crate::auth::KimiAuth) -> MvpAgent { use crate::agent::config::Config as AgentConfig; - use crate::auth::{AuthManager, GrokComConfig}; + use crate::auth::{AuthManager, KimiCodeConfig}; let temp_dir = tempfile::tempdir().unwrap(); let auth_manager = - std::sync::Arc::new(AuthManager::new(temp_dir.path(), GrokComConfig::default())); + std::sync::Arc::new(AuthManager::new(temp_dir.path(), KimiCodeConfig::default())); auth_manager.hot_swap(auth); let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); let gateway = GatewaySender::new(tx); @@ -1395,7 +1204,7 @@ fn build_agent_with_auth(auth: crate::auth::GrokAuth) -> MvpAgent { #[serial_test::serial] async fn ensure_plugin_registry_lazily_populates_snapshot() { use crate::agent::config::Config as AgentConfig; - use crate::auth::{AuthManager, GrokComConfig}; + use crate::auth::{AuthManager, KimiCodeConfig}; use kigi_test_support::EnvGuard; let kigi_home = tempfile::tempdir().unwrap(); let _env = EnvGuard::set("KIGI_SHARE_DIR", kigi_home.path()); @@ -1412,7 +1221,7 @@ async fn ensure_plugin_registry_lazily_populates_snapshot() { .unwrap(); let auth_home = tempfile::tempdir().unwrap(); let auth_manager = - std::sync::Arc::new(AuthManager::new(auth_home.path(), GrokComConfig::default())); + std::sync::Arc::new(AuthManager::new(auth_home.path(), KimiCodeConfig::default())); let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); let gateway = GatewaySender::new(tx); let mut cfg = AgentConfig::default(); @@ -1599,10 +1408,10 @@ fn drain_roster_changed( async fn push_roster_activity_delta_broadcasts_overridden_activity() { use crate::agent::config::Config as AgentConfig; use crate::agent::roster::RosterActivity; - use crate::auth::{AuthManager, GrokComConfig}; + use crate::auth::{AuthManager, KimiCodeConfig}; let temp_dir = tempfile::tempdir().unwrap(); let auth_manager = - std::sync::Arc::new(AuthManager::new(temp_dir.path(), GrokComConfig::default())); + std::sync::Arc::new(AuthManager::new(temp_dir.path(), KimiCodeConfig::default())); let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); let gateway = GatewaySender::new(tx); let cfg = AgentConfig::default(); @@ -2043,7 +1852,6 @@ async fn auth_type_session_based_no_current_returns_session_token() { for method_id in [ crate::agent::auth_method::CACHED_TOKEN_AUTH_METHOD_ID, crate::agent::auth_method::KIGI_COM_METHOD_ID, - crate::agent::auth_method::OIDC_METHOD_ID, ] { let agent = build_minimal_agent_for_tests(); agent.set_auth_method(acp::AuthMethodId::new(method_id)); @@ -2084,12 +1892,12 @@ async fn auth_type_xai_api_key_no_current_returns_api_key() { /// common case during a healthy session. #[tokio::test(flavor = "current_thread")] async fn auth_type_session_based_with_current_returns_session_token() { - use crate::auth::GrokAuth; + use crate::auth::KimiAuth; let agent = build_minimal_agent_for_tests(); agent.set_auth_method(acp::AuthMethodId::new( - crate::agent::auth_method::OIDC_METHOD_ID, + crate::agent::auth_method::KIGI_COM_METHOD_ID, )); - agent.auth_manager.hot_swap(GrokAuth::test_default()); + agent.auth_manager.hot_swap(KimiAuth::test_default()); assert!(agent.auth_manager.current().is_some()); assert_eq!(agent.auth_type(), kigi_chat_state::AuthType::SessionToken,); } @@ -2112,27 +1920,13 @@ async fn auth_type_no_method_id_no_current_returns_api_key() { /// here matches pre-fix behavior and keeps logging stable. #[tokio::test(flavor = "current_thread")] async fn auth_type_no_method_id_with_current_returns_session_token() { - use crate::auth::GrokAuth; + use crate::auth::KimiAuth; let agent = build_minimal_agent_for_tests(); - agent.auth_manager.hot_swap(GrokAuth::test_default()); + agent.auth_manager.hot_swap(KimiAuth::test_default()); assert!(agent.auth_method_id.load().is_none()); assert!(agent.auth_manager.current().is_some()); assert_eq!(agent.auth_type(), kigi_chat_state::AuthType::SessionToken,); } -/// Minimal agent whose `grok_com_config` engages the api-key kill switch -/// (`disable_api_key_auth = true`), mirroring a forced-IdP deployment. -fn build_agent_with_api_key_auth_disabled() -> MvpAgent { - use crate::agent::config::Config as AgentConfig; - use crate::auth::{AuthManager, GrokComConfig}; - let temp_dir = tempfile::tempdir().unwrap(); - let auth_manager = - std::sync::Arc::new(AuthManager::new(temp_dir.path(), GrokComConfig::default())); - let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); - let gateway = GatewaySender::new(tx); - let mut cfg = AgentConfig::default(); - cfg.grok_com_config.disable_api_key_auth = Some(true); - MvpAgent::new(gateway, &cfg, auth_manager, None).expect("valid test config") -} /// Deployment-key / managed-config user: `XAI_API_KEY` resolves and the kill /// switch is off, so a dead `cached_token` MUST fall through to `xai.api_key` /// (no browser). This is the exact regression the fallthrough fixes. @@ -2145,36 +1939,12 @@ async fn cached_token_fallthrough_prefers_api_key_for_deployment_key() { let _key = EnvGuard::set(XAI_API_KEY_ENV_VAR, "test-deployment-key"); let agent = build_minimal_agent_for_tests(); assert_eq!( - agent - .cached_token_fallthrough_method_id() - .as_ref() - .map(|id| id.0.as_ref()), - Some(XAI_API_KEY_METHOD_ID), + agent.cached_token_fallthrough_method_id().0.as_ref(), + XAI_API_KEY_METHOD_ID, "deployment-key user (XAI_API_KEY set, no kill switch) must fall \ through to xai.api_key on a dead cached_token -- not interactive login", ); } -/// Forced-IdP deployment: even with `XAI_API_KEY` present, the admin kill -/// switch keeps the fallthrough on interactive `grok.com` (api-key auth is -/// neither advertised nor an eligible fallthrough). -#[tokio::test(flavor = "current_thread")] -#[serial_test::serial] -async fn cached_token_fallthrough_respects_kill_switch() { - use crate::agent::auth_method::{KIGI_COM_METHOD_ID, XAI_API_KEY_ENV_VAR}; - use kigi_test_support::EnvGuard; - let _lockdown = EnvGuard::unset("KIGI_DISABLE_API_KEY_AUTH"); - let _key = EnvGuard::set(XAI_API_KEY_ENV_VAR, "test-deployment-key"); - let agent = build_agent_with_api_key_auth_disabled(); - assert_eq!( - agent - .cached_token_fallthrough_method_id() - .as_ref() - .map(|id| id.0.as_ref()), - Some(KIGI_COM_METHOD_ID), - "disable_api_key_auth must keep the cached_token fallthrough on \ - interactive grok.com so XAI_API_KEY can't bypass forced IdP login", - ); -} /// No advertiseable credentials at all (no env key, no kill switch): the user /// genuinely needs to log in, so the fallthrough is interactive `grok.com`. #[tokio::test(flavor = "current_thread")] @@ -2189,75 +1959,11 @@ async fn cached_token_fallthrough_falls_to_grok_com_without_credentials() { let _legacy = EnvGuard::unset(LEGACY_XAI_API_KEY_ENV_VAR); let agent = build_minimal_agent_for_tests(); assert_eq!( - agent - .cached_token_fallthrough_method_id() - .as_ref() - .map(|id| id.0.as_ref()), - Some(KIGI_COM_METHOD_ID), + agent.cached_token_fallthrough_method_id().0.as_ref(), + KIGI_COM_METHOD_ID, "no API-key creds and no kill switch -> interactive grok.com login", ); } -/// Verifies the 4-state matrix of `(disable_zdr_incompatible_tools, zdr_video_output_s3)`: -/// -/// | ZDR flag | S3 config | Result | -/// |----------|-----------|---------------------------------------------| -/// | false | None | Enabled, no S3 (normal non-ZDR mode) | -/// | true | None | Disabled (ZDR with no escape hatch) | -/// | false | Some | Enabled, S3 **not** threaded (non-ZDR) | -/// | true | Some | Enabled, S3 threaded (ZDR with upload path) | -#[tokio::test(flavor = "current_thread")] -async fn prepare_video_gen_config_disabled_when_zdr_flag_set() { - use kigi_tools::implementations::grok_build::video_gen::{ - S3AccessCredentials, VideoGenConfig, ZdrVideoOutputS3Config, - }; - fn zdr_s3() -> ZdrVideoOutputS3Config { - ZdrVideoOutputS3Config { - bucket: "team-videos".into(), - endpoint: "https://s3.example.com".into(), - region: "us-east-1".into(), - key_prefix: "grok-videos/".into(), - expires_secs: 900, - read_write: S3AccessCredentials { - access_key_id: "AKIA...".into(), - secret_access_key: "secret".into(), - }, - read_only: None, - } - } - let agent = build_minimal_agent_for_tests(); - agent.sampling_config.borrow_mut().api_key = Some("test-key".to_string()); - assert!(matches!( - agent.prepare_video_gen_config(), - VideoGenConfig::Enabled { .. } - )); - agent.cfg.borrow_mut().disable_zdr_incompatible_tools = true; - assert!(matches!( - agent.prepare_video_gen_config(), - VideoGenConfig::Disabled - )); - agent.cfg.borrow_mut().zdr_video_output_s3 = Some(zdr_s3()); - agent.cfg.borrow_mut().disable_zdr_incompatible_tools = false; - let VideoGenConfig::Enabled { - zdr_video_output_s3: s3_when_non_zdr, - .. - } = agent.prepare_video_gen_config() - else { - panic!("expected Enabled"); - }; - assert!( - s3_when_non_zdr.is_none(), - "S3 config must not be threaded when ZDR flag is off" - ); - agent.cfg.borrow_mut().disable_zdr_incompatible_tools = true; - let VideoGenConfig::Enabled { - zdr_video_output_s3, - .. - } = agent.prepare_video_gen_config() - else { - panic!("expected Enabled"); - }; - assert!(zdr_video_output_s3.as_ref().is_some_and(|c| c.is_valid())); -} /// The imagine tier gate fails **open**: with no resolved auth we can't confirm /// a restricted personal tier, so the tools stay advertised and un-flagged (the /// server 429 remains the authoritative backstop). Guards against accidentally @@ -2278,73 +1984,6 @@ async fn prepare_image_gen_config_fails_open_without_auth() { "no resolved auth ⇒ fail open (tools not tier-restricted)" ); } -#[tokio::test] -async fn data_collection_enabled_for_normal_user() { - let agent = build_agent_with_auth(crate::auth::GrokAuth::test_default()); - assert!( - !agent.is_data_collection_disabled(), - "normal user must have data collection enabled" - ); -} -#[tokio::test] -async fn data_collection_disabled_for_zdr_team() { - let agent = build_agent_with_auth(crate::auth::GrokAuth { - team_blocked_reasons: vec!["BLOCKED_REASON_NO_LOGS".into()], - ..crate::auth::GrokAuth::test_default() - }); - assert!( - agent.is_data_collection_disabled(), - "ZDR team must have data collection disabled" - ); -} -#[tokio::test] -async fn data_collection_disabled_for_zdr_moderated_team() { - let agent = build_agent_with_auth(crate::auth::GrokAuth { - team_blocked_reasons: vec!["BLOCKED_REASON_NO_LOGS_MODERATED".into()], - ..crate::auth::GrokAuth::test_default() - }); - assert!( - agent.is_data_collection_disabled(), - "ZDR-moderated team must have data collection disabled" - ); -} -#[tokio::test] -async fn data_collection_disabled_for_opted_out_team() { - let agent = build_agent_with_auth(crate::auth::GrokAuth { - coding_data_retention_opt_out: true, - ..crate::auth::GrokAuth::test_default() - }); - assert!( - agent.is_data_collection_disabled(), - "opted-out team must have data collection disabled" - ); -} -#[tokio::test] -async fn data_collection_disabled_for_zdr_plus_opt_out() { - let agent = build_agent_with_auth(crate::auth::GrokAuth { - team_blocked_reasons: vec!["BLOCKED_REASON_NO_LOGS".into()], - coding_data_retention_opt_out: true, - ..crate::auth::GrokAuth::test_default() - }); - assert!( - agent.is_data_collection_disabled(), - "ZDR + opt-out must have data collection disabled" - ); -} -#[tokio::test] -async fn data_collection_enabled_for_non_zdr_team_with_unrelated_blocks() { - let agent = build_agent_with_auth(crate::auth::GrokAuth { - team_blocked_reasons: vec![ - "BLOCKED_REASON_BILLING".into(), - "BLOCKED_REASON_SUSPENDED".into(), - ], - ..crate::auth::GrokAuth::test_default() - }); - assert!( - !agent.is_data_collection_disabled(), - "non-ZDR blocked reasons must not disable data collection" - ); -} /// `parse_session_kind` routes `session/load` to the gateway Chat path vs. the /// disk-backed Build path. Anything but an explicit `kind: "chat"` is Build. #[test] @@ -3098,10 +2737,10 @@ fn build_agent_with_gateway_rx() -> ( tokio::sync::mpsc::UnboundedReceiver, ) { use crate::agent::config::Config as AgentConfig; - use crate::auth::{AuthManager, GrokComConfig}; + use crate::auth::{AuthManager, KimiCodeConfig}; let temp_dir = tempfile::tempdir().unwrap(); let auth_manager = - std::sync::Arc::new(AuthManager::new(temp_dir.path(), GrokComConfig::default())); + std::sync::Arc::new(AuthManager::new(temp_dir.path(), KimiCodeConfig::default())); let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); let gateway = GatewaySender::new(tx); let cfg = AgentConfig::default(); @@ -3648,14 +3287,14 @@ mod soft_default_settings_emit { #[tokio::test] async fn emit_settings_update_carries_permission_mode_from_cfg() { use crate::agent::config::Config as AgentConfig; - use crate::auth::{AuthManager, GrokComConfig}; + use crate::auth::{AuthManager, KimiCodeConfig}; let local = tokio::task::LocalSet::new(); local .run_until(async { let temp_dir = tempfile::tempdir().unwrap(); let auth_manager = std::sync::Arc::new(AuthManager::new( temp_dir.path(), - GrokComConfig::default(), + KimiCodeConfig::default(), )); let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); let gateway = GatewaySender::new(tx); diff --git a/crates/codegen/kigi-shell/src/agent/session_registry_client.rs b/crates/codegen/kigi-shell/src/agent/session_registry_client.rs index a850587..e626911 100644 --- a/crates/codegen/kigi-shell/src/agent/session_registry_client.rs +++ b/crates/codegen/kigi-shell/src/agent/session_registry_client.rs @@ -141,7 +141,7 @@ pub struct SessionRegistryClient { raw_client: reqwest::Client, client: reqwest_middleware::ClientWithMiddleware, base_url: String, - credentials: crate::util::grok_auth_credentials::GrokAuthCredentials, + credentials: crate::util::kigi_auth_credentials::KigiAuthCredentials, session_id: Option, } @@ -152,7 +152,7 @@ impl SessionRegistryClient { raw_client: http_client.clone(), client: reqwest_middleware::ClientBuilder::new(http_client).build(), base_url: base_url.into(), - credentials: crate::util::grok_auth_credentials::GrokAuthCredentials::new(Some( + credentials: crate::util::kigi_auth_credentials::KigiAuthCredentials::new(Some( user_token.into(), )), session_id: None, @@ -549,7 +549,7 @@ mod tests { /// Verify per-request auth resolve picks up rotated tokens. #[tokio::test] async fn session_registry_client_uses_active_auth_for_each_request() { - use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig}; + use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig}; use axum::{Router, response::IntoResponse, routing::post}; use chrono::{Duration, Utc}; use std::net::SocketAddr; @@ -575,14 +575,14 @@ mod tests { tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); let dir = tempfile::tempdir().unwrap(); - let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - am.hot_swap(GrokAuth { + let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + am.hot_swap(KimiAuth { key: "fresh-from-auth-manager".into(), auth_mode: AuthMode::ApiKey, create_time: Utc::now(), user_id: "user-42".into(), expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }); let client = SessionRegistryClient::new(format!("http://{addr}"), "STALE-build-time-token") diff --git a/crates/codegen/kigi-shell/src/agent/subagent/mod.rs b/crates/codegen/kigi-shell/src/agent/subagent/mod.rs index 33feb7d..754913b 100644 --- a/crates/codegen/kigi-shell/src/agent/subagent/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/subagent/mod.rs @@ -156,7 +156,7 @@ pub(crate) struct SubagentSpawnContext { reason = "unused in production; remove expect when wired or delete the item" )] pub storage_mode: crate::config::StorageMode, - pub auth: Option, + pub auth: Option, pub parent_cwd: PathBuf, pub parent_session_id: String, pub yolo_mode: bool, diff --git a/crates/codegen/kigi-shell/src/agent/subscription_check.rs b/crates/codegen/kigi-shell/src/agent/subscription_check.rs deleted file mode 100644 index ddfffd1..0000000 --- a/crates/codegen/kigi-shell/src/agent/subscription_check.rs +++ /dev/null @@ -1,191 +0,0 @@ -//! Subscription check for paywall gate lift. -//! -//! Provides `single_check()` which queries `GET /user?include=subscription` -//! for the live subscription tier from the backend, independent of the JWT. -//! If a qualifying tier is detected, does a best-effort JWT refresh and -//! settings re-fetch, then returns an `UnblockResult` so the agent can -//! lift the gate. -//! -//! The pager drives the polling via `x.ai/auth/check_subscription`: the 5s -//! paywall chain, the free-tier watch, the refocus check, and -//! verify-before-paywall gate deferral (see the pager's `app::subscription` -//! module). -use crate::auth::AuthManager; -use crate::auth::UserInfo; -use crate::auth::manager::RefreshReason; -use crate::auth::token_type::TokenType; -use std::sync::Arc; -use std::time::Duration; -/// Subscription tiers that qualify for Grok Build access. -/// Any active subscription qualifies -- the access gate in remote settings -/// controls which tiers are actually allowed. -const QUALIFYING_TIERS: &[&str] = &[ - "SuperGrokPro", - "GrokPro", - "SuperGrokLite", - "XPremiumPlus", - "XPremium", - "XBasic", -]; -/// Successful subscription check result: confirmed qualifying tier + -/// optionally refreshed settings. -pub(crate) struct UnblockResult { - pub(crate) new_tier: String, - pub(crate) settings: Option, -} -/// Fetch `/user?include=subscription` and return the parsed `UserInfo`. -async fn fetch_user_info( - http_client: &reqwest::Client, - url: &str, - auth: &crate::auth::GrokAuth, - auth_manager: &AuthManager, - alpha_test_key: Option<&str>, -) -> Result { - let request = http_client - .get(url) - .timeout(Duration::from_secs(10)) - .header("Authorization", format!("Bearer {}", auth.key)) - .header( - "X-XAI-Token-Auth", - auth_manager.grok_com_config().token_header.as_str(), - ) - .header("x-grok-client-version", kigi_version::VERSION) - .header( - crate::http::CLIENT_MODE_HEADER, - crate::http::process_client_mode(), - ); - let _ = alpha_test_key; - match request.send().await { - Ok(resp) if resp.status().is_success() => { - resp.json::().await.map_err(|_| "parse") - } - Ok(_resp) => Err("http_status"), - Err(e) if e.is_timeout() => Err("timeout"), - Err(_) => Err("transport"), - } -} -/// Single-shot subscription check. Called by the pager every 5s while -/// the paywall is shown (`x.ai/auth/check_subscription`). -/// -/// Queries `/user?include=subscription` for the live tier. If a qualifying -/// tier is found, does a best-effort JWT refresh + settings re-fetch and -/// returns `Some(UnblockResult)`. Returns `None` if no qualifying -/// subscription exists or the request fails. -#[tracing::instrument(name = "paywall_check", skip_all, fields(user_id = %user_id))] -pub(crate) async fn single_check( - auth_manager: Arc, - proxy_base_url: &str, - alpha_test_key: Option<&str>, - user_id: &str, -) -> Option { - let user_url = format!("{}/user?include=subscription", proxy_base_url); - let http_client = crate::http::shared_client(); - let auth = auth_manager.current()?; - let user_info = match fetch_user_info( - &http_client, - &user_url, - &auth, - &auth_manager, - alpha_test_key, - ) - .await - { - Ok(ui) => ui, - Err(kind) => { - kigi_log::unified_log::warn( - "paywall_check_error", - None, - Some(serde_json::json!({ "user_id" : user_id, "kind" : kind })), - ); - return None; - } - }; - kigi_log::unified_log::info( - "paywall_check_result", - None, - Some(serde_json::json!( - { "user_id" : user_id, "subscription_tier" : user_info.subscription_tier, - } - )), - ); - let new_tier = match &user_info.subscription_tier { - Some(tier) if !tier.is_empty() => tier.clone(), - _ => return None, - }; - if !QUALIFYING_TIERS.contains(&new_tier.as_str()) { - return None; - } - kigi_log::unified_log::info( - "paywall_check_subscription_detected", - None, - Some(serde_json::json!({ "user_id" : user_id, "new_tier" : new_tier, })), - ); - if let Err(e) = auth_manager - .refresh_chain(TokenType::OidcSession, RefreshReason::ServerRejected) - .await - { - kigi_log::unified_log::warn( - "paywall_check_error", - None, - Some(serde_json::json!( - { "user_id" : user_id, "kind" : "refresh_failed", "detail" : e - .to_string(), } - )), - ); - } - let settings = if crate::util::config::resolve_remote_fetch_enabled() { - let base_url = proxy_base_url.to_string(); - let auth_for_settings = auth_manager.current().unwrap_or(auth); - let atk = alpha_test_key.map(str::to_string); - tokio::task::spawn_blocking(move || { - crate::remote::fetch_settings_blocking(&base_url, &auth_for_settings, atk.as_deref()) - }) - .await - .ok() - .flatten() - } else { - None - }; - kigi_log::unified_log::info( - "paywall_check_unblocked", - None, - Some(serde_json::json!({ "user_id" : user_id, "new_tier" : new_tier })), - ); - Some(UnblockResult { new_tier, settings }) -} -#[cfg(test)] -mod tests { - use super::*; - #[test] - fn qualifying_tiers_includes_all_paid_tiers() { - for tier in &[ - "SuperGrokPro", - "GrokPro", - "SuperGrokLite", - "XPremiumPlus", - "XPremium", - "XBasic", - ] { - assert!( - QUALIFYING_TIERS.contains(tier), - "{tier} must be in QUALIFYING_TIERS" - ); - } - } - #[test] - fn free_tier_is_not_qualifying() { - assert!(!QUALIFYING_TIERS.contains(&"Free")); - } - #[test] - fn empty_tier_is_not_qualifying() { - assert!(!QUALIFYING_TIERS.contains(&"")); - } - /// The subscription check only returns `Some` when `/user` reports a - /// qualifying tier. Verify the tier matching is exact (no prefix match). - #[test] - fn partial_tier_name_is_not_qualifying() { - assert!(!QUALIFYING_TIERS.contains(&"Super")); - assert!(!QUALIFYING_TIERS.contains(&"Grok")); - assert!(!QUALIFYING_TIERS.contains(&"XPremium+")); - } -} diff --git a/crates/codegen/kigi-shell/src/auth/attribution.rs b/crates/codegen/kigi-shell/src/auth/attribution.rs index ade8350..ea5b2d1 100644 --- a/crates/codegen/kigi-shell/src/auth/attribution.rs +++ b/crates/codegen/kigi-shell/src/auth/attribution.rs @@ -370,7 +370,7 @@ pub(crate) fn record_auth_401( /// This function performs **exactly one** read-side acquisition of /// [`AuthManager`]'s internal `RwLock` -- it calls /// [`AuthManager::current`] once and derives both `current_key_prefix` -/// and the mint/expiry fields from the resulting `GrokAuth`. +/// and the mint/expiry fields from the resulting `KimiAuth`. /// /// `is_stale_snapshot` is `true` only when the live `current()` token /// differs from the bearer the client sent. When `current()` returns @@ -390,7 +390,7 @@ fn compute_attribution_payload( // query can break down on this). let sent_prefix = sent_bearer.map(token_suffix).unwrap_or(""); - // Single read-lock acquisition: pull the live `GrokAuth` (or + // Single read-lock acquisition: pull the live `KimiAuth` (or // `None`) once and derive every other field from it. let current_auth = auth_manager.current(); let current_prefix_owned: Option = current_auth @@ -411,7 +411,7 @@ fn compute_attribution_payload( // // TODO: mirror the full External-with-ttl branch from // `AuthManager::is_token_expired` (uses - // `grok_com_config.auth_token_ttl` when `expires_at` is `None` + // `kimi_code_config.auth_token_ttl` when `expires_at` is `None` // and `auth_mode == External`). The current 2-branch fallback // (`expires_at` if Some else `create_time + TOKEN_TTL`) is good // enough for diagnostic metadata; the External-ttl branch is @@ -441,7 +441,7 @@ mod tests { use chrono::{Duration, Utc}; - use crate::auth::{AuthManager, GrokAuth, GrokComConfig}; + use crate::auth::{AuthManager, KimiAuth, KimiCodeConfig}; use super::*; @@ -449,17 +449,17 @@ mod tests { /// nothing from a developer's actual `~/.kigi/auth.json` leaks in. fn empty_auth_manager() -> (tempfile::TempDir, AuthManager) { let dir = tempfile::tempdir().expect("tempdir"); - let cfg = GrokComConfig::default(); + let cfg = KimiCodeConfig::default(); let am = AuthManager::new(dir.path(), cfg); (dir, am) } - fn fresh_auth(key: &str) -> GrokAuth { - GrokAuth { + fn fresh_auth(key: &str) -> KimiAuth { + KimiAuth { key: key.to_string(), create_time: Utc::now(), expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() } } @@ -553,12 +553,12 @@ mod tests { #[test] fn legacy_token_uses_two_branch_fallback() { let (_dir, am) = empty_auth_manager(); - let auth = GrokAuth { + let auth = KimiAuth { key: "k".into(), create_time: Utc::now() - Duration::seconds(60), // No expires_at => falls through to create_time + TOKEN_TTL // (= 30 days). - ..GrokAuth::test_default() + ..KimiAuth::test_default() }; am.hot_swap(auth); diff --git a/crates/codegen/kigi-shell/src/auth/config.rs b/crates/codegen/kigi-shell/src/auth/config.rs index c93cc96..e06a8ca 100644 --- a/crates/codegen/kigi-shell/src/auth/config.rs +++ b/crates/codegen/kigi-shell/src/auth/config.rs @@ -1,423 +1,43 @@ -use super::model::TEAM_PRINCIPAL_TYPE; +//! Kimi Code auth configuration. +//! +//! The wire endpoints come from [`kigi_env`] (`oauth_host()`, overridable via +//! `KIGI_OAUTH_HOST`) and the client id is fixed +//! ([`crate::auth::kimi_oauth::KIMI_CODE_CLIENT_ID`]), so this config carries +//! no per-deployment OAuth knobs. The struct is kept (deserialized from the +//! agent config TOML) as the extension point for future auth options. + use serde::{Deserialize, Serialize}; -// Transitional: the M1 auth rewrite (Kimi device flow) replaces this origin. -const AUTH_ORIGIN_DEFAULT: &str = "https://grok.com"; -fn default_oidc_scopes() -> Vec { - vec![ - "openid".into(), - "profile".into(), - "email".into(), - "offline_access".into(), - "api:access".into(), - ] -} -/// Default scopes for the xAI OAuth2 provider. Includes `grok-cli:access` -/// which authorizes the token for API proxy requests. -fn default_oauth2_scopes() -> Vec { - vec![ - "openid".into(), - "profile".into(), - "email".into(), - "offline_access".into(), - "grok-cli:access".into(), - "api:access".into(), - "conversations:read".into(), - "conversations:write".into(), - ] -} -fn default_team_oauth2_scopes() -> Vec { - vec![ - "profile".into(), - "offline_access".into(), - "grok-cli:access".into(), - "api:access".into(), - "team:read".into(), - "conversations:read".into(), - "conversations:write".into(), - ] -} -/// Pin automatic auth to one method (`[auth] preferred_method` in config.toml). -/// -/// When set, only that method is used for automatic selection; if it is -/// unavailable, auth fails (no silent fallthrough to the other method). -/// Unset keeps today's multi-method fallthrough (session preferred when both -/// exist). Config-toml only — not remote settings, settings UI, or env. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PreferredAuthMethod { - /// `XAI_API_KEY` / auth.json `xai::api_key` / per-model BYOK (`xai.api_key`). - ApiKey, - /// OIDC / OAuth2 session (`cached_token`, interactive `grok.com` / `oidc`, - /// including devbox-minted OIDC). - Oidc, -} -#[derive(Debug, Clone, Serialize, Deserialize)] + +/// Persisted-credential scope key for the Kimi Code OAuth session — both the +/// auth.json map key and the system-keyring entry name (service `kigi`). +pub const KIMI_CODE_OAUTH_SCOPE: &str = "oauth/kimi-code"; + +/// Auth configuration block (`[kimi_code_config]` in the agent config). +/// Currently empty: the OAuth host and client id are fixed by the +/// environment crate. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] #[serde(default)] -pub struct GrokComConfig { - /// Auth origin / login-host display (doubles as the legacy WS origin name). - pub grok_ws_origin: String, - pub token_header: String, - /// OIDC config for customer-provided IdPs. See [`OidcAuthConfig`]. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub oidc: Option, - /// OAuth2 provider config. When set, preferred over the legacy relay flow. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub oauth2: Option, - /// External auth provider command (stdout = token, stderr = user UX, exit 0 = success). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub auth_provider_command: Option, - /// Login button label (env: `KIGI_AUTH_PROVIDER_LABEL`). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub auth_provider_label: Option, - /// Token TTL in seconds for external auth providers that output bare - /// tokens without `expires_in`. Synthesizes `expires_at` so proactive - /// refresh works. Env: `KIGI_AUTH_TOKEN_TTL`. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub auth_token_ttl: Option, - /// Admin kill switch: when `Some(true)`, the `xai.api_key` auth method is - /// neither advertised nor accepted, so `XAI_API_KEY`/per-model credentials - /// can't bypass the deployment's IdP login. Env: `KIGI_DISABLE_API_KEY_AUTH`. - /// Parity with common force-login-method admin knobs. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub disable_api_key_auth: Option, - /// Restrict login to a specific team — the login token's team principal must - /// equal this. Put in `requirements.toml` to enforce as non-overridable policy. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub force_login_team_uuid: Option, - /// Pin automatic auth to `api_key` or `oidc`. When set and the chosen - /// method is unavailable, auth fails (no fallthrough). Unset keeps - /// multi-method fallthrough. Config.toml only (`[auth] preferred_method`). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub preferred_method: Option, -} -/// Team login restriction. TOML string or array; an empty array fails closed. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(untagged)] -pub enum ForceLoginTeam { - /// The only allowed team. - Single(String), - /// Allowed teams; empty = fail closed. - AnyOf(Vec), -} -/// Customer OIDC Identity Provider configuration (`[grok_com_config.oidc]`). -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct OidcAuthConfig { - pub issuer: String, - pub client_id: String, - #[serde(default = "default_oidc_scopes")] - pub scopes: Vec, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub audience: Option, -} -/// OAuth2 provider configuration (`KIGI_OAUTH2_ISSUER` / `KIGI_OAUTH2_CLIENT_ID`). -/// -/// Uses the standard OAuth 2.1 Auth Code + PKCE flow via [`OidcAuthConfig`]. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct OAuth2ProviderConfig { - pub issuer: String, - pub client_id: String, - #[serde(default = "default_oauth2_scopes")] - pub scopes: Vec, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub principal_type: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub principal_id: Option, - /// Client-supplied referrer for OAuth usage-attribution analytics. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub referrer: Option, -} -pub const XAI_OAUTH2_ISSUER: &str = "https://auth.x.ai"; -/// Production accounts-app origin allowlist — the only origins builds without -/// non-production builds accept. Lives in its own const, referenced by both -/// profiles below, so the frozen-contract test (monorepo CI compiles with -/// that feature enabled) still pins this production-origin const. -const PROD_ACCOUNTS_APP_ORIGINS: &[&str] = &["https://accounts.x.ai"]; -/// See the opt-in non-production feature variant above — builds without -/// the feature accept only the production accounts app. -pub fn allowed_accounts_app_origins() -> Vec { - PROD_ACCOUNTS_APP_ORIGINS - .iter() - .map(|o| o.to_string()) - .collect() -} -/// Build a CORS layer that accepts requests from the accounts-app deployments -/// listed in [`allowed_accounts_app_origins`] for the given HTTP method. -/// -/// Callers can chain additional configuration (e.g. `.allow_headers(...)` or -/// `.allow_private_network(true)`) onto the returned layer. -pub fn accounts_app_cors_layer(method: axum::http::Method) -> tower_http::cors::CorsLayer { - tower_http::cors::CorsLayer::new() - .allow_origin(tower_http::cors::AllowOrigin::list( - allowed_accounts_app_origins() - .iter() - .filter_map(|origin| match origin.parse() { - Ok(value) => Some(value), - Err(_) => { - tracing::warn!(origin, "skipping malformed accounts-app CORS origin"); - None - } - }), - )) - .allow_methods([method]) -} -/// Local-dev OAuth2 issuer (accounts-app running on localhost). -const XAI_OAUTH2_LOCAL_ISSUER: &str = "http://localhost:22255"; -const DEFAULT_OAUTH2_REFERRER: &str = "grok-build"; -/// Returns `true` when `KIGI_LOCAL_AUTH=1` is set, -/// indicating the local accounts-app should be used as the OAuth2 issuer. -pub fn use_local_auth() -> bool { - std::env::var("KIGI_LOCAL_AUTH") - .map(|v| !v.is_empty() && v != "0") - .unwrap_or(false) -} -/// Returns the active xAI OAuth2 issuer — the local-dev issuer when -/// `KIGI_LOCAL_AUTH=1` is set, otherwise the production issuer. -pub fn xai_oauth2_issuer() -> &'static str { - if use_local_auth() { - XAI_OAUTH2_LOCAL_ISSUER - } else { - XAI_OAUTH2_ISSUER - } -} -/// Returns `true` if `issuer` is a recognised xAI OAuth2 issuer -/// (production **or** local-dev). Use this instead of comparing against -/// [`XAI_OAUTH2_ISSUER`] directly so that local-dev sessions are still -/// treated as first-party xAI auth. -pub fn is_xai_oauth2_issuer(issuer: &str) -> bool { - issuer == XAI_OAUTH2_ISSUER || issuer == XAI_OAUTH2_LOCAL_ISSUER -} -/// auth.json scope key used by the pre-OIDC `grok login --legacy` flow. -/// Matches the key format produced by the original `accounts.x.ai` relay auth. -pub const LEGACY_AUTH_SCOPE: &str = "https://accounts.x.ai/sign-in"; -impl GrokComConfig { - /// Whether `xai.api_key` auth is disabled. Pinning a team - /// (`force_login_team_uuid`) implies this — team membership can't be verified - /// from a bare API key, so it must go through IdP login. The - /// `KIGI_DISABLE_API_KEY_AUTH` env lockdown is sticky: because the env value - /// seeds `default()` (the merge base), a lower-trust user `config.toml` could - /// otherwise set `disable_api_key_auth = false` and override it — so the env - /// is OR-ed in here and cannot be turned back off by a user layer. Trusted - /// `requirements.toml` already wins over `config.toml` via layer precedence. - pub fn api_key_auth_disabled(&self) -> bool { - self.disable_api_key_auth == Some(true) - || self.force_login_team_uuid.is_some() - || env_lockdown_forced() - } - /// When `preferred_method = api_key`, automatic OIDC paths (devbox mint, - /// interactive browser login, external auth provider) must not run — the - /// pin is fail-closed. Explicit `grok login --devbox` / `--api-key` bypass - /// this by not consulting automatic flow helpers. - pub fn blocks_automatic_oidc(&self) -> bool { - matches!(self.preferred_method, Some(PreferredAuthMethod::ApiKey)) - } - /// The auth.json scope key for this config. +pub struct KimiCodeConfig {} + +impl KimiCodeConfig { + /// The persisted-credential scope key for this configuration. pub fn auth_scope(&self) -> String { - if let Some(ref oidc) = self.oidc { - format!("{}::{}", oidc.issuer.trim_end_matches('/'), oidc.client_id) - } else if let Some(ref oauth2) = self.oauth2 { - oauth2.auth_scope() - } else { - unreachable!("oauth2 config is always present (xAI default or env override)") - } - } -} -impl OAuth2ProviderConfig { - pub fn is_team_principal(&self) -> bool { - self.principal_type.as_deref() == Some(TEAM_PRINCIPAL_TYPE) - } - pub fn from_env() -> Option { - let issuer = std::env::var("KIGI_OAUTH2_ISSUER").ok()?; - let client_id = std::env::var("KIGI_OAUTH2_CLIENT_ID").ok()?; - let principal_type = std::env::var("KIGI_OAUTH2_PRINCIPAL_TYPE").ok(); - let principal_id = std::env::var("KIGI_OAUTH2_PRINCIPAL_ID").ok(); - let default_scopes = match principal_type.as_deref() { - Some(TEAM_PRINCIPAL_TYPE) => default_team_oauth2_scopes(), - _ => default_oauth2_scopes(), - }; - Some(Self { - issuer, - client_id, - scopes: std::env::var("KIGI_OAUTH2_SCOPES") - .map(|s| s.split(',').map(|s| s.trim().to_owned()).collect()) - .unwrap_or(default_scopes), - principal_type, - principal_id, - referrer: Some( - std::env::var("KIGI_OAUTH2_REFERRER") - .unwrap_or_else(|_| DEFAULT_OAUTH2_REFERRER.to_owned()), - ), - }) - } - /// Convert to [`OidcAuthConfig`] to reuse the OIDC login flow. - pub fn as_oidc(&self) -> OidcAuthConfig { - OidcAuthConfig { - issuer: self.issuer.clone(), - client_id: self.client_id.clone(), - scopes: self.scopes.clone(), - audience: None, - } - } - pub fn base_auth_scope(&self) -> String { - format!("{}::{}", self.issuer.trim_end_matches('/'), self.client_id) - } - pub fn auth_scope(&self) -> String { - self.base_auth_scope() - } -} -impl Default for GrokComConfig { - fn default() -> Self { - let oidc = OidcAuthConfig::from_env(); - let oauth2 = if oidc.is_some() { - None - } else { - Some( - OAuth2ProviderConfig::from_env().unwrap_or_else(|| OAuth2ProviderConfig { - issuer: xai_oauth2_issuer().to_owned(), - client_id: obfstr::obfstr!("b1a00492-073a-47ea-816f-4c329264a828").to_owned(), - scopes: default_oauth2_scopes(), - principal_type: None, - principal_id: None, - referrer: Some(DEFAULT_OAUTH2_REFERRER.to_owned()), - }), - ) - }; - Self { - grok_ws_origin: std::env::var("KIGI_WS_ORIGIN") - .unwrap_or_else(|_| AUTH_ORIGIN_DEFAULT.to_owned()), - token_header: "xai-grok-cli".to_owned(), - oidc, - oauth2, - auth_provider_command: std::env::var("KIGI_AUTH_PROVIDER_COMMAND").ok(), - auth_provider_label: std::env::var("KIGI_AUTH_PROVIDER_LABEL").ok(), - auth_token_ttl: std::env::var("KIGI_AUTH_TOKEN_TTL") - .ok() - .and_then(|v| v.parse().ok()), - disable_api_key_auth: std::env::var("KIGI_DISABLE_API_KEY_AUTH") - .ok() - .map(|v| env_flag_enabled(&v)), - force_login_team_uuid: None, - preferred_method: None, - } - } -} -/// Parse a boolean env-var value for grok's on/off flags. A bare presence -/// enables the flag, but the common falsy spellings (`0`, `false`, `off`, -/// `no`, empty) count as disabled — so e.g. `KIGI_DISABLE_API_KEY_AUTH=false` -/// does NOT turn the kill switch on. -fn env_flag_enabled(value: &str) -> bool { - !matches!( - value.trim().to_ascii_lowercase().as_str(), - "" | "0" | "false" | "off" | "no" - ) -} -/// True when the admin has set `KIGI_DISABLE_API_KEY_AUTH` to a truthy value in -/// the process environment. Read live (call-time) and OR-ed into -/// `api_key_auth_disabled()` so the env lockdown is non-overridable by a -/// user-layer `config.toml`. -fn env_lockdown_forced() -> bool { - std::env::var("KIGI_DISABLE_API_KEY_AUTH") - .ok() - .is_some_and(|v| env_flag_enabled(&v)) -} -impl OidcAuthConfig { - pub fn from_env() -> Option { - let issuer = std::env::var("KIGI_OIDC_ISSUER").ok()?; - let client_id = std::env::var("KIGI_OIDC_CLIENT_ID").ok()?; - Some(Self { - issuer, - client_id, - scopes: std::env::var("KIGI_OIDC_SCOPES") - .map(|s| s.split(',').map(|s| s.trim().to_owned()).collect()) - .unwrap_or_else(|_| default_oidc_scopes()), - audience: std::env::var("KIGI_OIDC_AUDIENCE").ok(), - }) + KIMI_CODE_OAUTH_SCOPE.to_owned() } } + #[cfg(test)] mod tests { use super::*; + #[test] - fn team_auth_scope_is_base_scope() { - let cfg = OAuth2ProviderConfig { - issuer: "https://auth.x.ai".into(), - client_id: "client-123".into(), - scopes: default_team_oauth2_scopes(), - principal_type: Some("Team".into()), - principal_id: Some("team-abc".into()), - referrer: Some("grok-build".into()), - }; - assert_eq!(cfg.auth_scope(), "https://auth.x.ai::client-123"); + fn auth_scope_is_the_kimi_code_key() { + assert_eq!(KimiCodeConfig::default().auth_scope(), "oauth/kimi-code"); } + #[test] - fn env_flag_enabled_treats_falsy_spellings_as_off() { - for off in ["", " ", "0", "false", "FALSE", "off", "No", " false "] { - assert!(!env_flag_enabled(off), "{off:?} should be off"); - } - for on in ["1", "true", "yes", "on", "enabled"] { - assert!(env_flag_enabled(on), "{on:?} should be on"); - } - } - #[test] - fn personal_auth_scope_is_base_scope() { - let cfg = OAuth2ProviderConfig { - issuer: "https://auth.x.ai".into(), - client_id: "client-123".into(), - scopes: default_oauth2_scopes(), - principal_type: None, - principal_id: None, - referrer: Some("grok-build".into()), - }; - assert_eq!(cfg.auth_scope(), "https://auth.x.ai::client-123"); - } - /// FROZEN loopback contract: the accounts-app origins the CLI's loopback - /// callback server accepts cross-origin requests from. The consent page - /// (served from accounts.x.ai) delivers the code via `fetch(..., cors)`, so - /// removing an origin breaks loopback delivery for already-installed CLIs. - /// Keep in sync with the oauth2-provider / accounts-app deployments. - /// Non-production / local-dev origins are opt-in only. - #[test] - fn allowed_accounts_app_origins_are_frozen() { - assert_eq!(PROD_ACCOUNTS_APP_ORIGINS, &["https://accounts.x.ai"]); - assert_eq!(allowed_accounts_app_origins(), PROD_ACCOUNTS_APP_ORIGINS); - } - /// FROZEN client contract: the 8 scopes the xAI OAuth2 client requests. - /// The server must keep accepting all of them; existing tokens carry - /// exactly this set. Frozen OAuth client scope contract. - #[test] - fn default_oauth2_scopes_are_frozen() { - let scopes = default_oauth2_scopes(); - let scopes: Vec<&str> = scopes.iter().map(String::as_str).collect(); - assert_eq!( - scopes, - [ - "openid", - "profile", - "email", - "offline_access", - "grok-cli:access", - "api:access", - "conversations:read", - "conversations:write", - ] - ); - } - #[test] - fn preferred_method_deserializes_from_toml() { - let cfg: GrokComConfig = toml::from_str( - r#" - preferred_method = "api_key" - "#, - ) - .expect("parse"); - assert_eq!(cfg.preferred_method, Some(PreferredAuthMethod::ApiKey)); - let cfg: GrokComConfig = toml::from_str( - r#" - preferred_method = "oidc" - "#, - ) - .expect("parse"); - assert_eq!(cfg.preferred_method, Some(PreferredAuthMethod::Oidc)); - let cfg: GrokComConfig = toml::from_str("").expect("parse empty"); - assert_eq!(cfg.preferred_method, None); + fn deserializes_from_empty_toml() { + let cfg: KimiCodeConfig = toml::from_str("").expect("empty config parses"); + assert_eq!(cfg.auth_scope(), KIMI_CODE_OAUTH_SCOPE); } } diff --git a/crates/codegen/kigi-shell/src/auth/credential_provider.rs b/crates/codegen/kigi-shell/src/auth/credential_provider.rs index 340bdc4..d7df75e 100644 --- a/crates/codegen/kigi-shell/src/auth/credential_provider.rs +++ b/crates/codegen/kigi-shell/src/auth/credential_provider.rs @@ -1,5 +1,5 @@ use crate::auth::AuthManager; -use crate::util::grok_auth_credentials::GrokAuthCredentials; +use crate::util::kigi_auth_credentials::KigiAuthCredentials; use kigi_auth::{ AuthCredentialProvider, CredentialSnapshot, HttpAuth, StaticAuthCredentialProvider, }; @@ -7,7 +7,7 @@ use reqwest::RequestBuilder; use std::sync::Arc; /// `api_key.id` for the active credential: hash the stable API key, never the /// OIDC bearer (which rotates). `None` for non-API-key auth. -fn api_key_id_for(auth: Option<&crate::auth::GrokAuth>) -> Option { +fn api_key_id_for(auth: Option<&crate::auth::KimiAuth>) -> Option { auth.filter(|a| matches!(a.auth_mode, crate::auth::AuthMode::ApiKey)) .map(|a| crate::agent::config::deployment_id_from_key(&a.key)) } @@ -15,7 +15,7 @@ fn api_key_id_for(auth: Option<&crate::auth::GrokAuth>) -> Option { /// delegates to `AuthManager::unauthorized_recovery`. pub struct ShellAuthCredentialProvider { auth_manager: Arc, - static_credentials: GrokAuthCredentials, + static_credentials: KigiAuthCredentials, } impl ShellAuthCredentialProvider { pub(crate) fn new( @@ -23,7 +23,7 @@ impl ShellAuthCredentialProvider { deployment_key: Option, alpha_test_key: Option, ) -> Self { - let mut static_credentials = GrokAuthCredentials::new(None); + let mut static_credentials = KigiAuthCredentials::new(None); static_credentials.deployment_key = deployment_key; static_credentials.alpha_test_key = alpha_test_key; Self { @@ -61,18 +61,19 @@ impl AuthCredentialProvider for ShellAuthCredentialProvider { }; } let auth = self.auth_manager.current_or_expired(); - let user_id = auth.as_ref().map(|a| a.user_id.clone()); - let team_id = auth.as_ref().and_then(|a| a.team_id.clone()); - let organization_id = auth.as_ref().and_then(|a| a.organization_id.clone()); + // The Kimi token response carries no account info; `user_id` stays + // empty until a later feature surfaces it. + let user_id = auth + .as_ref() + .map(|a| a.user_id.clone()) + .filter(|id| !id.is_empty()); let api_key_id = api_key_id_for(auth.as_ref()); let token = auth.map(|a| a.key); CredentialSnapshot { token, user_id, - team_id, deployment_id: None, api_key_id, - organization_id, } } async fn refresh_after_unauthorized(&self) -> bool { @@ -81,15 +82,12 @@ impl AuthCredentialProvider for ShellAuthCredentialProvider { } self.auth_manager.try_recover_unauthorized().await } - fn needs_token_auth_header(&self) -> bool { - self.static_credentials.deployment_key.is_none() - } } #[cfg(test)] mod tests { use super::*; - use crate::auth::GrokAuth; - use crate::auth::GrokComConfig; + use crate::auth::KimiAuth; + use crate::auth::KimiCodeConfig; use crate::auth::manager::AuthManager; use chrono::{Duration as ChronoDuration, Utc}; use kigi_auth::AuthCredentialProvider; @@ -128,19 +126,19 @@ mod tests { } } } - fn make_auth(key: &str, expires_in: ChronoDuration) -> GrokAuth { - GrokAuth { + fn make_auth(key: &str, expires_in: ChronoDuration) -> KimiAuth { + KimiAuth { key: key.to_string(), user_id: "test-user".to_string(), create_time: Utc::now(), expires_at: Some(Utc::now() + expires_in), - ..GrokAuth::test_default() + ..KimiAuth::test_default() } } /// Build an `AuthManager` rooted at `dir`. Caller keeps `dir` alive for /// the duration of the test so the `TempDir` `Drop` actually cleans up. - fn make_manager(dir: &tempfile::TempDir, initial: Option) -> Arc { - let mgr = AuthManager::new(dir.path(), GrokComConfig::default()); + fn make_manager(dir: &tempfile::TempDir, initial: Option) -> Arc { + let mgr = AuthManager::new(dir.path(), KimiCodeConfig::default()); if let Some(auth) = initial { mgr.hot_swap(auth); } @@ -210,16 +208,16 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let mgr = Arc::new(AuthManager::new( dir.path(), - crate::auth::GrokComConfig::default(), + crate::auth::KimiCodeConfig::default(), )); - mgr.hot_swap(GrokAuth { + mgr.hot_swap(KimiAuth { key: "stale".into(), - auth_mode: crate::auth::AuthMode::Oidc, + auth_mode: crate::auth::AuthMode::OAuth, create_time: chrono::Utc::now() - ChronoDuration::hours(2), user_id: "u".into(), refresh_token: Some("rt-stale".into()), expires_at: Some(chrono::Utc::now() - ChronoDuration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }); struct OkRefresher { calls: Arc, @@ -231,14 +229,14 @@ mod tests { _r: crate::auth::manager::RefreshReason, ) -> crate::auth::refresh::RefreshOutcome { self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - crate::auth::refresh::RefreshOutcome::Success(Box::new(GrokAuth { + crate::auth::refresh::RefreshOutcome::Success(Box::new(KimiAuth { key: "fresh".into(), - auth_mode: crate::auth::AuthMode::Oidc, + auth_mode: crate::auth::AuthMode::OAuth, create_time: chrono::Utc::now(), user_id: "u".into(), refresh_token: Some("rt-new".into()), expires_at: Some(chrono::Utc::now() + ChronoDuration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() })) } } @@ -282,11 +280,11 @@ mod tests { Some(deployment_id_from_key("xai-token-EX").as_str()) ); assert!(dep.api_key_id.is_none()); - let api_auth = GrokAuth { + let api_auth = KimiAuth { key: "sk-apikey-xyz".into(), auth_mode: crate::auth::AuthMode::ApiKey, expires_at: Some(Utc::now() + ChronoDuration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }; let api = ShellAuthCredentialProvider::new(make_manager(&dir, Some(api_auth)), None, None) .snapshot(); diff --git a/crates/codegen/kigi-shell/src/auth/devbox_login_stub.rs b/crates/codegen/kigi-shell/src/auth/devbox_login_stub.rs deleted file mode 100644 index e6cc026..0000000 --- a/crates/codegen/kigi-shell/src/auth/devbox_login_stub.rs +++ /dev/null @@ -1,37 +0,0 @@ -//! Stub for builds without the devbox auth feature. -//! -//! Compiled instead of `devbox_login.rs` when the devbox auth feature is -//! off, so the remote devbox login helper is not reached. The API -//! mirrors the real module: `is_devbox_environment()` is always `false`, which -//! short-circuits every auto-recovery/migration call site, and the entry -//! points that can still be reached directly (`grok login --devbox`) return a -//! descriptive error. - -use super::manager::AuthManager; -use super::model::GrokAuth; - -const UNAVAILABLE: &str = - "devbox login is not available in this build (compiled without the `devbox-login` feature)"; - -/// Always `false` without the devbox auth feature; callers treat the -/// process as running outside a devbox environment. -pub(crate) fn is_devbox_environment() -> bool { - false -} - -/// Unreachable in practice (guarded by [`is_devbox_environment`]); errors -/// defensively if called. -pub(crate) async fn mint_devbox_auth(_auth_manager: &AuthManager) -> anyhow::Result { - anyhow::bail!(UNAVAILABLE) -} - -/// Unreachable in practice (guarded by [`is_devbox_environment`]); errors -/// defensively if called. -pub(super) async fn mint_devbox_auth_raw() -> anyhow::Result { - anyhow::bail!(UNAVAILABLE) -} - -/// `grok login --devbox` entry point: always errors in this build. -pub async fn run_devbox_login(_config: &crate::agent::config::Config) -> anyhow::Result { - anyhow::bail!(UNAVAILABLE) -} diff --git a/crates/codegen/kigi-shell/src/auth/device.rs b/crates/codegen/kigi-shell/src/auth/device.rs new file mode 100644 index 0000000..9cb86f9 --- /dev/null +++ b/crates/codegen/kigi-shell/src/auth/device.rs @@ -0,0 +1,297 @@ +//! Device identity headers for the Kimi Code OAuth endpoints. +//! +//! Every OAuth call (device authorization, token poll, refresh) carries three +//! headers identifying this installation (PRD F1): +//! +//! - `X-Msh-Device-Name` — the local hostname +//! - `X-Msh-Device-Model` — an honest local OS/arch string (e.g. +//! "macOS 15.5 arm64"), ported from kimi-cli's `_device_model()` +//! - `X-Msh-Device-Id` — a uuid4 hex persisted at `~/.kigi/device_id` +//! (owner-only), created on first use +//! +//! All values are ASCII-sanitized (ported from kimi-cli's +//! `_ascii_header_value`) since HTTP header values must be ASCII. + +use std::path::PathBuf; +use std::sync::OnceLock; + +use anyhow::Context as _; + +/// Sanitize a header value to ASCII: non-ASCII bytes are dropped; an empty +/// result falls back to `"unknown"`. Port of kimi-cli `_ascii_header_value`. +pub(crate) fn ascii_header_value(value: &str) -> String { + let sanitized: String = value.chars().filter(char::is_ascii).collect(); + let trimmed = sanitized.trim(); + if trimmed.is_empty() { + "unknown".to_owned() + } else { + trimmed.to_owned() + } +} + +/// The three device-identity headers sent on every OAuth call. +/// +/// Errors when the persistent device id cannot be created (e.g. read-only +/// `~/.kigi`): the OAuth endpoints require `X-Msh-Device-Id`, so login cannot +/// proceed without it. +pub(crate) fn device_headers() -> anyhow::Result<[(&'static str, String); 3]> { + Ok([ + ("X-Msh-Device-Name", ascii_header_value(&device_name())), + ("X-Msh-Device-Model", ascii_header_value(device_model())), + ("X-Msh-Device-Id", ascii_header_value(&device_id()?)), + ]) +} + +/// Local hostname (kimi-cli: `platform.node() or socket.gethostname()`). +fn device_name() -> String { + #[cfg(unix)] + { + let mut buf = [0u8; 256]; + // SAFETY: buf is a valid writable buffer of the passed length. + let rc = unsafe { libc::gethostname(buf.as_mut_ptr().cast(), buf.len()) }; + if rc == 0 { + let end = buf.iter().position(|&b| b == 0).unwrap_or(buf.len()); + let name = String::from_utf8_lossy(&buf[..end]).into_owned(); + if !name.trim().is_empty() { + return name; + } + } + "unknown".to_owned() + } + #[cfg(windows)] + { + std::env::var("COMPUTERNAME").unwrap_or_else(|_| "unknown".to_owned()) + } + #[cfg(not(any(unix, windows)))] + { + "unknown".to_owned() + } +} + +/// Honest local device-model string, computed once per process. Port of +/// kimi-cli `_device_model()`: +/// - macOS → `macOS {product_version} {arch}` (e.g. "macOS 15.5 arm64") +/// - Windows → `Windows {10|11} {arch}` (build ≥ 22000 reports 11) +/// - other → `{sysname} {kernel_release} {machine}` +pub(crate) fn device_model() -> &'static str { + static MODEL: OnceLock = OnceLock::new(); + MODEL.get_or_init(compute_device_model) +} + +fn compute_device_model() -> String { + #[cfg(target_os = "macos")] + { + // Match Python's platform.machine() spelling on macOS. + let arch = match std::env::consts::ARCH { + "aarch64" => "arm64", + other => other, + }; + match macos_product_version() { + Some(version) => format!("macOS {version} {arch}"), + None => format!("macOS {arch}"), + } + } + #[cfg(windows)] + { + let arch = std::env::consts::ARCH; + match windows_release() { + Some(release) => format!("Windows {release} {arch}"), + None => format!("Windows {arch}"), + } + } + #[cfg(not(any(target_os = "macos", windows)))] + { + let (sysname, release, machine) = uname_fields(); + match (release, machine) { + (Some(r), Some(m)) => format!("{sysname} {r} {m}"), + (Some(r), None) => format!("{sysname} {r}"), + (None, Some(m)) => format!("{sysname} {m}"), + (None, None) => sysname, + } + } +} + +/// macOS product version (e.g. "15.5") from the SystemVersion plist — the +/// same source Python's `platform.mac_ver()` reads. +#[cfg(target_os = "macos")] +fn macos_product_version() -> Option { + let plist = std::fs::read_to_string("/System/Library/CoreServices/SystemVersion.plist").ok()?; + plist_string_value(&plist, "ProductVersion") +} + +/// Extract `{key}value` from a plist XML body. +#[cfg(target_os = "macos")] +fn plist_string_value(plist: &str, key: &str) -> Option { + let key_tag = format!("{key}"); + let after_key = &plist[plist.find(&key_tag)? + key_tag.len()..]; + let start = after_key.find("")? + "".len(); + let end = after_key.find("")?; + (start <= end).then(|| after_key[start..end].trim().to_owned()) +} + +/// Windows major release ("10" or "11"), from the build number reported by +/// `cmd /c ver` (kimi-cli: `sys.getwindowsversion().build >= 22000` → 11). +#[cfg(windows)] +fn windows_release() -> Option { + let output = std::process::Command::new("cmd") + .args(["/c", "ver"]) + .output() + .ok()?; + let text = String::from_utf8_lossy(&output.stdout); + // "Microsoft Windows [Version 10.0.22631.3155]" + let version = text.split("Version").nth(1)?.trim(); + let mut parts = version.trim_end_matches(']').split('.'); + let major = parts.next()?.trim().to_owned(); + let _minor = parts.next()?; + let build: u32 = parts.next()?.trim().parse().ok()?; + if major == "10" && build >= 22000 { + Some("11".to_owned()) + } else { + Some(major) + } +} + +/// `uname(2)` sysname / release / machine for Linux and other Unix. +#[cfg(all(unix, not(target_os = "macos")))] +fn uname_fields() -> (String, Option, Option) { + // SAFETY: utsname is a plain-old-data struct; uname fills it in. + let mut uts: libc::utsname = unsafe { std::mem::zeroed() }; + if unsafe { libc::uname(&mut uts) } != 0 { + return (std::env::consts::OS.to_owned(), None, None); + } + fn field(raw: &[libc::c_char]) -> Option { + let bytes: Vec = raw + .iter() + .take_while(|&&c| c != 0) + .map(|&c| c as u8) + .collect(); + let s = String::from_utf8_lossy(&bytes).trim().to_owned(); + (!s.is_empty()).then_some(s) + } + ( + field(&uts.sysname).unwrap_or_else(|| std::env::consts::OS.to_owned()), + field(&uts.release), + field(&uts.machine), + ) +} + +#[cfg(not(unix))] +#[cfg(not(windows))] +fn uname_fields() -> (String, Option, Option) { + (std::env::consts::OS.to_owned(), None, None) +} + +/// Path of the persistent device id: `{kigi_home}/device_id`. +fn device_id_path() -> PathBuf { + kigi_config::kigi_home().join("device_id") +} + +/// Persistent uuid4-hex device id, created (owner-only, 0o600) on first use +/// and cached for the process lifetime. +pub(crate) fn device_id() -> anyhow::Result { + static DEVICE_ID: OnceLock = OnceLock::new(); + if let Some(id) = DEVICE_ID.get() { + return Ok(id.clone()); + } + let id = load_or_create_device_id(&device_id_path())?; + Ok(DEVICE_ID.get_or_init(|| id).clone()) +} + +/// Read `path`, or mint a uuid4 hex and persist it owner-only. +fn load_or_create_device_id(path: &std::path::Path) -> anyhow::Result { + if let Ok(existing) = std::fs::read_to_string(path) { + let trimmed = existing.trim(); + if !trimmed.is_empty() { + return Ok(trimmed.to_owned()); + } + } + let id = uuid::Uuid::new_v4().simple().to_string(); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent) + .with_context(|| format!("creating {} for device_id", parent.display()))?; + } + std::fs::write(path, &id) + .with_context(|| format!("writing device id to {}", path.display()))?; + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600)) + .with_context(|| format!("chmod 600 {}", path.display()))?; + } + tracing::info!(path = %path.display(), "auth: created persistent device id"); + Ok(id) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ascii_header_value_passes_ascii_through() { + assert_eq!(ascii_header_value("macOS 15.5 arm64"), "macOS 15.5 arm64"); + assert_eq!(ascii_header_value(" padded "), "padded"); + } + + #[test] + fn ascii_header_value_strips_non_ascii() { + assert_eq!(ascii_header_value("café-host"), "caf-host"); + assert_eq!(ascii_header_value("机器"), "unknown"); + assert_eq!(ascii_header_value(" "), "unknown"); + } + + #[test] + fn device_model_is_nonempty_ascii() { + let model = device_model(); + assert!(!model.is_empty()); + assert!(model.is_ascii(), "device model must be ASCII: {model:?}"); + // The honest local OS name must lead the string. + #[cfg(target_os = "macos")] + assert!(model.starts_with("macOS "), "got {model:?}"); + #[cfg(windows)] + assert!(model.starts_with("Windows"), "got {model:?}"); + } + + #[test] + fn load_or_create_device_id_roundtrips_and_is_owner_only() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("device_id"); + let created = load_or_create_device_id(&path).unwrap(); + assert_eq!(created.len(), 32, "uuid4 hex is 32 chars: {created:?}"); + assert!(created.chars().all(|c| c.is_ascii_hexdigit())); + // Second call reads the same id back. + let reread = load_or_create_device_id(&path).unwrap(); + assert_eq!(created, reread); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let mode = std::fs::metadata(&path).unwrap().permissions().mode(); + assert_eq!(mode & 0o777, 0o600, "device_id must be owner-only"); + } + } + + #[test] + fn load_or_create_device_id_ignores_empty_file() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("device_id"); + std::fs::write(&path, " \n").unwrap(); + let created = load_or_create_device_id(&path).unwrap(); + assert_eq!(created.len(), 32); + } + + #[cfg(target_os = "macos")] + #[test] + fn plist_string_value_extracts_product_version() { + let plist = r#" + + ProductBuildVersion + 24F74 + ProductVersion + 15.5 +"#; + assert_eq!( + plist_string_value(plist, "ProductVersion").as_deref(), + Some("15.5") + ); + assert_eq!(plist_string_value(plist, "Missing"), None); + } +} diff --git a/crates/codegen/kigi-shell/src/auth/device_code.rs b/crates/codegen/kigi-shell/src/auth/device_code.rs index 210bbdc..a69a69b 100644 --- a/crates/codegen/kigi-shell/src/auth/device_code.rs +++ b/crates/codegen/kigi-shell/src/auth/device_code.rs @@ -1,387 +1,105 @@ -//! RFC 8628 Device Authorization Grant -- CLI side. +//! Kimi Code device-code login (PRD F1). //! -//! Two-phase API: -//! 1. `request_device_code()` -- POST to server, get code + URL -//! 2. `complete_device_code_login()` -- poll until approved, persist credentials +//! Two-phase API mirroring kimi-cli's `login_kimi_code`: +//! 1. [`crate::auth::kimi_oauth::request_device_authorization`] — get a +//! user code + verification URL from the OAuth host +//! 2. [`complete_device_code_login`] — poll the token endpoint until +//! approved, then persist the token set via the `AuthManager` //! -//! Callers control what happens between the two phases (print to stderr, -//! show in TUI, display in IDE sidebar, etc.). +//! Poll semantics: the server-provided interval (default 5s, floored at 1s) +//! paces the loop; `slow_down` bumps it by 5s; `expired_token` restarts the +//! whole device authorization (fresh user code); every other non-200 outcome +//! (`authorization_pending`, unknown errors) waits and continues. use std::sync::Arc; -use chrono::{Duration, Utc}; -use serde::Deserialize; -use thiserror::Error; +use crate::auth::kimi_oauth::{ + DeviceAuthorization, DevicePollResult, poll_device_token, request_device_authorization, +}; +use crate::auth::{AuthChannels, AuthManager, AuthUrlInfo, AuthUrlMode, KimiAuth}; -use crate::auth::oidc::with_alpha_test_key; -use crate::auth::{AuthChannels, AuthManager, AuthMode, AuthUrlInfo, AuthUrlMode, GrokAuth}; +/// Extra wait added to the poll interval when the server answers `slow_down` +/// (OAuth-standard device-flow backpressure). +const SLOW_DOWN_INCREMENT_SECS: u64 = 5; -const DEVICE_GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type:device_code"; -const DEFAULT_DEVICE_POLL_INTERVAL_SECS: i32 = 5; -const DEVICE_SLOW_DOWN_INCREMENT_SECS: u64 = 5; -const MIN_DEVICE_CODE_EXPIRY_FALLBACK_SECS: i64 = 10 * 60; - -#[derive(Debug, Error)] -pub enum DeviceCodeError { - #[error( - "Device-code login is not available for this deployment. \ - Try `grok login` or set XAI_API_KEY instead." - )] - NotEnabled, - #[error(transparent)] - Other(#[from] anyhow::Error), -} - -impl From for DeviceCodeError { - fn from(e: reqwest::Error) -> Self { - Self::Other(e.into()) - } -} - -// --- Public types --- - -/// Low-cardinality client-surface hint sent to the OAuth2 provider as the -/// `x-grok-client-surface` header so device-flow metrics can separate logins a -/// human can actually finish (`Ui`, `Cli`) from headless automation -/// (`Headless`) that mints a device code but can never reach the browser -/// consent page — the traffic that otherwise pollutes the device-flow -/// conversion denominator. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ClientSurface { - /// An interactive front-end (TUI / IDE) renders the URL + code to a human. - Ui, - /// CLI attached to an interactive terminal (stderr is a TTY). - Cli, - /// No interactive surface (CI, container, script): no human can complete. - Headless, -} - -impl ClientSurface { - fn as_str(self) -> &'static str { - match self { - Self::Ui => "ui", - Self::Cli => "cli", - Self::Headless => "headless", - } - } -} - -/// Classify the CLI (non-TUI) surface: a TTY on stderr means a human is -/// watching the printed URL + code; otherwise we're headless (CI/container/ -/// script) and no one will complete the flow. -fn detect_cli_surface() -> ClientSurface { - use std::io::IsTerminal as _; - if std::io::stderr().is_terminal() { - ClientSurface::Cli - } else { - ClientSurface::Headless - } -} - -/// Result of requesting a device code from the server. -/// Callers display `verification_uri` + `user_code` to the user, -/// then pass this struct to `complete_device_code_login`. -#[derive(Debug, Clone)] -pub struct DeviceCode { - pub verification_uri: String, - pub verification_uri_complete: Option, - pub user_code: String, - device_code: String, - interval: i32, - expires_in: i64, -} - -// --- Wire types (serde) --- - -#[derive(Deserialize)] -struct DeviceCodeResponse { - device_code: String, - user_code: String, - verification_uri: String, - verification_uri_complete: Option, - expires_in: i64, - interval: Option, -} - -#[derive(Deserialize)] -struct TokenOk { - access_token: String, - refresh_token: Option, - expires_in: Option, - #[expect(dead_code, reason = "field retained for protocol compatibility")] - scope: Option, - id_token: Option, -} - -#[derive(Deserialize)] -struct TokenErr { - error: String, - error_description: Option, -} - -#[derive(Deserialize)] -struct IdTokenClaims { - sub: Option, - email: Option, -} - -// --- Phase 1: Request device code --- - -/// Request a device code + user code from the OAuth2 provider. -/// -/// This is a single HTTP POST. The caller is responsible for displaying -/// `DeviceCode::verification_uri` and `DeviceCode::user_code` to the user -/// before calling `complete_device_code_login`. -pub async fn request_device_code( - issuer: &str, - client_id: &str, - scopes: &[String], - surface: ClientSurface, -) -> Result { - let client = crate::http::shared_client(); - let url = format!("{}/oauth2/device/code", issuer.trim_end_matches('/')); - let scope_str = scopes.join(" "); - - let resp = with_alpha_test_key( - client - .post(&url) - // Lets oauth2-provider segment device-flow success by client version. - .header("x-grok-client-version", kigi_version::VERSION) - // Lets oauth2-provider separate human-completable logins from - // headless automation in the device-flow funnel metrics. - .header("x-grok-client-surface", surface.as_str()) - .form(&[ - ("client_id", client_id), - ("scope", scope_str.as_str()), - ("referrer", "grok-build"), - ]), - &url, - ) - .send() - .await?; - - if !resp.status().is_success() { - let status = resp.status(); - let body = resp.text().await.unwrap_or_default(); - if status.as_u16() == 404 { - return Err(DeviceCodeError::NotEnabled); - } - return Err(anyhow::anyhow!("Device code request failed (HTTP {status}): {body}").into()); - } - - let server_resp: DeviceCodeResponse = resp.json().await?; - - // Defend against control characters from a malicious issuer. - if !server_resp - .user_code - .chars() - .all(|c| c.is_ascii_alphanumeric() || c == '-') - { - return Err(anyhow::anyhow!( - "Server returned invalid user_code format (expected [A-Z0-9-])" - ) - .into()); - } - - validate_verification_uri(&server_resp.verification_uri)?; - if let Some(ref verification_uri_complete) = server_resp.verification_uri_complete { - validate_verification_uri(verification_uri_complete)?; - } - - Ok(DeviceCode { - verification_uri: server_resp.verification_uri, - verification_uri_complete: server_resp.verification_uri_complete, - user_code: server_resp.user_code, - device_code: server_resp.device_code, - interval: server_resp - .interval - .unwrap_or(DEFAULT_DEVICE_POLL_INTERVAL_SECS), - expires_in: server_resp.expires_in, - }) -} - -// --- Phase 2: Poll until approved --- - -/// Poll the token endpoint until the user approves (or denies / expires). -/// -/// On success, persists credentials to `~/.kigi/auth.json` and returns -/// the authenticated `GrokAuth`. -/// -/// Callers should have already displayed `device_code.verification_uri` -/// and `device_code.user_code` to the user before calling this. -pub async fn complete_device_code_login( - issuer: &str, - client_id: &str, - device_code: DeviceCode, - auth_manager: &Arc, - surface: ClientSurface, -) -> anyhow::Result<(GrokAuth, bool)> { - let client = crate::http::shared_client(); - let token_url = format!("{}/oauth2/token", issuer.trim_end_matches('/')); - let mut poll_interval = std::time::Duration::from_secs(device_code.interval.max(1) as u64); - let deadline = tokio::time::Instant::now() - + std::time::Duration::from_secs( - device_code - .expires_in - .max(MIN_DEVICE_CODE_EXPIRY_FALLBACK_SECS) as u64, - ); - - loop { - // Sleep first: an immediate poll on a fresh code only returns - // authorization_pending (and risks slow_down). - tokio::time::sleep(poll_interval).await; - - if tokio::time::Instant::now() > deadline { - anyhow::bail!("Device code expired. Run `grok login --device-auth` again."); - } - - let resp = with_alpha_test_key( - client - .post(&token_url) - .header("x-grok-client-version", kigi_version::VERSION) - .header("x-grok-client-surface", surface.as_str()) - .form(&[ - ("grant_type", DEVICE_GRANT_TYPE), - ("device_code", device_code.device_code.as_str()), - ("client_id", client_id), - ]), - &token_url, - ) - .send() - .await?; - - if resp.status().is_success() { - let tokens: TokenOk = resp.json().await?; - let auth = build_auth(&tokens, issuer, client_id, auth_manager).await?; - return Ok((auth, true)); - } - - let err: TokenErr = resp.json().await?; - let detail = err.error_description.as_deref().unwrap_or(&err.error); - match err.error.as_str() { - "authorization_pending" => { - // User hasn't acted yet -- keep polling. - continue; - } - "slow_down" => { - poll_interval += std::time::Duration::from_secs(DEVICE_SLOW_DOWN_INCREMENT_SECS); - continue; - } - "access_denied" => { - tracing::warn!(description = detail, "device auth authorization denied"); - anyhow::bail!("Authorization denied. The user rejected the request."); - } - "expired_token" => { - tracing::warn!(description = detail, "device auth token expired"); - anyhow::bail!("Device code expired. Run `grok login --device-auth` again."); - } - other => { - tracing::warn!( - error = other, - description = detail, - "device auth token exchange failed" - ); - anyhow::bail!("Token exchange error: {detail}"); - } - } - } +/// Outcome of one full poll loop over a single device authorization. +enum PollLoopOutcome { + /// Access token issued. + Done(Box), + /// The device code expired before the user approved — request a fresh + /// authorization and start over. + Restart, } /// Device-code login shared by the TUI and CLI. /// -/// With `channels` (TUI) the verification URL goes to `url_tx` and the browser -/// opens automatically; on failure the copyable URL is the fallback. Without -/// `channels` (CLI) the URL + code are printed to stderr via `prompt_and_poll`. -/// `code_rx` is unused here. The caller reports success (`✓ Signed in`). -/// -/// Takes `channels` by `&mut`, consuming it only after the device code is -/// obtained, so callers can reuse it for a loopback fallback on `NotEnabled`. +/// With `channels` (TUI) the verification URL goes to `url_tx` and the +/// browser opens automatically. Without `channels` (CLI) the URL + code are +/// printed to stderr. The caller reports success (`✓ Signed in`). pub async fn run_device_code_login_channels( - issuer: &str, - client_id: &str, - scopes: &[String], + host: &str, auth_manager: &Arc, channels: &mut Option, -) -> anyhow::Result<(GrokAuth, bool)> { - // A front-end (TUI/IDE) listening on `url_tx` renders the URL to a human, so - // it's `Ui`. Without one we're on the CLI: a TTY means a human can act - // (`Cli`), no TTY means headless automation (`Headless`) that will never - // complete. Computed before `take()` so the `request_device_code` call - // already carries the surface. - let surface = if channels.is_some() { - ClientSurface::Ui - } else { - detect_cli_surface() - }; +) -> anyhow::Result<(KimiAuth, bool)> { + let interactive_tui = channels.is_some(); + let mut channels = channels.take(); + loop { + let device_auth = request_device_authorization(host).await?; + let display_uri = device_auth.verification_uri_complete.clone(); - let device_code = request_device_code(issuer, client_id, scopes, surface).await?; + if interactive_tui { + // TUI: push the URL through the channel BEFORE opening the + // browser, so the UI isn't blocked on a slow/hanging browser + // launch (e.g. SSH/headless). + if let Some(tx) = channels.as_mut().and_then(|c| c.url_tx.take()) { + let _ = tx.send(AuthUrlInfo { + url: display_uri.clone(), + mode: AuthUrlMode::Device, + }); + } + open_browser_detached(&display_uri).await; + } else { + prompt_on_stderr(&device_auth).await; + } - let Some(channels) = channels.take() else { - // CLI: print the URL + code to stderr. - return prompt_and_poll(issuer, client_id, device_code, auth_manager, surface).await; - }; - - // TUI: push the URL through the channel BEFORE opening the browser, so - // `x.ai/auth/get_url` isn't blocked on a slow/hanging browser launch - // (e.g. SSH/headless). When the issuer omits `verification_uri_complete`, - // embed the code so the welcome screen can still show it (anti-phishing). - let display_uri = match device_code.verification_uri_complete.as_deref() { - Some(uri) => uri.to_owned(), - None => { - let sep = if device_code.verification_uri.contains('?') { - '&' - } else { - '?' - }; - format!( - "{}{}user_code={}", - device_code.verification_uri, sep, device_code.user_code - ) + match complete_device_code_login(host, &device_auth).await? { + PollLoopOutcome::Done(auth) => { + let auth = auth_manager + .update(*auth) + .await + .map_err(|e| anyhow::anyhow!("Failed to save credentials: {e}"))?; + return Ok((auth, true)); + } + PollLoopOutcome::Restart => { + tracing::info!("auth: device code expired, restarting device authorization"); + if !interactive_tui { + eprintln!("Device code expired — requesting a new one..."); + } + // The TUI already consumed url_tx; the restarted flow can + // only reach the user via the (re-)opened browser page. + continue; + } } - }; - if let Some(tx) = channels.url_tx { - let _ = tx.send(AuthUrlInfo { - url: display_uri.clone(), - mode: AuthUrlMode::Device, - }); } - open_browser_detached(&display_uri).await; - complete_device_code_login(issuer, client_id, device_code, auth_manager, surface).await } -/// Display the device code to stderr and poll until approved. -async fn prompt_and_poll( - issuer: &str, - client_id: &str, - device_code: DeviceCode, - auth_manager: &Arc, - surface: ClientSurface, -) -> anyhow::Result<(GrokAuth, bool)> { - let display_uri = device_code - .verification_uri_complete - .as_deref() - .unwrap_or(&device_code.verification_uri); - +/// Print the verification URL + user code to stderr and open the browser. +async fn prompt_on_stderr(device_auth: &DeviceAuthorization) { + let display_uri = &device_auth.verification_uri_complete; eprintln!(); eprintln!("To sign in, open this URL in your browser:"); eprintln!(); - eprintln!(" {}", display_uri); + eprintln!(" {display_uri}"); eprintln!(); - if !open_browser_detached(display_uri).await { eprintln!(" (Could not open browser automatically — open the URL above manually.)"); eprintln!(); } - - // Show the code to confirm it matches the browser (anti-phishing): a complete - // URL pre-fills it (just confirm), otherwise the user types it. - if device_code.verification_uri_complete.is_some() { - eprintln!("Confirm this code in your browser:"); - } else { - eprintln!("Then enter this code:"); - } + // Show the code so the user can confirm it matches the browser + // (anti-phishing): the complete URL pre-fills it. + eprintln!("Confirm this code in your browser:"); eprintln!(); - eprintln!(" {}", device_code.user_code); + eprintln!(" {}", device_auth.user_code); eprintln!(); eprintln!( "\x1b[90mOnly continue with a code you requested. \ @@ -389,10 +107,42 @@ async fn prompt_and_poll( ); eprintln!(); eprintln!("Waiting for authorization..."); +} - // The caller prints the `✓ Signed in` confirmation (it also owns the - // external-provider / devbox early-return paths that never reach here). - complete_device_code_login(issuer, client_id, device_code, auth_manager, surface).await +/// Poll the token endpoint until the user approves, the device code expires +/// (→ [`PollLoopOutcome::Restart`]), or the wire fails. +async fn complete_device_code_login( + host: &str, + device_auth: &DeviceAuthorization, +) -> anyhow::Result { + let mut poll_interval = std::time::Duration::from_secs(device_auth.interval.max(1) as u64); + loop { + // Sleep first: an immediate poll on a fresh code only returns + // authorization_pending (and risks slow_down). + tokio::time::sleep(poll_interval).await; + match poll_device_token(host, &device_auth.device_code).await? { + DevicePollResult::Success(auth) => { + tracing::info!("auth: device login authorized"); + return Ok(PollLoopOutcome::Done(auth)); + } + DevicePollResult::Expired => return Ok(PollLoopOutcome::Restart), + DevicePollResult::Pending { error, description } => { + if error == "slow_down" { + poll_interval += std::time::Duration::from_secs(SLOW_DOWN_INCREMENT_SECS); + tracing::info!( + new_interval_secs = poll_interval.as_secs(), + "auth: server asked to slow down device polling" + ); + } else { + tracing::debug!( + error = %error, + description = ?description, + "auth: device authorization pending" + ); + } + } + } + } } /// Open `url` in the browser off-thread: `webbrowser::open` is synchronous and @@ -414,477 +164,197 @@ async fn open_browser_detached(url: &str) -> bool { } } -// --- Internal helpers --- - -/// No id_token signature verification -- token arrives over a direct HTTPS -/// channel (no browser redirect), and is only used for display info (email). -async fn build_auth( - tokens: &TokenOk, - issuer: &str, - client_id: &str, - auth_manager: &Arc, -) -> anyhow::Result { - let (user_id, email) = if let Some(ref id_token) = tokens.id_token { - decode_jwt_claims(id_token) - } else { - (String::new(), None) - }; - - let (principal_type, principal_id, token_team_id) = - match crate::auth::oidc::peek_access_token_principal(&tokens.access_token) { - Some((pt, pid, tid)) => (Some(pt), Some(pid), tid), - None => (None, None, None), - }; - - // Device flow has no pre-selection; verify the token's principal here. - // Match the principal id even if `principal_type` is absent. - let principal_policy = - crate::auth::oidc::login_principal_policy(auth_manager.grok_com_config()); - crate::auth::oidc::enforce_login_principal( - principal_policy.as_ref(), - crate::auth::oidc::peek_access_token_principal_id(&tokens.access_token).as_deref(), - )?; - - let (user_id, email, team_id, organization_id) = - match (principal_type.as_deref(), principal_id.as_deref()) { - (Some(pt), Some(principal_id)) if pt == crate::auth::model::TEAM_PRINCIPAL_TYPE => ( - principal_id.to_owned(), - None, - Some(principal_id.to_owned()), - None, - ), - (Some("Organization"), Some(principal_id)) => ( - principal_id.to_owned(), - None, - None, - Some(principal_id.to_owned()), - ), - _ => (user_id, email, token_team_id, None), - }; - - let now = Utc::now(); - let mut auth = GrokAuth { - key: tokens.access_token.clone(), - auth_mode: AuthMode::Oidc, - create_time: now, - user_id, - email, - first_name: None, - last_name: None, - profile_image_asset_id: None, - principal_type, - principal_id, - organization_id, - organization_name: None, - organization_role: None, - team_id, - team_name: None, - team_role: None, - user_blocked_reason: None, - team_blocked_reasons: vec![], - coding_data_retention_opt_out: false, - has_grok_code_access: None, - refresh_token: tokens.refresh_token.clone(), - expires_at: tokens.expires_in.map(|s| now + Duration::seconds(s)), - oidc_issuer: Some(issuer.to_owned()), - oidc_client_id: Some(client_id.to_owned()), - }; - - auth_manager.enrich_auth_inline(&mut auth).await; - - auth_manager - .update(auth) - .await - .map_err(|e| anyhow::anyhow!("Failed to save credentials: {e}")) -} - -/// Decode JWT payload without signature verification. -/// Returns (sub, Option). -fn decode_jwt_claims(jwt: &str) -> (String, Option) { - use base64::Engine; - let parts: Vec<&str> = jwt.splitn(3, '.').collect(); - if parts.len() < 2 { - return (String::new(), None); - } - let payload = match base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(parts[1]) { - Ok(bytes) => bytes, - Err(_) => return (String::new(), None), - }; - let claims: IdTokenClaims = match serde_json::from_slice(&payload) { - Ok(claims) => claims, - Err(_) => return (String::new(), None), - }; - (claims.sub.unwrap_or_default(), claims.email) -} - -fn validate_verification_uri(uri: &str) -> anyhow::Result<()> { - if uri.chars().any(|c| c.is_ascii_control()) { - anyhow::bail!("Server returned invalid verification URI"); - } - - let parsed = url::Url::parse(uri) - .map_err(|_| anyhow::anyhow!("Server returned invalid verification URI"))?; - - match parsed.scheme() { - "https" => Ok(()), - "http" if matches!(parsed.host_str(), Some("localhost") | Some("127.0.0.1")) => Ok(()), - _ => anyhow::bail!("Server returned unsupported verification URI scheme"), - } -} - #[cfg(test)] -pub(crate) mod tests { - use std::sync::Arc; +mod tests { + use super::*; + use crate::auth::KimiCodeConfig; + use wiremock::matchers::{body_string_contains, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; - use super::{AuthManager, build_auth, validate_verification_uri}; - use crate::auth::{AuthMode, GrokComConfig}; - - #[test] - fn validate_verification_uri_rejects_unsupported_scheme() { - let err = validate_verification_uri("javascript:alert(1)").unwrap_err(); - assert_eq!( - "Server returned unsupported verification URI scheme", - err.to_string() - ); - } - - fn auth_manager_with_kigi_home( - kigi_home: &std::path::Path, - proxy_base_url: &str, - ) -> Arc { - Arc::new( - AuthManager::new(kigi_home, GrokComConfig::default()) - .with_proxy_base_url(proxy_base_url), - ) - } - - #[test] - fn build_auth_persists_credentials_without_proxy_fetch() { - let temp_dir = tempfile::tempdir().unwrap(); - let kigi_home = temp_dir.path().join(".kigi"); - std::fs::create_dir_all(&kigi_home).unwrap(); - let auth_manager = auth_manager_with_kigi_home(&kigi_home, "http://127.0.0.1:9"); - let tokens = super::TokenOk { - access_token: "access-token".to_string(), - refresh_token: Some("refresh-token".to_string()), - expires_in: Some(900), - scope: Some("openid email offline_access grok-cli:access".to_string()), - id_token: Some( - "eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0.eyJzdWIiOiJ1c2VyLTEyMyIsImVtYWlsIjoiZGV2aWNlLWF1dGhAbG9jYWwudGVzdCJ9.sig".to_string(), - ), - }; - - let auth = tokio::runtime::Runtime::new() - .unwrap() - .block_on(build_auth( - &tokens, - "http://localhost:22255", - "client-id", - &auth_manager, - )) - .unwrap(); - - assert_eq!("access-token", auth.key); - assert_eq!(AuthMode::Oidc, auth.auth_mode); - assert_eq!("user-123", auth.user_id); - assert_eq!(Some("device-auth@local.test".to_string()), auth.email); - assert_eq!(Some("refresh-token".to_string()), auth.refresh_token); - assert_eq!(Some("http://localhost:22255".to_string()), auth.oidc_issuer); - assert_eq!(Some("client-id".to_string()), auth.oidc_client_id); - assert!(auth_manager.current().is_some()); - } - - /// jsonwebtoken needs a process-level CryptoProvider; tests that encode - /// JWTs can't rely on another test having installed it first. - fn ensure_crypto_provider() { - let _ = jsonwebtoken::crypto::rust_crypto::DEFAULT_PROVIDER.install_default(); - } - - #[test] - fn build_auth_seeds_team_metadata_from_access_token() { - ensure_crypto_provider(); - let temp_dir = tempfile::tempdir().unwrap(); - let kigi_home = temp_dir.path().join(".kigi"); - std::fs::create_dir_all(&kigi_home).unwrap(); - let auth_manager = auth_manager_with_kigi_home(&kigi_home, "http://127.0.0.1:9"); - let header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::HS256); - let claims = serde_json::json!({ - "sub": "user-42", - "iss": "https://auth.x.ai", - "aud": "client-id", - "exp": 9999999999u64, - "iat": 1000000000u64, - "scope": "offline_access grok-cli:access team:read", - "principal_type": "Team", - "principal_id": "team-123", - "client_id": "client-id", - "jti": "token-1", - }); - let tokens = super::TokenOk { - access_token: jsonwebtoken::encode( - &header, - &claims, - &jsonwebtoken::EncodingKey::from_secret(b"test-secret"), - ) - .unwrap(), - refresh_token: Some("refresh-token".to_owned()), - expires_in: Some(900), - scope: Some("offline_access grok-cli:access team:read".to_owned()), - id_token: None, - }; - - let auth = tokio::runtime::Runtime::new() - .unwrap() - .block_on(build_auth( - &tokens, - "http://localhost:22255", - "client-id", - &auth_manager, - )) - .unwrap(); - - assert_eq!("team-123", auth.user_id); - assert_eq!(Some("Team".to_owned()), auth.principal_type); - assert_eq!(Some("team-123".to_owned()), auth.principal_id); - assert_eq!(Some("team-123".to_owned()), auth.team_id); - assert_eq!(None, auth.organization_id); - assert_eq!(None, auth.email); - } - - /// Team access token carrying `principal_id` (signature irrelevant — only - /// the principal claims are peeked). - fn team_access_token(principal_id: &str) -> super::TokenOk { - let header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::HS256); - let claims = serde_json::json!({ - "sub": "user-42", - "exp": 9999999999u64, - "principal_type": "Team", - "principal_id": principal_id, - }); - super::TokenOk { - access_token: jsonwebtoken::encode( - &header, - &claims, - &jsonwebtoken::EncodingKey::from_secret(b"test-secret"), - ) - .unwrap(), - refresh_token: Some("refresh-token".to_owned()), - expires_in: Some(900), - scope: None, - id_token: None, - } - } - - /// `build_auth` with `token_principal` must fail with `expected_err` and - /// persist nothing. - fn assert_build_auth_rejected(cfg: GrokComConfig, token_principal: &str, expected_err: &str) { - ensure_crypto_provider(); - let temp_dir = tempfile::tempdir().unwrap(); - let kigi_home = temp_dir.path().join(".kigi"); - std::fs::create_dir_all(&kigi_home).unwrap(); - let auth_manager = - Arc::new(AuthManager::new(&kigi_home, cfg).with_proxy_base_url("http://127.0.0.1:9")); - - let err = tokio::runtime::Runtime::new() - .unwrap() - .block_on(build_auth( - &team_access_token(token_principal), - "http://localhost:22255", - "client-id", - &auth_manager, - )) - .unwrap_err(); - - assert_eq!(err.to_string(), expected_err); - assert!( - auth_manager.current().is_none(), - "rejected login must not persist credentials", - ); - assert!( - !kigi_home.join("auth.json").exists(), - "rejected login must not write auth.json", - ); - } - - /// The legacy `oauth2.principal_id` only pre-selects a team; it must not - /// enforce a pin (only `force_login_team_uuid` does), so a different team's - /// token is accepted. - #[test] - fn build_auth_does_not_enforce_legacy_oauth2_principal_id() { - ensure_crypto_provider(); - let cfg = GrokComConfig { - oauth2: Some(crate::auth::OAuth2ProviderConfig { - issuer: "http://localhost:22255".into(), - client_id: "client-id".into(), - scopes: vec!["offline_access".into()], - principal_type: Some("Team".into()), - principal_id: Some("team-required".into()), - referrer: None, - }), - ..GrokComConfig::default() - }; - let temp_dir = tempfile::tempdir().unwrap(); - let kigi_home = temp_dir.path().join(".kigi"); - std::fs::create_dir_all(&kigi_home).unwrap(); - let auth_manager = - Arc::new(AuthManager::new(&kigi_home, cfg).with_proxy_base_url("http://127.0.0.1:9")); - - let auth = tokio::runtime::Runtime::new() - .unwrap() - .block_on(build_auth( - &team_access_token("team-other"), - "http://localhost:22255", - "client-id", - &auth_manager, - )) - .expect("legacy oauth2.principal_id must not enforce a pin"); - assert_eq!( - auth.team_id.as_deref(), - Some("team-other"), - "the token's own team is used; the legacy pre-select id does not gate it", - ); - } - - /// Persistence-seam enforcement via a `force_login_team_uuid` list. - #[test] - fn build_auth_rejects_token_outside_force_login_team_list() { - let cfg = GrokComConfig { - force_login_team_uuid: Some(crate::auth::ForceLoginTeam::AnyOf(vec![ - "team-a".into(), - "team-b".into(), - ])), - ..GrokComConfig::default() - }; - assert_build_auth_rejected( - cfg, - "team-other", - "This deployment requires logging into one of teams: team-a, team-b; \ - your login returned team-other", - ); - } - - // ── complete_device_code_login poll loop ──────────────────────────────── - - /// Spawn a mock `/oauth2/token` server that serves `responses` in order, - /// repeating the last entry. Returns the issuer base URL. - async fn spawn_token_server( - responses: Vec<(u16, serde_json::Value)>, - ) -> (String, tokio::task::JoinHandle<()>) { - use std::sync::atomic::{AtomicUsize, Ordering}; - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let issuer = format!("http://{}", listener.local_addr().unwrap()); - let counter = Arc::new(AtomicUsize::new(0)); - let responses = Arc::new(responses); - let app = axum::Router::new().route( - "/oauth2/token", - axum::routing::post(move || { - let counter = counter.clone(); - let responses = responses.clone(); - async move { - let idx = counter - .fetch_add(1, Ordering::SeqCst) - .min(responses.len() - 1); - let (status, body) = &responses[idx]; - ( - axum::http::StatusCode::from_u16(*status).unwrap(), - axum::Json(body.clone()), - ) - } - }), - ); - let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - (issuer, handle) - } - - fn device_code_for_test(interval: i32, expires_in: i64) -> super::DeviceCode { - super::DeviceCode { - verification_uri: "https://example.test/device".into(), - verification_uri_complete: Some( - "https://example.test/device?user_code=ABCD-EFGH".into(), - ), - user_code: "ABCD-EFGH".into(), - device_code: "dev-code-123".into(), - interval, - expires_in, - } - } - - // Real time (not `start_paused`: the shared client's 30s connect_timeout - // fires under auto-advance). Deadline-expiry isn't tested — the deadline is - // floored at 10 min (MIN_DEVICE_CODE_EXPIRY_FALLBACK_SECS). - async fn run_poll( - responses: Vec<(u16, serde_json::Value)>, - ) -> anyhow::Result<(super::GrokAuth, bool)> { - let (issuer, server) = spawn_token_server(responses).await; - let temp_dir = tempfile::tempdir().unwrap(); - let auth_manager = auth_manager_with_kigi_home(temp_dir.path(), "http://127.0.0.1:9"); - let device_code = device_code_for_test(1, 900); - let result = super::complete_device_code_login( - &issuer, - "client-id", - device_code, - &auth_manager, - super::ClientSurface::Cli, - ) - .await; - server.abort(); - result - } - - fn success_body() -> serde_json::Value { + fn device_auth_json(code: &str) -> serde_json::Value { serde_json::json!({ - "access_token": "mock-access-token", - "refresh_token": "mock-refresh-token", - "expires_in": 900, - "scope": "openid", + "user_code": "ABCD-1234", + "device_code": code, + "verification_uri": "https://auth.kimi.com/device", + "verification_uri_complete": "https://auth.kimi.com/device?code=ABCD-1234", + "expires_in": 600, + "interval": 0, // floored to 1s by the poll loop }) } + fn token_json(access: &str) -> serde_json::Value { + serde_json::json!({ + "access_token": access, + "refresh_token": "rt-1", + "expires_in": 3600, + "scope": "kimi-code", + "token_type": "bearer", + }) + } + + fn auth_manager(dir: &tempfile::TempDir) -> Arc { + Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())) + } + + /// End-to-end (mock server): authorization → pending → token, persisting + /// via the AuthManager. #[tokio::test] - async fn poll_succeeds_on_first_poll() { - let (auth, is_new) = run_poll(vec![(200, success_body())]) + async fn device_login_persists_token_after_pending() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/device_authorization")) + .respond_with(ResponseTemplate::new(200).set_body_json(device_auth_json("dev-1"))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with( + ResponseTemplate::new(400) + .set_body_json(serde_json::json!({ "error": "authorization_pending" })), + ) + .up_to_n_times(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .and(body_string_contains("device_code=dev-1")) + .respond_with(ResponseTemplate::new(200).set_body_json(token_json("at-done"))) + .mount(&server) + .await; + + let dir = tempfile::tempdir().unwrap(); + let mgr = auth_manager(&dir); + let mut channels = None; + let (auth, is_new) = run_device_code_login_channels(&server.uri(), &mgr, &mut channels) .await - .expect("should resolve to a token"); - assert_eq!(auth.key, "mock-access-token"); + .unwrap(); assert!(is_new); + assert_eq!(auth.key, "at-done"); + assert_eq!( + mgr.current_or_expired().map(|a| a.key), + Some("at-done".into()), + "login must land in the manager cache" + ); + assert!( + dir.path().join("auth.json").exists(), + "login must persist to the fallback file store" + ); } + /// `expired_token` during polling restarts the whole device + /// authorization (fresh device code), then completes. #[tokio::test] - async fn poll_succeeds_after_pending() { - let (auth, _) = run_poll(vec![ - (400, serde_json::json!({ "error": "authorization_pending" })), - (200, success_body()), - ]) - .await - .expect("should resolve to a token after pending"); - assert_eq!(auth.key, "mock-access-token"); - } + async fn expired_token_restarts_device_authorization() { + let server = MockServer::start().await; + // First authorization issues dev-1; second issues dev-2. + Mock::given(method("POST")) + .and(path("/api/oauth/device_authorization")) + .respond_with(ResponseTemplate::new(200).set_body_json(device_auth_json("dev-1"))) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/api/oauth/device_authorization")) + .respond_with(ResponseTemplate::new(200).set_body_json(device_auth_json("dev-2"))) + .expect(1) + .mount(&server) + .await; + // dev-1 polls expire; dev-2 polls succeed. + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .and(body_string_contains("device_code=dev-1")) + .respond_with( + ResponseTemplate::new(400) + .set_body_json(serde_json::json!({ "error": "expired_token" })), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .and(body_string_contains("device_code=dev-2")) + .respond_with(ResponseTemplate::new(200).set_body_json(token_json("at-restarted"))) + .expect(1) + .mount(&server) + .await; - #[tokio::test] - async fn poll_handles_slow_down_then_succeeds() { - // slow_down must be tolerated (interval bumped) without erroring. - let (auth, _) = run_poll(vec![ - (400, serde_json::json!({ "error": "slow_down" })), - (200, success_body()), - ]) - .await - .expect("slow_down should be retried, not fatal"); - assert_eq!(auth.key, "mock-access-token"); - } - - #[tokio::test] - async fn poll_maps_access_denied_to_error() { - let err = run_poll(vec![(400, serde_json::json!({ "error": "access_denied" }))]) + let dir = tempfile::tempdir().unwrap(); + let mgr = auth_manager(&dir); + let mut channels = None; + let (auth, _) = run_device_code_login_channels(&server.uri(), &mgr, &mut channels) .await - .expect_err("access_denied must be an error"); - assert!(err.to_string().contains("denied"), "got: {err}"); + .unwrap(); + assert_eq!(auth.key, "at-restarted"); } + /// `slow_down` is wait-and-continue (interval bumped), never fatal. + /// Unknown-error continuation is covered at the wire level by + /// `kimi_oauth::tests::poll_maps_pending_and_unknown_errors_to_pending`. #[tokio::test] - async fn poll_maps_expired_token_to_error() { - let err = run_poll(vec![(400, serde_json::json!({ "error": "expired_token" }))]) + async fn slow_down_keeps_polling_then_succeeds() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/device_authorization")) + .respond_with(ResponseTemplate::new(200).set_body_json(device_auth_json("dev-1"))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with( + ResponseTemplate::new(400) + .set_body_json(serde_json::json!({ "error": "slow_down" })), + ) + .up_to_n_times(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(token_json("at-patient"))) + .expect(1) + .mount(&server) + .await; + + let dir = tempfile::tempdir().unwrap(); + let mgr = auth_manager(&dir); + let mut channels = None; + let started = std::time::Instant::now(); + let (auth, _) = run_device_code_login_channels(&server.uri(), &mgr, &mut channels) .await - .expect_err("expired_token must be an error"); - assert!(err.to_string().contains("expired"), "got: {err}"); + .unwrap(); + assert_eq!(auth.key, "at-patient"); + assert!( + started.elapsed() >= std::time::Duration::from_secs(6), + "slow_down must bump the poll interval by {SLOW_DOWN_INCREMENT_SECS}s" + ); + } + + /// A 5xx from the token endpoint is a hard error (kimi-cli parity). + #[tokio::test] + async fn server_error_during_poll_fails_login() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/device_authorization")) + .respond_with(ResponseTemplate::new(200).set_body_json(device_auth_json("dev-1"))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with(ResponseTemplate::new(500)) + .mount(&server) + .await; + + let dir = tempfile::tempdir().unwrap(); + let mgr = auth_manager(&dir); + let mut channels = None; + let err = run_device_code_login_channels(&server.uri(), &mgr, &mut channels) + .await + .unwrap_err(); + assert!(err.to_string().contains("server error"), "{err}"); + assert!( + mgr.current_or_expired().is_none(), + "failed login must not persist credentials" + ); } } diff --git a/crates/codegen/kigi-shell/src/auth/error.rs b/crates/codegen/kigi-shell/src/auth/error.rs index efd9ace..291375f 100644 --- a/crates/codegen/kigi-shell/src/auth/error.rs +++ b/crates/codegen/kigi-shell/src/auth/error.rs @@ -3,30 +3,21 @@ use thiserror::Error; #[derive(Debug, Error)] #[non_exhaustive] pub enum AuthError { - #[error("Not logged in. Run `grok login`.")] + #[error("Not logged in. Run `kigi login`.")] NotLoggedIn, /// Token expired and no refresh authority available. - #[error("Token expired. Run `grok login` to re-authenticate.")] + #[error("Token expired. Run `kigi login` to re-authenticate.")] TokenExpiredNoRefresh, /// Server rejected the token (401) with no recovery path. - #[error("Authentication rejected by server. Run `grok login` to re-authenticate.")] + #[error("Authentication rejected by server. Run `kigi login` to re-authenticate.")] ServerRejectedNoRecovery, /// All recovery strategies exhausted. #[error("Auth recovery exhausted; re-authentication required.")] RecoveryExhausted, - /// A session's team principal violates the `force_login_team_uuid` pin. - /// `message` states which team is required vs. returned. - #[error("{message} Run `grok login` to sign in with the required team.")] - PinnedTeamMismatch { message: String }, - - /// Cached API-key session rejected because API-key auth is disabled. - #[error("API-key auth is disabled by your administrator. Run `grok login` to authenticate.")] - ApiKeyAuthDisabled, - /// Outcome of a refresh-authority attempt. Recoverability (and, for /// permanent failures, the reason) lives in [`RefreshTokenError`]. #[error(transparent)] @@ -38,7 +29,7 @@ pub enum AuthError { /// caller must make, so a future third state should break consumers loudly. #[derive(Debug, Error)] pub enum RefreshTokenError { - /// The credential is dead; the user must re-authenticate. + /// The credential was rejected; the tombstone cooldown gates re-attempts. #[error(transparent)] Permanent(#[from] RefreshTokenFailedError), /// Network / 5xx / unknown blip; safe to retry later. Carries the cause. @@ -48,12 +39,10 @@ pub enum RefreshTokenError { /// A retryable refresh failure, wrapping its cause. No public `From`: /// construct only via [`AuthError::transient`] / -/// [`AuthError::transient_source`], so a stray `?` on some error can't silently -/// classify a permanent failure as retryable (mirrors the dedicated -/// [`RefreshTokenFailedError`] on the permanent arm). Display frames the cause -/// as an auth-refresh failure so internal messages (lock timeout, sleep defer) -/// don't surface bare; the permanent arm derives its copy from -/// [`RefreshTokenFailedReason::user_message`] and is not prefixed. +/// [`AuthError::transient_source`], so a stray `?` on some error can't +/// silently classify a permanent failure as retryable. Display frames the +/// cause as an auth-refresh failure so internal messages (lock timeout, +/// sleep defer) don't surface bare. #[derive(Debug, Error)] #[error("auth refresh failed: {0}")] pub struct RefreshTransientError(#[source] Box); @@ -74,45 +63,31 @@ impl From for RefreshTokenFailedError { } } -/// Why a token refresh terminally failed, grounded in the OAuth2 error codes -/// our IdP actually emits. +/// Why a token refresh terminally failed. Both reasons carry the same +/// tombstone semantics (PRD F1): a 300s cooldown scoped to the rejected +/// refresh token, auto-cleared when the persisted refresh token differs +/// (another process rotated) or a fresh login lands. #[derive(Debug, Clone, Copy, PartialEq, Eq)] #[non_exhaustive] pub enum RefreshTokenFailedReason { - /// `invalid_grant` — the refresh token is no longer valid (expired, reused, - /// or revoked; the IdP does not distinguish these). + /// The OAuth host answered 401/403 — the refresh token is no longer + /// valid (expired, reused, or revoked). RefreshTokenRejected, - /// `invalid_client` — the client/app credential was rejected. - ClientRejected, - /// Escalation from repeated transient failures (OIDC) or a single - /// external-binary failure. Never a raw IdP code: an unrecognized terminal - /// code is classified transient, not `Other` (see `classify_terminal`). + /// Non-retryable terminal failure that isn't an explicit rejection + /// (malformed payload, unexpected 4xx). Other, } impl RefreshTokenFailedReason { - /// Sticky until the credential changes (never ages out): a revoked refresh - /// token never self-heals, whereas client rotation / transient escalation - /// recover, so those age out past the TTL. - pub(crate) fn is_sticky(self) -> bool { - match self { - Self::RefreshTokenRejected => true, - Self::ClientRejected | Self::Other => false, - } - } - - /// User-facing copy for a terminal refresh failure; the raw IdP code stays - /// in logs. + /// User-facing copy for a terminal refresh failure; the raw wire detail + /// stays in logs. pub(crate) fn user_message(self) -> &'static str { match self { Self::RefreshTokenRejected => { - "Your session has expired. Run `grok login` to sign in again." - } - Self::ClientRejected => { - "Authentication is temporarily unavailable. Run `grok login` if this persists." + "Your session has expired. Run `kigi login` to sign in again." } Self::Other => { - "Authentication could not be refreshed. Run `grok login` to sign in again." + "Authentication could not be refreshed. Run `kigi login` to sign in again." } } } diff --git a/crates/codegen/kigi-shell/src/auth/external_auth.rs b/crates/codegen/kigi-shell/src/auth/external_auth.rs deleted file mode 100644 index b23348e..0000000 --- a/crates/codegen/kigi-shell/src/auth/external_auth.rs +++ /dev/null @@ -1,283 +0,0 @@ -use crate::auth::{AuthMode, GrokAuth}; - -#[derive(serde::Deserialize)] -pub(crate) struct ExternalAuthOutput { - pub access_token: String, - #[serde(default)] - pub refresh_token: Option, - #[serde(default)] - pub expires_in: Option, - /// Token issuer. An xAI issuer marks the credential as first-party; - /// see [`GrokAuth::is_xai_auth`]. - #[serde(default)] - pub issuer: Option, -} - -/// Parse process output (stdout) into a `GrokAuth`. Accepts bare token or JSON. -pub(crate) fn parse_output(output: &std::process::Output) -> anyhow::Result { - if !output.status.success() { - anyhow::bail!("exited with {}", output.status); - } - - let stdout = String::from_utf8_lossy(&output.stdout).trim().to_owned(); - if stdout.is_empty() { - anyhow::bail!("produced no output on stdout"); - } - - let (token, refresh_token, expires_at, issuer) = - if let Ok(parsed) = serde_json::from_str::(&stdout) { - tracing::debug!( - has_refresh_token = parsed.refresh_token.is_some(), - expires_in = ?parsed.expires_in, - issuer = ?parsed.issuer, - "auth: parsed external provider output as JSON" - ); - let expires_at = parsed - .expires_in - .map(|secs| chrono::Utc::now() + chrono::Duration::seconds(secs as i64)); - let issuer = parsed - .issuer - .map(|i| i.trim().to_owned()) - .filter(|i| !i.is_empty()); - ( - parsed.access_token, - parsed.refresh_token, - expires_at, - issuer, - ) - } else { - tracing::debug!( - stdout_len = stdout.len(), - "auth: treating output as bare token" - ); - (stdout, None, None, None) - }; - - Ok(GrokAuth { - key: token, - auth_mode: AuthMode::External, - create_time: chrono::Utc::now(), - user_id: String::new(), - email: None, - first_name: None, - last_name: None, - profile_image_asset_id: None, - principal_type: None, - principal_id: None, - team_id: None, - team_name: None, - team_role: None, - organization_id: None, - organization_name: None, - organization_role: None, - user_blocked_reason: None, - team_blocked_reasons: vec![], - coding_data_retention_opt_out: false, - has_grok_code_access: None, - refresh_token, - expires_at, - oidc_issuer: issuer, - oidc_client_id: None, - }) -} - -/// Sync version for mid-session refresh. 5s timeout for refresh, 60s for initial. -pub(crate) fn run_external_auth_sync(command: &str, is_refresh: bool) -> Option { - use std::process::{Command, Stdio}; - - let timeout_secs = if is_refresh { 5 } else { 60 }; - - tracing::info!(cmd = %command, is_refresh, timeout_secs, "auth: running external auth provider (sync)"); - - let mut cmd = Command::new("sh"); - cmd.args(["-c", command]) - .stdin(Stdio::null()) - .stdout(Stdio::piped()) - // Pipe stderr — inherit would corrupt the TUI alternate screen. - .stderr(Stdio::piped()); - if is_refresh { - cmd.env("KIGI_AUTH_EXPIRED", "1"); - } - kigi_tools::util::detach_std_command(&mut cmd); - cmd.envs(kigi_tools::util::pager_env()); - let mut child = cmd.spawn() - .map_err(|e| { - tracing::warn!(error = %e, cmd = %command, "auth: failed to start external auth provider"); - e - }) - .ok()?; - - let timeout = std::time::Duration::from_secs(timeout_secs); - let start = std::time::Instant::now(); - loop { - match child.try_wait() { - Ok(Some(_status)) => break, - Ok(None) => { - if start.elapsed() > timeout { - tracing::warn!( - cmd = %command, - timeout_secs, - "auth: external auth provider timed out (likely needs interactive auth), killing" - ); - let _ = child.kill(); - let _ = child.wait(); - return None; - } - std::thread::sleep(std::time::Duration::from_millis(100)); - } - Err(e) => { - tracing::warn!(error = %e, "auth: error waiting for external auth provider"); - return None; - } - } - } - - let output = child - .wait_with_output() - .map_err(|e| { - tracing::warn!(error = %e, "auth: failed to read external auth provider output"); - e - }) - .ok()?; - - match parse_output(&output) { - Ok(auth) => { - tracing::info!("auth: external auth provider returned fresh token"); - Some(auth) - } - Err(e) => { - tracing::warn!(error = %e, "auth: external auth provider failed"); - None - } - } -} - -/// Run external auth provider, carrying forward `/user`-derived fields from previous auth. -pub(crate) fn refresh_with_command(command: &str, prev_auth: &GrokAuth) -> Option { - let mut auth = run_external_auth_sync(command, true)?; - auth.carry_user_profile_from(prev_auth); - Some(auth) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn parse_output_nonzero_exit_is_err() { - let output = std::process::Output { - status: std::process::Command::new("false").status().unwrap(), - stdout: b"token".to_vec(), - stderr: vec![], - }; - assert!(parse_output(&output).is_err()); - } - - #[test] - fn parse_output_empty_stdout_is_err() { - let output = std::process::Output { - status: std::process::Command::new("true").status().unwrap(), - stdout: b" \n".to_vec(), - stderr: vec![], - }; - assert!(parse_output(&output).is_err()); - } - - #[test] - fn parse_output_issuer_claim_enables_xai_auth() { - let ok = |stdout: &str| std::process::Output { - status: std::process::Command::new("true").status().unwrap(), - stdout: stdout.as_bytes().to_vec(), - stderr: vec![], - }; - - // x.ai issuer claim → first-party session (relay-eligible). - let auth = parse_output(&ok( - r#"{"access_token":"t","expires_in":900,"issuer":"https://auth.x.ai"}"#, - )) - .unwrap(); - assert_eq!(auth.oidc_issuer.as_deref(), Some("https://auth.x.ai")); - assert!(auth.is_xai_auth()); - - // Non-x.ai issuer is stored but stays third-party. - let auth = parse_output(&ok( - r#"{"access_token":"t","issuer":"https://idp.acme.example"}"#, - )) - .unwrap(); - assert_eq!( - auth.oidc_issuer.as_deref(), - Some("https://idp.acme.example") - ); - assert!(!auth.is_xai_auth()); - - // Missing / empty / whitespace issuer → None. - let auth = parse_output(&ok(r#"{"access_token":"t"}"#)).unwrap(); - assert_eq!(auth.oidc_issuer, None); - assert!(!auth.is_xai_auth()); - let auth = parse_output(&ok(r#"{"access_token":"t","issuer":" "}"#)).unwrap(); - assert_eq!(auth.oidc_issuer, None); - - // Bare-token output never carries an issuer. - let auth = parse_output(&ok("bare-token")).unwrap(); - assert_eq!(auth.oidc_issuer, None); - assert!(!auth.is_xai_auth()); - } - - #[test] - fn parse_output_malformed_json_falls_back_to_bare() { - let output = std::process::Output { - status: std::process::Command::new("true").status().unwrap(), - stdout: b"{not valid json}".to_vec(), - stderr: vec![], - }; - let auth = parse_output(&output).unwrap(); - assert_eq!(auth.key, "{not valid json}"); - } - - #[test] - fn sync_spawn_failure_returns_none() { - assert!(run_external_auth_sync("/nonexistent/binary", false).is_none()); - } - - #[test] - fn sync_sets_grok_auth_expired_env_on_refresh() { - let auth = run_external_auth_sync("echo $KIGI_AUTH_EXPIRED", true).unwrap(); - assert_eq!(auth.key, "1"); - } - - #[test] - fn refresh_carries_zdr_flags_forward() { - let prev = GrokAuth { - user_blocked_reason: Some("BLOCKED_REASON_OTHER".into()), - team_blocked_reasons: vec!["BLOCKED_REASON_NO_LOGS".into()], - coding_data_retention_opt_out: true, - organization_id: Some("org-1".into()), - ..GrokAuth::test_default() - }; - let auth = refresh_with_command("echo fresh-token", &prev).unwrap(); - assert_eq!(auth.key, "fresh-token"); - assert!(auth.is_zdr_team(), "ZDR flag must survive refresh"); - assert!(auth.coding_data_retention_opt_out); - assert_eq!( - auth.user_blocked_reason.as_deref(), - Some("BLOCKED_REASON_OTHER") - ); - assert_eq!(auth.user_id, "test-user", "profile must survive refresh"); - assert_eq!(auth.organization_id.as_deref(), Some("org-1")); - } - - #[test] - fn sync_refresh_interactive_times_out() { - // Binary writes link to stderr then blocks — 5s refresh timeout kills it. - let cmd = r#"echo 'Visit http://example.com/auth' >&2; sleep 20; echo token"#; - let start = std::time::Instant::now(); - let result = run_external_auth_sync(cmd, true); - let elapsed = start.elapsed(); - assert!(result.is_none(), "should timeout and return None"); - assert!( - elapsed.as_secs() < 10, - "refresh should use 5s timeout, not 60s (took {}s)", - elapsed.as_secs() - ); - } -} diff --git a/crates/codegen/kigi-shell/src/auth/flow.rs b/crates/codegen/kigi-shell/src/auth/flow.rs index a723d98..a65e95d 100644 --- a/crates/codegen/kigi-shell/src/auth/flow.rs +++ b/crates/codegen/kigi-shell/src/auth/flow.rs @@ -1,192 +1,27 @@ -use std::cell::RefCell; -use std::rc::Rc; +//! Auth-flow orchestration: cached credentials → silent refresh → the Kimi +//! Code device-code login (the only interactive login). + use std::sync::Arc; -use tokio::io::AsyncBufReadExt as _; use tokio::sync::{mpsc, oneshot}; -use crate::auth::config::LEGACY_AUTH_SCOPE; -use crate::auth::{AuthManager, GrokAuth, GrokComConfig, parse_output}; +use crate::auth::{AuthManager, KimiAuth, KimiCodeConfig}; use crate::util::kigi_home; -pub type StderrCallback = Box; - -/// Reject a cached credential for reuse if it lacks `oidc_issuer`, has a -/// mismatched issuer, or its team principal violates the `force_login_team_uuid` -/// pin — so interactive login starts fresh instead of reusing a stale/wrong-team -/// session. -fn is_cached_credential_compatible(auth: &GrokAuth, grok_com_config: &GrokComConfig) -> bool { - let expected_issuer = grok_com_config - .oidc - .as_ref() - .map(|c| c.issuer.as_str()) - .or_else(|| grok_com_config.oauth2.as_ref().map(|c| c.issuer.as_str())); - let issuer_compatible = match (auth.oidc_issuer.as_deref(), expected_issuer) { - (Some(actual), Some(expected)) => actual == expected, - (None, Some(_)) => false, - _ => true, - }; - if !issuer_compatible { - return false; - } - if let Some(policy) = crate::auth::oidc::login_principal_policy(grok_com_config) { - let actual = crate::auth::oidc::peek_access_token_principal_id(&auth.key); - if crate::auth::oidc::enforce_login_principal(Some(&policy), actual.as_deref()).is_err() { - return false; - } - } - true -} - -/// CLI-flag override for the interactive login transport. -/// -/// `--oauth` forces the loopback-callback flow; `--device-auth` forces the -/// device flow. `None` falls through to env / config / default. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] -pub enum LoginTransportOverride { - /// No CLI override — resolve from env / config / default. - #[default] - None, - /// `--oauth`: force the loopback-callback flow. - ForceLoopback, - /// `--device-auth`: force the RFC 8628 device flow. - ForceDevice, - /// Transport already resolved (and logged) upstream; the inner flow honors - /// the carried value (`true` = device, `false` = loopback) without - /// re-resolving, so it's never re-logged or mis-attributed to `cli`. - Preresolved(bool), -} - -impl LoginTransportOverride { - /// Resolve from the `--oauth` / `--device-auth` flags. `--oauth` wins if - /// both are somehow set. Single source of truth for both the CLI - /// (`run_cli_login`) and ACP (`AuthRequestMeta`) entry points. - pub fn from_flags(force_loopback: bool, force_device: bool) -> Self { - if force_loopback { - Self::ForceLoopback - } else if force_device { - Self::ForceDevice - } else { - Self::None - } - } - - /// Map to a `BoolFlag` CLI value (`Some(true)` = device, `Some(false)` = - /// loopback, `None` = no override). - fn as_cli_bool(self) -> Option { - match self { - Self::None => None, - Self::ForceLoopback => Some(false), - Self::ForceDevice => Some(true), - // Not a CLI decision — must never be reported as the `cli` tier. - Self::Preresolved(_) => None, - } - } -} - -/// `[auth] login_device_flow` from a config snapshot (shared with the proxy-URL read). -fn config_login_device_flow(effective: Option<&toml::Value>) -> Option { - effective.and_then(|cfg| cfg.get("auth")?.get("login_device_flow")?.as_bool()) -} - -/// Device-flow precedence: CLI > env > config > remote feature flag > loopback. -/// Returns the deciding tier so the caller can log which one chose the transport. -fn resolve_device_flow( - login_override: LoginTransportOverride, - config: Option, - remote: Option, -) -> crate::agent::config::Resolved { - crate::agent::config::BoolFlag::env("KIGI_LOGIN_DEVICE_FLOW") - .cli(login_override.as_cli_bool()) - .config(config) - .feature_flag(remote) - .default(false) - .resolve() -} - -/// Whether `run_cli_login` should use the device flow for `config`: only the -/// xAI OAuth2 provider supports it. Enterprise OIDC (`oidc=Some`) always uses -/// the loopback flow, mirroring `run_auth_flow_inner`'s precedence. -async fn cli_should_use_device( - config: &GrokComConfig, - login_override: LoginTransportOverride, -) -> bool { - !crate::auth::oidc::is_configured(config) && should_use_device_flow(login_override).await -} - -/// Whether interactive xAI OAuth2 login uses the RFC 8628 device flow (vs loopback). -/// -/// Precedence: CLI (`--oauth`/`--device-auth`) > `KIGI_LOGIN_DEVICE_FLOW` env > -/// `[auth] login_device_flow` config > `grok_build_login_device_flow` remote feature flag > loopback. -async fn should_use_device_flow(login_override: LoginTransportOverride) -> bool { - // Already resolved (and logged) upstream — honor it without re-resolving or - // emitting a second transport log. - if let LoginTransportOverride::Preresolved(use_device) = login_override { - return use_device; - } - let resolved = if login_override.as_cli_bool().is_some() { - // CLI flag wins outright, so skip the config load and the remote settings fetch. - resolve_device_flow(login_override, None, None) - } else { - // Read once to gate the fetch; resolve_device_flow reads it again for the decision. - let env = crate::agent::config::env_bool("KIGI_LOGIN_DEVICE_FLOW"); - // One config snapshot feeds both the `[auth]` tier and the proxy URL. - let effective = crate::config::load_effective_config().ok(); - let config = config_login_device_flow(effective.as_ref()); - // Only hit remote settings when env/config haven't already pinned the transport. - let remote = if env.is_none() && config.is_none() { - let proxy_url = effective - .as_ref() - .map(crate::agent::config::EndpointsConfig::from_config_value) - .unwrap_or_default() - .proxy_url(); - // Bound the whole fetch — including the one-time agent_id lookup — so a - // slow/hung agent_id or proxy can never stall login; time out to loopback. - tokio::time::timeout( - std::time::Duration::from_secs(2), - crate::remote::fetch_login_device_flow(&proxy_url), - ) - .await - .ok() - .flatten() - } else { - None - }; - resolve_device_flow(login_override, config, remote) - }; - tracing::info!( - transport = if resolved.value { "device" } else { "loopback" }, - source = %resolved.source, - "login: resolved interactive transport", - ); - resolved.value -} - -/// How login presents itself; surfaced to the TUI via `x.ai/auth/get_url`. +/// How login presents itself; surfaced to the TUI via the auth URL event. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum AuthUrlMode { - /// Loopback-callback flow — TUI shows a copyable URL + paste box. - Loopback, - /// External auth provider opened its own browser — TUI shows a waiting status. - Command, - /// RFC 8628 device flow — TUI shows the device code + copyable URL, no paste box. + /// Device flow — TUI shows the verification URL (user code pre-filled). Device, } impl AuthUrlMode { - /// Wire string for the `x.ai/auth/get_url` ACP response. + /// Wire string for the ACP auth-url response. pub fn as_wire_str(self) -> &'static str { match self { - Self::Loopback => "loopback", - Self::Command => "command", Self::Device => "device", } } - - /// Back-compat flag for older clients that only read `external_provider`. - pub fn is_external_provider(self) -> bool { - matches!(self, Self::Command) - } } /// Auth URL pushed from the auth flow to the TUI. @@ -196,313 +31,75 @@ pub struct AuthUrlInfo { } /// Channels for interactive login between the auth flow and the TUI/extension. +/// `code_rx` is unused by the device flow (the verification URL pre-fills the +/// user code) but kept so the ACP wiring stays uniform. pub struct AuthChannels { pub url_tx: Option>, pub code_rx: mpsc::Receiver, } -async fn run_external_auth_provider( - command: &str, - auth_manager: &Arc, - is_refresh: bool, - on_stderr: Option, -) -> anyhow::Result<(GrokAuth, bool)> { - let inherit_stderr = on_stderr.is_none(); - tracing::info!( - cmd = %command, - is_refresh, - inherit_stderr, - "auth: running external auth provider" - ); - - let mut cmd = tokio::process::Command::new("sh"); - cmd.args(["-c", command]) - .stdin(std::process::Stdio::null()) - .stdout(std::process::Stdio::piped()) - .kill_on_drop(true); - - // TUI: pipe stderr and forward via callback — inherit would corrupt the - // alternate screen. CLI / headless: inherit so URLs and progress appear in - // real time; piping without a reader hides output and can deadlock the child. - if inherit_stderr { - cmd.stderr(std::process::Stdio::inherit()); - } else { - cmd.stderr(std::process::Stdio::piped()); - } - - if is_refresh { - cmd.env("KIGI_AUTH_EXPIRED", "1"); - } - - kigi_tools::util::detach_command(&mut cmd); - cmd.envs(kigi_tools::util::pager_env()); - - let mut child = cmd - .spawn() - .map_err(|e| anyhow::anyhow!("failed to start auth provider `{command}`: {e}"))?; - - let stderr_task = if let Some(cb) = on_stderr { - let stderr = child.stderr.take().expect("stderr was set to piped"); - Some(tokio::task::spawn_local(async move { - let mut reader = tokio::io::BufReader::new(stderr); - let mut line = String::new(); - loop { - line.clear(); - match reader.read_line(&mut line).await { - Ok(0) => break, - Ok(_) => { - let trimmed = line.trim_end(); - tracing::debug!(line = trimmed, "auth: provider stderr"); - cb(trimmed); - } - Err(e) => { - tracing::warn!(error = %e, "auth: error reading provider stderr"); - break; - } - } - } - })) - } else { - None - }; - - let output = tokio::time::timeout( - std::time::Duration::from_secs(300), - child.wait_with_output(), - ) - .await - .map_err(|_| anyhow::anyhow!("external auth provider `{command}` timed out after 300s"))? - .map_err(|e| anyhow::anyhow!("external auth provider `{command}` IO error: {e}"))?; - - if let Some(task) = stderr_task { - let _ = task.await; - } - - let mut auth = parse_output(&output) - .map_err(|e| anyhow::anyhow!("external auth provider `{command}`: {e}"))?; - - // Verify the team pin before any persist (parity with the OIDC / device-code - // completion paths). A mismatch fails the login and writes nothing. - let principal_policy = - crate::auth::oidc::login_principal_policy(auth_manager.grok_com_config()); - crate::auth::oidc::enforce_login_principal( - principal_policy.as_ref(), - crate::auth::oidc::peek_access_token_principal_id(&auth.key).as_deref(), - )?; - - // Token output has no profile; carry it forward, or fetch it when reauth cleared prev. - match (is_refresh, auth_manager.current_or_expired()) { - (true, Some(prev)) => auth.carry_user_profile_from(&prev), - _ => auth_manager.enrich_auth_inline(&mut auth).await, - } - - let auth = auth_manager - .update(auth) - .await - .map_err(|e| anyhow::anyhow!("failed to save external auth credentials: {e}"))?; - - tracing::info!( - user_id = %auth.user_id, - email = ?auth.email, - "auth: external provider login complete" - ); - - Ok((auth, true)) -} - -/// GUI auth: bridges external provider stderr to `url_tx`, pipes code submission via `code_rx`. +/// GUI auth entry point (ACP `login` handler). pub async fn run_auth_flow_with_stderr_bridge( auth_manager: &Arc, - grok_com_config: &GrokComConfig, + kimi_code_config: &KimiCodeConfig, channels: AuthChannels, reauth: bool, force_interactive: bool, - login_override: LoginTransportOverride, -) -> anyhow::Result<(GrokAuth, bool)> { - let url_tx = Rc::new(RefCell::new(channels.url_tx)); - let stderr_lines: Rc>> = Rc::new(RefCell::new(Vec::new())); - - let writer = stderr_lines.clone(); - let on_stderr: StderrCallback = Box::new(move |line: &str| { - writer.borrow_mut().push(line.to_owned()); - }); - - let reader = stderr_lines.clone(); - let url_tx_bridge = url_tx.clone(); - let bridge = async move { - loop { - tokio::task::yield_now().await; - let content = { - let lines = reader.borrow(); - if lines.is_empty() { - None - } else { - Some(lines.join("\n")) - } - }; - if let Some(joined) = content - && let Some(tx) = url_tx_bridge.borrow_mut().take() - { - // The external binary may print preamble text alongside the - // URL (e.g. "Visit the following link to sign in: https://…"). - // Extract just the first https:// URL so the TUI displays a - // clean, clickable link. - let url = joined - .split_whitespace() - .find(|w| w.starts_with("https://")) - .map(|u| u.to_owned()) - .unwrap_or(joined); - let _ = tx.send(AuthUrlInfo { - url, - mode: AuthUrlMode::Command, - }); - } - tokio::time::sleep(std::time::Duration::from_millis(50)).await; - } - }; - - if force_interactive { - let auth = run_auth_flow_interactive( - auth_manager, - grok_com_config, - Some(on_stderr), - Some(url_tx), - Some(channels.code_rx), - login_override, - ); - tokio::select! { - r = auth => r, - _ = bridge => { - tracing::error!("auth stderr bridge exited unexpectedly during interactive login"); - Err(anyhow::anyhow!("Login failed. Please try again.")) - }, - } - } else { - let auth = run_auth_flow( - auth_manager, - grok_com_config, - reauth, - Some(on_stderr), - Some(url_tx), - Some(channels.code_rx), - login_override, - ); - tokio::select! { - r = auth => r, - _ = bridge => { - tracing::error!("auth stderr bridge exited unexpectedly during login"); - Err(anyhow::anyhow!("Login failed. Please try again.")) - }, - } - } +) -> anyhow::Result<(KimiAuth, bool)> { + run_auth_flow_inner( + auth_manager, + kimi_code_config, + reauth, + force_interactive, + Some(channels), + ) + .await } -/// Full auth chain: cache → refresh → external provider → interactive (OIDC/OAuth2/legacy). -/// When `url_tx` and `code_rx` are `None`, falls back to stderr/stdin (CLI mode). +/// Full auth chain: cache → silent refresh → device-code login. +/// When `channels` is `None`, login output goes to stderr (CLI mode). pub async fn run_auth_flow( auth_manager: &Arc, - grok_com_config: &GrokComConfig, + kimi_code_config: &KimiCodeConfig, reauth: bool, - on_stderr: Option, - url_tx: Option>>>>, - code_rx: Option>, - login_override: LoginTransportOverride, -) -> anyhow::Result<(GrokAuth, bool)> { - run_auth_flow_inner( - auth_manager, - grok_com_config, - reauth, - false, - on_stderr, - url_tx, - code_rx, - login_override, - ) - .await -} - -/// Like [`run_auth_flow`] but with `force_interactive`: skip cached -/// credentials without clearing them. Used by `/login` for mid-session -/// re-auth where abandoning the flow must not disrupt the session. -pub async fn run_auth_flow_interactive( - auth_manager: &Arc, - grok_com_config: &GrokComConfig, - on_stderr: Option, - url_tx: Option>>>>, - code_rx: Option>, - login_override: LoginTransportOverride, -) -> anyhow::Result<(GrokAuth, bool)> { - run_auth_flow_inner( - auth_manager, - grok_com_config, - false, - true, - on_stderr, - url_tx, - code_rx, - login_override, - ) - .await + channels: Option, +) -> anyhow::Result<(KimiAuth, bool)> { + run_auth_flow_inner(auth_manager, kimi_code_config, reauth, false, channels).await } async fn run_auth_flow_inner( auth_manager: &Arc, - grok_com_config: &GrokComConfig, + _kimi_code_config: &KimiCodeConfig, reauth: bool, force_interactive: bool, - on_stderr: Option, - url_tx: Option>>>>, - code_rx: Option>, - login_override: LoginTransportOverride, -) -> anyhow::Result<(GrokAuth, bool)> { - tracing::info!( - has_oidc = grok_com_config.oidc.is_some(), - has_oauth2 = grok_com_config.oauth2.is_some(), - has_external_auth = grok_com_config.auth_provider_command.is_some(), - reauth, - "auth: starting auth flow" - ); + channels: Option, +) -> anyhow::Result<(KimiAuth, bool)> { + tracing::info!(reauth, force_interactive, "auth: starting auth flow"); if reauth { auth_manager.clear()?; - // Also remove the legacy accounts.x.ai scope so stale tokens - // don't linger alongside the fresh OIDC credential. - let _ = auth_manager.remove_scope(LEGACY_AUTH_SCOPE); } if !force_interactive && let Some(auth) = auth_manager.current() { - if is_cached_credential_compatible(&auth, grok_com_config) { - tracing::info!(auth_mode = ?auth.auth_mode, "auth: using cached credentials"); - kigi_log::unified_log::info( - "auth: using cached credentials", - None, - Some(serde_json::json!({ "auth_mode": format!("{:?}", auth.auth_mode) })), - ); - return Ok((auth, false)); - } - tracing::info!( - auth_mode = ?auth.auth_mode, - "auth: cached credential incompatible with requested flow, proceeding to interactive login" + tracing::info!(auth_mode = ?auth.auth_mode, "auth: using cached credentials"); + kigi_log::unified_log::info( + "auth: using cached credentials", + None, + Some(serde_json::json!({ "auth_mode": format!("{:?}", auth.auth_mode) })), ); - // Remove the stale legacy credential from disk so it doesn't - // linger alongside the new OIDC entry after re-authentication. - if auth.auth_mode == super::AuthMode::WebLogin - && let Err(e) = auth_manager.remove_scope(LEGACY_AUTH_SCOPE) - { - tracing::warn!(error = ?e, "auth: failed to remove legacy scope entry (non-fatal)"); - } + return Ok((auth, false)); } if !force_interactive && !reauth && auth_manager.is_expired() { - // Acquire the cross-process file lock so we don't race with - // OidcRefresher instances in sibling processes. Without this, - // two processes can send the same refresh_token simultaneously, - // triggering IdP refresh-token-family revocation (reuse detection). + // Acquire the cross-process file lock so we don't race a refresher in + // a sibling process. Without this, two processes can send the same + // refresh_token simultaneously and trip server-side reuse detection. let _file_lock = auth_manager .try_lock_auth_file_async(crate::auth::manager::AUTH_LOCK_TIMEOUT) .await; - // Read disk first — another process may have already refreshed. + // Read the persisted store first — another process may have already + // refreshed. let disk_auth = auth_manager.read_disk_auth(); let disk_expired = disk_auth.as_ref().is_some_and(crate::auth::is_expired); kigi_log::unified_log::info( @@ -514,23 +111,20 @@ async fn run_auth_flow_inner( "disk_expired": disk_expired, })), ); - if disk_auth.as_ref().is_some_and(|d| { - !crate::auth::is_expired(d) && is_cached_credential_compatible(d, grok_com_config) - }) { + if let Some(d) = disk_auth.clone().filter(|d| !crate::auth::is_expired(d)) { kigi_log::unified_log::info("auth run_auth_flow using valid disk token", None, None); - let d = disk_auth.unwrap(); - let ret = d.clone(); - auth_manager.hot_swap(d); - return Ok((ret, false)); + auth_manager.hot_swap(d.clone()); + return Ok((d, false)); } - // Disk token not usable. Try the full auth() dispatcher which - // handles OIDC refresh, external binary, disk re-read — all - // through refresh_chain (single mutation point). + // Persisted token not usable. Try the full auth() dispatcher which + // handles refresh and sibling adoption — all through refresh_chain + // (single mutation point). match auth_manager.auth().await { Ok(fresh) => return Ok((fresh, false)), Err(e) => { - // Defer to consumer-level refresh if disk has a refresh_token. + // Defer to consumer-level refresh if the store has a + // refresh_token and the failure was transient. if let Some(d) = disk_auth.filter(|d| { matches!( &e, @@ -561,122 +155,30 @@ async fn run_auth_flow_inner( } } - if let Some(ref cmd) = grok_com_config.auth_provider_command { - let is_refresh = reauth || auth_manager.is_expired(); - match run_external_auth_provider(cmd, auth_manager, is_refresh, on_stderr).await { - Ok(result) => return Ok(result), - Err(e) => { - tracing::warn!( - error = %e, - "auth: external auth provider failed, falling through to interactive login" - ); - eprintln!("Signing in with browser instead..."); - } - } - } - - // Devbox auto-migration: before interactive login (which requires a - // browser and won't work on headless devboxes), try minting OIDC - // credentials via the remote devbox login helper. - // preferred_method=api_key: never auto-mint OIDC (fail-closed). Explicit - // `grok login --devbox` uses run_devbox_login and is not gated here. - if !grok_com_config.blocks_automatic_oidc() - && crate::auth::devbox_login::is_devbox_environment() - { - tracing::info!("auth: devbox detected, attempting devbox login before interactive flow"); - match crate::auth::devbox_login::mint_devbox_auth(auth_manager).await { - Ok(new_auth) => match auth_manager.save_without_enrichment(new_auth).await { - Ok(auth) => { - let _ = auth_manager.remove_scope(LEGACY_AUTH_SCOPE); - kigi_log::unified_log::info( - "auth: devbox migration in auth flow succeeded", - None, - Some(serde_json::json!({ - "user_id": auth.user_id, - "auth_mode": format!("{:?}", auth.auth_mode), - })), - ); - return Ok((auth, true)); - } - Err(e) => { - tracing::warn!(error = %e, "auth: devbox migration save failed in auth flow"); - } - }, - Err(e) => { - tracing::warn!(error = %e, "auth: devbox login failed, falling through to interactive"); - } - } - } - - let url_tx = url_tx.and_then(|rc| rc.borrow_mut().take()); - let mut channels = code_rx.map(|code_rx| AuthChannels { url_tx, code_rx }); - - // Enterprise OIDC keeps loopback (customer IdPs may lack a device endpoint). - // xAI OAuth2 also defaults to loopback; the device flow (robust on - // remote/SSH where the loopback redirect can't reach the CLI) is opt-in via - // --device-auth / KIGI_LOGIN_DEVICE_FLOW / [auth] login_device_flow. - if crate::auth::oidc::is_configured(grok_com_config) { - return crate::auth::oidc::run_login_flow(grok_com_config, auth_manager, channels).await; - } - - if let Some(ref oauth2_cfg) = grok_com_config.oauth2 { - if should_use_device_flow(login_override).await { - // On `NotEnabled` (no device endpoint) `channels` is untouched, - // so we can fall back to loopback below. - match crate::auth::device_code::run_device_code_login_channels( - &oauth2_cfg.issuer, - &oauth2_cfg.client_id, - &oauth2_cfg.scopes, - auth_manager, - &mut channels, - ) - .await - { - Err(e) - if matches!( - e.downcast_ref::(), - Some(crate::auth::device_code::DeviceCodeError::NotEnabled) - ) => - { - tracing::warn!( - "auth: device flow unavailable (404), falling back to loopback login" - ); - } - other => return other, - } - } - return crate::auth::oidc::run_login_flow_with_config( - &oauth2_cfg.as_oidc(), - auth_manager, - channels, - ) - .await; - } - - tracing::error!( - "auth: no OAuth2 configuration available (neither enterprise OIDC nor xAI OAuth2 configured)" - ); - anyhow::bail!( - "No OAuth2 configuration available. Run `grok login` to authenticate, or contact your administrator if you use enterprise SSO." + let mut channels = channels; + crate::auth::device_code::run_device_code_login_channels( + &kigi_env::oauth_host(), + auth_manager, + &mut channels, ) + .await } -/// Non-interactive auth refresh: returns valid credentials if available without -/// ever triggering interactive login (browser, device code, etc.). +/// Non-interactive auth refresh: returns valid credentials if available +/// without ever triggering an interactive login (browser, device code). /// /// Tries in order: /// 1. Cached credentials (non-expired) -/// 2. OIDC silent refresh (if expired token has a refresh_token) -/// 3. External auth provider command (if configured) +/// 2. Silent refresh via the persisted refresh token /// /// Returns `None` when no valid credentials can be obtained non-interactively. -pub async fn try_ensure_fresh_auth(grok_com_config: &GrokComConfig) -> Option { +pub async fn try_ensure_fresh_auth(kimi_code_config: &KimiCodeConfig) -> Option { let kigi_home = kigi_home::kigi_home(); - let auth_manager = std::sync::Arc::new(AuthManager::new(&kigi_home, grok_com_config.clone())); + let auth_manager = Arc::new(AuthManager::new(&kigi_home, kimi_code_config.clone())); - // auth() handles cached-valid (fast path), OIDC refresh, external - // binary -- all through refresh_chain (single mutation point). - auth_manager.configure_refresher(grok_com_config.auth_provider_command.clone()); + // auth() handles cached-valid (fast path) and refresh — all through + // refresh_chain (single mutation point). + auth_manager.configure_refresher(); match auth_manager.auth().await { Ok(auth) => Some(auth), Err(e) => { @@ -686,87 +188,35 @@ pub async fn try_ensure_fresh_auth(grok_com_config: &GrokComConfig) -> Option Option { - if let Some(auth) = try_ensure_fresh_auth(grok_com_config).await { + kimi_code_config: &KimiCodeConfig, +) -> Option { + if let Some(auth) = try_ensure_fresh_auth(kimi_code_config).await { return Some(auth); } let kigi_home = kigi_home::kigi_home(); - let auth_manager = Arc::new(AuthManager::new(&kigi_home, grok_com_config.clone())); + let auth_manager = Arc::new(AuthManager::new(&kigi_home, kimi_code_config.clone())); - // A refresh failure leaves the session on disk (credentials are retained; - // the verdict gates re-attempts). Return it so consumers self-recover on - // 401, rather than disabling the relay for the leader's lifetime. - if let Some(expired) = expired_refreshable_session(&auth_manager) { - return Some(expired); - } - - mint_session_noninteractive(&auth_manager, grok_com_config).await + // A refresh failure leaves the session persisted (credentials are + // retained; the tombstone gates re-attempts). Return it so consumers + // self-recover on 401. + expired_refreshable_session(&auth_manager) } -/// A cached, refreshable session (not BYOK/ApiKey). Reached only after fresh +/// A cached, refreshable session (not an API key). Reached only after fresh /// auth failed, so in practice the token is expired but recoverable on 401. -fn expired_refreshable_session(auth_manager: &AuthManager) -> Option { +fn expired_refreshable_session(auth_manager: &AuthManager) -> Option { auth_manager .current_or_expired() - .filter(|a| a.is_xai_auth() && a.refresh_token.is_some()) -} - -/// Cold-start mint via non-interactive providers (external command, devbox); -/// `None` when none is available. -async fn mint_session_noninteractive( - auth_manager: &Arc, - grok_com_config: &GrokComConfig, -) -> Option { - // preferred_method=api_key: never auto-mint OIDC (fail-closed). - if grok_com_config.blocks_automatic_oidc() { - tracing::debug!( - "mint_session_noninteractive: skipped (preferred_method=api_key blocks automatic OIDC)" - ); - return None; - } - - if let Some(cmd) = grok_com_config.auth_provider_command.as_deref() { - match run_external_auth_provider(cmd, auth_manager, false, None).await { - Ok((auth, _)) => return Some(auth), - Err(e) => { - tracing::debug!(error = %e, "mint_session_noninteractive: external provider failed"); - } - } - } - - if crate::auth::devbox_login::is_devbox_environment() { - match crate::auth::devbox_login::mint_devbox_auth(auth_manager).await { - Ok(new_auth) => return Some(persist_or_use_minted(auth_manager, new_auth).await), - Err(e) => { - tracing::debug!(error = %e, "mint_session_noninteractive: devbox mint failed"); - } - } - } - - None -} - -/// Persist a minted token; on persist failure, return it unpersisted rather -/// than dropping a valid credential. -async fn persist_or_use_minted(auth_manager: &AuthManager, new_auth: GrokAuth) -> GrokAuth { - match auth_manager.save_without_enrichment(new_auth.clone()).await { - Ok(auth) => { - let _ = auth_manager.remove_scope(LEGACY_AUTH_SCOPE); - auth - } - Err(e) => { - tracing::warn!(error = %e, "mint persist failed; using unpersisted token"); - new_auth - } - } + .filter(|a| a.is_session_auth() && a.refresh_token.is_some()) } /// Print the CLI "signed in" confirmation, clearing the spinner line first. -fn report_signed_in(auth: &GrokAuth) { +fn report_signed_in(auth: &KimiAuth) { eprint!("\r\x1b[K"); match auth.email { Some(ref email) => eprintln!("✓ Signed in as {email}"), @@ -776,56 +226,24 @@ fn report_signed_in(auth: &GrokAuth) { /// CLI auth entrypoint. For GUI, use `run_auth_flow_with_stderr_bridge`. pub async fn ensure_authenticated( - grok_com_config: &GrokComConfig, + kimi_code_config: &KimiCodeConfig, reauth: bool, message_prefix: Option<&str>, -) -> anyhow::Result { - ensure_authenticated_with_override( - grok_com_config, - reauth, - message_prefix, - LoginTransportOverride::None, - ) - .await -} - -/// Like [`ensure_authenticated`] but with an explicit login-transport override -/// (from `--oauth` / `--device-auth`). Used by `run_cli_login`. -pub async fn ensure_authenticated_with_override( - grok_com_config: &GrokComConfig, - reauth: bool, - message_prefix: Option<&str>, - login_override: LoginTransportOverride, -) -> anyhow::Result { +) -> anyhow::Result { let kigi_home = kigi_home::kigi_home(); - let auth_manager = Arc::new(AuthManager::new(&kigi_home, grok_com_config.clone())); + let auth_manager = Arc::new(AuthManager::new(&kigi_home, kimi_code_config.clone())); - // If not re-authing, accept any valid non-WebLogin credential. - // WebLogin tokens are always skipped — they must be migrated to OIDC. + // If not re-authing, accept any valid cached credential. if !reauth && let Some(auth) = auth_manager.current() { - if auth.auth_mode != super::AuthMode::WebLogin { - return Ok(auth); - } - tracing::info!("auth: skipping cached WebLogin credential, will migrate to OIDC"); - auth_manager.clear_in_memory(); - let _ = auth_manager.remove_scope(LEGACY_AUTH_SCOPE); + return Ok(auth); } - // Context only — the flow below prints the "Signing in…" line itself. + // Context only — the flow below prints the sign-in prompts itself. if let Some(msg) = message_prefix { eprintln!("{msg}"); } - let (auth, did_auth) = run_auth_flow( - &auth_manager, - grok_com_config, - reauth, - None, - None, - None, - login_override, - ) - .await?; + let (auth, did_auth) = run_auth_flow(&auth_manager, kimi_code_config, reauth, None).await?; if did_auth { report_signed_in(&auth); @@ -834,105 +252,40 @@ pub async fn ensure_authenticated_with_override( Ok(auth) } -/// Decides *whether to prompt* for an interactive login (the wire credential is -/// chosen separately by `ShellAuthCredentialProvider`). +/// Decides *whether to prompt* for an interactive login (the wire credential +/// is chosen separately by `ShellAuthCredentialProvider`). /// -/// With `has_noninteractive_auth`, only refresh a cached token best-effort (no -/// browser, no cold mint); otherwise require an interactive login. +/// With `has_noninteractive_auth`, only refresh a cached token best-effort +/// (no browser, no device prompt); otherwise require an interactive login. pub async fn ensure_authenticated_or_noninteractive( - grok_com_config: &GrokComConfig, + kimi_code_config: &KimiCodeConfig, has_noninteractive_auth: bool, message_prefix: Option<&str>, -) -> anyhow::Result> { +) -> anyhow::Result> { if has_noninteractive_auth { - Ok(try_ensure_fresh_auth(grok_com_config).await) + Ok(try_ensure_fresh_auth(kimi_code_config).await) } else { - ensure_authenticated(grok_com_config, false, message_prefix) + ensure_authenticated(kimi_code_config, false, message_prefix) .await .map(Some) } } -/// Unified `grok login` handler for CLI entry points (tui, pager). +/// `kigi login` handler for CLI entry points (tui, pager): the device-code +/// flow is THE login. /// -/// Precedence: `--oauth` forces loopback, `--device-auth` forces device, -/// otherwise `KIGI_LOGIN_DEVICE_FLOW` env / `[auth] login_device_flow` config / -/// loopback default. Both transports run through `run_auth_flow_inner` so the -/// external auth provider and devbox auto-migration are tried first. -pub async fn run_cli_login( - config: &crate::agent::config::Config, - oauth: bool, - device_auth: bool, - devbox: bool, -) -> anyhow::Result<()> { - let login_override = LoginTransportOverride::from_flags(oauth, device_auth); - - // Mirror `run_auth_flow_inner`'s precedence: enterprise OIDC (oidc=Some, - // oauth2=None) always uses the loopback flow; only the xAI OAuth2 provider - // supports the device flow. Without this guard, `grok login` on an - // enterprise-OIDC deployment would wrongly enter the device branch (which - // requires `oauth2`) and error. - let authenticated = if devbox { - super::devbox_login::run_devbox_login(config).await? - } else if cli_should_use_device(&config.grok_com_config, login_override).await { - if config.grok_com_config.oauth2.is_none() { - // No OIDC and no oauth2 here, so `--oauth` can't help. - anyhow::bail!("Sign-in is not available for this deployment. Set XAI_API_KEY instead."); - } - let kigi_home = kigi_home::kigi_home(); - let auth_manager = Arc::new(AuthManager::new(&kigi_home, config.grok_com_config.clone())); - // Route through the shared inner flow (not `run_device_code_login` - // directly) so the external auth provider and devbox auto-migration run - // before the interactive device login. `force_interactive` skips the - // up-front clear, so abandoning the device prompt doesn't log the user - // out; on `NotEnabled` it falls back to loopback. - // Already resolved/logged above; pass `Preresolved(true)` so the inner flow - // honors device without a second fetch or a duplicate `cli`-attributed log. - let (auth, did_auth) = run_auth_flow_interactive( - &auth_manager, - &config.grok_com_config, - None, - None, - None, - LoginTransportOverride::Preresolved(true), - ) - .await?; - if did_auth { - report_signed_in(&auth); - } - auth - } else { - // OIDC has no device endpoint, so `--device-auth` falls back here. - if device_auth && crate::auth::oidc::is_configured(&config.grok_com_config) { - eprintln!( - "Device-code login isn't available for your SSO provider; using browser sign-in." - ); - } - // Loopback. `reauth=true` clears creds up front (legacy-scope hygiene), - // so abandoning logs you out — unlike the device branch above. - // Already resolved/logged above; pass `Preresolved(false)` so the inner - // flow honors loopback without a duplicate `cli`-attributed log. - ensure_authenticated_with_override( - &config.grok_com_config, - true, - None, - LoginTransportOverride::Preresolved(false), - ) - .await? - }; - - // Sync this principal's config now rather than waiting for the background - // tick. Stay quiet about absence/failure during login — confirm only when - // config was actually applied; `grok setup` reports the no-config case. - let outcome = crate::managed_config::post_login_sync(Some(authenticated)).await; - match outcome { - crate::managed_config::ManagedConfigSync::Updated { is_team: true } => { - eprintln!("Applied your team's managed configuration."); - } - crate::managed_config::ManagedConfigSync::Updated { is_team: false } => { - eprintln!("Applied your deployment's managed configuration."); - } - _ => {} +/// Runs with `force_interactive` semantics — cached credentials are skipped +/// but not cleared, so abandoning the device prompt doesn't log the user out. +pub async fn run_cli_login(config: &crate::agent::config::Config) -> anyhow::Result<()> { + let kigi_home = kigi_home::kigi_home(); + let auth_manager = Arc::new(AuthManager::new( + &kigi_home, + config.kimi_code_config.clone(), + )); + let (auth, did_auth) = + run_auth_flow_inner(&auth_manager, &config.kimi_code_config, false, true, None).await?; + if did_auth { + report_signed_in(&auth); } Ok(()) } @@ -945,14 +298,15 @@ pub struct LogoutResult { pub was_logged_in: bool, /// Email of the session that was cleared (if available). pub email: Option, - /// `true` if `XAI_API_KEY` / `KIGI_CODE_XAI_API_KEY` env var is set. + /// `true` if an API-key env var is set. pub api_key_still_set: bool, } /// Core logout logic shared by the CLI subcommand and the ACP handler. /// /// When `scope` is `None`, clears the default scope (same as `/logout` -/// in the TUI). When `Some`, removes only that scope entry. +/// in the TUI). When `Some`, removes only that scope entry. The session +/// credential is removed from both stores (keyring + file). pub fn perform_logout( auth_manager: &AuthManager, scope: Option<&str>, @@ -961,7 +315,7 @@ pub fn perform_logout( let email = auth.as_ref().and_then(|a| a.email.clone()); let was_logged_in = auth.is_some(); // Intentional credential removal must be attributable in - // unified.jsonl, so a later "auth.json entry gone" can be + // unified.jsonl, so a later "auth entry gone" can be // distinguished from accidental loss (deleted/corrupt file). kigi_log::unified_log::info( "auth: logout", @@ -969,7 +323,6 @@ pub fn perform_logout( Some(serde_json::json!({ "was_logged_in": was_logged_in, "scope": scope.unwrap_or("(current)"), - "user_id": auth.as_ref().map(|a| a.user_id.clone()), })), ); if was_logged_in { @@ -978,9 +331,6 @@ pub fn perform_logout( } else { auth_manager.clear()?; } - // Clear the synced files if no principal remains to own them. A scoped - // logout that leaves a team (or a deployment key) signed in keeps them. - crate::managed_config::clear_orphan(); } Ok(LogoutResult { was_logged_in, @@ -989,17 +339,17 @@ pub fn perform_logout( }) } -/// `grok logout` CLI handler. Calls [`perform_logout`] and formats +/// `kigi logout` CLI handler. Calls [`perform_logout`] and formats /// the result to stderr. pub fn run_cli_logout(config: &crate::agent::config::Config) -> anyhow::Result<()> { let kigi_home = kigi_home::kigi_home(); - let auth_manager = AuthManager::new(&kigi_home, config.grok_com_config.clone()); + let auth_manager = AuthManager::new(&kigi_home, config.kimi_code_config.clone()); let result = perform_logout(&auth_manager, None) .map_err(|e| anyhow::anyhow!("Failed to clear auth: {e}"))?; if !result.was_logged_in { eprintln!("No cached session to log out of."); if result.api_key_still_set { - eprintln!("You are authenticated via XAI_API_KEY (environment variable)."); + eprintln!("You are authenticated via an API-key environment variable."); } return Ok(()); } @@ -1009,7 +359,7 @@ pub fn run_cli_logout(config: &crate::agent::config::Config) -> anyhow::Result<( eprintln!("Logged out"); } if result.api_key_still_set { - eprintln!("XAI_API_KEY is still set and will be used for authentication."); + eprintln!("An API-key environment variable is still set and will be used."); } Ok(()) } @@ -1018,43 +368,29 @@ pub fn run_cli_logout(config: &crate::agent::config::Config) -> anyhow::Result<( mod tests { use super::*; use crate::auth::AuthMode; - use crate::auth::config::XAI_OAUTH2_ISSUER; - use crate::env::EnvVarGuard; use chrono::Utc; - /// Run `f` with `KIGI_LOGIN_DEVICE_FLOW` set to `value` (unset for `None`). - /// `EnvVarGuard` serializes the process env and restores it on drop, so - /// `resolve_device_flow` reads the env tier from a known state. - fn with_device_flow_env(value: Option, f: impl FnOnce() -> T) -> T { - let _guard = match value { - Some(true) => EnvVarGuard::set("KIGI_LOGIN_DEVICE_FLOW", "true"), - Some(false) => EnvVarGuard::set("KIGI_LOGIN_DEVICE_FLOW", "false"), - None => EnvVarGuard::remove("KIGI_LOGIN_DEVICE_FLOW"), - }; - f() - } - - // A grok.com first-party (x.ai-issuer) OIDC session — `is_xai_auth()` true. - fn oidc_session(key: &str, refresh: Option<&str>) -> GrokAuth { - GrokAuth { + // A Kimi Code OAuth session credential. + fn oauth_session(key: &str, refresh: Option<&str>) -> KimiAuth { + KimiAuth { key: key.into(), - auth_mode: AuthMode::Oidc, - oidc_issuer: Some(XAI_OAUTH2_ISSUER.to_string()), + auth_mode: AuthMode::OAuth, refresh_token: refresh.map(str::to_string), - ..GrokAuth::test_default() + ..KimiAuth::test_default() } } #[test] fn expired_refreshable_session_gate() { let dir = tempfile::tempdir().unwrap(); - let mgr = AuthManager::new(dir.path(), GrokComConfig::default()); + let mgr = AuthManager::new(dir.path(), KimiCodeConfig::default()); - // Expired but refreshable → returned. Guards a `current_or_expired()` -> - // `current()` regression that would disable the relay on a transient blip. - mgr.hot_swap(GrokAuth { + // Expired but refreshable → returned. Guards a `current_or_expired()` + // -> `current()` regression that would disable detached consumers on + // a transient blip. + mgr.hot_swap(KimiAuth { expires_at: Some(Utc::now() - chrono::Duration::hours(1)), - ..oidc_session("expired-but-refreshable", Some("rt")) + ..oauth_session("expired-but-refreshable", Some("rt")) }); assert!( mgr.current().is_none(), @@ -1065,648 +401,55 @@ mod tests { Some("expired-but-refreshable".to_string()) ); - // No refresh token → rejected: never hand the relay a token it can't - // recover on 401 (the gate `for_session` doesn't check this). - mgr.hot_swap(oidc_session("no-rt", None)); + // No refresh token → rejected: never hand consumers a token they + // can't recover on 401. + mgr.hot_swap(oauth_session("no-rt", None)); assert!(expired_refreshable_session(&mgr).is_none()); - // An expired first-party *external* credential with a refresh token - // is likewise recoverable — 401 recovery re-runs the provider binary - // (the refresh token is a recoverability marker, not a grant input). - mgr.hot_swap(GrokAuth { - auth_mode: AuthMode::External, + // API keys are excluded (not a session). + mgr.hot_swap(KimiAuth { + auth_mode: AuthMode::ApiKey, expires_at: Some(Utc::now() - chrono::Duration::hours(1)), - ..oidc_session("expired-external", Some("rt")) - }); - assert_eq!( - expired_refreshable_session(&mgr).map(|a| a.key), - Some("expired-external".to_string()) - ); - - // Third-party external (no x.ai issuer) stays excluded. - mgr.hot_swap(GrokAuth { - oidc_issuer: None, - auth_mode: AuthMode::External, - expires_at: Some(Utc::now() - chrono::Duration::hours(1)), - ..oidc_session("expired-external-3p", Some("rt")) + ..oauth_session("api-key", Some("rt")) }); assert!(expired_refreshable_session(&mgr).is_none()); } - #[cfg(unix)] - #[tokio::test] - async fn persist_or_use_minted_returns_token_when_save_fails() { - use std::os::unix::fs::PermissionsExt; - // Read-only kigi_home: reading a missing auth.json succeeds (empty), but - // writing fails — exercising the save-failure path. - let dir = tempfile::tempdir().unwrap(); - std::fs::set_permissions(dir.path(), std::fs::Permissions::from_mode(0o500)).unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - let minted = oidc_session("minted-token", Some("rt")); + // ── run_auth_flow: expired path with persisted token ───────────── - let save = mgr.save_without_enrichment(minted.clone()).await; - // Root bypasses 0o500, so the write can't be forced to fail there — skip - // explicitly. Non-root MUST see the save fail (or this proves nothing). - if unsafe { libc::geteuid() } == 0 { - return; - } - assert!( - save.is_err(), - "non-root: save into a read-only dir must fail" - ); - let out = persist_or_use_minted(&mgr, minted).await; - assert_eq!( - out.key, "minted-token", - "must return the unpersisted minted token" - ); - } - - /// Proxy URL on a closed port: inline enrichment fails fast instead of - /// reaching outside the test. - fn dead_proxy_url() -> String { - let port = { - let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); - l.local_addr().unwrap().port() - }; - format!("http://127.0.0.1:{port}") - } - - #[tokio::test] - async fn mint_session_noninteractive_uses_external_provider() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig { - auth_provider_command: Some("printf '%s' xai-ext-token".to_string()), - ..GrokComConfig::default() - }; - let mgr = Arc::new( - AuthManager::new(dir.path(), cfg.clone()).with_proxy_base_url(&dead_proxy_url()), - ); - - let auth = mint_session_noninteractive(&mgr, &cfg).await; - assert_eq!(auth.map(|a| a.key), Some("xai-ext-token".to_string())); - } - - /// External-provider output is team-pinned before persist (parity with OIDC - /// / device-code): a wrong-team token is rejected and nothing is written. - #[tokio::test] - async fn external_provider_rejects_wrong_team_and_persists_nothing() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new( - AuthManager::new(dir.path(), pinned_cfg("team-good")) - .with_proxy_base_url(&dead_proxy_url()), - ); - let cmd = format!("printf '%s' {}", team_jwt("team-wrong")); - - assert!( - run_external_auth_provider(&cmd, &mgr, false, None) - .await - .is_err(), - "wrong-team external token must be rejected" - ); - assert!( - mgr.current_or_expired().is_none(), - "rejected external login must persist nothing" - ); - assert!( - !dir.path().join("auth.json").exists(), - "rejected external login must not write auth.json" - ); - } - - /// A matching-team external token is accepted and persisted. - #[tokio::test] - async fn external_provider_accepts_matching_team() { - let dir = tempfile::tempdir().unwrap(); - let jwt = team_jwt("team-good"); - let mgr = Arc::new( - AuthManager::new(dir.path(), pinned_cfg("team-good")) - .with_proxy_base_url(&dead_proxy_url()), - ); - let cmd = format!("printf '%s' {jwt}"); - - let (auth, _) = run_external_auth_provider(&cmd, &mgr, false, None) - .await - .expect("matching-team external token must be accepted"); - assert_eq!(auth.key, jwt); - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn external_reauth_without_prev_auth_enriches_inline() { - // Regression: reauth clears the manager before the provider runs with - // is_refresh=true; flags must then come from /user, not default empty. - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let app = axum::Router::new().route( - "/user", - axum::routing::get(|| async { - axum::Json(serde_json::json!({ - "userId": "u-1", - "teamBlockedReasons": ["BLOCKED_REASON_NO_LOGS"], - })) - }), - ); - tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new( - AuthManager::new(dir.path(), GrokComConfig::default()) - .with_proxy_base_url(&format!("http://127.0.0.1:{port}")), - ); - assert!(mgr.current_or_expired().is_none(), "precondition: no auth"); - - let (auth, _) = run_external_auth_provider("printf '%s' fresh-token", &mgr, true, None) - .await - .unwrap(); - assert_eq!(auth.key, "fresh-token"); - assert!(auth.is_zdr_team(), "flags must come from /user fetch"); - assert_eq!(auth.user_id, "u-1"); - } - - #[tokio::test] - async fn external_refresh_carries_profile_without_network() { - // Carry path must not need /user: dead proxy port, flags from prev. - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new( - AuthManager::new(dir.path(), GrokComConfig::default()) - .with_proxy_base_url(&dead_proxy_url()), - ); - mgr.hot_swap(GrokAuth { - team_blocked_reasons: vec!["BLOCKED_REASON_NO_LOGS".into()], - organization_id: Some("org-1".into()), - ..oidc_session("old-token", None) - }); - - let (auth, _) = run_external_auth_provider("printf '%s' fresh-token", &mgr, true, None) - .await - .unwrap(); - assert_eq!(auth.key, "fresh-token"); - assert!(auth.is_zdr_team(), "flags must carry from previous auth"); - assert_eq!(auth.user_id, "test-user"); - assert_eq!(auth.organization_id.as_deref(), Some("org-1")); - } - - #[tokio::test] - async fn device_flow_still_runs_external_provider() { - // Regression: with the device flow opted into (--device-auth), the - // external auth provider must still run first. `run_cli_login`'s device - // branch goes through `run_auth_flow_interactive`, so that path must - // pick up the provider instead of starting an interactive device login. - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig { - auth_provider_command: Some("printf '%s' xai-ext-token".to_string()), - // oauth2=Some, oidc=None → the device flow is available (opt-in). - ..GrokComConfig::default() - }; - assert!( - cli_should_use_device(&cfg, LoginTransportOverride::ForceDevice).await, - "precondition: --device-auth resolves to the device flow" - ); - let mgr = Arc::new( - AuthManager::new(dir.path(), cfg.clone()).with_proxy_base_url(&dead_proxy_url()), - ); - let (auth, did_auth) = run_auth_flow_interactive( - &mgr, - &cfg, - None, - None, - None, - LoginTransportOverride::ForceDevice, - ) - .await - .expect("external provider should satisfy login without device flow"); - assert_eq!( - auth.key, "xai-ext-token", - "external provider token must win" - ); - assert!(did_auth); - } - - #[test] - fn login_transport_override_maps_to_cli_bool() { - // `--oauth` → loopback, `--device-auth` → device, no flag → no override. - assert_eq!(LoginTransportOverride::None.as_cli_bool(), None); - assert_eq!( - LoginTransportOverride::ForceLoopback.as_cli_bool(), - Some(false) - ); - assert_eq!( - LoginTransportOverride::ForceDevice.as_cli_bool(), - Some(true) - ); - // Pre-resolved transports are honored upstream, never via the CLI tier, so - // they must not present a CLI value (which would mis-label the source as cli). - assert_eq!( - LoginTransportOverride::Preresolved(true).as_cli_bool(), - None - ); - assert_eq!( - LoginTransportOverride::Preresolved(false).as_cli_bool(), - None - ); - } - - #[tokio::test] - async fn preresolved_bypasses_resolver_and_is_never_cli() { - // Regression for the double-log / source=cli bug: the inner flow must - // honor `Preresolved` WITHOUT re-running the resolver. Each case pins the - // opposite env value, so a leak into the resolver would flip the result — - // returning the carried value proves the early return (and no second log). - { - let _guard = EnvVarGuard::set("KIGI_LOGIN_DEVICE_FLOW", "false"); - assert!( - should_use_device_flow(LoginTransportOverride::Preresolved(true)).await, - "Preresolved(true) honors device without re-resolving" - ); - } - { - let _guard = EnvVarGuard::set("KIGI_LOGIN_DEVICE_FLOW", "true"); - assert!( - !should_use_device_flow(LoginTransportOverride::Preresolved(false)).await, - "Preresolved(false) honors loopback without re-resolving" - ); - assert!( - should_use_device_flow(LoginTransportOverride::None).await, - "the resolver path still honors env (sole resolution)" - ); - } - // Even if it reached the resolver it carries no CLI value, so a remote - // decision is the remote tier, never cli. - with_device_flow_env(None, || { - assert_eq!( - resolve_device_flow(LoginTransportOverride::Preresolved(true), None, Some(true)) - .source, - crate::agent::config::ConfigSource::Remote, - "Preresolved must never resolve as the cli tier" - ); - }); - } - - #[test] - fn from_flags_prefers_oauth_over_device() { - // `--oauth` (loopback) wins if both are set — a defensive guard for the - // ACP meta path (clap already blocks both flags on the CLI). - assert_eq!( - LoginTransportOverride::from_flags(true, true), - LoginTransportOverride::ForceLoopback - ); - assert_eq!( - LoginTransportOverride::from_flags(true, false), - LoginTransportOverride::ForceLoopback - ); - assert_eq!( - LoginTransportOverride::from_flags(false, true), - LoginTransportOverride::ForceDevice - ); - assert_eq!( - LoginTransportOverride::from_flags(false, false), - LoginTransportOverride::None - ); - } - - #[tokio::test] - async fn enterprise_oidc_never_uses_device_flow() { - // oidc=Some, oauth2=None: `grok login` must use loopback, not device — - // even when --device-auth forces device (which would otherwise be true). - // ForceDevice short-circuits the remote settings fetch, so this stays hermetic. - let cfg = GrokComConfig { - oidc: Some(crate::auth::OidcAuthConfig { - issuer: "https://idp.example".into(), - client_id: "client".into(), - scopes: vec!["openid".into()], - audience: None, - }), - oauth2: None, - ..GrokComConfig::default() - }; - assert!( - !cli_should_use_device(&cfg, LoginTransportOverride::ForceDevice).await, - "enterprise OIDC must stay on loopback" - ); - // The xAI OAuth2 provider (oidc=None, oauth2=Some) does use device. - let xai = GrokComConfig::default(); - assert!(xai.oauth2.is_some() && xai.oidc.is_none()); - assert!(cli_should_use_device(&xai, LoginTransportOverride::ForceDevice).await); - } - - #[test] - fn device_flow_precedence_cli_beats_env_config_remote() { - // CLI flag wins over a *conflicting* env + config + remote feature flag. - with_device_flow_env(Some(true), || { - assert!( - !resolve_device_flow( - LoginTransportOverride::ForceLoopback, - Some(true), - Some(true) - ) - .value, - "--oauth must force loopback even when env+config+remote say device" - ); - }); - with_device_flow_env(Some(false), || { - assert!( - resolve_device_flow( - LoginTransportOverride::ForceDevice, - Some(false), - Some(false) - ) - .value, - "--device-auth must force device even when env+config+remote say loopback" - ); - }); - } - - #[test] - fn device_flow_precedence_env_beats_config() { - // No CLI flag: env wins over a conflicting config. - with_device_flow_env(Some(false), || { - assert!(!resolve_device_flow(LoginTransportOverride::None, Some(true), None).value); - }); - with_device_flow_env(Some(true), || { - assert!(resolve_device_flow(LoginTransportOverride::None, Some(false), None).value); - }); - } - - #[test] - fn device_flow_env_beats_remote() { - // env sits above the remote feature flag. - with_device_flow_env(Some(false), || { - assert!( - !resolve_device_flow(LoginTransportOverride::None, None, Some(true)).value, - "env=loopback must win over remote=device" - ); - }); - with_device_flow_env(Some(true), || { - assert!( - resolve_device_flow(LoginTransportOverride::None, None, Some(false)).value, - "env=device must win over remote=loopback" - ); - }); - } - - #[test] - fn device_flow_config_beats_remote() { - // Local config sits above the remote feature flag (env unset so config decides). - with_device_flow_env(None, || { - assert!( - !resolve_device_flow(LoginTransportOverride::None, Some(false), Some(true)).value, - "config=loopback must win over remote=device" - ); - assert!( - resolve_device_flow(LoginTransportOverride::None, Some(true), Some(false)).value, - "config=device must win over remote=loopback" - ); - }); - } - - #[test] - fn device_flow_precedence_config_then_default() { - // No CLI flag, no env: config decides; absent everything → loopback. - with_device_flow_env(None, || { - assert!(!resolve_device_flow(LoginTransportOverride::None, Some(false), None).value); - assert!(resolve_device_flow(LoginTransportOverride::None, Some(true), None).value); - assert!( - !resolve_device_flow(LoginTransportOverride::None, None, None).value, - "default is loopback" - ); - }); - } - - #[test] - fn device_flow_remote_then_default() { - // No CLI flag, no env, no config: the remote feature flag drives the rollout. - with_device_flow_env(None, || { - assert!( - resolve_device_flow(LoginTransportOverride::None, None, Some(true)).value, - "remote=device rolls device-auth in when nothing local is set" - ); - assert!( - !resolve_device_flow(LoginTransportOverride::None, None, Some(false)).value, - "remote=loopback keeps loopback when nothing local is set" - ); - // remote settings unavailable / flag unset → None → hardcoded loopback default. - assert!( - !resolve_device_flow(LoginTransportOverride::None, None, None).value, - "remote settings unavailable falls back to the loopback default" - ); - }); - } - - #[test] - fn device_flow_records_deciding_tier() { - // The resolver records which tier decided, so the rollout ramp can log it. - use crate::agent::config::ConfigSource; - with_device_flow_env(Some(false), || { - assert_eq!( - resolve_device_flow(LoginTransportOverride::ForceDevice, Some(false), None).source, - ConfigSource::Cli, - "an explicit CLI flag is reported as the cli tier" - ); - }); - with_device_flow_env(Some(true), || { - assert_eq!( - resolve_device_flow(LoginTransportOverride::None, None, Some(false)).source, - ConfigSource::Env - ); - }); - with_device_flow_env(None, || { - assert_eq!( - resolve_device_flow(LoginTransportOverride::None, Some(true), Some(false)).source, - ConfigSource::Config - ); - assert_eq!( - resolve_device_flow(LoginTransportOverride::None, None, Some(true)).source, - ConfigSource::Remote, - "the remote feature flag is reported as the remote tier" - ); - assert_eq!( - resolve_device_flow(LoginTransportOverride::None, None, None).source, - ConfigSource::Default - ); - }); - } - - fn legacy_auth() -> GrokAuth { - GrokAuth { - key: "k".into(), - auth_mode: AuthMode::WebLogin, - create_time: Utc::now(), - user_id: "u".into(), - email: None, - first_name: None, - last_name: None, - profile_image_asset_id: None, - principal_type: None, - principal_id: None, - team_id: None, - team_name: None, - team_role: None, - organization_id: None, - organization_name: None, - organization_role: None, - user_blocked_reason: None, - team_blocked_reasons: vec![], - coding_data_retention_opt_out: false, - has_grok_code_access: None, - refresh_token: None, - expires_at: None, - oidc_issuer: None, - oidc_client_id: None, - } - } - - fn oidc_auth(issuer: &str) -> GrokAuth { - GrokAuth { - oidc_issuer: Some(issuer.into()), - auth_mode: AuthMode::Oidc, - ..legacy_auth() - } - } - - #[test] - fn weblogin_cred_is_never_compatible() { - let cfg = GrokComConfig::default(); - assert!(!is_cached_credential_compatible(&legacy_auth(), &cfg)); - } - - #[test] - fn oidc_cred_with_matching_issuer_is_compatible() { - let cfg = GrokComConfig::default(); - assert!(is_cached_credential_compatible( - &oidc_auth(XAI_OAUTH2_ISSUER), - &cfg, - )); - } - - #[test] - fn external_cred_compatibility_follows_issuer() { - let cfg = GrokComConfig::default(); - - // A first-party external credential (provider emitted the issuer) is - // reused by interactive login like an OIDC session instead of - // re-running the provider. - assert!(is_cached_credential_compatible( - &GrokAuth { - auth_mode: AuthMode::External, - ..oidc_auth(XAI_OAUTH2_ISSUER) - }, - &cfg, - )); - - // Without an issuer (bare-token providers), external credentials stay - // incompatible and interactive login starts fresh, as before. - assert!(!is_cached_credential_compatible( - &GrokAuth { - auth_mode: AuthMode::External, - oidc_issuer: None, - ..legacy_auth() - }, - &cfg, - )); - } - - fn ensure_crypto_provider() { - let _ = jsonwebtoken::crypto::rust_crypto::DEFAULT_PROVIDER.install_default(); - } - - fn team_jwt(principal_id: &str) -> String { - ensure_crypto_provider(); - jsonwebtoken::encode( - &jsonwebtoken::Header::new(jsonwebtoken::Algorithm::HS256), - &serde_json::json!({ - "sub": "user-1", - "principal_type": "Team", - "principal_id": principal_id, - "exp": 9999999999u64, - }), - &jsonwebtoken::EncodingKey::from_secret(b"test-secret"), - ) - .unwrap() - } - - fn pinned_cfg(team: &str) -> GrokComConfig { - GrokComConfig { - force_login_team_uuid: Some(crate::auth::config::ForceLoginTeam::Single(team.into())), - ..GrokComConfig::default() - } - } - - /// Under a team pin, a cached session for a different team is not reused by - /// interactive login — it falls through to a fresh, compliant login. - #[test] - fn cached_cred_with_wrong_team_is_incompatible() { - let auth = GrokAuth { - key: team_jwt("team-wrong"), - ..oidc_auth(XAI_OAUTH2_ISSUER) - }; - assert!(!is_cached_credential_compatible( - &auth, - &pinned_cfg("team-good") - )); - } - - /// A cached session for the pinned team is reused normally. - #[test] - fn cached_cred_with_matching_team_is_compatible() { - let auth = GrokAuth { - key: team_jwt("team-good"), - ..oidc_auth(XAI_OAUTH2_ISSUER) - }; - assert!(is_cached_credential_compatible( - &auth, - &pinned_cfg("team-good") - )); - } - - // ── run_auth_flow: expired path with disk token ───────────────── - - /// When in-memory token is expired but disk has a valid token, - /// run_auth_flow should return the disk token without interactive login. + /// When the in-memory token is expired but the store has a valid token, + /// run_auth_flow should return the stored token without interactive login. #[tokio::test] async fn run_auth_flow_uses_valid_disk_token_when_expired() { let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); + let cfg = KimiCodeConfig::default(); - // Write a valid token to disk via a second AuthManager (simulates - // a sibling process that already refreshed). - let writer = Arc::new( - AuthManager::new(dir.path(), cfg.clone()).with_proxy_base_url("http://127.0.0.1:1"), - ); - let valid_disk = GrokAuth { + // Write a valid token via a second AuthManager (simulates a sibling + // process that already refreshed). + let writer = Arc::new(AuthManager::new(dir.path(), cfg.clone())); + let valid_disk = KimiAuth { key: "fresh-token-from-disk".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, expires_at: Some(Utc::now() + chrono::Duration::hours(1)), + expires_in: Some(3600), refresh_token: Some("new-rt".into()), - oidc_issuer: Some(XAI_OAUTH2_ISSUER.into()), - oidc_client_id: Some("client-1".into()), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }; writer.update(valid_disk).await.unwrap(); // Primary manager: in-memory token is expired let mgr = Arc::new(AuthManager::new(dir.path(), cfg.clone())); - let expired = GrokAuth { + let expired = KimiAuth { key: "expired-access-token".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, expires_at: Some(Utc::now() - chrono::Duration::hours(1)), refresh_token: Some("old-rt".into()), - oidc_issuer: Some(XAI_OAUTH2_ISSUER.into()), - oidc_client_id: Some("client-1".into()), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }; mgr.hot_swap(expired); assert!(mgr.is_expired()); - let (auth, is_new_login) = run_auth_flow( - &mgr, - &cfg, - false, // not reauth - None, - None, - None, - LoginTransportOverride::None, - ) - .await - .unwrap(); + let (auth, is_new_login) = run_auth_flow(&mgr, &cfg, false, None).await.unwrap(); assert_eq!(auth.key, "fresh-token-from-disk"); assert!(!is_new_login, "should not be a new login"); @@ -1714,35 +457,24 @@ mod tests { assert_eq!(mgr.current().unwrap().key, "fresh-token-from-disk"); } - /// When in-memory token is valid (not expired), run_auth_flow should - /// return it directly without checking disk. + /// When the in-memory token is valid (not expired), run_auth_flow should + /// return it directly without checking the store. #[tokio::test] async fn run_auth_flow_returns_cached_when_valid() { let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); + let cfg = KimiCodeConfig::default(); let mgr = Arc::new(AuthManager::new(dir.path(), cfg.clone())); - let valid = GrokAuth { + let valid = KimiAuth { key: "still-valid".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, expires_at: Some(Utc::now() + chrono::Duration::hours(1)), - oidc_issuer: Some(XAI_OAUTH2_ISSUER.into()), - oidc_client_id: Some("client-1".into()), - ..GrokAuth::test_default() + expires_in: Some(3600), + ..KimiAuth::test_default() }; mgr.hot_swap(valid); - let (auth, is_new_login) = run_auth_flow( - &mgr, - &cfg, - false, - None, - None, - None, - LoginTransportOverride::None, - ) - .await - .unwrap(); + let (auth, is_new_login) = run_auth_flow(&mgr, &cfg, false, None).await.unwrap(); assert_eq!(auth.key, "still-valid"); assert!(!is_new_login); @@ -1751,19 +483,15 @@ mod tests { #[tokio::test] async fn run_auth_flow_defers_to_consumer_refresh_on_transient_failure() { let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); + let cfg = KimiCodeConfig::default(); - let writer = Arc::new( - AuthManager::new(dir.path(), cfg.clone()).with_proxy_base_url("http://127.0.0.1:1"), - ); - let expired_with_rt = GrokAuth { + let writer = Arc::new(AuthManager::new(dir.path(), cfg.clone())); + let expired_with_rt = KimiAuth { key: "expired-access-token".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, expires_at: Some(Utc::now() - chrono::Duration::hours(1)), refresh_token: Some("valid-refresh-token".into()), - oidc_issuer: Some(XAI_OAUTH2_ISSUER.into()), - oidc_client_id: Some("client-1".into()), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }; writer.update(expired_with_rt.clone()).await.unwrap(); @@ -1772,123 +500,13 @@ mod tests { assert!(mgr.is_expired()); mgr.set_refresher(std::sync::Arc::new(AlwaysTransientRefresher)); - let (auth, is_new_login) = run_auth_flow( - &mgr, - &cfg, - false, // not reauth - None, - None, - None, - LoginTransportOverride::None, - ) - .await - .unwrap(); + let (auth, is_new_login) = run_auth_flow(&mgr, &cfg, false, None).await.unwrap(); assert_eq!(auth.key, "expired-access-token"); assert!(auth.refresh_token.is_some()); assert!(!is_new_login); } - #[tokio::test] - async fn run_auth_flow_falls_through_when_no_refresh_token() { - let dir = tempfile::tempdir().unwrap(); - // Point the OAuth2 issuer at a non-routable address so the OIDC - // discovery fails immediately without opening a browser window. - let mut cfg = GrokComConfig::default(); - cfg.oauth2.as_mut().unwrap().issuer = "http://127.0.0.1:1".into(); - - let writer = Arc::new( - AuthManager::new(dir.path(), cfg.clone()).with_proxy_base_url("http://127.0.0.1:1"), - ); - let expired_no_rt = GrokAuth { - key: "expired-legacy".into(), - auth_mode: AuthMode::WebLogin, - expires_at: Some(Utc::now() - chrono::Duration::hours(1)), - refresh_token: None, - ..GrokAuth::test_default() - }; - writer.update(expired_no_rt.clone()).await.unwrap(); - - let mgr = Arc::new(AuthManager::new(dir.path(), cfg.clone())); - mgr.hot_swap(expired_no_rt); - assert!(mgr.is_expired()); - - mgr.set_refresher(std::sync::Arc::new(AlwaysTransientRefresher)); - - // Force device explicitly so the assertion doesn't depend on ambient - // KIGI_LOGIN_DEVICE_FLOW / the real config file (the CLI override - // short-circuits the config read). - let result = run_auth_flow( - &mgr, - &cfg, - false, - None, - None, - None, - LoginTransportOverride::ForceDevice, - ) - .await; - - let err = result.unwrap_err(); - // Device flow fall-through hits the device-code endpoint (not OIDC - // discovery). - assert!( - err.to_string().contains("/oauth2/device/code"), - "expected device-code request error (proves flow fell through to interactive login), got: {err}" - ); - } - - #[test] - fn extract_url_from_external_provider_stderr() { - let extract = |input: &str| -> String { - input - .split_whitespace() - .find(|w| w.starts_with("https://")) - .map(|u| u.to_owned()) - .unwrap_or_else(|| input.to_owned()) - }; - - // Preamble text with URL - assert_eq!( - extract( - "Visit the following link to sign into Grok: https://auth.example.com/login?code=abc" - ), - "https://auth.example.com/login?code=abc" - ); - - // Multi-line with URL on second line - assert_eq!( - extract("Please sign in below\nhttps://auth.example.com/sso"), - "https://auth.example.com/sso" - ); - - // Just a bare URL - assert_eq!( - extract("https://auth.example.com/login"), - "https://auth.example.com/login" - ); - - // No URL at all — fallback to full content - assert_eq!(extract("some opaque output"), "some opaque output"); - } - - /// CLI `grok login` passes `on_stderr=None`; stderr must be inherited so - /// sign-in URLs appear in real time. Piped stderr with no reader deadlocks - /// once the child writes past the pipe buffer (~64 KiB). - #[tokio::test] - async fn external_provider_cli_path_does_not_deadlock_on_large_stderr() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new( - AuthManager::new(dir.path(), GrokComConfig::default()) - .with_proxy_base_url(&dead_proxy_url()), - ); - let cmd = r#"sh -c 'i=0; while [ $i -lt 2000 ]; do printf "%s" "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" >&2; i=$((i+1)); done; printf token'"#; - let (auth, _) = run_external_auth_provider(cmd, &mgr, false, None) - .await - .expect("CLI path must inherit stderr so large stderr does not deadlock"); - assert_eq!(auth.key, "token"); - } - struct AlwaysTransientRefresher; #[async_trait::async_trait] @@ -1902,54 +520,4 @@ mod tests { } } } - - /// Faithful reproduction of the cached-token bypass: the exact - /// repro JWT (wrong team) cached in `auth.json` under a pin, driven through - /// the same `AuthManager::new` + `auth()` engine `try_ensure_fresh_auth` - /// uses. Must be rejected and cleared; fails on the pre-fix tree. - #[tokio::test] - async fn noninteractive_auth_rejects_wrong_team_cached_token() { - // {"principal_id":"team-wrong","sub":"user-1"} — note: no principal_type. - const REPRO_JWT: &str = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJwcmluY2lwYWxfaWQiOiJ0ZWFtLXdyb25nIiwic3ViIjoidXNlci0xIn0.Signature"; - - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig { - force_login_team_uuid: Some(crate::auth::config::ForceLoginTeam::AnyOf(vec![ - "team-good".into(), - ])), - ..GrokComConfig::default() - }; - - // Persist the wrong-team session exactly as the repro's auth.json does. - let mut store = crate::auth::model::AuthStore::new(); - store.insert( - cfg.auth_scope(), - GrokAuth { - key: REPRO_JWT.into(), - auth_mode: AuthMode::Oidc, - team_id: Some("team-wrong".into()), - expires_at: chrono::DateTime::from_timestamp(9_999_999_999, 0), - ..GrokAuth::test_default() - }, - ); - let auth_path = dir.path().join("auth.json"); - crate::auth::storage::write_auth_json(&auth_path, &store).unwrap(); - - // Same engine as `try_ensure_fresh_auth`. - let auth_manager = Arc::new(AuthManager::new(dir.path(), cfg.clone())); - auth_manager.configure_refresher(cfg.auth_provider_command.clone()); - - assert!( - auth_manager.auth().await.is_err(), - "non-interactive auth must reject the wrong-team cached token" - ); - assert!( - auth_manager.current().is_none(), - "wrong-team token must not be usable via current()" - ); - assert!( - !auth_path.exists(), - "wrong-team auth.json must be cleared, forcing a compliant re-login" - ); - } } diff --git a/crates/codegen/kigi-shell/src/auth/jwt.rs b/crates/codegen/kigi-shell/src/auth/jwt.rs deleted file mode 100644 index bd77dd4..0000000 --- a/crates/codegen/kigi-shell/src/auth/jwt.rs +++ /dev/null @@ -1,45 +0,0 @@ -//! JWT expiration detection. Returns `None`/`false` for non-JWT tokens. - -use chrono::{DateTime, Duration, Utc}; -use serde::Deserialize; - -#[derive(Deserialize)] -struct Claims { - exp: Option, -} - -pub fn parse_jwt_expiration(token: &str) -> Option> { - jsonwebtoken::dangerous::insecure_decode::(token) - .ok() - .and_then(|data| data.claims.exp) - .and_then(|ts| DateTime::from_timestamp(ts, 0)) -} - -pub fn is_jwt_expired_or_near(token: &str, threshold: Duration) -> bool { - parse_jwt_expiration(token) - .map(|exp| exp <= Utc::now() + threshold) - .unwrap_or(false) -} - -#[cfg(test)] -mod tests { - use super::*; - - /// Tokens with an `aud` claim must parse successfully. - /// `jsonwebtoken::Validation::default()` enables audience validation which - /// silently rejects these tokens unless `validate_aud = false` is set. - #[test] - fn parses_jwt_with_aud_claim() { - let token = build_test_jwt(r#"{"aud":["some-audience"],"exp":1772575524}"#); - let exp = parse_jwt_expiration(&token); - assert_eq!(exp.unwrap().timestamp(), 1772575524); - } - - fn build_test_jwt(payload_json: &str) -> String { - use base64::Engine; - let enc = base64::engine::general_purpose::URL_SAFE_NO_PAD; - let header = enc.encode(r#"{"alg":"RS256","typ":"JWT"}"#); - let payload = enc.encode(payload_json); - format!("{header}.{payload}.fake-signature") - } -} diff --git a/crates/codegen/kigi-shell/src/auth/kimi_oauth.rs b/crates/codegen/kigi-shell/src/auth/kimi_oauth.rs new file mode 100644 index 0000000..d1d4659 --- /dev/null +++ b/crates/codegen/kigi-shell/src/auth/kimi_oauth.rs @@ -0,0 +1,626 @@ +//! Kimi Code OAuth wire protocol (PRD F1). +//! +//! Three calls against `{host}` (= `kigi_env::oauth_host()`), all +//! `application/x-www-form-urlencoded` POSTs carrying the device-identity +//! headers from [`super::device`]: +//! +//! - `POST /api/oauth/device_authorization` — form `client_id` +//! - `POST /api/oauth/token` (poll) — form `client_id` + `device_code` + +//! `grant_type=urn:ietf:params:oauth:grant-type:device_code` +//! - `POST /api/oauth/token` (refresh) — form `client_id` + +//! `grant_type=refresh_token` + `refresh_token`, with exponential backoff +//! over the retryable statuses {429, 500, 502, 503, 504} (3 tries) and +//! 401/403 mapped to [`RefreshError::Unauthorized`]. +//! +//! Ported from kimi-cli `auth/oauth.py` (the authoritative reference). + +use chrono::{Duration, Utc}; +use serde::Deserialize; + +use super::device::device_headers; +use super::model::{AuthMode, KimiAuth}; + +/// Kimi Code OAuth client id (fixed for the official device-flow client). +pub(crate) const KIMI_CODE_CLIENT_ID: &str = "17e5f671-d194-4dfb-9706-5516cb48c098"; + +const DEVICE_GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type:device_code"; +const REFRESH_GRANT_TYPE: &str = "refresh_token"; + +/// Refresh retry budget over the retryable statuses / network blips. +const MAX_REFRESH_RETRIES: u32 = 3; +/// HTTP statuses worth retrying a refresh for (kimi-cli parity). +const RETRYABLE_REFRESH_STATUSES: [u16; 5] = [429, 500, 502, 503, 504]; + +/// Result of `POST /api/oauth/device_authorization`. +#[derive(Debug, Clone)] +pub struct DeviceAuthorization { + pub user_code: String, + pub device_code: String, + /// Bare verification page (may be absent; the complete URI is required). + pub verification_uri: Option, + /// Verification page with the user code pre-filled — what we display + /// and open in the browser. + pub verification_uri_complete: String, + /// Device-code lifetime; `None` when the server omits it. + pub expires_in: Option, + /// Poll interval in seconds (server default 5; floored at 1 by callers). + pub interval: i64, +} + +#[derive(Deserialize)] +struct DeviceAuthorizationResponse { + user_code: String, + device_code: String, + #[serde(default)] + verification_uri: Option, + verification_uri_complete: String, + #[serde(default)] + expires_in: Option, + #[serde(default)] + interval: Option, +} + +/// Successful token payload (device grant and refresh grant share it). +#[derive(Debug, Deserialize)] +pub(crate) struct TokenResponse { + pub access_token: String, + pub refresh_token: String, + pub expires_in: i64, + #[serde(default)] + pub scope: Option, + #[serde(default)] + pub token_type: Option, +} + +impl TokenResponse { + /// Materialize the credential: `expires_at = now + expires_in`. + pub(crate) fn into_auth(self) -> KimiAuth { + let now = Utc::now(); + KimiAuth { + key: self.access_token, + auth_mode: AuthMode::OAuth, + create_time: now, + user_id: String::new(), + email: None, + refresh_token: Some(self.refresh_token), + expires_at: Some(now + Duration::seconds(self.expires_in)), + expires_in: Some(self.expires_in), + scope: self.scope, + token_type: self.token_type, + } + } +} + +#[derive(Deserialize, Default)] +struct OAuthErrorBody { + #[serde(default)] + error: Option, + #[serde(default)] + error_description: Option, +} + +/// One poll tick against the token endpoint. +#[derive(Debug)] +pub(crate) enum DevicePollResult { + /// 200 with an access token — login complete. + Success(Box), + /// `error == "expired_token"` — restart the whole device authorization. + Expired, + /// Any other non-200 outcome (`authorization_pending`, `slow_down`, + /// unknown errors) — wait and poll again. `slow_down` additionally bumps + /// the caller's interval. + Pending { + error: String, + description: Option, + }, +} + +fn oauth_url(host: &str, path: &str) -> String { + format!("{}{path}", host.trim_end_matches('/')) +} + +/// Attach the device-identity headers to a request. +fn with_device_headers( + mut builder: reqwest::RequestBuilder, +) -> anyhow::Result { + for (name, value) in device_headers()? { + builder = builder.header(name, value); + } + Ok(builder) +} + +/// Defend against control characters / non-https redirects from a +/// compromised or mis-configured OAuth host. +fn validate_verification_uri(uri: &str) -> anyhow::Result<()> { + if uri.chars().any(|c| c.is_ascii_control()) { + anyhow::bail!("Server returned invalid verification URI"); + } + let parsed = url::Url::parse(uri) + .map_err(|_| anyhow::anyhow!("Server returned invalid verification URI"))?; + match parsed.scheme() { + "https" => Ok(()), + "http" if matches!(parsed.host_str(), Some("localhost") | Some("127.0.0.1")) => Ok(()), + _ => anyhow::bail!("Server returned unsupported verification URI scheme"), + } +} + +/// `POST {host}/api/oauth/device_authorization` — start a device login. +pub(crate) async fn request_device_authorization( + host: &str, +) -> anyhow::Result { + let url = oauth_url(host, "/api/oauth/device_authorization"); + tracing::info!(url = %url, "auth: requesting device authorization"); + let resp = with_device_headers(crate::http::shared_client().post(&url))? + .form(&[("client_id", KIMI_CODE_CLIENT_ID)]) + .send() + .await?; + + let status = resp.status(); + if !status.is_success() { + let body = resp.text().await.unwrap_or_default(); + tracing::warn!(%status, "auth: device authorization failed"); + anyhow::bail!("Device authorization failed (HTTP {status}): {body}"); + } + let parsed: DeviceAuthorizationResponse = resp.json().await?; + + if !parsed + .user_code + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '-') + { + anyhow::bail!("Server returned invalid user_code format (expected [A-Z0-9-])"); + } + validate_verification_uri(&parsed.verification_uri_complete)?; + if let Some(ref uri) = parsed.verification_uri { + validate_verification_uri(uri)?; + } + + tracing::info!( + user_code = %parsed.user_code, + interval = parsed.interval.unwrap_or(5), + expires_in = ?parsed.expires_in, + "auth: device authorization issued" + ); + Ok(DeviceAuthorization { + user_code: parsed.user_code, + device_code: parsed.device_code, + verification_uri: parsed.verification_uri.filter(|u| !u.is_empty()), + verification_uri_complete: parsed.verification_uri_complete, + expires_in: parsed.expires_in.filter(|&e| e > 0), + interval: parsed.interval.unwrap_or(5), + }) +} + +/// One poll of `POST {host}/api/oauth/token` with the device grant. +/// +/// 5xx and network/decode failures are errors (kimi-cli parity: the login +/// loop surfaces them); everything else maps onto [`DevicePollResult`]. +pub(crate) async fn poll_device_token( + host: &str, + device_code: &str, +) -> anyhow::Result { + let url = oauth_url(host, "/api/oauth/token"); + let resp = with_device_headers(crate::http::shared_client().post(&url))? + .form(&[ + ("client_id", KIMI_CODE_CLIENT_ID), + ("device_code", device_code), + ("grant_type", DEVICE_GRANT_TYPE), + ]) + .send() + .await + .map_err(|e| anyhow::anyhow!("Token polling request failed: {e}"))?; + + let status = resp.status(); + if status.is_server_error() { + anyhow::bail!("Token polling server error: {status}"); + } + let body = resp.bytes().await?; + if status.is_success() { + if let Ok(tokens) = serde_json::from_slice::(&body) { + tracing::info!("auth: device poll succeeded, access token issued"); + return Ok(DevicePollResult::Success(Box::new(tokens.into_auth()))); + } + // 200 without an access token: treat as still-pending (kimi-cli + // requires "access_token" in the payload before accepting). + tracing::warn!("auth: device poll returned 200 without access_token; continuing"); + return Ok(DevicePollResult::Pending { + error: "missing_access_token".to_owned(), + description: None, + }); + } + let err: OAuthErrorBody = serde_json::from_slice(&body).unwrap_or_default(); + let error = err.error.unwrap_or_else(|| "unknown_error".to_owned()); + if error == "expired_token" { + tracing::info!("auth: device code expired; restarting device authorization"); + return Ok(DevicePollResult::Expired); + } + tracing::debug!(error = %error, "auth: device poll pending"); + Ok(DevicePollResult::Pending { + error, + description: err.error_description, + }) +} + +/// Why a refresh call terminally or transiently failed. +#[derive(Debug, thiserror::Error)] +pub(crate) enum RefreshError { + /// 401/403 — the refresh token was rejected. Triggers the tombstone + /// cooldown in the manager. + #[error("token refresh unauthorized (HTTP {status}): {description}")] + Unauthorized { status: u16, description: String }, + /// Non-retryable non-200 status. + #[error("token refresh failed (HTTP {status}): {description}")] + Fatal { status: u16, description: String }, + /// Retry budget exhausted over retryable statuses / network blips. + #[error("token refresh failed after {MAX_REFRESH_RETRIES} attempts: {last_error}")] + Exhausted { last_error: String }, + /// Local failure before the wire (e.g. device-id creation failed). + #[error(transparent)] + Local(#[from] anyhow::Error), +} + +/// `POST {host}/api/oauth/token` with `grant_type=refresh_token`. +/// +/// Retries the retryable statuses and network errors with exponential +/// backoff (`2^attempt` seconds); 401/403 returns immediately as +/// [`RefreshError::Unauthorized`]. +pub(crate) async fn refresh_token( + host: &str, + refresh_token: &str, +) -> Result { + let url = oauth_url(host, "/api/oauth/token"); + let mut last_error = String::from("no attempt made"); + for attempt in 0..MAX_REFRESH_RETRIES { + if attempt > 0 { + let backoff = std::time::Duration::from_secs(1 << (attempt - 1)); + tracing::warn!( + attempt, + backoff_secs = backoff.as_secs(), + last_error = %last_error, + "auth: retrying token refresh" + ); + tokio::time::sleep(backoff).await; + } + tracing::info!(attempt, "auth: token refresh attempt"); + let send_result = with_device_headers(crate::http::shared_client().post(&url))? + .form(&[ + ("client_id", KIMI_CODE_CLIENT_ID), + ("grant_type", REFRESH_GRANT_TYPE), + ("refresh_token", refresh_token), + ]) + .send() + .await; + + let resp = match send_result { + Ok(resp) => resp, + Err(e) => { + last_error = format!("network error: {e}"); + continue; + } + }; + let status = resp.status().as_u16(); + let body = resp.bytes().await.unwrap_or_default(); + if status == 401 || status == 403 { + let err: OAuthErrorBody = serde_json::from_slice(&body).unwrap_or_default(); + return Err(RefreshError::Unauthorized { + status, + description: err + .error_description + .unwrap_or_else(|| "Token refresh unauthorized.".to_owned()), + }); + } + if status == 200 { + return match serde_json::from_slice::(&body) { + Ok(tokens) => Ok(tokens.into_auth()), + Err(e) => Err(RefreshError::Fatal { + status, + description: format!("malformed token payload: {e}"), + }), + }; + } + let err: OAuthErrorBody = serde_json::from_slice(&body).unwrap_or_default(); + let description = err + .error_description + .unwrap_or_else(|| format!("Token refresh failed (HTTP {status}).")); + if RETRYABLE_REFRESH_STATUSES.contains(&status) { + last_error = description; + continue; + } + return Err(RefreshError::Fatal { + status, + description, + }); + } + Err(RefreshError::Exhausted { last_error }) +} + +#[cfg(test)] +mod tests { + use super::*; + use wiremock::matchers::{body_string_contains, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + fn token_json(access: &str, refresh: &str) -> serde_json::Value { + serde_json::json!({ + "access_token": access, + "refresh_token": refresh, + "expires_in": 3600, + "scope": "kimi-code", + "token_type": "bearer", + }) + } + + #[tokio::test] + async fn device_authorization_parses_wire_payload() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/device_authorization")) + .and(body_string_contains(format!( + "client_id={KIMI_CODE_CLIENT_ID}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "user_code": "ABCD-1234", + "device_code": "dev-code-1", + "verification_uri": "https://auth.kimi.com/device", + "verification_uri_complete": "https://auth.kimi.com/device?code=ABCD-1234", + "expires_in": 600, + "interval": 7, + }))) + .expect(1) + .mount(&server) + .await; + + let auth = request_device_authorization(&server.uri()).await.unwrap(); + assert_eq!(auth.user_code, "ABCD-1234"); + assert_eq!(auth.device_code, "dev-code-1"); + assert_eq!(auth.interval, 7); + assert_eq!(auth.expires_in, Some(600)); + assert_eq!( + auth.verification_uri_complete, + "https://auth.kimi.com/device?code=ABCD-1234" + ); + } + + #[tokio::test] + async fn device_authorization_sends_device_headers() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/device_authorization")) + .and(wiremock::matchers::header_exists("X-Msh-Device-Name")) + .and(wiremock::matchers::header_exists("X-Msh-Device-Model")) + .and(wiremock::matchers::header_exists("X-Msh-Device-Id")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "user_code": "AAAA", + "device_code": "d", + "verification_uri_complete": "https://auth.kimi.com/device?code=AAAA", + "interval": 5, + }))) + .expect(1) + .mount(&server) + .await; + request_device_authorization(&server.uri()).await.unwrap(); + } + + #[tokio::test] + async fn device_authorization_defaults_interval_to_five() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/device_authorization")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "user_code": "AAAA", + "device_code": "d", + "verification_uri_complete": "https://auth.kimi.com/device?code=AAAA", + }))) + .mount(&server) + .await; + let auth = request_device_authorization(&server.uri()).await.unwrap(); + assert_eq!(auth.interval, 5); + assert_eq!(auth.expires_in, None); + } + + #[tokio::test] + async fn device_authorization_rejects_bad_verification_uri() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/device_authorization")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "user_code": "AAAA", + "device_code": "d", + "verification_uri_complete": "javascript:alert(1)", + }))) + .mount(&server) + .await; + let err = request_device_authorization(&server.uri()) + .await + .unwrap_err(); + assert!(err.to_string().contains("verification URI"), "{err}"); + } + + #[tokio::test] + async fn device_authorization_surfaces_http_errors() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/device_authorization")) + .respond_with(ResponseTemplate::new(400).set_body_string("nope")) + .mount(&server) + .await; + let err = request_device_authorization(&server.uri()) + .await + .unwrap_err(); + assert!(err.to_string().contains("HTTP 400"), "{err}"); + } + + #[tokio::test] + async fn poll_success_builds_auth_with_expiry() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .and(body_string_contains("grant_type=urn")) + .and(body_string_contains("device_code=dev-1")) + .respond_with(ResponseTemplate::new(200).set_body_json(token_json("at-1", "rt-1"))) + .expect(1) + .mount(&server) + .await; + let result = poll_device_token(&server.uri(), "dev-1").await.unwrap(); + let DevicePollResult::Success(auth) = result else { + panic!("expected success, got {result:?}"); + }; + assert_eq!(auth.key, "at-1"); + assert_eq!(auth.refresh_token.as_deref(), Some("rt-1")); + assert_eq!(auth.expires_in, Some(3600)); + let remaining = auth.expires_at.unwrap() - chrono::Utc::now(); + assert!( + (3590..=3600).contains(&remaining.num_seconds()), + "expires_at must be ~now+expires_in, got {remaining:?}" + ); + assert_eq!(auth.scope.as_deref(), Some("kimi-code")); + assert_eq!(auth.token_type.as_deref(), Some("bearer")); + } + + #[tokio::test] + async fn poll_maps_expired_token_to_restart() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with( + ResponseTemplate::new(400) + .set_body_json(serde_json::json!({ "error": "expired_token" })), + ) + .mount(&server) + .await; + let result = poll_device_token(&server.uri(), "dev-1").await.unwrap(); + assert!(matches!(result, DevicePollResult::Expired), "{result:?}"); + } + + #[tokio::test] + async fn poll_maps_pending_and_unknown_errors_to_pending() { + let server = MockServer::start().await; + for error in ["authorization_pending", "slow_down", "surprise_error"] { + server.reset().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with( + ResponseTemplate::new(400).set_body_json(serde_json::json!({ "error": error })), + ) + .mount(&server) + .await; + let result = poll_device_token(&server.uri(), "dev-1").await.unwrap(); + match result { + DevicePollResult::Pending { error: got, .. } => assert_eq!(got, error), + other => panic!("expected pending for {error}, got {other:?}"), + } + } + } + + #[tokio::test] + async fn poll_server_error_is_an_error() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with(ResponseTemplate::new(502)) + .mount(&server) + .await; + let err = poll_device_token(&server.uri(), "dev-1").await.unwrap_err(); + assert!(err.to_string().contains("server error"), "{err}"); + } + + #[tokio::test] + async fn refresh_success_round_trip() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains("refresh_token=rt-old")) + .respond_with(ResponseTemplate::new(200).set_body_json(token_json("at-new", "rt-new"))) + .expect(1) + .mount(&server) + .await; + let auth = refresh_token(&server.uri(), "rt-old").await.unwrap(); + assert_eq!(auth.key, "at-new"); + assert_eq!(auth.refresh_token.as_deref(), Some("rt-new")); + } + + #[tokio::test] + async fn refresh_401_maps_to_unauthorized_without_retry() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with( + ResponseTemplate::new(401).set_body_json( + serde_json::json!({ "error_description": "refresh token revoked" }), + ), + ) + .expect(1) + .mount(&server) + .await; + let err = refresh_token(&server.uri(), "rt-dead").await.unwrap_err(); + match err { + RefreshError::Unauthorized { + status, + description, + } => { + assert_eq!(status, 401); + assert_eq!(description, "refresh token revoked"); + } + other => panic!("expected Unauthorized, got {other:?}"), + } + } + + #[tokio::test] + async fn refresh_retries_retryable_status_then_succeeds() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with(ResponseTemplate::new(503)) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(token_json("at-2", "rt-2"))) + .expect(1) + .mount(&server) + .await; + let auth = refresh_token(&server.uri(), "rt-old").await.unwrap(); + assert_eq!(auth.key, "at-2"); + } + + #[tokio::test] + async fn refresh_exhausts_after_three_retryable_failures() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with(ResponseTemplate::new(500)) + .expect(3) + .mount(&server) + .await; + let err = refresh_token(&server.uri(), "rt-old").await.unwrap_err(); + assert!(matches!(err, RefreshError::Exhausted { .. }), "{err:?}"); + } + + #[tokio::test] + async fn refresh_non_retryable_status_is_fatal_without_retry() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with( + ResponseTemplate::new(400) + .set_body_json(serde_json::json!({ "error_description": "bad request" })), + ) + .expect(1) + .mount(&server) + .await; + let err = refresh_token(&server.uri(), "rt-old").await.unwrap_err(); + match err { + RefreshError::Fatal { + status, + description, + } => { + assert_eq!(status, 400); + assert_eq!(description, "bad request"); + } + other => panic!("expected Fatal, got {other:?}"), + } + } +} diff --git a/crates/codegen/kigi-shell/src/auth/manager.rs b/crates/codegen/kigi-shell/src/auth/manager.rs index 641d413..0ed8796 100644 --- a/crates/codegen/kigi-shell/src/auth/manager.rs +++ b/crates/codegen/kigi-shell/src/auth/manager.rs @@ -1,8 +1,8 @@ //! `AuthManager` -- single source of truth for `auth.json` + the //! in-memory bearer cache. Mutations go through `refresh_chain` or `update`; lock -//! and enrichment helpers live in submodules. +//! and lock/sleep-gate helpers live in submodules. -use chrono::{Duration, Utc}; +use chrono::Duration; use parking_lot::RwLock; use std::path::{Path, PathBuf}; use std::sync::Arc; @@ -10,8 +10,6 @@ use std::time::Duration as StdDuration; use tokio_util::sync::CancellationToken; -#[path = "manager/enrichment.rs"] -mod enrichment; #[path = "manager/lock.rs"] mod lock; #[path = "manager/sleep_gate.rs"] @@ -20,31 +18,22 @@ mod sleep_gate; use lock::try_lock_auth_file_async; use sleep_gate::{GateRaise, InFlightGuard, SleepGate}; -use crate::auth::config::GrokComConfig; +use crate::auth::config::{KIMI_CODE_OAUTH_SCOPE, KimiCodeConfig}; use crate::auth::error::AuthError; use crate::auth::token_type::TokenType; -#[cfg(test)] -use super::model::UserInfo; -use super::model::{ - AuthMode, GrokAuth, early_invalidation, is_expired, is_expired_with_buffer, lookup_auth, - token_suffix, -}; +use super::model::{KimiAuth, is_expired, is_expired_with_buffer, lookup_auth, token_suffix}; use super::refresh::{RefreshOutcome, TokenRefresher, resolve_refresh_credential}; use super::storage::{ - AuthFileLock, read_auth_json, read_auth_json_or_empty_recovering_corrupt, write_auth_json, + AuthFileLock, KeyringRead, keyring_delete_session, keyring_enabled, keyring_read_session, + keyring_write_session, read_auth_json, read_auth_json_or_empty_recovering_corrupt, + write_auth_json, }; -#[cfg(test)] -use super::storage::read_auth_json_or_empty; -#[cfg(test)] -use chrono::DateTime; -#[cfg(test)] -use enrichment::apply_user_info_enrichment; - #[cfg(test)] use super::model::AuthStore; -use super::model::LEGACY_SCOPE; +#[cfg(test)] +use super::storage::read_auth_json_or_empty; /// Why a token refresh is being requested. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -56,7 +45,7 @@ pub(crate) enum RefreshReason { } /// Timeout for acquiring the advisory `auth.json.lock` file lock. -/// Used by advisory (non-critical) lock sites: `flow.rs`, `enrichment.rs`, +/// Used by advisory (non-critical) lock sites: `flow.rs`, /// `recovery.rs`. pub(crate) const AUTH_LOCK_TIMEOUT: StdDuration = StdDuration::from_secs(10); @@ -66,20 +55,18 @@ pub(crate) const AUTH_LOCK_TIMEOUT: StdDuration = StdDuration::from_secs(10); /// wait for the leader to finish rather than timing out and retrying. const REFRESH_LOCK_TIMEOUT: StdDuration = StdDuration::from_secs(45); -/// Long poll interval used by the proactive refresh task when no -/// productive refresh is possible (see [`compute_proactive_sleep`]). -/// Long enough to avoid CPU/log spam; short enough that a `hot_swap()` -/// or `configure_refresher()` is picked up in a reasonable window. -pub(crate) const BACKOFF_INTERVAL: StdDuration = StdDuration::from_secs(300); +/// Fixed cadence of the background refresh check (PRD F1: every 60s). +pub(crate) const PROACTIVE_REFRESH_INTERVAL: StdDuration = StdDuration::from_secs(60); + +/// A tick whose wall-clock gap exceeds `interval × this` indicates the +/// machine slept through timer ticks; the next tick forces a refresh +/// (kimi-cli `refreshing()` parity). +const SLEEP_WAKE_FORCE_FACTOR: u32 = 2; /// How long to wait after a file lock timeout before re-reading disk, /// giving the lock holder time to finish writing. const LOCK_TIMEOUT_WAIT: StdDuration = StdDuration::from_secs(2); -/// Maximum random jitter (seconds) added to the proactive refresh sleep -/// to stagger sibling processes and avoid thundering-herd IdP calls. -const JITTER_RANGE_SECS: i64 = 60; - /// `force_reload_from_disk` re-read budget. A single `auth.json` read can /// return `NotFound`/unreadable for reasons unrelated to logout — most /// notably the first read right after wake-from-sleep, where the filesystem @@ -92,28 +79,29 @@ const RELOAD_RETRY_TRIES: usize = 3; /// Only paid on the disk-anomaly branch, never on a healthy read. const RELOAD_RETRY_BACKOFF: StdDuration = StdDuration::from_millis(50); -/// Sticky permanent-refresh verdict, scoped to the credential that produced it -/// (`token_key`). The scope is what makes invalidation automatic: any other -/// credential reads through as "no failure", so no manual clearing is needed. +/// Refresh-rejection tombstone (PRD F1), scoped to the refresh-token value +/// the OAuth host rejected (`refresh_token_key`). The scope is what makes +/// invalidation automatic: once the persisted refresh token differs (another +/// process rotated it, or a fresh login landed), the tombstone reads through +/// as "no failure" without manual clearing. struct ScopedRefreshFailure { - token_key: String, + /// The rejected refresh-token value. + refresh_token_key: String, error: crate::auth::error::RefreshTokenFailedError, - /// Two-clock timestamp (see [`GateRaise`]): the TTL below is *real* time, - /// so it must keep counting across a system sleep. The monotonic clock - /// pauses during suspend — with it alone, a failure cached just before - /// sleep would still short-circuit `auth()` for a further + /// Two-clock timestamp (see [`GateRaise`]): the cooldown below is *real* + /// time, so it must keep counting across a system sleep. The monotonic + /// clock pauses during suspend — with it alone, a failure cached just + /// before sleep would still short-circuit `auth()` for a further /// [`PERMANENT_FAILURE_TTL`] of *awake* time after wake, exactly when the /// user comes back and expects a recovered session. recorded_at: GateRaise, } -/// Auto-expiry safety net for the recoverable reasons (`ClientRejected`, -/// `Other`): they self-heal without re-login even if the credential never -/// changes. `RefreshTokenRejected` is excluded (see `is_sticky`). Independent -/// of `BACKOFF_INTERVAL` (equal value is coincidental). Measured on both -/// clocks — expires once *either* the monotonic or the wall clock passes the -/// bound, so it means "5 real minutes", not "5 awake minutes" (a suspend -/// doesn't extend it). +/// Tombstone cooldown (PRD F1): a rejected refresh token is not re-sent for +/// this long; afterwards a retry is allowed (the server stays the authority). +/// Measured on both clocks — expires once *either* the monotonic or the wall +/// clock passes the bound, so it means "5 real minutes", not "5 awake +/// minutes" (a suspend doesn't extend it). const PERMANENT_FAILURE_TTL: StdDuration = StdDuration::from_secs(300); /// Single source of truth for `auth.json` + the in-memory bearer. @@ -121,19 +109,17 @@ const PERMANENT_FAILURE_TTL: StdDuration = StdDuration::from_secs(300); /// Lock order: `refresh_lock` (async) -> the sync locks (`inner` / `refresher` /// / `permanent_failure`), never co-held; `permanent_failure()` /// reads `permanent_failure` first and only then `inner` (via -/// `attempted_verdict_key`, when a verdict is stored), never co-held. Never hold +/// `attempted_tombstone_key`, when a tombstone is stored), never co-held. Never hold /// a `parking_lot` guard across `.await`. Refreshers return [`RefreshOutcome`] /// for `refresh_chain` to apply. pub struct AuthManager { /// In-memory bearer. Mutate via [`Self::with_inner_write`] or /// [`Self::refresh_chain`]; the closure helpers' sync return type - /// enforces "no `.await` while holding the lock". `Arc` - /// so the spawned `/user` enrichment task can write back. - inner: Arc>>, + /// enforces "no `.await` while holding the lock". + inner: Arc>>, path: PathBuf, scope: String, - grok_com_config: GrokComConfig, - proxy_base_url: String, + kimi_code_config: KimiCodeConfig, refresher: RwLock>>, /// Idempotency guard for `configure_refresher` so double-calls /// don't reset internal state. @@ -190,11 +176,6 @@ pub struct AuthManager { /// without a real macOS dark wake. `None` = consult the OS. #[cfg(test)] dark_wake_override: parking_lot::Mutex>, - /// Test-only override for [`AuthManager::is_devbox_environment`]. CI runs in - /// K8s pods where the real check is `true`, which would otherwise let - /// `DevboxRecovery` adopt a seeded valid token; `Some(_)` pins the result. - #[cfg(test)] - devbox_override: parking_lot::Mutex>, } /// Discriminated outcome of a disk read, for transition logging. @@ -249,16 +230,14 @@ impl ScopeRemoval { /// it without refreshing. enum LockOutcome { Held(AuthFileLock), - Adopted(Box), + Adopted(Box), } // ── Construction + builders ────────────────────────────────────────── impl AuthManager { - pub fn new(kigi_home: &Path, grok_com_config: GrokComConfig) -> Self { - let scope = grok_com_config.auth_scope(); - let proxy_base_url = - crate::agent::config::EndpointsConfig::from_effective_config().proxy_url(); + pub fn new(kigi_home: &Path, kimi_code_config: KimiCodeConfig) -> Self { + let scope = kimi_code_config.auth_scope(); kigi_log::unified_log::info( "AuthManager::new", @@ -270,18 +249,18 @@ impl AuthManager { "KIGI_SHARE_DIR": std::env::var("KIGI_SHARE_DIR").unwrap_or_else(|_| "(unset)".into()), "KIGI_AUTH_PATH": std::env::var("KIGI_AUTH_PATH").unwrap_or_else(|_| "(unset)".into()), "KIGI_AUTH": std::env::var("KIGI_AUTH").map(|_| "(set)".to_string()).unwrap_or_else(|_| "(unset)".into()), + "keyring_enabled": keyring_enabled(), })), ); // KIGI_AUTH: inline JSON credentials (highest priority, read-only). if let Ok(inline_json) = std::env::var("KIGI_AUTH") { - if let Ok(auth) = serde_json::from_str::(&inline_json) { + if let Ok(auth) = serde_json::from_str::(&inline_json) { return Self::assemble( Some(auth), kigi_home.join("auth.json"), scope, - grok_com_config, - proxy_base_url, + kimi_code_config, None, ); } @@ -293,30 +272,30 @@ impl AuthManager { .map(PathBuf::from) .unwrap_or_else(|_| kigi_home.join("auth.json")); + // Keyring first (PRD F1): the session credential's primary store. + if scope == KIMI_CODE_OAUTH_SCOPE + && let KeyringRead::Found(auth) = keyring_read_session() + { + kigi_log::unified_log::info( + "AuthManager::new loaded session from system keyring", + None, + Some(serde_json::json!({ + "key_prefix": token_suffix(&auth.key), + "is_expired": is_expired(&auth), + })), + ); + return Self::assemble( + Some(*auth), + path, + scope, + kimi_code_config, + Some(DiskAuthState::Ok), + ); + } + let (auth, auth_read_detail, initial_disk_state) = match read_auth_json(&path) { Ok(map) => { let found = lookup_auth(&map, &scope); - // If lookup_auth skipped a legacy WebLogin token, remove the - // stale scope entry from auth.json so it is not re-evaluated - // on every launch. - if found.is_none() - && map - .get(LEGACY_SCOPE) - .is_some_and(|a| a.auth_mode == AuthMode::WebLogin) - { - // Best-effort cleanup under advisory lock (consistent with - // other auth.json writers). Non-blocking: if the lock is - // held by a concurrent process, skip — retried next launch. - if let Some(_lock) = lock::try_lock_auth_file_nonblocking(&path) { - let mut cleaned = map.clone(); - cleaned.remove(LEGACY_SCOPE); - let _ = write_auth_json(&path, &cleaned); - tracing::debug!("auth: removed stale WebLogin scope from auth.json"); - // lock released on drop - } else { - tracing::debug!("auth: skipped WebLogin cleanup (lock unavailable)"); - } - } let detail = serde_json::json!({ "read": "ok", "resolved_path": path.display().to_string(), @@ -355,18 +334,13 @@ impl AuthManager { Some(auth_read_detail), ); - let manager = Self::assemble( + Self::assemble( auth, path, scope, - grok_com_config, - proxy_base_url, + kimi_code_config, Some(initial_disk_state), - ); - // Clear a wrong-team session left on disk before the pin was deployed, - // so the first launch forces a compliant login. - manager.enforce_pin_on_loaded_token(); - manager + ) } /// Single field-assembly point for [`Self::new`]'s two construction paths @@ -374,19 +348,17 @@ impl AuthManager { /// threaded fields. One literal means a newly added field can't be silently /// dropped from one branch. fn assemble( - inner: Option, + inner: Option, path: PathBuf, scope: String, - grok_com_config: GrokComConfig, - proxy_base_url: String, + kimi_code_config: KimiCodeConfig, disk_state: Option, ) -> Self { Self { inner: Arc::new(RwLock::new(inner)), path, scope, - grok_com_config, - proxy_base_url, + kimi_code_config, refresher: RwLock::new(None), refresher_configured: std::sync::atomic::AtomicBool::new(false), proactive_started: std::sync::atomic::AtomicBool::new(false), @@ -407,28 +379,9 @@ impl AuthManager { dark_wake_defer_since: parking_lot::RwLock::new(None), #[cfg(test)] dark_wake_override: parking_lot::Mutex::new(None), - #[cfg(test)] - devbox_override: parking_lot::Mutex::new(None), } } - /// Clear the disk-loaded token if it violates the team pin (startup only; - /// the read/dispense gates cover everything cached afterwards). - fn enforce_pin_on_loaded_token(&self) { - let loaded = self.inner.read().clone(); - if let Some(auth) = loaded - && let Some(e) = self.cached_token_policy_error(&auth) - { - self.reject_and_clear(&e); - } - } - - /// Override the proxy base URL (precedence over env var). - pub(crate) fn with_proxy_base_url(mut self, url: &str) -> Self { - self.proxy_base_url = url.to_owned(); - self - } - // ── State mutation (clear, hot_swap, update) ────────────────────── pub(crate) fn clear(&self) -> std::io::Result<()> { @@ -447,6 +400,14 @@ impl AuthManager { } fn remove_scope_impl(&self, scope: &str) -> std::io::Result<()> { + // Session credentials also live in the system keyring (primary + // store); drop that copy first so a file-side failure can't leave + // the token behind. + if scope == KIMI_CODE_OAUTH_SCOPE + && let Err(e) = keyring_delete_session() + { + tracing::warn!(error = %e, "auth: failed to remove session credential from keyring"); + } let disk_mutation = if let Some(_lock) = lock::try_lock_auth_file_nonblocking(&self.path) { self.write_scope_removal(scope)? // lock released on drop } else { @@ -531,16 +492,12 @@ impl AuthManager { // of the same broken refresh_token. DiskAuthState::Ok => { *self.inner.write() = auth; - // A re-read (e.g. relay reconnect) can adopt a wrong-team - // token a sibling wrote; clear it here, mirroring `new()`. - self.enforce_pin_on_loaded_token(); return; } // File readable, our scope genuinely absent: the trustworthy // logout / scope-removed signal. DiskAuthState::EntryMissing => { self.drop_in_memory_credentials("scope absent on readable auth.json"); - self.enforce_pin_on_loaded_token(); return; } // Disk anomaly: a transient (e.g. wake-time ENOENT) heals on a @@ -573,7 +530,6 @@ impl AuthManager { "disk anomaly; no live refresh token to retain (missing RT or permanent failure)", ); } - self.enforce_pin_on_loaded_token(); } /// Drop the in-memory credentials, loudly. Logs the discard (with `reason`) @@ -605,89 +561,26 @@ impl AuthManager { // | "have credentials at all?" | `is_expired()` | // | bypass memory, read disk | `read_disk_auth()` | - // ── Login-policy enforcement (force_login_team_uuid) ────────────── - // - // The team pin is enforced wherever the manager hands out a session token, - // not only on fresh login: sync reads (`current`/`expired_auth`) hide a - // violating token; the async gates (`auth`, recovery) and `new()` also - // clear `auth.json` to force a compliant re-login. - - /// `Some(error)` when a `force_login_team_uuid` pin is set and the token's - /// team principal isn't allowed; `None` when compliant or unpinned. - /// - /// Reads the principal from the token's own (unverified) JWT claim — - /// fail-fast defense-in-depth, not the security boundary (the server is - /// authoritative). An API-key session is rejected under the kill switch, - /// else allowed. - pub(crate) fn cached_token_policy_error(&self, auth: &GrokAuth) -> Option { - if auth.auth_mode == AuthMode::ApiKey { - // Else enforce_disable_api_key_auth swaps the key for itself (no-op). - return self - .grok_com_config - .api_key_auth_disabled() - .then_some(AuthError::ApiKeyAuthDisabled); - } - let policy = crate::auth::oidc::login_principal_policy(&self.grok_com_config)?; - let actual = crate::auth::oidc::peek_access_token_principal_id(&auth.key); - crate::auth::oidc::enforce_login_principal(Some(&policy), actual.as_deref()) - .err() - .map(|e| AuthError::PinnedTeamMismatch { - message: e.to_string(), - }) - } - - /// Log and clear a policy-violating session (disk + memory) so the next - /// launch forces a fresh, compliant login. - pub(crate) fn reject_and_clear(&self, error: &AuthError) { - let policy = match error { - AuthError::PinnedTeamMismatch { .. } => "team_pin", - AuthError::ApiKeyAuthDisabled => "api_key_disabled", - _ => "login_policy", - }; - kigi_log::unified_log::warn( - "auth: cached session rejected by login policy; clearing", - None, - Some(serde_json::json!({ "policy": policy, "reason": error.to_string() })), - ); - if let Err(e) = self.clear() { - tracing::warn!(error = %e, "auth: failed to clear policy-violating session"); - } - } - - /// Hide a cached token rejected by the login policy. No clear here (keeps - /// the sync read path lock-free); `auth()`/recovery/`new()` do the clearing. - fn vet_cached(&self, auth: GrokAuth) -> Option { - match self.cached_token_policy_error(&auth) { - None => Some(auth), - Some(e) => { - tracing::debug!(error = %e, "auth: hiding cached session rejected by login policy"); - None - } - } - } - - /// Cached in-memory token if outside the early-invalidation buffer. - pub(crate) fn current(&self) -> Option { - let auth = self - .inner + /// Cached in-memory token if outside the refresh-threshold buffer. + pub(crate) fn current(&self) -> Option { + self.inner .read() .as_ref() .filter(|a| !self.is_token_expired(a)) - .cloned()?; - self.vet_cached(auth) + .cloned() } /// Closure-scoped write. Sync return type prevents `.await` while /// the lock is held. Prefer this over `self.inner.write()`. #[inline] - pub(crate) fn with_inner_write(&self, f: impl FnOnce(&mut Option) -> R) -> R { + pub(crate) fn with_inner_write(&self, f: impl FnOnce(&mut Option) -> R) -> R { let mut guard = self.inner.write(); f(&mut guard) } /// Closure-scoped read counterpart to [`Self::with_inner_write`]. #[inline] - pub(crate) fn with_inner_read(&self, f: impl FnOnce(Option<&GrokAuth>) -> R) -> R { + pub(crate) fn with_inner_read(&self, f: impl FnOnce(Option<&KimiAuth>) -> R) -> R { let guard = self.inner.read(); f(guard.as_ref()) } @@ -700,89 +593,80 @@ impl AuthManager { .is_some_and(|a| self.is_token_expired(a)) } - /// In-memory bearer regardless of the early-invalidation buffer. + /// In-memory bearer regardless of the refresh-threshold buffer. /// Prefer [`Self::auth`] when `.await` is available. - pub(crate) fn current_or_expired(&self) -> Option { + pub(crate) fn current_or_expired(&self) -> Option { self.current().or_else(|| self.expired_auth()) } - /// `true` when data collection must be suppressed — the team has ZDR or - /// the user opted out of coding data retention. Reads - /// [`Self::current_or_expired`] because neither flag changes on token - /// expiry and `current()` returns `None` during the refresh window. - /// - /// Fail-open: no credential ⇒ `false` (not disabled). Collection paths - /// that must not act on unknown privacy state should use the fail-closed - /// [`Self::allows_data_collection`] instead. - pub(crate) fn is_data_collection_disabled(&self) -> bool { - self.current_or_expired() - .is_some_and(|a| a.is_data_collection_disabled()) - } - - /// Fail-closed collection predicate: `true` only when a credential - /// exists and carries no ZDR / retention-opt-out flag. Missing or - /// cleared auth (e.g. after a mid-session `/logout`) counts as - /// disabled — nothing may leave the machine while the privacy state is - /// unknown. - pub(crate) fn allows_data_collection(&self) -> bool { - self.current_or_expired() - .is_some_and(|a| !a.is_data_collection_disabled()) - } - /// Expired in-memory entry (for its `refresh_token`). - pub(crate) fn expired_auth(&self) -> Option { - let auth = self - .inner + pub(crate) fn expired_auth(&self) -> Option { + self.inner .read() .as_ref() .filter(|a| self.is_token_expired(a)) - .cloned()?; - self.vet_cached(auth) + .cloned() } - /// Expiry policy: `expires_at - early_invalidation` if present; - /// `External` with `auth_token_ttl` -> `create_time + ttl`; - /// fallback `create_time + 30d` (WebLogin-style). - fn is_token_expired(&self, auth: &GrokAuth) -> bool { - self.token_expired_with_buffer(auth, early_invalidation()) + /// Expiry policy (PRD F1): expiring-soon once the remaining lifetime + /// drops below `max(300, expires_in × 0.5)` seconds; credentials without + /// `expires_at` fall back to `create_time + 30d`. + fn is_token_expired(&self, auth: &KimiAuth) -> bool { + is_expired(auth) } - /// Actual (hard) expiry: the instant the proxy would actually reject the - /// token, with no early-invalidation margin. The export gate + /// Actual (hard) expiry: the instant the server would actually reject the + /// token, with no refresh-threshold margin. The export gate /// ([`Self::has_usable_token`]) uses this instead of [`Self::is_token_expired`] - /// because a token still inside the buffer is sent — and accepted — on the - /// wire via `current_or_expired()`, so it must not count as unusable. - fn is_token_hard_expired(&self, auth: &GrokAuth) -> bool { - self.token_expired_with_buffer(auth, Duration::zero()) + /// because a token still inside the threshold is sent — and accepted — on + /// the wire via `current_or_expired()`, so it must not count as unusable. + fn is_token_hard_expired(&self, auth: &KimiAuth) -> bool { + is_expired_with_buffer(auth, Duration::zero()) } - fn token_expired_with_buffer(&self, auth: &GrokAuth, buffer: Duration) -> bool { - if auth.expires_at.is_some() { - return is_expired_with_buffer(auth, buffer); - } - if auth.auth_mode == AuthMode::External - && let Some(ttl) = self.grok_com_config.auth_token_ttl - { - let age = Utc::now().signed_duration_since(auth.create_time); - return age >= Duration::seconds(ttl as i64) - buffer; - } - is_expired_with_buffer(auth, buffer) - } + // ── Persistence ─────────────────────────────────────────────────── - // ── Persistence + enrichment ────────────────────────────────────── - - /// Persist rotated tokens to disk + cache, then spawn `/user` enrichment. + /// Persist rotated tokens (keyring → file fallback) + cache. /// /// Invariants: - /// - **Disk write before any network I/O** (else a sibling process can - /// reuse the not-yet-rotated RT and the IdP returns `invalid_grant`). + /// - **Persist before any further network I/O** (else a sibling process + /// can reuse the not-yet-rotated RT and the OAuth host rejects it). /// - **Caller holds the `auth.json` file lock** (production callers: /// `refresh_chain` Success arm, `flow::run_auth_flow`). - /// - /// Returns the input `GrokAuth` BEFORE enrichment lands; callers - /// needing the post-enrichment view re-read `current()`. - pub(crate) async fn update(self: &Arc, auth: GrokAuth) -> std::io::Result { + pub(crate) async fn update(self: &Arc, auth: KimiAuth) -> std::io::Result { let update_started = std::time::Instant::now(); + + // Keyring first (PRD F1): the session credential's primary store. + if self.scope == KIMI_CODE_OAUTH_SCOPE && keyring_enabled() { + match keyring_write_session(&auth) { + Ok(()) => { + let elapsed_ms = update_started.elapsed().as_millis() as u64; + kigi_log::unified_log::info( + "auth update written to system keyring", + None, + Some(serde_json::json!({ + "rt_prefix": auth.refresh_token.as_deref().map(token_suffix), + "key_prefix": token_suffix(&auth.key), + "elapsed_ms": elapsed_ms, + })), + ); + // Drop any stale plaintext copy left from a fallback-era + // write so the two stores can't diverge. + self.strip_scope_from_file_best_effort(); + self.with_inner_write(|inner| *inner = Some(auth.clone())); + return Ok(auth); + } + Err(e) => { + tracing::warn!(error = %e, "auth: keyring write failed, falling back to file"); + kigi_log::unified_log::warn( + "auth update keyring write failed, using file fallback", + None, + Some(serde_json::json!({ "error": e.to_string() })), + ); + } + } + } + let map = match read_auth_json_or_empty_recovering_corrupt(&self.path) { Ok(map) => map, Err(e) => { @@ -794,12 +678,11 @@ impl AuthManager { Some(serde_json::json!({ "error": e.to_string() })), ); self.with_inner_write(|inner| *inner = Some(auth.clone())); - self.spawn_user_info_enrichment(auth.clone()); return Ok(auth); } }; let mut map = map; - // One entry per scope (personal and team share the scope key). + // One entry per scope. tracing::debug!(scope = %self.scope, "auth: storing token"); map.insert(self.scope.clone(), auth.clone()); let write_result = write_auth_json(&self.path, &map); @@ -829,78 +712,35 @@ impl AuthManager { // the stale/dead token in memory and the user is completely stuck. self.with_inner_write(|inner| *inner = Some(auth.clone())); - // Fire-and-forget enrichment. Off the critical path -- a slow - // `/user` would otherwise widen the sibling-process - // `invalid_grant` race window. - self.spawn_user_info_enrichment(auth.clone()); - write_result?; Ok(auth) } - /// Persist to disk and cache without spawning the background `/user` task - /// (already merged inline, or a stale fetch must not race a fresh write). - pub(crate) async fn save_without_enrichment( - &self, - auth: GrokAuth, - ) -> std::io::Result { - let started = std::time::Instant::now(); - let map = match read_auth_json_or_empty_recovering_corrupt(&self.path) { - Ok(map) => map, - Err(e) => { - // Non-recoverable error — keep conservative. - tracing::warn!(error = %e, "auth: read failed, updating in-memory only (no enrichment)"); - kigi_log::unified_log::warn( - "auth update skipped disk write (read failed, no enrichment)", - None, - Some(serde_json::json!({ "error": e.to_string() })), - ); - self.with_inner_write(|inner| *inner = Some(auth.clone())); - return Ok(auth); - } + /// Best-effort removal of this scope's entry from `auth.json` after a + /// successful keyring write, so a stale plaintext copy can't shadow the + /// keyring credential later. No lock escalation: callers already hold + /// the auth-file lock on the mutation paths that matter. + fn strip_scope_from_file_best_effort(&self) { + let Ok(mut map) = read_auth_json(&self.path) else { + return; // missing/corrupt file: nothing to strip }; - let mut map = map; - tracing::debug!(scope = %self.scope, "auth: storing token (no enrichment)"); - map.insert(self.scope.clone(), auth.clone()); - let write_result = write_auth_json(&self.path, &map); - let elapsed_ms = started.elapsed().as_millis() as u64; - match &write_result { - Ok(()) => kigi_log::unified_log::info( - "auth update disk written (no enrichment)", - None, - Some(serde_json::json!({ - "rt_prefix": auth.refresh_token.as_deref().map(token_suffix), - "key_prefix": token_suffix(&auth.key), - "elapsed_ms": elapsed_ms, - })), - ), - Err(e) => kigi_log::unified_log::error( - "auth update disk write failed (no enrichment)", - None, - Some(serde_json::json!({ - "error": e.to_string(), - "elapsed_ms": elapsed_ms, - })), - ), + if map.remove(&self.scope).is_none() { + return; + } + let result = if map.is_empty() { + std::fs::remove_file(&self.path) + } else { + write_auth_json(&self.path, &map) + }; + if let Err(e) = result { + tracing::warn!(error = %e, "auth: failed to strip stale file copy after keyring write"); + } else { + tracing::info!("auth: stripped stale file copy after keyring write"); } - // Always update in-memory, even if disk write failed (see update()). - self.with_inner_write(|inner| *inner = Some(auth.clone())); - write_result?; - Ok(auth) } - /// Spawn the `/user` enrichment task; body in the `enrichment` submodule. - fn spawn_user_info_enrichment(self: &Arc, auth: GrokAuth) { - enrichment::spawn(Arc::clone(self), auth); - } - - /// Blocking `/user` enrichment for login flows that exit before the background task lands. - pub(crate) async fn enrich_auth_inline(&self, auth: &mut GrokAuth) { - enrichment::enrich_inline(self, auth).await; - } - - pub(crate) fn grok_com_config(&self) -> &GrokComConfig { - &self.grok_com_config + pub(crate) fn kimi_code_config(&self) -> &KimiCodeConfig { + &self.kimi_code_config } /// Handle notified after every successful token refresh. @@ -930,15 +770,8 @@ impl AuthManager { post_key != pre_key } - /// Run the external auth command and parse its output. Pure: no - /// state mutation, no logging (refresher logs once on its arm). - pub(crate) fn run_external_refresh_command(&self, command: &str) -> Option { - let prev = self.inner_auth_or_external_default(); - crate::auth::refresh_with_command(command, &prev) - } - /// Hot-swap credentials (called by config watcher). Does NOT write to disk. - pub(crate) fn hot_swap(&self, new_auth: GrokAuth) { + pub(crate) fn hot_swap(&self, new_auth: KimiAuth) { self.with_inner_write(|inner| *inner = Some(new_auth)); } @@ -955,9 +788,9 @@ impl AuthManager { /// disk key must differ from in-memory (else no one refreshed). pub(crate) fn try_use_disk_token( &self, - disk_auth: Option<&GrokAuth>, + disk_auth: Option<&KimiAuth>, reason: RefreshReason, - ) -> Option { + ) -> Option { let disk_auth = disk_auth?; if self.is_token_expired(disk_auth) { return None; @@ -978,7 +811,7 @@ impl AuthManager { /// telemetry on success. Combines `read_disk_auth` + /// `try_use_disk_token` + the structured log that was previously /// duplicated at each callsite in `refresh_chain`. - fn try_adopt_disk_token(&self, reason: RefreshReason, msg: &str) -> Option { + fn try_adopt_disk_token(&self, reason: RefreshReason, msg: &str) -> Option { let disk_auth = self.read_disk_auth(); let refreshed = self.try_use_disk_token(disk_auth.as_ref(), reason)?; let adopted = token_suffix(&refreshed.key); @@ -995,21 +828,10 @@ impl AuthManager { Some(refreshed) } - /// Current auth or an `External`-defaulted placeholder. **External - /// path only** -- the placeholder's `auth_mode = External` would - /// mis-classify an OIDC token. Carries user fields forward into the - /// binary's freshly-minted token. - fn inner_auth_or_external_default(&self) -> GrokAuth { - self.inner.read().clone().unwrap_or_else(|| GrokAuth { - auth_mode: AuthMode::External, - ..Default::default() - }) - } - - /// Test-only hot_swap + disk write (skips proxy `/user`). + /// Test-only hot_swap + disk write (file store only). /// Production persistence routes through `update()`. #[cfg(test)] - fn persist_and_swap(&self, auth: GrokAuth) -> Option { + fn persist_and_swap(&self, auth: KimiAuth) -> Option { self.hot_swap(auth.clone()); let mut map = match read_auth_json_or_empty(&self.path) { Ok(m) => m, @@ -1048,16 +870,22 @@ impl AuthManager { mem_rt.as_deref() != Some(disk_rt) } - /// Re-read `auth.json` from disk without updating in-memory state. - pub(crate) fn read_disk_auth(&self) -> Option { + /// Re-read the persisted credential (keyring → file) without updating + /// in-memory state. + pub(crate) fn read_disk_auth(&self) -> Option { self.read_disk_auth_with_state().0 } - /// Disk read for the configured scope with NO observation side effects (no - /// `disk_state` write, no transition telemetry). For side-effect-free - /// getters like [`Self::attempted_verdict_key`]; prefer [`Self::read_disk_auth`] - /// when the read should drive transition logging. - fn read_disk_auth_silent(&self) -> Option { + /// Persisted read for the configured scope with NO observation side + /// effects (no `disk_state` write, no transition telemetry). For + /// side-effect-free getters like [`Self::attempted_tombstone_key`]; prefer + /// [`Self::read_disk_auth`] when the read should drive transition logging. + fn read_disk_auth_silent(&self) -> Option { + if self.scope == KIMI_CODE_OAUTH_SCOPE + && let KeyringRead::Found(auth) = keyring_read_session() + { + return Some(*auth); + } read_auth_json(&self.path) .ok() .and_then(|map| lookup_auth(&map, &self.scope)) @@ -1085,7 +913,15 @@ impl AuthManager { /// can tell a transient disk anomaly (`FileMissing`/`Unreadable`) apart from /// a genuine logout (`EntryMissing`). Observes the state for transition /// logging, exactly like `read_disk_auth`. - pub(crate) fn read_disk_auth_with_state(&self) -> (Option, DiskAuthState) { + pub(crate) fn read_disk_auth_with_state(&self) -> (Option, DiskAuthState) { + // Keyring first (PRD F1): a hit is authoritative for the session + // scope; a miss or an unavailable backend falls through to the file. + if self.scope == KIMI_CODE_OAUTH_SCOPE + && let KeyringRead::Found(auth) = keyring_read_session() + { + self.observe_disk_state(DiskAuthState::Ok, Some(&auth), None); + return (Some(*auth), DiskAuthState::Ok); + } let (auth, state, err_detail) = match read_auth_json(&self.path) { Ok(map) => { let found = lookup_auth(&map, &self.scope); @@ -1119,7 +955,7 @@ impl AuthManager { fn observe_disk_state( &self, new_state: DiskAuthState, - auth: Option<&GrokAuth>, + auth: Option<&KimiAuth>, err_detail: Option, ) { let prev = { @@ -1170,7 +1006,7 @@ impl AuthManager { /// per-session call sites don't reset refresher-internal state). /// Returns `true` if /// this call installed the refresher. - pub fn configure_refresher(self: &Arc, auth_provider_command: Option) -> bool { + pub fn configure_refresher(self: &Arc) -> bool { use std::sync::atomic::Ordering; // Idempotent: the AcqRel CAS publishes the subsequent // `refresher.write()` to any reader that observes @@ -1183,7 +1019,7 @@ impl AuthManager { tracing::debug!("auth: configure_refresher already wired; ignoring"); return false; } - let refresher = super::refresh::build_refresher(Arc::clone(self), auth_provider_command); + let refresher = super::refresh::build_refresher(Arc::clone(self)); *self.refresher.write() = Some(refresher); true } @@ -1208,8 +1044,7 @@ impl AuthManager { .load(std::sync::atomic::Ordering::SeqCst) } - /// `pub(super)` — for refresh dispatch only. External session - /// classification uses `is_session_based_method`. + /// `pub(super)` — for refresh dispatch only. pub(super) fn token_type(&self) -> TokenType { TokenType::from_auth(self.inner.read().as_ref()) } @@ -1218,23 +1053,15 @@ impl AuthManager { /// Pre-request entry point: per-`TokenType` dispatch. For just the key: /// [`Self::get_valid_token`]. - /// - /// Also the team-pin gate: a cached/refreshed wrong-team session is cleared - /// and rejected here, never handed to a consumer. #[tracing::instrument(skip(self), fields(token_type = tracing::field::Empty))] - pub async fn auth(self: &Arc) -> Result { - let auth = self.auth_dispatch().await?; - if let Some(e) = self.cached_token_policy_error(&auth) { - self.reject_and_clear(&e); - return Err(e); - } - Ok(auth) + pub async fn auth(self: &Arc) -> Result { + self.auth_dispatch().await } - async fn auth_dispatch(self: &Arc) -> Result { + async fn auth_dispatch(self: &Arc) -> Result { // Snapshot inner ONCE for dispatch atomicity (closes a TOCTOU // where a concurrent `clear()` raced `token_type()` + `inner.read()`). - let snapshot: Option = self.with_inner_read(|inner| inner.cloned()); + let snapshot: Option = self.with_inner_read(|inner| inner.cloned()); let token_type = TokenType::from_auth(snapshot.as_ref()); tracing::Span::current().record("token_type", tracing::field::debug(token_type)); @@ -1256,25 +1083,17 @@ impl AuthManager { return Ok(auth.clone()); } // A sibling process may have refreshed while we were in - // PermanentFailure. Check disk before giving up. + // PermanentFailure. Check the persisted store before giving up. if let Some(refreshed) = self.try_adopt_disk_token( RefreshReason::PreRequest, "auth: adopted sibling token during PermanentFailure in auth()", ) { return Ok(refreshed); } - // On devboxes, try minting fresh credentials before giving up. - // preferred_method=api_key forbids automatic OIDC mint. - if !self.grok_com_config.blocks_automatic_oidc() - && self.is_devbox_environment() - && let Ok(auth) = self.try_devbox_recovery().await - { - return Ok(auth); - } return Err(err); } - let result = match token_type { + match token_type { TokenType::None => Err(AuthError::NotLoggedIn), TokenType::ApiKey => { // The fast path above already returned for the valid case. @@ -1290,30 +1109,29 @@ impl AuthManager { Err(AuthError::NotLoggedIn) } } - TokenType::LegacySession => { - // Deliberate side effect: re-read auth.json under the - // assumption that a sibling process (`grok login` from - // another shell, the desktop app, etc.) may have refreshed - // the on-disk credentials. `pick_up_sibling_token` only - // mutates inner when the disk holds a *different valid* - // token, so the common cache-hit case is a single read. - // Documented at module level under "Lock ordering". + TokenType::SessionNoRefresh => { + // Deliberate side effect: re-read the persisted store under + // the assumption that a sibling process (`kigi login` from + // another shell) may have refreshed the credential. + // `pick_up_sibling_token` only mutates inner when the store + // holds a *different valid* token, so the common cache-hit + // case is a single read. self.pick_up_sibling_token(); self.current().ok_or(AuthError::TokenExpiredNoRefresh) } - TokenType::OidcSession | TokenType::ExternalBinary => { + TokenType::OAuthSession => { match self .refresh_chain(token_type, RefreshReason::PreRequest) .await { Ok(auth) => Ok(auth), Err(e) => { - // Grace: the early-invalidation buffer is OUR - // conservative estimate, not the IdP's actual - // expiry. If the cached token is still wire-valid + // Grace: the refresh threshold is OUR conservative + // estimate, not the server's actual expiry. If the + // cached token is still wire-valid // ([`Self::is_token_hard_expired`]), return it so a - // transient IdP blip during the buffer window - // is invisible to the user. + // transient OAuth-host blip during the threshold + // window is invisible to the user. if let Some(auth) = snapshot && !self.is_token_hard_expired(&auth) { @@ -1327,102 +1145,7 @@ impl AuthManager { } } } - }; - - // Devbox last-resort recovery: if all normal auth/refresh paths - // failed and we're on a devbox, try minting fresh credentials via - // the remote devbox login helper. Purges existing auth.json and writes only - // the new OIDC entry so we start from a clean state. - // preferred_method=api_key forbids automatic OIDC mint. - if result.is_err() - && !self.grok_com_config.blocks_automatic_oidc() - && self.is_devbox_environment() - && let Ok(auth) = self.try_devbox_recovery().await - { - return Ok(auth); } - - result - } - - /// Whether we're running inside a devbox environment. Wraps the free function - /// so tests can pin it per-instance (CI may run in a container where it is `true`). - pub(crate) fn is_devbox_environment(&self) -> bool { - #[cfg(test)] - if let Some(forced) = *self.devbox_override.lock() { - return forced; - } - crate::auth::devbox_login::is_devbox_environment() - } - - /// Force [`Self::is_devbox_environment`] in tests. - #[cfg(test)] - pub(crate) fn set_devbox_env_for_test(&self, is_devbox: bool) { - *self.devbox_override.lock() = Some(is_devbox); - } - - /// Last-resort devbox auth recovery: purge existing auth.json entirely - /// and mint fresh OIDC credentials via the remote devbox login helper. - /// Only callable on devboxes (where the local service-account token is - /// available). - /// - /// Fail-closed under `preferred_method=api_key` (no automatic OIDC mint), - /// including direct callers such as sampler 401 recovery. - pub(crate) async fn try_devbox_recovery(self: &Arc) -> Result { - if self.grok_com_config.blocks_automatic_oidc() { - tracing::debug!( - "auth: devbox recovery skipped (preferred_method=api_key blocks automatic OIDC)" - ); - kigi_log::unified_log::info( - "auth: devbox recovery skipped (preferred_method=api_key)", - None, - None, - ); - return Err(AuthError::NotLoggedIn); - } - - let _guard = self.refresh_lock.lock().await; - - // Double-check: another task may have recovered while we waited. - if let Some(auth) = self.current() { - return Ok(auth); - } - - tracing::info!("auth: attempting devbox recovery (purge + re-mint)"); - kigi_log::unified_log::info("auth: devbox recovery starting", None, None); - - // Raw mint: the `/user` merge would block up to 10s under refresh_lock. - let new_auth = super::devbox_login::mint_devbox_auth_raw() - .await - .map_err(|e| { - tracing::warn!(error = %e, "auth: devbox recovery mint failed"); - AuthError::transient_source(e) - })?; - - // Purge auth.json so we start clean — removes any corrupted, - // revoked, or legacy entries that caused the failure. - let _ = tokio::fs::remove_file(&self.path).await; - self.clear_inner(); - - let auth = self.save_without_enrichment(new_auth).await.map_err(|e| { - tracing::warn!(error = %e, "auth: devbox recovery save failed"); - AuthError::transient_source(e) - })?; - - // ZDR flags arrive via the background `/user` merge, off the lock. - self.spawn_user_info_enrichment(auth.clone()); - - kigi_log::unified_log::info( - "auth: devbox recovery succeeded", - None, - Some(serde_json::json!({ - "user_id": auth.user_id, - "has_refresh_token": auth.refresh_token.is_some(), - "expires_at": auth.expires_at.map(|e| e.to_rfc3339()), - })), - ); - - Ok(auth) } /// Return the current valid token string, or an error. @@ -1447,7 +1170,7 @@ impl AuthManager { self: &Arc, token_type: TokenType, reason: RefreshReason, - ) -> Result { + ) -> Result { // 0. Sticky permanent-failure short-circuit, checked BEFORE acquiring // the refresh lock so a backed-off chain doesn't block concurrent // traffic. Mirrors `auth()` so callers routing through @@ -1520,11 +1243,11 @@ impl AuthManager { return Err(AuthError::transient("no refresher configured")); }; - // Fallback verdict key, used only when the outcome carries no - // `tried_key` (external-binary flow). Captured before the IdP call so it + // Fallback tombstone key, used only when the outcome carries no + // `rejected_refresh_token`. Captured before the wire call so it // reflects the credential we resolved to send; see - // [`Self::attempted_verdict_key`]. - let attempted_key = self.attempted_verdict_key(reason); + // [`Self::attempted_tombstone_key`]. + let attempted_key = self.attempted_tombstone_key(reason); // 3a. Pre-IdP deferral guards (sleep / dark wake). self.check_refresh_deferral(reason)?; @@ -1706,8 +1429,9 @@ impl AuthManager { } /// Step 3c outcome handling: the only mutation point, persisting on success - /// and recording the verdict on permanent failure. `attempted_key` is the - /// fallback verdict scope (used when the outcome carries no `tried_key`). + /// and recording the tombstone on permanent failure. `attempted_key` is the + /// fallback tombstone scope (used when the outcome carries no + /// `rejected_refresh_token`). /// `_lock` is the held `auth.json` file lock: unused at runtime, threaded in /// to type-enforce that the persisting `update()` runs while the lock is held /// (so a future refactor can't drop it before persisting). @@ -1717,7 +1441,7 @@ impl AuthManager { reason: RefreshReason, attempted_key: Option, _lock: &AuthFileLock, - ) -> Result { + ) -> Result { let pre_key_prefix = attempted_key.as_deref().map(token_suffix); match outcome { RefreshOutcome::Success(new_auth) => match self.update(*new_auth).await { @@ -1747,7 +1471,10 @@ impl AuthManager { Err(AuthError::transient_source(e)) } }, - RefreshOutcome::PermanentFailure { error, tried_key } => { + RefreshOutcome::PermanentFailure { + error, + rejected_refresh_token, + } => { tracing::warn!(reason = ?error.reason, "auth.refresh.permanent_failure"); kigi_log::unified_log::warn( "auth.refresh.permanent_failure", @@ -1757,7 +1484,7 @@ impl AuthManager { })), ); // A sibling may have successfully refreshed while we got a 401. - // If disk has a valid token, adopt it instead. + // If the persisted store has a valid token, adopt it instead. if let Some(refreshed) = self.try_adopt_disk_token( reason, "auth: adopted sibling token after PermanentFailure", @@ -1768,12 +1495,12 @@ impl AuthManager { tracing::info!("auth: sibling-rotation detected; demoting to transient"); return Err(AuthError::transient(format!("sibling-rotation: {error}"))); } - // No clear: the verdict (+ TTL) gates re-attempts; the dead - // bearer is dropped only on explicit logout. Key on the - // credential the refresher actually sent (`tried_key`), falling - // back to our own resolution when the authority has no key. + // No clear: the tombstone (+ cooldown) gates re-attempts; the + // dead bearer is dropped only on explicit logout. Key on the + // refresh token the refresher actually sent, falling back to + // our own resolution when the authority has no key. let failed_reason = error.reason; - if let Some(key) = tried_key.or(attempted_key) { + if let Some(key) = rejected_refresh_token.or(attempted_key) { self.record_permanent_failure(key, error); } Err(AuthError::permanent(failed_reason)) @@ -1790,14 +1517,12 @@ impl AuthManager { } } - /// Re-read auth.json from disk and update the in-memory cache (used by the - /// refresh chains). Non-destructive: only updates in-memory if disk has a - /// different valid token (a sibling process wrote a fresher one). + /// Re-read the persisted credential and update the in-memory cache (used + /// by the refresh chains). Non-destructive: only updates in-memory if the + /// store has a different valid token (a sibling process wrote a fresher + /// one). pub(crate) fn pick_up_sibling_token(&self) { - let auth = match read_auth_json(&self.path) { - Ok(map) => lookup_auth(&map, &self.scope), - _ => None, - }; + let auth = self.read_disk_auth_silent(); if let Some(ref a) = auth && !self.is_token_expired(a) && self.is_different_token(a) @@ -1817,80 +1542,90 @@ impl AuthManager { } /// Check if a candidate auth has a different token than what's in memory. - pub(crate) fn is_different_token(&self, candidate: &GrokAuth) -> bool { + pub(crate) fn is_different_token(&self, candidate: &KimiAuth) -> bool { let current_key = self.inner.read().as_ref().map(|a| a.key.clone()); current_key.as_deref() != Some(&candidate.key) } - /// Record a permanent-failure verdict scoped to `token_key` (the rejected - /// credential). + /// Record a refresh-rejection tombstone scoped to `refresh_token_key` + /// (the rejected refresh-token value). PRD F1: 300s cooldown; auto-clears + /// when the persisted refresh token differs. pub(crate) fn record_permanent_failure( &self, - token_key: String, + refresh_token_key: String, error: crate::auth::error::RefreshTokenFailedError, ) { - // Don't advertise a TTL for a sticky (never-expiring) verdict. - let ttl_seconds = (!error.reason.is_sticky()).then(|| PERMANENT_FAILURE_TTL.as_secs()); kigi_log::unified_log::warn( - "auth.permanent_failure.set", + "auth.tombstone.set", None, Some(serde_json::json!({ "reason": format!("{:?}", error.reason), "message": error.reason.user_message(), - "ttl_seconds": ttl_seconds, + "rt_prefix": token_suffix(&refresh_token_key), + "cooldown_seconds": PERMANENT_FAILURE_TTL.as_secs(), })), ); + tracing::warn!( + rt_prefix = token_suffix(&refresh_token_key), + cooldown_secs = PERMANENT_FAILURE_TTL.as_secs(), + "auth: refresh-rejection tombstone set" + ); *self.permanent_failure.write() = Some(ScopedRefreshFailure { - token_key, + refresh_token_key, error, recorded_at: GateRaise::now(), }); } - /// Key the sticky verdict is scoped to: the credential a refresh for + /// Refresh token the tombstone is scoped to: the one a refresh for /// `reason` would send, via the shared [`resolve_refresh_credential`] (so - /// record and check can't drift). Does a synchronous `auth.json` read; that - /// read is load-bearing (it detects a sibling's freshly rotated token, so an - /// in-memory-only check could leave a stale verdict on a now-valid - /// credential). Called from [`Self::permanent_failure`] (only when a verdict - /// is stored) and once per active `refresh_chain` as the fallback verdict - /// key; both are pre-IdP paths where the read cost is bounded. - fn attempted_verdict_key(&self, reason: RefreshReason) -> Option { - resolve_refresh_credential(self, self.read_disk_auth_silent(), reason).map(|a| a.key) + /// record and check can't drift). Does a synchronous persisted-store read; + /// that read is load-bearing (it detects a sibling's freshly rotated + /// token, so an in-memory-only check could leave a stale tombstone on a + /// now-valid credential). Called from [`Self::permanent_failure`] (only + /// when a tombstone is stored) and once per active `refresh_chain` as the + /// fallback tombstone key; both are pre-wire paths where the read cost is + /// bounded. + fn attempted_tombstone_key(&self, reason: RefreshReason) -> Option { + resolve_refresh_credential(self, self.read_disk_auth_silent(), reason) + .and_then(|a| a.refresh_token) } - /// Sticky verdict for the *attempted* credential, or `None` once it changes - /// (key mismatch) or, for the recoverable reasons, ages out past - /// [`PERMANENT_FAILURE_TTL`]. Reads the stored verdict first (cheap lock): - /// the common no-verdict case returns before any disk I/O; only a stored - /// verdict triggers [`Self::attempted_verdict_key`]'s disk read. + /// Live tombstone for the *attempted* refresh token, or `None` once the + /// persisted refresh token changes (another process rotated it — the + /// tombstone auto-clears) or the 300s cooldown elapses. /// - /// TTL expiry is judged on *both* clocks (see [`GateRaise`]): the monotonic - /// clock pauses during a system suspend, so a wall-clock arm is required - /// for the TTL to elapse across sleep. Without it, a recoverable failure - /// cached just before the lid closes (e.g. a transient escalation while - /// the network was already down) would keep short-circuiting `auth()` — - /// surfacing "run /login" — for up to 5 *awake* minutes after wake, even - /// though the blip is long over. A genuine revocation simply re-caches on - /// the next refresh attempt, so expiring "early" costs one IdP roundtrip. + /// Cooldown expiry is judged on *both* clocks (see [`GateRaise`]): the + /// monotonic clock pauses during a system suspend, so a wall-clock arm is + /// required for the cooldown to elapse across sleep. Without it, a + /// rejection cached just before the lid closes would keep + /// short-circuiting `auth()` — surfacing "run `kigi login`" — for up to 5 + /// *awake* minutes after wake, even though the cooldown is long over. A + /// genuine revocation simply re-caches on the next refresh attempt, so + /// expiring "early" costs one OAuth-host roundtrip. pub(crate) fn permanent_failure(&self) -> Option { - let (token_key, reason) = { + let (refresh_token_key, reason) = { let guard = self.permanent_failure.read(); let pf = guard.as_ref()?; - if !pf.error.reason.is_sticky() { - let (mono, wall) = pf.recorded_at.elapsed(); - if mono >= PERMANENT_FAILURE_TTL || wall >= PERMANENT_FAILURE_TTL { - return None; - } + let (mono, wall) = pf.recorded_at.elapsed(); + if mono >= PERMANENT_FAILURE_TTL || wall >= PERMANENT_FAILURE_TTL { + tracing::info!("auth: refresh-rejection tombstone cooldown elapsed"); + return None; } - (pf.token_key.clone(), pf.error.reason) + (pf.refresh_token_key.clone(), pf.error.reason) }; - // Verdict exists: confirm it still scopes to the credential a refresh - // would attempt. Guard dropped above so `inner` isn't co-held. + // Tombstone exists: confirm it still scopes to the refresh token a + // refresh would attempt. Guard dropped above so `inner` isn't co-held. // Deliberately `ServerRejected` (the widest resolution) regardless of - // the caller's reason, so the read never misses a stored verdict. - (self.attempted_verdict_key(RefreshReason::ServerRejected)? == token_key) - .then(|| AuthError::permanent(reason)) + // the caller's reason, so the read never misses a stored tombstone. + let attempted = self.attempted_tombstone_key(RefreshReason::ServerRejected)?; + if attempted != refresh_token_key { + tracing::info!( + "auth: refresh-rejection tombstone cleared (persisted refresh token rotated)" + ); + return None; + } + Some(AuthError::permanent(reason)) } /// `true` iff [`Self::permanent_failure`] has a non-expired entry. Lets @@ -1944,7 +1679,7 @@ impl AuthManager { /// one-shot recovery off the live bearer, use `try_recover_unauthorized()`. pub(crate) fn unauthorized_recovery( self: &Arc, - rejected: Option, + rejected: Option, ) -> crate::auth::recovery::UnauthorizedRecovery { crate::auth::recovery::UnauthorizedRecovery::new(self.clone(), rejected) } @@ -1958,14 +1693,16 @@ impl AuthManager { // ── Proactive refresh ───────────────────────────────────────────── - /// Spawn a background task that proactively refreshes the token - /// ahead of expiry. Cancelled via `cancel`. + /// Spawn the background refresh task (PRD F1): a fixed + /// [`PROACTIVE_REFRESH_INTERVAL`] (60s) tick that refreshes when the + /// remaining lifetime drops below `max(300, expires_in × 0.5)` seconds + /// (enforced by [`Self::current`]'s dynamic threshold), plus sleep/wake + /// detection — a tick whose wall-clock gap exceeds twice the interval + /// forces a refresh regardless of the threshold, like kimi-cli's + /// `refreshing()`. Cancelled via `cancel`. /// /// Idempotent: a second call on the same `Arc` is a no-op (debug - /// log + return). Sleep duration and back-off conditions are - /// computed by [`compute_proactive_sleep`]; see its body for the - /// four non-busy-loop guards (permanent_failure, non-refreshable - /// type, no refresher, no expires_at). + /// log + return). pub(crate) fn start_proactive_refresh(self: &Arc, cancel: CancellationToken) { use std::sync::atomic::Ordering; // AcqRel/Acquire publishes the spawned task's captured Arc to @@ -1984,177 +1721,113 @@ impl AuthManager { let this = self.clone(); tokio::spawn(async move { loop { - let sleep_dur = compute_proactive_sleep(&this); - + let wall_before = std::time::SystemTime::now(); tokio::select! { _ = cancel.cancelled() => { tracing::debug!("auth: proactive refresh task cancelled"); return; } - _ = tokio::time::sleep(sleep_dur) => {} + _ = tokio::time::sleep(PROACTIVE_REFRESH_INTERVAL) => {} } #[cfg(test)] this.proactive_iter_count.fetch_add(1, Ordering::SeqCst); - // Re-check the back-off preconditions after the sleep so - // a concurrent `update()` / `hot_swap()` / - // `configure_refresher()` is observed before we attempt - // `auth()`. - if this.permanent_failure().is_some() { - // Try disk adoption — a sibling may have refreshed. - if let Some(_refreshed) = this.try_adopt_disk_token( - RefreshReason::PreRequest, - "auth: proactive refresh adopted sibling token during PermanentFailure", - ) { - // Fall through to the normal proactive-refresh sleep - // calculation which will schedule the next refresh - // based on the adopted token's expiry. - continue; - } - tracing::debug!( - "auth: skipping proactive refresh, permanent failure still set" - ); - continue; - } - if !this.token_type().is_refreshable() { - tracing::debug!( - "auth: skipping proactive refresh, token type is not refreshable" - ); - continue; - } - if this.refresher.read().is_none() { - tracing::debug!("auth: skipping proactive refresh, no refresher configured"); - continue; - } - - // Before calling the IdP, check if a sibling process - // already refreshed and wrote a valid token to disk. - // Combined with jitter, the first process to wake - // refreshes; later processes adopt the result here. - this.pick_up_sibling_token(); - if this.current().is_some() { - let adopted = this.current().map(|a| token_suffix(&a.key).to_owned()); - let expires_at = this - .inner - .read() - .as_ref() - .and_then(|a| a.expires_at.map(|e| e.to_rfc3339())); + // Sleep/wake detection: a 60s timer that took far longer on + // the wall clock means the machine was suspended; force a + // refresh so the session recovers immediately on wake. + let elapsed = wall_before.elapsed().unwrap_or_default(); + let force = elapsed > PROACTIVE_REFRESH_INTERVAL * SLEEP_WAKE_FORCE_FACTOR; + if force { tracing::info!( - "auth: proactive refresh skipped, adopted sibling token from disk" + elapsed_secs = elapsed.as_secs(), + "auth: detected possible sleep/wake, forcing token refresh" ); - kigi_log::unified_log::info( - "auth: proactive refresh adopted sibling token", - None, - Some(serde_json::json!({ - "adopted_key_prefix": adopted, - "expires_at": expires_at, - })), - ); - continue; } - tracing::info!("auth: proactive refresh starting"); - match this.auth().await { - Ok(auth) => { - tracing::info!("auth: proactive refresh succeeded"); - kigi_log::unified_log::info( - "auth: proactive refresh completed", - None, - Some(serde_json::json!({ - "result": "success", - "key_prefix": token_suffix(&auth.key), - "expires_at": auth.expires_at.map(|e| e.to_rfc3339()), - })), - ); - } - Err(e) => { - tracing::warn!(error = %e, "auth: proactive refresh failed"); - kigi_log::unified_log::warn( - "auth: proactive refresh completed", - None, - Some(serde_json::json!({ - "result": "failed", - "error": format!("{e}"), - })), - ); - } - } + this.proactive_tick(force).await; } }); } -} -/// Compute the sleep duration for the next iteration of the proactive -/// refresh loop. Pulled out of `start_proactive_refresh` so the gate -/// chain is testable in isolation and the spawned async block stays small. -pub(crate) fn compute_proactive_sleep(this: &AuthManager) -> StdDuration { - if this.permanent_failure().is_some() { - // A previous refresh failed permanently and the verdict is cached. - // `auth()` short-circuits every iteration -- back off until the cancel - // token fires or a fresh `update()` / `hot_swap()` clears it. Jitter so - // a synchronized fleet event (mass client rotation) doesn't re-hit the - // IdP in lockstep once the recoverable verdicts age out. - return BACKOFF_INTERVAL - + StdDuration::from_secs(rand::random_range(0..JITTER_RANGE_SECS) as u64); - } - if !this.token_type().is_refreshable() { - // `ApiKey`, `LegacySession`, and `None` cannot be refreshed - // silently. Without this gate, an expired token of these types - // produces `sleep_dur=0` -> `auth()` -> `TokenExpiredNoRefresh` - // -> repeat at 100% CPU and log spam (same shape as the pre-fix - // permanent-failure busy-loop). - return BACKOFF_INTERVAL; - } - if this.refresher.read().is_none() { - // Defensive: if `start_proactive_refresh` outraced - // `configure_refresher` at startup, `auth()` would emit a transient - // "no refresher configured" error every iteration (and busy-loop for - // a past-expiry token). Hold off here until the refresher is installed. - return BACKOFF_INTERVAL; - } - if this.is_sleep_gated() { - // System sleep is imminent: `refresh_chain` defers every attempt, so - // an expired token would otherwise busy-loop here (`sleep_dur=0` -> - // `auth()` -> transient defer -> repeat). Back off until the gate - // clears on wake or auto-expires (`SLEEP_GATE_MAX`). - return BACKOFF_INTERVAL; - } - if this.is_dark_wake() { - // Dark wake (maintenance / Power Nap): `refresh_chain` defers attempts - // (up to `DARK_WAKE_DEFER_MAX`) for the same reason, so back off instead - // of busy-looping until the next poll or a full wake. - return BACKOFF_INTERVAL; - } - match this.inner.read().as_ref().and_then(|a| a.expires_at) { - Some(expires_at) => { - let buffer = early_invalidation(); - // Add random jitter (0–60 s) so sibling processes don't all - // wake at the same instant and thundering-herd the IdP. The - // first process to wake refreshes and writes to disk; later - // processes pick up the sibling token via - // `pick_up_sibling_token` at the top of the loop. - let jitter = Duration::seconds(rand::random_range(0..JITTER_RANGE_SECS)); - let target = expires_at - buffer - jitter; - let delta = target.signed_duration_since(Utc::now()); - if delta <= Duration::zero() { - // Already past the early-invalidation boundary: - // `auth()` will enter `refresh_chain` immediately. - StdDuration::from_secs(0) - } else { - // The earlier `delta <= 0` branch already handled the - // negative case, so `to_std` cannot fail here. expect - // surfaces a clear panic message if a future change - // breaks the invariant. - delta - .to_std() - .expect("delta > 0 above; chrono::Duration -> std::Duration must succeed") + /// One iteration of the proactive refresh loop. `force` bypasses the + /// still-valid short-circuit (sleep/wake recovery). + pub(crate) async fn proactive_tick(self: &Arc, force: bool) { + // Back-off guards: skip ticks that cannot make progress. + if self.permanent_failure().is_some() { + // Tombstone live. A sibling may have rotated the credential — + // adopt it; otherwise wait out the cooldown. + if self + .try_adopt_disk_token( + RefreshReason::PreRequest, + "auth: proactive refresh adopted sibling token during tombstone cooldown", + ) + .is_none() + { + tracing::debug!("auth: skipping proactive refresh, tombstone cooldown active"); + } + return; + } + if !self.token_type().is_refreshable() { + tracing::debug!("auth: skipping proactive refresh, token type is not refreshable"); + return; + } + if self.refresher.read().is_none() { + tracing::debug!("auth: skipping proactive refresh, no refresher configured"); + return; + } + if self.is_sleep_gated() || self.is_dark_wake() { + tracing::debug!("auth: skipping proactive refresh, sleep gate / dark wake active"); + return; + } + + // Check the persisted store first: a sibling process may have + // already refreshed (its rotation is adopted instead of spending + // our refresh token). + self.pick_up_sibling_token(); + if !force && self.current().is_some() { + // Remaining lifetime is still above the dynamic threshold. + tracing::debug!("auth: proactive refresh not needed (above refresh threshold)"); + return; + } + + tracing::info!(force, "auth: proactive refresh starting"); + let result = if force { + // ServerRejected semantics = force: the refresh chain's + // double-check only short-circuits when another task already + // rotated the key, never on a merely-valid cached token. + self.refresh_chain(self.token_type(), RefreshReason::ServerRejected) + .await + } else { + self.auth().await + }; + match result { + Ok(auth) => { + tracing::info!("auth: proactive refresh succeeded"); + kigi_log::unified_log::info( + "auth: proactive refresh completed", + None, + Some(serde_json::json!({ + "result": "success", + "force": force, + "key_prefix": token_suffix(&auth.key), + "expires_at": auth.expires_at.map(|e| e.to_rfc3339()), + })), + ); + } + Err(e) => { + tracing::warn!(error = %e, "auth: proactive refresh failed"); + kigi_log::unified_log::warn( + "auth: proactive refresh completed", + None, + Some(serde_json::json!({ + "result": "failed", + "force": force, + "error": format!("{e}"), + })), + ); } } - // No expires_at (typical for external binaries): poll every - // BACKOFF_INTERVAL. Operators wanting tighter feedback set - // `[grok_com] auth_token_ttl` to drive a real schedule. - None => BACKOFF_INTERVAL, } } diff --git a/crates/codegen/kigi-shell/src/auth/manager/enrichment.rs b/crates/codegen/kigi-shell/src/auth/manager/enrichment.rs deleted file mode 100644 index ff7cb97..0000000 --- a/crates/codegen/kigi-shell/src/auth/manager/enrichment.rs +++ /dev/null @@ -1,256 +0,0 @@ -//! Background `/user` enrichment spawned by `AuthManager::update()`. - -use std::sync::Arc; -use std::time::Duration as StdDuration; - -use super::AuthManager; -use super::lock::try_lock_auth_file_async; -use crate::auth::manager::AUTH_LOCK_TIMEOUT; -use crate::auth::model::{GrokAuth, UserInfo, lookup_auth}; -use crate::auth::storage::{read_auth_json, write_auth_json}; - -/// `/user` fetch budget, shared by the inline (login) and background paths. -const USER_FETCH_TIMEOUT: StdDuration = StdDuration::from_secs(10); - -/// Logs `auth update enrichment dropped` if the task is cancelled -/// mid-flight. Disarmed on normal completion. -pub(super) struct EnrichmentExitGuard { - pub(super) started: std::time::Instant, - pub(super) armed: bool, -} - -impl EnrichmentExitGuard { - pub(super) fn disarm(&mut self) { - self.armed = false; - } -} - -impl Drop for EnrichmentExitGuard { - fn drop(&mut self) { - if !self.armed { - return; - } - kigi_log::unified_log::warn( - "auth update enrichment dropped", - None, - Some(serde_json::json!({ - "elapsed_ms": self.started.elapsed().as_millis() as u64, - })), - ); - } -} - -pub(super) fn spawn(manager: Arc, auth: GrokAuth) { - tokio::spawn(async move { - let mut exit_guard = EnrichmentExitGuard { - started: std::time::Instant::now(), - armed: true, - }; - run_user_info_enrichment(&manager, auth).await; - exit_guard.disarm(); - }); -} - -async fn fetch_user_info(manager: &AuthManager, key: &str, log_label: &str) -> Option { - let user_url = format!("{}/user", manager.proxy_base_url); - let token_header = &manager.grok_com_config.token_header; - let started = std::time::Instant::now(); - let http_client = crate::http::shared_client(); - let response = http_client - .get(&user_url) - .timeout(USER_FETCH_TIMEOUT) - .header("Authorization", format!("Bearer {}", key)) - .header("X-XAI-Token-Auth", token_header.as_str()) - .header("x-grok-client-version", kigi_version::VERSION) - .header( - crate::http::CLIENT_MODE_HEADER, - crate::http::process_client_mode(), - ) - .send() - .await; - - match response { - Ok(resp) if resp.status().is_success() => match resp.json::().await { - Ok(ui) if !ui.user_id.is_empty() => Some(ui), - Ok(_) => { - kigi_log::unified_log::warn( - &format!("{log_label} skipped"), - None, - Some(serde_json::json!({ - "reason": "empty_user_id", - "elapsed_ms": started.elapsed().as_millis() as u64, - })), - ); - None - } - Err(e) => { - kigi_log::unified_log::warn( - &format!("{log_label} failed"), - None, - Some(serde_json::json!({ - "reason": "parse", - "error": e.to_string(), - "elapsed_ms": started.elapsed().as_millis() as u64, - })), - ); - None - } - }, - Ok(resp) => { - kigi_log::unified_log::warn( - &format!("{log_label} failed"), - None, - Some(serde_json::json!({ - "reason": "http_status", - "http_status": resp.status().as_u16(), - "elapsed_ms": started.elapsed().as_millis() as u64, - })), - ); - None - } - Err(e) => { - kigi_log::unified_log::warn( - &format!("{log_label} failed"), - None, - Some(serde_json::json!({ - "reason": if e.is_timeout() { "timeout" } else { "transport" }, - "error": e.to_string(), - "elapsed_ms": started.elapsed().as_millis() as u64, - })), - ); - None - } - } -} - -/// Blocking login-time enrichment: merge `/user` fields before the first save. -pub(super) async fn enrich_inline(manager: &AuthManager, auth: &mut GrokAuth) { - let Some(ui) = fetch_user_info(manager, &auth.key, "auth login enrichment").await else { - return; - }; - apply_user_info_enrichment(auth, ui); -} - -async fn run_user_info_enrichment(manager: &AuthManager, auth: GrokAuth) { - let started = std::time::Instant::now(); - let Some(user_info) = fetch_user_info(manager, &auth.key, "auth update enrichment").await - else { - return; - }; - let user_elapsed_ms = started.elapsed().as_millis() as u64; - - // R-M-W file lock. On timeout, fall through to an unlocked write - // rather than drop the enrichment. - let lock_started = std::time::Instant::now(); - let lock_guard = try_lock_auth_file_async(&manager.path, AUTH_LOCK_TIMEOUT).await; - let lock_wait_ms = lock_started.elapsed().as_millis() as u64; - if lock_guard.is_none() { - tracing::warn!("auth: enrichment proceeding without auth.json.lock"); - } - - let Ok(mut map) = read_auth_json(&manager.path) else { - kigi_log::unified_log::warn( - "auth update enrichment skipped", - None, - Some(serde_json::json!({ "reason": "read_disk_failed" })), - ); - return; - }; - let Some(mut disk) = lookup_auth(&map, &manager.scope) else { - kigi_log::unified_log::info( - "auth update enrichment skipped", - None, - Some(serde_json::json!({ "reason": "no_disk_auth" })), - ); - return; - }; - // Sibling-stomp guard. If either the access token or refresh - // token on disk differs from the one we wrote, a sibling process - // rotated tokens since our update(). Skip enrichment to avoid - // writing stale profile data over the sibling's fresher entry. - // - // OR logic (not AND): a single-field rotation (key changes, RT - // stays) is the common case during concurrent refresh. The old - // AND logic required ALL three fields to differ, letting - // single-field rotations through. - // - // Team-login transitions (placeholder→real user_id) don't rotate - // tokens, so OR correctly allows enrichment for that case. - if disk.key != auth.key || disk.refresh_token != auth.refresh_token { - kigi_log::unified_log::info( - "auth update enrichment skipped", - None, - Some(serde_json::json!({ - "reason": "sibling_rotated", - "written_key_prefix": crate::auth::token_suffix(&auth.key), - "disk_key_prefix": crate::auth::token_suffix(&disk.key), - })), - ); - return; - } - - apply_user_info_enrichment(&mut disk, user_info); - - map.insert(manager.scope.clone(), disk.clone()); - let write_started = std::time::Instant::now(); - if let Err(e) = write_auth_json(&manager.path, &map) { - kigi_log::unified_log::error( - "auth update enrichment write failed", - None, - Some(serde_json::json!({ - "error": e.to_string(), - "user_ms": user_elapsed_ms, - "lock_wait_ms": lock_wait_ms, - "write_ms": write_started.elapsed().as_millis() as u64, - })), - ); - return; - } - manager.with_inner_write(|inner| *inner = Some(disk)); - kigi_log::unified_log::info( - "auth update enrichment done", - None, - Some(serde_json::json!({ - "user_ms": user_elapsed_ms, - "lock_wait_ms": lock_wait_ms, - "write_ms": write_started.elapsed().as_millis() as u64, - "total_ms": started.elapsed().as_millis() as u64, - })), - ); -} - -/// Merge enrichment fields into disk auth. Does not touch token fields. -pub(super) fn apply_user_info_enrichment(disk: &mut GrokAuth, user_info: UserInfo) { - disk.user_id = user_info.user_id; - disk.first_name = user_info.first_name.or(disk.first_name.take()); - disk.last_name = user_info.last_name.or(disk.last_name.take()); - disk.profile_image_asset_id = user_info - .profile_image_asset_id - .or(disk.profile_image_asset_id.take()); - disk.principal_type = user_info.principal_type.or(disk.principal_type.take()); - disk.principal_id = user_info.principal_id.or(disk.principal_id.take()); - disk.team_id = user_info.team_id.or(disk.team_id.take()); - disk.team_name = user_info.team_name.or(disk.team_name.take()); - disk.team_role = user_info.team_role.or(disk.team_role.take()); - disk.organization_id = user_info.organization_id.or(disk.organization_id.take()); - disk.organization_name = user_info - .organization_name - .or(disk.organization_name.take()); - disk.organization_role = user_info - .organization_role - .or(disk.organization_role.take()); - disk.user_blocked_reason = user_info - .user_blocked_reason - .or(disk.user_blocked_reason.take()); - if let Some(reasons) = user_info.team_blocked_reasons { - disk.team_blocked_reasons = reasons; - } - if let Some(opt_out) = user_info.coding_data_retention_opt_out { - disk.coding_data_retention_opt_out = opt_out; - } - if let Some(ref email) = user_info.email - && !email.is_empty() - { - disk.email = user_info.email; - } -} diff --git a/crates/codegen/kigi-shell/src/auth/manager_tests.rs b/crates/codegen/kigi-shell/src/auth/manager_tests.rs index ccd34c9..0c9047a 100644 --- a/crates/codegen/kigi-shell/src/auth/manager_tests.rs +++ b/crates/codegen/kigi-shell/src/auth/manager_tests.rs @@ -1,3912 +1,573 @@ -//! Unit tests for [`super::manager::AuthManager`]. Extracted from -//! `manager.rs` so the implementation reads top-to-bottom; wired in -//! via `#[path = "manager_tests.rs"] mod tests;` in manager.rs. +//! `AuthManager` behavior tests for the Kimi Code auth stack: +//! tombstone semantics (PRD F1), persistence (keyring + file fallback), +//! dynamic refresh threshold, dispatch, and the proactive-tick loop body. +//! +//! Cross-process lock behavior is covered in `manager/lock.rs`; sleep-gate +//! internals in `manager/sleep_gate.rs`; wire behavior in `kimi_oauth.rs` / +//! `refresh/kimi_refresher.rs`. use super::*; -use crate::auth::error::RefreshTokenError; -use std::sync::atomic::{AtomicU32, Ordering}; -use std::time::Instant; +use crate::auth::error::RefreshTokenFailedReason; +use crate::auth::model::AuthMode; +use chrono::Utc; +use std::sync::atomic::{AtomicU32, Ordering as AtomicOrdering}; -fn make_auth(expires_at: Option>, create_time: DateTime) -> GrokAuth { - GrokAuth { - auth_mode: AuthMode::External, - create_time, - user_id: String::new(), - expires_at, - ..GrokAuth::test_default() +fn mgr() -> (tempfile::TempDir, Arc) { + let dir = tempfile::tempdir().unwrap(); + let m = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + (dir, m) +} + +/// Session credential with `remaining_secs` of lifetime left out of a +/// `expires_in`-second grant. +fn session(key: &str, rt: &str, expires_in: i64, remaining_secs: i64) -> KimiAuth { + KimiAuth { + key: key.into(), + auth_mode: AuthMode::OAuth, + refresh_token: Some(rt.into()), + expires_at: Some(Utc::now() + Duration::seconds(remaining_secs)), + expires_in: Some(expires_in), + ..KimiAuth::test_default() } } -#[test] -fn expired_within_5min_buffer() { - let auth = make_auth(Some(Utc::now() + Duration::minutes(4)), Utc::now()); - assert!(is_expired(&auth)); +/// Counting refresher returning a fixed success. +struct OkRefresher { + calls: Arc, } - -#[test] -fn fallback_ttl_when_no_expires_at() { - let old = Utc::now() - Duration::days(30) + Duration::minutes(4); - let auth = make_auth(None, old); - assert!(is_expired(&auth)); - - let recent = Utc::now() - Duration::days(29); - let auth = make_auth(None, recent); - assert!(!is_expired(&auth)); -} - -#[test] -fn has_usable_disk_token_reads_disk_independent_of_memory() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - - assert!(!mgr.has_usable_disk_token()); - - let valid = make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()); - mgr.persist_and_swap(valid); - mgr.clear_in_memory(); - assert!(mgr.current().is_none(), "in-memory cleared"); - assert!( - mgr.has_usable_disk_token(), - "a valid token on disk is usable even when in-memory is empty" - ); - - let expired = make_auth(Some(Utc::now() - Duration::hours(1)), Utc::now()); - mgr.persist_and_swap(expired); - mgr.clear_in_memory(); - assert!( - !mgr.has_usable_disk_token(), - "an expired token on disk is not usable" - ); -} - -#[test] -fn has_usable_token_covers_memory_and_disk() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - - assert!(!mgr.has_usable_token(), "nothing in memory or on disk"); - - mgr.hot_swap(make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())); - assert!(!mgr.has_usable_disk_token(), "disk still empty"); - assert!(mgr.has_usable_token(), "valid in-memory token is usable"); - - mgr.persist_and_swap(make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())); - mgr.hot_swap(make_auth(Some(Utc::now() - Duration::hours(1)), Utc::now())); - assert!(mgr.current().is_none(), "in-memory token is expired"); - assert!(mgr.has_usable_token(), "fresh disk token keeps it usable"); - - mgr.persist_and_swap(make_auth(Some(Utc::now() - Duration::hours(1)), Utc::now())); - assert!( - !mgr.has_usable_token(), - "expired in memory and on disk is not usable" - ); -} - -#[test] -fn auth_scope_uses_oauth2_when_present() { - let cfg = GrokComConfig::default(); - // Default config always has oauth2 set to the xAI defaults. - assert_eq!( - cfg.auth_scope(), - format!( - "{}::{}", - crate::auth::config::XAI_OAUTH2_ISSUER, - obfstr::obfstr!("b1a00492-073a-47ea-816f-4c329264a828"), - ) - ); -} - -#[test] -fn legacy_scope_fallback_reads_old_auth_json() { - let dir = tempfile::tempdir().unwrap(); - let auth_path = dir.path().join("auth.json"); - - // Write auth.json with the legacy scope key (as `x setup` copies from - // a machine that was authenticated with an older grok version). - let legacy_auth = make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()); - let mut store = AuthStore::new(); - store.insert(LEGACY_SCOPE.to_string(), legacy_auth); - write_auth_json(&auth_path, &store).unwrap(); - - // AuthManager uses the new OAuth2 scope, but should still find the - // token under the legacy key. - let cfg = GrokComConfig::default(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - let current = mgr.current(); - assert!(current.is_some(), "should fall back to legacy scope key"); - assert_eq!(current.unwrap().key, "test-key"); -} - -#[test] -fn new_scope_takes_precedence_over_legacy() { - let dir = tempfile::tempdir().unwrap(); - let auth_path = dir.path().join("auth.json"); - - let legacy_auth = GrokAuth { - key: "legacy-key".into(), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - let new_auth = GrokAuth { - key: "new-key".into(), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - - let cfg = GrokComConfig::default(); - let scope = cfg.auth_scope(); - - let mut store = AuthStore::new(); - store.insert(LEGACY_SCOPE.to_string(), legacy_auth); - store.insert(scope, new_auth); - write_auth_json(&auth_path, &store).unwrap(); - - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - let current = mgr.current().expect("should find auth"); - assert_eq!(current.key, "new-key", "new scope should take precedence"); -} - -// -- Near-expiry (5-minute buffer) behavior ------------------------ - -/// Regression test: a token within the 5-minute early-invalidation buffer -/// must be invisible to `current()` (returns None) but visible to -/// `expired_auth()` so that callers can attempt a silent refresh. -#[test] -fn near_expiry_token_invisible_to_current_visible_to_expired_auth() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - - // Token expires in 3 minutes -- inside the 5-minute buffer. - let near_expiry = GrokAuth { - key: "near-expiry-key".into(), - user_id: "user-1".into(), - email: Some("user@test.com".into()), - refresh_token: Some("rt-valid".into()), - expires_at: Some(Utc::now() + Duration::minutes(3)), - oidc_issuer: Some("https://idp.example.com".into()), - oidc_client_id: Some("client-1".into()), - ..GrokAuth::test_default() - }; - mgr.hot_swap(near_expiry); - - // current() must return None (token is "expired" per buffer) - assert!( - mgr.current().is_none(), - "current() should return None for token within 5-min buffer" - ); - - // is_expired() must return true - assert!( - mgr.is_expired(), - "is_expired() should be true for token within 5-min buffer" - ); - - // expired_auth() must return the token so refresh can use it - let expired = mgr.expired_auth(); - assert!( - expired.is_some(), - "expired_auth() should return the near-expiry token" - ); - assert_eq!(expired.as_ref().unwrap().key, "near-expiry-key"); - assert_eq!( - expired.as_ref().unwrap().refresh_token.as_deref(), - Some("rt-valid"), - "refresh_token must be preserved for silent refresh" - ); -} - -#[tokio::test] -async fn update_preserves_other_scope_entries() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg.clone())); - - // Pre-populate with an external auth entry - let external = GrokAuth { - key: "external-key".into(), - auth_mode: AuthMode::External, - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - { - let mut map = AuthStore::new(); - map.insert("other-scope".into(), external); - write_auth_json(&dir.path().join("auth.json"), &map).unwrap(); - } - - // Now update via auth_manager - let new_auth = GrokAuth { - key: "oidc-token".into(), - auth_mode: AuthMode::Oidc, - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - mgr.update(new_auth).await.unwrap(); - - // Both entries should exist - let store = read_auth_json(&dir.path().join("auth.json")).unwrap(); - assert!(store.contains_key("other-scope")); - assert!(store.contains_key(&cfg.auth_scope())); -} - -/// Regression: when auth.json contains corrupt JSON, update() must not -/// clobber the file with a single-entry map. Instead it should update -/// in-memory only and leave the file untouched. -#[tokio::test] -async fn update_recovers_from_corrupt_auth_json_by_backing_up_old_file() { - let dir = tempfile::tempdir().unwrap(); - let auth_path = dir.path().join("auth.json"); - let cfg = GrokComConfig::default(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg.clone())); - - let bad_content = b"NOT VALID JSON {{{"; - std::fs::write(&auth_path, bad_content).unwrap(); - - let new_auth = GrokAuth { - key: "fresh-token".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("fresh-rt".into()), - user_id: "fresh-user".into(), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - - let result = mgr.update(new_auth).await; - assert!( - result.is_ok(), - "update must succeed and persist after corrupt recovery: {result:?}" - ); - - let current = mgr.current(); - assert_eq!( - current.as_ref().map(|a| a.key.as_str()), - Some("fresh-token") - ); - - let on_disk_raw = std::fs::read_to_string(&auth_path).unwrap(); - assert!( - on_disk_raw.contains("fresh-token"), - "auth.json must contain the new credential after recovery, got: {on_disk_raw}" - ); - let on_disk: AuthStore = - serde_json::from_str(&on_disk_raw).expect("auth.json must be valid JSON after recovery"); - assert!(on_disk.contains_key(&cfg.auth_scope())); - - let mut backup_found = None; - for entry in std::fs::read_dir(dir.path()).unwrap() { - let entry = entry.unwrap(); - let name = entry.file_name().to_string_lossy().into_owned(); - if name.starts_with("auth.json.corrupt.") { - backup_found = Some(entry.path()); - break; - } - } - let backup_path = backup_found.expect("a .corrupt.* backup file must have been created"); - let backup_content = std::fs::read_to_string(&backup_path).unwrap(); - assert!( - backup_content.contains("NOT VALID JSON"), - "backup must contain the original corrupt content, got: {backup_content}" - ); -} - -/// Regression test: update() must preserve team fields from the OIDC flow -/// when the proxy `/user` response does not include them. -#[tokio::test] -async fn update_preserves_team_fields_when_proxy_omits_them() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - // Point proxy_base_url to a non-existent server so the /user call - // fails and falls back to the auth-flow values. - let mgr = Arc::new(AuthManager::new(dir.path(), cfg).with_proxy_base_url("http://127.0.0.1:1")); - - let team_auth = GrokAuth { - key: "team-token".into(), - auth_mode: AuthMode::Oidc, - principal_type: Some("Team".into()), - principal_id: Some("team-xyz".into()), - team_id: Some("team-xyz".into()), - team_name: None, - team_role: None, - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - - let saved = mgr.update(team_auth).await.unwrap(); - - assert_eq!( - saved.principal_type.as_deref(), - Some("Team"), - "principal_type must survive proxy fallback" - ); - assert_eq!( - saved.principal_id.as_deref(), - Some("team-xyz"), - "principal_id must survive proxy fallback" - ); - assert_eq!( - saved.team_id.as_deref(), - Some("team-xyz"), - "team_id must survive proxy fallback" - ); - - // Verify on-disk too - let store = read_auth_json(&dir.path().join("auth.json")).unwrap(); - let on_disk = store.values().next().unwrap(); - assert_eq!(on_disk.principal_type.as_deref(), Some("Team")); - assert_eq!(on_disk.team_id.as_deref(), Some("team-xyz")); -} - -/// Team tokens are stored under the base scope key (same as personal). -/// There is at most one OAuth entry per issuer/client pair. -#[tokio::test] -async fn update_stores_team_token_under_base_scope() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let base_scope = cfg.auth_scope(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg).with_proxy_base_url("http://127.0.0.1:1")); - - let team_auth = GrokAuth { - key: "team-token".into(), - auth_mode: AuthMode::Oidc, - principal_type: Some("Team".into()), - principal_id: Some("team-abc".into()), - team_id: Some("team-abc".into()), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - - mgr.update(team_auth).await.unwrap(); - - let store = read_auth_json(&dir.path().join("auth.json")).unwrap(); - assert!( - store.contains_key(&base_scope), - "team token must be stored under base scope '{}', found keys: {:?}", - base_scope, - store.keys().collect::>() - ); - assert_eq!(store.get(&base_scope).unwrap().key, "team-token"); -} - -/// Logging in as personal must evict any existing team token -/// (at most one OAuth session per issuer/client pair). -#[tokio::test] -async fn team_login_then_personal_evicts_team_token() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let base_scope = cfg.auth_scope(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg).with_proxy_base_url("http://127.0.0.1:1")); - - // Step 1: login as team - let team_auth = GrokAuth { - key: "team-token".into(), - principal_type: Some("Team".into()), - principal_id: Some("team-abc".into()), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - mgr.update(team_auth).await.unwrap(); - - // Step 2: login as personal - let personal_auth = GrokAuth { - key: "personal-token".into(), - principal_type: None, - principal_id: None, - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - mgr.update(personal_auth).await.unwrap(); - - let store = read_auth_json(&dir.path().join("auth.json")).unwrap(); - assert_eq!( - store.len(), - 1, - "only one OAuth entry should remain, found: {:?}", - store.keys().collect::>() - ); - assert!(store.contains_key(&base_scope)); - assert_eq!(store.get(&base_scope).unwrap().key, "personal-token"); -} - -/// Regression test: clear() must only remove the current scope, not the -/// legacy scope. Previously, logging in with OAuth would also delete the -/// legacy `https://accounts.x.ai/sign-in` entry from auth.json. -#[test] -fn clear_does_not_remove_legacy_scope() { - let dir = tempfile::tempdir().unwrap(); - let auth_path = dir.path().join("auth.json"); - - let legacy_auth = GrokAuth { - key: "legacy-key".into(), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - let oauth_auth = GrokAuth { - key: "oauth-key".into(), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - - let cfg = GrokComConfig::default(); - let scope = cfg.auth_scope(); - - let mut store = AuthStore::new(); - store.insert(LEGACY_SCOPE.to_string(), legacy_auth); - store.insert(scope, oauth_auth); - write_auth_json(&auth_path, &store).unwrap(); - - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - // clear() should only remove the OAuth scope, not legacy - mgr.clear().unwrap(); - - let on_disk = read_auth_json(&auth_path).unwrap(); - assert!( - on_disk.contains_key(LEGACY_SCOPE), - "legacy scope should be preserved after clear()" - ); - assert!( - !on_disk.contains_key(&mgr.scope), - "current scope should be removed after clear()" - ); -} - -#[test] -fn is_data_collection_disabled_matrix() { - // (team_blocked_reasons, coding_data_retention_opt_out, expected) - let cases: &[(&[&str], bool, bool)] = &[ - // ZDR team alone - (&["BLOCKED_REASON_NO_LOGS"], false, true), - (&["BLOCKED_REASON_NO_LOGS_MODERATED"], false, true), - // Opt-out alone - (&[], true, true), - // Both - (&["BLOCKED_REASON_NO_LOGS"], true, true), - // Neither - (&[], false, false), - // Unrelated blocked reasons - ( - &["BLOCKED_REASON_BILLING", "BLOCKED_REASON_SUSPENDED"], - false, - false, - ), - (&["BLOCKED_REASON_BILLING"], true, true), - // ZDR mixed with other reasons - ( - &["BLOCKED_REASON_BILLING", "BLOCKED_REASON_NO_LOGS"], - false, - true, - ), - ]; - for (reasons, opt_out, expected) in cases { - let auth = GrokAuth { - team_blocked_reasons: reasons.iter().map(|s| (*s).into()).collect(), - coding_data_retention_opt_out: *opt_out, - ..GrokAuth::test_default() - }; - assert_eq!( - auth.is_data_collection_disabled(), - *expected, - "reasons={reasons:?} opt_out={opt_out} expected={expected}", - ); - } -} - -/// Fail-direction contract of the two `AuthManager` collection predicates: -/// `is_data_collection_disabled` fails open on missing credentials (legacy -/// semantics shared by telemetry/sync gates), `allows_data_collection` fails -/// closed (nothing may leave the machine while privacy state is unknown, -/// e.g. after a mid-session `/logout`). -#[test] -fn manager_collection_predicates_fail_directions() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - // No credential: disabled=false (fail-open), allows=false (fail-closed). - assert!(!mgr.is_data_collection_disabled()); - assert!( - !mgr.allows_data_collection(), - "missing credential must fail closed for collection" - ); - - // Normal user: both predicates allow collection. - mgr.hot_swap(GrokAuth::test_default()); - assert!(!mgr.is_data_collection_disabled()); - assert!(mgr.allows_data_collection()); - - // Opted-out user: both predicates suppress collection. - mgr.hot_swap(GrokAuth { - coding_data_retention_opt_out: true, - ..GrokAuth::test_default() - }); - assert!(mgr.is_data_collection_disabled()); - assert!(!mgr.allows_data_collection()); - - // Mid-session `/logout`: the fail-closed predicate flips back to - // "no collection" even after a previously permissive credential. - mgr.hot_swap(GrokAuth::test_default()); - assert!(mgr.allows_data_collection(), "precondition"); - mgr.clear_in_memory(); - assert!( - !mgr.allows_data_collection(), - "cleared credentials must close the collection gate" - ); -} - -// -- token_suffix ---------------------------------------------------------------- - -#[test] -fn token_suffix_matrix() { - let cases: &[(&str, &str)] = &[ - ("abcdefghijklmnop", "efghijklmnop"), // takes last 12 - ("short", "short"), // short unchanged - ("", ""), // empty - ("123456789012", "123456789012"), // exact 12 - ]; - for (input, expected) in cases { - assert_eq!(token_suffix(input), *expected, "input={input:?}"); - } -} - -// -- read_disk_auth ---------------------------------------------------------- - -// -- hot_swap / try_use_disk_token --------------------------------------- - -#[test] -fn hot_swap_updates_in_memory_without_disk() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - - assert!(mgr.current().is_none()); - let auth = GrokAuth { - key: "swapped".into(), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - mgr.hot_swap(auth); - assert_eq!(mgr.current().unwrap().key, "swapped"); - // Disk should NOT have the token - assert!(mgr.read_disk_auth().is_none()); -} - -#[test] -fn try_use_disk_token_accepts_valid_disk_token() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - - let valid_disk = GrokAuth { - key: "valid-disk".into(), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - let result = mgr.try_use_disk_token(Some(&valid_disk), RefreshReason::PreRequest); - assert_eq!(result.unwrap().key, "valid-disk"); - // Should also hot-swap into memory - assert_eq!(mgr.current().unwrap().key, "valid-disk"); -} - -#[test] -fn try_use_disk_token_rejects_expired_disk_token() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - - let expired_disk = make_auth(Some(Utc::now() - Duration::hours(1)), Utc::now()); - assert!( - mgr.try_use_disk_token(Some(&expired_disk), RefreshReason::PreRequest) - .is_none() - ); -} - -#[test] -fn try_use_disk_token_rejects_same_key_on_server_rejected() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - - let auth = GrokAuth { - key: "same-key".into(), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - mgr.hot_swap(auth.clone()); - - // ServerRejected should not accept a disk token with the same key - assert!( - mgr.try_use_disk_token(Some(&auth), RefreshReason::ServerRejected) - .is_none() - ); -} - -#[test] -fn try_use_disk_token_accepts_different_key_on_server_rejected() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - - let mem_auth = GrokAuth { - key: "old-key".into(), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - mgr.hot_swap(mem_auth); - - let disk_auth = GrokAuth { - key: "new-key".into(), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - let result = mgr.try_use_disk_token(Some(&disk_auth), RefreshReason::ServerRejected); - assert_eq!(result.unwrap().key, "new-key"); -} - -// -- File locking ---------------------------------------------------------- - -// -- Disk-refresh race simulation ------------------------------------------ - -/// Simulates the core scenario this PR fixes: an expired in-memory token -/// where another process has already refreshed on disk. The manager should -/// pick up the valid disk token via try_use_disk_token instead of -/// attempting its own refresh. -#[tokio::test] -async fn disk_refresh_wins_over_expired_in_memory() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let scope = cfg.auth_scope(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - - // Simulate: in-memory token is expired - let expired = GrokAuth { - key: "expired-key".into(), - refresh_token: Some("old-rt".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(expired); - assert!(mgr.is_expired()); - assert!(mgr.current().is_none()); - - // Simulate: another process wrote a valid token to disk - let fresh_disk = GrokAuth { - key: "fresh-key-from-sibling".into(), - refresh_token: Some("new-rt".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() - }; - let mut store = AuthStore::new(); - store.insert(scope, fresh_disk); - write_auth_json(&dir.path().join("auth.json"), &store).unwrap(); - - // Acquire lock + read disk (mirrors flow.rs logic) - let _lock = mgr - .try_lock_auth_file_async(StdDuration::from_secs(1)) - .await; - assert!(_lock.is_some()); - - let disk_auth = mgr.read_disk_auth(); - assert!(disk_auth.is_some()); - assert!(!is_expired(disk_auth.as_ref().unwrap())); - - // try_use_disk_token should accept it and hot-swap - let result = mgr.try_use_disk_token(disk_auth.as_ref(), RefreshReason::PreRequest); - assert_eq!(result.unwrap().key, "fresh-key-from-sibling"); - assert_eq!(mgr.current().unwrap().key, "fresh-key-from-sibling"); -} - -struct CountingRefresher { - call_count: Arc, - delay: StdDuration, -} - #[async_trait::async_trait] -impl TokenRefresher for CountingRefresher { - async fn refresh(&self, _reason: RefreshReason) -> crate::auth::refresh::RefreshOutcome { - self.call_count.fetch_add(1, Ordering::SeqCst); - tokio::time::sleep(self.delay).await; - let fresh = GrokAuth { - key: "fresh-token".into(), - expires_at: Some(Utc::now() + Duration::hours(1)), - refresh_token: Some("rt-new".into()), - ..GrokAuth::test_default() - }; - crate::auth::refresh::RefreshOutcome::Success(Box::new(fresh)) +impl TokenRefresher for OkRefresher { + async fn refresh(&self, _reason: RefreshReason) -> RefreshOutcome { + self.calls.fetch_add(1, AtomicOrdering::SeqCst); + RefreshOutcome::success(session("at-refreshed", "rt-refreshed", 3600, 3600)) } } -struct FailingRefresher { - call_count: Arc, +/// Counting refresher that always reports the refresh token rejected. +struct RejectRefresher { + calls: Arc, + rejected_rt: String, } - #[async_trait::async_trait] -impl TokenRefresher for FailingRefresher { - async fn refresh(&self, _reason: RefreshReason) -> crate::auth::refresh::RefreshOutcome { - self.call_count.fetch_add(1, Ordering::SeqCst); - crate::auth::refresh::RefreshOutcome::permanent( - crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected, - None, +impl TokenRefresher for RejectRefresher { + async fn refresh(&self, _reason: RefreshReason) -> RefreshOutcome { + self.calls.fetch_add(1, AtomicOrdering::SeqCst); + RefreshOutcome::permanent( + RefreshTokenFailedReason::RefreshTokenRejected, + Some(self.rejected_rt.clone()), ) } } -/// Record a permanent failure scoped to the auth manager's current (or expired) -/// credential key, mirroring what `refresh_chain` does in production. -fn record_permanent_failure( - auth_manager: &AuthManager, - reason: crate::auth::error::RefreshTokenFailedReason, -) { - let key = auth_manager - .current() - .or_else(|| auth_manager.expired_auth()) - .map(|a| a.key) - .unwrap_or_default(); - auth_manager.record_permanent_failure(key, reason.into()); -} - -/// Permanent-failure refresher that reports a specific `tried_key` (the -/// credential it claims to have sent to the IdP), letting tests assert the -/// verdict is keyed on the actually-tried credential. -struct TriedKeyFailRefresher { - tried_key: String, - call_count: Arc, -} - -#[async_trait::async_trait] -impl TokenRefresher for TriedKeyFailRefresher { - async fn refresh(&self, _reason: RefreshReason) -> crate::auth::refresh::RefreshOutcome { - self.call_count.fetch_add(1, Ordering::SeqCst); - crate::auth::refresh::RefreshOutcome::permanent( - crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected, - Some(self.tried_key.clone()), - ) - } -} - -/// With `inner == None` but a dead refresh-token on disk, the refresher still -/// exchanges that disk RT. The verdict must be keyed on the -/// credential actually tried (the disk RT), so repeated reactive refreshes -/// short-circuit on it instead of hammering the IdP. -#[tokio::test] -async fn storm_cap_engages_with_empty_inner_and_dead_disk_refresh_token() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let scope = cfg.auth_scope(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - - // Disk: an expired token carrying the (dead) refresh_token the OIDC - // refresher resolves. `inner` stays empty. - let dead = GrokAuth { - key: "disk-dead".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-dead".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }; - let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default(); - store.insert(scope, dead); - write_auth_json(&dir.path().join("auth.json"), &store).unwrap(); - assert!(mgr.current_or_expired().is_none(), "inner must be empty"); - +fn install_ok_refresher(m: &Arc) -> Arc { let calls = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(FailingRefresher { - call_count: calls.clone(), + m.set_refresher(Arc::new(OkRefresher { + calls: calls.clone(), })); - - for _ in 0..5 { - let _ = mgr - .refresh_chain(TokenType::OidcSession, RefreshReason::ServerRejected) - .await; - } - assert_eq!( - calls.load(Ordering::SeqCst), - 1, - "storm cap must hold the IdP to one call even with empty inner + dead disk RT", - ); + calls } -/// Record/check consistency: in-mem and disk are DIFFERENT stale credentials. -/// The refresher resolves & sends the DISK refresh token, so the verdict must be -/// keyed on THAT — proven by swapping the in-mem bearer afterward and confirming -/// the verdict still caps the storm (a verdict mis-keyed to the in-mem bearer -/// would read absent after the swap and re-hit the IdP). The `tried_key == None` -/// fallback (external-binary flow → `attempted_verdict_key`) is covered by -/// `storm_cap_engages_with_empty_inner_and_dead_disk_refresh_token`. -#[tokio::test] -async fn verdict_not_keyed_on_in_mem_bearer() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let scope = cfg.auth_scope(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); +// ── Dynamic refresh threshold (PRD: max(300, expires_in × 0.5)) ───────── - // in-mem: stale bearer K_mem (expired, with RT). - mgr.hot_swap(GrokAuth { - key: "mem-stale".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-mem".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }); - // disk: a DIFFERENT stale credential K_disk (expired, with RT) — what the - // refresher resolves first. - let disk = GrokAuth { - key: "disk-stale".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-disk".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }; - let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default(); - store.insert(scope, disk); - write_auth_json(&dir.path().join("auth.json"), &store).unwrap(); +/// A 7200s-lifetime token with 3000s left is inside the 3600s threshold: +/// `current()` hides it (refresh due) while `expired_auth()` still exposes +/// it (wire-valid bearer for senders). +#[test] +fn threshold_hides_current_but_keeps_expired_auth() { + let (_d, m) = mgr(); + m.hot_swap(session("at", "rt", 7200, 3000)); + assert!(m.current().is_none(), "inside threshold → refresh due"); + assert_eq!(m.expired_auth().map(|a| a.key), Some("at".into())); + assert!(m.is_expired()); + assert!(m.has_usable_token(), "still wire-valid"); - let calls = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(TriedKeyFailRefresher { - tried_key: "disk-stale".into(), - call_count: calls.clone(), - })); - - let _ = mgr - .refresh_chain(TokenType::OidcSession, RefreshReason::ServerRejected) - .await; - assert_eq!( - calls.load(Ordering::SeqCst), - 1, - "first call hits the IdP once" - ); - - // Swap the in-mem bearer to yet another stale key: a verdict mis-keyed to - // the old in-mem bearer would now read absent. - mgr.hot_swap(GrokAuth { - key: "mem-stale-2".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-mem-2".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }); - - let _ = mgr - .refresh_chain(TokenType::OidcSession, RefreshReason::ServerRejected) - .await; - assert_eq!( - calls.load(Ordering::SeqCst), - 1, - "verdict keyed on the tried disk credential must survive an in-mem swap", - ); + // 5000s left is outside the threshold → fresh. + m.hot_swap(session("at2", "rt", 7200, 5000)); + assert_eq!(m.current().map(|a| a.key), Some("at2".into())); + assert!(!m.is_expired()); } -/// Success → persist-failure → transient: a refresh that obtains a fresh token -/// but cannot write it to disk must surface `Transient` AND still swap the -/// in-memory bearer to the fresh token (the "always update in-memory even if the -/// disk write failed" invariant — without it a disk hiccup strands the session). -/// The write is failed deterministically (root-safe) by planting a *directory* -/// at the atomic-write temp path so `open_secure_file` hits `EISDIR`; the -/// auth.json read (file absent) and the file lock still succeed. +/// Hard expiry: a genuinely past-expiry token is not usable. +#[test] +fn hard_expired_token_is_not_usable() { + let (_d, m) = mgr(); + m.hot_swap(session("at", "rt", 3600, -10)); + assert!(m.current().is_none()); + assert!(m.expired_auth().is_some(), "kept for its refresh token"); + assert!(!m.has_usable_token()); +} + +// ── auth() dispatch ───────────────────────────────────────────────────── + #[tokio::test] -async fn refresh_persist_failure_is_transient_but_swaps_in_memory() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - // Expired in-mem bearer so the chain proceeds to the IdP (no early return). - mgr.hot_swap(GrokAuth { - key: "stale".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }); - - // `write_auth_json_atomic` writes `auth.json..tmp` then renames; a - // directory there makes the temp-file open fail with EISDIR (enforced even - // for root), so the persist fails while the read/lock paths are unaffected. - std::fs::create_dir( - dir.path() - .join(format!("auth.json.{}.tmp", std::process::id())), - ) - .unwrap(); - - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: Arc::new(AtomicU32::new(0)), - delay: StdDuration::ZERO, - })); - - let err = mgr - .refresh_chain(TokenType::OidcSession, RefreshReason::ServerRejected) - .await - .expect_err("persist failure must surface an error"); - assert!( - matches!(err, AuthError::Refresh(RefreshTokenError::Transient(_))), - "persist failure must be transient (retryable), got {err:?}", - ); - assert_eq!( - mgr.current().map(|a| a.key), - Some("fresh-token".to_string()), - "in-memory bearer must hold the fresh token despite the failed disk write", - ); +async fn auth_returns_not_logged_in_when_empty() { + let (_d, m) = mgr(); + let err = m.auth().await.unwrap_err(); + assert!(matches!(err, AuthError::NotLoggedIn), "got {err:?}"); } #[tokio::test] -async fn auth_concurrent_refresh_deduplicates() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - let expired = GrokAuth { - key: "expired-key".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-old".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(expired); - - let call_count = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: call_count.clone(), - delay: StdDuration::from_millis(50), - })); - - // Spawn 4 concurrent tasks that all call auth(). - let mut handles = Vec::new(); - for _ in 0..4 { - let m = mgr.clone(); - handles.push(tokio::spawn(async move { m.auth().await })); - } - - let mut results = Vec::new(); - for h in handles { - results.push(h.await.unwrap()); - } - - // All 4 should succeed with the same fresh token. - for r in &results { - assert_eq!( - r.as_ref().unwrap().key, - "fresh-token", - "all tasks must get the fresh token" - ); - } - - // The refresher should have been called exactly once. - assert_eq!( - call_count.load(Ordering::SeqCst), - 1, - "refresher must be called exactly once despite 4 concurrent callers" - ); +async fn auth_fast_path_returns_valid_cached_token() { + let (_d, m) = mgr(); + m.hot_swap(session("at-valid", "rt", 7200, 7200)); + let auth = m.auth().await.unwrap(); + assert_eq!(auth.key, "at-valid"); } #[tokio::test] -async fn auth_permanent_failure_stops_retries() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - let expired = GrokAuth { - key: "expired-key".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-old".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(expired); - - let call_count = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(FailingRefresher { - call_count: call_count.clone(), - })); - - // First auth(): refresher called, refresh_chain records permanent failure. - let err1 = mgr.auth().await.unwrap_err(); - assert!( - matches!(err1, AuthError::Refresh(RefreshTokenError::Permanent(_))), - "first call should return PermanentFailure, got: {err1:?}" - ); - - // Second auth(): permanent failure cached, refresher NOT called. - let err2 = mgr.auth().await.unwrap_err(); - assert!( - matches!(err2, AuthError::Refresh(RefreshTokenError::Permanent(_))), - "second call should return PermanentFailure, got: {err2:?}" - ); - - // Refresher must have been called exactly once. - assert_eq!( - call_count.load(Ordering::SeqCst), - 1, - "refresher must be called exactly once" - ); - - // hot_swap clears permanent failure; subsequent auth() succeeds. - let valid = GrokAuth { - key: "new-valid-key".into(), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(valid); - assert_eq!(mgr.auth().await.unwrap().key, "new-valid-key"); -} - -/// auth() re-reads disk via pick_up_sibling_token and returns the -/// sibling-written token when the in-memory token is stale. -#[tokio::test] -async fn auth_legacy_session_picks_up_sibling_disk_token() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let scope = cfg.auth_scope(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - - mgr.hot_swap(GrokAuth { - key: "stale-oidc".into(), - auth_mode: AuthMode::Oidc, - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }); - - // Sibling writes a valid token to disk. - let fresh = GrokAuth { - key: "fresh-from-sibling".into(), - auth_mode: AuthMode::Oidc, - expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() - }; - let mut store = AuthStore::new(); - store.insert(scope, fresh); - write_auth_json(&dir.path().join("auth.json"), &store).unwrap(); - - let auth = mgr.auth().await.expect("should pick up sibling token"); - assert_eq!(auth.key, "fresh-from-sibling"); -} - -/// refresh_chain returns TransientFailure when the refresher reports one. -#[tokio::test] -async fn refresh_chain_surfaces_transient_failure() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.hot_swap(GrokAuth { - key: "expired".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }); - - struct TransientRefresher; - #[async_trait::async_trait] - impl TokenRefresher for TransientRefresher { - async fn refresh(&self, _: RefreshReason) -> crate::auth::refresh::RefreshOutcome { - crate::auth::refresh::RefreshOutcome::TransientFailure { - message: "idp timeout".into(), - } - } - } - mgr.set_refresher(Arc::new(TransientRefresher)); - - let err = mgr.auth().await.unwrap_err(); - assert!( - matches!(err, AuthError::Refresh(RefreshTokenError::Transient(_))), - "TransientFailure should surface as a transient refresh error, got {err:?}" - ); -} - -/// Regression: `current()` and `auth()` must agree on whether an -/// expired API key is usable. Pre-fix, `current()` filtered with -/// `!is_token_expired()` (returning None) while the `auth()` -/// `TokenType::ApiKey` branch cloned the stale entry, so the UI saw -/// "logged out" while downstream consumers (trace upload, MCP, -/// embeddings) sent the stale key and hit 401. -#[tokio::test] -async fn auth_returns_expired_api_key_consistently_with_current() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - // Seed an API key that is past the 30-day TTL: `create_time` 60 - // days ago and no `expires_at`. `is_token_expired` falls through - // to the TTL check and reports `true`. - let expired_key = GrokAuth { - key: "stale-api-key".into(), +async fn auth_expired_api_key_surfaces_token_expired_no_refresh() { + let (_d, m) = mgr(); + m.hot_swap(KimiAuth { + key: "sk-old".into(), auth_mode: AuthMode::ApiKey, - create_time: Utc::now() - Duration::days(60), - expires_at: None, - refresh_token: None, - ..GrokAuth::test_default() - }; - mgr.hot_swap(expired_key); - - // UI / sync read path: the stale key is filtered out. - assert!( - mgr.current().is_none(), - "current() must hide the expired api_key (matches UI/login state)" - ); - - // Async path: must NOT clone the stale key for downstream - // consumers. Surface `TokenExpiredNoRefresh` so callers can - // funnel the user back through `grok login`. - let err = mgr.auth().await.unwrap_err(); - assert!( - matches!(err, AuthError::TokenExpiredNoRefresh), - "auth() must report TokenExpiredNoRefresh for expired api_key, got: {err:?}", - ); - assert!( - mgr.get_valid_token().await.is_err(), - "get_valid_token() must error rather than return the stale key" - ); - - // Sanity: a fresh API key restores both paths. - let fresh_key = GrokAuth { - key: "fresh-api-key".into(), - auth_mode: AuthMode::ApiKey, - create_time: Utc::now(), - expires_at: None, - refresh_token: None, - ..GrokAuth::test_default() - }; - mgr.hot_swap(fresh_key); - assert_eq!( - mgr.current().map(|a| a.key).as_deref(), - Some("fresh-api-key") - ); - assert_eq!( - mgr.get_valid_token().await.ok().as_deref(), - Some("fresh-api-key") - ); + create_time: Utc::now() - Duration::days(31), + ..KimiAuth::test_default() + }); + let err = m.auth().await.unwrap_err(); + assert!(matches!(err, AuthError::TokenExpiredNoRefresh), "{err:?}"); } -/// Regression: after a permanent refresh failure (e.g. `invalid_grant`), -/// the proactive refresh task must back off rather than hammer -/// `auth()` in a tight loop. Pre-fix, an expired token + cached -/// PermanentFailure caused `sleep_dur=0` -> `auth()` -> error -> repeat. -/// -/// Verified by observing the loop's iteration counter directly: in a -/// 300ms window we tolerate at most a few iterations (one for the -/// initial failure-recording pass, then back-off). Pre-fix the -/// counter would have been in the thousands. #[tokio::test] -async fn proactive_refresh_backs_off_on_permanent_failure() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - // Past-expiry OIDC token: without the backoff guard, the - // proactive loop computes sleep_dur=0 forever. - let expired = GrokAuth { - key: "expired".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(expired); - - // Refresher returns invalid_grant the first time it is called and - // counts every invocation. After the first call records the - // permanent failure, the proactive loop must skip subsequent - // calls until the failure is cleared. - let call_count = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(FailingRefresher { - call_count: call_count.clone(), - })); - - let cancel = CancellationToken::new(); - mgr.start_proactive_refresh(cancel.clone()); - - // Give the loop ample time to observe the failure and back off. - // The 300ms window is two orders of magnitude shorter than the - // 5-minute BACKOFF_INTERVAL, so a backed-off loop completes at - // most a couple of iterations: the initial pass that records the - // permanent failure, optionally a few re-check passes if the - // executor races, then sleeps for `BACKOFF_INTERVAL`. - tokio::time::sleep(StdDuration::from_millis(300)).await; - - let iterations = mgr.proactive_iteration_count(); - let after_failure = call_count.load(Ordering::SeqCst); - - // Direct observation of loop progress: a busy-loop produces - // hundreds-to-thousands of iterations in 300ms, the backed-off - // loop produces <= 5. - assert!( - iterations <= 5, - "proactive refresh busy-looped after permanent failure: \ - {iterations} iterations (refresher calls: {after_failure})", - ); - // Refresher invocation count is a secondary check: the - // permanent_failure short-circuit in `refresh_chain` (added in - // this PR) means at most 1 invocation here. - assert!( - after_failure <= 1, - "refresher must be invoked at most once before the permanent \ - failure is recorded, got {after_failure} calls" - ); - assert!( - mgr.permanent_failure().is_some(), - "permanent failure must be cached after invalid_grant", - ); - cancel.cancel(); +async fn auth_expired_session_refreshes_via_chain() { + let (_d, m) = mgr(); + m.hot_swap(session("at-old", "rt-old", 3600, -10)); + let calls = install_ok_refresher(&m); + let auth = m.auth().await.unwrap(); + assert_eq!(auth.key, "at-refreshed"); + assert_eq!(calls.load(AtomicOrdering::SeqCst), 1); + // Refresh persisted: a fresh manager on the same home adopts it. + assert_eq!(m.current().map(|a| a.key), Some("at-refreshed".into())); } -/// Regression: `start_proactive_refresh` must be -/// idempotent. Calling it twice on the same `Arc` was -/// previously valid (no guard) and would `tokio::spawn` two -/// background tasks racing on the same in-memory state. -/// -/// Asserting on `proactive_iteration_count` is not a meaningful signal -/// because the test fixture (ApiKey + expires_at: None) made every -/// spawned task sleep for `BACKOFF_INTERVAL` immediately. With or -/// without the guard the iteration counter stayed at 0, so that -/// assertion was vacuous (removing the guard left the test passing). The -/// fix is to assert on the new `proactive_start_count()` accessor, -/// which is bumped *inside* the `compare_exchange` success branch -/// in `start_proactive_refresh` -- so it is exactly 1 if the guard -/// fires and N otherwise. This directly observes the invariant -/// instead of inferring it from loop-iteration mechanics. +/// Refresh success persists to the store so a sibling manager adopts it. #[tokio::test] -async fn start_proactive_refresh_is_idempotent() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); +async fn refresh_success_is_visible_to_sibling_manager() { + let (dir, m) = mgr(); + m.hot_swap(session("at-old", "rt-old", 3600, -10)); + install_ok_refresher(&m); + m.auth().await.unwrap(); - let stale_api_key = GrokAuth { - key: "stale-api-key".into(), - auth_mode: AuthMode::ApiKey, - create_time: Utc::now() - Duration::days(60), - expires_at: None, - refresh_token: None, - ..GrokAuth::test_default() - }; - mgr.hot_swap(stale_api_key); - - let cancel = CancellationToken::new(); - // First call spawns the task; subsequent calls must be no-ops. - mgr.start_proactive_refresh(cancel.clone()); - mgr.start_proactive_refresh(cancel.clone()); - mgr.start_proactive_refresh(cancel.clone()); - - // Direct observation of the guard's behavior. Pre-fix: 3. - // Post-fix: exactly 1. + let sibling = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); assert_eq!( - mgr.proactive_start_count(), - 1, - "start_proactive_refresh idempotency guard failed; expected exactly \ - 1 spawn after 3 calls", + sibling.current().map(|a| a.key), + Some("at-refreshed".into()), + "sibling must load the rotated credential from the store" ); - - cancel.cancel(); } -/// Proactive path: near-expiry OIDC token -> background task fires -/// refresh_chain(PreRequest) -> consumer sees fresh token. +/// Refresh success wakes `wait_for_token_refresh` waiters. #[tokio::test] -async fn proactive_refresh_and_consumer_see_fresh_token_end_to_end() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - // expires_at inside the 5-min buffer -> proactive fires immediately. - mgr.hot_swap(GrokAuth { - key: "soon-to-expire".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-original".into()), - expires_at: Some(Utc::now() + Duration::seconds(2)), - ..GrokAuth::test_default() - }); - - let call_count = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: call_count.clone(), - delay: StdDuration::from_millis(0), - })); - - let cancel = CancellationToken::new(); - mgr.start_proactive_refresh(cancel.clone()); - tokio::time::sleep(StdDuration::from_millis(500)).await; - - assert!(call_count.load(Ordering::SeqCst) >= 1); - assert_eq!(mgr.get_valid_token().await.unwrap(), "fresh-token"); - - cancel.cancel(); -} - -/// Reactive path: expired OIDC token -> try_recover_unauthorized -> -/// refresh_chain(ServerRejected) -> refresher -> consumer sees fresh token. -#[tokio::test] -async fn reactive_401_recovery_produces_fresh_token_end_to_end() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - mgr.hot_swap(GrokAuth { - key: "expired-bearer".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-valid".into()), - expires_at: Some(Utc::now() - Duration::minutes(10)), - ..GrokAuth::test_default() - }); - - let call_count = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: call_count.clone(), - delay: StdDuration::from_millis(0), - })); - - assert!(mgr.try_recover_unauthorized().await); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - assert_eq!(mgr.get_valid_token().await.unwrap(), "fresh-token"); -} - -// refresh_chain permanent-failure short-circuit via recovery is tested -// in recovery::tests::refresh_authority_short_circuits_on_cached_permanent_failure. - -/// Different disk RT with expired AT: PermanentFailure is recorded -/// (not demoted to transient), stopping the retry loop. -#[tokio::test] -async fn refresh_chain_records_permanent_failure_when_disk_rt_differs_but_at_expired() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let scope = cfg.auth_scope(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - - // Memory has rt-old; disk has rt-new (different RT) but its - // access_token is also expired so try_use_disk_token rejects it - // and we fall through to the refresher. - let stale = GrokAuth { - key: "stale-key".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-old".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - oidc_issuer: Some("https://issuer.example".into()), - oidc_client_id: Some("client-1".into()), - ..GrokAuth::test_default() +async fn refresh_notifies_waiters() { + let (_d, m) = mgr(); + m.hot_swap(session("at-old", "rt-old", 3600, -10)); + install_ok_refresher(&m); + let waiter = { + let m = m.clone(); + tokio::spawn(async move { + m.wait_for_token_refresh(std::time::Duration::from_secs(5)) + .await + }) }; - mgr.hot_swap(stale); - - let sibling = GrokAuth { - key: "sibling-key".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-new".into()), - expires_at: Some(Utc::now() - Duration::minutes(30)), - oidc_issuer: Some("https://issuer.example".into()), - oidc_client_id: Some("client-1".into()), - ..GrokAuth::test_default() - }; - let mut store = AuthStore::new(); - store.insert(scope, sibling); - write_auth_json(&dir.path().join("auth.json"), &store).unwrap(); - - struct FailingRefresher; - #[async_trait::async_trait] - impl crate::auth::refresh::TokenRefresher for FailingRefresher { - async fn refresh( - &self, - _reason: crate::auth::manager::RefreshReason, - ) -> crate::auth::refresh::RefreshOutcome { - crate::auth::refresh::RefreshOutcome::permanent( - crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected, - None, - ) - } - } - mgr.set_refresher(Arc::new(FailingRefresher)); - - let err = mgr.auth().await.unwrap_err(); - // An expired disk AT means the sibling is dead too — the failure is - // permanent (not demoted to transient). Credentials are retained; the - // scoped verdict is cached and stops the retry storm. - assert!( - matches!(err, AuthError::Refresh(RefreshTokenError::Permanent(_))), - "must surface a permanent failure when disk AT is expired, got: {err:?}", - ); - assert!( - mgr.permanent_failure().is_some(), - "verdict must be cached (scoped to the retained credential)", - ); - // No-clear invariant: a refresh failure must NOT delete auth.json (a future - // regression that re-adds disk-clear-on-invalid_grant would fail here). - assert!( - mgr.read_disk_auth().is_some(), - "invalid_grant must not delete auth.json (no auto-clear)", - ); - // Second attempt short-circuits on the cached verdict — no extra IdP call. - assert!(matches!( - mgr.auth().await.unwrap_err(), - AuthError::Refresh(RefreshTokenError::Permanent(_)) - )); -} - -/// Different disk RT with valid AT: adopt the sibling's token directly. -#[tokio::test] -async fn refresh_chain_demotes_to_transient_when_disk_rt_differs_and_at_valid() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let scope = cfg.auth_scope(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - - let stale = GrokAuth { - key: "stale-key".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-old".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - oidc_issuer: Some("https://issuer.example".into()), - oidc_client_id: Some("client-1".into()), - ..GrokAuth::test_default() - }; - mgr.hot_swap(stale); - - let sibling = GrokAuth { - key: "sibling-key".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-new".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - oidc_issuer: Some("https://issuer.example".into()), - oidc_client_id: Some("client-1".into()), - ..GrokAuth::test_default() - }; - let mut store = AuthStore::new(); - store.insert(scope, sibling); - write_auth_json(&dir.path().join("auth.json"), &store).unwrap(); - - let calls = Arc::new(AtomicU32::new(0)); - struct CountingFailRefresher(Arc); - #[async_trait::async_trait] - impl crate::auth::refresh::TokenRefresher for CountingFailRefresher { - async fn refresh( - &self, - _reason: crate::auth::manager::RefreshReason, - ) -> crate::auth::refresh::RefreshOutcome { - self.0.fetch_add(1, Ordering::SeqCst); - crate::auth::refresh::RefreshOutcome::permanent( - crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected, - None, - ) - } - } - mgr.set_refresher(Arc::new(CountingFailRefresher(calls.clone()))); - - let result = mgr.auth().await; - assert!( - result.is_ok(), - "should adopt valid sibling token: {result:?}" - ); - assert_eq!(result.unwrap().key, "sibling-key"); - assert_eq!( - calls.load(Ordering::SeqCst), - 0, - "refresher must not be called when disk has a valid token" - ); -} - -/// Regression: after `clear()` the verdict must *read as absent* -/// — nothing drops it explicitly; it is scoped to the cleared credential and -/// reads through as `None` once that credential is gone — so subsequent -/// `auth()` reports the more useful `NotLoggedIn` (rather than the stale -/// `invalid_grant` from the just-cleared session). -#[tokio::test] -async fn permanent_failure_reads_absent_after_clear_so_auth_reports_not_logged_in() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - // Seed + record a permanent failure (as if invalid_grant fired). - let session = GrokAuth { - key: "broken-session".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-revoked".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(session); - record_permanent_failure( - &mgr, - crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected, - ); - assert!(mgr.permanent_failure().is_some()); - - // User runs `grok logout` which calls clear(). - mgr.clear().unwrap(); - - // The diagnostic the user now sees on the next request should be - // "Not logged in. Run `grok login`.", not the stale invalid_grant. - let err = mgr.auth().await.unwrap_err(); - assert!( - matches!(err, AuthError::NotLoggedIn), - "auth() after clear() must report NotLoggedIn, got: {err:?}", - ); - assert!( - mgr.permanent_failure().is_none(), - "the credential-scoped verdict must read as absent after clear()", - ); - - // Same check for the hot_swap_clear() path. - let session = GrokAuth { - key: "broken-2".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-2".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(session); - record_permanent_failure( - &mgr, - crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected, - ); - mgr.clear_in_memory(); - let err = mgr.auth().await.unwrap_err(); - assert!( - matches!(err, AuthError::NotLoggedIn), - "auth() after hot_swap_clear() must report NotLoggedIn, got: {err:?}", - ); -} - -/// `PERMANENT_FAILURE_TTL` means "5 *real* minutes", not "5 awake minutes": -/// a recoverable permanent failure cached just before a system suspend must -/// expire while the machine sleeps. The monotonic clock pauses across suspend, -/// so expiry is judged on both clocks (see `ScopedRefreshFailure::recorded_at`) -/// — this simulates the suspend by rewinding only the wall-clock arm and -/// asserts the failure no longer short-circuits `auth()` on wake. -#[tokio::test] -async fn permanent_failure_expires_on_wall_clock_across_sleep() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - // Seed a credential so the verdict scopes to it (an unscoped verdict - // reads through as absent), using the non-sticky `Other` reason — the - // "transient escalation just before lid close" case the TTL exists for. - mgr.hot_swap(GrokAuth { - key: "tok".into(), - ..GrokAuth::test_default() - }); - record_permanent_failure(&mgr, crate::auth::error::RefreshTokenFailedReason::Other); - assert!( - mgr.permanent_failure().is_some(), - "freshly recorded failure must be live on both clocks", - ); - - // Simulate a >TTL suspend: monotonic elapsed stays ~0 (paused during - // sleep), wall clock advanced past the TTL. - mgr.force_permanent_failure_wall_aged_out(); - - assert!( - mgr.permanent_failure().is_none(), - "a slept-through TTL must expire the cached permanent failure on wake", - ); - assert!( - !mgr.has_permanent_failure(), - "has_permanent_failure must agree with permanent_failure()", - ); -} - -// -- Regression: api_key in config.toml must not block OIDC refresh -- - -/// When a user has an OIDC session (auth.json) AND a model with api_key -/// in config.toml, the OIDC token must still be refreshable. auth() -/// checks TokenType (from AuthManager), not the global auth_method_id. -#[tokio::test] -async fn oidc_refresh_not_blocked_by_model_api_key() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - // Expired OIDC token (user has config.toml with api_key on another model). - let expired_oidc = GrokAuth { - key: "expired-session-token".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("valid-rt".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(expired_oidc); - - // TokenType is OidcSession regardless of what models exist in config. - assert_eq!(mgr.token_type(), TokenType::OidcSession); - - // auth() must attempt OIDC refresh, not short-circuit as ApiKey. - let call_count = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: call_count.clone(), - delay: StdDuration::from_millis(10), - })); - - let result = mgr.auth().await; - assert!(result.is_ok(), "auth() should succeed via OIDC refresh"); - assert_eq!(result.unwrap().key, "fresh-token"); - assert_eq!(call_count.load(Ordering::SeqCst), 1); -} - -// -- direct unit tests for `compute_proactive_sleep` -------- -// -// The proactive task's gate chain is a small pure function; testing -// it directly (rather than through `start_proactive_refresh` and a -// sleep window) gives us per-branch coverage that would have caught -// the original vacuity in seconds. Each test below pins one -// arm of `compute_proactive_sleep`. - -/// Permanent-failure cached -> backs off (>= BACKOFF_INTERVAL, plus jitter). -#[test] -fn compute_proactive_sleep_permanent_failure_returns_backoff() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - let oidc = GrokAuth { - key: "x".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(oidc); - record_permanent_failure( - &mgr, - crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected, - ); - let sleep = compute_proactive_sleep(&mgr); - assert!( - sleep >= BACKOFF_INTERVAL && sleep < BACKOFF_INTERVAL + StdDuration::from_secs(60), - "expected backoff + jitter, got {sleep:?}" - ); -} - -/// Non-refreshable types (LegacySession, ApiKey, None) -> BACKOFF_INTERVAL -/// even when expires_at is past. This is the gate the original -/// test failed to exercise. -#[test] -fn compute_proactive_sleep_non_refreshable_returns_backoff() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - // Inject a refresher so the "no refresher" branch doesn't mask - // the gate we're testing. - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: Arc::new(AtomicU32::new(0)), - delay: StdDuration::from_millis(0), - })); - - // (a) LegacySession (WebLogin) + Some(past) -- the canonical - // scenario where the absence of the gate produces a busy-loop. - mgr.hot_swap(GrokAuth { - key: "legacy".into(), - auth_mode: AuthMode::WebLogin, - create_time: Utc::now() - Duration::hours(2), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }); - assert_eq!(mgr.token_type(), TokenType::LegacySession); - assert_eq!(compute_proactive_sleep(&mgr), BACKOFF_INTERVAL); - - // (b) ApiKey + Some(past). - mgr.hot_swap(GrokAuth { - key: "api".into(), - auth_mode: AuthMode::ApiKey, - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }); - assert_eq!(mgr.token_type(), TokenType::ApiKey); - assert_eq!(compute_proactive_sleep(&mgr), BACKOFF_INTERVAL); - - // (c) None (no credentials loaded). - mgr.clear_in_memory(); - assert_eq!(mgr.token_type(), TokenType::None); - assert_eq!(compute_proactive_sleep(&mgr), BACKOFF_INTERVAL); -} - -/// Sleep gate raised -> BACKOFF_INTERVAL even for a refreshable token past -/// its expiry. Without this gate `refresh_chain` defers every attempt while -/// the proactive loop spins at `sleep_dur=0` (the busy-loop). -#[test] -fn compute_proactive_sleep_sleep_gated_returns_backoff() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: Arc::new(AtomicU32::new(0)), - delay: StdDuration::from_millis(0), - })); - // Refreshable OidcSession past the early-invalidation boundary: without - // the gate this returns 0 (would busy-loop). - mgr.hot_swap(GrokAuth { - key: "oidc".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }); - assert_eq!( - compute_proactive_sleep(&mgr), - StdDuration::from_secs(0), - "precondition: ungated expired refreshable token yields a 0 sleep" - ); - - mgr.set_system_sleep_imminent(true); - assert_eq!( - compute_proactive_sleep(&mgr), - BACKOFF_INTERVAL, - "sleep gate must back the proactive loop off instead of busy-looping" - ); -} - -/// Dark wake -> BACKOFF_INTERVAL even for a refreshable token past its expiry. -/// `refresh_chain` defers every attempt during a dark wake (to avoid an IdP -/// refresh straddling an unsignaled re-sleep), so the proactive loop must back -/// off rather than spin at `sleep_dur=0`. -#[test] -fn compute_proactive_sleep_dark_wake_returns_backoff() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: Arc::new(AtomicU32::new(0)), - delay: StdDuration::from_millis(0), - })); - // Refreshable OidcSession past the early-invalidation boundary: without - // the dark-wake gate this returns 0 (would busy-loop). - mgr.hot_swap(GrokAuth { - key: "oidc".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }); - assert_eq!( - compute_proactive_sleep(&mgr), - StdDuration::from_secs(0), - "precondition: non-dark-wake expired refreshable token yields a 0 sleep" - ); - - mgr.set_dark_wake_for_test(true); - assert_eq!( - compute_proactive_sleep(&mgr), - BACKOFF_INTERVAL, - "dark wake must back the proactive loop off instead of busy-looping" - ); - - // Returning to a full wake re-enables immediate refresh. - mgr.set_dark_wake_for_test(false); - assert_eq!( - compute_proactive_sleep(&mgr), - StdDuration::from_secs(0), - "full wake must allow the refresh to proceed again" - ); -} - -/// No refresher configured -> BACKOFF_INTERVAL even for refreshable -/// types. This is the startup-race guard. -#[test] -fn compute_proactive_sleep_no_refresher_returns_backoff() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.hot_swap(GrokAuth { - key: "oidc".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }); - // No `set_refresher` call -- the refresher slot is None. - assert!(mgr.refresher.read().is_none()); - assert_eq!(compute_proactive_sleep(&mgr), BACKOFF_INTERVAL); -} - -/// Refreshable type + no `expires_at` -> BACKOFF_INTERVAL (the -/// "external binary that doesn't return expiry" case). -#[test] -fn compute_proactive_sleep_refreshable_no_expiry_returns_backoff() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: Arc::new(AtomicU32::new(0)), - delay: StdDuration::from_millis(0), - })); - mgr.hot_swap(GrokAuth { - key: "external".into(), - auth_mode: AuthMode::External, - expires_at: None, - ..GrokAuth::test_default() - }); - assert_eq!(mgr.token_type(), TokenType::ExternalBinary); - assert_eq!(compute_proactive_sleep(&mgr), BACKOFF_INTERVAL); -} - -/// Refreshable type + `Some(past)` and gates pass -> sleep_dur = 0 -/// (refresh now). This is the "happy path" the gates don't block. -#[test] -fn compute_proactive_sleep_refreshable_past_expiry_returns_zero() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: Arc::new(AtomicU32::new(0)), - delay: StdDuration::from_millis(0), - })); - mgr.hot_swap(GrokAuth { - key: "oidc".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }); - assert_eq!(mgr.token_type(), TokenType::OidcSession); - assert_eq!(compute_proactive_sleep(&mgr), StdDuration::from_secs(0)); -} - -/// Refreshable type + `Some(future)` and gates pass -> sleep_dur ~= -/// expires_at - buffer (positive, <= delta). We use a 1-hour horizon -/// and assert the result is in a sane range rather than an exact value -/// (executor scheduling jitter). -#[test] -fn compute_proactive_sleep_refreshable_future_expiry_returns_delta() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: Arc::new(AtomicU32::new(0)), - delay: StdDuration::from_millis(0), - })); - let expires_at = Utc::now() + Duration::hours(1); - mgr.hot_swap(GrokAuth { - key: "oidc".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt".into()), - expires_at: Some(expires_at), - ..GrokAuth::test_default() - }); - let dur = compute_proactive_sleep(&mgr); - // Expected: 1h - 5min (early_invalidation) - jitter (0–60s) ≈ 54–55min. - // Range is generous (51–59min) to absorb both clock granularity and - // the random jitter added by `compute_proactive_sleep`. - assert!( - dur >= StdDuration::from_secs(51 * 60) && dur <= StdDuration::from_secs(59 * 60), - "expected ~55min, got {dur:?}", - ); -} - -/// `permanent_failure` cache auto-expires after `PERMANENT_FAILURE_TTL`, -/// so a misclassified transient IdP error (e.g. `invalid_client` during -/// an OAuth client rotation) doesn't permanently log the user out. -#[tokio::test] -async fn permanent_failure_expires_after_ttl() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.hot_swap(GrokAuth { - key: "tok".into(), - ..GrokAuth::test_default() - }); - record_permanent_failure( - &mgr, - crate::auth::error::RefreshTokenFailedReason::ClientRejected, - ); - assert!( - mgr.permanent_failure().is_some(), - "freshly recorded failure should be sticky" - ); - mgr.force_permanent_failure_aged_out(); - assert!( - mgr.permanent_failure().is_none(), - "aged-out recoverable failure should auto-expire so a retry can succeed" - ); - - // A revoked refresh token never self-heals: the verdict is sticky past the - // TTL (only a credential change clears it). Stops re-pinging a dead RT. - record_permanent_failure( - &mgr, - crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected, - ); - mgr.force_permanent_failure_aged_out(); - assert!( - mgr.permanent_failure().is_some(), - "RefreshTokenRejected must stay sticky past the TTL", - ); -} - -/// The sticky verdict is exempt from BOTH TTL clocks — the monotonic arm -/// (awake time) AND the wall arm (real time across a suspend, added by the -/// sleep-straddle fix). A revoked refresh token never self-heals with time: -/// re-pinging the IdP with it can only fail again, so no amount of aging on -/// either clock may expire the verdict. Only a credential change heals it — -/// the scoped read-through pinned by the `hot_swap` phase below. This is a -/// composition guard: the sticky/non-sticky split and the wall-clock arm -/// landed separately, so neither parent change could test their intersection. -#[tokio::test] -async fn sticky_verdict_survives_both_clocks_but_not_a_credential_change() { - // Guard against a vacuous pass: with < TTL of monotonic uptime the aging - // hook's `checked_sub` no-ops, and a *fresh* verdict would trivially - // satisfy the survival asserts below. - if std::time::Instant::now() - .checked_sub(PERMANENT_FAILURE_TTL + StdDuration::from_secs(1)) - .is_none() - { - eprintln!( - "skipping sticky_verdict_survives_both_clocks: host uptime < PERMANENT_FAILURE_TTL" - ); - return; - } - - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.hot_swap(GrokAuth { - key: "dead".into(), - ..GrokAuth::test_default() - }); - record_permanent_failure( - &mgr, - crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected, - ); - - // Age the verdict past the TTL on the monotonic clock AND rewind the - // wall-clock arm past it (what a >TTL suspend looks like to the reader). - mgr.force_permanent_failure_aged_out(); - mgr.force_permanent_failure_wall_aged_out(); - match mgr.permanent_failure() { - Some(AuthError::Refresh(RefreshTokenError::Permanent(e))) => assert_eq!( - e.reason, - crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected, - "the surviving verdict must carry the sticky reason", - ), - other => panic!("sticky verdict must survive both clocks aging out, got {other:?}"), - } - - // Time never heals it; a credential change does (read-through, no clear). - mgr.hot_swap(GrokAuth { - key: "fresh".into(), - ..GrokAuth::test_default() - }); - assert!( - mgr.permanent_failure().is_none(), - "stickiness must not outlive the credential it is scoped to", - ); -} - -/// The verdict is scoped to the credential that produced it: swapping in a -/// different credential makes it read through as absent, with no explicit -/// clear. -#[tokio::test] -async fn permanent_failure_is_scoped_to_its_credential() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - mgr.hot_swap(GrokAuth { - key: "dead".into(), - ..GrokAuth::test_default() - }); - record_permanent_failure( - &mgr, - crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected, - ); - assert!(mgr.permanent_failure().is_some()); - - // A different credential — no clear call — reads through as no failure. - mgr.hot_swap(GrokAuth { - key: "fresh".into(), - ..GrokAuth::test_default() - }); - assert!( - mgr.permanent_failure().is_none(), - "verdict must not apply to a different credential", - ); -} - -/// The verdict is about the *refresh* token: `auth()` must serve a cached -/// access token that is still within its real `expires_at` (buffer-expired -/// but wire-valid) despite a permanent verdict scoped to that credential, -/// without consulting the refresher. Once the same credential passes real -/// expiry, the bypass no longer applies and the permanent error surfaces. -#[tokio::test] -async fn auth_serves_wire_valid_token_despite_permanent_verdict() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - // CI runs in K8s pods where is_devbox_environment() is true; without this - // the past-expiry phase would mint via devbox recovery instead of - // surfacing the permanent error. - mgr.set_devbox_env_for_test(false); - - // Token in the 5-min buffer (1 min before real expiry): buffer-expired, - // still valid by the IdP's clock. - mgr.hot_swap(GrokAuth { - key: "wire-valid".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-dead".into()), - expires_at: Some(Utc::now() + Duration::minutes(1)), - ..GrokAuth::test_default() - }); - record_permanent_failure( - &mgr, - crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected, - ); - assert!( - mgr.permanent_failure().is_some(), - "verdict must scope to the live credential", - ); - - // A refresher is wired but must never be consulted: the verdict - // short-circuits the chain and the bypass serves the cached bearer. - let call_count = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: call_count.clone(), - delay: StdDuration::ZERO, - })); - - let served = mgr - .auth() - .await - .expect("a wire-valid token must be served despite the verdict"); - assert_eq!( - served.key, "wire-valid", - "auth() must return the cached wire-valid bearer", - ); - assert_eq!( - call_count.load(Ordering::SeqCst), - 0, - "the verdict must gate the refresher; serving the cached token is free", - ); - - // Same credential (same key, so the verdict still scopes to it) past its - // real expiry: the bypass no longer applies. - mgr.hot_swap(GrokAuth { - key: "wire-valid".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-dead".into()), - expires_at: Some(Utc::now() - Duration::minutes(1)), - ..GrokAuth::test_default() - }); - let err = mgr.auth().await.unwrap_err(); - assert!( - matches!(err, AuthError::Refresh(RefreshTokenError::Permanent(_))), - "past real expiry the verdict must surface, got: {err:?}", - ); - assert_eq!( - call_count.load(Ordering::SeqCst), - 0, - "the cached verdict must keep short-circuiting the refresher", - ); -} - -/// Refresh-failure grace: when the in-memory token is in the 5-min -/// early-invalidation buffer AND `refresh_chain` fails, `auth()` -/// returns the cached token if it's still within its real `expires_at`. -/// The user doesn't see a chat-turn failure for an IdP blip during -/// the buffer window. -#[tokio::test] -async fn auth_returns_cached_token_when_refresh_fails_within_real_expiry() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - // Point at an unreachable proxy so refresh_chain fails fast. - let mgr = Arc::new(AuthManager::new(dir.path(), cfg).with_proxy_base_url("http://127.0.0.1:1")); - - // Token in the 5-min buffer (1 min before real expiry) -- past - // the buffer threshold but still valid by the IdP's clock. - let in_buffer = GrokAuth { - key: "still-valid-by-idp".into(), - auth_mode: AuthMode::Oidc, - create_time: Utc::now() - Duration::minutes(55), - user_id: "user-42".into(), - refresh_token: Some("rt".into()), - // Real expiry 1 min away; our 5-min buffer marks it expired. - expires_at: Some(Utc::now() + Duration::minutes(1)), - oidc_issuer: Some("http://127.0.0.1:1".into()), - oidc_client_id: Some("client".into()), - ..GrokAuth::test_default() - }; - mgr.hot_swap(in_buffer); - - let result = mgr.auth().await.expect("grace should return cached token"); - assert_eq!( - result.key, "still-valid-by-idp", - "auth() must return the cached token when refresh fails within real expiry" - ); -} -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn update_writes_disk_before_user_enrichment() { - // Mock /user endpoint that blocks on a Notify before responding. - let release = Arc::new(tokio::sync::Notify::new()); - let release_for_handler = Arc::clone(&release); - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let app = axum::Router::new().route( - "/user", - axum::routing::get(move || { - let r = Arc::clone(&release_for_handler); - async move { - r.notified().await; - axum::Json(serde_json::json!({ - "userId": "enriched-user-id", - "email": "enriched@example.com", - "teamId": "enriched-team", - })) - } - }), - ); - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let mgr = Arc::new( - AuthManager::new(dir.path(), cfg).with_proxy_base_url(&format!("http://127.0.0.1:{port}")), - ); - - // user_id starts empty -- a freshly rotated OIDC token doesn't - // yet know its user_id; that's exactly what /user enriches. - // (If user_id were set AND mismatched the proxy's response, the - // enrichment would correctly bail with reason=user_changed.) - let new_auth = GrokAuth { - key: "rotated-key".into(), - refresh_token: Some("rotated-rt".into()), - user_id: String::new(), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - - // `update()` must return well before the `/user` timeout. The - // proxy handler is blocked on `release.notified()` until we say - // so; if `update()` was awaiting `/user` inline, this would - // hang. - let returned = tokio::time::timeout( - std::time::Duration::from_secs(2), - mgr.update(new_auth.clone()), - ) - .await - .expect("update() must not block on /user") - .expect("update() must succeed"); - assert_eq!(returned.key, "rotated-key"); - - // Disk must already reflect the rotated tokens, even though - // /user has not responded yet. - let on_disk_before = read_auth_json(&dir.path().join("auth.json")).unwrap(); - let entry_before = on_disk_before.values().next().expect("entry written"); - assert_eq!( - entry_before.key, "rotated-key", - "rotated key must be on disk before /user lands" - ); - assert_eq!( - entry_before.refresh_token.as_deref(), - Some("rotated-rt"), - "rotated refresh_token must be on disk before /user lands" - ); - assert_eq!( - entry_before.team_id, None, - "enrichment must not have landed yet" - ); - - // Now release the /user handler and wait for the enrichment - // task to merge into disk. Poll up to 5s. - release.notify_one(); - let auth_path = dir.path().join("auth.json"); - let mut enriched = None; - for _ in 0..50 { - tokio::time::sleep(std::time::Duration::from_millis(100)).await; - let store = read_auth_json(&auth_path).unwrap(); - let entry = store.values().next().unwrap().clone(); - if entry.team_id.is_some() { - enriched = Some(entry); - break; - } - } - let enriched = enriched.expect("enrichment must land within 5s"); - - // Enrichment must have merged in WITHOUT clobbering the rotated - // tokens. - assert_eq!(enriched.key, "rotated-key", "tokens preserved"); - assert_eq!( - enriched.refresh_token.as_deref(), - Some("rotated-rt"), - "refresh_token preserved" - ); - assert_eq!(enriched.team_id.as_deref(), Some("enriched-team")); - assert_eq!(enriched.user_id, "enriched-user-id"); - - server.abort(); -} - -/// Regression: back-to-back `update()` calls with different -/// `refresh_token`s must converge to the LATEST token on disk, even -/// though both spawned enrichment tasks read-modify-write disk -/// concurrently. This locks the property the spawn-task file lock -/// buys us; without it, the next "drop the lock for performance" -/// PR silently regresses (an interleaved enrichment write can -/// resurrect the older `refresh_token`, re-opening the -/// `invalid_grant` race). -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn enrichment_task_preserves_interleaved_token_rotation() { - // /user returns the SAME user_id for both calls so neither - // enrichment aborts via `user_changed`. The 50 ms latency keeps - // task v1 alive past the v2 update. - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let app = axum::Router::new().route( - "/user", - axum::routing::get(|| async { - tokio::time::sleep(std::time::Duration::from_millis(50)).await; - axum::Json(serde_json::json!({ - "userId": "stable-user", - "email": "user@corp.com", - "teamId": "team-alpha", - })) - }), - ); - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let mgr = Arc::new( - AuthManager::new(dir.path(), cfg).with_proxy_base_url(&format!("http://127.0.0.1:{port}")), - ); - - // Same user_id so neither enrichment aborts; only the rotated - // token fields differ -- the property under test. - let auth_v1 = GrokAuth { - key: "key-v1".into(), - refresh_token: Some("rt-v1".into()), - user_id: "stable-user".into(), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - let auth_v2 = GrokAuth { - key: "key-v2".into(), - refresh_token: Some("rt-v2".into()), - user_id: "stable-user".into(), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - - // Two rotations back-to-back. v2's update() lands while v1's - // spawned enrichment task is still in /user. - mgr.update(auth_v1).await.unwrap(); - mgr.update(auth_v2).await.unwrap(); - - // Wait for both spawned tasks to land. Each: 50ms /user + lock - // wait + write. We poll for the eventually-consistent state. - let auth_path = dir.path().join("auth.json"); - let mut final_state = None; - for _ in 0..30 { - tokio::time::sleep(std::time::Duration::from_millis(100)).await; - let store = read_auth_json(&auth_path).unwrap(); - let entry = store.values().next().unwrap().clone(); - // Both rotations done AND enrichment landed. - if entry.refresh_token.as_deref() == Some("rt-v2") && entry.team_id.is_some() { - final_state = Some(entry); - break; - } - } - let final_state = final_state.expect("v2 + enrichment must land within 3s"); - - // Core invariant: v2's tokens survive both enrichment writes. - assert_eq!( - final_state.refresh_token.as_deref(), - Some("rt-v2"), - "v2 refresh_token must survive v1's stale enrichment write" - ); - assert_eq!( - final_state.key, "key-v2", - "v2 access token must survive v1's stale enrichment write" - ); - // Enrichment actually ran. - assert_eq!(final_state.team_id.as_deref(), Some("team-alpha")); - assert_eq!(final_state.user_id, "stable-user"); - - server.abort(); -} - -/// Regression for the user-switch abort path: if disk's `user_id` -/// changes during an in-flight `/user` call (a different user -/// signed in via a sibling process), the spawned enrichment must -/// abort cleanly rather than overlay a previous user's -/// team/org/profile fields onto the new user's entry. -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn enrichment_aborts_when_disk_user_changes_mid_flight() { - // Slow /user so we have time to swap the disk entry mid-flight. - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let app = axum::Router::new().route( - "/user", - axum::routing::get(|| async { - tokio::time::sleep(std::time::Duration::from_millis(300)).await; - axum::Json(serde_json::json!({ - "userId": "fetched-user", - "email": "fetched@corp.com", - "teamId": "fetched-team", - })) - }), - ); - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let scope = cfg.auth_scope(); - let mgr = Arc::new( - AuthManager::new(dir.path(), cfg.clone()) - .with_proxy_base_url(&format!("http://127.0.0.1:{port}")), - ); - - // Initial entry's user_id matches what /user will return, so - // enrichment WOULD apply normally. - let initial = GrokAuth { - key: "initial-key".into(), - refresh_token: Some("initial-rt".into()), - user_id: "fetched-user".into(), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - mgr.update(initial).await.unwrap(); - - // Race: while /user is in-flight, a "different user" overwrites - // disk. The enrichment must NOT overlay onto this new entry. tokio::time::sleep(std::time::Duration::from_millis(50)).await; - let intruder = GrokAuth { - key: "intruder-key".into(), - refresh_token: Some("intruder-rt".into()), - user_id: "intruder-user".into(), - team_id: Some("intruder-team".into()), - email: Some("intruder@corp.com".into()), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - let mut store = AuthStore::new(); - store.insert(scope.clone(), intruder); - write_auth_json(&dir.path().join("auth.json"), &store).unwrap(); - - // The /user mock takes 300 ms; after that the spawned enrichment - // either writes (overlay path -- the regression we're guarding - // against) or aborts silently. Poll the disk over a 3 s window - // and fail fast at the first poll that shows an overlay -- a - // wall-clock `sleep(800ms)` would mask both slow-CI flakes and - // a real regression that just happens to land >800ms in. - let auth_path = dir.path().join("auth.json"); - for _ in 0..30 { - tokio::time::sleep(std::time::Duration::from_millis(100)).await; - let store = read_auth_json(&auth_path).unwrap(); - let entry = store.get(&scope).expect("entry exists"); - assert_eq!( - entry.user_id, "intruder-user", - "intruder's user_id must survive aborted enrichment" - ); - assert_eq!( - entry.refresh_token.as_deref(), - Some("intruder-rt"), - "intruder's refresh_token must survive aborted enrichment" - ); - assert_eq!( - entry.key, "intruder-key", - "intruder's access token must survive aborted enrichment" - ); - assert_eq!( - entry.team_id.as_deref(), - Some("intruder-team"), - "intruder's team must NOT be overwritten with fetched-team" - ); - assert_eq!( - entry.email.as_deref(), - Some("intruder@corp.com"), - "intruder's email must NOT be overwritten with fetched@corp.com" - ); - } - - server.abort(); + m.auth().await.unwrap(); + assert!( + waiter.await.unwrap(), + "waiter must observe the token change" + ); } -/// Regression: on initial Team-principal login, the OIDC flow -/// stamps `auth.user_id = team_id` as a placeholder so telemetry -/// can distinguish teams immediately (see `extract_user_info` in -/// `oidc.rs`). The `/user` enrichment then returns the *real* -/// user_id and must overlay it onto disk -- this is the entire -/// point of the enrichment call for Team logins. Earlier revisions -/// of this PR compared `disk.user_id` against `user_info.user_id` -/// and treated this legitimate placeholder->real swap as a -/// concurrent user-switch, throwing away the email / team_name / -/// org fields. The guard now compares against the user_id we -/// *wrote* (`auth.user_id`), which matches disk on the bootstrap -/// path and only diverges when a sibling actually stomped. -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn enrichment_overlays_team_login_placeholder_user_id() { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let app = axum::Router::new().route( - "/user", - axum::routing::get(|| async { - axum::Json(serde_json::json!({ - "userId": "real-user-id", - "email": "user@corp.com", - "firstName": "Real", - "lastName": "User", - "principalType": "Team", - "principalId": "team-xyz", - "teamId": "team-xyz", - "teamName": "Some Team", - "teamRole": "MEMBER", - "organizationId": "org-abc", - "organizationName": "Some Org", - "organizationRole": "ORGANIZATION_ROLE_MEMBER", - })) - }), - ); - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); +// ── Tombstone semantics (PRD F1) ──────────────────────────────────────── - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let mgr = Arc::new( - AuthManager::new(dir.path(), cfg).with_proxy_base_url(&format!("http://127.0.0.1:{port}")), - ); +/// A 401-rejected refresh sets a tombstone keyed by the rejected refresh +/// token; subsequent auth() calls short-circuit without hitting the wire. +#[tokio::test] +async fn rejected_refresh_sets_tombstone_and_short_circuits() { + let (_d, m) = mgr(); + m.hot_swap(session("at-old", "rt-dead", 3600, -10)); + let calls = Arc::new(AtomicU32::new(0)); + m.set_refresher(Arc::new(RejectRefresher { + calls: calls.clone(), + rejected_rt: "rt-dead".into(), + })); - // Mirrors what `extract_user_info` returns for a Team principal: - // user_id stamped with the team_id placeholder; email + profile - // + team_name + org_* all empty until /user lands. - let team_login = GrokAuth { - key: "team-key".into(), - refresh_token: Some("team-rt".into()), - user_id: "team-xyz".into(), - email: None, - first_name: None, - last_name: None, - principal_type: Some("Team".into()), - principal_id: Some("team-xyz".into()), - team_id: Some("team-xyz".into()), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) + let err = m.auth().await.unwrap_err(); + assert!( + matches!( + err, + AuthError::Refresh(crate::auth::error::RefreshTokenError::Permanent(_)) + ), + "{err:?}" + ); + assert_eq!(calls.load(AtomicOrdering::SeqCst), 1); + assert!(m.has_permanent_failure(), "tombstone must be live"); + + // Second attempt: short-circuit, no second wire call. + let err2 = m.auth().await.unwrap_err(); + assert!( + matches!( + err2, + AuthError::Refresh(crate::auth::error::RefreshTokenError::Permanent(_)) + ), + "{err2:?}" + ); + assert_eq!( + calls.load(AtomicOrdering::SeqCst), + 1, + "tombstone must prevent a second refresh attempt" + ); +} + +/// The tombstone auto-clears when the persisted refresh token differs +/// (another process rotated the credential). +#[tokio::test] +async fn tombstone_clears_when_persisted_refresh_token_rotates() { + let (dir, m) = mgr(); + m.hot_swap(session("at-old", "rt-dead", 3600, -10)); + m.record_permanent_failure( + "rt-dead".into(), + RefreshTokenFailedReason::RefreshTokenRejected.into(), + ); + assert!(m.has_permanent_failure()); + + // Sibling process rotates the persisted credential. + let sibling = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + sibling + .update(session("at-rotated", "rt-rotated", 3600, 3600)) + .await + .unwrap(); + + assert!( + !m.has_permanent_failure(), + "tombstone must auto-clear once the persisted refresh token differs" + ); + // And auth() adopts the sibling's credential. + let auth = m.auth().await.unwrap(); + assert_eq!(auth.key, "at-rotated"); +} + +/// The tombstone cooldown (300s) ages out on the monotonic clock. +#[tokio::test] +async fn tombstone_cooldown_ages_out() { + let (_d, m) = mgr(); + m.hot_swap(session("at-old", "rt-dead", 3600, -10)); + m.record_permanent_failure( + "rt-dead".into(), + RefreshTokenFailedReason::RefreshTokenRejected.into(), + ); + assert!(m.has_permanent_failure()); + m.force_permanent_failure_aged_out(); + assert!( + !m.has_permanent_failure(), + "cooldown elapsed → retry allowed" + ); +} + +/// The cooldown also elapses on the wall clock alone (system slept through +/// the cooldown; monotonic clock paused). +#[tokio::test] +async fn tombstone_cooldown_ages_out_across_suspend() { + let (_d, m) = mgr(); + m.hot_swap(session("at-old", "rt-dead", 3600, -10)); + m.record_permanent_failure( + "rt-dead".into(), + RefreshTokenFailedReason::RefreshTokenRejected.into(), + ); + m.force_permanent_failure_wall_aged_out(); + assert!( + !m.has_permanent_failure(), + "wall-clock aging alone must clear the cooldown" + ); +} + +/// While a valid-on-the-wire access token exists, a tombstone must not block +/// auth() (the tombstone is about the refresh token, not the bearer). +#[tokio::test] +async fn tombstone_does_not_block_wire_valid_bearer() { + let (_d, m) = mgr(); + // Inside threshold (refresh due) but not hard-expired. + m.hot_swap(session("at-usable", "rt-dead", 7200, 3000)); + m.record_permanent_failure( + "rt-dead".into(), + RefreshTokenFailedReason::RefreshTokenRejected.into(), + ); + let auth = m.auth().await.unwrap(); + assert_eq!(auth.key, "at-usable"); +} + +// ── Persistence: file fallback + keyring ──────────────────────────────── + +#[tokio::test] +async fn update_persists_to_file_when_keyring_disabled() { + let (dir, m) = mgr(); + m.update(session("at-1", "rt-1", 3600, 3600)).await.unwrap(); + let store = read_auth_json(&dir.path().join("auth.json")).unwrap(); + let entry = store.get(KIMI_CODE_OAUTH_SCOPE).expect("scope entry"); + assert_eq!(entry.key, "at-1"); + assert_eq!(entry.refresh_token.as_deref(), Some("rt-1")); +} + +#[tokio::test] +async fn remove_scope_deletes_file_entry_and_memory() { + let (dir, m) = mgr(); + m.update(session("at-1", "rt-1", 3600, 3600)).await.unwrap(); + m.clear().unwrap(); + assert!(m.current_or_expired().is_none()); + assert!( + !dir.path().join("auth.json").exists(), + "last scope removed → file deleted" + ); + let (auth, state) = m.read_disk_auth_with_state(); + assert!(auth.is_none()); + assert_eq!(state, DiskAuthState::FileMissing); +} + +#[cfg(any(target_os = "macos", windows))] +mod keyring_integration { + use super::*; + use crate::auth::storage::{ + disable_mock_keyring_for_test, enable_mock_keyring_for_test, keyring_read_session, }; - mgr.update(team_login).await.unwrap(); - // Wait for the spawned enrichment to land. - let auth_path = dir.path().join("auth.json"); - let mut enriched = None; - for _ in 0..50 { - tokio::time::sleep(std::time::Duration::from_millis(100)).await; - let store = read_auth_json(&auth_path).unwrap(); - let entry = store.values().next().expect("entry exists").clone(); - if entry.email.is_some() { - enriched = Some(entry); - break; + struct MockKeyringGuard; + impl MockKeyringGuard { + fn enable() -> Self { + enable_mock_keyring_for_test(); + Self + } + } + impl Drop for MockKeyringGuard { + fn drop(&mut self) { + disable_mock_keyring_for_test(); } } - let enriched = enriched.expect("enrichment must overlay onto Team login"); - // The whole point: real user_id replaces the team_id placeholder. - assert_eq!( - enriched.user_id, "real-user-id", - "team_id placeholder must be replaced by real user_id from /user" - ); - assert_eq!(enriched.email.as_deref(), Some("user@corp.com")); - assert_eq!(enriched.first_name.as_deref(), Some("Real")); - assert_eq!(enriched.last_name.as_deref(), Some("User")); - assert_eq!(enriched.team_name.as_deref(), Some("Some Team")); - assert_eq!(enriched.team_role.as_deref(), Some("MEMBER")); - assert_eq!(enriched.organization_id.as_deref(), Some("org-abc")); - assert_eq!(enriched.organization_name.as_deref(), Some("Some Org")); - assert_eq!( - enriched.organization_role.as_deref(), - Some("ORGANIZATION_ROLE_MEMBER") - ); - // Tokens and team-id-as-principal-id preserved. - assert_eq!(enriched.key, "team-key"); - assert_eq!(enriched.refresh_token.as_deref(), Some("team-rt")); - assert_eq!(enriched.principal_type.as_deref(), Some("Team")); - assert_eq!(enriched.team_id.as_deref(), Some("team-xyz")); + /// With the keyring available, update() writes the session there (not + /// the file), reads come back from the keyring, and logout removes it. + #[tokio::test] + #[serial_test::serial(kigi_keyring)] + async fn update_prefers_keyring_and_logout_clears_it() { + let _guard = MockKeyringGuard::enable(); + let (dir, m) = mgr(); + m.update(session("at-kr", "rt-kr", 3600, 3600)) + .await + .unwrap(); + assert!( + !dir.path().join("auth.json").exists(), + "session must NOT land in the plaintext file when the keyring is available" + ); + assert!(matches!( + keyring_read_session(), + crate::auth::storage::KeyringRead::Found(_) + )); + let (auth, state) = m.read_disk_auth_with_state(); + assert_eq!(auth.map(|a| a.key), Some("at-kr".into())); + assert_eq!(state, DiskAuthState::Ok); - server.abort(); + // A fresh manager (same process) loads from the keyring. + let sibling = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + assert_eq!(sibling.current().map(|a| a.key), Some("at-kr".into())); + + // Logout removes the keyring entry. + m.clear().unwrap(); + assert!(matches!( + keyring_read_session(), + crate::auth::storage::KeyringRead::Missing + )); + } + + /// A stale file copy left from fallback days is stripped on the next + /// keyring write, and the keyring copy wins on reads. + #[tokio::test] + #[serial_test::serial(kigi_keyring)] + async fn keyring_write_strips_stale_file_copy() { + let (dir, m) = mgr(); + // Keyring disabled: first write lands in the file. + m.update(session("at-file", "rt-file", 3600, 3600)) + .await + .unwrap(); + assert!(dir.path().join("auth.json").exists()); + + // Keyring becomes available: the next write moves the credential. + let _guard = MockKeyringGuard::enable(); + m.update(session("at-kr2", "rt-kr2", 3600, 3600)) + .await + .unwrap(); + assert!( + !dir.path().join("auth.json").exists(), + "stale plaintext copy must be stripped after the keyring write" + ); + assert_eq!( + m.read_disk_auth().map(|a| a.key), + Some("at-kr2".into()), + "keyring copy is authoritative" + ); + } +} + +// ── Sibling adoption + disk reload ────────────────────────────────────── + +#[tokio::test] +async fn pick_up_sibling_token_adopts_different_valid_token() { + let (dir, m) = mgr(); + m.hot_swap(session("at-mine", "rt-mine", 3600, -10)); + + let sibling = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + sibling + .update(session("at-sibling", "rt-sibling", 3600, 3600)) + .await + .unwrap(); + + m.pick_up_sibling_token(); + assert_eq!(m.current().map(|a| a.key), Some("at-sibling".into())); } -/// Type-system invariant: `apply_user_info_enrichment` must NEVER -/// touch `key`, `refresh_token`, `expires_at`, `oidc_issuer`, -/// `oidc_client_id`, `auth_mode`, `create_time`, or -/// `has_grok_code_access`. The `&mut GrokAuth` signature already -/// enforces this at the type level (you cannot construct a fresh -/// auth from a `UserInfo` -- there's no `From` impl), but a unit -/// test pins the exact list of preserved fields so a future -/// contributor adding a token-like field to both `GrokAuth` and -/// `UserInfo` is forced to look here. #[test] -fn apply_user_info_enrichment_preserves_token_fields() { - let mut disk = GrokAuth { - key: "ROT_KEY".into(), - refresh_token: Some("ROT_RT".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - oidc_issuer: Some("https://issuer.example".into()), - oidc_client_id: Some("client-xyz".into()), - auth_mode: AuthMode::Oidc, - create_time: Utc::now() - Duration::minutes(10), - has_grok_code_access: Some(true), - user_id: "old-user".into(), - email: Some("old@corp.com".into()), - team_id: Some("old-team".into()), - ..GrokAuth::test_default() - }; - let snapshot = disk.clone(); - - let user_info = UserInfo { - user_id: "new-user".into(), - email: Some("new@corp.com".into()), - first_name: Some("New".into()), - last_name: Some("User".into()), - profile_image_asset_id: None, - principal_type: None, - principal_id: None, - team_id: Some("new-team".into()), - team_name: Some("New Team".into()), - team_role: None, - organization_id: None, - organization_name: None, - organization_role: None, - user_blocked_reason: None, - team_blocked_reasons: None, - coding_data_retention_opt_out: None, - subscription_tier: None, - }; - - apply_user_info_enrichment(&mut disk, user_info); - - // Token fields and provenance untouched. - assert_eq!(disk.key, snapshot.key); - assert_eq!(disk.refresh_token, snapshot.refresh_token); - assert_eq!(disk.expires_at, snapshot.expires_at); - assert_eq!(disk.oidc_issuer, snapshot.oidc_issuer); - assert_eq!(disk.oidc_client_id, snapshot.oidc_client_id); - assert_eq!(disk.auth_mode, snapshot.auth_mode); - assert_eq!(disk.create_time, snapshot.create_time); - assert_eq!(disk.has_grok_code_access, snapshot.has_grok_code_access); - - // Enrichment fields updated. - assert_eq!(disk.user_id, "new-user"); - assert_eq!(disk.email.as_deref(), Some("new@corp.com")); - assert_eq!(disk.team_id.as_deref(), Some("new-team")); - assert_eq!(disk.team_name.as_deref(), Some("New Team")); - assert_eq!(disk.first_name.as_deref(), Some("New")); -} - -/// Regression: async provider calls must drive `auth()` so tool requests get refreshed tokens. -#[tokio::test] -async fn current_api_key_async_drives_refresh_chain() { - use kigi_tools::types::ApiKeyProvider; - - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.hot_swap(GrokAuth { - key: "expired-oidc".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }); - let call_count = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: call_count.clone(), - delay: StdDuration::from_millis(0), - })); - - let provider = super::SharedAuthKeyProvider(mgr.clone()); - assert_eq!(provider.current_api_key().as_deref(), Some("expired-oidc")); - let key = provider.current_api_key_async().await; - assert_eq!(key.as_deref(), Some("fresh-token")); - assert_eq!(call_count.load(Ordering::SeqCst), 1); -} - -/// Regression: empty or corrupt auth.json must be recoverable on login. -/// Previously the guard in `update()` would skip the disk write on any -/// non-NotFound error, leaving a working in-memory session but a broken file. -#[tokio::test] -async fn update_recovers_from_empty_auth_json() { - let dir = tempfile::tempdir().unwrap(); - let auth_path = dir.path().join("auth.json"); - let cfg = GrokComConfig::default(); - std::fs::write(&auth_path, b"").unwrap(); - assert_eq!(std::fs::metadata(&auth_path).unwrap().len(), 0); - - let mgr = Arc::new(AuthManager::new(dir.path(), cfg.clone())); - - let new_auth = GrokAuth { - key: "recovered-token".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("recovered-rt".into()), - user_id: "recovered-user".into(), - email: Some("user@example.com".into()), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - - let result = mgr.update(new_auth.clone()).await; +fn force_reload_drops_credentials_on_readable_entry_missing() { + let (dir, m) = mgr(); + m.hot_swap(session("at-mine", "rt-mine", 3600, 3600)); + // A readable auth.json without our scope = trustworthy logout signal. + let store = AuthStore::new(); + crate::auth::storage::write_auth_json(&dir.path().join("auth.json"), &store).unwrap(); + // Non-empty map required for EntryMissing (empty map is still readable). + m.force_reload_from_disk(); assert!( - result.is_ok(), - "update must succeed and write to disk: {result:?}" - ); - - let current = mgr.current(); - assert_eq!( - current.as_ref().map(|a| a.key.as_str()), - Some("recovered-token") - ); - - let on_disk_raw = std::fs::read_to_string(&auth_path).unwrap(); - assert!( - !on_disk_raw.is_empty(), - "auth.json must not be empty after recovery" - ); - let on_disk: AuthStore = - serde_json::from_str(&on_disk_raw).expect("auth.json must be valid JSON after recovery"); - assert!( - on_disk.contains_key(&cfg.auth_scope()), - "persisted scope must be present" - ); - assert_eq!( - on_disk.get(&cfg.auth_scope()).map(|a| a.key.as_str()), - Some("recovered-token") + m.current_or_expired().is_none(), + "scope absent on readable store must drop in-memory credentials" ); } -/// Same as above, but for whitespace-only content. -#[tokio::test] -async fn update_recovers_from_whitespace_only_auth_json() { - let dir = tempfile::tempdir().unwrap(); - let auth_path = dir.path().join("auth.json"); - let cfg = GrokComConfig::default(); - std::fs::write(&auth_path, b" \n\t ").unwrap(); - - let mgr = Arc::new(AuthManager::new(dir.path(), cfg.clone())); - - let new_auth = GrokAuth { - key: "ws-token".into(), - auth_mode: AuthMode::Oidc, - user_id: "ws-user".into(), - ..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()) - }; - - let result = mgr.update(new_auth).await; - assert!( - result.is_ok(), - "update must succeed for whitespace-only file: {result:?}" - ); - - let on_disk = std::fs::read_to_string(&auth_path).unwrap(); - assert!(on_disk.contains("ws-token"), "credential must be persisted"); -} - -// -- sibling_has_different_refresh_token ---------------------------------- - -/// Expired disk AT with different RT is not a live sibling. -#[tokio::test] -async fn sibling_different_rt_with_expired_at_is_not_treated_as_live() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg.clone())); - - // In-memory: the original RT (revoked via rotation), AT expired. - let original = GrokAuth { - key: "original-at".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-original".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(original); - - // Disk: the successor RT from rotation, AT also expired. - let successor = GrokAuth { - key: "successor-at".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-successor".into()), - expires_at: Some(Utc::now() - Duration::minutes(30)), - ..GrokAuth::test_default() - }; - let mut store = AuthStore::new(); - store.insert(cfg.auth_scope(), successor); - write_auth_json(&dir.path().join("auth.json"), &store).unwrap(); - - assert!( - !mgr.sibling_has_different_refresh_token(), - "expired disk token must not be treated as a live sibling" +#[test] +fn force_reload_retains_refresh_token_on_disk_anomaly() { + let (dir, m) = mgr(); + m.hot_swap(session("at-mine", "rt-mine", 3600, 3600)); + // No auth.json at all (FileMissing anomaly): the in-memory refresh token + // may be the only copy — retain it. + assert!(!dir.path().join("auth.json").exists()); + m.force_reload_from_disk(); + assert_eq!( + m.current_or_expired().map(|a| a.key), + Some("at-mine".into()), + "disk anomaly must not discard a live refresh token" ); } -/// Valid disk AT with different RT is a live sibling. +// ── Proactive tick (loop body) ────────────────────────────────────────── + #[tokio::test] -async fn sibling_different_rt_with_valid_at_is_treated_as_live() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg.clone())); - - let original = GrokAuth { - key: "original-at".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-original".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(original); - - // Disk: valid token from sibling process. - let sibling = GrokAuth { - key: "sibling-at".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-sibling".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() - }; - let mut store = AuthStore::new(); - store.insert(cfg.auth_scope(), sibling); - write_auth_json(&dir.path().join("auth.json"), &store).unwrap(); - - assert!( - mgr.sibling_has_different_refresh_token(), - "valid disk token with different RT must be treated as live sibling" - ); -} - -/// Regression: refresh_chain(ServerRejected) must bypass the "double-check" -/// early return when the in-memory token is still valid (not expired). -/// Without this, a JWT that is time-valid but missing a subscription claim -/// (post-purchase) is returned as-is and the IdP is never contacted. -#[tokio::test] -async fn refresh_chain_server_rejected_bypasses_valid_token_double_check() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - // Seed a valid (non-expired) token — simulates a JWT that is missing - // the subscription claim but is otherwise fine. - let valid_but_rejected = GrokAuth { - key: "pre-subscription-jwt".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-original".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(valid_but_rejected); - - let call_count = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: call_count.clone(), - delay: StdDuration::from_millis(0), - })); - - // Confirm the token is considered valid before refresh. - assert_eq!(mgr.current().unwrap().key, "pre-subscription-jwt"); - - // ServerRejected must force a real refresh despite the token being valid. - let result = mgr - .refresh_chain( - crate::auth::token_type::TokenType::OidcSession, - RefreshReason::ServerRejected, - ) - .await; - +async fn proactive_tick_skips_above_threshold() { + let (_d, m) = mgr(); + m.hot_swap(session("at-fresh", "rt", 7200, 7000)); + let calls = install_ok_refresher(&m); + m.proactive_tick(false).await; assert_eq!( - result.unwrap().key, - "fresh-token", - "refresh_chain(ServerRejected) must contact the IdP even with a valid token" - ); - assert_eq!( - call_count.load(Ordering::SeqCst), - 1, - "refresher must be called exactly once" - ); - assert_eq!( - mgr.current().unwrap().key, - "fresh-token", - "in-memory token must be updated to the refreshed one" - ); -} - -/// When two tasks both get 401 and call refresh_chain(ServerRejected) -/// concurrently, the second caller must return the already-refreshed token -/// without contacting the IdP again. This prevents the double-refresh race -/// where the second caller sends a rotated refresh token → invalid_grant. -#[tokio::test] -async fn refresh_chain_server_rejected_concurrent_skips_redundant_refresh() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - // Seed the "rejected" token that both tasks will see. - let rejected = GrokAuth { - key: "rejected-jwt".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-old".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(rejected); - - let call_count = Arc::new(AtomicU32::new(0)); - // Slow refresher so the second task blocks on the lock long enough - // to observe the first task's refresh result. - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: call_count.clone(), - delay: StdDuration::from_millis(50), - })); - - // Both tasks snapshot pre_lock_key = "rejected-jwt", then race for - // the lock. The first refreshes → "fresh-token". The second finds - // current() = "fresh-token" != pre_lock_key → returns early. - let mgr1 = mgr.clone(); - let mgr2 = mgr.clone(); - - let (r1, r2) = tokio::join!( - mgr1.refresh_chain( - crate::auth::token_type::TokenType::OidcSession, - RefreshReason::ServerRejected, - ), - mgr2.refresh_chain( - crate::auth::token_type::TokenType::OidcSession, - RefreshReason::ServerRejected, - ), - ); - - // Both must succeed with the refreshed token. - assert_eq!(r1.unwrap().key, "fresh-token"); - assert_eq!(r2.unwrap().key, "fresh-token"); - - // The IdP must be contacted exactly once, not twice. - assert_eq!( - call_count.load(Ordering::SeqCst), - 1, - "refresher must be called exactly once; second caller should \ - return the already-refreshed token via the double-check guard" - ); -} - -/// Counterpart: refresh_chain(PreRequest) with a valid token must -/// short-circuit and NOT call the refresher. -#[tokio::test] -async fn refresh_chain_pre_request_short_circuits_on_valid_token() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - let valid = GrokAuth { - key: "still-good".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(valid); - - let call_count = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: call_count.clone(), - delay: StdDuration::from_millis(0), - })); - - let result = mgr - .refresh_chain( - crate::auth::token_type::TokenType::OidcSession, - RefreshReason::PreRequest, - ) - .await; - - assert_eq!(result.unwrap().key, "still-good"); - assert_eq!( - call_count.load(Ordering::SeqCst), + calls.load(AtomicOrdering::SeqCst), 0, - "PreRequest must NOT call refresher when token is valid" - ); -} - -// -- login-time inline enrichment ------------------------------------------- - -/// Axum `/user` stub serving `body`; rejects requests missing `Bearer {token}`. -async fn spawn_user_stub(token: &'static str, body: &'static str) -> String { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let app = axum::Router::new().route( - "/user", - axum::routing::get(move |headers: axum::http::HeaderMap| async move { - let authz = headers - .get("authorization") - .and_then(|v| v.to_str().ok()) - .unwrap_or_default(); - if authz != format!("Bearer {token}") { - return Err(axum::http::StatusCode::UNAUTHORIZED); - } - Ok(([("content-type", "application/json")], body)) - }), - ); - tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - format!("http://127.0.0.1:{port}") -} - -#[tokio::test] -async fn enrich_auth_inline_populates_zdr_flags() { - let body = r#"{"userId":"u-1","teamBlockedReasons":["BLOCKED_REASON_NO_LOGS"],"codingDataRetentionOptOut":true}"#; - let base = spawn_user_stub("tok", body).await; - let dir = tempfile::tempdir().unwrap(); - let mgr = AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base); - - let mut auth = GrokAuth { - key: "tok".into(), - ..GrokAuth::test_default() - }; - assert!(!auth.is_data_collection_disabled(), "precondition"); - - mgr.enrich_auth_inline(&mut auth).await; - assert!(auth.is_zdr_team(), "team_blocked_reasons must be merged"); - assert!(auth.coding_data_retention_opt_out); - assert_eq!(auth.user_id, "u-1"); -} - -#[tokio::test] -async fn enrich_auth_inline_keeps_fields_absent_from_response() { - // `/user` omitting a field must not clear a value the login flow set. - let body = r#"{"userId":"u-1","teamBlockedReasons":["BLOCKED_REASON_NO_LOGS_MODERATED"]}"#; - let base = spawn_user_stub("tok", body).await; - let dir = tempfile::tempdir().unwrap(); - let mgr = AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base); - - let mut auth = GrokAuth { - key: "tok".into(), - principal_type: Some("Team".into()), - principal_id: Some("team-1".into()), - ..GrokAuth::test_default() - }; - - mgr.enrich_auth_inline(&mut auth).await; - assert_eq!(auth.user_id, "u-1"); - assert_eq!(auth.principal_type.as_deref(), Some("Team")); - assert_eq!(auth.principal_id.as_deref(), Some("team-1")); - assert!(auth.is_zdr_team()); - assert!( - !auth.coding_data_retention_opt_out, - "absent field stays unchanged" + "above the refresh threshold no refresh must run" ); } #[tokio::test] -async fn enrich_auth_inline_unreachable_server_leaves_auth_unchanged() { - let dir = tempfile::tempdir().unwrap(); - // Bind-then-drop to get a port that refuses connections. - let port = { - let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); - l.local_addr().unwrap().port() - }; - let mgr = AuthManager::new(dir.path(), GrokComConfig::default()) - .with_proxy_base_url(&format!("http://127.0.0.1:{port}")); - - let mut auth = GrokAuth { - key: "tok".into(), - ..GrokAuth::test_default() - }; - let before = auth.clone(); - mgr.enrich_auth_inline(&mut auth).await; - assert_eq!(auth.user_id, before.user_id); - assert!(!auth.is_data_collection_disabled()); +async fn proactive_tick_refreshes_inside_threshold() { + let (_d, m) = mgr(); + // 7200s lifetime, 3000s left → inside max(300, 3600) threshold. + m.hot_swap(session("at-aging", "rt", 7200, 3000)); + let calls = install_ok_refresher(&m); + m.proactive_tick(false).await; + assert_eq!(calls.load(AtomicOrdering::SeqCst), 1); + assert_eq!(m.current().map(|a| a.key), Some("at-refreshed".into())); } -// ── force_login_team_uuid spine enforcement ─────────────────────────── -// -// Regression coverage for the cached-token bypass: the pin must hold for every -// token the manager hands out (startup, sync reads, `auth()`), not just fresh -// login. Each test fails on the pre-fix tree. - -/// `jsonwebtoken` needs a process-level CryptoProvider; tests that encode -/// JWTs can't rely on another test having installed it first. -fn ensure_crypto_provider() { - let _ = jsonwebtoken::crypto::rust_crypto::DEFAULT_PROVIDER.install_default(); -} - -/// A signed (HS256) access token carrying a `Team` principal, matching the -/// shape `peek_access_token_principal` extracts in production. -fn team_jwt(principal_id: &str) -> String { - ensure_crypto_provider(); - jsonwebtoken::encode( - &jsonwebtoken::Header::new(jsonwebtoken::Algorithm::HS256), - &serde_json::json!({ - "sub": "user-1", - "principal_type": "Team", - "principal_id": principal_id, - "exp": 9999999999u64, - }), - &jsonwebtoken::EncodingKey::from_secret(b"test-secret"), - ) - .unwrap() -} - -/// An access token carrying `principal_id` but NO `principal_type`. -fn principal_id_only_jwt(principal_id: &str) -> String { - ensure_crypto_provider(); - jsonwebtoken::encode( - &jsonwebtoken::Header::new(jsonwebtoken::Algorithm::HS256), - &serde_json::json!({ - "sub": "user-1", - "principal_id": principal_id, - "exp": 9999999999u64, - }), - &jsonwebtoken::EncodingKey::from_secret(b"test-secret"), - ) - .unwrap() -} - -fn pinned_cfg(team: &str) -> GrokComConfig { - GrokComConfig { - force_login_team_uuid: Some(crate::auth::config::ForceLoginTeam::Single( - team.to_string(), - )), - ..GrokComConfig::default() - } -} - -/// A valid, non-expired OIDC session whose access token carries `principal_id`. -fn oidc_session_for_team(principal_id: &str) -> GrokAuth { - GrokAuth { - key: team_jwt(principal_id), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - oidc_issuer: Some(crate::auth::config::XAI_OAUTH2_ISSUER.to_string()), - oidc_client_id: Some("client".into()), - ..GrokAuth::test_default() - } -} - -/// The repro: a wrong-team session persisted to disk (e.g. logged in before -/// the pin was deployed) must be cleared at construction, not silently loaded. -#[test] -fn new_clears_wrong_team_token_loaded_from_disk() { - let dir = tempfile::tempdir().unwrap(); - let cfg = pinned_cfg("team-good"); - let scope = cfg.auth_scope(); - - let mut store = AuthStore::new(); - store.insert(scope, oidc_session_for_team("team-wrong")); - write_auth_json(&dir.path().join("auth.json"), &store).unwrap(); - - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - assert!(mgr.current().is_none(), "wrong-team token must be hidden"); - assert!( - mgr.current_or_expired().is_none(), - "wrong-team token must be cleared from memory, not just hidden" - ); - assert!( - !dir.path().join("auth.json").exists(), - "wrong-team auth.json must be cleared so the next launch re-logs in" - ); -} - -/// A matching-team session on disk is loaded normally (no false positive). -#[test] -fn new_keeps_matching_team_token_loaded_from_disk() { - let dir = tempfile::tempdir().unwrap(); - let cfg = pinned_cfg("team-good"); - let scope = cfg.auth_scope(); - let tok = oidc_session_for_team("team-good"); - - let mut store = AuthStore::new(); - store.insert(scope, tok.clone()); - write_auth_json(&dir.path().join("auth.json"), &store).unwrap(); - - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - assert_eq!(mgr.current().map(|a| a.key), Some(tok.key)); - assert!(dir.path().join("auth.json").exists()); -} - -/// `auth()` (the wire-bound chokepoint used by pager / MCP / -/// `try_ensure_fresh_auth`) rejects and clears a wrong-team cached token. +/// Sleep/wake force: a forced tick refreshes even a token comfortably above +/// the threshold (kimi-cli `refreshing()` parity). #[tokio::test] -async fn auth_rejects_and_clears_wrong_team_cached_token() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), pinned_cfg("team-good"))); - // hot_swap bypasses the pin (like a sibling adoption mid-session). - mgr.hot_swap(oidc_session_for_team("team-wrong")); - - assert!(mgr.current().is_none(), "sync read must hide the token"); - - let err = mgr.auth().await.unwrap_err(); - assert!( - matches!(err, AuthError::PinnedTeamMismatch { .. }), - "auth() must surface the policy violation, got {err:?}" - ); - assert!( - mgr.current_or_expired().is_none(), - "auth() must clear the violating session" +async fn proactive_tick_force_refreshes_valid_token() { + let (_d, m) = mgr(); + m.hot_swap(session("at-fresh", "rt", 7200, 7000)); + let calls = install_ok_refresher(&m); + m.proactive_tick(true).await; + assert_eq!( + calls.load(AtomicOrdering::SeqCst), + 1, + "force must bypass the threshold check" ); + assert_eq!(m.current().map(|a| a.key), Some("at-refreshed".into())); } -/// A matching-team cached token flows through `auth()` unchanged. #[tokio::test] -async fn auth_accepts_matching_team_cached_token() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), pinned_cfg("team-good"))); - let tok = oidc_session_for_team("team-good"); - mgr.hot_swap(tok.clone()); - - assert_eq!(mgr.current().map(|a| a.key.clone()), Some(tok.key.clone())); - assert_eq!(mgr.auth().await.unwrap().key, tok.key); -} - -/// No pin configured: any team is accepted (the enforcement is opt-in and -/// must not affect default deployments). -#[tokio::test] -async fn no_pin_accepts_any_team_cached_token() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - let tok = oidc_session_for_team("team-anything"); - mgr.hot_swap(tok.clone()); - - assert_eq!(mgr.current().map(|a| a.key.clone()), Some(tok.key.clone())); - assert_eq!(mgr.auth().await.unwrap().key, tok.key); -} - -/// A token that silently refreshes into a wrong-team principal is rejected by -/// `auth()` (the wrapper gates refresh results, not just the cached fast path). -#[tokio::test] -async fn auth_rejects_token_refreshed_into_wrong_team() { - struct WrongTeamRefresher { - jwt: String, - } - #[async_trait::async_trait] - impl TokenRefresher for WrongTeamRefresher { - async fn refresh(&self, _reason: RefreshReason) -> crate::auth::refresh::RefreshOutcome { - crate::auth::refresh::RefreshOutcome::Success(Box::new(GrokAuth { - key: self.jwt.clone(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-new".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() - })) - } - } - - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), pinned_cfg("team-good"))); - // Expired matching session forces a refresh; the refresher returns a - // wrong-team token (e.g. a re-pinned token family). - mgr.hot_swap(GrokAuth { - expires_at: Some(Utc::now() - Duration::minutes(10)), - ..oidc_session_for_team("team-good") - }); - mgr.set_refresher(Arc::new(WrongTeamRefresher { - jwt: team_jwt("team-wrong"), - })); - - let err = mgr.auth().await.unwrap_err(); - assert!( - matches!(err, AuthError::PinnedTeamMismatch { .. }), - "refreshed wrong-team token must be rejected, got {err:?}" - ); -} - -/// A sibling-written wrong-team token picked up by `force_reload_from_disk` -/// (relay reconnect) is cleared, not just hidden. -#[test] -fn force_reload_clears_wrong_team_token() { - let dir = tempfile::tempdir().unwrap(); - let cfg = pinned_cfg("team-good"); - let scope = cfg.auth_scope(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); // empty disk at startup - - let mut store = AuthStore::new(); - store.insert(scope, oidc_session_for_team("team-wrong")); - write_auth_json(&dir.path().join("auth.json"), &store).unwrap(); - - mgr.force_reload_from_disk(); - assert!( - mgr.current_or_expired().is_none(), - "reloaded wrong-team token must be cleared, not just hidden" - ); - assert!( - !dir.path().join("auth.json").exists(), - "force_reload must clear auth.json on a pin violation" - ); -} - -// -- force_reload_from_disk: transient disk anomaly vs real logout ---------- - -/// A real incident in miniature: a live in-memory OIDC session (RT -/// present, no permanent_failure) while `auth.json` transiently reads as -/// missing — e.g. the first read right after wake-from-sleep resolves the path -/// to `ENOENT`. The refresh token may exist nowhere else, so the reload must -/// RETAIN it, not discard it (the discard previously kicked off a -/// 401 -> reactive refresh -> suspend-straddle -> invalid_grant cascade). -#[test] -fn force_reload_retains_live_rt_on_transient_file_missing() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - let session = GrokAuth { - key: "live-session".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("live-rt".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(session); - assert!(mgr.permanent_failure().is_none()); - - // No auth.json on disk at all -> FileMissing on every read. - assert!(mgr.read_disk_auth().is_none()); - - // Zero backoff so the retry budget is exhausted instantly. - mgr.force_reload_from_disk_with(RELOAD_RETRY_TRIES, StdDuration::ZERO); - - let retained = mgr.current_or_expired(); - assert!( - retained.is_some(), - "a live RT must NOT be discarded on a transient FileMissing", - ); - let retained = retained.unwrap(); - assert_eq!(retained.key, "live-session"); - assert_eq!(retained.refresh_token.as_deref(), Some("live-rt")); -} - -/// Contrast with the retain case: once a `permanent_failure` is cached the RT -/// is known-dead, so a persistent FileMissing must drop it (and clear the -/// permanent_failure with it) so the next request reports `NotLoggedIn`. -#[tokio::test] -async fn force_reload_drops_rt_when_permanent_failure_set() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - let session = GrokAuth { - key: "broken".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-revoked".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(session); - record_permanent_failure( - &mgr, - crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected, - ); - assert!(mgr.permanent_failure().is_some()); - - mgr.force_reload_from_disk_with(RELOAD_RETRY_TRIES, StdDuration::ZERO); - - assert!( - mgr.current_or_expired().is_none(), - "a known-dead RT (permanent_failure set) must be dropped", - ); - assert!( - mgr.permanent_failure().is_none(), - "dropping creds must clear the cached permanent_failure", - ); - assert!(matches!( - mgr.auth().await.unwrap_err(), - AuthError::NotLoggedIn - )); -} - -/// A readable `auth.json` that simply lacks our scope is the trustworthy -/// "logged out / scope removed" signal (distinct from a missing file), so the -/// in-memory credentials are dropped even though an RT is present. -#[test] -fn force_reload_drops_creds_on_entry_missing() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - let session = GrokAuth { - key: "live-session".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("live-rt".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(session); - - // auth.json exists and is readable, but holds only an unrelated scope -> - // EntryMissing for this manager's scope. - let mut store = AuthStore::new(); - store.insert( - "https://example.invalid::nobody".to_string(), - make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()), - ); - write_auth_json(&dir.path().join("auth.json"), &store).unwrap(); - - mgr.force_reload_from_disk_with(RELOAD_RETRY_TRIES, StdDuration::ZERO); - - assert!( - mgr.current_or_expired().is_none(), - "scope absent on a readable auth.json is a real logout -> drop", - ); -} - -/// When disk holds a fresh token for our scope, the reload adopts it on the -/// first read (no retry) — the healthy path is unchanged. -#[test] -fn force_reload_adopts_fresh_disk_token() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let scope = cfg.auth_scope(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); - - let expired = GrokAuth { - key: "stale".into(), - refresh_token: Some("old-rt".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - }; - mgr.hot_swap(expired); - - let fresh = GrokAuth { - key: "fresh-from-disk".into(), - refresh_token: Some("new-rt".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() - }; - let mut store = AuthStore::new(); - store.insert(scope, fresh); - write_auth_json(&dir.path().join("auth.json"), &store).unwrap(); - - mgr.force_reload_from_disk_with(RELOAD_RETRY_TRIES, StdDuration::ZERO); - - assert_eq!(mgr.current().unwrap().key, "fresh-from-disk"); -} - -/// A token carrying `principal_id` without `principal_type` is matched on the -/// id alone: the pinned team is accepted, not falsely rejected. -#[tokio::test] -async fn pin_matches_principal_id_without_principal_type() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), pinned_cfg("team-good"))); - mgr.hot_swap(GrokAuth { - key: principal_id_only_jwt("team-good"), - auth_mode: AuthMode::Oidc, - expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() - }); - - assert!( - mgr.current().is_some(), - "matching team id must be accepted even without principal_type" - ); - assert!(mgr.auth().await.is_ok()); -} - -/// A cached `AuthMode::ApiKey` session is rejected under the kill switch (here -/// implied by a team pin), and honored when it's off. -#[tokio::test] -async fn cached_api_key_session_rejected_when_api_key_auth_disabled() { - let api_key_session = || GrokAuth { - key: "xai-cached-key".into(), +async fn proactive_tick_skips_non_refreshable_types() { + let (_d, m) = mgr(); + m.hot_swap(KimiAuth { + key: "sk-key".into(), auth_mode: AuthMode::ApiKey, - expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() - }; - - // Switch ON (via a team pin, which implies api_key_auth_disabled): reject. - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), pinned_cfg("team-good"))); - mgr.hot_swap(api_key_session()); - assert!( - mgr.current().is_none(), - "cached api-key session must be hidden under the kill switch" - ); - assert!( - matches!(mgr.auth().await, Err(AuthError::ApiKeyAuthDisabled)), - "auth() must reject a cached api-key session under the kill switch" - ); - - // Switch OFF (no pin / no disable): the api-key session is honored. - let dir2 = tempfile::tempdir().unwrap(); - let mgr2 = Arc::new(AuthManager::new(dir2.path(), GrokComConfig::default())); - mgr2.hot_swap(api_key_session()); - assert_eq!( - mgr2.current().map(|a| a.key), - Some("xai-cached-key".to_string()), - "api-key session must work normally when the switch is off" - ); + ..KimiAuth::test_default() + }); + let calls = install_ok_refresher(&m); + m.proactive_tick(true).await; + assert_eq!(calls.load(AtomicOrdering::SeqCst), 0); } #[tokio::test] -async fn shared_api_key_provider_resolves_live_bearer() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - let auth = GrokAuth { - key: "shared-provider-token".into(), - expires_at: Some(Utc::now() + Duration::hours(1)), - create_time: Utc::now(), - ..GrokAuth::test_default() - }; - mgr.hot_swap(auth); - - let provider = shared_api_key_provider(mgr.clone()); - - // Synchronous accessor surfaces the current (non-expired) bearer. - assert_eq!( - provider.current_api_key(), - Some("shared-provider-token".to_string()), - "shared_api_key_provider must expose the live bearer to out-of-crate consumers" +async fn proactive_tick_respects_tombstone_cooldown() { + let (_d, m) = mgr(); + m.hot_swap(session("at-old", "rt-dead", 3600, -10)); + m.record_permanent_failure( + "rt-dead".into(), + RefreshTokenFailedReason::RefreshTokenRejected.into(), ); - - // Async accessor resolves a valid bearer without a network refresh when - // the cached token is still fresh. + let calls = install_ok_refresher(&m); + m.proactive_tick(false).await; assert_eq!( - provider.current_api_key_async().await, - Some("shared-provider-token".to_string()), - "async accessor must resolve the current bearer for a fresh token" - ); - - // A hot-swap is reflected on the next resolution (no startup snapshot). - let rotated = GrokAuth { - key: "rotated-token".into(), - expires_at: Some(Utc::now() + Duration::hours(1)), - create_time: Utc::now(), - ..GrokAuth::test_default() - }; - mgr.hot_swap(rotated); - assert_eq!( - provider.current_api_key(), - Some("rotated-token".to_string()), - "provider must follow the manager's refresh chain rather than snapshot at startup" - ); -} - -fn expired_oidc() -> GrokAuth { - GrokAuth { - key: "expired-key".into(), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-old".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() - } -} - -/// Signals when it has started, then blocks until released. -struct BlockingRefresher { - started: Arc, - release: Arc, - call_count: Arc, -} - -#[async_trait::async_trait] -impl TokenRefresher for BlockingRefresher { - async fn refresh(&self, _reason: RefreshReason) -> crate::auth::refresh::RefreshOutcome { - self.call_count.fetch_add(1, Ordering::SeqCst); - self.started.notify_one(); - self.release.notified().await; - crate::auth::refresh::RefreshOutcome::Success(Box::new(GrokAuth { - key: "fresh-token".into(), - expires_at: Some(Utc::now() + Duration::hours(1)), - refresh_token: Some("rt-new".into()), - ..GrokAuth::test_default() - })) - } -} - -#[tokio::test] -async fn sleep_gate_defers_refresh_without_calling_idp() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.hot_swap(expired_oidc()); - let call_count = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: call_count.clone(), - delay: StdDuration::from_millis(0), - })); - - mgr.set_system_sleep_imminent(true); - - let err = mgr.auth().await.unwrap_err(); - assert!( - matches!(err, AuthError::Refresh(RefreshTokenError::Transient(_))), - "gated refresh must return a transient refresh error, got {err:?}" - ); - assert_eq!( - call_count.load(Ordering::SeqCst), + calls.load(AtomicOrdering::SeqCst), 0, - "the IdP refresher must NOT be called while the sleep gate is raised" + "tombstone cooldown must gate the proactive tick" ); } -/// A sleep-deferred refresh must not poison auth state: the deferral is a -/// typed transient (retryable on wake), never forces a manual re-login (a -/// lid close must never count as a forced re-login), and records -/// no permanent-failure verdict — even after more deferred attempts than the -/// refresher-level escalation budget tolerates (the transient-blip budget -/// lives in the refresher, which a deferral never reaches). -/// -/// Coverage depth: the gate is raised before the chain starts, so this drives -/// the step-3a deferral. The step-3c pre-IdP re-check (gate raised inside the -/// 3a→3c race window) returns the identical transient error and touches the -/// same state, but is not deterministically reachable without production test -/// hooks, so it is pinned only indirectly by these assertions. -#[tokio::test] -async fn sleep_deferred_refresh_is_transient_no_reauth_no_verdict() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - // Pin non-devbox so a deferred refresh surfaces the transient error - // instead of minting via devbox recovery (CI runs in K8s pods). - mgr.set_devbox_env_for_test(false); - mgr.hot_swap(expired_oidc()); - let call_count = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: call_count.clone(), - delay: StdDuration::from_millis(0), - })); - - mgr.set_system_sleep_imminent(true); - - // More attempts than MAX_CONSECUTIVE_TRANSIENT_FAILURES: deferrals must - // never accrue toward an escalated permanent verdict. - for _ in 0..4 { - let err = mgr.auth().await.unwrap_err(); - assert!( - matches!(err, AuthError::Refresh(RefreshTokenError::Transient(_))), - "a sleep-deferred refresh must be transient, got {err:?}" - ); - assert!( - !crate::auth::recovery::forces_manual_reauth(&err), - "a lid-close deferral must never force a manual re-login", - ); - } - assert!( - mgr.permanent_failure().is_none(), - "deferrals must not record a permanent-failure verdict", - ); - assert_eq!( - call_count.load(Ordering::SeqCst), - 0, - "the refresher must never run while the gate is raised", - ); - - // End-to-end through 401 recovery: a deferred recovery terminates with - // the transient error. - let mut rec = mgr.unauthorized_recovery(mgr.current_or_expired()); - let err = rec.next().await.unwrap_err(); - assert!( - matches!(err, AuthError::Refresh(RefreshTokenError::Transient(_))), - "deferred recovery must surface the transient deferral, got {err:?}" - ); -} +// ── Sleep gate integration ────────────────────────────────────────────── #[tokio::test] -async fn dark_wake_defers_refresh_without_calling_idp() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.hot_swap(expired_oidc()); - let call_count = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: call_count.clone(), - delay: StdDuration::from_millis(0), - })); - - mgr.set_dark_wake_for_test(true); - - let err = mgr.auth().await.unwrap_err(); +async fn refresh_chain_defers_when_sleep_imminent() { + let (_d, m) = mgr(); + m.hot_swap(session("at-old", "rt-old", 3600, -10)); + let calls = install_ok_refresher(&m); + m.set_system_sleep_imminent(true); + let err = m + .refresh_chain(m.token_type(), RefreshReason::PreRequest) + .await + .unwrap_err(); assert!( matches!( err, AuthError::Refresh(crate::auth::error::RefreshTokenError::Transient(_)) ), - "dark-wake refresh must return a transient refresh error, got {err:?}" + "sleep-gated refresh must defer transiently: {err:?}" ); + assert_eq!(calls.load(AtomicOrdering::SeqCst), 0); + m.set_system_sleep_imminent(false); + m.auth().await.unwrap(); assert_eq!( - call_count.load(Ordering::SeqCst), - 0, - "the IdP refresher must NOT be called during a dark wake (the refresh \ - token must not be sent into a possible re-sleep)" - ); - - // Returning to a full wake lets the refresh proceed and reach the IdP. - mgr.set_dark_wake_for_test(false); - assert_eq!(mgr.auth().await.unwrap().key, "fresh-token"); - assert_eq!( - call_count.load(Ordering::SeqCst), + calls.load(AtomicOrdering::SeqCst), 1, - "after a full wake the refresher must be invoked" + "wake resumes refresh" ); } -/// A machine stuck reporting a *continuous* dark wake (e.g. an interactive Mac -/// with no display) must not defer refresh forever — once the deferral budget -/// (`DARK_WAKE_DEFER_MAX`) is exhausted, one refresh is forced through. Without -/// this bound the user reaches the same logged-out state the dark-wake guard -/// was added to prevent. -#[tokio::test] -async fn dark_wake_defer_forces_refresh_after_max() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.hot_swap(expired_oidc()); - let call_count = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: call_count.clone(), - delay: StdDuration::from_millis(0), - })); - - mgr.set_dark_wake_for_test(true); - - // Backdate the start of the deferral run past the bound on both clocks, as - // if we had been continuously in dark wake longer than DARK_WAKE_DEFER_MAX. - let back = super::sleep_gate::DARK_WAKE_DEFER_MAX + StdDuration::from_secs(5); - let (Some(mono), Some(wall)) = ( - Instant::now().checked_sub(back), - std::time::SystemTime::now().checked_sub(back), - ) else { - return; // machine/clock can't represent the backdate — skip - }; - *mgr.dark_wake_defer_since.write() = Some(super::sleep_gate::GateRaise { mono, wall }); - - assert_eq!( - mgr.auth().await.unwrap().key, - "fresh-token", - "an exhausted dark-wake deferral budget must force the refresh through" - ); - assert_eq!( - call_count.load(Ordering::SeqCst), - 1, - "the IdP refresher must be invoked once the dark-wake defer budget is exhausted" - ); - assert!( - mgr.dark_wake_defer_since.read().is_none(), - "forcing a refresh through must reset the defer budget" - ); -} - -/// A `DidWake` (`SYSTEM_HAS_POWERED_ON`) event must not reset the dark-wake -/// defer budget while the system is *still* in a dark wake — macOS can deliver -/// powered-on events for dark wakes, and resetting then would stop the budget -/// from ever exhausting, so the forced refresh would never run. Only a genuine -/// full wake clears it. -#[test] -fn dark_wake_defer_budget_survives_powered_on_during_dark_wake() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - // Begin a deferral run. - mgr.set_dark_wake_for_test(true); - assert!( - mgr.should_defer_for_dark_wake(), - "a fresh dark wake should defer and start the budget" - ); - assert!(mgr.dark_wake_defer_since.read().is_some()); - - // A powered-on event arrives while still in a dark wake: the budget must - // persist so it can eventually exhaust and force a refresh through. - mgr.set_system_sleep_imminent(false); - assert!( - mgr.dark_wake_defer_since.read().is_some(), - "a powered-on event during a dark wake must not reset the defer budget" - ); - - // A genuine full wake clears the run. - mgr.set_dark_wake_for_test(false); - mgr.set_system_sleep_imminent(false); - assert!( - mgr.dark_wake_defer_since.read().is_none(), - "a full wake must clear the defer budget" - ); -} - -/// The `power_listener_started` guard in `is_dark_wake` must short-circuit to -/// `false` when no OS power listener was started (headless / datacenter), so -/// those processes never treat the OS power state as a dark wake. Exercises the -/// guard directly (no dark-wake override installed). -#[test] -fn is_dark_wake_false_when_power_listener_not_started() { - let dir = tempfile::tempdir().unwrap(); - let mgr = AuthManager::new(dir.path(), GrokComConfig::default()); - assert!( - !mgr.is_dark_wake(), - "is_dark_wake must be false when the power listener was never started" - ); -} +// ── Idempotency guards ────────────────────────────────────────────────── #[tokio::test] -async fn sleep_gate_cleared_on_wake_allows_refresh() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.hot_swap(expired_oidc()); - let call_count = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(CountingRefresher { - call_count: call_count.clone(), - delay: StdDuration::from_millis(0), - })); - - mgr.set_system_sleep_imminent(true); - mgr.set_system_sleep_imminent(false); // wake - - let auth = mgr.auth().await.expect("refresh should succeed after wake"); - assert_eq!(auth.key, "fresh-token"); - assert_eq!(call_count.load(Ordering::SeqCst), 1); +async fn start_proactive_refresh_is_idempotent_per_arc() { + let (_d, m) = mgr(); + let cancel = CancellationToken::new(); + m.start_proactive_refresh(cancel.clone()); + m.start_proactive_refresh(cancel.clone()); + assert_eq!(m.proactive_start_count(), 1, "second start must be a no-op"); + cancel.cancel(); } -#[tokio::test] -async fn sleep_gate_auto_expires_after_max() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - mgr.set_system_sleep_imminent(true); - assert!(mgr.is_sleep_gated(), "freshly-raised gate must be active"); - - // Simulate a missed wake while awake the whole time: both clocks were - // raised longer ago than the bound. - let back = super::sleep_gate::SLEEP_GATE_MAX + StdDuration::from_secs(5); - let (Some(mono), Some(wall)) = ( - Instant::now().checked_sub(back), - std::time::SystemTime::now().checked_sub(back), - ) else { - return; // machine/clock can't represent the backdate — not reproducible; skip - }; - *mgr.sleep_gate.raised_at.write() = Some(super::sleep_gate::GateRaise { mono, wall }); - - assert!( - !mgr.is_sleep_gated(), - "a gate older than SLEEP_GATE_MAX must auto-expire" - ); - assert!( - mgr.sleep_gate.raised_at.read().is_none(), - "auto-expiry must also lower the gate so a stale state can't linger" - ); -} - -/// Regression test for the dual-clock backstop: a gate that straddled a real -/// system sleep must auto-expire even though the monotonic clock is still -/// fresh, because the wall clock advanced past the bound during sleep. Before -/// the wall-clock arm this gate stayed shut and an expired token reached the -/// server — the 401 this fix targets. -#[tokio::test] -async fn sleep_gate_auto_expires_when_wall_clock_passes_during_sleep() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - mgr.set_system_sleep_imminent(true); - assert!(mgr.is_sleep_gated(), "freshly-raised gate must be active"); - - // Monotonic clock fresh (as if the machine just slept rather than spending - // the time awake); wall clock pushed past the bound (real time elapsed - // while asleep, where the monotonic clock is frozen). - let back = super::sleep_gate::SLEEP_GATE_MAX + StdDuration::from_secs(5); - let Some(wall) = std::time::SystemTime::now().checked_sub(back) else { - return; // clock can't represent the backdate — not reproducible; skip - }; - *mgr.sleep_gate.raised_at.write() = Some(super::sleep_gate::GateRaise { - mono: Instant::now(), - wall, - }); - - assert!( - !mgr.is_sleep_gated(), - "a gate whose wall-clock age exceeds SLEEP_GATE_MAX must auto-expire \ - even though the monotonic clock is still fresh" - ); - assert!( - mgr.sleep_gate.raised_at.read().is_none(), - "auto-expiry must also lower the gate so a stale state can't linger" - ); -} - -#[tokio::test] -async fn sleep_gate_lets_in_flight_refresh_complete() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.hot_swap(expired_oidc()); - - let started = Arc::new(tokio::sync::Notify::new()); - let release = Arc::new(tokio::sync::Notify::new()); - let call_count = Arc::new(AtomicU32::new(0)); - mgr.set_refresher(Arc::new(BlockingRefresher { - started: started.clone(), - release: release.clone(), - call_count: call_count.clone(), - })); - - let m = mgr.clone(); - let handle = tokio::spawn(async move { m.auth().await }); - - started.notified().await; - assert_eq!( - mgr.refresh_in_flight.load(Ordering::SeqCst), - 1, - "refresh must be counted as in flight while the IdP call is pending" - ); - // `set_system_sleep_imminent` now holds the OS sleep ack until the - // in-flight refresh drains. Drive it from a separate thread — as the real - // OS power-listener thread does — so the tokio runtime stays free to - // complete the refresh while the hold waits. - let sleeper = mgr.clone(); - let ack = std::thread::spawn(move || { - let start = Instant::now(); - sleeper.set_system_sleep_imminent(true); - start.elapsed() - }); - - release.notify_one(); - - let auth = tokio::time::timeout(StdDuration::from_secs(5), handle) - .await - .expect("auth() must return") - .unwrap() - .expect("in-flight refresh must complete, not abort"); - let ack_waited = ack.join().expect("ack thread panicked"); - - assert_eq!(auth.key, "fresh-token"); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - assert!(mgr.is_sleep_gated(), "WillSleep must raise the sleep gate"); - assert!( - ack_waited < super::sleep_gate::SLEEP_ACK_MAX_WAIT, - "the sleep-ack hold must release when the refresh drains, not wait out \ - SLEEP_ACK_MAX_WAIT; waited {ack_waited:?}" - ); - assert_eq!( - mgr.refresh_in_flight.load(Ordering::SeqCst), - 0, - "in-flight counter must be balanced after completion" - ); -} - -/// With nothing in flight, the sleep-ack hold must return promptly so the OS -/// suspend is never delayed unnecessarily. #[test] -fn sleep_ack_hold_returns_immediately_when_nothing_in_flight() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - - let start = Instant::now(); - mgr.test_hold_sleep_ack(StdDuration::from_secs(5)); - let waited = start.elapsed(); - - assert!( - waited < StdDuration::from_millis(250), - "no in-flight refresh must not delay the suspend; waited {waited:?}" - ); -} - -/// The sleep-ack hold must unblock as soon as the in-flight refresh drains, -/// well before the bound — this is the straddle the fix prevents. -#[test] -fn sleep_ack_hold_releases_when_in_flight_refresh_drains() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.test_enter_refresh_in_flight(); - - let releaser = mgr.clone(); - let drain = std::thread::spawn(move || { - std::thread::sleep(StdDuration::from_millis(120)); - releaser.test_exit_refresh_in_flight(); - }); - - let start = Instant::now(); - mgr.test_hold_sleep_ack(StdDuration::from_secs(5)); - let waited = start.elapsed(); - drain.join().unwrap(); - - assert!( - waited >= StdDuration::from_millis(100), - "must hold the ack until the refresh drains; waited only {waited:?}" - ); - assert!( - waited < StdDuration::from_secs(2), - "must release shortly after the drain, not near the bound; waited {waited:?}" - ); - assert_eq!(mgr.refresh_in_flight.load(Ordering::SeqCst), 0); -} - -/// A refresh that never drains must not pin the machine awake: the hold is -/// bounded and returns at the deadline, leaving the refresh running (never -/// aborted) for the existing straddle telemetry to catch. -#[test] -fn sleep_ack_hold_times_out_when_refresh_never_drains() { - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - mgr.test_enter_refresh_in_flight(); // never exits - - let start = Instant::now(); - mgr.test_hold_sleep_ack(StdDuration::from_millis(150)); - let waited = start.elapsed(); - - assert!( - waited >= StdDuration::from_millis(140), - "must wait out the bound; waited only {waited:?}" - ); - assert!( - waited < StdDuration::from_secs(1), - "must not exceed the bound by much; waited {waited:?}" - ); - assert_eq!( - mgr.refresh_in_flight.load(Ordering::SeqCst), - 1, - "the refresh is left running, not aborted, when the hold times out" - ); -} - -// ── manual re-auth classification ──────────────────────────── - -/// Truth table for `forces_manual_reauth`: only terminal failures that force -/// the user back to `/login` count; self-healing (TTL) reasons, transient / -/// no-credential errors, and API-key lockouts don't. -#[test] -fn forces_manual_reauth_maps_terminal_and_skips_non_forcing() { - use crate::auth::error::RefreshTokenFailedReason as Reason; - use crate::auth::recovery::forces_manual_reauth; - - let permanent = |reason: Reason| forces_manual_reauth(&AuthError::permanent(reason)); - // A revoked refresh token forces a re-login -> counts. - assert!(permanent(Reason::RefreshTokenRejected)); - // Every terminal pipeline error forces a re-login. - assert!(forces_manual_reauth(&AuthError::ServerRejectedNoRecovery)); - assert!(forces_manual_reauth(&AuthError::RecoveryExhausted)); - assert!(forces_manual_reauth(&AuthError::TokenExpiredNoRefresh)); - assert!(forces_manual_reauth(&AuthError::PinnedTeamMismatch { - message: String::new() - })); - // Self-healing (TTL) reasons, transient / no-credential, and API-key - // lockouts don't count. - assert!(!permanent(Reason::ClientRejected)); - assert!(!permanent(Reason::Other)); - assert!(!forces_manual_reauth(&AuthError::transient("x"))); - assert!(!forces_manual_reauth(&AuthError::NotLoggedIn)); - assert!(!forces_manual_reauth(&AuthError::ApiKeyAuthDisabled)); -} - -/// Truth table for `relay_should_cancel`: the relay gives up on any terminal -/// auth failure — including `ApiKeyAuthDisabled`, which deliberately doesn't -/// force a manual re-login — and keeps reconnecting through transient -/// blips, absent credentials, and the self-healing permanent reasons (those -/// age out via the TTL, so cancelling on them would orphan a session that -/// recovers minutes later). -#[test] -fn relay_should_cancel_gives_up_only_on_terminal_failures() { - use crate::auth::error::RefreshTokenFailedReason as Reason; - use crate::auth::recovery::relay_should_cancel; - - // Terminal: the handshake can't recover; stop reconnecting. - assert!(relay_should_cancel(&AuthError::permanent( - Reason::RefreshTokenRejected - ))); - assert!(relay_should_cancel(&AuthError::ServerRejectedNoRecovery)); - assert!(relay_should_cancel(&AuthError::RecoveryExhausted)); - assert!(relay_should_cancel(&AuthError::TokenExpiredNoRefresh)); - assert!(relay_should_cancel(&AuthError::PinnedTeamMismatch { - message: String::new() - })); - // Cancelled even though it never forces a re-login (a kill-switched API - // key means rotate the key, not `/login`). - assert!(relay_should_cancel(&AuthError::ApiKeyAuthDisabled)); - - // Recoverable: fall through and reconnect. - assert!(!relay_should_cancel(&AuthError::transient("network blip"))); - assert!(!relay_should_cancel(&AuthError::permanent( - Reason::ClientRejected - ))); - assert!(!relay_should_cancel(&AuthError::permanent(Reason::Other))); - assert!(!relay_should_cancel(&AuthError::NotLoggedIn)); +fn configure_refresher_is_idempotent() { + let (_d, m) = mgr(); + assert!(m.configure_refresher(), "first call installs"); + assert!(!m.configure_refresher(), "second call is a no-op"); + assert!(m.has_refresher_attached()); } diff --git a/crates/codegen/kigi-shell/src/auth/meta.rs b/crates/codegen/kigi-shell/src/auth/meta.rs index c2ae23a..30bfaa4 100644 --- a/crates/codegen/kigi-shell/src/auth/meta.rs +++ b/crates/codegen/kigi-shell/src/auth/meta.rs @@ -1,6 +1,8 @@ use serde::{Deserialize, Serialize}; -/// Access gate from `grok_build_access_gate`. +/// Access-gate copy resolved from remote settings (message + optional CTA). +/// Auth no longer produces gates (tier gating was an xAI concept); the pager +/// still renders one when remote settings carry a gate message. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct GateInfo { pub message: String, @@ -17,24 +19,6 @@ pub struct AuthMeta { pub email: Option, #[serde(default)] pub auth_mode: Option, - /// Team principal UUID when the session is a team login (`None` for personal). - #[serde(default)] - pub team_id: Option, - #[serde(default)] - pub team_name: Option, - #[serde(default)] - pub is_zdr: bool, - #[serde(default)] - pub team_role: Option, - #[serde(default)] - pub coding_data_retention_opt_out: bool, #[serde(default)] pub show_resolved_model: Option, - /// `Some` = user is blocked; `None` = user has access. - #[serde(default)] - pub gate: Option, - /// User-friendly display name for the current subscription tier - /// (e.g. "SuperGrok Heavy", "X Premium", "Free"). From CCP `/settings`. - #[serde(default)] - pub subscription_tier: Option, } diff --git a/crates/codegen/kigi-shell/src/auth/mod.rs b/crates/codegen/kigi-shell/src/auth/mod.rs index 624c894..6c60d36 100644 --- a/crates/codegen/kigi-shell/src/auth/mod.rs +++ b/crates/codegen/kigi-shell/src/auth/mod.rs @@ -1,42 +1,30 @@ pub(crate) mod attribution; mod config; pub mod credential_provider; -#[path = "devbox_login_stub.rs"] -pub(crate) mod devbox_login; +pub(crate) mod device; pub mod device_code; pub mod error; -mod external_auth; mod flow; -mod jwt; +pub(crate) mod kimi_oauth; pub(crate) mod manager; mod model; -pub mod oidc; pub(crate) mod recovery; pub(crate) mod refresh; mod storage; pub(crate) mod token_type; -pub(crate) use config::LEGACY_AUTH_SCOPE; -pub use config::{ - ForceLoginTeam, GrokComConfig, OAuth2ProviderConfig, OidcAuthConfig, PreferredAuthMethod, - XAI_OAUTH2_ISSUER, is_xai_oauth2_issuer, xai_oauth2_issuer, -}; -pub(crate) use external_auth::{parse_output, refresh_with_command}; -pub(crate) use flow::{ - AuthChannels, run_auth_flow, run_auth_flow_with_stderr_bridge, - try_ensure_session_noninteractive, -}; +pub use config::{KIMI_CODE_OAUTH_SCOPE, KimiCodeConfig}; +pub(crate) use flow::try_ensure_session_noninteractive; pub use flow::{ - AuthUrlInfo, AuthUrlMode, LoginTransportOverride, LogoutResult, ensure_authenticated, - ensure_authenticated_or_noninteractive, ensure_authenticated_with_override, perform_logout, - run_cli_login, run_cli_logout, try_ensure_fresh_auth, + AuthChannels, AuthUrlInfo, AuthUrlMode, LogoutResult, ensure_authenticated, + ensure_authenticated_or_noninteractive, perform_logout, run_auth_flow, + run_auth_flow_with_stderr_bridge, run_cli_login, run_cli_logout, try_ensure_fresh_auth, }; -pub use jwt::{is_jwt_expired_or_near, parse_jwt_expiration}; mod meta; pub use error::{AuthError, RefreshTokenError, RefreshTokenFailedReason}; pub use manager::{AuthManager, shared_api_key_provider}; pub use meta::{AuthMeta, GateInfo}; -pub use model::{AuthMode, GrokAuth, lookup_auth}; -pub(crate) use model::{TOKEN_TTL, UserInfo, is_expired, token_suffix}; +pub use model::{AuthMode, KimiAuth, lookup_auth}; +pub(crate) use model::{TOKEN_TTL, is_expired, token_suffix}; pub use storage::{ clear_api_key, read_api_key, read_auth_json, read_token_by_scope, store_api_key, }; diff --git a/crates/codegen/kigi-shell/src/auth/model.rs b/crates/codegen/kigi-shell/src/auth/model.rs index 90cd482..8bf6082 100644 --- a/crates/codegen/kigi-shell/src/auth/model.rs +++ b/crates/codegen/kigi-shell/src/auth/model.rs @@ -1,108 +1,75 @@ +//! Kimi Code auth data model: the persisted token set + expiry policy. + use chrono::{DateTime, Duration, Utc}; use serde::{Deserialize, Serialize}; use std::collections::BTreeMap; -use super::is_xai_oauth2_issuer; - +/// Fallback TTL for credentials without a server-provided expiry +/// (plain API keys). pub(crate) const TOKEN_TTL: Duration = Duration::days(30); -const DEFAULT_EARLY_INVALIDATION_SECS: u64 = 300; // 5 minutes -/// Legacy auth.json scope key. Fallback for old devbox auth files. -pub(super) const LEGACY_SCOPE: &str = "https://accounts.x.ai/sign-in"; +/// Minimum refresh threshold (PRD F1): refresh when the remaining lifetime +/// drops below `max(300, expires_in × 0.5)` seconds. +const DEFAULT_EARLY_INVALIDATION_SECS: u64 = 300; -/// auth.json scope key for plain API key auth (desktop login, `grok login --api-key`). -pub const API_KEY_SCOPE: &str = "xai::api_key"; +/// Fraction of `expires_in` that drives the dynamic refresh threshold. +const REFRESH_THRESHOLD_RATIO: f64 = 0.5; -const BLOCKED_REASON_NO_LOGS: &str = "BLOCKED_REASON_NO_LOGS"; -const BLOCKED_REASON_NO_LOGS_MODERATED: &str = "BLOCKED_REASON_NO_LOGS_MODERATED"; +/// auth.json scope key for plain API key auth (`kigi login --api-key`, F2). +pub const API_KEY_SCOPE: &str = "kigi::api_key"; -/// Token provenance (debugging/auth.json only -- no code branches on this). +/// How this credential was obtained. #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] #[serde(rename_all = "snake_case")] pub enum AuthMode { - /// Deprecated. Kept for deserializing old auth.json files. - #[serde(alias = "grok")] - WebLogin, - /// OIDC or OAuth2 interactive login via customer IdP - #[serde(alias = "oidc")] - Oidc, - /// External auth provider binary - External, - /// Plain API key (e.g. from grok-desktop login or `grok login --api-key`) + /// Kimi Code subscription OAuth (device-code flow). + #[serde(rename = "oauth")] + OAuth, + /// Plain API key. ApiKey, } -/// Wire value of `principal_type` for team OAuth principals (capitalized by -/// the auth service). Single source for every comparison site. -pub(crate) const TEAM_PRINCIPAL_TYPE: &str = "Team"; - +/// The Kimi Code credential: the OAuth token set (or a bare API key) plus +/// local bookkeeping. The Kimi token response carries no user info; `user_id` +/// / `email` stay empty until a later feature surfaces account info. #[derive(Clone, Serialize, Deserialize)] -pub struct GrokAuth { +pub struct KimiAuth { + /// The bearer sent on API calls (`Authorization: Bearer {key}`): + /// the OAuth access token, or the API key in `ApiKey` mode. pub key: String, pub auth_mode: AuthMode, pub create_time: DateTime, - pub user_id: String, - pub email: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub first_name: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub last_name: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub profile_image_asset_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub principal_type: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub principal_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub team_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub team_name: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub team_role: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub organization_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub organization_name: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub organization_role: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub user_blocked_reason: Option, - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub team_blocked_reasons: Vec, + /// Account id — the Kimi token response has none; empty until a later + /// feature surfaces it. #[serde(default)] - pub coding_data_retention_opt_out: bool, - - /// Deprecated. Kept for deserializing existing auth.json files. + pub user_id: String, + /// Account email — the Kimi token response has none; `None` until a + /// later feature surfaces it. #[serde(default, skip_serializing_if = "Option::is_none")] - pub has_grok_code_access: Option, - - /// Refresh token (OIDC/OAuth2 or external provider). + pub email: Option, + /// OAuth refresh token; `None` for API keys. #[serde(default, skip_serializing_if = "Option::is_none")] pub refresh_token: Option, - - /// Server-provided expiration (from OIDC `expires_in`). - /// When present, takes precedence over the hardcoded `TOKEN_TTL`. + /// `create_time + expires_in`, computed when the token was minted. #[serde(default, skip_serializing_if = "Option::is_none")] pub expires_at: Option>, - - /// Issuer URL that issued this token. For OIDC credentials it drives - /// refresh via discovery; for external-provider credentials it is the - /// provider's `issuer` claim. In both modes an x.ai issuer marks the - /// credential first-party (`is_xai_auth`). + /// Server-reported token lifetime in seconds; drives the dynamic + /// refresh threshold `max(300, expires_in × 0.5)`. #[serde(default, skip_serializing_if = "Option::is_none")] - pub oidc_issuer: Option, - - /// OIDC client_id used to obtain this token (needed for refresh). + pub expires_in: Option, + /// OAuth scope string as returned by the token endpoint. #[serde(default, skip_serializing_if = "Option::is_none")] - pub oidc_client_id: Option, + pub scope: Option, + /// Token type as returned by the token endpoint (e.g. "bearer"). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub token_type: Option, } -impl std::fmt::Debug for GrokAuth { +impl std::fmt::Debug for KimiAuth { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("GrokAuth") + f.debug_struct("KimiAuth") .field("key", &token_suffix(&self.key)) .field("auth_mode", &self.auth_mode) - .field("user_id", &self.user_id) .field("expires_at", &self.expires_at) .field( "refresh_token", @@ -112,128 +79,44 @@ impl std::fmt::Debug for GrokAuth { } } -impl GrokAuth { +impl KimiAuth { /// Seconds since this credential was minted. Negative when the local - /// clock stepped back past `create_time` (NTP correction, VM restore, or - /// a sibling machine's clock via an adopted auth.json) — `create_time` - /// is always stamped from the minting machine's local clock. + /// clock stepped back past `create_time` (NTP correction, VM restore). pub(crate) fn mint_age_seconds(&self) -> i64 { Utc::now() .signed_duration_since(self.create_time) .num_seconds() } - /// `true` when the token comes from a first-party xAI account — - /// either an OIDC login against https://auth.x.ai (or the local-dev - /// equivalent), or an external auth provider that declared an xAI - /// issuer for its token. - /// - /// The issuer is a client-side hint, not a trust assertion: everything - /// it unlocks still authenticates the actual token server-side, and it - /// never influences endpoints. - pub fn is_xai_auth(&self) -> bool { - match self.auth_mode { - AuthMode::Oidc | AuthMode::External => self - .oidc_issuer - .as_deref() - .is_some_and(is_xai_oauth2_issuer), - AuthMode::ApiKey | AuthMode::WebLogin => false, - } - } - - /// `true` when this auth can access grok.com managed MCP connectors. - pub fn is_managed_mcp_eligible(&self) -> bool { - self.is_xai_auth() || self.auth_mode == AuthMode::WebLogin - } - - /// Whether this credential can access `supported_in_api: false` models. - /// - /// Session logins (WebLogin, OIDC — including enterprise issuers) always - /// qualify; external-provider credentials qualify only when first-party - /// (`is_xai_auth`), matching the built-in devbox login they replace. - /// Plain API keys never do. + /// `true` for a refreshable subscription session (vs a bare API key). pub fn is_session_auth(&self) -> bool { - match self.auth_mode { - AuthMode::WebLogin | AuthMode::Oidc => true, - AuthMode::External => self.is_xai_auth(), - AuthMode::ApiKey => false, - } - } - - pub fn is_team_principal(&self) -> bool { - self.principal_type.as_deref() == Some(TEAM_PRINCIPAL_TYPE) && self.team_id.is_some() - } - - /// `true` when the team has Zero Data Retention (ZDR) enabled. - pub fn is_zdr_team(&self) -> bool { - self.team_blocked_reasons - .iter() - .any(|r| r == BLOCKED_REASON_NO_LOGS || r == BLOCKED_REASON_NO_LOGS_MODERATED) - } - - /// `true` when the team has ZDR or the user opted out of coding data - /// retention. Use this for trace-upload and research-data gates. - /// Product analytics (`telemetry_enabled`) and user-facing sync - /// features should use `is_zdr_team()` directly. - pub fn is_data_collection_disabled(&self) -> bool { - self.is_zdr_team() || self.coding_data_retention_opt_out - } - - /// Carry `/user`-derived fields from a previous auth so refresh rebuilds don't drop them. - pub(crate) fn carry_user_profile_from(&mut self, prev: &GrokAuth) { - self.user_id = prev.user_id.clone(); - self.email = prev.email.clone(); - self.principal_type = prev.principal_type.clone(); - self.principal_id = prev.principal_id.clone(); - self.team_id = prev.team_id.clone(); - self.team_name = prev.team_name.clone(); - self.team_role = prev.team_role.clone(); - self.organization_id = prev.organization_id.clone(); - self.organization_name = prev.organization_name.clone(); - self.organization_role = prev.organization_role.clone(); - self.user_blocked_reason = prev.user_blocked_reason.clone(); - self.team_blocked_reasons = prev.team_blocked_reasons.clone(); - self.coding_data_retention_opt_out = prev.coding_data_retention_opt_out; + self.auth_mode == AuthMode::OAuth } } -impl Default for GrokAuth { +impl Default for KimiAuth { fn default() -> Self { Self { key: String::new(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, create_time: Utc::now(), user_id: String::new(), email: None, - first_name: None, - last_name: None, - profile_image_asset_id: None, - principal_type: None, - principal_id: None, - team_id: None, - team_name: None, - team_role: None, - organization_id: None, - organization_name: None, - organization_role: None, - user_blocked_reason: None, - team_blocked_reasons: vec![], - coding_data_retention_opt_out: false, - has_grok_code_access: None, refresh_token: None, expires_at: None, - oidc_issuer: None, - oidc_client_id: None, + expires_in: None, + scope: None, + token_type: None, } } } #[cfg(test)] -impl GrokAuth { - /// Returns a `GrokAuth` with sensible defaults for tests. Override fields - /// with struct update syntax: +impl KimiAuth { + /// A `KimiAuth` with sensible defaults for tests. Override fields with + /// struct update syntax: /// ```ignore - /// GrokAuth { key: "my-key".into(), ..GrokAuth::test_default() } + /// KimiAuth { key: "my-key".into(), ..KimiAuth::test_default() } /// ``` pub fn test_default() -> Self { Self { @@ -244,82 +127,24 @@ impl GrokAuth { } } -pub(crate) type AuthStore = BTreeMap; +pub(crate) type AuthStore = BTreeMap; -/// User information from the cli-chat-proxy `GET /v1/user` endpoint. -#[derive(Debug, Clone, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct UserInfo { - pub(crate) user_id: String, - #[serde(default)] - pub(super) email: Option, - #[serde(default)] - pub(super) first_name: Option, - #[serde(default)] - pub(super) last_name: Option, - #[serde(default)] - pub(super) profile_image_asset_id: Option, - #[serde(default)] - pub(super) principal_type: Option, - #[serde(default)] - pub(super) principal_id: Option, - #[serde(default)] - pub(super) team_id: Option, - #[serde(default)] - pub(super) team_name: Option, - #[serde(default)] - pub(super) team_role: Option, - #[serde(default)] - pub(super) organization_id: Option, - #[serde(default)] - pub(super) organization_name: Option, - #[serde(default)] - pub(super) organization_role: Option, - #[serde(default)] - pub(super) user_blocked_reason: Option, - #[serde(default)] - pub(super) team_blocked_reasons: Option>, - #[serde(default)] - pub(super) coding_data_retention_opt_out: Option, - /// Live subscription tier from the backend (only present when - /// `?include=subscription` is passed to `/user`). - #[serde(default)] - pub(crate) subscription_tier: Option, -} - -/// Last 12 chars of a token string, safe for diagnostic logging. -/// Uses the tail because JWT access tokens all share the same base64 -/// header prefix (`eyJ0eXAiOiJh…`); the tail (signature bytes) is -/// unique per token and makes `key_changed` / `is_stale_snapshot` -/// diagnostics meaningful. +/// Last 12 chars of a token string, safe for diagnostic logging. Uses the +/// tail because token prefixes are shared across a family; the tail is +/// unique per token and makes `key_changed` diagnostics meaningful. pub(crate) fn token_suffix(t: &str) -> &str { let len = t.len(); if len > 12 { &t[len - 12..] } else { t } } /// Look up auth from the store by scope key. -/// -/// Legacy `WebLogin` tokens (from the pre-OIDC `grok login --legacy` -/// flow) are skipped — they are validated via a per-request DB lookup -/// server-side which fails at high volume. Skipping them here forces -/// affected users to re-authenticate via OIDC on next launch. -pub fn lookup_auth(map: &AuthStore, scope: &str) -> Option { - let auth = map.get(scope).cloned().or_else(|| { - if scope == LEGACY_SCOPE { - None - } else { - map.get(LEGACY_SCOPE).cloned() - } - })?; - if auth.auth_mode == AuthMode::WebLogin { - tracing::info!("auth: ignoring legacy WebLogin token — re-authentication required"); - return None; - } - Some(auth) +pub fn lookup_auth(map: &AuthStore, scope: &str) -> Option { + map.get(scope).cloned() } -/// Early-invalidation buffer. Override with `KIGI_AUTH_EARLY_INVALIDATION_SECS` -/// for testing (e.g. `=5` to shrink the buffer to 5 seconds). +/// Minimum refresh-threshold component. Override with +/// `KIGI_AUTH_EARLY_INVALIDATION_SECS` for testing (e.g. `=5` to shrink the +/// buffer to 5 seconds). pub(super) fn early_invalidation() -> Duration { std::env::var("KIGI_AUTH_EARLY_INVALIDATION_SECS") .ok() @@ -328,14 +153,30 @@ pub(super) fn early_invalidation() -> Duration { .unwrap_or_else(|| Duration::seconds(DEFAULT_EARLY_INVALIDATION_SECS as i64)) } -pub(crate) fn is_expired(auth: &GrokAuth) -> bool { - is_expired_with_buffer(auth, early_invalidation()) +/// Dynamic refresh threshold (PRD F1): `max(min_threshold, expires_in × 0.5)` +/// where `min_threshold` defaults to 300s. Credentials without a positive +/// `expires_in` use the minimum alone. +pub(crate) fn refresh_threshold(auth: &KimiAuth) -> Duration { + let min = early_invalidation(); + match auth.expires_in { + Some(expires_in) if expires_in > 0 => { + let ratio = Duration::seconds((expires_in as f64 * REFRESH_THRESHOLD_RATIO) as i64); + std::cmp::max(min, ratio) + } + _ => min, + } +} + +/// Whether the credential is inside its refresh threshold (i.e. should be +/// treated as expiring-soon for refresh scheduling). +pub(crate) fn is_expired(auth: &KimiAuth) -> bool { + is_expired_with_buffer(auth, refresh_threshold(auth)) } /// Like [`is_expired`] but with an explicit pre-expiry buffer. Pass /// `Duration::zero()` for actual (hard) expiry — the instant the token would -/// really be rejected on the wire, with no early-invalidation margin. -pub(crate) fn is_expired_with_buffer(auth: &GrokAuth, buffer: Duration) -> bool { +/// really be rejected on the wire. +pub(crate) fn is_expired_with_buffer(auth: &KimiAuth, buffer: Duration) -> bool { if let Some(expires_at) = auth.expires_at { Utc::now() >= (expires_at - buffer) } else { @@ -348,142 +189,108 @@ pub(crate) fn is_expired_with_buffer(auth: &GrokAuth, buffer: Duration) -> bool mod tests { use super::*; - fn make_auth(mode: AuthMode) -> GrokAuth { - GrokAuth { - key: "k".into(), - auth_mode: mode, - create_time: Utc::now(), - user_id: "u".into(), - email: None, - first_name: None, - last_name: None, - profile_image_asset_id: None, - principal_type: None, - principal_id: None, - team_id: None, - team_name: None, - team_role: None, - organization_id: None, - organization_name: None, - organization_role: None, - user_blocked_reason: None, - team_blocked_reasons: vec![], - coding_data_retention_opt_out: false, - has_grok_code_access: None, - refresh_token: None, - expires_at: None, - oidc_issuer: None, - oidc_client_id: None, + fn auth_with_lifetime(expires_in: i64, remaining_secs: i64) -> KimiAuth { + KimiAuth { + expires_in: Some(expires_in), + expires_at: Some(Utc::now() + Duration::seconds(remaining_secs)), + refresh_token: Some("rt".into()), + ..KimiAuth::test_default() } } + /// PRD threshold math: `max(300, expires_in × 0.5)`. #[test] - fn is_xai_auth_matrix() { - use crate::auth::XAI_OAUTH2_ISSUER; - let with_issuer = |mode: AuthMode, issuer: Option<&str>| GrokAuth { - oidc_issuer: issuer.map(str::to_owned), - ..make_auth(mode) + fn refresh_threshold_is_max_of_min_and_half_life() { + // Short-lived token: the 300s floor wins (600 × 0.5 = 300 → tie; 400 × 0.5 = 200 < 300). + let short = auth_with_lifetime(400, 400); + assert_eq!(refresh_threshold(&short).num_seconds(), 300); + // Long-lived token: half the lifetime wins (7200 × 0.5 = 3600). + let long = auth_with_lifetime(7200, 7200); + assert_eq!(refresh_threshold(&long).num_seconds(), 3600); + // No expires_in: the floor alone. + let bare = KimiAuth::test_default(); + assert_eq!(refresh_threshold(&bare).num_seconds(), 300); + // Non-positive expires_in must not produce a negative threshold. + let broken = KimiAuth { + expires_in: Some(-5), + ..KimiAuth::test_default() }; + assert_eq!(refresh_threshold(&broken).num_seconds(), 300); + } - // Only Oidc/External qualify, and only with an x.ai issuer. - assert!(with_issuer(AuthMode::Oidc, Some(XAI_OAUTH2_ISSUER)).is_xai_auth()); - assert!(with_issuer(AuthMode::External, Some(XAI_OAUTH2_ISSUER)).is_xai_auth()); - assert!(!with_issuer(AuthMode::Oidc, None).is_xai_auth()); - assert!(!with_issuer(AuthMode::External, None).is_xai_auth()); - assert!(!with_issuer(AuthMode::Oidc, Some("https://idp.acme.example")).is_xai_auth()); - assert!(!with_issuer(AuthMode::External, Some("https://idp.acme.example")).is_xai_auth()); + /// A token past its dynamic threshold counts as expiring-soon while a + /// token comfortably before it does not. + #[test] + fn is_expired_uses_dynamic_threshold() { + // 7200s lifetime → threshold 3600s. 3000s remaining < 3600 → expiring. + assert!(is_expired(&auth_with_lifetime(7200, 3000))); + // 5000s remaining > 3600 → fresh. + assert!(!is_expired(&auth_with_lifetime(7200, 5000))); + // Hard expiry ignores the buffer entirely. + assert!(!is_expired_with_buffer( + &auth_with_lifetime(7200, 3000), + Duration::zero() + )); + assert!(is_expired_with_buffer( + &auth_with_lifetime(7200, -1), + Duration::zero() + )); + } - // ApiKey / WebLogin stay false even with an x.ai issuer set. - assert!(!with_issuer(AuthMode::ApiKey, Some(XAI_OAUTH2_ISSUER)).is_xai_auth()); - assert!(!with_issuer(AuthMode::WebLogin, Some(XAI_OAUTH2_ISSUER)).is_xai_auth()); + /// Credentials without `expires_at` (API keys) age out via the 30-day TTL. + #[test] + fn no_expiry_falls_back_to_token_ttl() { + let fresh = KimiAuth::test_default(); + assert!(!is_expired(&fresh)); + let old = KimiAuth { + create_time: Utc::now() - Duration::days(31), + ..KimiAuth::test_default() + }; + assert!(is_expired(&old)); } #[test] - fn is_session_auth_requires_first_party_for_external() { - use crate::auth::XAI_OAUTH2_ISSUER; - let with_issuer = |mode: AuthMode, issuer: Option<&str>| GrokAuth { - oidc_issuer: issuer.map(str::to_owned), - ..make_auth(mode) + fn lookup_auth_finds_scope_entry() { + let mut map = AuthStore::new(); + map.insert("oauth/kimi-code".into(), KimiAuth::test_default()); + assert!(lookup_auth(&map, "oauth/kimi-code").is_some()); + assert!(lookup_auth(&map, "other").is_none()); + } + + #[test] + fn debug_redacts_tokens() { + let auth = KimiAuth { + key: "super-secret-access-token".into(), + refresh_token: Some("super-secret-refresh-token".into()), + ..KimiAuth::test_default() }; + let debug = format!("{auth:?}"); + assert!(!debug.contains("super-secret-access-token")); + assert!(!debug.contains("super-secret-refresh-token")); + } - // Session logins qualify regardless of issuer (incl. enterprise OIDC). - assert!(with_issuer(AuthMode::WebLogin, None).is_session_auth()); - assert!(with_issuer(AuthMode::Oidc, None).is_session_auth()); - assert!(with_issuer(AuthMode::Oidc, Some("https://idp.acme.example")).is_session_auth()); - - // External qualifies only when first-party (devbox-login parity). - assert!(with_issuer(AuthMode::External, Some(XAI_OAUTH2_ISSUER)).is_session_auth()); - assert!(!with_issuer(AuthMode::External, None).is_session_auth()); + #[test] + fn serde_roundtrip_preserves_token_set() { + let auth = KimiAuth { + key: "at".into(), + refresh_token: Some("rt".into()), + expires_at: Some(Utc::now()), + expires_in: Some(3600), + scope: Some("kimi-code".into()), + token_type: Some("bearer".into()), + ..KimiAuth::test_default() + }; + let json = serde_json::to_string(&auth).unwrap(); assert!( - !with_issuer(AuthMode::External, Some("https://idp.acme.example")).is_session_auth() + json.contains("\"oauth\""), + "wire spelling is \"oauth\": {json}" ); - - // Plain API keys never do. - assert!(!with_issuer(AuthMode::ApiKey, Some(XAI_OAUTH2_ISSUER)).is_session_auth()); - } - - #[test] - fn lookup_auth_skips_weblogin_on_primary_scope() { - let mut map = AuthStore::new(); - map.insert("scope".into(), make_auth(AuthMode::WebLogin)); - assert!(lookup_auth(&map, "scope").is_none()); - } - - #[test] - fn lookup_auth_skips_weblogin_on_legacy_fallback() { - let mut map = AuthStore::new(); - map.insert(LEGACY_SCOPE.into(), make_auth(AuthMode::WebLogin)); - assert!(lookup_auth(&map, "other-scope").is_none()); - } - - #[test] - fn lookup_auth_returns_oidc_token() { - let mut map = AuthStore::new(); - map.insert("scope".into(), make_auth(AuthMode::Oidc)); - assert!(lookup_auth(&map, "scope").is_some()); - } - - #[test] - fn lookup_auth_returns_api_key_token() { - let mut map = AuthStore::new(); - map.insert("scope".into(), make_auth(AuthMode::ApiKey)); - assert!(lookup_auth(&map, "scope").is_some()); - } - - /// subscriptionTier present → deserializes to Some. - #[test] - fn user_info_subscription_tier_present() { - let json = r#"{ - "userId": "u1", - "subscriptionTier": "SuperGrokPro" - }"#; - let info: UserInfo = serde_json::from_str(json).unwrap(); - assert_eq!(info.subscription_tier.as_deref(), Some("SuperGrokPro")); - } - - /// subscriptionTier absent → deserializes to None (backwards compat). - #[test] - fn user_info_subscription_tier_absent() { - let json = r#"{"userId": "u1"}"#; - let info: UserInfo = serde_json::from_str(json).unwrap(); - assert!(info.subscription_tier.is_none()); - } - - /// subscriptionTier null → deserializes to None. - #[test] - fn user_info_subscription_tier_null() { - let json = r#"{"userId": "u1", "subscriptionTier": null}"#; - let info: UserInfo = serde_json::from_str(json).unwrap(); - assert!(info.subscription_tier.is_none()); - } - - /// subscriptionTier empty string → deserializes to Some(""). - /// The paywall poller treats this as "no subscription" (line 230: - /// `Some(tier) if !tier.is_empty()`) and keeps polling. - #[test] - fn user_info_subscription_tier_empty_string() { - let json = r#"{"userId": "u1", "subscriptionTier": ""}"#; - let info: UserInfo = serde_json::from_str(json).unwrap(); - assert_eq!(info.subscription_tier.as_deref(), Some("")); + let back: KimiAuth = serde_json::from_str(&json).unwrap(); + assert_eq!(back.key, "at"); + assert_eq!(back.refresh_token.as_deref(), Some("rt")); + assert_eq!(back.expires_in, Some(3600)); + assert_eq!(back.scope.as_deref(), Some("kimi-code")); + assert_eq!(back.token_type.as_deref(), Some("bearer")); + assert_eq!(back.auth_mode, AuthMode::OAuth); } } diff --git a/crates/codegen/kigi-shell/src/auth/oidc/login.rs b/crates/codegen/kigi-shell/src/auth/oidc/login.rs deleted file mode 100644 index 643bdb8..0000000 --- a/crates/codegen/kigi-shell/src/auth/oidc/login.rs +++ /dev/null @@ -1,701 +0,0 @@ -//! Interactive login orchestration: callback HTTP server, browser -//! handoff, stdin paste fallback, race between the two. -//! -//! Cross-references [`super::protocol`] for OIDC mechanics and -//! [`super::super::AuthManager`] for credential persistence. - -use std::collections::HashMap; -use std::io::IsTerminal; -use std::sync::Arc; - -use axum::{ - Router, - extract::{Query, State}, - http::{Method, StatusCode}, - response::Html, - routing::get, -}; -use tokio::net::TcpListener; - -use super::super::config::{GrokComConfig, OidcAuthConfig}; -use super::super::{AuthManager, GrokAuth}; -use super::protocol::{ - OidcError, build_authorize_url, build_grok_auth, discover, enforce_login_principal, - exchange_code, extract_user_info, generate_pkce, login_principal_policy, - peek_access_token_principal, peek_access_token_principal_id, validate_state, -}; - -/// Maximum time to wait for the browser OAuth callback (or manual paste of the code). -/// 10 minutes is long enough for users who step away briefly during login. -const AUTH_CALLBACK_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(600); - -/// Parse user-pasted input into `(code, state)`. -/// -/// Accepts two formats: -/// 1. Full callback URL: `http://127.0.0.1:PORT/callback?code=XXX&state=YYY` -/// 2. Bare authorization code: `abc123` -fn parse_pasted_input(input: &str) -> Result { - let input = input.trim(); - if input.is_empty() { - return Err(OidcError::InvalidPastedInput("empty input".into())); - } - - if let Ok(url) = url::Url::parse(input) { - let params: HashMap = url.query_pairs().into_owned().collect(); - if let Some(code) = params.get("code") { - let state = params.get("state").cloned().unwrap_or_default(); - return Ok(Callback { - code: code.clone(), - state, - }); - } - if let Some(error) = params.get("error") { - let desc = params.get("error_description").cloned().unwrap_or_default(); - return Err(OidcError::CallbackAuthFailed(if desc.is_empty() { - error.clone() - } else { - format!("{error}: {desc}") - })); - } - return Err(OidcError::InvalidPastedInput( - "URL has no 'code' query parameter".into(), - )); - } - - Ok(Callback { - code: input.to_owned(), - state: String::new(), - }) -} - -/// Render a styled callback page shown in the browser after the OAuth redirect. -pub(crate) fn callback_page(title: &str, message: &str, is_success: bool) -> String { - let icon = if is_success { - // Grok logo - r#""# - } else { - // X circle - r#""# - }; - format!( - r#" - - - - - -{title} - - - -
- {icon} -

{title}

-

{message}

-
- -"#, - title = title, - icon = icon, - message = message, - ) -} - -/// Build the axum router for the OIDC loopback callback server. -fn build_callback_router(tx: tokio::sync::mpsc::Sender) -> Router { - let cors = - crate::auth::config::accounts_app_cors_layer(Method::GET).allow_private_network(true); - - Router::new() - .route("/callback", get(handle_callback)) - .layer(cors) - .with_state(tx) -} - -async fn handle_callback( - State(tx): State>, - Query(params): Query>, -) -> (StatusCode, Html) { - let result = parse_callback_params(¶ms); - let response = callback_response(&result); - if let Err(e) = tx.try_send(result) { - tracing::error!(?e, "OIDC: callback channel send failed; auth will time out"); - } - response -} - -fn parse_callback_params(params: &HashMap) -> CallbackResult { - if let Some(code) = params.get("code") { - let state = params.get("state").cloned().unwrap_or_default(); - tracing::debug!(state = %state, "OIDC: received code via loopback callback"); - return Ok(Callback { - code: code.clone(), - state, - }); - } - let error = params.get("error").cloned().unwrap_or_default(); - let desc = params.get("error_description").cloned().unwrap_or_default(); - tracing::error!(error = %error, desc = %desc, "OIDC: IdP returned error"); - Err(if desc.is_empty() { - error - } else { - format!("{error}: {desc}") - }) -} - -fn callback_response(result: &CallbackResult) -> (StatusCode, Html) { - let (title, message) = match result { - Ok(_) => ( - "Signed in", - "You can close this window and return to Grok Build.", - ), - Err(_) => ("Access denied", "Close this window and try again."), - }; - ( - StatusCode::OK, - Html(callback_page(title, message, result.is_ok())), - ) -} - -/// Wait until stdin has data or `tx` is closed. Returns `false` if closed. -#[cfg(unix)] -fn wait_for_stdin_or_closed( - stdin: &std::io::Stdin, - tx: &tokio::sync::mpsc::Sender, -) -> bool { - use std::os::unix::io::AsRawFd; - let fd = stdin.as_raw_fd(); - loop { - if tx.is_closed() { - return false; - } - let ready = unsafe { - let mut fds = std::mem::zeroed::(); - fds.fd = fd; - fds.events = libc::POLLIN; - libc::poll(&mut fds, 1, 200) - }; - if ready > 0 { - return true; - } - } -} - -fn spawn_stdin_reader(tx: tokio::sync::mpsc::Sender) { - tokio::task::spawn_blocking(move || { - use std::io::BufRead; - let stdin = std::io::stdin(); - let mut buf = String::new(); - loop { - #[cfg(unix)] - if !wait_for_stdin_or_closed(&stdin, &tx) { - tracing::debug!("OIDC: stdin reader exiting, channel closed"); - return; - } - #[cfg(not(unix))] - if tx.is_closed() { - tracing::debug!("OIDC: stdin reader exiting, channel closed"); - return; - } - - buf.clear(); - let mut handle = stdin.lock(); - match handle.read_line(&mut buf) { - Ok(0) => return, - Ok(_) => {} - Err(_) => return, - } - drop(handle); - - let trimmed = buf.trim().to_owned(); - if trimmed.is_empty() { - continue; - } - match parse_pasted_input(&trimmed) { - Ok(result) => { - tracing::debug!("OIDC: received code via stdin paste"); - let _ = tx.blocking_send(Ok(result)); - return; - } - Err(OidcError::InvalidPastedInput(msg)) => { - tracing::debug!(input = %msg, "OIDC: invalid stdin paste, retrying"); - eprintln!(" Invalid input: {msg}. Try again:"); - } - Err(e) => { - tracing::warn!(error = %e, "OIDC: stdin paste returned auth error"); - let _ = tx.blocking_send(Err(e.to_string())); - return; - } - } - } - }); -} - -/// Race loopback callback against manual paste from `code_rx`. -async fn race_callback_and_client_ui( - listener: TcpListener, - code_rx: &mut tokio::sync::mpsc::Receiver, -) -> anyhow::Result { - tracing::debug!("OIDC: waiting for auth code (loopback + client paste)"); - let (tx, mut rx) = tokio::sync::mpsc::channel::(1); - let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>(); - - let app = build_callback_router(tx.clone()); - let server = tokio::spawn(async move { - let _ = axum::serve(listener, app) - .with_graceful_shutdown(async { - let _ = shutdown_rx.await; - }) - .await; - }); - - // Bridge client paste input into the callback channel. - let client_tx = tx.clone(); - let client_bridge = async { - while let Some(code) = code_rx.recv().await { - match parse_pasted_input(&code) { - Ok(result) => { - tracing::debug!("OIDC: received code via client paste"); - let _ = client_tx.send(Ok(result)).await; - return; - } - Err(e) => { - tracing::debug!(error = %e, "OIDC: invalid client paste input"); - } - } - } - }; - - drop(tx); - - let result = tokio::select! { - r = tokio::time::timeout(AUTH_CALLBACK_TIMEOUT, rx.recv()) => { - r.map_err(|_| anyhow::Error::new(OidcError::CallbackTimeout))? - .ok_or_else(|| anyhow::Error::new(OidcError::CallbackChannelClosed))? - } - _ = client_bridge => { - rx.recv().await - .ok_or_else(|| anyhow::Error::new(OidcError::CallbackChannelClosed))? - } - }; - - let _ = shutdown_tx.send(()); - let _ = server.await; - - result.map_err(|e| anyhow::Error::new(OidcError::CallbackAuthFailed(e))) -} - -/// Race loopback callback against stdin paste. -async fn race_callback_and_stdin( - listener: TcpListener, - enable_stdin: bool, -) -> anyhow::Result { - tracing::debug!( - enable_stdin = enable_stdin, - "OIDC: waiting for auth code (loopback + stdin)" - ); - let (tx, mut rx) = tokio::sync::mpsc::channel::(1); - let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>(); - - let app = build_callback_router(tx.clone()); - let server = tokio::spawn(async move { - let _ = axum::serve(listener, app) - .with_graceful_shutdown(async { - let _ = shutdown_rx.await; - }) - .await; - }); - - if enable_stdin { - spawn_stdin_reader(tx.clone()); - } - - drop(tx); - - let result = tokio::time::timeout(AUTH_CALLBACK_TIMEOUT, rx.recv()) - .await - .map_err(|_| { - // "10 minutes" must match AUTH_CALLBACK_TIMEOUT above - tracing::error!("auth: timed out after 10 minutes waiting for auth code"); - anyhow::Error::new(OidcError::CallbackTimeout) - })? - .ok_or_else(|| { - tracing::error!( - "OIDC: callback channel closed, no code received from loopback or stdin" - ); - anyhow::Error::new(OidcError::CallbackChannelClosed) - })?; - - let _ = shutdown_tx.send(()); - let _ = server.await; - - result.map_err(|e| anyhow::Error::new(OidcError::CallbackAuthFailed(e))) -} - -/// Run the full OIDC login flow: discovery → PKCE → browser → callback → token exchange → persist. -pub async fn run_login_flow( - config: &GrokComConfig, - auth_manager: &Arc, - channels: Option, -) -> anyhow::Result<(GrokAuth, bool)> { - let oidc = config - .oidc - .as_ref() - .ok_or_else(|| anyhow::Error::new(OidcError::NotConfigured))?; - run_login_flow_with_config(oidc, auth_manager, channels).await -} - -/// Run the OIDC login flow with an explicit [`OidcAuthConfig`]. -/// -/// Also used by the OAuth2 provider path via [`OAuth2ProviderConfig::as_oidc`]. -/// -/// The flow races two input paths: -/// - **Path A**: A loopback HTTP server on `127.0.0.1` that receives the IdP redirect. -/// - **Path B**: Stdin paste — the user manually pastes the callback URL or bare auth code. -/// -/// Path B is essential for remote VMs where the browser runs on a different machine -/// and the `127.0.0.1` redirect cannot reach the CLI process. -/// * `channels` — `Some`: pushes the auth URL to the TUI and receives pasted codes. -/// `None`: prints to stderr / reads stdin (CLI mode). -pub async fn run_login_flow_with_config( - oidc: &OidcAuthConfig, - auth_manager: &Arc, - channels: Option, -) -> anyhow::Result<(GrokAuth, bool)> { - tracing::info!(issuer = %oidc.issuer, client_id = %oidc.client_id, "OIDC: starting login flow"); - - // Ensure jsonwebtoken CryptoProvider is installed (required for JWT validation). - jsonwebtoken::crypto::CryptoProvider::install_default( - &jsonwebtoken::crypto::rust_crypto::DEFAULT_PROVIDER, - ) - .ok(); - - let discovery = discover(&oidc.issuer).await?; - let pkce = generate_pkce(); - let state = uuid::Uuid::now_v7().to_string(); - let nonce = uuid::Uuid::now_v7().to_string(); - - // In local-dev mode, use a fixed callback port so the redirect_uri is stable - // and can be pre-registered with the local OAuth2 provider. In production the - // OS picks a random available port. - let callback_port: u16 = if super::super::config::use_local_auth() { - 56121 - } else { - 0 - }; - let listener = TcpListener::bind(("127.0.0.1", callback_port)) - .await - .map_err(|e| anyhow::Error::new(OidcError::BindLoopback(e.to_string())))?; - let port = listener.local_addr()?.port(); - let redirect_uri = format!("http://127.0.0.1:{}/callback", port); - let oauth2 = auth_manager.grok_com_config().oauth2.as_ref(); - let auth_url = build_authorize_url( - oidc, - oauth2, - &discovery, - &redirect_uri, - &pkce, - &state, - &nonce, - ); - tracing::debug!(port = port, redirect_uri = %redirect_uri, "OIDC: callback server bound"); - - let (url_tx, code_rx) = match channels { - Some(ch) => (ch.url_tx, Some(ch.code_rx)), - None => (None, None), - }; - let has_client_ui = code_rx.is_some(); - - if has_client_ui { - // Client provides its own auth UI; just open the browser. - if let Err(e) = webbrowser::open(&auth_url) { - tracing::debug!(error = %e, "OIDC: failed to open browser"); - } - } else { - // No client UI — print to stderr. - eprintln!(); - let provider_label = if oidc.issuer == super::super::config::XAI_OAUTH2_ISSUER { - "Grok".to_owned() - } else { - oidc.issuer.clone() - }; - eprintln!("Signing in with {}...", provider_label); - eprintln!(); - if let Err(e) = webbrowser::open(&auth_url) { - tracing::debug!(error = %e, "OIDC: failed to open browser"); - } - eprintln!("Open this URL to sign in:"); - eprintln!(" {}", auth_url); - } - - let use_stdin = !has_client_ui && std::io::stdin().is_terminal(); - if use_stdin { - eprintln!(); - eprintln!("Paste the URL here if it doesn't connect:"); - } - - // Push auth URL to the TUI via oneshot. - if let Some(tx) = url_tx { - let _ = tx.send(super::super::flow::AuthUrlInfo { - url: auth_url.clone(), - mode: super::super::flow::AuthUrlMode::Loopback, - }); - } - - let Callback { - code, - state: received_state, - } = if let Some(mut rx) = code_rx { - // Client UI: race loopback against manual paste via code_rx. - race_callback_and_client_ui(listener, &mut rx).await? - } else { - // No client UI: race loopback against stdin paste. - race_callback_and_stdin(listener, use_stdin).await? - }; - - // Validate state (skip for bare code paste where state is empty) - if !received_state.is_empty() { - validate_state(&state, &received_state)?; - } - - let tokens = exchange_code( - &discovery.token_endpoint, - &code, - &redirect_uri, - &oidc.client_id, - &pkce.code_verifier, - ) - .await?; - tracing::info!( - has_refresh = tokens.refresh_token.is_some(), - expires_in = ?tokens.expires_in, - "OIDC: token exchange complete" - ); - - // Resolve the actual principal chosen on the consent screen. - // - // The shell's config may not have principal_type set (personal login), - // but the user might pick "Team" on the consent screen. The server - // encodes the chosen principal in the access token JWT. If the config - // doesn't specify a principal, peek at the token to discover it. - let token_principal = peek_access_token_principal(&tokens.access_token); - - // The authorize URL only pre-selects; verify the token's principal here. - // Match the principal id even if `principal_type` is absent. - let principal_policy = login_principal_policy(auth_manager.grok_com_config()); - enforce_login_principal( - principal_policy.as_ref(), - peek_access_token_principal_id(&tokens.access_token).as_deref(), - )?; - - let (resolved_principal_type, resolved_principal_id, resolved_team_id) = { - let cfg_pt = oauth2.and_then(|cfg| cfg.principal_type.clone()); - let cfg_pid = oauth2.and_then(|cfg| cfg.principal_id.clone()); - if cfg_pt.is_some() { - (cfg_pt, cfg_pid, None) - } else if let Some((pt, pid, tid)) = token_principal { - tracing::info!( - principal_type = %pt, - principal_id = %pid, - team_id = ?tid, - "OIDC: resolved principal from access token" - ); - (Some(pt), Some(pid), tid) - } else { - (cfg_pt, cfg_pid, None) - } - }; - - let user_info = extract_user_info( - tokens.id_token.as_deref(), - &discovery, - &oidc.issuer, - &oidc.client_id, - &nonce, - resolved_principal_type.as_deref(), - resolved_principal_id.as_deref(), - resolved_team_id, - ) - .await?; - tracing::debug!(user_id = %user_info.user_id, "OIDC: extracted user info"); - - let mut auth = build_grok_auth(tokens, user_info, &oidc.issuer, &oidc.client_id); - auth_manager.enrich_auth_inline(&mut auth).await; - let auth = auth_manager - .update(auth) - .await - .map_err(|e| anyhow::Error::new(OidcError::SaveAuth(e.to_string())))?; - tracing::info!(user_id = %auth.user_id, "OIDC: login complete, credentials saved"); - - Ok((auth, true)) -} - -/// Successful OIDC callback payload. -#[derive(Debug, PartialEq, Eq)] -struct Callback { - code: String, - state: String, -} - -/// Result from the OIDC callback: either a [`Callback`] or an IdP error message. -type CallbackResult = Result; - -#[cfg(test)] -mod tests { - use super::super::test_helpers::*; - use super::*; - - /// End-to-end test: mock IdP + full login flow with code arriving via loopback. - /// Exercises discovery → PKCE → race_callback_and_stdin → token exchange → user info → persist. - #[tokio::test] - async fn full_login_flow_via_race() { - ensure_crypto_provider(); - let (issuer, idp_server) = start_mock_idp().await; - let temp_dir = tempfile::tempdir().unwrap(); - // Dead proxy port: inline `/user` enrichment fails fast in tests. - let dead_proxy = { - let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); - format!("http://127.0.0.1:{}", l.local_addr().unwrap().port()) - }; - let auth_manager = Arc::new( - AuthManager::new(temp_dir.path(), GrokComConfig::default()) - .with_proxy_base_url(&dead_proxy), - ); - - let oidc_cfg = OidcAuthConfig { - issuer: issuer.clone(), - client_id: TEST_CLIENT_ID.into(), - scopes: vec!["openid".into(), "email".into()], - audience: None, - }; - let discovery = discover(&oidc_cfg.issuer).await.unwrap(); - let pkce = generate_pkce(); - let state = "test-state".to_string(); - - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let redirect_uri = format!("http://127.0.0.1:{port}/callback"); - let _auth_url = build_authorize_url( - &oidc_cfg, - None, - &discovery, - &redirect_uri, - &pkce, - &state, - TEST_NONCE, - ); - - // Simulate browser callback via race_callback_and_stdin - let Callback { - code, - state: received_state, - } = tokio::join!(race_callback_and_stdin(listener, false), async { - tokio::time::sleep(std::time::Duration::from_millis(50)).await; - reqwest::get(format!( - "http://127.0.0.1:{port}/callback?code=mock-auth-code&state={state}" - )) - .await - .unwrap(); - }) - .0 - .unwrap(); - - assert_eq!(code, "mock-auth-code"); - assert_eq!(received_state, state); - - let tokens = exchange_code( - &discovery.token_endpoint, - &code, - &redirect_uri, - &oidc_cfg.client_id, - &pkce.code_verifier, - ) - .await - .unwrap(); - assert_eq!(tokens.access_token, "mock-access-token"); - - let user_info = extract_user_info( - tokens.id_token.as_deref(), - &discovery, - &oidc_cfg.issuer, - &oidc_cfg.client_id, - TEST_NONCE, - None, - None, - None, - ) - .await - .unwrap(); - let auth = build_grok_auth(tokens, user_info, &oidc_cfg.issuer, &oidc_cfg.client_id); - let auth = auth_manager.update(auth).await.unwrap(); - - assert_eq!(auth.key, "mock-access-token"); - assert_eq!(auth.refresh_token.as_deref(), Some("mock-refresh-token")); - assert_eq!(auth.user_id, "user-42"); - assert_eq!(auth.email.as_deref(), Some("test@corp.com")); - assert!(auth.principal_type.is_none()); - assert!(auth.principal_id.is_none()); - assert!(auth.expires_at.is_some()); - assert_eq!(auth.oidc_issuer.as_deref(), Some(issuer.as_str())); - - let auth_json = std::fs::read_to_string(temp_dir.path().join("auth.json")).unwrap(); - assert!(auth_json.contains("mock-access-token")); - assert!(auth_json.contains("user-42")); - - idp_server.abort(); - } - /// Parser matrix: full callback URL, bare code, error URL, empty. - /// Each case is one bug class: - /// - full URL: regression in URL extraction - /// - bare code: paste-friendly fallback - /// - error URL: surfaces IdP error to user - /// - empty: input validation - #[test] - fn parse_pasted_input_matrix() { - // (input, expected: Ok((code, state)) | Err substring) - let ok_cases: &[(&str, &str, &str)] = &[ - ( - "http://127.0.0.1:54321/callback?code=abc123&state=xyz789", - "abc123", - "xyz789", - ), - ("abc123def456", "abc123def456", ""), - ]; - for (input, code, state) in ok_cases { - let cb = - parse_pasted_input(input).unwrap_or_else(|e| panic!("parse {input:?} failed: {e}")); - assert_eq!(cb.code, *code, "code for {input:?}"); - assert_eq!(cb.state, *state, "state for {input:?}"); - } - - let err_cases: &[(&str, &str)] = &[ - ( - "http://127.0.0.1:54321/callback?error=access_denied&error_description=User+denied", - "access_denied", - ), - ("", ""), - (" ", ""), - ]; - for (input, expected_substr) in err_cases { - let err = parse_pasted_input(input).unwrap_err(); - if !expected_substr.is_empty() { - assert!( - err.to_string().contains(expected_substr), - "input {input:?} -> unexpected err: {err}" - ); - } - } - } -} diff --git a/crates/codegen/kigi-shell/src/auth/oidc/mod.rs b/crates/codegen/kigi-shell/src/auth/oidc/mod.rs deleted file mode 100644 index cb979fa..0000000 --- a/crates/codegen/kigi-shell/src/auth/oidc/mod.rs +++ /dev/null @@ -1,14 +0,0 @@ -//! OIDC authentication: protocol, login, and refresh submodules. - -mod login; -pub(crate) mod protocol; -pub(crate) mod refresh; -#[cfg(test)] -mod test_helpers; - -pub use login::{run_login_flow, run_login_flow_with_config}; -pub(crate) use protocol::{ - enforce_login_principal, is_configured, login_principal_policy, peek_access_token_principal, - peek_access_token_principal_id, with_alpha_test_key, -}; -pub(crate) use refresh::{OidcRefreshResult, oidc_token_exchange}; diff --git a/crates/codegen/kigi-shell/src/auth/oidc/protocol.rs b/crates/codegen/kigi-shell/src/auth/oidc/protocol.rs deleted file mode 100644 index 1a7a736..0000000 --- a/crates/codegen/kigi-shell/src/auth/oidc/protocol.rs +++ /dev/null @@ -1,1283 +0,0 @@ -//! Pure OIDC protocol mechanics: PKCE, discovery, token exchange, -//! refresh_tokens, JWT validation, principal extraction. -//! -//! No `AuthManager` mutation here. The login orchestration is in -//! [`super::login`]; refresh primitives are in [`super::refresh`]. -use super::super::config::{ForceLoginTeam, GrokComConfig, OAuth2ProviderConfig, OidcAuthConfig}; -use super::super::{AuthMode, GrokAuth}; -use base64::Engine; -use base64::engine::general_purpose::URL_SAFE_NO_PAD; -use chrono::{Duration, Utc}; -use parking_lot::RwLock; -use serde::Deserialize; -use sha2::{Digest, Sha256}; -use std::collections::HashMap; -use std::sync::LazyLock; -use std::time::{Duration as StdDuration, Instant}; -#[derive(Debug, Clone, thiserror::Error)] -pub(super) enum OidcError { - #[error("OIDC not configured")] - NotConfigured, - #[error("failed to bind OIDC loopback server: {0}")] - BindLoopback(String), - #[error("failed to save OIDC auth: {0}")] - SaveAuth(String), - #[error("OIDC discovery failed: HTTP {status} from {url}")] - DiscoveryHttp { status: u16, url: String }, - /// Keep the "10 minutes" text in sync with `AUTH_CALLBACK_TIMEOUT` in `login.rs`. - #[error("Login timed out after 10 minutes. Please try again.")] - CallbackTimeout, - #[error("OIDC callback channel closed unexpectedly")] - CallbackChannelClosed, - #[error("OIDC authentication failed: {0}")] - CallbackAuthFailed(String), - #[error("failed to parse pasted input: {0}")] - InvalidPastedInput(String), - #[error("OIDC token exchange failed: HTTP {status} — {body}")] - TokenExchangeHttp { status: u16, body: String }, - #[error("OIDC token refresh failed: HTTP {status} — {body}")] - TokenRefreshHttp { status: u16, body: String }, - #[error("OIDC authentication failed: state mismatch")] - StateMismatch, - #[error("OIDC id_token uses unsupported algorithm: {0}")] - UnsupportedAlg(String), - #[error("OIDC id_token alg {alg} is not in discovery supported list")] - AlgNotInDiscoverySupportedList { alg: String }, - #[error("OIDC id_token missing kid header")] - IdTokenMissingKid, - #[error("OIDC discovery missing jwks_uri")] - DiscoveryMissingJwksUri, - #[error("OIDC JWK not found for kid={kid}")] - JwkNotFound { kid: String }, - #[error("OIDC id_token issuer mismatch")] - IssuerMismatch, - #[error("OIDC id_token audience mismatch")] - AudienceMismatch, - #[error("OIDC id_token nonce mismatch")] - NonceMismatch, - #[error("OIDC token response missing id_token")] - MissingIdToken, - #[error("OIDC id_token validation failed: {0}")] - IdTokenValidationFailed(String), - #[error( - "This deployment requires logging into {expected}; your login returned {}", - actual.as_deref().unwrap_or("no team principal") - )] - PinnedPrincipalMismatch { - /// Pre-formatted requirement, e.g. `team ` or `one of teams: a, b`. - expected: String, - actual: Option, - }, - #[error( - "Login is blocked by your administrator: force_login_team_uuid is an empty \ - list, so no team is permitted to sign in" - )] - ForceLoginNoPrincipalsAllowed, -} -const ALLOWED_ID_TOKEN_ALGS: &[jsonwebtoken::Algorithm] = &[ - jsonwebtoken::Algorithm::RS256, - jsonwebtoken::Algorithm::RS384, - jsonwebtoken::Algorithm::RS512, - jsonwebtoken::Algorithm::PS256, - jsonwebtoken::Algorithm::PS384, - jsonwebtoken::Algorithm::PS512, - jsonwebtoken::Algorithm::ES256, - jsonwebtoken::Algorithm::ES384, - jsonwebtoken::Algorithm::EdDSA, -]; -/// Optionally attach an extra access header when the optional non-production -/// feature is enabled and the request targets a matching first-party host. -pub(crate) fn with_alpha_test_key( - builder: reqwest::RequestBuilder, - url: &str, -) -> reqwest::RequestBuilder { - let _ = url; - builder -} -pub fn is_configured(config: &GrokComConfig) -> bool { - config.oidc.is_some() -} -/// Peek at the unverified access token JWT to extract the `principal_type` -/// and `principal_id` chosen during the consent screen. -/// -/// When the user picks "Team" on the consent screen, the server strips -/// user-only scopes (`openid`, `email`) and issues the token with -/// `principal_type=Team`. The shell's config doesn't know which principal -/// the user picked, so we peek at the token to find out. -/// -/// Returns `(principal_type, principal_id)` or `None` if the token is not -/// a JWT or the claims can't be extracted. -pub(crate) fn peek_access_token_principal( - access_token: &str, -) -> Option<(String, String, Option)> { - #[derive(serde::Deserialize)] - struct MinimalClaims { - #[serde(default, alias = "principalType")] - principal_type: Option, - #[serde(default, alias = "principalId")] - principal_id: Option, - #[serde(default)] - team_id: Option, - } - let token_data = - jsonwebtoken::dangerous::insecure_decode::(access_token).ok()?; - let pt = token_data.claims.principal_type?; - let pid = token_data.claims.principal_id?; - if pt.is_empty() || pid.is_empty() { - return None; - } - let tid = token_data.claims.team_id.filter(|s| !s.is_empty()); - Some((pt, pid, tid)) -} -/// Extract just the `principal_id` claim for `force_login_team_uuid` matching, -/// regardless of whether `principal_type` is present. A token can carry the -/// team id in `principal_id` without a `principal_type`; the pin must still -/// match it. Matching the id alone is safe because a user id never collides -/// with a team uuid (distinct id spaces), and the server re-validates the -/// signed token anyway. Returns `None` only when no non-empty `principal_id` -/// is present (which `enforce_login_principal` treats as fail-closed). -pub(crate) fn peek_access_token_principal_id(access_token: &str) -> Option { - #[derive(serde::Deserialize)] - struct PrincipalIdClaim { - #[serde(default, alias = "principalId")] - principal_id: Option, - } - jsonwebtoken::dangerous::insecure_decode::(access_token) - .ok()? - .claims - .principal_id - .filter(|s| !s.is_empty()) -} -/// Resolved allowed-team set from the dedicated `force_login_team_uuid` lockdown -/// knob, or `None` (unrestricted). The legacy `oauth2.principal_id` is -/// intentionally NOT an enforcement gate — it only pre-selects the team on the -/// consent page — so deployments that set it for pre-selection keep letting -/// users pick a team (no surprise login failures on upgrade). Pure for testing. -pub(crate) fn resolve_login_principal_policy( - force_login_team_uuid: Option<&ForceLoginTeam>, -) -> Option { - force_login_team_uuid.cloned() -} -pub(crate) fn login_principal_policy(cfg: &GrokComConfig) -> Option { - resolve_login_principal_policy(cfg.force_login_team_uuid.as_ref()) -} -/// Reject a token whose principal isn't allowed, BEFORE persisting (no partial -/// state). A restriction also rejects a token with no principal (else picking -/// "personal" on the consent page defeats it); an empty `AnyOf` fails closed. -/// -/// The `actual` principal comes from the access-token claim -/// (`peek_access_token_principal`, an unverified `insecure_decode`). This -/// client-side check is fail-fast UX / defense-in-depth — NOT the security -/// boundary: the server re-validates the signed token on every API call and -/// is authoritative, so a locally tampered token still cannot reach the API. -pub(crate) fn enforce_login_principal( - policy: Option<&ForceLoginTeam>, - actual: Option<&str>, -) -> anyhow::Result<()> { - let allowed: &[String] = match policy { - None => return Ok(()), - Some(ForceLoginTeam::Single(id)) => std::slice::from_ref(id), - Some(ForceLoginTeam::AnyOf(ids)) if ids.is_empty() => { - tracing::warn!("OIDC: force_login_team_uuid is an empty list; failing closed"); - return Err(anyhow::Error::new(OidcError::ForceLoginNoPrincipalsAllowed)); - } - Some(ForceLoginTeam::AnyOf(ids)) => ids, - }; - if let Some(actual) = actual - && allowed.iter().any(|a| a == actual) - { - return Ok(()); - } - let expected = if allowed.len() == 1 { - format!("team {}", allowed[0]) - } else { - format!("one of teams: {}", allowed.join(", ")) - }; - tracing::warn!( - expected = % expected, actual = ? actual, - "OIDC: login principal does not satisfy required policy; rejecting" - ); - Err(anyhow::Error::new(OidcError::PinnedPrincipalMismatch { - expected, - actual: actual.map(str::to_owned), - })) -} -#[derive(Debug)] -pub(super) struct OidcUserInfo { - pub(super) user_id: String, - pub(super) email: Option, - pub(super) first_name: Option, - pub(super) last_name: Option, - pub(super) profile_image_asset_id: Option, - pub(super) principal_type: Option, - pub(super) principal_id: Option, - pub(super) team_id: Option, - pub(super) team_name: Option, - pub(super) team_role: Option, - pub(super) organization_id: Option, - pub(super) organization_name: Option, - pub(super) organization_role: Option, - pub(super) user_blocked_reason: Option, - pub(super) team_blocked_reasons: Vec, - pub(super) coding_data_retention_opt_out: bool, -} -pub(super) fn build_grok_auth( - tokens: TokenResponse, - user_info: OidcUserInfo, - issuer: &str, - client_id: &str, -) -> GrokAuth { - let now = Utc::now(); - GrokAuth { - key: tokens.access_token, - auth_mode: AuthMode::Oidc, - create_time: now, - user_id: user_info.user_id, - email: user_info.email, - first_name: user_info.first_name, - last_name: user_info.last_name, - profile_image_asset_id: user_info.profile_image_asset_id, - principal_type: user_info.principal_type, - principal_id: user_info.principal_id, - team_id: user_info.team_id, - team_name: user_info.team_name, - team_role: user_info.team_role, - organization_id: user_info.organization_id, - organization_name: user_info.organization_name, - organization_role: user_info.organization_role, - user_blocked_reason: user_info.user_blocked_reason, - team_blocked_reasons: user_info.team_blocked_reasons, - coding_data_retention_opt_out: user_info.coding_data_retention_opt_out, - has_grok_code_access: None, - refresh_token: tokens.refresh_token, - expires_at: tokens.expires_in.map(|s| now + Duration::seconds(s as i64)), - oidc_issuer: Some(issuer.to_owned()), - oidc_client_id: Some(client_id.to_owned()), - } -} -#[derive(Debug, Clone, Deserialize)] -pub(super) struct Discovery { - pub(super) authorization_endpoint: String, - pub(super) token_endpoint: String, - #[serde(default)] - pub(super) jwks_uri: Option, - #[serde(default)] - pub(super) id_token_signing_alg_values_supported: Option>, -} -/// RFC 8414 says discovery clients SHOULD cache. 1h is short enough -/// that an endpoint move propagates within an agent session, long -/// enough that a discovery-endpoint outage no longer blocks token -/// refresh once the doc is cached. -const DISCOVERY_CACHE_TTL: StdDuration = StdDuration::from_secs(3600); -/// Per-issuer cache of `(Discovery, fetched_at)`. Process-global -/// because the discovery doc is identity-free; multiple AuthManagers -/// pointed at the same IdP share one entry. -static DISCOVERY_CACHE: LazyLock>> = - LazyLock::new(|| RwLock::new(HashMap::new())); -pub(super) async fn discover(issuer: &str) -> anyhow::Result { - let issuer_key = issuer.trim_end_matches('/').to_owned(); - if let Some((doc, at)) = DISCOVERY_CACHE.read().get(&issuer_key) - && at.elapsed() < DISCOVERY_CACHE_TTL - { - return Ok(doc.clone()); - } - use backon::Retryable; - let key = issuer_key.clone(); - let doc = (|| { - let key = key.clone(); - async move { discover_once(&key).await } - }) - .retry(discovery_retry_policy()) - .await?; - DISCOVERY_CACHE - .write() - .insert(issuer_key, (doc.clone(), Instant::now())); - Ok(doc) -} -fn discovery_retry_policy() -> backon::ExponentialBuilder { - backon::ExponentialBuilder::default() - .with_max_times(2) - .with_min_delay(StdDuration::from_millis(500)) - .with_max_delay(StdDuration::from_secs(2)) - .with_jitter() -} -async fn discover_once(issuer_key: &str) -> anyhow::Result { - let url = format!("{issuer_key}/.well-known/openid-configuration"); - tracing::debug!(url = % url, "OIDC: fetching discovery document"); - let resp = with_alpha_test_key( - crate::http::shared_client() - .get(&url) - .timeout(StdDuration::from_secs(10)), - &url, - ) - .send() - .await?; - if !resp.status().is_success() { - return Err(anyhow::Error::new(OidcError::DiscoveryHttp { - status: resp.status().as_u16(), - url, - })); - } - let doc: Discovery = resp.json().await?; - tracing::debug!( - authorization_endpoint = % doc.authorization_endpoint, token_endpoint = % doc - .token_endpoint, jwks_uri = ? doc.jwks_uri, id_token_algs = ? doc - .id_token_signing_alg_values_supported, "OIDC: discovery complete" - ); - Ok(doc) -} -#[cfg(test)] -pub(super) fn clear_discovery_cache() { - DISCOVERY_CACHE.write().clear(); -} -pub(super) struct Pkce { - pub(super) code_verifier: String, - pub(super) code_challenge: String, -} -pub(super) fn generate_pkce() -> Pkce { - let random_bytes: [u8; 32] = rand::random(); - let code_verifier = URL_SAFE_NO_PAD.encode(random_bytes); - let code_challenge = URL_SAFE_NO_PAD.encode(Sha256::digest(code_verifier.as_bytes())); - Pkce { - code_verifier, - code_challenge, - } -} -pub(super) fn build_authorize_url( - config: &OidcAuthConfig, - oauth2: Option<&OAuth2ProviderConfig>, - discovery: &Discovery, - redirect_uri: &str, - pkce: &Pkce, - state: &str, - nonce: &str, -) -> String { - let scopes = config.scopes.join(" "); - let mut url = format!( - "{}?response_type=code&client_id={}&redirect_uri={}&scope={}\ - &code_challenge={}&code_challenge_method=S256&state={}&nonce={}", - discovery.authorization_endpoint, - urlencoding::encode(&config.client_id), - urlencoding::encode(redirect_uri), - urlencoding::encode(&scopes), - urlencoding::encode(&pkce.code_challenge), - urlencoding::encode(state), - urlencoding::encode(nonce), - ); - if let Some(ref audience) = config.audience { - url.push_str(&format!("&audience={}", urlencoding::encode(audience))); - } - if let Some(oauth2) = oauth2 { - if let Some(ref principal_type) = oauth2.principal_type { - url.push_str(&format!( - "&principal_type={}", - urlencoding::encode(principal_type) - )); - } - if let Some(ref principal_id) = oauth2.principal_id { - url.push_str(&format!( - "&principal_id={}", - urlencoding::encode(principal_id) - )); - } - } - let referrer = oauth2 - .and_then(|o| o.referrer.as_deref()) - .filter(|r| !r.is_empty()) - .unwrap_or("grok-build"); - url.push_str(&format!("&referrer={}", urlencoding::encode(referrer))); - url -} -#[derive(Debug, Deserialize)] -pub(super) struct TokenResponse { - pub(super) access_token: String, - #[serde(default)] - pub(super) refresh_token: Option, - #[serde(default)] - pub(super) id_token: Option, - #[serde(default)] - pub(super) expires_in: Option, -} -pub(super) async fn exchange_code( - token_endpoint: &str, - code: &str, - redirect_uri: &str, - client_id: &str, - code_verifier: &str, -) -> anyhow::Result { - tracing::debug!( - token_endpoint = % token_endpoint, "OIDC: exchanging code for tokens" - ); - let resp = with_alpha_test_key( - crate::http::shared_client() - .post(token_endpoint) - .header("x-grok-client-version", kigi_version::VERSION) - .form(&[ - ("grant_type", "authorization_code"), - ("code", code), - ("redirect_uri", redirect_uri), - ("client_id", client_id), - ("code_verifier", code_verifier), - ]) - .timeout(std::time::Duration::from_secs(15)), - token_endpoint, - ) - .send() - .await?; - if !resp.status().is_success() { - let status = resp.status().as_u16(); - let body = resp.text().await.unwrap_or_default(); - return Err(anyhow::Error::new(OidcError::TokenExchangeHttp { - status, - body, - })); - } - Ok(resp.json().await?) -} -/// Retry gate for `refresh_tokens`. Defers to `classify_terminal` (the single -/// source of truth): only a recognized terminal code (`invalid_grant`, -/// `invalid_client`) stops retries. Everything else (5xx, 429, bare 4xx, or an -/// unrecognized/RFC-transient code) is retried. -fn is_transient_refresh_error(err: &anyhow::Error) -> bool { - let Some(OidcError::TokenRefreshHttp { status, body }) = err.downcast_ref::() else { - return true; - }; - if *status >= 500 || *status == 429 { - return true; - } - let error_code = serde_json::from_str::(body) - .ok() - .and_then(|v| v.get("error")?.as_str().map(str::to_owned)); - error_code - .as_deref() - .and_then(super::refresh::classify_terminal) - .is_none() -} -/// Up to 3 attempts (1 + 2 retries), 200ms-2s jittered exponential -/// backoff. Bounded so a hard outage still surfaces to the user -/// promptly via the existing `RefreshOutcome::TransientFailure` path. -fn refresh_retry_policy() -> backon::ExponentialBuilder { - backon::ExponentialBuilder::default() - .with_max_times(2) - .with_min_delay(StdDuration::from_millis(200)) - .with_max_delay(StdDuration::from_secs(2)) - .with_jitter() -} -pub(super) async fn refresh_tokens( - token_endpoint: &str, - refresh_token: &str, - client_id: &str, - principal_type: Option<&str>, - principal_id: Option<&str>, -) -> anyhow::Result { - use backon::Retryable; - tracing::debug!( - token_endpoint = % token_endpoint, principal_type = ? principal_type, - principal_id = ? principal_id, "OIDC: refreshing token" - ); - (|| { - refresh_tokens_once( - token_endpoint, - refresh_token, - client_id, - principal_type, - principal_id, - ) - }) - .retry(refresh_retry_policy()) - .when(is_transient_refresh_error) - .await -} -/// One unretried POST to `token_endpoint`. Errors carry the typed -/// `OidcError::TokenRefreshHttp` so the retry classifier can read the -/// status code and OAuth2 `error` field without re-parsing. -async fn refresh_tokens_once( - token_endpoint: &str, - refresh_token: &str, - client_id: &str, - principal_type: Option<&str>, - principal_id: Option<&str>, -) -> anyhow::Result { - let mut params = vec![ - ("grant_type", "refresh_token"), - ("refresh_token", refresh_token), - ("client_id", client_id), - ]; - if let Some(pt) = principal_type { - params.push(("principal_type", pt)); - } - if let Some(pid) = principal_id { - params.push(("principal_id", pid)); - } - let resp = with_alpha_test_key( - crate::http::shared_client() - .post(token_endpoint) - .form(¶ms) - .timeout(StdDuration::from_secs(15)), - token_endpoint, - ) - .send() - .await?; - if !resp.status().is_success() { - let status = resp.status().as_u16(); - let body = resp.text().await.unwrap_or_default(); - let error_code = serde_json::from_str::(&body) - .ok() - .and_then(|v| v.get("error")?.as_str().map(str::to_owned)); - tracing::warn!( - http_status = status, oauth2_error = ? error_code, rt_prefix = crate - ::auth::token_suffix(refresh_token), client_id = % client_id, principal_type - = ? principal_type, "OIDC: token refresh HTTP error" - ); - return Err(anyhow::Error::new(OidcError::TokenRefreshHttp { - status, - body, - })); - } - Ok(resp.json().await?) -} -#[derive(Debug, Deserialize)] -pub(super) struct IdTokenClaims { - #[serde(default)] - pub(super) sub: Option, - #[serde(default)] - pub(super) email: Option, - #[serde(default)] - pub(super) iss: Option, - #[serde(default)] - pub(super) aud: Option, - #[serde(default)] - pub(super) nonce: Option, - #[serde(default, alias = "given_name")] - pub(super) first_name: Option, - #[serde(default, alias = "family_name")] - pub(super) last_name: Option, - #[serde(default)] - pub(super) picture: Option, -} -pub(super) fn aud_matches(aud: &serde_json::Value, expected: &str) -> bool { - match aud { - serde_json::Value::String(s) => s == expected, - serde_json::Value::Array(values) => values - .iter() - .any(|v| matches!(v, serde_json::Value::String(s) if s == expected)), - _ => false, - } -} -pub(super) fn validate_state(expected: &str, received: &str) -> anyhow::Result<()> { - if received != expected { - tracing::warn!( - expected = % expected, received = % received, "OIDC: state mismatch" - ); - return Err(anyhow::Error::new(OidcError::StateMismatch)); - } - Ok(()) -} -/// Explicit JWA name mapping — avoids coupling to `jsonwebtoken::Algorithm`'s `Debug` repr. -pub(super) fn alg_to_jwa_name(alg: jsonwebtoken::Algorithm) -> &'static str { - match alg { - jsonwebtoken::Algorithm::RS256 => "RS256", - jsonwebtoken::Algorithm::RS384 => "RS384", - jsonwebtoken::Algorithm::RS512 => "RS512", - jsonwebtoken::Algorithm::PS256 => "PS256", - jsonwebtoken::Algorithm::PS384 => "PS384", - jsonwebtoken::Algorithm::PS512 => "PS512", - jsonwebtoken::Algorithm::ES256 => "ES256", - jsonwebtoken::Algorithm::ES384 => "ES384", - jsonwebtoken::Algorithm::EdDSA => "EdDSA", - other => match other { - jsonwebtoken::Algorithm::HS256 => "HS256", - jsonwebtoken::Algorithm::HS384 => "HS384", - jsonwebtoken::Algorithm::HS512 => "HS512", - _ => "unknown", - }, - } -} -pub(super) fn ensure_alg_allowed( - alg: jsonwebtoken::Algorithm, - discovery_supported_algs: Option<&[String]>, -) -> anyhow::Result<()> { - let alg_name = alg_to_jwa_name(alg); - if !ALLOWED_ID_TOKEN_ALGS.contains(&alg) { - return Err(anyhow::Error::new(OidcError::UnsupportedAlg( - alg_name.to_owned(), - ))); - } - if let Some(supported) = discovery_supported_algs - && !supported.iter().any(|a| a == alg_name) - { - return Err(anyhow::Error::new( - OidcError::AlgNotInDiscoverySupportedList { - alg: alg_name.to_owned(), - }, - )); - } - Ok(()) -} -pub(super) async fn validate_and_extract_user_info( - token: &str, - discovery: &Discovery, - expected_issuer: &str, - expected_client_id: &str, - expected_nonce: &str, -) -> anyhow::Result { - let header = jsonwebtoken::decode_header(token)?; - let kid = header - .kid - .ok_or_else(|| anyhow::Error::new(OidcError::IdTokenMissingKid))?; - let jwks_uri = discovery - .jwks_uri - .as_ref() - .ok_or_else(|| anyhow::Error::new(OidcError::DiscoveryMissingJwksUri))?; - let jwks: jsonwebtoken::jwk::JwkSet = with_alpha_test_key( - crate::http::shared_client() - .get(jwks_uri) - .timeout(std::time::Duration::from_secs(10)), - jwks_uri, - ) - .send() - .await? - .error_for_status()? - .json() - .await?; - let jwk = jwks - .find(&kid) - .ok_or_else(|| anyhow::Error::new(OidcError::JwkNotFound { kid: kid.clone() }))?; - let decoding_key = jsonwebtoken::DecodingKey::from_jwk(jwk)?; - let alg = header.alg; - ensure_alg_allowed( - alg, - discovery.id_token_signing_alg_values_supported.as_deref(), - )?; - let mut validation = jsonwebtoken::Validation::new(alg); - validation.set_issuer(&[expected_issuer]); - validation.set_audience(&[expected_client_id]); - validation.validate_exp = true; - validation.validate_aud = true; - validation.required_spec_claims = ["sub", "iss", "aud", "exp"] - .into_iter() - .map(ToOwned::to_owned) - .collect(); - let token_data = jsonwebtoken::decode::(token, &decoding_key, &validation)?; - if token_data.claims.iss.as_deref() != Some(expected_issuer) { - return Err(anyhow::Error::new(OidcError::IssuerMismatch)); - } - if let Some(ref aud) = token_data.claims.aud - && !aud_matches(aud, expected_client_id) - { - return Err(anyhow::Error::new(OidcError::AudienceMismatch)); - } - if token_data.claims.nonce.as_deref() != Some(expected_nonce) { - return Err(anyhow::Error::new(OidcError::NonceMismatch)); - } - Ok(OidcUserInfo { - user_id: token_data - .claims - .sub - .unwrap_or_else(|| "unknown".to_string()), - email: token_data.claims.email, - first_name: token_data.claims.first_name, - last_name: token_data.claims.last_name, - profile_image_asset_id: token_data.claims.picture, - principal_type: None, - principal_id: None, - team_id: None, - team_name: None, - team_role: None, - organization_id: None, - organization_name: None, - organization_role: None, - user_blocked_reason: None, - team_blocked_reasons: vec![], - coding_data_retention_opt_out: false, - }) -} -pub(super) async fn extract_user_info( - id_token: Option<&str>, - discovery: &Discovery, - expected_issuer: &str, - expected_client_id: &str, - expected_nonce: &str, - principal_type: Option<&str>, - principal_id: Option<&str>, - fallback_team_id: Option, -) -> anyhow::Result { - if principal_type == Some(crate::auth::model::TEAM_PRINCIPAL_TYPE) { - let team_user_id = principal_id.unwrap_or("unknown").to_owned(); - return Ok(OidcUserInfo { - user_id: team_user_id, - email: None, - first_name: None, - last_name: None, - profile_image_asset_id: None, - principal_type: Some(crate::auth::model::TEAM_PRINCIPAL_TYPE.to_string()), - principal_id: principal_id.map(ToOwned::to_owned), - team_id: principal_id.map(ToOwned::to_owned).or(fallback_team_id), - team_name: None, - team_role: None, - organization_id: None, - organization_name: None, - organization_role: None, - user_blocked_reason: None, - team_blocked_reasons: vec![], - coding_data_retention_opt_out: false, - }); - } - let token = id_token.ok_or_else(|| anyhow::Error::new(OidcError::MissingIdToken))?; - validate_and_extract_user_info( - token, - discovery, - expected_issuer, - expected_client_id, - expected_nonce, - ) - .await - .map(|mut user_info| { - user_info.principal_type = principal_type.map(ToOwned::to_owned); - user_info.principal_id = principal_id.map(ToOwned::to_owned); - if user_info.team_id.is_none() { - user_info.team_id = fallback_team_id; - } - user_info - }) - .map_err(|e| anyhow::Error::new(OidcError::IdTokenValidationFailed(e.to_string()))) -} -#[cfg(test)] -mod tests { - use super::super::test_helpers::*; - use super::*; - #[test] - fn pkce_s256_challenge_matches_verifier() { - let pkce = generate_pkce(); - assert_eq!(pkce.code_verifier.len(), 43); - let expected = URL_SAFE_NO_PAD.encode(Sha256::digest(pkce.code_verifier.as_bytes())); - assert_eq!(pkce.code_challenge, expected); - } - #[test] - fn authorize_url_includes_required_oidc_params() { - let config = OidcAuthConfig { - issuer: "https://example.okta.com".into(), - client_id: TEST_CLIENT_ID.into(), - scopes: vec!["openid".into(), "profile".into()], - audience: Some("api://grok".into()), - }; - let discovery = Discovery { - authorization_endpoint: "https://example.okta.com/authorize".into(), - token_endpoint: "https://example.okta.com/token".into(), - jwks_uri: None, - id_token_signing_alg_values_supported: None, - }; - let pkce = Pkce { - code_verifier: "v".into(), - code_challenge: "c".into(), - }; - let url = build_authorize_url( - &config, - None, - &discovery, - "http://127.0.0.1:9999/callback", - &pkce, - "state123", - "nonce123", - ); - for required in [ - "response_type=code", - "client_id=test-client-id", - "code_challenge=c", - "code_challenge_method=S256", - "state=state123", - "nonce=nonce123", - "scope=openid", - "audience=api", - "referrer=grok-build", - ] { - assert!(url.contains(required), "missing param: {required}"); - } - assert_eq!( - url.matches("referrer=").count(), - 1, - "expected exactly one referrer param, got: {url}" - ); - } - #[test] - fn authorize_url_includes_team_principal_params() { - let config = OidcAuthConfig { - issuer: "https://auth.x.ai".into(), - client_id: TEST_CLIENT_ID.into(), - scopes: vec!["offline_access".into(), "grok-cli:access".into()], - audience: None, - }; - let oauth2 = OAuth2ProviderConfig { - issuer: "https://auth.x.ai".into(), - client_id: TEST_CLIENT_ID.into(), - scopes: vec!["offline_access".into(), "grok-cli:access".into()], - principal_type: Some("Team".into()), - principal_id: Some("team-123".into()), - referrer: Some("grok-build".into()), - }; - let discovery = Discovery { - authorization_endpoint: "https://auth.x.ai/authorize".into(), - token_endpoint: "https://auth.x.ai/token".into(), - jwks_uri: None, - id_token_signing_alg_values_supported: None, - }; - let pkce = Pkce { - code_verifier: "v".into(), - code_challenge: "c".into(), - }; - let url = build_authorize_url( - &config, - Some(&oauth2), - &discovery, - "http://127.0.0.1:9999/callback", - &pkce, - "state123", - "nonce123", - ); - assert!(url.contains("principal_type=Team")); - assert!(url.contains("principal_id=team-123")); - assert!(url.contains("referrer=grok-build")); - assert_eq!( - url.matches("referrer=").count(), - 1, - "expected exactly one referrer param, got: {url}" - ); - } - #[test] - fn authorize_url_uses_oauth2_referrer_override_once() { - let config = OidcAuthConfig { - issuer: "https://auth.x.ai".into(), - client_id: TEST_CLIENT_ID.into(), - scopes: vec!["offline_access".into(), "grok-cli:access".into()], - audience: None, - }; - let oauth2 = OAuth2ProviderConfig { - issuer: "https://auth.x.ai".into(), - client_id: TEST_CLIENT_ID.into(), - scopes: vec!["offline_access".into(), "grok-cli:access".into()], - principal_type: None, - principal_id: None, - referrer: Some("grok-desktop".into()), - }; - let discovery = Discovery { - authorization_endpoint: "https://auth.x.ai/authorize".into(), - token_endpoint: "https://auth.x.ai/token".into(), - jwks_uri: None, - id_token_signing_alg_values_supported: None, - }; - let pkce = Pkce { - code_verifier: "v".into(), - code_challenge: "c".into(), - }; - let url = build_authorize_url( - &config, - Some(&oauth2), - &discovery, - "http://127.0.0.1:9999/callback", - &pkce, - "state123", - "nonce123", - ); - assert!(url.contains("referrer=grok-desktop")); - assert!(!url.contains("referrer=grok-build")); - assert_eq!( - url.matches("referrer=").count(), - 1, - "expected exactly one referrer param, got: {url}" - ); - } - #[tokio::test] - async fn extract_user_info_allows_team_without_id_token() { - let discovery = Discovery { - authorization_endpoint: "https://example.okta.com/authorize".into(), - token_endpoint: "https://example.okta.com/token".into(), - jwks_uri: Some("https://example.okta.com/jwks".into()), - id_token_signing_alg_values_supported: Some(vec!["RS256".into()]), - }; - let user_info = extract_user_info( - None, - &discovery, - "https://example.okta.com", - "test-client", - "nonce123", - Some("Team"), - Some("team-123"), - None, - ) - .await - .expect("team login should not require id_token"); - assert_eq!( - user_info.user_id, "team-123", - "team user_id should be the principal_id" - ); - assert_eq!(user_info.principal_type.as_deref(), Some("Team")); - assert_eq!(user_info.principal_id.as_deref(), Some("team-123")); - assert_eq!(user_info.team_id.as_deref(), Some("team-123")); - assert!(user_info.email.is_none()); - } - #[test] - fn validate_state_rejects_mismatch() { - let err = validate_state("expected-state", "wrong-state").unwrap_err(); - assert!( - err.to_string().contains("state mismatch"), - "unexpected error: {err}" - ); - } - #[tokio::test] - async fn id_token_validation_fails_on_nonce_mismatch() { - ensure_crypto_provider(); - let (issuer, id_token, discovery, handle) = mock_idp_token().await; - let err = extract_user_info( - Some(&id_token), - &discovery, - &issuer, - TEST_CLIENT_ID, - "wrong-nonce", - None, - None, - None, - ) - .await - .unwrap_err(); - assert!( - err.to_string().contains("nonce mismatch"), - "unexpected error: {err}" - ); - handle.abort(); - } - #[test] - fn rejects_unsupported_id_token_alg() { - let err = ensure_alg_allowed(jsonwebtoken::Algorithm::HS256, Some(&["RS256".to_string()])) - .unwrap_err(); - assert!( - err.to_string().contains("unsupported algorithm"), - "unexpected error: {err}" - ); - } - #[tokio::test] - async fn id_token_validation_fails_on_audience_mismatch() { - ensure_crypto_provider(); - let (issuer, id_token, discovery, handle) = mock_idp_token().await; - let err = extract_user_info( - Some(&id_token), - &discovery, - &issuer, - "wrong-client", - TEST_NONCE, - None, - None, - None, - ) - .await - .unwrap_err(); - assert!( - err.to_string().contains("audience mismatch") - || err.to_string().contains("InvalidAudience"), - "unexpected error: {err}" - ); - handle.abort(); - } - /// JWT principal-extraction matrix: - /// - team JWT: extracts (Team, team_id, None) - /// - non-JWT garbage / empty: returns None - /// - JWT without principal_type/_id: returns None - #[test] - fn peek_access_token_principal_matrix() { - ensure_crypto_provider(); - fn make_jwt(claims: serde_json::Value) -> String { - let header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::HS256); - jsonwebtoken::encode( - &header, - &claims, - &jsonwebtoken::EncodingKey::from_secret(b"test-secret"), - ) - .unwrap() - } - let team_jwt = make_jwt(serde_json::json!( - { "sub" : "user-42", "iss" : "https://auth.x.ai", "aud" : "test-client", - "exp" : 9999999999u64, "iat" : 1000000000u64, "scope" : - "offline_access grok-cli:access api:access", "principal_type" : "Team", - "principal_id" : "team-abc-123", "client_id" : "test-client", "jti" : - "token-1", } - )); - let (pt, pid, tid) = peek_access_token_principal(&team_jwt).expect("team principal"); - assert_eq!(pt, "Team"); - assert_eq!(pid, "team-abc-123"); - assert_eq!(tid, None); - assert!(peek_access_token_principal("not-a-jwt-token").is_none()); - assert!(peek_access_token_principal("").is_none()); - let no_principal = make_jwt(serde_json::json!( - { "sub" : "user-42", "iss" : "https://auth.x.ai", "aud" : "test-client", - "exp" : 9999999999u64, "iat" : 1000000000u64, } - )); - assert!(peek_access_token_principal(&no_principal).is_none()); - } - /// `peek_access_token_principal_id` extracts the id even when - /// `principal_type` is absent, where the stricter - /// `peek_access_token_principal` returns `None`. - #[test] - fn peek_access_token_principal_id_does_not_require_type() { - ensure_crypto_provider(); - fn make_jwt(claims: serde_json::Value) -> String { - jsonwebtoken::encode( - &jsonwebtoken::Header::new(jsonwebtoken::Algorithm::HS256), - &claims, - &jsonwebtoken::EncodingKey::from_secret(b"test-secret"), - ) - .unwrap() - } - let id_only = make_jwt(serde_json::json!({ "principal_id" : "team-abc", "sub" : "u" })); - assert_eq!( - peek_access_token_principal_id(&id_only).as_deref(), - Some("team-abc"), - ); - assert!( - peek_access_token_principal(&id_only).is_none(), - "the strict peek still needs principal_type", - ); - let none = make_jwt(serde_json::json!({ "sub" : "u" })); - assert!(peek_access_token_principal_id(&none).is_none()); - assert!(peek_access_token_principal_id("not-a-jwt").is_none()); - } - /// Enforcement matrix: None passes; Single/AnyOf require a match (and reject - /// a no-principal token); empty AnyOf fails closed. - #[test] - fn enforce_login_principal_matrix() { - assert!(enforce_login_principal(None, None).is_ok()); - assert!(enforce_login_principal(None, Some("team-abc")).is_ok()); - let single = ForceLoginTeam::Single("team-abc".into()); - assert!(enforce_login_principal(Some(&single), Some("team-abc")).is_ok()); - let err = enforce_login_principal(Some(&single), Some("team-other")).unwrap_err(); - assert_eq!( - err.to_string(), - "This deployment requires logging into team team-abc; \ - your login returned team-other", - ); - let err = enforce_login_principal(Some(&single), None).unwrap_err(); - assert_eq!( - err.to_string(), - "This deployment requires logging into team team-abc; \ - your login returned no team principal", - ); - let any_of = ForceLoginTeam::AnyOf(vec!["team-a".into(), "team-b".into()]); - assert!(enforce_login_principal(Some(&any_of), Some("team-b")).is_ok()); - let err = enforce_login_principal(Some(&any_of), Some("team-c")).unwrap_err(); - assert_eq!( - err.to_string(), - "This deployment requires logging into one of teams: team-a, team-b; \ - your login returned team-c", - ); - let err = enforce_login_principal(Some(&ForceLoginTeam::AnyOf(vec![])), Some("team-a")) - .unwrap_err(); - assert_eq!( - err.to_string(), - "Login is blocked by your administrator: force_login_team_uuid is an empty \ - list, so no team is permitted to sign in", - ); - } - /// Only the dedicated `force_login_team_uuid` knob produces an enforcement - /// policy; the legacy `oauth2.principal_id` is pre-select-only and never an - /// enforcement gate (regression guard for the upgrade-behavior concern). - #[test] - fn resolve_login_principal_policy_uses_force_login_team_only() { - assert_eq!(resolve_login_principal_policy(None), None); - assert_eq!( - resolve_login_principal_policy(Some(&ForceLoginTeam::Single("team-locked".into()))), - Some(ForceLoginTeam::Single("team-locked".into())), - ); - assert_eq!( - resolve_login_principal_policy(Some(&ForceLoginTeam::AnyOf(vec![ - "a".into(), - "b".into() - ]))), - Some(ForceLoginTeam::AnyOf(vec!["a".into(), "b".into()])), - ); - } - /// Discovery is cached for `DISCOVERY_CACHE_TTL`: the second call - /// to `discover()` for the same issuer hits the cache and does not - /// fetch over HTTP. Without this, every refresh pays a discovery - /// round-trip and a discovery-endpoint blip blocks token refresh. - #[tokio::test] - async fn discover_uses_cache_within_ttl() { - use std::sync::atomic::{AtomicU32, Ordering}; - clear_discovery_cache(); - let hits = std::sync::Arc::new(AtomicU32::new(0)); - let hits_for_handler = hits.clone(); - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let issuer = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let issuer_for_handler = issuer.clone(); - let app = axum::Router::new().route( - "/.well-known/openid-configuration", - axum::routing::get(move || { - let b = issuer_for_handler.clone(); - let counter = hits_for_handler.clone(); - async move { - counter.fetch_add(1, Ordering::SeqCst); - axum::Json(serde_json::json!( - { "authorization_endpoint" : format!("{b}/authorize"), - "token_endpoint" : format!("{b}/token"), } - )) - } - }), - ); - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - let _ = discover(&issuer).await.unwrap(); - let _ = discover(&issuer).await.unwrap(); - let _ = discover(&issuer).await.unwrap(); - assert_eq!( - hits.load(Ordering::SeqCst), - 1, - "discover() must hit the network exactly once across 3 calls (cache TTL = 1h)" - ); - server.abort(); - } - /// `refresh_tokens` retries on a transient 503 and succeeds on the - /// next attempt. Without backon, a single IdP blip during refresh - /// surfaces to the user as a chat failure. - #[tokio::test] - async fn refresh_tokens_retries_on_transient_5xx() { - use std::sync::atomic::{AtomicU32, Ordering}; - let hits = std::sync::Arc::new(AtomicU32::new(0)); - let hits_for_handler = hits.clone(); - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let app = axum::Router::new().route( - "/token", - axum::routing::post(move || { - let counter = hits_for_handler.clone(); - async move { - let n = counter.fetch_add(1, Ordering::SeqCst) + 1; - if n == 1 { - ( - axum::http::StatusCode::SERVICE_UNAVAILABLE, - "upstream busy".to_string(), - ) - } else { - ( - axum::http::StatusCode::OK, - r#"{"access_token":"new-at","expires_in":3600}"#.to_string(), - ) - } - } - }), - ); - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - let token_endpoint = format!("http://127.0.0.1:{port}/token"); - let resp = refresh_tokens(&token_endpoint, "rt", "client", None, None) - .await - .expect("transient 5xx must be retried until success"); - assert_eq!(resp.access_token, "new-at"); - assert_eq!( - hits.load(Ordering::SeqCst), - 2, - "first attempt fails 503, second succeeds — exactly 2 hits" - ); - server.abort(); - } - /// Terminal OAuth2 errors (`invalid_grant`, `invalid_client`) MUST - /// NOT be retried -- retrying a revoked grant just wastes time and - /// risks rate-limit. Verifies `is_transient_refresh_error` correctly - /// classifies typed 4xx as terminal. - #[tokio::test] - async fn refresh_tokens_does_not_retry_terminal_invalid_grant() { - use std::sync::atomic::{AtomicU32, Ordering}; - let hits = std::sync::Arc::new(AtomicU32::new(0)); - let hits_for_handler = hits.clone(); - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let app = axum::Router::new().route( - "/token", - axum::routing::post(move || { - let counter = hits_for_handler.clone(); - async move { - counter.fetch_add(1, Ordering::SeqCst); - ( - axum::http::StatusCode::BAD_REQUEST, - r#"{"error":"invalid_grant","error_description":"refresh token revoked"}"# - .to_string(), - ) - } - }), - ); - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - let token_endpoint = format!("http://127.0.0.1:{port}/token"); - let err = refresh_tokens(&token_endpoint, "rt", "client", None, None) - .await - .expect_err("invalid_grant is terminal"); - assert!( - err.to_string().contains("400") || err.to_string().contains("invalid_grant"), - "error must surface the IdP rejection, got: {err}" - ); - assert_eq!( - hits.load(Ordering::SeqCst), - 1, - "terminal OAuth2 error must NOT be retried (exactly 1 hit)" - ); - server.abort(); - } - /// A 4xx carrying an OAuth2 code that is NOT a recognized terminal one - /// (e.g. RFC 6749 `temporarily_unavailable`) must be retried, not given up - /// on. The retry gate defers to `classify_terminal`, so only the recognized - /// terminal codes stop retries; everything else is transient. - #[tokio::test] - async fn refresh_tokens_retries_on_coded_transient_error() { - use std::sync::atomic::{AtomicU32, Ordering}; - let hits = std::sync::Arc::new(AtomicU32::new(0)); - let hits_for_handler = hits.clone(); - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let app = axum::Router::new().route( - "/token", - axum::routing::post(move || { - let counter = hits_for_handler.clone(); - async move { - let n = counter.fetch_add(1, Ordering::SeqCst) + 1; - if n == 1 { - ( - axum::http::StatusCode::BAD_REQUEST, - r#"{"error":"temporarily_unavailable"}"#.to_string(), - ) - } else { - ( - axum::http::StatusCode::OK, - r#"{"access_token":"new-at","expires_in":3600}"#.to_string(), - ) - } - } - }), - ); - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - let token_endpoint = format!("http://127.0.0.1:{port}/token"); - let resp = refresh_tokens(&token_endpoint, "rt", "client", None, None) - .await - .expect("a non-terminal coded 4xx must be retried until success"); - assert_eq!(resp.access_token, "new-at"); - assert_eq!( - hits.load(Ordering::SeqCst), - 2, - "temporarily_unavailable must be retried (1 fail + 1 success = 2 hits)" - ); - server.abort(); - } - #[test] - fn callback_timeout_error_is_user_friendly() { - let err: anyhow::Error = OidcError::CallbackTimeout.into(); - let msg = err.to_string(); - assert!( - msg.contains("Login timed out after 10 minutes"), - "expected friendly timeout message, got: {msg}" - ); - assert!( - msg.contains("Please try again"), - "expected 'Please try again' call to action, got: {msg}" - ); - assert!( - !msg.contains("OIDC"), - "should not leak internal 'OIDC' terminology to users, got: {msg}" - ); - assert!( - !msg.contains("300s"), - "should not mention raw seconds, got: {msg}" - ); - } -} diff --git a/crates/codegen/kigi-shell/src/auth/oidc/refresh.rs b/crates/codegen/kigi-shell/src/auth/oidc/refresh.rs deleted file mode 100644 index 3378611..0000000 --- a/crates/codegen/kigi-shell/src/auth/oidc/refresh.rs +++ /dev/null @@ -1,249 +0,0 @@ -//! Pure-data OIDC refresh. Talks to the IdP and returns -//! [`OidcRefreshResult`] without touching [`AuthManager`]. - -use super::super::GrokAuth; -use super::protocol::{OidcError, OidcUserInfo, build_grok_auth, discover, refresh_tokens}; -use crate::auth::error::RefreshTokenFailedReason; - -/// Outcome of a pure OIDC token refresh (no AuthManager mutations). -pub(crate) enum OidcRefreshResult { - /// Fresh token obtained. Caller must persist. - Success(Box), - /// Terminal error from the IdP, already classified into a reason. - TerminalError { reason: RefreshTokenFailedReason }, - /// Non-terminal failure (discovery failed, network error, etc.) - Failed, -} - -/// Classify an OAuth2 `error` code as a terminal refresh failure. `None` means -/// non-terminal (retryable). Single source of truth for which codes are fatal; -/// the retry gate (`protocol::is_transient_refresh_error`) defers to this too. -pub(super) fn classify_terminal(error_code: &str) -> Option { - match error_code { - "invalid_grant" => Some(RefreshTokenFailedReason::RefreshTokenRejected), - "invalid_client" => Some(RefreshTokenFailedReason::ClientRejected), - _ => None, - } -} - -/// `oauth2-provider` refresh-token rotation-grace window (ms). Only a clock -/// divergence past this bound is flagged as a suspected suspend-straddle, since -/// a longer suspend can turn a lost refresh response into a revoked RT. -const ROTATION_GRACE_MS: u64 = 60_000; - -/// Exchange a refresh_token for fresh tokens at the IdP. Pure data return, no -/// `AuthManager` mutations; the caller (`OidcRefresher`) routes the result -/// through `refresh_chain`. -pub(crate) async fn oidc_token_exchange(auth: &GrokAuth) -> OidcRefreshResult { - let has_rt = auth.refresh_token.is_some(); - let has_issuer = auth.oidc_issuer.is_some(); - let has_client_id = auth.oidc_client_id.is_some(); - tracing::debug!( - has_rt, - has_issuer, - has_client_id, - "oidc try_refresh_pure enter" - ); - if !has_rt || !has_issuer || !has_client_id { - kigi_log::unified_log::warn( - "oidc try_refresh skipped: missing fields", - None, - Some(serde_json::json!({ - "has_refresh_token": has_rt, - "has_issuer": has_issuer, - "has_client_id": has_client_id, - "auth_mode": format!("{:?}", auth.auth_mode), - })), - ); - } - let Some(refresh_tok) = auth.refresh_token.as_ref() else { - return OidcRefreshResult::Failed; - }; - let Some(issuer) = auth.oidc_issuer.as_ref() else { - return OidcRefreshResult::Failed; - }; - let Some(client_id) = auth.oidc_client_id.as_ref() else { - return OidcRefreshResult::Failed; - }; - - crate::unified_log::info( - "oidc try_refresh_pure enter", - None, - Some(serde_json::json!({ "issuer": issuer, "client_id": client_id })), - ); - - // Suspend probe: the monotonic clock pauses while the machine is asleep - // but the wall clock does not, so a large divergence around the IdP call - // means the process was suspended mid-refresh — the exact condition that - // can revoke the refresh token (response lost across sleep). - let started_mono = std::time::Instant::now(); - let started_wall = chrono::Utc::now(); - let timing = || { - let mono_ms = started_mono.elapsed().as_millis() as u64; - let wall_ms = (chrono::Utc::now() - started_wall) - .num_milliseconds() - .max(0) as u64; - let suspended_ms = wall_ms.saturating_sub(mono_ms); - ( - mono_ms, - wall_ms, - suspended_ms, - suspended_ms > ROTATION_GRACE_MS, - ) - }; - - let discovery = match discover(issuer).await { - Ok(d) => d, - Err(e) => { - let (mono_ms, wall_ms, suspended_ms, suspected_suspend) = timing(); - crate::unified_log::error( - "oidc try_refresh_pure discovery failed", - None, - Some(serde_json::json!({ - "error": format!("{e:#}"), - "mono_ms": mono_ms, - "wall_ms": wall_ms, - "suspended_ms": suspended_ms, - "suspected_suspend": suspected_suspend, - })), - ); - if suspected_suspend { - emit_suspend_spanned("discovery_failed", suspended_ms); - } - return OidcRefreshResult::Failed; - } - }; - let tokens = match refresh_tokens( - &discovery.token_endpoint, - refresh_tok, - client_id, - auth.principal_type.as_deref(), - auth.principal_id.as_deref(), - ) - .await - { - Ok(t) => t, - Err(e) => { - if let Some(OidcError::TokenRefreshHttp { body, .. }) = e.downcast_ref::() - && let Some(error_code) = serde_json::from_str::(body) - .ok() - .and_then(|v| v.get("error")?.as_str().map(str::to_owned)) - && let Some(reason) = classify_terminal(&error_code) - { - let (mono_ms, wall_ms, suspended_ms, suspected_suspend) = timing(); - let cred_age_secs = auth.mint_age_seconds(); - crate::unified_log::error( - "oidc try_refresh_pure terminal error", - None, - Some(serde_json::json!({ - "error_code": error_code, - "client_id": client_id, - "tried_rt_prefix": auth.refresh_token.as_deref().map(crate::auth::token_suffix), - "error_description": serde_json::from_str::(body) - .ok() - .and_then(|v| v.get("error_description").cloned()), - "mono_ms": mono_ms, - "wall_ms": wall_ms, - "suspended_ms": suspended_ms, - "suspected_suspend": suspected_suspend, - "cred_age_secs": cred_age_secs, - })), - ); - if suspected_suspend { - emit_suspend_spanned(&error_code, suspended_ms); - } - return OidcRefreshResult::TerminalError { reason }; - } - let http_status = e.downcast_ref::().and_then(|oe| match oe { - OidcError::TokenRefreshHttp { status, .. } => Some(*status), - _ => None, - }); - let (mono_ms, wall_ms, suspended_ms, suspected_suspend) = timing(); - crate::unified_log::error( - "oidc try_refresh_pure token exchange failed", - None, - Some(serde_json::json!({ - "error": e.to_string(), - "client_id": client_id, - "http_status": http_status, - "mono_ms": mono_ms, - "wall_ms": wall_ms, - "suspended_ms": suspended_ms, - "suspected_suspend": suspected_suspend, - })), - ); - tracing::warn!( - error = %e, - http_status = ?http_status, - client_id = %client_id, - issuer = %issuer, - "OIDC: token refresh failed" - ); - if suspected_suspend { - emit_suspend_spanned("transient_failed", suspended_ms); - } - return OidcRefreshResult::Failed; - } - }; - - // Reuse identity from original login; new id_token from refresh is intentionally skipped. - let user_info = OidcUserInfo { - user_id: auth.user_id.clone(), - email: auth.email.clone(), - first_name: auth.first_name.clone(), - last_name: auth.last_name.clone(), - profile_image_asset_id: auth.profile_image_asset_id.clone(), - principal_type: auth.principal_type.clone(), - principal_id: auth.principal_id.clone(), - team_id: auth.team_id.clone(), - team_name: auth.team_name.clone(), - team_role: auth.team_role.clone(), - organization_id: auth.organization_id.clone(), - organization_name: auth.organization_name.clone(), - organization_role: auth.organization_role.clone(), - user_blocked_reason: auth.user_blocked_reason.clone(), - team_blocked_reasons: auth.team_blocked_reasons.clone(), - coding_data_retention_opt_out: auth.coding_data_retention_opt_out, - }; - let mut new_auth = build_grok_auth(tokens, user_info, issuer, client_id); - let idp_rotated = new_auth.refresh_token.is_some(); - // Keep old refresh token if IdP didn't rotate it - if new_auth.refresh_token.is_none() { - new_auth.refresh_token = auth.refresh_token.clone(); - } - tracing::debug!( - idp_rotated, - key_prefix = crate::auth::token_suffix(&new_auth.key), - "oidc try_refresh_pure token obtained" - ); - let (mono_ms, wall_ms, suspended_ms, suspected_suspend) = timing(); - crate::unified_log::info( - "oidc try_refresh_pure succeeded", - None, - Some(serde_json::json!({ - "expires_at": new_auth.expires_at.map(|e| e.to_rfc3339()), - "mono_ms": mono_ms, - "wall_ms": wall_ms, - "suspended_ms": suspended_ms, - "suspected_suspend": suspected_suspend, - })), - ); - if suspected_suspend { - emit_suspend_spanned("ok", suspended_ms); - } - OidcRefreshResult::Success(Box::new(new_auth)) -} - -/// Alertable event: an OIDC refresh's network call spanned a suspend (wall -/// clock ran far ahead of the monotonic clock) — the precondition for a -/// lost-response refresh-token revocation. -fn emit_suspend_spanned(outcome: &str, suspended_ms: u64) { - crate::unified_log::warn( - "auth.refresh.suspend_spanned", - None, - Some(serde_json::json!({ - "outcome": outcome, - "suspended_ms": suspended_ms, - })), - ); -} diff --git a/crates/codegen/kigi-shell/src/auth/oidc/test_helpers.rs b/crates/codegen/kigi-shell/src/auth/oidc/test_helpers.rs deleted file mode 100644 index 09d26e4..0000000 --- a/crates/codegen/kigi-shell/src/auth/oidc/test_helpers.rs +++ /dev/null @@ -1,130 +0,0 @@ -//! Shared test helpers for `oidc::protocol::tests` and `oidc::login::tests`. -//! Both test modules need a mock IdP server (`start_mock_idp`), JWT -//! signing primitives (`generate_test_rsa_key`, `mock_idp_token`), and -//! the same constants. Extracted here so neither test mod has to -//! re-implement them. - -use base64::Engine; -use base64::engine::general_purpose::URL_SAFE_NO_PAD; - -use super::protocol::{Discovery, discover}; - -pub(super) const TEST_KID: &str = "test-kid"; -pub(super) const TEST_NONCE: &str = "test-nonce-value"; -pub(super) const TEST_CLIENT_ID: &str = "test-client-id"; -pub(super) fn ensure_crypto_provider() { - let _ = rustls::crypto::ring::default_provider().install_default(); - let _ = jsonwebtoken::crypto::rust_crypto::DEFAULT_PROVIDER.install_default(); -} -pub(super) fn generate_test_rsa_key() -> (String, String, String) { - use rsa::pkcs8::EncodePrivateKey; - use rsa::traits::PublicKeyParts; - let private_key = rsa::RsaPrivateKey::new(&mut rsa::rand_core::OsRng, 2048).unwrap(); - let pem = private_key - .to_pkcs8_pem(rsa::pkcs8::LineEnding::LF) - .unwrap() - .to_string(); - let jwk_n = URL_SAFE_NO_PAD.encode(private_key.n().to_bytes_be()); - let jwk_e = URL_SAFE_NO_PAD.encode(private_key.e().to_bytes_be()); - (pem, jwk_n, jwk_e) -} -pub(super) async fn mock_idp_token() -> (String, String, Discovery, tokio::task::JoinHandle<()>) { - let (issuer, handle) = start_mock_idp().await; - let discovery = discover(&issuer).await.unwrap(); - let resp: serde_json::Value = crate::http::shared_client() - .post(&discovery.token_endpoint) - .form(&[("grant_type", "authorization_code")]) - .send() - .await - .unwrap() - .json() - .await - .unwrap(); - let id_token = resp["id_token"] - .as_str() - .expect("mock missing id_token") - .to_string(); - (issuer, id_token, discovery, handle) -} -pub(super) async fn start_mock_idp() -> (String, tokio::task::JoinHandle<()>) { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let issuer = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let issuer_for_discovery = issuer.clone(); - let (rsa_pem, jwk_n, jwk_e) = generate_test_rsa_key(); - - #[derive(serde::Serialize)] - struct Claims { - sub: &'static str, - email: &'static str, - iss: String, - aud: &'static str, - nonce: &'static str, - exp: usize, - } - - let id_token = { - let mut hdr = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256); - hdr.kid = Some(TEST_KID.to_owned()); - jsonwebtoken::encode( - &hdr, - &Claims { - sub: "user-42", - email: "test@corp.com", - iss: issuer.clone(), - aud: TEST_CLIENT_ID, - nonce: TEST_NONCE, - exp: (chrono::Utc::now() + chrono::Duration::hours(1)).timestamp() as usize, - }, - &jsonwebtoken::EncodingKey::from_rsa_pem(rsa_pem.as_bytes()).unwrap(), - ) - .unwrap() - }; - - let app = axum::Router::new() - .route( - "/.well-known/openid-configuration", - axum::routing::get(move || { - let iss = issuer_for_discovery.clone(); - async move { - axum::Json(serde_json::json!({ - "authorization_endpoint": format!("{iss}/authorize"), - "token_endpoint": format!("{iss}/token"), - "jwks_uri": format!("{iss}/jwks"), - "id_token_signing_alg_values_supported": ["RS256"], - })) - } - }), - ) - .route( - "/jwks", - axum::routing::get(move || { - let n = jwk_n.clone(); - let e = jwk_e.clone(); - async move { - axum::Json(serde_json::json!({ - "keys": [{ - "kty": "RSA", "alg": "RS256", "kid": TEST_KID, - "n": n, "e": e, - }] - })) - } - }), - ) - .route( - "/token", - axum::routing::post(move || { - let tok = id_token.clone(); - async move { - axum::Json(serde_json::json!({ - "access_token": "mock-access-token", - "refresh_token": "mock-refresh-token", - "id_token": tok, - "expires_in": 3600, - })) - } - }), - ); - - let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - (issuer, handle) -} diff --git a/crates/codegen/kigi-shell/src/auth/recovery.rs b/crates/codegen/kigi-shell/src/auth/recovery.rs index cb6baae..311e783 100644 --- a/crates/codegen/kigi-shell/src/auth/recovery.rs +++ b/crates/codegen/kigi-shell/src/auth/recovery.rs @@ -3,71 +3,54 @@ //! When the server rejects a token, `UnauthorizedRecovery` walks through //! a sequence of recovery steps before giving up: //! -//! 1. **ReloadFromDisk** — re-read `auth.json` under a file lock; if the -//! on-disk token differs from the rejected one, accept it (another -//! process may have refreshed). -//! 2. **RefreshFromAuthority** — run the appropriate refresh chain -//! (OIDC token refresh, external binary, etc.) based on `TokenType`, -//! unless the live token was minted moments ago (fresh-mint guard). -//! 3. **DevboxRecovery** — on devboxes, purge `auth.json` and mint fresh -//! OIDC credentials. -//! 4. **Done** — all recovery strategies exhausted. +//! 1. **ReloadFromDisk** — re-read the persisted credential under a file +//! lock; if it differs from the rejected one, accept it (another process +//! may have refreshed). +//! 2. **RefreshFromAuthority** — run the refresh chain against the Kimi +//! OAuth host, unless the live token was minted moments ago (fresh-mint +//! guard). +//! 3. **Done** — all recovery strategies exhausted. use std::sync::Arc; use crate::auth::error::{AuthError, RefreshTokenError, RefreshTokenFailedReason}; use crate::auth::manager::AuthManager; -use crate::auth::model::GrokAuth; +use crate::auth::model::KimiAuth; use crate::auth::token_type::TokenType; -/// Whether a terminal `AuthError` forces a manual re-login (`None` cases +/// Whether a terminal `AuthError` forces a manual re-login (`false` cases /// self-heal or are transient). Lives here (not on `AuthError`) so the error /// model stays free of recovery policy. pub(crate) fn forces_manual_reauth(err: &AuthError) -> bool { match err { AuthError::Refresh(RefreshTokenError::Permanent(e)) => match e.reason { RefreshTokenFailedReason::RefreshTokenRejected => true, - // Self-healing via the TTL, not a manual re-auth. - RefreshTokenFailedReason::ClientRejected | RefreshTokenFailedReason::Other => false, + // Self-healing via the tombstone cooldown, not a manual re-auth. + RefreshTokenFailedReason::Other => false, }, AuthError::ServerRejectedNoRecovery | AuthError::RecoveryExhausted - | AuthError::TokenExpiredNoRefresh - | AuthError::PinnedTeamMismatch { .. } => true, - // API-key lockouts are out of scope: an admin disabling API-key auth - // means rotate the key, not `/login`. - AuthError::ApiKeyAuthDisabled - | AuthError::Refresh(RefreshTokenError::Transient(_)) - | AuthError::NotLoggedIn => false, + | AuthError::TokenExpiredNoRefresh => true, + AuthError::Refresh(RefreshTokenError::Transient(_)) | AuthError::NotLoggedIn => false, } } -/// Whether the relay should stop reconnecting on this recovery error. Its own -/// predicate rather than reusing `forces_manual_reauth`: the relay must give up -/// on any terminal auth failure, including `ApiKeyAuthDisabled` (a kill-switched -/// API key), which deliberately doesn't force a manual re-login. -pub(crate) fn relay_should_cancel(err: &AuthError) -> bool { - forces_manual_reauth(err) || matches!(err, AuthError::ApiKeyAuthDisabled) -} - /// Fresh-mint guard window (±) for `ServerRejected` refreshes /// ([`UnauthorizedRecovery::fresh_mint_guard`]). 120s outlasts in-flight /// requests sent with a previous key plus validation lag (observed stale -/// 401s land ~20s after mint), while `current()`'s 300s early-invalidation -/// buffer keeps any guard-returned token wire-valid. A genuinely-dead fresh -/// token waits at most this long to re-mint; the symmetric bound caps that -/// delay when the clock stepped back. +/// 401s land ~20s after mint), while the refresh-threshold buffer keeps any +/// guard-returned token wire-valid. A genuinely-dead fresh token waits at +/// most this long to re-mint; the symmetric bound caps that delay when the +/// clock stepped back. const FRESH_MINT_GUARD_SECS: i64 = 120; /// Which recovery step to attempt next. #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum RecoveryStep { - /// Re-read auth.json from disk (file-locked). + /// Re-read the persisted credential (file-locked). ReloadFromDisk, - /// Refresh via the authority (OIDC, external binary, etc.). + /// Refresh via the Kimi OAuth host. RefreshFromAuthority, - /// On devboxes: purge auth.json and mint fresh OIDC credentials. - DevboxRecovery, /// All strategies exhausted. Done, } @@ -79,8 +62,7 @@ pub struct UnauthorizedRecovery { rejected_token: String, /// Current step in the recovery sequence. step: RecoveryStep, - /// Error from `RefreshFromAuthority`, propagated as fallback when - /// devbox recovery doesn't apply. + /// Error from `RefreshFromAuthority`, propagated on exhaustion. authority_error: Option, /// Whether the last authority failure was transient. Kept past the /// `authority_error` handoff so exhaustion preserves the @@ -90,7 +72,7 @@ pub struct UnauthorizedRecovery { impl UnauthorizedRecovery { /// `rejected` is the credential the server rejected: its key drives recovery. - pub(crate) fn new(auth_manager: Arc, rejected: Option) -> Self { + pub(crate) fn new(auth_manager: Arc, rejected: Option) -> Self { let rejected_token = rejected.as_ref().map(|a| a.key.clone()).unwrap_or_default(); Self { auth_manager, @@ -102,14 +84,14 @@ impl UnauthorizedRecovery { } /// Attempt the next recovery step. Walks - /// `ReloadFromDisk -> RefreshFromAuthority -> DevboxRecovery -> Done`. + /// `ReloadFromDisk -> RefreshFromAuthority -> Done`. /// `token_type` span field is recorded lazily via /// `Span::is_disabled()` to avoid the lock when tracing is off. #[tracing::instrument( skip(self), fields(step = ?self.step, token_type = tracing::field::Empty), )] - pub async fn next(&mut self) -> Result { + pub async fn next(&mut self) -> Result { let span = tracing::Span::current(); if !span.is_disabled() { // Only acquire the inner-lock when tracing actually @@ -122,23 +104,10 @@ impl UnauthorizedRecovery { tracing::field::debug(self.auth_manager.token_type()), ); } - self.resolve_next().await + self.next_step_loop().await } - /// Walk the recovery steps and apply the team-pin policy gate. - async fn resolve_next(&mut self) -> Result { - // Team-pin gate: 401 recovery must not resurrect a wrong-team session - // (disk adoption / refresh / devbox mint) for the relay to reconnect - // with. Clear + reject on mismatch. - let auth = self.next_step_loop().await?; - if let Some(e) = self.auth_manager.cached_token_policy_error(&auth) { - self.auth_manager.reject_and_clear(&e); - return Err(e); - } - Ok(auth) - } - - async fn next_step_loop(&mut self) -> Result { + async fn next_step_loop(&mut self) -> Result { loop { match self.step { RecoveryStep::ReloadFromDisk => { @@ -148,30 +117,20 @@ impl UnauthorizedRecovery { } } RecoveryStep::RefreshFromAuthority => { - self.step = RecoveryStep::DevboxRecovery; + self.step = RecoveryStep::Done; match self.try_refresh_from_authority().await { Ok(auth) => return Ok(auth), Err(e) => { self.authority_was_transient = matches!(e, AuthError::Refresh(RefreshTokenError::Transient(_))); self.authority_error = Some(e); + return Err(self + .authority_error + .take() + .unwrap_or(AuthError::RecoveryExhausted)); } } } - RecoveryStep::DevboxRecovery => { - self.step = RecoveryStep::Done; - // preferred_method=api_key forbids automatic OIDC mint. - if !self.auth_manager.grok_com_config().blocks_automatic_oidc() - && self.auth_manager.is_devbox_environment() - && let Ok(auth) = self.auth_manager.try_devbox_recovery().await - { - return Ok(auth); - } - return Err(self - .authority_error - .take() - .unwrap_or(AuthError::RecoveryExhausted)); - } RecoveryStep::Done => { // Exhaustion after a *transient* authority failure stays // transient: `RecoveryExhausted` here would count a network @@ -187,9 +146,9 @@ impl UnauthorizedRecovery { } } - /// Re-read `auth.json` from disk. Accept the token only if it differs + /// Re-read the persisted credential. Accept the token only if it differs /// from the one that was rejected. - async fn try_reload_from_disk(&self) -> Option { + async fn try_reload_from_disk(&self) -> Option { let _lock = self .auth_manager .try_lock_auth_file_async(crate::auth::manager::AUTH_LOCK_TIMEOUT) @@ -203,13 +162,13 @@ impl UnauthorizedRecovery { // same-as-rejected / no entry): a silent arm hides which path // a recovery loop is taking. Debug level — the disk-state // *transition* is logged once by `read_disk_auth` itself. - kigi_log::unified_log::debug("auth recovery: no disk entry", None, None); + kigi_log::unified_log::debug("auth recovery: no persisted entry", None, None); return None; }; if crate::auth::is_expired(&disk_auth) { - tracing::debug!("auth recovery: disk token is expired, skipping"); + tracing::debug!("auth recovery: persisted token is expired, skipping"); kigi_log::unified_log::debug( - "auth recovery: disk token expired", + "auth recovery: persisted token expired", None, Some(serde_json::json!({ "disk_key_prefix": crate::auth::token_suffix(&disk_auth.key), @@ -219,9 +178,9 @@ impl UnauthorizedRecovery { return None; } if self.is_different_token(&disk_auth) { - tracing::info!("auth recovery: disk has a different token, accepting"); + tracing::info!("auth recovery: persisted store has a different token, accepting"); kigi_log::unified_log::info( - "auth recovery: adopted disk token", + "auth recovery: adopted persisted token", None, Some(serde_json::json!({ "adopted_key_prefix": crate::auth::token_suffix(&disk_auth.key), @@ -231,8 +190,12 @@ impl UnauthorizedRecovery { self.auth_manager.hot_swap(disk_auth.clone()); Some(disk_auth) } else { - tracing::debug!("auth recovery: disk token is same as rejected, skipping"); - kigi_log::unified_log::debug("auth recovery: disk token same as rejected", None, None); + tracing::debug!("auth recovery: persisted token is same as rejected, skipping"); + kigi_log::unified_log::debug( + "auth recovery: persisted token same as rejected", + None, + None, + ); None } } @@ -242,14 +205,12 @@ impl UnauthorizedRecovery { /// clock that stepped far back) falls through to a normal refresh. /// /// A 401 moments after a successful mint is a stale rejection (sent with - /// the previous key and mis-attributed — see `is_stale_snapshot`) or - /// validation lag on the new key — re-minting fixes neither, and a crash - /// between the IdP grant and persisting the response orphans the - /// replacement RT (forced re-login). Consumers retry with the returned - /// token; a genuinely-bad one refreshes once the window passes. Lives - /// here, not in `refresh_chain`, so paywall claims re-mints that call - /// `refresh_chain(ServerRejected)` directly are unaffected. - fn fresh_mint_guard(&self) -> Option { + /// the previous key) or validation lag on the new key — re-minting fixes + /// neither, and a crash between the token grant and persisting the + /// response orphans the replacement RT (forced re-login). Consumers retry + /// with the returned token; a genuinely-bad one refreshes once the window + /// passes. + fn fresh_mint_guard(&self) -> Option { let auth = self.auth_manager.current()?; let mint_age_seconds = auth.mint_age_seconds(); if !(-FRESH_MINT_GUARD_SECS..FRESH_MINT_GUARD_SECS).contains(&mint_age_seconds) { @@ -272,14 +233,14 @@ impl UnauthorizedRecovery { Some(auth) } - /// Dispatch to the correct refresh chain based on the current `TokenType`. + /// Dispatch to the refresh chain based on the current `TokenType`. /// /// Per-variant outcome: /// - /// - **OidcSession / ExternalBinary**: full refresh chain via the - /// authority, unless the live token is inside the fresh-mint guard - /// window ([`Self::fresh_mint_guard`]). - /// - **LegacySession / ApiKey**: no refresh authority for these + /// - **OAuthSession**: full refresh chain via the OAuth host, unless the + /// live token is inside the fresh-mint guard window + /// ([`Self::fresh_mint_guard`]). + /// - **SessionNoRefresh / ApiKey**: no refresh authority for these /// types. We've already tried `ReloadFromDisk` (the previous /// recovery step), so the server's 401 stands. Surface /// [`AuthError::ServerRejectedNoRecovery`] -- *not* @@ -289,10 +250,10 @@ impl UnauthorizedRecovery { /// reading the variant can distinguish "ran past local TTL" from /// "server actively rejected". /// - **None**: no credentials at all. - async fn try_refresh_from_authority(&self) -> Result { + async fn try_refresh_from_authority(&self) -> Result { let tt = self.auth_manager.token_type(); match tt { - TokenType::OidcSession | TokenType::ExternalBinary => { + TokenType::OAuthSession => { if let Some(auth) = self.fresh_mint_guard() { return Ok(auth); } @@ -325,7 +286,7 @@ impl UnauthorizedRecovery { } result } - TokenType::LegacySession | TokenType::ApiKey => { + TokenType::SessionNoRefresh | TokenType::ApiKey => { kigi_log::unified_log::warn( "auth recovery: no refresh authority for token type", None, @@ -338,7 +299,7 @@ impl UnauthorizedRecovery { } /// Check if a candidate token is different from the rejected one. - fn is_different_token(&self, candidate: &GrokAuth) -> bool { + fn is_different_token(&self, candidate: &KimiAuth) -> bool { candidate.key != self.rejected_token } } @@ -348,29 +309,28 @@ mod tests { //! State-machine matrix tests for `UnauthorizedRecovery`. //! //! Coverage targets: - //! - All 5 `TokenType` variants x dispatch in `try_refresh_from_authority`. + //! - All 4 `TokenType` variants x dispatch in `try_refresh_from_authority`. //! - `try_reload_from_disk`: same/different/no token on disk. //! - `next()` exhaustion (Done -> RecoveryExhausted). - //! - Fresh-mint guard: ±window bounds, ExternalBinary, verdict grace, - //! policy-hidden fall-through (fail closed). + //! - Fresh-mint guard: ±window bounds, tombstone grace. //! //! These tests use the same in-process `AuthManager` that production //! does and inject a counting refresher so we can observe whether the //! authority was consulted. use super::*; - use crate::auth::config::GrokComConfig; - use crate::auth::error::{RefreshTokenError, RefreshTokenFailedReason}; - use crate::auth::model::{AuthMode, GrokAuth}; + use crate::auth::config::KimiCodeConfig; + use crate::auth::error::RefreshTokenError; + use crate::auth::model::{AuthMode, KimiAuth}; use crate::auth::refresh::{RefreshOutcome, TokenRefresher}; use crate::auth::storage::{read_auth_json, write_auth_json}; use chrono::{Duration, Utc}; use std::sync::atomic::{AtomicU32, Ordering}; /// The rejected wire bearer these tests seed into the manager. - fn rejected_cred() -> Option { - Some(GrokAuth { + fn rejected_cred() -> Option { + Some(KimiAuth { key: "rejected-tok".into(), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }) } @@ -382,17 +342,17 @@ mod tests { impl TokenRefresher for OkRefresher { async fn refresh(&self, _reason: crate::auth::manager::RefreshReason) -> RefreshOutcome { self.calls.fetch_add(1, Ordering::SeqCst); - RefreshOutcome::Success(Box::new(GrokAuth { + RefreshOutcome::Success(Box::new(KimiAuth { key: "fresh-from-authority".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, refresh_token: Some("rt-new".into()), expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() })) } } - /// Refresher fake: returns PermanentFailure (invalid_grant). + /// Refresher fake: returns PermanentFailure (rejected refresh token). struct FailRefresher { calls: Arc, } @@ -406,19 +366,19 @@ mod tests { fn mgr() -> (tempfile::TempDir, Arc) { let dir = tempfile::tempdir().unwrap(); - let m = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); + let m = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); (dir, m) } fn seed(mgr: &AuthManager, mode: AuthMode, refresh_token: Option<&str>) { - let auth = GrokAuth { + let auth = KimiAuth { key: "rejected-tok".into(), auth_mode: mode, refresh_token: refresh_token.map(str::to_string), // Past expiry so `current()` returns None and the refresh // chain actually has to do work. expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }; mgr.hot_swap(auth); } @@ -426,9 +386,9 @@ mod tests { // -- TokenType dispatch matrix ---------------------------------------- #[tokio::test] - async fn dispatch_oidc_session_uses_refresh_chain() { + async fn dispatch_oauth_session_uses_refresh_chain() { let (_d, m) = mgr(); - seed(&m, AuthMode::Oidc, Some("rt")); + seed(&m, AuthMode::OAuth, Some("rt")); let calls = Arc::new(AtomicU32::new(0)); m.set_refresher(Arc::new(OkRefresher { calls: calls.clone(), @@ -441,39 +401,25 @@ mod tests { assert_eq!(calls.load(Ordering::SeqCst), 1); } - #[tokio::test] - async fn dispatch_external_binary_uses_refresh_chain() { - let (_d, m) = mgr(); - seed(&m, AuthMode::External, None); - let calls = Arc::new(AtomicU32::new(0)); - m.set_refresher(Arc::new(OkRefresher { - calls: calls.clone(), - })); - - let mut rec = m.unauthorized_recovery(rejected_cred()); - let auth = rec.next().await.expect("external-binary recovery succeeds"); - assert_eq!(auth.key, "fresh-from-authority"); - assert_eq!(calls.load(Ordering::SeqCst), 1); - } - // -- Fresh-mint guard -------------------------------------------------- /// Seed a *valid* (unexpired) in-memory token whose `create_time` lies /// `mint_age` in the past (negative = clock stepped back since mint). - fn seed_valid(mgr: &AuthManager, mode: AuthMode, mint_age: Duration) { - mgr.hot_swap(GrokAuth { + fn seed_valid(mgr: &AuthManager, mint_age: Duration) { + mgr.hot_swap(KimiAuth { key: "rejected-tok".into(), - auth_mode: mode, + auth_mode: AuthMode::OAuth, refresh_token: Some("rt".into()), create_time: Utc::now() - mint_age, expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() + expires_in: Some(3600), + ..KimiAuth::test_default() }); } /// Run one recovery against a counting refresher; return the outcome and /// how many times the authority was consulted. - async fn recover_with_ok_refresher(m: &Arc) -> (Result, u32) { + async fn recover_with_ok_refresher(m: &Arc) -> (Result, u32) { let calls = Arc::new(AtomicU32::new(0)); m.set_refresher(Arc::new(OkRefresher { calls: calls.clone(), @@ -484,9 +430,9 @@ mod tests { } #[tokio::test] - async fn fresh_mint_guard_skips_idp_for_freshly_minted_token() { + async fn fresh_mint_guard_skips_wire_for_freshly_minted_token() { let (_d, m) = mgr(); - seed_valid(&m, AuthMode::Oidc, Duration::seconds(10)); + seed_valid(&m, Duration::seconds(10)); let (result, calls) = recover_with_ok_refresher(&m).await; assert_eq!( result.expect("guard returns the live token").key, @@ -495,23 +441,11 @@ mod tests { assert_eq!(calls, 0, "a 10s-old token must not be re-minted"); } - #[tokio::test] - async fn fresh_mint_guard_applies_to_external_binary_tokens() { - let (_d, m) = mgr(); - seed_valid(&m, AuthMode::External, Duration::seconds(10)); - let (result, calls) = recover_with_ok_refresher(&m).await; - assert_eq!( - result.expect("guard returns the live token").key, - "rejected-tok" - ); - assert_eq!(calls, 0); - } - #[tokio::test] async fn fresh_mint_guard_treats_small_negative_age_as_fresh() { // Clock stepped back slightly since mint (NTP nudge). let (_d, m) = mgr(); - seed_valid(&m, AuthMode::Oidc, Duration::seconds(-60)); + seed_valid(&m, Duration::seconds(-60)); let (result, calls) = recover_with_ok_refresher(&m).await; assert_eq!( result.expect("guard returns the live token").key, @@ -525,19 +459,19 @@ mod tests { // A large backwards clock step must not wedge recovery for the whole // step: outside the ±window the guard stands down. let (_d, m) = mgr(); - seed_valid(&m, AuthMode::Oidc, Duration::hours(-1)); + seed_valid(&m, Duration::hours(-1)); let (result, calls) = recover_with_ok_refresher(&m).await; assert_eq!( result.expect("recovery should succeed").key, "fresh-from-authority" ); - assert_eq!(calls, 1, "far-negative mint age must reach the IdP"); + assert_eq!(calls, 1, "far-negative mint age must reach the wire"); } #[tokio::test] async fn fresh_mint_guard_lets_old_token_refresh() { let (_d, m) = mgr(); - seed_valid(&m, AuthMode::Oidc, Duration::minutes(10)); + seed_valid(&m, Duration::minutes(10)); let (result, calls) = recover_with_ok_refresher(&m).await; assert_eq!( result.expect("recovery should succeed").key, @@ -545,25 +479,25 @@ mod tests { ); assert_eq!( calls, 1, - "outside the guard window ServerRejected must reach the IdP" + "outside the guard window ServerRejected must reach the wire" ); } #[tokio::test] - async fn fresh_mint_guard_wins_over_cached_permanent_failure() { - // A fresh *valid* token is served even when a permanent-failure - // verdict is cached for it — mirrors `auth()`'s wire-valid grace arm; - // the verdict re-applies once the guard window passes. + async fn fresh_mint_guard_wins_over_cached_tombstone() { + // A fresh *valid* token is served even when a tombstone is cached + // for its refresh token — mirrors `auth()`'s wire-valid grace arm; + // the tombstone re-applies once the guard window passes. let (_d, m) = mgr(); - seed_valid(&m, AuthMode::Oidc, Duration::seconds(10)); + seed_valid(&m, Duration::seconds(10)); m.record_permanent_failure( - "rejected-tok".into(), + "rt".into(), RefreshTokenFailedReason::RefreshTokenRejected.into(), ); let (result, calls) = recover_with_ok_refresher(&m).await; assert_eq!( result - .expect("guard precedes the verdict short-circuit") + .expect("guard precedes the tombstone short-circuit") .key, "rejected-tok" ); @@ -571,65 +505,10 @@ mod tests { } #[tokio::test] - async fn fresh_mint_guard_never_returns_policy_hidden_token() { - // Wrong-team fresh token: `current()` hides it (vet_cached), so the - // guard must fall through to a normal refresh — fail closed. - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig { - force_login_team_uuid: Some(crate::auth::config::ForceLoginTeam::Single( - "team-good".into(), - )), - ..GrokComConfig::default() - }; - let m = Arc::new(AuthManager::new(dir.path(), cfg)); - m.hot_swap(GrokAuth { - key: team_jwt("team-wrong"), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt".into()), - create_time: Utc::now(), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() - }); - let calls = Arc::new(AtomicU32::new(0)); - m.set_refresher(Arc::new(OkRefresher { - calls: calls.clone(), - })); - - let mut rec = m.unauthorized_recovery(rejected_cred()); - let result = rec.next().await; - assert_eq!( - calls.load(Ordering::SeqCst), - 1, - "hidden token must not satisfy the guard" - ); - if let Ok(auth) = result { - assert_ne!( - auth.key, - team_jwt("team-wrong"), - "wrong-team token must never be returned" - ); - } - } - - #[tokio::test] - async fn dispatch_legacy_session_returns_server_rejected_no_recovery() { + async fn dispatch_session_without_refresh_token_returns_server_rejected_no_recovery() { + // OAuth without refresh_token classifies as SessionNoRefresh. let (_d, m) = mgr(); - // WebLogin (no refresh_token) -> LegacySession. - seed(&m, AuthMode::WebLogin, None); - - let mut rec = m.unauthorized_recovery(rejected_cred()); - let err = rec.next().await.unwrap_err(); - assert!( - matches!(err, AuthError::ServerRejectedNoRecovery), - "LegacySession recovery should surface ServerRejectedNoRecovery, got {err:?}", - ); - } - - #[tokio::test] - async fn dispatch_oidc_without_refresh_token_returns_server_rejected_no_recovery() { - // Oidc without refresh_token classifies as LegacySession. - let (_d, m) = mgr(); - seed(&m, AuthMode::Oidc, None); + seed(&m, AuthMode::OAuth, None); let mut rec = m.unauthorized_recovery(rejected_cred()); let err = rec.next().await.unwrap_err(); @@ -668,16 +547,17 @@ mod tests { #[tokio::test] async fn reload_from_disk_picks_up_different_token() { let (dir, m) = mgr(); - seed(&m, AuthMode::Oidc, Some("rt")); + seed(&m, AuthMode::OAuth, Some("rt")); // Sibling process wrote a different valid token to disk. - let scope = m.grok_com_config().auth_scope(); - let fresh = GrokAuth { + let scope = m.kimi_code_config().auth_scope(); + let fresh = KimiAuth { key: "fresh-from-disk".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, refresh_token: Some("rt-new".into()), expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() + expires_in: Some(3600), + ..KimiAuth::test_default() }; let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default(); store.insert(scope, fresh); @@ -694,16 +574,17 @@ mod tests { #[tokio::test] async fn reload_from_disk_skips_same_token_then_proceeds_to_authority() { let (dir, m) = mgr(); - seed(&m, AuthMode::Oidc, Some("rt")); + seed(&m, AuthMode::OAuth, Some("rt")); // Disk has the SAME token that was rejected -- skip, fall through. - let scope = m.grok_com_config().auth_scope(); - let same = GrokAuth { + let scope = m.kimi_code_config().auth_scope(); + let same = KimiAuth { key: "rejected-tok".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, refresh_token: Some("rt".into()), expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() + expires_in: Some(3600), + ..KimiAuth::test_default() }; let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default(); store.insert(scope, same); @@ -732,15 +613,11 @@ mod tests { #[tokio::test] async fn next_after_done_returns_recovery_exhausted() { let (_d, m) = mgr(); - seed(&m, AuthMode::Oidc, Some("rt")); + seed(&m, AuthMode::OAuth, Some("rt")); m.set_refresher(Arc::new(OkRefresher { calls: Arc::new(AtomicU32::new(0)), })); - // Pin non-devbox so DevboxRecovery can't adopt the seeded token (CI runs - // in K8s pods where is_devbox_environment() is true). - m.set_devbox_env_for_test(false); - let mut rec = m.unauthorized_recovery(rejected_cred()); let _ = rec.next().await.unwrap(); let err = loop { @@ -773,9 +650,8 @@ mod tests { } let (_d, m) = mgr(); - seed(&m, AuthMode::Oidc, Some("rt")); + seed(&m, AuthMode::OAuth, Some("rt")); m.set_refresher(Arc::new(TransientFailRefresher)); - m.set_devbox_env_for_test(false); let mut rec = m.unauthorized_recovery(rejected_cred()); // First next(): the authority's transient error propagates as-is. @@ -799,21 +675,17 @@ mod tests { !forces_manual_reauth(&err), "a transient exhaustion must not force a manual re-login", ); - assert!( - !relay_should_cancel(&err), - "the relay must reconnect (not cancel) on a transient exhaustion", - ); } - // -- Permanent failure short-circuit (cross-check) ------------ + // -- Tombstone short-circuit (cross-check) ------------ #[tokio::test] - async fn refresh_authority_short_circuits_on_cached_permanent_failure() { + async fn refresh_authority_short_circuits_on_cached_tombstone() { let (_d, m) = mgr(); - seed(&m, AuthMode::Oidc, Some("rt")); - // Pre-record a permanent failure scoped to the seeded credential. + seed(&m, AuthMode::OAuth, Some("rt")); + // Pre-record a tombstone scoped to the seeded refresh token. m.record_permanent_failure( - "rejected-tok".into(), + "rt".into(), RefreshTokenFailedReason::RefreshTokenRejected.into(), ); @@ -831,7 +703,7 @@ mod tests { assert_eq!( calls.load(Ordering::SeqCst), 0, - "refresher must not be invoked when permanent_failure is cached", + "refresher must not be invoked while the tombstone cooldown is live", ); } @@ -843,15 +715,15 @@ mod tests { #[tokio::test] async fn reload_from_disk_rejects_expired_different_token() { let (dir, m) = mgr(); - seed(&m, AuthMode::Oidc, Some("rt")); + seed(&m, AuthMode::OAuth, Some("rt")); - let scope = m.grok_com_config().auth_scope(); - let expired_different = GrokAuth { + let scope = m.kimi_code_config().auth_scope(); + let expired_different = KimiAuth { key: "different-but-expired".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, refresh_token: Some("rt-new".into()), expires_at: Some(Utc::now() - Duration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }; let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default(); store.insert(scope, expired_different); @@ -870,64 +742,4 @@ mod tests { ); assert_eq!(calls.load(Ordering::SeqCst), 1); } - - // -- force_login_team_uuid pin enforced on the 401-recovery path ------- - - fn ensure_crypto_provider() { - let _ = jsonwebtoken::crypto::rust_crypto::DEFAULT_PROVIDER.install_default(); - } - - fn team_jwt(principal_id: &str) -> String { - ensure_crypto_provider(); - jsonwebtoken::encode( - &jsonwebtoken::Header::new(jsonwebtoken::Algorithm::HS256), - &serde_json::json!({ - "sub": "user-1", - "principal_type": "Team", - "principal_id": principal_id, - "exp": 9999999999u64, - }), - &jsonwebtoken::EncodingKey::from_secret(b"test-secret"), - ) - .unwrap() - } - - /// A sibling writes a wrong-team token to disk; 401 recovery (relay path) - /// must reject + clear it at `next()`, not hand it back as a bearer. - #[tokio::test] - async fn recovery_rejects_wrong_team_adopted_disk_token() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig { - force_login_team_uuid: Some(crate::auth::config::ForceLoginTeam::Single( - "team-good".into(), - )), - ..GrokComConfig::default() - }; - let scope = cfg.auth_scope(); - let m = Arc::new(AuthManager::new(dir.path(), cfg)); - - // In-memory: the rejected (expired) session that triggered recovery. - seed(&m, AuthMode::Oidc, Some("rt")); - - // Disk: a different, non-expired, *wrong-team* token a sibling wrote. - let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default(); - store.insert( - scope, - GrokAuth { - key: team_jwt("team-wrong"), - auth_mode: AuthMode::Oidc, - refresh_token: Some("rt-sibling".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..GrokAuth::test_default() - }, - ); - write_auth_json(&dir.path().join("auth.json"), &store).unwrap(); - - let mut rec = m.unauthorized_recovery(rejected_cred()); - let err = rec.next().await.unwrap_err(); - assert!( - matches!(err, AuthError::PinnedTeamMismatch { .. }), - "recovery must reject a wrong-team disk token, got {err:?}" - ); - } } diff --git a/crates/codegen/kigi-shell/src/auth/refresh/auth_backend_contract_tests.rs b/crates/codegen/kigi-shell/src/auth/refresh/auth_backend_contract_tests.rs deleted file mode 100644 index 587098c..0000000 --- a/crates/codegen/kigi-shell/src/auth/refresh/auth_backend_contract_tests.rs +++ /dev/null @@ -1,351 +0,0 @@ -//! End-to-end auth-backend contract tests: a mock IdP whose `/token` response -//! is forced per case, asserting the refresh outcome, the storm cap, and the -//! terminal-error classification on the live recovery path. - -use super::*; -use crate::auth::error::RefreshTokenFailedReason; -use crate::auth::{GrokAuth, GrokComConfig}; -use chrono::{Duration, Utc}; -use std::sync::Arc; -use std::sync::atomic::{AtomicU32, Ordering}; - -/// Mock IdP: OIDC discovery + a `/token` endpoint returning a fixed -/// `(status, body)` and counting every hit, plus the `/user` endpoint -/// `AuthManager::update` calls after a successful refresh. `delay_ms` widens -/// the in-lock window so concurrent callers queue on `refresh_lock`. -async fn start_idp( - token_status: u16, - token_body: String, - hits: Arc, - delay_ms: u64, -) -> (String, tokio::task::JoinHandle<()>) { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let disco = base.clone(); - - let app = axum::Router::new() - .route( - "/.well-known/openid-configuration", - axum::routing::get(move || { - let b = disco.clone(); - async move { - axum::Json(serde_json::json!({ - "authorization_endpoint": format!("{b}/authorize"), - "token_endpoint": format!("{b}/token"), - })) - } - }), - ) - .route( - "/token", - axum::routing::post(move || { - let hits = hits.clone(); - let body = token_body.clone(); - async move { - hits.fetch_add(1, Ordering::SeqCst); - if delay_ms > 0 { - tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; - } - ( - axum::http::StatusCode::from_u16(token_status).unwrap(), - body, - ) - } - }), - ) - .route( - "/user", - axum::routing::get(|| async { - axum::Json(serde_json::json!({ "userId": "user-42", "email": "u@corp.com" })) - }), - ); - - let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - (base, handle) -} - -fn expired_oidc(base_url: &str) -> GrokAuth { - GrokAuth { - key: "expired-at".into(), - create_time: Utc::now() - Duration::hours(2), - user_id: "user-42".into(), - auth_mode: crate::auth::model::AuthMode::Oidc, - refresh_token: Some("rt-under-test".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - oidc_issuer: Some(base_url.to_owned()), - oidc_client_id: Some("client-under-test".into()), - ..GrokAuth::test_default() - } -} - -#[derive(Debug)] -enum Expect { - Success, - Permanent(RefreshTokenFailedReason), - Transient, -} - -/// The IdP token-endpoint contract: each response shape maps to one outcome. -/// `invalid_grant`/`invalid_client` are the only permanent verdicts; status -/// blips and unrecognized codes stay transient (never permanent-lock). -#[tokio::test] -async fn auth_backend_contract_token_responses_map_to_outcomes() { - use RefreshTokenFailedReason::{ClientRejected, RefreshTokenRejected}; - let cases: &[(&str, u16, &str, Expect)] = &[ - ( - "success", - 200, - r#"{"access_token":"fresh","refresh_token":"fresh-rt","expires_in":3600}"#, - Expect::Success, - ), - ( - "invalid_grant", - 400, - r#"{"error":"invalid_grant"}"#, - Expect::Permanent(RefreshTokenRejected), - ), - ( - "invalid_client", - 401, - r#"{"error":"invalid_client"}"#, - Expect::Permanent(ClientRejected), - ), - ("server_error_5xx", 503, "{}", Expect::Transient), - ("rate_limited_429", 429, "{}", Expect::Transient), - ( - "temporarily_unavailable", - 400, - r#"{"error":"temporarily_unavailable"}"#, - Expect::Transient, - ), - ("bare_4xx_no_body", 400, "", Expect::Transient), - ("malformed_body", 400, "not json", Expect::Transient), - // Proxy/WAF-mangled bodies must degrade to retry, never a false permanent - // lock: a nested error object or a non-string `error` is not a recognized - // top-level code, so it stays transient. - ( - "nested_error_object", - 400, - r#"{"error":{"code":"invalid_grant"}}"#, - Expect::Transient, - ), - ( - "non_string_error", - 400, - r#"{"error":123}"#, - Expect::Transient, - ), - ]; - - for (name, status, body, expect) in cases { - let hits = Arc::new(AtomicU32::new(0)); - let (base_url, server) = start_idp(*status, body.to_string(), hits.clone(), 0).await; - let dir = tempfile::tempdir().unwrap(); - let auth_manager = Arc::new( - AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url), - ); - auth_manager.hot_swap(expired_oidc(&base_url)); - - let refresher = OidcRefresher::new(auth_manager.clone()); - let result = refresher.refresh(RefreshReason::ServerRejected).await; - - match (expect, &result) { - (Expect::Success, RefreshOutcome::Success(_)) => {} - (Expect::Permanent(want), RefreshOutcome::PermanentFailure { error, .. }) => { - assert_eq!(error.reason, *want, "{name}: wrong permanent reason"); - } - (Expect::Transient, RefreshOutcome::TransientFailure { .. }) => {} - (exp, got) => panic!("{name}: expected {exp:?}, got {got:?}"), - } - server.abort(); - } -} - -/// A burst of concurrent 401s on the same revoked refresh token must hit the -/// IdP exactly once. The callers serialize on `refresh_lock`; the leader records -/// the verdict before releasing, so the in-lock re-check (`refresh_chain` step -/// 1b) short-circuits every follower. Delete step 1b and the count climbs to N. -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn auth_backend_contract_concurrent_401s_hit_idp_once() { - let hits = Arc::new(AtomicU32::new(0)); - // 100ms /token delay so every caller passes the pre-lock check and queues - // on refresh_lock before the leader records the verdict, exercising step 1b. - let (base_url, server) = start_idp( - 400, - r#"{"error":"invalid_grant"}"#.to_string(), - hits.clone(), - 100, - ) - .await; - let dir = tempfile::tempdir().unwrap(); - let auth_manager = Arc::new( - AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url), - ); - auth_manager.hot_swap(expired_oidc(&base_url)); - auth_manager.set_refresher(Arc::new(OidcRefresher::new(auth_manager.clone()))); - - let mut tasks = Vec::new(); - for _ in 0..6 { - let auth_manager = auth_manager.clone(); - tasks.push(tokio::spawn(async move { auth_manager.auth().await })); - } - for t in tasks { - let outcome = t.await.unwrap(); - assert!( - matches!( - outcome, - Err(crate::auth::AuthError::Refresh( - crate::auth::RefreshTokenError::Permanent(_) - )) - ), - "every concurrent caller must fail permanently on a revoked refresh token, got {outcome:?}", - ); - } - - assert_eq!( - hits.load(Ordering::SeqCst), - 1, - "concurrent 401s on one dead credential must hit the IdP exactly once", - ); - - server.abort(); -} - -/// The classification loop through the live recovery state machine: a dead -/// refresh token terminates recovery with an error that forces a manual -/// re-login; a refreshable token auto-refreshes. -#[tokio::test] -async fn auth_backend_contract_dead_token_forces_manual_reauth() { - // A dead refresh token terminates recovery with a forced-relogin error. - let hits = Arc::new(AtomicU32::new(0)); - let (url, server) = start_idp(400, r#"{"error":"invalid_grant"}"#.to_string(), hits, 0).await; - let dir = tempfile::tempdir().unwrap(); - let auth_manager = - Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&url)); - auth_manager.hot_swap(expired_oidc(&url)); - auth_manager.set_refresher(Arc::new(OidcRefresher::new(auth_manager.clone()))); - - let err = auth_manager - .unauthorized_recovery(auth_manager.current_or_expired()) - .next() - .await - .expect_err("a dead refresh token must fail recovery"); - assert!( - crate::auth::recovery::forces_manual_reauth(&err), - "a dead refresh token must be a forced-relogin error, got {err:?}", - ); - server.abort(); - - // Refreshable token: recovery auto-refreshes. - let ok_hits = Arc::new(AtomicU32::new(0)); - let (ok_url, ok_server) = start_idp( - 200, - r#"{"access_token":"fresh","refresh_token":"fresh-rt","expires_in":3600}"#.to_string(), - ok_hits, - 0, - ) - .await; - let ok_dir = tempfile::tempdir().unwrap(); - let ok_manager = Arc::new( - AuthManager::new(ok_dir.path(), GrokComConfig::default()).with_proxy_base_url(&ok_url), - ); - ok_manager.hot_swap(expired_oidc(&ok_url)); - ok_manager.set_refresher(Arc::new(OidcRefresher::new(ok_manager.clone()))); - - let refreshed = ok_manager - .unauthorized_recovery(ok_manager.current_or_expired()) - .next() - .await - .expect("a refreshable token must auto-refresh"); - assert_eq!( - refreshed.key, "fresh", - "recovery must return the fresh token" - ); - ok_server.abort(); -} - -/// Consecutive transient failures self-heal up to a bound, then escalate to a -/// non-sticky `Other` permanent failure (which ages out via the TTL). A -/// regression here would turn recoverable blips into a permanent `/login`. -#[tokio::test] -async fn auth_backend_contract_transient_failures_escalate_to_non_sticky_permanent() { - let hits = Arc::new(AtomicU32::new(0)); - // Persistent 503: every refresh attempt is transient. - let (base_url, server) = start_idp(503, "{}".to_string(), hits, 0).await; - let dir = tempfile::tempdir().unwrap(); - let auth_manager = Arc::new( - AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url), - ); - auth_manager.hot_swap(expired_oidc(&base_url)); - - // One refresher instance: it owns the consecutive-failure counter. - let refresher = OidcRefresher::new(auth_manager.clone()); - let mut outcomes = Vec::new(); - for _ in 0..3 { - outcomes.push(refresher.refresh(RefreshReason::ServerRejected).await); - } - - assert!( - matches!(outcomes[0], RefreshOutcome::TransientFailure { .. }), - "first blip is transient, not a lockout: {:?}", - outcomes[0], - ); - match &outcomes[2] { - RefreshOutcome::PermanentFailure { error, .. } => { - assert_eq!( - error.reason, - RefreshTokenFailedReason::Other, - "escalation must use the generic Other reason", - ); - assert!( - !error.reason.is_sticky(), - "an escalated transient must age out, not strand the user forever", - ); - } - other => panic!("repeated transients must escalate to a permanent Other, got {other:?}"), - } - - server.abort(); -} - -/// Two `AuthManager`s sharing one auth.json stand in for two CLI processes: the -/// auth.json flock must serialize their refreshes so the shared refresh token is -/// spent at the IdP exactly once. The loser adopts the rotated token from disk -/// instead of racing a second exchange (which the IdP could revoke as reuse). -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn auth_backend_contract_two_instances_share_one_idp_call() { - let hits = Arc::new(AtomicU32::new(0)); - let (url, server) = start_idp( - 200, - r#"{"access_token":"fresh","refresh_token":"fresh-rt","expires_in":3600}"#.to_string(), - hits.clone(), - 100, - ) - .await; - let dir = tempfile::tempdir().unwrap(); - - // Distinct managers, same on-disk auth.json (separate flock OFDs => they - // genuinely contend, like two processes). - let new_instance = || { - let m = Arc::new( - AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&url), - ); - m.hot_swap(expired_oidc(&url)); - m.set_refresher(Arc::new(OidcRefresher::new(m.clone()))); - m - }; - let a = new_instance(); - let b = new_instance(); - - let (ra, rb) = tokio::join!(a.auth(), b.auth()); - - assert_eq!(ra.expect("instance A must obtain a token").key, "fresh"); - assert_eq!(rb.expect("instance B must obtain a token").key, "fresh"); - assert_eq!( - hits.load(Ordering::SeqCst), - 1, - "two instances sharing auth.json must spend the refresh token at the IdP only once", - ); - - server.abort(); -} diff --git a/crates/codegen/kigi-shell/src/auth/refresh/external_refresher.rs b/crates/codegen/kigi-shell/src/auth/refresh/external_refresher.rs deleted file mode 100644 index 4457845..0000000 --- a/crates/codegen/kigi-shell/src/auth/refresh/external_refresher.rs +++ /dev/null @@ -1,176 +0,0 @@ -use std::sync::Arc; - -use crate::auth::error::RefreshTokenFailedReason; -use crate::auth::manager::RefreshReason; - -use super::{ExternalCommandRunner, RefreshOutcome, TokenRefresher}; - -/// Refreshes by re-running the operator's external auth binary via -/// `spawn_blocking`. Pure data return -- mutation lives in -/// `refresh_chain` (honors the [`TokenRefresher`] no-mutation contract). -pub(crate) struct ExternalBinaryRefresher { - runner: Arc, - command: String, - timeout: std::time::Duration, -} - -impl ExternalBinaryRefresher { - pub(crate) fn new(runner: Arc, command: String) -> Self { - Self { - runner, - command, - timeout: EXTERNAL_REFRESH_TIMEOUT, - } - } - - /// Override the binary timeout (tests use a short one to exercise the - /// timeout arm without a real 30s wait). - #[cfg(test)] - pub(crate) fn with_timeout(mut self, timeout: std::time::Duration) -> Self { - self.timeout = timeout; - self - } - - /// A failed binary run is a single-strike `Other` permanent failure; the - /// `PERMANENT_FAILURE_TTL` lets a flaky binary self-heal without `/login`. - /// No consecutive-blip tolerance like OIDC: a local binary failure is a - /// stronger signal than a network refresh blip. - fn record_failure(&self, message: String) -> RefreshOutcome { - tracing::warn!(%message, "auth: external binary refresh failed -> permanent"); - // No token key in the binary flow; the caller scopes the verdict. - RefreshOutcome::permanent(RefreshTokenFailedReason::Other, None) - } -} - -/// Timeout for the external auth binary. If the binary hangs, the -/// `spawn_blocking` thread is leaked (it cannot be interrupted), but this is -/// acceptable: the thread holds no locks and mutates no shared state. -const EXTERNAL_REFRESH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); - -#[async_trait::async_trait] -impl TokenRefresher for ExternalBinaryRefresher { - async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome { - tracing::debug!(?reason, "auth: external binary refresh starting"); - let runner = self.runner.clone(); - let cmd = self.command.clone(); - let timeout_ms = self.timeout.as_millis() as u64; - match tokio::time::timeout( - self.timeout, - tokio::task::spawn_blocking(move || runner.run_external_command(&cmd)), - ) - .await - { - Err(_elapsed) => { - tracing::warn!( - timeout_ms, - "auth: external binary refresh timed out (thread leaked)" - ); - crate::unified_log::warn( - "auth.refresh.external_timeout", - None, - Some(serde_json::json!({ "timeout_ms": timeout_ms })), - ); - self.record_failure(format!("external binary timed out after {timeout_ms}ms")) - } - Ok(Ok(Some(auth))) => { - crate::unified_log::info("auth: external binary refresh succeeded", None, None); - RefreshOutcome::success(auth) - } - Ok(Ok(None)) => { - crate::unified_log::warn( - "auth: external binary refresh returned no token", - None, - None, - ); - self.record_failure("external binary returned no token".into()) - } - Ok(Err(e)) => { - tracing::warn!(error = %e, "auth: external binary refresh task failed"); - self.record_failure(format!("external binary task failed: {e}")) - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::auth::GrokAuth; - - /// Minimal runner whose external command returns a fixed result. - struct FakeRunner { - external_result: Option, - } - impl ExternalCommandRunner for FakeRunner { - fn run_external_command(&self, _command: &str) -> Option { - self.external_result.clone() - } - } - - /// A failed binary run is a single-strike `Other` permanent failure that is - /// NON-sticky: it must age out via the TTL, never lock an external-binary - /// user out forever. (Flipping this to a sticky reason would be a silent - /// lockout regression.) - #[tokio::test] - async fn external_binary_failure_is_single_strike_non_sticky_permanent() { - let refresher = ExternalBinaryRefresher::new( - Arc::new(FakeRunner { - external_result: None, - }), - "auth-binary".into(), - ); - match refresher.refresh(RefreshReason::ServerRejected).await { - RefreshOutcome::PermanentFailure { error, .. } => { - assert_eq!(error.reason, RefreshTokenFailedReason::Other); - assert!( - !error.reason.is_sticky(), - "external-binary failure must age out, not strand the user forever", - ); - } - other => panic!("a failed binary run must be a permanent Other failure, got {other:?}"), - } - } - - /// A binary that outlives the (test-shortened) timeout hits the `Elapsed` - /// arm and maps to the same non-sticky `Other` permanent failure. - #[tokio::test] - async fn external_binary_timeout_is_non_sticky_permanent() { - struct SlowRunner; - impl ExternalCommandRunner for SlowRunner { - fn run_external_command(&self, _command: &str) -> Option { - std::thread::sleep(std::time::Duration::from_millis(50)); - Some(GrokAuth::test_default()) - } - } - let refresher = ExternalBinaryRefresher::new(Arc::new(SlowRunner), "auth-binary".into()) - .with_timeout(std::time::Duration::from_millis(5)); - match refresher.refresh(RefreshReason::ServerRejected).await { - RefreshOutcome::PermanentFailure { error, .. } => { - assert_eq!(error.reason, RefreshTokenFailedReason::Other); - assert!( - !error.reason.is_sticky(), - "timeout must age out, not strand" - ); - } - other => panic!("a timed-out binary must be a permanent Other failure, got {other:?}"), - } - } - - #[tokio::test] - async fn external_binary_success_returns_fresh_token() { - let token = GrokAuth { - key: "ext-fresh".into(), - ..GrokAuth::test_default() - }; - let refresher = ExternalBinaryRefresher::new( - Arc::new(FakeRunner { - external_result: Some(token), - }), - "auth-binary".into(), - ); - match refresher.refresh(RefreshReason::ServerRejected).await { - RefreshOutcome::Success(auth) => assert_eq!(auth.key, "ext-fresh"), - other => panic!("a successful binary run must return Success, got {other:?}"), - } - } -} diff --git a/crates/codegen/kigi-shell/src/auth/refresh/kimi_refresher.rs b/crates/codegen/kigi-shell/src/auth/refresh/kimi_refresher.rs new file mode 100644 index 0000000..bae9c37 --- /dev/null +++ b/crates/codegen/kigi-shell/src/auth/refresh/kimi_refresher.rs @@ -0,0 +1,356 @@ +//! Kimi Code token refresher: drives `POST /api/oauth/token` with +//! `grant_type=refresh_token` through the [`TokenRefresher`] seam. +//! +//! Ports kimi-cli `OAuthManager._refresh_tokens`' sibling-safety behavior: +//! the persisted credential is re-read before the wire call (adopt a +//! rotation instead of refreshing), and after a 401/403 the persisted +//! credential is re-read once more (with a 1s grace) so a concurrent +//! process's freshly rotated token is adopted instead of tombstoning it. + +use std::sync::Arc; + +use crate::auth::error::RefreshTokenFailedReason; +use crate::auth::kimi_oauth::{self, RefreshError}; +use crate::auth::manager::RefreshReason; + +use super::{AuthSnapshot, RefreshOutcome, TokenRefresher}; + +/// Grace period after a 401/403 before concluding the refresh token is dead: +/// a concurrent instance may still be persisting its rotated token +/// (kimi-cli parity: `await asyncio.sleep(1)`). +const POST_UNAUTHORIZED_GRACE: std::time::Duration = std::time::Duration::from_secs(1); + +pub(crate) struct KimiRefresher { + auth: Arc, + /// OAuth host; `kigi_env::oauth_host()` in production, injectable for + /// wiremock tests. + host: String, +} + +impl KimiRefresher { + pub(crate) fn new(auth: Arc, host: String) -> Self { + Self { auth, host } + } + + /// Post-401 sibling check (kimi-cli parity): wait a beat, re-read the + /// persisted credential, and adopt it when its refresh token differs + /// from the one the server just rejected. + async fn adopt_rotation_after_unauthorized(&self, tried_rt: &str) -> Option { + tokio::time::sleep(POST_UNAUTHORIZED_GRACE).await; + let latest = self.auth.read_disk_auth()?; + let latest_rt = latest.refresh_token.as_deref()?; + if latest_rt == tried_rt { + return None; + } + kigi_log::unified_log::info( + "auth.refresh.adopted_rotation_after_401", + None, + Some(serde_json::json!({ + "adopted_rt_prefix": crate::auth::token_suffix(latest_rt), + "rejected_rt_prefix": crate::auth::token_suffix(tried_rt), + })), + ); + Some(RefreshOutcome::success(latest)) + } +} + +#[async_trait::async_trait] +impl TokenRefresher for KimiRefresher { + async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome { + tracing::info!(?reason, "auth: kimi refresh attempt starting"); + + let disk_auth = self.auth.read_disk_auth(); + + // Sibling short-circuit: a valid persisted token whose key differs + // from in-memory means another process refreshed between the + // refresh_chain disk check (under lock) and here. Adopt directly. + if let Some(ref d) = disk_auth + && !crate::auth::is_expired(d) + && self.auth.current().map(|a| a.key).as_deref() != Some(&d.key) + { + kigi_log::unified_log::info( + "auth.refresh.adopted_sibling_token", + None, + Some(serde_json::json!({ + "disk_key_prefix": crate::auth::token_suffix(&d.key), + })), + ); + return RefreshOutcome::success(d.clone()); + } + + let Some(auth) = super::resolve_refresh_credential(self.auth.as_ref(), disk_auth, reason) + else { + tracing::warn!(?reason, "auth: no credential available for refresh"); + return RefreshOutcome::transient("no token with refresh_token available"); + }; + let Some(refresh_token) = auth.refresh_token.clone() else { + tracing::warn!(?reason, "auth: resolved credential has no refresh token"); + return RefreshOutcome::transient("credential has no refresh token"); + }; + + tracing::info!( + rt_prefix = crate::auth::token_suffix(&refresh_token), + expires_at = ?auth.expires_at, + "auth: sending refresh_token grant" + ); + + match kimi_oauth::refresh_token(&self.host, &refresh_token).await { + Ok(new_auth) => { + kigi_log::unified_log::info( + "auth.refresh.token_rotated", + None, + Some(serde_json::json!({ + "new_key_prefix": crate::auth::token_suffix(&new_auth.key), + "expires_at": new_auth.expires_at.map(|e| e.to_rfc3339()), + })), + ); + RefreshOutcome::success(new_auth) + } + Err(RefreshError::Unauthorized { + status, + description, + }) => { + tracing::warn!(status, %description, "auth: refresh token rejected"); + if let Some(adopted) = self.adopt_rotation_after_unauthorized(&refresh_token).await + { + return adopted; + } + kigi_log::unified_log::warn( + "auth.refresh.unauthorized", + None, + Some(serde_json::json!({ + "status": status, + "description": description, + "rt_prefix": crate::auth::token_suffix(&refresh_token), + })), + ); + RefreshOutcome::permanent( + RefreshTokenFailedReason::RefreshTokenRejected, + Some(refresh_token), + ) + } + // kimi-cli parity: non-401 failures never tombstone; the next + // 60s tick (or pre-request check) retries. + Err( + e @ (RefreshError::Exhausted { .. } + | RefreshError::Fatal { .. } + | RefreshError::Local(_)), + ) => { + tracing::warn!(error = %e, "auth: refresh attempt failed (transient)"); + kigi_log::unified_log::warn( + "auth.refresh.transient_wire_failure", + None, + Some(serde_json::json!({ "error": format!("{e}") })), + ); + RefreshOutcome::transient(format!("token refresh failed: {e}")) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::auth::model::KimiAuth; + use chrono::{Duration, Utc}; + use parking_lot::Mutex; + use wiremock::matchers::{body_string_contains, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + /// Scriptable snapshot: `disk` can be swapped mid-test to simulate a + /// sibling process rotating the persisted credential. + struct FakeSnapshot { + current: Mutex>, + disk: Mutex>, + } + + impl FakeSnapshot { + fn new(current: Option, disk: Option) -> Arc { + Arc::new(Self { + current: Mutex::new(current), + disk: Mutex::new(disk), + }) + } + } + + impl AuthSnapshot for FakeSnapshot { + fn current(&self) -> Option { + self.current + .lock() + .clone() + .filter(|a| !crate::auth::is_expired(a)) + } + fn expired_auth(&self) -> Option { + self.current.lock().clone().filter(crate::auth::is_expired) + } + fn read_disk_auth(&self) -> Option { + self.disk.lock().clone() + } + fn is_expired(&self) -> bool { + self.current + .lock() + .as_ref() + .is_some_and(crate::auth::is_expired) + } + } + + fn expired_session(key: &str, rt: &str) -> KimiAuth { + KimiAuth { + key: key.into(), + refresh_token: Some(rt.into()), + expires_at: Some(Utc::now() - Duration::hours(1)), + expires_in: Some(3600), + ..KimiAuth::test_default() + } + } + + fn valid_session(key: &str, rt: &str) -> KimiAuth { + KimiAuth { + key: key.into(), + refresh_token: Some(rt.into()), + expires_at: Some(Utc::now() + Duration::hours(2)), + expires_in: Some(7200), + ..KimiAuth::test_default() + } + } + + fn token_json(access: &str, refresh: &str) -> serde_json::Value { + serde_json::json!({ + "access_token": access, + "refresh_token": refresh, + "expires_in": 3600, + "scope": "kimi-code", + "token_type": "bearer", + }) + } + + #[tokio::test] + async fn refresh_success_returns_rotated_token() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .and(body_string_contains("refresh_token=rt-old")) + .respond_with(ResponseTemplate::new(200).set_body_json(token_json("at-new", "rt-new"))) + .expect(1) + .mount(&server) + .await; + let stale = expired_session("at-old", "rt-old"); + let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale)); + let refresher = KimiRefresher::new(snap, server.uri()); + + let outcome = refresher.refresh(RefreshReason::PreRequest).await; + let RefreshOutcome::Success(new_auth) = outcome else { + panic!("expected success, got {outcome:?}"); + }; + assert_eq!(new_auth.key, "at-new"); + assert_eq!(new_auth.refresh_token.as_deref(), Some("rt-new")); + } + + #[tokio::test] + async fn adopts_valid_sibling_token_without_wire_call() { + // Disk has a fresh token with a different key: adopt, no HTTP. + let server = MockServer::start().await; + // No mock mounted: any request would 404 and fail the refresh. + let snap = FakeSnapshot::new( + Some(expired_session("at-old", "rt-old")), + Some(valid_session("at-sibling", "rt-sibling")), + ); + let refresher = KimiRefresher::new(snap, server.uri()); + let outcome = refresher.refresh(RefreshReason::PreRequest).await; + let RefreshOutcome::Success(adopted) = outcome else { + panic!("expected sibling adoption, got {outcome:?}"); + }; + assert_eq!(adopted.key, "at-sibling"); + assert!( + server.received_requests().await.unwrap().is_empty(), + "sibling adoption must not consume a refresh token on the wire" + ); + } + + #[tokio::test] + async fn unauthorized_tombstones_the_tried_refresh_token() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with( + ResponseTemplate::new(401) + .set_body_json(serde_json::json!({ "error_description": "revoked" })), + ) + .expect(1) + .mount(&server) + .await; + let stale = expired_session("at-old", "rt-dead"); + let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale)); + let refresher = KimiRefresher::new(snap, server.uri()); + + let outcome = refresher.refresh(RefreshReason::PreRequest).await; + let RefreshOutcome::PermanentFailure { + error, + rejected_refresh_token, + } = outcome + else { + panic!("expected permanent failure, got {outcome:?}"); + }; + assert_eq!(error.reason, RefreshTokenFailedReason::RefreshTokenRejected); + assert_eq!(rejected_refresh_token.as_deref(), Some("rt-dead")); + } + + #[tokio::test] + async fn unauthorized_adopts_sibling_rotation_instead_of_tombstoning() { + // 401 lands, but by the time we re-check, a sibling has persisted a + // rotated credential — adopt it (the mutual-logout race guard). + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + let stale = expired_session("at-old", "rt-dead"); + let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale)); + let refresher = KimiRefresher::new(snap.clone(), server.uri()); + + // Swap the persisted credential while the wire call is in flight. + let rotator = { + let snap = snap.clone(); + tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + *snap.disk.lock() = Some(valid_session("at-rotated", "rt-rotated")); + }) + }; + let outcome = refresher.refresh(RefreshReason::PreRequest).await; + rotator.await.unwrap(); + let RefreshOutcome::Success(adopted) = outcome else { + panic!("expected rotation adoption, got {outcome:?}"); + }; + assert_eq!(adopted.key, "at-rotated"); + } + + #[tokio::test] + async fn wire_exhaustion_is_transient_not_tombstoned() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/oauth/token")) + .respond_with(ResponseTemplate::new(503)) + .expect(3) + .mount(&server) + .await; + let stale = expired_session("at-old", "rt-old"); + let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale)); + let refresher = KimiRefresher::new(snap, server.uri()); + let outcome = refresher.refresh(RefreshReason::PreRequest).await; + assert!( + matches!(outcome, RefreshOutcome::TransientFailure { .. }), + "5xx exhaustion must stay transient: {outcome:?}" + ); + } + + #[tokio::test] + async fn no_credential_is_transient() { + let server = MockServer::start().await; + let snap = FakeSnapshot::new(None, None); + let refresher = KimiRefresher::new(snap, server.uri()); + let outcome = refresher.refresh(RefreshReason::PreRequest).await; + assert!(matches!(outcome, RefreshOutcome::TransientFailure { .. })); + } +} diff --git a/crates/codegen/kigi-shell/src/auth/refresh/mod.rs b/crates/codegen/kigi-shell/src/auth/refresh/mod.rs index 8298dfe..26f9e1c 100644 --- a/crates/codegen/kigi-shell/src/auth/refresh/mod.rs +++ b/crates/codegen/kigi-shell/src/auth/refresh/mod.rs @@ -1,42 +1,38 @@ -mod external_refresher; -mod oidc_refresher; +mod kimi_refresher; -use std::future::Future; -use std::pin::Pin; use std::sync::Arc; use crate::auth::manager::AuthManager; pub(crate) use crate::auth::manager::RefreshReason; -use crate::auth::model::GrokAuth; +use crate::auth::model::KimiAuth; -use external_refresher::ExternalBinaryRefresher; -pub(crate) use oidc_refresher::OidcRefresher; +pub(crate) use kimi_refresher::KimiRefresher; /// Read-only view of `AuthManager` for refreshers. Enforces the /// no-mutation contract on *credential* state at the type level: refreshers /// hold `Arc` and physically cannot call `update()`, /// `clear()`, `hot_swap()`, or `refresh_chain()`. pub(crate) trait AuthSnapshot: Send + Sync { - /// Read the current in-memory bearer outside the early-invalidation buffer. - fn current(&self) -> Option; + /// Read the current in-memory bearer outside the refresh threshold. + fn current(&self) -> Option; /// Read the expired in-memory bearer (for its `refresh_token`). - fn expired_auth(&self) -> Option; - /// Re-read auth.json from disk for the configured scope. Read-only w.r.t. - /// credentials, but may advance disk-observation state and emit transition - /// telemetry (not credential mutation). - fn read_disk_auth(&self) -> Option; + fn expired_auth(&self) -> Option; + /// Re-read the persisted credential (keyring → file) for the configured + /// scope. Read-only w.r.t. credentials, but may advance disk-observation + /// state and emit transition telemetry (not credential mutation). + fn read_disk_auth(&self) -> Option; /// Whether the in-memory bearer is expired. fn is_expired(&self) -> bool; } impl AuthSnapshot for AuthManager { - fn current(&self) -> Option { + fn current(&self) -> Option { self.current() } - fn expired_auth(&self) -> Option { + fn expired_auth(&self) -> Option { self.expired_auth() } - fn read_disk_auth(&self) -> Option { + fn read_disk_auth(&self) -> Option { self.read_disk_auth() } fn is_expired(&self) -> bool { @@ -44,31 +40,18 @@ impl AuthSnapshot for AuthManager { } } -/// Capability to run the operator's external auth binary. Split out of -/// [`AuthSnapshot`] so OIDC refreshers (read-only) physically cannot reach it -/// (interface segregation); only [`ExternalBinaryRefresher`] depends on it. -pub(crate) trait ExternalCommandRunner: Send + Sync { - /// Run the external auth binary and return the parsed output. - fn run_external_command(&self, command: &str) -> Option; -} - -impl ExternalCommandRunner for AuthManager { - fn run_external_command(&self, command: &str) -> Option { - self.run_external_refresh_command(command) - } -} - -/// The credential a refresh would send to the IdP: disk refresh-token first, -/// then the expired in-mem bearer, then current (only on `ServerRejected`). -/// Single source of truth shared by [`OidcRefresher::refresh`] (the attempt) and -/// `AuthManager::attempted_verdict_key` (the verdict scope), so the two can't -/// drift. The caller supplies the disk read: the verdict path passes a -/// side-effect-free read, the refresher the observing one. +/// The credential a refresh would send to the OAuth host: persisted +/// refresh-token first, then the expired in-mem bearer, then current (only on +/// `ServerRejected`). Single source of truth shared by +/// [`KimiRefresher::refresh`] (the attempt) and +/// `AuthManager::attempted_tombstone_key` (the tombstone scope), so the two +/// can't drift. The caller supplies the persisted read: the tombstone path +/// passes a side-effect-free read, the refresher the observing one. pub(crate) fn resolve_refresh_credential( snap: &dyn AuthSnapshot, - disk_auth: Option, + disk_auth: Option, reason: RefreshReason, -) -> Option { +) -> Option { disk_auth .filter(|a| a.refresh_token.is_some()) .or_else(|| snap.expired_auth()) @@ -84,18 +67,16 @@ pub(crate) fn resolve_refresh_credential( #[must_use = "RefreshOutcome encodes a state transition; route it through refresh_chain"] pub(crate) enum RefreshOutcome { /// Authority returned a fresh token. Caller persists via `update()`. - Success(Box), - /// Terminal failure (e.g. invalid_grant), or a transient escalated to - /// `Other` after repeated blips. Caller records a verdict scoped to the - /// rejected credential and retains it (`RefreshTokenRejected` is sticky, - /// the rest age out past the TTL). + Success(Box), + /// Terminal failure (401/403 from the OAuth host). Caller records a + /// tombstone scoped to the rejected refresh token; the 300s cooldown (or + /// a rotated persisted refresh token) clears it. PermanentFailure { error: crate::auth::error::RefreshTokenFailedError, - /// Key of the credential the refresher actually sent to the IdP, so - /// `refresh_chain` scopes the verdict to it. `None` when the authority - /// has no token key (external binary flow); the caller falls back to - /// its own resolution. - tried_key: Option, + /// The refresh-token value the refresher actually sent, so + /// `refresh_chain` scopes the tombstone to it. `None` when the + /// attempt never reached the wire. + rejected_refresh_token: Option, }, /// Transient / unknown failure. Caller may retry later. Message-only: the /// underlying cause is logged structurally at the refresher, then flattened @@ -105,19 +86,19 @@ pub(crate) enum RefreshOutcome { impl RefreshOutcome { /// A fresh credential from the authority (hides the `Box`). - pub(crate) fn success(auth: GrokAuth) -> Self { + pub(crate) fn success(auth: KimiAuth) -> Self { Self::Success(Box::new(auth)) } - /// Terminal failure for an already-classified reason against the credential - /// `tried_key` (the one actually sent to the IdP). + /// Terminal failure for an already-classified reason against the + /// refresh token actually sent to the OAuth host. pub(crate) fn permanent( reason: crate::auth::error::RefreshTokenFailedReason, - tried_key: Option, + rejected_refresh_token: Option, ) -> Self { Self::PermanentFailure { error: reason.into(), - tried_key, + rejected_refresh_token, } } @@ -139,67 +120,8 @@ pub(crate) trait TokenRefresher: Send + Sync { async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome; } -pub(crate) fn build_refresher( - auth_manager: Arc, - auth_provider_command: Option, -) -> Arc { - match auth_provider_command { - Some(cmd) => { - let runner: Arc = auth_manager; - Arc::new(ExternalBinaryRefresher::new(runner, cmd)) - } - None => { - let snapshot: Arc = auth_manager; - Arc::new(OidcRefresher::new(snapshot)) - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::auth::{AuthMode, GrokAuth, GrokComConfig}; - use chrono::{Duration, Utc}; - - /// auth_token_ttl makes is_token_expired use create_time + ttl for - /// External tokens without expires_at, instead of the 30-day fallback. - #[test] - fn token_ttl_expires_external_token_by_create_time() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig { - auth_token_ttl: Some(3600), // 1 hour - ..GrokComConfig::default() - }; - let mgr = AuthManager::new(dir.path(), cfg); - - // Token created 2 hours ago, no expires_at. With auth_token_ttl=3600, - // is_token_expired should return true (age 2h > ttl 1h). - let old_token = GrokAuth { - key: "old-external-token".into(), - auth_mode: AuthMode::External, - create_time: Utc::now() - Duration::hours(2), - expires_at: None, - ..GrokAuth::test_default() - }; - mgr.hot_swap(old_token); - assert!( - mgr.current().is_none(), - "expired external token via auth_token_ttl" - ); - assert!(mgr.is_expired()); - - // Fresh token created just now — should be valid. - let new_token = GrokAuth { - key: "new-external-token".into(), - auth_mode: AuthMode::External, - create_time: Utc::now(), - expires_at: None, - ..GrokAuth::test_default() - }; - mgr.hot_swap(new_token); - assert!( - mgr.current().is_some(), - "fresh external token should be valid" - ); - } +/// Build the production refresher against `kigi_env::oauth_host()`. +pub(crate) fn build_refresher(auth_manager: Arc) -> Arc { + let snapshot: Arc = auth_manager; + Arc::new(KimiRefresher::new(snapshot, kigi_env::oauth_host())) } diff --git a/crates/codegen/kigi-shell/src/auth/refresh/oidc_refresher.rs b/crates/codegen/kigi-shell/src/auth/refresh/oidc_refresher.rs deleted file mode 100644 index 5c39285..0000000 --- a/crates/codegen/kigi-shell/src/auth/refresh/oidc_refresher.rs +++ /dev/null @@ -1,260 +0,0 @@ -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; - -use crate::auth::error::RefreshTokenFailedReason; -use crate::auth::manager::RefreshReason; -use crate::auth::oidc::OidcRefreshResult; - -use super::{AuthSnapshot, RefreshOutcome, TokenRefresher}; - -#[cfg(test)] -use crate::auth::manager::AuthManager; - -/// Escalate to `PermanentFailure` after this many consecutive transient -/// failures (then `PERMANENT_FAILURE_TTL` allows recovery). OIDC tolerates more -/// blips than `ExternalBinaryRefresher` (1) since network refreshes flake more -/// than a local binary. -const MAX_CONSECUTIVE_TRANSIENT_FAILURES: u32 = 3; - -/// Consecutive transient-failure budget, scoped to the credential it accrued -/// against. Held under one lock so the credential check, reset, and increment -/// are a single atomic step. -#[derive(Default)] -struct TransientBudget { - /// Credential the count belongs to. A different credential (e.g. after - /// re-login on this long-lived refresher) re-arms the budget so a fresh, - /// valid token never inherits a dead one's escalation. - key: Option, - count: u32, -} - -pub(crate) struct OidcRefresher { - auth: Arc, - transient_budget: parking_lot::Mutex, -} - -impl OidcRefresher { - pub(crate) fn new(auth: Arc) -> Self { - Self { - auth, - transient_budget: parking_lot::Mutex::new(TransientBudget::default()), - } - } - - /// Clear the transient-blip budget on refresh progress (a fresh token or an - /// adopted sibling token), so later blips start from a full budget. - fn note_refresh_progress(&self) { - *self.transient_budget.lock() = TransientBudget::default(); - } - - fn record_transient_failure( - &self, - message: String, - tried_key: Option, - ) -> RefreshOutcome { - let escalate = { - let mut budget = self.transient_budget.lock(); - // Re-arm when the credential changes so a fresh token never inherits - // a prior credential's accrued blips. - if budget.key != tried_key { - budget.key = tried_key.clone(); - budget.count = 0; - } - budget.count += 1; - let escalate = budget.count >= MAX_CONSECUTIVE_TRANSIENT_FAILURES; - // On escalation reset the count so the next TTL window gets the full - // budget (the verdict gates refresh() meanwhile). The key is left in - // place; a same-key retry resumes from zero, a new key re-arms. - if escalate { - budget.count = 0; - } - escalate - }; - if escalate { - tracing::warn!(%message, "auth: escalating consecutive transient failures to permanent"); - RefreshOutcome::permanent(RefreshTokenFailedReason::Other, tried_key) - } else { - RefreshOutcome::transient(message) - } - } - - /// One-shot retry with disk's RT after `invalid_grant`. - /// - /// If disk already has a valid (unexpired) AT with a different key, - /// adopt it directly, without consuming the disk's RT in another IdP - /// call. This prevents cascading `invalid_grant` when a sibling - /// already refreshed and wrote a valid token. - async fn retry_with_fresh_disk_token( - &self, - tried: &crate::auth::GrokAuth, - ) -> Option { - let disk_now = self.auth.read_disk_auth()?; - - // If disk has a valid AT that differs from what we tried, - // a sibling already refreshed. Adopt directly — no IdP call. - if !crate::auth::is_expired(&disk_now) && disk_now.key != tried.key { - crate::unified_log::info( - "oidc refresh: disk has valid AT, adopting instead of consuming RT", - None, - Some(serde_json::json!({ - "disk_key_prefix": crate::auth::token_suffix(&disk_now.key), - "tried_key_prefix": crate::auth::token_suffix(&tried.key), - })), - ); - self.note_refresh_progress(); - return Some(RefreshOutcome::success(disk_now)); - } - - if disk_now.refresh_token.is_none() - || disk_now.refresh_token.as_deref() == tried.refresh_token.as_deref() - { - return None; - } - - crate::unified_log::info( - "oidc refresh retrying with disk token", - None, - Some(serde_json::json!({ - "tried_rt_prefix": tried - .refresh_token - .as_deref() - .map(crate::auth::token_suffix), - "disk_rt_prefix": disk_now - .refresh_token - .as_deref() - .map(crate::auth::token_suffix), - })), - ); - - match crate::auth::oidc::oidc_token_exchange(&disk_now).await { - OidcRefreshResult::Success(new_auth) => { - self.note_refresh_progress(); - Some(RefreshOutcome::Success(new_auth)) - } - OidcRefreshResult::TerminalError { reason } => { - crate::unified_log::warn( - "oidc refresh disk retry exhausted", - None, - Some(serde_json::json!({ "reason": format!("{reason:?}") })), - ); - Some(RefreshOutcome::permanent( - reason, - Some(disk_now.key.clone()), - )) - } - OidcRefreshResult::Failed => { - Some(RefreshOutcome::transient("OIDC disk-retry refresh failed")) - } - } - } -} - -#[async_trait::async_trait] -impl TokenRefresher for OidcRefresher { - async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome { - crate::unified_log::debug( - "oidc refresh enter", - None, - Some(serde_json::json!({ - "reason": format!("{reason:?}"), - "has_current": self.auth.current().is_some(), - "is_expired": self.auth.is_expired(), - })), - ); - - let disk_auth = self.auth.read_disk_auth(); - - // Short-circuit: if disk has a valid unexpired AT that differs - // from in-memory, a sibling refreshed between refresh_chain - // step 2 (disk check under lock) and here. Adopt it directly, - // no IdP call needed. - if let Some(ref d) = disk_auth - && !crate::auth::is_expired(d) - && self.auth.current().map(|a| a.key).as_deref() != Some(&d.key) - { - crate::unified_log::info( - "oidc refresh: sibling refreshed, adopting valid disk AT", - None, - Some(serde_json::json!({ - "disk_key_prefix": crate::auth::token_suffix(&d.key), - })), - ); - self.note_refresh_progress(); - return RefreshOutcome::success(d.clone()); - } - - let auth = super::resolve_refresh_credential(self.auth.as_ref(), disk_auth, reason); - - let Some(auth) = auth else { - crate::unified_log::warn( - "oidc refresh no token available", - None, - Some(serde_json::json!({ "reason": format!("{reason:?}") })), - ); - return RefreshOutcome::transient("no token with refresh_token available"); - }; - - crate::unified_log::info( - "oidc refresh attempting idp", - None, - Some(serde_json::json!({ - "has_rt": auth.refresh_token.is_some(), - "issuer": auth.oidc_issuer, - "client_id": auth.oidc_client_id, - "expires_at": auth.expires_at.map(|e| e.to_rfc3339()), - })), - ); - - match crate::auth::oidc::oidc_token_exchange(&auth).await { - OidcRefreshResult::Success(new_auth) => { - self.note_refresh_progress(); - RefreshOutcome::Success(new_auth) - } - OidcRefreshResult::TerminalError { reason } => { - // Sibling-rotation race: disk may hold a - // fresher RT than the one we tried. One-shot retry. - if reason == RefreshTokenFailedReason::RefreshTokenRejected - && let Some(retry_outcome) = self.retry_with_fresh_disk_token(&auth).await - { - return retry_outcome; - } - - RefreshOutcome::permanent(reason, Some(auth.key.clone())) - } - OidcRefreshResult::Failed => { - tracing::warn!( - refresh_reason = ?reason, - user_id = %auth.user_id, - has_refresh_token = auth.refresh_token.is_some(), - issuer = ?auth.oidc_issuer, - client_id = ?auth.oidc_client_id, - expires_at = ?auth.expires_at, - "auth: OIDC token refresh failed" - ); - crate::unified_log::error( - "oidc refresh failed", - None, - Some(serde_json::json!({ - "has_refresh_token": auth.refresh_token.is_some(), - "auth_mode": format!("{:?}", auth.auth_mode), - "issuer": auth.oidc_issuer, - "client_id": auth.oidc_client_id, - "expires_at": auth.expires_at.map(|e| e.to_rfc3339()), - })), - ); - self.record_transient_failure( - "OIDC token refresh failed".into(), - Some(auth.key.clone()), - ) - } - } - } -} - -#[cfg(test)] -#[path = "oidc_refresher_tests.rs"] -mod tests; - -#[cfg(test)] -#[path = "auth_backend_contract_tests.rs"] -mod auth_backend_contract_tests; diff --git a/crates/codegen/kigi-shell/src/auth/refresh/oidc_refresher_tests.rs b/crates/codegen/kigi-shell/src/auth/refresh/oidc_refresher_tests.rs deleted file mode 100644 index 59fdd88..0000000 --- a/crates/codegen/kigi-shell/src/auth/refresh/oidc_refresher_tests.rs +++ /dev/null @@ -1,1311 +0,0 @@ -//! Unit tests for [`super::oidc_refresher::OidcRefresher`]. Extracted -//! from `oidc_refresher.rs` so the implementation reads top-to-bottom; -//! wired in via `#[path = "oidc_refresher_tests.rs"] mod tests;`. - -use super::*; -use crate::auth::{GrokAuth, GrokComConfig}; -use chrono::{Duration, Utc}; - -// ── OIDC refresh E2E with mock IdP ───────────────────────────────── - -/// Start a mock server that handles OIDC discovery, token refresh, and -/// the proxy /user endpoint (called by AuthManager::update). -async fn start_mock_oidc_and_proxy() -> (String, tokio::task::JoinHandle<()>) { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let base_for_discovery = base.clone(); - - let app = axum::Router::new() - .route( - "/.well-known/openid-configuration", - axum::routing::get(move || { - let b = base_for_discovery.clone(); - async move { - axum::Json(serde_json::json!({ - "authorization_endpoint": format!("{b}/authorize"), - "token_endpoint": format!("{b}/token"), - })) - } - }), - ) - .route( - "/token", - axum::routing::post( - |body: axum::extract::Form>| async move { - // Verify the request is a refresh_token grant. - let grant_type = body - .iter() - .find(|(k, _)| k == "grant_type") - .map(|(_, v)| v.as_str()); - assert_eq!( - grant_type, - Some("refresh_token"), - "expected refresh_token grant" - ); - - axum::Json(serde_json::json!({ - "access_token": "oidc-refreshed-token", - "refresh_token": "oidc-new-rt", - "expires_in": 3600, - })) - }, - ), - ) - .route( - "/user", - axum::routing::get(|| async { - axum::Json(serde_json::json!({ - "userId": "user-42", - "email": "test@corp.com", - })) - }), - ); - - let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - (base, handle) -} - -fn write_auth_to_disk(dir: &std::path::Path, scope: &str, auth: &GrokAuth) { - let path = dir.join("auth.json"); - let mut map = crate::auth::read_auth_json(&path).unwrap_or_default(); - map.insert(scope.to_owned(), auth.clone()); - let json = serde_json::to_string_pretty(&map).unwrap(); - std::fs::write(&path, json).unwrap(); -} - -#[tokio::test] -async fn oidc_refresher_e2e_full_refresh_cycle() { - let (base_url, server) = start_mock_oidc_and_proxy().await; - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new( - AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url), - ); - - // Seed an expired OIDC token with all required fields. - let expired = GrokAuth { - key: "old-expired-token".into(), - create_time: Utc::now() - Duration::hours(2), - user_id: "user-42".into(), - email: Some("test@corp.com".into()), - refresh_token: Some("old-refresh-token".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - oidc_issuer: Some(base_url.clone()), - oidc_client_id: Some("test-client".into()), - ..GrokAuth::test_default() - }; - mgr.hot_swap(expired); - - let refresher = OidcRefresher::new(mgr.clone()); - - // ServerRejected so we bypass the cache check (token is expired anyway). - let result = refresher.refresh(RefreshReason::ServerRejected).await; - let new_auth = match result { - RefreshOutcome::Success(auth) => auth, - other => panic!("expected Success, got: {other:?}"), - }; - assert_eq!(new_auth.key, "oidc-refreshed-token"); - assert_eq!(new_auth.refresh_token.as_deref(), Some("oidc-new-rt")); - assert_eq!(new_auth.user_id, "user-42"); - assert_eq!(new_auth.oidc_issuer.as_deref(), Some(base_url.as_str())); - assert!(new_auth.expires_at.is_some()); - - server.abort(); -} - -#[tokio::test] -async fn oidc_refresher_e2e_proactive_returns_cached_when_valid() { - let (base_url, server) = start_mock_oidc_and_proxy().await; - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new( - AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url), - ); - - // Seed a valid (not expired) OIDC token. - let valid = GrokAuth { - key: "still-valid-token".into(), - user_id: "user-42".into(), - email: Some("test@corp.com".into()), - refresh_token: Some("rt".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - oidc_issuer: Some(base_url.clone()), - oidc_client_id: Some("test-client".into()), - ..GrokAuth::test_default() - }; - mgr.hot_swap(valid); - - let refresher = OidcRefresher::new(mgr.clone()); - // PreRequest with a valid token: the refresher finds no expired_auth - // (token is valid) and no disk token, so it returns TransientFailure. - // The PreRequest fast-path is handled by refresh_chain (not the refresher). - let result = refresher.refresh(RefreshReason::PreRequest).await; - assert!( - matches!(result, RefreshOutcome::TransientFailure { .. }), - "PreRequest with valid token should return TransientFailure (refresh_chain handles fast-path)" - ); - - server.abort(); -} - -#[tokio::test] -async fn oidc_refresher_e2e_force_refreshes_locally_valid_token() { - let (base_url, server) = start_mock_oidc_and_proxy().await; - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new( - AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url), - ); - - // Seed a valid (not yet expired) OIDC token. force=true simulates - // the reactive 401 path — server rejected the token even though it - // looks locally valid (e.g. clock skew, server-side revocation). - // The refresher should still attempt an OIDC refresh. - let valid = GrokAuth { - key: "still-valid-token".into(), - user_id: "user-42".into(), - email: Some("test@corp.com".into()), - refresh_token: Some("rt".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - oidc_issuer: Some(base_url.clone()), - oidc_client_id: Some("test-client".into()), - ..GrokAuth::test_default() - }; - mgr.hot_swap(valid); - - let refresher = OidcRefresher::new(mgr.clone()); - // ServerRejected should refresh even though the token is locally valid. - let result = refresher.refresh(RefreshReason::ServerRejected).await; - let new_auth = match result { - RefreshOutcome::Success(auth) => auth, - other => panic!("expected Success, got: {other:?}"), - }; - assert_eq!( - new_auth.key, "oidc-refreshed-token", - "ServerRejected should refresh even when token is locally valid" - ); - - server.abort(); -} - -// ── Near-expiry (5-minute buffer) refresh scenarios ────────────── - -/// Regression test for token-expiry-window bug: when the token is within -/// the 5-minute early-invalidation buffer, current() returns None but -/// expired_auth() returns the token. The OidcRefresher must successfully -/// refresh it via the refresh_token grant — the exact path exercised by -/// initialize() in mvp_agent/mod.rs. -#[tokio::test] -async fn oidc_refresher_e2e_near_expiry_within_buffer_refreshes() { - let (base_url, server) = start_mock_oidc_and_proxy().await; - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new( - AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url), - ); - - // Token expires in 3 minutes — inside the 5-minute buffer. - // current() will return None, but expired_auth() will return it. - let near_expiry = GrokAuth { - key: "about-to-expire-token".into(), - user_id: "user-42".into(), - email: Some("test@corp.com".into()), - refresh_token: Some("rt-still-valid".into()), - expires_at: Some(Utc::now() + Duration::minutes(3)), - oidc_issuer: Some(base_url.clone()), - oidc_client_id: Some("test-client".into()), - ..GrokAuth::test_default() - }; - mgr.hot_swap(near_expiry); - - // Preconditions: confirm the bug scenario - assert!( - mgr.current().is_none(), - "current() should be None within buffer" - ); - assert!( - mgr.is_expired(), - "is_expired() should be true within buffer" - ); - assert!( - mgr.expired_auth().is_some(), - "expired_auth() should return the token" - ); - - // Simulate the initialize() refresh path: get expired auth, call try_refresh - mgr.set_refresher(std::sync::Arc::new(OidcRefresher::new(mgr.clone()))); - let refreshed = mgr.auth().await.ok(); - - assert!( - refreshed.is_some(), - "OIDC refresh should succeed for near-expiry token" - ); - let fresh = refreshed.unwrap(); - assert_eq!(fresh.key, "oidc-refreshed-token"); - assert_eq!(fresh.refresh_token.as_deref(), Some("oidc-new-rt")); - - // After refresh, current() should return the new valid token - let current = mgr.current(); - assert!( - current.is_some(), - "current() should return new token after refresh" - ); - assert_eq!(current.unwrap().key, "oidc-refreshed-token"); - - server.abort(); -} - -/// When the near-expiry token has a refresh_token but the IdP rejects -/// the refresh (e.g. refresh_token revoked), silent refresh must fail. -#[tokio::test] -async fn oidc_refresher_e2e_near_expiry_idp_rejects_refresh() { - // Start a mock that rejects refresh requests with 401 - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base_url = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let base_for_discovery = base_url.clone(); - - let app = axum::Router::new() - .route( - "/.well-known/openid-configuration", - axum::routing::get(move || { - let b = base_for_discovery.clone(); - async move { - axum::Json(serde_json::json!({ - "authorization_endpoint": format!("{b}/authorize"), - "token_endpoint": format!("{b}/token"), - })) - } - }), - ) - .route( - "/token", - axum::routing::post(|| async { - ( - axum::http::StatusCode::BAD_REQUEST, - axum::Json(serde_json::json!({"error": "invalid_grant"})), - ) - }), - ); - - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new( - AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url), - ); - - let near_expiry = GrokAuth { - key: "about-to-expire-token".into(), - user_id: "user-42".into(), - refresh_token: Some("rt-revoked".into()), - expires_at: Some(Utc::now() + Duration::minutes(3)), - oidc_issuer: Some(base_url.clone()), - oidc_client_id: Some("test-client".into()), - ..GrokAuth::test_default() - }; - mgr.hot_swap(near_expiry); - - // auth() dispatches to refresh_chain -> OidcRefresher -> invalid_grant. - // Because the token is still within real expires_at (3 min from now), - // the grace path returns the cached token as a fallback. - mgr.set_refresher(std::sync::Arc::new(OidcRefresher::new(mgr.clone()))); - let refreshed = mgr.auth().await; - assert!( - refreshed.is_ok(), - "grace path should return the cached token while within real expires_at" - ); - assert_eq!(refreshed.unwrap().key, "about-to-expire-token"); - - server.abort(); -} - -/// On `invalid_client` (client_id rotated, soft-deleted, or disabled), the -/// credential is retained and a permanent-failure verdict cached. Verdict + TTL -/// stop the retry loop; the bearer drops only on explicit logout, so a -/// transient client-rotation blip self-heals without a fleet re-login. -#[tokio::test] -async fn oidc_refresher_e2e_invalid_client_caches_verdict_and_retains_credentials() { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base_url = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let base_for_discovery = base_url.clone(); - - let app = axum::Router::new() - .route( - "/.well-known/openid-configuration", - axum::routing::get(move || { - let b = base_for_discovery.clone(); - async move { - axum::Json(serde_json::json!({ - "authorization_endpoint": format!("{b}/authorize"), - "token_endpoint": format!("{b}/token"), - })) - } - }), - ) - .route( - "/token", - axum::routing::post(|| async { - ( - axum::http::StatusCode::UNAUTHORIZED, - axum::Json(serde_json::json!({ - "error": "invalid_client", - "error_description": "Unknown client" - })), - ) - }), - ); - - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new( - AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url), - ); - - let expired = GrokAuth { - key: "old-token".into(), - create_time: Utc::now() - Duration::hours(2), - user_id: "user-42".into(), - refresh_token: Some("rt-valid".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - oidc_issuer: Some(base_url.clone()), - oidc_client_id: Some("deleted-client-id".into()), - ..GrokAuth::test_default() - }; - mgr.hot_swap(expired); - - mgr.set_refresher(std::sync::Arc::new(OidcRefresher::new(mgr.clone()))); - let refreshed = mgr.auth().await.ok(); - assert!( - refreshed.is_none(), - "refresh should fail when client is unknown" - ); - - // Credential retained (not cleared) — the bearer may be fine; the client - // credential isn't. - assert!( - mgr.expired_auth().is_some(), - "credentials must be retained after invalid_client", - ); - // The verdict is cached, scoped to the retained credential, and carries - // the non-sticky `ClientRejected` reason (so it ages out, not stuck-forever). - match mgr.permanent_failure() { - Some(crate::auth::AuthError::Refresh(crate::auth::RefreshTokenError::Permanent(e))) => { - assert_eq!( - e.reason, - crate::auth::RefreshTokenFailedReason::ClientRejected, - "invalid_client must map to ClientRejected", - ); - assert!( - !e.reason.is_sticky(), - "ClientRejected must age out past the TTL, not stick forever", - ); - } - other => panic!("invalid_client must cache a permanent-failure verdict, got {other:?}"), - } - - server.abort(); -} - -/// When the IdP would return `invalid_client` but disk auth.json already holds -/// a valid token with a different client_id (a sibling re-authenticated during -/// a client rotation), `auth()` adopts the sibling's disk token instead of -/// failing. -#[tokio::test] -async fn oidc_refresher_e2e_invalid_client_adopts_valid_sibling_disk_token() { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base_url = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let base_for_discovery = base_url.clone(); - - let app = axum::Router::new() - .route( - "/.well-known/openid-configuration", - axum::routing::get(move || { - let b = base_for_discovery.clone(); - async move { - axum::Json(serde_json::json!({ - "authorization_endpoint": format!("{b}/authorize"), - "token_endpoint": format!("{b}/token"), - })) - } - }), - ) - .route( - "/token", - axum::routing::post(|| async { - ( - axum::http::StatusCode::UNAUTHORIZED, - axum::Json(serde_json::json!({ - "error": "invalid_client", - "error_description": "Unknown client" - })), - ) - }), - ); - - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let scope = cfg.auth_scope(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg).with_proxy_base_url(&base_url)); - - // Pre-populate disk with auth that has a *different* client_id, - // simulating another process having re-authenticated. - let disk_auth = GrokAuth { - key: "disk-fresh-token".into(), - user_id: "user-42".into(), - refresh_token: Some("rt-disk".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - oidc_issuer: Some(base_url.clone()), - oidc_client_id: Some("rotated-new-client-id".into()), - ..GrokAuth::test_default() - }; - let mut store = std::collections::BTreeMap::new(); - store.insert(scope, disk_auth); - let json = serde_json::to_string_pretty(&store).unwrap(); - std::fs::write(dir.path().join("auth.json"), json).unwrap(); - - // In-memory auth has the OLD client_id that the server rejects. - let expired = GrokAuth { - key: "old-token".into(), - create_time: Utc::now() - Duration::hours(2), - user_id: "user-42".into(), - refresh_token: Some("rt-old".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - oidc_issuer: Some(base_url.clone()), - oidc_client_id: Some("deleted-client-id".into()), - ..GrokAuth::test_default() - }; - mgr.hot_swap(expired); - - // auth() picks up the valid disk token via try_use_disk_token - // (disk has a different, unexpired entry from a sibling process). - // This is BETTER than the old try_refresh path which ignored - // the valid disk token and hit the IdP. - mgr.set_refresher(std::sync::Arc::new(OidcRefresher::new(mgr.clone()))); - let refreshed = mgr.auth().await; - assert!( - refreshed.is_ok(), - "auth() should pick up the valid disk token, got: {refreshed:?}" - ); - assert_eq!( - refreshed.unwrap().oidc_client_id.as_deref(), - Some("rotated-new-client-id"), - "should use the sibling's rotated client_id from disk" - ); - - server.abort(); -} - -// The standalone `try_refresh_session_token` helper that previously -// lived in this module was removed when refresh was centralized in -// `AuthManager`. Its call sites now go through `AuthManager::auth()` -// / `AuthManager::unauthorized_recovery()`, both of which have their -// own coverage in `manager.rs`. The historical regression tests for -// the helper (`try_refresh_respects_auth_type`, -// `auth_type_must_be_session_token_after_session_key_set`) were -// dropped along with the function. The `resolve_credentials` -// invariant they also pinned remains covered by -// `agent::config::tests::{resolve_credentials_sets_auth_type, -// resolve_credentials_no_session_key_returns_api_key}`. - -/// When another process has already refreshed and written a valid token -/// to auth.json, `refresh_chain` (via `auth()`) should pick it up from -/// disk instead of hitting the IdP. -#[tokio::test] -async fn oidc_refresh_picks_up_valid_disk_token() { - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let scope = cfg.auth_scope(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg).with_proxy_base_url("http://127.0.0.1:1")); - - // Seed in-memory with an expired token (stale refresh_token). - let expired = GrokAuth { - key: "old-expired-token".into(), - create_time: Utc::now() - Duration::hours(2), - user_id: "user-42".into(), - refresh_token: Some("stale-rt".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - oidc_issuer: Some("https://idp.example.com".into()), - oidc_client_id: Some("client-1".into()), - ..GrokAuth::test_default() - }; - mgr.hot_swap(expired); - - // Simulate another process writing a valid token to disk. - let fresh_on_disk = GrokAuth { - key: "fresh-from-other-process".into(), - user_id: "user-42".into(), - email: Some("user@test.com".into()), - refresh_token: Some("new-rt".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - oidc_issuer: Some("https://idp.example.com".into()), - oidc_client_id: Some("client-1".into()), - ..GrokAuth::test_default() - }; - write_auth_to_disk(dir.path(), &scope, &fresh_on_disk); - - // Disk-token pickup is now refresh_chain's responsibility. - // Go through auth() which calls refresh_chain. - let result = mgr.auth().await; - assert_eq!( - result.unwrap().key, - "fresh-from-other-process", - "should use the valid token written by another process" - ); - assert_eq!( - mgr.current().unwrap().key, - "fresh-from-other-process", - "in-memory state should be updated" - ); -} - -/// When the disk token is also expired but has a newer refresh_token, -/// the OIDC refresher should use the disk's RT for the IdP call. -#[tokio::test] -async fn oidc_refresh_uses_disk_refresh_token() { - // Custom mock that captures the submitted refresh_token so we can - // assert the disk RT was sent, not the stale in-memory one. - let captured_rt = std::sync::Arc::new(parking_lot::Mutex::new(None::)); - let captured_for_handler = captured_rt.clone(); - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base_url = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let base_for_disc = base_url.clone(); - let app = axum::Router::new() - .route( - "/.well-known/openid-configuration", - axum::routing::get(move || { - let b = base_for_disc.clone(); - async move { - axum::Json(serde_json::json!({ - "authorization_endpoint": format!("{b}/authorize"), - "token_endpoint": format!("{b}/token"), - })) - } - }), - ) - .route( - "/token", - axum::routing::post(move |body: axum::extract::Form>| { - let captured = captured_for_handler.clone(); - async move { - let rt = body - .iter() - .find(|(k, _)| k == "refresh_token") - .map(|(_, v)| v.clone()); - *captured.lock() = rt; - axum::Json(serde_json::json!({ - "access_token": "oidc-refreshed-token", - "refresh_token": "oidc-new-rt", - "expires_in": 3600, - })) - } - }), - ) - .route( - "/user", - axum::routing::get(|| async { - axum::Json(serde_json::json!({ "userId": "user-42", "email": "test@corp.com" })) - }), - ); - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let scope = cfg.auth_scope(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg).with_proxy_base_url(&base_url)); - - mgr.hot_swap(GrokAuth { - key: "old-mem-token".into(), - create_time: Utc::now() - Duration::hours(2), - user_id: "user-42".into(), - email: Some("test@corp.com".into()), - refresh_token: Some("stale-rt-will-fail".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - oidc_issuer: Some(base_url.clone()), - oidc_client_id: Some("test-client".into()), - ..GrokAuth::test_default() - }); - - write_auth_to_disk( - dir.path(), - &scope, - &GrokAuth { - key: "old-disk-token".into(), - create_time: Utc::now() - Duration::hours(2), - user_id: "user-42".into(), - email: Some("test@corp.com".into()), - refresh_token: Some("disk-rt-valid".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - oidc_issuer: Some(base_url.clone()), - oidc_client_id: Some("test-client".into()), - ..GrokAuth::test_default() - }, - ); - - let refresher = OidcRefresher::new(mgr.clone()); - let result = refresher.refresh(RefreshReason::ServerRejected).await; - assert!(matches!(result, RefreshOutcome::Success(_))); - - assert_eq!( - captured_rt.lock().as_deref(), - Some("disk-rt-valid"), - "must send the disk token's RT to the IdP, not the stale in-memory one" - ); - - server.abort(); -} - -/// When the lock file is held by another process (simulated), the -/// refresher should fall through and still attempt the refresh. -/// (Lock is now managed by refresh_chain, but the refresher itself -/// should still succeed without a lock.) -#[tokio::test] -async fn lock_timeout_falls_through_to_refresh() { - let (base_url, server) = start_mock_oidc_and_proxy().await; - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let mgr = Arc::new(AuthManager::new(dir.path(), cfg).with_proxy_base_url(&base_url)); - - // Seed with expired token that has a valid refresh_token. - let expired = GrokAuth { - key: "old-token".into(), - create_time: Utc::now() - Duration::hours(2), - user_id: "user-42".into(), - email: Some("test@corp.com".into()), - refresh_token: Some("rt-valid".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - oidc_issuer: Some(base_url.clone()), - oidc_client_id: Some("test-client".into()), - ..GrokAuth::test_default() - }; - mgr.hot_swap(expired); - - // Hold the lock file externally so the refresher times out. - let lock_path = dir.path().join("auth.json.lock"); - let lock_file = std::fs::OpenOptions::new() - .read(true) - .write(true) - .create(true) - .truncate(false) - .open(&lock_path) - .unwrap(); - use fs2::FileExt; - lock_file.lock_exclusive().unwrap(); - - // Use a very short timeout so the test doesn't wait 30s. - let _lock = mgr - .try_lock_auth_file_async(std::time::Duration::from_millis(100)) - .await; - assert!(_lock.is_none(), "lock should timeout"); - - // The refresh should still succeed (refresher doesn't need the lock). - let refresher = OidcRefresher::new(mgr.clone()); - let result = refresher.refresh(RefreshReason::ServerRejected).await; - let new_auth = match result { - RefreshOutcome::Success(auth) => auth, - other => panic!("expected Success, got: {other:?}"), - }; - assert_eq!( - new_auth.key, "oidc-refreshed-token", - "refresh should succeed even when lock times out" - ); - - lock_file.unlock().unwrap(); - server.abort(); -} - -// ── Disk-token retry on invalid_grant ────────────────────── - -/// Mock IdP. `success_rts`: RT -> (access_token, new_rt). -/// `rotation_targets`: RT -> new disk RT written as a side effect -/// on `invalid_grant` (simulates sibling rotation). `attempts` -/// counts every POST so tests can assert one-shot. -async fn start_mock_oidc_with_disk_rotation( - success_rts: std::collections::HashMap<&'static str, (&'static str, &'static str)>, - rotation_targets: std::collections::HashMap<&'static str, &'static str>, - sibling_writes_disk: Option<(std::path::PathBuf, String)>, - attempts: Arc, -) -> (String, tokio::task::JoinHandle<()>) { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let base_for_discovery = base.clone(); - let attempts_for_handler = attempts.clone(); - - let app = axum::Router::new() - .route( - "/.well-known/openid-configuration", - axum::routing::get(move || { - let b = base_for_discovery.clone(); - async move { - axum::Json(serde_json::json!({ - "authorization_endpoint": format!("{b}/authorize"), - "token_endpoint": format!("{b}/token"), - })) - } - }), - ) - .route( - "/token", - axum::routing::post(move |body: axum::extract::Form>| { - let counter = attempts_for_handler.clone(); - let sibling_writes_disk = sibling_writes_disk.clone(); - let success_rts = success_rts.clone(); - let rotation_targets = rotation_targets.clone(); - async move { - counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - use axum::response::IntoResponse; - let rt = body - .iter() - .find(|(k, _)| k == "refresh_token") - .map(|(_, v)| v.as_str()) - .unwrap_or(""); - - if let Some((access, new_rt)) = success_rts.get(rt) { - return ( - axum::http::StatusCode::OK, - axum::Json(serde_json::json!({ - "access_token": access, - "refresh_token": new_rt, - "expires_in": 3600, - })), - ) - .into_response(); - } - - if let Some(rotate_to) = rotation_targets.get(rt) - && let Some((ref path, ref scope)) = sibling_writes_disk - { - let mut map = crate::auth::read_auth_json(path).unwrap_or_default(); - if let Some(entry) = map.get_mut(scope) { - entry.refresh_token = Some((*rotate_to).into()); - } - let json = serde_json::to_string_pretty(&map).unwrap(); - std::fs::write(path, json).unwrap(); - } - - ( - axum::http::StatusCode::BAD_REQUEST, - axum::Json(serde_json::json!({ - "error": "invalid_grant", - "error_description": - "Refresh token has been revoked", - })), - ) - .into_response() - } - }), - ) - .route( - "/user", - axum::routing::get(|| async { - axum::Json(serde_json::json!({ - "userId": "user-42", - "email": "test@corp.com", - })) - }), - ); - - let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - (base, handle) -} - -/// Sibling-rotation race: tried RT -> invalid_grant + sibling rotates disk; -/// retry with disk RT must succeed without surfacing failure. -#[tokio::test] -async fn refresher_retries_with_disk_token_after_invalid_grant() { - use std::sync::atomic::{AtomicU32, Ordering}; - - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let scope = cfg.auth_scope(); - let auth_path = dir.path().join("auth.json"); - - let attempts = Arc::new(AtomicU32::new(0)); - let success_rts = std::collections::HashMap::from([( - "rt-fresh-from-sibling", - ("fresh-access-token", "rt-newest"), - )]); - let rotation_targets = std::collections::HashMap::from([("rt-stale", "rt-fresh-from-sibling")]); - let (base_url, server) = start_mock_oidc_with_disk_rotation( - success_rts, - rotation_targets, - Some((auth_path.clone(), scope.clone())), - attempts.clone(), - ) - .await; - - let mgr = Arc::new(AuthManager::new(dir.path(), cfg).with_proxy_base_url(&base_url)); - - // Disk and memory both have rt-stale; mock rotates disk on - // the first invalid_grant so the retry sees the fresh RT. - let stale = GrokAuth { - key: "stale-access-token".into(), - create_time: Utc::now() - Duration::hours(2), - user_id: "user-42".into(), - refresh_token: Some("rt-stale".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - oidc_issuer: Some(base_url.clone()), - oidc_client_id: Some("test-client".into()), - ..GrokAuth::test_default() - }; - write_auth_to_disk(dir.path(), &scope, &stale); - mgr.hot_swap(stale); - - let refresher = OidcRefresher::new(mgr.clone()); - let outcome = refresher.refresh(RefreshReason::PreRequest).await; - - match outcome { - RefreshOutcome::Success(new_auth) => { - assert_eq!( - new_auth.key, "fresh-access-token", - "retry should return the access_token issued for the disk RT" - ); - assert_eq!( - new_auth.refresh_token.as_deref(), - Some("rt-newest"), - "retry should carry the newly-issued refresh_token forward" - ); - } - other => panic!("expected Success after disk-token retry, got: {other:?}"), - } - assert_eq!( - attempts.load(Ordering::SeqCst), - 2, - "exactly two IdP calls: stale RT then disk RT" - ); - - server.abort(); -} - -/// invalid_grant -> retry uses sibling-rotated disk RT -> invalid_client. -/// Both ATs expired: PermanentFailure is recorded (not demoted). -#[tokio::test] -async fn refresher_disk_retry_invalid_client_with_different_client_id_preserves_disk() { - use std::sync::atomic::{AtomicU32, Ordering}; - - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let scope = cfg.auth_scope(); - let auth_path = dir.path().join("auth.json"); - - let attempts = Arc::new(AtomicU32::new(0)); - let attempts_for_handler = attempts.clone(); - let auth_path_for_handler = auth_path.clone(); - let scope_for_handler = scope.clone(); - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base_url = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let base_for_discovery = base_url.clone(); - - let app = axum::Router::new() - .route( - "/.well-known/openid-configuration", - axum::routing::get(move || { - let b = base_for_discovery.clone(); - async move { - axum::Json(serde_json::json!({ - "authorization_endpoint": format!("{b}/authorize"), - "token_endpoint": format!("{b}/token"), - })) - } - }), - ) - .route( - "/token", - axum::routing::post(move |body: axum::extract::Form>| { - let attempts = attempts_for_handler.clone(); - let auth_path = auth_path_for_handler.clone(); - let scope = scope_for_handler.clone(); - async move { - use axum::response::IntoResponse; - let n = attempts.fetch_add(1, Ordering::SeqCst); - let rt = body - .iter() - .find(|(k, _)| k == "refresh_token") - .map(|(_, v)| v.as_str()) - .unwrap_or(""); - if n == 0 { - // First call: invalid_grant + sibling rotates disk. - assert_eq!(rt, "rt-stale"); - let mut map = crate::auth::read_auth_json(&auth_path).unwrap_or_default(); - if let Some(entry) = map.get_mut(&scope) { - entry.refresh_token = Some("rt-sibling".into()); - entry.oidc_client_id = Some("rotated-new-client-id".into()); - entry.key = "sibling-fresh-access".into(); - } - let json = serde_json::to_string_pretty(&map).unwrap(); - std::fs::write(&auth_path, json).unwrap(); - return ( - axum::http::StatusCode::BAD_REQUEST, - axum::Json(serde_json::json!({ - "error": "invalid_grant", - "error_description": "RT revoked", - })), - ) - .into_response(); - } - // Retry uses the sibling's RT -> invalid_client. - assert_eq!(rt, "rt-sibling"); - ( - axum::http::StatusCode::UNAUTHORIZED, - axum::Json(serde_json::json!({ - "error": "invalid_client", - "error_description": "Unknown client", - })), - ) - .into_response() - } - }), - ); - - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - - let mgr = Arc::new(AuthManager::new(dir.path(), cfg).with_proxy_base_url(&base_url)); - - let stale = GrokAuth { - key: "stale-access".into(), - create_time: Utc::now() - Duration::hours(2), - user_id: "user-42".into(), - refresh_token: Some("rt-stale".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - oidc_issuer: Some(base_url.clone()), - oidc_client_id: Some("client-stale".into()), - ..GrokAuth::test_default() - }; - write_auth_to_disk(dir.path(), &scope, &stale); - mgr.hot_swap(stale); - - let refresher: Arc = - Arc::new(OidcRefresher::new(mgr.clone())); - mgr.set_refresher(refresher); - - let result = mgr - .refresh_chain( - crate::auth::token_type::TokenType::OidcSession, - RefreshReason::ServerRejected, - ) - .await; - - match result { - Err(crate::auth::AuthError::Refresh(crate::auth::RefreshTokenError::Permanent(_))) => {} - other => panic!("expected PermanentFailure, got: {other:?}"), - } - - // Credential retained; the cached verdict (scoped to it) stops the storm. - assert!( - mgr.current_or_expired().is_some(), - "credential must be retained on permanent failure" - ); - assert!( - mgr.permanent_failure().is_some(), - "verdict must be cached to stop the retry storm" - ); - assert_eq!(attempts.load(Ordering::SeqCst), 2, "no recursion"); - - server.abort(); -} - -/// Both RTs revoked: retry is strictly one-shot (no third call); -/// refresh_chain's disk-RT-differs guard preserves disk creds. -#[tokio::test] -async fn refresher_disk_retry_is_one_shot() { - use std::sync::atomic::{AtomicU32, Ordering}; - - let dir = tempfile::tempdir().unwrap(); - let cfg = GrokComConfig::default(); - let scope = cfg.auth_scope(); - let auth_path = dir.path().join("auth.json"); - - let attempts = Arc::new(AtomicU32::new(0)); - // Empty success_rts; disk rotates after the first attempt - // so the retry fires but also fails -- exhausts cleanly. - let success_rts: std::collections::HashMap<&str, (&str, &str)> = - std::collections::HashMap::new(); - let rotation_targets = std::collections::HashMap::from([("rt-stale", "rt-also-revoked")]); - let (base_url, server) = start_mock_oidc_with_disk_rotation( - success_rts, - rotation_targets, - Some((auth_path.clone(), scope.clone())), - attempts.clone(), - ) - .await; - - let mgr = Arc::new(AuthManager::new(dir.path(), cfg).with_proxy_base_url(&base_url)); - - let stale = GrokAuth { - key: "stale-access-token".into(), - create_time: Utc::now() - Duration::hours(2), - user_id: "user-42".into(), - refresh_token: Some("rt-stale".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - oidc_issuer: Some(base_url.clone()), - oidc_client_id: Some("test-client".into()), - ..GrokAuth::test_default() - }; - write_auth_to_disk(dir.path(), &scope, &stale); - mgr.hot_swap(stale); - - let refresher = OidcRefresher::new(mgr.clone()); - let outcome = refresher.refresh(RefreshReason::PreRequest).await; - - match outcome { - RefreshOutcome::PermanentFailure { error, .. } => { - assert_eq!(error.reason, RefreshTokenFailedReason::RefreshTokenRejected); - } - other => panic!("expected PermanentFailure after exhausted retry, got: {other:?}"), - } - - assert_eq!( - attempts.load(Ordering::SeqCst), - 2, - "exactly two IdP calls — disk-token retry must NOT recurse" - ); - - // Disk auth must still be present (the refresher never clears). - assert!( - mgr.read_disk_auth().is_some(), - "refresher must not touch disk; clearing is refresh_chain's responsibility" - ); - - server.abort(); -} - -// ── Sleep-gate E2E (real OidcRefresher + mock IdP) ───────────────── - -/// Mock IdP that counts `/token` POSTs so a test can prove a deferred refresh -/// suppressed the network call rather than just changing the return value. -async fn start_counting_mock_oidc( - token_hits: Arc, -) -> (String, tokio::task::JoinHandle<()>) { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let base_for_discovery = base.clone(); - - let app = axum::Router::new() - .route( - "/.well-known/openid-configuration", - axum::routing::get(move || { - let b = base_for_discovery.clone(); - async move { - axum::Json(serde_json::json!({ - "authorization_endpoint": format!("{b}/authorize"), - "token_endpoint": format!("{b}/token"), - })) - } - }), - ) - .route( - "/token", - axum::routing::post(move || { - let hits = token_hits.clone(); - async move { - hits.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - axum::Json(serde_json::json!({ - "access_token": "oidc-refreshed-token", - "refresh_token": "oidc-new-rt", - "expires_in": 3600, - })) - } - }), - ) - .route( - "/user", - axum::routing::get(|| async { - axum::Json(serde_json::json!({ "userId": "user-42", "email": "test@corp.com" })) - }), - ); - - let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - (base, handle) -} - -fn expired_oidc_for(base_url: &str) -> GrokAuth { - GrokAuth { - key: "old-expired-token".into(), - create_time: Utc::now() - Duration::hours(2), - user_id: "user-42".into(), - email: Some("test@corp.com".into()), - refresh_token: Some("rt-valid".into()), - expires_at: Some(Utc::now() - Duration::hours(1)), - oidc_issuer: Some(base_url.to_owned()), - oidc_client_id: Some("test-client".into()), - ..GrokAuth::test_default() - } -} - -/// While sleep is imminent, `auth()` defers and never reaches the IdP; after -/// wake it recovers via a real OIDC refresh. Exercises the production -/// `OidcRefresher` against a mock IdP, not a stub. -#[tokio::test] -async fn sleep_gate_e2e_defers_then_recovers_on_wake() { - use std::sync::atomic::{AtomicU32, Ordering}; - - let token_hits = Arc::new(AtomicU32::new(0)); - let (base_url, server) = start_counting_mock_oidc(token_hits.clone()).await; - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new( - AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url), - ); - mgr.hot_swap(expired_oidc_for(&base_url)); - mgr.set_refresher(Arc::new(OidcRefresher::new(mgr.clone()))); - - mgr.set_system_sleep_imminent(true); - let err = mgr.auth().await.unwrap_err(); - assert!( - matches!( - err, - crate::auth::AuthError::Refresh(crate::auth::RefreshTokenError::Transient(_)) - ), - "gated refresh must return a transient refresh error, got {err:?}" - ); - assert_eq!( - token_hits.load(Ordering::SeqCst), - 0, - "a deferred refresh must not reach the IdP token endpoint" - ); - - mgr.set_system_sleep_imminent(false); - let fresh = mgr.auth().await.expect("refresh must succeed after wake"); - assert_eq!(fresh.key, "oidc-refreshed-token"); - assert_eq!( - token_hits.load(Ordering::SeqCst), - 1, - "exactly one IdP token call once the gate clears" - ); - - server.abort(); -} - -/// A refresh already in flight when sleep becomes imminent runs to completion -/// and persists its rotated token (no abort), proven through the real -/// `OidcRefresher` by holding the mock `/token` open until after the gate is -/// raised. The refresh token has already reached the IdP at that point, so -/// aborting would discard the rotated successor — the failure we guard against. -#[tokio::test] -async fn sleep_gate_e2e_in_flight_refresh_completes_across_imminent_sleep() { - let idp_hit = Arc::new(tokio::sync::Notify::new()); - let release = Arc::new(tokio::sync::Notify::new()); - let idp_hit_h = idp_hit.clone(); - let release_h = release.clone(); - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base_url = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let base_for_discovery = base_url.clone(); - let app = axum::Router::new() - .route( - "/.well-known/openid-configuration", - axum::routing::get(move || { - let b = base_for_discovery.clone(); - async move { - axum::Json(serde_json::json!({ - "authorization_endpoint": format!("{b}/authorize"), - "token_endpoint": format!("{b}/token"), - })) - } - }), - ) - .route( - "/token", - axum::routing::post(move || { - let idp_hit = idp_hit_h.clone(); - let release = release_h.clone(); - async move { - // Signal that the RT has reached the IdP, then block until - // released — this span is the in-flight window. - idp_hit.notify_one(); - release.notified().await; - axum::Json(serde_json::json!({ - "access_token": "oidc-refreshed-token", - "refresh_token": "oidc-new-rt", - "expires_in": 3600, - })) - } - }), - ) - .route( - "/user", - axum::routing::get(|| async { - axum::Json(serde_json::json!({ "userId": "user-42", "email": "test@corp.com" })) - }), - ); - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - - let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new( - AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url), - ); - mgr.hot_swap(expired_oidc_for(&base_url)); - mgr.set_refresher(Arc::new(OidcRefresher::new(mgr.clone()))); - - let m = mgr.clone(); - let handle = tokio::spawn(async move { m.auth().await }); - - idp_hit.notified().await; - - // `set_system_sleep_imminent` now holds the OS sleep ack until the in-flight - // refresh drains. Drive it off the runtime (as the real OS power-listener - // thread does) so the runtime can complete the refresh while it waits. - let sleeper = mgr.clone(); - let ack = std::thread::spawn(move || sleeper.set_system_sleep_imminent(true)); - release.notify_one(); - - let fresh = tokio::time::timeout(std::time::Duration::from_secs(5), handle) - .await - .expect("auth() must return") - .unwrap() - .expect("in-flight refresh must complete across imminent sleep"); - ack.join().expect("ack thread panicked"); - assert_eq!(fresh.key, "oidc-refreshed-token"); - assert_eq!( - mgr.current().map(|a| a.key), - Some("oidc-refreshed-token".to_owned()), - "the rotated token must be persisted, not discarded" - ); - - server.abort(); -} - -// ── Transient-blip budget is per-credential ───────────────────────── - -/// Minimal `AuthSnapshot` for exercising `record_transient_failure` in -/// isolation (it never reads credential state). -struct EmptySnapshot; -impl AuthSnapshot for EmptySnapshot { - fn current(&self) -> Option { - None - } - fn expired_auth(&self) -> Option { - None - } - fn read_disk_auth(&self) -> Option { - None - } - fn is_expired(&self) -> bool { - false - } -} - -/// A fresh credential (e.g. after re-login on this long-lived refresher) must -/// get the full blip budget instead of inheriting a dead credential's count, -/// so a valid token is never escalated to a permanent failure early. -#[test] -fn transient_blip_budget_is_scoped_to_the_credential() { - let refresher = OidcRefresher::new(Arc::new(EmptySnapshot)); - let key_a = Some("cred-a".to_owned()); - - // Accrue blips up to just under the escalation threshold on credential A. - for _ in 0..MAX_CONSECUTIVE_TRANSIENT_FAILURES - 1 { - assert!(matches!( - refresher.record_transient_failure("blip".into(), key_a.clone()), - RefreshOutcome::TransientFailure { .. } - )); - } - - // Credential B's first blip must stay transient, not escalate to permanent. - assert!( - matches!( - refresher.record_transient_failure("blip".into(), Some("cred-b".to_owned())), - RefreshOutcome::TransientFailure { .. } - ), - "a fresh credential must not inherit a prior credential's blip count", - ); -} diff --git a/crates/codegen/kigi-shell/src/auth/storage.rs b/crates/codegen/kigi-shell/src/auth/storage.rs index 15ad49c..1157791 100644 --- a/crates/codegen/kigi-shell/src/auth/storage.rs +++ b/crates/codegen/kigi-shell/src/auth/storage.rs @@ -2,7 +2,184 @@ use std::fs::File; use std::io::{Read, Write}; use std::path::{Path, PathBuf}; -use super::model::{API_KEY_SCOPE, AuthMode, AuthStore, GrokAuth, lookup_auth}; +use super::model::{API_KEY_SCOPE, AuthMode, AuthStore, KimiAuth, lookup_auth}; + +// ── System-keyring storage for the Kimi Code OAuth session ───────────── +// +// PRD F1: the OAuth token set lives in the system keyring (service `kigi`, +// entry `oauth/kimi-code`); when the keyring is unavailable we fall back to +// the file mechanism below (`auth.json`, owner-only, atomic writes). The +// official Kimi client's keyring entries (service `kimi-code`) and `~/.kimi` +// files are never touched. + +/// Keyring service name — deliberately distinct from the official client's +/// `kimi-code` service. +#[cfg(any(target_os = "macos", windows))] +pub(crate) const KEYRING_SERVICE: &str = "kigi"; + +/// Outcome of a keyring read for the session scope. +#[derive(Debug)] +pub(crate) enum KeyringRead { + /// Backend reachable and the entry exists. + Found(Box), + /// Backend reachable, no entry stored. + Missing, + /// Keyring disabled, unsupported on this platform, or the backend + /// errored — callers fall back to the file store. + Unavailable, +} + +/// Whether keyring storage participates for the session credential. +/// +/// Disabled when: +/// - the platform has no supported backend (non-macOS/Windows builds), +/// - `KIGI_DISABLE_KEYRING` is set to a truthy value, +/// - a non-default credential location is in use (`KIGI_SHARE_DIR` / +/// `KIGI_AUTH_PATH`): the keyring entry belongs to the default user +/// install; alternate profiles (and tests) stay file-scoped, and +/// - in unit-test builds, unless a test explicitly opted into the mock +/// keyring via [`enable_mock_keyring_for_test`]. +pub(crate) fn keyring_enabled() -> bool { + #[cfg(test)] + { + // Thread-local so keyring-specific tests (which opt in via + // `enable_mock_keyring_for_test`) can't leak the toggle into + // concurrently running persistence tests on other threads. + TEST_KEYRING_ENABLED.with(|flag| flag.get()) + } + #[cfg(not(test))] + { + #[cfg(not(any(target_os = "macos", windows)))] + { + false + } + #[cfg(any(target_os = "macos", windows))] + { + let disabled = std::env::var("KIGI_DISABLE_KEYRING") + .is_ok_and(|v| !matches!(v.trim(), "" | "0" | "false" | "off" | "no")); + !disabled + && std::env::var_os("KIGI_SHARE_DIR").is_none() + && std::env::var_os("KIGI_AUTH_PATH").is_none() + } + } +} + +#[cfg(test)] +thread_local! { + static TEST_KEYRING_ENABLED: std::cell::Cell = const { std::cell::Cell::new(false) }; +} + +/// Route keyring calls through an in-memory mock store for this process, +/// enable [`keyring_enabled`] on this thread, and clear any entry left by an +/// earlier test. Tests using this must serialize on the `kigi_keyring` key: +/// the mock entry is shared process-wide. +#[cfg(all(test, any(target_os = "macos", windows)))] +pub(crate) fn enable_mock_keyring_for_test() { + keyring::set_default_credential_builder(keyring::mock::default_credential_builder()); + TEST_KEYRING_ENABLED.with(|flag| flag.set(true)); + if let Err(e) = keyring_delete_session() { + panic!("mock keyring cleanup failed: {e}"); + } +} + +/// Disable the test keyring again (paired with +/// [`enable_mock_keyring_for_test`] in an RAII guard or test teardown). +#[cfg(test)] +pub(crate) fn disable_mock_keyring_for_test() { + TEST_KEYRING_ENABLED.with(|flag| flag.set(false)); +} + +/// The process-wide keyring entry handle. Cached so every reader/writer talks +/// to the same credential object: real backends read live state from the OS +/// store on each call, and the test mock keeps its state on the entry itself. +#[cfg(any(target_os = "macos", windows))] +fn keyring_entry() -> Result<&'static keyring::Entry, keyring::Error> { + static ENTRY: std::sync::OnceLock> = + std::sync::OnceLock::new(); + match ENTRY.get_or_init(|| { + keyring::Entry::new(KEYRING_SERVICE, crate::auth::config::KIMI_CODE_OAUTH_SCOPE) + }) { + Ok(entry) => Ok(entry), + // `keyring::Error` is not `Clone`; surface a stable equivalent. + Err(e) => { + tracing::warn!(error = %e, "auth: keyring entry construction failed"); + Err(keyring::Error::Invalid( + "keyring entry".into(), + e.to_string(), + )) + } + } +} + +/// Read the session credential from the system keyring. +#[cfg(any(target_os = "macos", windows))] +pub(crate) fn keyring_read_session() -> KeyringRead { + if !keyring_enabled() { + return KeyringRead::Unavailable; + } + let entry = match keyring_entry() { + Ok(entry) => entry, + Err(e) => { + tracing::warn!(error = %e, "auth: keyring entry unavailable, falling back to file"); + return KeyringRead::Unavailable; + } + }; + match entry.get_password() { + Ok(raw) => match serde_json::from_str::(&raw) { + Ok(auth) => KeyringRead::Found(Box::new(auth)), + Err(e) => { + tracing::warn!(error = %e, "auth: keyring entry is not valid JSON, ignoring"); + KeyringRead::Missing + } + }, + Err(keyring::Error::NoEntry) => KeyringRead::Missing, + Err(e) => { + tracing::warn!(error = %e, "auth: keyring read failed, falling back to file"); + KeyringRead::Unavailable + } + } +} + +#[cfg(not(any(target_os = "macos", windows)))] +pub(crate) fn keyring_read_session() -> KeyringRead { + KeyringRead::Unavailable +} + +/// Write the session credential to the system keyring. +#[cfg(any(target_os = "macos", windows))] +pub(crate) fn keyring_write_session(auth: &KimiAuth) -> anyhow::Result<()> { + anyhow::ensure!(keyring_enabled(), "keyring storage disabled"); + let payload = serde_json::to_string(auth)?; + keyring_entry()?.set_password(&payload)?; + tracing::info!("auth: session credential written to system keyring"); + Ok(()) +} + +#[cfg(not(any(target_os = "macos", windows)))] +pub(crate) fn keyring_write_session(_auth: &KimiAuth) -> anyhow::Result<()> { + anyhow::bail!("keyring storage is not supported on this platform") +} + +/// Delete the session credential from the system keyring (Ok when absent). +#[cfg(any(target_os = "macos", windows))] +pub(crate) fn keyring_delete_session() -> anyhow::Result<()> { + if !keyring_enabled() { + return Ok(()); + } + match keyring_entry()?.delete_credential() { + Ok(()) => { + tracing::info!("auth: session credential removed from system keyring"); + Ok(()) + } + Err(keyring::Error::NoEntry) => Ok(()), + Err(e) => Err(e.into()), + } +} + +#[cfg(not(any(target_os = "macos", windows)))] +pub(crate) fn keyring_delete_session() -> anyhow::Result<()> { + Ok(()) +} /// RAII guard for an exclusive advisory lock on `auth.json.lock`. /// The lock is released when the inner `File` is dropped (closing the FD). @@ -325,25 +502,23 @@ fn restore_prior_bytes(auth_file: &Path, bytes: &[u8]) -> std::io::Result<()> { } /// Read a single auth token from `auth.json` by scope key. -/// Falls back to the legacy `https://accounts.x.ai/sign-in` scope key -/// when the requested scope is not found (devbox auth.json migration). pub fn read_token_by_scope(kigi_home: &Path, scope: &str) -> anyhow::Result { let path = kigi_home.join("auth.json"); let store = - read_auth_json(&path).map_err(|_| anyhow::anyhow!("Not logged in. Run `grok login`."))?; + read_auth_json(&path).map_err(|_| anyhow::anyhow!("Not logged in. Run `kigi login`."))?; lookup_auth(&store, scope).map(|a| a.key).ok_or_else(|| { - anyhow::anyhow!("Your auth token is invalid. Run `grok login` to re-authenticate.") + anyhow::anyhow!("Your auth token is invalid. Run `kigi login` to re-authenticate.") }) } -/// Read the API key from the `xai::api_key` scope in auth.json. +/// Read the API key from the `kigi::api_key` scope in auth.json. pub fn read_api_key(kigi_home: &Path) -> Option { let path = kigi_home.join("auth.json"); let map = read_auth_json(&path).ok()?; map.get(API_KEY_SCOPE).map(|a| a.key.clone()) } -/// Store a plain API key in auth.json under the `xai::api_key` scope. +/// Store a plain API key in auth.json under the `kigi::api_key` scope. /// /// Uses the corrupt-recovery reader so a malformed auth.json (e.g. from a /// previous crash) can be healed when the user sets an API key. @@ -352,7 +527,7 @@ pub fn store_api_key(kigi_home: &Path, api_key: &str) -> std::io::Result<()> { let mut map = read_auth_json_or_empty_recovering_corrupt(&path)?; map.insert( API_KEY_SCOPE.to_owned(), - GrokAuth { + KimiAuth { key: api_key.to_owned(), auth_mode: AuthMode::ApiKey, ..Default::default() @@ -361,7 +536,7 @@ pub fn store_api_key(kigi_home: &Path, api_key: &str) -> std::io::Result<()> { write_auth_json(&path, &map) } -/// Remove the `xai::api_key` scope from auth.json. +/// Remove the `kigi::api_key` scope from auth.json. pub fn clear_api_key(kigi_home: &Path) -> std::io::Result<()> { let path = kigi_home.join("auth.json"); if let Ok(mut map) = read_auth_json(&path) { @@ -383,7 +558,7 @@ mod write_fallback_tests { let mut map = AuthStore::new(); map.insert( API_KEY_SCOPE.to_owned(), - GrokAuth { + KimiAuth { key: "secret-key".to_owned(), auth_mode: AuthMode::ApiKey, ..Default::default() @@ -481,7 +656,7 @@ mod write_fallback_tests { let mut replacement = AuthStore::new(); replacement.insert( API_KEY_SCOPE.to_owned(), - GrokAuth { + KimiAuth { key: "replacement-key".to_owned(), auth_mode: AuthMode::ApiKey, ..Default::default() @@ -510,3 +685,74 @@ mod write_fallback_tests { assert_eq!(mode & 0o777, 0o600, "restored file must stay 0o600"); } } + +#[cfg(all(test, any(target_os = "macos", windows)))] +mod keyring_tests { + use super::*; + use chrono::Utc; + + /// RAII teardown so a panicking test doesn't leave the process-global + /// test-keyring toggle enabled for later tests. + struct MockKeyringGuard; + impl MockKeyringGuard { + fn enable() -> Self { + enable_mock_keyring_for_test(); + Self + } + } + impl Drop for MockKeyringGuard { + fn drop(&mut self) { + disable_mock_keyring_for_test(); + } + } + + fn session_auth(key: &str, rt: &str) -> KimiAuth { + KimiAuth { + key: key.into(), + refresh_token: Some(rt.into()), + expires_at: Some(Utc::now() + chrono::Duration::seconds(3600)), + expires_in: Some(3600), + scope: Some("kimi-code".into()), + token_type: Some("bearer".into()), + ..KimiAuth::test_default() + } + } + + #[test] + #[serial_test::serial(kigi_keyring)] + fn keyring_session_roundtrip() { + let _guard = MockKeyringGuard::enable(); + // Fresh mock store: nothing there yet. + assert!(matches!(keyring_read_session(), KeyringRead::Missing)); + + keyring_write_session(&session_auth("at-1", "rt-1")).unwrap(); + let KeyringRead::Found(read) = keyring_read_session() else { + panic!("expected Found after write"); + }; + assert_eq!(read.key, "at-1"); + assert_eq!(read.refresh_token.as_deref(), Some("rt-1")); + assert_eq!(read.expires_in, Some(3600)); + + // Overwrite rotates in place. + keyring_write_session(&session_auth("at-2", "rt-2")).unwrap(); + let KeyringRead::Found(read) = keyring_read_session() else { + panic!("expected Found after rotate"); + }; + assert_eq!(read.key, "at-2"); + + // Delete is idempotent. + keyring_delete_session().unwrap(); + assert!(matches!(keyring_read_session(), KeyringRead::Missing)); + keyring_delete_session().unwrap(); + } + + #[test] + #[serial_test::serial(kigi_keyring)] + fn keyring_disabled_reads_unavailable() { + disable_mock_keyring_for_test(); + assert!(matches!(keyring_read_session(), KeyringRead::Unavailable)); + assert!(keyring_write_session(&session_auth("a", "r")).is_err()); + // Delete when disabled is a no-op success (logout stays best-effort). + keyring_delete_session().unwrap(); + } +} diff --git a/crates/codegen/kigi-shell/src/auth/token_type.rs b/crates/codegen/kigi-shell/src/auth/token_type.rs index 050d26d..814eb86 100644 --- a/crates/codegen/kigi-shell/src/auth/token_type.rs +++ b/crates/codegen/kigi-shell/src/auth/token_type.rs @@ -1,17 +1,13 @@ -use crate::auth::model::{AuthMode, GrokAuth}; +use crate::auth::model::{AuthMode, KimiAuth}; /// What kind of bearer is loaded right now. Dispatch key for /// `auth()`, `unauthorized_recovery()`, and proactive refresh. -/// -/// Not a session classifier — use `is_session_based_method` for that. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum TokenType { - /// OIDC/OAuth2 session with a refresh_token available. - OidcSession, - /// Legacy web-login session or OIDC without a refresh_token. - LegacySession, - /// External auth binary provides tokens. - ExternalBinary, + /// Kimi Code OAuth session with a refresh_token available. + OAuthSession, + /// OAuth session without a refresh_token (cannot be silently renewed). + SessionNoRefresh, /// Plain API key (no refresh possible). ApiKey, /// No credentials loaded. @@ -20,36 +16,56 @@ pub(crate) enum TokenType { impl TokenType { /// Classify the loaded credential (pure; no manager state). - pub(crate) fn from_auth(auth: Option<&GrokAuth>) -> Self { + pub(crate) fn from_auth(auth: Option<&KimiAuth>) -> Self { match auth { None => Self::None, - // Oidc without a refresh_token degrades to the unrefreshable LegacySession shape. Some(a) => match a.auth_mode { - AuthMode::Oidc if a.refresh_token.is_some() => Self::OidcSession, - AuthMode::Oidc | AuthMode::WebLogin => Self::LegacySession, - AuthMode::External => Self::ExternalBinary, + AuthMode::OAuth if a.refresh_token.is_some() => Self::OAuthSession, + AuthMode::OAuth => Self::SessionNoRefresh, AuthMode::ApiKey => Self::ApiKey, }, } } - /// `true` for types that can be silently refreshed (OIDC, external binary). + /// `true` for types that can be silently refreshed. pub(crate) fn is_refreshable(self) -> bool { - matches!(self, Self::OidcSession | Self::ExternalBinary) + matches!(self, Self::OAuthSession) } } #[cfg(test)] mod tests { - //! Per-variant matrix for `is_refreshable`. + //! Per-variant matrix for `is_refreshable` and classification. use super::*; #[test] fn is_refreshable_matrix() { - assert!(TokenType::OidcSession.is_refreshable()); - assert!(TokenType::ExternalBinary.is_refreshable()); - assert!(!TokenType::LegacySession.is_refreshable()); + assert!(TokenType::OAuthSession.is_refreshable()); + assert!(!TokenType::SessionNoRefresh.is_refreshable()); assert!(!TokenType::ApiKey.is_refreshable()); assert!(!TokenType::None.is_refreshable()); } + + #[test] + fn from_auth_classifies_by_mode_and_refresh_token() { + assert_eq!(TokenType::from_auth(None), TokenType::None); + let with_rt = KimiAuth { + refresh_token: Some("rt".into()), + ..KimiAuth::test_default() + }; + assert_eq!( + TokenType::from_auth(Some(&with_rt)), + TokenType::OAuthSession + ); + let no_rt = KimiAuth::test_default(); + assert_eq!( + TokenType::from_auth(Some(&no_rt)), + TokenType::SessionNoRefresh + ); + let api = KimiAuth { + auth_mode: AuthMode::ApiKey, + ..KimiAuth::test_default() + }; + assert_eq!(TokenType::from_auth(Some(&api)), TokenType::ApiKey); + } } diff --git a/crates/codegen/kigi-shell/src/cli_models.rs b/crates/codegen/kigi-shell/src/cli_models.rs index 6b1791d..976ef9f 100644 --- a/crates/codegen/kigi-shell/src/cli_models.rs +++ b/crates/codegen/kigi-shell/src/cli_models.rs @@ -22,28 +22,24 @@ impl AuthStatus { /// Banner status: env key → session → BYOK → deployment → none. /// /// Differs from sampling (`resolve_credentials`: BYOK → session → env) so a - /// logged-in user sees the login host. BYOK uses - /// [`crate::agent::auth_method::should_advertise_xai_api_key`] so - /// `disable_api_key_auth` is honored. + /// logged-in user sees the login host. pub fn resolve(agent_config: &AgentConfig) -> Self { if crate::agent::auth_method::has_xai_api_key_env() { return Self::ApiKey; } if agent_config.create_auth_manager().current().is_some() { - let origin = &agent_config.grok_com_config.grok_ws_origin; + let origin = kigi_env::oauth_host(); let host = origin .strip_prefix("https://") .or_else(|| origin.strip_prefix("http://")) - .unwrap_or(origin); + .unwrap_or(&origin); return Self::LoggedIn(host.to_owned()); } let models = crate::agent::config::resolve_model_list(agent_config, None); - if crate::agent::auth_method::should_advertise_xai_api_key( - agent_config.grok_com_config.api_key_auth_disabled(), - models.values(), - ) && let Some(name) = models - .iter() - .find_map(|(name, entry)| entry.has_own_credentials().then(|| name.clone())) + if crate::agent::auth_method::should_advertise_xai_api_key(models.values()) + && let Some(name) = models + .iter() + .find_map(|(name, entry)| entry.has_own_credentials().then(|| name.clone())) { return Self::ModelCredentials(name); } @@ -94,7 +90,7 @@ mod tests { use super::*; use crate::agent::auth_method::{LEGACY_XAI_API_KEY_ENV_VAR, XAI_API_KEY_ENV_VAR}; use crate::agent::config::Config; - use crate::auth::{AuthMode, GrokAuth}; + use crate::auth::{AuthMode, KimiAuth}; use kigi_test_support::EnvGuard; use serial_test::serial; @@ -155,17 +151,17 @@ mod tests { #[serial] fn resolve_oauth_session() { let (_dir, _g) = isolate_auth_sources(); - let token = GrokAuth { + let token = KimiAuth { key: "session-token".into(), - auth_mode: AuthMode::WebLogin, - ..GrokAuth::test_default() + auth_mode: AuthMode::OAuth, + ..KimiAuth::test_default() }; let json = serde_json::to_string(&token).unwrap(); let _auth = EnvGuard::set("KIGI_AUTH", &json); assert_eq!( AuthStatus::resolve(&Config::default()), - AuthStatus::LoggedIn("grok.com".to_owned()) + AuthStatus::LoggedIn("auth.kimi.com".to_owned()) ); } @@ -247,10 +243,10 @@ mod tests { #[serial] fn resolve_priority_session_over_byok_and_deployment() { let (_dir, _g) = isolate_auth_sources(); - let token = GrokAuth { + let token = KimiAuth { key: "session-token".into(), - auth_mode: AuthMode::WebLogin, - ..GrokAuth::test_default() + auth_mode: AuthMode::OAuth, + ..KimiAuth::test_default() }; let json = serde_json::to_string(&token).unwrap(); let _auth = EnvGuard::set("KIGI_AUTH", &json); @@ -259,7 +255,7 @@ mod tests { let cfg = config_from_toml(&byok_and_deployment_toml(dm)); assert_eq!( AuthStatus::resolve(&cfg), - AuthStatus::LoggedIn("grok.com".to_owned()) + AuthStatus::LoggedIn("auth.kimi.com".to_owned()) ); } @@ -275,45 +271,6 @@ mod tests { ); } - #[test] - #[serial] - fn resolve_disable_api_key_auth_suppresses_byok_banner() { - let (_dir, _g) = isolate_auth_sources(); - let dm = crate::models::default_model(); - let cfg = config_from_toml(&format!( - r#" - [grok_com_config] - disable_api_key_auth = true - - [model."{dm}"] - model = "{dm}" - api_key = "sk-byok" - "# - )); - assert_eq!(AuthStatus::resolve(&cfg), AuthStatus::NotAuthenticated); - } - - #[test] - #[serial] - fn resolve_disable_api_key_auth_falls_through_to_deployment() { - let (_dir, _g) = isolate_auth_sources(); - let dm = crate::models::default_model(); - let cfg = config_from_toml(&format!( - r#" - [grok_com_config] - disable_api_key_auth = true - - [endpoints] - deployment_key = "deploy-key" - - [model."{dm}"] - model = "{dm}" - api_key = "sk-byok" - "# - )); - assert_eq!(AuthStatus::resolve(&cfg), AuthStatus::DeploymentKey); - } - #[test] #[serial] fn resolve_model_credentials_uses_first_catalog_key() { diff --git a/crates/codegen/kigi-shell/src/config/reloader.rs b/crates/codegen/kigi-shell/src/config/reloader.rs index c15c3d2..0baa59f 100644 --- a/crates/codegen/kigi-shell/src/config/reloader.rs +++ b/crates/codegen/kigi-shell/src/config/reloader.rs @@ -7,7 +7,7 @@ use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; use tracing::{debug, error, info}; -use crate::auth::{GrokAuth, read_auth_json}; +use crate::auth::{KimiAuth, read_auth_json}; use super::watcher::ConfigChangeEvent; @@ -15,7 +15,7 @@ use super::watcher::ConfigChangeEvent; #[derive(Debug)] pub enum ConfigUpdate { /// New auth credentials from disk. - Auth(Box), + Auth(Box), /// Auth scope was removed (user logged out). AuthCleared, /// A **broadcast** MCP reload — applies to every active session @@ -535,14 +535,14 @@ fn extract_ui_fields(config: &toml::Value) -> (Option, bool, Option GrokAuth { - GrokAuth { + fn make_auth(key: &str) -> KimiAuth { + KimiAuth { key: key.to_string(), email: Some("test@test.com".to_string()), - ..GrokAuth::test_default() + ..KimiAuth::test_default() } } @@ -604,7 +604,7 @@ mod tests { reloader.reload_auth().unwrap(); let update = rx.try_recv().expect("should send Auth update"); assert!( - matches!(update, ConfigUpdate::Auth(a) if a.key == "new-key"), // a is Box, Deref coercion + matches!(update, ConfigUpdate::Auth(a) if a.key == "new-key"), // a is Box, Deref coercion "should contain new key" ); } diff --git a/crates/codegen/kigi-shell/src/extensions/auth.rs b/crates/codegen/kigi-shell/src/extensions/auth.rs index 166f89e..bcde58f 100644 --- a/crates/codegen/kigi-shell/src/extensions/auth.rs +++ b/crates/codegen/kigi-shell/src/extensions/auth.rs @@ -21,7 +21,6 @@ pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { "x.ai/auth/get_url" => handle_get_url(agent).await, "x.ai/auth/logout" => handle_logout(agent, args).await, "x.ai/auth/info" => handle_info(agent), - "x.ai/auth/check_subscription" => handle_check_subscription(agent).await, _ => Err(acp::Error::method_not_found()), } } @@ -111,7 +110,7 @@ async fn handle_get_url(agent: &MvpAgent) -> ExtResult { to_raw_response(&serde_json::json!({ "auth_url": auth_url, // `external_provider` kept for older clients; `mode` is authoritative. - "external_provider": mode.is_some_and(|m| m.is_external_provider()), + "external_provider": false, "mode": mode.map(|m| m.as_wire_str()), })) } @@ -141,43 +140,16 @@ async fn handle_logout(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { })) } -/// Single-shot subscription re-check (retry button on paywall screen). -/// -/// Calls `retry_subscription_check()`, then returns the updated auth -/// response with gate info so the pager can refresh the gate state. -async fn handle_check_subscription(agent: &MvpAgent) -> ExtResult { - agent.retry_subscription_check().await; - let response = agent.auth_response_with_meta(); - to_raw_response(&serde_json::json!({ - "authenticated": response.meta.is_some(), - "meta": response.meta, - })) -} - -/// Returns current auth method ID, user profile fields, and team/principal -/// metadata. +/// Returns current auth method ID and the account fields the Kimi flow +/// exposes (email/user id are empty until a later feature surfaces them). fn handle_info(agent: &MvpAgent) -> ExtResult { #[derive(Serialize)] #[serde(rename_all = "camelCase")] struct AuthInfoResponse { method_id: Option, email: Option, - first_name: Option, - last_name: Option, - /// `grok-asset://` URL resolved by the Electron protocol handler, - /// or a full `http(s)://` URL passed through unchanged. - profile_image_url: Option, - team_id: Option, - team_name: Option, - team_role: Option, - organization_id: Option, - organization_name: Option, - organization_role: Option, - principal_type: Option, - principal_id: Option, - user_blocked_reason: Option, - team_blocked_reasons: Vec, - coding_data_retention_opt_out: bool, + user_id: Option, + auth_mode: Option, } let method_id = agent @@ -186,40 +158,13 @@ fn handle_info(agent: &MvpAgent) -> ExtResult { .as_ref() .map(|m| m.0.to_string()); let auth = agent.auth_manager.current(); - let raw_asset_id = auth.as_ref().and_then(|a| a.profile_image_asset_id.clone()); - - // Return a grok-asset:// URL that the Electron renderer resolves at - // display time via a custom protocol handler. The handler proxies - // through cli-chat-proxy's /asset endpoint; Electron's HTTP cache - // handles reuse. No disk-cache or network call needed here. - let profile_image_url = match raw_asset_id.as_deref().filter(|k| !k.is_empty()) { - Some(key) if key.starts_with("http://") || key.starts_with("https://") => { - Some(key.to_owned()) - } - Some(key) => Some(format!("grok-asset:///{key}")), - None => None, - }; to_raw_response(&AuthInfoResponse { method_id, email: auth.as_ref().and_then(|a| a.email.clone()), - first_name: auth.as_ref().and_then(|a| a.first_name.clone()), - last_name: auth.as_ref().and_then(|a| a.last_name.clone()), - profile_image_url, - team_id: auth.as_ref().and_then(|a| a.team_id.clone()), - team_name: auth.as_ref().and_then(|a| a.team_name.clone()), - team_role: auth.as_ref().and_then(|a| a.team_role.clone()), - organization_id: auth.as_ref().and_then(|a| a.organization_id.clone()), - organization_name: auth.as_ref().and_then(|a| a.organization_name.clone()), - organization_role: auth.as_ref().and_then(|a| a.organization_role.clone()), - principal_type: auth.as_ref().and_then(|a| a.principal_type.clone()), - principal_id: auth.as_ref().and_then(|a| a.principal_id.clone()), - user_blocked_reason: auth.as_ref().and_then(|a| a.user_blocked_reason.clone()), - team_blocked_reasons: auth + user_id: auth .as_ref() - .map(|a| a.team_blocked_reasons.clone()) - .unwrap_or_default(), - coding_data_retention_opt_out: auth - .as_ref() - .is_some_and(|a| a.coding_data_retention_opt_out), + .map(|a| a.user_id.clone()) + .filter(|id| !id.is_empty()), + auth_mode: auth.as_ref().map(|a| format!("{:?}", a.auth_mode)), }) } diff --git a/crates/codegen/kigi-shell/src/extensions/auth_gate.rs b/crates/codegen/kigi-shell/src/extensions/auth_gate.rs index 0cc34c2..5a4d566 100644 --- a/crates/codegen/kigi-shell/src/extensions/auth_gate.rs +++ b/crates/codegen/kigi-shell/src/extensions/auth_gate.rs @@ -1,17 +1,17 @@ use agent_client_protocol as acp; -use crate::auth::{AuthManager, GrokAuth}; +use crate::auth::{AuthManager, KimiAuth}; -/// Require xAI auth from a sync context, accepting tokens in the client-side buffer window. +/// Require a Kimi Code session from a sync context, accepting tokens in the client-side buffer window. pub(crate) fn require_xai_auth( auth_manager: &AuthManager, missing_message: &'static str, non_xai_message: &'static str, -) -> Result { +) -> Result { let auth = auth_manager .current_or_expired() .ok_or_else(|| acp::Error::auth_required().data(missing_message))?; - if !auth.is_xai_auth() { + if !auth.is_session_auth() { return Err(acp::Error::auth_required().data(non_xai_message)); } Ok(auth) diff --git a/crates/codegen/kigi-shell/src/extensions/billing.rs b/crates/codegen/kigi-shell/src/extensions/billing.rs index 2c85cf0..2f4b1f2 100644 --- a/crates/codegen/kigi-shell/src/extensions/billing.rs +++ b/crates/codegen/kigi-shell/src/extensions/billing.rs @@ -213,10 +213,6 @@ async fn handle_get_billing(agent: &MvpAgent) -> ExtResult { let credits_resp = crate::http::shared_client() .get(&credits_url) .header("Authorization", format!("Bearer {}", auth.key)) - .header( - "X-XAI-Token-Auth", - crate::auth::GrokComConfig::default().token_header, - ) .header("x-userid", &auth.user_id) .header("x-grok-client-version", kigi_version::VERSION) .header( @@ -304,10 +300,6 @@ async fn handle_get_auto_topup_rule(agent: &MvpAgent) -> ExtResult { let response = crate::http::shared_client() .get(&url) .header("Authorization", format!("Bearer {}", auth.key)) - .header( - "X-XAI-Token-Auth", - crate::auth::GrokComConfig::default().token_header, - ) .header("x-userid", &auth.user_id) .header("x-grok-client-version", kigi_version::VERSION) .header( diff --git a/crates/codegen/kigi-shell/src/extensions/bundle.rs b/crates/codegen/kigi-shell/src/extensions/bundle.rs index 5c11740..0a0267c 100644 --- a/crates/codegen/kigi-shell/src/extensions/bundle.rs +++ b/crates/codegen/kigi-shell/src/extensions/bundle.rs @@ -484,37 +484,23 @@ mod tests { .insert("review".to_string(), "# Review skill\n".to_string()); bundle } - fn test_auth() -> crate::auth::GrokAuth { - crate::auth::GrokAuth { + fn test_auth() -> crate::auth::KimiAuth { + crate::auth::KimiAuth { key: "token".to_string(), - auth_mode: crate::auth::AuthMode::Oidc, + auth_mode: crate::auth::AuthMode::OAuth, create_time: chrono::Utc::now(), user_id: "user-1".to_string(), email: Some("test@example.com".to_string()), - first_name: None, - last_name: None, - profile_image_asset_id: None, - principal_type: None, - principal_id: None, - team_id: None, - team_name: None, - team_role: None, - organization_id: None, - organization_name: None, - organization_role: None, - user_blocked_reason: None, - team_blocked_reasons: vec![], - coding_data_retention_opt_out: false, - has_grok_code_access: None, refresh_token: None, expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), - oidc_issuer: None, - oidc_client_id: None, + expires_in: Some(3600), + scope: None, + token_type: None, } } fn test_auth_manager() -> Arc { let dir = tempfile::tempdir().unwrap(); - let mgr = crate::auth::AuthManager::new(dir.path(), crate::auth::GrokComConfig::default()); + let mgr = crate::auth::AuthManager::new(dir.path(), crate::auth::KimiCodeConfig::default()); mgr.hot_swap(test_auth()); std::mem::forget(dir); Arc::new(mgr) diff --git a/crates/codegen/kigi-shell/src/extensions/mod.rs b/crates/codegen/kigi-shell/src/extensions/mod.rs index 6c34447..9dca80e 100644 --- a/crates/codegen/kigi-shell/src/extensions/mod.rs +++ b/crates/codegen/kigi-shell/src/extensions/mod.rs @@ -17,7 +17,6 @@ pub mod memory; pub mod notification; pub mod plugins; pub mod pr; -pub mod privacy; pub mod prompt_history; pub mod prompt_meta; pub mod recap; diff --git a/crates/codegen/kigi-shell/src/extensions/privacy.rs b/crates/codegen/kigi-shell/src/extensions/privacy.rs deleted file mode 100644 index 2d41e7c..0000000 --- a/crates/codegen/kigi-shell/src/extensions/privacy.rs +++ /dev/null @@ -1,91 +0,0 @@ -//! `x.ai/privacy/setCodingDataRetention` extension handler. -//! -//! PUTs the new opt-out flag to cli-chat-proxy and updates local auth state -//! to match. The local update is fire-and-forget (best-effort cache refresh). - -use agent_client_protocol as acp; -use serde::Deserialize; - -use super::{ExtResult, parse_params, to_raw_response}; -use crate::agent::MvpAgent; - -#[tracing::instrument(skip_all, fields(method = %args.method))] -pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { - match args.method.as_ref() { - "x.ai/privacy/setCodingDataRetention" => handle_set(agent, args).await, - _ => Err(acp::Error::method_not_found()), - } -} - -async fn handle_set(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { - #[derive(Deserialize)] - #[serde(rename_all = "camelCase")] - struct Params { - coding_data_retention_opt_out: bool, - } - - let params: Params = parse_params(args)?; - - let auth = agent.auth_manager.auth().await.map_err(|e| { - tracing::warn!(error = %e, "privacy: auth resolution failed"); - acp::Error::auth_required() - .data("Authentication required. Run `grok login` to re-authenticate.") - })?; - - let proxy_url = agent.cfg.borrow().endpoints.proxy_url(); - let url = format!("{proxy_url}/privacy/coding-data-retention"); - let token_header = agent.auth_manager.grok_com_config().token_header.clone(); - - let body = serde_json::json!({ - "codingDataRetentionOptOut": params.coding_data_retention_opt_out, - }); - - let provider: std::sync::Arc = std::sync::Arc::new( - crate::auth::credential_provider::ShellAuthCredentialProvider::new( - agent.auth_manager.clone(), - None, - None, - ), - ); - let client = crate::http::with_auth_retry(crate::http::shared_client(), provider); - - let resp = client - .put(&url) - .header("X-XAI-Token-Auth", &token_header) - .header("x-grok-client-version", kigi_version::VERSION) - .header( - crate::http::CLIENT_MODE_HEADER, - crate::http::process_client_mode(), - ) - .json(&body) - .send() - .await - .map_err(|e| acp::Error::internal_error().data(format!("HTTP request failed: {e}")))?; - - if !resp.status().is_success() { - let status = resp.status().as_u16(); - let body = resp.text().await.unwrap_or_default(); - tracing::warn!(status, "setCodingDataRetention request failed"); - let friendly = serde_json::from_str::(&body) - .ok() - .and_then(|v| { - v.get("error") - .or_else(|| v.get("message")) - .and_then(|e| e.as_str().map(String::from)) - }) - .unwrap_or_else(|| format!("server returned HTTP {status}")); - return Err(acp::Error::internal_error().data(friendly)); - } - - // Update local auth state to reflect the change. - // Use save_without_enrichment to avoid a race: update() spawns a - // background GET /user enrichment that may read stale ACL state - // and overwrite the opt-out flag back to its previous value. - let mut updated = auth.clone(); - updated.coding_data_retention_opt_out = params.coding_data_retention_opt_out; - let _ = agent.auth_manager.save_without_enrichment(updated).await; - - to_raw_response(&serde_json::json!({ - "codingDataRetentionOptOut": params.coding_data_retention_opt_out, - })) -} diff --git a/crates/codegen/kigi-shell/src/extensions/session_admin.rs b/crates/codegen/kigi-shell/src/extensions/session_admin.rs index dd4adeb..3d23b2d 100644 --- a/crates/codegen/kigi-shell/src/extensions/session_admin.rs +++ b/crates/codegen/kigi-shell/src/extensions/session_admin.rs @@ -111,10 +111,7 @@ async fn handle_session_rename(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR // Send a SessionSummaryGenerated notification so the TUI updates its title notify_session_title(agent, session_id, &req.title).await; - if agent.is_writeback_storage() - && let Some(auth) = agent.current_auth() - && !auth.is_zdr_team() - { + if agent.is_writeback_storage() && agent.current_auth().is_some() { use crate::remote::client::BackendClient; use crate::session::export::ExportedMetadata; @@ -133,15 +130,7 @@ async fn handle_session_rename(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR // Hook 2: update session replica with summary (fire-and-forget) if let Some(client) = agent.session_registry_client() { let sid = req.session_id.to_string(); - let title = if agent - .auth_manager - .current_or_expired() - .is_some_and(|a| a.is_zdr_team()) - { - None - } else { - Some(req.title.clone()) - }; + let title = Some(req.title.clone()); tokio::spawn(async move { let update = crate::agent::session_registry_client::UpdateRequest { summary: title, @@ -248,8 +237,7 @@ async fn handle_session_delete(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR // For writeback storage (non-ZDR): remote delete is authoritative for // the cloud history and runs first; on failure no local bits are // touched so the pager does not remove the row or toast success. - let needs_remote = - agent.is_writeback_storage() && agent.current_auth().is_some_and(|a| !a.is_zdr_team()); + let needs_remote = agent.is_writeback_storage() && agent.current_auth().is_some(); // Shared delete: remote-first, then local disk + FTS eviction. // Mirrored by the `grok sessions delete ` CLI path. diff --git a/crates/codegen/kigi-shell/src/extensions/share.rs b/crates/codegen/kigi-shell/src/extensions/share.rs index 3d36d97..cc743a4 100644 --- a/crates/codegen/kigi-shell/src/extensions/share.rs +++ b/crates/codegen/kigi-shell/src/extensions/share.rs @@ -45,13 +45,6 @@ async fn handle_share_session(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtRe ); } - // Only block for ZDR teams (hard data-retention policy), not for - // coding-data-retention opt-out — sharing is user-initiated. - if auth.is_zdr_team() { - return Err(acp::Error::invalid_params() - .data("Session sharing is disabled for your team's data retention policy")); - } - // Find session info by searching through summaries let summaries = list_summaries(None).await.map_err(|e| { acp::Error::internal_error().data(format!("Failed to list sessions: {}", e)) @@ -94,7 +87,7 @@ async fn handle_share_session(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtRe fn require_xai_auth_for_share( auth_manager: &crate::auth::AuthManager, -) -> Result { +) -> Result { super::auth_gate::require_xai_auth( auth_manager, "Authentication required to share session", @@ -105,8 +98,8 @@ fn require_xai_auth_for_share( #[cfg(test)] mod tests { use super::*; - use crate::auth::GrokComConfig; - use crate::auth::{AuthMode, GrokAuth}; + use crate::auth::KimiCodeConfig; + use crate::auth::{AuthMode, KimiAuth}; use chrono::{Duration, Utc}; use std::sync::Arc; use tempfile::tempdir; @@ -117,7 +110,7 @@ mod tests { let dir = tempdir().expect("tempdir for share auth test"); let mgr = Arc::new(crate::auth::AuthManager::new( dir.path(), - GrokComConfig::default(), + KimiCodeConfig::default(), )); let expires_at = Utc::now() + ttl; @@ -126,9 +119,8 @@ mod tests { // Only OIDC tokens against https://auth.x.ai (or the local-dev equivalent) // return true from is_xai_auth(). This is required for the share tests to // exercise the happy path through require_xai_auth_for_share. - let auth = GrokAuth { - auth_mode: AuthMode::Oidc, - oidc_issuer: Some("https://auth.x.ai".to_string()), + let auth = KimiAuth { + auth_mode: AuthMode::OAuth, key: "test-key".into(), expires_at: Some(expires_at), create_time: Utc::now() - Duration::hours(1), @@ -168,7 +160,7 @@ mod tests { let dir = tempdir().expect("tempdir"); let mgr = Arc::new(crate::auth::AuthManager::new( dir.path(), - GrokComConfig::default(), + KimiCodeConfig::default(), )); assert!(require_xai_auth_for_share(&mgr).is_err()); } @@ -178,12 +170,12 @@ mod tests { let dir = tempdir().expect("tempdir"); let mgr = Arc::new(crate::auth::AuthManager::new( dir.path(), - GrokComConfig::default(), + KimiCodeConfig::default(), )); // API key is the simplest non-xAI credential (External and enterprise OIDC // are also rejected the same way). - let non_xai = GrokAuth { + let non_xai = KimiAuth { auth_mode: AuthMode::ApiKey, key: "xai-test-key".into(), create_time: Utc::now(), diff --git a/crates/codegen/kigi-shell/src/inspect/mod.rs b/crates/codegen/kigi-shell/src/inspect/mod.rs index 58e80a4..42a60a4 100644 --- a/crates/codegen/kigi-shell/src/inspect/mod.rs +++ b/crates/codegen/kigi-shell/src/inspect/mod.rs @@ -17,7 +17,6 @@ use std::path::{Path, PathBuf}; use serde::Serialize; -use crate::auth::ForceLoginTeam; use kigi_tools::types::config_source::ConfigSource; use kigi_tools::util::truncate::estimate_tokens; @@ -63,7 +62,6 @@ pub struct InspectReport { pub project_trusted: bool, pub project_instructions: Vec, pub permissions: PermissionsReport, - pub login_policy: LoginPolicyReport, pub hooks: Vec, pub skills: Vec, pub agents: Vec, @@ -139,20 +137,6 @@ pub struct SkippedRule { pub reason: String, } -/// Enterprise login-hardening policy resolved from `[grok_com_config]` -/// (TOML + env). Surfaced so admins can verify the deployment loaded it. -/// The team pin is admin policy, not a secret, so it is shown verbatim. -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct LoginPolicyReport { - /// Raw `disable_api_key_auth` knob (env `KIGI_DISABLE_API_KEY_AUTH`). - pub disable_api_key_auth: Option, - /// Configured team pin: single string, list, or null when unset. - pub force_login_team_uuid: Option, - /// Resolved verdict — true when either knob forces first-party login. - pub api_key_auth_disabled: bool, -} - #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] pub struct HookEntry { @@ -395,7 +379,6 @@ async fn build_report(cwd: &Path) -> InspectReport { project_trusted, project_instructions: instructions, permissions, - login_policy: login_policy_report(parsed_config.as_ref()), hooks, skills, agents, @@ -615,20 +598,6 @@ async fn list_permissions(cwd: &Path) -> PermissionsReport { } } -/// Resolves the enterprise login-hardening knobs from the merged config -/// (`[grok_com_config]`, the `[auth]` alias, and env overrides) so admins can -/// confirm the deployment's auth policy actually loaded. -fn login_policy_report(config: Option<&crate::agent::config::Config>) -> LoginPolicyReport { - let grok_com_config = config - .map(|c| c.grok_com_config.clone()) - .unwrap_or_default(); - LoginPolicyReport { - api_key_auth_disabled: grok_com_config.api_key_auth_disabled(), - disable_api_key_auth: grok_com_config.disable_api_key_auth, - force_login_team_uuid: grok_com_config.force_login_team_uuid, - } -} - /// Discovers hooks with every vendor enabled so compatibility can be annotated later. fn list_hooks( git_root: Option<&Path>, @@ -1174,19 +1143,6 @@ fn print_columns( } } -/// Render the team pin for the human view: single value, comma-joined list, -/// or an explicit empty-list marker (which fails closed at login). -fn format_force_login_team(team: &Option) -> String { - match team { - None => "(none)".to_string(), - Some(ForceLoginTeam::Single(s)) => s.clone(), - Some(ForceLoginTeam::AnyOf(list)) if list.is_empty() => { - "(empty -- fail closed)".to_string() - } - Some(ForceLoginTeam::AnyOf(list)) => list.join(", "), - } -} - /// Human label for an enforced setting. Uses product vocabulary, not the /// internal field names (no `ui.yolo` / `--yolo` / `permission_mode`). fn enforced_label(p: &EnforcedPolicy) -> String { @@ -1342,24 +1298,6 @@ fn print_human(r: &InspectReport) { } } - println!(); - println!(" Login Policy"); - println!( - " {TREE} disable_api_key_auth: {}", - match r.login_policy.disable_api_key_auth { - Some(v) => v.to_string(), - None => "(unset)".to_string(), - } - ); - println!( - " {TREE} force_login_team_uuid: {}", - format_force_login_team(&r.login_policy.force_login_team_uuid) - ); - println!( - " {TREE} api_key_auth_disabled: {}", - r.login_policy.api_key_auth_disabled - ); - print_columns( "Skills", &r.skills, diff --git a/crates/codegen/kigi-shell/src/leader/server.rs b/crates/codegen/kigi-shell/src/leader/server.rs index 4711da3..33d4dd4 100644 --- a/crates/codegen/kigi-shell/src/leader/server.rs +++ b/crates/codegen/kigi-shell/src/leader/server.rs @@ -210,24 +210,14 @@ impl AuthProvider for LeaderAuthProvider { AuthCredential::bearer(token) } /// Owner identity from the leader's `AuthManager`, surfaced on the auth - /// provider instead of a separate auth.json - /// read. Mirrors the in-process path (`mvp_agent`): prefer `GrokAuth.team_id` - /// (what shell telemetry/snapshot use) mapped onto a `"Team"` principal so - /// team attribution is derived; otherwise pass principal fields through. - /// `None` when no credential is available (identity resolution never blocks). + /// provider instead of a separate auth.json read. The Kimi credential + /// carries no principal metadata; only the (possibly empty) user id. fn identity(&self) -> Option { let a = self.auth_manager.current_or_expired()?; - Some(match a.team_id.filter(|t| !t.is_empty()) { - Some(team) => AuthIdentity { - user_id: a.user_id, - principal_type: Some("Team".to_string()), - principal_id: Some(team), - }, - None => AuthIdentity { - user_id: a.user_id, - principal_type: a.principal_type, - principal_id: a.principal_id, - }, + Some(AuthIdentity { + user_id: a.user_id, + principal_type: None, + principal_id: None, }) } } diff --git a/crates/codegen/kigi-shell/src/lib.rs b/crates/codegen/kigi-shell/src/lib.rs index 9a6b453..50d1bdb 100644 --- a/crates/codegen/kigi-shell/src/lib.rs +++ b/crates/codegen/kigi-shell/src/lib.rs @@ -34,7 +34,6 @@ pub mod session; pub mod terminal; #[cfg(test)] pub(crate) mod test_support; -pub mod tier; pub mod tools; pub mod trace_classifier; pub mod util; diff --git a/crates/codegen/kigi-shell/src/managed_config.rs b/crates/codegen/kigi-shell/src/managed_config.rs index f343c46..db18ede 100644 --- a/crates/codegen/kigi-shell/src/managed_config.rs +++ b/crates/codegen/kigi-shell/src/managed_config.rs @@ -3,7 +3,7 @@ mod response; -use crate::auth::GrokAuth; +use crate::auth::KimiAuth; pub use response::ManagedConfigError; use response::{ApplyOutcome, ManagedConfigResponse, ManagedConfigSource, verify_signed_envelope}; @@ -76,35 +76,10 @@ fn remove_managed_path(path: &std::path::Path) -> std::io::Result { } } -/// A team principal is eligible to fetch only if non-expired (an expired token -/// would just 401). -fn eligible_team_principal(auth: GrokAuth) -> Option { - (auth.is_team_principal() && !crate::auth::is_expired(&auth)).then_some(auth) -} - -/// The eligible team principal in `auth.json`, or `None`. Single-team: managed -/// config is a grok.com feature with one grok.com auth. -fn read_active_team_auth() -> Option { - let home = crate::util::kigi_home::kigi_home(); - let store = crate::auth::read_auth_json(&home.join("auth.json")).ok()?; - let team = store.values().find(|a| a.is_team_principal())?.clone(); - eligible_team_principal(team) -} - +/// Team principals were an xAI concept; the Kimi Code auth model has none, +/// so no team credential can ever serve managed config. pub(crate) fn has_active_team_auth() -> bool { - read_active_team_auth().is_some() -} - -/// Whether any team principal is signed in, **ignoring expiry** (a cold-start -/// expired token is not a logout). `Err` = `auth.json` unreadable: callers must -/// NOT treat that as a logout — it would wipe enforced policy on a read blip. -fn team_principal_signed_in() -> std::io::Result { - let home = crate::util::kigi_home::kigi_home(); - match crate::auth::read_auth_json(&home.join("auth.json")) { - Ok(store) => Ok(store.values().any(|a| a.is_team_principal())), - Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(false), - Err(e) => Err(e), - } + false } /// Clear the synced files when no principal could own them: no deployment key @@ -115,14 +90,6 @@ pub fn clear_orphan() { if resolve_deployment_key().is_some() { return; } - match team_principal_signed_in() { - Ok(true) => return, - Ok(false) => {} - Err(e) => { - tracing::warn!(error = %e, "auth.json unreadable; keeping managed config until it recovers"); - return; - } - } let home = crate::util::kigi_home::kigi_home(); let Some(_lock) = try_lock_managed_config(&home) else { return; // another process is syncing; retry next call @@ -427,7 +394,7 @@ pub fn is_fetch_enabled() -> bool { /// Fetch managed config + requirements and write to `~/.kigi/`, trying the /// deployment key first, then a signed-in team. `Ok(false)` when neither applies. pub async fn sync() -> Result { - Ok(sync_with_budget(SyncBudget::Standard, None).await?.wrote) + Ok(sync_with_budget(SyncBudget::Standard).await?.wrote) } struct SyncOutcome { @@ -479,25 +446,16 @@ impl SyncOutcome { /// Runs a sync under `budget`'s deadline, returning `None` when the deadline /// elapses first. -async fn sync_bounded( - budget: SyncBudget, - team_override: Option, -) -> Option> { - let sync = sync_with_budget(budget, team_override); +async fn sync_bounded(budget: SyncBudget) -> Option> { + let sync = sync_with_budget(budget); match budget.deadline() { Some(deadline) => tokio::time::timeout(deadline, sync).await.ok(), None => Some(sync.await), } } -/// `team_override` pins a specific team principal (the just-authenticated one, -/// post-login) instead of re-deriving the team from `auth.json`; `None` uses -/// [`read_active_team_auth`] (the current eligible team). -async fn sync_with_budget( - budget: SyncBudget, - team_override: Option, -) -> Result { - let outcome = sync_inner(budget, team_override).await?; +async fn sync_with_budget(budget: SyncBudget) -> Result { + let outcome = sync_inner(budget).await?; // Mark only when a principal was consulted AND the fetch wasn't signature-rejected — // a rejected fetch persisted nothing, so marking would claim an unwritten body. Lock // contention still marks (the holder persists the same config). @@ -519,73 +477,31 @@ enum FetchedConfig { key: String, body: ManagedConfigResponse, }, - Team { - auth: Box, - body: ManagedConfigResponse, - }, - /// No deployment key configured and no eligible team signed in. + /// No deployment key configured. NoPrincipal, } /// Fetches the configuration for the current principal without touching disk: /// the deployment key first, then a signed-in team. The installing sync and the /// read-only `grok setup --json` both build on this. -async fn fetch_for_principal( - budget: SyncBudget, - team_override: Option, -) -> Result { +async fn fetch_for_principal(budget: SyncBudget) -> Result { let max_attempts = budget.max_attempts(); - // Resolve from the merged config (managed_config_url > cli_chat_proxy_base_url, - // including the enterprise single-endpoint derivation) so endpoint overrides - // are honored and the bearer isn't sent to the public default. + // Resolve from the merged config (managed_config_url override) so endpoint + // overrides are honored and the bearer isn't sent to the public default. let url = crate::agent::config::EndpointsConfig::from_effective_config().resolve_managed_config_url(); - let team_auth = team_override.or_else(read_active_team_auth); - if let Some(dk) = resolve_deployment_key() { let source = ManagedConfigSource::DeploymentKey; - match fetch_managed_config(&url, &dk, source, max_attempts).await { - // A rejected dk (stale env/config) must not starve a valid team - // sign-in: fall through. Network/5xx do NOT — same unreachable - // server, double the latency for nothing. - Err(ManagedConfigError::DeploymentKeyRejected) if team_auth.is_some() => { - tracing::warn!("deployment key rejected; falling back to the team session token"); - } - Err(e) => return Err(e), - // Fall through to the team only when the dk has no config row: an apply - // converges disk to the served set, and the empty dk body must not delete - // the team's files. Gate on row existence, not content (which can serve empty). - Ok(body) if !body.config_exists() && team_auth.is_some() => { - tracing::debug!("deployment key has no config; trying the team principal"); - } - Ok(body) => return Ok(FetchedConfig::DeploymentKey { key: dk, body }), - } - } - - // The proxy resolves the team from the principal and returns its config. - if let Some(auth) = team_auth { - let body = fetch_managed_config( - &url, - &auth.key, - ManagedConfigSource::TeamOauth, - max_attempts, - ) - .await?; - return Ok(FetchedConfig::Team { - auth: Box::new(auth), - body, - }); + let body = fetch_managed_config(&url, &dk, source, max_attempts).await?; + return Ok(FetchedConfig::DeploymentKey { key: dk, body }); } Ok(FetchedConfig::NoPrincipal) } -async fn sync_inner( - budget: SyncBudget, - team_override: Option, -) -> Result { - match fetch_for_principal(budget, team_override).await? { +async fn sync_inner(budget: SyncBudget) -> Result { + match fetch_for_principal(budget).await? { FetchedConfig::DeploymentKey { key, body } => { let source = ManagedConfigSource::DeploymentKey; let fingerprint = deployment_key_fingerprint(&key); @@ -610,18 +526,6 @@ async fn sync_inner( &outcome, )) } - FetchedConfig::Team { auth, body } => { - let source = ManagedConfigSource::TeamOauth; - let outcome = apply_fetched(&body, source, auth.team_id.as_deref(), None)?; - // Team identity is bound via principal (team id), not a key fingerprint. - Ok(SyncOutcome::from_fetch( - &body, - source, - auth.team_id.clone(), - None, - &outcome, - )) - } FetchedConfig::NoPrincipal => Ok(SyncOutcome { wrote: false, served: false, @@ -720,7 +624,9 @@ fn evict_prior_managed_config(home: &std::path::Path) { fn credential_present(source: ManagedConfigSource) -> bool { match source { ManagedConfigSource::DeploymentKey => resolve_deployment_key().is_some(), - ManagedConfigSource::TeamOauth => team_principal_signed_in().unwrap_or(true), + // Team principals no longer exist; a cached team-sourced config has + // no live credential behind it. + ManagedConfigSource::TeamOauth => false, } } @@ -744,20 +650,15 @@ pub enum ManagedConfigSync { /// waiting for the background tick. `authenticated` pins the just-logged-in /// principal (`None` = on-disk team). Latency-bounded by [`SyncBudget::Login`]; /// failures are logged, not propagated (the background loop retries). -pub async fn post_login_sync(authenticated: Option) -> ManagedConfigSync { +pub async fn post_login_sync(_authenticated: Option) -> ManagedConfigSync { clear_orphan(); if !is_fetch_enabled() { return ManagedConfigSync::Skipped; } - // The just-authenticated team, else the on-disk one — reused for the gate - // and the sync (one auth.json read). With no team, only sync if due anyway. - let team = authenticated - .and_then(eligible_team_principal) - .or_else(read_active_team_auth); - if team.is_none() && !crate::config::is_managed_config_stale_for(¤t_serving_identity()) { + if !crate::config::is_managed_config_stale_for(¤t_serving_identity()) { return ManagedConfigSync::Skipped; } - match sync_bounded(SyncBudget::Login, team).await { + match sync_bounded(SyncBudget::Login).await { // Nothing was persisted for a rejected envelope — that's a failure to // report, not "no change" (the gate may refuse the next session). Some(Ok(SyncOutcome { @@ -791,14 +692,14 @@ pub async fn post_login_sync(authenticated: Option) -> ManagedConfigSy /// Whether a credential exists that `grok setup` could install config for. pub fn has_principal() -> bool { - resolve_deployment_key().is_some() || read_active_team_auth().is_some() + resolve_deployment_key().is_some() } /// Whether a managed identity owns this machine, IGNORING token expiry (unlike [`has_principal`]) so an /// expired/backdated `auth.json` can't disarm the gate. Unreadable → present (fail-safe; the gate ANDs this /// with [`crate::config::managed_policy_compromised_for`], which a personal user never satisfies). fn managed_principal_present() -> bool { - resolve_deployment_key().is_some() || team_principal_signed_in().unwrap_or(true) + resolve_deployment_key().is_some() } /// The serving identity for an optional team id: a configured deployment key always @@ -820,22 +721,12 @@ fn serving_identity_from(team_id: Option) -> crate::config::ServingIdent /// The identity to check the cache against for whoever serves now: a configured deployment key wins /// (else the active team, else none). pub fn current_serving_identity() -> crate::config::ServingIdentity { - serving_identity_from(read_active_team_auth().and_then(|a| a.team_id)) + serving_identity_from(None) } -/// The client's team_id, IGNORING token expiry (the binding must survive the cold-start -/// expired window). Must NOT special-case a configured deployment key — that would -/// disable envelope binding for a real team user. Used at fetch time to bind the envelope. +/// Team principals no longer exist in the Kimi Code auth model. pub fn active_team_id_any_expiry() -> Option { - let home = crate::util::kigi_home::kigi_home(); - let store = crate::auth::read_auth_json(&home.join("auth.json")).ok()?; - store - .values() - .find(|a| a.is_team_principal()) - .and_then(|a| a.team_id.clone()) - // A blank team_id (malformed auth.json) is unknown, not a distinct identity: it must not - // feed the gate's identity checks, the tenant-switch purge, or the envelope binding. - .filter(|id| !id.trim().is_empty()) + None } /// Like [`current_serving_identity`] but IGNORING token expiry, for the enforcement gate: @@ -854,36 +745,16 @@ pub async fn ensure_managed_policy_present( if !is_fetch_enabled() { return; } - // Cheap disk-only gates before any network token refresh, so the boot path doesn't pay - // an `auth()` in the common cases. A personal user (no deploy key, and no team in - // `auth.json` even ignoring expiry) skips entirely; a usable identity whose cache isn't - // hard-stale also skips. Only an expired-but-refreshable team token (identity reads - // `None` before the refresh) or a hard-stale cache falls through to `auth()` below. - // `auth.json` unreadable (`Err`) is NOT treated as "no principal" — that would skip - // enforcement on a transient read blip. - if resolve_deployment_key().is_none() && matches!(team_principal_signed_in(), Ok(false)) { - return; - } - let identity = current_serving_identity(); - if !matches!(identity, crate::config::ServingIdentity::None) - && !crate::config::is_managed_config_hard_stale_for(&identity) - { - return; - } - // Refresh before the heal so an expired-but-refreshable team token isn't dropped by - // the expiry filter. Bounded; deploy-key machines have no OAuth (auth() → None). - let team = tokio::time::timeout(SESSION_START_AUTH_DEADLINE, auth_manager.auth()) - .await - .ok() - .and_then(Result::ok) - .filter(GrokAuth::is_team_principal); - if !has_principal() { + // Cheap disk-only gates: only a configured deployment key can own managed + // policy now (team principals no longer exist). + let _ = auth_manager; + if resolve_deployment_key().is_none() { return; } if !crate::config::is_managed_config_hard_stale_for(¤t_serving_identity()) { return; } - match sync_bounded(SyncBudget::SessionStart, team).await { + match sync_bounded(SyncBudget::SessionStart).await { Some(Ok(_)) => {} Some(Err(e)) => tracing::warn!("session-start managed policy refresh failed: {e}"), None => tracing::warn!("session-start managed policy refresh timed out"), @@ -987,9 +858,8 @@ pub struct SetupReport { /// Fetches the report behind `grok setup --json` without writing anything: /// no artifacts, no signature sidecar, no sync marker. pub async fn fetch_setup_report() -> Result { - let (source, body) = match fetch_for_principal(SyncBudget::Standard, None).await? { + let (source, body) = match fetch_for_principal(SyncBudget::Standard).await? { FetchedConfig::DeploymentKey { body, .. } => (Some("deploymentKey"), body), - FetchedConfig::Team { body, .. } => (Some("teamOauth"), body), FetchedConfig::NoPrincipal => (None, ManagedConfigResponse::default()), }; // Match the installer's trust decision: a payload `grok setup` would refuse @@ -1015,7 +885,7 @@ pub async fn fetch_setup_report() -> Result { /// Run the `grok setup` sync for the current principal. The caller must check /// [`has_principal`] first and render the no-principal guidance. pub async fn run_setup() -> SetupOutcome { - match sync_with_budget(SyncBudget::Standard, None).await { + match sync_with_budget(SyncBudget::Standard).await { // A rejected envelope persisted nothing — reporting Installed would mask a // fetch the gate is about to refuse. Ok(SyncOutcome { diff --git a/crates/codegen/kigi-shell/src/mcp_doctor.rs b/crates/codegen/kigi-shell/src/mcp_doctor.rs index ab69d32..efb59d6 100644 --- a/crates/codegen/kigi-shell/src/mcp_doctor.rs +++ b/crates/codegen/kigi-shell/src/mcp_doctor.rs @@ -7,7 +7,7 @@ use std::sync::Arc; use kigi_tools::types::config_source::ConfigSource; use serde::Serialize; -use crate::auth::GrokComConfig; +use crate::auth::KimiCodeConfig; use crate::session::managed_mcp; use crate::session::mcp_servers; @@ -272,13 +272,13 @@ fn managed_found( /// Discover managed `grok_com_*` servers if the user has xAI auth on disk. async fn try_discover_managed_servers() -> (ConfigSourceStatus, Vec) { let kigi_home = kigi_tools::util::kigi_home::kigi_home(); - let grok_com_config = GrokComConfig::default(); - let auth_manager = Arc::new(crate::auth::AuthManager::new(&kigi_home, grok_com_config)); + let kimi_code_config = KimiCodeConfig::default(); + let auth_manager = Arc::new(crate::auth::AuthManager::new(&kigi_home, kimi_code_config)); let Some(snapshot) = auth_manager.current_or_expired() else { return managed_skipped("not logged in"); }; - if !snapshot.is_managed_mcp_eligible() { + if !snapshot.is_session_auth() { return managed_skipped(format!("{:?} auth (not xAI OIDC)", snapshot.auth_mode)); } diff --git a/crates/codegen/kigi-shell/src/remote/agent.rs b/crates/codegen/kigi-shell/src/remote/agent.rs index eeea49a..38bcb6a 100644 --- a/crates/codegen/kigi-shell/src/remote/agent.rs +++ b/crates/codegen/kigi-shell/src/remote/agent.rs @@ -5,7 +5,7 @@ use std::sync::Arc; -use crate::auth::{AuthManager, GrokComConfig}; +use crate::auth::{AuthManager, KimiCodeConfig}; use anyhow::{Context, Result, bail}; use serde::de::DeserializeOwned; @@ -64,7 +64,6 @@ impl SandboxClient { .context("failed to resolve sandbox auth")?; let mut builder = builder .header("Authorization", format!("Bearer {}", auth.key)) - .header("X-XAI-Token-Auth", GrokComConfig::default().token_header) .header("x-userid", &auth.user_id) .header("x-grok-client-version", kigi_version::VERSION); diff --git a/crates/codegen/kigi-shell/src/remote/chat_models_client.rs b/crates/codegen/kigi-shell/src/remote/chat_models_client.rs index f70d899..26c9518 100644 --- a/crates/codegen/kigi-shell/src/remote/chat_models_client.rs +++ b/crates/codegen/kigi-shell/src/remote/chat_models_client.rs @@ -121,10 +121,6 @@ impl ChatModelsClient { .post(&url) .json(&body) .header("Authorization", format!("Bearer {}", auth.key)) - .header( - "X-XAI-Token-Auth", - self.auth.grok_com_config().token_header.clone(), - ) .header("x-userid", &auth.user_id) .header("x-grok-client-version", kigi_version::VERSION) .header( diff --git a/crates/codegen/kigi-shell/src/remote/client.rs b/crates/codegen/kigi-shell/src/remote/client.rs index ca2f626..d9be243 100644 --- a/crates/codegen/kigi-shell/src/remote/client.rs +++ b/crates/codegen/kigi-shell/src/remote/client.rs @@ -1,5 +1,5 @@ //! HTTP client for backend CRUD operations. -use crate::auth::{GrokAuth, GrokComConfig}; +use crate::auth::{KimiAuth, KimiCodeConfig}; use crate::session::export::{ExportedMessage, ExportedMetadata, ExportedSession}; use indexmap::IndexMap; use prod_mc_cli_chat_proxy_types::SubagentBundle; @@ -16,13 +16,12 @@ pub fn share_url(permission_id: &str) -> String { } fn add_cli_chat_proxy_headers_blocking( builder: reqwest::blocking::RequestBuilder, - auth: &GrokAuth, + auth: &KimiAuth, alpha_test_key: Option<&str>, url: &str, ) -> reqwest::blocking::RequestBuilder { let mut builder = builder .header("Authorization", format!("Bearer {}", auth.key)) - .header("X-XAI-Token-Auth", GrokComConfig::default().token_header) .header("x-userid", &auth.user_id) .header("x-grok-client-version", kigi_version::VERSION); if let Some(email) = &auth.email { @@ -56,7 +55,7 @@ async fn add_bundle_fetch_headers( Some(am) => am.auth().await.ok(), None => None, }; - let mut credentials = crate::util::grok_auth_credentials::GrokAuthCredentials::new( + let mut credentials = crate::util::kigi_auth_credentials::KigiAuthCredentials::new( resolved_auth.as_ref().map(|auth| auth.key.clone()), ); credentials.deployment_key = deployment_key.map(str::to_owned); @@ -313,7 +312,7 @@ impl BackendClient { } } /// Attach a live `AuthManager` so every request resolves a fresh token - /// instead of requiring the caller to pass `&GrokAuth`. + /// instead of requiring the caller to pass `&KimiAuth`. pub fn with_auth_manager(mut self, manager: std::sync::Arc) -> Self { let credentials: std::sync::Arc = std::sync::Arc::new( @@ -328,7 +327,7 @@ impl BackendClient { self } /// Resolve auth from the attached `AuthManager`. - async fn resolve_auth(&self) -> Result { + async fn resolve_auth(&self) -> Result { let manager = self .auth_manager .as_ref() @@ -392,9 +391,7 @@ impl BackendClient { .await?; Ok(()) } - /// Build auth + identity headers. - /// Must include X-XAI-Token-Auth so nginx auth subrequest routes to authenticate_xai_grok_cli_token. - /// See: crates/codegen/kigi-shell/src/agent/app.rs:run_headless + /// Build auth + identity headers (plain bearer; no token-auth marker). async fn auth_header_map(&self) -> Result { use reqwest::header::{HeaderMap, HeaderValue}; let auth = self.resolve_auth().await?; @@ -403,10 +400,6 @@ impl BackendClient { HeaderValue::from_str(value) .map_err(|e| BackendError::Auth(format!("invalid {name} header: {e}"))) }; - headers.insert( - "X-XAI-Token-Auth", - required(&GrokComConfig::default().token_header, "X-XAI-Token-Auth")?, - ); headers.insert("x-userid", required(&auth.user_id, "x-userid")?); if let Some(email) = &auth.email && let Ok(v) = HeaderValue::from_str(email) @@ -556,7 +549,7 @@ impl BackendClient { /// network). 4xx and parse errors are not retried. pub fn fetch_settings_blocking( cli_chat_proxy_base_url: &str, - auth: &GrokAuth, + auth: &KimiAuth, alpha_test_key: Option<&str>, ) -> Option { let client = crate::http::shared_blocking_client(); @@ -713,7 +706,7 @@ pub struct FetchModelsResult { } pub(crate) fn fetch_models_blocking( endpoints: &crate::agent::config::EndpointsConfig, - auth: Option<&GrokAuth>, + auth: Option<&KimiAuth>, fetch_auth: crate::agent::models::ModelFetchAuth, ) -> Result { let client = crate::http::shared_blocking_client(); @@ -1279,37 +1272,23 @@ mod tests { let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); (format!("{base}/v1"), seen_headers, handle) } - fn test_auth() -> GrokAuth { - GrokAuth { + fn test_auth() -> KimiAuth { + KimiAuth { key: "token".to_string(), - auth_mode: crate::auth::AuthMode::Oidc, + auth_mode: crate::auth::AuthMode::OAuth, create_time: chrono::Utc::now(), user_id: "user-1".to_string(), email: Some("test@example.com".to_string()), - first_name: None, - last_name: None, - profile_image_asset_id: None, - principal_type: None, - principal_id: None, - team_id: None, - team_name: None, - team_role: None, - organization_id: None, - organization_name: None, - organization_role: None, - user_blocked_reason: None, - team_blocked_reasons: vec![], - coding_data_retention_opt_out: false, - has_grok_code_access: None, refresh_token: None, expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), - oidc_issuer: None, - oidc_client_id: None, + expires_in: Some(3600), + scope: None, + token_type: None, } } fn test_auth_manager() -> Arc { let dir = tempfile::tempdir().unwrap(); - let mgr = crate::auth::AuthManager::new(dir.path(), crate::auth::GrokComConfig::default()); + let mgr = crate::auth::AuthManager::new(dir.path(), crate::auth::KimiCodeConfig::default()); mgr.hot_swap(test_auth()); std::mem::forget(dir); Arc::new(mgr) @@ -1336,7 +1315,7 @@ mod tests { let headers = seen_headers.lock().unwrap(); let headers = headers.last().unwrap(); assert_eq!(headers.authorization.as_deref(), Some("Bearer token")); - assert_eq!(headers.token_auth.as_deref(), Some("xai-grok-cli")); + assert_eq!(headers.token_auth, None, "no token-auth marker header"); assert_eq!(headers.user_id.as_deref(), Some("user-1")); assert_eq!(headers.email.as_deref(), Some("test@example.com")); assert_eq!(headers.alpha_test_key, None); diff --git a/crates/codegen/kigi-shell/src/remote/conversations_client.rs b/crates/codegen/kigi-shell/src/remote/conversations_client.rs index a8b2f70..a84212b 100644 --- a/crates/codegen/kigi-shell/src/remote/conversations_client.rs +++ b/crates/codegen/kigi-shell/src/remote/conversations_client.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use serde::{Deserialize, Serialize}; -use crate::auth::{AuthManager, GrokAuth}; +use crate::auth::{AuthManager, KimiAuth}; const KIGI_WEB_URL: &str = "https://grok.com"; @@ -108,9 +108,9 @@ impl ConversationsClient { } } - async fn require_xai_auth(&self) -> Result { + async fn require_xai_auth(&self) -> Result { let auth = self.auth.auth().await.map_err(|_| ConvError::NoOauth)?; - if !auth.is_xai_auth() { + if !auth.is_session_auth() { return Err(ConvError::NoOauth); } Ok(auth) @@ -119,14 +119,10 @@ impl ConversationsClient { fn apply_auth_headers( &self, builder: reqwest::RequestBuilder, - auth: &GrokAuth, + auth: &KimiAuth, ) -> reqwest::RequestBuilder { let mut builder = builder .header("Authorization", format!("Bearer {}", auth.key)) - .header( - "X-XAI-Token-Auth", - self.auth.grok_com_config().token_header.clone(), - ) .header("x-userid", &auth.user_id) .header("x-grok-client-version", kigi_version::VERSION) .header( diff --git a/crates/codegen/kigi-shell/src/remote/pull_smoke_test.rs b/crates/codegen/kigi-shell/src/remote/pull_smoke_test.rs index d783431..9cbbc10 100644 --- a/crates/codegen/kigi-shell/src/remote/pull_smoke_test.rs +++ b/crates/codegen/kigi-shell/src/remote/pull_smoke_test.rs @@ -4,17 +4,17 @@ #[cfg(test)] mod tests { - use crate::auth::GrokAuth; + use crate::auth::KimiAuth; use crate::remote::client::BackendClient; use crate::session::storage::{JsonlStorageAdapter, StorageAdapter}; use std::collections::BTreeMap; use std::sync::Arc; - fn load_prod_auth() -> Option { + fn load_prod_auth() -> Option { let path = crate::util::kigi_home::kigi_home().join("auth.json"); let contents = std::fs::read_to_string(&path).ok()?; - let store: BTreeMap = serde_json::from_str(&contents).ok()?; - let scope = crate::auth::GrokComConfig::default().auth_scope(); + let store: BTreeMap = serde_json::from_str(&contents).ok()?; + let scope = crate::auth::KimiCodeConfig::default().auth_scope(); crate::auth::lookup_auth(&store, &scope) } @@ -33,7 +33,7 @@ mod tests { let auth = load_prod_auth().expect("No auth.json — run `grok login`"); let am = Arc::new(crate::auth::AuthManager::new( &crate::util::kigi_home::kigi_home(), - crate::auth::GrokComConfig::default(), + crate::auth::KimiCodeConfig::default(), )); am.hot_swap(auth); let client = BackendClient::new().with_auth_manager(am.clone()); diff --git a/crates/codegen/kigi-shell/src/remote/workspaces_client.rs b/crates/codegen/kigi-shell/src/remote/workspaces_client.rs index 3e9e4fa..a3ebce3 100644 --- a/crates/codegen/kigi-shell/src/remote/workspaces_client.rs +++ b/crates/codegen/kigi-shell/src/remote/workspaces_client.rs @@ -77,7 +77,7 @@ impl WorkspacesClient { pub async fn list_workspaces(&self, q: &WsQuery) -> Result { let auth = self.auth.auth().await.map_err(|_| WsError::NoOauth)?; - if !auth.is_xai_auth() { + if !auth.is_session_auth() { return Err(WsError::NoOauth); } @@ -98,10 +98,6 @@ impl WorkspacesClient { .get(&url) .query(&query) .header("Authorization", format!("Bearer {}", auth.key)) - .header( - "X-XAI-Token-Auth", - self.auth.grok_com_config().token_header.clone(), - ) .header("x-userid", &auth.user_id) .header("x-grok-client-version", kigi_version::VERSION) .header( diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs index 3d90ed2..f26b26d 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs @@ -328,7 +328,7 @@ impl SessionActor { .auth_manager .as_ref() .and_then(|am| am.current_or_expired()) - .filter(|a| a.is_xai_auth()) + .filter(|a| a.is_session_auth()) .map(|a| a.user_id), origin_client: self.origin_client.clone(), attribution_callback: self.attribution_callback.clone(), @@ -476,17 +476,11 @@ impl SessionActor { .and_then(|am| am.current_or_expired().map(|a| a.key.clone())); let models = self.models_manager.models(); let endpoints = self.models_manager.endpoints(); - let disable_api_key_auth = self - .auth_manager - .as_ref() - .map(|am| am.grok_com_config().api_key_auth_disabled()) - .unwrap_or(false); crate::agent::config::resolve_aux_model_sampling_config( slug, &models, &endpoints, session_key.as_deref(), - disable_api_key_auth, creds.alpha_test_key.clone(), creds.client_version.clone(), ) @@ -678,32 +672,6 @@ impl SessionActor { )), ); } - if auth_recovery_eligible - && crate::auth::devbox_login::is_devbox_environment() - && let Some(ref am) = self.auth_manager - { - match am.try_devbox_recovery().await { - Ok(auth) => { - tracing::info!( - session_id = % self.session_info.id.0, user_id = % auth.user_id, - "auth recovery: sampler 401, devbox re-mint, retrying" - ); - self.prepare_sampler_for_turn().await; - return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit); - } - Err(e) => { - tracing::warn!( - session_id = % self.session_info.id.0, error = % e, - "auth recovery: sampler 401, devbox re-mint failed" - ); - kigi_log::unified_log::warn( - "auth recovery: sampler 401, devbox re-mint failed", - Some(self.session_info.id.0.as_ref()), - Some(serde_json::json!({ "error" : format!("{e}") })), - ); - } - } - } if auth_recovery_eligible && let Some(ref am) = self.auth_manager { if am.try_recover_unauthorized().await { tracing::info!( @@ -762,24 +730,6 @@ impl SessionActor { .unwrap_or(crate::auth::AuthMode::ApiKey); let auth_mode_str = format!("{auth_mode:?}"); let client_version = kigi_version::VERSION; - if auth_mode == crate::auth::AuthMode::WebLogin { - let msg = format!( - "{detailed_message}\n\n\ - You are using a deprecated authentication method (WebLogin).\n\ - This auth method is no longer supported and will cause errors.\n\n\ - To fix: run `grok logout` then `grok login` to re-authenticate with OAuth2.\n\n\ - Version: {client_version}" - ); - self.log_terminal_failure("legacy_auth", error.status_code, &msg); - self.send_xai_notification(XaiSessionUpdate::RetryState( - crate::extensions::notification::RetryState::Failed { - error_type: "legacy_auth".to_string(), - message: msg.clone(), - }, - )) - .await; - return Err(acp::Error::internal_error().data(msg)); - } let is_model_404 = error.status_code == Some(404) && detailed_message.contains("does not exist"); let is_auth_401 = @@ -932,8 +882,10 @@ impl SessionActor { None, ); } - use crate::auth::{is_jwt_expired_or_near, parse_jwt_expiration}; - const REFRESH_THRESHOLD: chrono::Duration = chrono::Duration::minutes(5); + // BYOK path: pick up an externally rotated per-model key from + // config.toml. Kimi bearers are opaque (no client-side expiry + // probing); a changed on-disk key is adopted, an unchanged one is a + // no-op. let creds = self.chat_state_handle.get_credentials().await; let current_key = creds.api_key; let current_model_id = self @@ -943,41 +895,14 @@ impl SessionActor { .map(|c| c.model) .unwrap_or_default(); let Some(ref key) = current_key else { return }; - if !is_jwt_expired_or_near(key, REFRESH_THRESHOLD) { - if let Some(exp) = parse_jwt_expiration(key) { - let remaining_secs = (exp - chrono::Utc::now()).num_seconds(); - tracing::debug!( - model = % current_model_id, remaining_secs, - "JWT token valid, no refresh needed" - ); - } else { - tracing::debug!( - model = % current_model_id, key_len = key.len(), - "Token is not a JWT, expiry-based refresh not applicable" - ); - } - return; - } - let remaining_secs = - parse_jwt_expiration(key).map_or(0, |exp| (exp - chrono::Utc::now()).num_seconds()); - tracing::info!( - model = % current_model_id, remaining_secs, - "JWT near expiry, refreshing from config.toml" - ); let Some(new_key) = self.reload_api_key_from_config(¤t_model_id) else { return; }; if key == &new_key { - tracing::warn!( - model = % current_model_id, - "Config.toml returned same token (not yet rotated by external process?)" - ); return; } - let new_remaining_secs = parse_jwt_expiration(&new_key) - .map_or(0, |exp| (exp - chrono::Utc::now()).num_seconds()); tracing::info!( - model = % current_model_id, new_remaining_secs, key_len = new_key.len(), + model = % current_model_id, key_len = new_key.len(), "Refreshed API token from config.toml" ); let mut creds = self.chat_state_handle.get_credentials().await; diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs index abb4c9a..1fa3055 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs @@ -1,6 +1,6 @@ use super::support::*; use super::*; -use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig}; +use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig}; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use tokio::sync::mpsc; @@ -17,12 +17,12 @@ impl crate::auth::refresh::TokenRefresher for AlwaysSucceedRefresher { _reason: crate::auth::refresh::RefreshReason, ) -> crate::auth::refresh::RefreshOutcome { self.called.store(true, Ordering::SeqCst); - crate::auth::refresh::RefreshOutcome::Success(Box::new(GrokAuth { + crate::auth::refresh::RefreshOutcome::Success(Box::new(KimiAuth { key: "refreshed-test-token".to_string(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, refresh_token: Some("rt-new".into()), expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() })) } } @@ -34,13 +34,13 @@ fn auth_manager_with_refresher( refresher: Arc, ) -> (tempfile::TempDir, Arc) { let dir = tempfile::tempdir().expect("tempdir"); - let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - am.hot_swap(GrokAuth { + let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + am.hot_swap(KimiAuth { key: "initial-test-key".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, refresh_token: Some("rt".into()), expires_at: Some(chrono::Utc::now() - chrono::Duration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }); am.set_refresher(refresher); (dir, am) @@ -121,13 +121,13 @@ async fn make_actor_with_method_and_credentials( /// cache hit). The tempdir must outlive the manager (auth.json path). fn auth_manager_with_valid_token(key: &str) -> (tempfile::TempDir, Arc) { let dir = tempfile::tempdir().expect("tempdir"); - let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - am.hot_swap(GrokAuth { + let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + am.hot_swap(KimiAuth { key: key.into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, refresh_token: Some("rt".into()), expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }); (dir, am) } @@ -305,12 +305,12 @@ async fn proactive_refresh_makes_per_turn_refresh_a_cache_hit() { _: crate::auth::refresh::RefreshReason, ) -> crate::auth::refresh::RefreshOutcome { self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - crate::auth::refresh::RefreshOutcome::Success(Box::new(GrokAuth { + crate::auth::refresh::RefreshOutcome::Success(Box::new(KimiAuth { key: "proactive-fresh".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, refresh_token: Some("rt-new".into()), expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() })) } } @@ -318,14 +318,12 @@ async fn proactive_refresh_makes_per_turn_refresh_a_cache_hit() { }); let (_dir, am) = auth_manager_with_refresher(refresher); - let cancel = tokio_util::sync::CancellationToken::new(); - am.start_proactive_refresh(cancel.clone()); - - // Wait for proactive task to fire. - tokio::time::sleep(std::time::Duration::from_millis(500)).await; + // Drive one loop-body iteration directly (the production loop + // ticks on a fixed 60s cadence, far too slow for a unit test). + am.proactive_tick(false).await; assert!( call_count.load(Ordering::SeqCst) >= 1, - "proactive task must have fired" + "proactive tick must have refreshed" ); let count_after_proactive = call_count.load(Ordering::SeqCst); @@ -350,8 +348,6 @@ async fn proactive_refresh_makes_per_turn_refresh_a_cache_hit() { Some("proactive-fresh"), "per-turn refresh must pick up the proactively-refreshed token" ); - - cancel.cancel(); }) .await; } @@ -370,49 +366,6 @@ fn model_not_found_error() -> kigi_sampler::SamplingErrorInfo { } } -/// 404 model-not-found with a legacy WebLogin token appends a -/// "Legacy auth detected" hint to the error message. -#[tokio::test(flavor = "current_thread")] -async fn legacy_auth_hint_on_404_model_not_found() { - let local = tokio::task::LocalSet::new(); - local - .run_until(async { - let dir = tempfile::tempdir().expect("tempdir"); - let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - am.hot_swap(GrokAuth { - key: "legacy-token".into(), - auth_mode: AuthMode::WebLogin, - ..GrokAuth::test_default() - }); - - let (actor, _rx) = make_actor_with_auth_manager(Some(am)).await; - let result = actor.handle_sampling_failure(model_not_found_error()).await; - let err = match result { - Err(e) => e, - Ok(_) => panic!("expected Err from handle_sampling_failure"), - }; - let data = err.data.unwrap(); - let msg = data.as_str().unwrap(); - assert!( - msg.contains("deprecated authentication method"), - "404 with WebLogin must include deprecation message, got: {msg}" - ); - assert!( - msg.contains("grok logout"), - "hint must mention `grok logout`, got: {msg}" - ); - assert!( - msg.contains("grok login"), - "hint must mention `grok login`, got: {msg}" - ); - assert!( - msg.contains("Version:"), - "must show client version, got: {msg}" - ); - }) - .await; -} - /// Build a 401-shaped error that bypasses step 4b's auth recovery. /// /// In production, 401s arrive as `SamplingErrorKind::Auth` with @@ -437,47 +390,6 @@ fn unauthorized_401_error() -> kigi_sampler::SamplingErrorInfo { } } -/// 401 Unauthorized with a legacy WebLogin token appends a -/// "Legacy auth detected" hint to the error message. -#[tokio::test(flavor = "current_thread")] -async fn legacy_auth_hint_on_401_unauthorized() { - let local = tokio::task::LocalSet::new(); - local - .run_until(async { - let dir = tempfile::tempdir().expect("tempdir"); - let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - am.hot_swap(GrokAuth { - key: "legacy-token".into(), - auth_mode: AuthMode::WebLogin, - ..GrokAuth::test_default() - }); - - let (actor, _rx) = make_actor_with_auth_manager(Some(am)).await; - let result = actor - .handle_sampling_failure(unauthorized_401_error()) - .await; - let err = match result { - Err(e) => e, - Ok(_) => panic!("expected Err from handle_sampling_failure"), - }; - let data = err.data.unwrap(); - let msg = data.as_str().unwrap(); - assert!( - msg.contains("deprecated authentication method"), - "401 with WebLogin must include deprecation message, got: {msg}" - ); - assert!( - msg.contains("grok logout"), - "hint must mention `grok logout`, got: {msg}" - ); - assert!( - msg.contains("grok login"), - "hint must mention `grok login`, got: {msg}" - ); - }) - .await; -} - /// 401 with OIDC auth must NOT append the legacy hint. #[tokio::test(flavor = "current_thread")] async fn no_legacy_hint_on_401_for_oidc_auth() { @@ -485,13 +397,13 @@ async fn no_legacy_hint_on_401_for_oidc_auth() { local .run_until(async { let dir = tempfile::tempdir().expect("tempdir"); - let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - am.hot_swap(GrokAuth { + let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + am.hot_swap(KimiAuth { key: "oidc-token".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, refresh_token: Some("rt".into()), expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }); let (actor, _rx) = make_actor_with_auth_manager(Some(am)).await; @@ -513,7 +425,7 @@ async fn no_legacy_hint_on_401_for_oidc_auth() { "OIDC auth must NOT trigger WebLogin deprecation on 401, got: {msg}" ); assert!( - msg.contains("Auth: Oidc"), + msg.contains("Auth: OAuth"), "OIDC 401 must show auth mode in enriched message, got: {msg}" ); }) @@ -527,13 +439,13 @@ async fn no_legacy_hint_for_oidc_auth() { local .run_until(async { let dir = tempfile::tempdir().expect("tempdir"); - let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - am.hot_swap(GrokAuth { + let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + am.hot_swap(KimiAuth { key: "oidc-token".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, refresh_token: Some("rt".into()), expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }); let (actor, _rx) = make_actor_with_auth_manager(Some(am)).await; @@ -553,7 +465,7 @@ async fn no_legacy_hint_for_oidc_auth() { "OIDC auth must NOT trigger WebLogin deprecation, got: {msg}" ); assert!( - msg.contains("Auth: Oidc"), + msg.contains("Auth: OAuth"), "OIDC 404 must show auth mode in enriched message, got: {msg}" ); assert!( @@ -627,9 +539,9 @@ async fn sampler_401_session_method_with_stale_api_key_auth_type_still_recovers( .await; } -/// Same regression via the `oidc` method id (the other session-based variant). +/// Same regression via the interactive-login method id (the other session-based variant). #[tokio::test(flavor = "current_thread")] -async fn sampler_401_oidc_method_with_stale_api_key_auth_type_still_recovers() { +async fn sampler_401_login_method_with_stale_api_key_auth_type_still_recovers() { let local = tokio::task::LocalSet::new(); local .run_until(async { @@ -641,7 +553,7 @@ async fn sampler_401_oidc_method_with_stale_api_key_auth_type_still_recovers() { let (_dir, am) = auth_manager_with_refresher(refresher); let (actor, _rx) = make_actor_with_method_and_credentials( Some(am), - "oidc", + "grok.com", kigi_chat_state::AuthType::ApiKey, "stale-session-jwt".to_string(), ) @@ -651,7 +563,7 @@ async fn sampler_401_oidc_method_with_stale_api_key_auth_type_still_recovers() { assert!( matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)), - "oidc method must recover even when auth_type transiently reads ApiKey" + "interactive-login method must recover even when auth_type transiently reads ApiKey" ); assert!( called.load(Ordering::SeqCst), @@ -779,7 +691,9 @@ async fn session_born_on_api_key_recovers_after_oidc_login_without_restart() { // the shared handle this running actor already holds (no re-spawn). actor .auth_method_id - .store(Some(std::sync::Arc::new(acp::AuthMethodId::new("oidc")))); + .store(Some(std::sync::Arc::new(acp::AuthMethodId::new( + "cached_token", + )))); // The gate is recomputed each turn from the shared handle, so the // flip alone activates the live resolver on the very next turn -- diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/goal/goal_classifier_e2e_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/goal/goal_classifier_e2e_tests.rs index a87d318..00f2464 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/goal/goal_classifier_e2e_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/goal/goal_classifier_e2e_tests.rs @@ -2644,7 +2644,7 @@ fn test_auth_manager_for_models() -> std::sync::Arc { let tmp = tempfile::tempdir().expect("tempdir"); let mgr = std::sync::Arc::new(crate::auth::AuthManager::new( tmp.path(), - crate::auth::GrokComConfig::default(), + crate::auth::KimiCodeConfig::default(), )); std::mem::forget(tmp); mgr diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs index 2117583..f894422 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs @@ -132,13 +132,13 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() { let dir = tempfile::tempdir().unwrap(); let mgr = std::sync::Arc::new(crate::auth::AuthManager::new( dir.path(), - crate::auth::GrokComConfig::default(), + crate::auth::KimiCodeConfig::default(), )); - mgr.hot_swap(crate::auth::GrokAuth { - auth_mode: crate::auth::AuthMode::Oidc, + mgr.hot_swap(crate::auth::KimiAuth { + auth_mode: crate::auth::AuthMode::OAuth, refresh_token: Some("rt".into()), expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), - ..crate::auth::GrokAuth::test_default() + ..crate::auth::KimiAuth::test_default() }); std::mem::forget(dir); Some(mgr) diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs index 1306324..a6952ab 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs @@ -1256,13 +1256,13 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() { let dir = tempfile::tempdir().unwrap(); let mgr = std::sync::Arc::new(crate::auth::AuthManager::new( dir.path(), - crate::auth::GrokComConfig::default(), + crate::auth::KimiCodeConfig::default(), )); - mgr.hot_swap(crate::auth::GrokAuth { - auth_mode: crate::auth::AuthMode::Oidc, + mgr.hot_swap(crate::auth::KimiAuth { + auth_mode: crate::auth::AuthMode::OAuth, refresh_token: Some("rt".into()), expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), - ..crate::auth::GrokAuth::test_default() + ..crate::auth::KimiAuth::test_default() }); std::mem::forget(dir); Some(mgr) diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/media_gen_auth_retry_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/media_gen_auth_retry_tests.rs index 6a76cbc..85f8096 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/media_gen_auth_retry_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/media_gen_auth_retry_tests.rs @@ -1,17 +1,17 @@ use super::*; -use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig}; +use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig}; use kigi_tools::types::output::{ToolOutput, ToolRunResult}; use std::sync::atomic::{AtomicUsize, Ordering}; fn succeeding_am() -> Arc { let dir = tempfile::tempdir().unwrap(); - let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - am.hot_swap(GrokAuth { + let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + am.hot_swap(KimiAuth { key: "expired".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, refresh_token: Some("rt".into()), expires_at: Some(chrono::Utc::now() - chrono::Duration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }); struct Ok; #[async_trait::async_trait] @@ -20,11 +20,11 @@ fn succeeding_am() -> Arc { &self, _: crate::auth::refresh::RefreshReason, ) -> crate::auth::refresh::RefreshOutcome { - crate::auth::refresh::RefreshOutcome::Success(Box::new(GrokAuth { + crate::auth::refresh::RefreshOutcome::Success(Box::new(KimiAuth { key: "fresh".into(), expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), refresh_token: Some("rt-new".into()), - ..GrokAuth::test_default() + ..KimiAuth::test_default() })) } } @@ -36,13 +36,13 @@ fn succeeding_am() -> Arc { fn failing_am() -> Arc { let dir = tempfile::tempdir().unwrap(); - let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - am.hot_swap(GrokAuth { + let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + am.hot_swap(KimiAuth { key: "expired".into(), - auth_mode: AuthMode::Oidc, + auth_mode: AuthMode::OAuth, refresh_token: Some("rt".into()), expires_at: Some(chrono::Utc::now() - chrono::Duration::hours(1)), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }); struct Fail; #[async_trait::async_trait] diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_e2e_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_e2e_tests.rs index 79a55bc..5848b5b 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_e2e_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_e2e_tests.rs @@ -156,12 +156,12 @@ async fn actor_with_proxy( let home = tempfile::tempdir().expect("tempdir"); let auth_manager = Arc::new(crate::auth::AuthManager::new( home.path(), - crate::auth::GrokComConfig::default(), + crate::auth::KimiCodeConfig::default(), )); // Valid (1h) token in-memory only — `auth()` fast-paths it without network. - auth_manager.hot_swap(crate::auth::GrokAuth { + auth_manager.hot_swap(crate::auth::KimiAuth { expires_at: Some(Utc::now() + chrono::Duration::hours(1)), - ..crate::auth::GrokAuth::test_default() + ..crate::auth::KimiAuth::test_default() }); let cfg = crate::agent::config::Config { diff --git a/crates/codegen/kigi-shell/src/session/feedback_manager.rs b/crates/codegen/kigi-shell/src/session/feedback_manager.rs index d2fe71e..26d349b 100644 --- a/crates/codegen/kigi-shell/src/session/feedback_manager.rs +++ b/crates/codegen/kigi-shell/src/session/feedback_manager.rs @@ -981,23 +981,23 @@ mod tests { async fn test_is_auth_permanently_failed_reads_auth_manager() { use crate::agent::feedback_client::FeedbackClient; use crate::auth::error::RefreshTokenFailedReason; - use crate::auth::{AuthManager, GrokAuth, GrokComConfig}; + use crate::auth::{AuthManager, KimiAuth, KimiCodeConfig}; use std::sync::Arc; let dir = tempfile::tempdir().unwrap(); - let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); + let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); let client = FeedbackClient::new("http://example/v1", None).with_auth_manager(am.clone()); assert!(!client.is_auth_permanently_failed()); - // The verdict is scoped to the live credential's key. - am.hot_swap(GrokAuth { + // The tombstone is scoped to the live credential's refresh token. + am.hot_swap(KimiAuth { key: "tok".into(), - ..GrokAuth::test_default() + refresh_token: Some("rt".into()), + expires_at: Some(chrono::Utc::now() - chrono::Duration::hours(1)), + ..KimiAuth::test_default() }); - // Use a non-sticky reason: only recoverable verdicts age out (a sticky - // `RefreshTokenRejected` never expires), and this exercises the TTL path. - am.record_permanent_failure("tok".to_string(), RefreshTokenFailedReason::Other.into()); + am.record_permanent_failure("rt".to_string(), RefreshTokenFailedReason::Other.into()); assert!(client.is_auth_permanently_failed()); am.force_permanent_failure_aged_out(); @@ -1018,7 +1018,7 @@ mod tests { #[tokio::test] async fn test_has_token_refresher_requires_refresher_attached() { use crate::agent::feedback_client::FeedbackClient; - use crate::auth::{AuthManager, GrokComConfig}; + use crate::auth::{AuthManager, KimiCodeConfig}; use std::sync::Arc; struct NoOpRefresher; @@ -1035,7 +1035,7 @@ mod tests { } let dir = tempfile::tempdir().unwrap(); - let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); + let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); let bare = FeedbackClient::new("http://example/v1", None); assert!(!bare.has_token_refresher()); diff --git a/crates/codegen/kigi-shell/src/session/persistence.rs b/crates/codegen/kigi-shell/src/session/persistence.rs index 4261965..62896cc 100644 --- a/crates/codegen/kigi-shell/src/session/persistence.rs +++ b/crates/codegen/kigi-shell/src/session/persistence.rs @@ -1915,11 +1915,8 @@ fn init_remote_sync( "Writeback storage mode requires authentication. Run 'grok login' first.", ) })?; - if let Some(auth) = auth_manager.current_or_expired() { - if auth.is_zdr_team() { - tracing::debug!("ZDR team: skipping remote sync"); - return Ok(None); - } + if auth_manager.current_or_expired().is_some() { + // ZDR was an xAI team concept; nothing gates remote sync here. } else { tracing::warn!( "writeback: no auth loaded yet, ZDR check skipped (backend enforces server-side)" diff --git a/crates/codegen/kigi-shell/src/session/unified_list/mod.rs b/crates/codegen/kigi-shell/src/session/unified_list/mod.rs index 1b89bba..0bede93 100644 --- a/crates/codegen/kigi-shell/src/session/unified_list/mod.rs +++ b/crates/codegen/kigi-shell/src/session/unified_list/mod.rs @@ -565,13 +565,12 @@ mod tests { fn xai_auth_manager(dir: &std::path::Path) -> std::sync::Arc { let am = std::sync::Arc::new(crate::auth::AuthManager::new( dir, - crate::auth::GrokComConfig::default(), + crate::auth::KimiCodeConfig::default(), )); - am.hot_swap(crate::auth::GrokAuth { - auth_mode: crate::auth::AuthMode::Oidc, - oidc_issuer: Some(crate::auth::xai_oauth2_issuer().to_owned()), + am.hot_swap(crate::auth::KimiAuth { + auth_mode: crate::auth::AuthMode::OAuth, expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), - ..crate::auth::GrokAuth::test_default() + ..crate::auth::KimiAuth::test_default() }); am } @@ -648,9 +647,8 @@ mod tests { let home = tempfile::tempdir().expect("tempdir"); let auth = std::sync::Arc::new(crate::auth::AuthManager::new( home.path(), - crate::auth::GrokComConfig::default(), + crate::auth::KimiCodeConfig::default(), )); - auth.set_devbox_env_for_test(false); let client = ConversationsClient::new(auth); let mut req = ListReq::default(); force_kind_chat(&mut req); diff --git a/crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs b/crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs index 634ac28..2b8e067 100644 --- a/crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs +++ b/crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs @@ -125,7 +125,7 @@ pub(crate) fn ctx_with_toggle(toggle: HashMap) -> SubagentSpawnCon workspace_ops: kigi_workspace::WorkspaceOps::for_test(), auth_manager: Arc::new(crate::auth::AuthManager::new( std::path::Path::new("/tmp/nonexistent-grok-test"), - crate::auth::GrokComConfig::default(), + crate::auth::KimiCodeConfig::default(), )), attribution_callback: None, parent_agent_name: None, diff --git a/crates/codegen/kigi-shell/src/tier.rs b/crates/codegen/kigi-shell/src/tier.rs deleted file mode 100644 index 2fa1b4f..0000000 --- a/crates/codegen/kigi-shell/src/tier.rs +++ /dev/null @@ -1,58 +0,0 @@ -//! Subscription-tier classification shared across the shell and the pager. -//! -//! The subscription tier reaches the client as a free-form **display-name -//! string** (from CCP `/settings` `subscription_tier_display`, or the numeric -//! JWT `tier` claim mapped to a display-style string by -//! [`crate::agent::mvp_agent::jwt_tier_claim`]). There is no shared enum, so -//! gating decisions classify the string here in ONE place so the pager's -//! cosmetic slash-command gate and the shell's capability (toolset) gate can't -//! drift apart. -//! -//! "Restricted" tiers are the personal free tier and X Basic — the tiers the -//! server zero-limits on the Imagine and voice endpoints. Everything else -//! (SuperGrok, SuperGrok Heavy/Lite, X Premium/+, and any unknown future name) -//! is unrestricted (**fail-open**). - -/// Whether a **known** subscription-tier display name is a gated tier: the free -/// tier (CCP display "Free" or an empty string) or X Basic (CCP display -/// "X Basic"; JWT-claim fallback spelling "x_basic"). -/// -/// Case-insensitive and whitespace-trimmed. Callers decide the policy for an -/// *absent* tier (`None`): the pager treats absence as restricted (cosmetic, -/// recovers live on the next settings update), while the shell treats absence as -/// unrestricted (fail-open — the server authoritatively enforces per-tier -/// limits, so never withhold a capability on a guess). -pub fn is_restricted_tier_name(tier: &str) -> bool { - let t = tier.trim().to_ascii_lowercase(); - t.is_empty() || t == "free" || t == "x basic" || t == "x_basic" -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn restricted_names() { - assert!(is_restricted_tier_name("")); - assert!(is_restricted_tier_name(" ")); - assert!(is_restricted_tier_name("Free")); - assert!(is_restricted_tier_name("free")); - assert!(is_restricted_tier_name("X Basic")); - assert!(is_restricted_tier_name("x_basic")); - assert!(is_restricted_tier_name(" X BASIC ")); - } - - #[test] - fn unrestricted_names() { - assert!(!is_restricted_tier_name("SuperGrok")); - assert!(!is_restricted_tier_name("SuperGrok Heavy")); - assert!(!is_restricted_tier_name("supergrok_lite")); - assert!(!is_restricted_tier_name("X Premium")); - assert!(!is_restricted_tier_name("x_premium_plus")); - // API keys are not free-tier gated. - assert!(!is_restricted_tier_name("api_key")); - assert!(!is_restricted_tier_name("API Key")); - // Unknown future tiers fail open. - assert!(!is_restricted_tier_name("some_new_plan")); - } -} diff --git a/crates/codegen/kigi-shell/src/trace_classifier/mod.rs b/crates/codegen/kigi-shell/src/trace_classifier/mod.rs index 515e878..943b65f 100644 --- a/crates/codegen/kigi-shell/src/trace_classifier/mod.rs +++ b/crates/codegen/kigi-shell/src/trace_classifier/mod.rs @@ -1073,18 +1073,10 @@ pub async fn resolve_api_key(explicit: Option<&str>, kigi_home: &Path) -> Result /// * `Err(_)` — refresh attempt failed in a way the operator needs to /// see (network error, refresh_token rejected by the IdP, etc.). async fn non_interactive_auth_key(kigi_home: &Path) -> Result> { - use crate::auth::{AuthError, AuthManager, GrokComConfig}; + use crate::auth::{AuthError, AuthManager, KimiCodeConfig}; - // Production's `try_ensure_fresh_auth` clones the whole config to - // pass into `AuthManager::new` AND clones `auth_provider_command` - // again for `configure_refresher`. We extract the single field we - // need first, then move the rest of `config` into `AuthManager` - // — one `Option` clone instead of one full struct clone - // plus one Option clone. - let config = GrokComConfig::default(); - let auth_provider_command = config.auth_provider_command.clone(); - let manager = std::sync::Arc::new(AuthManager::new(kigi_home, config)); - manager.configure_refresher(auth_provider_command); + let manager = std::sync::Arc::new(AuthManager::new(kigi_home, KimiCodeConfig::default())); + manager.configure_refresher(); match manager.auth().await { Ok(auth) => { let trimmed = auth.key.trim(); @@ -2147,7 +2139,7 @@ mod tests { /// path entirely — useful for "plain key, no refresh wanted" /// fixtures. fn write_auth_json(kigi_home: &Path, key: &str) { - let scope = crate::auth::GrokComConfig::default().auth_scope(); + let scope = crate::auth::KimiCodeConfig::default().auth_scope(); let body = serde_json::json!({ scope: { "key": key, @@ -2172,11 +2164,11 @@ mod tests { /// returns it via the fast path; the refresher chain is NOT /// invoked, so no network call fires. fn write_fresh_oidc_auth_json(kigi_home: &Path, key: &str) { - let scope = crate::auth::GrokComConfig::default().auth_scope(); + let scope = crate::auth::KimiCodeConfig::default().auth_scope(); let body = serde_json::json!({ scope: { "key": key, - "auth_mode": "oidc", + "auth_mode": "oauth", "create_time": now_offset(0), "expires_at": now_offset(3600), "refresh_token": "test-refresh-token", @@ -2192,11 +2184,11 @@ mod tests { /// refresh chain has nothing to refresh against, so the auth /// call fails non-interactively. fn write_expired_oidc_auth_json_no_refresh(kigi_home: &Path, key: &str) { - let scope = crate::auth::GrokComConfig::default().auth_scope(); + let scope = crate::auth::KimiCodeConfig::default().auth_scope(); let body = serde_json::json!({ scope: { "key": key, - "auth_mode": "oidc", + "auth_mode": "oauth", "create_time": now_offset(-7200), "expires_at": now_offset(-3600), "user_id": "test-user", diff --git a/crates/codegen/kigi-shell/src/util/grok_auth_credentials.rs b/crates/codegen/kigi-shell/src/util/kigi_auth_credentials.rs similarity index 82% rename from crates/codegen/kigi-shell/src/util/grok_auth_credentials.rs rename to crates/codegen/kigi-shell/src/util/kigi_auth_credentials.rs index f8b9d13..62d9bda 100644 --- a/crates/codegen/kigi-shell/src/util/grok_auth_credentials.rs +++ b/crates/codegen/kigi-shell/src/util/kigi_auth_credentials.rs @@ -9,10 +9,10 @@ use std::sync::Arc; /// an `AuthManager` (visibility checks, bundle fetches, tests). /// /// Deployment key (enterprise) sends bare `Bearer`, routed to management key auth. -/// User token (xAI users) sends `Bearer` + `X-XAI-Token-Auth: xai-grok-cli`. +/// User tokens and deployment keys are both sent as a plain `Bearer`. /// Deployment key takes precedence when both are present. #[derive(Clone)] -pub struct GrokAuthCredentials { +pub struct KigiAuthCredentials { pub user_token: Option, pub deployment_key: Option, pub alpha_test_key: Option, @@ -20,9 +20,9 @@ pub struct GrokAuthCredentials { /// refresh chain; `resolve()` reads the in-memory cache. auth_manager: Option>, } -impl std::fmt::Debug for GrokAuthCredentials { +impl std::fmt::Debug for KigiAuthCredentials { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("GrokAuthCredentials") + f.debug_struct("KigiAuthCredentials") .field( "user_token", &self.user_token.as_ref().map(|_| ""), @@ -42,7 +42,7 @@ impl std::fmt::Debug for GrokAuthCredentials { .finish() } } -impl GrokAuthCredentials { +impl KigiAuthCredentials { /// Static credentials from a snapshot token. No refresh capability. pub fn new(user_token: Option) -> Self { Self { @@ -82,7 +82,7 @@ impl GrokAuthCredentials { /// Without this, the `resolve_async()` error fallback returns /// credentials with no token, causing requests to be sent without /// an Authorization header. - pub fn resolve(&self) -> GrokAuthCredentials { + pub fn resolve(&self) -> KigiAuthCredentials { if let Some(ref am) = self.auth_manager && let Some(auth) = am.current_or_expired() { @@ -97,7 +97,7 @@ impl GrokAuthCredentials { /// (memory -> disk -> active OIDC refresh). Falls back to sync /// `resolve()` on error so transient refresh failures don't drop /// the bearer. - pub async fn resolve_async(&self) -> GrokAuthCredentials { + pub async fn resolve_async(&self) -> KigiAuthCredentials { let Some(ref am) = self.auth_manager else { return self.clone(); }; @@ -120,12 +120,7 @@ impl GrokAuthCredentials { let builder = if let Some(ref key) = self.deployment_key { builder.header("Authorization", format!("Bearer {}", key)) } else if let Some(ref token) = self.user_token { - builder - .header("Authorization", format!("Bearer {}", token)) - .header( - obfstr::obfstr!("X-XAI-Token-Auth"), - obfstr::obfstr!("xai-grok-cli"), - ) + builder.header("Authorization", format!("Bearer {}", token)) } else { builder }; @@ -133,28 +128,28 @@ impl GrokAuthCredentials { builder } } -impl kigi_auth::HttpAuth for GrokAuthCredentials { +impl kigi_auth::HttpAuth for KigiAuthCredentials { fn apply(&self, builder: RequestBuilder, base_url: &str) -> RequestBuilder { - GrokAuthCredentials::apply(self, builder, base_url) + KigiAuthCredentials::apply(self, builder, base_url) } } #[cfg(test)] mod tests { use super::*; - use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig}; + use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig}; use chrono::{Duration, Utc}; use std::sync::Arc; fn make_manager_with_token( expires_at: chrono::DateTime, ) -> (Arc, tempfile::TempDir) { let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - let auth = GrokAuth { + let mgr = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + let auth = KimiAuth { key: "test-bearer-token".into(), - auth_mode: AuthMode::External, + auth_mode: AuthMode::OAuth, expires_at: Some(expires_at), create_time: Utc::now(), - ..GrokAuth::test_default() + ..KimiAuth::test_default() }; mgr.hot_swap(auth); (mgr, dir) @@ -162,14 +157,14 @@ mod tests { #[test] fn resolve_returns_token_when_not_expired() { let (mgr, _dir) = make_manager_with_token(Utc::now() + Duration::hours(1)); - let creds = GrokAuthCredentials::new(None).with_auth_manager(mgr); + let creds = KigiAuthCredentials::new(None).with_auth_manager(mgr); let resolved = creds.resolve(); assert_eq!(resolved.user_token.as_deref(), Some("test-bearer-token")); } #[test] fn resolve_returns_token_during_early_invalidation_window() { let (mgr, _dir) = make_manager_with_token(Utc::now() + Duration::minutes(3)); - let creds = GrokAuthCredentials::new(None).with_auth_manager(mgr.clone()); + let creds = KigiAuthCredentials::new(None).with_auth_manager(mgr.clone()); assert!(mgr.current().is_none()); assert!(mgr.current_or_expired().is_some()); assert_eq!( @@ -179,14 +174,14 @@ mod tests { } #[test] fn resolve_returns_static_token_when_no_auth_manager() { - let creds = GrokAuthCredentials::new(Some("static-token".into())); + let creds = KigiAuthCredentials::new(Some("static-token".into())); assert_eq!(creds.resolve().user_token.as_deref(), Some("static-token")); } #[test] fn resolve_returns_none_when_no_token_at_all() { let dir = tempfile::tempdir().unwrap(); - let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); - let creds = GrokAuthCredentials::new(None).with_auth_manager(mgr); + let mgr = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + let creds = KigiAuthCredentials::new(None).with_auth_manager(mgr); assert!(creds.resolve().user_token.is_none()); } } diff --git a/crates/codegen/kigi-shell/src/util/mod.rs b/crates/codegen/kigi-shell/src/util/mod.rs index 820e755..89c56ba 100644 --- a/crates/codegen/kigi-shell/src/util/mod.rs +++ b/crates/codegen/kigi-shell/src/util/mod.rs @@ -1,7 +1,7 @@ pub mod agent_id; pub mod config; -pub mod grok_auth_credentials; pub mod hooks; +pub mod kigi_auth_credentials; // The foundation utilities live in `kigi-shell-base` (upstream of this // crate so they build in parallel). Re-exported at the original paths so diff --git a/crates/codegen/kigi-shell/tests/signed_managed_config.rs b/crates/codegen/kigi-shell/tests/signed_managed_config.rs index 2881c6d..f98dbd9 100644 --- a/crates/codegen/kigi-shell/tests/signed_managed_config.rs +++ b/crates/codegen/kigi-shell/tests/signed_managed_config.rs @@ -12,8 +12,8 @@ mod common; use common::{ - MANAGED, REQUIREMENTS_FAIL_CLOSED, forged_team_body, install_test_key, reset, signed_team_body, - spawn_mock, team_identity, test_home, write_config, write_team_auth, + MANAGED, REQUIREMENTS_FAIL_CLOSED, dk_identity, forged_dk_body, install_test_key, reset, + signed_dk_body, spawn_mock, test_home, write_dk_config, }; use kigi_config::signed_policy; use serial_test::serial; @@ -27,20 +27,19 @@ async fn rejected_signature_persists_nothing_and_records_no_marker() { reset(&home); let (kp, _pubkey) = install_test_key(); - // Prior trusted state: team-b's files + marker (as if synced earlier). + // Prior trusted state: an earlier principal's files + marker. std::fs::write(home.join("managed_config.toml"), "[cli]\nprior = true\n").unwrap(); std::fs::write(home.join("requirements.toml"), "[features]\n").unwrap(); kigi_shell::config::mark_managed_config_synced(kigi_shell::config::SyncMarker { - principal: Some("team-b"), + principal: Some("dep-old"), had_managed_config: true, had_requirements: true, key_fingerprint: None, fail_closed: false, }); - let url = spawn_mock(forged_team_body(&kp, "team-007")); - write_config(&home, &url); - write_team_auth(&home, "team-007"); + let url = spawn_mock(forged_dk_body(&kp, "dep-42")); + write_dk_config(&home, &url, "dep-key-1"); let wrote = kigi_shell::managed_config::sync() .await @@ -61,7 +60,7 @@ async fn rejected_signature_persists_nothing_and_records_no_marker() { let v: serde_json::Value = serde_json::from_str(&marker).unwrap(); assert_eq!( v["principal"].as_str(), - Some("team-b"), + Some("dep-old"), "the marker must not be rewritten for a rejected fetch: {marker}" ); } @@ -75,14 +74,13 @@ async fn verified_envelope_persists_policy_and_sidecar() { reset(&home); let (kp, pubkey) = install_test_key(); - let url = spawn_mock(signed_team_body( + let url = spawn_mock(signed_dk_body( &kp, - "team-007", + "dep-42", Some(MANAGED), Some(REQUIREMENTS_FAIL_CLOSED), )); - write_config(&home, &url); - write_team_auth(&home, "team-007"); + write_dk_config(&home, &url, "dep-key-1"); let wrote = kigi_shell::managed_config::sync() .await @@ -114,7 +112,7 @@ async fn verified_envelope_persists_policy_and_sidecar() { assert!(payload.fail_closed, "the signed opt-in is carried"); assert!( - !kigi_shell::config::is_managed_config_hard_stale_for(&team_identity("team-007")), + !kigi_shell::config::is_managed_config_hard_stale_for(&dk_identity()), "a covered cache is not hard-stale" ); assert!( @@ -133,14 +131,13 @@ async fn deleted_sidecar_under_fail_closed_marker_refuses_at_gate() { reset(&home); let (kp, _pubkey) = install_test_key(); - let url = spawn_mock(signed_team_body( + let url = spawn_mock(signed_dk_body( &kp, - "team-007", + "dep-42", Some(MANAGED), Some(REQUIREMENTS_FAIL_CLOSED), )); - write_config(&home, &url); - write_team_auth(&home, "team-007"); + write_dk_config(&home, &url, "dep-key-1"); kigi_shell::managed_config::sync() .await .expect("initial sync should succeed"); @@ -152,11 +149,11 @@ async fn deleted_sidecar_under_fail_closed_marker_refuses_at_gate() { std::fs::remove_file(home.join("managed_config.sig.json")).unwrap(); assert!( - kigi_shell::config::is_managed_config_hard_stale_for(&team_identity("team-007")), + kigi_shell::config::is_managed_config_hard_stale_for(&dk_identity()), "a stripped sidecar must trigger the session-start refetch" ); assert!( - kigi_shell::config::is_managed_config_stale_for(&team_identity("team-007")), + kigi_shell::config::is_managed_config_stale_for(&dk_identity()), "the TIMER staleness sibling must fire too (background tick self-heal), even though the marker is timer-fresh" ); let gate = kigi_shell::managed_config::managed_policy_gate(); @@ -170,113 +167,3 @@ async fn deleted_sidecar_under_fail_closed_marker_refuses_at_gate() { "the refusal is the managed-policy gate message" ); } - -/// The keyed availability fix: after a fail_closed team-A install (signed sidecar + marker), an -/// OFFLINE switch to team B previously read Compromised (the authentic sidecar is bound to A) and -/// refused a legitimate switch. The gate's identity-change purge must shed team A's artifacts -/// INCLUDING the sidecar, PERMIT team B, and leave the cache hard-stale so the next online start -/// fetches team B's own policy. -#[tokio::test] -#[serial] -async fn offline_team_switch_purges_sidecar_and_permits_new_team() { - let home = test_home().clone(); - reset(&home); - let (kp, _pubkey) = install_test_key(); - - let url = spawn_mock(signed_team_body( - &kp, - "team-a", - Some(MANAGED), - Some(REQUIREMENTS_FAIL_CLOSED), - )); - write_config(&home, &url); - write_team_auth(&home, "team-a"); - kigi_shell::managed_config::sync() - .await - .expect("team A keyed sync should succeed"); - assert!( - home.join("managed_config.sig.json").exists(), - "the keyed sync persists a sidecar" - ); - assert!( - kigi_shell::managed_config::managed_policy_gate().is_ok(), - "team A's verified fail_closed policy must start" - ); - - // Switch the signed-in team to B; the gate is sync, so no fetch can rebind first. - write_team_auth(&home, "team-b"); - - // The bug this fixes: without the purge, team B evaluates against team A's - // foreign-bound sidecar → Compromised → a legitimate switch refused startup. - assert!( - kigi_shell::config::managed_policy_compromised_for(&team_identity("team-b")), - "pre-purge, the foreign-bound sidecar must read compromised for team B" - ); - - assert!( - kigi_shell::managed_config::managed_policy_gate().is_ok(), - "the gate must purge team A and permit the legitimate offline switch to team B" - ); - for f in [ - "requirements.toml", - "managed_config.toml", - "managed_config_cache.json", - "managed_config.sig.json", - ] { - assert!( - !home.join(f).exists(), - "{f} must be purged on the identity change" - ); - } - assert!( - kigi_shell::config::is_managed_config_hard_stale_for(&team_identity("team-b")), - "the purged cache must read hard-stale so the next online start fetches team B's policy" - ); -} - -/// A blank `team_id` in `auth.json` (a parse blip) over an authentic team-A-bound fail_closed -/// sidecar: the blank→None filter resolves the identity to None, the marker principal backstops -/// the signed binding (team-a vs team-a → Trusted), so the KEYED gate PERMITS — instead of -/// binding to "" and refusing as Compromised — and nothing is purged. -#[tokio::test] -#[serial] -async fn keyed_blank_team_id_is_not_refused_and_does_not_purge() { - let home = test_home().clone(); - reset(&home); - let (kp, _pubkey) = install_test_key(); - - let url = spawn_mock(signed_team_body( - &kp, - "team-a", - Some(MANAGED), - Some(REQUIREMENTS_FAIL_CLOSED), - )); - write_config(&home, &url); - write_team_auth(&home, "team-a"); - kigi_shell::managed_config::sync() - .await - .expect("team A keyed sync should succeed"); - assert!( - kigi_shell::managed_config::managed_policy_gate().is_ok(), - "team A's verified fail_closed policy must start" - ); - - // auth.json now carries a team principal with a BLANK team_id. - write_team_auth(&home, ""); - - assert!( - kigi_shell::managed_config::managed_policy_gate().is_ok(), - "a blank team_id must read as unknown, not a foreign binding that reads compromised" - ); - for f in [ - "requirements.toml", - "managed_config.toml", - "managed_config_cache.json", - "managed_config.sig.json", - ] { - assert!( - home.join(f).exists(), - "{f} must be retained on a blank team_id (a parse blip is not an identity change)" - ); - } -} diff --git a/crates/codegen/kigi-shell/tests/signed_managed_config/common.rs b/crates/codegen/kigi-shell/tests/signed_managed_config/common.rs index 3378d68..07b59af 100644 --- a/crates/codegen/kigi-shell/tests/signed_managed_config/common.rs +++ b/crates/codegen/kigi-shell/tests/signed_managed_config/common.rs @@ -93,17 +93,8 @@ pub fn spawn_mock(body: String) -> String { format!("http://{addr}/deployment/config") } -pub fn write_config(home: &std::path::Path, managed_config_url: &str) { - std::fs::write( - home.join("config.toml"), - format!("[endpoints]\nmanaged_config_url = \"{managed_config_url}\"\n"), - ) - .unwrap(); -} - -/// [`write_config`] plus a `deployment_key` (dead-code-allowed: compiled into -/// both binaries, called by one). -#[allow(dead_code)] +/// Endpoint config with a `deployment_key` (the only principal that can own +/// managed config now). pub fn write_dk_config(home: &std::path::Path, managed_config_url: &str, deployment_key: &str) { std::fs::write( home.join("config.toml"), @@ -114,22 +105,6 @@ pub fn write_dk_config(home: &std::path::Path, managed_config_url: &str, deploym .unwrap(); } -pub fn write_team_auth(home: &std::path::Path, team_id: &str) { - let scope = kigi_shell::auth::GrokComConfig::default().auth_scope(); - let auth = serde_json::json!({ - scope: { - "key": "team-session-token", - "auth_mode": "oidc", - "create_time": "2026-01-01T00:00:00Z", - "expires_at": "2099-01-01T00:00:00Z", - "user_id": "user-1", - "principal_type": "Team", - "team_id": team_id, - } - }); - std::fs::write(home.join("auth.json"), auth.to_string()).unwrap(); -} - /// A fresh Ed25519 keypair plus its raw public key, installed as the sole trusted /// key ([`TEST_KEY_ID`]) via the test seam. pub fn install_test_key() -> (ring::signature::Ed25519KeyPair, Vec) { @@ -163,18 +138,18 @@ pub fn sign_envelope( }) } -/// A team deployment-config response signed by `kp` under [`TEST_KEY_ID`]. The +/// A deployment-key config response signed by `kp` under [`TEST_KEY_ID`]. The /// body's legacy fields mirror the payload exactly (the client rejects a divergence). -pub fn signed_team_body( +pub fn signed_dk_body( kp: &ring::signature::Ed25519KeyPair, - team_id: &str, + deployment_id: &str, managed: Option<&str>, requirements: Option<&str>, ) -> String { let payload = SignedPayload { version: prod_mc_cli_chat_proxy_types::SIGNED_PAYLOAD_VERSION, - deployment_id: None, - team_id: Some(team_id.to_owned()), + deployment_id: Some(deployment_id.to_owned()), + team_id: None, managed_config: managed.map(str::to_owned), requirements: requirements.map(str::to_owned), fail_closed: requirements.is_some_and(kigi_config::fail_closed_flag_from_str), @@ -182,8 +157,8 @@ pub fn signed_team_body( key_id: TEST_KEY_ID.into(), }; serde_json::json!({ - "deployment_id": serde_json::Value::Null, - "team_id": team_id, + "deployment_id": deployment_id, + "team_id": serde_json::Value::Null, "managed_config": managed, "requirements": requirements, "signatures": [sign_envelope(kp, &payload)], @@ -191,19 +166,20 @@ pub fn signed_team_body( .to_string() } -/// A [`signed_team_body`] (managed config only) with the signature corrupted — +/// A [`signed_dk_body`] (managed config only) with the signature corrupted — /// valid base64, wrong bytes — so the verifier must reject the envelope. -pub fn forged_team_body(kp: &ring::signature::Ed25519KeyPair, team_id: &str) -> String { +pub fn forged_dk_body(kp: &ring::signature::Ed25519KeyPair, deployment_id: &str) -> String { let mut body: serde_json::Value = - serde_json::from_str(&signed_team_body(kp, team_id, Some(MANAGED), None)).unwrap(); + serde_json::from_str(&signed_dk_body(kp, deployment_id, Some(MANAGED), None)).unwrap(); body["signatures"][0]["signature"] = base64::engine::general_purpose::STANDARD .encode([0u8; 64]) .into(); body.to_string() } -pub fn team_identity(id: &str) -> kigi_shell::config::ServingIdentity { - kigi_shell::config::ServingIdentity::Team(id.to_owned()) +/// The live serving identity (the configured deployment key's fingerprint). +pub fn dk_identity() -> kigi_shell::config::ServingIdentity { + kigi_shell::managed_config::current_serving_identity() } /// True when `path` reads despite `chmod 000` (root / DAC bypass): chmod-based diff --git a/crates/codegen/kigi-shell/tests/signed_managed_config_extended.rs b/crates/codegen/kigi-shell/tests/signed_managed_config_extended.rs index a266d80..e4928e5 100644 --- a/crates/codegen/kigi-shell/tests/signed_managed_config_extended.rs +++ b/crates/codegen/kigi-shell/tests/signed_managed_config_extended.rs @@ -11,9 +11,8 @@ mod common; #[cfg(unix)] use common::skip_as_root; use common::{ - MANAGED, REQUIREMENTS_FAIL_CLOSED, TEST_EXPIRES_AT, TEST_KEY_ID, forged_team_body, - install_test_key, reset, sign_envelope, signed_team_body, spawn_mock, team_identity, test_home, - write_config, write_dk_config, write_team_auth, + MANAGED, REQUIREMENTS_FAIL_CLOSED, TEST_EXPIRES_AT, TEST_KEY_ID, dk_identity, forged_dk_body, + install_test_key, reset, sign_envelope, signed_dk_body, spawn_mock, test_home, write_dk_config, }; use kigi_config::signed_policy::{self, SignedPayload}; use serial_test::serial; @@ -21,14 +20,13 @@ use serial_test::serial; /// The healthy fail-closed starting state the tamper/heal scenarios mutate; /// the mock keeps serving the same body, so a healing sync can refetch it. async fn sync_fail_closed_policy(home: &std::path::Path, kp: &ring::signature::Ed25519KeyPair) { - let url = spawn_mock(signed_team_body( + let url = spawn_mock(signed_dk_body( kp, - "team-007", + "dep-42", Some(MANAGED), Some(REQUIREMENTS_FAIL_CLOSED), )); - write_config(home, &url); - write_team_auth(home, "team-007"); + write_dk_config(home, &url, "dep-key-1"); kigi_shell::managed_config::sync() .await .expect("initial sync should succeed"); @@ -94,9 +92,8 @@ async fn rejected_signature_surfaces_as_setup_and_login_failure() { reset(&home); let (kp, _pubkey) = install_test_key(); - let url = spawn_mock(forged_team_body(&kp, "team-007")); - write_config(&home, &url); - write_team_auth(&home, "team-007"); + let url = spawn_mock(forged_dk_body(&kp, "dep-42")); + write_dk_config(&home, &url, "dep-key-1"); let outcome = kigi_shell::managed_config::run_setup().await; assert!( @@ -130,8 +127,8 @@ async fn withdrawn_requirements_is_deleted_and_covered_by_the_new_sidecar() { sync_fail_closed_policy(&home, &kp).await; assert!(home.join("requirements.toml").exists()); - let url_partial = spawn_mock(signed_team_body(&kp, "team-007", Some(MANAGED), None)); - write_config(&home, &url_partial); + let url_partial = spawn_mock(signed_dk_body(&kp, "dep-42", Some(MANAGED), None)); + write_dk_config(&home, &url_partial, "dep-key-1"); let wrote = kigi_shell::managed_config::sync() .await .expect("withdrawing sync should succeed"); @@ -156,7 +153,7 @@ async fn withdrawn_requirements_is_deleted_and_covered_by_the_new_sidecar() { "the new sidecar covers the absence" ); assert!( - !kigi_shell::config::is_managed_config_hard_stale_for(&team_identity("team-007")), + !kigi_shell::config::is_managed_config_hard_stale_for(&dk_identity()), "the converged, covered cache is not hard-stale" ); assert!(kigi_shell::managed_config::managed_policy_gate().is_ok()); @@ -192,7 +189,7 @@ async fn directory_squat_reads_compromised_and_online_sync_heals() { "the refusal is the managed-policy gate message" ); assert!( - kigi_shell::config::is_managed_config_hard_stale_for(&team_identity("team-007")), + kigi_shell::config::is_managed_config_hard_stale_for(&dk_identity()), "the squat must trigger the refetch" ); @@ -239,7 +236,7 @@ async fn sidecar_read_blip_allows_session_and_triggers_refetch() { "a transient sidecar read blip must not refuse the session" ); assert!( - kigi_shell::config::is_managed_config_hard_stale_for(&team_identity("team-007")), + kigi_shell::config::is_managed_config_hard_stale_for(&dk_identity()), "the blip must trigger the refetch so the self-heal runs" ); // Restore so the tempdir (and later tests) stay clean. @@ -276,7 +273,7 @@ async fn sidecar_directory_squat_refuses_then_online_sync_heals() { "the refusal is the managed-policy gate message" ); assert!( - kigi_shell::config::is_managed_config_hard_stale_for(&team_identity("team-007")), + kigi_shell::config::is_managed_config_hard_stale_for(&dk_identity()), "the squat must trigger the refetch" ); diff --git a/crates/codegen/kigi-shell/tests/team_managed_config.rs b/crates/codegen/kigi-shell/tests/team_managed_config.rs deleted file mode 100644 index 4959550..0000000 --- a/crates/codegen/kigi-shell/tests/team_managed_config.rs +++ /dev/null @@ -1,1846 +0,0 @@ -//! End-to-end client tests for team-OAuth managed config against a mock -//! deployment-config endpoint. Proxy-side resolution is unit-tested in -//! the cli-chat-proxy deployment-config route. -//! -//! Every test here MUST be `#[serial]`: they share one process-global -//! `KIGI_SHARE_DIR` (the `kigi_home` `OnceLock` allows a single value per process) -//! and mutate that directory + process env, so concurrent tests would race. - -use std::io::{BufRead, BufReader, Write}; -use std::net::TcpListener; -use std::path::PathBuf; -use std::sync::{Arc, Mutex, OnceLock}; - -use kigi_shell::config::ServingIdentity; -use kigi_test_support::spawn_counting_server; -use serial_test::serial; - -/// The serving identity for a team id (the staleness checks key on this). -fn team_identity(id: &str) -> ServingIdentity { - ServingIdentity::Team(id.to_owned()) -} - -/// Shared temp dir used as KIGI_SHARE_DIR for the whole test binary (the kigi_home -/// `OnceLock` only allows one value per process). Also scrubs/installs the env -/// this suite depends on, before any test thread reads it. -fn test_home() -> &'static PathBuf { - static HOME: OnceLock = OnceLock::new(); - HOME.get_or_init(|| { - let path = tempfile::TempDir::new().unwrap().keep(); - // SAFETY: set once at init before other threads read the vars. - unsafe { - std::env::set_var("KIGI_SHARE_DIR", &path); - // Ambient env must not shadow the scenarios under test: a real - // deployment key, a managed-config opt-out, or a proxy that would - // intercept the 127.0.0.1 mocks. - for var in [ - "KIGI_DEPLOYMENT_KEY", - "KIGI_MANAGED_CONFIG", - "KIGI_DEPLOYMENT_CONFIG_REFRESH_INTERVAL_SECS", - "KIGI_DEPLOYMENT_CONFIG_CACHE_TTL_SECS", - "HTTP_PROXY", - "HTTPS_PROXY", - "ALL_PROXY", - "http_proxy", - "https_proxy", - "all_proxy", - ] { - std::env::remove_var(var); - } - // Real exponential backoff would add seconds per retry test. - std::env::set_var("KIGI_DEPLOYMENT_CONFIG_BACKOFF_MS", "10"); - } - path - }) -} - -fn reset(home: &std::path::Path) { - for f in [ - "config.toml", - "auth.json", - "managed_config.toml", - "requirements.toml", - "managed_config_cache.json", - "managed_config.lock", - ] { - let _ = std::fs::remove_file(home.join(f)); - } -} - -/// Read one HTTP request's header block (up to the blank line) and return the -/// `Authorization` header value, if any. Header-boundary-safe, unlike a single -/// fixed-size `read()`. -fn read_request_auth(stream: &mut std::net::TcpStream) -> Option { - let mut reader = BufReader::new(stream); - let mut auth = None; - loop { - let mut line = String::new(); - if reader.read_line(&mut line).unwrap_or(0) == 0 { - return auth; - } - let line = line.trim_end(); - if line.is_empty() { - return auth; - } - if let Some((name, value)) = line.split_once(':') - && name.eq_ignore_ascii_case("authorization") - { - auth = Some(value.trim().to_string()); - } - } -} - -/// Mock deployment-config server serving `body` to every request. Returns the -/// URL and the `Authorization` header of every request in order. -fn spawn_mock(body: String) -> (String, Arc>>) { - let (url, _count, auths) = spawn_mock_seq(vec![(200, body)]); - (url, auths) -} - -/// `(url, request_count, authorization_headers_in_order)`. -type MockHandle = (String, Arc>, Arc>>); - -/// Mock server that serves a sequence of `(status, body)` responses — response -/// `i` for request `i`, clamping to the last. The handle's request counter -/// backs retry/fail-fast assertions; the auth log backs credential-fallback -/// assertions. -fn spawn_mock_seq(responses: Vec<(u16, String)>) -> MockHandle { - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let addr = listener.local_addr().unwrap(); - let count: Arc> = Arc::new(Mutex::new(0)); - let counter = count.clone(); - let auths: Arc>> = Arc::new(Mutex::new(Vec::new())); - let seen_auths = auths.clone(); - std::thread::spawn(move || { - for stream in listener.incoming() { - let Ok(mut stream) = stream else { continue }; - if let Some(auth) = read_request_auth(&mut stream) { - seen_auths.lock().unwrap().push(auth); - } - let i = { - let mut c = counter.lock().unwrap(); - let idx = *c; - *c += 1; - idx - }; - let (status, body) = responses - .get(i) - .or_else(|| responses.last()) - .cloned() - .unwrap_or((200, "{}".to_string())); - let reason = match status { - 200 => "OK", - 401 => "Unauthorized", - 500 => "Internal Server Error", - _ => "Status", - }; - let resp = format!( - "HTTP/1.1 {status} {reason}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", - body.len(), - body - ); - let _ = stream.write_all(resp.as_bytes()); - let _ = stream.flush(); - } - }); - (format!("http://{addr}/deployment/config"), count, auths) -} - -/// Mock that abruptly closes its first `close_first` connections right after -/// reading the request — simulating a stale/poisoned keep-alive connection the -/// client reused — then serves `body` (HTTP 200) on every later connection. The -/// counter records accepted connections, backing the retry assertion. -fn spawn_mock_closing_first(close_first: usize, body: String) -> (String, Arc>) { - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let addr = listener.local_addr().unwrap(); - let count: Arc> = Arc::new(Mutex::new(0)); - let counter = count.clone(); - std::thread::spawn(move || { - for stream in listener.incoming() { - let Ok(mut stream) = stream else { continue }; - // Read the request first so the abort lands mid-response ("connection - // closed before message completed"), not as a connect/write failure. - let _ = read_request_auth(&mut stream); - let i = { - let mut c = counter.lock().unwrap(); - let idx = *c; - *c += 1; - idx - }; - if i < close_first { - drop(stream); - continue; - } - let resp = format!( - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", - body.len(), - body - ); - let _ = stream.write_all(resp.as_bytes()); - let _ = stream.flush(); - } - }); - (format!("http://{addr}/deployment/config"), count) -} - -/// Mock that, for its first `truncate_first` connections, writes a valid status line -/// + headers with an OVERSIZED `Content-Length` then closes WITHOUT the body — so the -/// client's body read fails mid-body ("connection closed before message completed", a -/// `reqwest` body-phase error, NOT a decode error). Every later connection serves the -/// full valid `body` (HTTP 200). The counter records accepted connections, backing the -/// retry assertion. -fn spawn_mock_truncating_body_first( - truncate_first: usize, - body: String, -) -> (String, Arc>) { - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let addr = listener.local_addr().unwrap(); - let count: Arc> = Arc::new(Mutex::new(0)); - let counter = count.clone(); - std::thread::spawn(move || { - for stream in listener.incoming() { - let Ok(mut stream) = stream else { continue }; - // Read the request first so the abort lands in the body phase, not as a - // connect/write failure. - let _ = read_request_auth(&mut stream); - let i = { - let mut c = counter.lock().unwrap(); - let idx = *c; - *c += 1; - idx - }; - if i < truncate_first { - // Promise 100000 body bytes, send none, then drop: the client reads the - // headers fine but fails collecting the body. - let headers = "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 100000\r\nConnection: close\r\n\r\n"; - let _ = stream.write_all(headers.as_bytes()); - let _ = stream.flush(); - drop(stream); - continue; - } - let resp = format!( - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", - body.len(), - body - ); - let _ = stream.write_all(resp.as_bytes()); - let _ = stream.flush(); - } - }); - (format!("http://{addr}/deployment/config"), count) -} - -/// Write a `config.toml` that routes the managed-config fetch at the mock. -fn write_config(home: &std::path::Path, managed_config_url: &str) { - std::fs::write( - home.join("config.toml"), - format!("[endpoints]\nmanaged_config_url = \"{managed_config_url}\"\n"), - ) - .unwrap(); -} - -/// Write an `auth.json` with a team OAuth principal under the active scope. -fn write_team_auth(home: &std::path::Path, team_id: &str) { - write_team_auth_expiry(home, team_id, "2099-01-01T00:00:00Z"); -} - -/// Like [`write_team_auth`] but with an explicit `expires_at`, so tests can -/// simulate a routine cold-start where the persisted access token is expired. -fn write_team_auth_expiry(home: &std::path::Path, team_id: &str, expires_at: &str) { - let scope = kigi_shell::auth::GrokComConfig::default().auth_scope(); - let auth = serde_json::json!({ - scope: { - "key": "team-session-token", - "auth_mode": "oidc", - "create_time": "2026-01-01T00:00:00Z", - "expires_at": expires_at, - "user_id": "user-1", - "principal_type": "Team", - "team_id": team_id, - } - }); - std::fs::write(home.join("auth.json"), auth.to_string()).unwrap(); -} - -/// Write an `auth.json` with an EXPIRED `external`-mode team principal, so a configured refresher -/// drives `AuthManager::auth()`. Models the cold-start where the persisted token is expired but refreshable. -fn write_expired_external_team_auth(home: &std::path::Path, team_id: &str) { - let scope = kigi_shell::auth::GrokComConfig::default().auth_scope(); - let auth = serde_json::json!({ - scope: { - "key": "stale-team-token", - "auth_mode": "external", - "create_time": "2026-01-01T00:00:00Z", - "expires_at": PAST, - "user_id": "user-1", - "principal_type": "Team", - "team_id": team_id, - "refresh_token": "rt-team", - } - }); - std::fs::write(home.join("auth.json"), auth.to_string()).unwrap(); -} - -const FAR_FUTURE: &str = "2099-01-01T00:00:00Z"; -const PAST: &str = "2000-01-01T00:00:00Z"; - -const TEAM_MANAGED: &str = "[[marketplace.sources]]\nname = \"internal\"\ngit = \"https://github.com/example/plugin-marketplace-internal\"\n"; -const TEAM_REQUIREMENTS: &str = - "[marketplace]\nallowlist = [\"https://github.com/example/plugin-marketplace-internal\"]\n"; - -fn team_config_body() -> String { - serde_json::json!({ - "deployment_id": serde_json::Value::Null, - "team_id": "team-007", - "managed_config": TEAM_MANAGED, - "requirements": TEAM_REQUIREMENTS, - }) - .to_string() -} - -#[tokio::test] -#[serial] -async fn team_sync_writes_files() { - let home = test_home().clone(); - reset(&home); - - let (url, auths) = spawn_mock(team_config_body()); - write_config(&home, &url); - write_team_auth(&home, "team-007"); - - let wrote = kigi_shell::managed_config::sync() - .await - .expect("sync should succeed"); - assert!(wrote, "expected team config to be written"); - - assert_eq!( - auths.lock().unwrap().last().map(String::as_str), - Some("Bearer team-session-token"), - "client must authenticate with the team session token" - ); - - let managed = std::fs::read_to_string(home.join("managed_config.toml")).unwrap(); - assert!( - managed.contains("plugin-marketplace-internal"), - "managed_config should contain the team marketplace source: {managed}" - ); - - let requirements = std::fs::read_to_string(home.join("requirements.toml")).unwrap(); - assert!( - requirements.contains("allowlist"), - "requirements should contain the enforced allowlist: {requirements}" - ); -} - -/// Switching the active team must not keep enforcing the prior team's policy: after B syncs, -/// A's artifacts are evicted and the marker records B served nothing. Fail-open. -#[tokio::test] -#[serial] -async fn team_switch_evicts_prior_teams_policy() { - let home = test_home().clone(); - reset(&home); - - // Team A serves both managed_config and requirements. - let (url_a, _auths_a) = spawn_mock(team_config_body()); - write_config(&home, &url_a); - write_team_auth(&home, "team-a"); - kigi_shell::managed_config::sync() - .await - .expect("team A sync should succeed"); - assert!(home.join("requirements.toml").exists()); - assert!(home.join("managed_config.toml").exists()); - - // Switch to team B, whose server returns a row (team_id) but no artifacts. - let body_b = serde_json::json!({ - "deployment_id": serde_json::Value::Null, - "team_id": "team-b", - "managed_config": serde_json::Value::Null, - "requirements": serde_json::Value::Null, - }) - .to_string(); - let (url_b, _auths_b) = spawn_mock(body_b); - write_config(&home, &url_b); - write_team_auth(&home, "team-b"); - - let wrote = kigi_shell::managed_config::sync() - .await - .expect("team B sync must not fail the session"); - assert!(!wrote, "team B serves no artifacts, so nothing is written"); - - // Team A's enforced policy is gone — team B does not inherit it. - assert!( - !home.join("requirements.toml").exists(), - "team A's requirements must be evicted on the switch to team B" - ); - assert!( - !home.join("managed_config.toml").exists(), - "team A's managed_config must be evicted on the switch to team B" - ); - - // The marker is now team B's and must not claim B served A's artifacts. - let marker = std::fs::read_to_string(home.join("managed_config_cache.json")).unwrap(); - let v: serde_json::Value = serde_json::from_str(&marker).unwrap(); - assert_eq!( - v["principal"].as_str(), - Some("team-b"), - "marker must rebind to team B: {marker}" - ); - assert_eq!( - v["had_requirements"].as_bool(), - Some(false), - "marker must not claim team B served requirements: {marker}" - ); - assert_eq!( - v["had_managed_config"].as_bool(), - Some(false), - "marker must not claim team B served managed_config: {marker}" - ); - - // Team B's cache reads fresh + identity-matched (no missing-artifact stale). - assert!(!kigi_shell::config::is_managed_config_stale_for( - &team_identity("team-b") - )); -} - -/// An artifact the server stops serving is removed on the next sync (disk converges -/// to the served set), and the marker stops claiming it — a withdrawn policy must not -/// keep enforcing from a stale file. -#[tokio::test] -#[serial] -async fn withdrawn_artifact_is_removed_on_next_sync() { - let home = test_home().clone(); - reset(&home); - - // Sync 1: both artifacts served. - let (url_full, _a) = spawn_mock(team_config_body()); - write_config(&home, &url_full); - write_team_auth(&home, "team-007"); - kigi_shell::managed_config::sync() - .await - .expect("initial sync should succeed"); - assert!(home.join("requirements.toml").exists()); - - // Sync 2: same team, requirements withdrawn. - let body = serde_json::json!({ - "deployment_id": serde_json::Value::Null, - "team_id": "team-007", - "managed_config": TEAM_MANAGED, - "requirements": serde_json::Value::Null, - }) - .to_string(); - let (url_partial, _a2) = spawn_mock(body); - write_config(&home, &url_partial); - let wrote = kigi_shell::managed_config::sync() - .await - .expect("second sync should succeed"); - assert!(wrote, "removing the withdrawn artifact is a change"); - - assert!( - home.join("managed_config.toml").exists(), - "the still-served artifact stays" - ); - assert!( - !home.join("requirements.toml").exists(), - "the withdrawn artifact is removed" - ); - let marker = std::fs::read_to_string(home.join("managed_config_cache.json")).unwrap(); - let v: serde_json::Value = serde_json::from_str(&marker).unwrap(); - assert_eq!( - v["had_requirements"].as_bool(), - Some(false), - "the marker stops claiming the withdrawn artifact: {marker}" - ); - assert!( - !kigi_shell::config::is_managed_config_stale_for(&team_identity("team-007")), - "the converged cache is not stale" - ); -} - -/// An empty dk response (`{}`) with a team signed in falls through to the team WITHOUT -/// applying: applying converges disk to the served (empty) set, which would delete the -/// team's files right before the team apply — observable when the team fetch then fails. -#[tokio::test] -#[serial] -async fn empty_dk_response_with_failing_team_leaves_team_policy_intact() { - let home = test_home().clone(); - reset(&home); - - // Seed the team's policy files + marker. - let (url, _a) = spawn_mock(team_config_body()); - write_config(&home, &url); - write_team_auth(&home, "team-007"); - kigi_shell::managed_config::sync() - .await - .expect("team seed sync should succeed"); - assert!(home.join("requirements.toml").exists()); - - // dk serves an empty row; the team fetch then fails (5xx for every retry). - let (url2, _c, auths) = spawn_mock_seq(vec![(200, "{}".into()), (500, "boom".into())]); - std::fs::write( - home.join("config.toml"), - format!("[endpoints]\nmanaged_config_url = \"{url2}\"\ndeployment_key = \"dep-key\"\n"), - ) - .unwrap(); - - let err = kigi_shell::managed_config::sync() - .await - .expect_err("the team fetch fails after the dk fallthrough"); - assert!(err.is_retryable(), "5xx is a transient failure: {err}"); - assert_eq!( - auths.lock().unwrap().first().map(String::as_str), - Some("Bearer dep-key"), - "the dk was consulted first" - ); - - assert!( - home.join("requirements.toml").exists() && home.join("managed_config.toml").exists(), - "the empty dk body must not be applied (it would delete the team's files)" - ); - let marker = std::fs::read_to_string(home.join("managed_config_cache.json")).unwrap(); - let v: serde_json::Value = serde_json::from_str(&marker).unwrap(); - assert_eq!( - v["principal"].as_str(), - Some("team-007"), - "the failed sync must not rewrite the marker: {marker}" - ); -} - -/// A served-then-deleted artifact reads stale for the active identity; the session-start refresh refetches it. -#[tokio::test] -#[serial] -async fn served_then_deleted_refetches_best_effort() { - let home = test_home().clone(); - reset(&home); - - let (url, _auths) = spawn_mock(team_config_body()); - write_config(&home, &url); - write_team_auth(&home, "team-007"); - kigi_shell::managed_config::sync() - .await - .expect("initial sync should succeed"); - assert!(home.join("requirements.toml").exists()); - assert!( - !kigi_shell::config::is_managed_config_stale_for(&team_identity("team-007")), - "a fresh, identity-matched, complete cache is not stale" - ); - - // Tamper: delete the served file but keep the fresh marker. - std::fs::remove_file(home.join("requirements.toml")).unwrap(); - assert!( - kigi_shell::config::is_managed_config_stale_for(&team_identity("team-007")), - "a served-but-now-missing artifact must read stale" - ); - // Identity mismatch also reads stale (team switch). - assert!(kigi_shell::config::is_managed_config_stale_for( - &team_identity("team-other") - )); - - // Best-effort refresh restores it; the session is never refused. Here the on-disk token is unexpired - // (cache hit). The expired-refreshable path is covered by `expired_refreshable_team_token_heals_after_auth_refresh`. - let auth_manager = std::sync::Arc::new(kigi_shell::auth::AuthManager::new( - &home, - kigi_shell::auth::GrokComConfig::default(), - )); - kigi_shell::managed_config::ensure_managed_policy_present(&auth_manager).await; - assert!( - home.join("requirements.toml").exists(), - "the best-effort refresh restored the deleted artifact" - ); - assert!(!kigi_shell::config::is_managed_config_stale_for( - &team_identity("team-007") - )); -} - -/// An expired-but-refreshable team token with a served-then-deleted artifact must heal at session start. -/// The heal drives `auth()` first so the refreshed principal can refetch; else the expiry filters drop it unhealed. -#[tokio::test] -#[serial] -async fn expired_refreshable_team_token_heals_after_auth_refresh() { - let home = test_home().clone(); - reset(&home); - - let (url, auths) = spawn_mock(team_config_body()); - // Point both the managed-config fetch and the proxy (post-refresh `/user`) at the mock — no production calls. - let base = url.trim_end_matches("/deployment/config"); - std::fs::write( - home.join("config.toml"), - format!( - "[endpoints]\nmanaged_config_url = \"{url}\"\ncli_chat_proxy_base_url = \"{base}\"\n" - ), - ) - .unwrap(); - - // Establish a served, identity-matched cache (writes the sync marker). - write_team_auth(&home, "team-007"); - kigi_shell::managed_config::sync() - .await - .expect("initial sync should succeed"); - assert!(home.join("requirements.toml").exists()); - - // Cold-start brick: expired-but-refreshable on-disk token, served requirements.toml gone but still in the marker. - write_expired_external_team_auth(&home, "team-007"); - std::fs::remove_file(home.join("requirements.toml")).unwrap(); - assert!( - !kigi_shell::managed_config::has_principal(), - "the expired token leaves no eligible managed principal — the expiry-filtered heal path can't see it (the brick)" - ); - assert!( - kigi_shell::config::is_managed_config_stale_for(&team_identity("team-007")), - "a served-but-now-missing artifact reads stale" - ); - - // A real AuthManager whose refresher mints a fresh team token, persisted by `auth()`. - let auth_manager = std::sync::Arc::new(kigi_shell::auth::AuthManager::new( - &home, - kigi_shell::auth::GrokComConfig::default(), - )); - auth_manager.configure_refresher(Some( - r#"echo '{"access_token":"refreshed-team-token","expires_in":3600}'"#.to_string(), - )); - - kigi_shell::managed_config::ensure_managed_policy_present(&auth_manager).await; - - // The refresh re-enabled the heal: policy restored, refetched with the fresh token. - assert!( - home.join("requirements.toml").exists(), - "the token refresh let the best-effort heal restore the deleted policy" - ); - assert!( - kigi_shell::managed_config::has_principal(), - "the refreshed token is a live, eligible managed (team) principal" - ); - assert!( - auths - .lock() - .unwrap() - .iter() - .any(|a| a == "Bearer refreshed-team-token"), - "the heal refetch authenticated with the refreshed token" - ); -} - -/// The boundary: when no refresh can succeed (offline / dead token), the expired token still fails closed — -/// `auth()` errs, the heal doesn't run, and the deleted policy is NOT restored. Unbricks ONLY on a real refresh. -#[tokio::test] -#[serial] -async fn expired_team_token_without_successful_refresh_stays_failed_closed() { - let home = test_home().clone(); - reset(&home); - - let (url, _auths) = spawn_mock(team_config_body()); - write_config(&home, &url); - write_team_auth(&home, "team-007"); - kigi_shell::managed_config::sync() - .await - .expect("initial sync should succeed"); - std::fs::remove_file(home.join("requirements.toml")).unwrap(); - write_expired_external_team_auth(&home, "team-007"); - - // A refresher that always fails -> `auth()` cannot produce a principal. - let auth_manager = std::sync::Arc::new(kigi_shell::auth::AuthManager::new( - &home, - kigi_shell::auth::GrokComConfig::default(), - )); - auth_manager.configure_refresher(Some("false".to_string())); - - kigi_shell::managed_config::ensure_managed_policy_present(&auth_manager).await; - - assert!( - !home.join("requirements.toml").exists(), - "with no successful refresh the expired token cannot heal (fail-closed)" - ); -} - -/// `managed_policy_gate` refuses a managed session when its served policy was deleted and the refetch can't -/// restore it (offline); intact or config-less is allowed. Exercises the real sync → marker → gate path. -#[tokio::test] -#[serial] -async fn managed_policy_gate_fails_closed_on_deleted_policy_offline() { - let home = test_home().clone(); - reset(&home); - - // Admin opts in by serving `fail_closed = true` (server-driven; no local env). - let body = serde_json::json!({ - "deployment_id": serde_json::Value::Null, - "team_id": "team-007", - "managed_config": TEAM_MANAGED, - "requirements": format!("fail_closed = true\n{TEAM_REQUIREMENTS}"), - }) - .to_string(); - let (url, _auths) = spawn_mock(body); - write_config(&home, &url); - write_team_auth(&home, "team-007"); - kigi_shell::managed_config::sync() - .await - .expect("initial sync should succeed"); - // Intact, identity-matched policy → the gate proceeds. - assert!( - kigi_shell::managed_config::managed_policy_gate().is_ok(), - "an intact served policy must not be refused" - ); - - // Tamper: delete the served file but keep the marker → gate fails closed. - std::fs::remove_file(home.join("requirements.toml")).unwrap(); - assert!( - kigi_shell::managed_config::managed_policy_gate().is_err(), - "a served-then-deleted policy must fail closed" - ); - - // Offline: the 5xx refetch can't restore the file → gate stays fail-closed (far-future token makes auth a cache hit). - let (err_url, _c, _a) = spawn_mock_seq(vec![(500, "{}".to_string())]); - write_config(&home, &err_url); - let auth_manager = std::sync::Arc::new(kigi_shell::auth::AuthManager::new( - &home, - kigi_shell::auth::GrokComConfig::default(), - )); - kigi_shell::managed_config::ensure_managed_policy_present(&auth_manager).await; - assert!( - !home.join("requirements.toml").exists(), - "a failed refetch cannot restore the deleted policy" - ); - assert!( - kigi_shell::managed_config::managed_policy_gate().is_err(), - "still missing after a failed refetch → gate stays fail-closed" - ); - - // A config-less principal (the server served nothing) is never refused. - reset(&home); - write_config(&home, &url); - write_team_auth(&home, "team-007"); - kigi_shell::config::mark_managed_config_synced(kigi_shell::config::SyncMarker { - principal: Some("team-007"), - had_managed_config: false, - had_requirements: false, - key_fingerprint: None, - fail_closed: false, - }); - assert!( - kigi_shell::managed_config::managed_policy_gate().is_ok(), - "a config-less principal must not be refused" - ); -} - -/// `bootstrap` must run the fail-closed gate (it is `bootstrap`'s first step): a compromised managed policy -/// must fail the whole bootstrap closed, not just the standalone `managed_policy_gate`. Guards against a -/// refactor that drops the gate call from `bootstrap` — which the gate's own tests would not catch. -#[tokio::test] -#[serial] -async fn bootstrap_fails_closed_when_managed_policy_compromised() { - let home = test_home().clone(); - reset(&home); - - // Provision a fail_closed team install (both artifacts served), then tamper by deleting the served policy. - let body = serde_json::json!({ - "deployment_id": serde_json::Value::Null, - "team_id": "team-007", - "managed_config": TEAM_MANAGED, - "requirements": format!("fail_closed = true\n{TEAM_REQUIREMENTS}"), - }) - .to_string(); - let (url, _auths) = spawn_mock(body); - write_config(&home, &url); - write_team_auth(&home, "team-007"); - kigi_shell::managed_config::sync() - .await - .expect("initial sync should succeed"); - std::fs::remove_file(home.join("requirements.toml")).unwrap(); - - // The gate is bootstrap's first step, so it refuses before any config/model work. - let cfg = kigi_shell::agent::config::Config::default(); - let auth_manager = std::sync::Arc::new(kigi_shell::auth::AuthManager::new( - &home, - kigi_shell::auth::GrokComConfig::default(), - )); - // `bootstrap`'s Ok type isn't `Debug`, so match rather than `expect_err`. - let err = match kigi_shell::agent::init::bootstrap(&cfg, &auth_manager, None) { - Err(e) => e, - Ok(_) => { - panic!("a compromised fail_closed policy must fail bootstrap closed, but it succeeded") - } - }; - assert!( - err.contains("Managed policy is required for this account"), - "bootstrap must fail via the managed-policy gate (proves bootstrap calls it); got: {err}" - ); -} - -/// Live wiring guard: an offline `KIGI_DEPLOYMENT_KEY` switch on a fail_closed install must FAIL CLOSED, else a -/// regression returning `None` silently disables deploy-key-switch detection. Same-key ALLOW checks the lib's own `blake3(KEY-AAA)` exactly. -#[tokio::test] -#[serial] -async fn managed_policy_gate_fails_closed_on_deployment_key_switch_offline() { - let home = test_home().clone(); - reset(&home); - - // Provision a fail_closed deploy install bound to key A; both artifacts written, so the only tamper signal is the key fingerprint. - let body = serde_json::json!({ - "deployment_id": "deploy-A", - "managed_config": TEAM_MANAGED, - "requirements": format!("fail_closed = true\n{TEAM_REQUIREMENTS}"), - }) - .to_string(); - let (url, _auths) = spawn_mock(body); - write_config(&home, &url); - - // SAFETY: #[serial] test; the env is restored before any assertion below. - unsafe { std::env::set_var("KIGI_DEPLOYMENT_KEY", "KEY-AAA") }; - kigi_shell::managed_config::sync() - .await - .expect("deployment-key sync should record the fail_closed marker"); - - // The marker records the lib-computed fingerprint (never the raw key), full blake3 hex. - let marker = std::fs::read_to_string(home.join("managed_config_cache.json")).unwrap(); - let v: serde_json::Value = serde_json::from_str(&marker).unwrap(); - let fp_a = v["key_fingerprint"] - .as_str() - .expect("deploy-key sync records a key fingerprint") - .to_string(); - let fail_closed_recorded = v["fail_closed"].as_bool().unwrap_or(false); - - // KIGI_MANAGED_CONFIG=0 disables any incidental background fetch (the gate is sync anyway). - // SAFETY: #[serial] test; restored before any assertion below. - unsafe { std::env::set_var("KIGI_MANAGED_CONFIG", "0") }; - - // Same key A → matching fingerprint → ALLOW (exact-equality vs recorded blake3). - let gate_same_key = kigi_shell::managed_config::managed_policy_gate(); - - // Switch to a different key B → REFUSE: exercises the full offline wiring. - // SAFETY: #[serial] test; restored immediately below. - unsafe { std::env::set_var("KIGI_DEPLOYMENT_KEY", "KEY-BBB") }; - let gate_switched_key = kigi_shell::managed_config::managed_policy_gate(); - - // SAFETY: #[serial] test; restore env BEFORE asserting so a failed assert can't leak it to later tests. - unsafe { - std::env::remove_var("KIGI_DEPLOYMENT_KEY"); - std::env::remove_var("KIGI_MANAGED_CONFIG"); - } - - // blake3-256 hex is exactly 64 hex chars — pins the recorded format. - assert!( - fp_a.len() == 64 && fp_a.chars().all(|c| c.is_ascii_hexdigit()), - "recorded key_fingerprint must be a full blake3 hex, not just non-empty: {marker}" - ); - assert!( - fail_closed_recorded, - "marker must record fail_closed = true: {marker}" - ); - assert!( - home.join("requirements.toml").exists() && home.join("managed_config.toml").exists(), - "both served artifacts must be present so the fingerprint is the only tamper signal" - ); - assert!( - !marker.contains("KEY-AAA"), - "the raw deployment key must never be written to disk: {marker}" - ); - assert!( - gate_same_key.is_ok(), - "same deployment key + intact fail_closed policy must be ALLOWED (proves recorded fp == fresh blake3(KEY-AAA))" - ); - assert!( - gate_switched_key.is_err(), - "a deployment-key switch (different fingerprint) on a fail_closed machine must FAIL CLOSED offline" - ); -} - -/// A leftover fail_closed marker from a prior managed stint must NOT lock out a user who has since signed out -/// (no deployment key, no team auth): the gate requires a present principal, so `managed_principal_present()` -/// short-circuits. Guards the worst-case regression — locking a normal user out of their own CLI. -#[test] -#[serial] -fn former_managed_user_signed_out_is_not_locked_out() { - let home = test_home().clone(); - reset(&home); - - // No deployment key in config, and `reset` removed auth.json → no principal. - std::fs::write(home.join("config.toml"), "[endpoints]\n").unwrap(); - // A stale, opted-in marker that reads tampered: it recorded a served requirements.toml that's now absent. - kigi_shell::config::mark_managed_config_synced(kigi_shell::config::SyncMarker { - principal: Some("team-007"), - had_managed_config: false, - had_requirements: true, - key_fingerprint: None, - fail_closed: true, - }); - - assert!( - kigi_shell::managed_config::managed_policy_gate().is_ok(), - "a signed-out user with a leftover fail_closed marker must not be refused (no principal to enforce)" - ); -} - -/// An unreadable auth.json makes `managed_principal_present()` fail safe to "present", but with NO fail_closed -/// marker the gate still allows — the fail-safe must never lock out a personal user who has no managed policy. -#[test] -#[serial] -fn unreadable_auth_without_marker_is_not_refused() { - let home = test_home().clone(); - reset(&home); - - std::fs::write(home.join("config.toml"), "[endpoints]\n").unwrap(); - std::fs::write(home.join("auth.json"), "{corrupt json").unwrap(); - // No managed_config_cache.json marker at all. - - assert!( - kigi_shell::managed_config::managed_policy_gate().is_ok(), - "unreadable auth (fail-safe present) with no fail_closed marker must not be refused" - ); -} - -/// Deploy-key online heal: a fail_closed deploy install whose served requirements.toml was deleted heals at -/// session start when the refetch succeeds, and the gate then allows — the key path's mirror of the team heal. -#[tokio::test] -#[serial] -async fn deployment_key_served_then_deleted_heals_online() { - let home = test_home().clone(); - reset(&home); - - let body = serde_json::json!({ - "deployment_id": "deploy-A", - "managed_config": TEAM_MANAGED, - "requirements": format!("fail_closed = true\n{TEAM_REQUIREMENTS}"), - }) - .to_string(); - let (url, _auths) = spawn_mock(body); - std::fs::write( - home.join("config.toml"), - format!("[endpoints]\nmanaged_config_url = \"{url}\"\ndeployment_key = \"KEY-AAA\"\n"), - ) - .unwrap(); - kigi_shell::managed_config::sync() - .await - .expect("initial deploy-key sync should succeed"); - assert!(home.join("requirements.toml").exists()); - - // Tamper: delete the served artifact (offline this would fail closed). - std::fs::remove_file(home.join("requirements.toml")).unwrap(); - - // The mock still serves, so the best-effort session-start refresh restores it; the gate then allows. - let auth_manager = std::sync::Arc::new(kigi_shell::auth::AuthManager::new( - &home, - kigi_shell::auth::GrokComConfig::default(), - )); - kigi_shell::managed_config::ensure_managed_policy_present(&auth_manager).await; - assert!( - home.join("requirements.toml").exists(), - "the online refetch must restore the deleted deploy-key policy" - ); - assert!( - kigi_shell::managed_config::managed_policy_gate().is_ok(), - "after a successful heal the deploy-key gate must allow" - ); -} - -/// A confirmed offline team switch (fail_closed team-A install, then team B signs in with no -/// network): the gate PURGES team A's now-foreign artifacts and marker, then PERMITS team B — -/// a legitimate switch is neither refused nor left running under team A's lingering policy. -#[tokio::test] -#[serial] -async fn identity_change_permits_offline_team_switch_and_purges_prior_team() { - let home = test_home().clone(); - reset(&home); - - let body = serde_json::json!({ - "deployment_id": serde_json::Value::Null, - "team_id": "team-a", - "managed_config": TEAM_MANAGED, - "requirements": format!("fail_closed = true\n{TEAM_REQUIREMENTS}"), - }) - .to_string(); - let (url, _auths) = spawn_mock(body); - write_config(&home, &url); - write_team_auth(&home, "team-a"); - kigi_shell::managed_config::sync() - .await - .expect("team-a sync should succeed"); - assert!(home.join("requirements.toml").exists()); - assert!( - kigi_shell::managed_config::managed_policy_gate().is_ok(), - "team A's intact fail_closed policy must start" - ); - - // Switch to team B while OFFLINE: the 5xx server means the session-start refresh cannot - // purge via the apply path; the gate's own purge must handle the switch. - let (err_url, _c, _a) = spawn_mock_seq(vec![(500, "{}".to_string())]); - write_config(&home, &err_url); - write_team_auth(&home, "team-b"); - let auth_manager = std::sync::Arc::new(kigi_shell::auth::AuthManager::new( - &home, - kigi_shell::auth::GrokComConfig::default(), - )); - kigi_shell::managed_config::ensure_managed_policy_present(&auth_manager).await; - - assert!( - kigi_shell::managed_config::managed_policy_gate().is_ok(), - "a legitimate offline team switch must not fail closed" - ); - assert!( - !home.join("requirements.toml").exists(), - "team A's enforced requirements must be purged on the switch" - ); - assert!( - !home.join("managed_config.toml").exists(), - "team A's managed_config must be purged on the switch" - ); - assert!( - !home.join("managed_config_cache.json").exists(), - "team A's sync marker must be purged on the switch" - ); -} - -/// The gate purge takes the managed-config lock best-effort and SKIPS on contention (the holder -/// owns the transition). While held, an A→B switch retains team A's files — and the gate still -/// permits (a pure identity mismatch is not gate-grade tamper); once released, the next gate -/// call purges — proving the skip was contention-driven, not a silent no-op. -#[tokio::test] -#[serial] -async fn gate_purge_skips_while_lock_contended() { - let home = test_home().clone(); - reset(&home); - - let body = serde_json::json!({ - "deployment_id": serde_json::Value::Null, - "team_id": "team-a", - "managed_config": TEAM_MANAGED, - "requirements": format!("fail_closed = true\n{TEAM_REQUIREMENTS}"), - }) - .to_string(); - let (url, _auths) = spawn_mock(body); - write_config(&home, &url); - write_team_auth(&home, "team-a"); - kigi_shell::managed_config::sync() - .await - .expect("team A sync should succeed"); - assert!(home.join("requirements.toml").exists()); - assert!(home.join("managed_config_cache.json").exists()); - - write_team_auth(&home, "team-b"); - - // Hold the managed-config flock (the same lock the gate purge tries), so the purge skips. - let lock = std::fs::OpenOptions::new() - .read(true) - .write(true) - .create(true) - .truncate(false) - .open(home.join("managed_config.lock")) - .unwrap(); - lock.lock().unwrap(); - - assert!( - kigi_shell::managed_config::managed_policy_gate().is_ok(), - "a contended purge skip leaves a pure identity mismatch, which must not refuse" - ); - assert!( - home.join("requirements.toml").exists(), - "team A requirements must be RETAINED while the lock is contended (purge skipped)" - ); - assert!( - home.join("managed_config_cache.json").exists(), - "team A marker must be RETAINED while the lock is contended" - ); - - // Release the lock: the next gate call acquires it and purges team A. - lock.unlock().unwrap(); - assert!( - kigi_shell::managed_config::managed_policy_gate().is_ok(), - "after the lock releases, the gate purges team A and team B starts" - ); - assert!( - !home.join("requirements.toml").exists(), - "the uncontended gate purges team A's requirements" - ); - assert!( - !home.join("managed_config_cache.json").exists(), - "the uncontended gate purges team A's marker" - ); -} - -/// A blank `team_id` in `auth.json` (a parse blip / malformed write) is "unknown", not a -/// distinct identity: the gate must NOT fail closed and the purge must NOT shed team A's -/// policy. Guards the blank→None map in `active_team_id_any_expiry` and the detector's -/// blank guard end to end. -#[tokio::test] -#[serial] -async fn blank_team_id_neither_fails_closed_nor_purges() { - let home = test_home().clone(); - reset(&home); - - let body = serde_json::json!({ - "deployment_id": serde_json::Value::Null, - "team_id": "team-a", - "managed_config": TEAM_MANAGED, - "requirements": format!("fail_closed = true\n{TEAM_REQUIREMENTS}"), - }) - .to_string(); - let (url, _auths) = spawn_mock(body); - write_config(&home, &url); - write_team_auth(&home, "team-a"); - kigi_shell::managed_config::sync() - .await - .expect("team A sync should succeed"); - assert!(home.join("requirements.toml").exists()); - - // auth.json now carries a team principal with a BLANK team_id. - write_team_auth(&home, ""); - - assert!( - kigi_shell::managed_config::managed_policy_gate().is_ok(), - "a blank team_id must read as unknown, not a foreign substituted cache" - ); - assert!( - home.join("requirements.toml").exists(), - "a parse blip must not purge team A's enforced policy" - ); - assert!( - home.join("managed_config_cache.json").exists(), - "the team A marker must be retained on a blank team_id" - ); -} - -/// The session-start gate reads no env: `KIGI_MANAGED_CONFIG_FAIL_CLOSED=0` must NOT disarm a fail_closed -/// refusal (unlike the requirements-layer version check, which the env can only tighten). No local bypass. -#[tokio::test] -#[serial] -async fn fail_closed_env_cannot_disarm_the_gate() { - let home = test_home().clone(); - reset(&home); - - let body = serde_json::json!({ - "deployment_id": "deploy-A", - "managed_config": TEAM_MANAGED, - "requirements": format!("fail_closed = true\n{TEAM_REQUIREMENTS}"), - }) - .to_string(); - let (url, _auths) = spawn_mock(body); - std::fs::write( - home.join("config.toml"), - format!("[endpoints]\nmanaged_config_url = \"{url}\"\ndeployment_key = \"KEY-AAA\"\n"), - ) - .unwrap(); - kigi_shell::managed_config::sync() - .await - .expect("deploy-key sync should succeed"); - - // Tamper: delete the served requirements (offline this fails closed). - std::fs::remove_file(home.join("requirements.toml")).unwrap(); - - // SAFETY: #[serial] test; both vars restored before the assertion below. - unsafe { - std::env::set_var("KIGI_MANAGED_CONFIG", "0"); // offline (the gate is sync anyway) - std::env::set_var("KIGI_MANAGED_CONFIG_FAIL_CLOSED", "0"); // attempt a local disarm - } - let gate = kigi_shell::managed_config::managed_policy_gate(); - unsafe { - std::env::remove_var("KIGI_MANAGED_CONFIG"); - std::env::remove_var("KIGI_MANAGED_CONFIG_FAIL_CLOSED"); - } - - assert!( - gate.is_err(), - "KIGI_MANAGED_CONFIG_FAIL_CLOSED=0 must not disarm the session-start gate (no local bypass)" - ); -} - -/// Full logout removes the team scope from `auth.json`; the post-logout clear -/// (what `perform_logout` runs) removes the orphaned team-sourced files. -#[tokio::test] -#[serial] -async fn logout_clears_team_config() { - let home = test_home().clone(); - reset(&home); - - let (url, _last_auth) = spawn_mock(team_config_body()); - write_config(&home, &url); - write_team_auth(&home, "team-007"); - kigi_shell::managed_config::sync() - .await - .expect("sync should succeed"); - assert!(home.join("managed_config.toml").exists()); - assert!( - home.join("managed_config_cache.json").exists(), - "a successful sync writes the sync-marker cache" - ); - - // `AuthManager::clear` deletes auth.json when the last scope is removed. - std::fs::remove_file(home.join("auth.json")).unwrap(); - kigi_shell::managed_config::clear_orphan(); - - assert!( - !home.join("managed_config.toml").exists(), - "team-sourced managed_config should be cleared on logout" - ); - assert!( - !home.join("requirements.toml").exists(), - "enforced requirements should be cleared on logout" - ); - assert!( - !home.join("managed_config_cache.json").exists(), - "the sync-marker cache should be cleared on logout too" - ); -} - -/// An expired token for a still-signed-in team is not a logout: cold-start -/// tokens are routinely expired before refresh, so the clear is expiry-agnostic. -#[test] -#[serial] -fn cold_start_expired_token_keeps_config() { - let home = test_home().clone(); - reset(&home); - - std::fs::write(home.join("managed_config.toml"), TEAM_MANAGED).unwrap(); - std::fs::write(home.join("requirements.toml"), TEAM_REQUIREMENTS).unwrap(); - write_team_auth_expiry(&home, "team-007", PAST); - kigi_shell::managed_config::clear_orphan(); - - assert!( - home.join("managed_config.toml").exists(), - "expired-but-present team token must not wipe enforced managed_config" - ); - assert!( - home.join("requirements.toml").exists(), - "expired-but-present team token must not wipe enforced requirements" - ); -} - -/// Fail-closed: an UNREADABLE (corrupt) auth.json is not a logout — the clear -/// must keep the team's enforced files until the read recovers. -#[test] -#[serial] -fn unreadable_auth_keeps_config() { - let home = test_home().clone(); - reset(&home); - - std::fs::write(home.join("requirements.toml"), TEAM_REQUIREMENTS).unwrap(); - std::fs::write(home.join("auth.json"), "{corrupt json").unwrap(); - kigi_shell::managed_config::clear_orphan(); - - assert!( - home.join("requirements.toml").exists(), - "an unreadable auth.json must not wipe enforced policy" - ); -} - -#[tokio::test] -#[serial] -async fn deployment_key_wins_over_team_when_both_present() { - let home = test_home().clone(); - reset(&home); - - let (url, auths) = spawn_mock(team_config_body()); - std::fs::write( - home.join("config.toml"), - format!("[endpoints]\nmanaged_config_url = \"{url}\"\ndeployment_key = \"dep-key-123\"\n"), - ) - .unwrap(); - write_team_auth(&home, "team-007"); - - let wrote = kigi_shell::managed_config::sync() - .await - .expect("sync should succeed"); - assert!(wrote); - assert_eq!( - auths.lock().unwrap().last().map(String::as_str), - Some("Bearer dep-key-123"), - "deployment key must win over the team token" - ); -} - -/// A successful deploy-key sync records the served `deployment_id` as `principal` and a non-empty -/// `key_fingerprint` (one-way hash), never the raw key — so a switched key stops serving the prior config. -#[tokio::test] -#[serial] -async fn deployment_key_sync_records_principal_and_key_fingerprint() { - let home = test_home().clone(); - reset(&home); - - let body = serde_json::json!({ - "deployment_id": "dep-42", - "managed_config": "[cli]\ntheme = \"dark\"\n", - "requirements": "[features]\nweb_fetch = false\n", - }) - .to_string(); - let (url, _auths) = spawn_mock(body); - std::fs::write( - home.join("config.toml"), - format!( - "[endpoints]\nmanaged_config_url = \"{url}\"\ndeployment_key = \"dep-key-secret\"\n" - ), - ) - .unwrap(); - - let wrote = kigi_shell::managed_config::sync() - .await - .expect("deployment-key sync should succeed"); - assert!(wrote); - - let marker = std::fs::read_to_string(home.join("managed_config_cache.json")).unwrap(); - let v: serde_json::Value = serde_json::from_str(&marker).unwrap(); - assert_eq!( - v["principal"].as_str(), - Some("dep-42"), - "deploy-key marker records the served deployment_id as principal: {marker}" - ); - let fp = v["key_fingerprint"] - .as_str() - .expect("deploy-key marker records a key fingerprint"); - assert!(!fp.is_empty(), "the key fingerprint must be non-empty"); - assert!( - !marker.contains("dep-key-secret"), - "the raw deployment key must never be written to disk: {marker}" - ); -} - -/// A configured deployment key keeps its files even with no team signed in — -/// the orphan clear must never delete a deployment-key install's config. -#[test] -#[serial] -fn deployment_key_config_survives_clear() { - let home = test_home().clone(); - reset(&home); - - // A deployment-key install: files present + the key persisted in config.toml. - std::fs::write( - home.join("config.toml"), - "[endpoints]\ndeployment_key = \"dep-key-123\"\n", - ) - .unwrap(); - std::fs::write( - home.join("managed_config.toml"), - "[cli]\ninstaller = \"internal\"\n", - ) - .unwrap(); - let _ = std::fs::remove_file(home.join("auth.json")); - - kigi_shell::managed_config::clear_orphan(); - - assert!( - home.join("managed_config.toml").exists(), - "deployment-key managed_config must survive the orphan clear" - ); -} - -/// A personal (non-team) OAuth login is not eligible: no bearer sent, nothing -/// written. Guards the `is_team_principal()` eligibility check. -#[tokio::test] -#[serial] -async fn personal_login_is_noop() { - let home = test_home().clone(); - reset(&home); - - let (url, count, _) = spawn_mock_seq(vec![(200, team_config_body())]); - write_config(&home, &url); - // A signed-in USER principal (no team_id, principal_type absent). - let scope = kigi_shell::auth::GrokComConfig::default().auth_scope(); - let auth = serde_json::json!({ - scope: { - "key": "personal-token", - "auth_mode": "oidc", - "create_time": "2026-01-01T00:00:00Z", - "expires_at": FAR_FUTURE, - "user_id": "user-1", - } - }); - std::fs::write(home.join("auth.json"), auth.to_string()).unwrap(); - - let wrote = kigi_shell::managed_config::sync() - .await - .expect("personal login → no-op, not an error"); - assert!(!wrote, "a personal login must not fetch team config"); - assert_eq!(*count.lock().unwrap(), 0, "no bearer sent to the endpoint"); - assert!(!home.join("managed_config.toml").exists()); -} - -/// Security guard: a lock-skipped apply (dk row HAS config) must not be read as -/// an empty row and fall through to the team token on a deployment-key machine. -#[tokio::test] -#[serial] -async fn lock_contention_does_not_fall_through_to_team() { - let home = test_home().clone(); - reset(&home); - - // dk row returns real config; a team principal is also present. - let (url, count, auths) = spawn_mock_seq(vec![(200, team_config_body())]); - std::fs::write( - home.join("config.toml"), - format!("[endpoints]\nmanaged_config_url = \"{url}\"\ndeployment_key = \"dep-key\"\n"), - ) - .unwrap(); - write_team_auth(&home, "team-007"); - - // Hold the managed-config lock (same flock the client uses) so apply_fetched - // skips and returns Ok(false). - let lock = std::fs::OpenOptions::new() - .read(true) - .write(true) - .create(true) - .truncate(false) - .open(home.join("managed_config.lock")) - .unwrap(); - lock.lock().unwrap(); - - let wrote = kigi_shell::managed_config::sync() - .await - .expect("contended sync is a no-op, not an error"); - lock.unlock().unwrap(); - - assert!(!wrote, "nothing applied while the lock is held"); - // The dk fetch happened; the team token must NOT have been tried as a - // fallthrough (that would fetch the team's config onto a dk machine). - assert_eq!( - auths.lock().unwrap().as_slice(), - ["Bearer dep-key"], - "contention must not trigger the dk->team fallthrough" - ); - assert_eq!(*count.lock().unwrap(), 1); -} - -/// `grok setup` with config served but the lock held by another writer reports -/// Installed (the holder is persisting it), not NothingConfigured. -#[tokio::test] -#[serial] -async fn setup_lock_skip_is_not_reported_as_no_config() { - let home = test_home().clone(); - reset(&home); - - let (url, _auths) = spawn_mock(team_config_body()); - write_config(&home, &url); - write_team_auth(&home, "team-007"); - - let lock = std::fs::OpenOptions::new() - .read(true) - .write(true) - .create(true) - .truncate(false) - .open(home.join("managed_config.lock")) - .unwrap(); - lock.lock().unwrap(); - - let outcome = kigi_shell::managed_config::run_setup().await; - lock.unlock().unwrap(); - - assert!( - matches!(outcome, kigi_shell::managed_config::SetupOutcome::Installed), - "served config with the lock held must not report NothingConfigured" - ); -} - -/// A transient (5xx) failure is retried; once the server recovers, the config -/// is written. (Backoff is overridden to 10ms in `test_home`.) -#[tokio::test] -#[serial] -async fn sync_retries_after_transient_error() { - let home = test_home().clone(); - reset(&home); - - let (url, count, _auths) = - spawn_mock_seq(vec![(500, "boom".into()), (200, team_config_body())]); - write_config(&home, &url); - write_team_auth(&home, "team-007"); - - let wrote = kigi_shell::managed_config::sync() - .await - .expect("should retry the 500 then succeed"); - assert!(wrote); - assert!(home.join("managed_config.toml").exists()); - assert!( - *count.lock().unwrap() >= 2, - "expected a retry after the transient 500" - ); -} - -/// A body-phase interruption — the server writes valid headers then drops before the -/// body completes ("connection closed before message completed", a `reqwest` body error, -/// NOT a decode error) — must be classified transient and recovered on a fresh -/// connection. Pre-fix this mapped to non-retryable `InvalidResponse` and would NOT -/// retry; here `sync()` succeeds. (Backoff is 10ms via `test_home`.) -#[tokio::test] -#[serial] -async fn sync_retries_after_body_phase_drop() { - let home = test_home().clone(); - reset(&home); - - let (url, count) = spawn_mock_truncating_body_first(1, team_config_body()); - write_config(&home, &url); - write_team_auth(&home, "team-007"); - - let wrote = kigi_shell::managed_config::sync() - .await - .expect("a mid-body drop must be retried on a fresh connection, then succeed"); - assert!(wrote); - assert!(home.join("managed_config.toml").exists()); - assert!( - *count.lock().unwrap() >= 2, - "expected a retry on a new connection after the first body read was interrupted" - ); -} - -/// Classification: an in-flight connection interruption must NOT be misreported as -/// an unreachable server — the error must not blame the user's network (the bug), -/// and must surface the transient-interruption wording instead. -#[tokio::test] -#[serial] -async fn connection_drop_is_not_reported_as_unreachable() { - let home = test_home().clone(); - reset(&home); - - // Every connection is dropped mid-flight, so all retries fail identically. - let (url, _count) = spawn_mock_closing_first(usize::MAX, team_config_body()); - write_config(&home, &url); - write_team_auth(&home, "team-007"); - - let err = kigi_shell::managed_config::sync() - .await - .expect_err("all connections dropped → the fetch fails"); - let msg = err.to_string().to_lowercase(); - assert!( - !msg.contains("check your network"), - "an in-flight interruption must not be misreported as unreachable: {msg}" - ); - assert!( - msg.contains("interrupted") || msg.contains("timed out"), - "the message must describe a transient connection interruption: {msg}" - ); -} - -/// The payload side of the body split: a 200 with a non-JSON body is a malformed payload, not a -/// transport interruption — it must fail TERMINALLY (`InvalidResponse`), not retry. Guards the -/// `from_slice` arm of the split; the transport arm is covered by `sync_retries_after_body_phase_drop`. -#[tokio::test] -#[serial] -async fn sync_fails_terminally_on_malformed_payload() { - let home = test_home().clone(); - reset(&home); - - // A 200 whose body is not JSON (e.g. an HTML error page slipped through a proxy). - let (url, count, _auths) = spawn_mock_seq(vec![(200, "not json".into())]); - write_config(&home, &url); - write_team_auth(&home, "team-007"); - - let err = kigi_shell::managed_config::sync() - .await - .expect_err("a malformed payload must fail, not write config"); - assert_eq!( - *count.lock().unwrap(), - 1, - "a malformed payload is terminal and must not be retried" - ); - let msg = err.to_string().to_lowercase(); - assert!( - msg.contains("unexpected response"), - "a malformed payload must surface as an unexpected-response error, not a transient one: {msg}" - ); - assert!(!home.join("managed_config.toml").exists()); -} - -// note: this asserts the POOLING half of the `shared_client()` tuning. Evicting/recovering a -// half-dead h2 connection (and the `send_with_retry_escaping_pool` fresh-final-attempt escape) isn't -// deterministically simulable with the std `TcpListener` harness, so it's covered by the e2e pass. - -/// The tuned pooled `shared_client()` reuses one TCP connection across back-to-back requests (pool -/// eviction is time-based, so two quick requests share a connection). Regression guard against a -/// future mis-tuning that disables pooling for the general-purpose client. -#[tokio::test] -#[serial] -async fn shared_client_reuses_pooled_connection() { - // Scrub the suite's env (notably HTTP(S)_PROXY) before `shared_client()` is built, so a proxy - // can't intercept the 127.0.0.1 mock and break the accept count regardless of test ordering. - let _ = test_home(); - let (base_url, accepts, _heads) = spawn_counting_server().await; - let client = kigi_shell::http::shared_client(); - - client - .get(base_url.as_str()) - .send() - .await - .expect("first request succeeds") - .bytes() - .await - .expect("first body reads to completion"); - // Brief pause so the idle connection is checked back into the pool before the second request. - tokio::time::sleep(std::time::Duration::from_millis(50)).await; - client - .get(base_url.as_str()) - .send() - .await - .expect("second request succeeds") - .bytes() - .await - .expect("second body reads to completion"); - - assert_eq!( - accepts.load(std::sync::atomic::Ordering::SeqCst), - 1, - "the pooled client must reuse one TCP connection across back-to-back requests" - ); -} - -/// An auth rejection is terminal: fail fast with the team-tailored message and -/// no retries (a bad credential won't fix itself by retrying). -#[tokio::test] -#[serial] -async fn sync_fails_fast_on_auth_error_without_retry() { - let home = test_home().clone(); - reset(&home); - - let (url, count, _auths) = spawn_mock_seq(vec![(401, "{\"error\":\"nope\"}".into())]); - write_config(&home, &url); - write_team_auth(&home, "team-007"); - - let err = kigi_shell::managed_config::sync() - .await - .expect_err("401 should be an error"); - assert_eq!(*count.lock().unwrap(), 1, "auth error must not be retried"); - let msg = err.to_string().to_lowercase(); - assert!(msg.contains("team sign-in"), "team-tailored message: {msg}"); -} - -/// A rejected deployment key (stale env/config leftover) must not starve a -/// valid team sign-in: the sync falls back to the team session token. -#[tokio::test] -#[serial] -async fn rejected_deployment_key_falls_back_to_team() { - let home = test_home().clone(); - reset(&home); - - let (url, count, auths) = spawn_mock_seq(vec![ - (401, "{\"error\":\"bad key\"}".into()), // deployment key attempt - (200, team_config_body()), // team fallback - ]); - std::fs::write( - home.join("config.toml"), - format!("[endpoints]\nmanaged_config_url = \"{url}\"\ndeployment_key = \"stale-key\"\n"), - ) - .unwrap(); - write_team_auth(&home, "team-007"); - - let wrote = kigi_shell::managed_config::sync() - .await - .expect("team fallback should succeed after the key is rejected"); - - assert!(wrote); - assert_eq!(*count.lock().unwrap(), 2, "key attempt + team attempt"); - assert_eq!( - auths.lock().unwrap().as_slice(), - ["Bearer stale-key", "Bearer team-session-token"], - "first the rejected key, then the team token" - ); -} - -/// A deployment key whose response is EMPTY (no row provisioned) must not -/// starve the signed-in team: the sync falls through to the team token. -#[tokio::test] -#[serial] -async fn empty_deployment_response_falls_through_to_team() { - let home = test_home().clone(); - reset(&home); - - let (url, count, auths) = spawn_mock_seq(vec![ - (200, "{}".into()), // deployment key: no row - (200, team_config_body()), // team fallback - ]); - std::fs::write( - home.join("config.toml"), - format!("[endpoints]\nmanaged_config_url = \"{url}\"\ndeployment_key = \"dep-key\"\n"), - ) - .unwrap(); - write_team_auth(&home, "team-007"); - - let wrote = kigi_shell::managed_config::sync() - .await - .expect("team fallthrough should succeed"); - - assert!(wrote); - assert_eq!(*count.lock().unwrap(), 2); - assert_eq!( - auths.lock().unwrap().as_slice(), - ["Bearer dep-key", "Bearer team-session-token"] - ); -} - -/// A deployment row with empty content (echoed `deployment_id`) still owns the -/// machine: fallthrough gates on existence, not content, so the team token isn't -/// tried. -#[tokio::test] -#[serial] -async fn empty_content_deployment_row_does_not_fall_through_to_team() { - let home = test_home().clone(); - reset(&home); - - // 200 with a row (deployment_id present) but empty content; team config is - // queued second so a wrong fallthrough would be observable. - let degraded = serde_json::json!({ - "deployment_id": "dep-1", - "managed_config": "", - "requirements": serde_json::Value::Null, - }) - .to_string(); - let (url, count, auths) = spawn_mock_seq(vec![(200, degraded), (200, team_config_body())]); - std::fs::write( - home.join("config.toml"), - format!("[endpoints]\nmanaged_config_url = \"{url}\"\ndeployment_key = \"dep-key\"\n"), - ) - .unwrap(); - write_team_auth(&home, "team-007"); - - let wrote = kigi_shell::managed_config::sync() - .await - .expect("degraded dk row is a no-op, not an error"); - - assert!(!wrote, "empty content writes nothing"); - assert_eq!( - *count.lock().unwrap(), - 1, - "a real dk row must not trigger the team fetch" - ); - assert_eq!( - auths.lock().unwrap().as_slice(), - ["Bearer dep-key"], - "a provisioned-but-empty dk row must not fall through to the team token" - ); - assert!(!home.join("managed_config.toml").exists()); -} - -/// `KIGI_MANAGED_CONFIG=0` is an explicit opt-out: the post-login sync must -/// make zero requests. -#[tokio::test] -#[serial] -async fn managed_config_opt_out_makes_no_requests() { - let home = test_home().clone(); - reset(&home); - - let (url, count, _) = spawn_mock_seq(vec![(200, team_config_body())]); - write_config(&home, &url); - write_team_auth(&home, "team-007"); - - // SAFETY: #[serial] test; restored before returning. - unsafe { std::env::set_var("KIGI_MANAGED_CONFIG", "0") }; - let outcome = kigi_shell::managed_config::post_login_sync(None).await; - unsafe { std::env::remove_var("KIGI_MANAGED_CONFIG") }; - - assert_eq!( - outcome, - kigi_shell::managed_config::ManagedConfigSync::Skipped - ); - - assert_eq!(*count.lock().unwrap(), 0, "opt-out must suppress the fetch"); - assert!(!home.join("managed_config.toml").exists()); -} - -/// Post-login pins the just-authenticated principal over a different on-disk -/// team — a login to team B can't sync team A's policy. -#[tokio::test] -#[serial] -async fn post_login_pins_authenticated_team_over_disk() { - let home = test_home().clone(); - reset(&home); - - let (url, auths) = spawn_mock(team_config_body()); - write_config(&home, &url); - // The on-disk "current" team uses the default "team-session-token". - write_team_auth(&home, "team-disk"); - - // But we just authenticated as a different team with a distinct token. - let pinned: kigi_shell::auth::GrokAuth = serde_json::from_value(serde_json::json!({ - "key": "pinned-token", - "auth_mode": "oidc", - "create_time": "2026-01-01T00:00:00Z", - "expires_at": FAR_FUTURE, - "user_id": "user-1", - "principal_type": "Team", - "team_id": "team-pinned", - })) - .unwrap(); - - let outcome = kigi_shell::managed_config::post_login_sync(Some(pinned)).await; - assert_eq!( - outcome, - kigi_shell::managed_config::ManagedConfigSync::Updated { is_team: true } - ); - assert_eq!( - auths.lock().unwrap().last().map(String::as_str), - Some("Bearer pinned-token"), - "must authenticate as the pinned principal, not the on-disk current team" - ); -} - -/// The login-path sync stops after its small retry budget (2), not the full -/// background budget (5), and never surfaces an error. -#[tokio::test] -#[serial] -async fn post_login_sync_is_latency_bounded() { - let home = test_home().clone(); - reset(&home); - - let (url, count, _auths) = spawn_mock_seq(vec![(500, "boom".into())]); - write_config(&home, &url); - write_team_auth(&home, "team-007"); - - let outcome = kigi_shell::managed_config::post_login_sync(None).await; - - assert_eq!( - outcome, - kigi_shell::managed_config::ManagedConfigSync::Failed - ); - assert_eq!( - *count.lock().unwrap(), - 2, - "login sync must stop after its bounded retry budget" - ); - assert!(!home.join("managed_config.toml").exists()); -} - -/// A deploy key is local config any process can write, not a signed-in identity — so on a dk -/// machine the gate purge must never fire, even when `auth.json` shows a confirmed team switch -/// underneath: purging would let any local process shed the key's policy offline. -#[tokio::test] -#[serial] -async fn deploy_key_machine_never_gate_purges_on_team_switch() { - let home = test_home().clone(); - reset(&home); - - // Sync as team A, then go offline and switch auth.json to team B — the purge-eligible shape. - let body = serde_json::json!({ - "deployment_id": serde_json::Value::Null, - "team_id": "team-a", - "managed_config": TEAM_MANAGED, - "requirements": TEAM_REQUIREMENTS, - }) - .to_string(); - let (url, _auths) = spawn_mock(body); - write_config(&home, &url); - write_team_auth(&home, "team-a"); - kigi_shell::managed_config::sync() - .await - .expect("team-a sync should succeed"); - assert!(home.join("managed_config.toml").exists()); - - // Same switch as the purging sibling test, but with a deployment key configured: the - // serving identity resolves to the key, so the Team-only purge path must not run. - let (err_url, _c, _a) = spawn_mock_seq(vec![(500, "{}".to_string())]); - std::fs::write( - home.join("config.toml"), - format!( - "[endpoints]\nmanaged_config_url = \"{err_url}\"\ndeployment_key = \"dk-under-test\"\n" - ), - ) - .unwrap(); - write_team_auth(&home, "team-b"); - let auth_manager = std::sync::Arc::new(kigi_shell::auth::AuthManager::new( - &home, - kigi_shell::auth::GrokComConfig::default(), - )); - kigi_shell::managed_config::ensure_managed_policy_present(&auth_manager).await; - - // The gate is the purge's only caller — without this call the guard is unexercised. - assert!( - kigi_shell::managed_config::managed_policy_gate().is_ok(), - "dk gate must permit" - ); - assert!( - home.join("managed_config.toml").exists(), - "a dk machine must keep its policy across an auth.json team flip" - ); - assert!( - home.join("managed_config_cache.json").exists(), - "the sync marker must survive too — the key, not the team, owns this machine's policy" - ); -} diff --git a/crates/codegen/kigi-shell/tests/test_settings_refresh.rs b/crates/codegen/kigi-shell/tests/test_settings_refresh.rs index 2b6131a..498bbaa 100644 --- a/crates/codegen/kigi-shell/tests/test_settings_refresh.rs +++ b/crates/codegen/kigi-shell/tests/test_settings_refresh.rs @@ -122,7 +122,7 @@ async fn test_fetch_settings_blocking_round_trip() { .expect("start mock server"); // Without settings configured: returns None (404 from mock) - let auth = kigi_shell::auth::GrokAuth { + let auth = kigi_shell::auth::KimiAuth { key: "test-key".into(), ..Default::default() }; diff --git a/crates/codegen/kigi-tui/src/acp/mod.rs b/crates/codegen/kigi-tui/src/acp/mod.rs index 07001d7..3f19cee 100644 --- a/crates/codegen/kigi-tui/src/acp/mod.rs +++ b/crates/codegen/kigi-tui/src/acp/mod.rs @@ -326,7 +326,7 @@ pub async fn connect_via_leader( // agent's disk-rotated token under the file lock (`try_adopt_disk_token`). let auth_manager = std::sync::Arc::new(kigi_shell::auth::AuthManager::new( &kigi_shell::util::kigi_home::kigi_home(), - agent_config.grok_com_config.clone(), + agent_config.kimi_code_config.clone(), )); Ok(AcpConnection { @@ -897,11 +897,7 @@ mod tests { // Realistic enterprise user: no cached session token, default `grok.com` // login (no enterprise OIDC). has_cached_token: false, - has_enterprise_oidc: false, - enterprise_oidc_issuer: None, login_label: None, - has_auth_provider_command: false, - preferred_method: None, }); let (needs, label, method_id, mode) = startup_auth_metadata(&built.methods); diff --git a/crates/codegen/kigi-tui/src/acp/spawn.rs b/crates/codegen/kigi-tui/src/acp/spawn.rs index aeeca7f..21546c1 100644 --- a/crates/codegen/kigi-tui/src/acp/spawn.rs +++ b/crates/codegen/kigi-tui/src/acp/spawn.rs @@ -40,9 +40,9 @@ pub async fn spawn_grok_shell( ) -> Result { let auth_manager = std::sync::Arc::new(AuthManager::new( &kigi_home(), - agent_config.grok_com_config.clone(), + agent_config.kimi_code_config.clone(), )); - auth_manager.configure_refresher(agent_config.grok_com_config.auth_provider_command.clone()); + auth_manager.configure_refresher(); // Pause token refreshes across system sleep so an OIDC refresh can't // straddle a suspend (which can revoke the refresh token and force // re-login). No-op where the OS listener is unavailable. diff --git a/crates/codegen/kigi-tui/src/app/acp_handler/tests/settings.rs b/crates/codegen/kigi-tui/src/app/acp_handler/tests/settings.rs index db8ddd7..89e6c03 100644 --- a/crates/codegen/kigi-tui/src/app/acp_handler/tests/settings.rs +++ b/crates/codegen/kigi-tui/src/app/acp_handler/tests/settings.rs @@ -19,7 +19,8 @@ )); assert!(!app.is_api_key_auth); assert!(app.usage_visible); - assert!(!app.tier_restricted_commands.is_empty()); + // Tier gating no longer exists; nothing gets re-restricted. + assert!(app.tier_restricted_commands.is_empty()); // A paid tier after API Key clears the api-key flag and tier limits. let mut app = make_app_with_agent("sess-paid-tier"); diff --git a/crates/codegen/kigi-tui/src/app/app_view.rs b/crates/codegen/kigi-tui/src/app/app_view.rs index bf4c55f..bffc134 100644 --- a/crates/codegen/kigi-tui/src/app/app_view.rs +++ b/crates/codegen/kigi-tui/src/app/app_view.rs @@ -403,16 +403,10 @@ pub(crate) const TIER_RESTRICTED_COMMANDS: &[&str] = &["usage", "imagine", "imag /// "x_basic"). Everything else — paid tiers and unknown future names — /// is unrestricted (fail-open). /// -/// The string classification is shared with the shell's capability -/// (toolset) gate via [`kigi_shell::tier::is_restricted_tier_name`] so -/// the two can't drift. The pager's *cosmetic* slash-command gate treats an -/// absent tier (`None`) as restricted (it recovers live on the next settings -/// update); the shell's capability gate treats absence as unrestricted. -fn is_restricted_tier(tier: Option<&str>) -> bool { - match tier { - None => true, - Some(t) => kigi_shell::tier::is_restricted_tier_name(t), - } +/// Tier gating was an xAI concept; the Kimi Code subscription has no +/// client-visible tier, so nothing is ever restricted. +fn is_restricted_tier(_tier: Option<&str>) -> bool { + false } /// True for API-key labels from shell/CCP: `"ApiKey"`, `"API Key"`, `"api_key"`. pub(crate) fn is_api_key_label(s: &str) -> bool { @@ -961,26 +955,18 @@ impl AppView { label: rs.gate_label.clone(), }) } - /// Apply typed auth metadata from the shell. + /// Apply typed auth metadata from the shell. The Kimi auth model carries + /// no team/tier/gate info; those fields only ever come from remote + /// settings now. pub fn apply_auth_meta(&mut self, meta: &kigi_shell::auth::AuthMeta) { self.pending_gate_verification = None; let was_gated = self.gate.is_some(); - self.team_id = meta.team_id.clone(); - self.team_name = meta.team_name.clone(); - self.is_zdr = meta.is_zdr; - self.team_role = meta.team_role.clone(); - self.coding_data_retention_opt_out = meta.coding_data_retention_opt_out; - self.gate = meta.gate.clone(); - if was_gated && self.gate.is_none() { + self.gate = None; + if was_gated { self.paywall_check_started = None; } - self.subscription_tier = meta.subscription_tier.clone(); - self.is_api_key_auth = meta.auth_mode.as_deref().is_some_and(is_api_key_label) - || meta - .subscription_tier - .as_deref() - .is_some_and(is_api_key_label); - self.usage_visible = meta.team_name.is_none() && !self.is_api_key_auth; + self.is_api_key_auth = meta.auth_mode.as_deref().is_some_and(is_api_key_label); + self.usage_visible = !self.is_api_key_auth; self.apply_tier_restrictions(); if let Some(show) = meta.show_resolved_model { self.show_resolved_model = show; @@ -5588,19 +5574,6 @@ pub(crate) mod tests { assert_eq!(counts.get("t_seen"), Some(&2)); } #[test] - fn apply_auth_meta_hides_usage_for_team_users() { - let mut app = test_app(); - assert!(app.usage_visible); - let meta = kigi_shell::auth::AuthMeta { - team_id: Some("team-uuid".into()), - team_name: Some("Acme Corp".into()), - ..Default::default() - }; - app.apply_auth_meta(&meta); - assert!(!app.usage_visible); - assert_eq!(app.team_id.as_deref(), Some("team-uuid")); - } - #[test] fn apply_auth_meta_shows_usage_for_personal_users() { let mut app = test_app(); app.usage_visible = false; @@ -5617,41 +5590,6 @@ pub(crate) mod tests { assert!(!app.is_api_key_auth); assert!(app.usage_visible); } - #[test] - fn apply_auth_meta_api_key_skips_tier_gate() { - let mut app = test_app(); - advertise_media_tools(&mut app); - app.apply_auth_meta(&kigi_shell::auth::AuthMeta { - auth_mode: Some("ApiKey".into()), - subscription_tier: Some("API Key".into()), - ..Default::default() - }); - assert!(app.is_api_key_auth); - assert!(!app.usage_visible); - assert!(app.tier_restricted_commands.is_empty()); - assert_tier_restricted_commands_present(&app); - let mut app = test_app(); - app.apply_auth_meta(&kigi_shell::auth::AuthMeta { - subscription_tier: Some("api_key".into()), - ..Default::default() - }); - assert!(app.is_api_key_auth); - assert!(app.tier_restricted_commands.is_empty()); - app.apply_auth_meta(&kigi_shell::auth::AuthMeta { - auth_mode: Some("Oidc".into()), - subscription_tier: Some("Free".into()), - ..Default::default() - }); - assert!(!app.is_api_key_auth); - assert!(app.usage_visible); - assert!(!app.tier_restricted_commands.is_empty()); - } - fn expected_tier_restricted_commands() -> Vec { - TIER_RESTRICTED_COMMANDS - .iter() - .map(|n| (*n).to_string()) - .collect() - } /// Make every tier-restricted command visible on the welcome prompt so the /// present/absent assertions exercise the deny list, not incidental /// fail-closed hiding: @@ -5668,16 +5606,6 @@ pub(crate) mod tests { .collect(), ); } - fn assert_tier_restricted_commands_absent(app: &AppView) { - let reg = app.welcome_prompt.slash_controller.registry(); - for name in TIER_RESTRICTED_COMMANDS { - assert!( - reg.get(name).is_none(), - "/{name} must be denied on a restricted tier" - ); - } - assert!(reg.get("cost").is_none(), "/cost alias must be denied"); - } fn assert_tier_restricted_commands_present(app: &AppView) { let reg = app.welcome_prompt.slash_controller.registry(); for name in TIER_RESTRICTED_COMMANDS { @@ -5688,103 +5616,31 @@ pub(crate) mod tests { } } #[test] - fn apply_auth_meta_restricts_usage_for_free_tier() { + fn apply_auth_meta_never_restricts_tiers() { let mut app = test_app(); advertise_media_tools(&mut app); app.apply_auth_meta(&kigi_shell::auth::AuthMeta::default()); - assert_eq!( - app.tier_restricted_commands, - expected_tier_restricted_commands() - ); - assert_tier_restricted_commands_absent(&app); - assert!(app.usage_visible); - } - #[test] - fn apply_auth_meta_restricts_usage_for_x_basic_tier() { - let mut app = test_app(); - advertise_media_tools(&mut app); - let meta = kigi_shell::auth::AuthMeta { - subscription_tier: Some("X Basic".into()), - ..Default::default() - }; - app.apply_auth_meta(&meta); - assert_eq!( - app.tier_restricted_commands, - expected_tier_restricted_commands() - ); - assert_tier_restricted_commands_absent(&app); - } - #[test] - fn apply_auth_meta_lifts_restrictions_for_paid_tiers_and_teams() { - let mut app = test_app(); - advertise_media_tools(&mut app); - let meta = kigi_shell::auth::AuthMeta { - subscription_tier: Some("SuperGrok".into()), - ..Default::default() - }; - app.apply_auth_meta(&meta); assert!(app.tier_restricted_commands.is_empty()); assert_tier_restricted_commands_present(&app); - let mut app = test_app(); - advertise_media_tools(&mut app); - app.apply_auth_meta(&kigi_shell::auth::AuthMeta::default()); - assert!(!app.tier_restricted_commands.is_empty()); - app.subscription_tier = Some("SuperGrok".into()); - app.apply_tier_restrictions(); - assert!(app.tier_restricted_commands.is_empty()); - assert_tier_restricted_commands_present(&app); - let mut app = test_app(); - let meta = kigi_shell::auth::AuthMeta { - team_id: Some("team-uuid".into()), - team_name: Some("Acme Corp".into()), - ..Default::default() - }; - app.apply_auth_meta(&meta); - assert!(app.tier_restricted_commands.is_empty()); } #[test] - fn is_restricted_tier_classification() { - assert!(is_restricted_tier(None)); - assert!(is_restricted_tier(Some(""))); - assert!(is_restricted_tier(Some("Free"))); - assert!(is_restricted_tier(Some("X Basic"))); - assert!(is_restricted_tier(Some("x_basic"))); - assert!(!is_restricted_tier(Some("SuperGrok"))); - assert!(!is_restricted_tier(Some("SuperGrok Heavy"))); - assert!(!is_restricted_tier(Some("X Premium"))); - assert!(!is_restricted_tier(Some("X Premium+"))); + fn is_restricted_tier_never_restricts() { + assert!(!is_restricted_tier(None)); + assert!(!is_restricted_tier(Some("Free"))); assert!(!is_restricted_tier(Some("SomeFutureTier"))); } #[test] - fn apply_auth_meta_clears_gate_on_subscription() { + fn apply_auth_meta_clears_gate_on_login() { let mut app = test_app(); app.gate = Some(kigi_shell::auth::GateInfo { - message: "Subscribe to use Grok Build".into(), - url: Some("https://grok.com/supergrok?referrer=grok-build".into()), - label: None, - }); - assert!(app.is_access_blocked()); - let meta = kigi_shell::auth::AuthMeta::default(); - app.apply_auth_meta(&meta); - assert!(app.gate.is_none()); - assert!(app.has_access()); - } - #[test] - fn apply_auth_meta_gate_unchanged_when_still_gated() { - let mut app = test_app(); - let gate = kigi_shell::auth::GateInfo { message: "Subscribe".into(), url: None, label: None, - }; - app.gate = Some(gate.clone()); - let meta = kigi_shell::auth::AuthMeta { - gate: Some(gate), - ..Default::default() - }; - app.apply_auth_meta(&meta); - assert!(app.gate.is_some()); + }); assert!(app.is_access_blocked()); + app.apply_auth_meta(&kigi_shell::auth::AuthMeta::default()); + assert!(app.gate.is_none()); + assert!(app.has_access()); } #[test] fn welcome_ctrl_q_requires_confirmation() { diff --git a/crates/codegen/kigi-tui/src/app/cli.rs b/crates/codegen/kigi-tui/src/app/cli.rs index 77189b5..b033052 100644 --- a/crates/codegen/kigi-tui/src/app/cli.rs +++ b/crates/codegen/kigi-tui/src/app/cli.rs @@ -19,29 +19,8 @@ pub enum Command { Leader(LeaderMgmtArgs), /// Sign out and clear cached credentials Logout, - /// Sign in - Login { - /// Ignored (kept for backwards compatibility). OAuth2 is now the only auth method. - #[arg(long, hide = true)] - legacy: bool, - /// Use Grok OAuth via auth.x.ai. - #[arg(long = "oauth", alias = "oidc", conflicts_with_all = ["device_auth"])] - oauth: bool, - /// Use device-code authentication for headless/remote environments. - #[arg( - long = "device-auth", - visible_alias = "device-code", - conflicts_with_all = ["oauth"] - )] - device_auth: bool, - /// Authenticate for remote development environments (hidden). - /// - /// Field is always present so match arms stay feature-unification-safe - /// across Bazel/cargo graphs; clap only registers `--devbox` when - /// `devbox-login` is enabled (`arg(skip)` otherwise → always false). - #[arg(skip)] - devbox: bool, - }, + /// Sign in with your Kimi Code subscription (device-code flow) + Login, /// Manage MCP server configurations Mcp(crate::mcp_cmd::McpArgs), /// Manage plugins diff --git a/crates/codegen/kigi-tui/src/app/dispatch/auth.rs b/crates/codegen/kigi-tui/src/app/dispatch/auth.rs index e01f511..ca4bdbd 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/auth.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/auth.rs @@ -45,15 +45,9 @@ pub(super) fn ensure_login_method(app: &mut AppView) { // No interactive method: leave login_method_id unset (fail-closed). } -/// Error when no interactive login method is available (empty auth_methods, -/// e.g. `preferred_method=api_key` with no credentials). Prefer the shell's -/// pin-unavailable copy when the list is empty. -fn no_login_method_error(app: &AppView) -> String { - if app.auth_methods.is_empty() { - kigi_shell::agent::auth_method::PREFERRED_API_KEY_UNAVAILABLE.to_string() - } else { - "No login method available".to_string() - } +/// Error when no interactive login method is available (empty auth_methods). +fn no_login_method_error(_app: &AppView) -> String { + "No login method available".to_string() } /// Log out, then start a new login flow in a single sequential task. diff --git a/crates/codegen/kigi-tui/src/app/dispatch/billing.rs b/crates/codegen/kigi-tui/src/app/dispatch/billing.rs index b36b65a..aff6c81 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/billing.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/billing.rs @@ -454,32 +454,21 @@ pub(super) fn handle_credit_limit_recheck_complete( agent_id: AgentId, meta: Option, ) -> Vec { - let old_tier = app.subscription_tier.clone(); if let Some(meta_val) = meta && let Ok(auth_meta) = serde_json::from_value::(meta_val) { app.apply_auth_meta(&auth_meta); } - let tier_changed = app.subscription_tier != old_tier && app.subscription_tier.is_some(); let Some(agent) = app.agents.get_mut(&agent_id) else { return vec![]; }; // If the user already submitted another prompt while the - // recheck was in flight, don't retry the stashed one — they've - // moved on. The tier update (above) still takes effect. + // recheck was in flight, don't show the upsell — they've moved on. let user_moved_on = !agent.session.state.is_idle() || !agent.session.pending_prompts.is_empty(); - if tier_changed && !user_moved_on { - if let Some(prompt) = agent.credit_limit_stashed_prompt.take() { - let tier_name = app.subscription_tier.as_deref().unwrap_or("a higher tier"); - agent.scrollback.push_block(RenderBlock::system(format!( - "Subscription upgraded to {tier_name}. Retrying\u{2026}" - ))); - agent.session.enqueue_in_flight_prompt_front(prompt); - } - } else if !user_moved_on { + if !user_moved_on { let balance = agent .credit_balance .as_ref() diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/auth.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/auth.rs index 6148a7a..60c087c 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/auth.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/auth.rs @@ -240,9 +240,9 @@ fn login_with_empty_auth_methods_fails_closed() { matches!( &app.auth_state, AuthState::Pending { error: Some(msg) } - if msg.contains("preferred_method=api_key") + if msg.contains("No login method available") ), - "must surface pin-unavailable error, got {:?}", + "must surface no-login-method error, got {:?}", app.auth_state ); assert!(app.login_method_id.is_none()); diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/billing.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/billing.rs index 17b84dc..2e347ab 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/billing.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/billing.rs @@ -60,50 +60,6 @@ fn dispatch_billing( ); } -#[test] -fn credit_limit_retry_preserves_image_submission_state() { - let mut app = test_app_with_agent(); - let mut image = crate::prompt_images::from_clipboard_data(&crate::clipboard::ImageData { - data: vec![1, 2, 3], - mime_type: "image/png".into(), - }); - image.display_number = 1; - let prompt = crate::app::agent::InFlightPrompt { - text: "retry [Image #1]".into(), - images: vec![image], - scrollback_entry: crate::scrollback::EntryId::new(0), - chip_elements: vec![crate::app::agent::ChipElement { - range: 6..16, - kind: crate::views::prompt_widget::KIND_IMAGE, - display: None, - }], - }; - app.agents - .get_mut(&AgentId(0)) - .unwrap() - .credit_limit_stashed_prompt = Some(prompt); - - let effects = dispatch( - Action::TaskComplete(TaskResult::CreditLimitRecheckComplete { - agent_id: AgentId(0), - meta: Some(serde_json::json!({"subscription_tier": "Upgraded"})), - }), - &mut app, - ); - assert!( - effects - .iter() - .any(|effect| matches!(effect, Effect::SendPromptBlocks { .. })) - ); - let in_flight = app.agents[&AgentId(0)] - .session - .in_flight_prompt - .as_ref() - .unwrap(); - assert_eq!(in_flight.images.len(), 1); - assert_eq!(in_flight.chip_elements.len(), 1); -} - #[test] fn is_max_tier_positive_match() { assert!(is_max_tier(Some("supergrok_heavy"))); diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/task_result.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/task_result.rs index b94bcec..c399650 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/task_result.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/task_result.rs @@ -1646,30 +1646,6 @@ fn verify_check_with_meta_resolves_pending_gate() { assert!(app.pending_gate_verification.is_none()); } -/// The live check confirmed the block (meta WITH a gate): the paywall -/// shows with the authoritative gate. -#[test] -fn verify_check_with_gated_meta_shows_gate() { - let mut app = test_app(); - let _effs = app.impose_gate(test_gate()); - - let meta = serde_json::to_value(kigi_shell::auth::AuthMeta { - gate: Some(test_gate()), - ..Default::default() - }) - .unwrap(); - dispatch_task_result( - TaskResult::CheckSubscriptionComplete { - verify: Some(app.gate_verify_gen), - meta: Some(meta), - }, - &mut app, - ); - - assert!(!app.has_access(), "verified gate must show"); - assert!(app.pending_gate_verification.is_none()); -} - /// The verification's own check failed (meta None) while its stale gate /// was deferred: err on blocking — the deferred gate is promoted. #[test] @@ -1857,59 +1833,6 @@ fn gate_verify_timeout_stale_generation_is_ignored() { ); } -/// A verified gate landing via `CheckSubscriptionComplete` (gated meta while -/// ungated) must arm the 5s paywall auto-check chain — verify-before-paywall -/// paths never went through the login-path chain start. -#[test] -fn verified_gate_via_check_complete_starts_paywall_chain() { - let mut app = test_app(); - let _effs = app.impose_gate(test_gate()); - - let meta = serde_json::to_value(kigi_shell::auth::AuthMeta { - gate: Some(test_gate()), - ..Default::default() - }) - .unwrap(); - let effects = dispatch_task_result( - TaskResult::CheckSubscriptionComplete { - verify: None, - meta: Some(meta), - }, - &mut app, - ); - - assert!(!app.has_access()); - assert!( - app.paywall_check_started.is_some(), - "verified gate must arm the paywall auto-check chain" - ); - assert!( - effects - .iter() - .any(|e| matches!(e, Effect::SchedulePaywallCheck)), - "verified gate must schedule the 5s chain; got: {effects:?}" - ); - - // Steady-state paywall-poller responses (already gated) must NOT fan - // out extra timers. - let meta = serde_json::to_value(kigi_shell::auth::AuthMeta { - gate: Some(test_gate()), - ..Default::default() - }) - .unwrap(); - let effects = dispatch_task_result( - TaskResult::CheckSubscriptionComplete { - verify: None, - meta: Some(meta), - }, - &mut app, - ); - assert!( - effects.is_empty(), - "already-gated check responses must not schedule more timers; got: {effects:?}" - ); -} - /// `GateRefreshed` with gate-free settings while a deferred gate awaits /// verification must drop the pending copy — the fresh settings are newer /// than the stale snapshot that produced it — and still run the lift diff --git a/crates/codegen/kigi-tui/src/app/effects/mod.rs b/crates/codegen/kigi-tui/src/app/effects/mod.rs index fe3da08..7890374 100644 --- a/crates/codegen/kigi-tui/src/app/effects/mod.rs +++ b/crates/codegen/kigi-tui/src/app/effects/mod.rs @@ -3521,7 +3521,7 @@ pub(crate) fn execute( &kigi_home.join("auth.json"), ) .ok()?; - let scope = kigi_shell::auth::GrokComConfig::default() + let scope = kigi_shell::auth::KimiCodeConfig::default() .auth_scope(); let auth = kigi_shell::auth::lookup_auth( &store, diff --git a/crates/codegen/kigi-tui/src/app/event_loop.rs b/crates/codegen/kigi-tui/src/app/event_loop.rs index 2311318..151442d 100644 --- a/crates/codegen/kigi-tui/src/app/event_loop.rs +++ b/crates/codegen/kigi-tui/src/app/event_loop.rs @@ -683,11 +683,8 @@ pub(crate) async fn run( // welcome/auth UI right away. let mut post_render_effects = if needs_interactive_login { if connection.auth_methods.is_empty() { - // preferred_method pin unavailable — no advertised method to start. app.auth_state = super::app_view::AuthState::Pending { - error: Some( - kigi_shell::agent::auth_method::PREFERRED_API_KEY_UNAVAILABLE.to_string(), - ), + error: Some("No login method available".to_string()), }; vec![] } else { diff --git a/crates/codegen/kigi-tui/src/app/mod.rs b/crates/codegen/kigi-tui/src/app/mod.rs index 2eba76d..7822e50 100644 --- a/crates/codegen/kigi-tui/src/app/mod.rs +++ b/crates/codegen/kigi-tui/src/app/mod.rs @@ -350,16 +350,16 @@ pub async fn run( let startup_start = std::time::Instant::now(); let raw_config = kigi_shell::config::load_effective_config() .map_err(|e| anyhow::anyhow!("Failed to load config: {e}"))?; - let grok_com_config = match kigi_shell::agent::config::Config::new_from_toml_cfg(&raw_config) { - Ok(c) => c.grok_com_config, + let kimi_code_config = match kigi_shell::agent::config::Config::new_from_toml_cfg(&raw_config) { + Ok(c) => c.kimi_code_config, Err(e) => { tracing::warn!( error = % e, "failed to parse config for auth refresh, using defaults" ); - kigi_shell::auth::GrokComConfig::default() + kigi_shell::auth::KimiCodeConfig::default() } }; - let refreshed_auth = kigi_shell::auth::try_ensure_fresh_auth(&grok_com_config).await; + let refreshed_auth = kigi_shell::auth::try_ensure_fresh_auth(&kimi_code_config).await; let early_prefetch = kigi_shell::agent::models::start_early_prefetch_with_auth(refreshed_auth); kigi_shell::agent::mvp_agent::warm_async_http_client(); tokio::task::spawn_blocking(|| {}); diff --git a/crates/codegen/kigi-tui/src/app/session_startup.rs b/crates/codegen/kigi-tui/src/app/session_startup.rs index 7dc464a..5f85e55 100644 --- a/crates/codegen/kigi-tui/src/app/session_startup.rs +++ b/crates/codegen/kigi-tui/src/app/session_startup.rs @@ -418,9 +418,9 @@ pub(crate) fn pre_acp_auth_manager( ) -> std::sync::Arc { let auth = std::sync::Arc::new(kigi_shell::auth::AuthManager::new( &kigi_shell::util::kigi_home::kigi_home(), - agent_config.grok_com_config.clone(), + agent_config.kimi_code_config.clone(), )); - auth.configure_refresher(agent_config.grok_com_config.auth_provider_command.clone()); + auth.configure_refresher(); auth } /// Preflight: preferred id must be a UUID and not a persisted session under `cwd`. @@ -621,7 +621,7 @@ async fn resolve_existing_session( use kigi_shell::util::kigi_home::kigi_home; let deployment_key = agent_config.endpoints.deployment_key.clone(); ensure_authenticated_or_noninteractive( - &agent_config.grok_com_config, + &agent_config.kimi_code_config, deployment_key.is_some(), None, ) @@ -629,7 +629,7 @@ async fn resolve_existing_session( .map_err(|e| anyhow::anyhow!("Failed to authenticate for session restore: {}", e))?; let auth_manager = std::sync::Arc::new(AuthManager::new( &kigi_home(), - agent_config.grok_com_config.clone(), + agent_config.kimi_code_config.clone(), )); let registry_client = SessionRegistryClient::new(agent_config.endpoints.proxy_url(), String::new()) diff --git a/crates/codegen/kigi-tui/src/sessions_cmd.rs b/crates/codegen/kigi-tui/src/sessions_cmd.rs index b45cd9f..15134da 100644 --- a/crates/codegen/kigi-tui/src/sessions_cmd.rs +++ b/crates/codegen/kigi-tui/src/sessions_cmd.rs @@ -40,11 +40,11 @@ pub async fn run(args: SessionsArgs, agent_config: &AgentConfig) -> Result<()> { // for these setups), any cached credential will be used. Otherwise we still // proceed so the SessionRegistryClient can use the deployment_key when // talking to the custom proxy. - let auth = try_ensure_fresh_auth(&agent_config.grok_com_config).await; + let auth = try_ensure_fresh_auth(&agent_config.kimi_code_config).await; let auth_manager = std::sync::Arc::new(AuthManager::new( &kigi_home(), - agent_config.grok_com_config.clone(), + agent_config.kimi_code_config.clone(), )); let client = kigi_shell::agent::session_registry_client::SessionRegistryClient::new( @@ -169,7 +169,7 @@ pub async fn run(args: SessionsArgs, agent_config: &AgentConfig) -> Result<()> { // backend delete is idempotent (a `404` is treated as success), // so this is safe for local-only sessions with no remote copy. // ZDR teams never upload, so there is nothing remote to delete. - let needs_remote = auth.as_ref().is_some_and(|a| !a.is_zdr_team()); + let needs_remote = auth.is_some(); // Pass `cwd = None` so the session is found by id regardless of // which workspace it was created in; the local delete still uses diff --git a/crates/codegen/kigi-tui/tests/pty_e2e/subscription_watch_and_gate_verify_pty.rs b/crates/codegen/kigi-tui/tests/pty_e2e/subscription_watch_and_gate_verify_pty.rs deleted file mode 100644 index 3268132..0000000 --- a/crates/codegen/kigi-tui/tests/pty_e2e/subscription_watch_and_gate_verify_pty.rs +++ /dev/null @@ -1,439 +0,0 @@ -// Per-test-case module for the `pty_e2e` integration test crate. -// -// End-to-end coverage for free→paid subscription auto-detection -// (`src/app/subscription.rs`). -#[allow(unused_imports)] -use super::common::*; - -/// Distinctive gate copy (unlikely to collide with welcome chrome). -const GATE_MSG: &str = "ZZSUBGATEMSG"; - -/// A tier in the shell's `QUALIFYING_TIERS` list. -const PAID_TIER: &str = "SuperGrokPro"; - -/// Display name delivered via `/settings` `subscription_tier_display`. -const PAID_TIER_DISPLAY: &str = "SuperGrok Pro"; - -/// Count of live subscription checks the client made against the mock -/// (`GET /v1/user?include=subscription`). Plain `/v1/user` enrichment -/// fetches are deliberately excluded. -fn user_check_count(content: &ContentController) -> usize { - content - .requests() - .iter() - .filter(|e| e.path.starts_with("/v1/user?") && e.path.contains("include=subscription")) - .count() -} - -/// Count of `GET /v1/settings` fetches (the qualifying-tier check refetches -/// settings, so a post-upgrade increase marks detection completing). -fn settings_count(content: &ContentController) -> usize { - content - .requests() - .iter() - .filter(|e| e.path == "/v1/settings") - .count() -} - -/// Count of `GET /v1/models` catalog fetches. Post-upgrade the shell must -/// re-fetch so tier-targeted models land without restart. -fn models_count(content: &ContentController) -> usize { - content - .requests() - .iter() - .filter(|e| e.path == "/v1/models" || e.path.starts_with("/v1/models?")) - .count() -} - -/// Paid-only model id used to prove the post-unblock catalog actually -/// replaced the free list in the picker (not merely that `/v1/models` was hit). -const PAID_ONLY_MODEL: &str = "composer-paid-only"; - -/// Minimal unsigned JWT with a `tier` claim matching [`PAID_TIER`]. -/// -/// Proto `prod_auth.SubscriptionTier`: 5 = `supergrok_heavy` = live -/// `/user` string `SuperGrokPro`. Must match -/// `jwt_claim_matches_user_subscription_tier` or post-unblock catalog -/// refresh treats the claim as stale and never re-fetches `/v1/models` -/// within the test timeout. -fn paid_tier_jwt() -> String { - use base64::Engine; - let enc = |v: &serde_json::Value| { - base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(v.to_string().as_bytes()) - }; - let header = enc(&json!({"alg": "none", "typ": "JWT"})); - let payload = enc(&json!({"sub": "pty-subwatch", "exp": 2_000_000_000u64, "tier": 5})); - format!("{header}.{payload}.sig") -} - -/// Bind the fixed local-dev OIDC issuer (`http://localhost:22255`) and return a -/// paid-tier JWT on refresh. Call **after** free-tier watch polling so early -/// checks still see connection-refused (hermetic free path) and only the -/// post-upgrade refresh succeeds with a paid token. -/// -/// Minimal raw HTTP (no axum dep in this crate): discovery + token only. -async fn start_local_oidc_paid_refresh() -> tokio::task::JoinHandle<()> { - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - - let listener = tokio::net::TcpListener::bind("127.0.0.1:22255") - .await - .expect("bind local OIDC issuer on :22255 for hermetic paid refresh"); - let paid_jwt = paid_tier_jwt(); - let discovery_body = serde_json::to_vec(&json!({ - "authorization_endpoint": "http://localhost:22255/authorize", - "token_endpoint": "http://localhost:22255/token", - })) - .expect("discovery json"); - let token_body = serde_json::to_vec(&json!({ - "access_token": paid_jwt, - "refresh_token": "pty-test-refresh-token-rotated", - "expires_in": 3600, - })) - .expect("token json"); - - tokio::spawn(async move { - loop { - let Ok((mut socket, _)) = listener.accept().await else { - break; - }; - let discovery_body = discovery_body.clone(); - let token_body = token_body.clone(); - tokio::spawn(async move { - let mut buf = vec![0u8; 4096]; - let n = match socket.read(&mut buf).await { - Ok(0) | Err(_) => return, - Ok(n) => n, - }; - let req = String::from_utf8_lossy(&buf[..n]); - let (status, body): (&str, &[u8]) = - if req.starts_with("GET /.well-known/openid-configuration") { - ("200 OK", discovery_body.as_slice()) - } else if req.starts_with("POST /token") { - ("200 OK", token_body.as_slice()) - } else { - ("404 Not Found", br"{}") - }; - let resp = format!( - "HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", - body.len() - ); - let _ = socket.write_all(resp.as_bytes()).await; - let _ = socket.write_all(body).await; - }); - } - }) -} - -/// Pump the PTY until `cond` holds or `timeout` elapses (panics with a -/// screen dump on timeout). -fn pump_until( - harness: &mut PtyHarness, - timeout: Duration, - mut cond: impl FnMut() -> bool, - what: &str, -) { - let deadline = Instant::now() + timeout; - while !cond() { - assert!( - Instant::now() < deadline, - "timed out waiting for {what}\nscreen:\n{}", - harness.screen_contents() - ); - harness.update(Duration::from_millis(100)); - } -} - -/// Like [`seed_fake_oauth`], but under the `KIGI_LOCAL_AUTH` dev issuer -/// (`http://localhost:22255`). Two reasons: `is_xai_oauth2_issuer()` accepts -/// the local issuer, so the subscription gate applies (an enterprise/unknown -/// issuer bypasses it); and the qualifying-tier JWT refresh then hits -/// `localhost:22255` — instant connection-refused instead of a real network -/// call to auth.x.ai (hermetic, no CI-network flake). Pair with -/// `KIGI_LOCAL_AUTH=1` in the spawn env so the shell's scope-key lookup -/// resolves this entry. -fn seed_fake_oauth_local_issuer(content: &ContentController, user: &str) { - let kigi_home = content.home().join(".kigi"); - std::fs::create_dir_all(&kigi_home).expect("create temp .kigi"); - std::fs::write( - kigi_home.join("auth.json"), - format!( - r#"{{ - "http://localhost:22255::b1a00492-073a-47ea-816f-4c329264a828": {{ - "key": "pty-test-oauth-token", - "auth_mode": "oidc", - "create_time": "2026-01-01T00:00:00Z", - "user_id": "{user}", - "email": "{user}@test.invalid", - "expires_at": "2030-01-01T00:00:00Z", - "refresh_token": "pty-test-refresh-token", - "oidc_issuer": "http://localhost:22255", - "oidc_client_id": "b1a00492-073a-47ea-816f-4c329264a828" - }} -}}"# - ), - ) - .expect("seed fake local-issuer oauth auth.json"); -} - -/// Spawn the pager with local-issuer session auth (see -/// [`seed_fake_oauth_local_issuer`]) plus `extra_env`. Does NOT wait for the -/// welcome screen — gate tests assert on the very first paint. -fn spawn_subscription_pager( - content: &ContentController, - oauth_user: &str, - extra_env: &[(&str, &str)], -) -> PtyHarness { - seed_fake_oauth_local_issuer(content, oauth_user); - let env = oauth_env_for_pager(content); - let mut env_refs: Vec<(&str, &str)> = - env.iter().map(|(k, v)| (k.as_str(), v.as_str())).collect(); - env_refs.push(("KIGI_LOCAL_AUTH", "1")); - env_refs.extend_from_slice(extra_env); - - let binary = pager_binary().expect("resolve pager binary"); - PtyHarness::new_in_dir( - &binary, - DEFAULT_ROWS, - DEFAULT_COLS, - &[], - &env_refs, - Some(content.home()), - ) - .expect("spawn pager with subscription session auth") -} - -/// [`spawn_subscription_pager`] driven into a live session -/// (welcome → prompt → mock response). -fn spawn_subscription_session( - content: &ContentController, - oauth_user: &str, - extra_env: &[(&str, &str)], -) -> PtyHarness { - let mut harness = spawn_subscription_pager(content, oauth_user, extra_env); - harness - .wait_for_text(WELCOME_SCREEN_SENTINEL, WELCOME_TIMEOUT) - .expect("welcome text"); - harness - .inject_keys(format!("{PROMPT}\r").as_bytes()) - .expect("submit prompt to enter session"); - harness - .wait_for_text(MOCK_RESPONSE_SENTINEL, Duration::from_secs(30)) - .expect("session response"); - harness -} - -/// Watch cadence while free, upgrade detection, then dormancy once paid. -/// -/// Also covers W-17: after free→paid unblock the shell refreshes the model -/// catalog with a **paid** JWT (mock IdP on `:22255`) and the paid-only model -/// id appears in the `/model` picker — not merely that `/v1/models` was called. -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -#[ignore = "PTY e2e; run the owning pty_e2e_* Cargo test with --ignored (see Cargo.toml)"] -async fn subscription_watch_polls_free_tier_then_goes_dormant_after_upgrade() { - // Start free-targeted (no paid-only model); swap after upgrade. - // OIDC mock is started only after free-phase polling so early refresh - // still connection-refuses (keeps the free watch path hermetic). - let content = ContentController::start_with_models(vec![MockModel::new("grok-3")]) - .await - .expect("start content"); - content.set_response(format!("{MOCK_RESPONSE_SENTINEL} watch cadence.")); - // Free tier: the mock's /v1/user returns no subscription tier by default. - let mut harness = spawn_subscription_session( - &content, - "pty-subwatch", - &[("KIGI_SUBSCRIPTION_WATCH_INTERVAL_SECS", "1")], - ); - - // While free, the watch fires repeatedly at the (test-shrunk) cadence. - pump_until( - &mut harness, - Duration::from_secs(30), - || user_check_count(&content) >= 3, - ">=3 live subscription checks while on the free tier", - ); - - // Now enable hermetic paid JWT refresh for the post-unblock catalog path. - let oidc = start_local_oidc_paid_refresh().await; - - // Server-side upgrade: the live tier flips to a qualifying value and - // settings now carry the paid display tier. Swap the model catalog *before* - // recording models_before so an in-flight free fetch cannot falsely satisfy - // the post-upgrade re-fetch wait. - let settings_before = settings_count(&content); - content.server().set_user_subscription_tier(Some(PAID_TIER)); - content.server().set_settings(json!({ - "allow_access": true, - "subscription_tier_display": PAID_TIER_DISPLAY, - })); - content.server().set_models(vec![ - MockModel::new("grok-3"), - MockModel::new(PAID_ONLY_MODEL), - ]); - let models_before = models_count(&content); - - // Detection: the qualifying check refetches /v1/settings (that's how the - // paid display tier reaches the client and disarms the watch). - pump_until( - &mut harness, - Duration::from_secs(30), - || settings_count(&content) > settings_before, - "settings refetch after the qualifying tier was detected", - ); - - // W-17: after gate lift + successful paid JWT refresh the shell must - // re-fetch /v1/models (fire-and-forget `on_auth_changed`). - pump_until( - &mut harness, - Duration::from_secs(30), - || models_count(&content) > models_before, - "model catalog re-fetch after free→paid subscription unblock", - ); - - // Stronger than a GET count: switch to the paid-only model id. Status bar - // shows it on success (same pattern as same_agent_type_switch_no_modal). - harness - .inject_keys(format!("/model {PAID_ONLY_MODEL}\r").as_bytes()) - .expect("switch to paid-only model"); - harness - .wait_for_text(PAID_ONLY_MODEL, Duration::from_secs(20)) - .expect("paid-only model applied after upgrade catalog refresh"); - - // Dormancy: once the paid tier lands, a full 6s quiet window (>=6 - // would-be ticks at the 1s cadence) passes with zero new checks. - let deadline = Instant::now() + Duration::from_secs(60); - loop { - let base = user_check_count(&content); - let window_end = Instant::now() + Duration::from_secs(6); - while Instant::now() < window_end { - harness.update(Duration::from_millis(200)); - } - if user_check_count(&content) == base { - break; - } - assert!( - Instant::now() < deadline, - "watch never went dormant after the upgrade (checks still firing)\nscreen:\n{}", - harness.screen_contents() - ); - } - - harness.quit().expect("clean quit"); - oidc.abort(); -} - -/// A genuinely-free user gated at startup still gets the paywall — but only -/// after a live subscription check confirmed the block (the gate is never -/// painted straight from the stale source). -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -#[ignore = "PTY e2e; run the owning pty_e2e_* Cargo test with --ignored (see Cargo.toml)"] -async fn startup_gate_shows_paywall_for_free_user_after_live_check() { - let content = ContentController::start().await.expect("start content"); - // Gated settings (no allow_access), free user (no subscriptionTier). - content.server().set_settings(json!({ - "gate_message": GATE_MSG, - "gate_url": "https://grok.com/supergrok?referrer=grok-build", - "gate_label": "Subscribe", - })); - - let mut harness = spawn_subscription_pager(&content, "pty-subgate-free", &[]); - - harness - .wait_for_text(WELCOME_SCREEN_SENTINEL, WELCOME_TIMEOUT) - .expect("welcome text"); - // The verified gate renders. Normally the check response resolves the - // deferral within seconds; the budget also covers the 30s hung-check - // safety net under full-suite contention. - harness - .wait_for_text(GATE_MSG, Duration::from_secs(45)) - .expect("gate copy renders for a genuinely-free user"); - - assert!( - user_check_count(&content) >= 1, - "a live subscription check must run before the paywall is shown; requests: {:?}", - content - .requests() - .iter() - .map(|e| e.path.clone()) - .collect::>() - ); - - harness.quit().expect("clean quit"); -} - -/// Verify-before-paywall: a user who ALREADY subscribed never sees a -/// paywall flash when a stale gated settings snapshot reaches the client. -/// -/// The stale snapshot is delivered via the `/new` settings refresh — the -/// only active `/v1/settings` consumer at that point (watch disabled via -/// env, gate poll only runs while gated, startup fetches settled). Queueing it at startup instead races -/// the shell's concurrent startup fetches: a slow gated fetch landing after -/// the verify check stores fresh settings can legitimately re-carry the -/// gate — a time-travel artifact of the scripted one-shot, not a client -/// bug (observed as a flake). -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -#[ignore = "PTY e2e; run the owning pty_e2e_* Cargo test with --ignored (see Cargo.toml)"] -async fn stale_gate_push_never_flashes_paywall_for_subscribed_user() { - let content = ContentController::start().await.expect("start content"); - // Live tier: already paid; steady settings allow access. - content.server().set_user_subscription_tier(Some(PAID_TIER)); - content.server().set_settings(json!({ - "allow_access": true, - "subscription_tier_display": PAID_TIER_DISPLAY, - })); - content.set_response(format!("{MOCK_RESPONSE_SENTINEL} paid path.")); - - // Watch disabled so the deferral's own VerifyPendingGate is the only - // subscription-check traffic (deferral does not depend on the watch). - let mut harness = spawn_subscription_session( - &content, - "pty-subgate-paid", - &[("KIGI_SUBSCRIPTION_WATCH_INTERVAL_SECS", "0")], - ); - - // Let startup fetches fully settle so the scripted one-shot below can - // only be consumed by the /new refresh. - harness.update(Duration::from_secs(2)); - let checks_before = user_check_count(&content); - - // One stale gated snapshot: the "remote settings stale moment". - content.enqueue_response( - "/v1/settings", - ScriptedResponse::json(200, json!({ "gate_message": GATE_MSG })), - ); - harness.inject_keys(b"/new\r").expect("run /new"); - - // Sample the screen across the deferral window: the gate copy must - // never appear. The deferred gate could only surface via a gated check - // result (impossible — the live tier is paid and the fresh settings - // allow) or the 30s hung-check net, which the resolving check disarms. - let end = Instant::now() + Duration::from_secs(8); - while Instant::now() < end { - harness.update(Duration::from_millis(150)); - assert!( - !harness.contains_text(GATE_MSG), - "paywall flashed for an already-subscribed user\nscreen:\n{}", - harness.screen_contents() - ); - } - - // Positive anchors: the deferral's live check actually ran, and the - // session is fully usable (prompt round-trips). - assert!( - user_check_count(&content) > checks_before, - "expected a live subscription check for the deferred gate; requests: {:?}", - content - .requests() - .iter() - .map(|e| e.path.clone()) - .collect::>() - ); - content.set_response(format!("{MOCK_RESPONSE_SENTINEL} still usable.")); - harness - .inject_keys(format!("{PROMPT}\r").as_bytes()) - .expect("submit prompt"); - harness - .wait_for_text(MOCK_RESPONSE_SENTINEL, Duration::from_secs(30)) - .expect("session usable for the subscribed user"); - - harness.quit().expect("clean quit"); -} diff --git a/crates/codegen/kigi-tui/tests/pty_e2e_config_ui.rs b/crates/codegen/kigi-tui/tests/pty_e2e_config_ui.rs index aa4c0ae..23f270e 100644 --- a/crates/codegen/kigi-tui/tests/pty_e2e_config_ui.rs +++ b/crates/codegen/kigi-tui/tests/pty_e2e_config_ui.rs @@ -36,8 +36,6 @@ mod reverse_agent_type_mismatch_cursor_to_default; mod same_agent_type_switch_no_modal; #[path = "pty_e2e/show_thinking_blocks_toggle_hides_existing_pty.rs"] mod show_thinking_blocks_toggle_hides_existing_pty; -#[path = "pty_e2e/subscription_watch_and_gate_verify_pty.rs"] -mod subscription_watch_and_gate_verify_pty; #[path = "pty_e2e/undo_tip_resets_each_new_session.rs"] mod undo_tip_resets_each_new_session; #[path = "pty_e2e/undo_tip_seen_count_never_persisted.rs"] diff --git a/crates/codegen/kigi-update/src/version.rs b/crates/codegen/kigi-update/src/version.rs index cbc43fc..a90a70c 100644 --- a/crates/codegen/kigi-update/src/version.rs +++ b/crates/codegen/kigi-update/src/version.rs @@ -48,7 +48,7 @@ impl UpdateConfig { pub fn from_environment() -> Self { Self { proxy_base_url: kigi_env::coding_api_base_url(), - auth_scope: kigi_shell::auth::GrokComConfig::default().auth_scope(), + auth_scope: kigi_shell::auth::KimiCodeConfig::default().auth_scope(), deployment_key: None, alpha_test_key: None, channel: "stable".to_string(), diff --git a/crates/codegen/kigi-workspace/src/hub_auth.rs b/crates/codegen/kigi-workspace/src/hub_auth.rs index c1a9513..c9170dc 100644 --- a/crates/codegen/kigi-workspace/src/hub_auth.rs +++ b/crates/codegen/kigi-workspace/src/hub_auth.rs @@ -300,7 +300,7 @@ mod tests { "scope": { "key": "eyJhbGciOiJFUzI1NiJ9.tok", "user_id": "u1", - "auth_mode": "oidc", + "auth_mode": "oauth", "create_time": "2026-01-01T00:00:00Z", "email": "test@x.ai", "first_name": "Test", diff --git a/crates/codegen/kigi-workspace/src/session/tool_config.rs b/crates/codegen/kigi-workspace/src/session/tool_config.rs index 293f2bd..8ed6013 100644 --- a/crates/codegen/kigi-workspace/src/session/tool_config.rs +++ b/crates/codegen/kigi-workspace/src/session/tool_config.rs @@ -500,13 +500,6 @@ fn build_proxy_headers(base_url: &str) -> indexmap::IndexMap { format!("kigi-workspace/{version}"), ); headers.insert("x-grok-client-version".to_string(), version.to_string()); - if base_url.contains("cli-chat-proxy") || base_url.contains("chat-proxy") { - headers.insert("X-XAI-Token-Auth".to_string(), "xai-grok-cli".to_string()); - headers.insert( - "x-authenticateresponse".to_string(), - "authenticate-response".to_string(), - ); - } headers } /// Build web fetch config. Enabled with default params unless diff --git a/deny.toml b/deny.toml index 56452dc..bb3d18c 100644 --- a/deny.toml +++ b/deny.toml @@ -39,6 +39,7 @@ allow = [ "CDLA-Permissive-2.0", "ISC", "MIT", + "MIT-0", "MPL-2.0", "OpenSSL", "Unicode-3.0",