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

M1/F1: Kimi Code OAuth device-code flow

Browse Source 
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.
This commit is contained in:
ZacharyZhang-NY
2026-07-17 07:37:29 -04:00
parent d6c20fc13f
commit 021b82443d
117 changed files with 4052 additions and 19900 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/codegen/kigi-tui/src/acp/mod.rs
+1 -5
View File
@@ -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);
crates/codegen/kigi-tui/src/acp/spawn.rs
+2 -2
View File
@@ -40,9 +40,9 @@ pub async fn spawn_grok_shell(
) -> Result<SpawnedAgent> {
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.
crates/codegen/kigi-tui/src/app/acp_handler/tests/settings.rs
+2 -1
View File
@@ -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");
crates/codegen/kigi-tui/src/app/app_view.rs
+20 -164
View File
@@ -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<String> {
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() {
crates/codegen/kigi-tui/src/app/cli.rs
+2 -23
View File
@@ -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
crates/codegen/kigi-tui/src/app/dispatch/auth.rs
+3 -9
View File
@@ -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.
crates/codegen/kigi-tui/src/app/dispatch/billing.rs
+2 -13
View File
@@ -454,32 +454,21 @@ pub(super) fn handle_credit_limit_recheck_complete(
agent_id: AgentId,
meta: Option<serde_json::Value>,
) -> Vec<Effect> {
let old_tier = app.subscription_tier.clone();
if let Some(meta_val) = meta
&& let Ok(auth_meta) = serde_json::from_value::<kigi_shell::auth::AuthMeta>(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()
crates/codegen/kigi-tui/src/app/dispatch/tests/auth.rs
+2 -2
View File
@@ -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());
crates/codegen/kigi-tui/src/app/dispatch/tests/billing.rs
-44
View File
@@ -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")));
crates/codegen/kigi-tui/src/app/dispatch/tests/task_result.rs
-77
View File
@@ -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
crates/codegen/kigi-tui/src/app/effects/mod.rs
+1 -1
View File
@@ -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,
crates/codegen/kigi-tui/src/app/event_loop.rs
+1 -4
View File
@@ -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 {
crates/codegen/kigi-tui/src/app/mod.rs
+4 -4
View File
@@ -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(|| {});
crates/codegen/kigi-tui/src/app/session_startup.rs
+4 -4
View File
@@ -418,9 +418,9 @@ pub(crate) fn pre_acp_auth_manager(
) -> std::sync::Arc<kigi_shell::auth::AuthManager> {
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())
crates/codegen/kigi-tui/src/sessions_cmd.rs
+3 -3
View File
@@ -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
crates/codegen/kigi-tui/tests/pty_e2e/subscription_watch_and_gate_verify_pty.rs
-439
View File
@@ -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::<Vec<_>>()
);
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::<Vec<_>>()
);
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");
}
crates/codegen/kigi-tui/tests/pty_e2e_config_ui.rs
-2
View File
@@ -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"]