diff --git a/src/helpers/config.test.ts b/src/helpers/config.test.ts index 6f5f714..24cd17b 100644 --- a/src/helpers/config.test.ts +++ b/src/helpers/config.test.ts @@ -1,7 +1,13 @@ import os from 'os'; import path from 'path'; import { expect, describe, it, vi, afterEach } from 'vitest'; -import { getConfig, setConfigs, showConfigUI } from './config'; +import { + getConfig, + hasOpenAiKey, + normalizeOpenAiKey, + setConfigs, + showConfigUI, +} from './config'; const mocks = vi.hoisted(() => { return { @@ -39,6 +45,29 @@ global.process = { ...realProcess, exit: mocks.exit }; const configFilePath = path.join(os.homedir(), '.micro-agent'); const newline = os.platform() === 'win32' ? '\r\n' : '\n'; +describe('normalizeOpenAiKey', () => { + it('should trim whitespace around the key', () => { + expect(normalizeOpenAiKey(' sk-test-key ')).toBe('sk-test-key'); + }); +}); + +describe('hasOpenAiKey', () => { + it('should return false for empty key values', () => { + expect(hasOpenAiKey('')).toBe(false); + expect(hasOpenAiKey(' ')).toBe(false); + expect(hasOpenAiKey(undefined)).toBe(false); + }); + + it('should return false when key is cancel', () => { + expect(hasOpenAiKey('cancel')).toBe(false); + expect(hasOpenAiKey(' CANCEL ')).toBe(false); + }); + + it('should return true for valid key values', () => { + expect(hasOpenAiKey('sk-test-key')).toBe(true); + }); +}); + describe('getConfig', () => { const defaultConfig = { ANTHROPIC_KEY: undefined, diff --git a/src/helpers/config.ts b/src/helpers/config.ts index b425620..206bb7d 100644 --- a/src/helpers/config.ts +++ b/src/helpers/config.ts @@ -11,9 +11,16 @@ const { hasOwnProperty } = Object.prototype; export const hasOwn = (object: unknown, key: PropertyKey) => hasOwnProperty.call(object, key); +export const normalizeOpenAiKey = (key?: string) => key?.trim(); + +export const hasOpenAiKey = (key?: string) => { + const normalizedKey = normalizeOpenAiKey(key); + return !!normalizedKey && normalizedKey.toLowerCase() !== 'cancel'; +}; + const configParsers = { OPENAI_KEY(key?: string) { - return key; + return normalizeOpenAiKey(key); }, ANTHROPIC_KEY(key?: string) { return key; @@ -155,13 +162,17 @@ export const showConfigUI = async () => { const key = await p.text({ message: 'Enter your OpenAI API key', validate: (value) => { - if (!value.length) { + const normalizedKey = normalizeOpenAiKey(value); + if (!normalizedKey) { return 'Please enter a key'; } + if (normalizedKey.toLowerCase() === 'cancel') { + return 'Press Ctrl+C to cancel'; + } }, }); if (p.isCancel(key)) return; - await setConfigs([['OPENAI_KEY', key]]); + await setConfigs([['OPENAI_KEY', normalizeOpenAiKey(key as string)!]]); } else if (choice === 'OPENAI_API_ENDPOINT') { const apiEndpoint = await p.text({ message: 'Enter your OpenAI API Endpoint', diff --git a/src/helpers/interactive-mode.ts b/src/helpers/interactive-mode.ts index 9a5dd89..58f1e84 100644 --- a/src/helpers/interactive-mode.ts +++ b/src/helpers/interactive-mode.ts @@ -3,7 +3,12 @@ import { intro, log, spinner, text } from '@clack/prompts'; import { glob } from 'glob'; import { RunOptions, runAll } from './run'; import { getFileSuggestion, getSimpleCompletion } from './llm'; -import { getConfig, setConfigs } from './config'; +import { + getConfig, + hasOpenAiKey, + normalizeOpenAiKey, + setConfigs, +} from './config'; import { readFile } from 'fs/promises'; import dedent from 'dedent'; import { formatMessage } from './test'; @@ -21,16 +26,27 @@ export async function interactiveMode(options: Partial) { const config = await getConfig(); - if (!config.OPENAI_KEY) { + if (!hasOpenAiKey(config.OPENAI_KEY)) { const openaiKey = exitOnCancel( await text({ message: `Welcome newcomer! What is your OpenAI key? ${gray( - '(this is kept private)' + '(this is kept private, press Ctrl+C to cancel)' )}`, + validate: (value) => { + const normalizedKey = normalizeOpenAiKey(value); + if (!normalizedKey) { + return 'Please enter an API key'; + } + if (normalizedKey.toLowerCase() === 'cancel') { + return 'Press Ctrl+C to cancel'; + } + }, }) ); - await setConfigs([['OPENAI_KEY', openaiKey as string]]); + await setConfigs([ + ['OPENAI_KEY', normalizeOpenAiKey(openaiKey as string) ?? ''], + ]); } const prompt = exitOnCancel( diff --git a/src/helpers/llm.test.ts b/src/helpers/llm.test.ts index 3b4643d..96217d4 100644 --- a/src/helpers/llm.test.ts +++ b/src/helpers/llm.test.ts @@ -15,6 +15,8 @@ const mocks = vi.hoisted(() => { return { openAIConstructor: vi.fn(), getConfig: vi.fn(), + hasOpenAiKey: vi.fn(), + normalizeOpenAiKey: vi.fn(), create: vi.fn(), }; }); @@ -22,6 +24,8 @@ const mocks = vi.hoisted(() => { vi.mock('./config', () => { return { getConfig: mocks.getConfig, + hasOpenAiKey: mocks.hasOpenAiKey, + normalizeOpenAiKey: mocks.normalizeOpenAiKey, }; }); @@ -41,6 +45,9 @@ mocks.openAIConstructor.mockImplementation(() => { }; }); +mocks.hasOpenAiKey.mockImplementation((key?: string) => !!key); +mocks.normalizeOpenAiKey.mockImplementation((key?: string) => key?.trim()); + const defaultConfig = { OPENAI_KEY: 'my-openai-key', OPENAI_API_ENDPOINT: 'https://api.openai.com/v1', @@ -51,15 +58,25 @@ describe('getOpenAi', () => { mocks.getConfig .mockResolvedValueOnce({ OPENAI_KEY: '' }) .mockResolvedValueOnce({ OPENAI_KEY: '' }); + mocks.hasOpenAiKey.mockReturnValueOnce(false).mockReturnValueOnce(false); await expect(getOpenAi()).rejects.toThrow(KnownError); await expect(getOpenAi()).rejects.toThrow( - 'Missing OpenAI key. Use `micro-agent config` to set it.' + 'Missing or invalid OpenAI key. Use `micro-agent config` to set it.' ); }); + it('should throw a KnownError if OPENAI_KEY is cancel', async () => { + mocks.getConfig.mockResolvedValueOnce({ OPENAI_KEY: 'cancel' }); + mocks.hasOpenAiKey.mockReturnValueOnce(false); + + await expect(getOpenAi()).rejects.toThrow(KnownError); + }); + it('should create a new OpenAI instance with the provided key and endpoint', async () => { mocks.getConfig.mockResolvedValueOnce(defaultConfig); + mocks.hasOpenAiKey.mockReturnValueOnce(true); + mocks.normalizeOpenAiKey.mockReturnValueOnce('my-openai-key'); await getOpenAi(); diff --git a/src/helpers/llm.ts b/src/helpers/llm.ts index 5db97ed..b531de0 100644 --- a/src/helpers/llm.ts +++ b/src/helpers/llm.ts @@ -1,5 +1,5 @@ import OpenAI, { AzureOpenAI } from 'openai'; -import { getConfig } from './config'; +import { getConfig, hasOpenAiKey, normalizeOpenAiKey } from './config'; import { KnownError } from './error'; import { commandName } from './constants'; import { systemPrompt } from './systemPrompt'; @@ -49,11 +49,12 @@ const supportsFunctionCalling = (model?: string) => { export const getOpenAi = async function () { const { OPENAI_KEY: openaiKey, OPENAI_API_ENDPOINT: endpoint } = await getConfig(); - if (!openaiKey) { + if (!hasOpenAiKey(openaiKey)) { throw new KnownError( - `Missing OpenAI key. Use \`${commandName} config\` to set it.` + `Missing or invalid OpenAI key. Use \`${commandName} config\` to set it.` ); } + const normalizedOpenAiKey = normalizeOpenAiKey(openaiKey)!; if (endpoint.indexOf('.openai.azure.com/openai') > 0) { const deploymentName = endpoint.split('/deployments/')[1].split('/')[0]; const apiVersion = endpoint.split('api-version=')[1]; @@ -66,14 +67,14 @@ export const getOpenAi = async function () { ); } return new AzureOpenAI({ - apiKey: openaiKey, + apiKey: normalizedOpenAiKey, endpoint: endpoint, deployment: `/deployments/${deploymentName}`, apiVersion: apiVersion, }); } else { return new OpenAI({ - apiKey: openaiKey, + apiKey: normalizedOpenAiKey, baseURL: endpoint, }); } diff --git a/src/tests/integration/interactive.test.ts b/src/tests/integration/interactive.test.ts index 5df5a17..ba0472e 100644 --- a/src/tests/integration/interactive.test.ts +++ b/src/tests/integration/interactive.test.ts @@ -1,9 +1,11 @@ import { execaCommand } from 'execa'; -import { lstat, writeFile } from 'fs/promises'; +import { lstat, readFile, writeFile } from 'fs/promises'; import { beforeAll, describe, expect, it } from 'vitest'; +const configPath = `${process.env.HOME}/.micro-agent`; + const checkConfigFileExists = async () => { - return await lstat(`${process.env.HOME}/.micro-agent`) + return await lstat(configPath) .then(() => true) .catch(() => false); }; @@ -13,7 +15,7 @@ describe('interactive cli', () => { const configFileExists = await checkConfigFileExists(); if (!configFileExists) { await writeFile( - `${process.env.HOME}/.micro-agent`, + configPath, 'OPENAI_KEY=sk-1234567890abcdef1234567890abcdef' ); } @@ -59,4 +61,25 @@ describe('interactive cli', () => { expect(output).toContain('What would you like to do?'); }); + + it('should ask for an OpenAI key if configured key is cancel', async () => { + const previousConfig = await readFile(configPath, 'utf8'); + await writeFile(configPath, 'OPENAI_KEY=cancel'); + + try { + const result = await execaCommand('jiti ./src/cli.ts', { + input: '\x03', + shell: process.env.SHELL || true, + env: { + ...process.env, + OPENAI_KEY: '', + }, + }); + + const output = result.stdout; + expect(output).toContain('Welcome newcomer! What is your OpenAI key?'); + } finally { + await writeFile(configPath, previousConfig); + } + }); });