From 6976a4c6e4c51c55e8217f8ce228192e41095bca Mon Sep 17 00:00:00 2001 From: Paulo Nascimento Date: Sat, 11 Jul 2026 19:42:56 -0400 Subject: [PATCH] Avoid MCP proxy task teardown abort --- .../Fixtures/test-suite-contract-ledger.tsv | 6 + Sources/RepoPromptMCP/main.swift | 134 +++++++-- .../MCP/MCPServiceProxyRaceTests.swift | 256 ++++++++++++++++++ 3 files changed, 366 insertions(+), 30 deletions(-) create mode 100644 Tests/RepoPromptTests/MCP/MCPServiceProxyRaceTests.swift diff --git a/Scripts/Fixtures/test-suite-contract-ledger.tsv b/Scripts/Fixtures/test-suite-contract-ledger.tsv index ded3e2ead..4ef51fe9f 100644 --- a/Scripts/Fixtures/test-suite-contract-ledger.tsv +++ b/Scripts/Fixtures/test-suite-contract-ledger.tsv @@ -239,6 +239,7 @@ root/RepoPromptTests.AgentModeViewModelInactiveRefreshTests/testActivationRepubl root/RepoPromptTests.AgentModeViewModelInactiveRefreshTests/testActiveCodexAssistantDeltaUsesPresentationOnlyRefresh root Tests/RepoPromptTests/AgentMode/AgentModeViewModelInactiveRefreshTests.swift RepoPromptTests.AgentModeViewModelInactiveRefreshTests testActiveCodexAssistantDeltaUsesPresentationOnlyRefresh AgentMode unreviewed unreviewed root_swiftpm routine 1 unreviewed unreviewed 0.005000 unreviewed retain_pending_review 0 initial census source line 376 root/RepoPromptTests.AgentModeViewModelInactiveRefreshTests/testActiveGenericAssistantDeltaPublishesTranscriptWithoutFullBindingSync root Tests/RepoPromptTests/AgentMode/AgentModeViewModelInactiveRefreshTests.swift RepoPromptTests.AgentModeViewModelInactiveRefreshTests testActiveGenericAssistantDeltaPublishesTranscriptWithoutFullBindingSync AgentMode unreviewed unreviewed root_swiftpm routine 1 unreviewed unreviewed 0.003000 unreviewed retain_pending_review 0 initial census source line 349 root/RepoPromptTests.AgentModeViewModelInactiveRefreshTests/testActiveRefreshCompactsSummaryOnlyToolResultSourceWhilePreservingRawRenderPayload root Tests/RepoPromptTests/AgentMode/AgentModeViewModelInactiveRefreshTests.swift RepoPromptTests.AgentModeViewModelInactiveRefreshTests testActiveRefreshCompactsSummaryOnlyToolResultSourceWhilePreservingRawRenderPayload AgentMode unreviewed unreviewed root_swiftpm routine 1 unreviewed unreviewed 0.003000 unreviewed retain_pending_review 0 initial census source line 7 +root/RepoPromptTests.AgentModeViewModelInactiveRefreshTests/testRetainedPayloadRefreshChangesOnlyMatchingResultRenderIdentity root Tests/RepoPromptTests/AgentMode/AgentModeViewModelInactiveRefreshTests.swift RepoPromptTests.AgentModeViewModelInactiveRefreshTests testRetainedPayloadRefreshChangesOnlyMatchingResultRenderIdentity AgentMode agent.transcript.raw_payload_item_identity tool_result,render_revision,item_identity,payload_retention deterministic_regression root_swiftpm routine 1 Refreshing one retained raw tool-result payload changes only that item's render revision and payload while leaving the sibling result unchanged. A tool-result refresh could invalidate or display another result's raw payload when multiple results share a transcript. 0.003000 test_case retain 0 Current-main ledger reconciliation for commit 1d783011. root/RepoPromptTests.AgentModeViewModelInactiveRefreshTests/testAmbiguousPersistentBindingFailsRoutingWithoutLeakingCandidates root Tests/RepoPromptTests/AgentMode/AgentModeViewModelInactiveRefreshTests.swift RepoPromptTests.AgentModeViewModelInactiveRefreshTests testAmbiguousPersistentBindingFailsRoutingWithoutLeakingCandidates AgentMode unreviewed unreviewed root_swiftpm routine 1 unreviewed unreviewed 0.001500 unreviewed retain_pending_review 0 initial census source line 597 root/RepoPromptTests.AgentModeViewModelInactiveRefreshTests/testAssistantPresentationRejectsStaleRevisionGenerationAndTabOwnership root Tests/RepoPromptTests/AgentMode/AgentModeViewModelInactiveRefreshTests.swift RepoPromptTests.AgentModeViewModelInactiveRefreshTests testAssistantPresentationRejectsStaleRevisionGenerationAndTabOwnership AgentMode unreviewed unreviewed root_swiftpm routine 1 unreviewed unreviewed 0.003000 unreviewed retain_pending_review 0 initial census source line 425 root/RepoPromptTests.AgentModeViewModelInactiveRefreshTests/testAssistantPresentationRequiresAuthoritativeHydratedBindingGeneration root Tests/RepoPromptTests/AgentMode/AgentModeViewModelInactiveRefreshTests.swift RepoPromptTests.AgentModeViewModelInactiveRefreshTests testAssistantPresentationRequiresAuthoritativeHydratedBindingGeneration AgentMode unreviewed unreviewed root_swiftpm routine 1 unreviewed unreviewed 0.004000 unreviewed retain_pending_review 0 initial census source line 473 @@ -1835,6 +1836,9 @@ root/RepoPromptTests.MCPProxyTerminalRecordTests/testLocalSocketReadFailureIsAtt root/RepoPromptTests.MCPProxyTerminalRecordTests/testProxyTaskGroupOutcomesClassifyCleanCompletionAndCancellation root Tests/RepoPromptTests/MCP/MCPProxyTerminalRecordTests.swift RepoPromptTests.MCPProxyTerminalRecordTests testProxyTaskGroupOutcomesClassifyCleanCompletionAndCancellation MCP mcp.proxy.task_group_outcome_classification proxy,cancellation,outcome_classification deterministic_regression root_swiftpm routine 1 MCPServiceProxyTaskGroupPolicy Clean transport completion stays clean while watcher and service cancellation outcomes normalize to host task-cancelled provenance. A canceled proxy could persist a false clean terminal record and hide lifecycle failures. 0.000000 test_case retain 0 PR #252 final proxy cancellation cleanup root/RepoPromptTests.MCPProxyTerminalRecordTests/testSocketReadFailureRemainsRetryablePeerTransportFailure root Tests/RepoPromptTests/MCP/MCPProxyTerminalRecordTests.swift RepoPromptTests.MCPProxyTerminalRecordTests testSocketReadFailureRemainsRetryablePeerTransportFailure MCP unreviewed unreviewed root_swiftpm routine 1 unreviewed unreviewed 0.000000 unreviewed retain_pending_review 0 initial census source line 60 root/RepoPromptTests.MCPProxyTerminalRecordTests/testTerminalRecordCopiesLiveLedgerSnapshotAndServerReason root Tests/RepoPromptTests/MCP/MCPProxyTerminalRecordTests.swift RepoPromptTests.MCPProxyTerminalRecordTests testTerminalRecordCopiesLiveLedgerSnapshotAndServerReason MCP unreviewed unreviewed root_swiftpm routine 1 unreviewed unreviewed 0.000000 unreviewed retain_pending_review 0 initial census source line 128 +root/RepoPromptTests.MCPServiceProxyRaceTests/testEachChildCanWinWithValueOrErrorAndDrainsLosers root Tests/RepoPromptTests/MCP/MCPServiceProxyRaceTests.swift RepoPromptTests.MCPServiceProxyRaceTests testEachChildCanWinWithValueOrErrorAndDrainsLosers MCP mcp.proxy.first_completion_priority proxy,cancellation,error_priority,task_drain deterministic_regression root_swiftpm fast 6 Each of the three proxy operations can win with its value or error, and both losing operations observe cancellation and complete before the race returns. Changing completion priority or failing to drain a suspended loser could misclassify terminal state or reintroduce the Swift release-runtime teardown abort. 0.001000 test_case retain 0 Issue #513 explicit-task first-completion race coverage. +root/RepoPromptTests.MCPServiceProxyRaceTests/testOuterCancellationCancelsAndDrainsAllChildren root Tests/RepoPromptTests/MCP/MCPServiceProxyRaceTests.swift RepoPromptTests.MCPServiceProxyRaceTests testOuterCancellationCancelsAndDrainsAllChildren MCP mcp.proxy.outer_cancellation_drain proxy,cancellation,task_drain deterministic_regression root_swiftpm fast 1 Caller cancellation reaches all three proxy operations, surfaces as CancellationError, and waits for every operation to complete before returning. Outer cancellation could leak or tear down a suspended proxy operation and persist incorrect terminal provenance. 0.001000 test_case retain 0 Issue #513 cancellation-handler teardown coverage. +root/RepoPromptTests.MCPServiceProxyRaceTests/testRaceDoesNotReturnUntilCancelledLosersComplete root Tests/RepoPromptTests/MCP/MCPServiceProxyRaceTests.swift RepoPromptTests.MCPServiceProxyRaceTests testRaceDoesNotReturnUntilCancelledLosersComplete MCP mcp.proxy.loser_completion_barrier proxy,cancellation,task_drain deterministic_regression root_swiftpm fast 1 A controlled loser-completion gate proves the proxy race remains suspended after cancellation until both losing operations finish. Returning after cancel without awaiting loser completion recreates the compiler-frame teardown shape implicated in swift_task_dealloc aborts. 0.001000 test_case retain 0 Issue #513 structural regression oracle; fails against cancel-without-await. root/RepoPromptTests.MCPReadFileAutoSelectionCoordinatorTests/testAuthoritativeSelectionPreservationRequiresFullCanonicalSuperset root Tests/RepoPromptTests/MCP/MCPReadFileAutoSelectionCoordinatorTests.swift RepoPromptTests.MCPReadFileAutoSelectionCoordinatorTests testAuthoritativeSelectionPreservationRequiresFullCanonicalSuperset MCP mcp.read_file.auto_selection.authoritative_selection_preservation exact,additive_superset,full_supersedes_slice,selected_path_loss,full_to_slice_demotion,auto_codemap_loss,slice_narrowing,codemap_mode_change deterministic_table_unit root_swiftpm routine 8 stored_selection Certificate mint revalidation accepts only a full canonical superset of the authoritative result across selected paths, auto-codemap paths, slices, and codemap mode. critical 0.000000 in_memory per_test_state retain 0 read-file selection rebasing full-state preservation guard root/RepoPromptTests.MCPReadFileAutoSelectionCoordinatorTests/testCancelledCanonicalDrainResumesPromptlyWithoutStoppingWorker root Tests/RepoPromptTests/MCP/MCPReadFileAutoSelectionCoordinatorTests.swift RepoPromptTests.MCPReadFileAutoSelectionCoordinatorTests testCancelledCanonicalDrainResumesPromptlyWithoutStoppingWorker MCP unreviewed unreviewed root_swiftpm routine 1 unreviewed unreviewed 0.004500 unreviewed retain_pending_review 0 initial census source line 141 root/RepoPromptTests.MCPReadFileAutoSelectionCoordinatorTests/testCancelledMirrorDrainResumesPromptlyWithoutStoppingWorker root Tests/RepoPromptTests/MCP/MCPReadFileAutoSelectionCoordinatorTests.swift RepoPromptTests.MCPReadFileAutoSelectionCoordinatorTests testCancelledMirrorDrainResumesPromptlyWithoutStoppingWorker MCP unreviewed unreviewed root_swiftpm routine 1 unreviewed unreviewed 0.004500 unreviewed retain_pending_review 0 initial census source line 175 @@ -2250,6 +2254,8 @@ root/RepoPromptTests.SettingsJSONOnlyPersistenceTests/testMissingGlobalSettingsC root/RepoPromptTests.SettingsJSONOnlyPersistenceTests/testObsoleteGitignoreJSONKeyIsIgnoredAndNeverEmitted root Tests/RepoPromptTests/SettingsJSONOnlyPersistenceTests.swift RepoPromptTests.SettingsJSONOnlyPersistenceTests testObsoleteGitignoreJSONKeyIsIgnoredAndNeverEmitted Root unreviewed unreviewed root_swiftpm routine 1 unreviewed unreviewed 0.001000 unreviewed retain_pending_review 0 initial census source line 32 root/RepoPromptTests.SettingsJSONOnlyPersistenceTests/testSentryScrubStringRedactsAuthorizationSchemesAndCredentialTails root Tests/RepoPromptTests/SettingsJSONOnlyPersistenceTests.swift RepoPromptTests.SettingsJSONOnlyPersistenceTests testSentryScrubStringRedactsAuthorizationSchemesAndCredentialTails Root telemetry.sentry_scrub.authorization_and_credential_tail_redaction sentry,privacy,authorization,bearer,basic,api_key privacy_redaction_unit root_swiftpm routine 6 Sentry scrubbing redacts Bearer and Basic authorization headers, lowercase authorization assignments, bare Bearer credential tails, X-Api-Key values, and preserves benign text. Telemetry breadcrumbs could leak authorization credentials or over-redact harmless prose. test_case retain 0 origin/main reconciliation after commits 5c1504f1/dc79e7a4: reviewed Sentry authorization-scheme redaction coverage. root/RepoPromptTests.SettingsJSONOnlyPersistenceTests/testSentryScrubStringRedactsSensitiveValues root Tests/RepoPromptTests/SettingsJSONOnlyPersistenceTests.swift RepoPromptTests.SettingsJSONOnlyPersistenceTests testSentryScrubStringRedactsSensitiveValues Root telemetry.sentry_scrub.sensitive_value_redaction sentry,privacy,token,password,user_path,ip_address privacy_redaction_unit root_swiftpm routine 4 Sentry scrubbing redacts token values, password values, the current user home path, and private IP addresses while preserving redaction markers. Telemetry payloads could expose local secrets, usernames, or network identifiers. test_case retain 0 origin/main reconciliation after commits 5c1504f1/dc79e7a4: reviewed Sentry sensitive-value redaction coverage. +root/RepoPromptTests.SettingsJSONOnlyPersistenceTests/testSentryScrubPayloadDropsNestedRequestAndGeoFieldsWithoutDroppingSafeSiblings root Tests/RepoPromptTests/SettingsJSONOnlyPersistenceTests.swift RepoPromptTests.SettingsJSONOnlyPersistenceTests testSentryScrubPayloadDropsNestedRequestAndGeoFieldsWithoutDroppingSafeSiblings Root telemetry.sentry_scrub.nested_sensitive_keys sentry,privacy,nested_payload,request,geo privacy_redaction_unit root_swiftpm routine 1 Recursive Sentry payload scrubbing drops nested URL, query-string, and user-geo fields while preserving safe sibling values. Nested telemetry dictionaries could leak request or location data, or privacy filtering could discard unrelated diagnostic fields. 0.001000 test_case retain 0 Current-main ledger reconciliation for commit 296f6787. +root/RepoPromptTests.SettingsJSONOnlyPersistenceTests/testSentryScrubPayloadDropsRequestGeoUserAndStableDeviceIdentifiers root Tests/RepoPromptTests/SettingsJSONOnlyPersistenceTests.swift RepoPromptTests.SettingsJSONOnlyPersistenceTests testSentryScrubPayloadDropsRequestGeoUserAndStableDeviceIdentifiers Root telemetry.sentry_scrub.identity_and_request_fields sentry,privacy,request,user,device_identifier privacy_redaction_unit root_swiftpm routine 1 Sentry payload scrubbing removes request, user, geo, installation, and stable device identifiers while preserving safe OS, model, app, and redacted-note fields. Telemetry events could expose stable user or device identity and request data, or over-redaction could remove safe operational context. 0.001000 test_case retain 0 Current-main ledger reconciliation for commit 296f6787. root/RepoPromptTests.SettingsJSONOnlyPersistenceTests/testShowDatesInMessageTimestampsDefaultsFalseWithoutPersisting root Tests/RepoPromptTests/SettingsJSONOnlyPersistenceTests.swift RepoPromptTests.SettingsJSONOnlyPersistenceTests testShowDatesInMessageTimestampsDefaultsFalseWithoutPersisting Settings settings.message_timestamp_dates.default_migration optional_json_field,default_false,no_read_time_persist deterministic_persistence root_swiftpm routine 1 temporary_settings_file,isolated_user_defaults A missing optional field reads false while the persisted JSON remains byte-for-byte unchanged. medium 0.001000 temporary_directory,filesystem,user_defaults method_local_values+defer retain 0 backward-compatible default contract root/RepoPromptTests.SettingsJSONOnlyPersistenceTests/testShowDatesInMessageTimestampsSavesAndLoadsTrue root Tests/RepoPromptTests/SettingsJSONOnlyPersistenceTests.swift RepoPromptTests.SettingsJSONOnlyPersistenceTests testShowDatesInMessageTimestampsSavesAndLoadsTrue Settings settings.message_timestamp_dates.round_trip json_round_trip,reload deterministic_persistence root_swiftpm routine 1 temporary_settings_file,isolated_user_defaults Enabling date context persists true and a newly initialized store reloads true. medium 0.001500 temporary_directory,filesystem,user_defaults method_local_values+defer retain 0 settings round-trip contract root/RepoPromptTests.SettingsJSONOnlyPersistenceTests/testTelemetryJSONValueOverridesStaleEnabledMirror root Tests/RepoPromptTests/SettingsJSONOnlyPersistenceTests.swift RepoPromptTests.SettingsJSONOnlyPersistenceTests testTelemetryJSONValueOverridesStaleEnabledMirror Root settings.telemetry.json_overrides_stale_enabled_mirror settings_json,telemetry,user_defaults_mirror persistence_regression root_swiftpm routine 1 temporary_settings_file,UserDefaults(suiteName:) A persisted telemetry JSON value of false overrides a stale legacy defaults mirror of true and rewrites the mirror to false. Legacy defaults could override the canonical JSON telemetry value. filesystem isolated_user_defaults test_case retain 0 origin/main reconciliation after commits 5c1504f1/dc79e7a4: reviewed telemetry JSON-over-mirror precedence coverage. diff --git a/Sources/RepoPromptMCP/main.swift b/Sources/RepoPromptMCP/main.swift index 36ff4fa60..0e8771997 100644 --- a/Sources/RepoPromptMCP/main.swift +++ b/Sources/RepoPromptMCP/main.swift @@ -783,6 +783,46 @@ enum MCPServiceProxyTaskGroupPolicy { } } +private final class MCPServiceProxySettlementState: @unchecked Sendable { + enum Outcome { + case completed(Result) + case callerCancelled + } + + private let lock = NSLock() + private var outcome: Outcome? + private var continuation: CheckedContinuation? + + func claim(_ candidate: Outcome) -> Bool { + lock.lock() + guard outcome == nil else { + lock.unlock() + return false + } + outcome = candidate + let continuation = continuation + self.continuation = nil + lock.unlock() + + continuation?.resume(returning: candidate) + return true + } + + func finish() async -> Outcome { + await withCheckedContinuation { continuation in + lock.lock() + if let outcome { + lock.unlock() + continuation.resume(returning: outcome) + } else { + precondition(self.continuation == nil) + self.continuation = continuation + lock.unlock() + } + } + } +} + // MARK: - Bootstrap Socket Proxy (CLI connects to App) /// Bootstrap socket proxy - connects to app's single socket server. @@ -2100,6 +2140,58 @@ actor MCPService: Service { killSignalContinuation = nil } + static func awaitFirstProxyOutcome( + killSignal: @escaping @Sendable () async throws -> MCPServiceProxyTaskOutcome, + watchdog: @escaping @Sendable () async throws -> MCPServiceProxyTaskOutcome, + transport: @escaping @Sendable () async throws -> MCPServiceProxyTaskOutcome + ) async throws -> MCPServiceProxyTaskOutcome { + let settlement = MCPServiceProxySettlementState() + let killSignalTask = Task { + do { + _ = try await settlement.claim(.completed(.success(killSignal()))) + } catch { + _ = settlement.claim(.completed(.failure(error))) + } + } + let watchdogTask = Task { + do { + _ = try await settlement.claim(.completed(.success(watchdog()))) + } catch { + _ = settlement.claim(.completed(.failure(error))) + } + } + let transportTask = Task { + do { + _ = try await settlement.claim(.completed(.success(transport()))) + } catch { + _ = settlement.claim(.completed(.failure(error))) + } + } + + let outcome = await withTaskCancellationHandler { + await settlement.finish() + } onCancel: { + _ = settlement.claim(.callerCancelled) + killSignalTask.cancel() + watchdogTask.cancel() + transportTask.cancel() + } + + killSignalTask.cancel() + watchdogTask.cancel() + transportTask.cancel() + await killSignalTask.value + await watchdogTask.value + await transportTask.value + + switch outcome { + case let .completed(result): + return try result.get() + case .callerCancelled: + throw CancellationError() + } + } + func run() async throws { // Set up kill signal watcher before starting transport setupKillSignalWatcher() @@ -2110,9 +2202,8 @@ actor MCPService: Service { do { // Race between transport loop, kill signal, and PPID watchdog - try await withThrowingTaskGroup(of: MCPServiceProxyTaskOutcome.self) { group in - // Kill signal monitor task - group.addTask { + let outcome = try await Self.awaitFirstProxyOutcome( + killSignal: { guard let signal = await self.waitForKillSignal() else { return .killSignalWaitCancelled } @@ -2120,38 +2211,21 @@ actor MCPService: Service { reason: signal.reason, message: signal.message ?? CLIKillSignal.messageForReason(signal.reason) )) - } - - // PPID watchdog - detect orphaned CLI when parent dies - group.addTask { + }, + watchdog: { try await self.runPPIDWatchdog(initialPPID: initialPPID) return .ppidWatchdogCancelled - } - - // Main transport task - group.addTask { + }, + transport: { try await self.runTransport() return .transportCompleted } - - // Wait for first to complete (either kill signal, orphan detection, or transport exit), - // then stop the watcher tasks so a clean transport exit can return promptly. - do { - guard let outcome = try await group.next() else { - group.cancelAll() - throw CancellationError() - } - group.cancelAll() - if let runtimeError = MCPServiceProxyTaskGroupPolicy.terminalRuntimeError( - firstCompletedOutcome: outcome, - serviceTaskIsCancelled: Task.isCancelled - ) { - throw runtimeError - } - } catch { - group.cancelAll() - throw error - } + ) + if let runtimeError = MCPServiceProxyTaskGroupPolicy.terminalRuntimeError( + firstCompletedOutcome: outcome, + serviceTaskIsCancelled: Task.isCancelled + ) { + throw runtimeError } try Task.checkCancellation() await persistProxyTerminalRecord( diff --git a/Tests/RepoPromptTests/MCP/MCPServiceProxyRaceTests.swift b/Tests/RepoPromptTests/MCP/MCPServiceProxyRaceTests.swift new file mode 100644 index 000000000..11ee15c25 --- /dev/null +++ b/Tests/RepoPromptTests/MCP/MCPServiceProxyRaceTests.swift @@ -0,0 +1,256 @@ +@testable import RepoPromptMCP +import XCTest + +final class MCPServiceProxyRaceTests: XCTestCase { + func testEachChildCanWinWithValueOrErrorAndDrainsLosers() async throws { + let outcomes: [MCPServiceProxyTaskOutcome] = [ + .killSignalWaitCancelled, + .ppidWatchdogCancelled, + .transportCompleted + ] + + for winnerID in outcomes.indices { + for winnerThrows in [false, true] { + let probe = MCPProxyRaceProbe(expectedStartCount: outcomes.count) + let winnerResult: Result = winnerThrows + ? .failure(MCPProxyRaceTestError.winner(winnerID)) + : .success(outcomes[winnerID]) + let operations = outcomes.indices.map { childID in + makeOperation( + childID: childID, + winnerID: winnerID, + winnerResult: winnerResult, + probe: probe + ) + } + + do { + let outcome = try await MCPService.awaitFirstProxyOutcome( + killSignal: operations[0], + watchdog: operations[1], + transport: operations[2] + ) + XCTAssertFalse(winnerThrows, "Expected child \(winnerID) to throw") + XCTAssertEqual(proxyOutcomeID(outcome), winnerID) + } catch let error as MCPProxyRaceTestError { + XCTAssertTrue(winnerThrows, "Unexpected child error: \(error)") + XCTAssertEqual(error, .winner(winnerID)) + } + + let snapshot = await probe.snapshot() + XCTAssertEqual(snapshot.completed, Set(outcomes.indices)) + XCTAssertEqual(snapshot.cancelled, Set(outcomes.indices.filter { $0 != winnerID })) + } + } + } + + func testRaceDoesNotReturnUntilCancelledLosersComplete() async throws { + let probe = MCPProxyRaceProbe(expectedStartCount: 3) + let drainGate = MCPProxyRaceGate() + let returned = MCPProxyRaceFlag() + let winnerResult = Result.success(.transportCompleted) + let operations = (0 ..< 3).map { childID in + makeOperation( + childID: childID, + winnerID: 2, + winnerResult: winnerResult, + probe: probe, + drainGate: drainGate + ) + } + let race = Task { + let outcome = try await MCPService.awaitFirstProxyOutcome( + killSignal: operations[0], + watchdog: operations[1], + transport: operations[2] + ) + await returned.set() + return outcome + } + + await probe.waitUntilCancelled(count: 2) + let didReturnBeforeLosersCompleted = await returned.value + XCTAssertFalse(didReturnBeforeLosersCompleted) + + await drainGate.release() + let outcome = try await race.value + XCTAssertEqual(proxyOutcomeID(outcome), 2) + let snapshot = await probe.snapshot() + XCTAssertEqual(snapshot.completed, Set(0 ..< 3)) + } + + func testOuterCancellationCancelsAndDrainsAllChildren() async throws { + let probe = MCPProxyRaceProbe(expectedStartCount: 3) + let drainGate = MCPProxyRaceGate() + let returned = MCPProxyRaceFlag() + let operations = (0 ..< 3).map { childID in + makeOperation( + childID: childID, + winnerID: nil, + winnerResult: .success(.transportCompleted), + probe: probe, + drainGate: drainGate + ) + } + let race = Task { + do { + let outcome = try await MCPService.awaitFirstProxyOutcome( + killSignal: operations[0], + watchdog: operations[1], + transport: operations[2] + ) + await returned.set() + return outcome + } catch { + await returned.set() + throw error + } + } + + await probe.waitUntilAllStarted() + race.cancel() + await probe.waitUntilCancelled(count: 3) + let didReturnBeforeChildrenCompleted = await returned.value + XCTAssertFalse(didReturnBeforeChildrenCompleted) + + await drainGate.release() + do { + _ = try await race.value + XCTFail("Expected outer cancellation") + } catch is CancellationError { + // Expected. + } + + let snapshot = await probe.snapshot() + XCTAssertEqual(snapshot.cancelled, Set(0 ..< 3)) + XCTAssertEqual(snapshot.completed, Set(0 ..< 3)) + } + + private func makeOperation( + childID: Int, + winnerID: Int?, + winnerResult: Result, + probe: MCPProxyRaceProbe, + drainGate: MCPProxyRaceGate? = nil + ) -> @Sendable () async throws -> MCPServiceProxyTaskOutcome { + { + await probe.markStarted(childID) + await probe.waitUntilAllStarted() + if childID == winnerID { + await probe.markCompleted(childID) + return try winnerResult.get() + } + + do { + try await Task.sleep(nanoseconds: UInt64.max) + throw MCPProxyRaceTestError.loserReturnedWithoutCancellation(childID) + } catch is CancellationError { + await probe.markCancelled(childID) + if let drainGate { + await drainGate.wait() + } + await probe.markCompleted(childID) + throw CancellationError() + } + } + } + + private func proxyOutcomeID(_ outcome: MCPServiceProxyTaskOutcome) -> Int { + switch outcome { + case .killSignalWaitCancelled: + 0 + case .ppidWatchdogCancelled: + 1 + case .transportCompleted: + 2 + } + } +} + +private enum MCPProxyRaceTestError: Error, Equatable { + case winner(Int) + case loserReturnedWithoutCancellation(Int) +} + +private actor MCPProxyRaceProbe { + struct Snapshot { + let cancelled: Set + let completed: Set + } + + private let expectedStartCount: Int + private var started: Set = [] + private var cancelled: Set = [] + private var completed: Set = [] + private var startWaiters: [CheckedContinuation] = [] + private var cancellationWaiters: [(count: Int, continuation: CheckedContinuation)] = [] + + init(expectedStartCount: Int) { + self.expectedStartCount = expectedStartCount + } + + func markStarted(_ childID: Int) { + started.insert(childID) + guard started.count == expectedStartCount else { return } + let waiters = startWaiters + startWaiters.removeAll() + waiters.forEach { $0.resume() } + } + + func waitUntilAllStarted() async { + guard started.count < expectedStartCount else { return } + await withCheckedContinuation { continuation in + startWaiters.append(continuation) + } + } + + func markCancelled(_ childID: Int) { + cancelled.insert(childID) + let ready = cancellationWaiters.filter { cancelled.count >= $0.count } + cancellationWaiters.removeAll { cancelled.count >= $0.count } + ready.forEach { $0.continuation.resume() } + } + + func waitUntilCancelled(count: Int) async { + guard cancelled.count < count else { return } + await withCheckedContinuation { continuation in + cancellationWaiters.append((count, continuation)) + } + } + + func markCompleted(_ childID: Int) { + completed.insert(childID) + } + + func snapshot() -> Snapshot { + Snapshot(cancelled: cancelled, completed: completed) + } +} + +private actor MCPProxyRaceGate { + private var isReleased = false + private var waiters: [CheckedContinuation] = [] + + func wait() async { + guard !isReleased else { return } + await withCheckedContinuation { continuation in + waiters.append(continuation) + } + } + + func release() { + guard !isReleased else { return } + isReleased = true + let waiters = waiters + self.waiters.removeAll() + waiters.forEach { $0.resume() } + } +} + +private actor MCPProxyRaceFlag { + private(set) var value = false + + func set() { + value = true + } +}