@copilotkitnext/sqlite-runner 0.0.22-alpha.0 → 0.0.22-alpha.2

package/src/sqlite-runner.ts

@@ -1,39 +1,77 @@
-import { AgentRunnerBase } from "@copilotkitnext/runtime";
-import { BaseEvent, RunAgentInput } from "@ag-ui/client";
-import Database from "better-sqlite3";
+import {
+  AgentRunner,
+  finalizeRunEvents,
+  type AgentRunnerConnectRequest,
+  type AgentRunnerIsRunningRequest,
+  type AgentRunnerRunRequest,
+  type AgentRunnerStopRequest,
+} from "@copilotkitnext/runtime";
 import { Observable, ReplaySubject } from "rxjs";
+import {
+  AbstractAgent,
+  BaseEvent,
+  RunAgentInput,
+  EventType,
+  RunStartedEvent,
+  compactEvents,
+} from "@ag-ui/client";
+import Database from "better-sqlite3";
 
 const SCHEMA_VERSION = 1;
 
+interface AgentRunRecord {
+  id: number;
+  thread_id: string;
+  run_id: string;
+  parent_run_id: string | null;
+  events: BaseEvent[];
+  input: RunAgentInput;
+  created_at: number;
+  version: number;
+}
+
 export interface SqliteAgentRunnerOptions {
   dbPath?: string;
 }
