diff --git a/Sources/CodexReview/ReviewRuntimeLifecycle.swift b/Sources/CodexReview/ReviewRuntimeLifecycle.swift new file mode 100644 index 00000000..d4de0944 --- /dev/null +++ b/Sources/CodexReview/ReviewRuntimeLifecycle.swift @@ -0,0 +1,243 @@ +import Foundation + +package struct ReviewRuntimeGeneration: Hashable, Sendable { + package let rawValue: UInt64 + + package init(rawValue: UInt64) { + self.rawValue = rawValue + } + + package func successor() -> Self { + .init(rawValue: rawValue + 1) + } +} + +package enum ReviewRuntimeTransitionPurpose: Equatable, Sendable { + case start + case restartSameAccount + case stop + case runtimeFailure +} + +package struct RuntimePublicationSnapshot: Sendable { + package let authentication: CodexReviewBackendModel.Auth.Snapshot + package let settings: CodexReviewSettings.Snapshot + + package init( + authentication: CodexReviewBackendModel.Auth.Snapshot, + settings: CodexReviewSettings.Snapshot + ) { + self.authentication = authentication + self.settings = settings + } +} + +@MainActor +package func applyRuntimeAuthenticationSnapshot( + _ snapshot: CodexReviewBackendModel.Auth.Snapshot, + to auth: CodexReviewAuthModel +) { + let observedAccounts = snapshot.accounts.compactMap { account -> CodexAccount? in + let label = account.label.trimmingCharacters(in: .whitespacesAndNewlines) + let accountKey = CodexAccount.normalizedEmail(account.id.rawValue) + guard label.isEmpty == false, accountKey.isEmpty == false else { + return nil + } + return CodexAccount( + accountKey: accountKey, + email: label, + planType: account.planType, + kind: account.kind, + capabilities: account.capabilities + ) + } + let activeAccountKey = snapshot.activeAccountID.map { + CodexAccount.normalizedEmail($0.rawValue) + } + var accounts = auth.persistedAccounts + for observedAccount in observedAccounts { + if let index = accounts.firstIndex(where: { + $0.accountKey == observedAccount.accountKey + }) { + accounts[index].updateEmail(observedAccount.email) + accounts[index].updateKind( + observedAccount.kind, + capabilities: observedAccount.capabilities + ) + accounts[index].updatePlanType(observedAccount.planType) + } else { + accounts.insert(observedAccount, at: 0) + } + } + auth.applyPersistedAccountStates( + accounts.map(savedAccountPayload(from:)), + activeAccountKey: activeAccountKey + ) + auth.selectPersistedAccount(activeAccountKey) + auth.updatePhase(.signedOut) +} + +@MainActor +package protocol RuntimeLifecycleHandle: AnyObject, Sendable { + func activate() async throws + func closeAdmission() async + func close(purpose: ReviewRuntimeTransitionPurpose) async throws + func waitUntilClosed() async throws +} + +package struct PreparedRuntime: Sendable { + package let snapshot: RuntimePublicationSnapshot + package let handle: any RuntimeLifecycleHandle + + package init( + snapshot: RuntimePublicationSnapshot, + handle: any RuntimeLifecycleHandle + ) { + self.snapshot = snapshot + self.handle = handle + } +} + +package final class PreparedMCPServer: @unchecked Sendable { + package init() {} +} + +package struct MCPServerPublicationSnapshot: Sendable { + package let serverURL: URL? + + package init(serverURL: URL?) { + self.serverURL = serverURL + } +} + +package struct RetainedMCPServer: Sendable { + package let serverURL: URL? + + package init(serverURL: URL?) { + self.serverURL = serverURL + } +} + +@MainActor +package final class RuntimeReplacementContext { + private var retiringRuntime: PreparedRuntime? + + package init(retiringRuntime: PreparedRuntime?) { + self.retiringRuntime = retiringRuntime + } + + package func takeRetiringRuntime() -> PreparedRuntime? { + defer { retiringRuntime = nil } + return retiringRuntime + } +} + +@MainActor +package final class RuntimeAcquisitionContext { + private var recyclingState: ReviewStoreRuntimeState? + + package init(recycling state: ReviewStoreRuntimeState? = nil) { + recyclingState = state + } + + package func takeRecyclingState() -> ReviewStoreRuntimeState? { + defer { recyclingState = nil } + return recyclingState + } +} + +package enum ReviewStoreRuntimeState { + case stopped(ReviewRuntimeGeneration) + case acquiring( + generation: ReviewRuntimeGeneration, + context: RuntimeAcquisitionContext, + task: Task + ) + case running( + generation: ReviewRuntimeGeneration, + runtime: PreparedRuntime, + mcp: RetainedMCPServer + ) + case replacing( + generation: ReviewRuntimeGeneration, + context: RuntimeReplacementContext, + retainedMCP: RetainedMCPServer, + task: Task + ) + case tearingDown( + generation: ReviewRuntimeGeneration, + cleanupIntent: ReviewRuntimeTeardownIntent, + finalIntent: ReviewRuntimeTeardownIntent, + task: Task + ) + case failed( + generation: ReviewRuntimeGeneration, + retainedMCP: RetainedMCPServer? + ) + + package var generation: ReviewRuntimeGeneration { + switch self { + case .stopped(let generation), + .acquiring(let generation, _, _), + .running(let generation, _, _), + .replacing(let generation, _, _, _), + .tearingDown(let generation, _, _, _), + .failed(let generation, _): + generation + } + } +} + +@MainActor +package protocol MCPServerLifecycleOwner: Sendable { + func prepare() async throws -> PreparedMCPServer + func activate( + _ preparation: PreparedMCPServer + ) async throws -> MCPServerPublicationSnapshot + func stop() async throws +} + +@MainActor +package final class NoMCPServerLifecycleOwner: MCPServerLifecycleOwner { + private enum State { + case stopped + case prepared(PreparedMCPServer) + case running(PreparedMCPServer) + } + + private var serverURL: URL? + private var state: State = .stopped + + package init(serverURL: URL? = nil) { + self.serverURL = serverURL + } + + package func updateServerURL(_ serverURL: URL?) { + self.serverURL = serverURL + } + + package func prepare() async throws -> PreparedMCPServer { + guard case .stopped = state else { + throw CancellationError() + } + let preparation = PreparedMCPServer() + state = .prepared(preparation) + return preparation + } + + package func activate( + _ preparation: PreparedMCPServer + ) async throws -> MCPServerPublicationSnapshot { + guard case .prepared(let current) = state, + current === preparation + else { + throw CancellationError() + } + state = .running(preparation) + return .init(serverURL: serverURL) + } + + package func stop() async throws { + state = .stopped + } +} diff --git a/Sources/CodexReview/Store/CodexReviewStore.swift b/Sources/CodexReview/Store/CodexReviewStore.swift index 271de613..c4715be3 100644 --- a/Sources/CodexReview/Store/CodexReviewStore.swift +++ b/Sources/CodexReview/Store/CodexReviewStore.swift @@ -4,13 +4,6 @@ import Observation @MainActor @Observable public final class CodexReviewStore { - private struct RuntimeTeardownOperation { - let id: UUID - let cleanupIntent: ReviewRuntimeTeardownIntent - var finalIntent: ReviewRuntimeTeardownIntent - let task: Task - } - package struct ReviewTerminalWaiter { package var id: UUID package var continuation: CheckedContinuation @@ -27,10 +20,13 @@ public final class CodexReviewStore { backend.seed.shouldAutoStartEmbeddedServer } package var runtimeTeardownFinalState: ReviewRuntimeTeardownIntent.FinalState? { - runtimeTeardownOperation?.finalIntent.finalState + guard case .tearingDown(_, _, let finalIntent, _) = runtimeState else { + return nil + } + return finalIntent.finalState } package var runtimeLifecycleAdmissionGeneration: UInt64 { - runtimeLifecycleGeneration + runtimeState.generation.rawValue } @ObservationIgnored package let diagnosticsURL: URL? @@ -50,8 +46,9 @@ public final class CodexReviewStore { @ObservationIgnored package var reviewTerminalWaiters: [String: [ReviewTerminalWaiter]] = [:] @ObservationIgnored package var closedSessions: Set = [] @ObservationIgnored package var accountRateLimitAutoRefreshDriver: CodexReviewStoreRateLimitAutoRefreshDriver? - @ObservationIgnored private var runtimeTeardownOperation: RuntimeTeardownOperation? - @ObservationIgnored private var runtimeLifecycleGeneration: UInt64 = 0 + @ObservationIgnored package var runtimeState: ReviewStoreRuntimeState = .stopped( + .init(rawValue: 0) + ) package init( backend: any CodexReviewStoreBackend = PreviewCodexReviewStoreBackend(), @@ -93,7 +90,14 @@ public final class CodexReviewStore { isolated deinit { accountRateLimitAutoRefreshDriver?.cancel() - runtimeTeardownOperation?.task.cancel() + switch runtimeState { + case .acquiring(_, _, let task), + .replacing(_, _, _, let task), + .tearingDown(_, _, _, let task): + task.cancel() + case .stopped, .running, .failed: + break + } for task in reviewWorkerTasks.values { task.cancel() } @@ -142,29 +146,55 @@ public final class CodexReviewStore { } public func start(forceRestartIfNeeded: Bool = false) async { - let admissionGeneration = admitRuntimeLifecycleRequest() - if let teardownTask = runtimeTeardownOperation?.task { - await teardownTask.value - guard admissionGeneration == runtimeLifecycleGeneration else { - return - } - } - switch serverState { - case .stopped, .failed: - break - case .starting: - return + let previousState = runtimeState + switch previousState { case .running where forceRestartIfNeeded == false: return - case .running: + case .acquiring(_, _, let task) where forceRestartIfNeeded == false, + .replacing(_, _, _, let task) where forceRestartIfNeeded == false: + await task.value + return + case .stopped, .acquiring, .running, .replacing, .tearingDown, .failed: break } - serverState = .starting - serverURL = nil - writeDiagnosticsIfNeeded() - await backend.start(store: self, forceRestartIfNeeded: forceRestartIfNeeded) - await settingsService.refreshIfRunning(serverState: serverState) - startAccountRateLimitAutoRefresh() + + let generation = previousState.generation.successor() + switch previousState { + case .running(_, let runtime, let mcp): + await beginRuntimeReplacement( + generation: generation, + context: .init(retiringRuntime: runtime), + retainedMCP: mcp + ) + case .replacing(_, let context, let retainedMCP, let task): + task.cancel() + await beginRuntimeReplacement( + generation: generation, + context: context, + retainedMCP: retainedMCP, + predecessor: task + ) + case .failed(_, let retainedMCP?): + await beginRuntimeReplacement( + generation: generation, + context: .init(retiringRuntime: nil), + retainedMCP: retainedMCP + ) + case .acquiring(_, let context, let task): + task.cancel() + await beginRuntimeAcquisition( + generation: generation, + context: context, + predecessor: task + ) + case .tearingDown(_, _, _, let task): + await beginRuntimeAcquisition( + generation: generation, + predecessor: task + ) + case .stopped, .failed: + await beginRuntimeAcquisition(generation: generation) + } } public func stop() async { @@ -172,7 +202,6 @@ public final class CodexReviewStore { } package func stop(intent: ReviewRuntimeTeardownIntent) async { - _ = admitRuntimeLifecycleRequest() let task = admitRuntimeTeardown(intent: intent) await task.value } @@ -180,27 +209,51 @@ public final class CodexReviewStore { package func requestRuntimeTeardown( intent: ReviewRuntimeTeardownIntent ) { - _ = admitRuntimeLifecycleRequest() _ = admitRuntimeTeardown(intent: intent) } - private func admitRuntimeLifecycleRequest() -> UInt64 { - runtimeLifecycleGeneration += 1 - return runtimeLifecycleGeneration + package func requestRuntimeFailure( + handle: any RuntimeLifecycleHandle, + cause: String + ) { + guard case .running(_, let runtime, _) = runtimeState, + runtime.handle === handle + else { + return + } + _ = admitRuntimeTeardown(intent: .unexpectedFailure(cause)) } private func admitRuntimeTeardown( intent: ReviewRuntimeTeardownIntent ) -> Task { - if var operation = runtimeTeardownOperation { - if intent.supersedesConcurrentFinalState { - operation.finalIntent = intent - runtimeTeardownOperation = operation + if case .tearingDown( + let currentGeneration, + let cleanupIntent, + let finalIntent, + let currentTask + ) = runtimeState { + guard intent.supersedesConcurrentFinalState, + finalIntent != intent + else { + return currentTask } - return operation.task + let generation = currentGeneration.successor() + let task = Task { @MainActor [weak self] in + await currentTask.value + self?.finishRuntimeTeardown(generation: generation) + } + runtimeState = .tearingDown( + generation: generation, + cleanupIntent: cleanupIntent, + finalIntent: intent, + task: task + ) + return task } - let operationID = UUID() + let previousState = runtimeState + let generation = previousState.generation.successor() if case .failed(let message) = intent.finalState { transitionToFailed(message) } @@ -208,11 +261,15 @@ public final class CodexReviewStore { guard let self else { return } - await self.performRuntimeTeardown(intent: intent) - self.finishRuntimeTeardown(operationID: operationID) + await self.performRuntimeTeardown( + previousState: previousState, + generation: generation, + intent: intent + ) + self.finishRuntimeTeardown(generation: generation) } - runtimeTeardownOperation = .init( - id: operationID, + runtimeState = .tearingDown( + generation: generation, cleanupIntent: intent, finalIntent: intent, task: task @@ -221,6 +278,448 @@ public final class CodexReviewStore { } private func performRuntimeTeardown( + previousState: ReviewStoreRuntimeState, + generation: ReviewRuntimeGeneration, + intent: ReviewRuntimeTeardownIntent + ) async { + switch previousState { + case .acquiring(_, let context, let task): + task.cancel() + await task.value + if let recyclingState = context.takeRecyclingState() { + await performRuntimeTeardown( + previousState: recyclingState, + generation: generation, + intent: intent + ) + } + + case .replacing(_, let context, _, let task): + task.cancel() + await task.value + if let retiringRuntime = context.takeRetiringRuntime() { + await closePublishedRuntimeForReplacement(retiringRuntime) + } + await stopMCPServer() + + case .running(_, let runtime, _): + await runtime.handle.closeAdmission() + await stopPublishedRuntimeSemantics(intent: intent) + await stopMCPServer() + await closeRuntime( + runtime, + purpose: intent == .explicitStop ? .stop : .runtimeFailure, + admissionAlreadyClosed: true + ) + + case .tearingDown(_, _, _, let task): + await task.value + + case .failed(_, let retainedMCP): + if retainedMCP != nil { + await stopMCPServer() + } + + case .stopped: + break + } + } + + private func finishRuntimeTeardown( + generation: ReviewRuntimeGeneration + ) { + guard case .tearingDown( + let currentGeneration, + _, + let finalIntent, + _ + ) = runtimeState, + currentGeneration == generation + else { + return + } + switch finalIntent.finalState { + case .stopped: + runtimeState = .stopped(generation) + transitionToStopped() + case .failed(let message): + runtimeState = .failed(generation: generation, retainedMCP: nil) + transitionToFailed(message) + } + } + + public func restart() async { + await start(forceRestartIfNeeded: true) + } + + package func recycleRuntimeAfterAccountChange() async { + await admitRuntimeRecycleAfterAccountChange()?.value + } + + package func admitRuntimeRecycleAfterAccountChange() -> Task? { + let previousState = runtimeState + let predecessor: Task? + let context: RuntimeAcquisitionContext + switch previousState { + case .acquiring(_, let currentContext, let task): + task.cancel() + predecessor = task + context = currentContext + case .replacing(_, _, _, let task): + task.cancel() + predecessor = task + context = .init(recycling: previousState) + case .running: + predecessor = nil + context = .init(recycling: previousState) + case .stopped, .tearingDown, .failed: + return nil + } + + let generation = previousState.generation.successor() + serverState = .starting + serverURL = nil + writeDiagnosticsIfNeeded() + let task = Task { @MainActor [weak self] in + if let predecessor { + await predecessor.value + } + guard let self, self.isCurrentAcquisition(generation) else { + return + } + await self.performRuntimeAcquisition( + generation: generation, + context: context + ) + } + runtimeState = .acquiring( + generation: generation, + context: context, + task: task + ) + return task + } + + public func waitUntilStopped() async { + if case .tearingDown(_, _, _, let task) = runtimeState { + await task.value + } + await backend.waitUntilStopped() + } + + private func beginRuntimeAcquisition( + generation: ReviewRuntimeGeneration, + context: RuntimeAcquisitionContext = .init(), + predecessor: Task? = nil + ) async { + if predecessor == nil { + serverState = .starting + serverURL = nil + writeDiagnosticsIfNeeded() + } + let task = Task { @MainActor [weak self] in + if let predecessor { + await predecessor.value + } + guard let self, self.isCurrentAcquisition(generation) else { + return + } + if predecessor != nil { + self.serverState = .starting + self.serverURL = nil + self.writeDiagnosticsIfNeeded() + } + await self.performRuntimeAcquisition( + generation: generation, + context: context + ) + } + runtimeState = .acquiring( + generation: generation, + context: context, + task: task + ) + await task.value + } + + private func performRuntimeAcquisition( + generation: ReviewRuntimeGeneration, + context: RuntimeAcquisitionContext + ) async { + guard isCurrentAcquisition(generation) else { + if let recyclingState = context.takeRecyclingState() { + await performRuntimeTeardown( + previousState: recyclingState, + generation: generation, + intent: .explicitStop + ) + } + return + } + var cutoverToken: CodexReviewSettingsService.RuntimeCutoverToken? + var settingsServiceOwnsCutover = false + var preparedRuntime: PreparedRuntime? + var preparedMCPServer: PreparedMCPServer? + + do { + let token = try await settingsService.beginRuntimeCutover() + cutoverToken = token + + if let previousState = context.takeRecyclingState() { + await performRuntimeTeardown( + previousState: previousState, + generation: generation, + intent: .explicitStop + ) + } + guard isCurrentAcquisition(generation) else { + cancelRuntimeCutover(token) + return + } + + let mcpPreparation = try await backend.mcpServerLifecycle.prepare() + preparedMCPServer = mcpPreparation + guard isCurrentAcquisition(generation) else { + await stopMCPServer() + cancelRuntimeCutover(token) + return + } + + let runtime = try await backend.prepareRuntime( + generation: generation, + purpose: .start + ) + preparedRuntime = runtime + guard isCurrentAcquisition(generation) else { + await closeRuntime(runtime, purpose: .start) + await stopMCPServer() + cancelRuntimeCutover(token) + return + } + + try await runtime.handle.activate() + guard isCurrentAcquisition(generation) else { + await closeRuntime(runtime, purpose: .start) + await stopMCPServer() + cancelRuntimeCutover(token) + return + } + + let mcpSnapshot = try await backend.mcpServerLifecycle.activate(mcpPreparation) + guard isCurrentAcquisition(generation) else { + await closeRuntime(runtime, purpose: .start) + await stopMCPServer() + cancelRuntimeCutover(token) + return + } + + settingsServiceOwnsCutover = true + try await settingsService.commitRuntimeSnapshot( + token: token, + snapshot: runtime.snapshot.settings + ) + guard isCurrentAcquisition(generation) else { + await closeRuntime(runtime, purpose: .start) + await stopMCPServer() + return + } + + try backend.commitRuntimePublication( + runtime.snapshot, + handle: runtime.handle, + auth: auth + ) + guard isCurrentAcquisition(generation) else { + await closeRuntime(runtime, purpose: .start) + await stopMCPServer() + return + } + + let retainedMCP = RetainedMCPServer(serverURL: mcpSnapshot.serverURL) + runtimeState = .running( + generation: generation, + runtime: runtime, + mcp: retainedMCP + ) + publishRuntime(serverURL: retainedMCP.serverURL) + await backend.waitForRuntimePublication(handle: runtime.handle) + } catch { + if let recyclingState = context.takeRecyclingState() { + await performRuntimeTeardown( + previousState: recyclingState, + generation: generation, + intent: .explicitStop + ) + } + if let preparedRuntime { + await closeRuntime(preparedRuntime, purpose: .start) + } + if preparedMCPServer != nil { + await stopMCPServer() + } + let isCurrentGeneration = isCurrentAcquisition(generation) + let wasIntentionallyCancelled = Task.isCancelled || isCurrentGeneration == false + if let cutoverToken, settingsServiceOwnsCutover == false { + if wasIntentionallyCancelled { + cancelRuntimeCutover(cutoverToken) + } else { + abortRuntimeCutover(cutoverToken, message: error.localizedDescription) + } + } + guard wasIntentionallyCancelled == false else { + return + } + runtimeState = .failed(generation: generation, retainedMCP: nil) + transitionToFailed(error.localizedDescription) + } + } + + private func beginRuntimeReplacement( + generation: ReviewRuntimeGeneration, + context: RuntimeReplacementContext, + retainedMCP: RetainedMCPServer, + predecessor: Task? = nil + ) async { + serverState = .starting + serverURL = retainedMCP.serverURL + writeDiagnosticsIfNeeded() + let task = Task { @MainActor [weak self] in + if let predecessor { + await predecessor.value + } + guard let self else { + return + } + await self.performRuntimeReplacement( + generation: generation, + context: context, + retainedMCP: retainedMCP + ) + } + runtimeState = .replacing( + generation: generation, + context: context, + retainedMCP: retainedMCP, + task: task + ) + await task.value + } + + private func performRuntimeReplacement( + generation: ReviewRuntimeGeneration, + context: RuntimeReplacementContext, + retainedMCP: RetainedMCPServer + ) async { + guard isCurrentReplacement(generation) else { + return + } + var cutoverToken: CodexReviewSettingsService.RuntimeCutoverToken? + var settingsServiceOwnsCutover = false + var closedPreviousRuntime = false + var preparedRuntime: PreparedRuntime? + + do { + let token = try await settingsService.beginRuntimeCutover() + cutoverToken = token + + if let retiringRuntime = context.takeRetiringRuntime() { + await closePublishedRuntimeForReplacement(retiringRuntime) + closedPreviousRuntime = true + } + guard isCurrentReplacement(generation) else { + cancelRuntimeCutover(token) + return + } + + let runtime = try await backend.prepareRuntime( + generation: generation, + purpose: .restartSameAccount + ) + preparedRuntime = runtime + guard isCurrentReplacement(generation) else { + await closeRuntime(runtime, purpose: .restartSameAccount) + cancelRuntimeCutover(token) + return + } + + try await runtime.handle.activate() + guard isCurrentReplacement(generation) else { + await closeRuntime(runtime, purpose: .restartSameAccount) + cancelRuntimeCutover(token) + return + } + + settingsServiceOwnsCutover = true + try await settingsService.commitRuntimeSnapshot( + token: token, + snapshot: runtime.snapshot.settings + ) + guard isCurrentReplacement(generation) else { + await closeRuntime(runtime, purpose: .restartSameAccount) + return + } + + try backend.commitRuntimePublication( + runtime.snapshot, + handle: runtime.handle, + auth: auth + ) + guard isCurrentReplacement(generation) else { + await closeRuntime(runtime, purpose: .restartSameAccount) + return + } + + runtimeState = .running( + generation: generation, + runtime: runtime, + mcp: retainedMCP + ) + publishRuntime(serverURL: retainedMCP.serverURL) + await backend.waitForRuntimePublication(handle: runtime.handle) + } catch { + if closedPreviousRuntime == false, + let retiringRuntime = context.takeRetiringRuntime() + { + await closePublishedRuntimeForReplacement(retiringRuntime) + } + if let preparedRuntime { + await closeRuntime(preparedRuntime, purpose: .restartSameAccount) + } + let isCurrentGeneration = isCurrentReplacement(generation) + let wasIntentionallyCancelled = Task.isCancelled || isCurrentGeneration == false + if let cutoverToken, settingsServiceOwnsCutover == false { + if wasIntentionallyCancelled { + cancelRuntimeCutover(cutoverToken) + } else { + abortRuntimeCutover(cutoverToken, message: error.localizedDescription) + } + } + guard wasIntentionallyCancelled == false else { + return + } + runtimeState = .failed( + generation: generation, + retainedMCP: retainedMCP + ) + serverURL = retainedMCP.serverURL + serverState = .failed(error.localizedDescription) + writeDiagnosticsIfNeeded() + } + } + + private func closePublishedRuntimeForReplacement( + _ runtime: PreparedRuntime + ) async { + await runtime.handle.closeAdmission() + await stopPublishedRuntimeSemantics(intent: .explicitStop) + await closeRuntime( + runtime, + purpose: .restartSameAccount, + admissionAlreadyClosed: true + ) + } + + private func stopPublishedRuntimeSemantics( intent: ReviewRuntimeTeardownIntent ) async { let locallyCancelledJobIDs: [String] @@ -241,28 +740,80 @@ public final class CodexReviewStore { ) } - private func finishRuntimeTeardown(operationID: UUID) { - guard let operation = runtimeTeardownOperation, - operation.id == operationID - else { - return + private func closeRuntime( + _ runtime: PreparedRuntime, + purpose: ReviewRuntimeTransitionPurpose, + admissionAlreadyClosed: Bool = false + ) async { + if admissionAlreadyClosed == false { + await runtime.handle.closeAdmission() } - runtimeTeardownOperation = nil - switch operation.finalIntent.finalState { - case .stopped: - transitionToStopped() - case .failed(let message): - transitionToFailed(message) + do { + try await runtime.handle.close(purpose: purpose) + } catch { + writeDiagnosticsIfNeeded() + } + do { + try await runtime.handle.waitUntilClosed() + } catch { + writeDiagnosticsIfNeeded() } } - public func restart() async { - await stop() - await start(forceRestartIfNeeded: true) + private func stopMCPServer() async { + do { + try await backend.mcpServerLifecycle.stop() + } catch { + writeDiagnosticsIfNeeded() + } } - public func waitUntilStopped() async { - await backend.waitUntilStopped() + private func abortRuntimeCutover( + _ token: CodexReviewSettingsService.RuntimeCutoverToken, + message: String + ) { + do { + try settingsService.abortRuntimeCutover(token: token, message: message) + } catch { + preconditionFailure( + "CodexReviewStore must consume its current runtime cutover token exactly once: \(error)" + ) + } + } + + private func cancelRuntimeCutover( + _ token: CodexReviewSettingsService.RuntimeCutoverToken + ) { + do { + try settingsService.cancelRuntimeCutover(token: token) + } catch { + preconditionFailure( + "CodexReviewStore must consume its current runtime cutover token exactly once: \(error)" + ) + } + } + + private func isCurrentAcquisition( + _ generation: ReviewRuntimeGeneration + ) -> Bool { + guard case .acquiring(let currentGeneration, _, _) = runtimeState else { + return false + } + return currentGeneration == generation + } + + private func isCurrentReplacement( + _ generation: ReviewRuntimeGeneration + ) -> Bool { + guard case .replacing(let currentGeneration, _, _, _) = runtimeState else { + return false + } + return currentGeneration == generation + } + + private func publishRuntime(serverURL: URL?) { + transitionToRunning(serverURL: serverURL) + startAccountRateLimitAutoRefresh() } public func refreshAuthentication() async { diff --git a/Sources/CodexReview/Store/CodexReviewStoreBackend.swift b/Sources/CodexReview/Store/CodexReviewStoreBackend.swift index 00f3bd9b..ec8ae857 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreBackend.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreBackend.swift @@ -24,13 +24,25 @@ package struct CodexReviewStoreSeed { } @MainActor -package protocol CodexReviewStoreBackend: CodexReviewSettingsBackend { +package protocol CodexReviewStoreBackend: CodexReviewSettingsBackend, Sendable { var seed: CodexReviewStoreSeed { get } var isActive: Bool { get } var handlesActiveReviewStopCleanup: Bool { get } + var mcpServerLifecycle: any MCPServerLifecycleOwner { get } func attachStore(_ store: CodexReviewStore) - func start(store: CodexReviewStore, forceRestartIfNeeded: Bool) async + func prepareRuntime( + generation: ReviewRuntimeGeneration, + purpose: ReviewRuntimeTransitionPurpose + ) async throws -> PreparedRuntime + func commitRuntimePublication( + _ snapshot: RuntimePublicationSnapshot, + handle: any RuntimeLifecycleHandle, + auth: CodexReviewAuthModel + ) throws + func waitForRuntimePublication( + handle: any RuntimeLifecycleHandle + ) async func stop(store: CodexReviewStore) async func stop( store: CodexReviewStore, @@ -76,6 +88,18 @@ extension CodexReviewStoreBackend { await stop(store: store) } + package func commitRuntimePublication( + _ snapshot: RuntimePublicationSnapshot, + handle _: any RuntimeLifecycleHandle, + auth: CodexReviewAuthModel + ) throws { + applyRuntimeAuthenticationSnapshot(snapshot.authentication, to: auth) + } + + package func waitForRuntimePublication( + handle _: any RuntimeLifecycleHandle + ) async {} + package func startReview( _ request: CodexReviewBackendModel.Review.Start ) async throws -> BackendReviewAttempt { diff --git a/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift b/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift index 3f059a4b..fd60169f 100644 --- a/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift +++ b/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift @@ -5,6 +5,7 @@ package class PreviewCodexReviewStoreBackend: CodexReviewStoreBackend { package let seed: CodexReviewStoreSeed package var isActive = false package var currentSettingsSnapshot: CodexReviewSettings.Snapshot + package let mcpServerLifecycle: any MCPServerLifecycleOwner = NoMCPServerLifecycleOwner() package init(seed: CodexReviewStoreSeed = .init()) { self.seed = seed @@ -17,9 +18,11 @@ package class PreviewCodexReviewStoreBackend: CodexReviewStoreBackend { package func attachStore(_: CodexReviewStore) {} - package func start(store: CodexReviewStore, forceRestartIfNeeded _: Bool) async { - isActive = true - store.transitionToFailed(Self.previewUnavailableMessage) + package func prepareRuntime( + generation _: ReviewRuntimeGeneration, + purpose _: ReviewRuntimeTransitionPurpose + ) async throws -> PreparedRuntime { + throw CodexReviewAPI.Error.io(Self.previewUnavailableMessage) } package func stop(store _: CodexReviewStore) async { diff --git a/Sources/CodexReviewHost/CodexReviewHost.swift b/Sources/CodexReviewHost/CodexReviewHost.swift index 1af2c8a8..9d3e987c 100644 --- a/Sources/CodexReviewHost/CodexReviewHost.swift +++ b/Sources/CodexReviewHost/CodexReviewHost.swift @@ -7,6 +7,7 @@ import CodexReviewMCPServer package final class CodexReviewHost { package let store: CodexReviewStore package let mcpServer: CodexReviewMCPServer + private let directBackend: DirectCodexReviewStoreBackend private let shutdown: @Sendable () async throws -> Void private var endpoint: URL? @@ -17,10 +18,15 @@ package final class CodexReviewHost { endpoint: URL? = nil, shutdown: @escaping @Sendable () async throws -> Void = {} ) { - self.shutdown = shutdown self.endpoint = endpoint + let directBackend = DirectCodexReviewStoreBackend( + backend: backend, + endpoint: endpoint + ) + self.directBackend = directBackend + self.shutdown = shutdown let store = CodexReviewStore( - backend: DirectCodexReviewStoreBackend(backend: backend), + backend: directBackend, clock: clock, idGenerator: idGenerator ) @@ -44,11 +50,15 @@ package final class CodexReviewHost { } package func start(endpoint: URL? = nil) async { + let endpointChanged = endpoint != nil && endpoint != self.endpoint if let endpoint { self.endpoint = endpoint } - store.transitionToRunning(serverURL: self.endpoint) - await store.refreshSettings() + directBackend.updateEndpoint(self.endpoint) + if endpointChanged { + await store.stop() + } + await store.start() } package func stop() async throws { @@ -60,7 +70,9 @@ package final class CodexReviewHost { @MainActor private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { let seed = CodexReviewStoreSeed() + let mcpServerLifecycle: any MCPServerLifecycleOwner private let backend: any CodexReviewBackend + private let mcpLifecycleOwner: NoMCPServerLifecycleOwner private var currentSettingsSnapshot = CodexReviewSettings.Snapshot() private var loginChallenge: CodexReviewBackendModel.Login.Challenge? private var active = false @@ -73,14 +85,33 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { currentSettingsSnapshot } - init(backend: any CodexReviewBackend) { + init(backend: any CodexReviewBackend, endpoint: URL?) { self.backend = backend + let mcpLifecycleOwner = NoMCPServerLifecycleOwner(serverURL: endpoint) + self.mcpLifecycleOwner = mcpLifecycleOwner + self.mcpServerLifecycle = mcpLifecycleOwner } func attachStore(_: CodexReviewStore) {} - func start(store _: CodexReviewStore, forceRestartIfNeeded _: Bool) async { - active = true + func updateEndpoint(_ endpoint: URL?) { + mcpLifecycleOwner.updateServerURL(endpoint) + } + + func prepareRuntime( + generation _: ReviewRuntimeGeneration, + purpose _: ReviewRuntimeTransitionPurpose + ) async throws -> PreparedRuntime { + let settings = try await Self.monitorSettings(from: backend.readSettings()) + let authentication = try await backend.readAuth() + return PreparedRuntime( + snapshot: .init(authentication: authentication, settings: settings), + handle: DirectRuntimeLifecycleHandle( + onActivate: { [weak self] in self?.active = true }, + onCloseAdmission: { [weak self] in self?.active = false }, + onClose: { [weak self] in self?.active = false } + ) + ) } func stop(store _: CodexReviewStore) async { @@ -90,6 +121,7 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { func waitUntilStopped() async {} func refreshSettings() async throws -> CodexReviewSettings.Snapshot { + guard active else { return currentSettingsSnapshot } currentSettingsSnapshot = try await Self.monitorSettings(from: backend.readSettings()) return currentSettingsSnapshot } @@ -101,6 +133,7 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { serviceTier: CodexReviewSettings.ServiceTier?, persistServiceTier: Bool ) async throws { + guard active else { return } var change = CodexReviewBackendModel.Settings.Change( model: model, updatesModel: true @@ -119,6 +152,7 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { func updateSettingsReasoningEffort( _ reasoningEffort: CodexReviewSettings.ReasoningEffort? ) async throws { + guard active else { return } currentSettingsSnapshot = try await Self.monitorSettings( from: backend.applySettings(.init( reasoningEffort: reasoningEffort?.rawValue, @@ -130,6 +164,7 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { func updateSettingsServiceTier( _ serviceTier: CodexReviewSettings.ServiceTier? ) async throws { + guard active else { return } currentSettingsSnapshot = try await Self.monitorSettings( from: backend.applySettings(.init( serviceTier: serviceTier?.rawValue, @@ -139,6 +174,7 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { } func refreshAuth(auth: CodexReviewAuthModel) async { + guard active else { return } do { Self.applyAuthSnapshot(try await backend.readAuth(), to: auth) } catch { @@ -147,6 +183,7 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { } func signIn(auth: CodexReviewAuthModel) async { + guard active else { return } do { let challenge = try await backend.startLogin(.init()) loginChallenge = challenge @@ -162,6 +199,7 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { } func addAccount(auth: CodexReviewAuthModel) async { + guard active else { return } await signIn(auth: auth) } @@ -237,7 +275,10 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { _ request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { - try await backend.startReview(request, admission: admission) + guard active else { + throw CodexReviewAPI.Error.io("Review runtime is not running.") + } + return try await backend.startReview(request, admission: admission) } func interruptReview( @@ -313,6 +354,51 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { } } +@MainActor +private final class DirectRuntimeLifecycleHandle: RuntimeLifecycleHandle { + private let onActivate: @MainActor @Sendable () -> Void + private let onCloseAdmission: @MainActor @Sendable () -> Void + private let onClose: @MainActor @Sendable () -> Void + private var isActivated = false + private var didClose = false + + init( + onActivate: @escaping @MainActor @Sendable () -> Void, + onCloseAdmission: @escaping @MainActor @Sendable () -> Void, + onClose: @escaping @MainActor @Sendable () -> Void + ) { + self.onActivate = onActivate + self.onCloseAdmission = onCloseAdmission + self.onClose = onClose + } + + func activate() async throws { + guard isActivated == false, didClose == false else { + throw CancellationError() + } + isActivated = true + onActivate() + } + + func closeAdmission() async { + onCloseAdmission() + } + + func close(purpose _: ReviewRuntimeTransitionPurpose) async throws { + guard didClose == false else { + return + } + didClose = true + onClose() + } + + func waitUntilClosed() async throws { + guard didClose else { + throw CancellationError() + } + } +} + extension CodexReviewBackendModel.Login.Challenge { func signInDetail(nativeAuthentication: Bool) -> String { if let userCode = userCode?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty { diff --git a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift index 39e36d1c..acc827d9 100644 --- a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift +++ b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift @@ -187,7 +187,7 @@ public extension CodexReviewStore { } @MainActor -private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { +private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend, MCPServerLifecycleOwner { typealias MCPHTTPServerFactory = @MainActor @Sendable ( CodexReviewStore, CodexReviewMCPHTTPServer.Configuration @@ -197,6 +197,8 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private var client: AppServerClient? private var appServerBackend: AppServerCodexReviewBackend? + private var activeRuntimeHandle: LiveRuntimeLifecycleHandle? + private var acceptsRuntimeRequests = false private var mcpHTTPServer: (any CodexReviewMCPHTTPServing)? private var loginChallenge: CodexReviewBackendModel.Login.Challenge? private var loginBackend: AppServerCodexReviewBackend? @@ -220,6 +222,9 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private let appServerRuntimeFactory: AppServerRuntimeFactory private let shutdownCleanupTimeout: Duration private weak var attachedStore: CodexReviewStore? + private var preparingMCPServer: PreparedMCPServer? + private var preparedMCPServer: PreparedMCPServer? + private var runningMCPServer: PreparedMCPServer? init( environment: [String: String] = ProcessInfo.processInfo.environment, @@ -269,7 +274,11 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } var isActive: Bool { - client != nil + acceptsRuntimeRequests + } + + var mcpServerLifecycle: any MCPServerLifecycleOwner { + self } var handlesActiveReviewStopCleanup: Bool { @@ -395,62 +404,209 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { attachedStore = store } - func start(store: CodexReviewStore, forceRestartIfNeeded: Bool) async { - logger.info("Starting review runtime; forceRestartIfNeeded=\(forceRestartIfNeeded, privacy: .public)") - if appServerBackend != nil, forceRestartIfNeeded == false { - logger.info("Review runtime already has an app-server backend") - store.transitionToRunning(serverURL: await mcpHTTPServer?.url) - return - } - if forceRestartIfNeeded { - await stop(store: store) + func prepare() async throws -> PreparedMCPServer { + guard mcpHTTPServer == nil, + preparingMCPServer == nil, + preparedMCPServer == nil, + runningMCPServer == nil + else { + throw CancellationError() } - - var startedClient: AppServerClient? - var startedBackend: AppServerCodexReviewBackend? - var startedHTTPServer: (any CodexReviewMCPHTTPServing)? + let preparation = PreparedMCPServer() + preparingMCPServer = preparation do { if mcpHTTPServerFactory != nil { try await mcpHTTPServerBindChecker(mcpHTTPServerConfiguration) } - let runtime = try await appServerRuntimeFactory(codexHomeURL) - let client = runtime.client - let backend = runtime.backend - startedClient = client - startedBackend = backend - self.client = client - self.appServerBackend = backend - observeAuthNotifications(client: client, backend: backend, store: store) + guard preparingMCPServer === preparation else { + throw CancellationError() + } if let mcpHTTPServerFactory { - let mcpHTTPServer = mcpHTTPServerFactory(store, mcpHTTPServerConfiguration) - startedHTTPServer = mcpHTTPServer - try await mcpHTTPServer.start() - self.mcpHTTPServer = mcpHTTPServer - } - store.transitionToRunning(serverURL: await self.mcpHTTPServer?.url) - let authSnapshot = try await backend.readAuth() - applyAuthSnapshot(authSnapshot, to: store.auth) - await refreshSelectedAccountRateLimits(auth: store.auth) - logger.info("Review runtime started") + guard let attachedStore else { + throw CancellationError() + } + mcpHTTPServer = mcpHTTPServerFactory( + attachedStore, + mcpHTTPServerConfiguration + ) + } + preparingMCPServer = nil + preparedMCPServer = preparation + return preparation } catch { - let failureMessage = await runtimeStartupFailureMessage(for: error) - logger.error("Review runtime failed to start: \(failureMessage, privacy: .public)") - self.client = nil - self.appServerBackend = nil - self.mcpHTTPServer = nil - authNotificationTask?.cancel() - authNotificationTask = nil - await stopMCPHTTPServer( - startedHTTPServer, - context: "runtime startup cleanup" + if preparingMCPServer === preparation { + preparingMCPServer = nil + } + throw CodexReviewAPI.Error.io(await runtimeStartupFailureMessage(for: error)) + } + } + + func activate( + _ preparation: PreparedMCPServer + ) async throws -> MCPServerPublicationSnapshot { + guard preparedMCPServer === preparation else { + throw CancellationError() + } + guard let server = mcpHTTPServer else { + preparedMCPServer = nil + runningMCPServer = preparation + return .init(serverURL: nil) + } + do { + try await server.start() + guard mcpHTTPServer === server, + preparedMCPServer === preparation + else { + throw CancellationError() + } + let serverURL = await server.url + guard mcpHTTPServer === server, + preparedMCPServer === preparation + else { + throw CancellationError() + } + preparedMCPServer = nil + runningMCPServer = preparation + return .init(serverURL: serverURL) + } catch { + if mcpHTTPServer === server { + mcpHTTPServer = nil + preparedMCPServer = nil + try? await server.stop() + } + throw error + } + } + + func stop() async throws { + preparingMCPServer = nil + preparedMCPServer = nil + runningMCPServer = nil + guard let server = mcpHTTPServer else { + return + } + mcpHTTPServer = nil + try await server.stop() + } + + func prepareRuntime( + generation _: ReviewRuntimeGeneration, + purpose _: ReviewRuntimeTransitionPurpose + ) async throws -> PreparedRuntime { + logger.info("Preparing review runtime") + let runtime = try await appServerRuntimeFactory(codexHomeURL) + do { + let authNotificationStream = await runtime.client.notificationStream() + let authentication = try await runtime.backend.readAuth() + let settings = try await Self.monitorSettings(from: runtime.backend.readSettings()) + let handle = LiveRuntimeLifecycleHandle( + owner: self, + client: runtime.client, + backend: runtime.backend, + authNotificationStream: authNotificationStream, + snapshot: .init( + authentication: authentication, + settings: settings + ) ) + logger.info("Review runtime prepared") + return .init(snapshot: handle.snapshot, handle: handle) + } catch { await closeAppServerRuntime( - backend: startedBackend, - fallbackClient: startedClient, - context: "runtime startup cleanup" + backend: runtime.backend, + fallbackClient: runtime.client, + context: "runtime preparation cleanup" + ) + throw error + } + } + + func commitRuntimePublication( + _ snapshot: RuntimePublicationSnapshot, + handle: any RuntimeLifecycleHandle, + auth: CodexReviewAuthModel + ) throws { + guard let handle = handle as? LiveRuntimeLifecycleHandle, + activeRuntimeHandle === handle, + let attachedStore + else { + throw CancellationError() + } + applyRuntimeAuthenticationSnapshot(snapshot.authentication, to: auth) + if let activeAccountID = snapshot.authentication.activeAccountID?.rawValue, + let account = auth.persistedAccounts.first(where: { + $0.accountKey == CodexAccount.normalizedEmail(activeAccountID) + }) + { + try? CodexReviewAccountRegistry.saveAccounts( + auth.persistedAccounts, + activeAccountKey: account.accountKey, + codexHomeURL: codexHomeURL + ) + try? CodexReviewAccountRegistry.saveSharedAuth( + for: account, + codexHomeURL: codexHomeURL ) - store.transitionToFailed(failureMessage) } + acceptsRuntimeRequests = true + observeAuthNotifications( + stream: handle.authNotificationStream, + backend: handle.backend, + handle: handle, + store: attachedStore + ) + handle.initialRateLimitTask = Task { @MainActor [weak self, weak auth] in + guard let self, let auth else { + return + } + await self.refreshSelectedAccountRateLimits( + auth: auth, + expectedRuntimeHandle: handle + ) + } + } + + func waitForRuntimePublication( + handle: any RuntimeLifecycleHandle + ) async { + guard let handle = handle as? LiveRuntimeLifecycleHandle else { + return + } + await handle.initialRateLimitTask?.value + } + + func activateRuntime(_ handle: LiveRuntimeLifecycleHandle) throws { + guard activeRuntimeHandle == nil, attachedStore != nil else { + throw CancellationError() + } + activeRuntimeHandle = handle + acceptsRuntimeRequests = false + client = handle.client + appServerBackend = handle.backend + settingsSnapshot = handle.snapshot.settings + } + + func closeRuntimeAdmission(_ handle: LiveRuntimeLifecycleHandle) { + guard activeRuntimeHandle === handle else { + return + } + acceptsRuntimeRequests = false + } + + func deactivateRuntime( + _ handle: LiveRuntimeLifecycleHandle + ) -> Task? { + guard activeRuntimeHandle === handle else { + return nil + } + activeRuntimeHandle = nil + acceptsRuntimeRequests = false + client = nil + appServerBackend = nil + let task = authNotificationTask + authNotificationTask = nil + task?.cancel() + return task } private func runtimeStartupFailureMessage(for error: Error) async -> String { @@ -509,15 +665,12 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { store: CodexReviewStore, intent: ReviewRuntimeTeardownIntent ) async { - let client = client let appServerBackend = appServerBackend - let mcpHTTPServer = mcpHTTPServer - let hasRuntimeState = client != nil || appServerBackend != nil || mcpHTTPServer != nil let loginCleanup = takeLoginRuntimeForCleanup() - guard hasRuntimeState || loginCleanup.isEmpty == false else { + guard appServerBackend != nil || loginCleanup.isEmpty == false else { return } - logger.info("Stopping review runtime for \(intent.diagnosticContext, privacy: .public)") + logger.info("Stopping review runtime semantic work for \(intent.diagnosticContext, privacy: .public)") if let appServerBackend { await cancelActiveReviewsForRuntimeTeardown( store: store, @@ -526,25 +679,14 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { timeoutWarning: intent.cleanupTimeoutWarning ) } - self.client = nil - self.mcpHTTPServer = nil - authNotificationTask?.cancel() - authNotificationTask = nil - await stopMCPHTTPServer(mcpHTTPServer, context: intent.diagnosticContext) - self.appServerBackend = nil await cleanupLoginRuntime(loginCleanup) - await closeAppServerRuntime( - backend: appServerBackend, - fallbackClient: client, - context: intent.diagnosticContext - ) - logger.info("Review runtime stopped after \(intent.diagnosticContext, privacy: .public)") + logger.info("Review runtime semantic work stopped after \(intent.diagnosticContext, privacy: .public)") } func waitUntilStopped() async {} func refreshSettings() async throws -> CodexReviewSettings.Snapshot { - guard let appServerBackend else { + guard activeRuntimeHandle != nil, let appServerBackend else { return settingsSnapshot } settingsSnapshot = try await Self.monitorSettings(from: appServerBackend.readSettings()) @@ -558,7 +700,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { serviceTier: CodexReviewSettings.ServiceTier?, persistServiceTier: Bool ) async throws { - guard let appServerBackend else { + guard activeRuntimeHandle != nil, let appServerBackend else { return } var change = CodexReviewBackendModel.Settings.Change( @@ -579,7 +721,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { func updateSettingsReasoningEffort( _ reasoningEffort: CodexReviewSettings.ReasoningEffort? ) async throws { - guard let appServerBackend else { + guard activeRuntimeHandle != nil, let appServerBackend else { return } settingsSnapshot = try await Self.monitorSettings( @@ -593,7 +735,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { func updateSettingsServiceTier( _ serviceTier: CodexReviewSettings.ServiceTier? ) async throws { - guard let appServerBackend else { + guard activeRuntimeHandle != nil, let appServerBackend else { return } settingsSnapshot = try await Self.monitorSettings( @@ -605,14 +747,29 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } func refreshAuth(auth: CodexReviewAuthModel) async { + guard acceptsRuntimeRequests, + let expectedRuntimeHandle = activeRuntimeHandle + else { + return + } do { guard let appServerBackend else { auth.updatePhase(.signedOut) return } let snapshot = try await appServerBackend.readAuth() + guard activeRuntimeHandle === expectedRuntimeHandle, + acceptsRuntimeRequests + else { + return + } applyAuthSnapshot(snapshot, to: auth) } catch { + guard activeRuntimeHandle === expectedRuntimeHandle, + acceptsRuntimeRequests + else { + return + } auth.updatePhase(.failed(message: error.localizedDescription)) } } @@ -680,12 +837,11 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { ) auth.selectPersistedAccount(auth.persistedAccounts.first(where: { $0.accountKey == accountKey })?.id) auth.updatePhase(.signedOut) - guard let attachedStore, appServerBackend != nil else { + guard let attachedStore else { return } await attachedStore.closeActiveReviewSessions(reason: .system(message: "Account switched.")) - await stop(store: attachedStore) - await start(store: attachedStore, forceRestartIfNeeded: true) + await attachedStore.recycleRuntimeAfterAccountChange() } func removeAccount(auth: CodexReviewAuthModel, accountKey: String) async throws { @@ -717,12 +873,11 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { if removedActiveAccount { auth.selectPersistedAccount(nil) auth.updatePhase(.signedOut) - guard let attachedStore, appServerBackend != nil else { + guard let attachedStore else { return } await attachedStore.closeActiveReviewSessions(reason: .system(message: "Account removed.")) - await stop(store: attachedStore) - await start(store: attachedStore, forceRestartIfNeeded: true) + await attachedStore.recycleRuntimeAfterAccountChange() } } @@ -755,7 +910,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { auth.selectPersistedAccount(nil) return } - let shouldRecycleRuntime = attachedStore != nil && appServerBackend != nil + let shouldRecycleRuntime = attachedStore != nil if shouldRecycleRuntime { await attachedStore?.closeActiveReviewSessions(reason: .system(message: "Signed out.")) } @@ -777,8 +932,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { auth.selectPersistedAccount(nil) auth.applyPersistedAccountStates(remaining.map(savedAccountPayload(from:)), activeAccountKey: nil) if shouldRecycleRuntime, let attachedStore { - await stop(store: attachedStore) - await start(store: attachedStore, forceRestartIfNeeded: true) + await attachedStore.recycleRuntimeAfterAccountChange() } } @@ -794,6 +948,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } private func startLogin(auth: CodexReviewAuthModel, activation: LoginActivation) async { + let expectedRuntimeHandle = activeRuntimeHandle var isolatedLoginClient: AppServerClient? var isolatedLoginCodexHomeURL: URL? do { @@ -803,6 +958,10 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { let loginClient = runtime.usesPrimaryRuntime ? nil : runtime.client isolatedLoginClient = loginClient isolatedLoginCodexHomeURL = loginCodexHomeURL + guard isCurrentRuntime(expectedRuntimeHandle) else { + await closeIsolatedLoginRuntime(client: loginClient, codexHomeURL: loginCodexHomeURL) + return + } guard runtime.usesPrimaryRuntime || self.appServerBackend != nil else { logger.error("Cannot start login because review runtime is not running") updateAuthenticationFailure( @@ -817,6 +976,14 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { let challenge = try await appServerBackend.startLogin(.init( nativeWebAuthenticationCallbackScheme: nativeAuthenticationConfiguration?.callbackScheme )) + guard let expectedRuntimeHandle, + activeRuntimeHandle === expectedRuntimeHandle, + acceptsRuntimeRequests + else { + try? await appServerBackend.cancelLogin(challenge) + await closeIsolatedLoginRuntime(client: loginClient, codexHomeURL: loginCodexHomeURL) + return + } loginChallenge = challenge loginBackend = appServerBackend self.loginClient = loginClient @@ -863,6 +1030,15 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { nativeAuthenticationConfiguration.browserSessionPolicy, nativeAuthenticationConfiguration.presentationAnchorProvider ) + guard activeRuntimeHandle === expectedRuntimeHandle, + acceptsRuntimeRequests, + loginChallenge?.id == challenge.id + else { + await session.cancel() + try? await appServerBackend.cancelLogin(challenge) + await closeIsolatedLoginRuntime(client: loginClient, codexHomeURL: loginCodexHomeURL) + return + } activeAuthenticationSession = session authenticationTask = Task { @MainActor [weak self, weak auth] in guard let self, let auth else { @@ -903,6 +1079,15 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } } + private func isCurrentRuntime( + _ expectedRuntimeHandle: LiveRuntimeLifecycleHandle? + ) -> Bool { + guard let expectedRuntimeHandle else { + return activeRuntimeHandle == nil + } + return activeRuntimeHandle === expectedRuntimeHandle && acceptsRuntimeRequests + } + private func monitorAuthenticationSession( challenge: CodexReviewBackendModel.Login.Challenge, session: any CodexReviewNativeAuthentication.WebSession, @@ -999,7 +1184,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private func loginRuntime(for activation: LoginActivation) async throws -> LoginRuntime { switch activation { case .activateAuthenticatedAccount: - guard let client, let appServerBackend else { + guard acceptsRuntimeRequests, let client, let appServerBackend else { throw CodexReviewAPI.Error.io("Review runtime is not running.") } return .init( @@ -1056,14 +1241,14 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { _ request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { - guard let appServerBackend else { + guard acceptsRuntimeRequests, let appServerBackend else { throw CodexReviewAPI.Error.io("Review runtime is not running.") } return try await appServerBackend.startReview(request, admission: admission) } func interruptReview(_ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason) async throws { - guard let appServerBackend else { + guard acceptsRuntimeRequests, let appServerBackend else { throw CodexReviewAPI.Error.io("Review runtime is not running.") } try await appServerBackend.interruptReview(run, reason: reason) @@ -1073,7 +1258,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { _ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { - guard let appServerBackend else { + guard acceptsRuntimeRequests, let appServerBackend else { throw CodexReviewAPI.Error.io("Review runtime is not running.") } return try await appServerBackend.beginReviewRecovery(run, reason: reason) @@ -1083,14 +1268,14 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { _ token: CodexReviewBackendModel.Review.RecoveryToken, request: CodexReviewBackendModel.Review.Start ) async throws -> BackendReviewAttempt { - guard let appServerBackend else { + guard acceptsRuntimeRequests, let appServerBackend else { throw CodexReviewAPI.Error.io("Review runtime is not running.") } return try await appServerBackend.resumeReviewRecovery(token, request: request) } func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { - guard let appServerBackend else { + guard activeRuntimeHandle != nil, let appServerBackend else { throw ReviewRuntimeCloseFailure.cleanup("Review runtime is not running.") } try await appServerBackend.cleanupReview(run) @@ -1160,8 +1345,9 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } private func observeAuthNotifications( - client: AppServerClient, + stream: AsyncThrowingStream, backend: AppServerCodexReviewBackend, + handle: LiveRuntimeLifecycleHandle, store: CodexReviewStore ) { authNotificationTask?.cancel() @@ -1169,52 +1355,72 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { guard let self, let store else { return } - let stream = await client.notificationStream() do { for try await notification in stream { + guard self.activeRuntimeHandle === handle, + self.acceptsRuntimeRequests + else { + return + } await self.handleAuthNotification( notification, backend: backend, + expectedRuntimeHandle: handle, auth: store.auth ) } } catch is CancellationError { } catch { logger.error("Auth notification stream ended: \(error.localizedDescription, privacy: .public)") - markRuntimeFailedAfterNotificationStreamError(error, store: store) + markRuntimeFailedAfterNotificationStreamError( + error, + handle: handle, + store: store + ) } } } private func markRuntimeFailedAfterNotificationStreamError( _ error: any Error, + handle: LiveRuntimeLifecycleHandle, store: CodexReviewStore ) { - guard client != nil - || appServerBackend != nil - || mcpHTTPServer != nil - || loginClient != nil - || loginCodexHomeURL != nil - || activeAuthenticationSession != nil - else { + guard activeRuntimeHandle === handle else { return } authNotificationTask = nil - store.requestRuntimeTeardown( - intent: .unexpectedFailure(error.localizedDescription) + store.requestRuntimeFailure( + handle: handle, + cause: error.localizedDescription ) } private func handleAuthNotification( _ notification: JSONRPC.Notification, backend: AppServerCodexReviewBackend, + expectedRuntimeHandle: LiveRuntimeLifecycleHandle, auth: CodexReviewAuthModel ) async { + guard activeRuntimeHandle === expectedRuntimeHandle, + acceptsRuntimeRequests + else { + return + } switch notification.method { case "account/login/completed": - await handleLoginCompletedNotification(notification, backend: backend, auth: auth) + await handleLoginCompletedNotification( + notification, + backend: backend, + expectedRuntimeHandle: expectedRuntimeHandle, + auth: auth + ) case "account/updated": - await handleAccountUpdatedNotification(backend: backend, auth: auth) + await handleAccountUpdatedNotification( + backend: backend, + expectedRuntimeHandle: expectedRuntimeHandle, + auth: auth + ) case "account/rateLimits/updated": await applyRateLimitsUpdatedNotification(notification, auth: auth) default: @@ -1268,10 +1474,22 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private func handleLoginCompletedNotification( _ notification: JSONRPC.Notification, backend: AppServerCodexReviewBackend, + expectedRuntimeHandle: LiveRuntimeLifecycleHandle? = nil, auth: CodexReviewAuthModel ) async { + if let expectedRuntimeHandle { + guard activeRuntimeHandle === expectedRuntimeHandle, + acceptsRuntimeRequests + else { + return + } + } guard notification.method == "account/login/completed" else { - await handleAccountUpdatedNotification(backend: backend, auth: auth) + await handleAccountUpdatedNotification( + backend: backend, + expectedRuntimeHandle: expectedRuntimeHandle, + auth: auth + ) return } do { @@ -1287,6 +1505,13 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { authenticationTask?.cancel() authenticationTask = nil await activeAuthenticationSession?.cancel() + if let expectedRuntimeHandle { + guard activeRuntimeHandle === expectedRuntimeHandle, + acceptsRuntimeRequests + else { + return + } + } guard payload.success else { updateAuthenticationFailure( payload.error ?? "Authentication failed.", @@ -1311,17 +1536,34 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private func handleAccountUpdatedNotification( backend: AppServerCodexReviewBackend, + expectedRuntimeHandle: LiveRuntimeLifecycleHandle? = nil, auth: CodexReviewAuthModel ) async { + if let expectedRuntimeHandle { + guard activeRuntimeHandle === expectedRuntimeHandle, + acceptsRuntimeRequests + else { + return + } + } guard isWaitingForLoginAccountUpdate else { - await refreshAuthAfterAccountNotification(backend: backend, auth: auth) + await refreshAuthAfterAccountNotification( + backend: backend, + expectedRuntimeHandle: expectedRuntimeHandle, + auth: auth + ) return } - await finishCompletedLoginAfterAccountUpdate(backend: backend, auth: auth) + await finishCompletedLoginAfterAccountUpdate( + backend: backend, + expectedRuntimeHandle: expectedRuntimeHandle, + auth: auth + ) } private func finishCompletedLoginAfterAccountUpdate( backend: AppServerCodexReviewBackend, + expectedRuntimeHandle: LiveRuntimeLifecycleHandle? = nil, auth: CodexReviewAuthModel ) async { let activation = loginActivation @@ -1335,8 +1577,16 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { authenticationTask?.cancel() authenticationTask = nil await activeAuthenticationSession?.cancel() + let snapshot = try await backend.readAuth() + if let expectedRuntimeHandle { + guard activeRuntimeHandle === expectedRuntimeHandle, + acceptsRuntimeRequests + else { + return + } + } let account = applyAuthSnapshot( - try await backend.readAuth(), + snapshot, to: auth, activation: activation, authSourceCodexHomeURL: loginCodexHomeURL @@ -1370,12 +1620,32 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private func refreshAuthAfterAccountNotification( backend: AppServerCodexReviewBackend, + expectedRuntimeHandle requestedRuntimeHandle: LiveRuntimeLifecycleHandle? = nil, auth: CodexReviewAuthModel ) async { + let expectedRuntimeHandle = requestedRuntimeHandle ?? activeRuntimeHandle + guard acceptsRuntimeRequests, + let expectedRuntimeHandle, + activeRuntimeHandle === expectedRuntimeHandle, + expectedRuntimeHandle.backend === backend + else { + return + } do { - applyAuthSnapshot(try await backend.readAuth(), to: auth) + let snapshot = try await backend.readAuth() + guard activeRuntimeHandle === expectedRuntimeHandle, + acceptsRuntimeRequests + else { + return + } + applyAuthSnapshot(snapshot, to: auth) await refreshSelectedAccountRateLimits(auth: auth) } catch { + guard activeRuntimeHandle === expectedRuntimeHandle, + acceptsRuntimeRequests + else { + return + } auth.updatePhase(.failed(message: error.localizedDescription)) } } @@ -1410,14 +1680,25 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } } - private func refreshSelectedAccountRateLimits(auth: CodexReviewAuthModel) async { + private func refreshSelectedAccountRateLimits( + auth: CodexReviewAuthModel, + expectedRuntimeHandle: LiveRuntimeLifecycleHandle? = nil + ) async { guard let selectedAccount = auth.selectedAccount else { return } - await refreshRateLimits(for: selectedAccount, auth: auth) + await refreshRateLimits( + for: selectedAccount, + auth: auth, + expectedRuntimeHandle: expectedRuntimeHandle + ) } - private func refreshRateLimits(for account: CodexAccount, auth: CodexReviewAuthModel) async { + private func refreshRateLimits( + for account: CodexAccount, + auth: CodexReviewAuthModel, + expectedRuntimeHandle requestedRuntimeHandle: LiveRuntimeLifecycleHandle? = nil + ) async { guard account.capabilities.supportsRateLimitRefresh else { return } @@ -1425,7 +1706,24 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { await refreshSavedAccountRateLimits(for: account) return } - let didRefresh = await refreshRateLimits(for: account, using: appServerBackend, source: "active-runtime") + let expectedRuntimeHandle = requestedRuntimeHandle ?? activeRuntimeHandle + guard acceptsRuntimeRequests, + let expectedRuntimeHandle, + activeRuntimeHandle === expectedRuntimeHandle + else { + return + } + let didRefresh = await refreshRateLimits( + for: account, + using: appServerBackend, + source: "active-runtime", + expectedRuntimeHandle: expectedRuntimeHandle + ) + guard activeRuntimeHandle === expectedRuntimeHandle, + acceptsRuntimeRequests + else { + return + } if didRefresh { persistRefreshedSharedAuth( from: codexHomeURL, @@ -1481,7 +1779,8 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private func refreshRateLimits( for account: CodexAccount, using backend: AppServerCodexReviewBackend?, - source: String + source: String, + expectedRuntimeHandle: LiveRuntimeLifecycleHandle? = nil ) async -> Bool { do { guard let backend else { @@ -1494,6 +1793,13 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { ) } let response = try await backend.readRateLimits() + if let expectedRuntimeHandle { + guard activeRuntimeHandle === expectedRuntimeHandle, + acceptsRuntimeRequests + else { + return false + } + } applyRateLimits( windows: response.codexRateLimitWindows, planType: response.codexPlanType, @@ -1505,6 +1811,11 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { ) return true } catch { + if let expectedRuntimeHandle, + (activeRuntimeHandle !== expectedRuntimeHandle || acceptsRuntimeRequests == false) + { + return false + } recordRateLimitRefreshFailure(error, account: account) try? CodexReviewAccountRegistry.updateCachedRateLimits( from: account, @@ -1578,22 +1889,6 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } } - private func stopMCPHTTPServer( - _ server: (any CodexReviewMCPHTTPServing)?, - context: String - ) async { - guard let server else { - return - } - do { - try await server.stop() - } catch { - logger.error( - "Failed to stop MCP HTTP server during \(context, privacy: .public): \(error.localizedDescription, privacy: .public)" - ) - } - } - private func closeClientRecordingFailure( _ client: AppServerClient?, context: String @@ -1722,6 +2017,82 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } } +@MainActor +private final class LiveRuntimeLifecycleHandle: RuntimeLifecycleHandle { + fileprivate let client: AppServerClient + fileprivate let backend: AppServerCodexReviewBackend + fileprivate let authNotificationStream: AsyncThrowingStream + fileprivate let snapshot: RuntimePublicationSnapshot + fileprivate var initialRateLimitTask: Task? + + private weak var owner: LiveCodexReviewStoreBackend? + private var isActivated = false + private var closeTask: Task, Never>? + + init( + owner: LiveCodexReviewStoreBackend, + client: AppServerClient, + backend: AppServerCodexReviewBackend, + authNotificationStream: AsyncThrowingStream, + snapshot: RuntimePublicationSnapshot + ) { + self.owner = owner + self.client = client + self.backend = backend + self.authNotificationStream = authNotificationStream + self.snapshot = snapshot + } + + func activate() async throws { + guard isActivated == false, closeTask == nil, let owner else { + throw CancellationError() + } + try owner.activateRuntime(self) + isActivated = true + } + + func closeAdmission() async { + owner?.closeRuntimeAdmission(self) + } + + func close(purpose _: ReviewRuntimeTransitionPurpose) async throws { + let task: Task, Never> + if let closeTask { + task = closeTask + } else { + let authObservationTask = owner?.deactivateRuntime(self) + let initialRateLimitTask = initialRateLimitTask + self.initialRateLimitTask = nil + let lifecycle = backend.runtimeOwnerLifecycleHandle + let created = Task, Never> { @MainActor in + authObservationTask?.cancel() + initialRateLimitTask?.cancel() + let result: Result + do { + await lifecycle.closeAdmission() + try await lifecycle.closeAndWait() + result = .success(()) + } catch { + result = .failure(error) + } + await authObservationTask?.value + await initialRateLimitTask?.value + return result + } + closeTask = created + task = created + } + try await task.value.get() + } + + func waitUntilClosed() async throws { + guard let closeTask else { + throw CancellationError() + } + try await closeTask.value.get() + } +} + @MainActor private struct AppServerRuntime: Sendable { var client: AppServerClient diff --git a/Sources/CodexReviewTesting/TestSupport.swift b/Sources/CodexReviewTesting/TestSupport.swift index 2bae5862..b22bbdb9 100644 --- a/Sources/CodexReviewTesting/TestSupport.swift +++ b/Sources/CodexReviewTesting/TestSupport.swift @@ -858,6 +858,150 @@ package struct StoreJobSnapshot: Sendable { package var cancellationRequested: Bool } +@MainActor +package final class TestingRuntimeLifecycleHandle: RuntimeLifecycleHandle { + package private(set) var activateCallCount = 0 + package private(set) var closeAdmissionCallCount = 0 + package private(set) var closePurposes: [ReviewRuntimeTransitionPurpose] = [] + package private(set) var waitUntilClosedCallCount = 0 + + private let onActivate: @MainActor @Sendable () -> Void + private let onClose: @MainActor @Sendable () -> Void + private var closeGate: AsyncGate? + private var closeStartedGate = AsyncGate() + private var didClose = false + + package init( + onActivate: @escaping @MainActor @Sendable () -> Void = {}, + onClose: @escaping @MainActor @Sendable () -> Void = {} + ) { + self.onActivate = onActivate + self.onClose = onClose + } + + package func activate() async throws { + activateCallCount += 1 + onActivate() + } + + package func closeAdmission() async { + closeAdmissionCallCount += 1 + } + + package func holdClose(with gate: AsyncGate) { + closeGate = gate + closeStartedGate = AsyncGate() + } + + package func waitForClose() async { + await closeStartedGate.wait() + } + + package func close(purpose: ReviewRuntimeTransitionPurpose) async throws { + closePurposes.append(purpose) + await closeStartedGate.open() + await closeGate?.waitIgnoringCancellation() + closeGate = nil + guard didClose == false else { + return + } + didClose = true + onClose() + } + + package func waitUntilClosed() async throws { + waitUntilClosedCallCount += 1 + guard didClose else { + throw CancellationError() + } + } +} + +@MainActor +package final class TestingMCPServerLifecycleOwner: MCPServerLifecycleOwner { + package private(set) var preparedServers: [PreparedMCPServer] = [] + package private(set) var activatedServers: [PreparedMCPServer] = [] + package private(set) var stopCallCount = 0 + + private let serverURL: URL? + private var preparedServer: PreparedMCPServer? + private var runningServer: PreparedMCPServer? + private var preparationGate: AsyncGate? + private var preparationStartedGate = AsyncGate() + private var preparationCancellationGate = AsyncGate() + private var stopGate: AsyncGate? + private var stopStartedGate = AsyncGate() + + package init(serverURL: URL? = nil) { + self.serverURL = serverURL + } + + package func holdPreparation(with gate: AsyncGate) { + preparationGate = gate + preparationStartedGate = AsyncGate() + preparationCancellationGate = AsyncGate() + } + + package func waitForPreparation() async { + await preparationStartedGate.wait() + } + + package func waitForPreparationCancellation() async { + await preparationCancellationGate.wait() + } + + package func holdStop(with gate: AsyncGate) { + stopGate = gate + stopStartedGate = AsyncGate() + } + + package func waitForStop() async { + await stopStartedGate.wait() + } + + package func prepare() async throws -> PreparedMCPServer { + let preparation = PreparedMCPServer() + preparedServer = preparation + preparedServers.append(preparation) + await preparationStartedGate.open() + if let preparationGate { + let cancellationGate = preparationCancellationGate + await withTaskCancellationHandler { + await preparationGate.waitIgnoringCancellation() + } onCancel: { + Task { await cancellationGate.open() } + } + self.preparationGate = nil + } + try Task.checkCancellation() + return preparation + } + + package func activate( + _ preparation: PreparedMCPServer + ) async throws -> MCPServerPublicationSnapshot { + guard preparedServer === preparation else { + throw CancellationError() + } + preparedServer = nil + runningServer = preparation + activatedServers.append(preparation) + return .init(serverURL: serverURL) + } + + package func stop() async throws { + guard preparedServer != nil || runningServer != nil else { + return + } + await stopStartedGate.open() + await stopGate?.waitIgnoringCancellation() + stopGate = nil + preparedServer = nil + runningServer = nil + stopCallCount += 1 + } +} + @MainActor package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { package let reviewBackend: FakeCodexReviewBackend @@ -865,14 +1009,23 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { package var currentSettingsSnapshot: CodexReviewSettings.Snapshot package private(set) var isActive = false package private(set) var startRequests: [Bool] = [] + package let mcpServerLifecycle: any MCPServerLifecycleOwner + package private(set) var lastPreparedRuntimeHandle: TestingRuntimeLifecycleHandle? + private var runtimePreparationGate: AsyncGate? + private var runtimePreparationFailureMessage: String? + private var throwsCancellationAfterHeldRuntimePreparation = false + private var runtimePreparationStartedGate = AsyncGate() + private var runtimePreparationCancellationGate = AsyncGate() package init( reviewBackend: FakeCodexReviewBackend, - seed: CodexReviewStoreSeed = .init() + seed: CodexReviewStoreSeed = .init(), + mcpServerLifecycle: (any MCPServerLifecycleOwner)? = nil ) { self.reviewBackend = reviewBackend self.seed = seed self.currentSettingsSnapshot = seed.initialSettingsSnapshot + self.mcpServerLifecycle = mcpServerLifecycle ?? NoMCPServerLifecycleOwner() } package var initialSettingsSnapshot: CodexReviewSettings.Snapshot { @@ -881,10 +1034,65 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { package func attachStore(_: CodexReviewStore) {} - package func start(store: CodexReviewStore, forceRestartIfNeeded: Bool) async { - startRequests.append(forceRestartIfNeeded) - isActive = true - store.transitionToRunning(serverURL: nil) + package func holdRuntimePreparation(with gate: AsyncGate) { + runtimePreparationGate = gate + runtimePreparationStartedGate = AsyncGate() + runtimePreparationCancellationGate = AsyncGate() + } + + package func waitForRuntimePreparation() async { + await runtimePreparationStartedGate.wait() + } + + package func waitForRuntimePreparationCancellation() async { + await runtimePreparationCancellationGate.wait() + } + + package func failNextRuntimePreparation(message: String) { + runtimePreparationFailureMessage = message + } + + package func throwCancellationAfterHeldRuntimePreparation() { + throwsCancellationAfterHeldRuntimePreparation = true + } + + package func prepareRuntime( + generation _: ReviewRuntimeGeneration, + purpose: ReviewRuntimeTransitionPurpose + ) async throws -> PreparedRuntime { + startRequests.append(purpose == .restartSameAccount) + let handle = TestingRuntimeLifecycleHandle( + onActivate: { [weak self] in self?.isActive = true }, + onClose: { [weak self] in self?.isActive = false } + ) + lastPreparedRuntimeHandle = handle + await runtimePreparationStartedGate.open() + if let runtimePreparationGate { + let cancellationGate = runtimePreparationCancellationGate + await withTaskCancellationHandler { + await runtimePreparationGate.waitIgnoringCancellation() + } onCancel: { + Task { await cancellationGate.open() } + } + self.runtimePreparationGate = nil + } + if throwsCancellationAfterHeldRuntimePreparation { + throwsCancellationAfterHeldRuntimePreparation = false + try Task.checkCancellation() + } + if let runtimePreparationFailureMessage { + self.runtimePreparationFailureMessage = nil + throw FakeCodexReviewBackendError(message: runtimePreparationFailureMessage) + } + let settings = try await monitoredSettingsSnapshot() + currentSettingsSnapshot = settings + return PreparedRuntime( + snapshot: .init( + authentication: try await reviewBackend.readAuth(), + settings: settings + ), + handle: handle + ) } package func stop(store _: CodexReviewStore) async { @@ -1023,15 +1231,19 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { } package func refreshSettings() async throws -> CodexReviewSettings.Snapshot { + currentSettingsSnapshot = try await monitoredSettingsSnapshot() + return currentSettingsSnapshot + } + + private func monitoredSettingsSnapshot() async throws -> CodexReviewSettings.Snapshot { let snapshot = try await reviewBackend.readSettings() - currentSettingsSnapshot = .init( + return .init( model: snapshot.model, fallbackModel: snapshot.fallbackModel, reasoningEffort: snapshot.reasoningEffort.flatMap(CodexReviewSettings.ReasoningEffort.init(rawValue:)), serviceTier: snapshot.serviceTier.flatMap(CodexReviewSettings.ServiceTier.init(rawValue:)), models: snapshot.models ) - return currentSettingsSnapshot } package func updateSettingsModel( diff --git a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift index 38e864b6..b0c840ee 100644 --- a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift +++ b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift @@ -382,6 +382,13 @@ struct CodexReviewHostTests { @Test func liveStoreStopsMCPServerAfterItsStartFails() async throws { let homeURL = try temporaryHome() let transport = FakeJSONRPCTransport() + try await transport.enqueue(AppServerAPI.Initialize.Response(), for: "initialize") + try await transport.enqueue(AppServerAPI.Account.Read.Response(), for: "account/read") + try await transport.enqueue( + AppServerAPI.Config.Read.Response(config: .init(model: "gpt-5")), + for: "config/read" + ) + try await transport.enqueue(AppServerAPI.Model.List.Response(data: []), for: "model/list") let server = ControlledMCPHTTPServer( endpoint: try #require(URL(string: "http://127.0.0.1:19434/mcp")), startFailure: .injected @@ -406,6 +413,44 @@ struct CodexReviewHostTests { #expect(server.stopCallCount == 1) } + @Test func liveSameAccountRestartRetainsMCPListenerAndURL() async throws { + let homeURL = try temporaryHome() + let firstTransport = FakeJSONRPCTransport() + let secondTransport = FakeJSONRPCTransport() + try await enqueueRuntimeStartResponses(firstTransport) + try await enqueueRuntimeStartResponses(secondTransport) + let server = ControlledMCPHTTPServer( + endpoint: try #require(URL(string: "http://127.0.0.1:19435/mcp")) + ) + var transports = [firstTransport, secondTransport] + var serverFactoryCallCount = 0 + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + mcpHTTPServerFactory: { _, _ in + serverFactoryCallCount += 1 + return server + }, + mcpHTTPServerBindChecker: { _ in }, + transportFactory: { _ in transports.removeFirst() } + ) + + await store.start() + let firstURL = store.serverURL + await store.restart() + + #expect(store.serverState == .running) + #expect(store.serverURL == firstURL) + #expect(serverFactoryCallCount == 1) + #expect(server.startCallCount == 1) + #expect(server.stopCallCount == 0) + #expect(await firstTransport.isClosedForTesting()) + #expect(transports.isEmpty) + + await store.stop() + #expect(server.stopCallCount == 1) + } + @Test func liveStoreReportsMCPPortOwnerWhenEndpointPortInUseAndDoesNotLaunchAppServer() async throws { let homeURL = try temporaryHome() let port = 54321 @@ -1224,6 +1269,99 @@ struct CodexReviewHostTests { #expect(await secondTransport.recordedRequests().map(\.method).contains("account/read")) } + @Test func accountSwitchDuringHeldLiveStartPublishesOnlyPostSwitchRuntime() async throws { + let homeURL = try temporaryHome() + let mainCodexHomeURL = homeURL.appendingPathComponent(".codex_review", isDirectory: true) + try writeRegistry( + homeURL: homeURL, + activeAccountKey: "first@example.com", + accounts: ["first@example.com", "second@example.com"] + ) + try writeSavedAccountAuth(homeURL: homeURL, accountKey: "first@example.com") + try writeSavedAccountAuth(homeURL: homeURL, accountKey: "second@example.com") + + let firstTransport = FakeJSONRPCTransport() + let secondTransport = FakeJSONRPCTransport() + try await enqueueRuntimeStartResponses(firstTransport, accountEmail: "first@example.com") + try await enqueueRuntimeStartResponses(secondTransport, accountEmail: "second@example.com") + let heldAuthRead = AsyncGate() + await firstTransport.holdNextIgnoringCancellation( + method: "account/read", + gate: heldAuthRead + ) + var transports = [firstTransport, secondTransport] + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + transportFactory: { codexHomeURL in + #expect(codexHomeURL == mainCodexHomeURL) + return transports.removeFirst() + } + ) + + let initialStart = Task { @MainActor in await store.start() } + await firstTransport.waitForRequestCount(2) + let accountSwitch = Task { @MainActor in + try await store.switchAccount(CodexAccount(email: "second@example.com")) + } + try #require(await waitUntil(timeout: .seconds(2)) { + store.auth.selectedAccount?.accountKey == "second@example.com" + }) + + await heldAuthRead.open() + try await accountSwitch.value + await initialStart.value + + #expect(store.serverState == .running) + #expect(store.auth.selectedAccount?.accountKey == "second@example.com") + #expect(await firstTransport.isClosedForTesting()) + #expect(transports.isEmpty) + await store.stop() + } + + @Test func staleAccountNotificationCannotOverwriteReplacementRuntime() async throws { + let homeURL = try temporaryHome() + let firstTransport = FakeJSONRPCTransport() + let secondTransport = FakeJSONRPCTransport() + try await enqueueRuntimeStartResponses(firstTransport, accountEmail: "first@example.com") + try await enqueueRuntimeStartResponses(secondTransport, accountEmail: "second@example.com") + var transports = [firstTransport, secondTransport] + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + transportFactory: { _ in transports.removeFirst() } + ) + await store.start() + await firstTransport.waitForNotificationStreamCount(1) + + try await firstTransport.enqueue( + AppServerAPI.Account.Read.Response( + account: .init(email: "stale@example.com", planType: "pro") + ), + for: "account/read" + ) + let staleReadGate = AsyncGate() + await firstTransport.holdNextIgnoringCancellation( + method: "account/read", + gate: staleReadGate + ) + let requestCount = await firstTransport.recordedRequests().count + try await firstTransport.emitServerNotification( + method: "account/updated", + params: EmptyResponse() + ) + await firstTransport.waitForRequestCount(requestCount + 1) + + let restart = Task { @MainActor in await store.restart() } + await staleReadGate.open() + await restart.value + + #expect(store.serverState == .running) + #expect(store.auth.selectedAccount?.accountKey == "second@example.com") + #expect(store.auth.accounts.contains { $0.accountKey == "stale@example.com" } == false) + await store.stop() + } + @Test func liveStoreSignOutRestartsRuntimeAndCancelsRunningReviews() async throws { let homeURL = try temporaryHome() let mainCodexHomeURL = homeURL.appendingPathComponent(".codex_review", isDirectory: true) @@ -2166,6 +2304,35 @@ private final class FakeWebAuthenticationSession: CodexReviewNativeAuthenticatio } } +private func enqueueRuntimeStartResponses( + _ transport: FakeJSONRPCTransport, + accountEmail: String? = nil +) async throws { + try await transport.enqueue(AppServerAPI.Initialize.Response(), for: "initialize") + if let accountEmail { + try await transport.enqueue( + AppServerAPI.Account.Read.Response( + account: .init(email: accountEmail, planType: "pro") + ), + for: "account/read" + ) + try await transport.enqueue( + AppServerAPI.Account.RateLimits.Response(rateLimits: .init( + limitID: "codex", + primary: .init(usedPercent: 10, windowDurationMins: 300) + )), + for: "account/rateLimits/read" + ) + } else { + try await transport.enqueue(AppServerAPI.Account.Read.Response(), for: "account/read") + } + try await transport.enqueue( + AppServerAPI.Config.Read.Response(config: .init(model: "gpt-5")), + for: "config/read" + ) + try await transport.enqueue(AppServerAPI.Model.List.Response(data: []), for: "model/list") +} + private func temporaryHome() throws -> URL { let url = FileManager.default.temporaryDirectory .appendingPathComponent("codex-review-host-tests-\(UUID().uuidString)", isDirectory: true) diff --git a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift new file mode 100644 index 00000000..1a75f9f7 --- /dev/null +++ b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift @@ -0,0 +1,465 @@ +import Foundation +import Testing +import CodexReview +import CodexReviewTesting + +@Suite("store runtime lifecycle", .serialized) +@MainActor +struct CodexReviewStoreLifecycleTests { + @Test func stopInvalidatesHeldRuntimePreparationAndClosesStaleHandleOnce() async throws { + let preparationGate = AsyncGate() + let mcpOwner = TestingMCPServerLifecycleOwner() + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend(), + mcpServerLifecycle: mcpOwner + ) + backend.holdRuntimePreparation(with: preparationGate) + let store = CodexReviewStore.makeTestingStore(backend: backend) + + let start = Task { @MainActor in await store.start() } + await backend.waitForRuntimePreparation() + let staleHandle = try #require(backend.lastPreparedRuntimeHandle) + let stop = Task { @MainActor in await store.stop() } + await backend.waitForRuntimePreparationCancellation() + + #expect(staleHandle.activateCallCount == 0) + #expect(store.serverState == .starting) + + await preparationGate.open() + await stop.value + await start.value + + #expect(store.serverState == .stopped) + #expect(staleHandle.activateCallCount == 0) + #expect(staleHandle.closeAdmissionCallCount == 1) + #expect(staleHandle.closePurposes == [.start]) + #expect(staleHandle.waitUntilClosedCallCount == 1) + #expect(mcpOwner.stopCallCount == 1) + #expect(backend.isActive == false) + #expect(store.settings.lastErrorMessage == nil) + } + + @Test func forceRestartDuringHeldInitialPreparationPublishesOnlyFreshRuntime() async throws { + let preparationGate = AsyncGate() + let mcpOwner = TestingMCPServerLifecycleOwner() + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend(), + mcpServerLifecycle: mcpOwner + ) + backend.holdRuntimePreparation(with: preparationGate) + let store = CodexReviewStore.makeTestingStore(backend: backend) + + let initialStart = Task { @MainActor in await store.start() } + await backend.waitForRuntimePreparation() + let staleHandle = try #require(backend.lastPreparedRuntimeHandle) + let restart = Task { @MainActor in + await store.start(forceRestartIfNeeded: true) + } + await backend.waitForRuntimePreparationCancellation() + + await preparationGate.open() + await restart.value + await initialStart.value + + let currentHandle = try #require(backend.lastPreparedRuntimeHandle) + #expect(currentHandle !== staleHandle) + #expect(staleHandle.activateCallCount == 0) + #expect(staleHandle.closePurposes == [.start]) + #expect(currentHandle.activateCallCount == 1) + #expect(store.serverState == .running) + #expect(mcpOwner.preparedServers.count == 2) + #expect(mcpOwner.activatedServers.count == 1) + #expect(mcpOwner.stopCallCount == 1) + #expect(store.settings.lastErrorMessage == nil) + await store.stop() + } + + @Test func acquisitionCancellationCatchConsumesCutoverWithoutError() async { + let preparationGate = AsyncGate() + let mcpOwner = TestingMCPServerLifecycleOwner() + mcpOwner.holdPreparation(with: preparationGate) + let store = CodexReviewStore.makeTestingStore(backend: TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend(), + mcpServerLifecycle: mcpOwner + )) + + let start = Task { @MainActor in await store.start() } + await mcpOwner.waitForPreparation() + let stop = Task { @MainActor in await store.stop() } + await mcpOwner.waitForPreparationCancellation() + await preparationGate.open() + await stop.value + await start.value + + #expect(store.serverState == .stopped) + #expect(store.settingsService.runtimeCutoverStatus == .awaitingRecovery) + #expect(store.settings.isLoading == false) + #expect(store.settings.lastErrorMessage == nil) + } + + @Test func accountRecycleDuringHeldInitialPreparationStartsFreshListenerAndRuntime() async throws { + let preparationGate = AsyncGate() + let mcpOwner = TestingMCPServerLifecycleOwner() + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend(), + mcpServerLifecycle: mcpOwner + ) + backend.holdRuntimePreparation(with: preparationGate) + let store = CodexReviewStore.makeTestingStore(backend: backend) + + let initialStart = Task { @MainActor in await store.start() } + await backend.waitForRuntimePreparation() + let staleHandle = try #require(backend.lastPreparedRuntimeHandle) + let recycle = Task { @MainActor in + await store.recycleRuntimeAfterAccountChange() + } + await backend.waitForRuntimePreparationCancellation() + + await preparationGate.open() + await recycle.value + await initialStart.value + + let currentHandle = try #require(backend.lastPreparedRuntimeHandle) + #expect(currentHandle !== staleHandle) + #expect(staleHandle.activateCallCount == 0) + #expect(currentHandle.activateCallCount == 1) + #expect(mcpOwner.preparedServers.count == 2) + #expect(mcpOwner.activatedServers.count == 1) + #expect(mcpOwner.stopCallCount == 1) + #expect(store.serverState == .running) + #expect(store.settings.lastErrorMessage == nil) + await store.stop() + } + + @Test func stopAdmittedBeforeRecycleTaskEntryInheritsCleanupOwnership() async throws { + let mcpOwner = TestingMCPServerLifecycleOwner() + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend(), + mcpServerLifecycle: mcpOwner + ) + let store = CodexReviewStore.makeTestingStore(backend: backend) + await store.start() + let retiringHandle = try #require(backend.lastPreparedRuntimeHandle) + + let recycle = try #require(store.admitRuntimeRecycleAfterAccountChange()) + store.requestRuntimeTeardown(intent: .explicitStop) + await store.stop() + await recycle.value + + #expect(store.serverState == .stopped) + #expect(retiringHandle.closePurposes == [.stop]) + #expect(retiringHandle.waitUntilClosedCallCount == 1) + #expect(mcpOwner.stopCallCount == 1) + #expect(store.settings.lastErrorMessage == nil) + } + + @Test func recycleSuccessorAdmittedBeforeTaskEntryInheritsCleanupOwnership() async throws { + let mcpOwner = TestingMCPServerLifecycleOwner() + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend(), + mcpServerLifecycle: mcpOwner + ) + let store = CodexReviewStore.makeTestingStore(backend: backend) + await store.start() + let retiringHandle = try #require(backend.lastPreparedRuntimeHandle) + + let recycle = try #require(store.admitRuntimeRecycleAfterAccountChange()) + let successor = try #require(store.admitRuntimeRecycleAfterAccountChange()) + await successor.value + await recycle.value + + let currentHandle = try #require(backend.lastPreparedRuntimeHandle) + #expect(currentHandle !== retiringHandle) + #expect(retiringHandle.closePurposes == [.stop]) + #expect(retiringHandle.waitUntilClosedCallCount == 1) + #expect(mcpOwner.preparedServers.count == 2) + #expect(mcpOwner.stopCallCount == 1) + #expect(store.serverState == .running) + #expect(store.settings.lastErrorMessage == nil) + await store.stop() + } + + @Test func sameAccountRestartRetainsMCPListenerAndURL() async throws { + let endpoint = try #require(URL(string: "http://127.0.0.1:19417/mcp")) + let mcpOwner = TestingMCPServerLifecycleOwner(serverURL: endpoint) + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend(), + mcpServerLifecycle: mcpOwner + ) + let store = CodexReviewStore.makeTestingStore(backend: backend) + + await store.start() + let firstHandle = try #require(backend.lastPreparedRuntimeHandle) + let listener = try #require(mcpOwner.preparedServers.first) + + await store.restart() + + let secondHandle = try #require(backend.lastPreparedRuntimeHandle) + #expect(secondHandle !== firstHandle) + #expect(firstHandle.closePurposes == [.restartSameAccount]) + #expect(secondHandle.activateCallCount == 1) + #expect(backend.startRequests == [false, true]) + #expect(store.serverState == .running) + #expect(store.serverURL == endpoint) + #expect(mcpOwner.preparedServers.count == 1) + #expect(mcpOwner.preparedServers.first === listener) + #expect(mcpOwner.activatedServers.first === listener) + #expect(mcpOwner.stopCallCount == 0) + await store.stop() + } + + @Test func stopInvalidatesHeldReplacementBeforePublication() async throws { + let endpoint = try #require(URL(string: "http://127.0.0.1:19422/mcp")) + let mcpOwner = TestingMCPServerLifecycleOwner(serverURL: endpoint) + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend(), + mcpServerLifecycle: mcpOwner + ) + let store = CodexReviewStore.makeTestingStore(backend: backend) + await store.start() + let firstHandle = try #require(backend.lastPreparedRuntimeHandle) + + let preparationGate = AsyncGate() + backend.holdRuntimePreparation(with: preparationGate) + let restart = Task { @MainActor in await store.restart() } + await backend.waitForRuntimePreparation() + let staleReplacement = try #require(backend.lastPreparedRuntimeHandle) + let stop = Task { @MainActor in await store.stop() } + await backend.waitForRuntimePreparationCancellation() + + #expect(firstHandle.closePurposes == [.restartSameAccount]) + #expect(staleReplacement.activateCallCount == 0) + #expect(store.serverURL == endpoint) + + await preparationGate.open() + await stop.value + await restart.value + + #expect(store.serverState == .stopped) + #expect(store.serverURL == nil) + #expect(staleReplacement.closeAdmissionCallCount == 1) + #expect(staleReplacement.closePurposes == [.restartSameAccount]) + #expect(staleReplacement.waitUntilClosedCallCount == 1) + #expect(mcpOwner.stopCallCount == 1) + #expect(store.settings.lastErrorMessage == nil) + } + + @Test func replacementCancellationCatchReplaysThroughFreshCutoverWithoutError() async throws { + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend() + ) + let store = CodexReviewStore.makeTestingStore(backend: backend) + await store.start() + + let preparationGate = AsyncGate() + backend.holdRuntimePreparation(with: preparationGate) + backend.throwCancellationAfterHeldRuntimePreparation() + let firstRestart = Task { @MainActor in await store.restart() } + await backend.waitForRuntimePreparation() + let secondRestart = Task { @MainActor in await store.restart() } + await backend.waitForRuntimePreparationCancellation() + await preparationGate.open() + await secondRestart.value + await firstRestart.value + + #expect(store.serverState == .running) + #expect(store.settingsService.runtimeCutoverStatus == .active) + #expect(store.settings.isLoading == false) + #expect(store.settings.lastErrorMessage == nil) + await store.stop() + } + + @Test func callerCancellationStillConsumesCutoverByPublishingCurrentRuntime() async throws { + let preparationGate = AsyncGate() + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend() + ) + backend.holdRuntimePreparation(with: preparationGate) + let store = CodexReviewStore.makeTestingStore(backend: backend) + + let caller = Task { @MainActor in await store.start() } + await backend.waitForRuntimePreparation() + caller.cancel() + await preparationGate.open() + await caller.value + + #expect(store.serverState == .running) + #expect(store.settingsService.runtimeCutoverStatus == .active) + #expect(store.settings.isLoading == false) + #expect(store.settings.lastErrorMessage == nil) + await store.stop() + } + + @Test func genuineRuntimePreparationFailureSurfacesSettingsError() async { + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend() + ) + backend.failNextRuntimePreparation(message: "Injected preparation failure.") + let store = CodexReviewStore.makeTestingStore(backend: backend) + + await store.start() + + #expect(store.serverState == .failed("Injected preparation failure.")) + #expect(store.settingsService.runtimeCutoverStatus == .awaitingRecovery) + #expect(store.settings.isLoading == false) + #expect(store.settings.lastErrorMessage == "Injected preparation failure.") + } + + @Test func serviceOwnedCommitFailureIsNotConsumedTwiceByStore() async throws { + let reviewBackend = FakeCodexReviewBackend() + let backend = TestingCodexReviewStoreBackend(reviewBackend: reviewBackend) + let store = CodexReviewStore.makeTestingStore(backend: backend) + await store.start() + + let preparationGate = AsyncGate() + backend.holdRuntimePreparation(with: preparationGate) + let restart = Task { @MainActor in await store.restart() } + await backend.waitForRuntimePreparation() + await store.updateSettingsModel("rejected-during-commit") + await reviewBackend.failNextSettingsUpdate(message: "Commit replay failed.") + await preparationGate.open() + await restart.value + + #expect(store.serverState == .failed("Commit replay failed.")) + #expect(store.settingsService.runtimeCutoverStatus == .awaitingRecovery) + #expect(store.settings.lastErrorMessage == "Commit replay failed.") + await store.stop() + } + + @Test func settingsCutoverDrainsOldWriteAndReplaysDeferredIntentOnce() async throws { + let reviewBackend = FakeCodexReviewBackend(settings: .init(model: "initial-model")) + let backend = TestingCodexReviewStoreBackend(reviewBackend: reviewBackend) + let store = CodexReviewStore.makeTestingStore(backend: backend) + await store.start() + let firstHandle = try #require(backend.lastPreparedRuntimeHandle) + + let writeGate = AsyncGate() + await reviewBackend.holdNextSettingsUpdate(with: writeGate) + let oldWrite = Task { @MainActor in + await store.updateSettingsModel("old-runtime-edit") + } + await reviewBackend.waitForSettingsUpdate() + + let restart = Task { @MainActor in await store.restart() } + try await waitForCutoverStatus(.draining, service: store.settingsService) + await store.updateSettingsModel("deferred-edit") + #expect(backend.lastPreparedRuntimeHandle === firstHandle) + #expect(await reviewBackend.recordedCommands().filter { + if case .applySettings = $0 { true } else { false } + }.count == 1) + + await writeGate.open() + await oldWrite.value + await restart.value + + #expect(store.settings.selectedModel == "deferred-edit") + #expect(await reviewBackend.settingsSnapshot().model == "deferred-edit") + #expect(await reviewBackend.recordedCommands().filter { + if case .applySettings = $0 { true } else { false } + }.count == 2) + #expect(store.settings.lastErrorMessage == nil) + await store.stop() + } + + @Test func stopDuringServiceOwnedCommitDrainPreservesDeferredIntentWithoutError() async throws { + let reviewBackend = FakeCodexReviewBackend(settings: .init(model: "initial-model")) + let backend = TestingCodexReviewStoreBackend(reviewBackend: reviewBackend) + let store = CodexReviewStore.makeTestingStore(backend: backend) + await store.start() + + let preparationGate = AsyncGate() + backend.holdRuntimePreparation(with: preparationGate) + let restart = Task { @MainActor in await store.restart() } + await backend.waitForRuntimePreparation() + await store.updateSettingsModel("deferred-during-commit") + + let commitGate = AsyncGate() + await reviewBackend.holdNextSettingsUpdateCheckingCancellationAfterGate( + with: commitGate + ) + await preparationGate.open() + await reviewBackend.waitForSettingsUpdate() + let stop = Task { @MainActor in await store.stop() } + + await commitGate.open() + await stop.value + await restart.value + + #expect(store.serverState == .stopped) + #expect(store.settings.selectedModel == "deferred-during-commit") + #expect(await reviewBackend.settingsSnapshot().model == "deferred-during-commit") + #expect(await reviewBackend.recordedCommands().filter { + if case .applySettings = $0 { true } else { false } + }.count == 1) + #expect(store.settings.lastErrorMessage == nil) + } + + @Test func staleRuntimeFailureCannotTearDownFreshGeneration() async throws { + let backend = TestingCodexReviewStoreBackend(reviewBackend: FakeCodexReviewBackend()) + let store = CodexReviewStore.makeTestingStore(backend: backend) + await store.start() + let staleHandle = try #require(backend.lastPreparedRuntimeHandle) + await store.restart() + let currentHandle = try #require(backend.lastPreparedRuntimeHandle) + + store.requestRuntimeFailure(handle: staleHandle, cause: "Stale stream failure.") + + #expect(currentHandle !== staleHandle) + #expect(store.serverState == .running) + #expect(currentHandle.closePurposes.isEmpty) + await store.stop() + } + + @Test func explicitStopSupersedesFailurePresentationButRetainsCleanupCause() async throws { + let backend = TestingCodexReviewStoreBackend(reviewBackend: FakeCodexReviewBackend()) + let store = CodexReviewStore.makeTestingStore(backend: backend) + await store.start() + let handle = try #require(backend.lastPreparedRuntimeHandle) + let closeGate = AsyncGate() + handle.holdClose(with: closeGate) + + store.requestRuntimeFailure(handle: handle, cause: "Injected runtime failure.") + await handle.waitForClose() + let expected = "Review runtime stopped unexpectedly: Injected runtime failure." + #expect(store.serverState == .failed(expected)) + + let stop = Task { @MainActor in await store.stop() } + try await waitForTeardownFinalState(.stopped, store: store) + await closeGate.open() + await stop.value + + #expect(store.serverState == .stopped) + #expect(handle.closePurposes == [.runtimeFailure]) + #expect(handle.waitUntilClosedCallCount == 1) + } +} + +@MainActor +private func waitForCutoverStatus( + _ expected: CodexReviewSettingsService.RuntimeCutoverStatus, + service: CodexReviewSettingsService +) async throws { + let clock = ContinuousClock() + let deadline = clock.now + .seconds(2) + while service.runtimeCutoverStatus != expected { + guard clock.now < deadline else { + throw CancellationError() + } + await Task.yield() + } +} + +@MainActor +private func waitForTeardownFinalState( + _ expected: ReviewRuntimeTeardownIntent.FinalState, + store: CodexReviewStore +) async throws { + let clock = ContinuousClock() + let deadline = clock.now + .seconds(2) + while store.runtimeTeardownFinalState != expected { + guard clock.now < deadline else { + throw CancellationError() + } + await Task.yield() + } +} diff --git a/Tests/ReviewUITests/ReviewUITests.swift b/Tests/ReviewUITests/ReviewUITests.swift index dba17b32..62fa78af 100644 --- a/Tests/ReviewUITests/ReviewUITests.swift +++ b/Tests/ReviewUITests/ReviewUITests.swift @@ -6695,12 +6695,15 @@ func makeStore(backend: AuthActionBackend) -> CodexReviewStore { final class CountingStartBackend: PreviewCodexReviewStoreBackend { private var startCalls = 0 - override func start( - store _: CodexReviewStore, - forceRestartIfNeeded _: Bool - ) async { - isActive = true + override func prepareRuntime( + generation: ReviewRuntimeGeneration, + purpose: ReviewRuntimeTransitionPurpose + ) async throws -> PreparedRuntime { startCalls += 1 + return try await super.prepareRuntime( + generation: generation, + purpose: purpose + ) } override func stop(store _: CodexReviewStore) async { @@ -6732,13 +6735,6 @@ final class AuthActionBackend: PreviewCodexReviewStoreBackend { ) } - override func start( - store _: CodexReviewStore, - forceRestartIfNeeded _: Bool - ) async { - isActive = true - } - override func stop(store _: CodexReviewStore) async { isActive = false } @@ -6775,12 +6771,6 @@ final class FailingCancellationBackend: PreviewCodexReviewStoreBackend { ) } - override func start( - store _: CodexReviewStore, - forceRestartIfNeeded _: Bool - ) async { - } - override func stop(store _: CodexReviewStore) async { } @@ -6824,12 +6814,6 @@ final class BlockingSettingsBackend: PreviewCodexReviewStoreBackend { ) } - override func start( - store _: CodexReviewStore, - forceRestartIfNeeded _: Bool - ) async { - } - override func stop(store _: CodexReviewStore) async { }