keystone-cli 1.3.0 → 2.0.0

package/src/db/memory-db.ts

@@ -3,7 +3,8 @@ import { randomUUID } from 'node:crypto';
 import { existsSync, mkdirSync } from 'node:fs';
 import { dirname, join } from 'node:path';
 import * as sqliteVec from 'sqlite-vec';
-import './sqlite-setup.ts';
+import { ConsoleLogger } from '../utils/logger';
+import { setupSqlite } from './sqlite-setup.ts';
 
 export interface MemoryEntry {
   id: string;
@@ -64,9 +65,16 @@ export class MemoryDb {
   private db: Database;
   // Cache connections by path to avoid reloading extensions
   private static connectionCache = new Map<string, { db: Database; refCount: number }>();
-  static readonly EMBEDDING_DIMENSION = 384;
+  private tableName: string;
 
-  constructor(public readonly dbPath = '.keystone/memory.db') {
+  constructor(
+    public readonly dbPath = '.keystone/memory.db',
+    private readonly embeddingDimension = 384
+  ) {
+    // Ensure SQLite is set up with custom library on macOS (idempotent)
+    setupSqlite();
+
+    this.tableName = `vec_memory_${embeddingDimension}`;
     const cached = MemoryDb.connectionCache.get(dbPath);
     if (cached) {
       cached.refCount++;
@@ -89,10 +97,36 @@ export class MemoryDb {
   }
 
   private initSchema(): void {
+    // Check if the legacy 'vec_memory' table exists and what its dimension is
+    const legacyTable = this.db
+      .prepare("SELECT sql FROM sqlite_master WHERE type='table' AND name='vec_memory'")
+      .get() as { sql: string } | undefined;
+
+    if (legacyTable) {
+      const match = legacyTable.sql.match(/FLOAT\[(\d+)\]/i);
+      if (match && Number.parseInt(match[1], 10) === this.embeddingDimension) {
+        // Legacy table exists and matches our dimension, reuse it
+        this.tableName = 'vec_memory';
+      } else {
+        // Mismatch or couldn't parse. We will use the specific table name `vec_memory_{dim}`.
+        // We log a warning to stdout since we don't have a logger instance here,
+        // but only if we haven't already created the specific table (to avoid spamming on every init).
+        const specificTableExists = this.db
+          .prepare(`SELECT name FROM sqlite_master WHERE type='table' AND name='${this.tableName}'`)
+          .get();
+        if (!specificTableExists) {
+          new ConsoleLogger().warn(
+            `\n⚠️  Vector DB: Found legacy table 'vec_memory' with dimension mismatch (expected ${this.embeddingDimension}).\n` +
+              `Using new table '${this.tableName}' instead. Old data is preserved in 'vec_memory'.\n`
+          );
+        }
+      }
+    }
+
     this.db.run(`
-      CREATE VIRTUAL TABLE IF NOT EXISTS vec_memory USING vec0(
+      CREATE VIRTUAL TABLE IF NOT EXISTS ${this.tableName} USING vec0(
         id TEXT PRIMARY KEY,
-        embedding FLOAT[${MemoryDb.EMBEDDING_DIMENSION}]
+        embedding FLOAT[${this.embeddingDimension}]
       );
     `);
 
@@ -106,10 +140,10 @@ export class MemoryDb {
     `);
   }
 
-  private static assertEmbeddingDimension(embedding: number[]): void {
-    if (embedding.length !== MemoryDb.EMBEDDING_DIMENSION) {
+  private assertEmbeddingDimension(embedding: number[]): void {
+    if (embedding.length !== this.embeddingDimension) {
       throw new Error(
-        `Embedding dimension mismatch: expected ${MemoryDb.EMBEDDING_DIMENSION}, got ${embedding.length}`
+        `Embedding dimension mismatch: expected ${this.embeddingDimension}, got ${embedding.length}`
       );
     }
   }
@@ -117,12 +151,8 @@ export class MemoryDb {
   /**
    * Store an embedding and its associated text/metadata.
    *
-   * Note: The async signature provides interface compatibility with potentially
-   * async backends (e.g., remote vector DBs). The current implementation uses
-   * synchronous bun:sqlite operations internally.
-   *
    * @param text - The text content to store
-   * @param embedding - The embedding vector (384 dimensions)
+   * @param embedding - The embedding vector
    * @param metadata - Optional metadata to associate with the entry
    * @returns The generated entry ID
    */
@@ -133,11 +163,11 @@ export class MemoryDb {
   ): Promise<string> {
     const id = randomUUID();
     const createdAt = new Date().toISOString();
-    MemoryDb.assertEmbeddingDimension(embedding);
+    this.assertEmbeddingDimension(embedding);
 
     // bun:sqlite transaction wrapper ensures atomicity synchronously
     const insertTransaction = this.db.transaction(() => {
-      this.db.run('INSERT INTO vec_memory(id, embedding) VALUES (?, ?)', [
+      this.db.run(`INSERT INTO ${this.tableName}(id, embedding) VALUES (?, ?)`, [
         id,
         new Float32Array(embedding),
       ]);
@@ -155,23 +185,19 @@ export class MemoryDb {
   /**
    * Search for similar embeddings using vector similarity.
    *
-   * Note: The async signature provides interface compatibility with potentially
-   * async backends (e.g., remote vector DBs). The current implementation uses
-   * synchronous bun:sqlite operations internally.
-   *
    * @param embedding - The query embedding vector
    * @param limit - Maximum number of results to return (default: 5)
    * @returns Array of matching entries with distance scores
    */
   async search(embedding: number[], limit = 5): Promise<MemoryEntry[]> {
-    MemoryDb.assertEmbeddingDimension(embedding);
+    this.assertEmbeddingDimension(embedding);
     const query = `
       SELECT 
         v.id, 
         v.distance,
         m.text,
         m.metadata
-      FROM vec_memory v
+      FROM ${this.tableName} v
       JOIN memory_metadata m ON v.id = m.id
       WHERE embedding MATCH ? AND k = ?
       ORDER BY distance

package/src/db/sqlite-setup.ts

@@ -2,7 +2,19 @@ import { Database } from 'bun:sqlite';
 import { existsSync } from 'node:fs';
 import { ConsoleLogger, type Logger } from '../utils/logger.ts';
 
-export function setupSqlite(logger: Logger = new ConsoleLogger()) {
+let sqliteSetupComplete = false;
+
+/**
+ * Setup SQLite with a custom library on macOS to support extensions.
+ * This is idempotent - calling it multiple times is safe.
+ */
+export function setupSqlite(logger: Logger = new ConsoleLogger()): void {
+  // Only run setup once
+  if (sqliteSetupComplete) {
+    return;
+  }
+  sqliteSetupComplete = true;
+
   // macOS typically comes with a system SQLite that doesn't support extensions
   // We need to try to load a custom one (e.g. from Homebrew) if on macOS
   if (process.platform === 'darwin') {
@@ -44,5 +56,16 @@ export function setupSqlite(logger: Logger = new ConsoleLogger()) {
   }
 }
 
-// Run setup immediately when imported
-setupSqlite();
+/**
+ * Reset SQLite setup state (mainly for testing).
+ */
+export function resetSqliteSetup(): void {
+  sqliteSetupComplete = false;
+}
+
+/**
+ * Check if SQLite setup has been completed.
+ */
+export function isSqliteSetupComplete(): boolean {
+  return sqliteSetupComplete;
+}

package/src/db/workflow-db.ts

@@ -2,7 +2,6 @@ import { Database, type Statement } from 'bun:sqlite';
 import { randomUUID } from 'node:crypto';
 import { existsSync, mkdirSync } from 'node:fs';
 import { dirname } from 'node:path';
-import './sqlite-setup.ts';
 import {
   StepStatus as StepStatusConst,
   type StepStatusType,
@@ -11,6 +10,7 @@ import {
 } from '../types/status';
 import { DB, LIMITS } from '../utils/constants';
 import { PathResolver } from '../utils/paths';
+import { setupSqlite } from './sqlite-setup.ts';
 
 export type RunStatus = WorkflowStatusType | 'pending';
 export type StepStatus = StepStatusType;
@@ -162,6 +162,9 @@ export class WorkflowDb {
   private isClosed = false;
 
   constructor(public readonly dbPath = PathResolver.resolveDbPath()) {
+    // Ensure SQLite is set up with custom library on macOS (idempotent)
+    setupSqlite();
+
     const dir = dirname(dbPath);
     if (!existsSync(dir)) {
       mkdirSync(dir, { recursive: true });
@@ -197,7 +200,11 @@ export class WorkflowDb {
       ORDER BY started_at DESC
       LIMIT ?
     `);
-    this.pruneRunsStmt = this.db.prepare('DELETE FROM workflow_runs WHERE started_at < ?');
+    this.pruneRunsStmt = this.db.prepare(`
+      DELETE FROM workflow_runs 
+      WHERE started_at < ? 
+      AND status IN ('success', 'failed', 'canceled')
+    `);
     this.createStepStmt = this.db.prepare(`
       INSERT INTO step_executions (id, run_id, step_id, iteration_index, status, retry_count)
       VALUES (?, ?, ?, ?, ?, ?)
@@ -448,7 +455,7 @@ export class WorkflowDb {
    * Uses exponential backoff with jitter to reduce contention.
    */
   private async withRetry<T>(operation: () => T, maxRetries = LIMITS.MAX_DB_RETRIES): Promise<T> {
-    let lastError: any;
+    let lastError: unknown;
 
     for (let attempt = 0; attempt < maxRetries; attempt++) {
       try {
@@ -477,13 +484,13 @@ export class WorkflowDb {
 
         // Wrap non-busy errors in DatabaseError
         const msg = error instanceof Error ? error.message : String(error);
-        const code = (error as any)?.code;
+        const code = (error as { code?: string | number })?.code;
         throw new DatabaseError(msg, code, false);
       }
     }
 
     const msg = lastError instanceof Error ? lastError.message : String(lastError);
-    const code = (lastError as any)?.code;
+    const code = (lastError as { code?: string | number })?.code;
     throw new DatabaseError(
       `SQLite operation failed after ${maxRetries} retries: ${msg}`,
       code,

package/src/parser/config-schema.ts

@@ -3,43 +3,41 @@ import { z } from 'zod';
 export const ConfigSchema = z.object({
   default_provider: z.string().default('openai'),
   default_model: z.string().optional(),
+  embedding_model: z.string().optional(),
+  embedding_dimension: z.number().int().positive().default(384),
   providers: z
     .record(
       z.object({
         type: z
-          .enum([
-            'openai',
-            'anthropic',
-            'anthropic-claude',
-            'copilot',
-            'openai-chatgpt',
-            'google-gemini',
-          ])
+          .enum(['openai', 'anthropic'])
+          .or(z.string()) // Allow custom types for BYOP
           .default('openai'),
         base_url: z.string().optional(),
         api_key_env: z.string().optional(),
         default_model: z.string().optional(),
         project_id: z.string().optional(),
+        embedding_dimension: z.number().int().positive().optional(),
+        // BYOP fields
+        package: z.string().optional(),
+        factory: z.string().optional(),
+        script: z.string().optional(),
       })
     )
     .default({
       openai: {
         type: 'openai',
+        package: '@ai-sdk/openai',
         base_url: 'https://api.openai.com/v1',
         api_key_env: 'OPENAI_API_KEY',
         default_model: 'gpt-4o',
       },
       anthropic: {
         type: 'anthropic',
+        package: '@ai-sdk/anthropic',
         base_url: 'https://api.anthropic.com/v1',
         api_key_env: 'ANTHROPIC_API_KEY',
         default_model: 'claude-3-5-sonnet-20240620',
       },
-      copilot: {
-        type: 'copilot',
-        base_url: 'https://api.githubcopilot.com',
-        default_model: 'gpt-4o',
-      },
     }),
   model_mappings: z.record(z.string()).default({
     'claude-*': 'anthropic',

package/src/parser/schema.ts

@@ -444,8 +444,10 @@ const DynamicStepSchema = BaseStepSchema.extend({
   allowInsecure: z.boolean().optional(), // Allow generated steps to use insecure commands (e.g. shell redirects)
 });
 
-// ===== Discriminated Union for Steps =====
-
+// Note: `as any` casts are required here because of circular type references:
+// BaseStepSchema.compensate → StepSchema → all step schemas → BaseStepSchema
+// TypeScript cannot infer types through this cycle, so we use z.ZodType<any>
+// and cast each schema. This is a known Zod limitation with recursive schemas.
 export const StepSchema: z.ZodType<any> = z.lazy(() =>
   z.discriminatedUnion('type', [
     ShellStepSchema as any,

package/src/runner/__test__/llm-mock-setup.ts

@@ -0,0 +1,173 @@
+/**
+ * Shared test mock setup for LLM adapter
+ *
+ * This file provides a unified mock model and setup utilities for tests
+ * that need to mock the LLM adapter without affecting other test files.
+ *
+ * Usage:
+ * 1. Import this at the top of your test file BEFORE any SUT imports
+ * 2. Call setupLlmAdapterMocks() before your tests
+ * 3. Use setCurrentChatFn() to control mock responses
+ */
+import { mock } from 'bun:test';
+
+// Mock response type
+export interface MockLLMResponse {
+  message: {
+    role: string;
+    content?: string | null;
+    tool_calls?: Array<{
+      id: string;
+      type: 'function';
+      function: { name: string; arguments: string };
+    }>;
+  };
+  usage?: { prompt_tokens: number; completion_tokens: number; total_tokens: number };
+}
+
+// Shared mock getModel function
+export const mockGetModel = mock();
+export const mockGetEmbeddingModel = mock();
+
+// Current chat function - set this in your test to control responses
+let _currentChatFn: (messages: any[], options?: any) => Promise<MockLLMResponse> = async () => ({
+  message: { role: 'assistant', content: 'Default mock response' },
+});
+
+export function setCurrentChatFn(fn: typeof _currentChatFn) {
+  _currentChatFn = fn;
+}
+
+export function getCurrentChatFn() {
+  return _currentChatFn;
+}
+
+/**
+ * Creates a unified mock model that simulates AI SDK LanguageModel behavior.
+ * This is used as the return value for mockGetModel.
+ */
+export function createUnifiedMockModel() {
+  return {
+    specificationVersion: 'v2',
+    provider: 'mock',
+    modelId: 'mock-model',
+    doStream: async (options: any) => {
+      // Convert AI SDK prompt format to our test format
+      const mapMessages = (prompt: any[]) =>
+        prompt.flatMap((m: any) => {
+          let content = m.content;
+          if (Array.isArray(m.content)) {
+            const toolResults = m.content.filter((p: any) => p.type === 'tool-result');
+            if (toolResults.length > 0) {
+              return toolResults.map((tr: any) => ({
+                role: 'tool',
+                tool_call_id: tr.toolCallId,
+                content: JSON.stringify(tr.result),
+              }));
+            }
+            const textParts = m.content
+              .filter((p: any) => p.type === 'text')
+              .map((p: any) => p.text)
+              .join('');
+            if (textParts) content = textParts;
+          }
+          return [
+            {
+              role: m.role,
+              content: typeof content === 'string' ? content : JSON.stringify(content),
+            },
+          ];
+        });
+
+      const messages = mapMessages(options.prompt || options.input);
+      const tools = (options.tools || options.mode?.tools)?.map((t: any) => ({
+        type: 'function',
+        function: {
+          name: t.name,
+          description: t.description,
+          parameters: t.parameters || t.inputSchema,
+        },
+      }));
+
+      const response = await _currentChatFn(messages, { tools });
+
+      const stream = new ReadableStream({
+        async start(controller) {
+          if (response.message.content) {
+            controller.enqueue({
+              type: 'text-delta',
+              delta: response.message.content,
+              text: response.message.content,
+            });
+          }
+
+          const toolCalls = response.message.tool_calls?.map((tc: any) => ({
+            type: 'tool-call',
+            toolCallId: tc.id,
+            toolName: tc.function.name,
+            args:
+              typeof tc.function.arguments === 'string'
+                ? JSON.parse(tc.function.arguments)
+                : tc.function.arguments,
+            id: tc.id,
+            name: tc.function.name,
+            input:
+              typeof tc.function.arguments === 'string'
+                ? tc.function.arguments
+                : JSON.stringify(tc.function.arguments),
+          }));
+
+          if (toolCalls?.length) {
+            for (const tc of toolCalls) {
+              controller.enqueue(tc);
+            }
+          }
+
+          controller.enqueue({
+            type: 'finish',
+            finishReason: toolCalls?.length ? 'tool-calls' : 'stop',
+            usage: { promptTokens: 10, completionTokens: 5 },
+          });
+
+          controller.close();
+        },
+      });
+
+      return { stream, rawResponse: { headers: {} } };
+    },
+  };
+}
+
+/**
+ * Sets up the LLM adapter module mocks.
+ * Call this at the TOP of your test file, before any imports of the SUT.
+ */
+export function setupLlmAdapterMocks() {
+  mock.module('../llm-adapter', () => ({
+    getModel: mockGetModel,
+    getEmbeddingModel: mockGetEmbeddingModel,
+    DynamicProviderRegistry: { getProvider: mock() },
+  }));
+
+  // Also mock with relative paths that might be used
+  mock.module('./llm-adapter', () => ({
+    getModel: mockGetModel,
+    getEmbeddingModel: mockGetEmbeddingModel,
+    DynamicProviderRegistry: { getProvider: mock() },
+  }));
+
+  // Reset mocks to use the unified model
+  mockGetModel.mockReset();
+  mockGetModel.mockResolvedValue(createUnifiedMockModel());
+}
+
+/**
+ * Resets all mocks to default state. Call in afterEach if needed.
+ */
+export function resetLlmMocks() {
+  mockGetModel.mockReset();
+  mockGetModel.mockResolvedValue(createUnifiedMockModel());
+  _currentChatFn = async () => ({
+    message: { role: 'assistant', content: 'Default mock response' },
+  });
+}