@@ -5,6 +5,9 @@ import { pickBeaconHost, pickGatewayPort } from "./discover.js";
55const acquireGatewayLock = vi.fn(async (_opts?: { port?: number }) => ({
66release: vi.fn(async () => {}),
77}));
8+const consumeGatewayRestartIntentPayloadSync = vi.fn<
9+() => { force?: boolean; waitMs?: number } | null
10+>(() => null);
811const consumeGatewaySigusr1RestartAuthorization = vi.fn(() => true);
912const consumeGatewayRestartIntentSync = vi.fn(() => false);
1013const isGatewaySigusr1RestartExternallyAllowed = vi.fn(() => false);
@@ -21,6 +24,17 @@ const scheduleGatewaySigusr1Restart = vi.fn((_opts?: { delayMs?: number; reason?
2124cooldownMsApplied: 0,
2225}));
2326const getActiveTaskCount = vi.fn(() => 0);
27+const getInspectableActiveTaskRestartBlockers = vi.fn(
28+() =>
29+[] as Array<{
30+taskId: string;
31+status: "queued" | "running";
32+runtime: "subagent" | "acp" | "cli" | "cron";
33+runId?: string;
34+label?: string;
35+title?: string;
36+}>,
37+);
2438const markGatewayDraining = vi.fn();
2539const waitForActiveTasks = vi.fn(async (_timeoutMs?: number) => ({ drained: true }));
2640const resetAllLanes = vi.fn();
@@ -64,6 +78,7 @@ vi.mock("../../infra/gateway-lock.js", () => ({
6478}));
65796680vi.mock("../../infra/restart.js", () => ({
81+consumeGatewayRestartIntentPayloadSync: () => consumeGatewayRestartIntentPayloadSync(),
6782consumeGatewaySigusr1RestartAuthorization: () => consumeGatewaySigusr1RestartAuthorization(),
6883consumeGatewayRestartIntentSync: () => consumeGatewayRestartIntentSync(),
6984isGatewaySigusr1RestartExternallyAllowed: () => isGatewaySigusr1RestartExternallyAllowed(),
@@ -103,6 +118,10 @@ vi.mock("../../tasks/runtime-internal.js", () => ({
103118reloadTaskRegistryFromStore: () => reloadTaskRegistryFromStore(),
104119}));
105120121+vi.mock("../../tasks/task-registry.maintenance.js", () => ({
122+getInspectableActiveTaskRestartBlockers: () => getInspectableActiveTaskRestartBlockers(),
123+}));
124+106125vi.mock("../../agents/pi-embedded-runner/runs.js", () => ({
107126abortEmbeddedPiRun: (sessionId?: string, opts?: { mode?: "all" | "compacting" }) =>
108127abortEmbeddedPiRun(sessionId, opts),
@@ -270,7 +289,7 @@ describe("runGatewayLoop", () => {
270289271290it("treats SIGTERM with a restart intent as a draining restart", async () => {
272291vi.clearAllMocks();
273-consumeGatewayRestartIntentSync.mockReturnValueOnce(true);
292+consumeGatewayRestartIntentPayloadSync.mockReturnValueOnce({});
274293getActiveTaskCount.mockReturnValueOnce(1).mockReturnValue(0);
275294276295await withIsolatedSignals(async ({ captureSignal }) => {
@@ -301,7 +320,7 @@ describe("runGatewayLoop", () => {
301320await new Promise<void>((resolve) => setImmediate(resolve));
302321await new Promise<void>((resolve) => setImmediate(resolve));
303322304-expect(consumeGatewayRestartIntentSync).toHaveBeenCalledOnce();
323+expect(consumeGatewayRestartIntentPayloadSync).toHaveBeenCalledOnce();
305324expect(markGatewayDraining).toHaveBeenCalledOnce();
306325expect(waitForActiveTasks).toHaveBeenCalledWith(90_000);
307326expect(closeFirst).toHaveBeenCalledWith({
@@ -321,6 +340,68 @@ describe("runGatewayLoop", () => {
321340});
322341});
323342343+it("uses restart intent wait overrides for SIGTERM drain", async () => {
344+vi.clearAllMocks();
345+consumeGatewayRestartIntentPayloadSync.mockReturnValueOnce({ waitMs: 2_500 });
346+getActiveTaskCount.mockReturnValueOnce(1).mockReturnValue(0);
347+348+await withIsolatedSignals(async ({ captureSignal }) => {
349+const { start, exited } = await createSignaledLoopHarness();
350+const sigterm = captureSignal("SIGTERM");
351+const sigint = captureSignal("SIGINT");
352+353+sigterm();
354+await new Promise<void>((resolve) => setImmediate(resolve));
355+await new Promise<void>((resolve) => setImmediate(resolve));
356+357+expect(waitForActiveTasks).toHaveBeenCalledWith(2_500);
358+expect(start).toHaveBeenCalledTimes(2);
359+360+sigint();
361+await expect(exited).resolves.toBe(0);
362+});
363+});
364+365+it("forces SIGTERM restarts without waiting for active task drain", async () => {
366+vi.clearAllMocks();
367+consumeGatewayRestartIntentPayloadSync.mockReturnValueOnce({ force: true });
368+getActiveTaskCount.mockReturnValueOnce(1).mockReturnValue(0);
369+getActiveEmbeddedRunCount.mockReturnValueOnce(1).mockReturnValue(0);
370+getInspectableActiveTaskRestartBlockers.mockReturnValueOnce([
371+{
372+taskId: "task-force",
373+runId: "run-force",
374+status: "running",
375+runtime: "cron",
376+label: "forced",
377+},
378+]);
379+380+await withIsolatedSignals(async ({ captureSignal }) => {
381+const { start, exited } = await createSignaledLoopHarness();
382+const sigterm = captureSignal("SIGTERM");
383+const sigint = captureSignal("SIGINT");
384+385+sigterm();
386+await new Promise<void>((resolve) => setImmediate(resolve));
387+await new Promise<void>((resolve) => setImmediate(resolve));
388+389+expect(waitForActiveTasks).not.toHaveBeenCalled();
390+expect(waitForActiveEmbeddedRuns).not.toHaveBeenCalled();
391+expect(abortEmbeddedPiRun).toHaveBeenCalledWith(undefined, { mode: "all" });
392+expect(gatewayLog.warn).toHaveBeenCalledWith(
393+expect.stringContaining("restart blocked by active task run(s): taskId=task-force"),
394+);
395+expect(gatewayLog.warn).toHaveBeenCalledWith(
396+"forced restart requested; skipping active work drain",
397+);
398+expect(start).toHaveBeenCalledTimes(2);
399+400+sigint();
401+await expect(exited).resolves.toBe(0);
402+});
403+});
404+324405it("restarts after SIGUSR1 even when drain times out, and resets runtime state for the new iteration", async () => {
325406vi.clearAllMocks();
326407loadConfig.mockReturnValue({
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。