diff --git a/backend/package.json b/backend/package.json index aeac65984..11cfc0433 100644 --- a/backend/package.json +++ b/backend/package.json @@ -23,6 +23,7 @@ "bun": ">=1.2.11" }, "dependencies": { + "@ai-sdk/anthropic": "1.0.8", "@ai-sdk/google-vertex": "3.0.6", "@ai-sdk/openai": "2.0.11", "@codebuff/billing": "workspace:*", @@ -56,4 +57,4 @@ "@types/express": "^4.17.13", "@types/ws": "^8.5.5" } -} +} \ No newline at end of file diff --git a/backend/src/llm-apis/message-cost-tracker.ts b/backend/src/llm-apis/message-cost-tracker.ts index 68e2cfa9e..bcf2517ef 100644 --- a/backend/src/llm-apis/message-cost-tracker.ts +++ b/backend/src/llm-apis/message-cost-tracker.ts @@ -579,7 +579,8 @@ export const saveMessage = async (value: { cacheReadInputTokens?: number finishedAt: Date latencyMs: number - usesUserApiKey?: boolean + usesUserApiKey?: boolean // Deprecated: use byokProvider instead + byokProvider?: 'anthropic' | 'gemini' | 'openai' | null chargeUser?: boolean costOverrideDollars?: number agentId?: string @@ -604,16 +605,21 @@ export const saveMessage = async (value: { // Default to 1 cent per credit const centsPerCredit = 1 + // Determine if user API key was used (support both old and new parameters) + const usesUserKey = value.byokProvider !== null && value.byokProvider !== undefined + ? !!value.byokProvider + : value.usesUserApiKey ?? false + const costInCents = value.chargeUser ?? true // default to true ? Math.max( - 0, - Math.round( - cost * - 100 * - (value.usesUserApiKey ? PROFIT_MARGIN : 1 + PROFIT_MARGIN), - ), - ) + 0, + Math.round( + cost * + 100 * + (usesUserKey ? PROFIT_MARGIN : 1 + PROFIT_MARGIN), + ), + ) : 0 const creditsUsed = Math.max(0, costInCents) diff --git a/backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts b/backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts index da2f708ba..1b1835e7b 100644 --- a/backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts +++ b/backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts @@ -1,5 +1,7 @@ import { google } from '@ai-sdk/google' import { openai } from '@ai-sdk/openai' +import { createAnthropic } from '@ai-sdk/anthropic' +import { env } from '@codebuff/internal' import { finetunedVertexModels, geminiModels, @@ -32,19 +34,56 @@ import type { import type { LanguageModel } from 'ai' import type { z } from 'zod/v4' +// User API keys for BYOK (Bring Your Own Key) +export interface UserApiKeys { + anthropic?: string + gemini?: string + openai?: string +} + +export type ByokMode = 'disabled' | 'prefer' | 'require' + export type StreamChunk = | { - type: 'text' - text: string - } + type: 'text' + text: string + } | { - type: 'reasoning' - text: string - } + type: 'reasoning' + text: string + } | { type: 'error'; message: string } -// TODO: We'll want to add all our models here! -const modelToAiSDKModel = (model: Model): LanguageModel => { +/** + * Helper function to determine if a model is an Anthropic model + */ +function isAnthropicModel(model: Model): boolean { + return model.startsWith('anthropic/') +} + +/** + * Helper function to determine which provider key was used for BYOK + */ +function determineByokProvider( + model: Model, + userApiKeys?: UserApiKeys, +): 'anthropic' | 'gemini' | 'openai' | null { + if (isAnthropicModel(model) && userApiKeys?.anthropic) return 'anthropic' + if (Object.values(geminiModels).includes(model as GeminiModel) && userApiKeys?.gemini) return 'gemini' + if (Object.values(openaiModels).includes(model as OpenAIModel) && userApiKeys?.openai) return 'openai' + return null +} + +/** + * Convert a model string to an AI SDK LanguageModel instance. + * Supports BYOK (Bring Your Own Key) for Anthropic, Gemini, and OpenAI. + */ +const modelToAiSDKModel = ( + model: Model, + userApiKeys?: UserApiKeys, + byokMode: ByokMode = 'prefer', +): LanguageModel => { + // Finetuned Vertex models if ( Object.values(finetunedVertexModels as Record).includes( model, @@ -52,16 +91,66 @@ const modelToAiSDKModel = (model: Model): LanguageModel => { ) { return vertexFinetuned(model) } + + // Gemini models - direct to Google if (Object.values(geminiModels).includes(model as GeminiModel)) { - return google.languageModel(model) + const apiKey = + byokMode === 'disabled' + ? env.GEMINI_API_KEY + : userApiKeys?.gemini ?? env.GEMINI_API_KEY + + if (byokMode === 'require' && !userApiKeys?.gemini) { + throw new Error('Gemini API key required but not provided (byokMode: require)') + } + + return google.languageModel(model, { apiKey }) } + + // OpenAI models - direct to OpenAI if (model === openaiModels.o3pro || model === openaiModels.o3) { - return openai.responses(model) + const apiKey = + byokMode === 'disabled' + ? env.OPENAI_API_KEY + : userApiKeys?.openai ?? env.OPENAI_API_KEY + + if (byokMode === 'require' && !userApiKeys?.openai) { + throw new Error('OpenAI API key required but not provided (byokMode: require)') + } + + return openai.responses(model, { apiKey }) } + if (Object.values(openaiModels).includes(model as OpenAIModel)) { - return openai.languageModel(model) + const apiKey = + byokMode === 'disabled' + ? env.OPENAI_API_KEY + : userApiKeys?.openai ?? env.OPENAI_API_KEY + + if (byokMode === 'require' && !userApiKeys?.openai) { + throw new Error('OpenAI API key required but not provided (byokMode: require)') + } + + return openai.languageModel(model, { apiKey }) } - // All other models go through OpenRouter + + // Anthropic models - direct to Anthropic (if user key provided) or OpenRouter + if (isAnthropicModel(model)) { + // If user has Anthropic key and byokMode allows it, use direct Anthropic API + if (byokMode !== 'disabled' && userApiKeys?.anthropic) { + const anthropic = createAnthropic({ apiKey: userApiKeys.anthropic }) + return anthropic.languageModel(model) + } + + // If byokMode is 'require', fail if no user key + if (byokMode === 'require') { + throw new Error('Anthropic API key required but not provided (byokMode: require)') + } + + // Otherwise, use OpenRouter with system key + return openRouterLanguageModel(model) + } + + // All other models go through OpenRouter with system key return openRouterLanguageModel(model) } @@ -82,6 +171,8 @@ export const promptAiSdkStream = async function* ( maxRetries?: number onCostCalculated?: (credits: number) => Promise includeCacheControl?: boolean + userApiKeys?: UserApiKeys + byokMode?: ByokMode } & Omit[0], 'model' | 'messages'>, ): AsyncGenerator { if ( @@ -103,7 +194,8 @@ export const promptAiSdkStream = async function* ( } const startTime = Date.now() - let aiSDKModel = modelToAiSDKModel(options.model) + const byokMode = options.byokMode ?? 'prefer' + let aiSDKModel = modelToAiSDKModel(options.model, options.userApiKeys, byokMode) const response = streamText({ ...options, @@ -156,8 +248,8 @@ export const promptAiSdkStream = async function* ( if ( ( options.providerOptions?.openrouter as - | OpenRouterProviderOptions - | undefined + | OpenRouterProviderOptions + | undefined )?.reasoning?.exclude ) { continue @@ -230,6 +322,7 @@ export const promptAiSdkStream = async function* ( } const messageId = (await response.response).id + const byokProvider = determineByokProvider(options.model, options.userApiKeys) const creditsUsedPromise = saveMessage({ messageId, userId: options.userId, @@ -246,6 +339,7 @@ export const promptAiSdkStream = async function* ( finishedAt: new Date(), latencyMs: Date.now() - startTime, chargeUser: options.chargeUser ?? true, + byokProvider, costOverrideDollars, agentId: options.agentId, }) @@ -273,6 +367,8 @@ export const promptAiSdk = async function ( onCostCalculated?: (credits: number) => Promise includeCacheControl?: boolean maxRetries?: number + userApiKeys?: UserApiKeys + byokMode?: ByokMode } & Omit[0], 'model' | 'messages'>, ): Promise { if ( @@ -294,7 +390,8 @@ export const promptAiSdk = async function ( } const startTime = Date.now() - let aiSDKModel = modelToAiSDKModel(options.model) + const byokMode = options.byokMode ?? 'prefer' + let aiSDKModel = modelToAiSDKModel(options.model, options.userApiKeys, byokMode) const response = await generateText({ ...options, @@ -305,6 +402,7 @@ export const promptAiSdk = async function ( const inputTokens = response.usage.inputTokens || 0 const outputTokens = response.usage.inputTokens || 0 + const byokProvider = determineByokProvider(options.model, options.userApiKeys) const creditsUsedPromise = saveMessage({ messageId: generateCompactId(), userId: options.userId, @@ -320,6 +418,7 @@ export const promptAiSdk = async function ( latencyMs: Date.now() - startTime, chargeUser: options.chargeUser ?? true, agentId: options.agentId, + byokProvider, }) // Call the cost callback if provided @@ -348,6 +447,8 @@ export const promptAiSdkStructured = async function (options: { onCostCalculated?: (credits: number) => Promise includeCacheControl?: boolean maxRetries?: number + userApiKeys?: UserApiKeys + byokMode?: ByokMode }): Promise { if ( !checkLiveUserInput( @@ -367,7 +468,8 @@ export const promptAiSdkStructured = async function (options: { return {} as T } const startTime = Date.now() - let aiSDKModel = modelToAiSDKModel(options.model) + const byokMode = options.byokMode ?? 'prefer' + let aiSDKModel = modelToAiSDKModel(options.model, options.userApiKeys, byokMode) const responsePromise = generateObject, 'object'>({ ...options, @@ -383,6 +485,7 @@ export const promptAiSdkStructured = async function (options: { const inputTokens = response.usage.inputTokens || 0 const outputTokens = response.usage.inputTokens || 0 + const byokProvider = determineByokProvider(options.model, options.userApiKeys) const creditsUsedPromise = saveMessage({ messageId: generateCompactId(), userId: options.userId, @@ -398,6 +501,7 @@ export const promptAiSdkStructured = async function (options: { latencyMs: Date.now() - startTime, chargeUser: options.chargeUser ?? true, agentId: options.agentId, + byokProvider, }) // Call the cost callback if provided diff --git a/backend/src/main-prompt.ts b/backend/src/main-prompt.ts index f37d610c4..3618033a0 100644 --- a/backend/src/main-prompt.ts +++ b/backend/src/main-prompt.ts @@ -8,6 +8,7 @@ import { getAgentTemplate } from './templates/agent-registry' import { logger } from './util/logger' import { expireMessages } from './util/messages' import { requestToolCall } from './websockets/websocket-action' +import { retrieveAndDecryptApiKey } from '@codebuff/common/api-keys/crypto' import type { AgentTemplate } from './templates/types' import type { ClientAction } from '@codebuff/common/actions' @@ -19,6 +20,7 @@ import type { AgentOutput, } from '@codebuff/common/types/session-state' import type { WebSocket } from 'ws' +import type { UserApiKeys, ByokMode } from './llm-apis/vercel-ai-sdk/ai-sdk' export interface MainPromptOptions { userId: string | undefined @@ -27,6 +29,38 @@ export interface MainPromptOptions { localAgentTemplates: Record } +/** + * Retrieves user API keys from the database for BYOK (Bring Your Own Key) + * Merges SDK-provided keys with database keys, with SDK keys taking precedence + */ +async function getUserApiKeys( + userId: string | undefined, + sdkKeys?: UserApiKeys, +): Promise { + if (!userId) { + return sdkKeys + } + + try { + // Retrieve keys from database + const [anthropicKey, geminiKey, openaiKey] = await Promise.all([ + retrieveAndDecryptApiKey(userId, 'anthropic'), + retrieveAndDecryptApiKey(userId, 'gemini'), + retrieveAndDecryptApiKey(userId, 'openai'), + ]) + + // Merge with SDK keys (SDK keys take precedence) + return { + anthropic: sdkKeys?.anthropic ?? anthropicKey ?? undefined, + gemini: sdkKeys?.gemini ?? geminiKey ?? undefined, + openai: sdkKeys?.openai ?? openaiKey ?? undefined, + } + } catch (error) { + logger.error({ error, userId }, 'Failed to retrieve user API keys') + return sdkKeys + } +} + export const mainPrompt = async ( ws: WebSocket, action: ClientAction<'prompt'>, @@ -47,9 +81,14 @@ export const mainPrompt = async ( promptId, agentId, promptParams, + userApiKeys: sdkUserApiKeys, + byokMode, } = action const { fileContext, mainAgentState } = sessionState + // Retrieve and merge user API keys (SDK keys take precedence over DB keys) + const userApiKeys = await getUserApiKeys(userId, sdkUserApiKeys) + if (prompt) { // Check if this is a direct terminal command const startTime = Date.now() @@ -203,6 +242,8 @@ export const mainPrompt = async ( clientSessionId, onResponseChunk, localAgentTemplates, + userApiKeys, + byokMode, }) logger.debug({ agentState, output }, 'Main prompt finished') diff --git a/backend/src/prompt-agent-stream.ts b/backend/src/prompt-agent-stream.ts index 774d3f6aa..4c4d72c93 100644 --- a/backend/src/prompt-agent-stream.ts +++ b/backend/src/prompt-agent-stream.ts @@ -15,6 +15,8 @@ export const getAgentStreamFromTemplate = (params: { onCostCalculated?: (credits: number) => Promise agentId?: string includeCacheControl?: boolean + userApiKeys?: import('./llm-apis/vercel-ai-sdk/ai-sdk').UserApiKeys + byokMode?: import('./llm-apis/vercel-ai-sdk/ai-sdk').ByokMode template: AgentTemplate }) => { @@ -26,6 +28,8 @@ export const getAgentStreamFromTemplate = (params: { onCostCalculated, agentId, includeCacheControl, + userApiKeys, + byokMode, template, } = params @@ -49,6 +53,8 @@ export const getAgentStreamFromTemplate = (params: { includeCacheControl, agentId, maxRetries: 3, + userApiKeys, + byokMode, } // Add Gemini-specific options if needed @@ -70,7 +76,7 @@ export const getAgentStreamFromTemplate = (params: { if (!options.providerOptions.openrouter) { options.providerOptions.openrouter = {} } - ;( + ; ( options.providerOptions.openrouter as OpenRouterProviderOptions ).reasoning = template.reasoningOptions diff --git a/backend/src/run-agent-step.ts b/backend/src/run-agent-step.ts index 877ca77be..eb383e483 100644 --- a/backend/src/run-agent-step.ts +++ b/backend/src/run-agent-step.ts @@ -250,6 +250,8 @@ export const runAgentStep = async ( userId, agentId: agentState.agentId, template: agentTemplate, + userApiKeys, + byokMode, onCostCalculated: async (credits: number) => { try { agentState.creditsUsed += credits @@ -455,6 +457,8 @@ export const loopAgentSteps = async ( clientSessionId, onResponseChunk, clearUserPromptMessagesAfterResponse = true, + userApiKeys, + byokMode, }: { userInputId: string agentType: AgentTemplateType @@ -470,6 +474,8 @@ export const loopAgentSteps = async ( userId: string | undefined clientSessionId: string onResponseChunk: (chunk: string | PrintModeEvent) => void + userApiKeys?: import('./llm-apis/vercel-ai-sdk/ai-sdk').UserApiKeys + byokMode?: import('./llm-apis/vercel-ai-sdk/ai-sdk').ByokMode }, ): Promise<{ agentState: AgentState @@ -497,27 +503,27 @@ export const loopAgentSteps = async ( // Get the instructions prompt if we have a prompt/params const instructionsPrompt = hasPrompt ? await getAgentPrompt({ - agentTemplate, - promptType: { type: 'instructionsPrompt' }, - fileContext, - agentState, - agentTemplates: localAgentTemplates, - additionalToolDefinitions: () => { - const additionalToolDefinitions = cloneDeep( - Object.fromEntries( - Object.entries(fileContext.customToolDefinitions).filter( - ([toolName]) => agentTemplate.toolNames.includes(toolName), - ), + agentTemplate, + promptType: { type: 'instructionsPrompt' }, + fileContext, + agentState, + agentTemplates: localAgentTemplates, + additionalToolDefinitions: () => { + const additionalToolDefinitions = cloneDeep( + Object.fromEntries( + Object.entries(fileContext.customToolDefinitions).filter( + ([toolName]) => agentTemplate.toolNames.includes(toolName), ), - ) - return getMCPToolData({ - ws, - toolNames: agentTemplate.toolNames, - mcpServers: agentTemplate.mcpServers, - writeTo: additionalToolDefinitions, - }) - }, - }) + ), + ) + return getMCPToolData({ + ws, + toolNames: agentTemplate.toolNames, + mcpServers: agentTemplate.mcpServers, + writeTo: additionalToolDefinitions, + }) + }, + }) : undefined // Build the initial message history with user prompt and instructions @@ -532,14 +538,14 @@ export const loopAgentSteps = async ( keepDuringTruncation: true, }, prompt && - prompt in additionalSystemPrompts && { - role: 'user' as const, - content: asSystemInstruction( - additionalSystemPrompts[ - prompt as keyof typeof additionalSystemPrompts - ], - ), - }, + prompt in additionalSystemPrompts && { + role: 'user' as const, + content: asSystemInstruction( + additionalSystemPrompts[ + prompt as keyof typeof additionalSystemPrompts + ], + ), + }, ], instructionsPrompt && { diff --git a/bun.lock b/bun.lock index 6a1c87ba1..b89d4e3b4 100644 --- a/bun.lock +++ b/bun.lock @@ -39,6 +39,7 @@ "name": "@codebuff/backend", "version": "1.0.0", "dependencies": { + "@ai-sdk/anthropic": "1.0.8", "@ai-sdk/google-vertex": "3.0.6", "@ai-sdk/openai": "2.0.11", "@codebuff/billing": "workspace:*", @@ -238,7 +239,7 @@ }, "sdk": { "name": "@codebuff/sdk", - "version": "0.3.8", + "version": "0.3.13", "dependencies": { "@vscode/tree-sitter-wasm": "0.1.4", "ai": "^5.0.0", @@ -370,7 +371,7 @@ "packages": { "@adobe/css-tools": ["@adobe/css-tools@4.4.4", "", {}, "sha512-Elp+iwUx5rN5+Y8xLt5/GRoG20WGoDCQ/1Fb+1LiGtvwbDavuSk0jhD/eZdckHAuzcDzccnkv+rEjyWfRx18gg=="], - "@ai-sdk/anthropic": ["@ai-sdk/anthropic@2.0.2", "", { "dependencies": { "@ai-sdk/provider": "2.0.0", "@ai-sdk/provider-utils": "3.0.2" }, "peerDependencies": { "zod": "^3.25.76 || ^4" } }, "sha512-R3xmEbbntgdKo/S3TDuW77RYALpo/OKQm4oSjQmryDAFiVGB6X6guZAr7FWt48C4fKGROScAu+y1MJTbzisfOQ=="], + "@ai-sdk/anthropic": ["@ai-sdk/anthropic@1.0.8", "", { "dependencies": { "@ai-sdk/provider": "1.0.4", "@ai-sdk/provider-utils": "2.0.7" }, "peerDependencies": { "zod": "^3.0.0" } }, "sha512-SruTs0JOZ5ZnVV2hzeu0XDzRrT9WHcgx9P1p5vpjJFJVr9FlVaTxgxisL+8tlhZy8FX68zAhtj09rAaL4gT+jA=="], "@ai-sdk/gateway": ["@ai-sdk/gateway@1.0.0", "", { "dependencies": { "@ai-sdk/provider": "2.0.0", "@ai-sdk/provider-utils": "3.0.0" }, "peerDependencies": { "zod": "^3.25.76 || ^4" } }, "sha512-VEm87DyRx1yIPywbTy8ntoyh4jEDv1rJ88m+2I7zOm08jJI5BhFtAWh0OF6YzZu1Vu4NxhOWO4ssGdsqydDQ3A=="], @@ -380,9 +381,9 @@ "@ai-sdk/openai": ["@ai-sdk/openai@2.0.11", "", { "dependencies": { "@ai-sdk/provider": "2.0.0", "@ai-sdk/provider-utils": "3.0.2" }, "peerDependencies": { "zod": "^3.25.76 || ^4" } }, "sha512-t4i+vS825EC0Gc2DdTsC5UkXIu1ScOi363noTD8DuFZp6WFPHRnW6HCyEQKxEm6cNjv3BW89rdXWqq932IFJhA=="], - "@ai-sdk/provider": ["@ai-sdk/provider@2.0.0", "", { "dependencies": { "json-schema": "^0.4.0" } }, "sha512-6o7Y2SeO9vFKB8lArHXehNuusnpddKPk7xqL7T2/b+OvXMRIXUO1rR4wcv1hAFUAT9avGZshty3Wlua/XA7TvA=="], + "@ai-sdk/provider": ["@ai-sdk/provider@1.0.4", "", { "dependencies": { "json-schema": "^0.4.0" } }, "sha512-lJi5zwDosvvZER3e/pB8lj1MN3o3S7zJliQq56BRr4e9V3fcRyFtwP0JRxaRS5vHYX3OJ154VezVoQNrk0eaKw=="], - "@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@3.0.2", "", { "dependencies": { "@ai-sdk/provider": "2.0.0", "@standard-schema/spec": "^1.0.0", "eventsource-parser": "^3.0.3", "zod-to-json-schema": "^3.24.1" }, "peerDependencies": { "zod": "^3.25.76 || ^4" } }, "sha512-0a5a6VafkV6+0irdpqnub8WE6qzG2VMsDBpXb9NQIz8c4TG8fI+GSTFIL9sqrLEwXrHdiRj7fwJsrir4jClL0w=="], + "@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@2.0.7", "", { "dependencies": { "@ai-sdk/provider": "1.0.4", "eventsource-parser": "^3.0.0", "nanoid": "^3.3.8", "secure-json-parse": "^2.7.0" }, "peerDependencies": { "zod": "^3.0.0" }, "optionalPeers": ["zod"] }, "sha512-4sfPlKEALHPXLmMFcPlYksst3sWBJXmCDZpIBJisRrmwGG6Nn3mq0N1Zu/nZaGcrWZoOY+HT2Wbxla1oTElYHQ=="], "@alloc/quick-lru": ["@alloc/quick-lru@5.2.0", "", {}, "sha512-UrcABB+4bUrFABwbluTIBErXwvbsU/V7TZWfmbgJfbkwiBuziS9gxdODUyuiecfdGQ85jglMW6juS3+z5TsKLw=="], @@ -3472,6 +3473,8 @@ "section-matter": ["section-matter@1.0.0", "", { "dependencies": { "extend-shallow": "^2.0.1", "kind-of": "^6.0.0" } }, "sha512-vfD3pmTzGpufjScBh50YHKzEu2lxBWhVEHsNGoEXmCmn2hKGfeNLYMzCJpe8cD7gqX7TJluOVpBkAequ6dgMmA=="], + "secure-json-parse": ["secure-json-parse@2.7.0", "", {}, "sha512-6aU+Rwsezw7VR8/nyvKTx8QpWH9FrcYiXXlqC4z5d5XQBDRqtbfsRjnwGyqbi3gddNtWHuEk9OANUotL26qKUw=="], + "seedrandom": ["seedrandom@3.0.5", "", {}, "sha512-8OwmbklUNzwezjGInmZ+2clQmExQPvomqjL7LFqOYqtmuxRgQYqOD3mHaU+MvZn5FLUeVxVfQjwLZW/n/JFuqg=="], "semver": ["semver@7.7.2", "", { "bin": { "semver": "bin/semver.js" } }, "sha512-RF0Fw+rO5AMf9MAyaRXI4AV0Ulj5lMHqVxxdSgiVbixSCXoEmmX/jk0CuJw4+3SqroYO9VoUh+HcuJivvtJemA=="], @@ -3968,8 +3971,26 @@ "zwitch": ["zwitch@2.0.4", "", {}, "sha512-bXE4cR/kVZhKZX/RjPEflHaKVhUVl85noU3v6b8apfQEc1x4A+zBxjZ4lN8LqGd6WZ3dl98pY4o717VFmoPp+A=="], + "@ai-sdk/gateway/@ai-sdk/provider": ["@ai-sdk/provider@2.0.0", "", { "dependencies": { "json-schema": "^0.4.0" } }, "sha512-6o7Y2SeO9vFKB8lArHXehNuusnpddKPk7xqL7T2/b+OvXMRIXUO1rR4wcv1hAFUAT9avGZshty3Wlua/XA7TvA=="], + "@ai-sdk/gateway/@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@3.0.0", "", { "dependencies": { "@ai-sdk/provider": "2.0.0", "@standard-schema/spec": "^1.0.0", "eventsource-parser": "^3.0.3", "zod-to-json-schema": "^3.24.1" }, "peerDependencies": { "zod": "^3.25.76 || ^4" } }, "sha512-BoQZtGcBxkeSH1zK+SRYNDtJPIPpacTeiMZqnG4Rv6xXjEwM0FH4MGs9c+PlhyEWmQCzjRM2HAotEydFhD4dYw=="], + "@ai-sdk/google/@ai-sdk/provider": ["@ai-sdk/provider@2.0.0", "", { "dependencies": { "json-schema": "^0.4.0" } }, "sha512-6o7Y2SeO9vFKB8lArHXehNuusnpddKPk7xqL7T2/b+OvXMRIXUO1rR4wcv1hAFUAT9avGZshty3Wlua/XA7TvA=="], + + "@ai-sdk/google/@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@3.0.2", "", { "dependencies": { "@ai-sdk/provider": "2.0.0", "@standard-schema/spec": "^1.0.0", "eventsource-parser": "^3.0.3", "zod-to-json-schema": "^3.24.1" }, "peerDependencies": { "zod": "^3.25.76 || ^4" } }, "sha512-0a5a6VafkV6+0irdpqnub8WE6qzG2VMsDBpXb9NQIz8c4TG8fI+GSTFIL9sqrLEwXrHdiRj7fwJsrir4jClL0w=="], + + "@ai-sdk/google-vertex/@ai-sdk/anthropic": ["@ai-sdk/anthropic@2.0.2", "", { "dependencies": { "@ai-sdk/provider": "2.0.0", "@ai-sdk/provider-utils": "3.0.2" }, "peerDependencies": { "zod": "^3.25.76 || ^4" } }, "sha512-R3xmEbbntgdKo/S3TDuW77RYALpo/OKQm4oSjQmryDAFiVGB6X6guZAr7FWt48C4fKGROScAu+y1MJTbzisfOQ=="], + + "@ai-sdk/google-vertex/@ai-sdk/provider": ["@ai-sdk/provider@2.0.0", "", { "dependencies": { "json-schema": "^0.4.0" } }, "sha512-6o7Y2SeO9vFKB8lArHXehNuusnpddKPk7xqL7T2/b+OvXMRIXUO1rR4wcv1hAFUAT9avGZshty3Wlua/XA7TvA=="], + + "@ai-sdk/google-vertex/@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@3.0.2", "", { "dependencies": { "@ai-sdk/provider": "2.0.0", "@standard-schema/spec": "^1.0.0", "eventsource-parser": "^3.0.3", "zod-to-json-schema": "^3.24.1" }, "peerDependencies": { "zod": "^3.25.76 || ^4" } }, "sha512-0a5a6VafkV6+0irdpqnub8WE6qzG2VMsDBpXb9NQIz8c4TG8fI+GSTFIL9sqrLEwXrHdiRj7fwJsrir4jClL0w=="], + + "@ai-sdk/openai/@ai-sdk/provider": ["@ai-sdk/provider@2.0.0", "", { "dependencies": { "json-schema": "^0.4.0" } }, "sha512-6o7Y2SeO9vFKB8lArHXehNuusnpddKPk7xqL7T2/b+OvXMRIXUO1rR4wcv1hAFUAT9avGZshty3Wlua/XA7TvA=="], + + "@ai-sdk/openai/@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@3.0.2", "", { "dependencies": { "@ai-sdk/provider": "2.0.0", "@standard-schema/spec": "^1.0.0", "eventsource-parser": "^3.0.3", "zod-to-json-schema": "^3.24.1" }, "peerDependencies": { "zod": "^3.25.76 || ^4" } }, "sha512-0a5a6VafkV6+0irdpqnub8WE6qzG2VMsDBpXb9NQIz8c4TG8fI+GSTFIL9sqrLEwXrHdiRj7fwJsrir4jClL0w=="], + + "@ai-sdk/provider-utils/nanoid": ["nanoid@3.3.11", "", { "bin": { "nanoid": "bin/nanoid.cjs" } }, "sha512-N8SpfPUnUp1bK+PMYW8qSWdl9U+wwNWI4QKxOYDy9JAro3WMX7p2OeVRF9v+347pnakNevPmiHhNmZ2HbFA76w=="], + "@ampproject/remapping/@jridgewell/trace-mapping": ["@jridgewell/trace-mapping@0.3.30", "", { "dependencies": { "@jridgewell/resolve-uri": "^3.1.0", "@jridgewell/sourcemap-codec": "^1.4.14" } }, "sha512-GQ7Nw5G2lTu/BtHTKfXhKHok2WGetd4XYcVKGx00SjAk8GMwgJM3zr6zORiPGuOE+/vkc90KtTosSSvaCjKb2Q=="], "@auth/core/jose": ["jose@6.1.0", "", {}, "sha512-TTQJyoEoKcC1lscpVDCSsVgYzUDg/0Bt3WE//WiTPK6uOCQC2KZS4MpugbMWt/zyjkopgZoXhZuCi00gLudfUA=="], @@ -4224,6 +4245,8 @@ "aceternity-ui/https-proxy-agent": ["https-proxy-agent@6.2.1", "", { "dependencies": { "agent-base": "^7.0.2", "debug": "4" } }, "sha512-ONsE3+yfZF2caH5+bJlcddtWqNI3Gvs5A38+ngvljxaBiRXRswym2c7yf8UAeFpRFKjFNHIFEHqR/OLAWJzyiA=="], + "ai/@ai-sdk/provider": ["@ai-sdk/provider@2.0.0", "", { "dependencies": { "json-schema": "^0.4.0" } }, "sha512-6o7Y2SeO9vFKB8lArHXehNuusnpddKPk7xqL7T2/b+OvXMRIXUO1rR4wcv1hAFUAT9avGZshty3Wlua/XA7TvA=="], + "ai/@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@3.0.0", "", { "dependencies": { "@ai-sdk/provider": "2.0.0", "@standard-schema/spec": "^1.0.0", "eventsource-parser": "^3.0.3", "zod-to-json-schema": "^3.24.1" }, "peerDependencies": { "zod": "^3.25.76 || ^4" } }, "sha512-BoQZtGcBxkeSH1zK+SRYNDtJPIPpacTeiMZqnG4Rv6xXjEwM0FH4MGs9c+PlhyEWmQCzjRM2HAotEydFhD4dYw=="], "autoprefixer/picocolors": ["picocolors@1.1.1", "", {}, "sha512-xceH2snhtb5M9liqDsmEw56le376mTZkEX/jEb/RxNFyegNul7eNslCXP9FDj/Lcu0X8KEyMceP2ntpaHrDEVA=="], diff --git a/common/src/actions.ts b/common/src/actions.ts index 4ab573460..1fc852152 100644 --- a/common/src/actions.ts +++ b/common/src/actions.ts @@ -41,6 +41,15 @@ export const CLIENT_ACTION_SCHEMA = z.discriminatedUnion('type', [ model: z.string().optional(), repoUrl: z.string().optional(), agentId: z.string().optional(), + // BYOK (Bring Your Own Key) support + userApiKeys: z + .object({ + anthropic: z.string().optional(), + gemini: z.string().optional(), + openai: z.string().optional(), + }) + .optional(), + byokMode: z.enum(['disabled', 'prefer', 'require']).optional(), }), z.object({ type: z.literal('read-files-response'), diff --git a/common/src/api-keys/crypto.ts b/common/src/api-keys/crypto.ts index 041a7fa9e..2e21821bd 100644 --- a/common/src/api-keys/crypto.ts +++ b/common/src/api-keys/crypto.ts @@ -198,6 +198,29 @@ export async function retrieveAndDecryptApiKey( } } +/** + * Validates an API key format based on its type. + * @param keyType The type of the API key (e.g., 'anthropic', 'gemini', 'openai'). + * @param apiKey The API key to validate. + * @returns True if the key format is valid, false otherwise. + */ +export function validateApiKey(keyType: ApiKeyType, apiKey: string): boolean { + const prefix = KEY_PREFIXES[keyType] + const length = KEY_LENGTHS[keyType] + + // Check prefix + if (prefix && !apiKey.startsWith(prefix)) { + return false + } + + // Check length + if (length && apiKey.length !== length) { + return false + } + + return true +} + /** * Deletes a specific API key entry for a given user and key type. * @param userId The ID of the user. diff --git a/sdk/README.md b/sdk/README.md index 2b6de4647..d6351fd8d 100644 --- a/sdk/README.md +++ b/sdk/README.md @@ -134,8 +134,101 @@ async function main() { main() ``` +### Example 3: Bring Your Own Key (BYOK) + +Use your own API keys for Anthropic, Gemini, or OpenAI models to pay directly for LLM costs with reduced or zero Codebuff markup. + +```typescript +import { CodebuffClient } from '@codebuff/sdk' + +async function main() { + const client = new CodebuffClient({ + // Option 1: Use only your provider keys (no Codebuff API key required) + userApiKeys: { + anthropic: process.env.ANTHROPIC_API_KEY, + gemini: process.env.GEMINI_API_KEY, + openai: process.env.OPENAI_API_KEY, + }, + byokMode: 'require', // Only use user keys, fail if missing + cwd: process.cwd(), + }) + + // Option 2: Use Codebuff API key with provider keys as fallback + const client2 = new CodebuffClient({ + apiKey: process.env.CODEBUFF_API_KEY, + userApiKeys: { + anthropic: process.env.ANTHROPIC_API_KEY, + }, + byokMode: 'prefer', // Use user keys when available, fallback to system keys (default) + cwd: process.cwd(), + }) + + // Option 3: Disable BYOK and always use system keys + const client3 = new CodebuffClient({ + apiKey: process.env.CODEBUFF_API_KEY, + byokMode: 'disabled', // Always use system keys + cwd: process.cwd(), + }) + + const run = await client.run({ + agent: 'codebuff/base@0.0.16', + prompt: 'Create a simple calculator class', + handleEvent: (event) => { + console.log('Codebuff Event', JSON.stringify(event)) + }, + }) +} + +main() +``` + +#### BYOK Modes + +- **`'disabled'`**: Always use Codebuff's system keys. Requires a Codebuff API key. +- **`'prefer'`** (default): Use your provider keys when available, fallback to system keys. Recommended for most users. +- **`'require'`**: Only use your provider keys. No Codebuff API key required. Fails if provider key is missing for the selected model. + +#### BYOK Benefits + +- **Lower Costs**: Pay only the provider's API costs with reduced Codebuff markup +- **Direct Billing**: Charges appear directly on your provider account +- **No Codebuff API Key Required**: When using `byokMode: 'require'`, you can use Codebuff without a Codebuff API key +- **Provider Choice**: Use your preferred provider's billing and rate limits + +#### Supported Providers + +- **Anthropic**: Claude models (e.g., `anthropic/claude-3.5-sonnet`) +- **Google Gemini**: Gemini models (e.g., `gemini-2.0-flash-exp`) +- **OpenAI**: GPT models (e.g., `gpt-4o`, `o1`, `o3-mini`) + +#### Security + +- API keys are encrypted at rest using AES-256-GCM +- Keys are validated before storage +- Keys are never logged or exposed in error messages + ## API Reference +### `new CodebuffClient(options)` + +Creates a new Codebuff client instance. + +#### Constructor Parameters + +- **`apiKey`** (string, optional): Your Codebuff API key. Get one at [codebuff.com/api-keys](https://www.codebuff.com/api-keys). Optional if using `byokMode: 'require'` with provider keys. + +- **`cwd`** (string, optional): Working directory for the agent. Defaults to `process.cwd()`. + +- **`userApiKeys`** (object, optional): Your own API keys for AI providers. Enables BYOK (Bring Your Own Key) mode. + - `anthropic` (string, optional): Anthropic API key (starts with `sk-ant-api03-`) + - `gemini` (string, optional): Google Gemini API key (starts with `AIzaSy`) + - `openai` (string, optional): OpenAI API key (starts with `sk-proj-`) + +- **`byokMode`** (string, optional): Controls how user API keys are used. Defaults to `'prefer'`. + - `'disabled'`: Always use Codebuff's system keys (requires Codebuff API key) + - `'prefer'`: Use user keys when available, fallback to system keys (default) + - `'require'`: Only use user keys, fail if missing (no Codebuff API key needed) + ### `client.run(options)` Runs a Codebuff agent with the specified options. diff --git a/sdk/knowledge.md b/sdk/knowledge.md new file mode 100644 index 000000000..6e29702b5 --- /dev/null +++ b/sdk/knowledge.md @@ -0,0 +1,223 @@ +# Codebuff SDK Knowledge Base + +## Architecture Overview + +The Codebuff SDK provides a TypeScript/JavaScript interface to the Codebuff AI coding agent platform. It handles communication with the Codebuff backend via WebSocket connections, manages agent state, and provides a simple API for running AI agents with custom tools. + +## BYOK (Bring Your Own Key) Architecture + +### Overview + +BYOK allows SDK users to provide their own API keys for Anthropic, Gemini, and OpenAI models. This enables users to: +- Pay directly for LLM API costs through their provider accounts +- Benefit from reduced or zero Codebuff markup +- Use Codebuff's agent infrastructure without a Codebuff API key (in `require` mode) + +### Key Components + +#### 1. SDK Layer (`sdk/src/`) + +**`client.ts`**: +- Accepts `userApiKeys` and `byokMode` in constructor options +- Validates authentication based on byokMode: + - `disabled`: Requires Codebuff API key + - `prefer`: Accepts either Codebuff API key or user keys + - `require`: Requires at least one user API key (no Codebuff key needed) + +**`run.ts`**: +- Passes `userApiKeys` and `byokMode` through WebSocket connection +- Includes these parameters in the CLIENT_ACTION_SCHEMA + +**`websocket-client.ts`**: +- Transmits user keys and mode to backend via WebSocket messages + +#### 2. Common Layer (`common/src/`) + +**`actions.ts`**: +- Defines CLIENT_ACTION_SCHEMA with optional `userApiKeys` and `byokMode` fields +- Validates action payloads before transmission + +**`api-keys/crypto.ts`**: +- `validateApiKey()`: Validates API key format (prefix and length) +- `encryptAndStoreApiKey()`: Encrypts keys using AES-256-GCM before storage +- `retrieveAndDecryptApiKey()`: Retrieves and decrypts keys from database +- `clearApiKey()`: Removes keys from database + +#### 3. Web Layer (`web/src/`) + +**`app/api/user-api-keys/route.ts`**: +- GET: Returns list of configured key types for authenticated user +- POST: Validates and stores encrypted API keys + +**`app/api/user-api-keys/[keyType]/route.ts`**: +- DELETE: Removes specific API key for authenticated user + +**`app/profile/components/user-api-keys-section.tsx`**: +- React component for managing provider API keys +- Card-based UI for each provider (Anthropic, Gemini, OpenAI) +- Shows configuration status, masked keys, input fields +- Handles save/update/remove operations + +#### 4. Backend Layer (`backend/src/`) + +**`main-prompt.ts`**: +- `getUserApiKeys()`: Retrieves user keys from database and merges with SDK-provided keys +- Key precedence: SDK keys > DB keys > system keys +- Passes merged keys to agent execution pipeline + +**`llm-apis/vercel-ai-sdk/ai-sdk.ts`**: +- `modelToAiSDKModel()`: Routes models to appropriate provider based on BYOK configuration +- `isAnthropicModel()`: Identifies Anthropic models +- `determineByokProvider()`: Determines which provider key was used +- Direct-to-provider routing: + - Anthropic models with user key → `@ai-sdk/anthropic` + - Gemini models with user key → `@ai-sdk/google` + - OpenAI models with user key → `@ai-sdk/openai` + - Models without user keys → OpenRouter (system keys) + +**`llm-apis/message-cost-tracker.ts`**: +- `saveMessage()`: Tracks costs per provider +- Applies reduced markup for BYOK usage: `PROFIT_MARGIN` vs `1 + PROFIT_MARGIN` +- `byokProvider` field indicates which provider key was used + +**`run-agent-step.ts`**: +- `loopAgentSteps()`: Passes BYOK parameters through agent execution loop + +**`prompt-agent-stream.ts`**: +- `getAgentStreamFromTemplate()`: Passes BYOK parameters to AI SDK functions + +### Data Flow + +1. **SDK Initialization**: + ``` + User → CodebuffClient(userApiKeys, byokMode) → Validation + ``` + +2. **Run Execution**: + ``` + client.run() → WebSocket → Backend → getUserApiKeys() → Merge Keys + ``` + +3. **Model Routing**: + ``` + Model Selection → modelToAiSDKModel(model, userApiKeys, byokMode) + → Direct Provider API or OpenRouter + ``` + +4. **Cost Tracking**: + ``` + API Response → determineByokProvider() → saveMessage(byokProvider) + → Reduced Markup Calculation + ``` + +### Key Precedence + +When determining which API key to use: +1. **SDK-provided keys** (passed in `client.run()` or constructor) +2. **Database keys** (stored via web UI) +3. **System keys** (Codebuff's keys) + +This allows users to override database keys on a per-run basis. + +### Security Considerations + +1. **Encryption**: All user API keys are encrypted at rest using AES-256-GCM +2. **Validation**: Keys are validated for correct format before storage +3. **No Logging**: Keys are never logged or exposed in error messages +4. **Secure Transmission**: Keys are transmitted over secure WebSocket connections +5. **Database Storage**: Keys stored in `encrypted_api_keys` table with composite primary key (user_id, type) + +### Provider Routing + +#### Anthropic Models +- **With User Key**: Direct to Anthropic API via `@ai-sdk/anthropic` +- **Without User Key**: Through OpenRouter with system keys +- **Model Format**: `anthropic/claude-3.5-sonnet`, etc. + +#### Gemini Models +- **With User Key**: Direct to Google API via `@ai-sdk/google` +- **Without User Key**: System Gemini key +- **Model Format**: `gemini-2.0-flash-exp`, etc. + +#### OpenAI Models +- **With User Key**: Direct to OpenAI API via `@ai-sdk/openai` +- **Without User Key**: System OpenAI key +- **Model Format**: `gpt-4o`, `o1`, `o3-mini`, etc. + +### Cost Calculation + +```typescript +// Without BYOK (system keys) +costInCents = cost * 100 * (1 + PROFIT_MARGIN) + +// With BYOK (user keys) +costInCents = cost * 100 * PROFIT_MARGIN +``` + +The reduced markup for BYOK reflects that users are paying for the LLM API costs directly. + +### Error Handling + +#### `byokMode: 'require'` +- Throws error if no user key available for selected model +- Example: "Anthropic API key required but not provided (byokMode: require)" + +#### `byokMode: 'prefer'` +- Falls back to system keys if user key unavailable +- No error thrown + +#### `byokMode: 'disabled'` +- Always uses system keys +- Requires Codebuff API key + +### Database Schema + +```sql +CREATE TABLE encrypted_api_keys ( + user_id TEXT NOT NULL, + type TEXT NOT NULL, -- 'anthropic' | 'gemini' | 'openai' + encrypted_key TEXT NOT NULL, + iv TEXT NOT NULL, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (user_id, type) +); +``` + +### API Key Validation + +Each provider has specific validation rules: + +- **Anthropic**: Prefix `sk-ant-api03-`, length 108 +- **Gemini**: Prefix `AIzaSy`, length 39 +- **OpenAI**: Prefix `sk-proj-`, length 164 + +### Future Enhancements + +Potential improvements to the BYOK system: +1. Support for additional providers (Azure OpenAI, AWS Bedrock, etc.) +2. Per-model key configuration +3. Key rotation and expiration +4. Usage analytics per provider +5. Cost alerts and budgets +6. Key sharing within organizations + +## Testing BYOK + +### Unit Tests +- Test key validation logic +- Test encryption/decryption +- Test key precedence + +### Integration Tests +- Test end-to-end flow with real provider keys +- Test fallback behavior +- Test error handling for missing keys + +### Manual Testing +1. Configure keys via web UI +2. Run agent with different byokMode settings +3. Verify correct provider routing +4. Check cost calculations +5. Test key removal and updates + diff --git a/sdk/src/client.ts b/sdk/src/client.ts index 61776b762..6e866427c 100644 --- a/sdk/src/client.ts +++ b/sdk/src/client.ts @@ -7,15 +7,35 @@ import type { RunState } from './run-state' export class CodebuffClient { public options: CodebuffClientOptions & { - apiKey: string + apiKey?: string fingerprintId: string } constructor(options: CodebuffClientOptions) { const foundApiKey = options.apiKey ?? process.env[API_KEY_ENV_VAR] - if (!foundApiKey) { + const hasUserApiKeys = + options.userApiKeys && + Object.values(options.userApiKeys).some((key) => key) + const byokMode = options.byokMode ?? 'prefer' + + // Authentication validation + if (byokMode === 'disabled' && !foundApiKey) { + throw new Error( + `Codebuff API key required when byokMode is 'disabled'. Please provide an apiKey in the constructor of CodebuffClient or set the ${API_KEY_ENV_VAR} environment variable.`, + ) + } + + if (byokMode === 'require' && !hasUserApiKeys) { + throw new Error( + `User API keys required when byokMode is 'require'. Please provide at least one provider API key in userApiKeys.`, + ) + } + + if (!foundApiKey && !hasUserApiKeys) { throw new Error( - `Codebuff API key not found. Please provide an apiKey in the constructor of CodebuffClient or set the ${API_KEY_ENV_VAR} environment variable.`, + `Authentication required: provide either a Codebuff API key or user provider API keys.\n\n` + + `Option 1: Provide apiKey in constructor or set ${API_KEY_ENV_VAR} environment variable.\n` + + `Option 2: Provide userApiKeys with at least one provider key (anthropic, gemini, or openai).`, ) } @@ -29,6 +49,7 @@ export class CodebuffClient { } }, fingerprintId: `codebuff-sdk-${Math.random().toString(36).substring(2, 15)}`, + byokMode, ...options, } } diff --git a/sdk/src/run.ts b/sdk/src/run.ts index 52a68991e..8e4d261b1 100644 --- a/sdk/src/run.ts +++ b/sdk/src/run.ts @@ -36,6 +36,7 @@ import type { SessionState } from '../../common/src/types/session-state' export type CodebuffClientOptions = { // Provide an API key or set the CODEBUFF_API_KEY environment variable. + // Optional if userApiKeys are provided. apiKey?: string cwd?: string @@ -60,6 +61,20 @@ export type CodebuffClientOptions = { } > customToolDefinitions?: CustomToolDefinition[] + + // BYOK (Bring Your Own Key) options + // User-provided API keys for direct provider access + userApiKeys?: { + anthropic?: string + gemini?: string + openai?: string + } + + // BYOK mode controls fallback behavior + // - 'disabled': Always use system keys (requires Codebuff apiKey) + // - 'prefer': Use user keys when available, fallback to system keys (default) + // - 'require': Only use user keys, fail if missing (no Codebuff apiKey needed) + byokMode?: 'disabled' | 'prefer' | 'require' } export type RunOptions = { @@ -103,7 +118,7 @@ export async function run({ } } - let resolve: (value: RunReturnType) => any = () => {} + let resolve: (value: RunReturnType) => any = () => { } const promise = new Promise((res) => { resolve = res }) @@ -114,8 +129,8 @@ export async function run({ onWebsocketError: (error) => { onError({ message: error.message }) }, - onWebsocketReconnect: () => {}, - onRequestReconnect: async () => {}, + onWebsocketReconnect: () => { }, + onRequestReconnect: async () => { }, onResponseError: async (error) => { onError({ message: error.message }) }, @@ -131,12 +146,12 @@ export async function run({ overrides: overrideTools ?? {}, customToolDefinitions: customToolDefinitions ? Object.fromEntries( - customToolDefinitions.map((def) => [def.toolName, def]), - ) + customToolDefinitions.map((def) => [def.toolName, def]), + ) : {}, cwd, }), - onCostResponse: async () => {}, + onCostResponse: async () => { }, onResponseChunk: async (action) => { const { userInputId, chunk } = action @@ -146,7 +161,7 @@ export async function run({ await handleEvent?.(chunk) } }, - onSubagentResponseChunk: async () => {}, + onSubagentResponseChunk: async () => { }, onPromptResponse: (action) => handlePromptResponse({ @@ -211,6 +226,8 @@ export async function run({ sessionState, toolResults: extraToolResults ?? [], agentId, + userApiKeys, + byokMode, }) const result = await promise @@ -322,9 +339,9 @@ async function handleToolCall({ value: { errorMessage: error && - typeof error === 'object' && - 'message' in error && - typeof error.message === 'string' + typeof error === 'object' && + 'message' in error && + typeof error.message === 'string' ? error.message : typeof error === 'string' ? error diff --git a/web/src/app/api/user-api-keys/[keyType]/route.ts b/web/src/app/api/user-api-keys/[keyType]/route.ts new file mode 100644 index 000000000..840f87e7c --- /dev/null +++ b/web/src/app/api/user-api-keys/[keyType]/route.ts @@ -0,0 +1,72 @@ +import { + API_KEY_TYPES, + type ApiKeyType, + READABLE_NAME, +} from '@codebuff/common/api-keys/constants' +import { clearApiKey } from '@codebuff/common/api-keys/crypto' +import { getServerSession } from 'next-auth' +import { NextResponse } from 'next/server' + +import type { NextRequest } from 'next/server' + +import { authOptions } from '@/app/api/auth/[...nextauth]/auth-options' +import { logger } from '@/util/logger' + +interface RouteParams { + params: { + keyType: string + } +} + +/** + * DELETE /api/user-api-keys/:keyType + * Removes a specific API key for the authenticated user + */ +export async function DELETE( + request: NextRequest, + { params }: RouteParams, +) { + const session = await getServerSession(authOptions) + if (!session?.user?.id) { + return NextResponse.json({ error: 'Unauthorized' }, { status: 401 }) + } + + const userId = session.user.id + const { keyType } = params + + // Validate keyType + if (!API_KEY_TYPES.includes(keyType as ApiKeyType)) { + return NextResponse.json( + { + error: 'Invalid key type', + message: `Key type must be one of: ${API_KEY_TYPES.join(', ')}`, + }, + { status: 400 }, + ) + } + + try { + await clearApiKey(userId, keyType as ApiKeyType) + + logger.info( + { userId, keyType }, + 'Successfully removed user API key', + ) + + return NextResponse.json({ + success: true, + message: `${READABLE_NAME[keyType as ApiKeyType]} API key removed successfully`, + }) + } catch (error) { + logger.error({ error, userId, keyType }, 'Error removing user API key') + return NextResponse.json( + { + error: 'Failed to remove API key', + message: + error instanceof Error ? error.message : 'Internal server error', + }, + { status: 500 }, + ) + } +} + diff --git a/web/src/app/api/user-api-keys/route.ts b/web/src/app/api/user-api-keys/route.ts new file mode 100644 index 000000000..4d89a00e1 --- /dev/null +++ b/web/src/app/api/user-api-keys/route.ts @@ -0,0 +1,135 @@ +import { + API_KEY_TYPES, + type ApiKeyType, + READABLE_NAME, +} from '@codebuff/common/api-keys/constants' +import { + encryptAndStoreApiKey, + validateApiKey, +} from '@codebuff/common/api-keys/crypto' +import db from '@codebuff/common/db' +import * as schema from '@codebuff/common/db/schema' +import { eq } from 'drizzle-orm' +import { getServerSession } from 'next-auth' +import { NextResponse } from 'next/server' +import { z } from 'zod/v4' + +import type { NextRequest } from 'next/server' + +import { authOptions } from '@/app/api/auth/[...nextauth]/auth-options' +import { logger } from '@/util/logger' + +/** + * GET /api/user-api-keys + * Returns a list of configured API key types for the authenticated user + */ +export async function GET(request: NextRequest) { + const session = await getServerSession(authOptions) + if (!session?.user?.id) { + return NextResponse.json({ error: 'Unauthorized' }, { status: 401 }) + } + + const userId = session.user.id + + try { + // Fetch all encrypted API keys for this user + const userKeys = await db.query.encryptedApiKeys.findMany({ + where: eq(schema.encryptedApiKeys.user_id, userId), + columns: { + type: true, + }, + }) + + // Create a map of configured keys + const configuredKeys = new Set(userKeys.map((k) => k.type)) + + // Build response with all key types + const keys = API_KEY_TYPES.map((keyType) => ({ + type: keyType, + name: READABLE_NAME[keyType], + configured: configuredKeys.has(keyType), + })) + + return NextResponse.json({ keys }) + } catch (error) { + logger.error( + { error, userId }, + 'Error fetching user API keys configuration', + ) + return NextResponse.json( + { error: 'Internal server error' }, + { status: 500 }, + ) + } +} + +/** + * POST /api/user-api-keys + * Stores or updates an API key for the authenticated user + */ +export async function POST(request: NextRequest) { + const session = await getServerSession(authOptions) + if (!session?.user?.id) { + return NextResponse.json({ error: 'Unauthorized' }, { status: 401 }) + } + + const userId = session.user.id + + try { + const body = await request.json() + + // Validate request body + const schema = z.object({ + keyType: z.enum(API_KEY_TYPES), + apiKey: z.string().min(1, 'API key cannot be empty'), + }) + + const parseResult = schema.safeParse(body) + if (!parseResult.success) { + return NextResponse.json( + { + error: 'Invalid request body', + details: parseResult.error.errors, + }, + { status: 400 }, + ) + } + + const { keyType, apiKey } = parseResult.data + + // Validate API key format + if (!validateApiKey(keyType, apiKey)) { + return NextResponse.json( + { + error: `Invalid ${READABLE_NAME[keyType]} API key format`, + message: `Please check that your API key is correct and matches the expected format for ${READABLE_NAME[keyType]}.`, + }, + { status: 400 }, + ) + } + + // Encrypt and store the API key + await encryptAndStoreApiKey(userId, keyType, apiKey) + + logger.info( + { userId, keyType }, + 'Successfully stored user API key', + ) + + return NextResponse.json({ + success: true, + message: `${READABLE_NAME[keyType]} API key stored successfully`, + }) + } catch (error) { + logger.error({ error, userId }, 'Error storing user API key') + return NextResponse.json( + { + error: 'Failed to store API key', + message: + error instanceof Error ? error.message : 'Internal server error', + }, + { status: 500 }, + ) + } +} + diff --git a/web/src/app/profile/components/user-api-keys-section.tsx b/web/src/app/profile/components/user-api-keys-section.tsx new file mode 100644 index 000000000..551ebd736 --- /dev/null +++ b/web/src/app/profile/components/user-api-keys-section.tsx @@ -0,0 +1,309 @@ +'use client' + +import { useState } from 'react' +import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query' +import { Button } from '@/components/ui/button' +import { Input } from '@/components/ui/input' +import { Label } from '@/components/ui/label' +import { + Card, + CardContent, + CardDescription, + CardHeader, + CardTitle, +} from '@/components/ui/card' +import { useToast } from '@/components/ui/use-toast' +import { Check, X, AlertCircle, Key } from 'lucide-react' +import { ConfirmationDialog } from '@/components/ui/confirmation-dialog' +import { ProfileSection } from './profile-section' +import { Alert, AlertDescription } from '@/components/ui/alert' + +interface UserApiKey { + type: string + name: string + configured: boolean +} + +async function fetchUserApiKeys(): Promise<{ keys: UserApiKey[] }> { + const res = await fetch('/api/user-api-keys') + if (!res.ok) throw new Error(await res.text()) + return res.json() +} + +export function UserApiKeysSection() { + const { toast } = useToast() + const queryClient = useQueryClient() + + const { + data: keysData, + isLoading: loadingKeys, + error: keysError, + refetch: refetchKeys, + } = useQuery({ + queryKey: ['user-api-keys'], + queryFn: fetchUserApiKeys, + }) + + const [editingKey, setEditingKey] = useState(null) + const [keyValues, setKeyValues] = useState>({}) + const [removeDialogOpen, setRemoveDialogOpen] = useState(false) + const [keyToRemove, setKeyToRemove] = useState(null) + + const saveKeyMutation = useMutation({ + mutationFn: async ({ + keyType, + apiKey, + }: { + keyType: string + apiKey: string + }) => { + const res = await fetch('/api/user-api-keys', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ keyType, apiKey }), + }) + if (!res.ok) { + const errorData = await res.json() + throw new Error(errorData.message || errorData.error || 'Failed to save key') + } + return res.json() + }, + onSuccess: async (data, variables) => { + await queryClient.invalidateQueries({ queryKey: ['user-api-keys'] }) + setEditingKey(null) + setKeyValues((prev) => ({ ...prev, [variables.keyType]: '' })) + toast({ title: data.message || 'API key saved successfully' }) + }, + onError: (e: any) => { + toast({ + title: 'Failed to save API key', + description: e.message ?? String(e), + variant: 'destructive' as any, + }) + }, + }) + + const removeKeyMutation = useMutation({ + mutationFn: async (keyType: string) => { + const res = await fetch(`/api/user-api-keys/${keyType}`, { + method: 'DELETE', + }) + if (!res.ok) { + const errorData = await res.json() + throw new Error(errorData.message || errorData.error || 'Failed to remove key') + } + return res.json() + }, + onSuccess: async (data) => { + await queryClient.invalidateQueries({ queryKey: ['user-api-keys'] }) + setRemoveDialogOpen(false) + setKeyToRemove(null) + toast({ title: data.message || 'API key removed successfully' }) + }, + onError: (e: any) => { + toast({ + title: 'Failed to remove API key', + description: e.message ?? String(e), + variant: 'destructive' as any, + }) + }, + }) + + const handleSave = (keyType: string) => { + const apiKey = keyValues[keyType] + if (!apiKey || apiKey.trim() === '') { + toast({ + title: 'Invalid input', + description: 'Please enter an API key', + variant: 'destructive' as any, + }) + return + } + saveKeyMutation.mutate({ keyType, apiKey }) + } + + const handleRemove = (keyType: string) => { + setKeyToRemove(keyType) + setRemoveDialogOpen(true) + } + + const confirmRemove = () => { + if (keyToRemove) { + removeKeyMutation.mutate(keyToRemove) + } + } + + const getKeyPlaceholder = (keyType: string) => { + switch (keyType) { + case 'anthropic': + return 'sk-ant-api03-...' + case 'gemini': + return 'AIzaSy...' + case 'openai': + return 'sk-proj-...' + default: + return 'Enter your API key' + } + } + + const getKeyDescription = (keyType: string) => { + switch (keyType) { + case 'anthropic': + return 'Use your own Anthropic API key for Claude models. Get one at console.anthropic.com' + case 'gemini': + return 'Use your own Google API key for Gemini models. Get one at aistudio.google.com' + case 'openai': + return 'Use your own OpenAI API key for GPT models. Get one at platform.openai.com' + default: + return 'Use your own API key for this provider' + } + } + + return ( + + + + + Bring Your Own Key (BYOK): When you provide your own + API keys, you pay only for actual API usage through your provider + accounts. Codebuff applies a reduced markup compared to using our + system keys. Your keys are encrypted at rest using AES-256-GCM. + + + + {keysError && ( + + + + Error loading API keys: {(keysError as any)?.message ?? 'Please try again.'} + + + + )} + + {loadingKeys ? ( +
+ {[1, 2, 3].map((i) => ( + + +
+
+ +
+
+
+ ))} +
+ ) : ( +
+ {keysData?.keys.map((key) => ( + + +
+
+ + {key.name} + {key.configured && ( + + + Configured + + )} +
+ {key.configured && ( + + )} +
+ {getKeyDescription(key.type)} +
+ + {key.configured && editingKey !== key.type ? ( +
+ + +
+ ) : ( +
+ +
+ + setKeyValues((prev) => ({ + ...prev, + [key.type]: e.target.value, + })) + } + className="flex-1" + /> + + {editingKey === key.type && ( + + )} +
+
+ )} +
+
+ ))} +
+ )} + + +
+ ) +} + diff --git a/web/src/app/profile/page.tsx b/web/src/app/profile/page.tsx index ebc6d8967..0673f4680 100644 --- a/web/src/app/profile/page.tsx +++ b/web/src/app/profile/page.tsx @@ -11,6 +11,7 @@ import { SecuritySection } from './components/security-section' import { ReferralsSection } from './components/referrals-section' import { UsageSection } from './components/usage-section' import { ApiKeysSection } from './components/api-keys-section' +import { UserApiKeysSection } from './components/user-api-keys-section' import { ProfileLoggedOut } from './components/logged-out' import { Button } from '@/components/ui/button' import { Sheet, SheetContent, SheetTrigger } from '@/components/ui/sheet' @@ -35,6 +36,12 @@ const sections = [ icon: Key, component: ApiKeysSection, }, + { + id: 'user-api-keys', + title: 'Provider API Keys', + icon: Key, + component: UserApiKeysSection, + }, { id: 'referrals', title: 'Referrals',