From 03489d1ebee6023a3af3da7c117b1577951fbd70 Mon Sep 17 00:00:00 2001 From: roboomp Date: Thu, 23 Jul 2026 19:00:08 +0000 Subject: [PATCH] fix(eval): coalesced python kernel replacement Tracked retained Python kernel generations and shared one replacement promise per dead generation. Reset and disposal now invalidate and drain replacement work before allowing a new session to take ownership. Added deterministic fake-kernel coverage for concurrent callers, cancellation, reset, owner/global disposal, and independent cwd keys. Fixes #6367 --- packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/src/eval/py/executor.ts | 136 ++++++++--- .../core/python-executor.lifecycle.test.ts | 223 +++++++++++++++++- 3 files changed, 326 insertions(+), 34 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index ad5f374bd..a4a2cb22e 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -4,6 +4,7 @@ ### Fixed +- Fixed concurrent dead-kernel replacement in persistent Python sessions starting multiple generations and orphaning the losing process; callers now share one generation-scoped replacement, while resets and disposal invalidate and drain stale replacements ([#6367](https://github.com/can1357/oh-my-pi/issues/6367)). - Fixed credential-shaped tokens (GitHub/GitLab/OpenAI/Anthropic key patterns) being redacted from outbound provider requests even with `secrets.enabled` off; the pattern redaction now follows the `secrets.enabled` ("Hide Secrets") setting like the secret obfuscator. - Fixed Ctrl-clicking a wrapped OAuth authorization URL opening only the clicked row's truncated fragment by preserving the complete hyperlink target on every rendered row. - Fixed used-only absolute usage amounts across output surfaces: CLI now renders `$123.45 used`; the TUI shows a neutral, width-bounded amount instead of a pending/dotted/account-count placeholder; and ACP preserves `123.45 usd used` while suppressing duplicate window suffixes such as `— extra`. ([#5575](https://github.com/can1357/oh-my-pi/issues/5575)) diff --git a/packages/coding-agent/src/eval/py/executor.ts b/packages/coding-agent/src/eval/py/executor.ts index fa44c074b..806e99a98 100644 --- a/packages/coding-agent/src/eval/py/executor.ts +++ b/packages/coding-agent/src/eval/py/executor.ts @@ -21,6 +21,7 @@ import { type KernelDisplayOutput, type KernelExecuteOptions, type KernelExecuteResult, + type KernelShutdownResult, PythonKernel, } from "./kernel"; import { resolveExplicitPythonRuntime } from "./runtime"; @@ -136,11 +137,19 @@ export interface PythonResult { // the same tuple; the kernel stays alive until the last owner detaches. // --------------------------------------------------------------------------- +interface SessionKernelReplacement { + generation: number; + deadlineMs?: number; + promise: Promise; +} + interface PythonSession { sessionKey: string; sessionId: string; cwd: string; kernel: PythonKernel; + generation: number; + replacement?: SessionKernelReplacement; ownerIds: Set; hasFallbackOwner: boolean; } @@ -254,6 +263,7 @@ async function acquireSession( sessionId, cwd, kernel, + generation: 0, ownerIds: new Set(), hasFallbackOwner: false, }; @@ -272,31 +282,99 @@ async function acquireSession( async function replaceSessionKernel( session: PythonSession, + kernel: PythonKernel, + generation: number, cwd: string, options: PythonExecutorOptions, -): Promise { - const old = session.kernel; - const remaining = getRemainingTimeoutMs(options.deadlineMs); - await old - .shutdown(remaining !== undefined ? { timeoutMs: Math.max(0, remaining) } : undefined) - .catch(() => undefined); - if (sessions.get(session.sessionKey) !== session) { +): Promise { + const inFlight = session.replacement; + if (inFlight?.generation === generation) { + if ( + inFlight.deadlineMs !== undefined && + (options.deadlineMs === undefined || options.deadlineMs > inFlight.deadlineMs) + ) { + inFlight.deadlineMs = options.deadlineMs; + } + return await waitForPromiseWithCancellation(inFlight.promise, options, PythonExecutionCancelledError); + } + if (sessions.get(session.sessionKey) !== session || session.generation !== generation || session.kernel !== kernel) { throw new PythonExecutionCancelledError(false); } - requireRemainingTimeoutMs(options.deadlineMs); - const next = await startKernel(cwd, options); - if (sessions.get(session.sessionKey) !== session) { - await next.shutdown().catch(() => undefined); - throw new PythonExecutionCancelledError(false); + + const deferred = Promise.withResolvers(); + const replacement: SessionKernelReplacement = { + generation, + deadlineMs: options.deadlineMs, + promise: deferred.promise, + }; + session.replacement = replacement; + void (async () => { + try { + const remaining = getRemainingTimeoutMs(options.deadlineMs); + await kernel + .shutdown(remaining !== undefined ? { timeoutMs: Math.max(0, remaining) } : undefined) + .catch(() => undefined); + if (replacement.deadlineMs !== undefined && replacement.deadlineMs <= Date.now()) { + throw new PythonExecutionCancelledError(true); + } + if ( + sessions.get(session.sessionKey) !== session || + session.generation !== generation || + session.kernel !== kernel + ) { + throw new PythonExecutionCancelledError(false); + } + const next = await startKernel(cwd, { + ...options, + signal: undefined, + deadlineMs: undefined, + }); + if ( + sessions.get(session.sessionKey) !== session || + session.generation !== generation || + session.kernel !== kernel + ) { + await next.shutdown().catch(() => undefined); + throw new PythonExecutionCancelledError(false); + } + session.kernel = next; + session.generation += 1; + deferred.resolve(next); + } catch (err) { + deferred.reject(err); + } finally { + if (session.replacement === replacement) session.replacement = undefined; + } + })(); + return await waitForPromiseWithCancellation(deferred.promise, options, PythonExecutionCancelledError); +} + +async function shutdownInvalidatedSession(session: PythonSession): Promise { + const replacement = session.replacement; + if (replacement) await replacement.promise.catch(() => undefined); + return await session.kernel.shutdown(); +} + +async function acquireLiveSessionKernel( + session: PythonSession, + cwd: string, + options: PythonExecutorOptions, +): Promise { + while (sessions.get(session.sessionKey) === session) { + const kernel = session.kernel; + const generation = session.generation; + if (kernel.isAlive()) return kernel; + await replaceSessionKernel(session, kernel, generation, cwd, options); } - session.kernel = next; + throw new PythonExecutionCancelledError(false); } async function resetSession(sessionKey: string): Promise { const existing = sessions.get(sessionKey) ?? (await startingSessions.get(sessionKey)?.catch(() => undefined)); if (!existing) return; + existing.generation += 1; sessions.delete(sessionKey); - await existing.kernel.shutdown().catch(() => undefined); + await shutdownInvalidatedSession(existing).catch(() => undefined); } // --------------------------------------------------------------------------- @@ -315,9 +393,10 @@ export async function disposeAllKernelSessions(): Promise { } } for (const [id, session] of all) { + session.generation += 1; if (sessions.get(id) === session) sessions.delete(id); } - const results = await Promise.allSettled(all.map(([, session]) => session.kernel.shutdown())); + const results = await Promise.allSettled(all.map(([, session]) => shutdownInvalidatedSession(session))); for (let i = 0; i < all.length; i += 1) { const [id, session] = all[i]; const result = results[i]; @@ -344,9 +423,10 @@ export async function disposeKernelSessionsByOwner(ownerId: string): Promise session.kernel.shutdown())); + const results = await Promise.allSettled(toShutdown.map(session => shutdownInvalidatedSession(session))); for (let i = 0; i < toShutdown.length; i += 1) { const session = toShutdown[i]; const result = results[i]; @@ -458,31 +538,21 @@ async function executeOnSession(code: string, cwd: string, options: PythonExecut isTimedOutCancellation(options.signal.reason, PythonExecutionCancelledError, options.signal), ); } - if (sessions.get(session.sessionKey) !== session) { + const kernel = await acquireLiveSessionKernel(session, cwd, options); + if (sessions.get(session.sessionKey) !== session || session.kernel !== kernel) { throw new PythonExecutionCancelledError(false); } - if (!session.kernel.isAlive()) { - await replaceSessionKernel(session, cwd, options); - if (sessions.get(session.sessionKey) !== session) { - throw new PythonExecutionCancelledError(false); - } - } const runOptions = { ...options, cwd }; try { - return await executeWithKernel(session.kernel, code, runOptions); + return await executeWithKernel(kernel, code, runOptions); } catch (err) { if (isCancellationError(err, PythonExecutionCancelledError) || options.signal?.aborted) throw err; - if (session.kernel.isAlive()) throw err; - if (sessions.get(session.sessionKey) !== session) { + if (kernel.isAlive()) throw err; + const retryKernel = await acquireLiveSessionKernel(session, cwd, options); + if (sessions.get(session.sessionKey) !== session || session.kernel !== retryKernel) { throw new PythonExecutionCancelledError(false); } - // Shared kernels are keyed by cwd, so a dead kernel can be recreated in place - // without risking cross-directory state bleed. - await replaceSessionKernel(session, cwd, options); - if (sessions.get(session.sessionKey) !== session) { - throw new PythonExecutionCancelledError(false); - } - return await executeWithKernel(session.kernel, code, runOptions); + return await executeWithKernel(retryKernel, code, runOptions); } } diff --git a/packages/coding-agent/test/core/python-executor.lifecycle.test.ts b/packages/coding-agent/test/core/python-executor.lifecycle.test.ts index b60e2d9ba..5e52e7058 100644 --- a/packages/coding-agent/test/core/python-executor.lifecycle.test.ts +++ b/packages/coding-agent/test/core/python-executor.lifecycle.test.ts @@ -1,5 +1,9 @@ import { afterEach, describe, expect, it } from "bun:test"; -import { disposeAllKernelSessions, executePython } from "@oh-my-pi/pi-coding-agent/eval/py/executor"; +import { + disposeAllKernelSessions, + disposeKernelSessionsByOwner, + executePython, +} from "@oh-my-pi/pi-coding-agent/eval/py/executor"; import { type KernelExecuteOptions, type KernelExecuteResult, @@ -29,6 +33,10 @@ class FakeKernel { return this.#alive; } + markDead(): void { + this.#alive = false; + } + async execute(code: string, options?: KernelExecuteOptions): Promise { this.executeCalls.push(code); this.#onExecute?.(options); @@ -53,6 +61,10 @@ const okResult: KernelExecuteResult = { stdinRequested: false, }; +async function flushMicrotasks(): Promise { + for (let i = 0; i < 5; i += 1) await Promise.resolve(); +} + describe("executePython session lifecycle", () => { const originalStart = PythonKernel.start; @@ -95,6 +107,215 @@ describe("executePython session lifecycle", () => { expect(liveKernel.executeCalls).toEqual(["print('restart')"]); }); + it("coalesces concurrent replacement of one dead session generation", async () => { + const deadKernel = new FakeKernel(okResult); + const replacementOne = new FakeKernel(okResult); + const replacementTwo = new FakeKernel(okResult); + const replacementStarted = Promise.withResolvers(); + const releaseReplacement = Promise.withResolvers(); + const replacements = [replacementOne, replacementTwo]; + let startCount = 0; + + PythonKernel.start = async () => { + startCount += 1; + if (startCount === 1) return deadKernel as unknown as PythonKernel; + replacementStarted.resolve(); + await releaseReplacement.promise; + return replacements.shift() as unknown as PythonKernel; + }; + + await executePython("print('setup')", { sessionId: "session-concurrent-restart" }); + deadKernel.markDead(); + + const first = executePython("print('first')", { sessionId: "session-concurrent-restart" }); + await replacementStarted.promise; + const second = executePython("print('second')", { sessionId: "session-concurrent-restart" }); + await Promise.resolve(); + await Promise.resolve(); + + expect(startCount).toBe(2); + + releaseReplacement.resolve(); + await Promise.all([first, second]); + + expect(replacementOne.executeCalls).toEqual(["print('first')", "print('second')"]); + expect(replacementTwo.executeCalls).toEqual([]); + + await disposeAllKernelSessions(); + expect(replacementOne.shutdownCalls).toBe(1); + expect(replacementTwo.shutdownCalls).toBe(0); + }); + + it("keeps a shared replacement alive when one caller cancels", async () => { + const deadKernel = new FakeKernel(okResult); + const replacement = new FakeKernel(okResult); + const replacementStarted = Promise.withResolvers(); + const releaseReplacement = Promise.withResolvers(); + let startCount = 0; + + PythonKernel.start = async () => { + startCount += 1; + if (startCount === 1) return deadKernel as unknown as PythonKernel; + replacementStarted.resolve(); + await releaseReplacement.promise; + return replacement as unknown as PythonKernel; + }; + + await executePython("print('setup')", { sessionId: "session-cancelled-restart" }); + deadKernel.markDead(); + + const abortController = new AbortController(); + const cancelled = executePython("print('cancelled')", { + sessionId: "session-cancelled-restart", + signal: abortController.signal, + }); + await replacementStarted.promise; + const retained = executePython("print('retained')", { sessionId: "session-cancelled-restart" }); + await flushMicrotasks(); + abortController.abort(Object.assign(new Error("replacement wait cancelled"), { name: "AbortError" })); + + expect((await cancelled).cancelled).toBe(true); + expect(startCount).toBe(2); + + releaseReplacement.resolve(); + expect((await retained).cancelled).toBe(false); + expect(replacement.executeCalls).toEqual(["print('retained')"]); + }); + + it("invalidates an in-flight replacement before resetting to a fresh generation", async () => { + const deadKernel = new FakeKernel(okResult); + const staleReplacement = new FakeKernel(okResult); + const freshKernel = new FakeKernel(okResult); + const replacementStarted = Promise.withResolvers(); + const releaseReplacement = Promise.withResolvers(); + let startCount = 0; + + PythonKernel.start = async () => { + startCount += 1; + if (startCount === 1) return deadKernel as unknown as PythonKernel; + if (startCount === 2) { + replacementStarted.resolve(); + await releaseReplacement.promise; + return staleReplacement as unknown as PythonKernel; + } + return freshKernel as unknown as PythonKernel; + }; + + await executePython("print('setup')", { sessionId: "session-reset-replacement" }); + deadKernel.markDead(); + + const obsolete = executePython("print('obsolete')", { sessionId: "session-reset-replacement" }); + await replacementStarted.promise; + const reset = executePython("print('reset')", { + sessionId: "session-reset-replacement", + reset: true, + }); + await flushMicrotasks(); + releaseReplacement.resolve(); + + expect((await obsolete).cancelled).toBe(true); + expect((await reset).cancelled).toBe(false); + expect(staleReplacement.executeCalls).toEqual([]); + expect(staleReplacement.shutdownCalls).toBe(1); + expect(freshKernel.executeCalls).toEqual(["print('reset')"]); + + await executePython("print('later')", { sessionId: "session-reset-replacement" }); + expect(startCount).toBe(3); + expect(freshKernel.executeCalls).toEqual(["print('reset')", "print('later')"]); + }); + + it("drains replacements invalidated by owner and global disposal", async () => { + const ownerKernel = new FakeKernel(okResult); + const globalKernel = new FakeKernel(okResult); + const ownerReplacement = new FakeKernel(okResult); + const globalReplacement = new FakeKernel(okResult); + const replacementsStarted = Promise.withResolvers(); + const releaseReplacements = Promise.withResolvers(); + const initialKernels = [ownerKernel, globalKernel]; + const replacementKernels = [ownerReplacement, globalReplacement]; + let replacementStartCount = 0; + + PythonKernel.start = async () => { + const initial = initialKernels.shift(); + if (initial) return initial as unknown as PythonKernel; + replacementStartCount += 1; + if (replacementStartCount === 2) replacementsStarted.resolve(); + await releaseReplacements.promise; + return replacementKernels.shift() as unknown as PythonKernel; + }; + + await executePython("print('owner setup')", { + sessionId: "session-owner-disposal-replacement", + kernelOwnerId: "replacement-owner", + }); + await executePython("print('global setup')", { sessionId: "session-global-disposal-replacement" }); + ownerKernel.markDead(); + globalKernel.markDead(); + + const ownerExecution = executePython("print('owner obsolete')", { + sessionId: "session-owner-disposal-replacement", + kernelOwnerId: "replacement-owner", + }); + const globalExecution = executePython("print('global obsolete')", { + sessionId: "session-global-disposal-replacement", + }); + await replacementsStarted.promise; + + const ownerDisposal = disposeKernelSessionsByOwner("replacement-owner"); + await flushMicrotasks(); + const globalDisposal = disposeAllKernelSessions(); + await flushMicrotasks(); + releaseReplacements.resolve(); + + expect((await ownerExecution).cancelled).toBe(true); + expect((await globalExecution).cancelled).toBe(true); + await Promise.all([ownerDisposal, globalDisposal]); + expect(ownerReplacement.executeCalls).toEqual([]); + expect(globalReplacement.executeCalls).toEqual([]); + expect(ownerReplacement.shutdownCalls).toBe(1); + expect(globalReplacement.shutdownCalls).toBe(1); + }); + + it("keeps replacement coordination independent across normalized cwd keys", async () => { + const deadOne = new FakeKernel(okResult); + const deadTwo = new FakeKernel(okResult); + const replacementOne = new FakeKernel(okResult); + const replacementTwo = new FakeKernel(okResult); + const kernels = [deadOne, deadTwo, replacementOne, replacementTwo]; + let startCount = 0; + + PythonKernel.start = async () => { + startCount += 1; + return kernels.shift() as unknown as PythonKernel; + }; + + await executePython("print('setup one')", { + cwd: "/tmp/replacement-key-one", + sessionId: "session-independent-replacement", + }); + await executePython("print('setup two')", { + cwd: "/tmp/replacement-key-two", + sessionId: "session-independent-replacement", + }); + deadOne.markDead(); + deadTwo.markDead(); + + await Promise.all([ + executePython("print('one')", { + cwd: "/tmp/replacement-key-one", + sessionId: "session-independent-replacement", + }), + executePython("print('two')", { + cwd: "/tmp/replacement-key-two", + sessionId: "session-independent-replacement", + }), + ]); + + expect(startCount).toBe(4); + expect(replacementOne.executeCalls).toEqual(["print('one')"]); + expect(replacementTwo.executeCalls).toEqual(["print('two')"]); + }); + it("resets the session kernel when requested", async () => { const firstKernel = new FakeKernel(okResult); const secondKernel = new FakeKernel(okResult);