From e7b565db98aa51f2599e110f86d40e8185590c39 Mon Sep 17 00:00:00 2001 From: Koji Date: Wed, 1 Oct 2025 22:31:33 -0300 Subject: [PATCH] Implements comprehensive BYOK (Bring Your Own Key) functionality, allowing SDK users to provide their own API keys for Anthropic, Gemini, and OpenAI models. Users can now use Codebuff's agent infrastructure while paying directly for LLM API costs through their own provider accounts, with reduced or zero Codebuff markup. --- backend/package.json | 3 +- backend/src/llm-apis/message-cost-tracker.ts | 22 +- backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts | 138 +++++++- backend/src/main-prompt.ts | 41 +++ backend/src/prompt-agent-stream.ts | 8 +- backend/src/run-agent-step.ts | 62 ++-- bun.lock | 31 +- common/src/actions.ts | 9 + common/src/api-keys/crypto.ts | 23 ++ sdk/README.md | 93 ++++++ sdk/knowledge.md | 223 +++++++++++++ sdk/src/client.ts | 27 +- sdk/src/run.ts | 37 ++- .../app/api/user-api-keys/[keyType]/route.ts | 72 ++++ web/src/app/api/user-api-keys/route.ts | 135 ++++++++ .../components/user-api-keys-section.tsx | 309 ++++++++++++++++++ web/src/app/profile/page.tsx | 7 + 17 files changed, 1168 insertions(+), 72 deletions(-) create mode 100644 sdk/knowledge.md create mode 100644 web/src/app/api/user-api-keys/[keyType]/route.ts create mode 100644 web/src/app/api/user-api-keys/route.ts create mode 100644 web/src/app/profile/components/user-api-keys-section.tsx 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 8e4b426df..f91af7486 100644 --- a/backend/src/llm-apis/message-cost-tracker.ts +++ b/backend/src/llm-apis/message-cost-tracker.ts @@ -573,7 +573,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 @@ -598,16 +599,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 43d49edfd..aa02493d7 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',