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))