diff --git a/.changeset/mcp-wait-timeout.md b/.changeset/mcp-wait-timeout.md new file mode 100644 index 0000000000..02d6c98231 --- /dev/null +++ b/.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. diff --git a/packages/cli-v3/src/mcp/config.ts b/packages/cli-v3/src/mcp/config.ts index 206b5910fa..5a1ec45cba 100644 --- a/packages/cli-v3/src/mcp/config.ts +++ b/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", diff --git a/packages/cli-v3/src/mcp/schemas.ts b/packages/cli-v3/src/mcp/schemas.ts index b98faca0da..8afb10f38f 100644 --- a/packages/cli-v3/src/mcp/schemas.ts +++ b/packages/cli-v3/src/mcp/schemas.ts @@ -123,6 +123,17 @@ export const CommonRunsInput = CommonProjectsInput.extend({ export type CommonRunsInput = z.output; +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; + export const GetRunDetailsInput = CommonRunsInput.extend({ maxTraceLines: z .number() diff --git a/packages/cli-v3/src/mcp/tools/runs.ts b/packages/cli-v3/src/mcp/tools/runs.ts index 13fe601da0..056544e3cd 100644 --- a/packages/cli-v3/src/mcp/tools/runs.ts +++ b/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,20 +87,35 @@ 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; } } @@ -108,8 +123,12 @@ export const waitForRunToCompleteTool = { 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) }], }; }), };