Skip to content

Commit 6edb774

Browse files
committed
fix: propagate sandbox timeout to nested tool calls
The sandbox's 5-minute timeout only interrupted QuickJS execution via the interrupt handler, but that handler doesn't fire during async suspension (when waiting for host functions like mux.bash()). Changes: - Add setTimeout timer that fires independently of QuickJS execution - Tools now use runtime.getAbortSignal() instead of AI SDK's signal - This signal is aborted both by timeout AND external abort requests - Add abortRequested flag to handle abort() before eval() starts - Fix edge case where addEventListener on already-aborted signal is a no-op When timeout fires during a tool call: 1. setTimeout aborts the runtime's AbortController 2. Tool sees aborted signal and can cancel (if it respects abort) 3. When tool returns, interrupt handler also stops QuickJS Note: JavaScript cannot preemptively cancel Promises - this is cooperative cancellation. Tools that check their abort signal will cancel; pure async operations complete but subsequent calls fail fast.
1 parent 905dc85 commit 6edb774

File tree

6 files changed

+106
-13
lines changed

6 files changed

+106
-13
lines changed

src/node/services/ptc/quickjsRuntime.test.ts

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -291,6 +291,44 @@ describe("QuickJSRuntime", () => {
291291
expect(result.success).toBe(false);
292292
expect(result.error).toContain("timeout");
293293
});
294+
295+
it("aborts signal when timeout fires during async host function", async () => {
296+
// This tests the setTimeout-based timeout's effect on the abort signal.
297+
// The interrupt handler only fires during QuickJS execution, but when
298+
// waiting for an async host function, the setTimeout aborts the signal.
299+
//
300+
// Important: The host function itself won't be cancelled mid-flight
301+
// (JavaScript can't interrupt Promises), but the signal will be aborted
302+
// so subsequent tool calls will see it and fail fast.
303+
let firstCallCompleted = false;
304+
305+
runtime.registerFunction("slowOp", async () => {
306+
// Sleep for 200ms
307+
await new Promise((resolve) => setTimeout(resolve, 200));
308+
firstCallCompleted = true;
309+
return "done";
310+
});
311+
312+
// Check the abort signal state from QuickJS (sync is fine, made async for type)
313+
runtime.registerFunction("checkAbortState", () => {
314+
return Promise.resolve({ aborted: runtime.getAbortSignal()?.aborted ?? false });
315+
});
316+
317+
runtime.setLimits({ timeoutMs: 100 }); // 100ms timeout
318+
319+
const result = await runtime.eval(`
320+
slowOp(); // Takes 200ms, timeout fires at 100ms
321+
checkAbortState(); // Should show aborted = true
322+
slowOp(); // This would start after abort
323+
return "finished";
324+
`);
325+
326+
// The first call completes (can't be interrupted mid-Promise)
327+
expect(firstCallCompleted).toBe(true);
328+
// But the overall execution fails due to timeout
329+
expect(result.success).toBe(false);
330+
expect(result.error).toContain("timeout");
331+
});
294332
});
295333

