fix(server): stop continuation publication sends at admission boundaries

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

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

@@ -45158,6 +45158,64 @@ describeEmbeddedPostgres("chat channel control-plane integration", () => {
     }
   });
 
+  it.each(["drain", "standby", "shutdown"] as const)("stops starting continuation publications during %s", async mode => {
+    const context = await telegramConfirmationRecoveryFixture();
+    await expect(context.onAction(context.actionEvent)).rejects.toThrow("injected post-interaction confirmation settlement crash");
+    const next = await issueThreadInteractionService(db).create(
+      { id: context.conversation.issueId, companyId: context.fixture.companyId },
+      { kind: "request_confirmation", continuationPolicy: "wake_assignee", title: "Next decision",
+        payload: { version: 1, prompt: "A second decision after the accepted response?" } },
+      { agentId: context.fixture.assignedAgentId },
+    );
+    const nextPublication = () => db.select({ state: chatPublications.state }).from(chatPublications).where(
+      eq(chatPublications.idempotencyKey, `interaction:${next.id}:${context.endpoint.id}`));
+    expect(await nextPublication()).toEqual([{ state: "pending" }]);
+    const provider = (context.service.runtime as unknown as FakeChatSdkRuntime).endpoints.get(context.endpoint.id)!;
+    const originalThread = provider.thread.bind(provider);
+    let entered = false;
+    let release!: () => void;
+    const gate = new Promise<void>(resolve => { release = resolve; });
+    const threadSpy = vi.spyOn(provider, "thread").mockImplementation(threadId => {
+      const thread = originalThread(threadId), edit = thread.adapter.editMessage.bind(thread.adapter);
+      thread.adapter.editMessage = async (...args) => { entered = true; await gate; return edit(...args); };
+      return thread;
+    });
+    let canRun = true;
+    const errors: unknown[] = [];
+    let coordinator = createDeliveryWorkCoordinator({ owner: db, canRun: () => canRun,
+      canReconcile: () => true, onError: error => errors.push(error) });
+    vi.useFakeTimers({ toFake: ["Date", "setTimeout", "clearTimeout"], shouldAdvanceTime: true });
+    const ready = context.service.registerInteractionContinuationWork(coordinator, () => canRun).ready;
+    let stopping: Promise<void> | undefined;
+    try {
+      await vi.waitFor(() => expect(entered).toBe(true), { timeout: 10_000 });
+      if (mode === "shutdown") stopping = coordinator.stop(); else canRun = false;
+      release(); await ready; await stopping;
+      await vi.waitFor(() => expect(idleWorkSnapshot().active).toBe(0));
+      expect(provider.edits).toHaveLength(1);
+      expect(await nextPublication()).toEqual([{ state: "pending" }]);
+      expect(context.wakeup).not.toHaveBeenCalled();
+      if (mode === "shutdown") {
+        coordinator = createDeliveryWorkCoordinator({ owner: db, canRun: () => canRun,
+          canReconcile: () => true, onError: error => errors.push(error) });
+        await context.service.registerInteractionContinuationWork(coordinator, () => canRun).ready;
+      } else {
+        canRun = true;
+        await vi.advanceTimersByTimeAsync(Math.max(0, coordinator.nextWakeAt()! - Date.now()));
+      }
+      await vi.waitFor(async () => {
+        await vi.advanceTimersByTimeAsync(10);
+        expect(context.wakeup).toHaveBeenCalledTimes(1);
+        expect(await nextPublication()).toEqual([{ state: "published" }]);
+        expect(idleWorkSnapshot().active).toBe(0); expect(coordinator.nextWakeAt()).toBeNull();
+      }, { timeout: 10_000 });
+      expect(errors).toEqual([]);
+    } finally {
+      release(); await ready; await stopping; await coordinator.stop(); vi.useRealTimers();
+      threadSpy.mockRestore(); await context.service.shutdown();
+    }
+  });
+
   it("defers continuation dispatch until the outer response transaction commits", async () => {
     const context = await telegramConfirmationRecoveryFixture();
     const errors: unknown[] = [];

server/src/__tests__/chat-interaction-publications.test.ts

@@ -27,7 +27,13 @@ import {
   startEmbeddedPostgresTestDatabase,
 } from "./helpers/embedded-postgres.js";
 
-const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport();
+const externalTestDatabaseUrl = process.env.PAPERCLIP_TEST_DATABASE_URL;
+if (externalTestDatabaseUrl && !new URL(externalTestDatabaseUrl).pathname.endsWith("_test")) {
+  throw new Error("PAPERCLIP_TEST_DATABASE_URL must name a disposable _test database");
+}
+const embeddedPostgresSupport = externalTestDatabaseUrl
+  ? { supported: true }
+  : await getEmbeddedPostgresTestSupport();
 const describeEmbeddedPostgres = embeddedPostgresSupport.supported
   ? describe
   : describe.skip;
@@ -43,16 +49,18 @@ describeEmbeddedPostgres(
 
     beforeAll(async () => {
       process.env.PAPERCLIP_PUBLIC_URL = "https://paperclip.example";
-      tempDb = await startEmbeddedPostgresTestDatabase(
-        "paperclip-terminal-chat-interaction-",
-      );
-      db = createDb(tempDb.connectionString);
+      if (externalTestDatabaseUrl) db = createDb(externalTestDatabaseUrl);
+      else {
+        tempDb = await startEmbeddedPostgresTestDatabase("paperclip-terminal-chat-interaction-");
+        db = createDb(tempDb.connectionString);
+      }
     }, 20_000);
 
     afterAll(async () => {
       if (previousPublicUrl === undefined)
         delete process.env.PAPERCLIP_PUBLIC_URL;
       else process.env.PAPERCLIP_PUBLIC_URL = previousPublicUrl;
+      await db?.$client.end({ timeout: 1 });
       await tempDb?.cleanup();
     });
 

server/src/services/chat-channels.ts

@@ -3479,7 +3479,7 @@ export function chatChannelService(db: Db, options: ChatChannelServiceOptions) {
   });
   const { processSelectedPublication } = publicationWorker;
   const { processPendingInteractionWakeups, registerInteractionContinuationWork } = createChatInteractionContinuations(db, {
-    preparePublications: scope => processPendingPublications(25, { maintenance: false, scope }),
+    preparePublications: (scope, canContinue) => processPendingPublications(25, { maintenance: false, scope, canContinue }),
     redactError, CONFIRMATION_WAKEUP_RETRY_BACKOFF_MS, get heartbeat() { return options.heartbeat; },
   });
   const { handleAction, handleModalSubmit, reconcileTerminalConfirmationActions, registerTerminalConfirmationWork } = createChatInteractionCallbacks(db, {

server/src/services/chat-channels/interactions/continuations.ts

@@ -39,7 +39,7 @@ import { issueThreadInteractionService } from "../../issue-thread-interactions.j
 
 type InteractionContinuationsOptions = {
   heartbeat: IssueAssignmentWakeupDeps;
-  preparePublications: (scope: { companyId: string; endpointId: string; conversationId: string }) => Promise<unknown>;
+  preparePublications: (scope: { companyId: string; endpointId: string; conversationId: string }, canContinue: () => boolean) => Promise<unknown>;
   redactError: (error: unknown) => string;
   CONFIRMATION_WAKEUP_RETRY_BACKOFF_MS: number;
 };
@@ -86,7 +86,7 @@ export function createChatInteractionContinuations(db: Db, options: InteractionC
     for (const candidate of candidates) {
       if (preparePublications && candidate.conversationId) {
         await options.preparePublications({ companyId: candidate.companyId,
-          endpointId: candidate.endpointId, conversationId: candidate.conversationId });
+          endpointId: candidate.endpointId, conversationId: candidate.conversationId }, canContinue);
       }
       if (!canContinue()) break;
       const writeAction = <T>(write: (tx: Parameters<Parameters<Db["transaction"]>[0]>[0]) => PromiseLike<T>) =>

server/src/services/chat-channels/publications/queue.ts

@@ -181,16 +181,17 @@ export function createChatPublicationQueue(db: Db, options: PublicationQueueOpti
 
   async function processPendingPublications(
     limit = 25,
-    { waitForCompletion = true, maintenance = true, scope }: {
-      waitForCompletion?: boolean; maintenance?: boolean;
+    { waitForCompletion = true, maintenance = true, scope, canContinue = () => true }: {
+      waitForCompletion?: boolean; maintenance?: boolean; canContinue?: () => boolean;
       scope?: { companyId: string; endpointId: string; conversationId: string };
     } = {},
   ) {
-    if (isShuttingDown()) return 0;
+    if (isShuttingDown() || !canContinue()) return 0;
     if (maintenance) {
       await enqueueInboundWakeupPublications();
       await reconcileTerminalConfirmationActions(limit);
     }
+    if (isShuttingDown() || !canContinue()) return 0;
     const scopeCondition = scope ? and(eq(chatPublications.companyId, scope.companyId),
       eq(chatPublications.endpointId, scope.endpointId), eq(chatPublications.conversationId, scope.conversationId)) : undefined;
     const pendingTasks = () => scope
@@ -287,7 +288,7 @@ export function createChatPublicationQueue(db: Db, options: PublicationQueueOpti
     let failed = false;
     let firstError: unknown;
     try {
-      while (attemptedIds.length < limit && !failed && !isShuttingDown()) {
+      while (attemptedIds.length < limit && !failed && !isShuttingDown() && canContinue()) {
         if (publicationEndpointTasks.size >= PUBLICATION_ENDPOINT_CONCURRENCY) {
           if (!waitForCompletion) break;
           await Promise.race(publicationEndpointTasks.values());
@@ -328,7 +329,7 @@ export function createChatPublicationQueue(db: Db, options: PublicationQueueOpti
             publicationEndpointTasks.size >= PUBLICATION_ENDPOINT_CONCURRENCY ||
             attemptedIds.length >= limit ||
             failed ||
-            isShuttingDown()
+            (isShuttingDown() || !canContinue())
           )
             break;
           if (publicationEndpointTasks.has(selectedPublication.endpointId))
@@ -340,6 +341,7 @@ export function createChatPublicationQueue(db: Db, options: PublicationQueueOpti
           // endpoint-exclusive credential fence instead of timing out behind it.
           const task = Promise.resolve()
             .then(async () => {
+              if (isShuttingDown() || !canContinue()) return;
               await processSelectedPublication(selectedPublication);
               await settleExternalConversation(db, selectedPublication.companyId, selectedPublication.issueId).catch((err) => {
                 logger.warn({ err, publicationId: selectedPublication.id }, "Slack conversation settlement deferred to reconciliation");