add in api key auth on proxy endpoint and pass it into the docs agent handler

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<void> {
   // 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<void> {
         // Map tool calls to message types
         let messageType: DocsAgentMessage['messageType'];
         let content = '';
-        let metadata = '';
+        const metadata = '';
 
         switch (part.toolName) {
           case 'sequentialThinking':

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<typeof ProxyConfigSchema>;
@@ -13,23 +14,40 @@ export type ProxyConfig = z.infer<typeof ProxyConfigSchema>;
  * 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<ProxyConfig> {
   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) {
+  if (creds.apiUrl) {
-    return ProxyConfigSchema.parse({ baseURL: 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,
+  });
 }

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) => {
+export const POST = new Hono().post(
-  try {
+  '/',
-    const { model, options } = c.req.valid('json');
+  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
+      // Get the gateway model
-    const modelInstance = gatewayModel(model);
+      const modelInstance = gatewayModel(model);
 
-    // Call the model's doStream method directly (this is a model-level operation)
+      // Call the model's doStream method directly (this is a model-level operation)
-    const result = await modelInstance.doStream(options);
+      const result = await modelInstance.doStream(options);
 
-    // Stream the LanguageModelV2StreamPart objects
+      // Stream the LanguageModelV2StreamPart objects
-    return stream(c, async (stream) => {
+      return stream(c, async (stream) => {
-      try {
+        try {
-        const reader = result.stream.getReader();
+          const reader = result.stream.getReader();
-        while (true) {
+          while (true) {
-          const { done, value } = await reader.read();
+            const { done, value } = await reader.read();
-          if (done) break;
+            if (done) break;
-          await stream.write(`${JSON.stringify(value)}\n`);
+            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);
+    } catch (error) {
-        throw streamError;
+      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);
   }
-});
+);

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 {
+import { bashExecute, createIdleTool } from '../../tools';
-  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<LanguageModelV2>().optional().describe('Custom language model to use (defaults to Sonnet4)'),
+  model: z
+    .custom<LanguageModelV2>()
+    .optional()
+    .describe('Custom language model to use (defaults to Sonnet4)'),
 });
 
 const DocsStreamOptionsSchema = z.object({

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<typeof ProxyModelConfigSchema>;
@@ -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,