296334
describe("abort", () => {

src/node/services/ptc/quickjsRuntime.ts

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ export class QuickJSRuntime implements IJSRuntime {
2828
private disposed = false;
2929
private eventHandler?: (event: PTCEvent) => void;
3030
private abortController?: AbortController;
31+
private abortRequested = false; // Track abort requests before eval() starts
3132
private limits: RuntimeLimits = {};
3233
private consoleSetup = false;
3334

@@ -243,6 +244,11 @@ export class QuickJSRuntime implements IJSRuntime {
243244
this.toolCalls = [];
244245
this.consoleOutput = [];
245246

247+
// Honor abort requests made before eval() was called
248+
if (this.abortRequested) {
249+
this.abortController.abort();
250+
}
251+
246252
// Set up console capturing (only once)
247253
if (!this.consoleSetup) {
248254
this.setupConsole();
@@ -264,6 +270,14 @@ export class QuickJSRuntime implements IJSRuntime {
264270
return false; // Continue execution
265271
});
266272

273+
// Set up a real timeout timer that fires even during async suspension.
274+
// The interrupt handler only runs during QuickJS execution, but when suspended
275+
// waiting for an async host function (e.g., mux.bash()), it never fires.
276+
// This timer ensures nested tools are cancelled when the deadline is exceeded.
277+
const timeoutId = setTimeout(() => {
278+
this.abortController?.abort();
279+
}, timeoutMs);
280+
267281
// Wrap code in function to allow return statements.
268282
// With asyncify, async host functions appear synchronous to QuickJS,
269283
// so we don't need an async IIFE. Using evalCodeAsync handles the suspension.
@@ -311,13 +325,20 @@ export class QuickJSRuntime implements IJSRuntime {
311325
consoleOutput: this.consoleOutput,
312326
duration_ms,
313327
};
328+
} finally {
329+
clearTimeout(timeoutId);
314330
}
315331
}
316332

317333
abort(): void {
334+
this.abortRequested = true;
318335
this.abortController?.abort();
319336
}
320337

338+
getAbortSignal(): AbortSignal | undefined {
339+
return this.abortController?.signal;
340+
}
341+
321342
dispose(): void {
322343
if (!this.disposed) {
323344
this.ctx.dispose();

src/node/services/ptc/runtime.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,13 @@ export interface IJSRuntime extends Disposable {
5959
*/
6060
abort(): void;
6161

62+
/**
63+
* Get the abort signal for the current execution.
64+
* This signal is aborted when the sandbox times out or abort() is called.
65+
* Used by tool bridge to propagate cancellation to nested tool calls.
66+
*/
67+
getAbortSignal(): AbortSignal | undefined;
68+
6269
/**
6370
* Clean up resources. Called automatically with `using` declarations.
6471
*/

src/node/services/ptc/toolBridge.test.ts

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ function createMockRuntime(overrides: Partial<IJSRuntime> = {}): IJSRuntime {
2828
setLimits: mock((_limits: RuntimeLimits) => undefined),
2929
onEvent: mock((_handler: (event: PTCEvent) => void) => undefined),
3030
abort: mock(() => undefined),
31+
getAbortSignal: mock(() => undefined),
3132
dispose: mock(() => undefined),
3233
[Symbol.dispose]: mock(() => undefined),
3334
...overrides,
@@ -191,7 +192,7 @@ describe("ToolBridge", () => {
191192
expect(result).toEqual({ error: "Result not JSON-serializable" });
192193
});
193194

194-
it("passes abortSignal to tool execute", async () => {
195+
it("uses runtime abort signal for tool cancellation", async () => {
195196
const mockExecute = mock((_args: unknown) => ({ result: "ok" }));
196197

197198
const tools: Record<string, Tool> = {
@@ -207,16 +208,20 @@ describe("ToolBridge", () => {
207208
return undefined;
208209
}
209210
);
210-
const mockRuntime = createMockRuntime({ registerObject: mockRegisterObject });
211-
211+
// Provide an abort signal via getAbortSignal
212212
const abortController = new AbortController();
213-
bridge.register(mockRuntime, abortController.signal);
213+
const mockRuntime = createMockRuntime({
214+
registerObject: mockRegisterObject,
215+
getAbortSignal: () => abortController.signal,
216+
});
217+
218+
bridge.register(mockRuntime);
214219

215220
await registeredMux.file_read({ filePath: "test.txt" });
216221
expect(mockExecute).toHaveBeenCalledTimes(1);
217222
});
218223

219-
it("throws if abortSignal is already aborted", async () => {
224+
it("throws if runtime abort signal is already aborted", async () => {
220225
const mockExecute = mock(() => ({ result: "ok" }));
221226

222227
const tools: Record<string, Tool> = {
@@ -232,11 +237,16 @@ describe("ToolBridge", () => {
232237
return undefined;
233238
}
234239
);
235-
const mockRuntime = createMockRuntime({ registerObject: mockRegisterObject });
236240

241+
// Pre-abort the signal
237242
const abortController = new AbortController();
238243
abortController.abort();
239-
bridge.register(mockRuntime, abortController.signal);
244+
const mockRuntime = createMockRuntime({
245+
registerObject: mockRegisterObject,
246+
getAbortSignal: () => abortController.signal,
247+
});
248+
249+
bridge.register(mockRuntime);
240250

241251
// Type assertion needed because Record indexing returns T | undefined for ESLint
242252
const fileRead = registeredMux.file_read as (...args: unknown[]) => Promise<unknown>;

src/node/services/ptc/toolBridge.ts

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,8 +65,17 @@ export class ToolBridge {
6565
return Object.fromEntries(this.nonBridgeableTools.entries());
6666
}
6767

68-
/** Register all bridgeable tools on the runtime under `mux` namespace */
69-
register(runtime: IJSRuntime, abortSignal?: AbortSignal): void {
68+
/**
69+
* Register all bridgeable tools on the runtime under `mux` namespace.
70+
*
71+
* Tools receive the runtime's abort signal, which is aborted when:
72+
* - The sandbox timeout is exceeded
73+
* - runtime.abort() is called (e.g., from the parent's abort signal)
74+
*
75+
* This ensures nested tool calls are cancelled when the sandbox times out,
76+
* not just when the parent stream is cancelled.
77+
*/
78+
register(runtime: IJSRuntime): void {
7079
const muxObj: Record<string, (...args: unknown[]) => Promise<unknown>> = {};
7180

7281
for (const [name, tool] of this.bridgeableTools) {
@@ -75,6 +84,9 @@ export class ToolBridge {
7584
const toolName = name;
7685

7786
muxObj[name] = async (args: unknown) => {
87+
// Get the runtime's abort signal - this is aborted on timeout or manual abort
88+
const abortSignal = runtime.getAbortSignal();
89+
7890
// Check if already aborted before executing
7991
if (abortSignal?.aborted) {
8092
throw new Error("Execution aborted");

src/node/services/tools/code_execution.ts

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -129,12 +129,17 @@ ${muxTypes}
129129
});
130130
}
131131

132-
// Register tools with abortSignal for mid-execution cancellation
133-
toolBridge.register(runtime, abortSignal);
132+
// Register tools - they'll use runtime.getAbortSignal() for cancellation
133+
toolBridge.register(runtime);
134134

135-
// Handle abort signal - interrupt sandbox at next checkpoint
135+
// Handle abort signal - interrupt sandbox and cancel nested tools
136136
if (abortSignal) {
137-
abortSignal.addEventListener("abort", () => runtime.abort(), { once: true });
137+
// If already aborted, abort runtime immediately
138+
if (abortSignal.aborted) {
139+
runtime.abort();
140+
} else {
141+
abortSignal.addEventListener("abort", () => runtime.abort(), { once: true });
142+
}
138143
}
139144

140145
// Execute the code

0 commit comments

Comments
 (0)