fix(gateway): keep in-process dispatch timeout budget

steipete · steipete · commit 6c00c17a4588 · 2026-06-07T10:42:09.000+01:00
diff --git a/src/gateway/server-plugins.test.ts b/src/gateway/server-plugins.test.ts
@@ -804,6 +804,62 @@ describe("loadGatewayPlugins", () => {
     ).resolves.toEqual({ status: "accepted", runId: "run-accepted" });
   });
 
+  test("uses one timeout budget across accepted and final in-process responses", async () => {
+    vi.useFakeTimers();
+    try {
+      serverPluginsModule.setFallbackGatewayContext(createTestContext("single-final-deadline"));
+      handleGatewayRequest.mockImplementationOnce(async (opts: HandleGatewayRequestOptions) => {
+        setTimeout(() => {
+          opts.respond(true, { status: "accepted", runId: "run-deadline" });
+        }, 7);
+        setTimeout(() => {
+          opts.respond(true, { status: "ok", runId: "run-deadline" });
+        }, 13);
+        await new Promise((resolve) => setTimeout(resolve, 13));
+      });
+
+      const result = expect(
+        serverPluginsModule.dispatchGatewayMethodInProcess(
+          "agent",
+          { sessionKey: "s-deadline" },
+          { expectFinal: true, timeoutMs: 10 },
+        ),
+      ).rejects.toThrow("gateway request timeout for agent");
+
+      await vi.advanceTimersByTimeAsync(10);
+      await result;
+      await vi.advanceTimersByTimeAsync(10);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  test("clears final-response timeout when handler rejects after accepted response", async () => {
+    vi.useFakeTimers();
+    try {
+      serverPluginsModule.setFallbackGatewayContext(createTestContext("accepted-then-error"));
+      handleGatewayRequest.mockImplementationOnce(async (opts: HandleGatewayRequestOptions) => {
+        opts.respond(true, { status: "accepted", runId: "run-error-after-accepted" });
+        await new Promise((resolve) => setTimeout(resolve, 5));
+        throw new Error("handler failed after accepted");
+      });
+
+      const result = expect(
+        serverPluginsModule.dispatchGatewayMethodInProcess(
+          "agent",
+          { sessionKey: "s-error-after-accepted" },
+          { expectFinal: true, timeoutMs: 1_000 },
+        ),
+      ).rejects.toThrow("handler failed after accepted");
+
+      await vi.advanceTimersByTimeAsync(5);
+      await result;
+      expect(vi.getTimerCount()).toBe(0);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
   test("filters connected plugin nodes locally without sending unsupported node.list params", async () => {
     loadOpenClawPlugins.mockReturnValue(createRegistry([]));
     loadGatewayStartupPluginsForTest();
diff --git a/src/gateway/server-plugins.ts b/src/gateway/server-plugins.ts
@@ -340,13 +340,24 @@ function resolveInProcessDispatchTimeoutMs(timeoutMs?: number): number | undefin
     : undefined;
 }
 
+function resolveInProcessDispatchDeadlineMs(timeoutMs?: number): number | undefined {
+  const safeTimeoutMs = resolveInProcessDispatchTimeoutMs(timeoutMs);
+  return safeTimeoutMs === undefined ? undefined : Date.now() + safeTimeoutMs;
+}
+
+function resolveRemainingInProcessDispatchTimeoutMs(deadlineMs?: number): number | undefined {
+  return deadlineMs === undefined
+    ? undefined
+    : resolveSafeTimeoutDelayMs(deadlineMs - Date.now(), { minMs: 0 });
+}
+
 async function waitForInProcessDispatch<T>(
   method: string,
   promise: Promise<T>,
-  timeoutMs?: number,
+  deadlineMs?: number,
 ): Promise<T> {
-  const safeTimeoutMs = resolveInProcessDispatchTimeoutMs(timeoutMs);
-  if (safeTimeoutMs === undefined) {
+  const remainingTimeoutMs = resolveRemainingInProcessDispatchTimeoutMs(deadlineMs);
+  if (remainingTimeoutMs === undefined) {
     return await promise;
   }
   let timeout: NodeJS.Timeout | undefined;
@@ -356,7 +367,7 @@ async function waitForInProcessDispatch<T>(
       new Promise<never>((_resolve, reject) => {
         timeout = setTimeout(() => {
           reject(new Error(`gateway request timeout for ${method}`));
-        }, safeTimeoutMs);
+        }, remainingTimeoutMs);
       }),
     ]);
   } finally {
@@ -396,6 +407,7 @@ export async function dispatchGatewayMethodInProcessRaw(
     resolveFirstResponse = resolve;
     rejectFirstResponse = reject;
   });
+  const deadlineMs = resolveInProcessDispatchDeadlineMs(options?.timeoutMs);
   const { handleGatewayRequest } = await import("./server-methods.js");
   const pluginRuntimeOwnerId =
     typeof options?.pluginRuntimeOwnerId === "string" && options.pluginRuntimeOwnerId.trim()
@@ -462,7 +474,7 @@ export async function dispatchGatewayMethodInProcessRaw(
       rejectFinalResponse?.(error);
     });
 
-  firstResponse = await waitForInProcessDispatch(method, firstResponsePromise, options?.timeoutMs);
+  firstResponse = await waitForInProcessDispatch(method, firstResponsePromise, deadlineMs);
   const firstPayload = firstResponse.payload as { status?: unknown } | undefined;
   if (options?.expectFinal !== true || firstPayload?.status !== "accepted") {
     return firstResponse;
@@ -474,32 +486,33 @@ export async function dispatchGatewayMethodInProcessRaw(
     finalResponse ??
     (await new Promise<GatewayMethodDispatchResponse>((resolve, reject) => {
       resolveFinalResponse = resolve;
-      rejectFinalResponse = reject;
-      const timeoutMs = resolveInProcessDispatchTimeoutMs(options.timeoutMs);
+      const timeoutMs = resolveRemainingInProcessDispatchTimeoutMs(deadlineMs);
       const timeout =
         timeoutMs === undefined
           ? undefined
           : setTimeout(() => {
               reject(new Error(`gateway request timeout for ${method}`));
             }, timeoutMs);
-      if (postFirstResponseError) {
+      const clearFinalTimeout = () => {
         if (timeout) {
           clearTimeout(timeout);
         }
-        reject(postFirstResponseError);
+      };
+      rejectFinalResponse = (err) => {
+        clearFinalTimeout();
+        reject(err);
+      };
+      if (postFirstResponseError) {
+        rejectFinalResponse(postFirstResponseError);
         return;
       }
       if (finalResponse) {
-        if (timeout) {
-          clearTimeout(timeout);
-        }
+        clearFinalTimeout();
         resolve(finalResponse);
         return;
       }
       resolveFinalResponse = (response) => {
-        if (timeout) {
-          clearTimeout(timeout);
-        }
+        clearFinalTimeout();
         resolve(response);
       };
     }));
