refactor(server): isolate pending tool review expiry

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

server/src/__tests__/server-startup-feedback-export.test.ts

@@ -67,6 +67,8 @@ const {
       },
       toolGateway: {
         sweepActionReviews: vi.fn(async () => ({ scanned: 0 })),
+        sweepExpiredActionReviews: vi.fn(async () => ({ scanned: 0 })),
+        nextActionReviewExpiryAt: vi.fn(async () => null),
         cleanupExpiredSessions: vi.fn(async () => ({ deletedCount: 0 })),
         nextCredentialExpiryAt: vi.fn(async () => null),
       },

server/src/index.ts

@@ -1,3 +1,4 @@
+import { registerToolReviewExpiryWork } from "./services/tool-review-expiry-work.js";
 import { registerGatewayCredentialExpiryWork } from "./services/gateway-credential-expiry-work.js";
 import { registerSecretProposalExpiryWork } from "./services/secret-proposal-expiry-work.js";
 import { registerCustomImageCleanupWork } from "./services/custom-image-cleanup-work.js";
@@ -1436,7 +1437,15 @@ async function startServerWithDatabaseTeardown(
       }));
   };
 
-  await timeStartupRecovery("tool_reviews", () => app.locals.toolGateway.sweepActionReviews()).catch((err: unknown) => logger.error({ err }, "startup tool review recovery failed"));
+  await timeStartupRecovery("tool_reviews", async () => {
+    if (config.heartbeatSchedulerEnabled) {
+      await registerToolReviewExpiryWork(deliveryWork, app.locals.toolGateway,
+        () => !isIdleTaskDrainActive() && !heartbeatSchedulerStopped && !isWarmStandby()).ready;
+    } else {
+      await app.locals.toolGateway.sweepExpiredActionReviews();
+    }
+    return app.locals.toolGateway.sweepActionReviews({ includeExpiry: false });
+  }).catch((err: unknown) => logger.error({ err }, "startup tool review recovery failed"));
   await timeStartupRecovery("gateway_sessions", () => config.heartbeatSchedulerEnabled
     ? registerGatewayCredentialExpiryWork(deliveryWork, app.locals.toolGateway,
       () => !isIdleTaskDrainActive() && !heartbeatSchedulerStopped && !isWarmStandby()).ready
@@ -1860,7 +1869,7 @@ async function startServerWithDatabaseTeardown(
             logger.error({ err }, "periodic tool connection health sweep failed");
           }));
 
