fix(gateway): own post-ready drain cancellation (#126855)

src/gateway/server-idle-task.test.ts

@@ -1,7 +1,13 @@
 import { afterEach, describe, expect, it, vi } from "vitest";
+import {
+  GatewayDrainingError,
+  markGatewayRestartDraining,
+  resetGatewayWorkAdmission,
+} from "../process/gateway-work-admission.js";
 import { scheduleGatewayIdleTask } from "./server-idle-task.js";
 
 afterEach(() => {
+  resetGatewayWorkAdmission();
   vi.useRealTimers();
 });
 
@@ -24,4 +30,48 @@ describe("scheduleGatewayIdleTask", () => {
     expect(run).toHaveBeenCalledOnce();
     handle.stop();
   });
+
+  it("quietly cancels idle work rejected by an active restart drain", async () => {
+    vi.useFakeTimers();
+    const run = vi.fn(async () => {});
+    const warn = vi.fn();
+    const handle = scheduleGatewayIdleTask({
+      delayMs: 10,
+      retryDelayMs: 5,
+      isClosing: () => false,
+      isBusy: () => false,
+      run,
+      log: { warn },
+      errorMessage: "idle task failed",
+    });
+
+    markGatewayRestartDraining();
+    await vi.advanceTimersByTimeAsync(10);
+
+    expect(run).not.toHaveBeenCalled();
+    expect(warn).not.toHaveBeenCalled();
+    handle.stop();
+  });
+
+  it("warns when idle work throws a draining error without an active restart", async () => {
+    vi.useFakeTimers();
+    const error = new GatewayDrainingError("unexpected task failure");
+    const warn = vi.fn();
+    const handle = scheduleGatewayIdleTask({
+      delayMs: 10,
+      retryDelayMs: 5,
+      isClosing: () => false,
+      isBusy: () => false,
+      run: async () => {
+        throw error;
+      },
+      log: { warn },
+      errorMessage: "idle task failed",
+    });
+
+    await vi.advanceTimersByTimeAsync(10);
+
+    expect(warn).toHaveBeenCalledWith(`idle task failed: ${String(error)}`);
+    handle.stop();
+  });
 });

src/gateway/server-idle-task.ts

@@ -1,4 +1,7 @@
-import { runWithGatewayIndependentRootWorkAdmission } from "../process/gateway-work-admission.js";
+import {
+  isGatewayRestartDrainError,
+  runWithGatewayIndependentRootWorkAdmission,
+} from "../process/gateway-work-admission.js";
 
 type GatewayIdleTaskLogger = {
   warn: (message: string) => void;
@@ -44,7 +47,11 @@ export function scheduleGatewayIdleTask(params: {
           return;
         }
         await params.run();
-      }).catch((error: unknown) => params.log.warn(`${params.errorMessage}: ${String(error)}`));
+      }).catch((error: unknown) => {
+        if (!isGatewayRestartDrainError(error)) {
+          params.log.warn(`${params.errorMessage}: ${String(error)}`);
+        }
+      });
     }, delayMs);
     timer.unref?.();
   };

src/gateway/server-startup-post-attach.test.ts

@@ -16,7 +16,9 @@ import { getPluginRuntimeGatewayRequestScope } from "../plugins/runtime/gateway-
 import type { PluginServicesHandle } from "../plugins/services.js";
 import type { OpenClawPluginServiceContext } from "../plugins/types.js";
 import {
+  GatewayDrainingError,
   getActiveGatewayRootWorkCount,
+  markGatewayRestartDraining,
   resetGatewayWorkAdmission,
   tryBeginGatewayRootWorkAdmission,
 } from "../process/gateway-work-admission.js";
@@ -1822,6 +1824,84 @@ describe("startGatewayPostAttachRuntime", () => {
     }
   });
 
