diff --git a/apps/api/src/handlers/mcp/unasked-task-tool-call.ts b/apps/api/src/handlers/mcp/unasked-task-tool-call.ts index eec4015ac..4e86b59d0 100644 --- a/apps/api/src/handlers/mcp/unasked-task-tool-call.ts +++ b/apps/api/src/handlers/mcp/unasked-task-tool-call.ts @@ -1,4 +1,5 @@ import { fingerprintIntegrationToolCall } from '@roomote/db/server'; +import { waitForIntegrationToolApproval } from '@roomote/sdk/tool-approval-wait'; import { INTEGRATION_TOOL_AUTO_PAUSED_AGENT_MESSAGE, describeIntegrationToolAutoAbsentDenial, @@ -9,7 +10,6 @@ import { describeProxyToolApprovalBlock, } from './tool-approval-enforcement'; -const APPROVAL_POLL_MS = 1_500; /** * Longer than an approval stays open, so a wait ends on the owner's decision * or on the approval expiring, never on this limit in normal operation. @@ -127,30 +127,28 @@ export async function decideUnaskedTaskToolCall(input: { break; } - const pollMs = input.pollMs ?? APPROVAL_POLL_MS; - const deadline = Date.now() + APPROVAL_MAX_WAIT_MS; - while (Date.now() < deadline) { - if (callerLeft()) return refused; - const status = await getTaskToolApprovalStatus({ - runId, - approvalId: result.approvalId, - }); + const decision = await waitForIntegrationToolApproval({ + readStatus: () => + getTaskToolApprovalStatus({ runId, approvalId: result.approvalId }), // The owner's approval is consumed here. - if (status === 'approved') return claim(); - if (status === 'expired') { - return { - allowed: false, - message: - 'The requester did not answer in time; the tool call was not run.', - }; - } - if (status !== 'pending') { - return { - allowed: false, - message: 'The requester rejected this tool call.', - }; - } - await new Promise((resolve) => setTimeout(resolve, pollMs)); + claimApproved: async () => (await claim()).allowed, + signal: input.signal, + pollMs: input.pollMs, + deadline: Date.now() + APPROVAL_MAX_WAIT_MS, + }); + if (decision === 'approved') return { allowed: true }; + if (decision === 'expired') { + return { + allowed: false, + message: + 'The requester did not answer in time; the tool call was not run.', + }; + } + if (decision === 'rejected') { + return { + allowed: false, + message: 'The requester rejected this tool call.', + }; } return refused; } diff --git a/apps/worker/src/sandbox-server/lib/harnesses/opencode-server/tool-approvals.ts b/apps/worker/src/sandbox-server/lib/harnesses/opencode-server/tool-approvals.ts index 6c7ef8ec1..cea8f915d 100644 --- a/apps/worker/src/sandbox-server/lib/harnesses/opencode-server/tool-approvals.ts +++ b/apps/worker/src/sandbox-server/lib/harnesses/opencode-server/tool-approvals.ts @@ -2,6 +2,7 @@ import { Client } from '@modelcontextprotocol/sdk/client/index.js'; import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'; import { sdk } from '@roomote/sdk/client'; +import { waitForIntegrationToolApproval } from '@roomote/sdk/tool-approval-wait'; import { INTEGRATION_TOOL_AUTO_PAUSED_AGENT_MESSAGE, describeIntegrationToolAutoAbsentDenial, @@ -14,7 +15,6 @@ import { parseDirectMcpConfig } from './mcp-config'; import type { OpenCodeServerClient } from './client'; import type { OpenCodeToolPart } from './types'; -const TOOL_APPROVAL_POLL_MS = 1_500; const AUTO_SERVER_TOOL_LIST_TIMEOUT_MS = 15_000; type TaskToolApprovalApi = Pick; @@ -171,7 +171,6 @@ export function createTaskToolApprovalRelay(options: { onPendingCountChange?: (pending: number) => void; }) { const api = options.api ?? sdk.toolApprovals; - const pollMs = options.pollMs ?? TOOL_APPROVAL_POLL_MS; const handled = new Set(); let pending = 0; @@ -245,27 +244,25 @@ export function createTaskToolApprovalRelay(options: { ); return; } - for (;;) { - if (options.signal.aborted) return; - const { status } = await api.status(result.approvalId); - if (status === 'approved') { - await reply(ask, 'once'); - return; - } - if (status === 'expired') { - await reply( - ask, - 'reject', - 'The requester did not answer in time; the tool call was not run.', - ); - return; - } - if (status !== 'pending') { - await reply(ask, 'reject', 'The requester rejected this tool call.'); - return; - } - await new Promise((resolve) => setTimeout(resolve, pollMs)); + const decision = await waitForIntegrationToolApproval({ + readStatus: async () => (await api.status(result.approvalId)).status, + signal: options.signal, + pollMs: options.pollMs, + }); + if (decision === 'aborted') return; + if (decision === 'approved') { + await reply(ask, 'once'); + return; + } + if (decision === 'expired') { + await reply( + ask, + 'reject', + 'The requester did not answer in time; the tool call was not run.', + ); + return; } + await reply(ask, 'reject', 'The requester rejected this tool call.'); }; /** Fire-and-forget from the event loop; the pause is the intended state. */ diff --git a/packages/cloud-agents/src/server/fast-agent/fast-agent-tool-approvals.ts b/packages/cloud-agents/src/server/fast-agent/fast-agent-tool-approvals.ts index 508e5fe76..00496c2b4 100644 --- a/packages/cloud-agents/src/server/fast-agent/fast-agent-tool-approvals.ts +++ b/packages/cloud-agents/src/server/fast-agent/fast-agent-tool-approvals.ts @@ -24,6 +24,10 @@ import { suspendIntegrationToolAutoForSession, } from '@roomote/db/server'; import { isSessionUserPresent } from '@roomote/redis'; +import { + waitForIntegrationToolApproval, + type IntegrationToolApprovalWaitResult, +} from '@roomote/sdk/tool-approval-wait'; import { INTEGRATION_TOOL_AUTO_PAUSED_AGENT_MESSAGE, describeIntegrationToolAutoAbsentDenial, @@ -70,9 +74,7 @@ import { buildFastAgentCodeModeServerNames } from './fast-agent-tool-policy'; * per ask, which binds the approval to the exact paused call; a repeated * call with changed arguments is a new ask by construction. */ -const INTEGRATION_TOOL_APPROVAL_POLL_MS = 1_500; - -type CardDecision = 'approved' | 'rejected' | 'expired' | 'invalid' | 'aborted'; +type CardDecision = IntegrationToolApprovalWaitResult; const SESSION_PRESENCE_LOOKUP_TIMEOUT_MS = 2_000; /** * An open session page renews its presence every 10 seconds, and a page that @@ -534,32 +536,21 @@ export function createFastAgentToolApprovalBridge(input: { notifiedApprovalIds.add(approval.approvalId); await input.notify(approval); } - const deadline = Date.parse(approval.expiresAt); - for (;;) { - if (input.signal?.aborted) return 'aborted'; - const row = await getIntegrationToolApproval(approval.approvalId); - if (!row || row.status === 'rejected' || row.status === 'cancelled') { - return 'rejected'; - } - if (row.status === 'expired') return 'expired'; - if (row.status === 'approved') { - // Consume before relaying: only the first relay of an approved, - // unclaimed decision reaches OpenCode; a cancelled or - // double-claimed row fails closed instead of executing twice. - const consumed = await markIntegrationToolApprovalConsumed({ + return waitForIntegrationToolApproval({ + readStatus: async () => + (await getIntegrationToolApproval(approval.approvalId))?.status ?? null, + // Consume before relaying: only the first relay of an approved, + // unclaimed decision reaches OpenCode; a cancelled or double-claimed + // row fails closed instead of executing twice. + claimApproved: () => + markIntegrationToolApprovalConsumed({ approvalId: approval.approvalId, requesterUserId: input.userId, - }); - return consumed ? 'approved' : 'invalid'; - } - if (Date.now() >= deadline) { - await expireIntegrationToolApproval(approval.approvalId); - return 'expired'; - } - await new Promise((resolve) => - setTimeout(resolve, INTEGRATION_TOOL_APPROVAL_POLL_MS), - ); - } + }), + signal: input.signal, + deadline: Date.parse(approval.expiresAt), + expire: () => expireIntegrationToolApproval(approval.approvalId), + }); }; // Once Auto stops in this turn the turn is ending: no other call runs or // leaves a card waiting. diff --git a/packages/sdk/package.json b/packages/sdk/package.json index 29926d400..416aa030e 100644 --- a/packages/sdk/package.json +++ b/packages/sdk/package.json @@ -16,6 +16,10 @@ "import": "./src/sandbox-router.ts", "require": "./src/sandbox-router.ts" }, + "./tool-approval-wait": { + "import": "./src/tool-approval-wait.ts", + "require": "./src/tool-approval-wait.ts" + }, "./server": { "import": "./src/server/index.ts", "require": "./src/server/index.ts" diff --git a/packages/sdk/src/tool-approval-wait.test.ts b/packages/sdk/src/tool-approval-wait.test.ts new file mode 100644 index 000000000..b2ba990af --- /dev/null +++ b/packages/sdk/src/tool-approval-wait.test.ts @@ -0,0 +1,128 @@ +import { waitForIntegrationToolApproval } from './tool-approval-wait'; + +describe('waitForIntegrationToolApproval', () => { + it.each([ + ['rejected'], + ['cancelled'], + ['consumed'], + ['auto_approved'], + ['auto_rejected'], + ['not_found'], + [null], + ] as const)('treats %s as a rejected terminal state', async (status) => { + await expect( + waitForIntegrationToolApproval({ + readStatus: async () => status, + pollMs: 0, + }), + ).resolves.toBe('rejected'); + }); + + it('polls pending approvals until they are approved', async () => { + const statuses: Array<'pending' | 'approved'> = [ + 'pending', + 'pending', + 'approved', + ]; + const readStatus = vi.fn(async () => statuses.shift() ?? 'approved'); + + await expect( + waitForIntegrationToolApproval({ readStatus, pollMs: 0 }), + ).resolves.toBe('approved'); + expect(readStatus).toHaveBeenCalledTimes(3); + }); + + it('returns expired without trying to claim', async () => { + const claimApproved = vi.fn(async () => true); + await expect( + waitForIntegrationToolApproval({ + readStatus: async () => 'expired', + claimApproved, + }), + ).resolves.toBe('expired'); + expect(claimApproved).not.toHaveBeenCalled(); + }); + + it('claims an approved decision exactly once and fails closed if unavailable', async () => { + const claimed = vi.fn(async () => true); + await expect( + waitForIntegrationToolApproval({ + readStatus: async () => 'approved', + claimApproved: claimed, + }), + ).resolves.toBe('approved'); + expect(claimed).toHaveBeenCalledTimes(1); + + await expect( + waitForIntegrationToolApproval({ + readStatus: async () => 'approved', + claimApproved: async () => false, + }), + ).resolves.toBe('invalid'); + }); + + it('expires an approval once its waiting deadline passes', async () => { + const expire = vi.fn(async () => undefined); + await expect( + waitForIntegrationToolApproval({ + readStatus: async () => 'pending', + deadline: Date.now() - 1, + expire, + }), + ).resolves.toBe('expired'); + expect(expire).toHaveBeenCalledTimes(1); + }); + + it('stops before reading or claiming after an abort', async () => { + const alreadyAborted = new AbortController(); + alreadyAborted.abort(); + const unread = vi.fn(async () => 'pending' as const); + await expect( + waitForIntegrationToolApproval({ + readStatus: unread, + signal: alreadyAborted.signal, + }), + ).resolves.toBe('aborted'); + expect(unread).not.toHaveBeenCalled(); + + const duringRead = new AbortController(); + const claimApproved = vi.fn(async () => true); + await expect( + waitForIntegrationToolApproval({ + readStatus: async () => { + duringRead.abort(); + return 'approved'; + }, + claimApproved, + signal: duringRead.signal, + }), + ).resolves.toBe('aborted'); + expect(claimApproved).not.toHaveBeenCalled(); + }); + + it('stays aborted when the caller leaves while an approval is claimed', async () => { + const controller = new AbortController(); + await expect( + waitForIntegrationToolApproval({ + readStatus: async () => 'approved', + claimApproved: async () => { + controller.abort(); + return true; + }, + signal: controller.signal, + }), + ).resolves.toBe('aborted'); + }); + + it('interrupts the polling delay when the caller leaves', async () => { + const controller = new AbortController(); + const waiting = waitForIntegrationToolApproval({ + readStatus: async () => 'pending', + signal: controller.signal, + pollMs: 60_000, + }); + + controller.abort(); + await expect(waiting).resolves.toBe('aborted'); + }); +}); diff --git a/packages/sdk/src/tool-approval-wait.ts b/packages/sdk/src/tool-approval-wait.ts new file mode 100644 index 000000000..3ce5483d8 --- /dev/null +++ b/packages/sdk/src/tool-approval-wait.ts @@ -0,0 +1,59 @@ +import type { IntegrationToolApprovalStatus } from '@roomote/types'; + +const DEFAULT_APPROVAL_POLL_MS = 1_500; + +export type IntegrationToolApprovalWaitResult = + | 'approved' + | 'rejected' + | 'expired' + | 'invalid' + | 'aborted'; + +/** + * Wait for a persisted approval to leave pending, optionally claiming an + * approval at the runtime's execution boundary. + */ +export async function waitForIntegrationToolApproval(input: { + readStatus: () => Promise; + claimApproved?: () => Promise; + signal?: AbortSignal; + pollMs?: number; + deadline?: number; + expire?: () => Promise; +}): Promise { + for (;;) { + if (input.signal?.aborted) return 'aborted'; + + const status = await input.readStatus(); + if (input.signal?.aborted) return 'aborted'; + + if (status === 'approved') { + if (!input.claimApproved) return 'approved'; + const claimed = await input.claimApproved(); + if (input.signal?.aborted) return 'aborted'; + return claimed ? 'approved' : 'invalid'; + } + if (status === 'expired') return 'expired'; + if (status !== 'pending') return 'rejected'; + + if (input.deadline !== undefined && Date.now() >= input.deadline) { + await input.expire?.(); + return 'expired'; + } + + await waitForPoll(input.pollMs ?? DEFAULT_APPROVAL_POLL_MS, input.signal); + } +} + +async function waitForPoll(ms: number, signal?: AbortSignal): Promise { + if (ms <= 0 || signal?.aborted) return; + await new Promise((resolve) => { + const done = () => { + clearTimeout(timer); + signal?.removeEventListener('abort', done); + resolve(); + }; + const timer = setTimeout(done, ms); + signal?.addEventListener('abort', done, { once: true }); + }); +}