A few improvements

Author: ericallam

packages/cli-v3/src/entryPoints/dev-run-worker.ts

@@ -409,27 +409,11 @@ const zodIpc = new ZodIpcConnection({
 
           _executionMeasurement = usage.start();
 
-          // This lives outside of the executor because this will eventually be moved to the controller level
           const timeoutController = timeout.abortAfterTimeout(execution.run.maxDuration);
 
-          timeoutController.signal.addEventListener("abort", () => {
-            if (_isCancelled) {
-              return;
-            }
-
-            if (cancelController.signal.aborted) {
-              return;
-            }
+          const signal = AbortSignal.any([cancelController.signal, timeoutController.signal]);
 
-            cancelController.abort(timeoutController.signal.reason);
-          });
-
-          const { result } = await executor.execute(
-            execution,
-            metadata,
-            traceContext,
-            cancelController.signal
-          );
+          const { result } = await executor.execute(execution, metadata, traceContext, signal);
 
           if (_isRunning && !_isCancelled) {
             const usageSample = usage.stop(_executionMeasurement);

packages/cli-v3/src/entryPoints/managed-run-worker.ts

@@ -404,27 +404,11 @@ const zodIpc = new ZodIpcConnection({
 
           _executionMeasurement = usage.start();
 
-          // This lives outside of the executor because this will eventually be moved to the controller level
           const timeoutController = timeout.abortAfterTimeout(execution.run.maxDuration);
 
-          timeoutController.signal.addEventListener("abort", () => {
-            if (_isCancelled) {
-              return;
-            }
-
-            if (cancelController.signal.aborted) {
-              return;
-            }
+          const signal = AbortSignal.any([cancelController.signal, timeoutController.signal]);
 
-            cancelController.abort(timeoutController.signal.reason);
-          });
-
-          const { result } = await executor.execute(
-            execution,
-            metadata,
-            traceContext,
-            cancelController.signal
-          );
+          const { result } = await executor.execute(execution, metadata, traceContext, signal);
 
           if (_isRunning && !_isCancelled) {
             const usageSample = usage.stop(_executionMeasurement);

packages/core/src/v3/timeout/usageTimeoutManager.ts

@@ -21,6 +21,10 @@ export class UsageTimeoutManager implements TimeoutManager {
       return this._abortController;
     }
 
+    if (this._intervalId) {
+      clearInterval(this._intervalId);
+    }
+
     // Now we need to start an interval that will measure usage and abort the signal if the usage is too high
     this._intervalId = setInterval(() => {
       const sample = this.usageManager.sample();

packages/core/src/v3/types/tasks.ts

@@ -117,6 +117,7 @@ export type CancelFnParams = Prettify<{
   ctx: Context;
   /** Abort signal that is aborted when a task run exceeds it's maxDuration or if the task run is cancelled. Can be used to automatically cancel downstream requests */
   signal: AbortSignal;
+  runPromise: Promise<unknown>;
   init?: InitOutput;
 }>;
 

references/d3-chat/src/trigger/chat.ts

@@ -231,6 +231,7 @@ export const interruptibleChat = schemaTask({
 
     // 👇 This is a global onCancel hook, but it's inside of the run function
     tasks.onCancel(async () => {
+      // We have access to the chunks here
       logger.info("interruptible-chat: task cancelled with chunks", { chunks });
     });
 
@@ -246,16 +247,6 @@ export const interruptibleChat = schemaTask({
         onChunk: ({ chunk }) => {
           chunks.push(chunk);
         },
-        onError: ({ error }) => {
-          if (error instanceof Error && error.name === "AbortError") {
-            logger.info("interruptible-chat: streamText aborted", { error });
-          } else {
-            logger.error("interruptible-chat: streamText error", { error });
-          }
-        },
-        onFinish: ({ finishReason }) => {
-          logger.info("interruptible-chat: streamText finished", { finishReason });
-        },
       });
 
       const textParts = [];
@@ -267,16 +258,10 @@ export const interruptibleChat = schemaTask({
       return textParts.join("");
     } catch (error) {
       if (error instanceof Error && error.name === "AbortError") {
-        logger.info("interruptible-chat: streamText aborted (inside catch)", { error });
+        // streamText will throw an AbortError if the signal is aborted, so we can handle it here
       } else {
-        logger.error("interruptible-chat: streamText error (inside catch)", { error });
+        throw error;
       }
     }
   },
-  onCancel: async ({ runPromise }) => {
-    //   👇 output is typed as `string` because that's the return type of the run function
-    const output = await runPromise;
-
-    logger.info("interruptible-chat: task cancelled with output", { output });
-  },
 });

references/hello-world/src/trigger/example.ts

@@ -212,25 +212,31 @@ export const hooksTask = task({
   },
 });
 
-export const cancelTask = task({
-  id: "cancel",
-  run: async (payload: { message: string }, { ctx, signal }) => {
-    logger.info("Hello, world from the cancel task", { message: payload.message });
+export const cancelExampleTask = task({
+  id: "cancel-example",
+  // Signal will be aborted when the task is cancelled 👇
+  run: async (payload: { timeoutInSeconds: number }, { signal }) => {
+    logger.info("Hello, world from the cancel task", {
+      timeoutInSeconds: payload.timeoutInSeconds,
+    });
 
+    // This is a global hook that will be called if the task is cancelled
     tasks.onCancel(async () => {
       logger.info("global task onCancel hook but inside of the run function baby!");
     });
 
-    await logger.trace("10s timeout", async (span) => {
+    await logger.trace("timeout", async (span) => {
       try {
-        await setTimeout(10_000, undefined, { signal });
+        // We pass the signal to setTimeout to abort the timeout if the task is cancelled
+        await setTimeout(payload.timeoutInSeconds * 1000, undefined, { signal });
       } catch (error) {
+        // If the timeout is aborted, this error will be thrown, we can handle it here
         logger.error("Timeout error", { error });
       }
     });
 
     logger.info("Hello, world from the cancel task after the timeout", {
-      message: payload.message,
+      timeoutInSeconds: payload.timeoutInSeconds,
     });
 
     return {
@@ -239,9 +245,14 @@ export const cancelTask = task({
   },
   onCancel: async ({ payload, runPromise }) => {
     logger.info("Hello, world from the onCancel hook", { payload });
+    // You can await the runPromise to get the output of the task
     const output = await runPromise;
+
     logger.info("Hello, world from the onCancel hook after the run", { payload, output });
+
+    // You can do work inside the onCancel hook, up to 30 seconds
     await setTimeout(10_000);
+
     logger.info("Hello, world from the onCancel hook after the timeout", { payload });
   },
 });