diff --git a/Sources/RepoPrompt/Infrastructure/AI/Agents/AgentToolTracker.swift b/Sources/RepoPrompt/Infrastructure/AI/Agents/AgentToolTracker.swift index 6ebd63ed7..0ceca5b53 100644 --- a/Sources/RepoPrompt/Infrastructure/AI/Agents/AgentToolTracker.swift +++ b/Sources/RepoPrompt/Infrastructure/AI/Agents/AgentToolTracker.swift @@ -5,6 +5,7 @@ import MCP /// Tool observers are registered by runID (not connectionID) to survive connection handovers. actor AgentToolTracker { private var trackedRunID: UUID? + private var observerToken: UUID? private var hasUnregistered = false /// Registers an enhanced tool event observer that receives args on call and result on completion. @@ -22,8 +23,9 @@ actor AgentToolTracker { onCalled: onCalled, onCompleted: onCompleted ) - await manager.registerToolEventObserver(for: runID, observer: observer) + let token = await manager.registerToolEventObserver(for: runID, observer: observer) trackedRunID = runID + observerToken = token hasUnregistered = false } @@ -86,9 +88,16 @@ actor AgentToolTracker { /// Unregister exactly once to avoid double-cleanup races private func unregisterObserverOnce(for runID: UUID, manager: ServerNetworkManager) async { guard !hasUnregistered, trackedRunID == runID else { return } + let token = observerToken hasUnregistered = true trackedRunID = nil - await manager.unregisterToolObservers(for: runID) + observerToken = nil + + if let token { + await manager.unregisterToolEventObserver(for: runID, token: token) + } else { + await manager.unregisterToolObservers(for: runID) + } } /// Detaches the observer if one was registered. diff --git a/Sources/RepoPrompt/Infrastructure/MCP/MCPConnectionManager.swift b/Sources/RepoPrompt/Infrastructure/MCP/MCPConnectionManager.swift index 93e4ded23..9b7356055 100644 --- a/Sources/RepoPrompt/Infrastructure/MCP/MCPConnectionManager.swift +++ b/Sources/RepoPrompt/Infrastructure/MCP/MCPConnectionManager.swift @@ -3350,6 +3350,36 @@ actor ServerNetworkManager { return token } + /// Unregister one tool event observer for a specific run. + /// + /// Owner-scoped teardown should use this token-specific path so another + /// observer registered for the same run remains active. If the observer was + /// already captured by a run-wide unregister, wait for that cleanup barrier + /// instead of returning before its in-flight delivery drains. + func unregisterToolEventObserver(for runID: UUID, token: UUID) async { + let removedObserver: ToolEventObserver? + if var observers = toolEventObservers[runID] { + removedObserver = observers.removeValue(forKey: token) + if observers.isEmpty { + toolEventObservers.removeValue(forKey: runID) + } else { + toolEventObservers[runID] = observers + } + } else { + removedObserver = nil + } + + if let removedObserver { + await removedObserver.deliveryBarrier.waitUntilIdle() + connectionLog("Unregistered tool event observer for runID: \(runID) token: \(token)") + return + } + + if let unregistration = toolObserverUnregistrationsByRunID[runID] { + await unregistration.task.value + } + } + /// Unregister all tool event observers for a specific run func unregisterToolEventObservers(for runID: UUID) async { await unregisterToolObservers(for: runID) diff --git a/Tests/RepoPromptTests/App/WindowCloseCoordinatorDecisionTests.swift b/Tests/RepoPromptTests/App/WindowCloseCoordinatorDecisionTests.swift index a77743108..e813a66ad 100644 --- a/Tests/RepoPromptTests/App/WindowCloseCoordinatorDecisionTests.swift +++ b/Tests/RepoPromptTests/App/WindowCloseCoordinatorDecisionTests.swift @@ -499,36 +499,20 @@ final class WindowCloseCoordinatorPolicyTests: XCTestCase { } } -private actor APISettingsInitialLoadGate { - private var hasEntered = false - private var hasReleased = false - private var entryWaiters: [CheckedContinuation] = [] - private var releaseWaiters: [CheckedContinuation] = [] +/// Initial-load fence. Shared `TestReleaseFence` with legacy method names. +private final class APISettingsInitialLoadGate: @unchecked Sendable { + private let fence = TestReleaseFence(name: "API settings initial load gate") func arriveAndWait() async { - hasEntered = true - let waiters = entryWaiters - entryWaiters.removeAll() - waiters.forEach { $0.resume() } - - guard !hasReleased else { return } - await withCheckedContinuation { continuation in - releaseWaiters.append(continuation) - } + await fence.enterAndWait() } - func waitUntilEntered() async { - guard !hasEntered else { return } - await withCheckedContinuation { continuation in - entryWaiters.append(continuation) - } + func waitUntilEntered(timeout: TimeInterval = TestFenceDefaults.enterWait) async { + _ = await fence.waitUntilEntered(timeout: timeout) } func release() { - hasReleased = true - let waiters = releaseWaiters - releaseWaiters.removeAll() - waiters.forEach { $0.resume() } + fence.release() } } @@ -540,81 +524,79 @@ private actor APISettingsProviderValidationProbe { } } -private actor GitContextRefreshGate { +/// Git refresh probe: shared release fence + cancellation observation counters. +private final class GitContextRefreshGate: @unchecked Sendable { struct Observation: Equatable { let entryCount: Int let cancellationCount: Int let activeCount: Int } - private var hasEntered = false - private var hasReleased = false - private var entryWaiters: [CheckedContinuation] = [] - private var releaseWaiters: [CheckedContinuation] = [] - private var cancellationWaiters: [CheckedContinuation] = [] + private let fence = TestReleaseFence(name: "git context refresh gate") + private let lock = NSLock() private var entryCount = 0 private var cancellationCount = 0 private var activeCount = 0 func refresh(rootPaths _: [String]) async -> [GitStatusActor.RepoDetection] { - hasEntered = true + lock.lock() entryCount += 1 activeCount += 1 - let waiters = entryWaiters - entryWaiters.removeAll() - waiters.forEach { $0.resume() } + lock.unlock() await withTaskCancellationHandler { - if !hasReleased { - await withCheckedContinuation { continuation in - releaseWaiters.append(continuation) - } - } + await fence.enterAndWaitIgnoringCancellationUntilRelease() } onCancel: { - Task { await self.recordCancellation() } + // Record cancellation, but keep the simulated refresh body parked until release. + self.recordCancellation() } + + lock.lock() activeCount -= 1 + lock.unlock() return [] } - func waitUntilEntered() async { - guard !hasEntered else { return } - await withCheckedContinuation { continuation in - entryWaiters.append(continuation) - } + func waitUntilEntered(timeout: TimeInterval = TestFenceDefaults.enterWait) async { + _ = await fence.waitUntilEntered(timeout: timeout) } func release() { - hasReleased = true - let waiters = releaseWaiters - releaseWaiters.removeAll() - waiters.forEach { $0.resume() } + fence.release() } - func waitUntilCancellationObserved() async { - guard cancellationCount == 0 else { return } - await withCheckedContinuation { continuation in - cancellationWaiters.append(continuation) + func waitUntilCancellationObserved(timeout: TimeInterval = TestFenceDefaults.enterWait) async { + if observedCancellation { return } + do { + try await AsyncTestWait.waitUntil( + "git context refresh cancellation observed", + timeout: timeout + ) { + self.observedCancellation + } + } catch { + XCTFail(error.localizedDescription) } } var observedCancellation: Bool { - cancellationCount > 0 + lock.withLock { cancellationCount > 0 } } var observation: Observation { - Observation( - entryCount: entryCount, - cancellationCount: cancellationCount, - activeCount: activeCount - ) + lock.withLock { + Observation( + entryCount: entryCount, + cancellationCount: cancellationCount, + activeCount: activeCount + ) + } } private func recordCancellation() { + lock.lock() cancellationCount += 1 - let waiters = cancellationWaiters - cancellationWaiters.removeAll() - waiters.forEach { $0.resume() } + lock.unlock() } } @@ -691,6 +673,7 @@ private actor WindowClosePollingClientSpy: CodexModelListingClient { } } } onCancel: { + // Sticky cancelledRequestIDs covers cancel-before-register; actor hop resumes parked request. Task { await self.cancelRequest(requestID) } } } diff --git a/Tests/RepoPromptTests/CodeMap/CodeMapArtifactRuntimeTests.swift b/Tests/RepoPromptTests/CodeMap/CodeMapArtifactRuntimeTests.swift index e6bdff6ea..a9b2f4910 100644 --- a/Tests/RepoPromptTests/CodeMap/CodeMapArtifactRuntimeTests.swift +++ b/Tests/RepoPromptTests/CodeMap/CodeMapArtifactRuntimeTests.swift @@ -736,11 +736,11 @@ private final class RuntimeProviderOverlapGate: @unchecked Sendable { let factoryEntryCount: Int } - private let condition = NSCondition() private let expectedCallerCount: Int + private let factoryFence = TestBlockingFence(name: "runtime provider factory fence") + private let condition = NSCondition() private var callerCount = 0 private var factoryEntryCount = 0 - private var factoryReleased = false init(expectedCallerCount: Int) { self.expectedCallerCount = expectedCallerCount @@ -759,17 +759,15 @@ private final class RuntimeProviderOverlapGate: @unchecked Sendable { condition.unlock() } - func factoryEnteredAndWaitForRelease() { + func factoryEnteredAndWaitForRelease(timeout: TimeInterval = TestFenceDefaults.releaseWait) { condition.lock() factoryEntryCount += 1 condition.broadcast() - while !factoryReleased { - condition.wait() - } condition.unlock() + factoryFence.enterAndWait(timeout: timeout) } - func waitForDeterministicOverlap(timeout: TimeInterval = 10) -> Bool { + func waitForDeterministicOverlap(timeout: TimeInterval = TestFenceDefaults.enterWait) -> Bool { condition.lock() defer { condition.unlock() } let deadline = Date().addingTimeInterval(timeout) @@ -780,10 +778,7 @@ private final class RuntimeProviderOverlapGate: @unchecked Sendable { } func releaseFactory() { - condition.lock() - factoryReleased = true - condition.broadcast() - condition.unlock() + factoryFence.release() } } diff --git a/Tests/RepoPromptTests/CodeMap/CodeMapRootManifestStoreTests.swift b/Tests/RepoPromptTests/CodeMap/CodeMapRootManifestStoreTests.swift index b3db33961..e2202932b 100644 --- a/Tests/RepoPromptTests/CodeMap/CodeMapRootManifestStoreTests.swift +++ b/Tests/RepoPromptTests/CodeMap/CodeMapRootManifestStoreTests.swift @@ -2707,8 +2707,8 @@ private final class ManifestAccessClock: @unchecked Sendable { private final class ManifestLockedMergeGate: @unchecked Sendable { private let condition = NSCondition() + private let firstWriterFence = TestBlockingFence(name: "manifest locked merge first-writer fence") private var firstWriterAcquired = false - private var firstWriterReleased = false private var secondWriterAttempted = false private var secondWriterAcquired = false @@ -2716,14 +2716,12 @@ private final class ManifestLockedMergeGate: @unchecked Sendable { condition.withLock { secondWriterAcquired } } - func firstWriterAcquiredLockAndWait() { + func firstWriterAcquiredLockAndWait(timeout: TimeInterval = TestFenceDefaults.releaseWait) { condition.lock() firstWriterAcquired = true condition.broadcast() - while !firstWriterReleased { - condition.wait() - } condition.unlock() + firstWriterFence.enterAndWait(timeout: timeout) } func secondWriterAttemptedLock() { @@ -2740,19 +2738,16 @@ private final class ManifestLockedMergeGate: @unchecked Sendable { } } - func waitUntilFirstWriterAcquiredLock(timeout: TimeInterval = 10) -> Bool { + func waitUntilFirstWriterAcquiredLock(timeout: TimeInterval = TestFenceDefaults.enterWait) -> Bool { wait(timeout: timeout) { firstWriterAcquired } } - func waitUntilSecondWriterAttemptedLock(timeout: TimeInterval = 10) -> Bool { + func waitUntilSecondWriterAttemptedLock(timeout: TimeInterval = TestFenceDefaults.enterWait) -> Bool { wait(timeout: timeout) { secondWriterAttempted } } func releaseFirstWriter() { - condition.withLock { - firstWriterReleased = true - condition.broadcast() - } + firstWriterFence.release() } private func wait(timeout: TimeInterval, condition predicate: () -> Bool) -> Bool { @@ -2766,30 +2761,20 @@ private final class ManifestLockedMergeGate: @unchecked Sendable { } } -private actor ManifestAccessRefreshGate { - private var blocked = false - private var released = false - private var blockedWaiters: [CheckedContinuation] = [] - private var releaseContinuation: CheckedContinuation? +/// Manifest access refresh block/release fence (shared `TestReleaseFence`). +private final class ManifestAccessRefreshGate: @unchecked Sendable { + private let fence = TestReleaseFence(name: "manifest access refresh gate") func block() async { - blocked = true - let waiters = blockedWaiters - blockedWaiters.removeAll() - waiters.forEach { $0.resume() } - if released { return } - await withCheckedContinuation { releaseContinuation = $0 } + await fence.enterAndWait() } - func waitUntilBlocked() async { - if blocked { return } - await withCheckedContinuation { blockedWaiters.append($0) } + func waitUntilBlocked(timeout: TimeInterval = TestFenceDefaults.enterWait) async { + _ = await fence.waitUntilEntered(timeout: timeout) } func release() { - released = true - releaseContinuation?.resume() - releaseContinuation = nil + fence.release() } } diff --git a/Tests/RepoPromptTests/CodeMap/GitBlobCodeMapLocatorStoreTests.swift b/Tests/RepoPromptTests/CodeMap/GitBlobCodeMapLocatorStoreTests.swift index d0c806654..3287fba7d 100644 --- a/Tests/RepoPromptTests/CodeMap/GitBlobCodeMapLocatorStoreTests.swift +++ b/Tests/RepoPromptTests/CodeMap/GitBlobCodeMapLocatorStoreTests.swift @@ -467,7 +467,8 @@ final class GitBlobCodeMapLocatorStoreTests: XCTestCase { ) let reader = try GitBlobCodeMapLocatorStore(rootURL: root) let task = Task { try await writer.write(association: association) } - await gate.waitUntilEntered() + let gateEntered = await gate.waitUntilEntered() + XCTAssertTrue(gateEntered) let readDuringPublish = try await reader.read(identity: identity) await gate.release() let publishResult = try await task.value @@ -733,7 +734,8 @@ final class GitBlobCodeMapLocatorStoreTests: XCTestCase { ) ) let truncateTask = Task { try await truncateReader.read(identity: identity) } - await truncateGate.waitUntilEntered() + let truncateEntered = await truncateGate.waitUntilEntered() + XCTAssertTrue(truncateEntered) let handle = try FileHandle(forWritingTo: recordURL) try handle.truncate(atOffset: 0) try handle.close() @@ -753,7 +755,8 @@ final class GitBlobCodeMapLocatorStoreTests: XCTestCase { ) ) let replaceTask = Task { try await replaceReader.read(identity: identity) } - await replaceGate.waitUntilEntered() + let replaceEntered = await replaceGate.waitUntilEntered() + XCTAssertTrue(replaceEntered) let replacement = recordURL.deletingLastPathComponent().appendingPathComponent("replacement") try original.write(to: replacement) XCTAssertEqual(chmod(replacement.path, 0o600), 0) @@ -1063,28 +1066,4 @@ private actor GitBlobLocatorAsyncBarrier { } } -private actor GitBlobLocatorAsyncGate { - private var entered = false - private var released = false - private var entryWaiters: [CheckedContinuation] = [] - private var releaseWaiters: [CheckedContinuation] = [] - - func enterAndWait() async { - entered = true - entryWaiters.forEach { $0.resume() } - entryWaiters.removeAll() - guard !released else { return } - await withCheckedContinuation { releaseWaiters.append($0) } - } - - func waitUntilEntered() async { - guard !entered else { return } - await withCheckedContinuation { entryWaiters.append($0) } - } - - func release() { - released = true - releaseWaiters.forEach { $0.resume() } - releaseWaiters.removeAll() - } -} +private typealias GitBlobLocatorAsyncGate = TestReleaseFence diff --git a/Tests/RepoPromptTests/ContextBuilder/ContextBuilderFollowUpFinalizationMonitorTests.swift b/Tests/RepoPromptTests/ContextBuilder/ContextBuilderFollowUpFinalizationMonitorTests.swift index e70ed6f6d..5e432f0ed 100644 --- a/Tests/RepoPromptTests/ContextBuilder/ContextBuilderFollowUpFinalizationMonitorTests.swift +++ b/Tests/RepoPromptTests/ContextBuilder/ContextBuilderFollowUpFinalizationMonitorTests.swift @@ -55,6 +55,7 @@ final class ContextBuilderFollowUpFinalizationMonitorTests: XCTestCase { func testStalledFakeQueryTimesOutWithAttributedSubphaseAndCancelsStream() async { let clock = ContextBuilderFinalizationTestClock() let cancellationRecorder = ContextBuilderFinalizationCancellationRecorder() + let finalizationGate = ContextBuilderCancellableFinalizationGate() let (events, continuation) = AsyncStream.makeStream() defer { continuation.finish() } @@ -72,7 +73,7 @@ final class ContextBuilderFollowUpFinalizationMonitorTests: XCTestCase { await Task.yield() }, waitForFinalization: { - try await Task.sleep(for: .seconds(60)) + try await finalizationGate.wait() }, cancelStreaming: { await cancellationRecorder.recordCancellation() @@ -89,6 +90,8 @@ final class ContextBuilderFollowUpFinalizationMonitorTests: XCTestCase { let cancellationCount = await cancellationRecorder.count() XCTAssertEqual(cancellationCount, 1) + let finalizationWasCancelled = await finalizationGate.wasCancelled() + XCTAssertTrue(finalizationWasCancelled) } func testTimeoutOutcomeWinsWhenCancellationAlsoCompletesFinalization() async { @@ -262,61 +265,71 @@ private actor ContextBuilderFinalizationProgressRecorder { } } -private actor ContextBuilderCancellableFinalizationGate { +/// Finalization wait with sticky cancel (lock-backed; no actor cancel hop). +private final class ContextBuilderCancellableFinalizationGate: @unchecked Sendable { + private let lock = NSLock() private var completed = false private var cancelled = false - private var continuation: CheckedContinuation? + private var storedWasCancelled = false + private var continuations: [UUID: CheckedContinuation] = [:] + private var cancelledWaiters = Set() func wait() async throws { + let waiterID = UUID() await withTaskCancellationHandler { await withCheckedContinuation { continuation in - Task { await self.register(continuation) } + lock.lock() + if completed || cancelled || cancelledWaiters.remove(waiterID) != nil { + lock.unlock() + continuation.resume() + return + } + continuations[waiterID] = continuation + lock.unlock() } } onCancel: { - Task { await self.cancel() } + cancel(waiterID: waiterID) } try Task.checkCancellation() } func complete() { + lock.lock() completed = true - continuation?.resume() - continuation = nil + let pending = Array(continuations.values) + continuations.removeAll() + lock.unlock() + for waiter in pending { + waiter.resume() + } } - private func register(_ continuation: CheckedContinuation) { - if completed || cancelled { - continuation.resume() - } else { - self.continuation = continuation - } + func wasCancelled() -> Bool { + lock.withLock { storedWasCancelled } } - private func cancel() { + private func cancel(waiterID: UUID) { + lock.lock() cancelled = true - continuation?.resume() - continuation = nil + storedWasCancelled = true + let pending = continuations.removeValue(forKey: waiterID) + if pending == nil { + cancelledWaiters.insert(waiterID) + } + lock.unlock() + pending?.resume() } } -private actor ContextBuilderFinalizationTestGate { - private var released = false - private var waiters: [CheckedContinuation] = [] +/// Finalization release fence (shared hang-hardened primitive). +private final class ContextBuilderFinalizationTestGate: @unchecked Sendable { + private let fence = TestReleaseFence(name: "context builder finalization test gate") func arriveAndWait() async { - guard !released else { return } - await withCheckedContinuation { continuation in - waiters.append(continuation) - } + await fence.enterAndWait() } func release() { - guard !released else { return } - released = true - let currentWaiters = waiters - waiters.removeAll() - for waiter in currentWaiters { - waiter.resume() - } + fence.release() } } diff --git a/Tests/RepoPromptTests/ContextBuilder/ContextBuilderMCPProgressTimelineTests.swift b/Tests/RepoPromptTests/ContextBuilder/ContextBuilderMCPProgressTimelineTests.swift index 6cb68c393..e15f1edc2 100644 --- a/Tests/RepoPromptTests/ContextBuilder/ContextBuilderMCPProgressTimelineTests.swift +++ b/Tests/RepoPromptTests/ContextBuilder/ContextBuilderMCPProgressTimelineTests.swift @@ -136,16 +136,19 @@ final class ContextBuilderMCPProgressTimelineTests: XCTestCase { } } - func testSoftStageBoundEmitsOnceWithoutFailingTimeline() async { + func testSoftStageBoundEmitsOnceWithoutFailingTimeline() async throws { let clock = ContextBuilderProgressTestClock() let recorder = ContextBuilderProgressEventRecorder() + let softBoundSleep = ContextBuilderSoftBoundSleepGate() let timeline = ContextBuilderMCPProgressTimeline( clock: { clock.now() }, - sleep: { _ in try await Task.sleep(for: .seconds(60)) }, + sleep: { seconds in try await softBoundSleep.sleep(seconds: seconds) }, sink: { event in await recorder.record(event) } ) await timeline.transition(to: .modelResolution) + let scheduledSoftBound = try await softBoundSleep.waitUntilSleeping() + XCTAssertEqual(scheduledSoftBound, 2, accuracy: 0.000_1) clock.advance(by: 2.5) await timeline.checkSoftBound() await timeline.checkSoftBound() @@ -160,6 +163,7 @@ final class ContextBuilderMCPProgressTimelineTests: XCTestCase { XCTAssertEqual(events.count { $0.kind == .softBoundExceeded }, 1) XCTAssertTrue(events[1].message.contains("soft bound 2.000s"), events[1].message) XCTAssertEqual(events.last?.phaseElapsed ?? 0, 2.5, accuracy: 0.000_1) + try await softBoundSleep.waitUntilCancelled() } @MainActor @@ -714,6 +718,158 @@ private final class ContextBuilderProgressTimelineReference: @unchecked Sendable } } +/// Soft-bound sleep gate: lock-backed sticky cancel (no `Task { await }` hop). +private final class ContextBuilderSoftBoundSleepGate: @unchecked Sendable { + private let lock = NSLock() + private var sleepingSeconds: TimeInterval? + private var cancelled = false + private var sleepWaitTerminalError: Error? + private var sleepContinuations: [UUID: CheckedContinuation] = [:] + private var cancelledSleepWaiters = Set() + private var sleepingWaiters: [CheckedContinuation] = [] + + func sleep(seconds: TimeInterval) async throws { + try Task.checkCancellation() + let waiterID = UUID() + try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + var result: Result? + lock.lock() + if cancelled || Task.isCancelled || cancelledSleepWaiters.remove(waiterID) != nil { + result = .failure(CancellationError()) + } else { + if sleepingSeconds == nil { + sleepingSeconds = seconds + } + sleepContinuations[waiterID] = continuation + let waiters = sleepingWaiters + sleepingWaiters.removeAll() + lock.unlock() + waiters.forEach { $0.resume(returning: seconds) } + return + } + lock.unlock() + if case let .failure(error) = result { + continuation.resume(throwing: error) + } + } + } onCancel: { + cancelSleep(waiterID: waiterID) + } + } + + func waitUntilSleeping(timeout: TimeInterval = TestFenceDefaults.enterWait) async throws -> TimeInterval { + if let seconds = peekSleepingSeconds { return seconds } + if isCancelled { throw CancellationError() } + if let terminal = terminalError { throw terminal } + + let timeoutError = AsyncTestConditionTimeout( + description: "context builder soft-bound sleep gate", + timeout: timeout + ) + return try await withThrowingTaskGroup(of: TimeInterval.self) { group in + group.addTask { + try await self.waitForSleepSignal() + } + group.addTask { + try await Task.sleep(nanoseconds: UInt64((timeout * 1_000_000_000).rounded())) + self.closeSleepWait(with: timeoutError) + throw timeoutError + } + defer { group.cancelAll() } + do { + if let value = try await group.next() { + return value + } + } catch { + if let seconds = self.peekSleepingSeconds { + return seconds + } + throw error + } + if let seconds = peekSleepingSeconds { + return seconds + } + throw timeoutError + } + } + + func waitUntilCancelled(timeout: TimeInterval = TestFenceDefaults.enterWait) async throws { + if isCancelled { return } + try await AsyncTestWait.waitUntil( + "context builder soft-bound cancellation", + timeout: timeout + ) { + self.isCancelled + } + } + + private var peekSleepingSeconds: TimeInterval? { + lock.withLock { sleepingSeconds } + } + + private var isCancelled: Bool { + lock.withLock { cancelled } + } + + private var terminalError: Error? { + lock.withLock { sleepWaitTerminalError } + } + + private func waitForSleepSignal() async throws -> TimeInterval { + try await withCheckedThrowingContinuation { continuation in + lock.lock() + if let sleepingSeconds { + lock.unlock() + continuation.resume(returning: sleepingSeconds) + return + } + if cancelled { + lock.unlock() + continuation.resume(throwing: CancellationError()) + return + } + if let sleepWaitTerminalError { + lock.unlock() + continuation.resume(throwing: sleepWaitTerminalError) + return + } + sleepingWaiters.append(continuation) + lock.unlock() + } + } + + private func closeSleepWait(with error: Error) { + lock.lock() + if sleepWaitTerminalError == nil { + sleepWaitTerminalError = error + } + let waiters = sleepingWaiters + sleepingWaiters.removeAll() + lock.unlock() + waiters.forEach { $0.resume(throwing: error) } + } + + private func cancelSleep(waiterID: UUID) { + lock.lock() + cancelled = true + let pending = Array(sleepContinuations.values) + let hadRegisteredWaiter = sleepContinuations.removeValue(forKey: waiterID) != nil + sleepContinuations.removeAll() + if !hadRegisteredWaiter { + cancelledSleepWaiters.insert(waiterID) + } + if sleepWaitTerminalError == nil { + sleepWaitTerminalError = CancellationError() + } + let waiters = sleepingWaiters + sleepingWaiters.removeAll() + lock.unlock() + pending.forEach { $0.resume(throwing: CancellationError()) } + waiters.forEach { $0.resume(throwing: CancellationError()) } + } +} + private actor ContextBuilderProgressEventRecorder { private var events: [ContextBuilderMCPProgressEvent] = [] diff --git a/Tests/RepoPromptTests/Helpers/AsyncTestCondition.swift b/Tests/RepoPromptTests/Helpers/AsyncTestCondition.swift index 259ff30fe..ce7bd3564 100644 --- a/Tests/RepoPromptTests/Helpers/AsyncTestCondition.swift +++ b/Tests/RepoPromptTests/Helpers/AsyncTestCondition.swift @@ -36,17 +36,24 @@ enum AsyncTestWait { maximumDelayNanoseconds: UInt64 = 25_000_000, condition: @escaping () async throws -> Bool ) async throws { - let deadline = Date().addingTimeInterval(timeout) - var delay = initialDelayNanoseconds + let timeoutDuration = Duration.seconds(max(0, timeout)) + let clock = ContinuousClock() + let deadline = clock.now.advanced(by: timeoutDuration) + let maximumDelay = max(1, min(maximumDelayNanoseconds, UInt64(Int64.max))) + var delay = max(1, min(initialDelayNanoseconds, maximumDelay)) + while true { if try await condition() { return } - let remaining = deadline.timeIntervalSinceNow - guard remaining > 0 else { + let now = clock.now + guard now < deadline else { throw AsyncTestConditionTimeout(description: description, timeout: timeout) } - let remainingNanoseconds = UInt64((remaining * 1_000_000_000).rounded(.up)) - try await Task.sleep(nanoseconds: min(delay, remainingNanoseconds)) - delay = min(delay * 2, maximumDelayNanoseconds) + let sleepDeadline = min( + deadline, + now.advanced(by: .nanoseconds(Int64(delay))) + ) + try await clock.sleep(until: sleepDeadline, tolerance: .zero) + delay = min(delay > maximumDelay / 2 ? maximumDelay : delay * 2, maximumDelay) } } } @@ -56,6 +63,10 @@ enum AsyncTestWait { /// This helper is intentionally small: tests mutate a protected snapshot and waiters /// are resumed exactly when a later mutation satisfies their predicate. The only /// sleep is the bounded timeout timer; it is not used for polling state. +/// +/// Hang guarantees: +/// - sticky cancel (timeout/cancel-before-register fails closed; late register cannot park forever) +/// - synchronous `onCancel` via lock-backed state (no `Task { await }` hop) final class AsyncTestCondition: @unchecked Sendable { private struct Waiter { let id: UUID @@ -63,9 +74,20 @@ final class AsyncTestCondition: @unchecked Sendable { let continuation: CheckedContinuation } + private enum WaiterState { + case pendingRegistration + case registered + case completed + } + private let lock = NSLock() private var value: Value private var waiters: [Waiter] = [] + /// Sticky cancel-before-register: timeout/cancel may race ahead of waiter registration. + private var cancelledWaiters: [UUID: Error] = [:] + /// Tracks terminal waiter IDs so cleanup after a successful wait cannot create + /// stale sticky-cancel entries for IDs that will never register again. + private var waiterStates: [UUID: WaiterState] = [:] init(_ value: Value) { self.value = value @@ -83,6 +105,7 @@ final class AsyncTestCondition: @unchecked Sendable { mutate(&value) waiters.removeAll { waiter in guard waiter.predicate(value) else { return false } + waiterStates[waiter.id] = .completed ready.append(waiter.continuation) return true } @@ -99,13 +122,20 @@ final class AsyncTestCondition: @unchecked Sendable { predicate: @escaping (Value) -> Bool ) async throws { let waiterID = UUID() + lock.lock() + waiterStates[waiterID] = .pendingRegistration + lock.unlock() + try await withThrowingTaskGroup(of: Void.self) { group in group.addTask { try await self.waitForSignal(id: waiterID, predicate: predicate) } group.addTask { - try await Task.sleep(nanoseconds: UInt64((timeout * 1_000_000_000).rounded())) - self.cancelWaiter(id: waiterID, error: AsyncTestConditionTimeout(description: description, timeout: timeout)) + try await Task.sleep(for: .seconds(max(0, timeout))) + self.cancelWaiter( + id: waiterID, + error: AsyncTestConditionTimeout(description: description, timeout: timeout) + ) throw AsyncTestConditionTimeout(description: description, timeout: timeout) } defer { @@ -117,27 +147,57 @@ final class AsyncTestCondition: @unchecked Sendable { } private func waitForSignal(id: UUID, predicate: @escaping (Value) -> Bool) async throws { - try await withCheckedThrowingContinuation { continuation in - var shouldResume = false - lock.lock() - if predicate(value) { - shouldResume = true - } else { - waiters.append(Waiter(id: id, predicate: predicate, continuation: continuation)) - } - lock.unlock() + try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { continuation in + var resumeResult: Result? + lock.lock() + if let error = cancelledWaiters.removeValue(forKey: id) { + waiterStates.removeValue(forKey: id) + resumeResult = .failure(error) + } else if Task.isCancelled { + waiterStates.removeValue(forKey: id) + resumeResult = .failure(CancellationError()) + } else if predicate(value) { + waiterStates[id] = .completed + resumeResult = .success(()) + } else { + waiters.append(Waiter(id: id, predicate: predicate, continuation: continuation)) + waiterStates[id] = .registered + } + lock.unlock() - if shouldResume { - continuation.resume() + switch resumeResult { + case .success: + continuation.resume() + case let .failure(error): + continuation.resume(throwing: error) + case nil: + break + } } + } onCancel: { + cancelWaiter(id: id, error: CancellationError()) } } private func cancelWaiter(id: UUID, error: Error) { var continuation: CheckedContinuation? lock.lock() - if let index = waiters.firstIndex(where: { $0.id == id }) { - continuation = waiters.remove(at: index).continuation + switch waiterStates[id] { + case .registered: + if let index = waiters.firstIndex(where: { $0.id == id }) { + continuation = waiters.remove(at: index).continuation + } + waiterStates[id] = .completed + case .pendingRegistration: + if cancelledWaiters[id] == nil { + cancelledWaiters[id] = error + } + case .completed: + waiterStates.removeValue(forKey: id) + cancelledWaiters.removeValue(forKey: id) + case nil: + cancelledWaiters.removeValue(forKey: id) } lock.unlock() continuation?.resume(throwing: error) diff --git a/Tests/RepoPromptTests/Helpers/CodemapBindingEngineTestSupport.swift b/Tests/RepoPromptTests/Helpers/CodemapBindingEngineTestSupport.swift index 246c90de2..0cb5d7e07 100644 --- a/Tests/RepoPromptTests/Helpers/CodemapBindingEngineTestSupport.swift +++ b/Tests/RepoPromptTests/Helpers/CodemapBindingEngineTestSupport.swift @@ -501,26 +501,64 @@ class CodemapBindingEngineTestCase: XCTestCase { final class EngineDemandResultTimeoutRace: @unchecked Sendable { private let lock = NSLock() private var continuation: CheckedContinuation? + private var observerTask: Task? + private var timeoutTask: Task? + private var observerCompleted = false private var resolved = false func wait( for task: Task, timeout: Duration ) async -> WorkspaceCodemapBindingDemandResult? { - await withCheckedContinuation { continuation in + let result = await withCheckedContinuation { continuation in lock.withLock { self.continuation = continuation } - Task { [self] in - await finish(task.value) + + let observer = Task { [weak self] in + let result = await task.value + self?.finish(result, observerCompleted: true) } - Task { [self] in + let timer = Task { [weak self] in try? await Task.sleep(for: timeout) - finish(nil) + guard !Task.isCancelled else { return } + task.cancel() + self?.finish(nil) + } + + lock.withLock { + observerTask = observer + timeoutTask = timer } } + + let tasks = lock.withLock { () -> (observer: Task?, timer: Task?) in + let tasks = (observerTask, timeoutTask) + observerTask = nil + timeoutTask = nil + return tasks + } + if let timer = tasks.timer { + timer.cancel() + await timer.value + } + if result == nil { + await assertObservedTaskDrainedAfterTimeout() + } + tasks.observer?.cancel() + return result + } + + private var hasObserverCompleted: Bool { + lock.withLock { observerCompleted } } - private func finish(_ result: WorkspaceCodemapBindingDemandResult?) { + private func finish( + _ result: WorkspaceCodemapBindingDemandResult?, + observerCompleted: Bool = false + ) { let continuation = lock.withLock { () -> CheckedContinuation? in + if observerCompleted { + self.observerCompleted = true + } guard !resolved else { return nil } resolved = true defer { self.continuation = nil } @@ -528,6 +566,19 @@ final class EngineDemandResultTimeoutRace: @unchecked Sendable { } continuation?.resume(returning: result) } + + private func assertObservedTaskDrainedAfterTimeout() async { + do { + try await AsyncTestWait.waitUntil( + "codemap binding demand task cancellation drain", + timeout: 1 + ) { + self.hasObserverCompleted + } + } catch { + XCTFail("Timed out waiting for cancelled codemap binding demand task to drain: \(error.localizedDescription)") + } + } } struct EngineFixture { @@ -670,76 +721,40 @@ final class EngineLockedFlag: @unchecked Sendable { } } -actor EngineBuildGate { - private var entered = false - private var released = false - private var continuation: CheckedContinuation? - private var enteredWaiters: [(id: UUID, continuation: CheckedContinuation)] = [] +/// Engine build-hook fence — named thin wrapper over `TestReleaseFence`. +final class EngineBuildGate: @unchecked Sendable { + private let fence = TestReleaseFence(name: "engine build gate") func enter() async { - entered = true - let waiters = enteredWaiters - enteredWaiters.removeAll() - for waiter in waiters { - waiter.continuation.resume() - } - if released { return } - await withCheckedContinuation { - if continuation != nil { - preconditionFailure("EngineBuildGate supports exactly one waiter") - } - continuation = $0 - } + await fence.enter() } - func waitUntilEntered(timeout: Duration = .seconds(10)) async -> Bool { - if entered { return true } - let waiterID = UUID() - let result = await withTaskGroup(of: Bool.self) { group in - group.addTask { [self] in - await waitForEnteredSignal(id: waiterID) - return true - } - group.addTask { - try? await Task.sleep(for: timeout) - return false - } - let result = await group.next() ?? false - if !result { - cancelEnteredWaiter(id: waiterID) - } - group.cancelAll() - return result - } - return result || entered + func enterAndWait() async { + await fence.enterAndWait() } - func release() { - released = true - continuation?.resume() - continuation = nil + func enterIgnoringCancellationUntilRelease() async { + await fence.enterAndWaitIgnoringCancellationUntilRelease() } - private func waitForEnteredSignal(id: UUID) async { - await withCheckedContinuation { continuation in - if entered { - continuation.resume() - } else { - enteredWaiters.append((id, continuation)) - } - } + @discardableResult + func waitUntilEntered( + timeout: Duration = TestFenceDefaults.enterWaitDuration, + failOnTimeout: Bool = true + ) async -> Bool { + await fence.waitUntilEntered(timeout: timeout, failOnTimeout: failOnTimeout) } - private func cancelEnteredWaiter(id: UUID) { - var cancelled: [CheckedContinuation] = [] - enteredWaiters.removeAll { waiter in - guard waiter.id == id else { return false } - cancelled.append(waiter.continuation) - return true - } - for continuation in cancelled { - continuation.resume() - } + @discardableResult + func waitUntilEnteredBlocking( + timeout: TimeInterval = TestFenceDefaults.enterWait, + failOnTimeout: Bool = true + ) -> Bool { + fence.waitUntilEnteredBlocking(timeout: timeout, failOnTimeout: failOnTimeout) + } + + func release() { + fence.release() } } @@ -779,85 +794,55 @@ actor EngineSecondCatalogResolutionMutation { } } +/// Sync writer-hook fence — named thin wrapper over `TestBlockingFence`. final class EngineBlockingGate: @unchecked Sendable { - private let condition = NSCondition() - private var entered = false - private var released = false + static let defaultEnterWaitTimeout = TestFenceDefaults.releaseWait - func enterAndWait() { - condition.lock() - entered = true - condition.broadcast() - while !released { - condition.wait() - } - condition.unlock() + private let fence = TestBlockingFence(name: "engine blocking gate") + + func enterAndWait(timeout: TimeInterval = defaultEnterWaitTimeout) { + fence.enterAndWait(timeout: timeout) } - func waitUntilEntered(timeout: TimeInterval = 10) -> Bool { - condition.lock() - defer { condition.unlock() } - let deadline = Date().addingTimeInterval(timeout) - while !entered { - guard condition.wait(until: deadline) else { return false } - } - return true + @discardableResult + func waitUntilEntered( + timeout: TimeInterval = TestFenceDefaults.enterWait, + failOnTimeout: Bool = true + ) -> Bool { + fence.waitUntilEntered(timeout: timeout, failOnTimeout: failOnTimeout) } func release() { - condition.lock() - released = true - condition.broadcast() - condition.unlock() + fence.release() } } +/// Async engine fence — named thin wrapper over `TestReleaseFence`. final class EngineAsyncGate: @unchecked Sendable { - private let condition = NSCondition() - private var entered = false - private var released = false - private var continuation: CheckedContinuation? + private let fence = TestReleaseFence(name: "engine async gate") func enterAndWait() async { - await withCheckedContinuation { continuation in - register(continuation) - } + await fence.enterAndWait() } - private func register(_ continuation: CheckedContinuation) { - condition.lock() - entered = true - condition.broadcast() - if released { - condition.unlock() - continuation.resume() - } else { - if self.continuation != nil { - preconditionFailure("EngineAsyncGate supports exactly one waiter") - } - self.continuation = continuation - condition.unlock() - } + @discardableResult + func waitUntilEnteredBlocking( + timeout: TimeInterval = TestFenceDefaults.enterWait, + failOnTimeout: Bool = true + ) -> Bool { + fence.waitUntilEnteredBlocking(timeout: timeout, failOnTimeout: failOnTimeout) } - func waitUntilEntered(timeout: TimeInterval = 10) -> Bool { - condition.lock() - defer { condition.unlock() } - let deadline = Date().addingTimeInterval(timeout) - while !entered { - guard condition.wait(until: deadline) else { return false } - } - return true + @discardableResult + func waitUntilEntered( + timeout: Duration = TestFenceDefaults.enterWaitDuration, + failOnTimeout: Bool = true + ) async -> Bool { + await fence.waitUntilEntered(timeout: timeout, failOnTimeout: failOnTimeout) } func release() { - condition.lock() - released = true - let continuation = continuation - self.continuation = nil - condition.broadcast() - condition.unlock() - continuation?.resume() + fence.release() } } @@ -876,164 +861,368 @@ enum EngineRegistrationInvalidationKind: CaseIterable { } actor EngineMultiEntryGate { - private var enteredCount = 0 - private var released = false - private var continuations: [CheckedContinuation] = [] - private var enteredWaiters: [(id: UUID, expectedCount: Int, continuation: CheckedContinuation)] = [] + private let state = EngineMultiEntryGateState() var count: Int { - enteredCount + state.count } func enter() async { - enteredCount += 1 - resumeSatisfiedEnteredWaiters() - if released { return } - await withCheckedContinuation { continuations.append($0) } + await state.enter() } func waitUntilEntered( _ expectedCount: Int, timeout: Duration = .seconds(10) ) async -> Bool { - if enteredCount >= expectedCount { return true } + do { + return try await state.waitUntilEntered( + expectedCount, + timeout: CodemapBindingEngineTestCase.timeInterval(timeout) + ) + } catch { + // Timeout sibling can win the task group even after the condition is met. + if state.count >= expectedCount { return true } + XCTFail(error.localizedDescription) + return false + } + } + + func releaseAll() { + state.releaseAll() + } +} + +private final class EngineMultiEntryGateState: @unchecked Sendable { + private struct EnteredWaiter { + let id: UUID + let expectedCount: Int + let continuation: CheckedContinuation + } + + private let lock = NSLock() + private var enteredCount = 0 + private var released = false + private var continuations: [UUID: CheckedContinuation] = [:] + private var cancelledWaiters = Set() + private var cancelledEnteredWaiters: [UUID: Error] = [:] + private var enteredWaiters: [EnteredWaiter] = [] + + var count: Int { + lock.withLock { enteredCount } + } + + func enter() async { let waiterID = UUID() - let result = await withTaskGroup(of: Bool.self) { group in - group.addTask { [self] in - await waitForEnteredSignal(id: waiterID, expectedCount: expectedCount) - return true - } - group.addTask { - try? await Task.sleep(for: timeout) - return false - } - let result = await group.next() ?? false - if !result { - cancelEnteredWaiter(id: waiterID) + let readyWaiters = lock.withLock { () -> [EnteredWaiter] in + enteredCount += 1 + return removeSatisfiedEnteredWaiters() + } + for waiter in readyWaiters { + waiter.continuation.resume() + } + await withTaskCancellationHandler { + await withCheckedContinuation { continuation in + var shouldResume = false + lock.lock() + if released || Task.isCancelled || cancelledWaiters.remove(waiterID) != nil { + shouldResume = true + } else { + continuations[waiterID] = continuation + } + lock.unlock() + if shouldResume { + continuation.resume() + } } - group.cancelAll() - return result + } onCancel: { + cancelEntryWaiter(id: waiterID) } - return result || enteredCount >= expectedCount } - func releaseAll() { - released = true - let pending = continuations - continuations.removeAll() - for continuation in pending { - continuation.resume() + func waitUntilEntered(_ expectedCount: Int, timeout: TimeInterval) async throws -> Bool { + if count >= expectedCount { return true } + let waiterID = UUID() + do { + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + try await self.waitForEnteredSignal(id: waiterID, expectedCount: expectedCount) + } + group.addTask { + try await Task.sleep(nanoseconds: UInt64((timeout * 1_000_000_000).rounded())) + self.cancelEnteredWaiter( + id: waiterID, + error: AsyncTestConditionTimeout(description: "engine multi-entry gate count \(expectedCount)", timeout: timeout) + ) + throw AsyncTestConditionTimeout(description: "engine multi-entry gate count \(expectedCount)", timeout: timeout) + } + defer { + group.cancelAll() + cancelEnteredWaiter(id: waiterID, error: CancellationError()) + } + _ = try await group.next() + } + } catch { + if count >= expectedCount { return true } + throw error } + return count >= expectedCount } - private func resumeSatisfiedEnteredWaiters() { - var ready: [CheckedContinuation] = [] - enteredWaiters.removeAll { waiter in - guard enteredCount >= waiter.expectedCount else { return false } - ready.append(waiter.continuation) - return true + func releaseAll() { + let pending = lock.withLock { () -> [CheckedContinuation] in + released = true + let pending = Array(continuations.values) + continuations.removeAll() + cancelledWaiters.removeAll() + cancelledEnteredWaiters.removeAll() + return pending } - for continuation in ready { + for continuation in pending { continuation.resume() } } - private func waitForEnteredSignal(id: UUID, expectedCount: Int) async { - await withCheckedContinuation { continuation in + private func waitForEnteredSignal(id: UUID, expectedCount: Int) async throws { + try await withCheckedThrowingContinuation { continuation in + var result: Result? + lock.lock() if enteredCount >= expectedCount { - continuation.resume() + result = .success(()) + } else if let error = cancelledEnteredWaiters.removeValue(forKey: id) { + result = .failure(error) + } else if Task.isCancelled { + result = .failure(CancellationError()) } else { - enteredWaiters.append((id, expectedCount, continuation)) + enteredWaiters.append(EnteredWaiter(id: id, expectedCount: expectedCount, continuation: continuation)) + } + lock.unlock() + + switch result { + case .success: + continuation.resume() + case let .failure(error): + continuation.resume(throwing: error) + case nil: + break } } } - private func cancelEnteredWaiter(id: UUID) { - var cancelled: [CheckedContinuation] = [] + private func removeSatisfiedEnteredWaiters() -> [EnteredWaiter] { + var ready: [EnteredWaiter] = [] enteredWaiters.removeAll { waiter in - guard waiter.id == id else { return false } - cancelled.append(waiter.continuation) + guard enteredCount >= waiter.expectedCount else { return false } + ready.append(waiter) return true } - for continuation in cancelled { - continuation.resume() + return ready + } + + private func cancelEntryWaiter(id: UUID) { + let continuation = lock.withLock { () -> CheckedContinuation? in + if let continuation = continuations.removeValue(forKey: id) { + return continuation + } + cancelledWaiters.insert(id) + return nil + } + continuation?.resume() + } + + private func cancelEnteredWaiter(id: UUID, error: Error) { + let continuation = lock.withLock { () -> CheckedContinuation? in + if let index = enteredWaiters.firstIndex(where: { $0.id == id }) { + return enteredWaiters.remove(at: index).continuation + } + // Sticky: timeout/cancel may race ahead of registration; late register must not park. + if cancelledEnteredWaiters[id] == nil { + cancelledEnteredWaiters[id] = error + } + return nil } + continuation?.resume(throwing: error) } } actor EngineFirstResolutionGate { - private(set) var resolutionCount = 0 - private var firstResolutionEntered = false + private let state = EngineFirstResolutionGateState() + + var resolutionCount: Int { + state.resolutionCount + } + + func enter() async { + await state.enter() + } + + func waitUntilFirstResolution(timeout: Duration = .seconds(10)) async -> Bool { + do { + return try await state.waitUntilFirstResolution( + timeout: CodemapBindingEngineTestCase.timeInterval(timeout) + ) + } catch { + if state.firstResolutionEntered { return true } + XCTFail(error.localizedDescription) + return false + } + } + + func releaseFirstResolution() { + state.releaseFirstResolution() + } +} + +private final class EngineFirstResolutionGateState: @unchecked Sendable { + private struct ResolutionWaiter { + let id: UUID + let continuation: CheckedContinuation + } + + private let lock = NSLock() + private var storedResolutionCount = 0 + private var storedFirstResolutionEntered = false private var firstResolutionReleased = false private var continuation: CheckedContinuation? - private var firstResolutionWaiters: [(id: UUID, continuation: CheckedContinuation)] = [] + private var cancelledWaiters = Set() + private var cancelledFirstResolutionWaiters: [UUID: Error] = [:] + private var firstResolutionWaiters: [ResolutionWaiter] = [] + + var resolutionCount: Int { + lock.withLock { storedResolutionCount } + } + + var firstResolutionEntered: Bool { + lock.withLock { storedFirstResolutionEntered } + } func enter() async { - resolutionCount += 1 - guard resolutionCount == 1 else { return } - firstResolutionEntered = true - let waiters = firstResolutionWaiters - firstResolutionWaiters.removeAll() - for waiter in waiters { + let entry = lock.withLock { () -> (readyWaiters: [ResolutionWaiter], shouldBlock: Bool) in + storedResolutionCount += 1 + guard storedResolutionCount == 1 else { return ([], false) } + storedFirstResolutionEntered = true + let waiters = firstResolutionWaiters + firstResolutionWaiters.removeAll() + return (waiters, true) + } + for waiter in entry.readyWaiters { waiter.continuation.resume() } - if firstResolutionReleased { return } - await withCheckedContinuation { - if continuation != nil { - preconditionFailure("EngineFirstResolutionGate supports exactly one waiter") + guard entry.shouldBlock else { return } + let waiterID = UUID() + await withTaskCancellationHandler { + await withCheckedContinuation { continuation in + var shouldResume = false + lock.lock() + if firstResolutionReleased || Task.isCancelled || cancelledWaiters.remove(waiterID) != nil { + shouldResume = true + } else if self.continuation != nil { + lock.unlock() + preconditionFailure("EngineFirstResolutionGate supports exactly one waiter") + } else { + self.continuation = continuation + } + lock.unlock() + + if shouldResume { + continuation.resume() + } } - continuation = $0 + } onCancel: { + cancelFirstResolutionBlocker(waiterID: waiterID) } } - func waitUntilFirstResolution(timeout: Duration = .seconds(10)) async -> Bool { + func waitUntilFirstResolution(timeout: TimeInterval) async throws -> Bool { if firstResolutionEntered { return true } let waiterID = UUID() - let result = await withTaskGroup(of: Bool.self) { group in - group.addTask { [self] in - await waitForFirstResolutionSignal(id: waiterID) - return true - } - group.addTask { - try? await Task.sleep(for: timeout) - return false - } - let result = await group.next() ?? false - if !result { - cancelFirstResolutionWaiter(id: waiterID) + do { + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + try await self.waitForFirstResolutionSignal(id: waiterID) + } + group.addTask { + try await Task.sleep(nanoseconds: UInt64((timeout * 1_000_000_000).rounded())) + self.cancelFirstResolutionWaiter( + id: waiterID, + error: AsyncTestConditionTimeout(description: "engine first resolution gate", timeout: timeout) + ) + throw AsyncTestConditionTimeout(description: "engine first resolution gate", timeout: timeout) + } + defer { + group.cancelAll() + cancelFirstResolutionWaiter(id: waiterID, error: CancellationError()) + } + _ = try await group.next() } - group.cancelAll() - return result + } catch { + if firstResolutionEntered { return true } + throw error } - return result || firstResolutionEntered + return firstResolutionEntered } func releaseFirstResolution() { - firstResolutionReleased = true - continuation?.resume() - continuation = nil - } + let pending = lock.withLock { () -> CheckedContinuation? in + firstResolutionReleased = true + let pending = self.continuation + self.continuation = nil + cancelledWaiters.removeAll() + cancelledFirstResolutionWaiters.removeAll() + return pending + } + pending?.resume() + } + + private func waitForFirstResolutionSignal(id: UUID) async throws { + try await withCheckedThrowingContinuation { continuation in + var result: Result? + lock.lock() + if storedFirstResolutionEntered { + result = .success(()) + } else if let error = cancelledFirstResolutionWaiters.removeValue(forKey: id) { + result = .failure(error) + } else if Task.isCancelled { + result = .failure(CancellationError()) + } else { + firstResolutionWaiters.append(ResolutionWaiter(id: id, continuation: continuation)) + } + lock.unlock() - private func waitForFirstResolutionSignal(id: UUID) async { - await withCheckedContinuation { continuation in - if firstResolutionEntered { + switch result { + case .success: continuation.resume() - } else { - firstResolutionWaiters.append((id, continuation)) + case let .failure(error): + continuation.resume(throwing: error) + case nil: + break } } } - private func cancelFirstResolutionWaiter(id: UUID) { - var cancelled: [CheckedContinuation] = [] - firstResolutionWaiters.removeAll { waiter in - guard waiter.id == id else { return false } - cancelled.append(waiter.continuation) - return true + private func cancelFirstResolutionBlocker(waiterID: UUID) { + let pending = lock.withLock { () -> CheckedContinuation? in + if let pending = self.continuation { + self.continuation = nil + return pending + } + cancelledWaiters.insert(waiterID) + return nil } - for continuation in cancelled { - continuation.resume() + pending?.resume() + } + + private func cancelFirstResolutionWaiter(id: UUID, error: Error) { + let continuation = lock.withLock { () -> CheckedContinuation? in + if let index = firstResolutionWaiters.firstIndex(where: { $0.id == id }) { + return firstResolutionWaiters.remove(at: index).continuation + } + if cancelledFirstResolutionWaiters[id] == nil { + cancelledFirstResolutionWaiters[id] = error + } + return nil } + continuation?.resume(throwing: error) } } diff --git a/Tests/RepoPromptTests/Helpers/CodemapSeamTestSupport.swift b/Tests/RepoPromptTests/Helpers/CodemapSeamTestSupport.swift index 30c0ef48d..53e4029f4 100644 --- a/Tests/RepoPromptTests/Helpers/CodemapSeamTestSupport.swift +++ b/Tests/RepoPromptTests/Helpers/CodemapSeamTestSupport.swift @@ -776,6 +776,7 @@ class WorkspaceFileContextStoreCodemapSeamTestSupport: XCTestCase { } return true } catch { + XCTFail(error.localizedDescription) return false } } @@ -1302,6 +1303,7 @@ final class CodemapSelectionGraphProbe: @unchecked Sendable { } return true } catch { + XCTFail(error.localizedDescription) return false } } @@ -1320,6 +1322,7 @@ final class CodemapSelectionGraphProbe: @unchecked Sendable { } return true } catch { + XCTFail(error.localizedDescription) return false } } @@ -1370,6 +1373,7 @@ final class CodemapSelectionGraphProbe: @unchecked Sendable { } return matchedKey } catch { + XCTFail(error.localizedDescription) return nil } } @@ -1394,6 +1398,7 @@ final class CodemapSelectionGraphProbe: @unchecked Sendable { } return matchedKey } catch { + XCTFail(error.localizedDescription) return nil } } @@ -1432,7 +1437,10 @@ final class CodemapSelectionGraphBuildGate: @unchecked Sendable { defer { condition.unlock() } let deadline = Date(timeIntervalSinceNow: 10) while blockedGenerations.isEmpty { - guard condition.wait(until: deadline) else { return nil } + guard condition.wait(until: deadline) else { + XCTFail("Timed out waiting for codemap selection graph build gate to block") + return nil + } } return blockedGenerations[0] } @@ -1442,7 +1450,10 @@ final class CodemapSelectionGraphBuildGate: @unchecked Sendable { defer { condition.unlock() } let deadline = Date(timeIntervalSinceNow: 10) while !blockedGenerations.contains(where: { $0 > generation }) { - guard condition.wait(until: deadline) else { return nil } + guard condition.wait(until: deadline) else { + XCTFail("Timed out waiting for codemap selection graph build gate after generation \(generation)") + return nil + } } return blockedGenerations.first(where: { $0 > generation }) } @@ -1472,11 +1483,29 @@ final class CodemapSelectionGraphBuildGate: @unchecked Sendable { if let autoReleaseTimeout { let deadline = Date(timeIntervalSinceNow: autoReleaseTimeout) while !isOpen, !releasedGenerations.contains(generation) { - guard condition.wait(until: deadline) else { break } + guard condition.wait(until: deadline) else { + XCTFail( + "Timed out waiting to release codemap selection graph generation \(generation) " + + "after \(String(format: "%.1f", autoReleaseTimeout))s. " + + "Pass a larger autoReleaseTimeout if this release is intentionally slow." + ) + // Fail open consistently with the nil-timeout path: open the gate so sibling + // parallel generations are not left blocked until their own deadlines. + isOpen = true + condition.broadcast() + break + } } } else { + let watchdogDeadline = Date(timeIntervalSinceNow: 10) while !isOpen, !releasedGenerations.contains(generation) { - condition.wait() + guard condition.wait(until: watchdogDeadline) else { + XCTFail("Codemap selection graph generation \(generation) is still blocked without release") + // Fail open: open the gate so the blocked builder (and siblings) can progress. + isOpen = true + condition.broadcast() + break + } } } condition.unlock() @@ -1624,13 +1653,9 @@ final class CodemapRetryTestClock: @unchecked Sendable { } actor CodemapAutomaticSelectionSequenceHarness { + private let waitState = CodemapAutomaticSelectionWaitState() private var demandTickets: [WorkspaceCodemapArtifactDemandTicket] = [] private var waiterInvocationCount = 0 - private var releasedWaits = Set() - private var releaseAllWaits = false - private var continuations: [Int: CheckedContinuation] = [:] - private var demandCountObservers: [UUID: AsyncStream<[WorkspaceCodemapArtifactDemandTicket]>.Continuation] = [:] - private var waitCountObservers: [UUID: AsyncStream.Continuation] = [:] var recordedTickets: [WorkspaceCodemapArtifactDemandTicket] { demandTickets @@ -1642,7 +1667,6 @@ actor CodemapAutomaticSelectionSequenceHarness { func recordDemand(_ ticket: WorkspaceCodemapArtifactDemandTicket) -> Int { demandTickets.append(ticket) - publishDemandTickets() return demandTickets.count } @@ -1650,427 +1674,445 @@ actor CodemapAutomaticSelectionSequenceHarness { try Task.checkCancellation() waiterInvocationCount += 1 let invocation = waiterInvocationCount - publishWaitCount() - guard !releaseAllWaits, !releasedWaits.contains(invocation) else { return } - try await withTaskCancellationHandler { - try await withCheckedThrowingContinuation { - (continuation: CheckedContinuation) in - if releaseAllWaits || releasedWaits.contains(invocation) || Task.isCancelled { - continuation.resume(throwing: CancellationError()) - } else { - continuations[invocation] = continuation - } - } - } onCancel: { - Task { await self.cancelWait(invocation) } - } + try await waitState.wait(invocation: invocation) } func releaseWait(_ invocation: Int) { - releasedWaits.insert(invocation) - continuations.removeValue(forKey: invocation)?.resume(returning: ()) + waitState.releaseWait(invocation) } func releaseAll() { - releaseAllWaits = true - let pending = Array(continuations.values) - continuations.removeAll() - let demandObservers = Array(demandCountObservers.values) - demandCountObservers.removeAll() - let waitObservers = Array(waitCountObservers.values) - waitCountObservers.removeAll() - for continuation in pending { - continuation.resume(returning: ()) - } - for observer in demandObservers { - observer.finish() - } - for observer in waitObservers { - observer.finish() - } + waitState.releaseAll() } func waitUntilDemandCount( - _ expectedCount: Int + _ expectedCount: Int, + timeout: Duration = .seconds(10) ) async -> [WorkspaceCodemapArtifactDemandTicket]? { if demandTickets.count >= expectedCount { return demandTickets } - let stream = demandTicketStream() - for await tickets in stream { - if tickets.count >= expectedCount { return tickets } - guard !Task.isCancelled else { return nil } + do { + try await AsyncTestWait.waitUntil( + "codemap automatic selection demand count \(expectedCount)", + timeout: WorkspaceFileContextStoreCodemapSeamTestSupport.timeInterval(timeout) + ) { + await self.recordedTickets.count >= expectedCount + } + return demandTickets + } catch { + XCTFail(error.localizedDescription) + return demandTickets.count >= expectedCount ? demandTickets : nil } - return demandTickets.count >= expectedCount ? demandTickets : nil } - func waitUntilWaitCount(_ expectedCount: Int) async -> Bool { + func waitUntilWaitCount(_ expectedCount: Int, timeout: Duration = .seconds(10)) async -> Bool { if waiterInvocationCount >= expectedCount { return true } - let stream = waitCountStream() - for await count in stream { - if count >= expectedCount { return true } - guard !Task.isCancelled else { return false } + do { + try await AsyncTestWait.waitUntil( + "codemap automatic selection wait count \(expectedCount)", + timeout: WorkspaceFileContextStoreCodemapSeamTestSupport.timeInterval(timeout) + ) { + await self.waitCount >= expectedCount + } + return true + } catch { + XCTFail(error.localizedDescription) + return waiterInvocationCount >= expectedCount } - return waiterInvocationCount >= expectedCount } +} - private func demandTicketStream() -> AsyncStream<[WorkspaceCodemapArtifactDemandTicket]> { - let id = UUID() - let (stream, continuation) = AsyncStream<[WorkspaceCodemapArtifactDemandTicket]>.makeStream( - bufferingPolicy: .bufferingNewest(1) - ) - demandCountObservers[id] = continuation - continuation.onTermination = { [weak self] _ in - Task { await self?.removeDemandObserver(id) } +private final class CodemapAutomaticSelectionWaitState: @unchecked Sendable { + private let lock = NSLock() + private var releasedWaits = Set() + private var cancelledWaits = Set() + private var releaseAllWaits = false + private var continuations: [Int: CheckedContinuation] = [:] + + func wait(invocation: Int) async throws { + try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + var result: Result? + lock.lock() + if releaseAllWaits || releasedWaits.contains(invocation) { + result = .success(()) + } else if Task.isCancelled || cancelledWaits.remove(invocation) != nil { + result = .failure(CancellationError()) + } else { + continuations[invocation] = continuation + } + lock.unlock() + + switch result { + case .success: + continuation.resume(returning: ()) + case let .failure(error): + continuation.resume(throwing: error) + case nil: + break + } + } + } onCancel: { + cancelWait(invocation) } - continuation.yield(demandTickets) - return stream } - private func waitCountStream() -> AsyncStream { - let id = UUID() - let (stream, continuation) = AsyncStream.makeStream( - bufferingPolicy: .bufferingNewest(1) - ) - waitCountObservers[id] = continuation - continuation.onTermination = { [weak self] _ in - Task { await self?.removeWaitObserver(id) } + func releaseWait(_ invocation: Int) { + let continuation = lock.withLock { () -> CheckedContinuation? in + releasedWaits.insert(invocation) + return continuations.removeValue(forKey: invocation) } - continuation.yield(waiterInvocationCount) - return stream + continuation?.resume(returning: ()) } - private func publishDemandTickets() { - for continuation in demandCountObservers.values { - continuation.yield(demandTickets) + func releaseAll() { + let pending = lock.withLock { () -> [CheckedContinuation] in + releaseAllWaits = true + let pending = Array(continuations.values) + continuations.removeAll() + cancelledWaits.removeAll() + return pending + } + for continuation in pending { + continuation.resume(returning: ()) } } - private func publishWaitCount() { - for continuation in waitCountObservers.values { - continuation.yield(waiterInvocationCount) + private func cancelWait(_ invocation: Int) { + let continuation = lock.withLock { () -> CheckedContinuation? in + if let continuation = continuations.removeValue(forKey: invocation) { + return continuation + } + cancelledWaits.insert(invocation) + return nil } + continuation?.resume(throwing: CancellationError()) } +} + +actor CodemapRetrySleepGate { + private let state = CodemapRetrySleepGateState() - private func removeDemandObserver(_ id: UUID) { - demandCountObservers.removeValue(forKey: id) + var delays: [UInt64] { + state.delays } - private func removeWaitObserver(_ id: UUID) { - waitCountObservers.removeValue(forKey: id) + func sleep(_ nanoseconds: UInt64) async throws { + try await state.sleep(nanoseconds) } - private func cancelWait(_ invocation: Int) { - continuations.removeValue(forKey: invocation)?.resume(throwing: CancellationError()) + func waitForFirstDelay(timeout: Duration = .seconds(10)) async -> UInt64? { + do { + return try await state.waitForFirstDelay( + timeout: WorkspaceFileContextStoreCodemapSeamTestSupport.timeInterval(timeout) + ) + } catch { + if let delay = delays.first { return delay } + XCTFail(error.localizedDescription) + return nil + } + } + + func releaseAll() { + state.releaseAll() } } -actor CodemapRetrySleepGate { - private(set) var delays: [UInt64] = [] +private final class CodemapRetrySleepGateState: @unchecked Sendable { + private struct DelayWaiter { + let id: UUID + let continuation: CheckedContinuation + } + + private let lock = NSLock() + private var storedDelays: [UInt64] = [] private var released = false - private var continuations: [UUID: CheckedContinuation] = [:] + private var sleepContinuations: [UUID: CheckedContinuation] = [:] + private var cancelledSleepWaiters = Set() + private var cancelledDelayWaiters: [UUID: Error] = [:] + private var delayWaiters: [DelayWaiter] = [] + + var delays: [UInt64] { + lock.withLock { storedDelays } + } func sleep(_ nanoseconds: UInt64) async throws { - delays.append(nanoseconds) + let readyDelayWaiters = lock.withLock { () -> [DelayWaiter] in + storedDelays.append(nanoseconds) + let waiters = delayWaiters + delayWaiters = [] + return waiters + } + for waiter in readyDelayWaiters { + waiter.continuation.resume(returning: nanoseconds) + } + try Task.checkCancellation() - guard !released else { return } let waiterID = UUID() try await withTaskCancellationHandler { try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in - if released || Task.isCancelled { - continuation.resume(throwing: CancellationError()) + var result: Result? + lock.lock() + // releaseAll() unblocks sleepers as success (proceed past the sleep), not cancellation. + if released { + result = .success(()) + } else if Task.isCancelled || cancelledSleepWaiters.remove(waiterID) != nil { + result = .failure(CancellationError()) } else { - continuations[waiterID] = continuation + sleepContinuations[waiterID] = continuation + } + lock.unlock() + + switch result { + case .success: + continuation.resume(returning: ()) + case let .failure(error): + continuation.resume(throwing: error) + case nil: + break } } } onCancel: { - Task { await self.cancel(waiterID) } + cancelSleep(waiterID) } } - func waitForFirstDelay(timeout: Duration = .seconds(10)) async -> UInt64? { - let clock = ContinuousClock() - let deadline = clock.now.advanced(by: timeout) - while delays.isEmpty, clock.now < deadline { - await Task.yield() + func waitForFirstDelay(timeout: TimeInterval) async throws -> UInt64? { + if let delay = delays.first { return delay } + let waiterID = UUID() + do { + return try await withThrowingTaskGroup(of: UInt64?.self) { group in + group.addTask { + try await self.waitForFirstDelaySignal(id: waiterID) + } + group.addTask { + try await Task.sleep(nanoseconds: UInt64((timeout * 1_000_000_000).rounded())) + self.cancelDelayWaiter( + id: waiterID, + error: AsyncTestConditionTimeout(description: "codemap retry first delay", timeout: timeout) + ) + throw AsyncTestConditionTimeout(description: "codemap retry first delay", timeout: timeout) + } + defer { + group.cancelAll() + cancelDelayWaiter(id: waiterID, error: CancellationError()) + } + return try await group.next() ?? nil + } + } catch { + if let delay = delays.first { return delay } + throw error } - return delays.first } func releaseAll() { - released = true - let pending = Array(continuations.values) - continuations.removeAll() - for continuation in pending { + // Intentionally resume sleep waiters with success so callers proceed past the sleep gate. + // Also drain first-delay observers so teardown does not sit on the waiter timeout. + let pending = lock.withLock { () -> ( + sleep: [CheckedContinuation], + delay: [CheckedContinuation] + ) in + released = true + let sleep = Array(sleepContinuations.values) + sleepContinuations.removeAll() + let delay = delayWaiters.map(\.continuation) + delayWaiters.removeAll() + cancelledSleepWaiters.removeAll() + cancelledDelayWaiters.removeAll() + return (sleep, delay) + } + for continuation in pending.sleep { continuation.resume(returning: ()) } + for continuation in pending.delay { + continuation.resume(returning: nil) + } } - private func cancel(_ waiterID: UUID) { - continuations.removeValue(forKey: waiterID)?.resume(throwing: CancellationError()) - } -} - -actor CodemapSuspensionGate { - private var entered = false - private var released = false - private var continuations: [UUID: CheckedContinuation] = [:] + private func waitForFirstDelaySignal(id: UUID) async throws -> UInt64? { + try await withCheckedThrowingContinuation { continuation in + var result: Result? + lock.lock() + if let delay = storedDelays.first { + result = .success(delay) + } else if let error = cancelledDelayWaiters.removeValue(forKey: id) { + result = .failure(error) + } else if Task.isCancelled { + result = .failure(CancellationError()) + } else { + delayWaiters.append(DelayWaiter(id: id, continuation: continuation)) + } + lock.unlock() - func enterAndWait() async { - entered = true - guard !released, !Task.isCancelled else { return } - let waiterID = UUID() - await withTaskCancellationHandler { - await withCheckedContinuation { continuation in - if released || Task.isCancelled { - continuation.resume() - } else { - continuations[waiterID] = continuation - } + switch result { + case let .success(value): + continuation.resume(returning: value) + case let .failure(error): + continuation.resume(throwing: error) + case nil: + break } - } onCancel: { - Task { await self.cancel(waiterID) } } } - func waitUntilEntered(timeout: Duration = .seconds(10)) async -> Bool { - do { - try await AsyncTestWait.waitUntil( - "codemap suspension gate entered", - timeout: WorkspaceFileContextStoreCodemapSeamTestSupport.timeInterval(timeout) - ) { - self.entered + private func cancelSleep(_ waiterID: UUID) { + let continuation = lock.withLock { () -> CheckedContinuation? in + if let continuation = sleepContinuations.removeValue(forKey: waiterID) { + return continuation } - return true - } catch { - return entered + cancelledSleepWaiters.insert(waiterID) + return nil } + continuation?.resume(throwing: CancellationError()) } - func release() { - released = true - let pending = Array(continuations.values) - continuations.removeAll() - for continuation in pending { - continuation.resume() + private func cancelDelayWaiter(id: UUID, error: Error) { + let continuation = lock.withLock { () -> CheckedContinuation? in + if let index = delayWaiters.firstIndex(where: { $0.id == id }) { + return delayWaiters.remove(at: index).continuation + } + if cancelledDelayWaiters[id] == nil { + cancelledDelayWaiters[id] = error + } + return nil } + continuation?.resume(throwing: error) } +} - private func cancel(_ waiterID: UUID) { - continuations.removeValue(forKey: waiterID)?.resume() +/// Shared hang-hardened enter/release fence for codemap suspension-style tests. +final class CodemapSuspensionGate: @unchecked Sendable { + private let fence = TestReleaseFence(name: "codemap suspension gate") + + func enterAndWait() async { + await fence.enterAndWait() + } + + @discardableResult + func waitUntilEntered( + timeout: Duration = TestFenceDefaults.enterWaitDuration, + failOnTimeout: Bool = true + ) async -> Bool { + await fence.waitUntilEntered(timeout: timeout, failOnTimeout: failOnTimeout) + } + + func release() { + fence.release() } } -actor CodemapArmableSuspensionGate { +/// Armable suspension: only parks after `arm()`. +final class CodemapArmableSuspensionGate: @unchecked Sendable { + private let fence = TestReleaseFence(name: "codemap armable suspension gate") + private let lock = NSLock() private var armed = false - private var entered = false - private var released = false - private var continuations: [UUID: CheckedContinuation] = [:] func arm() { + lock.lock() armed = true + lock.unlock() } func enterIfArmedAndWait() async { - guard armed else { return } - entered = true - guard !released, !Task.isCancelled else { return } - let waiterID = UUID() - await withTaskCancellationHandler { - await withCheckedContinuation { continuation in - if released || Task.isCancelled { - continuation.resume() - } else { - continuations[waiterID] = continuation - } - } - } onCancel: { - Task { await self.cancel(waiterID) } - } + let isArmed = lock.withLock { armed } + guard isArmed else { return } + await fence.enterAndWait() } - func waitUntilEntered(timeout: Duration = .seconds(10)) async -> Bool { - do { - try await AsyncTestWait.waitUntil( - "codemap suspension gate entered", - timeout: WorkspaceFileContextStoreCodemapSeamTestSupport.timeInterval(timeout) - ) { - self.entered - } - return true - } catch { - return entered - } + @discardableResult + func waitUntilEntered(timeout: Duration = TestFenceDefaults.enterWaitDuration) async -> Bool { + await fence.waitUntilEntered(timeout: timeout) } func release() { - released = true - let pending = Array(continuations.values) - continuations.removeAll() - for continuation in pending { - continuation.resume() - } - } - - private func cancel(_ waiterID: UUID) { - continuations.removeValue(forKey: waiterID)?.resume() + fence.release() } } -actor CodemapGraphPublicationGate { +/// Records publication root epochs while parking on a shared release fence. +final class CodemapGraphPublicationGate: @unchecked Sendable { + private let fence = TestReleaseFence(name: "codemap graph publication gate") + private let lock = NSLock() private var invocationRoots: [WorkspaceCodemapRootEpoch] = [] - private var isOpen = false - private var continuations: [UUID: CheckedContinuation] = [:] var invocationCount: Int { - invocationRoots.count + lock.withLock { invocationRoots.count } } func enterAndWait(_ rootEpoch: WorkspaceCodemapRootEpoch) async { - invocationRoots.append(rootEpoch) - guard !isOpen, !Task.isCancelled else { return } - let waiterID = UUID() - await withTaskCancellationHandler { - await withCheckedContinuation { continuation in - if isOpen || Task.isCancelled { - continuation.resume() - } else { - continuations[waiterID] = continuation - } - } - } onCancel: { - Task { await self.cancel(waiterID) } - } + lock.withLock { invocationRoots.append(rootEpoch) } + await fence.enterAndWait() } func waitUntilInvocationCount( _ expectedCount: Int, - timeout: Duration = .seconds(10) + timeout: Duration = TestFenceDefaults.enterWaitDuration ) async -> Bool { do { try await AsyncTestWait.waitUntil( "codemap graph publication gate invocation count", - timeout: WorkspaceFileContextStoreCodemapSeamTestSupport.timeInterval(timeout) + timeout: TestFenceDefaults.timeInterval(timeout) ) { - self.invocationRoots.count >= expectedCount + self.invocationCount >= expectedCount } return true } catch { - return invocationRoots.count >= expectedCount + XCTFail(error.localizedDescription) + return invocationCount >= expectedCount } } func release() { - isOpen = true - let pending = Array(continuations.values) - continuations.removeAll() - for continuation in pending { - continuation.resume() - } - } - - private func cancel(_ waiterID: UUID) { - continuations.removeValue(forKey: waiterID)?.resume() + fence.release() } } -actor CodemapRootSuspensionGate { +/// Parks only on the first root epoch enter. +final class CodemapRootSuspensionGate: @unchecked Sendable { + private let fence = TestReleaseFence(name: "codemap root suspension gate") + private let lock = NSLock() private var enteredRootEpoch: WorkspaceCodemapRootEpoch? - private var released = false - private var continuations: [UUID: CheckedContinuation] = [:] func enterAndWait(_ rootEpoch: WorkspaceCodemapRootEpoch) async { - guard enteredRootEpoch == nil else { return } - enteredRootEpoch = rootEpoch - guard !released, !Task.isCancelled else { return } - let waiterID = UUID() - await withTaskCancellationHandler { - await withCheckedContinuation { continuation in - if released || Task.isCancelled { - continuation.resume() - } else { - continuations[waiterID] = continuation - } - } - } onCancel: { - Task { await self.cancel(waiterID) } + let shouldPark = lock.withLock { () -> Bool in + guard enteredRootEpoch == nil else { return false } + enteredRootEpoch = rootEpoch + return true } + guard shouldPark else { return } + await fence.enterAndWait() } - func waitUntilEntered(timeout: Duration = .seconds(10)) async -> WorkspaceCodemapRootEpoch? { - do { - try await AsyncTestWait.waitUntil( - "codemap root suspension gate entered", - timeout: WorkspaceFileContextStoreCodemapSeamTestSupport.timeInterval(timeout) - ) { - self.enteredRootEpoch != nil - } - return enteredRootEpoch - } catch { - return enteredRootEpoch - } + func waitUntilEntered(timeout: Duration = TestFenceDefaults.enterWaitDuration) async -> WorkspaceCodemapRootEpoch? { + _ = await fence.waitUntilEntered(timeout: timeout) + return lock.withLock { enteredRootEpoch } } func release() { - released = true - let pending = Array(continuations.values) - continuations.removeAll() - for continuation in pending { - continuation.resume() - } - } - - private func cancel(_ waiterID: UUID) { - continuations.removeValue(forKey: waiterID)?.resume() + fence.release() } } -actor CodemapResolutionGate { - private var entered = false - private var released = false - private var continuations: [UUID: CheckedContinuation] = [:] - private(set) var resolutionCount = 0 +/// Resolution-count enter fence for codemap resolution hooks. +final class CodemapResolutionGate: @unchecked Sendable { + private let fence = TestReleaseFence(name: "codemap resolution gate") + private let lock = NSLock() + private var storedResolutionCount = 0 - func enterAndWait() async { - resolutionCount += 1 - entered = true - guard !released, !Task.isCancelled else { return } - let waiterID = UUID() - await withTaskCancellationHandler { - await withCheckedContinuation { continuation in - if released || Task.isCancelled { - continuation.resume() - } else { - continuations[waiterID] = continuation - } - } - } onCancel: { - Task { await self.cancel(waiterID) } - } + var resolutionCount: Int { + lock.withLock { storedResolutionCount } } - func waitUntilEntered(timeout: Duration = .seconds(10)) async -> Bool { - do { - try await AsyncTestWait.waitUntil( - "codemap suspension gate entered", - timeout: WorkspaceFileContextStoreCodemapSeamTestSupport.timeInterval(timeout) - ) { - self.entered - } - return true - } catch { - return entered - } + func enterAndWait() async { + lock.withLock { storedResolutionCount += 1 } + await fence.enterAndWait() } - func release() { - released = true - let pending = Array(continuations.values) - continuations.removeAll() - for continuation in pending { - continuation.resume() - } + @discardableResult + func waitUntilEntered(timeout: Duration = TestFenceDefaults.enterWaitDuration) async -> Bool { + await fence.waitUntilEntered(timeout: timeout) } - private func cancel(_ waiterID: UUID) { - continuations.removeValue(forKey: waiterID)?.resume() + func release() { + fence.release() } } diff --git a/Tests/RepoPromptTests/Helpers/TestGitCommandRunner.swift b/Tests/RepoPromptTests/Helpers/TestGitCommandRunner.swift index 88c74097f..552cbcadb 100644 --- a/Tests/RepoPromptTests/Helpers/TestGitCommandRunner.swift +++ b/Tests/RepoPromptTests/Helpers/TestGitCommandRunner.swift @@ -28,27 +28,33 @@ enum TestGitCommandRunner { return processEnvironment } + /// Runs git with the same default timeout as `TestProcessRunner` (30s). + /// Pass an explicit larger `timeout` for known-heavy fixtures (cold clone, large packs). static func runResult( _ arguments: [String], cwd: URL, - environment: Environment = .hermetic + environment: Environment = .hermetic, + timeout: TimeInterval = TestProcessRunner.defaultTimeout ) throws -> TestProcessResult { try TestProcessRunner.run( executableURL: executableURL, arguments: arguments, currentDirectoryURL: cwd, - environment: processEnvironment(environment) + environment: processEnvironment(environment), + timeout: timeout ) } + /// See `runResult` for timeout guidance. @discardableResult static func run( _ arguments: [String], cwd: URL, environment: Environment = .hermetic, + timeout: TimeInterval = TestProcessRunner.defaultTimeout, failureDomain: String ) throws -> String { - let result = try runResult(arguments, cwd: cwd, environment: environment) + let result = try runResult(arguments, cwd: cwd, environment: environment, timeout: timeout) guard result.terminationStatus == 0 else { throw NSError( domain: failureDomain, diff --git a/Tests/RepoPromptTests/Helpers/TestHangHardenedFences.swift b/Tests/RepoPromptTests/Helpers/TestHangHardenedFences.swift new file mode 100644 index 000000000..22457be23 --- /dev/null +++ b/Tests/RepoPromptTests/Helpers/TestHangHardenedFences.swift @@ -0,0 +1,427 @@ +import Foundation +import XCTest + +// MARK: - Defaults + +/// Shared timeout budgets for hang-hardened test fences. +/// Prefer these over ad-hoc 2s / 10s / 30s literals so suites stay consistent. +enum TestFenceDefaults { + /// Waiting for a producer to mark "entered" / "started". + static let enterWait: TimeInterval = 10 + /// Waiting for an explicit `release()` while parked in enter/block. + static let releaseWait: TimeInterval = 30 + + static let enterWaitDuration: Duration = .seconds(10) + static let releaseWaitDuration: Duration = .seconds(30) + + static func timeInterval(_ duration: Duration) -> TimeInterval { + let components = duration.components + return TimeInterval(components.seconds) + + TimeInterval(components.attoseconds) / 1_000_000_000_000_000_000 + } +} + +// MARK: - Async release fence + +/// Hang-hardened async enter/release fence. +/// +/// Guarantees: +/// - sticky cancel (cancel-before-register cannot park forever) +/// - synchronous `onCancel` via lock-backed state (no `Task { await }` hop) +/// - multi-waiter release parks (intentional; double-enter parks both until one `release()`) +/// - **cooperative** async `waitUntilEntered` (never blocks the calling executor) +/// - optional sync `waitUntilEntered` for true sync call sites only +/// +/// Prefer this over local `AsyncGate` / `*ReleaseGate` clones. +final class TestReleaseFence: @unchecked Sendable { + private let name: String + private let condition = NSCondition() + private var entered = false + private var released = false + private var continuations: [UUID: CheckedContinuation] = [:] + private var cancelledWaiters = Set() + private var timedOutWaiters = Set() + + init(name: String = "test release fence") { + self.name = name + } + + /// Mark entered, then park until `release()` (or sticky cancel / already released). + func enterAndWait() async { + let waiterID = UUID() + await withTaskCancellationHandler { + await withCheckedContinuation { continuation in + register(continuation, waiterID: waiterID) + } + } onCancel: { + cancel(waiterID: waiterID) + } + } + + /// Mark entered, then park until `release()` while ignoring task cancellation. + /// + /// Use this only for tests whose contract is that a cancelled owner does not + /// release the underlying body/resource until the test explicitly releases it. + /// A retained/cancelled timeout fails open so a missed release cannot wedge the suite. + func enterAndWaitIgnoringCancellationUntilRelease( + timeout: TimeInterval = TestFenceDefaults.releaseWait + ) async { + let waiterID = UUID() + let timeoutTask = Task.detached { [weak self] in + try? await Task.sleep(for: .seconds(max(0, timeout))) + guard !Task.isCancelled else { return } + self?.timeout(waiterID: waiterID, timeout: timeout) + } + await withCheckedContinuation { continuation in + registerIgnoringCancellation(continuation, waiterID: waiterID) + } + timeoutTask.cancel() + await timeoutTask.value + } + + /// Alias used by older engine call sites (`EngineBuildGate.enter`). + /// Multi-waiter: a second concurrent enter parks another waiter until `release()`. + func enter() async { + await enterAndWait() + } + + /// Park without flipping entered (for multi-phase fences that mark entered separately). + func waitUnlessReleased() async { + guard !Task.isCancelled else { return } + let waiterID = UUID() + await withTaskCancellationHandler { + await withCheckedContinuation { continuation in + registerParkOnly(continuation, waiterID: waiterID) + } + } onCancel: { + cancel(waiterID: waiterID) + } + } + + /// Synchronous enter wait for **true sync call sites only** + /// (e.g. from a DispatchQueue callback or another non-async context). + /// Do **not** call this from an `async` function that also needs the enter producer + /// on the same serial executor — use `await waitUntilEntered(...)` instead. + @discardableResult + func waitUntilEnteredBlocking( + timeout: TimeInterval = TestFenceDefaults.enterWait, + failOnTimeout: Bool = true + ) -> Bool { + condition.lock() + defer { condition.unlock() } + let deadline = Date().addingTimeInterval(timeout) + while !entered { + guard condition.wait(until: deadline) else { + if failOnTimeout { + XCTFail("Timed out waiting for \(name) to enter after \(String(format: "%.1f", timeout))s") + } + return false + } + } + return true + } + + /// Cooperative async enter wait — does **not** block the calling executor. + @discardableResult + func waitUntilEntered( + timeout: TimeInterval, + failOnTimeout: Bool = true + ) async -> Bool { + if hasEntered { return true } + do { + try await AsyncTestWait.waitUntil( + "\(name) entered", + timeout: timeout + ) { + self.hasEntered + } + return true + } catch { + if failOnTimeout { + XCTFail(error.localizedDescription) + } + return hasEntered + } + } + + @discardableResult + func waitUntilEntered( + timeout: Duration = TestFenceDefaults.enterWaitDuration, + failOnTimeout: Bool = true + ) async -> Bool { + await waitUntilEntered( + timeout: TestFenceDefaults.timeInterval(timeout), + failOnTimeout: failOnTimeout + ) + } + + func release() { + condition.lock() + released = true + let pending = Array(continuations.values) + continuations.removeAll() + cancelledWaiters.removeAll() + timedOutWaiters.removeAll() + condition.broadcast() + condition.unlock() + for continuation in pending { + continuation.resume() + } + } + + var hasEntered: Bool { + condition.lock() + defer { condition.unlock() } + return entered + } + + var isReleased: Bool { + condition.lock() + defer { condition.unlock() } + return released + } + + // MARK: Private + + private func register(_ continuation: CheckedContinuation, waiterID: UUID) { + condition.lock() + entered = true + condition.broadcast() + if released || Task.isCancelled || cancelledWaiters.remove(waiterID) != nil { + condition.unlock() + continuation.resume() + } else { + continuations[waiterID] = continuation + condition.unlock() + } + } + + private func registerIgnoringCancellation(_ continuation: CheckedContinuation, waiterID: UUID) { + condition.lock() + entered = true + condition.broadcast() + if released || timedOutWaiters.remove(waiterID) != nil { + condition.unlock() + continuation.resume() + } else { + continuations[waiterID] = continuation + condition.unlock() + } + } + + private func registerParkOnly(_ continuation: CheckedContinuation, waiterID: UUID) { + condition.lock() + if released || Task.isCancelled || cancelledWaiters.remove(waiterID) != nil { + condition.unlock() + continuation.resume() + } else { + continuations[waiterID] = continuation + condition.unlock() + } + } + + private func cancel(waiterID: UUID) { + condition.lock() + let continuation = continuations.removeValue(forKey: waiterID) + if continuation == nil { + cancelledWaiters.insert(waiterID) + } + condition.broadcast() + condition.unlock() + continuation?.resume() + } + + private func timeout(waiterID: UUID, timeout: TimeInterval) { + condition.lock() + let continuation = continuations.removeValue(forKey: waiterID) + let shouldFail = continuation != nil || !released + if continuation == nil, !released { + timedOutWaiters.insert(waiterID) + } + condition.unlock() + guard shouldFail else { return } + XCTFail("Timed out waiting for \(name) release after \(String(format: "%.1f", timeout))s") + continuation?.resume() + } +} + +// MARK: - Sync blocking fence + +/// Hang-hardened synchronous enter/release fence (`NSCondition`). +/// +/// Guarantees: +/// - bounded wait for `release()` +/// - XCTFail + fail-open on timeout so a missed `release()` cannot wedge the process +/// - `waitUntilEntered` defaults to XCTFail on timeout (aligned with `TestReleaseFence`) +final class TestBlockingFence: @unchecked Sendable { + private let name: String + private let condition = NSCondition() + private var entered = false + private var released = false + + init(name: String = "test blocking fence") { + self.name = name + } + + func enterAndWait(timeout: TimeInterval = TestFenceDefaults.releaseWait) { + condition.lock() + entered = true + condition.broadcast() + let deadline = Date().addingTimeInterval(timeout) + while !released { + guard condition.wait(until: deadline) else { + XCTFail( + "Timed out waiting for \(name) release after \(String(format: "%.1f", timeout))s" + ) + // Fail open: unblock the waiter so a missed release cannot hang later tests. + released = true + condition.broadcast() + break + } + } + condition.unlock() + } + + @discardableResult + func waitUntilEntered( + timeout: TimeInterval = TestFenceDefaults.enterWait, + failOnTimeout: Bool = true + ) -> Bool { + condition.lock() + defer { condition.unlock() } + let deadline = Date().addingTimeInterval(timeout) + while !entered { + guard condition.wait(until: deadline) else { + if failOnTimeout { + XCTFail( + "Timed out waiting for \(name) to enter after \(String(format: "%.1f", timeout))s" + ) + } + return false + } + } + return true + } + + func release() { + condition.lock() + released = true + condition.broadcast() + condition.unlock() + } +} + +// MARK: - Cancellation probe gate + +/// Hang-hardened cancellation handshake for tests that cancel cooperative work mid-flight. +/// +/// Guarantees: +/// - sticky cancel (cancel-before-register fails closed with `CancellationError`) +/// - synchronous `onCancel` via lock-backed state (no `Task { await }` hop) +/// - multi-waiter support (each waiter has its own sticky id) +/// - always marks entered before cancel branching so `waitUntilEntered` cannot hang +/// - bounded cooperative `waitUntilEntered` +/// - `forceCancel()` for teardown / observation +final class TestCancellationGate: @unchecked Sendable { + private let name: String + private let lock = NSLock() + private var entered = false + private var cancelled = false + private var storedCancellationCount = 0 + private var continuations: [UUID: CheckedContinuation] = [:] + private var cancelledWaiters = Set() + + init(name: String = "test cancellation gate") { + self.name = name + } + + var cancellationCount: Int { + lock.withLock { storedCancellationCount } + } + + /// Test/debug entry visibility for wrappers that rethrow typed timeout errors. + var hasEnteredForTesting: Bool { + lock.withLock { entered } + } + + func waitUntilCancelled() async throws { + try Task.checkCancellation() + let waiterID = UUID() + try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + // Registration only fails closed (CancellationError) or parks; never .success. + var failError: Error? + lock.lock() + entered = true + if cancelled || Task.isCancelled || cancelledWaiters.remove(waiterID) != nil { + failError = CancellationError() + } else { + continuations[waiterID] = continuation + } + lock.unlock() + + if let failError { + continuation.resume(throwing: failError) + } + } + } onCancel: { + cancel(waiterID: waiterID) + } + } + + @discardableResult + func waitUntilEntered( + timeout: TimeInterval = TestFenceDefaults.enterWait, + failOnTimeout: Bool = true + ) async -> Bool { + if hasEnteredForTesting { return true } + do { + try await AsyncTestWait.waitUntil( + "\(name) entered", + timeout: timeout + ) { + self.hasEnteredForTesting + } + return true + } catch { + if failOnTimeout { + XCTFail(error.localizedDescription) + } + return hasEnteredForTesting + } + } + + /// Explicit cancel for cleanup paths that are not driven by `Task.cancel()`. + /// Cancels all parked waiters and stamps sticky cancel for late registrants. + func forceCancel() { + lock.lock() + let alreadyCancelled = cancelled + cancelled = true + if !alreadyCancelled { + storedCancellationCount += 1 + } + let pending = Array(continuations.values) + continuations.removeAll() + lock.unlock() + for continuation in pending { + continuation.resume(throwing: CancellationError()) + } + } + + // MARK: Private + + private func cancel(waiterID: UUID) { + lock.lock() + let alreadyCancelled = cancelled + cancelled = true + if !alreadyCancelled { + storedCancellationCount += 1 + } + let continuation = continuations.removeValue(forKey: waiterID) + if continuation == nil { + cancelledWaiters.insert(waiterID) + } + lock.unlock() + continuation?.resume(throwing: CancellationError()) + } +} diff --git a/Tests/RepoPromptTests/Helpers/TestProcessRunner.swift b/Tests/RepoPromptTests/Helpers/TestProcessRunner.swift index 1334c46dd..2b87ca998 100644 --- a/Tests/RepoPromptTests/Helpers/TestProcessRunner.swift +++ b/Tests/RepoPromptTests/Helpers/TestProcessRunner.swift @@ -1,4 +1,10 @@ import Foundation +import XCTest +#if os(macOS) + import Darwin +#elseif os(Linux) + import Glibc +#endif struct TestProcessResult { let terminationStatus: Int32 @@ -9,13 +15,122 @@ struct TestProcessResult { } } +struct TestProcessTimeoutError: Error, LocalizedError, CustomStringConvertible { + let executableURL: URL + let arguments: [String] + let currentDirectoryURL: URL? + let timeout: TimeInterval + let output: Data + + var outputText: String { + String(decoding: output, as: UTF8.self) + } + + var errorDescription: String? { + description + } + + var description: String { + var parts = [ + "Process timed out after \(String(format: "%.3f", timeout))s:", + ([executableURL.path] + arguments).joined(separator: " ") + ] + if let currentDirectoryURL { + parts.append("cwd: \(currentDirectoryURL.path)") + } + if !output.isEmpty { + parts.append("captured output:\n\(outputText)") + } + return parts.joined(separator: "\n") + } +} + +/// Raised when the child exits within the deadline but stdout/stderr drain does not complete. +struct TestProcessOutputDrainTimeoutError: Error, LocalizedError, CustomStringConvertible { + let executableURL: URL + let arguments: [String] + let currentDirectoryURL: URL? + let drainTimeout: TimeInterval + let terminationStatus: Int32 + let output: Data + + var outputText: String { + String(decoding: output, as: UTF8.self) + } + + var errorDescription: String? { + description + } + + var description: String { + var parts = [ + "Process exited (status \(terminationStatus)) but output drain timed out after \(String(format: "%.3f", drainTimeout))s:", + ([executableURL.path] + arguments).joined(separator: " ") + ] + if let currentDirectoryURL { + parts.append("cwd: \(currentDirectoryURL.path)") + } + if !output.isEmpty { + parts.append("captured output:\n\(outputText)") + } + return parts.joined(separator: "\n") + } +} + enum TestProcessRunner { + /// Default wall-clock budget for `run`. Keep this modest so a wedged child cannot + /// stall the suite; known heavy call sites (cold git fixtures, large clones) should + /// pass an explicit larger `timeout:` rather than raising the global default. + static let defaultTimeout: TimeInterval = 30 + private static let terminationGraceInterval: TimeInterval = 1 + private static let outputDrainGraceInterval: TimeInterval = 1 + /// Tight budget for assistive pgrep child discovery (primary kill is process-group). + private static let childPIDQueryTimeout: TimeInterval = 0.2 + /// Cap emergency pgrep tree walks so hang recovery cannot spend unbounded time. + /// Process-group `kill(-pgid)` remains primary for the posix_spawn path; this walk is + /// assistive only for Foundation `Process` trees. + private static let maxProcessTreeWalkDepth = 3 + private static let maxProcessTreeWalkNodes = 16 + /// After kill grace, unreaped pids are tracked for best-effort `WNOHANG` reaping. + /// Residual zombies (D-state / unkillable) live only for the suite process lifetime. + private static let maxAbandonedPIDsBeforeFail = 16 + static func run( executableURL: URL, arguments: [String] = [], currentDirectoryURL: URL? = nil, - environment: [String: String]? = nil + environment: [String: String]? = nil, + timeout: TimeInterval = defaultTimeout ) throws -> TestProcessResult { + #if os(macOS) || os(Linux) + reapAbandonedChildren() + return try runWithSpawnedProcessGroup( + executableURL: executableURL, + arguments: arguments, + currentDirectoryURL: currentDirectoryURL, + environment: environment, + timeout: timeout + ) + #else + try runWithFoundationProcess( + executableURL: executableURL, + arguments: arguments, + currentDirectoryURL: currentDirectoryURL, + environment: environment, + timeout: timeout + ) + #endif + } + + private static func runWithFoundationProcess( + executableURL: URL, + arguments: [String], + currentDirectoryURL: URL?, + environment: [String: String]?, + timeout: TimeInterval + ) throws -> TestProcessResult { + precondition(timeout > 0, "TestProcessRunner timeout must be positive") + let process = Process() process.executableURL = executableURL process.arguments = arguments @@ -26,13 +141,612 @@ enum TestProcessRunner { process.standardOutput = output process.standardError = output - try process.run() - let outputData = output.fileHandleForReading.readDataToEndOfFile() - process.waitUntilExit() + let capturedOutput = LockedOutput() + let readerGroup = DispatchGroup() + let outputReader = output.fileHandleForReading + readerGroup.enter() + DispatchQueue.global(qos: .userInitiated).async { + defer { readerGroup.leave() } + while true { + guard let chunk = try? outputReader.read(upToCount: 64 * 1024), !chunk.isEmpty else { + break + } + capturedOutput.append(chunk) + } + } + + let terminationGroup = DispatchGroup() + terminationGroup.enter() + process.terminationHandler = { _ in + terminationGroup.leave() + } + + do { + try process.run() + } catch { + close(output.fileHandleForReading) + close(output.fileHandleForWriting) + readerGroup.wait() + throw error + } + + close(output.fileHandleForWriting) + + if terminationGroup.wait(timeout: .now() + timeout) == .timedOut { + terminate(process) + if terminationGroup.wait(timeout: .now() + terminationGraceInterval) == .timedOut { + forceTerminate(process) + _ = terminationGroup.wait(timeout: .now() + terminationGraceInterval) + } + + finishReadingAfterTimeout(output.fileHandleForReading, readerGroup: readerGroup) + throw TestProcessTimeoutError( + executableURL: executableURL, + arguments: arguments, + currentDirectoryURL: currentDirectoryURL, + timeout: timeout, + output: capturedOutput.data() + ) + } + + if finishReadingAfterProcessExit(output.fileHandleForReading, readerGroup: readerGroup) == false { + // Best-effort cleanup for any orphaned descendants still holding the pipe. + terminate(process) + forceTerminate(process) + throw TestProcessOutputDrainTimeoutError( + executableURL: executableURL, + arguments: arguments, + currentDirectoryURL: currentDirectoryURL, + drainTimeout: outputDrainGraceInterval, + terminationStatus: process.terminationStatus, + output: capturedOutput.data() + ) + } return TestProcessResult( terminationStatus: process.terminationStatus, - output: outputData + output: capturedOutput.data() ) } + + #if os(macOS) || os(Linux) + private static func runWithSpawnedProcessGroup( + executableURL: URL, + arguments: [String], + currentDirectoryURL: URL?, + environment: [String: String]?, + timeout: TimeInterval + ) throws -> TestProcessResult { + precondition(timeout > 0, "TestProcessRunner timeout must be positive") + + var outputPipe: [Int32] = [-1, -1] + guard pipe(&outputPipe) == 0 else { + throw POSIXError(POSIXErrorCode(rawValue: errno) ?? .EIO) + } + + func closePipe() { + if outputPipe[0] >= 0 { + systemClose(outputPipe[0]) + outputPipe[0] = -1 + } + if outputPipe[1] >= 0 { + systemClose(outputPipe[1]) + outputPipe[1] = -1 + } + } + + do { + try setCloseOnExec(outputPipe[0]) + try setCloseOnExec(outputPipe[1]) + } catch { + closePipe() + throw error + } + + #if os(macOS) + var fileActions: posix_spawn_file_actions_t? = nil + #else + var fileActions = posix_spawn_file_actions_t() + #endif + var result = posix_spawn_file_actions_init(&fileActions) + guard result == 0 else { + closePipe() + throw POSIXError(POSIXErrorCode(rawValue: result) ?? .EIO) + } + defer { posix_spawn_file_actions_destroy(&fileActions) } + + func checkFileAction(_ result: Int32) throws { + guard result == 0 else { + closePipe() + throw POSIXError(POSIXErrorCode(rawValue: result) ?? .EIO) + } + } + + // Under POSIX_SPAWN_CLOEXEC_DEFAULT (Darwin), fds not explicitly installed are closed. + // Open stdin to /dev/null so children match Foundation.Process (valid fd 0, EOF on read) + // rather than a closed stdin that some tools probe differently. + try checkFileAction( + posix_spawn_file_actions_addopen( + &fileActions, + STDIN_FILENO, + "/dev/null", + O_RDONLY, + 0 + ) + ) + try checkFileAction(posix_spawn_file_actions_adddup2(&fileActions, outputPipe[1], STDOUT_FILENO)) + try checkFileAction(posix_spawn_file_actions_adddup2(&fileActions, outputPipe[1], STDERR_FILENO)) + try checkFileAction(posix_spawn_file_actions_addclose(&fileActions, outputPipe[0])) + try checkFileAction(posix_spawn_file_actions_addclose(&fileActions, outputPipe[1])) + + if let currentDirectoryURL { + result = currentDirectoryURL.path.withCString { path in + posix_spawn_file_actions_addchdir_np(&fileActions, path) + } + try checkFileAction(result) + } + + #if os(macOS) + var attributes: posix_spawnattr_t? = nil + #else + var attributes = posix_spawnattr_t() + #endif + result = posix_spawnattr_init(&attributes) + guard result == 0 else { + closePipe() + throw POSIXError(POSIXErrorCode(rawValue: result) ?? .EIO) + } + defer { posix_spawnattr_destroy(&attributes) } + + var spawnFlags = Int16(POSIX_SPAWN_SETPGROUP) + #if canImport(Darwin) + spawnFlags |= Int16(POSIX_SPAWN_CLOEXEC_DEFAULT) + #endif + result = posix_spawnattr_setpgroup(&attributes, 0) + guard result == 0 else { + closePipe() + throw POSIXError(POSIXErrorCode(rawValue: result) ?? .EIO) + } + result = posix_spawnattr_setflags(&attributes, spawnFlags) + guard result == 0 else { + closePipe() + throw POSIXError(POSIXErrorCode(rawValue: result) ?? .EIO) + } + + var argv: [UnsafeMutablePointer?] = [] + argv.reserveCapacity(arguments.count + 2) + argv.append(strdup(executableURL.path)) + for argument in arguments { + argv.append(strdup(argument)) + } + argv.append(nil) + defer { + for pointer in argv where pointer != nil { + free(pointer) + } + } + + let processEnvironment = environment ?? ProcessInfo.processInfo.environment + var envp: [UnsafeMutablePointer?] = [] + envp.reserveCapacity(processEnvironment.count + 1) + for (key, value) in processEnvironment { + envp.append(strdup("\(key)=\(value)")) + } + envp.append(nil) + defer { + for pointer in envp where pointer != nil { + free(pointer) + } + } + + var pid: pid_t = 0 + result = posix_spawn( + &pid, + executableURL.path, + &fileActions, + &attributes, + argv, + envp + ) + guard result == 0 else { + closePipe() + throw POSIXError(POSIXErrorCode(rawValue: result) ?? .EIO) + } + + systemClose(outputPipe[1]) + outputPipe[1] = -1 + + let capturedOutput = LockedOutput() + let readerGroup = DispatchGroup() + let outputReader = FileHandle(fileDescriptor: outputPipe[0], closeOnDealloc: true) + outputPipe[0] = -1 + readerGroup.enter() + DispatchQueue.global(qos: .userInitiated).async { + defer { readerGroup.leave() } + while true { + guard let chunk = try? outputReader.read(upToCount: 64 * 1024), !chunk.isEmpty else { + break + } + capturedOutput.append(chunk) + } + } + + let terminationGroup = DispatchGroup() + let waitStatusBox = LockedWaitStatus() + let spawnedPID = pid + terminationGroup.enter() + DispatchQueue.global(qos: .userInitiated).async { + var status: Int32 = 0 + var waitResult: pid_t = 0 + repeat { + waitResult = waitpid(spawnedPID, &status, 0) + } while waitResult < 0 && errno == EINTR + + if waitResult == spawnedPID { + waitStatusBox.set(.exited(status)) + noteReaped(spawnedPID) + } else { + let errorNumber = errno + waitStatusBox.set(.failed(errorNumber)) + if errorNumber == ECHILD { + noteReaped(spawnedPID) + } + } + terminationGroup.leave() + } + + if terminationGroup.wait(timeout: .now() + timeout) == .timedOut { + signalProcessGroup(rootPID: pid, SIGTERM) + if terminationGroup.wait(timeout: .now() + terminationGraceInterval) == .timedOut { + signalProcessGroup(rootPID: pid, SIGKILL) + if terminationGroup.wait(timeout: .now() + terminationGraceInterval) == .timedOut { + // waitpid still blocked (D-state / unkillable). Track for best-effort reaping; + // residual zombies last only for the suite process lifetime. + noteAbandoned(spawnedPID) + } + } + finishReadingAfterTimeout(outputReader, readerGroup: readerGroup) + throw TestProcessTimeoutError( + executableURL: executableURL, + arguments: arguments, + currentDirectoryURL: currentDirectoryURL, + timeout: timeout, + output: capturedOutput.data() + ) + } + + let status: Int32 + switch waitStatusBox.value { + case let .exited(waitStatus): + status = terminationStatus(fromWaitStatus: waitStatus) + case let .failed(errorNumber): + throw POSIXError(POSIXErrorCode(rawValue: errorNumber) ?? .ECHILD) + case nil: + throw POSIXError(.ECHILD) + } + if finishReadingAfterProcessExit(outputReader, readerGroup: readerGroup) == false { + signalProcessGroup(rootPID: pid, SIGTERM) + signalProcessGroup(rootPID: pid, SIGKILL) + throw TestProcessOutputDrainTimeoutError( + executableURL: executableURL, + arguments: arguments, + currentDirectoryURL: currentDirectoryURL, + drainTimeout: outputDrainGraceInterval, + terminationStatus: status, + output: capturedOutput.data() + ) + } + + return TestProcessResult( + terminationStatus: status, + output: capturedOutput.data() + ) + } + #endif + + private static func terminate(_ process: Process) { + #if os(macOS) + signal(process, SIGTERM) + #elseif os(Linux) + signal(process, SIGTERM) + #endif + if process.isRunning { + process.terminate() + } + } + + private static func forceTerminate(_ process: Process) { + #if os(macOS) + signal(process, SIGKILL) + #elseif os(Linux) + signal(process, SIGKILL) + #endif + } + + private static func close(_ handle: FileHandle) { + do { + try handle.close() + } catch { + handle.closeFile() + } + } + + private static func finishReadingAfterProcessExit(_ handle: FileHandle, readerGroup: DispatchGroup) -> Bool { + if readerGroup.wait(timeout: .now() + outputDrainGraceInterval) == .timedOut { + close(handle) + _ = readerGroup.wait(timeout: .now() + outputDrainGraceInterval) + return false + } + close(handle) + return true + } + + private static func finishReadingAfterTimeout(_ handle: FileHandle, readerGroup: DispatchGroup) { + if readerGroup.wait(timeout: .now() + outputDrainGraceInterval) == .timedOut { + close(handle) + _ = readerGroup.wait(timeout: .now() + outputDrainGraceInterval) + } else { + close(handle) + } + } + + #if os(macOS) || os(Linux) + private static func systemClose(_ fd: Int32) { + #if os(macOS) + Darwin.close(fd) + #else + Glibc.close(fd) + #endif + } + + private static func systemKill(_ pid: pid_t, _ signal: Int32) { + #if os(macOS) + _ = Darwin.kill(pid, signal) + #else + _ = Glibc.kill(pid, signal) + #endif + } + + private static func setCloseOnExec(_ fd: Int32) throws { + #if os(macOS) + let flags = Darwin.fcntl(fd, F_GETFD) + guard flags >= 0 else { + throw POSIXError(POSIXErrorCode(rawValue: errno) ?? .EIO) + } + guard Darwin.fcntl(fd, F_SETFD, flags | FD_CLOEXEC) == 0 else { + throw POSIXError(POSIXErrorCode(rawValue: errno) ?? .EIO) + } + #else + let flags = Glibc.fcntl(fd, F_GETFD) + guard flags >= 0 else { + throw POSIXError(POSIXErrorCode(rawValue: errno) ?? .EIO) + } + guard Glibc.fcntl(fd, F_SETFD, flags | FD_CLOEXEC) == 0 else { + throw POSIXError(POSIXErrorCode(rawValue: errno) ?? .EIO) + } + #endif + } + + private static func signalProcessGroup(rootPID: pid_t, _ signal: Int32) { + guard rootPID > 0 else { return } + // Primary: process-group signal (posix_spawn SETPGROUP). Skip assistive pgrep tree walk + // here — `kill(-pgid)` already covers the spawn group; nested pgrep only adds teardown latency. + systemKill(-rootPID, signal) + } + + // MARK: Abandoned-child reaper + + private static let abandonedLock = NSLock() + private static var abandonedPIDs: [pid_t] = [] + + private static func noteAbandoned(_ pid: pid_t) { + guard pid > 0 else { return } + abandonedLock.lock() + if !abandonedPIDs.contains(pid) { + abandonedPIDs.append(pid) + } + let count = abandonedPIDs.count + abandonedLock.unlock() + if count > maxAbandonedPIDsBeforeFail { + XCTFail( + "TestProcessRunner has \(count) abandoned unreaped child pids " + + "(threshold \(maxAbandonedPIDsBeforeFail)). " + + "Residual zombies last only for the suite process lifetime; " + + "investigate unkillable/D-state children or raise kill coverage." + ) + } + } + + private static func noteReaped(_ pid: pid_t) { + guard pid > 0 else { return } + abandonedLock.lock() + abandonedPIDs.removeAll { $0 == pid } + abandonedLock.unlock() + } + + /// Best-effort `waitpid(..., WNOHANG)` over previously abandoned pids. + /// Safe when a background waitpid already reaped the child (ECHILD) or is still waiting (0). + private static func reapAbandonedChildren() { + abandonedLock.lock() + let candidates = abandonedPIDs + abandonedLock.unlock() + guard !candidates.isEmpty else { return } + + var stillAbandoned: [pid_t] = [] + for pid in candidates { + var status: Int32 = 0 + let result = waitpid(pid, &status, WNOHANG) + if result == pid || (result < 0 && errno == ECHILD) { + continue + } + stillAbandoned.append(pid) + } + + abandonedLock.lock() + // Merge: keep any newly abandoned during reap, drop successfully reaped. + let newlyAbandoned = abandonedPIDs.filter { !candidates.contains($0) } + abandonedPIDs = stillAbandoned + newlyAbandoned + abandonedLock.unlock() + } + + #if os(macOS) + private static func signal(_ process: Process, _ signal: Int32) { + let pid = process.processIdentifier + guard pid > 0 else { return } + var remainingNodes = maxProcessTreeWalkNodes + signalProcessTree(rootPID: pid, signal, depth: 0, remainingNodes: &remainingNodes) + } + + /// Kill a single Process without walking its tree. Used for nested query helpers + /// (e.g. pgrep) so hang recovery cannot re-enter signalProcessTree. + private static func killProcessDirectly(_ process: Process, _ signal: Int32) { + let pid = process.processIdentifier + if pid > 0 { + systemKill(pid, signal) + } + if signal != SIGKILL, process.isRunning { + process.terminate() + } + } + + /// Best-effort recursive kill assist for Foundation `Process` trees. + /// Primary kill for the posix_spawn path is `kill(-pgid)`; this walk is assistive only. + /// Depth/node caps bound teardown latency; after the final grace wait the runner abandons + /// any still-unreaped waitpid threads (D-state / unkillable edge cases). + private static func signalProcessTree( + rootPID: pid_t, + _ signal: Int32, + depth: Int, + remainingNodes: inout Int + ) { + remainingNodes -= 1 + if depth < maxProcessTreeWalkDepth, remainingNodes > 0 { + for childPID in childPIDs(of: rootPID) { + guard remainingNodes > 0 else { break } + signalProcessTree( + rootPID: childPID, + signal, + depth: depth + 1, + remainingNodes: &remainingNodes + ) + } + } + _ = Darwin.kill(rootPID, signal) + } + + private static func childPIDs(of parentPID: pid_t) -> [pid_t] { + let process = Process() + process.executableURL = URL(fileURLWithPath: "/usr/bin/pgrep") + process.arguments = ["-P", "\(parentPID)"] + + let output = Pipe() + process.standardOutput = output + process.standardError = FileHandle.nullDevice + + let readerGroup = DispatchGroup() + let captured = LockedOutput() + let outputReader = output.fileHandleForReading + readerGroup.enter() + DispatchQueue.global(qos: .userInitiated).async { + defer { readerGroup.leave() } + while true { + guard let chunk = try? outputReader.read(upToCount: 64 * 1024), !chunk.isEmpty else { + break + } + captured.append(chunk) + } + } + + let terminationGroup = DispatchGroup() + terminationGroup.enter() + process.terminationHandler = { _ in + terminationGroup.leave() + } + + do { + try process.run() + } catch { + close(output.fileHandleForReading) + close(output.fileHandleForWriting) + readerGroup.wait() + return [] + } + close(output.fileHandleForWriting) + + if terminationGroup.wait(timeout: .now() + childPIDQueryTimeout) == .timedOut { + // Direct kill only — do not call terminate/forceTerminate (those re-enter + // signalProcessTree via signal(_:_:) and can amplify hangs with nested pgrep). + killProcessDirectly(process, SIGTERM) + if terminationGroup.wait(timeout: .now() + terminationGraceInterval) == .timedOut { + killProcessDirectly(process, SIGKILL) + _ = terminationGroup.wait(timeout: .now() + terminationGraceInterval) + } + finishReadingAfterTimeout(output.fileHandleForReading, readerGroup: readerGroup) + return [] + } + + _ = finishReadingAfterProcessExit(output.fileHandleForReading, readerGroup: readerGroup) + guard process.terminationStatus == 0 else { + return [] + } + return String(decoding: captured.data(), as: UTF8.self) + .split(whereSeparator: \.isNewline) + .compactMap { pid_t($0.trimmingCharacters(in: .whitespacesAndNewlines)) } + } + + #elseif os(Linux) + private static func signal(_ process: Process, _ signal: Int32) { + let pid = process.processIdentifier + guard pid > 0 else { return } + systemKill(pid, signal) + } + #endif + + private static func terminationStatus(fromWaitStatus status: Int32) -> Int32 { + let signal = status & 0x7F + if signal == 0 { + return (status >> 8) & 0xFF + } + return 128 + signal + } + #endif +} + +private final class LockedOutput { + private let lock = NSLock() + private var storage = Data() + + func append(_ data: Data) { + lock.lock() + storage.append(data) + lock.unlock() + } + + func data() -> Data { + lock.lock() + defer { lock.unlock() } + return storage + } +} + +private final class LockedWaitStatus { + enum Value { + case exited(Int32) + case failed(Int32) + } + + private let lock = NSLock() + private var storage: Value? + + var value: Value? { + lock.lock() + defer { lock.unlock() } + return storage + } + + func set(_ status: Value) { + lock.lock() + storage = status + lock.unlock() + } } diff --git a/Tests/RepoPromptTests/Helpers/TestProcessRunnerTests.swift b/Tests/RepoPromptTests/Helpers/TestProcessRunnerTests.swift index d8e2636c3..97e64932a 100644 --- a/Tests/RepoPromptTests/Helpers/TestProcessRunnerTests.swift +++ b/Tests/RepoPromptTests/Helpers/TestProcessRunnerTests.swift @@ -5,10 +5,80 @@ final class TestProcessRunnerTests: XCTestCase { func testDrainsLargeOutputWhileChildIsRunning() throws { let result = try TestProcessRunner.run( executableURL: URL(fileURLWithPath: "/usr/bin/head"), - arguments: ["-c", "65536", "/dev/zero"] + arguments: ["-c", "131072", "/dev/zero"] ) XCTAssertEqual(result.terminationStatus, 0) - XCTAssertEqual(result.output.count, 65536) + XCTAssertEqual(result.output.count, 131_072) + } + + func testTimeoutTerminatesProcessAndReportsContext() throws { + let cwd = FileManager.default.temporaryDirectory + .appendingPathComponent("TestProcessRunnerTests-\(UUID().uuidString)", isDirectory: true) + try FileManager.default.createDirectory(at: cwd, withIntermediateDirectories: true) + addTeardownBlock { + try? FileManager.default.removeItem(at: cwd) + } + + do { + _ = try TestProcessRunner.run( + executableURL: URL(fileURLWithPath: "/bin/sh"), + arguments: ["-c", "printf started; exec /bin/sleep 5"], + currentDirectoryURL: cwd, + timeout: 0.25 + ) + XCTFail("Expected process timeout") + } catch let error as TestProcessTimeoutError { + XCTAssertEqual(error.executableURL.path, "/bin/sh") + XCTAssertEqual(error.arguments, ["-c", "printf started; exec /bin/sleep 5"]) + XCTAssertEqual(error.currentDirectoryURL, cwd) + XCTAssertEqual(error.timeout, 0.25) + XCTAssertEqual(error.outputText, "started") + XCTAssertTrue(error.description.contains("cwd: \(cwd.path)")) + } + } + + func testTimeoutReturnsWhenChildProcessKeepsPipeOpen() throws { + let startedAt = Date() + + do { + _ = try TestProcessRunner.run( + executableURL: URL(fileURLWithPath: "/bin/sh"), + arguments: ["-c", "printf parent-started; sleep 5 & wait"], + timeout: 0.25 + ) + XCTFail("Expected process timeout") + } catch let error as TestProcessTimeoutError { + XCTAssertEqual(error.outputText, "parent-started") + XCTAssertLessThan(Date().timeIntervalSince(startedAt), 3) + } + } + + func testTimeoutReturnsWhenExitedParentLeavesChildHoldingPipe() throws { + let startedAt = Date() + + // Hold the pipe write end with a child that ignores common signals so platform + // SIGHUP/termination of the orphan cannot make drain complete before the grace budget. + // Accept either drain-timeout (parent exits 0) or process-timeout as long as wall time + // stays bounded and the output prefix matches. + do { + _ = try TestProcessRunner.run( + executableURL: URL(fileURLWithPath: "/bin/sh"), + arguments: [ + "-c", + "printf parent-exited; (trap '' HUP INT TERM; exec /bin/sleep 30) & exit 0" + ], + timeout: 0.25 + ) + XCTFail("Expected a bounded timeout while orphaned child holds the pipe") + } catch let error as TestProcessOutputDrainTimeoutError { + XCTAssertEqual(error.outputText, "parent-exited") + XCTAssertEqual(error.terminationStatus, 0) + XCTAssertTrue(error.description.contains("output drain timed out")) + XCTAssertLessThan(Date().timeIntervalSince(startedAt), 3) + } catch let error as TestProcessTimeoutError { + XCTAssertEqual(error.outputText, "parent-exited") + XCTAssertLessThan(Date().timeIntervalSince(startedAt), 3) + } } } diff --git a/Tests/RepoPromptTests/MCP/AgentRunWorktreeStartTests.swift b/Tests/RepoPromptTests/MCP/AgentRunWorktreeStartTests.swift index 74c406250..f2094e443 100644 --- a/Tests/RepoPromptTests/MCP/AgentRunWorktreeStartTests.swift +++ b/Tests/RepoPromptTests/MCP/AgentRunWorktreeStartTests.swift @@ -872,28 +872,25 @@ final class AgentRunWorktreeStartTests: AgentRunWorktreeStartGitSeedTestCase { var composeTab = try XCTUnwrap(window.workspaceManager.composeTab(with: sourceTabID)) composeTab.activeAgentSessionID = parentSessionID window.workspaceManager.updateComposeTab(composeTab, markDirty: false) - _ = await window.selectionCoordinator.persistSelection( - storedSelection, - for: sourceIdentity, - source: .runtimeMutation, - mirrorToUIIfActive: false - ) // Wait for post-switch git-data root load to complete so any selection revision - // published by the async git-data load settles before we diverge the UI selection. + // published by the async git-data load settles before capture. await window.workspaceManager.waitUntilPostSwitchGitDataLoadComplete() - await window.selectionCoordinator.withApplyingSelectionMirror { - await window.workspaceFilesViewModel.applyStoredSelection(StoredSelection()) - } + await window.promptManager.createBlankComposeTab(createAgentSession: false) + let activeTabID = try XCTUnwrap(window.workspaceManager.activeWorkspace?.activeComposeTabID) + XCTAssertNotEqual(activeTabID, sourceTabID) let divergentUISelection = window.workspaceFilesViewModel.snapshotSelection() XCTAssertTrue(divergentUISelection.selectedPaths.isEmpty) XCTAssertNotEqual(divergentUISelection, storedSelection) + _ = await window.selectionCoordinator.persistSelection( storedSelection, for: sourceIdentity, source: .runtimeMutation, mirrorToUIIfActive: false ) - XCTAssertEqual(window.workspaceManager.composeTab(for: sourceIdentity)?.selection, storedSelection) + try await AsyncTestWait.waitUntil("source review selection remains stored", timeout: 5) { + window.workspaceManager.composeTab(for: sourceIdentity)?.selection == storedSelection + } let launchSnapshot = AgentRunOracleReviewLaunchSnapshot( route: .explicitTabContext, diff --git a/Tests/RepoPromptTests/MCP/Control/MCPSharedServerTestLease.swift b/Tests/RepoPromptTests/MCP/Control/MCPSharedServerTestLease.swift index 9b2c9b24b..ef5b16eec 100644 --- a/Tests/RepoPromptTests/MCP/Control/MCPSharedServerTestLease.swift +++ b/Tests/RepoPromptTests/MCP/Control/MCPSharedServerTestLease.swift @@ -1,3 +1,5 @@ +import Foundation + #if DEBUG import Foundation @@ -6,17 +8,12 @@ fileprivate init() {} } - private struct Waiter { - let id: UUID - let continuation: CheckedContinuation - } - static let shared = MCPSharedServerTestLease() private var occupied = false - private var waiters: [Waiter] = [] - private var cancelledWaiterIDs: Set = [] - private var grantedWaiterIDs: Set = [] + /// Lock-backed waiter queue so `onCancel` can remove/resume waiters **synchronously** + /// without an unstructured `Task { await }` hop. + private let waiterState = LeaseWaiterState() func withLease(_ operation: (Ownership) async throws -> T) async throws -> T { var ownsLease = false @@ -31,6 +28,10 @@ return try await operation(Ownership()) } + func waiterCountForTesting() -> Int { + waiterState.waiterCount + } + private func acquireLease() async throws { guard occupied else { occupied = true @@ -43,16 +44,7 @@ return } - let waiterID = UUID() - try await withTaskCancellationHandler { - try await withCheckedThrowingContinuation { continuation in - registerWaiter(id: waiterID, continuation: continuation) - } - } onCancel: { - Task { await Self.shared.cancelWaiter(id: waiterID) } - } - - grantedWaiterIDs.remove(waiterID) + try await waitForTurn() do { try Task.checkCancellation() } catch { @@ -61,31 +53,93 @@ } } - private func registerWaiter(id: UUID, continuation: CheckedContinuation) { - if cancelledWaiterIDs.remove(id) != nil { - continuation.resume(throwing: CancellationError()) - } else { - waiters.append(Waiter(id: id, continuation: continuation)) + private func waitForTurn() async throws { + let waiterID = waiterState.allocateWaiterID() + try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { continuation in + waiterState.enqueue(id: waiterID, continuation: continuation) + } + } onCancel: { + // Synchronous sticky cancel — must not hop through the actor executor. + waiterState.cancel(id: waiterID) } } private func releaseLease() { - if waiters.isEmpty { - occupied = false - } else { - let waiter = waiters.removeFirst() - grantedWaiterIDs.insert(waiter.id) - waiter.continuation.resume() + if let continuation = waiterState.dequeueNextReady() { + continuation.resume() + return } + occupied = false + } + } + + /// Shared lease waiter queue. `@unchecked Sendable` lock state so cancellation handlers + /// can run synchronously from any executor. + private final class LeaseWaiterState: @unchecked Sendable { + private let lock = NSLock() + private var nextWaiterID = 0 + private var pendingWaiterIDs: Set = [] + private var cancelledWaiterIDs: Set = [] + private var waiters: [(id: Int, continuation: CheckedContinuation)] = [] + + var waiterCount: Int { + lock.withLock { waiters.count } + } + + func allocateWaiterID() -> Int { + lock.lock() + let id = nextWaiterID + nextWaiterID += 1 + pendingWaiterIDs.insert(id) + lock.unlock() + return id } - private func cancelWaiter(id: UUID) { + func enqueue(id: Int, continuation: CheckedContinuation) { + lock.lock() + if cancelledWaiterIDs.remove(id) != nil { + pendingWaiterIDs.remove(id) + lock.unlock() + continuation.resume(throwing: CancellationError()) + return + } + waiters.append((id, continuation)) + lock.unlock() + } + + func cancel(id: Int) { + lock.lock() if let index = waiters.firstIndex(where: { $0.id == id }) { let waiter = waiters.remove(at: index) + pendingWaiterIDs.remove(id) + lock.unlock() waiter.continuation.resume(throwing: CancellationError()) - } else if !grantedWaiterIDs.contains(id) { + return + } + if pendingWaiterIDs.contains(id) { cancelledWaiterIDs.insert(id) } + lock.unlock() + } + + /// Returns the next non-cancelled waiter's continuation, or nil if the queue is empty. + func dequeueNextReady() -> CheckedContinuation? { + lock.lock() + while !waiters.isEmpty { + let waiter = waiters.removeFirst() + pendingWaiterIDs.remove(waiter.id) + if cancelledWaiterIDs.remove(waiter.id) != nil { + lock.unlock() + waiter.continuation.resume(throwing: CancellationError()) + lock.lock() + continue + } + lock.unlock() + return waiter.continuation + } + lock.unlock() + return nil } } #endif diff --git a/Tests/RepoPromptTests/MCP/Control/MCPToolExecutionWatchdogIntegrationTests.swift b/Tests/RepoPromptTests/MCP/Control/MCPToolExecutionWatchdogIntegrationTests.swift index 2a2b46803..470740d6e 100644 --- a/Tests/RepoPromptTests/MCP/Control/MCPToolExecutionWatchdogIntegrationTests.swift +++ b/Tests/RepoPromptTests/MCP/Control/MCPToolExecutionWatchdogIntegrationTests.swift @@ -1943,64 +1943,50 @@ import XCTest } } + /// Cooperative cancel probe: thin wrapper over shared `TestCancellationGate`. private actor MCPExecutionCooperativeCancellationGate { private static let synchronizationTimeout: Duration = .seconds(10) - private var entered = false - private var cancellationCount = 0 - private var continuation: CheckedContinuation? + private let gate = TestCancellationGate(name: "MCP execution cooperative cancellation gate") func enterAndWait() async throws { - entered = true - try await withTaskCancellationHandler { - try await withCheckedThrowingContinuation { continuation in - self.continuation = continuation - } - } onCancel: { - Task { await self.cancel() } - } + try await gate.waitUntilCancelled() } func waitUntilEntered( timeout: Duration = synchronizationTimeout ) async throws { - let clock = ContinuousClock() - let deadline = clock.now.advanced(by: timeout) - while !entered { - try Task.checkCancellation() - guard clock.now < deadline else { - throw MCPExecutionWatchdogIntegrationFixtureError.cooperativeGateDidNotEnter - } - try await Task.sleep(for: .milliseconds(10)) + let entered = await gate.waitUntilEntered( + timeout: TestFenceDefaults.timeInterval(timeout), + failOnTimeout: false + ) + guard entered else { + throw MCPExecutionWatchdogIntegrationFixtureError.cooperativeGateDidNotEnter } } func waitUntilCancellationObserved( timeout: Duration = synchronizationTimeout ) async throws { - let clock = ContinuousClock() - let deadline = clock.now.advanced(by: timeout) - while cancellationCount == 0 { - try Task.checkCancellation() - guard clock.now < deadline else { - throw MCPExecutionWatchdogIntegrationFixtureError.cooperativeGateCancellationNotObserved + let timeoutInterval = TestFenceDefaults.timeInterval(timeout) + do { + try await AsyncTestWait.waitUntil( + "MCP cooperative gate cancellation observed", + timeout: timeoutInterval + ) { + self.gate.cancellationCount > 0 } - try await Task.sleep(for: .milliseconds(10)) + } catch { + throw MCPExecutionWatchdogIntegrationFixtureError.cooperativeGateCancellationNotObserved } } - func observedCancellationCount() -> Int { - cancellationCount - } - - func cancelForCleanup() { - cancel() + func observedCancellationCount() async -> Int { + gate.cancellationCount } - private func cancel() { - cancellationCount += 1 - continuation?.resume(throwing: CancellationError()) - continuation = nil + func cancelForCleanup() async { + gate.forceCancel() } } diff --git a/Tests/RepoPromptTests/MCP/Control/MCPToolExecutionWatchdogTests.swift b/Tests/RepoPromptTests/MCP/Control/MCPToolExecutionWatchdogTests.swift index 01f233da7..59c060073 100644 --- a/Tests/RepoPromptTests/MCP/Control/MCPToolExecutionWatchdogTests.swift +++ b/Tests/RepoPromptTests/MCP/Control/MCPToolExecutionWatchdogTests.swift @@ -25,6 +25,7 @@ final class MCPToolExecutionWatchdogTests: XCTestCase { func testDeadlineCancelsCooperativeOperationAndReturnsSingleTimeout() async throws { let clock = ExecutionWatchdogManualClock() let events = ExecutionWatchdogEventRecorder() + let operationGate = ExecutionWatchdogCancellationGate() let task = Task { try await MCPToolExecutionWatchdog.execute( deadline: MCPTimeoutPolicy.boundedToolExecutionDeadline, @@ -32,7 +33,7 @@ final class MCPToolExecutionWatchdogTests: XCTestCase { environment: clock.environment, onEvent: { await events.append($0) } ) { - try await Task.sleep(for: .seconds(3600)) + try await operationGate.waitUntilCancelled() return 1 } } @@ -58,6 +59,7 @@ final class MCPToolExecutionWatchdogTests: XCTestCase { let clock = ExecutionWatchdogManualClock() let events = ExecutionWatchdogEventRecorder() let callbackGate = ExecutionWatchdogCallbackGate() + let operationGate = ExecutionWatchdogCancellationGate() let task = Task { try await MCPToolExecutionWatchdog.execute( deadline: MCPTimeoutPolicy.boundedToolExecutionDeadline, @@ -70,7 +72,7 @@ final class MCPToolExecutionWatchdogTests: XCTestCase { } } ) { - try await Task.sleep(for: .seconds(3600)) + try await operationGate.waitUntilCancelled() return 1 } } @@ -283,17 +285,14 @@ actor ExecutionWatchdogUncooperativeGate { } } +typealias ExecutionWatchdogCancellationGate = TestCancellationGate + actor ExecutionWatchdogManualClock { private static let synchronizationTimeout: Duration = .seconds(10) - private struct Sleeper { - let duration: Duration - let continuation: CheckedContinuation - } - private var elapsed: Duration = .zero - private var sleeperOrder: [UUID] = [] - private var sleepers: [UUID: Sleeper] = [:] + /// Lock-backed sleeper registry so `onCancel` can remove/resume synchronously. + private let sleeperState = ManualClockSleeperState() nonisolated var environment: MCPToolExecutionWatchdogEnvironment { MCPToolExecutionWatchdogEnvironment( @@ -311,16 +310,15 @@ actor ExecutionWatchdogManualClock { let id = UUID() try await withTaskCancellationHandler { try await withCheckedThrowingContinuation { continuation in - sleeperOrder.append(id) - sleepers[id] = Sleeper(duration: duration, continuation: continuation) + self.sleeperState.register(id: id, duration: duration, continuation: continuation) } } onCancel: { - Task { await self.cancelSleeper(id) } + sleeperState.cancel(id: id) } } func sleeperCount() -> Int { - sleepers.count + sleeperState.count } func waitForSleeperCount( @@ -329,10 +327,10 @@ actor ExecutionWatchdogManualClock { ) async throws { let clock = ContinuousClock() let deadline = clock.now.advanced(by: timeout) - while sleepers.count < count { + while sleeperState.count < count { try Task.checkCancellation() guard clock.now < deadline else { - throw ManualClockError.sleeperDidNotRegister(expected: count, actual: sleepers.count) + throw ManualClockError.sleeperDidNotRegister(expected: count, actual: sleeperState.count) } try await Task.sleep(for: .milliseconds(10)) } @@ -342,18 +340,14 @@ actor ExecutionWatchdogManualClock { guard duration > .zero else { throw ManualClockError.nonPositiveAdvance(duration) } - guard sleepers.isEmpty else { - throw ManualClockError.sleepersRegistered(sleepers.count) + guard sleeperState.count == 0 else { + throw ManualClockError.sleepersRegistered(sleeperState.count) } elapsed += duration } func advanceNext(expected: Duration) throws { - guard let id = sleeperOrder.first else { - throw ManualClockError.noSleeper - } - sleeperOrder.removeFirst() - guard let sleeper = sleepers.removeValue(forKey: id) else { + guard let sleeper = sleeperState.popNext() else { throw ManualClockError.noSleeper } guard sleeper.duration == expected else { @@ -363,11 +357,6 @@ actor ExecutionWatchdogManualClock { sleeper.continuation.resume() } - private func cancelSleeper(_ id: UUID) { - sleeperOrder.removeAll { $0 == id } - sleepers.removeValue(forKey: id)?.continuation.resume(throwing: CancellationError()) - } - private enum ManualClockError: Error { case noSleeper case nonPositiveAdvance(Duration) @@ -376,3 +365,56 @@ actor ExecutionWatchdogManualClock { case unexpectedDuration(expected: Duration, actual: Duration) } } + +/// Sync-cancelable sleeper registry for the manual watchdog clock. +private final class ManualClockSleeperState: @unchecked Sendable { + private struct Sleeper { + let duration: Duration + let continuation: CheckedContinuation + } + + private let lock = NSLock() + private var sleeperOrder: [UUID] = [] + private var sleepers: [UUID: Sleeper] = [:] + private var cancelledIDs = Set() + + var count: Int { + lock.withLock { sleepers.count } + } + + func register( + id: UUID, + duration: Duration, + continuation: CheckedContinuation + ) { + lock.lock() + if cancelledIDs.remove(id) != nil { + lock.unlock() + continuation.resume(throwing: CancellationError()) + return + } + sleeperOrder.append(id) + sleepers[id] = Sleeper(duration: duration, continuation: continuation) + lock.unlock() + } + + func cancel(id: UUID) { + lock.lock() + sleeperOrder.removeAll { $0 == id } + let sleeper = sleepers.removeValue(forKey: id) + if sleeper == nil { + cancelledIDs.insert(id) + } + lock.unlock() + sleeper?.continuation.resume(throwing: CancellationError()) + } + + func popNext() -> (duration: Duration, continuation: CheckedContinuation)? { + lock.lock() + defer { lock.unlock() } + guard let id = sleeperOrder.first else { return nil } + sleeperOrder.removeFirst() + guard let sleeper = sleepers.removeValue(forKey: id) else { return nil } + return (sleeper.duration, sleeper.continuation) + } +} diff --git a/Tests/RepoPromptTests/MCP/Control/PersistentAgentModeMCPReadFileConnectionTests.swift b/Tests/RepoPromptTests/MCP/Control/PersistentAgentModeMCPReadFileConnectionTests.swift index 87d4bf161..ee7fc7e69 100644 --- a/Tests/RepoPromptTests/MCP/Control/PersistentAgentModeMCPReadFileConnectionTests.swift +++ b/Tests/RepoPromptTests/MCP/Control/PersistentAgentModeMCPReadFileConnectionTests.swift @@ -4059,6 +4059,7 @@ final class PersistentAgentModeMCPReadFileConnectionTests: XCTestCase { private var fd: Int32 private var buffer = Data() private var nonMatchingFrames: [String] = [] + private let responseTimeout: TimeInterval = 10 init(fd: Int32) { self.fd = fd @@ -4103,45 +4104,58 @@ final class PersistentAgentModeMCPReadFileConnectionTests: XCTestCase { } private func response(matching expectedID: Int) async throws -> String { - let waiterHolder = SocketPairResponseWaiterHolder() + let waitState = SocketResponseWaitState() return try await withTaskCancellationHandler { try await withCheckedThrowingContinuation { continuation in - let responseWaiter = SocketPairResponseWaiter(continuation) - waiterHolder.install(responseWaiter) + guard waitState.install(continuation) else { return } + let deadline = Date().addingTimeInterval(responseTimeout) queue.async { + let result: Result do { - while true { - let line = try self.readLine() - let object = try JSONSerialization.jsonObject(with: line) as? [String: Any] - guard let object else { throw ClientError.invalidResponse } - if let rawID = object["id"] { - guard let responseID = (rawID as? NSNumber)?.intValue else { - throw ClientError.invalidResponse - } - guard responseID == expectedID else { - throw ClientError.invalidResponse - } - guard let rawJSON = String(data: line, encoding: .utf8) else { - throw ClientError.invalidResponse - } - responseWaiter.resume(returning: rawJSON) - return - } - guard object["method"] as? String != nil, - let rawJSON = String(data: line, encoding: .utf8) - else { - throw ClientError.invalidResponse - } - self.nonMatchingFrames.append(rawJSON) - } + result = try .success(self.readResponse( + matching: expectedID, + deadline: deadline, + isCancelled: { waitState.isCancelled } + )) } catch { - responseWaiter.resume(throwing: error) + result = .failure(error) } + waitState.resume(with: result) } } } onCancel: { - waiterHolder.resume(throwing: CancellationError()) - close() + waitState.cancel() + } + } + + private func readResponse( + matching expectedID: Int, + deadline: Date, + isCancelled: () -> Bool + ) throws -> String { + while true { + if isCancelled() { throw CancellationError() } + let line = try readLine(deadline: deadline, isCancelled: isCancelled) + let object = try JSONSerialization.jsonObject(with: line) as? [String: Any] + guard let object else { throw ClientError.invalidResponse } + if let rawID = object["id"] { + guard let responseID = (rawID as? NSNumber)?.intValue else { + throw ClientError.invalidResponse + } + guard responseID == expectedID else { + throw ClientError.invalidResponse + } + guard let rawJSON = String(data: line, encoding: .utf8) else { + throw ClientError.invalidResponse + } + return rawJSON + } + guard object["method"] as? String != nil, + let rawJSON = String(data: line, encoding: .utf8) + else { + throw ClientError.invalidResponse + } + nonMatchingFrames.append(rawJSON) } } @@ -4161,8 +4175,9 @@ final class PersistentAgentModeMCPReadFileConnectionTests: XCTestCase { } } - private func readLine() throws -> Data { + private func readLine(deadline: Date, isCancelled: () -> Bool) throws -> Data { while true { + if isCancelled() { throw CancellationError() } if let newline = buffer.firstIndex(of: 0x0A) { let line = Data(buffer[..= 0 else { throw ClientError.closed } var descriptor = pollfd(fd: fd, events: Int16(POLLIN), revents: 0) - let pollResult = Darwin.poll(&descriptor, 1, 10000) - if pollResult == 0 { throw ClientError.timedOut } + let remainingMilliseconds = Int32(deadline.timeIntervalSinceNow * 1000) + if remainingMilliseconds <= 0 { throw ClientError.timedOut } + let pollResult = Darwin.poll(&descriptor, 1, min(100, remainingMilliseconds)) + if pollResult == 0 { continue } if pollResult < 0 { if errno == EINTR { continue } throw ClientError.posix(operation: "poll", code: errno) @@ -4196,4 +4213,54 @@ final class PersistentAgentModeMCPReadFileConnectionTests: XCTestCase { } } } + + private final class SocketResponseWaitState: @unchecked Sendable { + private let lock = NSLock() + private var continuation: CheckedContinuation? + private var cancelled = false + private var completed = false + + var isCancelled: Bool { + lock.lock() + let result = cancelled + lock.unlock() + return result + } + + func install(_ continuation: CheckedContinuation) -> Bool { + lock.lock() + if cancelled || completed { + lock.unlock() + continuation.resume(throwing: CancellationError()) + return false + } + self.continuation = continuation + lock.unlock() + return true + } + + func cancel() { + let continuation = takeContinuation(cancelled: true) + continuation?.resume(throwing: CancellationError()) + } + + func resume(with result: Result) { + guard let continuation = takeContinuation(cancelled: false) else { return } + continuation.resume(with: result) + } + + private func takeContinuation(cancelled: Bool) -> CheckedContinuation? { + lock.lock() + if completed { + lock.unlock() + return nil + } + self.cancelled = self.cancelled || cancelled + completed = true + let continuation = continuation + self.continuation = nil + lock.unlock() + return continuation + } + } #endif diff --git a/Tests/RepoPromptTests/MCP/Control/PersistentMCPDistinctConnectionConcurrencyTests.swift b/Tests/RepoPromptTests/MCP/Control/PersistentMCPDistinctConnectionConcurrencyTests.swift index c97b0d7f9..3ad9fadc3 100644 --- a/Tests/RepoPromptTests/MCP/Control/PersistentMCPDistinctConnectionConcurrencyTests.swift +++ b/Tests/RepoPromptTests/MCP/Control/PersistentMCPDistinctConnectionConcurrencyTests.swift @@ -6,6 +6,38 @@ import XCTest @MainActor final class PersistentMCPDistinctConnectionConcurrencyTests: XCTestCase { + func testSharedServerLeaseCancellationRemovesQueuedWaiter() async throws { + #if DEBUG + try await MCPSharedServerTestLease.shared.withLease { _ in + let queuedWaiter = Task { + try await MCPSharedServerTestLease.shared.withLease { _ in + XCTFail("Cancelled lease waiter should not acquire the shared server lease.") + } + } + + let queued = await waitForSharedServerLeaseWaiterCount(1, timeout: .seconds(2)) + XCTAssertTrue(queued) + queuedWaiter.cancel() + + do { + try await queuedWaiter.value + XCTFail("Expected cancelled lease waiter to throw CancellationError.") + } catch is CancellationError { + // Expected cancellation path. + } + + let waiterCount = await MCPSharedServerTestLease.shared.waiterCountForTesting() + XCTAssertEqual(waiterCount, 0) + } + + try await withSharedServerLeaseTimeout(timeout: .seconds(2)) { + try await MCPSharedServerTestLease.shared.withLease { _ in } + } + #else + throw XCTSkip("Shared MCP server lease cancellation regression requires DEBUG diagnostics helpers.") + #endif + } + func testDistinctConnectionsOverlapWithoutCrossRoutingReadOrSearchResults() async throws { #if DEBUG try await MCPSharedServerTestLease.shared.withLease { lease in @@ -106,6 +138,52 @@ final class PersistentMCPDistinctConnectionConcurrencyTests: XCTestCase { case full } + func waitForSharedServerLeaseWaiterCount( + _ expectedCount: Int, + timeout: Duration + ) async -> Bool { + let clock = ContinuousClock() + let deadline = clock.now + timeout + while clock.now < deadline { + let waiterCount = await MCPSharedServerTestLease.shared.waiterCountForTesting() + if waiterCount == expectedCount { return true } + try? await Task.sleep(for: .milliseconds(10)) + } + return false + } + + func withSharedServerLeaseTimeout( + timeout: Duration, + operation: @escaping @Sendable () async throws -> T + ) async throws -> T { + let timeoutInterval = { + let components = timeout.components + return TimeInterval(components.seconds) + + (TimeInterval(components.attoseconds) / 1e18) + }() + return try await withThrowingTaskGroup(of: T.self) { group in + group.addTask { + try await operation() + } + group.addTask { + try await Task.sleep(for: timeout) + throw AsyncTestConditionTimeout( + description: "shared MCP server lease re-acquisition", + timeout: timeoutInterval + ) + } + guard let result = try await group.next() else { + group.cancelAll() + throw AsyncTestConditionTimeout( + description: "shared MCP server lease re-acquisition (empty task group)", + timeout: timeoutInterval + ) + } + group.cancelAll() + return result + } + } + func runReadSelectionPersistenceCheckpoint( fixture: PersistentMCPTestFixture, shape: ReadSelectionShape @@ -2002,7 +2080,10 @@ final class PersistentMCPDistinctConnectionConcurrencyTests: XCTestCase { guard activeFD >= 0 else { return } var descriptor = pollfd(fd: activeFD, events: Int16(POLLIN), revents: 0) let pollResult = Darwin.poll(&descriptor, 1, 100) - if pollResult == 0 { continue } + if pollResult == 0 { + if withStateLock({ isClosed }) { return } + continue + } if pollResult < 0 { if errno == EINTR { continue } if withStateLock({ isClosed }) { return } @@ -2015,6 +2096,7 @@ final class PersistentMCPDistinctConnectionConcurrencyTests: XCTestCase { close(with: ClientError.closed) return } + if withStateLock({ isClosed }) { return } var bytes = [UInt8](repeating: 0, count: 4096) let readCount = bytes.withUnsafeMutableBytes { storage in Darwin.read(activeFD, storage.baseAddress, storage.count) diff --git a/Tests/RepoPromptTests/MCP/Control/PersistentMCPResponseDeliveryTests.swift b/Tests/RepoPromptTests/MCP/Control/PersistentMCPResponseDeliveryTests.swift index 8f01b063e..c5167bdea 100644 --- a/Tests/RepoPromptTests/MCP/Control/PersistentMCPResponseDeliveryTests.swift +++ b/Tests/RepoPromptTests/MCP/Control/PersistentMCPResponseDeliveryTests.swift @@ -321,7 +321,7 @@ final class PersistentMCPResponseDeliveryTests: XCTestCase { Self.closeIfOpen(stdinPipe[1]) stdinPipe[1] = -1 let forwardedRequest = try await Task.detached { - try Self.readLine(from: sockets[1]) + try Self.readLine(from: sockets[1], timeout: 2) }.value XCTAssertEqual(forwardedRequest, requestFrame) let observedStdinClose = await poller.waitUntilStdinClosed() @@ -342,7 +342,7 @@ final class PersistentMCPResponseDeliveryTests: XCTestCase { await poller.resumeNext(.events(Int16(POLLIN))) try await bridgeTask.value - let deliveredResponse = try Self.readLine(from: stdoutPipe[0]) + let deliveredResponse = try Self.readLine(from: stdoutPipe[0], timeout: 2) XCTAssertEqual(deliveredResponse, responseFrame) let snapshot = await ledger.snapshot() XCTAssertEqual(snapshot.activeRequestCount, 0) @@ -395,7 +395,7 @@ final class PersistentMCPResponseDeliveryTests: XCTestCase { Self.closeIfOpen(stdinPipe[1]) stdinPipe[1] = -1 let forwardedRequest = try await Task.detached { - try Self.readLine(from: sockets[1]) + try Self.readLine(from: sockets[1], timeout: 2) }.value XCTAssertEqual(forwardedRequest, requestFrame) let observedStdinClose = await poller.waitUntilStdinClosed() @@ -1593,6 +1593,7 @@ private actor ManualBridgeSocketPoller { } } } onCancel: { + // Actor isolation requires a hop; sticky pendingPollCancelled covers cancel-before-register. Task { await self.cancelPendingPoll() } } } @@ -1988,31 +1989,20 @@ private actor WriteRecorder { } } -private actor AsyncGate { - private var entered = false - private var released = false - private var enteredWaiters: [CheckedContinuation] = [] - private var releaseWaiters: [CheckedContinuation] = [] +/// Response-delivery concurrency fence (shared `TestReleaseFence` with legacy names). +private final class AsyncGate: @unchecked Sendable { + private let fence = TestReleaseFence(name: "persistent MCP response delivery async gate") func markEnteredAndWait() async { - entered = true - let waiters = enteredWaiters - enteredWaiters.removeAll() - waiters.forEach { $0.resume() } - guard !released else { return } - await withCheckedContinuation { releaseWaiters.append($0) } + await fence.enterAndWait() } - func waitUntilEntered() async { - guard !entered else { return } - await withCheckedContinuation { enteredWaiters.append($0) } + func waitUntilEntered(timeout: TimeInterval = TestFenceDefaults.enterWait) async { + _ = await fence.waitUntilEntered(timeout: timeout) } func release() { - released = true - let waiters = releaseWaiters - releaseWaiters.removeAll() - waiters.forEach { $0.resume() } + fence.release() } } diff --git a/Tests/RepoPromptTests/MCP/MCPCodeStructureWorktreeTests.swift b/Tests/RepoPromptTests/MCP/MCPCodeStructureWorktreeTests.swift index a5230afbb..55d3c2960 100644 --- a/Tests/RepoPromptTests/MCP/MCPCodeStructureWorktreeTests.swift +++ b/Tests/RepoPromptTests/MCP/MCPCodeStructureWorktreeTests.swift @@ -640,6 +640,14 @@ final class MCPCodeStructureWorktreeTests: XCTestCase { var composeTab = try XCTUnwrap(window.workspaceManager.composeTab(with: tabID)) composeTab.selection = StoredSelection(selectedPaths: files.map(\.standardizedFullPath)) window.workspaceManager.updateComposeTab(composeTab, markDirty: false) + try await AsyncTestWait.waitUntil("selected code structure candidates are cataloged", timeout: 5) { + let resolution = await store.resolveSelectedCodeStructureFiles( + atPaths: files.map(\.standardizedFullPath), + rootScope: .visibleWorkspace, + maximumUniqueFileCount: 1 + ) + return resolution.didExceedLimit && resolution.visitedUniqueFileCount == 2 + } let connectionID = UUID() try window.mcpServer.bindTabForConnection( diff --git a/Tests/RepoPromptTests/Services/FileSystem/FileSystemContentLoadingConcurrencyTests.swift b/Tests/RepoPromptTests/Services/FileSystem/FileSystemContentLoadingConcurrencyTests.swift index b797b2887..e5a9dc37c 100644 --- a/Tests/RepoPromptTests/Services/FileSystem/FileSystemContentLoadingConcurrencyTests.swift +++ b/Tests/RepoPromptTests/Services/FileSystem/FileSystemContentLoadingConcurrencyTests.swift @@ -1047,10 +1047,11 @@ final class FileSystemContentLoadingConcurrencyTests: XCTestCase { XCTAssertEqual(afterError.foregroundActivityCount, 0) let started = AsyncSignal() + let cancellationGate = AsyncCancellationGate() let cancelled = Task { try await limiter.withForegroundActivity(kind: .interactiveRead) { await started.mark() - try await Task.sleep(for: .seconds(60)) + try await cancellationGate.waitUntilCancelled() } } let didStart = await started.waitUntilMarked() @@ -1303,33 +1304,20 @@ private enum ForegroundActivityTestError: Error { case expected } -private actor AsyncGate { - private var started = false - private var released = false - private var startWaiters: [CheckedContinuation] = [] - private var releaseWaiters: [CheckedContinuation] = [] +/// Content-loading concurrency fence (shared `TestReleaseFence` with legacy names). +private final class AsyncGate: @unchecked Sendable { + private let fence = TestReleaseFence(name: "file system content loading async gate") func markStartedAndWaitForRelease() async { - started = true - startWaiters.forEach { $0.resume() } - startWaiters.removeAll() - guard !released else { return } - await withCheckedContinuation { continuation in - releaseWaiters.append(continuation) - } + await fence.enterAndWaitIgnoringCancellationUntilRelease() } - func waitUntilStarted() async { - guard !started else { return } - await withCheckedContinuation { continuation in - startWaiters.append(continuation) - } + func waitUntilStarted(timeout: TimeInterval = TestFenceDefaults.enterWait) async { + _ = await fence.waitUntilEntered(timeout: timeout) } func release() { - released = true - releaseWaiters.forEach { $0.resume() } - releaseWaiters.removeAll() + fence.release() } } @@ -1341,6 +1329,7 @@ private actor AsyncCounter { private var count = 0 private var waiters: [UUID: Waiter] = [:] + private var cancelledWaiterIDs: Set = [] func incrementAndValue() -> Int { count += 1 @@ -1355,9 +1344,10 @@ private actor AsyncCounter { func waitUntilValue(atLeast target: Int, timeoutNanoseconds: UInt64 = 1_000_000_000) async -> Bool { guard count < target else { return true } let waiterID = UUID() + // Sticky cancel handled via cancelledWaiterIDs so cancel-before-register cannot hang. return await withTaskCancellationHandler { await withCheckedContinuation { continuation in - guard count < target, !Task.isCancelled else { + guard count < target, !Task.isCancelled, cancelledWaiterIDs.remove(waiterID) == nil else { continuation.resume(returning: count >= target) return } @@ -1368,7 +1358,7 @@ private actor AsyncCounter { } } } onCancel: { - Task { await self.finishWaiter(waiterID) } + Task { await self.finishWaiter(waiterID, fromCancel: true) } } } @@ -1381,9 +1371,14 @@ private actor AsyncCounter { } } - private func finishWaiter(_ waiterID: UUID) { - guard let waiter = waiters.removeValue(forKey: waiterID) else { return } - waiter.continuation.resume(returning: count >= waiter.target) + private func finishWaiter(_ waiterID: UUID, fromCancel: Bool = false) { + if let waiter = waiters.removeValue(forKey: waiterID) { + waiter.continuation.resume(returning: count >= waiter.target) + return + } + if fromCancel { + cancelledWaiterIDs.insert(waiterID) + } } } @@ -1402,6 +1397,7 @@ private actor AsyncValueRecorder { private actor AsyncSignal { private var marked = false private var waiters: [UUID: CheckedContinuation] = [:] + private var cancelledWaiterIDs: Set = [] func mark() { guard !marked else { return } @@ -1420,7 +1416,7 @@ private actor AsyncSignal { let waiterID = UUID() return await withTaskCancellationHandler { await withCheckedContinuation { continuation in - guard !marked, !Task.isCancelled else { + guard !marked, !Task.isCancelled, cancelledWaiterIDs.remove(waiterID) == nil else { continuation.resume(returning: marked) return } @@ -1431,12 +1427,19 @@ private actor AsyncSignal { } } } onCancel: { - Task { await self.finishWaiter(waiterID) } + Task { await self.finishWaiter(waiterID, fromCancel: true) } } } - private func finishWaiter(_ waiterID: UUID) { - guard let continuation = waiters.removeValue(forKey: waiterID) else { return } - continuation.resume(returning: marked) + private func finishWaiter(_ waiterID: UUID, fromCancel: Bool = false) { + if let continuation = waiters.removeValue(forKey: waiterID) { + continuation.resume(returning: marked) + return + } + if fromCancel { + cancelledWaiterIDs.insert(waiterID) + } } } + +private typealias AsyncCancellationGate = TestCancellationGate diff --git a/Tests/RepoPromptTests/Services/VCS/GitBlobSourceMaterializationServiceTests.swift b/Tests/RepoPromptTests/Services/VCS/GitBlobSourceMaterializationServiceTests.swift index 0755707b4..9c78cd3bb 100644 --- a/Tests/RepoPromptTests/Services/VCS/GitBlobSourceMaterializationServiceTests.swift +++ b/Tests/RepoPromptTests/Services/VCS/GitBlobSourceMaterializationServiceTests.swift @@ -261,8 +261,7 @@ final class GitBlobSourceMaterializationServiceTests: XCTestCase { client: GitBlobSourceMaterializationClient( size: { _, _ in UInt64(bytes.count) }, bytes: { _, _, _ in - await gate.enter() - try await Task.sleep(for: .seconds(60)) + try await gate.waitUntilCancelled() return bytes } ) @@ -361,22 +360,7 @@ final class GitBlobSourceMaterializationServiceTests: XCTestCase { } } -private actor MaterializationCancellationGate { - private var entered = false - private var waiters: [CheckedContinuation] = [] - - func enter() { - entered = true - let current = waiters - waiters.removeAll() - current.forEach { $0.resume() } - } - - func waitUntilEntered() async { - if entered { return } - await withCheckedContinuation { waiters.append($0) } - } -} +private typealias MaterializationCancellationGate = TestCancellationGate private final class MaterializationHashRecorder: @unchecked Sendable { private let lock = NSLock() @@ -395,27 +379,28 @@ private final class MaterializationHashRecorder: @unchecked Sendable { } } +/// Sync hash-hook fence: `oidForBytes` is synchronous, so use a hang-hardened +/// `TestBlockingFence` (bounded wait + XCTFail/fail-open) instead of unbounded semaphores. private final class MaterializationHashCancellationGate: @unchecked Sendable { - private let hashingEntered = DispatchSemaphore(value: 0) - private let hashingRelease = DispatchSemaphore(value: 0) + private let fence = TestBlockingFence(name: "materialization hash cancellation") func hash(bytes: Data, objectFormat: GitObjectFormat) -> GitBlobOID { - hashingEntered.signal() - hashingRelease.wait() + fence.enterAndWait() return GitBlobOID.blob(bytes: bytes, objectFormat: objectFormat) } func waitUntilHashing() async { - await withCheckedContinuation { continuation in + // `waitUntilEntered` is synchronous NSCondition; park off the calling executor. + await withCheckedContinuation { (continuation: CheckedContinuation) in DispatchQueue.global(qos: .userInitiated).async { - self.hashingEntered.wait() + _ = self.fence.waitUntilEntered() continuation.resume() } } } func releaseHashing() { - hashingRelease.signal() + fence.release() } } diff --git a/Tests/RepoPromptTests/Services/VCS/GitRawOutputSpoolTests.swift b/Tests/RepoPromptTests/Services/VCS/GitRawOutputSpoolTests.swift index c32b1064a..91e87e564 100644 --- a/Tests/RepoPromptTests/Services/VCS/GitRawOutputSpoolTests.swift +++ b/Tests/RepoPromptTests/Services/VCS/GitRawOutputSpoolTests.swift @@ -84,11 +84,13 @@ final class GitRawOutputSpoolTests: XCTestCase { ) try cancellationSpool.append(Data("payload".utf8)) let cancellationLease = try cancellationSpool.finish() + let cancellationGate = GitRawOutputCancellationGate() let task = Task { let reader = try cancellationLease.makeReader() - try await Task.sleep(for: .seconds(60)) + try await cancellationGate.waitUntilCancelled() return try reader.nextChunk() } + await cancellationGate.waitUntilEntered() task.cancel() do { _ = try await task.value @@ -138,3 +140,5 @@ final class GitRawOutputSpoolTests: XCTestCase { ) } } + +private typealias GitRawOutputCancellationGate = TestCancellationGate diff --git a/Tests/RepoPromptTests/Services/VCS/GitTargetEvidenceStreamingParserTests.swift b/Tests/RepoPromptTests/Services/VCS/GitTargetEvidenceStreamingParserTests.swift index 4c4934848..c2d1bad37 100644 --- a/Tests/RepoPromptTests/Services/VCS/GitTargetEvidenceStreamingParserTests.swift +++ b/Tests/RepoPromptTests/Services/VCS/GitTargetEvidenceStreamingParserTests.swift @@ -257,6 +257,7 @@ final class GitTargetEvidenceStreamingParserTests: XCTestCase { func testCancellationInterruptsDirectWriterEmission() async throws { let entered = GitTargetEvidenceAsyncSignal() + let cancellationGate = GitTargetEvidenceCancellationGate() let objectID = objectID let task = Task { var parser = try GitTargetIndexStreamingParser( @@ -264,7 +265,7 @@ final class GitTargetEvidenceStreamingParserTests: XCTestCase { rootPrefix: GitRepositoryRelativeRootPrefix("Root") ) { _ in await entered.signal() - try await Task.sleep(for: .seconds(60)) + try await cancellationGate.waitUntilCancelled() } try await parser.consume(Data("H 100644 \(objectID) 0\tRoot/file\0".utf8)) } @@ -351,6 +352,8 @@ private actor GitTargetEvidenceAsyncSignal { } } +private typealias GitTargetEvidenceCancellationGate = TestCancellationGate + private extension Data { func chunkedForParserTest(widths: [Int]) -> [Data] { precondition(!widths.isEmpty && widths.allSatisfy { $0 > 0 }) diff --git a/Tests/RepoPromptTests/Services/VCS/GitWorktreeCreationReceiptTests.swift b/Tests/RepoPromptTests/Services/VCS/GitWorktreeCreationReceiptTests.swift index e505f7859..58f7384fa 100644 --- a/Tests/RepoPromptTests/Services/VCS/GitWorktreeCreationReceiptTests.swift +++ b/Tests/RepoPromptTests/Services/VCS/GitWorktreeCreationReceiptTests.swift @@ -1850,39 +1850,20 @@ private actor MutationTokenBox { } #if DEBUG - private actor ReceiptMutationLockGate { - private var entered = false - private var released = false - private var enteredWaiters: [CheckedContinuation] = [] - private var releaseWaiters: [CheckedContinuation] = [] + /// Receipt mutation lock fence (shared `TestReleaseFence` with legacy names). + private final class ReceiptMutationLockGate: @unchecked Sendable { + private let fence = TestReleaseFence(name: "receipt mutation lock gate") func enterAndWaitForRelease() async { - entered = true - let waiters = enteredWaiters - enteredWaiters.removeAll() - for waiter in waiters { - waiter.resume() - } - guard !released else { return } - await withCheckedContinuation { continuation in - releaseWaiters.append(continuation) - } + await fence.enterAndWait() } - func waitUntilEntered() async { - guard !entered else { return } - await withCheckedContinuation { continuation in - enteredWaiters.append(continuation) - } + func waitUntilEntered(timeout: TimeInterval = TestFenceDefaults.enterWait) async { + _ = await fence.waitUntilEntered(timeout: timeout) } func release() { - released = true - let waiters = releaseWaiters - releaseWaiters.removeAll() - for waiter in waiters { - waiter.resume() - } + fence.release() } } #endif diff --git a/Tests/RepoPromptTests/WorkspaceContext/CodemapAutomaticSelectionBusyTests.swift b/Tests/RepoPromptTests/WorkspaceContext/CodemapAutomaticSelectionBusyTests.swift index a8fe1c520..dafb24ff5 100644 --- a/Tests/RepoPromptTests/WorkspaceContext/CodemapAutomaticSelectionBusyTests.swift +++ b/Tests/RepoPromptTests/WorkspaceContext/CodemapAutomaticSelectionBusyTests.swift @@ -1015,6 +1015,16 @@ final class CodemapAutomaticSelectionBusyTests: WorkspaceFileContextStoreCodemap let files = await store.files(inRoot: loaded.id).sorted { $0.standardizedRelativePath < $1.standardizedRelativePath } + XCTAssertEqual( + files.map(\.standardizedRelativePath), + ["Sources/First.swift", "Sources/Second.swift"] + ) + try await AsyncTestWait.waitUntil("automatic selection source identities are cataloged", timeout: 5) { + await store.codemapAutomaticSelectionSourceIdentities( + forFileIDs: files.map(\.id), + rootScope: .visibleWorkspace + ).count == 2 + } let demandCount = CodemapLockedCounter() let issuedTickets = CodemapLockedValues() let service = WorkspaceSelectionMutationService( diff --git a/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineInvalidationTests.swift b/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineInvalidationTests.swift index ebff422fe..cbce75317 100644 --- a/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineInvalidationTests.swift +++ b/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineInvalidationTests.swift @@ -458,7 +458,7 @@ final class CodemapBindingEngineInvalidationTests: CodemapBindingEngineTestCase root: root, runtime: runtime, capabilityHooks: WorkspaceCodemapGitCapabilityServiceHooks( - beforeResolution: { await gate.enter() } + beforeResolution: { await gate.enterIgnoringCancellationUntilRelease() } ) ) let registration = Task { await fixture.engine.registerRoot(fixture.registration) } diff --git a/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineManifestWriteTests.swift b/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineManifestWriteTests.swift index f7e1720c7..76c71aebf 100644 --- a/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineManifestWriteTests.swift +++ b/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineManifestWriteTests.swift @@ -24,7 +24,8 @@ final class CodemapBindingEngineManifestWriteTests: CodemapBindingEngineTestCase let demand = Task { await fixture.engine.demand(fixture.demand(path: "Sources/Shutdown.swift")) } - XCTAssertTrue(writeGate.waitUntilEntered()) + let writeEntered = await writeGate.waitUntilEntered() + XCTAssertTrue(writeEntered) let shutdownFinished = EngineCompletionFlag() let shutdown = Task { await fixture.engine.shutdown() @@ -68,7 +69,8 @@ final class CodemapBindingEngineManifestWriteTests: CodemapBindingEngineTestCase ) _ = await fixture.engine.registerRoot(fixture.registration) let first = Task { await fixture.engine.demand(fixture.demand(path: "Sources/One.swift")) } - XCTAssertTrue(writeGate.waitUntilEntered()) + let writeEntered = await writeGate.waitUntilEntered() + XCTAssertTrue(writeEntered) let second = Task { await fixture.engine.demand(fixture.demand(path: "Sources/Two.swift")) } XCTAssertTrue(hookEvents.wait(kind: .manifestRevisionQueued, numericValue: 2)) writeGate.release() @@ -209,7 +211,8 @@ final class CodemapBindingEngineManifestWriteTests: CodemapBindingEngineTestCase path: "Sources/First.swift" )) } - XCTAssertTrue(writerGate.waitUntilEntered()) + let writerEntered = await writerGate.waitUntilEntered() + XCTAssertTrue(writerEntered) defer { writerGate.release() } await engine.unloadRoot(rootEpoch: firstEpoch) guard case .cancelled = await first.value else { @@ -292,7 +295,8 @@ final class CodemapBindingEngineManifestWriteTests: CodemapBindingEngineTestCase let demand = Task { await fixture.engine.demand(fixture.demand(path: "Sources/Feature.swift")) } - XCTAssertTrue(writeGate.waitUntilEntered()) + let writeEntered = await writeGate.waitUntilEntered() + XCTAssertTrue(writeEntered) let invalidation = Task { await fixture.engine.invalidateModified( rootEpoch: fixture.rootEpoch, diff --git a/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineOverlayTests.swift b/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineOverlayTests.swift index 1fdb550bb..7b868897a 100644 --- a/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineOverlayTests.swift +++ b/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineOverlayTests.swift @@ -55,7 +55,8 @@ final class CodemapBindingEngineOverlayTests: CodemapBindingEngineTestCase { addTeardownBlock { publicationGate.release() } _ = await fixture.engine.registerRoot(fixture.registration) _ = await fixture.engine.scheduleProjectionPreload(rootEpoch: fixture.rootEpoch) - XCTAssertTrue(publicationGate.waitUntilEntered()) + let publicationEntered = await publicationGate.waitUntilEntered() + XCTAssertTrue(publicationEntered) guard case .ready = await fixture.engine.demand(fixture.demand(path: "Sources/Live.swift")) else { publicationGate.release() diff --git a/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineProjectionTests.swift b/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineProjectionTests.swift index 7744ffd43..70105ee3a 100644 --- a/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineProjectionTests.swift +++ b/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineProjectionTests.swift @@ -124,7 +124,8 @@ final class CodemapBindingEngineProjectionTests: CodemapBindingEngineTestCase { return XCTFail("Expected eligible registration.") } _ = await fixture.engine.scheduleProjectionPreload(rootEpoch: fixture.rootEpoch) - XCTAssertTrue(pageGate.waitUntilEntered()) + let pageEntered = await pageGate.waitUntilEntered() + XCTAssertTrue(pageEntered) let acquisition = await fixture.engine.acquireProjectionDemand( rootEpoch: fixture.rootEpoch, @@ -268,7 +269,8 @@ final class CodemapBindingEngineProjectionTests: CodemapBindingEngineTestCase { guard case let .acquired(replacementTicket, _) = replacement else { return XCTFail("Expected capacity to be reusable after expiry.") } - XCTAssertTrue(pageGate.waitUntilEntered()) + let pageEntered = await pageGate.waitUntilEntered() + XCTAssertTrue(pageEntered) clock.set(401) pageGate.release() let completedAfterDeadline = await waitForEngineCondition { @@ -392,7 +394,8 @@ final class CodemapBindingEngineProjectionTests: CodemapBindingEngineTestCase { } _ = await engine.scheduleProjectionPreload(rootEpoch: rootEpochs[0]) - XCTAssertTrue(blocker.waitUntilEntered()) + let blockerEntered = await blocker.waitUntilEntered() + XCTAssertTrue(blockerEntered) let later = await engine.acquireProjectionDemand( rootEpoch: rootEpochs[1], fileIDs: [UUID()], @@ -485,7 +488,8 @@ final class CodemapBindingEngineProjectionTests: CodemapBindingEngineTestCase { let foreground = Task { await fixture.engine.demand(fixture.demand(path: "Sources/Foreground.swift")) } - XCTAssertTrue(foregroundGate.waitUntilEntered()) + let foregroundEntered = await foregroundGate.waitUntilEntered() + XCTAssertTrue(foregroundEntered) defer { foreground.cancel() foregroundGate.release() @@ -1434,7 +1438,8 @@ final class CodemapBindingEngineProjectionTests: CodemapBindingEngineTestCase { ) _ = await fixture.engine.registerRoot(fixture.registration) _ = await fixture.engine.scheduleProjectionPreload(rootEpoch: fixture.rootEpoch) - XCTAssertTrue(gate.waitUntilEntered()) + let gateEntered = await gate.waitUntilEntered() + XCTAssertTrue(gateEntered) let unload = Task { await fixture.engine.unloadRoot(rootEpoch: fixture.rootEpoch) } gate.release() diff --git a/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineRootLeaseTests.swift b/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineRootLeaseTests.swift index a599fbc89..9064a1113 100644 --- a/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineRootLeaseTests.swift +++ b/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineRootLeaseTests.swift @@ -182,7 +182,7 @@ final class CodemapBindingEngineRootLeaseTests: CodemapBindingEngineTestCase { let runtime = try CodeMapArtifactRuntime( rootURL: artifactRoot, manifestStoreHooks: CodeMapRootManifestStoreHooks( - afterReadAdmission: { await loadGate.enter() } + afterReadAdmission: { await loadGate.enterIgnoringCancellationUntilRelease() } ) ) let fixture = try await makeEngineFixture(root: root, runtime: runtime) diff --git a/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineWarmManifestTests.swift b/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineWarmManifestTests.swift index 8f3bf5a8c..6299be843 100644 --- a/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineWarmManifestTests.swift +++ b/Tests/RepoPromptTests/WorkspaceContext/CodemapBindingEngineWarmManifestTests.swift @@ -329,7 +329,8 @@ final class CodemapBindingEngineWarmManifestTests: CodemapBindingEngineTestCase let racedLookup = Task { await fixture.engine.lookupPublishedArtifact(lookupRequest(UUID())) } - XCTAssertTrue(lookupCurrentnessGate.waitUntilEntered()) + let lookupCurrentnessEntered = await lookupCurrentnessGate.waitUntilEntered() + XCTAssertTrue(lookupCurrentnessEntered) let accountingBeforePostLookupInvalidation = await fixture.engine.accounting() let invalidation = await fixture.engine.invalidateModified( rootEpoch: fixture.rootEpoch, diff --git a/Tests/RepoPromptTests/WorkspaceContext/Search/SearchPathFilteringTests.swift b/Tests/RepoPromptTests/WorkspaceContext/Search/SearchPathFilteringTests.swift index 4d343220d..88fb5373c 100644 --- a/Tests/RepoPromptTests/WorkspaceContext/Search/SearchPathFilteringTests.swift +++ b/Tests/RepoPromptTests/WorkspaceContext/Search/SearchPathFilteringTests.swift @@ -100,13 +100,14 @@ final class SearchPathFilteringTests: XCTestCase { clauses: [.legacyPrefix(candidateLower: "sources")] ) - let gate = SearchCancellationGate() + let cancellationGate = SearchPathFilteringCancellationGate() let task = Task.detached(priority: .background) { - await gate.wait() + // Shared cancellation gate resumes with CancellationError; continue into filter. + try? await cancellationGate.waitUntilCancelled() return filterPathIndicesResult(snapshots: snapshots, spec: spec) } + await cancellationGate.waitUntilEntered() task.cancel() - await gate.open() let result = await task.value XCTAssertTrue(result.cancelled) @@ -114,21 +115,8 @@ final class SearchPathFilteringTests: XCTestCase { } } -private actor SearchCancellationGate { - private var isOpen = false - private var waiters: [CheckedContinuation] = [] - - func open() { - isOpen = true - let current = waiters - waiters.removeAll() - current.forEach { $0.resume() } - } - - func wait() async { - if isOpen { return } - await withCheckedContinuation { continuation in - waiters.append(continuation) - } - } -} +/// Search path filter cancel handshake (shared `TestCancellationGate`). +/// +/// Note: resumes with `CancellationError` on cancel (stricter than the old Never-continuation +/// version, which treated cancel as a non-throwing resume). Call sites only assert cancellation. +private typealias SearchPathFilteringCancellationGate = TestCancellationGate diff --git a/Tests/RepoPromptTests/WorkspaceContext/Search/StoreBackedWorkspaceSearchConcurrencyMatrixTests.swift b/Tests/RepoPromptTests/WorkspaceContext/Search/StoreBackedWorkspaceSearchConcurrencyMatrixTests.swift index bbb43b07a..ed9669075 100644 --- a/Tests/RepoPromptTests/WorkspaceContext/Search/StoreBackedWorkspaceSearchConcurrencyMatrixTests.swift +++ b/Tests/RepoPromptTests/WorkspaceContext/Search/StoreBackedWorkspaceSearchConcurrencyMatrixTests.swift @@ -353,65 +353,90 @@ import XCTest case sameStoreBroadSearchDidNotStart } - private actor SearchPermitGate { + /// Search permit fence: shared release fence + started-count polling. + private final class SearchPermitGate: @unchecked Sendable { + private let fence = TestReleaseFence(name: "search permit gate") + private let lock = NSLock() private var startedCount = 0 - private var released = false - private var releaseWaiters: [CheckedContinuation] = [] func markStartedAndWaitForRelease() async { - startedCount += 1 - guard !released else { return } - await withCheckedContinuation { continuation in - releaseWaiters.append(continuation) - } + lock.withLock { startedCount += 1 } + await fence.enterAndWait() } func waitUntilStartedCount( _ expectedCount: Int, timeout: Duration = .seconds(2) ) async -> Bool { - let clock = ContinuousClock() - let deadline = clock.now.advanced(by: timeout) - while startedCount < expectedCount, clock.now < deadline { - try? await Task.sleep(for: .milliseconds(5)) + let timeoutInterval = TestFenceDefaults.timeInterval(timeout) + do { + try await AsyncTestWait.waitUntil( + "search permit started count \(expectedCount)", + timeout: timeoutInterval + ) { + self.lock.withLock { self.startedCount >= expectedCount } + } + return true + } catch { + return lock.withLock { startedCount >= expectedCount } } - return startedCount >= expectedCount } func release() { - released = true - releaseWaiters.forEach { $0.resume() } - releaseWaiters.removeAll() + fence.release() } } - private actor KWayIsolationBarrier { + /// K-way barrier: wait for arrivals then park on a cancellable async release fence. + private final class KWayIsolationBarrier: @unchecked Sendable { private let expectedCount: Int + private let releaseFence = TestReleaseFence(name: "k-way isolation barrier") + private let lock = NSLock() private var arrivedCount = 0 - private var released = false init(expectedCount: Int) { self.expectedCount = expectedCount } - func arriveAndWaitForRelease() async { - arrivedCount += 1 - while !released { - try? await Task.sleep(for: .milliseconds(1)) + func arriveAndWaitForRelease(timeout: Duration = TestFenceDefaults.releaseWaitDuration) async { + lock.withLock { arrivedCount += 1 } + let timeoutInterval = TestFenceDefaults.timeInterval(timeout) + await withTaskGroup(of: Bool.self) { group in + group.addTask { + await self.releaseFence.enterAndWait() + return true + } + group.addTask { + try? await Task.sleep(nanoseconds: UInt64((timeoutInterval * 1_000_000_000).rounded())) + return false + } + let first = await group.next() ?? false + group.cancelAll() + if !first { + XCTFail("Timed out waiting for KWayIsolationBarrier release") + // Fail open so remaining waiters (and this cancelled task) can unwind. + self.releaseFence.release() + } } } func waitUntilAllArrived(timeout: Duration = .seconds(2)) async -> Bool { - let clock = ContinuousClock() - let deadline = clock.now.advanced(by: timeout) - while arrivedCount < expectedCount, clock.now < deadline { - try? await Task.sleep(for: .milliseconds(1)) + let timeoutInterval = TestFenceDefaults.timeInterval(timeout) + do { + try await AsyncTestWait.waitUntil( + "k-way isolation all arrived", + timeout: timeoutInterval + ) { + self.lock.withLock { self.arrivedCount >= self.expectedCount } + } + return true + } catch { + return lock.withLock { arrivedCount >= expectedCount } } - return arrivedCount == expectedCount } func release() { - released = true + releaseFence.release() } } #endif diff --git a/Tests/RepoPromptTests/WorkspaceContext/WorkspaceCodemapBindingIntegrationRegistryTests.swift b/Tests/RepoPromptTests/WorkspaceContext/WorkspaceCodemapBindingIntegrationRegistryTests.swift index 461b2e568..ca0acb59c 100644 --- a/Tests/RepoPromptTests/WorkspaceContext/WorkspaceCodemapBindingIntegrationRegistryTests.swift +++ b/Tests/RepoPromptTests/WorkspaceContext/WorkspaceCodemapBindingIntegrationRegistryTests.swift @@ -106,8 +106,10 @@ final class WorkspaceCodemapBindingIntegrationRegistryTests: XCTestCase { let catalogTask = Task { await catalog.resolveManifestBinding(rootEpoch, bindingIdentity.standardizedRelativePath) } - await sourceGate.waitUntilEntered() - await catalogGate.waitUntilEntered() + let sourceEntered = await sourceGate.waitUntilEntered() + let catalogEntered = await catalogGate.waitUntilEntered() + XCTAssertTrue(sourceEntered) + XCTAssertTrue(catalogEntered) let didUnregister = await registry.unregister(token) XCTAssertTrue(didUnregister) await sourceGate.release() @@ -177,8 +179,10 @@ final class WorkspaceCodemapBindingIntegrationRegistryTests: XCTestCase { let catalogTask = Task { await catalog.resolveManifestBinding(rootEpoch, identity.standardizedRelativePath) } - await sourceGate.waitUntilEntered() - await catalogGate.waitUntilEntered() + let sourceEntered = await sourceGate.waitUntilEntered() + let catalogEntered = await catalogGate.waitUntilEntered() + XCTAssertTrue(sourceEntered) + XCTAssertTrue(catalogEntered) routeEndpoint = nil XCTAssertNil(weakEndpoint.value) @@ -384,26 +388,4 @@ private enum RegistryTestError: Error { case expectedProjectionPage } -private actor RegistryRouteGate { - private var entered = false - private var released = false - private var continuation: CheckedContinuation? - - func enterAndWait() async { - entered = true - guard !released else { return } - await withCheckedContinuation { continuation = $0 } - } - - func waitUntilEntered() async { - while !entered { - await Task.yield() - } - } - - func release() { - released = true - continuation?.resume() - continuation = nil - } -} +private typealias RegistryRouteGate = TestReleaseFence diff --git a/Tests/RepoPromptTests/WorkspaceContext/WorkspaceCodemapLiveOverlayTests.swift b/Tests/RepoPromptTests/WorkspaceContext/WorkspaceCodemapLiveOverlayTests.swift index cd1736fb2..88e84f873 100644 --- a/Tests/RepoPromptTests/WorkspaceContext/WorkspaceCodemapLiveOverlayTests.swift +++ b/Tests/RepoPromptTests/WorkspaceContext/WorkspaceCodemapLiveOverlayTests.swift @@ -3115,30 +3115,27 @@ final class WorkspaceCodemapLiveOverlayTests: XCTestCase { } } -private actor OverlayFirstCommitGate { - private var firstCommitEntered = false - private var firstCommitReleased = false +/// First-commit-only fence: parks only on the first `enter()`. +private final class OverlayFirstCommitGate: @unchecked Sendable { + private let fence = TestReleaseFence(name: "overlay first commit gate") + private let lock = NSLock() private var commitCount = 0 - private var continuation: CheckedContinuation? func enter() async { - commitCount += 1 - guard commitCount == 1 else { return } - firstCommitEntered = true - if firstCommitReleased { return } - await withCheckedContinuation { continuation = $0 } + let shouldPark = lock.withLock { () -> Bool in + commitCount += 1 + return commitCount == 1 + } + guard shouldPark else { return } + await fence.enterAndWait() } - func waitUntilFirstCommit() async { - while !firstCommitEntered { - await Task.yield() - } + func waitUntilFirstCommit(timeout: TimeInterval = TestFenceDefaults.enterWait) async { + _ = await fence.waitUntilEntered(timeout: timeout) } func releaseFirstCommit() { - firstCommitReleased = true - continuation?.resume() - continuation = nil + fence.release() } } diff --git a/Tests/RepoPromptTests/WorkspaceContext/WorkspaceFileContextStoreExactCapabilityTests.swift b/Tests/RepoPromptTests/WorkspaceContext/WorkspaceFileContextStoreExactCapabilityTests.swift index 6cbd680af..5ee25824b 100644 --- a/Tests/RepoPromptTests/WorkspaceContext/WorkspaceFileContextStoreExactCapabilityTests.swift +++ b/Tests/RepoPromptTests/WorkspaceContext/WorkspaceFileContextStoreExactCapabilityTests.swift @@ -259,7 +259,8 @@ final class WorkspaceFileContextStoreExactCapabilityTests: XCTestCase { folderPolicy: .expandFolders ) } - await gate.waitUntilEntered() + let gateEntered = await gate.waitUntilEntered() + XCTAssertTrue(gateEntered) await store.releaseSessionWorktreeOwnership(ownerID: fixture.sessionID) let replacement = try await store.loadRoot( path: worktreeRoot.path, @@ -306,35 +307,5 @@ final class WorkspaceFileContextStoreExactCapabilityTests: XCTestCase { } #if DEBUG - private actor ContextBuilderCandidateGate { - private var entered = false - private var released = false - private var enteredWaiters: [CheckedContinuation] = [] - private var releaseWaiters: [CheckedContinuation] = [] - - func enterAndWait() async { - entered = true - let waiters = enteredWaiters - enteredWaiters.removeAll() - waiters.forEach { $0.resume() } - guard !released else { return } - await withCheckedContinuation { continuation in - releaseWaiters.append(continuation) - } - } - - func waitUntilEntered() async { - guard !entered else { return } - await withCheckedContinuation { continuation in - enteredWaiters.append(continuation) - } - } - - func release() { - released = true - let waiters = releaseWaiters - releaseWaiters.removeAll() - waiters.forEach { $0.resume() } - } - } + private typealias ContextBuilderCandidateGate = TestReleaseFence #endif diff --git a/Tests/RepoPromptTests/WorkspaceContext/WorkspaceRootTargetEvidenceCoordinatorTests.swift b/Tests/RepoPromptTests/WorkspaceContext/WorkspaceRootTargetEvidenceCoordinatorTests.swift index 4d7a40c85..45e0ea11c 100644 --- a/Tests/RepoPromptTests/WorkspaceContext/WorkspaceRootTargetEvidenceCoordinatorTests.swift +++ b/Tests/RepoPromptTests/WorkspaceContext/WorkspaceRootTargetEvidenceCoordinatorTests.swift @@ -240,12 +240,13 @@ final class WorkspaceRootTargetEvidenceCoordinatorTests: XCTestCase { func testLastWaiterCancellationCancelsProducerAndCleansFlight() async throws { let coordinator = WorkspaceRootTargetEvidenceCoordinator() let probe = TargetEvidenceCancellationProbe() + let cancellationGate = TargetEvidenceCancellationGate() let key = Self.flightKey(suffix: "last-waiter") let waiter = Task { try await coordinator.claim(for: key) { _, _ in await probe.markStarted() do { - try await Task.sleep(for: .seconds(30)) + try await cancellationGate.waitUntilCancelled() XCTFail("Producer should have been cancelled") return .sealed( handle: TargetEvidenceTestHandle(), @@ -315,13 +316,14 @@ final class WorkspaceRootTargetEvidenceCoordinatorTests: XCTestCase { let coordinator = WorkspaceRootTargetEvidenceCoordinator() let releases = TargetEvidenceReleaseCounter() let probe = TargetEvidenceCancellationProbe() + let cancellationGate = TargetEvidenceCancellationGate() let waiter = Task { try await coordinator.claim(for: Self.flightKey(suffix: "last-cancel-resource")) { _, context in try await context.retainResource( TargetEvidenceTestResource(attempt: context.attemptIndex, releases: releases) ) await probe.markStarted() - try await Task.sleep(for: .seconds(30)) + try await cancellationGate.waitUntilCancelled() return .sealed( handle: TargetEvidenceTestHandle(), authoritySnapshotIdentity: Data("authority".utf8) @@ -348,6 +350,7 @@ final class WorkspaceRootTargetEvidenceCoordinatorTests: XCTestCase { func testLastWaiterCancellationJoinsProducerBeforeReleasingAttemptResource() async throws { let coordinator = WorkspaceRootTargetEvidenceCoordinator() let ordering = TargetEvidenceCancellationOrderingProbe() + let cancellationGate = TargetEvidenceCancellationGate() let waiter = Task { try await coordinator.claim(for: Self.flightKey(suffix: "last-cancel-join-order")) { _, context in @@ -356,7 +359,7 @@ final class WorkspaceRootTargetEvidenceCoordinatorTests: XCTestCase { ) await ordering.markProducerStarted() do { - try await Task.sleep(for: .seconds(30)) + try await cancellationGate.waitUntilCancelled() return .sealed( handle: TargetEvidenceTestHandle(), authoritySnapshotIdentity: Data("authority".utf8) @@ -549,6 +552,8 @@ private actor TargetEvidenceCancellationProbe { } } +private typealias TargetEvidenceCancellationGate = TestCancellationGate + private actor TargetEvidenceCancellationOrderingProbe { enum Event: Equatable { case producerTerminated