fix(server): wake chat milestones on question claim transitions

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

server/src/__tests__/chat-run-milestone-work.test.ts

@@ -597,6 +597,73 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu
     expect((await publications()).map(row => row.idempotencyKey)).toContain(`run:${targetId}:queued:${f.endpointId}`);
     await quiet();
   });
+  it.each(["retry", "stale claim", "live claim", "rolled-back reset"] as const)("rechecks question prompt binding after %s", async (mode) => {
+    const source = await run();
+    const interactions = issueThreadInteractionService(db);
+    const interaction = await interactions.create({ id: f.issueId, companyId: f.companyId }, {
+      kind: "ask_user_questions", continuationPolicy: "wake_assignee", sourceRunId: source.id,
+      title: "Choose", payload: { version: 1, title: "Choose", questions: [{
+        id: "choice", prompt: "Continue?", selectionMode: "single", required: true,
+        options: [{ id: "yes", label: "Yes" }],
+      }] },
+    }, { agentId: f.agentId, runId: source.id });
+    await interactions.answerQuestions({ id: f.issueId, companyId: f.companyId, status: "in_progress" }, interaction.id,
+      { answers: [{ questionId: "choice", optionIds: ["yes"] }] }, { userId: f.userId });
+    const other = await fixture(f.companyId);
+    await db.update(chatEndpoints).set({ assignedAgentId: f.agentId }).where(eq(chatEndpoints.id, other.endpointId));
+    await db.update(chatConversations).set({ issueId: f.issueId }).where(eq(chatConversations.id, other.conversationId));
+    await db.insert(chatPublications).values({ companyId: f.companyId, endpointId: other.endpointId, conversationId: other.conversationId,
+      issueId: f.issueId, idempotencyKey: `interaction:${interaction.id}:${other.endpointId}`,
+      payload: { text: "Continue?", interactionId: interaction.id }, state: "published", providerMessageId: "other-question",
+    });
+    const events: string[] = [];
+    let claimEvents: string[] = [];
+    let duringClaimNotices: string[] | undefined;
+    const unsubscribe = subscribeDatabaseWork(db, "chat-run-milestones", event => events.push(event));
+    const delivery = questionResponseDeliveryService(db, { heartbeat: {
+      cancelRun: async () => {
+        claimEvents = [...events];
+        await start(); await flush();
+        duringClaimNotices = (await publications()).map(row => row.idempotencyKey);
+        throw new Error("temporary cancellation failure");
+      },
+      wakeup: async () => { throw new Error("Must retain the still-active source"); },
+    } });
+    try {
+      if (mode === "retry") {
+        await expect(delivery.deliver(interaction.id)).resolves.toBeNull();
+      } else {
+        await db.update(issueQuestionResponseDeliveries).set({ status: "delivering", attemptCount: 1,
+          lastAttemptAt: new Date(Date.now() - (mode === "live claim" ? 0 : 60_000)),
+        }).where(eq(issueQuestionResponseDeliveries.interactionId, interaction.id));
+        await start(); await flush();
+        duringClaimNotices = (await publications()).map(row => row.idempotencyKey);
+        if (mode === "rolled-back reset") {
+          await expect(db.transaction(async tx => {
+            await questionResponseDeliveryService(tx as unknown as typeof db, { heartbeat: {} as never }).sweepPending(0);
+            expect(events).toEqual(["intent"]);
+            throw new Error("rollback reset");
+          })).rejects.toThrow("rollback reset");
+        } else {
+          await delivery.sweepPending(0);
+        }
+      }
+      expect(duringClaimNotices).toBeDefined();
+      expect(duringClaimNotices!.some(key => key.startsWith(`run:${source.id}:`))).toBe(false);
+      expect((await db.select().from(issueQuestionResponseDeliveries).where(eq(issueQuestionResponseDeliveries.interactionId, interaction.id)))[0].status).toBe(mode === "live claim" || mode === "rolled-back reset" ? "delivering" : "pending");
+      await flush();
+      const noticeExists = (await publications()).some(row => row.idempotencyKey === `run:${source.id}:working:${f.endpointId}`);
+      expect(noticeExists).toBe(mode === "retry" || mode === "stale claim");
+      if (mode === "retry") expect(claimEvents).toEqual(["intent", "settled"]);
+      else if (mode === "live claim") expect(events).toEqual([]);
+      else {
+        expect(events[0]).toBe("intent");
+        expect(events.slice(1)).toContain("settled");
+        expect(events.filter(event => event === "intent")).toHaveLength(1);
+      }
+      await quiet();
+    } finally { unsubscribe(); }
+  });
   it("waits for outer commit and ignores rolled-back nested writes", async () => {
     await start(); await flush();
     await expect(db.transaction(async tx => {

server/src/__tests__/question-response-delivery.test.ts

@@ -26,6 +26,7 @@ import {
   getEmbeddedPostgresTestSupport,
   startEmbeddedPostgresTestDatabase,
 } from "./helpers/embedded-postgres.js";
+import { resolveExternalTestDatabaseUrl } from "./helpers/external-test-database.js";
 import { issueThreadInteractionService } from "../services/issue-thread-interactions.js";
 import {
   buildQuestionResponseDeliveryEnvelope,
@@ -33,7 +34,8 @@ import {
   questionResponseDeliveryService,
 } from "../services/question-response-delivery.js";
 
-const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport();
+const externalDatabaseUrl = resolveExternalTestDatabaseUrl("PAPERCLIP_QUESTION_DELIVERY_TEST_DATABASE_URL");
+const embeddedPostgresSupport = externalDatabaseUrl ? { supported: true } : await getEmbeddedPostgresTestSupport();
 const describeEmbeddedPostgres = embeddedPostgresSupport.supported
   ? describe
   : describe.skip;
@@ -64,10 +66,10 @@ describeEmbeddedPostgres("question response delivery", () => {
   > | null = null;
 
   beforeAll(async () => {
-    tempDb = await startEmbeddedPostgresTestDatabase(
-      "paperclip-question-delivery-",
-    );
-    db = createDb(tempDb.connectionString);
+    if (!externalDatabaseUrl) {
+      tempDb = await startEmbeddedPostgresTestDatabase("paperclip-question-delivery-");
+    }
+    db = createDb(externalDatabaseUrl ?? tempDb!.connectionString);
   }, 20_000);
 
   afterEach(async () => {
@@ -84,6 +86,7 @@ describeEmbeddedPostgres("question response delivery", () => {
   });
 
   afterAll(async () => {
+    await db?.$client.end();
     await tempDb?.cleanup();
   });
 

server/src/services/question-response-delivery.ts

@@ -455,6 +455,9 @@ export function questionResponseDeliveryService(
         current.lastAttemptAt.getTime() > claimAt.getTime() - claimStaleMs
       )
         return null;
+      if (current.status !== "delivering") {
+        await notifyChatRunMilestones(tx, { companyId: current.companyId, issueId: current.issueId });
+      }
       return tx
         .update(issueQuestionResponseDeliveries)
         .set({
@@ -593,25 +596,28 @@ export function questionResponseDeliveryService(
     const nextErrorCount = delivery.errorCount + (options.bounded ? 1 : 0);
     const exhausted =
       options.bounded && nextErrorCount >= MAX_DELIVERY_ATTEMPTS;
-    await db
-      .update(issueQuestionResponseDeliveries)
-      .set({
-        // Keep an exhausted claim owned until recordTerminal commits its outcome.
-        status: exhausted ? "delivering" : "pending",
-        ...(options.bounded ? { errorCount: nextErrorCount } : {}),
-        lastErrorCode: errorCode,
-        updatedAt: at,
-      })
-      .where(
-        and(
-          eq(issueQuestionResponseDeliveries.id, delivery.id),
-          eq(issueQuestionResponseDeliveries.status, "delivering"),
-          eq(
-            issueQuestionResponseDeliveries.attemptCount,
-            delivery.attemptCount,
-          ),
-        ),
+    await db.transaction(async (tx) => {
+      const ownedClaim = and(
+        eq(issueQuestionResponseDeliveries.id, delivery.id),
+        eq(issueQuestionResponseDeliveries.status, "delivering"),
+        eq(issueQuestionResponseDeliveries.attemptCount, delivery.attemptCount),
       );
+      const [current] = await tx.select({ id: issueQuestionResponseDeliveries.id })
+        .from(issueQuestionResponseDeliveries).where(ownedClaim).for("update");
+      if (!current) return;
+      if (!exhausted) {
+        await notifyChatRunMilestones(tx, { companyId: delivery.companyId, issueId: delivery.issueId });
+      }
+      await tx.update(issueQuestionResponseDeliveries)
+        .set({
+          // Keep an exhausted claim owned until recordTerminal commits its outcome.
+          status: exhausted ? "delivering" : "pending",
+          ...(options.bounded ? { errorCount: nextErrorCount } : {}),
+          lastErrorCode: errorCode,
+          updatedAt: at,
+        })
+        .where(ownedClaim);
+    });
     return exhausted;
   }
 
@@ -1303,21 +1309,32 @@ export function questionResponseDeliveryService(
   async function sweepPending(limit = 50) {
     const sweepAt = now();
     const staleAt = new Date(sweepAt.getTime() - claimStaleMs);
-    await db
-      .update(issueQuestionResponseDeliveries)
-      .set({
-        status: "pending",
-        updatedAt: sweepAt,
-      })
-      .where(
-        and(
-          eq(issueQuestionResponseDeliveries.status, "delivering"),
-          or(
-            isNull(issueQuestionResponseDeliveries.lastAttemptAt),
-            lte(issueQuestionResponseDeliveries.lastAttemptAt, staleAt),
-          ),
+    await db.transaction(async (tx) => {
+      const staleClaims = await tx.select({
+        id: issueQuestionResponseDeliveries.id,
+        companyId: issueQuestionResponseDeliveries.companyId,
+        issueId: issueQuestionResponseDeliveries.issueId,
+      }).from(issueQuestionResponseDeliveries).where(and(
+        eq(issueQuestionResponseDeliveries.status, "delivering"),
+        or(
+          isNull(issueQuestionResponseDeliveries.lastAttemptAt),
+          lte(issueQuestionResponseDeliveries.lastAttemptAt, staleAt),
         ),
-      );
+      )).orderBy(asc(issueQuestionResponseDeliveries.id)).for("update");
+      if (staleClaims.length === 0) return;
+      const notified = new Set<string>();
+      for (const claim of staleClaims) {
+        const scope = `${claim.companyId}:${claim.issueId}`;
+        if (notified.has(scope)) continue;
+        await notifyChatRunMilestones(tx, claim);
+        notified.add(scope);
+      }
+      for (let offset = 0; offset < staleClaims.length; offset += 500) {
+        await tx.update(issueQuestionResponseDeliveries)
+          .set({ status: "pending", updatedAt: sweepAt })
+          .where(inArray(issueQuestionResponseDeliveries.id, staleClaims.slice(offset, offset + 500).map(claim => claim.id)));
+      }
+    });
     const ids = await db
       .select({ interactionId: issueQuestionResponseDeliveries.interactionId })
       .from(issueQuestionResponseDeliveries)