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; }