-        trackHeartbeatSchedulerWork(app.locals.toolGateway.sweepActionReviews().catch((err: unknown) => logger.error({ err }, "tool review recovery failed")));
+        trackHeartbeatSchedulerWork(app.locals.toolGateway.sweepActionReviews({ includeExpiry: false }).catch((err: unknown) => logger.error({ err }, "tool review recovery failed")));
 
         if (isIdleTaskDrainActive() || heartbeatSchedulerStopped) return;
         if (!(await heartbeat.resolveSchedulingSuppression()).suppressed) {

server/src/services/delivery-work-notifications.ts

@@ -8,6 +8,7 @@ export const DELIVERY_QUEUES = {
   executionStatus: "execution-status",
   executionDisposition: "execution-disposition",
   environmentCleanup: "environment-cleanup",
+  toolReviewExpiry: "tool-review-expiry",
   gatewayCredentialExpiry: "gateway-credential-expiry",
   secretProposalExpiry: "secret-proposal-expiry",
   customImageExpiry: "custom-image-expiry",

server/src/services/issue-thread-interactions.ts

@@ -1,3 +1,4 @@
+import { notifyToolReviewChange } from "./tool-review-work-notifications.js";
 import { normalizeEscapedLineBreaks } from "@paperclipai/shared/validators/text";
 import { notifyDeliveryWork, DELIVERY_QUEUES } from "./delivery-work-notifications.js";
 import { historicalQuestionCondition } from "./issue-question-context.js";
@@ -1083,6 +1084,7 @@ async function resolveLinkedToolActionRequests(
   },
 ) {
   if (interaction.kind !== "request_confirmation") return;
+  await notifyToolReviewChange(db);
   await db
     .update(toolActionRequests)
     .set({

server/src/services/tool-access/activity/review-queue.ts

@@ -1,3 +1,4 @@
+import { writeToolReviewChange } from "../../tool-review-work-notifications.js";
 import { and, desc, eq, inArray } from "drizzle-orm";
 import type { Db } from "@paperclipai/db";
 import { toolApplications, toolActionRequests, toolCatalogEntries, toolConnections, toolInvocations } from "@paperclipai/db";
@@ -86,7 +87,7 @@ listActionRequests: async (
           if (!readable) invalidRequestIds.push(request.id);
         }
         if (invalidRequestIds.length > 0) {
-          await db
+          await writeToolReviewChange(db, tx => tx
             .update(toolActionRequests)
             .set({
               status: "cancelled",
@@ -99,7 +100,7 @@ listActionRequests: async (
                 eq(toolActionRequests.status, "pending"),
                 inArray(toolActionRequests.id, invalidRequestIds),
               ),
-            );
+            ));
         }
         const hiddenIds = new Set([
           ...invalidRequestIds,

server/src/services/tool-action-review.ts

@@ -1,3 +1,4 @@
+import { notifyToolReviewChange } from "./tool-review-work-notifications.js";
 import { notifyDeliveryWork, DELIVERY_QUEUES } from "./delivery-work-notifications.js";
 import { and, eq } from "drizzle-orm";
 import {
@@ -140,6 +141,7 @@ export async function commitToolActionReview(
         throw conflict("This review has already been resolved");
     }
     const now = new Date();
+    await notifyToolReviewChange(tx);
     const [updated] = await tx
       .update(toolActionRequests)
       .set({

server/src/services/tool-gateway.ts

@@ -1,3 +1,5 @@
+import { toolReviewExpiryService } from "./tool-review-expiry.js";
+import { writeToolReviewChange } from "./tool-review-work-notifications.js";
 import { restConnectorFor } from "./connectors/rest-registry.js";
 import type { runPromptInjectionProvider } from "./prompt-injection-provider.js";
 import { promptInjectionPreventionService } from "./prompt-injection-prevention.js";
@@ -2507,7 +2509,7 @@ export function createToolGatewayService(
         signingSecret: options.toolActionSigningSecret,
       });
     } catch (error) {
-      await db
+      await writeToolReviewChange(db, tx => tx
         .update(toolActionRequests)
         .set({
           status: "cancelled",
@@ -2519,7 +2521,7 @@ export function createToolGatewayService(
             eq(toolActionRequests.id, actionRequest.id),
             eq(toolActionRequests.status, "pending"),
           ),
-        );
+        ));
       if (error instanceof ToolActionSigningSecretMissingError) {
         await db
           .update(toolInvocations)
@@ -2666,7 +2668,7 @@ export function createToolGatewayService(
     // an expired row. When the guard matches no row, the row is already resolved.
     // Do not emit an approval card for a dead row; throw a retriable conflict so
     // the retry finds the live request or creates a fresh, signable one.
-    const signedRows = await db
+    const signedRows = await writeToolReviewChange(db, tx => tx
       .update(toolActionRequests)
       .set({
         interactionId: interaction.id,
@@ -2684,7 +2686,7 @@ export function createToolGatewayService(
           eq(toolActionRequests.status, "pending"),
         ),
       )
-      .returning({ id: toolActionRequests.id });
+      .returning({ id: toolActionRequests.id }));
     if (signedRows.length === 0) {
       await db
         .update(toolInvocations)
@@ -8496,7 +8498,7 @@ export function createToolGatewayService(
       pendingRequest.expiresAt.getTime() <= Date.now();
     if (pendingUnsigned || pendingExpired) {
       const now = new Date();
-      await db
+      await writeToolReviewChange(db, tx => tx
         .update(toolActionRequests)
         .set({ status: "expired", resolvedAt: now, updatedAt: now })
         .where(
@@ -8504,7 +8506,7 @@ export function createToolGatewayService(
             eq(toolActionRequests.id, match.actionRequest.id),
             eq(toolActionRequests.status, "pending"),
           ),
-        );
+        ));
       await db
         .update(toolInvocations)
         .set({
@@ -8707,6 +8709,11 @@ export function createToolGatewayService(
     };
   }
 
+  const reviewExpiry = toolReviewExpiryService(db, async expired => {
+    void emitConnectionInvoked(db, expired.invocationId);
+    await reflectToolActionInteractionLifecycle(expired);
+  });
+
   return {
     async recordRuntimeMcpDeliveryDiagnostic(input: {
       companyId: string;
@@ -9505,6 +9512,7 @@ export function createToolGatewayService(
             },
           );
         }
+        const pendingActionRequest = recorded.actionRequest;
         const canonicalArguments = canonicalToolArguments(requestedParameters);
         const canonicalArgumentsHash = toolRequestIdentityHash(requestedParameters, options.toolActionSigningSecret);
         const approvalSnapshot = await connectedRemoteApprovalSnapshot(
@@ -9525,7 +9533,7 @@ export function createToolGatewayService(
             signingSecret: options.toolActionSigningSecret,
           });
         } catch (error) {
-          await db
+          await writeToolReviewChange(db, tx => tx
             .update(toolActionRequests)
             .set({
               status: "cancelled",
@@ -9534,10 +9542,10 @@ export function createToolGatewayService(
             })
             .where(
               and(
-                eq(toolActionRequests.id, recorded.actionRequest.id),
+                eq(toolActionRequests.id, pendingActionRequest.id),
                 eq(toolActionRequests.status, "pending"),
               ),
-            );
+            ));
           if (error instanceof ToolActionSigningSecretMissingError) {
             await db
               .update(toolInvocations)
@@ -9566,7 +9574,7 @@ export function createToolGatewayService(
           tool,
           argumentsSummary: argumentValidation.summary,
         });
-        await db
+        await writeToolReviewChange(db, tx => tx
           .update(toolActionRequests)
           .set({
             canonicalArgumentsHash,
@@ -9576,7 +9584,7 @@ export function createToolGatewayService(
             expiresAt: new Date(Date.now() + 60 * 60 * 1000),
             updatedAt: new Date(),
           })
-          .where(eq(toolActionRequests.id, recorded.actionRequest.id));
+          .where(eq(toolActionRequests.id, pendingActionRequest.id)));
         await writeToolCallEvent({
           invocationId,
           actionRequestId: recorded.actionRequest.id,
@@ -9719,7 +9727,10 @@ export function createToolGatewayService(
       return buildTestCallStatus(actionRequest, invocation);
     },
 
-    async sweepActionReviews() {
+    sweepExpiredActionReviews: reviewExpiry.sweep,
+    nextActionReviewExpiryAt: reviewExpiry.next,
+
+    async sweepActionReviews(input: { includeExpiry?: boolean } = {}) {
       // A process may stop between persisting a provider outcome and updating the
       // feed projection. Reconcile from authoritative rows before delivering it.
       const unreflected = await db
@@ -9757,7 +9768,7 @@ export function createToolGatewayService(
       const now = new Date();
       const staleAt = new Date(now.getTime() - 10 * 60_000);
       let cursor: string | undefined;
-      let scanned = 0;
+      let scanned = input.includeExpiry === false ? 0 : (await reviewExpiry.sweep(now)).scanned;
       for (;;) {
         const rows = await db
           .select()
@@ -9765,10 +9776,6 @@ export function createToolGatewayService(
           .where(
             and(
               or(
-                and(
-                  eq(toolActionRequests.status, "pending"),
-                  lte(toolActionRequests.expiresAt, now),
-                ),
                 eq(toolActionRequests.status, "approved"),
                 and(
                   eq(toolActionRequests.status, "executing"),
@@ -9794,15 +9801,9 @@ export function createToolGatewayService(
             );
             continue;
           }
-          const status = row.status === "pending" ? "expired" : "failed";
-          const errorCode =
-            status === "failed"
-              ? "tool_execution_outcome_unknown"
-              : "action_expired";
-          const errorMessage =
-            status === "failed"
-              ? "Execution was interrupted; the external outcome is unknown. Inspect the provider before retrying."
-              : "The approval request expired before a decision.";
+          const status = "failed";
+          const errorCode = "tool_execution_outcome_unknown";
+          const errorMessage = "Execution was interrupted; the external outcome is unknown. Inspect the provider before retrying.";
           const [changed] = await db
             .update(toolActionRequests)
             .set({ status, resolvedAt: now, updatedAt: now })
@@ -9942,7 +9943,7 @@ export function createToolGatewayService(
       }
       if (!signedPayload) {
         if (actionRequest.status === "pending") {
-          await db
+          await writeToolReviewChange(db, tx => tx
             .update(toolActionRequests)
             .set({
               status: "cancelled",
@@ -9954,7 +9955,7 @@ export function createToolGatewayService(
                 eq(toolActionRequests.id, actionRequest.id),
                 eq(toolActionRequests.status, "pending"),
               ),
-            );
+            ));
         }
         await reflectToolActionInteractionLifecycle({
           actionRequestId: actionRequest.id,
@@ -10511,7 +10512,7 @@ export function createToolGatewayService(
           actionRequest.expiresAt.getTime() <= Date.now()
         ) {
           const expiredAt = new Date();
-          const [expired] = await db
+          const [expired] = await writeToolReviewChange(db, tx => tx
             .update(toolActionRequests)
             .set({
               status: "expired",
@@ -10524,7 +10525,7 @@ export function createToolGatewayService(
                 inArray(toolActionRequests.status, ["pending", "approved"]),
               ),
             )
-            .returning({ id: toolActionRequests.id });
+            .returning({ id: toolActionRequests.id }));
           if (expired) {
             await reflectToolActionInteractionLifecycle({
               actionRequestId: expired.id,
@@ -10559,7 +10560,7 @@ export function createToolGatewayService(
             interaction?.kind === "request_confirmation" &&
             interaction.status === "accepted"
           ) {
-            const [approved] = await db
+            const [approved] = await writeToolReviewChange(db, tx => tx
               .update(toolActionRequests)
               .set({
                 status: "approved",
@@ -10577,7 +10578,7 @@ export function createToolGatewayService(
                   eq(toolActionRequests.status, "pending"),
                 ),
               )
-              .returning();
+              .returning());
             if (!approved) {
               throw new ToolGatewayHttpError(
                 409,

server/src/services/tool-review-expiry-work.ts

@@ -0,0 +1,16 @@
+import type { createDeliveryWorkCoordinator } from "./delivery-work-coordinator.js";
+import { DELIVERY_QUEUES } from "./delivery-work-notifications.js";
+import { registerDeadlineWork } from "./deadline-work.js";
+import type { createToolGatewayService } from "./tool-gateway.js";
+
+export function registerToolReviewExpiryWork(
+  coordinator: ReturnType<typeof createDeliveryWorkCoordinator>,
+  gateway: Pick<ReturnType<typeof createToolGatewayService>, "sweepExpiredActionReviews" | "nextActionReviewExpiryAt">,
+  canRun: () => boolean,
+) {
+  return registerDeadlineWork(coordinator, DELIVERY_QUEUES.toolReviewExpiry, {
+    retryMs: 30_000, canRun,
+    run: () => gateway.sweepExpiredActionReviews(),
+    next: () => gateway.nextActionReviewExpiryAt(),
+  });
+}

server/src/services/tool-review-expiry.ts

@@ -0,0 +1,50 @@
+import { and, asc, eq, gt, lte, sql } from "drizzle-orm";
+import { toolActionRequests, toolInvocations, type Db } from "@paperclipai/db";
+import { writeToolReviewChange } from "./tool-review-work-notifications.js";
+
+export function toolReviewExpiryService(db: Db, onExpired: (input: {
+  actionRequestId: string;
+  invocationId: string;
+  status: "expired";
+  errorCode: string;
+  errorMessage: string;
+}) => Promise<void>) {
+  return {
+    async next(): Promise<number | null> {
+      const [row] = await db.select({ at: sql<string | Date | null>`min(${toolActionRequests.expiresAt})` })
+        .from(toolActionRequests).where(eq(toolActionRequests.status, "pending"));
+      return row?.at == null ? null : new Date(row.at).getTime();
+    },
+    async sweep(now = new Date()) {
+      let cursor: string | undefined;
+      let scanned = 0;
+      for (;;) {
+        const rows = await db.select().from(toolActionRequests).where(and(
+          eq(toolActionRequests.status, "pending"), lte(toolActionRequests.expiresAt, now),
+          cursor ? gt(toolActionRequests.id, cursor) : undefined,
+        )).orderBy(asc(toolActionRequests.id)).limit(100);
+        for (const row of rows) {
+          const errorCode = "action_expired";
+          const errorMessage = "The approval request expired before a decision.";
+          const changed = await writeToolReviewChange(db, async tx => {
+            const [updated] = await tx.update(toolActionRequests)
+              .set({ status: "expired", resolvedAt: now, updatedAt: now }).where(and(
+                eq(toolActionRequests.id, row.id), eq(toolActionRequests.status, "pending"),
+                eq(toolActionRequests.updatedAt, row.updatedAt),
+              )).returning({ id: toolActionRequests.id });
+            if (!updated) return false;
+            await tx.update(toolInvocations).set({ status: "failed", errorCode, errorMessage,
+              completedAt: now, updatedAt: now }).where(eq(toolInvocations.id, row.invocationId));
+            return true;
+          });
+          if (changed) await onExpired({ actionRequestId: row.id, invocationId: row.invocationId,
+            status: "expired", errorCode, errorMessage });
+        }
+        scanned += rows.length;
+        if (rows.length < 100) break;
+        cursor = rows[rows.length - 1].id;
+      }
+      return { scanned };
+    },
+  };
+}

server/src/services/tool-review-work-notifications.ts

@@ -0,0 +1,13 @@
+import type { Db } from "@paperclipai/db";
+import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
+
+export function notifyToolReviewChange(transaction: object) {
+  return notifyDeliveryWork(transaction, DELIVERY_QUEUES.toolReviewExpiry);
+}
+
+export function writeToolReviewChange<T>(db: Db, write: (transaction: Db) => Promise<T>): Promise<T> {
+  return db.transaction(async transaction => {
+    await notifyToolReviewChange(transaction);
+    return write(transaction as unknown as Db);
+  });
+}