From f637f32631e7aed7ec61bda6a93a0de0558ce2f7 Mon Sep 17 00:00:00 2001 From: Lyon <88232613+pikasTech@users.noreply.github.com> Date: Sat, 20 Jun 2026 10:29:00 +0800 Subject: [PATCH] fix: parallelize workbench read-model fact reads (#1680) --- internal/cloud/server-workbench-http.ts | 13 +-- internal/db/runtime-store.test.ts | 109 +++++++++++++++++++++++- internal/db/runtime-store.ts | 17 ++-- 3 files changed, 124 insertions(+), 15 deletions(-) diff --git a/internal/cloud/server-workbench-http.ts b/internal/cloud/server-workbench-http.ts index 75e1d381..57d26205 100644 --- a/internal/cloud/server-workbench-http.ts +++ b/internal/cloud/server-workbench-http.ts @@ -388,13 +388,14 @@ async function queryFactsForSessionSummaries(readModel, sessions = []) { const sessionIds = uniqueText(sessions.map(factSessionId)); const traceIds = uniqueText(sessions.map(factLastTraceId)); if (sessionIds.length === 0 && traceIds.length === 0) return { facts: emptyFactSet(), count: 0, durable: false, persistence: null }; - const bySession = sessionIds.length > 0 - ? await readModel.queryFacts({ sessionIds, families: ["messages", "parts", "turns"] }) - : { facts: emptyFactSet(), count: 0, durable: false, persistence: null }; + const bySessionPromise = sessionIds.length > 0 + ? readModel.queryFacts({ sessionIds, families: ["messages", "parts", "turns"] }) + : Promise.resolve({ facts: emptyFactSet(), count: 0, durable: false, persistence: null }); + const byTracePromise = traceIds.length > 0 + ? readModel.queryFacts({ traceIds, families: ["checkpoints"] }) + : Promise.resolve({ facts: emptyFactSet(), count: 0, durable: false, persistence: null }); + const [bySession, byTrace] = await Promise.all([bySessionPromise, byTracePromise]); if (bySession.error) return bySession; - const byTrace = traceIds.length > 0 - ? await readModel.queryFacts({ traceIds, families: ["checkpoints"] }) - : { facts: emptyFactSet(), count: 0, durable: false, persistence: bySession.persistence ?? null }; if (byTrace.error) return byTrace; const facts = mergeFactSets(bySession.facts, byTrace.facts); return { facts, count: Object.values(facts).reduce((sum, rows) => sum + rows.length, 0), durable: bySession.durable ?? byTrace.durable, persistence: bySession.persistence ?? byTrace.persistence ?? null }; diff --git a/internal/db/runtime-store.test.ts b/internal/db/runtime-store.test.ts index a339e432..ffd0761f 100644 --- a/internal/db/runtime-store.test.ts +++ b/internal/db/runtime-store.test.ts @@ -806,6 +806,106 @@ test("configured postgres runtime persists and queries Workbench durable facts", assert.ok(queryClient.calls.some((call) => call.sql.startsWith("SELECT session_json FROM workbench_sessions WHERE session_id = $1"))); }); +test("configured postgres runtime queries requested workbench fact families concurrently", async () => { + const expectedReads = 4; + let started = 0; + let releaseReads; + let resolveAllStarted; + const readGate = new Promise((resolve) => { + releaseReads = resolve; + }); + const allStarted = new Promise((resolve) => { + resolveAllStarted = resolve; + }); + const queryClient = createFakePostgresClient({ + migrationReady: true, + beforeWorkbenchFactRead: async (sql) => { + if (!/^SELECT (message_json|part_json|turn_json|checkpoint_json) FROM/u.test(sql)) return; + started += 1; + if (started === expectedReads) resolveAllStarted(); + await readGate; + } + }); + const store = createConfiguredCloudRuntimeStore({ + env: { + HWLAB_CLOUD_RUNTIME_ADAPTER: "postgres", + HWLAB_CLOUD_RUNTIME_DURABLE: "true" + }, + dbUrl: "postgres://hwlab_redacted@db.example.invalid:5432/hwlab", + queryClient, + now: () => "2026-06-20T10:00:00.000Z" + }); + + await store.writeWorkbenchFacts({ + facts: { + sessions: [], + messages: [{ + messageId: "msg_parallel", + sessionId: "ses_parallel", + turnId: "turn_parallel", + traceId: "trc_parallel", + role: "assistant", + status: "completed", + text: "done", + updatedAt: "2026-06-20T10:00:01.000Z" + }], + parts: [{ + partId: "part_parallel", + messageId: "msg_parallel", + sessionId: "ses_parallel", + turnId: "turn_parallel", + traceId: "trc_parallel", + partIndex: 0, + partType: "text", + status: "completed", + text: "done", + updatedAt: "2026-06-20T10:00:02.000Z" + }], + turns: [{ + turnId: "turn_parallel", + sessionId: "ses_parallel", + traceId: "trc_parallel", + messageId: "msg_parallel", + status: "completed", + finalResponse: { text: "done" }, + updatedAt: "2026-06-20T10:00:03.000Z" + }], + traceEvents: [], + checkpoints: [{ + traceId: "trc_parallel", + sessionId: "ses_parallel", + turnId: "turn_parallel", + runId: "run_parallel", + commandId: "cmd_parallel", + projectionStatus: "caught_up", + projectionHealth: "healthy", + updatedAt: "2026-06-20T10:00:04.000Z" + }] + } + }); + + const pending = store.queryWorkbenchFacts({ + sessionId: "ses_parallel", + traceId: "trc_parallel", + families: ["messages", "parts", "turns", "checkpoints"], + limit: 10 + }); + const startedConcurrently = await Promise.race([ + allStarted.then(() => true), + new Promise((resolve) => setTimeout(() => resolve(false), 150)) + ]); + releaseReads(); + const loaded = await pending; + + assert.equal(startedConcurrently, true); + assert.equal(started, expectedReads); + assert.equal(loaded.count, 4); + assert.equal(loaded.facts.messages[0].messageId, "msg_parallel"); + assert.equal(loaded.facts.parts[0].partId, "part_parallel"); + assert.equal(loaded.facts.turns[0].turnId, "turn_parallel"); + assert.equal(loaded.facts.checkpoints[0].traceId, "trc_parallel"); +}); + test("configured postgres runtime pushes trace event query filters into SQL", async () => { const queryClient = createFakePostgresClient({ migrationReady: true }); const store = createConfiguredCloudRuntimeStore({ @@ -834,7 +934,8 @@ function createFakePostgresClient({ migrationReady = true, migrationErrorCode = null, countErrorCode = null, - readErrorCode = null + readErrorCode = null, + beforeWorkbenchFactRead = null } = {}) { const state = { gateway_sessions: new Map(), @@ -952,21 +1053,27 @@ function createFakePostgresClient({ return { rows: rows.map((record) => ({ projection_json: record.projection_json })) }; } if (sql.startsWith("SELECT session_json FROM workbench_sessions")) { + if (beforeWorkbenchFactRead) await beforeWorkbenchFactRead(sql, params); return workbenchFactRows(state.workbench_sessions, "session_json", sql, params, readErrorCode); } if (sql.startsWith("SELECT message_json FROM workbench_messages")) { + if (beforeWorkbenchFactRead) await beforeWorkbenchFactRead(sql, params); return workbenchFactRows(state.workbench_messages, "message_json", sql, params, readErrorCode); } if (sql.startsWith("SELECT part_json FROM workbench_parts")) { + if (beforeWorkbenchFactRead) await beforeWorkbenchFactRead(sql, params); return workbenchFactRows(state.workbench_parts, "part_json", sql, params, readErrorCode); } if (sql.startsWith("SELECT turn_json FROM workbench_turns")) { + if (beforeWorkbenchFactRead) await beforeWorkbenchFactRead(sql, params); return workbenchFactRows(state.workbench_turns, "turn_json", sql, params, readErrorCode); } if (sql.startsWith("SELECT event_json FROM workbench_trace_events")) { + if (beforeWorkbenchFactRead) await beforeWorkbenchFactRead(sql, params); return workbenchFactRows(state.workbench_trace_events, "event_json", sql, params, readErrorCode); } if (sql.startsWith("SELECT checkpoint_json FROM workbench_projection_checkpoints")) { + if (beforeWorkbenchFactRead) await beforeWorkbenchFactRead(sql, params); return workbenchFactRows(state.workbench_projection_checkpoints, "checkpoint_json", sql, params, readErrorCode); } if (sql.startsWith("SELECT metadata_json FROM evidence_records")) { diff --git a/internal/db/runtime-store.ts b/internal/db/runtime-store.ts index 45c2bd3a..6837f79e 100644 --- a/internal/db/runtime-store.ts +++ b/internal/db/runtime-store.ts @@ -1215,14 +1215,15 @@ export class PostgresCloudRuntimeStore { async queryWorkbenchFacts(params = {}) { await this.assertReadyForDurableReads("workbench.facts.query"); const families = workbenchFactFamilySet(params); - const facts = { - sessions: families.has("sessions") ? await this.queryWorkbenchFactRows("workbench.sessions.query", "workbench_sessions", "session_json", params, { sessionId: "session_id", ownerUserId: "owner_user_id", projectId: "project_id", conversationId: "conversation_id", threadId: "thread_id", traceId: "last_trace_id", lastTraceId: "last_trace_id", status: "status" }) : [], - messages: families.has("messages") ? await this.queryWorkbenchFactRows("workbench.messages.query", "workbench_messages", "message_json", params, { messageId: "message_id", sessionId: "session_id", turnId: "turn_id", traceId: "trace_id", role: "role", status: "status" }) : [], - parts: families.has("parts") ? await this.queryWorkbenchFactRows("workbench.parts.query", "workbench_parts", "part_json", params, { partId: "part_id", messageId: "message_id", sessionId: "session_id", turnId: "turn_id", traceId: "trace_id", partType: "part_type", status: "status" }) : [], - turns: families.has("turns") ? await this.queryWorkbenchFactRows("workbench.turns.query", "workbench_turns", "turn_json", params, { turnId: "turn_id", sessionId: "session_id", traceId: "trace_id", messageId: "message_id", status: "status" }) : [], - traceEvents: families.has("traceEvents") ? await this.queryWorkbenchFactRows("workbench.trace-events.query", "workbench_trace_events", "event_json", params, { id: "id", traceId: "trace_id", sessionId: "session_id", turnId: "turn_id", messageId: "message_id", eventType: "event_type" }) : [], - checkpoints: families.has("checkpoints") ? await this.queryWorkbenchFactRows("workbench.projection-checkpoints.query", "workbench_projection_checkpoints", "checkpoint_json", params, { traceId: "trace_id", sessionId: "session_id", turnId: "turn_id", runId: "run_id", commandId: "command_id", projectionStatus: "projection_status", projectionHealth: "projection_health" }) : [] - }; + const entries = await Promise.all([ + ["sessions", families.has("sessions") ? this.queryWorkbenchFactRows("workbench.sessions.query", "workbench_sessions", "session_json", params, { sessionId: "session_id", ownerUserId: "owner_user_id", projectId: "project_id", conversationId: "conversation_id", threadId: "thread_id", traceId: "last_trace_id", lastTraceId: "last_trace_id", status: "status" }) : Promise.resolve([])], + ["messages", families.has("messages") ? this.queryWorkbenchFactRows("workbench.messages.query", "workbench_messages", "message_json", params, { messageId: "message_id", sessionId: "session_id", turnId: "turn_id", traceId: "trace_id", role: "role", status: "status" }) : Promise.resolve([])], + ["parts", families.has("parts") ? this.queryWorkbenchFactRows("workbench.parts.query", "workbench_parts", "part_json", params, { partId: "part_id", messageId: "message_id", sessionId: "session_id", turnId: "turn_id", traceId: "trace_id", partType: "part_type", status: "status" }) : Promise.resolve([])], + ["turns", families.has("turns") ? this.queryWorkbenchFactRows("workbench.turns.query", "workbench_turns", "turn_json", params, { turnId: "turn_id", sessionId: "session_id", traceId: "trace_id", messageId: "message_id", status: "status" }) : Promise.resolve([])], + ["traceEvents", families.has("traceEvents") ? this.queryWorkbenchFactRows("workbench.trace-events.query", "workbench_trace_events", "event_json", params, { id: "id", traceId: "trace_id", sessionId: "session_id", turnId: "turn_id", messageId: "message_id", eventType: "event_type" }) : Promise.resolve([])], + ["checkpoints", families.has("checkpoints") ? this.queryWorkbenchFactRows("workbench.projection-checkpoints.query", "workbench_projection_checkpoints", "checkpoint_json", params, { traceId: "trace_id", sessionId: "session_id", turnId: "turn_id", runId: "run_id", commandId: "command_id", projectionStatus: "projection_status", projectionHealth: "projection_health" }) : Promise.resolve([])] + ].map(async ([family, rowsPromise]) => [family, await rowsPromise])); + const facts = Object.fromEntries(entries); return withPersistence({ facts, count: Object.values(facts).reduce((sum, rows) => sum + rows.length, 0) }, this.summary()); }