From 15003c9226944a997f51c858634877426b73117e Mon Sep 17 00:00:00 2001 From: Pedro Filho Date: Tue, 8 Sep 2026 11:19:18 -0300 Subject: [PATCH] grok accounts --- CHANGELOG.md | 1 + README.md | 26 +-- package.json | 5 +- src/claude.ts | 10 +- src/cli.ts | 170 +++++++++---------- src/codex.ts | 50 +----- src/config-install.test.ts | 116 ++++++++++++- src/config-install.ts | 158 +++++++++++++++-- src/domain.ts | 20 ++- src/grok.test.ts | 298 +++++++++++++++++++++++++++++++++ src/grok.ts | 335 +++++++++++++++++++++++++++++++++++++ src/login.ts | 36 ++++ src/manager.ts | 35 ++-- src/paths.ts | 3 +- src/pricing.ts | 16 ++ src/proxy.ts | 13 +- src/ratelimit.test.ts | 21 +++ src/ratelimit.ts | 29 +++- src/storage.ts | 6 +- src/tui/dashboard.ts | 74 ++++---- src/tui/fixtures.ts | 66 ++++++-- src/ui.ts | 33 ++-- src/vault.ts | 8 + 23 files changed, 1262 insertions(+), 267 deletions(-) create mode 100644 src/grok.test.ts create mode 100644 src/grok.ts create mode 100644 src/login.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index 246de57..e14f587 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,4 @@ +- [2026-09-08] grok accounts - [2026-08-21] codex keeps its shell and apply_patch tools on the gpt-5.6 models - [2026-08-18] meter clients that hang up early - [2026-08-18] reclaim bare provider tables diff --git a/README.md b/README.md index 4354428..715adf0 100644 --- a/README.md +++ b/README.md @@ -4,9 +4,9 @@
-**One dashboard for all of your Codex and Claude Code accounts. Easily switch between them, and monitor usage.** +**One dashboard for all of your Codex, Claude Code and Grok accounts. Easily switch between them, and monitor usage.** -macOS · [Bun](https://bun.sh) · a [Rubric Labs](https://rubriclabs.com) project · not affiliated with OpenAI or Anthropic +macOS · [Bun](https://bun.sh) · a [Rubric Labs](https://rubriclabs.com) project · not affiliated with OpenAI, Anthropic or xAI

@@ -26,7 +26,7 @@ tokenmaxx #starts the dashboard ## What it does -You run a fleet of coding agents using multiple Codex or Claude accounts: +You run a fleet of coding agents using multiple Codex, Claude or Grok accounts:
a desktop full of parallel agent sessions burning tokens @@ -42,7 +42,7 @@ Run `tokenmaxx`. **Accounts** shows every account and its live rate-limit window every account and its live rate-limit windows, colored by pressure
-**Analytics** is combined token throughput across all accounts and both providers, with the ≈ cost of that usage at API list rates. Tokens are metered as responses stream by, never buffered, so every number is cross-checkable against your clients' own session logs. Press `m` for the full pricing breakdown per model. +**Analytics** is combined token throughput across all accounts and every provider, with the ≈ cost of that usage at API list rates. Tokens are metered as responses stream by, never buffered, so every number is cross-checkable against your clients' own session logs. Press `m` for the full pricing breakdown per model.
combined token throughput and API-list-price cost across all accounts @@ -59,7 +59,7 @@ Run `tokenmaxx`. **Accounts** shows every account and its live rate-limit window Turn it on and tokenmaxx watches the active account's rate-limit windows. When the fullest one crosses your threshold, it switches to whichever of your accounts has the most room. The default threshold is 90%, which leaves the last stretch of every window alone in case you want it later. If an account hits a hard limit in the middle of a request, the proxy retries that request on your next account with room. ```bash -tokenmaxx auto both on --threshold 90 # or: codex | claude … off +tokenmaxx auto all on --threshold 90 # or: codex | claude | grok … off ``` ## How it works @@ -67,19 +67,19 @@ tokenmaxx auto both on --threshold 90 # or: codex | claude … off A single loopback proxy on `127.0.0.1:8459`, and the clients you already use. - **Account per request.** The proxy reads which account is active for each request and attaches its credential. Switching lands on the next request. -- **Pressure read for free.** Both providers report rate-limit state on every response; the proxy reads it as traffic streams by, so it always knows how full the active account is, with zero extra requests. -- **Official apps only.** Your subscription login is for Claude Code and Codex. If you're building something custom, use an API key from the provider. tokenmaxx doesn't turn a subscription into an API plan. +- **Pressure read for free.** Codex and Claude report rate-limit state on every response; the proxy reads it as traffic streams by, so it always knows how full the active account is, with zero extra requests. Grok reports nothing until it says no, so a Grok account's meter is empty until a 429 fills it, and auto-rotation for Grok is driven by hard limits rather than a threshold. +- **Official apps only.** Your subscription login is for Claude Code, Codex and the grok CLI. If you're building something custom, use an API key from the provider. tokenmaxx doesn't turn a subscription into an API plan. ## Commands ```text tokenmaxx live dashboard -tokenmaxx login sign in; isolated, idempotent -tokenmaxx install route native codex & claude +tokenmaxx login sign in; isolated, idempotent +tokenmaxx install route native codex, claude & grok tokenmaxx uninstall restore native config -tokenmaxx switch make an account active -tokenmaxx logout [codex|claude] sign out; the credential is deleted -tokenmaxx auto [--threshold N] +tokenmaxx switch make an account active +tokenmaxx logout [codex|claude|grok] sign out; the credential is deleted +tokenmaxx auto [--threshold N] tokenmaxx list | status | refresh | doctor ``` @@ -94,7 +94,7 @@ tokenmaxx is for one person with accounts they pay for themselves. No account ge ## Not affiliated An independent [Rubric Labs](https://rubriclabs.com) project, not an official -product of, affiliated with, or endorsed by OpenAI or Anthropic. Inspired by +product of, affiliated with, or endorsed by OpenAI, Anthropic or xAI. Inspired by [codex-account-switcher](https://github.com/Sls0n/codex-account-switcher).
diff --git a/package.json b/package.json index 7ec8ed3..034e432 100644 --- a/package.json +++ b/package.json @@ -7,7 +7,7 @@ "@opentui/core": "^0.4.3", "zod": "^4.0.0" }, - "description": "Switch between your own Codex and Claude Code accounts from one local dashboard, with live token analytics.", + "description": "Switch between your own Codex, Claude Code and Grok accounts from one local dashboard, with live token analytics.", "devDependencies": { "@biomejs/biome": "^2.0.0", "@rubriclab/config": "*", @@ -29,6 +29,7 @@ "claude-code", "cli", "codex", + "grok", "proxy", "rate-limits", "rubric", @@ -59,5 +60,5 @@ "post-commit": "bun x @rubriclab/package post-commit" }, "type": "module", - "version": "0.0.66" + "version": "0.0.67" } diff --git a/src/claude.ts b/src/claude.ts index 4fec5bd..9772ad6 100644 --- a/src/claude.ts +++ b/src/claude.ts @@ -13,7 +13,7 @@ import { } from './domain.ts' import { ApplicationError, loginFailureMessage } from './errors.ts' import { type UpstreamInjection, upstreamFor } from './proxy.ts' -import { type CredentialVault, exclusive } from './vault.ts' +import { type CredentialVault, exclusive, readApiKey } from './vault.ts' const clientId = '9d1c250a-e61b-44d9-88ed-5944d1962f5e' const tokenEndpoint = 'https://console.anthropic.com/v1/oauth/token' @@ -203,14 +203,6 @@ async function readClaudeCredential( return ClaudeOauthSchema.parse(JSON.parse(serialized)) } -async function readApiKey(vault: CredentialVault, reference: string): Promise { - const key = await vault.read(reference) - if (key === null) { - throw new ApplicationError('CREDENTIAL_MISSING', `Missing credential ${reference}`) - } - return key -} - const anthropicVersion = '2023-06-01' async function validateAnthropicApiKey( diff --git a/src/cli.ts b/src/cli.ts index 356b520..92a7b6b 100644 --- a/src/cli.ts +++ b/src/cli.ts @@ -9,17 +9,16 @@ import { registerClaudeAccount, registerClaudeApiKeyAccount } from './claude.ts' import { registerCodexAccount, registerOpenAiApiKeyAccount } from './codex.ts' import { healInstalledConfigs, - installClaudeConfig, - installCodexConfig, installPiConfig, + installProviderConfig, installStatus, piStatus, - uninstallClaudeConfig, - uninstallCodexConfig, - uninstallPiConfig + uninstallPiConfig, + uninstallProviderConfig } from './config-install.ts' -import type { Account, ProviderId } from './domain.ts' +import { type Account, PROVIDERS, type ProviderId, ProviderIdSchema } from './domain.ts' import { ApplicationError, errorMessage } from './errors.ts' +import { registerGrokAccount, registerXaiApiKeyAccount } from './grok.ts' import { managerAvailable, managerRequest, @@ -194,7 +193,9 @@ interface ApplicationContext { store: StateStore } -function providerFromCli(value: string): 'openai' | 'anthropic' { +const cliWords = '' + +function providerFromCli(value: string): ProviderId { switch (value) { case 'codex': case 'openai': @@ -202,8 +203,14 @@ function providerFromCli(value: string): 'openai' | 'anthropic' { case 'claude': case 'anthropic': return 'anthropic' + case 'grok': + case 'xai': + return 'xai' default: - throw new ApplicationError('INVALID_PROVIDER', `Expected codex or claude, received ${value}`) + throw new ApplicationError( + 'INVALID_PROVIDER', + `Expected codex, claude or grok, received ${value}` + ) } } @@ -234,25 +241,25 @@ function help(): string { return [` ${accent(name)}`, ...lines.map(line => `${gutter}${dim(line)}`)].join('\n') } return [ - `${accent('tokenmaxx')} ${dim('— switch between your own Codex and Claude Code accounts')}`, + `${accent('tokenmaxx')} ${dim('— switch between your own Codex, Claude Code and Grok accounts')}`, '', `${head('Usage')} tokenmaxx [options] ${dim('run with no command for the dashboard')}`, '', head('Setup'), row( - 'login ', + `login ${cliWords}`, 'sign in an account · re-run to re-auth', 'add --api-key to use an API key instead' ), - row('install [pi]', 'route codex & claude, or pi, through tokenmaxx'), + row('install [pi]', 'route codex, claude & grok, or pi, through tokenmaxx'), row('uninstall [pi]', 'restore your original config'), '', head('Everyday'), row('list', 'accounts, health, and live usage'), - row('switch ', 'make an account active now'), - row('logout [codex|claude] ', 'sign out and delete the credential'), + row(`switch ${cliWords} `, 'make an account active now'), + row('logout [codex|claude|grok] ', 'sign out and delete the credential'), row( - 'auto ', + 'auto ', 'switch accounts at a usage threshold', 'optional: --threshold N (default 90)' ), @@ -273,7 +280,7 @@ function help(): string { dim(' Threshold switches hold for 5 minutes to avoid flapping; hard limits'), dim(' ignore the hold. Turning auto on is what authorizes the switching.'), '', - dim('Once installed, use codex and claude normally — a local proxy injects the'), + dim('Once installed, use codex, claude and grok normally — a local proxy injects the'), dim("active account's credential per request, so a switch takes effect on the"), dim('next request, even mid-turn, with no restart.') ].join('\n') @@ -312,7 +319,7 @@ async function runDaemon(context: ApplicationContext): Promise { }) if (healed.length > 0) { process.stdout.write( - `[${new Date().toISOString()}] re-applied ${healed.join(' and ')} routing for v${VERSION}\n` + `[${new Date().toISOString()}] re-applied ${healed.map(provider => PROVIDERS[provider].cli).join(', ')} routing for v${VERSION}\n` ) } const manager = new AccountManager({ @@ -460,22 +467,25 @@ async function ensureDaemon(context: ApplicationContext): Promise { } } -function registerIsolatedAccount(provider: 'openai' | 'anthropic'): Promise { +function registerIsolatedAccount(provider: ProviderId): Promise { switch (provider) { case 'openai': return registerCodexAccount({ vault: createMacOsKeychainVault() }) case 'anthropic': return registerClaudeAccount({ vault: createMacOsKeychainVault() }) + case 'xai': + return registerGrokAccount({ vault: createMacOsKeychainVault() }) } } const cliInstallHint: Record = { anthropic: 'npm install -g @anthropic-ai/claude-code', - openai: 'npm install -g @openai/codex' + openai: 'npm install -g @openai/codex', + xai: 'npm install -g @xai-official/grok' } function assertCliInstalled(provider: ProviderId): void { - const binary = provider === 'openai' ? 'codex' : 'claude' + const binary = PROVIDERS[provider].cli if (Bun.which(binary) === null) { throw new ApplicationError( 'CLI_MISSING', @@ -544,13 +554,13 @@ export function stripTerminalNoise(line: string): string { } async function registerApiKeyAccount( - provider: 'openai' | 'anthropic', + provider: ProviderId, keyArgument: string | undefined ): Promise { if (process.stdin.isTTY !== true && keyArgument === undefined) { throw new ApplicationError( 'USAGE', - 'Pass the key inline in non-interactive shells: tokenmaxx login --api-key ' + `Pass the key inline in non-interactive shells: tokenmaxx login ${cliWords} --api-key ` ) } await handTerminalBack() @@ -568,10 +578,15 @@ async function registerApiKeyAccount( if (label.trim().length === 0) { throw new ApplicationError('USAGE', 'The account needs a name') } - const vault = createMacOsKeychainVault() - return provider === 'openai' - ? registerOpenAiApiKeyAccount({ key, label: label.trim(), vault }) - : registerClaudeApiKeyAccount({ key, label: label.trim(), vault }) + const input = { key, label: label.trim(), vault: createMacOsKeychainVault() } + switch (provider) { + case 'openai': + return registerOpenAiApiKeyAccount(input) + case 'anthropic': + return registerClaudeApiKeyAccount(input) + case 'xai': + return registerXaiApiKeyAccount(input) + } } async function login( @@ -580,7 +595,7 @@ async function login( options: { apiKey: boolean; apiKeyValue?: string } = { apiKey: false } ): Promise { if (providerArgument === undefined) { - throw new ApplicationError('USAGE', 'Usage: tokenmaxx login [--api-key [key]]') + throw new ApplicationError('USAGE', `Usage: tokenmaxx login ${cliWords} [--api-key [key]]`) } const provider = providerFromCli(providerArgument) if (!options.apiKey) { @@ -621,9 +636,7 @@ async function login( : `Re-authenticated ${account.label}; live sessions pick it up on their next request.\n` ) if (existing === undefined) { - const status = await installStatus() - const alreadyRouted = provider === 'openai' ? status.codexRouted : status.claudeRouted - if (!alreadyRouted) { + if (!(await installStatus()).routed[provider]) { await setRouting(context, provider, true).catch(() => undefined) process.stdout.write( `tokenmaxx is on for ${providerArgument} — run ${providerArgument} as usual.\n` @@ -638,11 +651,7 @@ async function removeUnstoredAccount(account: Account): Promise { } } -function resolveAccount( - store: StateStore, - provider: 'openai' | 'anthropic', - reference: string -): Account { +function resolveAccount(store: StateStore, provider: ProviderId, reference: string): Account { const matches = store .listAccounts(provider) .filter(account => account.id === reference || account.label === reference) @@ -670,21 +679,19 @@ function listAccounts(context: ApplicationContext): void { const states = new Map(context.store.listProviderStates().map(state => [state.provider, state])) const accounts = context.store.listAccounts() if (accounts.length === 0) { - process.stdout.write( - 'No accounts yet. Sign in with: tokenmaxx login codex · tokenmaxx login claude\n' + const hints = ProviderIdSchema.options.map( + provider => `tokenmaxx login ${PROVIDERS[provider].cli}` ) + process.stdout.write(`No accounts yet. Sign in with: ${hints.join(' · ')}\n`) return } const width = Math.max(...accounts.map(account => account.label.length)) - for (const [provider, title] of [ - ['openai', 'codex'], - ['anthropic', 'claude'] - ] as const) { + for (const provider of ProviderIdSchema.options) { const group = accounts.filter(account => account.provider === provider) if (group.length === 0) { continue } - process.stdout.write(`\n${title}\n`) + process.stdout.write(`\n${PROVIDERS[provider].cli}\n`) for (const account of group) { const isActive = states.get(provider)?.activeAccountId === account.id process.stdout.write( @@ -695,14 +702,14 @@ function listAccounts(context: ApplicationContext): void { process.stdout.write('\n● = active\n') } -const providerWords = new Set(['codex', 'claude', 'openai', 'anthropic']) +const providerWords = new Set(['codex', 'claude', 'grok', 'openai', 'anthropic', 'xai']) async function logout(context: ApplicationContext, arguments_: readonly string[]): Promise { const qualified = arguments_[0] !== undefined && providerWords.has(arguments_[0]) const provider = qualified ? providerFromCli(arguments_[0] as string) : undefined const reference = qualified ? arguments_[1] : arguments_[0] if (reference === undefined) { - throw new ApplicationError('USAGE', 'Usage: tokenmaxx logout [codex|claude] ') + throw new ApplicationError('USAGE', 'Usage: tokenmaxx logout [codex|claude|grok] ') } const matches = context.store .listAccounts(provider) @@ -729,7 +736,7 @@ async function switchAccount( const providerArgument = arguments_[0] const accountReference = arguments_[1] if (providerArgument === undefined || accountReference === undefined) { - throw new ApplicationError('USAGE', 'Usage: tokenmaxx switch ') + throw new ApplicationError('USAGE', `Usage: tokenmaxx switch ${cliWords} `) } const provider = providerFromCli(providerArgument) const target = resolveAccount(context.store, provider, accountReference) @@ -747,13 +754,11 @@ async function configureAutomation( if (providerArgument === undefined || (mode !== 'on' && mode !== 'off')) { throw new ApplicationError( 'USAGE', - 'Usage: tokenmaxx auto [--threshold 95]' + 'Usage: tokenmaxx auto [--threshold 95]' ) } const providers = - providerArgument === 'both' - ? (['openai', 'anthropic'] as const) - : ([providerFromCli(providerArgument)] as const) + providerArgument === 'all' ? ProviderIdSchema.options : [providerFromCli(providerArgument)] const thresholdValue = option(arguments_, '--threshold') const thresholdPercent = thresholdValue === undefined ? undefined : Number(thresholdValue) if ( @@ -800,17 +805,18 @@ async function installConfig(context: ApplicationContext, targetArgument?: strin return } process.stdout.write( - `pi now has tokenmaxx-anthropic and tokenmaxx-openai providers (${result.path}).\n` + + `pi now has tokenmaxx-anthropic, tokenmaxx-openai and tokenmaxx-xai providers (${result.path}).\n` + 'Pick a tokenmaxx model with /model and requests route through the proxy.\n' + 'Undo any time with: tokenmaxx uninstall pi\n' ) return } - await installCodexConfig(context.paths) - await installClaudeConfig(context.paths) + for (const provider of ProviderIdSchema.options) { + await installProviderConfig(provider, context.paths) + } process.stdout.write( - 'Native codex and claude now route through tokenmaxx.\n' + - 'Just run `codex` or `claude` as usual — tokenmaxx injects the active account.\n' + + 'Native codex, claude and grok now route through tokenmaxx.\n' + + 'Just run `codex`, `claude` or `grok` as usual — tokenmaxx injects the active account.\n' + 'Undo any time with: tokenmaxx uninstall\n' ) } @@ -828,14 +834,15 @@ async function uninstallConfig(targetArgument?: string): Promise { ) return } - const codex = await uninstallCodexConfig() - const claude = await uninstallClaudeConfig() - if (codex === null && claude === null) { + const restored = await Promise.all( + ProviderIdSchema.options.map(provider => uninstallProviderConfig(provider)) + ) + if (restored.every(path => path === null)) { process.stdout.write('tokenmaxx was not installed; nothing to restore.\n') return } process.stdout.write( - 'Restored your original codex and claude config.\n' + + 'Restored your original codex, claude and grok config.\n' + 'Native clients no longer route through tokenmaxx. Re-enable with: tokenmaxx install\n' ) } @@ -845,18 +852,15 @@ async function setRouting( provider: ProviderId, enable: boolean ): Promise { - if (provider === 'openai') { - await (enable ? installCodexConfig(context.paths) : uninstallCodexConfig()) - } else { - await (enable ? installClaudeConfig(context.paths) : uninstallClaudeConfig()) - } + await (enable ? installProviderConfig(provider, context.paths) : uninstallProviderConfig(provider)) } async function doctor(context: ApplicationContext): Promise { const tools = [ ['bun', '1.2+'], ['codex', '0.144.1'], - ['claude', '2.1.206'] + ['claude', '2.1.206'], + ['grok', '1.0.13'] ] as const for (const [tool, testedVersion] of tools) { if (Bun.which(tool) === null) { @@ -893,22 +897,22 @@ async function doctor(context: ApplicationContext): Promise { ) } const routing = await installStatus() - process.stdout.write( - `${routing.codexRouted ? 'ok ' : 'note '} codex ${ - routing.codexRouted - ? 'config.toml selects the tokenmaxx provider' - : routing.codexStale - ? 'a tokenmaxx block exists but codex ignores it (top-level key was swallowed by a [table]) — run tokenmaxx install to repair' - : 'not routed — run tokenmaxx install' - }\n` - ) - process.stdout.write( - `${routing.claudeRouted ? 'ok ' : 'note '} claude ${ - routing.claudeRouted - ? 'settings.json routes ANTHROPIC_BASE_URL through tokenmaxx' + const routedText: Record = { + anthropic: 'settings.json routes ANTHROPIC_BASE_URL through tokenmaxx', + openai: 'config.toml selects the tokenmaxx provider', + xai: 'config.toml points cli_chat_proxy_base_url at tokenmaxx' + } + for (const provider of ProviderIdSchema.options) { + const routed = routing.routed[provider] + const detail = routed + ? routedText[provider] + : provider === 'openai' && routing.codexStale + ? 'a tokenmaxx block exists but codex ignores it (top-level key was swallowed by a [table]) — run tokenmaxx install to repair' : 'not routed — run tokenmaxx install' - }\n` - ) + process.stdout.write( + `${routed ? 'ok ' : 'note '} ${PROVIDERS[provider].cli.padEnd(8)} ${detail}\n` + ) + } const pi = await piStatus() if (pi.present) { process.stdout.write( @@ -956,7 +960,7 @@ export async function runCli(rawArguments: readonly string[]): Promise { now, timewarp: Number.isFinite(timewarp) && timewarp > 0 ? timewarp : 0 }, - routing: { anthropic: routed, openai: routed } + routing: { anthropic: routed, openai: routed, xai: routed } }) context.store.close() process.exit(0) @@ -964,10 +968,8 @@ export async function runCli(rawArguments: readonly string[]): Promise { await ensureDaemon(context) if (process.stdout.isTTY) { const { runTuiDashboard } = await import('./tui/dashboard.ts') - const readRouting = async (): Promise> => { - const status = await installStatus() - return { anthropic: status.claudeRouted, openai: status.codexRouted } - } + const readRouting = async (): Promise> => + (await installStatus()).routed let alert = '' for (;;) { const action = await runTuiDashboard(context.paths.managerSocket, { @@ -986,7 +988,7 @@ export async function runCli(rawArguments: readonly string[]): Promise { continue } if (action.kind === 'relogin' || action.kind === 'login' || action.kind === 'loginApiKey') { - const cli = action.provider === 'openai' ? 'codex' : 'claude' + const cli = PROVIDERS[action.provider].cli freshScreen(action.kind === 'loginApiKey' ? `add a ${cli} api key` : `sign in with ${cli}`) await login(context, cli, { apiKey: action.kind === 'loginApiKey' diff --git a/src/codex.ts b/src/codex.ts index 8edc6e2..c916399 100644 --- a/src/codex.ts +++ b/src/codex.ts @@ -1,5 +1,3 @@ -import { mkdtemp, readFile, rm } from 'node:fs/promises' -import { tmpdir } from 'node:os' import { join } from 'node:path' import { z } from 'zod' import { @@ -13,8 +11,9 @@ import { type UsageWindow } from './domain.ts' import { ApplicationError, loginFailureMessage } from './errors.ts' +import { defaultIsolatedLoginDependencies, type IsolatedLoginDependencies } from './login.ts' import { type UpstreamInjection, upstreamFor } from './proxy.ts' -import { type CredentialVault, exclusive } from './vault.ts' +import { type CredentialVault, exclusive, readApiKey } from './vault.ts' const clientId = 'app_EMoamEEZ73f0CkXaXp7hrann' const refreshEndpoint = 'https://auth.openai.com/oauth/token' @@ -75,16 +74,6 @@ interface CodexIdentity { accessExpiresAt: string | null } -interface CodexLoginDependencies { - run( - command: readonly string[], - environment: Record - ): Promise<{ exitCode: number; stderr: string }> - createTemporaryDirectory(prefix: string): Promise - read(path: string): Promise - remove(path: string): Promise -} - function base64UrlJson(segment: string): unknown { try { return JSON.parse(Buffer.from(segment, 'base64url').toString('utf8')) @@ -133,34 +122,11 @@ function codexIdentity(auth: CodexAuth): CodexIdentity { } } -function defaultCodexLoginDependencies(): CodexLoginDependencies { - return { - createTemporaryDirectory: prefix => mkdtemp(join(tmpdir(), prefix)), - read: path => readFile(path, 'utf8'), - remove: path => rm(path, { force: true, recursive: true }), - async run(command, environment) { - const child = Bun.spawn([...command], { - env: { ...process.env, ...environment }, - stderr: 'pipe', - stdin: 'inherit', - stdout: 'inherit' - }) - const decoder = new TextDecoder() - let stderr = '' - for await (const chunk of child.stderr) { - process.stderr.write(chunk) - stderr = `${stderr}${decoder.decode(chunk, { stream: true })}`.slice(-4_096) - } - return { exitCode: await child.exited, stderr } - } - } -} - export async function registerCodexAccount(input: { vault: CredentialVault - dependencies?: CodexLoginDependencies + dependencies?: IsolatedLoginDependencies }): Promise { - const dependencies = input.dependencies ?? defaultCodexLoginDependencies() + const dependencies = input.dependencies ?? defaultIsolatedLoginDependencies() const temporaryHome = await dependencies.createTemporaryDirectory('tokenmaxx-register-') try { const login = await dependencies.run( @@ -266,14 +232,6 @@ async function refreshCodexCredential(input: { }) } -async function readApiKey(vault: CredentialVault, reference: string): Promise { - const key = await vault.read(reference) - if (key === null) { - throw new ApplicationError('CREDENTIAL_MISSING', `Missing credential ${reference}`) - } - return key -} - const openAiApiBase = 'https://api.openai.com/v1' async function validateOpenAiApiKey( diff --git a/src/config-install.test.ts b/src/config-install.test.ts index a7da7c3..29f7195 100644 --- a/src/config-install.test.ts +++ b/src/config-install.test.ts @@ -7,10 +7,12 @@ import { healInstalledConfigs, installClaudeConfig, installCodexConfig, + installGrokConfig, installPiConfig, installStatus, uninstallClaudeConfig, uninstallCodexConfig, + uninstallGrokConfig, uninstallPiConfig } from './config-install.ts' import { applicationPaths } from './paths.ts' @@ -40,13 +42,16 @@ beforeEach(async () => { home = mkdtempSync(join(tmpdir(), 'tokenmaxx-config-')) process.env.CODEX_HOME = join(home, 'codex') process.env.CLAUDE_CONFIG_DIR = join(home, 'claude') + process.env.GROK_HOME = join(home, 'grok') await mkdir(process.env.CODEX_HOME, { recursive: true }) await mkdir(process.env.CLAUDE_CONFIG_DIR, { recursive: true }) + await mkdir(process.env.GROK_HOME, { recursive: true }) }) afterEach(() => { delete process.env.CODEX_HOME delete process.env.CLAUDE_CONFIG_DIR + delete process.env.GROK_HOME rmSync(home, { force: true, recursive: true }) }) @@ -77,7 +82,7 @@ describe('installCodexConfig', () => { test('installStatus flags the legacy swallowed block as stale, not routed', async () => { await writeCodexConfig(legacyBrokenConfig) const status = await installStatus() - expect(status.codexRouted).toBe(false) + expect(status.routed.openai).toBe(false) expect(status.codexStale).toBe(true) }) @@ -85,7 +90,7 @@ describe('installCodexConfig', () => { await writeCodexConfig(legacyBrokenConfig) await installCodexConfig(paths()) const status = await installStatus() - expect(status.codexRouted).toBe(true) + expect(status.routed.openai).toBe(true) expect(status.codexStale).toBe(false) }) @@ -172,7 +177,7 @@ describe('installClaudeConfig', () => { test('installStatus reports claude routing after install', async () => { await installClaudeConfig(paths()) - expect((await installStatus()).claudeRouted).toBe(true) + expect((await installStatus()).routed.anthropic).toBe(true) }) test('uninstall removes the routing but keeps a token the user set themselves', async () => { @@ -210,7 +215,7 @@ describe('healInstalledConfigs', () => { ANTHROPIC_BASE_URL: 'http://127.0.0.1:8459/anthropic' } }) - expect(await healInstalledConfigs(paths())).toEqual(['codex', 'claude']) + expect(await healInstalledConfigs(paths())).toEqual(['openai', 'anthropic']) const settings = await readClaudeSettings() expect(settings.env?.ANTHROPIC_AUTH_TOKEN).toBeUndefined() expect(settings.env?.ANTHROPIC_BASE_URL).toBe('http://127.0.0.1:8459/anthropic') @@ -219,7 +224,7 @@ describe('healInstalledConfigs', () => { test('never adds routing to an unrouted harness', async () => { expect(await healInstalledConfigs(paths())).toEqual([]) await expect(readClaudeSettings()).rejects.toThrow() - expect((await installStatus()).codexRouted).toBe(false) + expect((await installStatus()).routed.openai).toBe(false) }) test('runs once per version, not on every start', async () => { @@ -251,12 +256,14 @@ describe('pi install', () => { const config = JSON.parse(await readFile(modelsPath, 'utf8')) expect(config.providers['tokenmaxx-anthropic'].api).toBe('anthropic-messages') expect(config.providers['tokenmaxx-openai'].baseUrl).toContain('/openai') + expect(config.providers['tokenmaxx-xai'].baseUrl).toContain('/xai/v1') expect(config.providers.mine.baseUrl).toBe('https://example.com') const removed = await uninstallPiConfig() expect(removed.applied).toBe(true) const restored = JSON.parse(await readFile(modelsPath, 'utf8')) expect(restored.providers['tokenmaxx-anthropic']).toBeUndefined() expect(restored.providers['tokenmaxx-openai']).toBeUndefined() + expect(restored.providers['tokenmaxx-xai']).toBeUndefined() expect(restored.providers.mine.baseUrl).toBe('https://example.com') delete process.env.PI_CODING_AGENT_DIR }) @@ -269,7 +276,11 @@ describe('pi install', () => { const installed = await installPiConfig(applicationPaths()) expect(installed.applied).toBe(true) const config = JSON.parse(await readFile(installed.path, 'utf8')) - expect(Object.keys(config.providers)).toEqual(['tokenmaxx-anthropic', 'tokenmaxx-openai']) + expect(Object.keys(config.providers)).toEqual([ + 'tokenmaxx-anthropic', + 'tokenmaxx-openai', + 'tokenmaxx-xai' + ]) delete process.env.PI_CODING_AGENT_DIR }) @@ -286,6 +297,97 @@ describe('pi install', () => { }) }) +async function writeGrokConfig(content: string): Promise { + await writeFile(join(process.env.GROK_HOME ?? '', 'config.toml'), content) +} + +async function readGrokConfig(): Promise { + return readFile(join(process.env.GROK_HOME ?? '', 'config.toml'), 'utf8') +} + +const userGrokConfig = `[marketplace] +official_marketplace_auto_installed = true + +[[marketplace.sources]] +name = "xAI Official" +git = "https://github.com/xai-org/plugin-marketplace.git" +` + +describe('installGrokConfig', () => { + test('a fresh config gets a marked endpoints table', async () => { + await installGrokConfig(paths()) + const written = await readGrokConfig() + expect(written).toContain('[endpoints]') + expect(written).toContain('cli_chat_proxy_base_url = "http://127.0.0.1:8459/xai/v1"') + expect((await installStatus()).routed.xai).toBe(true) + }) + + test('the endpoint lands inside an existing endpoints table so TOML stays valid', async () => { + await writeGrokConfig( + `${userGrokConfig}\n[endpoints]\nmodels_list_url = "https://example.com/models"\n` + ) + await installGrokConfig(paths()) + const written = await readGrokConfig() + const parsed = Bun.TOML.parse(written) as { + endpoints?: Record + marketplace?: { sources?: unknown[] } + } + expect(parsed.endpoints?.cli_chat_proxy_base_url).toBe('http://127.0.0.1:8459/xai/v1') + expect(parsed.endpoints?.models_list_url).toBe('https://example.com/models') + expect(parsed.marketplace?.sources).toHaveLength(1) + expect(written.match(/\[endpoints\]/g)).toHaveLength(1) + }) + + test('a user-set endpoint is disabled on install and restored on uninstall', async () => { + await writeGrokConfig('[endpoints]\ncli_chat_proxy_base_url = "https://grok-proxy.acme.com/v1"\n') + await installGrokConfig(paths()) + const written = await readGrokConfig() + expect(written).toContain( + '# tokenmaxx-disabled: cli_chat_proxy_base_url = "https://grok-proxy.acme.com/v1"' + ) + expect((await installStatus()).routed.xai).toBe(true) + expect(await uninstallGrokConfig()).not.toBeNull() + expect(await readGrokConfig()).toBe( + '[endpoints]\ncli_chat_proxy_base_url = "https://grok-proxy.acme.com/v1"\n' + ) + expect((await installStatus()).routed.xai).toBe(false) + }) + + test('reinstall is idempotent', async () => { + await writeGrokConfig(userGrokConfig) + await installGrokConfig(paths()) + const once = await readGrokConfig() + await installGrokConfig(paths()) + expect(await readGrokConfig()).toBe(once) + expect(once.match(/cli_chat_proxy_base_url/g)).toHaveLength(1) + }) + + test('uninstall restores the user config and drops the table it added', async () => { + await writeGrokConfig(userGrokConfig) + await installGrokConfig(paths()) + expect(await uninstallGrokConfig()).not.toBeNull() + expect(await readGrokConfig()).toBe(userGrokConfig) + expect(await uninstallGrokConfig()).toBeNull() + }) + + test('a bare endpoint line grok re-serialized without our markers is reclaimed', async () => { + await writeGrokConfig( + `${userGrokConfig}\n[endpoints]\ncli_chat_proxy_base_url = "http://127.0.0.1:8459/xai/v1"\n` + ) + expect((await installStatus()).routed.xai).toBe(true) + await installGrokConfig(paths()) + expect((await readGrokConfig()).match(/cli_chat_proxy_base_url/g)).toHaveLength(1) + expect(await uninstallGrokConfig()).not.toBeNull() + expect(await readGrokConfig()).toBe(userGrokConfig) + }) + + test('a config with no grok routing is left alone on uninstall', async () => { + await writeGrokConfig(userGrokConfig) + expect(await uninstallGrokConfig()).toBeNull() + expect(await readGrokConfig()).toBe(userGrokConfig) + }) +}) + describe('codex-normalized configs', () => { test('install reclaims a bare provider table codex re-serialized without our markers', async () => { const configPath = join(process.env.CODEX_HOME ?? '', 'config.toml') @@ -324,7 +426,7 @@ describe('codex-normalized configs', () => { expect(parsed.notice?.hide_rate_limit_model_nudge).toBe(true) expect(written.match(/\[model_providers\.tokenmaxx\]/g)).toHaveLength(1) expect(written).not.toContain('# tokenmaxx-disabled: model_provider = "tokenmaxx"') - expect((await installStatus()).codexRouted).toBe(true) + expect((await installStatus()).routed.openai).toBe(true) }) test('uninstall removes a bare provider table even when the markers are gone', async () => { diff --git a/src/config-install.ts b/src/config-install.ts index fb10149..4c73a81 100644 --- a/src/config-install.ts +++ b/src/config-install.ts @@ -1,6 +1,7 @@ import { mkdir, readFile, stat, writeFile } from 'node:fs/promises' import { homedir } from 'node:os' import { dirname, join } from 'node:path' +import { type ProviderId, ProviderIdSchema } from './domain.ts' import type { ApplicationPaths } from './paths.ts' import { proxyBaseUrl } from './paths.ts' import { VERSION } from './version.ts' @@ -24,6 +25,10 @@ function claudeSettingsPath(): string { return join(process.env.CLAUDE_CONFIG_DIR ?? join(homedir(), '.claude'), 'settings.json') } +function grokConfigPath(): string { + return join(process.env.GROK_HOME ?? join(homedir(), '.grok'), 'config.toml') +} + async function readFileOrEmpty(path: string): Promise { return readFile(path, 'utf8').catch(() => '') } @@ -178,9 +183,116 @@ export async function uninstallClaudeConfig(): Promise { return path } +const endpointBeginMarker = '# >>> tokenmaxx endpoint (do not edit) >>>' +const endpointEndMarker = '# <<< tokenmaxx endpoint <<<' +const grokEndpointKey = 'cli_chat_proxy_base_url' +const ownEndpointLine = + /^\s*(?:endpoints\.)?cli_chat_proxy_base_url\s*=\s*"[^"]*127\.0\.0\.1:\d+\/xai\/v1"\s*$/ +const anyEndpointLine = /^\s*(?:endpoints\.)?cli_chat_proxy_base_url\s*=/ +const endpointsHeader = /^\s*\[endpoints\]\s*$/ +const tableHeader = /^\s*\[/ + +function withoutMarkedLines(lines: string[]): string[] { + const begin = lines.findIndex(line => line.trim() === endpointBeginMarker) + const end = lines.findIndex(line => line.trim() === endpointEndMarker) + return begin === -1 || end < begin ? lines : [...lines.slice(0, begin), ...lines.slice(end + 1)] +} + +function stripGrokManagedBlocks(content: string): string { + return withoutMarkedLines(content.split('\n')) + .filter(line => !ownEndpointLine.test(line.replace(disabledPrefix, ''))) + .map(line => (anyEndpointLine.test(line) ? `# tokenmaxx-disabled: ${line.trimStart()}` : line)) + .join('\n') + .replace(/\n{3,}/g, '\n\n') + .trim() +} + +function dropEmptyEndpointsTable(lines: string[]): string[] { + return lines.filter((line, index) => { + if (!endpointsHeader.test(line)) { + return true + } + const rest = lines.slice(index + 1) + const next = rest.findIndex(candidate => candidate.trim().length > 0) + return next !== -1 && !tableHeader.test(rest[next] ?? '') + }) +} + +function restoreGrokContent(content: string): string { + const restored = stripGrokManagedBlocks(content) + .split('\n') + .map(line => line.replace(disabledPrefix, '')) + return `${dropEmptyEndpointsTable(restored).join('\n').trimEnd()}\n` +} + +export async function installGrokConfig(paths: ApplicationPaths): Promise { + const path = grokConfigPath() + const base = stripGrokManagedBlocks(await readFileOrEmpty(path)) + const managedLine = `${grokEndpointKey} = "${proxyBaseUrl(paths, 'xai')}/v1"` + const lines = base.length === 0 ? [] : base.split('\n') + const header = lines.findIndex(line => endpointsHeader.test(line)) + const managed = + header === -1 + ? [ + ...lines, + ...(lines.length === 0 ? [] : ['']), + endpointBeginMarker, + '[endpoints]', + managedLine, + endpointEndMarker + ] + : [ + ...lines.slice(0, header + 1), + endpointBeginMarker, + managedLine, + endpointEndMarker, + ...lines.slice(header + 1) + ] + await mkdir(dirname(path), { recursive: true }) + await writeFile(path, `${managed.join('\n')}\n`, { mode: 0o600 }) + return path +} + +export async function uninstallGrokConfig(): Promise { + const path = grokConfigPath() + const existing = await readFile(path, 'utf8').catch(() => null) + const carriesOurConfig = (content: string): boolean => + content.includes(endpointBeginMarker) || + content.split('\n').some(line => ownEndpointLine.test(line) || disabledPrefix.test(line)) + if (existing === null || !carriesOurConfig(existing)) { + return null + } + await writeFile(path, restoreGrokContent(existing), { mode: 0o600 }) + return path +} + +export function installProviderConfig( + provider: ProviderId, + paths: ApplicationPaths +): Promise { + switch (provider) { + case 'openai': + return installCodexConfig(paths) + case 'anthropic': + return installClaudeConfig(paths) + case 'xai': + return installGrokConfig(paths) + } +} + +export function uninstallProviderConfig(provider: ProviderId): Promise { + switch (provider) { + case 'openai': + return uninstallCodexConfig() + case 'anthropic': + return uninstallClaudeConfig() + case 'xai': + return uninstallGrokConfig() + } +} + interface InstallStatus { - codexRouted: boolean - claudeRouted: boolean + routed: Record codexStale: boolean } @@ -210,26 +322,35 @@ export async function installStatus(): Promise { } catch { claudeRouted = false } - return { claudeRouted, codexRouted, codexStale } + + let grokRouted = false + try { + const parsed = Bun.TOML.parse(await readFileOrEmpty(grokConfigPath())) as { + endpoints?: { cli_chat_proxy_base_url?: unknown } + } + const baseUrl = parsed.endpoints?.cli_chat_proxy_base_url + grokRouted = typeof baseUrl === 'string' && baseUrl.includes('127.0.0.1') + } catch { + grokRouted = false + } + return { codexStale, routed: { anthropic: claudeRouted, openai: codexRouted, xai: grokRouted } } } // Configs written by an older version stay stale after an update (#17): re-apply // install for whatever is currently routed, once per version change. Never adds // routing — a harness the user uninstalled or never installed stays untouched. -export async function healInstalledConfigs(paths: ApplicationPaths): Promise { +export async function healInstalledConfigs(paths: ApplicationPaths): Promise { const stampPath = join(paths.root, 'healed-version') if ((await readFileOrEmpty(stampPath)).trim() === VERSION) { return [] } - const { claudeRouted, codexRouted } = await installStatus() - const healed: string[] = [] - if (codexRouted) { - await installCodexConfig(paths) - healed.push('codex') - } - if (claudeRouted) { - await installClaudeConfig(paths) - healed.push('claude') + const { routed } = await installStatus() + const healed: ProviderId[] = [] + for (const provider of ProviderIdSchema.options) { + if (routed[provider]) { + await installProviderConfig(provider, paths) + healed.push(provider) + } } await mkdir(paths.root, { recursive: true }) await writeFile(stampPath, `${VERSION}\n`) @@ -246,13 +367,14 @@ function piModelsPath(): string { return join(process.env.PI_CODING_AGENT_DIR ?? join(homedir(), '.pi', 'agent'), 'models.json') } -const piProviderKeys = ['tokenmaxx-anthropic', 'tokenmaxx-openai'] +const piProviderKeys = ['tokenmaxx-anthropic', 'tokenmaxx-openai', 'tokenmaxx-xai'] // The anthropic ids pair with an API-key account (subscription auth is not for // third-party harnesses); gpt-5.6-sol is the one id the ChatGPT codex backend // accepts for subscription accounts. const piAnthropicModelIds = ['claude-opus-4-8', 'claude-sonnet-4-6'] const piOpenaiModelIds = ['gpt-5.6-sol'] +const piXaiModelIds = ['grok-4.6'] function piProviders(paths: ApplicationPaths): Record { const models = (ids: readonly string[]) => ids.map(id => ({ id, reasoning: true })) @@ -268,6 +390,12 @@ function piProviders(paths: ApplicationPaths): Record { apiKey: dummyAuthToken, baseUrl: proxyBaseUrl(paths, 'openai'), models: models(piOpenaiModelIds) + }, + 'tokenmaxx-xai': { + api: 'openai-responses', + apiKey: dummyAuthToken, + baseUrl: `${proxyBaseUrl(paths, 'xai')}/v1`, + models: models(piXaiModelIds) } } } @@ -331,7 +459,7 @@ export async function uninstallPiConfig(): Promise { } return writePiProviders( null, - 'could not parse it as JSON — remove the tokenmaxx-anthropic and tokenmaxx-openai providers yourself' + `could not parse it as JSON — remove the ${piProviderKeys.join(', ')} providers yourself` ) } diff --git a/src/domain.ts b/src/domain.ts index 0677a20..1fb0f68 100644 --- a/src/domain.ts +++ b/src/domain.ts @@ -5,9 +5,15 @@ export type FetchImplementation = ( initialization?: RequestInit ) => Promise -export const ProviderIdSchema = z.enum(['openai', 'anthropic']) +export const ProviderIdSchema = z.enum(['openai', 'anthropic', 'xai']) export type ProviderId = z.infer +export const PROVIDERS: Record = { + anthropic: { app: 'Claude Code', cli: 'claude', vendor: 'Anthropic' }, + openai: { app: 'Codex', cli: 'codex', vendor: 'OpenAI' }, + xai: { app: 'Grok', cli: 'grok', vendor: 'xAI' } +} + export const AccountEmailSchema = z.string().trim().toLowerCase().email() const AccountNameSchema = z.string().trim().min(1) @@ -51,6 +57,12 @@ export const AccountSchema = z profilePath: z.string().trim().min(1).nullable(), provider: z.literal('anthropic'), secretReference: z.string().trim().min(1).nullable() + }).strict(), + AccountFieldsSchema.extend({ + externalUserId: z.null().default(null), + profilePath: z.null(), + provider: z.literal('xai'), + secretReference: z.string().trim().min(1) }).strict() ]) .refine(account => account.label === account.identity, { @@ -117,6 +129,12 @@ export const UsageSnapshotSchema = z.discriminatedUnion('provider', [ measuredSpendUsd: z.number().nonnegative().nullish().default(null), provider: z.literal('anthropic'), source: z.enum(['claudeUsageEndpoint', 'proxyResponseHeaders', 'apiKeyProbe']) + }).strict(), + UsageSnapshotFieldsSchema.extend({ + extraUsage: ExtraUsageSchema.nullish().default(null), + measuredSpendUsd: z.number().nonnegative().nullish().default(null), + provider: z.literal('xai'), + source: z.enum(['grokProbe', 'proxyResponseHeaders', 'apiKeyProbe']) }).strict() ]) export type UsageSnapshot = z.infer diff --git a/src/grok.test.ts b/src/grok.test.ts new file mode 100644 index 0000000..ff84ba7 --- /dev/null +++ b/src/grok.test.ts @@ -0,0 +1,298 @@ +import { describe, expect, test } from 'bun:test' +import type { Account, UsageSnapshot } from './domain.ts' +import { + type GrokAuth, + grokUpstream, + probeGrok, + refreshGrokCredential, + registerGrokAccount +} from './grok.ts' +import { xaiLimitWindow } from './ratelimit.ts' +import type { CredentialVault } from './vault.ts' + +function memoryVault(initial: Record): CredentialVault & { + items: Map +} { + const items = new Map(Object.entries(initial)) + return { + items, + read: async reference => items.get(reference) ?? null, + remove: async reference => { + items.delete(reference) + }, + write: async (reference, value) => { + items.set(reference, value) + } + } +} + +const reference = 'grok:test' +const stored: GrokAuth = { + auth_mode: 'oidc', + email: 'Dexter@RubricLabs.com', + expires_at: '2026-07-20T18:00:00.000Z', + key: 'old-key', + oidc_client_id: 'client-1', + oidc_issuer: 'https://auth.x.ai', + refresh_token: 'old-refresh', + user_id: 'user-1' +} +const authFile = JSON.stringify({ 'https://auth.x.ai::client-1': stored }) + +const account: Extract = { + auth: 'oauth', + createdAt: '2026-07-01T00:00:00.000Z', + enabled: true, + externalAccountId: 'user-1', + externalUserId: null, + health: 'ready', + id: '00000000-0000-4000-8000-000000000001', + identity: 'dexter@rubriclabs.com', + label: 'dexter@rubriclabs.com', + onThreshold: 'switch', + plan: null, + profilePath: null, + provider: 'xai', + secretReference: reference, + updatedAt: '2026-07-01T00:00:00.000Z' +} + +const tokenResponse = () => + Response.json({ access_token: 'new-key', expires_in: 21_600, refresh_token: 'new-refresh' }) + +function vaultWith(credential: GrokAuth = stored) { + return memoryVault({ [reference]: JSON.stringify(credential) }) +} + +describe('registerGrokAccount', () => { + test('imports the isolated login and stores the whole record', async () => { + const vault = memoryVault({}) + const removed: string[] = [] + const account = await registerGrokAccount({ + dependencies: { + createTemporaryDirectory: async () => '/tmp/grok-home', + read: async path => { + expect(path).toBe('/tmp/grok-home/auth.json') + return authFile + }, + remove: async path => { + removed.push(path) + }, + run: async (command, environment) => { + expect(command).toEqual(['grok', 'login']) + expect(environment.GROK_HOME).toBe('/tmp/grok-home') + return { exitCode: 0, stderr: '' } + } + }, + vault + }) + expect(account.provider).toBe('xai') + expect(account.identity).toBe('dexter@rubriclabs.com') + expect(account.externalAccountId).toBe('user-1') + expect(account.secretReference).toBe(`grok:${account.id}`) + expect(JSON.parse(vault.items.get(account.secretReference ?? '') ?? '{}')).toEqual(stored) + expect(removed).toEqual(['/tmp/grok-home']) + }) + + test('a failed login surfaces the cli error', async () => { + await expect( + registerGrokAccount({ + dependencies: { + createTemporaryDirectory: async () => '/tmp/grok-home', + read: async () => authFile, + remove: async () => undefined, + run: async () => ({ exitCode: 1, stderr: 'error: subscription required\n' }) + }, + vault: memoryVault({}) + }) + ).rejects.toThrow('grok login: error: subscription required') + }) +}) + +describe('refreshGrokCredential', () => { + test('posts the refresh grant and stores the rotated tokens', async () => { + const vault = vaultWith() + let body = '' + const refreshed = await refreshGrokCredential({ + fetchImplementation: async (input, initialization) => { + expect(String(input)).toBe('https://auth.x.ai/oauth2/token') + body = String(initialization?.body) + return tokenResponse() + }, + reference, + vault + }) + expect(new URLSearchParams(body).get('grant_type')).toBe('refresh_token') + expect(new URLSearchParams(body).get('client_id')).toBe('client-1') + expect(new URLSearchParams(body).get('refresh_token')).toBe('old-refresh') + expect(refreshed.key).toBe('new-key') + expect(refreshed.refresh_token).toBe('new-refresh') + expect(refreshed.user_id).toBe('user-1') + expect(JSON.parse(vault.items.get(reference) ?? '{}').key).toBe('new-key') + }) + + test('a rejected refresh token asks for a new login', async () => { + await expect( + refreshGrokCredential({ + fetchImplementation: async () => new Response('', { status: 401 }), + reference, + vault: vaultWith() + }) + ).rejects.toMatchObject({ code: 'REAUTHENTICATION_REQUIRED' }) + }) + + test('a stale caller does not refresh a credential someone else already rotated', async () => { + let calls = 0 + const refreshed = await refreshGrokCredential({ + fetchImplementation: async () => { + calls += 1 + return tokenResponse() + }, + reference, + staleKey: 'not-the-current-key', + vault: vaultWith() + }) + expect(calls).toBe(0) + expect(refreshed.key).toBe('old-key') + }) +}) + +describe('grokUpstream', () => { + test('a session account carries the cli token header to the chat proxy', async () => { + const injection = await grokUpstream({ + account, + forceRefresh: false, + now: () => Date.parse('2026-07-20T12:00:00.000Z'), + vault: vaultWith() + }) + expect(injection.baseUrl).toBe('https://cli-chat-proxy.grok.com') + expect(injection.headers.authorization).toBe('Bearer old-key') + expect(injection.headers['x-xai-token-auth']).toBe('xai-grok-cli') + }) + + test('a token inside the refresh margin is refreshed first', async () => { + const injection = await grokUpstream({ + account, + fetchImplementation: async () => tokenResponse(), + forceRefresh: false, + now: () => Date.parse('2026-07-20T17:59:00.000Z'), + vault: vaultWith() + }) + expect(injection.headers.authorization).toBe('Bearer new-key') + }) + + test('an api key account routes to the public api with its key', async () => { + const injection = await grokUpstream({ + account: { ...account, auth: 'apiKey', secretReference: 'grok-key:1' }, + forceRefresh: false, + vault: memoryVault({ 'grok-key:1': 'xai-test-123' }) + }) + expect(injection.baseUrl).toBe('https://api.x.ai') + expect(injection.headers.authorization).toBe('Bearer xai-test-123') + expect(injection.stripHeaders).toContain('x-xai-token-auth') + }) +}) + +describe('probeGrok', () => { + const userinfo = async () => Response.json({ email: 'dexter@rubriclabs.com', sub: 'user-1' }) + const now = () => new Date('2026-07-20T12:00:00.000Z') + const existing = (window: UsageSnapshot['windows'][number]): UsageSnapshot => ({ + accountId: account.id, + extraUsage: null, + hardLimitReached: true, + measuredSpendUsd: null, + observedAt: '2026-07-20T11:59:00.000Z', + provider: 'xai', + source: 'proxyResponseHeaders', + windows: [window] + }) + + test('reports an empty limit window when nothing is held', async () => { + const result = await probeGrok({ + account, + existing: null, + fetchImplementation: userinfo, + now, + vault: vaultWith() + }) + expect(result.account.health).toBe('ready') + expect(result.usage.source).toBe('grokProbe') + expect(result.usage.hardLimitReached).toBe(false) + expect(result.usage.windows).toEqual([xaiLimitWindow(0, null)]) + }) + + test('keeps a limit whose reset is still ahead', async () => { + const held = xaiLimitWindow(100, '2026-07-20T12:30:00.000Z') + const result = await probeGrok({ + account, + existing: existing(held), + fetchImplementation: userinfo, + now, + vault: vaultWith() + }) + expect(result.usage.hardLimitReached).toBe(true) + expect(result.usage.windows).toEqual([held]) + }) + + test('clears a limit once its reset has passed, and one with no reset at all', async () => { + for (const resetAt of ['2026-07-20T11:00:00.000Z', null]) { + const result = await probeGrok({ + account, + existing: existing(xaiLimitWindow(100, resetAt)), + fetchImplementation: userinfo, + now, + vault: vaultWith() + }) + expect(result.usage.hardLimitReached).toBe(false) + expect(result.usage.windows).toEqual([xaiLimitWindow(0, null)]) + } + }) + + test('a rejected token is refreshed once and the identity re-verified', async () => { + let userinfoCalls = 0 + const vault = memoryVault({ 'grok:rejected': JSON.stringify(stored) }) + const result = await probeGrok({ + account: { ...account, secretReference: 'grok:rejected' }, + existing: null, + fetchImplementation: async input => { + if (String(input).endsWith('/oauth2/token')) { + return tokenResponse() + } + userinfoCalls += 1 + return userinfoCalls === 1 ? new Response('', { status: 401 }) : userinfo() + }, + now, + vault + }) + expect(result.account.health).toBe('ready') + expect(userinfoCalls).toBe(2) + expect(JSON.parse(vault.items.get('grok:rejected') ?? '{}').key).toBe('new-key') + }) + + test('a credential for another account is refused', async () => { + await expect( + probeGrok({ + account: { ...account, secretReference: 'grok:other' }, + existing: null, + fetchImplementation: async () => Response.json({ email: 'someone@else.com' }), + now, + vault: memoryVault({ 'grok:other': JSON.stringify(stored) }) + }) + ).rejects.toMatchObject({ code: 'IDENTITY_CHANGED' }) + }) + + test('an api key account is validated against the models endpoint', async () => { + const result = await probeGrok({ + account: { ...account, auth: 'apiKey', secretReference: 'grok-key:1' }, + existing: null, + fetchImplementation: async input => { + expect(String(input)).toBe('https://api.x.ai/v1/models') + return Response.json({ data: [] }) + }, + now, + vault: memoryVault({ 'grok-key:1': 'xai-test-123' }) + }) + expect(result.usage.source).toBe('apiKeyProbe') + expect(result.usage.windows).toEqual([]) + }) +}) diff --git a/src/grok.ts b/src/grok.ts new file mode 100644 index 0000000..95a4c24 --- /dev/null +++ b/src/grok.ts @@ -0,0 +1,335 @@ +import { join } from 'node:path' +import { z } from 'zod' +import { + type Account, + AccountEmailSchema, + type FetchImplementation, + type ProviderProbeResult, + type UsageSnapshot, + type UsageWindow +} from './domain.ts' +import { ApplicationError, loginFailureMessage } from './errors.ts' +import { defaultIsolatedLoginDependencies, type IsolatedLoginDependencies } from './login.ts' +import { type UpstreamInjection, upstreamFor } from './proxy.ts' +import { xaiLimitWindow, xaiLimitWindowId } from './ratelimit.ts' +import { type CredentialVault, exclusive, readApiKey } from './vault.ts' + +const tokenEndpoint = 'https://auth.x.ai/oauth2/token' +const userinfoEndpoint = 'https://auth.x.ai/oauth2/userinfo' +const xaiApiBase = 'https://api.x.ai' +const sessionAuthHeaders = { 'x-xai-token-auth': 'xai-grok-cli' } +const refreshMarginMilliseconds = 120_000 + +const GrokAuthSchema = z + .object({ + email: z.string().min(1), + expires_at: z.iso.datetime(), + key: z.string().min(1), + oidc_client_id: z.string().min(1), + refresh_token: z.string().min(1), + user_id: z.string().min(1) + }) + .passthrough() +export type GrokAuth = z.infer + +const GrokAuthFileSchema = z.record(z.string(), z.unknown()) + +const TokenResponseSchema = z + .object({ + access_token: z.string().min(1), + expires_in: z.number().positive(), + refresh_token: z.string().min(1).optional() + }) + .passthrough() + +const UserinfoSchema = z.object({ email: z.string().optional() }).passthrough() + +function grokAuthFromFile(serialized: string): GrokAuth { + const entries = Object.values(GrokAuthFileSchema.parse(JSON.parse(serialized))) + for (const entry of entries) { + const parsed = GrokAuthSchema.safeParse(entry) + if (parsed.success) { + return parsed.data + } + } + throw new ApplicationError( + 'CREDENTIAL_MISSING', + 'grok login left no usable credential in auth.json' + ) +} + +export async function registerGrokAccount(input: { + vault: CredentialVault + dependencies?: IsolatedLoginDependencies +}): Promise { + const dependencies = input.dependencies ?? defaultIsolatedLoginDependencies() + const temporaryHome = await dependencies.createTemporaryDirectory('tokenmaxx-grok-') + try { + const login = await dependencies.run(['grok', 'login'], { GROK_HOME: temporaryHome }) + if (login.exitCode !== 0) { + throw new ApplicationError('LOGIN_FAILED', loginFailureMessage('grok login', login)) + } + const auth = grokAuthFromFile(await dependencies.read(join(temporaryHome, 'auth.json'))) + const email = AccountEmailSchema.safeParse(auth.email) + if (!email.success) { + throw new ApplicationError( + 'ACCOUNT_EMAIL_MISSING', + 'Grok did not return a verified account email; the login was not stored' + ) + } + const id = crypto.randomUUID() + const secretReference = `grok:${id}` + await input.vault.write(secretReference, JSON.stringify(auth)) + const now = new Date().toISOString() + return { + auth: 'oauth', + createdAt: now, + enabled: true, + externalAccountId: auth.user_id, + externalUserId: null, + health: 'ready', + id, + identity: email.data, + label: email.data, + onThreshold: 'switch', + plan: null, + profilePath: null, + provider: 'xai', + secretReference, + updatedAt: now + } + } finally { + await dependencies.remove(temporaryHome) + } +} + +async function readGrokCredential(vault: CredentialVault, reference: string): Promise { + const serialized = await vault.read(reference) + if (serialized === null) { + throw new ApplicationError('CREDENTIAL_MISSING', `Missing credential ${reference}`) + } + return GrokAuthSchema.parse(JSON.parse(serialized)) +} + +export async function refreshGrokCredential(input: { + reference: string + vault: CredentialVault + fetchImplementation?: FetchImplementation + staleKey?: string +}): Promise { + return exclusive(input.reference, async () => { + const current = await readGrokCredential(input.vault, input.reference) + if (input.staleKey !== undefined && current.key !== input.staleKey) { + return current + } + const response = await (input.fetchImplementation ?? fetch)(tokenEndpoint, { + body: new URLSearchParams({ + client_id: current.oidc_client_id, + grant_type: 'refresh_token', + refresh_token: current.refresh_token + }), + method: 'POST', + signal: AbortSignal.timeout(7_000) + }) + if (response.status === 400 || response.status === 401) { + throw new ApplicationError('REAUTHENTICATION_REQUIRED', 'Grok refresh token was rejected') + } + if (!response.ok) { + throw new ApplicationError( + 'PROVIDER_UNREACHABLE', + `Grok token refresh returned HTTP ${response.status}` + ) + } + const refreshed = TokenResponseSchema.parse(await response.json()) + const updated = GrokAuthSchema.parse({ + ...current, + expires_at: new Date(Date.now() + refreshed.expires_in * 1000).toISOString(), + key: refreshed.access_token, + refresh_token: refreshed.refresh_token ?? current.refresh_token + }) + await input.vault.write(input.reference, JSON.stringify(updated)) + return updated + }) +} + +async function validateXaiApiKey( + key: string, + fetchImplementation: FetchImplementation +): Promise { + const response = await fetchImplementation(`${xaiApiBase}/v1/models`, { + headers: { Authorization: `Bearer ${key}` }, + signal: AbortSignal.timeout(10_000) + }) + if (response.status === 401 || response.status === 403) { + throw new ApplicationError('ACCESS_TOKEN_REJECTED', 'xAI rejected this API key') + } + if (!response.ok && response.status !== 429) { + throw new ApplicationError('PROVIDER_UNREACHABLE', `xAI API returned HTTP ${response.status}`) + } +} + +export async function registerXaiApiKeyAccount(input: { + vault: CredentialVault + key: string + label: string + fetchImplementation?: FetchImplementation +}): Promise { + const key = input.key.trim() + if (key.length === 0) { + throw new ApplicationError('USAGE', 'The API key is empty') + } + await validateXaiApiKey(key, input.fetchImplementation ?? fetch) + const id = crypto.randomUUID() + const secretReference = `grok-key:${id}` + await input.vault.write(secretReference, key) + const now = new Date().toISOString() + return { + auth: 'apiKey', + createdAt: now, + enabled: true, + externalAccountId: null, + externalUserId: null, + health: 'ready', + id, + identity: input.label, + label: input.label, + onThreshold: 'switch', + plan: null, + profilePath: null, + provider: 'xai', + secretReference, + updatedAt: now + } +} + +export async function grokUpstream(input: { + account: Extract + vault: CredentialVault + fetchImplementation?: FetchImplementation + now?: () => number + forceRefresh: boolean +}): Promise { + const reference = input.account.secretReference + if (input.account.auth === 'apiKey') { + return { + accountId: input.account.id, + baseUrl: xaiApiBase, + headers: { authorization: `Bearer ${await readApiKey(input.vault, reference)}` }, + stripHeaders: ['x-xai-token-auth'] + } + } + let auth = await readGrokCredential(input.vault, reference) + const now = input.now ?? (() => Date.now()) + const stale = Date.parse(auth.expires_at) - now() <= refreshMarginMilliseconds + if (input.forceRefresh || stale) { + auth = await refreshGrokCredential({ + fetchImplementation: input.fetchImplementation, + reference, + vault: input.vault + }) + } + return { + accountId: input.account.id, + baseUrl: upstreamFor('xai'), + headers: { authorization: `Bearer ${auth.key}`, ...sessionAuthHeaders } + } +} + +async function fetchGrokEmail( + key: string, + fetchImplementation: FetchImplementation +): Promise { + const response = await fetchImplementation(userinfoEndpoint, { + headers: { Authorization: `Bearer ${key}` }, + signal: AbortSignal.timeout(7_000) + }) + if (response.status === 401) { + throw new ApplicationError('REAUTHENTICATION_REQUIRED', 'Grok credential was rejected') + } + if (!response.ok) { + throw new ApplicationError( + 'PROVIDER_UNREACHABLE', + `Grok userinfo endpoint returned HTTP ${response.status}` + ) + } + return UserinfoSchema.parse(await response.json()).email ?? null +} + +const verifiedIdentities = new Map() + +function heldLimit(existing: UsageSnapshot | null, now: Date): UsageWindow | null { + const held = existing?.windows.find(window => window.id === xaiLimitWindowId) + return held !== undefined && held.resetAt !== null && Date.parse(held.resetAt) > now.getTime() + ? held + : null +} + +export async function probeGrok(input: { + account: Extract + vault: CredentialVault + fetchImplementation: FetchImplementation + now(): Date + existing: UsageSnapshot | null +}): Promise { + const { account, vault, fetchImplementation } = input + const reference = account.secretReference + if (account.auth === 'apiKey') { + await validateXaiApiKey(await readApiKey(vault, reference), fetchImplementation) + return { + account: { ...account, health: 'ready', updatedAt: input.now().toISOString() }, + usage: { + accountId: account.id, + extraUsage: null, + hardLimitReached: false, + measuredSpendUsd: null, + observedAt: input.now().toISOString(), + provider: 'xai', + source: 'apiKeyProbe', + windows: [] + } + } + } + const refresh = (staleKey: string) => + refreshGrokCredential({ fetchImplementation, reference, staleKey, vault }) + let credential = await readGrokCredential(vault, reference) + if (Date.parse(credential.expires_at) <= input.now().getTime() + 300_000) { + credential = await refresh(credential.key) + } + const verifyIdentity = async (): Promise => { + const cached = verifiedIdentities.get(reference) + if (cached !== undefined && cached.key === credential.key) { + return cached.email + } + const email = await fetchGrokEmail(credential.key, fetchImplementation) + verifiedIdentities.set(reference, { email, key: credential.key }) + return email + } + const email = await verifyIdentity().catch(async error => { + if (!(error instanceof ApplicationError) || error.code !== 'REAUTHENTICATION_REQUIRED') { + throw error + } + credential = await refresh(credential.key) + return verifyIdentity() + }) + const verified = AccountEmailSchema.safeParse(email) + if (verified.success && verified.data !== account.identity) { + throw new ApplicationError( + 'IDENTITY_CHANGED', + 'Stored Grok credential belongs to a different account' + ) + } + const held = heldLimit(input.existing, input.now()) + return { + account: { ...account, health: 'ready', updatedAt: input.now().toISOString() }, + usage: { + accountId: account.id, + extraUsage: null, + hardLimitReached: held !== null, + measuredSpendUsd: null, + observedAt: input.now().toISOString(), + provider: 'xai', + source: 'grokProbe', + windows: [held ?? xaiLimitWindow(0, null)] + } + } +} diff --git a/src/login.ts b/src/login.ts new file mode 100644 index 0000000..8dfa146 --- /dev/null +++ b/src/login.ts @@ -0,0 +1,36 @@ +import { mkdtemp, readFile, rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' + +export interface IsolatedLoginDependencies { + run( + command: readonly string[], + environment: Record + ): Promise<{ exitCode: number; stderr: string }> + createTemporaryDirectory(prefix: string): Promise + read(path: string): Promise + remove(path: string): Promise +} + +export function defaultIsolatedLoginDependencies(): IsolatedLoginDependencies { + return { + createTemporaryDirectory: prefix => mkdtemp(join(tmpdir(), prefix)), + read: path => readFile(path, 'utf8'), + remove: path => rm(path, { force: true, recursive: true }), + async run(command, environment) { + const child = Bun.spawn([...command], { + env: { ...process.env, ...environment }, + stderr: 'pipe', + stdin: 'inherit', + stdout: 'inherit' + }) + const decoder = new TextDecoder() + let stderr = '' + for await (const chunk of child.stderr) { + process.stderr.write(chunk) + stderr = `${stderr}${decoder.decode(chunk, { stream: true })}`.slice(-4_096) + } + return { exitCode: await child.exited, stderr } + } + } +} diff --git a/src/manager.ts b/src/manager.ts index 3079ce5..3eee5d2 100644 --- a/src/manager.ts +++ b/src/manager.ts @@ -9,7 +9,9 @@ import { type Account, AutomationPolicySchema, type FetchImplementation, + PROVIDERS, type ProviderId, + ProviderIdSchema, type ProviderState, type ResetCreditsView, type ResetOutcome, @@ -20,6 +22,7 @@ import { type UsageWindow } from './domain.ts' import { ApplicationError, errorMessage, isNetworkFailure } from './errors.ts' +import { grokUpstream, probeGrok } from './grok.ts' import type { ApplicationPaths } from './paths.ts' import { costUsd } from './pricing.ts' import { @@ -63,7 +66,7 @@ function priceTokenTimeframe(aggregate: TokenTimeframeAggregate): TokenTimeframe const byProvider = new Map() const models: (TokenBreakdownAccumulator & { model: string; provider: ProviderId })[] = [] for (const entry of aggregate.byModel) { - const provider: ProviderId = entry.provider === 'anthropic' ? 'anthropic' : 'openai' + const provider = ProviderIdSchema.catch('openai').parse(entry.provider) const priced: TokenBreakdownAccumulator = { cacheCreation: entry.cacheCreation, cached: entry.cached, @@ -180,9 +183,14 @@ export class AccountManager { vault: this.#vault } try { - return account.provider === 'openai' - ? await codexUpstream({ account, ...shared }) - : await claudeUpstream({ account, ...shared }) + switch (account.provider) { + case 'openai': + return await codexUpstream({ account, ...shared }) + case 'anthropic': + return await claudeUpstream({ account, ...shared }) + case 'xai': + return await grokUpstream({ account, ...shared }) + } } catch (error) { const cause = error instanceof Error ? error : undefined if (isNetworkFailure(error)) { @@ -192,10 +200,9 @@ export class AccountManager { { cause } ) } - const cli = provider === 'openai' ? 'codex' : 'claude' throw new ApplicationError( 'ACTIVE_CREDENTIAL_UNUSABLE', - `${account.label} needs re-login — run: tokenmaxx login ${cli}`, + `${account.label} needs re-login — run: tokenmaxx login ${PROVIDERS[provider].cli}`, { cause } ) } @@ -369,10 +376,16 @@ export class AccountManager { now: () => this.#dependencies.now(), vault: this.#vault } - const result = - account.provider === 'anthropic' - ? await probeClaude({ account, ...shared }) - : await probeCodex({ account, ...shared }) + const result = await (() => { + switch (account.provider) { + case 'openai': + return probeCodex({ account, ...shared }) + case 'anthropic': + return probeClaude({ account, ...shared }) + case 'xai': + return probeGrok({ account, existing: this.#store.findUsage(account.id), ...shared }) + } + })() if (account.auth === 'apiKey') { const start = this.#dependencies.now().getTime() - 31 * 24 * 3_600_000 result.usage.measuredSpendUsd = this.#store @@ -515,7 +528,7 @@ export class AccountManager { if (this.#stopping) { return } - for (const provider of ['openai', 'anthropic'] as const) { + for (const provider of ProviderIdSchema.options) { if (this.#stopping) { return } diff --git a/src/paths.ts b/src/paths.ts index 93ea20d..d9eb64e 100644 --- a/src/paths.ts +++ b/src/paths.ts @@ -2,6 +2,7 @@ import { chmod, mkdir } from 'node:fs/promises' import { homedir } from 'node:os' import { join, resolve } from 'node:path' import { z } from 'zod' +import type { ProviderId } from './domain.ts' const ApplicationPathsSchema = z.object({ claudeProfiles: z.string().min(1), @@ -38,6 +39,6 @@ export async function ensureApplicationPaths(paths: ApplicationPaths): Promise chmod(directory, 0o700))) } -export function proxyBaseUrl(paths: ApplicationPaths, provider: 'openai' | 'anthropic'): string { +export function proxyBaseUrl(paths: ApplicationPaths, provider: ProviderId): string { return `http://127.0.0.1:${paths.proxyPort}/${provider}` } diff --git a/src/pricing.ts b/src/pricing.ts index b2858c8..2b55675 100644 --- a/src/pricing.ts +++ b/src/pricing.ts @@ -41,6 +41,22 @@ const PRICES: ReadonlyArray<{ match: string; price: ModelPrice }> = [ { match: 'o4', price: { cacheReadPerMTok: 0.275, cacheWritePerMTok: 0, inputPerMTok: 1.1, outputPerMTok: 4.4 } + }, + { + match: 'grok-4.6', + price: { cacheReadPerMTok: 0.5, cacheWritePerMTok: 0, inputPerMTok: 2, outputPerMTok: 6 } + }, + { + match: 'grok-4.5', + price: { cacheReadPerMTok: 0.3, cacheWritePerMTok: 0, inputPerMTok: 2, outputPerMTok: 6 } + }, + { + match: 'grok-build', + price: { cacheReadPerMTok: 0.2, cacheWritePerMTok: 0, inputPerMTok: 1, outputPerMTok: 2 } + }, + { + match: 'grok-4', + price: { cacheReadPerMTok: 0.2, cacheWritePerMTok: 0, inputPerMTok: 1.25, outputPerMTok: 2.5 } } ] diff --git a/src/proxy.ts b/src/proxy.ts index 50ab7cd..c91140c 100644 --- a/src/proxy.ts +++ b/src/proxy.ts @@ -1,4 +1,4 @@ -import type { FetchImplementation, ProviderId } from './domain.ts' +import { type FetchImplementation, PROVIDERS, type ProviderId, ProviderIdSchema } from './domain.ts' import { ApplicationError, errorMessage, isNetworkFailure } from './errors.ts' import { observeRateLimitHeaders, type RateLimitObservation } from './ratelimit.ts' @@ -283,12 +283,14 @@ const strippedRequestHeaders = [ ] const strippedResponseHeaders = ['content-encoding', 'content-length', 'transfer-encoding'] +const providerRoute = new RegExp(`^/(${ProviderIdSchema.options.join('|')})(/.*)?$`) + function routeProvider(pathname: string): { provider: ProviderId; rest: string } | null { - const match = pathname.match(/^\/(openai|anthropic)(\/.*)?$/) + const match = pathname.match(providerRoute) if (match === null) { return null } - return { provider: match[1] as ProviderId, rest: match[2] ?? '/' } + return { provider: ProviderIdSchema.parse(match[1]), rest: match[2] ?? '/' } } function forwardHeaders(incoming: Headers, injection: UpstreamInjection): Headers { @@ -387,7 +389,7 @@ function createProxyHandler(options: ProxyOptions): ProxyHandler { return { deliver, observation } } - const providerLabel = route.provider === 'anthropic' ? 'Anthropic' : 'OpenAI' + const providerLabel = PROVIDERS[route.provider].vendor let injection: UpstreamInjection | null try { injection = await options.source.resolve(route.provider) @@ -536,6 +538,7 @@ export function upstreamFor(provider: ProviderId): string { return 'https://chatgpt.com/backend-api/codex' case 'anthropic': return 'https://api.anthropic.com' + case 'xai': + return 'https://cli-chat-proxy.grok.com' } - throw new ApplicationError('UNKNOWN_PROVIDER', `No upstream for provider ${provider}`) } diff --git a/src/ratelimit.test.ts b/src/ratelimit.test.ts index 76886ff..9f9ece3 100644 --- a/src/ratelimit.test.ts +++ b/src/ratelimit.test.ts @@ -83,5 +83,26 @@ describe('observeRateLimitHeaders', () => { expect( observeRateLimitHeaders('openai', new Headers({ 'content-type': 'text/plain' }), 200) ).toBeNull() + expect(observeRateLimitHeaders('xai', new Headers(), 200)).toBeNull() + }) + + test('xai 429 fills the synthetic limit window and reads retry-after', () => { + const before = Date.now() + const observation = observeRateLimitHeaders('xai', new Headers({ 'retry-after': '90' }), 429) + expect(observation?.limited).toBe(true) + expect(observation?.windows).toHaveLength(1) + const window = observation?.windows[0] + expect(window?.id).toBe('limit') + expect(window?.kind).toBe('hard') + expect(window?.usedPercent).toBe(100) + const resetAt = Date.parse(window?.resetAt ?? '') + expect(resetAt).toBeGreaterThanOrEqual(before + 90_000) + expect(resetAt).toBeLessThanOrEqual(Date.now() + 90_000) + }) + + test('xai 429 without retry-after has no reset time', () => { + const observation = observeRateLimitHeaders('xai', new Headers(), 429) + expect(observation?.limited).toBe(true) + expect(observation?.windows[0]?.resetAt).toBeNull() }) }) diff --git a/src/ratelimit.ts b/src/ratelimit.ts index d71e00c..2b0c879 100644 --- a/src/ratelimit.ts +++ b/src/ratelimit.ts @@ -103,12 +103,35 @@ function codexObservation(headers: Headers, status: number): RateLimitObservatio return { limited: status === 429, windows } } +export const xaiLimitWindowId = 'limit' + +export function xaiLimitWindow(usedPercent: 0 | 100, resetAt: string | null): UsageWindow { + return { id: xaiLimitWindowId, kind: 'hard', label: 'rate limit', resetAt, usedPercent } +} + +function xaiObservation(headers: Headers, status: number): RateLimitObservation | null { + if (status !== 429) { + return null + } + const seconds = Number(headers.get('retry-after')) + const resetAt = + Number.isFinite(seconds) && seconds > 0 + ? new Date(Date.now() + seconds * 1000).toISOString() + : null + return { limited: true, windows: [xaiLimitWindow(100, resetAt)] } +} + export function observeRateLimitHeaders( provider: ProviderId, headers: Headers, status: number ): RateLimitObservation | null { - return provider === 'anthropic' - ? anthropicObservation(headers, status) - : codexObservation(headers, status) + switch (provider) { + case 'anthropic': + return anthropicObservation(headers, status) + case 'openai': + return codexObservation(headers, status) + case 'xai': + return xaiObservation(headers, status) + } } diff --git a/src/storage.ts b/src/storage.ts index c98daad..8bbc2aa 100644 --- a/src/storage.ts +++ b/src/storage.ts @@ -264,6 +264,9 @@ function migrate(database: Database): void { CREATE UNIQUE INDEX IF NOT EXISTS accounts_anthropic_external ON accounts(external_account_id) WHERE provider = 'anthropic' AND external_account_id IS NOT NULL; + CREATE UNIQUE INDEX IF NOT EXISTS accounts_xai_external + ON accounts(external_account_id) + WHERE provider = 'xai' AND external_account_id IS NOT NULL; `) } catch (error) { throw new ApplicationError( @@ -338,8 +341,7 @@ export function createStateStore(databasePath: string): StateStore { parsed.externalUserId !== null && candidate.externalAccountId === parsed.externalAccountId && candidate.externalUserId === parsed.externalUserId) || - (parsed.provider === 'anthropic' && - candidate.provider === 'anthropic' && + (parsed.provider !== 'openai' && parsed.externalAccountId !== null && candidate.externalAccountId === parsed.externalAccountId)) ) diff --git a/src/tui/dashboard.ts b/src/tui/dashboard.ts index 61209ad..7868777 100644 --- a/src/tui/dashboard.ts +++ b/src/tui/dashboard.ts @@ -1,16 +1,18 @@ import { Box, createCliRenderer, parseColor, type RGBA, Text } from '@opentui/core' import { installPiConfig, type PiStatus, piStatus, uninstallPiConfig } from '../config-install.ts' -import type { - Account, - AnalyticsSnapshot, - DashboardSnapshot, - ProviderId, - ProviderState, - ResetCreditsView, - ResetOutcome, - TokenTimeframe, - UsageSnapshot, - UsageWindow +import { + type Account, + type AnalyticsSnapshot, + type DashboardSnapshot, + PROVIDERS, + type ProviderId, + ProviderIdSchema, + type ProviderState, + type ResetCreditsView, + type ResetOutcome, + type TokenTimeframe, + type UsageSnapshot, + type UsageWindow } from '../domain.ts' import { readAnalytics, @@ -64,13 +66,11 @@ function rgb(hex: string): RGBA { return value } -const providerTitles: Record = { - anthropic: 'Anthropic · Claude Code', - openai: 'OpenAI · Codex' -} -const providerShort: Record = { anthropic: 'Claude Code', openai: 'Codex' } -const providerCli: Record = { anthropic: 'claude', openai: 'codex' } -const providerOrder: readonly ProviderId[] = ['openai', 'anthropic'] +const providerOrder = ProviderIdSchema.options +const providerTitle = (provider: ProviderId) => + `${PROVIDERS[provider].vendor} · ${PROVIDERS[provider].app}` +const providerShort = (provider: ProviderId) => PROVIDERS[provider].app +const providerCli = (provider: ProviderId) => PROVIDERS[provider].cli const fallbackTimeframe = TIMEFRAMES[2] as Timeframe interface Row { @@ -287,11 +287,11 @@ function addAccountLine(ctx: Ctx, provider: ProviderId, isSelected: boolean, sol Text({ content: ` ${isSelected ? '▸' : '＋'} `, fg: rgb(color) }), Text({ attributes: sole ? 1 : 0, - content: `add a ${providerShort[provider]} account`, + content: `add a ${providerShort(provider)} account`, fg: rgb(color) }), Text({ - content: installed ? ' ⏎' : ` · install ${providerCli[provider]} first`, + content: installed ? ' ⏎' : ` · install ${providerCli(provider)} first`, fg: rgb(installed ? ctx.theme.faint : ctx.theme.warn) }) ) @@ -416,8 +416,8 @@ function providerPanel( const routed = ctx.routing[provider] const auto = state?.policy.enabled ? `auto ${state.policy.thresholdPercent}%` : 'auto off' const title = routed - ? ` ${providerTitles[provider]} ● ${auto} ` - : ` ${providerTitles[provider]} ✗ off ` + ? ` ${providerTitle(provider)} ● ${auto} ` + : ` ${providerTitle(provider)} ✗ off ` const titleColor = !routed ? ctx.theme.warn : state?.policy.enabled @@ -429,7 +429,7 @@ function providerPanel( Box( { flexDirection: 'row', width: '100%' }, Text({ - content: ` tokenmaxx is off for ${providerCli[provider]} — turn it on in settings`, + content: ` tokenmaxx is off for ${providerCli(provider)} — turn it on in settings`, fg: rgb(ctx.theme.warn) }) ) @@ -500,7 +500,7 @@ function sessionResets(ctx: Ctx, snapshot: DashboardSnapshot): string | null { return [] } const label = account.label.length <= 22 ? account.label : `${account.label.slice(0, 21)}…` - return [`${providerCli[provider]} · ${label} · ↻ ${reset}`] + return [`${providerCli(provider)} · ${label} · ↻ ${reset}`] }) return parts.length === 0 ? null : parts.join(' ') } @@ -604,7 +604,7 @@ function metricsView(ctx: Ctx, tokens: TokenTimeframe, scroll: number) { for (const provider of tokens.byProvider) { body.push( metricRow( - { color: ctx.theme.fg, text: providerShort[provider.provider] }, + { color: ctx.theme.fg, text: providerShort(provider.provider) }, [ { color: ctx.theme.dim, text: num(provider.input) }, { color: ctx.theme.dim, text: num(provider.output) }, @@ -885,7 +885,7 @@ function settingsPanel( : 'shown' const hint = row.key === 'routing' - ? `run ${providerCli[row.provider]} through tokenmaxx` + ? `run ${providerCli(row.provider)} through tokenmaxx` : row.key === 'auto' ? 'switch accounts as the active one fills' : row.key === 'threshold' @@ -919,8 +919,8 @@ function settingsPanel( flexDirection: 'column', flexShrink: 0, title: routed - ? ` ${providerTitles[provider]} ${auto} ` - : ` ${providerTitles[provider]} ✗ off `, + ? ` ${providerTitle(provider)} ${auto} ` + : ` ${providerTitle(provider)} ✗ off `, titleColor: rgb(!routed ? ctx.theme.warn : policy?.enabled ? ctx.theme.good : ctx.theme.dim), width: '100%' }, @@ -988,8 +988,7 @@ function settingsBody(ctx: Ctx, snapshot: DashboardSnapshot, rows: SettingRow[], return column( ctx, [ - settingsPanel(ctx, snapshot, rows, 'openai', selected), - settingsPanel(ctx, snapshot, rows, 'anthropic', selected), + ...providerOrder.map(provider => settingsPanel(ctx, snapshot, rows, provider, selected)), displayPanel(ctx, rows, selected), harnessPanel(ctx, rows, selected) ], @@ -1003,8 +1002,7 @@ function accountsBody(ctx: Ctx, snapshot: DashboardSnapshot, rows: Row[], select return column( ctx, [ - providerPanel(ctx, snapshot, 'openai', rows, selected), - providerPanel(ctx, snapshot, 'anthropic', rows, selected), + ...providerOrder.map(provider => providerPanel(ctx, snapshot, provider, rows, selected)), ...(note === null ? [] : [note]) ], width + 2 @@ -1022,7 +1020,7 @@ interface AddConfirm { } function addConfirmBody(ctx: Ctx, confirm: AddConfirm) { - const cli = providerCli[confirm.provider] + const cli = providerCli(confirm.provider) const installed = ctx.cliPresent[confirm.provider] const line = (...children: ReturnType[]) => Box( @@ -1051,7 +1049,7 @@ function addConfirmBody(ctx: Ctx, confirm: AddConfirm) { borderColor: rgb(ctx.theme.accent), borderStyle: 'rounded', flexDirection: 'column', - title: ` Add a ${providerShort[confirm.provider]} account `, + title: ` Add a ${providerShort(confirm.provider)} account `, titleColor: rgb(ctx.theme.accent), width: '100%' }, @@ -1320,9 +1318,9 @@ export async function runTuiDashboard( try { process.stdin.setRawMode?.(true) } catch {} - const cliPresent: Record = live - ? { anthropic: Bun.which('claude') !== null, openai: Bun.which('codex') !== null } - : { anthropic: true, openai: true } + const cliPresent = Object.fromEntries( + providerOrder.map(provider => [provider, !live || Bun.which(providerCli(provider)) !== null]) + ) as Record const renderer = await createCliRenderer({ exitOnCtrlC: false, targetFps: 30 }) await renderer.waitForThemeMode(400).catch(() => null) const themeEnvironmentOverride = themeOverride(process.env) @@ -1560,7 +1558,7 @@ export async function runTuiDashboard( applyPolicy( provider, { enabled: enable }, - `auto-rotate ${providerCli[provider]} ${enable ? 'on' : 'off'}…` + `auto-rotate ${providerCli(provider)} ${enable ? 'on' : 'off'}…` ) } diff --git a/src/tui/fixtures.ts b/src/tui/fixtures.ts index c48e7de..733e788 100644 --- a/src/tui/fixtures.ts +++ b/src/tui/fixtures.ts @@ -31,11 +31,12 @@ function buildTokens(scale: number): TokenAnalytics { const buckets = raw.map(value => Math.round((value / rawSum) * target)) const totalTokens = buckets.reduce((sum, value) => sum + value, 0) const modelMix: { model: string; provider: ProviderId; share: number }[] = [ - { model: 'gpt-5.6-sol', provider: 'openai', share: 0.42 }, - { model: 'gpt-5.6-codex', provider: 'openai', share: 0.13 }, - { model: 'claude-opus-4-8', provider: 'anthropic', share: 0.3 }, - { model: 'claude-sonnet-4-6', provider: 'anthropic', share: 0.11 }, - { model: 'claude-haiku-4-5', provider: 'anthropic', share: 0.04 } + { model: 'gpt-5.6-sol', provider: 'openai', share: 0.38 }, + { model: 'gpt-5.6-codex', provider: 'openai', share: 0.12 }, + { model: 'claude-opus-4-8', provider: 'anthropic', share: 0.27 }, + { model: 'claude-sonnet-4-6', provider: 'anthropic', share: 0.1 }, + { model: 'claude-haiku-4-5', provider: 'anthropic', share: 0.04 }, + { model: 'grok-4.6', provider: 'xai', share: 0.09 } ] const models = modelMix .map(entry => { @@ -178,21 +179,32 @@ function account(seed: AccountSeed, now: number): Account { plan: seed.plan, updatedAt: new Date(now - 2 * MINUTE).toISOString() } as const - return seed.provider === 'openai' - ? { + switch (seed.provider) { + case 'openai': + return { ...base, externalUserId: `user_${seed.n}`, profilePath: null, provider: 'openai', secretReference: `codex:${base.externalAccountId}` } - : { + case 'anthropic': + return { ...base, externalUserId: null, profilePath: `/tmp/tokenmaxx/claude/${seed.n}`, provider: 'anthropic', secretReference: null } + case 'xai': + return { + ...base, + externalUserId: null, + profilePath: null, + provider: 'xai', + secretReference: `grok:${base.externalAccountId}` + } + } } function usage(seed: AccountSeed, now: number): UsageSnapshot { @@ -207,14 +219,19 @@ function usage(seed: AccountSeed, now: number): UsageSnapshot { ).toISOString(), windows } as const - return seed.provider === 'openai' - ? { + switch (seed.provider) { + case 'openai': + return { ...base, provider: 'openai', resetCredits: seed.resetCredits ?? null, source: 'codexUsageEndpoint' } - : { ...base, provider: 'anthropic', source: 'claudeUsageEndpoint' } + case 'anthropic': + return { ...base, provider: 'anthropic', source: 'claudeUsageEndpoint' } + case 'xai': + return { ...base, provider: 'xai', source: 'grokProbe' } + } } function policy( @@ -310,6 +327,16 @@ const claudeSession = (peak: number, nowFrac: number, seed: number): WindowSpec ...fiveHour(peak, nowFrac, seed), label: '5h session' }) +const grokLimit = (peak: number, nowFrac: number, seed: number): WindowSpec => ({ + fillFrac: 0.01, + id: 'limit', + label: 'rate limit', + nowFrac, + peak, + period: HOUR, + seed, + wobble: 0 +}) type ScenarioBuilder = (now: number) => AnalyticsSnapshot @@ -369,11 +396,26 @@ const cruising: ScenarioBuilder = now => n: 5, plan: null, provider: 'anthropic' + }, + { + email: 'dexter@rubriclabs.com', + n: 6, + plan: null, + provider: 'xai', + windows: [grokLimit(0, 0, 61)] + }, + { + email: 'ship@rubriclabs.com', + n: 7, + plan: null, + provider: 'xai', + windows: [grokLimit(100, 0.4, 62)] } ], [ { activeN: 1, auto: true, generation: 4, provider: 'openai', switchedMinutesAgo: 96 }, - { activeN: 3, auto: true, generation: 2, provider: 'anthropic', switchedMinutesAgo: 210 } + { activeN: 3, auto: true, generation: 2, provider: 'anthropic', switchedMinutesAgo: 210 }, + { activeN: 6, auto: true, generation: 1, provider: 'xai', switchedMinutesAgo: 12 } ] ) diff --git a/src/ui.ts b/src/ui.ts index fc09c0f..9357176 100644 --- a/src/ui.ts +++ b/src/ui.ts @@ -1,10 +1,12 @@ -import type { - Account, - DashboardSnapshot, - ProviderId, - ProviderState, - UsageSnapshot, - UsageWindow +import { + type Account, + type DashboardSnapshot, + PROVIDERS, + type ProviderId, + ProviderIdSchema, + type ProviderState, + type UsageSnapshot, + type UsageWindow } from './domain.ts' interface RenderOptions { @@ -107,16 +109,11 @@ function shortReset(resetAt: string | null, now: Date): string | null { } function providerTitle(provider: ProviderId): string { - switch (provider) { - case 'openai': - return 'OpenAI · Codex' - case 'anthropic': - return 'Anthropic · Claude Code' - } + return `${PROVIDERS[provider].vendor} · ${PROVIDERS[provider].app}` } function providerCliName(provider: ProviderId): string { - return provider === 'openai' ? 'codex' : 'claude' + return PROVIDERS[provider].cli } function sampleAge(observedAt: string, now: Date): string | null { @@ -255,10 +252,10 @@ export function renderDashboard( return [ header, '', - providerSection(paint, snapshot, 'openai', now), - '', - providerSection(paint, snapshot, 'anthropic', now), - '', + ...ProviderIdSchema.options.flatMap(provider => [ + providerSection(paint, snapshot, provider, now), + '' + ]), paint('● active — every request uses it · q quit · r refresh · tokenmaxx --help', 'dim') ].join('\n') } diff --git a/src/vault.ts b/src/vault.ts index aa2a8e7..8b0a1a6 100644 --- a/src/vault.ts +++ b/src/vault.ts @@ -6,6 +6,14 @@ export interface CredentialVault { remove(reference: string): Promise } +export async function readApiKey(vault: CredentialVault, reference: string): Promise { + const key = await vault.read(reference) + if (key === null) { + throw new ApplicationError('CREDENTIAL_MISSING', `Missing credential ${reference}`) + } + return key +} + const credentialLocks = new Map>() export async function exclusive(