From 5b856431d179b8e429b36708a7a0c94842e9caba Mon Sep 17 00:00:00 2001 From: Bestony Date: Tue, 28 Jul 2026 17:21:00 +0800 Subject: [PATCH] perf(server): replace inbox context and route-validation N+1 queries MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Inbox delivery and bound-agent route validation both scaled their query count with input size, inside paths that hold locks or run on every heartbeat. collectPrecedingContext ran one previous-notify query per chat plus one silent-context query per trigger, all inside the claiming transaction — a 50-entry drain meant 50+ sequential round trips holding FOR UPDATE locks. It now issues two statements regardless of batch size: a CROSS JOIN LATERAL over a jsonb_to_recordset of per-trigger bounds that selects and locks the silent rows, then one keyed read of the message payloads through the typed builder. Each trigger's lower bound comes from the preceding trigger in the same batch, falling back via COALESCE to the table cursor for a chat's first trigger, so trigger-relative windows, the 24h window, the MAX_ENTRIES cap and SKIP LOCKED semantics are preserved. ensureAgentStillRoutedHere probed one agent per query, and both inbox:ack and every heartbeat validate the socket's whole bound set with no hard agent-count cap. Route resolution is now split into a pure decision step plus a batched ensureAgentsStillRoutedHere that reads all candidates in one inArray query; the single-agent helper delegates to it so the two paths cannot drift. The heartbeat repair pass still re-validates after recordClientHeartbeat — the restore can flip an agent back to active, and only the post-restore row decides whether repair is safe — but does so in one query over the restored set. Tests: a multi-chat multi-trigger batch drain pinning each trigger's window, a statement-count assertion that a 6-trigger drain costs the same as a 2-trigger drain, and a heartbeat test asserting one route query for two bound agents. All three were verified to fail against the pre-fix shapes. Refs #1724 --- .../src/__tests__/inbox-ws-push.test.ts | 126 ++++++++++ .../__tests__/ws-client-branch-fake.test.ts | 100 +++++++- packages/server/src/api/agent/ws-client.ts | 101 ++++++-- packages/server/src/services/inbox.ts | 238 +++++++++++------- 4 files changed, 449 insertions(+), 116 deletions(-) diff --git a/packages/server/src/__tests__/inbox-ws-push.test.ts b/packages/server/src/__tests__/inbox-ws-push.test.ts index e4a3131a6..23871627c 100644 --- a/packages/server/src/__tests__/inbox-ws-push.test.ts +++ b/packages/server/src/__tests__/inbox-ws-push.test.ts @@ -1,8 +1,11 @@ import { and, asc, eq, inArray } from "drizzle-orm"; +import { drizzle } from "drizzle-orm/postgres-js"; import type { FastifyInstance } from "fastify"; +import postgres from "postgres"; import { describe, expect, it } from "vitest"; import { chatMembership } from "../db/schema/chat-membership.js"; import { inboxEntries } from "../db/schema/inbox-entries.js"; +import * as schema from "../db/schema/index.js"; import { createAgent, getAgent } from "../services/agent.js"; import { createChat } from "../services/chat.js"; import * as inboxService from "../services/inbox.js"; @@ -814,4 +817,127 @@ describe("inbox WS data-plane claim helpers", () => { const res = await inboxService.ackEntryByIdForBoundAgents(app.db, 1, []); expect(res).toEqual({ ok: false, reason: "not_found_or_not_bound" }); }); + + it("bundles trigger-relative context for every trigger of a multi-chat batch drain", async () => { + // PERF-061 moved preceding-context assembly from one query per trigger to a + // single set-based statement. The old shape walked chats sequentially and + // carried the previous-trigger cursor in a JS loop variable; the batched + // shape resolves every window in one LATERAL. This drains three triggers + // across two chats at once — the first multi-trigger, multi-chat claim in + // the suite — so a cursor that leaks across chats, or a chat's first + // trigger that forgets its already-delivered predecessor, fails here. + const app = getApp(); + const uid = crypto.randomUUID().slice(0, 6); + const human = await createTestAgent(app, { type: "human", name: `batch-h-${uid}` }); + const observer = await createTestAgent(app, { type: "agent", name: `batch-obs-${uid}` }); + const chatA = await createChat(app.db, human.agent.uuid, { + type: "group", + participantIds: [observer.agent.uuid], + }); + const chatB = await createChat(app.db, human.agent.uuid, { + type: "group", + participantIds: [observer.agent.uuid], + }); + + const silent = (chatId: string, content: string) => + sendMessage( + app.db, + chatId, + human.agent.uuid, + { source: "api", format: "text", content }, + { allowRecipientlessSend: true }, + ); + const trigger = (chatId: string, content: string) => + sendMessage(app.db, chatId, human.agent.uuid, { + source: "api", + format: "text", + content, + metadata: { mentions: [observer.agent.uuid] }, + }); + + // Chat A opens with a trigger claimed on its own, so the batch's first + // chat-A trigger has to resolve its lower bound from the table rather than + // from the batch — the COALESCE fallback in the batched statement. + await silent(chatA.id, "a-before-earlier-trigger"); + await trigger(chatA.id, "a-earlier-trigger"); + const earlier = await inboxService.claimBacklogForPush(app.db, observer.agent.inboxId, 1); + expect(earlier.map((e) => e.message.content)).toEqual(["a-earlier-trigger"]); + + await silent(chatA.id, "a-silent-1"); + await trigger(chatA.id, "a-trigger-1"); + await silent(chatA.id, "a-silent-2"); + await trigger(chatA.id, "a-trigger-2"); + await silent(chatB.id, "b-silent-1"); + await trigger(chatB.id, "b-trigger-1"); + + const claimed = await inboxService.claimBacklogForPush(app.db, observer.agent.inboxId, 10); + const precedingByTrigger = new Map( + claimed.map((entry) => [entry.message.content as string, entry.message.precedingMessages.map((p) => p.content)]), + ); + + expect([...precedingByTrigger.keys()].sort()).toEqual(["a-trigger-1", "a-trigger-2", "b-trigger-1"]); + // Bounded below by the already-delivered `a-earlier-trigger`, so the silent + // row that preceded it stays out. + expect(precedingByTrigger.get("a-trigger-1")).toEqual(["a-silent-1"]); + // Bounded below by the trigger claimed alongside it in this same batch. + expect(precedingByTrigger.get("a-trigger-2")).toEqual(["a-silent-2"]); + // Chat B has no earlier trigger at all, and chat A's cursor must not bleed + // into it. + expect(precedingByTrigger.get("b-trigger-1")).toEqual(["b-silent-1"]); + }); + + it("keeps drain statement count flat as the trigger batch grows", async () => { + // The N+1 guard for PERF-061: a drain's statement count must not scale with + // the number of claimed triggers. Asserting equality between a small and a + // large batch pins that directly — the pre-fix code issued one + // previous-notify query per chat plus one context query per trigger, so + // these two counts differed by four. + const app = getApp(); + const uid = crypto.randomUUID().slice(0, 6); + const human = await createTestAgent(app, { type: "human", name: `count-h-${uid}` }); + const observer = await createTestAgent(app, { type: "agent", name: `count-obs-${uid}` }); + const chat = await createChat(app.db, human.agent.uuid, { + type: "group", + participantIds: [observer.agent.uuid], + }); + + async function seedAndDrain(triggerCount: number): Promise { + for (let i = 0; i < triggerCount; i++) { + await sendMessage( + app.db, + chat.id, + human.agent.uuid, + { source: "api", format: "text", content: `silent-${uid}-${i}` }, + { allowRecipientlessSend: true }, + ); + await sendMessage(app.db, chat.id, human.agent.uuid, { + source: "api", + format: "text", + content: `trigger-${uid}-${i}`, + metadata: { mentions: [observer.agent.uuid] }, + }); + } + // Count on a dedicated connection so the shared app pool's traffic can't + // bleed into the sample. + const statements: string[] = []; + const client = postgres(process.env.DATABASE_URL ?? "", { + max: 1, + debug: (_connection, query) => { + statements.push(query); + }, + }); + try { + const db = Object.assign(drizzle(client, { schema }), { end: () => client.end() }); + const claimed = await inboxService.claimBacklogForPush(db as never, observer.agent.inboxId, 100); + expect(claimed).toHaveLength(triggerCount); + } finally { + await client.end(); + } + return statements.length; + } + + const smallBatch = await seedAndDrain(2); + const largeBatch = await seedAndDrain(6); + expect(largeBatch).toBe(smallBatch); + }); }); diff --git a/packages/server/src/__tests__/ws-client-branch-fake.test.ts b/packages/server/src/__tests__/ws-client-branch-fake.test.ts index b12fc549f..6f28db808 100644 --- a/packages/server/src/__tests__/ws-client-branch-fake.test.ts +++ b/packages/server/src/__tests__/ws-client-branch-fake.test.ts @@ -1,7 +1,7 @@ import { EventEmitter } from "node:events"; import { AUTH_REJECTED_CODES, type ClientMessage, type InboxEntryWithMessage } from "@first-tree/shared"; import { SignJWT } from "jose"; -import { afterEach, describe, expect, it, vi } from "vitest"; +import { afterEach, describe, expect, it, type Mock, vi } from "vitest"; import { clientWsRoutes } from "../api/agent/ws-client.js"; import type { inboxEntries } from "../db/schema/inbox-entries.js"; import * as activityService from "../services/activity.js"; @@ -60,6 +60,45 @@ function queuedDb(results: unknown[][]): unknown { }; } +/** Projection keys unique to the bound-agent route re-validation query. */ +function isRouteValidationProjection(projection: unknown): boolean { + if (!projection || typeof projection !== "object") return false; + const keys = Object.keys(projection); + return keys.includes("uuid") && keys.includes("runtimeProvider") && keys.includes("metadata"); +} + +/** + * Fake db that answers by projection shape rather than call order. + * + * `queuedDb`'s positional queue can't serve a multi-agent socket: the + * post-bind backlog drain runs detached, so its route-validation SELECT + * interleaves non-deterministically with the next `agent:bind` lookup and + * shifts every later slot. Routing by projection keeps each query family on + * its own queue. + */ +function projectionRoutedDb(opts: { + /** Auth-time user lookup — stable for the whole socket. */ + authUser: unknown[]; + /** Client row read at register and again on every `agent:bind`. */ + bindingClient: unknown[]; + /** Agent row per `agent:bind`, in bind order. */ + binds: unknown[][]; + /** Rows every bound-agent route re-validation resolves against. */ + routeRows: unknown[]; +}): unknown { + const binds = [...opts.binds]; + return { + select: vi.fn((projection?: unknown) => { + if (isRouteValidationProjection(projection)) return queryChain(opts.routeRows); + const keys = projection && typeof projection === "object" ? Object.keys(projection) : []; + if (keys.includes("displayName") && keys.includes("clientUserId")) return queryChain(binds.shift() ?? []); + if (keys.includes("retiredAt")) return queryChain(opts.bindingClient); + return queryChain(opts.authUser); + }), + update: vi.fn(() => queryChain([])), + }; +} + function throwingSelectDb(error: unknown): unknown { return { select: vi.fn(() => { @@ -326,6 +365,10 @@ describe("Agent client WS branch fakes", () => { function activeAgentRow(overrides: Record = {}): Record { return { id: "agent_1", + // Same column as `id` above (`agents.uuid`), under the alias the batched + // route re-validation selects it as — that path keys rows by uuid rather + // than taking the first row positionally. + uuid: "agent_1", displayName: "Agent", type: "agent", organizationId: "org_1", @@ -371,14 +414,24 @@ describe("Agent client WS branch fakes", () => { async function bindAgent(socket: FakeSocket, handler: WsHandler, ref = "bind-ok"): Promise { await authenticateAndRegister(socket, handler); + await emitBind(socket, "agent_1", ref); + } + + /** Bind an already-authenticated socket to one more agent. */ + async function emitBind(socket: FakeSocket, agentId: string, ref: string): Promise { await emitMessage(socket, { type: "agent:bind", - agentId: "agent_1", + agentId, ref, runtimeType: "claude-code", runtimeVersion: "test", }); - await waitUntil(() => socket.sent.some((frame) => (frame as { type?: string }).type === "agent:bound")); + await waitUntil(() => + socket.sent.some((frame) => { + const bound = frame as { type?: string; ref?: string }; + return bound.type === "agent:bound" && bound.ref === ref; + }), + ); } it("rejects first-bind races when the claim update returns no row", async () => { @@ -553,6 +606,47 @@ describe("Agent client WS branch fakes", () => { expect(socket.sent).toContainEqual({ type: "error", message: "Agent not bound" }); }); + it("re-validates every bound agent's route in one query per heartbeat", async () => { + // PERF-061: heartbeat used to probe each bound agent with its own SELECT, + // and a client has no hard agent-count cap. Pin the batched form — the + // assertion is on query count, so a regression to the sequential loop + // fails here even though both shapes produce the same routed set. + mockSuccessfulBindServices(); + // No restored agents → the repair pass is skipped, leaving exactly one + // route re-validation per heartbeat to assert on. + vi.spyOn(runtimeLivenessService, "recordClientHeartbeat").mockResolvedValue({ + clientUpdated: true, + restoredAgentIds: [], + }); + const agentRow2 = activeAgentRow({ id: "agent_2", uuid: "agent_2", inboxId: "inbox_2" }); + const db = projectionRoutedDb({ + authUser: [{ id: "user_1", status: "active" }], + bindingClient: [{ userId: "user_1", retiredAt: null }], + binds: [[activeAgentRow()], [agentRow2]], + // Every route re-validation resolves against both agents, so a batched + // caller gets one answer and a sequential caller would get two. + routeRows: [activeAgentRow(), agentRow2], + }) as { select: Mock }; + const { handler } = routeHarness(db); + const socket = new FakeSocket(); + await bindAgent(socket, handler, "bind-batch-1"); + await emitBind(socket, "agent_2", "bind-batch-2"); + + const routeValidationCalls = (): unknown[][] => + db.select.mock.calls.filter(([projection]) => isRouteValidationProjection(projection)); + + db.select.mockClear(); + await emitMessage(socket, { type: "heartbeat" }); + await waitUntil(() => socket.sent.some((frame) => (frame as { type?: string }).type === "heartbeat:ack")); + + expect(routeValidationCalls()).toHaveLength(1); + // …and that single statement really did validate both agents. + expect(runtimeLivenessService.recordClientHeartbeat).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ routedAgentIds: expect.arrayContaining(["agent_1", "agent_2"]) }), + ); + }); + it("throttles heartbeat-triggered inbox repair when the last repair is recent", async () => { mockSuccessfulBindServices(); const { handler } = routeHarness( diff --git a/packages/server/src/api/agent/ws-client.ts b/packages/server/src/api/agent/ws-client.ts index 3762af026..7cd8470d8 100644 --- a/packages/server/src/api/agent/ws-client.ts +++ b/packages/server/src/api/agent/ws-client.ts @@ -438,20 +438,29 @@ export function clientWsRoutes(notifier: Notifier, instanceId: string) { app.log.info({ clientId, agentId, reason }, "dropped stale local agent binding"); } - async function ensureAgentStillRoutedHere(agentId: string): Promise { - if (!isAgentStillRoutedHere(agentId)) return false; + /** Authoritative `agents` row shape consulted by the route re-validation. */ + type AuthoritativeRouteRow = { + clientId: string | null; + runtimeProvider: string; + status: string; + metadata: unknown; + }; + + /** + * Decide a single agent's fate against its authoritative row, dropping the + * local binding when the route has genuinely moved. Split out of + * `ensureAgentStillRoutedHere` so the single-agent and batched paths make + * byte-identical decisions — this is the piece that must not drift. + * + * A missing row (agent deleted) falls through to the drop, same as a row + * that now points at another client or runtime. The one exception is a + * suspended agent mid runtime-switch whose claim still names this client: + * that binding is being handed over deliberately, so we report "not routed + * here" without tearing down local state the switch will finish using. + */ + function resolveAuthoritativeRoute(agentId: string, row: AuthoritativeRouteRow | undefined): boolean { const info = boundAgents.get(agentId); if (!info || !clientId) return false; - const [row] = await app.db - .select({ - clientId: agents.clientId, - runtimeProvider: agents.runtimeProvider, - status: agents.status, - metadata: agents.metadata, - }) - .from(agents) - .where(eq(agents.uuid, agentId)) - .limit(1); if (row?.clientId === clientId && row.status === "active" && row.runtimeProvider === info.runtimeProvider) { return true; } @@ -468,6 +477,41 @@ export function clientWsRoutes(notifier: Notifier, instanceId: string) { return false; } + /** + * Batched route re-validation: one `WHERE uuid = ANY(...)` for the whole + * candidate set instead of one round trip per agent (PERF-061). ACK and + * every heartbeat validate the socket's full bound-agent set, and a client + * has no hard agent-count cap, so the sequential form put agent-count + * round trips on the heartbeat path. + * + * Returns the subset still authoritatively routed to this socket. + */ + async function ensureAgentsStillRoutedHere(agentIds: Iterable): Promise> { + const routed = new Set(); + if (!clientId) return routed; + const candidates = [...new Set(agentIds)].filter((id) => isAgentStillRoutedHere(id)); + if (candidates.length === 0) return routed; + const rows = await app.db + .select({ + uuid: agents.uuid, + clientId: agents.clientId, + runtimeProvider: agents.runtimeProvider, + status: agents.status, + metadata: agents.metadata, + }) + .from(agents) + .where(inArray(agents.uuid, candidates)); + const rowByAgentId = new Map(rows.map((row) => [row.uuid, row])); + for (const agentId of candidates) { + if (resolveAuthoritativeRoute(agentId, rowByAgentId.get(agentId))) routed.add(agentId); + } + return routed; + } + + async function ensureAgentStillRoutedHere(agentId: string): Promise { + return (await ensureAgentsStillRoutedHere([agentId])).has(agentId); + } + function inboxInFlightCount(agentId: string): number { const byChat = inboxInFlightByAgent.get(agentId); if (!byChat) return 0; @@ -1630,12 +1674,12 @@ export function clientWsRoutes(notifier: Notifier, instanceId: string) { await chainInboxDelivery("__socket", async () => { try { - const routedBoundAgents = []; - for (const agent of boundAgents.values()) { - if (await ensureAgentStillRoutedHere(agent.agentId)) { - routedBoundAgents.push(agent); - } - } + // `inbox:ack` carries no agentId, so every bound inbox is a + // candidate owner — validate the whole set in one query. + const routedAgentIds = await ensureAgentsStillRoutedHere(boundAgents.keys()); + const routedBoundAgents = [...boundAgents.values()].filter((agent) => + routedAgentIds.has(agent.agentId), + ); const ackResult = await inboxService.ackEntryByIdForBoundAgents( app.db, entryId, @@ -1768,10 +1812,7 @@ export function clientWsRoutes(notifier: Notifier, instanceId: string) { }); } else if (type === "heartbeat") { if (clientId && connectionManager.isActiveClientConnection(clientId, socket)) { - const routedAgentIds = []; - for (const id of boundAgents.keys()) { - if (await ensureAgentStillRoutedHere(id)) routedAgentIds.push(id); - } + const routedAgentIds = [...(await ensureAgentsStillRoutedHere(boundAgents.keys()))]; const pausedReason = msg && typeof msg === "object" && "pausedReason" in msg ? ((msg as { pausedReason?: "auth_rejected" | "auth_refresh_failed" | null }).pausedReason ?? null) @@ -1782,10 +1823,20 @@ export function clientWsRoutes(notifier: Notifier, instanceId: string) { routedAgentIds, pausedReason, }); + // Re-validate rather than reuse `routedAgentIds`: the heartbeat + // itself may have flipped a suspended agent back to active, and + // only the post-restore row decides whether backlog repair is + // safe. Scoped to the restored ∩ bound set, in one query. const repairableAgentIds = new Set(liveness.restoredAgentIds); - for (const info of boundAgents.values()) { - if (repairableAgentIds.has(info.agentId) && (await ensureAgentStillRoutedHere(info.agentId))) { - maybeRepairInboxBacklog(info.agentId, info.inboxId); + const repairCandidates = [...boundAgents.values()].filter((info) => + repairableAgentIds.has(info.agentId), + ); + if (repairCandidates.length > 0) { + const repairRoutedIds = await ensureAgentsStillRoutedHere( + repairCandidates.map((info) => info.agentId), + ); + for (const info of repairCandidates) { + if (repairRoutedIds.has(info.agentId)) maybeRepairInboxBacklog(info.agentId, info.inboxId); } } await reconcilePinnedAgentsForClient(); diff --git a/packages/server/src/services/inbox.ts b/packages/server/src/services/inbox.ts index 62ab8d04a..9e2b62d96 100644 --- a/packages/server/src/services/inbox.ts +++ b/packages/server/src/services/inbox.ts @@ -4,7 +4,7 @@ import { messageSourceSchema, type PrecedingMessage, } from "@first-tree/shared"; -import { and, asc, desc, eq, gt, inArray, isNull, lt, sql } from "drizzle-orm"; +import { and, asc, desc, eq, inArray, isNull, sql } from "drizzle-orm"; import type { PgDatabase, PgQueryResultHKT } from "drizzle-orm/pg-core"; import type { PostgresJsDatabase } from "drizzle-orm/postgres-js"; import type { Database } from "../db/connection.js"; @@ -29,7 +29,7 @@ export type AckEntryResult = | { ok: false; reason: "not_found_or_not_bound" | "non_notify" | "prefix_gap" }; /** Structurally-typed DB so both `Database` and transaction clients work. */ -type TxLike = Pick>, "select" | "update" | "delete" | "insert">; +type TxLike = Pick>, "select" | "update" | "delete" | "insert" | "execute">; /** Wider DB shape that matches both the concrete `Database` and the * `PgDatabase` widening used by sibling services (e.g. participant-mode). @@ -453,12 +453,38 @@ export async function claimBacklogForPushFair( ); } +/** One silent-context row as returned by the batched `collectPrecedingContext` + * statement. `db.execute` bypasses the driver's type parsers and hands back + * raw strings for every column, so this projection deliberately carries only + * ids — the message payload is re-read through the typed Drizzle builder. */ +type PrecedingContextRow = { + trigger_id: string; + message_id: string; +}; + +/** Per-trigger window bounds handed to SQL via `jsonb_to_recordset`. */ +type PrecedingContextBound = { + trigger_id: number; + chat_id: string; + /** Previous trigger in this same batch/chat, or `null` for the batch's first + * trigger in that chat — SQL then resolves the cursor from the table. */ + prev_trigger_id: number | null; + window_start: string; +}; + /** * Per claimed trigger: SELECT silent (notify=false) pending rows in the same * chat that occurred between the previous trigger in this batch (or beginning * of time) and this trigger, capped by `PRECEDING_CONTEXT_MAX_ENTRIES` and * `PRECEDING_CONTEXT_WINDOW_SECONDS`. Returned messages are oldest-first. * + * **Two statements for the whole batch, regardless of batch size** (PERF-061): + * one `CROSS JOIN LATERAL` over a `jsonb_to_recordset` of per-trigger bounds + * that selects + locks the silent rows, then one keyed read of the message + * payloads. Previously a 50-entry drain cost one previous-notify query per chat + * plus one query per trigger — every one of them a round trip taken inside the + * claiming transaction, holding `FOR UPDATE` locks for the whole sequence. + * * This function intentionally does not ACK silent rows. Bundling is not * consumption: recovery must be able to reset the notify trigger and rebuild * the same trigger-relative context window. Silent rows are drained only when @@ -480,97 +506,133 @@ async function collectPrecedingContext( byChat.set(t.chatId, list); } + // Within a chat the batch is its own cursor chain: trigger N's context starts + // after trigger N-1. Only the chat's FIRST trigger needs the table lookup for + // the preceding notify entry (which may have been delivered in an earlier, + // still-unacked batch) — expressed below as a COALESCE fallback so the whole + // chain resolves inside the single statement. + const bounds: PrecedingContextBound[] = []; for (const [chatId, chatTriggers] of byChat) { chatTriggers.sort((a, b) => a.id - b.id); - - const firstTrigger = chatTriggers[0]; - if (!firstTrigger) continue; - const [previousNotify] = await tx - .select({ id: inboxEntries.id }) - .from(inboxEntries) - .where( - and( - eq(inboxEntries.inboxId, inboxId), - eq(inboxEntries.chatId, chatId), - eq(inboxEntries.notify, true), - lt(inboxEntries.id, firstTrigger.id), - ), - ) - .orderBy(desc(inboxEntries.id)) - .limit(1); - - // For each trigger, fetch silent context strictly before it (and after - // the previous notify trigger cursor, even if that trigger was delivered - // in an earlier unacked batch). Window: 24h before the trigger. - // - // Order matters: when there are MORE than `PRECEDING_CONTEXT_MAX_ENTRIES` - // candidates, we want to keep the rows CLOSEST to the trigger (most - // contextually relevant) and drop the oldest. So select DESC + LIMIT, - // then reverse in JS to get chronological prompt-ready output. Selecting - // ASC + LIMIT would drop the recent rows; ACK-through later drains every - // silent row behind the consumed notify cursor, including rows excluded by - // this cap, so this delivery must choose the most relevant window now. - // - // We sort by `messages.createdAt` rather than `inboxEntries.createdAt` - // because `addParticipant`'s backfill writes 50 inbox rows in one - // `INSERT VALUES (...)` — they all share `statement_timestamp()`. The - // message rows themselves have distinct, monotonic timestamps (uuidv7 - // ids are time-ordered, `messages.created_at` is the authoritative - // chronology), so ordering by the joined message timestamp is the only - // stable contract the prompt-rendered context can rely on. - // - // Concurrency: `FOR UPDATE OF inboxEntries SKIP LOCKED` prevents two - // parallel polls on the same inbox from bundling the same silent row - // twice. Without it, poll A picking trigger T1 and poll B picking T2 - // (T2 > T1) would both include silent rows < T1 in their preceding - // context. With SKIP LOCKED, the second poll skips the rows the first - // has reserved. - let previousNotifyId: number | null = previousNotify?.id ?? null; + let prevTriggerId: number | null = null; for (const trigger of chatTriggers) { - const windowStart = new Date(trigger.createdAt.getTime() - PRECEDING_CONTEXT_WINDOW_SECONDS * 1000); - const rows = await tx - .select({ - messageId: messages.id, - senderId: messages.senderId, - format: messages.format, - content: messages.content, - metadata: messages.metadata, - source: messages.source, - createdAt: messages.createdAt, - }) - .from(inboxEntries) - .innerJoin(messages, eq(messages.id, inboxEntries.messageId)) - .where( - and( - eq(inboxEntries.inboxId, inboxId), - eq(inboxEntries.chatId, chatId), - eq(inboxEntries.status, "pending"), - eq(inboxEntries.notify, false), - lt(inboxEntries.id, trigger.id), - previousNotifyId === null ? undefined : gt(inboxEntries.id, previousNotifyId), - gt(inboxEntries.createdAt, windowStart), - ), - ) - .orderBy(desc(messages.createdAt)) - .limit(PRECEDING_CONTEXT_MAX_ENTRIES) - .for("update", { of: inboxEntries, skipLocked: true }); - - // Reverse so the prompt-rendered block reads oldest → newest. - const preceding: PrecedingMessage[] = rows - .map((r) => ({ - id: r.messageId, - senderId: r.senderId, - format: r.format, - content: r.content, - metadata: (r.metadata ?? {}) as Record, - source: messageSourceSchema.nullable().catch(null).parse(r.source), - createdAt: r.createdAt.toISOString(), - })) - .reverse(); - result.set(trigger.id, preceding); - previousNotifyId = trigger.id; + result.set(trigger.id, []); + bounds.push({ + trigger_id: trigger.id, + chat_id: chatId, + prev_trigger_id: prevTriggerId, + window_start: new Date(trigger.createdAt.getTime() - PRECEDING_CONTEXT_WINDOW_SECONDS * 1000).toISOString(), + }); + prevTriggerId = trigger.id; } } + if (bounds.length === 0) return result; + + // Order matters: when there are MORE than `PRECEDING_CONTEXT_MAX_ENTRIES` + // candidates, we want to keep the rows CLOSEST to the trigger (most + // contextually relevant) and drop the oldest. So the LATERAL selects DESC + + // LIMIT, and the outer ORDER BY flips it back to chronological prompt-ready + // output. Selecting ASC + LIMIT would drop the recent rows; ACK-through later + // drains every silent row behind the consumed notify cursor, including rows + // excluded by this cap, so this delivery must choose the most relevant window + // now. + // + // We sort by `messages.created_at` rather than `inbox_entries.created_at` + // because `addParticipant`'s backfill writes 50 inbox rows in one + // `INSERT VALUES (...)` — they all share `statement_timestamp()`. The + // message rows themselves have distinct, monotonic timestamps (uuidv7 + // ids are time-ordered, `messages.created_at` is the authoritative + // chronology), so ordering by the joined message timestamp is the only + // stable contract the prompt-rendered context can rely on. `inbox_entries.id` + // is the tie-breaker so the outer flip is an exact inverse of the inner cap. + // + // Concurrency: `FOR UPDATE OF inbox_entries SKIP LOCKED` prevents two + // parallel polls on the same inbox from bundling the same silent row + // twice. Without it, poll A picking trigger T1 and poll B picking T2 + // (T2 > T1) would both include silent rows < T1 in their preceding + // context. With SKIP LOCKED, the second poll skips the rows the first + // has reserved. The clause is legal inside a LATERAL subquery precisely + // because `CROSS JOIN` is an inner join — a `LEFT JOIN LATERAL` would be + // rejected by PostgreSQL as locking the nullable side of an outer join. + const rows = await tx.execute(sql` + WITH bound_input AS ( + SELECT t.trigger_id, t.chat_id, t.prev_trigger_id, t.window_start + FROM jsonb_to_recordset(${JSON.stringify(bounds)}::jsonb) + AS t(trigger_id bigint, chat_id text, prev_trigger_id bigint, window_start timestamptz) + ), + bounded AS ( + SELECT + b.trigger_id, + b.chat_id, + b.window_start, + COALESCE(b.prev_trigger_id, ( + SELECT max(p.id) + FROM inbox_entries p + WHERE p.inbox_id = ${inboxId} + AND p.chat_id = b.chat_id + AND p.notify = true + AND p.id < b.trigger_id + )) AS lower_id + FROM bound_input b + ) + SELECT + bounded.trigger_id::text AS trigger_id, + ctx.message_id + FROM bounded + CROSS JOIN LATERAL ( + SELECT m.id AS message_id, m.created_at, e.id AS entry_id + FROM inbox_entries e + INNER JOIN messages m ON m.id = e.message_id + WHERE e.inbox_id = ${inboxId} + AND e.chat_id = bounded.chat_id + AND e.status = 'pending' + AND e.notify = false + AND e.id < bounded.trigger_id + AND (bounded.lower_id IS NULL OR e.id > bounded.lower_id) + AND e.created_at > bounded.window_start + ORDER BY m.created_at DESC, e.id DESC + LIMIT ${PRECEDING_CONTEXT_MAX_ENTRIES} + FOR UPDATE OF e SKIP LOCKED + ) ctx + ORDER BY bounded.trigger_id ASC, ctx.created_at ASC, ctx.entry_id ASC + `); + if (rows.length === 0) return result; + + // Second and final statement: read the payloads through the typed builder so + // jsonb / timestamptz columns keep their Drizzle-mapped shapes. + const contextMessages = await tx + .select({ + id: messages.id, + senderId: messages.senderId, + format: messages.format, + content: messages.content, + metadata: messages.metadata, + source: messages.source, + createdAt: messages.createdAt, + }) + .from(messages) + // Deduped: replyTo routing can put one message under two triggers in + // different chats, and the bind list is capped at triggers × MAX_ENTRIES. + .where(inArray(messages.id, [...new Set(rows.map((r) => r.message_id))])); + const messageById = new Map(contextMessages.map((m) => [m.id, m])); + + for (const row of rows) { + const bucket = result.get(Number(row.trigger_id)); + const message = messageById.get(row.message_id); + // Unreachable in practice: trigger ids come from the seeded bounds and the + // message rows were just joined against inside the LATERAL. Skipping rather + // than throwing keeps a surprise here from failing a valid delivery. + if (!bucket || !message) continue; + bucket.push({ + id: message.id, + senderId: message.senderId, + format: message.format, + content: message.content, + metadata: message.metadata ?? {}, + source: messageSourceSchema.nullable().catch(null).parse(message.source), + createdAt: message.createdAt.toISOString(), + }); + } return result; }