fix: scope native response wakes to affected runs and tasks

Co-Authored-By: Paperclip <noreply@paperclip.ing>

server/src/__tests__/chat-channels.integration.test.ts

@@ -72675,6 +72675,34 @@ describeEmbeddedPostgres("chat channel control-plane integration", () => {
     } finally { worker.stop(); await coordinator.stop(); await replacement?.service.shutdown(); }
   });
 
+  it.each(["issue", "run"])("does not reload dormant accepted-response history for an unrelated %s wake", async scope => {
+    const context = await committedChatResponseRecoveryFixture("telegram");
+    await context.service.shutdown();
+    const errors: unknown[] = [];
+    const coordinator = createDeliveryWorkCoordinator({ owner: db, canRun: () => true, onError: error => errors.push(error) });
+    const work = nativeResponsePresentationWork(db);
+    const worker = work.register(coordinator, () => true);
+    try {
+      await worker.ready;
+      await expect.poll(() => coordinator.nextWakeAt(), { timeout: 10_000 }).toBeNull();
+      expect(await work.load({ companyId: context.fixture.companyId, runId: context.runId })).toHaveLength(1);
+      const previousDebug = db.$client.options.debug;
+      const queries: string[] = [];
+      db.$client.options.debug = (_connection, query) => {
+        if (query.startsWith("select") && query.includes("native_run_finalizations")) queries.push(query);
+      };
+      try {
+        await db.transaction(tx => notifyChatRunMilestones(tx, { companyId: context.fixture.companyId,
+          ...(scope === "issue" ? { issueId: randomUUID() } : { runId: randomUUID() }) }));
+        await expect.poll(() => ({ active: idleWorkSnapshot().active, next: coordinator.nextWakeAt() }),
+          { timeout: 5000 }).toEqual({ active: 0, next: null });
+        expect(queries).toHaveLength(1);
+      } finally { db.$client.options.debug = previousDebug; }
+      expect(await readChatPublicationsByEndpointId(context.endpoint.id)).toHaveLength(0);
+      expect(errors).toEqual([]);
+    } finally { worker.stop(); await coordinator.stop(); }
+  });
+
   it.each(["commit", "rollback"])("keeps accepted-response work behind the outer %s boundary", async outcome => {
     const context = await committedChatResponseRecoveryFixture("telegram");
     await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, context.runId));
