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/CodexReviewBackend.swift b/Sources/CodexReview/CodexReviewBackend.swift index e29cc64e..36f3dcf1 100644 --- a/Sources/CodexReview/CodexReviewBackend.swift +++ b/Sources/CodexReview/CodexReviewBackend.swift @@ -10,17 +10,34 @@ 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 beginReviewRecovery( - _ run: CodexReviewBackendModel.Review.Run, - reason: CodexReviewBackendModel.CancellationReason - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken + func forceCloseReviewConnection() async throws + func prepareReviewRecovery( + _ candidate: ReviewRecoveryCandidate + ) async throws -> ReviewRecoveryHandoff func resumeReviewRecovery( - _ token: CodexReviewBackendModel.Review.RecoveryToken, - request: CodexReviewBackendModel.Review.Start + _ handoff: ReviewRecoveryHandoff, + request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission ) 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 registered = try await admission.registerStart { admission in + try await self.startReview(request, admission: admission) + } + try await admission.activateStart(registered.id) + return try await registered.task.value + } } package struct BackendReviewAttempt: Sendable { @@ -36,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] = [] @@ -59,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 } } @@ -96,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() } @@ -166,8 +179,6 @@ package actor BackendReviewEventMailbox { switch terminal { case .finished: return .finished - case .cancelled: - return .cancelled case .failed(let message): return .failed(message) } @@ -183,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 new file mode 100644 index 00000000..26bdeb0e --- /dev/null +++ b/Sources/CodexReview/ReviewAttemptProcessor.swift @@ -0,0 +1,1486 @@ +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 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(ReviewTerminalRecord) + case stream(ReviewAttemptStreamFailure) + case localCancellation(ReviewCancellation) + + package var diagnosticDescription: String { + switch self { + case .canonical(let terminal): + "canonical terminal \(terminal.kind.rawValue)" + case .stream(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 + } +} + +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, + requestFailure: ReviewInterruptRequestFailure? + ) { + self.run = run + self.terminal = terminal + self.requestFailure = requestFailure + } +} + +package struct ReviewRecoveryCandidate: Equatable, Sendable { + package let resolved: ReviewResolvedAttemptTerminal + package let trigger: ReviewAttemptRecoveryTrigger + + fileprivate init( + resolved: ReviewResolvedAttemptTerminal, + trigger: ReviewAttemptRecoveryTrigger + ) { + self.resolved = resolved + self.trigger = trigger + } +} + +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 var resolvedAttempt: ReviewResolvedAttemptTerminal { + switch self { + case .productTerminal(let disposition): + disposition.resolved + case .replacement(let candidate): + candidate.resolved + } + } +} + +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. +package actor ReviewStartAdmission { + package enum RequestDispatch: Equatable, Sendable { + case notSent + case outcomeUnknown + } + + package enum Phase: Equatable, Sendable { + 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 + ) + 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 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? + 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? + 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? + private var graceDidExpire = false + 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>] = [] + + package init(closePolicy: ReviewRuntimeClosePolicy = .production) { + self.closePolicy = closePolicy + } + + package func registerStart( + _ operation: @escaping @Sendable (ReviewStartAdmission) async throws -> BackendReviewAttempt + ) 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 + } catch { + self.finishStart(with: .failure(error)) + throw error + } + } + startTask = 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 { + 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( + message: "Start activation handle \(id.generation) is not pending." + ) + } + phase = .activatedStart(id) + resolveStartActivation(.success(())) + } + + package func admitThreadStartDispatch() throws { + if let terminal { + throw startFailure(for: terminal) + } + if let requestedCancellation { + throw ReviewStartCancelledBeforeDispatch(cancellation: requestedCancellation) + } + 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." + ) + } + 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) + } + 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 + } + phase = .startingReview(preparedRun: run, dispatch: .notSent) + } + + package func admitReviewStartDispatch( + for preparedRun: CodexReviewBackendModel.Review.Run + ) 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 { + throw ReviewAttemptContractFailure( + message: "Review start dispatch requires its pending prepared attempt." + ) + } + phase = .startingReview(preparedRun: preparedRun, dispatch: .outcomeUnknown) + } + + package func recordActiveRun(_ run: CodexReviewBackendModel.Review.Run) { + guard terminal == nil else { + return + } + registeredRun = run + phase = .active(run) + resumeActiveRunWaiters(returning: run) + resumeStartResolutionWaitersIfNeeded() + } + + package func recordCanonicalTerminal( + _ 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(terminalRecord) + guard terminal == candidate else { + throw ReviewAttemptContractFailure( + message: "Conflicting terminal for review attempt \(run.attemptID)." + ) + } + return + } + receiveTerminal(.canonical(terminalRecord)) + } + + 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(candidate) + } + + 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 recoveryDispositionTask { + 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( + cancellation, + interrupt: interrupt, + forceClose: forceClose + ) + return try checkedCancellationResolution(resolution) + } + + package func beginRecovery( + trigger: ReviewAttemptRecoveryTrigger, + interrupt: @escaping @Sendable ( + CodexReviewBackendModel.Review.Run, + CodexReviewBackendModel.CancellationReason + ) async throws -> Void, + forceClose: @escaping @Sendable () async throws -> Void + ) 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( + trigger.cancellation, + interrupt: interrupt, + forceClose: forceClose + ) + 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( + _ 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) + } + if joinedTerminalCancellation == nil { + requestedCancellation = cancellation + resumeCancellationAdmissionWaiters(returning: 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 waitForActiveRun() async -> CodexReviewBackendModel.Review.Run? { + if let activeRun { + return activeRun + } + 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 || startFailed { + continuation.resume(returning: nil) + } else { + activeRunWaiters.append(continuation) + } + } + } + + package func cancellationRequest() -> ReviewCancellation? { + requestedCancellation + } + + package func waitForCancellationAdmission() async -> ReviewCancellation? { + if let terminalCancellation = joinedTerminalCancellation + ?? terminalCancellationPurpose { + return terminalCancellation + } + if recoveryDispositionTask == nil, let requestedCancellation { + return requestedCancellation + } + if terminal != nil { + return nil + } + if startFailed { + return nil + } + return await withCheckedContinuation { continuation in + 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) + } else { + cancellationAdmissionWaiters.append(continuation) + } + } + } + + 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 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? { + guard let cleanupTask = cleanupTasksByAttemptID[run.attemptID] else { + return nil + } + 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 = .preparingInitialRequest + } + + 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 + ) { + startDidFinish = true + switch result { + case .success(let attempt): + registeredRun = attempt.run + if terminal == nil { + phase = .active(attempt.run) + resumeActiveRunWaiters(returning: attempt.run) + } + case .failure(let error): + startFailed = true + if terminal == nil, + let cancellation = (error as? ReviewStartCancelledBeforeDispatch)?.cancellation { + receiveTerminal(.localCancellation(cancellation)) + } else if error is CancellationError, + let requestedCancellation { + switch phase { + case .registeredStart, .activatedStart, .preparingInitialRequest, + .preparingThread(.notSent): + receiveTerminal(.localCancellation(requestedCancellation)) + case .queued, .preparingThread(.outcomeUnknown), + .recoveryRollbackOutcomeUnknown, .startingReview, .active, + .interrupting, .finishing, .terminal: + break + } + } + resumeActiveRunWaiters(returning: nil) + resumeCancellationAdmissionWaiters(returning: nil) + resumeInterruptionAdmissionWaiters(returning: nil) + } + resumeStartResolutionWaitersIfNeeded() + } + + 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 .preparingInitialRequest = phase { + startTask?.cancel() + } else if case .preparingThread(.notSent) = phase { + 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) + } + + if startTask != nil, activeRun == nil, terminal == nil { + await waitForStartResolution() + } + + if case .failure(let error)? = cancellationResult { + await drainCancellationTasks() + if terminal == nil { + resetCancellationOperationForRetry() + } + throw error + } + + 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, + let requestFailure = error as? ReviewInterruptRequestFailure, + case .rejected = requestFailure.outcome { + resetRejectedCancellationForRetry() + } + throw error + } + } + + private func resetRejectedCancellationForRetry() { + resetCancellationOperationForRetry() + } + + private func resetCancellationOperationForRetry() { + requestedCancellation = nil + joinedTerminalCancellation = nil + interruptionPurpose = nil + cancellationTask = nil + interruptRequestTask = nil + terminalBarrierTask = nil + graceTask = nil + forceCloseTask = nil + requestResult = nil + forceCloseResult = nil + graceDidExpire = false + cancellationResult = nil + } + + 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: + graceDidExpire = true + case .forceClose(let result): + forceCloseResult = result + } + resolveCancellationIfPossible() + } + + 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) + for waiter in waiters { + waiter.resume(returning: terminal) + } + resumeActiveRunWaiters(returning: nil) + resumeCancellationAdmissionWaiters(returning: nil) + resumeInterruptionAdmissionWaiters(returning: nil) + resumeStartResolutionWaitersIfNeeded() + 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 resumeCancellationAdmissionWaiters( + returning cancellation: ReviewCancellation? + ) { + let waiters = cancellationAdmissionWaiters + cancellationAdmissionWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume(returning: cancellation) + } + } + + 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 + } + if graceDidExpire, forceCloseResult == nil { + return + } + if case .failure(let closeFailure)? = forceCloseResult { + 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 .stream(let streamFailure): + resolveCancellation(.success(.init( + terminal: terminal, + requestFailure: ReviewInterruptRequestFailure( + outcome: requestFailure.outcome, + secondaryBarrierDiagnostic: streamFailure.localizedDescription + ) + ))) + case .canonical, .localCancellation: + resolveCancellation(.success(.init( + terminal: terminal, + requestFailure: requestFailure + ))) + } + } + } + } + + private func resolveCancellation( + _ result: Result + ) { + cancellationResult = result + resumeStartResolutionWaitersIfNeeded() + 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 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) + } + + private func drainCancellationTasks() async { + if terminal == nil { + terminalBarrierTask?.cancel() + } + if case .failure? = forceCloseResult { + interruptRequestTask?.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, .terminal: + registeredRun + case .queued, .registeredStart, .activatedStart, .preparingInitialRequest, + .preparingThread, .recoveryRollbackOutcomeUnknown: + registeredRun + } + } + + private var activeRun: CodexReviewBackendModel.Review.Run? { + switch phase { + case .active(let run), .interrupting(let run): + run + case .queued, .registeredStart, .activatedStart, .preparingInitialRequest, + .preparingThread, .recoveryRollbackOutcomeUnknown, .startingReview, + .finishing, .terminal: + nil + } + } + + private static func isOutcomeUnknownStartPhase(_ phase: Phase) -> Bool { + switch phase { + case .preparingThread(.outcomeUnknown), .recoveryRollbackOutcomeUnknown, + .startingReview(_, .outcomeUnknown): + true + case .queued, .registeredStart, .activatedStart, + .preparingInitialRequest, .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/ReviewRuntimeLifecycle.swift b/Sources/CodexReview/ReviewRuntimeLifecycle.swift new file mode 100644 index 00000000..b0ae506e --- /dev/null +++ b/Sources/CodexReview/ReviewRuntimeLifecycle.swift @@ -0,0 +1,818 @@ +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 let closeRecord: RuntimeCloseRecord + + @MainActor + package init( + snapshot: RuntimePublicationSnapshot, + 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 + } + 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) + } + + 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)" + } +} + +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, registrationID: UUID) { + self.jobID = jobID + self.attemptID = attemptID + self.registrationID = registrationID + self.phase = .eligible + } +} + +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] = [] + + 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 + } + + @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 suppressParticipant(jobID: String, registrationID: UUID) { + guard let index = participants.firstIndex(where: { + $0.jobID == jobID && $0.registrationID == registrationID + }) else { + return + } + participants[index].phase = .suppressed + } + + 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( + _ 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( + generation: ReviewRuntimeGeneration, + task: Task, + record: ReviewRuntimeTransitionRecord + ) + case running( + generation: ReviewRuntimeGeneration, + runtime: PreparedRuntime, + mcpGeneration: MCPServerGeneration + ) + case transitioning( + generation: ReviewRuntimeGeneration, + purpose: ReviewRuntimeTransitionPurpose, + task: Task, + record: ReviewRuntimeTransitionRecord, + sourceRuntime: PreparedRuntime?, + recoveryReplacement: ReviewRuntimeRecoveryReplacement? + ) + case failed( + generation: ReviewRuntimeGeneration, + retainedMCPGeneration: MCPServerGeneration, + serverURL: URL?, + replacementFailure: ReviewRuntimeReplacementFailure? + ) + + package var generation: ReviewRuntimeGeneration { + switch self { + case .stopped(let generation), + .acquiring(let generation, _, _), + .running(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 +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 d7ab5001..75e847e6 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 @@ -25,18 +30,35 @@ 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 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 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] = [] @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? + @ObservationIgnored package let runtimeWorkerRegistry = ReviewRuntimeWorkerRegistry() package init( backend: any CodexReviewStoreBackend = PreviewCodexReviewStoreBackend(), @@ -45,11 +67,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 @@ -78,10 +102,22 @@ 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, _, _, _): + task.cancel() + case .stopped, .running, .failed: + break + } 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 { @@ -90,6 +126,9 @@ public final class CodexReviewStore { waiter.continuation.resume() } } + for waiter in closeCallerWaiters { + waiter.continuation.resume() + } } public static func makePreviewStore(diagnosticsURL: URL? = nil) -> CodexReviewStore { @@ -113,7 +152,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,67 +161,1013 @@ public final class CodexReviewStore { clock: clock, idGenerator: idGenerator, networkMonitor: networkMonitor, - networkRecoveryPolicy: networkRecoveryPolicy + networkRecoveryPolicy: networkRecoveryPolicy, + reviewRuntimeClosePolicy: reviewRuntimeClosePolicy ) } public func start(forceRestartIfNeeded: Bool = false) async { - switch serverState { - case .stopped, .failed: - break - case .starting: + guard case .open = lifetimeState else { + return + } + switch runtimeState { + case .acquiring: + return + case .transitioning(_, _, let task, _, _, _): + await task.value return case .running where forceRestartIfNeeded == false: return - case .running: + 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 } + 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() + let record = ReviewRuntimeTransitionRecord() 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, + record: record + ) + } + runtimeState = .acquiring( + generation: generation, + task: task, + record: record + ) + await task.value + } + + private func startRuntimeReplacement( + previousGeneration: ReviewRuntimeGeneration, + previousRuntime: PreparedRuntime?, + retainedMCPGeneration: MCPServerGeneration, + retainedServerURL: URL? + ) async { + 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( + replacement, + purpose: purpose, + record: record + ) + } + runtimeState = .transitioning( + generation: generation, + purpose: purpose, + task: task, + record: record, + sourceRuntime: sourceRuntime, + recoveryReplacement: replacement + ) + return (replacement, task) } public func stop() async { - let locallyCancelledJobIDs: [String] - if backend.handlesActiveReviewStopCleanup { - locallyCancelledJobIDs = [] + 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, _, _, _): + 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, + record: record + ) + } + runtimeState = .transitioning( + generation: invalidatedGeneration, + purpose: .stop, + task: task, + record: record, + sourceRuntime: previousState.runtimeForClose, + recoveryReplacement: nil + ) + 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): + 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) + 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 { - locallyCancelledJobIDs = await requestActiveReviewCancellationsForRuntimeStop() + 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, + record: ReviewRuntimeTransitionRecord + ) async { + defer { lastRuntimeTransitionRecord = record } + switch previousState { + case .acquiring(_, let task, let previousRecord): + task.cancel() + await stopPreparedMCPServer(record: record) + await task.value + record.merge(previousRecord) + case .running(_, let runtime, _): + await stopPublishedRuntime(runtime, record: record) + 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) + case .stopped: + break + } + guard case .transitioning(let currentGeneration, .stop, _, _, _, _) = runtimeState, + currentGeneration == invalidatedGeneration + else { + return + } + runtimeState = .stopped(invalidatedGeneration) + if let failureDescription = record.failureDescription { + transitionToFailed(failureDescription) + } else { + transitionToStopped() + } + } + + private func stopPublishedRuntime( + _ runtime: PreparedRuntime, + record: ReviewRuntimeTransitionRecord + ) async { + await performPublishedRuntimeSemanticStop(record: record) + await runtime.handle.closeAdmission() + await stopPreparedMCPServer(record: record) + await closeAppServerRuntime( + runtime, + purpose: .stop, + record: record + ) + } + + 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 ) - transitionToStopped() + 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, + record: ReviewRuntimeTransitionRecord + ) async { + let result = await runtime.closeRecord.closeAndWait( + handle: runtime.handle, + purpose: purpose + ) + record.record(contentsOf: runtime.closeRecord.consumeFailures()) + if result.failures.isEmpty == false { + 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, + record: ReviewRuntimeTransitionRecord + ) async { + defer { lastRuntimeTransitionRecord = record } + 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(record: record) + } + return + } + + let runtime = try await backend.prepareRuntime( + generation: generation, + purpose: purpose + ) + preparedRuntime = runtime + guard isCurrentAcquisition(generation) else { + await closeStaleRuntime( + runtime, + mcpServerWasPrepared: true, + purpose: purpose, + record: record + ) + return + } + + try await runtime.handle.activate() + guard isCurrentAcquisition(generation) else { + await closeStaleRuntime( + runtime, + mcpServerWasPrepared: true, + purpose: purpose, + record: record + ) + return + } + publishRuntimeSnapshot(runtime.snapshot) + + let mcpSnapshot = try await backend.mcpServerLifecycle.activate( + mcpServer.generation + ) + guard isCurrentAcquisition(generation) else { + await closeStaleRuntime( + runtime, + mcpServerWasPrepared: true, + purpose: purpose, + record: record + ) + return + } + + runtimeState = .running( + generation: generation, + runtime: runtime, + mcpGeneration: mcpServer.generation + ) + 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, + record: record + ) + } else if preparedMCPServer != nil, + currentTransitionOwnsMCPStop == false { + await stopPreparedMCPServer(record: record) + } + guard isCurrentAcquisition(generation) else { + return + } + runtimeState = .stopped(generation) + transitionToFailed(visibleFailureDescription) + } + } + + package func performRuntimeAcquisitionForTesting( + generation: ReviewRuntimeGeneration, + purpose: ReviewRuntimeTransitionPurpose + ) async { + await performRuntimeAcquisition( + generation: generation, + purpose: purpose, + record: ReviewRuntimeTransitionRecord() + ) + } + + private func performRuntimeReplacement( + _ replacement: ReviewRuntimeRecoveryReplacement, + purpose: ReviewRuntimeTransitionPurpose, + record: ReviewRuntimeTransitionRecord + ) async { + defer { lastRuntimeTransitionRecord = record } + var preparedRuntime: PreparedRuntime? + await runtimeReplacementEnrollmentPreparationForTesting?() + 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 consumedFailures.isEmpty == false { + record.markForceCloseFailureOwnership( + jobIDs: replacement.forceCloseObserverJobIDs + ) + record.record(contentsOf: consumedFailures) + applicationCloseFailureLedger?.importReceipts(from: record) + } + lastRuntimeTransitionRecord = record + if let firstFailure = closeResult.failures.first { + replacement.finishSourceClose(.failure( + reviewRuntimeCloseFailure(from: firstFailure) + )) + } else { + replacement.finishSourceClose(.success(())) + } + } else { + replacement.finishSourceClose(.success(())) + } + 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: replacement.replacementGeneration, + purpose: purpose + ) + preparedRuntime = runtime + guard isCurrentRuntimeReplacement(replacement, purpose: purpose) else { + await closeStaleRuntime( + runtime, + mcpServerWasPrepared: false, + purpose: purpose, + record: record + ) + replacement.finish(.superseded(currentSupersedingPurpose)) + return + } + + try await runtime.handle.activate() + guard isCurrentRuntimeReplacement(replacement, purpose: purpose) else { + await closeStaleRuntime( + runtime, + mcpServerWasPrepared: false, + purpose: purpose, + record: record + ) + replacement.finish(.superseded(currentSupersedingPurpose)) + return + } + + publishRuntimeSnapshot(runtime.snapshot) + runtimeState = .running( + generation: replacement.replacementGeneration, + runtime: runtime, + mcpGeneration: replacement.retainedMCPGeneration + ) + publishMCPServer(serverURL: replacement.retainedServerURL) + replacement.finish(.running(replacement.replacementGeneration)) + } catch { + if let preparedRuntime { + await closeStaleRuntime( + preparedRuntime, + mcpServerWasPrepared: false, + purpose: purpose, + record: record + ) + } + guard isCurrentRuntimeReplacement(replacement, purpose: purpose) else { + replacement.finish(.superseded(currentSupersedingPurpose)) + return + } + let failure = runtimeReplacementFailure(from: error) + record.record(.lifecycleResources(failure.resources)) + runtimeState = .failed( + generation: replacement.replacementGeneration, + retainedMCPGeneration: replacement.retainedMCPGeneration, + serverURL: replacement.retainedServerURL, + replacementFailure: failure + ) + 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( + _ generation: ReviewRuntimeGeneration + ) -> Bool { + guard case .open = lifetimeState else { + return false + } + guard case .acquiring(let currentGeneration, _, _) = runtimeState else { + return false + } + return currentGeneration == generation + } + + private func isCurrentTransition( + _ generation: ReviewRuntimeGeneration, + purpose: ReviewRuntimeTransitionPurpose + ) -> Bool { + guard case .open = lifetimeState else { + return false + } + guard case .transitioning( + let currentGeneration, + let currentPurpose, + _, + _, + _, + _ + ) = runtimeState else { + return false + } + return currentGeneration == generation && currentPurpose == purpose + } + + private var currentTransitionOwnsMCPStop: Bool { + switch lifetimeState { + case .closing, .closed: + return true + case .open: + break + } + guard case .transitioning(_, let purpose, _, _, _, _) = runtimeState else { + return false + } + return purpose == .stop || purpose == .applicationClose + } + + private func closeStaleRuntime( + _ runtime: PreparedRuntime, + mcpServerWasPrepared: Bool, + purpose: ReviewRuntimeTransitionPurpose, + record: ReviewRuntimeTransitionRecord + ) async { + await runtime.handle.closeAdmission() + 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(record: record) + } + } + + 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() + } + } + + 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 { - 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 { @@ -204,39 +1190,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 @@ -253,6 +1253,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 @@ -261,6 +1262,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 @@ -269,15 +1271,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", @@ -285,35 +1301,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) } @@ -335,19 +1370,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 { @@ -355,7 +1398,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 ad4d13d4..9378795e 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreBackend.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreBackend.swift @@ -24,14 +24,17 @@ package struct CodexReviewStoreSeed { } @MainActor -package protocol CodexReviewStoreBackend: CodexReviewSettingsBackend { +package protocol CodexReviewStoreBackend: CodexReviewSettingsBackend, Sendable { var seed: CodexReviewStoreSeed { get } var isActive: Bool { get } - var handlesActiveReviewStopCleanup: Bool { get } + var mcpServerLifecycle: any MCPServerLifecycleOwner { get } func attachStore(_ store: CodexReviewStore) - func start(store: CodexReviewStore, forceRestartIfNeeded: Bool) async - func stop(store: CodexReviewStore) async + func prepareRuntime( + generation: ReviewRuntimeGeneration, + purpose: ReviewRuntimeTransitionPurpose + ) async throws -> PreparedRuntime + func stop(store: CodexReviewStore) async throws func waitUntilStopped() async func refreshAuth(auth: CodexReviewAuthModel) async func signIn(auth: CodexReviewAuthModel) async @@ -44,21 +47,31 @@ 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 interruptReview(_ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason) async throws - func beginReviewRecovery( + func startReview( + _ request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission + ) async throws -> BackendReviewAttempt + func interruptReview( _ run: CodexReviewBackendModel.Review.Run, + admission: ReviewStartAdmission, reason: CodexReviewBackendModel.CancellationReason - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken + ) async throws + func forceCloseReviewConnection() async throws + func prepareReviewRecovery( + _ candidate: ReviewRecoveryCandidate + ) async throws -> ReviewRecoveryHandoff func resumeReviewRecovery( - _ token: CodexReviewBackendModel.Review.RecoveryToken, - request: CodexReviewBackendModel.Review.Start + _ handoff: ReviewRecoveryHandoff, + request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt - func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async -} - -extension CodexReviewStoreBackend { - package var handlesActiveReviewStopCleanup: Bool { - false - } + 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 fcac5bea..0cc8c669 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift @@ -1,57 +1,11 @@ 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 { - 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, @@ -145,16 +99,55 @@ extension CodexReviewStore { } } - package func requestActiveReviewCancellationsForRuntimeStop( - reason: ReviewCancellation = .system(message: "Review runtime stopped.") - ) async -> [String] { - let activeJobIDs = orderedJobs - .filter { $0.isTerminal == false } - .map(\.id) + package func requestActiveReviewCancellationsForApplicationClose( + reason: ReviewCancellation = .system(message: "Review Store closed."), + failureLedger: ReviewCloseFailureLedger + ) async -> ReviewCloseCancellationOutcome { + await reviewCancellationBarrierPreparationForTesting?() + let activeJobIDs = activeReviewJobIDsInRegistrationOrder + var failedJobIDs: Set = [] for jobID in activeJobIDs { - _ = try? await cancelReview(jobID: jobID, cancellation: reason) + do { + _ = try await cancelReview(jobID: jobID, cancellation: reason) + } catch { + 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) + } + } } - return activeJobIDs + return .init( + jobIDs: activeJobIDs, + failedJobIDs: failedJobIDs + ) + } + + private func closePrimaryFailure( + from error: any Error + ) -> 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 .attemptRuntime(.worker(error.localizedDescription)) } @discardableResult @@ -162,9 +155,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 [] } @@ -184,60 +175,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 } - activeRuns.removeValue(forKey: jobID) - reviewRecoveryWaitingJobIDs.remove(jobID) - startingJobIDs.remove(jobID) - startupCancellations.removeValue(forKey: jobID) + 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( @@ -268,4 +274,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 9efc8a58..3eae17a4 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift @@ -15,16 +15,13 @@ extension CodexReviewStore { sessionID: String, 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() - } - } + 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. + _ = try await awaitReview(sessionID: sessionID, jobID: jobID) + await reviewWorkerTasks[jobID]?.value + return try readReview(sessionID: sessionID, jobID: jobID) } @discardableResult @@ -33,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) } @@ -56,7 +53,48 @@ extension CodexReviewStore { private func beginReview( sessionID: String, request: CodexReviewAPI.Start.Request - ) throws -> String { + ) 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( + "Review runtime transition is in progress." + ) + case .failed: + throw CodexReviewAPI.Error.io("Review runtime is not running.") + case .stopped, .running: + break + } guard closedSessions.contains(sessionID) == false else { throw CodexReviewAPI.Error.invalidArguments("Review session \(sessionID) is closed.") } @@ -76,166 +114,219 @@ extension CodexReviewStore { ), logEntries: [] ) + let admission = ReviewStartAdmission(closePolicy: reviewRuntimeClosePolicy) + let startRequest = CodexReviewBackendModel.Review.Start( + jobID: jobID, + sessionID: sessionID, + request: validatedRequest, + 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) - startingJobIDs.insert(jobID) - launchReviewWorker(jobID: jobID, sessionID: sessionID, request: validatedRequest) + reviewAttemptOwnerships[jobID] = .initialStart(registered) + launchReviewWorker( + jobID: jobID, + startRequest: startRequest, + registeredStart: registered, + runtimeGeneration: runtimeGeneration + ) 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, - sessionID: String, - request: CodexReviewAPI.Start.Request + startRequest: CodexReviewBackendModel.Review.Start, + registeredStart: ReviewRegisteredStart, + runtimeGeneration: ReviewRuntimeGeneration ) { reviewWorkerTasks[jobID]?.cancel() reviewWorkerTasks[jobID] = Task { [weak self] in - await self?.runReviewWorker(jobID: jobID, sessionID: sessionID, request: request) + await self?.runReviewWorker( + jobID: jobID, + startRequest: startRequest, + registeredStart: registeredStart, + runtimeGeneration: runtimeGeneration + ) } } private func runReviewWorker( jobID: String, - sessionID: String, - request validatedRequest: CodexReviewAPI.Start.Request + startRequest: CodexReviewBackendModel.Review.Start, + registeredStart: ReviewRegisteredStart, + runtimeGeneration: ReviewRuntimeGeneration ) async { guard let job = job(id: jobID) else { - startingJobIDs.remove(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 backendAttempt = try await backend.startReview(startRequest) - let backendRun = backendAttempt.run - startingJobIDs.remove(jobID) - run = backendRun - if Task.isCancelled { - throw CancellationError() + 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." + ) } - applyBackendRun(backendRun, to: job) - if let startupCancellation = startupCancellations.removeValue(forKey: jobID) { - try? await backend.interruptReview( - backendRun, - reason: .init(message: startupCancellation.message) + let active = ReviewActiveAttempt( + run: backendAttempt.run, + admission: registeredStart.admission + ) + reviewAttemptOwnerships[jobID] = .active(active) + cleanupAttempt = active + applyBackendRun(backendAttempt.run, to: job) + + if job.isTerminal == false { + let completion = try await consumeReviewEvents( + for: backendAttempt, + job: job, + startRequest: startRequest, + runtimeGeneration: runtimeGeneration ) - if job.isTerminal == false { + cleanupAttempt = completion.cleanupAttempt + } + } catch let cancellation as ReviewStartCancelledBeforeDispatch { + if job.isTerminal == false { + do { try completeCancellationLocally( jobID: job.id, sessionID: job.sessionID, - cancellation: startupCancellation + cancellation: cancellation.cancellation ) + } catch { + markReviewFailed(job, message: error.localizedDescription) } - } 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) - activeRuns.removeValue(forKey: jobID) - reviewRecoveryWaitingJobIDs.remove(jobID) - } else { - let currentRun = try await consumeReviewEvents( - for: backendAttempt, - job: job, - startRequest: startRequest - ) - run = currentRun - await backend.cleanupReview(currentRun) - 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 { - try? completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: startupCancellation ?? job.core.lifecycle.cancellation ?? .system() + if let active = activeAttemptForCleanup(jobID: jobID) ?? cleanupAttempt { + cleanupAttempt = active + 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 { + markReviewFailed(job, message: error.localizedDescription) + } + 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) - } - 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 { - markReviewInterrupted( - job, - cause: .transport(message: transportFailure.message) - ) + if let active = activeAttemptForCleanup(jobID: jobID) ?? cleanupAttempt { + cleanupAttempt = active + } + if job.isTerminal == false, + let streamFailure = error as? ReviewAttemptStreamFailure { + if let cleanupAttempt { + do { + try await cleanupAttempt.admission.recordStreamTerminal(streamFailure) + } catch { + markReviewFailed(job, message: error.localizedDescription) + } + } + await reviewTerminalPublicationPreparationForTesting?() + await applyStreamProductTerminal(streamFailure, to: job) } else if job.isTerminal == false { markReviewFailed(job, message: error.localizedDescription) } } + + reviewAttemptOwnerships[jobID] = .terminal + if let cleanupAttempt { + await reviewCleanupPreparationForTesting?() + do { + try await cleanupReview( + cleanupAttempt.run, + admission: cleanupAttempt.admission + ) + } catch { + retainCleanupFailure(error, for: jobID) + } + } reviewWorkerTasks.removeValue(forKey: jobID) - runtimeStopDetachedReviewWorkerTasks.removeValue(forKey: jobID) + if case .terminal = reviewAttemptOwnerships[jobID] { + reviewAttemptOwnerships.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 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 + ) 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) { - activeRuns[job.id] = backendRun job.core.run = .init( reviewThreadID: backendRun.reviewThreadID, threadID: backendRun.threadID, @@ -259,7 +350,12 @@ extension CodexReviewStore { appendRecoveryProgress(networkRecoveryUnavailableMessage, to: job) } - private func reviewWorkerInputs(for attempt: BackendReviewAttempt) async -> ReviewWorkerInputs { + private func reviewWorkerInputs( + for attempt: BackendReviewAttempt, + owner: ReviewActiveAttempt, + jobID: String, + runtimeGeneration: ReviewRuntimeGeneration + ) async -> ReviewWorkerInputs { let networkMonitor = self.networkMonitor let policy = self.networkRecoveryPolicy let snapshots = networkMonitor.snapshots() @@ -271,16 +367,31 @@ extension CodexReviewStore { queue: queue ) let eventSource = ReviewWorkerEventSource(queue: queue) + let recoveryInterruptionSource = ReviewWorkerRecoveryInterruptionSource(queue: queue) + let runtimeWorkerRegistrationID = runtimeWorkerRegistry.register( + jobID: jobID, + attemptID: owner.run.attemptID, + runtimeGeneration: runtimeGeneration, + recoverySource: recoveryInterruptionSource + ) let networkTask = Task { for await snapshot in snapshots { 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, eventSource: eventSource, + recoveryInterruptionSource: recoveryInterruptionSource, + runtimeWorkerRegistry: runtimeWorkerRegistry, + runtimeWorkerRegistrationID: runtimeWorkerRegistrationID, + jobID: jobID, + initialRuntimeGeneration: runtimeGeneration, initialEventSubscriptionID: initialEventSubscriptionID, networkTask: networkTask, signalCoordinator: signalCoordinator @@ -403,83 +514,507 @@ extension CodexReviewStore { return .init(jobID: job.id, cancelled: false, core: job.core) } - recordCancellationRequest(cancellation, for: job) - - if job.core.lifecycle.status == .queued { - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation + 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: ownedCancellation ) - return .init(jobID: job.id, cancelled: true, core: job.core) + } catch { + if job.isTerminal == false { + try recordCancellationFailure( + jobID: job.id, + sessionID: job.sessionID, + message: error.localizedDescription + ) + } + throw error } + return .init( + jobID: job.id, + cancelled: job.core.lifecycle.status == .cancelled, + core: job.core + ) + } - if reviewRecoveryWaitingJobIDs.contains(jobID) { - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, + 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, + jobID: jobID + ) + 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, + jobID: jobID + ) + 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, + jobID: jobID + ) + 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() - return .init(jobID: job.id, cancelled: true, core: job.core) + await reviewWorkerTasks[jobID]?.value + case .preparingRecovery, .waitingForRecovery: + try await suppressRecoverySuccessor( + job: job, + 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, + cancellation: cancellation, + jobID: jobID + ) + 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." + ) + } } + } - if let run = activeRuns[jobID] { - do { + private func cancel( + admission: ReviewStartAdmission, + cancellation: ReviewCancellation, + jobID: String + ) async throws -> ReviewAttemptCancellationResolution { + let backend = self.backend + return try await admission.cancel( + cancellation, + interrupt: { run, reason in try await backend.interruptReview( run, - reason: .init(message: cancellation.message) + admission: admission, + reason: reason ) - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation + }, + 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, + trigger: .explicitCancellation(targetJobID: jobID) ) - reviewWorkerTasks[jobID]?.cancel() - } catch { - try recordCancellationFailure( - jobID: job.id, - sessionID: job.sessionID, - message: error.localizedDescription + } + ) + } + + 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 runtimeWorkerRegistry.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 + ) + 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 { + runtimeWorkerRegistry.suppressParticipant( + participant, + in: replacement ) - throw error + continue } - } else if let run = job.backendRun { - do { - try await backend.interruptReview( - run, - reason: .init(message: cancellation.message) + 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: + runtimeWorkerRegistry.suppressParticipant( + participant, + in: replacement ) - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation + continue + } + let recoveryTrigger: ReviewAttemptRecoveryTrigger = if case .recoverableNetwork( + let initiatingJobID + ) = replacement.trigger, initiatingJobID == participant.jobID { + .recoverableNetworkLoss + } else { + .sameAccountRestart + } + let didInstall = await runtimeWorkerRegistry.beginRecovery( + replacement: replacement, + participant: participant, + owner: active, + trigger: recoveryTrigger, + interrupt: { run, reason in + try await backend.interruptReview( + run, + admission: active.admission, + 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) + } + runtimeWorkerRegistry.suppressParticipant( + participant, + in: replacement ) - reviewWorkerTasks[jobID]?.cancel() - } catch { - try recordCancellationFailure( - jobID: job.id, - sessionID: job.sessionID, - message: error.localizedDescription + } + } + } + + 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 { + runtimeWorkerRegistry.suppressParticipant( + jobID: jobID, + in: replacement ) - throw error } - } else if startingJobIDs.contains(jobID) { - startupCancellations[jobID] = cancellation - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation + let sourceCloseResult = await replacement.waitForSourceClose() + await runtimeForceCloseReceiptRecordedForTesting?() + try sourceCloseResult.get() + return + } + + let runtime: PreparedRuntime + let record: ReviewRuntimeTransitionRecord + switch runtimeState { + case .running(_, let runningRuntime, _): + runtime = runningRuntime + record = ReviewRuntimeTransitionRecord() + case .transitioning(_, _, _, let transitionRecord, let sourceRuntime?, _): + runtime = sourceRuntime + 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 ) - return .init(jobID: job.id, cancelled: true, core: job.core) } else { - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, + record.recordForceCloseFailures(consumedFailures, jobID: jobID) + } + 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 + ) 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, + jobID: job.id + ) + 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, + jobID: job.id + ) + 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, + jobID: job.id + ) + 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 } - return .init(jobID: job.id, cancelled: true, core: job.core) + 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, + case .terminal = reviewAttemptOwnerships[jobID] + else { + return + } + reviewAttemptOwnerships.removeValue(forKey: jobID) } package func closeSession( @@ -492,12 +1027,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 } } @@ -549,6 +1094,9 @@ extension CodexReviewStore { workspaces.insert(workspace) } jobs.insert(job) + if reviewRegistrationOrder.contains(job.id) == false { + reviewRegistrationOrder.append(job.id) + } writeDiagnosticsIfNeeded() } @@ -605,247 +1153,642 @@ extension CodexReviewStore { private func consumeReviewEvents( for initialAttempt: BackendReviewAttempt, job: CodexReviewJob, - startRequest: CodexReviewBackendModel.Review.Start - ) async throws -> CodexReviewBackendModel.Review.Run { - let inputs = await reviewWorkerInputs(for: initialAttempt) - defer { - inputs.cancel() + startRequest: CodexReviewBackendModel.Review.Start, + runtimeGeneration: ReviewRuntimeGeneration + ) async throws -> ReviewWorkerAttemptCompletion { + 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." + ) } - var recoveryState = ReviewNetworkRecoveryLoopState(currentRun: initialAttempt.run) + 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 { + await inputs.cancel() + throw ReviewAttemptContractFailure( + message: "Initial event subscription became stale before publication completed." + ) + } + do { + let completion = try await consumeReviewEventLoop( + job: job, + startRequest: startRequest, + inputs: inputs, + initialCleanupAttempt: initialActive + ) + await inputs.cancel() + return completion + } catch { + if error is CancellationError || Task.isCancelled { + if let active = activeAttemptForCleanup(jobID: job.id) { + try await active.admission.recordStreamTerminal(.ownerCancellation) + } + } + await inputs.cancel() + throw error + } + } + + private func consumeReviewEventLoop( + job: CodexReviewJob, + startRequest: CodexReviewBackendModel.Review.Start, + inputs: ReviewWorkerInputs, + initialCleanupAttempt: ReviewActiveAttempt + ) async throws -> ReviewWorkerAttemptCompletion { + 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 recoveryState.currentRun + return .init(cleanupAttempt: cleanupAttempt) } switch input { case .reviewEvent(let event): guard activeEventSubscriptionID == event.subscriptionID, - recoveryState.shouldConsumeEvent(from: event.subscriptionRun) + let routed = routedAttempt( + jobID: job.id, + owner: event.owner + ) else { continue } - recoveryState.currentRun = handleReviewEvent( + if let terminal = reviewTerminalRecord(for: event.event, job: job) { + try await routed.active.admission.recordCanonicalTerminal( + terminal, + for: event.owner.run + ) + } + if routed.isResolvingRecovery { + guard event.event.supersedesNetworkRecovery else { + continue + } + } + _ = handleReviewEvent( event.event, job: job, - currentRun: recoveryState.currentRun + currentRun: routed.active.run ) if job.isTerminal { - return recoveryState.currentRun + return .init(cleanupAttempt: routed.active) } case .reviewEventsFinished(let finishedRun): - guard activeEventSubscriptionID == finishedRun.subscriptionID else { + guard activeEventSubscriptionID == finishedRun.subscriptionID, + let routed = routedAttempt(jobID: job.id, owner: finishedRun.owner) + else { continue } - if recoveryState.shouldIgnoreFinishedEvent(for: finishedRun.run) { + let failure = ReviewAttemptStreamFailure.workerContract(.init( + message: ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription + )) + try await routed.active.admission.recordStreamTerminal(failure) + if routed.isResolvingRecovery { + activeEventSubscriptionID = nil continue } - if try handleReviewEventsFinished( - job: job, - isWaitingForNetworkRecovery: recoveryState.isWaitingForNetworkRecovery - ) { - return recoveryState.currentRun + await reviewTerminalPublicationPreparationForTesting?() + if let productTerminal = await routed.active.admission + .terminalCancellationProductTerminal(for: failure) { + try applyRecoveryProductTerminal(productTerminal, to: job) + } else { + await applyStreamProductTerminal(failure, to: job) } + return .init(cleanupAttempt: routed.active) case .reviewEventsFailed(let failedRun): guard activeEventSubscriptionID == failedRun.subscriptionID, - recoveryState.shouldConsumeEvent(from: failedRun.run) + let routed = routedAttempt(jobID: job.id, owner: failedRun.owner) else { continue } - if failedRun.failure.isCancellation { - throw CancellationError() + if routed.isResolvingRecovery { + try await routed.active.admission.recordStreamTerminal(failedRun.failure) + activeEventSubscriptionID = nil + continue } - if await inputs.networkStatusTracker.currentStatus() != .satisfied { - recoveryState.recordPendingOutageStreamFailure(failedRun.failure) + if case .recoverableNetwork = failedRun.failure, + await inputs.networkStatusTracker.currentStatus() != .satisfied { + recoverySignals.recordPendingOutageStreamFailure( + failedRun.failure, + attemptID: routed.active.run.attemptID + ) activeEventSubscriptionID = nil await inputs.cancelActiveEventSubscription() continue } - try throwReviewEventStreamFailure(failedRun.failure) + 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 { + await applyStreamProductTerminal(failedRun.failure, to: job) + } + return .init(cleanupAttempt: routed.active) + case .recoveryBarrierResolved(let resolution): + guard case .resolvingRecovery(let resolving) = reviewAttemptOwnerships[job.id], + sameAttempt(resolving, resolution.owner) + else { + continue + } + switch resolution.result { + case .failure(let failure): + throw failure.underlying + case .success(let disposition): + reviewAttemptOwnerships[job.id] = .recoveryDisposition(disposition) + let candidate: ReviewRecoveryCandidate + switch disposition { + case .productTerminal(let product): + try applyRecoveryProductTerminal(product.productTerminal, to: job) + return .init(cleanupAttempt: resolving) + case .replacement(let replacement): + candidate = replacement + } + 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: cleanupAttempt) + } + throw ReviewAttemptContractFailure( + message: "Recovery preparation completed after its ownership changed." + ) + } + let handoff = try preparationResult.get() + reviewAttemptOwnerships[job.id] = .waitingForRecovery(handoff) + markReviewWaitingForNetworkRecovery(job) + activeEventSubscriptionID = nil + await inputs.cancelActiveEventSubscription() + guard case .waitingForRecovery(let currentHandoff) = reviewAttemptOwnerships[job.id], + currentHandoff == handoff + else { + if job.isTerminal { + return .init(cleanupAttempt: cleanupAttempt) + } + throw ReviewAttemptContractFailure( + 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 = recoveryState.takePendingOutageStreamFailureAfterTransientRecovery( - snapshot - ) { - try throwReviewEventStreamFailure(pendingFailure) + if let pendingFailure = recoverySignals + .takePendingOutageStreamFailureAfterTransientRecovery(snapshot), + case .active(let active) = reviewAttemptOwnerships[job.id], + active.run.attemptID == pendingFailure.attemptID { + throw pendingFailure.failure + } + let waitingHandoff: ReviewRecoveryHandoff? = if case .waitingForRecovery(let handoff) = reviewAttemptOwnerships[job.id] { + handoff + } else { + nil } - switch recoveryState.networkSnapshotEffect(snapshot, recoveryGeneration: recoveryGeneration) { + 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( + runtimeWorkerRegistry.recordNetworkRestoration( + registrationID: inputs.runtimeWorkerRegistrationID, + jobID: job.id + ) + switch try await restartReviewAfterRuntimeRecovery( job: job, startRequest: startRequest, inputs: inputs, - recoveryToken: recoveryState.recoveryToken + handoff: handoff, + sourceRuntimeGeneration: activeRuntimeGeneration ) { case .continueWaiting: - recoveryState.markWaitingForNetworkRecovery() continue case .finished: - reviewRecoveryWaitingJobIDs.remove(job.id) - return recoveryState.currentRun - case .recovered(let recoveredAttempt): - let recoveredRun = recoveredAttempt.run - applyBackendRun(recoveredRun, to: job) - recoveryState.markRecovered(with: recoveredRun) - reviewRecoveryWaitingJobIDs.remove(job.id) - activeEventSubscriptionID = await inputs.subscribe(to: recoveredAttempt) + 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 .networkOutageConfirmed: - guard recoveryState.isWaitingForNetworkRecovery == false, + guard case .active(let active) = reviewAttemptOwnerships[job.id], 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.markRecoveryToken(recoveryToken) + reviewAttemptOwnerships[job.id] = .resolvingRecovery(active) + let pendingFailure = recoverySignals + .takePendingOutageStreamFailureForConfirmedRecovery( + attemptID: active.run.attemptID + ) + let backend = self.backend + await inputs.beginRecoveryInterruption(for: active) { [self] in + try await active.admission.beginRecovery( + trigger: .recoverableNetworkLoss, + interrupt: { run, reason in + try await backend.interruptReview( + run, + admission: active.admission, + reason: reason + ) + }, + 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) + ) + } + ) + } + if let pendingFailure { + _ = await active.admission.waitForInterruptionAdmission() + guard case .resolvingRecovery(let current) = reviewAttemptOwnerships[job.id], + sameAttempt(current, active) + else { + continue + } + try await active.admission.recordStreamTerminal(pendingFailure.failure) + } } } if Task.isCancelled { throw CancellationError() } - if job.isTerminal == false { - if completePendingCancellationIfNeeded(for: job) { - return recoveryState.currentRun - } - markReviewFailed( - job, + if job.isTerminal == false, + case .active(let active) = reviewAttemptOwnerships[job.id] { + let failure = ReviewAttemptStreamFailure.workerContract(.init( 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 { + await applyStreamProductTerminal(failure, to: job) + } + return .init(cleanupAttempt: active) } - return recoveryState.currentRun - } - - private func handleReviewEventsFinished( - job: CodexReviewJob, - isWaitingForNetworkRecovery: Bool - ) throws -> Bool { - if Task.isCancelled { - throw CancellationError() + if job.isTerminal { + return .init(cleanupAttempt: cleanupAttempt) } + throw ReviewAttemptContractFailure( + message: "Review input queue finished without terminal attempt ownership." + ) + } - if isWaitingForNetworkRecovery { - return job.isTerminal || completePendingCancellationIfNeeded(for: job) + private func applyRecoveryProductTerminal( + _ terminal: ReviewTerminalRecord, + to job: CodexReviewJob + ) throws { + switch terminal { + case .completed: + if job.isTerminal == false { + markReviewFailed( + job, + message: "Canonical completion was not reduced before recovery disposition." + ) + } + 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) } + } - if job.isTerminal == false { - if completePendingCancellationIfNeeded(for: job) { - return true - } - markReviewFailed( + private func applyStreamProductTerminal( + _ failure: ReviewAttemptStreamFailure, + to job: CodexReviewJob + ) async { + 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, - message: ReviewIngestionError.streamEndedWithoutTerminal.localizedDescription + cause: .transport(message: failure.localizedDescription) ) } - return true } - private func throwReviewEventStreamFailure(_ failure: ReviewWorkerEventStreamFailure) throws -> Never { - switch failure { - case .cancelled: - throw CancellationError() - case .failed(let message): - throw ReviewWorkerInputQueueError(message: message) + 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 restartReviewAfterNetworkRestore( + 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 restartReviewAfterRuntimeRecovery( job: CodexReviewJob, startRequest: CodexReviewBackendModel.Review.Start, inputs: ReviewWorkerInputs, - recoveryToken: CodexReviewBackendModel.Review.RecoveryToken? + handoff: ReviewRecoveryHandoff, + sourceRuntimeGeneration: ReviewRuntimeGeneration ) 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 { + if handoff.candidate.trigger == .recoverableNetworkLoss, + await inputs.networkStatusTracker.currentStatus() != .satisfied { return .continueWaiting } - guard let recoveryToken else { - return .continueWaiting - } - let recoveredAttempt = try await backend.resumeReviewRecovery( - recoveryToken, - request: startRequest + let replacement = inputs.runtimeWorkerRegistry.replacement( + registrationID: inputs.runtimeWorkerRegistrationID, + jobID: job.id ) - let recoveredRun = recoveredAttempt.run - if try await stopRecoveredRunIfJobShouldNotResume(recoveredRun, job: job) { + let destinationRuntimeGeneration: ReviewRuntimeGeneration + if let replacement { + 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.runtimeWorkerRegistry.finishParticipant( + registrationID: inputs.runtimeWorkerRegistrationID, + jobID: job.id, + replacement: replacement + ) + return .finished + case .superseded: + inputs.runtimeWorkerRegistry.finishParticipant( + registrationID: inputs.runtimeWorkerRegistrationID, + jobID: job.id, + replacement: replacement + ) + 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 + try await backend.resumeReviewRecovery( + handoff, + request: startRequest, + admission: admission + ) + } + guard case .waitingForRecovery(let revalidatedHandoff) = reviewAttemptOwnerships[job.id], + revalidatedHandoff == handoff + else { + _ = try await cancel( + admission: recoveredAdmission, + cancellation: job.core.lifecycle.cancellation ?? .system(), + jobID: job.id + ) + _ = await registered.task.result return .finished } - return .recovered(recoveredAttempt) - } - - private func stopRecoveredRunIfJobShouldNotResume( - _ recoveredRun: CodexReviewBackendModel.Review.Run, - job: CodexReviewJob - ) async throws -> Bool { - if Task.isCancelled { - try? await backend.interruptReview( - recoveredRun, - reason: .init(message: job.core.lifecycle.cancellation?.message ?? "Cancellation requested.") + reviewAttemptOwnerships[job.id] = .replacementStart( + handoff: handoff, + start: registered + ) + 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 + } + throw ReviewAttemptContractFailure( + message: "Replacement start completed after its ownership changed." ) - await backend.cleanupReview(recoveredRun) - 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.") + 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 ) + } catch { + retainCleanupFailure(error, for: job.id) + if terminalizationError == nil { + terminalizationError = error + } } - await backend.cleanupReview(recoveredRun) - return true - } - - guard job.cancellationRequested else { - return false + if let terminalizationError { + throw terminalizationError + } + return .finished } - - 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 + try backend.commitResumedReviewRecovery( + handoff, + recoveredRun: recoveredAttempt.run ) } catch { - await backend.cleanupReview(recoveredRun) - try? recordCancellationFailure( + do { + try await backend.discardResumedReviewRecovery( + handoff, + recoveredRun: recoveredAttempt.run + ) + } catch { + retainCleanupFailure(error, for: job.id) + } + throw error + } + let active = ReviewActiveAttempt( + run: recoveredAttempt.run, + admission: recoveredAdmission + ) + reviewAttemptOwnerships[job.id] = .active(active) + applyBackendRun(recoveredAttempt.run, to: job) + if let replacement { + inputs.runtimeWorkerRegistry.finishParticipant( + registrationID: inputs.runtimeWorkerRegistrationID, jobID: job.id, - sessionID: job.sessionID, - message: error.localizedDescription + replacement: replacement ) - throw error } - await backend.cleanupReview(recoveredRun) - return true + inputs.runtimeWorkerRegistry.update( + registrationID: inputs.runtimeWorkerRegistrationID, + jobID: job.id, + attemptID: recoveredAttempt.run.attemptID, + runtimeGeneration: destinationRuntimeGeneration + ) + return .recovered( + recoveredAttempt, + active, + destinationRuntimeGeneration + ) } private func handleReviewEvent( @@ -856,10 +1799,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,17 +1869,23 @@ extension CodexReviewStore { return updatedRun } - private func completePendingCancellationIfNeeded(for job: CodexReviewJob) -> Bool { - guard job.cancellationRequested else { - return false + 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 } - let cancellation = job.core.lifecycle.cancellation ?? .system() - try? completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation - ) - return true } private func completeReview( @@ -1039,6 +1984,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 }) @@ -1155,6 +2112,31 @@ 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 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 { @@ -1196,48 +2178,61 @@ 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 failure: ReviewWorkerEventStreamFailure -} - -private enum ReviewWorkerEventStreamFailure: Sendable { - case cancelled - case failed(String) - - var isCancellation: Bool { - switch self { - case .cancelled: - true - case .failed: - false - } - } + var owner: ReviewActiveAttempt + var failure: ReviewAttemptStreamFailure } 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 owner: ReviewActiveAttempt + var result: Result +} + +private struct ReviewWorkerAttemptCompletion: Sendable { + var cleanupAttempt: ReviewActiveAttempt +} + +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, + ReviewActiveAttempt, + ReviewRuntimeGeneration + ) } private enum ReviewNetworkSnapshotEffect { @@ -1245,49 +2240,47 @@ private enum ReviewNetworkSnapshotEffect { case restartSettling } -private struct ReviewNetworkRecoveryLoopState { - var currentRun: CodexReviewBackendModel.Review.Run - private(set) var isWaitingForNetworkRecovery = false - private(set) var recoveryToken: CodexReviewBackendModel.Review.RecoveryToken? +private struct PendingOutageStreamFailure { + var attemptID: String + var failure: ReviewAttemptStreamFailure +} + +private struct ReviewNetworkRecoverySignals { private var isSettlingForNetworkRecovery = false private var recoverySettleGeneration: Int? - private var pendingOutageStreamFailure: ReviewWorkerEventStreamFailure? - let recoveryReason = CodexReviewBackendModel.CancellationReason(message: networkRecoveryUnavailableMessage) - - init(currentRun: CodexReviewBackendModel.Review.Run) { - self.currentRun = currentRun - } + private var recoverySettleHandoff: ReviewRecoveryHandoff? + private var pendingOutageStreamFailure: PendingOutageStreamFailure? - mutating func markWaitingForNetworkRecovery() { - isWaitingForNetworkRecovery = true + mutating func markRecovered() { isSettlingForNetworkRecovery = false recoverySettleGeneration = nil + recoverySettleHandoff = nil pendingOutageStreamFailure = nil } - mutating func markRecoveryToken(_ token: CodexReviewBackendModel.Review.RecoveryToken) { - recoveryToken = token - } - - mutating func markRecovered(with run: CodexReviewBackendModel.Review.Run) { - currentRun = run - isWaitingForNetworkRecovery = false - recoveryToken = nil - isSettlingForNetworkRecovery = false - recoverySettleGeneration = nil - pendingOutageStreamFailure = nil + mutating func recordPendingOutageStreamFailure( + _ failure: ReviewAttemptStreamFailure, + attemptID: String + ) { + pendingOutageStreamFailure = .init(attemptID: attemptID, failure: failure) } - mutating func recordPendingOutageStreamFailure(_ failure: ReviewWorkerEventStreamFailure) { - pendingOutageStreamFailure = failure + mutating func takePendingOutageStreamFailureForConfirmedRecovery( + attemptID: String + ) -> PendingOutageStreamFailure? { + guard pendingOutageStreamFailure?.attemptID == attemptID else { + return nil + } + defer { + pendingOutageStreamFailure = nil + } + return pendingOutageStreamFailure } mutating func takePendingOutageStreamFailureAfterTransientRecovery( _ snapshot: CodexReviewNetworkSnapshot - ) -> ReviewWorkerEventStreamFailure? { - guard snapshot.status == .satisfied, - isWaitingForNetworkRecovery == false - else { + ) -> PendingOutageStreamFailure? { + guard snapshot.status == .satisfied else { return nil } defer { @@ -1296,39 +2289,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 - && recoveryToken != nil - } - - func shouldConsumeEvent(from run: CodexReviewBackendModel.Review.Run) -> Bool { - isWaitingForNetworkRecovery == false && 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 } } @@ -1337,6 +2331,11 @@ private struct ReviewWorkerInputs { var queue: ReviewWorkerInputQueue var networkStatusTracker: ReviewNetworkStatusTracker var eventSource: ReviewWorkerEventSource + var recoveryInterruptionSource: ReviewWorkerRecoveryInterruptionSource + var runtimeWorkerRegistry: ReviewRuntimeWorkerRegistry + var runtimeWorkerRegistrationID: UUID + var jobID: String + var initialRuntimeGeneration: ReviewRuntimeGeneration var initialEventSubscriptionID: Int var networkTask: Task var signalCoordinator: ReviewNetworkSignalCoordinator @@ -1345,20 +2344,355 @@ 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 { await eventSource.cancelActiveSubscription() } - func cancel() { + func beginRecoveryInterruption( + for owner: ReviewActiveAttempt, + operation: @escaping @Sendable () async throws -> ReviewRecoveryDisposition + ) async { + await recoveryInterruptionSource.start(for: owner, 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() + await runtimeWorkerRegistry.finishRegisteredParticipant( + registrationID: runtimeWorkerRegistrationID, + jobID: jobID + ) + await runtimeWorkerRegistry.unregister( + registrationID: runtimeWorkerRegistrationID, + jobID: jobID + ) + } +} + +private actor ReviewWorkerRecoveryInterruptionSource { + private let queue: ReviewWorkerInputQueue + private var task: Task? + + init(queue: ReviewWorkerInputQueue) { + self.queue = queue + } + + func start( + for owner: ReviewActiveAttempt, + operation: @escaping @Sendable () async throws -> ReviewRecoveryDisposition + ) { + 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(owner: owner, result: result))) + self.finish() + } + } + + func cancel() async { + let task = task + task?.cancel() + await task?.value + self.task = nil + } + + private func finish() { + task = nil + } +} + +@MainActor +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 registrationWaitersByJobID: [String: [RegistrationWaiter]] = [:] + private var anyRegistrationWaitersByJobID: [ + String: [CheckedContinuation] + ] = [:] + + 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 + 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 let registration = workerRegistrationsByJobID[jobID], + registration.id == registrationID + else { + return + } + registration.attemptID = attemptID + registration.runtimeGeneration = runtimeGeneration + } + + fileprivate func unregister(registrationID: UUID, jobID: String) { + 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) + } + + 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, + registrationID: registration.id + ) + } + + fileprivate func replacement( + registrationID: UUID, + jobID: String + ) -> ReviewRuntimeRecoveryReplacement? { + 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, + 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.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, + interrupt: interrupt, + forceClose: forceClose + ) + } + return await owner.admission.waitForInterruptionAdmission() + == .recoverableTransition(trigger) + } + + fileprivate func recordNetworkRestoration( + registrationID: UUID, + jobID: String + ) { + replacement(registrationID: registrationID, jobID: jobID)? + .recordNetworkRestoration() + } + + 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, + in replacement: ReviewRuntimeRecoveryReplacement + ) { + guard let participant = replacement.participants.first(where: { + $0.jobID == jobID + && ($0.phase == .eligible || $0.phase == .recovering) + }) else { + return + } + 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( + registrationID: UUID, + jobID: String + ) { + guard let registration = workerRegistrationsByJobID[jobID], + registration.id == registrationID + else { + return + } + 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 + } } } } @@ -1447,14 +2781,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] = [:] @@ -1465,12 +2791,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 { @@ -1478,38 +2806,67 @@ 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 { - 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, + owner: owner, + subscriptionID: subscriptionID + ) } } 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( _ event: CodexReviewBackendModel.Review.Event, - run: CodexReviewBackendModel.Review.Run, + owner: ReviewActiveAttempt, subscriptionID: Int ) async { guard activeSubscriptionID == subscriptionID, @@ -1519,12 +2876,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 { @@ -1532,13 +2892,20 @@ private actor ReviewWorkerEventSource { } await queue.send(.reviewEventsFinished(.init( subscriptionID: subscriptionID, - run: run + owner: owner ))) } + private func finishTerminalDelivery(subscriptionID: Int) { + guard activeSubscriptionID == subscriptionID else { + return + } + eventTasks.removeValue(forKey: subscriptionID) + } + private func yieldEventsFailed( - _ error: any Error, - run: CodexReviewBackendModel.Review.Run, + _ failure: ReviewAttemptStreamFailure, + owner: ReviewActiveAttempt, subscriptionID: Int ) async { guard eventTasks.removeValue(forKey: subscriptionID) != nil else { @@ -1549,8 +2916,8 @@ private actor ReviewWorkerEventSource { } await queue.send(.reviewEventsFailed(.init( subscriptionID: subscriptionID, - run: run, - failure: error is CancellationError ? .cancelled : .failed(error.localizedDescription) + owner: owner, + failure: failure ))) } } @@ -1613,11 +2980,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/CodexReviewStoreTesting.swift b/Sources/CodexReview/Store/CodexReviewStoreTesting.swift index 851c7d42..ac995308 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreTesting.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreTesting.swift @@ -13,6 +13,64 @@ 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 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, @@ -77,6 +135,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) } @@ -84,7 +143,19 @@ extension CodexReviewStore { } package func cancelAndDrainReviewWorkersForTesting() async { - let tasks = Array(reviewWorkerTasks.values) + Array(runtimeStopDetachedReviewWorkerTasks.values) + 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) for task in tasks { task.cancel() } @@ -93,11 +164,8 @@ extension CodexReviewStore { } reviewWorkerTasks.removeAll(keepingCapacity: false) - runtimeStopDetachedReviewWorkerTasks.removeAll(keepingCapacity: false) - startingJobIDs.removeAll(keepingCapacity: false) - startupCancellations.removeAll(keepingCapacity: false) - activeRuns.removeAll(keepingCapacity: false) - reviewRecoveryWaitingJobIDs.removeAll(keepingCapacity: false) + reviewCleanupFailures.removeAll(keepingCapacity: false) + reviewAttemptOwnerships.removeAll(keepingCapacity: false) let waiters = reviewTerminalWaiters.values.flatMap { $0 } reviewTerminalWaiters.removeAll(keepingCapacity: false) diff --git a/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift b/Sources/CodexReview/Store/PreviewCodexReviewStoreBackend.swift index d75426d9..b61d5caa 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 { @@ -133,27 +136,48 @@ 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 beginReviewRecovery( + package func interruptReview( _: CodexReviewBackendModel.Review.Run, + admission _: ReviewStartAdmission, reason _: CodexReviewBackendModel.CancellationReason - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { + ) async throws {} + + package func forceCloseReviewConnection() async throws {} + + package func prepareReviewRecovery( + _: ReviewRecoveryCandidate + ) async throws -> ReviewRecoveryHandoff { throw CodexReviewAPI.Error.io(Self.previewUnavailableMessage) } package func resumeReviewRecovery( - _: CodexReviewBackendModel.Review.RecoveryToken, - request _: CodexReviewBackendModel.Review.Start + _: ReviewRecoveryHandoff, + request _: CodexReviewBackendModel.Review.Start, + admission _: ReviewStartAdmission ) async throws -> BackendReviewAttempt { throw CodexReviewAPI.Error.io(Self.previewUnavailableMessage) } - package func cleanupReview(_: CodexReviewBackendModel.Review.Run) async {} + 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." 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..e5f8a3f2 100644 --- a/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift +++ b/Sources/CodexReviewAppServer/AppServerCodexReviewBackend.swift @@ -15,13 +15,163 @@ private func makeAppServerReviewAttemptID() -> String { UUID().uuidString } +package struct AppServerRuntimeOwnerLifecycleHandle: Sendable { + private let closeAdmissionOperation: @Sendable () async -> Void + private let closeAndWaitOperation: @Sendable (ReviewRuntimeTransitionPurpose) async throws -> Void + + fileprivate init( + closeAdmissionOperation: @escaping @Sendable () async -> Void, + closeAndWaitOperation: @escaping @Sendable (ReviewRuntimeTransitionPurpose) async throws -> Void + ) { + self.closeAdmissionOperation = closeAdmissionOperation + self.closeAndWaitOperation = closeAndWaitOperation + } + + package func closeAdmission() async { + await closeAdmissionOperation() + } + + package func closeAndWait( + purpose: ReviewRuntimeTransitionPurpose = .stop + ) async throws { + try await closeAndWaitOperation(purpose) + } +} + 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,12 +180,16 @@ 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 private var diagnosedUnknownNotificationMethods: Set = [] + private var connectionStreamFailure: ReviewAttemptStreamFailure? package init( client: AppServerClient, @@ -45,6 +199,55 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { self.threadStartPermissionStrategy = threadStartPermissionStrategy } + package nonisolated var runtimeOwnerLifecycleHandle: AppServerRuntimeOwnerLifecycleHandle { + AppServerRuntimeOwnerLifecycleHandle( + closeAdmissionOperation: { [self] in + await closeAdmissionFromRuntimeOwner() + }, + closeAndWaitOperation: { [self] purpose in + try await closeFromRuntimeOwnerAndWait(purpose: purpose) + } + ) + } + + private func closeAdmissionFromRuntimeOwner() { + reviewOperationRegistry.closeAdmission() + } + + 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()) @@ -126,12 +329,18 @@ 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 { + 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) + let thread = try await startReviewThread(request, admission: admission) controlsByThreadID[thread.threadID] = control let attemptID = makeAppServerReviewAttemptID() let provisionalRun = CodexReviewBackendModel.Review.Run( @@ -147,6 +356,14 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { ) registerReviewEventSession(session, for: provisionalRun) control.recordThreadStarted(threadID: thread.threadID) + await admission.recordPreparedThread(provisionalRun) + + do { + try await admission.admitReviewStartDispatch(for: provisionalRun) + } catch { + try await cleanupReview(provisionalRun) + throw error + } let review: AppServerAPI.Review.Start.Response reviewStartRequestsInFlight += 1 @@ -157,7 +374,10 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } catch { reviewStartRequestsInFlight -= 1 discardUnmatchedReviewNotificationsIfIdle() - await cleanupReview(provisionalRun) + if let terminal = streamTerminal(for: error) { + try await admission.recordStreamTerminal(terminal) + } + try await cleanupReview(provisionalRun) throw error } let reviewThreadID = review.reviewThreadID ?? thread.threadID @@ -173,64 +393,88 @@ 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( + 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. - return try await startReviewThreadWithProfileSelectionPermissions(request) + try await admission.recordThreadStartRejectedForRetry() + 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( + try await admission.recordThreadStartRejectedForRetry() + 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 { + try await admission.admitThreadStartDispatch() + do { + return try await client.send(request) + } catch { + if let terminal = streamTerminal(for: error) { + try await admission.recordStreamTerminal(terminal) + } + throw error } } @@ -277,36 +521,174 @@ 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 func streamTerminal( + for error: any Error + ) -> ReviewAttemptStreamFailure? { + guard let jsonRPCError = error as? JSONRPC.Error else { + return nil + } + switch jsonRPCError { + 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)) + } + } + } + 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) - await finishReviewEventStream( - threadID: run.threadID, - cancellationMessage: reason.message, - buffersMissingContinuation: true - ) } catch { - await session.clearCancellationRequest() - throw error + let failure = Self.interruptRequestFailure(for: error) + if case .rejected = failure.outcome { + await session.clearCancellationRequest() + } + throw failure } } - package func beginReviewRecovery( - _ run: CodexReviewBackendModel.Review.Run, - reason _: CodexReviewBackendModel.CancellationReason - ) 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) + package func forceCloseReviewConnection() async throws { + connectionStreamFailure = .ownerForcedConnectionClose( + .connection("Review connection was force-closed by its attempt owner.") + ) + do { + try await client.close() + } catch { + throw Self.reviewRuntimeCloseFailure(for: error) + } + } + + // Only AppServerRuntimeOwnerLifecycleHandle can enter this transition. The + // notification router never owns that handle, so close cannot await itself. + private func closeFromRuntimeOwnerAndWait( + purpose: ReviewRuntimeTransitionPurpose + ) async throws { + lifecycleTestingObservation.recordCloseCaller() + let closeTask: Task + switch lifecycleState { + case .open: + if purpose == .recoveryReplacement || purpose == .restartSameAccount { + connectionStreamFailure = .ownerForcedConnectionClose( + .connection("Review connection was force-closed by its runtime owner.") + ) + } + closeAdmissionFromRuntimeOwner() + 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) + } + if case .failure(let error) = clientCloseResult { + throw Self.lifecycleCloseFailure(for: error) + } + } + 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() + } + + 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 { + 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( + _ candidate: ReviewRecoveryCandidate + ) async throws -> ReviewRecoveryHandoff { + let run = candidate.resolved.run + let interruption = AppServerReviewInterruption( + threadID: appServerTurnThreadID(for: run), + turnID: run.turnID ?? "" + ) + markTurnAbandoned(run.turnID) markAttemptAbandoned(run, interruption: interruption) if let session = unregisterReviewEventSession(for: run) { await session.abandon() @@ -315,22 +697,38 @@ 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, - request: CodexReviewBackendModel.Review.Start + _ handoff: ReviewRecoveryHandoff, + 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 - 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 @@ -347,6 +745,14 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { isRunFinalized: false ) registerReviewEventSession(session, for: provisionalRun) + await admission.recordPreparedThread(provisionalRun) + do { + try await admission.admitReviewStartDispatch(for: provisionalRun) + } catch { + _ = unregisterReviewEventSession(for: provisionalRun) + await session.abandon() + throw error + } let review: AppServerAPI.Review.Start.Response reviewStartRequestsInFlight += 1 @@ -357,6 +763,9 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } catch { reviewStartRequestsInFlight -= 1 _ = unregisterReviewEventSession(for: provisionalRun) + if let terminal = streamTerminal(for: error) { + try await admission.recordStreamTerminal(terminal) + } await session.abandon() discardUnmatchedReviewNotificationsIfIdle() throw error @@ -377,14 +786,14 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { ) await session.bufferStartupNotifications(takeUnmatchedReviewNotifications(for: recoveredRun)) await session.finalizeRun() + await admission.recordActiveRun(recoveredRun) reviewStartRequestsInFlight -= 1 discardUnmatchedReviewNotificationsIfIdle() 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,50 +804,45 @@ 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) - )) + guard case .open = lifecycleState else { + return } - for threadID in cleanupThreadIDs { - reviewEventSessionCanonicalThreadIDByThreadID.removeValue(forKey: threadID) + var failureMessages: [String] = [] + do { + let _: EmptyResponse = try await client.send(AppServerAPI.Thread.BackgroundTerminals.Clean.Request( + params: .init(threadID: run.threadID) + )) + } catch { + failureMessages.append( + "thread/backgroundTerminals/clean for \(run.threadID): \(error.localizedDescription)" + ) } - reviewThreadIDsForCleanupByThreadID.removeValue(forKey: run.threadID) - } - - package func cleanupActiveReviewsForShutdown(reason: CodexReviewBackendModel.CancellationReason) async { - let runs = await activeReviewRunsForShutdown() - guard runs.isEmpty == false else { - return + do { + let _: AppServerAPI.Thread.Unsubscribe.Response = try await client.send(AppServerAPI.Thread.Unsubscribe.Request( + params: .init(threadID: run.threadID) + )) + } catch { + failureMessages.append( + "thread/unsubscribe for \(run.threadID): \(error.localizedDescription)" + ) } - for run in runs { - if Task.isCancelled { - return - } - try? await interruptReview(run, reason: reason) - if Task.isCancelled { - return + for threadID in cleanupThreadIDs { + do { + let _: EmptyResponse = try await client.send(AppServerAPI.Thread.Delete.Request( + params: .init(threadID: threadID) + )) + } catch { + failureMessages.append( + "thread/delete for \(threadID): \(error.localizedDescription)" + ) } - await cleanupReview(run) } - } - - package func interruptActiveReviewsForShutdown(reason: CodexReviewBackendModel.CancellationReason) async { - let runs = await activeReviewRunsForShutdown() - guard runs.isEmpty == false else { - return + for threadID in cleanupThreadIDs { + reviewEventSessionCanonicalThreadIDByThreadID.removeValue(forKey: threadID) } - for run in runs { - if Task.isCancelled { - return - } - try? await interruptReview(run, reason: reason) + reviewThreadIDsForCleanupByThreadID.removeValue(forKey: run.threadID) + if failureMessages.isEmpty == false { + throw ReviewRuntimeCloseFailure.cleanup(failureMessages.joined(separator: "; ")) } } @@ -466,6 +870,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 @@ -474,17 +895,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 { @@ -501,6 +929,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 { @@ -541,6 +973,7 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { } } activeThreadIDsByAttemptID.removeValue(forKey: run.attemptID) + reviewEventSessionRegistrationOrdinalByAttemptID.removeValue(forKey: run.attemptID) return reviewEventSessionsByAttemptID.removeValue(forKey: run.attemptID) } @@ -558,14 +991,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 @@ -573,16 +998,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 @@ -610,26 +1025,11 @@ 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 + 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) @@ -668,22 +1068,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( @@ -693,9 +1126,11 @@ package actor AppServerCodexReviewBackend: CodexReviewBackend { for try await notification in notifications { await routeReviewNotification(notification) } - await finishAllReviewEventSessions(throwing: nil) + 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 } @@ -836,11 +1271,22 @@ 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)" ) - await client.close() - await finishAllReviewEventSessions(throwing: error) + do { + try await client.close() + } catch { + appServerBackendLogger.error( + "App-server connection close failed during routing containment: \(error.localizedDescription, privacy: .public)" + ) + } + await finishAllReviewEventSessions(throwing: .protocolViolation(.init( + message: error.localizedDescription + ))) } private func diagnoseUnknownNotificationMethod(_ method: String) { @@ -862,11 +1308,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] { @@ -1103,7 +1592,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 } @@ -1113,8 +1602,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() } @@ -1339,7 +1828,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 89c68172..b0ced6b4 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,13 @@ package actor AppServerProcessTransport: JSONRPC.Transport { private var stdoutReaderTask: Task? = nil private var stderrReaderTask: Task? = nil private var closed = false + private var closeTask: Task? + private var terminalError: JSONRPC.Error? - 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 +106,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)") @@ -135,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() @@ -150,42 +160,77 @@ 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, + error: JSONRPC.Error.transportTerminated(.ownerClose), readerTask: readerTask ) } private func closeTransport( terminateProcess: Bool, - error: any Error, + error: JSONRPC.Error, readerTask: ReaderTask? - ) async { - if closed { - if readerTask == nil { - await waitForReaderTasks(excluding: nil) + ) async throws { + let task: Task + if let closeTask { + task = closeTask + } else { + closed = true + terminalError = error + 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: JSONRPC.Error + ) async throws { + 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() + try await closeCompletionForTesting?() + } catch { + 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 } - finishAll(throwing: error) - await waitForReaderTasks(excluding: readerTask) } private func receiveStdout(_ event: AppServerPipeReadEvent) async { @@ -204,11 +249,16 @@ 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 as? JSONRPC.Error + ?? .invalidMessage(error.localizedDescription), + readerTask: .stdout + ) + } catch { + logger.error("Failed to close codex app-server process: \(error.localizedDescription, privacy: .public)") + } return } } @@ -242,15 +292,30 @@ 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 as? JSONRPC.Error + ?? .invalidMessage(error.localizedDescription), + 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, + 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)") + } } private func processMessage(_ data: Data) throws { @@ -409,7 +474,7 @@ package actor AppServerProcessTransport: JSONRPC.Transport { private func throwIfClosed() throws { if closed { - throw JSONRPC.Error.closed + throw terminalError ?? JSONRPC.Error.closed } } } @@ -729,7 +794,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 +804,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) { @@ -886,10 +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) @@ -897,6 +965,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..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,46 +81,11 @@ package final class AppServerReviewControl: @unchecked Sendable { private func sendInterrupt( threadID: String, - turnID: String, - willInterruptActiveTurn: (@Sendable (AppServerReviewInterruption) async -> Void)? + turnID: String ) 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..559c323a 100644 --- a/Sources/CodexReviewAppServer/JSONRPC.swift +++ b/Sources/CodexReviewAppServer/JSONRPC.swift @@ -27,13 +27,14 @@ 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 { 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 0b2ab211..5119f73e 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 -> 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 -> 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: { - await client.close() + try await backend.runtimeOwnerLifecycleHandle.closeAndWait() } ) } @@ -47,20 +53,28 @@ 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 { + package func stop() async throws { await store.stop() - await shutdown() } } @MainActor 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 @@ -73,14 +87,73 @@ 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) {} - func start(store _: CodexReviewStore, forceRestartIfNeeded _: Bool) async { - active = true + func prepareRuntime( + generation _: ReviewRuntimeGeneration, + purpose _: ReviewRuntimeTransitionPurpose + ) async throws -> PreparedRuntime { + let shutdown = shutdown + let handle = DirectRuntimeLifecycleHandle( + onActivate: { [weak self] in self?.active = true }, + 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 { @@ -233,33 +306,53 @@ 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( _ run: CodexReviewBackendModel.Review.Run, + admission _: ReviewStartAdmission, reason: CodexReviewBackendModel.CancellationReason ) async throws { try await backend.interruptReview(run, reason: reason) } - func beginReviewRecovery( - _ run: CodexReviewBackendModel.Review.Run, - reason: CodexReviewBackendModel.CancellationReason - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { - try await backend.beginReviewRecovery(run, reason: reason) + func forceCloseReviewConnection() async throws { + try await backend.forceCloseReviewConnection() + } + + func prepareReviewRecovery( + _ candidate: ReviewRecoveryCandidate + ) async throws -> ReviewRecoveryHandoff { + try await backend.prepareReviewRecovery(candidate) } func resumeReviewRecovery( - _ token: CodexReviewBackendModel.Review.RecoveryToken, - request: CodexReviewBackendModel.Review.Start + _ handoff: ReviewRecoveryHandoff, + request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { - try await backend.resumeReviewRecovery(token, request: request) + 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 { - await backend.cleanupReview(run) + func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async throws { + try await backend.cleanupReview(run) } private static func monitorSettings( @@ -310,6 +403,65 @@ private final class DirectCodexReviewStoreBackend: CodexReviewStoreBackend { } } +@MainActor +private final class DirectRuntimeLifecycleHandle: RuntimeLifecycleHandle { + private let onActivate: @MainActor @Sendable () -> Void + 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, + onCloseAdmission: @escaping @MainActor @Sendable () async -> Void, + onClose: @escaping @MainActor @Sendable () async throws -> Void + ) { + self.onActivate = onActivate + self.onCloseAdmission = onCloseAdmission + self.onClose = onClose + } + + func activate() async throws { + onActivate() + } + + func closeAdmission() async { + await onCloseAdmission() + } + + func close(purpose _: ReviewRuntimeTransitionPurpose) async throws { + 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 let closeTask else { + throw ReviewLifecycleResourceFailure.client( + "Direct runtime wait began before close." + ) + } + try await closeTask.value.get() + } +} + 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 0fb4fa89..1e99989d 100644 --- a/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift +++ b/Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift @@ -12,66 +12,18 @@ 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 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 } } @@ -93,11 +45,38 @@ 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 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 } func start() async throws - func stop() async + func closeAdmission() async + func waitForAdmittedHandlers() async + func stop() async throws +} + +extension CodexReviewMCPHTTPServing { + package func closeAdmission() async {} + package func waitForAdmittedHandlers() async {} } extension CodexReviewMCPHTTPServer: CodexReviewMCPHTTPServing {} @@ -124,9 +103,12 @@ public extension CodexReviewStore { externalURLOpener: @escaping @MainActor @Sendable (URL) -> Void = defaultExternalURLOpener, mcpPortOwnerResolver: CodexReviewMCPPortOwnerResolver? = nil, mcpHTTPServerBindChecker: CodexReviewMCPHTTPServerBindChecker? = nil, - shutdownCleanupTimeout: Duration = .seconds(2), + mcpLifecycleCallObserver: CodexReviewMCPLifecycleCallObserver? = nil, networkMonitor: any CodexReviewNetworkMonitoring = SystemCodexReviewNetworkMonitor(), networkRecoveryPolicy: CodexReviewNetworkRecoveryPolicy = .default, + reviewRuntimeClosePolicy: ReviewRuntimeClosePolicy = .production, + idGenerator: CodexReviewIDGenerator = .init(), + reviewRecoveryRoutingObserver: CodexReviewLiveRecoveryRoutingObserver? = nil, transport: any JSONRPC.Transport ) -> CodexReviewStore { makeLiveStoreForTesting( @@ -137,9 +119,12 @@ public extension CodexReviewStore { externalURLOpener: externalURLOpener, mcpPortOwnerResolver: mcpPortOwnerResolver, mcpHTTPServerBindChecker: mcpHTTPServerBindChecker, - shutdownCleanupTimeout: shutdownCleanupTimeout, + mcpLifecycleCallObserver: mcpLifecycleCallObserver, networkMonitor: networkMonitor, networkRecoveryPolicy: networkRecoveryPolicy, + reviewRuntimeClosePolicy: reviewRuntimeClosePolicy, + idGenerator: idGenerator, + reviewRecoveryRoutingObserver: reviewRecoveryRoutingObserver, transportFactory: { _ in transport } ) } @@ -156,9 +141,12 @@ public extension CodexReviewStore { ) -> any CodexReviewMCPHTTPServing)? = nil, mcpPortOwnerResolver: CodexReviewMCPPortOwnerResolver? = nil, mcpHTTPServerBindChecker: CodexReviewMCPHTTPServerBindChecker? = nil, - shutdownCleanupTimeout: Duration = .seconds(2), + mcpLifecycleCallObserver: CodexReviewMCPLifecycleCallObserver? = nil, 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( @@ -171,7 +159,8 @@ public extension CodexReviewStore { mcpHTTPServerFactory: mcpHTTPServerFactory, mcpPortOwnerResolver: mcpPortOwnerResolver, mcpHTTPServerBindChecker: mcpHTTPServerBindChecker, - shutdownCleanupTimeout: shutdownCleanupTimeout, + mcpLifecycleCallObserver: mcpLifecycleCallObserver, + reviewRecoveryRoutingObserver: reviewRecoveryRoutingObserver, appServerRuntimeFactory: { codexHomeURL in let client = AppServerClient(transport: try await transportFactory(codexHomeURL)) return .init( @@ -180,14 +169,43 @@ public extension CodexReviewStore { ) } ), + idGenerator: idGenerator, networkMonitor: networkMonitor, - networkRecoveryPolicy: networkRecoveryPolicy + networkRecoveryPolicy: networkRecoveryPolicy, + reviewRuntimeClosePolicy: reviewRuntimeClosePolicy ) } } @MainActor private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { + private struct AttemptRoute { + let generation: ReviewRuntimeGeneration + let runtime: LiveRuntimeLifecycleHandle + } + + 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 ( CodexReviewStore, CodexReviewMCPHTTPServer.Configuration @@ -196,8 +214,11 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { let seed: CodexReviewStoreSeed private var client: AppServerClient? - private var appServerBackend: AppServerCodexReviewBackend? - private var mcpHTTPServer: (any CodexReviewMCPHTTPServing)? + 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? private var loginClient: AppServerClient? @@ -206,21 +227,28 @@ 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 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 let reviewRecoveryRoutingObserver: CodexReviewLiveRecoveryRoutingObserver? private weak var attachedStore: CodexReviewStore? + private var appServerBackend: AppServerCodexReviewBackend? { + activeRuntimeHandle?.backend + } + init( environment: [String: String] = ProcessInfo.processInfo.environment, runtimePreferences: CodexReviewRuntime.Preferences = .defaults, @@ -235,7 +263,8 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { }, mcpPortOwnerResolver: CodexReviewMCPPortOwnerResolver? = nil, mcpHTTPServerBindChecker: CodexReviewMCPHTTPServerBindChecker? = nil, - shutdownCleanupTimeout: Duration = .seconds(2), + mcpLifecycleCallObserver: CodexReviewMCPLifecycleCallObserver? = nil, + reviewRecoveryRoutingObserver: CodexReviewLiveRecoveryRoutingObserver? = nil, appServerRuntimeFactory: AppServerRuntimeFactory? = nil ) { let runtimePreferences = runtimePreferences.normalized @@ -243,7 +272,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 @@ -251,10 +280,16 @@ 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 - self.shutdownCleanupTimeout = shutdownCleanupTimeout + self.reviewRecoveryRoutingObserver = reviewRecoveryRoutingObserver + let resolvedPortOwnerResolver = mcpPortOwnerResolver ?? Self.defaultMCPPortOwnerResolver + let resolvedBindChecker = mcpHTTPServerBindChecker ?? Self.defaultMCPHTTPServerBindChecker + self.mcpLifecycleOwner = LiveMCPServerLifecycleOwner( + configuration: mcpHTTPServerConfiguration, + factory: mcpHTTPServerFactory, + portOwnerResolver: resolvedPortOwnerResolver, + bindChecker: resolvedBindChecker, + lifecycleCallObserver: mcpLifecycleCallObserver + ) self.appServerRuntimeFactory = appServerRuntimeFactory ?? Self.makeAppServerRuntimeFactory( codexExecutablePath: runtimePreferences.codexExecutablePath ) @@ -272,8 +307,8 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend { client != nil } - var handlesActiveReviewStopCleanup: Bool { - true + var mcpServerLifecycle: any MCPServerLifecycleOwner { + mcpLifecycleOwner } var initialSettingsSnapshot: CodexReviewSettings.Snapshot { @@ -393,133 +428,138 @@ 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 - } - if forceRestartIfNeeded { - await stop(store: store) - } - - var startedClient: AppServerClient? - var startedHTTPServer: (any CodexReviewMCPHTTPServing)? + func prepareRuntime( + generation: ReviewRuntimeGeneration, + purpose _: ReviewRuntimeTransitionPurpose + ) async throws -> PreparedRuntime { + logger.info("Preparing review runtime") + let runtime = try await appServerRuntimeFactory(codexHomeURL) 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") + let authentication = try await runtime.backend.readAuth() + 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( + authentication: authentication, + settings: settings + ) + ) + logger.info("Review runtime prepared") + return PreparedRuntime(snapshot: handle.snapshot, handle: handle) } catch { - let failureMessage = await runtimeStartupFailureMessage(for: error) - logger.error("Review runtime failed to start: \(failureMessage, privacy: .public)") - await startedHTTPServer?.stop() - await startedClient?.close() - self.client = nil - self.appServerBackend = nil - self.mcpHTTPServer = nil - authNotificationTask?.cancel() - authNotificationTask = nil - store.transitionToFailed(failureMessage) + 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 } } - private func runtimeStartupFailureMessage(for error: Error) async -> String { - if let mcpHTTPServerError = error as? CodexReviewMCPHTTPServer.Error { - switch mcpHTTPServerError { - case .addressInUse: - return await mcpAddressInUseMessage() - } + func activateRuntime(_ handle: LiveRuntimeLifecycleHandle) throws { + guard activeRuntimeHandle == nil else { + throw ReviewLifecycleResourceFailure.client( + "A review runtime is already active." + ) } - return error.localizedDescription + guard let store = attachedStore else { + throw ReviewLifecycleResourceFailure.client( + "Review runtime activation requires its attached Store." + ) + } + activeRuntimeHandle = handle + acceptsRuntimeRequests = true + client = handle.client + settingsSnapshot = handle.snapshot.settings + observeAuthNotifications( + client: handle.client, + backend: handle.backend, + store: store + ) } - 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 deactivateRuntime( + _ handle: LiveRuntimeLifecycleHandle + ) -> Task? { + guard activeRuntimeHandle === handle else { + return nil } - message += ". Quit that process or change the MCP port in Settings, then reset the server." - return message + activeRuntimeHandle = nil + acceptsRuntimeRequests = false + client = nil + let task = authNotificationTask + authNotificationTask = nil + task?.cancel() + return task } - private func cancelActiveReviewsForRuntimeTeardown( - store: CodexReviewStore, - appServerBackend: AppServerCodexReviewBackend, - reason: ReviewCancellation, - timeoutWarning: String - ) async { - store.recordActiveReviewCancellationRequestsForRuntimeStop(reason: reason) - let didInterrupt = await runRuntimeShutdownCleanup(timeout: shutdownCleanupTimeout) { - await appServerBackend.interruptActiveReviewsForShutdown(reason: .init(message: reason.message)) - } - let locallyCancelledJobIDs = store.cancelActiveReviewsLocallyForRuntimeStop( - reason: reason, - cancelWorkers: false - ) - store.cancelAndDetachReviewWorkersForRuntimeStop(jobIDs: locallyCancelledJobIDs) - let didDrainReviewWorkers = await store.drainReviewWorkersForRuntimeStop( - timeout: shutdownCleanupTimeout - ) - if didInterrupt == false || didDrainReviewWorkers == false { - logger.warning("\(timeoutWarning, privacy: .public)") + func closeRuntimeAdmission(_ handle: LiveRuntimeLifecycleHandle) { + guard activeRuntimeHandle === handle else { + return } + acceptsRuntimeRequests = false } - func stop(store: CodexReviewStore) async { - let client = client + func stop(store: CodexReviewStore) async throws { 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 { + 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 let appServerBackend { - 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" + 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()) ) } - self.client = nil - self.mcpHTTPServer = nil - authNotificationTask?.cancel() - authNotificationTask = nil - await mcpHTTPServer?.stop() - self.appServerBackend = nil - await cleanupLoginRuntime(loginCleanup) - await client?.close() - logger.info("Review runtime stopped") } func waitUntilStopped() async {} @@ -613,28 +653,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 { @@ -643,13 +670,19 @@ 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 { guard auth.persistedAccounts.contains(where: { $0.accountKey == accountKey }) else { return } + let runtimeStore = appServerBackend == nil ? nil : attachedStore + if let runtimeStore { + try await runtimeStore.closeActiveReviewSessions( + reason: .system(message: "Account switched.") + ) + } try CodexReviewAccountRegistry.activateAccount( accountKey, accounts: auth.persistedAccounts, @@ -661,17 +694,21 @@ 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 runtimeStore.restart() } 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 { + try await runtimeStore.closeActiveReviewSessions( + reason: .system(message: "Account removed.") + ) + } if removedActiveAccount, let appServerBackend { _ = try? await appServerBackend.logout(.init(accountKey)) } @@ -698,12 +735,10 @@ 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 runtimeStore.restart() } } @@ -738,7 +773,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)) @@ -758,8 +793,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() } } @@ -845,10 +879,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, @@ -868,10 +907,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) } @@ -915,9 +952,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, @@ -952,9 +987,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, @@ -1026,52 +1059,221 @@ 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) } - func startReview(_ request: CodexReviewBackendModel.Review.Start) async throws -> BackendReviewAttempt { - guard let appServerBackend else { + func startReview( + _ request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission + ) async throws -> BackendReviewAttempt { + guard acceptsRuntimeRequests, let runtime = activeRuntimeHandle else { throw CodexReviewAPI.Error.io("Review runtime is not running.") } - return try await appServerBackend.startReview(request) - } - - func interruptReview(_ run: CodexReviewBackendModel.Review.Run, reason: CodexReviewBackendModel.CancellationReason) async throws { - guard let appServerBackend 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) + } } - try await appServerBackend.interruptReview(run, reason: reason) + let attempt = try await runtime.backend.startReview(request, admission: admission) + attemptRoutesByAttemptID[attempt.run.attemptID] = .init( + generation: runtime.generation, + runtime: runtime + ) + return attempt } - func beginReviewRecovery( + func interruptReview( _ run: CodexReviewBackendModel.Review.Run, + admission: ReviewStartAdmission, reason: CodexReviewBackendModel.CancellationReason - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { - guard let appServerBackend else { - throw CodexReviewAPI.Error.io("Review runtime is not running.") + ) async throws { + if let route = attemptRoutesByAttemptID[run.attemptID] { + try await route.runtime.backend.interruptReview(run, reason: reason) + return } - return try await appServerBackend.beginReviewRecovery(run, reason: reason) + 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 + } + throw CodexReviewAPI.Error.io("Review attempt route is unavailable.") + } + + func forceCloseReviewConnection() async throws { + guard acceptsRuntimeRequests, let appServerBackend else { + throw ReviewRuntimeCloseFailure.connection("Review runtime is not running.") + } + try await appServerBackend.forceCloseReviewConnection() + } + + func prepareReviewRecovery( + _ candidate: ReviewRecoveryCandidate + ) async throws -> ReviewRecoveryHandoff { + 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, + destination: nil + ) + return handoff } func resumeReviewRecovery( - _ token: CodexReviewBackendModel.Review.RecoveryToken, - request: CodexReviewBackendModel.Review.Start + _ handoff: ReviewRecoveryHandoff, + request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { - guard 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(token, request: request) + if handoff.candidate.trigger == .sameAccountRestart, + destination.generation == source.source.generation { + throw CodexReviewAPI.Error.io( + "Same-account recovery requires a replacement runtime generation." + ) + } + 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 cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async { - guard let appServerBackend else { - return + 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[recoveredRun.attemptID] = .init( + generation: destination.generation, + runtime: destination + ) + 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 { + 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) + } } - await appServerBackend.cleanupReview(run) + try await route.runtime.backend.cleanupReview(run) } @discardableResult @@ -1169,29 +1371,19 @@ 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)" - if let appServerBackend { - let reason = ReviewCancellation.system(message: message) - await cancelActiveReviewsForRuntimeTeardown( - store: store, - appServerBackend: appServerBackend, - reason: reason, - timeoutWarning: "Timed out cleaning active reviews after runtime failure" - ) - } let failedClient = client - let failedMCPHTTPServer = mcpHTTPServer - client = nil - appServerBackend = nil - mcpHTTPServer = nil + acceptsRuntimeRequests = false + if let authNotificationTask { + retiredAuthNotificationTasks.append(authNotificationTask) + } authNotificationTask = nil store.transitionToFailed(message) - await failedMCPHTTPServer?.stop() await cleanupLoginRuntime(loginCleanup) - await failedClient?.close() + await closeClientAfterFailure(failedClient) } private func handleAuthNotification( @@ -1216,7 +1408,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 @@ -1273,8 +1465,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( @@ -1286,8 +1477,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 } @@ -1321,8 +1511,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(), @@ -1352,8 +1541,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) } @@ -1537,40 +1725,119 @@ 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 + 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( @@ -1659,6 +1926,601 @@ 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 + + private weak var owner: LiveCodexReviewStoreBackend? + private var isActivated = false + private var closeTask: Task, Never>? + + 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 + } + + 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) + await backend.runtimeOwnerLifecycleHandle.closeAdmission() + } + + func close(purpose: ReviewRuntimeTransitionPurpose) async throws { + let task: Task, Never> + if let closeTask { + task = closeTask + } else { + let appServerLifecycle = backend.runtimeOwnerLifecycleHandle + let authObservationTask = owner?.deactivateRuntime(self) + let newTask = Task, Never> { @MainActor in + var failures: [ReviewLifecycleResourceFailure] = [] + authObservationTask?.cancel() + do { + try await appServerLifecycle.closeAndWait(purpose: purpose) + } catch { + 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 { + 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 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 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, + lifecycleCallObserver: CodexReviewMCPLifecycleCallObserver? + ) { + self.configuration = configuration + self.factory = factory + self.portOwnerResolver = portOwnerResolver + self.bindChecker = bindChecker + self.lifecycleCallObserver = lifecycleCallObserver + } + + func attachStore(_ store: CodexReviewStore) { + self.store = store + } + + func prepare() async throws -> PreparedMCPServer { + 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." + ) + case .closing, .closed: + throw ReviewLifecycleResourceFailure.mcpServer( + "MCP owner is closing or closed." + ) + } + + 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) + } + } + } + + func activate( + _ generation: MCPServerGeneration + ) async throws -> MCPServerPublicationSnapshot { + 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." + ) + } + + 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 { + 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 + } + } + + func drainAdmittedHandlers() async throws { + switch state { + case .running(let lease, _): + await lease.server?.waitForAdmittedHandlers() + case .stopped, .preparing, .prepared, .activating, .stopping, .closing, .closed: + return + } + } + + func stop() async throws { + stopCallerCount += 1 + lifecycleCallObserver?(.stop, stopCallerCount) + let operationID: UInt64 + let task: Task + switch state { + case .stopped: + return + 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 { + 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." + ) + } + } + + func close() async throws { + closeCallerCount += 1 + lifecycleCallObserver?(.close, closeCallerCount) + let operationID: UInt64 + let task: Task + switch state { + case .stopped: + 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 { + 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 { + 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(()) + } + } + + 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 { + 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..3cddb5ec 100644 --- a/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift +++ b/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift @@ -2,17 +2,13 @@ import Darwin import Foundation import MCP import OSLog +import CodexReview @preconcurrency import NIOCore @preconcurrency import NIOHTTP1 @preconcurrency import NIOPosix private let logger = Logger(subsystem: "CodexReviewKit", category: "mcp-http") -private struct TrackedHTTPResponse { - var response: HTTPResponse - var streamCompletion: ActiveRequestCompletion? = nil -} - package extension CodexReviewMCPHTTPServer { enum Error: Swift.Error, LocalizedError, Equatable, Sendable { case addressInUse(host: String, port: Int) @@ -98,13 +94,256 @@ package extension CodexReviewMCPHTTPServer { } } +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 + } + } +} + +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 SessionContext { - let server: Server + 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 networkResources: MCPHTTPNetworkResourceOwner + let task: Task + var admissionClosed = false + + init( + id: UInt64, + networkResources: MCPHTTPNetworkResourceOwner, + task: Task + ) { + self.id = id + self.networkResources = networkResources + 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 + let networkResources: MCPHTTPNetworkResourceOwner + var listenerCloseTask: Task, Never>? + + init( + id: UInt64, + listener: any Channel, + eventLoopGroup: MultiThreadedEventLoopGroup, + cleanupTask: Task, + boundURL: URL, + networkResources: MCPHTTPNetworkResourceOwner + ) { + self.id = id + self.listener = listener + self.eventLoopGroup = eventLoopGroup + self.cleanupTask = cleanupTask + self.boundURL = boundURL + self.networkResources = networkResources + } + } + + private enum LifecycleState { + case stopped([ReviewLifecycleResourceFailure]) + case starting(StartingGeneration) + case running(RunningGeneration) + case stopping( + id: UInt64, + resources: RunningGeneration?, + task: Task<[ReviewLifecycleResourceFailure], Never> + ) + } + + 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 { @@ -117,11 +356,21 @@ 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 nextSessionOrdinal: UInt64 = 0 + private var pendingCloseFailures: [ReviewLifecycleResourceFailure] = [] + private let handlerEntryGate = MCPHTTPHandlerEntryGate() + private let startCompletionGate = MCPHTTPStartCompletionGate() + 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, @@ -132,7 +381,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 { @@ -154,6 +409,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 @@ -170,97 +426,499 @@ 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 + let networkResources = MCPHTTPNetworkResourceOwner() + let task = Task { [self] in + await performStartGeneration( + id: id, + networkResources: networkResources + ) + } + let operation = StartingGeneration( + id: id, + networkResources: networkResources, + 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, + networkResources: MCPHTTPNetworkResourceOwner + ) async -> StartingGenerationResult { + 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)) + guard let connection = networkResources.admitConnection(channel) else { + return channel.close(mode: .all) + } + return channel.pipeline.configureHTTPServerPipeline( + withPipeliningAssistance: false + ).flatMap { + channel.pipeline.addHandler(CodexReviewMCPHTTPHandler( + server: self, + entryGate: handlerEntryGate, + networkResources: networkResources, + connection: connection + )) } } .childChannelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1) .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() - self.eventLoopGroup = group - self.channel = channel + listener = channel + await startCompletionGate.waitIfNeeded() + try Task.checkCancellation() let actualPort = channel.localAddress?.port - boundURL = configuration.url(boundPort: actualPort) - cleanupTask = Task { [weak self] in - await self?.sessionCleanupLoop() + let cleanupTask = Task { [weak self] in + guard let self else { return } + await self.sessionCleanupLoop() } - logger.info("MCP Streamable HTTP server listening at \(self.url.absoluteString, privacy: .public)") + return .success(RunningGeneration( + id: id, + listener: channel, + eventLoopGroup: group, + cleanupTask: cleanupTask, + boundURL: configuration.url(boundPort: actualPort), + networkResources: networkResources + )) } catch { - try? await group.shutdownGracefully() - throw CodexReviewMCPHTTPServer.Error.classifyStartError( - error, - configuration: configuration + networkResources.closeAdmission() + 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) + } + } + let closingGeneration = networkResources.beginClosing(.serverStop) + await closingGeneration.waitUntilClosed() + 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() + } + pendingCloseFailures.removeAll(keepingCapacity: false) + admittedHandlerDrainDidBegin = false + 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 { + 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 + resources.networkResources.closeAdmission() + 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) } - package func stop() async { - cleanupTask?.cancel() - cleanupTask = nil - await closeAllSessions() - try? await channel?.close() - channel = nil - if let eventLoopGroup { - try? await eventLoopGroup.shutdownGracefully() + private func closeStartingAdmission(_ operation: StartingGeneration) { + if operation.admissionClosed == false { + operation.admissionClosed = true + operation.networkResources.closeAdmission() + startCompletionGate.recordAdmissionClosed() + } + operation.task.cancel() + } + + private func performStopGeneration( + _ resources: RunningGeneration + ) async -> [ReviewLifecycleResourceFailure] { + resources.networkResources.closeAdmission() + var failures = pendingCloseFailures + pendingCloseFailures.removeAll(keepingCapacity: false) + 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) + } + } + await resources.cleanupTask.value + 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( + resources.networkResources.snapshot().isQuiescent, + "The MCP network generation owner must be quiescent 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") + return failures } - package func handleHTTPRequest(_ request: HTTPRequest) async -> HTTPResponse { - await handleTrackedHTTPRequest(request).response + 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() + } } - fileprivate func handleTrackedHTTPRequest(_ request: HTTPRequest) async -> TrackedHTTPResponse { + 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 { + 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) + 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 = listenerCloseTask(resources) + switch await task.value { + case .success: + return nil + case .failure(let failure): + return failure + } + } + + 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 + ) { + guard case .running(let current) = lifecycleState, + current === resources else { + return + } + if pendingCloseFailures.contains(failure) == false { + pendingCloseFailures.append(failure) + } + } + + package func waitForAdmittedHandlers() async { + admittedHandlerDrainDidBegin = true + let startWaiters = admittedHandlerDrainStartWaiters + admittedHandlerDrainStartWaiters.removeAll(keepingCapacity: false) + for waiter in startWaiters { + waiter.resume() + } + 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 + } + await networkResources?.waitForAdmittedHandlingWorkToDrain() + } + + fileprivate func handleAdmittedHTTPRequest( + _ request: HTTPRequest, + operation: MCPHTTPNetworkResourceOwner.RequestOperation, + networkResources: MCPHTTPNetworkResourceOwner + ) async -> HTTPResponse { + await performHTTPRequest( + request, + operation: operation, + networkResources: networkResources + ) + } + + 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) - 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( @@ -268,39 +926,77 @@ 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( adapter: adapter, defaultSessionID: sessionID, clientSession: clientSession, - boundedReviewWaitDuration: configuration.boundedReviewWaitDuration + boundedReviewWaitDuration: configuration.boundedReviewWaitDuration, + 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) - 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)") ) } } @@ -309,87 +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() + 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 - ) -> (response: TrackedHTTPResponse, didFinishRequest: Bool) { - switch response { - case .stream(let stream, let headers): - let completion = ActiveRequestCompletion { - Task { - await self.finishActiveRequest(sessionID: sessionID) - } - } - let trackedStream = AsyncThrowingStream(bufferingPolicy: .unbounded) { continuation in - let heartbeatTask = makeStreamHeartbeatTask(continuation: continuation) - let task = Task { - defer { - heartbeatTask?.cancel() - completion.finish() - } - do { - for try await chunk in stream { - continuation.yield(chunk) - } - continuation.finish() - } catch { - continuation.finish(throwing: error) - } - } - continuation.onTermination = { _ in - heartbeatTask?.cancel() - task.cancel() - completion.finish() - } - } - return ( - .init(response: .stream(trackedStream, headers: headers), streamCompletion: 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 completeScheduledSessionClose( + _ sessionID: String, + expectedOrdinal: UInt64 + ) async { + guard sessions[sessionID]?.ordinal == expectedOrdinal else { return } + await closeSession(sessionID) } - private func finishActiveRequest(sessionID: String) { - if var session = sessions[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 - } - return Task { - 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() } + 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 } } @@ -408,7 +1109,96 @@ 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 { + switch lifecycleState { + case .running(let resources), + .stopping(_, let resources?, _): + resources.listenerCloseTask == nil && resources.listener.isActive + case .stopped, .starting, .stopping: + false + } + } + + package func holdNextNetworkHandlerEntryForTesting() async { + await handlerEntryGate.holdNextEntry() + } + + package func releaseNetworkHandlerEntryForTesting() async { + 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 + } + await withCheckedContinuation { continuation in + if admittedHandlerDrainDidBegin { + continuation.resume() + } else { + admittedHandlerDrainStartWaiters.append(continuation) + } + } + } + + func networkSnapshotForTesting() -> MCPHTTPNetworkResourceOwner.Snapshot { + currentNetworkResources()?.snapshot() + ?? .init(revision: 0, phase: .closed, connections: []) + } + + 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 { + 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 { @@ -417,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) @@ -515,27 +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 CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked Sendable { typealias InboundIn = HTTPServerRequestPart typealias OutboundOut = HTTPServerResponsePart @@ -566,14 +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 entryGate: MCPHTTPHandlerEntryGate + private let networkResources: MCPHTTPNetworkResourceOwner + private let connection: MCPHTTPNetworkResourceOwner.Connection private var requestState: RequestState? - private var activeStreamTask: Task? - private var activeStreamID: UUID? - private var activeStreamCompletion: ActiveRequestCompletion? - init(server: CodexReviewMCPHTTPServer) { + init( + server: CodexReviewMCPHTTPServer, + entryGate: MCPHTTPHandlerEntryGate, + networkResources: MCPHTTPNetworkResourceOwner, + connection: MCPHTTPNetworkResourceOwner.Connection + ) { self.server = server + self.entryGate = entryGate + self.networkResources = networkResources + self.connection = connection } func channelRead(context: ChannelHandlerContext, data: NIOAny) { @@ -591,10 +1374,30 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked return } requestState = nil + let metadata = makeRequestMetadata(from: state) + guard let operation = connection.admitRequest(metadata: metadata), + let httpReservation = operation.beginHTTPHandling() + else { + writeAdmissionClosedResponse( + version: state.head.version, + context: context + ) + return + } nonisolated(unsafe) let context = context - Task { - await handleRequest(state: state, context: context) + let task = Task { + defer { httpReservation.acknowledge() } + await entryGate.waitIfNeeded() + guard Task.isCancelled == false else { + return + } + await handleRequest( + state: state, + operation: operation, + context: context + ) } + httpReservation.install(task) } } @@ -604,13 +1407,13 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked } func channelInactive(context: ChannelHandlerContext) { - finishActiveStream() + connection.beginClosing(.peerClosed) context.fireChannelInactive() } func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) { if case ChannelEvent.inputClosed = event { - finishActiveStream() + connection.beginClosing(.peerClosed) context.close(promise: nil) return } @@ -618,42 +1421,86 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked } func errorCaught(context: ChannelHandlerContext, error: any Error) { - finishActiveStream() + connection.beginClosing(.transportFailure(error.localizedDescription)) context.close(promise: nil) } - private func finishActiveStream() { - activeStreamTask?.cancel() - activeStreamCompletion?.finish() - activeStreamTask = nil - activeStreamID = nil - activeStreamCompletion = 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.handleTrackedHTTPRequest(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 makeHTTPRequest(from state: RequestState) -> HTTPRequest { + 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, + 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 { @@ -669,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( @@ -679,90 +1527,204 @@ 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 streamID = UUID() - let streamTask = Task { - 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() { - try Task.checkCancellation() - try await writeResponseBody(chunk, context: context, eventLoop: eventLoop) + 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)) + } + } + 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 { - trackedResponse.streamCompletion?.finish() - return - } catch { - trackedResponse.streamCompletion?.finish() - logger.error("MCP SSE stream failed: \(error.localizedDescription, privacy: .public)") - } - guard Task.isCancelled == false else { - return - } - 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 + 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 + } + } + group.cancelAll() + events.continuation.finish() } - self.activeStreamTask?.cancel() - self.activeStreamCompletion?.finish() - self.activeStreamTask = streamTask - self.activeStreamID = streamID - self.activeStreamCompletion = trackedResponse.streamCompletion - context.read() - } - await streamTask.value - eventLoop.execute { - if self.activeStreamID == streamID { - self.activeStreamTask = nil - self.activeStreamID = nil - self.activeStreamCompletion = nil + 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 } + } catch { + return .failed(error.localizedDescription) } 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) - if let body { - var buffer = context.channel.allocator.buffer(capacity: body.count) - buffer.writeBytes(body) - context.write(self.wrapOutboundOut(.body(.byteBuffer(buffer))), promise: nil) - } - context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil) + 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) + return .responded + } catch { + return .failed(error.localizedDescription) } } } diff --git a/Sources/CodexReviewMCPServer/CodexReviewMCPProtocolServer.swift b/Sources/CodexReviewMCPServer/CodexReviewMCPProtocolServer.swift index e606b639..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? @@ -30,11 +103,13 @@ 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, + foreignLifetimeLease: MCPProtocolServerForeignLifetimeWaiter.Lease ) async -> Server { let server = Server( name: "codex_review", @@ -45,60 +120,173 @@ 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) - ) + 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, + 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 - ) + 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)], + 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 is CancellationError { + throw CancellationError() + } catch { + return .init( + content: [.text(text: error.localizedDescription, 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) { [foreignLifetimeLease] _ in + try await performMCPDomainWork( + networkResources: networkResources, + foreignLifetimeLease: foreignLifetimeLease + ) { + .init(resources: helpResources.map(\.resource)) + } + } + + 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(ListResources.self) { _ in - .init(resources: helpResources.map(\.resource)) + await server.withMethodHandler(ListResourceTemplates.self) { [foreignLifetimeLease] _ in + try await performMCPDomainWork( + networkResources: networkResources, + foreignLifetimeLease: foreignLifetimeLease + ) { + .init(templates: helpResourceTemplates) + } } - 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")]) + 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." + ) } - await server.withMethodHandler(ListResourceTemplates.self) { _ in - .init(templates: helpResourceTemplates) + let task = Task { + let value = try await operation() + try Task.checkCancellation() + return value } + reservation.install(task) - return server + 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 { diff --git a/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift b/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift new file mode 100644 index 00000000..71f304a9 --- /dev/null +++ b/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift @@ -0,0 +1,1234 @@ +import Foundation +@preconcurrency import NIOCore + +final class MCPHTTPNetworkResourceOwner: @unchecked Sendable { + 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 + } + } + + struct RequestMetadata: Equatable, Sendable { + let method: String + let path: String + let jsonRPCID: String? + } + + enum CloseCause: Equatable, Sendable { + case sdkCancellation + case peerClosed + case sessionClosed + case serverStop + case transportFailure(String) + } + + enum OperationResult: Equatable, Sendable { + case responded + case cancelled(CloseCause) + case failed(String) + } + + 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.operation = operation + } + + func install( + _ task: Task + ) { + operation?.installCancellation({ task.cancel() }, for: self) + } + + 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 + } + } + + var isPending: Bool { + switch self { + case .reserved, .running: + true + case .notAdmitted, .responseReady, .turnGranted, .completed, .closed: + false + } + } + } + + private struct Slots { + var httpHandler: WorkSlot = .notAdmitted + var domain: WorkSlot = .notAdmitted + var source: WorkSlot = .notAdmitted + var writer: WorkSlot = .notAdmitted + + 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 + } + } + + private struct Admitted { + var common: Common + var slots: Slots + } + + private struct Handling { + var common: Common + var slots: Slots + } + + private struct Responding { + var common: Common + var slots: Slots + } + + private struct Closing { + var common: Common + var slots: Slots + let cause: CloseCause? + var pending: Set + let terminalResult: OperationResult + } + + private enum State { + case admitted(Admitted) + case handling(Handling) + case responding(Responding) + case closing(Closing) + case closed(Common, OperationResult, Slots) + } + + 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() + 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() + notifyChanged() + return reservation + } + + func admitDomainWork() -> WorkReservation? { + let reservation: WorkReservation + lock.lock() + 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 nil + } + lock.unlock() + notifyChanged() + return reservation + } + + func beginResponding() -> Bool { + lock.lock() + guard case .handling(let handling) = state else { + lock.unlock() + return false + } + state = .responding(.init(common: handling.common, slots: handling.slots)) + lock.unlock() + notifyChanged() + return true + } + + func bindResponseSource() -> WorkReservation? { + reserve(.source) + } + + func markResponseSourceNotRequired() { + lock.lock() + guard case .responding(var responding) = state, + case .notAdmitted = responding.slots.source + else { + lock.unlock() + return + } + responding.slots.source = .closed + state = .responding(responding) + lock.unlock() + notifyChanged() + } + + fileprivate func markResponseReady() -> Bool { + lock.lock() + 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() + 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) + } + } + } + + fileprivate func grantWriterTurn() -> Bool { + let waiter: CheckedContinuation? + lock.lock() + guard case .responding(var responding) = state else { + lock.unlock() + 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 + } + } + + fileprivate var isResponseReady: Bool { + lock.lock() + 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 + } + lock.unlock() + return result + } + + func bindWriter() -> WorkReservation? { + lock.lock() + guard case .responding(var responding) = state, + case .turnGranted = responding.slots.writer + else { + lock.unlock() + return nil + } + let reservation = WorkReservation(kind: .writer, operation: self) + responding.slots.writer = .reserved(reservation.id) + state = .responding(responding) + lock.unlock() + notifyChanged() + return reservation + } + + func bindSession(_ sessionID: String) { + lock.lock() + 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() + notifyChanged() + } + + func beginClosing(_ cause: CloseCause) { + transitionToClosing(cause: cause, result: .cancelled(cause), cancelPending: true) + } + + func acknowledgeResponseEnd() { + transitionToClosing(cause: nil, result: .responded, cancelPending: false) + } + + 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() + } + } + } + + fileprivate func waitForHandlingWorkToDrain() async { + await withCheckedContinuation { continuation in + lock.lock() + if handlingIsPendingLocked() == false { + lock.unlock() + continuation.resume() + } else { + handlingWaiters.append(continuation) + lock.unlock() + } + } + } + + 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() + for waiter in waiters { waiter.resume() } + notifyChanged() + } + + 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() + } + + 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() + } + + 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 + } + + 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 + } + + 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() + 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() + 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 { + 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 finishClosed( + _ closed: (OperationResult, [CheckedContinuation])? + ) { + guard let (result, waiters) = closed else { return } + for waiter in waiters { waiter.resume(returning: result) } + connection?.operationDidClose(self) + } + + private func notifyChanged() { + connection?.operationDidChange() + } + + 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 + } + } + + 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 + } + + 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 + } + } + } + + final class Connection: @unchecked Sendable { + private struct OpenState { + var queue: [RequestOperation] + var writerOperationID: UUID? + } + + private struct ClosingState { + var queue: [RequestOperation] + var writerOperationID: UUID? + var closeAcknowledged: Bool + } + + private enum State { + case accepting(OpenState) + case admissionClosed(OpenState) + case closing(ClosingState) + case closed + } + + 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] = [] + + 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() + } + } + + func admitRequest(metadata: RequestMetadata) -> RequestOperation? { + lock.lock() + 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 + } + + func supplyResponse(for operation: RequestOperation) async -> Bool { + guard operation.markResponseReady() else { return false } + pump() + return await operation.waitForWriterTurn() + } + + 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() + for operation in operations { operation.beginClosing(cause) } + if shouldCloseChannel { channel.close(mode: .all, promise: nil) } + owner?.changed() + } + + func waitUntilClosed() async { + await withCheckedContinuation { continuation in + lock.lock() + if case .closed = state { + lock.unlock(); continuation.resume() + } else { + closeWaiters.append(continuation); lock.unlock() + } + } + } + + 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() + } + + 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 operation + } + + 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 = [] + } + lock.unlock() + return operations + } + + fileprivate func operationDidChange() { + pump() + owner?.changed() + } + + 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() } + } + + 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() } + ) + } + + private func pump() { + var operation: RequestOperation? + lock.lock() + 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() + } + } + + private func acknowledgeChannelClose() { + var operationsToClose: [RequestOperation] = [] + var didClose = false + var waiters: [CheckedContinuation] = [] + lock.lock() + 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() } + } + } + + 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() + 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() + changed() + return connection + } + + func resolve(_ token: OperationToken) -> RequestOperation? { + lock.lock() + 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() + 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 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 { + 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 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 snapshot() -> Snapshot { + lock.lock() + 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 .init( + revision: currentRevision, + phase: phase, + connections: connections.sorted { $0.ordinal < $1.ordinal }.map { $0.snapshot() } + ) + } + + func nextSnapshot(after priorRevision: UInt64) async -> Snapshot { + await withCheckedContinuation { continuation in + lock.lock() + if revision > priorRevision { + lock.unlock(); continuation.resume(returning: snapshot()) + } else { + snapshotWaiters.append(.init(revision: priorRevision, continuation: continuation)) + lock.unlock() + } + } + } + + fileprivate func connectionDidClose(_ connection: Connection) { + var waiters: [CheckedContinuation] = [] + lock.lock() + 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() + } + + 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 case .closed = state { lock.unlock(); continuation.resume() } + else { closeWaiters.append(continuation); lock.unlock() } + } + } + + private func connectionsSnapshot() -> [Connection] { + lock.lock() + 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 = [] + } + lock.unlock() + return connections + } +} diff --git a/Sources/CodexReviewTesting/TestSupport.swift b/Sources/CodexReviewTesting/TestSupport.swift index 2cc7413a..8296c527 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) @@ -176,8 +182,9 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { case logout(CodexReviewBackendModel.Account.ID) case startReview(CodexReviewBackendModel.Review.Start) case interruptReview(CodexReviewBackendModel.Review.Run, CodexReviewBackendModel.CancellationReason) - case beginReviewRecovery(CodexReviewBackendModel.Review.Run, CodexReviewBackendModel.CancellationReason) - case resumeReviewRecovery(CodexReviewBackendModel.Review.RecoveryToken, CodexReviewBackendModel.Review.Start) + case forceCloseReviewConnection + case prepareReviewRecovery(ReviewRecoveryCandidate) + case resumeReviewRecovery(ReviewRecoveryHandoff, CodexReviewBackendModel.Review.Start) case cleanupReview(CodexReviewBackendModel.Review.Run) } @@ -186,14 +193,23 @@ 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? + private var authReadFailureMessage: String? + private var cleanupReviewGate: AsyncGate? + private let cleanupReviewStartedGate = AsyncGate() private var interruptReviewGate: AsyncGate? private var interruptReviewWaiters: [UUID: CheckedContinuation] = [:] - private var beginReviewRecoveryWaiters: [UUID: CheckedContinuation] = [:] + private var matchingInterruptReviewWaiters: [UUID: MatchingInterruptWaiter] = [:] + private var prepareReviewRecoveryWaiters: [UUID: CheckedContinuation] = [:] + private var prepareReviewRecoveryGate: AsyncGate? 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] = [:] @@ -221,6 +237,7 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { self.settings = settings self.auth = auth self.nextRun = nextRun + self.queuedRuns = [nextRun] } package func recordedCommands() -> [Command] { @@ -239,18 +256,52 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { recoveryFailureMessage = message } + package func failCleanup(message: String) { + 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 } 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) { + prepareReviewRecoveryGate = gate } package func setNextRecoveredRun(_ run: CodexReviewBackendModel.Review.Run) { 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 { @@ -327,9 +378,46 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { } } - package func waitForBeginReviewRecovery() async { + 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 waitForPrepareReviewRecovery() async { if commands.contains(where: { - if case .beginReviewRecovery = $0 { + if case .prepareReviewRecovery = $0 { true } else { false @@ -341,7 +429,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 @@ -349,19 +437,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() } } @@ -422,6 +510,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 } @@ -447,17 +538,36 @@ 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 { + let run = queuedRuns.isEmpty ? nextRun : queuedRuns.removeFirst() + try await admission.admitThreadStartDispatch() 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: run.attemptID, + threadID: run.threadID, + reviewThreadID: run.threadID, + model: run.model + ) + await admission.recordPreparedThread(provisionalRun) + do { + try await admission.admitReviewStartDispatch(for: provisionalRun) + } catch { + commands.append(.cleanupReview(provisionalRun)) + throw error + } if let startReviewGate { await startReviewGate.wait() } - 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 { @@ -467,62 +577,119 @@ 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() } if let interruptFailureMessage { - throw FakeCodexReviewBackendError(message: interruptFailureMessage) + throw ReviewInterruptRequestFailure( + outcome: .rejected(code: nil, message: interruptFailureMessage) + ) } } - package func beginReviewRecovery( - _ run: CodexReviewBackendModel.Review.Run, - reason: CodexReviewBackendModel.CancellationReason - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { - commands.append(.beginReviewRecovery(run, reason)) - let waiters = Array(beginReviewRecoveryWaiters.values) - beginReviewRecoveryWaiters.removeAll(keepingCapacity: false) + 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(.ownerForcedConnectionClose( + .connection("Connection force-closed.") + )) + } + } + + package func releaseHeldOperationsForCleanup() async { + await startReviewGate?.open() + await interruptReviewGate?.open() + await prepareReviewRecoveryGate?.open() + await resumeReviewRecoveryGate?.open() + } + + package func prepareReviewRecovery( + _ candidate: ReviewRecoveryCandidate + ) async throws -> ReviewRecoveryHandoff { + commands.append(.prepareReviewRecovery(candidate)) + let waiters = Array(prepareReviewRecoveryWaiters.values) + prepareReviewRecoveryWaiters.removeAll(keepingCapacity: false) for waiter in waiters { waiter.resume() } - if let interruptReviewGate { - await interruptReviewGate.wait() - } - if let interruptFailureMessage { - throw FakeCodexReviewBackendError(message: interruptFailureMessage) + if let prepareReviewRecoveryGate { + await prepareReviewRecoveryGate.wait() + try Task.checkCancellation() } - 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, - request: CodexReviewBackendModel.Review.Start + _ handoff: ReviewRecoveryHandoff, + request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { - commands.append(.resumeReviewRecovery(token, request)) + let token = handoff.token + let run = token.interruptedRun + let recoveredRun = queuedRecoveredRuns.isEmpty == false + ? queuedRecoveredRuns.removeFirst() + : 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) + try await admission.admitReviewStartDispatch(for: provisionalRun) + commands.append(.resumeReviewRecovery(handoff, request)) let waiters = Array(resumeReviewRecoveryWaiters.values) resumeReviewRecoveryWaiters.removeAll(keepingCapacity: false) for waiter in waiters { waiter.resume() } if let resumeReviewRecoveryGate { - await resumeReviewRecoveryGate.wait() + if resumeReviewRecoveryIgnoresCancellation { + await resumeReviewRecoveryGate.waitIgnoringCancellation() + } else { + 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)) } - package func cleanupReview(_ run: CodexReviewBackendModel.Review.Run) async { + 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 + } } package func yield(_ event: CodexReviewBackendModel.Review.Event, for run: CodexReviewBackendModel.Review.Run? = nil) async { @@ -533,8 +700,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 { @@ -567,8 +737,12 @@ package actor FakeCodexReviewBackend: CodexReviewBackend { interruptReviewWaiters.removeValue(forKey: id)?.resume() } - private func cancelBeginReviewRecoveryWaiter(id: UUID) { - beginReviewRecoveryWaiters.removeValue(forKey: id)?.resume() + private func cancelMatchingInterruptReviewWaiter(id: UUID) { + matchingInterruptReviewWaiters.removeValue(forKey: id)?.continuation.resume() + } + + private func cancelPrepareReviewRecoveryWaiter(id: UUID) { + prepareReviewRecoveryWaiters.removeValue(forKey: id)?.resume() } private func cancelResumeReviewRecoveryWaiter(id: UUID) { @@ -601,7 +775,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 ) } @@ -637,7 +813,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 } } @@ -658,6 +835,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 { @@ -678,10 +915,270 @@ 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 } +@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 + package private(set) var closePurposes: [ReviewRuntimeTransitionPurpose] = [] + + private let onActivate: @MainActor @Sendable () -> Void + 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 (ReviewRuntimeTransitionPurpose) async -> Void = { _ in } + ) { + self.onActivate = onActivate + self.onClose = onClose + } + + package func activate() async throws { + activateCallCount += 1 + onActivate() + } + + 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 + 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 + } + try await task.value.get() + } + + package func waitUntilClosed() async throws { + waitUntilClosedCallCount += 1 + guard let closeTask else { + throw ReviewLifecycleResourceFailure.client( + "Testing runtime wait began before close." + ) + } + try await closeTask.value.get() + } +} + +@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 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] = [] + + private let serverURL: URL? + private var nextGeneration: UInt64 = 0 + 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 + } + + 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 { + guard isClosed == false else { + throw ReviewLifecycleResourceFailure.mcpServer("Testing MCP owner is closed.") + } + 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 { + closeAdmissionCallCount += 1 + } + + package func holdHandlerDrain(with gate: AsyncGate) { + drainGate = gate + drainStartedGate = AsyncGate() + } + + package func waitForHandlerDrain() async { + await drainStartedGate.wait() + } + + 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 + } + + package func waitUntilStopped() async throws { + waitUntilStoppedCallCount += 1 + } + + 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 { + 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 package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { package let reviewBackend: FakeCodexReviewBackend @@ -689,13 +1186,27 @@ 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?] = [] + package let mcpServerLifecycle: any MCPServerLifecycleOwner + package private(set) var lastPreparedRuntimeHandle: TestingRuntimeLifecycleHandle? + private weak var store: CodexReviewStore? + 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, - seed: CodexReviewStoreSeed = .init() + seed: CodexReviewStoreSeed = .init(), + mcpServerLifecycle: any MCPServerLifecycleOwner = NoMCPServerLifecycleOwner() ) { self.reviewBackend = reviewBackend self.seed = seed + self.mcpServerLifecycle = mcpServerLifecycle self.currentSettingsSnapshot = seed.initialSettingsSnapshot } @@ -703,12 +1214,73 @@ package final class TestingCodexReviewStoreBackend: CodexReviewStoreBackend { currentSettingsSnapshot } - package func attachStore(_: CodexReviewStore) {} + package func attachStore(_ store: CodexReviewStore) { + self.store = store + } + + package func holdRuntimePreparation(with gate: AsyncGate) { + runtimePreparationGate = gate + runtimePreparationStartedGate = AsyncGate() + runtimePreparationCancellationGate = AsyncGate() + } - package func start(store: CodexReviewStore, forceRestartIfNeeded: Bool) async { - startRequests.append(forceRestartIfNeeded) - isActive = true - store.transitionToRunning(serverURL: nil) + package func waitForRuntimePreparation() async { + await runtimePreparationStartedGate.wait() + } + + package func waitForRuntimePreparationCancellation() async { + 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 + ) async throws -> PreparedRuntime { + startRequests.append(purpose == .restartSameAccount) + let handle = TestingRuntimeLifecycleHandle( + onActivate: { [weak self] in self?.isActive = true }, + onClose: { [weak self] purpose in + guard let self else { return } + self.isActive = false + if purpose == .recoveryReplacement || purpose == .restartSameAccount { + try? await self.reviewBackend.forceCloseReviewConnection() + } + await self.runtimeCloseOperation() + } + ) + 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 { @@ -718,6 +1290,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 @@ -814,45 +1391,79 @@ 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 { + reviewStartOwnershipSnapshots.append(attemptSnapshot( + from: store?.reviewAttemptOwnerships[request.jobID] + )) + return try await reviewBackend.startReview(request, admission: admission) } package func interruptReview( _ run: CodexReviewBackendModel.Review.Run, + admission _: ReviewStartAdmission, reason: CodexReviewBackendModel.CancellationReason ) async throws { try await reviewBackend.interruptReview(run, reason: reason) } - package func beginReviewRecovery( - _ run: CodexReviewBackendModel.Review.Run, - reason: CodexReviewBackendModel.CancellationReason - ) async throws -> CodexReviewBackendModel.Review.RecoveryToken { - try await reviewBackend.beginReviewRecovery(run, reason: reason) + package func forceCloseReviewConnection() async throws { + try await reviewBackend.forceCloseReviewConnection() + } + + package func prepareReviewRecovery( + _ candidate: ReviewRecoveryCandidate + ) async throws -> ReviewRecoveryHandoff { + try await reviewBackend.prepareReviewRecovery(candidate) } package func resumeReviewRecovery( - _ token: CodexReviewBackendModel.Review.RecoveryToken, - request: CodexReviewBackendModel.Review.Start + _ handoff: ReviewRecoveryHandoff, + request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission ) async throws -> BackendReviewAttempt { - try await reviewBackend.resumeReviewRecovery(token, request: request) + recoveryResumeOwnershipSnapshots.append(attemptSnapshot( + from: store?.reviewAttemptOwnerships[request.jobID] + )) + return try await reviewBackend.resumeReviewRecovery( + handoff, + request: request, + admission: admission + ) + } + + 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 { - 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 { + 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( @@ -912,9 +1523,16 @@ 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]] = [:] + private var deliveredResponseMethods: [String: Int] = [:] private var notificationStreamCountWaiters: [(Int, CheckedContinuation)] = [] private var closed = false + private var closeFailure: ReviewRuntimeCloseFailure? + private var closeCallCount = 0 + private var closeWaiters: [CheckedContinuation] = [] package init(responses: [String: [Data]] = [:]) { self.responses = responses @@ -948,12 +1566,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, @@ -961,9 +1584,16 @@ 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 { + throw JSONRPC.Error.closed + } + deliveredResponseMethods[request.method, default: 0] += 1 + resumeResponseMethodWaiters(request.method) if let queuedResponse { switch queuedResponse { case .success(let data): @@ -1004,12 +1634,27 @@ package actor FakeJSONRPCTransport: JSONRPC.Transport { } } - package func close() async { + 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 } + + activeRequestGates.values + 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 +1681,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 +1732,23 @@ package actor FakeJSONRPCTransport: JSONRPC.Transport { closed } + package func closeCallCountForTesting() -> Int { + 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 } @@ -1090,6 +1778,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..e584faff 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,16 +27,92 @@ 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 cleanupReview(_ attempt: BackendReviewAttempt) async { - await cleanupReview(attempt.run) + func prepareReviewRecovery( + _ run: CodexReviewBackendModel.Review.Run, + reason: CodexReviewBackendModel.CancellationReason + ) async throws -> ReviewRecoveryHandoff { + let candidate = try await makeRecoveryCandidate(backend: self, for: run, reason: reason) + return try await prepareReviewRecovery(candidate) + } + + func resumeReviewRecovery( + _ handoff: ReviewRecoveryHandoff, + request: CodexReviewBackendModel.Review.Start + ) async throws -> BackendReviewAttempt { + let admission = ReviewStartAdmission() + let registered = try await admission.registerStart { admission in + try await self.resumeReviewRecovery( + handoff, + request: request, + admission: admission + ) + } + try await admission.activateStart(registered.id) + return try await registered.task.value } + + func cleanupReview(_ attempt: BackendReviewAttempt) async throws { + try await cleanupReview(attempt.run) + } +} + +private func makeRecoveryCandidate( + backend: AppServerCodexReviewBackend, + for run: CodexReviewBackendModel.Review.Run, + reason: CodexReviewBackendModel.CancellationReason +) async throws -> ReviewRecoveryCandidate { + let admission = ReviewStartAdmission() + let registered = try await admission.registerStart { admission in + try await admission.admitThreadStartDispatch() + let provisionalRun = CodexReviewBackendModel.Review.Run( + attemptID: run.attemptID, + threadID: run.threadID, + reviewThreadID: run.threadID, + model: run.model + ) + await admission.recordPreparedThread(provisionalRun) + try await admission.admitReviewStartDispatch(for: provisionalRun) + await admission.recordActiveRun(run) + return .init(run: run) + } + try await admission.activateStart(registered.id) + _ = try await registered.task.value + + let requestOutcome = RecoveryRequestOutcomeProbe() + let recovery = Task { + try await admission.beginRecovery( + trigger: .recoverableNetworkLoss, + 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(.server(message: reason.message)), + for: run + ) + } + 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 { @@ -80,6 +156,205 @@ 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 + ) +} + +private func controlledReviewClosePolicy(gate: AsyncGate) -> ReviewRuntimeClosePolicy { + ReviewRuntimeClosePolicy(terminalGrace: .seconds(10)) { _ in + await gate.wait() + try Task.checkCancellation() + } +} + +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 { @@ -330,7 +605,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 @@ -339,8 +614,136 @@ 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() + } + } + + @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) + + let notifications = await transport.notificationStream() + var iterator = notifications.makeAsyncIterator() + await #expect(throws: JSONRPC.Error.transportTerminated(.processFailure( + failure.localizedDescription + ))) { + _ = try await iterator.next() + } + } + + @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)") + 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 + } } } @@ -376,7 +779,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") @@ -418,7 +821,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))) } } @@ -450,7 +855,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 +908,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 +1073,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 +1082,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 { @@ -1024,6 +1430,170 @@ 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 registeredStart = try await admission.registerStart { admission in + try await backend.startReview( + .init( + jobID: "job-1", + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ), + admission: admission + ) + } + try await admission.activateStart(registeredStart.id) + 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.") } + ) + } + #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) + await initializeGate.open() + + #expect(try await cancellation.value.terminal == .localCancellation( + .mcpClient(message: "Stop") + )) + await #expect(throws: ReviewStartCancelledBeforeDispatch.self) { + try await registeredStart.task.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 registeredStart = try await admission.registerStart { admission in + try await backend.startReview( + .init( + jobID: "job-1", + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ), + admission: admission + ) + } + try await admission.activateStart(registeredStart.id) + 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.") } + ) + } + #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) + await threadGate.open() + + #expect(try await cancellation.value.terminal == .localCancellation( + .mcpClient(message: "Stop") + )) + await #expect(throws: ReviewStartCancelledBeforeDispatch.self) { + try await registeredStart.task.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 registeredStart = try await admission.registerStart { admission in + try await backend.startReview( + .init( + jobID: "job-1", + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ), + admission: admission + ) + } + try await admission.activateStart(registeredStart.id) + 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() + } + ) + } + #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) + await reviewGate.open() + let attempt = try await registeredStart.task.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( + .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) @@ -1255,7 +1825,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 +1851,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 +1936,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) @@ -1385,6 +1970,472 @@ 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 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( + 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: ReviewLifecycleResourceFailureAggregate.self) { + try await close.value + } + await #expect(throws: ReviewLifecycleResourceFailureAggregate.self) { + 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) @@ -1406,10 +2457,75 @@ 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 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() } - await transport.close() } @Test func backendPreservesBufferedEventsBeforeNotificationStreamError() async throws { @@ -1432,7 +2548,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() } } @@ -1461,7 +2577,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 +2625,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 +2672,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) @@ -1692,10 +2823,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) @@ -1730,25 +2861,117 @@ struct AppServerClientTests { )) } - @Test func backendBuffersCancellationBeforeEventStreamRegistration() 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") + 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 run = try await backend.startReview(.init( + jobID: "job-1", + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + )) + + 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", + params: TestTurnNotification( + threadID: "thread-1", + turn: .init(id: "turn-1", status: "interrupted", error: .init(message: "Stop")) + ) + ) + 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 backendAdmitsRecoveryRollbackBeforeDispatch() 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") - 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 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") - let run = try await backend.startReview(.init( - jobID: "job-1", - sessionID: "session-1", - request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + #expect(await admission.currentPhase() == .recoveryRollbackOutcomeUnknown( + threadID: "review-thread-1" )) - try await backend.interruptReview(run, reason: .init(message: "Stop")) - - var iterator = await eventSequence(backend, run).makeAsyncIterator() - #expect(try await iterator.next() == .cancelled("Stop")) - #expect(try await iterator.next() == nil) + await rollbackGate.open() + let recovered = try await registered.task.value + #expect(recovered.turnID == "turn-2") } @Test func backendRecoverReviewRollsBackAndRestartsSameThread() async throws { @@ -1874,7 +3097,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( @@ -1998,7 +3221,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.") ) @@ -2059,7 +3282,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.") ) @@ -2096,7 +3319,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.") ) @@ -2135,7 +3358,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.") ) @@ -2145,16 +3368,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 +3393,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 +3403,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 +3412,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: ReviewInterruptRequestFailure(outcome: .rejected( + code: -32602, + message: "expected active turn id turn-old but found turn-active" + ))) { + try await backend.prepareReviewRecovery( + 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,12 +3475,10 @@ struct AppServerClientTests { )) } - @Test func backendIgnoresStaleTerminalWhileRecoveryInterruptIsInFlight() 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)) @@ -2319,18 +3489,15 @@ struct AppServerClientTests { model: "gpt-5" ) let initialEvents = await eventSequence(backend, run) + 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" } } @@ -2343,34 +3510,12 @@ 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 - #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 { @@ -2437,71 +3582,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 +3685,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 +3707,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 +4110,7 @@ struct AppServerClientTests { #expect(try await iterator.next() == nil) } - @Test func backendInterruptFinishesReviewEventStream() 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") @@ -3048,7 +4132,22 @@ 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")) + 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", + params: TestTurnNotification( + threadID: "thread-1", + turn: .init(id: "turn-1", status: "interrupted", error: .init(message: "Stop")) + ) + ) + try await cancellation.value #expect(try await iterator.next() == .cancelled("Stop")) #expect(try await iterator.next() == nil) @@ -3098,7 +4197,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 +4283,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 +4410,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 +4599,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 +5690,7 @@ struct AppServerClientTests { delta: "tail output\n" ) ) - await transport.close() + try await transport.close() #expect(try await iterator.next() == .logEntry( kind: .commandOutput, @@ -4572,7 +5699,9 @@ struct AppServerClientTests { replacesGroup: false, metadata: .init(sourceType: "commandExecution", title: "Command output", itemID: "cmd-1") )) - #expect(try await iterator.next() == nil) + await #expect(throws: ReviewAttemptStreamFailure.self) { + _ = try await iterator.next() + } } @Test func backendReviewExitCompletesMissingCommandCompletion() async throws { @@ -5288,6 +6417,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")) { @@ -5744,3 +6877,45 @@ 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 + } +} + +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/CodexReviewAppServerTests/CurrentV2ReviewContractTests.swift b/Tests/CodexReviewAppServerTests/CurrentV2ReviewContractTests.swift index 52aaea7b..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" @@ -1599,9 +1603,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 +1623,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 +1643,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 +1770,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 +1877,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/CodexReviewHostTests/CodexReviewHostTests.swift b/Tests/CodexReviewHostTests/CodexReviewHostTests.swift index 680fe822..86a6a509 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) } @@ -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: " ", @@ -166,73 +281,840 @@ struct CodexReviewHostTests { )) } - @Test func liveStoreUsesRuntimePreferenceCodexHome() async throws { + @Test func liveStoreUsesRuntimePreferenceCodexHome() async throws { + let homeURL = try temporaryHome() + let configuredCodexHomeURL = homeURL.appendingPathComponent("custom-codex-home", isDirectory: true) + 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 store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + runtimePreferences: .init(codexHomePath: configuredCodexHomeURL.path), + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + transportFactory: { codexHomeURL in + #expect(codexHomeURL == configuredCodexHomeURL) + return transport + } + ) + + await store.start(forceRestartIfNeeded: true) + + #expect(store.serverState == .running) + await store.stop() + } + + @Test func liveStorePassesRuntimePreferenceMCPPortAndPathToHTTPServerFactory() 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") + var capturedConfiguration: CodexReviewMCPHTTPServer.Configuration? + let store = CodexReviewStore.makeLiveStoreForTesting( + environment: ["HOME": homeURL.path], + runtimePreferences: .init( + mcpPort: 54321, + mcpPath: "custom-mcp" + ), + webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, + mcpHTTPServerFactory: { store, configuration in + capturedConfiguration = configuration + return CodexReviewMCPHTTPServer( + adapter: CodexReviewMCPServer(store: store), + configuration: .init( + host: configuration.host, + port: 0, + endpoint: configuration.endpoint + ) + ) + }, + mcpHTTPServerBindChecker: { _ in }, + transportFactory: { _ in transport } + ) + + await store.start(forceRestartIfNeeded: true) + let serverURL = try #require(store.serverURL) + + #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() + } + + @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 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 targetInterruptResponseGate = AsyncGate() + await firstTransport.holdNext( + method: "turn/interrupt", + gate: targetInterruptResponseGate + ) + + 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 graceGate = AsyncGate() + 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 await graceGate.wait() } + ), + 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 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", + 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() + 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 configuredCodexHomeURL = homeURL.appendingPathComponent("custom-codex-home", isDirectory: true) - 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" + let firstServer = ControlledMCPHTTPServer( + endpoint: try #require(URL(string: "http://127.0.0.1:19419/mcp")) ) - try await transport.enqueue(AppServerAPI.Model.List.Response(data: []), for: "model/list") + 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], - runtimePreferences: .init(codexHomePath: configuredCodexHomeURL.path), webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, - transportFactory: { codexHomeURL in - #expect(codexHomeURL == configuredCodexHomeURL) - return transport - } + 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() - await store.start(forceRestartIfNeeded: true) + 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) - #expect(store.serverState == .running) - await store.stop() + 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 liveStorePassesRuntimePreferenceMCPPortAndPathToHTTPServerFactory() async throws { + @Test func liveMCPOwnerCloseDuringActivationJoinsAndPreventsLatePublication() 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" + let server = ControlledMCPHTTPServer( + endpoint: try #require(URL(string: "http://127.0.0.1:19421/mcp")) ) - try await transport.enqueue(AppServerAPI.Model.List.Response(data: []), for: "model/list") - var capturedConfiguration: CodexReviewMCPHTTPServer.Configuration? + let lifecycleCalls = MCPLifecycleCallProbe() + let activationRelease = AsyncGate() + server.holdStart(with: activationRelease) + var factoryCallCount = 0 let store = CodexReviewStore.makeLiveStoreForTesting( environment: ["HOME": homeURL.path], - runtimePreferences: .init( - mcpPort: 54321, - mcpPath: "custom-mcp" - ), webAuthenticationSessionFactory: FakeWebAuthenticationSessions().makeSession, - mcpHTTPServerFactory: { store, configuration in - capturedConfiguration = configuration - return CodexReviewMCPHTTPServer( - adapter: CodexReviewMCPServer(store: store), - configuration: .init( - host: configuration.host, - port: 0, - endpoint: configuration.endpoint - ) - ) + mcpHTTPServerFactory: { _, _ in + factoryCallCount += 1 + return server }, mcpHTTPServerBindChecker: { _ in }, - transportFactory: { _ in transport } + 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() - await store.start(forceRestartIfNeeded: true) - let serverURL = try #require(store.serverURL) + 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) - #expect(capturedConfiguration?.port == 54321) - #expect(capturedConfiguration?.endpoint == "/custom-mcp") - #expect(serverURL.path == "/custom-mcp") - await store.stop() + 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 { @@ -1004,6 +1886,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 +1927,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) @@ -1053,6 +1952,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) @@ -1084,6 +2052,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 +2084,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 +2174,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 +2214,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 +2240,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 +2259,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 +2277,47 @@ 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: 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]) - let startedAt = Date() - await store.stop() - let elapsed = Date().timeIntervalSince(startedAt) - let resultBeforeRemoteCleanupUnblocked = try await waitForTaskValue(reviewRead, timeout: .seconds(1)) + 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() + #expect(await stopFinished.isCompleted() == false) await interruptGate.open() - let result = try #require(resultBeforeRemoteCleanupUnblocked) + 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 { @@ -1311,7 +2348,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 @@ -1325,11 +2361,28 @@ 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)) { + if case .waitingForRecovery = store.reviewAttemptOwnerships[jobID] { + true + } else { + false + } + }) let stopFinished = CompletionFlag() let stopTask = Task { @MainActor in @@ -1407,6 +2460,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( @@ -1463,6 +2517,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 { @@ -1981,6 +3045,121 @@ 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 { + 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 @@ -2058,3 +3237,113 @@ private actor CompletionFlag { completed } } + +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 { + 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) + ) + } +} diff --git a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift index 84f6c18e..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") @@ -102,46 +102,15 @@ 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) ) - 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: "review.local:9417", - 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:9417", - 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) - await server.stop() + try await server.start() + let sessionID = try await initializeSession(endpoint: await server.url) + #expect(sessionID.isEmpty == false) + await server.closeAdmission() + #expect(await server.listenerIsOpenForTesting() == false) + #expect((await server.networkSnapshotForTesting()).phase != .accepting) + try await server.stop() } @Test func streamableHTTPClassifiesAddressInUseBindError() { @@ -162,6 +131,73 @@ 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 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.networkSnapshotForTesting()).phase != .accepting) + + 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.networkSnapshotForTesting()).isQuiescent) + #expect(await recordedStopFailures(server) == expectedCleanupFailures) + } + @Test func streamableHTTPCallsReviewStartWithCustomTarget() async throws { let backend = FakeCodexReviewBackend() let store = CodexReviewStore.makeTestingStore( @@ -524,6 +560,53 @@ 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" + )] + ) + 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") + try await server.stop() + #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) == nil) + #expect((await server.networkSnapshotForTesting()).isQuiescent) + } + @Test func streamableHTTPScopesReviewReadToTransportSession() async throws { let backend = FakeCodexReviewBackend() let store = CodexReviewStore.makeTestingStore( @@ -756,23 +839,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 +866,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) @@ -857,6 +951,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, @@ -965,6 +1060,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, @@ -1046,6 +1142,326 @@ 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(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.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 waitForRequestOperation(on: server, jsonRPCID: "2") { snapshot in + snapshot.httpHandler == .running + } + + await server.releaseNetworkHandlerEntryForTesting() + await backend.waitForStartReview() + + let stopFinished = CompletionFlag() + let stopTask = Task { + try await server.stop() + await stopFinished.complete() + } + _ = await waitForRequestOperation(on: server, jsonRPCID: "2") { + $0.terminalCause == .serverStop + } + + #expect(await server.listenerIsOpenForTesting() == false) + #expect(await stopFinished.isCompleted() == false) + + await requestGate.open() + await backend.yield(.completed(summary: "Done", result: "review text")) + _ = try? await requestTask.value + try await stopTask.value + + #expect(await stopFinished.isCompleted()) + #expect(await server.sessionActiveRequestCountForTesting(sessionID: sessionID) == nil) + #expect((await server.networkSnapshotForTesting()).isQuiescent) + } + + @Test func stopFinishesSDKCancelledFinitePOSTWithoutNaturalTerminal() 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) + 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" }) + ) + 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 descriptor = try await openRawTCPConnection(endpoint: endpoint) + defer { Darwin.close(descriptor) } + + let firstBody = try makeJSONBody([ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": [ + "name": "review_start", + "arguments": [ + "cwd": "/tmp/project", + "target": ["type": "uncommittedChanges"], + ], + ], + ]) + let secondBody = try makeJSONBody([ + "jsonrpc": "2.0", + "id": 3, + "method": "tools/list", + ]) + try await sendRawPipelinedPOSTs( + descriptor: descriptor, + endpoint: endpoint, + sessionID: sessionID, + bodies: [firstBody, secondBody] + ) + + 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 firstResponseGate.open() + await backend.yield(.completed(summary: "Done", result: "review text")) + 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) + } + + #expect(responses.map(\.statusCode) == [200, 200]) + #expect(responseIDs == [2, 3]) + + try await server.stop() + #expect((await server.networkSnapshotForTesting()).isQuiescent) + } + + @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) + + 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 waitForNetworkSnapshot(on: server) { snapshot in + snapshot.connections.flatMap(\.operations).contains { + $0.metadata.method == "GET" && $0.writerIsRunning + } + } + try await server.stop() + _ = bytes + + #expect(await server.eventLoopGroupShutdownCountForTesting() == 1) + #expect((await server.networkSnapshotForTesting()).isQuiescent) + } + + @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 descriptor = try await openRawTCPConnection(endpoint: endpoint) + defer { Darwin.close(descriptor) } + _ = await waitForNetworkSnapshot(on: server) { + $0.connections.isEmpty == false + } + 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() + } + let restartFinished = CompletionFlag() + let restart = Task { + try await server.start() + await restartFinished.complete() + } + + 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) + #expect((await server.networkSnapshotForTesting()).connections.isEmpty) + + 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) + + try await openAndCloseRawEventStream(endpoint: endpoint, sessionID: sessionID) + _ = await waitForNetworkSnapshot(on: server) { snapshot in + snapshot.connections.flatMap(\.operations).contains { + $0.metadata.method == "GET" + } == false + } + try await server.stop() + + #expect(await server.eventLoopGroupShutdownCountForTesting() == 1) + #expect((await server.networkSnapshotForTesting()).isQuiescent) + } + @Test func streamableHTTPDoesNotExpireSessionWithOpenEventStream() async throws { let backend = FakeCodexReviewBackend() let store = CodexReviewStore.makeTestingStore( @@ -1181,8 +1597,13 @@ 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) + _ = try await store.awaitReview( + sessionID: sessionID, + jobID: running.id + ) #expect(response.statusCode == 200) #expect(running.core.lifecycle.status == .cancelled) @@ -1257,6 +1678,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), @@ -1271,10 +1705,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 } } @@ -1344,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...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 @@ -1423,6 +2049,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?, @@ -1480,6 +2146,89 @@ 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 struct MatchedRequestOperation { + let operation: MCPHTTPNetworkResourceOwner.OperationSnapshot + let connectionOperations: [MCPHTTPNetworkResourceOwner.OperationSnapshot] + + var admissionOrdinal: UInt64 { operation.admissionOrdinal } + var httpHandler: MCPHTTPNetworkResourceOwner.WorkStateSnapshot { + operation.httpHandler + } + var domainWorkIsPending: Bool { operation.domainWorkIsPending } + var responseIsReady: Bool { operation.responseIsReady } + var writerIsRunning: Bool { operation.writerIsRunning } + var terminalCause: MCPHTTPNetworkResourceOwner.CloseCause? { + operation.terminalCause + } + } + + private func waitForRequestOperation( + on server: CodexReviewMCPHTTPServer, + jsonRPCID: String, + satisfying condition: (MatchedRequestOperation) -> 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 { + private var completed = false + + func complete() { + completed = true + } + + func isCompleted() -> Bool { + completed + } } private nonisolated func currentPOSIXError() -> NSError { diff --git a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift index cd94bc1e..0ab4ca28 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) @@ -164,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) { @@ -732,11 +762,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 +810,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 +827,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() @@ -820,7 +886,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 @@ -830,7 +896,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 @@ -859,17 +925,332 @@ 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(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)) + + let commandsBeforeTerminal = await backend.recordedCommands() + #expect(commandsBeforeTerminal.contains { if case .prepareReviewRecovery = $0 { true } else { false } } == false) + #expect(isWaitingForRecovery(in: store, jobID: "job-1") == false) + + await backend.yield(.cancelled(reason.message), for: run) + try await backend.waitForPrepareReviewRecovery(timeout: .seconds(2)) + let barrierPublished = await waitUntil(timeout: .seconds(2)) { + isWaitingForRecovery(in: store, jobID: "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(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) + + let commandsBeforeAck = await backend.recordedCommands() + #expect(commandsBeforeAck.contains { if case .prepareReviewRecovery = $0 { true } else { false } } == false) + await interruptGate.open() + try await backend.waitForPrepareReviewRecovery(timeout: .seconds(2)) + let barrierPublished = await waitUntil(timeout: .seconds(2)) { + isWaitingForRecovery(in: store, jobID: "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(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)) + 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 .prepareReviewRecovery = $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(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( + .completed(summary: "Succeeded.", result: "natural review"), + for: run + ) + + let commandsBeforeAck = await backend.recordedCommands() + #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 .prepareReviewRecovery = $0 { true } else { false } } == false) + #expect(commands.contains { if case .resumeReviewRecovery = $0 { true } else { false } } == false) + } + } + + @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 cancellationDuringRecoveryPreparationPropagatesCleanupFailure() async throws { + let backend = FakeCodexReviewBackend() + await backend.failCleanup(message: "unsubscribe failed") + 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 + } + }) + + 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(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", @@ -901,11 +1282,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)) @@ -964,15 +1349,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)) @@ -1031,7 +1420,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 { @@ -1096,7 +1485,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)) @@ -1155,11 +1544,11 @@ 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 { - return run + if case .prepareReviewRecovery(let candidate) = command { + return candidate.resolved.run } return nil } @@ -1169,8 +1558,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 } @@ -1215,9 +1604,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 { @@ -1279,7 +1672,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)) @@ -1328,12 +1721,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 @@ -1362,8 +1754,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 }) @@ -1375,10 +1768,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)) @@ -1421,8 +1813,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)) @@ -1455,8 +1846,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 }) @@ -1467,18 +1859,51 @@ 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 backend.waitForBeginReviewRecovery(timeout: .seconds(2)) + try await completeNetworkRecoveryBarrier(backend: backend, store: store) 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) + 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", + 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( @@ -1516,13 +1941,13 @@ 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 { - 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") @@ -1533,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", @@ -1562,12 +1987,14 @@ 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) + await store.cancelAndAwaitReviewWorkersForRuntimeStop( + jobIDs: locallyCancelledJobIDs + ) #expect(store.reviewWorkerTasks["job-1"] == nil) - #expect(store.activeRuns["job-1"] == nil) + #expect(store.reviewAttemptOwnerships["job-1"] == nil) } } @@ -1601,6 +2028,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 @@ -1608,12 +2036,12 @@ 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) } } - @Test func runtimeStopDetachesNetworkRecoveryWaitingWorker() async throws { + @Test func runtimeStopAwaitsNetworkRecoveryWaitingWorker() async throws { let run = CodexReviewBackendModel.Review.Run( threadID: "thread-1", turnID: "turn-1", @@ -1636,22 +2064,23 @@ 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( reason: .system(message: "Review runtime stopped."), cancelWorkers: false ) - store.cancelAndDetachReviewWorkersForRuntimeStop(jobIDs: locallyCancelledJobIDs) + await store.cancelAndAwaitReviewWorkersForRuntimeStop( + jobIDs: locallyCancelledJobIDs + ) #expect(store.reviewWorkerTasks["job-1"] == nil) - #expect(store.activeRuns["job-1"] == nil) - #expect(store.reviewRecoveryWaitingJobIDs.contains("job-1") == false) + #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", @@ -1675,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) @@ -1704,15 +2134,19 @@ struct CodexReviewStoreCommandTests { reason: .system(message: "Review runtime stopped."), cancelWorkers: false ) - store.cancelAndDetachReviewWorkersForRuntimeStop(jobIDs: locallyCancelledJobIDs) - let resultBeforeStartReviewUnblocked = try await waitForTaskValue(running, timeout: .seconds(1)) + let drainTask = Task { @MainActor in + await store.cancelAndAwaitReviewWorkersForRuntimeStop( + jobIDs: locallyCancelledJobIDs + ) + } await startReviewGate.open() - let result = try #require(resultBeforeStartReviewUnblocked) + await drainTask.value + let result = try await running.value #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) } } @@ -1738,9 +2172,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) @@ -1778,9 +2211,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)) @@ -1814,15 +2251,20 @@ 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) let commands = await backend.recordedCommands() #expect(commands.contains { command in - if case .beginReviewRecovery = command { + if case .prepareReviewRecovery = command { true } else { false @@ -1855,7 +2297,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 @@ -1906,6 +2348,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", @@ -1941,6 +2384,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( @@ -1964,7 +2411,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 +2423,63 @@ 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: .unexpectedConnection( + .connection("Review event stream closed.") + )) + let cancel = try await cancellation + let read = try await result + + #expect(cancel.cancelled == false) + #expect(read.core.lifecycle.status == .failed) + #expect(read.core.lifecycle.terminal?.kind == .interrupted) + } + } + + @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() + } + ) + ) + await store.start() + 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.finishEvents(throwing: StreamClosedError()) + 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.output.summary == "Stop") + #expect(read.core.lifecycle.terminal == .interrupted(.requested( + .mcpClient(message: "Stop") + ))) + #expect(await backend.recordedCommands().contains(.forceCloseReviewConnection)) } } @@ -2001,7 +2496,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) @@ -2051,14 +2548,20 @@ struct CodexReviewStoreCommandTests { networkMonitor.yield(.init(status: .unsatisfied)) await outageSleepStarted.wait() - await backend.finishEvents(throwing: StreamClosedError(), for: initialRun) + await backend.finishEvents( + throwing: .recoverableNetwork(.connection("Verified network loss.")), + for: initialRun + ) let failedBeforeOutageConfirmed = await StoreSnapshotProbe(store: store) .waitUntilJobStatus(.failed, jobID: "job-1", timeout: .milliseconds(100)) != nil #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)) @@ -2071,7 +2574,53 @@ struct CodexReviewStoreCommandTests { } } - @Test func reviewStartCancellationInterruptsBackendRun() async throws { + @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( backend: TestingCodexReviewStoreBackend(reviewBackend: backend), @@ -2083,19 +2632,18 @@ 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 == .cancelled) + #expect(read.core.lifecycle.status == .failed) + #expect(read.core.lifecycle.terminal?.kind == .failed) + #expect(read.core.lifecycle.errorMessage == ReviewAttemptStreamFailure.ownerCancellation.localizedDescription) 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 +2657,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 +2680,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 +2697,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 +2709,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 +2742,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 +2750,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 +2770,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) @@ -2274,8 +2832,9 @@ struct CodexReviewStoreCommandTests { workspaces: [.init(cwd: "/tmp/project")], jobs: [running] ) + 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( @@ -2290,6 +2849,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( @@ -2509,18 +3094,111 @@ 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)) { + activeAttempt(in: store, jobID: jobID) != nil + } + guard becameActive else { + throw FakeCodexReviewBackendError(message: "Review did not become active before recovery.") + } + guard let run = activeAttempt(in: store, jobID: jobID)?.run 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) + do { + try await backend.waitForPrepareReviewRecovery(timeout: .seconds(2)) + } catch { + 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)) { + isWaitingForRecovery(in: store, jobID: 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.waitForPrepareReviewRecovery(timeout: .seconds(2)) + let barrierPublished = await waitUntil(timeout: .seconds(2)) { + 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, store: CodexReviewStore ) async { + await backend.releaseHeldOperationsForCleanup() await backend.finishEventMailboxes() await store.cancelAndDrainReviewWorkersForTesting() await backend.finishEventMailboxes() } -private struct StreamClosedError: Error {} - private actor ControlledTestSleeper { private let gate: AsyncGate private var shouldBlock = false diff --git a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift new file mode 100644 index 00000000..96a9a2cb --- /dev/null +++ b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift @@ -0,0 +1,1793 @@ +import Foundation +import Testing +import CodexReview +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, + record: ReviewRuntimeTransitionRecord(), + sourceRuntime: nil, + recoveryReplacement: nil + ) + + 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( + 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) + } + + @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 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 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() + #expect(try store.readReview(jobID: "job-target").core.lifecycle.cancellation == .mcpClient( + message: "Stop target" + )) + 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) + 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) + } + + @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"))) + let sourceCloseGate = AsyncGate() + runtime.holdClose(with: sourceCloseGate) + _ = 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 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) + + #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 + } +} + +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() + } + } +} 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() + } } diff --git a/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift new file mode 100644 index 00000000..c7017cf5 --- /dev/null +++ b/Tests/CodexReviewTests/ReviewAttemptProcessorTests.swift @@ -0,0 +1,1234 @@ +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( + .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(.completed))) + await requestGate.open() + let resolution = try await cancellation.value + + #expect(resolution.terminal == .canonical(.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) + + 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( + .interrupted(.requested(.mcpClient(message: "Stop again"))) + )) + } + + @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(.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( + .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() + try await admission.recordStreamTerminal(.unexpectedConnection(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(.stream(.unexpectedConnection(connection)))) + } + + @Test func recoveryAckWaitsForInterruptedTerminalAndReturnsBarrier() async throws { + let (admission, run) = try await makeActiveAdmission() + let requestAccepted = InvocationProbe() + let recovery = Task { + try await admission.beginRecovery( + trigger: .recoverableNetworkLoss, + interrupt: { _, _ in await requestAccepted.record() }, + forceClose: {} + ) + } + await requestAccepted.waitForInvocation() + #expect(await admission.currentPhase() == .interrupting(run)) + try await admission.recordCanonicalTerminal( + .interrupted(.server(message: "network recovery")), + for: run + ) + + 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")) + )) + } + + @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.beginRecovery( + trigger: .recoverableNetworkLoss, + interrupt: { _, _ in + await requestFailed.record() + throw requestFailure + }, + forceClose: {} + ) + } + await requestFailed.waitForInvocation() + try await admission.recordStreamTerminal(.recoverableNetwork(connection)) + + 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.beginRecovery( + trigger: .recoverableNetworkLoss, + interrupt: { _, _ in await requestAccepted.record() }, + forceClose: {} + ) + } + await requestAccepted.waitForInvocation() + try await admission.recordCanonicalTerminal(.completed, for: run) + + 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 { + let (admission, run) = try await makeActiveAdmission() + let rejection = ReviewInterruptRequestFailure( + outcome: .rejected(code: -32_000, message: "Not active") + ) + + await #expect(throws: rejection) { + try await admission.beginRecovery( + trigger: .recoverableNetworkLoss, + interrupt: { _, _ in throw rejection }, + forceClose: {} + ) + } + #expect(await admission.currentPhase() == .active(run)) + } + + @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() + try await admission.recordStreamTerminal( + .ownerForcedConnectionClose(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( + .stream(.ownerForcedConnectionClose(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 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 + #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 { + 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(.completed))) + } + + @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: {} + ) + } + await requestStarted.waitForInvocation() + let second = Task { + try await admission.cancel( + .system(message: "Runtime stop"), + interrupt: { _, _ in + Issue.record("Duplicate caller installed a second interrupt operation.") + }, + forceClose: {} + ) + } + try await admission.recordCanonicalTerminal( + .interrupted(.requested(.mcpClient(message: "Stop"))), + for: run + ) + await requestGate.open() + + #expect(try await first.value == second.value) + #expect(await admission.waitForInterruptionAdmission() == .terminalCancellation( + .mcpClient(message: "Stop") + )) + #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(.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(.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 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 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() + 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 + await #expect(throws: ReviewAttemptContractFailure.self) { + try await admission.activateStart(registered.id) + } + } + + @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")) + let duplicateCancellation = Task { + try await admission.cancel( + .system(message: "Runtime stop"), + interrupt: { _, _ in Issue.record("Duplicate joined cancellation sent a request.") }, + forceClose: {} + ) + } + 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, duplicateCancellation.value) + #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) + #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() + 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 = try await registerAndActivateStart(admission) { _ in + await startGate.waitIgnoringCancellation() + return .init(run: canonicalRun) + } + + try await admission.admitThreadStartDispatch() + await #expect(throws: ReviewAttemptContractFailure.self) { + try await admission.admitThreadStartDispatch() + } + + await startGate.open() + _ = try await startTask.value + } + + @Test func threadStartDispatchAdmissionAllowsVerifiedRejectionRetry() async throws { + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) + let startGate = AsyncGate() + let startTask = try await registerAndActivateStart(admission) { _ in + await startGate.waitIgnoringCancellation() + return .init(run: canonicalRun) + } + + try await admission.admitThreadStartDispatch() + try await admission.recordThreadStartRejectedForRetry() + try await admission.admitThreadStartDispatch() + await #expect(throws: ReviewAttemptContractFailure.self) { + try await admission.admitThreadStartDispatch() + } + + await startGate.open() + _ = try await startTask.value + } + + @Test func reviewStartDispatchAdmissionRejectsDirectDuplicate() async throws { + let admission = ReviewStartAdmission(closePolicy: controlledClosePolicy(gate: AsyncGate())) + let startGate = AsyncGate() + let startTask = try await registerAndActivateStart(admission) { _ in + await startGate.waitIgnoringCancellation() + return .init(run: canonicalRun) + } + try await admission.admitThreadStartDispatch() + await admission.recordPreparedThread(provisionalRun) + + try await admission.admitReviewStartDispatch(for: provisionalRun) + await #expect(throws: ReviewAttemptContractFailure.self) { + try await admission.admitReviewStartDispatch(for: provisionalRun) + } + + 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 = try await registerAndActivateStart(admission) { _ 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() + let dispatchGate = AsyncGate() + let startTask = try await registerAndActivateStart(admission) { admission in + await entered.record() + await dispatchGate.wait() + try Task.checkCancellation() + try await admission.admitThreadStartDispatch() + 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 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 + try await admission.admitThreadStartDispatch() + await threadDispatched.record() + await threadResponseGate.waitIgnoringCancellation() + let provisional = provisionalRun + await admission.recordPreparedThread(provisional) + try await admission.admitReviewStartDispatch(for: provisional) + 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: {} + ) + } + #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) + await threadResponseGate.open() + + #expect(try await cancellation.value.terminal == .localCancellation(.mcpClient(message: "Stop"))) + await #expect(throws: ReviewStartCancelledBeforeDispatch.self) { + try await startTask.value + } + } + + @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 = try await registerAndActivateStart(admission) { admission in + try 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() + 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 + } + } + + @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() + let reviewDispatchGate = AsyncGate() + let startTask = try await registerAndActivateStart(admission) { admission in + try await admission.admitThreadStartDispatch() + await admission.recordPreparedThread(provisionalRun) + await prepared.record() + await reviewDispatchGate.waitIgnoringCancellation() + try await admission.admitReviewStartDispatch(for: provisionalRun) + 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: {} + ) + } + #expect(await admission.waitForCancellationAdmission() == .mcpClient(message: "Stop")) + 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 = try await registerAndActivateStart(admission) { admission in + try await admission.admitThreadStartDispatch() + await admission.recordPreparedThread(provisionalRun) + try 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( + .interrupted(.requested(.mcpClient(message: "Stop"))) + )) + } + + @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 = try await registerAndActivateStart(admission) { admission in + try await admission.admitThreadStartDispatch() + await admission.recordPreparedThread(provisionalRun) + try 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() + 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 + } + } + + @Test func duplicateCleanupCallersJoinOneOwnedTask() async throws { + let (admission, run) = try await makeActiveAdmission() + let cleanupStarted = InvocationProbe() + let cleanupGate = AsyncGate() + + 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.value, second.value) + + #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 = try await registerAndActivateStart(admission) { admission in + try await admission.admitThreadStartDispatch() + await admission.recordPreparedThread(provisionalRun) + try await admission.admitReviewStartDispatch(for: provisionalRun) + await admission.recordActiveRun(canonicalRun) + return .init(run: canonicalRun) + } + _ = try await startTask.value + 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() + 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 } +} diff --git a/Tests/CodexReviewTests/ReviewTerminalContractTests.swift b/Tests/CodexReviewTests/ReviewTerminalContractTests.swift index c19c216d..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 streamEOFFailsWithoutAnAuthoritativeTerminal() async throws { + @Test func streamEOFUsesTypedWorkerContractFailureWithoutCanonicalTerminal() async throws { let (store, backend) = makeStore() async let started = store.startReview( sessionID: "session-1", @@ -192,6 +192,11 @@ struct ReviewTerminalContractTests { let result = try await started #expect(result.core.lifecycle.status == .failed) #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") == true) #expect(result.core.lifecycle.errorMessage?.contains("authoritative terminal") == true) } diff --git a/Tests/ReviewUITests/ReviewUITests.swift b/Tests/ReviewUITests/ReviewUITests.swift index dba17b32..77959081 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() @@ -1779,7 +1786,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 +1799,17 @@ 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.reviewAttemptOwnerships[job.id] = .active(.init( + run: run, + admission: admission + )) let viewController = ReviewMonitorSplitViewController(store: store, uiState: ReviewMonitorUIState(auth: store.auth)) viewController.loadViewIfNeeded() @@ -6695,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 { @@ -6732,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 } @@ -6775,19 +6789,19 @@ final class FailingCancellationBackend: PreviewCodexReviewStoreBackend { ) } - override func start( - store _: CodexReviewStore, - forceRestartIfNeeded _: Bool - ) async { - } - override func stop(store _: CodexReviewStore) async { } override func waitUntilStopped() async {} - override func interruptReview(_: CodexReviewBackendModel.Review.Run, reason _: CodexReviewBackendModel.CancellationReason) async throws { - throw CodexReviewAPI.Error.io("Cancellation failed.") + override func interruptReview( + _: CodexReviewBackendModel.Review.Run, + admission _: ReviewStartAdmission, + reason _: CodexReviewBackendModel.CancellationReason + ) async throws { + throw ReviewInterruptRequestFailure( + outcome: .rejected(code: nil, message: "Cancellation failed.") + ) } } @@ -6824,12 +6838,6 @@ final class BlockingSettingsBackend: PreviewCodexReviewStoreBackend { ) } - override func start( - store _: CodexReviewStore, - forceRestartIfNeeded _: Bool - ) async { - } - override func stop(store _: CodexReviewStore) async { } 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. 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": [ {