diff --git a/Sources/CodexReview/Store/CodexReviewStore.swift b/Sources/CodexReview/Store/CodexReviewStore.swift index fe07060..0e46eed 100644 --- a/Sources/CodexReview/Store/CodexReviewStore.swift +++ b/Sources/CodexReview/Store/CodexReviewStore.swift @@ -46,6 +46,7 @@ public final class CodexReviewStore { @ObservationIgnored package var reviewTerminalWaiters: [String: [ReviewTerminalWaiter]] = [:] @ObservationIgnored package var closedSessions: Set = [] @ObservationIgnored package var accountRateLimitAutoRefreshDriver: CodexReviewStoreRateLimitAutoRefreshDriver? + @ObservationIgnored package let storeWorkRegistry = ReviewStoreWorkRegistry() @ObservationIgnored package var runtimeState: ReviewStoreRuntimeState = .stopped( .init(rawValue: 0) ) @@ -90,6 +91,7 @@ public final class CodexReviewStore { isolated deinit { accountRateLimitAutoRefreshDriver?.cancel() + storeWorkRegistry.cancelWithoutWaiting() switch runtimeState { case .acquiring(_, _, let task), .replacing(_, _, _, let task), @@ -409,6 +411,98 @@ public final class CodexReviewStore { await backend.waitUntilStopped() } + package var storeWorkRegistryStatus: ReviewStoreWorkRegistryStatus { + storeWorkRegistry.status + } + + package func closeRegisteredStoreWork( + reason: ReviewCancellation + ) async -> ReviewStoreWorkDrainResult { + let operation = storeWorkRegistry.beginClosing { [self] in + recordActiveReviewCancellationRequestsForRuntimeStop(reason: reason) + accountRateLimitAutoRefreshDriver?.closeAdmission() + } + let result = await operation.task.value + await cancelAccountRateLimitAutoRefreshAndWait() + storeWorkRegistry.completeClosing(operation, result: result) + return result + } + + package func startRegisteredStoreWork( + kind: ReviewStoreWorkKind, + cancelledBeforeEntry: ReviewStoreWorkCancelledBeforeEntryPolicy = .skip, + operation: @escaping @MainActor @Sendable (CodexReviewStore) async -> Void + ) -> Task? { + guard let admission = storeWorkRegistry.register(kind) else { + return nil + } + let task = Task { @MainActor [weak self] in + defer { + self?.storeWorkRegistry.finish(admission) + } + guard let self else { + return + } + if Task.isCancelled || storeWorkRegistry.acceptsNewWork == false { + switch cancelledBeforeEntry { + case .skip: + return + case .runFinalizer(let finalizer): + finalizer(self) + return + } + } + await operation(self) + } + storeWorkRegistry.install(task, for: admission) + return task + } + + package func performRegisteredStoreWork( + kind: ReviewStoreWorkKind, + operation: @escaping @MainActor @Sendable (CodexReviewStore) async -> Void + ) async { + guard let task = startRegisteredStoreWork( + kind: kind, + operation: operation + ) else { + return + } + await withTaskCancellationHandler { + await task.value + } onCancel: { + task.cancel() + } + } + + package func performThrowingRegisteredStoreWork( + kind: ReviewStoreWorkKind, + operation: @escaping @MainActor @Sendable (CodexReviewStore) async throws -> Value + ) async throws -> Value { + guard let admission = storeWorkRegistry.register(kind) else { + throw CodexReviewAPI.Error.io("Review Store work admission is closed.") + } + let task = Task { @MainActor [weak self] in + guard let self else { + throw CancellationError() + } + try Task.checkCancellation() + if self.storeWorkRegistry.acceptsNewWork == false { + throw CancellationError() + } + return try await operation(self) + } + storeWorkRegistry.install(task, for: admission) + defer { + storeWorkRegistry.finish(admission) + } + return try await withTaskCancellationHandler { + try await task.value + } onCancel: { + task.cancel() + } + } + private func beginRuntimeAcquisition( generation: ReviewRuntimeGeneration, context: RuntimeAcquisitionContext = .init(), @@ -895,6 +989,9 @@ public final class CodexReviewStore { } package func requestSwitchAccount(_ account: CodexAccount, requiresConfirmation: Bool) { + guard storeWorkRegistry.acceptsNewWork else { + return + } auth.requestSwitchAccount(account, requiresConfirmation: requiresConfirmation) guard requiresConfirmation == false else { return @@ -903,6 +1000,9 @@ public final class CodexReviewStore { } package func requestSwitchAccountFromUserAction(_ account: CodexAccount) { + guard storeWorkRegistry.acceptsNewWork else { + return + } requestSwitchAccount( account, requiresConfirmation: hasRunningJobs @@ -911,6 +1011,9 @@ public final class CodexReviewStore { } package func requestSignOutActiveAccount(requiresConfirmation: Bool) { + guard storeWorkRegistry.acceptsNewWork else { + return + } auth.requestSignOutActiveAccount(requiresConfirmation: requiresConfirmation) guard requiresConfirmation == false else { return @@ -919,6 +1022,9 @@ public final class CodexReviewStore { } package func requestRemoveAccount(_ account: CodexAccount, requiresConfirmation: Bool) { + guard storeWorkRegistry.acceptsNewWork else { + return + } auth.requestRemoveAccount(account, requiresConfirmation: requiresConfirmation) guard requiresConfirmation == false else { return @@ -927,23 +1033,23 @@ public final class CodexReviewStore { } package func confirmPendingAccountAction() { + guard storeWorkRegistry.acceptsNewWork else { + return + } guard let action = auth.consumePendingAccountAction() else { return } - Task { @MainActor [weak self] in - guard let self else { - return - } + _ = startRegisteredStoreWork(kind: .accountAction) { @MainActor store in do { - try await self.executePendingAccountAction(action) - if let warningMessage = self.auth.warningMessage { - self.auth.presentAccountActionAlert( + try await store.executePendingAccountAction(action) + if let warningMessage = store.auth.warningMessage { + store.auth.presentAccountActionAlert( title: "Account Updated With Warning", message: warningMessage ) } } catch { - self.auth.presentAccountActionAlert( + store.auth.presentAccountActionAlert( title: action.failureTitle, message: error.localizedDescription ) @@ -952,10 +1058,16 @@ public final class CodexReviewStore { } package func cancelPendingAccountAction() { + guard storeWorkRegistry.acceptsNewWork else { + return + } auth.cancelPendingAccountAction() } package func dismissAccountActionAlert() { + guard storeWorkRegistry.acceptsNewWork else { + return + } auth.dismissAccountActionAlert() } diff --git a/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift b/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift index 3e29237..09f505b 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreCancellation.swift @@ -119,22 +119,33 @@ extension CodexReviewStore { let cancellation = ReviewCancellation.system( message: reason.nilIfEmpty ?? "Cancellation requested." ) + try await performThrowingRegisteredStoreWork( + kind: .reviewMutation("cancel-all") + ) { @MainActor store in + try await store.performCancelAllRunningJobs(cancellation: cancellation) + } + } + + private func performCancelAllRunningJobs( + cancellation: ReviewCancellation + ) async throws { let cancellableJobs = orderedJobs.filter { $0.isTerminal == false } var firstError: (any Error)? for job in cancellableJobs { do { - _ = try await cancelReview( + _ = try await performCancelReview( jobID: job.id, - sessionID: job.sessionID, cancellation: cancellation ) } catch { let message = error.localizedDescription.trimmingCharacters(in: .whitespacesAndNewlines) - try? recordCancellationFailure( - jobID: job.id, - sessionID: job.sessionID, - message: message.isEmpty ? "Failed to cancel review." : message - ) + if storeWorkRegistry.acceptsNewWork { + try? recordCancellationFailure( + jobID: job.id, + sessionID: job.sessionID, + message: message.isEmpty ? "Failed to cancel review." : message + ) + } if firstError == nil { firstError = error } diff --git a/Sources/CodexReview/Store/CodexReviewStoreRateLimitAutoRefresh.swift b/Sources/CodexReview/Store/CodexReviewStoreRateLimitAutoRefresh.swift index 560b289..30295b9 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 storeWorkRegistry.acceptsNewWork 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 storeWorkRegistry.acceptsNewWork else { + return + } startAccountRateLimitAutoRefresh() accountRateLimitAutoRefreshDriver?.refreshDueAccounts(now: now) } @@ -259,6 +273,7 @@ package final class CodexReviewStoreRateLimitAutoRefreshDriver { } private weak var store: CodexReviewStore? + private var admissionIsOpen = true private var observation: PortableObservationTracking.Token? private var scheduledWakeUp: ScheduledWakeUp? private var accountStates: [String: CodexReviewStoreRateLimitAutoRefreshAccountState] = [:] @@ -291,6 +306,9 @@ package final class CodexReviewStoreRateLimitAutoRefreshDriver { } func start() { + guard admissionIsOpen else { + return + } guard observation == nil else { syncLatestTargets() return @@ -306,6 +324,7 @@ package final class CodexReviewStoreRateLimitAutoRefreshDriver { } func cancel() { + admissionIsOpen = false observation?.cancel() observation = nil scheduledWakeUp?.task.cancel() @@ -316,15 +335,39 @@ package final class CodexReviewStoreRateLimitAutoRefreshDriver { accountStates.removeAll(keepingCapacity: false) } + func closeAdmission() { + admissionIsOpen = false + } + + func cancelAndWait() async { + admissionIsOpen = false + 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 { + guard admissionIsOpen, let store else { return } syncTargets(store.accountRateLimitAutoRefreshTargets(now: now), now: now) } private func syncLatestTargets() { - guard let store else { + guard admissionIsOpen, let store else { return } let now = store.clock.now() @@ -338,6 +381,9 @@ package final class CodexReviewStoreRateLimitAutoRefreshDriver { _ targets: [CodexReviewStoreRateLimitAutoRefreshTarget], now: Date ) { + guard admissionIsOpen else { + return + } let targets = scheduledTargets(from: targets, now: now) for target in targets { guard accountStates[target.accountKey]?.isRefreshing != true else { @@ -404,48 +450,64 @@ package final class CodexReviewStoreRateLimitAutoRefreshDriver { } scheduledWakeUp?.task.cancel() let delay = max(0, dueAt.timeIntervalSince(now)) + guard let store, + let task = store.startRegisteredStoreWork( + kind: .rateLimitWakeUp, + operation: { @MainActor [weak self] _ in + if delay > 0 { + try? await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000)) + } + guard Task.isCancelled == false else { + return + } + self?.scheduledWakeUp = nil + self?.syncLatestTargets() + } + ) + else { + scheduledWakeUp = nil + return + } scheduledWakeUp = .init( dueAt: dueAt, - task: Task { @MainActor [weak self] in - if delay > 0 { - try? await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000)) - } - guard Task.isCancelled == false else { - return - } - self?.scheduledWakeUp = nil - self?.syncLatestTargets() - } + task: task ) } private func startRefresh(accountKey: String) { - guard accountStates[accountKey]?.isRefreshing != true else { + guard admissionIsOpen, + accountStates[accountKey]?.isRefreshing != true, + let store, + let task = store.startRegisteredStoreWork( + kind: .rateLimitRefresh(accountKey: accountKey), + operation: { @MainActor [weak self] store in + guard let self else { + return + } + let lastFetchAtBeforeRefresh = store.auth.accounts + .first(where: { $0.accountKey == accountKey })? + .lastRateLimitFetchAt + await store.refreshAccountRateLimits(accountKey: accountKey) + guard Task.isCancelled == false else { + return + } + let lastFetchAtAfterRefresh = store.auth.accounts + .first(where: { $0.accountKey == accountKey })? + .lastRateLimitFetchAt + var state = self.accountStates[accountKey] ?? .init() + state.recordRefreshCompletion( + lastFetchAtBeforeRefresh: lastFetchAtBeforeRefresh, + lastFetchAtAfterRefresh: lastFetchAtAfterRefresh, + now: store.clock.now(), + retryDelay: Self.policy.noProgressRefreshRetryDelay + ) + self.accountStates[accountKey] = state + self.syncLatestTargets() + } + ) + else { return } - accountStates[accountKey, default: .init()].refreshTask = Task { @MainActor [weak self, weak store] in - guard let self, let store else { - return - } - let lastFetchAtBeforeRefresh = store.auth.accounts - .first(where: { $0.accountKey == accountKey })? - .lastRateLimitFetchAt - await store.refreshAccountRateLimits(accountKey: accountKey) - guard Task.isCancelled == false else { - return - } - let lastFetchAtAfterRefresh = store.auth.accounts - .first(where: { $0.accountKey == accountKey })? - .lastRateLimitFetchAt - var state = self.accountStates[accountKey] ?? .init() - state.recordRefreshCompletion( - lastFetchAtBeforeRefresh: lastFetchAtBeforeRefresh, - lastFetchAtAfterRefresh: lastFetchAtAfterRefresh, - now: store.clock.now(), - retryDelay: Self.policy.noProgressRefreshRetryDelay - ) - self.accountStates[accountKey] = state - self.syncLatestTargets() - } + accountStates[accountKey, default: .init()].refreshTask = task } } diff --git a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift index f35bd1a..71757a2 100644 --- a/Sources/CodexReview/Store/CodexReviewStoreReviews.swift +++ b/Sources/CodexReview/Store/CodexReviewStoreReviews.swift @@ -15,15 +15,14 @@ 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() - } + try await performThrowingRegisteredStoreWork( + kind: .reviewMutation("start") + ) { @MainActor store in + try await store.performStartReview( + sessionID: sessionID, + request: request, + waitTimeout: nil + ) } } @@ -32,9 +31,45 @@ extension CodexReviewStore { sessionID: String, request: CodexReviewAPI.Start.Request, waitTimeout: Duration + ) async throws -> CodexReviewAPI.Read.Result { + try await performThrowingRegisteredStoreWork( + kind: .reviewMutation("start") + ) { @MainActor store in + try await store.performStartReview( + sessionID: sessionID, + request: request, + waitTimeout: waitTimeout + ) + } + } + + private func performStartReview( + sessionID: String, + request: CodexReviewAPI.Start.Request, + waitTimeout: Duration? ) async throws -> CodexReviewAPI.Read.Result { let jobID = try beginReview(sessionID: sessionID, request: request) - return try await awaitReview(sessionID: sessionID, jobID: jobID, timeout: waitTimeout) + guard let waitTimeout else { + 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 workerTask = reviewWorkerTasks[jobID] + _ = try await awaitReview( + sessionID: sessionID, + jobID: jobID, + timeout: waitTimeout + ) + if storeWorkRegistry.acceptsNewWork == false { + await workerTask?.value + } + return try readReview(sessionID: sessionID, jobID: jobID) } package func awaitReview( @@ -76,24 +111,56 @@ extension CodexReviewStore { ), logEntries: [] ) + guard let workerTask = makeReviewWorker( + jobID: jobID, + sessionID: sessionID, + request: validatedRequest + ) else { + throw CodexReviewAPI.Error.io("Review Store work admission is closed.") + } insertReviewJob(job) markReviewRunning(job, startedAt: createdAt) startingJobIDs.insert(jobID) - launchReviewWorker(jobID: jobID, sessionID: sessionID, request: validatedRequest) + reviewWorkerTasks[jobID]?.cancel() + reviewWorkerTasks[jobID] = workerTask return jobID } - private func launchReviewWorker( + private func makeReviewWorker( jobID: String, sessionID: String, request: CodexReviewAPI.Start.Request - ) { - reviewWorkerTasks[jobID]?.cancel() - reviewWorkerTasks[jobID] = Task { [weak self] in - await self?.runReviewWorker(jobID: jobID, sessionID: sessionID, request: request) + ) -> Task? { + startRegisteredStoreWork( + kind: .reviewWorker(jobID: jobID), + cancelledBeforeEntry: .runFinalizer { store in + store.finishReviewWorkerCancelledBeforeStart(jobID: jobID) + } + ) { @MainActor store in + await store.runReviewWorker( + jobID: jobID, + sessionID: sessionID, + request: request + ) } } + private func finishReviewWorkerCancelledBeforeStart(jobID: String) { + startingJobIDs.remove(jobID) + startupCancellations.removeValue(forKey: jobID) + activeRuns.removeValue(forKey: jobID) + reviewRecoveryWaitingJobIDs.remove(jobID) + if let job = job(id: jobID), job.isTerminal == false { + try? completeCancellationLocally( + jobID: job.id, + sessionID: job.sessionID, + cancellation: job.core.lifecycle.cancellation ?? .system() + ) + } + reviewWorkerTasks.removeValue(forKey: jobID) + resumeReviewWaiters(for: jobID) + } + private func runReviewWorker( jobID: String, sessionID: String, @@ -127,21 +194,20 @@ extension CodexReviewStore { reason: .init(message: startupCancellation.message) ) if job.isTerminal == false { - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: startupCancellation + try completeCancellationAfterRegisteredWorkSuspension( + for: job, + requested: startupCancellation ) } } else if job.cancellationRequested { + let cancellation = job.core.lifecycle.cancellation ?? .system() try await backend.interruptReview( backendRun, - reason: .init(message: job.core.lifecycle.cancellation?.message ?? "Cancellation requested.") + reason: .init(message: cancellation.message) ) - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: job.core.lifecycle.cancellation ?? .system() + try completeCancellationAfterRegisteredWorkSuspension( + for: job, + requested: cancellation ) } @@ -167,10 +233,9 @@ extension CodexReviewStore { await interruptReviewAfterTaskCancellation(cleanupRun, job: job) await cleanupReviewAndRetainFailure(cleanupRun, for: job) } else if job.isTerminal == false || startupCancellation != nil { - try? completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: startupCancellation ?? job.core.lifecycle.cancellation ?? .system() + try? completeCancellationAfterRegisteredWorkSuspension( + for: job, + requested: startupCancellation ?? job.core.lifecycle.cancellation ?? .system() ) } activeRuns.removeValue(forKey: jobID) @@ -188,10 +253,9 @@ extension CodexReviewStore { activeRuns.removeValue(forKey: jobID) reviewRecoveryWaitingJobIDs.remove(jobID) if job.isTerminal == false, let startupCancellation { - try? completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: startupCancellation + try? completeCancellationAfterRegisteredWorkSuspension( + for: job, + requested: startupCancellation ) } else if job.isTerminal == false, let transportFailure = primaryError as? ReviewWorkerInputQueueError { @@ -247,6 +311,45 @@ extension CodexReviewStore { writeDiagnosticsIfNeeded() } + private func authoritativeCancellation( + for job: CodexReviewJob, + requested: ReviewCancellation + ) -> ReviewCancellation { + guard storeWorkRegistry.acceptsNewWork == false else { + return requested + } + return job.core.lifecycle.cancellation ?? requested + } + + private func completeCancellationAfterRegisteredWorkSuspension( + for job: CodexReviewJob, + requested: ReviewCancellation + ) throws { + try completeCancellationLocally( + jobID: job.id, + sessionID: job.sessionID, + cancellation: authoritativeCancellation(for: job, requested: requested) + ) + } + + private func recordCancellationFailureAfterRegisteredWorkSuspension( + for job: CodexReviewJob, + message: String + ) throws { + guard job.isTerminal == false else { + return + } + if storeWorkRegistry.acceptsNewWork { + try recordCancellationFailure( + jobID: job.id, + sessionID: job.sessionID, + message: message + ) + } else { + markReviewFailed(job, message: message) + } + } + private func interruptReviewAfterTaskCancellation(_ run: CodexReviewBackendModel.Review.Run, job: CodexReviewJob) async { guard job.isTerminal == false else { return @@ -261,15 +364,13 @@ extension CodexReviewStore { run, reason: .init(message: cancellation.message) ) - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation + try completeCancellationAfterRegisteredWorkSuspension( + for: job, + requested: cancellation ) } catch { - try? recordCancellationFailure( - jobID: job.id, - sessionID: job.sessionID, + try? recordCancellationFailureAfterRegisteredWorkSuspension( + for: job, message: error.localizedDescription ) } @@ -428,29 +529,54 @@ extension CodexReviewStore { sessionID: String, cancellation: ReviewCancellation = .system() ) async throws -> CodexReviewAPI.Cancel.Outcome { - guard let job = job(id: jobID), job.sessionID == sessionID else { - throw CodexReviewAPI.Error.jobNotFound("Job \(jobID) was not found.") + try await performThrowingRegisteredStoreWork( + kind: .reviewMutation("cancel") + ) { @MainActor store in + guard let job = store.job(id: jobID), job.sessionID == sessionID else { + throw CodexReviewAPI.Error.jobNotFound("Job \(jobID) was not found.") + } + return try await store.performCancelReview( + jobID: jobID, + cancellation: cancellation + ) } - return try await cancelReview(jobID: jobID, cancellation: cancellation) } @discardableResult package func cancelReview( jobID: String, cancellation: ReviewCancellation = .system() + ) async throws -> CodexReviewAPI.Cancel.Outcome { + try await performThrowingRegisteredStoreWork( + kind: .reviewMutation("cancel") + ) { @MainActor store in + try await store.performCancelReview( + jobID: jobID, + cancellation: cancellation + ) + } + } + + func performCancelReview( + jobID: String, + cancellation: ReviewCancellation ) async throws -> CodexReviewAPI.Cancel.Outcome { let job = try job(jobID: jobID) guard job.isTerminal == false else { return .init(jobID: job.id, cancelled: false, core: job.core) } - recordCancellationRequest(cancellation, for: job) + let requestedCancellation = authoritativeCancellation( + for: job, + requested: cancellation + ) + recordCancellationRequest(requestedCancellation, for: job) if job.core.lifecycle.status == .queued { try completeCancellationLocally( jobID: job.id, sessionID: job.sessionID, - cancellation: cancellation + cancellation: requestedCancellation ) return .init(jobID: job.id, cancelled: true, core: job.core) } @@ -459,7 +585,7 @@ extension CodexReviewStore { try completeCancellationLocally( jobID: job.id, sessionID: job.sessionID, - cancellation: cancellation + cancellation: requestedCancellation ) reviewWorkerTasks[jobID]?.cancel() return .init(jobID: job.id, cancelled: true, core: job.core) @@ -469,13 +595,12 @@ extension CodexReviewStore { do { try await backend.interruptReview( run, - reason: .init(message: cancellation.message) + reason: .init(message: requestedCancellation.message) ) if job.isTerminal == false { - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation + try completeCancellationAfterRegisteredWorkSuspension( + for: job, + requested: requestedCancellation ) reviewWorkerTasks[jobID]?.cancel() } @@ -487,9 +612,8 @@ extension CodexReviewStore { core: job.core ) } - try recordCancellationFailure( - jobID: job.id, - sessionID: job.sessionID, + try recordCancellationFailureAfterRegisteredWorkSuspension( + for: job, message: error.localizedDescription ) throw error @@ -498,13 +622,12 @@ extension CodexReviewStore { do { try await backend.interruptReview( run, - reason: .init(message: cancellation.message) + reason: .init(message: requestedCancellation.message) ) if job.isTerminal == false { - try completeCancellationLocally( - jobID: job.id, - sessionID: job.sessionID, - cancellation: cancellation + try completeCancellationAfterRegisteredWorkSuspension( + for: job, + requested: requestedCancellation ) reviewWorkerTasks[jobID]?.cancel() } @@ -516,26 +639,25 @@ extension CodexReviewStore { core: job.core ) } - try recordCancellationFailure( - jobID: job.id, - sessionID: job.sessionID, + try recordCancellationFailureAfterRegisteredWorkSuspension( + for: job, message: error.localizedDescription ) throw error } } else if startingJobIDs.contains(jobID) { - startupCancellations[jobID] = cancellation + startupCancellations[jobID] = requestedCancellation try completeCancellationLocally( jobID: job.id, sessionID: job.sessionID, - cancellation: cancellation + cancellation: requestedCancellation ) return .init(jobID: job.id, cancelled: true, core: job.core) } else { try completeCancellationLocally( jobID: job.id, sessionID: job.sessionID, - cancellation: cancellation + cancellation: requestedCancellation ) } return .init( @@ -548,19 +670,37 @@ extension CodexReviewStore { package func closeSession( _ sessionID: String, reason: ReviewCancellation = .sessionClosed() + ) async { + await performRegisteredStoreWork( + kind: .reviewMutation("close-session") + ) { @MainActor store in + await store.performCloseSession(sessionID, reason: reason) + } + } + + private func performCloseSession( + _ sessionID: String, + reason: ReviewCancellation ) async { closedSessions.insert(sessionID) for jobID in activeJobIDs(for: sessionID) { - _ = try? await cancelReview(jobID: jobID, cancellation: reason) + _ = try? await performCancelReview(jobID: jobID, cancellation: reason) } } package func closeActiveReviewSessions(reason: ReviewCancellation) async { - let jobIDs = orderedJobs - .filter { $0.isTerminal == false } - .map(\.id) - for jobID in jobIDs { - _ = try? await cancelReview(jobID: jobID, cancellation: reason) + await performRegisteredStoreWork( + kind: .reviewMutation("close-active-sessions") + ) { @MainActor store in + let jobIDs = store.orderedJobs + .filter { $0.isTerminal == false } + .map(\.id) + for jobID in jobIDs { + _ = try? await store.performCancelReview( + jobID: jobID, + cancellation: reason + ) + } } } @@ -671,9 +811,27 @@ extension CodexReviewStore { startRequest: CodexReviewBackendModel.Review.Start ) async throws -> CodexReviewBackendModel.Review.Run { let inputs = await reviewWorkerInputs(for: initialAttempt) - defer { - inputs.cancel() + do { + let result = try await consumeReviewEvents( + inputs: inputs, + initialAttempt: initialAttempt, + job: job, + startRequest: startRequest + ) + await inputs.cancelAndWait() + return result + } catch { + await inputs.cancelAndWait() + throw error } + } + + private func consumeReviewEvents( + inputs: ReviewWorkerInputs, + initialAttempt: BackendReviewAttempt, + job: CodexReviewJob, + startRequest: CodexReviewBackendModel.Review.Start + ) async throws -> CodexReviewBackendModel.Review.Run { var recoveryState = ReviewNetworkRecoveryLoopState(currentRun: initialAttempt.run) var activeEventSubscriptionID: Int? = inputs.initialEventSubscriptionID while let input = await inputs.next() { @@ -897,17 +1055,15 @@ extension CodexReviewStore { 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 completeCancellationAfterRegisteredWorkSuspension( + for: job, + requested: cancellation ) } catch { let primaryError = error let cleanupFailure = await cleanupReviewFailure(recoveredRun) - try? recordCancellationFailure( - jobID: job.id, - sessionID: job.sessionID, + try? recordCancellationFailureAfterRegisteredWorkSuspension( + for: job, message: primaryError.localizedDescription ) if let cleanupFailure { @@ -1079,8 +1235,10 @@ extension CodexReviewStore { return } let waiterID = UUID() - let timeoutTask = timeout.map { duration in - Task { @MainActor [weak self] in + let timeoutTask = timeout.flatMap { duration in + startRegisteredStoreWork( + kind: .reviewWaiter(jobID: jobID) + ) { @MainActor [weak self] _ in do { try await Task.sleep(for: duration) } catch { @@ -1439,13 +1597,12 @@ private struct ReviewWorkerInputs { await eventSource.cancelActiveSubscription() } - func cancel() { + func cancelAndWait() async { networkTask.cancel() - Task { - await eventSource.cancel() - await signalCoordinator.cancel() - await queue.finish() - } + await eventSource.cancelAndWait() + await signalCoordinator.cancelAndWait() + await queue.finish() + await networkTask.value } } @@ -1551,11 +1708,11 @@ private actor ReviewWorkerEventSource { self.queue = queue } - func subscribe(to attempt: BackendReviewAttempt) -> Int { + func subscribe(to attempt: BackendReviewAttempt) async -> Int { subscriptionID += 1 let subscriptionID = subscriptionID activeSubscriptionID = subscriptionID - cancelEventTasks() + await cancelEventTasksAndWait() let run = attempt.run let events = attempt.events eventTasks[subscriptionID] = Task { @@ -1574,23 +1731,27 @@ private actor ReviewWorkerEventSource { return subscriptionID } - func cancelActiveSubscription() { + func cancelActiveSubscription() async { subscriptionID += 1 activeSubscriptionID = nil - cancelEventTasks() + await cancelEventTasksAndWait() } - func cancel() { + func cancelAndWait() async { subscriptionID += 1 activeSubscriptionID = nil - cancelEventTasks() + await cancelEventTasksAndWait() } - private func cancelEventTasks() { - for task in eventTasks.values { + private func cancelEventTasksAndWait() async { + let tasks = eventTasks.sorted { $0.key < $1.key }.map(\.value) + eventTasks.removeAll(keepingCapacity: true) + for task in tasks { task.cancel() } - eventTasks.removeAll(keepingCapacity: true) + for task in tasks { + await task.value + } } private func yieldReviewEvent( @@ -1661,9 +1822,12 @@ private actor ReviewNetworkSignalCoordinator { private let policy: CodexReviewNetworkRecoveryPolicy private let tracker: ReviewNetworkStatusTracker private let queue: ReviewWorkerInputQueue - private var outageTask: Task? + private var admissionIsOpen = true + private var nextTaskID = 0 + private var ownedTasks: [Int: Task] = [:] + private var outageTaskID: Int? private var outageGeneration = 0 - private var recoveryTask: Task? + private var recoveryTaskID: Int? private var recoveryGeneration = 0 init( @@ -1677,50 +1841,71 @@ private actor ReviewNetworkSignalCoordinator { } func observe(_ snapshot: CodexReviewNetworkSnapshot) async { + guard admissionIsOpen else { + return + } await tracker.update(snapshot) + guard admissionIsOpen else { + return + } switch snapshot.status { case .satisfied: outageGeneration += 1 - outageTask?.cancel() - outageTask = nil + cancelOwnedTask(outageTaskID) + outageTaskID = nil recoveryGeneration += 1 let recoveryGeneration = recoveryGeneration - recoveryTask?.cancel() - recoveryTask = nil + cancelOwnedTask(recoveryTaskID) + recoveryTaskID = nil await queue.send(.networkSnapshot(snapshot, recoveryGeneration: recoveryGeneration)) scheduleRecoveryConfirmationIfNeeded(generation: recoveryGeneration) case .unsatisfied, .requiresConnection: recoveryGeneration += 1 let recoveryGeneration = recoveryGeneration - recoveryTask?.cancel() - recoveryTask = nil + cancelOwnedTask(recoveryTaskID) + recoveryTaskID = nil await queue.send(.networkSnapshot(snapshot, recoveryGeneration: recoveryGeneration)) scheduleOutageConfirmationIfNeeded() } } - func cancel() { - outageTask?.cancel() - outageTask = nil - recoveryTask?.cancel() - recoveryTask = nil + func cancelAndWait() async { + admissionIsOpen = false + outageGeneration += 1 + recoveryGeneration += 1 + outageTaskID = nil + recoveryTaskID = nil + let tasks = ownedTasks.sorted { $0.key < $1.key }.map(\.value) + for task in tasks { + task.cancel() + } + for task in tasks { + await task.value + } + ownedTasks.removeAll(keepingCapacity: false) } private func scheduleOutageConfirmationIfNeeded() { - guard outageTask == nil else { + guard admissionIsOpen, outageTaskID == nil else { return } let policy = policy outageGeneration += 1 let generation = outageGeneration - outageTask = Task { + nextTaskID += 1 + let taskID = nextTaskID + let task = Task { do { try await policy.sleep(policy.outageDebounce) } catch { + self.finishOwnedTask(taskID) return } await self.confirmOutageIfCurrent(generation: generation) + self.finishOwnedTask(taskID) } + ownedTasks[taskID] = task + outageTaskID = taskID } private func confirmOutageIfCurrent(generation: Int) async { @@ -1735,29 +1920,51 @@ private actor ReviewNetworkSignalCoordinator { } private func scheduleRecoveryConfirmationIfNeeded(generation: Int) { - guard recoveryTask == nil else { + guard admissionIsOpen, recoveryTaskID == nil else { return } let policy = policy - recoveryTask = Task { + nextTaskID += 1 + let taskID = nextTaskID + let task = Task { do { try await policy.sleep(policy.recoverySettle) } catch { + self.finishOwnedTask(taskID) return } await self.confirmRecoveryIfCurrent(generation: generation) + self.finishOwnedTask(taskID) } + ownedTasks[taskID] = task + recoveryTaskID = taskID } private func confirmRecoveryIfCurrent(generation: Int) async { guard generation == recoveryGeneration else { return } - recoveryTask = nil let latest = await tracker.latestSnapshot() guard latest.status == .satisfied else { return } await queue.send(.networkRecoverySettled(recoveryGeneration: generation)) } + + private func cancelOwnedTask(_ taskID: Int?) { + guard let taskID else { + return + } + ownedTasks[taskID]?.cancel() + } + + private func finishOwnedTask(_ taskID: Int) { + ownedTasks.removeValue(forKey: taskID) + if outageTaskID == taskID { + outageTaskID = nil + } + if recoveryTaskID == taskID { + recoveryTaskID = nil + } + } } diff --git a/Sources/CodexReview/Store/ReviewStoreLifecycle.swift b/Sources/CodexReview/Store/ReviewStoreLifecycle.swift new file mode 100644 index 0000000..21d28cb --- /dev/null +++ b/Sources/CodexReview/Store/ReviewStoreLifecycle.swift @@ -0,0 +1,282 @@ +import Foundation + +package enum ReviewStoreWorkKind: Hashable, Sendable { + case reviewMutation(String) + case reviewWorker(jobID: String) + case reviewWaiter(jobID: String) + case rateLimitWakeUp + case rateLimitRefresh(accountKey: String) + case accountAction + case testing(String) + + package var description: String { + switch self { + case .reviewMutation(let operation): + "Review mutation \(operation)" + case .reviewWorker(let jobID): + "Review worker \(jobID)" + case .reviewWaiter(let jobID): + "Review waiter \(jobID)" + case .rateLimitWakeUp: + "Rate-limit wake-up" + case .rateLimitRefresh(let accountKey): + "Rate-limit refresh \(accountKey)" + case .accountAction: + "Account action" + case .testing(let label): + label + } + } +} + +package enum ReviewStoreWorkCancelledBeforeEntryPolicy: Sendable { + case skip + case runFinalizer(@MainActor @Sendable (CodexReviewStore) -> Void) +} + +package enum ReviewStoreWorkFailureCause: LocalizedError, Equatable, Sendable { + case interruptRequest(ReviewInterruptRequestFailure) + case runtime(ReviewRuntimeCloseFailure) + case operation(String) + + package var errorDescription: String? { + switch self { + case .interruptRequest(let failure): + failure.localizedDescription + case .runtime(let failure): + failure.localizedDescription + case .operation(let message): + message + } + } +} + +package struct ReviewStoreWorkFailure: LocalizedError, Equatable, Sendable { + package let ordinal: UInt64 + package let kind: ReviewStoreWorkKind + package let cause: ReviewStoreWorkFailureCause + + package var errorDescription: String? { + "\(kind.description) failed: \(cause.localizedDescription)" + } +} + +package struct ReviewStoreWorkFailureAggregate: LocalizedError, Equatable, Sendable { + package let first: ReviewStoreWorkFailure + package let additionalInOrdinalOrder: [ReviewStoreWorkFailure] + + package var errorDescription: String? { + ([first] + additionalInOrdinalOrder) + .map(\.localizedDescription) + .joined(separator: "; ") + } +} + +package struct ReviewStoreWorkDrainResult: Equatable, Sendable { + package let failures: ReviewStoreWorkFailureAggregate? + + package static let success = ReviewStoreWorkDrainResult(failures: nil) +} + +package enum ReviewStoreWorkRegistryStatus: Equatable, Sendable { + case open + case closing + case closed +} + +@MainActor +package final class ReviewStoreWorkRegistry { + package struct Admission: Hashable, Sendable { + package let ordinal: UInt64 + package let kind: ReviewStoreWorkKind + } + + package struct CloseOperation { + package let id: UInt64 + package let task: Task + } + + private enum State { + case open + case closing(CloseOperation) + case closed(CloseOperation, ReviewStoreWorkDrainResult) + } + + @MainActor + private final class RegisteredTask { + let admission: Admission + private let cancelOperation: () -> Void + private let waitOperation: () async -> (any Error)? + + init( + admission: Admission, + task: Task + ) { + self.admission = admission + self.cancelOperation = { + task.cancel() + } + self.waitOperation = { + switch await task.result { + case .success: + return nil + case .failure(let error): + return error + } + } + } + + func cancel() { + cancelOperation() + } + + func waitForFailure() async -> ReviewStoreWorkFailure? { + guard let error = await waitOperation() else { + return nil + } + guard error is CancellationError == false else { + return nil + } + let cause: ReviewStoreWorkFailureCause + if let failure = error as? ReviewInterruptRequestFailure { + cause = .interruptRequest(failure) + } else if let failure = error as? ReviewRuntimeCloseFailure { + cause = .runtime(failure) + } else { + cause = .operation(error.localizedDescription) + } + return .init( + ordinal: admission.ordinal, + kind: admission.kind, + cause: cause + ) + } + } + + private var state: State = .open + private var admissionIsOpen = true + private var nextWorkOrdinal: UInt64 = 0 + private var nextCloseID: UInt64 = 0 + private var registeredTasks: [UInt64: RegisteredTask] = [:] + + package private(set) var closeTaskCreationCount = 0 + + package var status: ReviewStoreWorkRegistryStatus { + switch state { + case .open: + .open + case .closing: + .closing + case .closed: + .closed + } + } + + package var acceptsNewWork: Bool { + admissionIsOpen + } + + package var activeOrdinals: [UInt64] { + registeredTasks.keys.sorted() + } + + package func register(_ kind: ReviewStoreWorkKind) -> Admission? { + guard admissionIsOpen else { + return nil + } + guard nextWorkOrdinal < UInt64.max else { + preconditionFailure("ReviewStoreWorkRegistry work ordinal exhausted.") + } + nextWorkOrdinal += 1 + return .init(ordinal: nextWorkOrdinal, kind: kind) + } + + package func install( + _ task: Task, + for admission: Admission + ) { + precondition( + registeredTasks[admission.ordinal] == nil, + "ReviewStoreWorkRegistry owns exactly one Task for each admitted ordinal." + ) + registeredTasks[admission.ordinal] = .init( + admission: admission, + task: task + ) + } + + package func finish(_ admission: Admission) { + registeredTasks.removeValue(forKey: admission.ordinal) + } + + package func beginClosing( + onAdmissionClosed: () -> Void + ) -> CloseOperation { + switch state { + case .open: + admissionIsOpen = false + onAdmissionClosed() + let tasks = registeredTasks.values.sorted { + $0.admission.ordinal < $1.admission.ordinal + } + guard nextCloseID < UInt64.max else { + preconditionFailure("ReviewStoreWorkRegistry close ordinal exhausted.") + } + nextCloseID += 1 + let operation = CloseOperation( + id: nextCloseID, + task: Task { @MainActor in + for task in tasks { + task.cancel() + } + var failures: [ReviewStoreWorkFailure] = [] + for task in tasks { + if let failure = await task.waitForFailure() { + failures.append(failure) + } + } + guard let first = failures.first else { + return .success + } + return .init(failures: .init( + first: first, + additionalInOrdinalOrder: Array(failures.dropFirst()) + )) + } + ) + closeTaskCreationCount += 1 + state = .closing(operation) + return operation + + case .closing(let operation): + return operation + + case .closed(let operation, _): + return operation + } + } + + package func completeClosing( + _ operation: CloseOperation, + result: ReviewStoreWorkDrainResult + ) { + switch state { + case .closing(let current) where current.id == operation.id: + registeredTasks.removeAll(keepingCapacity: false) + state = .closed(current, result) + case .closed(let current, _) where current.id == operation.id: + break + case .open, .closing, .closed: + preconditionFailure( + "ReviewStoreWorkRegistry must complete its one recorded close operation." + ) + } + } + + package func cancelWithoutWaiting() { + admissionIsOpen = false + for task in registeredTasks.values { + task.cancel() + } + } +} diff --git a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift index 77484ab..10487a3 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreCommandTests.swift @@ -1633,6 +1633,137 @@ struct CodexReviewStoreCommandTests { } } + @Test func registeredWorkCloseFinalizesHeldInterruptFailureBeforeWorkerExit() async throws { + let interruptGate = AsyncGate() + let backend = HeldInterruptStoreBackend( + reviewBackend: FakeCodexReviewBackend(), + interruptGate: interruptGate, + failureMessage: "Interrupt failed" + ) + let store = CodexReviewStore.makeTestingStore( + backend: backend, + idGenerator: .init(next: { "job-1" }) + ) + let review = Task { @MainActor in + try await 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) + let reason = ReviewCancellation.system(message: "Store work owner closed.") + + let closeCompletion = StoreCommandTaskCompletion() + let close = Task { @MainActor in + let result = await store.closeRegisteredStoreWork(reason: reason) + await closeCompletion.complete() + return result + } + await backend.waitUntilInterruptStarts() + let held = try store.readReview(jobID: "job-1") + #expect(held.core.lifecycle.cancellation == reason) + #expect(await closeCompletion.isComplete() == false) + + await interruptGate.open() + #expect(await close.value == .success) + let failed = try await review.value + + #expect(failed.core.lifecycle.status == .failed) + #expect(failed.core.lifecycle.terminal == .failed(message: "Interrupt failed")) + #expect(failed.core.lifecycle.cancellation == reason) + #expect(failed.core.lifecycle.errorMessage == "Interrupt failed") + #expect(failed.core.output.summary == "Interrupt failed") + #expect(store.reviewWorkerTasks["job-1"] == nil) + #expect(store.storeWorkRegistry.activeOrdinals.isEmpty) + #expect(await closeCompletion.isComplete()) + } + + @Test func registeredWorkCloseReasonWinsInFlightCancellation() async throws { + let interruptGate = AsyncGate() + let backend = HeldInterruptStoreBackend( + reviewBackend: FakeCodexReviewBackend(), + interruptGate: interruptGate + ) + let store = CodexReviewStore.makeTestingStore(backend: backend) + let job = makeRunningReviewJob(id: "job-1") + store.loadForTesting( + serverState: .running, + workspaces: [.init(cwd: "/tmp/project")], + jobs: [job] + ) + let userReason = ReviewCancellation.mcpClient(message: "User cancellation.") + let cancellation = Task { @MainActor in + try await store.cancelReview(jobID: job.id, cancellation: userReason) + } + await backend.waitUntilInterruptStarts() + let closeReason = ReviewCancellation.system(message: "Store work owner closed.") + + let close = Task { @MainActor in + await store.closeRegisteredStoreWork(reason: closeReason) + } + let closing = await waitUntil { + store.storeWorkRegistryStatus == .closing + } + + #expect(closing) + #expect(job.core.lifecycle.cancellation == closeReason) + await interruptGate.open() + let outcome = try await cancellation.value + #expect(await close.value == .success) + + #expect(outcome.core.lifecycle.status == .cancelled) + #expect(outcome.core.lifecycle.cancellation == closeReason) + #expect(job.core.lifecycle.cancellation == closeReason) + #expect(store.storeWorkRegistry.activeOrdinals.isEmpty) + } + + @Test func registeredWorkCloseDrainsWholeBulkCancellation() async throws { + let interruptGate = AsyncGate() + let backend = HeldInterruptStoreBackend( + reviewBackend: FakeCodexReviewBackend(), + interruptGate: interruptGate + ) + let store = CodexReviewStore.makeTestingStore(backend: backend) + let firstJob = makeRunningReviewJob(id: "job-1") + let secondJob = makeRunningReviewJob(id: "job-2") + store.loadForTesting( + serverState: .running, + workspaces: [.init(cwd: "/tmp/project")], + jobs: [firstJob, secondJob] + ) + let cancellation = Task { @MainActor in + try await store.cancelAllRunningJobs(reason: "User cancelled all reviews.") + } + await backend.waitUntilInterruptStarts() + let closeReason = ReviewCancellation.system(message: "Store work owner closed.") + + let close = Task { @MainActor in + await store.closeRegisteredStoreWork(reason: closeReason) + } + let closing = await waitUntil { + store.storeWorkRegistryStatus == .closing + } + + #expect(closing) + #expect(firstJob.core.lifecycle.cancellation == closeReason) + #expect(secondJob.core.lifecycle.cancellation == closeReason) + await interruptGate.open() + let cancellationResult = await cancellation.result + #expect(await close.value == .success) + + if case .failure(let error) = cancellationResult { + Issue.record("Bulk cancellation failed after owner close: \(error)") + } + #expect(firstJob.core.lifecycle.status == .cancelled) + #expect(secondJob.core.lifecycle.status == .cancelled) + #expect(firstJob.core.lifecycle.cancellation == closeReason) + #expect(secondJob.core.lifecycle.cancellation == closeReason) + #expect(store.storeWorkRegistry.activeOrdinals.isEmpty) + } + @Test func runtimeStopDetachesNetworkRecoveryWaitingWorker() async throws { let run = CodexReviewBackendModel.Review.Run( threadID: "thread-1", @@ -1703,10 +1834,10 @@ struct CodexReviewStoreCommandTests { } } - @Test func runtimeStopDetachLetsStartReviewReturnWhenBackendStartIsStuck() async throws { + @Test func runtimeStopBoundedDrainReturnsWhileBackendStartIsStuck() async throws { let backend = FakeCodexReviewBackend() let startReviewGate = AsyncGate() - await backend.holdStartReview(with: startReviewGate) + await backend.holdStartReviewIgnoringCancellation(with: startReviewGate) let store = CodexReviewStore.makeTestingStore( backend: TestingCodexReviewStoreBackend(reviewBackend: backend), idGenerator: .init(next: { "job-1" }) @@ -1725,17 +1856,80 @@ struct CodexReviewStoreCommandTests { cancelWorkers: false ) store.cancelAndDetachReviewWorkersForRuntimeStop(jobIDs: locallyCancelledJobIDs) - let resultBeforeStartReviewUnblocked = try await waitForTaskValue(running, timeout: .seconds(1)) + let didDrain = await store.drainReviewWorkersForRuntimeStop( + timeout: .milliseconds(20) + ) + let resultBeforeStartReviewUnblocked = try await waitForTaskValue( + running, + timeout: .seconds(1) + ) await startReviewGate.open() let result = try #require(resultBeforeStartReviewUnblocked) #expect(locallyCancelledJobIDs == ["job-1"]) + #expect(didDrain == false) #expect(result.core.lifecycle.status == .cancelled) #expect(store.reviewWorkerTasks["job-1"] == nil) #expect(store.activeRuns["job-1"] == nil) } } + @Test func registeredWorkCloseAwaitsCurrentAndRetiringNetworkDebounceTasks() async throws { + let backend = FakeCodexReviewBackend() + let networkMonitor = ManualCodexReviewNetworkMonitor() + let sleepGate = AsyncGate() + let sleepProbe = StoreCommandSleepProbe() + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-1" }), + networkMonitor: networkMonitor, + networkRecoveryPolicy: .init( + outageDebounce: .seconds(1), + recoverySettle: .seconds(1), + sleep: { _ in + await sleepProbe.sleepIgnoringCancellation(on: sleepGate) + } + ) + ) + try await withStoreCommandTestCleanup(backend: backend, store: store) { + let review = Task { @MainActor in + try await store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .baseBranch("main")) + ) + } + try #require(await StoreSnapshotProbe(store: store) + .waitUntilJobStatus(.running, jobID: "job-1") != nil) + + await sleepProbe.waitForCount(1) + networkMonitor.yield(.init(status: .unsatisfied)) + await sleepProbe.waitForCount(2) + networkMonitor.yield(.satisfied()) + await sleepProbe.waitForCount(3) + + let closeCompletion = StoreCommandTaskCompletion() + let close = Task { @MainActor in + let result = await store.closeRegisteredStoreWork( + reason: .system(message: "Store work closed.") + ) + await closeCompletion.complete() + return result + } + let workerCancellationStarted = await waitUntil { + store.reviewWorkerTasks["job-1"]?.isCancelled == true + } + + #expect(workerCancellationStarted) + #expect(await closeCompletion.isComplete() == false) + await sleepGate.open() + #expect(await close.value == .success) + #expect(try await review.value.core.lifecycle.status == .cancelled) + #expect(await closeCompletion.isComplete()) + let finalSleepCount = await sleepProbe.count + #expect(finalSleepCount == 3) + } + } + @Test func cancellationDuringNetworkRecoveryStopsWhenEventStreamFinishes() async throws { let initialRun = CodexReviewBackendModel.Review.Run( threadID: "thread-1", @@ -2128,6 +2322,7 @@ struct CodexReviewStoreCommandTests { request: .init(cwd: "/tmp/project", target: .baseBranch("main")) ) } + try await backend.waitForStartReview(timeout: .seconds(2)) task.cancel() let read = try await task.value @@ -2492,6 +2687,31 @@ struct CodexReviewStoreCommandTests { #expect(backend.switchRequests == [displayedOtherAccount.accountKey]) } + @Test func registeredWorkCloseSkipsAccountActionCancelledBeforeEntry() async throws { + let selectedAccount = CodexAccount(email: "selected@example.com", planType: "pro") + let otherAccount = CodexAccount(email: "other@example.com", planType: "plus") + let backend = SwitchRecordingBackend() + let store = CodexReviewStore.makeTestingStore(backend: backend) + store.loadForTesting( + serverState: .running, + account: selectedAccount, + persistedAccounts: [selectedAccount, otherAccount], + workspaces: [] + ) + let displayedOtherAccount = try #require( + store.auth.persistedAccounts.first { $0.accountKey == otherAccount.accountKey } + ) + + store.requestSwitchAccountFromUserAction(displayedOtherAccount) + let closeOperation = store.storeWorkRegistry.beginClosing(onAdmissionClosed: {}) + let result = await closeOperation.task.value + store.storeWorkRegistry.completeClosing(closeOperation, result: result) + + #expect(result == .success) + #expect(backend.switchRequests.isEmpty) + #expect(store.auth.selectedAccount?.accountKey == selectedAccount.accountKey) + } + @Test func fakeBackendPreservesSettingsCatalogWhenApplyingOverrides() async throws { let model = CodexReviewSettings.ModelCatalogItem( id: "gpt-5.5", @@ -2586,6 +2806,62 @@ private final class SwitchRecordingBackend: PreviewCodexReviewStoreBackend { } } +@MainActor +private final class HeldInterruptStoreBackend: PreviewCodexReviewStoreBackend { + private let reviewBackend: FakeCodexReviewBackend + private let interruptGate: AsyncGate + private let failureMessage: String? + private let interruptStartedGate = AsyncGate() + + init( + reviewBackend: FakeCodexReviewBackend, + interruptGate: AsyncGate, + failureMessage: String? = nil + ) { + self.reviewBackend = reviewBackend + self.interruptGate = interruptGate + self.failureMessage = failureMessage + super.init() + } + + override func startReview( + _ request: CodexReviewBackendModel.Review.Start, + admission: ReviewStartAdmission + ) async throws -> BackendReviewAttempt { + try await reviewBackend.startReview(request, admission: admission) + } + + override func interruptReview( + _: CodexReviewBackendModel.Review.Run, + reason _: CodexReviewBackendModel.CancellationReason + ) async throws { + await interruptStartedGate.open() + await interruptGate.waitIgnoringCancellation() + if let failureMessage { + throw FakeCodexReviewBackendError(message: failureMessage) + } + } + + override func cleanupReview(_: CodexReviewBackendModel.Review.Run) async {} + + func waitUntilInterruptStarts() async { + await interruptStartedGate.wait() + } +} + +@MainActor +private func makeRunningReviewJob(id: String) -> CodexReviewJob { + CodexReviewJob.makeForTesting( + id: id, + cwd: "/tmp/project", + targetSummary: "Uncommitted changes", + threadID: "thread-\(id)", + turnID: "turn-\(id)", + status: .running, + summary: "Running" + ) +} + @MainActor private func waitUntil( timeout: Duration = .seconds(2), @@ -2630,6 +2906,50 @@ private func waitForTaskValue( } } +private actor StoreCommandTaskCompletion { + private var completed = false + + func complete() { + completed = true + } + + func isComplete() -> Bool { + completed + } +} + +private actor StoreCommandSleepProbe { + private var sleepCount = 0 + private var waiters: [(Int, CheckedContinuation)] = [] + + var count: Int { + sleepCount + } + + func sleepIgnoringCancellation(on gate: AsyncGate) async { + sleepCount += 1 + let ready = waiters.filter { sleepCount >= $0.0 } + waiters.removeAll { sleepCount >= $0.0 } + for (_, waiter) in ready { + waiter.resume() + } + await gate.waitIgnoringCancellation() + } + + func waitForCount(_ count: Int) async { + if sleepCount >= count { + return + } + await withCheckedContinuation { continuation in + if sleepCount >= count { + continuation.resume() + } else { + waiters.append((count, continuation)) + } + } + } +} + @MainActor private func withStoreCommandTestCleanup( backend: FakeCodexReviewBackend, diff --git a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift index b84a2b9..d6668d7 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreLifecycleTests.swift @@ -413,6 +413,298 @@ struct CodexReviewStoreLifecycleTests { #expect(store.settings.lastErrorMessage == nil) } + @Test func registeredWorkCloseJoinsOneTaskAndReplaysOneResult() async throws { + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend() + ) + ) + let entered = AsyncGate() + let release = AsyncGate() + let work = Task { @MainActor in + await store.performRegisteredStoreWork( + kind: .testing("held work") + ) { _ in + await entered.open() + await release.waitIgnoringCancellation() + } + } + await entered.wait() + + let reason = ReviewCancellation.system(message: "Store work closed.") + let firstClose = Task { @MainActor in + await store.closeRegisteredStoreWork(reason: reason) + } + try await waitForStoreWorkStatus(.closing, store: store) + let secondClose = Task { @MainActor in + await store.closeRegisteredStoreWork(reason: reason) + } + + #expect(store.storeWorkRegistry.closeTaskCreationCount == 1) + #expect(store.storeWorkRegistry.activeOrdinals == [1]) + #expect(store.startRegisteredStoreWork( + kind: .testing("rejected"), + operation: { _ in } + ) == nil) + + await release.open() + await work.value + let firstResult = await firstClose.value + let secondResult = await secondClose.value + let replayedResult = await store.closeRegisteredStoreWork(reason: reason) + + #expect(firstResult == .success) + #expect(secondResult == firstResult) + #expect(replayedResult == firstResult) + #expect(store.storeWorkRegistryStatus == .closed) + #expect(store.storeWorkRegistry.closeTaskCreationCount == 1) + #expect(store.storeWorkRegistry.activeOrdinals.isEmpty) + await #expect(throws: CodexReviewAPI.Error.self) { + _ = try await store.startReview( + sessionID: "closed-session", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + } + + @Test func registeredWorkFailuresStayInAdmissionOrder() async throws { + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend() + ) + ) + let firstEntered = AsyncGate() + let firstRelease = AsyncGate() + let secondEntered = AsyncGate() + let secondRelease = AsyncGate() + let first = Task { @MainActor in + try? await store.performThrowingRegisteredStoreWork( + kind: .testing("first work") + ) { _ in + await firstEntered.open() + await firstRelease.waitIgnoringCancellation() + throw StoreWorkTestFailure.first + } + } + await firstEntered.wait() + let second = Task { @MainActor in + try? await store.performThrowingRegisteredStoreWork( + kind: .testing("second work") + ) { _ in + await secondEntered.open() + await secondRelease.waitIgnoringCancellation() + throw StoreWorkTestFailure.second + } + } + await secondEntered.wait() + + let close = Task { @MainActor in + await store.closeRegisteredStoreWork( + reason: .system(message: "Store work closed.") + ) + } + try await waitForStoreWorkStatus(.closing, store: store) + await secondRelease.open() + await firstRelease.open() + await first.value + await second.value + let result = await close.value + let failures = try #require(result.failures) + + #expect(failures.first.ordinal == 1) + #expect(failures.first.kind == .testing("first work")) + #expect(failures.first.cause == .operation("first work failed")) + #expect(failures.additionalInOrdinalOrder.count == 1) + #expect(failures.additionalInOrdinalOrder[0].ordinal == 2) + #expect(failures.additionalInOrdinalOrder[0].kind == .testing("second work")) + #expect(failures.additionalInOrdinalOrder[0].cause == .operation("second work failed")) + #expect(await store.closeRegisteredStoreWork( + reason: .system(message: "ignored replay reason") + ) == result) + } + + @Test func registeredWorkCloseAwaitsReviewWorkerWithExactReason() async throws { + let backend = FakeCodexReviewBackend() + let startGate = AsyncGate() + await backend.holdStartReviewIgnoringCancellation(with: startGate) + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-1" }) + ) + await store.start() + let review = Task { @MainActor in + try await store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges) + ) + } + try await backend.waitForStartReview(timeout: .seconds(2)) + let awaitCompletion = StoreWorkCompletion() + let awaiter = Task { @MainActor in + let result = try await store.awaitReview( + sessionID: "session-1", + jobID: "job-1" + ) + await awaitCompletion.complete() + return result + } + try await waitForReviewWaiterCount(2, jobID: "job-1", store: store) + let reason = ReviewCancellation.system(message: "Store work owner closed.") + + let close = Task { @MainActor in + await store.closeRegisteredStoreWork(reason: reason) + } + try await waitForStoreWorkStatus(.closing, store: store) + try await waitForReviewWorkerCancellation(jobID: "job-1", store: store) + #expect(store.reviewWorkerTasks["job-1"]?.isCancelled == true) + #expect(store.storeWorkRegistry.activeOrdinals.isEmpty == false) + #expect(await awaitCompletion.isComplete() == false) + + await startGate.open() + #expect(await close.value == .success) + let result = try await review.value + let awaitedResult = try await awaiter.value + + #expect(result.core.lifecycle.status == .cancelled) + #expect(awaitedResult.core.lifecycle.status == .cancelled) + #expect(result.core.lifecycle.cancellation == reason) + #expect(store.reviewWorkerTasks["job-1"] == nil) + #expect(store.storeWorkRegistry.activeOrdinals.isEmpty) + await store.stop() + } + + @Test func registeredWorkCloseDoesNotReturnTimedStartBeforeWorkerFinalization() async throws { + let backend = FakeCodexReviewBackend() + let startGate = AsyncGate() + await backend.holdStartReviewIgnoringCancellation(with: startGate) + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend), + idGenerator: .init(next: { "job-1" }) + ) + await store.start() + let reviewCompletion = StoreWorkCompletion() + let review = Task { @MainActor in + let result = try await store.startReview( + sessionID: "session-1", + request: .init(cwd: "/tmp/project", target: .uncommittedChanges), + waitTimeout: .seconds(30) + ) + await reviewCompletion.complete() + return result + } + try await backend.waitForStartReview(timeout: .seconds(2)) + try await waitForReviewWaiterCount(1, jobID: "job-1", store: store) + let reason = ReviewCancellation.system(message: "Store work owner closed.") + + let close = Task { @MainActor in + await store.closeRegisteredStoreWork(reason: reason) + } + try await waitForStoreWorkStatus(.closing, store: store) + try await waitForReviewWorkerCancellation(jobID: "job-1", store: store) + let returnedBeforeWorkerFinalization = await waitForStoreWorkCompletion( + reviewCompletion, + timeout: .milliseconds(500) + ) + + #expect(returnedBeforeWorkerFinalization == false) + await startGate.open() + #expect(await close.value == .success) + let result = try await review.value + + #expect(result.core.lifecycle.status == .cancelled) + #expect(result.core.lifecycle.cancellation == reason) + #expect(store.reviewWorkerTasks["job-1"] == nil) + #expect(store.storeWorkRegistry.activeOrdinals.isEmpty) + await store.stop() + } + + @Test func registeredWorkCloseAppliesPreEntryCancellationPolicy() async throws { + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend() + ) + ) + var skippedOperationRan = false + var finalizerRan = false + var finalizedOperationRan = false + let skippedTask = try #require(store.startRegisteredStoreWork( + kind: .testing("skip before entry") + ) { _ in + skippedOperationRan = true + }) + let finalizedTask = try #require(store.startRegisteredStoreWork( + kind: .testing("finalize before entry"), + cancelledBeforeEntry: .runFinalizer { _ in + finalizerRan = true + } + ) { _ in + finalizedOperationRan = true + }) + let closeOperation = store.storeWorkRegistry.beginClosing(onAdmissionClosed: {}) + let result = await closeOperation.task.value + store.storeWorkRegistry.completeClosing(closeOperation, result: result) + await skippedTask.value + await finalizedTask.value + + #expect(result == .success) + #expect(skippedOperationRan == false) + #expect(finalizerRan) + #expect(finalizedOperationRan == false) + #expect(store.storeWorkRegistry.activeOrdinals.isEmpty) + } + + @Test func individualCancellationAppliesRegisteredWorkPreEntryPolicy() async { + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend( + reviewBackend: FakeCodexReviewBackend() + ) + ) + var skippedOperationRan = false + var finalizerRunCount = 0 + var finalizedOperationRan = false + let skippedTask = store.startRegisteredStoreWork( + kind: .testing("individually cancelled skip") + ) { _ in + skippedOperationRan = true + } + let finalizedTask = store.startRegisteredStoreWork( + kind: .testing("individually cancelled finalizer"), + cancelledBeforeEntry: .runFinalizer { _ in + finalizerRunCount += 1 + } + ) { _ in + finalizedOperationRan = true + } + skippedTask?.cancel() + finalizedTask?.cancel() + + await skippedTask?.value + await finalizedTask?.value + + #expect(skippedOperationRan == false) + #expect(finalizerRunCount == 1) + #expect(finalizedOperationRan == false) + #expect(store.storeWorkRegistryStatus == .open) + #expect(store.storeWorkRegistry.activeOrdinals.isEmpty) + + var throwingOperationRan = false + let throwingTask = Task { @MainActor in + try await store.performThrowingRegisteredStoreWork( + kind: .testing("individually cancelled throwing work") + ) { _ in + throwingOperationRan = true + } + } + throwingTask.cancel() + + await #expect(throws: CancellationError.self) { + try await throwingTask.value + } + #expect(throwingOperationRan == false) + #expect(store.storeWorkRegistryStatus == .open) + #expect(store.storeWorkRegistry.activeOrdinals.isEmpty) + } + @Test func staleRuntimeFailureCannotTearDownFreshGeneration() async throws { let backend = TestingCodexReviewStoreBackend(reviewBackend: FakeCodexReviewBackend()) let store = CodexReviewStore.makeTestingStore(backend: backend) @@ -454,6 +746,48 @@ struct CodexReviewStoreLifecycleTests { } } +private enum StoreWorkTestFailure: LocalizedError, Sendable { + case first + case second + + var errorDescription: String? { + switch self { + case .first: + "first work failed" + case .second: + "second work failed" + } + } +} + +private actor StoreWorkCompletion { + private var completed = false + + func complete() { + completed = true + } + + func isComplete() -> Bool { + completed + } +} + +@MainActor +private func waitForStoreWorkCompletion( + _ completion: StoreWorkCompletion, + timeout: Duration +) async -> Bool { + let clock = ContinuousClock() + let deadline = clock.now + timeout + while await completion.isComplete() == false { + guard clock.now < deadline else { + return false + } + await Task.yield() + } + return true +} + @MainActor private func waitForCutoverStatus( _ expected: CodexReviewSettingsService.RuntimeCutoverStatus, @@ -483,3 +817,49 @@ private func waitForTeardownFinalState( await Task.yield() } } + +@MainActor +private func waitForStoreWorkStatus( + _ expected: ReviewStoreWorkRegistryStatus, + store: CodexReviewStore +) async throws { + let clock = ContinuousClock() + let deadline = clock.now + .seconds(2) + while store.storeWorkRegistryStatus != expected { + guard clock.now < deadline else { + throw CancellationError() + } + await Task.yield() + } +} + +@MainActor +private func waitForReviewWorkerCancellation( + jobID: String, + store: CodexReviewStore +) async throws { + let clock = ContinuousClock() + let deadline = clock.now + .seconds(2) + while store.reviewWorkerTasks[jobID]?.isCancelled != true { + guard clock.now < deadline else { + throw CancellationError() + } + await Task.yield() + } +} + +@MainActor +private func waitForReviewWaiterCount( + _ count: Int, + jobID: String, + store: CodexReviewStore +) async throws { + let clock = ContinuousClock() + let deadline = clock.now + .seconds(2) + while store.reviewTerminalWaiters[jobID]?.count != count { + guard clock.now < deadline else { + throw CancellationError() + } + await Task.yield() + } +} diff --git a/Tests/CodexReviewTests/CodexReviewStoreRateLimitAutoRefreshTests.swift b/Tests/CodexReviewTests/CodexReviewStoreRateLimitAutoRefreshTests.swift index 4104210..8b4e1b2 100644 --- a/Tests/CodexReviewTests/CodexReviewStoreRateLimitAutoRefreshTests.swift +++ b/Tests/CodexReviewTests/CodexReviewStoreRateLimitAutoRefreshTests.swift @@ -214,6 +214,69 @@ struct CodexReviewStoreRateLimitAutoRefreshTests { } } + @Test func registeredWorkCloseCancelsAndAwaitsRateLimitRefresh() 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 close = Task { @MainActor in + await store.closeRegisteredStoreWork( + reason: .system(message: "Store work closed.") + ) + } + try await waitForCondition { + store.storeWorkRegistryStatus == .closing + } + + #expect(backend.refreshCompletionCount == 0) + #expect(store.accountRateLimitAutoRefreshInFlightAccountKeys == [account.accountKey]) + + await backend.releaseRefresh() + #expect(await close.value == .success) + + #expect(backend.refreshCompletionCount == 1) + #expect(store.accountRateLimitAutoRefreshInFlightAccountKeys.isEmpty) + store.refreshDueAccountRateLimits(now: now) + #expect(backend.refreshedAccountKeys == [account.accountKey]) + } + + @Test func registeredWorkCloseSkipsRateLimitRefreshCancelledBeforeEntry() async throws { + let account = makeAccount(lastFetchAt: now.addingTimeInterval(-15 * 60)) + let backend = NoProgressRateLimitRefreshBackend(account: account) + let store = CodexReviewStore.makeTestingStore( + backend: backend, + clock: .init(now: { now }) + ) + store.loadForTesting( + serverState: .running, + authPhase: .signedOut, + account: account, + persistedAccounts: [account], + workspaces: [] + ) + + store.refreshDueAccountRateLimits(now: now) + let closeOperation = store.storeWorkRegistry.beginClosing { + store.accountRateLimitAutoRefreshDriver?.closeAdmission() + } + let result = await closeOperation.task.value + await store.cancelAccountRateLimitAutoRefreshAndWait() + store.storeWorkRegistry.completeClosing(closeOperation, result: result) + + #expect(result == .success) + #expect(backend.refreshedAccountKeys.isEmpty) + #expect(store.accountRateLimitAutoRefreshInFlightAccountKeys.isEmpty) + } + @Test func noProgressRefreshDoesNotImmediatelyRestartSameAccount() async throws { let account = makeAccount(lastFetchAt: now.addingTimeInterval(-15 * 60)) let backend = NoProgressRateLimitRefreshBackend(account: account) @@ -379,6 +442,7 @@ private final class BlockingRateLimitRefreshBackend: PreviewCodexReviewStoreBack private let startedGate = AsyncGate() private let releaseGate = AsyncGate() private(set) var refreshedAccountKeys: [String] = [] + private(set) var refreshCompletionCount = 0 init(account: CodexAccount) { super.init(seed: .init( @@ -393,7 +457,8 @@ private final class BlockingRateLimitRefreshBackend: PreviewCodexReviewStoreBack ) async { refreshedAccountKeys.append(accountKey) await startedGate.open() - await releaseGate.wait() + await releaseGate.waitIgnoringCancellation() + refreshCompletionCount += 1 } func waitUntilRefreshStarts() async {