perf(server): replace inbox context and route-validation N+1 queries

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<number> {
+      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);
+  });
 });

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<string, unknown> = {}): Record<string, unknown> {
     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<void> {
     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<void> {
     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(