feat(mcp): add timeout parameter to wait_for_run_to_complete tool

Author: ericallam

.changeset/mcp-wait-timeout.md

@@ -0,0 +1,5 @@
+---
+"trigger.dev": patch
+---
+
+Add optional `timeoutInSeconds` parameter to the `wait_for_run_to_complete` MCP tool. Defaults to 60 seconds. If the run doesn't complete within the timeout, the current state of the run is returned instead of waiting indefinitely.

packages/cli-v3/src/mcp/config.ts

@@ -68,7 +68,7 @@ export const toolsMetadata = {
     name: "wait_for_run_to_complete",
     title: "Wait for Run to Complete",
     description:
-      "Wait for a run to complete. The run ID is the ID of the run that was triggered. It starts with run_",
+      "Wait for a run to complete. The run ID is the ID of the run that was triggered. It starts with run_. Has an optional timeoutInSeconds parameter (default 60s) - if the run doesn't complete within that time, the current state of the run will be returned.",
   },
   cancel_run: {
     name: "cancel_run",

packages/cli-v3/src/mcp/schemas.ts

@@ -123,6 +123,17 @@ export const CommonRunsInput = CommonProjectsInput.extend({
 
 export type CommonRunsInput = z.output<typeof CommonRunsInput>;
 
+export const WaitForRunInput = CommonRunsInput.extend({
+  timeoutInSeconds: z
+    .number()
+    .describe(
+      "The maximum time in seconds to wait for the run to complete. If the run doesn't complete within this time, the current state of the run will be returned. Defaults to 60 seconds."
+    )
+    .default(60),
+});
+
+export type WaitForRunInput = z.output<typeof WaitForRunInput>;
+
 export const GetRunDetailsInput = CommonRunsInput.extend({
   maxTraceLines: z
     .number()

packages/cli-v3/src/mcp/tools/runs.ts

@@ -1,7 +1,7 @@
 import { AnyRunShape } from "@trigger.dev/core/v3";
 import { toolsMetadata } from "../config.js";
 import { formatRun, formatRunList, formatRunShape, formatRunTrace } from "../formatters.js";
-import { CommonRunsInput, GetRunDetailsInput, ListRunsInput } from "../schemas.js";
+import { CommonRunsInput, GetRunDetailsInput, ListRunsInput, WaitForRunInput } from "../schemas.js";
 import { respondWithError, toolHandler } from "../utils.js";
 
 export const getRunDetailsTool = {
@@ -65,8 +65,8 @@ export const waitForRunToCompleteTool = {
   name: toolsMetadata.wait_for_run_to_complete.name,
   title: toolsMetadata.wait_for_run_to_complete.title,
   description: toolsMetadata.wait_for_run_to_complete.description,
-  inputSchema: CommonRunsInput.shape,
-  handler: toolHandler(CommonRunsInput.shape, async (input, { ctx, signal }) => {
+  inputSchema: WaitForRunInput.shape,
+  handler: toolHandler(WaitForRunInput.shape, async (input, { ctx, signal }) => {
     ctx.logger?.log("calling wait_for_run_to_complete", { input });
 
     if (ctx.options.devOnly && input.environment !== "dev") {
@@ -87,29 +87,48 @@ export const waitForRunToCompleteTool = {
       branch: input.branch,
     });
 
-    const runSubscription = apiClient.subscribeToRun(input.runId, { signal });
+    const timeoutMs = input.timeoutInSeconds * 1000;
+    const timeoutSignal = AbortSignal.timeout(timeoutMs);
+    const combinedSignal = signal
+      ? AbortSignal.any([signal, timeoutSignal])
+      : timeoutSignal;
+
+    const runSubscription = apiClient.subscribeToRun(input.runId, { signal: combinedSignal });
     const readableStream = runSubscription.getReader();
 
     let run: AnyRunShape | null = null;
-
-    while (true) {
-      const { done, value } = await readableStream.read();
-      if (done) {
-        break;
+    let timedOut = false;
+
+    try {
+      while (true) {
+        const { done, value } = await readableStream.read();
+        if (done) {
+          break;
+        }
+        run = value;
+
+        if (value.isCompleted) {
+          break;
+        }
       }
-      run = value;
-
-      if (value.isCompleted) {
-        break;
+    } catch (error) {
+      if (timeoutSignal.aborted) {
+        timedOut = true;
+      } else {
+        throw error;
       }
     }
 
     if (!run) {
       return respondWithError("Run not found");
     }
 
+    const prefix = timedOut
+      ? `Timed out after ${input.timeoutInSeconds}s. Returning current run state:\n\n`
+      : "";
+
     return {
-      content: [{ type: "text", text: formatRunShape(run) }],
+      content: [{ type: "text", text: prefix + formatRunShape(run) }],
     };
   }),
 };