From 73432a644a8502501b36a4c26cfce83b6379fce3 Mon Sep 17 00:00:00 2001 From: Tobi Okedeji Date: Thu, 12 Feb 2026 01:01:08 +0100 Subject: [PATCH] feat: add --anthropic-key option for direct Anthropic API access (BYOK) --- bin.ts | 5 ++ src/__tests__/cli.test.ts | 53 +++++++++++++++++++ src/lib/__tests__/agent-interface.test.ts | 64 ++++++++++++++++++++++- src/lib/agent-interface.ts | 45 +++++++++++----- src/lib/agent-runner.ts | 8 ++- src/run.ts | 6 +++ src/utils/types.ts | 5 ++ 7 files changed, 170 insertions(+), 16 deletions(-) diff --git a/bin.ts b/bin.ts index 7c82ae0..07ad155 100644 --- a/bin.ts +++ b/bin.ts @@ -79,6 +79,11 @@ yargs(hideBin(process.argv)) 'PostHog personal API key (phx_xxx) for authentication\nenv: POSTHOG_WIZARD_API_KEY', type: 'string', }, + 'anthropic-key': { + describe: + 'Anthropic API key for direct API access (bypasses LLM gateway)\nenv: POSTHOG_WIZARD_ANTHROPIC_KEY', + type: 'string', + }, }) .command( ['$0'], diff --git a/src/__tests__/cli.test.ts b/src/__tests__/cli.test.ts index 98b965b..ee8d30d 100644 --- a/src/__tests__/cli.test.ts +++ b/src/__tests__/cli.test.ts @@ -25,6 +25,7 @@ describe('CLI argument parsing', () => { delete process.env.POSTHOG_WIZARD_DEFAULT; delete process.env.POSTHOG_WIZARD_CI; delete process.env.POSTHOG_WIZARD_API_KEY; + delete process.env.POSTHOG_WIZARD_ANTHROPIC_KEY; delete process.env.POSTHOG_WIZARD_INSTALL_DIR; // Mock process.exit to prevent test runner from exiting @@ -254,6 +255,58 @@ describe('CLI argument parsing', () => { }); }); + describe('--anthropic-key flag', () => { + test('is undefined when not specified', async () => { + await runCLI([]); + + const args = getLastCallArgs(mockRunWizard); + expect(args.anthropicKey).toBeUndefined(); + }); + + test('passes --anthropic-key to runWizard', async () => { + await runCLI(['--anthropic-key', 'sk-ant-test123']); + + const args = getLastCallArgs(mockRunWizard); + expect(args.anthropicKey).toBe('sk-ant-test123'); + }); + + test('works alongside --ci mode', async () => { + await runCLI([ + '--ci', + '--region', + 'us', + '--api-key', + 'phx_test', + '--install-dir', + '/tmp/test', + '--anthropic-key', + 'sk-ant-test123', + ]); + + const args = getLastCallArgs(mockRunWizard); + expect(args.ci).toBe(true); + expect(args.anthropicKey).toBe('sk-ant-test123'); + }); + + test('respects POSTHOG_WIZARD_ANTHROPIC_KEY env var', async () => { + process.env.POSTHOG_WIZARD_ANTHROPIC_KEY = 'sk-ant-env123'; + + await runCLI([]); + + const args = getLastCallArgs(mockRunWizard); + expect(args.anthropicKey).toBe('sk-ant-env123'); + }); + + test('CLI arg overrides POSTHOG_WIZARD_ANTHROPIC_KEY env var', async () => { + process.env.POSTHOG_WIZARD_ANTHROPIC_KEY = 'sk-ant-env123'; + + await runCLI(['--anthropic-key', 'sk-ant-cli456']); + + const args = getLastCallArgs(mockRunWizard); + expect(args.anthropicKey).toBe('sk-ant-cli456'); + }); + }); + describe('CI environment variables', () => { test('respects POSTHOG_WIZARD_CI', async () => { process.env.POSTHOG_WIZARD_CI = 'true'; diff --git a/src/lib/__tests__/agent-interface.test.ts b/src/lib/__tests__/agent-interface.test.ts index beed455..095a070 100644 --- a/src/lib/__tests__/agent-interface.test.ts +++ b/src/lib/__tests__/agent-interface.test.ts @@ -1,10 +1,14 @@ -import { runAgent } from '../agent-interface'; +import { runAgent, initializeAgent } from '../agent-interface'; import type { WizardOptions } from '../../utils/types'; // Mock dependencies jest.mock('../../utils/clack'); jest.mock('../../utils/analytics'); jest.mock('../../utils/debug'); +jest.mock('../env-file-tools', () => ({ + createEnvFileServer: jest.fn().mockResolvedValue({}), + ENV_FILE_TOOL_NAMES: [], +})); // Mock the SDK module const mockQuery = jest.fn(); @@ -61,6 +65,64 @@ describe('runAgent', () => { }; }); + describe('BYOK vs gateway mode', () => { + const defaultConfig = { + workingDirectory: '/test/dir', + posthogMcpUrl: 'http://localhost:8787/mcp', + posthogApiKey: 'phx_test_key', + posthogApiHost: 'http://localhost:8000', + }; + + const originalEnv = { ...process.env }; + + afterEach(() => { + process.env = { ...originalEnv }; + }); + + it('should configure direct Anthropic API in BYOK mode', async () => { + const options: WizardOptions = { + ...defaultOptions, + anthropicKey: 'sk-ant-test123', + }; + + await initializeAgent(defaultConfig, options); + + expect(process.env.ANTHROPIC_BASE_URL).toBe('https://api.anthropic.com'); + expect(process.env.ANTHROPIC_API_KEY).toBe('sk-ant-test123'); + expect(process.env.ANTHROPIC_AUTH_TOKEN).toBeUndefined(); + expect(process.env.CLAUDE_CODE_OAUTH_TOKEN).toBeUndefined(); + }); + + it('should configure LLM gateway when no anthropicKey provided', async () => { + await initializeAgent(defaultConfig, defaultOptions); + + expect(process.env.ANTHROPIC_BASE_URL).toBe( + 'http://localhost:3308/wizard', + ); + expect(process.env.ANTHROPIC_AUTH_TOKEN).toBe('phx_test_key'); + expect(process.env.CLAUDE_CODE_OAUTH_TOKEN).toBe('phx_test_key'); + // LLM gateway doesn't support experimental betas + expect(process.env.CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS).toBe('true'); + expect(process.env.ANTHROPIC_API_KEY).toBeUndefined(); + }); + + it('should clear gateway env vars when switching to BYOK mode', async () => { + // Simulate a prior gateway run that set these env vars + process.env.ANTHROPIC_AUTH_TOKEN = 'phx_old_token'; + process.env.CLAUDE_CODE_OAUTH_TOKEN = 'phx_old_token'; + + const options: WizardOptions = { + ...defaultOptions, + anthropicKey: 'sk-ant-test123', + }; + + await initializeAgent(defaultConfig, options); + + expect(process.env.ANTHROPIC_AUTH_TOKEN).toBeUndefined(); + expect(process.env.CLAUDE_CODE_OAUTH_TOKEN).toBeUndefined(); + }); + }); + describe('race condition handling', () => { it('should return success when agent completes successfully then SDK cleanup fails', async () => { // This simulates the race condition: diff --git a/src/lib/agent-interface.ts b/src/lib/agent-interface.ts index 8c36715..d67525b 100644 --- a/src/lib/agent-interface.ts +++ b/src/lib/agent-interface.ts @@ -335,16 +335,31 @@ export async function initializeAgent( clack.log.step('Initializing Claude agent...'); try { - // Configure LLM gateway environment variables (inherited by SDK subprocess) - const gatewayUrl = getLlmGatewayUrlFromHost(config.posthogApiHost); - process.env.ANTHROPIC_BASE_URL = gatewayUrl; - process.env.ANTHROPIC_AUTH_TOKEN = config.posthogApiKey; - // Use CLAUDE_CODE_OAUTH_TOKEN to override any stored /login credentials - process.env.CLAUDE_CODE_OAUTH_TOKEN = config.posthogApiKey; - // Disable experimental betas (like input_examples) that the LLM gateway doesn't support - process.env.CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS = 'true'; - - logToFile('Configured LLM gateway:', gatewayUrl); + // Configure LLM environment variables (inherited by SDK subprocess) + let baseUrl: string; + + if (options.anthropicKey) { + // BYOK: direct Anthropic API + baseUrl = 'https://api.anthropic.com'; + process.env.ANTHROPIC_BASE_URL = baseUrl; + process.env.ANTHROPIC_API_KEY = options.anthropicKey; + delete process.env.ANTHROPIC_AUTH_TOKEN; + delete process.env.CLAUDE_CODE_OAUTH_TOKEN; + + logToFile('Configured BYOK mode:', baseUrl); + } else { + // Default: LLM gateway + baseUrl = getLlmGatewayUrlFromHost(config.posthogApiHost); + process.env.ANTHROPIC_BASE_URL = baseUrl; + process.env.ANTHROPIC_AUTH_TOKEN = config.posthogApiKey; + // Use CLAUDE_CODE_OAUTH_TOKEN to override any stored /login credentials + process.env.CLAUDE_CODE_OAUTH_TOKEN = config.posthogApiKey; + // Disable experimental betas (like input_examples) that the LLM gateway doesn't support + process.env.CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS = 'true'; + delete process.env.ANTHROPIC_API_KEY; + + logToFile('Configured LLM gateway:', baseUrl); + } // Configure MCP server with PostHog authentication const mcpServers: McpServersConfig = { @@ -375,7 +390,8 @@ export async function initializeAgent( logToFile('Agent config:', { workingDirectory: agentRunConfig.workingDirectory, posthogMcpUrl: config.posthogMcpUrl, - gatewayUrl, + mode: options.anthropicKey ? 'byok' : 'gateway', + baseUrl, apiKeyPresent: !!config.posthogApiKey, }); @@ -383,7 +399,8 @@ export async function initializeAgent( debug('Agent config:', { workingDirectory: agentRunConfig.workingDirectory, posthogMcpUrl: config.posthogMcpUrl, - gatewayUrl, + mode: options.anthropicKey ? 'byok' : 'gateway', + baseUrl, apiKeyPresent: !!config.posthogApiKey, }); } @@ -535,10 +552,10 @@ export async function runAgent( settingSources: ['project'], // Explicitly enable required tools including Skill allowedTools, + // Env var cleanup (ANTHROPIC_API_KEY, AUTH_TOKEN, etc.) is already handled + // by initializeAgent based on BYOK vs gateway mode env: { ...process.env, - // Prevent user's Anthropic API key from overriding the wizard's OAuth token - ANTHROPIC_API_KEY: undefined, }, canUseTool: (toolName: string, input: unknown) => { logToFile('canUseTool called:', { toolName, input }); diff --git a/src/lib/agent-runner.ts b/src/lib/agent-runner.ts index 46ed98b..5509790 100644 --- a/src/lib/agent-runner.ts +++ b/src/lib/agent-runner.ts @@ -85,7 +85,9 @@ export async function runAgentWizard( } clack.log.info( - `We're about to read your project using our LLM gateway.\n\n.env* file contents will not leave your machine.\n\nOther files will be read and edited to provide a fully-custom PostHog integration.`, + options.anthropicKey + ? `We're about to read your project using your Anthropic API key.\n\n.env* file contents will not leave your machine.\n\nOther files will be read and edited to provide a fully-custom PostHog integration.` + : `We're about to read your project using our LLM gateway.\n\n.env* file contents will not leave your machine.\n\nOther files will be read and edited to provide a fully-custom PostHog integration.`, ); const aiConsent = await askForAIConsent(options); @@ -135,6 +137,10 @@ export async function runAgentWizard( analytics.setTag(`${config.metadata.integration}-version`, versionBucket); } + if (options.anthropicKey) { + analytics.setTag('byok', true); + } + analytics.capture(WIZARD_INTERACTION_EVENT_NAME, { action: 'started agent integration', integration: config.metadata.integration, diff --git a/src/run.ts b/src/run.ts index e5e5311..e9a4166 100644 --- a/src/run.ts +++ b/src/run.ts @@ -26,6 +26,7 @@ type Args = { localMcp?: boolean; ci?: boolean; apiKey?: string; + anthropicKey?: string; menu?: boolean; }; @@ -56,6 +57,7 @@ export async function runWizard(argv: Args) { localMcp: finalArgs.localMcp ?? false, ci: finalArgs.ci ?? false, apiKey: finalArgs.apiKey, + anthropicKey: finalArgs.anthropicKey, menu: finalArgs.menu ?? false, }; @@ -65,6 +67,10 @@ export async function runWizard(argv: Args) { clack.log.info(chalk.dim('Running in CI mode')); } + if (wizardOptions.anthropicKey) { + clack.log.info(chalk.dim('Using your Anthropic API key (BYOK mode)')); + } + const integration = finalArgs.integration ?? (await getIntegrationForSetup(wizardOptions)); diff --git a/src/utils/types.ts b/src/utils/types.ts index 7379156..c9dc0dd 100644 --- a/src/utils/types.ts +++ b/src/utils/types.ts @@ -56,6 +56,11 @@ export type WizardOptions = { */ apiKey?: string; + /** + * Anthropic API key for direct API access, bypasses PostHog LLM gateway + */ + anthropicKey?: string; + /** * Whether to show the menu for manual integration selection instead of auto-detecting. */