History

Switch worker approval waits to Convex onUpdate

Move Cloudflare sandbox approval handling from retry polling to Convex websocket updates so approved tool calls resume immediately without hammering callback RPCs.

Rhys Sullivan committed 9 months ago

Commit 8e4b94e1 parent 55a88e94 files · +115 −32Browse files

Unified · Split

4 files
  1. executor/packages/convex/executor_node/external_tool_call.ts +1 −1
  2. executor/packages/convex/runtimeCallbacks.ts +22 −1
  3. executor/packages/runner-sandbox-host/src/index.ts +88 −16
  4. executor/packages/runner-sandbox-host/src/isolate/harness.isolate.js +4 −14
executor/packages/convex/executor_node/external_tool_call.ts+1−1
@@ -43,7 +43,7 @@ export async function handleExternalToolCallRequest( · 42 unmodified lines
4343 ok: false,
4444 kind: "pending",
4545 approvalId,
46− retryAfterMs: 500,
46+ retryAfterMs: 0,
4747 error: "Approval pending",
4848 };
4949 }
executor/packages/convex/runtimeCallbacks.ts+22−1
@@ -1,6 +1,6 @@
11 import { v } from "convex/values";
22 import { internal } from "./_generated/api";
3−import { action, mutation } from "./_generated/server";
3+import { action, mutation, query } from "./_generated/server";
44 import type { ToolCallResult } from "../core/src/types";
55
66 function requireInternalSecret(secret: string): void {
@@ -55,3 +55,24 @@ export const completeRun = mutation({ · 48 unmodified lines
5555 });
5656 },
5757 });
58+
59+export const getApprovalStatus = query({
60+ args: {
61+ internalSecret: v.string(),
62+ runId: v.string(),
63+ approvalId: v.string(),
64+ },
65+ handler: async (ctx, args) => {
66+ requireInternalSecret(args.internalSecret);
67+
68+ const approval = await ctx.runQuery(internal.database.getApproval, {
69+ approvalId: args.approvalId,
70+ });
71+
72+ if (!approval || approval.taskId !== args.runId) {
73+ return { status: "missing" as const };
74+ }
75+
76+ return { status: approval.status };
77+ },
78+});
executor/packages/runner-sandbox-host/src/index.ts+88−16
@@ -36,7 +36,7 @@ · 35 unmodified lines
3636 import { Result } from "better-result";
3737 import { WorkerEntrypoint } from "cloudflare:workers";
3838 import { api } from "@executor/convex/_generated/api";
39−import { ConvexHttpClient } from "convex/browser";
39+import { ConvexClient, ConvexHttpClient } from "convex/browser";
4040
4141 // Import isolate modules as raw text — these are loaded as JS modules inside
4242 // the dynamic isolate, NOT executed in the host worker. The *.isolate.js
@@ -113,6 +113,8 @@ interface BridgeProps { · 70 unmodified lines
113113 taskId: string;
114114 }
115115
116+const APPROVAL_SUBSCRIPTION_TIMEOUT_MS = 10 * 60 * 1000;
117+
116118 // ── Auth ──────────────────────────────────────────────────────────────────────
117119
118120 /** Constant-time string comparison to prevent timing side-channels. */
@@ -168,6 +170,58 @@ export class ToolBridge extends WorkerEntrypoint<Env> { · 49 unmodified lines
168170 });
169171 }
170172
173+ private createRealtimeClient(): ConvexClient {
174+ return new ConvexClient(this.props.callbackConvexUrl, {
175+ skipConvexDeploymentUrlCheck: true,
176+ });
177+ }
178+
179+ private async waitForApprovalUpdate(approvalId: string): Promise<void> {
180+ const client = this.createRealtimeClient();
181+
182+ await new Promise<void>((resolve, reject) => {
183+ let settled = false;
184+
185+ const timer = setTimeout(() => {
186+ if (settled) return;
187+ settled = true;
188+ unsubscribe();
189+ client.close();
190+ reject(new Error(`Timed out waiting for approval update: ${approvalId}`));
191+ }, APPROVAL_SUBSCRIPTION_TIMEOUT_MS);
192+
193+ const unsubscribe = client.onUpdate(
194+ (api as any).runtimeCallbacks.getApprovalStatus,
195+ {
196+ internalSecret: this.props.callbackInternalSecret,
197+ runId: this.props.taskId,
198+ approvalId,
199+ },
200+ (value: { status?: "pending" | "approved" | "denied" | "missing" } | null | undefined) => {
201+ const status = value?.status;
202+ if (!status || status === "pending") {
203+ return;
204+ }
205+ if (status === "missing") {
206+ if (settled) return;
207+ settled = true;
208+ clearTimeout(timer);
209+ unsubscribe();
210+ client.close();
211+ reject(new Error(`Approval not found: ${approvalId}`));
212+ return;
213+ }
214+ if (settled) return;
215+ settled = true;
216+ clearTimeout(timer);
217+ unsubscribe();
218+ client.close();
219+ resolve();
220+ },
221+ );
222+ });
223+ }
224+
171225 /** Forward a tool call to the Convex callback RPC action. */
172226 async callTool(toolPath: string, input: unknown, callId?: string): Promise<ToolCallResult> {
173227 const { callbackInternalSecret, taskId } = this.props;
@@ -175,24 +229,42 @@ export class ToolBridge extends WorkerEntrypoint<Env> { · 1 unmodified line
175229 ? callId
176230 : `call_${crypto.randomUUID()}`;
177231
178− const response = await Result.tryPromise(async () => {
179− const convex = this.createConvexClient();
180− return await convex.action(api.runtimeCallbacks.handleToolCall, {
181− internalSecret: callbackInternalSecret,
182− runId: taskId,
183− callId: effectiveCallId,
184− toolPath,
185− input,
232+ while (true) {
233+ const response = await Result.tryPromise(async () => {
234+ const convex = this.createConvexClient();
235+ return await convex.action(api.runtimeCallbacks.handleToolCall, {
236+ internalSecret: callbackInternalSecret,
237+ runId: taskId,
238+ callId: effectiveCallId,
239+ toolPath,
240+ input,
241+ });
186242 });
187− });
188243
189− if (response.isErr()) {
190− const cause = response.error.cause;
191− const message = cause instanceof Error ? cause.message : String(cause);
192− return { ok: false, kind: "failed", error: `Tool callback failed: ${message}` };
244+ if (response.isErr()) {
245+ const cause = response.error.cause;
246+ const message = cause instanceof Error ? cause.message : String(cause);
247+ return { ok: false, kind: "failed", error: `Tool callback failed: ${message}` };
248+ }
249+
250+ const result = response.value as ToolCallResult;
251+ if (!result.ok && result.kind === "pending") {
252+ if (!result.approvalId) {
253+ return { ok: false, kind: "failed", error: "Approval pending without approvalId" };
254+ }
255+
256+ const approvalId = result.approvalId;
257+ const wait = await Result.tryPromise(() => this.waitForApprovalUpdate(approvalId));
258+ if (wait.isErr()) {
259+ const cause = wait.error.cause;
260+ const message = cause instanceof Error ? cause.message : String(cause);
261+ return { ok: false, kind: "failed", error: `Approval subscription failed: ${message}` };
262+ }
263+ continue;
264+ }
265+
266+ return result;
193267 }
194−
195− return response.value as ToolCallResult;
196268 }
197269 }
198270
executor/packages/runner-sandbox-host/src/isolate/harness.isolate.js+4−14
@@ -3,10 +3,6 @@ import { run } from "./user-code.js"; · 2 unmodified lines
33
44 const APPROVAL_DENIED_PREFIX = "APPROVAL_DENIED:";
55
6−async function sleep(ms) {
7− await new Promise((resolve) => setTimeout(resolve, ms));
8−}
9−
106 function createToolsProxy(bridge, path = []) {
117 const callable = () => {};
128 return new Proxy(callable, {
@@ -21,16 +17,10 @@ function createToolsProxy(bridge, path = []) { · 8 unmodified lines
2117 const input = args.length > 0 ? args[0] : {};
2218 const callId = "call_" + crypto.randomUUID();
2319
24− while (true) {
25− const result = await bridge.callTool(toolPath, input, callId);
26− if (result.ok) return result.value;
27− if (result.kind === "pending") {
28− await sleep(Math.max(50, result.retryAfterMs ?? 500));
29− continue;
30− }
31− if (result.kind === "denied") throw new Error(APPROVAL_DENIED_PREFIX + result.error);
32− throw new Error(result.error);
33− }
20+ const result = await bridge.callTool(toolPath, input, callId);
21+ if (result.ok) return result.value;
22+ if (result.kind === "denied") throw new Error(APPROVAL_DENIED_PREFIX + result.error);
23+ throw new Error(result.error);
3424 },
3525 });
3626 }