+  it("owns a queued provider auth rewarm rejected by restart drain without warning", async () => {
+    vi.useFakeTimers();
+    const log = { info: vi.fn(), warn: vi.fn() };
+    const unhandledRejections: unknown[] = [];
+    const onUnhandledRejection = (reason: unknown) => {
+      unhandledRejections.push(reason);
+    };
+    process.on("unhandledRejection", onUnhandledRejection);
+
+    const sidecar = testing.scheduleProviderAuthStatePrewarm({
+      getConfig: () => ({}) as never,
+      log,
+      startupWarmEnabled: false,
+    });
+
+    try {
+      await vi.dynamicImportSettled();
+      await waitForGatewayTestState(() => {
+        expect(hoisted.setAuthProfileFailureHook).toHaveBeenCalledOnce();
+      });
+      const failureHook = hoisted.setAuthProfileFailureHook.mock.calls[0]?.[0] as
+        | (() => void)
+        | undefined;
+      if (!failureHook) {
+        throw new Error("Expected provider auth failure hook to be registered");
+      }
+
+      failureHook();
+      markGatewayRestartDraining();
+      await vi.advanceTimersByTimeAsync(1_000);
+      await vi.dynamicImportSettled();
+
+      expect(hoisted.warmCurrentProviderAuthStateOffMainThread).not.toHaveBeenCalled();
+      expect(log.warn).not.toHaveBeenCalled();
+      expect(unhandledRejections).toStrictEqual([]);
+    } finally {
+      await sidecar.stop();
+      process.off("unhandledRejection", onUnhandledRejection);
+      resetGatewayWorkAdmission();
+      vi.useRealTimers();
+    }
+  });
+
+  it.each([
+    { label: "ordinary failure", error: new Error("provider warm failed") },
+    { label: "draining error outside restart", error: new GatewayDrainingError("not draining") },
+  ])("warns for a queued provider auth rewarm $label", async ({ error }) => {
+    vi.useFakeTimers();
+    const log = { info: vi.fn(), warn: vi.fn() };
+    hoisted.warmCurrentProviderAuthStateOffMainThread.mockRejectedValueOnce(error);
+    const sidecar = testing.scheduleProviderAuthStatePrewarm({
+      getConfig: () => ({}) as never,
+      log,
+      startupWarmEnabled: false,
+    });
+
+    try {
+      await vi.dynamicImportSettled();
+      await waitForGatewayTestState(() => {
+        expect(hoisted.setAuthProfileFailureHook).toHaveBeenCalledOnce();
+      });
+      const failureHook = hoisted.setAuthProfileFailureHook.mock.calls[0]?.[0] as
+        | (() => void)
+        | undefined;
+      if (!failureHook) {
+        throw new Error("Expected provider auth failure hook to be registered");
+      }
+
+      failureHook();
+      await vi.advanceTimersByTimeAsync(1_000);
+
+      expect(log.warn).toHaveBeenCalledWith(`provider auth state rewarm failed: ${String(error)}`);
+    } finally {
+      await sidecar.stop();
+      vi.useRealTimers();
+    }
+  });
+
   it("delays explicit provider auth prewarm beyond the early post-ready window", async () => {
     expect(testing.providerAuthPrewarmStartDelayMs).toBe(5_000);
   });

src/gateway/server-startup-post-attach.ts

@@ -22,7 +22,10 @@ import { getPluginModuleLoaderStats } from "../plugins/plugin-module-loader-cach
 import type { PluginRegistry } from "../plugins/registry.js";
 import { withPluginRuntimeRegistryScope } from "../plugins/runtime/gateway-request-scope.js";
 import type { PluginServicesHandle } from "../plugins/services.js";
-import { runWithGatewayIndependentRootWorkAdmission } from "../process/gateway-work-admission.js";
+import {
+  isGatewayRestartDrainError,
+  runWithGatewayIndependentRootWorkAdmission,
+} from "../process/gateway-work-admission.js";
 import { sweepSessionStateWatchNotices } from "../sessions/session-state-events.js";
 import { createDeferredCore } from "../shared/deferred.js";
 import { createLazyRuntimeModule } from "../shared/lazy-runtime.js";
