-
Notifications
You must be signed in to change notification settings - Fork 0
Update Codex runtime while preserving queued reviews and MCP sessions #387
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
Show all changes
8 commits
Select commit
Hold shift + click to select a range
ddfa08b
Update Codex runtime while retaining queued reviews and MCP sessions
lynnswap 6941e40
Preserve update recovery across retries and history deletion
lynnswap 7854079
Use one publication path for stopped-runtime updates
lynnswap 558a0f9
Cancel queued reviews independently of runtime teardown state
lynnswap 7018857
Join Codex updates after startup history loading
lynnswap b90e17c
Cancel update dependency waits and honor pending runtime stops
lynnswap bb0c3da
Confirm current process closure when retrying failed updates
lynnswap 8ff02db
Request runtime closure before confirming recovery
lynnswap File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Large diffs are not rendered by default.
Oops, something went wrong.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,200 @@ | ||
| import Foundation | ||
|
|
||
| extension CodexReviewStore { | ||
| @_spi(ApplicationHostSupport) public enum CodexUpdateTiming: Sendable { | ||
| case afterCurrentReviews | ||
| case immediately | ||
| } | ||
|
|
||
| @_spi(ApplicationHostSupport) public enum CodexUpdateState: Equatable, Sendable { | ||
| case idle | ||
| case waitingForReviews | ||
| case stoppingRuntime | ||
| case installing | ||
| case restarting | ||
| case failed(String) | ||
| } | ||
|
|
||
| /// Keeps MCP sessions and accepted jobs alive while replacing the Codex runtime. | ||
| /// Concurrent callers join the existing operation; only its installation closure runs. | ||
| /// An installation failure is thrown even when the runtime recovers and jobs resume. | ||
| @_spi(ApplicationHostSupport) public func updateCodex( | ||
| when timing: CodexUpdateTiming, | ||
| install: @escaping @MainActor @Sendable () async throws -> Void | ||
| ) async throws { | ||
| if let task = codexUpdateTask { | ||
| try await task.value | ||
| return | ||
| } | ||
| guard applicationShutdownRequested == false, pendingRuntimeStopCount == 0 else { throw CancellationError() } | ||
| let previousAccountOperations = Array(runtimeAccountOperations.values) | ||
| let previousRuntimeTask: Task<Void, Never>? = switch runtimeState { | ||
| case .acquiring(_, _, let task), .replacing(_, let task), .tearingDown(_, _, _, _, let task): task | ||
| case .running, .stopped, .failed: nil | ||
| } | ||
| suspendReviewStarts() | ||
| codexUpdate = .waitingForReviews | ||
| let task = Task { @MainActor [self] in | ||
| defer { | ||
| codexUpdateTask = nil | ||
| if Task.isCancelled == false { resumeReviewsAfterCodexUpdateIfPossible() } | ||
| } | ||
| do { | ||
| try await Self.waitForPrecedingRuntimeWork( | ||
| accountOperations: previousAccountOperations, runtimeTask: previousRuntimeTask | ||
| ) | ||
| try await performCodexUpdate(when: timing, install: install) | ||
| codexUpdate = .idle | ||
| } catch { | ||
| if error is CancellationError { | ||
| codexUpdate = .idle | ||
| } else { | ||
| codexUpdate = .failed(error.localizedDescription) | ||
| } | ||
| throw error | ||
| } | ||
| } | ||
| codexUpdateTask = task | ||
| try await task.value | ||
| } | ||
|
|
||
| private static func waitForPrecedingRuntimeWork( | ||
| accountOperations: [Task<Void, any Error>], | ||
| runtimeTask: Task<Void, Never>? | ||
| ) async throws { | ||
| // Cancelling the update abandons its join, not the predecessor's lifecycle ownership. | ||
| let (completion, continuation) = AsyncStream<Void>.makeStream() | ||
| let waiter = Task { | ||
| for operation in accountOperations { | ||
| _ = try? await operation.value | ||
| guard Task.isCancelled == false else { return } | ||
| } | ||
| await runtimeTask?.value | ||
| continuation.yield(()) | ||
| continuation.finish() | ||
| } | ||
| defer { | ||
| waiter.cancel() | ||
| continuation.finish() | ||
| } | ||
| var iterator = completion.makeAsyncIterator() | ||
| _ = await iterator.next() | ||
| try Task.checkCancellation() | ||
| } | ||
|
|
||
| private func performCodexUpdate( | ||
| when timing: CodexUpdateTiming, | ||
| install: @escaping @MainActor @Sendable () async throws -> Void | ||
| ) async throws { | ||
| // A dispatcher already saving an execution start belongs to the current batch. | ||
| await queuedReviewDispatchTask?.value | ||
| try Task.checkCancellation() | ||
| if timing == .afterCurrentReviews { | ||
| let executingIDs = jobs.filter { | ||
| $0.isTerminal == false && queuedReviewStarts[$0.id] == nil | ||
| }.map(\.id) | ||
| for id in executingIDs { | ||
| // Another completed review can be removed while we await this batch. | ||
| guard job(id: id)?.isTerminal == false else { continue } | ||
| _ = try await awaitReview(sessionID: nil, jobID: id) | ||
| try Task.checkCancellation() | ||
| } | ||
| _ = await drainReviewWorkersForRuntimeStop(timeout: backend.shutdownCleanupTimeout) | ||
| try Task.checkCancellation() | ||
| } | ||
| guard applicationShutdownRequested == false else { throw CancellationError() } | ||
| codexUpdate = .stoppingRuntime | ||
| var installationFailure: String? | ||
| let installation: @MainActor @Sendable () async -> Void = { [self] in | ||
| codexUpdate = .installing | ||
| do { try await install() } catch { installationFailure = error.localizedDescription } | ||
| codexUpdate = .restarting | ||
| } | ||
| let operation: RuntimeStartOperation | ||
| switch runtimeState { | ||
| case .running(let generation, let runtime, let mcp): | ||
| operation = admitRuntimeReplacement( | ||
| sourceGeneration: generation, retiringRuntime: runtime, retainedMCP: mcp, | ||
| preservingQueuedReviews: true, | ||
| install: installation | ||
| ) | ||
| case .failed(let generation, let mcp?, _): | ||
| operation = admitRuntimeReplacement( | ||
| sourceGeneration: generation, retiringRuntime: unclosedCodexUpdateRuntime, retainedMCP: mcp, | ||
| preservingQueuedReviews: true, | ||
| install: installation | ||
| ) | ||
| case .stopped(let generation), .failed(let generation, nil, _): | ||
| try await closeUnclosedCodexUpdateRuntimeIfNeeded() | ||
| await installation() | ||
| let preparation = try await backend.mcpServerLifecycle.prepare() | ||
| let snapshot: MCPServerPublicationSnapshot | ||
| do { | ||
| snapshot = try await backend.mcpServerLifecycle.activate(preparation) | ||
| } catch { | ||
| var message = error.localizedDescription | ||
| do { try await backend.mcpServerLifecycle.stop() } | ||
| catch { message += "; MCP cleanup failed: \(error.localizedDescription)" } | ||
| throw CodexReviewAPI.Error.io(message) | ||
| } | ||
| operation = admitRuntimeReplacement( | ||
| sourceGeneration: generation, retiringRuntime: nil, | ||
| retainedMCP: RetainedMCPServer(serverURL: snapshot.serverURL), | ||
| preservingQueuedReviews: true | ||
| ) | ||
| case .acquiring, .replacing, .tearingDown: | ||
| throw CodexReviewAPI.Error.io("The Codex runtime changed while waiting to update.") | ||
| } | ||
| await operation.task.value | ||
| let failures = [installationFailure, serverState.failureMessage].compactMap { $0 } | ||
| if failures.isEmpty == false { | ||
| throw CodexReviewAPI.Error.io(failures.joined(separator: "; ")) | ||
| } | ||
| guard case .running = runtimeState else { | ||
| throw CodexReviewAPI.Error.io("Codex update finished without an available runtime.") | ||
| } | ||
| } | ||
|
|
||
| package func resumeReviewsAfterCodexUpdateIfPossible() { | ||
| guard runtimeAccountOperations.isEmpty, applicationShutdownRequested == false, | ||
| pendingRuntimeStopCount == 0, | ||
| case .running = runtimeState else { return } | ||
| resumeReviewStarts() | ||
| } | ||
|
|
||
| package func performRuntimeAuthentication( | ||
| _ operation: @escaping @MainActor @Sendable (CodexReviewStore) async -> Void | ||
| ) async { | ||
| do { try await performRuntimeAccountChange { store in await operation(store) } } | ||
| catch is CancellationError { } | ||
| catch { auth.updatePhase(.failed(message: error.localizedDescription)) } | ||
| } | ||
|
|
||
| package func performRuntimeAccountChange( | ||
| _ operation: @escaping @MainActor @Sendable (CodexReviewStore) async throws -> Void | ||
| ) async throws { | ||
| let update = codexUpdateTask | ||
| let id = UUID() | ||
| let task = Task { @MainActor [self] in | ||
| _ = try? await update?.value | ||
| defer { | ||
| runtimeAccountOperations.removeValue(forKey: id) | ||
| if codexUpdateTask == nil { resumeReviewsAfterCodexUpdateIfPossible() } | ||
| } | ||
| try Task.checkCancellation() | ||
| guard applicationShutdownRequested == false else { throw CancellationError() } | ||
| try await operation(self) | ||
| } | ||
| runtimeAccountOperations[id] = task | ||
| try await withTaskCancellationHandler { | ||
| try await task.value | ||
| } onCancel: { | ||
| task.cancel() | ||
| } | ||
| } | ||
|
|
||
| package var hasQueuedCodexUpdateRecovery: Bool { | ||
| if case .failed = codexUpdateState { return reviewStartsAreSuspended } | ||
| return false | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.