Harden test harness against hangs and waiter leaks

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.

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)

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<Void, Never>] = []
-    private var releaseWaiters: [CheckedContinuation<Void, Never>] = []
+/// 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<Void, Never>] = []
-    private var releaseWaiters: [CheckedContinuation<Void, Never>] = []
-    private var cancellationWaiters: [CheckedContinuation<Void, Never>] = []
+    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) }
         }
     }

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()
     }
 }
 

Tests/RepoPromptTests/CodeMap/CodeMapRootManifestStoreTests.swift

@@ -2707,23 +2707,21 @@ 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
 
     var didSecondWriterAcquireLock: Bool {
         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<Void, Never>] = []
-    private var releaseContinuation: CheckedContinuation<Void, Never>?
+/// 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()
     }
 }