diff --git a/.gitleaks.toml b/.gitleaks.toml new file mode 100644 index 000000000..2bdbaf727 --- /dev/null +++ b/.gitleaks.toml @@ -0,0 +1,18 @@ +# Roomote gitleaks configuration. +# Extends default rules; only add allowlist entries with a concrete reason. +# +# Scope allowlists to the specific value, never to a whole file: a path +# allowlist exempts every future line in that file from every rule. + +title = "Roomote" + +[allowlist] +description = "False positives and intentional public identifiers" + +# Public Grok CLI OAuth client id reused by peer tools for SuperGrok +# device-code login. This is not a secret; it is the same client_id published +# by ecosystem agents (Hermes, OpenClaw, CC Switch) for auth.x.ai. It is +# UUID-shaped, which is the only reason the default rules flag it. +regexes = [ + '''b1a00492-073a-47ea-816f-4c329264a828''', +] diff --git a/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts b/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts index 0f9dacd32..432f63ea4 100644 --- a/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts +++ b/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts @@ -8,12 +8,14 @@ const { mockResolveModelProviderEnvValue, mockGetFreshChatGptAccessToken, mockGetGitHubCopilotAccessToken, + mockGetFreshXaiAccessToken, mockRecordLlmUsage, } = vi.hoisted(() => ({ mockFindTaskRun: vi.fn(), mockResolveModelProviderEnvValue: vi.fn(), mockGetFreshChatGptAccessToken: vi.fn(), mockGetGitHubCopilotAccessToken: vi.fn(), + mockGetFreshXaiAccessToken: vi.fn(), mockRecordLlmUsage: vi.fn(), })); @@ -28,6 +30,7 @@ vi.mock('@roomote/db/server', () => ({ resolveModelProviderEnvValue: mockResolveModelProviderEnvValue, getFreshChatGptAccessToken: mockGetFreshChatGptAccessToken, getGitHubCopilotAccessToken: mockGetGitHubCopilotAccessToken, + getFreshXaiAccessToken: mockGetFreshXaiAccessToken, })); vi.mock('@roomote/sdk/server', () => ({ @@ -114,6 +117,7 @@ describe('inference gateway', () => { vi.unstubAllGlobals(); mockFindTaskRun.mockResolvedValue({ id: 42 }); mockGetGitHubCopilotAccessToken.mockResolvedValue(null); + mockGetFreshXaiAccessToken.mockResolvedValue(null); mockResolveModelProviderEnvValue.mockImplementation( async (names: string | readonly string[]) => { const nameList = typeof names === 'string' ? [names] : names; @@ -776,6 +780,74 @@ describe('inference gateway', () => { }); }); + describe('xAI Grok subscription (xai-oauth strategy)', () => { + async function postXai( + app: Hono<{ Variables: Variables }>, + path = '/api/inference/xai/v1/chat/completions', + ) { + return app.request(path, { + method: 'POST', + headers: { + 'content-type': 'application/json', + authorization: 'Bearer run-token-value', + }, + body: JSON.stringify({ model: 'grok-4.5', messages: [] }), + }); + } + + it('prefers a connected OAuth access token over XAI_API_KEY', async () => { + const fetchMock = stubUpstreamFetch(); + mockGetFreshXaiAccessToken.mockResolvedValue({ + access: 'xai-oauth-access', + refresh: 'xai-oauth-refresh', + expires: Date.now() + 3_600_000, + }); + + const response = await postXai(createApp(createRunToken())); + + expect(response.status).toBe(200); + const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]; + expect(url).toBe('https://api.x.ai/v1/chat/completions'); + const headers = new Headers(init.headers); + expect(headers.get('authorization')).toBe('Bearer xai-oauth-access'); + }); + + it('falls back to XAI_API_KEY when no subscription is connected', async () => { + const fetchMock = stubUpstreamFetch(); + mockGetFreshXaiAccessToken.mockResolvedValue(null); + mockResolveModelProviderEnvValue.mockImplementation( + async (names: string | readonly string[]) => { + const nameList = typeof names === 'string' ? [names] : names; + if (nameList.includes('XAI_API_KEY')) { + return 'xai-api-key-value'; + } + return undefined; + }, + ); + + const response = await postXai(createApp(createRunToken())); + + expect(response.status).toBe(200); + const [, init] = fetchMock.mock.calls[0] as [string, RequestInit]; + const headers = new Headers(init.headers); + expect(headers.get('authorization')).toBe('Bearer xai-api-key-value'); + }); + + it('returns 404 when disconnected with no API key', async () => { + const fetchMock = stubUpstreamFetch(); + mockGetFreshXaiAccessToken.mockResolvedValue(null); + mockResolveModelProviderEnvValue.mockResolvedValue(undefined); + + const response = await postXai(createApp(createRunToken())); + + expect(response.status).toBe(404); + expect(await response.json()).toMatchObject({ + error: expect.stringMatching(/xAI|subscription|XAI_API_KEY/i), + }); + expect(fetchMock).not.toHaveBeenCalled(); + }); + }); + it('returns 404 when the provider key is not configured', async () => { const fetchMock = stubUpstreamFetch(); mockResolveModelProviderEnvValue.mockResolvedValue(undefined); diff --git a/apps/api/src/handlers/inference/registry.ts b/apps/api/src/handlers/inference/registry.ts index fdf8a325f..8cd9f3548 100644 --- a/apps/api/src/handlers/inference/registry.ts +++ b/apps/api/src/handlers/inference/registry.ts @@ -6,6 +6,7 @@ import { } from '@roomote/types'; import { getFreshChatGptAccessToken, + getFreshXaiAccessToken, getGitHubCopilotAccessToken, resolveModelProviderEnvValue, } from '@roomote/db/server'; @@ -32,6 +33,8 @@ export type GatewayUpstreamResolution = * - `api-key`: inject the deployment's static key at the request path. * - `chatgpt-oauth`: mint a fresh subscription access token, add the * account-id header, and collapse the request onto the Codex backend. + * - `xai-oauth`: prefer a connected Grok subscription access token, then + * fall back to `XAI_API_KEY`. */ export async function resolveGatewayUpstream( provider: InferenceGatewayProvider, @@ -46,6 +49,10 @@ export async function resolveGatewayUpstream( return resolveGitHubCopilotUpstream(provider, upstreamPath, search); } + if (provider.authStrategy === 'xai-oauth') { + return resolveXaiUpstream(provider, upstreamPath, search); + } + const [apiKey, upstreamBaseUrl] = await Promise.all([ resolveModelProviderEnvValue(provider.envVarNames), resolveProviderUpstreamBaseUrl(provider), @@ -76,6 +83,49 @@ export async function resolveGatewayUpstream( }; } +/** + * xAI supports both SuperGrok OAuth and a BYOK API key. Prefer a connected + * subscription (fresh access token) so subscription users never need a key; + * fall back to the deployment key when only that is configured. + */ +async function resolveXaiUpstream( + provider: InferenceGatewayProvider, + upstreamPath: string, + search: string, +): Promise { + const [oauthToken, apiKey, upstreamBaseUrl] = await Promise.all([ + getFreshXaiAccessToken(), + resolveModelProviderEnvValue(provider.envVarNames), + resolveProviderUpstreamBaseUrl(provider), + ]); + + const bearer = oauthToken?.access ?? apiKey; + + if (!bearer) { + return { + ok: false, + status: 404, + error: + 'No connected xAI Grok subscription or XAI_API_KEY is available for this deployment', + }; + } + + return { + ok: true, + resolved: { + upstreamUrl: `${upstreamBaseUrl}${upstreamPath}${search}`, + headers: provider.authHeader + ? { + [provider.authHeader.name]: formatProviderAuthHeaderValue( + provider, + bearer, + ), + } + : {}, + }, + }; +} + async function resolveGitHubCopilotUpstream( provider: InferenceGatewayProvider, upstreamPath: string, diff --git a/apps/docs/environment-variables.mdx b/apps/docs/environment-variables.mdx index 358b04856..5e6212e93 100644 --- a/apps/docs/environment-variables.mdx +++ b/apps/docs/environment-variables.mdx @@ -180,7 +180,8 @@ as per-task auth tokens or workspace paths. | `AI_GATEWAY_API_KEY` | Provider key | Vercel AI Gateway API key. | | `OPENAI_API_KEY` | Provider key | OpenAI API key. | | `ANTHROPIC_API_KEY` | Provider key | Anthropic API key. | -| `XAI_API_KEY` | Provider key | xAI / Grok API key. | +| `XAI_API_KEY` | Provider key | xAI / Grok API key. Optional when a SuperGrok / eligible X Premium+ subscription is connected from **Settings > Models**. | +| `XAI_OAUTH_CLIENT_ID` | Optional | Override for the public Grok CLI OAuth client id used by the SuperGrok device-code connect flow. Defaults to the public CLI client id. | | `MOONSHOT_API_KEY` | Provider key | Moonshot AI / Kimi Open Platform API key. | | `KIMI_API_KEY` | Provider key | Kimi for Coding API key for `kimi-for-coding/...` models. | | `MINIMAX_API_KEY` | Provider key | MiniMax API key. | diff --git a/apps/docs/models.mdx b/apps/docs/models.mdx index 31373bace..d93dc6060 100644 --- a/apps/docs/models.mdx +++ b/apps/docs/models.mdx @@ -17,9 +17,9 @@ Configure models from **Settings > Models**. An inference provider is the service that hosts or routes model calls. Roomote supports providers such as OpenRouter, Vercel AI Gateway, Baseten, Together AI, OpenAI, Anthropic, Moonshot AI, Kimi for Coding, MiniMax, Z.AI, -Z.AI Coding Plan, OpenCode, Amazon Bedrock, Google Gemini, xAI, GitHub Copilot, -ChatGPT subscriptions, and OpenAI-compatible endpoints such as LiteLLM, Ollama, -and vLLM. +Z.AI Coding Plan, OpenCode, Amazon Bedrock, Google Gemini, xAI (API key or +Grok subscription), GitHub Copilot, ChatGPT subscriptions, and +OpenAI-compatible endpoints such as LiteLLM, Ollama, and vLLM. You can connect more than one inference provider in the same deployment. That lets you mix and match models by provider instead of betting the whole @@ -138,12 +138,16 @@ list. - **Models** controls the provider/model pairs that are available, the default model, and specialized model roles. -For subscription-based providers (ChatGPT, GitHub Copilot, and Kimi for -Coding), the provider row also shows current plan usage when the provider -reports it, such as remaining Copilot premium requests or the percent used of -the ChatGPT 5-hour and weekly limits. These numbers come from unofficial -provider endpoints, so the usage line is hidden whenever the provider does not -return usable data. +For subscription-based providers (ChatGPT, GitHub Copilot, Kimi for Coding, and +xAI Grok / SuperGrok), the provider row also shows current plan usage when the +provider reports it, such as remaining Copilot premium requests, the percent +used of the ChatGPT 5-hour and weekly limits, or Grok included-usage windows. +These numbers come from unofficial provider endpoints, so the usage line is +hidden whenever the provider does not return usable data. + +xAI can be connected with an API key, a SuperGrok or eligible X Premium+ +subscription, or both. When both are configured, the subscription is preferred +at runtime for `xai/` models. Admins can enable or disable models from the task model list. The default model must stay enabled, because Roomote uses it when a task does not request a diff --git a/apps/web/src/app/(onboarding)/setup/StepInferenceProvider.client.test.tsx b/apps/web/src/app/(onboarding)/setup/StepInferenceProvider.client.test.tsx index d04e6306a..937606842 100644 --- a/apps/web/src/app/(onboarding)/setup/StepInferenceProvider.client.test.tsx +++ b/apps/web/src/app/(onboarding)/setup/StepInferenceProvider.client.test.tsx @@ -35,6 +35,11 @@ vi.mock('@/trpc/client', () => ({ queryOptions: () => ({ queryKey: ['chatgptSubscription.status'] }), }, }, + xaiSubscription: { + status: { + queryOptions: () => ({ queryKey: ['xaiSubscription.status'] }), + }, + }, taskModels: { discoverProviderModels: { mutationOptions: (options: Record) => options, @@ -78,6 +83,23 @@ vi.mock('@/components/settings/GitHubCopilotConnectDialog', () => ({ GitHubCopilotConnectDialog: () => null, })); +vi.mock('@/components/settings/XaiConnectDialog', () => ({ + XaiConnectDialog: ({ + open, + onOpenChange, + }: { + open: boolean; + onOpenChange: (open: boolean) => void; + }) => + open ? ( +
+ +
+ ) : null, +})); + vi.mock('@/components/system', () => ({ ArrowRight: (props: SVGProps) => , Button: ({ @@ -171,6 +193,31 @@ function githubCopilotProviderStatus( }; } +function xaiProviderStatus(options: { + keyConnected?: boolean; + subscriptionConnected?: boolean; +}): SetupModelStatus['providers'][number] { + const keyConnected = options.keyConnected ?? false; + const subscriptionConnected = options.subscriptionConnected ?? false; + return { + id: 'xai', + label: 'xAI', + envVarName: 'XAI_API_KEY', + defaultRoomoteModel: 'xai/grok-4.5', + authKind: 'api-key', + suggestedTaskModels: [], + additionalEnvFields: [], + additionalEnvValues: {}, + runtimeApiKeySatisfied: false, + savedApiKeySatisfied: keyConnected || subscriptionConnected, + credentialHelp: { + text: 'Use an API key, or connect a SuperGrok / eligible X Premium+ subscription.', + href: 'https://console.x.ai', + linkLabel: 'xAI console', + }, + }; +} + function openrouterProviderStatus(): SetupModelStatus['providers'][number] { return { id: 'openrouter', @@ -234,16 +281,36 @@ function setupMutationMock() { function setupQueryMocks(options: { chatgptConnected: boolean; chatgptEmail?: string; + xaiConnected?: boolean; + xaiEmail?: string; }) { - mockUseQuery.mockReturnValue({ - data: options.chatgptConnected - ? { - connected: true, - status: 'connected', - email: options.chatgptEmail ?? 'owner@example.com', - } - : { connected: false, status: 'disconnected' }, - } as unknown as ReturnType); + mockUseQuery.mockImplementation((queryOptions) => { + const key = JSON.stringify( + (queryOptions as { queryKey?: unknown }).queryKey ?? [], + ); + if (key.includes('xaiSubscription')) { + return { + data: options.xaiConnected + ? { + connected: true, + status: 'connected', + email: options.xaiEmail, + } + : { connected: false, status: 'disconnected' }, + isPending: false, + } as never; + } + return { + data: options.chatgptConnected + ? { + connected: true, + status: 'connected', + email: options.chatgptEmail, + } + : { connected: false, status: 'disconnected' }, + isPending: false, + } as never; + }); } function selectProvider(provider: SetupModelProviderId) { @@ -581,6 +648,74 @@ describe('StepInferenceProvider ChatGPT subscription', () => { ).not.toBeInTheDocument(); }); + it('offers SuperGrok subscription connect alongside the xAI API key field', () => { + render( + , + ); + + selectProvider('xai'); + + expect(screen.getByPlaceholderText(/API key for xAI/i)).toBeInTheDocument(); + expect( + screen.getByRole('button', { name: /Connect Grok subscription/i }), + ).toBeInTheDocument(); + expect(screen.getByRole('button', { name: /continue/i })).toBeDisabled(); + }); + + it('enables Continue when only an xAI Grok subscription is connected', () => { + setupQueryMocks({ chatgptConnected: false, xaiConnected: true }); + render( + , + ); + + selectProvider('xai'); + + expect(screen.getByText(/Connected to a SuperGrok/i)).toBeInTheDocument(); + expect(screen.getByRole('button', { name: /continue/i })).toBeEnabled(); + expect( + screen.queryByRole('button', { name: /Connect Grok subscription/i }), + ).not.toBeInTheDocument(); + }); + + it('opens the xAI connect dialog from the subscription button', () => { + render( + , + ); + + selectProvider('xai'); + + expect(screen.queryByTestId('xai-connect-dialog')).not.toBeInTheDocument(); + fireEvent.click( + screen.getByRole('button', { name: /Connect Grok subscription/i }), + ); + expect(screen.getByTestId('xai-connect-dialog')).toBeInTheDocument(); + }); + it('opens the ChatGPT connect dialog from the connect button', () => { render( { @@ -149,6 +163,7 @@ export function StepInferenceProvider({ setEditingSavedValue(false); setIsChatGptDialogOpen(false); setIsGitHubCopilotDialogOpen(false); + setIsXaiDialogOpen(false); }, [selectedProvider]); useEffect(() => { @@ -179,6 +194,9 @@ export function StepInferenceProvider({ const isEndpointProvider = selectedProviderStatus?.authKind === 'endpoint'; const chatgptConnected = Boolean(modelSetup.chatgptConnected); const githubCopilotConnected = Boolean(modelSetup.githubCopilotConnected); + const xaiSubscriptionConnected = Boolean( + modelSetup.xaiSubscriptionConnected || xaiStatus?.connected, + ); const hasRuntimeProviderKey = selectedProviderStatus?.runtimeApiKeySatisfied === true; const hasSavedProviderKey = @@ -203,7 +221,14 @@ export function StepInferenceProvider({ !editingSavedValue; const shouldShowConfiguredMask = hasRuntimeProviderKey || shouldShowSavedValueMask; - const canContinueWithoutApiKey = hasRuntimeProviderKey || hasSavedProviderKey; + // SuperGrok covers the `xai` surface without an API key, same pattern as + // ChatGPT covering openai/. Continue when either path is satisfied. + const canContinueWithoutApiKey = + hasRuntimeProviderKey || + hasSavedProviderKey || + (isXaiSurface && xaiSubscriptionConnected) || + (isChatGptProvider && chatgptConnected) || + (isGitHubCopilotProvider && githubCopilotConnected); const hasMissingRequiredFields = !canContinueWithoutApiKey && additionalEnvFields.some( @@ -388,9 +413,31 @@ export function StepInferenceProvider({ ) : null} - {(hasRuntimeProviderKey || hasSavedProviderKey) && } + {isXaiSurface && !hasRuntimeProviderKey && !xaiSubscriptionConnected ? ( + + ) : null} + + {(hasRuntimeProviderKey || + hasSavedProviderKey || + (isXaiSurface && xaiSubscriptionConnected)) && } + {isXaiSurface && !hasRuntimeProviderKey && xaiSubscriptionConnected ? ( + + {xaiStatus?.email + ? `Connected as ${xaiStatus.email}` + : 'Connected to a SuperGrok / X Premium+ account.'} + + ) : null} + {requiresConnectionName && !hasRuntimeProviderKey ? (
@@ -534,6 +581,15 @@ export function StepInferenceProvider({ }); }} /> + { + await queryClient.invalidateQueries({ + queryKey: trpc.setupNew.status.queryKey(), + }); + }} + /> ) : ( + ) : null} + +
+ + ); +} + function GitHubCopilotSubscriptionRow({ errored, errorMessage, @@ -749,6 +817,7 @@ export function InferenceProviderSection({ const [isChatGptDialogOpen, setIsChatGptDialogOpen] = useState(false); const [isGitHubCopilotDialogOpen, setIsGitHubCopilotDialogOpen] = useState(false); + const [isXaiDialogOpen, setIsXaiDialogOpen] = useState(false); const chatgptStatusQuery = useQuery( trpc.chatgptSubscription.status.queryOptions(), @@ -758,6 +827,8 @@ export function InferenceProviderSection({ trpc.githubCopilotSubscription.status.queryOptions(), ); const githubCopilotStatus = githubCopilotStatusQuery.data ?? null; + const xaiStatusQuery = useQuery(trpc.xaiSubscription.status.queryOptions()); + const xaiStatus = xaiStatusQuery.data ?? null; // Usage endpoints are unofficial upstream surfaces; a provider missing from // the result just means no usage line is rendered for it. @@ -893,6 +964,25 @@ export function InferenceProviderSection({ onError: (error) => toast.error(error.message), }), ); + const disconnectXai = useMutation( + trpc.xaiSubscription.disconnect.mutationOptions({ + onSuccess: async () => { + toast.success('Disconnected xAI Grok subscription.'); + await Promise.all([ + queryClient.invalidateQueries({ + queryKey: trpc.taskModels.providerSetup.queryKey(), + }), + queryClient.invalidateQueries({ + queryKey: trpc.taskModels.launchOptions.queryKey(), + }), + queryClient.invalidateQueries({ + queryKey: trpc.xaiSubscription.status.queryKey(), + }), + ]); + }, + onError: (error) => toast.error(error.message), + }), + ); const deleteProvider = useMutation( trpc.taskModels.deleteProvider.mutationOptions({ @@ -946,6 +1036,9 @@ export function InferenceProviderSection({ const handleDisconnectGitHubCopilot = async () => { await disconnectGitHubCopilot.mutateAsync(); }; + const handleDisconnectXai = async () => { + await disconnectXai.mutateAsync(); + }; const handleDeleteProvider = async () => { if (!deleteProviderId) { @@ -974,23 +1067,41 @@ export function InferenceProviderSection({ (githubCopilotStatusQuery.isPending && Boolean(providerSetup?.githubCopilotConnected)); const githubCopilotErrored = githubCopilotStatus?.status === 'error'; - - // The ChatGPT provider is OAuth-backed: it renders as a connected row with - // Disconnect/Reconnect when a subscription record exists, and otherwise - // joins the add-provider dropdown with a Connect button instead of an API - // key field. Exclude it from the API-key connected rows either way. - const apiKeyConnectedProviders = connectedProviders.filter( - (provider) => - provider.id !== CHATGPT_SUBSCRIPTION_PROVIDER_ID && - provider.id !== 'github-copilot', - ); - const addableProviders = availableProviders.filter((provider) => - provider.id === CHATGPT_SUBSCRIPTION_PROVIDER_ID - ? !chatgptHasRecord - : provider.id === 'github-copilot' - ? !githubCopilotHasRecord - : provider.authKind === 'api-key' || provider.authKind === 'endpoint', - ); + const xaiHasRecord = + xaiStatus?.status === 'connected' || + xaiStatus?.status === 'error' || + (xaiStatusQuery.isPending && + Boolean(providerSetup?.xaiSubscriptionConnected)); + const xaiErrored = xaiStatus?.status === 'error'; + const xaiHasApiKey = Boolean(providerSetup?.xaiApiKeyConnected); + + // Subscription oauth rows are rendered separately; API-key xAI stays in the + // key list and can coexist with the SuperGrok subscription provider. + const apiKeyConnectedProviders = connectedProviders.filter((provider) => { + if ( + provider.id === CHATGPT_SUBSCRIPTION_PROVIDER_ID || + provider.id === 'github-copilot' || + provider.id === XAI_SUBSCRIPTION_PROVIDER_ID + ) { + return false; + } + return true; + }); + const addableProviders = availableProviders.filter((provider) => { + if (provider.id === CHATGPT_SUBSCRIPTION_PROVIDER_ID) { + return !chatgptHasRecord; + } + if (provider.id === 'github-copilot') { + return !githubCopilotHasRecord; + } + if (provider.id === XAI_SUBSCRIPTION_PROVIDER_ID) { + return !xaiHasRecord; + } + if (provider.id === 'xai') { + return !xaiHasApiKey; + } + return provider.authKind === 'api-key' || provider.authKind === 'endpoint'; + }); const sortedApiKeyConnectedProviders = useMemo( () => [...apiKeyConnectedProviders].sort((left, right) => @@ -1070,12 +1181,18 @@ export function InferenceProviderSection({ const hasConnectedApiKeys = sortedApiKeyConnectedProviders.length > 0; const hasConnectedProviders = - hasConnectedApiKeys || chatgptHasRecord || githubCopilotHasRecord; + hasConnectedApiKeys || + chatgptHasRecord || + githubCopilotHasRecord || + xaiHasRecord; const canAddProvider = sortedAddableProviders.length > 0; + // Count key rows and subscription rows independently so dual-path xAI + // (API key + SuperGrok) can delete the key while the subscription remains. const connectedProviderCount = sortedApiKeyConnectedProviders.length + (chatgptHasRecord ? 1 : 0) + - (githubCopilotHasRecord ? 1 : 0); + (githubCopilotHasRecord ? 1 : 0) + + (xaiHasRecord ? 1 : 0); return (
@@ -1094,6 +1211,14 @@ export function InferenceProviderSection({

) : null} + {xaiHasRecord && xaiHasApiKey ? ( +

+ Both an xAI API key and a Grok subscription are configured. The + subscription is preferred at runtime when an xai/ model + is selected. +

+ ) : null} + + {providerDialog ? ( ) : null} + {xaiHasRecord ? ( + setIsXaiDialogOpen(true)} + onDisconnect={handleDisconnectXai} + isDisconnecting={disconnectXai.isPending} + /> + ) : null} {sortedApiKeyConnectedProviders.map((provider) => ( + provider.id === XAI_SUBSCRIPTION_PROVIDER_ID && + provider.savedApiKeySatisfied, + ); + const xaiConnected = sortedConnectedProviders.some( + (provider) => + provider.id === 'xai' && + (provider.savedApiKeySatisfied || provider.runtimeApiKeySatisfied), + ); const activeNewModelProvider = useMemo( () => sortedConnectedProviders.find( @@ -1104,29 +1115,26 @@ export function ModelSettingsSection({ return options; }, [settingsData]); + const groupOptions = useMemo( + () => ({ + chatgptConnected, + openaiConnected, + xaiSubscriptionConnected, + xaiConnected, + }), + [chatgptConnected, openaiConnected, xaiSubscriptionConnected, xaiConnected], + ); const codingModelGroups = useMemo( - () => - groupModelsByDisplayProvider(codingModelOptions, { - chatgptConnected, - openaiConnected, - }), - [codingModelOptions, chatgptConnected, openaiConnected], + () => groupModelsByDisplayProvider(codingModelOptions, groupOptions), + [codingModelOptions, groupOptions], ); const helperModelGroups = useMemo( - () => - groupModelsByDisplayProvider(helperModelOptions, { - chatgptConnected, - openaiConnected, - }), - [helperModelOptions, chatgptConnected, openaiConnected], + () => groupModelsByDisplayProvider(helperModelOptions, groupOptions), + [helperModelOptions, groupOptions], ); const modelGroups = useMemo( - () => - groupModelsByDisplayProvider(models, { - chatgptConnected, - openaiConnected, - }), - [models, chatgptConnected, openaiConnected], + () => groupModelsByDisplayProvider(models, groupOptions), + [models, groupOptions], ); const roleOptionGroups: Record< TaskModelRole, diff --git a/apps/web/src/components/settings/SubscriptionUsageLine.tsx b/apps/web/src/components/settings/SubscriptionUsageLine.tsx index 0635412c9..b7c758c71 100644 --- a/apps/web/src/components/settings/SubscriptionUsageLine.tsx +++ b/apps/web/src/components/settings/SubscriptionUsageLine.tsx @@ -44,6 +44,9 @@ function formatUsageWindow(window: SubscriptionUsageWindow): string | null { value = `${window.used.toLocaleString()} of ${window.limit.toLocaleString()} used`; } else if (window.usedPercent !== undefined) { value = `${Math.round(window.usedPercent)}% used`; + } else if (window.remaining !== undefined) { + // Prepaid credit balance with no hard cap (e.g. Grok prepaidBalance). + value = `${window.remaining.toLocaleString()} left`; } else { return null; } diff --git a/apps/web/src/components/settings/XaiConnectDialog.test.tsx b/apps/web/src/components/settings/XaiConnectDialog.test.tsx new file mode 100644 index 000000000..2470302fa --- /dev/null +++ b/apps/web/src/components/settings/XaiConnectDialog.test.tsx @@ -0,0 +1,137 @@ +import { act, render, screen } from '@testing-library/react'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +const startMutate = vi.hoisted(() => vi.fn()); +const pollMutateAsync = vi.hoisted(() => vi.fn()); +const invalidateQueries = vi.hoisted(() => vi.fn()); + +vi.mock('@tanstack/react-query', () => ({ + useMutation: (options: { + mutationKey?: unknown; + onSuccess?: (result: unknown) => void; + onError?: (error: Error) => void; + }) => { + // Distinguish start vs poll by whether onSuccess is provided (start has it). + if (options.onSuccess) { + return { + mutate: () => { + startMutate(); + options.onSuccess?.({ + deviceCode: 'device-1', + userCode: 'ABCD-EFGH', + verificationUrl: 'https://accounts.x.ai/device', + intervalMs: 1_000, + expiresInMs: 5_000, + }); + }, + mutateAsync: vi.fn(), + isPending: false, + isError: false, + reset: vi.fn(), + }; + } + return { + mutate: vi.fn(), + mutateAsync: pollMutateAsync, + isPending: false, + isError: false, + reset: vi.fn(), + }; + }, + useQueryClient: () => ({ invalidateQueries }), +})); + +vi.mock('sonner', () => ({ + toast: { success: vi.fn(), error: vi.fn() }, +})); + +vi.mock('@/trpc/client', () => ({ + useTRPC: () => ({ + taskModels: { + providerSetup: { queryKey: () => ['taskModels', 'providerSetup'] }, + launchOptions: { queryKey: () => ['taskModels', 'launchOptions'] }, + }, + xaiSubscription: { + status: { queryKey: () => ['xaiSubscription', 'status'] }, + startDeviceAuth: { + mutationOptions: (options?: { + onSuccess?: (result: unknown) => void; + onError?: (error: Error) => void; + }) => options ?? {}, + }, + pollDeviceAuth: { + mutationOptions: () => ({}), + }, + }, + subscriptionUsage: { + list: { queryKey: () => ['subscriptionUsage', 'list'] }, + }, + }), +})); + +import { XaiConnectDialog } from './XaiConnectDialog'; + +describe('XaiConnectDialog', () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.useFakeTimers(); + pollMutateAsync.mockResolvedValue({ status: 'pending' }); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it('stops polling when the device code expires and offers Restart', async () => { + const start = new Date('2026-01-01T00:00:00.000Z'); + vi.setSystemTime(start); + + render(); + + await act(async () => { + await Promise.resolve(); + }); + + expect(screen.getByDisplayValue('ABCD-EFGH')).toBeInTheDocument(); + expect(pollMutateAsync).toHaveBeenCalled(); + + // Jump past expiresInMs and flush the sleep scheduled after the first poll. + await act(async () => { + vi.setSystemTime(new Date(start.getTime() + 6_000)); + await vi.advanceTimersByTimeAsync(6_000); + }); + + expect( + screen.getByText(/device authorization code expired/i), + ).toBeInTheDocument(); + expect( + screen.getByRole('button', { name: /Restart/i }), + ).toBeInTheDocument(); + }); + + it('replaces the poll interval on slow_down instead of accumulating', async () => { + pollMutateAsync + .mockResolvedValueOnce({ status: 'pending', intervalMs: 3_000 }) + .mockResolvedValue({ status: 'pending' }); + + render(); + + await act(async () => { + await Promise.resolve(); + }); + + expect(pollMutateAsync).toHaveBeenCalledTimes(1); + + // After first poll, sleep should be 3000 (replaced), not 1000+3000=4000. + // Advance 2999ms: second poll must not have fired yet. + await act(async () => { + await vi.advanceTimersByTimeAsync(2_999); + }); + expect(pollMutateAsync).toHaveBeenCalledTimes(1); + + await act(async () => { + await vi.advanceTimersByTimeAsync(2); + }); + expect(pollMutateAsync).toHaveBeenCalledTimes(2); + }); +}); diff --git a/apps/web/src/components/settings/XaiConnectDialog.tsx b/apps/web/src/components/settings/XaiConnectDialog.tsx new file mode 100644 index 000000000..310af3be0 --- /dev/null +++ b/apps/web/src/components/settings/XaiConnectDialog.tsx @@ -0,0 +1,245 @@ +'use client'; + +import { useEffect, useRef, useState } from 'react'; +import { useMutation, useQueryClient } from '@tanstack/react-query'; +import { toast } from 'sonner'; + +import { useTRPC } from '@/trpc/client'; +import { + Button, + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, + ExternalLink, + Input, + Spinner, +} from '@/components/system'; + +type DeviceAuth = { + deviceCode: string; + userCode: string; + verificationUrl: string; + intervalMs: number; + expiresInMs: number; +}; + +export function XaiConnectDialog({ + open, + onOpenChange, + onConnected, +}: { + open: boolean; + onOpenChange: (open: boolean) => void; + onConnected?: () => void | Promise; +}) { + const trpc = useTRPC(); + const queryClient = useQueryClient(); + const [deviceAuth, setDeviceAuth] = useState(null); + const [error, setError] = useState(null); + /** + * Monotonic generation bumped whenever a poll loop must die (close, restart, + * unmount). A shared boolean is not enough: cancel + reopen can flip the + * flag back to true and revive a stale loop that still holds the previous + * device code. + */ + const pollGenerationRef = useRef(0); + + const startMutation = useMutation( + trpc.xaiSubscription.startDeviceAuth.mutationOptions({ + onSuccess: (result) => { + setDeviceAuth(result); + setError(null); + }, + onError: (mutationError) => setError(mutationError.message), + }), + ); + const pollMutation = useMutation( + trpc.xaiSubscription.pollDeviceAuth.mutationOptions(), + ); + + function invalidatePollLoops() { + pollGenerationRef.current += 1; + } + + function close(next: boolean) { + if (!next) { + invalidatePollLoops(); + } + onOpenChange(next); + } + + useEffect(() => { + if (!open) { + invalidatePollLoops(); + setDeviceAuth(null); + setError(null); + return; + } + if (!deviceAuth && !startMutation.isPending && !startMutation.isError) { + startMutation.mutate(); + } + }, [deviceAuth, open, startMutation]); + + useEffect(() => { + if (!open || !deviceAuth) { + return; + } + + const generation = ++pollGenerationRef.current; + const activeDeviceCode = deviceAuth.deviceCode; + const expiresAt = Date.now() + deviceAuth.expiresInMs; + + const poll = async () => { + let intervalMs = deviceAuth.intervalMs; + while (pollGenerationRef.current === generation) { + if (Date.now() >= expiresAt) { + setError( + 'xAI device authorization code expired. Restart the connection.', + ); + return; + } + + const result = await pollMutation.mutateAsync({ + deviceCode: activeDeviceCode, + }); + // Another open/close/restart may have started while we awaited. + if (pollGenerationRef.current !== generation) { + return; + } + if (result.status === 'success') { + toast.success('xAI Grok subscription connected.'); + await Promise.all([ + queryClient.invalidateQueries({ + queryKey: trpc.taskModels.providerSetup.queryKey(), + }), + queryClient.invalidateQueries({ + queryKey: trpc.taskModels.launchOptions.queryKey(), + }), + queryClient.invalidateQueries({ + queryKey: trpc.xaiSubscription.status.queryKey(), + }), + queryClient.invalidateQueries({ + queryKey: trpc.subscriptionUsage.list.queryKey(), + }), + ]); + if (pollGenerationRef.current !== generation) { + return; + } + await onConnected?.(); + if (pollGenerationRef.current !== generation) { + return; + } + close(false); + return; + } + if (result.status === 'failed') { + setError(result.error); + return; + } + // slow_down returns the new absolute poll interval, not a delta. + if (result.intervalMs) { + intervalMs = result.intervalMs; + } + + const remainingMs = expiresAt - Date.now(); + if (remainingMs <= 0) { + setError( + 'xAI device authorization code expired. Restart the connection.', + ); + return; + } + await new Promise((resolve) => + setTimeout(resolve, Math.min(intervalMs, remainingMs)), + ); + } + }; + + void poll().catch((pollError: unknown) => { + if (pollGenerationRef.current !== generation) { + return; + } + setError( + pollError instanceof Error + ? pollError.message + : 'xAI authorization polling failed.', + ); + }); + return () => { + // Invalidate only this effect's generation when deviceAuth/open change. + if (pollGenerationRef.current === generation) { + pollGenerationRef.current += 1; + } + }; + // The polling lifecycle is intentionally keyed only to the active device + // flow; mutation/query objects are recreated by hooks between renders. + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [deviceAuth, open]); + + return ( + + + + Connect xAI Grok subscription + + Sign in with a SuperGrok or eligible X Premium+ account to run Grok + models on your subscription instead of an API key. + + + + {deviceAuth ? ( +
+

+ Open xAI, then enter this one-time code: +

+ + +

+ Waiting for authorization… +

+
+ ) : startMutation.isPending ? ( +
+ +
+ ) : null} + + {error ?

{error}

: null} + + + + {error ? ( + + ) : null} + + + + ); +} diff --git a/apps/web/src/components/tasks/ModelSelect.tsx b/apps/web/src/components/tasks/ModelSelect.tsx index f2e4fd169..65b4fa620 100644 --- a/apps/web/src/components/tasks/ModelSelect.tsx +++ b/apps/web/src/components/tasks/ModelSelect.tsx @@ -46,8 +46,16 @@ export function ModelSelect({ return groupModelsByDisplayProvider(sortedModels, { chatgptConnected: data?.chatgptConnected, openaiConnected: data?.openaiConnected, + xaiSubscriptionConnected: data?.xaiSubscriptionConnected, + xaiConnected: data?.xaiConnected, }); - }, [data?.chatgptConnected, data?.openaiConnected, data?.models]); + }, [ + data?.chatgptConnected, + data?.openaiConnected, + data?.xaiSubscriptionConnected, + data?.xaiConnected, + data?.models, + ]); const showProviderHeaders = modelGroups.length > 1; return ( diff --git a/apps/web/src/trpc/commands/setup-new/index.test.ts b/apps/web/src/trpc/commands/setup-new/index.test.ts index 3dd29bd4a..d83d98b0c 100644 --- a/apps/web/src/trpc/commands/setup-new/index.test.ts +++ b/apps/web/src/trpc/commands/setup-new/index.test.ts @@ -127,6 +127,9 @@ vi.mock('@roomote/db/server', () => ({ sql: vi.fn(), users: {}, workItems: {}, + isChatGptSubscriptionConnected: vi.fn(async () => false), + isGitHubCopilotSubscriptionConnected: vi.fn(async () => false), + isXaiSubscriptionConnected: vi.fn(async () => false), })); vi.mock('@/lib/server', () => ({ diff --git a/apps/web/src/trpc/commands/setup-new/index.ts b/apps/web/src/trpc/commands/setup-new/index.ts index 6c5a3f8b2..28929c7cc 100644 --- a/apps/web/src/trpc/commands/setup-new/index.ts +++ b/apps/web/src/trpc/commands/setup-new/index.ts @@ -36,6 +36,7 @@ import { resolveDiscordRuntimeCredentials, isChatGptSubscriptionConnected, isGitHubCopilotSubscriptionConnected, + isXaiSubscriptionConnected, type DatabaseOrTransaction, } from '@roomote/db/server'; import { @@ -54,6 +55,7 @@ import { buildSetupModelStatus, buildSetupSourceControlStatus, CHATGPT_SUBSCRIPTION_PROVIDER_ID, + XAI_SUBSCRIPTION_PROVIDER_ID, OPENAI_COMPATIBLE_PROVIDER_ID, collectSetupModelProviderCredentialValues, createEmptyDeploymentModelConfig, @@ -1291,6 +1293,7 @@ export async function getSetupNewStatusCommand(auth: UserAuthSuccess) { nonSecretSourceControlEnvValues, chatgptConnected, githubCopilotConnected, + xaiSubscriptionConnected, ] = await Promise.all([ getSetupBaseStatus(auth), getSetupSlackAccessStatus({ userId }), @@ -1316,6 +1319,7 @@ export async function getSetupNewStatusCommand(auth: UserAuthSuccess) { ]), isChatGptSubscriptionConnected(), isGitHubCopilotSubscriptionConnected(), + isXaiSubscriptionConnected(), ]); let setupNewState = normalizeSetupNewState(baseStatus.setupNewState); @@ -1403,6 +1407,7 @@ export async function getSetupNewStatusCommand(auth: UserAuthSuccess) { selectedProvider: setupNewState.modelProvider, chatgptConnected, githubCopilotConnected, + xaiSubscriptionConnected, }); const computeSetup = buildSetupComputeStatus({ runtimeEnv: process.env, @@ -1510,16 +1515,20 @@ export async function saveSetupNewModelConfigCommand( const provider = getSetupModelProvider(providerId); const isOauthProvider = provider.authKind === 'oauth'; - const [chatgptConnected, githubCopilotConnected] = await Promise.all([ - isChatGptSubscriptionConnected(), - isGitHubCopilotSubscriptionConnected(), - ]); + const [chatgptConnected, githubCopilotConnected, xaiSubscriptionConnected] = + await Promise.all([ + isChatGptSubscriptionConnected(), + isGitHubCopilotSubscriptionConnected(), + isXaiSubscriptionConnected(), + ]); const selectedOauthConnected = provider.id === CHATGPT_SUBSCRIPTION_PROVIDER_ID ? chatgptConnected : provider.id === 'github-copilot' ? githubCopilotConnected - : false; + : provider.id === XAI_SUBSCRIPTION_PROVIDER_ID + ? xaiSubscriptionConnected + : false; // OAuth providers are connected through their // device-code flow rather than an API key. The setup wizard's Continue @@ -1623,6 +1632,7 @@ export async function saveSetupNewModelConfigCommand( persistedEnvVarNames, chatgptConnected, githubCopilotConnected, + xaiSubscriptionConnected, }), ]); const autoAdd = buildAutoAddedTaskModelSettings({ diff --git a/apps/web/src/trpc/commands/subscription-usage/index.ts b/apps/web/src/trpc/commands/subscription-usage/index.ts index fda8dccfa..77a70f7b7 100644 --- a/apps/web/src/trpc/commands/subscription-usage/index.ts +++ b/apps/web/src/trpc/commands/subscription-usage/index.ts @@ -9,7 +9,8 @@ function assertAdmin(auth: UserAuthSuccess): void { /** * Usage/quota for connected subscription providers (ChatGPT, GitHub Copilot, - * Kimi for Coding, Z.AI / Z.AI Coding Plan). Providers without a configured + * Kimi for Coding, xAI Grok subscription, Z.AI / Z.AI Coding Plan). + * Providers without a configured * credential or whose usage endpoint is unavailable are simply absent from * the result. */ diff --git a/apps/web/src/trpc/commands/task-models/auto-add-models.test.ts b/apps/web/src/trpc/commands/task-models/auto-add-models.test.ts index 318f52d3a..455637b77 100644 --- a/apps/web/src/trpc/commands/task-models/auto-add-models.test.ts +++ b/apps/web/src/trpc/commands/task-models/auto-add-models.test.ts @@ -337,4 +337,17 @@ describe('collectConnectedTaskModelProviderIds', () => { expect(connected.has('openai')).toBe(true); expect(connected.has('google')).toBe(false); }); + + it('includes xai when SuperGrok is connected without an xAI API key', () => { + const connected = collectConnectedTaskModelProviderIds({ + runtimeEnv: {}, + persistedEnvVarNames: [], + chatgptConnected: false, + xaiSubscriptionConnected: true, + }); + + expect(connected.has('xai-subscription')).toBe(true); + // SuperGrok serves xai/ model ids, same alias pattern as ChatGPT→openai. + expect(connected.has('xai')).toBe(true); + }); }); diff --git a/apps/web/src/trpc/commands/task-models/auto-add-models.ts b/apps/web/src/trpc/commands/task-models/auto-add-models.ts index 00f115102..ca7404178 100644 --- a/apps/web/src/trpc/commands/task-models/auto-add-models.ts +++ b/apps/web/src/trpc/commands/task-models/auto-add-models.ts @@ -22,19 +22,22 @@ function deriveDisplayNameFromModelId(modelId: string): string { /** * Model-id prefixes of the providers currently connected (saved or runtime * env). A connected ChatGPT subscription serves `openai/` model ids, so it - * contributes `openai` alongside its own `chatgpt` catalog id. + * contributes `openai` alongside its own `chatgpt` catalog id. SuperGrok + * similarly contributes `xai` alongside `xai-subscription`. */ export function collectConnectedTaskModelProviderIds(options: { runtimeEnv: Partial>; persistedEnvVarNames: Iterable; chatgptConnected: boolean; githubCopilotConnected?: boolean; + xaiSubscriptionConnected?: boolean; }): Set { const status = buildSetupModelStatus({ runtimeEnv: options.runtimeEnv, persistedEnvVarNames: options.persistedEnvVarNames, chatgptConnected: options.chatgptConnected, githubCopilotConnected: options.githubCopilotConnected, + xaiSubscriptionConnected: options.xaiSubscriptionConnected, }); return new Set([ @@ -45,6 +48,7 @@ export function collectConnectedTaskModelProviderIds(options: { ) .map((provider) => provider.id), ...(options.chatgptConnected ? ['openai'] : []), + ...(options.xaiSubscriptionConnected ? ['xai'] : []), ]); } diff --git a/apps/web/src/trpc/commands/task-models/index.test.ts b/apps/web/src/trpc/commands/task-models/index.test.ts index 87c580dea..4d81f4c53 100644 --- a/apps/web/src/trpc/commands/task-models/index.test.ts +++ b/apps/web/src/trpc/commands/task-models/index.test.ts @@ -15,6 +15,7 @@ const { mockUpsertDeploymentEnvironmentVariables, mockIsChatGptSubscriptionConnected, mockIsGitHubCopilotSubscriptionConnected, + mockIsXaiSubscriptionConnected, } = vi.hoisted(() => ({ mockFindDeploymentSettings: vi.fn(), mockInsertDeploymentSettings: vi.fn(), @@ -27,6 +28,7 @@ const { mockUpsertDeploymentEnvironmentVariables: vi.fn(), mockIsChatGptSubscriptionConnected: vi.fn(), mockIsGitHubCopilotSubscriptionConnected: vi.fn(), + mockIsXaiSubscriptionConnected: vi.fn(), })); vi.mock('@roomote/db/server', () => ({ @@ -56,6 +58,7 @@ vi.mock('@roomote/db/server', () => ({ isChatGptSubscriptionConnected: mockIsChatGptSubscriptionConnected, isGitHubCopilotSubscriptionConnected: mockIsGitHubCopilotSubscriptionConnected, + isXaiSubscriptionConnected: mockIsXaiSubscriptionConnected, isNull: vi.fn((column) => ({ isNull: column })), })); @@ -158,6 +161,7 @@ describe('lookupTaskModelCommand', () => { delete process.env.R_PLANNING_MODEL_REASONING_EFFORT; mockIsChatGptSubscriptionConnected.mockResolvedValue(false); mockIsGitHubCopilotSubscriptionConnected.mockResolvedValue(false); + mockIsXaiSubscriptionConnected.mockResolvedValue(false); mockGetPersistedEnvironmentVariableValues.mockResolvedValue({}); mockFindDeploymentSettings.mockImplementation(async (options) => { const columns = (options as { columns?: Record }) @@ -1199,6 +1203,7 @@ describe('task model provider commands', () => { mockGetPersistedEnvironmentVariableValues.mockResolvedValue({}); mockIsChatGptSubscriptionConnected.mockResolvedValue(false); mockIsGitHubCopilotSubscriptionConnected.mockResolvedValue(false); + mockIsXaiSubscriptionConnected.mockResolvedValue(false); mockPersistedSetupNewState({}); }); @@ -1620,6 +1625,7 @@ describe('task model provider commands', () => { // allowed because the subscription keeps the deployment usable. mockIsChatGptSubscriptionConnected.mockResolvedValue(true); mockIsGitHubCopilotSubscriptionConnected.mockResolvedValue(false); + mockIsXaiSubscriptionConnected.mockResolvedValue(false); mockGetPersistedEnvironmentVariableNames.mockResolvedValue([ 'ANTHROPIC_API_KEY', ]); @@ -1631,6 +1637,66 @@ describe('task model provider commands', () => { expect(mockTxDelete).toHaveBeenCalled(); }); + it('removes only the xAI API key when a Grok subscription remains connected', async () => { + // Dual-path xAI shares catalog id `xai`. Deleting the key must not strip + // xai/* models while SuperGrok is still connected. + mockIsChatGptSubscriptionConnected.mockResolvedValue(false); + mockIsGitHubCopilotSubscriptionConnected.mockResolvedValue(false); + mockIsXaiSubscriptionConnected.mockResolvedValue(true); + + const persistedRow = { + taskModelSettings: { + models: [ + { + id: 'xai/grok-4.5', + displayName: 'Grok 4.5', + family: 'Grok', + }, + ], + allowedModelIds: ['xai/grok-4.5'], + defaultModelId: 'xai/grok-4.5', + }, + runtimeModelConfig: { + roomoteModel: 'xai/grok-4.5', + roomoteSmallModel: 'xai/grok-4.5', + roomoteVisionModel: null, + roomoteCodeReviewModel: null, + roomoteExploreModel: null, + roomotePlanningModel: null, + roomoteModelReasoningEffort: null, + roomoteSmallModelReasoningEffort: null, + roomoteVisionModelReasoningEffort: null, + roomoteCodeReviewModelReasoningEffort: null, + roomoteExploreModelReasoningEffort: null, + roomotePlanningModelReasoningEffort: null, + }, + }; + mockFindDeploymentSettings.mockImplementation(async (options) => { + const columns = (options as { columns?: Record }) + ?.columns; + return { + ...(columns?.taskModelSettings + ? { taskModelSettings: persistedRow.taskModelSettings } + : {}), + ...(columns?.runtimeModelConfig + ? { runtimeModelConfig: persistedRow.runtimeModelConfig } + : {}), + }; + }); + mockGetPersistedEnvironmentVariableNames.mockResolvedValue(['XAI_API_KEY']); + mockPersistedSetupNewState({ modelProvider: 'xai' }); + + await deleteTaskModelProviderCommand(buildMockAuth(), { + provider: 'xai', + }); + + const { inArray } = await import('@roomote/db/server'); + expect(mockTxDelete).toHaveBeenCalled(); + expect(inArray).toHaveBeenCalledWith('env.name', ['XAI_API_KEY']); + // Models and runtime stay; only the key was removed. + expect(txOnConflictDoUpdate).not.toHaveBeenCalled(); + }); + it('rejects deleting the last connected provider', async () => { mockGetPersistedEnvironmentVariableNames.mockResolvedValue([ 'ANTHROPIC_API_KEY', diff --git a/apps/web/src/trpc/commands/task-models/index.ts b/apps/web/src/trpc/commands/task-models/index.ts index 4f277a51f..7c97e9a7e 100644 --- a/apps/web/src/trpc/commands/task-models/index.ts +++ b/apps/web/src/trpc/commands/task-models/index.ts @@ -7,6 +7,7 @@ import { inArray, isChatGptSubscriptionConnected, isGitHubCopilotSubscriptionConnected, + isXaiSubscriptionConnected, isNull, type DatabaseOrTransaction, } from '@roomote/db/server'; @@ -298,12 +299,14 @@ export async function getTaskModelSettingsCommand( persistedEnvVarNames, chatgptConnected, githubCopilotConnected, + xaiSubscriptionConnected, ] = await Promise.all([ getDeploymentTaskModelSettings(), getDeploymentRuntimeModelConfig(), getPersistedEnvironmentVariableNames(), isChatGptSubscriptionConnected(), isGitHubCopilotSubscriptionConnected(), + isXaiSubscriptionConnected(), ]); // The Available Models list always shows the full recommended set for // every connected provider; entries that are not persisted yet render @@ -313,6 +316,7 @@ export async function getTaskModelSettingsCommand( persistedEnvVarNames, chatgptConnected, githubCopilotConnected, + xaiSubscriptionConnected, }); // Models selected for a runtime role (or as the default coding model) stay // listed and active even when a release drops them from the recommended @@ -414,12 +418,14 @@ export async function getTaskModelProviderSetupCommand( setupNewState, chatgptConnected, githubCopilotConnected, + xaiSubscriptionConnected, ] = await Promise.all([ getDeploymentRuntimeModelConfig(), getPersistedEnvironmentVariableNames(), getDeploymentSetupNewState(), isChatGptSubscriptionConnected(), isGitHubCopilotSubscriptionConnected(), + isXaiSubscriptionConnected(), ]); // Include non-secret OpenAI-compatible env values (base URLs + connection @@ -459,6 +465,7 @@ export async function getTaskModelProviderSetupCommand( selectedProvider: setupNewState.modelProvider, chatgptConnected, githubCopilotConnected, + xaiSubscriptionConnected, }); return { providerSetup }; @@ -554,14 +561,17 @@ export async function saveTaskModelProviderCommand( if (provider.authKind === 'oauth') { throw new Error( - `${provider.label} is connected with a ChatGPT account from the Models settings page and does not use an API key.`, + `${provider.label} is connected with a subscription account from the Models settings page and does not use an API key.`, ); } // When remapping a newly named OpenAI-compatible connection,, rewrite the // primary base URL key by treating apiKey as the template primary value and // collecting against the named descriptor. - const chatgptConnected = await isChatGptSubscriptionConnected(); + const [chatgptConnected, xaiSubscriptionConnected] = await Promise.all([ + isChatGptSubscriptionConnected(), + isXaiSubscriptionConnected(), + ]); let addedRecommendedModelCount = 0; @@ -614,6 +624,7 @@ export async function saveTaskModelProviderCommand( runtimeEnv: process.env, persistedEnvVarNames, chatgptConnected, + xaiSubscriptionConnected, }), ]); const autoAdd = buildAutoAddedTaskModelSettings({ @@ -827,7 +838,7 @@ export async function deleteTaskModelProviderCommand( if (provider.authKind === 'oauth') { throw new Error( - `${provider.label} is connected with a ChatGPT account and cannot be deleted here.`, + `${provider.label} is connected with a subscription account and cannot be deleted here.`, ); } @@ -837,12 +848,14 @@ export async function deleteTaskModelProviderCommand( persistedEnvVarNames, setupNewState, chatgptConnected, + xaiSubscriptionConnected, persistedTaskModelSettings, ] = await Promise.all([ getDeploymentRuntimeModelConfig(), getPersistedEnvironmentVariableNames(tx), getDeploymentSetupNewState(tx), isChatGptSubscriptionConnected(), + isXaiSubscriptionConnected(), getDeploymentTaskModelSettings(), ]); @@ -852,16 +865,30 @@ export async function deleteTaskModelProviderCommand( persistedEnvVarNames, selectedProvider: setupNewState.modelProvider, chatgptConnected, + xaiSubscriptionConnected, }); const providerStatus = providerSetup.providers.find( (candidate) => candidate.id === provider.id, ); + // xAI dual-path: API key and SuperGrok share the `xai` catalog id. When + // deleting the key while a subscription remains, only strip env vars and + // leave models/runtime intact. + const xaiKeyOnlyDelete = + provider.id === 'xai' && + xaiSubscriptionConnected && + Boolean(providerSetup.xaiApiKeyConnected); + if (!providerStatus?.savedApiKeySatisfied) { throw new Error(`${provider.label} does not have saved credentials.`); } - if (getConnectedModelProviderCount(providerSetup) <= 1) { + // Key-only xAI delete is allowed even when the catalog shows a single + // connected `xai` entry, because the subscription remains as a provider. + if ( + !xaiKeyOnlyDelete && + getConnectedModelProviderCount(providerSetup) <= 1 + ) { throw new Error('Keep at least one inference provider connected.'); } @@ -878,6 +905,11 @@ export async function deleteTaskModelProviderCommand( ); } + if (xaiKeyOnlyDelete) { + // Subscription still covers xai/* models; do not cascade model removal. + return; + } + const nextModelState = removeTaskModelsForProvider({ settings: persistedTaskModelSettings, runtimeModelConfig: persistedRuntimeModelConfig, @@ -920,11 +952,13 @@ export async function getLaunchTaskModelsCommand(_auth: UserAuthSuccess) { settings, chatgptConnected, githubCopilotConnected, + xaiSubscriptionConnected, persistedEnvVarNames, ] = await Promise.all([ getDeploymentTaskModelSettings(), isChatGptSubscriptionConnected(), isGitHubCopilotSubscriptionConnected(), + isXaiSubscriptionConnected(), getPersistedEnvironmentVariableNames(), ]); const providerSetup = buildSetupModelStatus({ @@ -932,21 +966,27 @@ export async function getLaunchTaskModelsCommand(_auth: UserAuthSuccess) { persistedEnvVarNames, chatgptConnected, githubCopilotConnected, + xaiSubscriptionConnected, }); - const openaiConnected = Boolean( - providerSetup.providers.find( - (provider) => - provider.id === 'openai' && - (provider.savedApiKeySatisfied || provider.runtimeApiKeySatisfied), - ), - ); + const isApiKeyProviderConnected = (providerId: SetupModelProviderId) => + Boolean( + providerSetup.providers.find( + (provider) => + provider.id === providerId && + (provider.savedApiKeySatisfied || provider.runtimeApiKeySatisfied), + ), + ); const enabledModels = getEnabledTaskModels(settings); const defaultModel = getDefaultTaskModel(settings); return { defaultModelId: defaultModel.id, chatgptConnected, - openaiConnected, + openaiConnected: isApiKeyProviderConnected('openai'), + // Display-grouping inputs: `xai/` models group under the Grok + // subscription only when the API-key sibling is not also connected. + xaiSubscriptionConnected, + xaiConnected: isApiKeyProviderConnected('xai'), models: enabledModels.map((option) => ({ ...option, isDefault: option.id === defaultModel.id, diff --git a/apps/web/src/trpc/commands/xai-subscription/index.test.ts b/apps/web/src/trpc/commands/xai-subscription/index.test.ts new file mode 100644 index 000000000..76bec8bb6 --- /dev/null +++ b/apps/web/src/trpc/commands/xai-subscription/index.test.ts @@ -0,0 +1,103 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const { + mockDisconnectXaiSubscription, + mockGetXaiSubscriptionStatus, + mockIsXaiSubscriptionConnected, + mockPollXaiDeviceAuth, + mockStartXaiDeviceAuth, +} = vi.hoisted(() => ({ + mockDisconnectXaiSubscription: vi.fn(), + mockGetXaiSubscriptionStatus: vi.fn(), + mockIsXaiSubscriptionConnected: vi.fn(), + mockPollXaiDeviceAuth: vi.fn(), + mockStartXaiDeviceAuth: vi.fn(), +})); + +vi.mock('@roomote/db/server', () => ({ + disconnectXaiSubscription: mockDisconnectXaiSubscription, + getXaiSubscriptionStatus: mockGetXaiSubscriptionStatus, + isXaiSubscriptionConnected: mockIsXaiSubscriptionConnected, + pollXaiDeviceAuth: mockPollXaiDeviceAuth, + startXaiDeviceAuth: mockStartXaiDeviceAuth, +})); + +import { + disconnectXaiSubscriptionCommand, + pollXaiDeviceAuthCommand, + startXaiDeviceAuthCommand, +} from './index'; + +const adminAuth = { + userId: 'user-1', + isAdmin: true, +} as never; + +const nonAdminAuth = { + userId: 'user-2', + isAdmin: false, +} as never; + +describe('xai subscription commands', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('rejects non-admin operators', async () => { + await expect(startXaiDeviceAuthCommand(nonAdminAuth)).rejects.toThrow( + 'Unauthorized', + ); + await expect( + pollXaiDeviceAuthCommand(nonAdminAuth, { deviceCode: 'device-1' }), + ).rejects.toThrow('Unauthorized'); + await expect( + disconnectXaiSubscriptionCommand(nonAdminAuth), + ).rejects.toThrow('Unauthorized'); + }); + + it('never returns access or refresh tokens from poll success', async () => { + mockPollXaiDeviceAuth.mockResolvedValue({ + status: 'success', + record: { + refresh: 'refresh-secret', + access: 'access-secret', + expires: Date.now() + 60_000, + status: 'connected', + connectedAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-01T00:00:00.000Z', + }, + }); + + const result = await pollXaiDeviceAuthCommand(adminAuth, { + deviceCode: 'device-1', + }); + + expect(result).toEqual({ status: 'success' }); + expect(result).not.toHaveProperty('record'); + expect(result).not.toHaveProperty('access'); + expect(result).not.toHaveProperty('refresh'); + expect(JSON.stringify(result)).not.toContain('access-secret'); + expect(JSON.stringify(result)).not.toContain('refresh-secret'); + }); + + it('forwards pending and failed poll results without token fields', async () => { + mockPollXaiDeviceAuth.mockResolvedValueOnce({ + status: 'pending', + intervalMs: 7_000, + }); + await expect( + pollXaiDeviceAuthCommand(adminAuth, { deviceCode: 'device-1' }), + ).resolves.toEqual({ status: 'pending', intervalMs: 7_000 }); + + mockPollXaiDeviceAuth.mockResolvedValueOnce({ + status: 'failed', + error: 'xAI device authorization was denied.', + }); + await expect( + pollXaiDeviceAuthCommand(adminAuth, { deviceCode: 'device-1' }), + ).resolves.toEqual({ + status: 'failed', + error: 'xAI device authorization was denied.', + }); + }); +}); diff --git a/apps/web/src/trpc/commands/xai-subscription/index.ts b/apps/web/src/trpc/commands/xai-subscription/index.ts new file mode 100644 index 000000000..300ede794 --- /dev/null +++ b/apps/web/src/trpc/commands/xai-subscription/index.ts @@ -0,0 +1,71 @@ +import { + disconnectXaiSubscription, + getXaiSubscriptionStatus, + isXaiSubscriptionConnected, + pollXaiDeviceAuth, + startXaiDeviceAuth, + type XaiSubscriptionPublicStatus, +} from '@roomote/db/server'; + +import type { UserAuthSuccess } from '@/types'; + +function assertAdmin(auth: UserAuthSuccess): void { + if (!auth.isAdmin) { + throw new Error('Unauthorized'); + } +} + +export async function getXaiSubscriptionStatusCommand( + auth: UserAuthSuccess, +): Promise { + assertAdmin(auth); + return getXaiSubscriptionStatus(); +} + +export async function isXaiSubscriptionConnectedCommand( + auth: UserAuthSuccess, +): Promise { + assertAdmin(auth); + return isXaiSubscriptionConnected(); +} + +export async function startXaiDeviceAuthCommand(auth: UserAuthSuccess) { + assertAdmin(auth); + return startXaiDeviceAuth(); +} + +type PollXaiDeviceAuthResult = + | { status: 'pending'; intervalMs?: number } + | { status: 'success' } + | { status: 'failed'; error: string }; + +export async function pollXaiDeviceAuthCommand( + auth: UserAuthSuccess, + input: { deviceCode: string }, +): Promise { + assertAdmin(auth); + const result = await pollXaiDeviceAuth(input); + + // Never surface the stored OAuth record (which contains refresh/access + // tokens) to the client. Return a token-free status shape only. + if (result.status === 'success') { + return { status: 'success' }; + } + + if (result.status === 'failed') { + return { status: 'failed', error: result.error }; + } + + return { + status: 'pending', + ...(result.intervalMs !== undefined && { intervalMs: result.intervalMs }), + }; +} + +export async function disconnectXaiSubscriptionCommand( + auth: UserAuthSuccess, +): Promise<{ success: true }> { + assertAdmin(auth); + await disconnectXaiSubscription(); + return { success: true }; +} diff --git a/apps/web/src/trpc/routers/_app.ts b/apps/web/src/trpc/routers/_app.ts index 311d2bdf5..11ee9176e 100644 --- a/apps/web/src/trpc/routers/_app.ts +++ b/apps/web/src/trpc/routers/_app.ts @@ -331,6 +331,13 @@ import { pollGitHubCopilotDeviceAuthCommand, startGitHubCopilotDeviceAuthCommand, } from '../commands/github-copilot-subscription'; +import { + disconnectXaiSubscriptionCommand, + getXaiSubscriptionStatusCommand, + isXaiSubscriptionConnectedCommand, + pollXaiDeviceAuthCommand, + startXaiDeviceAuthCommand, +} from '../commands/xai-subscription'; import { getSubscriptionProviderUsageCommand } from '../commands/subscription-usage'; import { getProviderCreditBalancesCommand } from '../commands/provider-credits'; import { @@ -2023,6 +2030,26 @@ export const appRouter = createRouter({ ), }), + xaiSubscription: createRouter({ + status: protectedProcedure.query(({ ctx: { auth } }) => + getXaiSubscriptionStatusCommand(auth), + ), + isConnected: protectedProcedure.query(({ ctx: { auth } }) => + isXaiSubscriptionConnectedCommand(auth), + ), + startDeviceAuth: protectedProcedure.mutation(({ ctx: { auth } }) => + startXaiDeviceAuthCommand(auth), + ), + pollDeviceAuth: protectedProcedure + .input(z.object({ deviceCode: z.string().min(1) })) + .mutation(({ ctx: { auth }, input }) => + pollXaiDeviceAuthCommand(auth, input), + ), + disconnect: protectedProcedure.mutation(({ ctx: { auth } }) => + disconnectXaiSubscriptionCommand(auth), + ), + }), + subscriptionUsage: createRouter({ list: protectedProcedure.query(({ ctx: { auth } }) => getSubscriptionProviderUsageCommand(auth), diff --git a/apps/worker/src/run-task/agent-home.test.ts b/apps/worker/src/run-task/agent-home.test.ts index 9a344f88c..a1ca2df84 100644 --- a/apps/worker/src/run-task/agent-home.test.ts +++ b/apps/worker/src/run-task/agent-home.test.ts @@ -282,6 +282,25 @@ describe('generateOpenCodeConfig provider support', () => { ).toBeUndefined(); }); + it('rebases xAI onto the gateway when the Grok subscription marker is set', () => { + const result = generateOpenCodeConfig({ + homeDir: createHomeDir(), + runtimeEnv: { + R_MODEL: 'xai/grok-4.5', + R_INFERENCE_GATEWAY_URL: 'https://api.example.com/api/inference', + R_INFERENCE_GATEWAY_XAI: '1', + }, + }); + const config = JSON.parse(result.configContent) as { + provider: Record }>; + }; + + expect(config.provider.xai?.options).toMatchObject({ + baseURL: 'https://api.example.com/api/inference/xai/v1', + apiKey: '{env:ROOMOTE_CLOUD_TOKEN}', + }); + }); + it('leaves the openai provider alone when a ChatGPT subscription is present', () => { const result = generateOpenCodeConfig({ homeDir: createHomeDir(), diff --git a/apps/worker/src/run-task/agent-home.ts b/apps/worker/src/run-task/agent-home.ts index 2467c3596..2064a9262 100644 --- a/apps/worker/src/run-task/agent-home.ts +++ b/apps/worker/src/run-task/agent-home.ts @@ -21,6 +21,8 @@ import { INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME, INFERENCE_GATEWAY_REGION_PATTERN, INFERENCE_GATEWAY_URL_ENV_VAR_NAME, + INFERENCE_GATEWAY_XAI_ENV_VAR_NAME, + XAI_OPENCODE_PROVIDER_ID, type InferenceGatewayProvider, isConfiguredEnvValue, isOpenAiCompatibleProviderId, @@ -912,6 +914,8 @@ function mergeInferenceGatewayProviderConfig( runtimeEnv[INFERENCE_GATEWAY_CHATGPT_ENV_VAR_NAME] === '1'; const routeGitHubCopilotThroughGateway = runtimeEnv[INFERENCE_GATEWAY_GITHUB_COPILOT_ENV_VAR_NAME] === '1'; + const routeXaiThroughGateway = + runtimeEnv[INFERENCE_GATEWAY_XAI_ENV_VAR_NAME] === '1'; if (!gatewayUrl) { return providerConfig; @@ -981,6 +985,21 @@ function mergeInferenceGatewayProviderConfig( } } + // xAI Grok subscription served by the gateway: rebase even when no + // XAI_API_KEY was withheld (OAuth-only setups). + if (routeXaiThroughGateway) { + const xaiProvider = getInferenceGatewayProvider(XAI_OPENCODE_PROVIDER_ID); + + if (xaiProvider) { + merged = rebaseProviderOntoGateway( + merged, + XAI_OPENCODE_PROVIDER_ID, + gatewayUrl, + xaiProvider, + ); + } + } + // These providers are configured explicitly because OpenCode has no catalog // entries for their arbitrary model IDs. Rebase only registered gateway // providers: vLLM stays direct until the gateway declares its route. diff --git a/apps/worker/src/run-task/env.ts b/apps/worker/src/run-task/env.ts index 15c20a5cb..b31501768 100644 --- a/apps/worker/src/run-task/env.ts +++ b/apps/worker/src/run-task/env.ts @@ -26,6 +26,7 @@ const MODEL_RUNTIME_ENV_KEYS = [ 'R_INFERENCE_GATEWAY_KEYS', 'R_INFERENCE_GATEWAY_CHATGPT', 'R_INFERENCE_GATEWAY_GITHUB_COPILOT', + 'R_INFERENCE_GATEWAY_XAI', 'R_MODEL', 'R_SMALL_MODEL', 'R_VISION_MODEL', diff --git a/apps/worker/src/run-task/run-task.ts b/apps/worker/src/run-task/run-task.ts index 009c68197..8b69c989c 100644 --- a/apps/worker/src/run-task/run-task.ts +++ b/apps/worker/src/run-task/run-task.ts @@ -5,6 +5,7 @@ import { DISABLED_MODEL_PROVIDER_ENV_VAR_NAMES, INFERENCE_GATEWAY_CHATGPT_ENV_VAR_NAME, INFERENCE_GATEWAY_GITHUB_COPILOT_ENV_VAR_NAME, + INFERENCE_GATEWAY_XAI_ENV_VAR_NAME, INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME, INFERENCE_GATEWAY_URL_ENV_VAR_NAME, isTaskModelIdDisabled, @@ -710,11 +711,14 @@ export const runTask = async ({ unsanitizedEnv[INFERENCE_GATEWAY_CHATGPT_ENV_VAR_NAME] === '1'; const inferenceGatewayGitHubCopilot = unsanitizedEnv[INFERENCE_GATEWAY_GITHUB_COPILOT_ENV_VAR_NAME] === '1'; + const inferenceGatewayXai = + unsanitizedEnv[INFERENCE_GATEWAY_XAI_ENV_VAR_NAME] === '1'; if ( inferenceGatewayServedKeys.length > 0 || inferenceGatewayChatGpt || - inferenceGatewayGitHubCopilot + inferenceGatewayGitHubCopilot || + inferenceGatewayXai ) { for (const servedKey of inferenceGatewayServedKeys) { delete runtimeEnv[servedKey]; @@ -744,11 +748,18 @@ export const runTask = async ({ } else { delete runtimeEnv[INFERENCE_GATEWAY_GITHUB_COPILOT_ENV_VAR_NAME]; } + if (inferenceGatewayXai) { + runtimeEnv[INFERENCE_GATEWAY_XAI_ENV_VAR_NAME] = '1'; + delete runtimeEnv[OPENCODE_AUTH_CONTENT_ENV_VAR_NAME]; + } else { + delete runtimeEnv[INFERENCE_GATEWAY_XAI_ENV_VAR_NAME]; + } } else { delete runtimeEnv[INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME]; delete runtimeEnv[INFERENCE_GATEWAY_URL_ENV_VAR_NAME]; delete runtimeEnv[INFERENCE_GATEWAY_CHATGPT_ENV_VAR_NAME]; delete runtimeEnv[INFERENCE_GATEWAY_GITHUB_COPILOT_ENV_VAR_NAME]; + delete runtimeEnv[INFERENCE_GATEWAY_XAI_ENV_VAR_NAME]; } const workerHomeDir = runtimeEnv.HOME ?? sanitizedEnv.HOME ?? ''; diff --git a/apps/worker/src/sandbox-server/lib/harnesses/opencode-server/bootstrap.ts b/apps/worker/src/sandbox-server/lib/harnesses/opencode-server/bootstrap.ts index 910e345c2..32c646476 100644 --- a/apps/worker/src/sandbox-server/lib/harnesses/opencode-server/bootstrap.ts +++ b/apps/worker/src/sandbox-server/lib/harnesses/opencode-server/bootstrap.ts @@ -341,8 +341,13 @@ async function materializeOpenCodeAuthJson(options: { commandEnv[INFERENCE_GATEWAY_CHATGPT_ENV_VAR_NAME] === '1'; const routeGitHubCopilotThroughGateway = commandEnv.R_INFERENCE_GATEWAY_GITHUB_COPILOT === '1'; + const routeXaiThroughGateway = commandEnv.R_INFERENCE_GATEWAY_XAI === '1'; - if (routeChatGptThroughGateway || routeGitHubCopilotThroughGateway) { + if ( + routeChatGptThroughGateway || + routeGitHubCopilotThroughGateway || + routeXaiThroughGateway + ) { // Defense in depth: dequeue and run-task already omit this value, but a // conflicting deployment env must not survive into the OpenCode process. delete commandEnv[OPENCODE_AUTH_CONTENT_ENV_VAR_NAME]; diff --git a/packages/db/src/lib/__tests__/subscription-provider-usage.test.ts b/packages/db/src/lib/__tests__/subscription-provider-usage.test.ts index 96aa39c05..6251f9fef 100644 --- a/packages/db/src/lib/__tests__/subscription-provider-usage.test.ts +++ b/packages/db/src/lib/__tests__/subscription-provider-usage.test.ts @@ -1,17 +1,31 @@ -import { describe, expect, it, vi } from 'vitest'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const { mockGetFreshXaiAccessToken } = vi.hoisted(() => ({ + mockGetFreshXaiAccessToken: vi.fn(), +})); vi.mock('../../encryption', () => ({ encryptJSON: (value: unknown) => JSON.stringify(value), decryptSecrets: async (value: string) => JSON.parse(value) as unknown, })); +vi.mock('../xai-subscription', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + getFreshXaiAccessToken: mockGetFreshXaiAccessToken, + }; +}); + import { fetchChatGptUsage, fetchGitHubCopilotUsage, fetchKimiForCodingUsage, + fetchXaiSubscriptionUsage, fetchZaiCodingPlanUsage, fetchZaiUsage, getSubscriptionProviderUsage, + parseXaiSubscriptionUsage, parseZaiQuotaUsage, } from '../subscription-provider-usage'; @@ -412,6 +426,216 @@ describe('getSubscriptionProviderUsage', () => { }); }); +describe('parseXaiSubscriptionUsage', () => { + it('normalizes included usage, credits, and on-demand windows', () => { + const windows = parseXaiSubscriptionUsage( + { + included_usage: { + used_percent: 42.5, + used: 425, + remaining: 575, + limit: 1000, + resets_at: '2026-08-01T00:00:00.000Z', + }, + credits: { balance: 12.5 }, + on_demand: { used: 3, limit: 10 }, + }, + Date.parse('2026-07-01T00:00:00.000Z'), + ); + + expect(windows).toEqual([ + { + label: 'Included usage', + usedPercent: 42.5, + used: 425, + remaining: 575, + limit: 1000, + resetsAt: '2026-08-01T00:00:00.000Z', + }, + { + label: 'Credits', + remaining: 12.5, + }, + { + label: 'On-demand', + used: 3, + limit: 10, + usedPercent: 30, + }, + ]); + }); + + it('parses the live cli-chat-proxy format=credits payload', () => { + const windows = parseXaiSubscriptionUsage( + { + config: { + currentPeriod: { + type: 'USAGE_PERIOD_TYPE_WEEKLY', + start: '2026-07-24T01:21:57.505552+00:00', + end: '2026-07-31T01:21:57.505552+00:00', + }, + creditUsagePercent: 27, + onDemandCap: { val: 0 }, + onDemandUsed: { val: 0 }, + productUsage: [ + { product: 'GrokBuild', usagePercent: 26 }, + { product: 'Api', usagePercent: 1 }, + { product: 'GrokChat' }, + { product: 'GrokVoice' }, + ], + prepaidBalance: { val: 0 }, + billingPeriodEnd: '2026-07-31T01:21:57.505552+00:00', + }, + }, + Date.parse('2026-07-26T00:00:00.000Z'), + ); + + // Only the aggregate pool renders: the per-product entries slice the + // same shared pool (duplicate-looking bars), and a zero on-demand + // credit balance is the idle state, not a warning. + expect(windows).toEqual([ + { + label: 'Included usage', + usedPercent: 27, + resetsAt: '2026-07-31T01:21:57.505Z', + }, + ]); + }); + + it('shows the credit balance only when on-demand credits exist', () => { + const windows = parseXaiSubscriptionUsage( + { + config: { + creditUsagePercent: 100, + prepaidBalance: { val: 12.5 }, + billingPeriodEnd: '2026-07-31T01:21:57.505552+00:00', + }, + }, + Date.parse('2026-07-26T00:00:00.000Z'), + ); + + // Pool exhausted with a positive balance: tasks continue on metered + // spend, which is exactly when the operator needs to see the number. + expect(windows).toEqual([ + { + label: 'Included usage', + usedPercent: 100, + resetsAt: '2026-07-31T01:21:57.505Z', + }, + { + label: 'Credits', + remaining: 12.5, + }, + ]); + }); + + it('parses the live monthly billing payload', () => { + const windows = parseXaiSubscriptionUsage( + { + config: { + monthlyLimit: { val: 150_000 }, + used: { val: 73_743 }, + billingPeriodEnd: '2026-08-01T00:00:00+00:00', + }, + }, + Date.parse('2026-07-26T00:00:00.000Z'), + ); + + expect(windows).toEqual([ + { + label: 'Included usage', + used: 73_743, + limit: 150_000, + usedPercent: 49.162, + resetsAt: '2026-08-01T00:00:00.000Z', + }, + ]); + }); + + it('returns an empty list for unusable payloads', () => { + expect(parseXaiSubscriptionUsage(null)).toEqual([]); + expect(parseXaiSubscriptionUsage({})).toEqual([]); + expect(parseXaiSubscriptionUsage({ plan: 'pro' })).toEqual([]); + }); +}); + +describe('fetchXaiSubscriptionUsage', () => { + beforeEach(() => { + mockGetFreshXaiAccessToken.mockReset(); + }); + + it('fetches user then billing and returns normalized windows', async () => { + mockGetFreshXaiAccessToken.mockResolvedValue({ + access: 'xai-access', + expires: Date.now() + 3_600_000, + }); + const fetchImpl = vi + .fn() + .mockResolvedValueOnce(jsonResponse({ userId: 'user-1' })) + .mockResolvedValueOnce( + jsonResponse({ + includedUsage: { + usedPercent: 10, + limit: 100, + remaining: 90, + }, + }), + ); + + const usage = await fetchXaiSubscriptionUsage({ fetchImpl }); + + expect(usage).toMatchObject({ + providerId: 'xai-subscription', + windows: [ + { + label: 'Included usage', + usedPercent: 10, + limit: 100, + remaining: 90, + }, + ], + }); + expect(fetchImpl).toHaveBeenNthCalledWith( + 1, + 'https://cli-chat-proxy.grok.com/v1/user', + expect.objectContaining({ + headers: expect.objectContaining({ + authorization: 'Bearer xai-access', + }), + }), + ); + expect(fetchImpl).toHaveBeenNthCalledWith( + 2, + 'https://cli-chat-proxy.grok.com/v1/billing?format=credits', + expect.objectContaining({ + headers: expect.objectContaining({ + authorization: 'Bearer xai-access', + }), + }), + ); + }); + + it('returns null when no subscription is connected', async () => { + mockGetFreshXaiAccessToken.mockResolvedValue(null); + const fetchImpl = vi.fn(); + await expect(fetchXaiSubscriptionUsage({ fetchImpl })).resolves.toBeNull(); + expect(fetchImpl).not.toHaveBeenCalled(); + }); + + it('returns null when the billing payload cannot be parsed', async () => { + mockGetFreshXaiAccessToken.mockResolvedValue({ + access: 'xai-access', + expires: Date.now() + 3_600_000, + }); + const fetchImpl = vi + .fn() + .mockResolvedValueOnce(jsonResponse({ id: 'user-1' })) + .mockResolvedValueOnce(jsonResponse({ status: 'ok' })); + + await expect(fetchXaiSubscriptionUsage({ fetchImpl })).resolves.toBeNull(); + }); +}); + describe('parseZaiQuotaUsage', () => { it('normalizes live monitor quota windows and plan level', () => { const parsed = parseZaiQuotaUsage({ diff --git a/packages/db/src/lib/__tests__/xai-subscription.test.ts b/packages/db/src/lib/__tests__/xai-subscription.test.ts new file mode 100644 index 000000000..e5cbc6e82 --- /dev/null +++ b/packages/db/src/lib/__tests__/xai-subscription.test.ts @@ -0,0 +1,870 @@ +import { describe, expect, it, vi } from 'vitest'; + +vi.mock('../../encryption', () => ({ + encryptJSON: (value: unknown) => JSON.stringify(value), + decryptSecrets: async (value: string) => JSON.parse(value) as unknown, +})); + +import { + XAI_OAUTH_DEVICE_CODE_ENDPOINT, + XAI_OAUTH_SCOPE, + XAI_OAUTH_TOKEN_ENDPOINT, + XAI_REFRESH_SAFETY_MARGIN_MS, +} from '@roomote/types'; + +import { + disconnectXaiSubscription, + getFreshXaiAccessToken, + getXaiSubscriptionStatus, + pollXaiDeviceAuth, + resolveXaiOAuthClientId, + startXaiDeviceAuth, + XAI_SUBSCRIPTION_INTERNAL, +} from '../xai-subscription'; + +function secretNameFromWhere(clause: { + queryChunks?: unknown[]; +}): string | null { + const chunks = clause.queryChunks ?? []; + for (const chunk of chunks) { + if (typeof chunk === 'string' && chunk.startsWith('XAI_')) { + return chunk; + } + } + return null; +} + +function makeExecutor( + initial: Record = {}, + options: { withTransaction?: boolean } = {}, +) { + const store = new Map( + Object.entries(initial).map(([name, value]) => [ + name, + JSON.stringify(value), + ]), + ); + + const onConflictDoUpdate = vi.fn().mockResolvedValue(undefined); + const values = vi + .fn() + .mockImplementation((row: { name: string; value: string }) => { + store.set(row.name, row.value); + return { onConflictDoUpdate }; + }); + const insert = vi.fn().mockReturnValue({ values }); + + const limit = vi + .fn() + .mockImplementation( + async function limitImpl(this: { __secretName?: string | null }) { + const name = this.__secretName; + if (!name) { + return []; + } + const value = store.get(name); + return value ? [{ value }] : []; + }, + ); + + const where = vi + .fn() + .mockImplementation((clause: { queryChunks?: unknown[] }) => { + const secretName = secretNameFromWhere(clause); + return { + limit: limit.bind({ __secretName: secretName }), + // delete().where() resolves directly + then: undefined, + }; + }); + + // delete uses where that resolves as a promise when awaited + const delWhere = vi + .fn() + .mockImplementation(async (clause: { queryChunks?: unknown[] }) => { + const secretName = secretNameFromWhere(clause); + if (secretName) { + store.delete(secretName); + } + }); + const del = vi.fn().mockReturnValue({ where: delWhere }); + + const from = vi.fn().mockReturnValue({ where }); + const select = vi.fn().mockReturnValue({ from }); + const execute = vi.fn().mockResolvedValue(undefined); + + const base = { insert, select, delete: del, execute, store }; + + if (options.withTransaction) { + const transaction = vi.fn(async (fn: (tx: unknown) => Promise) => { + await fn(base); + }); + return { + executor: { ...base, transaction } as never, + values, + insert, + select, + del, + delWhere, + store, + transaction, + }; + } + + return { + executor: base as never, + values, + insert, + select, + del, + delWhere, + store, + }; +} + +function activePending( + deviceCode = 'device-1', + expiresAt = Date.now() + 900_000, + claimId = 'claim-1', +) { + return { + claimId, + deviceCode, + expiresAt, + startedAt: new Date().toISOString(), + }; +} + +function makeJwt(payload: object): string { + const header = Buffer.from(JSON.stringify({ alg: 'none' })).toString( + 'base64url', + ); + const body = Buffer.from(JSON.stringify(payload)).toString('base64url'); + return `${header}.${body}.sig`; +} + +describe('xAI OAuth client id', () => { + it('defaults to the public Grok CLI client id and accepts an override', () => { + expect(resolveXaiOAuthClientId({})).toBe( + 'b1a00492-073a-47ea-816f-4c329264a828', + ); + expect( + resolveXaiOAuthClientId({ + XAI_OAUTH_CLIENT_ID: ' roomote-xai-client ', + }), + ).toBe('roomote-xai-client'); + }); +}); + +describe('startXaiDeviceAuth', () => { + it('starts the Grok CLI device-code flow with the public client id and scopes', async () => { + const { executor, store } = makeExecutor(); + const fetchImpl = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + device_code: 'device-1', + user_code: 'ABCD-EFGH', + verification_uri: 'https://accounts.x.ai/device', + verification_uri_complete: + 'https://accounts.x.ai/device?user_code=ABCD-EFGH', + interval: 5, + expires_in: 900, + }), + { status: 200 }, + ), + ); + + const now = 1_700_000_000_000; + await expect( + startXaiDeviceAuth(fetchImpl, { executor, now }), + ).resolves.toEqual({ + deviceCode: 'device-1', + userCode: 'ABCD-EFGH', + verificationUrl: 'https://accounts.x.ai/device?user_code=ABCD-EFGH', + intervalMs: 5_000, + expiresInMs: 900_000, + }); + + expect(fetchImpl).toHaveBeenCalledWith( + XAI_OAUTH_DEVICE_CODE_ENDPOINT, + expect.objectContaining({ + method: 'POST', + }), + ); + const [, init] = fetchImpl.mock.calls[0] as [string, RequestInit]; + const body = new URLSearchParams(String(init.body)); + expect(body.get('client_id')).toBe('b1a00492-073a-47ea-816f-4c329264a828'); + expect(body.get('scope')).toBe(XAI_OAUTH_SCOPE); + + const pending = JSON.parse( + store.get(XAI_SUBSCRIPTION_INTERNAL.pendingSecretName)!, + ) as { + claimId: string; + deviceCode: string; + expiresAt: number; + startedAt: string; + }; + expect(pending.deviceCode).toBe('device-1'); + expect(pending.expiresAt).toBe(now + 900_000); + expect(pending.startedAt).toBe(new Date(now).toISOString()); + expect(pending.claimId).toEqual(expect.any(String)); + }); + + it('supersedes a prior pending device code on restart', async () => { + const { executor, store } = makeExecutor({ + [XAI_SUBSCRIPTION_INTERNAL.pendingSecretName]: + activePending('old-device'), + }); + const fetchImpl = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + device_code: 'new-device', + user_code: 'WXYZ-1234', + verification_uri: 'https://accounts.x.ai/device', + interval: 5, + expires_in: 900, + }), + { status: 200 }, + ), + ); + + await startXaiDeviceAuth(fetchImpl, { executor, now: Date.now() }); + + const pending = JSON.parse( + store.get(XAI_SUBSCRIPTION_INTERNAL.pendingSecretName)!, + ) as { deviceCode: string }; + expect(pending.deviceCode).toBe('new-device'); + }); + + it('does not let a slower earlier start overwrite a newer claim', async () => { + const { executor, store } = makeExecutor(); + let releaseEarlyHttp: (() => void) | undefined; + const earlyHttpGate = new Promise((resolve) => { + releaseEarlyHttp = resolve; + }); + + const earlyFetch = vi.fn().mockImplementation(async () => { + await earlyHttpGate; + return new Response( + JSON.stringify({ + device_code: 'stale-device', + user_code: 'STALE-CODE', + verification_uri: 'https://accounts.x.ai/device', + interval: 5, + expires_in: 900, + }), + { status: 200 }, + ); + }); + + const lateFetch = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + device_code: 'fresh-device', + user_code: 'FRESH-CODE', + verification_uri: 'https://accounts.x.ai/device', + interval: 5, + expires_in: 900, + }), + { status: 200 }, + ), + ); + + const earlyStart = startXaiDeviceAuth(earlyFetch, { + executor, + now: Date.now(), + }); + // Allow the early start to reserve its claim before the later start. + await Promise.resolve(); + await Promise.resolve(); + + const lateResult = await startXaiDeviceAuth(lateFetch, { + executor, + now: Date.now() + 1, + }); + expect(lateResult.deviceCode).toBe('fresh-device'); + + releaseEarlyHttp?.(); + await expect(earlyStart).rejects.toThrow( + XAI_SUBSCRIPTION_INTERNAL.supersededDeviceFlowError, + ); + + const pending = JSON.parse( + store.get(XAI_SUBSCRIPTION_INTERNAL.pendingSecretName)!, + ) as { deviceCode: string }; + expect(pending.deviceCode).toBe('fresh-device'); + }); + + it('fails closed when the device-code endpoint returns an error', async () => { + const { executor, store } = makeExecutor(); + const fetchImpl = vi + .fn() + .mockResolvedValue(new Response('nope', { status: 503 })); + + await expect(startXaiDeviceAuth(fetchImpl, { executor })).rejects.toThrow( + 'Failed to initiate xAI device authorization: 503', + ); + // Failed start clears its provisional claim. + expect(store.has(XAI_SUBSCRIPTION_INTERNAL.pendingSecretName)).toBe(false); + }); +}); + +describe('pollXaiDeviceAuth', () => { + it('reports pending authorization without storing a credential', async () => { + const { executor, values, store } = makeExecutor({ + [XAI_SUBSCRIPTION_INTERNAL.pendingSecretName]: activePending('device-1'), + }); + const result = await pollXaiDeviceAuth( + { deviceCode: 'device-1' }, + { + executor, + fetchImpl: vi.fn().mockResolvedValue( + new Response(JSON.stringify({ error: 'authorization_pending' }), { + status: 400, + }), + ), + }, + ); + + expect(result).toEqual({ status: 'pending' }); + // values may be unused for pending; subscription secret must stay empty + expect(store.has(XAI_SUBSCRIPTION_INTERNAL.secretName)).toBe(false); + expect(values).not.toHaveBeenCalled(); + }); + + it('honors slow_down with an optional interval bump', async () => { + const { executor, values } = makeExecutor({ + [XAI_SUBSCRIPTION_INTERNAL.pendingSecretName]: activePending('device-1'), + }); + const result = await pollXaiDeviceAuth( + { deviceCode: 'device-1' }, + { + executor, + fetchImpl: vi.fn().mockResolvedValue( + new Response(JSON.stringify({ error: 'slow_down', interval: 10 }), { + status: 400, + }), + ), + }, + ); + + expect(result).toEqual({ status: 'pending', intervalMs: 10_000 }); + expect(values).not.toHaveBeenCalled(); + }); + + it('encrypts and stores tokens after approval without returning them as public status', async () => { + const { executor, values, store } = makeExecutor({ + [XAI_SUBSCRIPTION_INTERNAL.pendingSecretName]: activePending('device-1'), + }); + const idToken = makeJwt({ email: 'user@example.com' }); + const result = await pollXaiDeviceAuth( + { deviceCode: 'device-1' }, + { + executor, + fetchImpl: vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + access_token: 'access-secret', + refresh_token: 'refresh-secret', + expires_in: 3600, + id_token: idToken, + }), + { status: 200 }, + ), + ), + }, + ); + + expect(result.status).toBe('success'); + if (result.status !== 'success') { + throw new Error('expected success'); + } + // The internal record has tokens for persistence, but the public status + // shape used by tRPC never includes them. + expect(result.record).toMatchObject({ + access: 'access-secret', + refresh: 'refresh-secret', + status: 'connected', + email: 'user@example.com', + }); + expect(values).toHaveBeenCalled(); + const stored = JSON.parse( + store.get(XAI_SUBSCRIPTION_INTERNAL.secretName)!, + ) as { access: string; refresh: string }; + expect(stored.access).toBe('access-secret'); + expect(stored.refresh).toBe('refresh-secret'); + // Pending flow cleared after successful connect. + expect(store.has(XAI_SUBSCRIPTION_INTERNAL.pendingSecretName)).toBe(false); + }); + + it('does not persist tokens for a superseded device code', async () => { + const { executor, store } = makeExecutor({ + // Active flow is a newer restart; old poll still holds device-1. + [XAI_SUBSCRIPTION_INTERNAL.pendingSecretName]: activePending('device-2'), + }); + + const result = await pollXaiDeviceAuth( + { deviceCode: 'device-1' }, + { + executor, + fetchImpl: vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + access_token: 'stale-access', + refresh_token: 'stale-refresh', + expires_in: 3600, + }), + { status: 200 }, + ), + ), + }, + ); + + expect(result).toEqual({ + status: 'failed', + error: XAI_SUBSCRIPTION_INTERNAL.supersededDeviceFlowError, + }); + expect(store.has(XAI_SUBSCRIPTION_INTERNAL.secretName)).toBe(false); + // Active newer pending remains so the restart can complete. + expect( + JSON.parse(store.get(XAI_SUBSCRIPTION_INTERNAL.pendingSecretName)!) + .deviceCode, + ).toBe('device-2'); + }); + + it('does not persist when a newer start wins while the token exchange is in flight', async () => { + const { executor, store } = makeExecutor({ + [XAI_SUBSCRIPTION_INTERNAL.pendingSecretName]: activePending('device-1'), + }); + + const fetchImpl = vi.fn().mockImplementation(async () => { + // Simulate cancel/reopen registering a newer device code mid-request. + store.set( + XAI_SUBSCRIPTION_INTERNAL.pendingSecretName, + JSON.stringify(activePending('device-2')), + ); + return new Response( + JSON.stringify({ + access_token: 'old-access', + refresh_token: 'old-refresh', + expires_in: 3600, + }), + { status: 200 }, + ); + }); + + const result = await pollXaiDeviceAuth( + { deviceCode: 'device-1' }, + { executor, fetchImpl }, + ); + + expect(result).toEqual({ + status: 'failed', + error: XAI_SUBSCRIPTION_INTERNAL.supersededDeviceFlowError, + }); + expect(store.has(XAI_SUBSCRIPTION_INTERNAL.secretName)).toBe(false); + expect( + JSON.parse(store.get(XAI_SUBSCRIPTION_INTERNAL.pendingSecretName)!) + .deviceCode, + ).toBe('device-2'); + }); + + it('fails when the grant has no refresh token', async () => { + const { executor } = makeExecutor({ + [XAI_SUBSCRIPTION_INTERNAL.pendingSecretName]: activePending('device-1'), + }); + const result = await pollXaiDeviceAuth( + { deviceCode: 'device-1' }, + { + executor, + fetchImpl: vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + access_token: 'access-only', + }), + { status: 200 }, + ), + ), + }, + ); + + expect(result).toEqual({ + status: 'failed', + error: 'xAI device authorization returned no refresh_token.', + }); + }); + + it('fails on access_denied and expired_token', async () => { + const { executor: executorDenied, store: storeDenied } = makeExecutor({ + [XAI_SUBSCRIPTION_INTERNAL.pendingSecretName]: activePending('device-1'), + }); + + await expect( + pollXaiDeviceAuth( + { deviceCode: 'device-1' }, + { + executor: executorDenied, + fetchImpl: vi.fn().mockResolvedValue( + new Response(JSON.stringify({ error: 'access_denied' }), { + status: 400, + }), + ), + }, + ), + ).resolves.toMatchObject({ status: 'failed' }); + expect(storeDenied.has(XAI_SUBSCRIPTION_INTERNAL.pendingSecretName)).toBe( + false, + ); + + const { executor: executorExpired } = makeExecutor({ + [XAI_SUBSCRIPTION_INTERNAL.pendingSecretName]: activePending('device-1'), + }); + + await expect( + pollXaiDeviceAuth( + { deviceCode: 'device-1' }, + { + executor: executorExpired, + fetchImpl: vi.fn().mockResolvedValue( + new Response(JSON.stringify({ error: 'expired_token' }), { + status: 400, + }), + ), + }, + ), + ).resolves.toMatchObject({ + status: 'failed', + error: expect.stringContaining('expired'), + }); + }); +}); + +describe('getXaiSubscriptionStatus', () => { + it('reports disconnected when no subscription record exists', async () => { + const { executor } = makeExecutor(); + await expect(getXaiSubscriptionStatus(executor)).resolves.toEqual({ + connected: false, + status: 'disconnected', + }); + }); + + it('reports connected without access or refresh tokens', async () => { + const { executor } = makeExecutor({ + [XAI_SUBSCRIPTION_INTERNAL.secretName]: { + refresh: 'rt', + access: 'at', + expires: Date.now() + 60_000, + status: 'connected', + email: 'a@b.com', + connectedAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-01T00:00:00.000Z', + }, + }); + + const status = await getXaiSubscriptionStatus(executor); + expect(status).toEqual({ + connected: true, + status: 'connected', + email: 'a@b.com', + connectedAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-01T00:00:00.000Z', + }); + expect(status).not.toHaveProperty('access'); + expect(status).not.toHaveProperty('refresh'); + }); +}); + +describe('disconnectXaiSubscription', () => { + it('deletes the encrypted deployment secret so status becomes disconnected', async () => { + const connectedRecord = { + refresh: 'rt', + access: 'at', + expires: Date.now() + 60_000, + status: 'connected' as const, + connectedAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-01T00:00:00.000Z', + }; + const { executor, del, delWhere, store } = makeExecutor({ + [XAI_SUBSCRIPTION_INTERNAL.secretName]: connectedRecord, + [XAI_SUBSCRIPTION_INTERNAL.pendingSecretName]: activePending('device-1'), + }); + + // Before disconnect the record is still loadable. + await expect(getXaiSubscriptionStatus(executor)).resolves.toMatchObject({ + connected: true, + status: 'connected', + }); + + await disconnectXaiSubscription(executor); + + expect(del).toHaveBeenCalled(); + expect(delWhere).toHaveBeenCalled(); + expect(store.has(XAI_SUBSCRIPTION_INTERNAL.secretName)).toBe(false); + expect(store.has(XAI_SUBSCRIPTION_INTERNAL.pendingSecretName)).toBe(false); + + await expect(getXaiSubscriptionStatus(executor)).resolves.toEqual({ + connected: false, + status: 'disconnected', + }); + }); + + it('wins over an in-flight poll that already received tokens', async () => { + const { executor, store } = makeExecutor({ + [XAI_SUBSCRIPTION_INTERNAL.pendingSecretName]: activePending('device-1'), + }); + + const fetchImpl = vi.fn().mockImplementation(async () => { + // Operator disconnects while the device poll is awaiting xAI. + await disconnectXaiSubscription(executor); + return new Response( + JSON.stringify({ + access_token: 'should-not-persist', + refresh_token: 'should-not-persist', + expires_in: 3600, + }), + { status: 200 }, + ); + }); + + const result = await pollXaiDeviceAuth( + { deviceCode: 'device-1' }, + { executor, fetchImpl }, + ); + + expect(result).toEqual({ + status: 'failed', + error: XAI_SUBSCRIPTION_INTERNAL.supersededDeviceFlowError, + }); + expect(store.has(XAI_SUBSCRIPTION_INTERNAL.secretName)).toBe(false); + expect(store.has(XAI_SUBSCRIPTION_INTERNAL.pendingSecretName)).toBe(false); + }); +}); + +describe('getFreshXaiAccessToken', () => { + it('returns the existing access token when it is outside the safety margin', async () => { + const expires = Date.now() + XAI_REFRESH_SAFETY_MARGIN_MS + 60_000; + const { executor } = makeExecutor({ + [XAI_SUBSCRIPTION_INTERNAL.secretName]: { + refresh: 'rt', + access: 'still-valid', + expires, + status: 'connected', + connectedAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-01T00:00:00.000Z', + }, + }); + const fetchImpl = vi.fn(); + + const fresh = await getFreshXaiAccessToken({ + executor, + fetchImpl, + now: Date.now(), + }); + + expect(fresh).toEqual({ + access: 'still-valid', + expires, + }); + expect(fresh).not.toHaveProperty('refresh'); + expect(fetchImpl).not.toHaveBeenCalled(); + }); + + it('refreshes an expiring access token and replaces the stored record', async () => { + const now = Date.now(); + const expires = now + 30_000; // well inside the safety margin + const existing = { + refresh: 'old-rt', + access: 'old-at', + expires, + status: 'connected' as const, + connectedAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-01T00:00:00.000Z', + }; + + const { executor, values } = makeExecutor( + { [XAI_SUBSCRIPTION_INTERNAL.secretName]: existing }, + { withTransaction: true }, + ); + + const fetchImpl = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + access_token: 'new-at', + refresh_token: 'new-rt', + expires_in: 7200, + }), + { status: 200 }, + ), + ); + + const fresh = await getFreshXaiAccessToken({ + executor, + fetchImpl, + now, + }); + + expect(fresh).toMatchObject({ + access: 'new-at', + }); + expect(fresh).not.toHaveProperty('refresh'); + expect(fetchImpl).toHaveBeenCalledWith( + XAI_OAUTH_TOKEN_ENDPOINT, + expect.objectContaining({ method: 'POST' }), + ); + const [, init] = fetchImpl.mock.calls[0] as [string, RequestInit]; + const body = new URLSearchParams(String(init.body)); + expect(body.get('grant_type')).toBe('refresh_token'); + expect(body.get('refresh_token')).toBe('old-rt'); + expect(values).toHaveBeenCalled(); + }); + + it('does not clobber a newer subscription when refresh HTTP is still in flight', async () => { + const now = Date.now(); + const expires = now + 30_000; + const existing = { + refresh: 'old-rt', + access: 'old-at', + expires, + status: 'connected' as const, + connectedAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-01T00:00:00.000Z', + }; + const reconnected = { + refresh: 'fresh-rt', + access: 'fresh-at', + expires: now + 3_600_000, + status: 'connected' as const, + connectedAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-01T00:01:00.000Z', + }; + + const { executor, store } = makeExecutor( + { [XAI_SUBSCRIPTION_INTERNAL.secretName]: existing }, + { withTransaction: true }, + ); + + const fetchImpl = vi.fn().mockImplementation(async () => { + // A concurrent device-code connect replaces the secret mid-refresh. + store.set( + XAI_SUBSCRIPTION_INTERNAL.secretName, + JSON.stringify(reconnected), + ); + return new Response( + JSON.stringify({ + access_token: 'stale-refresh-at', + refresh_token: 'stale-refresh-rt', + expires_in: 7200, + }), + { status: 200 }, + ); + }); + + const fresh = await getFreshXaiAccessToken({ + executor, + fetchImpl, + now, + }); + + expect(fresh).toMatchObject({ + access: 'fresh-at', + }); + expect(fresh).not.toHaveProperty('refresh'); + const stored = JSON.parse( + store.get(XAI_SUBSCRIPTION_INTERNAL.secretName)!, + ) as { access: string; refresh: string }; + expect(stored.access).toBe('fresh-at'); + expect(stored.refresh).toBe('fresh-rt'); + }); + + it('returns null when no connected subscription exists', async () => { + const { executor } = makeExecutor(); + await expect( + getFreshXaiAccessToken({ executor, fetchImpl: vi.fn() }), + ).resolves.toBeNull(); + }); + + it('marks the record errored and returns null when refresh fails terminally', async () => { + const now = Date.now(); + const expires = now + 30_000; // inside the safety margin + const existing = { + refresh: 'stale-rt', + access: 'stale-at', + expires, + status: 'connected' as const, + connectedAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-01T00:00:00.000Z', + }; + + const { executor, values, store } = makeExecutor( + { [XAI_SUBSCRIPTION_INTERNAL.secretName]: existing }, + { withTransaction: true }, + ); + + const fetchImpl = vi + .fn() + .mockResolvedValue(new Response('invalid_grant', { status: 401 })); + + const fresh = await getFreshXaiAccessToken({ + executor, + fetchImpl, + now, + }); + + expect(fresh).toBeNull(); + expect(values).toHaveBeenCalled(); + const stored = JSON.parse( + store.get(XAI_SUBSCRIPTION_INTERNAL.secretName)!, + ) as { status: string; error?: string; access: string }; + expect(stored.status).toBe('error'); + expect(stored.error).toMatch(/xAI token refresh failed: 401/); + // Keep the prior tokens so reconnect can be distinguished from wipe. + expect(stored.access).toBe('stale-at'); + }); + + it('keeps connected status on a transient refresh failure and serves remaining access', async () => { + const now = Date.now(); + const expires = now + 30_000; // inside the safety margin but still valid + const existing = { + refresh: 'rt', + access: 'still-usable', + expires, + status: 'connected' as const, + connectedAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-01T00:00:00.000Z', + }; + + const { executor, store } = makeExecutor( + { [XAI_SUBSCRIPTION_INTERNAL.secretName]: existing }, + { withTransaction: true }, + ); + + const fetchImpl = vi + .fn() + .mockResolvedValue(new Response('upstream unavailable', { status: 503 })); + + const fresh = await getFreshXaiAccessToken({ + executor, + fetchImpl, + now, + }); + + expect(fresh).toMatchObject({ + access: 'still-usable', + }); + expect(fresh).not.toHaveProperty('refresh'); + const stored = JSON.parse( + store.get(XAI_SUBSCRIPTION_INTERNAL.secretName)!, + ) as { status: string; access: string }; + expect(stored.status).toBe('connected'); + expect(stored.access).toBe('still-usable'); + }); +}); + +describe('XAI_SUBSCRIPTION_INTERNAL', () => { + it('exposes the deployment secret name for ops tooling', () => { + expect(XAI_SUBSCRIPTION_INTERNAL.secretName).toBe('XAI_SUBSCRIPTION_OAUTH'); + expect(XAI_SUBSCRIPTION_INTERNAL.pendingSecretName).toBe( + 'XAI_SUBSCRIPTION_PENDING_DEVICE', + ); + }); +}); diff --git a/packages/db/src/lib/model-runtime-config.test.ts b/packages/db/src/lib/model-runtime-config.test.ts index 8a6805526..2bb9ea8cc 100644 --- a/packages/db/src/lib/model-runtime-config.test.ts +++ b/packages/db/src/lib/model-runtime-config.test.ts @@ -4,12 +4,14 @@ const { mockEnvironmentVariablesFindMany, mockResolveOpenCodeAuthContent, mockResolveGitHubCopilotOpenCodeAuthContent, + mockGetFreshXaiAccessToken, } = vi.hoisted(() => ({ mockDecryptSecrets: vi.fn(), mockDeploymentSettingsFindFirst: vi.fn(), mockEnvironmentVariablesFindMany: vi.fn(), mockResolveOpenCodeAuthContent: vi.fn(), mockResolveGitHubCopilotOpenCodeAuthContent: vi.fn(), + mockGetFreshXaiAccessToken: vi.fn(), })); vi.mock('../encryption', () => ({ @@ -45,6 +47,11 @@ vi.mock('./github-copilot-subscription', () => ({ mockResolveGitHubCopilotOpenCodeAuthContent(...args), })); +vi.mock('./xai-subscription', () => ({ + getFreshXaiAccessToken: (...args: unknown[]) => + mockGetFreshXaiAccessToken(...args), +})); + vi.mock('../schema', () => ({ deploymentSettings: { id: 'deploymentSettings.id' }, eq: vi.fn(), @@ -62,6 +69,7 @@ describe('resolveEffectiveModelRuntimeEnv', () => { mockEnvironmentVariablesFindMany.mockResolvedValue([]); mockResolveGitHubCopilotOpenCodeAuthContent.mockResolvedValue(null); mockResolveOpenCodeAuthContent.mockResolvedValue(null); + mockGetFreshXaiAccessToken.mockResolvedValue(null); }); it('prefers real runtime env values over persisted deployment config', async () => { @@ -691,6 +699,92 @@ describe('resolveEffectiveModelRuntimeEnv', () => { expect(env).not.toHaveProperty('OPENCODE_AUTH_CONTENT'); }); + it('emits the xAI gateway marker for OAuth-only setups without shipping tokens or a real key', async () => { + mockDeploymentSettingsFindFirst.mockResolvedValue({ + runtimeModelConfig: { + roomoteModel: 'xai/grok-4.5', + }, + }); + mockGetFreshXaiAccessToken.mockResolvedValue({ + access: 'xai-oauth-access', + expires: Date.now() + 3_600_000, + }); + + const env = await resolveSandboxModelRuntimeEnv({ + runtimeEnv: {}, + // No XAI_API_KEY: subscription alone must cover the gateway path. + deploymentEnvVars: {}, + }); + + expect(mockGetFreshXaiAccessToken).toHaveBeenCalled(); + // OAuth record stays on the control plane; marker drives worker rebase. + expect(env).not.toHaveProperty('OPENCODE_AUTH_CONTENT'); + expect(env).not.toHaveProperty('XAI_API_KEY'); + expect(env.R_INFERENCE_GATEWAY_XAI).toBe('1'); + // Advertise XAI_API_KEY in served keys so the worker rebases xai even + // when only the subscription is connected. + expect(env.R_INFERENCE_GATEWAY_KEYS?.split(',')).toContain('XAI_API_KEY'); + }); + + it('does not emit the xAI gateway marker when no subscription is connected', async () => { + mockDeploymentSettingsFindFirst.mockResolvedValue({ + runtimeModelConfig: { + roomoteModel: 'xai/grok-4.5', + }, + }); + mockGetFreshXaiAccessToken.mockResolvedValue(null); + + const env = await resolveSandboxModelRuntimeEnv({ + runtimeEnv: {}, + deploymentEnvVars: {}, + }); + + expect(env).not.toHaveProperty('R_INFERENCE_GATEWAY_XAI'); + expect(env).not.toHaveProperty('OPENCODE_AUTH_CONTENT'); + }); + + it('injects a mint access token as XAI_API_KEY for non-gateway OAuth-only control plane', async () => { + mockDeploymentSettingsFindFirst.mockResolvedValue({ + runtimeModelConfig: { + roomoteModel: 'xai/grok-4.5', + }, + }); + mockGetFreshXaiAccessToken.mockResolvedValue({ + access: 'xai-oauth-access-only', + expires: Date.now() + 3_600_000, + }); + + const env = await resolveEffectiveModelRuntimeEnv({ + runtimeEnv: {}, + deploymentEnvVars: {}, + }); + + // OpenCode's xAI provider is API-key shaped: inject the access token as + // XAI_API_KEY and never ship the refresh token or oauth JSON. + expect(env.XAI_API_KEY).toBe('xai-oauth-access-only'); + expect(env).not.toHaveProperty('OPENCODE_AUTH_CONTENT'); + }); + + it('prefers a connected OAuth access token over BYOK XAI_API_KEY on the control plane', async () => { + mockDeploymentSettingsFindFirst.mockResolvedValue({ + runtimeModelConfig: { + roomoteModel: 'xai/grok-4.5', + }, + }); + mockGetFreshXaiAccessToken.mockResolvedValue({ + access: 'xai-oauth-access-only', + expires: Date.now() + 3_600_000, + }); + + const env = await resolveEffectiveModelRuntimeEnv({ + runtimeEnv: {}, + deploymentEnvVars: { XAI_API_KEY: 'sk-byok-key' }, + }); + + // Match gateway precedence: subscription wins when connected. + expect(env.XAI_API_KEY).toBe('xai-oauth-access-only'); + }); + it('does not inject OPENCODE_AUTH_CONTENT when no openai/ model is used', async () => { mockDeploymentSettingsFindFirst.mockResolvedValue({ runtimeModelConfig: { roomoteModel: 'anthropic/claude-sonnet-4' }, diff --git a/packages/db/src/lib/model-runtime-config.ts b/packages/db/src/lib/model-runtime-config.ts index bb47b847c..ff7293a56 100644 --- a/packages/db/src/lib/model-runtime-config.ts +++ b/packages/db/src/lib/model-runtime-config.ts @@ -10,18 +10,21 @@ import { INFERENCE_GATEWAY_CHATGPT_ENV_VAR_NAME, INFERENCE_GATEWAY_GITHUB_COPILOT_ENV_VAR_NAME, INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME, + INFERENCE_GATEWAY_XAI_ENV_VAR_NAME, isConfiguredEnvValue, isInferenceGatewayCoveredEnvVar, normalizeDeploymentModelConfig, normalizeOptionalReasoningEffort, parseModelProviderEnvKeys, resolveSetupModelProviderIdFromModel, + XAI_OPENCODE_PROVIDER_ID, type TaskModelOption, } from '@roomote/types'; import { decryptSecrets } from '../encryption'; import { resolveOpenCodeAuthContent } from './chatgpt-subscription'; import { resolveGitHubCopilotOpenCodeAuthContent } from './github-copilot-subscription'; +import { getFreshXaiAccessToken } from './xai-subscription'; import { type DatabaseOrTransaction, db } from '../db'; import { deploymentSettings } from '../schema'; @@ -200,9 +203,9 @@ export async function resolveEffectiveModelRuntimeEnv( /** * Resolve model runtime env for a task sandbox: the configured provider keys * the inference gateway can serve stay on the control plane and their names - * are advertised via `R_INFERENCE_GATEWAY_KEYS` instead, and a connected - * ChatGPT subscription is routed through the gateway rather than - * materializing `OPENCODE_AUTH_CONTENT` in the sandbox. + * are advertised via `R_INFERENCE_GATEWAY_KEYS` instead. Connected ChatGPT, + * GitHub Copilot, and xAI Grok subscriptions are routed through the gateway + * via markers rather than materializing OAuth credentials in the sandbox. */ export async function resolveSandboxModelRuntimeEnv( options: ModelRuntimeEnvOptions = {}, @@ -438,6 +441,30 @@ async function resolveModelRuntimeEnv( : null; const routeGitHubCopilotThroughGateway = inferenceGateway && githubCopilotAuthContent != null; + // xAI Grok subscription: when a connected OAuth record exists and any role + // or switchable model uses `xai/`, route through the gateway (marker only) + // or mint a fresh access token for non-gateway control-plane use. + const usesXaiModel = [ + ...resolvedRoleModels, + ...gatewaySwitchableModelIds, + ].some( + (modelId) => + typeof modelId === 'string' && + modelId.startsWith(`${XAI_OPENCODE_PROVIDER_ID}/`), + ); + const xaiAccessToken = usesXaiModel + ? await getFreshXaiAccessToken({ executor }) + : null; + const routeXaiThroughGateway = inferenceGateway && xaiAccessToken != null; + // OpenCode's xAI provider is API-key shaped, not oauth Auth.Info. For + // non-gateway control-plane inference (titles, summaries, routing), mint a + // fresh access token and inject it as XAI_API_KEY. Never put the refresh + // token into OpenCode env. Prefer a connected subscription over BYOK so + // control-plane and gateway dual-path use the same credential precedence. + const xaiApiKeyFromOAuth: Record = + xaiAccessToken && !routeXaiThroughGateway + ? { XAI_API_KEY: xaiAccessToken.access } + : {}; const directOpenCodeAuthContent = [ injectedOpenCodeAuthContent, githubCopilotAuthContent, @@ -445,6 +472,15 @@ async function resolveModelRuntimeEnv( return content ? { ...merged, ...JSON.parse(content) } : merged; }, {}); + // When only the subscription is connected (no XAI_API_KEY), still advertise + // the provider key name so the worker knows xai models are gateway-covered. + // The marker alone drives rebase; R_INFERENCE_GATEWAY_KEYS stays key-based. + const xaiKeyAlreadyServed = gatewayServedKeyNameSet.has('XAI_API_KEY'); + const effectiveGatewayServedKeyNames = + routeXaiThroughGateway && !xaiKeyAlreadyServed + ? [...gatewayServedKeyNames, 'XAI_API_KEY'] + : gatewayServedKeyNames; + return { ...(resolvedRoomoteModel && { R_MODEL: resolvedRoomoteModel }), ...(resolvedRoomoteSmallModel && { @@ -488,8 +524,10 @@ async function resolveModelRuntimeEnv( R_MODEL_ENV_KEYS: providerKeyNames.join(','), }), ...resolvedProviderKeyValues, - ...(gatewayServedKeyNames.length > 0 && { - [INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME]: gatewayServedKeyNames.join(','), + ...xaiApiKeyFromOAuth, + ...(effectiveGatewayServedKeyNames.length > 0 && { + [INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME]: + effectiveGatewayServedKeyNames.join(','), }), ...(routeChatGptThroughGateway ? { [INFERENCE_GATEWAY_CHATGPT_ENV_VAR_NAME]: '1' } @@ -497,6 +535,9 @@ async function resolveModelRuntimeEnv( ...(routeGitHubCopilotThroughGateway ? { [INFERENCE_GATEWAY_GITHUB_COPILOT_ENV_VAR_NAME]: '1' } : {}), + ...(routeXaiThroughGateway + ? { [INFERENCE_GATEWAY_XAI_ENV_VAR_NAME]: '1' } + : {}), ...(!routeChatGptThroughGateway && !routeGitHubCopilotThroughGateway && Object.keys(directOpenCodeAuthContent).length > 0 diff --git a/packages/db/src/lib/subscription-provider-usage.ts b/packages/db/src/lib/subscription-provider-usage.ts index f95223975..5d666f222 100644 --- a/packages/db/src/lib/subscription-provider-usage.ts +++ b/packages/db/src/lib/subscription-provider-usage.ts @@ -3,6 +3,8 @@ import { CHATGPT_USAGE_ENDPOINT, GITHUB_COPILOT_USAGE_ENDPOINT, KIMI_FOR_CODING_USAGE_ENDPOINTS, + XAI_USAGE_BILLING_ENDPOINT, + XAI_USAGE_USER_ENDPOINT, ZAI_USAGE_HOSTS, ZAI_USAGE_QUOTA_PATH, type SubscriptionProviderUsage, @@ -14,13 +16,14 @@ import { type DatabaseOrTransaction } from '../db'; import { getFreshChatGptAccessToken } from './chatgpt-subscription'; import { getGitHubCopilotAccessToken } from './github-copilot-subscription'; import { resolveModelProviderEnvValue } from './model-runtime-config'; +import { getFreshXaiAccessToken } from './xai-subscription'; /** * Server-side usage/quota lookups for subscription-style inference providers, - * using the same credentials the inference gateway already holds. Upstream - * endpoints are undocumented (each is what the provider's own CLI polls), so - * every fetcher parses defensively and resolves to `null` on any failure — - * the settings UI omits the usage line rather than erroring. + * using the same credentials the inference gateway already holds. All three + * upstream endpoints are undocumented (each is what the provider's own CLI + * polls), so every fetcher parses defensively and resolves to `null` on any + * failure — the settings UI omits the usage line rather than erroring. */ const USAGE_FETCH_TIMEOUT_MS = 10_000; @@ -37,21 +40,35 @@ function asRecord(value: unknown): Record | undefined { : undefined; } -/** Kimi's usage payload serializes numbers as strings, so coerce both. */ +/** + * Coerce a bare number/string, or Grok's `{ val: N }` credit wrappers, into a + * finite number. Kimi also serializes numbers as strings. + */ +function asFiniteNumber(value: unknown): number | undefined { + if (typeof value === 'number' && Number.isFinite(value)) { + return value; + } + if (typeof value === 'string' && value.trim().length > 0) { + const parsed = Number(value); + if (Number.isFinite(parsed)) { + return parsed; + } + } + const record = asRecord(value); + if (record && 'val' in record) { + return asFiniteNumber(record.val); + } + return undefined; +} + function firstNumber( source: Record | undefined, keys: readonly string[], ): number | undefined { for (const key of keys) { - const value = source?.[key]; - if (typeof value === 'number' && Number.isFinite(value)) { - return value; - } - if (typeof value === 'string' && value.trim().length > 0) { - const parsed = Number(value); - if (Number.isFinite(parsed)) { - return parsed; - } + const parsed = asFiniteNumber(source?.[key]); + if (parsed !== undefined) { + return parsed; } } return undefined; @@ -532,6 +549,280 @@ export async function fetchKimiForCodingUsage( return null; } +// --- xAI Grok subscription ------------------------------------------------ + +/** + * Live `GET /v1/billing?format=credits` nests usage under `config` with + * `{ val }` wrappers and product-level percents. Older/synthetic shapes put + * fields at the root. Accept both so the settings bar keeps working. + */ +function resolveXaiBillingConfig( + payload: unknown, +): Record | undefined { + const root = asRecord(payload); + if (!root) { + return undefined; + } + return asRecord(root.config) ?? root; +} + +function resolveXaiResetsAt( + config: Record, + now: number, +): string | undefined { + const period = asRecord(config.currentPeriod ?? config.current_period); + return ( + toResetIso( + period?.end ?? + period?.ends_at ?? + period?.endsAt ?? + config.billingPeriodEnd ?? + config.billing_period_end ?? + config.resets_at ?? + config.reset_at, + ) ?? + resetFromRelativeSeconds( + firstNumber(config, ['reset_in', 'resets_in_seconds']), + now, + ) + ); +} + +/** + * Parse the unofficial Grok billing payload. Shape is not documented; accept + * the live cli-chat-proxy `config` form plus earlier included-usage/credits + * guesses. Exported for unit tests. + */ +export function parseXaiSubscriptionUsage( + payload: unknown, + now: number = Date.now(), +): SubscriptionUsageWindow[] { + const root = asRecord(payload); + if (!root) { + return []; + } + + const config = resolveXaiBillingConfig(payload); + if (!config) { + return []; + } + + const windows: SubscriptionUsageWindow[] = []; + const resetsAt = resolveXaiResetsAt(config, now); + + // Live format=credits: paid Grok plans share one weekly pool across all + // Grok products, so only the aggregate percent is shown. The payload's + // per-product breakdown slices the same pool and would render as duplicate + // bars whenever one product dominates (an API-only deployment shows + // "Api: N%" identical to the aggregate); the breakdown stays on grok.com's + // own usage page. + const creditUsagePercent = firstNumber(config, [ + 'creditUsagePercent', + 'credit_usage_percent', + 'used_percent', + 'usedPercent', + 'included_used_percent', + 'includedUsedPercent', + ]); + + if (creditUsagePercent !== undefined) { + windows.push({ + label: 'Included usage', + usedPercent: clampPercent(creditUsagePercent), + ...(resetsAt && { resetsAt }), + }); + } + + // Live monthly billing (`/v1/billing` without format=credits). + const monthlyLimit = firstNumber(config, [ + 'monthlyLimit', + 'monthly_limit', + 'limit', + 'allowance', + 'total', + 'entitlement', + ]); + const monthlyUsed = firstNumber(config, ['used', 'consumed', 'includedUsed']); + if ( + windows.length === 0 && + (monthlyLimit !== undefined || monthlyUsed !== undefined) + ) { + windows.push({ + label: 'Included usage', + ...(monthlyUsed !== undefined && { used: monthlyUsed }), + ...(monthlyLimit !== undefined && { limit: monthlyLimit }), + ...(monthlyUsed !== undefined && + monthlyLimit !== undefined && + monthlyLimit > 0 && { + usedPercent: clampPercent((monthlyUsed / monthlyLimit) * 100), + }), + ...(resetsAt && { resetsAt }), + }); + } + + // Legacy / synthetic nested shapes (kept for unit tests and older proxies). + const included = asRecord( + config.included_usage ?? + config.includedUsage ?? + config.included ?? + root.included_usage ?? + root.includedUsage ?? + root.included, + ); + if (windows.length === 0 && included) { + const includedPercent = firstNumber(included, [ + 'used_percent', + 'usedPercent', + 'percent_used', + 'percentUsed', + ]); + const includedLimit = firstNumber(included, [ + 'limit', + 'allowance', + 'total', + 'entitlement', + ]); + const includedRemaining = firstNumber(included, [ + 'remaining', + 'remaining_credits', + 'remainingCredits', + ]); + const includedUsed = firstNumber(included, ['used', 'consumed']); + const includedResetsAt = + toResetIso( + included['resets_at'] ?? included['reset_at'] ?? included['resetTime'], + ) ?? resetsAt; + + if ( + includedPercent !== undefined || + includedLimit !== undefined || + includedRemaining !== undefined || + includedUsed !== undefined + ) { + windows.push({ + label: 'Included usage', + ...(includedPercent !== undefined && { + usedPercent: clampPercent(includedPercent), + }), + ...(includedUsed !== undefined && { used: includedUsed }), + ...(includedRemaining !== undefined && { + remaining: includedRemaining, + }), + ...(includedLimit !== undefined && { limit: includedLimit }), + ...(includedResetsAt && { resetsAt: includedResetsAt }), + }); + } + } + + const credits = asRecord(config.credits ?? config.credit ?? root.credits); + const creditBalance = + firstNumber(credits, [ + 'balance', + 'remaining', + 'available', + 'prepaid_balance', + 'prepaidBalance', + ]) ?? + firstNumber(config, ['prepaidBalance', 'prepaid_balance', 'creditBalance']); + // On-demand credits are the overflow path once the included pool is + // exhausted: a positive balance means tasks keep running on metered spend, + // which operators need to see. A zero balance is the common idle state and + // reads as an error ("Credits: 0 left"), so omit it like On-demand below. + if (creditBalance !== undefined && creditBalance > 0) { + windows.push({ + label: 'Credits', + remaining: creditBalance, + }); + } + + const onDemand = asRecord( + config.on_demand ?? + config.onDemand ?? + config.ondemand ?? + root.on_demand ?? + root.onDemand, + ); + const onDemandUsed = + firstNumber(onDemand, ['used', 'consumed']) ?? + firstNumber(config, ['onDemandUsed', 'on_demand_used']); + const onDemandLimit = + firstNumber(onDemand, ['limit', 'cap', 'allowance']) ?? + firstNumber(config, ['onDemandCap', 'on_demand_cap']); + if ( + (onDemandUsed !== undefined && onDemandUsed > 0) || + (onDemandLimit !== undefined && onDemandLimit > 0) + ) { + windows.push({ + label: 'On-demand', + ...(onDemandUsed !== undefined && { used: onDemandUsed }), + ...(onDemandLimit !== undefined && { limit: onDemandLimit }), + ...(onDemandUsed !== undefined && + onDemandLimit !== undefined && + onDemandLimit > 0 && { + usedPercent: clampPercent((onDemandUsed / onDemandLimit) * 100), + }), + }); + } + + return windows; +} + +export async function fetchXaiSubscriptionUsage( + options: UsageFetchOptions = {}, +): Promise { + const fetchImpl = options.fetchImpl ?? fetch; + const token = await getFreshXaiAccessToken({ + ...(options.executor && { executor: options.executor }), + fetchImpl, + }); + + if (!token) { + return null; + } + + const authHeaders = { + authorization: `Bearer ${token.access}`, + accept: 'application/json', + 'user-agent': 'roomote', + }; + + // Identity-first: resolve user, then request billing for that session. + // Fail closed without surfacing identity or raw bodies to callers. + const userResult = await fetchJson( + fetchImpl, + XAI_USAGE_USER_ENDPOINT, + authHeaders, + ); + if (userResult.status !== 200 || !asRecord(userResult.payload)?.userId) { + // Some deployments return `id` instead of `userId`. + const user = asRecord(userResult.payload); + if ( + userResult.status !== 200 || + (!firstString(user, ['userId', 'user_id', 'id']) && + !firstNumber(user, ['userId', 'id'])) + ) { + return null; + } + } + + const billingResult = await fetchJson( + fetchImpl, + XAI_USAGE_BILLING_ENDPOINT, + authHeaders, + ); + const windows = parseXaiSubscriptionUsage(billingResult.payload, Date.now()); + + if (windows.length === 0) { + return null; + } + + return { + providerId: 'xai-subscription', + windows, + fetchedAt: new Date().toISOString(), + }; +} + // --- Z.AI / Z.AI Coding Plan ---------------------------------------------- /** @@ -762,6 +1053,7 @@ export async function getSubscriptionProviderUsage( fetchChatGptUsage(options), fetchGitHubCopilotUsage(options), fetchKimiForCodingUsage(options), + fetchXaiSubscriptionUsage(options), fetchZaiUsage(options), fetchZaiCodingPlanUsage(options), ]); diff --git a/packages/db/src/lib/xai-subscription.ts b/packages/db/src/lib/xai-subscription.ts new file mode 100644 index 000000000..e4893372b --- /dev/null +++ b/packages/db/src/lib/xai-subscription.ts @@ -0,0 +1,863 @@ +import { randomUUID } from 'node:crypto'; + +import { sql } from 'drizzle-orm'; + +import { + DEFAULT_XAI_OAUTH_CLIENT_ID, + XAI_DEFAULT_ACCESS_TOKEN_TTL_MS, + XAI_OAUTH_CLIENT_ID_ENV_VAR_NAME, + XAI_OAUTH_DEVICE_CODE_ENDPOINT, + XAI_OAUTH_DEVICE_VERIFICATION_URL, + XAI_OAUTH_SCOPE, + XAI_OAUTH_TOKEN_ENDPOINT, + XAI_POLL_SLOW_DOWN_MS, + XAI_REFRESH_SAFETY_MARGIN_MS, +} from '@roomote/types'; + +import { type DatabaseOrTransaction, db } from '../db'; +import { decryptSecrets, encryptJSON } from '../encryption'; +import { deploymentSecrets } from '../schema'; + +/** + * Server-side lifecycle for the deployment-wide xAI Grok subscription OAuth + * credential. Roomote owns token refresh; the inference gateway mints a fresh + * access token per request so the OAuth record never enters sandboxes. + */ + +const XAI_SUBSCRIPTION_SECRET_NAME = 'XAI_SUBSCRIPTION_OAUTH'; +const XAI_SUBSCRIPTION_PENDING_SECRET_NAME = 'XAI_SUBSCRIPTION_PENDING_DEVICE'; +const XAI_SUBSCRIPTION_ADVISORY_LOCK_KEY = 'roomote-xai-subscription-refresh'; +const XAI_SUBSCRIPTION_DEVICE_LOCK_KEY = 'roomote-xai-subscription-device'; + +/** + * Tracks the single deployment-wide device-code flow that is allowed to + * persist tokens. A newer start supersedes older codes so an in-flight poll + * from a cancelled/restarted dialog cannot overwrite the active connection. + * + * `claimId` is reserved before the device-code HTTP call so a slower earlier + * start cannot register after a newer start. `deviceCode` is set only after + * that HTTP call completes while the claim is still current. + */ +interface XaiPendingDeviceFlow { + claimId: string; + /** Set after the device-code endpoint returns for the active claim. */ + deviceCode?: string; + /** Epoch milliseconds when the device code expires. */ + expiresAt: number; + startedAt: string; +} + +const SUPERSEDED_DEVICE_FLOW_ERROR = + 'This xAI device authorization is no longer active. Restart the connection.'; + +export type XaiSubscriptionStatusValue = 'connected' | 'error' | 'disconnected'; + +export interface XaiSubscriptionRecord { + refresh: string; + access: string; + /** Epoch milliseconds when the access token expires. */ + expires: number; + email?: string; + status: XaiSubscriptionStatusValue; + /** Last refresh/connect error message, populated when status is 'error'. */ + error?: string; + connectedAt: string; + updatedAt: string; +} + +/** + * Public, token-free status for the Settings UI. Never includes `refresh` or + * `access`. + */ +export interface XaiSubscriptionPublicStatus { + connected: boolean; + status: XaiSubscriptionStatusValue; + email?: string; + error?: string; + connectedAt?: string; + updatedAt?: string; +} + +interface TokenResponse { + access_token: string; + refresh_token?: string; + expires_in?: number; + id_token?: string; + error?: string; + error_description?: string; + interval?: number; +} + +interface DeviceCodeResponse { + device_code: string; + user_code: string; + verification_uri?: string; + verification_uri_complete?: string; + interval?: number; + expires_in?: number; +} + +interface IdTokenClaims { + email?: string; +} + +export function resolveXaiOAuthClientId( + runtimeEnv: Partial> = process.env, +): string { + return ( + runtimeEnv[XAI_OAUTH_CLIENT_ID_ENV_VAR_NAME]?.trim() || + DEFAULT_XAI_OAUTH_CLIENT_ID + ); +} + +function parseJwtClaims(token: string): IdTokenClaims | undefined { + const parts = token.split('.'); + if (parts.length !== 3) { + return undefined; + } + try { + return JSON.parse( + Buffer.from(parts[1]!, 'base64url').toString('utf-8'), + ) as IdTokenClaims; + } catch { + return undefined; + } +} + +function extractEmailFromTokens(tokens: TokenResponse): string | undefined { + if (tokens.id_token) { + const claims = parseJwtClaims(tokens.id_token); + if (claims?.email) { + return claims.email; + } + } + return undefined; +} + +function resolveAccessTokenExpires(tokens: TokenResponse, now: number): number { + const ttlMs = + tokens.expires_in && tokens.expires_in > 0 + ? tokens.expires_in * 1000 + : XAI_DEFAULT_ACCESS_TOKEN_TTL_MS; + return now + ttlMs; +} + +async function loadSecretValue( + executor: DatabaseOrTransaction, + secretName: string, +): Promise { + const [row] = await executor + .select({ value: deploymentSecrets.value }) + .from(deploymentSecrets) + .where(sql`${deploymentSecrets.name} = ${secretName}`) + .limit(1); + + if (!row) { + return null; + } + + return decryptSecrets(row.value); +} + +async function persistSecretValue( + executor: DatabaseOrTransaction, + secretName: string, + value: unknown, +): Promise { + const encrypted = encryptJSON(value); + + await executor + .insert(deploymentSecrets) + .values({ + name: secretName, + value: encrypted, + }) + .onConflictDoUpdate({ + target: deploymentSecrets.name, + set: { + value: encrypted, + updatedAt: new Date(), + }, + }); +} + +async function deleteSecretValue( + executor: DatabaseOrTransaction, + secretName: string, +): Promise { + await executor + .delete(deploymentSecrets) + .where(sql`${deploymentSecrets.name} = ${secretName}`); +} + +async function loadRecord( + executor: DatabaseOrTransaction = db, +): Promise { + return loadSecretValue( + executor, + XAI_SUBSCRIPTION_SECRET_NAME, + ); +} + +async function persistRecord( + executor: DatabaseOrTransaction, + record: XaiSubscriptionRecord, +): Promise { + await persistSecretValue(executor, XAI_SUBSCRIPTION_SECRET_NAME, record); +} + +async function loadPendingDeviceFlow( + executor: DatabaseOrTransaction = db, +): Promise { + return loadSecretValue( + executor, + XAI_SUBSCRIPTION_PENDING_SECRET_NAME, + ); +} + +async function persistPendingDeviceFlow( + executor: DatabaseOrTransaction, + pending: XaiPendingDeviceFlow, +): Promise { + await persistSecretValue( + executor, + XAI_SUBSCRIPTION_PENDING_SECRET_NAME, + pending, + ); +} + +async function clearPendingDeviceFlow( + executor: DatabaseOrTransaction, +): Promise { + await deleteSecretValue(executor, XAI_SUBSCRIPTION_PENDING_SECRET_NAME); +} + +function isActivePendingDeviceFlow( + pending: XaiPendingDeviceFlow | null, + deviceCode: string, + now: number = Date.now(), +): boolean { + return Boolean( + pending && + pending.deviceCode && + pending.deviceCode === deviceCode && + pending.expiresAt > now, + ); +} + +type XaiLockKind = 'device' | 'refresh' | 'both'; + +/** + * Serialize subscription mutations. Unit-test executors may omit + * `transaction`; fall back to running the body directly on the executor. + * + * Lock order is always refresh then device when both are needed, to avoid + * deadlocks between disconnect, refresh, and device-code connect. + */ +async function withXaiSubscriptionLock( + executor: DatabaseOrTransaction, + kind: XaiLockKind, + body: (tx: DatabaseOrTransaction) => Promise, +): Promise { + const maybeTx = executor as DatabaseOrTransaction & { + transaction?: ( + fn: (tx: DatabaseOrTransaction) => Promise, + ) => Promise; + execute?: (query: unknown) => Promise; + }; + + if (typeof maybeTx.transaction === 'function') { + await maybeTx.transaction(async (tx) => { + if (typeof tx.execute === 'function') { + if (kind === 'refresh' || kind === 'both') { + await tx.execute( + sql`SELECT pg_advisory_xact_lock(hashtext(${XAI_SUBSCRIPTION_ADVISORY_LOCK_KEY}))`, + ); + } + if (kind === 'device' || kind === 'both') { + await tx.execute( + sql`SELECT pg_advisory_xact_lock(hashtext(${XAI_SUBSCRIPTION_DEVICE_LOCK_KEY}))`, + ); + } + } + await body(tx); + }); + return; + } + + await body(executor); +} + +async function withDeviceAuthLock( + executor: DatabaseOrTransaction, + body: (tx: DatabaseOrTransaction) => Promise, +): Promise { + await withXaiSubscriptionLock(executor, 'device', body); +} + +/** + * Terminal IdP failures should mark the subscription `error` and force + * reconnect. Transient failures (5xx / network) keep `connected` so a brief + * outage does not brick OAuth-only deploys while access may still be valid. + */ +function isTerminalXaiRefreshFailure(error: unknown): boolean { + const message = error instanceof Error ? error.message : String(error ?? ''); + if (/invalid_grant/i.test(message)) { + return true; + } + // 401/403 are auth failures; 400 with invalid_grant handled above. + if (/xAI token refresh failed: (401|403)\b/.test(message)) { + return true; + } + if (/no access_token/i.test(message)) { + return true; + } + return false; +} + +export async function getXaiSubscription( + executor: DatabaseOrTransaction = db, +): Promise { + return loadRecord(executor); +} + +export async function isXaiSubscriptionConnected( + executor: DatabaseOrTransaction = db, +): Promise { + const record = await loadRecord(executor); + return Boolean(record && record.status === 'connected'); +} + +export async function getXaiSubscriptionStatus( + executor: DatabaseOrTransaction = db, +): Promise { + const record = await loadRecord(executor); + + if (!record) { + return { connected: false, status: 'disconnected' }; + } + + return { + connected: record.status === 'connected', + status: record.status, + ...(record.email && { email: record.email }), + ...(record.error && { error: record.error }), + ...(record.connectedAt && { connectedAt: record.connectedAt }), + ...(record.updatedAt && { updatedAt: record.updatedAt }), + }; +} + +export async function saveXaiSubscription( + tokens: TokenResponse, + executor: DatabaseOrTransaction = db, + existing?: XaiSubscriptionRecord | null, +): Promise { + const now = Date.now(); + const email = extractEmailFromTokens(tokens) ?? existing?.email; + const refresh = tokens.refresh_token ?? existing?.refresh; + + if (!refresh) { + throw new Error('xAI token response did not include a refresh_token.'); + } + + const record: XaiSubscriptionRecord = { + refresh, + access: tokens.access_token, + expires: resolveAccessTokenExpires(tokens, now), + ...(email && { email }), + status: 'connected', + connectedAt: existing?.connectedAt ?? new Date(now).toISOString(), + updatedAt: new Date(now).toISOString(), + }; + + await persistRecord(executor, record); + return record; +} + +export async function disconnectXaiSubscription( + executor: DatabaseOrTransaction = db, +): Promise { + // Take both locks so neither an in-flight device poll nor a concurrent + // token refresh can re-persist credentials after the operator disconnects. + await withXaiSubscriptionLock(executor, 'both', async (tx) => { + await deleteSecretValue(tx, XAI_SUBSCRIPTION_SECRET_NAME); + await clearPendingDeviceFlow(tx); + }); +} + +export type FreshXaiAccessToken = { + access: string; + expires: number; + email?: string; +}; + +function toPublicAccessToken( + record: XaiSubscriptionRecord, +): FreshXaiAccessToken { + return { + access: record.access, + expires: record.expires, + ...(record.email && { email: record.email }), + }; +} + +/** + * Refresh the access token when it is expiring within + * `XAI_REFRESH_SAFETY_MARGIN_MS`. Lock critical sections only around reads and + * writes; IdP HTTP runs outside the lock. CAS on the original refresh token + * so a concurrent reconnect cannot be overwritten. On terminal failure, marks + * the record `status: 'error'` instead of deleting it. + * + * Returns access credentials only (never the refresh token) so callers cannot + * accidentally log or forward refresh material. + */ +export async function getFreshXaiAccessToken( + options: { + executor?: DatabaseOrTransaction; + fetchImpl?: typeof fetch; + now?: number; + } = {}, +): Promise { + const executor = options.executor ?? db; + const fetchImpl = options.fetchImpl ?? fetch; + const now = options.now ?? Date.now(); + const existing = await loadRecord(executor); + + if (!existing || existing.status !== 'connected') { + return null; + } + + if (existing.expires > now + XAI_REFRESH_SAFETY_MARGIN_MS) { + return toPublicAccessToken(existing); + } + + // Snapshot under lock, then refresh outside so disconnect/device-start are + // not blocked on IdP latency. Box the mutable state so TypeScript does not + // treat closure assignments as non-narrowing dead state. + type RefreshSnapshot = { + refresh: string; + access: string; + expires: number; + email?: string; + connectedAt: string; + }; + const lockPhase: { + refreshSnapshot: RefreshSnapshot | null; + alreadyFresh: XaiSubscriptionRecord | null; + } = { + refreshSnapshot: null, + alreadyFresh: null, + }; + + await withXaiSubscriptionLock(executor, 'both', async (tx) => { + const locked = await loadRecord(tx); + + if (!locked || locked.status !== 'connected') { + return; + } + + if (locked.expires > now + XAI_REFRESH_SAFETY_MARGIN_MS) { + lockPhase.alreadyFresh = locked; + return; + } + + lockPhase.refreshSnapshot = { + refresh: locked.refresh, + access: locked.access, + expires: locked.expires, + ...(locked.email && { email: locked.email }), + connectedAt: locked.connectedAt, + }; + }); + + if (lockPhase.alreadyFresh) { + return toPublicAccessToken(lockPhase.alreadyFresh); + } + + const snapshot = lockPhase.refreshSnapshot; + if (!snapshot) { + return null; + } + + type RefreshHttpResult = + | { ok: true; tokens: TokenResponse } + | { ok: false; error: Error }; + + let httpResult: RefreshHttpResult; + try { + const clientId = resolveXaiOAuthClientId(); + const response = await fetchImpl(XAI_OAUTH_TOKEN_ENDPOINT, { + method: 'POST', + headers: { + 'Content-Type': 'application/x-www-form-urlencoded', + Accept: 'application/json', + 'User-Agent': 'roomote', + }, + body: new URLSearchParams({ + grant_type: 'refresh_token', + refresh_token: snapshot.refresh, + client_id: clientId, + }).toString(), + }); + + let bodyText = ''; + try { + bodyText = await response.text(); + } catch { + bodyText = ''; + } + + let tokens: Partial = {}; + if (bodyText) { + try { + tokens = JSON.parse(bodyText) as Partial; + } catch { + tokens = {}; + } + } + + if (!response.ok) { + const detail = + tokens.error_description ?? + tokens.error ?? + `xAI token refresh failed: ${response.status}`; + throw new Error( + tokens.error + ? `xAI token refresh failed: ${response.status} ${tokens.error}` + : detail.startsWith('xAI token refresh failed:') + ? detail + : `xAI token refresh failed: ${response.status}`, + ); + } + + if (!tokens.access_token) { + throw new Error('xAI token refresh returned no access_token.'); + } + + httpResult = { ok: true, tokens: tokens as TokenResponse }; + } catch (error) { + httpResult = { + ok: false, + error: error instanceof Error ? error : new Error(String(error)), + }; + } + + const writePhase: { record: XaiSubscriptionRecord | null } = { + record: null, + }; + + await withXaiSubscriptionLock(executor, 'both', async (tx) => { + const current = await loadRecord(tx); + + if ( + !current || + current.status !== 'connected' || + current.refresh !== snapshot.refresh + ) { + // Replaced by a newer connect/disconnect while we awaited xAI. + writePhase.record = current?.status === 'connected' ? current : null; + return; + } + + if (httpResult.ok) { + writePhase.record = await saveXaiSubscription( + httpResult.tokens, + tx, + current, + ); + return; + } + + const message = httpResult.error.message; + + if (isTerminalXaiRefreshFailure(httpResult.error)) { + const errored: XaiSubscriptionRecord = { + ...current, + status: 'error', + error: message, + updatedAt: new Date().toISOString(), + }; + await persistRecord(tx, errored); + return; + } + + // Transient failure: keep status connected. Serve remaining access if + // it has not expired yet so OAuth-only deploys survive brief IdP blips. + if (current.expires > now) { + writePhase.record = current; + } + }); + + if (!writePhase.record || writePhase.record.status !== 'connected') { + return null; + } + + return toPublicAccessToken(writePhase.record); +} + +/** + * Initiate the xAI device-code authorization flow. Returns the user code and + * verification URL the operator enters in a browser. The caller polls + * {@link pollXaiDeviceAuth} with the returned `deviceCode` until it resolves. + * + * Reserves a claim before the device-code HTTP request, then registers the + * returned code only if that claim is still current. A slower earlier start + * cannot overwrite a newer dialog's pending flow. + */ +export async function startXaiDeviceAuth( + fetchImpl: typeof fetch = fetch, + options: { + executor?: DatabaseOrTransaction; + now?: number; + } = {}, +): Promise<{ + deviceCode: string; + userCode: string; + verificationUrl: string; + intervalMs: number; + expiresInMs: number; +}> { + const executor = options.executor ?? db; + const now = options.now ?? Date.now(); + const clientId = resolveXaiOAuthClientId(); + const claimId = randomUUID(); + // Provisional window until the device-code endpoint returns real expiry. + const provisionalExpiresAt = now + 15 * 60 * 1000; + + // Reserve the active claim before any network I/O so concurrent starts order + // by claim time, not by which HTTP response finishes first. + await withDeviceAuthLock(executor, async (tx) => { + await persistPendingDeviceFlow(tx, { + claimId, + expiresAt: provisionalExpiresAt, + startedAt: new Date(now).toISOString(), + }); + }); + + const response = await fetchImpl(XAI_OAUTH_DEVICE_CODE_ENDPOINT, { + method: 'POST', + headers: { + 'Content-Type': 'application/x-www-form-urlencoded', + Accept: 'application/json', + 'User-Agent': 'roomote', + }, + body: new URLSearchParams({ + client_id: clientId, + scope: XAI_OAUTH_SCOPE, + }).toString(), + }); + + if (!response.ok) { + // Drop this claim if it is still active so a failed start does not block + // a later retry indefinitely under the provisional expiry. + await withDeviceAuthLock(executor, async (tx) => { + const pending = await loadPendingDeviceFlow(tx); + if (pending?.claimId === claimId) { + await clearPendingDeviceFlow(tx); + } + }); + throw new Error( + `Failed to initiate xAI device authorization: ${response.status}`, + ); + } + + const data = (await response.json()) as DeviceCodeResponse; + const expiresInMs = Math.max(data.expires_in ?? 900, 1) * 1000; + const result = { + deviceCode: data.device_code, + userCode: data.user_code, + verificationUrl: + data.verification_uri_complete ?? + data.verification_uri ?? + XAI_OAUTH_DEVICE_VERIFICATION_URL, + intervalMs: Math.max(data.interval ?? 5, 1) * 1000, + expiresInMs, + }; + + let registered = false; + await withDeviceAuthLock(executor, async (tx) => { + const pending = await loadPendingDeviceFlow(tx); + if (!pending || pending.claimId !== claimId) { + // A newer start claimed the flow while we awaited xAI. + return; + } + await persistPendingDeviceFlow(tx, { + claimId, + deviceCode: result.deviceCode, + expiresAt: now + expiresInMs, + startedAt: pending.startedAt, + }); + registered = true; + }); + + if (!registered) { + throw new Error(SUPERSEDED_DEVICE_FLOW_ERROR); + } + + return result; +} + +export type XaiDevicePollResult = + | { status: 'pending'; intervalMs?: number } + | { status: 'success'; record: XaiSubscriptionRecord } + | { status: 'failed'; error: string }; + +/** + * Poll the xAI token endpoint once for a device-code grant. Returns `pending` + * while the operator has not finished authorization, `success` with the + * persisted record once tokens are exchanged, or `failed` on a terminal error. + * + * Token persistence is gated on the deployment's currently active pending + * device code (set by {@link startXaiDeviceAuth}). A poll for a superseded + * code never writes subscription tokens, even if xAI still returns them. + */ +export async function pollXaiDeviceAuth( + input: { deviceCode: string }, + options: { + executor?: DatabaseOrTransaction; + fetchImpl?: typeof fetch; + now?: number; + } = {}, +): Promise { + const executor = options.executor ?? db; + const fetchImpl = options.fetchImpl ?? fetch; + const now = options.now ?? Date.now(); + const clientId = resolveXaiOAuthClientId(); + + // Cheap reject before talking to xAI when this code is already superseded. + const pendingBefore = await loadPendingDeviceFlow(executor); + if (!isActivePendingDeviceFlow(pendingBefore, input.deviceCode, now)) { + return { + status: 'failed', + error: SUPERSEDED_DEVICE_FLOW_ERROR, + }; + } + + const response = await fetchImpl(XAI_OAUTH_TOKEN_ENDPOINT, { + method: 'POST', + headers: { + 'Content-Type': 'application/x-www-form-urlencoded', + Accept: 'application/json', + 'User-Agent': 'roomote', + }, + body: new URLSearchParams({ + grant_type: 'urn:ietf:params:oauth:grant-type:device_code', + client_id: clientId, + device_code: input.deviceCode, + }).toString(), + }); + + let data: TokenResponse; + try { + data = (await response.json()) as TokenResponse; + } catch { + return { + status: 'failed', + error: `xAI device authorization failed: ${response.status}`, + }; + } + + if (data.access_token) { + if (!data.refresh_token) { + return { + status: 'failed', + error: 'xAI device authorization returned no refresh_token.', + }; + } + + // Re-check under both locks: cancel/reopen may have registered a newer + // code, and a concurrent refresh of the prior credential must not write + // after this persist without seeing the new record. + let result: XaiDevicePollResult = { + status: 'failed', + error: SUPERSEDED_DEVICE_FLOW_ERROR, + }; + + await withXaiSubscriptionLock(executor, 'both', async (tx) => { + const pending = await loadPendingDeviceFlow(tx); + if (!isActivePendingDeviceFlow(pending, input.deviceCode, now)) { + return; + } + + const record = await saveXaiSubscription(data, tx); + await clearPendingDeviceFlow(tx); + result = { status: 'success', record }; + }); + + return result; + } + + // Terminal failures for a superseded code should still fail closed without + // clearing a newer pending flow. + const pendingAfter = await loadPendingDeviceFlow(executor); + if (!isActivePendingDeviceFlow(pendingAfter, input.deviceCode, now)) { + return { + status: 'failed', + error: SUPERSEDED_DEVICE_FLOW_ERROR, + }; + } + + if (data.error === 'authorization_pending') { + return { status: 'pending' }; + } + + if (data.error === 'slow_down') { + return { + status: 'pending', + intervalMs: + typeof data.interval === 'number' && data.interval > 0 + ? data.interval * 1000 + : XAI_POLL_SLOW_DOWN_MS, + }; + } + + if (data.error === 'expired_token') { + await withDeviceAuthLock(executor, async (tx) => { + const pending = await loadPendingDeviceFlow(tx); + if (isActivePendingDeviceFlow(pending, input.deviceCode, now)) { + await clearPendingDeviceFlow(tx); + } + }); + return { + status: 'failed', + error: 'xAI device authorization code expired. Restart the connection.', + }; + } + + if (data.error === 'access_denied') { + await withDeviceAuthLock(executor, async (tx) => { + const pending = await loadPendingDeviceFlow(tx); + if (isActivePendingDeviceFlow(pending, input.deviceCode, now)) { + await clearPendingDeviceFlow(tx); + } + }); + return { + status: 'failed', + error: 'xAI device authorization was denied.', + }; + } + + if (!response.ok && !data.error) { + return { + status: 'failed', + error: `xAI device authorization failed: ${response.status}`, + }; + } + + return { + status: 'failed', + error: + data.error_description ?? + (data.error + ? `xAI device authorization failed: ${data.error}` + : 'xAI device authorization returned no access token.'), + }; +} + +export const XAI_SUBSCRIPTION_INTERNAL = { + secretName: XAI_SUBSCRIPTION_SECRET_NAME, + pendingSecretName: XAI_SUBSCRIPTION_PENDING_SECRET_NAME, + defaultAccessTokenTtlMs: XAI_DEFAULT_ACCESS_TOKEN_TTL_MS, + tokenEndpoint: XAI_OAUTH_TOKEN_ENDPOINT, + supersededDeviceFlowError: SUPERSEDED_DEVICE_FLOW_ERROR, +}; diff --git a/packages/db/src/server.ts b/packages/db/src/server.ts index ddd77bc0a..689bcd57a 100644 --- a/packages/db/src/server.ts +++ b/packages/db/src/server.ts @@ -64,6 +64,7 @@ export * from './lib/compute-runtime-config'; export * from './lib/model-runtime-config'; export * from './lib/chatgpt-subscription'; export * from './lib/github-copilot-subscription'; +export * from './lib/xai-subscription'; export * from './lib/subscription-provider-usage'; export * from './lib/provider-credit-balance'; export * from './lib/preview-runtime-config'; diff --git a/packages/types/src/__tests__/inference-gateway.test.ts b/packages/types/src/__tests__/inference-gateway.test.ts index 64cdb88b3..3cfe84f43 100644 --- a/packages/types/src/__tests__/inference-gateway.test.ts +++ b/packages/types/src/__tests__/inference-gateway.test.ts @@ -188,6 +188,19 @@ describe('inference gateway key lookups', () => { ).toBe('https://api.example.com/api/inference/github-copilot'); }); + it('registers xAI with the hybrid xai-oauth strategy and API host', () => { + const provider = getInferenceGatewayProvider('xai'); + expect(provider).toMatchObject({ + authStrategy: 'xai-oauth', + envVarNames: ['XAI_API_KEY'], + upstreamBaseUrl: 'https://api.x.ai', + openCodeBaseUrlSuffix: '/v1', + }); + expect(provider?.allowedPaths).toEqual( + expect.arrayContaining(['/v1/chat/completions', '/v1/responses']), + ); + }); + it('offers exactly the regions its gateway providers hold base URLs for', () => { const regionProviders = INFERENCE_GATEWAY_PROVIDERS.filter( (provider) => provider.region?.baseUrls, diff --git a/packages/types/src/index.ts b/packages/types/src/index.ts index 8b0fb1eee..2aadcabc6 100644 --- a/packages/types/src/index.ts +++ b/packages/types/src/index.ts @@ -11,6 +11,7 @@ export * from './pr-review-action'; export * from './task-runs'; export * from './chatgpt-subscription'; export * from './github-copilot-subscription'; +export * from './xai-subscription'; export * from './communication'; export * from './communication-message-prompt'; export * from './identity-display-name'; diff --git a/packages/types/src/inference-gateway.ts b/packages/types/src/inference-gateway.ts index dc159a501..ef815ba9f 100644 --- a/packages/types/src/inference-gateway.ts +++ b/packages/types/src/inference-gateway.ts @@ -66,6 +66,12 @@ export const INFERENCE_GATEWAY_CHATGPT_ENV_VAR_NAME = export const INFERENCE_GATEWAY_GITHUB_COPILOT_ENV_VAR_NAME = 'R_INFERENCE_GATEWAY_GITHUB_COPILOT'; +/** + * Marker indicating that the gateway holds an xAI Grok subscription OAuth + * token (and should serve `xai/` models even when `XAI_API_KEY` is absent). + */ +export const INFERENCE_GATEWAY_XAI_ENV_VAR_NAME = 'R_INFERENCE_GATEWAY_XAI'; + interface InferenceGatewayAuthHeader { name: string; scheme?: 'bearer'; @@ -76,11 +82,14 @@ interface InferenceGatewayAuthHeader { * - `api-key`: inject a static deployment key (the default). * - `chatgpt-oauth`: mint a fresh ChatGPT-subscription access token * server-side and add the account-id header; the sandbox holds no key. + * - `xai-oauth`: prefer a connected Grok subscription access token, falling + * back to `XAI_API_KEY` when no subscription is present. */ export type InferenceGatewayAuthStrategy = | 'api-key' | 'chatgpt-oauth' - | 'github-copilot-oauth'; + | 'github-copilot-oauth' + | 'xai-oauth'; export interface InferenceGatewayProvider { /** @@ -322,6 +331,9 @@ export const INFERENCE_GATEWAY_PROVIDERS: readonly InferenceGatewayProvider[] = { id: 'xai', name: 'xAI', + // Prefer SuperGrok / Premium+ OAuth when connected; fall back to the + // deployment API key so BYOK setups keep working unchanged. + authStrategy: 'xai-oauth', envVarNames: ['XAI_API_KEY'], upstreamBaseUrl: 'https://api.x.ai', authHeader: { name: 'authorization', scheme: 'bearer' }, diff --git a/packages/types/src/model-provider-config.test.ts b/packages/types/src/model-provider-config.test.ts index 4846a0959..eb041c71f 100644 --- a/packages/types/src/model-provider-config.test.ts +++ b/packages/types/src/model-provider-config.test.ts @@ -238,6 +238,7 @@ describe('SETUP_MODEL_PROVIDER_CATALOG', () => { 'ollama', 'vllm', 'chatgpt', + 'xai-subscription', ], ); }); @@ -249,14 +250,16 @@ describe('SETUP_MODEL_PROVIDER_CATALOG', () => { expect(provider.suggestedTaskModels).toEqual([]); continue; } - // ChatGPT serves openai/ models; the Bedrock setup surface serves the - // worker's custom bedrock-mantle/ OpenCode provider. + // ChatGPT serves openai/ models; SuperGrok serves xai/; Bedrock uses + // the worker's custom bedrock-mantle/ OpenCode provider. const expectedPrefix = provider.id === 'chatgpt' ? 'openai/' - : provider.id === 'amazon-bedrock' - ? 'bedrock-mantle/' - : `${provider.id}/`; + : provider.id === 'xai-subscription' + ? 'xai/' + : provider.id === 'amazon-bedrock' + ? 'bedrock-mantle/' + : `${provider.id}/`; expect(provider.defaultRoomoteModel.startsWith(expectedPrefix)).toBe( true, @@ -282,9 +285,11 @@ describe('SETUP_MODEL_PROVIDER_CATALOG', () => { const expectedPrefix = provider.id === 'chatgpt' ? 'openai/' - : provider.id === 'amazon-bedrock' - ? 'bedrock-mantle/' - : `${provider.id}/`; + : provider.id === 'xai-subscription' + ? 'xai/' + : provider.id === 'amazon-bedrock' + ? 'bedrock-mantle/' + : `${provider.id}/`; const suggestedModelIds = new Set( provider.suggestedTaskModels.map((suggestion) => suggestion.id), ); @@ -453,6 +458,29 @@ describe('SETUP_MODEL_PROVIDER_CATALOG', () => { ).toBe('openai'); }); + it('groups xai/ models under the Grok subscription when it is the only xAI connection', () => { + expect( + getDisplayModelProviderId('xai/grok-4.5', { + xaiSubscriptionConnected: true, + }), + ).toBe('xai-subscription'); + expect(getDisplayModelProviderId('xai/grok-4.5', {})).toBe('xai'); + expect( + getDisplayModelProviderId('openai/gpt-5.6-terra', { + xaiSubscriptionConnected: true, + }), + ).toBe('openai'); + }); + + it('keeps xai/ models under xAI when an API key is also connected', () => { + expect( + getDisplayModelProviderId('xai/grok-4.5', { + xaiSubscriptionConnected: true, + xaiConnected: true, + }), + ).toBe('xai'); + }); + it('groups model chooser options by display provider and catalog order', () => { const groups = groupModelsByDisplayProvider( [ @@ -572,6 +600,28 @@ describe('SETUP_MODEL_PROVIDER_CATALOG', () => { expect(disconnected?.runtimeApiKeySatisfied).toBe(false); expect(disconnected?.savedApiKeySatisfied).toBe(false); }); + + it('marks xAI Grok subscription connected as its own OAuth provider without an API key', () => { + const status = buildSetupModelStatus({ + xaiSubscriptionConnected: true, + }); + const subscription = status.providers.find( + (provider) => provider.id === 'xai-subscription', + ); + const xaiKey = status.providers.find((provider) => provider.id === 'xai'); + expect(subscription?.savedApiKeySatisfied).toBe(true); + expect(xaiKey?.savedApiKeySatisfied).toBe(false); + expect(status.xaiSubscriptionConnected).toBe(true); + expect(status.xaiApiKeyConnected).toBe(false); + }); + + it('reports xaiApiKeyConnected when XAI_API_KEY is present', () => { + const status = buildSetupModelStatus({ + persistedEnvVarNames: ['XAI_API_KEY'], + }); + expect(status.xaiApiKeyConnected).toBe(true); + expect(status.xaiSubscriptionConnected).toBe(false); + }); }); describe('buildRecommendedDeploymentModelConfig', () => { diff --git a/packages/types/src/model-provider-config.ts b/packages/types/src/model-provider-config.ts index 17fb8fc46..341a9f431 100644 --- a/packages/types/src/model-provider-config.ts +++ b/packages/types/src/model-provider-config.ts @@ -35,10 +35,18 @@ import { */ export const CHATGPT_SUBSCRIPTION_PROVIDER_ID = 'chatgpt' as const; +/** + * Setup-catalog id for SuperGrok / eligible X Premium+ OAuth. Models keep the + * `xai/` id prefix (opencode's xAI provider), so this id is a + * configuration/connect surface only — not a model-id prefix. + */ +export const XAI_SUBSCRIPTION_PROVIDER_ID = 'xai-subscription' as const; + export const SETUP_MODEL_PROVIDER_IDS = [ 'openrouter', ...ENABLED_DIRECT_TASK_MODEL_PROVIDER_IDS, CHATGPT_SUBSCRIPTION_PROVIDER_ID, + XAI_SUBSCRIPTION_PROVIDER_ID, ] as const; /** @@ -510,12 +518,18 @@ export const SETUP_MODEL_PROVIDER_CATALOG = [ }, { // Provider id matches the models.dev/opencode `xai` provider so - // `xai/` slugs resolve at runtime. + // `xai/` slugs resolve at runtime. API-key path only; SuperGrok + // OAuth is a separate catalog entry (`xai-subscription`). id: 'xai', label: 'xAI', envVarName: 'XAI_API_KEY', defaultRoomoteModel: 'xai/grok-4.5', authKind: 'api-key', + credentialHelp: { + text: 'Paste an xAI API key from the console.', + href: 'https://console.x.ai', + linkLabel: 'xAI console', + }, suggestedTaskModels: mapRecommendedTaskModels({ 'grok-4-5': 'xai/grok-4.5', }), @@ -689,6 +703,16 @@ export const SETUP_MODEL_PROVIDER_CATALOG = [ planning: 'openai/gpt-5.6-sol', }, }, + { + id: XAI_SUBSCRIPTION_PROVIDER_ID, + label: 'xAI (Grok subscription)', + envVarName: undefined, + defaultRoomoteModel: 'xai/grok-4.5', + authKind: 'oauth', + suggestedTaskModels: mapRecommendedTaskModels({ + 'grok-4-5': 'xai/grok-4.5', + }), + }, ] as const satisfies readonly SetupModelProviderDescriptor[]; const EXTRA_MODEL_PROVIDER_ENV_KEYS_BY_PROVIDER = { @@ -966,6 +990,19 @@ export type SetupModelStatus = { chatgptConnected: boolean; /** Whether a GitHub Copilot subscription OAuth record is connected. */ githubCopilotConnected?: boolean; + /** + * Whether an xAI Grok subscription OAuth record is connected. The `xai` + * catalog entry still accepts `XAI_API_KEY`; this flag marks the provider + * connected when only the subscription is present (and lets the gateway + * prefer the OAuth bearer when both exist). + */ + xaiSubscriptionConnected?: boolean; + /** + * Whether an `XAI_API_KEY` is configured (runtime or saved), independent of + * the SuperGrok subscription. Used by Settings to show both paths without + * double-counting an OAuth-only connection as an API-key row. + */ + xaiApiKeyConnected?: boolean; /** * True when both `OPENAI_API_KEY` and a ChatGPT subscription are * configured. opencode's Codex plugin prefers OAuth auth when both are @@ -1145,6 +1182,8 @@ export function getDisplayModelProviderId( options?: { chatgptConnected?: boolean; openaiConnected?: boolean; + xaiSubscriptionConnected?: boolean; + xaiConnected?: boolean; }, ): string | null { const normalizedModelId = normalizeOptionalString(modelId); @@ -1163,6 +1202,17 @@ export function getDisplayModelProviderId( return CHATGPT_SUBSCRIPTION_PROVIDER_ID; } + // Mirror the ChatGPT/OpenAI relationship: `xai/` models group under the + // Grok subscription when it is the only xAI-facing connection, and stay + // under the native xAI group when an API key is also present. + if ( + runtimeProviderId === 'xai' && + options?.xaiSubscriptionConnected && + !options?.xaiConnected + ) { + return XAI_SUBSCRIPTION_PROVIDER_ID; + } + return runtimeProviderId; } @@ -1185,16 +1235,14 @@ export function groupModelsByDisplayProvider( options?: { chatgptConnected?: boolean; openaiConnected?: boolean; + xaiSubscriptionConnected?: boolean; + xaiConnected?: boolean; }, ): DisplayModelProviderGroup[] { const groups = new Map(); for (const item of items) { - const providerId = - getDisplayModelProviderId(item.id, { - chatgptConnected: options?.chatgptConnected, - openaiConnected: options?.openaiConnected, - }) ?? 'other'; + const providerId = getDisplayModelProviderId(item.id, options) ?? 'other'; const groupItems = groups.get(providerId) ?? []; groupItems.push(item); groups.set(providerId, groupItems); @@ -1463,10 +1511,12 @@ export function buildSetupModelStatus(input: { */ chatgptConnected?: boolean; githubCopilotConnected?: boolean; + xaiSubscriptionConnected?: boolean; }): SetupModelStatus { const runtimeEnv = input.runtimeEnv ?? {}; const chatgptConnected = Boolean(input.chatgptConnected); const githubCopilotConnected = Boolean(input.githubCopilotConnected); + const xaiSubscriptionConnected = Boolean(input.xaiSubscriptionConnected); const persistedModelConfig = normalizeDeploymentModelConfig( input.persistedModelConfig, ); @@ -1517,7 +1567,9 @@ export function buildSetupModelStatus(input: { ? chatgptConnected : provider.id === 'github-copilot' ? githubCopilotConnected - : false; + : provider.id === XAI_SUBSCRIPTION_PROVIDER_ID + ? xaiSubscriptionConnected + : false; return { ...provider, additionalEnvValues: {}, @@ -1563,17 +1615,20 @@ export function buildSetupModelStatus(input: { .filter((entry): entry is [string, string] => entry[1] !== null), ); + const runtimeApiKeySatisfied = + hasRequiredEnvVars && requiredEnvVarNames.every(isRuntimeConfigured); + const savedApiKeySatisfied = + hasRequiredEnvVars && + requiredEnvVarNames.every( + (name) => isPersisted(name) || isRuntimeConfigured(name), + ) && + requiredEnvVarNames.some(isPersisted); + return { ...provider, additionalEnvValues, - runtimeApiKeySatisfied: - hasRequiredEnvVars && requiredEnvVarNames.every(isRuntimeConfigured), - savedApiKeySatisfied: - hasRequiredEnvVars && - requiredEnvVarNames.every( - (name) => isPersisted(name) || isRuntimeConfigured(name), - ) && - requiredEnvVarNames.some(isPersisted), + runtimeApiKeySatisfied, + savedApiKeySatisfied, }; }; @@ -1655,27 +1710,36 @@ export function buildSetupModelStatus(input: { // A ChatGPT subscription covers `openai/` models, so it can satisfy the // runtime/persisted provider status for the OpenAI API-key provider even - // without OPENAI_API_KEY. + // without OPENAI_API_KEY. SuperGrok covers `xai/` the same way. const chatgptCoversOpenAi = chatgptConnected; + const xaiSubscriptionCoversXai = xaiSubscriptionConnected; const openaiRuntimeSatisfied = runtimeProviderStatus?.id === 'openai' ? runtimeProviderStatus.runtimeApiKeySatisfied || chatgptCoversOpenAi - : (runtimeProviderStatus?.runtimeApiKeySatisfied ?? false); + : runtimeProviderStatus?.id === 'xai' + ? runtimeProviderStatus.runtimeApiKeySatisfied || + xaiSubscriptionCoversXai + : (runtimeProviderStatus?.runtimeApiKeySatisfied ?? false); const openaiPersistedSatisfied = persistedProviderStatus?.id === 'openai' ? persistedProviderStatus.savedApiKeySatisfied || persistedProviderStatus.runtimeApiKeySatisfied || chatgptCoversOpenAi - : (persistedProviderStatus?.savedApiKeySatisfied ?? - persistedProviderStatus?.runtimeApiKeySatisfied ?? - false); + : persistedProviderStatus?.id === 'xai' + ? persistedProviderStatus.savedApiKeySatisfied || + persistedProviderStatus.runtimeApiKeySatisfied || + xaiSubscriptionCoversXai + : (persistedProviderStatus?.savedApiKeySatisfied ?? + persistedProviderStatus?.runtimeApiKeySatisfied ?? + false); const setupSatisfiedByRuntimeEnv = runtimeRoomoteModelSatisfied && openaiRuntimeSatisfied; // The persisted model works with a key from either source: the effective // model runtime env resolves keys runtime-first with a DB fallback, so a // saved model choice plus a runtime-env key is a complete setup. A - // ChatGPT subscription also satisfies persisted openai/ models. + // ChatGPT subscription also satisfies persisted openai/ models; SuperGrok + // does the same for xai/. // Symmetric: a runtime env model plus a saved provider key is also complete, // because task runtime resolution layer loads keys the same way. const setupSatisfiedByRuntimeModelAndSavedKey = @@ -1684,7 +1748,10 @@ export function buildSetupModelStatus(input: { runtimeProviderStatus && (runtimeProviderStatus.id === 'openai' ? runtimeProviderStatus.savedApiKeySatisfied || chatgptCoversOpenAi - : runtimeProviderStatus.savedApiKeySatisfied), + : runtimeProviderStatus.id === 'xai' + ? runtimeProviderStatus.savedApiKeySatisfied || + xaiSubscriptionCoversXai + : runtimeProviderStatus.savedApiKeySatisfied), ); const setupSatisfied = setupSatisfiedByRuntimeEnv || @@ -1693,6 +1760,7 @@ export function buildSetupModelStatus(input: { persistedRoomoteModel && (openaiPersistedSatisfied || (persistedProviderStatus?.id !== 'openai' && + persistedProviderStatus?.id !== 'xai' && (persistedProviderStatus?.savedApiKeySatisfied || persistedProviderStatus?.runtimeApiKeySatisfied))), ); @@ -1706,6 +1774,26 @@ export function buildSetupModelStatus(input: { openaiProvider.savedApiKeySatisfied), ); + // Key presence for xAI is computed from env only so the UI can show the + // API-key path separately from SuperGrok OAuth. + const xaiApiKeyConnected = (() => { + const required = getSetupModelProviderRequiredEnvVarNames( + getSetupModelProvider('xai'), + ); + if (required.length === 0) { + return false; + } + const isRuntimeConfigured = (name: string) => + isConfiguredEnvValue(runtimeEnv[name]); + const isPersisted = (name: string) => persistedEnvVarNameSet.has(name); + return ( + required.every( + (name) => isPersisted(name) || isRuntimeConfigured(name), + ) && + required.some((name) => isPersisted(name) || isRuntimeConfigured(name)) + ); + })(); + return { runtimeRoomoteModel, runtimeRoomoteModelSatisfied, @@ -1718,6 +1806,8 @@ export function buildSetupModelStatus(input: { setupSatisfiedByRuntimeEnv, chatgptConnected, githubCopilotConnected, + xaiSubscriptionConnected, + xaiApiKeyConnected, openaiAndChatGptBothConfigured, }; } diff --git a/packages/types/src/subscription-provider-usage.ts b/packages/types/src/subscription-provider-usage.ts index eec7c1943..920cb6197 100644 --- a/packages/types/src/subscription-provider-usage.ts +++ b/packages/types/src/subscription-provider-usage.ts @@ -15,6 +15,7 @@ export type SubscriptionUsageProviderId = | 'chatgpt' | 'github-copilot' | 'kimi-for-coding' + | 'xai-subscription' | 'zai' | 'zai-coding-plan'; @@ -88,6 +89,15 @@ export const KIMI_FOR_CODING_USAGE_ENDPOINTS = [ 'https://api.kimi.com/coding/v1/usage', ] as const; +/** + * Unofficial Grok CLI session proxy used for SuperGrok / Premium+ billing + * lookups (same surface Grok Build and peer OAuth tools poll). Not a stable + * public API; parse defensively and omit usage on failure. + */ +export const XAI_CLI_CHAT_PROXY_BASE_URL = 'https://cli-chat-proxy.grok.com'; +export const XAI_USAGE_USER_ENDPOINT = `${XAI_CLI_CHAT_PROXY_BASE_URL}/v1/user`; +export const XAI_USAGE_BILLING_ENDPOINT = `${XAI_CLI_CHAT_PROXY_BASE_URL}/v1/billing?format=credits`; + /** * Undocumented Z.AI / BigModel monitor quota endpoint polled by Coding Plan * tools and community usage bars. Host is region-specific; path is shared. diff --git a/packages/types/src/xai-subscription.ts b/packages/types/src/xai-subscription.ts new file mode 100644 index 000000000..27aea4d1d --- /dev/null +++ b/packages/types/src/xai-subscription.ts @@ -0,0 +1,46 @@ +/** + * xAI Grok subscription OAuth constants. + * + * These reuse the public Grok CLI OAuth client identity and device-code + * endpoints used by peer agents (Hermes, OpenClaw, CC Switch). Roomote does + * not register its own xAI OAuth app; tokens authenticate chat/coding + * inference against `https://api.x.ai` under the existing `xai/` model-id + * prefix. + */ + +/** Public Grok CLI client id shared by ecosystem tools. */ +export const DEFAULT_XAI_OAUTH_CLIENT_ID = + 'b1a00492-073a-47ea-816f-4c329264a828'; +export const XAI_OAUTH_CLIENT_ID_ENV_VAR_NAME = 'XAI_OAUTH_CLIENT_ID'; + +export const XAI_OAUTH_ISSUER = 'https://auth.x.ai'; +export const XAI_OAUTH_DEVICE_CODE_ENDPOINT = `${XAI_OAUTH_ISSUER}/oauth2/device/code`; +export const XAI_OAUTH_TOKEN_ENDPOINT = `${XAI_OAUTH_ISSUER}/oauth2/token`; +/** Fallback verification URL when the device-code response omits one. */ +export const XAI_OAUTH_DEVICE_VERIFICATION_URL = 'https://accounts.x.ai'; + +/** + * Scopes requested by the official Grok CLI for SuperGrok / eligible Premium+ + * API access, including offline_access so Roomote can refresh server-side. + */ +export const XAI_OAUTH_SCOPE = + 'openid profile email offline_access grok-cli:access api:access'; + +/** opencode / models.dev provider id for Grok models (`xai/`). */ +export const XAI_OPENCODE_PROVIDER_ID = 'xai'; + +/** + * Refresh an access token this many milliseconds before its stated expiry. + * xAI access tokens are short-lived (~6h); one hour of margin keeps long + * gateway/cron workloads from hitting brief credential-expiry gaps. + */ +export const XAI_REFRESH_SAFETY_MARGIN_MS = 60 * 60 * 1000; + +/** + * Default access-token lifetime when the token endpoint omits `expires_in`. + * Matches observed SuperGrok device-code grants (~6 hours). + */ +export const XAI_DEFAULT_ACCESS_TOKEN_TTL_MS = 6 * 3600 * 1000; + +/** GitHub-style slow_down bump when the token endpoint asks for it. */ +export const XAI_POLL_SLOW_DOWN_MS = 5_000;