diff --git a/apps/cli/src/services/docs-agent-handler.ts b/apps/cli/src/services/docs-agent-handler.ts index bbc1a5b44..fc5720ecf 100644 --- a/apps/cli/src/services/docs-agent-handler.ts +++ b/apps/cli/src/services/docs-agent-handler.ts @@ -1,7 +1,7 @@ +import { randomUUID } from 'node:crypto'; import { createDocsAgent } from '@buster/ai/agents/docs-agent/docs-agent'; import { createProxyModel } from '@buster/ai/llm/providers/proxy-model'; import type { ModelMessage } from 'ai'; -import { randomUUID } from 'node:crypto'; import { getProxyConfig } from '../utils/ai-proxy'; export interface DocsAgentMessage { @@ -32,6 +32,7 @@ export async function runDocsAgent(params: RunDocsAgentParams): Promise { // Create proxy model that routes through server const proxyModel = createProxyModel({ baseURL: proxyConfig.baseURL, + apiKey: proxyConfig.apiKey, modelId: 'anthropic/claude-4-sonnet-20250514', }); @@ -71,7 +72,7 @@ export async function runDocsAgent(params: RunDocsAgentParams): Promise { // Map tool calls to message types let messageType: DocsAgentMessage['messageType']; let content = ''; - let metadata = ''; + const metadata = ''; switch (part.toolName) { case 'sequentialThinking': diff --git a/apps/cli/src/utils/ai-proxy.ts b/apps/cli/src/utils/ai-proxy.ts index 5c3e20866..fad40a6eb 100644 --- a/apps/cli/src/utils/ai-proxy.ts +++ b/apps/cli/src/utils/ai-proxy.ts @@ -2,6 +2,7 @@ import { z } from 'zod'; const ProxyConfigSchema = z.object({ baseURL: z.string().url().describe('Base URL for the AI proxy endpoint'), + apiKey: z.string().min(1).describe('API key for authentication'), }); export type ProxyConfig = z.infer; @@ -13,23 +14,40 @@ export type ProxyConfig = z.infer; * 1. BUSTER_AI_PROXY_URL environment variable * 2. Saved credentials apiUrl from ~/.buster/credentials.json * 3. Default to localhost:3002 for local development + * + * API key comes from credentials (required) */ export async function getProxyConfig(): Promise { const { getCredentials } = await import('./credentials'); const creds = await getCredentials(); + if (!creds?.apiKey) { + throw new Error( + 'API key not found. Please run "buster login" or set BUSTER_API_KEY environment variable' + ); + } + // Check for AI proxy-specific URL (highest priority) const proxyUrl = process.env.BUSTER_AI_PROXY_URL; if (proxyUrl) { - return ProxyConfigSchema.parse({ baseURL: proxyUrl }); + return ProxyConfigSchema.parse({ + baseURL: proxyUrl, + apiKey: creds.apiKey, + }); } // Fall back to regular API URL from credentials - if (creds?.apiUrl) { - return ProxyConfigSchema.parse({ baseURL: creds.apiUrl }); + if (creds.apiUrl) { + return ProxyConfigSchema.parse({ + baseURL: creds.apiUrl, + apiKey: creds.apiKey, + }); } // Default to localhost for development - return ProxyConfigSchema.parse({ baseURL: 'http://localhost:3002' }); + return ProxyConfigSchema.parse({ + baseURL: 'http://localhost:3002', + apiKey: creds.apiKey, + }); } diff --git a/apps/server/src/api/v2/llm/proxy/POST.ts b/apps/server/src/api/v2/llm/proxy/POST.ts index 631cc6139..e05be1e2a 100644 --- a/apps/server/src/api/v2/llm/proxy/POST.ts +++ b/apps/server/src/api/v2/llm/proxy/POST.ts @@ -3,40 +3,46 @@ import { zValidator } from '@hono/zod-validator'; import { Hono } from 'hono'; import { stream } from 'hono/streaming'; import { z } from 'zod'; +import { createApiKeyAuthMiddleware } from '../../../../middleware/api-key-auth'; const ProxyRequestSchema = z.object({ model: z.string().describe('Model ID to use'), options: z.any().describe('LanguageModelV2CallOptions from AI SDK'), }); -export const POST = new Hono().post('/', zValidator('json', ProxyRequestSchema), async (c) => { - try { - const { model, options } = c.req.valid('json'); +export const POST = new Hono().post( + '/', + createApiKeyAuthMiddleware(), + zValidator('json', ProxyRequestSchema), + async (c) => { + try { + const { model, options } = c.req.valid('json'); - console.info('[PROXY] Request received', { model }); + console.info('[PROXY] Request received', { model }); - // Get the gateway model - const modelInstance = gatewayModel(model); + // Get the gateway model + const modelInstance = gatewayModel(model); - // Call the model's doStream method directly (this is a model-level operation) - const result = await modelInstance.doStream(options); + // Call the model's doStream method directly (this is a model-level operation) + const result = await modelInstance.doStream(options); - // Stream the LanguageModelV2StreamPart objects - return stream(c, async (stream) => { - try { - const reader = result.stream.getReader(); - while (true) { - const { done, value } = await reader.read(); - if (done) break; - await stream.write(`${JSON.stringify(value)}\n`); + // Stream the LanguageModelV2StreamPart objects + return stream(c, async (stream) => { + try { + const reader = result.stream.getReader(); + while (true) { + const { done, value } = await reader.read(); + if (done) break; + await stream.write(`${JSON.stringify(value)}\n`); + } + } catch (streamError) { + console.error('[PROXY] Stream error:', streamError); + throw streamError; } - } catch (streamError) { - console.error('[PROXY] Stream error:', streamError); - throw streamError; - } - }); - } catch (error) { - console.error('[PROXY] Endpoint error:', error); - return c.json({ error: error instanceof Error ? error.message : 'Unknown error' }, 500); + }); + } catch (error) { + console.error('[PROXY] Endpoint error:', error); + return c.json({ error: error instanceof Error ? error.message : 'Unknown error' }, 500); + } } -}); +); diff --git a/packages/ai/src/agents/docs-agent/docs-agent.ts b/packages/ai/src/agents/docs-agent/docs-agent.ts index e7b8c597d..30b28f9f5 100644 --- a/packages/ai/src/agents/docs-agent/docs-agent.ts +++ b/packages/ai/src/agents/docs-agent/docs-agent.ts @@ -1,14 +1,11 @@ -import type { Sandbox } from '@buster/sandbox'; import type { LanguageModelV2 } from '@ai-sdk/provider'; +import type { Sandbox } from '@buster/sandbox'; import { type ModelMessage, hasToolCall, stepCountIs, streamText } from 'ai'; import { wrapTraced } from 'braintrust'; import z from 'zod'; import { DEFAULT_ANTHROPIC_OPTIONS } from '../../llm/providers/gateway'; import { Sonnet4 } from '../../llm/sonnet-4'; -import { - bashExecute, - createIdleTool, -} from '../../tools'; +import { bashExecute, createIdleTool } from '../../tools'; import { type AgentContext, repairToolCall } from '../../utils/tool-call-repair'; import { getDocsAgentSystemPrompt } from './get-docs-agent-system-prompt'; @@ -31,7 +28,10 @@ const DocsAgentOptionsSchema = z.object({ { message: 'Invalid Sandbox instance' } ) .optional(), - model: z.custom().optional().describe('Custom language model to use (defaults to Sonnet4)'), + model: z + .custom() + .optional() + .describe('Custom language model to use (defaults to Sonnet4)'), }); const DocsStreamOptionsSchema = z.object({ diff --git a/packages/ai/src/llm/providers/proxy-model.ts b/packages/ai/src/llm/providers/proxy-model.ts index 987b657a2..e6b34ef5e 100644 --- a/packages/ai/src/llm/providers/proxy-model.ts +++ b/packages/ai/src/llm/providers/proxy-model.ts @@ -8,6 +8,7 @@ import { z } from 'zod'; const ProxyModelConfigSchema = z.object({ baseURL: z.string().describe('Base URL of the proxy server'), modelId: z.string().describe('Model ID to proxy requests to'), + apiKey: z.string().min(1).describe('API key for authentication'), }); type ProxyModelConfig = z.infer; @@ -34,7 +35,10 @@ export function createProxyModel(config: ProxyModelConfig): LanguageModelV2 { async doGenerate(options: LanguageModelV2CallOptions) { const response = await fetch(`${validated.baseURL}/api/v2/llm/proxy`, { method: 'POST', - headers: { 'Content-Type': 'application/json' }, + headers: { + 'Content-Type': 'application/json', + Authorization: `Bearer ${validated.apiKey}`, + }, body: JSON.stringify({ model: validated.modelId, options, @@ -139,7 +143,10 @@ export function createProxyModel(config: ProxyModelConfig): LanguageModelV2 { async doStream(options: LanguageModelV2CallOptions) { const response = await fetch(`${validated.baseURL}/api/v2/llm/proxy`, { method: 'POST', - headers: { 'Content-Type': 'application/json' }, + headers: { + 'Content-Type': 'application/json', + Authorization: `Bearer ${validated.apiKey}`, + }, body: JSON.stringify({ model: validated.modelId, options,