-export class SqliteAgentRunner extends AgentRunnerBase {
+
+interface ActiveConnectionContext {
+  subject: ReplaySubject<BaseEvent>;
+  agent?: AbstractAgent;
+  runSubject?: ReplaySubject<BaseEvent>;
+  currentEvents?: BaseEvent[];
+  stopRequested?: boolean;
+}
+
+// Active connections for streaming events and stop support
+const ACTIVE_CONNECTIONS = new Map<string, ActiveConnectionContext>();
+
+export class SqliteAgentRunner extends AgentRunner {
   private db: any;
-  private channels = new Map<string, ReplaySubject<BaseEvent>>();
 
   constructor(options: SqliteAgentRunnerOptions = {}) {
     super();
     const dbPath = options.dbPath ?? ":memory:";
+    
     if (!Database) {
       throw new Error(
         'better-sqlite3 is required for SqliteAgentRunner but was not found.\n' +
-          'Please install it in your project:\n' +
-          '  npm install better-sqlite3\n' +
-          '  or\n' +
-          '  pnpm add better-sqlite3\n' +
-          '  or\n' +
-          '  yarn add better-sqlite3\n\n' +
-          "If you don't need persistence, use InMemoryAgentRunner instead.",
+        'Please install it in your project:\n' +
+        '  npm install better-sqlite3\n' +
+        '  or\n' +
+        '  pnpm add better-sqlite3\n' +
+        '  or\n' +
+        '  yarn add better-sqlite3\n\n' +
+        'If you don\'t need persistence, use InMemoryAgentRunner instead.'
       );
     }
-    const db = new Database(dbPath);
-    SqliteAgentRunner.initializeSchema(db);
-    this.db = db;
+    
+    this.db = new Database(dbPath);
+    this.initializeSchema();
   }
 
-  private static initializeSchema(db: any): void {
-    db.exec(`
+  private initializeSchema(): void {
+    // Create the agent_runs table
+    this.db.exec(`
       CREATE TABLE IF NOT EXISTS agent_runs (
         id INTEGER PRIMARY KEY AUTOINCREMENT,
         thread_id TEXT NOT NULL,
@@ -46,7 +84,8 @@ export class SqliteAgentRunner extends AgentRunnerBase {
       )
     `);
 
-    db.exec(`
+    // Create run_state table to track active runs
+    this.db.exec(`
       CREATE TABLE IF NOT EXISTS run_state (
         thread_id TEXT PRIMARY KEY,
         is_running INTEGER DEFAULT 0,
@@ -55,179 +94,404 @@ export class SqliteAgentRunner extends AgentRunnerBase {
       )
     `);
 
-    db.exec(`
+    // Create indexes for efficient queries
+    this.db.exec(`
       CREATE INDEX IF NOT EXISTS idx_thread_id ON agent_runs(thread_id);
       CREATE INDEX IF NOT EXISTS idx_parent_run_id ON agent_runs(parent_run_id);
     `);
 
-    db.exec(`
+    // Create schema version table
+    this.db.exec(`
       CREATE TABLE IF NOT EXISTS schema_version (
         version INTEGER PRIMARY KEY,
         applied_at INTEGER NOT NULL
       )
     `);
 
-    const currentVersion = db
+    // Check and set schema version
+    const currentVersion = this.db
       .prepare("SELECT version FROM schema_version ORDER BY version DESC LIMIT 1")
       .get() as { version: number } | undefined;
+
     if (!currentVersion || currentVersion.version < SCHEMA_VERSION) {
-      db
+      this.db
         .prepare("INSERT OR REPLACE INTO schema_version (version, applied_at) VALUES (?, ?)")
         .run(SCHEMA_VERSION, Date.now());
     }
   }
 
-  // Hooks implementation using SQLite
-  protected async acquireRun(threadId: string, runId: string): Promise<boolean> {
-    const row = this.db.prepare("SELECT is_running FROM run_state WHERE thread_id = ?").get(threadId) as
-      | { is_running: number }
-      | undefined;
-    if (row?.is_running === 1) return false;
-    this.db
-      .prepare(
-        "INSERT OR REPLACE INTO run_state (thread_id, is_running, current_run_id, updated_at) VALUES (?, ?, ?, ?)",
-      )
-      .run(threadId, 1, runId, Date.now());
-    return true;
-  }
-
-  protected async releaseRun(threadId: string): Promise<void> {
-    this.db
-      .prepare("INSERT OR REPLACE INTO run_state (thread_id, is_running, current_run_id, updated_at) VALUES (?, 0, NULL, ?)")
-      .run(threadId, Date.now());
-  }
-
-  protected async isRunningState(threadId: string): Promise<boolean> {
-    const row = this.db.prepare("SELECT is_running FROM run_state WHERE thread_id = ?").get(threadId) as
-      | { is_running: number }
-      | undefined;
-    return row?.is_running === 1;
-  }
-
-  protected async listRuns(threadId: string): Promise<Array<{ runId: string; events: BaseEvent[]; createdAt: number }>> {
-    const rows = this.db
-      .prepare(
-        "SELECT run_id, events, created_at FROM agent_runs WHERE thread_id = ? ORDER BY created_at ASC",
-      )
-      .all(threadId) as Array<{ run_id: string; events: string; created_at: number }>;
-    return rows.map((r) => ({ runId: r.run_id, events: JSON.parse(r.events), createdAt: r.created_at }));
-  }
-
-  protected async saveRun(
+  private storeRun(
     threadId: string,
     runId: string,
     events: BaseEvent[],
     input: RunAgentInput,
-    parentRunId: string | null,
-  ): Promise<void> {
-    const stmt = this.db.prepare(
-      "INSERT INTO agent_runs (thread_id, run_id, parent_run_id, events, input, created_at, version) VALUES (?, ?, ?, ?, ?, ?, ?)",
-    );
+    parentRunId?: string | null
+  ): void {
+    // Compact ONLY the events from this run
+    const compactedEvents = compactEvents(events);
+    
+    const stmt = this.db.prepare(`
+      INSERT INTO agent_runs (thread_id, run_id, parent_run_id, events, input, created_at, version)
+      VALUES (?, ?, ?, ?, ?, ?, ?)
+    `);
+
     stmt.run(
       threadId,
       runId,
       parentRunId ?? null,
-      JSON.stringify(events),
+      JSON.stringify(compactedEvents), // Store only this run's compacted events
       JSON.stringify(input),
       Date.now(),
-      SCHEMA_VERSION,
+      SCHEMA_VERSION
     );
   }
 
-  protected async pageThreads(params: { scope?: any; limit?: number; offset?: number }): Promise<{ threadIds: string[]; total: number }> {
-    const limit = params.limit ?? 20;
-    const offset = params.offset ?? 0;
-    const totalRow = this.db
-      .prepare("SELECT COUNT(DISTINCT thread_id) AS total FROM agent_runs")
-      .get() as { total: number };
-    const rows = this.db
-      .prepare(
-        "SELECT thread_id, MAX(created_at) AS last FROM agent_runs GROUP BY thread_id ORDER BY last DESC LIMIT ? OFFSET ?",
+  private getHistoricRuns(threadId: string): AgentRunRecord[] {
+    const stmt = this.db.prepare(`
+      WITH RECURSIVE run_chain AS (
+        -- Base case: find the root runs (those without parent)
+        SELECT * FROM agent_runs 
+        WHERE thread_id = ? AND parent_run_id IS NULL
+        
+        UNION ALL
+        
+        -- Recursive case: find children of current level
+        SELECT ar.* FROM agent_runs ar
+        INNER JOIN run_chain rc ON ar.parent_run_id = rc.run_id
+        WHERE ar.thread_id = ?
       )
-      .all(limit, offset) as Array<{ thread_id: string; last: number }>;
-    return { threadIds: rows.map((r) => r.thread_id), total: totalRow.total };
-  }
+      SELECT * FROM run_chain
+      ORDER BY created_at ASC
+    `);
 
-  protected async deleteThreadStorage(threadId: string): Promise<void> {
-    this.db.prepare("DELETE FROM agent_runs WHERE thread_id = ?").run(threadId);
-    this.db.prepare("DELETE FROM run_state WHERE thread_id = ?").run(threadId);
+    const rows = stmt.all(threadId, threadId) as any[];
+    
+    return rows.map(row => ({
+      id: row.id,
+      thread_id: row.thread_id,
+      run_id: row.run_id,
+      parent_run_id: row.parent_run_id,
+      events: JSON.parse(row.events),
+      input: JSON.parse(row.input),
+      created_at: row.created_at,
+      version: row.version
+    }));
   }
 
-  private ensureChannel(threadId: string): ReplaySubject<BaseEvent> {
-    let s = this.channels.get(threadId);
-    if (!s || s.closed) {
-      s = new ReplaySubject<BaseEvent>(Infinity);
-      this.channels.set(threadId, s);
-    }
-    return s;
-  }
+  private getLatestRunId(threadId: string): string | null {
+    const stmt = this.db.prepare(`
+      SELECT run_id FROM agent_runs 
+      WHERE thread_id = ? 
+      ORDER BY created_at DESC 
+      LIMIT 1
+    `);
 
-  protected publishLive(threadId: string, event: BaseEvent): void {
-    // Reset the live channel at the start of each run to avoid replaying
-    // previous runs' events to new concurrent connections.
-    if (event.type === 'RUN_STARTED') {
-      const s = new ReplaySubject<BaseEvent>(Infinity);
-      this.channels.set(threadId, s);
-    }
-    this.ensureChannel(threadId).next(event);
+    const result = stmt.get(threadId) as { run_id: string } | undefined;
+    return result?.run_id ?? null;
   }
 
-  protected completeLive(threadId: string): void {
-    const s = this.ensureChannel(threadId);
-    if (!s.closed) s.complete();
+  private setRunState(threadId: string, isRunning: boolean, runId?: string): void {
+    const stmt = this.db.prepare(`
+      INSERT OR REPLACE INTO run_state (thread_id, is_running, current_run_id, updated_at)
+      VALUES (?, ?, ?, ?)
+    `);
+    stmt.run(threadId, isRunning ? 1 : 0, runId ?? null, Date.now());
   }
 
-  protected subscribeLive(threadId: string): Observable<BaseEvent> {
-    return this.ensureChannel(threadId).asObservable();
+  private getRunState(threadId: string): { isRunning: boolean; currentRunId: string | null } {
+    const stmt = this.db.prepare(`
+      SELECT is_running, current_run_id FROM run_state WHERE thread_id = ?
+    `);
+    const result = stmt.get(threadId) as { is_running: number; current_run_id: string | null } | undefined;
+    
+    return {
+      isRunning: result?.is_running === 1,
+      currentRunId: result?.current_run_id ?? null
+    };
   }
 
-  protected async closeLive(threadId: string): Promise<void> {
-    const s = this.channels.get(threadId);
-    if (s && !s.closed) s.complete();
-    this.channels.delete(threadId);
-  }
+  run(request: AgentRunnerRunRequest): Observable<BaseEvent> {
+    // Check if thread is already running in database
+    const runState = this.getRunState(request.threadId);
+    if (runState.isRunning) {
+      throw new Error("Thread already running");
+    }
+
+    // Mark thread as running in database
+    this.setRunState(request.threadId, true, request.input.runId);
 
-  // Override connect to preserve original sqlite-runner ordering without re-compacting across runs
-  connect(request: { threadId: string }) {
-    const { threadId } = request;
-    const subject = new ReplaySubject<BaseEvent>(Infinity);
-    void (async () => {
-      // Load runs and emit stored events in order
-      const runs = await this.listRuns(threadId);
-      const emittedMessageIds = new Set<string>();
-      for (const r of runs) {
-        for (const e of r.events) {
-          subject.next(e);
-          const id = (e as any).messageId;
-          if (typeof id === "string") emittedMessageIds.add(id);
+    // Track seen message IDs and current run events in memory for this run
+    const seenMessageIds = new Set<string>();
+    const currentRunEvents: BaseEvent[] = [];
+    
+    // Get all previously seen message IDs from historic runs
+    const historicRuns = this.getHistoricRuns(request.threadId);
+    const historicMessageIds = new Set<string>();
+    for (const run of historicRuns) {
+      for (const event of run.events) {
+        if ('messageId' in event && typeof event.messageId === 'string') {
+          historicMessageIds.add(event.messageId);
+        }
+        if (event.type === EventType.RUN_STARTED) {
+          const runStarted = event as RunStartedEvent;
+          const messages = runStarted.input?.messages ?? [];
+          for (const message of messages) {
+            historicMessageIds.add(message.id);
+          }
         }
       }
+    }
+
+    // Get or create subject for this thread's connections
+    const nextSubject = new ReplaySubject<BaseEvent>(Infinity);
+    const prevConnection = ACTIVE_CONNECTIONS.get(request.threadId);
+    const prevSubject = prevConnection?.subject;
+    
+    // Create a subject for run() return value
+    const runSubject = new ReplaySubject<BaseEvent>(Infinity);
+
+    // Update the active connection for this thread
+    ACTIVE_CONNECTIONS.set(request.threadId, {
+      subject: nextSubject,
+      agent: request.agent,
+      runSubject,
+      currentEvents: currentRunEvents,
+      stopRequested: false,
+    });
 
-      const running = await this.isRunningState(threadId);
-      if (!running) {
-        subject.complete();
-        return;
+    // Helper function to run the agent and handle errors
+    const runAgent = async () => {
+      // Get parent run ID for chaining
+      const parentRunId = this.getLatestRunId(request.threadId);
+      
+      try {
+        await request.agent.runAgent(request.input, {
+          onEvent: ({ event }) => {
+            let processedEvent: BaseEvent = event;
+            if (event.type === EventType.RUN_STARTED) {
+              const runStartedEvent = event as RunStartedEvent;
+              if (!runStartedEvent.input) {
+                const sanitizedMessages = request.input.messages
+                  ? request.input.messages.filter(
+                      (message) => !historicMessageIds.has(message.id),
+                    )
+                  : undefined;
+                const updatedInput = {
+                  ...request.input,
+                  ...(sanitizedMessages !== undefined
+                    ? { messages: sanitizedMessages }
+                    : {}),
+                };
+                processedEvent = {
+                  ...runStartedEvent,
+                  input: updatedInput,
+                } as RunStartedEvent;
+              }
+            }
+
+            runSubject.next(processedEvent); // For run() return - only agent events
+            nextSubject.next(processedEvent); // For connect() / store - all events
+            currentRunEvents.push(processedEvent); // Accumulate for database storage
+          },
+          onNewMessage: ({ message }) => {
+            // Called for each new message
+            if (!seenMessageIds.has(message.id)) {
+              seenMessageIds.add(message.id);
+            }
+          },
+          onRunStartedEvent: () => {
+            // Mark input messages as seen without emitting duplicates
+            if (request.input.messages) {
+              for (const message of request.input.messages) {
+                if (!seenMessageIds.has(message.id)) {
+                  seenMessageIds.add(message.id);
+                }
+              }
+            }
+          },
+        });
+        
+        const connection = ACTIVE_CONNECTIONS.get(request.threadId);
+        const appendedEvents = finalizeRunEvents(currentRunEvents, {
+          stopRequested: connection?.stopRequested ?? false,
+        });
+        for (const event of appendedEvents) {
+          runSubject.next(event);
+          nextSubject.next(event);
+        }
+
+        // Store the run in database
+        this.storeRun(
+          request.threadId,
+          request.input.runId,
+          currentRunEvents,
+          request.input,
+          parentRunId
+        );
+        
+        // Mark run as complete in database
+        this.setRunState(request.threadId, false);
+
+        if (connection) {
+          connection.agent = undefined;
+          connection.runSubject = undefined;
+          connection.currentEvents = undefined;
+          connection.stopRequested = false;
+        }
+
+        // Complete the subjects
+        runSubject.complete();
+        nextSubject.complete();
+
+        ACTIVE_CONNECTIONS.delete(request.threadId);
+      } catch {
+        const connection = ACTIVE_CONNECTIONS.get(request.threadId);
+        const appendedEvents = finalizeRunEvents(currentRunEvents, {
+          stopRequested: connection?.stopRequested ?? false,
+        });
+        for (const event of appendedEvents) {
+          runSubject.next(event);
+          nextSubject.next(event);
+        }
+
+        // Store the run even if it failed (partial events)
+        if (currentRunEvents.length > 0) {
+          this.storeRun(
+            request.threadId,
+            request.input.runId,
+            currentRunEvents,
+            request.input,
+            parentRunId
+          );
+        }
+        
+        // Mark run as complete in database
+        this.setRunState(request.threadId, false);
+
+        if (connection) {
+          connection.agent = undefined;
+          connection.runSubject = undefined;
+          connection.currentEvents = undefined;
+          connection.stopRequested = false;
+        }
+
+        // Don't emit error to the subject, just complete it
+        // This allows subscribers to get events emitted before the error
+        runSubject.complete();
+        nextSubject.complete();
+
+        ACTIVE_CONNECTIONS.delete(request.threadId);
       }
+    };
+
+    // Bridge previous events if they exist
+    if (prevSubject) {
+      prevSubject.subscribe({
+        next: (e) => nextSubject.next(e),
+        error: (err) => nextSubject.error(err),
+        complete: () => {
+          // Don't complete nextSubject here - it needs to stay open for new events
+        },
+      });
+    }
+
+    // Start the agent execution immediately (not lazily)
+    runAgent();
+
+    // Return the run subject (only agent events, no injected messages)
+    return runSubject.asObservable();
+  }
+
+  connect(request: AgentRunnerConnectRequest): Observable<BaseEvent> {
+    const connectionSubject = new ReplaySubject<BaseEvent>(Infinity);
 
-      const live$ = this.subscribeLive(threadId);
-      let completed = false;
-      live$.subscribe({
-        next: (e) => {
-          const id = (e as any).messageId;
-          if (typeof id === "string" && emittedMessageIds.has(id)) return;
-          subject.next(e);
-          if (e.type === 'RUN_FINISHED' || e.type === 'RUN_ERROR') {
-            if (!completed) { completed = true; subject.complete(); }
+    // Load historic runs from database
+    const historicRuns = this.getHistoricRuns(request.threadId);
+    
+    // Collect all historic events from database
+    const allHistoricEvents: BaseEvent[] = [];
+    for (const run of historicRuns) {
+      allHistoricEvents.push(...run.events);
+    }
+    
+    // Compact all events together before emitting
+    const compactedEvents = compactEvents(allHistoricEvents);
+    
+    // Emit compacted events and track message IDs
+    const emittedMessageIds = new Set<string>();
+    for (const event of compactedEvents) {
+      connectionSubject.next(event);
+      if ('messageId' in event && typeof event.messageId === 'string') {
+        emittedMessageIds.add(event.messageId);
+      }
+    }
+    
+    // Bridge active run to connection if exists
+    const activeConnection = ACTIVE_CONNECTIONS.get(request.threadId);
+    const runState = this.getRunState(request.threadId);
+
+    if (activeConnection && (runState.isRunning || activeConnection.stopRequested)) {
+      activeConnection.subject.subscribe({
+        next: (event) => {
+          // Skip message events that we've already emitted from historic
+          if ('messageId' in event && typeof event.messageId === 'string' && emittedMessageIds.has(event.messageId)) {
+            return;
           }
+          connectionSubject.next(event);
         },
-        error: (err) => subject.error(err),
-        complete: () => { if (!completed) { completed = true; subject.complete(); } },
+        complete: () => connectionSubject.complete(),
+        error: (err) => connectionSubject.error(err)
       });
-    })();
-    return subject.asObservable();
+    } else {
+      // No active run, complete after historic events
+      connectionSubject.complete();
+    }
+    
+    return connectionSubject.asObservable();
   }
 
+  isRunning(request: AgentRunnerIsRunningRequest): Promise<boolean> {
+    const runState = this.getRunState(request.threadId);
+    return Promise.resolve(runState.isRunning);
+  }
+
+  stop(request: AgentRunnerStopRequest): Promise<boolean | undefined> {
+    const runState = this.getRunState(request.threadId);
+    if (!runState.isRunning) {
+      return Promise.resolve(false);
+    }
+
+    const connection = ACTIVE_CONNECTIONS.get(request.threadId);
+    const agent = connection?.agent;
+
+    if (!connection || !agent) {
+      return Promise.resolve(false);
+    }
+
+    if (connection.stopRequested) {
+      return Promise.resolve(false);
+    }
+
+    connection.stopRequested = true;
+    this.setRunState(request.threadId, false);
+
+    try {
+      agent.abortRun();
+      return Promise.resolve(true);
+    } catch (error) {
+      console.error("Failed to abort sqlite agent run", error);
+      connection.stopRequested = false;
+      this.setRunState(request.threadId, true);
+      return Promise.resolve(false);
+    }
+  }
+
+  /**
+   * Close the database connection (for cleanup)
+   */
   close(): void {
-    if (this.db) this.db.close();
+    if (this.db) {
+      this.db.close();
+    }
   }
 }

package/.turbo/turbo-test.log

@@ -1,32 +0,0 @@
-
-
-> @copilotkitnext/sqlite-runner@0.0.21 test /Users/mme/Projects/CopilotKit2/threads-final-review/packages/sqlite-runner
-> vitest run
-
-[?25l
- RUN  v3.2.4 /Users/mme/Projects/CopilotKit2/threads-final-review/packages/sqlite-runner
-
-[?2026h
- ❯ src/__tests__/sqlite-runner.test.ts [queued]
-
- Test Files 0 passed (2)
-      Tests 0 passed (0)
-   Start at 17:14:55
-   Duration 311ms
-[?2026l[?2026h
- ❯ src/__tests__/sqlite-runner.e2e.test.ts [queued]
- ❯ src/__tests__/sqlite-runner.test.ts 0/7
-
- Test Files 0 passed (2)
-      Tests 0 passed (7)
-   Start at 17:14:55
-   Duration 513ms
-[?2026l ✓ src/__tests__/sqlite-runner.e2e.test.ts (8 tests) 22ms
- ✓ src/__tests__/sqlite-runner.test.ts (7 tests) 88ms
-
- Test Files  2 passed (2)
-      Tests  15 passed (15)
-   Start at  17:14:55
-   Duration  677ms (transform 134ms, setup 0ms, collect 433ms, tests 110ms, environment 0ms, prepare 150ms)
-
-[?25h