From e8b4e39a97bc7949c396272eec8d33ac1bcff99e Mon Sep 17 00:00:00 2001 From: Vincent Koc Date: Mon, 27 Apr 2026 21:19:21 -0700 Subject: [PATCH] fix(gateway): clear fallback context on close Fixes gateway fallback request context cleanup on close/startup failure and shards the full gateway Vitest lane to avoid the observed memory hang.\n\nValidation:\n- Testbox: OPENCLAW_TESTBOX=1 pnpm check:changed\n- Testbox: env OPENCLAW_VITEST_MAX_WORKERS=1 /usr/bin/time -v pnpm test:gateway (254 files, 2950 tests, max RSS 4144692 KB) --- package.json | 4 +- src/gateway/server-plugins.lifecycle.test.ts | 28 +++++++++++ src/gateway/server-plugins.test.ts | 49 ++++++++++++++++++++ src/gateway/server-plugins.ts | 27 ++++++++++- src/gateway/server-restart-sentinel.test.ts | 3 ++ src/gateway/server.impl.ts | 39 +++++++++++----- src/infra/vitest-config.test.ts | 2 +- test/vitest/vitest.gateway.config.ts | 16 ++++++- 8 files changed, 151 insertions(+), 17 deletions(-) create mode 100644 src/gateway/server-plugins.lifecycle.test.ts diff --git a/package.json b/package.json index e1e8e9e6ce4..01e25a5e0e4 100644 --- a/package.json +++ b/package.json @@ -1533,7 +1533,7 @@ "test:extensions:package-boundary:compile": "node scripts/check-extension-package-tsc-boundary.mjs --mode=compile", "test:fast": "node scripts/run-vitest.mjs run --config test/vitest/vitest.unit.config.ts", "test:force": "node --import tsx scripts/test-force.ts", - "test:gateway": "node scripts/run-vitest.mjs run --config test/vitest/vitest.gateway.config.ts", + "test:gateway": "OPENCLAW_GATEWAY_PROJECT_SHARDS=1 node scripts/run-vitest.mjs run --config test/vitest/vitest.gateway.config.ts", "test:gateway:watch-regression": "node scripts/check-gateway-watch-regression.mjs", "test:install:e2e": "bash scripts/test-install-sh-e2e-docker.sh", "test:install:e2e:anthropic": "OPENCLAW_E2E_MODELS=anthropic bash scripts/test-install-sh-e2e-docker.sh", @@ -1564,7 +1564,7 @@ "test:perf:imports:changed": "OPENCLAW_VITEST_IMPORT_DURATIONS=1 OPENCLAW_VITEST_PRINT_IMPORT_BREAKDOWN=1 node scripts/test-projects.mjs --changed origin/main", "test:perf:profile:main": "node scripts/run-vitest-profile.mjs main", "test:perf:profile:runner": "node scripts/run-vitest-profile.mjs runner", - "test:sectriage": "node scripts/run-vitest.mjs run --config test/vitest/vitest.gateway.config.ts && node scripts/run-vitest.mjs run --config test/vitest/vitest.unit.config.ts --exclude src/daemon/launchd.integration.test.ts --exclude src/process/exec.test.ts", + "test:sectriage": "OPENCLAW_GATEWAY_PROJECT_SHARDS=1 node scripts/run-vitest.mjs run --config test/vitest/vitest.gateway.config.ts && node scripts/run-vitest.mjs run --config test/vitest/vitest.unit.config.ts --exclude src/daemon/launchd.integration.test.ts --exclude src/process/exec.test.ts", "test:serial": "OPENCLAW_TEST_PROJECTS_SERIAL=1 OPENCLAW_VITEST_MAX_WORKERS=1 node scripts/test-projects.mjs", "test:stability:gateway": "OPENCLAW_VITEST_MAX_WORKERS=1 node scripts/run-vitest.mjs run --config test/vitest/vitest.gateway.config.ts src/gateway/gateway-stability.test.ts && OPENCLAW_VITEST_MAX_WORKERS=1 node scripts/run-vitest.mjs run --config test/vitest/vitest.logging.config.ts src/logging/diagnostic-stability-bundle.test.ts && OPENCLAW_VITEST_MAX_WORKERS=1 node scripts/run-vitest.mjs run --config test/vitest/vitest.infra.config.ts src/infra/fatal-error-hooks.test.ts", "test:startup:bench": "node --import tsx scripts/bench-cli-startup.ts", diff --git a/src/gateway/server-plugins.lifecycle.test.ts b/src/gateway/server-plugins.lifecycle.test.ts new file mode 100644 index 00000000000..96b2532b066 --- /dev/null +++ b/src/gateway/server-plugins.lifecycle.test.ts @@ -0,0 +1,28 @@ +import { afterEach, describe, expect, it } from "vitest"; +import { clearFallbackGatewayContext, createGatewaySubagentRuntime } from "./server-plugins.js"; +import { installGatewayTestHooks, startServer } from "./test-helpers.server.js"; + +installGatewayTestHooks(); + +afterEach(() => { + clearFallbackGatewayContext(); +}); + +describe("gateway plugin fallback context lifecycle", () => { + it("clears the fallback gateway context after server close", async () => { + const runtime = createGatewaySubagentRuntime(); + const started = await startServer(); + + try { + await expect( + runtime.getSessionMessages({ sessionKey: "agent:main:main", limit: 1 }), + ).resolves.toEqual({ messages: [] }); + } finally { + await started.server.close({ reason: "fallback context lifecycle test done" }); + } + + await expect( + runtime.getSessionMessages({ sessionKey: "agent:main:main", limit: 1 }), + ).rejects.toThrow("No scope set and no fallback context available"); + }); +}); diff --git a/src/gateway/server-plugins.test.ts b/src/gateway/server-plugins.test.ts index f6e54ad3682..ac7fd771bed 100644 --- a/src/gateway/server-plugins.test.ts +++ b/src/gateway/server-plugins.test.ts @@ -344,6 +344,7 @@ beforeEach(() => { }); afterEach(() => { + serverPluginsModule.clearFallbackGatewayContext(); runtimeModule.clearGatewaySubagentRuntime(); runtimeRegistryModule.resetPluginRuntimeStateForTest(); }); @@ -1388,4 +1389,52 @@ describe("loadGatewayPlugins", () => { await runtime.run({ sessionKey: "s-5", message: "prefer resolver" }); expect(getLastDispatchedContext()).toBe(freshContext); }); + + test("clears fallback context snapshots when a resolver is registered", async () => { + const serverPlugins = serverPluginsModule; + const runtime = await createSubagentRuntime(serverPlugins); + const staleContext = createTestContext("stale-snapshot"); + + serverPlugins.setFallbackGatewayContext(staleContext); + serverPlugins.setFallbackGatewayContextResolver(() => undefined); + + await expect(runtime.run({ sessionKey: "s-6", message: "stale fallback" })).rejects.toThrow( + "No scope set and no fallback context available", + ); + }); + + test("clears fallback context and resolver state", async () => { + const serverPlugins = serverPluginsModule; + const runtime = await createSubagentRuntime(serverPlugins); + const context = createTestContext("clear-context"); + + serverPlugins.setFallbackGatewayContextResolver(() => context); + await runtime.run({ sessionKey: "s-7", message: "before clear" }); + expect(getLastDispatchedContext()).toBe(context); + + serverPlugins.clearFallbackGatewayContext(); + + await expect(runtime.run({ sessionKey: "s-7", message: "after clear" })).rejects.toThrow( + "No scope set and no fallback context available", + ); + }); + + test("resolver cleanup only clears the resolver it registered", async () => { + const serverPlugins = serverPluginsModule; + const runtime = await createSubagentRuntime(serverPlugins); + const firstContext = createTestContext("first-owner"); + const secondContext = createTestContext("second-owner"); + + const clearFirst = serverPlugins.setFallbackGatewayContextResolver(() => firstContext); + const clearSecond = serverPlugins.setFallbackGatewayContextResolver(() => secondContext); + + clearFirst(); + await runtime.run({ sessionKey: "s-8", message: "after first cleanup" }); + expect(getLastDispatchedContext()).toBe(secondContext); + + clearSecond(); + await expect( + runtime.run({ sessionKey: "s-8", message: "after second cleanup" }), + ).rejects.toThrow("No scope set and no fallback context available"); + }); }); diff --git a/src/gateway/server-plugins.ts b/src/gateway/server-plugins.ts index ca87d0e7225..96e6ce85c4b 100644 --- a/src/gateway/server-plugins.ts +++ b/src/gateway/server-plugins.ts @@ -44,17 +44,40 @@ const getFallbackGatewayContextState = () => resolveContext: undefined, })); -export function setFallbackGatewayContext(ctx: GatewayRequestContext): void { +export function setFallbackGatewayContext(ctx: GatewayRequestContext): () => void { const fallbackGatewayContextState = getFallbackGatewayContextState(); fallbackGatewayContextState.context = ctx; fallbackGatewayContextState.resolveContext = undefined; + return () => { + const currentFallbackGatewayContextState = getFallbackGatewayContextState(); + if ( + currentFallbackGatewayContextState.context === ctx && + currentFallbackGatewayContextState.resolveContext === undefined + ) { + currentFallbackGatewayContextState.context = undefined; + } + }; } export function setFallbackGatewayContextResolver( resolveContext: () => GatewayRequestContext | undefined, -): void { +): () => void { const fallbackGatewayContextState = getFallbackGatewayContextState(); + fallbackGatewayContextState.context = undefined; fallbackGatewayContextState.resolveContext = resolveContext; + return () => { + const currentFallbackGatewayContextState = getFallbackGatewayContextState(); + if (currentFallbackGatewayContextState.resolveContext === resolveContext) { + currentFallbackGatewayContextState.context = undefined; + currentFallbackGatewayContextState.resolveContext = undefined; + } + }; +} + +export function clearFallbackGatewayContext(): void { + const fallbackGatewayContextState = getFallbackGatewayContextState(); + fallbackGatewayContextState.context = undefined; + fallbackGatewayContextState.resolveContext = undefined; } function getFallbackGatewayContext(): GatewayRequestContext | undefined { diff --git a/src/gateway/server-restart-sentinel.test.ts b/src/gateway/server-restart-sentinel.test.ts index cb3288f62c3..f5cc8ff1980 100644 --- a/src/gateway/server-restart-sentinel.test.ts +++ b/src/gateway/server-restart-sentinel.test.ts @@ -121,6 +121,9 @@ const mocks = vi.hoisted(() => { }; }); +vi.unmock("./server-restart-sentinel.js"); +vi.resetModules(); + vi.mock("../agents/agent-scope.js", () => ({ resolveSessionAgentId: mocks.resolveSessionAgentId, })); diff --git a/src/gateway/server.impl.ts b/src/gateway/server.impl.ts index 971768f9c23..cfd8756092a 100644 --- a/src/gateway/server.impl.ts +++ b/src/gateway/server.impl.ts @@ -726,9 +726,14 @@ export async function startGatewayServer( httpServers, })(opts); }; + let clearFallbackGatewayContextForServer = () => {}; const closeOnStartupFailure = async () => { - await runClosePrelude(); - await createCloseHandler()({ reason: "gateway startup failed" }); + try { + await runClosePrelude(); + await createCloseHandler()({ reason: "gateway startup failed" }); + } finally { + clearFallbackGatewayContextForServer(); + } }; const broadcastVoiceWakeRoutingChanged = (config: VoiceWakeRoutingConfig) => { broadcast("voicewake.routing.changed", { config }, { dropIfSlow: true }); @@ -888,7 +893,15 @@ export async function startGatewayServer( broadcastVoiceWakeRoutingChanged, }); - setFallbackGatewayContextResolver(() => gatewayRequestContext); + const fallbackGatewayContextCleanup: unknown = setFallbackGatewayContextResolver( + () => gatewayRequestContext, + ); + clearFallbackGatewayContextForServer = + typeof fallbackGatewayContextCleanup === "function" + ? () => { + fallbackGatewayContextCleanup(); + } + : () => {}; if (!minimalTestGateway) { if (deferredConfiguredChannelPluginIds.length > 0) { @@ -1039,14 +1052,18 @@ export async function startGatewayServer( return { close: async (opts) => { - // Run gateway_stop plugin hook before shutdown - await runGlobalGatewayStopSafely({ - event: { reason: opts?.reason ?? "gateway stopping" }, - ctx: { port }, - onError: (err) => log.warn(`gateway_stop hook failed: ${String(err)}`), - }); - await runClosePrelude(); - await close(opts); + try { + // Run gateway_stop plugin hook before shutdown + await runGlobalGatewayStopSafely({ + event: { reason: opts?.reason ?? "gateway stopping" }, + ctx: { port }, + onError: (err) => log.warn(`gateway_stop hook failed: ${String(err)}`), + }); + await runClosePrelude(); + await close(opts); + } finally { + clearFallbackGatewayContextForServer(); + } }, }; } diff --git a/src/infra/vitest-config.test.ts b/src/infra/vitest-config.test.ts index 2d7b9868b48..d9903a92058 100644 --- a/src/infra/vitest-config.test.ts +++ b/src/infra/vitest-config.test.ts @@ -237,7 +237,7 @@ describe("test scripts", () => { expect(pkg.scripts?.["test"]).toBe("node scripts/test-projects.mjs"); expect(pkg.scripts?.["test:force"]).toBe("node --import tsx scripts/test-force.ts"); expect(pkg.scripts?.["test:gateway"]).toBe( - "node scripts/run-vitest.mjs run --config test/vitest/vitest.gateway.config.ts", + "OPENCLAW_GATEWAY_PROJECT_SHARDS=1 node scripts/run-vitest.mjs run --config test/vitest/vitest.gateway.config.ts", ); expect(pkg.scripts?.["test:single"]).toBeUndefined(); }); diff --git a/test/vitest/vitest.gateway.config.ts b/test/vitest/vitest.gateway.config.ts index 78398facb56..ba2cc6c47d4 100644 --- a/test/vitest/vitest.gateway.config.ts +++ b/test/vitest/vitest.gateway.config.ts @@ -1,5 +1,13 @@ +import { createProjectShardVitestConfig } from "./vitest.project-shard-config.ts"; import { createScopedVitestConfig } from "./vitest.scoped-config.ts"; +const gatewayProjectConfigs = [ + "test/vitest/vitest.gateway-core.config.ts", + "test/vitest/vitest.gateway-client.config.ts", + "test/vitest/vitest.gateway-methods.config.ts", + "test/vitest/vitest.gateway-server.config.ts", +] as const; + export function createGatewayVitestConfig(env?: Record) { return createScopedVitestConfig(["src/gateway/**/*.test.ts"], { dir: "src/gateway", @@ -13,4 +21,10 @@ export function createGatewayVitestConfig(env?: Record