From fd2d848a7ba5e9a23e8bf079a5ca81799933408c Mon Sep 17 00:00:00 2001 From: Dotta Date: Sat, 24 Oct 2026 10:44:45 -0700 Subject: [PATCH] fix(server): schedule interrupted tool claims by their deadlines Co-Authored-By: Paperclip --- packages/db/src/client.ts | 1 + .../src/migrations/0381_open_madame_web.sql | 1 + ...{0376_snapshot.json => 0381_snapshot.json} | 79 ++++++- packages/db/src/migrations/meta/_journal.json | 7 + packages/db/src/schema/tool_access.ts | 1 + .../server-startup-feedback-export.test.ts | 2 + .../__tests__/tool-gateway-service.test.ts | 211 +++++++++++++++--- server/src/index.ts | 11 +- .../services/delivery-work-notifications.ts | 1 + server/src/services/tool-gateway.ts | 88 ++------ .../services/tool-review-interruption-work.ts | 16 ++ .../src/services/tool-review-interruption.ts | 61 +++++ .../tool-review-work-notifications.ts | 1 + 13 files changed, 381 insertions(+), 99 deletions(-) create mode 100644 packages/db/src/migrations/0381_open_madame_web.sql rename packages/db/src/migrations/meta/{0376_snapshot.json => 0381_snapshot.json} (99%) create mode 100644 server/src/services/tool-review-interruption-work.ts create mode 100644 server/src/services/tool-review-interruption.ts diff --git a/packages/db/src/client.ts b/packages/db/src/client.ts index 5dac4f14540..30c09ed5550 100644 --- a/packages/db/src/client.ts +++ b/packages/db/src/client.ts @@ -506,6 +506,7 @@ const MIGRATIONS_WITH_TOP_LEVEL_STATEMENTS = new Set([ "0369_late_logan.sql", "0378_mean_scream.sql", "0379_narrow_the_executioner.sql", + "0381_open_madame_web.sql", "0380_funny_adam_warlock.sql", ]); diff --git a/packages/db/src/migrations/0381_open_madame_web.sql b/packages/db/src/migrations/0381_open_madame_web.sql new file mode 100644 index 00000000000..19ce52b5cdb --- /dev/null +++ b/packages/db/src/migrations/0381_open_madame_web.sql @@ -0,0 +1 @@ +CREATE INDEX CONCURRENTLY IF NOT EXISTS "tool_action_requests_executing_deadline_idx" ON "tool_action_requests" USING btree ("updated_at","id") WHERE "tool_action_requests"."status" = 'executing'; \ No newline at end of file diff --git a/packages/db/src/migrations/meta/0376_snapshot.json b/packages/db/src/migrations/meta/0381_snapshot.json similarity index 99% rename from packages/db/src/migrations/meta/0376_snapshot.json rename to packages/db/src/migrations/meta/0381_snapshot.json index 1015edaf69f..5e15d7e4c76 100644 --- a/packages/db/src/migrations/meta/0376_snapshot.json +++ b/packages/db/src/migrations/meta/0381_snapshot.json @@ -1,6 +1,6 @@ { - "id": "7832df38-fa1b-4687-bd06-eccbc29bd28c", - "prevId": "3e17203b-ec95-41e9-8363-685982450199", + "id": "afdc5e70-19b3-416d-917a-beb9fdc1e2e1", + "prevId": "e78284d4-7983-44c9-a285-0b68f22c3326", "version": "7", "dialect": "postgresql", "tables": { @@ -52787,6 +52787,66 @@ "method": "btree", "with": {} }, + "tool_action_requests_pending_expiry_idx": { + "name": "tool_action_requests_pending_expiry_idx", + "columns": [ + { + "expression": "expires_at", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"tool_action_requests\".\"status\" = 'pending'", + "concurrently": false, + "method": "btree", + "with": {} + }, + "tool_action_requests_executing_deadline_idx": { + "name": "tool_action_requests_executing_deadline_idx", + "columns": [ + { + "expression": "updated_at", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"tool_action_requests\".\"status\" = 'executing'", + "concurrently": false, + "method": "btree", + "with": {} + }, + "tool_action_requests_approved_idx": { + "name": "tool_action_requests_approved_idx", + "columns": [ + { + "expression": "id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"tool_action_requests\".\"status\" = 'approved'", + "concurrently": false, + "method": "btree", + "with": {} + }, "tool_action_requests_invocation_idx": { "name": "tool_action_requests_invocation_idx", "columns": [ @@ -55197,6 +55257,21 @@ "method": "btree", "with": {} }, + "tool_gateway_sessions_expiry_idx": { + "name": "tool_gateway_sessions_expiry_idx", + "columns": [ + { + "expression": "expires_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, "tool_gateway_sessions_run_idx": { "name": "tool_gateway_sessions_run_idx", "columns": [ diff --git a/packages/db/src/migrations/meta/_journal.json b/packages/db/src/migrations/meta/_journal.json index f4191cc18e2..9b6437c673a 100644 --- a/packages/db/src/migrations/meta/_journal.json +++ b/packages/db/src/migrations/meta/_journal.json @@ -2654,6 +2654,13 @@ "when": 1792863024087, "tag": "0380_funny_adam_warlock", "breakpoints": true + }, + { + "idx": 381, + "version": "7", + "when": 1792863626931, + "tag": "0381_open_madame_web", + "breakpoints": true } ] } \ No newline at end of file diff --git a/packages/db/src/schema/tool_access.ts b/packages/db/src/schema/tool_access.ts index f6b87c17913..1b7d2484b49 100644 --- a/packages/db/src/schema/tool_access.ts +++ b/packages/db/src/schema/tool_access.ts @@ -830,6 +830,7 @@ export const toolActionRequests = pgTable( index("tool_action_requests_company_status_idx").on(table.companyId, table.status), index("tool_action_requests_pending_expiry_idx").on(table.expiresAt, table.id) .where(sql`${table.status} = 'pending'`), + index("tool_action_requests_executing_deadline_idx").on(table.updatedAt, table.id).where(sql`${table.status} = 'executing'`), index("tool_action_requests_approved_idx").on(table.id).where(sql`${table.status} = 'approved'`), index("tool_action_requests_invocation_idx").on(table.invocationId), index("tool_action_requests_issue_idx").on(table.companyId, table.issueId), diff --git a/server/src/__tests__/server-startup-feedback-export.test.ts b/server/src/__tests__/server-startup-feedback-export.test.ts index 790fd64ab54..c055bf3cb7f 100644 --- a/server/src/__tests__/server-startup-feedback-export.test.ts +++ b/server/src/__tests__/server-startup-feedback-export.test.ts @@ -66,6 +66,8 @@ const { stop: vi.fn(async () => {}), }, toolGateway: { + settleInterruptedActionReviews: vi.fn(async () => ({ scanned: 0 })), + nextActionReviewInterruptionAt: vi.fn(async () => null), recoverApprovedActionReviews: vi.fn(async () => ({ scanned: 0, retry: false })), sweepActionReviews: vi.fn(async () => ({ scanned: 0 })), sweepExpiredActionReviews: vi.fn(async () => ({ scanned: 0 })), diff --git a/server/src/__tests__/tool-gateway-service.test.ts b/server/src/__tests__/tool-gateway-service.test.ts index 166f8c3dbaa..c8eb1e1bf0b 100644 --- a/server/src/__tests__/tool-gateway-service.test.ts +++ b/server/src/__tests__/tool-gateway-service.test.ts @@ -1,3 +1,5 @@ +import { registerToolReviewInterruptionWork } from "../services/tool-review-interruption-work.js"; +import { writeToolReviewChange } from "../services/tool-review-work-notifications.js"; import { subscribeDeliveryWork } from "../services/delivery-work-notifications.js"; import { DELIVERY_QUEUES } from "../services/delivery-work-notifications.js"; import { randomUUID } from "node:crypto"; @@ -310,6 +312,186 @@ describeEmbeddedPostgres("tool gateway service", () => { expect(await db.select().from(toolCallEvents).where(eq(toolCallEvents.reasonCode, "approved_action_executed"))).toEqual([]); } + async function testRequest(provider: NonNullable) { + await db.insert(companyMemberships).values({ companyId: fixture.company.id, principalType: "user", + principalId: "board", status: "active", membershipRole: "owner" }); + const remote = await createRemoteMcpToolFixture(db, fixture.company.id); + await db.update(toolCatalogEntries).set({ isReadOnly: false, riskLevel: "write" }) + .where(eq(toolCatalogEntries.id, remote.catalogEntry.id)); + await db.insert(toolPolicies).values({ companyId: fixture.company.id, name: "Review connected writes", + policyType: "require_approval", priority: 10, selectors: { connectionId: remote.connection.id } }); + gateway = createTestToolGatewayService(db, { remoteHttpRequest: provider }); + const result = await gateway.executeTestCall({ companyId: fixture.company.id, connectionId: remote.connection.id, + agentId: fixture.agent.id, userId: "board", toolName: "needs_input", parameters: { value: "reviewed" } }); + expect(result.decision).toBe("ask_first"); + const [pending] = await db.select().from(toolActionRequests); + return { pending, remote }; + } + function providerFixture() { + let release!: () => void, entered!: () => void; + const gate = new Promise(resolve => { release = resolve; }); + const started = new Promise(resolve => { entered = resolve; }); + const provider = vi.fn(async (_url, init) => { + const payload = JSON.parse(String(init.body)); + expect(payload.method).toBe("tools/call"); + entered(); await gate; + return Response.json({ jsonrpc: "2.0", id: payload.id, + result: { content: [{ type: "text", text: "Saved once" }] } }); + }); + return { provider, release, started }; + } + + describe("interrupted review recovery worker", () => { + async function startInterruption() { + coordinator = createDeliveryWorkCoordinator({ owner: db, canRun: () => !draining, + canReconcile: () => true, onError: error => errors.push(error) }); + await registerToolReviewInterruptionWork(coordinator, gateway, () => !draining).ready; + } + async function savedClaim(source = db, age = 0) { + const pending = await request(); + const claimedAt = new Date(Date.now() - age); + await writeToolReviewChange(source, async tx => { + await tx.update(toolInvocations).set({ status: "executing", startedAt: claimedAt, updatedAt: claimedAt }) + .where(eq(toolInvocations.id, pending.invocationId)); + await tx.update(toolActionRequests).set({ status: "executing", updatedAt: claimedAt }) + .where(eq(toolActionRequests.id, pending.id)); + }); + return { ...pending, updatedAt: claimedAt }; + } + async function unknown(pending: Awaited>) { + expect((await db.select().from(toolActionRequests).where(eq(toolActionRequests.id, pending.id)))[0].status).toBe("failed"); + expect((await db.select().from(toolInvocations).where(eq(toolInvocations.id, pending.invocationId)))[0]) + .toMatchObject({ status: "failed", errorCode: "tool_execution_outcome_unknown" }); + } + + it("arms from a real execution claim and disarms after the provider result", async () => { + await startInterruption(); await quiet(); + const remote = providerFixture(); + const { pending } = await testRequest(remote.provider); + const running = gateway.approveActionRequest({ companyId: fixture.company.id, + actionRequestId: pending.id, actor: { userId: "board" } }); + await Promise.race([remote.started, running]); + try { + await settle(); + const [claim] = await db.select().from(toolActionRequests); + expect(claim.status).toBe("executing"); + expect(coordinator!.nextWakeAt()).toBe(claim.updatedAt.getTime() + 10 * 60_000); + const reads = vi.spyOn(db, "select"); + await settle(60_000); expect(reads).not.toHaveBeenCalled(); reads.mockRestore(); + } finally { remote.release(); } + await running; await quiet(); + expect((await db.select().from(toolActionRequests))[0].status).toBe("executed"); + expect(remote.provider).toHaveBeenCalledTimes(1); + }); + + it("reconstructs the saved deadline and quarantines only when due", async () => { + const pending = await savedClaim(); await startInterruption(); + expect(coordinator!.nextWakeAt()).toBe(pending.updatedAt.getTime() + 10 * 60_000); + await settle(9 * 60_000); + expect((await db.select().from(toolActionRequests))[0].status).toBe("executing"); + await settle(2 * 60_000); await unknown(pending); await quiet(); + expect(await db.select().from(toolCallEvents).where(eq(toolCallEvents.reasonCode, "approved_action_executed"))).toEqual([]); + }); + + it.each(["succeeded", "failed", "denied", "timed_out", "rate_limited"] as const)("preserves a saved %s invocation while repairing the request", async status => { + const pending = await savedClaim(db, 11 * 60_000); + const completedAt = new Date(Date.now() - 5 * 60_000); + await db.update(toolInvocations).set({ status, completedAt, resultSummary: { summary: "Saved provider result" }, + errorCode: status === "succeeded" ? null : "saved_provider_failure" }).where(eq(toolInvocations.id, pending.invocationId)); + await startInterruption(); await quiet(); + expect((await db.select().from(toolActionRequests))[0]).toMatchObject({ + status: status === "succeeded" ? "executed" : "failed", resolvedAt: completedAt }); + expect((await db.select().from(toolInvocations))[0]).toMatchObject({ status, completedAt, + errorCode: status === "succeeded" ? null : "saved_provider_failure", resultSummary: { summary: "Saved provider result" } }); + }); + + it("retains a first claim committed during drain without passive database scans", async () => { + await startInterruption(); await quiet(); draining = true; + const pending = await savedClaim(); await settle(); + const reads = vi.spyOn(db, "select"); + await settle(60 * 60_000); expect(reads).not.toHaveBeenCalled(); reads.mockRestore(); + draining = false; await settle(); await unknown(pending); await quiet(); + }); + + it.each([false, true])("waits for outer execution commit (rollback %s)", async rollback => { + const pending = await request(); await startInterruption(); await quiet(); + let entered!: () => void, release!: () => void; + const ready = new Promise(resolve => { entered = resolve; }); + const gate = new Promise(resolve => { release = resolve; }); + const aborted = new Error("fixture rollback"); + const running = db.transaction(async tx => { + await writeToolReviewChange(tx as unknown as typeof db, async nested => { + await nested.update(toolActionRequests).set({ status: "executing", updatedAt: new Date() }) + .where(eq(toolActionRequests.id, pending.id)); + }); + entered(); await gate; if (rollback) throw aborted; + }).catch(error => { if (error !== aborted) throw error; }); + await ready; + try { + await vi.advanceTimersByTimeAsync(11 * 60_000); + expect((await db.select().from(toolActionRequests))[0].status).toBe("pending"); + } finally { release(); await running; } + await vi.waitFor(async () => { + await vi.advanceTimersByTimeAsync(30_000); + if (rollback) expect((await db.select().from(toolActionRequests))[0].status).toBe("pending"); + else await unknown(pending); + expect(idleWorkSnapshot().active).toBe(0); + }, { timeout: 5000 }); + await quiet(); + }); + + it("rolls back both outcome rows on failure and retries after drain", async () => { + const pending = await savedClaim(db, 11 * 60_000); + await db.execute(sql`create function sleep_review_fail() returns trigger language plpgsql as $$ + begin raise exception 'interruption fixture failure'; end $$`); + await db.execute(sql`create trigger sleep_review_fail before update on tool_invocations + for each row when (new.error_code = 'tool_execution_outcome_unknown') execute function sleep_review_fail()`); + await startInterruption(); expect(errors.length).toBeGreaterThan(0); errors.length = 0; + expect((await db.select().from(toolActionRequests))[0].status).toBe("executing"); + expect((await db.select().from(toolInvocations))[0].status).toBe("executing"); + draining = true; + await db.execute(sql`drop trigger sleep_review_fail on tool_invocations`); + await settle(60 * 60_000); + draining = false; + await vi.waitFor(async () => { + await vi.advanceTimersByTimeAsync(30_000); + await unknown(pending); + }, { timeout: 5000 }); + await quiet(); + }); + + it("rechecks a changed claim after selecting a stale candidate", async () => { + const pending = await savedClaim(db, 11 * 60_000); + const transaction = db.transaction.bind(db); + let entered!: () => void, release!: () => void; + const ready = new Promise(resolve => { entered = resolve; }); + const gate = new Promise(resolve => { release = resolve; }); + const intercepted = vi.spyOn(db, "transaction").mockImplementationOnce(async callback => { + entered(); await gate; return transaction(callback); + }); + const sweep = gateway.settleInterruptedActionReviews(); + await ready; + const updatedAt = new Date(); + try { + await writeToolReviewChange(db, tx => tx.update(toolActionRequests).set({ updatedAt }) + .where(eq(toolActionRequests.id, pending.id))); + } finally { release(); } + await sweep; intercepted.mockRestore(); + expect((await db.select().from(toolActionRequests))[0].status).toBe("executing"); + await startInterruption(); + expect(coordinator!.nextWakeAt()).toBe(updatedAt.getTime() + 10 * 60_000); + }); + + it("settles more than one page of interrupted claims", async () => { + const pending = await savedClaim(db, 11 * 60_000); + await db.insert(toolActionRequests).values(Array.from({ length: 101 }, () => ({ ...pending, + id: randomUUID(), interactionId: null, status: "executing" as const }))); + await startInterruption(); await quiet(); + const rows = await db.select().from(toolActionRequests); + expect(rows).toHaveLength(102); expect(rows.every(row => row.status === "failed")).toBe(true); + }); + }); + describe("approved review recovery worker", () => { async function startRecovery() { coordinator = createDeliveryWorkCoordinator({ owner: db, canRun: () => !draining, @@ -386,35 +568,6 @@ describeEmbeddedPostgres("tool gateway service", () => { draining = false; await settle(); await executed(pending); await quiet(); }); - async function testRequest(provider: NonNullable) { - await db.insert(companyMemberships).values({ companyId: fixture.company.id, principalType: "user", - principalId: "board", status: "active", membershipRole: "owner" }); - const remote = await createRemoteMcpToolFixture(db, fixture.company.id); - await db.update(toolCatalogEntries).set({ isReadOnly: false, riskLevel: "write" }) - .where(eq(toolCatalogEntries.id, remote.catalogEntry.id)); - await db.insert(toolPolicies).values({ companyId: fixture.company.id, name: "Review connected writes", - policyType: "require_approval", priority: 10, selectors: { connectionId: remote.connection.id } }); - gateway = createTestToolGatewayService(db, { remoteHttpRequest: provider }); - const result = await gateway.executeTestCall({ companyId: fixture.company.id, connectionId: remote.connection.id, - agentId: fixture.agent.id, userId: "board", toolName: "needs_input", parameters: { value: "reviewed" } }); - expect(result.decision).toBe("ask_first"); - const [pending] = await db.select().from(toolActionRequests); - return { pending, remote }; - } - function providerFixture() { - let release!: () => void, entered!: () => void; - const gate = new Promise(resolve => { release = resolve; }); - const started = new Promise(resolve => { entered = resolve; }); - const provider = vi.fn(async (_url, init) => { - const payload = JSON.parse(String(init.body)); - expect(payload.method).toBe("tools/call"); - entered(); await gate; - return Response.json({ jsonrpc: "2.0", id: payload.id, - result: { content: [{ type: "text", text: "Saved once" }] } }); - }); - return { provider, release, started }; - } - it.each(["foreground", "recovery"] as const)("claims a Test-tab action once when %s wins", async winner => { const remoteCall = providerFixture(); const { pending } = await testRequest(remoteCall.provider); diff --git a/server/src/index.ts b/server/src/index.ts index 163eb8e5cbc..c02b2a252c5 100644 --- a/server/src/index.ts +++ b/server/src/index.ts @@ -1,3 +1,4 @@ +import { registerToolReviewInterruptionWork } from "./services/tool-review-interruption-work.js"; import { registerToolReviewRecoveryWork } from "./services/tool-review-recovery-work.js"; import { registerToolReviewExpiryWork } from "./services/tool-review-expiry-work.js"; import { registerGatewayCredentialExpiryWork } from "./services/gateway-credential-expiry-work.js"; @@ -1451,7 +1452,13 @@ async function startServerWithDatabaseTeardown( } else { await app.locals.toolGateway.recoverApprovedActionReviews(); } - return app.locals.toolGateway.sweepActionReviews({ includeExpiry: false, includeApproved: false }); + if (config.heartbeatSchedulerEnabled) { + await registerToolReviewInterruptionWork(deliveryWork, app.locals.toolGateway, + () => !isIdleTaskDrainActive() && !heartbeatSchedulerStopped && !isWarmStandby()).ready; + } else { + await app.locals.toolGateway.settleInterruptedActionReviews(); + } + return app.locals.toolGateway.sweepActionReviews({ includeExpiry: false, includeApproved: false, includeInterrupted: false }); }).catch((err: unknown) => logger.error({ err }, "startup tool review recovery failed")); await timeStartupRecovery("gateway_sessions", () => config.heartbeatSchedulerEnabled ? registerGatewayCredentialExpiryWork(deliveryWork, app.locals.toolGateway, @@ -1876,7 +1883,7 @@ async function startServerWithDatabaseTeardown( logger.error({ err }, "periodic tool connection health sweep failed"); })); - trackHeartbeatSchedulerWork(app.locals.toolGateway.sweepActionReviews({ includeExpiry: false, includeApproved: false }).catch((err: unknown) => logger.error({ err }, "tool review recovery failed"))); + trackHeartbeatSchedulerWork(app.locals.toolGateway.sweepActionReviews({ includeExpiry: false, includeApproved: false, includeInterrupted: false }).catch((err: unknown) => logger.error({ err }, "tool review recovery failed"))); if (isIdleTaskDrainActive() || heartbeatSchedulerStopped) return; if (!(await heartbeat.resolveSchedulingSuppression()).suppressed) { diff --git a/server/src/services/delivery-work-notifications.ts b/server/src/services/delivery-work-notifications.ts index d77c801ca5b..35073a92fb6 100644 --- a/server/src/services/delivery-work-notifications.ts +++ b/server/src/services/delivery-work-notifications.ts @@ -8,6 +8,7 @@ export const DELIVERY_QUEUES = { executionStatus: "execution-status", executionDisposition: "execution-disposition", environmentCleanup: "environment-cleanup", + toolReviewInterruption: "tool-review-interruption", toolReviewRecovery: "tool-review-recovery", toolReviewExpiry: "tool-review-expiry", gatewayCredentialExpiry: "gateway-credential-expiry", diff --git a/server/src/services/tool-gateway.ts b/server/src/services/tool-gateway.ts index 4846e53bdf1..8ac013da29e 100644 --- a/server/src/services/tool-gateway.ts +++ b/server/src/services/tool-gateway.ts @@ -1,3 +1,4 @@ +import { toolReviewInterruptionService } from "./tool-review-interruption.js"; import { claimToolTestApproval, persistToolTestOutcome, startToolTestInvocation, type ToolTestApprovalClaim } from "./tool-test-approval.js"; import { recoverApprovedToolReviews } from "./tool-review-recovery.js"; import { toolReviewExpiryService } from "./tool-review-expiry.js"; @@ -44,7 +45,6 @@ import { asc, desc, eq, - gt, inArray, isNull, lte, @@ -7767,7 +7767,7 @@ export function createToolGatewayService( const message = input.error instanceof Error ? input.error.message : String(input.error); const now = new Date(); - const settled = await db.transaction(async (tx) => { + const settled = await writeToolReviewChange(db, async (tx) => { // Lock in the same invocation -> request order used by the normal // execution settlement path. The claim timestamp is the ownership token: // a consumer that merely read an approved row cannot settle another @@ -8286,10 +8286,10 @@ export function createToolGatewayService( }) .where(eq(toolInvocations.id, invocation.id)); void emitConnectionInvoked(db, invocation.id); - await db + await writeToolReviewChange(db, tx => tx .update(toolActionRequests) .set({ status: "executed", resolvedAt: now, updatedAt: now }) - .where(eq(toolActionRequests.id, claimed.id)); + .where(eq(toolActionRequests.id, claimed.id))); await reflectToolActionInteractionLifecycle({ actionRequestId: claimed.id, status: "executed", @@ -8628,6 +8628,11 @@ export function createToolGatewayService( }; } + const reviewInterruption = toolReviewInterruptionService(db, async settled => { + void emitConnectionInvoked(db, settled.invocationId); + await reflectToolActionInteractionLifecycle(settled); + }); + const reviewExpiry = toolReviewExpiryService(db, async expired => { void emitConnectionInvoked(db, expired.invocationId); await reflectToolActionInteractionLifecycle(expired); @@ -9646,6 +9651,8 @@ export function createToolGatewayService( return buildTestCallStatus(actionRequest, invocation); }, + settleInterruptedActionReviews: reviewInterruption.sweep, + nextActionReviewInterruptionAt: reviewInterruption.next, sweepExpiredActionReviews: reviewExpiry.sweep, nextActionReviewExpiryAt: reviewExpiry.next, @@ -9673,7 +9680,7 @@ export function createToolGatewayService( }); }, - async sweepActionReviews(input: { includeExpiry?: boolean; includeApproved?: boolean } = {}) { + async sweepActionReviews(input: { includeExpiry?: boolean; includeApproved?: boolean; includeInterrupted?: 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 @@ -9708,62 +9715,9 @@ export function createToolGatewayService( errorMessage: invocation.errorMessage, resultSummary: invocation.resultSummary?.summary, }); - const now = new Date(); - const staleAt = new Date(now.getTime() - 10 * 60_000); - let cursor: string | undefined; - let scanned = input.includeExpiry === false ? 0 : (await reviewExpiry.sweep(now)).scanned; + let scanned = input.includeExpiry === false ? 0 : (await reviewExpiry.sweep()).scanned; if (input.includeApproved !== false) scanned += (await this.recoverApprovedActionReviews()).scanned; - for (;;) { - const rows = await db - .select() - .from(toolActionRequests) - .where( - and( - eq(toolActionRequests.status, "executing"), - lte(toolActionRequests.updatedAt, staleAt), - cursor ? gt(toolActionRequests.id, cursor) : undefined, - ), - ) - .orderBy(asc(toolActionRequests.id)) - .limit(100); - for (const row of rows) { - 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 }) - .where( - and( - eq(toolActionRequests.id, row.id), - eq(toolActionRequests.status, row.status), - eq(toolActionRequests.updatedAt, row.updatedAt), - ), - ) - .returning(); - if (!changed) continue; - await db - .update(toolInvocations) - .set({ - status: "failed", - errorCode, - errorMessage, - completedAt: now, - updatedAt: now, - }) - .where(eq(toolInvocations.id, row.invocationId)); - void emitConnectionInvoked(db, row.invocationId); - await reflectToolActionInteractionLifecycle({ - actionRequestId: row.id, - status, - errorCode, - errorMessage, - }); - } - scanned += rows.length; - if (rows.length < 100) break; - cursor = rows[rows.length - 1].id; - } + if (input.includeInterrupted !== false) scanned += (await reviewInterruption.sweep()).scanned; return { scanned }; }, @@ -10961,7 +10915,8 @@ export function createToolGatewayService( }) .where(eq(toolInvocations.id, invocationId)); if (input.approvedActionRequestId) { - const [executedRequest] = await db + const actionRequestId = input.approvedActionRequestId; + const [executedRequest] = await writeToolReviewChange(db, tx => tx .update(toolActionRequests) .set({ status: "executed", @@ -10970,11 +10925,11 @@ export function createToolGatewayService( }) .where( and( - eq(toolActionRequests.id, input.approvedActionRequestId), + eq(toolActionRequests.id, actionRequestId), eq(toolActionRequests.status, "executing"), ), ) - .returning({ id: toolActionRequests.id }); + .returning({ id: toolActionRequests.id })); if (executedRequest) { await reflectToolActionInteractionLifecycle({ actionRequestId: executedRequest.id, @@ -11094,7 +11049,8 @@ export function createToolGatewayService( .where(eq(toolInvocations.id, invocationId)); void emitConnectionInvoked(db, invocationId); if (input.approvedActionRequestId) { - const [failedRequest] = await db + const actionRequestId = input.approvedActionRequestId; + const [failedRequest] = await writeToolReviewChange(db, tx => tx .update(toolActionRequests) .set({ status: "failed", @@ -11103,11 +11059,11 @@ export function createToolGatewayService( }) .where( and( - eq(toolActionRequests.id, input.approvedActionRequestId), + eq(toolActionRequests.id, actionRequestId), eq(toolActionRequests.status, "executing"), ), ) - .returning({ id: toolActionRequests.id }); + .returning({ id: toolActionRequests.id })); if (failedRequest) { await reflectToolActionInteractionLifecycle({ actionRequestId: failedRequest.id, diff --git a/server/src/services/tool-review-interruption-work.ts b/server/src/services/tool-review-interruption-work.ts new file mode 100644 index 00000000000..204ce092fb7 --- /dev/null +++ b/server/src/services/tool-review-interruption-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 registerToolReviewInterruptionWork( + coordinator: ReturnType, + gateway: Pick, "settleInterruptedActionReviews" | "nextActionReviewInterruptionAt">, + canRun: () => boolean, +) { + return registerDeadlineWork(coordinator, DELIVERY_QUEUES.toolReviewInterruption, { + retryMs: 30_000, canRun, + run: () => gateway.settleInterruptedActionReviews(), + next: () => gateway.nextActionReviewInterruptionAt(), + }); +} diff --git a/server/src/services/tool-review-interruption.ts b/server/src/services/tool-review-interruption.ts new file mode 100644 index 00000000000..bbe2c492d2a --- /dev/null +++ b/server/src/services/tool-review-interruption.ts @@ -0,0 +1,61 @@ +import { and, asc, eq, gt, lte } from "drizzle-orm"; +import { toolActionRequests, toolInvocations, type Db } from "@paperclipai/db"; +import { writeToolReviewChange } from "./tool-review-work-notifications.js"; + +const STALE_EXECUTION_MS = 10 * 60_000; + +export function toolReviewInterruptionService(db: Db, onSettled: (input: { + actionRequestId: string; + invocationId: string; + status: "executed" | "failed"; + errorCode: string | null; + errorMessage: string | null; + resultSummary?: string; +}) => Promise) { + return { + async next(): Promise { + const [row] = await db.select({ at: toolActionRequests.updatedAt }).from(toolActionRequests) + .where(eq(toolActionRequests.status, "executing")).orderBy(asc(toolActionRequests.updatedAt)).limit(1); + return row ? row.at.getTime() + STALE_EXECUTION_MS : null; + }, + async sweep(now = new Date()) { + const staleAt = new Date(now.getTime() - STALE_EXECUTION_MS); + let cursor: string | undefined; + let scanned = 0; + for (;;) { + const rows = await db.select().from(toolActionRequests).where(and( + eq(toolActionRequests.status, "executing"), lte(toolActionRequests.updatedAt, staleAt), + cursor ? gt(toolActionRequests.id, cursor) : undefined, + )).orderBy(asc(toolActionRequests.id)).limit(100); + for (const row of rows) { + const settled = await writeToolReviewChange(db, async tx => { + const [invocation] = await tx.select().from(toolInvocations) + .where(eq(toolInvocations.id, row.invocationId)).for("update"); + const [request] = await tx.select().from(toolActionRequests) + .where(eq(toolActionRequests.id, row.id)).for("update"); + if (!invocation || request?.status !== "executing" || + request.updatedAt.getTime() !== row.updatedAt.getTime()) return null; + const completed = ["succeeded", "failed", "denied", "timed_out", "rate_limited"].includes(invocation.status); + const status = invocation.status === "succeeded" ? "executed" as const : "failed" as const; + const errorCode = completed ? invocation.errorCode : "tool_execution_outcome_unknown"; + const errorMessage = completed ? invocation.errorMessage + : "Execution was interrupted; the external outcome is unknown. Inspect the provider before retrying."; + if (!completed) { + await tx.update(toolInvocations).set({ status: "failed", errorCode, errorMessage, + completedAt: now, updatedAt: now }).where(eq(toolInvocations.id, row.invocationId)); + } + await tx.update(toolActionRequests).set({ status, resolvedAt: invocation.completedAt ?? now, updatedAt: now }) + .where(eq(toolActionRequests.id, row.id)); + return { actionRequestId: row.id, invocationId: row.invocationId, status, + errorCode, errorMessage, resultSummary: invocation.resultSummary?.summary }; + }); + if (settled) await onSettled(settled); + } + scanned += rows.length; + if (rows.length < 100) break; + cursor = rows[rows.length - 1].id; + } + return { scanned }; + }, + }; +} diff --git a/server/src/services/tool-review-work-notifications.ts b/server/src/services/tool-review-work-notifications.ts index dd36fbb30af..728fad6d15f 100644 --- a/server/src/services/tool-review-work-notifications.ts +++ b/server/src/services/tool-review-work-notifications.ts @@ -4,6 +4,7 @@ import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notificatio export async function notifyToolReviewChange(transaction: object) { await notifyDeliveryWork(transaction, DELIVERY_QUEUES.toolReviewExpiry); await notifyDeliveryWork(transaction, DELIVERY_QUEUES.toolReviewRecovery); + await notifyDeliveryWork(transaction, DELIVERY_QUEUES.toolReviewInterruption); } export function writeToolReviewChange(db: Db, write: (transaction: Db) => Promise): Promise {