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

This commit is contained in:
dal 2025-09-30 20:59:53 -06:00
parent 95ca83d441
commit cce3eaf009
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
5 changed files with 71 additions and 39 deletions

View File

@ -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':

View File

@ -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) {
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,
});
}

View File

@ -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);
}
}
});
);

View File

@ -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<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({

View File

@ -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,