chore: sync with upstream 2026-06-20

packages/control-plane/src/db/automation-store.ts

@@ -447,14 +447,17 @@ export class AutomationStore {
   }
 
   // --- Recovery sweep queries ---
+  // Backed by partial indexes (migration 0024); `status` must stay a literal, not
+  // a bound param, or the planner skips the index and full-scans automation_runs.
+  static readonly ORPHANED_STARTING_RUNS_SQL =
+    "SELECT * FROM automation_runs WHERE status = 'starting' AND created_at < ?";
+  static readonly TIMED_OUT_RUNNING_RUNS_SQL =
+    "SELECT * FROM automation_runs WHERE status = 'running' AND started_at IS NOT NULL AND started_at < ?";
 
   async getOrphanedStartingRuns(thresholdMs: number): Promise<AutomationRunRow[]> {
     const cutoff = Date.now() - thresholdMs;
     const result = await this.db
-      .prepare(
-        `SELECT * FROM automation_runs
-         WHERE status = 'starting' AND created_at < ?`
-      )
+      .prepare(AutomationStore.ORPHANED_STARTING_RUNS_SQL)
       .bind(cutoff)
       .all<AutomationRunRow>();
     return result.results || [];
@@ -463,10 +466,7 @@ export class AutomationStore {
   async getTimedOutRunningRuns(executionTimeoutMs: number): Promise<AutomationRunRow[]> {
     const cutoff = Date.now() - executionTimeoutMs;
     const result = await this.db
-      .prepare(
-        `SELECT * FROM automation_runs
-         WHERE status = 'running' AND started_at IS NOT NULL AND started_at < ?`
-      )
+      .prepare(AutomationStore.TIMED_OUT_RUNNING_RUNS_SQL)
       .bind(cutoff)
       .all<AutomationRunRow>();
     return result.results || [];

packages/control-plane/src/session/durable-object.ts

@@ -72,6 +72,7 @@ import { DOFetcherAdapter } from "../scheduler/do-fetcher-adapter";
 import { PresenceService } from "./presence-service";
 import { SessionMessageQueue } from "./message-queue";
 import { SessionSandboxEventProcessor } from "./sandbox-events";
+import { SessionEventStream } from "./event-stream";
 import { createSessionInternalRoutes } from "./http/routes";
 import { createMessagesHandler, type MessagesHandler } from "./http/handlers/messages.handler";
 import {
@@ -139,6 +140,7 @@ export class SessionDO extends DurableObject<Env> {
   private _messageQueue: SessionMessageQueue | null = null;
   // Message service (lazily initialized)
   private _messageService: MessageService | null = null;
+  private _eventStream: SessionEventStream | null = null;
   // Messages handler (lazily initialized)
   private _messagesHandler: MessagesHandler | null = null;
   // Child sessions handler (lazily initialized)
@@ -347,6 +349,14 @@ export class SessionDO extends DurableObject<Env> {
     return this._messageService;
   }
 
+  private get eventStream(): SessionEventStream {
+    if (!this._eventStream) {
+      this._eventStream = new SessionEventStream(this.repository);
+    }
+
+    return this._eventStream;
+  }
+
   private get messagesHandler(): MessagesHandler {
     if (!this._messagesHandler) {
       this._messagesHandler = createMessagesHandler({
@@ -1254,7 +1264,7 @@ export class SessionDO extends DurableObject<Env> {
     const sandbox = this.getSandbox();
     const state = await this.getSessionState(sandbox);
     const artifacts = this.messageService.listArtifacts();
-    const replay = this.getReplayData();
+    const replay = this.eventStream.getReplay();
 
     this.safeSend(ws, {
       type: "subscribed",
@@ -1281,33 +1291,6 @@ export class SessionDO extends DurableObject<Env> {
     this.presenceService.broadcastPresence();
   }
 
-  /**
-   * Collect historical events for replay.
-   * Returns parsed events and pagination metadata for inclusion in the subscribed message.
-   */
-  private getReplayData(): {
-    events: SandboxEvent[];
-    hasMore: boolean;
-    cursor: { timestamp: number; id: string } | null;
-  } {
-    const REPLAY_LIMIT = 500;
-    const rows = this.repository.getEventsForReplay(REPLAY_LIMIT);
-    const hasMore = rows.length >= REPLAY_LIMIT;
-
-    const events: SandboxEvent[] = [];
-    for (const row of rows) {
-      try {
-        events.push(JSON.parse(row.data));
-      } catch {
-        // Skip malformed events
-      }
-    }
-
-    const cursor = rows.length > 0 ? { timestamp: rows[0].created_at, id: rows[0].id } : null;
-
-    return { events, hasMore, cursor };
-  }
-
   /**
    * Get client info for a WebSocket, reconstructing from storage if needed after hibernation.
    */
@@ -1400,30 +1383,16 @@ export class SessionDO extends DurableObject<Env> {
     }
     client.lastFetchHistoryAt = now;
 
-    const rawLimit = typeof data.limit === "number" ? data.limit : 200;
-    const limit = Math.max(1, Math.min(rawLimit, 500));
-    const page = this.repository.getEventTimelinePage({
-      cursor: { kind: "timeline", createdAt: data.cursor.timestamp, id: data.cursor.id },
-      excludeTypes: ["heartbeat"],
-      limit,
+    const page = this.eventStream.getHistoryPage({
+      cursor: data.cursor,
+      limit: data.limit,
     });
 
-    const items: SandboxEvent[] = [];
-    for (const event of page.events) {
-      try {
-        items.push(JSON.parse(event.data));
-      } catch {
-        // Skip malformed events
-      }
-    }
-
     this.safeSend(ws, {
       type: "history_page",
-      items,
+      items: page.items,
       hasMore: page.hasMore,
-      cursor: page.nextCursor
-        ? { timestamp: page.nextCursor.createdAt, id: page.nextCursor.id }
-        : null,
+      cursor: page.cursor,
     } as ServerMessage);
   }
 

packages/control-plane/src/session/event-stream.test.ts

@@ -0,0 +1,205 @@
+import { describe, expect, it, vi } from "vitest";
+import { SessionEventStream, type SessionEventStreamRepository } from "./event-stream";
+import type { EventRow } from "./types";
+
+function createStream() {
+  const repository = {
+    getEventsForReplay: vi.fn(),
+    getEventTimelinePage: vi.fn(),
+    listEventPage: vi.fn(),
+  } as unknown as SessionEventStreamRepository;
+
+  return {
+    stream: new SessionEventStream(repository),
+    repository,
+  };
+}
+
+function eventRow(
+  id: string,
+  type: EventRow["type"],
+  data: Record<string, unknown> | string,
+  createdAt: number
+): EventRow {
+  return {
+    id,
+    type,
+    data: typeof data === "string" ? data : JSON.stringify(data),
+    message_id: null,
+    created_at: createdAt,
+  };
+}
+
+describe("SessionEventStream", () => {
+  describe("getReplay", () => {
+    it("loads replay rows with the default replay limit", () => {
+      const { stream, repository } = createStream();
+      vi.mocked(repository.getEventsForReplay).mockReturnValue([]);
+
+      stream.getReplay();
+
+      expect(repository.getEventsForReplay).toHaveBeenCalledWith(500);
+    });
+
+    it("returns parsed replay events and the oldest cursor from the loaded window", () => {
+      const { stream, repository } = createStream();
+      vi.mocked(repository.getEventsForReplay).mockReturnValue([
+        eventRow("e1", "tool_call", { type: "tool_call", tool: "read_file" }, 1000),
+        eventRow("e2", "tool_result", { type: "tool_result", result: "ok" }, 2000),
+      ]);
+
+      const replay = stream.getReplay();
+
+      expect(replay).toEqual({
+        events: [
+          { type: "tool_call", tool: "read_file" },
+          { type: "tool_result", result: "ok" },
+        ],
+        hasMore: false,
+        cursor: { timestamp: 1000, id: "e1" },
+      });
+    });
+
+    it("marks replay as having more when the loaded window reaches the limit", () => {
+      const { stream, repository } = createStream();
+      vi.mocked(repository.getEventsForReplay).mockReturnValue([
+        eventRow("e1", "token", { type: "token", content: "a" }, 1000),
+        eventRow("e2", "token", { type: "token", content: "b" }, 2000),
+      ]);
+
+      const replay = stream.getReplay(2);
+
+      expect(replay.hasMore).toBe(true);
+    });
+
+    it("skips malformed replay event JSON", () => {
+      const { stream, repository } = createStream();
+      vi.mocked(repository.getEventsForReplay).mockReturnValue([
+        eventRow("bad", "tool_call", "{bad", 1000),
+        eventRow("good", "tool_result", { type: "tool_result", result: "ok" }, 2000),
+      ]);
+
+      const replay = stream.getReplay();
+
+      expect(replay.events).toEqual([{ type: "tool_result", result: "ok" }]);
+      expect(replay.cursor).toEqual({ timestamp: 1000, id: "bad" });
+    });
+  });
+
+  describe("getHistoryPage", () => {
+    it("loads history after a client cursor while excluding heartbeats", () => {
+      const { stream, repository } = createStream();
+      vi.mocked(repository.getEventTimelinePage).mockReturnValue({
+        events: [eventRow("e1", "tool_call", { type: "tool_call", tool: "write_file" }, 1000)],
+        hasMore: false,
+        nextCursor: { kind: "timeline", createdAt: 1000, id: "e1" },
+      });
+
+      const page = stream.getHistoryPage({
+        cursor: { timestamp: 2000, id: "cursor-id" },
+        limit: 100,
+      });
+
+      expect(repository.getEventTimelinePage).toHaveBeenCalledWith({
+        cursor: { kind: "timeline", createdAt: 2000, id: "cursor-id" },
+        excludeTypes: ["heartbeat"],
+        limit: 100,
+      });
+      expect(page).toEqual({
+        items: [{ type: "tool_call", tool: "write_file" }],
+        hasMore: false,
+        cursor: { timestamp: 1000, id: "e1" },
+      });
+    });
+
+    it("clamps history limits to the supported range", () => {
+      const { stream, repository } = createStream();
+      vi.mocked(repository.getEventTimelinePage).mockReturnValue({
+        events: [],
+        hasMore: false,
+        nextCursor: null,
+      });
+
+      stream.getHistoryPage({ cursor: { timestamp: 2000, id: "cursor-id" }, limit: 999 });
+      stream.getHistoryPage({ cursor: { timestamp: 2000, id: "cursor-id" }, limit: 0 });
+      stream.getHistoryPage({ cursor: { timestamp: 2000, id: "cursor-id" } });
+
+      expect(repository.getEventTimelinePage).toHaveBeenNthCalledWith(1, {
+        cursor: { kind: "timeline", createdAt: 2000, id: "cursor-id" },
+        excludeTypes: ["heartbeat"],
+        limit: 500,
+      });
+      expect(repository.getEventTimelinePage).toHaveBeenNthCalledWith(2, {
+        cursor: { kind: "timeline", createdAt: 2000, id: "cursor-id" },
+        excludeTypes: ["heartbeat"],
+        limit: 1,
+      });
+      expect(repository.getEventTimelinePage).toHaveBeenNthCalledWith(3, {
+        cursor: { kind: "timeline", createdAt: 2000, id: "cursor-id" },
+        excludeTypes: ["heartbeat"],
+        limit: 200,
+      });
+    });
+
+    it("skips malformed history event JSON", () => {
+      const { stream, repository } = createStream();
+      vi.mocked(repository.getEventTimelinePage).mockReturnValue({
+        events: [
+          eventRow("bad", "tool_call", "{bad", 1000),
+          eventRow("good", "tool_result", { type: "tool_result", result: "ok" }, 2000),
+        ],
+        hasMore: true,
+        nextCursor: { kind: "timeline", createdAt: 1000, id: "bad" },
+      });
+
+      const page = stream.getHistoryPage({
+        cursor: { timestamp: 3000, id: "cursor-id" },
+        limit: 10,
+      });
+
+      expect(page).toEqual({
+        items: [{ type: "tool_result", result: "ok" }],
+        hasMore: true,
+        cursor: { timestamp: 1000, id: "bad" },
+      });
+    });
+  });
+
+  describe("listEvents", () => {
+    it("projects event rows to the shared HTTP response shape", () => {
+      const { stream, repository } = createStream();
+      vi.mocked(repository.listEventPage).mockReturnValue({
+        events: [eventRow("e1", "token", { type: "token", content: "hello" }, 1000)],
+        hasMore: true,
+        nextCursor: { kind: "timeline", createdAt: 1000, id: "e1" },
+      });
+
+      const page = stream.listEvents({
+        cursor: null,
+        limit: 10,
+        type: "token",
+        messageId: "m1",
+      });
+
+      expect(repository.listEventPage).toHaveBeenCalledWith({
+        cursor: null,
+        limit: 10,
+        type: "token",
+        messageId: "m1",
+      });
+      expect(page).toEqual({
+        events: [
+          {
+            id: "e1",
+            type: "token",
+            data: { type: "token", content: "hello" },
+            messageId: null,
+            createdAt: 1000,
+          },
+        ],
+        cursor: "1000:e1",
+        hasMore: true,
+      });
+    });
+  });
+});