Skip to content

Commit 70baa6f

Browse files
authored
feat: add Grok (xAI) API adapter with custom model mapping (#152)
Add xAI Grok as a new API provider. Reuses OpenAI-compatible message/tool converters and stream adapter with Grok-specific client and model mapping. Default model mapping: opus → grok-4.20-reasoning sonnet → grok-3-mini-fast haiku → grok-3-mini-fast Users can customize mapping via: - GROK_MODEL env var (override all) - GROK_MODEL_MAP env var (JSON family map, e.g. {"opus":"grok-4"}) - GROK_DEFAULT_{FAMILY}_MODEL env vars Activation: CLAUDE_CODE_USE_GROK=1 or modelType: "grok" in settings.json Also integrates with /provider command for runtime switching.
1 parent dfa7aa1 commit 70baa6f

9 files changed

Lines changed: 488 additions & 6 deletions

File tree

src/commands/provider.ts

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@ function getEnvVarForProvider(provider: string): string {
1515
return 'CLAUDE_CODE_USE_FOUNDRY'
1616
case 'gemini':
1717
return 'CLAUDE_CODE_USE_GEMINI'
18+
case 'grok':
19+
return 'CLAUDE_CODE_USE_GROK'
1820
default:
1921
throw new Error(`Unknown provider: ${provider}`)
2022
}
@@ -48,6 +50,7 @@ const call: LocalCommandCall = async (args, context) => {
4850
delete process.env.CLAUDE_CODE_USE_FOUNDRY
4951
delete process.env.CLAUDE_CODE_USE_OPENAI
5052
delete process.env.CLAUDE_CODE_USE_GEMINI
53+
delete process.env.CLAUDE_CODE_USE_GROK
5154
return {
5255
type: 'text',
5356
value: 'API provider cleared (will use environment variables).',
@@ -59,6 +62,7 @@ const call: LocalCommandCall = async (args, context) => {
5962
'anthropic',
6063
'openai',
6164
'gemini',
65+
'grok',
6266
'bedrock',
6367
'vertex',
6468
'foundry',
@@ -87,6 +91,19 @@ const call: LocalCommandCall = async (args, context) => {
8791
}
8892
}
8993

94+
// Check env vars when switching to grok (including settings.env)
95+
if (arg === 'grok') {
96+
const mergedEnv = getMergedEnv()
97+
const hasKey = !!(mergedEnv.GROK_API_KEY || mergedEnv.XAI_API_KEY)
98+
if (!hasKey) {
99+
updateSettingsForSource('userSettings', { modelType: 'grok' })
100+
return {
101+
type: 'text',
102+
value: `Switched to Grok provider.\nWarning: Missing env var: GROK_API_KEY (or XAI_API_KEY)\nConfigure it via settings.json env or set manually.`,
103+
}
104+
}
105+
}
106+
90107
// Check env vars when switching to gemini (including settings.env)
91108
if (arg === 'gemini') {
92109
const mergedEnv = getMergedEnv()
@@ -104,13 +121,14 @@ const call: LocalCommandCall = async (args, context) => {
104121
// Handle different provider types
105122
// - 'anthropic', 'openai', 'gemini' are stored in settings.json (persistent)
106123
// - 'bedrock', 'vertex', 'foundry' are env-only (do NOT touch settings.json)
107-
if (arg === 'anthropic' || arg === 'openai' || arg === 'gemini') {
124+
if (arg === 'anthropic' || arg === 'openai' || arg === 'gemini' || arg === 'grok') {
108125
// Clear any cloud provider env vars to avoid conflicts
109126
delete process.env.CLAUDE_CODE_USE_BEDROCK
110127
delete process.env.CLAUDE_CODE_USE_VERTEX
111128
delete process.env.CLAUDE_CODE_USE_FOUNDRY
112129
delete process.env.CLAUDE_CODE_USE_OPENAI
113130
delete process.env.CLAUDE_CODE_USE_GEMINI
131+
delete process.env.CLAUDE_CODE_USE_GROK
114132
// Update settings.json
115133
updateSettingsForSource('userSettings', { modelType: arg })
116134
// Ensure settings.env gets applied to process.env
@@ -122,6 +140,7 @@ const call: LocalCommandCall = async (args, context) => {
122140
delete process.env.OPENAI_API_KEY
123141
delete process.env.OPENAI_BASE_URL
124142
delete process.env.CLAUDE_CODE_USE_GEMINI
143+
delete process.env.CLAUDE_CODE_USE_GROK
125144
process.env[getEnvVarForProvider(arg)] = '1'
126145
// Do not modify settings.json - cloud providers controlled solely by env vars
127146
applyConfigEnvironmentVariables()
@@ -136,9 +155,9 @@ const provider = {
136155
type: 'local',
137156
name: 'provider',
138157
description:
139-
'Switch API provider (anthropic/openai/gemini/bedrock/vertex/foundry)',
158+
'Switch API provider (anthropic/openai/gemini/grok/bedrock/vertex/foundry)',
140159
aliases: ['api'],
141-
argumentHint: '[anthropic|openai|gemini|bedrock|vertex|foundry|unset]',
160+
argumentHint: '[anthropic|openai|gemini|grok|bedrock|vertex|foundry|unset]',
142161
supportsNonInteractive: true,
143162
load: () => Promise.resolve({ call }),
144163
} satisfies Command

src/services/api/claude.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1350,6 +1350,12 @@ async function* queryModel(
13501350
return
13511351
}
13521352

1353+
if (getAPIProvider() === 'grok') {
1354+
const { queryModelGrok } = await import('./grok/index.js')
1355+
yield* queryModelGrok(messagesForAPI, systemPrompt, filteredTools, signal, options)
1356+
return
1357+
}
1358+
13531359
// Instrumentation: Track message count after normalization
13541360
logEvent('tengu_api_after_normalize', {
13551361
postNormalizedMessageCount: messagesForAPI.length,
Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
import { describe, expect, test, beforeEach, afterEach } from 'bun:test'
2+
import { getGrokClient, clearGrokClientCache } from '../client.js'
3+
4+
describe('getGrokClient', () => {
5+
const originalEnv = { ...process.env }
6+
7+
beforeEach(() => {
8+
clearGrokClientCache()
9+
process.env.GROK_API_KEY = 'test-key'
10+
delete process.env.GROK_BASE_URL
11+
})
12+
13+
afterEach(() => {
14+
clearGrokClientCache()
15+
process.env = { ...originalEnv }
16+
})
17+
18+
test('creates client with default base URL', () => {
19+
const client = getGrokClient()
20+
expect(client).toBeDefined()
21+
expect(client.baseURL).toBe('https://api.x.ai/v1')
22+
})
23+
24+
test('uses GROK_BASE_URL when set', () => {
25+
process.env.GROK_BASE_URL = 'https://custom.grok.api/v1'
26+
clearGrokClientCache()
27+
const client = getGrokClient()
28+
expect(client.baseURL).toBe('https://custom.grok.api/v1')
29+
})
30+
31+
test('returns cached client on second call', () => {
32+
const client1 = getGrokClient()
33+
const client2 = getGrokClient()
34+
expect(client1).toBe(client2)
35+
})
36+
37+
test('clearGrokClientCache resets cache', () => {
38+
const client1 = getGrokClient()
39+
clearGrokClientCache()
40+
process.env.GROK_BASE_URL = 'https://other.api/v1'
41+
const client2 = getGrokClient()
42+
expect(client1).not.toBe(client2)
43+
})
44+
})
Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
import { describe, expect, test, beforeEach, afterEach } from 'bun:test'
2+
import { resolveGrokModel } from '../modelMapping.js'
3+
4+
describe('resolveGrokModel', () => {
5+
const originalEnv = { ...process.env }
6+
7+
beforeEach(() => {
8+
delete process.env.GROK_MODEL
9+
delete process.env.GROK_MODEL_MAP
10+
delete process.env.GROK_DEFAULT_SONNET_MODEL
11+
delete process.env.GROK_DEFAULT_OPUS_MODEL
12+
delete process.env.GROK_DEFAULT_HAIKU_MODEL
13+
delete process.env.ANTHROPIC_DEFAULT_SONNET_MODEL
14+
delete process.env.ANTHROPIC_DEFAULT_OPUS_MODEL
15+
delete process.env.ANTHROPIC_DEFAULT_HAIKU_MODEL
16+
})
17+
18+
afterEach(() => {
19+
process.env = { ...originalEnv }
20+
})
21+
22+
test('GROK_MODEL env var takes highest priority', () => {
23+
process.env.GROK_MODEL = 'grok-custom'
24+
expect(resolveGrokModel('claude-sonnet-4-6')).toBe('grok-custom')
25+
})
26+
27+
test('maps opus models to grok-4.20-reasoning', () => {
28+
expect(resolveGrokModel('claude-opus-4-6')).toBe('grok-4.20-reasoning')
29+
})
30+
31+
test('maps sonnet models to grok-3-mini-fast', () => {
32+
expect(resolveGrokModel('claude-sonnet-4-6')).toBe('grok-3-mini-fast')
33+
})
34+
35+
test('maps haiku models to grok-3-mini-fast', () => {
36+
expect(resolveGrokModel('claude-haiku-4-5-20251001')).toBe('grok-3-mini-fast')
37+
})
38+
39+
test('GROK_MODEL_MAP overrides family mapping', () => {
40+
process.env.GROK_MODEL_MAP = '{"opus":"grok-4","sonnet":"grok-3","haiku":"grok-mini"}'
41+
expect(resolveGrokModel('claude-opus-4-6')).toBe('grok-4')
42+
expect(resolveGrokModel('claude-sonnet-4-6')).toBe('grok-3')
43+
expect(resolveGrokModel('claude-haiku-4-5-20251001')).toBe('grok-mini')
44+
})
45+
46+
test('GROK_MODEL_MAP ignores invalid JSON', () => {
47+
process.env.GROK_MODEL_MAP = 'not-json'
48+
expect(resolveGrokModel('claude-opus-4-6')).toBe('grok-4.20-reasoning')
49+
})
50+
51+
test('GROK_DEFAULT_{FAMILY}_MODEL overrides default map', () => {
52+
process.env.GROK_DEFAULT_OPUS_MODEL = 'grok-2-latest'
53+
expect(resolveGrokModel('claude-opus-4-6')).toBe('grok-2-latest')
54+
})
55+
56+
test('passes through unknown model names', () => {
57+
expect(resolveGrokModel('some-unknown-model')).toBe('some-unknown-model')
58+
})
59+
60+
test('strips [1m] suffix before lookup', () => {
61+
expect(resolveGrokModel('claude-sonnet-4-6[1m]')).toBe('grok-3-mini-fast')
62+
})
63+
64+
test('falls back to family default for unlisted model', () => {
65+
expect(resolveGrokModel('claude-opus-99-20300101')).toBe('grok-4.20-reasoning')
66+
})
67+
})

src/services/api/grok/client.ts

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
import OpenAI from 'openai'
2+
import { getProxyFetchOptions } from 'src/utils/proxy.js'
3+
4+
/**
5+
* Environment variables:
6+
*
7+
* GROK_API_KEY (or XAI_API_KEY): Required. API key for the xAI Grok endpoint.
8+
* GROK_BASE_URL: Optional. Defaults to https://api.x.ai/v1.
9+
*/
10+
11+
const DEFAULT_BASE_URL = 'https://api.x.ai/v1'
12+
13+
let cachedClient: OpenAI | null = null
14+
15+
export function getGrokClient(options?: {
16+
maxRetries?: number
17+
fetchOverride?: typeof fetch
18+
source?: string
19+
}): OpenAI {
20+
if (cachedClient) return cachedClient
21+
22+
const apiKey = process.env.GROK_API_KEY || process.env.XAI_API_KEY || ''
23+
const baseURL = process.env.GROK_BASE_URL || DEFAULT_BASE_URL
24+
25+
const client = new OpenAI({
26+
apiKey,
27+
baseURL,
28+
maxRetries: options?.maxRetries ?? 0,
29+
timeout: parseInt(process.env.API_TIMEOUT_MS || String(600 * 1000), 10),
30+
dangerouslyAllowBrowser: true,
31+
fetchOptions: getProxyFetchOptions({ forAnthropicAPI: false }) as RequestInit,
32+
...(options?.fetchOverride && { fetch: options.fetchOverride }),
33+
})
34+
35+
if (!options?.fetchOverride) {
36+
cachedClient = client
37+
}
38+
39+
return client
40+
}
41+
42+
export function clearGrokClientCache(): void {
43+
cachedClient = null
44+
}

0 commit comments

Comments
 (0)