From 6ed9560f00fe29fbe333928fc952109a02ccbf2d Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Sat, 22 Aug 2026 18:31:04 +0900 Subject: [PATCH] fix(runtime): close admission before transitions Close the published runtime admission synchronously when Store admits restart, account recycle, stop, or failure. Keep semantic and physical teardown asynchronous while preventing stale runtime completions from publishing into a successor generation. --- .../CodexReview/ReviewRuntimeLifecycle.swift | 2 +- .../CodexReview/Store/CodexReviewStore.swift | 16 ++++++++++++--- Sources/CodexReviewHost/CodexReviewHost.swift | 2 +- .../LiveCodexReviewStoreBackend.swift | 2 +- Sources/CodexReviewTesting/TestSupport.swift | 2 +- .../CodexReviewHostTests.swift | 4 ++++ .../CodexReviewStoreLifecycleTests.swift | 20 +++++++++++++++++++ 7 files changed, 41 insertions(+), 7 deletions(-) diff --git a/Sources/CodexReview/ReviewRuntimeLifecycle.swift b/Sources/CodexReview/ReviewRuntimeLifecycle.swift index d4de094..f8e0592 100644 --- a/Sources/CodexReview/ReviewRuntimeLifecycle.swift +++ b/Sources/CodexReview/ReviewRuntimeLifecycle.swift @@ -80,7 +80,7 @@ package func applyRuntimeAuthenticationSnapshot( @MainActor package protocol RuntimeLifecycleHandle: AnyObject, Sendable { func activate() async throws - func closeAdmission() async + func closeAdmission() func close(purpose: ReviewRuntimeTransitionPurpose) async throws func waitUntilClosed() async throws } diff --git a/Sources/CodexReview/Store/CodexReviewStore.swift b/Sources/CodexReview/Store/CodexReviewStore.swift index c4715be..fe07060 100644 --- a/Sources/CodexReview/Store/CodexReviewStore.swift +++ b/Sources/CodexReview/Store/CodexReviewStore.swift @@ -158,6 +158,7 @@ public final class CodexReviewStore { break } + closePublishedRuntimeAdmission(in: previousState) let generation = previousState.generation.successor() switch previousState { case .running(_, let runtime, let mcp): @@ -253,6 +254,7 @@ public final class CodexReviewStore { } let previousState = runtimeState + closePublishedRuntimeAdmission(in: previousState) let generation = previousState.generation.successor() if case .failed(let message) = intent.finalState { transitionToFailed(message) @@ -303,7 +305,6 @@ public final class CodexReviewStore { await stopMCPServer() case .running(_, let runtime, _): - await runtime.handle.closeAdmission() await stopPublishedRuntimeSemantics(intent: intent) await stopMCPServer() await closeRuntime( @@ -358,6 +359,7 @@ public final class CodexReviewStore { package func admitRuntimeRecycleAfterAccountChange() -> Task? { let previousState = runtimeState + closePublishedRuntimeAdmission(in: previousState) let predecessor: Task? let context: RuntimeAcquisitionContext switch previousState { @@ -710,7 +712,6 @@ public final class CodexReviewStore { private func closePublishedRuntimeForReplacement( _ runtime: PreparedRuntime ) async { - await runtime.handle.closeAdmission() await stopPublishedRuntimeSemantics(intent: .explicitStop) await closeRuntime( runtime, @@ -746,7 +747,7 @@ public final class CodexReviewStore { admissionAlreadyClosed: Bool = false ) async { if admissionAlreadyClosed == false { - await runtime.handle.closeAdmission() + runtime.handle.closeAdmission() } do { try await runtime.handle.close(purpose: purpose) @@ -760,6 +761,15 @@ public final class CodexReviewStore { } } + private func closePublishedRuntimeAdmission( + in state: ReviewStoreRuntimeState + ) { + guard case .running(_, let runtime, _) = state else { + return + } + runtime.handle.closeAdmission() + } + private func stopMCPServer() async { do { try await backend.mcpServerLifecycle.stop() diff --git a/Sources/CodexReviewHost/CodexReviewHost.swift b/Sources/CodexReviewHost/CodexReviewHost.swift index 9d3e987..55a06fd 100644 --- a/Sources/CodexReviewHost/CodexReviewHost.swift +++ b/Sources/CodexReviewHost/CodexReviewHost.swift @@ -380,7 +380,7 @@ private final class DirectRuntimeLifecycleHandle: RuntimeLifecycleHandle { onActivate() } - func closeAdmission() async { + func closeAdmission() { onCloseAdmission() } diff --git a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift index acc827d..786e86d 100644 --- a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift +++ b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift @@ -2051,7 +2051,7 @@ private final class LiveRuntimeLifecycleHandle: RuntimeLifecycleHandle { isActivated = true } - func closeAdmission() async { + func closeAdmission() { owner?.closeRuntimeAdmission(self) } diff --git a/Sources/CodexReviewTesting/TestSupport.swift b/Sources/CodexReviewTesting/TestSupport.swift index b22bbdb..2d718ca 100644 --- a/Sources/CodexReviewTesting/TestSupport.swift +++ b/Sources/CodexReviewTesting/TestSupport.swift @@ -884,7 +884,7 @@ package final class TestingRuntimeLifecycleHandle: RuntimeLifecycleHandle { onActivate() } - package func closeAdmission() async { + package func closeAdmission() { closeAdmissionCallCount += 1 } diff --git a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift index b0c840e..bd68b7f 100644 --- a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift +++ b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift @@ -1352,7 +1352,11 @@ struct CodexReviewHostTests { ) await firstTransport.waitForRequestCount(requestCount + 1) + let generationBeforeRestart = store.runtimeLifecycleAdmissionGeneration let restart = Task { @MainActor in await store.restart() } + try #require(await waitUntil(timeout: .seconds(2)) { + store.runtimeLifecycleAdmissionGeneration > generationBeforeRestart + }) await staleReadGate.open() await restart.value diff --git a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift index 1a75f9f..b84a2b9 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift @@ -142,6 +142,7 @@ struct CodexReviewStoreLifecycleTests { let retiringHandle = try #require(backend.lastPreparedRuntimeHandle) let recycle = try #require(store.admitRuntimeRecycleAfterAccountChange()) + #expect(retiringHandle.closeAdmissionCallCount == 1) store.requestRuntimeTeardown(intent: .explicitStop) await store.stop() await recycle.value @@ -164,6 +165,7 @@ struct CodexReviewStoreLifecycleTests { let retiringHandle = try #require(backend.lastPreparedRuntimeHandle) let recycle = try #require(store.admitRuntimeRecycleAfterAccountChange()) + #expect(retiringHandle.closeAdmissionCallCount == 1) let successor = try #require(store.admitRuntimeRecycleAfterAccountChange()) await successor.value await recycle.value @@ -179,6 +181,22 @@ struct CodexReviewStoreLifecycleTests { await store.stop() } + @Test func explicitStopAdmissionClosesPublishedRuntimeBeforeTeardownTaskEntry() async throws { + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend() + ) + let store = CodexReviewStore.makeTestingStore(backend: backend) + await store.start() + let handle = try #require(backend.lastPreparedRuntimeHandle) + + store.requestRuntimeTeardown(intent: .explicitStop) + + #expect(handle.closeAdmissionCallCount == 1) + await store.stop() + #expect(handle.closePurposes == [.stop]) + #expect(handle.waitUntilClosedCallCount == 1) + } + @Test func sameAccountRestartRetainsMCPListenerAndURL() async throws { let endpoint = try #require(URL(string: "http://127.0.0.1:19417/mcp")) let mcpOwner = TestingMCPServerLifecycleOwner(serverURL: endpoint) @@ -342,6 +360,7 @@ struct CodexReviewStoreLifecycleTests { let restart = Task { @MainActor in await store.restart() } try await waitForCutoverStatus(.draining, service: store.settingsService) + #expect(firstHandle.closeAdmissionCallCount == 1) await store.updateSettingsModel("deferred-edit") #expect(backend.lastPreparedRuntimeHandle === firstHandle) #expect(await reviewBackend.recordedCommands().filter { @@ -419,6 +438,7 @@ struct CodexReviewStoreLifecycleTests { handle.holdClose(with: closeGate) store.requestRuntimeFailure(handle: handle, cause: "Injected runtime failure.") + #expect(handle.closeAdmissionCallCount == 1) await handle.waitForClose() let expected = "Review runtime stopped unexpectedly: Injected runtime failure." #expect(store.serverState == .failed(expected))