Skip to content
Open
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
7 changes: 7 additions & 0 deletions src/crates/adapters/ai-adapters/AGENTS.md
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,13 @@ For the auth/discovery path, use `cargo test -p openbitfun-ai-adapters --feature
subscription-auth --lib`. Device-grant timing tests use the dev-only Tokio
test clock and synthetic tokens; they do not authorize real accounts.

xAI OAuth resolves device and token endpoints from its OpenID discovery document
at the start of each login or refresh. Keep the issuer and endpoint origin pinned
to `https://auth.x.ai`; discovery failure must remain explicit. A pending device
grant retains its discovered token endpoint throughout polling. Discovery uses
the host's subscription HTTP proxy options and does not change inference routes
or persisted credential shapes.

```bash
cargo test -p openbitfun-agent-stream
cargo test -p openbitfun-ai-adapters
Expand Down
134 changes: 126 additions & 8 deletions src/crates/adapters/ai-adapters/src/subscription_auth/grok.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@ use std::time::Duration;
use tokio_util::sync::CancellationToken;

const CLIENT_ID: &str = "b1a00492-073a-47ea-816f-4c329264a828";
const DEVICE_AUTHORIZATION_URL: &str = "https://auth.x.ai/oauth2/device/code";
const TOKEN_URL: &str = "https://auth.x.ai/oauth2/token";
const ISSUER: &str = "https://auth.x.ai";
const DISCOVERY_URL: &str = "https://auth.x.ai/.well-known/openid-configuration";
const DEVICE_CODE_GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type:device_code";
const SCOPE: &str = "openid profile email offline_access grok-cli:access api:access";
const XAI_BASE_URL: &str = "https://api.x.ai/v1";
Expand All @@ -30,6 +30,63 @@ const SHORT_TOKEN_REFRESH_LEEWAY_MS: i64 = 2 * 60 * 1000;
const LONG_TOKEN_REFRESH_LEEWAY_MS: i64 = 60 * 60 * 1000;
const SHORT_TOKEN_THRESHOLD_MS: i64 = 45 * 60 * 1000;

#[derive(Debug, Deserialize)]
struct DiscoveryDocument {
issuer: String,
device_authorization_endpoint: String,
token_endpoint: String,
}

struct OAuthEndpoints {
device_authorization: reqwest::Url,
token: reqwest::Url,
}

fn trusted_auth_endpoint(value: &str) -> Result<reqwest::Url> {
let url = reqwest::Url::parse(value).context("parse xAI authentication endpoint")?;
if value.chars().any(|character| character.is_ascii_control())
|| url.scheme() != "https"
|| url.host_str() != Some("auth.x.ai")
|| url.port_or_known_default() != Some(443)
|| !url.username().is_empty()
|| url.password().is_some()
|| url.fragment().is_some()
{
return Err(anyhow!(
"xAI discovery returned an untrusted authentication endpoint"
));
}
Ok(url)
}

fn discovery_endpoints(document: DiscoveryDocument) -> Result<OAuthEndpoints> {
if document.issuer.trim_end_matches('/') != ISSUER {
return Err(anyhow!(
"xAI discovery issuer does not match the expected issuer"
));
}
Ok(OAuthEndpoints {
device_authorization: trusted_auth_endpoint(&document.device_authorization_endpoint)?,
token: trusted_auth_endpoint(&document.token_endpoint)?,
})
}

async fn discover_endpoints(options: &SubscriptionHttpOptions) -> Result<OAuthEndpoints> {
let response = oauth_request(http_client(options)?.get(DISCOVERY_URL))
.send()
.await
.context("fetch xAI OpenID discovery document")?;
if !response.status().is_success() {
return Err(anyhow!("xAI discovery failed: HTTP {}", response.status()));
}
discovery_endpoints(
response
.json()
.await
.context("parse xAI discovery document")?,
)
}

#[derive(Debug, Deserialize)]
struct DeviceCodeResponse {
device_code: String,
Expand Down Expand Up @@ -141,9 +198,12 @@ fn validate_verification_url(url: &str) -> Result<()> {
Ok(())
}

async fn request_device_code(options: &SubscriptionHttpOptions) -> Result<DeviceCodeResponse> {
async fn request_device_code(
options: &SubscriptionHttpOptions,
endpoints: &OAuthEndpoints,
) -> Result<DeviceCodeResponse> {
let client = http_client(options)?;
let response = oauth_request(client.post(DEVICE_AUTHORIZATION_URL))
let response = oauth_request(client.post(endpoints.device_authorization.clone()))
.form(&[
("client_id", CLIENT_ID),
("scope", SCOPE),
Expand Down Expand Up @@ -210,9 +270,10 @@ fn classify_device_poll_error(
async fn poll_once(
device_code: &str,
options: &SubscriptionHttpOptions,
endpoints: &OAuthEndpoints,
) -> Result<DevicePoll<TokenResponse>> {
let client = http_client(options)?;
let response = oauth_request(client.post(TOKEN_URL))
let response = oauth_request(client.post(endpoints.token.clone()))
.form(&[
("grant_type", DEVICE_CODE_GRANT_TYPE),
("client_id", CLIENT_ID),
Expand Down Expand Up @@ -295,8 +356,9 @@ async fn persist_tokens(tokens: TokenResponse, expected_revision: u64) -> Result
}

async fn refresh(refresh_token: &str, options: &SubscriptionHttpOptions) -> Result<TokenResponse> {
let endpoints = discover_endpoints(options).await?;
let client = http_client(options)?;
let response = oauth_request(client.post(TOKEN_URL))
let response = oauth_request(client.post(endpoints.token))
.form(&[
("grant_type", "refresh_token"),
("refresh_token", refresh_token),
Expand All @@ -323,7 +385,8 @@ pub(crate) async fn begin_login(
expected_revision: u64,
options: SubscriptionHttpOptions,
) -> Result<StartedLogin> {
let device = request_device_code(&options).await?;
let endpoints = discover_endpoints(&options).await?;
let device = request_device_code(&options, &endpoints).await?;
let interval = positive_seconds(device.interval, DEFAULT_POLL_INTERVAL_SECS);
let expires_in = positive_seconds(device.expires_in, DEFAULT_DEVICE_LIFETIME_SECS)
.min(super::LOGIN_TIMEOUT.as_secs() as i64);
Expand All @@ -344,7 +407,7 @@ pub(crate) async fn begin_login(
Duration::from_secs(expires_in as u64),
Duration::from_secs(3),
true,
|| poll_once(&device_code, &options),
|| poll_once(&device_code, &options, &endpoints),
)
.await
.context("complete xAI device authorization")
Expand Down Expand Up @@ -485,6 +548,61 @@ mod tests {
XAI_BASE_URL, XAI_REQUEST_URL,
};

#[test]
fn discovery_accepts_changed_paths_and_ignores_unrelated_metadata() {
let document = serde_json::from_value(serde_json::json!({
"issuer": "https://auth.x.ai/",
"device_authorization_endpoint": "https://auth.x.ai/new/device",
"token_endpoint": "https://auth.x.ai/new/token",
"authorization_endpoint": "https://auth.x.ai/authorize"
}))
.unwrap();
let endpoints = super::discovery_endpoints(document).unwrap();
assert_eq!(
endpoints.device_authorization.as_str(),
"https://auth.x.ai/new/device"
);
assert_eq!(endpoints.token.as_str(), "https://auth.x.ai/new/token");
}

#[test]
fn discovery_rejects_wrong_issuer_missing_fields_and_untrusted_endpoints() {
for invalid in [
"http://auth.x.ai/token",
"https://auth.x.ai:8443/token",
"https://user@auth.x.ai/token",
"https://auth.x.ai.evil.test/token",
"https://attacker.example/token",
"https://auth.x.ai/token#fragment",
"https://auth.x.ai/\ntoken",
] {
for field in ["token_endpoint", "device_authorization_endpoint"] {
let mut value = serde_json::json!({
"issuer": "https://auth.x.ai",
"device_authorization_endpoint": "https://auth.x.ai/device",
"token_endpoint": "https://auth.x.ai/token"
});
value[field] = serde_json::json!(invalid);
assert!(
super::discovery_endpoints(serde_json::from_value(value).unwrap()).is_err()
);
}
}
let document = serde_json::from_value(serde_json::json!({
"issuer": "https://attacker.example",
"device_authorization_endpoint": "https://auth.x.ai/device",
"token_endpoint": "https://auth.x.ai/token"
}))
.unwrap();
assert!(super::discovery_endpoints(document).is_err());
assert!(
serde_json::from_value::<super::DiscoveryDocument>(serde_json::json!({
"issuer": "https://auth.x.ai", "token_endpoint": "https://auth.x.ai/token"
}))
.is_err()
);
}

#[test]
fn accepts_https_verification_urls_and_rejects_unsafe_urls() {
validate_verification_url("https://accounts.x.ai/oauth2/device?user_code=ABCD-EFGH")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,7 @@ import { getActiveSurfaceScope } from '@/infrastructure/peer-device/deviceSurfac
import { LONG_CONTEXT_WARNING_THRESHOLD_TOKENS } from '@/shared/constants/modelContext';
import {
preferredSubscriptionLoginMethod,
subscriptionLoginMethodsForSurface,
settleSubscriptionLoginStart,
subscriptionLoginRequiresLocalDevice,
SubscriptionLoginCoordinator,
Expand Down Expand Up @@ -1185,7 +1186,10 @@ const ModelSettingsPage: React.FC = () => {
};
}, []);

const handleSubscriptionLogin = useCallback(async (provider: SubscriptionProvider) => {
const handleSubscriptionLogin = useCallback(async (
provider: SubscriptionProvider,
method?: SubscriptionLoginMethod,
) => {
if ((!isTauriRuntime() || isPeerDeviceModeActive()) && subscriptionLoginRequiresLocalDevice(provider)) {
notification.error(t('subscriptionAuth.peerLoginRequiresLocalDevice'));
return;
Expand All @@ -1195,7 +1199,7 @@ const ModelSettingsPage: React.FC = () => {
// newer provider's state or leaving an undiscoverable backend session.
const operation = loginCoordinatorRef.current.begin(provider);
if (!operation) return;
const requestedMethod = preferredSubscriptionLoginMethod(
const requestedMethod = method ?? preferredSubscriptionLoginMethod(
provider,
isTauriRuntime() && !isPeerDeviceModeActive(),
);
Expand Down Expand Up @@ -3629,6 +3633,12 @@ const ModelSettingsPage: React.FC = () => {
const isRefreshing = refreshingSubscriptionProviders.has(account.provider);
const isLoggingIn = loggingInProvider === account.provider;
const anyLoginInProgress = loggingInProvider !== null;
const loginMethods = subscriptionLoginMethodsForSurface(
account.provider,
account.login_methods,
isTauriRuntime() && !isPeerDeviceModeActive(),
);
const loginChoices = loginMethods.length ? loginMethods : [undefined];
const loginPanel = subscriptionLoginPanel?.provider === account.provider
? subscriptionLoginPanel
: null;
Expand Down Expand Up @@ -3693,17 +3703,26 @@ const ModelSettingsPage: React.FC = () => {
{t('subscriptionAuth.retryVault')}
</Button>
) : (
<Button
size="sm"
variant="primary"
loading={isLoggingIn}
disabled={anyLoginInProgress}
onClick={() => void handleSubscriptionLogin(account.provider)}
>
{t(loginPanel?.status === 'failed'
? 'subscriptionAuth.retryLogin'
: 'subscriptionAuth.login')}
</Button>
<>
{loginChoices.map((method, index) => (
<Button
key={method ?? 'default'}
size="sm"
variant={index === 0 ? 'primary' : 'outline'}
loading={isLoggingIn && subscriptionLoginPanel?.method === method}
disabled={anyLoginInProgress}
onClick={() => void handleSubscriptionLogin(account.provider, method)}
>
{t(method === 'device'
? 'subscriptionAuth.deviceLogin'
: method === 'browser'
? 'subscriptionAuth.browserLogin'
: loginPanel?.status === 'failed'
? 'subscriptionAuth.retryLogin'
: 'subscriptionAuth.login')}
</Button>
))}
</>
)}
{isLoggingIn && (
<Button
Expand Down
Original file line number Diff line number Diff line change
@@ -1,11 +1,35 @@
import { describe, expect, it } from 'vitest';
import {
preferredSubscriptionLoginMethod,
subscriptionLoginMethodsForSurface,
settleSubscriptionLoginStart,
subscriptionLoginRequiresLocalDevice,
SubscriptionLoginCoordinator,
} from './subscriptionLoginCoordinator';

describe('subscriptionLoginMethodsForSurface', () => {
it('exposes both advertised Codex methods locally and device login through a peer', () => {
expect(subscriptionLoginMethodsForSurface('codex', ['browser', 'device'], true))
.toEqual(['browser', 'device']);
expect(subscriptionLoginMethodsForSurface('codex', ['browser', 'device'], false))
.toEqual(['device']);
});

it('preserves the device-only providers and browser-only Antigravity contract', () => {
for (const provider of ['opencode', 'grok', 'hermes'] as const) {
expect(subscriptionLoginMethodsForSurface(provider, ['device'], true)).toEqual(['device']);
expect(subscriptionLoginMethodsForSurface(provider, ['device'], false)).toEqual(['device']);
}
expect(subscriptionLoginMethodsForSurface('antigravity', ['browser'], true)).toEqual(['browser']);
});

it('keeps legacy backends on default login without inventing unadvertised capabilities', () => {
expect(subscriptionLoginMethodsForSurface('codex', undefined, true)).toEqual([]);
expect(subscriptionLoginMethodsForSurface('codex', [], false)).toEqual([]);
expect(subscriptionLoginMethodsForSurface('codex', ['browser'], true)).toEqual(['browser']);
});
});

describe('preferredSubscriptionLoginMethod', () => {
it('uses Codex device authorization when the callback browser is not local', () => {
expect(preferredSubscriptionLoginMethod('codex', false)).toBe('device');
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,15 @@
import type { SubscriptionProvider } from '../types';
import type { SubscriptionLoginMethod } from '@/infrastructure/api/service-api/AIApi';

export function subscriptionLoginMethodsForSurface(
provider: SubscriptionProvider,
methods: readonly SubscriptionLoginMethod[] | undefined,
localBrowserCallbackReachable: boolean,
): SubscriptionLoginMethod[] {
return (methods ?? []).filter((method) => (
provider !== 'codex' || method !== 'browser' || localBrowserCallbackReachable
));
}

/**
* Keeps the settings UI to one sign-in action while selecting the Codex flow
Expand Down
Loading