@@ -149,6 +152,11 @@ function scheduleProviderAuthStatePrewarm(params: {
   let pendingRewarmReason: string | undefined;
   const isStopped = () => stopped;
   const delayMs = params.delayMs ?? PROVIDER_AUTH_PREWARM_START_DELAY_MS;
+  const logProviderAuthWarmFailure = (operation: string, error: unknown) => {
+    if (!isGatewayRestartDrainError(error)) {
+      params.log.warn(`provider auth state ${operation} failed: ${String(error)}`);
+    }
+  };
   void runWithGatewayIndependentRootWorkAdmission(async () => {
     const [{ setAuthProfileFailureHook }, { clearCurrentProviderAuthState }] = await Promise.all([
       import("../agents/auth-profiles/failure-hook.js"),
@@ -174,7 +182,7 @@ function scheduleProviderAuthStatePrewarm(params: {
             `provider auth state re-warmed (${reason}) ${formatProviderAuthWarmMetrics(metrics)}`,
           );
         } catch (err) {
-          params.log.warn(`provider auth state rewarm failed: ${String(err)}`);
+          logProviderAuthWarmFailure("rewarm", err);
         } finally {
           rewarmInFlight = false;
           const nextReason = pendingRewarmReason;
@@ -199,7 +207,9 @@ function scheduleProviderAuthStatePrewarm(params: {
         rewarmTimer = undefined;
         const nextReason = pendingRewarmReason ?? reason;
         pendingRewarmReason = undefined;
-        void runRewarm(nextReason);
+        void runRewarm(nextReason).catch((error: unknown) =>
+          logProviderAuthWarmFailure("rewarm", error),
+        );
       }, PROVIDER_AUTH_REWARM_DELAY_MS);
       rewarmTimer.unref?.();
     };
@@ -235,16 +245,12 @@ function scheduleProviderAuthStatePrewarm(params: {
           params.log.info(
             `provider auth state pre-warmed ${formatProviderAuthWarmMetrics(metrics)}`,
           );
-        }).catch((err: unknown) => {
-          params.log.warn(`provider auth state pre-warm failed: ${String(err)}`);
-        });
+        }).catch((error: unknown) => logProviderAuthWarmFailure("pre-warm", error));
       },
       Math.max(0, delayMs),
     );
     startupTimer.unref?.();
-  }).catch((err: unknown) => {
-    params.log.warn(`provider auth state pre-warm setup failed: ${String(err)}`);
-  });
+  }).catch((error: unknown) => logProviderAuthWarmFailure("pre-warm setup", error));
   return {
     stop: () => {
       stopped = true;

src/process/gateway-work-admission.test.ts

@@ -5,6 +5,7 @@ import {
   beginGatewayRootWorkAdmissionWhenOpen,
   GatewayDrainingError,
   getActiveGatewayRootWorkCount,
+  isGatewayRestartDrainError,
   isGatewaySubordinateWorkAdmissionClosed,
   isGatewayWorkAdmissionClosed,
   markGatewayRestartDraining,
@@ -23,6 +24,26 @@ import { runWithGatewayRootWorkAdmissionForTest } from "./gateway-work-admission
 beforeEach(resetGatewayWorkAdmission);
 afterEach(resetGatewayWorkAdmission);
 
+it("classifies draining errors only while an authoritative restart signal or drain is active", () => {
+  const error = new GatewayDrainingError();
+
+  expect(isGatewayRestartDrainError(error)).toBe(false);
+  expect(isGatewayRestartDrainError(new Error("GatewayDrainingError"))).toBe(false);
+
+  const suspension = tryBeginGatewaySuspendAdmission(() => {});
+  expect(isGatewayRestartDrainError(error)).toBe(false);
+  expect(suspension?.rollback()).toBe(true);
+
+  const signal = beginGatewayRestartSignalAdmission();
+  expect(isGatewayRestartDrainError(error)).toBe(true);
+  expect(isGatewayRestartDrainError(new Error("gateway is draining for restart"))).toBe(false);
+  expect(signal?.rollback()).toBe(true);
+  expect(isGatewayRestartDrainError(error)).toBe(false);
+
+  markGatewayRestartDraining();
+  expect(isGatewayRestartDrainError(error)).toBe(true);
+});
+
 it("counts one nested root chain once and excludes the preparing caller", async () => {
   const outer = tryBeginGatewayRootWorkAdmission();
   expect(outer).not.toBeNull();

src/process/gateway-work-admission.ts

@@ -179,6 +179,10 @@ export function isGatewayRestartDraining(): boolean {
   );
 }
 
+export function isGatewayRestartDrainError(error: unknown): error is GatewayDrainingError {
+  return error instanceof GatewayDrainingError && isGatewayRestartDraining();
+}
+
 /** Restart drain is one-way until the in-process restart resets runtime state. */
 export function markGatewayRestartDraining(): void {
   if (GATEWAY_WORK_ADMISSION_STATE.restartDraining) {