From 59b3bea725603cc9199c20104daf7bb9494c9be9 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 16:46:49 +0900 Subject: [PATCH 01/38] fix(review): await canonical interrupt terminals --- Sources/CodexReview/CodexReviewBackend.swift | 20 +- .../CodexReview/ReviewAttemptProcessor.swift | 690 ++++++++++++++++++ .../CodexReview/Store/CodexReviewStore.swift | 15 +- .../Store/CodexReviewStoreBackend.swift | 10 +- .../Store/CodexReviewStoreCancellation.swift | 3 +- .../Store/CodexReviewStoreReviews.swift | 346 ++++----- .../Store/CodexReviewStoreTesting.swift | 4 +- .../PreviewCodexReviewStoreBackend.swift | 9 +- .../AppServerClient.swift | 4 +- .../AppServerCodexReviewBackend.swift | 225 ++++-- .../AppServerProcessTransport.swift | 65 +- .../AppServerReviewControl.swift | 44 +- Sources/CodexReviewAppServer/JSONRPC.swift | 2 +- Sources/CodexReviewHost/CodexReviewHost.swift | 25 +- .../LiveCodexReviewStoreBackend.swift | 49 +- Sources/CodexReviewTesting/TestSupport.swift | 135 +++- .../AppServerClientTests.swift | 357 ++++----- .../CodexReviewHostTests.swift | 4 +- .../CodexReviewStoreCommandTests.swift | 134 +++- .../ReviewAttemptProcessorTests.swift | 556 ++++++++++++++ 20 files changed, 2153 insertions(+), 544 deletions(-) create mode 100644 Sources/CodexReview/ReviewAttemptProcessor.swift create mode 100644 Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift diff --git a/Sources/CodexReview/CodexReviewBackend.swift b/Sources/CodexReview/CodexReviewBackend.swift index e29cc64e..82710137 100644 --- a/Sources/CodexReview/CodexReviewBackend.swift +++ b/Sources/CodexReview/CodexReviewBackend.swift @@ -10,8 +10,12 @@ package protocol CodexReviewBackend: Sendable { func completeLogin(_ response: CodexReviewBackendModel.Login.Response) async throws -> CodexReviewBackendModel.Auth.Snapshot func logout(_ account: CodexReviewBackendModel.Account.ID) async throws -> CodexReviewBackendModel.Auth.Snapshot - func startReview(_ request: CodexReviewBackendModel.Review.Start) async throws -> BackendReviewAttempt + func startReview( + _ request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission + ) async throws -> BackendReviewAttempt func interruptReview(_ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason) async throws + func forceCloseReviewConnection() async throws func beginReviewRecovery( _ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason @@ -20,7 +24,19 @@ package protocol CodexReviewBackend: Sendable { _ token: CodexReviewBackendModel.Review.RecoveryToken, request: CodexReviewBackendModel.Review.Start ) async throws -> BackendReviewAttempt - func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async + func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws +} + +package extension CodexReviewBackend { + func startReview( + _ request: CodexReviewBackendModel.Review.Start + ) async throws -> BackendReviewAttempt { + let admission = ReviewStartAdmission() + let task = await admission.start { admission in + try await self.startReview(request, admission: admission) + } + return try await task.value + } } package struct BackendReviewAttempt: Sendable { diff --git a/Sources/CodexReview/ReviewAttemptProcessor.swift b/Sources/CodexReview/ReviewAttemptProcessor.swift new file mode 100644 index 00000000..ab4126fc --- /dev/null +++ b/Sources/CodexReview/ReviewAttemptProcessor.swift @@ -0,0 +1,690 @@ +import Foundation + +package struct ReviewRuntimeClosePolicy: Sendable { + package var terminalGrace: Duration + package var sleep: @Sendable (Duration) async throws -> Void + + package init( + terminalGrace: Duration, + sleep: @escaping @Sendable (Duration) async throws -> Void + ) { + self.terminalGrace = terminalGrace + self.sleep = sleep + } + + package static let production = Self( + terminalGrace: .seconds(10), + sleep: { try await Task.sleep(for: $0) } + ) +} + +package enum ReviewInterruptRequestOutcome: Equatable, Sendable { + case rejected(code: Int?, message: String) + case outcomeUnknown(message: String) +} + +package struct ReviewInterruptRequestFailure: LocalizedError, Equatable, Sendable { + package var outcome: ReviewInterruptRequestOutcome + package var secondaryBarrierDiagnostic: String? + + package init( + outcome: ReviewInterruptRequestOutcome, + secondaryBarrierDiagnostic: String? = nil + ) { + self.outcome = outcome + self.secondaryBarrierDiagnostic = secondaryBarrierDiagnostic + } + + package var errorDescription: String? { + switch outcome { + case .rejected(_, let message), .outcomeUnknown(let message): + message + } + } +} + +package enum ReviewRuntimeCloseFailure: LocalizedError, Equatable, Sendable { + case connection(String) + case process(String) + case worker(String) + case cleanup(String) + case mcpHandlerDrain(String) + + package var errorDescription: String? { + switch self { + case .connection(let message): + "App-server connection close failed: \(message)" + case .process(let message): + "App-server process close failed: \(message)" + case .worker(let message): + "Review worker close failed: \(message)" + case .cleanup(let message): + "Review cleanup failed: \(message)" + case .mcpHandlerDrain(let message): + "MCP handler drain failed: \(message)" + } + } +} + +package struct ReviewAttemptContractFailure: LocalizedError, Equatable, Sendable { + package var message: String + + package init(message: String) { + self.message = message + } + + package var errorDescription: String? { message } +} + +package struct ReviewStartCancelledBeforeDispatch: LocalizedError, Equatable, Sendable { + package var cancellation: ReviewCancellation + + package init(cancellation: ReviewCancellation) { + self.cancellation = cancellation + } + + package var errorDescription: String? { cancellation.message } +} + +package enum ReviewAttemptBarrierTerminal: Equatable, Sendable { + case canonical( + run: CodexReviewBackendModel.Review.Run, + terminal: ReviewTerminalRecord + ) + case connection(ReviewRuntimeCloseFailure) + case localCancellation(ReviewCancellation) + + package var diagnosticDescription: String { + switch self { + case .canonical(let run, let terminal): + "canonical terminal \(terminal.kind.rawValue) for attempt \(run.attemptID)" + case .connection(let failure): + failure.localizedDescription + case .localCancellation(let cancellation): + cancellation.message + } + } +} + +package struct ReviewAttemptCancellationResolution: Equatable, Sendable { + package var terminal: ReviewAttemptBarrierTerminal + package var requestFailure: ReviewInterruptRequestFailure? + + package init( + terminal: ReviewAttemptBarrierTerminal, + requestFailure: ReviewInterruptRequestFailure? = nil + ) { + self.terminal = terminal + self.requestFailure = requestFailure + } +} + +/// Owns one review attempt from the first dispatch admission through terminal and cleanup. +/// Store cancellation and backend request dispatch both consult this actor; there is no +/// call-site startup-cancellation mirror. +package actor ReviewStartAdmission { + package enum RequestDispatch: Equatable, Sendable { + case notSent + case outcomeUnknown + } + + package enum Phase: Equatable, Sendable { + case queued + case preparingThread(RequestDispatch) + case startingReview( + preparedRun: CodexReviewBackendModel.Review.Run, + dispatch: RequestDispatch + ) + case active(CodexReviewBackendModel.Review.Run) + case interrupting(CodexReviewBackendModel.Review.Run) + case finishing(ReviewAttemptBarrierTerminal) + case terminal(ReviewAttemptBarrierTerminal) + } + + private enum CancellationTaskEvent { + case request(Result) + case barrier(ReviewAttemptBarrierTerminal) + case graceExpired + case forceClose(Result) + } + + private let closePolicy: ReviewRuntimeClosePolicy + private var phase: Phase = .queued + private var requestedCancellation: ReviewCancellation? + private var startTask: Task? + private var cancellationTask: Task? + private var interruptRequestTask: Task? + private var terminalBarrierTask: Task? + private var graceTask: Task? + private var forceCloseTask: Task? + private var cleanupTasksByAttemptID: [String: Task] = [:] + private var terminal: ReviewAttemptBarrierTerminal? + private var requestResult: Result? + private var forceCloseResult: Result? + private var cancellationResult: Result? + private var terminalWaiters: [UUID: CheckedContinuation] = [:] + private var cancellationWaiters: [CheckedContinuation, Never>] = [] + + package init(closePolicy: ReviewRuntimeClosePolicy = .production) { + self.closePolicy = closePolicy + } + + package func start( + _ operation: @escaping @Sendable (ReviewStartAdmission) async throws -> BackendReviewAttempt + ) -> Task { + precondition( + startTask == nil, + "ReviewStartAdmission owns exactly one registered start Task." + ) + phase = .preparingThread(.notSent) + let task = Task { + do { + let attempt = try await operation(self) + self.finishStart(with: .success(attempt)) + return attempt + } catch { + self.finishStart(with: .failure(error)) + throw error + } + } + startTask = task + return task + } + + package func admitThreadStartDispatch() -> Bool { + guard requestedCancellation == nil else { + return false + } + switch phase { + case .preparingThread: + phase = .preparingThread(.outcomeUnknown) + return true + case .queued, .startingReview, .active, .interrupting, .finishing, .terminal: + return false + } + } + + package func recordPreparedThread(_ run: CodexReviewBackendModel.Review.Run) { + guard terminal == nil else { + return + } + phase = .startingReview(preparedRun: run, dispatch: .notSent) + } + + package func admitReviewStartDispatch( + for preparedRun: CodexReviewBackendModel.Review.Run + ) -> Bool { + guard requestedCancellation == nil else { + return false + } + guard case .startingReview(let currentRun, _) = phase, + currentRun.attemptID == preparedRun.attemptID + else { + return false + } + phase = .startingReview(preparedRun: preparedRun, dispatch: .outcomeUnknown) + return true + } + + package func recordActiveRun(_ run: CodexReviewBackendModel.Review.Run) { + guard terminal == nil else { + return + } + phase = .active(run) + } + + package func recordCanonicalTerminal( + _ terminalRecord: ReviewTerminalRecord, + for run: CodexReviewBackendModel.Review.Run + ) throws { + if let terminal { + let candidate = ReviewAttemptBarrierTerminal.canonical(run: run, terminal: terminalRecord) + guard terminal == candidate else { + throw ReviewAttemptContractFailure( + message: "Conflicting terminal for review attempt \(run.attemptID)." + ) + } + return + } + guard let canonicalRun = canonicalRunForTerminal, + Self.matchesCanonicalPair(run, canonicalRun) + else { + return + } + receiveTerminal(.canonical(run: run, terminal: terminalRecord)) + } + + package func recordConnectionTerminal(_ failure: ReviewRuntimeCloseFailure) { + guard terminal == nil else { + return + } + receiveTerminal(.connection(failure)) + } + + package func cancel( + _ cancellation: ReviewCancellation, + interrupt: @escaping @Sendable ( + CodexReviewBackendModel.Review.Run, + CodexReviewBackendModel.CancellationReason + ) async throws -> Void, + forceClose: @escaping @Sendable () async throws -> Void + ) async throws -> ReviewAttemptCancellationResolution { + if let cancellationTask { + return try await cancellationTask.value + } + if let terminal { + return .init(terminal: terminal) + } + requestedCancellation = cancellation + let task = Task { + try await self.performCancellation( + cancellation, + interrupt: interrupt, + forceClose: forceClose + ) + } + cancellationTask = task + return try await task.value + } + + package func cleanup( + run: CodexReviewBackendModel.Review.Run, + _ operation: @escaping @Sendable () async throws -> Void + ) async throws { + if let cleanupTask = cleanupTasksByAttemptID[run.attemptID] { + return try await cleanupTask.value + } + let task = Task { + try await operation() + } + cleanupTasksByAttemptID[run.attemptID] = task + return try await task.value + } + + package func currentPhase() -> Phase { phase } + + package func cancellationRequest() -> ReviewCancellation? { + requestedCancellation + } + + package func recordedCleanupResult( + for run: CodexReviewBackendModel.Review.Run + ) async -> Result? { + guard let cleanupTask = cleanupTasksByAttemptID[run.attemptID] else { + return nil + } + return await cleanupTask.result + } + + private func finishStart( + with result: Result + ) { + switch result { + case .success(let attempt): + if terminal == nil { + phase = .active(attempt.run) + } + case .failure(let error): + if let cancellation = (error as? ReviewStartCancelledBeforeDispatch)?.cancellation { + receiveTerminal(.localCancellation(cancellation)) + } else if error is CancellationError, + let requestedCancellation, + case .preparingThread(.notSent) = phase { + receiveTerminal(.localCancellation(requestedCancellation)) + } + } + } + + private func performCancellation( + _ cancellation: ReviewCancellation, + interrupt: @escaping @Sendable ( + CodexReviewBackendModel.Review.Run, + CodexReviewBackendModel.CancellationReason + ) async throws -> Void, + forceClose: @escaping @Sendable () async throws -> Void + ) async throws -> ReviewAttemptCancellationResolution { + if case .preparingThread(.notSent) = phase { + startTask?.cancel() + } else if case .queued = phase { + receiveTerminal(.localCancellation(cancellation)) + } else if Self.isOutcomeUnknownStartPhase(phase) { + installGraceTask(forceClose: forceClose) + } + + if let startTask, activeRun == nil, terminal == nil { + _ = await startTask.result + } + + if let terminal { + return try await finishCancellationAfterTerminal(terminal) + } + guard let run = activeRun else { + if let startTask { + switch await startTask.result { + case .success(let attempt): + return try await beginActiveCancellation( + run: attempt.run, + cancellation: cancellation, + interrupt: interrupt, + forceClose: forceClose + ) + case .failure(let error): + throw error + } + } + let local = ReviewAttemptBarrierTerminal.localCancellation(cancellation) + receiveTerminal(local) + return .init(terminal: local) + } + return try await beginActiveCancellation( + run: run, + cancellation: cancellation, + interrupt: interrupt, + forceClose: forceClose + ) + } + + private func beginActiveCancellation( + run: CodexReviewBackendModel.Review.Run, + cancellation: ReviewCancellation, + interrupt: @escaping @Sendable ( + CodexReviewBackendModel.Review.Run, + CodexReviewBackendModel.CancellationReason + ) async throws -> Void, + forceClose: @escaping @Sendable () async throws -> Void + ) async throws -> ReviewAttemptCancellationResolution { + phase = .interrupting(run) + installTerminalBarrierTask() + installInterruptRequestTask( + run: run, + cancellation: cancellation, + interrupt: interrupt + ) + installGraceTask(forceClose: forceClose) + let result = await withCheckedContinuation { continuation in + if let cancellationResult { + continuation.resume(returning: cancellationResult) + } else { + cancellationWaiters.append(continuation) + resolveCancellationIfPossible() + } + } + await drainCancellationTasks() + switch result { + case .success(let resolution): + phase = .terminal(resolution.terminal) + return resolution + case .failure(let error): + if let terminal { + phase = .terminal(terminal) + } else if case .active = phase { + requestedCancellation = nil + } + throw error + } + } + + private func finishCancellationAfterTerminal( + _ terminal: ReviewAttemptBarrierTerminal + ) async throws -> ReviewAttemptCancellationResolution { + if interruptRequestTask != nil || terminalBarrierTask != nil { + let result = await withCheckedContinuation { continuation in + if let cancellationResult { + continuation.resume(returning: cancellationResult) + } else { + cancellationWaiters.append(continuation) + resolveCancellationIfPossible() + } + } + await drainCancellationTasks() + return try result.get() + } + await drainCancellationTasks() + phase = .terminal(terminal) + return .init(terminal: terminal) + } + + private func installInterruptRequestTask( + run: CodexReviewBackendModel.Review.Run, + cancellation: ReviewCancellation, + interrupt: @escaping @Sendable ( + CodexReviewBackendModel.Review.Run, + CodexReviewBackendModel.CancellationReason + ) async throws -> Void + ) { + guard interruptRequestTask == nil else { + return + } + interruptRequestTask = Task { + let result: Result + do { + try await interrupt(run, .init(message: cancellation.message)) + result = .success(()) + } catch let failure as ReviewInterruptRequestFailure { + result = .failure(failure) + } catch { + result = .failure(.init(outcome: .outcomeUnknown( + message: error.localizedDescription + ))) + } + self.receive(.request(result)) + } + } + + private func installTerminalBarrierTask() { + guard terminalBarrierTask == nil else { + return + } + terminalBarrierTask = Task { + if let terminal = await self.waitForTerminal() { + self.receive(.barrier(terminal)) + } + } + } + + private func installGraceTask( + forceClose: @escaping @Sendable () async throws -> Void + ) { + guard graceTask == nil else { + return + } + let policy = closePolicy + graceTask = Task { + do { + try await policy.sleep(policy.terminalGrace) + } catch { + return + } + self.receive(.graceExpired) + self.installForceCloseTask(forceClose) + } + } + + private func installForceCloseTask( + _ forceClose: @escaping @Sendable () async throws -> Void + ) { + guard forceCloseTask == nil else { + return + } + forceCloseTask = Task { + let result: Result + do { + try await forceClose() + result = .success(()) + } catch let failure as ReviewRuntimeCloseFailure { + result = .failure(failure) + } catch { + result = .failure(.connection(error.localizedDescription)) + } + self.receive(.forceClose(result)) + } + } + + private func receive(_ event: CancellationTaskEvent) { + switch event { + case .request(let result): + requestResult = result + case .barrier(let terminal): + if self.terminal == nil { + self.terminal = terminal + } + case .graceExpired: + break + case .forceClose(let result): + forceCloseResult = result + } + resolveCancellationIfPossible() + } + + private func receiveTerminal(_ terminal: ReviewAttemptBarrierTerminal) { + self.terminal = terminal + phase = cancellationTask == nil ? .terminal(terminal) : .finishing(terminal) + let waiters = Array(terminalWaiters.values) + terminalWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume(returning: terminal) + } + resolveCancellationIfPossible() + } + + private func resolveCancellationIfPossible() { + guard cancellationResult == nil else { + return + } + if case .failure(let closeFailure)? = forceCloseResult, terminal == nil { + resolveCancellation(.failure(closeFailure)) + return + } + guard let requestResult else { + return + } + switch requestResult { + case .success: + guard let terminal else { + return + } + resolveCancellation(.success(.init(terminal: terminal))) + case .failure(let requestFailure): + switch requestFailure.outcome { + case .rejected: + if let terminal { + resolveCancellation(.success(.init( + terminal: terminal, + requestFailure: requestFailure + ))) + } else { + requestedCancellation = nil + if case .interrupting(let run) = phase { + phase = .active(run) + } + resolveCancellation(.failure(requestFailure)) + } + case .outcomeUnknown: + guard let terminal else { + return + } + switch terminal { + case .connection(let connectionFailure): + resolveCancellation(.failure(ReviewInterruptRequestFailure( + outcome: requestFailure.outcome, + secondaryBarrierDiagnostic: connectionFailure.localizedDescription + ))) + case .canonical, .localCancellation: + resolveCancellation(.success(.init( + terminal: terminal, + requestFailure: requestFailure + ))) + } + } + } + } + + private func resolveCancellation( + _ result: Result + ) { + cancellationResult = result + let waiters = cancellationWaiters + cancellationWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume(returning: result) + } + } + + private func waitForTerminal() async -> ReviewAttemptBarrierTerminal? { + if let terminal { + return terminal + } + let waiterID = UUID() + return await withTaskCancellationHandler { + await withCheckedContinuation { continuation in + if let terminal { + continuation.resume(returning: terminal) + } else { + terminalWaiters[waiterID] = continuation + } + } + } onCancel: { + Task { + await self.cancelTerminalWaiter(waiterID) + } + } + } + + private func cancelTerminalWaiter(_ id: UUID) { + terminalWaiters.removeValue(forKey: id)?.resume(returning: nil) + } + + private func drainCancellationTasks() async { + if terminal == nil { + terminalBarrierTask?.cancel() + } + graceTask?.cancel() + await interruptRequestTask?.value + await terminalBarrierTask?.value + await graceTask?.value + await forceCloseTask?.value + } + + private var canonicalRunForTerminal: CodexReviewBackendModel.Review.Run? { + switch phase { + case .startingReview(let run, _), .active(let run), .interrupting(let run): + run + case .finishing(let terminal), .terminal(let terminal): + if case .canonical(let run, _) = terminal { + run + } else { + nil + } + case .queued, .preparingThread: + nil + } + } + + private var activeRun: CodexReviewBackendModel.Review.Run? { + switch phase { + case .active(let run), .interrupting(let run): + run + case .queued, .preparingThread, .startingReview, .finishing, .terminal: + nil + } + } + + private static func isOutcomeUnknownStartPhase(_ phase: Phase) -> Bool { + switch phase { + case .preparingThread(.outcomeUnknown), .startingReview(_, .outcomeUnknown): + true + case .queued, .preparingThread(.notSent), .startingReview(_, .notSent), + .active, .interrupting, .finishing, .terminal: + false + } + } + + private static func matchesCanonicalPair( + _ lhs: CodexReviewBackendModel.Review.Run, + _ rhs: CodexReviewBackendModel.Review.Run + ) -> Bool { + lhs.attemptID == rhs.attemptID + && lhs.reviewThreadID == rhs.reviewThreadID + && lhs.turnID == rhs.turnID + } +} diff --git a/Sources/CodexReview/Store/CodexReviewStore.swift b/Sources/CodexReview/Store/CodexReviewStore.swift index d7ab5001..e5f898b2 100644 --- a/Sources/CodexReview/Store/CodexReviewStore.swift +++ b/Sources/CodexReview/Store/CodexReviewStore.swift @@ -25,13 +25,14 @@ public final class CodexReviewStore { @ObservationIgnored package let backend: any CodexReviewStoreBackend @ObservationIgnored package let networkMonitor: any CodexReviewNetworkMonitoring @ObservationIgnored package let networkRecoveryPolicy: CodexReviewNetworkRecoveryPolicy + @ObservationIgnored package let reviewRuntimeClosePolicy: ReviewRuntimeClosePolicy @ObservationIgnored package var previewSupportRetainer: AnyObject? @ObservationIgnored package let clock: CodexReviewClock @ObservationIgnored package let idGenerator: CodexReviewIDGenerator @ObservationIgnored package var activeRuns: [String: CodexReviewBackendModel.Review.Run] = [:] @ObservationIgnored package var reviewRecoveryWaitingJobIDs: Set = [] - @ObservationIgnored package var startingJobIDs: Set = [] - @ObservationIgnored package var startupCancellations: [String: ReviewCancellation] = [:] + @ObservationIgnored package var reviewStartAdmissions: [String: ReviewStartAdmission] = [:] + @ObservationIgnored package var reviewCleanupFailures: [String: ReviewRuntimeCloseFailure] = [:] @ObservationIgnored package var reviewWorkerTasks: [String: Task] = [:] @ObservationIgnored package var runtimeStopDetachedReviewWorkerTasks: [String: Task] = [:] @ObservationIgnored package var reviewTerminalWaiters: [String: [ReviewTerminalWaiter]] = [:] @@ -45,11 +46,13 @@ public final class CodexReviewStore { clock: CodexReviewClock = .init(), idGenerator: CodexReviewIDGenerator = .init(), networkMonitor: any CodexReviewNetworkMonitoring = SystemCodexReviewNetworkMonitor(), - networkRecoveryPolicy: CodexReviewNetworkRecoveryPolicy = .default + networkRecoveryPolicy: CodexReviewNetworkRecoveryPolicy = .default, + reviewRuntimeClosePolicy: ReviewRuntimeClosePolicy = .production ) { self.backend = backend self.networkMonitor = networkMonitor self.networkRecoveryPolicy = networkRecoveryPolicy + self.reviewRuntimeClosePolicy = reviewRuntimeClosePolicy self.diagnosticsURL = diagnosticsURL self.clock = clock self.idGenerator = idGenerator @@ -113,7 +116,8 @@ public final class CodexReviewStore { clock: CodexReviewClock = .init(), idGenerator: CodexReviewIDGenerator = .init(), networkMonitor: any CodexReviewNetworkMonitoring = StaticCodexReviewNetworkMonitor(), - networkRecoveryPolicy: CodexReviewNetworkRecoveryPolicy = .default + networkRecoveryPolicy: CodexReviewNetworkRecoveryPolicy = .default, + reviewRuntimeClosePolicy: ReviewRuntimeClosePolicy = .production ) -> CodexReviewStore { CodexReviewStore( backend: backend, @@ -121,7 +125,8 @@ public final class CodexReviewStore { clock: clock, idGenerator: idGenerator, networkMonitor: networkMonitor, - networkRecoveryPolicy: networkRecoveryPolicy + networkRecoveryPolicy: networkRecoveryPolicy, + reviewRuntimeClosePolicy: reviewRuntimeClosePolicy ) } diff --git a/Sources/CodexReview/Store/CodexReviewStoreBackend.swift b/Sources/CodexReview/Store/CodexReviewStoreBackend.swift index ad4d13d4..12631c80 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreBackend.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreBackend.swift @@ -24,7 +24,7 @@ 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 } @@ -44,8 +44,12 @@ package protocol CodexReviewStoreBackend: CodexReviewSettingsBackend { func refreshAccountRateLimits(auth: CodexReviewAuthModel, accountKey: String) async func requiresCurrentSessionRecovery(auth: CodexReviewAuthModel, accountKey: String) -> Bool - func startReview(_ request: CodexReviewBackendModel.Review.Start) async throws -> BackendReviewAttempt + func startReview( + _ request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission + ) async throws -> BackendReviewAttempt func interruptReview(_ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason) async throws + func forceCloseReviewConnection() async throws func beginReviewRecovery( _ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason @@ -54,7 +58,7 @@ package protocol CodexReviewStoreBackend: CodexReviewSettingsBackend { _ token: CodexReviewBackendModel.Review.RecoveryToken, request: CodexReviewBackendModel.Review.Start ) async throws -> BackendReviewAttempt - func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async + func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws } extension CodexReviewStoreBackend { diff --git a/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift b/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift index fcac5bea..eccaf094 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift @@ -192,8 +192,7 @@ extension CodexReviewStore { } activeRuns.removeValue(forKey: jobID) reviewRecoveryWaitingJobIDs.remove(jobID) - startingJobIDs.remove(jobID) - startupCancellations.removeValue(forKey: jobID) + reviewStartAdmissions.removeValue(forKey: jobID) } } diff --git a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift index 9efc8a58..c79d4480 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift @@ -16,15 +16,9 @@ extension CodexReviewStore { request: CodexReviewAPI.Start.Request ) async throws -> CodexReviewAPI.Read.Result { let jobID = try beginReview(sessionID: sessionID, request: request) - return try await withTaskCancellationHandler { - _ = try await awaitReview(sessionID: sessionID, jobID: jobID) - await reviewWorkerTasks[jobID]?.value - return try readReview(sessionID: sessionID, jobID: jobID) - } onCancel: { - Task { @MainActor [weak self] in - self?.reviewWorkerTasks[jobID]?.cancel() - } - } + _ = try await awaitReview(sessionID: sessionID, jobID: jobID) + await reviewWorkerTasks[jobID]?.value + return try readReview(sessionID: sessionID, jobID: jobID) } @discardableResult @@ -78,29 +72,42 @@ extension CodexReviewStore { ) insertReviewJob(job) markReviewRunning(job, startedAt: createdAt) - startingJobIDs.insert(jobID) - launchReviewWorker(jobID: jobID, sessionID: sessionID, request: validatedRequest) + let admission = ReviewStartAdmission(closePolicy: reviewRuntimeClosePolicy) + reviewStartAdmissions[jobID] = admission + launchReviewWorker( + jobID: jobID, + sessionID: sessionID, + request: validatedRequest, + admission: admission + ) return jobID } private func launchReviewWorker( jobID: String, sessionID: String, - request: CodexReviewAPI.Start.Request + request: CodexReviewAPI.Start.Request, + admission: ReviewStartAdmission ) { reviewWorkerTasks[jobID]?.cancel() reviewWorkerTasks[jobID] = Task { [weak self] in - await self?.runReviewWorker(jobID: jobID, sessionID: sessionID, request: request) + await self?.runReviewWorker( + jobID: jobID, + sessionID: sessionID, + request: request, + admission: admission + ) } } private func runReviewWorker( jobID: String, sessionID: String, - request validatedRequest: CodexReviewAPI.Start.Request + request validatedRequest: CodexReviewAPI.Start.Request, + admission: ReviewStartAdmission ) async { guard let job = job(id: jobID) else { - startingJobIDs.remove(jobID) + reviewStartAdmissions.removeValue(forKey: jobID) reviewWorkerTasks.removeValue(forKey: jobID) resumeReviewWaiters(for: jobID) return @@ -113,84 +120,80 @@ extension CodexReviewStore { ) var run: CodexReviewBackendModel.Review.Run? do { - let backendAttempt = try await backend.startReview(startRequest) + let backend = self.backend + let startTask = await admission.start { admission in + try await backend.startReview(startRequest, admission: admission) + } + let backendAttempt = try await startTask.value let backendRun = backendAttempt.run - startingJobIDs.remove(jobID) run = backendRun - if Task.isCancelled { - throw CancellationError() - } applyBackendRun(backendRun, to: job) - if let startupCancellation = startupCancellations.removeValue(forKey: jobID) { - try? await backend.interruptReview( - backendRun, - reason: .init(message: startupCancellation.message) - ) - if job.isTerminal == false { - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: startupCancellation - ) - } - } else if job.cancellationRequested { - try await backend.interruptReview( - backendRun, - reason: .init(message: job.core.lifecycle.cancellation?.message ?? "Cancellation requested.") - ) - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: job.core.lifecycle.cancellation ?? .system() - ) - } if job.isTerminal { - await backend.cleanupReview(backendRun) + do { + try await cleanupReview(backendRun, admission: admission) + } catch { + retainCleanupFailure(error, for: jobID) + } activeRuns.removeValue(forKey: jobID) reviewRecoveryWaitingJobIDs.remove(jobID) } else { let currentRun = try await consumeReviewEvents( for: backendAttempt, job: job, - startRequest: startRequest + startRequest: startRequest, + admission: admission ) run = currentRun - await backend.cleanupReview(currentRun) + do { + try await cleanupReview(currentRun, admission: admission) + } catch { + retainCleanupFailure(error, for: jobID) + } activeRuns.removeValue(forKey: jobID) reviewRecoveryWaitingJobIDs.remove(jobID) } - } catch let error where error is CancellationError || Task.isCancelled { - startingJobIDs.remove(jobID) - let startupCancellation = startupCancellations.removeValue(forKey: jobID) - if let cleanupRun = activeRuns[jobID] ?? run { - await interruptReviewAfterTaskCancellation(cleanupRun, job: job) - await backend.cleanupReview(cleanupRun) - } else if job.isTerminal == false || startupCancellation != nil { + } catch let cancellation as ReviewStartCancelledBeforeDispatch { + if job.isTerminal == false { try? completeCancellationLocally( jobID: job.id, sessionID: job.sessionID, - cancellation: startupCancellation ?? job.core.lifecycle.cancellation ?? .system() + cancellation: cancellation.cancellation ) } + } catch let error where error is CancellationError || Task.isCancelled { + if let cleanupRun = activeRuns[jobID] ?? run { + let failure = ReviewRuntimeCloseFailure.worker( + "Review worker was cancelled before a canonical terminal." + ) + await admission.recordConnectionTerminal(failure) + do { + try await cleanupReview(cleanupRun, admission: admission) + } catch { + retainCleanupFailure(error, for: jobID) + } + if job.isTerminal == false { + markReviewInterrupted(job, cause: .transport(message: failure.localizedDescription)) + } + } else if job.isTerminal == false { + markReviewFailed(job, message: error.localizedDescription) + } activeRuns.removeValue(forKey: jobID) reviewRecoveryWaitingJobIDs.remove(jobID) } catch { - startingJobIDs.remove(jobID) - let startupCancellation = startupCancellations.removeValue(forKey: jobID) if let cleanupRun = activeRuns[jobID] ?? run { - await backend.cleanupReview(cleanupRun) + do { + try await cleanupReview(cleanupRun, admission: admission) + } catch { + retainCleanupFailure(error, for: jobID) + } } activeRuns.removeValue(forKey: jobID) reviewRecoveryWaitingJobIDs.remove(jobID) - if job.isTerminal == false, let startupCancellation { - try? completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: startupCancellation - ) - } else if job.isTerminal == false, - let transportFailure = error as? ReviewWorkerInputQueueError { + if job.isTerminal == false, + let transportFailure = error as? ReviewWorkerInputQueueError { + let failure = ReviewRuntimeCloseFailure.connection(transportFailure.message) + await admission.recordConnectionTerminal(failure) markReviewInterrupted( job, cause: .transport(message: transportFailure.message) @@ -201,37 +204,34 @@ extension CodexReviewStore { } reviewWorkerTasks.removeValue(forKey: jobID) runtimeStopDetachedReviewWorkerTasks.removeValue(forKey: jobID) + if reviewCleanupFailures[jobID] == nil { + reviewStartAdmissions.removeValue(forKey: jobID) + } if job.isTerminal { resumeReviewWaiters(for: jobID) } } - private func interruptReviewAfterTaskCancellation(_ run: CodexReviewBackendModel.Review.Run, job: CodexReviewJob) async { - guard job.isTerminal == false else { + private func cleanupReview( + _ run: CodexReviewBackendModel.Review.Run, + admission: ReviewStartAdmission + ) async throws { + let backend = self.backend + try await admission.cleanup(run: run) { + try await backend.cleanupReview(run) + } + } + + private func retainCleanupFailure(_ error: any Error, for jobID: String) { + guard reviewCleanupFailures[jobID] == nil else { return } - let cancellation = job.core.lifecycle.cancellation ?? .system() - job.cancellationRequested = true - job.core.lifecycle.cancellation = cancellation - job.core.output.summary = cancellation.message - job.core.lifecycle.errorMessage = cancellation.message - do { - try await backend.interruptReview( - run, - reason: .init(message: cancellation.message) - ) - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation - ) - } catch { - try? recordCancellationFailure( - jobID: job.id, - sessionID: job.sessionID, - message: error.localizedDescription - ) + if let failure = error as? ReviewRuntimeCloseFailure { + reviewCleanupFailures[jobID] = failure + } else { + reviewCleanupFailures[jobID] = .cleanup(error.localizedDescription) } + writeDiagnosticsIfNeeded() } private func applyBackendRun(_ backendRun: CodexReviewBackendModel.Review.Run, to job: CodexReviewJob) { @@ -405,81 +405,70 @@ extension CodexReviewStore { recordCancellationRequest(cancellation, for: job) - if job.core.lifecycle.status == .queued { + if reviewRecoveryWaitingJobIDs.contains(jobID) { try completeCancellationLocally( jobID: job.id, sessionID: job.sessionID, cancellation: cancellation ) + reviewWorkerTasks[jobID]?.cancel() return .init(jobID: job.id, cancelled: true, core: job.core) } - if reviewRecoveryWaitingJobIDs.contains(jobID) { + guard let admission = reviewStartAdmissions[jobID] else { try completeCancellationLocally( jobID: job.id, sessionID: job.sessionID, cancellation: cancellation ) - reviewWorkerTasks[jobID]?.cancel() return .init(jobID: job.id, cancelled: true, core: job.core) } - if let run = activeRuns[jobID] { - do { - try await backend.interruptReview( - run, - reason: .init(message: cancellation.message) - ) + let backend = self.backend + do { + let resolution = try await admission.cancel( + cancellation, + interrupt: { run, reason in + try await backend.interruptReview(run, reason: reason) + }, + forceClose: { + try await backend.forceCloseReviewConnection() + } + ) + await reviewWorkerTasks[jobID]?.value + if job.isTerminal == false, + case .localCancellation = resolution.terminal { try completeCancellationLocally( jobID: job.id, sessionID: job.sessionID, cancellation: cancellation ) - reviewWorkerTasks[jobID]?.cancel() - } catch { - try recordCancellationFailure( - jobID: job.id, - sessionID: job.sessionID, - message: error.localizedDescription - ) - throw error } - } else if let run = job.backendRun { - do { - try await backend.interruptReview( - run, - reason: .init(message: cancellation.message) - ) - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation - ) - reviewWorkerTasks[jobID]?.cancel() - } catch { + if let run = job.backendRun, + let cleanupResult = await admission.recordedCleanupResult(for: run) { + try cleanupResult.get() + } + } catch { + let phase = await admission.currentPhase() + if case .finishing = phase { + await reviewWorkerTasks[jobID]?.value + } else if case .terminal = phase { + await reviewWorkerTasks[jobID]?.value + } + if job.isTerminal == false { try recordCancellationFailure( jobID: job.id, sessionID: job.sessionID, message: error.localizedDescription ) - throw error } - } else if startingJobIDs.contains(jobID) { - startupCancellations[jobID] = cancellation - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation - ) - return .init(jobID: job.id, cancelled: true, core: job.core) - } else { - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation - ) + throw error } - return .init(jobID: job.id, cancelled: true, core: job.core) + return .init( + jobID: job.id, + cancelled: job.core.lifecycle.status == .cancelled, + core: job.core + ) } package func closeSession( @@ -605,7 +594,8 @@ extension CodexReviewStore { private func consumeReviewEvents( for initialAttempt: BackendReviewAttempt, job: CodexReviewJob, - startRequest: CodexReviewBackendModel.Review.Start + startRequest: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission ) async throws -> CodexReviewBackendModel.Review.Run { let inputs = await reviewWorkerInputs(for: initialAttempt) defer { @@ -624,6 +614,12 @@ extension CodexReviewStore { else { continue } + if let terminal = reviewTerminalRecord(for: event.event, job: job) { + try await admission.recordCanonicalTerminal( + terminal, + for: recoveryState.currentRun + ) + } recoveryState.currentRun = handleReviewEvent( event.event, job: job, @@ -639,9 +635,10 @@ extension CodexReviewStore { if recoveryState.shouldIgnoreFinishedEvent(for: finishedRun.run) { continue } - if try handleReviewEventsFinished( + if await handleReviewEventsFinished( job: job, - isWaitingForNetworkRecovery: recoveryState.isWaitingForNetworkRecovery + isWaitingForNetworkRecovery: recoveryState.isWaitingForNetworkRecovery, + admission: admission ) { return recoveryState.currentRun } @@ -683,7 +680,8 @@ extension CodexReviewStore { job: job, startRequest: startRequest, inputs: inputs, - recoveryToken: recoveryState.recoveryToken + recoveryToken: recoveryState.recoveryToken, + admission: admission ) { case .continueWaiting: recoveryState.markWaitingForNetworkRecovery() @@ -693,6 +691,7 @@ extension CodexReviewStore { return recoveryState.currentRun case .recovered(let recoveredAttempt): let recoveredRun = recoveredAttempt.run + await admission.recordActiveRun(recoveredRun) applyBackendRun(recoveredRun, to: job) recoveryState.markRecovered(with: recoveredRun) reviewRecoveryWaitingJobIDs.remove(job.id) @@ -723,23 +722,22 @@ extension CodexReviewStore { throw CancellationError() } if job.isTerminal == false { - if completePendingCancellationIfNeeded(for: job) { - return recoveryState.currentRun - } - markReviewFailed( - job, - message: ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription + let failure = ReviewRuntimeCloseFailure.connection( + ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription ) + await admission.recordConnectionTerminal(failure) + markReviewInterrupted(job, cause: .transport(message: failure.localizedDescription)) } return recoveryState.currentRun } private func handleReviewEventsFinished( job: CodexReviewJob, - isWaitingForNetworkRecovery: Bool - ) throws -> Bool { + isWaitingForNetworkRecovery: Bool, + admission: ReviewStartAdmission + ) async -> Bool { if Task.isCancelled { - throw CancellationError() + return true } if isWaitingForNetworkRecovery { @@ -747,13 +745,11 @@ extension CodexReviewStore { } if job.isTerminal == false { - if completePendingCancellationIfNeeded(for: job) { - return true - } - markReviewFailed( - job, - message: ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription + let failure = ReviewRuntimeCloseFailure.connection( + ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription ) + await admission.recordConnectionTerminal(failure) + markReviewInterrupted(job, cause: .transport(message: failure.localizedDescription)) } return true } @@ -771,7 +767,8 @@ extension CodexReviewStore { job: CodexReviewJob, startRequest: CodexReviewBackendModel.Review.Start, inputs: ReviewWorkerInputs, - recoveryToken: CodexReviewBackendModel.Review.RecoveryToken? + recoveryToken: CodexReviewBackendModel.Review.RecoveryToken?, + admission: ReviewStartAdmission ) async throws -> NetworkRestoreRestartResult { if job.isTerminal || completePendingCancellationIfNeeded(for: job) { return .finished @@ -793,7 +790,12 @@ extension CodexReviewStore { request: startRequest ) let recoveredRun = recoveredAttempt.run - if try await stopRecoveredRunIfJobShouldNotResume(recoveredRun, job: job) { + await admission.recordActiveRun(recoveredRun) + if try await stopRecoveredRunIfJobShouldNotResume( + recoveredRun, + job: job, + admission: admission + ) { return .finished } return .recovered(recoveredAttempt) @@ -801,14 +803,15 @@ extension CodexReviewStore { private func stopRecoveredRunIfJobShouldNotResume( _ recoveredRun: CodexReviewBackendModel.Review.Run, - job: CodexReviewJob + job: CodexReviewJob, + admission: ReviewStartAdmission ) async throws -> Bool { if Task.isCancelled { try? await backend.interruptReview( recoveredRun, reason: .init(message: job.core.lifecycle.cancellation?.message ?? "Cancellation requested.") ) - await backend.cleanupReview(recoveredRun) + try await cleanupReview(recoveredRun, admission: admission) throw CancellationError() } @@ -819,7 +822,7 @@ extension CodexReviewStore { reason: .init(message: job.core.lifecycle.cancellation?.message ?? "Cancellation requested.") ) } - await backend.cleanupReview(recoveredRun) + try await cleanupReview(recoveredRun, admission: admission) return true } @@ -836,7 +839,7 @@ extension CodexReviewStore { cancellation: cancellation ) } catch { - await backend.cleanupReview(recoveredRun) + try await cleanupReview(recoveredRun, admission: admission) try? recordCancellationFailure( jobID: job.id, sessionID: job.sessionID, @@ -844,7 +847,7 @@ extension CodexReviewStore { ) throw error } - await backend.cleanupReview(recoveredRun) + try await cleanupReview(recoveredRun, admission: admission) return true } @@ -856,10 +859,6 @@ extension CodexReviewStore { guard job.isTerminal == false else { return currentRun } - if event.completesReviewRun, completePendingCancellationIfNeeded(for: job) { - writeDiagnosticsIfNeeded() - return currentRun - } let updatedRun = currentRun switch event { case .started: @@ -930,6 +929,25 @@ extension CodexReviewStore { return updatedRun } + private func reviewTerminalRecord( + for event: CodexReviewBackendModel.Review.Event, + job: CodexReviewJob + ) -> ReviewTerminalRecord? { + switch event { + case .completed: + return .completed + case .failed(let message): + return .failed(message: message?.nilIfEmpty) + case .cancelled(let message): + if let cancellation = job.core.lifecycle.cancellation { + return .interrupted(.requested(cancellation)) + } + return .interrupted(.server(message: message?.nilIfEmpty)) + case .started, .message, .messageDelta, .log, .logEntry: + return nil + } + } + private func completePendingCancellationIfNeeded(for job: CodexReviewJob) -> Bool { guard job.cancellationRequested else { return false diff --git a/Sources/CodexReview/Store/CodexReviewStoreTesting.swift b/Sources/CodexReview/Store/CodexReviewStoreTesting.swift index 851c7d42..33b9f10e 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreTesting.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreTesting.swift @@ -94,8 +94,8 @@ extension CodexReviewStore { reviewWorkerTasks.removeAll(keepingCapacity: false) runtimeStopDetachedReviewWorkerTasks.removeAll(keepingCapacity: false) - startingJobIDs.removeAll(keepingCapacity: false) - startupCancellations.removeAll(keepingCapacity: false) + reviewStartAdmissions.removeAll(keepingCapacity: false) + reviewCleanupFailures.removeAll(keepingCapacity: false) activeRuns.removeAll(keepingCapacity: false) reviewRecoveryWaitingJobIDs.removeAll(keepingCapacity: false) diff --git a/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift b/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift index d75426d9..4ef1d95d 100644 --- a/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift +++ b/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift @@ -133,12 +133,17 @@ package class PreviewCodexReviewStoreBackend: CodexReviewStoreBackend { false } - package func startReview(_: CodexReviewBackendModel.Review.Start) async throws -> BackendReviewAttempt { + package func startReview( + _: CodexReviewBackendModel.Review.Start, + admission _: ReviewStartAdmission + ) async throws -> BackendReviewAttempt { throw CodexReviewAPI.Error.io(Self.previewUnavailableMessage) } package func interruptReview(_: CodexReviewBackendModel.Review.Run, reason _: CodexReviewBackendModel.CancellationReason) async throws {} + package func forceCloseReviewConnection() async throws {} + package func beginReviewRecovery( _: CodexReviewBackendModel.Review.Run, reason _: CodexReviewBackendModel.CancellationReason @@ -153,7 +158,7 @@ package class PreviewCodexReviewStoreBackend: CodexReviewStoreBackend { throw CodexReviewAPI.Error.io(Self.previewUnavailableMessage) } - package func cleanupReview(_: CodexReviewBackendModel.Review.Run) async {} + package func cleanupReview(_: CodexReviewBackendModel.Review.Run) async throws {} fileprivate static let previewUnavailableMessage = "Embedded server is unavailable in preview mode." fileprivate static let previewAuthenticationFailureMessage = "Authentication is unavailable in preview mode." diff --git a/Sources/CodexReviewAppServer/AppServerClient.swift b/Sources/CodexReviewAppServer/AppServerClient.swift index e3ba083c..6022e691 100644 --- a/Sources/CodexReviewAppServer/AppServerClient.swift +++ b/Sources/CodexReviewAppServer/AppServerClient.swift @@ -133,8 +133,8 @@ package actor AppServerClient { await transport.notificationStream() } - package func close() async { - await transport.close() + package func close() async throws { + try await transport.close() } private func allocateRequestID() -> Int { diff --git a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift index 9a3feb9a..9ec67e3b 100644 --- a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift +++ b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift @@ -126,12 +126,15 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { return try await readAuth() } - package func startReview(_ request: CodexReviewBackendModel.Review.Start) async throws -> BackendReviewAttempt { + package func startReview( + _ request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission + ) async throws -> BackendReviewAttempt { _ = try await client.initialize() await ensureNotificationRouterStarted() let control = AppServerReviewControl(client: client) - let thread = try await startReviewThread(request) + let thread = try await startReviewThread(request, admission: admission) controlsByThreadID[thread.threadID] = control let attemptID = makeAppServerReviewAttemptID() let provisionalRun = CodexReviewBackendModel.Review.Run( @@ -147,6 +150,14 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { ) registerReviewEventSession(session, for: provisionalRun) control.recordThreadStarted(threadID: thread.threadID) + await admission.recordPreparedThread(provisionalRun) + + guard await admission.admitReviewStartDispatch(for: provisionalRun) else { + try await cleanupReview(provisionalRun) + throw ReviewStartCancelledBeforeDispatch( + cancellation: await admission.cancellationRequest() ?? .system() + ) + } let review: AppServerAPI.Review.Start.Response reviewStartRequestsInFlight += 1 @@ -157,7 +168,10 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } catch { reviewStartRequestsInFlight -= 1 discardUnmatchedReviewNotificationsIfIdle() - await cleanupReview(provisionalRun) + if let terminal = Self.connectionTerminal(for: error) { + await admission.recordConnectionTerminal(terminal) + } + try await cleanupReview(provisionalRun) throw error } let reviewThreadID = review.reviewThreadID ?? thread.threadID @@ -173,64 +187,89 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { control.recordReviewStarted(turnThreadID: appServerTurnThreadID(for: run), turnID: review.turnID) await session.bufferStartupNotifications(takeUnmatchedReviewNotifications(for: run)) await session.finalizeRun() + await admission.recordActiveRun(run) reviewStartRequestsInFlight -= 1 discardUnmatchedReviewNotificationsIfIdle() return await session.attempt() } - private func startReviewThread(_ request: CodexReviewBackendModel.Review.Start) async throws -> AppServerAPI.Thread.Start.Response { + private func startReviewThread( + _ request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission + ) async throws -> AppServerAPI.Thread.Start.Response { if threadStartPermissionStrategy == .legacySandbox { // Deprecated compatibility: installed Codex builds without the app-server v2 // session-source flag can ignore permissions without failing the request. - return try await client.send(AppServerAPI.Thread.Start.Request( + return try await sendThreadStart(AppServerAPI.Thread.Start.Request( params: threadStartParamsWithLegacySandbox(request) - )) + ), admission: admission) } do { - return try await startReviewThreadWithProfileIDPermissions(request) + return try await startReviewThreadWithProfileIDPermissions(request, admission: admission) } catch let error as JSONRPC.Error where Self.shouldRetryThreadStartWithLegacySandbox(error) { // Deprecated compatibility: some builds accept the permissions field shape // without registering the danger-full-access built-in profile. - return try await client.send(AppServerAPI.Thread.Start.Request( + return try await sendThreadStart(AppServerAPI.Thread.Start.Request( params: threadStartParamsWithLegacySandbox(request) - )) + ), admission: admission) } catch let error as JSONRPC.Error where Self.shouldRetryThreadStartWithObjectPermissions(error) { // Deprecated compatibility: installed Codex builds can require object-shaped // permissions while the latest local app-server source accepts a profile ID string. - return try await startReviewThreadWithProfileSelectionPermissions(request) + return try await startReviewThreadWithProfileSelectionPermissions(request, admission: admission) } } private func startReviewThreadWithProfileIDPermissions( - _ request: CodexReviewBackendModel.Review.Start + _ request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission ) async throws -> AppServerAPI.Thread.Start.Response { - try await client.send(AppServerAPI.Thread.Start.Request( + try await sendThreadStart(AppServerAPI.Thread.Start.Request( params: threadStartParams( request, permissions: .profileID(Self.reviewPermissionProfileID) ) - )) + ), admission: admission) } private func startReviewThreadWithProfileSelectionPermissions( - _ request: CodexReviewBackendModel.Review.Start + _ request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission ) async throws -> AppServerAPI.Thread.Start.Response { do { - return try await client.send(AppServerAPI.Thread.Start.Request( + return try await sendThreadStart(AppServerAPI.Thread.Start.Request( params: threadStartParams( request, permissions: .profileSelection(.init(id: Self.reviewPermissionProfileID)) ) - )) + ), admission: admission) } catch let error as JSONRPC.Error where Self.shouldRetryThreadStartWithLegacySandbox(error) { // Deprecated compatibility: installed Codex builds can know the permissions // object shape without registering the danger-full-access built-in profile. - return try await client.send(AppServerAPI.Thread.Start.Request( + return try await sendThreadStart(AppServerAPI.Thread.Start.Request( params: threadStartParamsWithLegacySandbox(request) - )) + ), admission: admission) + } + } + + private func sendThreadStart( + _ request: AppServerAPI.Thread.Start.Request, + admission: ReviewStartAdmission + ) async throws -> AppServerAPI.Thread.Start.Response { + guard await admission.admitThreadStartDispatch() else { + throw ReviewStartCancelledBeforeDispatch( + cancellation: await admission.cancellationRequest() ?? .system() + ) + } + do { + return try await client.send(request) + } catch { + if let terminal = Self.connectionTerminal(for: error) { + await admission.recordConnectionTerminal(terminal) + } + throw error } } @@ -277,6 +316,29 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { || message.contains("default_permissions refers to unknown") } + private nonisolated static func interruptRequestFailure( + for error: any Error + ) -> ReviewInterruptRequestFailure { + if case JSONRPC.Error.responseError(let code, let message) = error { + return .init(outcome: .rejected(code: code, message: message)) + } + return .init(outcome: .outcomeUnknown(message: error.localizedDescription)) + } + + private nonisolated static func connectionTerminal( + for error: any Error + ) -> ReviewRuntimeCloseFailure? { + guard let jsonRPCError = error as? JSONRPC.Error else { + return nil + } + switch jsonRPCError { + case .closed, .invalidMessage: + return .connection(jsonRPCError.localizedDescription) + case .responseError: + return nil + } + } + package func interruptReview(_ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason) async throws { _ = try await client.initialize() guard abandonedReviewAttemptIDs.contains(run.attemptID) == false else { @@ -286,14 +348,25 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { await session.requestCancellation(message: reason.message) do { _ = try await sendTurnInterrupt(for: run) - await finishReviewEventStream( - threadID: run.threadID, - cancellationMessage: reason.message, - buffersMissingContinuation: true - ) + await session.waitForTerminalBarrier() } catch { - await session.clearCancellationRequest() - throw error + let failure = Self.interruptRequestFailure(for: error) + if case .rejected = failure.outcome { + await session.clearCancellationRequest() + throw failure + } + await session.waitForTerminalBarrier() + throw failure + } + } + + package func forceCloseReviewConnection() async throws { + do { + try await client.close() + } catch let failure as ReviewRuntimeCloseFailure { + throw failure + } catch { + throw ReviewRuntimeCloseFailure.connection(error.localizedDescription) } } @@ -303,10 +376,10 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { _ = try await client.initialize() await ensureNotificationRouterStarted() - markTurnAbandoned(run.turnID) let interruption = try await sendTurnInterrupt(for: run) { retryInterruption in await self.markInterruptionTurnAbandoned(retryInterruption, canonicalThreadID: run.threadID) } + markTurnAbandoned(run.turnID) markAttemptAbandoned(run, interruption: interruption) if let session = unregisterReviewEventSession(for: run) { await session.abandon() @@ -383,8 +456,7 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { return await session.attempt() } - package func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async { - _ = try? await client.initialize() + package func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { controlsByThreadID.removeValue(forKey: run.threadID) var cleanupThreadIDs = cleanupThreadIDs(for: run) if let session = unregisterReviewEventSession(for: run) { @@ -395,24 +467,40 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { completedReviewEventSessionMetricsByThreadID[threadID] = metrics } } - let _: EmptyResponse? = try? await client.send(AppServerAPI.Thread.BackgroundTerminals.Clean.Request( - params: .init(threadID: run.threadID) - )) - let _: AppServerAPI.Thread.Unsubscribe.Response? = try? await client.send(AppServerAPI.Thread.Unsubscribe.Request( - params: .init(threadID: run.threadID) - )) - for threadID in cleanupThreadIDs { - let _: EmptyResponse? = try? await client.send(AppServerAPI.Thread.Delete.Request( - params: .init(threadID: threadID) + var firstFailure: ReviewRuntimeCloseFailure? + do { + let _: EmptyResponse = try await client.send(AppServerAPI.Thread.BackgroundTerminals.Clean.Request( + params: .init(threadID: run.threadID) + )) + } catch { + firstFailure = .cleanup("thread/backgroundTerminals/clean for \(run.threadID): \(error.localizedDescription)") + } + do { + let _: AppServerAPI.Thread.Unsubscribe.Response = try await client.send(AppServerAPI.Thread.Unsubscribe.Request( + params: .init(threadID: run.threadID) )) + } catch where firstFailure == nil { + firstFailure = .cleanup("thread/unsubscribe for \(run.threadID): \(error.localizedDescription)") + } catch {} + for threadID in cleanupThreadIDs { + do { + let _: EmptyResponse = try await client.send(AppServerAPI.Thread.Delete.Request( + params: .init(threadID: threadID) + )) + } catch where firstFailure == nil { + firstFailure = .cleanup("thread/delete for \(threadID): \(error.localizedDescription)") + } catch {} } for threadID in cleanupThreadIDs { reviewEventSessionCanonicalThreadIDByThreadID.removeValue(forKey: threadID) } reviewThreadIDsForCleanupByThreadID.removeValue(forKey: run.threadID) + if let firstFailure { + throw firstFailure + } } - package func cleanupActiveReviewsForShutdown(reason: CodexReviewBackendModel.CancellationReason) async { + package func cleanupActiveReviewsForShutdown(reason: CodexReviewBackendModel.CancellationReason) async throws { let runs = await activeReviewRunsForShutdown() guard runs.isEmpty == false else { return @@ -421,15 +509,15 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { if Task.isCancelled { return } - try? await interruptReview(run, reason: reason) + try await interruptReview(run, reason: reason) if Task.isCancelled { return } - await cleanupReview(run) + try await cleanupReview(run) } } - package func interruptActiveReviewsForShutdown(reason: CodexReviewBackendModel.CancellationReason) async { + package func interruptActiveReviewsForShutdown(reason: CodexReviewBackendModel.CancellationReason) async throws { let runs = await activeReviewRunsForShutdown() guard runs.isEmpty == false else { return @@ -438,7 +526,7 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { if Task.isCancelled { return } - try? await interruptReview(run, reason: reason) + try await interruptReview(run, reason: reason) } } @@ -610,20 +698,6 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { unmatchedReviewNotificationsByThreadID.removeAll(keepingCapacity: true) } - private func finishReviewEventStream( - threadID: String, - cancellationMessage: String?, - buffersMissingContinuation: Bool = false - ) async { - guard let session = reviewEventSession(forThreadID: threadID) else { - return - } - await session.finish( - cancellationMessage: cancellationMessage, - buffersMissingContinuation: buffersMissingContinuation - ) - } - private func sendTurnInterrupt( for run: CodexReviewBackendModel.Review.Run, willInterruptActiveTurn: (@Sendable (AppServerReviewInterruption) async -> Void)? = nil @@ -693,7 +767,7 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { for try await notification in notifications { await routeReviewNotification(notification) } - await finishAllReviewEventSessions(throwing: nil) + await finishAllReviewEventSessions(throwing: JSONRPC.Error.closed) } catch { await finishAllReviewEventSessions(throwing: error) } @@ -839,7 +913,13 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { appServerBackendLogger.error( "Closing app-server connection after review routing failure: \(error.localizedDescription, privacy: .public)" ) - await client.close() + do { + try await client.close() + } catch { + appServerBackendLogger.error( + "App-server connection close failed during routing containment: \(error.localizedDescription, privacy: .public)" + ) + } await finishAllReviewEventSessions(throwing: error) } @@ -1007,6 +1087,8 @@ private actor AppServerReviewEventSession { private var isRunFinalized: Bool private var isDrainingStartupNotifications = false private var pendingStartupNotifications: [AppServerRoutedReviewNotification] = [] + private var terminalBarrierResolved = false + private var terminalBarrierWaiters: [CheckedContinuation] = [] private var metrics = AppServerReviewEventSessionMetrics() init( @@ -1074,6 +1156,19 @@ private actor AppServerReviewEventSession { cancellationRequestedMessage = nil } + func waitForTerminalBarrier() async { + if terminalBarrierResolved { + return + } + await withCheckedContinuation { continuation in + if terminalBarrierResolved { + continuation.resume() + } else { + terminalBarrierWaiters.append(continuation) + } + } + } + func receive(_ notification: AppServerRoutedReviewNotification) async { metrics.routed += 1 guard finished == false else { @@ -1114,6 +1209,7 @@ private actor AppServerReviewEventSession { pendingStartupNotifications.removeAll(keepingCapacity: true) await emitPrecedingEvents(precedingEvents) if let error { + resolveTerminalBarrier() await mailbox.fail(error) } else { await mailbox.finish() @@ -1401,11 +1497,24 @@ private actor AppServerReviewEventSession { event = .failed(message) } noteEmission(event) + resolveTerminalBarrier() await mailbox.append(event) recordReviewEvent(event, controlThreadID: controlThreadID) await mailbox.finish() } + private func resolveTerminalBarrier() { + guard terminalBarrierResolved == false else { + return + } + terminalBarrierResolved = true + let waiters = terminalBarrierWaiters + terminalBarrierWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume() + } + } + private func shouldCloseActiveCommandsBeforeEvents( notification: AppServerRoutedReviewNotification, decoded: DecodedReviewNotification diff --git a/Sources/CodexReviewAppServer/AppServerProcessTransport.swift b/Sources/CodexReviewAppServer/AppServerProcessTransport.swift index 89c68172..1e0f6a00 100644 --- a/Sources/CodexReviewAppServer/AppServerProcessTransport.swift +++ b/Sources/CodexReviewAppServer/AppServerProcessTransport.swift @@ -150,15 +150,15 @@ package actor AppServerProcessTransport: JSONRPC.Transport { } } - package func close() async { - await closeTransport(terminateProcess: true, readerTask: nil) + package func close() async throws { + try await closeTransport(terminateProcess: true, readerTask: nil) } private func closeTransport( terminateProcess: Bool, readerTask: ReaderTask? - ) async { - await closeTransport( + ) async throws { + try await closeTransport( terminateProcess: terminateProcess, error: JSONRPC.Error.closed, readerTask: readerTask @@ -169,7 +169,7 @@ package actor AppServerProcessTransport: JSONRPC.Transport { terminateProcess: Bool, error: any Error, readerTask: ReaderTask? - ) async { + ) async throws { if closed { if readerTask == nil { await waitForReaderTasks(excluding: nil) @@ -180,12 +180,20 @@ package actor AppServerProcessTransport: JSONRPC.Transport { stdoutEvents.cancel() stderrEvents.cancel() try? stdin.fileHandleForWriting.close() + var processCloseError: (any Error)? if terminateProcess { logger.info("Terminating codex app-server pid \(self.process.processIdentifier, privacy: .public)") - await process.terminateAndWait() + do { + try await process.terminateAndWait() + } catch { + processCloseError = error + } } - finishAll(throwing: error) + finishAll(throwing: processCloseError ?? error) await waitForReaderTasks(excluding: readerTask) + if let processCloseError { + throw processCloseError + } } private func receiveStdout(_ event: AppServerPipeReadEvent) async { @@ -204,11 +212,15 @@ package actor AppServerProcessTransport: JSONRPC.Transport { try processMessage(message) } catch { logger.error("Closing codex app-server after invalid JSON-RPC framing: \(error.localizedDescription, privacy: .public)") - await closeTransport( - terminateProcess: true, - error: error, - readerTask: .stdout - ) + do { + try await closeTransport( + terminateProcess: true, + error: error, + readerTask: .stdout + ) + } catch { + logger.error("Failed to close codex app-server process: \(error.localizedDescription, privacy: .public)") + } return } } @@ -242,15 +254,23 @@ package actor AppServerProcessTransport: JSONRPC.Transport { try processMessage(message) } catch { logger.error("Closing codex app-server after invalid trailing JSON-RPC framing: \(error.localizedDescription, privacy: .public)") - await closeTransport( - terminateProcess: true, - error: error, - readerTask: .stdout - ) + do { + try await closeTransport( + terminateProcess: true, + error: error, + readerTask: .stdout + ) + } catch { + logger.error("Failed to close codex app-server process: \(error.localizedDescription, privacy: .public)") + } return } } - await closeTransport(terminateProcess: true, readerTask: .stdout) + do { + try await closeTransport(terminateProcess: true, readerTask: .stdout) + } catch { + logger.error("Failed to close codex app-server process after stdout EOF: \(error.localizedDescription, privacy: .public)") + } } private func processMessage(_ data: Data) throws { @@ -729,7 +749,7 @@ private final class AppServerSpawnedProcess: @unchecked Sendable { func terminateAndWait( graceDuration: Duration = .seconds(2), killDuration: Duration = .seconds(1) - ) async { + ) async throws { let trackedProcessIDs = descendantProcessIDs() guard isFullyTerminated(trackedProcessIDs: trackedProcessIDs) == false else { return @@ -739,7 +759,9 @@ private final class AppServerSpawnedProcess: @unchecked Sendable { return } signalProcessTree(SIGKILL, trackedProcessIDs: trackedProcessIDs) - _ = await waitUntilExit(timeout: killDuration, trackedProcessIDs: trackedProcessIDs) + guard await waitUntilExit(timeout: killDuration, trackedProcessIDs: trackedProcessIDs) else { + throw AppServerProcessTransportError.processDidNotTerminate(processIdentifier) + } } private func signalProcessTree(_ signal: Int32, trackedProcessIDs: Set) { @@ -888,6 +910,7 @@ private final class AppServerSpawnedProcess: @unchecked Sendable { private enum AppServerProcessTransportError: LocalizedError { case executableNotFound(command: String, path: String?) + case processDidNotTerminate(pid_t) var errorDescription: String? { switch self { @@ -897,6 +920,8 @@ private enum AppServerProcessTransportError: LocalizedError { return "Unable to locate \(command) executable in PATH: \(resolvedPath)" } return "Unable to locate \(command) executable. Set PATH so codex can be found." + case .processDidNotTerminate(let processIdentifier): + return "Codex app-server process \(processIdentifier) did not terminate after SIGKILL." } } } diff --git a/Sources/CodexReviewAppServer/AppServerReviewControl.swift b/Sources/CodexReviewAppServer/AppServerReviewControl.swift index 1e8f60a9..11142a9f 100644 --- a/Sources/CodexReviewAppServer/AppServerReviewControl.swift +++ b/Sources/CodexReviewAppServer/AppServerReviewControl.swift @@ -86,45 +86,11 @@ package final class AppServerReviewControl: @unchecked Sendable { private func sendInterrupt( threadID: String, turnID: String, - willInterruptActiveTurn: (@Sendable (AppServerReviewInterruption) async -> Void)? + willInterruptActiveTurn _: (@Sendable (AppServerReviewInterruption) async -> Void)? ) async throws -> AppServerReviewInterruption { - do { - let _: EmptyResponse = try await client.send(AppServerAPI.Turn.Interrupt.Request( - params: .init(threadID: threadID, turnID: turnID) - )) - return .init(threadID: threadID, turnID: turnID) - } catch { - guard let activeTurnID = Self.activeTurnID(from: error), - activeTurnID != turnID - else { - throw error - } - let activeInterruption = AppServerReviewInterruption(threadID: threadID, turnID: activeTurnID) - if let willInterruptActiveTurn { - await willInterruptActiveTurn(activeInterruption) - } - let _: EmptyResponse = try await client.send(AppServerAPI.Turn.Interrupt.Request( - params: .init(threadID: threadID, turnID: activeTurnID) - )) - setPhase(.reviewStarted(turnThreadID: threadID, turnID: activeTurnID)) - return activeInterruption - } - } - - private func setPhase(_ phase: Phase) { - phaseLock.lock() - defer { phaseLock.unlock() } - self.phase = phase - } - - private static func activeTurnID(from error: Error) -> String? { - guard case JSONRPC.Error.responseError(_, let message) = error, - let range = message.range(of: " but found ") - else { - return nil - } - return String(message[range.upperBound...]) - .trimmingCharacters(in: CharacterSet(charactersIn: "` ").union(.whitespacesAndNewlines)) - .nilIfEmpty + let _: EmptyResponse = try await client.send(AppServerAPI.Turn.Interrupt.Request( + params: .init(threadID: threadID, turnID: turnID) + )) + return .init(threadID: threadID, turnID: turnID) } } diff --git a/Sources/CodexReviewAppServer/JSONRPC.swift b/Sources/CodexReviewAppServer/JSONRPC.swift index c7314648..6c0ff932 100644 --- a/Sources/CodexReviewAppServer/JSONRPC.swift +++ b/Sources/CodexReviewAppServer/JSONRPC.swift @@ -27,7 +27,7 @@ package enum JSONRPC { func send(_ request: Request) async throws -> Data func notify(_ notification: Notification) async throws func notificationStream() async -> AsyncThrowingStream - func close() async + func close() async throws } package enum Error: Swift.Error, Equatable, Sendable, LocalizedError { diff --git a/Sources/CodexReviewHost/CodexReviewHost.swift b/Sources/CodexReviewHost/CodexReviewHost.swift index 0b2ab211..4d92aec5 100644 --- a/Sources/CodexReviewHost/CodexReviewHost.swift +++ b/Sources/CodexReviewHost/CodexReviewHost.swift @@ -7,7 +7,7 @@ import CodexReviewMCPServer package final class CodexReviewHost { package let store: CodexReviewStore package let mcpServer: CodexReviewMCPServer - private let shutdown: @Sendable () async -> Void + private let shutdown: @Sendable () async throws -> Void private var endpoint: URL? package init( @@ -15,7 +15,7 @@ package final class CodexReviewHost { clock: CodexReviewClock = .init(), idGenerator: CodexReviewIDGenerator = .init(), endpoint: URL? = nil, - shutdown: @escaping @Sendable () async -> Void = {} + shutdown: @escaping @Sendable () async throws -> Void = {} ) { self.shutdown = shutdown self.endpoint = endpoint @@ -38,7 +38,7 @@ package final class CodexReviewHost { backend: backend, endpoint: endpoint, shutdown: { - await client.close() + try await client.close() } ) } @@ -51,9 +51,9 @@ package final class CodexReviewHost { await store.refreshSettings() } - package func stop() async { + package func stop() async throws { await store.stop() - await shutdown() + try await shutdown() } } @@ -233,8 +233,11 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { false } - func startReview(_ request: CodexReviewBackendModel.Review.Start) async throws -> BackendReviewAttempt { - try await backend.startReview(request) + func startReview( + _ request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission + ) async throws -> BackendReviewAttempt { + try await backend.startReview(request, admission: admission) } func interruptReview( @@ -244,6 +247,10 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { try await backend.interruptReview(run, reason: reason) } + func forceCloseReviewConnection() async throws { + try await backend.forceCloseReviewConnection() + } + func beginReviewRecovery( _ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason @@ -258,8 +265,8 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { try await backend.resumeReviewRecovery(token, request: request) } - func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async { - await backend.cleanupReview(run) + func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { + try await backend.cleanupReview(run) } private static func monitorSettings( diff --git a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift index 0fb4fa89..7879b63c 100644 --- a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift +++ b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift @@ -434,7 +434,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { let failureMessage = await runtimeStartupFailureMessage(for: error) logger.error("Review runtime failed to start: \(failureMessage, privacy: .public)") await startedHTTPServer?.stop() - await startedClient?.close() + await closeClientAfterFailure(startedClient) self.client = nil self.appServerBackend = nil self.mcpHTTPServer = nil @@ -477,7 +477,13 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { ) async { store.recordActiveReviewCancellationRequestsForRuntimeStop(reason: reason) let didInterrupt = await runRuntimeShutdownCleanup(timeout: shutdownCleanupTimeout) { - await appServerBackend.interruptActiveReviewsForShutdown(reason: .init(message: reason.message)) + do { + try await appServerBackend.interruptActiveReviewsForShutdown( + reason: .init(message: reason.message) + ) + } catch { + logger.error("Failed to interrupt active reviews during runtime teardown: \(error.localizedDescription, privacy: .public)") + } } let locallyCancelledJobIDs = store.cancelActiveReviewsLocallyForRuntimeStop( reason: reason, @@ -518,7 +524,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { await mcpHTTPServer?.stop() self.appServerBackend = nil await cleanupLoginRuntime(loginCleanup) - await client?.close() + await closeClientAfterFailure(client) logger.info("Review runtime stopped") } @@ -1033,11 +1039,14 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { await closeIsolatedLoginRuntime(client: loginClient, codexHomeURL: loginCodexHomeURL) } - func startReview(_ request: CodexReviewBackendModel.Review.Start) async throws -> BackendReviewAttempt { + func startReview( + _ request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission + ) async throws -> BackendReviewAttempt { guard let appServerBackend else { throw CodexReviewAPI.Error.io("Review runtime is not running.") } - return try await appServerBackend.startReview(request) + return try await appServerBackend.startReview(request, admission: admission) } func interruptReview(_ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason) async throws { @@ -1047,6 +1056,13 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { try await appServerBackend.interruptReview(run, reason: reason) } + func forceCloseReviewConnection() async throws { + guard let appServerBackend else { + throw ReviewRuntimeCloseFailure.connection("Review runtime is not running.") + } + try await appServerBackend.forceCloseReviewConnection() + } + func beginReviewRecovery( _ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason @@ -1067,11 +1083,11 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { return try await appServerBackend.resumeReviewRecovery(token, request: request) } - func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async { + func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { guard let appServerBackend else { - return + throw ReviewRuntimeCloseFailure.cleanup("Review runtime is not running.") } - await appServerBackend.cleanupReview(run) + try await appServerBackend.cleanupReview(run) } @discardableResult @@ -1191,7 +1207,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { store.transitionToFailed(message) await failedMCPHTTPServer?.stop() await cleanupLoginRuntime(loginCleanup) - await failedClient?.close() + await closeClientAfterFailure(failedClient) } private func handleAuthNotification( @@ -1537,16 +1553,27 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private func closeIsolatedLoginRuntime(client: AppServerClient?, codexHomeURL: URL?) async { guard let codexHomeURL else { - await client?.close() + await closeClientAfterFailure(client) return } guard codexHomeURL != self.codexHomeURL else { return } - await client?.close() + await closeClientAfterFailure(client) try? FileManager.default.removeItem(at: codexHomeURL) } + private func closeClientAfterFailure(_ client: AppServerClient?) async { + guard let client else { + return + } + do { + try await client.close() + } catch { + logger.error("Failed to close app-server client: \(error.localizedDescription, privacy: .public)") + } + } + private func takeLoginRuntimeForCleanup() -> PendingLoginRuntimeCleanup { loginChallenge = nil loginBackend = nil diff --git a/Sources/CodexReviewTesting/TestSupport.swift b/Sources/CodexReviewTesting/TestSupport.swift index 2cc7413a..f61aee48 100644 --- a/Sources/CodexReviewTesting/TestSupport.swift +++ b/Sources/CodexReviewTesting/TestSupport.swift @@ -176,6 +176,7 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { case logout(CodexReviewBackendModel.Account.ID) case startReview(CodexReviewBackendModel.Review.Start) case interruptReview(CodexReviewBackendModel.Review.Run, CodexReviewBackendModel.CancellationReason) + case forceCloseReviewConnection case beginReviewRecovery(CodexReviewBackendModel.Review.Run, CodexReviewBackendModel.CancellationReason) case resumeReviewRecovery(CodexReviewBackendModel.Review.RecoveryToken, CodexReviewBackendModel.Review.Start) case cleanupReview(CodexReviewBackendModel.Review.Run) @@ -188,6 +189,7 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { private var nextRecoveredRun: CodexReviewBackendModel.Review.Run? private var interruptFailureMessage: String? private var recoveryFailureMessage: String? + private var cleanupFailure: ReviewRuntimeCloseFailure? private var interruptReviewGate: AsyncGate? private var interruptReviewWaiters: [UUID: CheckedContinuation] = [:] private var beginReviewRecoveryWaiters: [UUID: CheckedContinuation] = [:] @@ -239,6 +241,10 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { recoveryFailureMessage = message } + package func failCleanup(message: String) { + cleanupFailure = .cleanup(message) + } + package func holdInterruptReview(with gate: AsyncGate) { interruptReviewGate = gate } @@ -447,16 +453,38 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { return auth } - package func startReview(_ request: CodexReviewBackendModel.Review.Start) async throws -> BackendReviewAttempt { + package func startReview( + _ request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission + ) async throws -> BackendReviewAttempt { + guard await admission.admitThreadStartDispatch() else { + throw ReviewStartCancelledBeforeDispatch( + cancellation: await admission.cancellationRequest() ?? .system() + ) + } commands.append(.startReview(request)) let waiters = Array(startReviewWaiters.values) startReviewWaiters.removeAll(keepingCapacity: false) for waiter in waiters { waiter.resume() } + let provisionalRun = CodexReviewBackendModel.Review.Run( + attemptID: nextRun.attemptID, + threadID: nextRun.threadID, + reviewThreadID: nextRun.threadID, + model: nextRun.model + ) + await admission.recordPreparedThread(provisionalRun) + guard await admission.admitReviewStartDispatch(for: provisionalRun) else { + commands.append(.cleanupReview(provisionalRun)) + throw ReviewStartCancelledBeforeDispatch( + cancellation: await admission.cancellationRequest() ?? .system() + ) + } if let startReviewGate { await startReviewGate.wait() } + await admission.recordActiveRun(nextRun) return .init(run: nextRun, events: eventMailbox(for: nextRun)) } @@ -471,7 +499,19 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { await interruptReviewGate.wait() } if let interruptFailureMessage { - throw FakeCodexReviewBackendError(message: interruptFailureMessage) + throw ReviewInterruptRequestFailure( + outcome: .rejected(code: nil, message: interruptFailureMessage) + ) + } + } + + package func forceCloseReviewConnection() async throws { + commands.append(.forceCloseReviewConnection) + await startReviewGate?.open() + await interruptReviewGate?.open() + let mailboxes = Array(eventMailboxes.values) + for mailbox in mailboxes { + await mailbox.fail(ReviewRuntimeCloseFailure.connection("Connection force-closed.")) } } @@ -521,8 +561,11 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { return .init(run: recoveredRun, events: eventMailbox(for: recoveredRun)) } - package func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async { + package func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { commands.append(.cleanupReview(run)) + if let cleanupFailure { + throw cleanupFailure + } } package func yield(_ event: CodexReviewBackendModel.Review.Event, for run: CodexReviewBackendModel.Review.Run? = nil) async { @@ -814,8 +857,11 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { false } - package func startReview(_ request: CodexReviewBackendModel.Review.Start) async throws -> BackendReviewAttempt { - try await reviewBackend.startReview(request) + package func startReview( + _ request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission + ) async throws -> BackendReviewAttempt { + try await reviewBackend.startReview(request, admission: admission) } package func interruptReview( @@ -825,6 +871,10 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { try await reviewBackend.interruptReview(run, reason: reason) } + package func forceCloseReviewConnection() async throws { + try await reviewBackend.forceCloseReviewConnection() + } + package func beginReviewRecovery( _ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason @@ -839,8 +889,8 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { try await reviewBackend.resumeReviewRecovery(token, request: request) } - package func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async { - await reviewBackend.cleanupReview(run) + package func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { + try await reviewBackend.cleanupReview(run) } package func refreshSettings() async throws -> CodexReviewSettings.Snapshot { @@ -913,8 +963,13 @@ package actor FakeJSONRPCTransport: JSONRPC.Transport { private var gatesByMethod: [String: RequestGate] = [:] private var oneShotGatesByMethod: [String: [RequestGate]] = [:] private var requestCountWaiters: [(Int, CheckedContinuation)] = [] + private var requestMethodWaiters: [String: [CheckedContinuation]] = [:] + private var responseMethodWaiters: [String: [CheckedContinuation]] = [:] + private var deliveredResponseMethods: [String: Int] = [:] private var notificationStreamCountWaiters: [(Int, CheckedContinuation)] = [] private var closed = false + private var closeFailure: ReviewRuntimeCloseFailure? + private var closeCallCount = 0 package init(responses: [String: [Data]] = [:]) { self.responses = responses @@ -948,12 +1003,17 @@ package actor FakeJSONRPCTransport: JSONRPC.Transport { oneShotGatesByMethod[method, default: []].append(.init(gate: gate, ignoresCancellation: true)) } + package func failClose(with failure: ReviewRuntimeCloseFailure) { + closeFailure = failure + } + package func send(_ request: JSONRPC.Request) async throws -> Data { guard closed == false else { throw JSONRPC.Error.closed } requests.append(request) resumeRequestCountWaiters() + resumeRequestMethodWaiters(request.method) activeByMethod[request.method, default: 0] += 1 maxActiveByMethod[request.method] = max( maxActiveByMethod[request.method] ?? 0, @@ -964,6 +1024,11 @@ package actor FakeJSONRPCTransport: JSONRPC.Transport { await gate.wait() } activeByMethod[request.method, default: 1] -= 1 + guard closed == false else { + throw JSONRPC.Error.closed + } + deliveredResponseMethods[request.method, default: 0] += 1 + resumeResponseMethodWaiters(request.method) if let queuedResponse { switch queuedResponse { case .success(let data): @@ -1004,12 +1069,20 @@ package actor FakeJSONRPCTransport: JSONRPC.Transport { } } - package func close() async { + package func close() async throws { + closeCallCount += 1 closed = true + let gates = Array(gatesByMethod.values) + oneShotGatesByMethod.values.flatMap { $0 } + for gate in gates { + await gate.gate.open() + } for continuation in serverNotificationContinuations { - continuation.finish() + continuation.finish(throwing: JSONRPC.Error.closed) } serverNotificationContinuations.removeAll() + if let closeFailure { + throw closeFailure + } } package func finishNotificationStreams(throwing error: any Error) { @@ -1036,6 +1109,32 @@ package actor FakeJSONRPCTransport: JSONRPC.Transport { } } + package func waitForRequest(method: String) async { + if requests.contains(where: { $0.method == method }) { + return + } + await withCheckedContinuation { continuation in + if requests.contains(where: { $0.method == method }) { + continuation.resume() + } else { + requestMethodWaiters[method, default: []].append(continuation) + } + } + } + + package func waitForResponseDelivery(method: String) async { + if deliveredResponseMethods[method, default: 0] > 0 { + return + } + await withCheckedContinuation { continuation in + if deliveredResponseMethods[method, default: 0] > 0 { + continuation.resume() + } else { + responseMethodWaiters[method, default: []].append(continuation) + } + } + } + package func recordedNotifications() -> [JSONRPC.Notification] { notifications } @@ -1061,6 +1160,10 @@ package actor FakeJSONRPCTransport: JSONRPC.Transport { closed } + package func closeCallCountForTesting() -> Int { + closeCallCount + } + package func maxActiveCount(for method: String) -> Int { maxActiveByMethod[method] ?? 0 } @@ -1090,6 +1193,20 @@ package actor FakeJSONRPCTransport: JSONRPC.Transport { requestCountWaiters = remaining } + private func resumeRequestMethodWaiters(_ method: String) { + let waiters = requestMethodWaiters.removeValue(forKey: method) ?? [] + for waiter in waiters { + waiter.resume() + } + } + + private func resumeResponseMethodWaiters(_ method: String) { + let waiters = responseMethodWaiters.removeValue(forKey: method) ?? [] + for waiter in waiters { + waiter.resume() + } + } + private func resumeNotificationStreamCountWaiters() { var remaining: [(Int, CheckedContinuation)] = [] for waiter in notificationStreamCountWaiters { diff --git a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift index 5c85e151..b18df189 100644 --- a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift +++ b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift @@ -34,8 +34,8 @@ private extension AppServerCodexReviewBackend { try await beginReviewRecovery(attempt.run, reason: reason) } - func cleanupReview(_ attempt: BackendReviewAttempt) async { - await cleanupReview(attempt.run) + func cleanupReview(_ attempt: BackendReviewAttempt) async throws { + try await cleanupReview(attempt.run) } } @@ -80,6 +80,42 @@ private func eventSequence( return BackendReviewEventSequence(mailbox: attempt.events) } +private func interruptAndDeliverCanonicalTerminal( + backend: AppServerCodexReviewBackend, + transport: FakeJSONRPCTransport, + run: CodexReviewBackendModel.Review.Run, + reason: CodexReviewBackendModel.CancellationReason +) async throws { + async let cancellation: Void = backend.interruptReview(run, reason: reason) + await transport.waitForResponseDelivery(method: "turn/interrupt") + try await transport.emitServerNotification( + method: "turn/completed", + params: TestTurnNotification( + threadID: run.reviewThreadID ?? run.threadID, + turn: .init( + id: run.turnID ?? "", + status: "interrupted", + error: .init(message: reason.message) + ) + ) + ) + try await cancellation +} + +private func interruptAndDeliverCanonicalTerminal( + backend: AppServerCodexReviewBackend, + transport: FakeJSONRPCTransport, + run attempt: BackendReviewAttempt, + reason: CodexReviewBackendModel.CancellationReason +) async throws { + try await interruptAndDeliverCanonicalTerminal( + backend: backend, + transport: transport, + run: attempt.run, + reason: reason + ) +} + @Suite("app-server client") struct AppServerClientTests { @Test func processTransportConfigurationResolvesCodexFromProvidedPath() throws { @@ -330,7 +366,7 @@ struct AppServerClientTests { } #expect(Darwin.kill(childPID, 0) == 0) - await transport.close() + try await transport.close() let childExited = await waitUntil(timeout: .seconds(2)) { Darwin.kill(childPID, 0) != 0 && errno == ESRCH @@ -376,7 +412,7 @@ struct AppServerClientTests { method: "test/request", params: Data("{}".utf8) )) - await transport.close() + try await transport.close() let object = try #require(JSONSerialization.jsonObject(with: data) as? [String: Any]) #expect(object["value"] as? String == "done") @@ -450,7 +486,7 @@ struct AppServerClientTests { var iterator = notifications.makeAsyncIterator() let notification = try #require(try await iterator.next()) - await transport.close() + try await transport.close() #expect(notification.method == "item/completed") #expect(try JSONSerialization.jsonObject( @@ -503,7 +539,7 @@ struct AppServerClientTests { let notificationWritten = await waitUntil(timeout: .seconds(2)) { FileManager.default.fileExists(atPath: notificationFile.path) } - await transport.close() + try await transport.close() #expect(notificationWritten) let request = try #require(JSONSerialization.jsonObject( @@ -668,7 +704,7 @@ struct AppServerClientTests { #expect(params.turnID == "turn-1") } - @Test func runningInterruptRetriesWithCurrentActiveTurnID() async throws { + @Test func runningInterruptRejectsCanonicalMismatchWithoutErrorTextRebinding() async throws { let transport = FakeJSONRPCTransport() await transport.enqueueFailure( .responseError( @@ -677,20 +713,21 @@ struct AppServerClientTests { ), for: "turn/interrupt" ) - try await transport.enqueue(EmptyResponse(), for: "turn/interrupt") let client = AppServerClient(transport: transport) let control = AppServerReviewControl(client: client) control.recordReviewStarted(turnThreadID: "thread-1", turnID: "turn-old") - let interruption = try await control.interrupt() - #expect(interruption == .init(threadID: "thread-1", turnID: "turn-new")) + await #expect(throws: JSONRPC.Error.responseError( + code: -32602, + message: "expected active turn id turn-old but found turn-new" + )) { + try await control.interrupt() + } let requests = await transport.recordedRequests() - #expect(requests.map(\.method) == ["turn/interrupt", "turn/interrupt"]) + #expect(requests.map(\.method) == ["turn/interrupt"]) let first = try JSONDecoder().decode(AppServerAPI.Turn.Interrupt.Params.self, from: requests[0].params) - let second = try JSONDecoder().decode(AppServerAPI.Turn.Interrupt.Params.self, from: requests[1].params) #expect(first.turnID == "turn-old") - #expect(second.turnID == "turn-new") } @Test func initializeSendsInitializedNotificationOnce() async throws { @@ -1255,7 +1292,12 @@ struct AppServerClientTests { var iterator = events.makeAsyncIterator() #expect(try await iterator.next() == .started(turnID: "turn-old", reviewThreadID: "review-thread", model: nil)) - try await backend.interruptReview(run, reason: .init()) + try await interruptAndDeliverCanonicalTerminal( + backend: backend, + transport: transport, + run: run, + reason: .init() + ) let request = try #require(await transport.recordedRequests().last) let params = try JSONDecoder().decode(AppServerAPI.Turn.Interrupt.Params.self, from: request.params) @@ -1276,7 +1318,12 @@ struct AppServerClientTests { sessionID: "session-1", request: .init(cwd: "/tmp/project", target: .uncommittedChanges) )) - try await backend.interruptReview(run, reason: .init(message: "Stop")) + try await interruptAndDeliverCanonicalTerminal( + backend: backend, + transport: transport, + run: run, + reason: .init(message: "Stop") + ) let request = try #require(await transport.recordedRequests().last) #expect(request.method == "turn/interrupt") @@ -1356,7 +1403,12 @@ struct AppServerClientTests { #expect(try await iterator.next() == .messageDelta("review text", itemID: "message-1")) #expect(await backend.reviewEventSessionMetricsForTesting(threadID: "review-thread")?.routed == 2) - try await backend.interruptReview(run, reason: .init(message: "Stop")) + try await interruptAndDeliverCanonicalTerminal( + backend: backend, + transport: transport, + run: run, + reason: .init(message: "Stop") + ) let interruptRequest = try #require(await transport.recordedRequests().last) #expect(interruptRequest.method == "turn/interrupt") let interruptParams = try JSONDecoder().decode(AppServerAPI.Turn.Interrupt.Params.self, from: interruptRequest.params) @@ -1409,7 +1461,7 @@ struct AppServerClientTests { await #expect(throws: BackendReviewEventMailboxError.self) { _ = try await iterator.next() } - await transport.close() + try await transport.close() } @Test func backendPreservesBufferedEventsBeforeNotificationStreamError() async throws { @@ -1461,7 +1513,12 @@ struct AppServerClientTests { #expect(try await iterator.next() == .started(turnID: "turn-new", reviewThreadID: "review-thread", model: nil)) #expect(try await iterator.next() == .messageDelta("review text", itemID: "message-1")) - try await backend.interruptReview(run, reason: .init(message: "Stop")) + try await interruptAndDeliverCanonicalTerminal( + backend: backend, + transport: transport, + run: run, + reason: .init(message: "Stop") + ) let interruptRequest = try #require(await transport.recordedRequests().last) #expect(interruptRequest.method == "turn/interrupt") let interruptParams = try JSONDecoder().decode(AppServerAPI.Turn.Interrupt.Params.self, from: interruptRequest.params) @@ -1504,7 +1561,12 @@ struct AppServerClientTests { #expect(try await iterator.next() == .started(turnID: "turn-old", reviewThreadID: "review-thread", model: nil)) #expect(await backend.reviewEventSessionMetricsForTesting(threadID: "review-thread")?.routed == 1) - try await backend.interruptReview(run, reason: .init(message: "Stop")) + try await interruptAndDeliverCanonicalTerminal( + backend: backend, + transport: transport, + run: run, + reason: .init(message: "Stop") + ) let interruptRequest = try #require(await transport.recordedRequests().last) #expect(interruptRequest.method == "turn/interrupt") let interruptParams = try JSONDecoder().decode(AppServerAPI.Turn.Interrupt.Params.self, from: interruptRequest.params) @@ -1546,7 +1608,12 @@ struct AppServerClientTests { var iterator = await eventSequence(backend, run).makeAsyncIterator() #expect(try await iterator.next() == .started(turnID: "turn-response", reviewThreadID: "thread-1", model: nil)) - try await backend.interruptReview(run, reason: .init(message: "Stop")) + try await interruptAndDeliverCanonicalTerminal( + backend: backend, + transport: transport, + run: run, + reason: .init(message: "Stop") + ) let interruptRequest = try #require(await transport.recordedRequests().last) #expect(interruptRequest.method == "turn/interrupt") let interruptParams = try JSONDecoder().decode(AppServerAPI.Turn.Interrupt.Params.self, from: interruptRequest.params) @@ -1730,7 +1797,7 @@ struct AppServerClientTests { )) } - @Test func backendBuffersCancellationBeforeEventStreamRegistration() async throws { + @Test func backendInterruptAckDoesNotSynthesizeTerminalBeforeEventStreamRegistration() async throws { let transport = FakeJSONRPCTransport() try await enqueueInitialize(transport) try await transport.enqueue(AppServerAPI.Thread.Start.Response(threadID: "thread-1", model: "gpt-5"), for: "thread/start") @@ -1744,7 +1811,17 @@ struct AppServerClientTests { request: .init(cwd: "/tmp/project", target: .uncommittedChanges) )) - try await backend.interruptReview(run, reason: .init(message: "Stop")) + async let cancellation: Void = backend.interruptReview(run, reason: .init(message: "Stop")) + await transport.waitForResponseDelivery(method: "turn/interrupt") + #expect(await run.events.isFinished() == false) + try await transport.emitServerNotification( + method: "turn/completed", + params: TestTurnNotification( + threadID: "thread-1", + turn: .init(id: "turn-1", status: "interrupted", error: .init(message: "Stop")) + ) + ) + try await cancellation var iterator = await eventSequence(backend, run).makeAsyncIterator() #expect(try await iterator.next() == .cancelled("Stop")) @@ -2145,16 +2222,23 @@ struct AppServerClientTests { #expect(recoveryInterruptRequested) async let cancellation: Void = backend.interruptReview(run, reason: .init(message: "Stop")) - let cancellationInterruptRequested = await waitUntil { - await transport.recordedRequests().filter { $0.method == "turn/interrupt" }.count == 2 - } - #expect(cancellationInterruptRequested) - await interruptGate.open() do { _ = try await recovery Issue.record("Expected recovery interrupt to fail.") } catch {} + let cancellationInterruptRequested = await waitUntil { + await transport.recordedRequests().filter { $0.method == "turn/interrupt" }.count == 2 + } + #expect(cancellationInterruptRequested) + + try await transport.emitServerNotification( + method: "turn/completed", + params: TestTurnNotification( + threadID: "thread-1", + turn: .init(id: "turn-1", status: "interrupted", error: .init(message: "Stop")) + ) + ) try await cancellation #expect(try await iterator.next() == .cancelled("Stop")) @@ -2163,7 +2247,7 @@ struct AppServerClientTests { #expect(interruptRequests.count == 2) } - @Test func backendSuppressesRecoveryInterruptRetriedToActiveTurn() async throws { + @Test func backendRecoveryInterruptRejectsMismatchWithoutRebinding() async throws { let transport = FakeJSONRPCTransport() try await enqueueInitialize(transport) await transport.enqueueFailure( @@ -2173,7 +2257,6 @@ struct AppServerClientTests { ), for: "turn/interrupt" ) - try await transport.enqueue(EmptyResponse(), for: "turn/interrupt") let backend = AppServerCodexReviewBackend(client: .init(transport: transport)) let run = CodexReviewBackendModel.Review.Run( threadID: "thread-1", @@ -2183,79 +2266,22 @@ struct AppServerClientTests { ) let events = await eventSequence(backend, run) - _ = try await backend.beginReviewRecovery( - run, - reason: .init(message: "Network unavailable; waiting to reconnect.") - ) + await #expect(throws: JSONRPC.Error.responseError( + code: -32602, + message: "expected active turn id turn-old but found turn-active" + )) { + try await backend.beginReviewRecovery( + run, + reason: .init(message: "Network unavailable; waiting to reconnect.") + ) + } let requests = await transport.recordedRequests() let interruptRequests = requests.filter { $0.method == "turn/interrupt" } let interruptTurnIDs = try interruptRequests.map { request in try JSONDecoder().decode(AppServerAPI.Turn.Interrupt.Params.self, from: request.params).turnID } - #expect(interruptTurnIDs == ["turn-old", "turn-active"]) - - try await transport.emitServerNotification( - method: "turn/completed", - params: TestTurnNotification( - threadID: "thread-1", - turn: .init(id: "turn-active", status: "interrupted", error: .init(message: "Network unavailable")) - ) - ) - let ignoredInterruptedTurn = await waitUntil { - await backend.notificationRouterMetricsForTesting().ignored == 1 - } - #expect(ignoredInterruptedTurn) - _ = events - } - - @Test func backendSuppressesActiveTurnTerminalWhileRecoveryRetryInterruptIsInFlight() async throws { - let transport = FakeJSONRPCTransport() - try await enqueueInitialize(transport) - await transport.enqueueFailure( - .responseError( - code: -32602, - message: "expected active turn id turn-old but found turn-active" - ), - for: "turn/interrupt" - ) - try await transport.enqueue(EmptyResponse(), for: "turn/interrupt") - let firstInterruptGate = AsyncGate() - await firstInterruptGate.open() - let retryInterruptGate = AsyncGate() - await transport.holdNext(method: "turn/interrupt", gate: firstInterruptGate) - await transport.holdNext(method: "turn/interrupt", gate: retryInterruptGate) - let backend = AppServerCodexReviewBackend(client: .init(transport: transport)) - let run = CodexReviewBackendModel.Review.Run( - threadID: "thread-1", - turnID: "turn-old", - reviewThreadID: "thread-1", - model: "gpt-5" - ) - let events = await eventSequence(backend, run) - - async let recovery: CodexReviewBackendModel.Review.RecoveryToken = backend.beginReviewRecovery( - run, - reason: .init(message: "Network unavailable; waiting to reconnect.") - ) - let retryInterruptRequested = await waitUntil { - await transport.recordedRequests().filter { $0.method == "turn/interrupt" }.count == 2 - } - #expect(retryInterruptRequested) - - try await transport.emitServerNotification( - method: "turn/completed", - params: TestTurnNotification( - threadID: "thread-1", - turn: .init(id: "turn-active", status: "interrupted", error: .init(message: "Network unavailable")) - ) - ) - let ignoredTerminal = await waitUntil { - await backend.notificationRouterMetricsForTesting().ignored == 1 - } - #expect(ignoredTerminal) + #expect(interruptTurnIDs == ["turn-old"]) - await retryInterruptGate.open() - _ = try await recovery _ = events } @@ -2303,7 +2329,7 @@ struct AppServerClientTests { )) } - @Test func backendIgnoresStaleTerminalWhileRecoveryInterruptIsInFlight() async throws { + @Test func backendCanonicalTerminalWinsWhileRecoveryInterruptIsInFlight() async throws { let transport = FakeJSONRPCTransport() try await enqueueInitialize(transport) try await transport.enqueue(EmptyResponse(), for: "turn/interrupt") @@ -2319,6 +2345,7 @@ struct AppServerClientTests { model: "gpt-5" ) let initialEvents = await eventSequence(backend, run) + var initialIterator = initialEvents.makeAsyncIterator() defer { withExtendedLifetime(initialEvents) {} } async let recovered = backend.resumeReviewRecovery( @@ -2343,10 +2370,7 @@ struct AppServerClientTests { turn: .init(id: "turn-1", status: "failed", error: .init(message: "Old turn failed")) ) ) - let ignoredStaleTerminal = await waitUntil { - await backend.notificationRouterMetricsForTesting().ignored == 1 - } - #expect(ignoredStaleTerminal) + #expect(try await initialIterator.next() == .failed("Old turn failed")) await interruptGate.open() let recoveredRun = try await recovered @@ -2437,71 +2461,6 @@ struct AppServerClientTests { )) } - @Test func backendSuppressesRetriedActiveTurnNotificationsDuringRollback() async throws { - let transport = FakeJSONRPCTransport() - try await enqueueInitialize(transport) - await transport.enqueueFailure( - .responseError( - code: -32602, - message: "expected active turn id turn-old but found turn-active" - ), - for: "turn/interrupt" - ) - try await transport.enqueue(EmptyResponse(), for: "turn/interrupt") - try await transport.enqueue(EmptyResponse(), for: "thread/rollback") - try await transport.enqueue(AppServerAPI.Review.Start.Response(turnID: "turn-2", reviewThreadID: "thread-1"), for: "review/start") - let rollbackGate = AsyncGate() - await transport.holdNext(method: "thread/rollback", gate: rollbackGate) - let backend = AppServerCodexReviewBackend(client: .init(transport: transport)) - let run = CodexReviewBackendModel.Review.Run( - threadID: "thread-1", - turnID: "turn-old", - reviewThreadID: "thread-1", - model: "gpt-5" - ) - let initialEvents = await eventSequence(backend, run) - defer { withExtendedLifetime(initialEvents) {} } - - async let recovered = backend.resumeReviewRecovery( - run, - request: CodexReviewBackendModel.Review.Start( - jobID: "job-1", - sessionID: "session-1", - request: .init(cwd: "/tmp/project", target: .baseBranch("main")), - model: "gpt-5" - ), - reason: .init(message: "Network unavailable; waiting to reconnect.") - ) - let rollbackRequested = await waitUntil { - await transport.recordedRequests().contains { $0.method == "thread/rollback" } - } - #expect(rollbackRequested) - - try await transport.emitServerNotification( - method: "turn/started", - params: TestTurnNotification(threadID: "thread-1", turn: .init(id: "turn-active")) - ) - let ignoredStaleNotification = await waitUntil { - await backend.notificationRouterMetricsForTesting().ignored == 1 - } - #expect(ignoredStaleNotification) - - await rollbackGate.open() - let recoveredRun = try await recovered - #expect(recoveredRun.turnID == "turn-2") - let recoveredEvents = await eventSequence(backend, recoveredRun) - var iterator = recoveredEvents.makeAsyncIterator() - try await transport.emitServerNotification( - method: "turn/started", - params: TestTurnNotification(threadID: "thread-1", turn: .init(id: "turn-2")) - ) - #expect(try await iterator.next() == .started( - turnID: "turn-2", - reviewThreadID: "thread-1", - model: nil - )) - } - @Test func backendRecoveryClearsInterruptedCommandStateBeforeReplayingRecoveredTurn() async throws { let transport = FakeJSONRPCTransport() try await enqueueInitialize(transport) @@ -2605,6 +2564,10 @@ struct AppServerClientTests { try await transport.enqueue(EmptyResponse(), for: "turn/interrupt") try await transport.enqueue(EmptyResponse(), for: "thread/rollback") try await transport.enqueue(AppServerAPI.Review.Start.Response(turnID: "turn-2", reviewThreadID: "review-thread-2"), for: "review/start") + try await transport.enqueue( + AppServerAPI.Thread.Unsubscribe.Response(status: .unsubscribed), + for: "thread/unsubscribe" + ) let backend = AppServerCodexReviewBackend(client: .init(transport: transport)) let run = CodexReviewBackendModel.Review.Run( threadID: "thread-1", @@ -2623,7 +2586,7 @@ struct AppServerClientTests { ), reason: .init(message: "Network unavailable; waiting to reconnect.") ) - await backend.cleanupReview(recovered) + try await backend.cleanupReview(recovered) let deleteThreadIDs = try await transport.recordedRequests() .filter { $0.method == "thread/delete" } @@ -3026,7 +2989,7 @@ struct AppServerClientTests { #expect(try await iterator.next() == nil) } - @Test func backendInterruptFinishesReviewEventStream() async throws { + @Test func backendInterruptAckWaitsForCanonicalTerminalToFinishEventStream() async throws { let transport = FakeJSONRPCTransport() try await enqueueInitialize(transport) try await transport.enqueue(AppServerAPI.Thread.Start.Response(threadID: "thread-1", model: "gpt-5"), for: "thread/start") @@ -3048,7 +3011,17 @@ struct AppServerClientTests { #expect(try await iterator.next() == .started(turnID: "turn-1", reviewThreadID: "thread-1", model: nil)) - try await backend.interruptReview(run, reason: .init(message: "Stop")) + async let cancellation: Void = backend.interruptReview(run, reason: .init(message: "Stop")) + await transport.waitForResponseDelivery(method: "turn/interrupt") + #expect(await run.events.isFinished() == false) + try await transport.emitServerNotification( + method: "turn/completed", + params: TestTurnNotification( + threadID: "thread-1", + turn: .init(id: "turn-1", status: "interrupted", error: .init(message: "Stop")) + ) + ) + try await cancellation #expect(try await iterator.next() == .cancelled("Stop")) #expect(try await iterator.next() == nil) @@ -3098,7 +3071,16 @@ struct AppServerClientTests { ) )) - try await backend.interruptReview(run, reason: .init(message: "Stop")) + async let cancellation: Void = backend.interruptReview(run, reason: .init(message: "Stop")) + await transport.waitForResponseDelivery(method: "turn/interrupt") + try await transport.emitServerNotification( + method: "turn/completed", + params: TestTurnNotification( + threadID: "thread-1", + turn: .init(id: "turn-1", status: "interrupted", error: .init(message: "Stop")) + ) + ) + try await cancellation guard case .logEntry(let kind, let text, let groupID, let replacesGroup, let metadata) = try await iterator.next() else { @@ -3175,7 +3157,16 @@ struct AppServerClientTests { } #expect(routedOutput) - try await backend.interruptReview(run, reason: .init(message: "Stop")) + async let cancellation: Void = backend.interruptReview(run, reason: .init(message: "Stop")) + await transport.waitForResponseDelivery(method: "turn/interrupt") + try await transport.emitServerNotification( + method: "turn/completed", + params: TestTurnNotification( + threadID: "thread-1", + turn: .init(id: "turn-1", status: "interrupted", error: .init(message: "Stop")) + ) + ) + try await cancellation var sawClosedOutput = false while let event = try await iterator.next() { @@ -3293,7 +3284,12 @@ struct AppServerClientTests { var iterator = events.makeAsyncIterator() #expect(try await iterator.next() == .started(turnID: "turn-old", reviewThreadID: "thread-1", model: nil)) - try await backend.interruptReview(run, reason: .init()) + try await interruptAndDeliverCanonicalTerminal( + backend: backend, + transport: transport, + run: run, + reason: .init() + ) let request = try #require(await transport.recordedRequests().last) let params = try JSONDecoder().decode(AppServerAPI.Turn.Interrupt.Params.self, from: request.params) @@ -3477,7 +3473,12 @@ struct AppServerClientTests { )) #expect(try await iterator.next() == .messageDelta("review output", itemID: "message-1")) - try await backend.interruptReview(run, reason: .init()) + try await interruptAndDeliverCanonicalTerminal( + backend: backend, + transport: transport, + run: run, + reason: .init() + ) let params = try JSONDecoder().decode( AppServerAPI.Turn.Interrupt.Params.self, @@ -4563,7 +4564,7 @@ struct AppServerClientTests { delta: "tail output\n" ) ) - await transport.close() + try await transport.close() #expect(try await iterator.next() == .logEntry( kind: .commandOutput, @@ -4572,7 +4573,9 @@ struct AppServerClientTests { replacesGroup: false, metadata: .init(sourceType: "commandExecution", title: "Command output", itemID: "cmd-1") )) - #expect(try await iterator.next() == nil) + await #expect(throws: BackendReviewEventMailboxError.self) { + _ = try await iterator.next() + } } @Test func backendReviewExitCompletesMissingCommandCompletion() async throws { @@ -5288,6 +5291,10 @@ struct AppServerClientTests { try await enqueueInitialize(transport) try await transport.enqueue(AppServerAPI.Thread.Start.Response(threadID: "thread-1"), for: "thread/start") await transport.enqueueFailure(.responseError(code: -32602, message: "invalid target"), for: "review/start") + try await transport.enqueue( + AppServerAPI.Thread.Unsubscribe.Response(status: .unsubscribed), + for: "thread/unsubscribe" + ) let backend = AppServerCodexReviewBackend(client: .init(transport: transport)) await #expect(throws: JSONRPC.Error.responseError(code: -32602, message: "invalid target")) { diff --git a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift index 680fe822..f30af30f 100644 --- a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift +++ b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift @@ -11,7 +11,7 @@ import CodexReviewTesting @Suite("host composition") @MainActor struct CodexReviewHostTests { - @Test func hostStartsAndStopsRuntimeWithFakeBackend() async { + @Test func hostStartsAndStopsRuntimeWithFakeBackend() async throws { let backend = FakeCodexReviewBackend() let host = CodexReviewHost( backend: backend, @@ -22,7 +22,7 @@ struct CodexReviewHostTests { #expect(host.store.serverState == .running) #expect(host.store.serverURL == URL(string: "http://localhost:9417/mcp")) - await host.stop() + try await host.stop() #expect(host.store.serverState == .stopped) } diff --git a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift index cd94bc1e..1b77f960 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift @@ -112,10 +112,15 @@ struct CodexReviewStoreCommandTests { jobID: "job-1", timeout: .seconds(1) ) - _ = try await store.cancelReview( - jobID: "job-1", - cancellation: .mcpClient(message: "Stop") - ) + let cancellation = Task { @MainActor in + try await store.cancelReview( + jobID: "job-1", + cancellation: .mcpClient(message: "Stop") + ) + } + await backend.waitForInterruptReview() + await backend.yield(.cancelled("Stop")) + _ = try await cancellation.value let final = try await awaited #expect(final.core.lifecycle.status == .cancelled) @@ -732,11 +737,13 @@ struct CodexReviewStoreCommandTests { request: .init(cwd: "/tmp/project", target: .baseBranch("main")) ) try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running, jobID: "job-1") != nil) - let cancel = try await store.cancelReview( + async let pendingCancel = store.cancelReview( jobID: "job-1", cancellation: .mcpClient(message: "Stop") ) + await backend.waitForInterruptReview() await backend.yield(.cancelled("Stop")) + let cancel = try await pendingCancel _ = try await result #expect(cancel.cancelled) @@ -778,11 +785,13 @@ struct CodexReviewStoreCommandTests { #expect(await waitUntil { store.job(id: "job-1")?.logText.hasSuffix(delta) == true }) - _ = try await store.cancelReview( + async let cancel = store.cancelReview( jobID: "job-1", cancellation: .mcpClient(message: "Stop") ) + await backend.waitForInterruptReview() await backend.yield(.cancelled("Stop")) + _ = try await cancel let read = try await result let job = try #require(store.job(id: "job-1")) @@ -793,6 +802,38 @@ struct CodexReviewStoreCommandTests { } } + @Test func cancellationPreservesTypedCleanupFailureAfterTerminal() async throws { + let backend = FakeCodexReviewBackend() + await backend.failCleanup(message: "unsubscribe failed") + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-1" }) + ) + try await withStoreCommandTestCleanup(backend: backend, store: store) { + async let result = store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running, jobID: "job-1") != nil) + let cancellation = Task { @MainActor in + try await store.cancelReview( + jobID: "job-1", + cancellation: .mcpClient(message: "Stop") + ) + } + await backend.waitForInterruptReview() + await backend.yield(.cancelled("Stop")) + + await #expect(throws: ReviewRuntimeCloseFailure.cleanup("unsubscribe failed")) { + try await cancellation.value + } + let read = try await result + + #expect(read.core.lifecycle.status == .cancelled) + #expect(store.reviewCleanupFailures["job-1"] == .cleanup("unsubscribe failed")) + } + } + @Test func transientNetworkOutageDoesNotRecoverReview() async throws { let backend = FakeCodexReviewBackend() let networkMonitor = ManualCodexReviewNetworkMonitor() @@ -1601,6 +1642,7 @@ struct CodexReviewStoreCommandTests { let inFlight = try store.readReview(jobID: "job-1") #expect(inFlight.core.lifecycle.status == .running) + await backend.yield(.cancelled("Review runtime stopped."), for: run) await interruptGate.open() await stopTask.value @@ -1814,9 +1856,14 @@ struct CodexReviewStoreCommandTests { try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running, jobID: "job-1") != nil) networkMonitor.yield(.init(status: .unsatisfied)) - _ = try await store.cancelReview(jobID: "job-1", cancellation: .mcpClient(message: "Stop")) - await debounceGate.open() + async let cancellation = store.cancelReview( + jobID: "job-1", + cancellation: .mcpClient(message: "Stop") + ) + await backend.waitForInterruptReview() await backend.yield(.cancelled("Stop")) + _ = try await cancellation + await debounceGate.open() let read = try await result #expect(read.core.lifecycle.status == .cancelled) @@ -1964,7 +2011,7 @@ struct CodexReviewStoreCommandTests { } } - @Test func cancelledReviewStaysCancelledWhenStreamClosesWithError() async throws { + @Test func acceptedInterruptFollowedByConnectionLossUsesTransportTerminal() async throws { let backend = FakeCodexReviewBackend() let store = CodexReviewStore.makeTestingStore( backend: TestingCodexReviewStoreBackend(reviewBackend: backend), @@ -1976,15 +2023,18 @@ struct CodexReviewStoreCommandTests { request: .init(cwd: "/tmp/project", target: .baseBranch("main")) ) try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running, jobID: "job-1") != nil) - _ = try await store.cancelReview( + async let cancellation = store.cancelReview( jobID: "job-1", cancellation: .mcpClient(message: "Stop") ) + await backend.waitForInterruptReview() await backend.finishEvents(throwing: StreamClosedError()) + let cancel = try await cancellation let read = try await result - #expect(read.core.lifecycle.status == .cancelled) - #expect(read.core.output.summary == "Stop") + #expect(cancel.cancelled == false) + #expect(read.core.lifecycle.status == .failed) + #expect(read.core.lifecycle.terminal?.kind == .interrupted) } } @@ -2071,7 +2121,7 @@ struct CodexReviewStoreCommandTests { } } - @Test func reviewStartCancellationInterruptsBackendRun() async throws { + @Test func eventConsumerCancellationIsAWorkerTerminalNotARequestedInterrupt() async throws { let backend = FakeCodexReviewBackend() let store = CodexReviewStore.makeTestingStore( backend: TestingCodexReviewStoreBackend(reviewBackend: backend), @@ -2086,16 +2136,14 @@ struct CodexReviewStoreCommandTests { await backend.finishEvents(throwing: CancellationError()) let read = try await result - #expect(read.core.lifecycle.status == .cancelled) + #expect(read.core.lifecycle.status == .failed) + #expect(read.core.lifecycle.terminal?.kind == .interrupted) let commands = await backend.recordedCommands() - #expect(commands.contains(.interruptReview( - .init(threadID: "thread-1", turnID: "turn-1", reviewThreadID: "review-thread-1"), - .init(message: "Cancellation requested.") - ))) + #expect(commands.contains { if case .interruptReview = $0 { true } else { false } } == false) } } - @Test func reviewStartTaskCancellationInterruptsBackendRun() async throws { + @Test func callerTaskCancellationDoesNotAbandonOrTerminalizeBackendRun() async throws { let backend = FakeCodexReviewBackend() let store = CodexReviewStore.makeTestingStore( backend: TestingCodexReviewStoreBackend(reviewBackend: backend), @@ -2109,14 +2157,13 @@ struct CodexReviewStoreCommandTests { ) } task.cancel() + await backend.waitForStartReview() + await backend.yield(.completed(summary: "Succeeded.", result: "review text")) let read = try await task.value - #expect(read.core.lifecycle.status == .cancelled) + #expect(read.core.lifecycle.status == .succeeded) let commands = await backend.recordedCommands() - #expect(commands.contains(.interruptReview( - .init(threadID: "thread-1", turnID: "turn-1", reviewThreadID: "review-thread-1"), - .init(message: "Cancellation requested.") - ))) + #expect(commands.contains { if case .interruptReview = $0 { true } else { false } } == false) } } @@ -2133,7 +2180,7 @@ struct CodexReviewStoreCommandTests { request: .init(cwd: "/tmp/project", target: .baseBranch("main")) ) try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running, jobID: "job-1") != nil) - await #expect(throws: FakeCodexReviewBackendError.self) { + await #expect(throws: ReviewInterruptRequestFailure.self) { try await store.cancelReview( jobID: "job-1", cancellation: .mcpClient(message: "Stop") @@ -2150,7 +2197,7 @@ struct CodexReviewStoreCommandTests { } } - @Test func cancelledReviewIgnoresBufferedTerminalEvents() async throws { + @Test func canonicalCompletionWinsWhileCancellationIsPending() async throws { let backend = FakeCodexReviewBackend() let store = CodexReviewStore.makeTestingStore( backend: TestingCodexReviewStoreBackend(reviewBackend: backend), @@ -2162,21 +2209,28 @@ struct CodexReviewStoreCommandTests { request: .init(cwd: "/tmp/project", target: .baseBranch("main")) ) try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running, jobID: "job-1") != nil) - _ = try await store.cancelReview( + async let cancellation = store.cancelReview( jobID: "job-1", cancellation: .mcpClient(message: "Stop") ) + await backend.waitForInterruptReview() await backend.yield(.completed(summary: "Succeeded.", result: "late result")) + let cancel = try await cancellation let read = try await result - #expect(read.core.lifecycle.status == .cancelled) - #expect(read.core.output.summary == "Stop") - #expect(read.core.output.lastAgentMessage == nil) + #expect(cancel.cancelled == false) + #expect(read.core.lifecycle.status == .succeeded) + #expect(read.core.output.lastAgentMessage == "late result") } } - @Test func terminalEventDuringPendingCancellationKeepsCancelledState() async throws { - let backend = FakeCodexReviewBackend() + @Test func canonicalFailureBeforeInterruptAckWins() async throws { + let run = CodexReviewBackendModel.Review.Run( + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "review-thread-1" + ) + let backend = FakeCodexReviewBackend(nextRun: run) let interruptGate = AsyncGate() await backend.holdInterruptReview(with: interruptGate) let store = CodexReviewStore.makeTestingStore( @@ -2188,6 +2242,7 @@ struct CodexReviewStoreCommandTests { sessionID: "session-1", request: .init(cwd: "/tmp/project", target: .uncommittedChanges) ) + try #require(await waitForRunAttemptActivation(store: store, run: run)) async let cancel = store.cancelReview(jobID: "job-1", cancellation: .mcpClient(message: "Stop")) try await backend.waitForInterruptReview(timeout: .seconds(2)) await backend.yield(.completed(summary: "Reviewer failed to output a response.", result: nil)) @@ -2195,8 +2250,8 @@ struct CodexReviewStoreCommandTests { _ = try await cancel let read = try await result - #expect(read.core.lifecycle.status == .cancelled) - #expect(read.core.output.summary == "Stop") + #expect(read.core.lifecycle.status == .failed) + #expect(read.core.output.summary == ReviewIngestionError.missingFinalReview.localizedDescription) #expect(read.core.output.hasFinalReview == false) } } @@ -2215,11 +2270,14 @@ struct CodexReviewStoreCommandTests { request: .init(cwd: "/tmp/project", target: .uncommittedChanges) ) try await backend.waitForStartReview(timeout: .seconds(2)) - let cancel = try await store.cancelReview(jobID: "job-1", cancellation: .mcpClient(message: "Stop")) - let cancelledDuringStartup = try #require(store.jobs.first) - #expect(cancel.core.lifecycle.status == .cancelled) - #expect(cancelledDuringStartup.core.lifecycle.status == .cancelled) + async let pendingCancel = store.cancelReview( + jobID: "job-1", + cancellation: .mcpClient(message: "Stop") + ) await gate.open() + await backend.waitForInterruptReview() + await backend.yield(.cancelled("Stop")) + let cancel = try await pendingCancel let read = try await result #expect(cancel.cancelled) diff --git a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift new file mode 100644 index 00000000..f901d417 --- /dev/null +++ b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift @@ -0,0 +1,556 @@ +import Foundation +import Testing +@testable import CodexReview +import CodexReviewTesting + +@Suite("Review attempt processor") +struct ReviewAttemptProcessorTests { + @Test func interruptAckBeforeTerminalRemainsInterrupting() async throws { + let (admission, run) = try await makeActiveAdmission() + let requestReturned = InvocationProbe() + + let cancellation = Task { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in await requestReturned.record() }, + forceClose: {} + ) + } + await requestReturned.waitForInvocation() + + #expect(await admission.currentPhase() == .interrupting(run)) + try await admission.recordCanonicalTerminal( + .interrupted(.requested(.mcpClient(message: "Stop"))), + for: run + ) + let resolution = try await cancellation.value + + #expect(resolution.terminal == .canonical( + run: run, + terminal: .interrupted(.requested(.mcpClient(message: "Stop"))) + )) + #expect(await admission.currentPhase() == .terminal(resolution.terminal)) + } + + @Test func terminalBeforeInterruptAckRetainsAndDrainsRequestTask() async throws { + let (admission, run) = try await makeActiveAdmission() + let requestStarted = InvocationProbe() + let requestGate = AsyncGate() + + let cancellation = Task { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in + await requestStarted.record() + await requestGate.waitIgnoringCancellation() + }, + forceClose: {} + ) + } + await requestStarted.waitForInvocation() + try await admission.recordCanonicalTerminal(.completed, for: run) + + #expect(await admission.currentPhase() == .finishing(.canonical( + run: run, + terminal: .completed + ))) + await requestGate.open() + let resolution = try await cancellation.value + + #expect(resolution.terminal == .canonical(run: run, terminal: .completed)) + } + + @Test func explicitRejectionReturnsAttemptToActive() async throws { + let (admission, run) = try await makeActiveAdmission() + let rejection = ReviewInterruptRequestFailure( + outcome: .rejected(code: -32_000, message: "Not active") + ) + + await #expect(throws: rejection) { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in throw rejection }, + forceClose: {} + ) + } + + #expect(await admission.currentPhase() == .active(run)) + #expect(await admission.cancellationRequest() == nil) + } + + @Test func rejectionAfterTerminalCannotRewriteTerminal() async throws { + let (admission, run) = try await makeActiveAdmission() + let requestStarted = InvocationProbe() + let requestGate = AsyncGate() + let rejection = ReviewInterruptRequestFailure( + outcome: .rejected(code: -32_000, message: "Already finished") + ) + + let cancellation = Task { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in + await requestStarted.record() + await requestGate.waitIgnoringCancellation() + throw rejection + }, + forceClose: {} + ) + } + await requestStarted.waitForInvocation() + try await admission.recordCanonicalTerminal(.completed, for: run) + await requestGate.open() + let resolution = try await cancellation.value + + #expect(resolution.terminal == .canonical(run: run, terminal: .completed)) + #expect(resolution.requestFailure == rejection) + } + + @Test func outcomeUnknownFollowedByCanonicalTerminalReturnsTerminal() async throws { + let (admission, run) = try await makeActiveAdmission() + let requestFailed = InvocationProbe() + let failure = ReviewInterruptRequestFailure( + outcome: .outcomeUnknown(message: "Response lost") + ) + + let cancellation = Task { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in + await requestFailed.record() + throw failure + }, + forceClose: {} + ) + } + await requestFailed.waitForInvocation() + try await admission.recordCanonicalTerminal( + .interrupted(.requested(.mcpClient(message: "Stop"))), + for: run + ) + let resolution = try await cancellation.value + + #expect(resolution.requestFailure == failure) + #expect(resolution.terminal == .canonical( + run: run, + terminal: .interrupted(.requested(.mcpClient(message: "Stop"))) + )) + } + + @Test func outcomeUnknownFollowedByConnectionTerminalPreservesBothDiagnostics() async throws { + let (admission, _) = try await makeActiveAdmission() + let requestFailed = InvocationProbe() + let failure = ReviewInterruptRequestFailure( + outcome: .outcomeUnknown(message: "Response lost") + ) + let connection = ReviewRuntimeCloseFailure.connection("Connection ended") + + let cancellation = Task { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in + await requestFailed.record() + throw failure + }, + forceClose: {} + ) + } + await requestFailed.waitForInvocation() + await admission.recordConnectionTerminal(connection) + + do { + _ = try await cancellation.value + Issue.record("Expected outcome-unknown cancellation failure.") + } catch let received as ReviewInterruptRequestFailure { + #expect(received.outcome == failure.outcome) + #expect(received.secondaryBarrierDiagnostic == connection.localizedDescription) + } + #expect(await admission.currentPhase() == .terminal(.connection(connection))) + } + + @Test func graceExpiryForceClosesOnceAndAwaitsConnectionAndRequestCompletion() async throws { + let graceGate = AsyncGate() + let (admission, _) = try await makeActiveAdmission( + closePolicy: controlledClosePolicy(gate: graceGate) + ) + let requestStarted = InvocationProbe() + let requestGate = AsyncGate() + let forceClose = InvocationProbe() + let connection = ReviewRuntimeCloseFailure.connection("Forced close") + + let cancellation = Task { + try await admission.cancel( + .system(message: "Stop"), + interrupt: { _, _ in + await requestStarted.record() + await requestGate.waitIgnoringCancellation() + throw ReviewInterruptRequestFailure( + outcome: .outcomeUnknown(message: "Connection closed before response") + ) + }, + forceClose: { + await forceClose.record() + await admission.recordConnectionTerminal(connection) + await requestGate.open() + } + ) + } + await requestStarted.waitForInvocation() + await graceGate.open() + await forceClose.waitForInvocation() + + await #expect(throws: ReviewInterruptRequestFailure.self) { + try await cancellation.value + } + #expect(await forceClose.invocationCount() == 1) + #expect(await admission.currentPhase() == .terminal(.connection(connection))) + } + + @Test func forceCloseFailureRemainsTypedAfterOutcomeUnknownRequestCompletes() async throws { + let graceGate = AsyncGate() + let (admission, _) = try await makeActiveAdmission( + closePolicy: controlledClosePolicy(gate: graceGate) + ) + let requestFailed = InvocationProbe() + let forceFailure = ReviewRuntimeCloseFailure.process("Process remained alive") + + let cancellation = Task { + try await admission.cancel( + .system(message: "Stop"), + interrupt: { _, _ in + await requestFailed.record() + throw ReviewInterruptRequestFailure( + outcome: .outcomeUnknown(message: "Response lost") + ) + }, + forceClose: { + throw forceFailure + } + ) + } + await requestFailed.waitForInvocation() + await graceGate.open() + + await #expect(throws: forceFailure) { + try await cancellation.value + } + } + + @Test func duplicateCancellationCallersJoinOneRequest() async throws { + let (admission, run) = try await makeActiveAdmission() + let requestStarted = InvocationProbe() + let requestGate = AsyncGate() + + let first = Task { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in + await requestStarted.record() + await requestGate.waitIgnoringCancellation() + }, + forceClose: {} + ) + } + let second = Task { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in + Issue.record("Duplicate caller installed a second interrupt operation.") + }, + forceClose: {} + ) + } + await requestStarted.waitForInvocation() + try await admission.recordCanonicalTerminal( + .interrupted(.requested(.mcpClient(message: "Stop"))), + for: run + ) + await requestGate.open() + + #expect(try await first.value == second.value) + #expect(await requestStarted.invocationCount() == 1) + } + + @Test func staleAndCrossTurnTerminalsCannotSatisfyBarrier() async throws { + let (admission, run) = try await makeActiveAdmission() + let requestReturned = InvocationProbe() + let cancellation = Task { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in await requestReturned.record() }, + forceClose: {} + ) + } + await requestReturned.waitForInvocation() + + var stale = run + stale.attemptID = "attempt-stale" + try await admission.recordCanonicalTerminal(.completed, for: stale) + var child = run + child.reviewThreadID = "review-child" + try await admission.recordCanonicalTerminal(.completed, for: child) + var crossTurn = run + crossTurn.turnID = "turn-other" + try await admission.recordCanonicalTerminal(.completed, for: crossTurn) + + #expect(await admission.currentPhase() == .interrupting(run)) + try await admission.recordCanonicalTerminal(.completed, for: run) + #expect(try await cancellation.value.terminal == .canonical(run: run, terminal: .completed)) + } + + @Test func conflictingDuplicateTerminalFailsWithoutRewrite() async throws { + let (admission, run) = try await makeActiveAdmission() + try await admission.recordCanonicalTerminal(.completed, for: run) + + await #expect(throws: ReviewAttemptContractFailure.self) { + try await admission.recordCanonicalTerminal( + .failed(message: "conflict"), + for: run + ) + } + #expect(await admission.currentPhase() == .terminal(.canonical( + run: run, + terminal: .completed + ))) + } + + @Test func queuedCancellationCompletesLocallyWithoutDispatch() async throws { + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) + + let resolution = try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in Issue.record("Queued cancellation dispatched interrupt.") }, + forceClose: { Issue.record("Queued cancellation force-closed connection.") } + ) + + #expect(resolution.terminal == .localCancellation(.mcpClient(message: "Stop"))) + } + + @Test func cancellationBeforeThreadDispatchRefusesWrite() async throws { + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) + let entered = InvocationProbe() + let dispatchGate = AsyncGate() + let startTask = await admission.start { admission in + await entered.record() + await dispatchGate.wait() + try Task.checkCancellation() + guard await admission.admitThreadStartDispatch() else { + throw ReviewStartCancelledBeforeDispatch( + cancellation: await admission.cancellationRequest() ?? .system() + ) + } + Issue.record("Thread request was dispatched after cancellation.") + return .init(run: canonicalRun) + } + await entered.waitForInvocation() + + let resolution = try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in Issue.record("Pre-dispatch cancellation interrupted a turn.") }, + forceClose: {} + ) + + #expect(resolution.terminal == .localCancellation(.mcpClient(message: "Stop"))) + await #expect(throws: CancellationError.self) { + try await startTask.value + } + } + + @Test func cancellationAfterThreadDispatchRefusesReviewDispatchAfterResponse() async throws { + let graceGate = AsyncGate() + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: graceGate)) + let threadDispatched = InvocationProbe() + let threadResponseGate = AsyncGate() + let startTask = await admission.start { admission in + #expect(await admission.admitThreadStartDispatch()) + await threadDispatched.record() + await threadResponseGate.waitIgnoringCancellation() + let provisional = provisionalRun + await admission.recordPreparedThread(provisional) + guard await admission.admitReviewStartDispatch(for: provisional) else { + throw ReviewStartCancelledBeforeDispatch( + cancellation: await admission.cancellationRequest() ?? .system() + ) + } + Issue.record("Review request was dispatched after cancellation.") + return .init(run: canonicalRun) + } + await threadDispatched.waitForInvocation() + + let cancellation = Task { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in Issue.record("Thread-only attempt interrupted an empty turn.") }, + forceClose: {} + ) + } + await threadResponseGate.open() + + #expect(try await cancellation.value.terminal == .localCancellation(.mcpClient(message: "Stop"))) + await #expect(throws: ReviewStartCancelledBeforeDispatch.self) { + try await startTask.value + } + } + + @Test func cancellationAfterThreadResponseRefusesNotSentReviewDispatch() async throws { + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) + let prepared = InvocationProbe() + let reviewDispatchGate = AsyncGate() + let startTask = await admission.start { admission in + #expect(await admission.admitThreadStartDispatch()) + await admission.recordPreparedThread(provisionalRun) + await prepared.record() + await reviewDispatchGate.waitIgnoringCancellation() + guard await admission.admitReviewStartDispatch(for: provisionalRun) else { + throw ReviewStartCancelledBeforeDispatch( + cancellation: await admission.cancellationRequest() ?? .system() + ) + } + Issue.record("Review request was dispatched after cancellation.") + return .init(run: canonicalRun) + } + await prepared.waitForInvocation() + + let cancellation = Task { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in Issue.record("Not-sent review was interrupted.") }, + forceClose: {} + ) + } + await reviewDispatchGate.open() + + #expect(try await cancellation.value.terminal == .localCancellation(.mcpClient(message: "Stop"))) + await #expect(throws: ReviewStartCancelledBeforeDispatch.self) { + try await startTask.value + } + } + + @Test func cancellationAfterReviewDispatchJoinsResponseThenInterruptsCanonicalRun() async throws { + let graceGate = AsyncGate() + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: graceGate)) + let reviewDispatched = InvocationProbe() + let reviewResponseGate = AsyncGate() + let interruptCalled = InvocationProbe() + let startTask = await admission.start { admission in + #expect(await admission.admitThreadStartDispatch()) + await admission.recordPreparedThread(provisionalRun) + #expect(await admission.admitReviewStartDispatch(for: provisionalRun)) + await reviewDispatched.record() + await reviewResponseGate.waitIgnoringCancellation() + await admission.recordActiveRun(canonicalRun) + return .init(run: canonicalRun) + } + await reviewDispatched.waitForInvocation() + + let cancellation = Task { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { run, _ in + #expect(run == canonicalRun) + await interruptCalled.record() + }, + forceClose: {} + ) + } + await reviewResponseGate.open() + _ = try await startTask.value + await interruptCalled.waitForInvocation() + try await admission.recordCanonicalTerminal( + .interrupted(.requested(.mcpClient(message: "Stop"))), + for: canonicalRun + ) + + #expect(try await cancellation.value.terminal == .canonical( + run: canonicalRun, + terminal: .interrupted(.requested(.mcpClient(message: "Stop"))) + )) + } + + @Test func duplicateCleanupCallersJoinOneOwnedTask() async throws { + let (admission, run) = try await makeActiveAdmission() + let cleanupStarted = InvocationProbe() + let cleanupGate = AsyncGate() + + async let first: Void = admission.cleanup(run: run) { + await cleanupStarted.record() + await cleanupGate.waitIgnoringCancellation() + } + async let second: Void = admission.cleanup(run: run) { + Issue.record("Duplicate cleanup caller installed a second cleanup Task.") + } + await cleanupStarted.waitForInvocation() + await cleanupGate.open() + _ = try await (first, second) + + #expect(await cleanupStarted.invocationCount() == 1) + } +} + +private let canonicalRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-1", + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "review-thread-1", + model: "gpt-5" +) + +private let provisionalRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-1", + threadID: "thread-1", + reviewThreadID: "thread-1", + model: "gpt-5" +) + +private func makeActiveAdmission( + closePolicy: ReviewRuntimeClosePolicy? = nil +) async throws -> (ReviewStartAdmission, CodexReviewBackendModel.Review.Run) { + let admission = ReviewStartAdmission( + closePolicy: closePolicy ?? controlledClosePolicy(gate: AsyncGate()) + ) + let startTask = await admission.start { admission in + #expect(await admission.admitThreadStartDispatch()) + await admission.recordPreparedThread(provisionalRun) + #expect(await admission.admitReviewStartDispatch(for: provisionalRun)) + await admission.recordActiveRun(canonicalRun) + return .init(run: canonicalRun) + } + _ = try await startTask.value + return (admission, canonicalRun) +} + +private func controlledClosePolicy(gate: AsyncGate) -> ReviewRuntimeClosePolicy { + ReviewRuntimeClosePolicy(terminalGrace: .seconds(10)) { _ in + await gate.wait() + try Task.checkCancellation() + } +} + +private actor InvocationProbe { + private var count = 0 + private var waiters: [CheckedContinuation] = [] + + func record() { + count += 1 + let waiters = waiters + self.waiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume() + } + } + + func waitForInvocation() async { + if count > 0 { + return + } + await withCheckedContinuation { continuation in + if count > 0 { + continuation.resume() + } else { + waiters.append(continuation) + } + } + } + + func invocationCount() -> Int { count } +} From 19f20550a96528090ce834eb0f2366814462b212 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 17:24:23 +0900 Subject: [PATCH 02/38] test(review): cover interrupt lifecycle races --- .../CodexReview/ReviewAttemptProcessor.swift | 32 ++++ .../CodexReview/Store/CodexReviewStore.swift | 6 +- .../Store/CodexReviewStoreCancellation.swift | 14 +- .../LiveCodexReviewStoreBackend.swift | 114 ++++-------- Sources/CodexReviewTesting/TestSupport.swift | 26 ++- .../AppServerClientTests.swift | 166 ++++++++++++++++++ .../CodexReviewHostTests.swift | 127 ++++++++++++-- 7 files changed, 384 insertions(+), 101 deletions(-) diff --git a/Sources/CodexReview/ReviewAttemptProcessor.swift b/Sources/CodexReview/ReviewAttemptProcessor.swift index ab4126fc..a031ffac 100644 --- a/Sources/CodexReview/ReviewAttemptProcessor.swift +++ b/Sources/CodexReview/ReviewAttemptProcessor.swift @@ -163,6 +163,7 @@ package actor ReviewStartAdmission { private var forceCloseResult: Result? private var cancellationResult: Result? private var terminalWaiters: [UUID: CheckedContinuation] = [:] + private var activeRunWaiters: [CheckedContinuation] = [] private var cancellationWaiters: [CheckedContinuation, Never>] = [] package init(closePolicy: ReviewRuntimeClosePolicy = .production) { @@ -231,6 +232,7 @@ package actor ReviewStartAdmission { return } phase = .active(run) + resumeActiveRunWaiters(returning: run) } package func recordCanonicalTerminal( @@ -303,6 +305,24 @@ package actor ReviewStartAdmission { package func currentPhase() -> Phase { phase } + package func waitForActiveRun() async -> CodexReviewBackendModel.Review.Run? { + if let activeRun { + return activeRun + } + if terminal != nil { + return nil + } + return await withCheckedContinuation { continuation in + if let activeRun { + continuation.resume(returning: activeRun) + } else if terminal != nil { + continuation.resume(returning: nil) + } else { + activeRunWaiters.append(continuation) + } + } + } + package func cancellationRequest() -> ReviewCancellation? { requestedCancellation } @@ -323,6 +343,7 @@ package actor ReviewStartAdmission { case .success(let attempt): if terminal == nil { phase = .active(attempt.run) + resumeActiveRunWaiters(returning: attempt.run) } case .failure(let error): if let cancellation = (error as? ReviewStartCancelledBeforeDispatch)?.cancellation { @@ -544,9 +565,20 @@ package actor ReviewStartAdmission { for waiter in waiters { waiter.resume(returning: terminal) } + resumeActiveRunWaiters(returning: nil) resolveCancellationIfPossible() } + private func resumeActiveRunWaiters( + returning run: CodexReviewBackendModel.Review.Run? + ) { + let waiters = activeRunWaiters + activeRunWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume(returning: run) + } + } + private func resolveCancellationIfPossible() { guard cancellationResult == nil else { return diff --git a/Sources/CodexReview/Store/CodexReviewStore.swift b/Sources/CodexReview/Store/CodexReviewStore.swift index e5f898b2..0be95602 100644 --- a/Sources/CodexReview/Store/CodexReviewStore.swift +++ b/Sources/CodexReview/Store/CodexReviewStore.swift @@ -154,7 +154,11 @@ public final class CodexReviewStore { if backend.handlesActiveReviewStopCleanup { locallyCancelledJobIDs = [] } else { - locallyCancelledJobIDs = await requestActiveReviewCancellationsForRuntimeStop() + do { + locallyCancelledJobIDs = try await requestActiveReviewCancellationsForRuntimeStop() + } catch { + locallyCancelledJobIDs = [] + } } await backend.stop(store: self) let remainingLocallyCancelledJobIDs = cancelActiveReviewsLocallyForRuntimeStop(cancelWorkers: false) diff --git a/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift b/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift index eccaf094..b29bfa2b 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift @@ -147,12 +147,22 @@ extension CodexReviewStore { package func requestActiveReviewCancellationsForRuntimeStop( reason: ReviewCancellation = .system(message: "Review runtime stopped.") - ) async -> [String] { + ) async throws -> [String] { let activeJobIDs = orderedJobs .filter { $0.isTerminal == false } .map(\.id) + var firstError: (any Error)? for jobID in activeJobIDs { - _ = try? await cancelReview(jobID: jobID, cancellation: reason) + do { + _ = try await cancelReview(jobID: jobID, cancellation: reason) + } catch { + if firstError == nil { + firstError = error + } + } + } + if let firstError { + throw firstError } return activeJobIDs } diff --git a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift index 7879b63c..b7584ad5 100644 --- a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift +++ b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift @@ -12,59 +12,6 @@ private let defaultExternalURLOpener: ExternalURLOpener = { url in _ = NSWorkspace.shared.open(url) } -private actor RuntimeShutdownCleanupRace { - private var result: Bool? - private var continuation: CheckedContinuation? - - func finish(_ value: Bool) { - guard result == nil else { - return - } - result = value - continuation?.resume(returning: value) - continuation = nil - } - - func wait() async -> Bool { - if let result { - return result - } - return await withCheckedContinuation { continuation in - if let result { - continuation.resume(returning: result) - } else { - self.continuation = continuation - } - } - } -} - -private func runRuntimeShutdownCleanup( - timeout: Duration, - operation: @escaping @Sendable () async -> Void -) async -> Bool { - let race = RuntimeShutdownCleanupRace() - let operationTask = Task { - await operation() - await race.finish(true) - } - let timeoutTask = Task { - do { - try await Task.sleep(for: timeout) - } catch { - return - } - await race.finish(false) - } - let result = await race.wait() - if result { - timeoutTask.cancel() - } else { - operationTask.cancel() - } - return result -} - private struct PendingLoginRuntimeCleanup { var client: AppServerClient? var codexHomeURL: URL? @@ -127,6 +74,7 @@ public extension CodexReviewStore { shutdownCleanupTimeout: Duration = .seconds(2), networkMonitor: any CodexReviewNetworkMonitoring = SystemCodexReviewNetworkMonitor(), networkRecoveryPolicy: CodexReviewNetworkRecoveryPolicy = .default, + reviewRuntimeClosePolicy: ReviewRuntimeClosePolicy = .production, transport: any JSONRPC.Transport ) -> CodexReviewStore { makeLiveStoreForTesting( @@ -140,6 +88,7 @@ public extension CodexReviewStore { shutdownCleanupTimeout: shutdownCleanupTimeout, networkMonitor: networkMonitor, networkRecoveryPolicy: networkRecoveryPolicy, + reviewRuntimeClosePolicy: reviewRuntimeClosePolicy, transportFactory: { _ in transport } ) } @@ -159,6 +108,7 @@ public extension CodexReviewStore { shutdownCleanupTimeout: Duration = .seconds(2), networkMonitor: any CodexReviewNetworkMonitoring = SystemCodexReviewNetworkMonitor(), networkRecoveryPolicy: CodexReviewNetworkRecoveryPolicy = .default, + reviewRuntimeClosePolicy: ReviewRuntimeClosePolicy = .production, transportFactory: @escaping @MainActor @Sendable (URL) async throws -> any JSONRPC.Transport ) -> CodexReviewStore { CodexReviewStore( @@ -181,7 +131,8 @@ public extension CodexReviewStore { } ), networkMonitor: networkMonitor, - networkRecoveryPolicy: networkRecoveryPolicy + networkRecoveryPolicy: networkRecoveryPolicy, + reviewRuntimeClosePolicy: reviewRuntimeClosePolicy ) } } @@ -471,29 +422,20 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private func cancelActiveReviewsForRuntimeTeardown( store: CodexReviewStore, - appServerBackend: AppServerCodexReviewBackend, reason: ReviewCancellation, timeoutWarning: String ) async { - store.recordActiveReviewCancellationRequestsForRuntimeStop(reason: reason) - let didInterrupt = await runRuntimeShutdownCleanup(timeout: shutdownCleanupTimeout) { - do { - try await appServerBackend.interruptActiveReviewsForShutdown( - reason: .init(message: reason.message) - ) - } catch { - logger.error("Failed to interrupt active reviews during runtime teardown: \(error.localizedDescription, privacy: .public)") - } + var cancellationFailure: (any Error)? + do { + _ = try await store.requestActiveReviewCancellationsForRuntimeStop(reason: reason) + } catch { + cancellationFailure = error + logger.error("Failed to cancel active reviews during runtime teardown: \(error.localizedDescription, privacy: .public)") } - let locallyCancelledJobIDs = store.cancelActiveReviewsLocallyForRuntimeStop( - reason: reason, - cancelWorkers: false - ) - store.cancelAndDetachReviewWorkersForRuntimeStop(jobIDs: locallyCancelledJobIDs) let didDrainReviewWorkers = await store.drainReviewWorkersForRuntimeStop( timeout: shutdownCleanupTimeout ) - if didInterrupt == false || didDrainReviewWorkers == false { + if cancellationFailure != nil || didDrainReviewWorkers == false { logger.warning("\(timeoutWarning, privacy: .public)") } } @@ -508,11 +450,10 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { return } logger.info("Stopping review runtime") - if let appServerBackend { + if appServerBackend != nil { let reason = ReviewCancellation.system(message: "Review runtime stopped.") await cancelActiveReviewsForRuntimeTeardown( store: store, - appServerBackend: appServerBackend, reason: reason, timeoutWarning: "Timed out cleaning active reviews before stopping runtime" ) @@ -656,6 +597,12 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { guard auth.persistedAccounts.contains(where: { $0.accountKey == accountKey }) else { return } + let runtimeStore = appServerBackend == nil ? nil : attachedStore + if let runtimeStore { + await runtimeStore.closeActiveReviewSessions( + reason: .system(message: "Account switched.") + ) + } try CodexReviewAccountRegistry.activateAccount( accountKey, accounts: auth.persistedAccounts, @@ -667,17 +614,22 @@ 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 runtimeStore else { return } - await attachedStore.closeActiveReviewSessions(reason: .system(message: "Account switched.")) - await stop(store: attachedStore) - await start(store: attachedStore, forceRestartIfNeeded: true) + await stop(store: runtimeStore) + await start(store: runtimeStore, forceRestartIfNeeded: true) } func removeAccount(auth: CodexReviewAuthModel, accountKey: String) async throws { let removedActiveAccount = auth.selectedAccount?.accountKey == accountKey || auth.persistedActiveAccountKey == accountKey + let runtimeStore = removedActiveAccount && appServerBackend != nil ? attachedStore : nil + if let runtimeStore { + await runtimeStore.closeActiveReviewSessions( + reason: .system(message: "Account removed.") + ) + } if removedActiveAccount, let appServerBackend { _ = try? await appServerBackend.logout(.init(accountKey)) } @@ -704,12 +656,11 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { if removedActiveAccount { auth.selectPersistedAccount(nil) auth.updatePhase(.signedOut) - guard let attachedStore, appServerBackend != nil else { + guard let runtimeStore else { return } - await attachedStore.closeActiveReviewSessions(reason: .system(message: "Account removed.")) - await stop(store: attachedStore) - await start(store: attachedStore, forceRestartIfNeeded: true) + await stop(store: runtimeStore) + await start(store: runtimeStore, forceRestartIfNeeded: true) } } @@ -1189,11 +1140,10 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { return } let message = "Review runtime stopped unexpectedly: \(error.localizedDescription)" - if let appServerBackend { + if appServerBackend != nil { let reason = ReviewCancellation.system(message: message) await cancelActiveReviewsForRuntimeTeardown( store: store, - appServerBackend: appServerBackend, reason: reason, timeoutWarning: "Timed out cleaning active reviews after runtime failure" ) diff --git a/Sources/CodexReviewTesting/TestSupport.swift b/Sources/CodexReviewTesting/TestSupport.swift index f61aee48..abf47221 100644 --- a/Sources/CodexReviewTesting/TestSupport.swift +++ b/Sources/CodexReviewTesting/TestSupport.swift @@ -962,6 +962,7 @@ package actor FakeJSONRPCTransport: JSONRPC.Transport { private var maxActiveByMethod: [String: Int] = [:] private var gatesByMethod: [String: RequestGate] = [:] private var oneShotGatesByMethod: [String: [RequestGate]] = [:] + private var activeRequestGates: [Int: RequestGate] = [:] private var requestCountWaiters: [(Int, CheckedContinuation)] = [] private var requestMethodWaiters: [String: [CheckedContinuation]] = [:] private var responseMethodWaiters: [String: [CheckedContinuation]] = [:] @@ -970,6 +971,7 @@ package actor FakeJSONRPCTransport: JSONRPC.Transport { private var closed = false private var closeFailure: ReviewRuntimeCloseFailure? private var closeCallCount = 0 + private var closeWaiters: [CheckedContinuation] = [] package init(responses: [String: [Data]] = [:]) { self.responses = responses @@ -1021,7 +1023,9 @@ package actor FakeJSONRPCTransport: JSONRPC.Transport { ) let queuedResponse = dequeueResponse(for: request.method) if let gate = dequeueOneShotGate(for: request.method) ?? gatesByMethod[request.method] { + activeRequestGates[request.id] = gate await gate.wait() + activeRequestGates.removeValue(forKey: request.id) } activeByMethod[request.method, default: 1] -= 1 guard closed == false else { @@ -1071,8 +1075,15 @@ package actor FakeJSONRPCTransport: JSONRPC.Transport { package func close() async throws { closeCallCount += 1 + let closeWaiters = closeWaiters + self.closeWaiters.removeAll(keepingCapacity: false) + for waiter in closeWaiters { + waiter.resume() + } closed = true - let gates = Array(gatesByMethod.values) + oneShotGatesByMethod.values.flatMap { $0 } + let gates = Array(gatesByMethod.values) + + oneShotGatesByMethod.values.flatMap { $0 } + + activeRequestGates.values for gate in gates { await gate.gate.open() } @@ -1164,6 +1175,19 @@ package actor FakeJSONRPCTransport: JSONRPC.Transport { closeCallCount } + package func waitForCloseCall() async { + if closeCallCount > 0 { + return + } + await withCheckedContinuation { continuation in + if closeCallCount > 0 { + continuation.resume() + } else { + closeWaiters.append(continuation) + } + } + } + package func maxActiveCount(for method: String) -> Int { maxActiveByMethod[method] ?? 0 } diff --git a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift index b18df189..a1f2f289 100644 --- a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift +++ b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift @@ -116,6 +116,13 @@ private func interruptAndDeliverCanonicalTerminal( ) } +private func controlledReviewClosePolicy(gate: AsyncGate) -> ReviewRuntimeClosePolicy { + ReviewRuntimeClosePolicy(terminalGrace: .seconds(10)) { _ in + await gate.wait() + try Task.checkCancellation() + } +} + @Suite("app-server client") struct AppServerClientTests { @Test func processTransportConfigurationResolvesCodexFromProvidedPath() throws { @@ -1061,6 +1068,165 @@ struct AppServerClientTests { #expect(object["sandbox"] == nil) } + @Test func backendAdmissionRefusesThreadStartBeforeDispatch() async throws { + let transport = FakeJSONRPCTransport() + let initializeGate = AsyncGate() + await transport.holdNextIgnoringCancellation(method: "initialize", gate: initializeGate) + try await enqueueInitialize(transport) + let backend = AppServerCodexReviewBackend(client: .init(transport: transport)) + let admission = ReviewStartAdmission( + closePolicy: controlledReviewClosePolicy(gate: AsyncGate()) + ) + let startTask = await admission.start { admission in + try await backend.startReview( + .init( + jobID: "job-1", + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ), + admission: admission + ) + } + await transport.waitForRequest(method: "initialize") + + let cancellation = Task { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in Issue.record("Pre-dispatch cancellation sent interrupt.") }, + forceClose: { Issue.record("Pre-dispatch cancellation force-closed connection.") } + ) + } + await initializeGate.open() + + #expect(try await cancellation.value.terminal == .localCancellation( + .mcpClient(message: "Stop") + )) + await #expect(throws: ReviewStartCancelledBeforeDispatch.self) { + try await startTask.value + } + #expect(await transport.recordedRequests().map(\.method) == ["initialize"]) + } + + @Test func backendAdmissionCleansPreparedThreadBeforeRefusingReviewDispatch() async throws { + let transport = FakeJSONRPCTransport() + try await enqueueInitialize(transport) + let threadGate = AsyncGate() + await transport.holdNextIgnoringCancellation(method: "thread/start", gate: threadGate) + try await transport.enqueue( + AppServerAPI.Thread.Start.Response(threadID: "thread-1", model: "gpt-5"), + for: "thread/start" + ) + try await transport.enqueue( + AppServerAPI.Thread.Unsubscribe.Response(status: .unsubscribed), + for: "thread/unsubscribe" + ) + let backend = AppServerCodexReviewBackend(client: .init(transport: transport)) + let admission = ReviewStartAdmission( + closePolicy: controlledReviewClosePolicy(gate: AsyncGate()) + ) + let startTask = await admission.start { admission in + try await backend.startReview( + .init( + jobID: "job-1", + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ), + admission: admission + ) + } + await transport.waitForRequest(method: "thread/start") + + let cancellation = Task { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in Issue.record("Thread-only cancellation sent empty-turn interrupt.") }, + forceClose: { Issue.record("Prepared-thread cancellation force-closed connection.") } + ) + } + await threadGate.open() + + #expect(try await cancellation.value.terminal == .localCancellation( + .mcpClient(message: "Stop") + )) + await #expect(throws: ReviewStartCancelledBeforeDispatch.self) { + try await startTask.value + } + let methods = await transport.recordedRequests().map(\.method) + #expect(methods.contains("review/start") == false) + #expect(Array(methods.suffix(3)) == [ + "thread/backgroundTerminals/clean", + "thread/unsubscribe", + "thread/delete", + ]) + } + + @Test func backendAdmissionJoinsOutcomeUnknownReviewStartBeforeInterrupt() async throws { + let transport = FakeJSONRPCTransport() + try await enqueueInitialize(transport) + try await transport.enqueue( + AppServerAPI.Thread.Start.Response(threadID: "thread-1", model: "gpt-5"), + for: "thread/start" + ) + let reviewGate = AsyncGate() + await transport.holdNextIgnoringCancellation(method: "review/start", gate: reviewGate) + try await transport.enqueue( + AppServerAPI.Review.Start.Response(turnID: "turn-1", reviewThreadID: "thread-1"), + for: "review/start" + ) + try await transport.enqueue(EmptyResponse(), for: "turn/interrupt") + let backend = AppServerCodexReviewBackend(client: .init(transport: transport)) + let admission = ReviewStartAdmission( + closePolicy: controlledReviewClosePolicy(gate: AsyncGate()) + ) + let startTask = await admission.start { admission in + try await backend.startReview( + .init( + jobID: "job-1", + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ), + admission: admission + ) + } + await transport.waitForRequest(method: "review/start") + + let cancellation = Task { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { run, reason in + try await backend.interruptReview(run, reason: reason) + }, + forceClose: { + try await backend.forceCloseReviewConnection() + } + ) + } + await reviewGate.open() + let attempt = try await startTask.value + await transport.waitForResponseDelivery(method: "turn/interrupt") + try await transport.emitServerNotification( + method: "turn/completed", + params: TestTurnNotification( + threadID: "thread-1", + turn: .init(id: "turn-1", status: "interrupted", error: .init(message: "Stop")) + ) + ) + try await admission.recordCanonicalTerminal( + .interrupted(.requested(.mcpClient(message: "Stop"))), + for: attempt.run + ) + let resolution = try await cancellation.value + + #expect(resolution.terminal == .canonical( + run: attempt.run, + terminal: .interrupted(.requested(.mcpClient(message: "Stop"))) + )) + let methods = await transport.recordedRequests().map(\.method) + let reviewIndex = try #require(methods.firstIndex(of: "review/start")) + let interruptIndex = try #require(methods.firstIndex(of: "turn/interrupt")) + #expect(reviewIndex < interruptIndex) + } + @Test func backendUsesLegacySandboxWhenProcessDoesNotSupportModernSessionSource() async throws { let transport = FakeJSONRPCTransport() try await enqueueInitialize(transport) diff --git a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift index f30af30f..73718737 100644 --- a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift +++ b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift @@ -1004,6 +1004,10 @@ struct CodexReviewHostTests { try await firstTransport.enqueue(AppServerAPI.Thread.Start.Response(threadID: "thread-first", model: "gpt-5"), for: "thread/start") try await firstTransport.enqueue(AppServerAPI.Review.Start.Response(turnID: "turn-first"), for: "review/start") try await firstTransport.enqueue(EmptyResponse(), for: "turn/interrupt") + try await firstTransport.enqueue( + AppServerAPI.Thread.Unsubscribe.Response(status: .unsubscribed), + for: "thread/unsubscribe" + ) let secondTransport = FakeJSONRPCTransport() try await secondTransport.enqueue(AppServerAPI.Initialize.Response(), for: "initialize") @@ -1041,7 +1045,20 @@ struct CodexReviewHostTests { ) await waitUntil { store.jobs.first?.core.run.turnID == "turn-first" } - try await store.switchAccount(CodexAccount(email: "second@example.com")) + let switchTask = Task { @MainActor in + try await store.switchAccount(CodexAccount(email: "second@example.com")) + } + await firstTransport.waitForResponseDelivery(method: "turn/interrupt") + try await firstTransport.emitServerNotification( + method: "turn/completed", + params: HostTurnNotification( + threadID: "thread-first", + turnID: "turn-first", + status: "interrupted", + errorMessage: "Account switched." + ) + ) + try await switchTask.value let result = try await reviewRead await secondTransport.waitForRequestCount(2) await firstTransport.waitForRequestCount(8) @@ -1084,6 +1101,10 @@ struct CodexReviewHostTests { try await firstTransport.enqueue(AppServerAPI.Thread.Start.Response(threadID: "thread-active", model: "gpt-5"), for: "thread/start") try await firstTransport.enqueue(AppServerAPI.Review.Start.Response(turnID: "turn-active"), for: "review/start") try await firstTransport.enqueue(EmptyResponse(), for: "turn/interrupt") + try await firstTransport.enqueue( + AppServerAPI.Thread.Unsubscribe.Response(status: .unsubscribed), + for: "thread/unsubscribe" + ) try await firstTransport.enqueue(EmptyResponse(), for: "account/logout") let secondTransport = FakeJSONRPCTransport() @@ -1112,7 +1133,20 @@ struct CodexReviewHostTests { ) await waitUntil { store.jobs.first?.core.run.turnID == "turn-active" } - await store.logout() + let logoutTask = Task { @MainActor in + await store.logout() + } + await firstTransport.waitForResponseDelivery(method: "turn/interrupt") + try await firstTransport.emitServerNotification( + method: "turn/completed", + params: HostTurnNotification( + threadID: "thread-active", + turnID: "turn-active", + status: "interrupted", + errorMessage: "Signed out." + ) + ) + await logoutTask.value let result = try await reviewRead await secondTransport.waitForRequestCount(2) @@ -1189,6 +1223,10 @@ struct CodexReviewHostTests { try await transport.enqueue(AppServerAPI.Thread.Start.Response(threadID: "thread-1", model: "gpt-5"), for: "thread/start") try await transport.enqueue(AppServerAPI.Review.Start.Response(turnID: "turn-1"), for: "review/start") try await transport.enqueue(EmptyResponse(), for: "turn/interrupt") + try await transport.enqueue( + AppServerAPI.Thread.Unsubscribe.Response(status: .unsubscribed), + for: "thread/unsubscribe" + ) let store = CodexReviewStore.makeLiveStoreForTesting( environment: ["HOME": homeURL.path], webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, @@ -1225,6 +1263,16 @@ struct CodexReviewHostTests { #expect(jobBeforeInterruptCompletes.cancellationRequested) #expect(jobBeforeInterruptCompletes.core.lifecycle.cancellation?.message == "Review runtime stopped.") await interruptGate.open() + await transport.waitForResponseDelivery(method: "turn/interrupt") + try await transport.emitServerNotification( + method: "turn/completed", + params: HostTurnNotification( + threadID: "thread-1", + turnID: "turn-1", + status: "interrupted", + errorMessage: "Review runtime stopped." + ) + ) await stopTask.value let result = try await reviewRead @@ -1241,11 +1289,13 @@ struct CodexReviewHostTests { #expect(interruptIndex < deleteIndex) } - @Test func liveStoreStopBoundsStuckReviewCancellationCleanup() async throws { + @Test func liveStoreStopGraceForceClosesAndAwaitsConnectionTerminal() async throws { let homeURL = try temporaryHome() let interruptGate = AsyncGate() + let graceStarted = AsyncGate() + let graceGate = AsyncGate() let transport = FakeJSONRPCTransport() - await transport.holdNext(method: "turn/interrupt", gate: interruptGate) + await transport.holdNextIgnoringCancellation(method: "turn/interrupt", gate: interruptGate) try await transport.enqueue(AppServerAPI.Initialize.Response(), for: "initialize") try await transport.enqueue(AppServerAPI.Account.Read.Response(), for: "account/read") try await transport.enqueue( @@ -1258,7 +1308,14 @@ struct CodexReviewHostTests { let store = CodexReviewStore.makeLiveStoreForTesting( environment: ["HOME": homeURL.path], webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, - shutdownCleanupTimeout: .milliseconds(20), + reviewRuntimeClosePolicy: .init( + terminalGrace: .seconds(10), + sleep: { _ in + await graceStarted.open() + await graceGate.wait() + try Task.checkCancellation() + } + ), transport: transport ) @@ -1269,18 +1326,30 @@ struct CodexReviewHostTests { request: .init(cwd: "/tmp/project", target: .uncommittedChanges) ) } - await waitUntil { store.jobs.first?.core.run.turnID == "turn-1" } + await transport.waitForRequest(method: "review/start") + let jobID = try #require(store.jobs.first?.id) + let admission = try #require(store.reviewStartAdmissions[jobID]) + let run = try #require(await admission.waitForActiveRun()) + #expect(run.turnID == "turn-1") + let worker = try #require(store.reviewWorkerTasks[jobID]) - let startedAt = Date() - await store.stop() - let elapsed = Date().timeIntervalSince(startedAt) - let resultBeforeRemoteCleanupUnblocked = try await waitForTaskValue(reviewRead, timeout: .seconds(1)) - await interruptGate.open() - let result = try #require(resultBeforeRemoteCleanupUnblocked) + let stopFinished = CompletionFlag() + let stopTask = Task { @MainActor in + await store.stop() + await stopFinished.complete() + } + await graceStarted.wait() + #expect(await stopFinished.isCompleted() == false) + await graceGate.open() + await transport.waitForCloseCall() + await worker.value + await stopTask.value + let result = try await reviewRead.value - #expect(elapsed < 1) - #expect(result.core.lifecycle.status == .cancelled) - #expect(await transport.recordedRequests().map(\.method).contains("turn/interrupt")) + #expect(result.core.lifecycle.status == .failed) + #expect(result.core.lifecycle.terminal?.kind == .interrupted) + #expect(await transport.closeCallCountForTesting() >= 1) + #expect(await stopFinished.isCompleted()) } @Test func liveStoreStopDrainsRecoveryWaitingWorkerCleanupBeforeDroppingBackend() async throws { @@ -2058,3 +2127,31 @@ private actor CompletionFlag { completed } } + +private struct HostTurnNotification: Encodable, Sendable { + struct Turn: Encodable, Sendable { + struct TurnError: Encodable, Sendable { + var message: String + } + + var id: String + var items: [String] + var itemsView: String + var status: String + var error: TurnError + } + + var threadId: String + var turn: Turn + + init(threadID: String, turnID: String, status: String, errorMessage: String) { + self.threadId = threadID + self.turn = Turn( + id: turnID, + items: [], + itemsView: "notLoaded", + status: status, + error: .init(message: errorMessage) + ) + } +} From 7e62c3fd84bc135caa30c1d0f08b809dcb08e138 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 17:25:22 +0900 Subject: [PATCH 03/38] fix(app-server): retain cleanup failures --- .../AppServerCodexReviewBackend.swift | 26 ++++++++++++------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift index 9ec67e3b..f0e9ed2f 100644 --- a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift +++ b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift @@ -467,36 +467,42 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { completedReviewEventSessionMetricsByThreadID[threadID] = metrics } } - var firstFailure: ReviewRuntimeCloseFailure? + var failureMessages: [String] = [] do { let _: EmptyResponse = try await client.send(AppServerAPI.Thread.BackgroundTerminals.Clean.Request( params: .init(threadID: run.threadID) )) } catch { - firstFailure = .cleanup("thread/backgroundTerminals/clean for \(run.threadID): \(error.localizedDescription)") + failureMessages.append( + "thread/backgroundTerminals/clean for \(run.threadID): \(error.localizedDescription)" + ) } do { let _: AppServerAPI.Thread.Unsubscribe.Response = try await client.send(AppServerAPI.Thread.Unsubscribe.Request( params: .init(threadID: run.threadID) )) - } catch where firstFailure == nil { - firstFailure = .cleanup("thread/unsubscribe for \(run.threadID): \(error.localizedDescription)") - } catch {} + } catch { + failureMessages.append( + "thread/unsubscribe for \(run.threadID): \(error.localizedDescription)" + ) + } for threadID in cleanupThreadIDs { do { let _: EmptyResponse = try await client.send(AppServerAPI.Thread.Delete.Request( params: .init(threadID: threadID) )) - } catch where firstFailure == nil { - firstFailure = .cleanup("thread/delete for \(threadID): \(error.localizedDescription)") - } catch {} + } catch { + failureMessages.append( + "thread/delete for \(threadID): \(error.localizedDescription)" + ) + } } for threadID in cleanupThreadIDs { reviewEventSessionCanonicalThreadIDByThreadID.removeValue(forKey: threadID) } reviewThreadIDsForCleanupByThreadID.removeValue(forKey: run.threadID) - if let firstFailure { - throw firstFailure + if failureMessages.isEmpty == false { + throw ReviewRuntimeCloseFailure.cleanup(failureMessages.joined(separator: "; ")) } } From b82b23462ebfd6770dbc37a6cee379202ee8c14a Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 17:33:11 +0900 Subject: [PATCH 04/38] test(review): drive cancellation through attempts --- .../CodexReviewMCPHTTPServerTests.swift | 67 +++++++++++-------- Tests/ReviewUITests/ReviewUITests.swift | 14 +++- 2 files changed, 51 insertions(+), 30 deletions(-) diff --git a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift index 84f6c18e..ee561073 100644 --- a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift +++ b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift @@ -756,23 +756,20 @@ struct CodexReviewMCPHTTPServerTests { } @Test func streamableHTTPCancelsReviewByTransportScopedSelector() async throws { - let backend = FakeCodexReviewBackend() + let run = CodexReviewBackendModel.Review.Run( + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "thread-1", + model: "gpt-5" + ) + let backend = FakeCodexReviewBackend(nextRun: run) let store = CodexReviewStore.makeTestingStore( - backend: TestingCodexReviewStoreBackend(reviewBackend: backend) + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-running" }) ) try await withHTTPServer(store: store) { server in let sessionID = try await initializeSession(endpoint: await server.url) - let running = CodexReviewJob.makeForTesting( - id: "job-running", - sessionID: sessionID, - cwd: "/tmp/project", - targetSummary: "Uncommitted changes", - threadID: "thread-1", - turnID: "turn-1", - status: .running, - summary: "Running" - ) let otherSession = CodexReviewJob.makeForTesting( id: "job-other-session", sessionID: "other-session", @@ -786,26 +783,40 @@ struct CodexReviewMCPHTTPServerTests { store.loadForTesting( serverState: .running, workspaces: [.init(cwd: "/tmp/project")], - jobs: [running, otherSession] + jobs: [otherSession] ) - let response = try await postJSONRPC( - endpoint: await server.url, + async let started = store.startReview( sessionID: sessionID, - body: [ - "jsonrpc": "2.0", - "id": 2, - "method": "tools/call", - "params": [ - "name": "review_cancel", - "arguments": [ - "sessionID": "other-session", - "cwd": "/tmp/project", - "statuses": ["running"], - "reason": "Stop from MCP", - ], + request: .init(cwd: "/tmp/project", target: .uncommittedChanges), + waitTimeout: .zero + ) + await backend.waitForStartReview() + _ = try await started + let running = try #require(store.job(id: "job-running")) + + let endpoint = await server.url + let requestBody = try makeJSONBody([ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": [ + "name": "review_cancel", + "arguments": [ + "sessionID": "other-session", + "cwd": "/tmp/project", + "statuses": ["running"], + "reason": "Stop from MCP", ], - ] + ], + ]) + async let responseData = postJSONRPCData( + endpoint: endpoint, + sessionID: sessionID, + bodyData: requestBody ) + await backend.waitForInterruptReview() + await backend.yield(.cancelled("Stop from MCP"), for: run) + let response = try decodeSSEJSON(from: try await responseData) #expect(response.value(for: ["result", "structuredContent", "jobId"]) as? String == "job-running") #expect(response.value(for: ["result", "structuredContent", "cancelled"]) as? Bool == true) diff --git a/Tests/ReviewUITests/ReviewUITests.swift b/Tests/ReviewUITests/ReviewUITests.swift index dba17b32..f79a971c 100644 --- a/Tests/ReviewUITests/ReviewUITests.swift +++ b/Tests/ReviewUITests/ReviewUITests.swift @@ -1779,7 +1779,7 @@ struct ReviewUITests { #expect(job.core.lifecycle.endedAt != nil) } - @Test func cancellationFailureUpdatesJobErrorState() async { + @Test func cancellationFailureUpdatesJobErrorState() async throws { let job = makeJob( id: "job-running", cwd: "/tmp/workspace-alpha", @@ -1792,6 +1792,14 @@ struct ReviewUITests { serverState: .running, content: makeSidebarContent(from: [job]) ) + let run = CodexReviewBackendModel.Review.Run( + threadID: try #require(job.core.run.threadID), + turnID: job.core.run.turnID, + reviewThreadID: job.core.run.reviewThreadID + ) + let admission = ReviewStartAdmission() + await admission.recordActiveRun(run) + store.reviewStartAdmissions[job.id] = admission let viewController = ReviewMonitorSplitViewController(store: store, uiState: ReviewMonitorUIState(auth: store.auth)) viewController.loadViewIfNeeded() @@ -6787,7 +6795,9 @@ final class FailingCancellationBackend: PreviewCodexReviewStoreBackend { override func waitUntilStopped() async {} override func interruptReview(_: CodexReviewBackendModel.Review.Run, reason _: CodexReviewBackendModel.CancellationReason) async throws { - throw CodexReviewAPI.Error.io("Cancellation failed.") + throw ReviewInterruptRequestFailure( + outcome: .rejected(code: nil, message: "Cancellation failed.") + ) } } From f20d8eb6e1e70c337905c5e1cc7b170cd860e367 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 17:35:09 +0900 Subject: [PATCH 05/38] test(review): expect typed EOF interruption --- Tests/CodexReviewTests/ReviewTerminalContractTests.swift | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/Tests/CodexReviewTests/ReviewTerminalContractTests.swift b/Tests/CodexReviewTests/ReviewTerminalContractTests.swift index c19c216d..17aa1b09 100644 --- a/Tests/CodexReviewTests/ReviewTerminalContractTests.swift +++ b/Tests/CodexReviewTests/ReviewTerminalContractTests.swift @@ -180,7 +180,7 @@ struct ReviewTerminalContractTests { #expect(result.core.reviewText != "arbitrary last message") } - @Test func streamEOFFailsWithoutAnAuthoritativeTerminal() async throws { + @Test func streamEOFUsesTypedTransportInterruptionWithoutCanonicalTerminal() async throws { let (store, backend) = makeStore() async let started = store.startReview( sessionID: "session-1", @@ -191,7 +191,12 @@ struct ReviewTerminalContractTests { let result = try await started #expect(result.core.lifecycle.status == .failed) - #expect(result.core.lifecycle.terminal?.kind == .failed) + #expect(result.core.lifecycle.terminal?.kind == .interrupted) + guard case .interrupted(.transport(let message)) = result.core.lifecycle.terminal else { + Issue.record("Expected a typed transport interruption.") + return + } + #expect(message.contains("authoritative terminal")) #expect(result.core.lifecycle.errorMessage?.contains("authoritative terminal") == true) } From 8cc1b8b556edce2413351ce6af68d1f9708b2ff0 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 17:47:53 +0900 Subject: [PATCH 06/38] fix(review): permit cancellation retry after rejection --- .../CodexReview/ReviewAttemptProcessor.swift | 18 ++++++++++++++++-- .../ReviewAttemptProcessorTests.swift | 19 +++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/Sources/CodexReview/ReviewAttemptProcessor.swift b/Sources/CodexReview/ReviewAttemptProcessor.swift index a031ffac..897c8d86 100644 --- a/Sources/CodexReview/ReviewAttemptProcessor.swift +++ b/Sources/CodexReview/ReviewAttemptProcessor.swift @@ -438,13 +438,27 @@ package actor ReviewStartAdmission { case .failure(let error): if let terminal { phase = .terminal(terminal) - } else if case .active = phase { - requestedCancellation = nil + } else if case .active = phase, + let requestFailure = error as? ReviewInterruptRequestFailure, + case .rejected = requestFailure.outcome { + resetRejectedCancellationForRetry() } throw error } } + private func resetRejectedCancellationForRetry() { + requestedCancellation = nil + cancellationTask = nil + interruptRequestTask = nil + terminalBarrierTask = nil + graceTask = nil + forceCloseTask = nil + requestResult = nil + forceCloseResult = nil + cancellationResult = nil + } + private func finishCancellationAfterTerminal( _ terminal: ReviewAttemptBarrierTerminal ) async throws -> ReviewAttemptCancellationResolution { diff --git a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift index f901d417..36529687 100644 --- a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift +++ b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift @@ -76,6 +76,25 @@ struct ReviewAttemptProcessorTests { #expect(await admission.currentPhase() == .active(run)) #expect(await admission.cancellationRequest() == nil) + + let retryRequested = InvocationProbe() + let retry = Task { + try await admission.cancel( + .mcpClient(message: "Stop again"), + interrupt: { _, _ in await retryRequested.record() }, + forceClose: {} + ) + } + await retryRequested.waitForInvocation() + try await admission.recordCanonicalTerminal( + .interrupted(.requested(.mcpClient(message: "Stop again"))), + for: run + ) + + #expect(try await retry.value.terminal == .canonical( + run: run, + terminal: .interrupted(.requested(.mcpClient(message: "Stop again"))) + )) } @Test func rejectionAfterTerminalCannotRewriteTerminal() async throws { From d5d89c4f669c93b75bee3f4c20afc7bf842d7487 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 17:54:33 +0900 Subject: [PATCH 07/38] test(review): gate duplicate cancellation owner --- Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift index 36529687..91cafeb2 100644 --- a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift +++ b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift @@ -270,6 +270,7 @@ struct ReviewAttemptProcessorTests { forceClose: {} ) } + await requestStarted.waitForInvocation() let second = Task { try await admission.cancel( .mcpClient(message: "Stop"), @@ -279,7 +280,6 @@ struct ReviewAttemptProcessorTests { forceClose: {} ) } - await requestStarted.waitForInvocation() try await admission.recordCanonicalTerminal( .interrupted(.requested(.mcpClient(message: "Stop"))), for: run From 0c6c28e31421fcba05e0f3e54a2bcadf44aa128d Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 17:56:57 +0900 Subject: [PATCH 08/38] test(review): gate duplicate cleanup owner --- .../ReviewAttemptProcessorTests.swift | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift index 91cafeb2..60e86e16 100644 --- a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift +++ b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift @@ -491,16 +491,20 @@ struct ReviewAttemptProcessorTests { let cleanupStarted = InvocationProbe() let cleanupGate = AsyncGate() - async let first: Void = admission.cleanup(run: run) { - await cleanupStarted.record() - await cleanupGate.waitIgnoringCancellation() - } - async let second: Void = admission.cleanup(run: run) { - Issue.record("Duplicate cleanup caller installed a second cleanup Task.") + let first = Task { + try await admission.cleanup(run: run) { + await cleanupStarted.record() + await cleanupGate.waitIgnoringCancellation() + } } await cleanupStarted.waitForInvocation() + let second = Task { + try await admission.cleanup(run: run) { + Issue.record("Duplicate cleanup caller installed a second cleanup Task.") + } + } await cleanupGate.open() - _ = try await (first, second) + _ = try await (first.value, second.value) #expect(await cleanupStarted.invocationCount() == 1) } From 466f12bfaf97c635396080f2d334f563927e0975 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:08:11 +0900 Subject: [PATCH 09/38] fix(review): drain startup after connection terminal --- .../CodexReview/ReviewAttemptProcessor.swift | 3 + .../ReviewAttemptProcessorTests.swift | 74 +++++++++++++++++++ 2 files changed, 77 insertions(+) diff --git a/Sources/CodexReview/ReviewAttemptProcessor.swift b/Sources/CodexReview/ReviewAttemptProcessor.swift index 897c8d86..d20ff7b6 100644 --- a/Sources/CodexReview/ReviewAttemptProcessor.swift +++ b/Sources/CodexReview/ReviewAttemptProcessor.swift @@ -260,6 +260,9 @@ package actor ReviewStartAdmission { guard terminal == nil else { return } + if Self.isOutcomeUnknownStartPhase(phase) { + startTask?.cancel() + } receiveTerminal(.connection(failure)) } diff --git a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift index 60e86e16..b5887557 100644 --- a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift +++ b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift @@ -411,6 +411,42 @@ struct ReviewAttemptProcessorTests { } } + @Test func cancellationAfterOutcomeUnknownThreadDispatchDrainsThroughForcedConnectionTerminal() async throws { + let graceGate = AsyncGate() + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: graceGate)) + let threadDispatched = InvocationProbe() + let threadResponseGate = AsyncGate() + let forceClose = InvocationProbe() + let connection = ReviewRuntimeCloseFailure.connection("Forced close") + let startTask = await admission.start { admission in + #expect(await admission.admitThreadStartDispatch()) + await threadDispatched.record() + await threadResponseGate.wait() + try Task.checkCancellation() + Issue.record("Thread request outlived its typed connection terminal.") + return .init(run: canonicalRun) + } + await threadDispatched.waitForInvocation() + + let cancellation = Task { + try await admission.cancel( + .system(message: "Stop"), + interrupt: { _, _ in Issue.record("Thread-only attempt interrupted a turn.") }, + forceClose: { + await forceClose.record() + await admission.recordConnectionTerminal(connection) + } + ) + } + await graceGate.open() + await forceClose.waitForInvocation() + + #expect(try await cancellation.value.terminal == .connection(connection)) + await #expect(throws: CancellationError.self) { + try await startTask.value + } + } + @Test func cancellationAfterThreadResponseRefusesNotSentReviewDispatch() async throws { let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) let prepared = InvocationProbe() @@ -486,6 +522,44 @@ struct ReviewAttemptProcessorTests { )) } + @Test func cancellationAfterOutcomeUnknownReviewDispatchDrainsThroughForcedConnectionTerminal() async throws { + let graceGate = AsyncGate() + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: graceGate)) + let reviewDispatched = InvocationProbe() + let reviewResponseGate = AsyncGate() + let forceClose = InvocationProbe() + let connection = ReviewRuntimeCloseFailure.connection("Forced close") + let startTask = await admission.start { admission in + #expect(await admission.admitThreadStartDispatch()) + await admission.recordPreparedThread(provisionalRun) + #expect(await admission.admitReviewStartDispatch(for: provisionalRun)) + await reviewDispatched.record() + await reviewResponseGate.wait() + try Task.checkCancellation() + Issue.record("Review request outlived its typed connection terminal.") + return .init(run: canonicalRun) + } + await reviewDispatched.waitForInvocation() + + let cancellation = Task { + try await admission.cancel( + .system(message: "Stop"), + interrupt: { _, _ in Issue.record("Unresolved review request interrupted a turn.") }, + forceClose: { + await forceClose.record() + await admission.recordConnectionTerminal(connection) + } + ) + } + await graceGate.open() + await forceClose.waitForInvocation() + + #expect(try await cancellation.value.terminal == .connection(connection)) + await #expect(throws: CancellationError.self) { + try await startTask.value + } + } + @Test func duplicateCleanupCallersJoinOneOwnedTask() async throws { let (admission, run) = try await makeActiveAdmission() let cleanupStarted = InvocationProbe() From 9c6b33e24e9897a9bfc1595f549177157fb71e1e Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:12:38 +0900 Subject: [PATCH 10/38] test(review): await startup cancellation admission --- .../CodexReview/ReviewAttemptProcessor.swift | 31 +++++++++++++++++++ .../ReviewAttemptProcessorTests.swift | 4 +++ 2 files changed, 35 insertions(+) diff --git a/Sources/CodexReview/ReviewAttemptProcessor.swift b/Sources/CodexReview/ReviewAttemptProcessor.swift index d20ff7b6..09c0a46e 100644 --- a/Sources/CodexReview/ReviewAttemptProcessor.swift +++ b/Sources/CodexReview/ReviewAttemptProcessor.swift @@ -164,6 +164,7 @@ package actor ReviewStartAdmission { private var cancellationResult: Result? private var terminalWaiters: [UUID: CheckedContinuation] = [:] private var activeRunWaiters: [CheckedContinuation] = [] + private var cancellationAdmissionWaiters: [CheckedContinuation] = [] private var cancellationWaiters: [CheckedContinuation, Never>] = [] package init(closePolicy: ReviewRuntimeClosePolicy = .production) { @@ -281,6 +282,7 @@ package actor ReviewStartAdmission { return .init(terminal: terminal) } requestedCancellation = cancellation + resumeCancellationAdmissionWaiters(returning: cancellation) let task = Task { try await self.performCancellation( cancellation, @@ -330,6 +332,24 @@ package actor ReviewStartAdmission { requestedCancellation } + package func waitForCancellationAdmission() async -> ReviewCancellation? { + if let requestedCancellation { + return requestedCancellation + } + if terminal != nil { + return nil + } + return await withCheckedContinuation { continuation in + if let requestedCancellation { + continuation.resume(returning: requestedCancellation) + } else if terminal != nil { + continuation.resume(returning: nil) + } else { + cancellationAdmissionWaiters.append(continuation) + } + } + } + package func recordedCleanupResult( for run: CodexReviewBackendModel.Review.Run ) async -> Result? { @@ -583,6 +603,7 @@ package actor ReviewStartAdmission { waiter.resume(returning: terminal) } resumeActiveRunWaiters(returning: nil) + resumeCancellationAdmissionWaiters(returning: nil) resolveCancellationIfPossible() } @@ -596,6 +617,16 @@ package actor ReviewStartAdmission { } } + private func resumeCancellationAdmissionWaiters( + returning cancellation: ReviewCancellation? + ) { + let waiters = cancellationAdmissionWaiters + cancellationAdmissionWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume(returning: cancellation) + } + } + private func resolveCancellationIfPossible() { guard cancellationResult == nil else { return diff --git a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift index b5887557..3f2d0353 100644 --- a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift +++ b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift @@ -403,6 +403,7 @@ struct ReviewAttemptProcessorTests { forceClose: {} ) } + #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) await threadResponseGate.open() #expect(try await cancellation.value.terminal == .localCancellation(.mcpClient(message: "Stop"))) @@ -438,6 +439,7 @@ struct ReviewAttemptProcessorTests { } ) } + #expect(await admission.waitForCancellationAdmission() == .system(message: "Stop")) await graceGate.open() await forceClose.waitForInvocation() @@ -473,6 +475,7 @@ struct ReviewAttemptProcessorTests { forceClose: {} ) } + #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) await reviewDispatchGate.open() #expect(try await cancellation.value.terminal == .localCancellation(.mcpClient(message: "Stop"))) @@ -551,6 +554,7 @@ struct ReviewAttemptProcessorTests { } ) } + #expect(await admission.waitForCancellationAdmission() == .system(message: "Stop")) await graceGate.open() await forceClose.waitForInvocation() From 24a6a46f551f52a9440cc79b56f955145c8828ba Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:38:24 +0900 Subject: [PATCH 11/38] fix(review): serialize start dispatch admission --- .../CodexReview/ReviewAttemptProcessor.swift | 38 +++++++++-- .../AppServerCodexReviewBackend.swift | 3 + .../ReviewAttemptProcessorTests.swift | 63 +++++++++++++++++++ 3 files changed, 99 insertions(+), 5 deletions(-) diff --git a/Sources/CodexReview/ReviewAttemptProcessor.swift b/Sources/CodexReview/ReviewAttemptProcessor.swift index 09c0a46e..6c603be9 100644 --- a/Sources/CodexReview/ReviewAttemptProcessor.swift +++ b/Sources/CodexReview/ReviewAttemptProcessor.swift @@ -159,6 +159,7 @@ package actor ReviewStartAdmission { private var forceCloseTask: Task? private var cleanupTasksByAttemptID: [String: Task] = [:] private var terminal: ReviewAttemptBarrierTerminal? + private var startFailed = false private var requestResult: Result? private var forceCloseResult: Result? private var cancellationResult: Result? @@ -198,14 +199,32 @@ package actor ReviewStartAdmission { return false } switch phase { - case .preparingThread: + case .preparingThread(.notSent): phase = .preparingThread(.outcomeUnknown) return true - case .queued, .startingReview, .active, .interrupting, .finishing, .terminal: + case .queued, .preparingThread(.outcomeUnknown), .startingReview, + .active, .interrupting, .finishing, .terminal: return false } } + package func recordThreadStartRejectedForRetry() throws { + if let requestedCancellation { + throw ReviewStartCancelledBeforeDispatch(cancellation: requestedCancellation) + } + guard terminal == nil else { + throw ReviewAttemptContractFailure( + message: "Thread start retry cannot follow an attempt terminal." + ) + } + guard case .preparingThread(.outcomeUnknown) = phase else { + throw ReviewAttemptContractFailure( + message: "Thread start retry requires one rejected dispatched request." + ) + } + phase = .preparingThread(.notSent) + } + package func recordPreparedThread(_ run: CodexReviewBackendModel.Review.Run) { guard terminal == nil else { return @@ -219,7 +238,7 @@ package actor ReviewStartAdmission { guard requestedCancellation == nil else { return false } - guard case .startingReview(let currentRun, _) = phase, + guard case .startingReview(let currentRun, .notSent) = phase, currentRun.attemptID == preparedRun.attemptID else { return false @@ -317,10 +336,13 @@ package actor ReviewStartAdmission { if terminal != nil { return nil } + if startFailed { + return nil + } return await withCheckedContinuation { continuation in if let activeRun { continuation.resume(returning: activeRun) - } else if terminal != nil { + } else if terminal != nil || startFailed { continuation.resume(returning: nil) } else { activeRunWaiters.append(continuation) @@ -339,10 +361,13 @@ package actor ReviewStartAdmission { if terminal != nil { return nil } + if startFailed { + return nil + } return await withCheckedContinuation { continuation in if let requestedCancellation { continuation.resume(returning: requestedCancellation) - } else if terminal != nil { + } else if terminal != nil || startFailed { continuation.resume(returning: nil) } else { cancellationAdmissionWaiters.append(continuation) @@ -369,6 +394,7 @@ package actor ReviewStartAdmission { resumeActiveRunWaiters(returning: attempt.run) } case .failure(let error): + startFailed = true if let cancellation = (error as? ReviewStartCancelledBeforeDispatch)?.cancellation { receiveTerminal(.localCancellation(cancellation)) } else if error is CancellationError, @@ -376,6 +402,8 @@ package actor ReviewStartAdmission { case .preparingThread(.notSent) = phase { receiveTerminal(.localCancellation(requestedCancellation)) } + resumeActiveRunWaiters(returning: nil) + resumeCancellationAdmissionWaiters(returning: nil) } } diff --git a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift index f0e9ed2f..24fd2cd9 100644 --- a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift +++ b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift @@ -210,12 +210,14 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } catch let error as JSONRPC.Error where Self.shouldRetryThreadStartWithLegacySandbox(error) { // Deprecated compatibility: some builds accept the permissions field shape // without registering the danger-full-access built-in profile. + try await admission.recordThreadStartRejectedForRetry() return try await sendThreadStart(AppServerAPI.Thread.Start.Request( params: threadStartParamsWithLegacySandbox(request) ), admission: admission) } catch let error as JSONRPC.Error where Self.shouldRetryThreadStartWithObjectPermissions(error) { // Deprecated compatibility: installed Codex builds can require object-shaped // permissions while the latest local app-server source accepts a profile ID string. + try await admission.recordThreadStartRejectedForRetry() return try await startReviewThreadWithProfileSelectionPermissions(request, admission: admission) } } @@ -248,6 +250,7 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { { // Deprecated compatibility: installed Codex builds can know the permissions // object shape without registering the danger-full-access built-in profile. + try await admission.recordThreadStartRejectedForRetry() return try await sendThreadStart(AppServerAPI.Thread.Start.Request( params: threadStartParamsWithLegacySandbox(request) ), admission: admission) diff --git a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift index 3f2d0353..9744ac4f 100644 --- a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift +++ b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift @@ -345,6 +345,69 @@ struct ReviewAttemptProcessorTests { #expect(resolution.terminal == .localCancellation(.mcpClient(message: "Stop"))) } + @Test func threadStartDispatchAdmissionRejectsDirectDuplicate() async throws { + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) + let startGate = AsyncGate() + let startTask = await admission.start { _ in + await startGate.waitIgnoringCancellation() + return .init(run: canonicalRun) + } + + #expect(await admission.admitThreadStartDispatch()) + #expect(await admission.admitThreadStartDispatch() == false) + + await startGate.open() + _ = try await startTask.value + } + + @Test func threadStartDispatchAdmissionAllowsVerifiedRejectionRetry() async throws { + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) + let startGate = AsyncGate() + let startTask = await admission.start { _ in + await startGate.waitIgnoringCancellation() + return .init(run: canonicalRun) + } + + #expect(await admission.admitThreadStartDispatch()) + try await admission.recordThreadStartRejectedForRetry() + #expect(await admission.admitThreadStartDispatch()) + #expect(await admission.admitThreadStartDispatch() == false) + + await startGate.open() + _ = try await startTask.value + } + + @Test func reviewStartDispatchAdmissionRejectsDirectDuplicate() async throws { + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) + let startGate = AsyncGate() + let startTask = await admission.start { _ in + await startGate.waitIgnoringCancellation() + return .init(run: canonicalRun) + } + #expect(await admission.admitThreadStartDispatch()) + await admission.recordPreparedThread(provisionalRun) + + #expect(await admission.admitReviewStartDispatch(for: provisionalRun)) + #expect(await admission.admitReviewStartDispatch(for: provisionalRun) == false) + + await startGate.open() + _ = try await startTask.value + } + + @Test func generalStartFailureEndsAdmissionWaiters() async throws { + let failure = ReviewAttemptContractFailure(message: "Start failed") + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) + let startTask = await admission.start { _ in + throw failure + } + + await #expect(throws: failure) { + try await startTask.value + } + #expect(await admission.waitForActiveRun() == nil) + #expect(await admission.waitForCancellationAdmission() == nil) + } + @Test func cancellationBeforeThreadDispatchRefusesWrite() async throws { let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) let entered = InvocationProbe() From cf28075be94cd7bed353938443b3aee6a77cca2b Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:41:14 +0900 Subject: [PATCH 12/38] fix(review): surface force-close failure after terminal --- .../CodexReview/ReviewAttemptProcessor.swift | 12 ++++- .../ReviewAttemptProcessorTests.swift | 47 +++++++++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/Sources/CodexReview/ReviewAttemptProcessor.swift b/Sources/CodexReview/ReviewAttemptProcessor.swift index 6c603be9..a9dad2c7 100644 --- a/Sources/CodexReview/ReviewAttemptProcessor.swift +++ b/Sources/CodexReview/ReviewAttemptProcessor.swift @@ -162,6 +162,7 @@ package actor ReviewStartAdmission { private var startFailed = false private var requestResult: Result? private var forceCloseResult: Result? + private var graceDidExpire = false private var cancellationResult: Result? private var terminalWaiters: [UUID: CheckedContinuation] = [:] private var activeRunWaiters: [CheckedContinuation] = [] @@ -507,6 +508,7 @@ package actor ReviewStartAdmission { forceCloseTask = nil requestResult = nil forceCloseResult = nil + graceDidExpire = false cancellationResult = nil } @@ -615,7 +617,7 @@ package actor ReviewStartAdmission { self.terminal = terminal } case .graceExpired: - break + graceDidExpire = true case .forceClose(let result): forceCloseResult = result } @@ -659,7 +661,10 @@ package actor ReviewStartAdmission { guard cancellationResult == nil else { return } - if case .failure(let closeFailure)? = forceCloseResult, terminal == nil { + if graceDidExpire, forceCloseResult == nil { + return + } + if case .failure(let closeFailure)? = forceCloseResult { resolveCancellation(.failure(closeFailure)) return } @@ -746,6 +751,9 @@ package actor ReviewStartAdmission { if terminal == nil { terminalBarrierTask?.cancel() } + if case .failure? = forceCloseResult { + interruptRequestTask?.cancel() + } graceTask?.cancel() await interruptRequestTask?.value await terminalBarrierTask?.value diff --git a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift index 9744ac4f..7b4edb62 100644 --- a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift +++ b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift @@ -255,6 +255,53 @@ struct ReviewAttemptProcessorTests { } } + @Test func terminalFirstForceCloseFailureCancelsPendingRequestAndPreservesTerminal() async throws { + let graceGate = AsyncGate() + let (admission, run) = try await makeActiveAdmission( + closePolicy: controlledClosePolicy(gate: graceGate) + ) + let requestStarted = InvocationProbe() + let requestGate = AsyncGate() + let requestCancelled = InvocationProbe() + let forceCloseStarted = InvocationProbe() + let forceFailure = ReviewRuntimeCloseFailure.process("Process remained alive") + + let cancellation = Task { + try await admission.cancel( + .system(message: "Stop"), + interrupt: { _, _ in + await requestStarted.record() + do { + await requestGate.wait() + try Task.checkCancellation() + } catch { + if error is CancellationError { + await requestCancelled.record() + } + throw error + } + }, + forceClose: { + await forceCloseStarted.record() + throw forceFailure + } + ) + } + await requestStarted.waitForInvocation() + try await admission.recordCanonicalTerminal(.completed, for: run) + await graceGate.open() + await forceCloseStarted.waitForInvocation() + + await #expect(throws: forceFailure) { + try await cancellation.value + } + #expect(await requestCancelled.invocationCount() == 1) + #expect(await admission.currentPhase() == .terminal(.canonical( + run: run, + terminal: .completed + ))) + } + @Test func duplicateCancellationCallersJoinOneRequest() async throws { let (admission, run) = try await makeActiveAdmission() let requestStarted = InvocationProbe() From 58f2cb1e614fd2d7a1dbc02bbc27ee9ca080c43e Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:45:51 +0900 Subject: [PATCH 13/38] fix(app-server): join repeated transport close --- .../AppServerProcessTransport.swift | 43 +++++++--- .../AppServerClientTests.swift | 85 +++++++++++++++++++ 2 files changed, 118 insertions(+), 10 deletions(-) diff --git a/Sources/CodexReviewAppServer/AppServerProcessTransport.swift b/Sources/CodexReviewAppServer/AppServerProcessTransport.swift index 1e0f6a00..6d8ac9c4 100644 --- a/Sources/CodexReviewAppServer/AppServerProcessTransport.swift +++ b/Sources/CodexReviewAppServer/AppServerProcessTransport.swift @@ -61,6 +61,7 @@ package actor AppServerProcessTransport: JSONRPC.Transport { private let stderr: Pipe private let stdoutEvents: AppServerPipeReadEventSource private let stderrEvents: AppServerPipeReadEventSource + private let closeCompletionForTesting: (@Sendable () async throws -> Void)? private var framer = JSONRPC.Framer() private var pending: [Int: PendingResponse] = [:] private var notificationContinuations: [UUID: AsyncThrowingStream.Continuation] = [:] @@ -68,8 +69,12 @@ package actor AppServerProcessTransport: JSONRPC.Transport { private var stdoutReaderTask: Task? = nil private var stderrReaderTask: Task? = nil private var closed = false + private var closeTask: Task? - package init(configuration: Configuration = .init()) throws { + package init( + configuration: Configuration = .init(), + closeCompletionForTesting: (@Sendable () async throws -> Void)? = nil + ) throws { guard FileManager.default.isExecutableFile(atPath: configuration.executable) else { throw AppServerProcessTransportError.executableNotFound( command: configuration.executable, @@ -100,6 +105,7 @@ package actor AppServerProcessTransport: JSONRPC.Transport { ) self.stdoutEvents = stdoutEvents self.stderrEvents = stderrEvents + self.closeCompletionForTesting = closeCompletionForTesting logger.info("Launching codex app-server: \(configuration.executable, privacy: .public) \(configuration.arguments.joined(separator: " "), privacy: .public)") logger.info("Using codex app-server home: \(configuration.codexHomeURL.path, privacy: .public)") logger.info("codex app-server launched with pid \(process.processIdentifier, privacy: .public)") @@ -170,27 +176,44 @@ package actor AppServerProcessTransport: JSONRPC.Transport { error: any Error, readerTask: ReaderTask? ) async throws { - if closed { - if readerTask == nil { - await waitForReaderTasks(excluding: nil) + let task: Task + if let closeTask { + task = closeTask + } else { + closed = true + stdoutEvents.cancel() + stderrEvents.cancel() + try? stdin.fileHandleForWriting.close() + let newTask = Task { + try await self.performCloseTransport( + terminateProcess: terminateProcess, + error: error + ) } - return + closeTask = newTask + task = newTask } - closed = true - stdoutEvents.cancel() - stderrEvents.cancel() - try? stdin.fileHandleForWriting.close() + + let result = await task.result + await waitForReaderTasks(excluding: readerTask) + try result.get() + } + + private func performCloseTransport( + terminateProcess: Bool, + error: any Error + ) async throws { var processCloseError: (any Error)? if terminateProcess { logger.info("Terminating codex app-server pid \(self.process.processIdentifier, privacy: .public)") do { try await process.terminateAndWait() + try await closeCompletionForTesting?() } catch { processCloseError = error } } finishAll(throwing: processCloseError ?? error) - await waitForReaderTasks(excluding: readerTask) if let processCloseError { throw processCloseError } diff --git a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift index a1f2f289..c0fa5273 100644 --- a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift +++ b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift @@ -387,6 +387,79 @@ struct AppServerClientTests { } } + @Test func concurrentProcessTransportCloseCallersJoinOneFailure() async throws { + let directory = FileManager.default.temporaryDirectory + .appending(path: "codex-review-concurrent-close-\(UUID().uuidString)") + defer { + try? FileManager.default.removeItem(at: directory) + } + let closeStarted = AsyncGate() + let closeGate = AsyncGate() + let closeCompletions = CallCounter() + let failure = ReviewRuntimeCloseFailure.process("Injected close failure") + let transport = try AppServerProcessTransport( + configuration: .init( + executable: "/bin/cat", + arguments: [], + environment: [ + "HOME": directory.path, + "PATH": "/bin:/usr/bin", + ] + ), + closeCompletionForTesting: { + await closeCompletions.record() + await closeStarted.open() + await closeGate.waitIgnoringCancellation() + throw failure + } + ) + + let first = Task { try await transport.close() } + await closeStarted.wait() + let second = Task { try await transport.close() } + await closeGate.open() + + await #expect(throws: failure) { + try await first.value + } + await #expect(throws: failure) { + try await second.value + } + #expect(await closeCompletions.value() == 1) + } + + @Test func repeatedProcessTransportCloseRethrowsRecordedFailure() async throws { + let directory = FileManager.default.temporaryDirectory + .appending(path: "codex-review-repeated-close-\(UUID().uuidString)") + defer { + try? FileManager.default.removeItem(at: directory) + } + let closeCompletions = CallCounter() + let failure = ReviewRuntimeCloseFailure.process("Injected close failure") + let transport = try AppServerProcessTransport( + configuration: .init( + executable: "/bin/cat", + arguments: [], + environment: [ + "HOME": directory.path, + "PATH": "/bin:/usr/bin", + ] + ), + closeCompletionForTesting: { + await closeCompletions.record() + throw failure + } + ) + + await #expect(throws: failure) { + try await transport.close() + } + await #expect(throws: failure) { + try await transport.close() + } + #expect(await closeCompletions.value() == 1) + } + @Test func processTransportProcessesChunkedStdoutBeforeEOF() async throws { let directory = FileManager.default.temporaryDirectory .appending(path: "codex-review-stdout-order-\(UUID().uuidString)") @@ -5917,3 +5990,15 @@ private struct TestErrorNotification: Encodable, Sendable { try container.encode(willRetry, forKey: .willRetry) } } + +private actor CallCounter { + private var count = 0 + + func record() { + count += 1 + } + + func value() -> Int { + count + } +} From c98319756dd91a4c4ce9da2eaddadfa728318870 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 19:36:11 +0900 Subject: [PATCH 14/38] fix(review): await recovery attempt barrier --- Sources/CodexReview/CodexReviewBackend.swift | 6 +- .../CodexReview/ReviewAttemptProcessor.swift | 102 ++++- .../Store/CodexReviewStoreBackend.swift | 6 +- .../Store/CodexReviewStoreReviews.swift | 418 +++++++++++++----- .../PreviewCodexReviewStoreBackend.swift | 6 +- .../AppServerCodexReviewBackend.swift | 113 +---- .../AppServerReviewControl.swift | 13 +- Sources/CodexReviewHost/CodexReviewHost.swift | 10 +- .../LiveCodexReviewStoreBackend.swift | 14 +- Sources/CodexReviewTesting/TestSupport.swift | 125 +++++- .../AppServerClientTests.swift | 193 ++++++-- .../CodexReviewHostTests.swift | 13 + .../CodexReviewStoreCommandTests.swift | 310 +++++++++++-- .../ReviewAttemptProcessorTests.swift | 91 ++++ 14 files changed, 1075 insertions(+), 345 deletions(-) diff --git a/Sources/CodexReview/CodexReviewBackend.swift b/Sources/CodexReview/CodexReviewBackend.swift index 82710137..887b8212 100644 --- a/Sources/CodexReview/CodexReviewBackend.swift +++ b/Sources/CodexReview/CodexReviewBackend.swift @@ -17,12 +17,12 @@ package protocol CodexReviewBackend: Sendable { func interruptReview(_ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason) async throws func forceCloseReviewConnection() async throws func beginReviewRecovery( - _ run: CodexReviewBackendModel.Review.Run, - reason: CodexReviewBackendModel.CancellationReason + _ barrier: ReviewAttemptRecoveryBarrier ) async throws -> CodexReviewBackendModel.Review.RecoveryToken func resumeReviewRecovery( _ token: CodexReviewBackendModel.Review.RecoveryToken, - request: CodexReviewBackendModel.Review.Start + request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws } diff --git a/Sources/CodexReview/ReviewAttemptProcessor.swift b/Sources/CodexReview/ReviewAttemptProcessor.swift index a9dad2c7..112b22d8 100644 --- a/Sources/CodexReview/ReviewAttemptProcessor.swift +++ b/Sources/CodexReview/ReviewAttemptProcessor.swift @@ -119,6 +119,37 @@ package struct ReviewAttemptCancellationResolution: Equatable, Sendable { } } +package struct ReviewAttemptRecoveryBarrier: Equatable, Sendable { + package var run: CodexReviewBackendModel.Review.Run + package var terminal: ReviewAttemptBarrierTerminal + package var cancellation: ReviewCancellation + package var requestFailure: ReviewInterruptRequestFailure? + + fileprivate init( + run: CodexReviewBackendModel.Review.Run, + terminal: ReviewAttemptBarrierTerminal, + cancellation: ReviewCancellation, + requestFailure: ReviewInterruptRequestFailure? + ) { + self.run = run + self.terminal = terminal + self.cancellation = cancellation + self.requestFailure = requestFailure + } +} + +package struct ReviewRecoverySupersededByTerminal: LocalizedError, Equatable, Sendable { + package var terminal: ReviewAttemptBarrierTerminal + + package init(terminal: ReviewAttemptBarrierTerminal) { + self.terminal = terminal + } + + package var errorDescription: String? { + "Recovery was superseded by \(terminal.diagnosticDescription)." + } +} + /// Owns one review attempt from the first dispatch admission through terminal and cleanup. /// Store cancellation and backend request dispatch both consult this actor; there is no /// call-site startup-cancellation mirror. @@ -159,6 +190,7 @@ package actor ReviewStartAdmission { private var forceCloseTask: Task? private var cleanupTasksByAttemptID: [String: Task] = [:] private var terminal: ReviewAttemptBarrierTerminal? + private var registeredRun: CodexReviewBackendModel.Review.Run? private var startFailed = false private var requestResult: Result? private var forceCloseResult: Result? @@ -252,6 +284,7 @@ package actor ReviewStartAdmission { guard terminal == nil else { return } + registeredRun = run phase = .active(run) resumeActiveRunWaiters(returning: run) } @@ -294,6 +327,65 @@ package actor ReviewStartAdmission { CodexReviewBackendModel.CancellationReason ) async throws -> Void, forceClose: @escaping @Sendable () async throws -> Void + ) async throws -> ReviewAttemptCancellationResolution { + let resolution = try await joinedCancellationResolution( + cancellation, + interrupt: interrupt, + forceClose: forceClose + ) + if let requestFailure = resolution.requestFailure, + case .outcomeUnknown = requestFailure.outcome, + case .connection(let connectionFailure) = resolution.terminal { + throw ReviewInterruptRequestFailure( + outcome: requestFailure.outcome, + secondaryBarrierDiagnostic: connectionFailure.localizedDescription + ) + } + return resolution + } + + package func interruptForRecovery( + _ cancellation: ReviewCancellation, + interrupt: @escaping @Sendable ( + CodexReviewBackendModel.Review.Run, + CodexReviewBackendModel.CancellationReason + ) async throws -> Void, + forceClose: @escaping @Sendable () async throws -> Void + ) async throws -> ReviewAttemptRecoveryBarrier { + guard let run = activeRun ?? canonicalRunForTerminal ?? registeredRun else { + throw ReviewAttemptContractFailure( + message: "Recovery interruption requires one canonical review run." + ) + } + let resolution = try await joinedCancellationResolution( + cancellation, + interrupt: interrupt, + forceClose: forceClose + ) + switch resolution.terminal { + case .canonical(_, let terminal) where terminal.kind != .interrupted: + throw ReviewRecoverySupersededByTerminal(terminal: resolution.terminal) + case .canonical, .connection: + return .init( + run: run, + terminal: resolution.terminal, + cancellation: cancellation, + requestFailure: resolution.requestFailure + ) + case .localCancellation: + throw ReviewAttemptContractFailure( + message: "A dispatched review recovery cannot complete locally." + ) + } + } + + private func joinedCancellationResolution( + _ cancellation: ReviewCancellation, + interrupt: @escaping @Sendable ( + CodexReviewBackendModel.Review.Run, + CodexReviewBackendModel.CancellationReason + ) async throws -> Void, + forceClose: @escaping @Sendable () async throws -> Void ) async throws -> ReviewAttemptCancellationResolution { if let cancellationTask { return try await cancellationTask.value @@ -390,6 +482,7 @@ package actor ReviewStartAdmission { ) { switch result { case .success(let attempt): + registeredRun = attempt.run if terminal == nil { phase = .active(attempt.run) resumeActiveRunWaiters(returning: attempt.run) @@ -698,9 +791,12 @@ package actor ReviewStartAdmission { } switch terminal { case .connection(let connectionFailure): - resolveCancellation(.failure(ReviewInterruptRequestFailure( - outcome: requestFailure.outcome, - secondaryBarrierDiagnostic: connectionFailure.localizedDescription + resolveCancellation(.success(.init( + terminal: terminal, + requestFailure: ReviewInterruptRequestFailure( + outcome: requestFailure.outcome, + secondaryBarrierDiagnostic: connectionFailure.localizedDescription + ) ))) case .canonical, .localCancellation: resolveCancellation(.success(.init( diff --git a/Sources/CodexReview/Store/CodexReviewStoreBackend.swift b/Sources/CodexReview/Store/CodexReviewStoreBackend.swift index 12631c80..5ce3b8ec 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreBackend.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreBackend.swift @@ -51,12 +51,12 @@ package protocol CodexReviewStoreBackend: CodexReviewSettingsBackend, Sendable { func interruptReview(_ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason) async throws func forceCloseReviewConnection() async throws func beginReviewRecovery( - _ run: CodexReviewBackendModel.Review.Run, - reason: CodexReviewBackendModel.CancellationReason + _ barrier: ReviewAttemptRecoveryBarrier ) async throws -> CodexReviewBackendModel.Review.RecoveryToken func resumeReviewRecovery( _ token: CodexReviewBackendModel.Review.RecoveryToken, - request: CodexReviewBackendModel.Review.Start + request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws } diff --git a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift index c79d4480..d417485d 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift @@ -138,15 +138,15 @@ extension CodexReviewStore { activeRuns.removeValue(forKey: jobID) reviewRecoveryWaitingJobIDs.remove(jobID) } else { - let currentRun = try await consumeReviewEvents( + let completion = try await consumeReviewEvents( for: backendAttempt, job: job, startRequest: startRequest, admission: admission ) - run = currentRun + run = completion.run do { - try await cleanupReview(currentRun, admission: admission) + try await cleanupReview(completion.run, admission: completion.admission) } catch { retainCleanupFailure(error, for: jobID) } @@ -163,12 +163,13 @@ extension CodexReviewStore { } } catch let error where error is CancellationError || Task.isCancelled { if let cleanupRun = activeRuns[jobID] ?? run { + let cleanupAdmission = reviewStartAdmissions[jobID] ?? admission let failure = ReviewRuntimeCloseFailure.worker( "Review worker was cancelled before a canonical terminal." ) - await admission.recordConnectionTerminal(failure) + await cleanupAdmission.recordConnectionTerminal(failure) do { - try await cleanupReview(cleanupRun, admission: admission) + try await cleanupReview(cleanupRun, admission: cleanupAdmission) } catch { retainCleanupFailure(error, for: jobID) } @@ -182,8 +183,9 @@ extension CodexReviewStore { reviewRecoveryWaitingJobIDs.remove(jobID) } catch { if let cleanupRun = activeRuns[jobID] ?? run { + let cleanupAdmission = reviewStartAdmissions[jobID] ?? admission do { - try await cleanupReview(cleanupRun, admission: admission) + try await cleanupReview(cleanupRun, admission: cleanupAdmission) } catch { retainCleanupFailure(error, for: jobID) } @@ -193,7 +195,8 @@ extension CodexReviewStore { if job.isTerminal == false, let transportFailure = error as? ReviewWorkerInputQueueError { let failure = ReviewRuntimeCloseFailure.connection(transportFailure.message) - await admission.recordConnectionTerminal(failure) + let currentAdmission = reviewStartAdmissions[jobID] ?? admission + await currentAdmission.recordConnectionTerminal(failure) markReviewInterrupted( job, cause: .transport(message: transportFailure.message) @@ -259,6 +262,10 @@ extension CodexReviewStore { appendRecoveryProgress(networkRecoveryUnavailableMessage, to: job) } + private func recordReviewRecoveryBarrier(for jobID: String) { + reviewRecoveryWaitingJobIDs.insert(jobID) + } + private func reviewWorkerInputs(for attempt: BackendReviewAttempt) async -> ReviewWorkerInputs { let networkMonitor = self.networkMonitor let policy = self.networkRecoveryPolicy @@ -271,6 +278,7 @@ extension CodexReviewStore { queue: queue ) let eventSource = ReviewWorkerEventSource(queue: queue) + let recoveryInterruptionSource = ReviewWorkerRecoveryInterruptionSource(queue: queue) let networkTask = Task { for await snapshot in snapshots { await signalCoordinator.observe(snapshot) @@ -281,6 +289,7 @@ extension CodexReviewStore { queue: queue, networkStatusTracker: tracker, eventSource: eventSource, + recoveryInterruptionSource: recoveryInterruptionSource, initialEventSubscriptionID: initialEventSubscriptionID, networkTask: networkTask, signalCoordinator: signalCoordinator @@ -596,21 +605,48 @@ extension CodexReviewStore { job: CodexReviewJob, startRequest: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission - ) async throws -> CodexReviewBackendModel.Review.Run { + ) async throws -> ReviewWorkerAttemptCompletion { let inputs = await reviewWorkerInputs(for: initialAttempt) - defer { - inputs.cancel() + do { + let completion = try await consumeReviewEventLoop( + for: initialAttempt, + job: job, + startRequest: startRequest, + initialAdmission: admission, + inputs: inputs + ) + await inputs.cancel() + return completion + } catch { + if error is CancellationError || Task.isCancelled { + let currentAdmission = reviewStartAdmissions[job.id] ?? admission + await currentAdmission.recordConnectionTerminal(.worker( + "Review worker stopped while recovery interruption was pending." + )) + } + await inputs.cancel() + throw error } + } + + private func consumeReviewEventLoop( + for initialAttempt: BackendReviewAttempt, + job: CodexReviewJob, + startRequest: CodexReviewBackendModel.Review.Start, + initialAdmission: ReviewStartAdmission, + inputs: ReviewWorkerInputs + ) async throws -> ReviewWorkerAttemptCompletion { + var admission = initialAdmission var recoveryState = ReviewNetworkRecoveryLoopState(currentRun: initialAttempt.run) var activeEventSubscriptionID: Int? = inputs.initialEventSubscriptionID while let input = await inputs.next() { if job.isTerminal { - return recoveryState.currentRun + return .init(run: recoveryState.currentRun, admission: admission) } switch input { case .reviewEvent(let event): guard activeEventSubscriptionID == event.subscriptionID, - recoveryState.shouldConsumeEvent(from: event.subscriptionRun) + recoveryState.shouldRouteAttemptInput(from: event.subscriptionRun) else { continue } @@ -620,16 +656,31 @@ extension CodexReviewStore { for: recoveryState.currentRun ) } + if recoveryState.isInterruptingForNetworkRecovery { + guard event.event.supersedesNetworkRecovery else { + continue + } + } recoveryState.currentRun = handleReviewEvent( event.event, job: job, currentRun: recoveryState.currentRun ) if job.isTerminal { - return recoveryState.currentRun + return .init(run: recoveryState.currentRun, admission: admission) } case .reviewEventsFinished(let finishedRun): - guard activeEventSubscriptionID == finishedRun.subscriptionID else { + guard activeEventSubscriptionID == finishedRun.subscriptionID, + recoveryState.shouldRouteAttemptInput(from: finishedRun.run) + else { + continue + } + if recoveryState.isInterruptingForNetworkRecovery { + let failure = ReviewRuntimeCloseFailure.connection( + ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription + ) + await admission.recordConnectionTerminal(failure) + activeEventSubscriptionID = nil continue } if recoveryState.shouldIgnoreFinishedEvent(for: finishedRun.run) { @@ -640,17 +691,25 @@ extension CodexReviewStore { isWaitingForNetworkRecovery: recoveryState.isWaitingForNetworkRecovery, admission: admission ) { - return recoveryState.currentRun + return .init(run: recoveryState.currentRun, admission: admission) } case .reviewEventsFailed(let failedRun): guard activeEventSubscriptionID == failedRun.subscriptionID, - recoveryState.shouldConsumeEvent(from: failedRun.run) + recoveryState.shouldRouteAttemptInput(from: failedRun.run) else { continue } if failedRun.failure.isCancellation { throw CancellationError() } + if recoveryState.isInterruptingForNetworkRecovery { + let failure = ReviewRuntimeCloseFailure.connection( + failedRun.failure.message + ) + await admission.recordConnectionTerminal(failure) + activeEventSubscriptionID = nil + continue + } if await inputs.networkStatusTracker.currentStatus() != .satisfied { recoveryState.recordPendingOutageStreamFailure(failedRun.failure) activeEventSubscriptionID = nil @@ -658,6 +717,29 @@ extension CodexReviewStore { continue } try throwReviewEventStreamFailure(failedRun.failure) + case .recoveryBarrierResolved(let resolution): + guard recoveryState.isInterruptingForNetworkRecovery, + resolution.run.attemptID == recoveryState.currentRun.attemptID + else { + continue + } + switch resolution.result { + case .failure(let failure): + if failure.underlying is ReviewRecoverySupersededByTerminal { + continue + } + throw failure.underlying + case .success(let ready): + if job.isTerminal || completePendingCancellationIfNeeded(for: job) { + return .init(run: recoveryState.currentRun, admission: admission) + } + recoveryState.markWaitingForNetworkRecovery() + markReviewWaitingForNetworkRecovery(job) + recordReviewRecoveryBarrier(for: job.id) + activeEventSubscriptionID = nil + await inputs.cancelActiveEventSubscription() + recoveryState.markRecoveryReady(ready) + } case .networkSnapshot(let snapshot, let recoveryGeneration): if let pendingFailure = recoveryState.takePendingOutageStreamFailureAfterTransientRecovery( snapshot @@ -680,18 +762,17 @@ extension CodexReviewStore { job: job, startRequest: startRequest, inputs: inputs, - recoveryToken: recoveryState.recoveryToken, - admission: admission + recoveryReady: recoveryState.recoveryReady ) { case .continueWaiting: recoveryState.markWaitingForNetworkRecovery() continue case .finished: reviewRecoveryWaitingJobIDs.remove(job.id) - return recoveryState.currentRun - case .recovered(let recoveredAttempt): + return .init(run: recoveryState.currentRun, admission: admission) + case .recovered(let recoveredAttempt, let recoveredAdmission): let recoveredRun = recoveredAttempt.run - await admission.recordActiveRun(recoveredRun) + admission = recoveredAdmission applyBackendRun(recoveredRun, to: job) recoveryState.markRecovered(with: recoveredRun) reviewRecoveryWaitingJobIDs.remove(job.id) @@ -699,22 +780,40 @@ extension CodexReviewStore { } case .networkOutageConfirmed: guard recoveryState.isWaitingForNetworkRecovery == false, + recoveryState.isInterruptingForNetworkRecovery == false, job.isTerminal == false, job.cancellationRequested == false, await inputs.networkStatusTracker.currentStatus() != .satisfied else { continue } - recoveryState.markWaitingForNetworkRecovery() - markReviewWaitingForNetworkRecovery(job) - reviewRecoveryWaitingJobIDs.insert(job.id) - activeEventSubscriptionID = nil - await inputs.cancelActiveEventSubscription() - let recoveryToken = try await backend.beginReviewRecovery( - recoveryState.currentRun, - reason: recoveryState.recoveryReason + recoveryState.markInterruptingForNetworkRecovery() + let recoveryRun = recoveryState.currentRun + let recoveryAdmission = admission + let pendingFailure = recoveryState.takePendingOutageStreamFailureForConfirmedRecovery() + let recoveryCancellation = ReviewCancellation.system( + message: recoveryState.recoveryReason.message ) - recoveryState.markRecoveryToken(recoveryToken) + let backend = self.backend + await inputs.beginRecoveryInterruption(for: recoveryRun) { + let barrier = try await recoveryAdmission.interruptForRecovery( + recoveryCancellation, + interrupt: { run, reason in + try await backend.interruptReview(run, reason: reason) + }, + forceClose: { + try await backend.forceCloseReviewConnection() + } + ) + let token = try await backend.beginReviewRecovery(barrier) + return .init(barrier: barrier, token: token) + } + if let pendingFailure { + _ = await recoveryAdmission.waitForCancellationAdmission() + await recoveryAdmission.recordConnectionTerminal(.connection( + pendingFailure.message + )) + } } } @@ -728,7 +827,7 @@ extension CodexReviewStore { await admission.recordConnectionTerminal(failure) markReviewInterrupted(job, cause: .transport(message: failure.localizedDescription)) } - return recoveryState.currentRun + return .init(run: recoveryState.currentRun, admission: admission) } private func handleReviewEventsFinished( @@ -767,8 +866,7 @@ extension CodexReviewStore { job: CodexReviewJob, startRequest: CodexReviewBackendModel.Review.Start, inputs: ReviewWorkerInputs, - recoveryToken: CodexReviewBackendModel.Review.RecoveryToken?, - admission: ReviewStartAdmission + recoveryReady: ReviewWorkerRecoveryReady? ) async throws -> NetworkRestoreRestartResult { if job.isTerminal || completePendingCancellationIfNeeded(for: job) { return .finished @@ -782,73 +880,22 @@ extension CodexReviewStore { guard await inputs.networkStatusTracker.currentStatus() == .satisfied else { return .continueWaiting } - guard let recoveryToken else { + guard let recoveryReady else { return .continueWaiting } - let recoveredAttempt = try await backend.resumeReviewRecovery( - recoveryToken, - request: startRequest - ) - let recoveredRun = recoveredAttempt.run - await admission.recordActiveRun(recoveredRun) - if try await stopRecoveredRunIfJobShouldNotResume( - recoveredRun, - job: job, - admission: admission - ) { - return .finished - } - return .recovered(recoveredAttempt) - } - - private func stopRecoveredRunIfJobShouldNotResume( - _ recoveredRun: CodexReviewBackendModel.Review.Run, - job: CodexReviewJob, - admission: ReviewStartAdmission - ) async throws -> Bool { - if Task.isCancelled { - try? await backend.interruptReview( - recoveredRun, - reason: .init(message: job.core.lifecycle.cancellation?.message ?? "Cancellation requested.") - ) - try await cleanupReview(recoveredRun, admission: admission) - throw CancellationError() - } - - if job.isTerminal { - if job.core.lifecycle.status == .cancelled { - try? await backend.interruptReview( - recoveredRun, - reason: .init(message: job.core.lifecycle.cancellation?.message ?? "Cancellation requested.") - ) - } - try await cleanupReview(recoveredRun, admission: admission) - return true - } - - guard job.cancellationRequested else { - return false - } - - let cancellation = job.core.lifecycle.cancellation ?? .system() - do { - try await backend.interruptReview(recoveredRun, reason: .init(message: cancellation.message)) - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation - ) - } catch { - try await cleanupReview(recoveredRun, admission: admission) - try? recordCancellationFailure( - jobID: job.id, - sessionID: job.sessionID, - message: error.localizedDescription + let recoveredAdmission = ReviewStartAdmission(closePolicy: reviewRuntimeClosePolicy) + reviewStartAdmissions[job.id] = recoveredAdmission + reviewRecoveryWaitingJobIDs.remove(job.id) + let backend = self.backend + let startTask = await recoveredAdmission.start { admission in + try await backend.resumeReviewRecovery( + recoveryReady.token, + request: startRequest, + admission: admission ) - throw error } - try await cleanupReview(recoveredRun, admission: admission) - return true + let recoveredAttempt = try await startTask.value + return .recovered(recoveredAttempt, recoveredAdmission) } private func handleReviewEvent( @@ -1173,6 +1220,15 @@ private extension CodexReviewBackendModel.Review.Event { false } } + + var supersedesNetworkRecovery: Bool { + switch self { + case .completed, .failed: + true + case .cancelled, .started, .message, .messageDelta, .log, .logEntry: + false + } + } } private extension CodexReviewJob { @@ -1241,21 +1297,54 @@ private enum ReviewWorkerEventStreamFailure: Sendable { false } } + + var message: String { + switch self { + case .cancelled: + "Review event stream was cancelled." + case .failed(let message): + message + } + } } private enum ReviewWorkerInput: Sendable { case reviewEvent(ReviewWorkerReviewEvent) case reviewEventsFinished(ReviewWorkerEventStreamFinished) case reviewEventsFailed(ReviewWorkerEventStreamFailed) + case recoveryBarrierResolved(ReviewWorkerRecoveryBarrierResolution) case networkSnapshot(CodexReviewNetworkSnapshot, recoveryGeneration: Int) case networkOutageConfirmed case networkRecoverySettled(recoveryGeneration: Int) } +private struct ReviewWorkerRecoveryBarrierResolution: Sendable { + var run: CodexReviewBackendModel.Review.Run + var result: Result +} + +private struct ReviewWorkerRecoveryReady: Sendable { + var barrier: ReviewAttemptRecoveryBarrier + var token: CodexReviewBackendModel.Review.RecoveryToken +} + +private struct ReviewWorkerAttemptCompletion: Sendable { + var run: CodexReviewBackendModel.Review.Run + var admission: ReviewStartAdmission +} + +private struct ReviewWorkerRecoveryFailure: LocalizedError, @unchecked Sendable { + var underlying: any Error + + var errorDescription: String? { + underlying.localizedDescription + } +} + private enum NetworkRestoreRestartResult { case continueWaiting case finished - case recovered(BackendReviewAttempt) + case recovered(BackendReviewAttempt, ReviewStartAdmission) } private enum ReviewNetworkSnapshotEffect { @@ -1263,10 +1352,16 @@ private enum ReviewNetworkSnapshotEffect { case restartSettling } +private enum ReviewNetworkRecoveryPhase { + case active + case interrupting + case waiting +} + private struct ReviewNetworkRecoveryLoopState { var currentRun: CodexReviewBackendModel.Review.Run - private(set) var isWaitingForNetworkRecovery = false - private(set) var recoveryToken: CodexReviewBackendModel.Review.RecoveryToken? + private(set) var recoveryPhase = ReviewNetworkRecoveryPhase.active + private(set) var recoveryReady: ReviewWorkerRecoveryReady? private var isSettlingForNetworkRecovery = false private var recoverySettleGeneration: Int? private var pendingOutageStreamFailure: ReviewWorkerEventStreamFailure? @@ -1276,21 +1371,35 @@ private struct ReviewNetworkRecoveryLoopState { self.currentRun = currentRun } + var isInterruptingForNetworkRecovery: Bool { + recoveryPhase == .interrupting + } + + var isWaitingForNetworkRecovery: Bool { + recoveryPhase == .waiting + } + + mutating func markInterruptingForNetworkRecovery() { + recoveryPhase = .interrupting + isSettlingForNetworkRecovery = false + recoverySettleGeneration = nil + } + mutating func markWaitingForNetworkRecovery() { - isWaitingForNetworkRecovery = true + recoveryPhase = .waiting isSettlingForNetworkRecovery = false recoverySettleGeneration = nil pendingOutageStreamFailure = nil } - mutating func markRecoveryToken(_ token: CodexReviewBackendModel.Review.RecoveryToken) { - recoveryToken = token + mutating func markRecoveryReady(_ ready: ReviewWorkerRecoveryReady) { + recoveryReady = ready } mutating func markRecovered(with run: CodexReviewBackendModel.Review.Run) { currentRun = run - isWaitingForNetworkRecovery = false - recoveryToken = nil + recoveryPhase = .active + recoveryReady = nil isSettlingForNetworkRecovery = false recoverySettleGeneration = nil pendingOutageStreamFailure = nil @@ -1300,6 +1409,13 @@ private struct ReviewNetworkRecoveryLoopState { pendingOutageStreamFailure = failure } + mutating func takePendingOutageStreamFailureForConfirmedRecovery() -> ReviewWorkerEventStreamFailure? { + defer { + pendingOutageStreamFailure = nil + } + return pendingOutageStreamFailure + } + mutating func takePendingOutageStreamFailureAfterTransientRecovery( _ snapshot: CodexReviewNetworkSnapshot ) -> ReviewWorkerEventStreamFailure? { @@ -1322,11 +1438,11 @@ private struct ReviewNetworkRecoveryLoopState { isWaitingForNetworkRecovery && isSettlingForNetworkRecovery && recoverySettleGeneration == recoveryGeneration - && recoveryToken != nil + && recoveryReady != nil } - func shouldConsumeEvent(from run: CodexReviewBackendModel.Review.Run) -> Bool { - isWaitingForNetworkRecovery == false && run.attemptID == currentRun.attemptID + func shouldRouteAttemptInput(from run: CodexReviewBackendModel.Review.Run) -> Bool { + recoveryPhase != .waiting && run.attemptID == currentRun.attemptID } mutating func networkSnapshotEffect( @@ -1355,6 +1471,7 @@ private struct ReviewWorkerInputs { var queue: ReviewWorkerInputQueue var networkStatusTracker: ReviewNetworkStatusTracker var eventSource: ReviewWorkerEventSource + var recoveryInterruptionSource: ReviewWorkerRecoveryInterruptionSource var initialEventSubscriptionID: Int var networkTask: Task var signalCoordinator: ReviewNetworkSignalCoordinator @@ -1371,14 +1488,60 @@ private struct ReviewWorkerInputs { await eventSource.cancelActiveSubscription() } - func cancel() { + func beginRecoveryInterruption( + for run: CodexReviewBackendModel.Review.Run, + operation: @escaping @Sendable () async throws -> ReviewWorkerRecoveryReady + ) async { + await recoveryInterruptionSource.start(for: run, operation: operation) + } + + func cancel() async { networkTask.cancel() - Task { - await eventSource.cancel() - await signalCoordinator.cancel() - await queue.finish() + await recoveryInterruptionSource.cancel() + await eventSource.cancel() + await signalCoordinator.cancel() + await networkTask.value + await queue.finish() + } +} + +private actor ReviewWorkerRecoveryInterruptionSource { + private let queue: ReviewWorkerInputQueue + private var task: Task? + + init(queue: ReviewWorkerInputQueue) { + self.queue = queue + } + + func start( + for run: CodexReviewBackendModel.Review.Run, + operation: @escaping @Sendable () async throws -> ReviewWorkerRecoveryReady + ) { + guard task == nil else { + return + } + task = Task { + let result: Result + do { + result = .success(try await operation()) + } catch { + result = .failure(.init(underlying: error)) + } + await queue.send(.recoveryBarrierResolved(.init(run: run, result: result))) + self.finish() } } + + func cancel() async { + let task = task + task?.cancel() + await task?.value + self.task = nil + } + + private func finish() { + task = nil + } } private actor ReviewWorkerInputQueue { @@ -1506,23 +1669,32 @@ private actor ReviewWorkerEventSource { return subscriptionID } - func cancelActiveSubscription() { + func cancelActiveSubscription() async { subscriptionID += 1 activeSubscriptionID = nil - cancelEventTasks() + let tasks = cancelEventTasks() + for task in tasks { + await task.value + } } - func cancel() { + func cancel() async { subscriptionID += 1 activeSubscriptionID = nil - cancelEventTasks() + let tasks = cancelEventTasks() + for task in tasks { + await task.value + } } - private func cancelEventTasks() { - for task in eventTasks.values { + @discardableResult + private func cancelEventTasks() -> [Task] { + let tasks = Array(eventTasks.values) + for task in tasks { task.cancel() } eventTasks.removeAll(keepingCapacity: true) + return tasks } private func yieldReviewEvent( @@ -1631,11 +1803,15 @@ private actor ReviewNetworkSignalCoordinator { } } - func cancel() { - outageTask?.cancel() + func cancel() async { + let pendingOutageTask = outageTask + let pendingRecoveryTask = recoveryTask + pendingOutageTask?.cancel() outageTask = nil - recoveryTask?.cancel() + pendingRecoveryTask?.cancel() recoveryTask = nil + await pendingOutageTask?.value + await pendingRecoveryTask?.value } private func scheduleOutageConfirmationIfNeeded() { diff --git a/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift b/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift index 4ef1d95d..b6fb4f36 100644 --- a/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift +++ b/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift @@ -145,15 +145,15 @@ package class PreviewCodexReviewStoreBackend: CodexReviewStoreBackend { package func forceCloseReviewConnection() async throws {} package func beginReviewRecovery( - _: CodexReviewBackendModel.Review.Run, - reason _: CodexReviewBackendModel.CancellationReason + _: ReviewAttemptRecoveryBarrier ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { throw CodexReviewAPI.Error.io(Self.previewUnavailableMessage) } package func resumeReviewRecovery( _: CodexReviewBackendModel.Review.RecoveryToken, - request _: CodexReviewBackendModel.Review.Start + request _: CodexReviewBackendModel.Review.Start, + admission _: ReviewStartAdmission ) async throws -> BackendReviewAttempt { throw CodexReviewAPI.Error.io(Self.previewUnavailableMessage) } diff --git a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift index 24fd2cd9..2a621e08 100644 --- a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift +++ b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift @@ -351,14 +351,11 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { await session.requestCancellation(message: reason.message) do { _ = try await sendTurnInterrupt(for: run) - await session.waitForTerminalBarrier() } catch { let failure = Self.interruptRequestFailure(for: error) if case .rejected = failure.outcome { await session.clearCancellationRequest() - throw failure } - await session.waitForTerminalBarrier() throw failure } } @@ -374,14 +371,13 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } package func beginReviewRecovery( - _ run: CodexReviewBackendModel.Review.Run, - reason _: CodexReviewBackendModel.CancellationReason + _ barrier: ReviewAttemptRecoveryBarrier ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { - _ = try await client.initialize() - await ensureNotificationRouterStarted() - let interruption = try await sendTurnInterrupt(for: run) { retryInterruption in - await self.markInterruptionTurnAbandoned(retryInterruption, canonicalThreadID: run.threadID) - } + let run = barrier.run + let interruption = AppServerReviewInterruption( + threadID: appServerTurnThreadID(for: run), + turnID: run.turnID ?? "" + ) markTurnAbandoned(run.turnID) markAttemptAbandoned(run, interruption: interruption) if let session = unregisterReviewEventSession(for: run) { @@ -399,7 +395,8 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { package func resumeReviewRecovery( _ token: CodexReviewBackendModel.Review.RecoveryToken, - request: CodexReviewBackendModel.Review.Start + request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { _ = try await client.initialize() await ensureNotificationRouterStarted() @@ -423,6 +420,14 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { isRunFinalized: false ) registerReviewEventSession(session, for: provisionalRun) + await admission.recordPreparedThread(provisionalRun) + guard await admission.admitReviewStartDispatch(for: provisionalRun) else { + _ = unregisterReviewEventSession(for: provisionalRun) + await session.abandon() + throw ReviewStartCancelledBeforeDispatch( + cancellation: await admission.cancellationRequest() ?? .system() + ) + } let review: AppServerAPI.Review.Start.Response reviewStartRequestsInFlight += 1 @@ -433,6 +438,9 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } catch { reviewStartRequestsInFlight -= 1 _ = unregisterReviewEventSession(for: provisionalRun) + if let terminal = Self.connectionTerminal(for: error) { + await admission.recordConnectionTerminal(terminal) + } await session.abandon() discardUnmatchedReviewNotificationsIfIdle() throw error @@ -453,6 +461,7 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { ) await session.bufferStartupNotifications(takeUnmatchedReviewNotifications(for: recoveredRun)) await session.finalizeRun() + await admission.recordActiveRun(recoveredRun) reviewStartRequestsInFlight -= 1 discardUnmatchedReviewNotificationsIfIdle() @@ -509,36 +518,6 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } } - package func cleanupActiveReviewsForShutdown(reason: CodexReviewBackendModel.CancellationReason) async throws { - let runs = await activeReviewRunsForShutdown() - guard runs.isEmpty == false else { - return - } - for run in runs { - if Task.isCancelled { - return - } - try await interruptReview(run, reason: reason) - if Task.isCancelled { - return - } - try await cleanupReview(run) - } - } - - package func interruptActiveReviewsForShutdown(reason: CodexReviewBackendModel.CancellationReason) async throws { - let runs = await activeReviewRunsForShutdown() - guard runs.isEmpty == false else { - return - } - for run in runs { - if Task.isCancelled { - return - } - try await interruptReview(run, reason: reason) - } - } - package func notificationRouterMetricsForTesting() -> AppServerNotificationRouterMetrics { notificationRouterMetrics } @@ -655,14 +634,6 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { noteReviewThreadIDForCleanup(interruption.threadID, canonicalThreadID: run.threadID) } - private func markInterruptionTurnAbandoned( - _ interruption: AppServerReviewInterruption, - canonicalThreadID: String - ) { - markTurnAbandoned(interruption.turnID) - noteReviewThreadIDForCleanup(interruption.threadID, canonicalThreadID: canonicalThreadID) - } - private func markTurnAbandoned(_ turnID: String?) { guard let turnID = turnID?.nilIfEmpty else { return @@ -670,16 +641,6 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { abandonedTurnIDs.insert(turnID) } - private func activeReviewRunsForShutdown() async -> [CodexReviewBackendModel.Review.Run] { - let sessions = Array(reviewEventSessionsByAttemptID.values) - var runsByAttemptID: [String: CodexReviewBackendModel.Review.Run] = [:] - for session in sessions { - let run = await session.currentRun() - runsByAttemptID[run.attemptID] = run - } - return Array(runsByAttemptID.values) - } - private func bufferUnmatchedReviewNotification(_ notification: AppServerRoutedReviewNotification) -> Bool { guard reviewStartRequestsInFlight > 0, let threadID = notification.payload.threadID @@ -708,11 +669,10 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } private func sendTurnInterrupt( - for run: CodexReviewBackendModel.Review.Run, - willInterruptActiveTurn: (@Sendable (AppServerReviewInterruption) async -> Void)? = nil + for run: CodexReviewBackendModel.Review.Run ) async throws -> AppServerReviewInterruption { if let control = controlsByThreadID[run.threadID], - let interruption = try await control.interrupt(willInterruptActiveTurn: willInterruptActiveTurn) { + let interruption = try await control.interrupt() { return interruption } let threadID = appServerTurnThreadID(for: run) @@ -1096,8 +1056,6 @@ private actor AppServerReviewEventSession { private var isRunFinalized: Bool private var isDrainingStartupNotifications = false private var pendingStartupNotifications: [AppServerRoutedReviewNotification] = [] - private var terminalBarrierResolved = false - private var terminalBarrierWaiters: [CheckedContinuation] = [] private var metrics = AppServerReviewEventSessionMetrics() init( @@ -1165,19 +1123,6 @@ private actor AppServerReviewEventSession { cancellationRequestedMessage = nil } - func waitForTerminalBarrier() async { - if terminalBarrierResolved { - return - } - await withCheckedContinuation { continuation in - if terminalBarrierResolved { - continuation.resume() - } else { - terminalBarrierWaiters.append(continuation) - } - } - } - func receive(_ notification: AppServerRoutedReviewNotification) async { metrics.routed += 1 guard finished == false else { @@ -1218,7 +1163,6 @@ private actor AppServerReviewEventSession { pendingStartupNotifications.removeAll(keepingCapacity: true) await emitPrecedingEvents(precedingEvents) if let error { - resolveTerminalBarrier() await mailbox.fail(error) } else { await mailbox.finish() @@ -1506,24 +1450,11 @@ private actor AppServerReviewEventSession { event = .failed(message) } noteEmission(event) - resolveTerminalBarrier() await mailbox.append(event) recordReviewEvent(event, controlThreadID: controlThreadID) await mailbox.finish() } - private func resolveTerminalBarrier() { - guard terminalBarrierResolved == false else { - return - } - terminalBarrierResolved = true - let waiters = terminalBarrierWaiters - terminalBarrierWaiters.removeAll(keepingCapacity: false) - for waiter in waiters { - waiter.resume() - } - } - private func shouldCloseActiveCommandsBeforeEvents( notification: AppServerRoutedReviewNotification, decoded: DecodedReviewNotification diff --git a/Sources/CodexReviewAppServer/AppServerReviewControl.swift b/Sources/CodexReviewAppServer/AppServerReviewControl.swift index 11142a9f..bda0ad81 100644 --- a/Sources/CodexReviewAppServer/AppServerReviewControl.swift +++ b/Sources/CodexReviewAppServer/AppServerReviewControl.swift @@ -49,9 +49,7 @@ package final class AppServerReviewControl: @unchecked Sendable { } @discardableResult - package func interrupt( - willInterruptActiveTurn: (@Sendable (AppServerReviewInterruption) async -> Void)? = nil - ) async throws -> AppServerReviewInterruption? { + package func interrupt() async throws -> AppServerReviewInterruption? { let currentPhase = phaseSnapshot() switch currentPhase { case .preparing, .finished: @@ -59,14 +57,12 @@ package final class AppServerReviewControl: @unchecked Sendable { case .threadStarted(let threadID): return try await sendInterrupt( threadID: threadID, - turnID: "", - willInterruptActiveTurn: willInterruptActiveTurn + turnID: "" ) case .reviewStarted(let turnThreadID, let turnID): return try await sendInterrupt( threadID: turnThreadID, - turnID: turnID, - willInterruptActiveTurn: willInterruptActiveTurn + turnID: turnID ) } } @@ -85,8 +81,7 @@ package final class AppServerReviewControl: @unchecked Sendable { private func sendInterrupt( threadID: String, - turnID: String, - willInterruptActiveTurn _: (@Sendable (AppServerReviewInterruption) async -> Void)? + turnID: String ) async throws -> AppServerReviewInterruption { let _: EmptyResponse = try await client.send(AppServerAPI.Turn.Interrupt.Request( params: .init(threadID: threadID, turnID: turnID) diff --git a/Sources/CodexReviewHost/CodexReviewHost.swift b/Sources/CodexReviewHost/CodexReviewHost.swift index 4d92aec5..54a3d6e7 100644 --- a/Sources/CodexReviewHost/CodexReviewHost.swift +++ b/Sources/CodexReviewHost/CodexReviewHost.swift @@ -252,17 +252,17 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { } func beginReviewRecovery( - _ run: CodexReviewBackendModel.Review.Run, - reason: CodexReviewBackendModel.CancellationReason + _ barrier: ReviewAttemptRecoveryBarrier ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { - try await backend.beginReviewRecovery(run, reason: reason) + try await backend.beginReviewRecovery(barrier) } func resumeReviewRecovery( _ token: CodexReviewBackendModel.Review.RecoveryToken, - request: CodexReviewBackendModel.Review.Start + request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { - try await backend.resumeReviewRecovery(token, request: request) + try await backend.resumeReviewRecovery(token, request: request, admission: admission) } func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { diff --git a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift index b7584ad5..a13a8d02 100644 --- a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift +++ b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift @@ -1015,23 +1015,27 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } func beginReviewRecovery( - _ run: CodexReviewBackendModel.Review.Run, - reason: CodexReviewBackendModel.CancellationReason + _ barrier: ReviewAttemptRecoveryBarrier ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { guard let appServerBackend else { throw CodexReviewAPI.Error.io("Review runtime is not running.") } - return try await appServerBackend.beginReviewRecovery(run, reason: reason) + return try await appServerBackend.beginReviewRecovery(barrier) } func resumeReviewRecovery( _ token: CodexReviewBackendModel.Review.RecoveryToken, - request: CodexReviewBackendModel.Review.Start + request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { guard let appServerBackend else { throw CodexReviewAPI.Error.io("Review runtime is not running.") } - return try await appServerBackend.resumeReviewRecovery(token, request: request) + return try await appServerBackend.resumeReviewRecovery( + token, + request: request, + admission: admission + ) } func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { diff --git a/Sources/CodexReviewTesting/TestSupport.swift b/Sources/CodexReviewTesting/TestSupport.swift index abf47221..98b1f1e7 100644 --- a/Sources/CodexReviewTesting/TestSupport.swift +++ b/Sources/CodexReviewTesting/TestSupport.swift @@ -145,7 +145,7 @@ package struct FakeCodexReviewBackendTimeout: LocalizedError, Sendable { } } -private func withFakeBackendTimeout( +package func withFakeBackendTimeout( operation: String, timeout: Duration, wait: @escaping @Sendable () async -> Void @@ -166,6 +166,12 @@ private func withFakeBackendTimeout( } package actor FakeCodexReviewBackend: CodexReviewBackend { + private struct MatchingInterruptWaiter { + var run: CodexReviewBackendModel.Review.Run + var reason: CodexReviewBackendModel.CancellationReason + var continuation: CheckedContinuation + } + package enum Command: Equatable, Sendable { case readSettings case applySettings(CodexReviewBackendModel.Settings.Change) @@ -192,6 +198,7 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { private var cleanupFailure: ReviewRuntimeCloseFailure? private var interruptReviewGate: AsyncGate? private var interruptReviewWaiters: [UUID: CheckedContinuation] = [:] + private var matchingInterruptReviewWaiters: [UUID: MatchingInterruptWaiter] = [:] private var beginReviewRecoveryWaiters: [UUID: CheckedContinuation] = [:] private var startReviewGate: AsyncGate? private var startReviewWaiters: [UUID: CheckedContinuation] = [:] @@ -333,6 +340,43 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { } } + package func waitForInterruptReview( + run: CodexReviewBackendModel.Review.Run, + reason: CodexReviewBackendModel.CancellationReason + ) async { + if commands.contains(.interruptReview(run, reason)) { + return + } + let waiterID = UUID() + await withTaskCancellationHandler { + await withCheckedContinuation { continuation in + if commands.contains(.interruptReview(run, reason)) { + continuation.resume() + } else { + matchingInterruptReviewWaiters[waiterID] = .init( + run: run, + reason: reason, + continuation: continuation + ) + } + } + } onCancel: { + Task { + await self.cancelMatchingInterruptReviewWaiter(id: waiterID) + } + } + } + + package func waitForInterruptReview( + run: CodexReviewBackendModel.Review.Run, + reason: CodexReviewBackendModel.CancellationReason, + timeout: Duration + ) async throws { + try await withFakeBackendTimeout(operation: "interruptReview for \(run.attemptID)", timeout: timeout) { + await self.waitForInterruptReview(run: run, reason: reason) + } + } + package func waitForBeginReviewRecovery() async { if commands.contains(where: { if case .beginReviewRecovery = $0 { @@ -495,6 +539,13 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { for waiter in waiters { waiter.resume() } + let matchingWaiters = matchingInterruptReviewWaiters.filter { + $0.value.run == run && $0.value.reason == reason + } + for (waiterID, waiter) in matchingWaiters { + matchingInterruptReviewWaiters.removeValue(forKey: waiterID) + waiter.continuation.resume() + } if let interruptReviewGate { await interruptReviewGate.wait() } @@ -515,29 +566,53 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { } } + package func releaseHeldOperationsForCleanup() async { + await startReviewGate?.open() + await interruptReviewGate?.open() + await resumeReviewRecoveryGate?.open() + } + package func beginReviewRecovery( - _ run: CodexReviewBackendModel.Review.Run, - reason: CodexReviewBackendModel.CancellationReason + _ barrier: ReviewAttemptRecoveryBarrier ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { - commands.append(.beginReviewRecovery(run, reason)) + commands.append(.beginReviewRecovery( + barrier.run, + .init(message: barrier.cancellation.message) + )) let waiters = Array(beginReviewRecoveryWaiters.values) beginReviewRecoveryWaiters.removeAll(keepingCapacity: false) for waiter in waiters { waiter.resume() } - if let interruptReviewGate { - await interruptReviewGate.wait() - } - if let interruptFailureMessage { - throw FakeCodexReviewBackendError(message: interruptFailureMessage) - } + let run = barrier.run return .init(interruptedRun: run, rollbackThreadID: run.reviewThreadID ?? run.threadID) } package func resumeReviewRecovery( _ token: CodexReviewBackendModel.Review.RecoveryToken, - request: CodexReviewBackendModel.Review.Start + request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { + let run = token.interruptedRun + let recoveredRun = nextRecoveredRun ?? .init( + attemptID: "attempt-recovered", + threadID: run.threadID, + turnID: "turn-recovered", + reviewThreadID: run.reviewThreadID, + model: run.model ?? request.model + ) + let provisionalRun = CodexReviewBackendModel.Review.Run( + attemptID: recoveredRun.attemptID, + threadID: recoveredRun.threadID, + reviewThreadID: recoveredRun.threadID, + model: recoveredRun.model + ) + await admission.recordPreparedThread(provisionalRun) + guard await admission.admitReviewStartDispatch(for: provisionalRun) else { + throw ReviewStartCancelledBeforeDispatch( + cancellation: await admission.cancellationRequest() ?? .system() + ) + } commands.append(.resumeReviewRecovery(token, request)) let waiters = Array(resumeReviewRecoveryWaiters.values) resumeReviewRecoveryWaiters.removeAll(keepingCapacity: false) @@ -546,18 +621,12 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { } if let resumeReviewRecoveryGate { await resumeReviewRecoveryGate.wait() + try Task.checkCancellation() } if let recoveryFailureMessage { throw FakeCodexReviewBackendError(message: recoveryFailureMessage) } - let run = token.interruptedRun - let recoveredRun = nextRecoveredRun ?? .init( - attemptID: "attempt-recovered", - threadID: run.threadID, - turnID: "turn-recovered", - reviewThreadID: run.reviewThreadID, - model: run.model ?? request.model - ) + await admission.recordActiveRun(recoveredRun) return .init(run: recoveredRun, events: eventMailbox(for: recoveredRun)) } @@ -610,6 +679,10 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { interruptReviewWaiters.removeValue(forKey: id)?.resume() } + private func cancelMatchingInterruptReviewWaiter(id: UUID) { + matchingInterruptReviewWaiters.removeValue(forKey: id)?.continuation.resume() + } + private func cancelBeginReviewRecoveryWaiter(id: UUID) { beginReviewRecoveryWaiters.removeValue(forKey: id)?.resume() } @@ -876,17 +949,21 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { } package func beginReviewRecovery( - _ run: CodexReviewBackendModel.Review.Run, - reason: CodexReviewBackendModel.CancellationReason + _ barrier: ReviewAttemptRecoveryBarrier ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { - try await reviewBackend.beginReviewRecovery(run, reason: reason) + try await reviewBackend.beginReviewRecovery(barrier) } package func resumeReviewRecovery( _ token: CodexReviewBackendModel.Review.RecoveryToken, - request: CodexReviewBackendModel.Review.Start + request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { - try await reviewBackend.resumeReviewRecovery(token, request: request) + try await reviewBackend.resumeReviewRecovery( + token, + request: request, + admission: admission + ) } package func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { diff --git a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift index c0fa5273..bfbafe53 100644 --- a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift +++ b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift @@ -34,11 +34,81 @@ private extension AppServerCodexReviewBackend { try await beginReviewRecovery(attempt.run, reason: reason) } + func beginReviewRecovery( + _ run: CodexReviewBackendModel.Review.Run, + reason: CodexReviewBackendModel.CancellationReason + ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { + let barrier = try await makeRecoveryBarrier(backend: self, for: run, reason: reason) + return try await beginReviewRecovery(barrier) + } + + func resumeReviewRecovery( + _ token: CodexReviewBackendModel.Review.RecoveryToken, + request: CodexReviewBackendModel.Review.Start + ) async throws -> BackendReviewAttempt { + let admission = ReviewStartAdmission() + let task = await admission.start { admission in + try await self.resumeReviewRecovery( + token, + request: request, + admission: admission + ) + } + return try await task.value + } + func cleanupReview(_ attempt: BackendReviewAttempt) async throws { try await cleanupReview(attempt.run) } } +private func makeRecoveryBarrier( + backend: AppServerCodexReviewBackend, + for run: CodexReviewBackendModel.Review.Run, + reason: CodexReviewBackendModel.CancellationReason +) async throws -> ReviewAttemptRecoveryBarrier { + let admission = ReviewStartAdmission() + let startTask = await admission.start { admission in + #expect(await admission.admitThreadStartDispatch()) + let provisionalRun = CodexReviewBackendModel.Review.Run( + attemptID: run.attemptID, + threadID: run.threadID, + reviewThreadID: run.threadID, + model: run.model + ) + await admission.recordPreparedThread(provisionalRun) + #expect(await admission.admitReviewStartDispatch(for: provisionalRun)) + await admission.recordActiveRun(run) + return .init(run: run) + } + _ = try await startTask.value + + let requestOutcome = RecoveryRequestOutcomeProbe() + let cancellation = ReviewCancellation.system(message: reason.message) + let recovery = Task { + try await admission.interruptForRecovery( + cancellation, + interrupt: { run, reason in + do { + try await backend.interruptReview(run, reason: reason) + await requestOutcome.record(.success(())) + } catch { + await requestOutcome.record(.failure(error)) + throw error + } + }, + forceClose: {} + ) + } + if case .success = await requestOutcome.wait() { + try await admission.recordCanonicalTerminal( + .interrupted(.requested(cancellation)), + for: run + ) + } + return try await recovery.value +} + private extension BackendReviewAttempt { var attemptID: String { run.attemptID } var threadID: String { run.threadID } @@ -2036,7 +2106,7 @@ struct AppServerClientTests { )) } - @Test func backendInterruptAckDoesNotSynthesizeTerminalBeforeEventStreamRegistration() async throws { + @Test func backendInterruptReturnsAfterAckWhileCanonicalTerminalRemainsPending() async throws { let transport = FakeJSONRPCTransport() try await enqueueInitialize(transport) try await transport.enqueue(AppServerAPI.Thread.Start.Response(threadID: "thread-1", model: "gpt-5"), for: "thread/start") @@ -2050,8 +2120,13 @@ struct AppServerClientTests { request: .init(cwd: "/tmp/project", target: .uncommittedChanges) )) - async let cancellation: Void = backend.interruptReview(run, reason: .init(message: "Stop")) + let interruptReturned = AsyncGate() + let cancellation = Task { + try await backend.interruptReview(run, reason: .init(message: "Stop")) + await interruptReturned.open() + } await transport.waitForResponseDelivery(method: "turn/interrupt") + await interruptReturned.wait() #expect(await run.events.isFinished() == false) try await transport.emitServerNotification( method: "turn/completed", @@ -2060,13 +2135,35 @@ struct AppServerClientTests { turn: .init(id: "turn-1", status: "interrupted", error: .init(message: "Stop")) ) ) - try await cancellation + try await cancellation.value var iterator = await eventSequence(backend, run).makeAsyncIterator() #expect(try await iterator.next() == .cancelled("Stop")) #expect(try await iterator.next() == nil) } + @Test func backendInterruptOutcomeUnknownReturnsBeforeSessionTerminal() async throws { + let transport = FakeJSONRPCTransport() + try await enqueueInitialize(transport) + await transport.enqueueFailure(JSONRPC.Error.closed, for: "turn/interrupt") + let backend = AppServerCodexReviewBackend(client: .init(transport: transport)) + let run = CodexReviewBackendModel.Review.Run( + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "thread-1", + model: "gpt-5" + ) + let events = await eventSequence(backend, run) + + do { + try await backend.interruptReview(run, reason: .init(message: "Stop")) + Issue.record("Expected outcome-unknown interrupt failure.") + } catch let failure as ReviewInterruptRequestFailure { + #expect(failure.outcome == .outcomeUnknown(message: JSONRPC.Error.closed.localizedDescription)) + } + #expect(await events.mailbox.isFinished() == false) + } + @Test func backendRecoverReviewRollsBackAndRestartsSameThread() async throws { let transport = FakeJSONRPCTransport() try await enqueueInitialize(transport) @@ -2505,10 +2602,10 @@ struct AppServerClientTests { ) let events = await eventSequence(backend, run) - await #expect(throws: JSONRPC.Error.responseError( + await #expect(throws: ReviewInterruptRequestFailure(outcome: .rejected( code: -32602, message: "expected active turn id turn-old but found turn-active" - )) { + ))) { try await backend.beginReviewRecovery( run, reason: .init(message: "Network unavailable; waiting to reconnect.") @@ -2568,12 +2665,10 @@ struct AppServerClientTests { )) } - @Test func backendCanonicalTerminalWinsWhileRecoveryInterruptIsInFlight() async throws { + @Test func backendCanonicalTerminalIsDeliveredWhileInterruptAckIsInFlight() async throws { let transport = FakeJSONRPCTransport() try await enqueueInitialize(transport) try await transport.enqueue(EmptyResponse(), for: "turn/interrupt") - try await transport.enqueue(EmptyResponse(), for: "thread/rollback") - try await transport.enqueue(AppServerAPI.Review.Start.Response(turnID: "turn-2", reviewThreadID: "thread-1"), for: "review/start") let interruptGate = AsyncGate() await transport.holdNext(method: "turn/interrupt", gate: interruptGate) let backend = AppServerCodexReviewBackend(client: .init(transport: transport)) @@ -2587,16 +2682,12 @@ struct AppServerClientTests { var initialIterator = initialEvents.makeAsyncIterator() defer { withExtendedLifetime(initialEvents) {} } - async let recovered = backend.resumeReviewRecovery( - run, - request: CodexReviewBackendModel.Review.Start( - jobID: "job-1", - sessionID: "session-1", - request: .init(cwd: "/tmp/project", target: .baseBranch("main")), - model: "gpt-5" - ), - reason: .init(message: "Network unavailable; waiting to reconnect.") - ) + let interruption = Task { + try await backend.interruptReview( + run, + reason: .init(message: "Network unavailable; waiting to reconnect.") + ) + } let interruptRequested = await waitUntil { await transport.recordedRequests().contains { $0.method == "turn/interrupt" } } @@ -2612,28 +2703,9 @@ struct AppServerClientTests { #expect(try await initialIterator.next() == .failed("Old turn failed")) await interruptGate.open() - let recoveredRun = try await recovered - #expect(recoveredRun.turnID == "turn-2") - let recoveredEvents = await eventSequence(backend, recoveredRun) - var iterator = recoveredEvents.makeAsyncIterator() - - try await transport.emitServerNotification( - method: "turn/started", - params: TestTurnNotification(threadID: "thread-1", turn: .init(id: "turn-2")) - ) - #expect(try await iterator.next() == .started( - turnID: "turn-2", - reviewThreadID: "thread-1", - model: nil - )) - - try await transport.emitServerNotification( - method: "turn/completed", - params: TestTurnNotification(threadID: "thread-1", turn: .init(id: "turn-2", status: "completed")) - ) - #expect(try await iterator.next() == .failed( - ReviewIngestionError.missingFinalReview.localizedDescription - )) + try await interruption.value + let methods = await transport.recordedRequests().map(\.method) + #expect(methods == ["initialize", "turn/interrupt"]) } @Test func backendIgnoresStaleInterruptedTurnNotificationsWhileRollbackIsInFlight() async throws { @@ -3228,7 +3300,7 @@ struct AppServerClientTests { #expect(try await iterator.next() == nil) } - @Test func backendInterruptAckWaitsForCanonicalTerminalToFinishEventStream() async throws { + @Test func backendInterruptAckReturnsBeforeCanonicalTerminalWithActiveStream() async throws { let transport = FakeJSONRPCTransport() try await enqueueInitialize(transport) try await transport.enqueue(AppServerAPI.Thread.Start.Response(threadID: "thread-1", model: "gpt-5"), for: "thread/start") @@ -3250,8 +3322,13 @@ struct AppServerClientTests { #expect(try await iterator.next() == .started(turnID: "turn-1", reviewThreadID: "thread-1", model: nil)) - async let cancellation: Void = backend.interruptReview(run, reason: .init(message: "Stop")) + let interruptReturned = AsyncGate() + let cancellation = Task { + try await backend.interruptReview(run, reason: .init(message: "Stop")) + await interruptReturned.open() + } await transport.waitForResponseDelivery(method: "turn/interrupt") + await interruptReturned.wait() #expect(await run.events.isFinished() == false) try await transport.emitServerNotification( method: "turn/completed", @@ -3260,7 +3337,7 @@ struct AppServerClientTests { turn: .init(id: "turn-1", status: "interrupted", error: .init(message: "Stop")) ) ) - try await cancellation + try await cancellation.value #expect(try await iterator.next() == .cancelled("Stop")) #expect(try await iterator.next() == nil) @@ -6002,3 +6079,33 @@ private actor CallCounter { count } } + +private actor RecoveryRequestOutcomeProbe { + private var outcome: Result? + private var waiters: [CheckedContinuation, Never>] = [] + + func record(_ outcome: Result) { + guard self.outcome == nil else { + return + } + self.outcome = outcome + let waiters = waiters + self.waiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume(returning: outcome) + } + } + + func wait() async -> Result { + if let outcome { + return outcome + } + return await withCheckedContinuation { continuation in + if let outcome { + continuation.resume(returning: outcome) + } else { + waiters.append(continuation) + } + } + } +} diff --git a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift index 73718737..1301fcaf 100644 --- a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift +++ b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift @@ -1394,11 +1394,24 @@ struct CodexReviewHostTests { ) } try #require(await waitUntil(timeout: .seconds(2)) { store.jobs.first?.core.run.turnID == "turn-1" }) + let jobID = try #require(store.jobs.first?.id) networkMonitor.yield(.init(status: .unsatisfied)) try #require(await waitUntil(timeout: .seconds(2)) { await transport.recordedRequests().map(\.method).contains("turn/interrupt") }) + try await transport.emitServerNotification( + method: "turn/completed", + params: HostTurnNotification( + threadID: "review-thread-1", + turnID: "turn-1", + status: "interrupted", + errorMessage: "Network unavailable; waiting to reconnect." + ) + ) + try #require(await waitUntil(timeout: .seconds(2)) { + store.reviewRecoveryWaitingJobIDs.contains(jobID) + }) let stopFinished = CompletionFlag() let stopTask = Task { @MainActor in diff --git a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift index 1b77f960..1ede9aec 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift @@ -900,17 +900,179 @@ struct CodexReviewStoreCommandTests { ) networkMonitor.yield(.init(status: .unsatisfied)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await completeNetworkRecoveryBarrier(backend: backend, store: store) let running = try store.readReview(jobID: "job-1") #expect(running.core.lifecycle.status == .running) #expect(running.core.output.summary == "Network unavailable; waiting to reconnect.") _ = try await store.cancelReview(jobID: "job-1", cancellation: .mcpClient(message: "Stop")) - await backend.yield(.cancelled("Stop")) _ = try await result } } + @Test func networkRecoveryAckDoesNotTokenizeBeforeCanonicalTerminal() async throws { + let backend = FakeCodexReviewBackend() + let networkMonitor = ManualCodexReviewNetworkMonitor() + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-1" }), + networkMonitor: networkMonitor, + networkRecoveryPolicy: .init(sleep: { _ in }) + ) + try await withStoreCommandTestCleanup(backend: backend, store: store) { + async let result = store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .baseBranch("main")) + ) + try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running) != nil) + let run = try #require(store.activeRuns["job-1"]) + let reason = CodexReviewBackendModel.CancellationReason( + message: "Network unavailable; waiting to reconnect." + ) + + networkMonitor.yield(.init(status: .unsatisfied)) + try await backend.waitForInterruptReview(run: run, reason: reason, timeout: .seconds(2)) + + let commandsBeforeTerminal = await backend.recordedCommands() + #expect(commandsBeforeTerminal.contains { if case .beginReviewRecovery = $0 { true } else { false } } == false) + #expect(store.reviewRecoveryWaitingJobIDs.contains("job-1") == false) + + await backend.yield(.cancelled(reason.message), for: run) + try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + let barrierPublished = await waitUntil(timeout: .seconds(2)) { + store.reviewRecoveryWaitingJobIDs.contains("job-1") + } + #expect(barrierPublished) + let waiting = try store.readReview(jobID: "job-1") + #expect(waiting.core.lifecycle.status == .running) + #expect(waiting.core.lifecycle.terminal == nil) + + _ = try await store.cancelReview( + jobID: "job-1", + cancellation: .mcpClient(message: "Stop") + ) + _ = try await result + } + } + + @Test func networkRecoveryTerminalBeforeAckDrainsRequestBeforeTokenization() async throws { + let backend = FakeCodexReviewBackend() + let interruptGate = AsyncGate() + await backend.holdInterruptReview(with: interruptGate) + let networkMonitor = ManualCodexReviewNetworkMonitor() + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-1" }), + networkMonitor: networkMonitor, + networkRecoveryPolicy: .init(sleep: { _ in }) + ) + try await withStoreCommandTestCleanup(backend: backend, store: store) { + async let result = store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .baseBranch("main")) + ) + try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running) != nil) + let run = try #require(store.activeRuns["job-1"]) + let reason = CodexReviewBackendModel.CancellationReason( + message: "Network unavailable; waiting to reconnect." + ) + + networkMonitor.yield(.init(status: .unsatisfied)) + try await backend.waitForInterruptReview(run: run, reason: reason, timeout: .seconds(2)) + await backend.yield(.cancelled(reason.message), for: run) + + let commandsBeforeAck = await backend.recordedCommands() + #expect(commandsBeforeAck.contains { if case .beginReviewRecovery = $0 { true } else { false } } == false) + await interruptGate.open() + try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + let barrierPublished = await waitUntil(timeout: .seconds(2)) { + store.reviewRecoveryWaitingJobIDs.contains("job-1") + } + #expect(barrierPublished) + + _ = try await store.cancelReview( + jobID: "job-1", + cancellation: .mcpClient(message: "Stop") + ) + _ = try await result + } + } + + @Test func networkRecoveryRejectionMintsNoTokenAndFailsProduct() async throws { + let backend = FakeCodexReviewBackend() + await backend.failInterrupts(message: "Recovery rejected") + let networkMonitor = ManualCodexReviewNetworkMonitor() + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-1" }), + networkMonitor: networkMonitor, + networkRecoveryPolicy: .init(sleep: { _ in }) + ) + try await withStoreCommandTestCleanup(backend: backend, store: store) { + async let result = store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .baseBranch("main")) + ) + try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running) != nil) + let run = try #require(store.activeRuns["job-1"]) + let reason = CodexReviewBackendModel.CancellationReason( + message: "Network unavailable; waiting to reconnect." + ) + + networkMonitor.yield(.init(status: .unsatisfied)) + try await backend.waitForInterruptReview(run: run, reason: reason, timeout: .seconds(2)) + let read = try await result + + #expect(read.core.lifecycle.status == .failed) + #expect(read.core.lifecycle.errorMessage == "Recovery rejected") + let commands = await backend.recordedCommands() + #expect(commands.contains { if case .beginReviewRecovery = $0 { true } else { false } } == false) + #expect(commands.contains { if case .resumeReviewRecovery = $0 { true } else { false } } == false) + } + } + + @Test func naturalCompletionDuringRecoveryInterruptSuppressesReplacement() async throws { + let backend = FakeCodexReviewBackend() + let interruptGate = AsyncGate() + await backend.holdInterruptReview(with: interruptGate) + let networkMonitor = ManualCodexReviewNetworkMonitor() + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-1" }), + networkMonitor: networkMonitor, + networkRecoveryPolicy: .init(sleep: { _ in }) + ) + try await withStoreCommandTestCleanup(backend: backend, store: store) { + async let result = store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .baseBranch("main")) + ) + try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running) != nil) + let run = try #require(store.activeRuns["job-1"]) + let reason = CodexReviewBackendModel.CancellationReason( + message: "Network unavailable; waiting to reconnect." + ) + + networkMonitor.yield(.init(status: .unsatisfied)) + try await backend.waitForInterruptReview(run: run, reason: reason, timeout: .seconds(2)) + await backend.yield( + .completed(summary: "Succeeded.", result: "natural review"), + for: run + ) + + let commandsBeforeAck = await backend.recordedCommands() + #expect(commandsBeforeAck.contains { if case .beginReviewRecovery = $0 { true } else { false } } == false) + await interruptGate.open() + let read = try await result + + #expect(read.core.lifecycle.status == .succeeded) + #expect(read.core.output.lastAgentMessage == "natural review") + let commands = await backend.recordedCommands() + #expect(commands.contains { if case .beginReviewRecovery = $0 { true } else { false } } == false) + #expect(commands.contains { if case .resumeReviewRecovery = $0 { true } else { false } } == false) + } + } + @Test func networkRecoveryWaitDiscardsOldAttemptCompletion() async throws { let initialRun = CodexReviewBackendModel.Review.Run( threadID: "thread-1", @@ -942,11 +1104,15 @@ struct CodexReviewStoreCommandTests { ) networkMonitor.yield(.init(status: .unsatisfied)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await completeNetworkRecoveryBarrier( + backend: backend, + store: store, + beforeTerminal: { run in + await backend.yield(.message("completed review text"), for: run) + } + ) _ = try await running - await backend.yield(.message("completed review text"), for: initialRun) - await backend.yield(.completed(summary: "Succeeded.", result: nil), for: initialRun) networkMonitor.yield(.satisfied()) try await backend.waitForResumeReviewRecovery(timeout: .seconds(2)) try #require(await waitForRunAttemptActivation(store: store, run: recoveredRun)) @@ -1005,15 +1171,19 @@ struct CodexReviewStoreCommandTests { ) networkMonitor.yield(.init(status: .unsatisfied)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await completeNetworkRecoveryBarrier( + backend: backend, + store: store, + beforeTerminal: { run in + await backend.yield(.message("completed during settle"), for: run) + } + ) await sleeper.blockFutureSleeps() networkMonitor.yield(.satisfied()) #expect(await waitUntil { store.job(id: "job-1")?.core.output.summary == "Network restored; restarting review." }) - await backend.yield(.message("completed during settle"), for: initialRun) - await backend.yield(.completed(summary: "Succeeded.", result: nil), for: initialRun) await settleGate.open() try await backend.waitForResumeReviewRecovery(timeout: .seconds(2)) try #require(await waitForRunAttemptActivation(store: store, run: recoveredRun)) @@ -1072,7 +1242,7 @@ struct CodexReviewStoreCommandTests { ) networkMonitor.yield(.init(status: .unsatisfied)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await completeNetworkRecoveryBarrier(backend: backend, store: store) await sleeper.blockFutureSleeps() networkMonitor.yield(.satisfied()) #expect(await waitUntil { @@ -1137,7 +1307,7 @@ struct CodexReviewStoreCommandTests { ), for: initialRun) networkMonitor.yield(.init(status: .unsatisfied)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await completeNetworkRecoveryBarrier(backend: backend, store: store) _ = try await running networkMonitor.yield(.satisfied()) try await backend.waitForResumeReviewRecovery(timeout: .seconds(2)) @@ -1196,7 +1366,7 @@ struct CodexReviewStoreCommandTests { }) networkMonitor.yield(.init(status: .unsatisfied)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await completeNetworkRecoveryBarrier(backend: backend, store: store) let commandsAfterInterrupt = await backend.recordedCommands() let interruptedRuns = commandsAfterInterrupt.compactMap { command -> CodexReviewBackendModel.Review.Run? in if case .beginReviewRecovery(let run, _) = command { @@ -1256,9 +1426,13 @@ struct CodexReviewStoreCommandTests { ) networkMonitor.yield(.init(status: .unsatisfied)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) - await backend.yield(.message("stale aborted output"), for: initialRun) - await backend.yield(.cancelled("Network lost"), for: initialRun) + try await completeNetworkRecoveryBarrier( + backend: backend, + store: store, + beforeTerminal: { run in + await backend.yield(.message("stale aborted output"), for: run) + } + ) networkMonitor.yield(.satisfied()) try await backend.waitForResumeReviewRecovery(timeout: .seconds(2)) #expect(await waitUntil { @@ -1320,7 +1494,7 @@ struct CodexReviewStoreCommandTests { } != nil) networkMonitor.yield(.init(status: .unsatisfied)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await completeNetworkRecoveryBarrier(backend: backend, store: store) networkMonitor.yield(.satisfied()) try await backend.waitForResumeReviewRecovery(timeout: .seconds(2)) try #require(await waitForRunAttemptActivation(store: store, run: recoveredRun)) @@ -1369,12 +1543,11 @@ struct CodexReviewStoreCommandTests { ) networkMonitor.yield(.init(status: .unsatisfied)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await completeNetworkRecoveryBarrier(backend: backend, store: store) networkMonitor.yield(.satisfied()) try await backend.waitForResumeReviewRecovery(timeout: .seconds(2)) try #require(await waitForRunAttemptActivation(store: store, run: recoveredRun)) - await backend.yield(.completed(summary: "Succeeded.", result: "stale review"), for: initialRun) await backend.yield(.completed(summary: "Succeeded.", result: "recovered review"), for: recoveredRun) let read = try await result @@ -1416,10 +1589,9 @@ struct CodexReviewStoreCommandTests { ) networkMonitor.yield(.init(status: .unsatisfied)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await completeNetworkRecoveryBarrier(backend: backend, store: store) networkMonitor.yield(.satisfied()) try await backend.waitForResumeReviewRecovery(timeout: .seconds(2)) - await backend.yield(.cancelled("Network lost"), for: initialRun) await recoverGate.open() try #require(await waitForRunAttemptActivation(store: store, run: recoveredRun)) @@ -1462,8 +1634,7 @@ struct CodexReviewStoreCommandTests { ) networkMonitor.yield(.init(status: .unsatisfied)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) - await backend.finishEvents(for: initialRun) + try await completeNetworkRecoveryBarrier(backend: backend, store: store) networkMonitor.yield(.satisfied()) try await backend.waitForResumeReviewRecovery(timeout: .seconds(2)) try #require(await waitForRunAttemptActivation(store: store, run: recoveredRun)) @@ -1509,17 +1680,34 @@ struct CodexReviewStoreCommandTests { ) networkMonitor.yield(.init(status: .unsatisfied)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await completeNetworkRecoveryBarrier(backend: backend, store: store) + let oldAdmission = try #require(store.reviewStartAdmissions["job-1"]) networkMonitor.yield(.satisfied()) try await backend.waitForResumeReviewRecovery(timeout: .seconds(2)) - let cancel = try await store.cancelReview(jobID: "job-1", cancellation: .mcpClient(message: "Stop")) - #expect(cancel.cancelled) + let recoveredAdmission = try #require(store.reviewStartAdmissions["job-1"]) + #expect(recoveredAdmission !== oldAdmission) + let cancelTask = Task { @MainActor in + try await store.cancelReview( + jobID: "job-1", + cancellation: .mcpClient(message: "Stop") + ) + } + #expect(await recoveredAdmission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) await recoverGate.open() + _ = try #require(await recoveredAdmission.waitForActiveRun()) + try await backend.waitForInterruptReview( + run: recoveredRun, + reason: .init(message: "Stop"), + timeout: .seconds(2) + ) + await backend.yield(.cancelled("Stop"), for: recoveredRun) + let cancel = try await cancelTask.value + #expect(cancel.cancelled) let read = try await result #expect(read.core.lifecycle.status == .cancelled) - #expect(read.core.run.turnID == "turn-1") + #expect(read.core.run.turnID == "turn-2") let commands = await backend.recordedCommands() #expect(commands.contains(.interruptReview( @@ -1557,9 +1745,8 @@ struct CodexReviewStoreCommandTests { ) networkMonitor.yield(.init(status: .unsatisfied)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await completeNetworkRecoveryBarrier(backend: backend, store: store) _ = try await running - await backend.finishEvents(for: initialRun) let cancel = try await store.cancelReview(jobID: "job-1", cancellation: .mcpClient(message: "Stop")) let cleanedUp = await waitUntil { @@ -1678,7 +1865,7 @@ struct CodexReviewStoreCommandTests { ) networkMonitor.yield(.init(status: .unsatisfied)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await completeNetworkRecoveryBarrier(backend: backend, store: store) _ = try await running let locallyCancelledJobIDs = store.cancelActiveReviewsLocallyForRuntimeStop( @@ -1780,9 +1967,8 @@ struct CodexReviewStoreCommandTests { ) networkMonitor.yield(.init(status: .unsatisfied)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await completeNetworkRecoveryBarrier(backend: backend, store: store) _ = try await store.cancelReview(jobID: "job-1", cancellation: .mcpClient(message: "Stop")) - await backend.finishEvents(for: initialRun) let read = try await result #expect(read.core.lifecycle.status == .cancelled) @@ -1820,9 +2006,13 @@ struct CodexReviewStoreCommandTests { ) networkMonitor.yield(.init(status: .unsatisfied)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) - await backend.yield(.message("stale old attempt output"), for: initialRun) - await backend.yield(.completed(summary: "Succeeded.", result: nil), for: initialRun) + try await completeNetworkRecoveryBarrier( + backend: backend, + store: store, + beforeTerminal: { run in + await backend.yield(.message("stale old attempt output"), for: run) + } + ) networkMonitor.yield(.satisfied()) try await backend.waitForResumeReviewRecovery(timeout: .seconds(2)) @@ -1902,7 +2092,7 @@ struct CodexReviewStoreCommandTests { ) networkMonitor.yield(.init(status: .requiresConnection)) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await completeNetworkRecoveryBarrier(backend: backend, store: store) networkMonitor.yield(.satisfied()) let read = try await result @@ -2108,7 +2298,10 @@ struct CodexReviewStoreCommandTests { #expect(failedBeforeOutageConfirmed == false) await debounceGate.open() - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await completeNetworkRecoveryAfterConnectionTerminal( + backend: backend, + store: store + ) networkMonitor.yield(.satisfied()) try await backend.waitForResumeReviewRecovery(timeout: .seconds(2)) try #require(await waitForRunAttemptActivation(store: store, run: recoveredRun)) @@ -2567,11 +2760,58 @@ private func withStoreCommandTestCleanup( await cleanupStoreCommandTest(backend: backend, store: store) } +@MainActor +private func completeNetworkRecoveryBarrier( + backend: FakeCodexReviewBackend, + store: CodexReviewStore, + jobID: String = "job-1", + beforeTerminal: @MainActor (CodexReviewBackendModel.Review.Run) async -> Void = { _ in } +) async throws { + let becameActive = await waitUntil(timeout: .seconds(2)) { + store.activeRuns[jobID] != nil + } + guard becameActive else { + throw FakeCodexReviewBackendError(message: "Review did not become active before recovery.") + } + guard let run = store.activeRuns[jobID] else { + throw FakeCodexReviewBackendError(message: "Missing active run for recovery barrier.") + } + let reason = CodexReviewBackendModel.CancellationReason( + message: "Network unavailable; waiting to reconnect." + ) + try await backend.waitForInterruptReview(run: run, reason: reason, timeout: .seconds(2)) + await beforeTerminal(run) + await backend.yield(.cancelled(reason.message), for: run) + try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + let barrierPublished = await waitUntil(timeout: .seconds(2)) { + store.reviewRecoveryWaitingJobIDs.contains(jobID) + } + guard barrierPublished else { + throw FakeCodexReviewBackendError(message: "Store did not publish recovery barrier.") + } +} + +@MainActor +private func completeNetworkRecoveryAfterConnectionTerminal( + backend: FakeCodexReviewBackend, + store: CodexReviewStore, + jobID: String = "job-1" +) async throws { + try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + let barrierPublished = await waitUntil(timeout: .seconds(2)) { + store.reviewRecoveryWaitingJobIDs.contains(jobID) + } + guard barrierPublished else { + throw FakeCodexReviewBackendError(message: "Store did not publish connection recovery barrier.") + } +} + @MainActor private func cleanupStoreCommandTest( backend: FakeCodexReviewBackend, store: CodexReviewStore ) async { + await backend.releaseHeldOperationsForCleanup() await backend.finishEventMailboxes() await store.cancelAndDrainReviewWorkersForTesting() await backend.finishEventMailboxes() diff --git a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift index 7b4edb62..cfd27511 100644 --- a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift +++ b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift @@ -187,6 +187,97 @@ struct ReviewAttemptProcessorTests { #expect(await admission.currentPhase() == .terminal(.connection(connection))) } + @Test func recoveryAckWaitsForInterruptedTerminalAndReturnsBarrier() async throws { + let (admission, run) = try await makeActiveAdmission() + let requestAccepted = InvocationProbe() + let recoveryCancellation = ReviewCancellation.system(message: "Recover") + + let recovery = Task { + try await admission.interruptForRecovery( + recoveryCancellation, + interrupt: { _, _ in await requestAccepted.record() }, + forceClose: {} + ) + } + await requestAccepted.waitForInvocation() + #expect(await admission.currentPhase() == .interrupting(run)) + try await admission.recordCanonicalTerminal( + .interrupted(.requested(recoveryCancellation)), + for: run + ) + + let barrier = try await recovery.value + #expect(barrier.run == run) + #expect(barrier.terminal == .canonical( + run: run, + terminal: .interrupted(.requested(recoveryCancellation)) + )) + } + + @Test func recoveryOutcomeUnknownConnectionReturnsBarrierDiagnostic() async throws { + let (admission, run) = try await makeActiveAdmission() + let requestFailed = InvocationProbe() + let requestFailure = ReviewInterruptRequestFailure( + outcome: .outcomeUnknown(message: "Response lost") + ) + let connection = ReviewRuntimeCloseFailure.connection("Connection ended") + + let recovery = Task { + try await admission.interruptForRecovery( + .system(message: "Recover"), + interrupt: { _, _ in + await requestFailed.record() + throw requestFailure + }, + forceClose: {} + ) + } + await requestFailed.waitForInvocation() + await admission.recordConnectionTerminal(connection) + + let barrier = try await recovery.value + #expect(barrier.run == run) + #expect(barrier.terminal == .connection(connection)) + #expect(barrier.requestFailure?.outcome == requestFailure.outcome) + #expect(barrier.requestFailure?.secondaryBarrierDiagnostic == connection.localizedDescription) + } + + @Test func recoveryNaturalTerminalSupersedesReplacement() async throws { + let (admission, run) = try await makeActiveAdmission() + let requestAccepted = InvocationProbe() + let recovery = Task { + try await admission.interruptForRecovery( + .system(message: "Recover"), + interrupt: { _, _ in await requestAccepted.record() }, + forceClose: {} + ) + } + await requestAccepted.waitForInvocation() + try await admission.recordCanonicalTerminal(.completed, for: run) + + await #expect(throws: ReviewRecoverySupersededByTerminal( + terminal: .canonical(run: run, terminal: .completed) + )) { + try await recovery.value + } + } + + @Test func recoveryRejectionReturnsAttemptToActive() async throws { + let (admission, run) = try await makeActiveAdmission() + let rejection = ReviewInterruptRequestFailure( + outcome: .rejected(code: -32_000, message: "Not active") + ) + + await #expect(throws: rejection) { + try await admission.interruptForRecovery( + .system(message: "Recover"), + interrupt: { _, _ in throw rejection }, + forceClose: {} + ) + } + #expect(await admission.currentPhase() == .active(run)) + } + @Test func graceExpiryForceClosesOnceAndAwaitsConnectionAndRequestCompletion() async throws { let graceGate = AsyncGate() let (admission, _) = try await makeActiveAdmission( From 879077e11780f3771f3f81cdc0fa78c67ff68ed7 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 19:43:55 +0900 Subject: [PATCH 15/38] docs(review): clarify caller cancellation ownership --- Sources/CodexReview/Store/CodexReviewStoreReviews.swift | 3 +++ 1 file changed, 3 insertions(+) diff --git a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift index d417485d..528abc33 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift @@ -16,6 +16,9 @@ extension CodexReviewStore { request: CodexReviewAPI.Start.Request ) async throws -> CodexReviewAPI.Read.Result { let jobID = try beginReview(sessionID: sessionID, request: request) + // Caller Task cancellation is not a review-cancellation command: only the + // attempt admission can distinguish not-sent from outcome-unknown dispatch. + // Session owners must use cancelReview/closeSession so the canonical barrier drains. _ = try await awaitReview(sessionID: sessionID, jobID: jobID) await reviewWorkerTasks[jobID]?.value return try readReview(sessionID: sessionID, jobID: jobID) From 6e335be2ca781d50915b89ad18c3b94e1ce6afe8 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 20:47:07 +0900 Subject: [PATCH 16/38] fix(review): type attempt recovery disposition --- Sources/CodexReview/CodexReviewBackend.swift | 43 +- .../CodexReview/ReviewAttemptProcessor.swift | 601 +++++++++++++++--- .../Store/CodexReviewStoreBackend.swift | 8 +- .../Store/CodexReviewStoreReviews.swift | 225 +++---- .../PreviewCodexReviewStoreBackend.swift | 8 +- .../AppServerCodexReviewBackend.swift | 123 +++- .../AppServerProcessTransport.swift | 42 +- Sources/CodexReviewAppServer/JSONRPC.swift | 18 + Sources/CodexReviewHost/CodexReviewHost.swift | 12 +- .../LiveCodexReviewStoreBackend.swift | 12 +- Sources/CodexReviewTesting/TestSupport.swift | 57 +- .../AppServerClientTests.swift | 186 ++++-- .../CurrentV2ReviewContractTests.swift | 36 +- .../CodexReviewStoreCommandTests.swift | 51 +- .../ReviewAttemptProcessorTests.swift | 308 +++++++-- 15 files changed, 1287 insertions(+), 443 deletions(-) diff --git a/Sources/CodexReview/CodexReviewBackend.swift b/Sources/CodexReview/CodexReviewBackend.swift index 887b8212..36f3dcf1 100644 --- a/Sources/CodexReview/CodexReviewBackend.swift +++ b/Sources/CodexReview/CodexReviewBackend.swift @@ -16,11 +16,11 @@ package protocol CodexReviewBackend: Sendable { ) async throws -> BackendReviewAttempt func interruptReview(_ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason) async throws func forceCloseReviewConnection() async throws - func beginReviewRecovery( - _ barrier: ReviewAttemptRecoveryBarrier - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken + func prepareReviewRecovery( + _ candidate: ReviewRecoveryCandidate + ) async throws -> ReviewRecoveryHandoff func resumeReviewRecovery( - _ token: CodexReviewBackendModel.Review.RecoveryToken, + _ handoff: ReviewRecoveryHandoff, request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt @@ -32,10 +32,11 @@ package extension CodexReviewBackend { _ request: CodexReviewBackendModel.Review.Start ) async throws -> BackendReviewAttempt { let admission = ReviewStartAdmission() - let task = await admission.start { admission in + let registered = try await admission.registerStart { admission in try await self.startReview(request, admission: admission) } - return try await task.value + try await admission.activateStart(registered.id) + return try await registered.task.value } } @@ -52,15 +53,13 @@ package struct BackendReviewAttempt: Sendable { package actor BackendReviewEventMailbox { private enum Terminal { case finished - case cancelled - case failed(String) + case failed(ReviewAttemptStreamFailure) } private enum Delivery { case event(CodexReviewBackendModel.Review.Event) case finished - case cancelled - case failed(String) + case failed(ReviewAttemptStreamFailure) } private var bufferedEvents: [CodexReviewBackendModel.Review.Event] = [] @@ -75,10 +74,8 @@ package actor BackendReviewEventMailbox { return event case .finished: return nil - case .cancelled: - throw CancellationError() - case .failed(let message): - throw BackendReviewEventMailboxError(message: message) + case .failed(let failure): + throw failure } } @@ -112,11 +109,11 @@ package actor BackendReviewEventMailbox { resumeWaitersForTerminal() } - package func fail(_ error: any Error) { + package func fail(_ failure: ReviewAttemptStreamFailure) { guard terminal == nil else { return } - terminal = error is CancellationError ? .cancelled : .failed(error.localizedDescription) + terminal = .failed(failure) resumeWaitersForTerminal() } @@ -182,8 +179,6 @@ package actor BackendReviewEventMailbox { switch terminal { case .finished: return .finished - case .cancelled: - return .cancelled case .failed(let message): return .failed(message) } @@ -199,18 +194,6 @@ package actor BackendReviewEventMailbox { } } -package struct BackendReviewEventMailboxError: LocalizedError, Sendable { - package var message: String - - package init(message: String) { - self.message = message - } - - package var errorDescription: String? { - message - } -} - package struct CodexReviewClock: Sendable { package var now: @Sendable () -> Date diff --git a/Sources/CodexReview/ReviewAttemptProcessor.swift b/Sources/CodexReview/ReviewAttemptProcessor.swift index 112b22d8..b47f5921 100644 --- a/Sources/CodexReview/ReviewAttemptProcessor.swift +++ b/Sources/CodexReview/ReviewAttemptProcessor.swift @@ -86,19 +86,81 @@ package struct ReviewStartCancelledBeforeDispatch: LocalizedError, Equatable, Se package var errorDescription: String? { cancellation.message } } +package struct ReviewStartTerminatedBeforeActivation: LocalizedError, Equatable, Sendable { + package let terminal: ReviewAttemptBarrierTerminal + + package init(terminal: ReviewAttemptBarrierTerminal) { + self.terminal = terminal + } + + package var errorDescription: String? { + "Review start was resolved before activation: \(terminal.diagnosticDescription)." + } +} + +package enum ReviewAttemptRecoveryTrigger: Equatable, Sendable { + case sameAccountRestart + case recoverableNetworkLoss + + fileprivate var cancellation: ReviewCancellation { + switch self { + case .sameAccountRestart: + .system(message: "Review runtime is restarting.") + case .recoverableNetworkLoss: + .system(message: "Network unavailable; waiting to reconnect.") + } + } +} + +package enum ReviewAttemptInterruptionPurpose: Equatable, Sendable { + case terminalCancellation(ReviewCancellation) + case recoverableTransition(ReviewAttemptRecoveryTrigger) +} + +package enum ReviewAttemptStreamFailure: LocalizedError, Equatable, Sendable { + case recoverableNetwork(ReviewRuntimeCloseFailure) + case ownerForcedConnectionClose(ReviewRuntimeCloseFailure) + case unexpectedConnection(ReviewRuntimeCloseFailure) + case process(ReviewRuntimeCloseFailure) + case protocolViolation(ReviewAttemptContractFailure) + case workerContract(ReviewAttemptContractFailure) + case ownerCancellation + + package var errorDescription: String? { + switch self { + case .recoverableNetwork(let failure), + .ownerForcedConnectionClose(let failure), + .unexpectedConnection(let failure), + .process(let failure): + failure.localizedDescription + case .protocolViolation(let failure), .workerContract(let failure): + failure.localizedDescription + case .ownerCancellation: + "Review event owner was cancelled before an attempt terminal." + } + } + + package var permitsRecoveryReplacement: Bool { + switch self { + case .recoverableNetwork, .ownerForcedConnectionClose: + true + case .unexpectedConnection, .process, .protocolViolation, + .workerContract, .ownerCancellation: + false + } + } +} + package enum ReviewAttemptBarrierTerminal: Equatable, Sendable { - case canonical( - run: CodexReviewBackendModel.Review.Run, - terminal: ReviewTerminalRecord - ) - case connection(ReviewRuntimeCloseFailure) + case canonical(ReviewTerminalRecord) + case stream(ReviewAttemptStreamFailure) case localCancellation(ReviewCancellation) package var diagnosticDescription: String { switch self { - case .canonical(let run, let terminal): - "canonical terminal \(terminal.kind.rawValue) for attempt \(run.attemptID)" - case .connection(let failure): + case .canonical(let terminal): + "canonical terminal \(terminal.kind.rawValue)" + case .stream(let failure): failure.localizedDescription case .localCancellation(let cancellation): cancellation.message @@ -119,37 +181,120 @@ package struct ReviewAttemptCancellationResolution: Equatable, Sendable { } } -package struct ReviewAttemptRecoveryBarrier: Equatable, Sendable { - package var run: CodexReviewBackendModel.Review.Run - package var terminal: ReviewAttemptBarrierTerminal - package var cancellation: ReviewCancellation - package var requestFailure: ReviewInterruptRequestFailure? +package struct ReviewResolvedAttemptTerminal: Equatable, Sendable { + package let run: CodexReviewBackendModel.Review.Run + package let terminal: ReviewAttemptBarrierTerminal + package let requestFailure: ReviewInterruptRequestFailure? fileprivate init( run: CodexReviewBackendModel.Review.Run, terminal: ReviewAttemptBarrierTerminal, - cancellation: ReviewCancellation, requestFailure: ReviewInterruptRequestFailure? ) { self.run = run self.terminal = terminal - self.cancellation = cancellation self.requestFailure = requestFailure } } -package struct ReviewRecoverySupersededByTerminal: LocalizedError, Equatable, Sendable { - package var terminal: ReviewAttemptBarrierTerminal +package struct ReviewRecoveryCandidate: Equatable, Sendable { + package let resolved: ReviewResolvedAttemptTerminal + package let trigger: ReviewAttemptRecoveryTrigger - package init(terminal: ReviewAttemptBarrierTerminal) { - self.terminal = terminal + fileprivate init( + resolved: ReviewResolvedAttemptTerminal, + trigger: ReviewAttemptRecoveryTrigger + ) { + self.resolved = resolved + self.trigger = trigger } +} - package var errorDescription: String? { - "Recovery was superseded by \(terminal.diagnosticDescription)." +package struct ReviewProductTerminalDisposition: Equatable, Sendable { + package let resolved: ReviewResolvedAttemptTerminal + package let productTerminal: ReviewTerminalRecord + + fileprivate init( + resolved: ReviewResolvedAttemptTerminal, + productTerminal: ReviewTerminalRecord + ) { + self.resolved = resolved + self.productTerminal = productTerminal + } +} + +package enum ReviewRecoveryDisposition: Equatable, Sendable { + case productTerminal(ReviewProductTerminalDisposition) + case replacement(ReviewRecoveryCandidate) +} + +package struct ReviewRecoveryHandoff: Equatable, Sendable { + package let candidate: ReviewRecoveryCandidate + package let token: CodexReviewBackendModel.Review.RecoveryToken + + package init( + candidate: ReviewRecoveryCandidate, + token: CodexReviewBackendModel.Review.RecoveryToken + ) { + self.candidate = candidate + self.token = token + } +} + +package struct ReviewActiveAttempt: Sendable { + package let run: CodexReviewBackendModel.Review.Run + package let admission: ReviewStartAdmission + + package init( + run: CodexReviewBackendModel.Review.Run, + admission: ReviewStartAdmission + ) { + self.run = run + self.admission = admission + } +} + +package struct ReviewStartHandleID: Hashable, Sendable { + package let generation: UInt64 + + package init(generation: UInt64) { + self.generation = generation } } +package struct ReviewRegisteredStart: Sendable { + package let id: ReviewStartHandleID + package let admission: ReviewStartAdmission + package let task: Task + + package init( + id: ReviewStartHandleID, + admission: ReviewStartAdmission, + task: Task + ) { + self.id = id + self.admission = admission + self.task = task + } +} + +package enum ReviewAttemptOwnership: Sendable { + case initialStart(ReviewRegisteredStart) + case active(ReviewActiveAttempt) + case resolvingRecovery(ReviewActiveAttempt) + case recoveryDisposition(ReviewRecoveryDisposition) + case preparingRecovery( + candidate: ReviewRecoveryCandidate, + preparationTask: Task + ) + case waitingForRecovery(ReviewRecoveryHandoff) + case replacementStart( + handoff: ReviewRecoveryHandoff, + start: ReviewRegisteredStart + ) + case terminal +} + /// Owns one review attempt from the first dispatch admission through terminal and cleanup. /// Store cancellation and backend request dispatch both consult this actor; there is no /// call-site startup-cancellation mirror. @@ -161,6 +306,8 @@ package actor ReviewStartAdmission { package enum Phase: Equatable, Sendable { case queued + case registeredStart(ReviewStartHandleID) + case activatedStart(ReviewStartHandleID) case preparingThread(RequestDispatch) case startingReview( preparedRun: CodexReviewBackendModel.Review.Run, @@ -182,8 +329,16 @@ package actor ReviewStartAdmission { private let closePolicy: ReviewRuntimeClosePolicy private var phase: Phase = .queued private var requestedCancellation: ReviewCancellation? + private var joinedTerminalCancellation: ReviewCancellation? + private var interruptionPurpose: ReviewAttemptInterruptionPurpose? private var startTask: Task? + private var nextStartGeneration: UInt64 = 0 + private var registeredStartID: ReviewStartHandleID? + private var startActivationResult: Result? + private var startActivationWaiters: [CheckedContinuation, Never>] = [] private var cancellationTask: Task? + private var recoveryDispositionTask: Task? + private var installedRecoveryDisposition: ReviewRecoveryDisposition? private var interruptRequestTask: Task? private var terminalBarrierTask: Task? private var graceTask: Task? @@ -199,22 +354,35 @@ package actor ReviewStartAdmission { private var terminalWaiters: [UUID: CheckedContinuation] = [:] private var activeRunWaiters: [CheckedContinuation] = [] private var cancellationAdmissionWaiters: [CheckedContinuation] = [] + private var interruptionAdmissionWaiters: [CheckedContinuation] = [] private var cancellationWaiters: [CheckedContinuation, Never>] = [] package init(closePolicy: ReviewRuntimeClosePolicy = .production) { self.closePolicy = closePolicy } - package func start( + package func registerStart( _ operation: @escaping @Sendable (ReviewStartAdmission) async throws -> BackendReviewAttempt - ) -> Task { - precondition( - startTask == nil, - "ReviewStartAdmission owns exactly one registered start Task." - ) - phase = .preparingThread(.notSent) + ) throws -> ReviewRegisteredStart { + guard startTask == nil else { + throw ReviewAttemptContractFailure( + message: "ReviewStartAdmission already owns a registered start Task." + ) + } + nextStartGeneration &+= 1 + let id = ReviewStartHandleID(generation: nextStartGeneration) + registeredStartID = id + if let terminal { + startActivationResult = .failure(startFailure(for: terminal)) + } else if let requestedCancellation { + receiveTerminal(.localCancellation(requestedCancellation)) + } else { + phase = .registeredStart(id) + } let task = Task { do { + try await self.waitForStartActivation(id) + try self.beginActivatedStart(id) let attempt = try await operation(self) self.finishStart(with: .success(attempt)) return attempt @@ -224,7 +392,25 @@ package actor ReviewStartAdmission { } } startTask = task - return task + return ReviewRegisteredStart(id: id, admission: self, task: task) + } + + package func activateStart(_ id: ReviewStartHandleID) throws { + guard registeredStartID == id else { + throw ReviewAttemptContractFailure( + message: "Start activation handle \(id.generation) is stale or belongs to another attempt." + ) + } + if let startActivationResult { + return try startActivationResult.get() + } + guard case .registeredStart(id) = phase else { + throw ReviewAttemptContractFailure( + message: "Start activation handle \(id.generation) is not pending." + ) + } + phase = .activatedStart(id) + resolveStartActivation(.success(())) } package func admitThreadStartDispatch() -> Bool { @@ -235,7 +421,8 @@ package actor ReviewStartAdmission { case .preparingThread(.notSent): phase = .preparingThread(.outcomeUnknown) return true - case .queued, .preparingThread(.outcomeUnknown), .startingReview, + case .queued, .registeredStart, .activatedStart, + .preparingThread(.outcomeUnknown), .startingReview, .active, .interrupting, .finishing, .terminal: return false } @@ -293,8 +480,13 @@ package actor ReviewStartAdmission { _ terminalRecord: ReviewTerminalRecord, for run: CodexReviewBackendModel.Review.Run ) throws { + guard let canonicalRun = registeredRun ?? canonicalRunForTerminal, + Self.matchesCanonicalPair(run, canonicalRun) + else { + return + } if let terminal { - let candidate = ReviewAttemptBarrierTerminal.canonical(run: run, terminal: terminalRecord) + let candidate = ReviewAttemptBarrierTerminal.canonical(terminalRecord) guard terminal == candidate else { throw ReviewAttemptContractFailure( message: "Conflicting terminal for review attempt \(run.attemptID)." @@ -302,22 +494,23 @@ package actor ReviewStartAdmission { } return } - guard let canonicalRun = canonicalRunForTerminal, - Self.matchesCanonicalPair(run, canonicalRun) - else { - return - } - receiveTerminal(.canonical(run: run, terminal: terminalRecord)) + receiveTerminal(.canonical(terminalRecord)) } - package func recordConnectionTerminal(_ failure: ReviewRuntimeCloseFailure) { - guard terminal == nil else { + package func recordStreamTerminal(_ failure: ReviewAttemptStreamFailure) throws { + let candidate = ReviewAttemptBarrierTerminal.stream(failure) + if let terminal { + guard terminal == candidate else { + throw ReviewAttemptContractFailure( + message: "Conflicting stream terminal for the review attempt." + ) + } return } if Self.isOutcomeUnknownStartPhase(phase) { startTask?.cancel() } - receiveTerminal(.connection(failure)) + receiveTerminal(candidate) } package func cancel( @@ -328,6 +521,15 @@ package actor ReviewStartAdmission { ) async throws -> Void, forceClose: @escaping @Sendable () async throws -> Void ) async throws -> ReviewAttemptCancellationResolution { + if let recoveryDispositionTask { + joinedTerminalCancellation = cancellation + requestedCancellation = cancellation + resumeCancellationAdmissionWaiters(returning: cancellation) + let disposition = try await recoveryDispositionTask.value + return try cancellationResolution(for: disposition) + } + interruptionPurpose = .terminalCancellation(cancellation) + resumeInterruptionAdmissionWaiters(returning: interruptionPurpose) let resolution = try await joinedCancellationResolution( cancellation, interrupt: interrupt, @@ -335,48 +537,169 @@ package actor ReviewStartAdmission { ) if let requestFailure = resolution.requestFailure, case .outcomeUnknown = requestFailure.outcome, - case .connection(let connectionFailure) = resolution.terminal { + case .stream(let streamFailure) = resolution.terminal { throw ReviewInterruptRequestFailure( outcome: requestFailure.outcome, - secondaryBarrierDiagnostic: connectionFailure.localizedDescription + secondaryBarrierDiagnostic: streamFailure.localizedDescription ) } return resolution } - package func interruptForRecovery( - _ cancellation: ReviewCancellation, + package func beginRecovery( + trigger: ReviewAttemptRecoveryTrigger, interrupt: @escaping @Sendable ( CodexReviewBackendModel.Review.Run, CodexReviewBackendModel.CancellationReason ) async throws -> Void, forceClose: @escaping @Sendable () async throws -> Void - ) async throws -> ReviewAttemptRecoveryBarrier { - guard let run = activeRun ?? canonicalRunForTerminal ?? registeredRun else { + ) async throws -> ReviewRecoveryDisposition { + if let recoveryDispositionTask { + return try await recoveryDispositionTask.value + } + guard cancellationTask == nil else { + throw ReviewAttemptContractFailure( + message: "Recovery cannot replace an admitted terminal cancellation." + ) + } + guard activeRun != nil else { throw ReviewAttemptContractFailure( message: "Recovery interruption requires one canonical review run." ) } + interruptionPurpose = .recoverableTransition(trigger) + resumeInterruptionAdmissionWaiters(returning: interruptionPurpose) + let task = Task { + try await self.performRecovery( + trigger: trigger, + interrupt: interrupt, + forceClose: forceClose + ) + } + recoveryDispositionTask = task + do { + return try await task.value + } catch { + recoveryDispositionTask = nil + interruptionPurpose = nil + throw error + } + } + + private func performRecovery( + trigger: ReviewAttemptRecoveryTrigger, + interrupt: @escaping @Sendable ( + CodexReviewBackendModel.Review.Run, + CodexReviewBackendModel.CancellationReason + ) async throws -> Void, + forceClose: @escaping @Sendable () async throws -> Void + ) async throws -> ReviewRecoveryDisposition { + guard let run = activeRun else { + throw ReviewAttemptContractFailure( + message: "Recovery interruption lost its active review run." + ) + } let resolution = try await joinedCancellationResolution( - cancellation, + trigger.cancellation, interrupt: interrupt, forceClose: forceClose ) - switch resolution.terminal { - case .canonical(_, let terminal) where terminal.kind != .interrupted: - throw ReviewRecoverySupersededByTerminal(terminal: resolution.terminal) - case .canonical, .connection: - return .init( - run: run, - terminal: resolution.terminal, - cancellation: cancellation, - requestFailure: resolution.requestFailure - ) - case .localCancellation: - throw ReviewAttemptContractFailure( - message: "A dispatched review recovery cannot complete locally." + let resolved = ReviewResolvedAttemptTerminal( + run: run, + terminal: resolution.terminal, + requestFailure: resolution.requestFailure + ) + let disposition = makeRecoveryDisposition(resolved, trigger: trigger) + installedRecoveryDisposition = disposition + return disposition + } + + private func makeRecoveryDisposition( + _ resolved: ReviewResolvedAttemptTerminal, + trigger: ReviewAttemptRecoveryTrigger + ) -> ReviewRecoveryDisposition { + switch resolved.terminal { + case .canonical(let terminal): + switch terminal { + case .completed, .failed: + return .productTerminal(.init( + resolved: resolved, + productTerminal: terminal + )) + case .interrupted: + if let joinedTerminalCancellation { + return .productTerminal(.init( + resolved: resolved, + productTerminal: .interrupted(.requested(joinedTerminalCancellation)) + )) + } + return .replacement(.init(resolved: resolved, trigger: trigger)) + } + case .stream(let failure): + if let joinedTerminalCancellation { + if resolved.requestFailure == nil { + return .productTerminal(.init( + resolved: resolved, + productTerminal: .interrupted(.requested(joinedTerminalCancellation)) + )) + } + return .productTerminal(.init( + resolved: resolved, + productTerminal: productTerminal(for: failure) + )) + } + if failure.permitsRecoveryReplacement { + return .replacement(.init(resolved: resolved, trigger: trigger)) + } + return .productTerminal(.init( + resolved: resolved, + productTerminal: productTerminal(for: failure) + )) + case .localCancellation(let cancellation): + return .productTerminal(.init( + resolved: resolved, + productTerminal: .interrupted(.requested(cancellation)) + )) + } + } + + private func productTerminal( + for failure: ReviewAttemptStreamFailure + ) -> ReviewTerminalRecord { + switch failure { + case .process: + .interrupted(.previousProcessExit) + case .protocolViolation(let failure), .workerContract(let failure): + .failed(message: failure.localizedDescription) + case .ownerCancellation: + .failed(message: failure.localizedDescription) + case .recoverableNetwork, .ownerForcedConnectionClose, + .unexpectedConnection: + .interrupted(.transport(message: failure.localizedDescription)) + } + } + + private func cancellationResolution( + for disposition: ReviewRecoveryDisposition + ) throws -> ReviewAttemptCancellationResolution { + let resolved: ReviewResolvedAttemptTerminal = switch disposition { + case .productTerminal(let product): + product.resolved + case .replacement(let candidate): + candidate.resolved + } + if let requestFailure = resolved.requestFailure, + case .outcomeUnknown = requestFailure.outcome, + case .stream(let failure) = resolved.terminal { + throw ReviewInterruptRequestFailure( + outcome: requestFailure.outcome, + secondaryBarrierDiagnostic: failure.localizedDescription ) } + return .init( + terminal: resolved.terminal, + requestFailure: resolved.requestFailure + ) } private func joinedCancellationResolution( @@ -393,8 +716,10 @@ package actor ReviewStartAdmission { if let terminal { return .init(terminal: terminal) } - requestedCancellation = cancellation - resumeCancellationAdmissionWaiters(returning: cancellation) + if joinedTerminalCancellation == nil { + requestedCancellation = cancellation + resumeCancellationAdmissionWaiters(returning: cancellation) + } let task = Task { try await self.performCancellation( cancellation, @@ -448,7 +773,11 @@ package actor ReviewStartAdmission { } package func waitForCancellationAdmission() async -> ReviewCancellation? { - if let requestedCancellation { + if let terminalCancellation = joinedTerminalCancellation + ?? terminalCancellationPurpose { + return terminalCancellation + } + if recoveryDispositionTask == nil, let requestedCancellation { return requestedCancellation } if terminal != nil { @@ -458,7 +787,10 @@ package actor ReviewStartAdmission { return nil } return await withCheckedContinuation { continuation in - if let requestedCancellation { + if let terminalCancellation = joinedTerminalCancellation + ?? terminalCancellationPurpose { + continuation.resume(returning: terminalCancellation) + } else if recoveryDispositionTask == nil, let requestedCancellation { continuation.resume(returning: requestedCancellation) } else if terminal != nil || startFailed { continuation.resume(returning: nil) @@ -468,6 +800,24 @@ package actor ReviewStartAdmission { } } + package func waitForInterruptionAdmission() async -> ReviewAttemptInterruptionPurpose? { + if let interruptionPurpose { + return interruptionPurpose + } + if terminal != nil || startFailed { + return nil + } + return await withCheckedContinuation { continuation in + if let interruptionPurpose { + continuation.resume(returning: interruptionPurpose) + } else if terminal != nil || startFailed { + continuation.resume(returning: nil) + } else { + interruptionAdmissionWaiters.append(continuation) + } + } + } + package func recordedCleanupResult( for run: CodexReviewBackendModel.Review.Run ) async -> Result? { @@ -477,6 +827,70 @@ package actor ReviewStartAdmission { return await cleanupTask.result } + private func waitForStartActivation(_ id: ReviewStartHandleID) async throws { + guard registeredStartID == id else { + throw ReviewAttemptContractFailure( + message: "Start handle \(id.generation) became stale before activation." + ) + } + if let startActivationResult { + return try startActivationResult.get() + } + let result = await withCheckedContinuation { continuation in + if let startActivationResult { + continuation.resume(returning: startActivationResult) + } else { + startActivationWaiters.append(continuation) + } + } + try result.get() + } + + private func beginActivatedStart(_ id: ReviewStartHandleID) throws { + guard registeredStartID == id else { + throw ReviewAttemptContractFailure( + message: "Start handle \(id.generation) became stale before dispatch." + ) + } + if let terminal { + throw startFailure(for: terminal) + } + if let requestedCancellation { + let terminal = ReviewAttemptBarrierTerminal.localCancellation(requestedCancellation) + receiveTerminal(terminal) + throw startFailure(for: terminal) + } + guard case .activatedStart(id) = phase else { + throw ReviewAttemptContractFailure( + message: "Start handle \(id.generation) was not activated for dispatch." + ) + } + phase = .preparingThread(.notSent) + } + + private func resolveStartActivation(_ result: Result) { + guard startActivationResult == nil else { + return + } + startActivationResult = result + let waiters = startActivationWaiters + startActivationWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume(returning: result) + } + } + + private func startFailure(for terminal: ReviewAttemptBarrierTerminal) -> any Error { + switch terminal { + case .localCancellation(let cancellation): + ReviewStartCancelledBeforeDispatch(cancellation: cancellation) + case .stream(let failure): + failure + case .canonical: + ReviewStartTerminatedBeforeActivation(terminal: terminal) + } + } + private func finishStart( with result: Result ) { @@ -492,12 +906,18 @@ package actor ReviewStartAdmission { if let cancellation = (error as? ReviewStartCancelledBeforeDispatch)?.cancellation { receiveTerminal(.localCancellation(cancellation)) } else if error is CancellationError, - let requestedCancellation, - case .preparingThread(.notSent) = phase { - receiveTerminal(.localCancellation(requestedCancellation)) + let requestedCancellation { + switch phase { + case .registeredStart, .activatedStart, .preparingThread(.notSent): + receiveTerminal(.localCancellation(requestedCancellation)) + case .queued, .preparingThread(.outcomeUnknown), .startingReview, + .active, .interrupting, .finishing, .terminal: + break + } } resumeActiveRunWaiters(returning: nil) resumeCancellationAdmissionWaiters(returning: nil) + resumeInterruptionAdmissionWaiters(returning: nil) } } @@ -513,6 +933,10 @@ package actor ReviewStartAdmission { startTask?.cancel() } else if case .queued = phase { receiveTerminal(.localCancellation(cancellation)) + } else if case .registeredStart = phase { + receiveTerminal(.localCancellation(cancellation)) + } else if case .activatedStart = phase { + receiveTerminal(.localCancellation(cancellation)) } else if Self.isOutcomeUnknownStartPhase(phase) { installGraceTask(forceClose: forceClose) } @@ -719,6 +1143,9 @@ package actor ReviewStartAdmission { private func receiveTerminal(_ terminal: ReviewAttemptBarrierTerminal) { self.terminal = terminal + if registeredStartID != nil, startActivationResult == nil { + resolveStartActivation(.failure(startFailure(for: terminal))) + } phase = cancellationTask == nil ? .terminal(terminal) : .finishing(terminal) let waiters = Array(terminalWaiters.values) terminalWaiters.removeAll(keepingCapacity: false) @@ -727,6 +1154,7 @@ package actor ReviewStartAdmission { } resumeActiveRunWaiters(returning: nil) resumeCancellationAdmissionWaiters(returning: nil) + resumeInterruptionAdmissionWaiters(returning: nil) resolveCancellationIfPossible() } @@ -750,6 +1178,23 @@ package actor ReviewStartAdmission { } } + private func resumeInterruptionAdmissionWaiters( + returning purpose: ReviewAttemptInterruptionPurpose? + ) { + let waiters = interruptionAdmissionWaiters + interruptionAdmissionWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume(returning: purpose) + } + } + + private var terminalCancellationPurpose: ReviewCancellation? { + guard case .terminalCancellation(let cancellation) = interruptionPurpose else { + return nil + } + return cancellation + } + private func resolveCancellationIfPossible() { guard cancellationResult == nil else { return @@ -790,12 +1235,12 @@ package actor ReviewStartAdmission { return } switch terminal { - case .connection(let connectionFailure): + case .stream(let streamFailure): resolveCancellation(.success(.init( terminal: terminal, requestFailure: ReviewInterruptRequestFailure( outcome: requestFailure.outcome, - secondaryBarrierDiagnostic: connectionFailure.localizedDescription + secondaryBarrierDiagnostic: streamFailure.localizedDescription ) ))) case .canonical, .localCancellation: @@ -861,14 +1306,10 @@ package actor ReviewStartAdmission { switch phase { case .startingReview(let run, _), .active(let run), .interrupting(let run): run - case .finishing(let terminal), .terminal(let terminal): - if case .canonical(let run, _) = terminal { - run - } else { - nil - } - case .queued, .preparingThread: - nil + case .finishing, .terminal: + registeredRun + case .queued, .registeredStart, .activatedStart, .preparingThread: + registeredRun } } @@ -876,7 +1317,8 @@ package actor ReviewStartAdmission { switch phase { case .active(let run), .interrupting(let run): run - case .queued, .preparingThread, .startingReview, .finishing, .terminal: + case .queued, .registeredStart, .activatedStart, .preparingThread, + .startingReview, .finishing, .terminal: nil } } @@ -885,7 +1327,8 @@ package actor ReviewStartAdmission { switch phase { case .preparingThread(.outcomeUnknown), .startingReview(_, .outcomeUnknown): true - case .queued, .preparingThread(.notSent), .startingReview(_, .notSent), + case .queued, .registeredStart, .activatedStart, + .preparingThread(.notSent), .startingReview(_, .notSent), .active, .interrupting, .finishing, .terminal: false } diff --git a/Sources/CodexReview/Store/CodexReviewStoreBackend.swift b/Sources/CodexReview/Store/CodexReviewStoreBackend.swift index 5ce3b8ec..1d6d75df 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreBackend.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreBackend.swift @@ -50,11 +50,11 @@ package protocol CodexReviewStoreBackend: CodexReviewSettingsBackend, Sendable { ) async throws -> BackendReviewAttempt func interruptReview(_ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason) async throws func forceCloseReviewConnection() async throws - func beginReviewRecovery( - _ barrier: ReviewAttemptRecoveryBarrier - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken + func prepareReviewRecovery( + _ candidate: ReviewRecoveryCandidate + ) async throws -> ReviewRecoveryHandoff func resumeReviewRecovery( - _ token: CodexReviewBackendModel.Review.RecoveryToken, + _ handoff: ReviewRecoveryHandoff, request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt diff --git a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift index 528abc33..b27fe1cc 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift @@ -124,10 +124,11 @@ extension CodexReviewStore { var run: CodexReviewBackendModel.Review.Run? do { let backend = self.backend - let startTask = await admission.start { admission in + let registered = try await admission.registerStart { admission in try await backend.startReview(startRequest, admission: admission) } - let backendAttempt = try await startTask.value + try await admission.activateStart(registered.id) + let backendAttempt = try await registered.task.value let backendRun = backendAttempt.run run = backendRun applyBackendRun(backendRun, to: job) @@ -170,7 +171,11 @@ extension CodexReviewStore { let failure = ReviewRuntimeCloseFailure.worker( "Review worker was cancelled before a canonical terminal." ) - await cleanupAdmission.recordConnectionTerminal(failure) + do { + try await cleanupAdmission.recordStreamTerminal(.ownerCancellation) + } catch { + markReviewFailed(job, message: error.localizedDescription) + } do { try await cleanupReview(cleanupRun, admission: cleanupAdmission) } catch { @@ -196,13 +201,16 @@ extension CodexReviewStore { activeRuns.removeValue(forKey: jobID) reviewRecoveryWaitingJobIDs.remove(jobID) if job.isTerminal == false, - let transportFailure = error as? ReviewWorkerInputQueueError { - let failure = ReviewRuntimeCloseFailure.connection(transportFailure.message) + let streamFailure = error as? ReviewAttemptStreamFailure { let currentAdmission = reviewStartAdmissions[jobID] ?? admission - await currentAdmission.recordConnectionTerminal(failure) + do { + try await currentAdmission.recordStreamTerminal(streamFailure) + } catch { + markReviewFailed(job, message: error.localizedDescription) + } markReviewInterrupted( job, - cause: .transport(message: transportFailure.message) + cause: .transport(message: streamFailure.localizedDescription) ) } else if job.isTerminal == false { markReviewFailed(job, message: error.localizedDescription) @@ -623,9 +631,7 @@ extension CodexReviewStore { } catch { if error is CancellationError || Task.isCancelled { let currentAdmission = reviewStartAdmissions[job.id] ?? admission - await currentAdmission.recordConnectionTerminal(.worker( - "Review worker stopped while recovery interruption was pending." - )) + try await currentAdmission.recordStreamTerminal(.ownerCancellation) } await inputs.cancel() throw error @@ -679,10 +685,10 @@ extension CodexReviewStore { continue } if recoveryState.isInterruptingForNetworkRecovery { - let failure = ReviewRuntimeCloseFailure.connection( - ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription - ) - await admission.recordConnectionTerminal(failure) + let failure = ReviewAttemptStreamFailure.workerContract(.init( + message: ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription + )) + try await admission.recordStreamTerminal(failure) activeEventSubscriptionID = nil continue } @@ -702,24 +708,19 @@ extension CodexReviewStore { else { continue } - if failedRun.failure.isCancellation { - throw CancellationError() - } if recoveryState.isInterruptingForNetworkRecovery { - let failure = ReviewRuntimeCloseFailure.connection( - failedRun.failure.message - ) - await admission.recordConnectionTerminal(failure) + try await admission.recordStreamTerminal(failedRun.failure) activeEventSubscriptionID = nil continue } - if await inputs.networkStatusTracker.currentStatus() != .satisfied { + if case .unexpectedConnection = failedRun.failure, + await inputs.networkStatusTracker.currentStatus() != .satisfied { recoveryState.recordPendingOutageStreamFailure(failedRun.failure) activeEventSubscriptionID = nil await inputs.cancelActiveEventSubscription() continue } - try throwReviewEventStreamFailure(failedRun.failure) + throw failedRun.failure case .recoveryBarrierResolved(let resolution): guard recoveryState.isInterruptingForNetworkRecovery, resolution.run.attemptID == recoveryState.currentRun.attemptID @@ -728,26 +729,35 @@ extension CodexReviewStore { } switch resolution.result { case .failure(let failure): - if failure.underlying is ReviewRecoverySupersededByTerminal { - continue - } throw failure.underlying - case .success(let ready): - if job.isTerminal || completePendingCancellationIfNeeded(for: job) { + case .success(let disposition): + if job.isTerminal { + return .init(run: recoveryState.currentRun, admission: admission) + } + let candidate: ReviewRecoveryCandidate + switch disposition { + case .productTerminal(let product): + applyRecoveryProductTerminal(product.productTerminal, to: job) + return .init(run: recoveryState.currentRun, admission: admission) + case .replacement(let replacement): + candidate = replacement + } + if completePendingCancellationIfNeeded(for: job) { return .init(run: recoveryState.currentRun, admission: admission) } + let handoff = try await self.backend.prepareReviewRecovery(candidate) recoveryState.markWaitingForNetworkRecovery() markReviewWaitingForNetworkRecovery(job) recordReviewRecoveryBarrier(for: job.id) activeEventSubscriptionID = nil await inputs.cancelActiveEventSubscription() - recoveryState.markRecoveryReady(ready) + recoveryState.markRecoveryReady(handoff) } case .networkSnapshot(let snapshot, let recoveryGeneration): if let pendingFailure = recoveryState.takePendingOutageStreamFailureAfterTransientRecovery( snapshot ) { - try throwReviewEventStreamFailure(pendingFailure) + throw pendingFailure } switch recoveryState.networkSnapshotEffect(snapshot, recoveryGeneration: recoveryGeneration) { case .none: @@ -794,13 +804,10 @@ extension CodexReviewStore { let recoveryRun = recoveryState.currentRun let recoveryAdmission = admission let pendingFailure = recoveryState.takePendingOutageStreamFailureForConfirmedRecovery() - let recoveryCancellation = ReviewCancellation.system( - message: recoveryState.recoveryReason.message - ) let backend = self.backend await inputs.beginRecoveryInterruption(for: recoveryRun) { - let barrier = try await recoveryAdmission.interruptForRecovery( - recoveryCancellation, + try await recoveryAdmission.beginRecovery( + trigger: .recoverableNetworkLoss, interrupt: { run, reason in try await backend.interruptReview(run, reason: reason) }, @@ -808,14 +815,17 @@ extension CodexReviewStore { try await backend.forceCloseReviewConnection() } ) - let token = try await backend.beginReviewRecovery(barrier) - return .init(barrier: barrier, token: token) } if let pendingFailure { - _ = await recoveryAdmission.waitForCancellationAdmission() - await recoveryAdmission.recordConnectionTerminal(.connection( - pendingFailure.message - )) + _ = await recoveryAdmission.waitForInterruptionAdmission() + let failure: ReviewAttemptStreamFailure + switch pendingFailure { + case .unexpectedConnection(let closeFailure): + failure = .recoverableNetwork(closeFailure) + default: + failure = pendingFailure + } + try await recoveryAdmission.recordStreamTerminal(failure) } } } @@ -824,10 +834,10 @@ extension CodexReviewStore { throw CancellationError() } if job.isTerminal == false { - let failure = ReviewRuntimeCloseFailure.connection( - ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription - ) - await admission.recordConnectionTerminal(failure) + let failure = ReviewAttemptStreamFailure.workerContract(.init( + message: ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription + )) + try await admission.recordStreamTerminal(failure) markReviewInterrupted(job, cause: .transport(message: failure.localizedDescription)) } return .init(run: recoveryState.currentRun, admission: admission) @@ -847,21 +857,36 @@ extension CodexReviewStore { } if job.isTerminal == false { - let failure = ReviewRuntimeCloseFailure.connection( - ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription - ) - await admission.recordConnectionTerminal(failure) + let failure = ReviewAttemptStreamFailure.workerContract(.init( + message: ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription + )) + do { + try await admission.recordStreamTerminal(failure) + } catch { + markReviewFailed(job, message: error.localizedDescription) + return true + } markReviewInterrupted(job, cause: .transport(message: failure.localizedDescription)) } return true } - private func throwReviewEventStreamFailure(_ failure: ReviewWorkerEventStreamFailure) throws -> Never { - switch failure { - case .cancelled: - throw CancellationError() + private func applyRecoveryProductTerminal( + _ terminal: ReviewTerminalRecord, + to job: CodexReviewJob + ) { + switch terminal { + case .completed: + if job.isTerminal == false { + markReviewFailed( + job, + message: "Canonical completion was not reduced before recovery disposition." + ) + } case .failed(let message): - throw ReviewWorkerInputQueueError(message: message) + markReviewFailed(job, message: message, terminal: terminal) + case .interrupted(let cause): + markReviewInterrupted(job, cause: cause) } } @@ -869,7 +894,7 @@ extension CodexReviewStore { job: CodexReviewJob, startRequest: CodexReviewBackendModel.Review.Start, inputs: ReviewWorkerInputs, - recoveryReady: ReviewWorkerRecoveryReady? + recoveryReady: ReviewRecoveryHandoff? ) async throws -> NetworkRestoreRestartResult { if job.isTerminal || completePendingCancellationIfNeeded(for: job) { return .finished @@ -890,14 +915,15 @@ extension CodexReviewStore { reviewStartAdmissions[job.id] = recoveredAdmission reviewRecoveryWaitingJobIDs.remove(job.id) let backend = self.backend - let startTask = await recoveredAdmission.start { admission in + let registered = try await recoveredAdmission.registerStart { admission in try await backend.resumeReviewRecovery( - recoveryReady.token, + recoveryReady, request: startRequest, admission: admission ) } - let recoveredAttempt = try await startTask.value + try await recoveredAdmission.activateStart(registered.id) + let recoveredAttempt = try await registered.task.value return .recovered(recoveredAttempt, recoveredAdmission) } @@ -1285,30 +1311,7 @@ private struct ReviewWorkerEventStreamFinished: Sendable { private struct ReviewWorkerEventStreamFailed: Sendable { var subscriptionID: Int var run: CodexReviewBackendModel.Review.Run - var failure: ReviewWorkerEventStreamFailure -} - -private enum ReviewWorkerEventStreamFailure: Sendable { - case cancelled - case failed(String) - - var isCancellation: Bool { - switch self { - case .cancelled: - true - case .failed: - false - } - } - - var message: String { - switch self { - case .cancelled: - "Review event stream was cancelled." - case .failed(let message): - message - } - } + var failure: ReviewAttemptStreamFailure } private enum ReviewWorkerInput: Sendable { @@ -1323,12 +1326,7 @@ private enum ReviewWorkerInput: Sendable { private struct ReviewWorkerRecoveryBarrierResolution: Sendable { var run: CodexReviewBackendModel.Review.Run - var result: Result -} - -private struct ReviewWorkerRecoveryReady: Sendable { - var barrier: ReviewAttemptRecoveryBarrier - var token: CodexReviewBackendModel.Review.RecoveryToken + var result: Result } private struct ReviewWorkerAttemptCompletion: Sendable { @@ -1364,10 +1362,10 @@ private enum ReviewNetworkRecoveryPhase { private struct ReviewNetworkRecoveryLoopState { var currentRun: CodexReviewBackendModel.Review.Run private(set) var recoveryPhase = ReviewNetworkRecoveryPhase.active - private(set) var recoveryReady: ReviewWorkerRecoveryReady? + private(set) var recoveryReady: ReviewRecoveryHandoff? private var isSettlingForNetworkRecovery = false private var recoverySettleGeneration: Int? - private var pendingOutageStreamFailure: ReviewWorkerEventStreamFailure? + private var pendingOutageStreamFailure: ReviewAttemptStreamFailure? let recoveryReason = CodexReviewBackendModel.CancellationReason(message: networkRecoveryUnavailableMessage) init(currentRun: CodexReviewBackendModel.Review.Run) { @@ -1395,7 +1393,7 @@ private struct ReviewNetworkRecoveryLoopState { pendingOutageStreamFailure = nil } - mutating func markRecoveryReady(_ ready: ReviewWorkerRecoveryReady) { + mutating func markRecoveryReady(_ ready: ReviewRecoveryHandoff) { recoveryReady = ready } @@ -1408,11 +1406,11 @@ private struct ReviewNetworkRecoveryLoopState { pendingOutageStreamFailure = nil } - mutating func recordPendingOutageStreamFailure(_ failure: ReviewWorkerEventStreamFailure) { + mutating func recordPendingOutageStreamFailure(_ failure: ReviewAttemptStreamFailure) { pendingOutageStreamFailure = failure } - mutating func takePendingOutageStreamFailureForConfirmedRecovery() -> ReviewWorkerEventStreamFailure? { + mutating func takePendingOutageStreamFailureForConfirmedRecovery() -> ReviewAttemptStreamFailure? { defer { pendingOutageStreamFailure = nil } @@ -1421,7 +1419,7 @@ private struct ReviewNetworkRecoveryLoopState { mutating func takePendingOutageStreamFailureAfterTransientRecovery( _ snapshot: CodexReviewNetworkSnapshot - ) -> ReviewWorkerEventStreamFailure? { + ) -> ReviewAttemptStreamFailure? { guard snapshot.status == .satisfied, isWaitingForNetworkRecovery == false else { @@ -1493,7 +1491,7 @@ private struct ReviewWorkerInputs { func beginRecoveryInterruption( for run: CodexReviewBackendModel.Review.Run, - operation: @escaping @Sendable () async throws -> ReviewWorkerRecoveryReady + operation: @escaping @Sendable () async throws -> ReviewRecoveryDisposition ) async { await recoveryInterruptionSource.start(for: run, operation: operation) } @@ -1518,13 +1516,13 @@ private actor ReviewWorkerRecoveryInterruptionSource { func start( for run: CodexReviewBackendModel.Review.Run, - operation: @escaping @Sendable () async throws -> ReviewWorkerRecoveryReady + operation: @escaping @Sendable () async throws -> ReviewRecoveryDisposition ) { guard task == nil else { return } task = Task { - let result: Result + let result: Result do { result = .success(try await operation()) } catch { @@ -1631,14 +1629,6 @@ private actor ReviewWorkerInputQueue { } } -private struct ReviewWorkerInputQueueError: LocalizedError, Sendable { - var message: String - - var errorDescription: String? { - message - } -} - private actor ReviewWorkerEventSource { private let queue: ReviewWorkerInputQueue private var eventTasks: [Int: Task] = [:] @@ -1663,10 +1653,22 @@ private actor ReviewWorkerEventSource { return } await self.yieldReviewEvent(event, run: run, subscriptionID: subscriptionID) + if event.completesReviewRun { + self.finishTerminalDelivery(subscriptionID: subscriptionID) + return + } } await self.yieldEventsFinished(run: run, subscriptionID: subscriptionID) } catch { - await self.yieldEventsFailed(error, run: run, subscriptionID: subscriptionID) + let failure: ReviewAttemptStreamFailure + if let typed = error as? ReviewAttemptStreamFailure { + failure = typed + } else if error is CancellationError { + failure = .ownerCancellation + } else { + failure = .workerContract(.init(message: error.localizedDescription)) + } + await self.yieldEventsFailed(failure, run: run, subscriptionID: subscriptionID) } } return subscriptionID @@ -1729,8 +1731,15 @@ private actor ReviewWorkerEventSource { ))) } + private func finishTerminalDelivery(subscriptionID: Int) { + guard activeSubscriptionID == subscriptionID else { + return + } + eventTasks.removeValue(forKey: subscriptionID) + } + private func yieldEventsFailed( - _ error: any Error, + _ failure: ReviewAttemptStreamFailure, run: CodexReviewBackendModel.Review.Run, subscriptionID: Int ) async { @@ -1743,7 +1752,7 @@ private actor ReviewWorkerEventSource { await queue.send(.reviewEventsFailed(.init( subscriptionID: subscriptionID, run: run, - failure: error is CancellationError ? .cancelled : .failed(error.localizedDescription) + failure: failure ))) } } diff --git a/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift b/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift index b6fb4f36..c4c58067 100644 --- a/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift +++ b/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift @@ -144,14 +144,14 @@ package class PreviewCodexReviewStoreBackend: CodexReviewStoreBackend { package func forceCloseReviewConnection() async throws {} - package func beginReviewRecovery( - _: ReviewAttemptRecoveryBarrier - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { + package func prepareReviewRecovery( + _: ReviewRecoveryCandidate + ) async throws -> ReviewRecoveryHandoff { throw CodexReviewAPI.Error.io(Self.previewUnavailableMessage) } package func resumeReviewRecovery( - _: CodexReviewBackendModel.Review.RecoveryToken, + _: ReviewRecoveryHandoff, request _: CodexReviewBackendModel.Review.Start, admission _: ReviewStartAdmission ) async throws -> BackendReviewAttempt { diff --git a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift index 2a621e08..b670e848 100644 --- a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift +++ b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift @@ -36,6 +36,7 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { private var notificationRouterMetrics = AppServerNotificationRouterMetrics() private var reviewStartRequestsInFlight = 0 private var diagnosedUnknownNotificationMethods: Set = [] + private var connectionStreamFailure: ReviewAttemptStreamFailure? package init( client: AppServerClient, @@ -168,8 +169,8 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } catch { reviewStartRequestsInFlight -= 1 discardUnmatchedReviewNotificationsIfIdle() - if let terminal = Self.connectionTerminal(for: error) { - await admission.recordConnectionTerminal(terminal) + if let terminal = streamTerminal(for: error) { + try await admission.recordStreamTerminal(terminal) } try await cleanupReview(provisionalRun) throw error @@ -269,8 +270,8 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { do { return try await client.send(request) } catch { - if let terminal = Self.connectionTerminal(for: error) { - await admission.recordConnectionTerminal(terminal) + if let terminal = streamTerminal(for: error) { + try await admission.recordStreamTerminal(terminal) } throw error } @@ -328,17 +329,28 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { return .init(outcome: .outcomeUnknown(message: error.localizedDescription)) } - private nonisolated static func connectionTerminal( + private func streamTerminal( for error: any Error - ) -> ReviewRuntimeCloseFailure? { + ) -> ReviewAttemptStreamFailure? { guard let jsonRPCError = error as? JSONRPC.Error else { return nil } switch jsonRPCError { - case .closed, .invalidMessage: - return .connection(jsonRPCError.localizedDescription) + case .closed: + return connectionStreamFailure + ?? .unexpectedConnection(.connection(jsonRPCError.localizedDescription)) + case .invalidMessage: + return .protocolViolation(.init(message: jsonRPCError.localizedDescription)) case .responseError: return nil + case .transportTerminated(let termination): + switch termination { + case .ownerClose: + return connectionStreamFailure + ?? .unexpectedConnection(.connection(jsonRPCError.localizedDescription)) + case .processExit, .processFailure: + return .process(.process(jsonRPCError.localizedDescription)) + } } } @@ -361,6 +373,9 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } package func forceCloseReviewConnection() async throws { + connectionStreamFailure = .ownerForcedConnectionClose( + .connection("Review connection was force-closed by its attempt owner.") + ) do { try await client.close() } catch let failure as ReviewRuntimeCloseFailure { @@ -370,10 +385,10 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } } - package func beginReviewRecovery( - _ barrier: ReviewAttemptRecoveryBarrier - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { - let run = barrier.run + package func prepareReviewRecovery( + _ candidate: ReviewRecoveryCandidate + ) async throws -> ReviewRecoveryHandoff { + let run = candidate.resolved.run let interruption = AppServerReviewInterruption( threadID: appServerTurnThreadID(for: run), turnID: run.turnID ?? "" @@ -387,19 +402,23 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { completedReviewEventSessionMetricsByThreadID[threadID] = metrics } } - return CodexReviewBackendModel.Review.RecoveryToken( - interruptedRun: run, - rollbackThreadID: interruption.threadID + return ReviewRecoveryHandoff( + candidate: candidate, + token: .init( + interruptedRun: run, + rollbackThreadID: interruption.threadID + ) ) } package func resumeReviewRecovery( - _ token: CodexReviewBackendModel.Review.RecoveryToken, + _ handoff: ReviewRecoveryHandoff, request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { _ = try await client.initialize() await ensureNotificationRouterStarted() + let token = handoff.token let interruptedRun = token.interruptedRun let _: EmptyResponse = try await client.send(AppServerAPI.Thread.Rollback.Request( params: .init(threadID: token.rollbackThreadID, numTurns: 1) @@ -438,8 +457,8 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } catch { reviewStartRequestsInFlight -= 1 _ = unregisterReviewEventSession(for: provisionalRun) - if let terminal = Self.connectionTerminal(for: error) { - await admission.recordConnectionTerminal(terminal) + if let terminal = streamTerminal(for: error) { + try await admission.recordStreamTerminal(terminal) } await session.abandon() discardUnmatchedReviewNotificationsIfIdle() @@ -736,9 +755,11 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { for try await notification in notifications { await routeReviewNotification(notification) } - await finishAllReviewEventSessions(throwing: JSONRPC.Error.closed) + await finishAllReviewEventSessions(throwing: .workerContract(.init( + message: "App-server notification stream ended without a transport terminal." + ))) } catch { - await finishAllReviewEventSessions(throwing: error) + await finishAllReviewEventSessions(throwing: streamFailure(for: error)) } notificationRouterTask = nil } @@ -879,6 +900,9 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { private func failConnection(_ error: ReviewIngestionError) async { notificationRouterMetrics.connectionFailures += 1 + connectionStreamFailure = .protocolViolation(.init( + message: error.localizedDescription + )) appServerBackendLogger.error( "Closing app-server connection after review routing failure: \(error.localizedDescription, privacy: .public)" ) @@ -889,7 +913,9 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { "App-server connection close failed during routing containment: \(error.localizedDescription, privacy: .public)" ) } - await finishAllReviewEventSessions(throwing: error) + await finishAllReviewEventSessions(throwing: .protocolViolation(.init( + message: error.localizedDescription + ))) } private func diagnoseUnknownNotificationMethod(_ method: String) { @@ -911,11 +937,54 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { ) } - private func finishAllReviewEventSessions(throwing error: (any Error)?) async { + private func finishAllReviewEventSessions( + throwing failure: ReviewAttemptStreamFailure? + ) async { let sessions = Array(reviewEventSessionsByAttemptID.values) for session in sessions { - await session.finish(throwing: error) + await session.finish(throwing: failure) + } + } + + private func streamFailure( + for error: any Error + ) -> ReviewAttemptStreamFailure { + if let failure = error as? ReviewAttemptStreamFailure { + return failure + } + if let failure = error as? ReviewRuntimeCloseFailure { + switch failure { + case .process: + return .process(failure) + case .connection: + return .unexpectedConnection(failure) + case .worker, .cleanup, .mcpHandlerDrain: + return .workerContract(.init(message: failure.localizedDescription)) + } + } + if let jsonRPCError = error as? JSONRPC.Error { + switch jsonRPCError { + case .invalidMessage: + return .protocolViolation(.init(message: jsonRPCError.localizedDescription)) + case .closed: + return connectionStreamFailure + ?? .unexpectedConnection(.connection(jsonRPCError.localizedDescription)) + case .responseError: + return .workerContract(.init(message: jsonRPCError.localizedDescription)) + case .transportTerminated(let termination): + switch termination { + case .ownerClose: + return connectionStreamFailure + ?? .unexpectedConnection(.connection(jsonRPCError.localizedDescription)) + case .processExit, .processFailure: + return .process(.process(jsonRPCError.localizedDescription)) + } + } + } + if error is CancellationError { + return .ownerCancellation } + return .workerContract(.init(message: error.localizedDescription)) } private func readModelCatalog() async throws -> [CodexReviewSettings.ModelCatalogItem] { @@ -1152,7 +1221,7 @@ private actor AppServerReviewEventSession { await finish(precedingEvents: precedingEvents, cancellationMessage: cancellationMessage) } - func finish(throwing error: (any Error)?) async { + func finish(throwing failure: ReviewAttemptStreamFailure?) async { guard finished == false else { return } @@ -1162,8 +1231,8 @@ private actor AppServerReviewEventSession { commandLifecycleByItemID.removeAll(keepingCapacity: true) pendingStartupNotifications.removeAll(keepingCapacity: true) await emitPrecedingEvents(precedingEvents) - if let error { - await mailbox.fail(error) + if let failure { + await mailbox.fail(failure) } else { await mailbox.finish() } @@ -1388,7 +1457,7 @@ private actor AppServerReviewEventSession { cancelPendingStreamedLogFlush() pendingStartupNotifications.removeAll(keepingCapacity: true) await emitPrecedingEvents(precedingEvents) - await emitTerminal(.failed(message: error.localizedDescription)) + await mailbox.fail(.protocolViolation(.init(message: error.localizedDescription))) finished = true } diff --git a/Sources/CodexReviewAppServer/AppServerProcessTransport.swift b/Sources/CodexReviewAppServer/AppServerProcessTransport.swift index 6d8ac9c4..66173a50 100644 --- a/Sources/CodexReviewAppServer/AppServerProcessTransport.swift +++ b/Sources/CodexReviewAppServer/AppServerProcessTransport.swift @@ -70,6 +70,7 @@ package actor AppServerProcessTransport: JSONRPC.Transport { private var stderrReaderTask: Task? = nil private var closed = false private var closeTask: Task? + private var terminalError: JSONRPC.Error? package init( configuration: Configuration = .init(), @@ -141,11 +142,14 @@ package actor AppServerProcessTransport: JSONRPC.Transport { try stdin.fileHandleForWriting.write(contentsOf: payload) } - package func notificationStream() -> AsyncThrowingStream { + package func notificationStream() async -> AsyncThrowingStream { ensureReaderTasksStarted() + if let closeTask { + _ = await closeTask.result + } return AsyncThrowingStream(bufferingPolicy: .unbounded) { continuation in if closed { - continuation.finish(throwing: JSONRPC.Error.closed) + continuation.finish(throwing: terminalError ?? JSONRPC.Error.closed) return } let id = UUID() @@ -166,14 +170,14 @@ package actor AppServerProcessTransport: JSONRPC.Transport { ) async throws { try await closeTransport( terminateProcess: terminateProcess, - error: JSONRPC.Error.closed, + error: JSONRPC.Error.transportTerminated(.ownerClose), readerTask: readerTask ) } private func closeTransport( terminateProcess: Bool, - error: any Error, + error: JSONRPC.Error, readerTask: ReaderTask? ) async throws { let task: Task @@ -181,6 +185,7 @@ package actor AppServerProcessTransport: JSONRPC.Transport { task = closeTask } else { closed = true + terminalError = error stdoutEvents.cancel() stderrEvents.cancel() try? stdin.fileHandleForWriting.close() @@ -201,7 +206,7 @@ package actor AppServerProcessTransport: JSONRPC.Transport { private func performCloseTransport( terminateProcess: Bool, - error: any Error + error: JSONRPC.Error ) async throws { var processCloseError: (any Error)? if terminateProcess { @@ -213,7 +218,16 @@ package actor AppServerProcessTransport: JSONRPC.Transport { processCloseError = error } } - finishAll(throwing: processCloseError ?? error) + if let processCloseError { + let typedError = JSONRPC.Error.transportTerminated(.processFailure( + processCloseError.localizedDescription + )) + terminalError = typedError + finishAll(throwing: typedError) + } else { + terminalError = error + finishAll(throwing: error) + } if let processCloseError { throw processCloseError } @@ -238,7 +252,8 @@ package actor AppServerProcessTransport: JSONRPC.Transport { do { try await closeTransport( terminateProcess: true, - error: error, + error: error as? JSONRPC.Error + ?? .invalidMessage(error.localizedDescription), readerTask: .stdout ) } catch { @@ -280,7 +295,8 @@ package actor AppServerProcessTransport: JSONRPC.Transport { do { try await closeTransport( terminateProcess: true, - error: error, + error: error as? JSONRPC.Error + ?? .invalidMessage(error.localizedDescription), readerTask: .stdout ) } catch { @@ -290,7 +306,13 @@ package actor AppServerProcessTransport: JSONRPC.Transport { } } do { - try await closeTransport(terminateProcess: true, readerTask: .stdout) + try await closeTransport( + terminateProcess: true, + error: .transportTerminated(.processExit( + "Codex app-server process exited after stdout reached EOF." + )), + readerTask: .stdout + ) } catch { logger.error("Failed to close codex app-server process after stdout EOF: \(error.localizedDescription, privacy: .public)") } @@ -452,7 +474,7 @@ package actor AppServerProcessTransport: JSONRPC.Transport { private func throwIfClosed() throws { if closed { - throw JSONRPC.Error.closed + throw terminalError ?? JSONRPC.Error.closed } } } diff --git a/Sources/CodexReviewAppServer/JSONRPC.swift b/Sources/CodexReviewAppServer/JSONRPC.swift index 6c0ff932..559c323a 100644 --- a/Sources/CodexReviewAppServer/JSONRPC.swift +++ b/Sources/CodexReviewAppServer/JSONRPC.swift @@ -34,6 +34,7 @@ package enum JSONRPC { case closed case invalidMessage(String) case responseError(code: Int, message: String) + case transportTerminated(TransportTermination) package var errorDescription: String? { switch self { @@ -43,6 +44,23 @@ package enum JSONRPC { "Invalid JSON-RPC message: \(message)" case .responseError(_, let message): message + case .transportTerminated(let termination): + termination.localizedDescription + } + } + } + + package enum TransportTermination: Equatable, Sendable, LocalizedError { + case ownerClose + case processExit(String) + case processFailure(String) + + package var errorDescription: String? { + switch self { + case .ownerClose: + "JSON-RPC transport was closed by its owner." + case .processExit(let message), .processFailure(let message): + message } } } diff --git a/Sources/CodexReviewHost/CodexReviewHost.swift b/Sources/CodexReviewHost/CodexReviewHost.swift index 54a3d6e7..2702790c 100644 --- a/Sources/CodexReviewHost/CodexReviewHost.swift +++ b/Sources/CodexReviewHost/CodexReviewHost.swift @@ -251,18 +251,18 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { try await backend.forceCloseReviewConnection() } - func beginReviewRecovery( - _ barrier: ReviewAttemptRecoveryBarrier - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { - try await backend.beginReviewRecovery(barrier) + func prepareReviewRecovery( + _ candidate: ReviewRecoveryCandidate + ) async throws -> ReviewRecoveryHandoff { + try await backend.prepareReviewRecovery(candidate) } func resumeReviewRecovery( - _ token: CodexReviewBackendModel.Review.RecoveryToken, + _ handoff: ReviewRecoveryHandoff, request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { - try await backend.resumeReviewRecovery(token, request: request, admission: admission) + try await backend.resumeReviewRecovery(handoff, request: request, admission: admission) } func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { diff --git a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift index a13a8d02..fd165f31 100644 --- a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift +++ b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift @@ -1014,17 +1014,17 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { try await appServerBackend.forceCloseReviewConnection() } - func beginReviewRecovery( - _ barrier: ReviewAttemptRecoveryBarrier - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { + func prepareReviewRecovery( + _ candidate: ReviewRecoveryCandidate + ) async throws -> ReviewRecoveryHandoff { guard let appServerBackend else { throw CodexReviewAPI.Error.io("Review runtime is not running.") } - return try await appServerBackend.beginReviewRecovery(barrier) + return try await appServerBackend.prepareReviewRecovery(candidate) } func resumeReviewRecovery( - _ token: CodexReviewBackendModel.Review.RecoveryToken, + _ handoff: ReviewRecoveryHandoff, request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { @@ -1032,7 +1032,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { throw CodexReviewAPI.Error.io("Review runtime is not running.") } return try await appServerBackend.resumeReviewRecovery( - token, + handoff, request: request, admission: admission ) diff --git a/Sources/CodexReviewTesting/TestSupport.swift b/Sources/CodexReviewTesting/TestSupport.swift index 98b1f1e7..29d914ab 100644 --- a/Sources/CodexReviewTesting/TestSupport.swift +++ b/Sources/CodexReviewTesting/TestSupport.swift @@ -183,8 +183,8 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { case startReview(CodexReviewBackendModel.Review.Start) case interruptReview(CodexReviewBackendModel.Review.Run, CodexReviewBackendModel.CancellationReason) case forceCloseReviewConnection - case beginReviewRecovery(CodexReviewBackendModel.Review.Run, CodexReviewBackendModel.CancellationReason) - case resumeReviewRecovery(CodexReviewBackendModel.Review.RecoveryToken, CodexReviewBackendModel.Review.Start) + case prepareReviewRecovery(ReviewRecoveryCandidate) + case resumeReviewRecovery(ReviewRecoveryHandoff, CodexReviewBackendModel.Review.Start) case cleanupReview(CodexReviewBackendModel.Review.Run) } @@ -379,7 +379,7 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { package func waitForBeginReviewRecovery() async { if commands.contains(where: { - if case .beginReviewRecovery = $0 { + if case .prepareReviewRecovery = $0 { true } else { false @@ -391,7 +391,7 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { await withTaskCancellationHandler { await withCheckedContinuation { continuation in if commands.contains(where: { - if case .beginReviewRecovery = $0 { + if case .prepareReviewRecovery = $0 { true } else { false @@ -562,7 +562,9 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { await interruptReviewGate?.open() let mailboxes = Array(eventMailboxes.values) for mailbox in mailboxes { - await mailbox.fail(ReviewRuntimeCloseFailure.connection("Connection force-closed.")) + await mailbox.fail(.ownerForcedConnectionClose( + .connection("Connection force-closed.") + )) } } @@ -572,27 +574,31 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { await resumeReviewRecoveryGate?.open() } - package func beginReviewRecovery( - _ barrier: ReviewAttemptRecoveryBarrier - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { - commands.append(.beginReviewRecovery( - barrier.run, - .init(message: barrier.cancellation.message) - )) + package func prepareReviewRecovery( + _ candidate: ReviewRecoveryCandidate + ) async throws -> ReviewRecoveryHandoff { + commands.append(.prepareReviewRecovery(candidate)) let waiters = Array(beginReviewRecoveryWaiters.values) beginReviewRecoveryWaiters.removeAll(keepingCapacity: false) for waiter in waiters { waiter.resume() } - let run = barrier.run - return .init(interruptedRun: run, rollbackThreadID: run.reviewThreadID ?? run.threadID) + let run = candidate.resolved.run + return .init( + candidate: candidate, + token: .init( + interruptedRun: run, + rollbackThreadID: run.reviewThreadID ?? run.threadID + ) + ) } package func resumeReviewRecovery( - _ token: CodexReviewBackendModel.Review.RecoveryToken, + _ handoff: ReviewRecoveryHandoff, request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { + let token = handoff.token let run = token.interruptedRun let recoveredRun = nextRecoveredRun ?? .init( attemptID: "attempt-recovered", @@ -613,7 +619,7 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { cancellation: await admission.cancellationRequest() ?? .system() ) } - commands.append(.resumeReviewRecovery(token, request)) + commands.append(.resumeReviewRecovery(handoff, request)) let waiters = Array(resumeReviewRecoveryWaiters.values) resumeReviewRecoveryWaiters.removeAll(keepingCapacity: false) for waiter in waiters { @@ -645,8 +651,11 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { await eventMailbox(for: run ?? nextRun).finish() } - package func finishEvents(throwing error: any Error, for run: CodexReviewBackendModel.Review.Run? = nil) async { - await eventMailbox(for: run ?? nextRun).fail(error) + package func finishEvents( + throwing failure: ReviewAttemptStreamFailure, + for run: CodexReviewBackendModel.Review.Run? = nil + ) async { + await eventMailbox(for: run ?? nextRun).fail(failure) } package func finishEventMailboxes() async { @@ -948,19 +957,19 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { try await reviewBackend.forceCloseReviewConnection() } - package func beginReviewRecovery( - _ barrier: ReviewAttemptRecoveryBarrier - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { - try await reviewBackend.beginReviewRecovery(barrier) + package func prepareReviewRecovery( + _ candidate: ReviewRecoveryCandidate + ) async throws -> ReviewRecoveryHandoff { + try await reviewBackend.prepareReviewRecovery(candidate) } package func resumeReviewRecovery( - _ token: CodexReviewBackendModel.Review.RecoveryToken, + _ handoff: ReviewRecoveryHandoff, request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { try await reviewBackend.resumeReviewRecovery( - token, + handoff, request: request, admission: admission ) diff --git a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift index bfbafe53..7462a8b4 100644 --- a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift +++ b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift @@ -11,8 +11,8 @@ private extension AppServerCodexReviewBackend { request: CodexReviewBackendModel.Review.Start, reason: CodexReviewBackendModel.CancellationReason ) async throws -> BackendReviewAttempt { - let token = try await beginReviewRecovery(run, reason: reason) - return try await resumeReviewRecovery(token, request: request) + let handoff = try await prepareReviewRecovery(run, reason: reason) + return try await resumeReviewRecovery(handoff, request: request) } func resumeReviewRecovery( @@ -27,34 +27,35 @@ private extension AppServerCodexReviewBackend { try await interruptReview(attempt.run, reason: reason) } - func beginReviewRecovery( + func prepareReviewRecovery( _ attempt: BackendReviewAttempt, reason: CodexReviewBackendModel.CancellationReason - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { - try await beginReviewRecovery(attempt.run, reason: reason) + ) async throws -> ReviewRecoveryHandoff { + try await prepareReviewRecovery(attempt.run, reason: reason) } - func beginReviewRecovery( + func prepareReviewRecovery( _ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { - let barrier = try await makeRecoveryBarrier(backend: self, for: run, reason: reason) - return try await beginReviewRecovery(barrier) + ) async throws -> ReviewRecoveryHandoff { + let candidate = try await makeRecoveryCandidate(backend: self, for: run, reason: reason) + return try await prepareReviewRecovery(candidate) } func resumeReviewRecovery( - _ token: CodexReviewBackendModel.Review.RecoveryToken, + _ handoff: ReviewRecoveryHandoff, request: CodexReviewBackendModel.Review.Start ) async throws -> BackendReviewAttempt { let admission = ReviewStartAdmission() - let task = await admission.start { admission in + let registered = try await admission.registerStart { admission in try await self.resumeReviewRecovery( - token, + handoff, request: request, admission: admission ) } - return try await task.value + try await admission.activateStart(registered.id) + return try await registered.task.value } func cleanupReview(_ attempt: BackendReviewAttempt) async throws { @@ -62,13 +63,13 @@ private extension AppServerCodexReviewBackend { } } -private func makeRecoveryBarrier( +private func makeRecoveryCandidate( backend: AppServerCodexReviewBackend, for run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason -) async throws -> ReviewAttemptRecoveryBarrier { +) async throws -> ReviewRecoveryCandidate { let admission = ReviewStartAdmission() - let startTask = await admission.start { admission in + let registered = try await admission.registerStart { admission in #expect(await admission.admitThreadStartDispatch()) let provisionalRun = CodexReviewBackendModel.Review.Run( attemptID: run.attemptID, @@ -81,13 +82,13 @@ private func makeRecoveryBarrier( await admission.recordActiveRun(run) return .init(run: run) } - _ = try await startTask.value + try await admission.activateStart(registered.id) + _ = try await registered.task.value let requestOutcome = RecoveryRequestOutcomeProbe() - let cancellation = ReviewCancellation.system(message: reason.message) let recovery = Task { - try await admission.interruptForRecovery( - cancellation, + try await admission.beginRecovery( + trigger: .recoverableNetworkLoss, interrupt: { run, reason in do { try await backend.interruptReview(run, reason: reason) @@ -102,11 +103,16 @@ private func makeRecoveryBarrier( } if case .success = await requestOutcome.wait() { try await admission.recordCanonicalTerminal( - .interrupted(.requested(cancellation)), + .interrupted(.server(message: reason.message)), for: run ) } - return try await recovery.value + guard case .replacement(let candidate) = try await recovery.value else { + throw ReviewAttemptContractFailure( + message: "Test recovery helper expected one replacement candidate." + ) + } + return candidate } private extension BackendReviewAttempt { @@ -452,7 +458,7 @@ struct AppServerClientTests { let notifications = await transport.notificationStream() var iterator = notifications.makeAsyncIterator() - await #expect(throws: (any Error).self) { + await #expect(throws: JSONRPC.Error.transportTerminated(.ownerClose)) { _ = try await iterator.next() } } @@ -528,6 +534,53 @@ struct AppServerClientTests { try await transport.close() } #expect(await closeCompletions.value() == 1) + + let notifications = await transport.notificationStream() + var iterator = notifications.makeAsyncIterator() + await #expect(throws: JSONRPC.Error.transportTerminated(.processFailure( + failure.localizedDescription + ))) { + _ = try await iterator.next() + } + } + + @Test func spontaneousProcessExitReplaysTypedCauseToLateSubscriber() async throws { + let directory = FileManager.default.temporaryDirectory + .appending(path: "codex-review-process-exit-\(UUID().uuidString)") + try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: directory) } + let executable = directory.appending(path: "app-server-stub.sh") + try Data("#!/bin/sh\nexit 0\n".utf8).write(to: executable) + try FileManager.default.setAttributes([.posixPermissions: 0o755], ofItemAtPath: executable.path) + let transport = try AppServerProcessTransport(configuration: .init( + executable: executable.path, + arguments: [], + environment: ["HOME": directory.path, "PATH": "/bin:/usr/bin"] + )) + + let notifications = await transport.notificationStream() + var iterator = notifications.makeAsyncIterator() + do { + _ = try await iterator.next() + Issue.record("Expected a typed process-exit terminal.") + } catch let error as JSONRPC.Error { + guard case .transportTerminated(.processExit) = error else { + Issue.record("Expected processExit, received \(error).") + return + } + } + + let replay = await transport.notificationStream() + var replayIterator = replay.makeAsyncIterator() + do { + _ = try await replayIterator.next() + Issue.record("Expected the recorded process-exit terminal.") + } catch let error as JSONRPC.Error { + guard case .transportTerminated(.processExit) = error else { + Issue.record("Expected replayed processExit, received \(error).") + return + } + } } @Test func processTransportProcessesChunkedStdoutBeforeEOF() async throws { @@ -604,7 +657,9 @@ struct AppServerClientTests { return } } - await #expect(throws: JSONRPC.Error.closed) { + await #expect(throws: JSONRPC.Error.invalidMessage( + "app-server emitted invalid JSON" + )) { try await transport.notify(.init(method: "initialized", params: Data("{}".utf8))) } } @@ -1220,7 +1275,7 @@ struct AppServerClientTests { let admission = ReviewStartAdmission( closePolicy: controlledReviewClosePolicy(gate: AsyncGate()) ) - let startTask = await admission.start { admission in + let registeredStart = try await admission.registerStart { admission in try await backend.startReview( .init( jobID: "job-1", @@ -1230,6 +1285,7 @@ struct AppServerClientTests { admission: admission ) } + try await admission.activateStart(registeredStart.id) await transport.waitForRequest(method: "initialize") let cancellation = Task { @@ -1245,7 +1301,7 @@ struct AppServerClientTests { .mcpClient(message: "Stop") )) await #expect(throws: ReviewStartCancelledBeforeDispatch.self) { - try await startTask.value + try await registeredStart.task.value } #expect(await transport.recordedRequests().map(\.method) == ["initialize"]) } @@ -1267,7 +1323,7 @@ struct AppServerClientTests { let admission = ReviewStartAdmission( closePolicy: controlledReviewClosePolicy(gate: AsyncGate()) ) - let startTask = await admission.start { admission in + let registeredStart = try await admission.registerStart { admission in try await backend.startReview( .init( jobID: "job-1", @@ -1277,6 +1333,7 @@ struct AppServerClientTests { admission: admission ) } + try await admission.activateStart(registeredStart.id) await transport.waitForRequest(method: "thread/start") let cancellation = Task { @@ -1292,7 +1349,7 @@ struct AppServerClientTests { .mcpClient(message: "Stop") )) await #expect(throws: ReviewStartCancelledBeforeDispatch.self) { - try await startTask.value + try await registeredStart.task.value } let methods = await transport.recordedRequests().map(\.method) #expect(methods.contains("review/start") == false) @@ -1321,7 +1378,7 @@ struct AppServerClientTests { let admission = ReviewStartAdmission( closePolicy: controlledReviewClosePolicy(gate: AsyncGate()) ) - let startTask = await admission.start { admission in + let registeredStart = try await admission.registerStart { admission in try await backend.startReview( .init( jobID: "job-1", @@ -1331,6 +1388,7 @@ struct AppServerClientTests { admission: admission ) } + try await admission.activateStart(registeredStart.id) await transport.waitForRequest(method: "review/start") let cancellation = Task { @@ -1345,7 +1403,7 @@ struct AppServerClientTests { ) } await reviewGate.open() - let attempt = try await startTask.value + let attempt = try await registeredStart.task.value await transport.waitForResponseDelivery(method: "turn/interrupt") try await transport.emitServerNotification( method: "turn/completed", @@ -1361,8 +1419,7 @@ struct AppServerClientTests { let resolution = try await cancellation.value #expect(resolution.terminal == .canonical( - run: attempt.run, - terminal: .interrupted(.requested(.mcpClient(message: "Stop"))) + .interrupted(.requested(.mcpClient(message: "Stop"))) )) let methods = await transport.recordedRequests().map(\.method) let reviewIndex = try #require(methods.firstIndex(of: "review/start")) @@ -1767,12 +1824,55 @@ struct AppServerClientTests { #expect(routerStopped) var iterator = await eventSequence(backend, run).makeAsyncIterator() - await #expect(throws: BackendReviewEventMailboxError.self) { + await #expect(throws: ReviewAttemptStreamFailure.unexpectedConnection( + .connection(JSONRPC.Error.closed.localizedDescription) + )) { _ = try await iterator.next() } try await transport.close() } + @Test func backendMapsTypedProcessTerminationWithoutRecoverableFallback() async throws { + let transport = FakeJSONRPCTransport() + let backend = AppServerCodexReviewBackend(client: .init(transport: transport)) + let run = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-1", + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "thread-1" + ) + var iterator = await eventSequence(backend, run).makeAsyncIterator() + let message = "Codex app-server process exited." + + await transport.finishNotificationStreams( + throwing: JSONRPC.Error.transportTerminated(.processExit(message)) + ) + + await #expect(throws: ReviewAttemptStreamFailure.process(.process(message))) { + _ = try await iterator.next() + } + } + + @Test func backendCorrelatesOwnerForcedCloseWithAdmittedCloseOperation() async throws { + let transport = FakeJSONRPCTransport() + let backend = AppServerCodexReviewBackend(client: .init(transport: transport)) + let run = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-1", + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "thread-1" + ) + var iterator = await eventSequence(backend, run).makeAsyncIterator() + + try await backend.forceCloseReviewConnection() + + await #expect(throws: ReviewAttemptStreamFailure.ownerForcedConnectionClose( + .connection("Review connection was force-closed by its attempt owner.") + )) { + _ = try await iterator.next() + } + } + @Test func backendPreservesBufferedEventsBeforeNotificationStreamError() async throws { let transport = FakeJSONRPCTransport() let backend = AppServerCodexReviewBackend(client: .init(transport: transport)) @@ -1793,7 +1893,7 @@ struct AppServerClientTests { var iterator = events.makeAsyncIterator() #expect(try await iterator.next() == .started(turnID: "turn-1", reviewThreadID: "thread-1", model: nil)) #expect(try await iterator.next() == .messageDelta("partial review", itemID: "message-1")) - await #expect(throws: BackendReviewEventMailboxError.self) { + await #expect(throws: ReviewAttemptStreamFailure.self) { _ = try await iterator.next() } } @@ -2068,10 +2168,10 @@ struct AppServerClientTests { params: TestErrorNotification(message: "App-server failed.", willRetry: false) ) - await #expect(throws: BackendReviewEventMailboxError.self) { + await #expect(throws: ReviewAttemptStreamFailure.self) { _ = try await firstIterator.next() } - await #expect(throws: BackendReviewEventMailboxError.self) { + await #expect(throws: ReviewAttemptStreamFailure.self) { _ = try await secondIterator.next() } #expect(await backend.notificationRouterMetricsForTesting().connectionFailures == 1) @@ -2287,7 +2387,7 @@ struct AppServerClientTests { let currentRun = startedRun let reason = CodexReviewBackendModel.CancellationReason(message: "Network unavailable; waiting to reconnect.") - let token = try await backend.beginReviewRecovery(currentRun, reason: reason) + let token = try await backend.prepareReviewRecovery(currentRun, reason: reason) let recovered = try await backend.resumeReviewRecovery( token, request: .init( @@ -2411,7 +2511,7 @@ struct AppServerClientTests { let initialEvents = await eventSequence(backend, run) defer { withExtendedLifetime(initialEvents) {} } - let token = try await backend.beginReviewRecovery( + let token = try await backend.prepareReviewRecovery( run, reason: .init(message: "Network unavailable; waiting to reconnect.") ) @@ -2472,7 +2572,7 @@ struct AppServerClientTests { let events = await eventSequence(backend, run) var iterator = events.makeAsyncIterator() - _ = try await backend.beginReviewRecovery( + _ = try await backend.prepareReviewRecovery( run, reason: .init(message: "Network unavailable; waiting to reconnect.") ) @@ -2509,7 +2609,7 @@ struct AppServerClientTests { let events = await eventSequence(backend, run) var iterator = events.makeAsyncIterator() - _ = try await backend.beginReviewRecovery( + _ = try await backend.prepareReviewRecovery( run, reason: .init(message: "Network unavailable; waiting to reconnect.") ) @@ -2548,7 +2648,7 @@ struct AppServerClientTests { let events = await eventSequence(backend, run) var iterator = events.makeAsyncIterator() - async let recovery: CodexReviewBackendModel.Review.RecoveryToken = backend.beginReviewRecovery( + async let recovery: ReviewRecoveryHandoff = backend.prepareReviewRecovery( run, reason: .init(message: "Network unavailable; waiting to reconnect.") ) @@ -2606,7 +2706,7 @@ struct AppServerClientTests { code: -32602, message: "expected active turn id turn-old but found turn-active" ))) { - try await backend.beginReviewRecovery( + try await backend.prepareReviewRecovery( run, reason: .init(message: "Network unavailable; waiting to reconnect.") ) @@ -4889,7 +4989,7 @@ struct AppServerClientTests { replacesGroup: false, metadata: .init(sourceType: "commandExecution", title: "Command output", itemID: "cmd-1") )) - await #expect(throws: BackendReviewEventMailboxError.self) { + await #expect(throws: ReviewAttemptStreamFailure.self) { _ = try await iterator.next() } } diff --git a/Tests/CodexReviewAppServerTests/CurrentV2ReviewContractTests.swift b/Tests/CodexReviewAppServerTests/CurrentV2ReviewContractTests.swift index 52aaea7b..98dd7c75 100644 --- a/Tests/CodexReviewAppServerTests/CurrentV2ReviewContractTests.swift +++ b/Tests/CodexReviewAppServerTests/CurrentV2ReviewContractTests.swift @@ -1599,9 +1599,7 @@ struct CurrentV2ReviewRoutingIntegrationTests { try await transport.emitServerNotification(method: "item/completed", params: 1) - await #expect(throws: BackendReviewEventMailboxError.self) { - _ = try await attempt.events.next() - } + await expectProtocolViolation(from: attempt.events) #expect(await backend.notificationRouterMetricsForTesting().connectionFailures == 1) #expect(await transport.isClosedForTesting()) } @@ -1621,9 +1619,7 @@ struct CurrentV2ReviewRoutingIntegrationTests { params: V2WarningNotification(message: "Guardian warning") ) - await #expect(throws: BackendReviewEventMailboxError.self) { - _ = try await attempt.events.next() - } + await expectProtocolViolation(from: attempt.events) #expect(await backend.notificationRouterMetricsForTesting().connectionFailures == 1) #expect(await transport.isClosedForTesting()) } @@ -1643,9 +1639,7 @@ struct CurrentV2ReviewRoutingIntegrationTests { params: V2ContextCompactedNotification(threadID: "thread-review") ) - await #expect(throws: BackendReviewEventMailboxError.self) { - _ = try await attempt.events.next() - } + await expectProtocolViolation(from: attempt.events) #expect(await backend.notificationRouterMetricsForTesting().connectionFailures == 1) #expect(await transport.isClosedForTesting()) } @@ -1772,12 +1766,8 @@ struct CurrentV2ReviewRoutingIntegrationTests { ) ) - await #expect(throws: BackendReviewEventMailboxError.self) { - _ = try await first.events.next() - } - await #expect(throws: BackendReviewEventMailboxError.self) { - _ = try await second.events.next() - } + await expectProtocolViolation(from: first.events) + await expectProtocolViolation(from: second.events) #expect(await backend.notificationRouterMetricsForTesting().connectionFailures == 1) #expect(await transport.isClosedForTesting()) } @@ -1883,6 +1873,22 @@ struct CurrentV2ReviewRoutingIntegrationTests { return events } + private func expectProtocolViolation( + from mailbox: BackendReviewEventMailbox + ) async { + do { + _ = try await mailbox.next() + Issue.record("Expected a typed protocol violation.") + } catch let failure as ReviewAttemptStreamFailure { + guard case .protocolViolation = failure else { + Issue.record("Expected protocolViolation, received \(failure).") + return + } + } catch { + Issue.record("Expected ReviewAttemptStreamFailure, received \(error).") + } + } + private func outputTexts( in events: [CodexReviewBackendModel.Review.Event] ) -> [String] { diff --git a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift index 1ede9aec..6181239e 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift @@ -861,7 +861,7 @@ struct CodexReviewStoreCommandTests { let attemptedRecovery = await waitUntil(timeout: .milliseconds(100)) { let commands = await backend.recordedCommands() return commands.contains { command in - if case .beginReviewRecovery = command { + if case .prepareReviewRecovery = command { true } else { false @@ -871,7 +871,7 @@ struct CodexReviewStoreCommandTests { #expect(attemptedRecovery == false) let commands = await backend.recordedCommands() #expect(commands.contains { command in - if case .beginReviewRecovery = command { + if case .prepareReviewRecovery = command { true } else { false @@ -934,7 +934,7 @@ struct CodexReviewStoreCommandTests { try await backend.waitForInterruptReview(run: run, reason: reason, timeout: .seconds(2)) let commandsBeforeTerminal = await backend.recordedCommands() - #expect(commandsBeforeTerminal.contains { if case .beginReviewRecovery = $0 { true } else { false } } == false) + #expect(commandsBeforeTerminal.contains { if case .prepareReviewRecovery = $0 { true } else { false } } == false) #expect(store.reviewRecoveryWaitingJobIDs.contains("job-1") == false) await backend.yield(.cancelled(reason.message), for: run) @@ -982,7 +982,7 @@ struct CodexReviewStoreCommandTests { await backend.yield(.cancelled(reason.message), for: run) let commandsBeforeAck = await backend.recordedCommands() - #expect(commandsBeforeAck.contains { if case .beginReviewRecovery = $0 { true } else { false } } == false) + #expect(commandsBeforeAck.contains { if case .prepareReviewRecovery = $0 { true } else { false } } == false) await interruptGate.open() try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) let barrierPublished = await waitUntil(timeout: .seconds(2)) { @@ -1026,7 +1026,7 @@ struct CodexReviewStoreCommandTests { #expect(read.core.lifecycle.status == .failed) #expect(read.core.lifecycle.errorMessage == "Recovery rejected") let commands = await backend.recordedCommands() - #expect(commands.contains { if case .beginReviewRecovery = $0 { true } else { false } } == false) + #expect(commands.contains { if case .prepareReviewRecovery = $0 { true } else { false } } == false) #expect(commands.contains { if case .resumeReviewRecovery = $0 { true } else { false } } == false) } } @@ -1061,14 +1061,14 @@ struct CodexReviewStoreCommandTests { ) let commandsBeforeAck = await backend.recordedCommands() - #expect(commandsBeforeAck.contains { if case .beginReviewRecovery = $0 { true } else { false } } == false) + #expect(commandsBeforeAck.contains { if case .prepareReviewRecovery = $0 { true } else { false } } == false) await interruptGate.open() let read = try await result #expect(read.core.lifecycle.status == .succeeded) #expect(read.core.output.lastAgentMessage == "natural review") let commands = await backend.recordedCommands() - #expect(commands.contains { if case .beginReviewRecovery = $0 { true } else { false } } == false) + #expect(commands.contains { if case .prepareReviewRecovery = $0 { true } else { false } } == false) #expect(commands.contains { if case .resumeReviewRecovery = $0 { true } else { false } } == false) } } @@ -1369,8 +1369,8 @@ struct CodexReviewStoreCommandTests { try await completeNetworkRecoveryBarrier(backend: backend, store: store) let commandsAfterInterrupt = await backend.recordedCommands() let interruptedRuns = commandsAfterInterrupt.compactMap { command -> CodexReviewBackendModel.Review.Run? in - if case .beginReviewRecovery(let run, _) = command { - return run + if case .prepareReviewRecovery(let candidate) = command { + return candidate.resolved.run } return nil } @@ -1380,8 +1380,8 @@ struct CodexReviewStoreCommandTests { try await backend.waitForResumeReviewRecovery(timeout: .seconds(2)) let commandsAfterRecovery = await backend.recordedCommands() let recoveredFromRuns = commandsAfterRecovery.compactMap { command -> CodexReviewBackendModel.Review.Run? in - if case .resumeReviewRecovery(let token, _) = command { - return token.interruptedRun + if case .resumeReviewRecovery(let handoff, _) = command { + return handoff.candidate.resolved.run } return nil } @@ -2059,7 +2059,7 @@ struct CodexReviewStoreCommandTests { #expect(read.core.lifecycle.status == .cancelled) let commands = await backend.recordedCommands() #expect(commands.contains { command in - if case .beginReviewRecovery = command { + if case .prepareReviewRecovery = command { true } else { false @@ -2218,7 +2218,9 @@ struct CodexReviewStoreCommandTests { cancellation: .mcpClient(message: "Stop") ) await backend.waitForInterruptReview() - await backend.finishEvents(throwing: StreamClosedError()) + await backend.finishEvents(throwing: .unexpectedConnection( + .connection("Review event stream closed.") + )) let cancel = try await cancellation let read = try await result @@ -2241,7 +2243,9 @@ struct CodexReviewStoreCommandTests { ) try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running, jobID: "job-1") != nil) await backend.yield(.message("partial review")) - await backend.finishEvents(throwing: StreamClosedError()) + await backend.finishEvents(throwing: .unexpectedConnection( + .connection("Review event stream closed.") + )) let read = try await result #expect(read.core.lifecycle.status == .failed) @@ -2291,7 +2295,10 @@ struct CodexReviewStoreCommandTests { networkMonitor.yield(.init(status: .unsatisfied)) await outageSleepStarted.wait() - await backend.finishEvents(throwing: StreamClosedError(), for: initialRun) + await backend.finishEvents( + throwing: .unexpectedConnection(.connection("Review event stream closed.")), + for: initialRun + ) let failedBeforeOutageConfirmed = await StoreSnapshotProbe(store: store) .waitUntilJobStatus(.failed, jobID: "job-1", timeout: .milliseconds(100)) != nil @@ -2326,7 +2333,7 @@ struct CodexReviewStoreCommandTests { request: .init(cwd: "/tmp/project", target: .baseBranch("main")) ) try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running, jobID: "job-1") != nil) - await backend.finishEvents(throwing: CancellationError()) + await backend.finishEvents(throwing: .ownerCancellation) let read = try await result #expect(read.core.lifecycle.status == .failed) @@ -2782,7 +2789,15 @@ private func completeNetworkRecoveryBarrier( try await backend.waitForInterruptReview(run: run, reason: reason, timeout: .seconds(2)) await beforeTerminal(run) await backend.yield(.cancelled(reason.message), for: run) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + do { + try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + } catch { + let phase = await store.reviewStartAdmissions[jobID]?.currentPhase() + let commands = await backend.recordedCommands() + throw FakeCodexReviewBackendError( + message: "Recovery preparation was not admitted; phase=\(String(describing: phase)), commands=\(commands)." + ) + } let barrierPublished = await waitUntil(timeout: .seconds(2)) { store.reviewRecoveryWaitingJobIDs.contains(jobID) } @@ -2817,8 +2832,6 @@ private func cleanupStoreCommandTest( await backend.finishEventMailboxes() } -private struct StreamClosedError: Error {} - private actor ControlledTestSleeper { private let gate: AsyncGate private var shouldBlock = false diff --git a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift index cfd27511..82beee9d 100644 --- a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift +++ b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift @@ -26,8 +26,7 @@ struct ReviewAttemptProcessorTests { let resolution = try await cancellation.value #expect(resolution.terminal == .canonical( - run: run, - terminal: .interrupted(.requested(.mcpClient(message: "Stop"))) + .interrupted(.requested(.mcpClient(message: "Stop"))) )) #expect(await admission.currentPhase() == .terminal(resolution.terminal)) } @@ -50,14 +49,11 @@ struct ReviewAttemptProcessorTests { await requestStarted.waitForInvocation() try await admission.recordCanonicalTerminal(.completed, for: run) - #expect(await admission.currentPhase() == .finishing(.canonical( - run: run, - terminal: .completed - ))) + #expect(await admission.currentPhase() == .finishing(.canonical(.completed))) await requestGate.open() let resolution = try await cancellation.value - #expect(resolution.terminal == .canonical(run: run, terminal: .completed)) + #expect(resolution.terminal == .canonical(.completed)) } @Test func explicitRejectionReturnsAttemptToActive() async throws { @@ -92,8 +88,7 @@ struct ReviewAttemptProcessorTests { ) #expect(try await retry.value.terminal == .canonical( - run: run, - terminal: .interrupted(.requested(.mcpClient(message: "Stop again"))) + .interrupted(.requested(.mcpClient(message: "Stop again"))) )) } @@ -121,7 +116,7 @@ struct ReviewAttemptProcessorTests { await requestGate.open() let resolution = try await cancellation.value - #expect(resolution.terminal == .canonical(run: run, terminal: .completed)) + #expect(resolution.terminal == .canonical(.completed)) #expect(resolution.requestFailure == rejection) } @@ -151,8 +146,7 @@ struct ReviewAttemptProcessorTests { #expect(resolution.requestFailure == failure) #expect(resolution.terminal == .canonical( - run: run, - terminal: .interrupted(.requested(.mcpClient(message: "Stop"))) + .interrupted(.requested(.mcpClient(message: "Stop"))) )) } @@ -175,7 +169,7 @@ struct ReviewAttemptProcessorTests { ) } await requestFailed.waitForInvocation() - await admission.recordConnectionTerminal(connection) + try await admission.recordStreamTerminal(.unexpectedConnection(connection)) do { _ = try await cancellation.value @@ -184,17 +178,15 @@ struct ReviewAttemptProcessorTests { #expect(received.outcome == failure.outcome) #expect(received.secondaryBarrierDiagnostic == connection.localizedDescription) } - #expect(await admission.currentPhase() == .terminal(.connection(connection))) + #expect(await admission.currentPhase() == .terminal(.stream(.unexpectedConnection(connection)))) } @Test func recoveryAckWaitsForInterruptedTerminalAndReturnsBarrier() async throws { let (admission, run) = try await makeActiveAdmission() let requestAccepted = InvocationProbe() - let recoveryCancellation = ReviewCancellation.system(message: "Recover") - let recovery = Task { - try await admission.interruptForRecovery( - recoveryCancellation, + try await admission.beginRecovery( + trigger: .recoverableNetworkLoss, interrupt: { _, _ in await requestAccepted.record() }, forceClose: {} ) @@ -202,15 +194,17 @@ struct ReviewAttemptProcessorTests { await requestAccepted.waitForInvocation() #expect(await admission.currentPhase() == .interrupting(run)) try await admission.recordCanonicalTerminal( - .interrupted(.requested(recoveryCancellation)), + .interrupted(.server(message: "network recovery")), for: run ) - let barrier = try await recovery.value - #expect(barrier.run == run) - #expect(barrier.terminal == .canonical( - run: run, - terminal: .interrupted(.requested(recoveryCancellation)) + guard case .replacement(let candidate) = try await recovery.value else { + Issue.record("Interrupted recovery terminal must create a replacement candidate.") + return + } + #expect(candidate.resolved.run == run) + #expect(candidate.resolved.terminal == .canonical( + .interrupted(.server(message: "network recovery")) )) } @@ -223,8 +217,8 @@ struct ReviewAttemptProcessorTests { let connection = ReviewRuntimeCloseFailure.connection("Connection ended") let recovery = Task { - try await admission.interruptForRecovery( - .system(message: "Recover"), + try await admission.beginRecovery( + trigger: .recoverableNetworkLoss, interrupt: { _, _ in await requestFailed.record() throw requestFailure @@ -233,21 +227,24 @@ struct ReviewAttemptProcessorTests { ) } await requestFailed.waitForInvocation() - await admission.recordConnectionTerminal(connection) + try await admission.recordStreamTerminal(.recoverableNetwork(connection)) - let barrier = try await recovery.value - #expect(barrier.run == run) - #expect(barrier.terminal == .connection(connection)) - #expect(barrier.requestFailure?.outcome == requestFailure.outcome) - #expect(barrier.requestFailure?.secondaryBarrierDiagnostic == connection.localizedDescription) + guard case .replacement(let candidate) = try await recovery.value else { + Issue.record("Recoverable connection terminal must create a replacement candidate.") + return + } + #expect(candidate.resolved.run == run) + #expect(candidate.resolved.terminal == .stream(.recoverableNetwork(connection))) + #expect(candidate.resolved.requestFailure?.outcome == requestFailure.outcome) + #expect(candidate.resolved.requestFailure?.secondaryBarrierDiagnostic == connection.localizedDescription) } @Test func recoveryNaturalTerminalSupersedesReplacement() async throws { let (admission, run) = try await makeActiveAdmission() let requestAccepted = InvocationProbe() let recovery = Task { - try await admission.interruptForRecovery( - .system(message: "Recover"), + try await admission.beginRecovery( + trigger: .recoverableNetworkLoss, interrupt: { _, _ in await requestAccepted.record() }, forceClose: {} ) @@ -255,11 +252,13 @@ struct ReviewAttemptProcessorTests { await requestAccepted.waitForInvocation() try await admission.recordCanonicalTerminal(.completed, for: run) - await #expect(throws: ReviewRecoverySupersededByTerminal( - terminal: .canonical(run: run, terminal: .completed) - )) { - try await recovery.value + guard case .productTerminal(let product) = try await recovery.value else { + Issue.record("Natural terminal must supersede replacement.") + return } + #expect(product.resolved.run == run) + #expect(product.resolved.terminal == .canonical(.completed)) + #expect(product.productTerminal == .completed) } @Test func recoveryRejectionReturnsAttemptToActive() async throws { @@ -269,8 +268,8 @@ struct ReviewAttemptProcessorTests { ) await #expect(throws: rejection) { - try await admission.interruptForRecovery( - .system(message: "Recover"), + try await admission.beginRecovery( + trigger: .recoverableNetworkLoss, interrupt: { _, _ in throw rejection }, forceClose: {} ) @@ -300,7 +299,9 @@ struct ReviewAttemptProcessorTests { }, forceClose: { await forceClose.record() - await admission.recordConnectionTerminal(connection) + try await admission.recordStreamTerminal( + .ownerForcedConnectionClose(connection) + ) await requestGate.open() } ) @@ -313,7 +314,9 @@ struct ReviewAttemptProcessorTests { try await cancellation.value } #expect(await forceClose.invocationCount() == 1) - #expect(await admission.currentPhase() == .terminal(.connection(connection))) + #expect(await admission.currentPhase() == .terminal( + .stream(.ownerForcedConnectionClose(connection)) + )) } @Test func forceCloseFailureRemainsTypedAfterOutcomeUnknownRequestCompletes() async throws { @@ -387,10 +390,7 @@ struct ReviewAttemptProcessorTests { try await cancellation.value } #expect(await requestCancelled.invocationCount() == 1) - #expect(await admission.currentPhase() == .terminal(.canonical( - run: run, - terminal: .completed - ))) + #expect(await admission.currentPhase() == .terminal(.canonical(.completed))) } @Test func duplicateCancellationCallersJoinOneRequest() async throws { @@ -452,7 +452,7 @@ struct ReviewAttemptProcessorTests { #expect(await admission.currentPhase() == .interrupting(run)) try await admission.recordCanonicalTerminal(.completed, for: run) - #expect(try await cancellation.value.terminal == .canonical(run: run, terminal: .completed)) + #expect(try await cancellation.value.terminal == .canonical(.completed)) } @Test func conflictingDuplicateTerminalFailsWithoutRewrite() async throws { @@ -465,10 +465,7 @@ struct ReviewAttemptProcessorTests { for: run ) } - #expect(await admission.currentPhase() == .terminal(.canonical( - run: run, - terminal: .completed - ))) + #expect(await admission.currentPhase() == .terminal(.canonical(.completed))) } @Test func queuedCancellationCompletesLocallyWithoutDispatch() async throws { @@ -483,10 +480,169 @@ struct ReviewAttemptProcessorTests { #expect(resolution.terminal == .localCancellation(.mcpClient(message: "Stop"))) } + @Test func registeredStartIsBackendInertUntilExactActivation() async throws { + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) + let backendWrite = InvocationProbe() + let registered = try await admission.registerStart { _ in + await backendWrite.record() + return .init(run: canonicalRun) + } + + await Task.yield() + #expect(await backendWrite.invocationCount() == 0) + #expect(await admission.currentPhase() == .registeredStart(registered.id)) + + try await admission.activateStart(registered.id) + _ = try await registered.task.value + #expect(await backendWrite.invocationCount() == 1) + } + + @Test func cancellationBetweenRegistrationAndActivationCompletesTaskWithoutWrite() async throws { + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) + let backendWrite = InvocationProbe() + let cancellation = ReviewCancellation.mcpClient(message: "Stop") + let registered = try await admission.registerStart { _ in + await backendWrite.record() + return .init(run: canonicalRun) + } + + let resolution = try await admission.cancel( + cancellation, + interrupt: { _, _ in Issue.record("Unactivated start dispatched interrupt.") }, + forceClose: { Issue.record("Unactivated start force-closed a connection.") } + ) + + #expect(resolution.terminal == .localCancellation(cancellation)) + await #expect(throws: ReviewStartCancelledBeforeDispatch(cancellation: cancellation)) { + try await registered.task.value + } + #expect(await backendWrite.invocationCount() == 0) + } + + @Test func cancellationBeforeRegistrationSurvivesRegistrationWithoutWrite() async throws { + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) + let cancellation = ReviewCancellation.mcpClient(message: "Stop") + _ = try await admission.cancel( + cancellation, + interrupt: { _, _ in Issue.record("Queued cancellation dispatched interrupt.") }, + forceClose: {} + ) + let backendWrite = InvocationProbe() + + let registered = try await admission.registerStart { _ in + await backendWrite.record() + return .init(run: canonicalRun) + } + + await #expect(throws: ReviewStartCancelledBeforeDispatch(cancellation: cancellation)) { + try await registered.task.value + } + #expect(await backendWrite.invocationCount() == 0) + } + + @Test func activationIsIdempotentForLiveHandleAndRejectsWrongHandle() async throws { + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) + let operationGate = AsyncGate() + let registered = try await admission.registerStart { _ in + await operationGate.waitIgnoringCancellation() + return .init(run: canonicalRun) + } + let wrong = ReviewStartHandleID(generation: registered.id.generation + 1) + + await #expect(throws: ReviewAttemptContractFailure.self) { + try await admission.activateStart(wrong) + } + try await admission.activateStart(registered.id) + try await admission.activateStart(registered.id) + await operationGate.open() + _ = try await registered.task.value + } + + @Test func joinedExplicitCancellationInstallsProductDispositionBeforeRecoveryPreparation() async throws { + let (admission, run) = try await makeActiveAdmission() + let requestGate = AsyncGate() + let requestStarted = InvocationProbe() + let recovery = Task { + try await admission.beginRecovery( + trigger: .recoverableNetworkLoss, + interrupt: { _, _ in + await requestStarted.record() + await requestGate.waitIgnoringCancellation() + }, + forceClose: {} + ) + } + await requestStarted.waitForInvocation() + let cancellation = Task { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in Issue.record("Joined cancellation sent a second request.") }, + forceClose: {} + ) + } + #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) + try await admission.recordCanonicalTerminal( + .interrupted(.server(message: "network recovery")), + for: run + ) + await requestGate.open() + + let disposition = try await recovery.value + guard case .productTerminal(let product) = disposition else { + Issue.record("Joined explicit cancellation must suppress replacement.") + return + } + #expect(product.resolved.run == run) + #expect(product.productTerminal == .interrupted(.requested(.mcpClient(message: "Stop")))) + _ = try await cancellation.value + #expect(await requestStarted.invocationCount() == 1) + } + + @Test func recoveryClassifiesTypedStreamFailureBeforeTokenization() async throws { + let (recoverableAdmission, recoverableRun) = try await makeActiveAdmission() + let recoverableRequest = InvocationProbe() + let recoverable = Task { + try await recoverableAdmission.beginRecovery( + trigger: .recoverableNetworkLoss, + interrupt: { _, _ in await recoverableRequest.record() }, + forceClose: {} + ) + } + await recoverableRequest.waitForInvocation() + try await recoverableAdmission.recordStreamTerminal( + .recoverableNetwork(.connection("offline")) + ) + guard case .replacement(let candidate) = try await recoverable.value else { + Issue.record("Verified network failure must produce a replacement candidate.") + return + } + #expect(candidate.resolved.run == recoverableRun) + + let (protocolAdmission, protocolRun) = try await makeActiveAdmission() + let protocolFailure = ReviewAttemptContractFailure(message: "invalid notification") + let protocolRequest = InvocationProbe() + let nonrecoverable = Task { + try await protocolAdmission.beginRecovery( + trigger: .recoverableNetworkLoss, + interrupt: { _, _ in await protocolRequest.record() }, + forceClose: {} + ) + } + await protocolRequest.waitForInvocation() + try await protocolAdmission.recordStreamTerminal(.protocolViolation(protocolFailure)) + guard case .productTerminal(let product) = try await nonrecoverable.value else { + Issue.record("Protocol failure must terminalize the product.") + return + } + #expect(product.resolved.run == protocolRun) + #expect(product.resolved.terminal == .stream(.protocolViolation(protocolFailure))) + #expect(product.productTerminal == .failed(message: protocolFailure.message)) + } + @Test func threadStartDispatchAdmissionRejectsDirectDuplicate() async throws { let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) let startGate = AsyncGate() - let startTask = await admission.start { _ in + let startTask = try await registerAndActivateStart(admission) { _ in await startGate.waitIgnoringCancellation() return .init(run: canonicalRun) } @@ -501,7 +657,7 @@ struct ReviewAttemptProcessorTests { @Test func threadStartDispatchAdmissionAllowsVerifiedRejectionRetry() async throws { let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) let startGate = AsyncGate() - let startTask = await admission.start { _ in + let startTask = try await registerAndActivateStart(admission) { _ in await startGate.waitIgnoringCancellation() return .init(run: canonicalRun) } @@ -518,7 +674,7 @@ struct ReviewAttemptProcessorTests { @Test func reviewStartDispatchAdmissionRejectsDirectDuplicate() async throws { let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) let startGate = AsyncGate() - let startTask = await admission.start { _ in + let startTask = try await registerAndActivateStart(admission) { _ in await startGate.waitIgnoringCancellation() return .init(run: canonicalRun) } @@ -535,7 +691,7 @@ struct ReviewAttemptProcessorTests { @Test func generalStartFailureEndsAdmissionWaiters() async throws { let failure = ReviewAttemptContractFailure(message: "Start failed") let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) - let startTask = await admission.start { _ in + let startTask = try await registerAndActivateStart(admission) { _ in throw failure } @@ -550,7 +706,7 @@ struct ReviewAttemptProcessorTests { let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) let entered = InvocationProbe() let dispatchGate = AsyncGate() - let startTask = await admission.start { admission in + let startTask = try await registerAndActivateStart(admission) { admission in await entered.record() await dispatchGate.wait() try Task.checkCancellation() @@ -581,7 +737,7 @@ struct ReviewAttemptProcessorTests { let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: graceGate)) let threadDispatched = InvocationProbe() let threadResponseGate = AsyncGate() - let startTask = await admission.start { admission in + let startTask = try await registerAndActivateStart(admission) { admission in #expect(await admission.admitThreadStartDispatch()) await threadDispatched.record() await threadResponseGate.waitIgnoringCancellation() @@ -620,7 +776,7 @@ struct ReviewAttemptProcessorTests { let threadResponseGate = AsyncGate() let forceClose = InvocationProbe() let connection = ReviewRuntimeCloseFailure.connection("Forced close") - let startTask = await admission.start { admission in + let startTask = try await registerAndActivateStart(admission) { admission in #expect(await admission.admitThreadStartDispatch()) await threadDispatched.record() await threadResponseGate.wait() @@ -636,7 +792,9 @@ struct ReviewAttemptProcessorTests { interrupt: { _, _ in Issue.record("Thread-only attempt interrupted a turn.") }, forceClose: { await forceClose.record() - await admission.recordConnectionTerminal(connection) + try await admission.recordStreamTerminal( + .ownerForcedConnectionClose(connection) + ) } ) } @@ -644,7 +802,9 @@ struct ReviewAttemptProcessorTests { await graceGate.open() await forceClose.waitForInvocation() - #expect(try await cancellation.value.terminal == .connection(connection)) + #expect(try await cancellation.value.terminal == .stream( + .ownerForcedConnectionClose(connection) + )) await #expect(throws: CancellationError.self) { try await startTask.value } @@ -654,7 +814,7 @@ struct ReviewAttemptProcessorTests { let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) let prepared = InvocationProbe() let reviewDispatchGate = AsyncGate() - let startTask = await admission.start { admission in + let startTask = try await registerAndActivateStart(admission) { admission in #expect(await admission.admitThreadStartDispatch()) await admission.recordPreparedThread(provisionalRun) await prepared.record() @@ -691,7 +851,7 @@ struct ReviewAttemptProcessorTests { let reviewDispatched = InvocationProbe() let reviewResponseGate = AsyncGate() let interruptCalled = InvocationProbe() - let startTask = await admission.start { admission in + let startTask = try await registerAndActivateStart(admission) { admission in #expect(await admission.admitThreadStartDispatch()) await admission.recordPreparedThread(provisionalRun) #expect(await admission.admitReviewStartDispatch(for: provisionalRun)) @@ -721,8 +881,7 @@ struct ReviewAttemptProcessorTests { ) #expect(try await cancellation.value.terminal == .canonical( - run: canonicalRun, - terminal: .interrupted(.requested(.mcpClient(message: "Stop"))) + .interrupted(.requested(.mcpClient(message: "Stop"))) )) } @@ -733,7 +892,7 @@ struct ReviewAttemptProcessorTests { let reviewResponseGate = AsyncGate() let forceClose = InvocationProbe() let connection = ReviewRuntimeCloseFailure.connection("Forced close") - let startTask = await admission.start { admission in + let startTask = try await registerAndActivateStart(admission) { admission in #expect(await admission.admitThreadStartDispatch()) await admission.recordPreparedThread(provisionalRun) #expect(await admission.admitReviewStartDispatch(for: provisionalRun)) @@ -751,7 +910,9 @@ struct ReviewAttemptProcessorTests { interrupt: { _, _ in Issue.record("Unresolved review request interrupted a turn.") }, forceClose: { await forceClose.record() - await admission.recordConnectionTerminal(connection) + try await admission.recordStreamTerminal( + .ownerForcedConnectionClose(connection) + ) } ) } @@ -759,7 +920,9 @@ struct ReviewAttemptProcessorTests { await graceGate.open() await forceClose.waitForInvocation() - #expect(try await cancellation.value.terminal == .connection(connection)) + #expect(try await cancellation.value.terminal == .stream( + .ownerForcedConnectionClose(connection) + )) await #expect(throws: CancellationError.self) { try await startTask.value } @@ -810,7 +973,7 @@ private func makeActiveAdmission( let admission = ReviewStartAdmission( closePolicy: closePolicy ?? controlledClosePolicy(gate: AsyncGate()) ) - let startTask = await admission.start { admission in + let startTask = try await registerAndActivateStart(admission) { admission in #expect(await admission.admitThreadStartDispatch()) await admission.recordPreparedThread(provisionalRun) #expect(await admission.admitReviewStartDispatch(for: provisionalRun)) @@ -821,6 +984,15 @@ private func makeActiveAdmission( return (admission, canonicalRun) } +private func registerAndActivateStart( + _ admission: ReviewStartAdmission, + operation: @escaping @Sendable (ReviewStartAdmission) async throws -> BackendReviewAttempt +) async throws -> Task { + let registered = try await admission.registerStart(operation) + try await admission.activateStart(registered.id) + return registered.task +} + private func controlledClosePolicy(gate: AsyncGate) -> ReviewRuntimeClosePolicy { ReviewRuntimeClosePolicy(terminalGrace: .seconds(10)) { _ in await gate.wait() From 030b21d21c86bce7a2a1d3c6cf22a51b05a49884 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 21:34:55 +0900 Subject: [PATCH 17/38] fix(review): unify store attempt ownership --- .../CodexReview/ReviewAttemptProcessor.swift | 50 +- .../CodexReview/Store/CodexReviewStore.swift | 4 +- .../Store/CodexReviewStoreCancellation.swift | 27 - .../Store/CodexReviewStoreReviews.swift | 1088 +++++++++++------ .../Store/CodexReviewStoreTesting.swift | 16 +- Sources/CodexReviewTesting/TestSupport.swift | 118 +- .../CodexReviewHostTests.swift | 23 +- .../CodexReviewStoreCommandTests.swift | 358 +++++- .../ReviewAttemptProcessorTests.swift | 64 + .../ReviewTerminalContractTests.swift | 10 +- Tests/ReviewUITests/ReviewUITests.swift | 12 +- 11 files changed, 1302 insertions(+), 468 deletions(-) diff --git a/Sources/CodexReview/ReviewAttemptProcessor.swift b/Sources/CodexReview/ReviewAttemptProcessor.swift index b47f5921..d0b4cff0 100644 --- a/Sources/CodexReview/ReviewAttemptProcessor.swift +++ b/Sources/CodexReview/ReviewAttemptProcessor.swift @@ -226,6 +226,15 @@ package struct ReviewProductTerminalDisposition: Equatable, Sendable { package enum ReviewRecoveryDisposition: Equatable, Sendable { case productTerminal(ReviewProductTerminalDisposition) case replacement(ReviewRecoveryCandidate) + + package var resolvedAttempt: ReviewResolvedAttemptTerminal { + switch self { + case .productTerminal(let disposition): + disposition.resolved + case .replacement(let candidate): + candidate.resolved + } + } } package struct ReviewRecoveryHandoff: Equatable, Sendable { @@ -332,6 +341,7 @@ package actor ReviewStartAdmission { private var joinedTerminalCancellation: ReviewCancellation? private var interruptionPurpose: ReviewAttemptInterruptionPurpose? private var startTask: Task? + private var startDidFinish = false private var nextStartGeneration: UInt64 = 0 private var registeredStartID: ReviewStartHandleID? private var startActivationResult: Result? @@ -402,7 +412,17 @@ package actor ReviewStartAdmission { ) } if let startActivationResult { - return try startActivationResult.get() + switch startActivationResult { + case .failure(let error): + throw error + case .success: + guard startDidFinish == false else { + throw ReviewAttemptContractFailure( + message: "Start activation handle \(id.generation) is stale." + ) + } + return + } } guard case .registeredStart(id) = phase else { throw ReviewAttemptContractFailure( @@ -818,6 +838,33 @@ package actor ReviewStartAdmission { } } + package func recoveryDispositionIfInstalled() -> ReviewRecoveryDisposition? { + installedRecoveryDisposition + } + + package func terminalCancellationProductTerminal( + for failure: ReviewAttemptStreamFailure + ) async -> ReviewTerminalRecord? { + guard case .terminalCancellation(let cancellation) = interruptionPurpose, + let cancellationTask + else { + return nil + } + switch await cancellationTask.result { + case .success(let resolution): + guard resolution.terminal == .stream(failure) else { + return nil + } + if resolution.requestFailure == nil, + case .ownerForcedConnectionClose = failure { + return .interrupted(.requested(cancellation)) + } + return productTerminal(for: failure) + case .failure: + return productTerminal(for: failure) + } + } + package func recordedCleanupResult( for run: CodexReviewBackendModel.Review.Run ) async -> Result? { @@ -894,6 +941,7 @@ package actor ReviewStartAdmission { private func finishStart( with result: Result ) { + startDidFinish = true switch result { case .success(let attempt): registeredRun = attempt.run diff --git a/Sources/CodexReview/Store/CodexReviewStore.swift b/Sources/CodexReview/Store/CodexReviewStore.swift index 0be95602..6a3713f4 100644 --- a/Sources/CodexReview/Store/CodexReviewStore.swift +++ b/Sources/CodexReview/Store/CodexReviewStore.swift @@ -29,9 +29,7 @@ public final class CodexReviewStore { @ObservationIgnored package var previewSupportRetainer: AnyObject? @ObservationIgnored package let clock: CodexReviewClock @ObservationIgnored package let idGenerator: CodexReviewIDGenerator - @ObservationIgnored package var activeRuns: [String: CodexReviewBackendModel.Review.Run] = [:] - @ObservationIgnored package var reviewRecoveryWaitingJobIDs: Set = [] - @ObservationIgnored package var reviewStartAdmissions: [String: ReviewStartAdmission] = [:] + @ObservationIgnored package var reviewAttemptOwnerships: [String: ReviewAttemptOwnership] = [:] @ObservationIgnored package var reviewCleanupFailures: [String: ReviewRuntimeCloseFailure] = [:] @ObservationIgnored package var reviewWorkerTasks: [String: Task] = [:] @ObservationIgnored package var runtimeStopDetachedReviewWorkerTasks: [String: Task] = [:] diff --git a/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift b/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift index b29bfa2b..8d28d375 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift @@ -28,30 +28,6 @@ private actor RuntimeStopDetachedReviewWorkerDrainRace { } extension CodexReviewStore { - package func recordCancellationRequest( - _ cancellation: ReviewCancellation, - for job: CodexReviewJob - ) { - guard job.isTerminal == false else { - return - } - job.cancellationRequested = true - job.core.lifecycle.cancellation = cancellation - job.core.output.summary = cancellation.message - job.core.lifecycle.errorMessage = cancellation.message - } - - @discardableResult - package func recordActiveReviewCancellationRequestsForRuntimeStop( - reason: ReviewCancellation = .system(message: "Review runtime stopped.") - ) -> [String] { - let jobs = orderedJobs.filter { $0.isTerminal == false } - for job in jobs { - recordCancellationRequest(reason, for: job) - } - return jobs.map(\.id) - } - package func completeCancellationLocally( jobID: String, sessionID: String, @@ -200,9 +176,6 @@ extension CodexReviewStore { task.cancel() runtimeStopDetachedReviewWorkerTasks[jobID] = task } - activeRuns.removeValue(forKey: jobID) - reviewRecoveryWaitingJobIDs.remove(jobID) - reviewStartAdmissions.removeValue(forKey: jobID) } } diff --git a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift index b27fe1cc..29fc11fb 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift @@ -15,7 +15,7 @@ extension CodexReviewStore { sessionID: String, request: CodexReviewAPI.Start.Request ) async throws -> CodexReviewAPI.Read.Result { - let jobID = try beginReview(sessionID: sessionID, request: request) + let jobID = try await beginReview(sessionID: sessionID, request: request) // Caller Task cancellation is not a review-cancellation command: only the // attempt admission can distinguish not-sent from outcome-unknown dispatch. // Session owners must use cancelReview/closeSession so the canonical barrier drains. @@ -30,7 +30,7 @@ extension CodexReviewStore { request: CodexReviewAPI.Start.Request, waitTimeout: Duration ) async throws -> CodexReviewAPI.Read.Result { - let jobID = try beginReview(sessionID: sessionID, request: request) + let jobID = try await beginReview(sessionID: sessionID, request: request) return try await awaitReview(sessionID: sessionID, jobID: jobID, timeout: waitTimeout) } @@ -53,7 +53,7 @@ extension CodexReviewStore { private func beginReview( sessionID: String, request: CodexReviewAPI.Start.Request - ) throws -> String { + ) async throws -> String { guard closedSessions.contains(sessionID) == false else { throw CodexReviewAPI.Error.invalidArguments("Review session \(sessionID) is closed.") } @@ -73,159 +73,160 @@ extension CodexReviewStore { ), logEntries: [] ) - insertReviewJob(job) - markReviewRunning(job, startedAt: createdAt) let admission = ReviewStartAdmission(closePolicy: reviewRuntimeClosePolicy) - reviewStartAdmissions[jobID] = admission - launchReviewWorker( + let startRequest = CodexReviewBackendModel.Review.Start( jobID: jobID, sessionID: sessionID, request: validatedRequest, - admission: admission + model: settings.effectiveModel + ) + let backend = self.backend + let registered = try await admission.registerStart { admission in + try await backend.startReview(startRequest, admission: admission) + } + insertReviewJob(job) + markReviewRunning(job, startedAt: createdAt) + reviewAttemptOwnerships[jobID] = .initialStart(registered) + launchReviewWorker( + jobID: jobID, + startRequest: startRequest, + registeredStart: registered ) return jobID } private func launchReviewWorker( jobID: String, - sessionID: String, - request: CodexReviewAPI.Start.Request, - admission: ReviewStartAdmission + startRequest: CodexReviewBackendModel.Review.Start, + registeredStart: ReviewRegisteredStart ) { reviewWorkerTasks[jobID]?.cancel() reviewWorkerTasks[jobID] = Task { [weak self] in await self?.runReviewWorker( jobID: jobID, - sessionID: sessionID, - request: request, - admission: admission + startRequest: startRequest, + registeredStart: registeredStart ) } } private func runReviewWorker( jobID: String, - sessionID: String, - request validatedRequest: CodexReviewAPI.Start.Request, - admission: ReviewStartAdmission + startRequest: CodexReviewBackendModel.Review.Start, + registeredStart: ReviewRegisteredStart ) async { guard let job = job(id: jobID) else { - reviewStartAdmissions.removeValue(forKey: jobID) + reviewAttemptOwnerships.removeValue(forKey: jobID) reviewWorkerTasks.removeValue(forKey: jobID) resumeReviewWaiters(for: jobID) return } - let startRequest = CodexReviewBackendModel.Review.Start( - jobID: jobID, - sessionID: sessionID, - request: validatedRequest, - model: settings.effectiveModel - ) - var run: CodexReviewBackendModel.Review.Run? + var cleanupAttempt: ReviewActiveAttempt? do { - let backend = self.backend - let registered = try await admission.registerStart { admission in - try await backend.startReview(startRequest, admission: admission) + try await registeredStart.admission.activateStart(registeredStart.id) + let backendAttempt = try await registeredStart.task.value + guard case .initialStart(let currentStart) = reviewAttemptOwnerships[jobID], + currentStart.id == registeredStart.id + else { + throw ReviewAttemptContractFailure( + message: "Initial review start completed after its ownership changed." + ) } - try await admission.activateStart(registered.id) - let backendAttempt = try await registered.task.value - let backendRun = backendAttempt.run - run = backendRun - applyBackendRun(backendRun, to: job) + let active = ReviewActiveAttempt( + run: backendAttempt.run, + admission: registeredStart.admission + ) + reviewAttemptOwnerships[jobID] = .active(active) + cleanupAttempt = active + applyBackendRun(backendAttempt.run, to: job) - if job.isTerminal { - do { - try await cleanupReview(backendRun, admission: admission) - } catch { - retainCleanupFailure(error, for: jobID) - } - activeRuns.removeValue(forKey: jobID) - reviewRecoveryWaitingJobIDs.remove(jobID) - } else { + if job.isTerminal == false { let completion = try await consumeReviewEvents( for: backendAttempt, job: job, - startRequest: startRequest, - admission: admission + startRequest: startRequest ) - run = completion.run - do { - try await cleanupReview(completion.run, admission: completion.admission) - } catch { - retainCleanupFailure(error, for: jobID) - } - activeRuns.removeValue(forKey: jobID) - reviewRecoveryWaitingJobIDs.remove(jobID) + cleanupAttempt = completion.cleanupAttempt } } catch let cancellation as ReviewStartCancelledBeforeDispatch { if job.isTerminal == false { - try? completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation.cancellation - ) + do { + try completeCancellationLocally( + jobID: job.id, + sessionID: job.sessionID, + cancellation: cancellation.cancellation + ) + } catch { + markReviewFailed(job, message: error.localizedDescription) + } } } catch let error where error is CancellationError || Task.isCancelled { - if let cleanupRun = activeRuns[jobID] ?? run { - let cleanupAdmission = reviewStartAdmissions[jobID] ?? admission + if let active = activeAttemptForCleanup(jobID: jobID) ?? cleanupAttempt { + cleanupAttempt = active let failure = ReviewRuntimeCloseFailure.worker( "Review worker was cancelled before a canonical terminal." ) do { - try await cleanupAdmission.recordStreamTerminal(.ownerCancellation) + try await active.admission.recordStreamTerminal(.ownerCancellation) } catch { markReviewFailed(job, message: error.localizedDescription) } - do { - try await cleanupReview(cleanupRun, admission: cleanupAdmission) - } catch { - retainCleanupFailure(error, for: jobID) - } if job.isTerminal == false { markReviewInterrupted(job, cause: .transport(message: failure.localizedDescription)) } } else if job.isTerminal == false { markReviewFailed(job, message: error.localizedDescription) } - activeRuns.removeValue(forKey: jobID) - reviewRecoveryWaitingJobIDs.remove(jobID) } catch { - if let cleanupRun = activeRuns[jobID] ?? run { - let cleanupAdmission = reviewStartAdmissions[jobID] ?? admission - do { - try await cleanupReview(cleanupRun, admission: cleanupAdmission) - } catch { - retainCleanupFailure(error, for: jobID) - } + if let active = activeAttemptForCleanup(jobID: jobID) ?? cleanupAttempt { + cleanupAttempt = active } - activeRuns.removeValue(forKey: jobID) - reviewRecoveryWaitingJobIDs.remove(jobID) if job.isTerminal == false, let streamFailure = error as? ReviewAttemptStreamFailure { - let currentAdmission = reviewStartAdmissions[jobID] ?? admission - do { - try await currentAdmission.recordStreamTerminal(streamFailure) - } catch { - markReviewFailed(job, message: error.localizedDescription) + if let cleanupAttempt { + do { + try await cleanupAttempt.admission.recordStreamTerminal(streamFailure) + } catch { + markReviewFailed(job, message: error.localizedDescription) + } } - markReviewInterrupted( - job, - cause: .transport(message: streamFailure.localizedDescription) - ) + applyStreamProductTerminal(streamFailure, to: job) } else if job.isTerminal == false { markReviewFailed(job, message: error.localizedDescription) } } + + reviewAttemptOwnerships[jobID] = .terminal + if let cleanupAttempt { + do { + try await cleanupReview( + cleanupAttempt.run, + admission: cleanupAttempt.admission + ) + } catch { + retainCleanupFailure(error, for: jobID) + } + } reviewWorkerTasks.removeValue(forKey: jobID) runtimeStopDetachedReviewWorkerTasks.removeValue(forKey: jobID) - if reviewCleanupFailures[jobID] == nil { - reviewStartAdmissions.removeValue(forKey: jobID) + if case .terminal = reviewAttemptOwnerships[jobID] { + reviewAttemptOwnerships.removeValue(forKey: jobID) } if job.isTerminal { resumeReviewWaiters(for: jobID) } } + private func activeAttemptForCleanup(jobID: String) -> ReviewActiveAttempt? { + switch reviewAttemptOwnerships[jobID] { + case .active(let active), .resolvingRecovery(let active): + active + case .initialStart, .recoveryDisposition, .preparingRecovery, + .waitingForRecovery, .replacementStart, .terminal, nil: + nil + } + } + private func cleanupReview( _ run: CodexReviewBackendModel.Review.Run, admission: ReviewStartAdmission @@ -249,7 +250,6 @@ extension CodexReviewStore { } private func applyBackendRun(_ backendRun: CodexReviewBackendModel.Review.Run, to job: CodexReviewJob) { - activeRuns[job.id] = backendRun job.core.run = .init( reviewThreadID: backendRun.reviewThreadID, threadID: backendRun.threadID, @@ -273,11 +273,10 @@ extension CodexReviewStore { appendRecoveryProgress(networkRecoveryUnavailableMessage, to: job) } - private func recordReviewRecoveryBarrier(for jobID: String) { - reviewRecoveryWaitingJobIDs.insert(jobID) - } - - private func reviewWorkerInputs(for attempt: BackendReviewAttempt) async -> ReviewWorkerInputs { + private func reviewWorkerInputs( + for attempt: BackendReviewAttempt, + owner: ReviewActiveAttempt + ) async -> ReviewWorkerInputs { let networkMonitor = self.networkMonitor let policy = self.networkRecoveryPolicy let snapshots = networkMonitor.snapshots() @@ -295,7 +294,10 @@ extension CodexReviewStore { await signalCoordinator.observe(snapshot) } } - let initialEventSubscriptionID = await eventSource.subscribe(to: attempt) + let initialEventSubscriptionID = await eventSource.subscribe( + to: attempt, + owner: owner + ) return .init( queue: queue, networkStatusTracker: tracker, @@ -423,58 +425,16 @@ extension CodexReviewStore { return .init(jobID: job.id, cancelled: false, core: job.core) } - recordCancellationRequest(cancellation, for: job) - - if reviewRecoveryWaitingJobIDs.contains(jobID) { - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation - ) - reviewWorkerTasks[jobID]?.cancel() - return .init(jobID: job.id, cancelled: true, core: job.core) - } - - guard let admission = reviewStartAdmissions[jobID] else { - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation - ) - return .init(jobID: job.id, cancelled: true, core: job.core) - } - - let backend = self.backend + job.cancellationRequested = true + job.core.lifecycle.cancellation = cancellation + job.core.output.summary = cancellation.message + job.core.lifecycle.errorMessage = cancellation.message do { - let resolution = try await admission.cancel( - cancellation, - interrupt: { run, reason in - try await backend.interruptReview(run, reason: reason) - }, - forceClose: { - try await backend.forceCloseReviewConnection() - } + try await cancelOwnedReviewAttempt( + job: job, + cancellation: cancellation ) - await reviewWorkerTasks[jobID]?.value - if job.isTerminal == false, - case .localCancellation = resolution.terminal { - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation - ) - } - if let run = job.backendRun, - let cleanupResult = await admission.recordedCleanupResult(for: run) { - try cleanupResult.get() - } } catch { - let phase = await admission.currentPhase() - if case .finishing = phase { - await reviewWorkerTasks[jobID]?.value - } else if case .terminal = phase { - await reviewWorkerTasks[jobID]?.value - } if job.isTerminal == false { try recordCancellationFailure( jobID: job.id, @@ -491,6 +451,246 @@ extension CodexReviewStore { ) } + private func cancelOwnedReviewAttempt( + job: CodexReviewJob, + cancellation: ReviewCancellation + ) async throws { + let jobID = job.id + guard let ownership = reviewAttemptOwnerships[jobID] else { + throw ReviewAttemptContractFailure( + message: "Nonterminal review \(jobID) has no attempt ownership." + ) + } + switch ownership { + case .initialStart(let start): + let resolution = try await cancel( + admission: start.admission, + cancellation: cancellation + ) + if case .localCancellation = resolution.terminal, + job.isTerminal == false { + try finishOwnedCancellation(job: job, cancellation: cancellation) + } + _ = await start.task.result + await reviewWorkerTasks[jobID]?.value + removeTerminalOwnershipWithoutWorker(jobID: jobID) + try await rethrowCleanupFailureIfPresent( + admission: start.admission, + run: job.backendRun + ) + case .active(let active): + let resolution = try await cancel( + admission: active.admission, + cancellation: cancellation + ) + try commitAcknowledgedForcedCancellationIfNeeded( + resolution, + job: job, + cancellation: cancellation + ) + await reviewWorkerTasks[jobID]?.value + try await rethrowCleanupFailureIfPresent( + admission: active.admission, + run: active.run + ) + case .resolvingRecovery(let active): + _ = try await cancel( + admission: active.admission, + cancellation: cancellation + ) + if case .resolvingRecovery(let current) = reviewAttemptOwnerships[jobID], + sameAttempt(current, active), + let disposition = await active.admission.recoveryDispositionIfInstalled() { + if disposition.isNaturalCanonicalProductTerminal { + await reviewWorkerTasks[jobID]?.value + return + } + reviewAttemptOwnerships[jobID] = .recoveryDisposition(disposition) + } + try await suppressRecoverySuccessor( + job: job, + cancellation: cancellation + ) + await reviewWorkerTasks[jobID]?.value + case .recoveryDisposition(let disposition): + try finishRecoveryDispositionForCancellation( + disposition, + job: job, + cancellation: cancellation + ) + reviewWorkerTasks[jobID]?.cancel() + await reviewWorkerTasks[jobID]?.value + case .preparingRecovery, .waitingForRecovery: + try await suppressRecoverySuccessor( + job: job, + cancellation: cancellation + ) + await reviewWorkerTasks[jobID]?.value + case .replacementStart(_, let start): + let resolution = try await cancel( + admission: start.admission, + cancellation: cancellation + ) + if case .localCancellation = resolution.terminal, + job.isTerminal == false { + try finishOwnedCancellation(job: job, cancellation: cancellation) + } + _ = await start.task.result + await reviewWorkerTasks[jobID]?.value + removeTerminalOwnershipWithoutWorker(jobID: jobID) + case .terminal: + guard job.isTerminal else { + throw ReviewAttemptContractFailure( + message: "Terminal attempt ownership has a nonterminal product review." + ) + } + } + } + + private func cancel( + admission: ReviewStartAdmission, + cancellation: ReviewCancellation + ) async throws -> ReviewAttemptCancellationResolution { + let backend = self.backend + return try await admission.cancel( + cancellation, + interrupt: { run, reason in + try await backend.interruptReview(run, reason: reason) + }, + forceClose: { + try await backend.forceCloseReviewConnection() + } + ) + } + + private func suppressRecoverySuccessor( + job: CodexReviewJob, + cancellation: ReviewCancellation + ) async throws { + switch reviewAttemptOwnerships[job.id] { + case .resolvingRecovery: + throw ReviewAttemptContractFailure( + message: "Recovery cancellation completed before disposition installation." + ) + case .recoveryDisposition(let disposition): + try finishRecoveryDispositionForCancellation( + disposition, + job: job, + cancellation: cancellation + ) + reviewWorkerTasks[job.id]?.cancel() + case .preparingRecovery(_, let task): + try finishOwnedCancellation(job: job, cancellation: cancellation) + task.cancel() + _ = await task.result + reviewWorkerTasks[job.id]?.cancel() + case .waitingForRecovery: + try finishOwnedCancellation(job: job, cancellation: cancellation) + reviewWorkerTasks[job.id]?.cancel() + case .replacementStart(_, let start): + let resolution = try await cancel( + admission: start.admission, + cancellation: cancellation + ) + if case .localCancellation = resolution.terminal, + job.isTerminal == false { + try finishOwnedCancellation(job: job, cancellation: cancellation) + } + _ = await start.task.result + await reviewWorkerTasks[job.id]?.value + case .active(let active): + let resolution = try await cancel( + admission: active.admission, + cancellation: cancellation + ) + try commitAcknowledgedForcedCancellationIfNeeded( + resolution, + job: job, + cancellation: cancellation + ) + await reviewWorkerTasks[job.id]?.value + case .initialStart(let start): + let resolution = try await cancel( + admission: start.admission, + cancellation: cancellation + ) + if case .localCancellation = resolution.terminal, + job.isTerminal == false { + try finishOwnedCancellation(job: job, cancellation: cancellation) + } + _ = await start.task.result + await reviewWorkerTasks[job.id]?.value + case .terminal: + return + case nil: + throw ReviewAttemptContractFailure( + message: "Recovery cancellation lost attempt ownership." + ) + } + } + + private func finishOwnedCancellation( + job: CodexReviewJob, + cancellation: ReviewCancellation + ) throws { + try completeCancellationLocally( + jobID: job.id, + sessionID: job.sessionID, + cancellation: cancellation + ) + reviewAttemptOwnerships[job.id] = .terminal + } + + private func commitAcknowledgedForcedCancellationIfNeeded( + _ resolution: ReviewAttemptCancellationResolution, + job: CodexReviewJob, + cancellation: ReviewCancellation + ) throws { + guard job.isTerminal == false, + resolution.requestFailure == nil, + case .stream(.ownerForcedConnectionClose) = resolution.terminal + else { + return + } + try finishOwnedCancellation(job: job, cancellation: cancellation) + } + + private func finishRecoveryDispositionForCancellation( + _ disposition: ReviewRecoveryDisposition, + job: CodexReviewJob, + cancellation: ReviewCancellation + ) throws { + switch disposition { + case .productTerminal(let product): + try applyRecoveryProductTerminal(product.productTerminal, to: job) + reviewAttemptOwnerships[job.id] = .terminal + case .replacement: + try finishOwnedCancellation(job: job, cancellation: cancellation) + } + } + + private func rethrowCleanupFailureIfPresent( + admission: ReviewStartAdmission, + run: CodexReviewBackendModel.Review.Run? + ) async throws { + guard let run, + let cleanupResult = await admission.recordedCleanupResult(for: run) + else { + return + } + try cleanupResult.get() + } + + private func removeTerminalOwnershipWithoutWorker(jobID: String) { + guard reviewWorkerTasks[jobID] == nil, + runtimeStopDetachedReviewWorkerTasks[jobID] == nil, + case .terminal = reviewAttemptOwnerships[jobID] + else { + return + } + reviewAttemptOwnerships.removeValue(forKey: jobID) + } + package func closeSession( _ sessionID: String, reason: ReviewCancellation = .sessionClosed() @@ -614,24 +814,37 @@ extension CodexReviewStore { private func consumeReviewEvents( for initialAttempt: BackendReviewAttempt, job: CodexReviewJob, - startRequest: CodexReviewBackendModel.Review.Start, - admission: ReviewStartAdmission + startRequest: CodexReviewBackendModel.Review.Start ) async throws -> ReviewWorkerAttemptCompletion { - let inputs = await reviewWorkerInputs(for: initialAttempt) + guard case .active(let initialActive) = reviewAttemptOwnerships[job.id], + sameAttempt(initialActive, run: initialAttempt.run) + else { + throw ReviewAttemptContractFailure( + message: "Initial event subscription requires the published active attempt." + ) + } + let inputs = await reviewWorkerInputs(for: initialAttempt, owner: initialActive) + guard case .active(let revalidatedActive) = reviewAttemptOwnerships[job.id], + sameAttempt(revalidatedActive, initialActive) + else { + await inputs.cancel() + throw ReviewAttemptContractFailure( + message: "Initial event subscription became stale before publication completed." + ) + } do { let completion = try await consumeReviewEventLoop( - for: initialAttempt, job: job, startRequest: startRequest, - initialAdmission: admission, inputs: inputs ) await inputs.cancel() return completion } catch { if error is CancellationError || Task.isCancelled { - let currentAdmission = reviewStartAdmissions[job.id] ?? admission - try await currentAdmission.recordStreamTerminal(.ownerCancellation) + if let active = activeAttemptForCleanup(jobID: job.id) { + try await active.admission.recordStreamTerminal(.ownerCancellation) + } } await inputs.cancel() throw error @@ -639,91 +852,98 @@ extension CodexReviewStore { } private func consumeReviewEventLoop( - for initialAttempt: BackendReviewAttempt, job: CodexReviewJob, startRequest: CodexReviewBackendModel.Review.Start, - initialAdmission: ReviewStartAdmission, inputs: ReviewWorkerInputs ) async throws -> ReviewWorkerAttemptCompletion { - var admission = initialAdmission - var recoveryState = ReviewNetworkRecoveryLoopState(currentRun: initialAttempt.run) + var recoverySignals = ReviewNetworkRecoverySignals() var activeEventSubscriptionID: Int? = inputs.initialEventSubscriptionID while let input = await inputs.next() { if job.isTerminal { - return .init(run: recoveryState.currentRun, admission: admission) + return .init(cleanupAttempt: activeAttemptForCleanup(jobID: job.id)) } switch input { case .reviewEvent(let event): guard activeEventSubscriptionID == event.subscriptionID, - recoveryState.shouldRouteAttemptInput(from: event.subscriptionRun) + let routed = routedAttempt( + jobID: job.id, + owner: event.owner + ) else { continue } if let terminal = reviewTerminalRecord(for: event.event, job: job) { - try await admission.recordCanonicalTerminal( + try await routed.active.admission.recordCanonicalTerminal( terminal, - for: recoveryState.currentRun + for: event.owner.run ) } - if recoveryState.isInterruptingForNetworkRecovery { + if routed.isResolvingRecovery { guard event.event.supersedesNetworkRecovery else { continue } } - recoveryState.currentRun = handleReviewEvent( + _ = handleReviewEvent( event.event, job: job, - currentRun: recoveryState.currentRun + currentRun: routed.active.run ) if job.isTerminal { - return .init(run: recoveryState.currentRun, admission: admission) + return .init(cleanupAttempt: routed.active) } case .reviewEventsFinished(let finishedRun): guard activeEventSubscriptionID == finishedRun.subscriptionID, - recoveryState.shouldRouteAttemptInput(from: finishedRun.run) + let routed = routedAttempt(jobID: job.id, owner: finishedRun.owner) else { continue } - if recoveryState.isInterruptingForNetworkRecovery { - let failure = ReviewAttemptStreamFailure.workerContract(.init( - message: ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription - )) - try await admission.recordStreamTerminal(failure) + let failure = ReviewAttemptStreamFailure.workerContract(.init( + message: ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription + )) + try await routed.active.admission.recordStreamTerminal(failure) + if routed.isResolvingRecovery { activeEventSubscriptionID = nil continue } - if recoveryState.shouldIgnoreFinishedEvent(for: finishedRun.run) { - continue - } - if await handleReviewEventsFinished( - job: job, - isWaitingForNetworkRecovery: recoveryState.isWaitingForNetworkRecovery, - admission: admission - ) { - return .init(run: recoveryState.currentRun, admission: admission) + if let productTerminal = await routed.active.admission + .terminalCancellationProductTerminal(for: failure) { + try applyRecoveryProductTerminal(productTerminal, to: job) + } else { + applyStreamProductTerminal(failure, to: job) } + return .init(cleanupAttempt: routed.active) case .reviewEventsFailed(let failedRun): guard activeEventSubscriptionID == failedRun.subscriptionID, - recoveryState.shouldRouteAttemptInput(from: failedRun.run) + let routed = routedAttempt(jobID: job.id, owner: failedRun.owner) else { continue } - if recoveryState.isInterruptingForNetworkRecovery { - try await admission.recordStreamTerminal(failedRun.failure) + if routed.isResolvingRecovery { + try await routed.active.admission.recordStreamTerminal(failedRun.failure) activeEventSubscriptionID = nil continue } - if case .unexpectedConnection = failedRun.failure, + if case .recoverableNetwork = failedRun.failure, await inputs.networkStatusTracker.currentStatus() != .satisfied { - recoveryState.recordPendingOutageStreamFailure(failedRun.failure) + recoverySignals.recordPendingOutageStreamFailure( + failedRun.failure, + attemptID: routed.active.run.attemptID + ) activeEventSubscriptionID = nil await inputs.cancelActiveEventSubscription() continue } - throw failedRun.failure + try await routed.active.admission.recordStreamTerminal(failedRun.failure) + if let productTerminal = await routed.active.admission + .terminalCancellationProductTerminal(for: failedRun.failure) { + try applyRecoveryProductTerminal(productTerminal, to: job) + } else { + applyStreamProductTerminal(failedRun.failure, to: job) + } + return .init(cleanupAttempt: routed.active) case .recoveryBarrierResolved(let resolution): - guard recoveryState.isInterruptingForNetworkRecovery, - resolution.run.attemptID == recoveryState.currentRun.attemptID + guard case .resolvingRecovery(let resolving) = reviewAttemptOwnerships[job.id], + sameAttempt(resolving, resolution.owner) else { continue } @@ -731,82 +951,125 @@ extension CodexReviewStore { case .failure(let failure): throw failure.underlying case .success(let disposition): - if job.isTerminal { - return .init(run: recoveryState.currentRun, admission: admission) - } + reviewAttemptOwnerships[job.id] = .recoveryDisposition(disposition) let candidate: ReviewRecoveryCandidate switch disposition { case .productTerminal(let product): - applyRecoveryProductTerminal(product.productTerminal, to: job) - return .init(run: recoveryState.currentRun, admission: admission) + try applyRecoveryProductTerminal(product.productTerminal, to: job) + return .init(cleanupAttempt: resolving) case .replacement(let replacement): candidate = replacement } - if completePendingCancellationIfNeeded(for: job) { - return .init(run: recoveryState.currentRun, admission: admission) + let backend = self.backend + let preparationTask = Task { + try await backend.prepareReviewRecovery(candidate) + } + reviewAttemptOwnerships[job.id] = .preparingRecovery( + candidate: candidate, + preparationTask: preparationTask + ) + let preparationResult = await preparationTask.result + guard case .preparingRecovery(let currentCandidate, _) = reviewAttemptOwnerships[job.id], + currentCandidate == candidate + else { + if job.isTerminal { + return .init(cleanupAttempt: nil) + } + throw ReviewAttemptContractFailure( + message: "Recovery preparation completed after its ownership changed." + ) } - let handoff = try await self.backend.prepareReviewRecovery(candidate) - recoveryState.markWaitingForNetworkRecovery() + let handoff = try preparationResult.get() + reviewAttemptOwnerships[job.id] = .waitingForRecovery(handoff) markReviewWaitingForNetworkRecovery(job) - recordReviewRecoveryBarrier(for: job.id) activeEventSubscriptionID = nil await inputs.cancelActiveEventSubscription() - recoveryState.markRecoveryReady(handoff) + guard case .waitingForRecovery(let currentHandoff) = reviewAttemptOwnerships[job.id], + currentHandoff == handoff + else { + if job.isTerminal { + return .init(cleanupAttempt: nil) + } + throw ReviewAttemptContractFailure( + message: "Recovery handoff changed while detaching the old subscription." + ) + } } case .networkSnapshot(let snapshot, let recoveryGeneration): - if let pendingFailure = recoveryState.takePendingOutageStreamFailureAfterTransientRecovery( - snapshot - ) { - throw pendingFailure + if let pendingFailure = recoverySignals + .takePendingOutageStreamFailureAfterTransientRecovery(snapshot), + case .active(let active) = reviewAttemptOwnerships[job.id], + active.run.attemptID == pendingFailure.attemptID { + throw pendingFailure.failure } - switch recoveryState.networkSnapshotEffect(snapshot, recoveryGeneration: recoveryGeneration) { + let waitingHandoff: ReviewRecoveryHandoff? = if case .waitingForRecovery(let handoff) = reviewAttemptOwnerships[job.id] { + handoff + } else { + nil + } + switch recoverySignals.networkSnapshotEffect( + snapshot, + recoveryGeneration: recoveryGeneration, + waitingHandoff: waitingHandoff + ) { case .none: continue case .restartSettling: appendRecoveryProgress(networkRecoveryRestoredMessage, to: job) } case .networkRecoverySettled(let recoveryGeneration): - guard recoveryState.shouldRestartReviewAfterRecoverySettle( - recoveryGeneration: recoveryGeneration - ) else { + guard case .waitingForRecovery(let handoff) = reviewAttemptOwnerships[job.id], + recoverySignals.shouldRestartReviewAfterRecoverySettle( + recoveryGeneration: recoveryGeneration, + handoff: handoff + ) else { continue } switch try await restartReviewAfterNetworkRestore( job: job, startRequest: startRequest, inputs: inputs, - recoveryReady: recoveryState.recoveryReady + handoff: handoff ) { case .continueWaiting: - recoveryState.markWaitingForNetworkRecovery() continue case .finished: - reviewRecoveryWaitingJobIDs.remove(job.id) - return .init(run: recoveryState.currentRun, admission: admission) - case .recovered(let recoveredAttempt, let recoveredAdmission): - let recoveredRun = recoveredAttempt.run - admission = recoveredAdmission - applyBackendRun(recoveredRun, to: job) - recoveryState.markRecovered(with: recoveredRun) - reviewRecoveryWaitingJobIDs.remove(job.id) - activeEventSubscriptionID = await inputs.subscribe(to: recoveredAttempt) + return .init(cleanupAttempt: nil) + case .recovered(let recoveredAttempt, let active): + let subscriptionID = await inputs.subscribe( + to: recoveredAttempt, + owner: active + ) + guard case .active(let current) = reviewAttemptOwnerships[job.id], + sameAttempt(current, active) + else { + await inputs.cancelActiveEventSubscription() + if job.isTerminal { + return .init(cleanupAttempt: nil) + } + throw ReviewAttemptContractFailure( + message: "Recovered subscription completed after its active attempt changed." + ) + } + activeEventSubscriptionID = subscriptionID + recoverySignals.markRecovered() } case .networkOutageConfirmed: - guard recoveryState.isWaitingForNetworkRecovery == false, - recoveryState.isInterruptingForNetworkRecovery == false, + guard case .active(let active) = reviewAttemptOwnerships[job.id], job.isTerminal == false, job.cancellationRequested == false, await inputs.networkStatusTracker.currentStatus() != .satisfied else { continue } - recoveryState.markInterruptingForNetworkRecovery() - let recoveryRun = recoveryState.currentRun - let recoveryAdmission = admission - let pendingFailure = recoveryState.takePendingOutageStreamFailureForConfirmedRecovery() + reviewAttemptOwnerships[job.id] = .resolvingRecovery(active) + let pendingFailure = recoverySignals + .takePendingOutageStreamFailureForConfirmedRecovery( + attemptID: active.run.attemptID + ) let backend = self.backend - await inputs.beginRecoveryInterruption(for: recoveryRun) { - try await recoveryAdmission.beginRecovery( + await inputs.beginRecoveryInterruption(for: active) { + try await active.admission.beginRecovery( trigger: .recoverableNetworkLoss, interrupt: { run, reason in try await backend.interruptReview(run, reason: reason) @@ -817,15 +1080,13 @@ extension CodexReviewStore { ) } if let pendingFailure { - _ = await recoveryAdmission.waitForInterruptionAdmission() - let failure: ReviewAttemptStreamFailure - switch pendingFailure { - case .unexpectedConnection(let closeFailure): - failure = .recoverableNetwork(closeFailure) - default: - failure = pendingFailure + _ = await active.admission.waitForInterruptionAdmission() + guard case .resolvingRecovery(let current) = reviewAttemptOwnerships[job.id], + sameAttempt(current, active) + else { + continue } - try await recoveryAdmission.recordStreamTerminal(failure) + try await active.admission.recordStreamTerminal(pendingFailure.failure) } } } @@ -833,48 +1094,32 @@ extension CodexReviewStore { if Task.isCancelled { throw CancellationError() } - if job.isTerminal == false { + if job.isTerminal == false, + case .active(let active) = reviewAttemptOwnerships[job.id] { let failure = ReviewAttemptStreamFailure.workerContract(.init( message: ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription )) - try await admission.recordStreamTerminal(failure) - markReviewInterrupted(job, cause: .transport(message: failure.localizedDescription)) - } - return .init(run: recoveryState.currentRun, admission: admission) - } - - private func handleReviewEventsFinished( - job: CodexReviewJob, - isWaitingForNetworkRecovery: Bool, - admission: ReviewStartAdmission - ) async -> Bool { - if Task.isCancelled { - return true - } - - if isWaitingForNetworkRecovery { - return job.isTerminal || completePendingCancellationIfNeeded(for: job) - } - - if job.isTerminal == false { - let failure = ReviewAttemptStreamFailure.workerContract(.init( - message: ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription - )) - do { - try await admission.recordStreamTerminal(failure) - } catch { - markReviewFailed(job, message: error.localizedDescription) - return true + try await active.admission.recordStreamTerminal(failure) + if let productTerminal = await active.admission + .terminalCancellationProductTerminal(for: failure) { + try applyRecoveryProductTerminal(productTerminal, to: job) + } else { + applyStreamProductTerminal(failure, to: job) } - markReviewInterrupted(job, cause: .transport(message: failure.localizedDescription)) + return .init(cleanupAttempt: active) } - return true + if job.isTerminal { + return .init(cleanupAttempt: activeAttemptForCleanup(jobID: job.id)) + } + throw ReviewAttemptContractFailure( + message: "Review input queue finished without terminal attempt ownership." + ) } private func applyRecoveryProductTerminal( _ terminal: ReviewTerminalRecord, to job: CodexReviewJob - ) { + ) throws { switch terminal { case .completed: if job.isTerminal == false { @@ -885,46 +1130,136 @@ extension CodexReviewStore { } case .failed(let message): markReviewFailed(job, message: message, terminal: terminal) + case .interrupted(.requested(let cancellation)): + try completeCancellationLocally( + jobID: job.id, + sessionID: job.sessionID, + cancellation: cancellation + ) case .interrupted(let cause): markReviewInterrupted(job, cause: cause) } } + private func applyStreamProductTerminal( + _ failure: ReviewAttemptStreamFailure, + to job: CodexReviewJob + ) { + switch failure { + case .process: + markReviewInterrupted(job, cause: .previousProcessExit) + case .protocolViolation(let failure), .workerContract(let failure): + markReviewFailed(job, message: failure.localizedDescription) + case .ownerCancellation: + markReviewFailed(job, message: failure.localizedDescription) + case .recoverableNetwork, .ownerForcedConnectionClose, + .unexpectedConnection: + markReviewInterrupted( + job, + cause: .transport(message: failure.localizedDescription) + ) + } + } + + private func routedAttempt( + jobID: String, + owner: ReviewActiveAttempt + ) -> RoutedReviewAttempt? { + switch reviewAttemptOwnerships[jobID] { + case .active(let active): + guard sameAttempt(active, owner) else { return nil } + return .init(active: active, isResolvingRecovery: false) + case .resolvingRecovery(let active): + guard sameAttempt(active, owner) else { return nil } + return .init(active: active, isResolvingRecovery: true) + case .initialStart, .recoveryDisposition, .preparingRecovery, + .waitingForRecovery, .replacementStart, .terminal, nil: + return nil + } + } + + private func sameAttempt( + _ active: ReviewActiveAttempt, + run: CodexReviewBackendModel.Review.Run + ) -> Bool { + active.run.attemptID == run.attemptID + } + + private func sameAttempt( + _ lhs: ReviewActiveAttempt, + _ rhs: ReviewActiveAttempt + ) -> Bool { + lhs.run.attemptID == rhs.run.attemptID + && lhs.admission === rhs.admission + } + private func restartReviewAfterNetworkRestore( job: CodexReviewJob, startRequest: CodexReviewBackendModel.Review.Start, inputs: ReviewWorkerInputs, - recoveryReady: ReviewRecoveryHandoff? + handoff: ReviewRecoveryHandoff ) async throws -> NetworkRestoreRestartResult { - if job.isTerminal || completePendingCancellationIfNeeded(for: job) { - return .finished - } - if Task.isCancelled { - throw CancellationError() + guard case .waitingForRecovery(let currentHandoff) = reviewAttemptOwnerships[job.id], + currentHandoff == handoff + else { + if job.isTerminal { + return .finished + } + throw ReviewAttemptContractFailure( + message: "Recovery restart requires its exact waiting handoff." + ) } - if job.isTerminal || completePendingCancellationIfNeeded(for: job) { + if Task.isCancelled || job.isTerminal { return .finished } guard await inputs.networkStatusTracker.currentStatus() == .satisfied else { return .continueWaiting } - guard let recoveryReady else { - return .continueWaiting - } let recoveredAdmission = ReviewStartAdmission(closePolicy: reviewRuntimeClosePolicy) - reviewStartAdmissions[job.id] = recoveredAdmission - reviewRecoveryWaitingJobIDs.remove(job.id) let backend = self.backend let registered = try await recoveredAdmission.registerStart { admission in try await backend.resumeReviewRecovery( - recoveryReady, + handoff, request: startRequest, admission: admission ) } + guard case .waitingForRecovery(let revalidatedHandoff) = reviewAttemptOwnerships[job.id], + revalidatedHandoff == handoff + else { + _ = try await recoveredAdmission.cancel( + job.core.lifecycle.cancellation ?? .system(), + interrupt: { _, _ in }, + forceClose: {} + ) + _ = await registered.task.result + return .finished + } + reviewAttemptOwnerships[job.id] = .replacementStart( + handoff: handoff, + start: registered + ) try await recoveredAdmission.activateStart(registered.id) - let recoveredAttempt = try await registered.task.value - return .recovered(recoveredAttempt, recoveredAdmission) + let result = await registered.task.result + guard case .replacementStart(let currentHandoff, let currentStart) = reviewAttemptOwnerships[job.id], + currentHandoff == handoff, + currentStart.id == registered.id + else { + if job.isTerminal { + return .finished + } + throw ReviewAttemptContractFailure( + message: "Replacement start completed after its ownership changed." + ) + } + let recoveredAttempt = try result.get() + let active = ReviewActiveAttempt( + run: recoveredAttempt.run, + admission: recoveredAdmission + ) + reviewAttemptOwnerships[job.id] = .active(active) + applyBackendRun(recoveredAttempt.run, to: job) + return .recovered(recoveredAttempt, active) } private func handleReviewEvent( @@ -1024,19 +1359,6 @@ extension CodexReviewStore { } } - private func completePendingCancellationIfNeeded(for job: CodexReviewJob) -> Bool { - guard job.cancellationRequested else { - return false - } - let cancellation = job.core.lifecycle.cancellation ?? .system() - try? completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation - ) - return true - } - private func completeReview( _ job: CodexReviewJob, summary: String, @@ -1260,6 +1582,22 @@ private extension CodexReviewBackendModel.Review.Event { } } +private extension ReviewRecoveryDisposition { + var isNaturalCanonicalProductTerminal: Bool { + guard case .productTerminal(let product) = self, + case .canonical(let terminal) = product.resolved.terminal + else { + return false + } + switch terminal { + case .completed, .failed: + return true + case .interrupted: + return false + } + } +} + private extension CodexReviewJob { var backendRun: CodexReviewBackendModel.Review.Run? { guard let threadID = core.run.threadID else { @@ -1299,18 +1637,23 @@ private extension CodexReviewJob { private struct ReviewWorkerReviewEvent: Sendable { var subscriptionID: Int - var subscriptionRun: CodexReviewBackendModel.Review.Run + var owner: ReviewActiveAttempt var event: CodexReviewBackendModel.Review.Event } +private struct RoutedReviewAttempt: Sendable { + var active: ReviewActiveAttempt + var isResolvingRecovery: Bool +} + private struct ReviewWorkerEventStreamFinished: Sendable { var subscriptionID: Int - var run: CodexReviewBackendModel.Review.Run + var owner: ReviewActiveAttempt } private struct ReviewWorkerEventStreamFailed: Sendable { var subscriptionID: Int - var run: CodexReviewBackendModel.Review.Run + var owner: ReviewActiveAttempt var failure: ReviewAttemptStreamFailure } @@ -1325,13 +1668,12 @@ private enum ReviewWorkerInput: Sendable { } private struct ReviewWorkerRecoveryBarrierResolution: Sendable { - var run: CodexReviewBackendModel.Review.Run + var owner: ReviewActiveAttempt var result: Result } private struct ReviewWorkerAttemptCompletion: Sendable { - var run: CodexReviewBackendModel.Review.Run - var admission: ReviewStartAdmission + var cleanupAttempt: ReviewActiveAttempt? } private struct ReviewWorkerRecoveryFailure: LocalizedError, @unchecked Sendable { @@ -1345,7 +1687,7 @@ private struct ReviewWorkerRecoveryFailure: LocalizedError, @unchecked Sendable private enum NetworkRestoreRestartResult { case continueWaiting case finished - case recovered(BackendReviewAttempt, ReviewStartAdmission) + case recovered(BackendReviewAttempt, ReviewActiveAttempt) } private enum ReviewNetworkSnapshotEffect { @@ -1353,64 +1695,37 @@ private enum ReviewNetworkSnapshotEffect { case restartSettling } -private enum ReviewNetworkRecoveryPhase { - case active - case interrupting - case waiting +private struct PendingOutageStreamFailure { + var attemptID: String + var failure: ReviewAttemptStreamFailure } -private struct ReviewNetworkRecoveryLoopState { - var currentRun: CodexReviewBackendModel.Review.Run - private(set) var recoveryPhase = ReviewNetworkRecoveryPhase.active - private(set) var recoveryReady: ReviewRecoveryHandoff? +private struct ReviewNetworkRecoverySignals { private var isSettlingForNetworkRecovery = false private var recoverySettleGeneration: Int? - private var pendingOutageStreamFailure: ReviewAttemptStreamFailure? - let recoveryReason = CodexReviewBackendModel.CancellationReason(message: networkRecoveryUnavailableMessage) - - init(currentRun: CodexReviewBackendModel.Review.Run) { - self.currentRun = currentRun - } - - var isInterruptingForNetworkRecovery: Bool { - recoveryPhase == .interrupting - } - - var isWaitingForNetworkRecovery: Bool { - recoveryPhase == .waiting - } - - mutating func markInterruptingForNetworkRecovery() { - recoveryPhase = .interrupting - isSettlingForNetworkRecovery = false - recoverySettleGeneration = nil - } - - mutating func markWaitingForNetworkRecovery() { - recoveryPhase = .waiting - isSettlingForNetworkRecovery = false - recoverySettleGeneration = nil - pendingOutageStreamFailure = nil - } - - mutating func markRecoveryReady(_ ready: ReviewRecoveryHandoff) { - recoveryReady = ready - } + private var recoverySettleHandoff: ReviewRecoveryHandoff? + private var pendingOutageStreamFailure: PendingOutageStreamFailure? - mutating func markRecovered(with run: CodexReviewBackendModel.Review.Run) { - currentRun = run - recoveryPhase = .active - recoveryReady = nil + mutating func markRecovered() { isSettlingForNetworkRecovery = false recoverySettleGeneration = nil + recoverySettleHandoff = nil pendingOutageStreamFailure = nil } - mutating func recordPendingOutageStreamFailure(_ failure: ReviewAttemptStreamFailure) { - pendingOutageStreamFailure = failure + mutating func recordPendingOutageStreamFailure( + _ failure: ReviewAttemptStreamFailure, + attemptID: String + ) { + pendingOutageStreamFailure = .init(attemptID: attemptID, failure: failure) } - mutating func takePendingOutageStreamFailureForConfirmedRecovery() -> ReviewAttemptStreamFailure? { + mutating func takePendingOutageStreamFailureForConfirmedRecovery( + attemptID: String + ) -> PendingOutageStreamFailure? { + guard pendingOutageStreamFailure?.attemptID == attemptID else { + return nil + } defer { pendingOutageStreamFailure = nil } @@ -1419,10 +1734,8 @@ private struct ReviewNetworkRecoveryLoopState { mutating func takePendingOutageStreamFailureAfterTransientRecovery( _ snapshot: CodexReviewNetworkSnapshot - ) -> ReviewAttemptStreamFailure? { - guard snapshot.status == .satisfied, - isWaitingForNetworkRecovery == false - else { + ) -> PendingOutageStreamFailure? { + guard snapshot.status == .satisfied else { return nil } defer { @@ -1431,39 +1744,40 @@ private struct ReviewNetworkRecoveryLoopState { return pendingOutageStreamFailure } - func shouldIgnoreFinishedEvent(for run: CodexReviewBackendModel.Review.Run) -> Bool { - isWaitingForNetworkRecovery || run.attemptID != currentRun.attemptID - } - - func shouldRestartReviewAfterRecoverySettle(recoveryGeneration: Int) -> Bool { - isWaitingForNetworkRecovery - && isSettlingForNetworkRecovery + func shouldRestartReviewAfterRecoverySettle( + recoveryGeneration: Int, + handoff: ReviewRecoveryHandoff + ) -> Bool { + isSettlingForNetworkRecovery && recoverySettleGeneration == recoveryGeneration - && recoveryReady != nil - } - - func shouldRouteAttemptInput(from run: CodexReviewBackendModel.Review.Run) -> Bool { - recoveryPhase != .waiting && run.attemptID == currentRun.attemptID + && recoverySettleHandoff == handoff } mutating func networkSnapshotEffect( _ snapshot: CodexReviewNetworkSnapshot, - recoveryGeneration: Int + recoveryGeneration: Int, + waitingHandoff: ReviewRecoveryHandoff? ) -> ReviewNetworkSnapshotEffect { - guard isWaitingForNetworkRecovery else { + guard let waitingHandoff else { + isSettlingForNetworkRecovery = false + recoverySettleGeneration = nil + recoverySettleHandoff = nil return .none } guard snapshot.status == .satisfied else { isSettlingForNetworkRecovery = false recoverySettleGeneration = nil + recoverySettleHandoff = nil return .none } guard isSettlingForNetworkRecovery == false else { recoverySettleGeneration = recoveryGeneration + recoverySettleHandoff = waitingHandoff return .none } isSettlingForNetworkRecovery = true recoverySettleGeneration = recoveryGeneration + recoverySettleHandoff = waitingHandoff return .restartSettling } } @@ -1481,8 +1795,11 @@ private struct ReviewWorkerInputs { await queue.next() } - func subscribe(to attempt: BackendReviewAttempt) async -> Int { - await eventSource.subscribe(to: attempt) + func subscribe( + to attempt: BackendReviewAttempt, + owner: ReviewActiveAttempt + ) async -> Int { + await eventSource.subscribe(to: attempt, owner: owner) } func cancelActiveEventSubscription() async { @@ -1490,10 +1807,10 @@ private struct ReviewWorkerInputs { } func beginRecoveryInterruption( - for run: CodexReviewBackendModel.Review.Run, + for owner: ReviewActiveAttempt, operation: @escaping @Sendable () async throws -> ReviewRecoveryDisposition ) async { - await recoveryInterruptionSource.start(for: run, operation: operation) + await recoveryInterruptionSource.start(for: owner, operation: operation) } func cancel() async { @@ -1515,7 +1832,7 @@ private actor ReviewWorkerRecoveryInterruptionSource { } func start( - for run: CodexReviewBackendModel.Review.Run, + for owner: ReviewActiveAttempt, operation: @escaping @Sendable () async throws -> ReviewRecoveryDisposition ) { guard task == nil else { @@ -1528,7 +1845,7 @@ private actor ReviewWorkerRecoveryInterruptionSource { } catch { result = .failure(.init(underlying: error)) } - await queue.send(.recoveryBarrierResolved(.init(run: run, result: result))) + await queue.send(.recoveryBarrierResolved(.init(owner: owner, result: result))) self.finish() } } @@ -1639,12 +1956,14 @@ private actor ReviewWorkerEventSource { self.queue = queue } - func subscribe(to attempt: BackendReviewAttempt) -> Int { + func subscribe( + to attempt: BackendReviewAttempt, + owner: ReviewActiveAttempt + ) -> Int { subscriptionID += 1 let subscriptionID = subscriptionID activeSubscriptionID = subscriptionID cancelEventTasks() - let run = attempt.run let events = attempt.events eventTasks[subscriptionID] = Task { do { @@ -1652,13 +1971,17 @@ private actor ReviewWorkerEventSource { guard Task.isCancelled == false else { return } - await self.yieldReviewEvent(event, run: run, subscriptionID: subscriptionID) + await self.yieldReviewEvent( + event, + owner: owner, + subscriptionID: subscriptionID + ) if event.completesReviewRun { self.finishTerminalDelivery(subscriptionID: subscriptionID) return } } - await self.yieldEventsFinished(run: run, subscriptionID: subscriptionID) + await self.yieldEventsFinished(owner: owner, subscriptionID: subscriptionID) } catch { let failure: ReviewAttemptStreamFailure if let typed = error as? ReviewAttemptStreamFailure { @@ -1668,7 +1991,11 @@ private actor ReviewWorkerEventSource { } else { failure = .workerContract(.init(message: error.localizedDescription)) } - await self.yieldEventsFailed(failure, run: run, subscriptionID: subscriptionID) + await self.yieldEventsFailed( + failure, + owner: owner, + subscriptionID: subscriptionID + ) } } return subscriptionID @@ -1704,7 +2031,7 @@ private actor ReviewWorkerEventSource { private func yieldReviewEvent( _ event: CodexReviewBackendModel.Review.Event, - run: CodexReviewBackendModel.Review.Run, + owner: ReviewActiveAttempt, subscriptionID: Int ) async { guard activeSubscriptionID == subscriptionID, @@ -1714,12 +2041,15 @@ private actor ReviewWorkerEventSource { } await queue.send(.reviewEvent(.init( subscriptionID: subscriptionID, - subscriptionRun: run, + owner: owner, event: event ))) } - private func yieldEventsFinished(run: CodexReviewBackendModel.Review.Run, subscriptionID: Int) async { + private func yieldEventsFinished( + owner: ReviewActiveAttempt, + subscriptionID: Int + ) async { guard activeSubscriptionID == subscriptionID, eventTasks.removeValue(forKey: subscriptionID) != nil else { @@ -1727,7 +2057,7 @@ private actor ReviewWorkerEventSource { } await queue.send(.reviewEventsFinished(.init( subscriptionID: subscriptionID, - run: run + owner: owner ))) } @@ -1740,7 +2070,7 @@ private actor ReviewWorkerEventSource { private func yieldEventsFailed( _ failure: ReviewAttemptStreamFailure, - run: CodexReviewBackendModel.Review.Run, + owner: ReviewActiveAttempt, subscriptionID: Int ) async { guard eventTasks.removeValue(forKey: subscriptionID) != nil else { @@ -1751,7 +2081,7 @@ private actor ReviewWorkerEventSource { } await queue.send(.reviewEventsFailed(.init( subscriptionID: subscriptionID, - run: run, + owner: owner, failure: failure ))) } diff --git a/Sources/CodexReview/Store/CodexReviewStoreTesting.swift b/Sources/CodexReview/Store/CodexReviewStoreTesting.swift index 33b9f10e..6aa63b48 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreTesting.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreTesting.swift @@ -84,6 +84,18 @@ extension CodexReviewStore { } package func cancelAndDrainReviewWorkersForTesting() async { + for job in orderedJobs where job.isTerminal == false { + do { + _ = try await cancelReview( + jobID: job.id, + cancellation: .system(message: "Test cleanup requested.") + ) + } catch { + reviewCleanupFailures[job.id] = .worker( + "Test cleanup cancellation failed: \(error.localizedDescription)" + ) + } + } let tasks = Array(reviewWorkerTasks.values) + Array(runtimeStopDetachedReviewWorkerTasks.values) for task in tasks { task.cancel() @@ -94,10 +106,8 @@ extension CodexReviewStore { reviewWorkerTasks.removeAll(keepingCapacity: false) runtimeStopDetachedReviewWorkerTasks.removeAll(keepingCapacity: false) - reviewStartAdmissions.removeAll(keepingCapacity: false) reviewCleanupFailures.removeAll(keepingCapacity: false) - activeRuns.removeAll(keepingCapacity: false) - reviewRecoveryWaitingJobIDs.removeAll(keepingCapacity: false) + reviewAttemptOwnerships.removeAll(keepingCapacity: false) let waiters = reviewTerminalWaiters.values.flatMap { $0 } reviewTerminalWaiters.removeAll(keepingCapacity: false) diff --git a/Sources/CodexReviewTesting/TestSupport.swift b/Sources/CodexReviewTesting/TestSupport.swift index 29d914ab..86f00f66 100644 --- a/Sources/CodexReviewTesting/TestSupport.swift +++ b/Sources/CodexReviewTesting/TestSupport.swift @@ -199,7 +199,8 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { private var interruptReviewGate: AsyncGate? private var interruptReviewWaiters: [UUID: CheckedContinuation] = [:] private var matchingInterruptReviewWaiters: [UUID: MatchingInterruptWaiter] = [:] - private var beginReviewRecoveryWaiters: [UUID: CheckedContinuation] = [:] + private var prepareReviewRecoveryWaiters: [UUID: CheckedContinuation] = [:] + private var prepareReviewRecoveryGate: AsyncGate? private var startReviewGate: AsyncGate? private var startReviewWaiters: [UUID: CheckedContinuation] = [:] private var resumeReviewRecoveryGate: AsyncGate? @@ -260,6 +261,10 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { resumeReviewRecoveryGate = gate } + package func holdPrepareReviewRecovery(with gate: AsyncGate) { + prepareReviewRecoveryGate = gate + } + package func setNextRecoveredRun(_ run: CodexReviewBackendModel.Review.Run) { nextRecoveredRun = run } @@ -377,7 +382,7 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { } } - package func waitForBeginReviewRecovery() async { + package func waitForPrepareReviewRecovery() async { if commands.contains(where: { if case .prepareReviewRecovery = $0 { true @@ -399,19 +404,19 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { }) { continuation.resume() } else { - beginReviewRecoveryWaiters[waiterID] = continuation + prepareReviewRecoveryWaiters[waiterID] = continuation } } } onCancel: { Task { - await self.cancelBeginReviewRecoveryWaiter(id: waiterID) + await self.cancelPrepareReviewRecoveryWaiter(id: waiterID) } } } - package func waitForBeginReviewRecovery(timeout: Duration = .seconds(2)) async throws { - try await withFakeBackendTimeout(operation: "beginReviewRecovery", timeout: timeout) { - await self.waitForBeginReviewRecovery() + package func waitForPrepareReviewRecovery(timeout: Duration = .seconds(2)) async throws { + try await withFakeBackendTimeout(operation: "prepareReviewRecovery", timeout: timeout) { + await self.waitForPrepareReviewRecovery() } } @@ -571,6 +576,7 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { package func releaseHeldOperationsForCleanup() async { await startReviewGate?.open() await interruptReviewGate?.open() + await prepareReviewRecoveryGate?.open() await resumeReviewRecoveryGate?.open() } @@ -578,11 +584,15 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { _ candidate: ReviewRecoveryCandidate ) async throws -> ReviewRecoveryHandoff { commands.append(.prepareReviewRecovery(candidate)) - let waiters = Array(beginReviewRecoveryWaiters.values) - beginReviewRecoveryWaiters.removeAll(keepingCapacity: false) + let waiters = Array(prepareReviewRecoveryWaiters.values) + prepareReviewRecoveryWaiters.removeAll(keepingCapacity: false) for waiter in waiters { waiter.resume() } + if let prepareReviewRecoveryGate { + await prepareReviewRecoveryGate.wait() + try Task.checkCancellation() + } let run = candidate.resolved.run return .init( candidate: candidate, @@ -692,8 +702,8 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { matchingInterruptReviewWaiters.removeValue(forKey: id)?.continuation.resume() } - private func cancelBeginReviewRecoveryWaiter(id: UUID) { - beginReviewRecoveryWaiters.removeValue(forKey: id)?.resume() + private func cancelPrepareReviewRecoveryWaiter(id: UUID) { + prepareReviewRecoveryWaiters.removeValue(forKey: id)?.resume() } private func cancelResumeReviewRecoveryWaiter(id: UUID) { @@ -726,7 +736,9 @@ package final class StoreSnapshotProbe { lastAgentMessage: job.core.output.lastAgentMessage, logs: job.logEntries, run: job.core.run, - activeRun: store.activeRuns[job.id], + attempt: attemptSnapshot( + from: store.reviewAttemptOwnerships[job.id] + ), cancellationRequested: job.cancellationRequested ) } @@ -762,7 +774,8 @@ package final class StoreSnapshotProbe { timeout: Duration = .seconds(2) ) async -> StoreSnapshot? { await waitUntil(timeout: timeout) { snapshot in - snapshot.job(jobID)?.activeRun?.attemptID == attemptID + snapshot.job(jobID)?.attempt?.phase == .active + && snapshot.job(jobID)?.attempt?.attemptID == attemptID } } @@ -783,6 +796,66 @@ package final class StoreSnapshotProbe { try? await Task.sleep(for: .milliseconds(10)) } } + +} + +@MainActor +private func attemptSnapshot( + from ownership: ReviewAttemptOwnership? +) -> StoreAttemptSnapshot? { + switch ownership { + case .initialStart(let start): + .init(phase: .initialStart, startID: start.id, attemptID: nil) + case .active(let active): + .init(phase: .active, startID: nil, attemptID: active.run.attemptID) + case .resolvingRecovery(let active): + .init(phase: .resolvingRecovery, startID: nil, attemptID: active.run.attemptID) + case .recoveryDisposition(let disposition): + .init( + phase: .recoveryDisposition, + startID: nil, + attemptID: disposition.resolvedAttempt.run.attemptID + ) + case .preparingRecovery(let candidate, _): + .init( + phase: .preparingRecovery, + startID: nil, + attemptID: candidate.resolved.run.attemptID + ) + case .waitingForRecovery(let handoff): + .init( + phase: .waitingForRecovery, + startID: nil, + attemptID: handoff.candidate.resolved.run.attemptID + ) + case .replacementStart(let handoff, let start): + .init( + phase: .replacementStart, + startID: start.id, + attemptID: handoff.candidate.resolved.run.attemptID + ) + case .terminal: + .init(phase: .terminal, startID: nil, attemptID: nil) + case nil: + nil + } +} + +package enum StoreAttemptPhase: Equatable, Sendable { + case initialStart + case active + case resolvingRecovery + case recoveryDisposition + case preparingRecovery + case waitingForRecovery + case replacementStart + case terminal +} + +package struct StoreAttemptSnapshot: Equatable, Sendable { + package var phase: StoreAttemptPhase + package var startID: ReviewStartHandleID? + package var attemptID: String? } package struct StoreSnapshot: Sendable { @@ -803,7 +876,7 @@ package struct StoreJobSnapshot: Sendable { package var lastAgentMessage: String? package var logs: [ReviewLogEntry] package var run: ReviewJobCore.Run - package var activeRun: CodexReviewBackendModel.Review.Run? + package var attempt: StoreAttemptSnapshot? package var cancellationRequested: Bool } @@ -814,6 +887,9 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { package var currentSettingsSnapshot: CodexReviewSettings.Snapshot package private(set) var isActive = false package private(set) var startRequests: [Bool] = [] + package private(set) var reviewStartOwnershipSnapshots: [StoreAttemptSnapshot?] = [] + package private(set) var recoveryResumeOwnershipSnapshots: [StoreAttemptSnapshot?] = [] + private weak var store: CodexReviewStore? package init( reviewBackend: FakeCodexReviewBackend, @@ -828,7 +904,9 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { currentSettingsSnapshot } - package func attachStore(_: CodexReviewStore) {} + package func attachStore(_ store: CodexReviewStore) { + self.store = store + } package func start(store: CodexReviewStore, forceRestartIfNeeded: Bool) async { startRequests.append(forceRestartIfNeeded) @@ -943,7 +1021,10 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { _ request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { - try await reviewBackend.startReview(request, admission: admission) + reviewStartOwnershipSnapshots.append(attemptSnapshot( + from: store?.reviewAttemptOwnerships[request.jobID] + )) + return try await reviewBackend.startReview(request, admission: admission) } package func interruptReview( @@ -968,7 +1049,10 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { - try await reviewBackend.resumeReviewRecovery( + recoveryResumeOwnershipSnapshots.append(attemptSnapshot( + from: store?.reviewAttemptOwnerships[request.jobID] + )) + return try await reviewBackend.resumeReviewRecovery( handoff, request: request, admission: admission diff --git a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift index 1301fcaf..706f0058 100644 --- a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift +++ b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift @@ -1328,8 +1328,23 @@ struct CodexReviewHostTests { } await transport.waitForRequest(method: "review/start") let jobID = try #require(store.jobs.first?.id) - let admission = try #require(store.reviewStartAdmissions[jobID]) + let admission: ReviewStartAdmission + switch store.reviewAttemptOwnerships[jobID] { + case .initialStart(let start): + admission = start.admission + case .active(let active): + admission = active.admission + default: + Issue.record("Expected initial or active attempt ownership before runtime stop.") + return + } let run = try #require(await admission.waitForActiveRun()) + try #require(await waitUntil(timeout: .seconds(2)) { + guard case .active(let active) = store.reviewAttemptOwnerships[jobID] else { + return false + } + return active.run.attemptID == run.attemptID + }) #expect(run.turnID == "turn-1") let worker = try #require(store.reviewWorkerTasks[jobID]) @@ -1410,7 +1425,11 @@ struct CodexReviewHostTests { ) ) try #require(await waitUntil(timeout: .seconds(2)) { - store.reviewRecoveryWaitingJobIDs.contains(jobID) + if case .waitingForRecovery = store.reviewAttemptOwnerships[jobID] { + true + } else { + false + } }) let stopFinished = CompletionFlag() diff --git a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift index 6181239e..e9639cd1 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift @@ -169,10 +169,35 @@ struct CodexReviewStoreCommandTests { } } + @Test func initialStartPublishesOwnershipBeforeBackendDispatch() async throws { + let backend = FakeCodexReviewBackend() + let storeBackend = TestingCodexReviewStoreBackend(reviewBackend: backend) + let store = CodexReviewStore.makeTestingStore( + backend: storeBackend, + idGenerator: .init(next: { "job-1" }) + ) + try await withStoreCommandTestCleanup(backend: backend, store: store) { + async let result = store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + try await backend.waitForStartReview(timeout: .seconds(2)) + + let snapshot = try #require(storeBackend.reviewStartOwnershipSnapshots.last ?? nil) + #expect(snapshot.phase == .initialStart) + #expect(snapshot.startID != nil) + #expect(snapshot.attemptID == nil) + + await backend.yield(.completed(summary: "Succeeded.", result: "review text")) + _ = try await result + } + } + @Test func reviewStartPreservesCanonicalResponseTurnAndMergesAgentMessageDeltas() async throws { let backend = FakeCodexReviewBackend() + let storeBackend = TestingCodexReviewStoreBackend(reviewBackend: backend) let store = CodexReviewStore.makeTestingStore( - backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + backend: storeBackend, idGenerator: .init(next: { "job-1" }) ) try await withStoreCommandTestCleanup(backend: backend, store: store) { @@ -925,7 +950,7 @@ struct CodexReviewStoreCommandTests { request: .init(cwd: "/tmp/project", target: .baseBranch("main")) ) try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running) != nil) - let run = try #require(store.activeRuns["job-1"]) + let run = try #require(activeAttempt(in: store, jobID: "job-1")?.run) let reason = CodexReviewBackendModel.CancellationReason( message: "Network unavailable; waiting to reconnect." ) @@ -935,12 +960,12 @@ struct CodexReviewStoreCommandTests { let commandsBeforeTerminal = await backend.recordedCommands() #expect(commandsBeforeTerminal.contains { if case .prepareReviewRecovery = $0 { true } else { false } } == false) - #expect(store.reviewRecoveryWaitingJobIDs.contains("job-1") == false) + #expect(isWaitingForRecovery(in: store, jobID: "job-1") == false) await backend.yield(.cancelled(reason.message), for: run) - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await backend.waitForPrepareReviewRecovery(timeout: .seconds(2)) let barrierPublished = await waitUntil(timeout: .seconds(2)) { - store.reviewRecoveryWaitingJobIDs.contains("job-1") + isWaitingForRecovery(in: store, jobID: "job-1") } #expect(barrierPublished) let waiting = try store.readReview(jobID: "job-1") @@ -972,7 +997,7 @@ struct CodexReviewStoreCommandTests { request: .init(cwd: "/tmp/project", target: .baseBranch("main")) ) try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running) != nil) - let run = try #require(store.activeRuns["job-1"]) + let run = try #require(activeAttempt(in: store, jobID: "job-1")?.run) let reason = CodexReviewBackendModel.CancellationReason( message: "Network unavailable; waiting to reconnect." ) @@ -984,9 +1009,9 @@ struct CodexReviewStoreCommandTests { let commandsBeforeAck = await backend.recordedCommands() #expect(commandsBeforeAck.contains { if case .prepareReviewRecovery = $0 { true } else { false } } == false) await interruptGate.open() - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await backend.waitForPrepareReviewRecovery(timeout: .seconds(2)) let barrierPublished = await waitUntil(timeout: .seconds(2)) { - store.reviewRecoveryWaitingJobIDs.contains("job-1") + isWaitingForRecovery(in: store, jobID: "job-1") } #expect(barrierPublished) @@ -1014,7 +1039,7 @@ struct CodexReviewStoreCommandTests { request: .init(cwd: "/tmp/project", target: .baseBranch("main")) ) try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running) != nil) - let run = try #require(store.activeRuns["job-1"]) + let run = try #require(activeAttempt(in: store, jobID: "job-1")?.run) let reason = CodexReviewBackendModel.CancellationReason( message: "Network unavailable; waiting to reconnect." ) @@ -1048,7 +1073,7 @@ struct CodexReviewStoreCommandTests { request: .init(cwd: "/tmp/project", target: .baseBranch("main")) ) try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running) != nil) - let run = try #require(store.activeRuns["job-1"]) + let run = try #require(activeAttempt(in: store, jobID: "job-1")?.run) let reason = CodexReviewBackendModel.CancellationReason( message: "Network unavailable; waiting to reconnect." ) @@ -1073,6 +1098,124 @@ struct CodexReviewStoreCommandTests { } } + @Test func naturalCompletionWinsCancellationJoinedToRecoveryBarrier() async throws { + let backend = FakeCodexReviewBackend() + let interruptGate = AsyncGate() + await backend.holdInterruptReview(with: interruptGate) + let networkMonitor = ManualCodexReviewNetworkMonitor() + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-1" }), + networkMonitor: networkMonitor, + networkRecoveryPolicy: .init(sleep: { _ in }) + ) + try await withStoreCommandTestCleanup(backend: backend, store: store) { + async let result = store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .baseBranch("main")) + ) + try #require(await waitUntil(timeout: .seconds(2)) { + activeAttempt(in: store, jobID: "job-1") != nil + }) + let run = try #require(activeAttempt(in: store, jobID: "job-1")?.run) + let reason = CodexReviewBackendModel.CancellationReason( + message: "Network unavailable; waiting to reconnect." + ) + + networkMonitor.yield(.init(status: .unsatisfied)) + try await backend.waitForInterruptReview( + run: run, + reason: reason, + timeout: .seconds(2) + ) + guard case .resolvingRecovery(let resolving) = + store.reviewAttemptOwnerships["job-1"] + else { + Issue.record("Expected resolvingRecovery ownership.") + return + } + let cancellation = Task { @MainActor in + try await store.cancelReview( + jobID: "job-1", + cancellation: .mcpClient(message: "Stop") + ) + } + #expect(await resolving.admission.waitForCancellationAdmission() == + .mcpClient(message: "Stop")) + await backend.yield( + .completed(summary: "Succeeded.", result: "natural review"), + for: run + ) + await interruptGate.open() + + let cancel = try await cancellation.value + let read = try await result + let commands = await backend.recordedCommands() + + #expect(cancel.cancelled == false) + #expect(read.core.lifecycle.status == .succeeded) + #expect(read.core.output.lastAgentMessage == "natural review") + #expect(commands.contains { + if case .prepareReviewRecovery = $0 { true } else { false } + } == false) + } + } + + @Test func cancellationDuringRecoveryPreparationDiscardsHandoffBeforeResume() async throws { + let backend = FakeCodexReviewBackend() + let preparationGate = AsyncGate() + await backend.holdPrepareReviewRecovery(with: preparationGate) + let networkMonitor = ManualCodexReviewNetworkMonitor() + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-1" }), + networkMonitor: networkMonitor, + networkRecoveryPolicy: .init(sleep: { _ in }) + ) + try await withStoreCommandTestCleanup(backend: backend, store: store) { + async let result = store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .baseBranch("main")) + ) + try #require(await waitUntil(timeout: .seconds(2)) { + activeAttempt(in: store, jobID: "job-1") != nil + }) + let run = try #require(activeAttempt(in: store, jobID: "job-1")?.run) + let reason = CodexReviewBackendModel.CancellationReason( + message: "Network unavailable; waiting to reconnect." + ) + + networkMonitor.yield(.init(status: .unsatisfied)) + try await backend.waitForInterruptReview( + run: run, + reason: reason, + timeout: .seconds(2) + ) + await backend.yield(.cancelled(reason.message), for: run) + try await backend.waitForPrepareReviewRecovery(timeout: .seconds(2)) + try #require(await waitUntil(timeout: .seconds(2)) { + if case .preparingRecovery = store.reviewAttemptOwnerships["job-1"] { + true + } else { + false + } + }) + + let cancel = try await store.cancelReview( + jobID: "job-1", + cancellation: .mcpClient(message: "Stop") + ) + let read = try await result + let commands = await backend.recordedCommands() + + #expect(cancel.cancelled) + #expect(read.core.lifecycle.status == .cancelled) + #expect(commands.contains { + if case .resumeReviewRecovery = $0 { true } else { false } + } == false) + } + } + @Test func networkRecoveryWaitDiscardsOldAttemptCompletion() async throws { let initialRun = CodexReviewBackendModel.Review.Run( threadID: "thread-1", @@ -1576,8 +1719,9 @@ struct CodexReviewStoreCommandTests { let recoverGate = AsyncGate() await backend.holdResumeReviewRecovery(with: recoverGate) let networkMonitor = ManualCodexReviewNetworkMonitor() + let storeBackend = TestingCodexReviewStoreBackend(reviewBackend: backend) let store = CodexReviewStore.makeTestingStore( - backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + backend: storeBackend, idGenerator: .init(next: { "job-1" }), networkMonitor: networkMonitor, networkRecoveryPolicy: .init(sleep: { _ in }) @@ -1667,8 +1811,9 @@ struct CodexReviewStoreCommandTests { let recoverGate = AsyncGate() await backend.holdResumeReviewRecovery(with: recoverGate) let networkMonitor = ManualCodexReviewNetworkMonitor() + let storeBackend = TestingCodexReviewStoreBackend(reviewBackend: backend) let store = CodexReviewStore.makeTestingStore( - backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + backend: storeBackend, idGenerator: .init(next: { "job-1" }), networkMonitor: networkMonitor, networkRecoveryPolicy: .init(sleep: { _ in }) @@ -1679,14 +1824,30 @@ struct CodexReviewStoreCommandTests { request: .init(cwd: "/tmp/project", target: .baseBranch("main")) ) + try #require(await waitUntil(timeout: .seconds(2)) { + activeAttempt(in: store, jobID: "job-1") != nil + }) + let oldAdmission = try #require(activeAttempt(in: store, jobID: "job-1")?.admission) networkMonitor.yield(.init(status: .unsatisfied)) try await completeNetworkRecoveryBarrier(backend: backend, store: store) - let oldAdmission = try #require(store.reviewStartAdmissions["job-1"]) networkMonitor.yield(.satisfied()) try await backend.waitForResumeReviewRecovery(timeout: .seconds(2)) - let recoveredAdmission = try #require(store.reviewStartAdmissions["job-1"]) + guard case .replacementStart(let handoff, let recoveredStart) = + store.reviewAttemptOwnerships["job-1"] + else { + Issue.record("Expected coherent replacementStart ownership during rollback.") + return + } + let recoveredAdmission = recoveredStart.admission #expect(recoveredAdmission !== oldAdmission) + #expect(handoff.candidate.resolved.run.attemptID == initialRun.attemptID) + #expect(activeAttempt(in: store, jobID: "job-1") == nil) + let resumeSnapshot = try #require( + storeBackend.recoveryResumeOwnershipSnapshots.last ?? nil + ) + #expect(resumeSnapshot.phase == .replacementStart) + #expect(resumeSnapshot.startID == recoveredStart.id) let cancelTask = Task { @MainActor in try await store.cancelReview( jobID: "job-1", @@ -1750,7 +1911,8 @@ struct CodexReviewStoreCommandTests { let cancel = try await store.cancelReview(jobID: "job-1", cancellation: .mcpClient(message: "Stop")) let cleanedUp = await waitUntil { - store.reviewWorkerTasks["job-1"] == nil && store.activeRuns["job-1"] == nil + store.reviewWorkerTasks["job-1"] == nil + && store.reviewAttemptOwnerships["job-1"] == nil } let read = try store.readReview(jobID: "job-1") @@ -1790,12 +1952,12 @@ struct CodexReviewStoreCommandTests { #expect(locallyCancelledJobIDs == ["job-1"]) #expect(cancelled.core.lifecycle.status == .cancelled) #expect(store.reviewWorkerTasks["job-1"] != nil) - #expect(store.activeRuns["job-1"] == run) + #expect(activeAttempt(in: store, jobID: "job-1")?.run == run) store.cancelAndDetachReviewWorkersForRuntimeStop(jobIDs: locallyCancelledJobIDs) #expect(store.reviewWorkerTasks["job-1"] == nil) - #expect(store.activeRuns["job-1"] == nil) + #expect(activeAttempt(in: store, jobID: "job-1")?.run == run) } } @@ -1837,7 +1999,7 @@ struct CodexReviewStoreCommandTests { let commands = await backend.recordedCommands() #expect(commands.contains(.interruptReview(run, .init(message: "Review runtime stopped.")))) #expect(stopped.core.lifecycle.status == .cancelled) - #expect(store.activeRuns["job-1"] == nil) + #expect(store.reviewAttemptOwnerships["job-1"] == nil) #expect(store.reviewWorkerTasks["job-1"] == nil) } } @@ -1875,8 +2037,7 @@ struct CodexReviewStoreCommandTests { store.cancelAndDetachReviewWorkersForRuntimeStop(jobIDs: locallyCancelledJobIDs) #expect(store.reviewWorkerTasks["job-1"] == nil) - #expect(store.activeRuns["job-1"] == nil) - #expect(store.reviewRecoveryWaitingJobIDs.contains("job-1") == false) + #expect(isWaitingForRecovery(in: store, jobID: "job-1")) } } @@ -1935,13 +2096,15 @@ struct CodexReviewStoreCommandTests { ) store.cancelAndDetachReviewWorkersForRuntimeStop(jobIDs: locallyCancelledJobIDs) let resultBeforeStartReviewUnblocked = try await waitForTaskValue(running, timeout: .seconds(1)) + #expect(store.reviewAttemptOwnerships["job-1"] != nil) await startReviewGate.open() let result = try #require(resultBeforeStartReviewUnblocked) + #expect(await store.drainRuntimeStopDetachedReviewWorkers(timeout: .seconds(2))) #expect(locallyCancelledJobIDs == ["job-1"]) #expect(result.core.lifecycle.status == .cancelled) #expect(store.reviewWorkerTasks["job-1"] == nil) - #expect(store.activeRuns["job-1"] == nil) + #expect(store.reviewAttemptOwnerships["job-1"] == nil) } } @@ -2143,6 +2306,7 @@ struct CodexReviewStoreCommandTests { workspaces: [.init(cwd: "/tmp/project")], jobs: [running] ) + try await seedQueuedAttemptOwnership(in: store, for: running) let cancel = try await store.cancelReview( jobID: "job-1", @@ -2178,6 +2342,10 @@ struct CodexReviewStoreCommandTests { sessionID: "session-1", request: .init(cwd: "/tmp/project", target: .baseBranch("main")) ) + try #require(await StoreSnapshotProbe(store: store).waitUntilRunAttempt( + "attempt-1", + jobID: "job-1" + ) != nil) await #expect(throws: (any Error).self) { try await store.cancelReview( @@ -2230,6 +2398,48 @@ struct CodexReviewStoreCommandTests { } } + @Test func acknowledgedCancellationForceCloseCommitsRequestedTerminal() async throws { + let backend = FakeCodexReviewBackend() + let graceGate = AsyncGate() + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-1" }), + reviewRuntimeClosePolicy: .init( + terminalGrace: .seconds(10), + sleep: { _ in + await graceGate.wait() + try Task.checkCancellation() + } + ) + ) + try await withStoreCommandTestCleanup(backend: backend, store: store) { + async let result = store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .baseBranch("main")) + ) + try #require(await StoreSnapshotProbe(store: store).waitUntilRunAttempt( + "attempt-1", + jobID: "job-1" + ) != nil) + async let cancellation = store.cancelReview( + jobID: "job-1", + cancellation: .mcpClient(message: "Stop") + ) + await backend.waitForInterruptReview() + await graceGate.open() + + let cancel = try await cancellation + let read = try await result + + #expect(cancel.cancelled) + #expect(read.core.lifecycle.status == .cancelled) + #expect(read.core.lifecycle.terminal == .interrupted(.requested( + .mcpClient(message: "Stop") + ))) + #expect(await backend.recordedCommands().contains(.forceCloseReviewConnection)) + } + } + @Test func failedReviewPreservesBufferedEventsBeforeStreamError() async throws { let backend = FakeCodexReviewBackend() let store = CodexReviewStore.makeTestingStore( @@ -2296,7 +2506,7 @@ struct CodexReviewStoreCommandTests { networkMonitor.yield(.init(status: .unsatisfied)) await outageSleepStarted.wait() await backend.finishEvents( - throwing: .unexpectedConnection(.connection("Review event stream closed.")), + throwing: .recoverableNetwork(.connection("Verified network loss.")), for: initialRun ) @@ -2321,6 +2531,52 @@ struct CodexReviewStoreCommandTests { } } + @Test func unsatisfiedNetworkDoesNotRecoverTypedNonrecoverableFailures() async throws { + let failures: [ReviewAttemptStreamFailure] = [ + .unexpectedConnection(.connection("Unexpected connection loss.")), + .ownerForcedConnectionClose(.connection("Uncorrelated owner close.")), + .process(.process("Process exited.")), + .protocolViolation(.init(message: "Malformed notification.")), + .workerContract(.init(message: "Worker contract failed.")), + .ownerCancellation, + ] + + for (index, failure) in failures.enumerated() { + let backend = FakeCodexReviewBackend() + let networkMonitor = ManualCodexReviewNetworkMonitor() + let outageGate = AsyncGate() + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-\(index)" }), + networkMonitor: networkMonitor, + networkRecoveryPolicy: .init(sleep: { _ in await outageGate.wait() }) + ) + try await withStoreCommandTestCleanup(backend: backend, store: store) { + async let result = store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .baseBranch("main")) + ) + try #require(await StoreSnapshotProbe(store: store).waitUntilRunAttempt( + "attempt-1", + jobID: "job-\(index)" + ) != nil) + + networkMonitor.yield(.init(status: .unsatisfied)) + await backend.finishEvents(throwing: failure) + let read = try await result + let commands = await backend.recordedCommands() + + #expect(read.core.lifecycle.status == .failed) + #expect(commands.contains { + if case .prepareReviewRecovery = $0 { true } else { false } + } == false) + #expect(commands.contains { + if case .resumeReviewRecovery = $0 { true } else { false } + } == false) + } + } + } + @Test func eventConsumerCancellationIsAWorkerTerminalNotARequestedInterrupt() async throws { let backend = FakeCodexReviewBackend() let store = CodexReviewStore.makeTestingStore( @@ -2337,7 +2593,8 @@ struct CodexReviewStoreCommandTests { let read = try await result #expect(read.core.lifecycle.status == .failed) - #expect(read.core.lifecycle.terminal?.kind == .interrupted) + #expect(read.core.lifecycle.terminal?.kind == .failed) + #expect(read.core.lifecycle.errorMessage == ReviewAttemptStreamFailure.ownerCancellation.localizedDescription) let commands = await backend.recordedCommands() #expect(commands.contains { if case .interruptReview = $0 { true } else { false } } == false) } @@ -2532,6 +2789,7 @@ struct CodexReviewStoreCommandTests { workspaces: [.init(cwd: "/tmp/project")], jobs: [running] ) + try await seedQueuedAttemptOwnership(in: store, for: running) await store.closeActiveReviewSessions(reason: .system(message: "Account switched.")) @@ -2775,12 +3033,12 @@ private func completeNetworkRecoveryBarrier( beforeTerminal: @MainActor (CodexReviewBackendModel.Review.Run) async -> Void = { _ in } ) async throws { let becameActive = await waitUntil(timeout: .seconds(2)) { - store.activeRuns[jobID] != nil + activeAttempt(in: store, jobID: jobID) != nil } guard becameActive else { throw FakeCodexReviewBackendError(message: "Review did not become active before recovery.") } - guard let run = store.activeRuns[jobID] else { + guard let run = activeAttempt(in: store, jobID: jobID)?.run else { throw FakeCodexReviewBackendError(message: "Missing active run for recovery barrier.") } let reason = CodexReviewBackendModel.CancellationReason( @@ -2790,16 +3048,16 @@ private func completeNetworkRecoveryBarrier( await beforeTerminal(run) await backend.yield(.cancelled(reason.message), for: run) do { - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await backend.waitForPrepareReviewRecovery(timeout: .seconds(2)) } catch { - let phase = await store.reviewStartAdmissions[jobID]?.currentPhase() + let phase = StoreSnapshotProbe(store: store).snapshot().job(jobID)?.attempt?.phase let commands = await backend.recordedCommands() throw FakeCodexReviewBackendError( message: "Recovery preparation was not admitted; phase=\(String(describing: phase)), commands=\(commands)." ) } let barrierPublished = await waitUntil(timeout: .seconds(2)) { - store.reviewRecoveryWaitingJobIDs.contains(jobID) + isWaitingForRecovery(in: store, jobID: jobID) } guard barrierPublished else { throw FakeCodexReviewBackendError(message: "Store did not publish recovery barrier.") @@ -2812,15 +3070,55 @@ private func completeNetworkRecoveryAfterConnectionTerminal( store: CodexReviewStore, jobID: String = "job-1" ) async throws { - try await backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await backend.waitForPrepareReviewRecovery(timeout: .seconds(2)) let barrierPublished = await waitUntil(timeout: .seconds(2)) { - store.reviewRecoveryWaitingJobIDs.contains(jobID) + isWaitingForRecovery(in: store, jobID: jobID) } guard barrierPublished else { throw FakeCodexReviewBackendError(message: "Store did not publish connection recovery barrier.") } } +@MainActor +private func activeAttempt( + in store: CodexReviewStore, + jobID: String +) -> ReviewActiveAttempt? { + switch store.reviewAttemptOwnerships[jobID] { + case .active(let active), .resolvingRecovery(let active): + active + case .initialStart, .recoveryDisposition, .preparingRecovery, + .waitingForRecovery, .replacementStart, .terminal, nil: + nil + } +} + +@MainActor +private func isWaitingForRecovery( + in store: CodexReviewStore, + jobID: String +) -> Bool { + if case .waitingForRecovery = store.reviewAttemptOwnerships[jobID] { + true + } else { + false + } +} + +@MainActor +private func seedQueuedAttemptOwnership( + in store: CodexReviewStore, + for job: CodexReviewJob +) async throws { + let admission = ReviewStartAdmission() + let registered = try await admission.registerStart { _ in + throw ReviewAttemptContractFailure( + message: "A queued fixture start must remain backend-inert." + ) + } + store.reviewAttemptOwnerships[job.id] = .initialStart(registered) +} + @MainActor private func cleanupStoreCommandTest( backend: FakeCodexReviewBackend, diff --git a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift index 82beee9d..71558df2 100644 --- a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift +++ b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift @@ -540,6 +540,25 @@ struct ReviewAttemptProcessorTests { #expect(await backendWrite.invocationCount() == 0) } + @Test func streamTerminalBeforeRegistrationSurvivesRegistrationWithoutWrite() async throws { + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) + let failure = ReviewAttemptStreamFailure.unexpectedConnection( + .connection("Connection ended") + ) + try await admission.recordStreamTerminal(failure) + let backendWrite = InvocationProbe() + + let registered = try await admission.registerStart { _ in + await backendWrite.record() + return .init(run: canonicalRun) + } + + await #expect(throws: failure) { + _ = try await registered.task.value + } + #expect(await backendWrite.invocationCount() == 0) + } + @Test func activationIsIdempotentForLiveHandleAndRejectsWrongHandle() async throws { let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) let operationGate = AsyncGate() @@ -556,6 +575,9 @@ struct ReviewAttemptProcessorTests { try await admission.activateStart(registered.id) await operationGate.open() _ = try await registered.task.value + await #expect(throws: ReviewAttemptContractFailure.self) { + try await admission.activateStart(registered.id) + } } @Test func joinedExplicitCancellationInstallsProductDispositionBeforeRecoveryPreparation() async throws { @@ -598,6 +620,48 @@ struct ReviewAttemptProcessorTests { #expect(await requestStarted.invocationCount() == 1) } + @Test func joinedCancellationOutcomeUnknownConnectionRetainsTransportProduct() async throws { + let (admission, run) = try await makeActiveAdmission() + let requestFailed = InvocationProbe() + let requestFailure = ReviewInterruptRequestFailure( + outcome: .outcomeUnknown(message: "Response lost") + ) + let connection = ReviewRuntimeCloseFailure.connection("Connection ended") + let recovery = Task { + try await admission.beginRecovery( + trigger: .recoverableNetworkLoss, + interrupt: { _, _ in + await requestFailed.record() + throw requestFailure + }, + forceClose: {} + ) + } + await requestFailed.waitForInvocation() + let cancellation = Task { + try await admission.cancel( + .mcpClient(message: "Stop"), + interrupt: { _, _ in Issue.record("Joined cancellation sent a second request.") }, + forceClose: {} + ) + } + #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) + try await admission.recordStreamTerminal(.recoverableNetwork(connection)) + + guard case .productTerminal(let product) = try await recovery.value else { + Issue.record("Outcome-unknown connection plus joined cancel must terminalize product.") + return + } + #expect(product.resolved.run == run) + #expect(product.resolved.requestFailure?.outcome == requestFailure.outcome) + #expect(product.productTerminal == .interrupted(.transport( + message: connection.localizedDescription + ))) + await #expect(throws: ReviewInterruptRequestFailure.self) { + _ = try await cancellation.value + } + } + @Test func recoveryClassifiesTypedStreamFailureBeforeTokenization() async throws { let (recoverableAdmission, recoverableRun) = try await makeActiveAdmission() let recoverableRequest = InvocationProbe() diff --git a/Tests/CodexReviewTests/ReviewTerminalContractTests.swift b/Tests/CodexReviewTests/ReviewTerminalContractTests.swift index 17aa1b09..b4028a06 100644 --- a/Tests/CodexReviewTests/ReviewTerminalContractTests.swift +++ b/Tests/CodexReviewTests/ReviewTerminalContractTests.swift @@ -180,7 +180,7 @@ struct ReviewTerminalContractTests { #expect(result.core.reviewText != "arbitrary last message") } - @Test func streamEOFUsesTypedTransportInterruptionWithoutCanonicalTerminal() async throws { + @Test func streamEOFUsesTypedWorkerContractFailureWithoutCanonicalTerminal() async throws { let (store, backend) = makeStore() async let started = store.startReview( sessionID: "session-1", @@ -191,12 +191,12 @@ struct ReviewTerminalContractTests { let result = try await started #expect(result.core.lifecycle.status == .failed) - #expect(result.core.lifecycle.terminal?.kind == .interrupted) - guard case .interrupted(.transport(let message)) = result.core.lifecycle.terminal else { - Issue.record("Expected a typed transport interruption.") + #expect(result.core.lifecycle.terminal?.kind == .failed) + guard case .failed(let message) = result.core.lifecycle.terminal else { + Issue.record("Expected a nonrecoverable worker-contract failure.") return } - #expect(message.contains("authoritative terminal")) + #expect(message?.contains("authoritative terminal") == true) #expect(result.core.lifecycle.errorMessage?.contains("authoritative terminal") == true) } diff --git a/Tests/ReviewUITests/ReviewUITests.swift b/Tests/ReviewUITests/ReviewUITests.swift index f79a971c..b8cce071 100644 --- a/Tests/ReviewUITests/ReviewUITests.swift +++ b/Tests/ReviewUITests/ReviewUITests.swift @@ -1765,6 +1765,13 @@ struct ReviewUITests { serverState: .running, content: makeSidebarContent(from: [job]) ) + let admission = ReviewStartAdmission() + let registered = try await admission.registerStart { _ in + throw ReviewAttemptContractFailure( + message: "Sidebar cancellation fixture must remain backend-inert." + ) + } + store.reviewAttemptOwnerships[job.id] = .initialStart(registered) let viewController = ReviewMonitorSplitViewController(store: store, uiState: ReviewMonitorUIState(auth: store.auth)) viewController.loadViewIfNeeded() @@ -1799,7 +1806,10 @@ struct ReviewUITests { ) let admission = ReviewStartAdmission() await admission.recordActiveRun(run) - store.reviewStartAdmissions[job.id] = admission + store.reviewAttemptOwnerships[job.id] = .active(.init( + run: run, + admission: admission + )) let viewController = ReviewMonitorSplitViewController(store: store, uiState: ReviewMonitorUIState(auth: store.auth)) viewController.loadViewIfNeeded() From 8908dc410982360f3591734f20c40be3219c711a Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 21:37:04 +0900 Subject: [PATCH 18/38] test(app-server): expect typed routing failures --- .../CurrentV2ReviewContractTests.swift | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/Tests/CodexReviewAppServerTests/CurrentV2ReviewContractTests.swift b/Tests/CodexReviewAppServerTests/CurrentV2ReviewContractTests.swift index 98dd7c75..7d88b771 100644 --- a/Tests/CodexReviewAppServerTests/CurrentV2ReviewContractTests.swift +++ b/Tests/CodexReviewAppServerTests/CurrentV2ReviewContractTests.swift @@ -1214,9 +1214,11 @@ struct CurrentV2ReviewRoutingIntegrationTests { ) ) - #expect(try await failedAttempt.events.next() == .failed( - "Malformed app-server notification item/completed: id must be a nonempty string" - )) + await #expect(throws: ReviewAttemptStreamFailure.protocolViolation(.init( + message: "Malformed app-server notification item/completed: id must be a nonempty string" + ))) { + _ = try await failedAttempt.events.next() + } #expect(try await collectEvents(from: healthyAttempt.events).last == .completed( summary: "Succeeded.", result: "Healthy review" @@ -1265,9 +1267,11 @@ struct CurrentV2ReviewRoutingIntegrationTests { ) ) - #expect(try await failedAttempt.events.next() == .failed( - "Unsupported app-server item type futureItem in item/completed." - )) + await #expect(throws: ReviewAttemptStreamFailure.protocolViolation(.init( + message: "Unsupported app-server item type futureItem in item/completed." + ))) { + _ = try await failedAttempt.events.next() + } #expect(try await collectEvents(from: healthyAttempt.events).last == .completed( summary: "Succeeded.", result: "Healthy review" From c47b4d0a805ff3832b9d9442885cbea3a8a61d02 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 20 Aug 2026 21:39:27 +0900 Subject: [PATCH 19/38] test(mcp): seed attempt ownership for cancellation --- .../CodexReviewMCPHTTPServerTests.swift | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift index ee561073..c6a35eb2 100644 --- a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift +++ b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift @@ -868,6 +868,7 @@ struct CodexReviewMCPHTTPServerTests { workspaces: [.init(cwd: "/tmp/project")], jobs: [completed, running] ) + try await seedQueuedAttemptOwnership(in: store, for: running) let response = try await postJSONRPC( endpoint: await server.url, @@ -976,6 +977,7 @@ struct CodexReviewMCPHTTPServerTests { workspaces: [.init(cwd: "/tmp/project")], jobs: [running] ) + try await seedQueuedAttemptOwnership(in: store, for: running) let response = try await postJSONRPC( endpoint: await server.url, @@ -1192,6 +1194,7 @@ struct CodexReviewMCPHTTPServerTests { workspaces: [.init(cwd: "/tmp/project")], jobs: [running] ) + try await seedQueuedAttemptOwnership(in: store, for: running) let response = try await deleteSession(endpoint: await server.url, sessionID: sessionID) @@ -1491,6 +1494,19 @@ struct CodexReviewMCPHTTPServerTests { } return true } + + private func seedQueuedAttemptOwnership( + in store: CodexReviewStore, + for job: CodexReviewJob + ) async throws { + let admission = ReviewStartAdmission() + let registered = try await admission.registerStart { _ in + throw ReviewAttemptContractFailure( + message: "An MCP cancellation fixture must remain backend-inert." + ) + } + store.reviewAttemptOwnerships[job.id] = .initialStart(registered) + } } private nonisolated func currentPOSIXError() -> NSError { From 2bb11f1dde11d95f2a9c244c2d74ad72930fe10a Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 01:42:12 +0900 Subject: [PATCH 20/38] fix(review): close attempt cancellation state gaps --- .../CodexReview/ReviewAttemptProcessor.swift | 118 +++++++++++---- .../AppServerCodexReviewBackend.swift | 22 ++- Sources/CodexReviewTesting/TestSupport.swift | 20 +-- .../AppServerClientTests.swift | 4 +- .../ReviewAttemptProcessorTests.swift | 138 +++++++++++++----- 5 files changed, 208 insertions(+), 94 deletions(-) diff --git a/Sources/CodexReview/ReviewAttemptProcessor.swift b/Sources/CodexReview/ReviewAttemptProcessor.swift index d0b4cff0..63671798 100644 --- a/Sources/CodexReview/ReviewAttemptProcessor.swift +++ b/Sources/CodexReview/ReviewAttemptProcessor.swift @@ -363,6 +363,7 @@ package actor ReviewStartAdmission { private var cancellationResult: Result? private var terminalWaiters: [UUID: CheckedContinuation] = [:] private var activeRunWaiters: [CheckedContinuation] = [] + private var startResolutionWaiters: [CheckedContinuation] = [] private var cancellationAdmissionWaiters: [CheckedContinuation] = [] private var interruptionAdmissionWaiters: [CheckedContinuation] = [] private var cancellationWaiters: [CheckedContinuation, Never>] = [] @@ -433,19 +434,19 @@ package actor ReviewStartAdmission { resolveStartActivation(.success(())) } - package func admitThreadStartDispatch() -> Bool { - guard requestedCancellation == nil else { - return false + package func admitThreadStartDispatch() throws { + if let terminal { + throw startFailure(for: terminal) } - switch phase { - case .preparingThread(.notSent): - phase = .preparingThread(.outcomeUnknown) - return true - case .queued, .registeredStart, .activatedStart, - .preparingThread(.outcomeUnknown), .startingReview, - .active, .interrupting, .finishing, .terminal: - return false + if let requestedCancellation { + throw ReviewStartCancelledBeforeDispatch(cancellation: requestedCancellation) + } + guard case .preparingThread(.notSent) = phase else { + throw ReviewAttemptContractFailure( + message: "Thread start dispatch requires one pending not-sent request." + ) } + phase = .preparingThread(.outcomeUnknown) } package func recordThreadStartRejectedForRetry() throws { @@ -474,17 +475,21 @@ package actor ReviewStartAdmission { package func admitReviewStartDispatch( for preparedRun: CodexReviewBackendModel.Review.Run - ) -> Bool { - guard requestedCancellation == nil else { - return false + ) throws { + if let terminal { + throw startFailure(for: terminal) + } + if let requestedCancellation { + throw ReviewStartCancelledBeforeDispatch(cancellation: requestedCancellation) } guard case .startingReview(let currentRun, .notSent) = phase, currentRun.attemptID == preparedRun.attemptID else { - return false + throw ReviewAttemptContractFailure( + message: "Review start dispatch requires its pending prepared attempt." + ) } phase = .startingReview(preparedRun: preparedRun, dispatch: .outcomeUnknown) - return true } package func recordActiveRun(_ run: CodexReviewBackendModel.Review.Run) { @@ -494,6 +499,7 @@ package actor ReviewStartAdmission { registeredRun = run phase = .active(run) resumeActiveRunWaiters(returning: run) + resumeStartResolutionWaitersIfNeeded() } package func recordCanonicalTerminal( @@ -542,12 +548,17 @@ package actor ReviewStartAdmission { forceClose: @escaping @Sendable () async throws -> Void ) async throws -> ReviewAttemptCancellationResolution { if let recoveryDispositionTask { - joinedTerminalCancellation = cancellation - requestedCancellation = cancellation - resumeCancellationAdmissionWaiters(returning: cancellation) + if joinedTerminalCancellation == nil { + joinedTerminalCancellation = cancellation + requestedCancellation = cancellation + resumeCancellationAdmissionWaiters(returning: cancellation) + } let disposition = try await recoveryDispositionTask.value return try cancellationResolution(for: disposition) } + if let cancellationTask { + return try checkedCancellationResolution(try await cancellationTask.value) + } interruptionPurpose = .terminalCancellation(cancellation) resumeInterruptionAdmissionWaiters(returning: interruptionPurpose) let resolution = try await joinedCancellationResolution( @@ -555,15 +566,7 @@ package actor ReviewStartAdmission { interrupt: interrupt, forceClose: forceClose ) - if let requestFailure = resolution.requestFailure, - case .outcomeUnknown = requestFailure.outcome, - case .stream(let streamFailure) = resolution.terminal { - throw ReviewInterruptRequestFailure( - outcome: requestFailure.outcome, - secondaryBarrierDiagnostic: streamFailure.localizedDescription - ) - } - return resolution + return try checkedCancellationResolution(resolution) } package func beginRecovery( @@ -951,7 +954,8 @@ package actor ReviewStartAdmission { } case .failure(let error): startFailed = true - if let cancellation = (error as? ReviewStartCancelledBeforeDispatch)?.cancellation { + if terminal == nil, + let cancellation = (error as? ReviewStartCancelledBeforeDispatch)?.cancellation { receiveTerminal(.localCancellation(cancellation)) } else if error is CancellationError, let requestedCancellation { @@ -967,6 +971,7 @@ package actor ReviewStartAdmission { resumeCancellationAdmissionWaiters(returning: nil) resumeInterruptionAdmissionWaiters(returning: nil) } + resumeStartResolutionWaitersIfNeeded() } private func performCancellation( @@ -989,8 +994,13 @@ package actor ReviewStartAdmission { installGraceTask(forceClose: forceClose) } - if let startTask, activeRun == nil, terminal == nil { - _ = await startTask.result + if startTask != nil, activeRun == nil, terminal == nil { + await waitForStartResolution() + } + + if case .failure(let error)? = cancellationResult { + await drainCancellationTasks() + throw error } if let terminal { @@ -1066,6 +1076,8 @@ package actor ReviewStartAdmission { private func resetRejectedCancellationForRetry() { requestedCancellation = nil + joinedTerminalCancellation = nil + interruptionPurpose = nil cancellationTask = nil interruptRequestTask = nil terminalBarrierTask = nil @@ -1203,6 +1215,7 @@ package actor ReviewStartAdmission { resumeActiveRunWaiters(returning: nil) resumeCancellationAdmissionWaiters(returning: nil) resumeInterruptionAdmissionWaiters(returning: nil) + resumeStartResolutionWaitersIfNeeded() resolveCancellationIfPossible() } @@ -1305,6 +1318,7 @@ package actor ReviewStartAdmission { _ result: Result ) { cancellationResult = result + resumeStartResolutionWaitersIfNeeded() let waiters = cancellationWaiters cancellationWaiters.removeAll(keepingCapacity: false) for waiter in waiters { @@ -1332,6 +1346,48 @@ package actor ReviewStartAdmission { } } + private func waitForStartResolution() async { + if hasStartResolution { + return + } + await withCheckedContinuation { continuation in + if hasStartResolution { + continuation.resume() + } else { + startResolutionWaiters.append(continuation) + } + } + } + + private func resumeStartResolutionWaitersIfNeeded() { + guard hasStartResolution else { + return + } + let waiters = startResolutionWaiters + startResolutionWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume() + } + } + + private var hasStartResolution: Bool { + activeRun != nil || terminal != nil || startDidFinish || cancellationResult != nil + } + + private func checkedCancellationResolution( + _ resolution: ReviewAttemptCancellationResolution + ) throws -> ReviewAttemptCancellationResolution { + if let requestFailure = resolution.requestFailure, + case .outcomeUnknown = requestFailure.outcome, + case .stream(let streamFailure) = resolution.terminal { + throw ReviewInterruptRequestFailure( + outcome: requestFailure.outcome, + secondaryBarrierDiagnostic: streamFailure.localizedDescription + ) + } + return resolution + } + private func cancelTerminalWaiter(_ id: UUID) { terminalWaiters.removeValue(forKey: id)?.resume(returning: nil) } diff --git a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift index b670e848..ae8c2dae 100644 --- a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift +++ b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift @@ -153,11 +153,11 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { control.recordThreadStarted(threadID: thread.threadID) await admission.recordPreparedThread(provisionalRun) - guard await admission.admitReviewStartDispatch(for: provisionalRun) else { + do { + try await admission.admitReviewStartDispatch(for: provisionalRun) + } catch { try await cleanupReview(provisionalRun) - throw ReviewStartCancelledBeforeDispatch( - cancellation: await admission.cancellationRequest() ?? .system() - ) + throw error } let review: AppServerAPI.Review.Start.Response @@ -262,11 +262,7 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { _ request: AppServerAPI.Thread.Start.Request, admission: ReviewStartAdmission ) async throws -> AppServerAPI.Thread.Start.Response { - guard await admission.admitThreadStartDispatch() else { - throw ReviewStartCancelledBeforeDispatch( - cancellation: await admission.cancellationRequest() ?? .system() - ) - } + try await admission.admitThreadStartDispatch() do { return try await client.send(request) } catch { @@ -440,12 +436,12 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { ) registerReviewEventSession(session, for: provisionalRun) await admission.recordPreparedThread(provisionalRun) - guard await admission.admitReviewStartDispatch(for: provisionalRun) else { + do { + try await admission.admitReviewStartDispatch(for: provisionalRun) + } catch { _ = unregisterReviewEventSession(for: provisionalRun) await session.abandon() - throw ReviewStartCancelledBeforeDispatch( - cancellation: await admission.cancellationRequest() ?? .system() - ) + throw error } let review: AppServerAPI.Review.Start.Response diff --git a/Sources/CodexReviewTesting/TestSupport.swift b/Sources/CodexReviewTesting/TestSupport.swift index 86f00f66..f3399f7a 100644 --- a/Sources/CodexReviewTesting/TestSupport.swift +++ b/Sources/CodexReviewTesting/TestSupport.swift @@ -506,11 +506,7 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { _ request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { - guard await admission.admitThreadStartDispatch() else { - throw ReviewStartCancelledBeforeDispatch( - cancellation: await admission.cancellationRequest() ?? .system() - ) - } + try await admission.admitThreadStartDispatch() commands.append(.startReview(request)) let waiters = Array(startReviewWaiters.values) startReviewWaiters.removeAll(keepingCapacity: false) @@ -524,11 +520,11 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { model: nextRun.model ) await admission.recordPreparedThread(provisionalRun) - guard await admission.admitReviewStartDispatch(for: provisionalRun) else { + do { + try await admission.admitReviewStartDispatch(for: provisionalRun) + } catch { commands.append(.cleanupReview(provisionalRun)) - throw ReviewStartCancelledBeforeDispatch( - cancellation: await admission.cancellationRequest() ?? .system() - ) + throw error } if let startReviewGate { await startReviewGate.wait() @@ -624,11 +620,7 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { model: recoveredRun.model ) await admission.recordPreparedThread(provisionalRun) - guard await admission.admitReviewStartDispatch(for: provisionalRun) else { - throw ReviewStartCancelledBeforeDispatch( - cancellation: await admission.cancellationRequest() ?? .system() - ) - } + try await admission.admitReviewStartDispatch(for: provisionalRun) commands.append(.resumeReviewRecovery(handoff, request)) let waiters = Array(resumeReviewRecoveryWaiters.values) resumeReviewRecoveryWaiters.removeAll(keepingCapacity: false) diff --git a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift index 7462a8b4..736db6b7 100644 --- a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift +++ b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift @@ -70,7 +70,7 @@ private func makeRecoveryCandidate( ) async throws -> ReviewRecoveryCandidate { let admission = ReviewStartAdmission() let registered = try await admission.registerStart { admission in - #expect(await admission.admitThreadStartDispatch()) + try await admission.admitThreadStartDispatch() let provisionalRun = CodexReviewBackendModel.Review.Run( attemptID: run.attemptID, threadID: run.threadID, @@ -78,7 +78,7 @@ private func makeRecoveryCandidate( model: run.model ) await admission.recordPreparedThread(provisionalRun) - #expect(await admission.admitReviewStartDispatch(for: provisionalRun)) + try await admission.admitReviewStartDispatch(for: provisionalRun) await admission.recordActiveRun(run) return .init(run: run) } diff --git a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift index 71558df2..b6bdc976 100644 --- a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift +++ b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift @@ -349,6 +349,45 @@ struct ReviewAttemptProcessorTests { } } + @Test func startupForceCloseFailureReturnsWithoutWaitingForStartResponse() async throws { + let graceGate = AsyncGate() + let startResponseGate = AsyncGate() + let startDispatched = InvocationProbe() + let forceCloseStarted = InvocationProbe() + let forceFailure = ReviewRuntimeCloseFailure.process("Process remained alive") + let admission = ReviewStartAdmission( + closePolicy: controlledClosePolicy(gate: graceGate) + ) + let registered = try await admission.registerStart { admission in + try await admission.admitThreadStartDispatch() + await startDispatched.record() + await startResponseGate.waitIgnoringCancellation() + return .init(run: canonicalRun) + } + try await admission.activateStart(registered.id) + await startDispatched.waitForInvocation() + + let cancellation = Task { + try await admission.cancel( + .system(message: "Stop"), + interrupt: { _, _ in Issue.record("A startup-only attempt was interrupted.") }, + forceClose: { + await forceCloseStarted.record() + throw forceFailure + } + ) + } + #expect(await admission.waitForCancellationAdmission() == .system(message: "Stop")) + await graceGate.open() + await forceCloseStarted.waitForInvocation() + + await #expect(throws: forceFailure) { + try await cancellation.value + } + await startResponseGate.open() + _ = try await registered.task.value + } + @Test func terminalFirstForceCloseFailureCancelsPendingRequestAndPreservesTerminal() async throws { let graceGate = AsyncGate() let (admission, run) = try await makeActiveAdmission( @@ -411,7 +450,7 @@ struct ReviewAttemptProcessorTests { await requestStarted.waitForInvocation() let second = Task { try await admission.cancel( - .mcpClient(message: "Stop"), + .system(message: "Runtime stop"), interrupt: { _, _ in Issue.record("Duplicate caller installed a second interrupt operation.") }, @@ -425,6 +464,9 @@ struct ReviewAttemptProcessorTests { await requestGate.open() #expect(try await first.value == second.value) + #expect(await admission.waitForInterruptionAdmission() == .terminalCancellation( + .mcpClient(message: "Stop") + )) #expect(await requestStarted.invocationCount() == 1) } @@ -602,6 +644,13 @@ struct ReviewAttemptProcessorTests { forceClose: {} ) } + let duplicateCancellation = Task { + try await admission.cancel( + .system(message: "Runtime stop"), + interrupt: { _, _ in Issue.record("Duplicate joined cancellation sent a request.") }, + forceClose: {} + ) + } #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) try await admission.recordCanonicalTerminal( .interrupted(.server(message: "network recovery")), @@ -616,7 +665,8 @@ struct ReviewAttemptProcessorTests { } #expect(product.resolved.run == run) #expect(product.productTerminal == .interrupted(.requested(.mcpClient(message: "Stop")))) - _ = try await cancellation.value + _ = try await (cancellation.value, duplicateCancellation.value) + #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) #expect(await requestStarted.invocationCount() == 1) } @@ -711,8 +761,10 @@ struct ReviewAttemptProcessorTests { return .init(run: canonicalRun) } - #expect(await admission.admitThreadStartDispatch()) - #expect(await admission.admitThreadStartDispatch() == false) + try await admission.admitThreadStartDispatch() + await #expect(throws: ReviewAttemptContractFailure.self) { + try await admission.admitThreadStartDispatch() + } await startGate.open() _ = try await startTask.value @@ -726,10 +778,12 @@ struct ReviewAttemptProcessorTests { return .init(run: canonicalRun) } - #expect(await admission.admitThreadStartDispatch()) + try await admission.admitThreadStartDispatch() try await admission.recordThreadStartRejectedForRetry() - #expect(await admission.admitThreadStartDispatch()) - #expect(await admission.admitThreadStartDispatch() == false) + try await admission.admitThreadStartDispatch() + await #expect(throws: ReviewAttemptContractFailure.self) { + try await admission.admitThreadStartDispatch() + } await startGate.open() _ = try await startTask.value @@ -742,11 +796,13 @@ struct ReviewAttemptProcessorTests { await startGate.waitIgnoringCancellation() return .init(run: canonicalRun) } - #expect(await admission.admitThreadStartDispatch()) + try await admission.admitThreadStartDispatch() await admission.recordPreparedThread(provisionalRun) - #expect(await admission.admitReviewStartDispatch(for: provisionalRun)) - #expect(await admission.admitReviewStartDispatch(for: provisionalRun) == false) + try await admission.admitReviewStartDispatch(for: provisionalRun) + await #expect(throws: ReviewAttemptContractFailure.self) { + try await admission.admitReviewStartDispatch(for: provisionalRun) + } await startGate.open() _ = try await startTask.value @@ -774,11 +830,7 @@ struct ReviewAttemptProcessorTests { await entered.record() await dispatchGate.wait() try Task.checkCancellation() - guard await admission.admitThreadStartDispatch() else { - throw ReviewStartCancelledBeforeDispatch( - cancellation: await admission.cancellationRequest() ?? .system() - ) - } + try await admission.admitThreadStartDispatch() Issue.record("Thread request was dispatched after cancellation.") return .init(run: canonicalRun) } @@ -796,22 +848,44 @@ struct ReviewAttemptProcessorTests { } } + @Test func streamTerminalAfterActivationBeforeDispatchRemainsAuthoritative() async throws { + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) + let operationEntered = InvocationProbe() + let dispatchGate = AsyncGate() + let failure = ReviewAttemptStreamFailure.unexpectedConnection( + .connection("Connection ended before dispatch") + ) + let registered = try await admission.registerStart { admission in + await operationEntered.record() + await dispatchGate.waitIgnoringCancellation() + try await admission.admitThreadStartDispatch() + Issue.record("A terminal attempt admitted a backend write.") + return .init(run: canonicalRun) + } + try await admission.activateStart(registered.id) + await operationEntered.waitForInvocation() + + try await admission.recordStreamTerminal(failure) + await dispatchGate.open() + + await #expect(throws: failure) { + try await registered.task.value + } + #expect(await admission.currentPhase() == .terminal(.stream(failure))) + } + @Test func cancellationAfterThreadDispatchRefusesReviewDispatchAfterResponse() async throws { let graceGate = AsyncGate() let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: graceGate)) let threadDispatched = InvocationProbe() let threadResponseGate = AsyncGate() let startTask = try await registerAndActivateStart(admission) { admission in - #expect(await admission.admitThreadStartDispatch()) + try await admission.admitThreadStartDispatch() await threadDispatched.record() await threadResponseGate.waitIgnoringCancellation() let provisional = provisionalRun await admission.recordPreparedThread(provisional) - guard await admission.admitReviewStartDispatch(for: provisional) else { - throw ReviewStartCancelledBeforeDispatch( - cancellation: await admission.cancellationRequest() ?? .system() - ) - } + try await admission.admitReviewStartDispatch(for: provisional) Issue.record("Review request was dispatched after cancellation.") return .init(run: canonicalRun) } @@ -841,7 +915,7 @@ struct ReviewAttemptProcessorTests { let forceClose = InvocationProbe() let connection = ReviewRuntimeCloseFailure.connection("Forced close") let startTask = try await registerAndActivateStart(admission) { admission in - #expect(await admission.admitThreadStartDispatch()) + try await admission.admitThreadStartDispatch() await threadDispatched.record() await threadResponseGate.wait() try Task.checkCancellation() @@ -879,15 +953,11 @@ struct ReviewAttemptProcessorTests { let prepared = InvocationProbe() let reviewDispatchGate = AsyncGate() let startTask = try await registerAndActivateStart(admission) { admission in - #expect(await admission.admitThreadStartDispatch()) + try await admission.admitThreadStartDispatch() await admission.recordPreparedThread(provisionalRun) await prepared.record() await reviewDispatchGate.waitIgnoringCancellation() - guard await admission.admitReviewStartDispatch(for: provisionalRun) else { - throw ReviewStartCancelledBeforeDispatch( - cancellation: await admission.cancellationRequest() ?? .system() - ) - } + try await admission.admitReviewStartDispatch(for: provisionalRun) Issue.record("Review request was dispatched after cancellation.") return .init(run: canonicalRun) } @@ -916,9 +986,9 @@ struct ReviewAttemptProcessorTests { let reviewResponseGate = AsyncGate() let interruptCalled = InvocationProbe() let startTask = try await registerAndActivateStart(admission) { admission in - #expect(await admission.admitThreadStartDispatch()) + try await admission.admitThreadStartDispatch() await admission.recordPreparedThread(provisionalRun) - #expect(await admission.admitReviewStartDispatch(for: provisionalRun)) + try await admission.admitReviewStartDispatch(for: provisionalRun) await reviewDispatched.record() await reviewResponseGate.waitIgnoringCancellation() await admission.recordActiveRun(canonicalRun) @@ -957,9 +1027,9 @@ struct ReviewAttemptProcessorTests { let forceClose = InvocationProbe() let connection = ReviewRuntimeCloseFailure.connection("Forced close") let startTask = try await registerAndActivateStart(admission) { admission in - #expect(await admission.admitThreadStartDispatch()) + try await admission.admitThreadStartDispatch() await admission.recordPreparedThread(provisionalRun) - #expect(await admission.admitReviewStartDispatch(for: provisionalRun)) + try await admission.admitReviewStartDispatch(for: provisionalRun) await reviewDispatched.record() await reviewResponseGate.wait() try Task.checkCancellation() @@ -1038,9 +1108,9 @@ private func makeActiveAdmission( closePolicy: closePolicy ?? controlledClosePolicy(gate: AsyncGate()) ) let startTask = try await registerAndActivateStart(admission) { admission in - #expect(await admission.admitThreadStartDispatch()) + try await admission.admitThreadStartDispatch() await admission.recordPreparedThread(provisionalRun) - #expect(await admission.admitReviewStartDispatch(for: provisionalRun)) + try await admission.admitReviewStartDispatch(for: provisionalRun) await admission.recordActiveRun(canonicalRun) return .init(run: canonicalRun) } From cce9354a0c0538c7be00b6bcd5505faf03d0adb0 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 01:45:54 +0900 Subject: [PATCH 21/38] fix(review): preserve close barrier failures --- .../Store/CodexReviewStoreReviews.swift | 14 +++- .../AppServerCodexReviewBackend.swift | 26 ++++++- .../AppServerProcessTransport.swift | 4 +- .../LiveCodexReviewStoreBackend.swift | 6 +- .../AppServerClientTests.swift | 8 +++ .../CodexReviewHostTests.swift | 69 +++++++++++++++++++ .../CodexReviewStoreCommandTests.swift | 28 +++++++- 7 files changed, 144 insertions(+), 11 deletions(-) diff --git a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift index 29fc11fb..33c9a85c 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift @@ -701,12 +701,22 @@ extension CodexReviewStore { } } - package func closeActiveReviewSessions(reason: ReviewCancellation) async { + package func closeActiveReviewSessions(reason: ReviewCancellation) async throws { let jobIDs = orderedJobs .filter { $0.isTerminal == false } .map(\.id) + var firstError: (any Error)? for jobID in jobIDs { - _ = try? await cancelReview(jobID: jobID, cancellation: reason) + do { + _ = try await cancelReview(jobID: jobID, cancellation: reason) + } catch { + if firstError == nil { + firstError = error + } + } + } + if let firstError { + throw firstError } } diff --git a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift index ae8c2dae..8c6ba9cb 100644 --- a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift +++ b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift @@ -374,11 +374,31 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { ) do { try await client.close() - } catch let failure as ReviewRuntimeCloseFailure { - throw failure } catch { - throw ReviewRuntimeCloseFailure.connection(error.localizedDescription) + throw Self.reviewRuntimeCloseFailure(for: error) + } + } + + package static func reviewRuntimeCloseFailure( + for error: any Error + ) -> ReviewRuntimeCloseFailure { + if let failure = error as? ReviewRuntimeCloseFailure { + return failure + } + if let processError = error as? AppServerProcessTransportError, + case .processDidNotTerminate = processError { + return .process(processError.localizedDescription) + } + if let jsonRPCError = error as? JSONRPC.Error, + case .transportTerminated(let termination) = jsonRPCError { + switch termination { + case .processExit, .processFailure: + return .process(jsonRPCError.localizedDescription) + case .ownerClose: + break + } } + return .connection(error.localizedDescription) } package func prepareReviewRecovery( diff --git a/Sources/CodexReviewAppServer/AppServerProcessTransport.swift b/Sources/CodexReviewAppServer/AppServerProcessTransport.swift index 66173a50..b0ced6b4 100644 --- a/Sources/CodexReviewAppServer/AppServerProcessTransport.swift +++ b/Sources/CodexReviewAppServer/AppServerProcessTransport.swift @@ -953,11 +953,11 @@ private final class AppServerSpawnedProcess: @unchecked Sendable { } } -private enum AppServerProcessTransportError: LocalizedError { +package enum AppServerProcessTransportError: LocalizedError, Equatable, Sendable { case executableNotFound(command: String, path: String?) case processDidNotTerminate(pid_t) - var errorDescription: String? { + package var errorDescription: String? { switch self { case .executableNotFound(let command, let path): let resolvedPath = path?.trimmingCharacters(in: .whitespacesAndNewlines) diff --git a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift index fd165f31..f5f05e19 100644 --- a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift +++ b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift @@ -599,7 +599,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } let runtimeStore = appServerBackend == nil ? nil : attachedStore if let runtimeStore { - await runtimeStore.closeActiveReviewSessions( + try await runtimeStore.closeActiveReviewSessions( reason: .system(message: "Account switched.") ) } @@ -626,7 +626,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { || auth.persistedActiveAccountKey == accountKey let runtimeStore = removedActiveAccount && appServerBackend != nil ? attachedStore : nil if let runtimeStore { - await runtimeStore.closeActiveReviewSessions( + try await runtimeStore.closeActiveReviewSessions( reason: .system(message: "Account removed.") ) } @@ -695,7 +695,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } let shouldRecycleRuntime = attachedStore != nil && appServerBackend != nil if shouldRecycleRuntime { - await attachedStore?.closeActiveReviewSessions(reason: .system(message: "Signed out.")) + try await attachedStore?.closeActiveReviewSessions(reason: .system(message: "Signed out.")) } if let appServerBackend { _ = try await appServerBackend.logout(.init(account.accountKey)) diff --git a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift index 736db6b7..c9aaf2cc 100644 --- a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift +++ b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift @@ -544,6 +544,14 @@ struct AppServerClientTests { } } + @Test func backendClassifiesProcessCloseFailureAsProcessFailure() { + let error = AppServerProcessTransportError.processDidNotTerminate(42) + + #expect(AppServerCodexReviewBackend.reviewRuntimeCloseFailure(for: error) == .process( + error.localizedDescription + )) + } + @Test func spontaneousProcessExitReplaysTypedCauseToLateSubscriber() async throws { let directory = FileManager.default.temporaryDirectory .appending(path: "codex-review-process-exit-\(UUID().uuidString)") diff --git a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift index 706f0058..498f9429 100644 --- a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift +++ b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift @@ -1070,6 +1070,75 @@ struct CodexReviewHostTests { #expect(await secondTransport.recordedRequests().map(\.method).contains("account/read")) } + @Test func liveStoreSwitchAccountPreservesCredentialsWhenReviewBarrierFails() async throws { + let homeURL = try temporaryHome() + try writeRegistry( + homeURL: homeURL, + activeAccountKey: "first@example.com", + accounts: ["first@example.com", "second@example.com"] + ) + try writeSavedAccountAuth(homeURL: homeURL, accountKey: "second@example.com") + + let transport = FakeJSONRPCTransport() + try await transport.enqueue(AppServerAPI.Initialize.Response(), for: "initialize") + try await transport.enqueue( + AppServerAPI.Account.Read.Response(account: .init(email: "first@example.com", planType: "pro")), + 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") + try await transport.enqueue( + AppServerAPI.Account.RateLimits.Response(rateLimits: .init( + limitID: "codex", + primary: .init(usedPercent: 10, windowDurationMins: 300) + )), + for: "account/rateLimits/read" + ) + try await transport.enqueue( + AppServerAPI.Thread.Start.Response(threadID: "thread-first", model: "gpt-5"), + for: "thread/start" + ) + try await transport.enqueue( + AppServerAPI.Review.Start.Response(turnID: "turn-first"), + for: "review/start" + ) + await transport.enqueueFailure( + .responseError(code: -32_000, message: "Interrupt rejected"), + for: "turn/interrupt" + ) + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + transport: transport + ) + + await store.start(forceRestartIfNeeded: true) + async let reviewRead = store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + await waitUntil { store.jobs.first?.core.run.turnID == "turn-first" } + + await #expect(throws: ReviewInterruptRequestFailure.self) { + try await store.switchAccount(CodexAccount(email: "second@example.com")) + } + #expect(store.auth.selectedAccount?.accountKey == "first@example.com") + #expect(try activeAccountKey(homeURL: homeURL) == "first@example.com") + try await transport.emitServerNotification( + method: "turn/completed", + params: HostTurnNotification( + threadID: "thread-first", + turnID: "turn-first", + status: "interrupted", + errorMessage: "Review ended." + ) + ) + _ = try await reviewRead + } + @Test func liveStoreSignOutRestartsRuntimeAndCancelsRunningReviews() async throws { let homeURL = try temporaryHome() let mainCodexHomeURL = homeURL.appendingPathComponent(".codex_review", isDirectory: true) diff --git a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift index e9639cd1..3da6e7f2 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift @@ -2791,7 +2791,7 @@ struct CodexReviewStoreCommandTests { ) try await seedQueuedAttemptOwnership(in: store, for: running) - await store.closeActiveReviewSessions(reason: .system(message: "Account switched.")) + try await store.closeActiveReviewSessions(reason: .system(message: "Account switched.")) #expect(running.core.lifecycle.status == .cancelled) async let result = store.startReview( @@ -2806,6 +2806,32 @@ struct CodexReviewStoreCommandTests { } } + @Test func closeActiveReviewSessionsPropagatesBarrierFailure() async throws { + let backend = FakeCodexReviewBackend() + await backend.failInterrupts(message: "Interrupt rejected") + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-1" }) + ) + try await withStoreCommandTestCleanup(backend: backend, store: store) { + async let result = store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + try #require(await StoreSnapshotProbe(store: store).waitUntilJobStatus(.running, jobID: "job-1") != nil) + + await #expect(throws: ReviewInterruptRequestFailure.self) { + try await store.closeActiveReviewSessions( + reason: .system(message: "Account switched.") + ) + } + #expect(try store.readReview(jobID: "job-1").cancellable) + + await backend.yield(.completed(summary: "Succeeded.", result: "review text")) + _ = try await result + } + } + @Test func authAndSettingsUseSingleBackendContract() async throws { let backend = FakeCodexReviewBackend(settings: .init(model: "gpt-5")) let store = CodexReviewStore.makeTestingStore( From d90be363162692ca34b55cd9c2da77109ec852cf Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 01:52:55 +0900 Subject: [PATCH 22/38] fix(review): retry failed startup cancellation --- .../CodexReview/ReviewAttemptProcessor.swift | 7 ++++++ .../ReviewAttemptProcessorTests.swift | 22 +++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/Sources/CodexReview/ReviewAttemptProcessor.swift b/Sources/CodexReview/ReviewAttemptProcessor.swift index 63671798..4d566455 100644 --- a/Sources/CodexReview/ReviewAttemptProcessor.swift +++ b/Sources/CodexReview/ReviewAttemptProcessor.swift @@ -1000,6 +1000,9 @@ package actor ReviewStartAdmission { if case .failure(let error)? = cancellationResult { await drainCancellationTasks() + if terminal == nil { + resetCancellationOperationForRetry() + } throw error } @@ -1075,6 +1078,10 @@ package actor ReviewStartAdmission { } private func resetRejectedCancellationForRetry() { + resetCancellationOperationForRetry() + } + + private func resetCancellationOperationForRetry() { requestedCancellation = nil joinedTerminalCancellation = nil interruptionPurpose = nil diff --git a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift index b6bdc976..f7ba968d 100644 --- a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift +++ b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift @@ -386,6 +386,28 @@ struct ReviewAttemptProcessorTests { } await startResponseGate.open() _ = try await registered.task.value + #expect(await admission.currentPhase() == .active(canonicalRun)) + + let retryStarted = InvocationProbe() + let retry = Task { + try await admission.cancel( + .mcpClient(message: "Retry stop"), + interrupt: { run, _ in + #expect(run == canonicalRun) + await retryStarted.record() + }, + forceClose: {} + ) + } + await retryStarted.waitForInvocation() + try await admission.recordCanonicalTerminal( + .interrupted(.requested(.mcpClient(message: "Retry stop"))), + for: canonicalRun + ) + + #expect(try await retry.value.terminal == .canonical( + .interrupted(.requested(.mcpClient(message: "Retry stop"))) + )) } @Test func terminalFirstForceCloseFailureCancelsPendingRequestAndPreservesTerminal() async throws { From 4b5ef19d872e0a4c7ed774c0db15ef20076c81c3 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 02:22:29 +0900 Subject: [PATCH 23/38] fix: drain recovery cancellation barriers --- .../CodexReview/ReviewAttemptProcessor.swift | 54 ++++++++++++++---- .../Store/CodexReviewStoreReviews.swift | 25 ++++++--- .../AppServerCodexReviewBackend.swift | 14 ++++- .../AppServerClientTests.swift | 55 +++++++++++++++++++ .../CodexReviewStoreCommandTests.swift | 47 ++++++++++++++-- .../ReviewAttemptProcessorTests.swift | 48 ++++++++++++++++ 6 files changed, 213 insertions(+), 30 deletions(-) diff --git a/Sources/CodexReview/ReviewAttemptProcessor.swift b/Sources/CodexReview/ReviewAttemptProcessor.swift index 4d566455..26bdeb0e 100644 --- a/Sources/CodexReview/ReviewAttemptProcessor.swift +++ b/Sources/CodexReview/ReviewAttemptProcessor.swift @@ -317,7 +317,9 @@ package actor ReviewStartAdmission { case queued case registeredStart(ReviewStartHandleID) case activatedStart(ReviewStartHandleID) + case preparingInitialRequest case preparingThread(RequestDispatch) + case recoveryRollbackOutcomeUnknown(threadID: String) case startingReview( preparedRun: CodexReviewBackendModel.Review.Run, dispatch: RequestDispatch @@ -441,7 +443,12 @@ package actor ReviewStartAdmission { if let requestedCancellation { throw ReviewStartCancelledBeforeDispatch(cancellation: requestedCancellation) } - guard case .preparingThread(.notSent) = phase else { + switch phase { + case .preparingInitialRequest, .preparingThread(.notSent): + break + case .queued, .registeredStart, .activatedStart, + .preparingThread(.outcomeUnknown), .recoveryRollbackOutcomeUnknown, + .startingReview, .active, .interrupting, .finishing, .terminal: throw ReviewAttemptContractFailure( message: "Thread start dispatch requires one pending not-sent request." ) @@ -449,6 +456,21 @@ package actor ReviewStartAdmission { phase = .preparingThread(.outcomeUnknown) } + package func admitRecoveryRollbackDispatch(threadID: String) throws { + if let terminal { + throw startFailure(for: terminal) + } + if let requestedCancellation { + throw ReviewStartCancelledBeforeDispatch(cancellation: requestedCancellation) + } + guard case .preparingInitialRequest = phase else { + throw ReviewAttemptContractFailure( + message: "Recovery rollback dispatch requires one pending recovery request." + ) + } + phase = .recoveryRollbackOutcomeUnknown(threadID: threadID) + } + package func recordThreadStartRejectedForRetry() throws { if let requestedCancellation { throw ReviewStartCancelledBeforeDispatch(cancellation: requestedCancellation) @@ -915,7 +937,7 @@ package actor ReviewStartAdmission { message: "Start handle \(id.generation) was not activated for dispatch." ) } - phase = .preparingThread(.notSent) + phase = .preparingInitialRequest } private func resolveStartActivation(_ result: Result) { @@ -960,10 +982,12 @@ package actor ReviewStartAdmission { } else if error is CancellationError, let requestedCancellation { switch phase { - case .registeredStart, .activatedStart, .preparingThread(.notSent): + case .registeredStart, .activatedStart, .preparingInitialRequest, + .preparingThread(.notSent): receiveTerminal(.localCancellation(requestedCancellation)) - case .queued, .preparingThread(.outcomeUnknown), .startingReview, - .active, .interrupting, .finishing, .terminal: + case .queued, .preparingThread(.outcomeUnknown), + .recoveryRollbackOutcomeUnknown, .startingReview, .active, + .interrupting, .finishing, .terminal: break } } @@ -982,7 +1006,9 @@ package actor ReviewStartAdmission { ) async throws -> Void, forceClose: @escaping @Sendable () async throws -> Void ) async throws -> ReviewAttemptCancellationResolution { - if case .preparingThread(.notSent) = phase { + if case .preparingInitialRequest = phase { + startTask?.cancel() + } else if case .preparingThread(.notSent) = phase { startTask?.cancel() } else if case .queued = phase { receiveTerminal(.localCancellation(cancellation)) @@ -1419,7 +1445,8 @@ package actor ReviewStartAdmission { run case .finishing, .terminal: registeredRun - case .queued, .registeredStart, .activatedStart, .preparingThread: + case .queued, .registeredStart, .activatedStart, .preparingInitialRequest, + .preparingThread, .recoveryRollbackOutcomeUnknown: registeredRun } } @@ -1428,19 +1455,22 @@ package actor ReviewStartAdmission { switch phase { case .active(let run), .interrupting(let run): run - case .queued, .registeredStart, .activatedStart, .preparingThread, - .startingReview, .finishing, .terminal: + case .queued, .registeredStart, .activatedStart, .preparingInitialRequest, + .preparingThread, .recoveryRollbackOutcomeUnknown, .startingReview, + .finishing, .terminal: nil } } private static func isOutcomeUnknownStartPhase(_ phase: Phase) -> Bool { switch phase { - case .preparingThread(.outcomeUnknown), .startingReview(_, .outcomeUnknown): + case .preparingThread(.outcomeUnknown), .recoveryRollbackOutcomeUnknown, + .startingReview(_, .outcomeUnknown): true case .queued, .registeredStart, .activatedStart, - .preparingThread(.notSent), .startingReview(_, .notSent), - .active, .interrupting, .finishing, .terminal: + .preparingInitialRequest, .preparingThread(.notSent), + .startingReview(_, .notSent), .active, .interrupting, .finishing, + .terminal: false } } diff --git a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift index 33c9a85c..13b75fed 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift @@ -526,6 +526,9 @@ extension CodexReviewStore { cancellation: cancellation ) await reviewWorkerTasks[jobID]?.value + if let cleanupFailure = reviewCleanupFailures[jobID] { + throw cleanupFailure + } case .replacementStart(_, let start): let resolution = try await cancel( admission: start.admission, @@ -846,7 +849,8 @@ extension CodexReviewStore { let completion = try await consumeReviewEventLoop( job: job, startRequest: startRequest, - inputs: inputs + inputs: inputs, + initialCleanupAttempt: initialActive ) await inputs.cancel() return completion @@ -864,13 +868,15 @@ extension CodexReviewStore { private func consumeReviewEventLoop( job: CodexReviewJob, startRequest: CodexReviewBackendModel.Review.Start, - inputs: ReviewWorkerInputs + inputs: ReviewWorkerInputs, + initialCleanupAttempt: ReviewActiveAttempt ) async throws -> ReviewWorkerAttemptCompletion { var recoverySignals = ReviewNetworkRecoverySignals() var activeEventSubscriptionID: Int? = inputs.initialEventSubscriptionID + var cleanupAttempt = initialCleanupAttempt while let input = await inputs.next() { if job.isTerminal { - return .init(cleanupAttempt: activeAttemptForCleanup(jobID: job.id)) + return .init(cleanupAttempt: cleanupAttempt) } switch input { case .reviewEvent(let event): @@ -983,7 +989,7 @@ extension CodexReviewStore { currentCandidate == candidate else { if job.isTerminal { - return .init(cleanupAttempt: nil) + return .init(cleanupAttempt: cleanupAttempt) } throw ReviewAttemptContractFailure( message: "Recovery preparation completed after its ownership changed." @@ -998,7 +1004,7 @@ extension CodexReviewStore { currentHandoff == handoff else { if job.isTerminal { - return .init(cleanupAttempt: nil) + return .init(cleanupAttempt: cleanupAttempt) } throw ReviewAttemptContractFailure( message: "Recovery handoff changed while detaching the old subscription." @@ -1044,8 +1050,9 @@ extension CodexReviewStore { case .continueWaiting: continue case .finished: - return .init(cleanupAttempt: nil) + return .init(cleanupAttempt: cleanupAttempt) case .recovered(let recoveredAttempt, let active): + cleanupAttempt = active let subscriptionID = await inputs.subscribe( to: recoveredAttempt, owner: active @@ -1055,7 +1062,7 @@ extension CodexReviewStore { else { await inputs.cancelActiveEventSubscription() if job.isTerminal { - return .init(cleanupAttempt: nil) + return .init(cleanupAttempt: cleanupAttempt) } throw ReviewAttemptContractFailure( message: "Recovered subscription completed after its active attempt changed." @@ -1119,7 +1126,7 @@ extension CodexReviewStore { return .init(cleanupAttempt: active) } if job.isTerminal { - return .init(cleanupAttempt: activeAttemptForCleanup(jobID: job.id)) + return .init(cleanupAttempt: cleanupAttempt) } throw ReviewAttemptContractFailure( message: "Review input queue finished without terminal attempt ownership." @@ -1683,7 +1690,7 @@ private struct ReviewWorkerRecoveryBarrierResolution: Sendable { } private struct ReviewWorkerAttemptCompletion: Sendable { - var cleanupAttempt: ReviewActiveAttempt? + var cleanupAttempt: ReviewActiveAttempt } private struct ReviewWorkerRecoveryFailure: LocalizedError, @unchecked Sendable { diff --git a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift index 8c6ba9cb..b149f859 100644 --- a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift +++ b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift @@ -436,9 +436,17 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { await ensureNotificationRouterStarted() let token = handoff.token let interruptedRun = token.interruptedRun - let _: EmptyResponse = try await client.send(AppServerAPI.Thread.Rollback.Request( - params: .init(threadID: token.rollbackThreadID, numTurns: 1) - )) + try await admission.admitRecoveryRollbackDispatch(threadID: token.rollbackThreadID) + do { + let _: EmptyResponse = try await client.send(AppServerAPI.Thread.Rollback.Request( + params: .init(threadID: token.rollbackThreadID, numTurns: 1) + )) + } catch { + if let terminal = streamTerminal(for: error) { + try await admission.recordStreamTerminal(terminal) + } + throw error + } let control = controlsByThreadID[interruptedRun.threadID] ?? AppServerReviewControl(client: client) controlsByThreadID[interruptedRun.threadID] = control diff --git a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift index c9aaf2cc..0b71e5e2 100644 --- a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift +++ b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift @@ -2272,6 +2272,61 @@ struct AppServerClientTests { #expect(await events.mailbox.isFinished() == false) } + @Test func backendAdmitsRecoveryRollbackBeforeDispatch() async throws { + let transport = FakeJSONRPCTransport() + try await enqueueInitialize(transport) + try await transport.enqueue(EmptyResponse(), for: "turn/interrupt") + let rollbackGate = AsyncGate() + await transport.holdNextIgnoringCancellation(method: "thread/rollback", gate: rollbackGate) + try await transport.enqueue(EmptyResponse(), for: "thread/rollback") + try await transport.enqueue( + AppServerAPI.Review.Start.Response( + turnID: "turn-2", + reviewThreadID: "review-thread-1" + ), + for: "review/start" + ) + let backend = AppServerCodexReviewBackend(client: .init(transport: transport)) + let interruptedRun = CodexReviewBackendModel.Review.Run( + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "review-thread-1", + model: "gpt-5" + ) + let reason = CodexReviewBackendModel.CancellationReason( + message: "Network unavailable; waiting to reconnect." + ) + let candidate = try await makeRecoveryCandidate( + backend: backend, + for: interruptedRun, + reason: reason + ) + let handoff = try await backend.prepareReviewRecovery(candidate) + let admission = ReviewStartAdmission() + let registered = try await admission.registerStart { admission in + try await backend.resumeReviewRecovery( + handoff, + request: .init( + jobID: "job-1", + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .baseBranch("main")), + model: "gpt-5" + ), + admission: admission + ) + } + try await admission.activateStart(registered.id) + await transport.waitForRequest(method: "thread/rollback") + + #expect(await admission.currentPhase() == .recoveryRollbackOutcomeUnknown( + threadID: "review-thread-1" + )) + + await rollbackGate.open() + let recovered = try await registered.task.value + #expect(recovered.turnID == "turn-2") + } + @Test func backendRecoverReviewRollsBackAndRestartsSameThread() async throws { let transport = FakeJSONRPCTransport() try await enqueueInitialize(transport) diff --git a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift index 3da6e7f2..9142a988 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift @@ -1161,8 +1161,9 @@ struct CodexReviewStoreCommandTests { } } - @Test func cancellationDuringRecoveryPreparationDiscardsHandoffBeforeResume() async throws { + @Test func cancellationDuringRecoveryPreparationPropagatesCleanupFailure() async throws { let backend = FakeCodexReviewBackend() + await backend.failCleanup(message: "unsubscribe failed") let preparationGate = AsyncGate() await backend.holdPrepareReviewRecovery(with: preparationGate) let networkMonitor = ManualCodexReviewNetworkMonitor() @@ -1201,21 +1202,55 @@ struct CodexReviewStoreCommandTests { } }) - let cancel = try await store.cancelReview( - jobID: "job-1", - cancellation: .mcpClient(message: "Stop") - ) + await #expect(throws: ReviewRuntimeCloseFailure.cleanup("unsubscribe failed")) { + try await store.cancelReview( + jobID: "job-1", + cancellation: .mcpClient(message: "Stop") + ) + } let read = try await result let commands = await backend.recordedCommands() - #expect(cancel.cancelled) #expect(read.core.lifecycle.status == .cancelled) + #expect(store.reviewCleanupFailures["job-1"] == .cleanup("unsubscribe failed")) #expect(commands.contains { if case .resumeReviewRecovery = $0 { true } else { false } } == false) } } + @Test func cancellationWhileWaitingForRecoveryPropagatesCleanupFailure() async throws { + let backend = FakeCodexReviewBackend() + await backend.failCleanup(message: "unsubscribe failed") + let networkMonitor = ManualCodexReviewNetworkMonitor() + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-1" }), + networkMonitor: networkMonitor, + networkRecoveryPolicy: .init(sleep: { _ in }) + ) + try await withStoreCommandTestCleanup(backend: backend, store: store) { + async let result = store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .baseBranch("main")) + ) + + networkMonitor.yield(.init(status: .unsatisfied)) + try await completeNetworkRecoveryBarrier(backend: backend, store: store) + + await #expect(throws: ReviewRuntimeCloseFailure.cleanup("unsubscribe failed")) { + try await store.cancelReview( + jobID: "job-1", + cancellation: .mcpClient(message: "Stop") + ) + } + let read = try await result + + #expect(read.core.lifecycle.status == .cancelled) + #expect(store.reviewCleanupFailures["job-1"] == .cleanup("unsubscribe failed")) + } + } + @Test func networkRecoveryWaitDiscardsOldAttemptCompletion() async throws { let initialRun = CodexReviewBackendModel.Review.Run( threadID: "thread-1", diff --git a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift index f7ba968d..f31035ce 100644 --- a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift +++ b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift @@ -970,6 +970,54 @@ struct ReviewAttemptProcessorTests { } } + @Test func cancellationAfterOutcomeUnknownRecoveryRollbackDrainsThroughForcedConnectionTerminal() async throws { + let graceGate = AsyncGate() + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: graceGate)) + let rollbackDispatched = InvocationProbe() + let rollbackResponseGate = AsyncGate() + let forceClose = InvocationProbe() + let connection = ReviewRuntimeCloseFailure.connection("Forced close") + let startTask = try await registerAndActivateStart(admission) { admission in + try await admission.admitRecoveryRollbackDispatch(threadID: "review-thread-1") + await rollbackDispatched.record() + await rollbackResponseGate.wait() + try Task.checkCancellation() + Issue.record("Recovery rollback outlived its typed connection terminal.") + return .init(run: canonicalRun) + } + await rollbackDispatched.waitForInvocation() + + #expect(await admission.currentPhase() == .recoveryRollbackOutcomeUnknown( + threadID: "review-thread-1" + )) + await #expect(throws: ReviewAttemptContractFailure.self) { + try await admission.admitThreadStartDispatch() + } + let cancellation = Task { + try await admission.cancel( + .system(message: "Stop"), + interrupt: { _, _ in Issue.record("Rollback-only attempt interrupted a turn.") }, + forceClose: { + await forceClose.record() + try await admission.recordStreamTerminal( + .ownerForcedConnectionClose(connection) + ) + } + ) + } + #expect(await admission.waitForCancellationAdmission() == .system(message: "Stop")) + await graceGate.open() + await forceClose.waitForInvocation() + + #expect(try await cancellation.value.terminal == .stream( + .ownerForcedConnectionClose(connection) + )) + await #expect(throws: CancellationError.self) { + try await startTask.value + } + #expect(await forceClose.invocationCount() == 1) + } + @Test func cancellationAfterThreadResponseRefusesNotSentReviewDispatch() async throws { let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) let prepared = InvocationProbe() From d72958bc61b85b743147d916efee036811b96d82 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 04:26:03 +0900 Subject: [PATCH 24/38] refactor(runtime): own prepared runtime generations --- .../CodexReview/ReviewRuntimeLifecycle.swift | 323 ++++++++++++ .../CodexReview/Store/CodexReviewStore.swift | 283 ++++++++++- .../Store/CodexReviewStoreBackend.swift | 6 +- .../Store/CodexReviewStoreReviews.swift | 8 + .../PreviewCodexReviewStoreBackend.swift | 9 +- Sources/CodexReviewHost/CodexReviewHost.swift | 55 ++- .../LiveCodexReviewStoreBackend.swift | 460 ++++++++++++++---- .../CodexReviewMCPHTTPServer.swift | 62 ++- Sources/CodexReviewTesting/TestSupport.swift | 103 +++- .../CodexReviewHostTests.swift | 51 ++ .../CodexReviewStoreLifecycleTests.swift | 68 +++ Tests/ReviewUITests/ReviewUITests.swift | 32 +- 12 files changed, 1309 insertions(+), 151 deletions(-) create mode 100644 Sources/CodexReview/ReviewRuntimeLifecycle.swift create mode 100644 Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift diff --git a/Sources/CodexReview/ReviewRuntimeLifecycle.swift b/Sources/CodexReview/ReviewRuntimeLifecycle.swift new file mode 100644 index 00000000..799e7591 --- /dev/null +++ b/Sources/CodexReview/ReviewRuntimeLifecycle.swift @@ -0,0 +1,323 @@ +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 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 + } +} + +package struct MCPServerGeneration: Hashable, Sendable { + package let rawValue: UInt64 + + package init(rawValue: UInt64) { + self.rawValue = rawValue + } +} + +package struct MCPServerPublicationSnapshot: Sendable { + package let serverURL: URL? + + package init(serverURL: URL?) { + self.serverURL = serverURL + } +} + +package struct PreparedMCPServer: Sendable { + package let generation: MCPServerGeneration + + package init(generation: MCPServerGeneration) { + self.generation = generation + } +} + +package enum ReviewRuntimeTransitionPurpose: Equatable, Sendable { + case stop + case restartSameAccount + case accountTransition + case applicationClose + case recoveryReplacement +} + +@MainActor +package protocol RuntimeLifecycleHandle: 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 enum ReviewStoreRuntimeState { + case stopped(ReviewRuntimeGeneration) + case acquiring( + generation: ReviewRuntimeGeneration, + task: Task + ) + case running( + generation: ReviewRuntimeGeneration, + runtime: PreparedRuntime, + mcpGeneration: MCPServerGeneration + ) + case transitioning( + generation: ReviewRuntimeGeneration, + task: Task + ) + + package var generation: ReviewRuntimeGeneration { + switch self { + case .stopped(let generation), + .acquiring(let generation, _), + .running(let generation, _, _), + .transitioning(let generation, _): + generation + } + } +} + +@MainActor +package protocol MCPServerLifecycleOwner: Sendable { + func prepare() async throws -> PreparedMCPServer + func activate( + _ generation: MCPServerGeneration + ) async throws -> MCPServerPublicationSnapshot + func closeAdmission() async + func drainAdmittedHandlers() async throws + func stop() async throws + func waitUntilStopped() async throws + func close() async throws + func waitUntilClosed() async throws +} + +package enum ReviewLifecycleResourceFailure: LocalizedError, Equatable, Sendable { + case client(String) + case process(String) + case authenticationObservation(String) + case reader(String) + case router(String) + case session(String) + case rateLimit(String) + case mcpHandlerDrain(String) + case mcpServer(String) + + package var errorDescription: String? { + switch self { + case .client(let message): + "App-server client close failed: \(message)" + case .process(let message): + "App-server process close failed: \(message)" + case .authenticationObservation(let message): + "Authentication observation close failed: \(message)" + case .reader(let message): + "App-server reader close failed: \(message)" + case .router(let message): + "App-server router close failed: \(message)" + case .session(let message): + "App-server session close failed: \(message)" + case .rateLimit(let message): + "Rate-limit refresh close failed: \(message)" + case .mcpHandlerDrain(let message): + "MCP handler drain failed: \(message)" + case .mcpServer(let message): + "MCP server close failed: \(message)" + } + } +} + +package struct ReviewLifecycleResourceFailureAggregate: LocalizedError, Sendable { + package let first: ReviewLifecycleResourceFailure + package let additionalInLifecycleOrder: [ReviewLifecycleResourceFailure] + + package init( + first: ReviewLifecycleResourceFailure, + additionalInLifecycleOrder: [ReviewLifecycleResourceFailure] = [] + ) { + self.first = first + self.additionalInLifecycleOrder = additionalInLifecycleOrder + } + + package var errorDescription: String? { + ([first] + additionalInLifecycleOrder) + .map(\.localizedDescription) + .joined(separator: "; ") + } +} + +package enum ReviewPersistenceError: LocalizedError, Sendable { + case open(String) + case migration(String) + case read(String) + case write(String) + case close(String) + + package var errorDescription: String? { + switch self { + case .open(let message): "Review history open failed: \(message)" + case .migration(let message): "Review history migration failed: \(message)" + case .read(let message): "Review history read failed: \(message)" + case .write(let message): "Review history write failed: \(message)" + case .close(let message): "Review history close failed: \(message)" + } + } +} + +package enum ReviewClosePrimaryFailure: LocalizedError, Sendable { + case interruptRequest(ReviewInterruptRequestFailure) + case attemptRuntime(ReviewRuntimeCloseFailure) + case lifecycleResources(ReviewLifecycleResourceFailureAggregate) + case persistence(ReviewPersistenceError) + + package var errorDescription: String? { + switch self { + case .interruptRequest(let failure): failure.localizedDescription + case .attemptRuntime(let failure): failure.localizedDescription + case .lifecycleResources(let failure): failure.localizedDescription + case .persistence(let failure): failure.localizedDescription + } + } +} + +package struct ReviewCloseFailureAggregate: LocalizedError, Sendable { + package let first: ReviewClosePrimaryFailure + package let additionalInLifecycleOrder: [ReviewClosePrimaryFailure] + + package init( + first: ReviewClosePrimaryFailure, + additionalInLifecycleOrder: [ReviewClosePrimaryFailure] = [] + ) { + self.first = first + self.additionalInLifecycleOrder = additionalInLifecycleOrder + } + + package var errorDescription: String? { + ([first] + additionalInLifecycleOrder) + .map(\.localizedDescription) + .joined(separator: "; ") + } +} + +package struct ReviewCloseError: LocalizedError, Sendable { + package let failures: ReviewCloseFailureAggregate + package let secondaryPhysicalDatabaseClose: ReviewPersistenceError? + + package init( + failures: ReviewCloseFailureAggregate, + secondaryPhysicalDatabaseClose: ReviewPersistenceError? = nil + ) { + self.failures = failures + self.secondaryPhysicalDatabaseClose = secondaryPhysicalDatabaseClose + } + + package var errorDescription: String? { + guard let secondaryPhysicalDatabaseClose else { + return failures.localizedDescription + } + return "\(failures.localizedDescription); \(secondaryPhysicalDatabaseClose.localizedDescription)" + } +} + +@MainActor +package final class NoMCPServerLifecycleOwner: MCPServerLifecycleOwner { + private enum State { + case stopped + case prepared(MCPServerGeneration) + case running(MCPServerGeneration) + case closed + } + + private var state: State = .stopped + private var nextGeneration: UInt64 = 0 + + package init() {} + + package func prepare() async throws -> PreparedMCPServer { + guard case .stopped = state else { + throw ReviewLifecycleResourceFailure.mcpServer( + "No-MCP owner preparation requires stopped state." + ) + } + nextGeneration &+= 1 + let generation = MCPServerGeneration(rawValue: nextGeneration) + state = .prepared(generation) + return .init(generation: generation) + } + + package func activate( + _ generation: MCPServerGeneration + ) async throws -> MCPServerPublicationSnapshot { + guard case .prepared(generation) = state else { + throw ReviewLifecycleResourceFailure.mcpServer( + "No-MCP activation requires its exact prepared generation." + ) + } + state = .running(generation) + return .init(serverURL: nil) + } + + package func closeAdmission() async {} + + package func drainAdmittedHandlers() async throws {} + + package func stop() async throws { + switch state { + case .stopped: + return + case .prepared, .running: + state = .stopped + case .closed: + throw ReviewLifecycleResourceFailure.mcpServer( + "No-MCP owner is closed." + ) + } + } + + package func waitUntilStopped() async throws { + guard case .stopped = state else { + throw ReviewLifecycleResourceFailure.mcpServer( + "No-MCP owner did not stop." + ) + } + } + + package func close() async throws { + state = .closed + } + + package func waitUntilClosed() async throws { + guard case .closed = state else { + throw ReviewLifecycleResourceFailure.mcpServer( + "No-MCP owner did not close." + ) + } + } +} diff --git a/Sources/CodexReview/Store/CodexReviewStore.swift b/Sources/CodexReview/Store/CodexReviewStore.swift index 6a3713f4..97e11274 100644 --- a/Sources/CodexReview/Store/CodexReviewStore.swift +++ b/Sources/CodexReview/Store/CodexReviewStore.swift @@ -36,6 +36,9 @@ public final class CodexReviewStore { @ObservationIgnored package var reviewTerminalWaiters: [String: [ReviewTerminalWaiter]] = [:] @ObservationIgnored package var closedSessions: Set = [] @ObservationIgnored package var accountRateLimitAutoRefreshDriver: CodexReviewStoreRateLimitAutoRefreshDriver? + @ObservationIgnored package var runtimeState: ReviewStoreRuntimeState = .stopped( + .init(rawValue: 0) + ) package init( backend: any CodexReviewStoreBackend = PreviewCodexReviewStoreBackend(), @@ -79,6 +82,12 @@ public final class CodexReviewStore { isolated deinit { accountRateLimitAutoRefreshDriver?.cancel() + switch runtimeState { + case .acquiring(_, let task), .transitioning(_, let task): + task.cancel() + case .stopped, .running: + break + } for task in reviewWorkerTasks.values { task.cancel() } @@ -129,25 +138,91 @@ public final class CodexReviewStore { } public func start(forceRestartIfNeeded: Bool = false) async { - switch serverState { - case .stopped, .failed: - break - case .starting: + switch runtimeState { + case .acquiring, .transitioning: return case .running where forceRestartIfNeeded == false: return case .running: + await stop() + case .stopped: break } + let purpose: ReviewRuntimeTransitionPurpose = forceRestartIfNeeded + ? .restartSameAccount + : .stop + await startRuntime(purpose: purpose) + } + + private func startRuntime(purpose: ReviewRuntimeTransitionPurpose) async { + guard case .stopped(let previousGeneration) = runtimeState else { + return + } + let generation = previousGeneration.successor() serverState = .starting serverURL = nil writeDiagnosticsIfNeeded() - await backend.start(store: self, forceRestartIfNeeded: forceRestartIfNeeded) - await settingsService.refreshIfRunning(serverState: serverState) - startAccountRateLimitAutoRefresh() + let task = Task { @MainActor [weak self] in + guard let self else { return } + await self.performRuntimeAcquisition( + generation: generation, + purpose: purpose + ) + } + runtimeState = .acquiring(generation: generation, task: task) + await task.value } public func stop() async { + let previousState = runtimeState + switch previousState { + case .stopped: + transitionToStopped() + return + case .transitioning(_, let task): + await task.value + return + case .acquiring, .running: + break + } + let invalidatedGeneration = previousState.generation.successor() + let task = Task { @MainActor [weak self] in + guard let self else { return } + await self.performRuntimeStop( + previousState: previousState, + invalidatedGeneration: invalidatedGeneration + ) + } + runtimeState = .transitioning( + generation: invalidatedGeneration, + task: task + ) + await task.value + } + + private func performRuntimeStop( + previousState: ReviewStoreRuntimeState, + invalidatedGeneration: ReviewRuntimeGeneration + ) async { + switch previousState { + case .acquiring(_, let task): + task.cancel() + await task.value + case .running(_, let runtime, _): + await stopPublishedRuntime(runtime) + case .stopped, .transitioning: + break + } + guard case .transitioning(let currentGeneration, _) = runtimeState, + currentGeneration == invalidatedGeneration + else { + return + } + runtimeState = .stopped(invalidatedGeneration) + transitionToStopped() + } + + private func stopPublishedRuntime(_ runtime: PreparedRuntime) async { let locallyCancelledJobIDs: [String] if backend.handlesActiveReviewStopCleanup { locallyCancelledJobIDs = [] @@ -163,16 +238,206 @@ public final class CodexReviewStore { cancelAndDetachReviewWorkersForRuntimeStop( jobIDs: Array(Set(locallyCancelledJobIDs + remainingLocallyCancelledJobIDs)) ) - transitionToStopped() + await runtime.handle.closeAdmission() + do { + try await backend.mcpServerLifecycle.stop() + } catch { + writeDiagnosticsIfNeeded() + } + do { + try await backend.mcpServerLifecycle.waitUntilStopped() + } catch { + writeDiagnosticsIfNeeded() + } + do { + try await runtime.handle.close(purpose: .stop) + } catch { + writeDiagnosticsIfNeeded() + } + do { + try await runtime.handle.waitUntilClosed() + } catch { + writeDiagnosticsIfNeeded() + } } public func restart() async { - await stop() await start(forceRestartIfNeeded: true) } public func waitUntilStopped() async { + if case .transitioning(_, let task) = runtimeState { + await task.value + } await backend.waitUntilStopped() + try? await backend.mcpServerLifecycle.waitUntilStopped() + } + + private func performRuntimeAcquisition( + generation: ReviewRuntimeGeneration, + purpose: ReviewRuntimeTransitionPurpose + ) async { + var preparedMCPServer: PreparedMCPServer? + var preparedRuntime: PreparedRuntime? + do { + let mcpServer = try await backend.mcpServerLifecycle.prepare() + preparedMCPServer = mcpServer + + let runtime = try await backend.prepareRuntime( + generation: generation, + purpose: purpose + ) + preparedRuntime = runtime + guard isCurrentAcquisition(generation) else { + await closeStaleRuntime( + runtime, + mcpServerWasPrepared: true, + purpose: purpose + ) + return + } + + try await runtime.handle.activate() + guard isCurrentAcquisition(generation) else { + await closeStaleRuntime( + runtime, + mcpServerWasPrepared: true, + purpose: purpose + ) + return + } + publishRuntimeSnapshot(runtime.snapshot) + + let mcpSnapshot = try await backend.mcpServerLifecycle.activate( + mcpServer.generation + ) + guard isCurrentAcquisition(generation) else { + await closeStaleRuntime( + runtime, + mcpServerWasPrepared: true, + purpose: purpose + ) + return + } + + runtimeState = .running( + generation: generation, + runtime: runtime, + mcpGeneration: mcpServer.generation + ) + publishMCPServer(serverURL: mcpSnapshot.serverURL) + } catch { + if let preparedRuntime { + await closeStaleRuntime( + preparedRuntime, + mcpServerWasPrepared: preparedMCPServer != nil, + purpose: purpose + ) + } else if preparedMCPServer != nil { + await stopPreparedMCPServer() + } + guard isCurrentAcquisition(generation) else { + return + } + runtimeState = .stopped(generation) + transitionToFailed(error.localizedDescription) + } + } + + private func isCurrentAcquisition( + _ generation: ReviewRuntimeGeneration + ) -> Bool { + guard case .acquiring(let currentGeneration, _) = runtimeState else { + return false + } + return currentGeneration == generation + } + + private func closeStaleRuntime( + _ runtime: PreparedRuntime, + mcpServerWasPrepared: Bool, + purpose: ReviewRuntimeTransitionPurpose + ) async { + await runtime.handle.closeAdmission() + do { + try await runtime.handle.close(purpose: purpose) + } catch { + writeDiagnosticsIfNeeded() + } + do { + try await runtime.handle.waitUntilClosed() + } catch { + writeDiagnosticsIfNeeded() + } + if mcpServerWasPrepared { + await stopPreparedMCPServer() + } + } + + private func stopPreparedMCPServer() async { + do { + try await backend.mcpServerLifecycle.stop() + } catch { + writeDiagnosticsIfNeeded() + } + do { + try await backend.mcpServerLifecycle.waitUntilStopped() + } catch { + writeDiagnosticsIfNeeded() + } + } + + private func publishRuntimeSnapshot(_ snapshot: RuntimePublicationSnapshot) { + settings.apply(snapshot: snapshot.settings) + applyRuntimeAuthenticationSnapshot(snapshot.authentication) + } + + private func publishMCPServer(serverURL: URL?) { + transitionToRunning(serverURL: serverURL) + startAccountRateLimitAutoRefresh() + } + + private func applyRuntimeAuthenticationSnapshot( + _ snapshot: CodexReviewBackendModel.Auth.Snapshot + ) { + 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) } public func refreshAuthentication() async { diff --git a/Sources/CodexReview/Store/CodexReviewStoreBackend.swift b/Sources/CodexReview/Store/CodexReviewStoreBackend.swift index 1d6d75df..ff3a3577 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreBackend.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreBackend.swift @@ -28,9 +28,13 @@ 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 stop(store: CodexReviewStore) async func waitUntilStopped() async func refreshAuth(auth: CodexReviewAuthModel) async diff --git a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift index 13b75fed..5fe51202 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift @@ -54,6 +54,14 @@ extension CodexReviewStore { sessionID: String, request: CodexReviewAPI.Start.Request ) async throws -> String { + switch runtimeState { + case .acquiring, .transitioning: + throw CodexReviewAPI.Error.io( + "Review runtime transition is in progress." + ) + case .stopped, .running: + break + } guard closedSessions.contains(sessionID) == false else { throw CodexReviewAPI.Error.invalidArguments("Review session \(sessionID) is closed.") } diff --git a/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift b/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift index c4c58067..d3b4ea32 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 2702790c..44d495fa 100644 --- a/Sources/CodexReviewHost/CodexReviewHost.swift +++ b/Sources/CodexReviewHost/CodexReviewHost.swift @@ -60,6 +60,7 @@ package final class CodexReviewHost { @MainActor private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { let seed = CodexReviewStoreSeed() + let mcpServerLifecycle: any MCPServerLifecycleOwner = NoMCPServerLifecycleOwner() private let backend: any CodexReviewBackend private var currentSettingsSnapshot = CodexReviewSettings.Snapshot() private var loginChallenge: CodexReviewBackendModel.Login.Challenge? @@ -79,8 +80,23 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { func attachStore(_: CodexReviewStore) {} - func start(store _: CodexReviewStore, forceRestartIfNeeded _: Bool) async { - active = true + func prepareRuntime( + generation _: ReviewRuntimeGeneration, + purpose _: ReviewRuntimeTransitionPurpose + ) async throws -> PreparedRuntime { + let authentication = try await backend.readAuth() + let settings = try await Self.monitorSettings(from: backend.readSettings()) + let handle = DirectRuntimeLifecycleHandle( + onActivate: { [weak self] in self?.active = true }, + onClose: { [weak self] in self?.active = false } + ) + return .init( + snapshot: .init( + authentication: authentication, + settings: settings + ), + handle: handle + ) } func stop(store _: CodexReviewStore) async { @@ -317,6 +333,41 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { } } +@MainActor +private final class DirectRuntimeLifecycleHandle: RuntimeLifecycleHandle { + private let onActivate: @MainActor @Sendable () -> Void + private let onClose: @MainActor @Sendable () -> Void + private var didClose = false + + init( + onActivate: @escaping @MainActor @Sendable () -> Void, + onClose: @escaping @MainActor @Sendable () -> Void + ) { + self.onActivate = onActivate + self.onClose = onClose + } + + func activate() async throws { + onActivate() + } + + func closeAdmission() async {} + + func close(purpose _: ReviewRuntimeTransitionPurpose) async throws { + guard didClose == false else { return } + didClose = true + onClose() + } + + func waitUntilClosed() async throws { + guard didClose else { + throw ReviewLifecycleResourceFailure.client( + "Direct runtime wait began before close." + ) + } + } +} + 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 f5f05e19..da52bf03 100644 --- a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift +++ b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift @@ -44,9 +44,16 @@ package protocol CodexReviewMCPHTTPServing: AnyObject, Sendable { var url: URL { get async } func start() async throws + func closeAdmission() async + func waitForAdmittedHandlers() async func stop() async } +extension CodexReviewMCPHTTPServing { + package func closeAdmission() async {} + package func waitForAdmittedHandlers() async {} +} + extension CodexReviewMCPHTTPServer: CodexReviewMCPHTTPServing {} @MainActor @@ -148,7 +155,8 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private var client: AppServerClient? private var appServerBackend: AppServerCodexReviewBackend? - private var mcpHTTPServer: (any CodexReviewMCPHTTPServing)? + private var activeRuntimeHandle: LiveRuntimeLifecycleHandle? + private var acceptsRuntimeRequests = false private var loginChallenge: CodexReviewBackendModel.Login.Challenge? private var loginBackend: AppServerCodexReviewBackend? private var loginClient: AppServerClient? @@ -161,13 +169,10 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private var loginNotificationTask: Task? private var settingsSnapshot = CodexReviewSettings.Snapshot() private let codexHomeURL: URL - private let mcpHTTPServerConfiguration: CodexReviewMCPHTTPServer.Configuration private let nativeAuthenticationConfiguration: CodexReviewNativeAuthentication.Configuration? private let webAuthenticationSessionFactory: CodexReviewNativeAuthentication.WebSessionFactory private let externalURLOpener: ExternalURLOpener - private let mcpHTTPServerFactory: MCPHTTPServerFactory? - private let mcpPortOwnerResolver: CodexReviewMCPPortOwnerResolver - private let mcpHTTPServerBindChecker: CodexReviewMCPHTTPServerBindChecker + private let mcpLifecycleOwner: LiveMCPServerLifecycleOwner private let appServerRuntimeFactory: AppServerRuntimeFactory private let shutdownCleanupTimeout: Duration private weak var attachedStore: CodexReviewStore? @@ -194,7 +199,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { runtimePreferences: runtimePreferences, environment: environment ) - self.mcpHTTPServerConfiguration = .init( + let mcpHTTPServerConfiguration = CodexReviewMCPHTTPServer.Configuration( host: runtimePreferences.mcpHost, port: runtimePreferences.mcpPort, endpoint: runtimePreferences.mcpPath @@ -202,9 +207,14 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { self.nativeAuthenticationConfiguration = nativeAuthenticationConfiguration self.webAuthenticationSessionFactory = webAuthenticationSessionFactory self.externalURLOpener = externalURLOpener - self.mcpHTTPServerFactory = mcpHTTPServerFactory - self.mcpPortOwnerResolver = mcpPortOwnerResolver ?? Self.defaultMCPPortOwnerResolver - self.mcpHTTPServerBindChecker = mcpHTTPServerBindChecker ?? Self.defaultMCPHTTPServerBindChecker + let resolvedPortOwnerResolver = mcpPortOwnerResolver ?? Self.defaultMCPPortOwnerResolver + let resolvedBindChecker = mcpHTTPServerBindChecker ?? Self.defaultMCPHTTPServerBindChecker + self.mcpLifecycleOwner = LiveMCPServerLifecycleOwner( + configuration: mcpHTTPServerConfiguration, + factory: mcpHTTPServerFactory, + portOwnerResolver: resolvedPortOwnerResolver, + bindChecker: resolvedBindChecker + ) self.shutdownCleanupTimeout = shutdownCleanupTimeout self.appServerRuntimeFactory = appServerRuntimeFactory ?? Self.makeAppServerRuntimeFactory( codexExecutablePath: runtimePreferences.codexExecutablePath @@ -223,6 +233,10 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { client != nil } + var mcpServerLifecycle: any MCPServerLifecycleOwner { + mcpLifecycleOwner + } + var handlesActiveReviewStopCleanup: Bool { true } @@ -344,80 +358,79 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { func attachStore(_ store: CodexReviewStore) { attachedStore = store + mcpLifecycleOwner.attachStore(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 + func prepareRuntime( + generation _: ReviewRuntimeGeneration, + purpose _: ReviewRuntimeTransitionPurpose + ) async throws -> PreparedRuntime { + logger.info("Preparing review runtime") + let runtime = try await appServerRuntimeFactory(codexHomeURL) + do { + 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, + snapshot: .init( + authentication: authentication, + settings: settings + ) + ) + logger.info("Review runtime prepared") + return PreparedRuntime(snapshot: handle.snapshot, handle: handle) + } catch { + await closeClientAfterFailure(runtime.client) + throw error } - if forceRestartIfNeeded { - await stop(store: store) + } + + func activateRuntime(_ handle: LiveRuntimeLifecycleHandle) throws { + guard activeRuntimeHandle == nil else { + throw ReviewLifecycleResourceFailure.client( + "A review runtime is already active." + ) + } + guard let store = attachedStore else { + throw ReviewLifecycleResourceFailure.client( + "Review runtime activation requires its attached Store." + ) } + activeRuntimeHandle = handle + acceptsRuntimeRequests = true + client = handle.client + appServerBackend = handle.backend + settingsSnapshot = handle.snapshot.settings + observeAuthNotifications( + client: handle.client, + backend: handle.backend, + store: store + ) + } - var startedClient: AppServerClient? - var startedHTTPServer: (any CodexReviewMCPHTTPServing)? - do { - if mcpHTTPServerFactory != nil { - try await mcpHTTPServerBindChecker(mcpHTTPServerConfiguration) - } - let runtime = try await appServerRuntimeFactory(codexHomeURL) - let client = runtime.client - let backend = runtime.backend - startedClient = client - self.client = client - self.appServerBackend = backend - observeAuthNotifications(client: client, backend: backend, store: store) - if let mcpHTTPServerFactory { - let mcpHTTPServer = mcpHTTPServerFactory(store, mcpHTTPServerConfiguration) - try await mcpHTTPServer.start() - startedHTTPServer = mcpHTTPServer - 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") - } catch { - let failureMessage = await runtimeStartupFailureMessage(for: error) - logger.error("Review runtime failed to start: \(failureMessage, privacy: .public)") - await startedHTTPServer?.stop() - await closeClientAfterFailure(startedClient) - self.client = nil - self.appServerBackend = nil - self.mcpHTTPServer = nil - authNotificationTask?.cancel() - authNotificationTask = nil - store.transitionToFailed(failureMessage) - } - } - - private func runtimeStartupFailureMessage(for error: Error) async -> String { - if let mcpHTTPServerError = error as? CodexReviewMCPHTTPServer.Error { - switch mcpHTTPServerError { - case .addressInUse: - return await mcpAddressInUseMessage() - } + func deactivateRuntime( + _ handle: LiveRuntimeLifecycleHandle + ) -> Task? { + guard activeRuntimeHandle === handle else { + return nil } - return error.localizedDescription + activeRuntimeHandle = nil + acceptsRuntimeRequests = false + client = nil + appServerBackend = nil + let task = authNotificationTask + authNotificationTask = nil + task?.cancel() + return task } - private func mcpAddressInUseMessage() async -> String { - let endpoint = mcpHTTPServerConfiguration.url() - var message = "MCP endpoint \(endpoint.absoluteString) is already in use" - if let owner = await mcpPortOwnerResolver(mcpHTTPServerConfiguration) { - message += " by PID \(owner.processIdentifier)" - if let command = owner.command?.trimmingCharacters(in: .whitespacesAndNewlines), - command.isEmpty == false - { - message += " (\(command))" - } + func closeRuntimeAdmission(_ handle: LiveRuntimeLifecycleHandle) { + guard activeRuntimeHandle === handle else { + return } - message += ". Quit that process or change the MCP port in Settings, then reset the server." - return message + acceptsRuntimeRequests = false } private func cancelActiveReviewsForRuntimeTeardown( @@ -441,10 +454,8 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } func stop(store: CodexReviewStore) async { - let client = client let appServerBackend = appServerBackend - let mcpHTTPServer = mcpHTTPServer - let hasRuntimeState = client != nil || appServerBackend != nil || mcpHTTPServer != nil + let hasRuntimeState = client != nil || appServerBackend != nil let loginCleanup = takeLoginRuntimeForCleanup() guard hasRuntimeState || loginCleanup.isEmpty == false else { return @@ -458,15 +469,8 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { timeoutWarning: "Timed out cleaning active reviews before stopping runtime" ) } - self.client = nil - self.mcpHTTPServer = nil - authNotificationTask?.cancel() - authNotificationTask = nil - await mcpHTTPServer?.stop() - self.appServerBackend = nil await cleanupLoginRuntime(loginCleanup) - await closeClientAfterFailure(client) - logger.info("Review runtime stopped") + logger.info("Review runtime semantic work stopped") } func waitUntilStopped() async {} @@ -617,8 +621,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { guard let runtimeStore else { return } - await stop(store: runtimeStore) - await start(store: runtimeStore, forceRestartIfNeeded: true) + await runtimeStore.restart() } func removeAccount(auth: CodexReviewAuthModel, accountKey: String) async throws { @@ -659,8 +662,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { guard let runtimeStore else { return } - await stop(store: runtimeStore) - await start(store: runtimeStore, forceRestartIfNeeded: true) + await runtimeStore.restart() } } @@ -715,8 +717,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.restart() } } @@ -994,21 +995,21 @@ 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) } func forceCloseReviewConnection() async throws { - guard let appServerBackend else { + guard acceptsRuntimeRequests, let appServerBackend else { throw ReviewRuntimeCloseFailure.connection("Review runtime is not running.") } try await appServerBackend.forceCloseReviewConnection() @@ -1017,7 +1018,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { func prepareReviewRecovery( _ candidate: ReviewRecoveryCandidate ) async throws -> ReviewRecoveryHandoff { - guard let appServerBackend else { + guard acceptsRuntimeRequests, let appServerBackend else { throw CodexReviewAPI.Error.io("Review runtime is not running.") } return try await appServerBackend.prepareReviewRecovery(candidate) @@ -1028,7 +1029,7 @@ 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.resumeReviewRecovery( @@ -1039,7 +1040,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { - guard let appServerBackend else { + guard acceptsRuntimeRequests, let appServerBackend else { throw ReviewRuntimeCloseFailure.cleanup("Review runtime is not running.") } try await appServerBackend.cleanupReview(run) @@ -1140,7 +1141,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { store: CodexReviewStore ) async { let loginCleanup = takeLoginRuntimeForCleanup() - guard client != nil || appServerBackend != nil || mcpHTTPServer != nil || loginCleanup.isEmpty == false else { + guard client != nil || appServerBackend != nil || loginCleanup.isEmpty == false else { return } let message = "Review runtime stopped unexpectedly: \(error.localizedDescription)" @@ -1153,13 +1154,9 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { ) } let failedClient = client - let failedMCPHTTPServer = mcpHTTPServer - client = nil - appServerBackend = nil - mcpHTTPServer = nil + acceptsRuntimeRequests = false authNotificationTask = nil store.transitionToFailed(message) - await failedMCPHTTPServer?.stop() await cleanupLoginRuntime(loginCleanup) await closeClientAfterFailure(failedClient) } @@ -1640,6 +1637,263 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } } +@MainActor +private final class LiveRuntimeLifecycleHandle: RuntimeLifecycleHandle { + fileprivate let client: AppServerClient + fileprivate let backend: AppServerCodexReviewBackend + fileprivate let snapshot: RuntimePublicationSnapshot + + private weak var owner: LiveCodexReviewStoreBackend? + private var isActivated = false + private var closeTask: Task, Never>? + + init( + owner: LiveCodexReviewStoreBackend, + client: AppServerClient, + backend: AppServerCodexReviewBackend, + snapshot: RuntimePublicationSnapshot + ) { + self.owner = owner + self.client = client + self.backend = backend + self.snapshot = snapshot + } + + func activate() async throws { + guard isActivated == false, closeTask == nil else { + throw ReviewLifecycleResourceFailure.client( + "Runtime activation requires one inert open handle." + ) + } + guard let owner else { + throw ReviewLifecycleResourceFailure.client( + "Runtime activation lost its Host owner." + ) + } + 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 client = client + let authObservationTask = owner?.deactivateRuntime(self) + let newTask = Task, Never> { @MainActor in + var failures: [ReviewLifecycleResourceFailure] = [] + authObservationTask?.cancel() + do { + try await client.close() + } catch { + failures.append(.client(error.localizedDescription)) + } + await authObservationTask?.value + if let first = failures.first { + return Result.failure(.init( + first: first, + additionalInLifecycleOrder: Array(failures.dropFirst()) + )) + } + return Result.success(()) + } + closeTask = newTask + task = newTask + } + try await task.value.get() + } + + func waitUntilClosed() async throws { + guard let closeTask else { + throw ReviewLifecycleResourceFailure.client( + "Runtime completion wait began before close." + ) + } + try await closeTask.value.get() + } +} + +@MainActor +private final class LiveMCPServerLifecycleOwner: MCPServerLifecycleOwner { + typealias Factory = LiveCodexReviewStoreBackend.MCPHTTPServerFactory + + private enum State { + case stopped + case prepared(MCPServerGeneration, (any CodexReviewMCPHTTPServing)?) + case running(MCPServerGeneration, (any CodexReviewMCPHTTPServing)?) + case stopping + case closing + case closed + } + + private let configuration: CodexReviewMCPHTTPServer.Configuration + private let factory: Factory? + private let portOwnerResolver: CodexReviewMCPPortOwnerResolver + private let bindChecker: CodexReviewMCPHTTPServerBindChecker + private weak var store: CodexReviewStore? + private var state: State = .stopped + private var nextGeneration: UInt64 = 0 + + init( + configuration: CodexReviewMCPHTTPServer.Configuration, + factory: Factory?, + portOwnerResolver: @escaping CodexReviewMCPPortOwnerResolver, + bindChecker: @escaping CodexReviewMCPHTTPServerBindChecker + ) { + self.configuration = configuration + self.factory = factory + self.portOwnerResolver = portOwnerResolver + self.bindChecker = bindChecker + } + + func attachStore(_ store: CodexReviewStore) { + self.store = store + } + + func prepare() async throws -> PreparedMCPServer { + guard case .stopped = state else { + throw ReviewLifecycleResourceFailure.mcpServer( + "MCP preparation requires stopped state." + ) + } + nextGeneration &+= 1 + let generation = MCPServerGeneration(rawValue: nextGeneration) + guard let factory else { + state = .prepared(generation, nil) + return .init(generation: generation) + } + guard let store else { + throw ReviewLifecycleResourceFailure.mcpServer( + "MCP preparation requires its attached Store." + ) + } + do { + try await bindChecker(configuration) + } catch { + throw await mappedPreparationFailure(error) + } + let server = factory(store, configuration) + state = .prepared(generation, server) + return .init(generation: generation) + } + + func activate( + _ generation: MCPServerGeneration + ) async throws -> MCPServerPublicationSnapshot { + guard case .prepared(generation, let server) = state else { + throw ReviewLifecycleResourceFailure.mcpServer( + "MCP activation requires its exact prepared generation." + ) + } + guard let server else { + state = .running(generation, nil) + return .init(serverURL: nil) + } + do { + try await server.start() + let url = await server.url + state = .running(generation, server) + return .init(serverURL: url) + } catch { + await server.stop() + state = .stopped + throw error + } + } + + func closeAdmission() async { + guard case .running(_, let server) = state else { + return + } + await server?.closeAdmission() + } + + func drainAdmittedHandlers() async throws { + guard case .running(_, let server) = state else { + return + } + await server?.waitForAdmittedHandlers() + } + + func stop() async throws { + switch state { + case .stopped: + return + case .prepared(_, let server), .running(_, let server): + state = .stopping + await server?.closeAdmission() + await server?.stop() + state = .stopped + case .stopping: + return + case .closing, .closed: + throw ReviewLifecycleResourceFailure.mcpServer( + "MCP owner is closing or closed." + ) + } + } + + func waitUntilStopped() async throws { + guard case .stopped = state else { + throw ReviewLifecycleResourceFailure.mcpServer( + "MCP owner did not stop." + ) + } + } + + func close() async throws { + switch state { + case .closed: + return + case .prepared(_, let server), .running(_, let server): + state = .closing + await server?.closeAdmission() + await server?.stop() + state = .closed + case .stopped: + state = .closed + case .stopping, .closing: + throw ReviewLifecycleResourceFailure.mcpServer( + "MCP owner already has an in-flight lifecycle transition." + ) + } + } + + func waitUntilClosed() async throws { + guard case .closed = state else { + throw ReviewLifecycleResourceFailure.mcpServer( + "MCP owner did not close." + ) + } + } + + private func mappedPreparationFailure( + _ error: any Error + ) async -> ReviewLifecycleResourceFailure { + guard let mcpError = error as? CodexReviewMCPHTTPServer.Error, + case .addressInUse = mcpError + else { + return .mcpServer(error.localizedDescription) + } + let endpoint = configuration.url() + var message = "MCP endpoint \(endpoint.absoluteString) is already in use" + if let owner = await portOwnerResolver(configuration) { + message += " by PID \(owner.processIdentifier)" + if let command = owner.command?.trimmingCharacters(in: .whitespacesAndNewlines), + command.isEmpty == false { + message += " (\(command))" + } + } + message += ". Quit that process or change the MCP port in Settings, then reset the server." + return .mcpServer(message) + } +} + @MainActor private struct AppServerRuntime: Sendable { var client: AppServerClient diff --git a/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift b/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift index 484614ac..47b66026 100644 --- a/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift +++ b/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift @@ -122,6 +122,9 @@ package actor CodexReviewMCPHTTPServer { private var sessions: [String: SessionContext] = [:] private var cleanupTask: Task? private var boundURL: URL? + private var acceptsRequests = false + private var admittedRequestCount = 0 + private var admittedRequestDrainWaiters: [CheckedContinuation] = [] package init( adapter: CodexReviewMCPServer, @@ -194,6 +197,7 @@ package actor CodexReviewMCPHTTPServer { ).get() self.eventLoopGroup = group self.channel = channel + acceptsRequests = true let actualPort = channel.localAddress?.port boundURL = configuration.url(boundPort: actualPort) cleanupTask = Task { [weak self] in @@ -210,11 +214,12 @@ package actor CodexReviewMCPHTTPServer { } package func stop() async { + await closeAdmission() cleanupTask?.cancel() + await cleanupTask?.value cleanupTask = nil await closeAllSessions() - try? await channel?.close() - channel = nil + await waitForAdmittedHandlers() if let eventLoopGroup { try? await eventLoopGroup.shutdownGracefully() } @@ -223,11 +228,46 @@ package actor CodexReviewMCPHTTPServer { logger.info("MCP Streamable HTTP server stopped") } + package func closeAdmission() async { + guard acceptsRequests || channel != nil else { + return + } + acceptsRequests = false + try? await channel?.close() + channel = nil + } + + package func waitForAdmittedHandlers() async { + guard admittedRequestCount > 0 else { + return + } + await withCheckedContinuation { continuation in + if admittedRequestCount == 0 { + continuation.resume() + } else { + admittedRequestDrainWaiters.append(continuation) + } + } + } + package func handleHTTPRequest(_ request: HTTPRequest) async -> HTTPResponse { - await handleTrackedHTTPRequest(request).response + await performTrackedHTTPRequest(request).response } fileprivate func handleTrackedHTTPRequest(_ request: HTTPRequest) async -> TrackedHTTPResponse { + guard acceptsRequests else { + return .init(response: .error( + statusCode: 503, + .internalError("MCP server is not accepting requests.") + )) + } + admittedRequestCount += 1 + let tracked = await performTrackedHTTPRequest(request) + finishAdmittedRequest() + return tracked + } + + private func performTrackedHTTPRequest(_ request: HTTPRequest) async -> TrackedHTTPResponse { let sessionID = request.header(HTTPHeaderName.sessionID) if let sessionID, var session = sessions[sessionID] { @@ -260,6 +300,22 @@ package actor CodexReviewMCPHTTPServer { ) } + private func finishAdmittedRequest() { + precondition( + admittedRequestCount > 0, + "CodexReviewMCPHTTPServer owns one completion per admitted request." + ) + admittedRequestCount -= 1 + guard admittedRequestCount == 0 else { + return + } + let waiters = admittedRequestDrainWaiters + admittedRequestDrainWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume() + } + } + private func createSessionAndHandle(_ request: HTTPRequest) async -> TrackedHTTPResponse { let sessionID = UUID().uuidString let clientSession = MCPClientSessionState() diff --git a/Sources/CodexReviewTesting/TestSupport.swift b/Sources/CodexReviewTesting/TestSupport.swift index f3399f7a..373618da 100644 --- a/Sources/CodexReviewTesting/TestSupport.swift +++ b/Sources/CodexReviewTesting/TestSupport.swift @@ -872,6 +872,53 @@ 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 closeCallCount = 0 + package private(set) var waitUntilClosedCallCount = 0 + + private let onActivate: @MainActor @Sendable () -> Void + private let onClose: @MainActor @Sendable () -> Void + 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 close(purpose _: ReviewRuntimeTransitionPurpose) async throws { + closeCallCount += 1 + guard didClose == false else { + return + } + didClose = true + onClose() + } + + package func waitUntilClosed() async throws { + waitUntilClosedCallCount += 1 + guard didClose else { + throw ReviewLifecycleResourceFailure.client( + "Testing runtime wait began before close." + ) + } + } +} + @MainActor package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { package let reviewBackend: FakeCodexReviewBackend @@ -881,7 +928,12 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { package private(set) var startRequests: [Bool] = [] package private(set) var reviewStartOwnershipSnapshots: [StoreAttemptSnapshot?] = [] package private(set) var recoveryResumeOwnershipSnapshots: [StoreAttemptSnapshot?] = [] + package let mcpServerLifecycle: any MCPServerLifecycleOwner = NoMCPServerLifecycleOwner() + package private(set) var lastPreparedRuntimeHandle: TestingRuntimeLifecycleHandle? private weak var store: CodexReviewStore? + private var runtimePreparationGate: AsyncGate? + private let runtimePreparationStartedGate = AsyncGate() + private let runtimePreparationCancellationGate = AsyncGate() package init( reviewBackend: FakeCodexReviewBackend, @@ -900,10 +952,45 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { self.store = store } - 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 + } + + package func waitForRuntimePreparation() async { + await runtimePreparationStartedGate.wait() + } + + package func waitForRuntimePreparationCancellation() async { + await runtimePreparationCancellationGate.wait() + } + + 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 runtimePreparationCancellationGate = runtimePreparationCancellationGate + await withTaskCancellationHandler { + await runtimePreparationGate.waitIgnoringCancellation() + } onCancel: { + Task { await runtimePreparationCancellationGate.open() } + } + self.runtimePreparationGate = nil + } + return PreparedRuntime( + snapshot: .init( + authentication: try await reviewBackend.readAuth(), + settings: try await monitoredSettingsSnapshot() + ), + handle: handle + ) } package func stop(store _: CodexReviewStore) async { @@ -1056,15 +1143,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 498f9429..4388daea 100644 --- a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift +++ b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift @@ -232,6 +232,57 @@ struct CodexReviewHostTests { #expect(capturedConfiguration?.port == 54321) #expect(capturedConfiguration?.endpoint == "/custom-mcp") #expect(serverURL.path == "/custom-mcp") + #expect(serverURL.port != 0) + await store.stop() + } + + @Test func liveStoreStopThenStartRebindsMCPWithStableOwner() async throws { + let homeURL = try temporaryHome() + let firstTransport = FakeJSONRPCTransport() + let secondTransport = FakeJSONRPCTransport() + for transport in [firstTransport, secondTransport] { + 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" + ) + } + var transports = [firstTransport, secondTransport] + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + mcpHTTPServerFactory: { store, configuration in + CodexReviewMCPHTTPServer( + adapter: CodexReviewMCPServer(store: store), + configuration: .init( + host: configuration.host, + port: 0, + endpoint: configuration.endpoint + ) + ) + }, + mcpHTTPServerBindChecker: { _ in }, + transportFactory: { _ in transports.removeFirst() } + ) + + await store.start() + let firstURL = try #require(store.serverURL) + #expect(firstURL.port != 0) + + await store.stop() + #expect(store.serverState == .stopped) + + await store.start() + let secondURL = try #require(store.serverURL) + #expect(secondURL.port != 0) + #expect(store.serverState == .running) + #expect(transports.isEmpty) + await store.stop() } diff --git a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift new file mode 100644 index 00000000..23838c2b --- /dev/null +++ b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift @@ -0,0 +1,68 @@ +import Testing +import CodexReview +import CodexReviewTesting + +@Suite("store runtime lifecycle") +@MainActor +struct CodexReviewStoreLifecycleTests { + @Test func stopInvalidatesHeldRuntimePreparationAndClosesStaleHandleOnce() async throws { + let preparationGate = AsyncGate() + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend() + ) + backend.holdRuntimePreparation(with: preparationGate) + let store = CodexReviewStore.makeTestingStore(backend: backend) + + let startTask = Task { @MainActor in + await store.start() + } + await backend.waitForRuntimePreparation() + let handle = try #require(backend.lastPreparedRuntimeHandle) + + let stopTask = Task { @MainActor in + await store.stop() + } + await backend.waitForRuntimePreparationCancellation() + + #expect(handle.activateCallCount == 0) + #expect(handle.closeAdmissionCallCount == 0) + #expect(handle.closeCallCount == 0) + #expect(store.serverState == .starting) + + await preparationGate.open() + await stopTask.value + await startTask.value + + #expect(store.serverState == .stopped) + #expect(store.serverURL == nil) + #expect(handle.activateCallCount == 0) + #expect(handle.closeAdmissionCallCount == 1) + #expect(handle.closeCallCount == 1) + #expect(handle.waitUntilClosedCallCount == 1) + #expect(backend.isActive == false) + } + + @Test func stoppedStoreCanPrepareAndPublishANewRuntimeGeneration() async throws { + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend() + ) + let store = CodexReviewStore.makeTestingStore(backend: backend) + + await store.start() + let firstHandle = try #require(backend.lastPreparedRuntimeHandle) + #expect(store.serverState == .running) + #expect(store.serverURL == nil) + #expect(firstHandle.activateCallCount == 1) + + await store.stop() + #expect(store.serverState == .stopped) + #expect(firstHandle.closeCallCount == 1) + + await store.start() + let secondHandle = try #require(backend.lastPreparedRuntimeHandle) + #expect(secondHandle !== firstHandle) + #expect(secondHandle.activateCallCount == 1) + #expect(store.serverState == .running) + #expect(store.serverURL == nil) + } +} diff --git a/Tests/ReviewUITests/ReviewUITests.swift b/Tests/ReviewUITests/ReviewUITests.swift index b8cce071..939355fc 100644 --- a/Tests/ReviewUITests/ReviewUITests.swift +++ b/Tests/ReviewUITests/ReviewUITests.swift @@ -6713,12 +6713,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 { @@ -6750,13 +6753,6 @@ final class AuthActionBackend: PreviewCodexReviewStoreBackend { ) } - override func start( - store _: CodexReviewStore, - forceRestartIfNeeded _: Bool - ) async { - isActive = true - } - override func stop(store _: CodexReviewStore) async { isActive = false } @@ -6793,12 +6789,6 @@ final class FailingCancellationBackend: PreviewCodexReviewStoreBackend { ) } - override func start( - store _: CodexReviewStore, - forceRestartIfNeeded _: Bool - ) async { - } - override func stop(store _: CodexReviewStore) async { } @@ -6844,12 +6834,6 @@ final class BlockingSettingsBackend: PreviewCodexReviewStoreBackend { ) } - override func start( - store _: CodexReviewStore, - forceRestartIfNeeded _: Bool - ) async { - } - override func stop(store _: CodexReviewStore) async { } From 88c8cee8b2b42993d4d8e00df52c1f14eb93f4fb Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 05:23:21 +0900 Subject: [PATCH 25/38] fix(runtime): serialize replacement and MCP lifecycle --- .../CodexReview/ReviewRuntimeLifecycle.swift | 9 +- .../CodexReview/Store/CodexReviewStore.swift | 214 +++++++- .../Store/CodexReviewStoreReviews.swift | 2 + .../LiveCodexReviewStoreBackend.swift | 463 ++++++++++++++--- .../CodexReviewMCPHTTPServer.swift | 424 +++++++++++---- Sources/CodexReviewTesting/TestSupport.swift | 90 +++- .../CodexReviewHostTests.swift | 481 ++++++++++++++++++ .../CodexReviewMCPHTTPServerTests.swift | 93 +++- .../CodexReviewStoreLifecycleTests.swift | 161 ++++++ 9 files changed, 1754 insertions(+), 183 deletions(-) diff --git a/Sources/CodexReview/ReviewRuntimeLifecycle.swift b/Sources/CodexReview/ReviewRuntimeLifecycle.swift index 799e7591..74f3ae57 100644 --- a/Sources/CodexReview/ReviewRuntimeLifecycle.swift +++ b/Sources/CodexReview/ReviewRuntimeLifecycle.swift @@ -91,15 +91,22 @@ package enum ReviewStoreRuntimeState { ) case transitioning( generation: ReviewRuntimeGeneration, + purpose: ReviewRuntimeTransitionPurpose, task: Task ) + case failed( + generation: ReviewRuntimeGeneration, + retainedMCPGeneration: MCPServerGeneration, + serverURL: URL? + ) package var generation: ReviewRuntimeGeneration { switch self { case .stopped(let generation), .acquiring(let generation, _), .running(let generation, _, _), - .transitioning(let generation, _): + .transitioning(let generation, _, _), + .failed(let generation, _, _): generation } } diff --git a/Sources/CodexReview/Store/CodexReviewStore.swift b/Sources/CodexReview/Store/CodexReviewStore.swift index 97e11274..2119f514 100644 --- a/Sources/CodexReview/Store/CodexReviewStore.swift +++ b/Sources/CodexReview/Store/CodexReviewStore.swift @@ -83,9 +83,9 @@ public final class CodexReviewStore { isolated deinit { accountRateLimitAutoRefreshDriver?.cancel() switch runtimeState { - case .acquiring(_, let task), .transitioning(_, let task): + case .acquiring(_, let task), .transitioning(_, _, let task): task.cancel() - case .stopped, .running: + case .stopped, .running, .failed: break } for task in reviewWorkerTasks.values { @@ -139,12 +139,29 @@ public final class CodexReviewStore { public func start(forceRestartIfNeeded: Bool = false) async { switch runtimeState { - case .acquiring, .transitioning: + case .acquiring: + return + case .transitioning(_, _, let task): + await task.value return case .running where forceRestartIfNeeded == false: return - case .running: - await stop() + case .running(let generation, let runtime, let mcpGeneration): + await startRuntimeReplacement( + previousGeneration: generation, + previousRuntime: runtime, + retainedMCPGeneration: mcpGeneration, + retainedServerURL: serverURL + ) + return + case .failed(let generation, let mcpGeneration, let retainedServerURL): + await startRuntimeReplacement( + previousGeneration: generation, + previousRuntime: nil, + retainedMCPGeneration: mcpGeneration, + retainedServerURL: retainedServerURL + ) + return case .stopped: break } @@ -173,16 +190,42 @@ public final class CodexReviewStore { await task.value } + private func startRuntimeReplacement( + previousGeneration: ReviewRuntimeGeneration, + previousRuntime: PreparedRuntime?, + retainedMCPGeneration: MCPServerGeneration, + retainedServerURL: URL? + ) async { + let generation = previousGeneration.successor() + serverState = .starting + writeDiagnosticsIfNeeded() + let task = Task { @MainActor [weak self] in + guard let self else { return } + await self.performRuntimeReplacement( + generation: generation, + previousRuntime: previousRuntime, + retainedMCPGeneration: retainedMCPGeneration, + retainedServerURL: retainedServerURL + ) + } + runtimeState = .transitioning( + generation: generation, + purpose: .restartSameAccount, + task: task + ) + await task.value + } + public func stop() async { let previousState = runtimeState switch previousState { case .stopped: transitionToStopped() return - case .transitioning(_, let task): + case .transitioning(_, .stop, let task): await task.value return - case .acquiring, .running: + case .acquiring, .running, .transitioning, .failed: break } let invalidatedGeneration = previousState.generation.successor() @@ -195,6 +238,7 @@ public final class CodexReviewStore { } runtimeState = .transitioning( generation: invalidatedGeneration, + purpose: .stop, task: task ) await task.value @@ -207,13 +251,20 @@ public final class CodexReviewStore { switch previousState { case .acquiring(_, let task): task.cancel() + await stopPreparedMCPServer() await task.value case .running(_, let runtime, _): await stopPublishedRuntime(runtime) - case .stopped, .transitioning: + case .transitioning(_, _, let task): + task.cancel() + await task.value + await stopPreparedMCPServer() + case .failed: + await stopPreparedMCPServer() + case .stopped: break } - guard case .transitioning(let currentGeneration, _) = runtimeState, + guard case .transitioning(let currentGeneration, .stop, _) = runtimeState, currentGeneration == invalidatedGeneration else { return @@ -223,6 +274,13 @@ public final class CodexReviewStore { } private func stopPublishedRuntime(_ runtime: PreparedRuntime) async { + await performPublishedRuntimeSemanticStop() + await runtime.handle.closeAdmission() + await stopPreparedMCPServer() + await closeAppServerRuntime(runtime, purpose: .stop) + } + + private func performPublishedRuntimeSemanticStop() async { let locallyCancelledJobIDs: [String] if backend.handlesActiveReviewStopCleanup { locallyCancelledJobIDs = [] @@ -238,19 +296,14 @@ public final class CodexReviewStore { cancelAndDetachReviewWorkersForRuntimeStop( jobIDs: Array(Set(locallyCancelledJobIDs + remainingLocallyCancelledJobIDs)) ) - await runtime.handle.closeAdmission() - do { - try await backend.mcpServerLifecycle.stop() - } catch { - writeDiagnosticsIfNeeded() - } - do { - try await backend.mcpServerLifecycle.waitUntilStopped() - } catch { - writeDiagnosticsIfNeeded() - } + } + + private func closeAppServerRuntime( + _ runtime: PreparedRuntime, + purpose: ReviewRuntimeTransitionPurpose + ) async { do { - try await runtime.handle.close(purpose: .stop) + try await runtime.handle.close(purpose: purpose) } catch { writeDiagnosticsIfNeeded() } @@ -266,7 +319,7 @@ public final class CodexReviewStore { } public func waitUntilStopped() async { - if case .transitioning(_, let task) = runtimeState { + if case .transitioning(_, _, let task) = runtimeState { await task.value } await backend.waitUntilStopped() @@ -277,11 +330,20 @@ public final class CodexReviewStore { generation: ReviewRuntimeGeneration, purpose: ReviewRuntimeTransitionPurpose ) async { + guard isCurrentAcquisition(generation) else { + return + } var preparedMCPServer: PreparedMCPServer? var preparedRuntime: PreparedRuntime? do { let mcpServer = try await backend.mcpServerLifecycle.prepare() preparedMCPServer = mcpServer + guard isCurrentAcquisition(generation) else { + if currentTransitionOwnsMCPStop == false { + await stopPreparedMCPServer() + } + return + } let runtime = try await backend.prepareRuntime( generation: generation, @@ -333,7 +395,8 @@ public final class CodexReviewStore { mcpServerWasPrepared: preparedMCPServer != nil, purpose: purpose ) - } else if preparedMCPServer != nil { + } else if preparedMCPServer != nil, + currentTransitionOwnsMCPStop == false { await stopPreparedMCPServer() } guard isCurrentAcquisition(generation) else { @@ -344,6 +407,88 @@ public final class CodexReviewStore { } } + package func performRuntimeAcquisitionForTesting( + generation: ReviewRuntimeGeneration, + purpose: ReviewRuntimeTransitionPurpose + ) async { + await performRuntimeAcquisition( + generation: generation, + purpose: purpose + ) + } + + private func performRuntimeReplacement( + generation: ReviewRuntimeGeneration, + previousRuntime: PreparedRuntime?, + retainedMCPGeneration: MCPServerGeneration, + retainedServerURL: URL? + ) async { + var preparedRuntime: PreparedRuntime? + if let previousRuntime { + await performPublishedRuntimeSemanticStop() + await previousRuntime.handle.closeAdmission() + await closeAppServerRuntime( + previousRuntime, + purpose: .restartSameAccount + ) + } + guard isCurrentTransition(generation, purpose: .restartSameAccount) else { + return + } + do { + let runtime = try await backend.prepareRuntime( + generation: generation, + purpose: .restartSameAccount + ) + preparedRuntime = runtime + guard isCurrentTransition(generation, purpose: .restartSameAccount) else { + await closeStaleRuntime( + runtime, + mcpServerWasPrepared: false, + purpose: .restartSameAccount + ) + return + } + + try await runtime.handle.activate() + guard isCurrentTransition(generation, purpose: .restartSameAccount) else { + await closeStaleRuntime( + runtime, + mcpServerWasPrepared: false, + purpose: .restartSameAccount + ) + return + } + + publishRuntimeSnapshot(runtime.snapshot) + runtimeState = .running( + generation: generation, + runtime: runtime, + mcpGeneration: retainedMCPGeneration + ) + publishMCPServer(serverURL: retainedServerURL) + } catch { + if let preparedRuntime { + await closeStaleRuntime( + preparedRuntime, + mcpServerWasPrepared: false, + purpose: .restartSameAccount + ) + } + guard isCurrentTransition(generation, purpose: .restartSameAccount) else { + return + } + runtimeState = .failed( + generation: generation, + retainedMCPGeneration: retainedMCPGeneration, + serverURL: retainedServerURL + ) + serverURL = retainedServerURL + serverState = .failed(error.localizedDescription) + writeDiagnosticsIfNeeded() + } + } + private func isCurrentAcquisition( _ generation: ReviewRuntimeGeneration ) -> Bool { @@ -353,6 +498,27 @@ public final class CodexReviewStore { return currentGeneration == generation } + private func isCurrentTransition( + _ generation: ReviewRuntimeGeneration, + purpose: ReviewRuntimeTransitionPurpose + ) -> Bool { + guard case .transitioning( + let currentGeneration, + let currentPurpose, + _ + ) = runtimeState else { + return false + } + return currentGeneration == generation && currentPurpose == purpose + } + + private var currentTransitionOwnsMCPStop: Bool { + guard case .transitioning(_, let purpose, _) = runtimeState else { + return false + } + return purpose == .stop || purpose == .applicationClose + } + private func closeStaleRuntime( _ runtime: PreparedRuntime, mcpServerWasPrepared: Bool, @@ -369,7 +535,7 @@ public final class CodexReviewStore { } catch { writeDiagnosticsIfNeeded() } - if mcpServerWasPrepared { + if mcpServerWasPrepared, currentTransitionOwnsMCPStop == false { await stopPreparedMCPServer() } } diff --git a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift index 5fe51202..d1dc800a 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift @@ -59,6 +59,8 @@ extension CodexReviewStore { throw CodexReviewAPI.Error.io( "Review runtime transition is in progress." ) + case .failed: + throw CodexReviewAPI.Error.io("Review runtime is not running.") case .stopped, .running: break } diff --git a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift index da52bf03..0578ac59 100644 --- a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift +++ b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift @@ -40,6 +40,16 @@ package typealias CodexReviewMCPHTTPServerBindChecker = @MainActor @Sendable ( CodexReviewMCPHTTPServer.Configuration ) async throws -> Void +package enum CodexReviewMCPLifecycleCall: Hashable, Sendable { + case stop + case close +} + +package typealias CodexReviewMCPLifecycleCallObserver = @MainActor @Sendable ( + CodexReviewMCPLifecycleCall, + Int +) -> Void + package protocol CodexReviewMCPHTTPServing: AnyObject, Sendable { var url: URL { get async } @@ -78,6 +88,7 @@ public extension CodexReviewStore { externalURLOpener: @escaping @MainActor @Sendable (URL) -> Void = defaultExternalURLOpener, mcpPortOwnerResolver: CodexReviewMCPPortOwnerResolver? = nil, mcpHTTPServerBindChecker: CodexReviewMCPHTTPServerBindChecker? = nil, + mcpLifecycleCallObserver: CodexReviewMCPLifecycleCallObserver? = nil, shutdownCleanupTimeout: Duration = .seconds(2), networkMonitor: any CodexReviewNetworkMonitoring = SystemCodexReviewNetworkMonitor(), networkRecoveryPolicy: CodexReviewNetworkRecoveryPolicy = .default, @@ -92,6 +103,7 @@ public extension CodexReviewStore { externalURLOpener: externalURLOpener, mcpPortOwnerResolver: mcpPortOwnerResolver, mcpHTTPServerBindChecker: mcpHTTPServerBindChecker, + mcpLifecycleCallObserver: mcpLifecycleCallObserver, shutdownCleanupTimeout: shutdownCleanupTimeout, networkMonitor: networkMonitor, networkRecoveryPolicy: networkRecoveryPolicy, @@ -112,6 +124,7 @@ public extension CodexReviewStore { ) -> any CodexReviewMCPHTTPServing)? = nil, mcpPortOwnerResolver: CodexReviewMCPPortOwnerResolver? = nil, mcpHTTPServerBindChecker: CodexReviewMCPHTTPServerBindChecker? = nil, + mcpLifecycleCallObserver: CodexReviewMCPLifecycleCallObserver? = nil, shutdownCleanupTimeout: Duration = .seconds(2), networkMonitor: any CodexReviewNetworkMonitoring = SystemCodexReviewNetworkMonitor(), networkRecoveryPolicy: CodexReviewNetworkRecoveryPolicy = .default, @@ -128,6 +141,7 @@ public extension CodexReviewStore { mcpHTTPServerFactory: mcpHTTPServerFactory, mcpPortOwnerResolver: mcpPortOwnerResolver, mcpHTTPServerBindChecker: mcpHTTPServerBindChecker, + mcpLifecycleCallObserver: mcpLifecycleCallObserver, shutdownCleanupTimeout: shutdownCleanupTimeout, appServerRuntimeFactory: { codexHomeURL in let client = AppServerClient(transport: try await transportFactory(codexHomeURL)) @@ -191,6 +205,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { }, mcpPortOwnerResolver: CodexReviewMCPPortOwnerResolver? = nil, mcpHTTPServerBindChecker: CodexReviewMCPHTTPServerBindChecker? = nil, + mcpLifecycleCallObserver: CodexReviewMCPLifecycleCallObserver? = nil, shutdownCleanupTimeout: Duration = .seconds(2), appServerRuntimeFactory: AppServerRuntimeFactory? = nil ) { @@ -213,7 +228,8 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { configuration: mcpHTTPServerConfiguration, factory: mcpHTTPServerFactory, portOwnerResolver: resolvedPortOwnerResolver, - bindChecker: resolvedBindChecker + bindChecker: resolvedBindChecker, + lifecycleCallObserver: mcpLifecycleCallObserver ) self.shutdownCleanupTimeout = shutdownCleanupTimeout self.appServerRuntimeFactory = appServerRuntimeFactory ?? Self.makeAppServerRuntimeFactory( @@ -1722,33 +1738,69 @@ private final class LiveRuntimeLifecycleHandle: RuntimeLifecycleHandle { private final class LiveMCPServerLifecycleOwner: MCPServerLifecycleOwner { typealias Factory = LiveCodexReviewStoreBackend.MCPHTTPServerFactory + private struct Lease: Sendable { + let generation: MCPServerGeneration + let server: (any CodexReviewMCPHTTPServing)? + } + + private struct Activation: Sendable { + let lease: Lease + let snapshot: MCPServerPublicationSnapshot + } + + private typealias PreparationResult = Result + private typealias ActivationResult = Result + private typealias LifecycleResult = Result + private enum State { case stopped - case prepared(MCPServerGeneration, (any CodexReviewMCPHTTPServing)?) - case running(MCPServerGeneration, (any CodexReviewMCPHTTPServing)?) - case stopping - case closing - case closed + case preparing( + operationID: UInt64, + generation: MCPServerGeneration, + task: Task + ) + case prepared(Lease) + case activating( + operationID: UInt64, + lease: Lease, + task: Task + ) + case running(Lease, MCPServerPublicationSnapshot) + case stopping( + operationID: UInt64, + task: Task + ) + case closing( + operationID: UInt64, + task: Task + ) + case closed(LifecycleResult) } private let configuration: CodexReviewMCPHTTPServer.Configuration private let factory: Factory? private let portOwnerResolver: CodexReviewMCPPortOwnerResolver private let bindChecker: CodexReviewMCPHTTPServerBindChecker + private let lifecycleCallObserver: CodexReviewMCPLifecycleCallObserver? private weak var store: CodexReviewStore? private var state: State = .stopped private var nextGeneration: UInt64 = 0 + private var nextOperationID: UInt64 = 0 + private var stopCallerCount = 0 + private var closeCallerCount = 0 init( configuration: CodexReviewMCPHTTPServer.Configuration, factory: Factory?, portOwnerResolver: @escaping CodexReviewMCPPortOwnerResolver, - bindChecker: @escaping CodexReviewMCPHTTPServerBindChecker + bindChecker: @escaping CodexReviewMCPHTTPServerBindChecker, + lifecycleCallObserver: CodexReviewMCPLifecycleCallObserver? ) { self.configuration = configuration self.factory = factory self.portOwnerResolver = portOwnerResolver self.bindChecker = bindChecker + self.lifecycleCallObserver = lifecycleCallObserver } func attachStore(_ store: CodexReviewStore) { @@ -1756,90 +1808,191 @@ private final class LiveMCPServerLifecycleOwner: MCPServerLifecycleOwner { } func prepare() async throws -> PreparedMCPServer { - guard case .stopped = state else { + let operationID: UInt64 + let generation: MCPServerGeneration + let task: Task + switch state { + case .stopped: + nextGeneration &+= 1 + generation = MCPServerGeneration(rawValue: nextGeneration) + operationID = makeOperationID() + task = makePreparationTask(generation: generation) + state = .preparing( + operationID: operationID, + generation: generation, + task: task + ) + case .preparing(let currentOperationID, let currentGeneration, let currentTask): + operationID = currentOperationID + generation = currentGeneration + task = currentTask + case .prepared(let lease): + return .init(generation: lease.generation) + case .activating, .running, .stopping: throw ReviewLifecycleResourceFailure.mcpServer( "MCP preparation requires stopped state." ) - } - nextGeneration &+= 1 - let generation = MCPServerGeneration(rawValue: nextGeneration) - guard let factory else { - state = .prepared(generation, nil) - return .init(generation: generation) - } - guard let store else { + case .closing, .closed: throw ReviewLifecycleResourceFailure.mcpServer( - "MCP preparation requires its attached Store." + "MCP owner is closing or closed." ) } - do { - try await bindChecker(configuration) - } catch { - throw await mappedPreparationFailure(error) + + let result = await task.value + switch result { + case .failure(let failure): + if case .preparing(let currentOperationID, let currentGeneration, _) = state, + currentOperationID == operationID, + currentGeneration == generation { + state = .stopped + } + throw failure + case .success(let lease): + switch state { + case .preparing(let currentOperationID, let currentGeneration, _) + where currentOperationID == operationID && currentGeneration == generation: + state = .prepared(lease) + return .init(generation: generation) + case .prepared(let currentLease) where currentLease.generation == generation: + return .init(generation: generation) + default: + throw supersededFailure("preparation", generation: generation) + } } - let server = factory(store, configuration) - state = .prepared(generation, server) - return .init(generation: generation) } func activate( _ generation: MCPServerGeneration ) async throws -> MCPServerPublicationSnapshot { - guard case .prepared(generation, let server) = state else { + let operationID: UInt64 + let lease: Lease + let task: Task + switch state { + case .prepared(let preparedLease) where preparedLease.generation == generation: + lease = preparedLease + operationID = makeOperationID() + task = makeActivationTask(lease: lease) + state = .activating( + operationID: operationID, + lease: lease, + task: task + ) + case .activating(let currentOperationID, let currentLease, let currentTask) + where currentLease.generation == generation: + operationID = currentOperationID + lease = currentLease + task = currentTask + case .running(let currentLease, let snapshot) + where currentLease.generation == generation: + return snapshot + default: throw ReviewLifecycleResourceFailure.mcpServer( "MCP activation requires its exact prepared generation." ) } - guard let server else { - state = .running(generation, nil) - return .init(serverURL: nil) - } - do { - try await server.start() - let url = await server.url - state = .running(generation, server) - return .init(serverURL: url) - } catch { - await server.stop() - state = .stopped - throw error + + let result = await task.value + switch result { + case .success(let activation): + switch state { + case .activating(let currentOperationID, let currentLease, _) + where currentOperationID == operationID && currentLease.generation == generation: + state = .running(activation.lease, activation.snapshot) + return activation.snapshot + case .running(let currentLease, let snapshot) + where currentLease.generation == generation: + return snapshot + default: + throw supersededFailure("activation", generation: generation) + } + case .failure(let failure): + let cleanupTask = activationFailureCleanupTask( + operationID: operationID, + lease: lease + ) + if let cleanupTask { + let cleanupResult = await cleanupTask.task.value + finishStoppingIfCurrent( + cleanupTask.operationID, + result: cleanupResult + ) + } + throw failure } } func closeAdmission() async { - guard case .running(_, let server) = state else { + switch state { + case .prepared(let lease), .running(let lease, _): + await lease.server?.closeAdmission() + case .activating(_, let lease, _): + await lease.server?.closeAdmission() + case .stopped, .preparing, .stopping, .closing, .closed: return } - await server?.closeAdmission() } func drainAdmittedHandlers() async throws { - guard case .running(_, let server) = state else { + switch state { + case .running(let lease, _): + await lease.server?.waitForAdmittedHandlers() + case .stopped, .preparing, .prepared, .activating, .stopping, .closing, .closed: return } - await server?.waitForAdmittedHandlers() } func stop() async throws { + stopCallerCount += 1 + lifecycleCallObserver?(.stop, stopCallerCount) + let operationID: UInt64 + let task: Task switch state { case .stopped: return - case .prepared(_, let server), .running(_, let server): - state = .stopping - await server?.closeAdmission() - await server?.stop() - state = .stopped - case .stopping: - return - case .closing, .closed: - throw ReviewLifecycleResourceFailure.mcpServer( - "MCP owner is closing or closed." + case .preparing(_, _, let preparationTask): + operationID = makeOperationID() + task = makeLifecycleTask(preparationTask: preparationTask) + state = .stopping(operationID: operationID, task: task) + case .prepared(let lease), .running(let lease, _): + operationID = makeOperationID() + task = makeLifecycleTask(lease: lease) + state = .stopping(operationID: operationID, task: task) + case .activating(_, let lease, let activationTask): + operationID = makeOperationID() + task = makeLifecycleTask( + activationTask: activationTask, + lease: lease ) + state = .stopping(operationID: operationID, task: task) + case .stopping(let currentOperationID, let currentTask): + operationID = currentOperationID + task = currentTask + case .closing(_, let closeTask): + try await closeTask.value.get() + return + case .closed(let result): + try result.get() + return } + + let result = await task.value + finishStoppingIfCurrent(operationID, result: result) + try result.get() } func waitUntilStopped() async throws { - guard case .stopped = state else { + switch state { + case .stopped: + return + case .stopping(let operationID, let task): + let result = await task.value + finishStoppingIfCurrent(operationID, result: result) + try result.get() + case .closing(_, let task): + try await task.value.get() + case .closed(let result): + try result.get() + case .preparing, .prepared, .activating, .running: throw ReviewLifecycleResourceFailure.mcpServer( "MCP owner did not stop." ) @@ -1847,31 +2000,211 @@ private final class LiveMCPServerLifecycleOwner: MCPServerLifecycleOwner { } func close() async throws { + closeCallerCount += 1 + lifecycleCallObserver?(.close, closeCallerCount) + let operationID: UInt64 + let task: Task switch state { - case .closed: - return - case .prepared(_, let server), .running(_, let server): - state = .closing - await server?.closeAdmission() - await server?.stop() - state = .closed case .stopped: - state = .closed - case .stopping, .closing: - throw ReviewLifecycleResourceFailure.mcpServer( - "MCP owner already has an in-flight lifecycle transition." + operationID = makeOperationID() + task = makeLifecycleTask() + state = .closing(operationID: operationID, task: task) + case .preparing(_, _, let preparationTask): + operationID = makeOperationID() + task = makeLifecycleTask(preparationTask: preparationTask) + state = .closing(operationID: operationID, task: task) + case .prepared(let lease), .running(let lease, _): + operationID = makeOperationID() + task = makeLifecycleTask(lease: lease) + state = .closing(operationID: operationID, task: task) + case .activating(_, let lease, let activationTask): + operationID = makeOperationID() + task = makeLifecycleTask( + activationTask: activationTask, + lease: lease ) + state = .closing(operationID: operationID, task: task) + case .stopping(_, let stopTask): + operationID = makeOperationID() + task = makeLifecycleTask(lifecycleTask: stopTask) + state = .closing(operationID: operationID, task: task) + case .closing(let currentOperationID, let currentTask): + operationID = currentOperationID + task = currentTask + case .closed(let result): + try result.get() + return } + + let result = await task.value + finishClosingIfCurrent(operationID, result: result) + try result.get() } func waitUntilClosed() async throws { - guard case .closed = state else { + switch state { + case .closing(let operationID, let task): + let result = await task.value + finishClosingIfCurrent(operationID, result: result) + try result.get() + case .closed(let result): + try result.get() + case .stopped, .preparing, .prepared, .activating, .running, .stopping: throw ReviewLifecycleResourceFailure.mcpServer( "MCP owner did not close." ) } } + private func makeOperationID() -> UInt64 { + nextOperationID &+= 1 + return nextOperationID + } + + private func makePreparationTask( + generation: MCPServerGeneration + ) -> Task { + let configuration = configuration + let bindChecker = bindChecker + let factory = factory + let store = store + return Task { @MainActor [weak self] in + guard let factory else { + return .success(.init(generation: generation, server: nil)) + } + guard let store else { + return .failure(.mcpServer( + "MCP preparation requires its attached Store." + )) + } + do { + try await bindChecker(configuration) + try Task.checkCancellation() + return .success(.init( + generation: generation, + server: factory(store, configuration) + )) + } catch { + guard let self else { + return .failure(.mcpServer(error.localizedDescription)) + } + return .failure(await self.mappedPreparationFailure(error)) + } + } + } + + private func makeActivationTask( + lease: Lease + ) -> Task { + Task { @MainActor in + guard let server = lease.server else { + return .success(.init( + lease: lease, + snapshot: .init(serverURL: nil) + )) + } + do { + try await server.start() + try Task.checkCancellation() + return .success(.init( + lease: lease, + snapshot: .init(serverURL: await server.url) + )) + } catch { + return .failure(.mcpServer(error.localizedDescription)) + } + } + } + + private func makeLifecycleTask( + preparationTask: Task? = nil, + activationTask: Task? = nil, + lifecycleTask: Task? = nil, + lease initialLease: Lease? = nil + ) -> Task { + Task { @MainActor in + preparationTask?.cancel() + activationTask?.cancel() + var lease = initialLease + if let lease { + await lease.server?.closeAdmission() + } + if let preparationTask, + case .success(let preparedLease) = await preparationTask.value { + lease = preparedLease + await preparedLease.server?.closeAdmission() + } + if let activationTask { + _ = await activationTask.value + } + if let lifecycleTask { + _ = await lifecycleTask.value + } else { + await lease?.server?.stop() + } + return .success(()) + } + } + + private func activationFailureCleanupTask( + operationID: UInt64, + lease: Lease + ) -> (operationID: UInt64, task: Task)? { + switch state { + case .activating(let currentOperationID, let currentLease, _) + where currentOperationID == operationID + && currentLease.generation == lease.generation: + let cleanupOperationID = makeOperationID() + let task = makeLifecycleTask(lease: lease) + state = .stopping(operationID: cleanupOperationID, task: task) + return (cleanupOperationID, task) + case .stopping(let currentOperationID, let task): + return (currentOperationID, task) + case .closing(let currentOperationID, let task): + return (currentOperationID, task) + case .stopped, .preparing, .prepared, .activating, .running, .closed: + return nil + } + } + + private func finishStoppingIfCurrent( + _ operationID: UInt64, + result: LifecycleResult + ) { + guard case .stopping(let currentOperationID, _) = state, + currentOperationID == operationID + else { + return + } + switch result { + case .success: + state = .stopped + case .failure: + state = .stopped + } + } + + private func finishClosingIfCurrent( + _ operationID: UInt64, + result: LifecycleResult + ) { + guard case .closing(let currentOperationID, _) = state, + currentOperationID == operationID + else { + return + } + state = .closed(result) + } + + private func supersededFailure( + _ operation: String, + generation: MCPServerGeneration + ) -> ReviewLifecycleResourceFailure { + .mcpServer( + "MCP \(operation) for generation \(generation.rawValue) was superseded by a lifecycle transition." + ) + } + private func mappedPreparationFailure( _ error: any Error ) async -> ReviewLifecycleResourceFailure { diff --git a/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift b/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift index 47b66026..6686344a 100644 --- a/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift +++ b/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift @@ -98,6 +98,161 @@ package extension CodexReviewMCPHTTPServer { } } +private final class MCPHTTPAdmissionRegistry: @unchecked Sendable { + struct Admission: Sendable { + fileprivate let id: UUID + } + + private struct AdmissionCountWaiter { + let targetCount: Int + let continuation: CheckedContinuation + } + + private let lock = NSLock() + private var acceptsRequests = false + private var admittedRequestIDs: Set = [] + private var drainWaiters: [CheckedContinuation] = [] + private var totalAdmissionCount = 0 + private var admissionCountWaiters: [UUID: AdmissionCountWaiter] = [:] + + func open() { + lock.lock() + precondition( + admittedRequestIDs.isEmpty, + "MCPHTTPAdmissionRegistry must drain one listener generation before reopening." + ) + acceptsRequests = true + lock.unlock() + } + + func close() { + lock.lock() + acceptsRequests = false + lock.unlock() + } + + func admit() -> Admission? { + let waiters: [CheckedContinuation] + lock.lock() + guard acceptsRequests else { + lock.unlock() + return nil + } + let admission = Admission(id: UUID()) + admittedRequestIDs.insert(admission.id) + totalAdmissionCount += 1 + let completedWaiterIDs = admissionCountWaiters.compactMap { id, waiter in + totalAdmissionCount >= waiter.targetCount ? id : nil + } + waiters = completedWaiterIDs.compactMap { + admissionCountWaiters.removeValue(forKey: $0)?.continuation + } + lock.unlock() + for waiter in waiters { + waiter.resume() + } + return admission + } + + func finish(_ admission: Admission) { + let waiters: [CheckedContinuation] + lock.lock() + precondition( + admittedRequestIDs.remove(admission.id) != nil, + "MCPHTTPAdmissionRegistry owns exactly one completion per admitted request." + ) + if admittedRequestIDs.isEmpty { + waiters = drainWaiters + drainWaiters.removeAll(keepingCapacity: false) + } else { + waiters = [] + } + lock.unlock() + for waiter in waiters { + waiter.resume() + } + } + + func waitUntilDrained() async { + await withCheckedContinuation { continuation in + lock.lock() + if admittedRequestIDs.isEmpty { + lock.unlock() + continuation.resume() + } else { + drainWaiters.append(continuation) + lock.unlock() + } + } + } + + func admissionCount() -> Int { + lock.lock() + let count = totalAdmissionCount + lock.unlock() + return count + } + + func waitForAdmissionCount(_ targetCount: Int) async { + await withCheckedContinuation { continuation in + lock.lock() + if totalAdmissionCount >= targetCount { + lock.unlock() + continuation.resume() + } else { + admissionCountWaiters[UUID()] = .init( + targetCount: targetCount, + continuation: continuation + ) + lock.unlock() + } + } + } +} + +private actor MCPHTTPHandlerEntryGate { + private var shouldHoldNextEntry = false + private var releaseWasRequested = false + private var continuation: CheckedContinuation? + + func holdNextEntry() { + precondition( + shouldHoldNextEntry == false && continuation == nil, + "MCPHTTPHandlerEntryGate owns at most one held test entry." + ) + shouldHoldNextEntry = true + releaseWasRequested = false + } + + func waitIfNeeded() async { + guard shouldHoldNextEntry else { + return + } + if releaseWasRequested { + shouldHoldNextEntry = false + releaseWasRequested = false + return + } + await withCheckedContinuation { continuation in + self.continuation = continuation + } + shouldHoldNextEntry = false + releaseWasRequested = false + } + + func release() { + guard shouldHoldNextEntry else { + return + } + if let continuation { + self.continuation = nil + continuation.resume() + } else { + releaseWasRequested = true + } + } +} + package actor CodexReviewMCPHTTPServer { private struct SessionContext { let server: Server @@ -122,9 +277,10 @@ package actor CodexReviewMCPHTTPServer { private var sessions: [String: SessionContext] = [:] private var cleanupTask: Task? private var boundURL: URL? - private var acceptsRequests = false - private var admittedRequestCount = 0 - private var admittedRequestDrainWaiters: [CheckedContinuation] = [] + private let admissionRegistry = MCPHTTPAdmissionRegistry() + private let handlerEntryGate = MCPHTTPHandlerEntryGate() + private var admittedHandlerDrainDidBegin = false + private var admittedHandlerDrainStartWaiters: [CheckedContinuation] = [] package init( adapter: CodexReviewMCPServer, @@ -177,13 +333,19 @@ package actor CodexReviewMCPHTTPServer { return } + let admissionRegistry = admissionRegistry + let handlerEntryGate = handlerEntryGate let group = MultiThreadedEventLoopGroup(numberOfThreads: 1) let bootstrap = ServerBootstrap(group: group) .serverChannelOption(ChannelOptions.backlog, value: 128) .serverChannelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1) .childChannelInitializer { channel in channel.pipeline.configureHTTPServerPipeline().flatMap { - channel.pipeline.addHandler(CodexReviewMCPHTTPHandler(server: self)) + channel.pipeline.addHandler(CodexReviewMCPHTTPHandler( + server: self, + admissionRegistry: admissionRegistry, + entryGate: handlerEntryGate + )) } } .childChannelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1) @@ -195,9 +357,15 @@ package actor CodexReviewMCPHTTPServer { host: configuration.host, port: configuration.port ).get() + guard Task.isCancelled == false else { + try? await channel.close() + try? await group.shutdownGracefully() + throw CancellationError() + } self.eventLoopGroup = group self.channel = channel - acceptsRequests = true + admissionRegistry.open() + admittedHandlerDrainDidBegin = false let actualPort = channel.localAddress?.port boundURL = configuration.url(boundPort: actualPort) cleanupTask = Task { [weak self] in @@ -216,10 +384,13 @@ package actor CodexReviewMCPHTTPServer { package func stop() async { await closeAdmission() cleanupTask?.cancel() + let cleanupTask = cleanupTask + self.cleanupTask = nil + await waitForAdmittedHandlers() await cleanupTask?.value - cleanupTask = nil await closeAllSessions() - await waitForAdmittedHandlers() + try? await channel?.close() + channel = nil if let eventLoopGroup { try? await eventLoopGroup.shutdownGracefully() } @@ -229,42 +400,37 @@ package actor CodexReviewMCPHTTPServer { } package func closeAdmission() async { - guard acceptsRequests || channel != nil else { - return - } - acceptsRequests = false + admissionRegistry.close() try? await channel?.close() channel = nil } package func waitForAdmittedHandlers() async { - guard admittedRequestCount > 0 else { - return - } - await withCheckedContinuation { continuation in - if admittedRequestCount == 0 { - continuation.resume() - } else { - admittedRequestDrainWaiters.append(continuation) - } + admittedHandlerDrainDidBegin = true + let startWaiters = admittedHandlerDrainStartWaiters + admittedHandlerDrainStartWaiters.removeAll(keepingCapacity: false) + for waiter in startWaiters { + waiter.resume() } + await admissionRegistry.waitUntilDrained() } package func handleHTTPRequest(_ request: HTTPRequest) async -> HTTPResponse { - await performTrackedHTTPRequest(request).response - } - - fileprivate func handleTrackedHTTPRequest(_ request: HTTPRequest) async -> TrackedHTTPResponse { - guard acceptsRequests else { - return .init(response: .error( + guard let admission = admissionRegistry.admit() else { + return .error( statusCode: 503, .internalError("MCP server is not accepting requests.") - )) + ) } - admittedRequestCount += 1 - let tracked = await performTrackedHTTPRequest(request) - finishAdmittedRequest() - return tracked + let response = await performTrackedHTTPRequest(request).response + admissionRegistry.finish(admission) + return response + } + + fileprivate func handleAdmittedHTTPRequest( + _ request: HTTPRequest + ) async -> TrackedHTTPResponse { + await performTrackedHTTPRequest(request) } private func performTrackedHTTPRequest(_ request: HTTPRequest) async -> TrackedHTTPResponse { @@ -300,22 +466,6 @@ package actor CodexReviewMCPHTTPServer { ) } - private func finishAdmittedRequest() { - precondition( - admittedRequestCount > 0, - "CodexReviewMCPHTTPServer owns one completion per admitted request." - ) - admittedRequestCount -= 1 - guard admittedRequestCount == 0 else { - return - } - let waiters = admittedRequestDrainWaiters - admittedRequestDrainWaiters.removeAll(keepingCapacity: false) - for waiter in waiters { - waiter.resume() - } - } - private func createSessionAndHandle(_ request: HTTPRequest) async -> TrackedHTTPResponse { let sessionID = UUID().uuidString let clientSession = MCPClientSessionState() @@ -467,6 +617,41 @@ package actor CodexReviewMCPHTTPServer { sessions[sessionID]?.activeRequestCount } + package func listenerIsOpenForTesting() -> Bool { + channel != nil + } + + package func admittedNetworkRequestCountForTesting() -> Int { + admissionRegistry.admissionCount() + } + + package func waitForAdmittedNetworkRequestCountForTesting( + _ count: Int + ) async { + await admissionRegistry.waitForAdmissionCount(count) + } + + package func holdNextNetworkHandlerEntryForTesting() async { + await handlerEntryGate.holdNextEntry() + } + + package func releaseNetworkHandlerEntryForTesting() async { + await handlerEntryGate.release() + } + + package func waitForAdmittedHandlerDrainToBeginForTesting() async { + if admittedHandlerDrainDidBegin { + return + } + await withCheckedContinuation { continuation in + if admittedHandlerDrainDidBegin { + continuation.resume() + } else { + admittedHandlerDrainStartWaiters.append(continuation) + } + } + } + private func closeExpiredSessions(now: Date) async { var expiredSessionIDs: [String] = [] for (sessionID, context) in sessions { @@ -623,13 +808,21 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked } private let server: CodexReviewMCPHTTPServer + private let admissionRegistry: MCPHTTPAdmissionRegistry + private let entryGate: MCPHTTPHandlerEntryGate private var requestState: RequestState? private var activeStreamTask: Task? private var activeStreamID: UUID? private var activeStreamCompletion: ActiveRequestCompletion? - init(server: CodexReviewMCPHTTPServer) { + init( + server: CodexReviewMCPHTTPServer, + admissionRegistry: MCPHTTPAdmissionRegistry, + entryGate: MCPHTTPHandlerEntryGate + ) { self.server = server + self.admissionRegistry = admissionRegistry + self.entryGate = entryGate } func channelRead(context: ChannelHandlerContext, data: NIOAny) { @@ -647,9 +840,21 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked return } requestState = nil + guard let admission = admissionRegistry.admit() else { + writeAdmissionClosedResponse( + version: state.head.version, + context: context + ) + return + } nonisolated(unsafe) let context = context Task { - await handleRequest(state: state, context: context) + await entryGate.waitIfNeeded() + await handleRequest( + state: state, + admission: admission, + context: context + ) } } } @@ -688,8 +893,10 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked private func handleRequest( state: RequestState, + admission: MCPHTTPAdmissionRegistry.Admission, context: ChannelHandlerContext ) async { + defer { admissionRegistry.finish(admission) } let head = state.head let path = head.uri.split(separator: "?").first.map(String.init) ?? head.uri let endpoint = await server.endpoint @@ -703,10 +910,41 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked } let request = makeHTTPRequest(from: state) - let response = await server.handleTrackedHTTPRequest(request) + let response = await server.handleAdmittedHTTPRequest(request) await writeResponse(response, version: head.version, context: context) } + private func writeAdmissionClosedResponse( + version: HTTPVersion, + context: ChannelHandlerContext + ) { + let response = HTTPResponse.error( + statusCode: 503, + .internalError("MCP server is not accepting requests.") + ) + var head = HTTPResponseHead( + version: version, + status: .serviceUnavailable + ) + for (name, value) in response.headers { + head.headers.add(name: name, value: value) + } + let body = response.bodyData + if let body { + head.headers.add(name: "Content-Length", value: "\(body.count)") + } + context.write(wrapOutboundOut(.head(head)), promise: nil) + if let body { + var buffer = context.channel.allocator.buffer(capacity: body.count) + buffer.writeBytes(body) + context.write( + wrapOutboundOut(.body(.byteBuffer(buffer))), + promise: nil + ) + } + context.writeAndFlush(wrapOutboundOut(.end(nil)), promise: nil) + } + private func makeHTTPRequest(from state: RequestState) -> HTTPRequest { var headers: [String: String] = [:] for (name, value) in state.head.headers { @@ -749,58 +987,74 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked switch response { case .stream(let stream, _): let streamID = UUID() - let streamTask = Task { - var head = HTTPResponseHead(version: version, status: status) - for (name, value) in headers { - head.headers.add(name: name, value: value) + let registration = eventLoop.makePromise(of: Void.self) + eventLoop.execute { + guard context.channel.isActive else { + trackedResponse.streamCompletion?.finish() + registration.succeed(()) + return } + let streamTask = Task { + defer { + eventLoop.execute { + if self.activeStreamID == streamID { + self.activeStreamTask = nil + self.activeStreamID = nil + self.activeStreamCompletion = nil + } + } + } + var head = HTTPResponseHead(version: version, status: status) + for (name, value) in headers { + head.headers.add(name: name, value: value) + } - var iterator = stream.makeAsyncIterator() - do { - try Task.checkCancellation() - try await writeResponsePart(.head(head), context: context, eventLoop: eventLoop) - while let chunk = try await iterator.next() { + var iterator = stream.makeAsyncIterator() + do { try Task.checkCancellation() - try await writeResponseBody(chunk, context: context, eventLoop: eventLoop) + try await self.writeResponsePart( + .head(head), + context: context, + eventLoop: eventLoop + ) + while let chunk = try await iterator.next() { + try Task.checkCancellation() + try await self.writeResponseBody( + chunk, + context: context, + eventLoop: eventLoop + ) + } + } catch is CancellationError { + trackedResponse.streamCompletion?.finish() + return + } catch { + trackedResponse.streamCompletion?.finish() + logger.error("MCP SSE stream failed: \(error.localizedDescription, privacy: .public)") } - } catch is CancellationError { - trackedResponse.streamCompletion?.finish() - return - } catch { - trackedResponse.streamCompletion?.finish() - logger.error("MCP SSE stream failed: \(error.localizedDescription, privacy: .public)") - } - guard Task.isCancelled == false else { - return + guard Task.isCancelled == false else { + return + } + try? await self.writeResponsePart( + .end(nil), + context: context, + eventLoop: eventLoop + ) } - try? await writeResponsePart(.end(nil), context: context, eventLoop: eventLoop) - } - eventLoop.execute { context.channel.closeFuture.whenComplete { _ in trackedResponse.streamCompletion?.finish() streamTask.cancel() } - guard context.channel.isActive else { - trackedResponse.streamCompletion?.finish() - streamTask.cancel() - return - } self.activeStreamTask?.cancel() self.activeStreamCompletion?.finish() self.activeStreamTask = streamTask self.activeStreamID = streamID self.activeStreamCompletion = trackedResponse.streamCompletion context.read() + registration.succeed(()) } - await streamTask.value - eventLoop.execute { - if self.activeStreamID == streamID { - self.activeStreamTask = nil - self.activeStreamID = nil - self.activeStreamCompletion = nil - } - } + try? await registration.futureResult.get() default: let body = response.bodyData diff --git a/Sources/CodexReviewTesting/TestSupport.swift b/Sources/CodexReviewTesting/TestSupport.swift index 373618da..bdc2618b 100644 --- a/Sources/CodexReviewTesting/TestSupport.swift +++ b/Sources/CodexReviewTesting/TestSupport.swift @@ -878,6 +878,7 @@ package final class TestingRuntimeLifecycleHandle: RuntimeLifecycleHandle { package private(set) var closeAdmissionCallCount = 0 package private(set) var closeCallCount = 0 package private(set) var waitUntilClosedCallCount = 0 + package private(set) var closePurposes: [ReviewRuntimeTransitionPurpose] = [] private let onActivate: @MainActor @Sendable () -> Void private let onClose: @MainActor @Sendable () -> Void @@ -900,7 +901,8 @@ package final class TestingRuntimeLifecycleHandle: RuntimeLifecycleHandle { closeAdmissionCallCount += 1 } - package func close(purpose _: ReviewRuntimeTransitionPurpose) async throws { + package func close(purpose: ReviewRuntimeTransitionPurpose) async throws { + closePurposes.append(purpose) closeCallCount += 1 guard didClose == false else { return @@ -919,6 +921,80 @@ package final class TestingRuntimeLifecycleHandle: RuntimeLifecycleHandle { } } +@MainActor +package final class TestingMCPServerLifecycleOwner: MCPServerLifecycleOwner { + package private(set) var prepareCallCount = 0 + package private(set) var activateCallCount = 0 + package private(set) var stopCallCount = 0 + package private(set) var waitUntilStoppedCallCount = 0 + package private(set) var preparedGenerations: [MCPServerGeneration] = [] + package private(set) var activatedGenerations: [MCPServerGeneration] = [] + + private let serverURL: URL? + private var nextGeneration: UInt64 = 0 + private var preparationGate: AsyncGate? + private let preparationStartedGate = AsyncGate() + private let preparationCancellationGate = AsyncGate() + + package init(serverURL: URL? = nil) { + self.serverURL = serverURL + } + + package func holdPreparation(with gate: AsyncGate) { + preparationGate = gate + } + + package func waitForPreparation() async { + await preparationStartedGate.wait() + } + + package func waitForPreparationCancellation() async { + await preparationCancellationGate.wait() + } + + package func prepare() async throws -> PreparedMCPServer { + prepareCallCount += 1 + nextGeneration &+= 1 + await preparationStartedGate.open() + if let preparationGate { + let cancellationGate = preparationCancellationGate + await withTaskCancellationHandler { + await preparationGate.waitIgnoringCancellation() + } onCancel: { + Task { await cancellationGate.open() } + } + self.preparationGate = nil + } + let generation = MCPServerGeneration(rawValue: nextGeneration) + preparedGenerations.append(generation) + return .init(generation: generation) + } + + package func activate( + _ generation: MCPServerGeneration + ) async throws -> MCPServerPublicationSnapshot { + activateCallCount += 1 + activatedGenerations.append(generation) + return .init(serverURL: serverURL) + } + + package func closeAdmission() async {} + + package func drainAdmittedHandlers() async throws {} + + package func stop() async throws { + stopCallCount += 1 + } + + package func waitUntilStopped() async throws { + waitUntilStoppedCallCount += 1 + } + + package func close() async throws {} + + package func waitUntilClosed() async throws {} +} + @MainActor package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { package let reviewBackend: FakeCodexReviewBackend @@ -928,19 +1004,21 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { package private(set) var startRequests: [Bool] = [] package private(set) var reviewStartOwnershipSnapshots: [StoreAttemptSnapshot?] = [] package private(set) var recoveryResumeOwnershipSnapshots: [StoreAttemptSnapshot?] = [] - package let mcpServerLifecycle: any MCPServerLifecycleOwner = NoMCPServerLifecycleOwner() + package let mcpServerLifecycle: any MCPServerLifecycleOwner package private(set) var lastPreparedRuntimeHandle: TestingRuntimeLifecycleHandle? private weak var store: CodexReviewStore? private var runtimePreparationGate: AsyncGate? - private let runtimePreparationStartedGate = AsyncGate() - private let runtimePreparationCancellationGate = AsyncGate() + private var runtimePreparationStartedGate = AsyncGate() + private var runtimePreparationCancellationGate = AsyncGate() package init( reviewBackend: FakeCodexReviewBackend, - seed: CodexReviewStoreSeed = .init() + seed: CodexReviewStoreSeed = .init(), + mcpServerLifecycle: any MCPServerLifecycleOwner = NoMCPServerLifecycleOwner() ) { self.reviewBackend = reviewBackend self.seed = seed + self.mcpServerLifecycle = mcpServerLifecycle self.currentSettingsSnapshot = seed.initialSettingsSnapshot } @@ -954,6 +1032,8 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { package func holdRuntimePreparation(with gate: AsyncGate) { runtimePreparationGate = gate + runtimePreparationStartedGate = AsyncGate() + runtimePreparationCancellationGate = AsyncGate() } package func waitForRuntimePreparation() async { diff --git a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift index 4388daea..6de5835b 100644 --- a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift +++ b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift @@ -286,6 +286,385 @@ struct CodexReviewHostTests { await store.stop() } + @Test func liveStoreRestartReplacesOnlyAppServerAndRetainsMCPListener() async throws { + let homeURL = try temporaryHome() + let firstTransport = FakeJSONRPCTransport() + let secondTransport = FakeJSONRPCTransport() + for transport in [firstTransport, secondTransport] { + 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" + ) + } + var transports = [firstTransport, secondTransport] + var mcpServerFactoryCallCount = 0 + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + mcpHTTPServerFactory: { store, configuration in + mcpServerFactoryCallCount += 1 + return CodexReviewMCPHTTPServer( + adapter: CodexReviewMCPServer(store: store), + configuration: .init( + host: configuration.host, + port: 0, + endpoint: configuration.endpoint + ) + ) + }, + mcpHTTPServerBindChecker: { _ in }, + transportFactory: { _ in transports.removeFirst() } + ) + + await store.start() + let initialURL = try #require(store.serverURL) + + await store.restart() + + #expect(store.serverState == .running) + #expect(store.serverURL == initialURL) + #expect(initialURL.port != 0) + #expect(mcpServerFactoryCallCount == 1) + #expect(transports.isEmpty) + + await store.stop() + } + + @Test func liveMCPOwnerStopDuringPreparationJoinsAndAllowsLaterStart() async throws { + let homeURL = try temporaryHome() + let preparationStarted = AsyncGate() + let preparationCancelled = AsyncGate() + let preparationRelease = AsyncGate() + let server = ControlledMCPHTTPServer( + endpoint: try #require(URL(string: "http://127.0.0.1:19418/mcp")) + ) + let lifecycleCalls = MCPLifecycleCallProbe() + var factoryCallCount = 0 + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + mcpHTTPServerFactory: { _, _ in + factoryCallCount += 1 + return server + }, + mcpHTTPServerBindChecker: { _ in + await preparationStarted.open() + await withTaskCancellationHandler { + await preparationRelease.waitIgnoringCancellation() + } onCancel: { + Task { await preparationCancelled.open() } + } + }, + mcpLifecycleCallObserver: { call, count in + lifecycleCalls.record(call, count: count) + }, + transportFactory: { _ in FakeJSONRPCTransport() } + ) + let owner = store.backend.mcpServerLifecycle + + let prepareTask = Task { try await owner.prepare() } + await preparationStarted.wait() + let firstStop = Task { try await owner.stop() } + await preparationCancelled.wait() + let secondStop = Task { try await owner.stop() } + await lifecycleCalls.waitFor(.stop, count: 2) + + #expect(factoryCallCount == 0) + await preparationRelease.open() + await #expect(throws: ReviewLifecycleResourceFailure.self) { + _ = try await prepareTask.value + } + try await firstStop.value + try await secondStop.value + try await owner.waitUntilStopped() + #expect(factoryCallCount == 0) + + let prepared = try await owner.prepare() + let repeatedPreparation = try await owner.prepare() + let snapshot = try await owner.activate(prepared.generation) + let repeatedSnapshot = try await owner.activate(prepared.generation) + let serverURL = await server.url + #expect(repeatedPreparation.generation == prepared.generation) + #expect(snapshot.serverURL == serverURL) + #expect(repeatedSnapshot.serverURL == snapshot.serverURL) + #expect(factoryCallCount == 1) + #expect(server.startCallCount == 1) + try await owner.stop() + try await owner.waitUntilStopped() + } + + @Test func liveStoreStopSignalsHeldMCPPreparationBeforeJoiningAcquisition() async throws { + let homeURL = try temporaryHome() + let preparationStarted = AsyncGate() + let preparationCancelled = AsyncGate() + let preparationRelease = AsyncGate() + var mcpFactoryCallCount = 0 + var appServerFactoryCallCount = 0 + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + mcpHTTPServerFactory: { _, configuration in + mcpFactoryCallCount += 1 + return NoopMCPHTTPServer(endpoint: configuration.url()) + }, + mcpHTTPServerBindChecker: { _ in + await preparationStarted.open() + await withTaskCancellationHandler { + await preparationRelease.waitIgnoringCancellation() + } onCancel: { + Task { await preparationCancelled.open() } + } + }, + transportFactory: { _ in + appServerFactoryCallCount += 1 + return FakeJSONRPCTransport() + } + ) + + let startTask = Task { @MainActor in + await store.start() + } + await preparationStarted.wait() + let stopFinished = CompletionFlag() + let stopTask = Task { @MainActor in + await store.stop() + await stopFinished.complete() + } + await preparationCancelled.wait() + + #expect(await stopFinished.isCompleted() == false) + #expect(mcpFactoryCallCount == 0) + #expect(appServerFactoryCallCount == 0) + + await preparationRelease.open() + await stopTask.value + await startTask.value + + #expect(await stopFinished.isCompleted()) + #expect(store.serverState == .stopped) + #expect(store.serverURL == nil) + #expect(mcpFactoryCallCount == 0) + #expect(appServerFactoryCallCount == 0) + } + + @Test func liveStoreStopSignalsHeldMCPActivationBeforeJoiningAcquisition() 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 activationRelease = AsyncGate() + let server = ControlledMCPHTTPServer( + endpoint: try #require(URL(string: "http://127.0.0.1:19424/mcp")) + ) + server.holdStart(with: activationRelease) + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + mcpHTTPServerFactory: { _, _ in server }, + mcpHTTPServerBindChecker: { _ in }, + transportFactory: { _ in transport } + ) + + let startTask = Task { @MainActor in + await store.start() + } + await server.waitForStart() + let stopFinished = CompletionFlag() + let stopTask = Task { @MainActor in + await store.stop() + await stopFinished.complete() + } + await server.waitForStartCancellation() + + #expect(await stopFinished.isCompleted() == false) + #expect(store.serverURL == nil) + #expect(server.stopCallCount == 0) + + await activationRelease.open() + await stopTask.value + await startTask.value + + #expect(await stopFinished.isCompleted()) + #expect(store.serverState == .stopped) + #expect(store.serverURL == nil) + #expect(server.stopCallCount == 1) + } + + @Test func liveMCPOwnerJoinerCancellationDoesNotCancelSharedOperations() async throws { + let homeURL = try temporaryHome() + let preparationStarted = AsyncGate() + let preparationRelease = AsyncGate() + let activationRelease = AsyncGate() + let server = ControlledMCPHTTPServer( + endpoint: try #require(URL(string: "http://127.0.0.1:19423/mcp")) + ) + server.holdStart(with: activationRelease) + var bindCheckCallCount = 0 + var factoryCallCount = 0 + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + mcpHTTPServerFactory: { _, _ in + factoryCallCount += 1 + return server + }, + mcpHTTPServerBindChecker: { _ in + bindCheckCallCount += 1 + await preparationStarted.open() + await preparationRelease.waitIgnoringCancellation() + }, + transportFactory: { _ in FakeJSONRPCTransport() } + ) + let owner = store.backend.mcpServerLifecycle + + let firstPreparation = Task { try await owner.prepare() } + await preparationStarted.wait() + let secondPreparation = Task { try await owner.prepare() } + firstPreparation.cancel() + await preparationRelease.open() + let firstPrepared = try await firstPreparation.value + let secondPrepared = try await secondPreparation.value + + #expect(firstPrepared.generation == secondPrepared.generation) + #expect(bindCheckCallCount == 1) + #expect(factoryCallCount == 1) + + let firstActivation = Task { + try await owner.activate(firstPrepared.generation) + } + await server.waitForStart() + let secondActivation = Task { + try await owner.activate(firstPrepared.generation) + } + firstActivation.cancel() + await activationRelease.open() + let firstSnapshot = try await firstActivation.value + let secondSnapshot = try await secondActivation.value + + #expect(firstSnapshot.serverURL == secondSnapshot.serverURL) + #expect(server.startCallCount == 1) + try await owner.stop() + } + + @Test func liveMCPOwnerStopDuringActivationJoinsAndStartsNewLeaseLater() async throws { + let homeURL = try temporaryHome() + let firstServer = ControlledMCPHTTPServer( + endpoint: try #require(URL(string: "http://127.0.0.1:19419/mcp")) + ) + let secondServer = ControlledMCPHTTPServer( + endpoint: try #require(URL(string: "http://127.0.0.1:19420/mcp")) + ) + let lifecycleCalls = MCPLifecycleCallProbe() + let activationRelease = AsyncGate() + firstServer.holdStart(with: activationRelease) + var servers = [firstServer, secondServer] + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + mcpHTTPServerFactory: { _, _ in servers.removeFirst() }, + mcpHTTPServerBindChecker: { _ in }, + mcpLifecycleCallObserver: { call, count in + lifecycleCalls.record(call, count: count) + }, + transportFactory: { _ in FakeJSONRPCTransport() } + ) + let owner = store.backend.mcpServerLifecycle + let prepared = try await owner.prepare() + let activationTask = Task { + try await owner.activate(prepared.generation) + } + await firstServer.waitForStart() + + let firstStop = Task { try await owner.stop() } + await firstServer.waitForStartCancellation() + let secondStop = Task { try await owner.stop() } + await lifecycleCalls.waitFor(.stop, count: 2) + #expect(firstServer.stopCallCount == 0) + + await activationRelease.open() + await #expect(throws: ReviewLifecycleResourceFailure.self) { + _ = try await activationTask.value + } + try await firstStop.value + try await secondStop.value + try await owner.waitUntilStopped() + #expect(firstServer.stopCallCount == 1) + + let nextPrepared = try await owner.prepare() + let nextSnapshot = try await owner.activate(nextPrepared.generation) + let repeatedSnapshot = try await owner.activate(nextPrepared.generation) + let secondServerURL = await secondServer.url + #expect(nextPrepared.generation != prepared.generation) + #expect(nextSnapshot.serverURL == secondServerURL) + #expect(repeatedSnapshot.serverURL == nextSnapshot.serverURL) + #expect(secondServer.startCallCount == 1) + #expect(servers.isEmpty) + try await owner.stop() + } + + @Test func liveMCPOwnerCloseDuringActivationJoinsAndPreventsLatePublication() async throws { + let homeURL = try temporaryHome() + let server = ControlledMCPHTTPServer( + endpoint: try #require(URL(string: "http://127.0.0.1:19421/mcp")) + ) + let lifecycleCalls = MCPLifecycleCallProbe() + let activationRelease = AsyncGate() + server.holdStart(with: activationRelease) + var factoryCallCount = 0 + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + mcpHTTPServerFactory: { _, _ in + factoryCallCount += 1 + return server + }, + mcpHTTPServerBindChecker: { _ in }, + mcpLifecycleCallObserver: { call, count in + lifecycleCalls.record(call, count: count) + }, + transportFactory: { _ in FakeJSONRPCTransport() } + ) + let owner = store.backend.mcpServerLifecycle + let prepared = try await owner.prepare() + let activationTask = Task { + try await owner.activate(prepared.generation) + } + await server.waitForStart() + + let firstClose = Task { try await owner.close() } + await server.waitForStartCancellation() + let secondClose = Task { try await owner.close() } + await lifecycleCalls.waitFor(.close, count: 2) + #expect(server.stopCallCount == 0) + + await activationRelease.open() + await #expect(throws: ReviewLifecycleResourceFailure.self) { + _ = try await activationTask.value + } + try await firstClose.value + try await secondClose.value + try await owner.waitUntilClosed() + #expect(server.stopCallCount == 1) + #expect(factoryCallCount == 1) + await #expect(throws: ReviewLifecycleResourceFailure.self) { + _ = try await owner.prepare() + } + } + @Test func liveStoreReportsMCPPortOwnerWhenEndpointPortInUseAndDoesNotLaunchAppServer() async throws { let homeURL = try temporaryHome() let port = 54321 @@ -2202,6 +2581,108 @@ private func failedMessage(from phase: CodexReviewAuthModel.Phase) -> String? { return message } +@MainActor +private final class MCPLifecycleCallProbe { + private struct Waiter { + let call: CodexReviewMCPLifecycleCall + let count: Int + let continuation: CheckedContinuation + } + + private var counts: [CodexReviewMCPLifecycleCall: Int] = [:] + private var waiters: [UUID: Waiter] = [:] + + func record(_ call: CodexReviewMCPLifecycleCall, count: Int) { + counts[call] = count + let completedWaiterIDs = waiters.compactMap { id, waiter in + waiter.call == call && count >= waiter.count ? id : nil + } + for waiterID in completedWaiterIDs { + waiters.removeValue(forKey: waiterID)?.continuation.resume() + } + } + + func waitFor(_ call: CodexReviewMCPLifecycleCall, count: Int) async { + if counts[call, default: 0] >= count { + return + } + await withCheckedContinuation { continuation in + if counts[call, default: 0] >= count { + continuation.resume() + } else { + waiters[UUID()] = .init( + call: call, + count: count, + continuation: continuation + ) + } + } + } +} + +@MainActor +private final class ControlledMCPHTTPServer: CodexReviewMCPHTTPServing { + private let endpoint: URL + private var startGate: AsyncGate? + private var startStartedGate = AsyncGate() + private var startCancellationGate = AsyncGate() + + private(set) var startCallCount = 0 + private(set) var closeAdmissionCallCount = 0 + private(set) var waitForAdmittedHandlersCallCount = 0 + private(set) var stopCallCount = 0 + + init(endpoint: URL) { + self.endpoint = endpoint + } + + var url: URL { + get async { + endpoint + } + } + + func holdStart(with gate: AsyncGate) { + startGate = gate + startStartedGate = AsyncGate() + startCancellationGate = AsyncGate() + } + + func waitForStart() async { + await startStartedGate.wait() + } + + func waitForStartCancellation() async { + await startCancellationGate.wait() + } + + func start() async throws { + startCallCount += 1 + await startStartedGate.open() + if let startGate { + let cancellationGate = startCancellationGate + await withTaskCancellationHandler { + await startGate.waitIgnoringCancellation() + } onCancel: { + Task { await cancellationGate.open() } + } + self.startGate = nil + } + } + + func closeAdmission() async { + closeAdmissionCallCount += 1 + } + + func waitForAdmittedHandlers() async { + waitForAdmittedHandlersCallCount += 1 + } + + func stop() async { + stopCallCount += 1 + } +} + private final class NoopMCPHTTPServer: CodexReviewMCPHTTPServing, @unchecked Sendable { private let endpoint: URL diff --git a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift index c6a35eb2..77564a0b 100644 --- a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift +++ b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift @@ -102,8 +102,10 @@ struct CodexReviewMCPHTTPServerTests { ) let server = CodexReviewMCPHTTPServer( adapter: CodexReviewMCPServer(store: store), - configuration: .init(host: "review.local", port: 9417) + configuration: .init(host: "127.0.0.1", port: 0) ) + try await server.start() + let boundPort = try #require(await server.url.port) let initializeBody = try makeJSONBody([ "jsonrpc": "2.0", "id": 1, @@ -120,7 +122,7 @@ struct CodexReviewMCPHTTPServerTests { let response = await server.handleHTTPRequest(HTTPRequest( method: "POST", headers: [ - HTTPHeaderName.host: "review.local:9417", + HTTPHeaderName.host: "127.0.0.1:\(boundPort)", HTTPHeaderName.accept: "text/event-stream, application/json", HTTPHeaderName.contentType: "application/json", ], @@ -130,7 +132,7 @@ struct CodexReviewMCPHTTPServerTests { let denied = await server.handleHTTPRequest(HTTPRequest( method: "POST", headers: [ - HTTPHeaderName.host: "other.local:9417", + HTTPHeaderName.host: "other.local:\(boundPort)", HTTPHeaderName.accept: "text/event-stream, application/json", HTTPHeaderName.contentType: "application/json", ], @@ -141,6 +143,18 @@ struct CodexReviewMCPHTTPServerTests { #expect(response.statusCode == 200) #expect(response.headers[HTTPHeaderName.sessionID]?.isEmpty == false) #expect(denied.statusCode == 421) + await server.closeAdmission() + let afterAdmissionClose = await server.handleHTTPRequest(HTTPRequest( + method: "POST", + headers: [ + HTTPHeaderName.host: "127.0.0.1:\(boundPort)", + HTTPHeaderName.accept: "text/event-stream, application/json", + HTTPHeaderName.contentType: "application/json", + ], + body: initializeBody, + path: "/mcp" + )) + #expect(afterAdmissionClose.statusCode == 503) await server.stop() } @@ -1059,6 +1073,67 @@ struct CodexReviewMCPHTTPServerTests { } } + @Test func stopDrainsAdmittedHandlerBeforeClosingItsSession() async throws { + let backend = FakeCodexReviewBackend() + let requestGate = AsyncGate() + await backend.holdStartReview(with: requestGate) + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-1" }) + ) + let server = CodexReviewMCPHTTPServer( + adapter: CodexReviewMCPServer(store: store), + configuration: .init(port: 0) + ) + try await server.start() + let endpoint = await server.url + let sessionID = try await initializeSession(endpoint: endpoint) + let priorAdmissionCount = await server.admittedNetworkRequestCountForTesting() + await server.holdNextNetworkHandlerEntryForTesting() + let requestBody = try makeJSONBody([ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": [ + "name": "review_start", + "arguments": [ + "cwd": "/tmp/project", + "target": ["type": "uncommittedChanges"], + ], + ], + ]) + let requestTask = Task { + try await postJSONRPCData( + endpoint: endpoint, + sessionID: sessionID, + bodyData: requestBody + ) + } + await server.waitForAdmittedNetworkRequestCountForTesting( + priorAdmissionCount + 1 + ) + + let stopFinished = CompletionFlag() + let stopTask = Task { + await server.stop() + await stopFinished.complete() + } + await server.waitForAdmittedHandlerDrainToBeginForTesting() + + #expect(await server.listenerIsOpenForTesting() == false) + #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) != nil) + #expect(await stopFinished.isCompleted() == false) + + await server.releaseNetworkHandlerEntryForTesting() + await requestGate.open() + await backend.yield(.completed(summary: "Done", result: "review text")) + _ = try? await requestTask.value + await stopTask.value + + #expect(await stopFinished.isCompleted()) + #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) == nil) + } + @Test func streamableHTTPDoesNotExpireSessionWithOpenEventStream() async throws { let backend = FakeCodexReviewBackend() let store = CodexReviewStore.makeTestingStore( @@ -1509,6 +1584,18 @@ struct CodexReviewMCPHTTPServerTests { } } +private actor CompletionFlag { + private var completed = false + + func complete() { + completed = true + } + + func isCompleted() -> Bool { + completed + } +} + private nonisolated func currentPOSIXError() -> NSError { NSError(domain: NSPOSIXErrorDomain, code: Int(errno)) } diff --git a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift index 23838c2b..399c98fa 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift @@ -1,3 +1,4 @@ +import Foundation import Testing import CodexReview import CodexReviewTesting @@ -5,6 +6,32 @@ import CodexReviewTesting @Suite("store runtime lifecycle") @MainActor struct CodexReviewStoreLifecycleTests { + @Test func invalidatedGenerationBeforeAcquisitionEntryPerformsNoMCPWork() async { + let mcpOwner = TestingMCPServerLifecycleOwner() + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend(), + mcpServerLifecycle: mcpOwner + ) + let store = CodexReviewStore.makeTestingStore(backend: backend) + let generation = ReviewRuntimeGeneration(rawValue: 1) + let transitionTask = Task { @MainActor in } + store.runtimeState = .transitioning( + generation: generation.successor(), + purpose: .stop, + task: transitionTask + ) + + await store.performRuntimeAcquisitionForTesting( + generation: generation, + purpose: .stop + ) + + #expect(mcpOwner.prepareCallCount == 0) + #expect(mcpOwner.activateCallCount == 0) + #expect(backend.lastPreparedRuntimeHandle == nil) + #expect(backend.startRequests.isEmpty) + } + @Test func stopInvalidatesHeldRuntimePreparationAndClosesStaleHandleOnce() async throws { let preparationGate = AsyncGate() let backend = TestingCodexReviewStoreBackend( @@ -65,4 +92,138 @@ struct CodexReviewStoreLifecycleTests { #expect(store.serverState == .running) #expect(store.serverURL == nil) } + + @Test func staleMCPPreparationDoesNotAcquireAnAppServerRuntime() async throws { + let preparationGate = AsyncGate() + let mcpOwner = TestingMCPServerLifecycleOwner() + mcpOwner.holdPreparation(with: preparationGate) + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend(), + mcpServerLifecycle: mcpOwner + ) + let store = CodexReviewStore.makeTestingStore(backend: backend) + + let startTask = Task { @MainActor in + await store.start() + } + await mcpOwner.waitForPreparation() + let stopTask = Task { @MainActor in + await store.stop() + } + await mcpOwner.waitForPreparationCancellation() + + #expect(backend.lastPreparedRuntimeHandle == nil) + #expect(backend.startRequests.isEmpty) + + await preparationGate.open() + await stopTask.value + await startTask.value + + #expect(store.serverState == .stopped) + #expect(backend.lastPreparedRuntimeHandle == nil) + #expect(mcpOwner.stopCallCount == 1) + #expect(mcpOwner.waitUntilStoppedCallCount == 1) + #expect(mcpOwner.activateCallCount == 0) + } + + @Test func sameAccountRestartRetainsOneMCPGenerationAndReplacesOnlyAppServer() 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 firstMCPGeneration = try #require(mcpOwner.preparedGenerations.first) + #expect(store.serverURL == endpoint) + + let replacementGate = AsyncGate() + backend.holdRuntimePreparation(with: replacementGate) + let restartTask = Task { @MainActor in + await store.restart() + } + await backend.waitForRuntimePreparation() + + let ownsReplacementTask: Bool + if case .transitioning(_, .restartSameAccount, _) = store.runtimeState { + ownsReplacementTask = true + } else { + ownsReplacementTask = false + } + #expect(ownsReplacementTask) + #expect(store.serverURL == endpoint) + #expect(mcpOwner.prepareCallCount == 1) + #expect(mcpOwner.activateCallCount == 1) + #expect(mcpOwner.stopCallCount == 0) + #expect(mcpOwner.waitUntilStoppedCallCount == 0) + #expect(firstHandle.closePurposes == [.restartSameAccount]) + + await replacementGate.open() + await restartTask.value + + let secondHandle = try #require(backend.lastPreparedRuntimeHandle) + #expect(secondHandle !== firstHandle) + #expect(secondHandle.activateCallCount == 1) + #expect(backend.startRequests == [false, true]) + #expect(store.serverState == .running) + #expect(store.serverURL == endpoint) + #expect(mcpOwner.preparedGenerations == [firstMCPGeneration]) + #expect(mcpOwner.activatedGenerations == [firstMCPGeneration]) + #expect(mcpOwner.stopCallCount == 0) + + guard case .running(_, _, let retainedMCPGeneration) = store.runtimeState else { + Issue.record("Replacement must publish the new AppServer runtime.") + return + } + #expect(retainedMCPGeneration == firstMCPGeneration) + + await store.stop() + } + + @Test func stopInvalidatesHeldRestartBeforeReplacementCanPublish() 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 replacementGate = AsyncGate() + backend.holdRuntimePreparation(with: replacementGate) + let restartTask = Task { @MainActor in + await store.restart() + } + await backend.waitForRuntimePreparation() + let staleReplacement = try #require(backend.lastPreparedRuntimeHandle) + let stopTask = Task { @MainActor in + await store.stop() + } + await backend.waitForRuntimePreparationCancellation() + + #expect(firstHandle.closePurposes == [.restartSameAccount]) + #expect(staleReplacement !== firstHandle) + #expect(staleReplacement.activateCallCount == 0) + #expect(store.serverURL == endpoint) + + await replacementGate.open() + await stopTask.value + await restartTask.value + + #expect(store.serverState == .stopped) + #expect(store.serverURL == nil) + #expect(staleReplacement.activateCallCount == 0) + #expect(staleReplacement.closeAdmissionCallCount == 1) + #expect(staleReplacement.closeCallCount == 1) + #expect(staleReplacement.waitUntilClosedCallCount == 1) + #expect(mcpOwner.prepareCallCount == 1) + #expect(mcpOwner.activateCallCount == 1) + #expect(mcpOwner.stopCallCount == 1) + #expect(mcpOwner.waitUntilStoppedCallCount == 1) + } } From b66d78868859a27bdb4058be18eca689808ea9ff Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 05:23:57 +0900 Subject: [PATCH 26/38] fix(app-server): join admitted lifecycle operations --- .../AppServerCodexReviewBackend.swift | 329 +++++++++- .../AppServerClientTests.swift | 600 ++++++++++++++++++ 2 files changed, 914 insertions(+), 15 deletions(-) diff --git a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift index b149f859..584ce098 100644 --- a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift +++ b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift @@ -15,13 +15,154 @@ private func makeAppServerReviewAttemptID() -> String { UUID().uuidString } +package struct AppServerRuntimeOwnerLifecycleHandle: Sendable { + private let closeAndWaitOperation: @Sendable () async throws -> Void + + fileprivate init( + closeAndWaitOperation: @escaping @Sendable () async throws -> Void + ) { + self.closeAndWaitOperation = closeAndWaitOperation + } + + package func closeAndWait() async throws { + try await closeAndWaitOperation() + } +} + package actor AppServerCodexReviewBackend: CodexReviewBackend { + private struct AdmittedReviewOperationID: Hashable { + let rawValue: Int + } + + private struct ReviewOperationRegistry { + enum Admission { + case open + case closed + } + + var admission: Admission = .open + var nextID = 0 + var admitted: Set = [] + var drainWaiters: [CheckedContinuation] = [] + + mutating func register() -> AdmittedReviewOperationID? { + guard case .open = admission else { + return nil + } + let id = AdmittedReviewOperationID(rawValue: nextID) + nextID += 1 + admitted.insert(id) + return id + } + + mutating func closeAdmission() { + admission = .closed + resumeDrainWaitersIfNeeded() + } + + mutating func finish(_ id: AdmittedReviewOperationID) { + admitted.remove(id) + resumeDrainWaitersIfNeeded() + } + + func contains(_ id: AdmittedReviewOperationID) -> Bool { + admitted.contains(id) + } + + mutating func appendDrainWaiter( + _ continuation: CheckedContinuation + ) { + if admitted.isEmpty { + continuation.resume() + } else { + drainWaiters.append(continuation) + } + } + + private mutating func resumeDrainWaitersIfNeeded() { + guard case .closed = admission, admitted.isEmpty else { + return + } + let waiters = drainWaiters + drainWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume() + } + } + } + + private struct OwnedLifecycleSnapshot: Sendable { + let routerStartTask: Task? + let routerTask: Task? + let sessions: [AppServerReviewEventSession] + } + + // Test-only acknowledgements emitted by the lifecycle owner. They do not + // participate in admission, close ordering, or result assembly. + private struct LifecycleTestingObservation { + var closeCallerCount = 0 + var closeCallerWaiters: [(Int, CheckedContinuation)] = [] + var clientCloseResultRecorded = false + var clientCloseResultWaiters: [CheckedContinuation] = [] + + mutating func recordCloseCaller() { + closeCallerCount += 1 + let count = closeCallerCount + let ready = closeCallerWaiters.filter { count >= $0.0 } + closeCallerWaiters.removeAll { count >= $0.0 } + for (_, waiter) in ready { + waiter.resume() + } + } + + mutating func appendCloseCallerWaiter( + count: Int, + continuation: CheckedContinuation + ) { + if closeCallerCount >= count { + continuation.resume() + } else { + closeCallerWaiters.append((count, continuation)) + } + } + + mutating func recordClientCloseResult() { + guard clientCloseResultRecorded == false else { + return + } + clientCloseResultRecorded = true + let waiters = clientCloseResultWaiters + clientCloseResultWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume() + } + } + + mutating func appendClientCloseResultWaiter( + _ continuation: CheckedContinuation + ) { + if clientCloseResultRecorded { + continuation.resume() + } else { + clientCloseResultWaiters.append(continuation) + } + } + } + + private enum LifecycleState { + case open + case closing(Task) + case closed(Result) + } + private static let reviewPermissionProfileID = ":danger-full-access" private let client: AppServerClient private let threadStartPermissionStrategy: AppServerAPI.Thread.Start.PermissionStrategy private var controlsByThreadID: [String: AppServerReviewControl] = [:] private var reviewEventSessionsByAttemptID: [String: AppServerReviewEventSession] = [:] + private var reviewEventSessionRegistrationOrdinalByAttemptID: [String: Int] = [:] + private var nextReviewEventSessionRegistrationOrdinal = 0 private var activeReviewAttemptIDsByThreadID: [String: Set] = [:] private var activeThreadIDsByAttemptID: [String: Set] = [:] private var reviewEventSessionCanonicalThreadIDByThreadID: [String: String] = [:] @@ -30,8 +171,11 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { private var abandonedTurnIDs: Set = [] private var unmatchedReviewNotificationsByThreadID: [String: [AppServerRoutedReviewNotification]] = [:] private var completedReviewEventSessionMetricsByThreadID: [String: AppServerReviewEventSessionMetrics] = [:] + private var lifecycleState: LifecycleState = .open + private var lifecycleTestingObservation = LifecycleTestingObservation() + private var reviewOperationRegistry = ReviewOperationRegistry() + private var notificationRouterStartTask: Task? private var notificationRouterTask: Task? - private var isNotificationRouterStarting = false private var reviewNotificationSequence = 0 private var notificationRouterMetrics = AppServerNotificationRouterMetrics() private var reviewStartRequestsInFlight = 0 @@ -46,6 +190,46 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { self.threadStartPermissionStrategy = threadStartPermissionStrategy } + package nonisolated var runtimeOwnerLifecycleHandle: AppServerRuntimeOwnerLifecycleHandle { + AppServerRuntimeOwnerLifecycleHandle { [self] in + try await closeFromRuntimeOwnerAndWait() + } + } + + private func admitReviewOperation() throws -> AdmittedReviewOperationID { + guard let id = reviewOperationRegistry.register() else { + throw JSONRPC.Error.closed + } + return id + } + + private func finishReviewOperation(_ id: AdmittedReviewOperationID) { + reviewOperationRegistry.finish(id) + } + + private func waitForAdmittedReviewOperations() async { + if reviewOperationRegistry.admitted.isEmpty { + return + } + await withCheckedContinuation { continuation in + reviewOperationRegistry.appendDrainWaiter(continuation) + } + } + + private func ownedLifecycleSnapshot() -> OwnedLifecycleSnapshot { + let sessions = reviewEventSessionsByAttemptID + .sorted { + reviewEventSessionRegistrationOrdinalByAttemptID[$0.key, default: .max] + < reviewEventSessionRegistrationOrdinalByAttemptID[$1.key, default: .max] + } + .map(\.value) + return OwnedLifecycleSnapshot( + routerStartTask: notificationRouterStartTask, + routerTask: notificationRouterTask, + sessions: sessions + ) + } + package func readSettings() async throws -> CodexReviewBackendModel.Settings.Snapshot { _ = try await client.initialize() let response = try await client.send(AppServerAPI.Config.Read.Request()) @@ -131,8 +315,11 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { _ request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { + let operationID = try admitReviewOperation() + defer { finishReviewOperation(operationID) } + _ = try await client.initialize() - await ensureNotificationRouterStarted() + await ensureNotificationRouterStarted(for: operationID) let control = AppServerReviewControl(client: client) let thread = try await startReviewThread(request, admission: admission) @@ -351,11 +538,14 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } package func interruptReview(_ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason) async throws { + let operationID = try admitReviewOperation() + defer { finishReviewOperation(operationID) } + _ = try await client.initialize() guard abandonedReviewAttemptIDs.contains(run.attemptID) == false else { return } - let session = await reviewEventSession(for: run) + let session = await reviewEventSession(for: run, admittedBy: operationID) await session.requestCancellation(message: reason.message) do { _ = try await sendTurnInterrupt(for: run) @@ -379,6 +569,50 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } } + // Only AppServerRuntimeOwnerLifecycleHandle can enter this transition. The + // notification router never owns that handle, so close cannot await itself. + private func closeFromRuntimeOwnerAndWait() async throws { + lifecycleTestingObservation.recordCloseCaller() + let closeTask: Task + switch lifecycleState { + case .open: + reviewOperationRegistry.closeAdmission() + let client = client + let task = Task { + let clientCloseTask = Task, Never> { + do { + try await client.close() + return .success(()) + } catch { + return .failure(error) + } + } + + await self.waitForAdmittedReviewOperations() + let clientCloseResult = await clientCloseTask.value + let ownedLifecycle = self.ownedLifecycleSnapshot() + self.lifecycleTestingObservation.recordClientCloseResult() + await ownedLifecycle.routerStartTask?.value + await ownedLifecycle.routerTask?.value + for session in ownedLifecycle.sessions { + await session.finish(throwing: .ownerCancellation) + } + try clientCloseResult.get() + } + lifecycleState = .closing(task) + closeTask = task + case .closing(let task): + closeTask = task + case .closed(let result): + try result.get() + return + } + + let result = await closeTask.result + lifecycleState = .closed(result) + try result.get() + } + package static func reviewRuntimeCloseFailure( for error: any Error ) -> ReviewRuntimeCloseFailure { @@ -432,8 +666,11 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { + let operationID = try admitReviewOperation() + defer { finishReviewOperation(operationID) } + _ = try await client.initialize() - await ensureNotificationRouterStarted() + await ensureNotificationRouterStarted(for: operationID) let token = handoff.token let interruptedRun = token.interruptedRun try await admission.admitRecoveryRollbackDispatch(threadID: token.rollbackThreadID) @@ -585,6 +822,23 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { notificationRouterTask != nil } + // These continuation-backed observations acknowledge owner transitions; + // tests never infer them from Task scheduling or resource-side callbacks. + package func waitForRuntimeOwnerCloseCallersForTesting(_ count: Int) async { + await withCheckedContinuation { continuation in + lifecycleTestingObservation.appendCloseCallerWaiter( + count: count, + continuation: continuation + ) + } + } + + package func waitForClientCloseResultBeforeRouterWaitForTesting() async { + await withCheckedContinuation { continuation in + lifecycleTestingObservation.appendClientCloseResultWaiter(continuation) + } + } + package func detachReviewEventStreamForTesting(threadID: String, subscriptionID: Int) async { guard let session = reviewEventSession(forThreadID: threadID) else { return @@ -593,17 +847,24 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } package func reviewAttemptForTesting(_ run: CodexReviewBackendModel.Review.Run) async -> BackendReviewAttempt { - let session = await reviewEventSession(for: run) + guard let operationID = reviewOperationRegistry.register() else { + preconditionFailure("Review event-session admission is closed.") + } + defer { finishReviewOperation(operationID) } + let session = await reviewEventSession(for: run, admittedBy: operationID) return await session.attempt() } - private func reviewEventSession(for run: CodexReviewBackendModel.Review.Run) async -> AppServerReviewEventSession { - await ensureNotificationRouterStarted() + private func reviewEventSession( + for run: CodexReviewBackendModel.Review.Run, + admittedBy operationID: AdmittedReviewOperationID + ) async -> AppServerReviewEventSession { if let session = reviewEventSessionsByAttemptID[run.attemptID] { await session.updateRun(run) registerReviewEventSession(session, for: run) return session } + await ensureNotificationRouterStarted(for: operationID) let control = controlsByThreadID[run.threadID] ?? AppServerReviewControl(client: client) controlsByThreadID[run.threadID] = control if let turnID = run.turnID { @@ -620,6 +881,10 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { _ session: AppServerReviewEventSession, for run: CodexReviewBackendModel.Review.Run ) { + if reviewEventSessionRegistrationOrdinalByAttemptID[run.attemptID] == nil { + reviewEventSessionRegistrationOrdinalByAttemptID[run.attemptID] = nextReviewEventSessionRegistrationOrdinal + nextReviewEventSessionRegistrationOrdinal += 1 + } reviewEventSessionsByAttemptID[run.attemptID] = session let activeThreadIDs = Set([run.threadID, run.reviewThreadID].compactMap { $0?.nilIfEmpty }) for threadID in activeThreadIDsByAttemptID[run.attemptID] ?? [] where activeThreadIDs.contains(threadID) == false { @@ -660,6 +925,7 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } } activeThreadIDsByAttemptID.removeValue(forKey: run.attemptID) + reviewEventSessionRegistrationOrdinalByAttemptID.removeValue(forKey: run.attemptID) return reviewEventSessionsByAttemptID.removeValue(forKey: run.attemptID) } @@ -754,22 +1020,55 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { return merged } - private func ensureNotificationRouterStarted() async { + private func ensureNotificationRouterStarted( + for operationID: AdmittedReviewOperationID + ) async { + precondition( + reviewOperationRegistry.contains(operationID), + "Review operation registry must own router-start admission." + ) if notificationRouterTask != nil { return } - while isNotificationRouterStarting { - await Task.yield() - if notificationRouterTask != nil { - return + + let startTask: Task + if let notificationRouterStartTask { + startTask = notificationRouterStartTask + } else { + let client = client + let task = Task { [client] in + let notifications = await client.notificationStream() + self.installNotificationRouter( + notifications, + admittedBy: operationID + ) } + notificationRouterStartTask = task + startTask = task + } + await startTask.value + precondition( + reviewOperationRegistry.contains(operationID) + && notificationRouterTask != nil, + "Review operation registry must retain admission through router start." + ) + } + + private func installNotificationRouter( + _ notifications: AsyncThrowingStream, + admittedBy operationID: AdmittedReviewOperationID + ) { + notificationRouterStartTask = nil + guard notificationRouterTask == nil else { + return } - isNotificationRouterStarting = true - let notifications = await client.notificationStream() + precondition( + reviewOperationRegistry.contains(operationID), + "Review operation registry must own router installation." + ) notificationRouterTask = Task { [notifications] in await self.consumeReviewNotifications(notifications) } - isNotificationRouterStarting = false } private func consumeReviewNotifications( diff --git a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift index 0b71e5e2..654e7404 100644 --- a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift +++ b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift @@ -199,6 +199,162 @@ private func controlledReviewClosePolicy(gate: AsyncGate) -> ReviewRuntimeCloseP } } +private actor RequestBarrier { + private let releaseGate = AsyncGate() + private var entered = false + private var entryWaiters: [CheckedContinuation] = [] + + func enterAndWait() async { + entered = true + let waiters = entryWaiters + entryWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume() + } + await releaseGate.waitIgnoringCancellation() + } + + func waitUntilEntered() async { + if entered { + return + } + await withCheckedContinuation { continuation in + if entered { + continuation.resume() + } else { + entryWaiters.append(continuation) + } + } + } + + func open() async { + await releaseGate.open() + } +} + +private actor DeferredNotificationCloseTransport: JSONRPC.Transport { + private let closeFailure: ReviewRuntimeCloseFailure? + private var responses: [String: [Data]] = [:] + private var requestBarriersByMethod: [String: [RequestBarrier]] = [:] + private var notificationStreamGate: AsyncGate? + private var notificationStreamRequested = false + private var notificationStreamWaiters: [CheckedContinuation] = [] + private var notificationContinuation: AsyncThrowingStream.Continuation? + private var closeCallCount = 0 + private var sendCallCount = 0 + + init(closeFailure: ReviewRuntimeCloseFailure? = nil) { + self.closeFailure = closeFailure + } + + func enqueue( + _ response: Response, + for method: String + ) throws { + responses[method, default: []].append(try JSONEncoder().encode(response)) + } + + func holdNext(method: String, barrier: RequestBarrier) { + requestBarriersByMethod[method, default: []].append(barrier) + } + + func holdNotificationStream(on gate: AsyncGate) { + notificationStreamGate = gate + } + + func send(_ request: JSONRPC.Request) async throws -> Data { + sendCallCount += 1 + if var barriers = requestBarriersByMethod[request.method], + barriers.isEmpty == false + { + let barrier = barriers.removeFirst() + requestBarriersByMethod[request.method] = barriers + await barrier.enterAndWait() + } + guard var methodResponses = responses[request.method], + methodResponses.isEmpty == false + else { + return try JSONEncoder().encode(EmptyResponse()) + } + let response = methodResponses.removeFirst() + responses[request.method] = methodResponses + return response + } + + func notify(_: JSONRPC.Notification) async throws {} + + func notificationStream() async -> AsyncThrowingStream { + notificationStreamRequested = true + let waiters = notificationStreamWaiters + notificationStreamWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume() + } + if let notificationStreamGate { + await notificationStreamGate.waitIgnoringCancellation() + } + return AsyncThrowingStream(bufferingPolicy: .unbounded) { continuation in + notificationContinuation = continuation + } + } + + func close() async throws { + closeCallCount += 1 + if let closeFailure { + throw closeFailure + } + } + + func waitForNotificationStreamRequest() async { + if notificationStreamRequested { + return + } + await withCheckedContinuation { continuation in + if notificationStreamRequested { + continuation.resume() + } else { + notificationStreamWaiters.append(continuation) + } + } + } + + func emitServerNotification( + method: String, + params: Params + ) throws { + let notification = JSONRPC.Notification( + method: method, + params: try JSONEncoder().encode(params) + ) + notificationContinuation?.yield(notification) + } + + func finishNotificationStream(throwing error: any Error) { + notificationContinuation?.finish(throwing: error) + notificationContinuation = nil + } + + func recordedSendCallCount() -> Int { + sendCallCount + } + + func recordedCloseCallCount() -> Int { + closeCallCount + } +} + +private actor CompletionProbe { + private var completed = false + + func recordCompletion() { + completed = true + } + + func hasCompleted() -> Bool { + completed + } +} + @Suite("app-server client") struct AppServerClientTests { @Test func processTransportConfigurationResolvesCodexFromProvidedPath() throws { @@ -1811,6 +1967,450 @@ struct AppServerClientTests { #expect(try await iterator.next() == .messageDelta("review text", itemID: "message-1")) } + @Test func clientCloseCanReturnBeforeBackendRouterAndEventSessionFinish() async throws { + let transport = DeferredNotificationCloseTransport() + let client = AppServerClient(transport: transport) + let backend = AppServerCodexReviewBackend(client: client) + let run = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-1", + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "thread-1" + ) + var iterator = await eventSequence(backend, run).makeAsyncIterator() + + try await client.close() + + try await transport.emitServerNotification( + method: "item/agentMessage/delta", + params: TestDeltaNotification( + threadID: "thread-1", + turnID: "turn-1", + itemID: "message-1", + delta: "delivered after client close" + ) + ) + #expect(try await iterator.next() == .started( + turnID: "turn-1", + reviewThreadID: "thread-1", + model: nil + )) + #expect(try await iterator.next() == .messageDelta( + "delivered after client close", + itemID: "message-1" + )) + + await transport.finishNotificationStream(throwing: JSONRPC.Error.closed) + await #expect(throws: ReviewAttemptStreamFailure.unexpectedConnection( + .connection(JSONRPC.Error.closed.localizedDescription) + )) { + _ = try await iterator.next() + } + } + + @Test func backendLifecycleCloseJoinsOwnedRouterAndEventSessions() async throws { + let transport = DeferredNotificationCloseTransport() + let backend = AppServerCodexReviewBackend( + client: AppServerClient(transport: transport) + ) + let lifecycle = backend.runtimeOwnerLifecycleHandle + let run = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-1", + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "thread-1" + ) + var iterator = await eventSequence(backend, run).makeAsyncIterator() + let firstCompletion = CompletionProbe() + let secondCompletion = CompletionProbe() + + let firstClose = Task { + do { + try await lifecycle.closeAndWait() + await firstCompletion.recordCompletion() + } catch { + await firstCompletion.recordCompletion() + throw error + } + } + let secondClose = Task { + do { + try await lifecycle.closeAndWait() + await secondCompletion.recordCompletion() + } catch { + await secondCompletion.recordCompletion() + throw error + } + } + await backend.waitForRuntimeOwnerCloseCallersForTesting(2) + await backend.waitForClientCloseResultBeforeRouterWaitForTesting() + #expect(await firstCompletion.hasCompleted() == false) + #expect(await secondCompletion.hasCompleted() == false) + #expect(await transport.recordedCloseCallCount() == 1) + await #expect(throws: JSONRPC.Error.closed) { + _ = try await backend.startReview(.init( + jobID: "job-2", + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + )) + } + #expect(await transport.recordedSendCallCount() == 0) + + try await transport.emitServerNotification( + method: "item/agentMessage/delta", + params: TestDeltaNotification( + threadID: "thread-1", + turnID: "turn-1", + itemID: "message-1", + delta: "delivered while lifecycle close waits" + ) + ) + #expect(try await iterator.next() == .started( + turnID: "turn-1", + reviewThreadID: "thread-1", + model: nil + )) + #expect(try await iterator.next() == .messageDelta( + "delivered while lifecycle close waits", + itemID: "message-1" + )) + + await transport.finishNotificationStream(throwing: JSONRPC.Error.closed) + try await firstClose.value + try await secondClose.value + try await lifecycle.closeAndWait() + + #expect(await transport.recordedCloseCallCount() == 1) + #expect(await backend.notificationRouterIsRunningForTesting() == false) + await #expect(throws: ReviewAttemptStreamFailure.unexpectedConnection( + .connection(JSONRPC.Error.closed.localizedDescription) + )) { + _ = try await iterator.next() + } + + } + + @Test func backendLifecycleCloseJoinsOwnedTasksBeforeReplayingClientCloseFailure() async throws { + let closeFailure = ReviewRuntimeCloseFailure.connection("close failed") + let transport = DeferredNotificationCloseTransport(closeFailure: closeFailure) + let backend = AppServerCodexReviewBackend( + client: AppServerClient(transport: transport) + ) + let lifecycle = backend.runtimeOwnerLifecycleHandle + let run = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-1", + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "thread-1" + ) + var iterator = await eventSequence(backend, run).makeAsyncIterator() + let completion = CompletionProbe() + + let close = Task { + do { + try await lifecycle.closeAndWait() + await completion.recordCompletion() + } catch { + await completion.recordCompletion() + throw error + } + } + await backend.waitForRuntimeOwnerCloseCallersForTesting(1) + await backend.waitForClientCloseResultBeforeRouterWaitForTesting() + #expect(await completion.hasCompleted() == false) + await transport.finishNotificationStream(throwing: JSONRPC.Error.closed) + + await #expect(throws: closeFailure) { + try await close.value + } + await #expect(throws: closeFailure) { + try await lifecycle.closeAndWait() + } + #expect(await transport.recordedCloseCallCount() == 1) + #expect(await backend.notificationRouterIsRunningForTesting() == false) + await #expect(throws: ReviewAttemptStreamFailure.unexpectedConnection( + .connection(JSONRPC.Error.closed.localizedDescription) + )) { + _ = try await iterator.next() + } + } + + @Test func backendLifecycleCloseJoinsInFlightRouterStart() async throws { + let transport = DeferredNotificationCloseTransport() + let notificationStreamGate = AsyncGate() + await transport.holdNotificationStream(on: notificationStreamGate) + let backend = AppServerCodexReviewBackend( + client: AppServerClient(transport: transport) + ) + let lifecycle = backend.runtimeOwnerLifecycleHandle + let run = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-1", + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "thread-1" + ) + let operationCompletion = CompletionProbe() + let closeCompletion = CompletionProbe() + + let operation = Task { + let attempt = await backend.reviewAttemptForTesting(run) + await operationCompletion.recordCompletion() + return attempt + } + await transport.waitForNotificationStreamRequest() + let close = Task { + do { + try await lifecycle.closeAndWait() + await closeCompletion.recordCompletion() + } catch { + await closeCompletion.recordCompletion() + throw error + } + } + await backend.waitForRuntimeOwnerCloseCallersForTesting(1) + #expect(await operationCompletion.hasCompleted() == false) + #expect(await closeCompletion.hasCompleted() == false) + + await notificationStreamGate.open() + let attempt = await operation.value + await backend.waitForClientCloseResultBeforeRouterWaitForTesting() + #expect(await operationCompletion.hasCompleted()) + #expect(await closeCompletion.hasCompleted() == false) + + await transport.finishNotificationStream(throwing: JSONRPC.Error.closed) + try await close.value + #expect(await closeCompletion.hasCompleted()) + #expect(await backend.notificationRouterIsRunningForTesting() == false) + + var iterator = BackendReviewEventSequence(mailbox: attempt.events).makeAsyncIterator() + await #expect(throws: ReviewAttemptStreamFailure.unexpectedConnection( + .connection(JSONRPC.Error.closed.localizedDescription) + )) { + _ = try await iterator.next() + } + } + + @Test func backendLifecycleCloseJoinsAdmittedThreadStartBeforeSessionSnapshot() async throws { + let transport = DeferredNotificationCloseTransport() + try await transport.enqueue(AppServerAPI.Initialize.Response(), for: "initialize") + try await transport.enqueue( + AppServerAPI.Thread.Start.Response(threadID: "thread-1", model: "gpt-5"), + for: "thread/start" + ) + try await transport.enqueue( + AppServerAPI.Review.Start.Response(turnID: "turn-1", reviewThreadID: "thread-1"), + for: "review/start" + ) + let threadStartBarrier = RequestBarrier() + await transport.holdNext(method: "thread/start", barrier: threadStartBarrier) + let backend = AppServerCodexReviewBackend( + client: AppServerClient(transport: transport) + ) + let lifecycle = backend.runtimeOwnerLifecycleHandle + let startCompletion = CompletionProbe() + let closeCompletion = CompletionProbe() + + let start = Task { + do { + let attempt = try await backend.startReview( + .init( + jobID: "job-1", + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + ) + await startCompletion.recordCompletion() + return attempt + } catch { + await startCompletion.recordCompletion() + throw error + } + } + await threadStartBarrier.waitUntilEntered() + let close = Task { + do { + try await lifecycle.closeAndWait() + await closeCompletion.recordCompletion() + } catch { + await closeCompletion.recordCompletion() + throw error + } + } + await backend.waitForRuntimeOwnerCloseCallersForTesting(1) + #expect(await startCompletion.hasCompleted() == false) + #expect(await closeCompletion.hasCompleted() == false) + + await threadStartBarrier.open() + let attempt = try await start.value + await backend.waitForClientCloseResultBeforeRouterWaitForTesting() + #expect(await startCompletion.hasCompleted()) + #expect(await closeCompletion.hasCompleted() == false) + + await transport.finishNotificationStream(throwing: JSONRPC.Error.closed) + try await close.value + #expect(await closeCompletion.hasCompleted()) + #expect(await backend.notificationRouterIsRunningForTesting() == false) + + var iterator = BackendReviewEventSequence(mailbox: attempt.events).makeAsyncIterator() + await #expect(throws: ReviewAttemptStreamFailure.unexpectedConnection( + .connection(JSONRPC.Error.closed.localizedDescription) + )) { + _ = try await iterator.next() + } + } + + @Test func backendLifecycleCloseJoinsAdmittedRecoveryAcrossRollbackAndReviewStart() async throws { + let transport = DeferredNotificationCloseTransport() + try await transport.enqueue(AppServerAPI.Initialize.Response(), for: "initialize") + try await transport.enqueue( + AppServerAPI.Thread.Start.Response(threadID: "thread-1", model: "gpt-5"), + for: "thread/start" + ) + try await transport.enqueue( + AppServerAPI.Review.Start.Response(turnID: "turn-1", reviewThreadID: "thread-1"), + for: "review/start" + ) + try await transport.enqueue( + AppServerAPI.Review.Start.Response(turnID: "turn-2", reviewThreadID: "thread-1"), + for: "review/start" + ) + let backend = AppServerCodexReviewBackend( + client: AppServerClient(transport: transport) + ) + let lifecycle = backend.runtimeOwnerLifecycleHandle + let initialAttempt = try await backend.startReview(.init( + jobID: "job-1", + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + )) + let handoff = try await backend.prepareReviewRecovery( + initialAttempt, + reason: .init(message: "Restart runtime") + ) + let rollbackBarrier = RequestBarrier() + let reviewStartBarrier = RequestBarrier() + await transport.holdNext(method: "thread/rollback", barrier: rollbackBarrier) + await transport.holdNext(method: "review/start", barrier: reviewStartBarrier) + let recoveryCompletion = CompletionProbe() + let closeCompletion = CompletionProbe() + + let recovery = Task { + do { + let attempt = try await backend.resumeReviewRecovery( + handoff, + request: .init( + jobID: "job-1", + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + ) + await recoveryCompletion.recordCompletion() + return attempt + } catch { + await recoveryCompletion.recordCompletion() + throw error + } + } + await rollbackBarrier.waitUntilEntered() + let close = Task { + do { + try await lifecycle.closeAndWait() + await closeCompletion.recordCompletion() + } catch { + await closeCompletion.recordCompletion() + throw error + } + } + await backend.waitForRuntimeOwnerCloseCallersForTesting(1) + #expect(await recoveryCompletion.hasCompleted() == false) + #expect(await closeCompletion.hasCompleted() == false) + + await rollbackBarrier.open() + await reviewStartBarrier.waitUntilEntered() + #expect(await recoveryCompletion.hasCompleted() == false) + #expect(await closeCompletion.hasCompleted() == false) + + await reviewStartBarrier.open() + let recoveredAttempt = try await recovery.value + await backend.waitForClientCloseResultBeforeRouterWaitForTesting() + #expect(await recoveryCompletion.hasCompleted()) + #expect(await closeCompletion.hasCompleted() == false) + + await transport.finishNotificationStream(throwing: JSONRPC.Error.closed) + try await close.value + #expect(await closeCompletion.hasCompleted()) + #expect(await backend.notificationRouterIsRunningForTesting() == false) + + var iterator = BackendReviewEventSequence(mailbox: recoveredAttempt.events).makeAsyncIterator() + await #expect(throws: ReviewAttemptStreamFailure.unexpectedConnection( + .connection(JSONRPC.Error.closed.localizedDescription) + )) { + _ = try await iterator.next() + } + } + + @Test func backendLifecycleCloseJoinsAdmittedInterrupt() async throws { + let transport = DeferredNotificationCloseTransport() + try await transport.enqueue(AppServerAPI.Initialize.Response(), for: "initialize") + let interruptBarrier = RequestBarrier() + await transport.holdNext(method: "turn/interrupt", barrier: interruptBarrier) + let backend = AppServerCodexReviewBackend( + client: AppServerClient(transport: transport) + ) + let lifecycle = backend.runtimeOwnerLifecycleHandle + let run = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-1", + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "thread-1" + ) + var iterator = await eventSequence(backend, run).makeAsyncIterator() + let interruptCompletion = CompletionProbe() + let closeCompletion = CompletionProbe() + + let interrupt = Task { + do { + try await backend.interruptReview( + run, + reason: .init(message: "Stop") + ) + await interruptCompletion.recordCompletion() + } catch { + await interruptCompletion.recordCompletion() + throw error + } + } + await interruptBarrier.waitUntilEntered() + let close = Task { + do { + try await lifecycle.closeAndWait() + await closeCompletion.recordCompletion() + } catch { + await closeCompletion.recordCompletion() + throw error + } + } + await backend.waitForRuntimeOwnerCloseCallersForTesting(1) + #expect(await interruptCompletion.hasCompleted() == false) + #expect(await closeCompletion.hasCompleted() == false) + + await interruptBarrier.open() + try await interrupt.value + await backend.waitForClientCloseResultBeforeRouterWaitForTesting() + #expect(await interruptCompletion.hasCompleted()) + #expect(await closeCompletion.hasCompleted() == false) + + await transport.finishNotificationStream(throwing: JSONRPC.Error.closed) + try await close.value + #expect(await closeCompletion.hasCompleted()) + #expect(await backend.notificationRouterIsRunningForTesting() == false) + await #expect(throws: ReviewAttemptStreamFailure.unexpectedConnection( + .connection(JSONRPC.Error.closed.localizedDescription) + )) { + _ = try await iterator.next() + } + } + @Test func backendPreservesNotificationStreamErrorForLateEventStreamSubscriber() async throws { let transport = FakeJSONRPCTransport() try await enqueueInitialize(transport) From c76ff88c5c9f0c466d6af899ee989ed1bba0ed29 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 04:40:19 +0900 Subject: [PATCH 27/38] test(review): gate joined cancellation owner --- Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift index f31035ce..c7017cf5 100644 --- a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift +++ b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift @@ -666,6 +666,7 @@ struct ReviewAttemptProcessorTests { forceClose: {} ) } + #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) let duplicateCancellation = Task { try await admission.cancel( .system(message: "Runtime stop"), @@ -673,7 +674,6 @@ struct ReviewAttemptProcessorTests { forceClose: {} ) } - #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) try await admission.recordCanonicalTerminal( .interrupted(.server(message: "network recovery")), for: run From 6a05d14a67932ffddb2f98791855df4a7a6a93a3 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 07:36:42 +0900 Subject: [PATCH 28/38] Add joined Store application close lifecycle --- .../CodexReviewKitProductConsumer/main.swift | 3 +- .../CodexReview/ReviewRuntimeLifecycle.swift | 281 +++++++- .../CodexReview/Store/CodexReviewStore.swift | 669 +++++++++++++++--- .../Store/CodexReviewStoreBackend.swift | 9 +- .../Store/CodexReviewStoreCancellation.swift | 189 +++-- ...CodexReviewStoreRateLimitAutoRefresh.swift | 33 + .../Store/CodexReviewStoreReviews.swift | 196 ++++- .../Store/CodexReviewStoreTesting.swift | 28 +- .../AppServerCodexReviewBackend.swift | 46 +- Sources/CodexReviewHost/CodexReviewHost.swift | 131 +++- .../LiveCodexReviewStoreBackend.swift | 298 +++++--- .../CodexReviewMCPHTTPServer.swift | 91 ++- Sources/CodexReviewTesting/TestSupport.swift | 216 +++++- .../AppServerClientTests.swift | 26 +- .../CodexReviewHostTests.swift | 144 +++- .../CodexReviewMCPHTTPServerTests.swift | 10 +- .../CodexReviewStoreCommandTests.swift | 40 +- .../CodexReviewStoreLifecycleTests.swift | 595 +++++++++++++++- ...ReviewStoreRateLimitAutoRefreshTests.swift | 60 +- 19 files changed, 2648 insertions(+), 417 deletions(-) diff --git a/Fixtures/CodexReviewKitProductConsumer/Sources/CodexReviewKitProductConsumer/main.swift b/Fixtures/CodexReviewKitProductConsumer/Sources/CodexReviewKitProductConsumer/main.swift index 25fa1388..1dc9ca4d 100644 --- a/Fixtures/CodexReviewKitProductConsumer/Sources/CodexReviewKitProductConsumer/main.swift +++ b/Fixtures/CodexReviewKitProductConsumer/Sources/CodexReviewKitProductConsumer/main.swift @@ -7,7 +7,7 @@ import TextTransitions @main @MainActor struct CodexReviewKitProductConsumer { - static func main() { + static func main() async throws { let lifecycleInitializer: ( ReviewJobState, Int?, @@ -73,6 +73,7 @@ struct CodexReviewKitProductConsumer { fatalError("TextTransitions public rendering contract drifted.") } + try await store.close() withExtendedLifetime((store, windowController, transitionView)) {} print("CodexReviewKit public product consumer passed.") } diff --git a/Sources/CodexReview/ReviewRuntimeLifecycle.swift b/Sources/CodexReview/ReviewRuntimeLifecycle.swift index 74f3ae57..435a5148 100644 --- a/Sources/CodexReview/ReviewRuntimeLifecycle.swift +++ b/Sources/CodexReview/ReviewRuntimeLifecycle.swift @@ -68,13 +68,191 @@ package protocol RuntimeLifecycleHandle: Sendable { package struct PreparedRuntime: Sendable { package let snapshot: RuntimePublicationSnapshot package let handle: any RuntimeLifecycleHandle + package let closeRecord: RuntimeCloseRecord + @MainActor package init( snapshot: RuntimePublicationSnapshot, - handle: any RuntimeLifecycleHandle + handle: any RuntimeLifecycleHandle, + closeRecord: RuntimeCloseRecord = RuntimeCloseRecord() ) { self.snapshot = snapshot self.handle = handle + self.closeRecord = closeRecord + } +} + +@MainActor +package final class RuntimeCloseRecord { + package struct JoinResult { + package let failures: [ReviewClosePrimaryFailure] + package let installedClose: Bool + } + + private enum State { + case open + case closing(Task<[ReviewClosePrimaryFailure], Never>) + case closed([ReviewClosePrimaryFailure]) + } + + private var state: State = .open + private var failuresWereConsumed = false + + package init() {} + + package func closeAndWait( + handle: any RuntimeLifecycleHandle, + purpose: ReviewRuntimeTransitionPurpose + ) async -> JoinResult { + let task: Task<[ReviewClosePrimaryFailure], Never> + let installedClose: Bool + switch state { + case .open: + let newTask = Task<[ReviewClosePrimaryFailure], Never> { @MainActor in + let record = ReviewRuntimeTransitionRecord() + var closeFailureWasRecorded = false + do { + try await handle.close(purpose: purpose) + } catch { + closeFailureWasRecorded = true + record.record( + error, + fallback: .client(error.localizedDescription) + ) + } + do { + try await handle.waitUntilClosed() + } catch { + if closeFailureWasRecorded == false { + record.record( + error, + fallback: .client(error.localizedDescription) + ) + } + } + return record.failures + } + state = .closing(newTask) + task = newTask + installedClose = true + case .closing(let existingTask): + task = existingTask + installedClose = false + case .closed(let failures): + return .init(failures: failures, installedClose: false) + } + + let failures = await task.value + state = .closed(failures) + return .init(failures: failures, installedClose: installedClose) + } + + package func consumeFailures() -> [ReviewClosePrimaryFailure] { + guard failuresWereConsumed == false else { + return [] + } + guard case .closed(let failures) = state else { + return [] + } + failuresWereConsumed = true + return failures + } +} + +@MainActor +package final class ReviewCloseFailureLedger { + package private(set) var failures: [ReviewClosePrimaryFailure] = [] + package private(set) var consumedReviewCleanupJobIDs: Set = [] + private var forceCloseFailureJobIDs: Set = [] + + package init() {} + + package func record( + _ error: any Error, + fallback: ReviewLifecycleResourceFailure + ) { + if let failure = error as? ReviewInterruptRequestFailure { + failures.append(.interruptRequest(failure)) + } else if let failure = error as? ReviewRuntimeCloseFailure { + failures.append(.attemptRuntime(failure)) + } else if let aggregate = error as? ReviewLifecycleResourceFailureAggregate { + failures.append(.lifecycleResources(aggregate)) + } else if let failure = error as? ReviewLifecycleResourceFailure { + failures.append(.lifecycleResources(.init(first: failure))) + } else { + failures.append(.lifecycleResources(.init(first: fallback))) + } + } + + package func record(_ failure: ReviewClosePrimaryFailure) { + failures.append(failure) + } + + package func record(contentsOf failures: [ReviewClosePrimaryFailure]) { + self.failures.append(contentsOf: failures) + } + + package func recordReviewCleanupFailure( + _ failure: ReviewRuntimeCloseFailure, + jobID: String + ) { + guard consumedReviewCleanupJobIDs.insert(jobID).inserted else { + return + } + failures.append(.attemptRuntime(failure)) + } + + package func recordForceCloseFailures( + _ failures: [ReviewClosePrimaryFailure], + jobID: String + ) { + guard failures.isEmpty == false else { + return + } + forceCloseFailureJobIDs.insert(jobID) + self.failures.append(contentsOf: failures) + } + + package func ownsForceCloseFailure(for jobID: String) -> Bool { + forceCloseFailureJobIDs.contains(jobID) + } + + package func merge(_ other: ReviewCloseFailureLedger) { + failures.append(contentsOf: other.failures) + importReceipts(from: other) + } + + package func importReceipts(from other: ReviewCloseFailureLedger) { + consumedReviewCleanupJobIDs.formUnion(other.consumedReviewCleanupJobIDs) + forceCloseFailureJobIDs.formUnion(other.forceCloseFailureJobIDs) + } + + package var failureDescription: String? { + failures.first.map { + ReviewCloseFailureAggregate( + first: $0, + additionalInLifecycleOrder: Array(failures.dropFirst()) + ).localizedDescription + } + } +} + +package typealias ReviewRuntimeTransitionRecord = ReviewCloseFailureLedger + +package struct ReviewRuntimePreparationFailure: LocalizedError, Sendable { + package let preparationDescription: String + package let cleanupFailures: ReviewLifecycleResourceFailureAggregate + + package init( + preparationError: any Error, + cleanupFailures: ReviewLifecycleResourceFailureAggregate + ) { + self.preparationDescription = preparationError.localizedDescription + self.cleanupFailures = cleanupFailures + } + + package var errorDescription: String? { + "\(preparationDescription); \(cleanupFailures.localizedDescription)" } } @@ -82,7 +260,8 @@ package enum ReviewStoreRuntimeState { case stopped(ReviewRuntimeGeneration) case acquiring( generation: ReviewRuntimeGeneration, - task: Task + task: Task, + record: ReviewRuntimeTransitionRecord ) case running( generation: ReviewRuntimeGeneration, @@ -92,7 +271,9 @@ package enum ReviewStoreRuntimeState { case transitioning( generation: ReviewRuntimeGeneration, purpose: ReviewRuntimeTransitionPurpose, - task: Task + task: Task, + record: ReviewRuntimeTransitionRecord, + sourceRuntime: PreparedRuntime? ) case failed( generation: ReviewRuntimeGeneration, @@ -103,13 +284,103 @@ package enum ReviewStoreRuntimeState { package var generation: ReviewRuntimeGeneration { switch self { case .stopped(let generation), - .acquiring(let generation, _), + .acquiring(let generation, _, _), .running(let generation, _, _), - .transitioning(let generation, _, _), + .transitioning(let generation, _, _, _, _), .failed(let generation, _, _): generation } } + + package var runtimeForClose: PreparedRuntime? { + switch self { + case .running(_, let runtime, _): + return runtime + case .transitioning(_, _, _, _, let sourceRuntime): + return sourceRuntime + case .stopped, .acquiring, .failed: + return nil + } + } +} + +package enum ReviewStoreLifetimeState { + case open + case closing(Task, Never>) + case closed(Result) +} + +package struct ReviewStoreCommandRegistry { + package enum Admission { + case open + case closed + } + + package struct DrainWaiter { + let continuation: CheckedContinuation + } + + private(set) var admission: Admission = .open + private(set) var nextID: UInt64 = 0 + private(set) var activeIDs: Set = [] + private(set) var ownedTasks: [UInt64: Task] = [:] + private var drainWaiters: [DrainWaiter] = [] + + package mutating func register() -> UInt64? { + guard case .open = admission else { + return nil + } + nextID &+= 1 + activeIDs.insert(nextID) + return nextID + } + + package mutating func installOwnedTask( + _ task: Task, + for id: UInt64 + ) { + precondition( + activeIDs.contains(id), + "ReviewStoreCommandRegistry must register a command before installing its Task." + ) + ownedTasks[id] = task + } + + package mutating func closeAdmission() { + admission = .closed + resumeDrainWaitersIfNeeded() + } + + package mutating func finish(_ id: UInt64) { + activeIDs.remove(id) + ownedTasks.removeValue(forKey: id) + resumeDrainWaitersIfNeeded() + } + + package func ownedTaskSnapshot() -> [(UInt64, Task)] { + ownedTasks.sorted { $0.key < $1.key } + } + + package mutating func appendDrainWaiter( + _ continuation: CheckedContinuation + ) { + if activeIDs.isEmpty { + continuation.resume() + } else { + drainWaiters.append(.init(continuation: continuation)) + } + } + + private mutating func resumeDrainWaitersIfNeeded() { + guard case .closed = admission, activeIDs.isEmpty else { + return + } + let waiters = drainWaiters + drainWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.continuation.resume() + } + } } @MainActor diff --git a/Sources/CodexReview/Store/CodexReviewStore.swift b/Sources/CodexReview/Store/CodexReviewStore.swift index 2119f514..f8f3ad46 100644 --- a/Sources/CodexReview/Store/CodexReviewStore.swift +++ b/Sources/CodexReview/Store/CodexReviewStore.swift @@ -10,6 +10,11 @@ public final class CodexReviewStore { package var timeoutTask: Task? } + package struct CloseCallerWaiter { + package var targetCount: Int + package var continuation: CheckedContinuation + } + public package(set) var serverState: CodexReviewServerState = .stopped public let auth: CodexReviewAuthModel package let settings: SettingsStore @@ -30,15 +35,27 @@ public final class CodexReviewStore { @ObservationIgnored package let clock: CodexReviewClock @ObservationIgnored package let idGenerator: CodexReviewIDGenerator @ObservationIgnored package var reviewAttemptOwnerships: [String: ReviewAttemptOwnership] = [:] + @ObservationIgnored package var reviewRegistrationOrder: [String] = [] @ObservationIgnored package var reviewCleanupFailures: [String: ReviewRuntimeCloseFailure] = [:] @ObservationIgnored package var reviewWorkerTasks: [String: Task] = [:] - @ObservationIgnored package var runtimeStopDetachedReviewWorkerTasks: [String: Task] = [:] + @ObservationIgnored package var nextReviewMutationID: UInt64 = 0 + @ObservationIgnored package var reviewMutationTasks: [UInt64: Task] = [:] + @ObservationIgnored package var reviewMutationPreparationForTesting: (@MainActor @Sendable () async -> Void)? + @ObservationIgnored package var reviewCleanupPreparationForTesting: (@MainActor @Sendable () async -> Void)? + @ObservationIgnored package var reviewTerminalPublicationPreparationForTesting: (@MainActor @Sendable () async -> Void)? + @ObservationIgnored package var runtimeForceCloseReceiptRecordedForTesting: (@MainActor @Sendable () async -> Void)? + @ObservationIgnored package var storeCommandRegistry = ReviewStoreCommandRegistry() + @ObservationIgnored package var closeCallerCount = 0 + @ObservationIgnored package var closeCallerWaiters: [CloseCallerWaiter] = [] @ObservationIgnored package var reviewTerminalWaiters: [String: [ReviewTerminalWaiter]] = [:] @ObservationIgnored package var closedSessions: Set = [] @ObservationIgnored package var accountRateLimitAutoRefreshDriver: CodexReviewStoreRateLimitAutoRefreshDriver? + @ObservationIgnored package var lifetimeState: ReviewStoreLifetimeState = .open + @ObservationIgnored package var applicationCloseFailureLedger: ReviewCloseFailureLedger? @ObservationIgnored package var runtimeState: ReviewStoreRuntimeState = .stopped( .init(rawValue: 0) ) + @ObservationIgnored package var lastRuntimeTransitionRecord: ReviewRuntimeTransitionRecord? package init( backend: any CodexReviewStoreBackend = PreviewCodexReviewStoreBackend(), @@ -82,8 +99,11 @@ public final class CodexReviewStore { isolated deinit { accountRateLimitAutoRefreshDriver?.cancel() + if case .closing(let task) = lifetimeState { + task.cancel() + } switch runtimeState { - case .acquiring(_, let task), .transitioning(_, _, let task): + case .acquiring(_, let task, _), .transitioning(_, _, let task, _, _): task.cancel() case .stopped, .running, .failed: break @@ -91,7 +111,10 @@ public final class CodexReviewStore { for task in reviewWorkerTasks.values { task.cancel() } - for task in runtimeStopDetachedReviewWorkerTasks.values { + for task in reviewMutationTasks.values { + task.cancel() + } + for (_, task) in storeCommandRegistry.ownedTaskSnapshot() { task.cancel() } for waiters in reviewTerminalWaiters.values { @@ -100,6 +123,9 @@ public final class CodexReviewStore { waiter.continuation.resume() } } + for waiter in closeCallerWaiters { + waiter.continuation.resume() + } } public static func makePreviewStore(diagnosticsURL: URL? = nil) -> CodexReviewStore { @@ -138,10 +164,13 @@ public final class CodexReviewStore { } public func start(forceRestartIfNeeded: Bool = false) async { + guard case .open = lifetimeState else { + return + } switch runtimeState { case .acquiring: return - case .transitioning(_, _, let task): + case .transitioning(_, _, let task, _, _): await task.value return case .running where forceRestartIfNeeded == false: @@ -176,6 +205,7 @@ public final class CodexReviewStore { return } let generation = previousGeneration.successor() + let record = ReviewRuntimeTransitionRecord() serverState = .starting serverURL = nil writeDiagnosticsIfNeeded() @@ -183,10 +213,15 @@ public final class CodexReviewStore { guard let self else { return } await self.performRuntimeAcquisition( generation: generation, - purpose: purpose + purpose: purpose, + record: record ) } - runtimeState = .acquiring(generation: generation, task: task) + runtimeState = .acquiring( + generation: generation, + task: task, + record: record + ) await task.value } @@ -197,6 +232,7 @@ public final class CodexReviewStore { retainedServerURL: URL? ) async { let generation = previousGeneration.successor() + let record = ReviewRuntimeTransitionRecord() serverState = .starting writeDiagnosticsIfNeeded() let task = Task { @MainActor [weak self] in @@ -205,111 +241,414 @@ public final class CodexReviewStore { generation: generation, previousRuntime: previousRuntime, retainedMCPGeneration: retainedMCPGeneration, - retainedServerURL: retainedServerURL + retainedServerURL: retainedServerURL, + record: record ) } runtimeState = .transitioning( generation: generation, purpose: .restartSameAccount, - task: task + task: task, + record: record, + sourceRuntime: previousRuntime ) await task.value } public func stop() async { + switch lifetimeState { + case .open: + break + case .closing(let task): + _ = await task.value + return + case .closed: + return + } let previousState = runtimeState switch previousState { case .stopped: transitionToStopped() return - case .transitioning(_, .stop, let task): + case .transitioning(_, .stop, let task, _, _): await task.value return case .acquiring, .running, .transitioning, .failed: break } let invalidatedGeneration = previousState.generation.successor() + let record = ReviewRuntimeTransitionRecord() let task = Task { @MainActor [weak self] in guard let self else { return } await self.performRuntimeStop( previousState: previousState, - invalidatedGeneration: invalidatedGeneration + invalidatedGeneration: invalidatedGeneration, + record: record ) } runtimeState = .transitioning( generation: invalidatedGeneration, purpose: .stop, - task: task + task: task, + record: record, + sourceRuntime: previousState.runtimeForClose ) await task.value } + /// Permanently closes this Store and awaits every resource it owns. + /// + /// The first call closes mutation admission and records one application-lifetime + /// close operation. Concurrent and later calls join or replay that exact result. + /// After close begins, `start(forceRestartIfNeeded:)` and `restart()` are no-ops + /// and new review mutations are rejected. A thrown error reports the recorded + /// lifecycle failures after all close stages have been attempted. + public func close() async throws { + closeCallerCount += 1 + let completedCloseCallerWaiters = closeCallerWaiters.filter { + closeCallerCount >= $0.targetCount + } + closeCallerWaiters.removeAll { + closeCallerCount >= $0.targetCount + } + for waiter in completedCloseCallerWaiters { + waiter.continuation.resume() + } + let task: Task, Never> + switch lifetimeState { + case .open: + storeCommandRegistry.closeAdmission() + let previousRuntimeState = runtimeState + let invalidatedGeneration = previousRuntimeState.generation.successor() + let failureLedger = ReviewCloseFailureLedger() + applicationCloseFailureLedger = failureLedger + let newTask = Task, Never> { @MainActor [self] in + return await self.performApplicationClose( + previousRuntimeState: previousRuntimeState, + invalidatedGeneration: invalidatedGeneration, + failureLedger: failureLedger + ) + } + lifetimeState = .closing(newTask) + task = newTask + case .closing(let existingTask): + task = existingTask + case .closed(let result): + try result.get() + return + } + + let result = await task.value + lifetimeState = .closed(result) + applicationCloseFailureLedger = nil + try result.get() + } + + package func waitForCloseCallersForTesting(_ count: Int) async { + if closeCallerCount >= count { + return + } + await withCheckedContinuation { continuation in + if closeCallerCount >= count { + continuation.resume() + } else { + closeCallerWaiters.append(.init( + targetCount: count, + continuation: continuation + )) + } + } + } + + package func registerStoreCommand() -> UInt64? { + storeCommandRegistry.register() + } + + package func installOwnedStoreCommandTask( + _ task: Task, + for id: UInt64 + ) { + storeCommandRegistry.installOwnedTask(task, for: id) + } + + package func finishStoreCommand(_ id: UInt64) { + storeCommandRegistry.finish(id) + } + + package func waitForAdmittedStoreCommands() async { + if storeCommandRegistry.activeIDs.isEmpty { + return + } + await withCheckedContinuation { continuation in + storeCommandRegistry.appendDrainWaiter(continuation) + } + } + + package func performAdmittedStoreCommand( + _ operation: @MainActor () async -> Void + ) async { + guard let commandID = registerStoreCommand() else { + return + } + defer { finishStoreCommand(commandID) } + await operation() + } + + package func performThrowingAdmittedStoreCommand( + _ operation: @MainActor () async throws -> Value + ) async throws -> Value { + guard let commandID = registerStoreCommand() else { + throw CodexReviewAPI.Error.io("Review Store is closed.") + } + defer { finishStoreCommand(commandID) } + return try await operation() + } + + private func performApplicationClose( + previousRuntimeState: ReviewStoreRuntimeState, + invalidatedGeneration: ReviewRuntimeGeneration, + failureLedger: ReviewCloseFailureLedger + ) async -> Result { + let runningRuntime = previousRuntimeState.runtimeForClose + if let lastRuntimeTransitionRecord { + failureLedger.importReceipts(from: lastRuntimeTransitionRecord) + } + + await backend.mcpServerLifecycle.closeAdmission() + await cancelAndAwaitAllReviewMutationTasks() + + let cancellation = await requestActiveReviewCancellationsForApplicationClose( + failureLedger: failureLedger + ) + + await runningRuntime?.handle.closeAdmission() + var runtimePhysicalCloseCompleted = false + let unresolvedCancellationJobIDs = cancellation.jobIDs.filter { jobID in + cancellation.failedJobIDs.contains(jobID) + && (job(id: jobID)?.isTerminal == false || reviewWorkerTasks[jobID] != nil) + } + if unresolvedCancellationJobIDs.isEmpty == false { + if let runningRuntime { + await applicationCloseRuntime( + runningRuntime, + failureLedger: failureLedger + ) + runtimePhysicalCloseCompleted = true + await awaitReviewWorkers(jobIDs: unresolvedCancellationJobIDs) + } else { + await cancelAndAwaitReviewWorkersForRuntimeStop( + jobIDs: unresolvedCancellationJobIDs + ) + } + } + await finishAllReviewWaitersForStoreClose() + + do { + try await backend.mcpServerLifecycle.drainAdmittedHandlers() + } catch { + failureLedger.record(closePrimaryFailure( + from: error, + fallback: .mcpHandlerDrain(error.localizedDescription) + )) + } + + var mcpCloseFailureWasRecorded = false + do { + try await backend.mcpServerLifecycle.close() + } catch { + mcpCloseFailureWasRecorded = true + failureLedger.record(closePrimaryFailure( + from: error, + fallback: .mcpServer(error.localizedDescription) + )) + } + do { + try await backend.mcpServerLifecycle.waitUntilClosed() + } catch { + if mcpCloseFailureWasRecorded == false { + failureLedger.record(closePrimaryFailure( + from: error, + fallback: .mcpServer(error.localizedDescription) + )) + } + } + + await cancelAccountRateLimitAutoRefreshAndWait() + switch previousRuntimeState { + case .acquiring(_, let task, let record), + .transitioning(_, _, let task, let record, _): + task.cancel() + await task.value + failureLedger.merge(record) + case .stopped, .running, .failed: + if let lastRuntimeTransitionRecord { + failureLedger.merge(lastRuntimeTransitionRecord) + } + break + } + + await cancelAndAwaitAllReviewMutationTasks() + await cancelAndAwaitOwnedStoreCommandTasks() + await waitForAdmittedStoreCommands() + await awaitAllReviewWorkers() + await finishAllReviewWaitersForStoreClose() + do { + try await backend.stop(store: self) + } catch { + failureLedger.record(closePrimaryFailure( + from: error, + fallback: .client(error.localizedDescription) + )) + } + await backend.waitUntilStopped() + + for jobID in reviewRegistrationOrder + where failureLedger.consumedReviewCleanupJobIDs.contains(jobID) == false { + if let failure = reviewCleanupFailures[jobID] { + failureLedger.recordReviewCleanupFailure(failure, jobID: jobID) + } + } + + if let runningRuntime, runtimePhysicalCloseCompleted == false { + await applicationCloseRuntime( + runningRuntime, + failureLedger: failureLedger + ) + } + + runtimeState = .stopped(invalidatedGeneration) + let result: Result + if let first = failureLedger.failures.first { + let closeError = ReviewCloseError(failures: .init( + first: first, + additionalInLifecycleOrder: Array(failureLedger.failures.dropFirst()) + )) + transitionToFailed(closeError.localizedDescription) + result = .failure(closeError) + } else { + transitionToStopped() + result = .success(()) + } + return result + } + + private func applicationCloseRuntime( + _ runtime: PreparedRuntime, + failureLedger: ReviewCloseFailureLedger + ) async { + _ = await runtime.closeRecord.closeAndWait( + handle: runtime.handle, + purpose: .applicationClose + ) + failureLedger.record(contentsOf: runtime.closeRecord.consumeFailures()) + } + + private func closePrimaryFailure( + from error: any Error, + fallback: ReviewLifecycleResourceFailure + ) -> ReviewClosePrimaryFailure { + if let failure = error as? ReviewInterruptRequestFailure { + return .interruptRequest(failure) + } + if let failure = error as? ReviewRuntimeCloseFailure { + return .attemptRuntime(failure) + } + if let aggregate = error as? ReviewLifecycleResourceFailureAggregate { + return .lifecycleResources(aggregate) + } + if let failure = error as? ReviewLifecycleResourceFailure { + return .lifecycleResources(.init(first: failure)) + } + return .lifecycleResources(.init(first: fallback)) + } + private func performRuntimeStop( previousState: ReviewStoreRuntimeState, - invalidatedGeneration: ReviewRuntimeGeneration + invalidatedGeneration: ReviewRuntimeGeneration, + record: ReviewRuntimeTransitionRecord ) async { + defer { lastRuntimeTransitionRecord = record } switch previousState { - case .acquiring(_, let task): + case .acquiring(_, let task, let previousRecord): task.cancel() - await stopPreparedMCPServer() + await stopPreparedMCPServer(record: record) await task.value + record.merge(previousRecord) case .running(_, let runtime, _): - await stopPublishedRuntime(runtime) - case .transitioning(_, _, let task): + await stopPublishedRuntime(runtime, record: record) + case .transitioning(_, _, let task, let previousRecord, _): task.cancel() await task.value - await stopPreparedMCPServer() + record.merge(previousRecord) + await stopPreparedMCPServer(record: record) case .failed: - await stopPreparedMCPServer() + await stopPreparedMCPServer(record: record) case .stopped: break } - guard case .transitioning(let currentGeneration, .stop, _) = runtimeState, + guard case .transitioning(let currentGeneration, .stop, _, _, _) = runtimeState, currentGeneration == invalidatedGeneration else { return } runtimeState = .stopped(invalidatedGeneration) - transitionToStopped() + if let failureDescription = record.failureDescription { + transitionToFailed(failureDescription) + } else { + transitionToStopped() + } } - private func stopPublishedRuntime(_ runtime: PreparedRuntime) async { - await performPublishedRuntimeSemanticStop() + private func stopPublishedRuntime( + _ runtime: PreparedRuntime, + record: ReviewRuntimeTransitionRecord + ) async { + await performPublishedRuntimeSemanticStop(record: record) await runtime.handle.closeAdmission() - await stopPreparedMCPServer() - await closeAppServerRuntime(runtime, purpose: .stop) + await stopPreparedMCPServer(record: record) + await closeAppServerRuntime( + runtime, + purpose: .stop, + record: record + ) } - private func performPublishedRuntimeSemanticStop() async { - let locallyCancelledJobIDs: [String] - if backend.handlesActiveReviewStopCleanup { - locallyCancelledJobIDs = [] - } else { - do { - locallyCancelledJobIDs = try await requestActiveReviewCancellationsForRuntimeStop() - } catch { - locallyCancelledJobIDs = [] - } + private func performPublishedRuntimeSemanticStop( + record: ReviewRuntimeTransitionRecord + ) async { + let cancellation = await requestActiveReviewCancellationsForApplicationClose( + reason: .system(message: "Review runtime stopped."), + failureLedger: record + ) + do { + try await backend.stop(store: self) + } catch { + record.record(error, fallback: .client(error.localizedDescription)) } - await backend.stop(store: self) let remainingLocallyCancelledJobIDs = cancelActiveReviewsLocallyForRuntimeStop(cancelWorkers: false) - cancelAndDetachReviewWorkersForRuntimeStop( - jobIDs: Array(Set(locallyCancelledJobIDs + remainingLocallyCancelledJobIDs)) + await cancelAndAwaitReviewWorkersForRuntimeStop( + jobIDs: cancellation.jobIDs + remainingLocallyCancelledJobIDs ) + for jobID in cancellation.jobIDs + where record.consumedReviewCleanupJobIDs.contains(jobID) == false { + if let failure = reviewCleanupFailures[jobID] { + record.recordReviewCleanupFailure(failure, jobID: jobID) + } + } + await cancelAccountRateLimitAutoRefreshAndWait() } private func closeAppServerRuntime( _ runtime: PreparedRuntime, - purpose: ReviewRuntimeTransitionPurpose + purpose: ReviewRuntimeTransitionPurpose, + record: ReviewRuntimeTransitionRecord ) async { - do { - try await runtime.handle.close(purpose: purpose) - } catch { - writeDiagnosticsIfNeeded() - } - do { - try await runtime.handle.waitUntilClosed() - } catch { + let result = await runtime.closeRecord.closeAndWait( + handle: runtime.handle, + purpose: purpose + ) + record.record(contentsOf: runtime.closeRecord.consumeFailures()) + if result.failures.isEmpty == false { writeDiagnosticsIfNeeded() } } @@ -319,7 +658,7 @@ public final class CodexReviewStore { } public func waitUntilStopped() async { - if case .transitioning(_, _, let task) = runtimeState { + if case .transitioning(_, _, let task, _, _) = runtimeState { await task.value } await backend.waitUntilStopped() @@ -328,8 +667,10 @@ public final class CodexReviewStore { private func performRuntimeAcquisition( generation: ReviewRuntimeGeneration, - purpose: ReviewRuntimeTransitionPurpose + purpose: ReviewRuntimeTransitionPurpose, + record: ReviewRuntimeTransitionRecord ) async { + defer { lastRuntimeTransitionRecord = record } guard isCurrentAcquisition(generation) else { return } @@ -340,7 +681,7 @@ public final class CodexReviewStore { preparedMCPServer = mcpServer guard isCurrentAcquisition(generation) else { if currentTransitionOwnsMCPStop == false { - await stopPreparedMCPServer() + await stopPreparedMCPServer(record: record) } return } @@ -354,7 +695,8 @@ public final class CodexReviewStore { await closeStaleRuntime( runtime, mcpServerWasPrepared: true, - purpose: purpose + purpose: purpose, + record: record ) return } @@ -364,7 +706,8 @@ public final class CodexReviewStore { await closeStaleRuntime( runtime, mcpServerWasPrepared: true, - purpose: purpose + purpose: purpose, + record: record ) return } @@ -377,7 +720,8 @@ public final class CodexReviewStore { await closeStaleRuntime( runtime, mcpServerWasPrepared: true, - purpose: purpose + purpose: purpose, + record: record ) return } @@ -389,21 +733,29 @@ public final class CodexReviewStore { ) publishMCPServer(serverURL: mcpSnapshot.serverURL) } catch { + let visibleFailureDescription: String + if let preparationFailure = error as? ReviewRuntimePreparationFailure { + record.record(.lifecycleResources(preparationFailure.cleanupFailures)) + visibleFailureDescription = preparationFailure.preparationDescription + } else { + visibleFailureDescription = error.localizedDescription + } if let preparedRuntime { await closeStaleRuntime( preparedRuntime, mcpServerWasPrepared: preparedMCPServer != nil, - purpose: purpose + purpose: purpose, + record: record ) } else if preparedMCPServer != nil, currentTransitionOwnsMCPStop == false { - await stopPreparedMCPServer() + await stopPreparedMCPServer(record: record) } guard isCurrentAcquisition(generation) else { return } runtimeState = .stopped(generation) - transitionToFailed(error.localizedDescription) + transitionToFailed(visibleFailureDescription) } } @@ -413,7 +765,8 @@ public final class CodexReviewStore { ) async { await performRuntimeAcquisition( generation: generation, - purpose: purpose + purpose: purpose, + record: ReviewRuntimeTransitionRecord() ) } @@ -421,15 +774,18 @@ public final class CodexReviewStore { generation: ReviewRuntimeGeneration, previousRuntime: PreparedRuntime?, retainedMCPGeneration: MCPServerGeneration, - retainedServerURL: URL? + retainedServerURL: URL?, + record: ReviewRuntimeTransitionRecord ) async { + defer { lastRuntimeTransitionRecord = record } var preparedRuntime: PreparedRuntime? if let previousRuntime { - await performPublishedRuntimeSemanticStop() + await performPublishedRuntimeSemanticStop(record: record) await previousRuntime.handle.closeAdmission() await closeAppServerRuntime( previousRuntime, - purpose: .restartSameAccount + purpose: .restartSameAccount, + record: record ) } guard isCurrentTransition(generation, purpose: .restartSameAccount) else { @@ -445,7 +801,8 @@ public final class CodexReviewStore { await closeStaleRuntime( runtime, mcpServerWasPrepared: false, - purpose: .restartSameAccount + purpose: .restartSameAccount, + record: record ) return } @@ -455,7 +812,8 @@ public final class CodexReviewStore { await closeStaleRuntime( runtime, mcpServerWasPrepared: false, - purpose: .restartSameAccount + purpose: .restartSameAccount, + record: record ) return } @@ -472,7 +830,8 @@ public final class CodexReviewStore { await closeStaleRuntime( preparedRuntime, mcpServerWasPrepared: false, - purpose: .restartSameAccount + purpose: .restartSameAccount, + record: record ) } guard isCurrentTransition(generation, purpose: .restartSameAccount) else { @@ -492,7 +851,10 @@ public final class CodexReviewStore { private func isCurrentAcquisition( _ generation: ReviewRuntimeGeneration ) -> Bool { - guard case .acquiring(let currentGeneration, _) = runtimeState else { + guard case .open = lifetimeState else { + return false + } + guard case .acquiring(let currentGeneration, _, _) = runtimeState else { return false } return currentGeneration == generation @@ -502,9 +864,14 @@ public final class CodexReviewStore { _ generation: ReviewRuntimeGeneration, purpose: ReviewRuntimeTransitionPurpose ) -> Bool { + guard case .open = lifetimeState else { + return false + } guard case .transitioning( let currentGeneration, let currentPurpose, + _, + _, _ ) = runtimeState else { return false @@ -513,7 +880,13 @@ public final class CodexReviewStore { } private var currentTransitionOwnsMCPStop: Bool { - guard case .transitioning(_, let purpose, _) = runtimeState else { + switch lifetimeState { + case .closing, .closed: + return true + case .open: + break + } + guard case .transitioning(_, let purpose, _, _, _) = runtimeState else { return false } return purpose == .stop || purpose == .applicationClose @@ -522,33 +895,46 @@ public final class CodexReviewStore { private func closeStaleRuntime( _ runtime: PreparedRuntime, mcpServerWasPrepared: Bool, - purpose: ReviewRuntimeTransitionPurpose + purpose: ReviewRuntimeTransitionPurpose, + record: ReviewRuntimeTransitionRecord ) async { await runtime.handle.closeAdmission() - do { - try await runtime.handle.close(purpose: purpose) - } catch { - writeDiagnosticsIfNeeded() - } - do { - try await runtime.handle.waitUntilClosed() - } catch { + let result = await runtime.closeRecord.closeAndWait( + handle: runtime.handle, + purpose: purpose + ) + record.record(contentsOf: runtime.closeRecord.consumeFailures()) + if result.failures.isEmpty == false { writeDiagnosticsIfNeeded() } if mcpServerWasPrepared, currentTransitionOwnsMCPStop == false { - await stopPreparedMCPServer() + await stopPreparedMCPServer(record: record) } } - private func stopPreparedMCPServer() async { + private func stopPreparedMCPServer( + record: ReviewRuntimeTransitionRecord + ) async { + var stopFailureWasRecorded = false do { try await backend.mcpServerLifecycle.stop() } catch { + stopFailureWasRecorded = true + record.record( + error, + fallback: .mcpServer(error.localizedDescription) + ) writeDiagnosticsIfNeeded() } do { try await backend.mcpServerLifecycle.waitUntilStopped() } catch { + if stopFailureWasRecorded == false { + record.record( + error, + fallback: .mcpServer(error.localizedDescription) + ) + } writeDiagnosticsIfNeeded() } } @@ -607,19 +993,27 @@ public final class CodexReviewStore { } public func refreshAuthentication() async { - await backend.refreshAuth(auth: auth) + await performAdmittedStoreCommand { + await backend.refreshAuth(auth: auth) + } } public func signIn() async { - await backend.signIn(auth: auth) + await performAdmittedStoreCommand { + await backend.signIn(auth: auth) + } } public func addAccount() async { - await backend.addAccount(auth: auth) + await performAdmittedStoreCommand { + await backend.addAccount(auth: auth) + } } public func cancelAuthentication() async { - await backend.cancelAuthentication(auth: auth) + await performAdmittedStoreCommand { + await backend.cancelAuthentication(auth: auth) + } } package func performPrimaryAuthenticationAction() async { @@ -642,39 +1036,53 @@ public final class CodexReviewStore { } public func logout() async { - if auth.isAuthenticating, auth.selectedAccount == nil { - await cancelAuthentication() - return - } - do { - try await signOutActiveAccount() - } catch { - if auth.errorMessage == nil, auth.isAuthenticated { - auth.updatePhase(.failed(message: error.localizedDescription)) + await performAdmittedStoreCommand { + if auth.isAuthenticating, auth.selectedAccount == nil { + await backend.cancelAuthentication(auth: auth) + return + } + do { + try await performSignOutActiveAccount() + } catch { + if auth.errorMessage == nil, auth.isAuthenticated { + auth.updatePhase(.failed(message: error.localizedDescription)) + } } } } public func signOutActiveAccount() async throws { + try await performThrowingAdmittedStoreCommand { + try await performSignOutActiveAccount() + } + } + + private func performSignOutActiveAccount() async throws { try await backend.signOutActiveAccount(auth: auth) } package func switchAccount(_ account: CodexAccount) async throws { - guard canSwitchAccount(account) else { - return - } - let targetAccount = auth.persistedAccounts.first(where: { $0.accountKey == account.accountKey }) - if auth.persistedAccounts.contains(where: { $0.isSwitching }) || auth.selectedAccount?.isSwitching == true { - return - } - targetAccount?.updateIsSwitching(true) - defer { - targetAccount?.updateIsSwitching(false) + try await performThrowingAdmittedStoreCommand { + guard canSwitchAccount(account) else { + return + } + let targetAccount = auth.persistedAccounts.first(where: { + $0.accountKey == account.accountKey + }) + if auth.persistedAccounts.contains(where: { $0.isSwitching }) + || auth.selectedAccount?.isSwitching == true { + return + } + targetAccount?.updateIsSwitching(true) + defer { + targetAccount?.updateIsSwitching(false) + } + try await backend.switchAccount(auth: auth, accountKey: account.accountKey) } - try await backend.switchAccount(auth: auth, accountKey: account.accountKey) } package func requestSwitchAccount(_ account: CodexAccount, requiresConfirmation: Bool) { + guard case .open = lifetimeState else { return } auth.requestSwitchAccount(account, requiresConfirmation: requiresConfirmation) guard requiresConfirmation == false else { return @@ -691,6 +1099,7 @@ public final class CodexReviewStore { } package func requestSignOutActiveAccount(requiresConfirmation: Bool) { + guard case .open = lifetimeState else { return } auth.requestSignOutActiveAccount(requiresConfirmation: requiresConfirmation) guard requiresConfirmation == false else { return @@ -699,6 +1108,7 @@ public final class CodexReviewStore { } package func requestRemoveAccount(_ account: CodexAccount, requiresConfirmation: Bool) { + guard case .open = lifetimeState else { return } auth.requestRemoveAccount(account, requiresConfirmation: requiresConfirmation) guard requiresConfirmation == false else { return @@ -707,15 +1117,29 @@ public final class CodexReviewStore { } package func confirmPendingAccountAction() { + guard case .open = lifetimeState else { + return + } guard let action = auth.consumePendingAccountAction() else { return } - Task { @MainActor [weak self] in + guard let commandID = registerStoreCommand() else { + return + } + let task = Task { @MainActor [weak self] in guard let self else { return } + defer { + self.finishStoreCommand(commandID) + } do { try await self.executePendingAccountAction(action) + guard Task.isCancelled == false, + case .open = self.lifetimeState + else { + return + } if let warningMessage = self.auth.warningMessage { self.auth.presentAccountActionAlert( title: "Account Updated With Warning", @@ -723,35 +1147,54 @@ public final class CodexReviewStore { ) } } catch { + guard Task.isCancelled == false, + case .open = self.lifetimeState + else { + return + } self.auth.presentAccountActionAlert( title: action.failureTitle, message: error.localizedDescription ) } } + installOwnedStoreCommandTask(task, for: commandID) } package func cancelPendingAccountAction() { + guard case .open = lifetimeState else { return } auth.cancelPendingAccountAction() } package func dismissAccountActionAlert() { + guard case .open = lifetimeState else { return } auth.dismissAccountActionAlert() } package func removeAccount(accountKey: String) async throws { - try await backend.removeAccount(auth: auth, accountKey: accountKey) + try await performThrowingAdmittedStoreCommand { + try await backend.removeAccount(auth: auth, accountKey: accountKey) + } } package func reorderPersistedAccount(accountKey: String, toIndex: Int) async throws { - try await backend.reorderPersistedAccount(auth: auth, accountKey: accountKey, toIndex: toIndex) + try await performThrowingAdmittedStoreCommand { + try await backend.reorderPersistedAccount( + auth: auth, + accountKey: accountKey, + toIndex: toIndex + ) + } } package func refreshAccountRateLimits(accountKey: String) async { - await backend.refreshAccountRateLimits(auth: auth, accountKey: accountKey) + await performAdmittedStoreCommand { + await backend.refreshAccountRateLimits(auth: auth, accountKey: accountKey) + } } package func startStartupAuthRefresh() { + guard case .open = lifetimeState else { return } if auth.selectedAccount == nil { auth.updatePhase(.signedOut) } @@ -773,19 +1216,27 @@ public final class CodexReviewStore { } package func refreshSettings() async { - await settingsService.refresh() + await performAdmittedStoreCommand { + await settingsService.refresh() + } } package func updateSettingsModel(_ model: String) async { - await settingsService.updateModel(model) + await performAdmittedStoreCommand { + await settingsService.updateModel(model) + } } package func clearSettingsModelOverride() async { - await settingsService.clearModelOverride() + await performAdmittedStoreCommand { + await settingsService.clearModelOverride() + } } package func updateSettingsReasoningEffort(_ reasoningEffort: CodexReviewSettings.ReasoningEffort?) async { - await settingsService.updateReasoningEffort(reasoningEffort) + await performAdmittedStoreCommand { + await settingsService.updateReasoningEffort(reasoningEffort) + } } package func clearSettingsReasoningEffort() async { @@ -793,7 +1244,9 @@ public final class CodexReviewStore { } package func updateSettingsServiceTier(_ serviceTier: CodexReviewSettings.ServiceTier?) async { - await settingsService.updateServiceTier(serviceTier) + await performAdmittedStoreCommand { + await settingsService.updateServiceTier(serviceTier) + } } package func transitionToRunning(serverURL: URL?) { diff --git a/Sources/CodexReview/Store/CodexReviewStoreBackend.swift b/Sources/CodexReview/Store/CodexReviewStoreBackend.swift index ff3a3577..611fc975 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreBackend.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreBackend.swift @@ -27,7 +27,6 @@ package struct CodexReviewStoreSeed { 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) @@ -35,7 +34,7 @@ package protocol CodexReviewStoreBackend: CodexReviewSettingsBackend, Sendable { generation: ReviewRuntimeGeneration, purpose: ReviewRuntimeTransitionPurpose ) async throws -> PreparedRuntime - func stop(store: CodexReviewStore) async + func stop(store: CodexReviewStore) async throws func waitUntilStopped() async func refreshAuth(auth: CodexReviewAuthModel) async func signIn(auth: CodexReviewAuthModel) async @@ -64,9 +63,3 @@ package protocol CodexReviewStoreBackend: CodexReviewSettingsBackend, Sendable { ) async throws -> BackendReviewAttempt func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws } - -extension CodexReviewStoreBackend { - package var handlesActiveReviewStopCleanup: Bool { - false - } -} diff --git a/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift b/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift index 8d28d375..c64e6ba3 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift @@ -1,30 +1,8 @@ import Foundation -private actor RuntimeStopDetachedReviewWorkerDrainRace { - private var result: Bool? - private var continuation: CheckedContinuation? - - func finish(_ value: Bool) { - guard result == nil else { - return - } - result = value - continuation?.resume(returning: value) - continuation = nil - } - - func wait() async -> Bool { - if let result { - return result - } - return await withCheckedContinuation { continuation in - if let result { - continuation.resume(returning: result) - } else { - self.continuation = continuation - } - } - } +package struct ReviewCloseCancellationOutcome { + package let jobIDs: [String] + package let failedJobIDs: Set } extension CodexReviewStore { @@ -121,26 +99,54 @@ extension CodexReviewStore { } } - package func requestActiveReviewCancellationsForRuntimeStop( - reason: ReviewCancellation = .system(message: "Review runtime stopped.") - ) async throws -> [String] { - let activeJobIDs = orderedJobs - .filter { $0.isTerminal == false } - .map(\.id) - var firstError: (any Error)? + package func requestActiveReviewCancellationsForApplicationClose( + reason: ReviewCancellation = .system(message: "Review Store closed."), + failureLedger: ReviewCloseFailureLedger + ) async -> ReviewCloseCancellationOutcome { + let activeJobIDs = activeReviewJobIDsInRegistrationOrder + var failedJobIDs: Set = [] for jobID in activeJobIDs { do { _ = try await cancelReview(jobID: jobID, cancellation: reason) } catch { - if firstError == nil { - firstError = error + failedJobIDs.insert(jobID) + guard failureLedger.ownsForceCloseFailure(for: jobID) == false else { + continue + } + let failure = closePrimaryFailure(from: error) + if case .attemptRuntime(let runtimeFailure) = failure, + reviewCleanupFailures[jobID] == runtimeFailure { + failureLedger.recordReviewCleanupFailure( + runtimeFailure, + jobID: jobID + ) + } else { + failureLedger.record(failure) } } } - if let firstError { - throw firstError + return .init( + jobIDs: activeJobIDs, + failedJobIDs: failedJobIDs + ) + } + + private func closePrimaryFailure( + from error: any Error + ) -> ReviewClosePrimaryFailure { + if let failure = error as? ReviewInterruptRequestFailure { + return .interruptRequest(failure) } - return activeJobIDs + if let failure = error as? ReviewRuntimeCloseFailure { + return .attemptRuntime(failure) + } + if let aggregate = error as? ReviewLifecycleResourceFailureAggregate { + return .lifecycleResources(aggregate) + } + if let failure = error as? ReviewLifecycleResourceFailure { + return .lifecycleResources(.init(first: failure)) + } + return .attemptRuntime(.worker(error.localizedDescription)) } @discardableResult @@ -148,9 +154,7 @@ extension CodexReviewStore { reason: ReviewCancellation = .system(message: "Review runtime stopped."), cancelWorkers: Bool = true ) -> [String] { - let activeJobIDs = orderedJobs - .filter { $0.isTerminal == false } - .map(\.id) + let activeJobIDs = activeReviewJobIDsInRegistrationOrder guard activeJobIDs.isEmpty == false else { return [] } @@ -170,56 +174,75 @@ extension CodexReviewStore { return activeJobIDs } - package func cancelAndDetachReviewWorkersForRuntimeStop(jobIDs: [String]) { - for jobID in jobIDs { - if let task = reviewWorkerTasks.removeValue(forKey: jobID) { - task.cancel() - runtimeStopDetachedReviewWorkerTasks[jobID] = task + package func cancelAndAwaitReviewWorkersForRuntimeStop( + jobIDs: [String] + ) async { + var seenJobIDs: Set = [] + let orderedTasks = jobIDs.compactMap { jobID -> Task? in + guard seenJobIDs.insert(jobID).inserted else { + return nil } + return reviewWorkerTasks[jobID] + } + for task in orderedTasks { + task.cancel() + } + for task in orderedTasks { + await task.value } } - package func drainRuntimeStopDetachedReviewWorkers(timeout: Duration) async -> Bool { - let tasks = Array(runtimeStopDetachedReviewWorkerTasks.values) - return await drainReviewWorkerTasksForRuntimeStop(tasks, timeout: timeout) + package func awaitReviewWorkers(jobIDs: [String]) async { + var seenJobIDs: Set = [] + let orderedTasks = jobIDs.compactMap { jobID -> Task? in + guard seenJobIDs.insert(jobID).inserted else { + return nil + } + return reviewWorkerTasks[jobID] + } + for task in orderedTasks { + await task.value + } } - package func drainReviewWorkersForRuntimeStop(timeout: Duration) async -> Bool { - let tasks = Array(reviewWorkerTasks.values) + Array(runtimeStopDetachedReviewWorkerTasks.values) - return await drainReviewWorkerTasksForRuntimeStop(tasks, timeout: timeout) + package func awaitAllReviewWorkers() async { + let orderedJobIDs = reviewRegistrationOrder + let remainingJobIDs = reviewWorkerTasks.keys + .filter { orderedJobIDs.contains($0) == false } + .sorted() + await awaitReviewWorkers(jobIDs: orderedJobIDs + remainingJobIDs) } - private func drainReviewWorkerTasksForRuntimeStop( - _ tasks: [Task], - timeout: Duration - ) async -> Bool { - guard tasks.isEmpty == false else { - return true - } + package func cancelAndAwaitAllReviewWorkers() async { + let orderedJobIDs = reviewRegistrationOrder + let remainingJobIDs = reviewWorkerTasks.keys + .filter { orderedJobIDs.contains($0) == false } + .sorted() + await cancelAndAwaitReviewWorkersForRuntimeStop( + jobIDs: orderedJobIDs + remainingJobIDs + ) + } - let race = RuntimeStopDetachedReviewWorkerDrainRace() - let drainTask = Task { - for task in tasks { - await task.value - } - await race.finish(true) + package func cancelAndAwaitAllReviewMutationTasks() async { + let entries = reviewMutationTasks.sorted { $0.key < $1.key } + for (_, task) in entries { + task.cancel() } - let timeoutTask = Task { - do { - try await Task.sleep(for: timeout) - } catch { - return - } - await race.finish(false) + for (id, task) in entries { + _ = await task.result + reviewMutationTasks.removeValue(forKey: id) } + } - let didDrain = await race.wait() - if didDrain { - timeoutTask.cancel() - } else { - drainTask.cancel() + package func cancelAndAwaitOwnedStoreCommandTasks() async { + let entries = storeCommandRegistry.ownedTaskSnapshot() + for (_, task) in entries { + task.cancel() + } + for (id, task) in entries { + await task.value + finishStoreCommand(id) } - return didDrain } package func terminateAllRunningJobsLocally( @@ -250,4 +273,16 @@ extension CodexReviewStore { resumeReviewWaiters(for: jobID) } } + + private var activeReviewJobIDsInRegistrationOrder: [String] { + let registered = reviewRegistrationOrder.filter { + job(id: $0)?.isTerminal == false + } + let registeredSet = Set(registered) + let missing = orderedJobs + .filter { $0.isTerminal == false && registeredSet.contains($0.id) == false } + .map(\.id) + .sorted() + return registered + missing + } } diff --git a/Sources/CodexReview/Store/CodexReviewStoreRateLimitAutoRefresh.swift b/Sources/CodexReview/Store/CodexReviewStoreRateLimitAutoRefresh.swift index 560b289c..385652f5 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreRateLimitAutoRefresh.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreRateLimitAutoRefresh.swift @@ -223,12 +223,23 @@ private struct CodexReviewStoreRateLimitAutoRefreshAccountState { @MainActor extension CodexReviewStore { package func startAccountRateLimitAutoRefresh() { + guard case .open = lifetimeState else { + return + } if accountRateLimitAutoRefreshDriver == nil { accountRateLimitAutoRefreshDriver = CodexReviewStoreRateLimitAutoRefreshDriver(store: self) } accountRateLimitAutoRefreshDriver?.start() } + package func cancelAccountRateLimitAutoRefreshAndWait() async { + guard let driver = accountRateLimitAutoRefreshDriver else { + return + } + accountRateLimitAutoRefreshDriver = nil + await driver.cancelAndWait() + } + package func accountRateLimitAutoRefreshTargets(now: Date) -> [CodexReviewStoreRateLimitAutoRefreshTarget] { CodexReviewStoreRateLimitAutoRefreshDriver.targets( accounts: auth.accounts, @@ -240,6 +251,9 @@ extension CodexReviewStore { } package func refreshDueAccountRateLimits(now: Date) { + guard case .open = lifetimeState else { + return + } startAccountRateLimitAutoRefresh() accountRateLimitAutoRefreshDriver?.refreshDueAccounts(now: now) } @@ -316,6 +330,25 @@ package final class CodexReviewStoreRateLimitAutoRefreshDriver { accountStates.removeAll(keepingCapacity: false) } + func cancelAndWait() async { + observation?.cancel() + observation = nil + let wakeUpTask = scheduledWakeUp?.task + scheduledWakeUp = nil + let refreshTasks = accountStates.keys.sorted().compactMap { + accountStates[$0]?.refreshTask + } + wakeUpTask?.cancel() + for task in refreshTasks { + task.cancel() + } + await wakeUpTask?.value + for task in refreshTasks { + await task.value + } + accountStates.removeAll(keepingCapacity: false) + } + func refreshDueAccounts(now: Date) { guard let store else { return diff --git a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift index d1dc800a..e71a988b 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift @@ -54,6 +54,37 @@ extension CodexReviewStore { sessionID: String, request: CodexReviewAPI.Start.Request ) async throws -> String { + guard case .open = lifetimeState else { + throw CodexReviewAPI.Error.io("Review Store is closed.") + } + nextReviewMutationID &+= 1 + let mutationID = nextReviewMutationID + let runtimeGeneration = runtimeState.generation + let task = Task { @MainActor [weak self] in + guard let self else { + throw CodexReviewAPI.Error.io("Review Store was released.") + } + return try await self.performBeginReview( + sessionID: sessionID, + request: request, + runtimeGeneration: runtimeGeneration + ) + } + reviewMutationTasks[mutationID] = task + defer { + reviewMutationTasks.removeValue(forKey: mutationID) + } + return try await task.value + } + + private func performBeginReview( + sessionID: String, + request: CodexReviewAPI.Start.Request, + runtimeGeneration: ReviewRuntimeGeneration + ) async throws -> String { + guard case .open = lifetimeState else { + throw CodexReviewAPI.Error.io("Review Store is closed.") + } switch runtimeState { case .acquiring, .transitioning: throw CodexReviewAPI.Error.io( @@ -91,9 +122,21 @@ extension CodexReviewStore { model: settings.effectiveModel ) let backend = self.backend + await reviewMutationPreparationForTesting?() let registered = try await admission.registerStart { admission in try await backend.startReview(startRequest, admission: admission) } + guard isReviewMutationCurrent(runtimeGeneration) else { + _ = try? await cancel( + admission: admission, + cancellation: .system(message: "Review Store closed."), + jobID: jobID + ) + _ = await registered.task.result + throw CodexReviewAPI.Error.io( + "Review runtime changed before review publication." + ) + } insertReviewJob(job) markReviewRunning(job, startedAt: createdAt) reviewAttemptOwnerships[jobID] = .initialStart(registered) @@ -105,6 +148,21 @@ extension CodexReviewStore { return jobID } + private func isReviewMutationCurrent( + _ generation: ReviewRuntimeGeneration + ) -> Bool { + guard case .open = lifetimeState else { + return false + } + switch runtimeState { + case .stopped(let currentGeneration), + .running(let currentGeneration, _, _): + return currentGeneration == generation + case .acquiring, .transitioning, .failed: + return false + } + } + private func launchReviewWorker( jobID: String, startRequest: CodexReviewBackendModel.Review.Start, @@ -176,6 +234,9 @@ extension CodexReviewStore { let failure = ReviewRuntimeCloseFailure.worker( "Review worker was cancelled before a canonical terminal." ) + if job.isTerminal == false { + retainCleanupFailure(failure, for: jobID) + } do { try await active.admission.recordStreamTerminal(.ownerCancellation) } catch { @@ -200,7 +261,8 @@ extension CodexReviewStore { markReviewFailed(job, message: error.localizedDescription) } } - applyStreamProductTerminal(streamFailure, to: job) + await reviewTerminalPublicationPreparationForTesting?() + await applyStreamProductTerminal(streamFailure, to: job) } else if job.isTerminal == false { markReviewFailed(job, message: error.localizedDescription) } @@ -208,6 +270,7 @@ extension CodexReviewStore { reviewAttemptOwnerships[jobID] = .terminal if let cleanupAttempt { + await reviewCleanupPreparationForTesting?() do { try await cleanupReview( cleanupAttempt.run, @@ -218,7 +281,6 @@ extension CodexReviewStore { } } reviewWorkerTasks.removeValue(forKey: jobID) - runtimeStopDetachedReviewWorkerTasks.removeValue(forKey: jobID) if case .terminal = reviewAttemptOwnerships[jobID] { reviewAttemptOwnerships.removeValue(forKey: jobID) } @@ -475,7 +537,8 @@ extension CodexReviewStore { case .initialStart(let start): let resolution = try await cancel( admission: start.admission, - cancellation: cancellation + cancellation: cancellation, + jobID: jobID ) if case .localCancellation = resolution.terminal, job.isTerminal == false { @@ -491,7 +554,8 @@ extension CodexReviewStore { case .active(let active): let resolution = try await cancel( admission: active.admission, - cancellation: cancellation + cancellation: cancellation, + jobID: jobID ) try commitAcknowledgedForcedCancellationIfNeeded( resolution, @@ -506,7 +570,8 @@ extension CodexReviewStore { case .resolvingRecovery(let active): _ = try await cancel( admission: active.admission, - cancellation: cancellation + cancellation: cancellation, + jobID: jobID ) if case .resolvingRecovery(let current) = reviewAttemptOwnerships[jobID], sameAttempt(current, active), @@ -542,7 +607,8 @@ extension CodexReviewStore { case .replacementStart(_, let start): let resolution = try await cancel( admission: start.admission, - cancellation: cancellation + cancellation: cancellation, + jobID: jobID ) if case .localCancellation = resolution.terminal, job.isTerminal == false { @@ -562,7 +628,8 @@ extension CodexReviewStore { private func cancel( admission: ReviewStartAdmission, - cancellation: ReviewCancellation + cancellation: ReviewCancellation, + jobID: String ) async throws -> ReviewAttemptCancellationResolution { let backend = self.backend return try await admission.cancel( @@ -570,12 +637,87 @@ extension CodexReviewStore { interrupt: { run, reason in try await backend.interruptReview(run, reason: reason) }, - forceClose: { - try await backend.forceCloseReviewConnection() + forceClose: { @MainActor [weak self] in + guard let self else { + throw ReviewRuntimeCloseFailure.connection( + "Review Store was released before runtime force-close." + ) + } + try await self.forceCloseCurrentRuntimeForAttempt(jobID: jobID) } ) } + private func forceCloseCurrentRuntimeForAttempt(jobID: String) async throws { + let generation: ReviewRuntimeGeneration + let runtime: PreparedRuntime + let mcpGeneration: MCPServerGeneration? + let record: ReviewRuntimeTransitionRecord + switch runtimeState { + case .running( + let runningGeneration, + let runningRuntime, + let runningMCPGeneration + ): + generation = runningGeneration + runtime = runningRuntime + mcpGeneration = runningMCPGeneration + record = ReviewRuntimeTransitionRecord() + case .transitioning( + let transitionGeneration, + _, + _, + let transitionRecord, + let sourceRuntime? + ): + generation = transitionGeneration + runtime = sourceRuntime + mcpGeneration = nil + record = transitionRecord + case .stopped, .acquiring, .failed, .transitioning: + throw ReviewRuntimeCloseFailure.connection( + "Review runtime is not running." + ) + } + await runtime.handle.closeAdmission() + let closeResult = await runtime.closeRecord.closeAndWait( + handle: runtime.handle, + purpose: .recoveryReplacement + ) + let consumedFailures = runtime.closeRecord.consumeFailures() + if let applicationCloseFailureLedger { + applicationCloseFailureLedger.recordForceCloseFailures( + consumedFailures, + jobID: jobID + ) + } else { + record.recordForceCloseFailures(consumedFailures, jobID: jobID) + } + if let mcpGeneration, closeResult.installedClose { + if applicationCloseFailureLedger == nil { + lastRuntimeTransitionRecord = record + } + runtimeState = .failed( + generation: generation.successor(), + retainedMCPGeneration: mcpGeneration, + serverURL: serverURL + ) + } + await runtimeForceCloseReceiptRecordedForTesting?() + if let firstFailure = closeResult.failures.first { + switch firstFailure { + case .attemptRuntime(let failure): + throw failure + case .lifecycleResources(let failure): + throw failure + case .interruptRequest(let failure): + throw ReviewRuntimeCloseFailure.connection(failure.localizedDescription) + case .persistence(let failure): + throw ReviewRuntimeCloseFailure.connection(failure.localizedDescription) + } + } + } + private func suppressRecoverySuccessor( job: CodexReviewJob, cancellation: ReviewCancellation @@ -603,7 +745,8 @@ extension CodexReviewStore { case .replacementStart(_, let start): let resolution = try await cancel( admission: start.admission, - cancellation: cancellation + cancellation: cancellation, + jobID: job.id ) if case .localCancellation = resolution.terminal, job.isTerminal == false { @@ -614,7 +757,8 @@ extension CodexReviewStore { case .active(let active): let resolution = try await cancel( admission: active.admission, - cancellation: cancellation + cancellation: cancellation, + jobID: job.id ) try commitAcknowledgedForcedCancellationIfNeeded( resolution, @@ -625,7 +769,8 @@ extension CodexReviewStore { case .initialStart(let start): let resolution = try await cancel( admission: start.admission, - cancellation: cancellation + cancellation: cancellation, + jobID: job.id ) if case .localCancellation = resolution.terminal, job.isTerminal == false { @@ -696,7 +841,6 @@ extension CodexReviewStore { private func removeTerminalOwnershipWithoutWorker(jobID: String) { guard reviewWorkerTasks[jobID] == nil, - runtimeStopDetachedReviewWorkerTasks[jobID] == nil, case .terminal = reviewAttemptOwnerships[jobID] else { return @@ -781,6 +925,9 @@ extension CodexReviewStore { workspaces.insert(workspace) } jobs.insert(job) + if reviewRegistrationOrder.contains(job.id) == false { + reviewRegistrationOrder.append(job.id) + } writeDiagnosticsIfNeeded() } @@ -931,11 +1078,12 @@ extension CodexReviewStore { activeEventSubscriptionID = nil continue } + await reviewTerminalPublicationPreparationForTesting?() if let productTerminal = await routed.active.admission .terminalCancellationProductTerminal(for: failure) { try applyRecoveryProductTerminal(productTerminal, to: job) } else { - applyStreamProductTerminal(failure, to: job) + await applyStreamProductTerminal(failure, to: job) } return .init(cleanupAttempt: routed.active) case .reviewEventsFailed(let failedRun): @@ -960,11 +1108,12 @@ extension CodexReviewStore { continue } try await routed.active.admission.recordStreamTerminal(failedRun.failure) + await reviewTerminalPublicationPreparationForTesting?() if let productTerminal = await routed.active.admission .terminalCancellationProductTerminal(for: failedRun.failure) { try applyRecoveryProductTerminal(productTerminal, to: job) } else { - applyStreamProductTerminal(failedRun.failure, to: job) + await applyStreamProductTerminal(failedRun.failure, to: job) } return .init(cleanupAttempt: routed.active) case .recoveryBarrierResolved(let resolution): @@ -1127,11 +1276,12 @@ extension CodexReviewStore { message: ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription )) try await active.admission.recordStreamTerminal(failure) + await reviewTerminalPublicationPreparationForTesting?() if let productTerminal = await active.admission .terminalCancellationProductTerminal(for: failure) { try applyRecoveryProductTerminal(productTerminal, to: job) } else { - applyStreamProductTerminal(failure, to: job) + await applyStreamProductTerminal(failure, to: job) } return .init(cleanupAttempt: active) } @@ -1171,7 +1321,7 @@ extension CodexReviewStore { private func applyStreamProductTerminal( _ failure: ReviewAttemptStreamFailure, to job: CodexReviewJob - ) { + ) async { switch failure { case .process: markReviewInterrupted(job, cause: .previousProcessExit) @@ -1482,6 +1632,18 @@ extension CodexReviewStore { } } + package func finishAllReviewWaitersForStoreClose() async { + let waiters = reviewTerminalWaiters.values.flatMap { $0 } + reviewTerminalWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.timeoutTask?.cancel() + waiter.continuation.resume() + } + for waiter in waiters { + await waiter.timeoutTask?.value + } + } + private func resumeReviewWaiter(jobID: String, waiterID: UUID) { guard var waiters = reviewTerminalWaiters[jobID], let index = waiters.firstIndex(where: { $0.id == waiterID }) diff --git a/Sources/CodexReview/Store/CodexReviewStoreTesting.swift b/Sources/CodexReview/Store/CodexReviewStoreTesting.swift index 6aa63b48..cf42424b 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreTesting.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreTesting.swift @@ -13,6 +13,30 @@ extension CodexReviewStore { set { requestCancellationDelay = newValue } } + package func setReviewMutationPreparationForTesting( + _ operation: (@MainActor @Sendable () async -> Void)? + ) { + reviewMutationPreparationForTesting = operation + } + + package func setReviewCleanupPreparationForTesting( + _ operation: (@MainActor @Sendable () async -> Void)? + ) { + reviewCleanupPreparationForTesting = operation + } + + package func setReviewTerminalPublicationPreparationForTesting( + _ operation: (@MainActor @Sendable () async -> Void)? + ) { + reviewTerminalPublicationPreparationForTesting = operation + } + + package func setRuntimeForceCloseReceiptRecordedForTesting( + _ operation: (@MainActor @Sendable () async -> Void)? + ) { + runtimeForceCloseReceiptRecordedForTesting = operation + } + package func loadForTesting( serverState: CodexReviewServerState, authPhase: CodexReviewAuthModel.Phase = .signedOut, @@ -77,6 +101,7 @@ extension CodexReviewStore { job.sortOrder = Double(workspaceJobs.count - index - 1) } self.jobs = Set(resolvedJobs) + reviewRegistrationOrder = resolvedJobs.map(\.id) if let settingsSnapshot { settings.loadForTesting(snapshot: settingsSnapshot) } @@ -96,7 +121,7 @@ extension CodexReviewStore { ) } } - let tasks = Array(reviewWorkerTasks.values) + Array(runtimeStopDetachedReviewWorkerTasks.values) + let tasks = Array(reviewWorkerTasks.values) for task in tasks { task.cancel() } @@ -105,7 +130,6 @@ extension CodexReviewStore { } reviewWorkerTasks.removeAll(keepingCapacity: false) - runtimeStopDetachedReviewWorkerTasks.removeAll(keepingCapacity: false) reviewCleanupFailures.removeAll(keepingCapacity: false) reviewAttemptOwnerships.removeAll(keepingCapacity: false) diff --git a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift index 584ce098..ebd53646 100644 --- a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift +++ b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift @@ -16,14 +16,21 @@ private func makeAppServerReviewAttemptID() -> String { } package struct AppServerRuntimeOwnerLifecycleHandle: Sendable { + private let closeAdmissionOperation: @Sendable () async -> Void private let closeAndWaitOperation: @Sendable () async throws -> Void fileprivate init( + closeAdmissionOperation: @escaping @Sendable () async -> Void, closeAndWaitOperation: @escaping @Sendable () async throws -> Void ) { + self.closeAdmissionOperation = closeAdmissionOperation self.closeAndWaitOperation = closeAndWaitOperation } + package func closeAdmission() async { + await closeAdmissionOperation() + } + package func closeAndWait() async throws { try await closeAndWaitOperation() } @@ -191,9 +198,18 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } package nonisolated var runtimeOwnerLifecycleHandle: AppServerRuntimeOwnerLifecycleHandle { - AppServerRuntimeOwnerLifecycleHandle { [self] in - try await closeFromRuntimeOwnerAndWait() - } + AppServerRuntimeOwnerLifecycleHandle( + closeAdmissionOperation: { [self] in + await closeAdmissionFromRuntimeOwner() + }, + closeAndWaitOperation: { [self] in + try await closeFromRuntimeOwnerAndWait() + } + ) + } + + private func closeAdmissionFromRuntimeOwner() { + reviewOperationRegistry.closeAdmission() } private func admitReviewOperation() throws -> AdmittedReviewOperationID { @@ -576,7 +592,7 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { let closeTask: Task switch lifecycleState { case .open: - reviewOperationRegistry.closeAdmission() + closeAdmissionFromRuntimeOwner() let client = client let task = Task { let clientCloseTask = Task, Never> { @@ -597,7 +613,9 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { for session in ownedLifecycle.sessions { await session.finish(throwing: .ownerCancellation) } - try clientCloseResult.get() + if case .failure(let error) = clientCloseResult { + throw Self.lifecycleCloseFailure(for: error) + } } lifecycleState = .closing(task) closeTask = task @@ -613,6 +631,24 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { try result.get() } + private static func lifecycleCloseFailure( + for error: any Error + ) -> ReviewLifecycleResourceFailureAggregate { + if let aggregate = error as? ReviewLifecycleResourceFailureAggregate { + return aggregate + } + if let failure = error as? ReviewLifecycleResourceFailure { + return .init(first: failure) + } + let resourceFailure: ReviewLifecycleResourceFailure + if case .process(let message) = reviewRuntimeCloseFailure(for: error) { + resourceFailure = .process(message) + } else { + resourceFailure = .client(error.localizedDescription) + } + return .init(first: resourceFailure) + } + package static func reviewRuntimeCloseFailure( for error: any Error ) -> ReviewRuntimeCloseFailure { diff --git a/Sources/CodexReviewHost/CodexReviewHost.swift b/Sources/CodexReviewHost/CodexReviewHost.swift index 44d495fa..1aa48534 100644 --- a/Sources/CodexReviewHost/CodexReviewHost.swift +++ b/Sources/CodexReviewHost/CodexReviewHost.swift @@ -7,7 +7,6 @@ import CodexReviewMCPServer package final class CodexReviewHost { package let store: CodexReviewStore package let mcpServer: CodexReviewMCPServer - private let shutdown: @Sendable () async throws -> Void private var endpoint: URL? package init( @@ -15,12 +14,16 @@ package final class CodexReviewHost { clock: CodexReviewClock = .init(), idGenerator: CodexReviewIDGenerator = .init(), endpoint: URL? = nil, - shutdown: @escaping @Sendable () async throws -> Void = {} + closeAdmission: @escaping @MainActor @Sendable () async -> Void = {}, + shutdown: @escaping @MainActor @Sendable () async throws -> Void = {} ) { - self.shutdown = shutdown self.endpoint = endpoint let store = CodexReviewStore( - backend: DirectCodexReviewStoreBackend(backend: backend), + backend: DirectCodexReviewStoreBackend( + backend: backend, + closeAdmission: closeAdmission, + shutdown: shutdown + ), clock: clock, idGenerator: idGenerator ) @@ -37,8 +40,11 @@ package final class CodexReviewHost { self.init( backend: backend, endpoint: endpoint, + closeAdmission: { + await backend.runtimeOwnerLifecycleHandle.closeAdmission() + }, shutdown: { - try await client.close() + try await backend.runtimeOwnerLifecycleHandle.closeAndWait() } ) } @@ -47,13 +53,18 @@ package final class CodexReviewHost { if let endpoint { self.endpoint = endpoint } + await store.start() + guard case .open = store.lifetimeState, + case .running = store.runtimeState + else { + return + } store.transitionToRunning(serverURL: self.endpoint) await store.refreshSettings() } package func stop() async throws { await store.stop() - try await shutdown() } } @@ -62,6 +73,8 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { let seed = CodexReviewStoreSeed() let mcpServerLifecycle: any MCPServerLifecycleOwner = NoMCPServerLifecycleOwner() private let backend: any CodexReviewBackend + private let closeAdmission: @MainActor @Sendable () async -> Void + private let shutdown: @MainActor @Sendable () async throws -> Void private var currentSettingsSnapshot = CodexReviewSettings.Snapshot() private var loginChallenge: CodexReviewBackendModel.Login.Challenge? private var active = false @@ -74,8 +87,14 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { currentSettingsSnapshot } - init(backend: any CodexReviewBackend) { + init( + backend: any CodexReviewBackend, + closeAdmission: @escaping @MainActor @Sendable () async -> Void, + shutdown: @escaping @MainActor @Sendable () async throws -> Void + ) { self.backend = backend + self.closeAdmission = closeAdmission + self.shutdown = shutdown } func attachStore(_: CodexReviewStore) {} @@ -84,19 +103,57 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { generation _: ReviewRuntimeGeneration, purpose _: ReviewRuntimeTransitionPurpose ) async throws -> PreparedRuntime { - let authentication = try await backend.readAuth() - let settings = try await Self.monitorSettings(from: backend.readSettings()) + let shutdown = shutdown let handle = DirectRuntimeLifecycleHandle( onActivate: { [weak self] in self?.active = true }, - onClose: { [weak self] in self?.active = false } - ) - return .init( - snapshot: .init( - authentication: authentication, - settings: settings - ), - handle: handle + onCloseAdmission: closeAdmission, + onClose: { [weak self] in + defer { self?.active = false } + try await shutdown() + } ) + let closeRecord = RuntimeCloseRecord() + do { + let authentication = try await backend.readAuth() + let settings = try await Self.monitorSettings(from: backend.readSettings()) + return .init( + snapshot: .init( + authentication: authentication, + settings: settings + ), + handle: handle, + closeRecord: closeRecord + ) + } catch { + await handle.closeAdmission() + _ = await closeRecord.closeAndWait( + handle: handle, + purpose: .stop + ) + let cleanupFailures = closeRecord.consumeFailures().flatMap { + failure -> [ReviewLifecycleResourceFailure] in + switch failure { + case .lifecycleResources(let aggregate): + return [aggregate.first] + aggregate.additionalInLifecycleOrder + case .attemptRuntime(let failure): + return [.client(failure.localizedDescription)] + case .interruptRequest(let failure): + return [.client(failure.localizedDescription)] + case .persistence(let failure): + return [.client(failure.localizedDescription)] + } + } + if let first = cleanupFailures.first { + throw ReviewRuntimePreparationFailure( + preparationError: error, + cleanupFailures: .init( + first: first, + additionalInLifecycleOrder: Array(cleanupFailures.dropFirst()) + ) + ) + } + throw error + } } func stop(store _: CodexReviewStore) async { @@ -336,14 +393,17 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { @MainActor private final class DirectRuntimeLifecycleHandle: RuntimeLifecycleHandle { private let onActivate: @MainActor @Sendable () -> Void - private let onClose: @MainActor @Sendable () -> Void - private var didClose = false + private let onCloseAdmission: @MainActor @Sendable () async -> Void + private let onClose: @MainActor @Sendable () async throws -> Void + private var closeTask: Task, Never>? init( onActivate: @escaping @MainActor @Sendable () -> Void, - onClose: @escaping @MainActor @Sendable () -> Void + onCloseAdmission: @escaping @MainActor @Sendable () async -> Void, + onClose: @escaping @MainActor @Sendable () async throws -> Void ) { self.onActivate = onActivate + self.onCloseAdmission = onCloseAdmission self.onClose = onClose } @@ -351,20 +411,41 @@ private final class DirectRuntimeLifecycleHandle: RuntimeLifecycleHandle { onActivate() } - func closeAdmission() async {} + func closeAdmission() async { + await onCloseAdmission() + } func close(purpose _: ReviewRuntimeTransitionPurpose) async throws { - guard didClose == false else { return } - didClose = true - onClose() + let task: Task, Never> + if let closeTask { + task = closeTask + } else { + let onClose = onClose + let newTask = Task, Never> { @MainActor in + do { + try await onClose() + return .success(()) + } catch let aggregate as ReviewLifecycleResourceFailureAggregate { + return .failure(aggregate) + } catch let failure as ReviewLifecycleResourceFailure { + return .failure(.init(first: failure)) + } catch { + return .failure(.init(first: .client(error.localizedDescription))) + } + } + closeTask = newTask + task = newTask + } + try await task.value.get() } func waitUntilClosed() async throws { - guard didClose else { + guard let closeTask else { throw ReviewLifecycleResourceFailure.client( "Direct runtime wait began before close." ) } + try await closeTask.value.get() } } diff --git a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift index 0578ac59..9e1598fc 100644 --- a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift +++ b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift @@ -14,11 +14,16 @@ private let defaultExternalURLOpener: ExternalURLOpener = { url in private struct PendingLoginRuntimeCleanup { var client: AppServerClient? + var lifecycle: AppServerRuntimeOwnerLifecycleHandle? var codexHomeURL: URL? var authenticationSession: (any CodexReviewNativeAuthentication.WebSession)? + var authenticationTask: Task? + var notificationTask: Task? var isEmpty: Bool { - client == nil && codexHomeURL == nil && authenticationSession == nil + client == nil && lifecycle == nil && codexHomeURL == nil + && authenticationSession == nil && authenticationTask == nil + && notificationTask == nil } } @@ -56,7 +61,7 @@ package protocol CodexReviewMCPHTTPServing: AnyObject, Sendable { func start() async throws func closeAdmission() async func waitForAdmittedHandlers() async - func stop() async + func stop() async throws } extension CodexReviewMCPHTTPServing { @@ -89,7 +94,6 @@ public extension CodexReviewStore { mcpPortOwnerResolver: CodexReviewMCPPortOwnerResolver? = nil, mcpHTTPServerBindChecker: CodexReviewMCPHTTPServerBindChecker? = nil, mcpLifecycleCallObserver: CodexReviewMCPLifecycleCallObserver? = nil, - shutdownCleanupTimeout: Duration = .seconds(2), networkMonitor: any CodexReviewNetworkMonitoring = SystemCodexReviewNetworkMonitor(), networkRecoveryPolicy: CodexReviewNetworkRecoveryPolicy = .default, reviewRuntimeClosePolicy: ReviewRuntimeClosePolicy = .production, @@ -104,7 +108,6 @@ public extension CodexReviewStore { mcpPortOwnerResolver: mcpPortOwnerResolver, mcpHTTPServerBindChecker: mcpHTTPServerBindChecker, mcpLifecycleCallObserver: mcpLifecycleCallObserver, - shutdownCleanupTimeout: shutdownCleanupTimeout, networkMonitor: networkMonitor, networkRecoveryPolicy: networkRecoveryPolicy, reviewRuntimeClosePolicy: reviewRuntimeClosePolicy, @@ -125,7 +128,6 @@ public extension CodexReviewStore { mcpPortOwnerResolver: CodexReviewMCPPortOwnerResolver? = nil, mcpHTTPServerBindChecker: CodexReviewMCPHTTPServerBindChecker? = nil, mcpLifecycleCallObserver: CodexReviewMCPLifecycleCallObserver? = nil, - shutdownCleanupTimeout: Duration = .seconds(2), networkMonitor: any CodexReviewNetworkMonitoring = SystemCodexReviewNetworkMonitor(), networkRecoveryPolicy: CodexReviewNetworkRecoveryPolicy = .default, reviewRuntimeClosePolicy: ReviewRuntimeClosePolicy = .production, @@ -142,7 +144,6 @@ public extension CodexReviewStore { mcpPortOwnerResolver: mcpPortOwnerResolver, mcpHTTPServerBindChecker: mcpHTTPServerBindChecker, mcpLifecycleCallObserver: mcpLifecycleCallObserver, - shutdownCleanupTimeout: shutdownCleanupTimeout, appServerRuntimeFactory: { codexHomeURL in let client = AppServerClient(transport: try await transportFactory(codexHomeURL)) return .init( @@ -179,8 +180,14 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private var isWaitingForLoginAccountUpdate = false private var activeAuthenticationSession: (any CodexReviewNativeAuthentication.WebSession)? private var authenticationTask: Task? + private var authenticationTaskID: UInt64? + private var nextAuthenticationTaskID: UInt64 = 0 + private var retiredAuthenticationTasks: [Task] = [] private var authNotificationTask: Task? + private var retiredAuthNotificationTasks: [Task] = [] + private var pendingLifecycleFailures: [ReviewLifecycleResourceFailure] = [] private var loginNotificationTask: Task? + private var retiredLoginNotificationTasks: [Task] = [] private var settingsSnapshot = CodexReviewSettings.Snapshot() private let codexHomeURL: URL private let nativeAuthenticationConfiguration: CodexReviewNativeAuthentication.Configuration? @@ -188,7 +195,6 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private let externalURLOpener: ExternalURLOpener private let mcpLifecycleOwner: LiveMCPServerLifecycleOwner private let appServerRuntimeFactory: AppServerRuntimeFactory - private let shutdownCleanupTimeout: Duration private weak var attachedStore: CodexReviewStore? init( @@ -206,7 +212,6 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { mcpPortOwnerResolver: CodexReviewMCPPortOwnerResolver? = nil, mcpHTTPServerBindChecker: CodexReviewMCPHTTPServerBindChecker? = nil, mcpLifecycleCallObserver: CodexReviewMCPLifecycleCallObserver? = nil, - shutdownCleanupTimeout: Duration = .seconds(2), appServerRuntimeFactory: AppServerRuntimeFactory? = nil ) { let runtimePreferences = runtimePreferences.normalized @@ -231,7 +236,6 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { bindChecker: resolvedBindChecker, lifecycleCallObserver: mcpLifecycleCallObserver ) - self.shutdownCleanupTimeout = shutdownCleanupTimeout self.appServerRuntimeFactory = appServerRuntimeFactory ?? Self.makeAppServerRuntimeFactory( codexExecutablePath: runtimePreferences.codexExecutablePath ) @@ -253,10 +257,6 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { mcpLifecycleOwner } - var handlesActiveReviewStopCleanup: Bool { - true - } - var initialSettingsSnapshot: CodexReviewSettings.Snapshot { settingsSnapshot } @@ -398,7 +398,26 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { logger.info("Review runtime prepared") return PreparedRuntime(snapshot: handle.snapshot, handle: handle) } catch { - await closeClientAfterFailure(runtime.client) + let lifecycle = runtime.backend.runtimeOwnerLifecycleHandle + await lifecycle.closeAdmission() + do { + try await lifecycle.closeAndWait() + } catch let cleanupFailures as ReviewLifecycleResourceFailureAggregate { + throw ReviewRuntimePreparationFailure( + preparationError: error, + cleanupFailures: cleanupFailures + ) + } catch let cleanupFailure as ReviewLifecycleResourceFailure { + throw ReviewRuntimePreparationFailure( + preparationError: error, + cleanupFailures: .init(first: cleanupFailure) + ) + } catch { + throw ReviewRuntimePreparationFailure( + preparationError: error, + cleanupFailures: .init(first: .client(error.localizedDescription)) + ) + } throw error } } @@ -435,7 +454,6 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { activeRuntimeHandle = nil acceptsRuntimeRequests = false client = nil - appServerBackend = nil let task = authNotificationTask authNotificationTask = nil task?.cancel() @@ -449,44 +467,45 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { acceptsRuntimeRequests = false } - private func cancelActiveReviewsForRuntimeTeardown( - store: CodexReviewStore, - reason: ReviewCancellation, - timeoutWarning: String - ) async { - var cancellationFailure: (any Error)? - do { - _ = try await store.requestActiveReviewCancellationsForRuntimeStop(reason: reason) - } catch { - cancellationFailure = error - logger.error("Failed to cancel active reviews during runtime teardown: \(error.localizedDescription, privacy: .public)") - } - let didDrainReviewWorkers = await store.drainReviewWorkersForRuntimeStop( - timeout: shutdownCleanupTimeout - ) - if cancellationFailure != nil || didDrainReviewWorkers == false { - logger.warning("\(timeoutWarning, privacy: .public)") - } - } - - func stop(store: CodexReviewStore) async { + func stop(store: CodexReviewStore) async throws { let appServerBackend = appServerBackend let hasRuntimeState = client != nil || appServerBackend != nil let loginCleanup = takeLoginRuntimeForCleanup() - guard hasRuntimeState || loginCleanup.isEmpty == false else { + var failures = takePendingLifecycleFailures() + guard hasRuntimeState || loginCleanup.isEmpty == false + || retiredAuthenticationTasks.isEmpty == false + || retiredLoginNotificationTasks.isEmpty == false + || retiredAuthNotificationTasks.isEmpty == false + || failures.isEmpty == false + else { return } logger.info("Stopping review runtime") - if appServerBackend != nil { - let reason = ReviewCancellation.system(message: "Review runtime stopped.") - await cancelActiveReviewsForRuntimeTeardown( - store: store, - reason: reason, - timeoutWarning: "Timed out cleaning active reviews before stopping runtime" - ) - } await cleanupLoginRuntime(loginCleanup) + failures.append(contentsOf: takePendingLifecycleFailures()) + let retiredLoginNotificationTasks = retiredLoginNotificationTasks + self.retiredLoginNotificationTasks.removeAll(keepingCapacity: false) + for task in retiredLoginNotificationTasks { + await task.value + } + let retiredAuthenticationTasks = retiredAuthenticationTasks + self.retiredAuthenticationTasks.removeAll(keepingCapacity: false) + for task in retiredAuthenticationTasks { + await task.value + } + let retiredAuthNotificationTasks = retiredAuthNotificationTasks + self.retiredAuthNotificationTasks.removeAll(keepingCapacity: false) + for task in retiredAuthNotificationTasks { + await task.value + } + failures.append(contentsOf: takePendingLifecycleFailures()) logger.info("Review runtime semantic work stopped") + if let first = failures.first { + throw ReviewLifecycleResourceFailureAggregate( + first: first, + additionalInLifecycleOrder: Array(failures.dropFirst()) + ) + } } func waitUntilStopped() async {} @@ -580,28 +599,15 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } func cancelAuthentication(auth: CodexReviewAuthModel) async { - let activeAuthenticationSession = activeAuthenticationSession - self.activeAuthenticationSession = nil - authenticationTask?.cancel() - authenticationTask = nil - loginNotificationTask?.cancel() - loginNotificationTask = nil let loginBackend = loginBackend - self.loginBackend = nil - isWaitingForLoginAccountUpdate = false - let loginClient = loginClient - self.loginClient = nil - let loginCodexHomeURL = loginCodexHomeURL - self.loginCodexHomeURL = nil - defer { - loginChallenge = nil - } - await activeAuthenticationSession?.cancel() + let loginChallenge = loginChallenge + let cleanup = takeLoginRuntimeForCleanup() + await cleanup.authenticationSession?.cancel() guard let loginBackend, let loginChallenge else { if auth.selectedAccount == nil { auth.updatePhase(.signedOut) } - await closeIsolatedLoginRuntime(client: loginClient, codexHomeURL: loginCodexHomeURL) + await cleanupLoginRuntime(cleanup) return } do { @@ -610,7 +616,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } catch { auth.updatePhase(.failed(message: error.localizedDescription)) } - await closeIsolatedLoginRuntime(client: loginClient, codexHomeURL: loginCodexHomeURL) + await cleanupLoginRuntime(cleanup) } func switchAccount(auth: CodexReviewAuthModel, accountKey: String) async throws { @@ -819,10 +825,15 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { nativeAuthenticationConfiguration.presentationAnchorProvider ) activeAuthenticationSession = session + nextAuthenticationTaskID &+= 1 + let taskID = nextAuthenticationTaskID + authenticationTaskID = taskID authenticationTask = Task { @MainActor [weak self, weak auth] in - guard let self, let auth else { + guard let self else { return } + defer { self.finishAuthenticationTask(taskID) } + guard let auth else { return } await self.monitorAuthenticationSession( challenge: challenge, session: session, @@ -842,10 +853,8 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { let loginCodexHomeURL = loginCodexHomeURL ?? isolatedLoginCodexHomeURL self.loginCodexHomeURL = nil activeAuthenticationSession = nil - authenticationTask?.cancel() - authenticationTask = nil - loginNotificationTask?.cancel() - loginNotificationTask = nil + retireAuthenticationTask() + retireLoginNotificationTask() if let pendingLoginBackend, let pendingLoginChallenge { try? await pendingLoginBackend.cancelLogin(pendingLoginChallenge) } @@ -889,9 +898,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { self.loginClient = nil self.loginCodexHomeURL = nil activeAuthenticationSession = nil - authenticationTask = nil - loginNotificationTask?.cancel() - loginNotificationTask = nil + retireLoginNotificationTask() let account = applyAuthSnapshot( snapshot, to: auth, @@ -926,9 +933,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { self.loginClient = nil self.loginCodexHomeURL = nil activeAuthenticationSession = nil - authenticationTask = nil - loginNotificationTask?.cancel() - loginNotificationTask = nil + retireLoginNotificationTask() await closeIsolatedLoginRuntime(client: loginClient, codexHomeURL: loginCodexHomeURL) updateAuthenticationFailure( error.localizedDescription, @@ -1000,9 +1005,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { self.loginClient = nil self.loginCodexHomeURL = nil activeAuthenticationSession = nil - authenticationTask = nil - loginNotificationTask?.cancel() - loginNotificationTask = nil + retireLoginNotificationTask() auth.updatePhase(.signedOut) await closeIsolatedLoginRuntime(client: loginClient, codexHomeURL: loginCodexHomeURL) } @@ -1034,7 +1037,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { func prepareReviewRecovery( _ candidate: ReviewRecoveryCandidate ) async throws -> ReviewRecoveryHandoff { - guard acceptsRuntimeRequests, let appServerBackend else { + guard let appServerBackend else { throw CodexReviewAPI.Error.io("Review runtime is not running.") } return try await appServerBackend.prepareReviewRecovery(candidate) @@ -1056,7 +1059,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { - guard acceptsRuntimeRequests, let appServerBackend else { + guard let appServerBackend else { throw ReviewRuntimeCloseFailure.cleanup("Review runtime is not running.") } try await appServerBackend.cleanupReview(run) @@ -1161,16 +1164,11 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { return } let message = "Review runtime stopped unexpectedly: \(error.localizedDescription)" - if appServerBackend != nil { - let reason = ReviewCancellation.system(message: message) - await cancelActiveReviewsForRuntimeTeardown( - store: store, - reason: reason, - timeoutWarning: "Timed out cleaning active reviews after runtime failure" - ) - } let failedClient = client acceptsRuntimeRequests = false + if let authNotificationTask { + retiredAuthNotificationTasks.append(authNotificationTask) + } authNotificationTask = nil store.transitionToFailed(message) await cleanupLoginRuntime(loginCleanup) @@ -1199,7 +1197,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { backend: AppServerCodexReviewBackend, auth: CodexReviewAuthModel ) { - loginNotificationTask?.cancel() + retireLoginNotificationTask() loginNotificationTask = Task { @MainActor [weak self, weak auth] in guard let self, let auth else { return @@ -1256,8 +1254,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { let loginCodexHomeURL = loginCodexHomeURL let activeAuthenticationSession = activeAuthenticationSession self.activeAuthenticationSession = nil - authenticationTask?.cancel() - authenticationTask = nil + retireAuthenticationTask() await activeAuthenticationSession?.cancel() guard payload.success else { updateAuthenticationFailure( @@ -1269,8 +1266,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { isWaitingForLoginAccountUpdate = false self.loginClient = nil self.loginCodexHomeURL = nil - loginNotificationTask?.cancel() - loginNotificationTask = nil + retireLoginNotificationTask() await closeIsolatedLoginRuntime(client: loginClient, codexHomeURL: loginCodexHomeURL) return } @@ -1304,8 +1300,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { do { loginChallenge = nil self.activeAuthenticationSession = nil - authenticationTask?.cancel() - authenticationTask = nil + retireAuthenticationTask() await activeAuthenticationSession?.cancel() let account = applyAuthSnapshot( try await backend.readAuth(), @@ -1335,8 +1330,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { self.loginClient = nil self.loginCodexHomeURL = nil isWaitingForLoginAccountUpdate = false - loginNotificationTask?.cancel() - loginNotificationTask = nil + retireLoginNotificationTask() await closeIsolatedLoginRuntime(client: loginClient, codexHomeURL: loginCodexHomeURL) } @@ -1543,28 +1537,96 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private func takeLoginRuntimeForCleanup() -> PendingLoginRuntimeCleanup { loginChallenge = nil + let loginClient = loginClient + let loginLifecycle = loginClient == nil + ? nil + : loginBackend?.runtimeOwnerLifecycleHandle loginBackend = nil isWaitingForLoginAccountUpdate = false - let loginClient = loginClient self.loginClient = nil let loginCodexHomeURL = loginCodexHomeURL self.loginCodexHomeURL = nil let activeAuthenticationSession = activeAuthenticationSession self.activeAuthenticationSession = nil + let authenticationTask = authenticationTask authenticationTask?.cancel() - authenticationTask = nil + self.authenticationTask = nil + authenticationTaskID = nil + let loginNotificationTask = loginNotificationTask loginNotificationTask?.cancel() - loginNotificationTask = nil + self.loginNotificationTask = nil return .init( client: loginClient, + lifecycle: loginLifecycle, codexHomeURL: loginCodexHomeURL, - authenticationSession: activeAuthenticationSession + authenticationSession: activeAuthenticationSession, + authenticationTask: authenticationTask, + notificationTask: loginNotificationTask ) } - private func cleanupLoginRuntime(_ cleanup: PendingLoginRuntimeCleanup) async { + private func cleanupLoginRuntime( + _ cleanup: PendingLoginRuntimeCleanup + ) async { + var failures: [ReviewLifecycleResourceFailure] = [] await cleanup.authenticationSession?.cancel() - await closeIsolatedLoginRuntime(client: cleanup.client, codexHomeURL: cleanup.codexHomeURL) + await cleanup.notificationTask?.value + await cleanup.authenticationTask?.value + if let lifecycle = cleanup.lifecycle { + do { + await lifecycle.closeAdmission() + try await lifecycle.closeAndWait() + } catch let aggregate as ReviewLifecycleResourceFailureAggregate { + failures.append(aggregate.first) + failures.append(contentsOf: aggregate.additionalInLifecycleOrder) + } catch let failure as ReviewLifecycleResourceFailure { + failures.append(failure) + } catch { + failures.append(.client(error.localizedDescription)) + } + if let codexHomeURL = cleanup.codexHomeURL { + try? FileManager.default.removeItem(at: codexHomeURL) + } + } else { + await closeIsolatedLoginRuntime( + client: cleanup.client, + codexHomeURL: cleanup.codexHomeURL + ) + } + pendingLifecycleFailures.append(contentsOf: failures) + } + + private func takePendingLifecycleFailures() -> [ReviewLifecycleResourceFailure] { + let failures = pendingLifecycleFailures + pendingLifecycleFailures.removeAll(keepingCapacity: false) + return failures + } + + private func retireAuthenticationTask() { + guard let authenticationTask else { + return + } + authenticationTask.cancel() + retiredAuthenticationTasks.append(authenticationTask) + self.authenticationTask = nil + authenticationTaskID = nil + } + + private func retireLoginNotificationTask() { + guard let loginNotificationTask else { + return + } + loginNotificationTask.cancel() + retiredLoginNotificationTasks.append(loginNotificationTask) + self.loginNotificationTask = nil + } + + private func finishAuthenticationTask(_ taskID: UInt64) { + guard authenticationTaskID == taskID else { + return + } + authenticationTask = nil + authenticationTaskID = nil } private func applyRateLimits( @@ -1692,6 +1754,7 @@ private final class LiveRuntimeLifecycleHandle: RuntimeLifecycleHandle { func closeAdmission() async { owner?.closeRuntimeAdmission(self) + await backend.runtimeOwnerLifecycleHandle.closeAdmission() } func close(purpose _: ReviewRuntimeTransitionPurpose) async throws { @@ -1699,15 +1762,25 @@ private final class LiveRuntimeLifecycleHandle: RuntimeLifecycleHandle { if let closeTask { task = closeTask } else { - let client = client + let appServerLifecycle = backend.runtimeOwnerLifecycleHandle let authObservationTask = owner?.deactivateRuntime(self) let newTask = Task, Never> { @MainActor in var failures: [ReviewLifecycleResourceFailure] = [] authObservationTask?.cancel() do { - try await client.close() + try await appServerLifecycle.closeAndWait() } catch { - failures.append(.client(error.localizedDescription)) + if let aggregate = error as? ReviewLifecycleResourceFailureAggregate { + failures.append(aggregate.first) + failures.append(contentsOf: aggregate.additionalInLifecycleOrder) + } else if let failure = error as? ReviewLifecycleResourceFailure { + failures.append(failure) + } else if let failure = error as? ReviewRuntimeCloseFailure, + case .process(let message) = failure { + failures.append(.process(message)) + } else { + failures.append(.client(error.localizedDescription)) + } } await authObservationTask?.value if let first = failures.first { @@ -1750,7 +1823,7 @@ private final class LiveMCPServerLifecycleOwner: MCPServerLifecycleOwner { private typealias PreparationResult = Result private typealias ActivationResult = Result - private typealias LifecycleResult = Result + private typealias LifecycleResult = Result private enum State { case stopped @@ -2138,9 +2211,16 @@ private final class LiveMCPServerLifecycleOwner: MCPServerLifecycleOwner { _ = await activationTask.value } if let lifecycleTask { - _ = await lifecycleTask.value - } else { - await lease?.server?.stop() + return await lifecycleTask.value + } + do { + try await lease?.server?.stop() + } catch let aggregate as ReviewLifecycleResourceFailureAggregate { + return .failure(aggregate) + } catch let failure as ReviewLifecycleResourceFailure { + return .failure(.init(first: failure)) + } catch { + return .failure(.init(first: .mcpServer(error.localizedDescription))) } return .success(()) } diff --git a/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift b/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift index 6686344a..2c57d1b3 100644 --- a/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift +++ b/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift @@ -2,6 +2,7 @@ import Darwin import Foundation import MCP import OSLog +import CodexReview @preconcurrency import NIOCore @preconcurrency import NIOHTTP1 @preconcurrency import NIOPosix @@ -277,6 +278,8 @@ package actor CodexReviewMCPHTTPServer { private var sessions: [String: SessionContext] = [:] private var cleanupTask: Task? private var boundURL: URL? + private var pendingCloseFailures: [ReviewLifecycleResourceFailure] = [] + private var listenerCloseTask: Task, Never>? private let admissionRegistry = MCPHTTPAdmissionRegistry() private let handlerEntryGate = MCPHTTPHandlerEntryGate() private var admittedHandlerDrainDidBegin = false @@ -364,6 +367,8 @@ package actor CodexReviewMCPHTTPServer { } self.eventLoopGroup = group self.channel = channel + pendingCloseFailures.removeAll(keepingCapacity: false) + listenerCloseTask = nil admissionRegistry.open() admittedHandlerDrainDidBegin = false let actualPort = channel.localAddress?.port @@ -381,7 +386,7 @@ package actor CodexReviewMCPHTTPServer { } } - package func stop() async { + package func stop() async throws { await closeAdmission() cleanupTask?.cancel() let cleanupTask = cleanupTask @@ -389,20 +394,54 @@ package actor CodexReviewMCPHTTPServer { await waitForAdmittedHandlers() await cleanupTask?.value await closeAllSessions() - try? await channel?.close() - channel = nil if let eventLoopGroup { - try? await eventLoopGroup.shutdownGracefully() + do { + try await eventLoopGroup.shutdownGracefully() + } catch { + pendingCloseFailures.append(.mcpServer(error.localizedDescription)) + } } eventLoopGroup = nil boundURL = nil logger.info("MCP Streamable HTTP server stopped") + if let first = pendingCloseFailures.first { + let aggregate = ReviewLifecycleResourceFailureAggregate( + first: first, + additionalInLifecycleOrder: Array(pendingCloseFailures.dropFirst()) + ) + pendingCloseFailures.removeAll(keepingCapacity: false) + throw aggregate + } } package func closeAdmission() async { admissionRegistry.close() - try? await channel?.close() - channel = nil + guard let channel else { + return + } + let task: Task, Never> + if let listenerCloseTask { + task = listenerCloseTask + } else { + let newTask = Task, Never> { + do { + try await channel.close() + return .success(()) + } catch { + return .failure(.mcpServer(error.localizedDescription)) + } + } + listenerCloseTask = newTask + task = newTask + } + switch await task.value { + case .success: + self.channel = nil + case .failure(let failure): + if pendingCloseFailures.contains(failure) == false { + pendingCloseFailures.append(failure) + } + } } package func waitForAdmittedHandlers() async { @@ -1058,21 +1097,35 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked default: let body = response.bodyData - eventLoop.execute { - var head = HTTPResponseHead(version: version, status: status) - for (name, value) in headers { - head.headers.add(name: name, value: value) - } - if let body { - head.headers.add(name: "Content-Length", value: "\(body.count)") - } - context.write(self.wrapOutboundOut(.head(head)), promise: nil) + var head = HTTPResponseHead(version: version, status: status) + for (name, value) in headers { + head.headers.add(name: name, value: value) + } + if let body { + head.headers.add(name: "Content-Length", value: "\(body.count)") + } + do { + try await writeResponsePart( + .head(head), + context: context, + eventLoop: eventLoop + ) if let body { - var buffer = context.channel.allocator.buffer(capacity: body.count) - buffer.writeBytes(body) - context.write(self.wrapOutboundOut(.body(.byteBuffer(buffer))), promise: nil) + try await writeResponseBody( + body, + context: context, + eventLoop: eventLoop + ) } - context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil) + try await writeResponsePart( + .end(nil), + context: context, + eventLoop: eventLoop + ) + } catch { + logger.debug( + "MCP HTTP response ended during connection shutdown: \(error.localizedDescription, privacy: .public)" + ) } } } diff --git a/Sources/CodexReviewTesting/TestSupport.swift b/Sources/CodexReviewTesting/TestSupport.swift index bdc2618b..f0b0d6b9 100644 --- a/Sources/CodexReviewTesting/TestSupport.swift +++ b/Sources/CodexReviewTesting/TestSupport.swift @@ -196,6 +196,9 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { private var interruptFailureMessage: String? private var recoveryFailureMessage: String? private var cleanupFailure: ReviewRuntimeCloseFailure? + private var authReadFailureMessage: String? + private var cleanupReviewGate: AsyncGate? + private let cleanupReviewStartedGate = AsyncGate() private var interruptReviewGate: AsyncGate? private var interruptReviewWaiters: [UUID: CheckedContinuation] = [:] private var matchingInterruptReviewWaiters: [UUID: MatchingInterruptWaiter] = [:] @@ -253,6 +256,18 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { cleanupFailure = .cleanup(message) } + package func failAuthRead(message: String) { + authReadFailureMessage = message + } + + package func holdCleanupReview(with gate: AsyncGate) { + cleanupReviewGate = gate + } + + package func waitForCleanupReview() async { + await cleanupReviewStartedGate.wait() + } + package func holdInterruptReview(with gate: AsyncGate) { interruptReviewGate = gate } @@ -477,6 +492,9 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { package func readAuth() async throws -> CodexReviewBackendModel.Auth.Snapshot { commands.append(.readAuth) + if let authReadFailureMessage { + throw FakeCodexReviewBackendError(message: authReadFailureMessage) + } return auth } @@ -640,6 +658,10 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { package func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { commands.append(.cleanupReview(run)) + await cleanupReviewStartedGate.open() + if let cleanupReviewGate { + await cleanupReviewGate.waitIgnoringCancellation() + } if let cleanupFailure { throw cleanupFailure } @@ -881,12 +903,16 @@ package final class TestingRuntimeLifecycleHandle: RuntimeLifecycleHandle { package private(set) var closePurposes: [ReviewRuntimeTransitionPurpose] = [] private let onActivate: @MainActor @Sendable () -> Void - private let onClose: @MainActor @Sendable () -> Void - private var didClose = false + private let onClose: @MainActor @Sendable (ReviewRuntimeTransitionPurpose) async -> Void + private var closeAdmissionStartedGate = AsyncGate() + private var closeGate: AsyncGate? + private var closeStartedGate = AsyncGate() + private var closeFailure: ReviewLifecycleResourceFailureAggregate? + private var closeTask: Task, Never>? package init( onActivate: @escaping @MainActor @Sendable () -> Void = {}, - onClose: @escaping @MainActor @Sendable () -> Void = {} + onClose: @escaping @MainActor @Sendable (ReviewRuntimeTransitionPurpose) async -> Void = { _ in } ) { self.onActivate = onActivate self.onClose = onClose @@ -899,25 +925,63 @@ package final class TestingRuntimeLifecycleHandle: RuntimeLifecycleHandle { package func closeAdmission() async { closeAdmissionCallCount += 1 + await closeAdmissionStartedGate.open() + } + + package func waitForCloseAdmission() async { + await closeAdmissionStartedGate.wait() + } + + package func holdClose(with gate: AsyncGate) { + closeGate = gate + closeStartedGate = AsyncGate() + } + + package func waitForClose() async { + await closeStartedGate.wait() + } + + package func failClose(with failure: ReviewLifecycleResourceFailureAggregate) { + closeFailure = failure } package func close(purpose: ReviewRuntimeTransitionPurpose) async throws { closePurposes.append(purpose) closeCallCount += 1 - guard didClose == false else { - return + let task: Task, Never> + if let closeTask { + task = closeTask + } else { + let closeGate = closeGate + let closeStartedGate = closeStartedGate + let closeFailure = closeFailure + let onClose = onClose + let purpose = purpose + let newTask = Task, Never> { @MainActor in + await closeStartedGate.open() + if let closeGate { + await closeGate.waitIgnoringCancellation() + } + await onClose(purpose) + if let closeFailure { + return .failure(closeFailure) + } + return .success(()) + } + closeTask = newTask + task = newTask } - didClose = true - onClose() + try await task.value.get() } package func waitUntilClosed() async throws { waitUntilClosedCallCount += 1 - guard didClose else { + guard let closeTask else { throw ReviewLifecycleResourceFailure.client( "Testing runtime wait began before close." ) } + try await closeTask.value.get() } } @@ -927,6 +991,10 @@ package final class TestingMCPServerLifecycleOwner: MCPServerLifecycleOwner { package private(set) var activateCallCount = 0 package private(set) var stopCallCount = 0 package private(set) var waitUntilStoppedCallCount = 0 + package private(set) var closeAdmissionCallCount = 0 + package private(set) var drainAdmittedHandlersCallCount = 0 + package private(set) var closeCallCount = 0 + package private(set) var waitUntilClosedCallCount = 0 package private(set) var preparedGenerations: [MCPServerGeneration] = [] package private(set) var activatedGenerations: [MCPServerGeneration] = [] @@ -935,6 +1003,15 @@ package final class TestingMCPServerLifecycleOwner: MCPServerLifecycleOwner { private var preparationGate: AsyncGate? private let preparationStartedGate = AsyncGate() private let preparationCancellationGate = AsyncGate() + private var drainGate: AsyncGate? + private var drainStartedGate = AsyncGate() + private var closeGate: AsyncGate? + private var closeStartedGate = AsyncGate() + private var drainFailure: ReviewLifecycleResourceFailure? + private var closeFailure: ReviewLifecycleResourceFailure? + private var waitUntilClosedFailure: ReviewLifecycleResourceFailure? + private var closeTask: Task, Never>? + private var isClosed = false package init(serverURL: URL? = nil) { self.serverURL = serverURL @@ -953,6 +1030,9 @@ package final class TestingMCPServerLifecycleOwner: MCPServerLifecycleOwner { } package func prepare() async throws -> PreparedMCPServer { + guard isClosed == false else { + throw ReviewLifecycleResourceFailure.mcpServer("Testing MCP owner is closed.") + } prepareCallCount += 1 nextGeneration &+= 1 await preparationStartedGate.open() @@ -978,9 +1058,33 @@ package final class TestingMCPServerLifecycleOwner: MCPServerLifecycleOwner { return .init(serverURL: serverURL) } - package func closeAdmission() async {} + package func closeAdmission() async { + closeAdmissionCallCount += 1 + } + + package func holdHandlerDrain(with gate: AsyncGate) { + drainGate = gate + drainStartedGate = AsyncGate() + } + + package func waitForHandlerDrain() async { + await drainStartedGate.wait() + } - package func drainAdmittedHandlers() async throws {} + package func failHandlerDrain(with failure: ReviewLifecycleResourceFailure) { + drainFailure = failure + } + + package func drainAdmittedHandlers() async throws { + drainAdmittedHandlersCallCount += 1 + await drainStartedGate.open() + if let drainGate { + await drainGate.waitIgnoringCancellation() + } + if let drainFailure { + throw drainFailure + } + } package func stop() async throws { stopCallCount += 1 @@ -990,9 +1094,64 @@ package final class TestingMCPServerLifecycleOwner: MCPServerLifecycleOwner { waitUntilStoppedCallCount += 1 } - package func close() async throws {} + package func holdClose(with gate: AsyncGate) { + closeGate = gate + closeStartedGate = AsyncGate() + } + + package func waitForClose() async { + await closeStartedGate.wait() + } + + package func failClose(with failure: ReviewLifecycleResourceFailure) { + closeFailure = failure + } + + package func failWaitUntilClosed(with failure: ReviewLifecycleResourceFailure) { + waitUntilClosedFailure = failure + } + + package func close() async throws { + closeCallCount += 1 + let task: Task, Never> + if let closeTask { + task = closeTask + } else { + let closeGate = closeGate + let closeStartedGate = closeStartedGate + let closeFailure = closeFailure + let newTask = Task, Never> { @MainActor in + await closeStartedGate.open() + if let closeGate { + await closeGate.waitIgnoringCancellation() + } + if let closeFailure { + return .failure(closeFailure) + } + return .success(()) + } + closeTask = newTask + task = newTask + } + let result = await task.value + isClosed = true + try result.get() + } - package func waitUntilClosed() async throws {} + package func waitUntilClosed() async throws { + waitUntilClosedCallCount += 1 + guard let closeTask else { + throw ReviewLifecycleResourceFailure.mcpServer( + "Testing MCP completion wait began before close." + ) + } + let result = await closeTask.value + isClosed = true + try result.get() + if let waitUntilClosedFailure { + throw waitUntilClosedFailure + } + } } @MainActor @@ -1010,6 +1169,10 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { private var runtimePreparationGate: AsyncGate? private var runtimePreparationStartedGate = AsyncGate() private var runtimePreparationCancellationGate = AsyncGate() + private var runtimeCloseOperation: @MainActor @Sendable () async -> Void = {} + private var authRefreshGate: AsyncGate? + private var authRefreshStartedGate = AsyncGate() + package private(set) var authRefreshCallCount = 0 package init( reviewBackend: FakeCodexReviewBackend, @@ -1044,6 +1207,21 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { await runtimePreparationCancellationGate.wait() } + package func setRuntimeCloseOperation( + _ operation: @escaping @MainActor @Sendable () async -> Void + ) { + runtimeCloseOperation = operation + } + + package func holdAuthRefresh(with gate: AsyncGate) { + authRefreshGate = gate + authRefreshStartedGate = AsyncGate() + } + + package func waitForAuthRefresh() async { + await authRefreshStartedGate.wait() + } + package func prepareRuntime( generation _: ReviewRuntimeGeneration, purpose: ReviewRuntimeTransitionPurpose @@ -1051,7 +1229,14 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { startRequests.append(purpose == .restartSameAccount) let handle = TestingRuntimeLifecycleHandle( onActivate: { [weak self] in self?.isActive = true }, - onClose: { [weak self] in self?.isActive = false } + onClose: { [weak self] purpose in + guard let self else { return } + self.isActive = false + if purpose == .recoveryReplacement { + try? await self.reviewBackend.forceCloseReviewConnection() + } + await self.runtimeCloseOperation() + } ) lastPreparedRuntimeHandle = handle await runtimePreparationStartedGate.open() @@ -1080,6 +1265,11 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { package func waitUntilStopped() async {} package func refreshAuth(auth: CodexReviewAuthModel) async { + authRefreshCallCount += 1 + await authRefreshStartedGate.open() + if let authRefreshGate { + await authRefreshGate.waitIgnoringCancellation() + } do { let snapshot = try await reviewBackend.readAuth() let accounts = snapshot.accounts.compactMap { account -> CodexAccount? in diff --git a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift index 654e7404..d60e3bde 100644 --- a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift +++ b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift @@ -2008,6 +2008,28 @@ struct AppServerClientTests { } } + @Test func backendLifecycleAdmissionCanCloseBeforePhysicalClientClose() async throws { + let transport = FakeJSONRPCTransport() + let backend = AppServerCodexReviewBackend( + client: AppServerClient(transport: transport) + ) + let lifecycle = backend.runtimeOwnerLifecycleHandle + + await lifecycle.closeAdmission() + + await #expect(throws: JSONRPC.Error.closed) { + _ = try await backend.startReview(.init( + jobID: "job-1", + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + )) + } + #expect(await transport.closeCallCountForTesting() == 0) + + try await lifecycle.closeAndWait() + #expect(await transport.closeCallCountForTesting() == 1) + } + @Test func backendLifecycleCloseJoinsOwnedRouterAndEventSessions() async throws { let transport = DeferredNotificationCloseTransport() let backend = AppServerCodexReviewBackend( @@ -2120,10 +2142,10 @@ struct AppServerClientTests { #expect(await completion.hasCompleted() == false) await transport.finishNotificationStream(throwing: JSONRPC.Error.closed) - await #expect(throws: closeFailure) { + await #expect(throws: ReviewLifecycleResourceFailureAggregate.self) { try await close.value } - await #expect(throws: closeFailure) { + await #expect(throws: ReviewLifecycleResourceFailureAggregate.self) { try await lifecycle.closeAndWait() } #expect(await transport.recordedCloseCallCount() == 1) diff --git a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift index 6de5835b..0e805024 100644 --- a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift +++ b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift @@ -40,7 +40,7 @@ struct CodexReviewHostTests { await backend.waitForStartReview() let commands = await backend.recordedCommands() - #expect(commands.first == .readSettings) + #expect(Array(commands.prefix(2)) == [.readAuth, .readSettings]) let startReview = try #require(commands.compactMap { command -> CodexReviewBackendModel.Review.Start? in if case .startReview(let request) = command { request @@ -72,6 +72,121 @@ struct CodexReviewHostTests { #expect(host.store.auth.persistedActiveAccountKey == "review@example.com") } + @Test func hostStoreCloseOwnsDirectAppServerShutdown() async throws { + 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 host = CodexReviewHost(appServerTransport: transport) + + await host.start() + try await host.store.close() + try await host.store.close() + await host.start(endpoint: URL(string: "http://localhost:19425/mcp")) + + #expect(await transport.closeCallCountForTesting() == 1) + #expect(host.store.serverState == .stopped) + #expect(host.store.serverURL == nil) + } + + @Test func directPreparationFailureRetainsShutdownFailureForStoreClose() async throws { + let backend = FakeCodexReviewBackend() + await backend.failAuthRead(message: "auth read failed") + let host = CodexReviewHost( + backend: backend, + shutdown: { + throw ReviewLifecycleResourceFailure.client("direct shutdown failed") + } + ) + + await host.start() + let closeError = try #require(await captureStoreCloseError(host.store)) + + guard case .lifecycleResources(let lifecycle) = closeError.failures.first else { + Issue.record("Direct preparation cleanup must remain a lifecycle failure.") + return + } + #expect(lifecycle.first == .client("direct shutdown failed")) + } + + @Test func liveStoreCloseReplaysAppServerOwnerFailureOnce() 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" + ) + await transport.failClose(with: .connection("close failed")) + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + mcpHTTPServerFactory: nil, + transportFactory: { _ in transport } + ) + await store.start() + + let firstError = await captureStoreCloseError(store) + let secondError = await captureStoreCloseError(store) + + #expect(firstError?.localizedDescription == secondError?.localizedDescription) + #expect(await transport.closeCallCountForTesting() == 1) + let closeError = try #require(firstError) + guard case .lifecycleResources(let lifecycle) = closeError.failures.first else { + Issue.record("AppServer owner close failure must remain a lifecycle failure.") + return + } + guard case .client(let message) = lifecycle.first else { + Issue.record("AppServer connection close must map to client lifecycle failure.") + return + } + #expect(message.contains("close failed")) + } + + @Test func livePreparationFailureRetainsAppServerCleanupFailureForStoreClose() async throws { + let homeURL = try temporaryHome() + let transport = FakeJSONRPCTransport() + try await transport.enqueue(AppServerAPI.Initialize.Response(), for: "initialize") + await transport.enqueueFailure( + .responseError(code: -32_000, message: "auth read failed"), + for: "account/read" + ) + await transport.failClose(with: .connection("preparation cleanup failed")) + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + mcpHTTPServerFactory: nil, + transportFactory: { _ in transport } + ) + + await store.start() + let closeError = try #require(await captureStoreCloseError(store)) + + #expect(await transport.closeCallCountForTesting() == 1) + #expect(closeError.failures.additionalInLifecycleOrder.isEmpty) + guard case .lifecycleResources(let lifecycle) = closeError.failures.first else { + Issue.record("Live preparation cleanup must remain a lifecycle failure.") + return + } + guard case .client(let message) = lifecycle.first else { + Issue.record("Live preparation client close must retain its resource kind.") + return + } + #expect(message.contains("preparation cleanup failed")) + } + @Test func runtimePreferencesNormalizeInvalidValues() { let preferences = CodexReviewRuntime.Preferences( codexHomePath: " ", @@ -1856,6 +1971,8 @@ struct CodexReviewHostTests { #expect(await stopFinished.isCompleted() == false) await graceGate.open() await transport.waitForCloseCall() + #expect(await stopFinished.isCompleted() == false) + await interruptGate.open() await worker.value await stopTask.value let result = try await reviewRead.value @@ -1894,7 +2011,6 @@ struct CodexReviewHostTests { let store = CodexReviewStore.makeLiveStoreForTesting( environment: ["HOME": homeURL.path], webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, - shutdownCleanupTimeout: .seconds(1), networkMonitor: networkMonitor, networkRecoveryPolicy: .init(sleep: { _ in }), transport: transport @@ -2007,6 +2123,7 @@ struct CodexReviewHostTests { ) try await mainTransport.enqueue(AppServerAPI.Model.List.Response(data: []), for: "model/list") let loginTransport = FakeJSONRPCTransport() + await loginTransport.failClose(with: .connection("isolated login close failed")) try await loginTransport.enqueue(AppServerAPI.Initialize.Response(), for: "initialize") try await loginTransport.enqueue( AppServerAPI.Account.Login.Response.chatgpt( @@ -2063,6 +2180,16 @@ struct CodexReviewHostTests { await #expect(throws: JSONRPC.Error.closed) { _ = try await loginTransport.send(JSONRPC.Request(id: 99, method: "ping", params: Data())) } + let closeError = try #require(await captureStoreCloseError(store)) + guard case .lifecycleResources(let lifecycle) = closeError.failures.first else { + Issue.record("Retired isolated-login cleanup must reach Store close.") + return + } + guard case .client(let message) = lifecycle.first else { + Issue.record("Isolated-login connection close must remain a client failure.") + return + } + #expect(message.contains("isolated login close failed")) } @Test func liveStoreRemovingActiveAccountClearsSharedAuthAndRestartsSignedOutRuntime() async throws { @@ -2581,6 +2708,19 @@ private func failedMessage(from phase: CodexReviewAuthModel.Phase) -> String? { return message } +@MainActor +private func captureStoreCloseError(_ store: CodexReviewStore) async -> ReviewCloseError? { + do { + try await store.close() + return nil + } catch let error as ReviewCloseError { + return error + } catch { + Issue.record("Unexpected Store close error: \(error)") + return nil + } +} + @MainActor private final class MCPLifecycleCallProbe { private struct Waiter { diff --git a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift index 77564a0b..90b75b05 100644 --- a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift +++ b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift @@ -155,7 +155,7 @@ struct CodexReviewMCPHTTPServerTests { path: "/mcp" )) #expect(afterAdmissionClose.statusCode == 503) - await server.stop() + try await server.stop() } @Test func streamableHTTPClassifiesAddressInUseBindError() { @@ -1115,7 +1115,7 @@ struct CodexReviewMCPHTTPServerTests { let stopFinished = CompletionFlag() let stopTask = Task { - await server.stop() + try await server.stop() await stopFinished.complete() } await server.waitForAdmittedHandlerDrainToBeginForTesting() @@ -1128,7 +1128,7 @@ struct CodexReviewMCPHTTPServerTests { await requestGate.open() await backend.yield(.completed(summary: "Done", result: "review text")) _ = try? await requestTask.value - await stopTask.value + try await stopTask.value #expect(await stopFinished.isCompleted()) #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) == nil) @@ -1360,10 +1360,10 @@ struct CodexReviewMCPHTTPServerTests { try await server.start() do { let result = try await operation(server) - await server.stop() + try await server.stop() return result } catch { - await server.stop() + try? await server.stop() throw error } } diff --git a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift index 9142a988..0ab4ca28 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift @@ -1958,7 +1958,7 @@ struct CodexReviewStoreCommandTests { } } - @Test func runtimeStopLocalCancellationDetachesWorker() async throws { + @Test func runtimeStopLocalCancellationAwaitsWorker() async throws { let run = CodexReviewBackendModel.Review.Run( threadID: "thread-1", turnID: "turn-1", @@ -1989,10 +1989,12 @@ struct CodexReviewStoreCommandTests { #expect(store.reviewWorkerTasks["job-1"] != nil) #expect(activeAttempt(in: store, jobID: "job-1")?.run == run) - store.cancelAndDetachReviewWorkersForRuntimeStop(jobIDs: locallyCancelledJobIDs) + await store.cancelAndAwaitReviewWorkersForRuntimeStop( + jobIDs: locallyCancelledJobIDs + ) #expect(store.reviewWorkerTasks["job-1"] == nil) - #expect(activeAttempt(in: store, jobID: "job-1")?.run == run) + #expect(store.reviewAttemptOwnerships["job-1"] == nil) } } @@ -2039,7 +2041,7 @@ struct CodexReviewStoreCommandTests { } } - @Test func runtimeStopDetachesNetworkRecoveryWaitingWorker() async throws { + @Test func runtimeStopAwaitsNetworkRecoveryWaitingWorker() async throws { let run = CodexReviewBackendModel.Review.Run( threadID: "thread-1", turnID: "turn-1", @@ -2069,14 +2071,16 @@ struct CodexReviewStoreCommandTests { reason: .system(message: "Review runtime stopped."), cancelWorkers: false ) - store.cancelAndDetachReviewWorkersForRuntimeStop(jobIDs: locallyCancelledJobIDs) + await store.cancelAndAwaitReviewWorkersForRuntimeStop( + jobIDs: locallyCancelledJobIDs + ) #expect(store.reviewWorkerTasks["job-1"] == nil) - #expect(isWaitingForRecovery(in: store, jobID: "job-1")) + #expect(store.reviewAttemptOwnerships["job-1"] == nil) } } - @Test func runtimeStopCanDrainDetachedWorkerCleanup() async throws { + @Test func runtimeStopAwaitsWorkerCleanup() async throws { let run = CodexReviewBackendModel.Review.Run( threadID: "thread-1", turnID: "turn-1", @@ -2100,15 +2104,16 @@ struct CodexReviewStoreCommandTests { reason: .system(message: "Review runtime stopped."), cancelWorkers: false ) - store.cancelAndDetachReviewWorkersForRuntimeStop(jobIDs: locallyCancelledJobIDs) + await store.cancelAndAwaitReviewWorkersForRuntimeStop( + jobIDs: locallyCancelledJobIDs + ) - #expect(await store.drainRuntimeStopDetachedReviewWorkers(timeout: .seconds(2))) - #expect(store.runtimeStopDetachedReviewWorkerTasks["job-1"] == nil) + #expect(store.reviewWorkerTasks["job-1"] == nil) #expect(await backend.recordedCommands().contains(.cleanupReview(run))) } } - @Test func runtimeStopDetachLetsStartReviewReturnWhenBackendStartIsStuck() async throws { + @Test func runtimeStopAwaitJoinsStartAfterBackendUnblocks() async throws { let backend = FakeCodexReviewBackend() let startReviewGate = AsyncGate() await backend.holdStartReview(with: startReviewGate) @@ -2129,12 +2134,14 @@ struct CodexReviewStoreCommandTests { reason: .system(message: "Review runtime stopped."), cancelWorkers: false ) - store.cancelAndDetachReviewWorkersForRuntimeStop(jobIDs: locallyCancelledJobIDs) - let resultBeforeStartReviewUnblocked = try await waitForTaskValue(running, timeout: .seconds(1)) - #expect(store.reviewAttemptOwnerships["job-1"] != nil) + let drainTask = Task { @MainActor in + await store.cancelAndAwaitReviewWorkersForRuntimeStop( + jobIDs: locallyCancelledJobIDs + ) + } await startReviewGate.open() - let result = try #require(resultBeforeStartReviewUnblocked) - #expect(await store.drainRuntimeStopDetachedReviewWorkers(timeout: .seconds(2))) + await drainTask.value + let result = try await running.value #expect(locallyCancelledJobIDs == ["job-1"]) #expect(result.core.lifecycle.status == .cancelled) @@ -2447,6 +2454,7 @@ struct CodexReviewStoreCommandTests { } ) ) + await store.start() try await withStoreCommandTestCleanup(backend: backend, store: store) { async let result = store.startReview( sessionID: "session-1", diff --git a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift index 399c98fa..553987ef 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift @@ -18,7 +18,9 @@ struct CodexReviewStoreLifecycleTests { store.runtimeState = .transitioning( generation: generation.successor(), purpose: .stop, - task: transitionTask + task: transitionTask, + record: ReviewRuntimeTransitionRecord(), + sourceRuntime: nil ) await store.performRuntimeAcquisitionForTesting( @@ -148,7 +150,7 @@ struct CodexReviewStoreLifecycleTests { await backend.waitForRuntimePreparation() let ownsReplacementTask: Bool - if case .transitioning(_, .restartSameAccount, _) = store.runtimeState { + if case .transitioning(_, .restartSameAccount, _, _, _) = store.runtimeState { ownsReplacementTask = true } else { ownsReplacementTask = false @@ -226,4 +228,593 @@ struct CodexReviewStoreLifecycleTests { #expect(mcpOwner.stopCallCount == 1) #expect(mcpOwner.waitUntilStoppedCallCount == 1) } + + @Test func concurrentAndRepeatedCloseJoinOneRecordedResultAndRejectLaterMutation() async throws { + let mcpOwner = TestingMCPServerLifecycleOwner() + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend(), + mcpServerLifecycle: mcpOwner + ) + let store = CodexReviewStore.makeTestingStore(backend: backend) + await store.start() + let runtime = try #require(backend.lastPreparedRuntimeHandle) + let runtimeCloseGate = AsyncGate() + runtime.holdClose(with: runtimeCloseGate) + + let firstClose = Task { @MainActor in + try await store.close() + } + await runtime.waitForClose() + let secondClose = Task { @MainActor in + try await store.close() + } + await store.waitForCloseCallersForTesting(2) + + #expect(runtime.closeCallCount == 1) + #expect(runtime.waitUntilClosedCallCount == 0) + #expect(mcpOwner.closeAdmissionCallCount == 1) + #expect(mcpOwner.drainAdmittedHandlersCallCount == 1) + #expect(mcpOwner.closeCallCount == 1) + #expect(mcpOwner.waitUntilClosedCallCount == 1) + + await runtimeCloseGate.open() + try await firstClose.value + try await secondClose.value + try await store.close() + + #expect(runtime.closeCallCount == 1) + #expect(runtime.waitUntilClosedCallCount == 1) + #expect(runtime.closePurposes == [.applicationClose]) + guard case .closed(.success) = store.lifetimeState else { + Issue.record("Store close must record one successful result.") + return + } + + let startRequests = backend.startRequests + await store.start() + await store.restart() + await store.stop() + #expect(backend.startRequests == startRequests) + await #expect(throws: CodexReviewAPI.Error.self) { + _ = try await store.startReview( + sessionID: "session-after-close", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + } + + @Test func closeWaitsForAdmittedMCPHandlerBeforeRuntimePhysicalClose() async throws { + let mcpOwner = TestingMCPServerLifecycleOwner() + let handlerGate = AsyncGate() + mcpOwner.holdHandlerDrain(with: handlerGate) + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend(), + mcpServerLifecycle: mcpOwner + ) + let store = CodexReviewStore.makeTestingStore(backend: backend) + await store.start() + let runtime = try #require(backend.lastPreparedRuntimeHandle) + + let closeTask = Task { @MainActor in + try await store.close() + } + await mcpOwner.waitForHandlerDrain() + + #expect(mcpOwner.closeAdmissionCallCount == 1) + #expect(mcpOwner.closeCallCount == 0) + #expect(runtime.closeCallCount == 0) + + await handlerGate.open() + try await closeTask.value + + #expect(mcpOwner.closeCallCount == 1) + #expect(mcpOwner.waitUntilClosedCallCount == 1) + #expect(runtime.closeCallCount == 1) + } + + @Test func closeInvalidatesHeldRuntimePreparationAndNeverPublishesIt() 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 startTask = Task { @MainActor in + await store.start() + } + await backend.waitForRuntimePreparation() + let staleRuntime = try #require(backend.lastPreparedRuntimeHandle) + let closeTask = Task { @MainActor in + try await store.close() + } + await backend.waitForRuntimePreparationCancellation() + + #expect(staleRuntime.activateCallCount == 0) + #expect(store.serverState == .starting) + + await preparationGate.open() + try await closeTask.value + await startTask.value + + #expect(staleRuntime.activateCallCount == 0) + #expect(staleRuntime.closeAdmissionCallCount == 1) + #expect(staleRuntime.closeCallCount == 1) + #expect(staleRuntime.waitUntilClosedCallCount == 1) + #expect(store.serverState == .stopped) + #expect(store.serverURL == nil) + } + + @Test func closeAggregatesLifecycleFailuresOnceAndReplaysThem() async throws { + let mcpOwner = TestingMCPServerLifecycleOwner() + mcpOwner.failHandlerDrain(with: .mcpHandlerDrain("handler drain failed")) + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend(), + mcpServerLifecycle: mcpOwner + ) + let store = CodexReviewStore.makeTestingStore(backend: backend) + await store.start() + let runtime = try #require(backend.lastPreparedRuntimeHandle) + runtime.failClose(with: .init(first: .client("runtime close failed"))) + + let firstError = await capturedCloseError(from: store) + let secondError = await capturedCloseError(from: store) + + #expect(firstError?.localizedDescription == secondError?.localizedDescription) + #expect(runtime.closeCallCount == 1) + #expect(runtime.waitUntilClosedCallCount == 1) + #expect(mcpOwner.drainAdmittedHandlersCallCount == 1) + let failures = try #require(firstError).failures + guard case .lifecycleResources(let firstLifecycle) = failures.first else { + Issue.record("MCP handler drain must be the first close lifecycle failure.") + return + } + #expect(firstLifecycle.first == .mcpHandlerDrain("handler drain failed")) + guard case .lifecycleResources(let secondLifecycle) = failures.additionalInLifecycleOrder.first else { + Issue.record("Runtime close must follow MCP lifecycle failures.") + return + } + #expect(secondLifecycle.first == .client("runtime close failed")) + #expect(failures.additionalInLifecycleOrder.count == 1) + } + + @Test func cancellationFailureUsesRuntimeTerminalBeforeJoiningWorker() async throws { + let run = CodexReviewBackendModel.Review.Run( + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "review-thread-1" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: run) + await reviewBackend.failInterrupts(message: "interrupt rejected") + let backend = TestingCodexReviewStoreBackend(reviewBackend: reviewBackend) + backend.setRuntimeCloseOperation { + try? await reviewBackend.forceCloseReviewConnection() + } + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { "job-1" }) + ) + await store.start() + let runtime = try #require(backend.lastPreparedRuntimeHandle) + let runtimeCloseGate = AsyncGate() + runtime.holdClose(with: runtimeCloseGate) + let initial = try await store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges), + waitTimeout: .milliseconds(20) + ) + #expect(initial.core.lifecycle.status == .running) + + let closeTask = Task { @MainActor in + await capturedCloseError(from: store) + } + await runtime.waitForClose() + + let beforeRuntimeClose = try store.readReview(jobID: "job-1") + #expect(beforeRuntimeClose.core.lifecycle.status == .running) + #expect(store.reviewWorkerTasks["job-1"] != nil) + + await runtimeCloseGate.open() + let closeError = try #require(await closeTask.value) + + let terminal = try store.readReview(jobID: "job-1") + #expect(terminal.core.lifecycle.status == .failed) + #expect(terminal.core.lifecycle.cancellation == nil) + #expect(store.reviewWorkerTasks["job-1"] == nil) + guard case .interruptRequest(let failure) = closeError.failures.first else { + Issue.record("Interrupt failure must remain first in the close aggregate.") + return + } + #expect(failure.outcome == .rejected(code: nil, message: "interrupt rejected")) + } + + @Test func successfulCancellationCleanupFinishesBeforeAppServerAdmissionCloses() async throws { + let run = CodexReviewBackendModel.Review.Run( + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "review-thread-1" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: run) + let cleanupGate = AsyncGate() + await reviewBackend.holdCleanupReview(with: cleanupGate) + let mcpOwner = TestingMCPServerLifecycleOwner() + let backend = TestingCodexReviewStoreBackend( + reviewBackend: reviewBackend, + mcpServerLifecycle: mcpOwner + ) + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { "job-1" }) + ) + await store.start() + let runtime = try #require(backend.lastPreparedRuntimeHandle) + _ = try await store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges), + waitTimeout: .milliseconds(20) + ) + + let closeTask = Task { @MainActor in + try await store.close() + } + await reviewBackend.waitForInterruptReview() + await reviewBackend.yield( + .cancelled("Review Store closed."), + for: run + ) + await reviewBackend.waitForCleanupReview() + + #expect(runtime.closeAdmissionCallCount == 0) + #expect(runtime.closeCallCount == 0) + #expect(mcpOwner.drainAdmittedHandlersCallCount == 0) + + await cleanupGate.open() + try await closeTask.value + + #expect(runtime.closeAdmissionCallCount == 1) + #expect(runtime.closeCallCount == 1) + #expect(store.reviewWorkerTasks["job-1"] == nil) + } + + @Test func closeAwaitsPrePublicationReviewMutationAndPreventsBackendWrite() async throws { + let reviewBackend = FakeCodexReviewBackend() + let backend = TestingCodexReviewStoreBackend(reviewBackend: reviewBackend) + let store = CodexReviewStore.makeTestingStore(backend: backend) + await store.start() + let mutationEntered = AsyncGate() + let mutationRelease = AsyncGate() + store.setReviewMutationPreparationForTesting { + await mutationEntered.open() + await mutationRelease.waitIgnoringCancellation() + } + let startTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await mutationEntered.wait() + let closeCompletion = StoreCloseCompletion() + let closeTask = Task { @MainActor in + try await store.close() + await closeCompletion.complete() + } + await store.waitForCloseCallersForTesting(1) + + let commandsBeforeRelease = await reviewBackend.recordedCommands() + #expect(await closeCompletion.isComplete() == false) + #expect(store.reviewMutationTasks.count == 1) + #expect(commandsBeforeRelease.contains { command in + if case .startReview = command { true } else { false } + } == false) + + await mutationRelease.open() + await #expect(throws: CodexReviewAPI.Error.self) { + _ = try await startTask.value + } + try await closeTask.value + + #expect(await closeCompletion.isComplete()) + #expect(store.reviewMutationTasks.isEmpty) + #expect(store.jobs.isEmpty) + } + + @Test func closeAwaitsAdmittedAuthenticationMutationAndRejectsLaterMutation() async throws { + let authRefreshGate = AsyncGate() + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend() + ) + backend.holdAuthRefresh(with: authRefreshGate) + let store = CodexReviewStore.makeTestingStore(backend: backend) + await store.start() + let runtime = try #require(backend.lastPreparedRuntimeHandle) + + let refreshTask = Task { @MainActor in + await store.refreshAuthentication() + } + await backend.waitForAuthRefresh() + let closeCompletion = StoreCloseCompletion() + let closeTask = Task { @MainActor in + try await store.close() + await closeCompletion.complete() + } + await store.waitForCloseCallersForTesting(1) + + #expect(await closeCompletion.isComplete() == false) + #expect(runtime.closeCallCount == 0) + #expect(backend.authRefreshCallCount == 1) + + await authRefreshGate.open() + await refreshTask.value + try await closeTask.value + + #expect(await closeCompletion.isComplete()) + #expect(runtime.closeCallCount == 1) + await store.refreshAuthentication() + #expect(backend.authRefreshCallCount == 1) + } + + @Test func closeReplaysFailureFromACompletedRuntimeTransition() async throws { + let backend = TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend() + ) + let store = CodexReviewStore.makeTestingStore(backend: backend) + await store.start() + let runtime = try #require(backend.lastPreparedRuntimeHandle) + runtime.failClose(with: .init(first: .client("stop close failed"))) + + await store.stop() + let firstError = try #require(await capturedCloseError(from: store)) + let secondError = try #require(await capturedCloseError(from: store)) + + #expect(firstError.localizedDescription == secondError.localizedDescription) + guard case .lifecycleResources(let lifecycle) = firstError.failures.first else { + Issue.record("The completed transition failure must remain owned by Store close.") + return + } + #expect(lifecycle.first == .client("stop close failed")) + #expect(runtime.closeCallCount == 1) + #expect(runtime.waitUntilClosedCallCount == 1) + } + + @Test func closePreservesInterruptFailureFromCompletedStopTransition() async throws { + let run = CodexReviewBackendModel.Review.Run( + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "review-thread-1" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: run) + await reviewBackend.failInterrupts(message: "stop interrupt rejected") + let backend = TestingCodexReviewStoreBackend(reviewBackend: reviewBackend) + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { "job-1" }) + ) + await store.start() + _ = try await store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges), + waitTimeout: .milliseconds(20) + ) + + await store.stop() + let closeError = try #require(await capturedCloseError(from: store)) + + guard case .interruptRequest(let failure) = closeError.failures.first else { + Issue.record("The completed stop transition must retain its interrupt failure.") + return + } + #expect(failure.outcome == .rejected(code: nil, message: "stop interrupt rejected")) + } + + @Test func applicationGraceCloseConsumesRuntimeFailureExactlyOnce() async throws { + let run = CodexReviewBackendModel.Review.Run( + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "review-thread-1" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: run) + let backend = TestingCodexReviewStoreBackend(reviewBackend: reviewBackend) + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { "job-1" }), + reviewRuntimeClosePolicy: .init( + terminalGrace: .seconds(10), + sleep: { _ in } + ) + ) + await store.start() + let runtime = try #require(backend.lastPreparedRuntimeHandle) + runtime.failClose(with: .init(first: .client("forced runtime close failed"))) + _ = try await store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges), + waitTimeout: .milliseconds(20) + ) + + let closeError = try #require(await capturedCloseError(from: store)) + + #expect(closeError.failures.additionalInLifecycleOrder.isEmpty) + guard case .lifecycleResources(let lifecycle) = closeError.failures.first else { + Issue.record("Forced runtime close must retain its lifecycle failure.") + return + } + #expect(lifecycle.first == .client("forced runtime close failed")) + #expect(runtime.closeCallCount == 1) + #expect(runtime.waitUntilClosedCallCount == 1) + } + + @Test func closeImportsPriorGraceFailureReceiptBeforeRecancellingActiveJob() async throws { + let run = CodexReviewBackendModel.Review.Run( + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "review-thread-1" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: run) + let backend = TestingCodexReviewStoreBackend(reviewBackend: reviewBackend) + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { "job-1" }), + reviewRuntimeClosePolicy: .init( + terminalGrace: .seconds(10), + sleep: { _ in } + ) + ) + let terminalPublicationEntered = AsyncGate() + let terminalPublicationRelease = AsyncGate() + let forceCloseReceiptRecorded = AsyncGate() + store.setReviewTerminalPublicationPreparationForTesting { + await terminalPublicationEntered.open() + await terminalPublicationRelease.waitIgnoringCancellation() + } + store.setRuntimeForceCloseReceiptRecordedForTesting { + await forceCloseReceiptRecorded.open() + } + await store.start() + let runtime = try #require(backend.lastPreparedRuntimeHandle) + runtime.failClose(with: .init(first: .client("prior forced close failed"))) + _ = try await store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges), + waitTimeout: .milliseconds(20) + ) + + let cancellationTask = Task { @MainActor in + try await store.cancelReview( + jobID: "job-1", + cancellation: .mcpClient(message: "Stop") + ) + } + await forceCloseReceiptRecorded.wait() + await terminalPublicationEntered.wait() + #expect(try store.readReview(jobID: "job-1").core.lifecycle.status == .running) + + let closeTask = Task { @MainActor in + await capturedCloseError(from: store) + } + await store.waitForCloseCallersForTesting(1) + await terminalPublicationRelease.open() + _ = await cancellationTask.result + let closeError = try #require(await closeTask.value) + + #expect(closeError.failures.additionalInLifecycleOrder.isEmpty) + guard case .lifecycleResources(let lifecycle) = closeError.failures.first else { + Issue.record("The prior force-close receipt must retain lifecycle ownership.") + return + } + #expect(lifecycle.first == .client("prior forced close failed")) + #expect(runtime.closeCallCount == 1) + #expect(runtime.waitUntilClosedCallCount == 1) + } + + @Test func closeDoesNotRepeatCleanupFailureConsumedByCompletedStop() async throws { + let run = CodexReviewBackendModel.Review.Run( + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "review-thread-1" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: run) + await reviewBackend.failCleanup(message: "stop cleanup failed") + let backend = TestingCodexReviewStoreBackend(reviewBackend: reviewBackend) + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { "job-1" }) + ) + await store.start() + _ = try await store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges), + waitTimeout: .milliseconds(20) + ) + let stopTask = Task { @MainActor in + await store.stop() + } + await reviewBackend.waitForInterruptReview() + await reviewBackend.yield(.cancelled("Review runtime stopped."), for: run) + await stopTask.value + + let closeError = try #require(await capturedCloseError(from: store)) + + #expect(closeError.failures.additionalInLifecycleOrder.isEmpty) + guard case .attemptRuntime(let failure) = closeError.failures.first else { + Issue.record("Stop cleanup failure must remain an attempt-runtime failure.") + return + } + #expect(failure == .cleanup("stop cleanup failed")) + } + + @Test func naturallyTerminalWorkerFinishesCleanupBeforeRuntimePhysicalClose() async throws { + let run = CodexReviewBackendModel.Review.Run( + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "review-thread-1" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: run) + let backend = TestingCodexReviewStoreBackend(reviewBackend: reviewBackend) + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { "job-1" }) + ) + let cleanupEntered = AsyncGate() + let cleanupRelease = AsyncGate() + store.setReviewCleanupPreparationForTesting { + await cleanupEntered.open() + await cleanupRelease.waitIgnoringCancellation() + } + await store.start() + let runtime = try #require(backend.lastPreparedRuntimeHandle) + _ = try await store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges), + waitTimeout: .milliseconds(20) + ) + await reviewBackend.yield( + .completed(summary: "Done", result: "review text"), + for: run + ) + await cleanupEntered.wait() + #expect(try store.readReview(jobID: "job-1").core.lifecycle.status == .succeeded) + + let closeTask = Task { @MainActor in + try await store.close() + } + await runtime.waitForCloseAdmission() + + #expect(runtime.closeCallCount == 0) + #expect(await reviewBackend.recordedCommands().contains(.cleanupReview(run)) == false) + + await cleanupRelease.open() + try await closeTask.value + + #expect(await reviewBackend.recordedCommands().contains(.cleanupReview(run))) + #expect(runtime.closeCallCount == 1) + #expect(runtime.waitUntilClosedCallCount == 1) + } +} + +@MainActor +private func capturedCloseError(from store: CodexReviewStore) async -> ReviewCloseError? { + do { + try await store.close() + return nil + } catch let error as ReviewCloseError { + return error + } catch { + Issue.record("Unexpected close error: \(error)") + return nil + } +} + +private actor StoreCloseCompletion { + private var completeValue = false + + func complete() { + completeValue = true + } + + func isComplete() -> Bool { + completeValue + } } diff --git a/Tests/CodexReviewTests/CodexReviewStoreRateLimitAutoRefreshTests.swift b/Tests/CodexReviewTests/CodexReviewStoreRateLimitAutoRefreshTests.swift index 41042100..fff75e46 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreRateLimitAutoRefreshTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreRateLimitAutoRefreshTests.swift @@ -268,6 +268,40 @@ struct CodexReviewStoreRateLimitAutoRefreshTests { } } + @Test func storeCloseCancelsAndAwaitsHeldRateLimitRefresh() async throws { + let account = makeAccount(lastFetchAt: now.addingTimeInterval(-15 * 60)) + let backend = BlockingRateLimitRefreshBackend(account: account) + let store = CodexReviewStore.makeTestingStore(backend: backend) + store.loadForTesting( + serverState: .running, + authPhase: .signedOut, + account: account, + persistedAccounts: [account], + workspaces: [] + ) + store.refreshDueAccountRateLimits(now: now) + await backend.waitUntilRefreshStarts() + + let closeCompletion = RateLimitCloseCompletion() + let closeTask = Task { @MainActor in + try await store.close() + await closeCompletion.complete() + } + await backend.waitUntilRefreshIsCancelled() + + #expect(await closeCompletion.isComplete() == false) + #expect(backend.refreshCompletionCount == 0) + + await backend.releaseRefresh() + try await closeTask.value + + #expect(await closeCompletion.isComplete()) + #expect(backend.refreshCompletionCount == 1) + #expect(store.accountRateLimitAutoRefreshInFlightAccountKeys.isEmpty) + store.refreshDueAccountRateLimits(now: now) + #expect(backend.refreshedAccountKeys == [account.accountKey]) + } + private func makeAccount( accountKey: String? = nil, email: String = "review@example.com", @@ -347,6 +381,18 @@ private struct TestFailure: Error { } } +private actor RateLimitCloseCompletion { + private var completeValue = false + + func complete() { + completeValue = true + } + + func isComplete() -> Bool { + completeValue + } +} + private final class MutableTestClock: @unchecked Sendable { nonisolated(unsafe) var now: Date @@ -378,7 +424,9 @@ private final class NoProgressRateLimitRefreshBackend: PreviewCodexReviewStoreBa private final class BlockingRateLimitRefreshBackend: PreviewCodexReviewStoreBackend { private let startedGate = AsyncGate() private let releaseGate = AsyncGate() + private let cancellationGate = AsyncGate() private(set) var refreshedAccountKeys: [String] = [] + private(set) var refreshCompletionCount = 0 init(account: CodexAccount) { super.init(seed: .init( @@ -393,7 +441,13 @@ private final class BlockingRateLimitRefreshBackend: PreviewCodexReviewStoreBack ) async { refreshedAccountKeys.append(accountKey) await startedGate.open() - await releaseGate.wait() + let cancellationGate = cancellationGate + await withTaskCancellationHandler { + await releaseGate.waitIgnoringCancellation() + } onCancel: { + Task { await cancellationGate.open() } + } + refreshCompletionCount += 1 } func waitUntilRefreshStarts() async { @@ -403,4 +457,8 @@ private final class BlockingRateLimitRefreshBackend: PreviewCodexReviewStoreBack func releaseRefresh() async { await releaseGate.open() } + + func waitUntilRefreshIsCancelled() async { + await cancellationGate.wait() + } } From 846fed0178f25bb2613e1df1ad9b7e4477490359 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 07:37:42 +0900 Subject: [PATCH 29/38] Accept Store close API baseline --- .../v0.6.2/metadata.json | 6 ++--- .../v0.6.2/public-api.json | 22 +++++++++++++++++++ 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/scripts/compatibility-baselines/v0.6.2/metadata.json b/scripts/compatibility-baselines/v0.6.2/metadata.json index 1d912bf2..4c1b8249 100644 --- a/scripts/compatibility-baselines/v0.6.2/metadata.json +++ b/scripts/compatibility-baselines/v0.6.2/metadata.json @@ -1,8 +1,8 @@ { "architecture": "arm64", "baseline": "v0.6.2 public Swift products", - "baselineSHA256": "465956e4233ae545980a42dd03bd34a4de2bedea93c4ce735f15b68a37ded17d", - "captureRevision": "d6b877752f8d5b101aa36382ba64d7d6fe786d71", + "baselineSHA256": "04517f712053061a8f662752ecbb33d367738354aa2eef10609ebea57e0e501b", + "captureRevision": "6a05d14a67932ffddb2f98791855df4a7a6a93a3", "digesterInvocation": "xcrun swift-api-digester -dump-sdk -module CodexReview -module CodexReviewHost -module ReviewUI -module TextTransitions -swift-only -avoid-location -avoid-tool-args -abort-on-module-fail -sdk -target arm64-apple-macosx26.0 -swift-version 6 ", "macOSSDKVersion": "27.0", "modules": [ @@ -12,7 +12,7 @@ "TextTransitions" ], "publishedReleaseRevision": "82bddbcb1310a091eff742b36ab90781a4cbee5a", - "recoveryDesignRevision": "18499aa795d6af7716a5d94c92d4c7210fc3b821", + "recoveryDesignRevision": "d0ec991ec9c553b30f49ac5689e8b3b159f6f1ae", "swiftPMBuildSystem": "swiftbuild", "swiftAPIDigesterSHA256": "63f3ca45b1cc6ee81ccf2b6a11e9de3e7deeef85438915403ad5fa0bd5bccde8", "swiftVersion": "swift-driver version: 1.168.5 Apple Swift version 6.4 (swiftlang-6.4.0.27.1 clang-2100.3.27.1)", diff --git a/scripts/compatibility-baselines/v0.6.2/public-api.json b/scripts/compatibility-baselines/v0.6.2/public-api.json index f6519d19..b4a780f9 100644 --- a/scripts/compatibility-baselines/v0.6.2/public-api.json +++ b/scripts/compatibility-baselines/v0.6.2/public-api.json @@ -16645,6 +16645,28 @@ "printedName": "stop()", "usr": "s:11CodexReview0aB5StoreC4stopyyYaF" }, + { + "children": [ + { + "kind": "TypeNominal", + "name": "Void", + "printedName": "()" + } + ], + "declAttributes": [ + "Final", + "Custom" + ], + "declKind": "Func", + "funcSelfKind": "NonMutating", + "kind": "Function", + "mangledName": "$s11CodexReview0aB5StoreC5closeyyYaKF", + "moduleName": "CodexReview", + "name": "close", + "printedName": "close()", + "throwing": true, + "usr": "s:11CodexReview0aB5StoreC5closeyyYaKF" + }, { "children": [ { From f449ab388fe1e18b1244fb347212fb9538504bf6 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 07:38:32 +0900 Subject: [PATCH 30/38] Document accepted Store close baseline --- scripts/compatibility-baselines/README.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/scripts/compatibility-baselines/README.md b/scripts/compatibility-baselines/README.md index 1e5d33da..e8007d9f 100644 --- a/scripts/compatibility-baselines/README.md +++ b/scripts/compatibility-baselines/README.md @@ -36,5 +36,13 @@ file; the checker never rewrites it. - Replacement baseline SHA-256: `465956e4233ae545980a42dd03bd34a4de2bedea93c4ce735f15b68a37ded17d` - Accepted baseline update: [`f377ba37dee59eace8237f28661a27c121707d16`](https://github.com/lynnswap/CodexReviewKit/commit/f377ba37dee59eace8237f28661a27c121707d16) +### v0.6.2 Store close addition + +- Reviewed contract change: additive `CodexReviewStore.close() async throws` +- Capture revision: `6a05d14a67932ffddb2f98791855df4a7a6a93a3` +- Recovery design acceptance: [`d0ec991ec9c553b30f49ac5689e8b3b159f6f1ae`](https://github.com/lynnswap/CodexReviewKit/commit/d0ec991ec9c553b30f49ac5689e8b3b159f6f1ae) +- Replacement baseline SHA-256: `04517f712053061a8f662752ecbb33d367738354aa2eef10609ebea57e0e501b` +- Accepted baseline update: [`846fed0178f25bb2613e1df1ad9b7e4477490359`](https://github.com/lynnswap/CodexReviewKit/commit/846fed0178f25bb2613e1df1ad9b7e4477490359) + Future entries must identify the reviewed contract change, the replacement baseline checksum, and the commit that accepted both. From 22874f30856202d126a2020a27abb39d2db5c04e Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 05:30:10 +0900 Subject: [PATCH 31/38] test(app-server): gate pre-dispatch cancellation --- Tests/CodexReviewAppServerTests/AppServerClientTests.swift | 3 +++ 1 file changed, 3 insertions(+) diff --git a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift index d60e3bde..7c4c275a 100644 --- a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift +++ b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift @@ -1459,6 +1459,7 @@ struct AppServerClientTests { forceClose: { Issue.record("Pre-dispatch cancellation force-closed connection.") } ) } + #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) await initializeGate.open() #expect(try await cancellation.value.terminal == .localCancellation( @@ -1507,6 +1508,7 @@ struct AppServerClientTests { forceClose: { Issue.record("Prepared-thread cancellation force-closed connection.") } ) } + #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) await threadGate.open() #expect(try await cancellation.value.terminal == .localCancellation( @@ -1566,6 +1568,7 @@ struct AppServerClientTests { } ) } + #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) await reviewGate.open() let attempt = try await registeredStart.task.value await transport.waitForResponseDelivery(method: "turn/interrupt") From 74b6470d62419ae6a8de97a455d90f7860fecae5 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 07:58:36 +0900 Subject: [PATCH 32/38] Route runtime replacement close cause --- .../AppServerCodexReviewBackend.swift | 23 +++++++++++++------ .../LiveCodexReviewStoreBackend.swift | 4 ++-- .../AppServerClientTests.swift | 22 ++++++++++++++++++ 3 files changed, 40 insertions(+), 9 deletions(-) diff --git a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift index ebd53646..17619f52 100644 --- a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift +++ b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift @@ -17,11 +17,11 @@ private func makeAppServerReviewAttemptID() -> String { package struct AppServerRuntimeOwnerLifecycleHandle: Sendable { private let closeAdmissionOperation: @Sendable () async -> Void - private let closeAndWaitOperation: @Sendable () async throws -> Void + private let closeAndWaitOperation: @Sendable (ReviewRuntimeTransitionPurpose) async throws -> Void fileprivate init( closeAdmissionOperation: @escaping @Sendable () async -> Void, - closeAndWaitOperation: @escaping @Sendable () async throws -> Void + closeAndWaitOperation: @escaping @Sendable (ReviewRuntimeTransitionPurpose) async throws -> Void ) { self.closeAdmissionOperation = closeAdmissionOperation self.closeAndWaitOperation = closeAndWaitOperation @@ -31,8 +31,10 @@ package struct AppServerRuntimeOwnerLifecycleHandle: Sendable { await closeAdmissionOperation() } - package func closeAndWait() async throws { - try await closeAndWaitOperation() + package func closeAndWait( + purpose: ReviewRuntimeTransitionPurpose = .stop + ) async throws { + try await closeAndWaitOperation(purpose) } } @@ -202,8 +204,8 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { closeAdmissionOperation: { [self] in await closeAdmissionFromRuntimeOwner() }, - closeAndWaitOperation: { [self] in - try await closeFromRuntimeOwnerAndWait() + closeAndWaitOperation: { [self] purpose in + try await closeFromRuntimeOwnerAndWait(purpose: purpose) } ) } @@ -587,11 +589,18 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { // Only AppServerRuntimeOwnerLifecycleHandle can enter this transition. The // notification router never owns that handle, so close cannot await itself. - private func closeFromRuntimeOwnerAndWait() async throws { + private func closeFromRuntimeOwnerAndWait( + purpose: ReviewRuntimeTransitionPurpose + ) async throws { lifecycleTestingObservation.recordCloseCaller() let closeTask: Task switch lifecycleState { case .open: + if purpose == .recoveryReplacement { + connectionStreamFailure = .ownerForcedConnectionClose( + .connection("Review connection was force-closed by its runtime owner.") + ) + } closeAdmissionFromRuntimeOwner() let client = client let task = Task { diff --git a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift index 9e1598fc..60c1f3fc 100644 --- a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift +++ b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift @@ -1757,7 +1757,7 @@ private final class LiveRuntimeLifecycleHandle: RuntimeLifecycleHandle { await backend.runtimeOwnerLifecycleHandle.closeAdmission() } - func close(purpose _: ReviewRuntimeTransitionPurpose) async throws { + func close(purpose: ReviewRuntimeTransitionPurpose) async throws { let task: Task, Never> if let closeTask { task = closeTask @@ -1768,7 +1768,7 @@ private final class LiveRuntimeLifecycleHandle: RuntimeLifecycleHandle { var failures: [ReviewLifecycleResourceFailure] = [] authObservationTask?.cancel() do { - try await appServerLifecycle.closeAndWait() + try await appServerLifecycle.closeAndWait(purpose: purpose) } catch { if let aggregate = error as? ReviewLifecycleResourceFailureAggregate { failures.append(aggregate.first) diff --git a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift index 7c4c275a..e584faff 100644 --- a/Tests/CodexReviewAppServerTests/AppServerClientTests.swift +++ b/Tests/CodexReviewAppServerTests/AppServerClientTests.swift @@ -2506,6 +2506,28 @@ struct AppServerClientTests { } } + @Test func backendRuntimeReplacementClosePublishesOwnerForcedTerminal() async throws { + let transport = FakeJSONRPCTransport() + let backend = AppServerCodexReviewBackend(client: .init(transport: transport)) + let run = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-1", + threadID: "thread-1", + turnID: "turn-1", + reviewThreadID: "thread-1" + ) + var iterator = await eventSequence(backend, run).makeAsyncIterator() + + let lifecycle = backend.runtimeOwnerLifecycleHandle + await lifecycle.closeAdmission() + try await lifecycle.closeAndWait(purpose: .recoveryReplacement) + + await #expect(throws: ReviewAttemptStreamFailure.ownerForcedConnectionClose( + .connection("Review connection was force-closed by its runtime owner.") + )) { + _ = try await iterator.next() + } + } + @Test func backendPreservesBufferedEventsBeforeNotificationStreamError() async throws { let transport = FakeJSONRPCTransport() let backend = AppServerCodexReviewBackend(client: .init(transport: transport)) From 86342c15916a65383b19dd949c194ba99df29a18 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 08:21:48 +0900 Subject: [PATCH 33/38] Own shared runtime recovery replacement --- .../CodexReview/ReviewRuntimeLifecycle.swift | 179 +++++- .../CodexReview/Store/CodexReviewStore.swift | 240 ++++++-- .../Store/CodexReviewStoreReviews.swift | 521 ++++++++++++++++-- .../AppServerCodexReviewBackend.swift | 2 +- .../LiveCodexReviewStoreBackend.swift | 88 ++- Sources/CodexReviewTesting/TestSupport.swift | 30 +- .../CodexReviewStoreLifecycleTests.swift | 109 +++- 7 files changed, 1039 insertions(+), 130 deletions(-) diff --git a/Sources/CodexReview/ReviewRuntimeLifecycle.swift b/Sources/CodexReview/ReviewRuntimeLifecycle.swift index 435a5148..83e1790c 100644 --- a/Sources/CodexReview/ReviewRuntimeLifecycle.swift +++ b/Sources/CodexReview/ReviewRuntimeLifecycle.swift @@ -256,6 +256,173 @@ package struct ReviewRuntimePreparationFailure: LocalizedError, Sendable { } } +package struct ReviewRuntimeReplacementParticipant: Equatable, Sendable { + package let jobID: String + package let attemptID: String + + package init(jobID: String, attemptID: String) { + self.jobID = jobID + self.attemptID = attemptID + } +} + +package struct ReviewRuntimeReplacementFailure: LocalizedError, Sendable { + package let resources: ReviewLifecycleResourceFailureAggregate + + package init(resources: ReviewLifecycleResourceFailureAggregate) { + self.resources = resources + } + + package init(_ failure: ReviewLifecycleResourceFailure) { + self.resources = .init(first: failure) + } + + package var errorDescription: String? { + resources.localizedDescription + } +} + +@MainActor +package final class ReviewRuntimeRecoveryReplacement { + package enum Trigger: Equatable, Sendable { + case explicitCancellation(targetJobID: String) + case sameAccountRestart + case recoverableNetwork(initiatingJobID: String) + + package var targetJobID: String? { + guard case .explicitCancellation(let targetJobID) = self else { + return nil + } + return targetJobID + } + + package var requiresNetworkRestoration: Bool { + if case .recoverableNetwork = self { true } else { false } + } + } + + package enum Outcome: Sendable { + case running(ReviewRuntimeGeneration) + case failed(ReviewRuntimeReplacementFailure) + case superseded(ReviewRuntimeTransitionPurpose) + } + + package let sourceGeneration: ReviewRuntimeGeneration + package let replacementGeneration: ReviewRuntimeGeneration + package let sourceRuntime: PreparedRuntime? + package let retainedMCPGeneration: MCPServerGeneration + package let retainedServerURL: URL? + package let trigger: Trigger + package private(set) var participants: [ReviewRuntimeReplacementParticipant] + + private var sourceCloseResult: Result? + private var sourceCloseWaiters: [CheckedContinuation, Never>] = [] + private var outcome: Outcome? + private var outcomeContinuations: [AsyncStream.Continuation] = [] + private var networkRestorationWasRecorded = false + private var networkRestorationContinuations: [AsyncStream.Continuation] = [] + private var remainingParticipantJobIDs: Set + + package init( + sourceGeneration: ReviewRuntimeGeneration, + replacementGeneration: ReviewRuntimeGeneration, + sourceRuntime: PreparedRuntime?, + retainedMCPGeneration: MCPServerGeneration, + retainedServerURL: URL?, + trigger: Trigger, + participants: [ReviewRuntimeReplacementParticipant] + ) { + self.sourceGeneration = sourceGeneration + self.replacementGeneration = replacementGeneration + self.sourceRuntime = sourceRuntime + self.retainedMCPGeneration = retainedMCPGeneration + self.retainedServerURL = retainedServerURL + self.trigger = trigger + self.participants = participants + self.remainingParticipantJobIDs = Set(participants.map(\.jobID)) + } + + package func suppressParticipant(jobID: String) { + participants.removeAll { $0.jobID == jobID } + remainingParticipantJobIDs.remove(jobID) + } + + package func finishParticipant(jobID: String) { + remainingParticipantJobIDs.remove(jobID) + } + + package var hasRemainingParticipants: Bool { + remainingParticipantJobIDs.isEmpty == false + } + + package func finishSourceClose( + _ result: Result + ) { + guard sourceCloseResult == nil else { return } + sourceCloseResult = result + let waiters = sourceCloseWaiters + sourceCloseWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume(returning: result) + } + } + + package func waitForSourceClose() async -> Result { + if let sourceCloseResult { return sourceCloseResult } + return await withCheckedContinuation { continuation in + if let sourceCloseResult { + continuation.resume(returning: sourceCloseResult) + } else { + sourceCloseWaiters.append(continuation) + } + } + } + + package func replacementOutcomes() -> AsyncStream { + AsyncStream(bufferingPolicy: .bufferingNewest(1)) { continuation in + if let outcome { + continuation.yield(outcome) + continuation.finish() + } else { + outcomeContinuations.append(continuation) + } + } + } + + package func finish(_ outcome: Outcome) { + guard self.outcome == nil else { return } + self.outcome = outcome + let continuations = outcomeContinuations + outcomeContinuations.removeAll(keepingCapacity: false) + for continuation in continuations { + continuation.yield(outcome) + continuation.finish() + } + } + + package func recordNetworkRestoration() { + guard networkRestorationWasRecorded == false else { return } + networkRestorationWasRecorded = true + let continuations = networkRestorationContinuations + networkRestorationContinuations.removeAll(keepingCapacity: false) + for continuation in continuations { + continuation.yield(()) + continuation.finish() + } + } + + package func networkRestorations() -> AsyncStream { + AsyncStream(bufferingPolicy: .bufferingNewest(1)) { continuation in + if trigger.requiresNetworkRestoration == false || networkRestorationWasRecorded { + continuation.yield(()) + continuation.finish() + } else { + networkRestorationContinuations.append(continuation) + } + } + } +} + package enum ReviewStoreRuntimeState { case stopped(ReviewRuntimeGeneration) case acquiring( @@ -273,12 +440,14 @@ package enum ReviewStoreRuntimeState { purpose: ReviewRuntimeTransitionPurpose, task: Task, record: ReviewRuntimeTransitionRecord, - sourceRuntime: PreparedRuntime? + sourceRuntime: PreparedRuntime?, + recoveryReplacement: ReviewRuntimeRecoveryReplacement? ) case failed( generation: ReviewRuntimeGeneration, retainedMCPGeneration: MCPServerGeneration, - serverURL: URL? + serverURL: URL?, + replacementFailure: ReviewRuntimeReplacementFailure? ) package var generation: ReviewRuntimeGeneration { @@ -286,8 +455,8 @@ package enum ReviewStoreRuntimeState { case .stopped(let generation), .acquiring(let generation, _, _), .running(let generation, _, _), - .transitioning(let generation, _, _, _, _), - .failed(let generation, _, _): + .transitioning(let generation, _, _, _, _, _), + .failed(let generation, _, _, _): generation } } @@ -296,7 +465,7 @@ package enum ReviewStoreRuntimeState { switch self { case .running(_, let runtime, _): return runtime - case .transitioning(_, _, _, _, let sourceRuntime): + case .transitioning(_, _, _, _, let sourceRuntime, _): return sourceRuntime case .stopped, .acquiring, .failed: return nil diff --git a/Sources/CodexReview/Store/CodexReviewStore.swift b/Sources/CodexReview/Store/CodexReviewStore.swift index f8f3ad46..efba1d9e 100644 --- a/Sources/CodexReview/Store/CodexReviewStore.swift +++ b/Sources/CodexReview/Store/CodexReviewStore.swift @@ -56,6 +56,7 @@ public final class CodexReviewStore { .init(rawValue: 0) ) @ObservationIgnored package var lastRuntimeTransitionRecord: ReviewRuntimeTransitionRecord? + @ObservationIgnored package let runtimeReplacementCoordinator = ReviewRuntimeReplacementCoordinator() package init( backend: any CodexReviewStoreBackend = PreviewCodexReviewStoreBackend(), @@ -103,7 +104,7 @@ public final class CodexReviewStore { task.cancel() } switch runtimeState { - case .acquiring(_, let task, _), .transitioning(_, _, let task, _, _): + case .acquiring(_, let task, _), .transitioning(_, _, let task, _, _, _): task.cancel() case .stopped, .running, .failed: break @@ -170,7 +171,7 @@ public final class CodexReviewStore { switch runtimeState { case .acquiring: return - case .transitioning(_, _, let task, _, _): + case .transitioning(_, _, let task, _, _, _): await task.value return case .running where forceRestartIfNeeded == false: @@ -183,7 +184,7 @@ public final class CodexReviewStore { retainedServerURL: serverURL ) return - case .failed(let generation, let mcpGeneration, let retainedServerURL): + case .failed(let generation, let mcpGeneration, let retainedServerURL, _): await startRuntimeReplacement( previousGeneration: generation, previousRuntime: nil, @@ -231,28 +232,55 @@ public final class CodexReviewStore { retainedMCPGeneration: MCPServerGeneration, retainedServerURL: URL? ) async { - let generation = previousGeneration.successor() - let record = ReviewRuntimeTransitionRecord() serverState = .starting writeDiagnosticsIfNeeded() + let (_, task) = installRuntimeReplacement( + previousGeneration: previousGeneration, + sourceRuntime: previousRuntime, + retainedMCPGeneration: retainedMCPGeneration, + retainedServerURL: retainedServerURL, + trigger: .sameAccountRestart, + purpose: .restartSameAccount + ) + await task.value + } + + @discardableResult + package func installRuntimeReplacement( + previousGeneration: ReviewRuntimeGeneration, + sourceRuntime: PreparedRuntime?, + retainedMCPGeneration: MCPServerGeneration, + retainedServerURL: URL?, + trigger: ReviewRuntimeRecoveryReplacement.Trigger, + purpose: ReviewRuntimeTransitionPurpose + ) -> (ReviewRuntimeRecoveryReplacement, Task) { + let generation = previousGeneration.successor() + let record = ReviewRuntimeTransitionRecord() + let replacement = makeRuntimeRecoveryReplacement( + sourceGeneration: previousGeneration, + replacementGeneration: generation, + sourceRuntime: sourceRuntime, + retainedMCPGeneration: retainedMCPGeneration, + retainedServerURL: retainedServerURL, + trigger: trigger + ) let task = Task { @MainActor [weak self] in guard let self else { return } await self.performRuntimeReplacement( - generation: generation, - previousRuntime: previousRuntime, - retainedMCPGeneration: retainedMCPGeneration, - retainedServerURL: retainedServerURL, + replacement, + purpose: purpose, record: record ) } runtimeState = .transitioning( generation: generation, - purpose: .restartSameAccount, + purpose: purpose, task: task, record: record, - sourceRuntime: previousRuntime + sourceRuntime: sourceRuntime, + recoveryReplacement: replacement ) - await task.value + return (replacement, task) } public func stop() async { @@ -270,12 +298,16 @@ public final class CodexReviewStore { case .stopped: transitionToStopped() return - case .transitioning(_, .stop, let task, _, _): + case .transitioning(_, .stop, let task, _, _, _): await task.value return case .acquiring, .running, .transitioning, .failed: break } + if case .transitioning(_, _, _, _, _, let replacement) = previousState { + replacement?.finish(.superseded(.stop)) + replacement?.recordNetworkRestoration() + } let invalidatedGeneration = previousState.generation.successor() let record = ReviewRuntimeTransitionRecord() let task = Task { @MainActor [weak self] in @@ -291,7 +323,8 @@ public final class CodexReviewStore { purpose: .stop, task: task, record: record, - sourceRuntime: previousState.runtimeForClose + sourceRuntime: previousState.runtimeForClose, + recoveryReplacement: nil ) await task.value } @@ -410,6 +443,10 @@ public final class CodexReviewStore { failureLedger: ReviewCloseFailureLedger ) async -> Result { let runningRuntime = previousRuntimeState.runtimeForClose + if case .transitioning(_, _, _, _, _, let replacement) = previousRuntimeState { + replacement?.finish(.superseded(.applicationClose)) + replacement?.recordNetworkRestoration() + } if let lastRuntimeTransitionRecord { failureLedger.importReceipts(from: lastRuntimeTransitionRecord) } @@ -475,8 +512,13 @@ public final class CodexReviewStore { await cancelAccountRateLimitAutoRefreshAndWait() switch previousRuntimeState { - case .acquiring(_, let task, let record), - .transitioning(_, _, let task, let record, _): + case .acquiring(_, let task, let record): + task.cancel() + await task.value + failureLedger.merge(record) + case .transitioning(_, _, let task, let record, _, let replacement): + replacement?.finish(.superseded(.applicationClose)) + replacement?.recordNetworkRestoration() task.cancel() await task.value failureLedger.merge(record) @@ -576,7 +618,9 @@ public final class CodexReviewStore { record.merge(previousRecord) case .running(_, let runtime, _): await stopPublishedRuntime(runtime, record: record) - case .transitioning(_, _, let task, let previousRecord, _): + case .transitioning(_, _, let task, let previousRecord, _, let replacement): + replacement?.finish(.superseded(.stop)) + replacement?.recordNetworkRestoration() task.cancel() await task.value record.merge(previousRecord) @@ -586,7 +630,7 @@ public final class CodexReviewStore { case .stopped: break } - guard case .transitioning(let currentGeneration, .stop, _, _, _) = runtimeState, + guard case .transitioning(let currentGeneration, .stop, _, _, _, _) = runtimeState, currentGeneration == invalidatedGeneration else { return @@ -658,7 +702,7 @@ public final class CodexReviewStore { } public func waitUntilStopped() async { - if case .transitioning(_, _, let task, _, _) = runtimeState { + if case .transitioning(_, _, let task, _, _, _) = runtimeState { await task.value } await backend.waitUntilStopped() @@ -771,81 +815,180 @@ public final class CodexReviewStore { } private func performRuntimeReplacement( - generation: ReviewRuntimeGeneration, - previousRuntime: PreparedRuntime?, - retainedMCPGeneration: MCPServerGeneration, - retainedServerURL: URL?, + _ replacement: ReviewRuntimeRecoveryReplacement, + purpose: ReviewRuntimeTransitionPurpose, record: ReviewRuntimeTransitionRecord ) async { defer { lastRuntimeTransitionRecord = record } var preparedRuntime: PreparedRuntime? - if let previousRuntime { - await performPublishedRuntimeSemanticStop(record: record) - await previousRuntime.handle.closeAdmission() - await closeAppServerRuntime( - previousRuntime, - purpose: .restartSameAccount, - record: record + await enrollRuntimeReplacementParticipants(replacement) + if let sourceRuntime = replacement.sourceRuntime { + await sourceRuntime.handle.closeAdmission() + let closeResult = await sourceRuntime.closeRecord.closeAndWait( + handle: sourceRuntime.handle, + purpose: purpose ) + let consumedFailures = sourceRuntime.closeRecord.consumeFailures() + if let targetJobID = replacement.trigger.targetJobID { + record.recordForceCloseFailures( + consumedFailures, + jobID: targetJobID + ) + } else { + record.record(contentsOf: consumedFailures) + } + if let firstFailure = closeResult.failures.first { + replacement.finishSourceClose(.failure( + reviewRuntimeCloseFailure(from: firstFailure) + )) + } else { + replacement.finishSourceClose(.success(())) + } + } else { + replacement.finishSourceClose(.success(())) } - guard isCurrentTransition(generation, purpose: .restartSameAccount) else { + guard isCurrentRuntimeReplacement(replacement, purpose: purpose) else { + replacement.finish(.superseded(currentSupersedingPurpose)) return } + if replacement.trigger.requiresNetworkRestoration { + var iterator = replacement.networkRestorations().makeAsyncIterator() + guard await iterator.next() != nil, + isCurrentRuntimeReplacement(replacement, purpose: purpose) + else { + replacement.finish(.superseded(currentSupersedingPurpose)) + return + } + } do { let runtime = try await backend.prepareRuntime( - generation: generation, - purpose: .restartSameAccount + generation: replacement.replacementGeneration, + purpose: purpose ) preparedRuntime = runtime - guard isCurrentTransition(generation, purpose: .restartSameAccount) else { + guard isCurrentRuntimeReplacement(replacement, purpose: purpose) else { await closeStaleRuntime( runtime, mcpServerWasPrepared: false, - purpose: .restartSameAccount, + purpose: purpose, record: record ) + replacement.finish(.superseded(currentSupersedingPurpose)) return } try await runtime.handle.activate() - guard isCurrentTransition(generation, purpose: .restartSameAccount) else { + guard isCurrentRuntimeReplacement(replacement, purpose: purpose) else { await closeStaleRuntime( runtime, mcpServerWasPrepared: false, - purpose: .restartSameAccount, + purpose: purpose, record: record ) + replacement.finish(.superseded(currentSupersedingPurpose)) return } publishRuntimeSnapshot(runtime.snapshot) runtimeState = .running( - generation: generation, + generation: replacement.replacementGeneration, runtime: runtime, - mcpGeneration: retainedMCPGeneration + mcpGeneration: replacement.retainedMCPGeneration ) - publishMCPServer(serverURL: retainedServerURL) + publishMCPServer(serverURL: replacement.retainedServerURL) + replacement.finish(.running(replacement.replacementGeneration)) } catch { if let preparedRuntime { await closeStaleRuntime( preparedRuntime, mcpServerWasPrepared: false, - purpose: .restartSameAccount, + purpose: purpose, record: record ) } - guard isCurrentTransition(generation, purpose: .restartSameAccount) else { + guard isCurrentRuntimeReplacement(replacement, purpose: purpose) else { + replacement.finish(.superseded(currentSupersedingPurpose)) return } + let failure = runtimeReplacementFailure(from: error) + record.record(.lifecycleResources(failure.resources)) runtimeState = .failed( - generation: generation, - retainedMCPGeneration: retainedMCPGeneration, - serverURL: retainedServerURL + generation: replacement.replacementGeneration, + retainedMCPGeneration: replacement.retainedMCPGeneration, + serverURL: replacement.retainedServerURL, + replacementFailure: failure ) - serverURL = retainedServerURL - serverState = .failed(error.localizedDescription) + serverURL = replacement.retainedServerURL + serverState = .failed(failure.localizedDescription) writeDiagnosticsIfNeeded() + replacement.finish(.failed(failure)) + } + } + + private func isCurrentRuntimeReplacement( + _ replacement: ReviewRuntimeRecoveryReplacement, + purpose: ReviewRuntimeTransitionPurpose + ) -> Bool { + guard case .open = lifetimeState, + case .transitioning( + let generation, + let currentPurpose, + _, + _, + _, + let currentReplacement + ) = runtimeState + else { + return false + } + return generation == replacement.replacementGeneration + && currentPurpose == purpose + && currentReplacement === replacement + } + + private var currentSupersedingPurpose: ReviewRuntimeTransitionPurpose { + if case .closing = lifetimeState { return .applicationClose } + if case .transitioning(_, let purpose, _, _, _, _) = runtimeState { + return purpose } + return .stop + } + + private func reviewRuntimeCloseFailure( + from failure: ReviewClosePrimaryFailure + ) -> ReviewRuntimeCloseFailure { + switch failure { + case .attemptRuntime(let failure): + return failure + case .lifecycleResources(let failure): + return .connection(failure.localizedDescription) + case .interruptRequest(let failure): + return .connection(failure.localizedDescription) + case .persistence(let failure): + return .connection(failure.localizedDescription) + } + } + + private func runtimeReplacementFailure( + from error: any Error + ) -> ReviewRuntimeReplacementFailure { + if let failure = error as? ReviewRuntimeReplacementFailure { + return failure + } + if let failure = error as? ReviewRuntimePreparationFailure { + return .init(resources: .init( + first: .client(failure.preparationDescription), + additionalInLifecycleOrder: [failure.cleanupFailures.first] + + failure.cleanupFailures.additionalInLifecycleOrder + )) + } + if let failure = error as? ReviewLifecycleResourceFailureAggregate { + return .init(resources: failure) + } + if let failure = error as? ReviewLifecycleResourceFailure { + return .init(failure) + } + return .init(.client(error.localizedDescription)) } private func isCurrentAcquisition( @@ -872,6 +1015,7 @@ public final class CodexReviewStore { let currentPurpose, _, _, + _, _ ) = runtimeState else { return false @@ -886,7 +1030,7 @@ public final class CodexReviewStore { case .open: break } - guard case .transitioning(_, let purpose, _, _, _) = runtimeState else { + guard case .transitioning(_, let purpose, _, _, _, _) = runtimeState else { return false } return purpose == .stop || purpose == .applicationClose diff --git a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift index e71a988b..13e421ac 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift @@ -143,7 +143,8 @@ extension CodexReviewStore { launchReviewWorker( jobID: jobID, startRequest: startRequest, - registeredStart: registered + registeredStart: registered, + runtimeGeneration: runtimeGeneration ) return jobID } @@ -166,14 +167,16 @@ extension CodexReviewStore { private func launchReviewWorker( jobID: String, startRequest: CodexReviewBackendModel.Review.Start, - registeredStart: ReviewRegisteredStart + registeredStart: ReviewRegisteredStart, + runtimeGeneration: ReviewRuntimeGeneration ) { reviewWorkerTasks[jobID]?.cancel() reviewWorkerTasks[jobID] = Task { [weak self] in await self?.runReviewWorker( jobID: jobID, startRequest: startRequest, - registeredStart: registeredStart + registeredStart: registeredStart, + runtimeGeneration: runtimeGeneration ) } } @@ -181,7 +184,8 @@ extension CodexReviewStore { private func runReviewWorker( jobID: String, startRequest: CodexReviewBackendModel.Review.Start, - registeredStart: ReviewRegisteredStart + registeredStart: ReviewRegisteredStart, + runtimeGeneration: ReviewRuntimeGeneration ) async { guard let job = job(id: jobID) else { reviewAttemptOwnerships.removeValue(forKey: jobID) @@ -212,7 +216,8 @@ extension CodexReviewStore { let completion = try await consumeReviewEvents( for: backendAttempt, job: job, - startRequest: startRequest + startRequest: startRequest, + runtimeGeneration: runtimeGeneration ) cleanupAttempt = completion.cleanupAttempt } @@ -347,7 +352,9 @@ extension CodexReviewStore { private func reviewWorkerInputs( for attempt: BackendReviewAttempt, - owner: ReviewActiveAttempt + owner: ReviewActiveAttempt, + jobID: String, + runtimeGeneration: ReviewRuntimeGeneration ) async -> ReviewWorkerInputs { let networkMonitor = self.networkMonitor let policy = self.networkRecoveryPolicy @@ -361,6 +368,12 @@ extension CodexReviewStore { ) let eventSource = ReviewWorkerEventSource(queue: queue) let recoveryInterruptionSource = ReviewWorkerRecoveryInterruptionSource(queue: queue) + let runtimeReplacementRegistrationID = runtimeReplacementCoordinator.register( + jobID: jobID, + attemptID: owner.run.attemptID, + runtimeGeneration: runtimeGeneration, + recoverySource: recoveryInterruptionSource + ) let networkTask = Task { for await snapshot in snapshots { await signalCoordinator.observe(snapshot) @@ -375,6 +388,10 @@ extension CodexReviewStore { networkStatusTracker: tracker, eventSource: eventSource, recoveryInterruptionSource: recoveryInterruptionSource, + runtimeReplacementCoordinator: runtimeReplacementCoordinator, + runtimeReplacementRegistrationID: runtimeReplacementRegistrationID, + jobID: jobID, + initialRuntimeGeneration: runtimeGeneration, initialEventSubscriptionID: initialEventSubscriptionID, networkTask: networkTask, signalCoordinator: signalCoordinator @@ -643,36 +660,167 @@ extension CodexReviewStore { "Review Store was released before runtime force-close." ) } - try await self.forceCloseCurrentRuntimeForAttempt(jobID: jobID) + try await self.forceCloseCurrentRuntimeForAttempt( + jobID: jobID, + trigger: .explicitCancellation(targetJobID: jobID) + ) } ) } - private func forceCloseCurrentRuntimeForAttempt(jobID: String) async throws { - let generation: ReviewRuntimeGeneration + package func makeRuntimeRecoveryReplacement( + sourceGeneration: ReviewRuntimeGeneration, + replacementGeneration: ReviewRuntimeGeneration, + sourceRuntime: PreparedRuntime?, + retainedMCPGeneration: MCPServerGeneration, + retainedServerURL: URL?, + trigger: ReviewRuntimeRecoveryReplacement.Trigger + ) -> ReviewRuntimeRecoveryReplacement { + let participants: [ReviewRuntimeReplacementParticipant] = reviewRegistrationOrder.compactMap { jobID in + guard jobID != trigger.targetJobID, + let job = job(id: jobID), + job.isTerminal == false, + job.cancellationRequested == false + else { + return nil + } + let active: ReviewActiveAttempt + switch reviewAttemptOwnerships[jobID] { + case .active(let current): + active = current + case .resolvingRecovery(let current): + guard case .recoverableNetwork(let initiatingJobID) = trigger, + initiatingJobID == jobID + else { + return nil + } + active = current + case .initialStart, .recoveryDisposition, .preparingRecovery, + .waitingForRecovery, .replacementStart, .terminal, nil: + return nil + } + return runtimeReplacementCoordinator.participant( + jobID: jobID, + attemptID: active.run.attemptID, + sourceGeneration: sourceGeneration + ) + } + let replacement = ReviewRuntimeRecoveryReplacement( + sourceGeneration: sourceGeneration, + replacementGeneration: replacementGeneration, + sourceRuntime: sourceRuntime, + retainedMCPGeneration: retainedMCPGeneration, + retainedServerURL: retainedServerURL, + trigger: trigger, + participants: participants + ) + runtimeReplacementCoordinator.install(replacement) + return replacement + } + + package func enrollRuntimeReplacementParticipants( + _ replacement: ReviewRuntimeRecoveryReplacement + ) async { + let backend = self.backend + for participant in replacement.participants { + guard let job = job(id: participant.jobID), + job.isTerminal == false, + job.cancellationRequested == false + else { + replacement.suppressParticipant(jobID: participant.jobID) + continue + } + let active: ReviewActiveAttempt + switch reviewAttemptOwnerships[participant.jobID] { + case .active(let current) where current.run.attemptID == participant.attemptID: + active = current + reviewAttemptOwnerships[participant.jobID] = .resolvingRecovery(current) + case .resolvingRecovery(let current) where current.run.attemptID == participant.attemptID: + active = current + default: + replacement.suppressParticipant(jobID: participant.jobID) + continue + } + let recoveryTrigger: ReviewAttemptRecoveryTrigger = if case .recoverableNetwork( + let initiatingJobID + ) = replacement.trigger, initiatingJobID == participant.jobID { + .recoverableNetworkLoss + } else { + .sameAccountRestart + } + let didInstall = await runtimeReplacementCoordinator.beginRecovery( + participant: participant, + owner: active, + trigger: recoveryTrigger, + interrupt: { run, reason in + try await backend.interruptReview(run, reason: reason) + }, + forceClose: { + try await replacement.waitForSourceClose().get() + } + ) + if didInstall == false { + if case .resolvingRecovery(let current) = reviewAttemptOwnerships[participant.jobID], + sameAttempt(current, active) { + reviewAttemptOwnerships[participant.jobID] = .active(active) + } + replacement.suppressParticipant(jobID: participant.jobID) + } + } + } + + private func forceCloseCurrentRuntimeForAttempt( + jobID: String, + trigger: ReviewRuntimeRecoveryReplacement.Trigger + ) async throws { + if case .open = lifetimeState, + case .running( + let generation, + let runtime, + let mcpGeneration + ) = runtimeState { + serverState = .starting + writeDiagnosticsIfNeeded() + let (replacement, _) = installRuntimeReplacement( + previousGeneration: generation, + sourceRuntime: runtime, + retainedMCPGeneration: mcpGeneration, + retainedServerURL: serverURL, + trigger: trigger, + purpose: .recoveryReplacement + ) + let sourceCloseResult = await replacement.waitForSourceClose() + await runtimeForceCloseReceiptRecordedForTesting?() + try sourceCloseResult.get() + return + } + + if case .transitioning( + _, + let purpose, + _, + _, + _, + let replacement? + ) = runtimeState, + purpose == .recoveryReplacement || purpose == .restartSameAccount { + if case .explicitCancellation = trigger { + replacement.suppressParticipant(jobID: jobID) + } + let sourceCloseResult = await replacement.waitForSourceClose() + await runtimeForceCloseReceiptRecordedForTesting?() + try sourceCloseResult.get() + return + } + let runtime: PreparedRuntime - let mcpGeneration: MCPServerGeneration? let record: ReviewRuntimeTransitionRecord switch runtimeState { - case .running( - let runningGeneration, - let runningRuntime, - let runningMCPGeneration - ): - generation = runningGeneration + case .running(_, let runningRuntime, _): runtime = runningRuntime - mcpGeneration = runningMCPGeneration record = ReviewRuntimeTransitionRecord() - case .transitioning( - let transitionGeneration, - _, - _, - let transitionRecord, - let sourceRuntime? - ): - generation = transitionGeneration + case .transitioning(_, _, _, let transitionRecord, let sourceRuntime?, _): runtime = sourceRuntime - mcpGeneration = nil record = transitionRecord case .stopped, .acquiring, .failed, .transitioning: throw ReviewRuntimeCloseFailure.connection( @@ -693,16 +841,6 @@ extension CodexReviewStore { } else { record.recordForceCloseFailures(consumedFailures, jobID: jobID) } - if let mcpGeneration, closeResult.installedClose { - if applicationCloseFailureLedger == nil { - lastRuntimeTransitionRecord = record - } - runtimeState = .failed( - generation: generation.successor(), - retainedMCPGeneration: mcpGeneration, - serverURL: serverURL - ) - } await runtimeForceCloseReceiptRecordedForTesting?() if let firstFailure = closeResult.failures.first { switch firstFailure { @@ -984,7 +1122,8 @@ extension CodexReviewStore { private func consumeReviewEvents( for initialAttempt: BackendReviewAttempt, job: CodexReviewJob, - startRequest: CodexReviewBackendModel.Review.Start + startRequest: CodexReviewBackendModel.Review.Start, + runtimeGeneration: ReviewRuntimeGeneration ) async throws -> ReviewWorkerAttemptCompletion { guard case .active(let initialActive) = reviewAttemptOwnerships[job.id], sameAttempt(initialActive, run: initialAttempt.run) @@ -993,7 +1132,12 @@ extension CodexReviewStore { message: "Initial event subscription requires the published active attempt." ) } - let inputs = await reviewWorkerInputs(for: initialAttempt, owner: initialActive) + let inputs = await reviewWorkerInputs( + for: initialAttempt, + owner: initialActive, + jobID: job.id, + runtimeGeneration: runtimeGeneration + ) guard case .active(let revalidatedActive) = reviewAttemptOwnerships[job.id], sameAttempt(revalidatedActive, initialActive) else { @@ -1031,6 +1175,7 @@ extension CodexReviewStore { var recoverySignals = ReviewNetworkRecoverySignals() var activeEventSubscriptionID: Int? = inputs.initialEventSubscriptionID var cleanupAttempt = initialCleanupAttempt + var activeRuntimeGeneration = inputs.initialRuntimeGeneration while let input = await inputs.next() { if job.isTerminal { return .init(cleanupAttempt: cleanupAttempt) @@ -1169,6 +1314,46 @@ extension CodexReviewStore { message: "Recovery handoff changed while detaching the old subscription." ) } + if candidate.trigger == .sameAccountRestart { + switch try await restartReviewAfterRuntimeRecovery( + job: job, + startRequest: startRequest, + inputs: inputs, + handoff: handoff, + sourceRuntimeGeneration: activeRuntimeGeneration + ) { + case .continueWaiting: + throw ReviewAttemptContractFailure( + message: "Same-account recovery unexpectedly waited for network restoration." + ) + case .finished: + return .init(cleanupAttempt: cleanupAttempt) + case .recovered( + let recoveredAttempt, + let active, + let runtimeGeneration + ): + cleanupAttempt = active + let subscriptionID = await inputs.subscribe( + to: recoveredAttempt, + owner: active + ) + guard case .active(let current) = reviewAttemptOwnerships[job.id], + sameAttempt(current, active) + else { + await inputs.cancelActiveEventSubscription() + if job.isTerminal { + return .init(cleanupAttempt: cleanupAttempt) + } + throw ReviewAttemptContractFailure( + message: "Recovered subscription completed after its active attempt changed." + ) + } + activeEventSubscriptionID = subscriptionID + activeRuntimeGeneration = runtimeGeneration + recoverySignals.markRecovered() + } + } } case .networkSnapshot(let snapshot, let recoveryGeneration): if let pendingFailure = recoverySignals @@ -1200,17 +1385,25 @@ extension CodexReviewStore { ) else { continue } - switch try await restartReviewAfterNetworkRestore( + runtimeReplacementCoordinator.recordNetworkRestoration( + sourceGeneration: activeRuntimeGeneration + ) + switch try await restartReviewAfterRuntimeRecovery( job: job, startRequest: startRequest, inputs: inputs, - handoff: handoff + handoff: handoff, + sourceRuntimeGeneration: activeRuntimeGeneration ) { case .continueWaiting: continue case .finished: return .init(cleanupAttempt: cleanupAttempt) - case .recovered(let recoveredAttempt, let active): + case .recovered( + let recoveredAttempt, + let active, + let runtimeGeneration + ): cleanupAttempt = active let subscriptionID = await inputs.subscribe( to: recoveredAttempt, @@ -1228,6 +1421,7 @@ extension CodexReviewStore { ) } activeEventSubscriptionID = subscriptionID + activeRuntimeGeneration = runtimeGeneration recoverySignals.markRecovered() } case .networkOutageConfirmed: @@ -1244,14 +1438,22 @@ extension CodexReviewStore { attemptID: active.run.attemptID ) let backend = self.backend - await inputs.beginRecoveryInterruption(for: active) { + await inputs.beginRecoveryInterruption(for: active) { [self] in try await active.admission.beginRecovery( trigger: .recoverableNetworkLoss, interrupt: { run, reason in try await backend.interruptReview(run, reason: reason) }, - forceClose: { - try await backend.forceCloseReviewConnection() + forceClose: { @MainActor [weak self] in + guard let self else { + throw ReviewRuntimeCloseFailure.connection( + "Review Store was released before network runtime replacement." + ) + } + try await self.forceCloseCurrentRuntimeForAttempt( + jobID: job.id, + trigger: .recoverableNetwork(initiatingJobID: job.id) + ) } ) } @@ -1370,11 +1572,12 @@ extension CodexReviewStore { && lhs.admission === rhs.admission } - private func restartReviewAfterNetworkRestore( + private func restartReviewAfterRuntimeRecovery( job: CodexReviewJob, startRequest: CodexReviewBackendModel.Review.Start, inputs: ReviewWorkerInputs, - handoff: ReviewRecoveryHandoff + handoff: ReviewRecoveryHandoff, + sourceRuntimeGeneration: ReviewRuntimeGeneration ) async throws -> NetworkRestoreRestartResult { guard case .waitingForRecovery(let currentHandoff) = reviewAttemptOwnerships[job.id], currentHandoff == handoff @@ -1389,9 +1592,44 @@ extension CodexReviewStore { if Task.isCancelled || job.isTerminal { return .finished } - guard await inputs.networkStatusTracker.currentStatus() == .satisfied else { + if handoff.candidate.trigger == .recoverableNetworkLoss, + await inputs.networkStatusTracker.currentStatus() != .satisfied { return .continueWaiting } + let destinationRuntimeGeneration: ReviewRuntimeGeneration + if let replacement = inputs.runtimeReplacementCoordinator.replacement( + sourceGeneration: sourceRuntimeGeneration + ) { + var outcomes = replacement.replacementOutcomes().makeAsyncIterator() + guard let outcome = await outcomes.next() else { + return .finished + } + switch outcome { + case .running(let generation): + destinationRuntimeGeneration = generation + case .failed(let failure): + markReviewFailed(job, message: failure.localizedDescription) + reviewAttemptOwnerships[job.id] = .terminal + inputs.runtimeReplacementCoordinator.finishParticipant( + jobID: job.id, + sourceGeneration: sourceRuntimeGeneration + ) + return .finished + case .superseded: + inputs.runtimeReplacementCoordinator.finishParticipant( + jobID: job.id, + sourceGeneration: sourceRuntimeGeneration + ) + throw CancellationError() + } + } else { + guard handoff.candidate.trigger == .recoverableNetworkLoss else { + throw ReviewAttemptContractFailure( + message: "Same-account recovery lost its runtime replacement owner." + ) + } + destinationRuntimeGeneration = sourceRuntimeGeneration + } let recoveredAdmission = ReviewStartAdmission(closePolicy: reviewRuntimeClosePolicy) let backend = self.backend let registered = try await recoveredAdmission.registerStart { admission in @@ -1430,13 +1668,40 @@ extension CodexReviewStore { ) } let recoveredAttempt = try result.get() + if destinationRuntimeGeneration != sourceRuntimeGeneration { + guard case .running(let currentGeneration, _, _) = runtimeState, + currentGeneration == destinationRuntimeGeneration, + case .open = lifetimeState + else { + _ = try? await recoveredAdmission.cancel( + job.core.lifecycle.cancellation ?? .system(), + interrupt: { _, _ in }, + forceClose: {} + ) + return .finished + } + } let active = ReviewActiveAttempt( run: recoveredAttempt.run, admission: recoveredAdmission ) reviewAttemptOwnerships[job.id] = .active(active) applyBackendRun(recoveredAttempt.run, to: job) - return .recovered(recoveredAttempt, active) + inputs.runtimeReplacementCoordinator.update( + registrationID: inputs.runtimeReplacementRegistrationID, + jobID: job.id, + attemptID: recoveredAttempt.run.attemptID, + runtimeGeneration: destinationRuntimeGeneration + ) + inputs.runtimeReplacementCoordinator.finishParticipant( + jobID: job.id, + sourceGeneration: sourceRuntimeGeneration + ) + return .recovered( + recoveredAttempt, + active, + destinationRuntimeGeneration + ) } private func handleReviewEvent( @@ -1876,7 +2141,11 @@ private struct ReviewWorkerRecoveryFailure: LocalizedError, @unchecked Sendable private enum NetworkRestoreRestartResult { case continueWaiting case finished - case recovered(BackendReviewAttempt, ReviewActiveAttempt) + case recovered( + BackendReviewAttempt, + ReviewActiveAttempt, + ReviewRuntimeGeneration + ) } private enum ReviewNetworkSnapshotEffect { @@ -1976,6 +2245,10 @@ private struct ReviewWorkerInputs { var networkStatusTracker: ReviewNetworkStatusTracker var eventSource: ReviewWorkerEventSource var recoveryInterruptionSource: ReviewWorkerRecoveryInterruptionSource + var runtimeReplacementCoordinator: ReviewRuntimeReplacementCoordinator + var runtimeReplacementRegistrationID: UUID + var jobID: String + var initialRuntimeGeneration: ReviewRuntimeGeneration var initialEventSubscriptionID: Int var networkTask: Task var signalCoordinator: ReviewNetworkSignalCoordinator @@ -2009,6 +2282,14 @@ private struct ReviewWorkerInputs { await signalCoordinator.cancel() await networkTask.value await queue.finish() + await runtimeReplacementCoordinator.finishRegisteredParticipant( + registrationID: runtimeReplacementRegistrationID, + jobID: jobID + ) + await runtimeReplacementCoordinator.unregister( + registrationID: runtimeReplacementRegistrationID, + jobID: jobID + ) } } @@ -2051,6 +2332,148 @@ private actor ReviewWorkerRecoveryInterruptionSource { } } +@MainActor +package final class ReviewRuntimeReplacementCoordinator { + private struct WorkerRegistration { + let id: UUID + let jobID: String + var attemptID: String + var runtimeGeneration: ReviewRuntimeGeneration + let recoverySource: ReviewWorkerRecoveryInterruptionSource + } + + private var workerRegistrationsByJobID: [String: WorkerRegistration] = [:] + private var replacementsBySourceGeneration: [ + ReviewRuntimeGeneration: ReviewRuntimeRecoveryReplacement + ] = [:] + + package init() {} + + fileprivate func register( + jobID: String, + attemptID: String, + runtimeGeneration: ReviewRuntimeGeneration, + recoverySource: ReviewWorkerRecoveryInterruptionSource + ) -> UUID { + let registration = WorkerRegistration( + id: UUID(), + jobID: jobID, + attemptID: attemptID, + runtimeGeneration: runtimeGeneration, + recoverySource: recoverySource + ) + workerRegistrationsByJobID[jobID] = registration + return registration.id + } + + fileprivate func update( + registrationID: UUID, + jobID: String, + attemptID: String, + runtimeGeneration: ReviewRuntimeGeneration + ) { + guard var registration = workerRegistrationsByJobID[jobID], + registration.id == registrationID + else { + return + } + registration.attemptID = attemptID + registration.runtimeGeneration = runtimeGeneration + workerRegistrationsByJobID[jobID] = registration + } + + fileprivate func unregister(registrationID: UUID, jobID: String) { + guard workerRegistrationsByJobID[jobID]?.id == registrationID else { + return + } + workerRegistrationsByJobID.removeValue(forKey: jobID) + } + + fileprivate func participant( + jobID: String, + attemptID: String, + sourceGeneration: ReviewRuntimeGeneration + ) -> ReviewRuntimeReplacementParticipant? { + guard let registration = workerRegistrationsByJobID[jobID], + registration.attemptID == attemptID, + registration.runtimeGeneration == sourceGeneration + else { + return nil + } + return .init(jobID: jobID, attemptID: attemptID) + } + + fileprivate func install(_ replacement: ReviewRuntimeRecoveryReplacement) { + replacementsBySourceGeneration[replacement.sourceGeneration] = replacement + } + + fileprivate func replacement( + sourceGeneration: ReviewRuntimeGeneration + ) -> ReviewRuntimeRecoveryReplacement? { + replacementsBySourceGeneration[sourceGeneration] + } + + fileprivate func beginRecovery( + participant: ReviewRuntimeReplacementParticipant, + owner: ReviewActiveAttempt, + trigger: ReviewAttemptRecoveryTrigger, + interrupt: @escaping @Sendable ( + CodexReviewBackendModel.Review.Run, + CodexReviewBackendModel.CancellationReason + ) async throws -> Void, + forceClose: @escaping @Sendable () async throws -> Void + ) async -> Bool { + guard let registration = workerRegistrationsByJobID[participant.jobID], + registration.attemptID == participant.attemptID + else { + return false + } + await registration.recoverySource.start(for: owner) { + try await owner.admission.beginRecovery( + trigger: trigger, + interrupt: interrupt, + forceClose: forceClose + ) + } + return await owner.admission.waitForInterruptionAdmission() + == .recoverableTransition(trigger) + } + + fileprivate func recordNetworkRestoration( + sourceGeneration: ReviewRuntimeGeneration + ) { + replacementsBySourceGeneration[sourceGeneration]?.recordNetworkRestoration() + } + + fileprivate func finishParticipant( + jobID: String, + sourceGeneration: ReviewRuntimeGeneration + ) { + guard let replacement = replacementsBySourceGeneration[sourceGeneration] else { + return + } + replacement.finishParticipant(jobID: jobID) + if replacement.hasRemainingParticipants == false { + replacementsBySourceGeneration.removeValue(forKey: sourceGeneration) + } + } + + fileprivate func finishRegisteredParticipant( + registrationID: UUID, + jobID: String + ) { + guard let registration = workerRegistrationsByJobID[jobID], + registration.id == registrationID + else { + return + } + finishParticipant( + jobID: jobID, + sourceGeneration: registration.runtimeGeneration + ) + } +} + private actor ReviewWorkerInputQueue { private enum Delivery { case input(ReviewWorkerInput) diff --git a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift index 17619f52..020de9f0 100644 --- a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift +++ b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift @@ -596,7 +596,7 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { let closeTask: Task switch lifecycleState { case .open: - if purpose == .recoveryReplacement { + if purpose == .recoveryReplacement || purpose == .restartSameAccount { connectionStreamFailure = .ownerForcedConnectionClose( .connection("Review connection was force-closed by its runtime owner.") ) diff --git a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift index 60c1f3fc..acd23de9 100644 --- a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift +++ b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift @@ -161,6 +161,16 @@ public extension CodexReviewStore { @MainActor private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { + private struct AttemptRoute { + let generation: ReviewRuntimeGeneration + let runtime: LiveRuntimeLifecycleHandle + } + + private struct RecoveryRoute { + let handoff: ReviewRecoveryHandoff + let source: AttemptRoute + } + typealias MCPHTTPServerFactory = @MainActor @Sendable ( CodexReviewStore, CodexReviewMCPHTTPServer.Configuration @@ -169,8 +179,9 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { let seed: CodexReviewStoreSeed private var client: AppServerClient? - private var appServerBackend: AppServerCodexReviewBackend? private var activeRuntimeHandle: LiveRuntimeLifecycleHandle? + private var attemptRoutesByAttemptID: [String: AttemptRoute] = [:] + private var recoveryRoutesByAttemptID: [String: RecoveryRoute] = [:] private var acceptsRuntimeRequests = false private var loginChallenge: CodexReviewBackendModel.Login.Challenge? private var loginBackend: AppServerCodexReviewBackend? @@ -197,6 +208,10 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private let appServerRuntimeFactory: AppServerRuntimeFactory private weak var attachedStore: CodexReviewStore? + private var appServerBackend: AppServerCodexReviewBackend? { + activeRuntimeHandle?.backend + } + init( environment: [String: String] = ProcessInfo.processInfo.environment, runtimePreferences: CodexReviewRuntime.Preferences = .defaults, @@ -378,7 +393,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { } func prepareRuntime( - generation _: ReviewRuntimeGeneration, + generation: ReviewRuntimeGeneration, purpose _: ReviewRuntimeTransitionPurpose ) async throws -> PreparedRuntime { logger.info("Preparing review runtime") @@ -388,6 +403,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { let settings = try await Self.monitorSettings(from: runtime.backend.readSettings()) let handle = LiveRuntimeLifecycleHandle( owner: self, + generation: generation, client: runtime.client, backend: runtime.backend, snapshot: .init( @@ -436,7 +452,6 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { activeRuntimeHandle = handle acceptsRuntimeRequests = true client = handle.client - appServerBackend = handle.backend settingsSnapshot = handle.snapshot.settings observeAuthNotifications( client: handle.client, @@ -1014,17 +1029,22 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { _ request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { - guard acceptsRuntimeRequests, let appServerBackend else { + guard acceptsRuntimeRequests, let runtime = activeRuntimeHandle else { throw CodexReviewAPI.Error.io("Review runtime is not running.") } - return try await appServerBackend.startReview(request, admission: admission) + let attempt = try await runtime.backend.startReview(request, admission: admission) + attemptRoutesByAttemptID[attempt.run.attemptID] = .init( + generation: runtime.generation, + runtime: runtime + ) + return attempt } func interruptReview(_ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason) async throws { - guard acceptsRuntimeRequests, let appServerBackend else { - throw CodexReviewAPI.Error.io("Review runtime is not running.") + guard let route = attemptRoutesByAttemptID[run.attemptID] else { + throw CodexReviewAPI.Error.io("Review attempt route is unavailable.") } - try await appServerBackend.interruptReview(run, reason: reason) + try await route.runtime.backend.interruptReview(run, reason: reason) } func forceCloseReviewConnection() async throws { @@ -1037,10 +1057,16 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { func prepareReviewRecovery( _ candidate: ReviewRecoveryCandidate ) async throws -> ReviewRecoveryHandoff { - guard let appServerBackend else { - throw CodexReviewAPI.Error.io("Review runtime is not running.") - } - return try await appServerBackend.prepareReviewRecovery(candidate) + let attemptID = candidate.resolved.run.attemptID + guard let source = attemptRoutesByAttemptID[attemptID] else { + throw CodexReviewAPI.Error.io("Review recovery source route is unavailable.") + } + let handoff = try await source.runtime.backend.prepareReviewRecovery(candidate) + recoveryRoutesByAttemptID[attemptID] = .init( + handoff: handoff, + source: source + ) + return handoff } func resumeReviewRecovery( @@ -1048,21 +1074,46 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { - guard acceptsRuntimeRequests, let appServerBackend else { + let sourceAttemptID = handoff.candidate.resolved.run.attemptID + guard let source = recoveryRoutesByAttemptID[sourceAttemptID], + source.handoff == handoff + else { + throw CodexReviewAPI.Error.io("Review recovery handoff route is unavailable.") + } + guard acceptsRuntimeRequests, let destination = activeRuntimeHandle else { throw CodexReviewAPI.Error.io("Review runtime is not running.") } - return try await appServerBackend.resumeReviewRecovery( + if handoff.candidate.trigger == .sameAccountRestart, + destination.generation == source.source.generation { + throw CodexReviewAPI.Error.io( + "Same-account recovery requires a replacement runtime generation." + ) + } + let attempt = try await destination.backend.resumeReviewRecovery( handoff, request: request, admission: admission ) + attemptRoutesByAttemptID.removeValue(forKey: sourceAttemptID) + recoveryRoutesByAttemptID.removeValue(forKey: sourceAttemptID) + attemptRoutesByAttemptID[attempt.run.attemptID] = .init( + generation: destination.generation, + runtime: destination + ) + return attempt } func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { - guard let appServerBackend else { - throw ReviewRuntimeCloseFailure.cleanup("Review runtime is not running.") + guard let route = attemptRoutesByAttemptID[run.attemptID] else { + throw ReviewRuntimeCloseFailure.cleanup("Review attempt route is unavailable.") + } + defer { + if attemptRoutesByAttemptID[run.attemptID]?.runtime === route.runtime { + attemptRoutesByAttemptID.removeValue(forKey: run.attemptID) + recoveryRoutesByAttemptID.removeValue(forKey: run.attemptID) + } } - try await appServerBackend.cleanupReview(run) + try await route.runtime.backend.cleanupReview(run) } @discardableResult @@ -1717,6 +1768,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { @MainActor private final class LiveRuntimeLifecycleHandle: RuntimeLifecycleHandle { + fileprivate let generation: ReviewRuntimeGeneration fileprivate let client: AppServerClient fileprivate let backend: AppServerCodexReviewBackend fileprivate let snapshot: RuntimePublicationSnapshot @@ -1727,11 +1779,13 @@ private final class LiveRuntimeLifecycleHandle: RuntimeLifecycleHandle { init( owner: LiveCodexReviewStoreBackend, + generation: ReviewRuntimeGeneration, client: AppServerClient, backend: AppServerCodexReviewBackend, snapshot: RuntimePublicationSnapshot ) { self.owner = owner + self.generation = generation self.client = client self.backend = backend self.snapshot = snapshot diff --git a/Sources/CodexReviewTesting/TestSupport.swift b/Sources/CodexReviewTesting/TestSupport.swift index f0b0d6b9..3e831ff7 100644 --- a/Sources/CodexReviewTesting/TestSupport.swift +++ b/Sources/CodexReviewTesting/TestSupport.swift @@ -193,6 +193,8 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { private var commands: [Command] = [] private var nextRun: CodexReviewBackendModel.Review.Run private var nextRecoveredRun: CodexReviewBackendModel.Review.Run? + private var queuedRuns: [CodexReviewBackendModel.Review.Run] = [] + private var queuedRecoveredRuns: [CodexReviewBackendModel.Review.Run] = [] private var interruptFailureMessage: String? private var recoveryFailureMessage: String? private var cleanupFailure: ReviewRuntimeCloseFailure? @@ -234,6 +236,7 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { self.settings = settings self.auth = auth self.nextRun = nextRun + self.queuedRuns = [nextRun] } package func recordedCommands() -> [Command] { @@ -284,6 +287,14 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { nextRecoveredRun = run } + package func enqueueRun(_ run: CodexReviewBackendModel.Review.Run) { + queuedRuns.append(run) + } + + package func enqueueRecoveredRun(_ run: CodexReviewBackendModel.Review.Run) { + queuedRecoveredRuns.append(run) + } + package func waitForStartReview() async { if commands.contains(where: { if case .startReview = $0 { @@ -524,6 +535,7 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { _ request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { + let run = queuedRuns.isEmpty ? nextRun : queuedRuns.removeFirst() try await admission.admitThreadStartDispatch() commands.append(.startReview(request)) let waiters = Array(startReviewWaiters.values) @@ -532,10 +544,10 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { waiter.resume() } let provisionalRun = CodexReviewBackendModel.Review.Run( - attemptID: nextRun.attemptID, - threadID: nextRun.threadID, - reviewThreadID: nextRun.threadID, - model: nextRun.model + attemptID: run.attemptID, + threadID: run.threadID, + reviewThreadID: run.threadID, + model: run.model ) await admission.recordPreparedThread(provisionalRun) do { @@ -547,8 +559,8 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { if let startReviewGate { await startReviewGate.wait() } - await admission.recordActiveRun(nextRun) - return .init(run: nextRun, events: eventMailbox(for: nextRun)) + await admission.recordActiveRun(run) + return .init(run: run, events: eventMailbox(for: run)) } package func interruptReview(_ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason) async throws { @@ -624,7 +636,9 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { ) async throws -> BackendReviewAttempt { let token = handoff.token let run = token.interruptedRun - let recoveredRun = nextRecoveredRun ?? .init( + let recoveredRun = queuedRecoveredRuns.isEmpty == false + ? queuedRecoveredRuns.removeFirst() + : nextRecoveredRun ?? .init( attemptID: "attempt-recovered", threadID: run.threadID, turnID: "turn-recovered", @@ -1232,7 +1246,7 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { onClose: { [weak self] purpose in guard let self else { return } self.isActive = false - if purpose == .recoveryReplacement { + if purpose == .recoveryReplacement || purpose == .restartSameAccount { try? await self.reviewBackend.forceCloseReviewConnection() } await self.runtimeCloseOperation() diff --git a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift index 553987ef..880c9cfa 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift @@ -20,7 +20,8 @@ struct CodexReviewStoreLifecycleTests { purpose: .stop, task: transitionTask, record: ReviewRuntimeTransitionRecord(), - sourceRuntime: nil + sourceRuntime: nil, + recoveryReplacement: nil ) await store.performRuntimeAcquisitionForTesting( @@ -150,7 +151,7 @@ struct CodexReviewStoreLifecycleTests { await backend.waitForRuntimePreparation() let ownsReplacementTask: Bool - if case .transitioning(_, .restartSameAccount, _, _, _) = store.runtimeState { + if case .transitioning(_, .restartSameAccount, _, _, _, _) = store.runtimeState { ownsReplacementTask = true } else { ownsReplacementTask = false @@ -185,6 +186,95 @@ struct CodexReviewStoreLifecycleTests { await store.stop() } + @Test func graceForceCloseCancelsTargetAndResumesSiblingOnOneReplacement() async throws { + let targetRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-target", + threadID: "thread-target", + turnID: "turn-target", + reviewThreadID: "review-target" + ) + let siblingRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-sibling", + threadID: "thread-sibling", + turnID: "turn-sibling", + reviewThreadID: "review-sibling" + ) + let recoveredSiblingRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-sibling-recovered", + threadID: "thread-sibling", + turnID: "turn-sibling-recovered", + reviewThreadID: "review-sibling" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: targetRun) + await reviewBackend.enqueueRun(siblingRun) + await reviewBackend.enqueueRecoveredRun(recoveredSiblingRun) + let mcpOwner = TestingMCPServerLifecycleOwner( + serverURL: URL(string: "http://127.0.0.1:19431/mcp") + ) + let backend = TestingCodexReviewStoreBackend( + reviewBackend: reviewBackend, + mcpServerLifecycle: mcpOwner + ) + let jobIDs = SequentialJobIDs(["job-target", "job-sibling"]) + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { jobIDs.next() }), + reviewRuntimeClosePolicy: .init( + terminalGrace: .seconds(10), + sleep: { _ in } + ) + ) + await store.start() + let sourceRuntime = try #require(backend.lastPreparedRuntimeHandle) + _ = try await store.startReview( + sessionID: "session-target", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges), + waitTimeout: .milliseconds(20) + ) + _ = try await store.startReview( + sessionID: "session-sibling", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges), + waitTimeout: .milliseconds(20) + ) + + let cancellation = Task { @MainActor in + try await store.cancelReview( + jobID: "job-target", + cancellation: .mcpClient(message: "Stop target") + ) + } + await reviewBackend.waitForResumeReviewRecovery() + await reviewBackend.yield( + .completed(summary: "Done", result: "sibling result"), + for: recoveredSiblingRun + ) + let targetCancellation = try await cancellation.value + let sibling = try await store.awaitReview( + sessionID: "session-sibling", + jobID: "job-sibling" + ) + + #expect(targetCancellation.cancelled) + #expect(try store.readReview(jobID: "job-target").core.lifecycle.terminal == .interrupted( + .requested(.mcpClient(message: "Stop target")) + )) + #expect(sibling.core.lifecycle.status == .succeeded) + #expect(sibling.core.run.turnID == "turn-sibling-recovered") + #expect(sourceRuntime.closeCallCount == 1) + #expect(sourceRuntime.waitUntilClosedCallCount == 1) + #expect(sourceRuntime.closePurposes == [.recoveryReplacement]) + #expect(backend.startRequests == [false, false]) + #expect(mcpOwner.prepareCallCount == 1) + #expect(mcpOwner.activateCallCount == 1) + #expect(mcpOwner.stopCallCount == 0) + let commands = await reviewBackend.recordedCommands() + #expect(commands.filter { if case .forceCloseReviewConnection = $0 { true } else { false } }.count == 1) + #expect(commands.filter { if case .prepareReviewRecovery = $0 { true } else { false } }.count == 1) + #expect(commands.filter { if case .resumeReviewRecovery = $0 { true } else { false } }.count == 1) + + await store.stop() + } + @Test func stopInvalidatesHeldRestartBeforeReplacementCanPublish() async throws { let endpoint = try #require(URL(string: "http://127.0.0.1:19422/mcp")) let mcpOwner = TestingMCPServerLifecycleOwner(serverURL: endpoint) @@ -818,3 +908,18 @@ private actor StoreCloseCompletion { completeValue } } + +private final class SequentialJobIDs: @unchecked Sendable { + private let lock = NSLock() + private var values: [String] + + init(_ values: [String]) { + self.values = values + } + + func next() -> String { + lock.withLock { + values.removeFirst() + } + } +} From 54b04629f6bfdaabef2d8e3cc88269a9ee344d21 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 09:48:48 +0900 Subject: [PATCH 34/38] Complete shared runtime recovery replacement --- .../CodexReview/ReviewRuntimeLifecycle.swift | 70 +- .../CodexReview/Store/CodexReviewStore.swift | 40 +- .../Store/CodexReviewStoreBackend.swift | 14 +- .../Store/CodexReviewStoreCancellation.swift | 1 + .../Store/CodexReviewStoreReviews.swift | 364 ++++++-- .../Store/CodexReviewStoreTesting.swift | 34 + .../PreviewCodexReviewStoreBackend.swift | 18 +- .../AppServerCodexReviewBackend.swift | 3 + Sources/CodexReviewHost/CodexReviewHost.swift | 13 + .../LiveCodexReviewStoreBackend.swift | 182 +++- Sources/CodexReviewTesting/TestSupport.swift | 28 +- .../CodexReviewHostTests.swift | 409 +++++++++ .../CodexReviewStoreLifecycleTests.swift | 869 +++++++++++++++++- Tests/ReviewUITests/ReviewUITests.swift | 6 +- 14 files changed, 1931 insertions(+), 120 deletions(-) diff --git a/Sources/CodexReview/ReviewRuntimeLifecycle.swift b/Sources/CodexReview/ReviewRuntimeLifecycle.swift index 83e1790c..b0ae506e 100644 --- a/Sources/CodexReview/ReviewRuntimeLifecycle.swift +++ b/Sources/CodexReview/ReviewRuntimeLifecycle.swift @@ -209,10 +209,14 @@ package final class ReviewCloseFailureLedger { guard failures.isEmpty == false else { return } - forceCloseFailureJobIDs.insert(jobID) + markForceCloseFailureOwnership(jobIDs: [jobID]) self.failures.append(contentsOf: failures) } + package func markForceCloseFailureOwnership(jobIDs: [String]) { + forceCloseFailureJobIDs.formUnion(jobIDs) + } + package func ownsForceCloseFailure(for jobID: String) -> Bool { forceCloseFailureJobIDs.contains(jobID) } @@ -257,12 +261,23 @@ package struct ReviewRuntimePreparationFailure: LocalizedError, Sendable { } package struct ReviewRuntimeReplacementParticipant: Equatable, Sendable { + package enum Phase: Equatable, Sendable { + case eligible + case recovering + case suppressed + case finished + } + package let jobID: String package let attemptID: String + package let registrationID: UUID + package var phase: Phase - package init(jobID: String, attemptID: String) { + package init(jobID: String, attemptID: String, registrationID: UUID) { self.jobID = jobID self.attemptID = attemptID + self.registrationID = registrationID + self.phase = .eligible } } @@ -321,7 +336,6 @@ package final class ReviewRuntimeRecoveryReplacement { private var outcomeContinuations: [AsyncStream.Continuation] = [] private var networkRestorationWasRecorded = false private var networkRestorationContinuations: [AsyncStream.Continuation] = [] - private var remainingParticipantJobIDs: Set package init( sourceGeneration: ReviewRuntimeGeneration, @@ -339,20 +353,54 @@ package final class ReviewRuntimeRecoveryReplacement { self.retainedServerURL = retainedServerURL self.trigger = trigger self.participants = participants - self.remainingParticipantJobIDs = Set(participants.map(\.jobID)) } - package func suppressParticipant(jobID: String) { - participants.removeAll { $0.jobID == jobID } - remainingParticipantJobIDs.remove(jobID) + @discardableResult + package func beginParticipantRecovery( + jobID: String, + registrationID: UUID + ) -> Bool { + guard let index = participants.firstIndex(where: { + $0.jobID == jobID && $0.registrationID == registrationID + }), participants[index].phase == .eligible else { + return false + } + participants[index].phase = .recovering + return true } - package func finishParticipant(jobID: String) { - remainingParticipantJobIDs.remove(jobID) + package func suppressParticipant(jobID: String, registrationID: UUID) { + guard let index = participants.firstIndex(where: { + $0.jobID == jobID && $0.registrationID == registrationID + }) else { + return + } + participants[index].phase = .suppressed } - package var hasRemainingParticipants: Bool { - remainingParticipantJobIDs.isEmpty == false + package func finishParticipant(jobID: String, registrationID: UUID) { + guard let index = participants.firstIndex(where: { + $0.jobID == jobID && $0.registrationID == registrationID + }) else { + return + } + participants[index].phase = .finished + } + + package var forceCloseObserverJobIDs: [String] { + var jobIDs: [String] = [] + switch trigger { + case .explicitCancellation(let targetJobID): + jobIDs.append(targetJobID) + case .recoverableNetwork(let initiatingJobID): + jobIDs.append(initiatingJobID) + case .sameAccountRestart: + break + } + for participant in participants where jobIDs.contains(participant.jobID) == false { + jobIDs.append(participant.jobID) + } + return jobIDs } package func finishSourceClose( diff --git a/Sources/CodexReview/Store/CodexReviewStore.swift b/Sources/CodexReview/Store/CodexReviewStore.swift index efba1d9e..75e847e6 100644 --- a/Sources/CodexReview/Store/CodexReviewStore.swift +++ b/Sources/CodexReview/Store/CodexReviewStore.swift @@ -44,6 +44,8 @@ public final class CodexReviewStore { @ObservationIgnored package var reviewCleanupPreparationForTesting: (@MainActor @Sendable () async -> Void)? @ObservationIgnored package var reviewTerminalPublicationPreparationForTesting: (@MainActor @Sendable () async -> Void)? @ObservationIgnored package var runtimeForceCloseReceiptRecordedForTesting: (@MainActor @Sendable () async -> Void)? + @ObservationIgnored package var runtimeReplacementEnrollmentPreparationForTesting: (@MainActor @Sendable () async -> Void)? + @ObservationIgnored package var reviewCancellationBarrierPreparationForTesting: (@MainActor @Sendable () async -> Void)? @ObservationIgnored package var storeCommandRegistry = ReviewStoreCommandRegistry() @ObservationIgnored package var closeCallerCount = 0 @ObservationIgnored package var closeCallerWaiters: [CloseCallerWaiter] = [] @@ -56,7 +58,7 @@ public final class CodexReviewStore { .init(rawValue: 0) ) @ObservationIgnored package var lastRuntimeTransitionRecord: ReviewRuntimeTransitionRecord? - @ObservationIgnored package let runtimeReplacementCoordinator = ReviewRuntimeReplacementCoordinator() + @ObservationIgnored package let runtimeWorkerRegistry = ReviewRuntimeWorkerRegistry() package init( backend: any CodexReviewStoreBackend = PreviewCodexReviewStoreBackend(), @@ -304,10 +306,6 @@ public final class CodexReviewStore { case .acquiring, .running, .transitioning, .failed: break } - if case .transitioning(_, _, _, _, _, let replacement) = previousState { - replacement?.finish(.superseded(.stop)) - replacement?.recordNetworkRestoration() - } let invalidatedGeneration = previousState.generation.successor() let record = ReviewRuntimeTransitionRecord() let task = Task { @MainActor [weak self] in @@ -443,10 +441,6 @@ public final class CodexReviewStore { failureLedger: ReviewCloseFailureLedger ) async -> Result { let runningRuntime = previousRuntimeState.runtimeForClose - if case .transitioning(_, _, _, _, _, let replacement) = previousRuntimeState { - replacement?.finish(.superseded(.applicationClose)) - replacement?.recordNetworkRestoration() - } if let lastRuntimeTransitionRecord { failureLedger.importReceipts(from: lastRuntimeTransitionRecord) } @@ -618,12 +612,27 @@ public final class CodexReviewStore { record.merge(previousRecord) case .running(_, let runtime, _): await stopPublishedRuntime(runtime, record: record) - case .transitioning(_, _, let task, let previousRecord, _, let replacement): + case .transitioning( + _, + _, + let task, + let previousRecord, + let sourceRuntime, + let replacement + ): + await performPublishedRuntimeSemanticStop(record: record) replacement?.finish(.superseded(.stop)) replacement?.recordNetworkRestoration() task.cancel() await task.value record.merge(previousRecord) + if let sourceRuntime { + await closeAppServerRuntime( + sourceRuntime, + purpose: .stop, + record: record + ) + } await stopPreparedMCPServer(record: record) case .failed: await stopPreparedMCPServer(record: record) @@ -821,6 +830,7 @@ public final class CodexReviewStore { ) async { defer { lastRuntimeTransitionRecord = record } var preparedRuntime: PreparedRuntime? + await runtimeReplacementEnrollmentPreparationForTesting?() await enrollRuntimeReplacementParticipants(replacement) if let sourceRuntime = replacement.sourceRuntime { await sourceRuntime.handle.closeAdmission() @@ -829,14 +839,14 @@ public final class CodexReviewStore { purpose: purpose ) let consumedFailures = sourceRuntime.closeRecord.consumeFailures() - if let targetJobID = replacement.trigger.targetJobID { - record.recordForceCloseFailures( - consumedFailures, - jobID: targetJobID + if consumedFailures.isEmpty == false { + record.markForceCloseFailureOwnership( + jobIDs: replacement.forceCloseObserverJobIDs ) - } else { record.record(contentsOf: consumedFailures) + applicationCloseFailureLedger?.importReceipts(from: record) } + lastRuntimeTransitionRecord = record if let firstFailure = closeResult.failures.first { replacement.finishSourceClose(.failure( reviewRuntimeCloseFailure(from: firstFailure) diff --git a/Sources/CodexReview/Store/CodexReviewStoreBackend.swift b/Sources/CodexReview/Store/CodexReviewStoreBackend.swift index 611fc975..9378795e 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreBackend.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreBackend.swift @@ -51,7 +51,11 @@ package protocol CodexReviewStoreBackend: CodexReviewSettingsBackend, Sendable { _ request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt - func interruptReview(_ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason) async throws + func interruptReview( + _ run: CodexReviewBackendModel.Review.Run, + admission: ReviewStartAdmission, + reason: CodexReviewBackendModel.CancellationReason + ) async throws func forceCloseReviewConnection() async throws func prepareReviewRecovery( _ candidate: ReviewRecoveryCandidate @@ -61,5 +65,13 @@ package protocol CodexReviewStoreBackend: CodexReviewSettingsBackend, Sendable { request: CodexReviewBackendModel.Review.Start, admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt + func commitResumedReviewRecovery( + _ handoff: ReviewRecoveryHandoff, + recoveredRun: CodexReviewBackendModel.Review.Run + ) throws + func discardResumedReviewRecovery( + _ handoff: ReviewRecoveryHandoff, + recoveredRun: CodexReviewBackendModel.Review.Run + ) async throws func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws } diff --git a/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift b/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift index c64e6ba3..0cc8c669 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift @@ -103,6 +103,7 @@ extension CodexReviewStore { reason: ReviewCancellation = .system(message: "Review Store closed."), failureLedger: ReviewCloseFailureLedger ) async -> ReviewCloseCancellationOutcome { + await reviewCancellationBarrierPreparationForTesting?() let activeJobIDs = activeReviewJobIDsInRegistrationOrder var failedJobIDs: Set = [] for jobID in activeJobIDs { diff --git a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift index 13e421ac..22b40ef9 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift @@ -368,7 +368,7 @@ extension CodexReviewStore { ) let eventSource = ReviewWorkerEventSource(queue: queue) let recoveryInterruptionSource = ReviewWorkerRecoveryInterruptionSource(queue: queue) - let runtimeReplacementRegistrationID = runtimeReplacementCoordinator.register( + let runtimeWorkerRegistrationID = runtimeWorkerRegistry.register( jobID: jobID, attemptID: owner.run.attemptID, runtimeGeneration: runtimeGeneration, @@ -388,8 +388,8 @@ extension CodexReviewStore { networkStatusTracker: tracker, eventSource: eventSource, recoveryInterruptionSource: recoveryInterruptionSource, - runtimeReplacementCoordinator: runtimeReplacementCoordinator, - runtimeReplacementRegistrationID: runtimeReplacementRegistrationID, + runtimeWorkerRegistry: runtimeWorkerRegistry, + runtimeWorkerRegistrationID: runtimeWorkerRegistrationID, jobID: jobID, initialRuntimeGeneration: runtimeGeneration, initialEventSubscriptionID: initialEventSubscriptionID, @@ -652,7 +652,11 @@ extension CodexReviewStore { return try await admission.cancel( cancellation, interrupt: { run, reason in - try await backend.interruptReview(run, reason: reason) + try await backend.interruptReview( + run, + admission: admission, + reason: reason + ) }, forceClose: { @MainActor [weak self] in guard let self else { @@ -699,7 +703,7 @@ extension CodexReviewStore { .waitingForRecovery, .replacementStart, .terminal, nil: return nil } - return runtimeReplacementCoordinator.participant( + return runtimeWorkerRegistry.participant( jobID: jobID, attemptID: active.run.attemptID, sourceGeneration: sourceGeneration @@ -714,7 +718,6 @@ extension CodexReviewStore { trigger: trigger, participants: participants ) - runtimeReplacementCoordinator.install(replacement) return replacement } @@ -727,7 +730,10 @@ extension CodexReviewStore { job.isTerminal == false, job.cancellationRequested == false else { - replacement.suppressParticipant(jobID: participant.jobID) + runtimeWorkerRegistry.suppressParticipant( + participant, + in: replacement + ) continue } let active: ReviewActiveAttempt @@ -738,7 +744,10 @@ extension CodexReviewStore { case .resolvingRecovery(let current) where current.run.attemptID == participant.attemptID: active = current default: - replacement.suppressParticipant(jobID: participant.jobID) + runtimeWorkerRegistry.suppressParticipant( + participant, + in: replacement + ) continue } let recoveryTrigger: ReviewAttemptRecoveryTrigger = if case .recoverableNetwork( @@ -748,12 +757,17 @@ extension CodexReviewStore { } else { .sameAccountRestart } - let didInstall = await runtimeReplacementCoordinator.beginRecovery( + let didInstall = await runtimeWorkerRegistry.beginRecovery( + replacement: replacement, participant: participant, owner: active, trigger: recoveryTrigger, interrupt: { run, reason in - try await backend.interruptReview(run, reason: reason) + try await backend.interruptReview( + run, + admission: active.admission, + reason: reason + ) }, forceClose: { try await replacement.waitForSourceClose().get() @@ -764,7 +778,10 @@ extension CodexReviewStore { sameAttempt(current, active) { reviewAttemptOwnerships[participant.jobID] = .active(active) } - replacement.suppressParticipant(jobID: participant.jobID) + runtimeWorkerRegistry.suppressParticipant( + participant, + in: replacement + ) } } } @@ -805,7 +822,10 @@ extension CodexReviewStore { ) = runtimeState, purpose == .recoveryReplacement || purpose == .restartSameAccount { if case .explicitCancellation = trigger { - replacement.suppressParticipant(jobID: jobID) + runtimeWorkerRegistry.suppressParticipant( + jobID: jobID, + in: replacement + ) } let sourceCloseResult = await replacement.waitForSourceClose() await runtimeForceCloseReceiptRecordedForTesting?() @@ -1385,8 +1405,9 @@ extension CodexReviewStore { ) else { continue } - runtimeReplacementCoordinator.recordNetworkRestoration( - sourceGeneration: activeRuntimeGeneration + runtimeWorkerRegistry.recordNetworkRestoration( + registrationID: inputs.runtimeWorkerRegistrationID, + jobID: job.id ) switch try await restartReviewAfterRuntimeRecovery( job: job, @@ -1442,7 +1463,11 @@ extension CodexReviewStore { try await active.admission.beginRecovery( trigger: .recoverableNetworkLoss, interrupt: { run, reason in - try await backend.interruptReview(run, reason: reason) + try await backend.interruptReview( + run, + admission: active.admission, + reason: reason + ) }, forceClose: { @MainActor [weak self] in guard let self else { @@ -1596,10 +1621,12 @@ extension CodexReviewStore { await inputs.networkStatusTracker.currentStatus() != .satisfied { return .continueWaiting } + let replacement = inputs.runtimeWorkerRegistry.replacement( + registrationID: inputs.runtimeWorkerRegistrationID, + jobID: job.id + ) let destinationRuntimeGeneration: ReviewRuntimeGeneration - if let replacement = inputs.runtimeReplacementCoordinator.replacement( - sourceGeneration: sourceRuntimeGeneration - ) { + if let replacement { var outcomes = replacement.replacementOutcomes().makeAsyncIterator() guard let outcome = await outcomes.next() else { return .finished @@ -1610,15 +1637,17 @@ extension CodexReviewStore { case .failed(let failure): markReviewFailed(job, message: failure.localizedDescription) reviewAttemptOwnerships[job.id] = .terminal - inputs.runtimeReplacementCoordinator.finishParticipant( + inputs.runtimeWorkerRegistry.finishParticipant( + registrationID: inputs.runtimeWorkerRegistrationID, jobID: job.id, - sourceGeneration: sourceRuntimeGeneration + replacement: replacement ) return .finished case .superseded: - inputs.runtimeReplacementCoordinator.finishParticipant( + inputs.runtimeWorkerRegistry.finishParticipant( + registrationID: inputs.runtimeWorkerRegistrationID, jobID: job.id, - sourceGeneration: sourceRuntimeGeneration + replacement: replacement ) throw CancellationError() } @@ -1642,10 +1671,10 @@ extension CodexReviewStore { guard case .waitingForRecovery(let revalidatedHandoff) = reviewAttemptOwnerships[job.id], revalidatedHandoff == handoff else { - _ = try await recoveredAdmission.cancel( - job.core.lifecycle.cancellation ?? .system(), - interrupt: { _, _ in }, - forceClose: {} + _ = try await cancel( + admission: recoveredAdmission, + cancellation: job.core.lifecycle.cancellation ?? .system(), + jobID: job.id ) _ = await registered.task.result return .finished @@ -1656,10 +1685,20 @@ extension CodexReviewStore { ) try await recoveredAdmission.activateStart(registered.id) let result = await registered.task.result + let recoveredAttempt = try result.get() guard case .replacementStart(let currentHandoff, let currentStart) = reviewAttemptOwnerships[job.id], currentHandoff == handoff, currentStart.id == registered.id else { + do { + try await backend.discardResumedReviewRecovery( + handoff, + recoveredRun: recoveredAttempt.run + ) + } catch { + retainCleanupFailure(error, for: job.id) + throw error + } if job.isTerminal { return .finished } @@ -1667,19 +1706,53 @@ extension CodexReviewStore { message: "Replacement start completed after its ownership changed." ) } - let recoveredAttempt = try result.get() - if destinationRuntimeGeneration != sourceRuntimeGeneration { - guard case .running(let currentGeneration, _, _) = runtimeState, - currentGeneration == destinationRuntimeGeneration, - case .open = lifetimeState - else { - _ = try? await recoveredAdmission.cancel( - job.core.lifecycle.cancellation ?? .system(), - interrupt: { _, _ in }, - forceClose: {} + guard isReviewMutationCurrent(destinationRuntimeGeneration) else { + let ownerCancellation = ReviewAttemptStreamFailure.ownerForcedConnectionClose( + .connection("Recovered review was superseded before Store publication.") + ) + var terminalizationError: (any Error)? + do { + try await recoveredAdmission.recordStreamTerminal(ownerCancellation) + if let terminal = await recoveredAdmission + .terminalCancellationProductTerminal(for: ownerCancellation) { + try applyRecoveryProductTerminal(terminal, to: job) + } else { + await applyStreamProductTerminal(ownerCancellation, to: job) + } + } catch { + terminalizationError = error + } + do { + try await backend.discardResumedReviewRecovery( + handoff, + recoveredRun: recoveredAttempt.run ) - return .finished + } catch { + retainCleanupFailure(error, for: job.id) + if terminalizationError == nil { + terminalizationError = error + } + } + if let terminalizationError { + throw terminalizationError + } + return .finished + } + do { + try backend.commitResumedReviewRecovery( + handoff, + recoveredRun: recoveredAttempt.run + ) + } catch { + do { + try await backend.discardResumedReviewRecovery( + handoff, + recoveredRun: recoveredAttempt.run + ) + } catch { + retainCleanupFailure(error, for: job.id) } + throw error } let active = ReviewActiveAttempt( run: recoveredAttempt.run, @@ -1687,16 +1760,19 @@ extension CodexReviewStore { ) reviewAttemptOwnerships[job.id] = .active(active) applyBackendRun(recoveredAttempt.run, to: job) - inputs.runtimeReplacementCoordinator.update( - registrationID: inputs.runtimeReplacementRegistrationID, + if let replacement { + inputs.runtimeWorkerRegistry.finishParticipant( + registrationID: inputs.runtimeWorkerRegistrationID, + jobID: job.id, + replacement: replacement + ) + } + inputs.runtimeWorkerRegistry.update( + registrationID: inputs.runtimeWorkerRegistrationID, jobID: job.id, attemptID: recoveredAttempt.run.attemptID, runtimeGeneration: destinationRuntimeGeneration ) - inputs.runtimeReplacementCoordinator.finishParticipant( - jobID: job.id, - sourceGeneration: sourceRuntimeGeneration - ) return .recovered( recoveredAttempt, active, @@ -2245,8 +2321,8 @@ private struct ReviewWorkerInputs { var networkStatusTracker: ReviewNetworkStatusTracker var eventSource: ReviewWorkerEventSource var recoveryInterruptionSource: ReviewWorkerRecoveryInterruptionSource - var runtimeReplacementCoordinator: ReviewRuntimeReplacementCoordinator - var runtimeReplacementRegistrationID: UUID + var runtimeWorkerRegistry: ReviewRuntimeWorkerRegistry + var runtimeWorkerRegistrationID: UUID var jobID: String var initialRuntimeGeneration: ReviewRuntimeGeneration var initialEventSubscriptionID: Int @@ -2282,12 +2358,12 @@ private struct ReviewWorkerInputs { await signalCoordinator.cancel() await networkTask.value await queue.finish() - await runtimeReplacementCoordinator.finishRegisteredParticipant( - registrationID: runtimeReplacementRegistrationID, + await runtimeWorkerRegistry.finishRegisteredParticipant( + registrationID: runtimeWorkerRegistrationID, jobID: jobID ) - await runtimeReplacementCoordinator.unregister( - registrationID: runtimeReplacementRegistrationID, + await runtimeWorkerRegistry.unregister( + registrationID: runtimeWorkerRegistrationID, jobID: jobID ) } @@ -2333,18 +2409,39 @@ private actor ReviewWorkerRecoveryInterruptionSource { } @MainActor -package final class ReviewRuntimeReplacementCoordinator { - private struct WorkerRegistration { +package final class ReviewRuntimeWorkerRegistry { + private struct RegistrationWaiter { + let attemptID: String + let continuation: CheckedContinuation + } + + private final class WorkerRegistration { let id: UUID let jobID: String var attemptID: String var runtimeGeneration: ReviewRuntimeGeneration let recoverySource: ReviewWorkerRecoveryInterruptionSource + var replacement: ReviewRuntimeRecoveryReplacement? + + init( + id: UUID, + jobID: String, + attemptID: String, + runtimeGeneration: ReviewRuntimeGeneration, + recoverySource: ReviewWorkerRecoveryInterruptionSource + ) { + self.id = id + self.jobID = jobID + self.attemptID = attemptID + self.runtimeGeneration = runtimeGeneration + self.recoverySource = recoverySource + } } private var workerRegistrationsByJobID: [String: WorkerRegistration] = [:] - private var replacementsBySourceGeneration: [ - ReviewRuntimeGeneration: ReviewRuntimeRecoveryReplacement + private var registrationWaitersByJobID: [String: [RegistrationWaiter]] = [:] + private var anyRegistrationWaitersByJobID: [ + String: [CheckedContinuation] ] = [:] package init() {} @@ -2363,29 +2460,83 @@ package final class ReviewRuntimeReplacementCoordinator { recoverySource: recoverySource ) workerRegistrationsByJobID[jobID] = registration + let anyWaiters = anyRegistrationWaitersByJobID.removeValue(forKey: jobID) ?? [] + for waiter in anyWaiters { + waiter.resume(returning: attemptID) + } + let waiters = registrationWaitersByJobID.removeValue(forKey: jobID) ?? [] + for waiter in waiters { + if waiter.attemptID == attemptID { + waiter.continuation.resume() + } else { + registrationWaitersByJobID[jobID, default: []].append(waiter) + } + } return registration.id } + package func waitForRegistrationForTesting( + jobID: String, + attemptID: String + ) async { + if let registration = workerRegistrationsByJobID[jobID], + registration.attemptID == attemptID { + return + } + await withCheckedContinuation { continuation in + if let registration = workerRegistrationsByJobID[jobID], + registration.attemptID == attemptID { + continuation.resume() + } else { + registrationWaitersByJobID[jobID, default: []].append(.init( + attemptID: attemptID, + continuation: continuation + )) + } + } + } + + package func waitForRegistrationForTesting(jobID: String) async -> String { + if let registration = workerRegistrationsByJobID[jobID] { + return registration.attemptID + } + return await withCheckedContinuation { continuation in + if let registration = workerRegistrationsByJobID[jobID] { + continuation.resume(returning: registration.attemptID) + } else { + anyRegistrationWaitersByJobID[jobID, default: []].append(continuation) + } + } + } + fileprivate func update( registrationID: UUID, jobID: String, attemptID: String, runtimeGeneration: ReviewRuntimeGeneration ) { - guard var registration = workerRegistrationsByJobID[jobID], + guard let registration = workerRegistrationsByJobID[jobID], registration.id == registrationID else { return } registration.attemptID = attemptID registration.runtimeGeneration = runtimeGeneration - workerRegistrationsByJobID[jobID] = registration } fileprivate func unregister(registrationID: UUID, jobID: String) { - guard workerRegistrationsByJobID[jobID]?.id == registrationID else { + guard let registration = workerRegistrationsByJobID[jobID], + registration.id == registrationID + else { return } + if let replacement = registration.replacement { + replacement.finishParticipant( + jobID: jobID, + registrationID: registrationID + ) + registration.replacement = nil + } workerRegistrationsByJobID.removeValue(forKey: jobID) } @@ -2400,20 +2551,27 @@ package final class ReviewRuntimeReplacementCoordinator { else { return nil } - return .init(jobID: jobID, attemptID: attemptID) - } - - fileprivate func install(_ replacement: ReviewRuntimeRecoveryReplacement) { - replacementsBySourceGeneration[replacement.sourceGeneration] = replacement + return .init( + jobID: jobID, + attemptID: attemptID, + registrationID: registration.id + ) } fileprivate func replacement( - sourceGeneration: ReviewRuntimeGeneration + registrationID: UUID, + jobID: String ) -> ReviewRuntimeRecoveryReplacement? { - replacementsBySourceGeneration[sourceGeneration] + guard let registration = workerRegistrationsByJobID[jobID], + registration.id == registrationID + else { + return nil + } + return registration.replacement } fileprivate func beginRecovery( + replacement: ReviewRuntimeRecoveryReplacement, participant: ReviewRuntimeReplacementParticipant, owner: ReviewActiveAttempt, trigger: ReviewAttemptRecoveryTrigger, @@ -2424,10 +2582,16 @@ package final class ReviewRuntimeReplacementCoordinator { forceClose: @escaping @Sendable () async throws -> Void ) async -> Bool { guard let registration = workerRegistrationsByJobID[participant.jobID], - registration.attemptID == participant.attemptID + registration.id == participant.registrationID, + registration.attemptID == participant.attemptID, + replacement.beginParticipantRecovery( + jobID: participant.jobID, + registrationID: participant.registrationID + ) else { return false } + registration.replacement = replacement await registration.recoverySource.start(for: owner) { try await owner.admission.beginRecovery( trigger: trigger, @@ -2440,22 +2604,59 @@ package final class ReviewRuntimeReplacementCoordinator { } fileprivate func recordNetworkRestoration( - sourceGeneration: ReviewRuntimeGeneration + registrationID: UUID, + jobID: String ) { - replacementsBySourceGeneration[sourceGeneration]?.recordNetworkRestoration() + replacement(registrationID: registrationID, jobID: jobID)? + .recordNetworkRestoration() } - fileprivate func finishParticipant( + fileprivate func suppressParticipant( + _ participant: ReviewRuntimeReplacementParticipant, + in replacement: ReviewRuntimeRecoveryReplacement + ) { + replacement.suppressParticipant( + jobID: participant.jobID, + registrationID: participant.registrationID + ) + guard let registration = workerRegistrationsByJobID[participant.jobID], + registration.id == participant.registrationID, + registration.replacement === replacement + else { + return + } + registration.replacement = nil + } + + fileprivate func suppressParticipant( jobID: String, - sourceGeneration: ReviewRuntimeGeneration + in replacement: ReviewRuntimeRecoveryReplacement ) { - guard let replacement = replacementsBySourceGeneration[sourceGeneration] else { + guard let participant = replacement.participants.first(where: { + $0.jobID == jobID + && ($0.phase == .eligible || $0.phase == .recovering) + }) else { return } - replacement.finishParticipant(jobID: jobID) - if replacement.hasRemainingParticipants == false { - replacementsBySourceGeneration.removeValue(forKey: sourceGeneration) + suppressParticipant(participant, in: replacement) + } + + fileprivate func finishParticipant( + registrationID: UUID, + jobID: String, + replacement: ReviewRuntimeRecoveryReplacement + ) { + replacement.finishParticipant( + jobID: jobID, + registrationID: registrationID + ) + guard let registration = workerRegistrationsByJobID[jobID], + registration.id == registrationID, + registration.replacement === replacement + else { + return } + registration.replacement = nil } fileprivate func finishRegisteredParticipant( @@ -2467,10 +2668,21 @@ package final class ReviewRuntimeReplacementCoordinator { else { return } - finishParticipant( - jobID: jobID, - sourceGeneration: registration.runtimeGeneration - ) + if let replacement = registration.replacement { + finishParticipant( + registrationID: registrationID, + jobID: jobID, + replacement: replacement + ) + } + } + + package var activeReplacementReceiptCountForTesting: Int { + workerRegistrationsByJobID.values.reduce(into: 0) { count, registration in + if registration.replacement != nil { + count += 1 + } + } } } diff --git a/Sources/CodexReview/Store/CodexReviewStoreTesting.swift b/Sources/CodexReview/Store/CodexReviewStoreTesting.swift index cf42424b..ac995308 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreTesting.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreTesting.swift @@ -37,6 +37,40 @@ extension CodexReviewStore { runtimeForceCloseReceiptRecordedForTesting = operation } + package func setRuntimeReplacementEnrollmentPreparationForTesting( + _ operation: (@MainActor @Sendable () async -> Void)? + ) { + runtimeReplacementEnrollmentPreparationForTesting = operation + } + + package func setReviewCancellationBarrierPreparationForTesting( + _ operation: (@MainActor @Sendable () async -> Void)? + ) { + reviewCancellationBarrierPreparationForTesting = operation + } + + package func waitForRuntimeReplacementRegistrationForTesting( + jobID: String, + attemptID: String + ) async { + await runtimeWorkerRegistry.waitForRegistrationForTesting( + jobID: jobID, + attemptID: attemptID + ) + } + + package func waitForRuntimeReplacementRegistrationForTesting( + jobID: String + ) async -> String { + await runtimeWorkerRegistry.waitForRegistrationForTesting( + jobID: jobID + ) + } + + package var activeRuntimeReplacementReceiptCountForTesting: Int { + runtimeWorkerRegistry.activeReplacementReceiptCountForTesting + } + package func loadForTesting( serverState: CodexReviewServerState, authPhase: CodexReviewAuthModel.Phase = .signedOut, diff --git a/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift b/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift index d3b4ea32..b61d5caa 100644 --- a/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift +++ b/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift @@ -143,7 +143,11 @@ package class PreviewCodexReviewStoreBackend: CodexReviewStoreBackend { throw CodexReviewAPI.Error.io(Self.previewUnavailableMessage) } - package func interruptReview(_: CodexReviewBackendModel.Review.Run, reason _: CodexReviewBackendModel.CancellationReason) async throws {} + package func interruptReview( + _: CodexReviewBackendModel.Review.Run, + admission _: ReviewStartAdmission, + reason _: CodexReviewBackendModel.CancellationReason + ) async throws {} package func forceCloseReviewConnection() async throws {} @@ -161,6 +165,18 @@ package class PreviewCodexReviewStoreBackend: CodexReviewStoreBackend { throw CodexReviewAPI.Error.io(Self.previewUnavailableMessage) } + package func commitResumedReviewRecovery( + _: ReviewRecoveryHandoff, + recoveredRun _: CodexReviewBackendModel.Review.Run + ) throws { + throw CodexReviewAPI.Error.io(Self.previewUnavailableMessage) + } + + package func discardResumedReviewRecovery( + _: ReviewRecoveryHandoff, + recoveredRun _: CodexReviewBackendModel.Review.Run + ) async throws {} + package func cleanupReview(_: CodexReviewBackendModel.Review.Run) async throws {} fileprivate static let previewUnavailableMessage = "Embedded server is unavailable in preview mode." diff --git a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift index 020de9f0..e5f8a3f2 100644 --- a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift +++ b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift @@ -804,6 +804,9 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { completedReviewEventSessionMetricsByThreadID[threadID] = metrics } } + guard case .open = lifecycleState else { + return + } var failureMessages: [String] = [] do { let _: EmptyResponse = try await client.send(AppServerAPI.Thread.BackgroundTerminals.Clean.Request( diff --git a/Sources/CodexReviewHost/CodexReviewHost.swift b/Sources/CodexReviewHost/CodexReviewHost.swift index 1aa48534..5119f73e 100644 --- a/Sources/CodexReviewHost/CodexReviewHost.swift +++ b/Sources/CodexReviewHost/CodexReviewHost.swift @@ -315,6 +315,7 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { func interruptReview( _ run: CodexReviewBackendModel.Review.Run, + admission _: ReviewStartAdmission, reason: CodexReviewBackendModel.CancellationReason ) async throws { try await backend.interruptReview(run, reason: reason) @@ -338,6 +339,18 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { try await backend.resumeReviewRecovery(handoff, request: request, admission: admission) } + func commitResumedReviewRecovery( + _: ReviewRecoveryHandoff, + recoveredRun _: CodexReviewBackendModel.Review.Run + ) throws {} + + func discardResumedReviewRecovery( + _: ReviewRecoveryHandoff, + recoveredRun: CodexReviewBackendModel.Review.Run + ) async throws { + try await backend.cleanupReview(recoveredRun) + } + func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { try await backend.cleanupReview(run) } diff --git a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift index acd23de9..1e99989d 100644 --- a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift +++ b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift @@ -55,6 +55,16 @@ package typealias CodexReviewMCPLifecycleCallObserver = @MainActor @Sendable ( Int ) -> Void +package enum CodexReviewLiveRecoveryRoutingEvent: Equatable, Sendable { + case staged(sourceAttemptID: String, recoveredAttemptID: String) + case committed(sourceAttemptID: String, recoveredAttemptID: String) + case discarded(sourceAttemptID: String, recoveredAttemptID: String) +} + +package typealias CodexReviewLiveRecoveryRoutingObserver = @MainActor @Sendable ( + CodexReviewLiveRecoveryRoutingEvent +) -> Void + package protocol CodexReviewMCPHTTPServing: AnyObject, Sendable { var url: URL { get async } @@ -97,6 +107,8 @@ public extension CodexReviewStore { networkMonitor: any CodexReviewNetworkMonitoring = SystemCodexReviewNetworkMonitor(), networkRecoveryPolicy: CodexReviewNetworkRecoveryPolicy = .default, reviewRuntimeClosePolicy: ReviewRuntimeClosePolicy = .production, + idGenerator: CodexReviewIDGenerator = .init(), + reviewRecoveryRoutingObserver: CodexReviewLiveRecoveryRoutingObserver? = nil, transport: any JSONRPC.Transport ) -> CodexReviewStore { makeLiveStoreForTesting( @@ -111,6 +123,8 @@ public extension CodexReviewStore { networkMonitor: networkMonitor, networkRecoveryPolicy: networkRecoveryPolicy, reviewRuntimeClosePolicy: reviewRuntimeClosePolicy, + idGenerator: idGenerator, + reviewRecoveryRoutingObserver: reviewRecoveryRoutingObserver, transportFactory: { _ in transport } ) } @@ -131,6 +145,8 @@ public extension CodexReviewStore { networkMonitor: any CodexReviewNetworkMonitoring = SystemCodexReviewNetworkMonitor(), networkRecoveryPolicy: CodexReviewNetworkRecoveryPolicy = .default, reviewRuntimeClosePolicy: ReviewRuntimeClosePolicy = .production, + idGenerator: CodexReviewIDGenerator = .init(), + reviewRecoveryRoutingObserver: CodexReviewLiveRecoveryRoutingObserver? = nil, transportFactory: @escaping @MainActor @Sendable (URL) async throws -> any JSONRPC.Transport ) -> CodexReviewStore { CodexReviewStore( @@ -144,6 +160,7 @@ public extension CodexReviewStore { mcpPortOwnerResolver: mcpPortOwnerResolver, mcpHTTPServerBindChecker: mcpHTTPServerBindChecker, mcpLifecycleCallObserver: mcpLifecycleCallObserver, + reviewRecoveryRoutingObserver: reviewRecoveryRoutingObserver, appServerRuntimeFactory: { codexHomeURL in let client = AppServerClient(transport: try await transportFactory(codexHomeURL)) return .init( @@ -152,6 +169,7 @@ public extension CodexReviewStore { ) } ), + idGenerator: idGenerator, networkMonitor: networkMonitor, networkRecoveryPolicy: networkRecoveryPolicy, reviewRuntimeClosePolicy: reviewRuntimeClosePolicy @@ -169,6 +187,23 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private struct RecoveryRoute { let handoff: ReviewRecoveryHandoff let source: AttemptRoute + var destination: RecoveryDestination? + } + + private enum RecoveryDestination { + case resuming( + runtime: LiveRuntimeLifecycleHandle, + admissionID: ObjectIdentifier + ) + case staged( + runtime: LiveRuntimeLifecycleHandle, + admissionID: ObjectIdentifier, + run: CodexReviewBackendModel.Review.Run + ) + case discarding( + runtime: LiveRuntimeLifecycleHandle, + run: CodexReviewBackendModel.Review.Run + ) } typealias MCPHTTPServerFactory = @MainActor @Sendable ( @@ -182,6 +217,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private var activeRuntimeHandle: LiveRuntimeLifecycleHandle? private var attemptRoutesByAttemptID: [String: AttemptRoute] = [:] private var recoveryRoutesByAttemptID: [String: RecoveryRoute] = [:] + private var pendingRuntimeByAdmissionID: [ObjectIdentifier: LiveRuntimeLifecycleHandle] = [:] private var acceptsRuntimeRequests = false private var loginChallenge: CodexReviewBackendModel.Login.Challenge? private var loginBackend: AppServerCodexReviewBackend? @@ -206,6 +242,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { private let externalURLOpener: ExternalURLOpener private let mcpLifecycleOwner: LiveMCPServerLifecycleOwner private let appServerRuntimeFactory: AppServerRuntimeFactory + private let reviewRecoveryRoutingObserver: CodexReviewLiveRecoveryRoutingObserver? private weak var attachedStore: CodexReviewStore? private var appServerBackend: AppServerCodexReviewBackend? { @@ -227,6 +264,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { mcpPortOwnerResolver: CodexReviewMCPPortOwnerResolver? = nil, mcpHTTPServerBindChecker: CodexReviewMCPHTTPServerBindChecker? = nil, mcpLifecycleCallObserver: CodexReviewMCPLifecycleCallObserver? = nil, + reviewRecoveryRoutingObserver: CodexReviewLiveRecoveryRoutingObserver? = nil, appServerRuntimeFactory: AppServerRuntimeFactory? = nil ) { let runtimePreferences = runtimePreferences.normalized @@ -242,6 +280,7 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { self.nativeAuthenticationConfiguration = nativeAuthenticationConfiguration self.webAuthenticationSessionFactory = webAuthenticationSessionFactory self.externalURLOpener = externalURLOpener + self.reviewRecoveryRoutingObserver = reviewRecoveryRoutingObserver let resolvedPortOwnerResolver = mcpPortOwnerResolver ?? Self.defaultMCPPortOwnerResolver let resolvedBindChecker = mcpHTTPServerBindChecker ?? Self.defaultMCPHTTPServerBindChecker self.mcpLifecycleOwner = LiveMCPServerLifecycleOwner( @@ -1032,6 +1071,13 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { guard acceptsRuntimeRequests, let runtime = activeRuntimeHandle else { throw CodexReviewAPI.Error.io("Review runtime is not running.") } + let admissionID = ObjectIdentifier(admission) + pendingRuntimeByAdmissionID[admissionID] = runtime + defer { + if pendingRuntimeByAdmissionID[admissionID] === runtime { + pendingRuntimeByAdmissionID.removeValue(forKey: admissionID) + } + } let attempt = try await runtime.backend.startReview(request, admission: admission) attemptRoutesByAttemptID[attempt.run.attemptID] = .init( generation: runtime.generation, @@ -1040,11 +1086,32 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { return attempt } - func interruptReview(_ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason) async throws { - guard let route = attemptRoutesByAttemptID[run.attemptID] else { - throw CodexReviewAPI.Error.io("Review attempt route is unavailable.") + func interruptReview( + _ run: CodexReviewBackendModel.Review.Run, + admission: ReviewStartAdmission, + reason: CodexReviewBackendModel.CancellationReason + ) async throws { + if let route = attemptRoutesByAttemptID[run.attemptID] { + try await route.runtime.backend.interruptReview(run, reason: reason) + return + } + let admissionID = ObjectIdentifier(admission) + if let pendingRuntime = pendingRuntimeByAdmissionID[admissionID] { + try await pendingRuntime.backend.interruptReview(run, reason: reason) + return + } + for recovery in recoveryRoutesByAttemptID.values { + guard case .staged(let runtime, let stagedAdmissionID, let stagedRun) = + recovery.destination, + stagedAdmissionID == admissionID, + stagedRun.attemptID == run.attemptID + else { + continue + } + try await runtime.backend.interruptReview(run, reason: reason) + return } - try await route.runtime.backend.interruptReview(run, reason: reason) + throw CodexReviewAPI.Error.io("Review attempt route is unavailable.") } func forceCloseReviewConnection() async throws { @@ -1064,7 +1131,8 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { let handoff = try await source.runtime.backend.prepareReviewRecovery(candidate) recoveryRoutesByAttemptID[attemptID] = .init( handoff: handoff, - source: source + source: source, + destination: nil ) return handoff } @@ -1089,18 +1157,110 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { "Same-account recovery requires a replacement runtime generation." ) } - let attempt = try await destination.backend.resumeReviewRecovery( - handoff, - request: request, - admission: admission + let admissionID = ObjectIdentifier(admission) + recoveryRoutesByAttemptID[sourceAttemptID]?.destination = .resuming( + runtime: destination, + admissionID: admissionID ) + pendingRuntimeByAdmissionID[admissionID] = destination + do { + let attempt = try await destination.backend.resumeReviewRecovery( + handoff, + request: request, + admission: admission + ) + guard let current = recoveryRoutesByAttemptID[sourceAttemptID], + current.handoff == handoff + else { + try await destination.backend.cleanupReview(attempt.run) + throw CodexReviewAPI.Error.io("Review recovery route changed before staging.") + } + recoveryRoutesByAttemptID[sourceAttemptID]?.destination = .staged( + runtime: destination, + admissionID: admissionID, + run: attempt.run + ) + reviewRecoveryRoutingObserver?(.staged( + sourceAttemptID: sourceAttemptID, + recoveredAttemptID: attempt.run.attemptID + )) + pendingRuntimeByAdmissionID.removeValue(forKey: admissionID) + return attempt + } catch { + if pendingRuntimeByAdmissionID[admissionID] === destination { + pendingRuntimeByAdmissionID.removeValue(forKey: admissionID) + } + if let current = recoveryRoutesByAttemptID[sourceAttemptID], + current.handoff == handoff, + case .resuming(let runtime, let currentAdmissionID) = current.destination, + runtime === destination, + currentAdmissionID == admissionID { + recoveryRoutesByAttemptID[sourceAttemptID]?.destination = nil + } + throw error + } + } + + func commitResumedReviewRecovery( + _ handoff: ReviewRecoveryHandoff, + recoveredRun: CodexReviewBackendModel.Review.Run + ) throws { + let sourceAttemptID = handoff.candidate.resolved.run.attemptID + guard let recovery = recoveryRoutesByAttemptID[sourceAttemptID], + recovery.handoff == handoff, + case .staged(let destination, _, let stagedRun) = recovery.destination, + stagedRun.attemptID == recoveredRun.attemptID, + activeRuntimeHandle === destination + else { + throw ReviewAttemptContractFailure( + message: "Review recovery destination route is not staged for commit." + ) + } attemptRoutesByAttemptID.removeValue(forKey: sourceAttemptID) recoveryRoutesByAttemptID.removeValue(forKey: sourceAttemptID) - attemptRoutesByAttemptID[attempt.run.attemptID] = .init( + attemptRoutesByAttemptID[recoveredRun.attemptID] = .init( generation: destination.generation, runtime: destination ) - return attempt + reviewRecoveryRoutingObserver?(.committed( + sourceAttemptID: sourceAttemptID, + recoveredAttemptID: recoveredRun.attemptID + )) + } + + func discardResumedReviewRecovery( + _ handoff: ReviewRecoveryHandoff, + recoveredRun: CodexReviewBackendModel.Review.Run + ) async throws { + let sourceAttemptID = handoff.candidate.resolved.run.attemptID + guard let recovery = recoveryRoutesByAttemptID[sourceAttemptID], + recovery.handoff == handoff, + case .staged(let destination, let admissionID, let stagedRun) = recovery.destination, + stagedRun.attemptID == recoveredRun.attemptID + else { + throw ReviewRuntimeCloseFailure.cleanup( + "Review recovery destination route is not staged for discard." + ) + } + recoveryRoutesByAttemptID[sourceAttemptID]?.destination = .discarding( + runtime: destination, + run: stagedRun + ) + reviewRecoveryRoutingObserver?(.discarded( + sourceAttemptID: sourceAttemptID, + recoveredAttemptID: recoveredRun.attemptID + )) + pendingRuntimeByAdmissionID.removeValue(forKey: admissionID) + defer { + if let current = recoveryRoutesByAttemptID[sourceAttemptID], + current.handoff == handoff, + case .discarding(let runtime, let currentRun) = current.destination, + runtime === destination, + currentRun.attemptID == recoveredRun.attemptID { + recoveryRoutesByAttemptID[sourceAttemptID]?.destination = nil + } + } + try await destination.backend.cleanupReview(recoveredRun) } func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { diff --git a/Sources/CodexReviewTesting/TestSupport.swift b/Sources/CodexReviewTesting/TestSupport.swift index 3e831ff7..8296c527 100644 --- a/Sources/CodexReviewTesting/TestSupport.swift +++ b/Sources/CodexReviewTesting/TestSupport.swift @@ -209,6 +209,7 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { private var startReviewGate: AsyncGate? private var startReviewWaiters: [UUID: CheckedContinuation] = [:] private var resumeReviewRecoveryGate: AsyncGate? + private var resumeReviewRecoveryIgnoresCancellation = false private var resumeReviewRecoveryWaiters: [UUID: CheckedContinuation] = [:] private var eventMailboxes: [EventMailboxKey: BackendReviewEventMailbox] = [:] @@ -277,6 +278,12 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { package func holdResumeReviewRecovery(with gate: AsyncGate) { resumeReviewRecoveryGate = gate + resumeReviewRecoveryIgnoresCancellation = false + } + + package func holdResumeReviewRecoveryIgnoringCancellation(with gate: AsyncGate) { + resumeReviewRecoveryGate = gate + resumeReviewRecoveryIgnoresCancellation = true } package func holdPrepareReviewRecovery(with gate: AsyncGate) { @@ -660,8 +667,12 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { waiter.resume() } if let resumeReviewRecoveryGate { - await resumeReviewRecoveryGate.wait() - try Task.checkCancellation() + if resumeReviewRecoveryIgnoresCancellation { + await resumeReviewRecoveryGate.waitIgnoringCancellation() + } else { + await resumeReviewRecoveryGate.wait() + try Task.checkCancellation() + } } if let recoveryFailureMessage { throw FakeCodexReviewBackendError(message: recoveryFailureMessage) @@ -1392,6 +1403,7 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { package func interruptReview( _ run: CodexReviewBackendModel.Review.Run, + admission _: ReviewStartAdmission, reason: CodexReviewBackendModel.CancellationReason ) async throws { try await reviewBackend.interruptReview(run, reason: reason) @@ -1422,6 +1434,18 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { ) } + package func commitResumedReviewRecovery( + _: ReviewRecoveryHandoff, + recoveredRun _: CodexReviewBackendModel.Review.Run + ) throws {} + + package func discardResumedReviewRecovery( + _: ReviewRecoveryHandoff, + recoveredRun: CodexReviewBackendModel.Review.Run + ) async throws { + try await reviewBackend.cleanupReview(recoveredRun) + } + package func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { try await reviewBackend.cleanupReview(run) } diff --git a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift index 0e805024..da844471 100644 --- a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift +++ b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift @@ -451,6 +451,333 @@ struct CodexReviewHostTests { await store.stop() } + @Test func liveGraceForceCloseResumesSiblingOnlyOnReplacementBackend() async throws { + let homeURL = try temporaryHome() + let firstTransport = FakeJSONRPCTransport() + try await firstTransport.enqueue(AppServerAPI.Initialize.Response(), for: "initialize") + try await firstTransport.enqueue(AppServerAPI.Account.Read.Response(), for: "account/read") + try await firstTransport.enqueue( + AppServerAPI.Config.Read.Response(config: .init(model: "gpt-5")), + for: "config/read" + ) + try await firstTransport.enqueue(AppServerAPI.Model.List.Response(data: []), for: "model/list") + try await firstTransport.enqueue( + AppServerAPI.Thread.Start.Response(threadID: "thread-target", model: "gpt-5"), + for: "thread/start" + ) + try await firstTransport.enqueue( + AppServerAPI.Review.Start.Response( + turnID: "turn-target", + reviewThreadID: "review-target" + ), + for: "review/start" + ) + try await firstTransport.enqueue( + AppServerAPI.Thread.Start.Response(threadID: "thread-sibling", model: "gpt-5"), + for: "thread/start" + ) + try await firstTransport.enqueue( + AppServerAPI.Review.Start.Response( + turnID: "turn-sibling", + reviewThreadID: "review-sibling" + ), + for: "review/start" + ) + try await firstTransport.enqueue(EmptyResponse(), for: "turn/interrupt") + try await firstTransport.enqueue(EmptyResponse(), for: "turn/interrupt") + + let secondTransport = FakeJSONRPCTransport() + try await secondTransport.enqueue(AppServerAPI.Initialize.Response(), for: "initialize") + try await secondTransport.enqueue(AppServerAPI.Account.Read.Response(), for: "account/read") + try await secondTransport.enqueue( + AppServerAPI.Config.Read.Response(config: .init(model: "gpt-5")), + for: "config/read" + ) + try await secondTransport.enqueue(AppServerAPI.Model.List.Response(data: []), for: "model/list") + try await secondTransport.enqueue(EmptyResponse(), for: "thread/rollback") + try await secondTransport.enqueue( + AppServerAPI.Review.Start.Response( + turnID: "turn-sibling-recovered", + reviewThreadID: "review-sibling" + ), + for: "review/start" + ) + + let mcpServer = ControlledMCPHTTPServer( + endpoint: try #require(URL(string: "http://127.0.0.1:19432/mcp")) + ) + var transports = [firstTransport, secondTransport] + var mcpFactoryCallCount = 0 + let routingProbe = HostRecoveryRoutingProbe() + let jobIDs = HostSequentialIDs(["job-target", "job-sibling"]) + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + mcpHTTPServerFactory: { _, _ in + mcpFactoryCallCount += 1 + return mcpServer + }, + mcpHTTPServerBindChecker: { _ in }, + reviewRuntimeClosePolicy: .init( + terminalGrace: .seconds(10), + sleep: { _ in } + ), + idGenerator: .init(next: { jobIDs.next() }), + reviewRecoveryRoutingObserver: { event in + routingProbe.record(event) + }, + transportFactory: { _ in transports.removeFirst() } + ) + await store.start() + let initialURL = store.serverURL + let targetTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-target", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + _ = await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-target" + ) + let siblingTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-sibling", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + _ = await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-sibling" + ) + + let cancellation = Task { @MainActor in + try await store.cancelReview( + jobID: "job-target", + cancellation: .mcpClient(message: "Stop target") + ) + } + await secondTransport.waitForRequest(method: "review/start") + try await secondTransport.emitServerNotification( + method: "item/completed", + params: HostCompletedReviewItemNotification( + threadID: "review-sibling", + turnID: "turn-sibling-recovered", + result: "replacement result" + ) + ) + try await secondTransport.emitServerNotification( + method: "turn/completed", + params: HostCompletedReviewTurnNotification( + threadID: "review-sibling", + turnID: "turn-sibling-recovered", + result: "replacement result" + ) + ) + + let target = try await targetTask.value + let sibling = try await siblingTask.value + #expect(try await cancellation.value.cancelled) + #expect(target.core.lifecycle.terminal == .interrupted( + .requested(.mcpClient(message: "Stop target")) + )) + #expect(sibling.core.lifecycle.status == .succeeded) + #expect(sibling.core.run.turnID == "turn-sibling-recovered") + #expect(store.serverURL == initialURL) + #expect(mcpFactoryCallCount == 1) + #expect(mcpServer.startCallCount == 1) + #expect(transports.isEmpty) + let oldMethods = await firstTransport.recordedRequests().map(\.method) + let newMethods = await secondTransport.recordedRequests().map(\.method) + #expect(oldMethods.filter { $0 == "review/start" }.count == 2) + #expect(oldMethods.contains("thread/rollback") == false) + #expect(newMethods.filter { $0 == "review/start" }.count == 1) + #expect(newMethods.contains("thread/rollback")) + #expect(routingProbe.events == [ + .staged( + sourceAttemptID: routingProbe.sourceAttemptID, + recoveredAttemptID: routingProbe.recoveredAttemptID + ), + .committed( + sourceAttemptID: routingProbe.sourceAttemptID, + recoveredAttemptID: routingProbe.recoveredAttemptID + ), + ]) + + await store.stop() + } + + @Test func liveHeldRecoveryResumeIsDiscardedWhenStopWins() async throws { + try await exerciseHeldRecoveryResumeDiscard(applicationClose: false) + } + + @Test func liveHeldRecoveryResumeIsDiscardedWhenCloseWins() async throws { + try await exerciseHeldRecoveryResumeDiscard(applicationClose: true) + } + + private func exerciseHeldRecoveryResumeDiscard( + applicationClose: Bool + ) async throws { + let homeURL = try temporaryHome() + let firstTransport = FakeJSONRPCTransport() + try await firstTransport.enqueue(AppServerAPI.Initialize.Response(), for: "initialize") + try await firstTransport.enqueue(AppServerAPI.Account.Read.Response(), for: "account/read") + try await firstTransport.enqueue( + AppServerAPI.Config.Read.Response(config: .init(model: "gpt-5")), + for: "config/read" + ) + try await firstTransport.enqueue(AppServerAPI.Model.List.Response(data: []), for: "model/list") + try await firstTransport.enqueue( + AppServerAPI.Thread.Start.Response(threadID: "thread-target", model: "gpt-5"), + for: "thread/start" + ) + try await firstTransport.enqueue( + AppServerAPI.Review.Start.Response( + turnID: "turn-target", + reviewThreadID: "review-target" + ), + for: "review/start" + ) + try await firstTransport.enqueue( + AppServerAPI.Thread.Start.Response(threadID: "thread-sibling", model: "gpt-5"), + for: "thread/start" + ) + try await firstTransport.enqueue( + AppServerAPI.Review.Start.Response( + turnID: "turn-sibling", + reviewThreadID: "review-sibling" + ), + for: "review/start" + ) + try await firstTransport.enqueue(EmptyResponse(), for: "turn/interrupt") + try await firstTransport.enqueue(EmptyResponse(), for: "turn/interrupt") + + let secondTransport = FakeJSONRPCTransport() + let recoveredStartGate = AsyncGate() + await secondTransport.holdNextIgnoringCancellation( + method: "review/start", + gate: recoveredStartGate + ) + try await secondTransport.enqueue(AppServerAPI.Initialize.Response(), for: "initialize") + try await secondTransport.enqueue(AppServerAPI.Account.Read.Response(), for: "account/read") + try await secondTransport.enqueue( + AppServerAPI.Config.Read.Response(config: .init(model: "gpt-5")), + for: "config/read" + ) + try await secondTransport.enqueue(AppServerAPI.Model.List.Response(data: []), for: "model/list") + try await secondTransport.enqueue(EmptyResponse(), for: "thread/rollback") + try await secondTransport.enqueue( + AppServerAPI.Review.Start.Response( + turnID: "turn-sibling-recovered", + reviewThreadID: "review-sibling" + ), + for: "review/start" + ) + try await secondTransport.enqueue(EmptyResponse(), for: "turn/interrupt") + + let mcpServer = ControlledMCPHTTPServer( + endpoint: try #require(URL(string: "http://127.0.0.1:19436/mcp")) + ) + let routingProbe = HostRecoveryRoutingProbe() + var transports = [firstTransport, secondTransport] + let jobIDs = HostSequentialIDs(["job-target", "job-sibling"]) + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + mcpHTTPServerFactory: { _, _ in mcpServer }, + mcpHTTPServerBindChecker: { _ in }, + reviewRuntimeClosePolicy: .init( + terminalGrace: .seconds(10), + sleep: { _ in } + ), + idGenerator: .init(next: { jobIDs.next() }), + reviewRecoveryRoutingObserver: { event in + routingProbe.record(event) + }, + transportFactory: { _ in transports.removeFirst() } + ) + await store.start() + let targetTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-target", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + _ = await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-target" + ) + let siblingTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-sibling", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + _ = await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-sibling" + ) + let targetCancellation = Task { @MainActor in + try await store.cancelReview( + jobID: "job-target", + cancellation: .mcpClient(message: "Stop target") + ) + } + await secondTransport.waitForRequest(method: "review/start") + + let cancellationBarrierEntered = AsyncGate() + store.setReviewCancellationBarrierPreparationForTesting { + await cancellationBarrierEntered.open() + } + let shutdownTask = Task { @MainActor in + do { + if applicationClose { + try await store.close() + } else { + await store.stop() + } + return Result.success(()) + } catch { + return Result.failure(error) + } + } + await cancellationBarrierEntered.wait() + await recoveredStartGate.open() + await secondTransport.waitForRequest(method: "turn/interrupt") + try await shutdownTask.value.get() + let target = try await targetTask.value + let sibling = try await siblingTask.value + _ = try await targetCancellation.value + + let expectedSiblingCancellation: ReviewCancellation = applicationClose + ? .system(message: "Review Store closed.") + : .system(message: "Review runtime stopped.") + #expect(target.core.lifecycle.terminal == .interrupted( + .requested(.mcpClient(message: "Stop target")) + )) + #expect(sibling.core.lifecycle.terminal == .interrupted( + .requested(expectedSiblingCancellation) + )) + #expect(sibling.core.run.turnID == "turn-sibling") + #expect(store.reviewWorkerTasks.isEmpty) + #expect(store.activeRuntimeReplacementReceiptCountForTesting == 0) + let oldMethods = await firstTransport.recordedRequests().map(\.method) + let newMethods = await secondTransport.recordedRequests().map(\.method) + #expect(oldMethods.contains("thread/rollback") == false) + #expect(newMethods.contains("thread/rollback")) + #expect(newMethods.filter { $0 == "review/start" }.count == 1) + #expect(newMethods.contains("turn/interrupt")) + #expect(routingProbe.events == [ + .staged( + sourceAttemptID: routingProbe.sourceAttemptID, + recoveredAttemptID: routingProbe.recoveredAttemptID + ), + .discarded( + sourceAttemptID: routingProbe.sourceAttemptID, + recoveredAttemptID: routingProbe.recoveredAttemptID + ), + ]) + #expect(await secondTransport.closeCallCountForTesting() == 1) + #expect(store.serverState == .stopped) + #expect(store.serverURL == nil) + } + @Test func liveMCPOwnerStopDuringPreparationJoinsAndAllowsLaterStart() async throws { let homeURL = try temporaryHome() let preparationStarted = AsyncGate() @@ -2901,6 +3228,88 @@ private actor CompletionFlag { } } +private final class HostSequentialIDs: @unchecked Sendable { + private let lock = NSLock() + private var values: [String] + + init(_ values: [String]) { + self.values = values + } + + func next() -> String { + lock.withLock { values.removeFirst() } + } +} + +@MainActor +private final class HostRecoveryRoutingProbe { + private(set) var events: [CodexReviewLiveRecoveryRoutingEvent] = [] + + var sourceAttemptID: String { + guard case .staged(let sourceAttemptID, _) = events.first else { + return "missing-source" + } + return sourceAttemptID + } + + var recoveredAttemptID: String { + guard case .staged(_, let recoveredAttemptID) = events.first else { + return "missing-recovered" + } + return recoveredAttemptID + } + + func record(_ event: CodexReviewLiveRecoveryRoutingEvent) { + events.append(event) + } +} + +private struct HostCompletedReviewTurnNotification: Encodable, Sendable { + private struct Turn: Encodable, Sendable { + struct Item: Encodable, Sendable { + var type = "exitedReviewMode" + var id = "final-review" + var review: String + } + + var id: String + var items: [Item] + var itemsView = "full" + var status = "completed" + var error: String? = nil + } + + private var threadId: String + private var turn: Turn + + init(threadID: String, turnID: String, result: String) { + self.threadId = threadID + self.turn = Turn( + id: turnID, + items: [.init(review: result)] + ) + } +} + +private struct HostCompletedReviewItemNotification: Encodable, Sendable { + private struct Item: Encodable, Sendable { + var type = "exitedReviewMode" + var id = "final-review" + var review: String + } + + private var threadId: String + private var turnId: String + private var item: Item + private var completedAtMs: Int64 = 0 + + init(threadID: String, turnID: String, result: String) { + self.threadId = threadID + self.turnId = turnID + self.item = Item(review: result) + } +} + private struct HostTurnNotification: Encodable, Sendable { struct Turn: Encodable, Sendable { struct TurnError: Encodable, Sendable { diff --git a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift index 880c9cfa..bb1ad4d7 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift @@ -275,6 +275,867 @@ struct CodexReviewStoreLifecycleTests { await store.stop() } + @Test func manualRestartRecoversActiveReviewThroughTheSharedCoordinator() async throws { + let initialRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-active", + threadID: "thread-active", + turnID: "turn-active", + reviewThreadID: "review-active" + ) + let recoveredRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-active-recovered", + threadID: "thread-active", + turnID: "turn-active-recovered", + reviewThreadID: "review-active" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: initialRun) + await reviewBackend.enqueueRecoveredRun(recoveredRun) + let endpoint = try #require(URL(string: "http://127.0.0.1:19435/mcp")) + let mcpOwner = TestingMCPServerLifecycleOwner(serverURL: endpoint) + let backend = TestingCodexReviewStoreBackend( + reviewBackend: reviewBackend, + mcpServerLifecycle: mcpOwner + ) + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { "job-active" }) + ) + await store.start() + let sourceRuntime = try #require(backend.lastPreparedRuntimeHandle) + let reviewTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-active", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-active", + attemptID: initialRun.attemptID + ) + + let restartTask = Task { @MainActor in + await store.restart() + } + await reviewBackend.waitForResumeReviewRecovery() + await reviewBackend.yield( + .completed(summary: "Done", result: "manual restart recovered"), + for: recoveredRun + ) + await restartTask.value + let review = try await reviewTask.value + + #expect(review.core.lifecycle.status == .succeeded) + #expect(review.core.run.turnID == recoveredRun.turnID) + #expect(sourceRuntime.closeCallCount == 1) + #expect(sourceRuntime.closePurposes == [.restartSameAccount]) + #expect(backend.startRequests == [false, true]) + #expect(store.serverURL == endpoint) + #expect(mcpOwner.prepareCallCount == 1) + #expect(mcpOwner.activateCallCount == 1) + #expect(mcpOwner.stopCallCount == 0) + let commands = await reviewBackend.recordedCommands() + #expect(commands.filter { if case .prepareReviewRecovery = $0 { true } else { false } }.count == 1) + #expect(commands.filter { if case .resumeReviewRecovery = $0 { true } else { false } }.count == 1) + + await store.stop() + } + + @Test func graceForceCloseWithNoSiblingStillPublishesOneReplacementRuntime() async throws { + let targetRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-target", + threadID: "thread-target", + turnID: "turn-target", + reviewThreadID: "review-target" + ) + let nextRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-next", + threadID: "thread-next", + turnID: "turn-next", + reviewThreadID: "review-next" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: targetRun) + await reviewBackend.enqueueRun(nextRun) + let endpoint = try #require(URL(string: "http://127.0.0.1:19432/mcp")) + let mcpOwner = TestingMCPServerLifecycleOwner(serverURL: endpoint) + let backend = TestingCodexReviewStoreBackend( + reviewBackend: reviewBackend, + mcpServerLifecycle: mcpOwner + ) + let jobIDs = SequentialJobIDs(["job-target", "job-next"]) + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { jobIDs.next() }), + reviewRuntimeClosePolicy: .init( + terminalGrace: .seconds(10), + sleep: { _ in } + ) + ) + await store.start() + let sourceRuntime = try #require(backend.lastPreparedRuntimeHandle) + let targetTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-target", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-target", + attemptID: targetRun.attemptID + ) + + let replacementGate = AsyncGate() + backend.holdRuntimePreparation(with: replacementGate) + let cancellation = Task { @MainActor in + try await store.cancelReview( + jobID: "job-target", + cancellation: .mcpClient(message: "Stop target") + ) + } + await backend.waitForRuntimePreparation() + let replacementRuntime = try #require(backend.lastPreparedRuntimeHandle) + + #expect(sourceRuntime.closeCallCount == 1) + #expect(sourceRuntime.closePurposes == [.recoveryReplacement]) + #expect(replacementRuntime !== sourceRuntime) + #expect(replacementRuntime.activateCallCount == 0) + #expect(backend.startRequests == [false, false]) + #expect(store.serverURL == endpoint) + #expect(mcpOwner.prepareCallCount == 1) + #expect(mcpOwner.activateCallCount == 1) + #expect(mcpOwner.stopCallCount == 0) + + await replacementGate.open() + await store.start() + let cancelled = try await cancellation.value + _ = try await targetTask.value + + #expect(cancelled.cancelled) + #expect(replacementRuntime.activateCallCount == 1) + #expect(store.serverState == .running) + #expect(store.serverURL == endpoint) + #expect(store.activeRuntimeReplacementReceiptCountForTesting == 0) + + let nextTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-next", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-next", + attemptID: nextRun.attemptID + ) + await reviewBackend.yield( + .completed(summary: "Done", result: "replacement runtime usable"), + for: nextRun + ) + let next = try await nextTask.value + #expect(next.core.lifecycle.status == .succeeded) + #expect(next.core.run.turnID == nextRun.turnID) + + await store.stop() + } + + @Test func siblingCanonicalCompletionBeforeEnrollmentSuppressesItsRecoverySuccessor() async throws { + let targetRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-target", + threadID: "thread-target", + turnID: "turn-target", + reviewThreadID: "review-target" + ) + let siblingRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-sibling", + threadID: "thread-sibling", + turnID: "turn-sibling", + reviewThreadID: "review-sibling" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: targetRun) + await reviewBackend.enqueueRun(siblingRun) + let backend = TestingCodexReviewStoreBackend(reviewBackend: reviewBackend) + let jobIDs = SequentialJobIDs(["job-target", "job-sibling"]) + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { jobIDs.next() }), + reviewRuntimeClosePolicy: .init( + terminalGrace: .seconds(10), + sleep: { _ in } + ) + ) + await store.start() + let enrollmentEntered = AsyncGate() + let enrollmentRelease = AsyncGate() + store.setRuntimeReplacementEnrollmentPreparationForTesting { + await enrollmentEntered.open() + await enrollmentRelease.waitIgnoringCancellation() + } + let targetTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-target", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-target", + attemptID: targetRun.attemptID + ) + let siblingTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-sibling", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-sibling", + attemptID: siblingRun.attemptID + ) + + let cancellation = Task { @MainActor in + try await store.cancelReview( + jobID: "job-target", + cancellation: .mcpClient(message: "Stop target") + ) + } + await enrollmentEntered.wait() + await reviewBackend.yield( + .completed(summary: "Done", result: "natural sibling result"), + for: siblingRun + ) + let sibling = try await siblingTask.value + + await enrollmentRelease.open() + store.setRuntimeReplacementEnrollmentPreparationForTesting(nil) + let cancelled = try await cancellation.value + _ = try await targetTask.value + await store.start() + + #expect(cancelled.cancelled) + #expect(sibling.core.lifecycle.status == .succeeded) + #expect(sibling.core.run.turnID == siblingRun.turnID) + #expect(backend.startRequests == [false, false]) + #expect(store.activeRuntimeReplacementReceiptCountForTesting == 0) + let commands = await reviewBackend.recordedCommands() + #expect(commands.filter { if case .prepareReviewRecovery = $0 { true } else { false } }.isEmpty) + #expect(commands.filter { if case .resumeReviewRecovery = $0 { true } else { false } }.isEmpty) + + await store.stop() + } + + @Test func replacementPreparationFailureFailsAllEligibleSiblingsAndRetainsMCP() async throws { + let targetRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-target", + threadID: "thread-target", + turnID: "turn-target", + reviewThreadID: "review-target" + ) + let firstSiblingRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-sibling-1", + threadID: "thread-sibling-1", + turnID: "turn-sibling-1", + reviewThreadID: "review-sibling-1" + ) + let secondSiblingRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-sibling-2", + threadID: "thread-sibling-2", + turnID: "turn-sibling-2", + reviewThreadID: "review-sibling-2" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: targetRun) + await reviewBackend.enqueueRun(firstSiblingRun) + await reviewBackend.enqueueRun(secondSiblingRun) + let endpoint = try #require(URL(string: "http://127.0.0.1:19433/mcp")) + let mcpOwner = TestingMCPServerLifecycleOwner(serverURL: endpoint) + let backend = TestingCodexReviewStoreBackend( + reviewBackend: reviewBackend, + mcpServerLifecycle: mcpOwner + ) + let jobIDs = SequentialJobIDs(["job-target", "job-sibling-1", "job-sibling-2"]) + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { jobIDs.next() }), + reviewRuntimeClosePolicy: .init( + terminalGrace: .seconds(10), + sleep: { _ in } + ) + ) + await store.start() + let targetTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-target", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-target", + attemptID: targetRun.attemptID + ) + let firstSiblingTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-sibling-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-sibling-1", + attemptID: firstSiblingRun.attemptID + ) + let secondSiblingTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-sibling-2", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-sibling-2", + attemptID: secondSiblingRun.attemptID + ) + await reviewBackend.failAuthRead(message: "Replacement authentication unavailable.") + + let cancellation = Task { @MainActor in + try await store.cancelReview( + jobID: "job-target", + cancellation: .mcpClient(message: "Stop target") + ) + } + let cancelled = try await cancellation.value + let target = try await targetTask.value + let firstSibling = try await firstSiblingTask.value + let secondSibling = try await secondSiblingTask.value + + #expect(cancelled.cancelled) + #expect(target.core.lifecycle.terminal == .interrupted( + .requested(.mcpClient(message: "Stop target")) + )) + guard case .failed(let firstMessage) = firstSibling.core.lifecycle.terminal, + case .failed(let secondMessage) = secondSibling.core.lifecycle.terminal + else { + Issue.record("Replacement failure must terminalize every eligible sibling.") + await store.stop() + return + } + #expect(firstMessage == secondMessage) + #expect(firstMessage?.contains("Replacement authentication unavailable.") == true) + guard case .failed = store.serverState else { + Issue.record("Replacement preparation failure must fail the Store runtime.") + await store.stop() + return + } + #expect(store.serverURL == endpoint) + #expect(backend.startRequests == [false, false]) + #expect(mcpOwner.prepareCallCount == 1) + #expect(mcpOwner.activateCallCount == 1) + #expect(mcpOwner.stopCallCount == 0) + await #expect(throws: CodexReviewAPI.Error.self) { + _ = try await store.startReview( + sessionID: "session-after-failure", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + + await store.stop() + } + + @Test func siblingCancellationDuringHeldHandoffSuppressesResume() async throws { + let targetRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-target", + threadID: "thread-target", + turnID: "turn-target", + reviewThreadID: "review-target" + ) + let siblingRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-sibling", + threadID: "thread-sibling", + turnID: "turn-sibling", + reviewThreadID: "review-sibling" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: targetRun) + await reviewBackend.enqueueRun(siblingRun) + let handoffGate = AsyncGate() + await reviewBackend.holdPrepareReviewRecovery(with: handoffGate) + let backend = TestingCodexReviewStoreBackend(reviewBackend: reviewBackend) + let jobIDs = SequentialJobIDs(["job-target", "job-sibling"]) + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { jobIDs.next() }), + reviewRuntimeClosePolicy: .init( + terminalGrace: .seconds(10), + sleep: { _ in } + ) + ) + await store.start() + let targetTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-target", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-target", + attemptID: targetRun.attemptID + ) + let siblingTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-sibling", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-sibling", + attemptID: siblingRun.attemptID + ) + + let targetCancellation = Task { @MainActor in + try await store.cancelReview( + jobID: "job-target", + cancellation: .mcpClient(message: "Stop target") + ) + } + await reviewBackend.waitForPrepareReviewRecovery() + let siblingCancellation = Task { @MainActor in + try await store.cancelReview( + jobID: "job-sibling", + cancellation: .system(message: "Stop sibling during handoff") + ) + } + let siblingCancelled = try await siblingCancellation.value + await handoffGate.open() + let targetCancelled = try await targetCancellation.value + let target = try await targetTask.value + let sibling = try await siblingTask.value + await store.start() + + #expect(targetCancelled.cancelled) + #expect(siblingCancelled.cancelled) + #expect(target.core.lifecycle.terminal == .interrupted( + .requested(.mcpClient(message: "Stop target")) + )) + #expect(sibling.core.lifecycle.terminal == .interrupted( + .requested(.system(message: "Stop sibling during handoff")) + )) + let commands = await reviewBackend.recordedCommands() + #expect(commands.filter { if case .prepareReviewRecovery = $0 { true } else { false } }.count == 1) + #expect(commands.filter { if case .resumeReviewRecovery = $0 { true } else { false } }.isEmpty) + #expect(backend.startRequests == [false, false]) + + await store.stop() + } + + @Test func applicationCloseSupersedesHeldRecoveryReplacementWithoutLatePublication() async throws { + let targetRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-target", + threadID: "thread-target", + turnID: "turn-target", + reviewThreadID: "review-target" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: targetRun) + let mcpOwner = TestingMCPServerLifecycleOwner() + let backend = TestingCodexReviewStoreBackend( + reviewBackend: reviewBackend, + mcpServerLifecycle: mcpOwner + ) + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { "job-target" }), + reviewRuntimeClosePolicy: .init( + terminalGrace: .seconds(10), + sleep: { _ in } + ) + ) + await store.start() + let sourceRuntime = try #require(backend.lastPreparedRuntimeHandle) + let targetTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-target", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-target", + attemptID: targetRun.attemptID + ) + + let replacementGate = AsyncGate() + backend.holdRuntimePreparation(with: replacementGate) + let cancellation = Task { @MainActor in + try await store.cancelReview( + jobID: "job-target", + cancellation: .mcpClient(message: "Stop target") + ) + } + await backend.waitForRuntimePreparation() + let staleReplacement = try #require(backend.lastPreparedRuntimeHandle) + let closeTask = Task { @MainActor in + try await store.close() + } + await backend.waitForRuntimePreparationCancellation() + + #expect(sourceRuntime.closeCallCount == 1) + #expect(sourceRuntime.closePurposes == [.recoveryReplacement]) + #expect(staleReplacement.activateCallCount == 0) + #expect(mcpOwner.closeAdmissionCallCount == 1) + + await replacementGate.open() + try await closeTask.value + _ = try await cancellation.value + _ = try await targetTask.value + + #expect(staleReplacement.activateCallCount == 0) + #expect(staleReplacement.closeCallCount == 1) + #expect(staleReplacement.waitUntilClosedCallCount == 1) + #expect(sourceRuntime.closeCallCount == 1) + #expect(sourceRuntime.waitUntilClosedCallCount == 1) + #expect(mcpOwner.closeCallCount == 1) + #expect(mcpOwner.waitUntilClosedCallCount == 1) + #expect(store.serverState == .stopped) + #expect(store.serverURL == nil) + } + + @Test func stopSupersedingHeldReplacementCancelsAndJoinsEligibleSibling() async throws { + let targetRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-target", + threadID: "thread-target", + turnID: "turn-target", + reviewThreadID: "review-target" + ) + let siblingRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-sibling", + threadID: "thread-sibling", + turnID: "turn-sibling", + reviewThreadID: "review-sibling" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: targetRun) + await reviewBackend.enqueueRun(siblingRun) + let mcpOwner = TestingMCPServerLifecycleOwner() + let backend = TestingCodexReviewStoreBackend( + reviewBackend: reviewBackend, + mcpServerLifecycle: mcpOwner + ) + let jobIDs = SequentialJobIDs(["job-target", "job-sibling"]) + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { jobIDs.next() }), + reviewRuntimeClosePolicy: .init( + terminalGrace: .seconds(10), + sleep: { _ in } + ) + ) + await store.start() + let targetTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-target", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-target", + attemptID: targetRun.attemptID + ) + let siblingTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-sibling", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-sibling", + attemptID: siblingRun.attemptID + ) + + let replacementGate = AsyncGate() + backend.holdRuntimePreparation(with: replacementGate) + let targetCancellation = Task { @MainActor in + try await store.cancelReview( + jobID: "job-target", + cancellation: .mcpClient(message: "Stop target") + ) + } + await backend.waitForRuntimePreparation() + let stopTask = Task { @MainActor in + await store.stop() + } + await backend.waitForRuntimePreparationCancellation() + await replacementGate.open() + await stopTask.value + let target = try await targetTask.value + let sibling = try await siblingTask.value + _ = try await targetCancellation.value + + #expect(target.core.lifecycle.terminal == .interrupted( + .requested(.mcpClient(message: "Stop target")) + )) + #expect(sibling.core.lifecycle.terminal == .interrupted( + .requested(.system(message: "Review runtime stopped.")) + )) + #expect(store.reviewWorkerTasks.isEmpty) + #expect(store.activeRuntimeReplacementReceiptCountForTesting == 0) + #expect(store.serverState == .stopped) + #expect(mcpOwner.stopCallCount == 1) + #expect(mcpOwner.waitUntilStoppedCallCount == 1) + } + + @Test func recoverableNetworkForceCloseWaitsForRestorationAndResumesOnReplacement() async throws { + let initialRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-network", + threadID: "thread-network", + turnID: "turn-network", + reviewThreadID: "review-network" + ) + let recoveredRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-network-recovered", + threadID: "thread-network", + turnID: "turn-network-recovered", + reviewThreadID: "review-network" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: initialRun) + await reviewBackend.enqueueRecoveredRun(recoveredRun) + let endpoint = try #require(URL(string: "http://127.0.0.1:19434/mcp")) + let mcpOwner = TestingMCPServerLifecycleOwner(serverURL: endpoint) + let backend = TestingCodexReviewStoreBackend( + reviewBackend: reviewBackend, + mcpServerLifecycle: mcpOwner + ) + let networkMonitor = ManualCodexReviewNetworkMonitor() + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { "job-network" }), + networkMonitor: networkMonitor, + networkRecoveryPolicy: .init(sleep: { _ in }), + reviewRuntimeClosePolicy: .init( + terminalGrace: .seconds(10), + sleep: { _ in } + ) + ) + await store.start() + let sourceRuntime = try #require(backend.lastPreparedRuntimeHandle) + let reviewTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-network", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-network", + attemptID: initialRun.attemptID + ) + + let replacementGate = AsyncGate() + backend.holdRuntimePreparation(with: replacementGate) + networkMonitor.yield(.init(status: .unsatisfied)) + await reviewBackend.waitForInterruptReview( + run: initialRun, + reason: .init(message: "Network unavailable; waiting to reconnect.") + ) + await sourceRuntime.waitForClose() + await reviewBackend.waitForPrepareReviewRecovery() + + #expect(sourceRuntime.closeCallCount == 1) + #expect(sourceRuntime.closePurposes == [.recoveryReplacement]) + #expect(backend.startRequests == [false]) + #expect(store.serverURL == endpoint) + #expect(mcpOwner.prepareCallCount == 1) + #expect(mcpOwner.activateCallCount == 1) + #expect(mcpOwner.stopCallCount == 0) + + networkMonitor.yield(.satisfied()) + await backend.waitForRuntimePreparation() + let replacementRuntime = try #require(backend.lastPreparedRuntimeHandle) + #expect(replacementRuntime !== sourceRuntime) + #expect(replacementRuntime.activateCallCount == 0) + #expect(backend.startRequests == [false, false]) + + await replacementGate.open() + await reviewBackend.waitForResumeReviewRecovery() + await reviewBackend.yield( + .completed(summary: "Done", result: "network recovered"), + for: recoveredRun + ) + let review = try await reviewTask.value + + #expect(review.core.lifecycle.status == .succeeded) + #expect(review.core.run.turnID == recoveredRun.turnID) + #expect(replacementRuntime.activateCallCount == 1) + #expect(store.serverState == .running) + #expect(store.serverURL == endpoint) + #expect(mcpOwner.prepareCallCount == 1) + #expect(mcpOwner.activateCallCount == 1) + #expect(mcpOwner.stopCallCount == 0) + + await store.stop() + } + + @Test func networkForceCloseFailureHasOneApplicationCloseReceipt() async throws { + let initialRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-network", + threadID: "thread-network", + turnID: "turn-network", + reviewThreadID: "review-network" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: initialRun) + let backend = TestingCodexReviewStoreBackend(reviewBackend: reviewBackend) + let networkMonitor = ManualCodexReviewNetworkMonitor() + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { "job-network" }), + networkMonitor: networkMonitor, + networkRecoveryPolicy: .init(sleep: { _ in }), + reviewRuntimeClosePolicy: .init( + terminalGrace: .seconds(10), + sleep: { _ in } + ) + ) + let forceReceiptRecorded = AsyncGate() + store.setRuntimeForceCloseReceiptRecordedForTesting { + await forceReceiptRecorded.open() + } + await store.start() + let sourceRuntime = try #require(backend.lastPreparedRuntimeHandle) + sourceRuntime.failClose(with: .init(first: .client("network source close failed"))) + let reviewTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-network", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-network", + attemptID: initialRun.attemptID + ) + + networkMonitor.yield(.init(status: .unsatisfied)) + await forceReceiptRecorded.wait() + let review = try await reviewTask.value + let closeError = try #require(await capturedCloseError(from: store)) + + #expect(review.core.lifecycle.status == .failed) + #expect(sourceRuntime.closeCallCount == 1) + #expect(sourceRuntime.waitUntilClosedCallCount == 1) + #expect(closeError.failures.additionalInLifecycleOrder.isEmpty) + guard case .lifecycleResources(let lifecycle) = closeError.failures.first else { + Issue.record("Network force close failure must retain lifecycle ownership.") + return + } + #expect(lifecycle.first == .client("network source close failed")) + } + + @Test func sameGenerationNetworkResumeIsDiscardedWhenStopWins() async throws { + try await exerciseSameGenerationNetworkResumeDiscard(applicationClose: false) + } + + @Test func sameGenerationNetworkResumeIsDiscardedWhenCloseWins() async throws { + try await exerciseSameGenerationNetworkResumeDiscard(applicationClose: true) + } + + @Test func staleNetworkResumeTerminalizesBeforeCloseCancellationAdmission() async throws { + try await exerciseSameGenerationNetworkResumeDiscard( + applicationClose: true, + terminalBeforeCancellationAdmission: true + ) + } + + private func exerciseSameGenerationNetworkResumeDiscard( + applicationClose: Bool, + terminalBeforeCancellationAdmission: Bool = false + ) async throws { + let initialRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-network", + threadID: "thread-network", + turnID: "turn-network", + reviewThreadID: "review-network" + ) + let recoveredRun = CodexReviewBackendModel.Review.Run( + attemptID: "attempt-network-recovered", + threadID: "thread-network", + turnID: "turn-network-recovered", + reviewThreadID: "review-network" + ) + let reviewBackend = FakeCodexReviewBackend(nextRun: initialRun) + await reviewBackend.enqueueRecoveredRun(recoveredRun) + let resumeGate = AsyncGate() + await reviewBackend.holdResumeReviewRecoveryIgnoringCancellation( + with: resumeGate + ) + let backend = TestingCodexReviewStoreBackend(reviewBackend: reviewBackend) + let networkMonitor = ManualCodexReviewNetworkMonitor() + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { "job-network" }), + networkMonitor: networkMonitor, + networkRecoveryPolicy: .init(sleep: { _ in }) + ) + await store.start() + let runtime = try #require(backend.lastPreparedRuntimeHandle) + let reviewTask = Task { @MainActor in + try await store.startReview( + sessionID: "session-network", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + await store.waitForRuntimeReplacementRegistrationForTesting( + jobID: "job-network", + attemptID: initialRun.attemptID + ) + + let reason = CodexReviewBackendModel.CancellationReason( + message: "Network unavailable; waiting to reconnect." + ) + networkMonitor.yield(.init(status: .unsatisfied)) + await reviewBackend.waitForInterruptReview(run: initialRun, reason: reason) + await reviewBackend.yield(.cancelled(reason.message), for: initialRun) + await reviewBackend.waitForPrepareReviewRecovery() + networkMonitor.yield(.satisfied()) + await reviewBackend.waitForResumeReviewRecovery() + + let cancellationBarrierEntered = AsyncGate() + let cancellationBarrierRelease = AsyncGate() + store.setReviewCancellationBarrierPreparationForTesting { + await cancellationBarrierEntered.open() + if terminalBeforeCancellationAdmission { + await cancellationBarrierRelease.waitIgnoringCancellation() + } + } + let shutdownTask = Task { @MainActor in + do { + if applicationClose { + try await store.close() + } else { + await store.stop() + } + return Result.success(()) + } catch { + return Result.failure(error) + } + } + await cancellationBarrierEntered.wait() + await resumeGate.open() + if terminalBeforeCancellationAdmission { + await reviewBackend.waitForCleanupReview() + await cancellationBarrierRelease.open() + } + try await shutdownTask.value.get() + let review = try await reviewTask.value + + if terminalBeforeCancellationAdmission { + #expect(review.core.lifecycle.status == .failed) + #expect(review.core.lifecycle.terminal?.kind == .interrupted) + #expect(review.core.lifecycle.cancellation == nil) + } else { + let expectedCancellation: ReviewCancellation = applicationClose + ? .system(message: "Review Store closed.") + : .system(message: "Review runtime stopped.") + #expect(review.core.lifecycle.terminal == .interrupted( + .requested(expectedCancellation) + )) + } + #expect(review.core.run.turnID == initialRun.turnID) + #expect(store.reviewWorkerTasks.isEmpty) + #expect(store.activeRuntimeReplacementReceiptCountForTesting == 0) + #expect(runtime.closeCallCount == 1) + let commands = await reviewBackend.recordedCommands() + #expect(commands.contains(.cleanupReview(recoveredRun))) + #expect(commands.contains(.cleanupReview(initialRun))) + #expect(store.serverState == .stopped) + #expect(store.serverURL == nil) + } + @Test func stopInvalidatesHeldRestartBeforeReplacementCanPublish() async throws { let endpoint = try #require(URL(string: "http://127.0.0.1:19422/mcp")) let mcpOwner = TestingMCPServerLifecycleOwner(serverURL: endpoint) @@ -765,6 +1626,8 @@ struct CodexReviewStoreLifecycleTests { await store.start() let runtime = try #require(backend.lastPreparedRuntimeHandle) runtime.failClose(with: .init(first: .client("prior forced close failed"))) + let sourceCloseGate = AsyncGate() + runtime.holdClose(with: sourceCloseGate) _ = try await store.startReview( sessionID: "session-1", request: .init(cwd: "/tmp/project", target: .uncommittedChanges), @@ -777,14 +1640,16 @@ struct CodexReviewStoreLifecycleTests { cancellation: .mcpClient(message: "Stop") ) } - await forceCloseReceiptRecorded.wait() - await terminalPublicationEntered.wait() + await runtime.waitForClose() #expect(try store.readReview(jobID: "job-1").core.lifecycle.status == .running) let closeTask = Task { @MainActor in await capturedCloseError(from: store) } await store.waitForCloseCallersForTesting(1) + await sourceCloseGate.open() + await forceCloseReceiptRecorded.wait() + await terminalPublicationEntered.wait() await terminalPublicationRelease.open() _ = await cancellationTask.result let closeError = try #require(await closeTask.value) diff --git a/Tests/ReviewUITests/ReviewUITests.swift b/Tests/ReviewUITests/ReviewUITests.swift index 939355fc..77959081 100644 --- a/Tests/ReviewUITests/ReviewUITests.swift +++ b/Tests/ReviewUITests/ReviewUITests.swift @@ -6794,7 +6794,11 @@ final class FailingCancellationBackend: PreviewCodexReviewStoreBackend { override func waitUntilStopped() async {} - override func interruptReview(_: CodexReviewBackendModel.Review.Run, reason _: CodexReviewBackendModel.CancellationReason) async throws { + override func interruptReview( + _: CodexReviewBackendModel.Review.Run, + admission _: ReviewStartAdmission, + reason _: CodexReviewBackendModel.CancellationReason + ) async throws { throw ReviewInterruptRequestFailure( outcome: .rejected(code: nil, message: "Cancellation failed.") ) From 57c615967232bd9464bfa33aaf14a04e144ed391 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:29:53 +0900 Subject: [PATCH 35/38] Fix MCP shutdown ownership and cancellation precedence --- .../Store/CodexReviewStoreReviews.swift | 21 +- .../CodexReviewMCPHTTPServer.swift | 904 ++++++++++++++++-- .../CodexReviewMCPProtocolServer.swift | 87 +- .../MCPHTTPNetworkResourceOwner.swift | 522 ++++++++++ .../CodexReviewMCPHTTPServerTests.swift | 382 +++++++- .../CodexReviewStoreLifecycleTests.swift | 3 + 6 files changed, 1770 insertions(+), 149 deletions(-) create mode 100644 Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift diff --git a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift index 22b40ef9..3eae17a4 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift @@ -514,14 +514,25 @@ extension CodexReviewStore { return .init(jobID: job.id, cancelled: false, core: job.core) } - job.cancellationRequested = true - job.core.lifecycle.cancellation = cancellation - job.core.output.summary = cancellation.message - job.core.lifecycle.errorMessage = cancellation.message + let ownedCancellation: ReviewCancellation + if job.cancellationRequested { + guard let recorded = job.core.lifecycle.cancellation else { + throw ReviewAttemptContractFailure( + message: "A requested review cancellation must retain its exact owner." + ) + } + ownedCancellation = recorded + } else { + job.cancellationRequested = true + job.core.lifecycle.cancellation = cancellation + job.core.output.summary = cancellation.message + job.core.lifecycle.errorMessage = cancellation.message + ownedCancellation = cancellation + } do { try await cancelOwnedReviewAttempt( job: job, - cancellation: cancellation + cancellation: ownedCancellation ) } catch { if job.isTerminal == false { diff --git a/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift b/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift index 2c57d1b3..5477c208 100644 --- a/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift +++ b/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift @@ -12,6 +12,7 @@ private let logger = Logger(subsystem: "CodexReviewKit", category: "mcp-http") private struct TrackedHTTPResponse { var response: HTTPResponse var streamCompletion: ActiveRequestCompletion? = nil + var isFiniteResponseStream = false } package extension CodexReviewMCPHTTPServer { @@ -254,7 +255,182 @@ private actor MCPHTTPHandlerEntryGate { } } +private final class MCPHTTPStartCompletionGate: @unchecked Sendable { + private let lock = NSLock() + private var shouldHoldNextCompletion = false + private var isHoldingCompletion = false + private var releaseWasRequested = false + private var continuation: CheckedContinuation? + private var holdWaiters: [CheckedContinuation] = [] + private var admissionClosed = false + private var admissionCloseWaiters: [CheckedContinuation] = [] + + func holdNextCompletion() { + lock.lock() + precondition( + shouldHoldNextCompletion == false && continuation == nil, + "MCPHTTPStartCompletionGate owns at most one held start." + ) + shouldHoldNextCompletion = true + isHoldingCompletion = false + releaseWasRequested = false + admissionClosed = false + lock.unlock() + } + + func waitIfNeeded() async { + await withCheckedContinuation { continuation in + let waiters: [CheckedContinuation] + lock.lock() + guard shouldHoldNextCompletion else { + lock.unlock() + continuation.resume() + return + } + isHoldingCompletion = true + waiters = holdWaiters + holdWaiters.removeAll(keepingCapacity: false) + if releaseWasRequested { + resetLocked() + lock.unlock() + for waiter in waiters { + waiter.resume() + } + continuation.resume() + return + } + self.continuation = continuation + lock.unlock() + for waiter in waiters { + waiter.resume() + } + } + } + + func waitUntilHolding() async { + await withCheckedContinuation { continuation in + lock.lock() + if isHoldingCompletion { + lock.unlock() + continuation.resume() + } else { + holdWaiters.append(continuation) + lock.unlock() + } + } + } + + func release() { + let continuation: CheckedContinuation? + lock.lock() + guard shouldHoldNextCompletion else { + lock.unlock() + return + } + if let held = self.continuation { + self.continuation = nil + continuation = held + resetLocked() + } else { + releaseWasRequested = true + continuation = nil + } + lock.unlock() + continuation?.resume() + } + + func recordAdmissionClosed() { + let waiters: [CheckedContinuation] + lock.lock() + admissionClosed = true + waiters = admissionCloseWaiters + admissionCloseWaiters.removeAll(keepingCapacity: false) + lock.unlock() + for waiter in waiters { + waiter.resume() + } + } + + func waitUntilAdmissionClosed() async { + await withCheckedContinuation { continuation in + lock.lock() + if admissionClosed { + lock.unlock() + continuation.resume() + } else { + admissionCloseWaiters.append(continuation) + lock.unlock() + } + } + } + + private func resetLocked() { + shouldHoldNextCompletion = false + isHoldingCompletion = false + releaseWasRequested = false + } +} + + package actor CodexReviewMCPHTTPServer { + private struct StartingGenerationFailure: Swift.Error, Sendable { + let primary: any Swift.Error + let cleanupFailures: [ReviewLifecycleResourceFailure] + } + + private typealias StartingGenerationResult = Result< + RunningGeneration, + StartingGenerationFailure + > + + private final class StartingGeneration { + let id: UInt64 + let task: Task + var admissionClosed = false + + init( + id: UInt64, + task: Task + ) { + self.id = id + self.task = task + } + } + + private final class RunningGeneration: @unchecked Sendable { + let id: UInt64 + let listener: any Channel + let eventLoopGroup: MultiThreadedEventLoopGroup + let cleanupTask: Task + let boundURL: URL + var listenerCloseTask: Task, Never>? + + init( + id: UInt64, + listener: any Channel, + eventLoopGroup: MultiThreadedEventLoopGroup, + cleanupTask: Task, + boundURL: URL + ) { + self.id = id + self.listener = listener + self.eventLoopGroup = eventLoopGroup + self.cleanupTask = cleanupTask + self.boundURL = boundURL + } + } + + private enum LifecycleState { + case stopped([ReviewLifecycleResourceFailure]) + case starting(StartingGeneration) + case running(RunningGeneration) + case stopping( + id: UInt64, + resources: RunningGeneration?, + task: Task<[ReviewLifecycleResourceFailure], Never> + ) + } + private struct SessionContext { let server: Server let transport: StatefulHTTPServerTransport @@ -273,17 +449,22 @@ package actor CodexReviewMCPHTTPServer { private let adapter: CodexReviewMCPServer private let configuration: CodexReviewMCPHTTPServer.Configuration - private var channel: Channel? - private var eventLoopGroup: MultiThreadedEventLoopGroup? + private var lifecycleState = LifecycleState.stopped([]) + private var nextGenerationID: UInt64 = 0 private var sessions: [String: SessionContext] = [:] - private var cleanupTask: Task? - private var boundURL: URL? private var pendingCloseFailures: [ReviewLifecycleResourceFailure] = [] - private var listenerCloseTask: Task, Never>? private let admissionRegistry = MCPHTTPAdmissionRegistry() private let handlerEntryGate = MCPHTTPHandlerEntryGate() + private let startCompletionGate = MCPHTTPStartCompletionGate() + private let networkResources = MCPHTTPNetworkResourceOwner() private var admittedHandlerDrainDidBegin = false private var admittedHandlerDrainStartWaiters: [CheckedContinuation] = [] + private var eventLoopGroupShutdownCount = 0 + private var eventLoopGroupShutdownWaiters: [ + (count: Int, continuation: CheckedContinuation) + ] = [] + private var nextStartListenerCleanupFailureForTesting: ReviewLifecycleResourceFailure? + private var nextStartEventLoopGroupCleanupFailureForTesting: ReviewLifecycleResourceFailure? package init( adapter: CodexReviewMCPServer, @@ -294,7 +475,13 @@ package actor CodexReviewMCPHTTPServer { } package var url: URL { - boundURL ?? configuration.url() + switch lifecycleState { + case .running(let resources), + .stopping(_, let resources?, _): + resources.boundURL + case .stopped, .starting, .stopping: + configuration.url() + } } package var endpoint: String { @@ -316,6 +503,7 @@ package actor CodexReviewMCPHTTPServer { .childChannelOption(ChannelOptions.allowRemoteHalfClosure, value: true) do { + try Task.checkCancellation() let channel = try await bootstrap.bind( host: configuration.host, port: configuration.port @@ -332,22 +520,57 @@ package actor CodexReviewMCPHTTPServer { } package func start() async throws { - guard channel == nil else { - return + while true { + switch lifecycleState { + case .stopped(let failures): + try throwLifecycleFailures(failures) + nextGenerationID &+= 1 + let id = nextGenerationID + networkResources.open() + let task = Task { [self] in + await performStartGeneration(id: id) + } + let operation = StartingGeneration(id: id, task: task) + lifecycleState = .starting(operation) + let result = await task.value + try await publishStartResult(result, operation: operation) + return + + case .starting(let operation): + let result = await operation.task.value + try await publishStartResult(result, operation: operation) + return + + case .running: + return + + case .stopping(let id, _, let task): + let failures = await task.value + finishStopIfCurrent(id: id, failures: failures) + try throwLifecycleFailures(failures) + } } + } + private func performStartGeneration(id: UInt64) async -> StartingGenerationResult { let admissionRegistry = admissionRegistry let handlerEntryGate = handlerEntryGate + let networkResources = networkResources let group = MultiThreadedEventLoopGroup(numberOfThreads: 1) let bootstrap = ServerBootstrap(group: group) .serverChannelOption(ChannelOptions.backlog, value: 128) .serverChannelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1) .childChannelInitializer { channel in - channel.pipeline.configureHTTPServerPipeline().flatMap { + guard let childRegistration = networkResources.registerChild(channel) else { + return channel.close(mode: .all) + } + return channel.pipeline.configureHTTPServerPipeline().flatMap { channel.pipeline.addHandler(CodexReviewMCPHTTPHandler( server: self, admissionRegistry: admissionRegistry, - entryGate: handlerEntryGate + entryGate: handlerEntryGate, + networkResources: networkResources, + childRegistration: childRegistration )) } } @@ -355,92 +578,309 @@ package actor CodexReviewMCPHTTPServer { .childChannelOption(ChannelOptions.maxMessagesPerRead, value: 1) .childChannelOption(ChannelOptions.allowRemoteHalfClosure, value: true) + var listener: (any Channel)? do { + try Task.checkCancellation() let channel = try await bootstrap.bind( host: configuration.host, port: configuration.port ).get() - guard Task.isCancelled == false else { - try? await channel.close() - try? await group.shutdownGracefully() + listener = channel + await startCompletionGate.waitIfNeeded() + try Task.checkCancellation() + let actualPort = channel.localAddress?.port + let cleanupTask = Task { [weak self] in + guard let self else { return } + await self.sessionCleanupLoop() + } + return .success(RunningGeneration( + id: id, + listener: channel, + eventLoopGroup: group, + cleanupTask: cleanupTask, + boundURL: configuration.url(boundPort: actualPort) + )) + } catch { + networkResources.closeChildAdmission() + var cleanupFailures: [ReviewLifecycleResourceFailure] = [] + if let listener { + do { + try await listener.close() + } catch { + cleanupFailures.append(.mcpServer(error.localizedDescription)) + } + if let injected = nextStartListenerCleanupFailureForTesting { + nextStartListenerCleanupFailureForTesting = nil + cleanupFailures.append(injected) + } + } + await networkResources.closeAndDrainChildren() + await networkResources.closeTaskAdmissionCancelAndDrain() + do { + try await group.shutdownGracefully() + } catch { + cleanupFailures.append(.mcpServer(error.localizedDescription)) + } + if let injected = nextStartEventLoopGroupCleanupFailureForTesting { + nextStartEventLoopGroupCleanupFailureForTesting = nil + cleanupFailures.append(injected) + } + return .failure(.init( + primary: CodexReviewMCPHTTPServer.Error.classifyStartError( + error, + configuration: configuration + ), + cleanupFailures: cleanupFailures + )) + } + } + + private func publishStartResult( + _ result: StartingGenerationResult, + operation: StartingGeneration + ) async throws { + switch result { + case .success(let resources): + if operation.admissionClosed { + if case .starting(let current) = lifecycleState, + current === operation { + lifecycleState = .running(resources) + } + if let failure = await closeListener(resources) { + recordPendingListenerCloseFailureIfRunning( + failure, + resources: resources + ) + } + throw CancellationError() + } + if case .running(let current) = lifecycleState, + current === resources { + return + } + guard case .starting(let current) = lifecycleState, + current === operation else { throw CancellationError() } - self.eventLoopGroup = group - self.channel = channel pendingCloseFailures.removeAll(keepingCapacity: false) - listenerCloseTask = nil admissionRegistry.open() admittedHandlerDrainDidBegin = false - let actualPort = channel.localAddress?.port - boundURL = configuration.url(boundPort: actualPort) - cleanupTask = Task { [weak self] in - await self?.sessionCleanupLoop() - } - logger.info("MCP Streamable HTTP server listening at \(self.url.absoluteString, privacy: .public)") - } catch { - try? await group.shutdownGracefully() - throw CodexReviewMCPHTTPServer.Error.classifyStartError( - error, - configuration: configuration + lifecycleState = .running(resources) + logger.info( + "MCP Streamable HTTP server listening at \(resources.boundURL.absoluteString, privacy: .public)" ) + case .failure(let failure): + if case .starting(let current) = lifecycleState, + current === operation { + lifecycleState = .stopped(failure.cleanupFailures) + } + throw failure.primary } } package func stop() async throws { - await closeAdmission() - cleanupTask?.cancel() - let cleanupTask = cleanupTask - self.cleanupTask = nil + admissionRegistry.close() + networkResources.closeChildAdmission() + networkResources.closeTaskAdmission(kind: .domainHandler) + let id: UInt64 + let task: Task<[ReviewLifecycleResourceFailure], Never> + switch lifecycleState { + case .stopped(let failures): + try throwLifecycleFailures(failures) + return + case .stopping(let currentID, _, let currentTask): + id = currentID + task = currentTask + case .running(let resources): + id = resources.id + let newTask = Task<[ReviewLifecycleResourceFailure], Never> { [self] in + await performStopGeneration(resources) + } + lifecycleState = .stopping( + id: id, + resources: resources, + task: newTask + ) + task = newTask + case .starting(let operation): + closeStartingAdmission(operation) + id = operation.id + let newTask = Task<[ReviewLifecycleResourceFailure], Never> { [self] in + switch await operation.task.value { + case .success(let resources): + return await performStopGeneration(resources) + case .failure(let failure): + return failure.cleanupFailures + } + } + lifecycleState = .stopping(id: id, resources: nil, task: newTask) + task = newTask + } + + let failures = await task.value + finishStopIfCurrent(id: id, failures: failures) + try throwLifecycleFailures(failures) + } + + private func closeStartingAdmission(_ operation: StartingGeneration) { + if operation.admissionClosed == false { + operation.admissionClosed = true + startCompletionGate.recordAdmissionClosed() + } + operation.task.cancel() + } + + private func performStopGeneration( + _ resources: RunningGeneration + ) async -> [ReviewLifecycleResourceFailure] { + admissionRegistry.close() + networkResources.closeChildAdmission() + networkResources.closeTaskAdmission(kind: .domainHandler) + var failures = pendingCloseFailures + pendingCloseFailures.removeAll(keepingCapacity: false) + if let listenerFailure = await closeListener(resources) { + if failures.contains(listenerFailure) == false { + failures.append(listenerFailure) + } + } + resources.cleanupTask.cancel() await waitForAdmittedHandlers() - await cleanupTask?.value + await resources.cleanupTask.value + await networkResources.waitForTasksDrained(kind: .domainHandler) + await networkResources.waitForTasksDrained(kind: .finiteResponseSource) + await networkResources.waitForTasksDrained(kind: .finiteResponseWriter) await closeAllSessions() - if let eventLoopGroup { - do { - try await eventLoopGroup.shutdownGracefully() - } catch { - pendingCloseFailures.append(.mcpServer(error.localizedDescription)) - } + await networkResources.closeAndDrainChildren() + await networkResources.closeTaskAdmissionCancelAndDrain() + let resourceCounts = networkResources.resourceCountsForTesting() + precondition( + resourceCounts.children == 0 && resourceCounts.tasks == 0, + "MCPHTTPNetworkResourceOwner must drain children and Tasks before EventLoopGroup shutdown." + ) + eventLoopGroupShutdownCount += 1 + let shutdownWaiters = eventLoopGroupShutdownWaiters.filter { + eventLoopGroupShutdownCount >= $0.count + } + eventLoopGroupShutdownWaiters.removeAll { + eventLoopGroupShutdownCount >= $0.count + } + for waiter in shutdownWaiters { + waiter.continuation.resume() + } + do { + try await resources.eventLoopGroup.shutdownGracefully() + } catch { + failures.append(.mcpServer(error.localizedDescription)) } - eventLoopGroup = nil - boundURL = nil logger.info("MCP Streamable HTTP server stopped") - if let first = pendingCloseFailures.first { - let aggregate = ReviewLifecycleResourceFailureAggregate( - first: first, - additionalInLifecycleOrder: Array(pendingCloseFailures.dropFirst()) - ) - pendingCloseFailures.removeAll(keepingCapacity: false) - throw aggregate + return failures + } + + private func finishStopIfCurrent( + id: UInt64, + failures: [ReviewLifecycleResourceFailure] + ) { + guard case .stopping(let currentID, _, _) = lifecycleState, + currentID == id else { + return + } + lifecycleState = .stopped(failures) + } + + private func throwLifecycleFailures( + _ failures: [ReviewLifecycleResourceFailure] + ) throws { + guard let first = failures.first else { + return } + throw ReviewLifecycleResourceFailureAggregate( + first: first, + additionalInLifecycleOrder: Array(failures.dropFirst()) + ) } package func closeAdmission() async { admissionRegistry.close() - guard let channel else { + networkResources.closeChildAdmission() + networkResources.closeTaskAdmission(kind: .domainHandler) + var resources: RunningGeneration? + var startingOperation: StartingGeneration? + switch lifecycleState { + case .running(let running): + resources = running + case .stopping(_, let stopping, _): + resources = stopping + case .starting(let operation): + closeStartingAdmission(operation) + startingOperation = operation + resources = nil + case .stopped: + resources = nil + } + if let startingOperation { + switch await startingOperation.task.value { + case .success(let prepared): + if case .starting(let current) = lifecycleState, + current === startingOperation { + lifecycleState = .running(prepared) + } + resources = prepared + case .failure(let failure): + if case .starting(let current) = lifecycleState, + current === startingOperation { + lifecycleState = .stopped(failure.cleanupFailures) + } + return + } + } + guard let resources else { return } + if let failure = await closeListener(resources) { + recordPendingListenerCloseFailureIfRunning( + failure, + resources: resources + ) + } + } + + private func closeListener( + _ resources: RunningGeneration + ) async -> ReviewLifecycleResourceFailure? { let task: Task, Never> - if let listenerCloseTask { - task = listenerCloseTask + if let existing = resources.listenerCloseTask { + task = existing } else { + let listener = resources.listener let newTask = Task, Never> { do { - try await channel.close() + try await listener.close() return .success(()) } catch { return .failure(.mcpServer(error.localizedDescription)) } } - listenerCloseTask = newTask + resources.listenerCloseTask = newTask task = newTask } switch await task.value { case .success: - self.channel = nil + return nil case .failure(let failure): - if pendingCloseFailures.contains(failure) == false { - pendingCloseFailures.append(failure) - } + return failure + } + } + + private func recordPendingListenerCloseFailureIfRunning( + _ failure: ReviewLifecycleResourceFailure, + resources: RunningGeneration + ) { + guard case .running(let current) = lifecycleState, + current === resources else { + return + } + if pendingCloseFailures.contains(failure) == false { + pendingCloseFailures.append(failure) } } @@ -480,7 +920,11 @@ package actor CodexReviewMCPHTTPServer { session.activeRequestCount += 1 sessions[sessionID] = session let response = await session.transport.handleRequest(request) - let (trackedResponse, didFinishRequest) = trackActiveRequest(response, sessionID: sessionID) + let (trackedResponse, didFinishRequest) = trackActiveRequest( + response, + sessionID: sessionID, + method: request.method + ) if didFinishRequest, request.method.uppercased() == "DELETE", trackedResponse.response.statusCode == 200 { await closeSession(sessionID) } @@ -519,7 +963,8 @@ package actor CodexReviewMCPHTTPServer { adapter: adapter, defaultSessionID: sessionID, clientSession: clientSession, - boundedReviewWaitDuration: configuration.boundedReviewWaitDuration + boundedReviewWaitDuration: configuration.boundedReviewWaitDuration, + networkResources: networkResources ) try await server.start(transport: transport) { clientInfo, _ in await clientSession.update(clientInfo: clientInfo) @@ -533,7 +978,11 @@ package actor CodexReviewMCPHTTPServer { ) let response = await transport.handleRequest(request) - let (trackedResponse, didFinishRequest) = trackActiveRequest(response, sessionID: sessionID) + let (trackedResponse, didFinishRequest) = trackActiveRequest( + response, + sessionID: sessionID, + method: request.method + ) if didFinishRequest, case .error = trackedResponse.response { sessions.removeValue(forKey: sessionID) await transport.disconnect() @@ -555,27 +1004,48 @@ package actor CodexReviewMCPHTTPServer { return } await session.transport.disconnect() + await session.server.waitUntilCompleted() + await session.server.stop() await adapter.closeSession(sessionID) logger.info("Closed MCP HTTP session \(sessionID, privacy: .public)") } private func trackActiveRequest( _ response: HTTPResponse, - sessionID: String + sessionID: String, + method: String ) -> (response: TrackedHTTPResponse, didFinishRequest: Bool) { switch response { case .stream(let stream, let headers): + let networkResources = networkResources + let isFiniteResponseStream = method.uppercased() == "POST" let completion = ActiveRequestCompletion { - Task { + guard let receipt = networkResources.registerTask( + kind: .streamCompletion + ) else { + return + } + let task = Task { await self.finishActiveRequest(sessionID: sessionID) + receipt.finish() } + receipt.install(task) } let trackedStream = AsyncThrowingStream(bufferingPolicy: .unbounded) { continuation in let heartbeatTask = makeStreamHeartbeatTask(continuation: continuation) + guard let receipt = networkResources.registerTask( + kind: isFiniteResponseStream ? .finiteResponseSource : .streamBridge + ) else { + heartbeatTask?.cancel() + completion.finish() + continuation.finish() + return + } let task = Task { defer { heartbeatTask?.cancel() completion.finish() + receipt.finish() } do { for try await chunk in stream { @@ -586,14 +1056,21 @@ package actor CodexReviewMCPHTTPServer { continuation.finish(throwing: error) } } + receipt.install(task) continuation.onTermination = { _ in heartbeatTask?.cancel() - task.cancel() - completion.finish() + if isFiniteResponseStream == false { + task.cancel() + completion.finish() + } } } return ( - .init(response: .stream(trackedStream, headers: headers), streamCompletion: completion), + .init( + response: .stream(trackedStream, headers: headers), + streamCompletion: isFiniteResponseStream ? nil : completion, + isFiniteResponseStream: isFiniteResponseStream + ), false ) @@ -617,7 +1094,13 @@ package actor CodexReviewMCPHTTPServer { guard let interval = configuration.streamHeartbeatInterval else { return nil } - return Task { + guard let receipt = networkResources.registerTask( + kind: .streamHeartbeat + ) else { + return nil + } + let task = Task { + defer { receipt.finish() } while Task.isCancelled == false { do { try await Task.sleep(for: interval) @@ -630,6 +1113,8 @@ package actor CodexReviewMCPHTTPServer { continuation.yield(Data(": keep-alive\n\n".utf8)) } } + receipt.install(task) + return task } private func closeAllSessions() async { @@ -657,7 +1142,13 @@ package actor CodexReviewMCPHTTPServer { } package func listenerIsOpenForTesting() -> Bool { - channel != nil + switch lifecycleState { + case .running(let resources), + .stopping(_, let resources?, _): + resources.listenerCloseTask == nil && resources.listener.isActive + case .stopped, .starting, .stopping: + false + } } package func admittedNetworkRequestCountForTesting() -> Int { @@ -678,6 +1169,30 @@ package actor CodexReviewMCPHTTPServer { await handlerEntryGate.release() } + package func holdNextStartCompletionForTesting() async { + startCompletionGate.holdNextCompletion() + } + + package func waitForHeldStartCompletionForTesting() async { + await startCompletionGate.waitUntilHolding() + } + + package func releaseHeldStartCompletionForTesting() async { + startCompletionGate.release() + } + + package func waitForHeldStartAdmissionCloseForTesting() async { + await startCompletionGate.waitUntilAdmissionClosed() + } + + package func failNextStartCleanupForTesting( + listener message: String, + eventLoopGroup groupMessage: String + ) { + nextStartListenerCleanupFailureForTesting = .mcpServer(message) + nextStartEventLoopGroupCleanupFailureForTesting = .mcpServer(groupMessage) + } + package func waitForAdmittedHandlerDrainToBeginForTesting() async { if admittedHandlerDrainDidBegin { return @@ -691,6 +1206,95 @@ package actor CodexReviewMCPHTTPServer { } } + package func networkResourceCountsForTesting() -> (children: Int, tasks: Int) { + networkResources.resourceCountsForTesting() + } + + package func networkTaskRegistrationCountForTesting( + streamWriter: Bool + ) -> Int { + networkResources.taskCountForTesting( + kind: streamWriter ? .streamWriter : .response + ) + } + + package func waitForNetworkTaskRegistrationCountForTesting( + streamWriter: Bool, + count: Int + ) async { + await networkResources.waitForTaskCountForTesting( + kind: streamWriter ? .streamWriter : .response, + count: count + ) + } + + package func holdNextNetworkTaskCompletionForTesting( + streamWriter: Bool + ) { + networkResources.holdNextTaskCompletionForTesting( + kind: streamWriter ? .streamWriter : .response + ) + } + + package func holdNextFiniteResponseSourceCompletionForTesting() { + networkResources.holdNextTaskCompletionForTesting( + kind: .finiteResponseSource + ) + } + + package func finiteResponseSourceCompletionIsHeldForTesting() -> Bool { + networkResources.hasHeldTaskCompletionForTesting() + } + + package func waitForHeldNetworkTaskCompletionForTesting() async { + await networkResources.waitForHeldTaskCompletionForTesting() + } + + package func releaseHeldNetworkTaskCompletionForTesting() { + networkResources.releaseHeldTaskCompletionForTesting() + } + + package func childChannelRegistrationCountForTesting() -> Int { + networkResources.childCountForTesting() + } + + package func waitForChildChannelRegistrationCountForTesting( + _ count: Int + ) async { + await networkResources.waitForChildCountForTesting(count) + } + + package func holdNextChildCloseAcknowledgementForTesting() { + networkResources.holdNextChildCloseAcknowledgementForTesting() + } + + package func waitForHeldChildCloseAcknowledgementForTesting() async { + await networkResources.waitForHeldChildCloseAcknowledgementForTesting() + } + + package func releaseHeldChildCloseAcknowledgementForTesting() { + networkResources.releaseHeldChildCloseAcknowledgementForTesting() + } + + package func eventLoopGroupShutdownCountForTesting() -> Int { + eventLoopGroupShutdownCount + } + + package func waitForEventLoopGroupShutdownCountForTesting( + _ count: Int + ) async { + if eventLoopGroupShutdownCount >= count { + return + } + await withCheckedContinuation { continuation in + if eventLoopGroupShutdownCount >= count { + continuation.resume() + } else { + eventLoopGroupShutdownWaiters.append((count, continuation)) + } + } + } + private func closeExpiredSessions(now: Date) async { var expiredSessionIDs: [String] = [] for (sessionID, context) in sessions { @@ -816,6 +1420,59 @@ private final class ActiveRequestCompletion: @unchecked Sendable { } } +private final class MCPHTTPStreamOwnership: @unchecked Sendable { + private enum TerminationOwner { + case active + case channel + case task + } + + let id = UUID() + let receipt: MCPHTTPNetworkResourceOwner.TaskReceipt + let completion: ActiveRequestCompletion? + private let lock = NSLock() + private var terminationOwner = TerminationOwner.active + + init( + receipt: MCPHTTPNetworkResourceOwner.TaskReceipt, + completion: ActiveRequestCompletion? + ) { + self.receipt = receipt + self.completion = completion + } + + func install(_ task: Task) { + receipt.install(task) + } + + func terminateFromChannel() { + lock.lock() + guard terminationOwner == .active else { + lock.unlock() + return + } + terminationOwner = .channel + lock.unlock() + completion?.finish() + receipt.cancel() + } + + func claimTaskTermination() -> Bool { + lock.lock() + guard terminationOwner == .active else { + lock.unlock() + return false + } + terminationOwner = .task + lock.unlock() + return true + } + + func finishTask() { + receipt.finish() + } +} + private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked Sendable { typealias InboundIn = HTTPServerRequestPart typealias OutboundOut = HTTPServerResponsePart @@ -849,19 +1506,23 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked private let server: CodexReviewMCPHTTPServer private let admissionRegistry: MCPHTTPAdmissionRegistry private let entryGate: MCPHTTPHandlerEntryGate + private let networkResources: MCPHTTPNetworkResourceOwner + private let childRegistration: MCPHTTPNetworkResourceOwner.ChildRegistration private var requestState: RequestState? - private var activeStreamTask: Task? - private var activeStreamID: UUID? - private var activeStreamCompletion: ActiveRequestCompletion? + private var activeStreamOwnership: MCPHTTPStreamOwnership? init( server: CodexReviewMCPHTTPServer, admissionRegistry: MCPHTTPAdmissionRegistry, - entryGate: MCPHTTPHandlerEntryGate + entryGate: MCPHTTPHandlerEntryGate, + networkResources: MCPHTTPNetworkResourceOwner, + childRegistration: MCPHTTPNetworkResourceOwner.ChildRegistration ) { self.server = server self.admissionRegistry = admissionRegistry self.entryGate = entryGate + self.networkResources = networkResources + self.childRegistration = childRegistration } func channelRead(context: ChannelHandlerContext, data: NIOAny) { @@ -886,15 +1547,33 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked ) return } + guard let taskReceipt = networkResources.registerTask( + kind: .response, + child: childRegistration + ) else { + admissionRegistry.finish(admission) + writeAdmissionClosedResponse( + version: state.head.version, + context: context + ) + return + } nonisolated(unsafe) let context = context - Task { + let task = Task { + defer { + admissionRegistry.finish(admission) + taskReceipt.finish() + } await entryGate.waitIfNeeded() + guard Task.isCancelled == false else { + return + } await handleRequest( state: state, - admission: admission, context: context ) } + taskReceipt.install(task) } } @@ -905,12 +1584,14 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked func channelInactive(context: ChannelHandlerContext) { finishActiveStream() + networkResources.cancelTasks(for: childRegistration) context.fireChannelInactive() } func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) { if case ChannelEvent.inputClosed = event { finishActiveStream() + networkResources.cancelTasks(for: childRegistration) context.close(promise: nil) return } @@ -919,23 +1600,19 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked func errorCaught(context: ChannelHandlerContext, error: any Error) { finishActiveStream() + networkResources.cancelTasks(for: childRegistration) context.close(promise: nil) } private func finishActiveStream() { - activeStreamTask?.cancel() - activeStreamCompletion?.finish() - activeStreamTask = nil - activeStreamID = nil - activeStreamCompletion = nil + activeStreamOwnership?.terminateFromChannel() + activeStreamOwnership = nil } private func handleRequest( state: RequestState, - admission: MCPHTTPAdmissionRegistry.Admission, context: ChannelHandlerContext ) async { - defer { admissionRegistry.finish(admission) } let head = state.head let path = head.uri.split(separator: "?").first.map(String.init) ?? head.uri let endpoint = await server.endpoint @@ -1025,24 +1702,28 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked switch response { case .stream(let stream, _): - let streamID = UUID() + guard let streamReceipt = networkResources.registerTask( + kind: trackedResponse.isFiniteResponseStream + ? .finiteResponseWriter + : .streamWriter, + child: childRegistration + ) else { + trackedResponse.streamCompletion?.finish() + return + } + let ownership = MCPHTTPStreamOwnership( + receipt: streamReceipt, + completion: trackedResponse.streamCompletion + ) let registration = eventLoop.makePromise(of: Void.self) eventLoop.execute { guard context.channel.isActive else { trackedResponse.streamCompletion?.finish() + streamReceipt.finish() registration.succeed(()) return } let streamTask = Task { - defer { - eventLoop.execute { - if self.activeStreamID == streamID { - self.activeStreamTask = nil - self.activeStreamID = nil - self.activeStreamCompletion = nil - } - } - } var head = HTTPResponseHead(version: version, status: status) for (name, value) in headers { head.headers.add(name: name, value: value) @@ -1065,31 +1746,29 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked ) } } catch is CancellationError { - trackedResponse.streamCompletion?.finish() - return } catch { - trackedResponse.streamCompletion?.finish() logger.error("MCP SSE stream failed: \(error.localizedDescription, privacy: .public)") } - guard Task.isCancelled == false else { - return + if Task.isCancelled == false { + try? await self.writeResponsePart( + .end(nil), + context: context, + eventLoop: eventLoop + ) } - try? await self.writeResponsePart( - .end(nil), + await self.finishStreamTask( + ownership, context: context, eventLoop: eventLoop ) } context.channel.closeFuture.whenComplete { _ in - trackedResponse.streamCompletion?.finish() - streamTask.cancel() + ownership.terminateFromChannel() } - self.activeStreamTask?.cancel() - self.activeStreamCompletion?.finish() - self.activeStreamTask = streamTask - self.activeStreamID = streamID - self.activeStreamCompletion = trackedResponse.streamCompletion + self.activeStreamOwnership?.terminateFromChannel() + self.activeStreamOwnership = ownership + ownership.install(streamTask) context.read() registration.succeed(()) } @@ -1130,6 +1809,25 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked } } + private func finishStreamTask( + _ ownership: MCPHTTPStreamOwnership, + context: ChannelHandlerContext, + eventLoop: any EventLoop + ) async { + if ownership.claimTaskTermination() { + let completion = eventLoop.makePromise(of: Void.self) + eventLoop.execute { + if self.activeStreamOwnership === ownership { + self.activeStreamOwnership = nil + } + ownership.completion?.finish() + completion.succeed(()) + } + try? await completion.futureResult.get() + } + ownership.finishTask() + } + private func writeResponsePart( _ part: HTTPServerResponsePart, context: ChannelHandlerContext, diff --git a/Sources/CodexReviewMCPServer/CodexReviewMCPProtocolServer.swift b/Sources/CodexReviewMCPServer/CodexReviewMCPProtocolServer.swift index e606b639..34db427e 100644 --- a/Sources/CodexReviewMCPServer/CodexReviewMCPProtocolServer.swift +++ b/Sources/CodexReviewMCPServer/CodexReviewMCPProtocolServer.swift @@ -30,11 +30,12 @@ package actor MCPClientSessionState { } @MainActor -package func makeMCPProtocolServer( +func makeMCPProtocolServer( adapter: CodexReviewMCPServer, defaultSessionID: String? = nil, clientSession: MCPClientSessionState = .init(), - boundedReviewWaitDuration: Duration = .seconds(540) + boundedReviewWaitDuration: Duration = .seconds(540), + networkResources: MCPHTTPNetworkResourceOwner ) async -> Server { let server = Server( name: "codex_review", @@ -46,56 +47,66 @@ package func makeMCPProtocolServer( ) await server.withMethodHandler(ListTools.self) { _ in - let tools = await adapter.tools.map { descriptor in - Tool( - name: descriptor.name.rawValue, - description: descriptor.description, - inputSchema: schema(for: descriptor.name) - ) + try await networkResources.performTask(kind: .domainHandler) { + let tools = await adapter.tools.map { descriptor in + Tool( + name: descriptor.name.rawValue, + description: descriptor.description, + inputSchema: schema(for: descriptor.name) + ) + } + return .init(tools: tools) } - return .init(tools: tools) } await server.withMethodHandler(CallTool.self) { params in - guard let tool = CodexReviewMCP.Tool.Name(rawValue: params.name) else { - return .init( - content: [.text(text: "Unknown tool: \(params.name)", annotations: nil, _meta: nil)], - isError: true - ) - } - - do { - let httpContext = Server.currentHandlerContext?.httpContext - let useBoundedReviewStart = await clientSession.usesBoundedReviewStart(httpContext: httpContext) - let request = try toolRequest( - tool: tool, - arguments: params.arguments ?? [:], - defaultSessionID: defaultSessionID, - boundedReviewWaitDuration: boundedReviewWaitDuration, - useBoundedReviewStart: useBoundedReviewStart - ) - let response = try await adapter.handle(request) - return try toolResult(tool: tool, response: response) - } catch { - return .init( - content: [.text(text: error.localizedDescription, annotations: nil, _meta: nil)], - isError: true - ) + try await networkResources.performTask(kind: .domainHandler) { + guard let tool = CodexReviewMCP.Tool.Name(rawValue: params.name) else { + return .init( + content: [.text(text: "Unknown tool: \(params.name)", annotations: nil, _meta: nil)], + isError: true + ) + } + + do { + let httpContext = Server.currentHandlerContext?.httpContext + let useBoundedReviewStart = await clientSession.usesBoundedReviewStart(httpContext: httpContext) + let request = try toolRequest( + tool: tool, + arguments: params.arguments ?? [:], + defaultSessionID: defaultSessionID, + boundedReviewWaitDuration: boundedReviewWaitDuration, + useBoundedReviewStart: useBoundedReviewStart + ) + let response = try await adapter.handle(request) + return try toolResult(tool: tool, response: response) + } catch { + return .init( + content: [.text(text: error.localizedDescription, annotations: nil, _meta: nil)], + isError: true + ) + } } } await server.withMethodHandler(ListResources.self) { _ in - .init(resources: helpResources.map(\.resource)) + try await networkResources.performTask(kind: .domainHandler) { + .init(resources: helpResources.map(\.resource)) + } } await server.withMethodHandler(ReadResource.self) { params in - let content = helpResources.first { $0.uri == params.uri }?.content - ?? "Resource not found: \(params.uri)" - return .init(contents: [.text(content, uri: params.uri, mimeType: "text/markdown")]) + try await networkResources.performTask(kind: .domainHandler) { + let content = helpResources.first { $0.uri == params.uri }?.content + ?? "Resource not found: \(params.uri)" + return .init(contents: [.text(content, uri: params.uri, mimeType: "text/markdown")]) + } } await server.withMethodHandler(ListResourceTemplates.self) { _ in - .init(templates: helpResourceTemplates) + try await networkResources.performTask(kind: .domainHandler) { + .init(templates: helpResourceTemplates) + } } return server diff --git a/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift b/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift new file mode 100644 index 00000000..87a1ff0c --- /dev/null +++ b/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift @@ -0,0 +1,522 @@ +import Foundation +@preconcurrency import NIOCore + +final class MCPHTTPNetworkResourceOwner: @unchecked Sendable { + struct TaskAdmissionClosed: Error, Sendable { + let kind: TaskKind + } + + enum TaskKind: Hashable, Sendable { + case domainHandler + case response + case finiteResponseSource + case finiteResponseWriter + case streamBridge + case streamHeartbeat + case streamWriter + case streamCompletion + } + + struct ChildRegistration: Hashable, Sendable { + fileprivate let id: UUID + } + + final class TaskReceipt: @unchecked Sendable { + fileprivate let id: UUID + fileprivate let kind: TaskKind + fileprivate let childID: UUID? + private weak var owner: MCPHTTPNetworkResourceOwner? + private let lock = NSLock() + private var cancelTask: (@Sendable () -> Void)? + private var cancellationWasRequested = false + private var didFinish = false + + fileprivate init( + id: UUID, + kind: TaskKind, + childID: UUID?, + owner: MCPHTTPNetworkResourceOwner + ) { + self.id = id + self.kind = kind + self.childID = childID + self.owner = owner + } + + func install( + _ task: Task + ) { + let shouldCancel: Bool + lock.lock() + if didFinish { + shouldCancel = false + } else { + cancelTask = { task.cancel() } + shouldCancel = cancellationWasRequested + } + lock.unlock() + if shouldCancel { + task.cancel() + } + } + + func cancel() { + let cancelTask: (@Sendable () -> Void)? + lock.lock() + cancellationWasRequested = true + cancelTask = self.cancelTask + lock.unlock() + cancelTask?() + } + + func finish() { + let owner: MCPHTTPNetworkResourceOwner? + lock.lock() + guard didFinish == false else { + lock.unlock() + return + } + didFinish = true + cancelTask = nil + owner = self.owner + lock.unlock() + owner?.finishTask(id: id, kind: kind) + } + } + + private final class ChildResource: @unchecked Sendable { + let registration: ChildRegistration + let channel: any Channel + + init(registration: ChildRegistration, channel: any Channel) { + self.registration = registration + self.channel = channel + } + } + + private struct TaskCountWaiter { + let kind: TaskKind + let targetCount: Int + let continuation: CheckedContinuation + } + + private struct ChildCountWaiter { + let targetCount: Int + let continuation: CheckedContinuation + } + + private let lock = NSLock() + private var acceptsChildren = false + private var acceptsTasks = false + private var closedTaskAdmissionKinds: Set = [] + private var children: [UUID: ChildResource] = [:] + private var tasks: [UUID: TaskReceipt] = [:] + private var childDrainWaiters: [CheckedContinuation] = [] + private var taskDrainWaiters: [CheckedContinuation] = [] + private var taskKindDrainWaiters: [ + TaskKind: [CheckedContinuation] + ] = [:] + private var totalTaskCounts: [TaskKind: Int] = [:] + private var taskCountWaiters: [UUID: TaskCountWaiter] = [:] + private var totalChildCount = 0 + private var childCountWaiters: [UUID: ChildCountWaiter] = [:] + private var heldTaskCompletionKind: TaskKind? + private var heldTaskCompletionIDs: Set = [] + private var taskCompletionHoldWaiters: [CheckedContinuation] = [] + private var shouldHoldNextChildCloseAcknowledgement = false + private var heldChildCloseAcknowledgementIDs: Set = [] + private var childCloseHoldWaiters: [CheckedContinuation] = [] + + func open() { + lock.lock() + precondition( + children.isEmpty && tasks.isEmpty, + "MCP network generation must drain resources before reopening." + ) + acceptsChildren = true + acceptsTasks = true + closedTaskAdmissionKinds.removeAll(keepingCapacity: true) + lock.unlock() + } + + func registerChild(_ channel: any Channel) -> ChildRegistration? { + let registration: ChildRegistration + let completedWaiters: [CheckedContinuation] + lock.lock() + guard acceptsChildren else { + lock.unlock() + return nil + } + registration = .init(id: UUID()) + children[registration.id] = ChildResource( + registration: registration, + channel: channel + ) + totalChildCount += 1 + let completedIDs = childCountWaiters.compactMap { id, waiter in + totalChildCount >= waiter.targetCount ? id : nil + } + completedWaiters = completedIDs.compactMap { + childCountWaiters.removeValue(forKey: $0)?.continuation + } + lock.unlock() + for waiter in completedWaiters { + waiter.resume() + } + channel.closeFuture.whenComplete { [weak self] _ in + self?.acknowledgeChildClose(registration) + } + return registration + } + + func closeChildAdmission() { + lock.lock() + acceptsChildren = false + lock.unlock() + } + + func closeAndDrainChildren() async { + let channels = closeChildAdmissionAndSnapshot() + for channel in channels { + channel.close(mode: .all, promise: nil) + } + await withCheckedContinuation { continuation in + lock.lock() + if children.isEmpty { + lock.unlock() + continuation.resume() + } else { + childDrainWaiters.append(continuation) + lock.unlock() + } + } + } + + private func closeChildAdmissionAndSnapshot() -> [any Channel] { + lock.lock() + acceptsChildren = false + let channels = children.values.map(\.channel) + lock.unlock() + return channels + } + + func registerTask( + kind: TaskKind, + child: ChildRegistration? = nil + ) -> TaskReceipt? { + let receipt: TaskReceipt + let completedWaiters: [CheckedContinuation] + lock.lock() + guard acceptsTasks, closedTaskAdmissionKinds.contains(kind) == false else { + lock.unlock() + return nil + } + receipt = TaskReceipt( + id: UUID(), + kind: kind, + childID: child?.id, + owner: self + ) + tasks[receipt.id] = receipt + totalTaskCounts[kind, default: 0] += 1 + let completedIDs = taskCountWaiters.compactMap { id, waiter in + waiter.kind == kind && totalTaskCounts[kind, default: 0] >= waiter.targetCount + ? id + : nil + } + completedWaiters = completedIDs.compactMap { + taskCountWaiters.removeValue(forKey: $0)?.continuation + } + lock.unlock() + for waiter in completedWaiters { + waiter.resume() + } + return receipt + } + + func performTask( + kind: TaskKind, + operation: @escaping @Sendable () async throws -> Success + ) async throws -> Success { + guard let receipt = registerTask(kind: kind) else { + throw TaskAdmissionClosed(kind: kind) + } + let task = Task { + try await operation() + } + receipt.install(task) + return try await withTaskCancellationHandler { + defer { receipt.finish() } + return try await task.value + } onCancel: { + receipt.cancel() + } + } + + func closeTaskAdmission(kind: TaskKind) { + lock.lock() + closedTaskAdmissionKinds.insert(kind) + lock.unlock() + } + + func closeTaskAdmissionCancelAndDrain() async { + let receipts = closeTaskAdmissionAndSnapshot() + for receipt in receipts { + receipt.cancel() + } + await withCheckedContinuation { continuation in + lock.lock() + if tasks.isEmpty { + lock.unlock() + continuation.resume() + } else { + taskDrainWaiters.append(continuation) + lock.unlock() + } + } + } + + private func closeTaskAdmissionAndSnapshot() -> [TaskReceipt] { + lock.lock() + acceptsTasks = false + let receipts = Array(tasks.values) + lock.unlock() + return receipts + } + + func cancelTasks(for child: ChildRegistration) { + let receipts: [TaskReceipt] + lock.lock() + receipts = tasks.values.filter { $0.childID == child.id } + lock.unlock() + for receipt in receipts { + receipt.cancel() + } + } + + private func acknowledgeChildClose(_ registration: ChildRegistration) { + let holdWaiters: [CheckedContinuation] + lock.lock() + guard children[registration.id] != nil else { + lock.unlock() + return + } + if shouldHoldNextChildCloseAcknowledgement { + shouldHoldNextChildCloseAcknowledgement = false + heldChildCloseAcknowledgementIDs.insert(registration.id) + holdWaiters = childCloseHoldWaiters + childCloseHoldWaiters.removeAll(keepingCapacity: false) + lock.unlock() + for waiter in holdWaiters { + waiter.resume() + } + return + } + let waiters = finishChildLocked(id: registration.id) + lock.unlock() + for waiter in waiters { + waiter.resume() + } + } + + private func finishChildLocked( + id: UUID + ) -> [CheckedContinuation] { + children.removeValue(forKey: id) + guard children.isEmpty else { + return [] + } + let waiters = childDrainWaiters + childDrainWaiters.removeAll(keepingCapacity: false) + return waiters + } + + fileprivate func finishTask(id: UUID, kind: TaskKind) { + let holdWaiters: [CheckedContinuation] + lock.lock() + guard tasks[id] != nil else { + lock.unlock() + return + } + if heldTaskCompletionKind == kind { + heldTaskCompletionKind = nil + heldTaskCompletionIDs.insert(id) + holdWaiters = taskCompletionHoldWaiters + taskCompletionHoldWaiters.removeAll(keepingCapacity: false) + lock.unlock() + for waiter in holdWaiters { + waiter.resume() + } + return + } + let waiters = finishTaskLocked(id: id) + lock.unlock() + for waiter in waiters { + waiter.resume() + } + } + + private func finishTaskLocked( + id: UUID + ) -> [CheckedContinuation] { + guard let finished = tasks.removeValue(forKey: id) else { + return [] + } + var waiters: [CheckedContinuation] = [] + if tasks.values.contains(where: { $0.kind == finished.kind }) == false { + waiters.append(contentsOf: taskKindDrainWaiters.removeValue( + forKey: finished.kind + ) ?? []) + } + if tasks.isEmpty { + waiters.append(contentsOf: taskDrainWaiters) + taskDrainWaiters.removeAll(keepingCapacity: false) + } + return waiters + } + + func waitForTasksDrained(kind: TaskKind) async { + await withCheckedContinuation { continuation in + lock.lock() + if tasks.values.contains(where: { $0.kind == kind }) == false { + lock.unlock() + continuation.resume() + } else { + taskKindDrainWaiters[kind, default: []].append(continuation) + lock.unlock() + } + } + } + + func holdNextTaskCompletionForTesting(kind: TaskKind) { + lock.lock() + precondition( + heldTaskCompletionKind == nil && heldTaskCompletionIDs.isEmpty, + "MCP task completion test gate owns one held completion." + ) + heldTaskCompletionKind = kind + lock.unlock() + } + + func waitForHeldTaskCompletionForTesting() async { + await withCheckedContinuation { continuation in + lock.lock() + if heldTaskCompletionIDs.isEmpty == false { + lock.unlock() + continuation.resume() + } else { + taskCompletionHoldWaiters.append(continuation) + lock.unlock() + } + } + } + + func hasHeldTaskCompletionForTesting() -> Bool { + lock.lock() + let hasHeldCompletion = heldTaskCompletionIDs.isEmpty == false + lock.unlock() + return hasHeldCompletion + } + + func releaseHeldTaskCompletionForTesting() { + let waiters: [CheckedContinuation] + lock.lock() + let heldIDs = heldTaskCompletionIDs + heldTaskCompletionIDs.removeAll(keepingCapacity: false) + waiters = heldIDs.flatMap { finishTaskLocked(id: $0) } + lock.unlock() + for waiter in waiters { + waiter.resume() + } + } + + func waitForTaskCountForTesting(kind: TaskKind, count: Int) async { + await withCheckedContinuation { continuation in + lock.lock() + if totalTaskCounts[kind, default: 0] >= count { + lock.unlock() + continuation.resume() + } else { + taskCountWaiters[UUID()] = .init( + kind: kind, + targetCount: count, + continuation: continuation + ) + lock.unlock() + } + } + } + + func taskCountForTesting(kind: TaskKind) -> Int { + lock.lock() + let count = totalTaskCounts[kind, default: 0] + lock.unlock() + return count + } + + func childCountForTesting() -> Int { + lock.lock() + let count = totalChildCount + lock.unlock() + return count + } + + func waitForChildCountForTesting(_ count: Int) async { + await withCheckedContinuation { continuation in + lock.lock() + if totalChildCount >= count { + lock.unlock() + continuation.resume() + } else { + childCountWaiters[UUID()] = .init( + targetCount: count, + continuation: continuation + ) + lock.unlock() + } + } + } + + func holdNextChildCloseAcknowledgementForTesting() { + lock.lock() + precondition( + shouldHoldNextChildCloseAcknowledgement == false + && heldChildCloseAcknowledgementIDs.isEmpty, + "MCP child close test gate owns one acknowledgement." + ) + shouldHoldNextChildCloseAcknowledgement = true + lock.unlock() + } + + func waitForHeldChildCloseAcknowledgementForTesting() async { + await withCheckedContinuation { continuation in + lock.lock() + if heldChildCloseAcknowledgementIDs.isEmpty == false { + lock.unlock() + continuation.resume() + } else { + childCloseHoldWaiters.append(continuation) + lock.unlock() + } + } + } + + func releaseHeldChildCloseAcknowledgementForTesting() { + let waiters: [CheckedContinuation] + lock.lock() + let heldIDs = heldChildCloseAcknowledgementIDs + heldChildCloseAcknowledgementIDs.removeAll(keepingCapacity: false) + waiters = heldIDs.flatMap { finishChildLocked(id: $0) } + lock.unlock() + for waiter in waiters { + waiter.resume() + } + } + + func resourceCountsForTesting() -> (children: Int, tasks: Int) { + lock.lock() + let counts = (children.count, tasks.count) + lock.unlock() + return counts + } +} diff --git a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift index 90b75b05..ee681e9d 100644 --- a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift +++ b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift @@ -176,6 +176,86 @@ struct CodexReviewMCPHTTPServerTests { )) } + @Test(arguments: [false, true]) + func startingGenerationAdmissionCloseCannotBeReopened( + useStop: Bool + ) async throws { + let backend = FakeCodexReviewBackend() + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend) + ) + let server = CodexReviewMCPHTTPServer( + adapter: CodexReviewMCPServer(store: store), + configuration: .init(host: "127.0.0.1", port: 0) + ) + let configuredURL = await server.url + let priorAdmissionCount = await server.admittedNetworkRequestCountForTesting() + let priorChildCount = await server.childChannelRegistrationCountForTesting() + let expectedCleanupFailures: [ReviewLifecycleResourceFailure] = [ + .mcpServer("starting listener cleanup failed"), + .mcpServer("starting event-loop group cleanup failed"), + ] + await server.failNextStartCleanupForTesting( + listener: "starting listener cleanup failed", + eventLoopGroup: "starting event-loop group cleanup failed" + ) + await server.holdNextStartCompletionForTesting() + let startFinished = CompletionFlag() + let startTask = Task { + let wasCancelled: Bool + do { + try await server.start() + wasCancelled = false + } catch is CancellationError { + wasCancelled = true + } catch { + wasCancelled = false + } + await startFinished.complete() + return wasCancelled + } + await server.waitForHeldStartCompletionForTesting() + + let closeFinished = CompletionFlag() + let closeTask = Task { + if useStop == false { + await server.closeAdmission() + } + let failures = await recordedStopFailures(server) + await closeFinished.complete() + return failures + } + await server.waitForHeldStartAdmissionCloseForTesting() + + #expect(await startFinished.isCompleted() == false) + #expect(await closeFinished.isCompleted() == false) + #expect(await server.listenerIsOpenForTesting() == false) + #expect(await server.url == configuredURL) + #expect(await server.admittedNetworkRequestCountForTesting() == priorAdmissionCount) + #expect(await server.childChannelRegistrationCountForTesting() == priorChildCount) + + await server.releaseHeldStartCompletionForTesting() + #expect(await closeTask.value == expectedCleanupFailures) + #expect(await startTask.value) + + #expect(await startFinished.isCompleted()) + #expect(await closeFinished.isCompleted()) + #expect(await server.listenerIsOpenForTesting() == false) + #expect(await server.url == configuredURL) + #expect(await server.admittedNetworkRequestCountForTesting() == priorAdmissionCount) + #expect(await server.childChannelRegistrationCountForTesting() == priorChildCount) + let denied = await server.handleHTTPRequest(HTTPRequest( + method: "GET", + headers: [:], + path: "/mcp" + )) + #expect(denied.statusCode == 503) + let resources = await server.networkResourceCountsForTesting() + #expect(resources.children == 0) + #expect(resources.tasks == 0) + #expect(await recordedStopFailures(server) == expectedCleanupFailures) + } + @Test func streamableHTTPCallsReviewStartWithCustomTarget() async throws { let backend = FakeCodexReviewBackend() let store = CodexReviewStore.makeTestingStore( @@ -1083,11 +1163,15 @@ struct CodexReviewMCPHTTPServerTests { ) let server = CodexReviewMCPHTTPServer( adapter: CodexReviewMCPServer(store: store), - configuration: .init(port: 0) + configuration: .init(host: "127.0.0.1", port: 0) ) try await server.start() let endpoint = await server.url let sessionID = try await initializeSession(endpoint: endpoint) + let priorResponseTaskCount = await server.networkTaskRegistrationCountForTesting( + streamWriter: false + ) + await server.holdNextNetworkTaskCompletionForTesting(streamWriter: false) let priorAdmissionCount = await server.admittedNetworkRequestCountForTesting() await server.holdNextNetworkHandlerEntryForTesting() let requestBody = try makeJSONBody([ @@ -1113,6 +1197,9 @@ struct CodexReviewMCPHTTPServerTests { priorAdmissionCount + 1 ) + await server.releaseNetworkHandlerEntryForTesting() + await backend.waitForStartReview() + let stopFinished = CompletionFlag() let stopTask = Task { try await server.stop() @@ -1124,14 +1211,249 @@ struct CodexReviewMCPHTTPServerTests { #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) != nil) #expect(await stopFinished.isCompleted() == false) - await server.releaseNetworkHandlerEntryForTesting() await requestGate.open() await backend.yield(.completed(summary: "Done", result: "review text")) _ = try? await requestTask.value + await server.waitForNetworkTaskRegistrationCountForTesting( + streamWriter: false, + count: priorResponseTaskCount + 1 + ) + await server.waitForHeldNetworkTaskCompletionForTesting() + #expect(await stopFinished.isCompleted() == false) + #expect(await server.eventLoopGroupShutdownCountForTesting() == 0) + await server.releaseHeldNetworkTaskCompletionForTesting() + try await stopTask.value + + #expect(await stopFinished.isCompleted()) + #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) == nil) + let resources = await server.networkResourceCountsForTesting() + #expect(resources.children == 0) + #expect(resources.tasks == 0) + } + + @Test func stopDrainsFinitePOSTResponseSourceAfterClientCloses() async throws { + let backend = FakeCodexReviewBackend() + let requestGate = AsyncGate() + await backend.holdStartReview(with: requestGate) + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-1" }) + ) + let server = CodexReviewMCPHTTPServer( + adapter: CodexReviewMCPServer(store: store), + configuration: .init(host: "127.0.0.1", port: 0) + ) + try await server.start() + let endpoint = await server.url + let sessionID = try await initializeSession(endpoint: endpoint) + await server.holdNextFiniteResponseSourceCompletionForTesting() + let requestBody = try makeJSONBody([ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": [ + "name": "review_start", + "arguments": [ + "cwd": "/tmp/project", + "target": ["type": "uncommittedChanges"], + ], + ], + ]) + let requestTask = Task { + try await postJSONRPCData( + endpoint: endpoint, + sessionID: sessionID, + bodyData: requestBody + ) + } + await backend.waitForStartReview() + requestTask.cancel() + _ = try? await requestTask.value + + let stopFinished = CompletionFlag() + let stopTask = Task { + try await server.stop() + await stopFinished.complete() + } + await server.waitForAdmittedHandlerDrainToBeginForTesting() + + #expect(await server.finiteResponseSourceCompletionIsHeldForTesting() == false) + #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) != nil) + #expect(await stopFinished.isCompleted() == false) + + await requestGate.open() + await backend.yield(.completed(summary: "Done", result: "review text")) + await server.waitForHeldNetworkTaskCompletionForTesting() + + #expect(await server.finiteResponseSourceCompletionIsHeldForTesting()) + #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) != nil) + #expect(await stopFinished.isCompleted() == false) + #expect(await server.eventLoopGroupShutdownCountForTesting() == 0) + + await server.releaseHeldNetworkTaskCompletionForTesting() try await stopTask.value #expect(await stopFinished.isCompleted()) #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) == nil) + let resources = await server.networkResourceCountsForTesting() + #expect(resources.children == 0) + #expect(resources.tasks == 0) + } + + @Test func stopAwaitsSSEWriterCompletionBeforeEventLoopShutdown() async throws { + let backend = FakeCodexReviewBackend() + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend) + ) + let server = CodexReviewMCPHTTPServer( + adapter: CodexReviewMCPServer(store: store), + configuration: .init(port: 0) + ) + try await server.start() + let endpoint = await server.url + let sessionID = try await initializeSession(endpoint: endpoint) + let priorWriterCount = await server.networkTaskRegistrationCountForTesting( + streamWriter: true + ) + await server.holdNextNetworkTaskCompletionForTesting(streamWriter: true) + + var request = URLRequest(url: endpoint) + request.httpMethod = "GET" + request.setValue("text/event-stream, application/json", forHTTPHeaderField: "Accept") + request.setValue(sessionID, forHTTPHeaderField: "MCP-Session-Id") + let (bytes, response) = try await URLSession.shared.bytes(for: request) + let httpResponse = try #require(response as? HTTPURLResponse) + #expect(httpResponse.statusCode == 200) + await server.waitForNetworkTaskRegistrationCountForTesting( + streamWriter: true, + count: priorWriterCount + 1 + ) + + let stopFinished = CompletionFlag() + let stopTask = Task { + try await server.stop() + await stopFinished.complete() + } + await server.waitForHeldNetworkTaskCompletionForTesting() + + #expect(await stopFinished.isCompleted() == false) + #expect(await server.eventLoopGroupShutdownCountForTesting() == 0) + #expect(await server.networkResourceCountsForTesting().tasks > 0) + + await server.releaseHeldNetworkTaskCompletionForTesting() + try await stopTask.value + _ = bytes + + #expect(await stopFinished.isCompleted()) + #expect(await server.eventLoopGroupShutdownCountForTesting() == 1) + let resources = await server.networkResourceCountsForTesting() + #expect(resources.children == 0) + #expect(resources.tasks == 0) + } + + @Test func concurrentStopAndRestartAwaitAcceptedChildCloseAcknowledgement() async throws { + let backend = FakeCodexReviewBackend() + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend) + ) + let server = CodexReviewMCPHTTPServer( + adapter: CodexReviewMCPServer(store: store), + configuration: .init(host: "127.0.0.1", port: 0) + ) + try await server.start() + let endpoint = await server.url + let priorChildCount = await server.childChannelRegistrationCountForTesting() + let descriptor = try await openRawTCPConnection(endpoint: endpoint) + defer { Darwin.close(descriptor) } + await server.waitForChildChannelRegistrationCountForTesting(priorChildCount + 1) + await server.holdNextChildCloseAcknowledgementForTesting() + let priorShutdownCount = await server.eventLoopGroupShutdownCountForTesting() + + let firstStopFinished = CompletionFlag() + let firstStop = Task { + try await server.stop() + await firstStopFinished.complete() + } + let secondStopFinished = CompletionFlag() + let secondStop = Task { + try await server.stop() + await secondStopFinished.complete() + } + await server.waitForHeldChildCloseAcknowledgementForTesting() + let restartFinished = CompletionFlag() + let restart = Task { + try await server.start() + await restartFinished.complete() + } + + #expect(await rawConnectionReachedEOF(descriptor: descriptor)) + #expect(await firstStopFinished.isCompleted() == false) + #expect(await secondStopFinished.isCompleted() == false) + #expect(await restartFinished.isCompleted() == false) + #expect(await server.eventLoopGroupShutdownCountForTesting() == priorShutdownCount) + + await server.releaseHeldChildCloseAcknowledgementForTesting() + try await firstStop.value + try await secondStop.value + try await restart.value + + #expect(await firstStopFinished.isCompleted()) + #expect(await secondStopFinished.isCompleted()) + #expect(await restartFinished.isCompleted()) + #expect(await server.listenerIsOpenForTesting()) + #expect(await server.eventLoopGroupShutdownCountForTesting() == priorShutdownCount + 1) + let restartedResources = await server.networkResourceCountsForTesting() + #expect(restartedResources.children == 0) + #expect(restartedResources.tasks == 0) + + try await server.stop() + #expect(await server.eventLoopGroupShutdownCountForTesting() == priorShutdownCount + 2) + } + + @Test func channelCloseOwnsSSETerminationWithoutLateEventLoopCleanup() async throws { + let backend = FakeCodexReviewBackend() + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend) + ) + let server = CodexReviewMCPHTTPServer( + adapter: CodexReviewMCPServer(store: store), + configuration: .init( + host: "127.0.0.1", + port: 0, + streamHeartbeatInterval: .milliseconds(50) + ) + ) + try await server.start() + let endpoint = await server.url + let sessionID = try await initializeSession(endpoint: endpoint) + let priorWriterCount = await server.networkTaskRegistrationCountForTesting( + streamWriter: true + ) + await server.holdNextNetworkTaskCompletionForTesting(streamWriter: true) + + try await openAndCloseRawEventStream(endpoint: endpoint, sessionID: sessionID) + await server.waitForNetworkTaskRegistrationCountForTesting( + streamWriter: true, + count: priorWriterCount + 1 + ) + await server.waitForHeldNetworkTaskCompletionForTesting() + + let stopFinished = CompletionFlag() + let stopTask = Task { + try await server.stop() + await stopFinished.complete() + } + #expect(await stopFinished.isCompleted() == false) + #expect(await server.eventLoopGroupShutdownCountForTesting() == 0) + + await server.releaseHeldNetworkTaskCompletionForTesting() + try await stopTask.value + + #expect(await stopFinished.isCompleted()) + #expect(await server.eventLoopGroupShutdownCountForTesting() == 1) + let resources = await server.networkResourceCountsForTesting() + #expect(resources.children == 0) + #expect(resources.tasks == 0) } @Test func streamableHTTPDoesNotExpireSessionWithOpenEventStream() async throws { @@ -1346,6 +1668,19 @@ struct CodexReviewMCPHTTPServerTests { } } + private func recordedStopFailures( + _ server: CodexReviewMCPHTTPServer + ) async -> [ReviewLifecycleResourceFailure] { + do { + try await server.stop() + return [] + } catch let aggregate as ReviewLifecycleResourceFailureAggregate { + return [aggregate.first] + aggregate.additionalInLifecycleOrder + } catch { + return [.mcpServer("Unexpected stop error: \(error.localizedDescription)")] + } + } + private func withHTTPServer( store: CodexReviewStore, configuration: CodexReviewMCPHTTPServer.Configuration = .init(port: 0), @@ -1450,7 +1785,8 @@ struct CodexReviewMCPHTTPServerTests { address.sin_len = UInt8(MemoryLayout.size) address.sin_family = sa_family_t(AF_INET) address.sin_port = in_port_t(port).bigEndian - guard inet_pton(AF_INET, host, &address.sin_addr) == 1 else { + let ipv4Host = host == "localhost" ? "127.0.0.1" : host + guard inet_pton(AF_INET, ipv4Host, &address.sin_addr) == 1 else { throw testError("Unable to resolve IPv4 loopback host \(host)") } let connected = withUnsafePointer(to: &address) { pointer in @@ -1512,6 +1848,46 @@ struct CodexReviewMCPHTTPServerTests { }.value } + private nonisolated func openRawTCPConnection(endpoint: URL) async throws -> Int32 { + try await Task.detached { + let components = try #require(URLComponents(url: endpoint, resolvingAgainstBaseURL: false)) + let host = try #require(components.host) + let port = try #require(components.port) + let descriptor = Darwin.socket(AF_INET, SOCK_STREAM, 0) + guard descriptor >= 0 else { + throw currentPOSIXError() + } + + var address = sockaddr_in() + address.sin_len = UInt8(MemoryLayout.size) + address.sin_family = sa_family_t(AF_INET) + address.sin_port = in_port_t(port).bigEndian + let ipv4Host = host == "localhost" ? "127.0.0.1" : host + guard inet_pton(AF_INET, ipv4Host, &address.sin_addr) == 1 else { + Darwin.close(descriptor) + throw testError("Unable to resolve IPv4 loopback host \(host)") + } + let connected = withUnsafePointer(to: &address) { pointer in + pointer.withMemoryRebound(to: sockaddr.self, capacity: 1) { + Darwin.connect(descriptor, $0, socklen_t(MemoryLayout.size)) + } + } + guard connected == 0 else { + let error = currentPOSIXError() + Darwin.close(descriptor) + throw error + } + return descriptor + }.value + } + + private nonisolated func rawConnectionReachedEOF(descriptor: Int32) async -> Bool { + await Task.detached { + var byte: UInt8 = 0 + return Darwin.recv(descriptor, &byte, 1, 0) == 0 + }.value + } + private func postJSONRPCResponse( endpoint: URL, sessionID: String?, diff --git a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift index bb1ad4d7..96a9a2cb 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift @@ -849,6 +849,9 @@ struct CodexReviewStoreLifecycleTests { ) } await backend.waitForRuntimePreparation() + #expect(try store.readReview(jobID: "job-target").core.lifecycle.cancellation == .mcpClient( + message: "Stop target" + )) let stopTask = Task { @MainActor in await store.stop() } From 6a29a53a142d036a55fe4f3b4da5fe76e4764ede Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 12:14:24 +0900 Subject: [PATCH 36/38] Fix finite MCP response shutdown race --- .../CodexReviewMCPHTTPServer.swift | 231 ++++++++++++++---- .../MCPHTTPNetworkResourceOwner.swift | 200 ++++++++++++++- .../CodexReviewMCPHTTPServerTests.swift | 76 +++++- 3 files changed, 451 insertions(+), 56 deletions(-) diff --git a/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift b/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift index 5477c208..679ce34a 100644 --- a/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift +++ b/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift @@ -9,10 +9,14 @@ import CodexReview private let logger = Logger(subsystem: "CodexReviewKit", category: "mcp-http") +private enum TrackedHTTPStreamLifetime { + case open(ActiveRequestCompletion) + case finite(MCPHTTPNetworkResourceOwner.FiniteResponseOperation) +} + private struct TrackedHTTPResponse { var response: HTTPResponse - var streamCompletion: ActiveRequestCompletion? = nil - var isFiniteResponseStream = false + var streamLifetime: TrackedHTTPStreamLifetime? = nil } package extension CodexReviewMCPHTTPServer { @@ -459,6 +463,8 @@ package actor CodexReviewMCPHTTPServer { private let networkResources = MCPHTTPNetworkResourceOwner() private var admittedHandlerDrainDidBegin = false private var admittedHandlerDrainStartWaiters: [CheckedContinuation] = [] + private var finiteResponseDrainDidBegin = false + private var finiteResponseDrainStartWaiters: [CheckedContinuation] = [] private var eventLoopGroupShutdownCount = 0 private var eventLoopGroupShutdownWaiters: [ (count: Int, continuation: CheckedContinuation) @@ -665,6 +671,7 @@ package actor CodexReviewMCPHTTPServer { pendingCloseFailures.removeAll(keepingCapacity: false) admissionRegistry.open() admittedHandlerDrainDidBegin = false + finiteResponseDrainDidBegin = false lifecycleState = .running(resources) logger.info( "MCP Streamable HTTP server listening at \(resources.boundURL.absoluteString, privacy: .public)" @@ -747,8 +754,8 @@ package actor CodexReviewMCPHTTPServer { await waitForAdmittedHandlers() await resources.cleanupTask.value await networkResources.waitForTasksDrained(kind: .domainHandler) - await networkResources.waitForTasksDrained(kind: .finiteResponseSource) - await networkResources.waitForTasksDrained(kind: .finiteResponseWriter) + beginFiniteResponseDrain() + await networkResources.waitForTasksDrained(kind: .finiteResponse) await closeAllSessions() await networkResources.closeAndDrainChildren() await networkResources.closeTaskAdmissionCancelAndDrain() @@ -894,6 +901,15 @@ package actor CodexReviewMCPHTTPServer { await admissionRegistry.waitUntilDrained() } + private func beginFiniteResponseDrain() { + finiteResponseDrainDidBegin = true + let waiters = finiteResponseDrainStartWaiters + finiteResponseDrainStartWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume() + } + } + package func handleHTTPRequest(_ request: HTTPRequest) async -> HTTPResponse { guard let admission = admissionRegistry.admit() else { return .error( @@ -901,9 +917,12 @@ package actor CodexReviewMCPHTTPServer { .internalError("MCP server is not accepting requests.") ) } - let response = await performTrackedHTTPRequest(request).response + let trackedResponse = await performTrackedHTTPRequest(request) + if case .finite(let operation) = trackedResponse.streamLifetime { + operation.waiveWriter() + } admissionRegistry.finish(admission) - return response + return trackedResponse.response } fileprivate func handleAdmittedHTTPRequest( @@ -1031,35 +1050,77 @@ package actor CodexReviewMCPHTTPServer { } receipt.install(task) } - let trackedStream = AsyncThrowingStream(bufferingPolicy: .unbounded) { continuation in - let heartbeatTask = makeStreamHeartbeatTask(continuation: continuation) - guard let receipt = networkResources.registerTask( - kind: isFiniteResponseStream ? .finiteResponseSource : .streamBridge - ) else { - heartbeatTask?.cancel() + let finiteResponseOperation: MCPHTTPNetworkResourceOwner.FiniteResponseOperation? + if isFiniteResponseStream { + guard let operation = networkResources.registerFiniteResponseOperation() else { completion.finish() - continuation.finish() - return + return ( + .init(response: .error( + statusCode: 503, + .internalError("MCP server is not accepting responses.") + )), + true + ) } - let task = Task { - defer { + finiteResponseOperation = operation + } else { + finiteResponseOperation = nil + } + let trackedStream = AsyncThrowingStream(bufferingPolicy: .unbounded) { continuation in + let heartbeatTask = makeStreamHeartbeatTask(continuation: continuation) + if isFiniteResponseStream { + guard let finiteResponseOperation else { heartbeatTask?.cancel() completion.finish() - receipt.finish() + continuation.finish() + return } - do { - for try await chunk in stream { - continuation.yield(chunk) + let task = Task { + defer { + heartbeatTask?.cancel() + completion.finish() + finiteResponseOperation.finishSource() + } + do { + for try await chunk in stream { + continuation.yield(chunk) + } + continuation.finish() + } catch { + continuation.finish(throwing: error) } + } + finiteResponseOperation.installSource(task) + continuation.onTermination = { _ in + heartbeatTask?.cancel() + } + } else { + guard let receipt = networkResources.registerTask( + kind: .streamBridge + ) else { + heartbeatTask?.cancel() + completion.finish() continuation.finish() - } catch { - continuation.finish(throwing: error) + return } - } - receipt.install(task) - continuation.onTermination = { _ in - heartbeatTask?.cancel() - if isFiniteResponseStream == false { + let task = Task { + defer { + heartbeatTask?.cancel() + completion.finish() + receipt.finish() + } + do { + for try await chunk in stream { + continuation.yield(chunk) + } + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + receipt.install(task) + continuation.onTermination = { _ in + heartbeatTask?.cancel() task.cancel() completion.finish() } @@ -1068,8 +1129,9 @@ package actor CodexReviewMCPHTTPServer { return ( .init( response: .stream(trackedStream, headers: headers), - streamCompletion: isFiniteResponseStream ? nil : completion, - isFiniteResponseStream: isFiniteResponseStream + streamLifetime: finiteResponseOperation.map { + .finite($0) + } ?? .open(completion) ), false ) @@ -1206,6 +1268,23 @@ package actor CodexReviewMCPHTTPServer { } } + package func waitForFiniteResponseDrainToBeginForTesting() async { + if finiteResponseDrainDidBegin { + return + } + await withCheckedContinuation { continuation in + if finiteResponseDrainDidBegin { + continuation.resume() + } else { + finiteResponseDrainStartWaiters.append(continuation) + } + } + } + + package func waitForFiniteResponsesToDrainForTesting() async { + await networkResources.waitForTasksDrained(kind: .finiteResponse) + } + package func networkResourceCountsForTesting() -> (children: Int, tasks: Int) { networkResources.resourceCountsForTesting() } @@ -1236,13 +1315,13 @@ package actor CodexReviewMCPHTTPServer { ) } - package func holdNextFiniteResponseSourceCompletionForTesting() { + package func holdNextFiniteResponseCompletionForTesting() { networkResources.holdNextTaskCompletionForTesting( - kind: .finiteResponseSource + kind: .finiteResponse ) } - package func finiteResponseSourceCompletionIsHeldForTesting() -> Bool { + package func finiteResponseCompletionIsHeldForTesting() -> Bool { networkResources.hasHeldTaskCompletionForTesting() } @@ -1473,6 +1552,53 @@ private final class MCPHTTPStreamOwnership: @unchecked Sendable { } } +private enum MCPHTTPActiveStreamOwnership: Sendable { + case open(MCPHTTPStreamOwnership) + case finite(MCPHTTPNetworkResourceOwner.FiniteResponseOperation) + + var id: UUID { + switch self { + case .open(let ownership): ownership.id + case .finite(let operation): operation.id + } + } + + var completion: ActiveRequestCompletion? { + switch self { + case .open(let ownership): ownership.completion + case .finite: nil + } + } + + func install(_ task: Task) { + switch self { + case .open(let ownership): ownership.install(task) + case .finite(let operation): operation.installWriter(task) + } + } + + func terminateFromChannel() { + switch self { + case .open(let ownership): ownership.terminateFromChannel() + case .finite(let operation): operation.terminateWriterFromChannel() + } + } + + func claimTaskTermination() -> Bool { + switch self { + case .open(let ownership): ownership.claimTaskTermination() + case .finite(let operation): operation.claimWriterTaskTermination() + } + } + + func finishTask() { + switch self { + case .open(let ownership): ownership.finishTask() + case .finite(let operation): operation.finishWriter() + } + } +} + private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked Sendable { typealias InboundIn = HTTPServerRequestPart typealias OutboundOut = HTTPServerResponsePart @@ -1509,7 +1635,7 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked private let networkResources: MCPHTTPNetworkResourceOwner private let childRegistration: MCPHTTPNetworkResourceOwner.ChildRegistration private var requestState: RequestState? - private var activeStreamOwnership: MCPHTTPStreamOwnership? + private var activeStreamOwnership: MCPHTTPActiveStreamOwnership? init( server: CodexReviewMCPHTTPServer, @@ -1702,24 +1828,33 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked switch response { case .stream(let stream, _): - guard let streamReceipt = networkResources.registerTask( - kind: trackedResponse.isFiniteResponseStream - ? .finiteResponseWriter - : .streamWriter, - child: childRegistration - ) else { - trackedResponse.streamCompletion?.finish() + let ownership: MCPHTTPActiveStreamOwnership + switch trackedResponse.streamLifetime { + case .finite(let operation): + guard operation.claimWriter() else { + return + } + ownership = .finite(operation) + case .open(let completion): + guard let streamReceipt = networkResources.registerTask( + kind: .streamWriter, + child: childRegistration + ) else { + completion.finish() + return + } + ownership = .open(MCPHTTPStreamOwnership( + receipt: streamReceipt, + completion: completion + )) + case nil: return } - let ownership = MCPHTTPStreamOwnership( - receipt: streamReceipt, - completion: trackedResponse.streamCompletion - ) let registration = eventLoop.makePromise(of: Void.self) eventLoop.execute { guard context.channel.isActive else { - trackedResponse.streamCompletion?.finish() - streamReceipt.finish() + ownership.completion?.finish() + ownership.finishTask() registration.succeed(()) return } @@ -1810,14 +1945,14 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked } private func finishStreamTask( - _ ownership: MCPHTTPStreamOwnership, + _ ownership: MCPHTTPActiveStreamOwnership, context: ChannelHandlerContext, eventLoop: any EventLoop ) async { if ownership.claimTaskTermination() { let completion = eventLoop.makePromise(of: Void.self) eventLoop.execute { - if self.activeStreamOwnership === ownership { + if self.activeStreamOwnership?.id == ownership.id { self.activeStreamOwnership = nil } ownership.completion?.finish() diff --git a/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift b/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift index 87a1ff0c..d044cf98 100644 --- a/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift +++ b/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift @@ -9,8 +9,7 @@ final class MCPHTTPNetworkResourceOwner: @unchecked Sendable { enum TaskKind: Hashable, Sendable { case domainHandler case response - case finiteResponseSource - case finiteResponseWriter + case finiteResponse case streamBridge case streamHeartbeat case streamWriter @@ -45,18 +44,24 @@ final class MCPHTTPNetworkResourceOwner: @unchecked Sendable { func install( _ task: Task + ) { + installCancellation { task.cancel() } + } + + fileprivate func installCancellation( + _ cancelTask: @escaping @Sendable () -> Void ) { let shouldCancel: Bool lock.lock() if didFinish { shouldCancel = false } else { - cancelTask = { task.cancel() } + self.cancelTask = cancelTask shouldCancel = cancellationWasRequested } lock.unlock() if shouldCancel { - task.cancel() + cancelTask() } } @@ -84,6 +89,184 @@ final class MCPHTTPNetworkResourceOwner: @unchecked Sendable { } } + final class FiniteResponseOperation: @unchecked Sendable { + private enum WriterDisposition: Equatable { + case unresolved + case required + case waived + } + + private enum WriterTerminationOwner: Equatable { + case active + case channel + case task + } + + let id = UUID() + private let receipt: TaskReceipt + private let lock = NSLock() + private var sourceCancellation: (@Sendable () -> Void)? + private var writerCancellation: (@Sendable () -> Void)? + private var sourceFinished = false + private var writerFinished = false + private var writerDisposition = WriterDisposition.unresolved + private var writerTerminationOwner = WriterTerminationOwner.active + private var cancellationWasRequested = false + private var writerCancellationWasRequested = false + private var didFinish = false + + fileprivate init(receipt: TaskReceipt) { + self.receipt = receipt + } + + fileprivate func installOwnerCancellation() { + receipt.installCancellation { [self] in + cancelAll() + } + } + + func installSource(_ task: Task) { + let shouldCancel: Bool + lock.lock() + if sourceFinished { + shouldCancel = false + } else { + sourceCancellation = { task.cancel() } + shouldCancel = cancellationWasRequested + } + lock.unlock() + if shouldCancel { + task.cancel() + } + } + + func finishSource() { + let shouldFinish: Bool + lock.lock() + guard sourceFinished == false else { + lock.unlock() + return + } + sourceFinished = true + sourceCancellation = nil + shouldFinish = claimCompletionIfReadyLocked() + lock.unlock() + if shouldFinish { + receipt.finish() + } + } + + func claimWriter() -> Bool { + lock.lock() + guard writerDisposition == .unresolved else { + lock.unlock() + return false + } + writerDisposition = .required + lock.unlock() + return true + } + + func waiveWriter() { + let shouldFinish: Bool + lock.lock() + guard writerDisposition == .unresolved else { + lock.unlock() + return + } + writerDisposition = .waived + shouldFinish = claimCompletionIfReadyLocked() + lock.unlock() + if shouldFinish { + receipt.finish() + } + } + + func installWriter(_ task: Task) { + let shouldCancel: Bool + lock.lock() + if writerFinished { + shouldCancel = false + } else { + writerCancellation = { task.cancel() } + shouldCancel = cancellationWasRequested + || writerCancellationWasRequested + } + lock.unlock() + if shouldCancel { + task.cancel() + } + } + + func terminateWriterFromChannel() { + let cancelWriter: (@Sendable () -> Void)? + lock.lock() + guard writerTerminationOwner == .active else { + lock.unlock() + return + } + writerTerminationOwner = .channel + writerCancellationWasRequested = true + cancelWriter = writerCancellation + lock.unlock() + cancelWriter?() + } + + func claimWriterTaskTermination() -> Bool { + lock.lock() + guard writerTerminationOwner == .active else { + lock.unlock() + return false + } + writerTerminationOwner = .task + lock.unlock() + return true + } + + func finishWriter() { + let shouldFinish: Bool + lock.lock() + guard writerFinished == false else { + lock.unlock() + return + } + writerFinished = true + writerCancellation = nil + shouldFinish = claimCompletionIfReadyLocked() + lock.unlock() + if shouldFinish { + receipt.finish() + } + } + + private func cancelAll() { + let cancelSource: (@Sendable () -> Void)? + let cancelWriter: (@Sendable () -> Void)? + lock.lock() + cancellationWasRequested = true + cancelSource = sourceCancellation + cancelWriter = writerCancellation + lock.unlock() + cancelSource?() + cancelWriter?() + } + + private func claimCompletionIfReadyLocked() -> Bool { + guard didFinish == false, sourceFinished else { + return false + } + switch writerDisposition { + case .unresolved: + return false + case .required where writerFinished == false: + return false + case .required, .waived: + didFinish = true + return true + } + } + } + private final class ChildResource: @unchecked Sendable { let registration: ChildRegistration let channel: any Channel @@ -234,6 +417,15 @@ final class MCPHTTPNetworkResourceOwner: @unchecked Sendable { return receipt } + func registerFiniteResponseOperation() -> FiniteResponseOperation? { + guard let receipt = registerTask(kind: .finiteResponse) else { + return nil + } + let operation = FiniteResponseOperation(receipt: receipt) + operation.installOwnerCancellation() + return operation + } + func performTask( kind: TaskKind, operation: @escaping @Sendable () async throws -> Success diff --git a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift index ee681e9d..8553d976 100644 --- a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift +++ b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift @@ -618,6 +618,74 @@ struct CodexReviewMCPHTTPServerTests { } } + @Test func stopRacingReturnedFiniteReviewListResponseJoinsOneOperation() async throws { + let backend = FakeCodexReviewBackend() + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend) + ) + let server = CodexReviewMCPHTTPServer( + adapter: CodexReviewMCPServer(store: store), + configuration: .init(host: "127.0.0.1", port: 0) + ) + try await server.start() + let endpoint = await server.url + let sessionID = try await initializeSession(endpoint: endpoint) + store.loadForTesting( + serverState: .running, + workspaces: [.init(cwd: "/tmp/project")], + jobs: [CodexReviewJob.makeForTesting( + id: "running", + sessionID: sessionID, + cwd: "/tmp/project", + targetSummary: "Running", + status: .running, + summary: "Running" + )] + ) + await server.waitForFiniteResponsesToDrainForTesting() + await server.holdNextFiniteResponseCompletionForTesting() + + let response = try await postJSONRPC( + endpoint: endpoint, + sessionID: sessionID, + body: [ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": [ + "name": "review_list", + "arguments": ["limit": 20], + ], + ] + ) + let items = try #require( + response.value(for: ["result", "structuredContent", "items"]) + as? [[String: Any]] + ) + #expect(items.first?["jobId"] as? String == "running") + await server.waitForHeldNetworkTaskCompletionForTesting() + + let stopFinished = CompletionFlag() + let stopTask = Task { + try await server.stop() + await stopFinished.complete() + } + await server.waitForFiniteResponseDrainToBeginForTesting() + + #expect(await stopFinished.isCompleted() == false) + #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) != nil) + #expect(await server.eventLoopGroupShutdownCountForTesting() == 0) + + await server.releaseHeldNetworkTaskCompletionForTesting() + try await stopTask.value + + #expect(await stopFinished.isCompleted()) + #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) == nil) + let resources = await server.networkResourceCountsForTesting() + #expect(resources.children == 0) + #expect(resources.tasks == 0) + } + @Test func streamableHTTPScopesReviewReadToTransportSession() async throws { let backend = FakeCodexReviewBackend() let store = CodexReviewStore.makeTestingStore( @@ -1231,7 +1299,7 @@ struct CodexReviewMCPHTTPServerTests { #expect(resources.tasks == 0) } - @Test func stopDrainsFinitePOSTResponseSourceAfterClientCloses() async throws { + @Test func stopDrainsFinitePOSTResponseAfterClientCloses() async throws { let backend = FakeCodexReviewBackend() let requestGate = AsyncGate() await backend.holdStartReview(with: requestGate) @@ -1246,7 +1314,7 @@ struct CodexReviewMCPHTTPServerTests { try await server.start() let endpoint = await server.url let sessionID = try await initializeSession(endpoint: endpoint) - await server.holdNextFiniteResponseSourceCompletionForTesting() + await server.holdNextFiniteResponseCompletionForTesting() let requestBody = try makeJSONBody([ "jsonrpc": "2.0", "id": 2, @@ -1277,7 +1345,7 @@ struct CodexReviewMCPHTTPServerTests { } await server.waitForAdmittedHandlerDrainToBeginForTesting() - #expect(await server.finiteResponseSourceCompletionIsHeldForTesting() == false) + #expect(await server.finiteResponseCompletionIsHeldForTesting() == false) #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) != nil) #expect(await stopFinished.isCompleted() == false) @@ -1285,7 +1353,7 @@ struct CodexReviewMCPHTTPServerTests { await backend.yield(.completed(summary: "Done", result: "review text")) await server.waitForHeldNetworkTaskCompletionForTesting() - #expect(await server.finiteResponseSourceCompletionIsHeldForTesting()) + #expect(await server.finiteResponseCompletionIsHeldForTesting()) #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) != nil) #expect(await stopFinished.isCompleted() == false) #expect(await server.eventLoopGroupShutdownCountForTesting() == 0) From c6122e87f01a806776843f3d50fe303703f3b1ae Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 12:31:14 +0900 Subject: [PATCH 37/38] Make grace replacement host test deterministic --- .../CodexReviewHostTests/CodexReviewHostTests.swift | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift index da844471..86a6a509 100644 --- a/Tests/CodexReviewHostTests/CodexReviewHostTests.swift +++ b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift @@ -485,6 +485,11 @@ struct CodexReviewHostTests { ) try await firstTransport.enqueue(EmptyResponse(), for: "turn/interrupt") try await firstTransport.enqueue(EmptyResponse(), for: "turn/interrupt") + let targetInterruptResponseGate = AsyncGate() + await firstTransport.holdNext( + method: "turn/interrupt", + gate: targetInterruptResponseGate + ) let secondTransport = FakeJSONRPCTransport() try await secondTransport.enqueue(AppServerAPI.Initialize.Response(), for: "initialize") @@ -510,6 +515,7 @@ struct CodexReviewHostTests { var mcpFactoryCallCount = 0 let routingProbe = HostRecoveryRoutingProbe() let jobIDs = HostSequentialIDs(["job-target", "job-sibling"]) + let graceGate = AsyncGate() let store = CodexReviewStore.makeLiveStoreForTesting( environment: ["HOME": homeURL.path], webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, @@ -520,7 +526,7 @@ struct CodexReviewHostTests { mcpHTTPServerBindChecker: { _ in }, reviewRuntimeClosePolicy: .init( terminalGrace: .seconds(10), - sleep: { _ in } + sleep: { _ in await graceGate.wait() } ), idGenerator: .init(next: { jobIDs.next() }), reviewRecoveryRoutingObserver: { event in @@ -555,6 +561,10 @@ struct CodexReviewHostTests { cancellation: .mcpClient(message: "Stop target") ) } + await firstTransport.waitForRequest(method: "turn/interrupt") + await targetInterruptResponseGate.open() + await firstTransport.waitForResponseDelivery(method: "turn/interrupt") + await graceGate.open() await secondTransport.waitForRequest(method: "review/start") try await secondTransport.emitServerNotification( method: "item/completed", From 71f289c38831e52c9b55da0474afff5d232b5482 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Fri, 21 Aug 2026 14:24:25 +0900 Subject: [PATCH 38/38] Own MCP requests through connection shutdown --- .../CodexReviewMCPHTTPServer.swift | 1314 ++++++-------- .../CodexReviewMCPProtocolServer.swift | 199 ++- .../MCPHTTPNetworkResourceOwner.swift | 1556 +++++++++++------ .../CodexReviewMCPHTTPServerTests.swift | 611 ++++--- 4 files changed, 2167 insertions(+), 1513 deletions(-) diff --git a/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift b/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift index 679ce34a..3cddb5ec 100644 --- a/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift +++ b/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift @@ -9,16 +9,6 @@ import CodexReview private let logger = Logger(subsystem: "CodexReviewKit", category: "mcp-http") -private enum TrackedHTTPStreamLifetime { - case open(ActiveRequestCompletion) - case finite(MCPHTTPNetworkResourceOwner.FiniteResponseOperation) -} - -private struct TrackedHTTPResponse { - var response: HTTPResponse - var streamLifetime: TrackedHTTPStreamLifetime? = nil -} - package extension CodexReviewMCPHTTPServer { enum Error: Swift.Error, LocalizedError, Equatable, Sendable { case addressInUse(host: String, port: Int) @@ -104,118 +94,6 @@ package extension CodexReviewMCPHTTPServer { } } -private final class MCPHTTPAdmissionRegistry: @unchecked Sendable { - struct Admission: Sendable { - fileprivate let id: UUID - } - - private struct AdmissionCountWaiter { - let targetCount: Int - let continuation: CheckedContinuation - } - - private let lock = NSLock() - private var acceptsRequests = false - private var admittedRequestIDs: Set = [] - private var drainWaiters: [CheckedContinuation] = [] - private var totalAdmissionCount = 0 - private var admissionCountWaiters: [UUID: AdmissionCountWaiter] = [:] - - func open() { - lock.lock() - precondition( - admittedRequestIDs.isEmpty, - "MCPHTTPAdmissionRegistry must drain one listener generation before reopening." - ) - acceptsRequests = true - lock.unlock() - } - - func close() { - lock.lock() - acceptsRequests = false - lock.unlock() - } - - func admit() -> Admission? { - let waiters: [CheckedContinuation] - lock.lock() - guard acceptsRequests else { - lock.unlock() - return nil - } - let admission = Admission(id: UUID()) - admittedRequestIDs.insert(admission.id) - totalAdmissionCount += 1 - let completedWaiterIDs = admissionCountWaiters.compactMap { id, waiter in - totalAdmissionCount >= waiter.targetCount ? id : nil - } - waiters = completedWaiterIDs.compactMap { - admissionCountWaiters.removeValue(forKey: $0)?.continuation - } - lock.unlock() - for waiter in waiters { - waiter.resume() - } - return admission - } - - func finish(_ admission: Admission) { - let waiters: [CheckedContinuation] - lock.lock() - precondition( - admittedRequestIDs.remove(admission.id) != nil, - "MCPHTTPAdmissionRegistry owns exactly one completion per admitted request." - ) - if admittedRequestIDs.isEmpty { - waiters = drainWaiters - drainWaiters.removeAll(keepingCapacity: false) - } else { - waiters = [] - } - lock.unlock() - for waiter in waiters { - waiter.resume() - } - } - - func waitUntilDrained() async { - await withCheckedContinuation { continuation in - lock.lock() - if admittedRequestIDs.isEmpty { - lock.unlock() - continuation.resume() - } else { - drainWaiters.append(continuation) - lock.unlock() - } - } - } - - func admissionCount() -> Int { - lock.lock() - let count = totalAdmissionCount - lock.unlock() - return count - } - - func waitForAdmissionCount(_ targetCount: Int) async { - await withCheckedContinuation { continuation in - lock.lock() - if totalAdmissionCount >= targetCount { - lock.unlock() - continuation.resume() - } else { - admissionCountWaiters[UUID()] = .init( - targetCount: targetCount, - continuation: continuation - ) - lock.unlock() - } - } - } -} - private actor MCPHTTPHandlerEntryGate { private var shouldHoldNextEntry = false private var releaseWasRequested = false @@ -389,14 +267,17 @@ package actor CodexReviewMCPHTTPServer { private final class StartingGeneration { let id: UInt64 + let networkResources: MCPHTTPNetworkResourceOwner let task: Task var admissionClosed = false init( id: UInt64, + networkResources: MCPHTTPNetworkResourceOwner, task: Task ) { self.id = id + self.networkResources = networkResources self.task = task } } @@ -407,6 +288,7 @@ package actor CodexReviewMCPHTTPServer { let eventLoopGroup: MultiThreadedEventLoopGroup let cleanupTask: Task let boundURL: URL + let networkResources: MCPHTTPNetworkResourceOwner var listenerCloseTask: Task, Never>? init( @@ -414,13 +296,15 @@ package actor CodexReviewMCPHTTPServer { listener: any Channel, eventLoopGroup: MultiThreadedEventLoopGroup, cleanupTask: Task, - boundURL: URL + boundURL: URL, + networkResources: MCPHTTPNetworkResourceOwner ) { self.id = id self.listener = listener self.eventLoopGroup = eventLoopGroup self.cleanupTask = cleanupTask self.boundURL = boundURL + self.networkResources = networkResources } } @@ -435,12 +319,31 @@ package actor CodexReviewMCPHTTPServer { ) } - private struct SessionContext { - let server: Server + private final class SessionContext: @unchecked Sendable { + let ordinal: UInt64 + var server: Server? let transport: StatefulHTTPServerTransport + let foreignLifetimeWaiter: MCPProtocolServerForeignLifetimeWaiter let createdAt: Date var lastAccessedAt: Date - var activeRequestCount: Int + var semanticCloseTask: Task? + + init( + ordinal: UInt64, + server: Server, + transport: StatefulHTTPServerTransport, + foreignLifetimeWaiter: MCPProtocolServerForeignLifetimeWaiter, + createdAt: Date, + lastAccessedAt: Date + ) { + self.ordinal = ordinal + self.server = server + self.transport = transport + self.foreignLifetimeWaiter = foreignLifetimeWaiter + self.createdAt = createdAt + self.lastAccessedAt = lastAccessedAt + self.semanticCloseTask = nil + } } private struct FixedSessionIDGenerator: SessionIDGenerator { @@ -456,15 +359,12 @@ package actor CodexReviewMCPHTTPServer { private var lifecycleState = LifecycleState.stopped([]) private var nextGenerationID: UInt64 = 0 private var sessions: [String: SessionContext] = [:] + private var nextSessionOrdinal: UInt64 = 0 private var pendingCloseFailures: [ReviewLifecycleResourceFailure] = [] - private let admissionRegistry = MCPHTTPAdmissionRegistry() private let handlerEntryGate = MCPHTTPHandlerEntryGate() private let startCompletionGate = MCPHTTPStartCompletionGate() - private let networkResources = MCPHTTPNetworkResourceOwner() private var admittedHandlerDrainDidBegin = false private var admittedHandlerDrainStartWaiters: [CheckedContinuation] = [] - private var finiteResponseDrainDidBegin = false - private var finiteResponseDrainStartWaiters: [CheckedContinuation] = [] private var eventLoopGroupShutdownCount = 0 private var eventLoopGroupShutdownWaiters: [ (count: Int, continuation: CheckedContinuation) @@ -532,11 +432,18 @@ package actor CodexReviewMCPHTTPServer { try throwLifecycleFailures(failures) nextGenerationID &+= 1 let id = nextGenerationID - networkResources.open() + let networkResources = MCPHTTPNetworkResourceOwner() let task = Task { [self] in - await performStartGeneration(id: id) + await performStartGeneration( + id: id, + networkResources: networkResources + ) } - let operation = StartingGeneration(id: id, task: task) + let operation = StartingGeneration( + id: id, + networkResources: networkResources, + task: task + ) lifecycleState = .starting(operation) let result = await task.value try await publishStartResult(result, operation: operation) @@ -558,25 +465,27 @@ package actor CodexReviewMCPHTTPServer { } } - private func performStartGeneration(id: UInt64) async -> StartingGenerationResult { - let admissionRegistry = admissionRegistry + private func performStartGeneration( + id: UInt64, + networkResources: MCPHTTPNetworkResourceOwner + ) async -> StartingGenerationResult { let handlerEntryGate = handlerEntryGate - let networkResources = networkResources let group = MultiThreadedEventLoopGroup(numberOfThreads: 1) let bootstrap = ServerBootstrap(group: group) .serverChannelOption(ChannelOptions.backlog, value: 128) .serverChannelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1) .childChannelInitializer { channel in - guard let childRegistration = networkResources.registerChild(channel) else { + guard let connection = networkResources.admitConnection(channel) else { return channel.close(mode: .all) } - return channel.pipeline.configureHTTPServerPipeline().flatMap { + return channel.pipeline.configureHTTPServerPipeline( + withPipeliningAssistance: false + ).flatMap { channel.pipeline.addHandler(CodexReviewMCPHTTPHandler( server: self, - admissionRegistry: admissionRegistry, entryGate: handlerEntryGate, networkResources: networkResources, - childRegistration: childRegistration + connection: connection )) } } @@ -604,10 +513,11 @@ package actor CodexReviewMCPHTTPServer { listener: channel, eventLoopGroup: group, cleanupTask: cleanupTask, - boundURL: configuration.url(boundPort: actualPort) + boundURL: configuration.url(boundPort: actualPort), + networkResources: networkResources )) } catch { - networkResources.closeChildAdmission() + networkResources.closeAdmission() var cleanupFailures: [ReviewLifecycleResourceFailure] = [] if let listener { do { @@ -620,8 +530,8 @@ package actor CodexReviewMCPHTTPServer { cleanupFailures.append(injected) } } - await networkResources.closeAndDrainChildren() - await networkResources.closeTaskAdmissionCancelAndDrain() + let closingGeneration = networkResources.beginClosing(.serverStop) + await closingGeneration.waitUntilClosed() do { try await group.shutdownGracefully() } catch { @@ -669,9 +579,7 @@ package actor CodexReviewMCPHTTPServer { throw CancellationError() } pendingCloseFailures.removeAll(keepingCapacity: false) - admissionRegistry.open() admittedHandlerDrainDidBegin = false - finiteResponseDrainDidBegin = false lifecycleState = .running(resources) logger.info( "MCP Streamable HTTP server listening at \(resources.boundURL.absoluteString, privacy: .public)" @@ -686,9 +594,6 @@ package actor CodexReviewMCPHTTPServer { } package func stop() async throws { - admissionRegistry.close() - networkResources.closeChildAdmission() - networkResources.closeTaskAdmission(kind: .domainHandler) let id: UInt64 let task: Task<[ReviewLifecycleResourceFailure], Never> switch lifecycleState { @@ -700,6 +605,7 @@ package actor CodexReviewMCPHTTPServer { task = currentTask case .running(let resources): id = resources.id + resources.networkResources.closeAdmission() let newTask = Task<[ReviewLifecycleResourceFailure], Never> { [self] in await performStopGeneration(resources) } @@ -732,6 +638,7 @@ package actor CodexReviewMCPHTTPServer { private func closeStartingAdmission(_ operation: StartingGeneration) { if operation.admissionClosed == false { operation.admissionClosed = true + operation.networkResources.closeAdmission() startCompletionGate.recordAdmissionClosed() } operation.task.cancel() @@ -740,29 +647,46 @@ package actor CodexReviewMCPHTTPServer { private func performStopGeneration( _ resources: RunningGeneration ) async -> [ReviewLifecycleResourceFailure] { - admissionRegistry.close() - networkResources.closeChildAdmission() - networkResources.closeTaskAdmission(kind: .domainHandler) + resources.networkResources.closeAdmission() var failures = pendingCloseFailures pendingCloseFailures.removeAll(keepingCapacity: false) - if let listenerFailure = await closeListener(resources) { - if failures.contains(listenerFailure) == false { - failures.append(listenerFailure) + let listenerCloseTask = listenerCloseTask(resources) + resources.cleanupTask.cancel() + + let closingSessions = sessions + .map { (id: $0.key, context: $0.value) } + .sorted { $0.context.ordinal < $1.context.ordinal } + sessions.removeAll(keepingCapacity: false) + let disconnectTasks = closingSessions.map { session in + Task { + await session.context.transport.disconnect() + } + } + for task in disconnectTasks { + await task.value + } + let closingGeneration = resources.networkResources.beginClosing(.serverStop) + + switch await listenerCloseTask.value { + case .success: + break + case .failure(let failure): + if failures.contains(failure) == false { + failures.append(failure) } } - resources.cleanupTask.cancel() - await waitForAdmittedHandlers() await resources.cleanupTask.value - await networkResources.waitForTasksDrained(kind: .domainHandler) - beginFiniteResponseDrain() - await networkResources.waitForTasksDrained(kind: .finiteResponse) - await closeAllSessions() - await networkResources.closeAndDrainChildren() - await networkResources.closeTaskAdmissionCancelAndDrain() - let resourceCounts = networkResources.resourceCountsForTesting() + await closingGeneration.waitUntilClosed() + for session in closingSessions { + await session.context.semanticCloseTask?.value + } + await stopProtocolServersAndReleaseForeignLifetimes(closingSessions) + for session in closingSessions { + await adapter.closeSession(session.id) + } precondition( - resourceCounts.children == 0 && resourceCounts.tasks == 0, - "MCPHTTPNetworkResourceOwner must drain children and Tasks before EventLoopGroup shutdown." + resources.networkResources.snapshot().isQuiescent, + "The MCP network generation owner must be quiescent before EventLoopGroup shutdown." ) eventLoopGroupShutdownCount += 1 let shutdownWaiters = eventLoopGroupShutdownWaiters.filter { @@ -783,6 +707,23 @@ package actor CodexReviewMCPHTTPServer { return failures } + private func stopProtocolServersAndReleaseForeignLifetimes( + _ sessions: [(id: String, context: SessionContext)] + ) async { + for session in sessions { + var server = session.context.server + if let current = server { + await current.waitUntilCompleted() + await current.stop() + } + session.context.server = nil + server = nil + } + for session in sessions { + await session.context.foreignLifetimeWaiter.wait() + } + } + private func finishStopIfCurrent( id: UInt64, failures: [ReviewLifecycleResourceFailure] @@ -807,15 +748,14 @@ package actor CodexReviewMCPHTTPServer { } package func closeAdmission() async { - admissionRegistry.close() - networkResources.closeChildAdmission() - networkResources.closeTaskAdmission(kind: .domainHandler) var resources: RunningGeneration? var startingOperation: StartingGeneration? switch lifecycleState { case .running(let running): + running.networkResources.closeAdmission() resources = running case .stopping(_, let stopping, _): + stopping?.networkResources.closeAdmission() resources = stopping case .starting(let operation): closeStartingAdmission(operation) @@ -854,22 +794,7 @@ package actor CodexReviewMCPHTTPServer { private func closeListener( _ resources: RunningGeneration ) async -> ReviewLifecycleResourceFailure? { - let task: Task, Never> - if let existing = resources.listenerCloseTask { - task = existing - } else { - let listener = resources.listener - let newTask = Task, Never> { - do { - try await listener.close() - return .success(()) - } catch { - return .failure(.mcpServer(error.localizedDescription)) - } - } - resources.listenerCloseTask = newTask - task = newTask - } + let task = listenerCloseTask(resources) switch await task.value { case .success: return nil @@ -878,6 +803,25 @@ package actor CodexReviewMCPHTTPServer { } } + private func listenerCloseTask( + _ resources: RunningGeneration + ) -> Task, Never> { + if let existing = resources.listenerCloseTask { + return existing + } + let listener = resources.listener + let newTask = Task, Never> { + do { + try await listener.close() + return .success(()) + } catch { + return .failure(.mcpServer(error.localizedDescription)) + } + } + resources.listenerCloseTask = newTask + return newTask + } + private func recordPendingListenerCloseFailureIfRunning( _ failure: ReviewLifecycleResourceFailure, resources: RunningGeneration @@ -898,77 +842,83 @@ package actor CodexReviewMCPHTTPServer { for waiter in startWaiters { waiter.resume() } - await admissionRegistry.waitUntilDrained() - } - - private func beginFiniteResponseDrain() { - finiteResponseDrainDidBegin = true - let waiters = finiteResponseDrainStartWaiters - finiteResponseDrainStartWaiters.removeAll(keepingCapacity: false) - for waiter in waiters { - waiter.resume() - } - } - - package func handleHTTPRequest(_ request: HTTPRequest) async -> HTTPResponse { - guard let admission = admissionRegistry.admit() else { - return .error( - statusCode: 503, - .internalError("MCP server is not accepting requests.") - ) - } - let trackedResponse = await performTrackedHTTPRequest(request) - if case .finite(let operation) = trackedResponse.streamLifetime { - operation.waiveWriter() + let networkResources: MCPHTTPNetworkResourceOwner? + switch lifecycleState { + case .starting(let operation): + networkResources = operation.networkResources + case .running(let resources), .stopping(_, let resources?, _): + networkResources = resources.networkResources + case .stopped, .stopping: + networkResources = nil } - admissionRegistry.finish(admission) - return trackedResponse.response + await networkResources?.waitForAdmittedHandlingWorkToDrain() } fileprivate func handleAdmittedHTTPRequest( - _ request: HTTPRequest - ) async -> TrackedHTTPResponse { - await performTrackedHTTPRequest(request) + _ request: HTTPRequest, + operation: MCPHTTPNetworkResourceOwner.RequestOperation, + networkResources: MCPHTTPNetworkResourceOwner + ) async -> HTTPResponse { + await performHTTPRequest( + request, + operation: operation, + networkResources: networkResources + ) } - private func performTrackedHTTPRequest(_ request: HTTPRequest) async -> TrackedHTTPResponse { + private func performHTTPRequest( + _ request: HTTPRequest, + operation: MCPHTTPNetworkResourceOwner.RequestOperation, + networkResources: MCPHTTPNetworkResourceOwner + ) async -> HTTPResponse { let sessionID = request.header(HTTPHeaderName.sessionID) - if let sessionID, var session = sessions[sessionID] { + if let sessionID, let session = sessions[sessionID] { + guard isCurrentGeneration(networkResources) else { + operation.beginClosing(.serverStop) + return .error( + statusCode: 503, + .internalError("MCP server is stopping.") + ) + } + operation.bindSession(sessionID) session.lastAccessedAt = Date() - session.activeRequestCount += 1 - sessions[sessionID] = session let response = await session.transport.handleRequest(request) - let (trackedResponse, didFinishRequest) = trackActiveRequest( - response, - sessionID: sessionID, - method: request.method - ) - if didFinishRequest, request.method.uppercased() == "DELETE", trackedResponse.response.statusCode == 200 { - await closeSession(sessionID) + if request.method.uppercased() == "DELETE", response.statusCode == 200 { + scheduleSessionClose( + sessionID, + after: operation, + expected: session + ) } - return trackedResponse + return response } if request.method.uppercased() == "POST", let body = request.body, Self.isInitializeRequest(body) { - return await createSessionAndHandle(request) + return await createSessionAndHandle( + request, + operation: operation, + networkResources: networkResources + ) } if sessionID != nil { - return .init(response: .error(statusCode: 404, .invalidRequest("Not Found: Session not found or expired"))) + return .error(statusCode: 404, .invalidRequest("Not Found: Session not found or expired")) } - return .init( - response: .error( - statusCode: 400, - .invalidRequest("Bad Request: Missing \(HTTPHeaderName.sessionID) header") - ) + return .error( + statusCode: 400, + .invalidRequest("Bad Request: Missing \(HTTPHeaderName.sessionID) header") ) } - private func createSessionAndHandle(_ request: HTTPRequest) async -> TrackedHTTPResponse { + private func createSessionAndHandle( + _ request: HTTPRequest, + operation: MCPHTTPNetworkResourceOwner.RequestOperation, + networkResources: MCPHTTPNetworkResourceOwner + ) async -> HTTPResponse { let sessionID = UUID().uuidString let clientSession = MCPClientSessionState() let transport = StatefulHTTPServerTransport( @@ -976,6 +926,14 @@ package actor CodexReviewMCPHTTPServer { validationPipeline: makeValidationPipeline(), retryInterval: configuration.retryInterval ) + let foreignLifetimeWaiter: MCPProtocolServerForeignLifetimeWaiter + var foreignLifetimeLease: MCPProtocolServerForeignLifetimeWaiter.Lease? + do { + let pair = MCPProtocolServerForeignLifetimeWaiter.makePair() + foreignLifetimeWaiter = pair.waiter + foreignLifetimeLease = pair.lease + } + var protocolServer: Server? do { let server = await makeMCPProtocolServer( @@ -983,37 +941,62 @@ package actor CodexReviewMCPHTTPServer { defaultSessionID: sessionID, clientSession: clientSession, boundedReviewWaitDuration: configuration.boundedReviewWaitDuration, - networkResources: networkResources + networkResources: networkResources, + foreignLifetimeLease: foreignLifetimeLease! ) + foreignLifetimeLease = nil + protocolServer = server + guard isCurrentGeneration(networkResources) else { + throw CancellationError() + } try await server.start(transport: transport) { clientInfo, _ in await clientSession.update(clientInfo: clientInfo) } + guard isCurrentGeneration(networkResources) else { + throw CancellationError() + } + operation.bindSession(sessionID) + nextSessionOrdinal &+= 1 sessions[sessionID] = SessionContext( + ordinal: nextSessionOrdinal, server: server, transport: transport, + foreignLifetimeWaiter: foreignLifetimeWaiter, createdAt: Date(), - lastAccessedAt: Date(), - activeRequestCount: 1 + lastAccessedAt: Date() ) let response = await transport.handleRequest(request) - let (trackedResponse, didFinishRequest) = trackActiveRequest( - response, - sessionID: sessionID, - method: request.method - ) - if didFinishRequest, case .error = trackedResponse.response { + guard isCurrentGeneration(networkResources) else { + if let session = sessions.removeValue(forKey: sessionID) { + await physicallyCloseSession(session) + } + operation.beginClosing(.serverStop) + return .error(statusCode: 503, .internalError("MCP server is stopping.")) + } + if case .error = response { sessions.removeValue(forKey: sessionID) - await transport.disconnect() + await physicallyCloseSessionContext( + server: &protocolServer, + transport: transport, + foreignLifetimeWaiter: foreignLifetimeWaiter + ) } - return trackedResponse + return response } catch { - await transport.disconnect() - return .init( - response: .error( - statusCode: 500, - .internalError("Failed to create MCP session: \(error.localizedDescription)") - ) + foreignLifetimeLease = nil + await physicallyCloseSessionContext( + server: &protocolServer, + transport: transport, + foreignLifetimeWaiter: foreignLifetimeWaiter + ) + if error is CancellationError { + operation.beginClosing(.serverStop) + return .error(statusCode: 503, .internalError("MCP server is stopping.")) + } + return .error( + statusCode: 500, + .internalError("Failed to create MCP session: \(error.localizedDescription)") ) } } @@ -1022,166 +1005,92 @@ package actor CodexReviewMCPHTTPServer { guard let session = sessions.removeValue(forKey: sessionID) else { return } + let matchingOperations = currentNetworkResources()?.snapshot().connections + .flatMap(\.operations) + .filter { $0.boundSessionID == sessionID } ?? [] await session.transport.disconnect() - await session.server.waitUntilCompleted() - await session.server.stop() + for snapshot in matchingOperations { + currentNetworkResources()?.resolve(snapshot.token)?.beginClosing(.sessionClosed) + } + for snapshot in matchingOperations { + _ = await currentNetworkResources()?.resolve(snapshot.token)?.waitUntilClosed() + } + await physicallyCloseSession(session) await adapter.closeSession(sessionID) logger.info("Closed MCP HTTP session \(sessionID, privacy: .public)") } - private func trackActiveRequest( - _ response: HTTPResponse, - sessionID: String, - method: String - ) -> (response: TrackedHTTPResponse, didFinishRequest: Bool) { - switch response { - case .stream(let stream, let headers): - let networkResources = networkResources - let isFiniteResponseStream = method.uppercased() == "POST" - let completion = ActiveRequestCompletion { - guard let receipt = networkResources.registerTask( - kind: .streamCompletion - ) else { - return - } - let task = Task { - await self.finishActiveRequest(sessionID: sessionID) - receipt.finish() - } - receipt.install(task) - } - let finiteResponseOperation: MCPHTTPNetworkResourceOwner.FiniteResponseOperation? - if isFiniteResponseStream { - guard let operation = networkResources.registerFiniteResponseOperation() else { - completion.finish() - return ( - .init(response: .error( - statusCode: 503, - .internalError("MCP server is not accepting responses.") - )), - true - ) - } - finiteResponseOperation = operation - } else { - finiteResponseOperation = nil - } - let trackedStream = AsyncThrowingStream(bufferingPolicy: .unbounded) { continuation in - let heartbeatTask = makeStreamHeartbeatTask(continuation: continuation) - if isFiniteResponseStream { - guard let finiteResponseOperation else { - heartbeatTask?.cancel() - completion.finish() - continuation.finish() - return - } - let task = Task { - defer { - heartbeatTask?.cancel() - completion.finish() - finiteResponseOperation.finishSource() - } - do { - for try await chunk in stream { - continuation.yield(chunk) - } - continuation.finish() - } catch { - continuation.finish(throwing: error) - } - } - finiteResponseOperation.installSource(task) - continuation.onTermination = { _ in - heartbeatTask?.cancel() - } - } else { - guard let receipt = networkResources.registerTask( - kind: .streamBridge - ) else { - heartbeatTask?.cancel() - completion.finish() - continuation.finish() - return - } - let task = Task { - defer { - heartbeatTask?.cancel() - completion.finish() - receipt.finish() - } - do { - for try await chunk in stream { - continuation.yield(chunk) - } - continuation.finish() - } catch { - continuation.finish(throwing: error) - } - } - receipt.install(task) - continuation.onTermination = { _ in - heartbeatTask?.cancel() - task.cancel() - completion.finish() - } - } - } - return ( - .init( - response: .stream(trackedStream, headers: headers), - streamLifetime: finiteResponseOperation.map { - .finite($0) - } ?? .open(completion) - ), - false + private func scheduleSessionClose( + _ sessionID: String, + after operation: MCPHTTPNetworkResourceOwner.RequestOperation, + expected session: SessionContext + ) { + guard session.semanticCloseTask == nil else { return } + let sessionOrdinal = session.ordinal + let task = Task { [weak self] in + _ = await operation.waitUntilClosed() + await self?.completeScheduledSessionClose( + sessionID, + expectedOrdinal: sessionOrdinal ) - - default: - finishActiveRequest(sessionID: sessionID) - return (.init(response: response), true) } + session.semanticCloseTask = task } - private func finishActiveRequest(sessionID: String) { - if var session = sessions[sessionID] { + private func completeScheduledSessionClose( + _ sessionID: String, + expectedOrdinal: UInt64 + ) async { + guard sessions[sessionID]?.ordinal == expectedOrdinal else { return } + await closeSession(sessionID) + } + + fileprivate func requestOperationDidFinish(sessionID: String?) { + if let sessionID, let session = sessions[sessionID] { session.lastAccessedAt = Date() - session.activeRequestCount = max(0, session.activeRequestCount - 1) - sessions[sessionID] = session } } - private func makeStreamHeartbeatTask( - continuation: AsyncThrowingStream.Continuation - ) -> Task? { - guard let interval = configuration.streamHeartbeatInterval else { - return nil - } - guard let receipt = networkResources.registerTask( - kind: .streamHeartbeat - ) else { - return nil - } - let task = Task { - defer { receipt.finish() } - while Task.isCancelled == false { - do { - try await Task.sleep(for: interval) - } catch { - return - } - guard Task.isCancelled == false else { - return - } - continuation.yield(Data(": keep-alive\n\n".utf8)) - } + fileprivate var responseHeartbeatInterval: Duration? { + configuration.streamHeartbeatInterval + } + + private func physicallyCloseSession(_ session: SessionContext) async { + var server = session.server + session.server = nil + await physicallyCloseSessionContext( + server: &server, + transport: session.transport, + foreignLifetimeWaiter: session.foreignLifetimeWaiter + ) + } + + private func physicallyCloseSessionContext( + server: inout Server?, + transport: StatefulHTTPServerTransport, + foreignLifetimeWaiter: MCPProtocolServerForeignLifetimeWaiter + ) async { + await transport.disconnect() + if let current = server { + await current.waitUntilCompleted() + await current.stop() } - receipt.install(task) - return task + server = nil + await foreignLifetimeWaiter.wait() } - private func closeAllSessions() async { - for sessionID in sessions.keys { - await closeSession(sessionID) + private func isCurrentGeneration( + _ networkResources: MCPHTTPNetworkResourceOwner + ) -> Bool { + guard case .running(let resources) = lifecycleState else { return false } + return resources.networkResources === networkResources + } + + private func currentNetworkResources() -> MCPHTTPNetworkResourceOwner? { + switch lifecycleState { + case .starting(let operation): operation.networkResources + case .running(let resources), .stopping(_, let resources?, _): resources.networkResources + case .stopped, .stopping: nil } } @@ -1200,7 +1109,8 @@ package actor CodexReviewMCPHTTPServer { } package func sessionActiveRequestCountForTesting(sessionID: String) -> Int? { - sessions[sessionID]?.activeRequestCount + guard sessions[sessionID] != nil else { return nil } + return currentNetworkResources()?.liveOperationCount(boundTo: sessionID) ?? 0 } package func listenerIsOpenForTesting() -> Bool { @@ -1213,16 +1123,6 @@ package actor CodexReviewMCPHTTPServer { } } - package func admittedNetworkRequestCountForTesting() -> Int { - admissionRegistry.admissionCount() - } - - package func waitForAdmittedNetworkRequestCountForTesting( - _ count: Int - ) async { - await admissionRegistry.waitForAdmissionCount(count) - } - package func holdNextNetworkHandlerEntryForTesting() async { await handlerEntryGate.holdNextEntry() } @@ -1268,91 +1168,18 @@ package actor CodexReviewMCPHTTPServer { } } - package func waitForFiniteResponseDrainToBeginForTesting() async { - if finiteResponseDrainDidBegin { - return - } - await withCheckedContinuation { continuation in - if finiteResponseDrainDidBegin { - continuation.resume() - } else { - finiteResponseDrainStartWaiters.append(continuation) - } - } - } - - package func waitForFiniteResponsesToDrainForTesting() async { - await networkResources.waitForTasksDrained(kind: .finiteResponse) - } - - package func networkResourceCountsForTesting() -> (children: Int, tasks: Int) { - networkResources.resourceCountsForTesting() + func networkSnapshotForTesting() -> MCPHTTPNetworkResourceOwner.Snapshot { + currentNetworkResources()?.snapshot() + ?? .init(revision: 0, phase: .closed, connections: []) } - package func networkTaskRegistrationCountForTesting( - streamWriter: Bool - ) -> Int { - networkResources.taskCountForTesting( - kind: streamWriter ? .streamWriter : .response - ) - } - - package func waitForNetworkTaskRegistrationCountForTesting( - streamWriter: Bool, - count: Int - ) async { - await networkResources.waitForTaskCountForTesting( - kind: streamWriter ? .streamWriter : .response, - count: count - ) - } - - package func holdNextNetworkTaskCompletionForTesting( - streamWriter: Bool - ) { - networkResources.holdNextTaskCompletionForTesting( - kind: streamWriter ? .streamWriter : .response - ) - } - - package func holdNextFiniteResponseCompletionForTesting() { - networkResources.holdNextTaskCompletionForTesting( - kind: .finiteResponse - ) - } - - package func finiteResponseCompletionIsHeldForTesting() -> Bool { - networkResources.hasHeldTaskCompletionForTesting() - } - - package func waitForHeldNetworkTaskCompletionForTesting() async { - await networkResources.waitForHeldTaskCompletionForTesting() - } - - package func releaseHeldNetworkTaskCompletionForTesting() { - networkResources.releaseHeldTaskCompletionForTesting() - } - - package func childChannelRegistrationCountForTesting() -> Int { - networkResources.childCountForTesting() - } - - package func waitForChildChannelRegistrationCountForTesting( - _ count: Int - ) async { - await networkResources.waitForChildCountForTesting(count) - } - - package func holdNextChildCloseAcknowledgementForTesting() { - networkResources.holdNextChildCloseAcknowledgementForTesting() - } - - package func waitForHeldChildCloseAcknowledgementForTesting() async { - await networkResources.waitForHeldChildCloseAcknowledgementForTesting() - } - - package func releaseHeldChildCloseAcknowledgementForTesting() { - networkResources.releaseHeldChildCloseAcknowledgementForTesting() + func nextNetworkSnapshotForTesting( + after revision: UInt64 + ) async -> MCPHTTPNetworkResourceOwner.Snapshot { + guard let resources = currentNetworkResources() else { + return .init(revision: revision &+ 1, phase: .closed, connections: []) + } + return await resources.nextSnapshot(after: revision) } package func eventLoopGroupShutdownCountForTesting() -> Int { @@ -1380,18 +1207,17 @@ package actor CodexReviewMCPHTTPServer { guard now.timeIntervalSince(context.lastAccessedAt) > configuration.sessionTimeout else { continue } - if context.activeRequestCount > 0 { + if currentNetworkResources()?.liveOperationCount(boundTo: sessionID) ?? 0 > 0 { continue } if await adapter.hasActiveReviews(in: sessionID) { - if var session = sessions[sessionID] { + if let session = sessions[sessionID] { session.lastAccessedAt = Date() - sessions[sessionID] = session } continue } if let current = sessions[sessionID], - current.activeRequestCount == 0, + currentNetworkResources()?.liveOperationCount(boundTo: sessionID) ?? 0 == 0, now.timeIntervalSince(current.lastAccessedAt) > configuration.sessionTimeout { expiredSessionIDs.append(sessionID) @@ -1478,127 +1304,6 @@ package actor CodexReviewMCPHTTPServer { } } -private final class ActiveRequestCompletion: @unchecked Sendable { - private let lock = NSLock() - private let onFinish: @Sendable () -> Void - private var didFinish = false - - init(onFinish: @escaping @Sendable () -> Void) { - self.onFinish = onFinish - } - - func finish() { - lock.lock() - if didFinish { - lock.unlock() - return - } - didFinish = true - lock.unlock() - onFinish() - } -} - -private final class MCPHTTPStreamOwnership: @unchecked Sendable { - private enum TerminationOwner { - case active - case channel - case task - } - - let id = UUID() - let receipt: MCPHTTPNetworkResourceOwner.TaskReceipt - let completion: ActiveRequestCompletion? - private let lock = NSLock() - private var terminationOwner = TerminationOwner.active - - init( - receipt: MCPHTTPNetworkResourceOwner.TaskReceipt, - completion: ActiveRequestCompletion? - ) { - self.receipt = receipt - self.completion = completion - } - - func install(_ task: Task) { - receipt.install(task) - } - - func terminateFromChannel() { - lock.lock() - guard terminationOwner == .active else { - lock.unlock() - return - } - terminationOwner = .channel - lock.unlock() - completion?.finish() - receipt.cancel() - } - - func claimTaskTermination() -> Bool { - lock.lock() - guard terminationOwner == .active else { - lock.unlock() - return false - } - terminationOwner = .task - lock.unlock() - return true - } - - func finishTask() { - receipt.finish() - } -} - -private enum MCPHTTPActiveStreamOwnership: Sendable { - case open(MCPHTTPStreamOwnership) - case finite(MCPHTTPNetworkResourceOwner.FiniteResponseOperation) - - var id: UUID { - switch self { - case .open(let ownership): ownership.id - case .finite(let operation): operation.id - } - } - - var completion: ActiveRequestCompletion? { - switch self { - case .open(let ownership): ownership.completion - case .finite: nil - } - } - - func install(_ task: Task) { - switch self { - case .open(let ownership): ownership.install(task) - case .finite(let operation): operation.installWriter(task) - } - } - - func terminateFromChannel() { - switch self { - case .open(let ownership): ownership.terminateFromChannel() - case .finite(let operation): operation.terminateWriterFromChannel() - } - } - - func claimTaskTermination() -> Bool { - switch self { - case .open(let ownership): ownership.claimTaskTermination() - case .finite(let operation): operation.claimWriterTaskTermination() - } - } - - func finishTask() { - switch self { - case .open(let ownership): ownership.finishTask() - case .finite(let operation): operation.finishWriter() - } - } -} - private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked Sendable { typealias InboundIn = HTTPServerRequestPart typealias OutboundOut = HTTPServerResponsePart @@ -1629,26 +1334,29 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked var bodyBuffer: ByteBuffer } + private enum WriterEvent: Sendable { + case body(Data) + case heartbeat + case sourceFinished + case sourceFailed(String) + } + private let server: CodexReviewMCPHTTPServer - private let admissionRegistry: MCPHTTPAdmissionRegistry private let entryGate: MCPHTTPHandlerEntryGate private let networkResources: MCPHTTPNetworkResourceOwner - private let childRegistration: MCPHTTPNetworkResourceOwner.ChildRegistration + private let connection: MCPHTTPNetworkResourceOwner.Connection private var requestState: RequestState? - private var activeStreamOwnership: MCPHTTPActiveStreamOwnership? init( server: CodexReviewMCPHTTPServer, - admissionRegistry: MCPHTTPAdmissionRegistry, entryGate: MCPHTTPHandlerEntryGate, networkResources: MCPHTTPNetworkResourceOwner, - childRegistration: MCPHTTPNetworkResourceOwner.ChildRegistration + connection: MCPHTTPNetworkResourceOwner.Connection ) { self.server = server - self.admissionRegistry = admissionRegistry self.entryGate = entryGate self.networkResources = networkResources - self.childRegistration = childRegistration + self.connection = connection } func channelRead(context: ChannelHandlerContext, data: NIOAny) { @@ -1666,18 +1374,10 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked return } requestState = nil - guard let admission = admissionRegistry.admit() else { - writeAdmissionClosedResponse( - version: state.head.version, - context: context - ) - return - } - guard let taskReceipt = networkResources.registerTask( - kind: .response, - child: childRegistration - ) else { - admissionRegistry.finish(admission) + let metadata = makeRequestMetadata(from: state) + guard let operation = connection.admitRequest(metadata: metadata), + let httpReservation = operation.beginHTTPHandling() + else { writeAdmissionClosedResponse( version: state.head.version, context: context @@ -1686,20 +1386,18 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked } nonisolated(unsafe) let context = context let task = Task { - defer { - admissionRegistry.finish(admission) - taskReceipt.finish() - } + defer { httpReservation.acknowledge() } await entryGate.waitIfNeeded() guard Task.isCancelled == false else { return } await handleRequest( state: state, + operation: operation, context: context ) } - taskReceipt.install(task) + httpReservation.install(task) } } @@ -1709,15 +1407,13 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked } func channelInactive(context: ChannelHandlerContext) { - finishActiveStream() - networkResources.cancelTasks(for: childRegistration) + connection.beginClosing(.peerClosed) context.fireChannelInactive() } func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) { if case ChannelEvent.inputClosed = event { - finishActiveStream() - networkResources.cancelTasks(for: childRegistration) + connection.beginClosing(.peerClosed) context.close(promise: nil) return } @@ -1725,35 +1421,42 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked } func errorCaught(context: ChannelHandlerContext, error: any Error) { - finishActiveStream() - networkResources.cancelTasks(for: childRegistration) + connection.beginClosing(.transportFailure(error.localizedDescription)) context.close(promise: nil) } - private func finishActiveStream() { - activeStreamOwnership?.terminateFromChannel() - activeStreamOwnership = nil - } - private func handleRequest( state: RequestState, + operation: MCPHTTPNetworkResourceOwner.RequestOperation, context: ChannelHandlerContext ) async { let head = state.head let path = head.uri.split(separator: "?").first.map(String.init) ?? head.uri let endpoint = await server.endpoint + let response: HTTPResponse guard path == endpoint else { - await writeResponse( - .init(response: .error(statusCode: 404, .invalidRequest("Not Found"))), + response = .error(statusCode: 404, .invalidRequest("Not Found")) + await prepareAndQueueResponse( + response, + operation: operation, version: head.version, context: context ) return } - let request = makeHTTPRequest(from: state) - let response = await server.handleAdmittedHTTPRequest(request) - await writeResponse(response, version: head.version, context: context) + let request = makeHTTPRequest(from: state, token: operation.token) + response = await server.handleAdmittedHTTPRequest( + request, + operation: operation, + networkResources: networkResources + ) + await prepareAndQueueResponse( + response, + operation: operation, + version: head.version, + context: context + ) } private func writeAdmissionClosedResponse( @@ -1787,9 +1490,17 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked context.writeAndFlush(wrapOutboundOut(.end(nil)), promise: nil) } - private func makeHTTPRequest(from state: RequestState) -> HTTPRequest { + private func makeHTTPRequest( + from state: RequestState, + token: MCPHTTPNetworkResourceOwner.OperationToken + ) -> HTTPRequest { var headers: [String: String] = [:] for (name, value) in state.head.headers { + guard name.caseInsensitiveCompare( + MCPHTTPNetworkResourceOwner.operationTokenHeaderName + ) != .orderedSame else { + continue + } if let existing = headers[name] { headers[name] = "\(existing), \(value)" } else { @@ -1805,6 +1516,7 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked } else { body = nil } + headers[MCPHTTPNetworkResourceOwner.operationTokenHeaderName] = token.headerValue let path = String(state.head.uri.split(separator: "?").first ?? Substring(state.head.uri)) return HTTPRequest( @@ -1815,152 +1527,206 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked ) } - private func writeResponse( - _ trackedResponse: TrackedHTTPResponse, + private func makeRequestMetadata( + from state: RequestState + ) -> MCPHTTPNetworkResourceOwner.RequestMetadata { + let path = String(state.head.uri.split(separator: "?").first ?? Substring(state.head.uri)) + let jsonRPCID: String? + if state.bodyBuffer.readableBytes > 0, + let bytes = state.bodyBuffer.getBytes(at: 0, length: state.bodyBuffer.readableBytes), + let object = try? JSONSerialization.jsonObject(with: Data(bytes)) as? [String: Any], + let id = object["id"] + { + if let string = id as? String { + jsonRPCID = string + } else if let number = id as? NSNumber { + jsonRPCID = number.stringValue + } else { + jsonRPCID = nil + } + } else { + jsonRPCID = nil + } + return .init( + method: state.head.method.rawValue, + path: path, + jsonRPCID: jsonRPCID + ) + } + + private func prepareAndQueueResponse( + _ response: HTTPResponse, + operation: MCPHTTPNetworkResourceOwner.RequestOperation, version: HTTPVersion, context: ChannelHandlerContext ) async { + guard operation.beginResponding() else { return } + let preparedResponse: HTTPResponse + switch response { + case .stream(let source, let headers): + guard let sourceReservation = operation.bindResponseSource() else { + operation.beginClosing(.transportFailure("Response source admission was closed.")) + return + } + let tracked = AsyncThrowingStream(bufferingPolicy: .unbounded) { continuation in + let sourceTask = Task { + do { + for try await chunk in source { + try Task.checkCancellation() + continuation.yield(chunk) + } + continuation.finish() + sourceReservation.acknowledge() + } catch is CancellationError { + continuation.finish() + sourceReservation.acknowledge() + } catch { + continuation.finish(throwing: error) + sourceReservation.acknowledge(.failed(error.localizedDescription)) + } + } + sourceReservation.install(sourceTask) + continuation.onTermination = { _ in sourceTask.cancel() } + } + preparedResponse = .stream(tracked, headers: headers) + default: + operation.markResponseSourceNotRequired() + preparedResponse = response + } + + guard await connection.supplyResponse(for: operation), + let writerReservation = operation.bindWriter() + else { return } nonisolated(unsafe) let context = context + let heartbeatInterval = await server.responseHeartbeatInterval + let writerTask = Task { + let completion = await self.writeResponse( + preparedResponse, + operation: operation, + version: version, + context: context, + heartbeatInterval: heartbeatInterval + ) + switch completion { + case .responded: + operation.acknowledgeResponseEnd() + writerReservation.acknowledge() + case .cancelled: + writerReservation.acknowledge() + case .failed(let message): + writerReservation.acknowledge(.failed(message)) + } + await self.server.requestOperationDidFinish( + sessionID: operation.snapshot().boundSessionID + ) + } + writerReservation.install(writerTask) + } + + private enum WriterCompletion { + case responded + case cancelled + case failed(String) + } + + private func writeResponse( + _ response: HTTPResponse, + operation: MCPHTTPNetworkResourceOwner.RequestOperation, + version: HTTPVersion, + context: ChannelHandlerContext, + heartbeatInterval: Duration? + ) async -> WriterCompletion { let eventLoop = context.eventLoop - let response = trackedResponse.response let status = HTTPResponseStatus(statusCode: response.statusCode) let headers = response.headers + var head = HTTPResponseHead(version: version, status: status) + for (name, value) in headers { head.headers.add(name: name, value: value) } switch response { case .stream(let stream, _): - let ownership: MCPHTTPActiveStreamOwnership - switch trackedResponse.streamLifetime { - case .finite(let operation): - guard operation.claimWriter() else { - return - } - ownership = .finite(operation) - case .open(let completion): - guard let streamReceipt = networkResources.registerTask( - kind: .streamWriter, - child: childRegistration - ) else { - completion.finish() - return - } - ownership = .open(MCPHTTPStreamOwnership( - receipt: streamReceipt, - completion: completion - )) - case nil: - return - } - let registration = eventLoop.makePromise(of: Void.self) - eventLoop.execute { - guard context.channel.isActive else { - ownership.completion?.finish() - ownership.finishTask() - registration.succeed(()) - return - } - let streamTask = Task { - var head = HTTPResponseHead(version: version, status: status) - for (name, value) in headers { - head.headers.add(name: name, value: value) + do { + try await writeResponsePart(.head(head), context: context, eventLoop: eventLoop) + let events = AsyncStream.makeStream(bufferingPolicy: .unbounded) + await withTaskGroup(of: Void.self) { group in + group.addTask { + do { + for try await chunk in stream { + try Task.checkCancellation() + events.continuation.yield(.body(chunk)) + } + events.continuation.yield(.sourceFinished) + } catch is CancellationError { + events.continuation.yield(.sourceFinished) + } catch { + events.continuation.yield(.sourceFailed(error.localizedDescription)) + } } - - var iterator = stream.makeAsyncIterator() - do { - try Task.checkCancellation() - try await self.writeResponsePart( - .head(head), - context: context, - eventLoop: eventLoop - ) - while let chunk = try await iterator.next() { - try Task.checkCancellation() - try await self.writeResponseBody( - chunk, - context: context, - eventLoop: eventLoop - ) + if let heartbeatInterval { + group.addTask { + while Task.isCancelled == false { + do { try await Task.sleep(for: heartbeatInterval) } + catch { return } + guard Task.isCancelled == false else { return } + events.continuation.yield(.heartbeat) + } } - } catch is CancellationError { - } catch { - logger.error("MCP SSE stream failed: \(error.localizedDescription, privacy: .public)") } - if Task.isCancelled == false { - try? await self.writeResponsePart( - .end(nil), - context: context, - eventLoop: eventLoop - ) + eventLoopLoop: for await event in events.stream { + switch event { + case .body(let data): + do { try await writeResponseBody(data, context: context, eventLoop: eventLoop) } + catch { + operation.beginClosing(.transportFailure(error.localizedDescription)) + break eventLoopLoop + } + case .heartbeat: + do { + try await writeResponseBody( + Data(": keep-alive\n\n".utf8), + context: context, + eventLoop: eventLoop + ) + } catch { + operation.beginClosing(.transportFailure(error.localizedDescription)) + break eventLoopLoop + } + case .sourceFinished: + break eventLoopLoop + case .sourceFailed(let message): + logger.error("MCP SSE stream failed: \(message, privacy: .public)") + operation.beginClosing(.transportFailure(message)) + break eventLoopLoop + } } - await self.finishStreamTask( - ownership, - context: context, - eventLoop: eventLoop - ) + group.cancelAll() + events.continuation.finish() } - context.channel.closeFuture.whenComplete { _ in - ownership.terminateFromChannel() + try await writeResponsePart(.end(nil), context: context, eventLoop: eventLoop) + return .responded + } catch is CancellationError { + do { + try await writeResponsePart(.end(nil), context: context, eventLoop: eventLoop) + return .responded + } catch { + return .cancelled } - self.activeStreamOwnership?.terminateFromChannel() - self.activeStreamOwnership = ownership - ownership.install(streamTask) - context.read() - registration.succeed(()) + } catch { + return .failed(error.localizedDescription) } - try? await registration.futureResult.get() default: let body = response.bodyData - var head = HTTPResponseHead(version: version, status: status) - for (name, value) in headers { - head.headers.add(name: name, value: value) - } - if let body { - head.headers.add(name: "Content-Length", value: "\(body.count)") - } + if let body { head.headers.add(name: "Content-Length", value: "\(body.count)") } do { - try await writeResponsePart( - .head(head), - context: context, - eventLoop: eventLoop - ) - if let body { - try await writeResponseBody( - body, - context: context, - eventLoop: eventLoop - ) - } - try await writeResponsePart( - .end(nil), - context: context, - eventLoop: eventLoop - ) + try await writeResponsePart(.head(head), context: context, eventLoop: eventLoop) + if let body { try await writeResponseBody(body, context: context, eventLoop: eventLoop) } + try await writeResponsePart(.end(nil), context: context, eventLoop: eventLoop) + return .responded } catch { - logger.debug( - "MCP HTTP response ended during connection shutdown: \(error.localizedDescription, privacy: .public)" - ) - } - } - } - - private func finishStreamTask( - _ ownership: MCPHTTPActiveStreamOwnership, - context: ChannelHandlerContext, - eventLoop: any EventLoop - ) async { - if ownership.claimTaskTermination() { - let completion = eventLoop.makePromise(of: Void.self) - eventLoop.execute { - if self.activeStreamOwnership?.id == ownership.id { - self.activeStreamOwnership = nil - } - ownership.completion?.finish() - completion.succeed(()) + return .failed(error.localizedDescription) } - try? await completion.futureResult.get() } - ownership.finishTask() } private func writeResponsePart( diff --git a/Sources/CodexReviewMCPServer/CodexReviewMCPProtocolServer.swift b/Sources/CodexReviewMCPServer/CodexReviewMCPProtocolServer.swift index 34db427e..b4ce7f25 100644 --- a/Sources/CodexReviewMCPServer/CodexReviewMCPProtocolServer.swift +++ b/Sources/CodexReviewMCPServer/CodexReviewMCPProtocolServer.swift @@ -1,7 +1,80 @@ import Foundation import MCP +import OSLog import CodexReview +private let protocolServerLogger = Logger( + subsystem: "CodexReviewKit", + category: "mcp-protocol" +) + +final class MCPProtocolServerForeignLifetimeWaiter: @unchecked Sendable { + final class Lease: @unchecked Sendable { + private let state: State + + fileprivate init(state: State) { + self.state = state + } + + deinit { + state.finish() + } + } + + fileprivate final class State: @unchecked Sendable { + private let lock = NSLock() + private var didFinish = false + private var waiters: [CheckedContinuation] = [] + + func finish() { + let waiters: [CheckedContinuation] + lock.lock() + guard didFinish == false else { + lock.unlock() + return + } + didFinish = true + waiters = self.waiters + self.waiters.removeAll(keepingCapacity: false) + lock.unlock() + for waiter in waiters { + waiter.resume() + } + } + + func wait() async { + await withCheckedContinuation { continuation in + lock.lock() + if didFinish { + lock.unlock() + continuation.resume() + } else { + waiters.append(continuation) + lock.unlock() + } + } + } + } + + private let state: State + + private init(state: State) { + self.state = state + } + + static func makePair() -> (lease: Lease, waiter: MCPProtocolServerForeignLifetimeWaiter) { + let state = State() + return ( + Lease(state: state), + MCPProtocolServerForeignLifetimeWaiter(state: state) + ) + } + + func wait() async { + await state.wait() + } +} + package actor MCPClientSessionState { private var clientInfo: Client.Info? @@ -35,7 +108,8 @@ func makeMCPProtocolServer( defaultSessionID: String? = nil, clientSession: MCPClientSessionState = .init(), boundedReviewWaitDuration: Duration = .seconds(540), - networkResources: MCPHTTPNetworkResourceOwner + networkResources: MCPHTTPNetworkResourceOwner, + foreignLifetimeLease: MCPProtocolServerForeignLifetimeWaiter.Lease ) async -> Server { let server = Server( name: "codex_review", @@ -46,8 +120,11 @@ func makeMCPProtocolServer( ) ) - await server.withMethodHandler(ListTools.self) { _ in - try await networkResources.performTask(kind: .domainHandler) { + await server.withMethodHandler(ListTools.self) { [foreignLifetimeLease] _ in + try await performMCPDomainWork( + networkResources: networkResources, + foreignLifetimeLease: foreignLifetimeLease + ) { let tools = await adapter.tools.map { descriptor in Tool( name: descriptor.name.rawValue, @@ -59,8 +136,11 @@ func makeMCPProtocolServer( } } - await server.withMethodHandler(CallTool.self) { params in - try await networkResources.performTask(kind: .domainHandler) { + await server.withMethodHandler(CallTool.self) { [foreignLifetimeLease] params in + try await performMCPDomainWork( + networkResources: networkResources, + foreignLifetimeLease: foreignLifetimeLease + ) { guard let tool = CodexReviewMCP.Tool.Name(rawValue: params.name) else { return .init( content: [.text(text: "Unknown tool: \(params.name)", annotations: nil, _meta: nil)], @@ -80,6 +160,8 @@ func makeMCPProtocolServer( ) let response = try await adapter.handle(request) return try toolResult(tool: tool, response: response) + } catch is CancellationError { + throw CancellationError() } catch { return .init( content: [.text(text: error.localizedDescription, annotations: nil, _meta: nil)], @@ -89,22 +171,31 @@ func makeMCPProtocolServer( } } - await server.withMethodHandler(ListResources.self) { _ in - try await networkResources.performTask(kind: .domainHandler) { + await server.withMethodHandler(ListResources.self) { [foreignLifetimeLease] _ in + try await performMCPDomainWork( + networkResources: networkResources, + foreignLifetimeLease: foreignLifetimeLease + ) { .init(resources: helpResources.map(\.resource)) } } - await server.withMethodHandler(ReadResource.self) { params in - try await networkResources.performTask(kind: .domainHandler) { + await server.withMethodHandler(ReadResource.self) { [foreignLifetimeLease] params in + try await performMCPDomainWork( + networkResources: networkResources, + foreignLifetimeLease: foreignLifetimeLease + ) { let content = helpResources.first { $0.uri == params.uri }?.content ?? "Resource not found: \(params.uri)" return .init(contents: [.text(content, uri: params.uri, mimeType: "text/markdown")]) } } - await server.withMethodHandler(ListResourceTemplates.self) { _ in - try await networkResources.performTask(kind: .domainHandler) { + await server.withMethodHandler(ListResourceTemplates.self) { [foreignLifetimeLease] _ in + try await performMCPDomainWork( + networkResources: networkResources, + foreignLifetimeLease: foreignLifetimeLease + ) { .init(templates: helpResourceTemplates) } } @@ -112,6 +203,92 @@ func makeMCPProtocolServer( return server } +private func performMCPDomainWork( + networkResources: MCPHTTPNetworkResourceOwner, + foreignLifetimeLease: MCPProtocolServerForeignLifetimeWaiter.Lease, + operation: @escaping @Sendable () async throws -> Success +) async throws -> Success { + _ = foreignLifetimeLease + let requestOperation = try resolveMCPRequestOperation( + networkResources: networkResources + ) + guard let reservation = requestOperation.admitDomainWork() else { + return try unavailableMCPDomainWork( + networkResources: networkResources, + reason: "The HTTP request operation did not admit domain work." + ) + } + + let task = Task { + let value = try await operation() + try Task.checkCancellation() + return value + } + reservation.install(task) + + do { + let value = try await withTaskCancellationHandler { + let value = try await task.value + try Task.checkCancellation() + return value + } onCancel: { + requestOperation.beginClosing(.sdkCancellation) + task.cancel() + } + reservation.acknowledge() + return value + } catch is CancellationError { + requestOperation.beginClosing(.sdkCancellation) + reservation.acknowledge() + throw CancellationError() + } catch { + reservation.acknowledge(.failed(error.localizedDescription)) + throw error + } +} + +private func resolveMCPRequestOperation( + networkResources: MCPHTTPNetworkResourceOwner +) throws -> MCPHTTPNetworkResourceOwner.RequestOperation { + guard let encodedToken = Server.currentHandlerContext? + .httpContext? + .header(MCPHTTPNetworkResourceOwner.operationTokenHeaderName) + else { + return try unavailableMCPDomainWork( + networkResources: networkResources, + reason: "The HTTP request did not contain an operation token." + ) + } + guard let token = MCPHTTPNetworkResourceOwner.OperationToken( + headerValue: encodedToken + ) else { + return try unavailableMCPDomainWork( + networkResources: networkResources, + reason: "The HTTP request contained a malformed operation token." + ) + } + guard let requestOperation = networkResources.resolve(token) else { + return try unavailableMCPDomainWork( + networkResources: networkResources, + reason: "The HTTP request operation token did not resolve in its generation." + ) + } + return requestOperation +} + +private func unavailableMCPDomainWork( + networkResources: MCPHTTPNetworkResourceOwner, + reason: String +) throws -> Success { + switch networkResources.snapshot().phase { + case .accepting: + protocolServerLogger.error("\(reason, privacy: .public)") + throw MCPError.internalError("MCP request ownership could not be resolved.") + case .admissionClosed, .closing, .closed: + throw CancellationError() + } +} + private func schema(for tool: CodexReviewMCP.Tool.Name) -> Value { switch tool { case .reviewStart: diff --git a/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift b/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift index d044cf98..71f304a9 100644 --- a/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift +++ b/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift @@ -2,713 +2,1233 @@ import Foundation @preconcurrency import NIOCore final class MCPHTTPNetworkResourceOwner: @unchecked Sendable { - struct TaskAdmissionClosed: Error, Sendable { - let kind: TaskKind + static let operationTokenHeaderName = "X-CodexReview-Request-Operation" + + struct OperationToken: Hashable, Sendable { + fileprivate let connectionID: UUID + fileprivate let operationID: UUID + + var headerValue: String { + "\(connectionID.uuidString.lowercased()):\(operationID.uuidString.lowercased())" + } + + init?(headerValue: String) { + let parts = headerValue.split(separator: ":", omittingEmptySubsequences: false) + guard parts.count == 2, + let connectionID = UUID(uuidString: String(parts[0])), + let operationID = UUID(uuidString: String(parts[1])) + else { + return nil + } + self.connectionID = connectionID + self.operationID = operationID + } + + fileprivate init(connectionID: UUID, operationID: UUID) { + self.connectionID = connectionID + self.operationID = operationID + } } - enum TaskKind: Hashable, Sendable { - case domainHandler - case response - case finiteResponse - case streamBridge - case streamHeartbeat - case streamWriter - case streamCompletion + struct RequestMetadata: Equatable, Sendable { + let method: String + let path: String + let jsonRPCID: String? } - struct ChildRegistration: Hashable, Sendable { - fileprivate let id: UUID + enum CloseCause: Equatable, Sendable { + case sdkCancellation + case peerClosed + case sessionClosed + case serverStop + case transportFailure(String) } - final class TaskReceipt: @unchecked Sendable { - fileprivate let id: UUID - fileprivate let kind: TaskKind - fileprivate let childID: UUID? - private weak var owner: MCPHTTPNetworkResourceOwner? - private let lock = NSLock() - private var cancelTask: (@Sendable () -> Void)? - private var cancellationWasRequested = false - private var didFinish = false + enum OperationResult: Equatable, Sendable { + case responded + case cancelled(CloseCause) + case failed(String) + } - fileprivate init( - id: UUID, - kind: TaskKind, - childID: UUID?, - owner: MCPHTTPNetworkResourceOwner - ) { - self.id = id + enum WorkKind: CaseIterable, Sendable { + case httpHandler + case domain + case source + case writer + } + + enum WorkCompletion: Sendable { + case succeeded + case failed(String) + } + + enum WorkStateSnapshot: Equatable, Sendable { + case notAdmitted + case responseReady + case turnGranted + case reserved + case running + case completed + case closed + } + + enum OperationPhaseSnapshot: Equatable, Sendable { + case admitted + case handling + case responding + case closing(CloseCause?, OperationResult) + case closed(OperationResult) + } + + struct OperationSnapshot: Equatable, Sendable { + let token: OperationToken + let admissionOrdinal: UInt64 + let metadata: RequestMetadata + let phase: OperationPhaseSnapshot + let boundSessionID: String? + let httpHandler: WorkStateSnapshot + let domain: WorkStateSnapshot + let source: WorkStateSnapshot + let writer: WorkStateSnapshot + + var responseIsReady: Bool { + switch writer { + case .responseReady, .turnGranted, .reserved, .running, .completed: + true + case .notAdmitted, .closed: + false + } + } + + var writerIsRunning: Bool { writer == .running } + + var domainWorkIsPending: Bool { + domain == .reserved || domain == .running + } + + var terminalCause: CloseCause? { + switch phase { + case .closing(let cause, _): + cause + case .closed(.cancelled(let cause)): + cause + case .admitted, .handling, .responding, .closed: + nil + } + } + } + + enum ConnectionPhaseSnapshot: Equatable, Sendable { + case accepting + case admissionClosed + case closing + case closed + } + + struct ConnectionSnapshot: Equatable, Sendable { + let id: UUID + let ordinal: UInt64 + let phase: ConnectionPhaseSnapshot + let operations: [OperationSnapshot] + } + + enum GenerationPhaseSnapshot: Equatable, Sendable { + case accepting + case admissionClosed + case closing + case closed + } + + struct Snapshot: Equatable, Sendable { + let revision: UInt64 + let phase: GenerationPhaseSnapshot + let connections: [ConnectionSnapshot] + + var isQuiescent: Bool { + phase == .closed && connections.isEmpty + } + } + + final class WorkReservation: @unchecked Sendable { + fileprivate let id = UUID() + fileprivate let kind: WorkKind + private weak var operation: RequestOperation? + + fileprivate init(kind: WorkKind, operation: RequestOperation) { self.kind = kind - self.childID = childID - self.owner = owner + self.operation = operation } func install( _ task: Task ) { - installCancellation { task.cancel() } + operation?.installCancellation({ task.cancel() }, for: self) } - fileprivate func installCancellation( - _ cancelTask: @escaping @Sendable () -> Void - ) { - let shouldCancel: Bool - lock.lock() - if didFinish { - shouldCancel = false - } else { - self.cancelTask = cancelTask - shouldCancel = cancellationWasRequested + func acknowledge(_ completion: WorkCompletion = .succeeded) { + operation?.acknowledge(self, completion) + } + } + + final class RequestOperation: @unchecked Sendable { + private struct Common { + let metadata: RequestMetadata + var boundSessionID: String? + } + + private enum WorkSlot { + case notAdmitted + case responseReady(CheckedContinuation?) + case turnGranted + case reserved(UUID) + case running(UUID, @Sendable () -> Void) + case completed + case closed + + var snapshot: WorkStateSnapshot { + switch self { + case .notAdmitted: .notAdmitted + case .responseReady: .responseReady + case .turnGranted: .turnGranted + case .reserved: .reserved + case .running: .running + case .completed: .completed + case .closed: .closed + } } - lock.unlock() - if shouldCancel { - cancelTask() + + var isPending: Bool { + switch self { + case .reserved, .running: + true + case .notAdmitted, .responseReady, .turnGranted, .completed, .closed: + false + } } } - func cancel() { - let cancelTask: (@Sendable () -> Void)? - lock.lock() - cancellationWasRequested = true - cancelTask = self.cancelTask - lock.unlock() - cancelTask?() - } + private struct Slots { + var httpHandler: WorkSlot = .notAdmitted + var domain: WorkSlot = .notAdmitted + var source: WorkSlot = .notAdmitted + var writer: WorkSlot = .notAdmitted - func finish() { - let owner: MCPHTTPNetworkResourceOwner? - lock.lock() - guard didFinish == false else { - lock.unlock() - return + subscript(kind: WorkKind) -> WorkSlot { + get { + switch kind { + case .httpHandler: httpHandler + case .domain: domain + case .source: source + case .writer: writer + } + } + set { + switch kind { + case .httpHandler: httpHandler = newValue + case .domain: domain = newValue + case .source: source = newValue + case .writer: writer = newValue + } + } + } + + var handlingIsPending: Bool { + httpHandler.isPending || domain.isPending } - didFinish = true - cancelTask = nil - owner = self.owner - lock.unlock() - owner?.finishTask(id: id, kind: kind) } - } - final class FiniteResponseOperation: @unchecked Sendable { - private enum WriterDisposition: Equatable { - case unresolved - case required - case waived + private struct Admitted { + var common: Common + var slots: Slots } - private enum WriterTerminationOwner: Equatable { - case active - case channel - case task + private struct Handling { + var common: Common + var slots: Slots } - let id = UUID() - private let receipt: TaskReceipt - private let lock = NSLock() - private var sourceCancellation: (@Sendable () -> Void)? - private var writerCancellation: (@Sendable () -> Void)? - private var sourceFinished = false - private var writerFinished = false - private var writerDisposition = WriterDisposition.unresolved - private var writerTerminationOwner = WriterTerminationOwner.active - private var cancellationWasRequested = false - private var writerCancellationWasRequested = false - private var didFinish = false + private struct Responding { + var common: Common + var slots: Slots + } - fileprivate init(receipt: TaskReceipt) { - self.receipt = receipt + private struct Closing { + var common: Common + var slots: Slots + let cause: CloseCause? + var pending: Set + let terminalResult: OperationResult } - fileprivate func installOwnerCancellation() { - receipt.installCancellation { [self] in - cancelAll() - } + private enum State { + case admitted(Admitted) + case handling(Handling) + case responding(Responding) + case closing(Closing) + case closed(Common, OperationResult, Slots) } - func installSource(_ task: Task) { - let shouldCancel: Bool + let token: OperationToken + let admissionOrdinal: UInt64 + let metadata: RequestMetadata + private weak var connection: Connection? + private let lock = NSLock() + private var state: State + private var closeWaiters: [CheckedContinuation] = [] + private var handlingWaiters: [CheckedContinuation] = [] + + fileprivate init( + token: OperationToken, + admissionOrdinal: UInt64, + metadata: RequestMetadata, + connection: Connection + ) { + self.token = token + self.admissionOrdinal = admissionOrdinal + self.metadata = metadata + self.connection = connection + state = .admitted(.init(common: .init(metadata: metadata), slots: .init())) + } + + func beginHTTPHandling() -> WorkReservation? { lock.lock() - if sourceFinished { - shouldCancel = false - } else { - sourceCancellation = { task.cancel() } - shouldCancel = cancellationWasRequested + guard case .admitted(var admitted) = state, + case .notAdmitted = admitted.slots.httpHandler + else { + lock.unlock() + return nil } + let reservation = WorkReservation(kind: .httpHandler, operation: self) + admitted.slots.httpHandler = .reserved(reservation.id) + state = .handling(.init(common: admitted.common, slots: admitted.slots)) lock.unlock() - if shouldCancel { - task.cancel() - } + notifyChanged() + return reservation } - func finishSource() { - let shouldFinish: Bool + func admitDomainWork() -> WorkReservation? { + let reservation: WorkReservation lock.lock() - guard sourceFinished == false else { + switch state { + case .handling(var handling): + guard case .notAdmitted = handling.slots.domain else { + lock.unlock() + return nil + } + reservation = WorkReservation(kind: .domain, operation: self) + handling.slots.domain = .reserved(reservation.id) + state = .handling(handling) + case .responding(var responding): + guard case .notAdmitted = responding.slots.domain else { + lock.unlock() + return nil + } + reservation = WorkReservation(kind: .domain, operation: self) + responding.slots.domain = .reserved(reservation.id) + state = .responding(responding) + case .admitted, .closing, .closed: lock.unlock() - return + return nil } - sourceFinished = true - sourceCancellation = nil - shouldFinish = claimCompletionIfReadyLocked() lock.unlock() - if shouldFinish { - receipt.finish() - } + notifyChanged() + return reservation } - func claimWriter() -> Bool { + func beginResponding() -> Bool { lock.lock() - guard writerDisposition == .unresolved else { + guard case .handling(let handling) = state else { lock.unlock() return false } - writerDisposition = .required + state = .responding(.init(common: handling.common, slots: handling.slots)) lock.unlock() + notifyChanged() return true } - func waiveWriter() { - let shouldFinish: Bool + func bindResponseSource() -> WorkReservation? { + reserve(.source) + } + + func markResponseSourceNotRequired() { lock.lock() - guard writerDisposition == .unresolved else { + guard case .responding(var responding) = state, + case .notAdmitted = responding.slots.source + else { lock.unlock() return } - writerDisposition = .waived - shouldFinish = claimCompletionIfReadyLocked() + responding.slots.source = .closed + state = .responding(responding) lock.unlock() - if shouldFinish { - receipt.finish() - } + notifyChanged() } - func installWriter(_ task: Task) { - let shouldCancel: Bool + fileprivate func markResponseReady() -> Bool { lock.lock() - if writerFinished { - shouldCancel = false - } else { - writerCancellation = { task.cancel() } - shouldCancel = cancellationWasRequested - || writerCancellationWasRequested + guard case .responding(var responding) = state, + case .notAdmitted = responding.slots.writer + else { + lock.unlock() + return false } + responding.slots.writer = .responseReady(nil) + state = .responding(responding) lock.unlock() - if shouldCancel { - task.cancel() + notifyChanged() + return true + } + + fileprivate func waitForWriterTurn() async -> Bool { + await withCheckedContinuation { continuation in + lock.lock() + guard case .responding(var responding) = state else { + lock.unlock() + continuation.resume(returning: false) + return + } + switch responding.slots.writer { + case .turnGranted: + lock.unlock() + continuation.resume(returning: true) + case .responseReady(nil): + responding.slots.writer = .responseReady(continuation) + state = .responding(responding) + lock.unlock() + default: + lock.unlock() + continuation.resume(returning: false) + } } } - func terminateWriterFromChannel() { - let cancelWriter: (@Sendable () -> Void)? + fileprivate func grantWriterTurn() -> Bool { + let waiter: CheckedContinuation? lock.lock() - guard writerTerminationOwner == .active else { + guard case .responding(var responding) = state else { lock.unlock() - return + return false + } + switch responding.slots.writer { + case .responseReady(let continuation): + waiter = continuation + responding.slots.writer = .turnGranted + state = .responding(responding) + lock.unlock() + waiter?.resume(returning: true) + notifyChanged() + return true + default: + lock.unlock() + return false } - writerTerminationOwner = .channel - writerCancellationWasRequested = true - cancelWriter = writerCancellation - lock.unlock() - cancelWriter?() } - func claimWriterTaskTermination() -> Bool { + fileprivate var isResponseReady: Bool { lock.lock() - guard writerTerminationOwner == .active else { - lock.unlock() - return false + let result: Bool + if case .responding(let responding) = state { + switch responding.slots.writer { + case .responseReady, .turnGranted, .reserved, .running: + result = true + case .notAdmitted, .completed, .closed: + result = false + } + } else { + result = false } - writerTerminationOwner = .task lock.unlock() - return true + return result } - func finishWriter() { - let shouldFinish: Bool + func bindWriter() -> WorkReservation? { lock.lock() - guard writerFinished == false else { + guard case .responding(var responding) = state, + case .turnGranted = responding.slots.writer + else { lock.unlock() - return + return nil } - writerFinished = true - writerCancellation = nil - shouldFinish = claimCompletionIfReadyLocked() + let reservation = WorkReservation(kind: .writer, operation: self) + responding.slots.writer = .reserved(reservation.id) + state = .responding(responding) lock.unlock() - if shouldFinish { - receipt.finish() - } + notifyChanged() + return reservation } - private func cancelAll() { - let cancelSource: (@Sendable () -> Void)? - let cancelWriter: (@Sendable () -> Void)? + func bindSession(_ sessionID: String) { lock.lock() - cancellationWasRequested = true - cancelSource = sourceCancellation - cancelWriter = writerCancellation + switch state { + case .admitted(var admitted): + if admitted.common.boundSessionID == nil { + admitted.common.boundSessionID = sessionID + state = .admitted(admitted) + } + case .handling(var handling): + if handling.common.boundSessionID == nil { + handling.common.boundSessionID = sessionID + state = .handling(handling) + } + case .responding(var responding): + if responding.common.boundSessionID == nil { + responding.common.boundSessionID = sessionID + state = .responding(responding) + } + case .closing(var closing): + if closing.common.boundSessionID == nil { + closing.common.boundSessionID = sessionID + state = .closing(closing) + } + case .closed: + break + } lock.unlock() - cancelSource?() - cancelWriter?() + notifyChanged() } - private func claimCompletionIfReadyLocked() -> Bool { - guard didFinish == false, sourceFinished else { - return false - } - switch writerDisposition { - case .unresolved: - return false - case .required where writerFinished == false: - return false - case .required, .waived: - didFinish = true - return true - } + func beginClosing(_ cause: CloseCause) { + transitionToClosing(cause: cause, result: .cancelled(cause), cancelPending: true) } - } - private final class ChildResource: @unchecked Sendable { - let registration: ChildRegistration - let channel: any Channel - - init(registration: ChildRegistration, channel: any Channel) { - self.registration = registration - self.channel = channel + func acknowledgeResponseEnd() { + transitionToClosing(cause: nil, result: .responded, cancelPending: false) } - } - private struct TaskCountWaiter { - let kind: TaskKind - let targetCount: Int - let continuation: CheckedContinuation - } - - private struct ChildCountWaiter { - let targetCount: Int - let continuation: CheckedContinuation - } + func waitUntilClosed() async -> OperationResult { + await withCheckedContinuation { continuation in + lock.lock() + if case .closed(_, let result, _) = state { + lock.unlock() + continuation.resume(returning: result) + } else { + closeWaiters.append(continuation) + lock.unlock() + } + } + } - private let lock = NSLock() - private var acceptsChildren = false - private var acceptsTasks = false - private var closedTaskAdmissionKinds: Set = [] - private var children: [UUID: ChildResource] = [:] - private var tasks: [UUID: TaskReceipt] = [:] - private var childDrainWaiters: [CheckedContinuation] = [] - private var taskDrainWaiters: [CheckedContinuation] = [] - private var taskKindDrainWaiters: [ - TaskKind: [CheckedContinuation] - ] = [:] - private var totalTaskCounts: [TaskKind: Int] = [:] - private var taskCountWaiters: [UUID: TaskCountWaiter] = [:] - private var totalChildCount = 0 - private var childCountWaiters: [UUID: ChildCountWaiter] = [:] - private var heldTaskCompletionKind: TaskKind? - private var heldTaskCompletionIDs: Set = [] - private var taskCompletionHoldWaiters: [CheckedContinuation] = [] - private var shouldHoldNextChildCloseAcknowledgement = false - private var heldChildCloseAcknowledgementIDs: Set = [] - private var childCloseHoldWaiters: [CheckedContinuation] = [] - - func open() { - lock.lock() - precondition( - children.isEmpty && tasks.isEmpty, - "MCP network generation must drain resources before reopening." - ) - acceptsChildren = true - acceptsTasks = true - closedTaskAdmissionKinds.removeAll(keepingCapacity: true) - lock.unlock() - } + fileprivate func waitForHandlingWorkToDrain() async { + await withCheckedContinuation { continuation in + lock.lock() + if handlingIsPendingLocked() == false { + lock.unlock() + continuation.resume() + } else { + handlingWaiters.append(continuation) + lock.unlock() + } + } + } - func registerChild(_ channel: any Channel) -> ChildRegistration? { - let registration: ChildRegistration - let completedWaiters: [CheckedContinuation] - lock.lock() - guard acceptsChildren else { + fileprivate func closeDomainAdmission() { + var waiters: [CheckedContinuation] = [] + lock.lock() + switch state { + case .handling(var handling): + if case .notAdmitted = handling.slots.domain { + handling.slots.domain = .closed + state = .handling(handling) + } + case .responding(var responding): + if case .notAdmitted = responding.slots.domain { + responding.slots.domain = .closed + state = .responding(responding) + } + case .admitted, .closing, .closed: + break + } + if handlingIsPendingLocked() == false { + waiters = handlingWaiters + handlingWaiters.removeAll(keepingCapacity: false) + } lock.unlock() - return nil + for waiter in waiters { waiter.resume() } + notifyChanged() } - registration = .init(id: UUID()) - children[registration.id] = ChildResource( - registration: registration, - channel: channel - ) - totalChildCount += 1 - let completedIDs = childCountWaiters.compactMap { id, waiter in - totalChildCount >= waiter.targetCount ? id : nil - } - completedWaiters = completedIDs.compactMap { - childCountWaiters.removeValue(forKey: $0)?.continuation - } - lock.unlock() - for waiter in completedWaiters { - waiter.resume() + + fileprivate func installCancellation( + _ cancellation: @escaping @Sendable () -> Void, + for reservation: WorkReservation + ) { + var shouldCancel = false + lock.lock() + switch state { + case .handling(var handling): + _ = Self.install(cancellation, reservation: reservation, slots: &handling.slots) + state = .handling(handling) + case .responding(var responding): + _ = Self.install(cancellation, reservation: reservation, slots: &responding.slots) + state = .responding(responding) + case .closing(var closing): + if Self.install(cancellation, reservation: reservation, slots: &closing.slots) { + shouldCancel = true + } + state = .closing(closing) + case .admitted, .closed: + break + } + lock.unlock() + if shouldCancel { cancellation() } + notifyChanged() } - channel.closeFuture.whenComplete { [weak self] _ in - self?.acknowledgeChildClose(registration) + + fileprivate func acknowledge( + _ reservation: WorkReservation, + _ completion: WorkCompletion + ) { + var closed: (OperationResult, [CheckedContinuation])? + var handlingWaitersToResume: [CheckedContinuation] = [] + var failure: String? + lock.lock() + switch state { + case .handling(var handling): + if Self.complete(reservation, slots: &handling.slots) { + state = .handling(handling) + if case .failed(let message) = completion, + reservation.kind == .source || reservation.kind == .writer + { failure = message } + } + case .responding(var responding): + if Self.complete(reservation, slots: &responding.slots) { + state = .responding(responding) + if case .failed(let message) = completion, + reservation.kind == .source || reservation.kind == .writer + { failure = message } + } + case .closing(var closing): + if Self.complete(reservation, slots: &closing.slots) { + closing.pending.remove(reservation.kind) + if closing.pending.isEmpty { + let result = closing.terminalResult + let waiters = closeWaiters + closeWaiters.removeAll(keepingCapacity: false) + state = .closed(closing.common, result, closing.slots) + closed = (result, waiters) + } else { + state = .closing(closing) + } + } + case .admitted, .closed: + break + } + if handlingIsPendingLocked() == false { + handlingWaitersToResume = handlingWaiters + handlingWaiters.removeAll(keepingCapacity: false) + } + lock.unlock() + for waiter in handlingWaitersToResume { waiter.resume() } + if let failure { + transitionToClosing(cause: nil, result: .failed(failure), cancelPending: true) + return + } + finishClosed(closed) + notifyChanged() } - return registration - } - func closeChildAdmission() { - lock.lock() - acceptsChildren = false - lock.unlock() - } + func snapshot() -> OperationSnapshot { + lock.lock() + let common: Common + let slots: Slots + let phase: OperationPhaseSnapshot + switch state { + case .admitted(let admitted): + common = admitted.common; slots = admitted.slots; phase = .admitted + case .handling(let handling): + common = handling.common; slots = handling.slots; phase = .handling + case .responding(let responding): + common = responding.common; slots = responding.slots; phase = .responding + case .closing(let closing): + common = closing.common; slots = closing.slots + phase = .closing(closing.cause, closing.terminalResult) + case .closed(let closedCommon, let result, let closedSlots): + common = closedCommon; slots = closedSlots; phase = .closed(result) + } + let result = OperationSnapshot( + token: token, + admissionOrdinal: admissionOrdinal, + metadata: common.metadata, + phase: phase, + boundSessionID: common.boundSessionID, + httpHandler: slots.httpHandler.snapshot, + domain: slots.domain.snapshot, + source: slots.source.snapshot, + writer: slots.writer.snapshot + ) + lock.unlock() + return result + } - func closeAndDrainChildren() async { - let channels = closeChildAdmissionAndSnapshot() - for channel in channels { - channel.close(mode: .all, promise: nil) + private func reserve(_ kind: WorkKind) -> WorkReservation? { + lock.lock() + guard case .responding(var responding) = state, + case .notAdmitted = responding.slots[kind] + else { + lock.unlock() + return nil + } + let reservation = WorkReservation(kind: kind, operation: self) + responding.slots[kind] = .reserved(reservation.id) + state = .responding(responding) + lock.unlock() + notifyChanged() + return reservation } - await withCheckedContinuation { continuation in + + private func transitionToClosing( + cause: CloseCause?, + result: OperationResult, + cancelPending: Bool + ) { + var cancellations: [@Sendable () -> Void] = [] + var writerWaiter: CheckedContinuation? + var closed: (OperationResult, [CheckedContinuation])? + var handlingWaitersToResume: [CheckedContinuation] = [] lock.lock() - if children.isEmpty { + let common: Common + var slots: Slots + switch state { + case .admitted(let admitted): common = admitted.common; slots = admitted.slots + case .handling(let handling): common = handling.common; slots = handling.slots + case .responding(let responding): common = responding.common; slots = responding.slots + case .closing, .closed: lock.unlock() - continuation.resume() + return + } + if case .responseReady(let waiter) = slots.writer { writerWaiter = waiter } + var pending: Set = [] + for kind in WorkKind.allCases { + switch slots[kind] { + case .reserved: + pending.insert(kind) + case .running(_, let cancel): + pending.insert(kind) + if cancelPending { cancellations.append(cancel) } + case .notAdmitted, .responseReady, .turnGranted: + slots[kind] = .closed + case .completed, .closed: + break + } + } + if pending.isEmpty { + let waiters = closeWaiters + closeWaiters.removeAll(keepingCapacity: false) + state = .closed(common, result, slots) + closed = (result, waiters) } else { - childDrainWaiters.append(continuation) - lock.unlock() + state = .closing(.init( + common: common, + slots: slots, + cause: cause, + pending: pending, + terminalResult: result + )) } + handlingWaitersToResume = handlingWaiters + handlingWaiters.removeAll(keepingCapacity: false) + lock.unlock() + writerWaiter?.resume(returning: false) + for cancellation in cancellations { cancellation() } + for waiter in handlingWaitersToResume { waiter.resume() } + finishClosed(closed) + notifyChanged() } - } - - private func closeChildAdmissionAndSnapshot() -> [any Channel] { - lock.lock() - acceptsChildren = false - let channels = children.values.map(\.channel) - lock.unlock() - return channels - } - func registerTask( - kind: TaskKind, - child: ChildRegistration? = nil - ) -> TaskReceipt? { - let receipt: TaskReceipt - let completedWaiters: [CheckedContinuation] - lock.lock() - guard acceptsTasks, closedTaskAdmissionKinds.contains(kind) == false else { - lock.unlock() - return nil + private func finishClosed( + _ closed: (OperationResult, [CheckedContinuation])? + ) { + guard let (result, waiters) = closed else { return } + for waiter in waiters { waiter.resume(returning: result) } + connection?.operationDidClose(self) } - receipt = TaskReceipt( - id: UUID(), - kind: kind, - childID: child?.id, - owner: self - ) - tasks[receipt.id] = receipt - totalTaskCounts[kind, default: 0] += 1 - let completedIDs = taskCountWaiters.compactMap { id, waiter in - waiter.kind == kind && totalTaskCounts[kind, default: 0] >= waiter.targetCount - ? id - : nil + + private func notifyChanged() { + connection?.operationDidChange() } - completedWaiters = completedIDs.compactMap { - taskCountWaiters.removeValue(forKey: $0)?.continuation + + private func handlingIsPendingLocked() -> Bool { + switch state { + case .admitted(let admitted): admitted.slots.handlingIsPending + case .handling(let handling): handling.slots.handlingIsPending + case .responding(let responding): responding.slots.handlingIsPending + case .closing(let closing): + closing.pending.contains(.httpHandler) || closing.pending.contains(.domain) + case .closed: false + } } - lock.unlock() - for waiter in completedWaiters { - waiter.resume() + + private static func install( + _ cancellation: @escaping @Sendable () -> Void, + reservation: WorkReservation, + slots: inout Slots + ) -> Bool { + guard case .reserved(let id) = slots[reservation.kind], id == reservation.id else { + return false + } + slots[reservation.kind] = .running(reservation.id, cancellation) + return true } - return receipt - } - func registerFiniteResponseOperation() -> FiniteResponseOperation? { - guard let receipt = registerTask(kind: .finiteResponse) else { - return nil + private static func complete( + _ reservation: WorkReservation, + slots: inout Slots + ) -> Bool { + switch slots[reservation.kind] { + case .reserved(let id), .running(let id, _): + guard id == reservation.id else { return false } + slots[reservation.kind] = .completed + return true + case .notAdmitted, .responseReady, .turnGranted, .completed, .closed: + return false + } } - let operation = FiniteResponseOperation(receipt: receipt) - operation.installOwnerCancellation() - return operation } - func performTask( - kind: TaskKind, - operation: @escaping @Sendable () async throws -> Success - ) async throws -> Success { - guard let receipt = registerTask(kind: kind) else { - throw TaskAdmissionClosed(kind: kind) + final class Connection: @unchecked Sendable { + private struct OpenState { + var queue: [RequestOperation] + var writerOperationID: UUID? } - let task = Task { - try await operation() + + private struct ClosingState { + var queue: [RequestOperation] + var writerOperationID: UUID? + var closeAcknowledged: Bool } - receipt.install(task) - return try await withTaskCancellationHandler { - defer { receipt.finish() } - return try await task.value - } onCancel: { - receipt.cancel() + + private enum State { + case accepting(OpenState) + case admissionClosed(OpenState) + case closing(ClosingState) + case closed } - } - func closeTaskAdmission(kind: TaskKind) { - lock.lock() - closedTaskAdmissionKinds.insert(kind) - lock.unlock() - } + let id = UUID() + let ordinal: UInt64 + private weak var owner: MCPHTTPNetworkResourceOwner? + private let channel: any Channel + private let lock = NSLock() + private var state = State.accepting(.init(queue: [], writerOperationID: nil)) + private var nextOperationOrdinal: UInt64 = 0 + private var closeWaiters: [CheckedContinuation] = [] - func closeTaskAdmissionCancelAndDrain() async { - let receipts = closeTaskAdmissionAndSnapshot() - for receipt in receipts { - receipt.cancel() + fileprivate init( + ordinal: UInt64, + channel: any Channel, + owner: MCPHTTPNetworkResourceOwner + ) { + self.ordinal = ordinal + self.channel = channel + self.owner = owner + channel.closeFuture.whenComplete { [weak self] _ in + self?.acknowledgeChannelClose() + } } - await withCheckedContinuation { continuation in + + func admitRequest(metadata: RequestMetadata) -> RequestOperation? { lock.lock() - if tasks.isEmpty { - lock.unlock() - continuation.resume() - } else { - taskDrainWaiters.append(continuation) + guard case .accepting(var open) = state else { lock.unlock() + return nil } + nextOperationOrdinal &+= 1 + let operationID = UUID() + let operation = RequestOperation( + token: .init(connectionID: id, operationID: operationID), + admissionOrdinal: nextOperationOrdinal, + metadata: metadata, + connection: self + ) + open.queue.append(operation) + state = .accepting(open) + lock.unlock() + owner?.changed() + return operation } - } - - private func closeTaskAdmissionAndSnapshot() -> [TaskReceipt] { - lock.lock() - acceptsTasks = false - let receipts = Array(tasks.values) - lock.unlock() - return receipts - } - func cancelTasks(for child: ChildRegistration) { - let receipts: [TaskReceipt] - lock.lock() - receipts = tasks.values.filter { $0.childID == child.id } - lock.unlock() - for receipt in receipts { - receipt.cancel() + func supplyResponse(for operation: RequestOperation) async -> Bool { + guard operation.markResponseReady() else { return false } + pump() + return await operation.waitForWriterTurn() } - } - private func acknowledgeChildClose(_ registration: ChildRegistration) { - let holdWaiters: [CheckedContinuation] - lock.lock() - guard children[registration.id] != nil else { + func beginClosing(_ cause: CloseCause) { + let operations: [RequestOperation] + var shouldCloseChannel = false + lock.lock() + switch state { + case .accepting(let open), .admissionClosed(let open): + operations = open.queue + state = .closing(.init( + queue: open.queue, + writerOperationID: open.writerOperationID, + closeAcknowledged: false + )) + shouldCloseChannel = true + case .closing, .closed: + operations = [] + } lock.unlock() - return + for operation in operations { operation.beginClosing(cause) } + if shouldCloseChannel { channel.close(mode: .all, promise: nil) } + owner?.changed() } - if shouldHoldNextChildCloseAcknowledgement { - shouldHoldNextChildCloseAcknowledgement = false - heldChildCloseAcknowledgementIDs.insert(registration.id) - holdWaiters = childCloseHoldWaiters - childCloseHoldWaiters.removeAll(keepingCapacity: false) - lock.unlock() - for waiter in holdWaiters { - waiter.resume() + + func waitUntilClosed() async { + await withCheckedContinuation { continuation in + lock.lock() + if case .closed = state { + lock.unlock(); continuation.resume() + } else { + closeWaiters.append(continuation); lock.unlock() + } } - return - } - let waiters = finishChildLocked(id: registration.id) - lock.unlock() - for waiter in waiters { - waiter.resume() } - } - private func finishChildLocked( - id: UUID - ) -> [CheckedContinuation] { - children.removeValue(forKey: id) - guard children.isEmpty else { - return [] + fileprivate func closeAdmission() { + let operations: [RequestOperation] + lock.lock() + switch state { + case .accepting(let open): state = .admissionClosed(open); operations = open.queue + case .admissionClosed(let open): operations = open.queue + case .closing(let closing): operations = closing.queue + case .closed: operations = [] + } + lock.unlock() + for operation in operations { operation.closeDomainAdmission() } + owner?.changed() } - let waiters = childDrainWaiters - childDrainWaiters.removeAll(keepingCapacity: false) - return waiters - } - fileprivate func finishTask(id: UUID, kind: TaskKind) { - let holdWaiters: [CheckedContinuation] - lock.lock() - guard tasks[id] != nil else { + fileprivate func resolve(operationID: UUID) -> RequestOperation? { + lock.lock() + let operation: RequestOperation? + switch state { + case .accepting(let open), .admissionClosed(let open): + operation = open.queue.first { $0.token.operationID == operationID } + case .closing(let closing): + operation = closing.queue.first { $0.token.operationID == operationID } + case .closed: operation = nil + } lock.unlock() - return + return operation } - if heldTaskCompletionKind == kind { - heldTaskCompletionKind = nil - heldTaskCompletionIDs.insert(id) - holdWaiters = taskCompletionHoldWaiters - taskCompletionHoldWaiters.removeAll(keepingCapacity: false) - lock.unlock() - for waiter in holdWaiters { - waiter.resume() + + fileprivate func operationsSnapshot() -> [RequestOperation] { + lock.lock() + let operations: [RequestOperation] + switch state { + case .accepting(let open), .admissionClosed(let open): operations = open.queue + case .closing(let closing): operations = closing.queue + case .closed: operations = [] } - return - } - let waiters = finishTaskLocked(id: id) - lock.unlock() - for waiter in waiters { - waiter.resume() + lock.unlock() + return operations } - } - private func finishTaskLocked( - id: UUID - ) -> [CheckedContinuation] { - guard let finished = tasks.removeValue(forKey: id) else { - return [] + fileprivate func operationDidChange() { + pump() + owner?.changed() } - var waiters: [CheckedContinuation] = [] - if tasks.values.contains(where: { $0.kind == finished.kind }) == false { - waiters.append(contentsOf: taskKindDrainWaiters.removeValue( - forKey: finished.kind - ) ?? []) + + fileprivate func operationDidClose(_ operation: RequestOperation) { + var didClose = false + var waiters: [CheckedContinuation] = [] + lock.lock() + switch state { + case .accepting(var open): + open.queue.removeAll { $0 === operation } + if open.writerOperationID == operation.token.operationID { open.writerOperationID = nil } + state = .accepting(open) + case .admissionClosed(var open): + open.queue.removeAll { $0 === operation } + if open.writerOperationID == operation.token.operationID { open.writerOperationID = nil } + state = .admissionClosed(open) + case .closing(var closing): + closing.queue.removeAll { $0 === operation } + if closing.writerOperationID == operation.token.operationID { closing.writerOperationID = nil } + if closing.queue.isEmpty, closing.closeAcknowledged { + state = .closed + waiters = closeWaiters + closeWaiters.removeAll(keepingCapacity: false) + didClose = true + } else { + state = .closing(closing) + } + case .closed: break + } + lock.unlock() + for waiter in waiters { waiter.resume() } + if didClose { owner?.connectionDidClose(self) } + else { pump(); owner?.changed() } } - if tasks.isEmpty { - waiters.append(contentsOf: taskDrainWaiters) - taskDrainWaiters.removeAll(keepingCapacity: false) + + fileprivate func snapshot() -> ConnectionSnapshot { + let phase: ConnectionPhaseSnapshot + let operations: [RequestOperation] + lock.lock() + switch state { + case .accepting(let open): phase = .accepting; operations = open.queue + case .admissionClosed(let open): phase = .admissionClosed; operations = open.queue + case .closing(let closing): phase = .closing; operations = closing.queue + case .closed: phase = .closed; operations = [] + } + lock.unlock() + return .init( + id: id, + ordinal: ordinal, + phase: phase, + operations: operations.map { $0.snapshot() } + ) } - return waiters - } - func waitForTasksDrained(kind: TaskKind) async { - await withCheckedContinuation { continuation in + private func pump() { + var operation: RequestOperation? lock.lock() - if tasks.values.contains(where: { $0.kind == kind }) == false { - lock.unlock() - continuation.resume() - } else { - taskKindDrainWaiters[kind, default: []].append(continuation) + switch state { + case .accepting(var open): + if open.writerOperationID == nil, let head = open.queue.first, head.isResponseReady { + open.writerOperationID = head.token.operationID; operation = head + } + state = .accepting(open) + case .admissionClosed(var open): + if open.writerOperationID == nil, let head = open.queue.first, head.isResponseReady { + open.writerOperationID = head.token.operationID; operation = head + } + state = .admissionClosed(open) + case .closing, .closed: break + } + lock.unlock() + guard let operation else { return } + if operation.grantWriterTurn() == false { + lock.lock() + switch state { + case .accepting(var open): + if open.writerOperationID == operation.token.operationID { open.writerOperationID = nil } + state = .accepting(open) + case .admissionClosed(var open): + if open.writerOperationID == operation.token.operationID { open.writerOperationID = nil } + state = .admissionClosed(open) + case .closing, .closed: break + } lock.unlock() } } - } - - func holdNextTaskCompletionForTesting(kind: TaskKind) { - lock.lock() - precondition( - heldTaskCompletionKind == nil && heldTaskCompletionIDs.isEmpty, - "MCP task completion test gate owns one held completion." - ) - heldTaskCompletionKind = kind - lock.unlock() - } - func waitForHeldTaskCompletionForTesting() async { - await withCheckedContinuation { continuation in + private func acknowledgeChannelClose() { + var operationsToClose: [RequestOperation] = [] + var didClose = false + var waiters: [CheckedContinuation] = [] lock.lock() - if heldTaskCompletionIDs.isEmpty == false { - lock.unlock() - continuation.resume() - } else { - taskCompletionHoldWaiters.append(continuation) - lock.unlock() + switch state { + case .accepting(let open), .admissionClosed(let open): + operationsToClose = open.queue + state = .closing(.init( + queue: open.queue, + writerOperationID: open.writerOperationID, + closeAcknowledged: true + )) + if open.queue.isEmpty { + state = .closed; didClose = true + waiters = closeWaiters; closeWaiters.removeAll(keepingCapacity: false) + } + case .closing(var closing): + closing.closeAcknowledged = true + if closing.queue.isEmpty { + state = .closed; didClose = true + waiters = closeWaiters; closeWaiters.removeAll(keepingCapacity: false) + } else { state = .closing(closing) } + case .closed: break } + lock.unlock() + for operation in operationsToClose { operation.beginClosing(.peerClosed) } + for waiter in waiters { waiter.resume() } + if didClose { owner?.connectionDidClose(self) } else { owner?.changed() } } } - func hasHeldTaskCompletionForTesting() -> Bool { + final class ClosingGeneration: @unchecked Sendable { + private weak var owner: MCPHTTPNetworkResourceOwner? + fileprivate init(owner: MCPHTTPNetworkResourceOwner) { self.owner = owner } + func waitUntilClosed() async { await owner?.waitUntilClosed() } + } + + private struct GenerationState { var connections: [UUID: Connection] } + private enum State { + case accepting(GenerationState) + case admissionClosed(GenerationState) + case closing(GenerationState) + case closed + } + private struct SnapshotWaiter { + let revision: UInt64 + let continuation: CheckedContinuation + } + + private let lock = NSLock() + private var state = State.accepting(.init(connections: [:])) + private var nextConnectionOrdinal: UInt64 = 0 + private var revision: UInt64 = 0 + private var closeWaiters: [CheckedContinuation] = [] + private var snapshotWaiters: [SnapshotWaiter] = [] + + func admitConnection(_ channel: any Channel) -> Connection? { lock.lock() - let hasHeldCompletion = heldTaskCompletionIDs.isEmpty == false + guard case .accepting(var accepting) = state else { lock.unlock(); return nil } + nextConnectionOrdinal &+= 1 + let connection = Connection(ordinal: nextConnectionOrdinal, channel: channel, owner: self) + accepting.connections[connection.id] = connection + state = .accepting(accepting) lock.unlock() - return hasHeldCompletion + changed() + return connection } - func releaseHeldTaskCompletionForTesting() { - let waiters: [CheckedContinuation] + func resolve(_ token: OperationToken) -> RequestOperation? { lock.lock() - let heldIDs = heldTaskCompletionIDs - heldTaskCompletionIDs.removeAll(keepingCapacity: false) - waiters = heldIDs.flatMap { finishTaskLocked(id: $0) } + let connection: Connection? + switch state { + case .accepting(let current), .admissionClosed(let current), .closing(let current): + connection = current.connections[token.connectionID] + case .closed: connection = nil + } lock.unlock() - for waiter in waiters { - waiter.resume() + return connection?.resolve(operationID: token.operationID) + } + + func closeAdmission() { + let connections: [Connection] + lock.lock() + switch state { + case .accepting(let accepting): + state = .admissionClosed(accepting); connections = Array(accepting.connections.values) + case .admissionClosed(let current), .closing(let current): + connections = Array(current.connections.values) + case .closed: connections = [] } + lock.unlock() + for connection in connections { connection.closeAdmission() } + changed() } - func waitForTaskCountForTesting(kind: TaskKind, count: Int) async { - await withCheckedContinuation { continuation in - lock.lock() - if totalTaskCounts[kind, default: 0] >= count { - lock.unlock() - continuation.resume() + func beginClosing(_ cause: CloseCause) -> ClosingGeneration { + let connections: [Connection] + var waiters: [CheckedContinuation] = [] + lock.lock() + switch state { + case .accepting(let current), .admissionClosed(let current): + if current.connections.isEmpty { + state = .closed; connections = [] + waiters = closeWaiters; closeWaiters.removeAll(keepingCapacity: false) } else { - taskCountWaiters[UUID()] = .init( - kind: kind, - targetCount: count, - continuation: continuation - ) - lock.unlock() + state = .closing(current) + connections = current.connections.values.sorted { $0.ordinal < $1.ordinal } } + case .closing(let current): + connections = current.connections.values.sorted { $0.ordinal < $1.ordinal } + case .closed: connections = [] } + lock.unlock() + for connection in connections { connection.beginClosing(cause) } + for waiter in waiters { waiter.resume() } + changed() + return ClosingGeneration(owner: self) } - func taskCountForTesting(kind: TaskKind) -> Int { - lock.lock() - let count = totalTaskCounts[kind, default: 0] - lock.unlock() - return count + func waitForAdmittedHandlingWorkToDrain() async { + let operations = connectionsSnapshot().flatMap { $0.operationsSnapshot() } + for operation in operations { await operation.waitForHandlingWorkToDrain() } + } + + func liveOperationCount(boundTo sessionID: String) -> Int { + connectionsSnapshot() + .flatMap { $0.operationsSnapshot() } + .map { $0.snapshot() } + .filter { $0.boundSessionID == sessionID } + .count } - func childCountForTesting() -> Int { + func snapshot() -> Snapshot { lock.lock() - let count = totalChildCount + let currentRevision = revision + let phase: GenerationPhaseSnapshot + let connections: [Connection] + switch state { + case .accepting(let current): phase = .accepting; connections = Array(current.connections.values) + case .admissionClosed(let current): phase = .admissionClosed; connections = Array(current.connections.values) + case .closing(let current): phase = .closing; connections = Array(current.connections.values) + case .closed: phase = .closed; connections = [] + } lock.unlock() - return count + return .init( + revision: currentRevision, + phase: phase, + connections: connections.sorted { $0.ordinal < $1.ordinal }.map { $0.snapshot() } + ) } - func waitForChildCountForTesting(_ count: Int) async { + func nextSnapshot(after priorRevision: UInt64) async -> Snapshot { await withCheckedContinuation { continuation in lock.lock() - if totalChildCount >= count { - lock.unlock() - continuation.resume() + if revision > priorRevision { + lock.unlock(); continuation.resume(returning: snapshot()) } else { - childCountWaiters[UUID()] = .init( - targetCount: count, - continuation: continuation - ) + snapshotWaiters.append(.init(revision: priorRevision, continuation: continuation)) lock.unlock() } } } - func holdNextChildCloseAcknowledgementForTesting() { + fileprivate func connectionDidClose(_ connection: Connection) { + var waiters: [CheckedContinuation] = [] lock.lock() - precondition( - shouldHoldNextChildCloseAcknowledgement == false - && heldChildCloseAcknowledgementIDs.isEmpty, - "MCP child close test gate owns one acknowledgement." - ) - shouldHoldNextChildCloseAcknowledgement = true + switch state { + case .accepting(var current): + current.connections.removeValue(forKey: connection.id); state = .accepting(current) + case .admissionClosed(var current): + current.connections.removeValue(forKey: connection.id); state = .admissionClosed(current) + case .closing(var current): + current.connections.removeValue(forKey: connection.id) + if current.connections.isEmpty { + state = .closed; waiters = closeWaiters; closeWaiters.removeAll(keepingCapacity: false) + } else { state = .closing(current) } + case .closed: break + } lock.unlock() + for waiter in waiters { waiter.resume() } + changed() } - func waitForHeldChildCloseAcknowledgementForTesting() async { + fileprivate func changed() { + let waiters: [CheckedContinuation] + lock.lock() + revision &+= 1 + let ready = snapshotWaiters.filter { revision > $0.revision } + snapshotWaiters.removeAll { revision > $0.revision } + waiters = ready.map(\.continuation) + lock.unlock() + guard waiters.isEmpty == false else { return } + let current = snapshot() + for waiter in waiters { waiter.resume(returning: current) } + } + + private func waitUntilClosed() async { await withCheckedContinuation { continuation in lock.lock() - if heldChildCloseAcknowledgementIDs.isEmpty == false { - lock.unlock() - continuation.resume() - } else { - childCloseHoldWaiters.append(continuation) - lock.unlock() - } + if case .closed = state { lock.unlock(); continuation.resume() } + else { closeWaiters.append(continuation); lock.unlock() } } } - func releaseHeldChildCloseAcknowledgementForTesting() { - let waiters: [CheckedContinuation] + private func connectionsSnapshot() -> [Connection] { lock.lock() - let heldIDs = heldChildCloseAcknowledgementIDs - heldChildCloseAcknowledgementIDs.removeAll(keepingCapacity: false) - waiters = heldIDs.flatMap { finishChildLocked(id: $0) } - lock.unlock() - for waiter in waiters { - waiter.resume() + let connections: [Connection] + switch state { + case .accepting(let current), .admissionClosed(let current), .closing(let current): + connections = current.connections.values.sorted { $0.ordinal < $1.ordinal } + case .closed: connections = [] } - } - - func resourceCountsForTesting() -> (children: Int, tasks: Int) { - lock.lock() - let counts = (children.count, tasks.count) lock.unlock() - return counts + return connections } } diff --git a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift index 8553d976..d5511091 100644 --- a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift +++ b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift @@ -4,7 +4,7 @@ import MCP @preconcurrency import NIOCore import Testing @_spi(Testing) @testable import CodexReview -import CodexReviewMCPServer +@testable import CodexReviewMCPServer import CodexReviewTesting @Suite("MCP Streamable HTTP server") @@ -105,56 +105,11 @@ struct CodexReviewMCPHTTPServerTests { configuration: .init(host: "127.0.0.1", port: 0) ) try await server.start() - let boundPort = try #require(await server.url.port) - let initializeBody = try makeJSONBody([ - "jsonrpc": "2.0", - "id": 1, - "method": "initialize", - "params": [ - "protocolVersion": "2025-11-25", - "capabilities": [:], - "clientInfo": [ - "name": "CodexReviewKitTests", - "version": "0.0.0", - ], - ], - ]) - let response = await server.handleHTTPRequest(HTTPRequest( - method: "POST", - headers: [ - HTTPHeaderName.host: "127.0.0.1:\(boundPort)", - HTTPHeaderName.accept: "text/event-stream, application/json", - HTTPHeaderName.contentType: "application/json", - ], - body: initializeBody, - path: "/mcp" - )) - let denied = await server.handleHTTPRequest(HTTPRequest( - method: "POST", - headers: [ - HTTPHeaderName.host: "other.local:\(boundPort)", - HTTPHeaderName.accept: "text/event-stream, application/json", - HTTPHeaderName.contentType: "application/json", - ], - body: initializeBody, - path: "/mcp" - )) - - #expect(response.statusCode == 200) - #expect(response.headers[HTTPHeaderName.sessionID]?.isEmpty == false) - #expect(denied.statusCode == 421) + let sessionID = try await initializeSession(endpoint: await server.url) + #expect(sessionID.isEmpty == false) await server.closeAdmission() - let afterAdmissionClose = await server.handleHTTPRequest(HTTPRequest( - method: "POST", - headers: [ - HTTPHeaderName.host: "127.0.0.1:\(boundPort)", - HTTPHeaderName.accept: "text/event-stream, application/json", - HTTPHeaderName.contentType: "application/json", - ], - body: initializeBody, - path: "/mcp" - )) - #expect(afterAdmissionClose.statusCode == 503) + #expect(await server.listenerIsOpenForTesting() == false) + #expect((await server.networkSnapshotForTesting()).phase != .accepting) try await server.stop() } @@ -189,8 +144,6 @@ struct CodexReviewMCPHTTPServerTests { configuration: .init(host: "127.0.0.1", port: 0) ) let configuredURL = await server.url - let priorAdmissionCount = await server.admittedNetworkRequestCountForTesting() - let priorChildCount = await server.childChannelRegistrationCountForTesting() let expectedCleanupFailures: [ReviewLifecycleResourceFailure] = [ .mcpServer("starting listener cleanup failed"), .mcpServer("starting event-loop group cleanup failed"), @@ -231,8 +184,7 @@ struct CodexReviewMCPHTTPServerTests { #expect(await closeFinished.isCompleted() == false) #expect(await server.listenerIsOpenForTesting() == false) #expect(await server.url == configuredURL) - #expect(await server.admittedNetworkRequestCountForTesting() == priorAdmissionCount) - #expect(await server.childChannelRegistrationCountForTesting() == priorChildCount) + #expect((await server.networkSnapshotForTesting()).phase != .accepting) await server.releaseHeldStartCompletionForTesting() #expect(await closeTask.value == expectedCleanupFailures) @@ -242,17 +194,7 @@ struct CodexReviewMCPHTTPServerTests { #expect(await closeFinished.isCompleted()) #expect(await server.listenerIsOpenForTesting() == false) #expect(await server.url == configuredURL) - #expect(await server.admittedNetworkRequestCountForTesting() == priorAdmissionCount) - #expect(await server.childChannelRegistrationCountForTesting() == priorChildCount) - let denied = await server.handleHTTPRequest(HTTPRequest( - method: "GET", - headers: [:], - path: "/mcp" - )) - #expect(denied.statusCode == 503) - let resources = await server.networkResourceCountsForTesting() - #expect(resources.children == 0) - #expect(resources.tasks == 0) + #expect((await server.networkSnapshotForTesting()).isQuiescent) #expect(await recordedStopFailures(server) == expectedCleanupFailures) } @@ -642,9 +584,6 @@ struct CodexReviewMCPHTTPServerTests { summary: "Running" )] ) - await server.waitForFiniteResponsesToDrainForTesting() - await server.holdNextFiniteResponseCompletionForTesting() - let response = try await postJSONRPC( endpoint: endpoint, sessionID: sessionID, @@ -663,27 +602,9 @@ struct CodexReviewMCPHTTPServerTests { as? [[String: Any]] ) #expect(items.first?["jobId"] as? String == "running") - await server.waitForHeldNetworkTaskCompletionForTesting() - - let stopFinished = CompletionFlag() - let stopTask = Task { - try await server.stop() - await stopFinished.complete() - } - await server.waitForFiniteResponseDrainToBeginForTesting() - - #expect(await stopFinished.isCompleted() == false) - #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) != nil) - #expect(await server.eventLoopGroupShutdownCountForTesting() == 0) - - await server.releaseHeldNetworkTaskCompletionForTesting() - try await stopTask.value - - #expect(await stopFinished.isCompleted()) + try await server.stop() #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) == nil) - let resources = await server.networkResourceCountsForTesting() - #expect(resources.children == 0) - #expect(resources.tasks == 0) + #expect((await server.networkSnapshotForTesting()).isQuiescent) } @Test func streamableHTTPScopesReviewReadToTransportSession() async throws { @@ -1236,11 +1157,6 @@ struct CodexReviewMCPHTTPServerTests { try await server.start() let endpoint = await server.url let sessionID = try await initializeSession(endpoint: endpoint) - let priorResponseTaskCount = await server.networkTaskRegistrationCountForTesting( - streamWriter: false - ) - await server.holdNextNetworkTaskCompletionForTesting(streamWriter: false) - let priorAdmissionCount = await server.admittedNetworkRequestCountForTesting() await server.holdNextNetworkHandlerEntryForTesting() let requestBody = try makeJSONBody([ "jsonrpc": "2.0", @@ -1261,9 +1177,9 @@ struct CodexReviewMCPHTTPServerTests { bodyData: requestBody ) } - await server.waitForAdmittedNetworkRequestCountForTesting( - priorAdmissionCount + 1 - ) + _ = await waitForRequestOperation(on: server, jsonRPCID: "2") { snapshot in + snapshot.httpHandler == .running + } await server.releaseNetworkHandlerEntryForTesting() await backend.waitForStartReview() @@ -1273,36 +1189,104 @@ struct CodexReviewMCPHTTPServerTests { try await server.stop() await stopFinished.complete() } - await server.waitForAdmittedHandlerDrainToBeginForTesting() + _ = await waitForRequestOperation(on: server, jsonRPCID: "2") { + $0.terminalCause == .serverStop + } #expect(await server.listenerIsOpenForTesting() == false) - #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) != nil) #expect(await stopFinished.isCompleted() == false) await requestGate.open() await backend.yield(.completed(summary: "Done", result: "review text")) _ = try? await requestTask.value - await server.waitForNetworkTaskRegistrationCountForTesting( - streamWriter: false, - count: priorResponseTaskCount + 1 - ) - await server.waitForHeldNetworkTaskCompletionForTesting() - #expect(await stopFinished.isCompleted() == false) - #expect(await server.eventLoopGroupShutdownCountForTesting() == 0) - await server.releaseHeldNetworkTaskCompletionForTesting() try await stopTask.value #expect(await stopFinished.isCompleted()) #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) == nil) - let resources = await server.networkResourceCountsForTesting() - #expect(resources.children == 0) - #expect(resources.tasks == 0) + #expect((await server.networkSnapshotForTesting()).isQuiescent) } - @Test func stopDrainsFinitePOSTResponseAfterClientCloses() async throws { + @Test func stopFinishesSDKCancelledFinitePOSTWithoutNaturalTerminal() async throws { let backend = FakeCodexReviewBackend() - let requestGate = AsyncGate() - await backend.holdStartReview(with: requestGate) + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend) + ) + let server = CodexReviewMCPHTTPServer( + adapter: CodexReviewMCPServer(store: store), + configuration: .init(host: "127.0.0.1", port: 0) + ) + try await server.start() + let endpoint = await server.url + let sessionID = try await initializeSession(endpoint: endpoint) + let running = CodexReviewJob.makeForTesting( + id: "job-running", + sessionID: sessionID, + cwd: "/tmp/project", + targetSummary: "Working tree changes", + threadID: "thread-running", + turnID: "turn-running", + status: .running, + summary: "Running" + ) + store.loadForTesting( + serverState: .running, + workspaces: [.init(cwd: "/tmp/project")], + jobs: [running] + ) + + let responseTask = Task { + try await postJSONRPCData( + endpoint: endpoint, + sessionID: sessionID, + bodyData: makeJSONBody([ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": [ + "name": "review_await", + "arguments": ["jobID": running.id], + ], + ]) + ) + } + _ = await waitForRequestOperation( + on: server, + jsonRPCID: "2" + ) { $0.domainWorkIsPending } + + let cancellationResponse = try await sendJSONRPCNotification( + endpoint: endpoint, + sessionID: sessionID, + body: [ + "jsonrpc": "2.0", + "method": "notifications/cancelled", + "params": [ + "requestId": 2, + "reason": "Cancel the held finite response", + ], + ] + ) + #expect(cancellationResponse.statusCode == 202) + _ = await waitForRequestOperation( + on: server, + jsonRPCID: "2" + ) { $0.terminalCause == .sdkCancellation } + + let stopTask = Task { + try await server.stop() + } + _ = try await responseTask.value + try await stopTask.value + + let snapshot = await server.networkSnapshotForTesting() + #expect(snapshot.isQuiescent) + #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) == nil) + } + + @Test func oneConnectionEmitsPipelinedPOSTResponsesInAdmissionOrder() async throws { + let backend = FakeCodexReviewBackend() + let firstResponseGate = AsyncGate() + await backend.holdStartReview(with: firstResponseGate) let store = CodexReviewStore.makeTestingStore( backend: TestingCodexReviewStoreBackend(reviewBackend: backend), idGenerator: .init(next: { "job-1" }) @@ -1314,8 +1298,10 @@ struct CodexReviewMCPHTTPServerTests { try await server.start() let endpoint = await server.url let sessionID = try await initializeSession(endpoint: endpoint) - await server.holdNextFiniteResponseCompletionForTesting() - let requestBody = try makeJSONBody([ + let descriptor = try await openRawTCPConnection(endpoint: endpoint) + defer { Darwin.close(descriptor) } + + let firstBody = try makeJSONBody([ "jsonrpc": "2.0", "id": 2, "method": "tools/call", @@ -1327,45 +1313,42 @@ struct CodexReviewMCPHTTPServerTests { ], ], ]) - let requestTask = Task { - try await postJSONRPCData( - endpoint: endpoint, - sessionID: sessionID, - bodyData: requestBody - ) - } - await backend.waitForStartReview() - requestTask.cancel() - _ = try? await requestTask.value - - let stopFinished = CompletionFlag() - let stopTask = Task { - try await server.stop() - await stopFinished.complete() - } - await server.waitForAdmittedHandlerDrainToBeginForTesting() + let secondBody = try makeJSONBody([ + "jsonrpc": "2.0", + "id": 3, + "method": "tools/list", + ]) + try await sendRawPipelinedPOSTs( + descriptor: descriptor, + endpoint: endpoint, + sessionID: sessionID, + bodies: [firstBody, secondBody] + ) - #expect(await server.finiteResponseCompletionIsHeldForTesting() == false) - #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) != nil) - #expect(await stopFinished.isCompleted() == false) + await backend.waitForStartReview() + let secondReady = await waitForRequestOperation( + on: server, + jsonRPCID: "3" + ) { $0.responseIsReady } + let firstHeld = try #require( + secondReady.connectionOperations.first { $0.jsonRPCID == "2" } + ) + #expect(firstHeld.admissionOrdinal < secondReady.admissionOrdinal) + #expect(secondReady.writerIsRunning == false) - await requestGate.open() + await firstResponseGate.open() await backend.yield(.completed(summary: "Done", result: "review text")) - await server.waitForHeldNetworkTaskCompletionForTesting() - - #expect(await server.finiteResponseCompletionIsHeldForTesting()) - #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) != nil) - #expect(await stopFinished.isCompleted() == false) - #expect(await server.eventLoopGroupShutdownCountForTesting() == 0) + let responseBytes = try await readTwoRawHTTPResponses(descriptor: descriptor) + let responses = try parseRawHTTPResponses(responseBytes, expectedCount: 2) + let responseIDs = try responses.map { + try #require(decodeSSEJSON(from: $0.body)["id"] as? Int) + } - await server.releaseHeldNetworkTaskCompletionForTesting() - try await stopTask.value + #expect(responses.map(\.statusCode) == [200, 200]) + #expect(responseIDs == [2, 3]) - #expect(await stopFinished.isCompleted()) - #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) == nil) - let resources = await server.networkResourceCountsForTesting() - #expect(resources.children == 0) - #expect(resources.tasks == 0) + try await server.stop() + #expect((await server.networkSnapshotForTesting()).isQuiescent) } @Test func stopAwaitsSSEWriterCompletionBeforeEventLoopShutdown() async throws { @@ -1380,10 +1363,6 @@ struct CodexReviewMCPHTTPServerTests { try await server.start() let endpoint = await server.url let sessionID = try await initializeSession(endpoint: endpoint) - let priorWriterCount = await server.networkTaskRegistrationCountForTesting( - streamWriter: true - ) - await server.holdNextNetworkTaskCompletionForTesting(streamWriter: true) var request = URLRequest(url: endpoint) request.httpMethod = "GET" @@ -1392,31 +1371,16 @@ struct CodexReviewMCPHTTPServerTests { let (bytes, response) = try await URLSession.shared.bytes(for: request) let httpResponse = try #require(response as? HTTPURLResponse) #expect(httpResponse.statusCode == 200) - await server.waitForNetworkTaskRegistrationCountForTesting( - streamWriter: true, - count: priorWriterCount + 1 - ) - - let stopFinished = CompletionFlag() - let stopTask = Task { - try await server.stop() - await stopFinished.complete() + _ = await waitForNetworkSnapshot(on: server) { snapshot in + snapshot.connections.flatMap(\.operations).contains { + $0.metadata.method == "GET" && $0.writerIsRunning + } } - await server.waitForHeldNetworkTaskCompletionForTesting() - - #expect(await stopFinished.isCompleted() == false) - #expect(await server.eventLoopGroupShutdownCountForTesting() == 0) - #expect(await server.networkResourceCountsForTesting().tasks > 0) - - await server.releaseHeldNetworkTaskCompletionForTesting() - try await stopTask.value + try await server.stop() _ = bytes - #expect(await stopFinished.isCompleted()) #expect(await server.eventLoopGroupShutdownCountForTesting() == 1) - let resources = await server.networkResourceCountsForTesting() - #expect(resources.children == 0) - #expect(resources.tasks == 0) + #expect((await server.networkSnapshotForTesting()).isQuiescent) } @Test func concurrentStopAndRestartAwaitAcceptedChildCloseAcknowledgement() async throws { @@ -1430,11 +1394,11 @@ struct CodexReviewMCPHTTPServerTests { ) try await server.start() let endpoint = await server.url - let priorChildCount = await server.childChannelRegistrationCountForTesting() let descriptor = try await openRawTCPConnection(endpoint: endpoint) defer { Darwin.close(descriptor) } - await server.waitForChildChannelRegistrationCountForTesting(priorChildCount + 1) - await server.holdNextChildCloseAcknowledgementForTesting() + _ = await waitForNetworkSnapshot(on: server) { + $0.connections.isEmpty == false + } let priorShutdownCount = await server.eventLoopGroupShutdownCountForTesting() let firstStopFinished = CompletionFlag() @@ -1447,32 +1411,23 @@ struct CodexReviewMCPHTTPServerTests { try await server.stop() await secondStopFinished.complete() } - await server.waitForHeldChildCloseAcknowledgementForTesting() let restartFinished = CompletionFlag() let restart = Task { try await server.start() await restartFinished.complete() } - #expect(await rawConnectionReachedEOF(descriptor: descriptor)) - #expect(await firstStopFinished.isCompleted() == false) - #expect(await secondStopFinished.isCompleted() == false) - #expect(await restartFinished.isCompleted() == false) - #expect(await server.eventLoopGroupShutdownCountForTesting() == priorShutdownCount) - - await server.releaseHeldChildCloseAcknowledgementForTesting() try await firstStop.value try await secondStop.value try await restart.value + #expect(await rawConnectionReachedEOF(descriptor: descriptor)) #expect(await firstStopFinished.isCompleted()) #expect(await secondStopFinished.isCompleted()) #expect(await restartFinished.isCompleted()) #expect(await server.listenerIsOpenForTesting()) #expect(await server.eventLoopGroupShutdownCountForTesting() == priorShutdownCount + 1) - let restartedResources = await server.networkResourceCountsForTesting() - #expect(restartedResources.children == 0) - #expect(restartedResources.tasks == 0) + #expect((await server.networkSnapshotForTesting()).connections.isEmpty) try await server.stop() #expect(await server.eventLoopGroupShutdownCountForTesting() == priorShutdownCount + 2) @@ -1494,34 +1449,17 @@ struct CodexReviewMCPHTTPServerTests { try await server.start() let endpoint = await server.url let sessionID = try await initializeSession(endpoint: endpoint) - let priorWriterCount = await server.networkTaskRegistrationCountForTesting( - streamWriter: true - ) - await server.holdNextNetworkTaskCompletionForTesting(streamWriter: true) try await openAndCloseRawEventStream(endpoint: endpoint, sessionID: sessionID) - await server.waitForNetworkTaskRegistrationCountForTesting( - streamWriter: true, - count: priorWriterCount + 1 - ) - await server.waitForHeldNetworkTaskCompletionForTesting() - - let stopFinished = CompletionFlag() - let stopTask = Task { - try await server.stop() - await stopFinished.complete() + _ = await waitForNetworkSnapshot(on: server) { snapshot in + snapshot.connections.flatMap(\.operations).contains { + $0.metadata.method == "GET" + } == false } - #expect(await stopFinished.isCompleted() == false) - #expect(await server.eventLoopGroupShutdownCountForTesting() == 0) - - await server.releaseHeldNetworkTaskCompletionForTesting() - try await stopTask.value + try await server.stop() - #expect(await stopFinished.isCompleted()) #expect(await server.eventLoopGroupShutdownCountForTesting() == 1) - let resources = await server.networkResourceCountsForTesting() - #expect(resources.children == 0) - #expect(resources.tasks == 0) + #expect((await server.networkSnapshotForTesting()).isQuiescent) } @Test func streamableHTTPDoesNotExpireSessionWithOpenEventStream() async throws { @@ -1662,6 +1600,10 @@ struct CodexReviewMCPHTTPServerTests { try await seedQueuedAttemptOwnership(in: store, for: running) let response = try await deleteSession(endpoint: await server.url, sessionID: sessionID) + _ = try await store.awaitReview( + sessionID: sessionID, + jobID: running.id + ) #expect(response.statusCode == 200) #expect(running.core.lifecycle.status == .cancelled) @@ -1836,6 +1778,197 @@ struct CodexReviewMCPHTTPServerTests { return data } + private nonisolated func sendJSONRPCNotification( + endpoint: URL, + sessionID: String, + body: [String: Any] + ) async throws -> HTTPURLResponse { + var request = URLRequest(url: endpoint) + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("text/event-stream, application/json", forHTTPHeaderField: "Accept") + request.setValue(sessionID, forHTTPHeaderField: "MCP-Session-Id") + request.httpBody = try JSONSerialization.data(withJSONObject: body) + let (_, response) = try await URLSession.shared.data(for: request) + return try #require(response as? HTTPURLResponse) + } + + private nonisolated func sendRawPipelinedPOSTs( + descriptor: Int32, + endpoint: URL, + sessionID: String, + bodies: [Data] + ) async throws { + try await Task.detached { + let components = try #require(URLComponents( + url: endpoint, + resolvingAgainstBaseURL: false + )) + let host = try #require(components.host) + let port = try #require(components.port) + var requestBytes = Data() + for body in bodies { + let lines = [ + "POST \(endpoint.path) HTTP/1.1", + "Host: \(host):\(port)", + "Accept: text/event-stream, application/json", + "Content-Type: application/json", + "MCP-Session-Id: \(sessionID)", + "MCP-Protocol-Version: 2025-11-25", + "Content-Length: \(body.count)", + "Connection: keep-alive", + "", + "", + ] + requestBytes.append(Data(lines.joined(separator: "\r\n").utf8)) + requestBytes.append(body) + } + try requestBytes.withUnsafeBytes { rawBuffer in + guard let baseAddress = rawBuffer.baseAddress else { + throw testError("Empty pipelined HTTP request") + } + var sent = 0 + while sent < rawBuffer.count { + let count = Darwin.send( + descriptor, + baseAddress.advanced(by: sent), + rawBuffer.count - sent, + 0 + ) + guard count > 0 else { + throw currentPOSIXError() + } + sent += count + } + } + }.value + } + + private nonisolated func readTwoRawHTTPResponses( + descriptor: Int32 + ) async throws -> Data { + try await Task.detached { + var response = Data() + var buffer = [UInt8](repeating: 0, count: 4096) + while true { + let text = String(decoding: response, as: UTF8.self) + let responseCount = text.components( + separatedBy: "HTTP/1.1 200" + ).count - 1 + let terminalCount = text.components( + separatedBy: "\r\n0\r\n\r\n" + ).count - 1 + if responseCount >= 2, terminalCount >= 2 { + return response + } + let count = Darwin.recv(descriptor, &buffer, buffer.count, 0) + guard count > 0 else { + throw testError("Pipelined HTTP connection closed before both responses ended") + } + response.append(contentsOf: buffer.prefix(count)) + guard response.count <= 2 * 1024 * 1024 else { + throw testError("Pipelined HTTP responses exceeded the test bound") + } + } + }.value + } + + private struct RawHTTPResponse { + let statusCode: Int + let body: Data + } + + private nonisolated func parseRawHTTPResponses( + _ data: Data, + expectedCount: Int + ) throws -> [RawHTTPResponse] { + let bytes = Array(data) + let crlf = Array("\r\n".utf8) + let headerTerminal = Array("\r\n\r\n".utf8) + var cursor = 0 + var responses: [RawHTTPResponse] = [] + + func rangeOf( + _ needle: [UInt8], + from start: Int + ) -> Range? { + guard needle.isEmpty == false, start <= bytes.count - needle.count else { + return nil + } + for index in start...(bytes.count - needle.count) { + if bytes[index..<(index + needle.count)].elementsEqual(needle) { + return index..<(index + needle.count) + } + } + return nil + } + + while responses.count < expectedCount { + let headerRange = try #require(rangeOf(headerTerminal, from: cursor)) + let headerData = Data(bytes[cursor.. Bool + ) async -> MatchedRequestOperation { + var snapshot = await server.networkSnapshotForTesting() + while true { + for connection in snapshot.connections { + if let operation = connection.operations.first(where: { + $0.metadata.jsonRPCID == jsonRPCID + }) { + let matched = MatchedRequestOperation( + operation: operation, + connectionOperations: connection.operations + ) + if condition(matched) { + return matched + } + } + } + snapshot = await server.nextNetworkSnapshotForTesting( + after: snapshot.revision + ) + } + } + + private func waitForNetworkSnapshot( + on server: CodexReviewMCPHTTPServer, + satisfying condition: (MCPHTTPNetworkResourceOwner.Snapshot) -> Bool + ) async -> MCPHTTPNetworkResourceOwner.Snapshot { + var snapshot = await server.networkSnapshotForTesting() + while condition(snapshot) == false { + snapshot = await server.nextNetworkSnapshotForTesting( + after: snapshot.revision + ) + } + return snapshot + } +} + +private extension MCPHTTPNetworkResourceOwner.OperationSnapshot { + var jsonRPCID: String? { metadata.jsonRPCID } } private actor CompletionFlag {