Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
119 changes: 82 additions & 37 deletions rust/src/providers/grok/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,8 @@ const CLI_SETTINGS_ENDPOINT: &str = "https://cli-chat-proxy.grok.com/v1/settings
const RESET_CREDITS_TIMEOUT: Duration = Duration::from_secs(2);
const RESET_CREDITS_JOIN_GRACE: Duration = Duration::from_millis(250);
const GROK_CLI_ACP_ARGS: [&str; 2] = ["agent", "stdio"];
/// Shown to the user when no usable Grok login is left after every auto step.
const GROK_SIGN_IN_HINT: &str = "Grok sign-in required. Run `grok login`, then refresh CodexBar.";

/// Reset-window pace (upstream `GrokProviderDescriptor.pace`, v0.70.0): weekly
/// windows inside their length, see [`reset_window_pace_matches`].
Expand Down Expand Up @@ -737,49 +739,59 @@ impl Provider for GrokProvider {
}

async fn fetch_usage(&self, ctx: &FetchContext) -> Result<ProviderFetchResult, ProviderError> {
match ctx.source_mode {
SourceMode::Auto => self.fetch_auto(ctx).await,
SourceMode::Web if ctx.token_account_isolated => {
match ctx
.manual_cookie_header
.as_deref()
.filter(|cookie| !cookie.trim().is_empty())
{
Some(cookie_header) => self.fetch_with_cookie(cookie_header, ctx).await,
None => Err(isolated_account_missing_credential()),
// The sign-in hint is applied to the final result only: auto steps
// inside the block still match on `AuthRequired` to fall through.
let result: Result<ProviderFetchResult, ProviderError> = async {
match ctx.source_mode {
SourceMode::Auto => self.fetch_auto(ctx).await,
SourceMode::Web if ctx.token_account_isolated => {
match ctx
.manual_cookie_header
.as_deref()
.filter(|cookie| !cookie.trim().is_empty())
{
Some(cookie_header) => self.fetch_with_cookie(cookie_header, ctx).await,
None => Err(isolated_account_missing_credential()),
}
}
}
SourceMode::Web => {
if let Some(cookie_header) = &ctx.manual_cookie_header {
return self.fetch_with_cookie(cookie_header, ctx).await;
SourceMode::Web => {
if let Some(cookie_header) = &ctx.manual_cookie_header {
return self.fetch_with_cookie(cookie_header, ctx).await;
}
self.fetch_with_cookie_refresh(ctx).await
}
self.fetch_with_cookie_refresh(ctx).await
}
SourceMode::Cli | SourceMode::OAuth if ctx.token_account_isolated => {
self.fetch_isolated_account(ctx).await
}
SourceMode::Cli => match Self::load_credentials(GrokAuthKind::Cli).await {
Ok(credentials) => {
self.fetch_with_auth(&credentials, GrokAuthKind::Cli, ctx)
.await
SourceMode::Cli | SourceMode::OAuth if ctx.token_account_isolated => {
self.fetch_isolated_account(ctx).await
}
Err(error) => self.fetch_cli_acp_without_login(error, ctx).await,
},
SourceMode::OAuth => {
// Prefer the switched ~/.grok/auth.json over a leftover token
// account so Weekly/notifications follow Grok account Switch.
let credentials = match Self::load_credentials(GrokAuthKind::OAuth).await {
Ok(credentials) => credentials,
Err(error) => {
let Some(token) = ctx.api_key.as_deref() else {
return Err(error);
};
GrokCredentials::from_bearer(token)
SourceMode::Cli => match Self::load_credentials(GrokAuthKind::Cli).await {
Ok(credentials) => {
self.fetch_with_auth(&credentials, GrokAuthKind::Cli, ctx)
.await
}
};
self.fetch_with_oauth_fallback(&credentials, ctx).await
Err(error) => self.fetch_cli_acp_without_login(error, ctx).await,
},
SourceMode::OAuth => {
// Prefer the switched ~/.grok/auth.json over a leftover token
// account so Weekly/notifications follow Grok account Switch.
let credentials = match Self::load_credentials(GrokAuthKind::OAuth).await {
Ok(credentials) => credentials,
Err(error) => {
let Some(token) = ctx.api_key.as_deref() else {
return Err(error);
};
GrokCredentials::from_bearer(token)
}
};
self.fetch_with_oauth_fallback(&credentials, ctx).await
}
}
}
.await;
if sign_in_hint_applies(ctx) {
result.map_err(sign_in_hint_for_terminal_error)
} else {
result
}
}

fn available_sources(&self) -> Vec<SourceMode> {
Expand Down Expand Up @@ -1147,6 +1159,39 @@ fn is_missing_or_expired_login(err: &ProviderError) -> bool {
)
}

/// Maps a sign-in failure that reaches the caller to an error that names the
/// fix. `OAuth` keeps the NeedsAuthentication state and the auth error
/// category (the shape Claude uses for a missing login), unlike `NotInstalled`.
fn sign_in_hint_for_terminal_error(error: ProviderError) -> ProviderError {
match error {
ProviderError::AuthRequired => ProviderError::OAuth(GROK_SIGN_IN_HINT.to_string()),
other => other,
}
}

/// Whether a terminal `AuthRequired` came from the ambient Grok login, so
/// `grok login` is the fix. A supplied API key or manual cookie answers last on
/// its route (auto steps return its result; OAuth falls back to the key), and
/// an isolated account never reads the ambient login, so those keep the
/// original error. The presence checks mirror each route's own checks.
fn sign_in_hint_applies(ctx: &FetchContext) -> bool {
if ctx.token_account_isolated {
return false;
}
let has_api_key = || ctx.api_key.as_deref().is_some();
let has_trimmed = |value: &Option<String>| {
value
.as_deref()
.is_some_and(|value| !value.trim().is_empty())
};
match ctx.source_mode {
SourceMode::Cli => true,
SourceMode::OAuth => !has_api_key(),
SourceMode::Web => ctx.manual_cookie_header.is_none(),
SourceMode::Auto => !has_trimmed(&ctx.api_key) && !has_trimmed(&ctx.manual_cookie_header),
}
}

/// Whether a cookie-path error should invalidate the cached browser session.
fn is_cookie_authentication_failure(err: &ProviderError) -> bool {
matches!(err, ProviderError::AuthRequired)
Expand Down
66 changes: 66 additions & 0 deletions rust/src/providers/grok/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -507,3 +507,69 @@ fn account_usage_marks_informational_windows_unavailable() {
assert!(!usage.usage_available);
assert_eq!(usage.used_percent, None);
}

#[test]
fn terminal_sign_in_failure_names_the_grok_login_fix() {
let error = sign_in_hint_for_terminal_error(ProviderError::AuthRequired);

assert!(error.to_string().contains("grok login"), "{error}");
assert_eq!(
error.state_kind(),
crate::core::ProviderStateKind::NeedsAuthentication
);
assert!(!matches!(error, ProviderError::AuthRequired));
}

#[test]
fn terminal_mapping_leaves_other_errors_unchanged() {
let error = sign_in_hint_for_terminal_error(ProviderError::Timeout);

assert!(matches!(error, ProviderError::Timeout), "{error}");
let other = sign_in_hint_for_terminal_error(ProviderError::Other("kept".to_string()));
assert_eq!(other.to_string(), "kept");
}

fn hint_context(
source_mode: SourceMode,
api_key: Option<&str>,
manual_cookie_header: Option<&str>,
token_account_isolated: bool,
) -> FetchContext {
FetchContext {
source_mode,
api_key: api_key.map(str::to_string),
manual_cookie_header: manual_cookie_header.map(str::to_string),
token_account_isolated,
..FetchContext::default()
}
}

#[test]
fn sign_in_hint_applies_to_ambient_login_routes() {
for ctx in [
hint_context(SourceMode::Cli, None, None, false),
hint_context(SourceMode::Cli, Some("key"), Some("sso=x"), false),
hint_context(SourceMode::OAuth, None, Some("sso=x"), false),
hint_context(SourceMode::Web, Some("key"), None, false),
hint_context(SourceMode::Auto, None, None, false),
hint_context(SourceMode::Auto, Some(" "), Some(" "), false),
] {
assert!(sign_in_hint_applies(&ctx), "{:?}", ctx.source_mode);
}
}

#[test]
fn sign_in_hint_skips_supplied_and_isolated_credentials() {
for ctx in [
hint_context(SourceMode::Web, None, Some("sso=x"), false),
hint_context(SourceMode::OAuth, Some("key"), None, false),
hint_context(SourceMode::Auto, Some("key"), None, false),
hint_context(SourceMode::Auto, None, Some("sso=x"), false),
hint_context(SourceMode::Auto, Some("key"), None, true),
hint_context(SourceMode::Web, None, Some("sso=x"), true),
hint_context(SourceMode::Cli, Some("key"), None, true),
hint_context(SourceMode::OAuth, None, Some("sso=x"), true),
] {
assert!(!sign_in_hint_applies(&ctx), "{:?}", ctx.source_mode);
}
}