Implements comprehensive BYOK (Bring Your Own Key) functionality, all…

[walkerRnD]

🔑 Add BYOK (Bring Your Own Key) Support for AI Provider APIs

Summary

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.

Motivation

• Cost Savings: Users pay only for actual API usage with reduced markup
• Flexibility: Users can leverage their existing provider accounts and rate limits
• No Codebuff API Key Required: In require mode, users can authenticate solely with provider keys
• Direct Billing: Charges appear directly on provider accounts for better cost tracking

Changes

🎨 Frontend (Web UI)

New Files:

• web/src/app/api/user-api-keys/route.ts - REST API for managing user keys (GET, POST)
• web/src/app/api/user-api-keys/[keyType]/route.ts - REST API for deleting keys (DELETE)
• web/src/app/profile/components/user-api-keys-section.tsx - Card-based UI for key management

Modified Files:

• web/src/app/profile/page.tsx - Added "Provider API Keys" section to profile

📦 SDK Layer

Modified Files:

• sdk/src/client.ts - Added userApiKeys and byokMode options with authentication logic
• sdk/src/run.ts - Extended CodebuffClientOptions to accept BYOK parameters
• sdk/README.md - Added comprehensive BYOK documentation and examples

New Files:

• sdk/knowledge.md - Detailed architecture documentation for BYOK

🔧 Backend

Modified Files:

• backend/package.json - Added @ai-sdk/[email protected] dependency
• backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts - Implemented direct-to-provider routing with BYOK support
• backend/src/llm-apis/message-cost-tracker.ts - Added per-provider cost tracking with reduced markup
• backend/src/main-prompt.ts - Added getUserApiKeys() for DB key retrieval and merging
• backend/src/run-agent-step.ts - Pass BYOK parameters through agent execution
• backend/src/prompt-agent-stream.ts - Pass BYOK parameters to AI SDK functions

🔐 Common Layer

Modified Files:

• common/src/actions.ts - Added userApiKeys and byokMode to CLIENT_ACTION_SCHEMA
• common/src/api-keys/crypto.ts - Added validateApiKey() function for format validation

📝 Documentation

New Files:

• .plan/byok-implementation-summary.md - Complete implementation summary

Key Features

1. Three BYOK Modes

// 'disabled' - Always use system keys (requires Codebuff API key)
const client = new CodebuffClient({
  apiKey: process.env.CODEBUFF_API_KEY,
  byokMode: 'disabled',
})

// 'prefer' - Use user keys when available, fallback to system (default)
const client = new CodebuffClient({
  apiKey: process.env.CODEBUFF_API_KEY,
  userApiKeys: { anthropic: process.env.ANTHROPIC_API_KEY },
  byokMode: 'prefer',
})

// 'require' - Only use user keys, no Codebuff API key needed
const client = new CodebuffClient({
  userApiKeys: { anthropic: process.env.ANTHROPIC_API_KEY },
  byokMode: 'require',
})

2. Direct-to-Provider Routing

• Anthropic models with user key → @ai-sdk/anthropic (direct)
• Gemini models with user key → @ai-sdk/google (direct)
• OpenAI models with user key → @ai-sdk/openai (direct)
• All models without user key → System keys or OpenRouter

3. Key Precedence

1. SDK-provided keys (highest priority)
2. Database-stored keys (medium priority)
3. System keys (fallback)

4. Security

• ✅ AES-256-GCM encryption for keys at rest
• ✅ Format validation (prefix and length) per provider
• ✅ Keys never logged or exposed in error messages
• ✅ Secure WebSocket transmission

5. Cost Tracking

// Without BYOK (system keys)
costInCents = cost * 100 * (1 + PROFIT_MARGIN)

// With BYOK (user keys)
costInCents = cost * 100 * PROFIT_MARGIN  // Reduced markup

API Changes

SDK Constructor

new CodebuffClient({
  apiKey?: string,  // Now optional if userApiKeys provided
  userApiKeys?: {
    anthropic?: string,  // sk-ant-api03-...
    gemini?: string,     // AIzaSy...
    openai?: string,     // sk-proj-...
  },
  byokMode?: 'disabled' | 'prefer' | 'require',  // Default: 'prefer'
  cwd?: string,
})

Web API Endpoints

• GET /api/user-api-keys - List configured key types
• POST /api/user-api-keys - Add/update API key
• DELETE /api/user-api-keys/[keyType] - Remove API key

Database Schema

Uses existing encrypted_api_keys table:

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

Testing

Manual Testing Steps

1. Navigate to /profile?tab=user-api-keys
2. Add provider API keys via web UI
3. Run SDK with different byokMode settings
4. Verify correct provider routing in logs
5. Check cost calculations show reduced markup
6. Test key removal and updates

Recommended Unit Tests

• Test validateApiKey() with valid/invalid keys
• Test encryption/decryption functions
• Test key precedence in getUserApiKeys()
• Test determineByokProvider() logic
• Test cost calculation with/without BYOK

Recommended Integration Tests

• End-to-end flow with real provider keys
• All three byokMode settings
• Fallback behavior when keys missing
• Error handling for invalid keys
• Key storage/retrieval via web API

Breaking Changes

None. This is a fully backward-compatible addition. Existing code continues to work without any changes.

Migration Guide

No migration needed. To adopt BYOK:

1. Via Web UI: Add keys at /profile?tab=user-api-keys
2. Via SDK: Pass userApiKeys in constructor or run() options

Documentation

• ✅ Updated sdk/README.md with usage examples
• ✅ Created sdk/knowledge.md with architecture details
• ✅ Created .plan/byok-implementation-summary.md with implementation summary

Future Enhancements

• Support for additional providers (Azure OpenAI, AWS Bedrock)
• Per-model key configuration
• Key rotation and expiration
• Usage analytics per provider
• Cost alerts and budgets
• Organization-level key sharing

Checklist

• Phase 1: Backend API Endpoints
• Phase 2: Web UI Components
• Phase 3: SDK Configuration
• Phase 4: Backend Integration
• Phase 5: Documentation
• Unit tests
• Integration tests
• Manual testing
• Staging deployment

[jroth55]

Looking forward to this!