From 4d62b1be57e3f94e72f6d14ad644d306e546d704 Mon Sep 17 00:00:00 2001 From: lyon Date: Fri, 19 Jun 2026 12:30:43 +0800 Subject: [PATCH] fix: reduce durable workbench trace read pressure --- internal/cloud/access-control.ts | 9 +- .../migrations/0001_cloud_core_skeleton.sql | 3 + internal/db/runtime-store.test.ts | 36 +++++++- internal/db/runtime-store.ts | 83 ++++++++++++++++++- 4 files changed, 125 insertions(+), 6 deletions(-) diff --git a/internal/cloud/access-control.ts b/internal/cloud/access-control.ts index 26aadcf6..c5460670 100644 --- a/internal/cloud/access-control.ts +++ b/internal/cloud/access-control.ts @@ -1301,7 +1301,14 @@ class PostgresAccessStore extends MemoryAccessStore { return pgAgentSession(result.rows?.[0]); } async getAgentSession(id) { await this.ensureSchema(); const result = await this.query("SELECT id, project_id, agent_id, status, started_at, ended_at, owner_user_id, conversation_id, thread_id, last_trace_id, session_json, updated_at FROM agent_sessions WHERE id = $1 LIMIT 1", [id]); return pgAgentSession(result.rows?.[0]); } - async getAgentSessionByTraceId(traceId) { await this.ensureSchema(); const result = await this.query("SELECT id, project_id, agent_id, status, started_at, ended_at, owner_user_id, conversation_id, thread_id, last_trace_id, session_json, updated_at FROM agent_sessions WHERE last_trace_id = $1 OR session_json LIKE $2 ORDER BY updated_at DESC NULLS LAST LIMIT 1", [traceId, `%${traceId}%`]); return pgAgentSession(result.rows?.[0]); } + async getAgentSessionByTraceId(traceId) { + await this.ensureSchema(); + const exact = await this.query("SELECT id, project_id, agent_id, status, started_at, ended_at, owner_user_id, conversation_id, thread_id, last_trace_id, session_json, updated_at FROM agent_sessions WHERE last_trace_id = $1 ORDER BY updated_at DESC NULLS LAST LIMIT 1", [traceId]); + const exactSession = pgAgentSession(exact.rows?.[0]); + if (exactSession) return exactSession; + const fallback = await this.query("SELECT id, project_id, agent_id, status, started_at, ended_at, owner_user_id, conversation_id, thread_id, last_trace_id, session_json, updated_at FROM agent_sessions WHERE session_json LIKE $1 ORDER BY updated_at DESC NULLS LAST LIMIT 1", [`%${traceId}%`]); + return pgAgentSession(fallback.rows?.[0]); + } async listAgentSessionsForUser(input = {}) { await this.ensureSchema(); const role = textOr(input.actorRole, "user"); diff --git a/internal/db/migrations/0001_cloud_core_skeleton.sql b/internal/db/migrations/0001_cloud_core_skeleton.sql index 1d2cddbd..47a04ce1 100644 --- a/internal/db/migrations/0001_cloud_core_skeleton.sql +++ b/internal/db/migrations/0001_cloud_core_skeleton.sql @@ -162,6 +162,9 @@ CREATE TABLE IF NOT EXISTS agent_trace_events ( event_json TEXT NOT NULL DEFAULT '{}', occurred_at TEXT NOT NULL ); +CREATE INDEX IF NOT EXISTS idx_agent_trace_events_trace_order ON agent_trace_events(trace_id, occurred_at, id); +CREATE INDEX IF NOT EXISTS idx_agent_trace_events_session_order ON agent_trace_events(agent_session_id, occurred_at, id); +CREATE INDEX IF NOT EXISTS idx_agent_trace_events_worker_order ON agent_trace_events(worker_session_id, occurred_at, id); CREATE TABLE IF NOT EXISTS evidence_records ( id TEXT PRIMARY KEY, diff --git a/internal/db/runtime-store.test.ts b/internal/db/runtime-store.test.ts index ec15734d..09f8caec 100644 --- a/internal/db/runtime-store.test.ts +++ b/internal/db/runtime-store.test.ts @@ -642,9 +642,34 @@ test("configured postgres runtime persists and queries Code Agent trace events", assert.equal(events.events[0].traceId, traceId); assert.equal(events.events[0].terminal, true); assert.equal(readiness.counts.agentTraceEvents, 1); + assert.ok(queryClient.calls.some((call) => call.sql.startsWith("CREATE INDEX IF NOT EXISTS idx_agent_trace_events_trace_order"))); assert.ok(queryClient.calls.some((call) => call.sql.startsWith("INSERT INTO agent_trace_events"))); }); +test("configured postgres runtime pushes trace event query filters into SQL", async () => { + const queryClient = createFakePostgresClient({ migrationReady: true }); + 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-17T03:50:00.000Z" + }); + + await store.writeAgentTraceEvent({ event: { traceId: "trc_sql_filter_a", sessionId: "ses_filter_a", seq: 1, level: "info", label: "a", createdAt: "2026-06-17T03:50:01.000Z" } }); + await store.writeAgentTraceEvent({ event: { traceId: "trc_sql_filter_b", sessionId: "ses_filter_b", seq: 1, level: "info", label: "b", createdAt: "2026-06-17T03:50:02.000Z" } }); + + const events = await store.queryAgentTraceEvents({ sessionId: "ses_filter_b" }); + + assert.equal(events.count, 1); + assert.equal(events.events[0].traceId, "trc_sql_filter_b"); + const select = queryClient.calls.findLast((call) => call.sql.startsWith("SELECT event_json FROM agent_trace_events")); + assert.match(select.sql, /WHERE agent_session_id = \$1/u); + assert.deepEqual(select.params, ["ses_filter_b"]); +}); + function createFakePostgresClient({ migrationReady = true, migrationErrorCode = null, @@ -695,6 +720,9 @@ function createFakePostgresClient({ const row = state.hwlab_schema_migrations.get(params[0]); return { rows: row ? [row] : [] }; } + if (sql.startsWith("CREATE INDEX IF NOT EXISTS idx_agent_trace_events_")) { + return { rows: [] }; + } if (sql.startsWith("SELECT gateway_session_json FROM gateway_sessions")) { return jsonSelect(state.gateway_sessions, params[0], "gateway_session_json"); } @@ -718,9 +746,15 @@ function createFakePostgresClient({ error.code = readErrorCode; throw error; } - const traceId = sql.includes("WHERE trace_id = $1") ? params[0] : null; + const traceId = sql.includes("trace_id = $1") ? params[0] : null; + const sessionId = sql.includes("agent_session_id = $1") ? params[0] : null; + const workerSessionId = sql.includes("worker_session_id = $1") ? params[0] : null; + const level = sql.includes("level = $1") ? params[0] : null; const rows = [...state.agent_trace_events.values()] .filter((record) => !traceId || record.trace_id === traceId) + .filter((record) => !sessionId || record.agent_session_id === sessionId) + .filter((record) => !workerSessionId || record.worker_session_id === workerSessionId) + .filter((record) => !level || record.level === level) .sort((left, right) => String(left.occurred_at).localeCompare(String(right.occurred_at)) || String(left.id).localeCompare(String(right.id))) .map((record) => ({ event_json: record.event_json })); return { rows }; diff --git a/internal/db/runtime-store.ts b/internal/db/runtime-store.ts index 4c3017d9..d5d4e26c 100644 --- a/internal/db/runtime-store.ts +++ b/internal/db/runtime-store.ts @@ -54,6 +54,12 @@ const postgresCountTables = Object.freeze([ "agent_trace_events" ]); +const postgresRuntimeReadIndexes = Object.freeze([ + "CREATE INDEX IF NOT EXISTS idx_agent_trace_events_trace_order ON agent_trace_events(trace_id, occurred_at, id)", + "CREATE INDEX IF NOT EXISTS idx_agent_trace_events_session_order ON agent_trace_events(agent_session_id, occurred_at, id)", + "CREATE INDEX IF NOT EXISTS idx_agent_trace_events_worker_order ON agent_trace_events(worker_session_id, occurred_at, id)" +]); + export function createCloudRuntimeStore(options = {}) { return new CloudRuntimeStore(options); } @@ -629,6 +635,8 @@ export class PostgresCloudRuntimeStore { this.queryClient = queryClient; this.pgModuleLoader = pgModuleLoader; this.pool = null; + this.runtimeReadIndexesReady = false; + this.runtimeReadIndexesReadyPromise = null; this.memory = new CloudRuntimeStore({ now }); this.lastReadiness = this.blockedSummary({ blocker: dbUrl || queryClient ? RUNTIME_DURABLE_ADAPTER_QUERY_BLOCKED : RUNTIME_DURABLE_ADAPTER_UNCONFIGURED, @@ -753,6 +761,37 @@ export class PostgresCloudRuntimeStore { return this.summary(); } + try { + await this.ensureRuntimeReadIndexes(); + } catch (error) { + const classified = classifyRuntimeDbError(error); + this.lastReadiness = this.blockedSummary({ + ...classified, + schema, + migration, + reason: "Postgres runtime schema is present, but Workbench trace read indexes could not be ensured", + connection: { + queryAttempted: true, + queryResult: classified.connection?.queryResult ?? "index_blocked", + errorCode: classified.connection?.errorCode ?? null + }, + gates: runtimeGates({ + ssl: classified.blocker === RUNTIME_DURABLE_ADAPTER_SSL_BLOCKED + ? blockedGate({ checked: true, blocker: classified.blocker }) + : readyGate(), + auth: classified.blocker === RUNTIME_DURABLE_ADAPTER_AUTH_BLOCKED + ? blockedGate({ checked: true, blocker: classified.blocker }) + : classified.blocker === RUNTIME_DURABLE_ADAPTER_SSL_BLOCKED + ? notCheckedGate() + : readyGate(), + schema: readyGate(), + migration: readyGate(), + durability: blockedGate({ checked: true, blocker: classified.blocker }) + }) + }); + return this.summary(); + } + let counts; try { counts = await this.readCounts(); @@ -948,10 +987,29 @@ export class PostgresCloudRuntimeStore { async queryAgentTraceEvents(params = {}) { await this.assertReadyForDurableReads("agent.trace-events.query"); const traceId = textOr(params.traceId, ""); - const sql = traceId - ? "SELECT event_json FROM agent_trace_events WHERE trace_id = $1 ORDER BY occurred_at ASC, id ASC" - : "SELECT event_json FROM agent_trace_events ORDER BY occurred_at ASC, id ASC"; - const result = await this.queryDurableReadRows("agent.trace-events.query", sql, traceId ? [traceId] : []); + const sessionId = textOr(params.sessionId, ""); + const workerSessionId = textOr(params.workerSessionId, ""); + const level = textOr(params.level, ""); + const clauses = []; + const queryParams = []; + if (traceId) { + queryParams.push(traceId); + clauses.push(`trace_id = $${queryParams.length}`); + } + if (sessionId) { + queryParams.push(sessionId); + clauses.push(`agent_session_id = $${queryParams.length}`); + } + if (workerSessionId) { + queryParams.push(workerSessionId); + clauses.push(`worker_session_id = $${queryParams.length}`); + } + if (level) { + queryParams.push(level); + clauses.push(`level = $${queryParams.length}`); + } + const sql = `SELECT event_json FROM agent_trace_events${clauses.length ? ` WHERE ${clauses.join(" AND ")}` : ""} ORDER BY occurred_at ASC, id ASC`; + const result = await this.queryDurableReadRows("agent.trace-events.query", sql, queryParams); const events = result.rows .map((row) => parseJsonColumn(row.event_json, null)) .filter(Boolean) @@ -1200,6 +1258,23 @@ export class PostgresCloudRuntimeStore { }; } + async ensureRuntimeReadIndexes() { + if (this.runtimeReadIndexesReady) return; + if (!this.runtimeReadIndexesReadyPromise) { + this.runtimeReadIndexesReadyPromise = (async () => { + for (const sql of postgresRuntimeReadIndexes) { + await this.query(sql, []); + } + this.runtimeReadIndexesReady = true; + })(); + } + try { + await this.runtimeReadIndexesReadyPromise; + } finally { + if (!this.runtimeReadIndexesReady) this.runtimeReadIndexesReadyPromise = null; + } + } + async query(sql, params = []) { const client = await this.getQueryClient(); return client.query(sql, params);