@@ -72686,7 +72714,7 @@ describeEmbeddedPostgres("chat channel control-plane integration", () => {
       await expect.poll(() => coordinator.nextWakeAt(), { timeout: 5000 }).toBeNull();
       const transaction = db.transaction(async tx => {
         await tx.transaction(async nested => {
-          await notifyChatRunMilestones(nested, { companyId: context.fixture.companyId, issueId: context.issue.id });
+          await notifyChatRunMilestones(nested, { companyId: context.fixture.companyId, issueId: context.issue.id, runId: randomUUID() });
           await nested.update(heartbeatRuns).set({ status: "failed" }).where(eq(heartbeatRuns.id, context.runId));
         });
         expect(await readChatPublicationsByEndpointId(context.endpoint.id)).toHaveLength(0);
@@ -72886,7 +72914,7 @@ describeEmbeddedPostgres("chat channel control-plane integration", () => {
     const worker = nativeResponsePresentationWork(db).register(coordinator, () => true);
     try {
       await worker.ready;
-      expect(coordinator.nextWakeAt()).toBeNull();
+      await expect.poll(() => coordinator.nextWakeAt(), { timeout: 10_000 }).toBeNull();
       expect(await readChatPublicationsByEndpointId(context.endpoint.id)).toHaveLength(0);
       const port = new PaperclipControlPlanePort(db, {
         companyId: context.binding.companyId, issueId: context.binding.issueId,

server/src/services/chat-run-milestone-notifications.ts

@@ -6,10 +6,11 @@ import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notificatio
 
 /** Lineage changes can authorize descendants, so invalidate the issue rather
  * than just the run whose row changed. Configuration may invalidate a company. */
-export type ChatRunMilestoneScope = { companyId: string; issueId?: string };
+export type ChatRunMilestoneScope = { companyId: string; issueId?: string; runId?: string };
 
 export async function notifyChatRunMilestones(tx: object, scope: ChatRunMilestoneScope) {
-  await notifyNativeResponsePresentation(tx, scope);
+  await notifyNativeResponsePresentation(tx, { companyId: scope.companyId, issueId: scope.issueId,
+    runId: scope.issueId ? undefined : scope.runId });
   await notifyDeliveryWork(tx, DELIVERY_QUEUES.chatGitHubReviewChecks, JSON.stringify({ companyId: scope.companyId }));
   await notifyDeliveryWork(tx, DELIVERY_QUEUES.chatGitHubTerminalResponses, JSON.stringify({ companyId: scope.companyId, issueId: scope.issueId }));
   await notifyDeliveryWork(tx, DELIVERY_QUEUES.chatRunMilestones, JSON.stringify({ companyId: scope.companyId, issueId: scope.issueId }));

server/src/services/heartbeat/queue.ts

@@ -974,7 +974,8 @@ export function createHeartbeatQueue(db: Db, dependencies: HeartbeatQueueDepende
             .where(and(eq(heartbeatRuns.id, run.id), eq(heartbeatRuns.companyId, run.companyId), eq(heartbeatRuns.status, "queued")))
             .for("update");
           if (!currentRun) return null;
-          await notifyChatRunMilestones(tx, { companyId: run.companyId });
+          await notifyChatRunMilestones(tx, { companyId: run.companyId, runId: run.id,
+            issueId: run.nativeIssueId ?? run.issueId ?? readNonEmptyString(parseObject(run.contextSnapshot).issueId) ?? undefined });
           await notifyInstructionCollection(tx, { companyId: run.companyId, runId: run.id });
           await notifyCancelledWorkspacePreservation(tx, { companyId: run.companyId, runId: run.id });
           await notifyUnsafeWorkspaceExportRunChange(tx, run);
@@ -1623,7 +1624,7 @@ export function createHeartbeatQueue(db: Db, dependencies: HeartbeatQueueDepende
             )).for("update");
             if (!lockedRun || lockedRun.status !== "queued") return null;
             await withExternalAdmissionGuard(claimTx, run.companyId, run.agentId, () => assertNoExternalOverlap(claimTx, run.companyId, run.agentId));
-            await notifyChatRunMilestones(claimTx, { companyId: run.companyId, issueId: issueId ?? undefined });
+            await notifyChatRunMilestones(claimTx, { companyId: run.companyId, runId: run.id, issueId: issueId ?? undefined });
             const claimedRun = await claimTx.update(heartbeatRuns).set(claimValues).where(and(
               eq(heartbeatRuns.id, run.id), eq(heartbeatRuns.status, "queued"),
             )).returning().then((rows) => rows[0] ?? null);
@@ -1724,7 +1725,7 @@ export function createHeartbeatQueue(db: Db, dependencies: HeartbeatQueueDepende
       const [candidate] = await tx.select({ companyId: heartbeatRuns.companyId, contextSnapshot: heartbeatRuns.contextSnapshot })
         .from(heartbeatRuns).where(and(eq(heartbeatRuns.id, runId), eq(heartbeatRuns.status, "running"))).for("update");
       if (!candidate) return;
-      await notifyChatRunMilestones(tx, { companyId: candidate.companyId,
+      await notifyChatRunMilestones(tx, { companyId: candidate.companyId, runId,
         issueId: readNonEmptyString(parseObject(candidate.contextSnapshot).issueId) ?? undefined });
       const released = await tx
         .update(heartbeatRuns)

server/src/services/heartbeat/run-completion.ts

@@ -891,7 +891,7 @@ export function createHeartbeatRunCompletion(db: Db, dependencies: HeartbeatRunC
         }
 
         await db.transaction(async tx => {
-          await notifyChatRunMilestones(tx, { companyId: livenessRun.companyId, issueId: issueId ?? undefined });
+          await notifyChatRunMilestones(tx, { companyId: livenessRun.companyId, runId: livenessRun.id, issueId: issueId ?? undefined });
           await tx.update(heartbeatRuns).set({
             resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) || ${JSON.stringify({ presentationDecision })}::jsonb`,
             updatedAt: new Date(),

server/src/services/heartbeat/run-lifecycle.test.ts

@@ -1,7 +1,7 @@
 import { resolveExternalTestDatabaseUrl } from "../../__tests__/helpers/external-test-database.js";
 import { randomBytes } from "node:crypto";
 import type { Db } from "@paperclipai/db";
-import { agents, agentWakeupRequests, companies, createDb, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, heartbeatRunEvents, heartbeatRuns } from "@paperclipai/db";
+import { agents, agentWakeupRequests, companies, createDb, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, heartbeatRunEvents, heartbeatRuns, issues, subscribeDatabaseWork } from "@paperclipai/db";
 import { eq, sql } from "drizzle-orm";
 import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
 import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../../__tests__/helpers/embedded-postgres.js";
@@ -241,6 +241,23 @@ describe.skipIf(!support.supported)("heartbeat lifecycle database wiring", () =>
     } finally { unsubscribe(); }
   });
 
+  it.each(["native", "legacy", "unbound"] as const)("keeps %s run-status response recovery scoped to its task or run", async binding => {
+    const { run } = await fixture();
+    const [issue] = await db.insert(issues).values({ companyId: run.companyId, title: "Scope owner" }).returning();
+    await db.update(heartbeatRuns).set(binding === "native" ? { nativeIssueId: issue.id }
+      : binding === "legacy" ? { contextSnapshot: { issueId: issue.id } } : { contextSnapshot: {} }).where(eq(heartbeatRuns.id, run.id));
+    const hints: unknown[] = [];
+    const unsubscribe = subscribeDatabaseWork(db, DELIVERY_QUEUES.nativeResponsePresentation, (event, scopes) => {
+      if (event === "settled") for (const hint of scopes ?? []) hints.push(JSON.parse(hint));
+    });
+    try {
+      await createHeartbeatLifecycle(db, callbacks(db)).setRunStatusIfRunning(run.id, "succeeded", { finishedAt: new Date() });
+      expect(hints).toEqual([binding === "unbound" ? { companyId: run.companyId, runId: run.id }
+        : { companyId: run.companyId, issueId: issue.id }]);
+      expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, run.id)))[0].status).toBe("succeeded");
+    } finally { unsubscribe(); }
+  });
+
   it("does not request accounting for an unchanged running status", async () => {
     const { run } = await fixture();
     const settled = vi.fn();

server/src/services/heartbeat/run-lifecycle.ts

@@ -1016,7 +1016,7 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
     return recordHeartbeatRunRuntimeProgress(currentRun, update, issueId);
   }
 
-  async function writeRunStatus<T>(status: string, previous: Pick<typeof heartbeatRuns.$inferSelect, "id" | "companyId" | "status" | "nativeIssueId" | "contextSnapshot"> | null,
+  async function writeRunStatus<T>(status: string, previous: Pick<typeof heartbeatRuns.$inferSelect, "id" | "companyId" | "status" | "nativeIssueId" | "issueId" | "contextSnapshot"> | null,
     write: (tx: Db) => Promise<T>) {
     const companyId = previous?.companyId;
     // Register on the outer transaction before the terminal write. An unknown
@@ -1030,7 +1030,8 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
       if (previous) await notifyCancelledWorkspacePreservation(tx, { companyId: previous.companyId, runId: previous.id });
       if (isHeartbeatRunTerminalStatus(status)) await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, companyId);
       await notifyExecutionRunWork(tx);
-      if (companyId) await notifyChatRunMilestones(tx, { companyId });
+      if (previous) await notifyChatRunMilestones(tx, { companyId: previous.companyId, runId: previous.id,
+        issueId: previous.nativeIssueId ?? previous.issueId ?? readNonEmptyString(parseObject(previous.contextSnapshot).issueId) ?? undefined });
       return write(tx);
     });
   }