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-shell/src/agent/mvp_agent/acp_agent.rs
+18 -269
View File
@@ -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
}
crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs
+30 -129
View File
@@ -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<crate::auth::GrokAuth> {
pub(super) fn current_or_buffered_auth(&self) -> Option<crate::auth::KimiAuth> {
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<u64> {
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<crate::auth::GrokAuth> {
/// Return the current KimiAuth credentials, if authenticated and not expired.
pub(crate) fn current_auth(&self) -> Option<crate::auth::KimiAuth> {
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<acp::AuthMethodId> {
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<AuthenticateResponse, acp::Error> {
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<crate::util::config::RemoteSettings> {
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<crate::http::OriginClientInfo>,
) -> 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<SamplingConfig> {
@@ -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
crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs
+10 -479
View File
@@ -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<String> {
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<String>,
auth: Option<&crate::auth::GrokAuth>,
) -> Option<String> {
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<Vec<()>> {
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<RefCell<std::collections::HashSet<PathBuf>>>,
/// 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<bool>,
/// 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<std::sync::atomic::AtomicBool>,
/// 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<std::sync::atomic::AtomicBool>,
/// 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<std::sync::atomic::AtomicBool>,
}
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<crate::auth::AuthManager>,
models_manager: crate::agent::models::ModelsManager,
in_flight: Arc<std::sync::atomic::AtomicBool>,
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(
crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs
+25 -386
View File
@@ -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<kigi_acp_lib::AcpClientMessage>,
) {
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);