perf(sessions): batch runtime event sequence allocation

apps/api/src/modules/sessions/infrastructure/session-runtime-event-store.repository.ts

@@ -1,6 +1,6 @@
 import { sessionEventsTable, sessionRunsTable, sessionsTable } from "@mosoo/db";
 import { createPlatformId } from "@mosoo/id";
-import type { RuntimeEventId, SessionId, SessionRunId } from "@mosoo/id";
+import type { AgentId, RuntimeEventId, SessionId, SessionRunId } from "@mosoo/id";
 import { and, eq, inArray, isNull, sql } from "drizzle-orm";
 
 import { getAppDatabase } from "../../../platform/db/drizzle";
@@ -37,6 +37,8 @@ export type {
 const MAX_SESSION_RUNTIME_EVENT_INSERT_ATTEMPTS = 5;
 // D1 accepts at most 100 bound parameters; each session_event row binds 21.
 const MAX_SESSION_EVENT_ROWS_PER_INSERT = 4;
+// The terminal allocation query binds four statuses, leaving room for 96 session ids.
+const MAX_SESSION_EVENT_ALLOCATIONS_PER_UPDATE = 96;
 const WRITABLE_SESSION_STATUSES = ["IDLE", "RUNNING", "RESCHEDULING"] as const;
 const TERMINAL_LIFECYCLE_WRITABLE_SESSION_STATUSES = [
   ...WRITABLE_SESSION_STATUSES,
@@ -318,24 +320,39 @@ async function allocateOneRuntimeEventPerSession(
   database: D1Database,
   records: readonly OneRuntimeEventPerSessionInput[],
 ): Promise<Map<SessionId, OneRuntimeEventPerSessionAllocation>> {
-  const sessionIds = [...new Set(records.map((record) => record.sessionId))];
-  const recordsBySessionId = new Map(records.map((record) => [record.sessionId, record]));
-  const allocations = new Map<SessionId, OneRuntimeEventPerSessionAllocation>();
   const appDb = getAppDatabase(database);
+  const ordinarySessionIds: SessionId[] = [];
+  const terminalSessionIds: SessionId[] = [];
+
+  for (const record of records) {
+    const sessionIds = canWriteAfterTerminatedSession([record])
+      ? terminalSessionIds
+      : ordinarySessionIds;
+    sessionIds.push(record.sessionId);
+  }
+
+  const statements: D1PreparedStatement[] = [];
 
-  for (const sessionId of sessionIds) {
-    const record = recordsBySessionId.get(sessionId);
-    const allowTerminatedSession =
-      record === undefined ? false : canWriteAfterTerminatedSession([record]);
-    const session =
-      (await appDb
+  for (const [sessionIds, allowTerminatedSession] of [
+    [ordinarySessionIds, false],
+    [terminalSessionIds, true],
+  ] as const) {
+    for (
+      let index = 0;
+      index < sessionIds.length;
+      index += MAX_SESSION_EVENT_ALLOCATIONS_PER_UPDATE
+    ) {
+      const query = appDb
         .update(sessionsTable)
         .set({
           runtimeEventSeqCursor: sql`${sessionsTable.runtimeEventSeqCursor} + 1`,
         })
         .where(
           and(
-            eq(sessionsTable.id, sessionId),
+            inArray(
+              sessionsTable.id,
+              sessionIds.slice(index, index + MAX_SESSION_EVENT_ALLOCATIONS_PER_UPDATE),
+            ),
             isNull(sessionsTable.archivedAt),
             inArray(sessionsTable.status, sessionWritableStatusValues(allowTerminatedSession)),
           ),
@@ -345,13 +362,25 @@ async function allocateOneRuntimeEventPerSession(
           seq: sessionsTable.runtimeEventSeqCursor,
           sessionId: sessionsTable.id,
         })
-        .get()) ?? null;
+        .toSQL();
 
-    if (session !== null) {
-      allocations.set(session.sessionId, {
-        agentId: session.agentId,
-        seq: session.seq,
-        sessionId: session.sessionId,
+      statements.push(database.prepare(query.sql).bind(...query.params));
+    }
+  }
+
+  const allocations = new Map<SessionId, OneRuntimeEventPerSessionAllocation>();
+  const results = await database.batch<{
+    agent_id: AgentId;
+    id: SessionId;
+    runtime_event_seq_cursor: number;
+  }>(statements);
+
+  for (const result of results) {
+    for (const row of result.results) {
+      allocations.set(row.id, {
+        agentId: row.agent_id,
+        seq: row.runtime_event_seq_cursor,
+        sessionId: row.id,
       });
     }
   }

apps/api/tests/session-runtime-event-store.test.ts

@@ -208,6 +208,55 @@ describe("session runtime event store", () => {
     expect(count?.count).toBe(0);
   });
 
+  test("batches one-event allocations within D1's bound parameter limit", async () => {
+    const database = createRuntimeEventStoreDatabase({ maxBoundParams: 100 });
+    const originalBatch = database.batch.bind(database);
+    const batchSizes: number[] = [];
+    database.batch = async <T = unknown>(statements: D1PreparedStatement[]) => {
+      batchSizes.push(statements.length);
+      return originalBatch<T>(statements);
+    };
+
+    for (let index = 2; index <= 98; index += 1) {
+      await database
+        .prepare("INSERT INTO session (id, agent_id, archived_at, status) VALUES (?, ?, NULL, ?)")
+        .bind(`session-${index}`, "01J00000000000000000000009", "TERMINATED")
+        .run();
+    }
+
+    const records = Array.from({ length: 98 }, (_, index) => {
+      const sessionId = `session-${index + 1}`;
+      const terminated = index > 0;
+
+      return {
+        event: runtimeEvent({
+          id: `event-${index + 1}`,
+          kind: terminated ? "session.lifecycle.updated" : "agent.task.updated",
+          occurredAtMs: 4_000 + index,
+          payload: terminated
+            ? { status: "TERMINATED" }
+            : {
+                agentId: "01J00000000000000000000009",
+                operation: "restart",
+                startedAt: new Date(4_000 + index).toISOString(),
+                status: "running",
+              },
+          sessionId,
+        }),
+        occurredAt: 4_000 + index,
+        sessionId,
+      };
+    });
+    const result = await persistOneRuntimeEventPerSession(database, { records });
+    const allocated = await database
+      .prepare("SELECT COUNT(*) AS count FROM session WHERE runtime_event_seq_cursor = 1")
+      .first<{ count: number }>();
+
+    expect(result).toEqual({ persistedCount: 98, skippedSessionIds: [] });
+    expect(allocated?.count).toBe(98);
+    expect(batchSizes[0]).toBe(3);
+  });
+
   test("persists mixed source ids and skips source replays before allocating sequence", async () => {
     const database = createRuntimeEventStoreDatabase();