feat(providers): add xAI to hosted key rotation pool (#5574)

* feat(providers): add xAI to hosted key rotation pool

Wires xai into the same hosted-key mechanism as OpenAI, Anthropic,
and Z.ai so Sim can serve Grok models without users bringing their
own key.

* fix(pi): include xai in Pi cloud-mode workspace BYOK read-back

xai was fully wired as a Pi-supported provider but missing from
WORKSPACE_BYOK_PROVIDERS, so a stored workspace xAI key was never
read back for cloud-mode Pi runs.

* fix(byok): add xai settings UI row

xai is both hosted (Pi block hides its inline API key field for
hosted models) and Pi-supported (cloud mode requires a user key),
so without a Settings > BYOK row users had no way to supply an xai
key for Pi cloud runs.
This commit is contained in:
Waleed
2026-07-11 14:23:27 -04:00
committed by GitHub
parent def2d5299a
commit 7c2de1d426
13 changed files with 74 additions and 7 deletions
@@ -34,6 +34,7 @@ import {
SerperIcon,
TogetherIcon,
WizaIcon,
xAIIcon,
ZeroBounceIcon,
} from '@/components/icons'
import { MAX_BYOK_KEYS_PER_PROVIDER } from '@/lib/api/contracts/byok-keys'
@@ -76,6 +77,13 @@ const PROVIDERS: (BYOKManagerProvider & { id: BYOKProviderId })[] = [
description: 'LLM calls and Knowledge Base OCR',
placeholder: 'Enter your API key',
},
{
id: 'xai',
name: 'xAI',
icon: xAIIcon,
description: 'LLM calls',
placeholder: 'xai-...',
},
{
id: 'fireworks',
name: 'Fireworks',
@@ -287,6 +295,7 @@ const PROVIDER_SECTIONS: BYOKProviderSection[] = [
'anthropic',
'google',
'mistral',
'xai',
'fireworks',
'together',
'baseten',
@@ -134,6 +134,21 @@ describe('resolvePiModelKey', () => {
expect(mockGetApiKeyWithBYOK).not.toHaveBeenCalled()
})
it('cloud mode falls back to a stored workspace key for xAI', async () => {
mockGetProviderFromModel.mockReturnValue('xai')
mockGetBYOKKey.mockResolvedValue({ apiKey: 'xai-workspace-key', isBYOK: true })
const result = await resolvePiModelKey({
model: 'grok-4.5',
mode: 'cloud',
workspaceId: 'ws-1',
})
expect(result).toEqual({ providerId: 'xai', apiKey: 'xai-workspace-key', isBYOK: true })
expect(mockGetBYOKKey).toHaveBeenCalledWith('ws-1', 'xai')
expect(mockGetApiKeyWithBYOK).not.toHaveBeenCalled()
})
it('cloud mode rejects when no user key is available (never a hosted key)', async () => {
mockGetProviderFromModel.mockReturnValue('anthropic')
mockGetBYOKKey.mockResolvedValue(null)
+7 -1
View File
@@ -33,7 +33,13 @@ interface ResolvePiModelKeyParams {
}
/** Providers whose key Sim can store as a workspace BYOK key (read back for cloud). */
const WORKSPACE_BYOK_PROVIDERS = new Set<string>(['anthropic', 'openai', 'google', 'mistral'])
const WORKSPACE_BYOK_PROVIDERS = new Set<string>([
'anthropic',
'openai',
'google',
'mistral',
'xai',
])
/** Resolves the provider and a usable API key for the selected model. */
export async function resolvePiModelKey(params: ResolvePiModelKeyParams): Promise<PiKeyResolution> {
+2 -1
View File
@@ -205,13 +205,14 @@ export async function getApiKeyWithBYOK(
const isGeminiModel = provider === 'google'
const isMistralModel = provider === 'mistral'
const isZaiModel = provider === 'zai'
const isXaiModel = provider === 'xai'
const byokProviderId = isGeminiModel ? 'google' : (provider as BYOKProviderId)
if (
isHosted &&
workspaceId &&
(isOpenAIModel || isClaudeModel || isGeminiModel || isMistralModel || isZaiModel)
(isOpenAIModel || isClaudeModel || isGeminiModel || isMistralModel || isZaiModel || isXaiModel)
) {
const hostedModels = getHostedModels()
const isModelHosted = hostedModels.some((m) => m.toLowerCase() === model.toLowerCase())
+1
View File
@@ -7,6 +7,7 @@ export const byokProviderIdSchema = z.enum([
'google',
'mistral',
'zai',
'xai',
'fireworks',
'together',
'baseten',
+6 -1
View File
@@ -12,7 +12,8 @@ export function getRotatingApiKey(provider: string): string {
provider !== 'anthropic' &&
provider !== 'gemini' &&
provider !== 'cohere' &&
provider !== 'zai'
provider !== 'zai' &&
provider !== 'xai'
) {
throw new Error(`No rotation implemented for provider: ${provider}`)
}
@@ -39,6 +40,10 @@ export function getRotatingApiKey(provider: string): string {
if (env.ZAI_API_KEY_1) keys.push(env.ZAI_API_KEY_1)
if (env.ZAI_API_KEY_2) keys.push(env.ZAI_API_KEY_2)
if (env.ZAI_API_KEY_3) keys.push(env.ZAI_API_KEY_3)
} else if (provider === 'xai') {
if (env.XAI_API_KEY_1) keys.push(env.XAI_API_KEY_1)
if (env.XAI_API_KEY_2) keys.push(env.XAI_API_KEY_2)
if (env.XAI_API_KEY_3) keys.push(env.XAI_API_KEY_3)
}
if (keys.length === 0) {
+3
View File
@@ -147,6 +147,9 @@ export const env = createEnv({
ZAI_API_KEY_1: z.string().min(1).optional(), // Primary Z.ai API key for load balancing
ZAI_API_KEY_2: z.string().min(1).optional(), // Additional Z.ai API key for load balancing
ZAI_API_KEY_3: z.string().min(1).optional(), // Additional Z.ai API key for load balancing
XAI_API_KEY_1: z.string().min(1).optional(), // Primary xAI API key for load balancing
XAI_API_KEY_2: z.string().min(1).optional(), // Additional xAI API key for load balancing
XAI_API_KEY_3: z.string().min(1).optional(), // Additional xAI API key for load balancing
OLLAMA_URL: z.string().url().optional(), // Ollama local LLM server URL
VLLM_BASE_URL: z.string().url().optional(), // vLLM self-hosted base URL (OpenAI-compatible)
VLLM_API_KEY: z.string().optional(), // Optional bearer token for vLLM
+8
View File
@@ -43,6 +43,9 @@ vi.mock('@/lib/core/config/env', () =>
GEMINI_API_KEY_1: 'test-gemini-key-1',
GEMINI_API_KEY_2: 'test-gemini-key-2',
GEMINI_API_KEY_3: 'test-gemini-key-3',
XAI_API_KEY_1: 'test-xai-key-1',
XAI_API_KEY_2: 'test-xai-key-2',
XAI_API_KEY_3: 'test-xai-key-3',
})
)
@@ -327,6 +330,11 @@ describe('getRotatingApiKey', () => {
expect(result).toMatch(/^test-gemini-key-[1-3]$/)
})
it.concurrent('should return xAI API key based on current minute', () => {
const result = getRotatingApiKey('xai')
expect(result).toMatch(/^test-xai-key-[1-3]$/)
})
it.concurrent('should throw error for unsupported provider', () => {
expect(() => getRotatingApiKey('unsupported')).toThrow('No rotation implemented for provider')
})
+14
View File
@@ -218,3 +218,17 @@ describe('zai provider definition', () => {
expect(getHostedModels()).toContain('glm-4.6')
})
})
describe('xai provider definition', () => {
const xai = PROVIDER_DEFINITIONS.xai
it('is registered with grok-4.5 as the default model', () => {
expect(xai).toBeDefined()
expect(xai.id).toBe('xai')
expect(xai.defaultModel).toBe('grok-4.5')
})
it('is included in getHostedModels since Sim provides the xAI key server-side', () => {
expect(getHostedModels()).toContain('grok-4.5')
})
})
+1
View File
@@ -3921,6 +3921,7 @@ export function getHostedModels(): string[] {
...getProviderModels('anthropic'),
...getProviderModels('google'),
...getProviderModels('zai'),
...getProviderModels('xai'),
]
}
+5 -3
View File
@@ -865,7 +865,7 @@ describe('Cost Calculation', () => {
})
describe('getHostedModels', () => {
it.concurrent('should return OpenAI, Anthropic, and Google models as hosted', () => {
it.concurrent('should return OpenAI, Anthropic, Google, and xAI models as hosted', () => {
const hostedModels = getHostedModels()
expect(hostedModels).toContain('gpt-4o')
@@ -877,8 +877,9 @@ describe('getHostedModels', () => {
expect(hostedModels).toContain('gemini-2.5-pro')
expect(hostedModels).toContain('gemini-2.5-flash')
expect(hostedModels).toContain('grok-4.5')
expect(hostedModels).not.toContain('deepseek-v3')
expect(hostedModels).not.toContain('grok-4-latest')
})
it.concurrent('should return an array of strings', () => {
@@ -902,11 +903,12 @@ describe('shouldBillModelUsage', () => {
expect(shouldBillModelUsage('gemini-2.5-pro')).toBe(true)
expect(shouldBillModelUsage('gemini-2.5-flash')).toBe(true)
expect(shouldBillModelUsage('grok-4.5')).toBe(true)
})
it.concurrent('should return false for non-hosted models', () => {
expect(shouldBillModelUsage('deepseek-v3')).toBe(false)
expect(shouldBillModelUsage('grok-4-latest')).toBe(false)
expect(shouldBillModelUsage('unknown-model')).toBe(false)
})
+2 -1
View File
@@ -903,8 +903,9 @@ export function getApiKey(provider: string, model: string, userProvidedKey?: str
const isClaudeModel = provider === 'anthropic'
const isGeminiModel = provider === 'google'
const isZaiModel = provider === 'zai'
const isXaiModel = provider === 'xai'
if (isHosted && (isOpenAIModel || isClaudeModel || isGeminiModel || isZaiModel)) {
if (isHosted && (isOpenAIModel || isClaudeModel || isGeminiModel || isZaiModel || isXaiModel)) {
const hostedModels = getHostedModels()
const isModelHosted = hostedModels.some((m) => m.toLowerCase() === model.toLowerCase())
+1
View File
@@ -8,6 +8,7 @@ export type BYOKProviderId =
| 'google'
| 'mistral'
| 'zai'
| 'xai'
| 'fireworks'
| 'together'
| 'baseten'