From df510d3b1df8cfb6dbbfd4dbfcfb0902fe412ed9 Mon Sep 17 00:00:00 2001 From: roboomp Date: Tue, 16 Jun 2026 04:49:58 +0000 Subject: [PATCH] fix(cli): sped up exit shutdown handlers Run session_shutdown extension handlers concurrently under the existing shutdown cap so /exit and /quit do not wait one full timeout per hanging extension. Fixes #2736 --- packages/coding-agent/CHANGELOG.md | 4 ++ .../src/extensibility/extensions/runner.ts | 18 +++++++- .../repro-issue-2600-shutdown-timeout.test.ts | 43 +++++++++++++++++-- 3 files changed, 60 insertions(+), 5 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index ce58d520d..d73b2ea4b 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed `/exit` and `/quit` waiting one shutdown timeout per hanging extension by running `session_shutdown` handlers within a shared shutdown window ([#2736](https://github.com/can1357/oh-my-pi/issues/2736)). + ## [16.0.1] - 2026-06-15 ### Breaking Changes diff --git a/packages/coding-agent/src/extensibility/extensions/runner.ts b/packages/coding-agent/src/extensibility/extensions/runner.ts index 5b6731e92..d94f34400 100644 --- a/packages/coding-agent/src/extensibility/extensions/runner.ts +++ b/packages/coding-agent/src/extensibility/extensions/runner.ts @@ -543,7 +543,9 @@ export class ExtensionRunner { event.type === "session_before_tree" ); } - + #isSessionShutdownEvent(event: RunnerEmitEvent): event is Extract { + return event.type === "session_shutdown"; + } async #runHandlerWithTimeout( handler: (event: TEvent, ctx: ExtensionContext) => Promise | TResult | undefined, event: TEvent, @@ -588,6 +590,20 @@ export class ExtensionRunner { const ctx = this.createContext(); let result: SessionBeforeEventResult | SessionCompactingResult | undefined; + if (this.#isSessionShutdownEvent(event)) { + const timeoutMs = handlerTimeoutForEvent(event.type); + const promises: Promise[] = []; + for (const ext of this.extensions) { + const handlers = ext.handlers.get(event.type); + if (!handlers || handlers.length === 0) continue; + for (const handler of handlers) { + promises.push(this.#runHandlerWithTimeout(handler, event, ctx, ext, timeoutMs)); + } + } + await Promise.all(promises); + return result as RunnerEmitResult; + } + for (const ext of this.extensions) { const handlers = ext.handlers.get(event.type); if (!handlers || handlers.length === 0) continue; diff --git a/packages/coding-agent/test/repro-issue-2600-shutdown-timeout.test.ts b/packages/coding-agent/test/repro-issue-2600-shutdown-timeout.test.ts index 2c7784a0a..027b7d98b 100644 --- a/packages/coding-agent/test/repro-issue-2600-shutdown-timeout.test.ts +++ b/packages/coding-agent/test/repro-issue-2600-shutdown-timeout.test.ts @@ -60,19 +60,27 @@ describe("issue #2600 - session_shutdown handler timeout", () => { testSetSessionShutdownHandlerTimeoutMs(SESSION_SHUTDOWN_HANDLER_TIMEOUT_MS); }); - async function buildRunnerWithHangingShutdown(): Promise<{ + async function buildRunnerWithHangingShutdown(count = 1): Promise<{ runner: ExtensionRunner; hangExtensionPath: string; + hangExtensionPaths: string[]; cleanup: () => void; }> { + if (count < 1) throw new Error("count must be positive"); const tempDir = TempDir.createSync("@pi-issue-2600-test-"); const extensionsDir = path.join(getProjectAgentDir(tempDir.path()), "extensions"); fs.mkdirSync(extensionsDir, { recursive: true }); - const hangExtensionPath = path.join(tempDir.path(), "hang-session-shutdown.ts"); - fs.writeFileSync(hangExtensionPath, HANG_EXTENSION_SRC); + const hangExtensionPaths: string[] = []; + for (let i = 0; i < count; i++) { + const hangExtensionPath = path.join(tempDir.path(), `hang-session-shutdown-${i}.ts`); + fs.writeFileSync(hangExtensionPath, HANG_EXTENSION_SRC); + hangExtensionPaths.push(hangExtensionPath); + } + const hangExtensionPath = hangExtensionPaths[0]; + if (!hangExtensionPath) throw new Error("missing hanging extension"); const sessionManager = SessionManager.inMemory(); - const result = await discoverAndLoadExtensions([extensionsDir, hangExtensionPath], tempDir.path()); + const result = await discoverAndLoadExtensions([extensionsDir, ...hangExtensionPaths], tempDir.path()); const runner = new ExtensionRunner( result.extensions, result.runtime, @@ -83,10 +91,37 @@ describe("issue #2600 - session_shutdown handler timeout", () => { return { runner, hangExtensionPath, + hangExtensionPaths, cleanup: () => tempDir.removeSync(), }; } + it("runs multiple session_shutdown handlers within one cap", async () => { + const { runner, hangExtensionPaths, cleanup } = await buildRunnerWithHangingShutdown(4); + const warnSpy = vi.spyOn(logger, "warn").mockImplementation(() => {}); + try { + testSetSessionShutdownHandlerTimeoutMs(100); + + const startedAt = performance.now(); + await runner.emit({ type: "session_shutdown" }); + const elapsedMs = performance.now() - startedAt; + + // Multiple hung shutdown handlers must share the cap. Sequential + // dispatch would consume roughly count × cap and keep `/exit` slow. + expect(elapsedMs).toBeLessThan(350); + for (const hangExtensionPath of hangExtensionPaths) { + expect(warnSpy).toHaveBeenCalledWith("Extension handler timed out", { + extensionPath: hangExtensionPath, + event: "session_shutdown", + timeoutMs: 100, + }); + } + } finally { + warnSpy.mockRestore(); + cleanup(); + } + }); + it("defaults the session_shutdown cap to ≤ 5s, never the generic 30s budget", () => { expect(SESSION_SHUTDOWN_HANDLER_TIMEOUT_MS).toBeLessThanOrEqual(5_000); expect(SESSION_SHUTDOWN_HANDLER_TIMEOUT_MS).toBeLessThan(EXTENSION_HANDLER_TIMEOUT_MS);