Make control-plane and repair waits deterministic

Sources/XcodeMCPProxyRuntime/Session/ControlPlane/ControlPlaneCoordinator.swift

@@ -131,6 +131,9 @@ actor ControlPlaneCoordinator {
         if let rawResult = cachedToolsCatalog() {
             return rawResult
         }
+        guard deadlineExceeded(deadlineUptimeNs) == false else {
+            throw TimeoutError()
+        }
 
         let requestedTimeout = sharedRequestTimeout(for: deadlineUptimeNs)
         let requestedPromotionDeadlineUptimeNs = promotionDeadlineUptimeNs(
@@ -166,6 +169,9 @@ actor ControlPlaneCoordinator {
         guard acceptsNewLoads else {
             throw CancellationError()
         }
+        guard deadlineExceeded(deadlineUptimeNs) == false else {
+            throw TimeoutError()
+        }
         let requestedTimeout = sharedRequestTimeout(for: deadlineUptimeNs)
         let requestedPromotionDeadlineUptimeNs = promotionDeadlineUptimeNs(
             forWaiterDeadlineUptimeNs: deadlineUptimeNs
@@ -202,6 +208,9 @@ actor ControlPlaneCoordinator {
             syncDebug()
             return rawResult
         }
+        guard deadlineExceeded(deadlineUptimeNs) == false else {
+            return nil
+        }
         guard
             toolsCatalogLoad == nil,
             prewarmToolsCatalogLoad == nil
@@ -408,6 +417,11 @@ actor ControlPlaneCoordinator {
             return
         }
         if deadlineExceeded(deadlineUptimeNs) {
+            if load.waiters.isEmpty {
+                clearToolsCatalogLoadState(loadID: loadID)
+                cancelToolsCatalogLoad(load, error: TimeoutError())
+                syncDebug()
+            }
             continuation.resume(throwing: TimeoutError())
             return
         }
@@ -439,6 +453,11 @@ actor ControlPlaneCoordinator {
             return
         }
         if deadlineExceeded(deadlineUptimeNs) {
+            if load.waiters.isEmpty {
+                windowLoads.removeValue(forKey: route)
+                cancelWindowLoad(load, error: TimeoutError())
+                syncDebug()
+            }
             continuation.resume(throwing: TimeoutError())
             return
         }

Sources/XcodeMCPProxyRuntime/Session/DocumentationProvider.swift

@@ -913,41 +913,42 @@ private final class DocumentationSearchServiceRepairWaiter: @unchecked Sendable
     private var resolved = false
 
     func wait(
-        for task: Task<DocumentationSearchServiceRepairResult, Never>,
+        operation: @escaping @Sendable () async -> DocumentationSearchServiceRepairResult,
         timeout: TimeAmount,
         clock: ClockClient
     ) async -> Result {
         await withTaskCancellationHandler {
             await withCheckedContinuation { continuation in
-                setContinuation(continuation)
+                guard setContinuation(continuation) else {
+                    return
+                }
                 addTask(Task {
-                    let result = await task.value
+                    let result = await operation()
                     self.resume(Result(repairResult: result, timedOut: false))
                 })
                 addTask(Task {
                     await clock.sleep(.nanoseconds(timeout.nanoseconds))
                     guard Task.isCancelled == false else {
                         return
                     }
-                    task.cancel()
                     self.resume(Result(repairResult: nil, timedOut: true))
                 })
             }
         } onCancel: {
-            task.cancel()
             resume(Result(repairResult: nil, timedOut: true))
         }
     }
 
-    private func setContinuation(_ continuation: CheckedContinuation<Result, Never>) {
+    private func setContinuation(_ continuation: CheckedContinuation<Result, Never>) -> Bool {
         lock.lock()
         if resolved {
             lock.unlock()
             continuation.resume(returning: Result(repairResult: nil, timedOut: true))
-            return
+            return false
         }
         self.continuation = continuation
         lock.unlock()
+        return true
     }
 
     private func addTask(_ task: Task<Void, Never>) {
@@ -4183,8 +4184,9 @@ actor DocumentationProviderManager: DocumentationProviderManaging {
         guard let timeout, timeout.nanoseconds > 0 else {
             return await serviceRepairer.repairDocumentationSearch(for: target)
         }
+        let serviceRepairer = serviceRepairer
         let result = await DocumentationSearchServiceRepairWaiter().wait(
-            for: Task {
+            operation: {
                 await serviceRepairer.repairDocumentationSearch(for: target)
             },
             timeout: timeout,

Tests/XcodeMCPProxyRuntimeTests/RuntimeCoordinatorTestSupport.swift

@@ -600,7 +600,63 @@ func seedProcessToolCatalogs(
     }
 }
 
+struct ControlPlaneLoadTestSnapshot: Sendable {
+    let loadID: UUID
+    let waiterCount: Int
+    let foregroundWaiterCount: Int
+    let rpcHandle: ControlPlane.RPCHandle
+}
+
 extension ControlPlaneCoordinator {
+    func requestToolsCatalogLoadSnapshotForTesting() -> ControlPlaneLoadTestSnapshot? {
+        toolsCatalogLoad.map(loadSnapshotForTesting)
+    }
+
+    func prewarmToolsCatalogLoadSnapshotForTesting() -> ControlPlaneLoadTestSnapshot? {
+        prewarmToolsCatalogLoad.map(loadSnapshotForTesting)
+    }
+
+    func windowLoadSnapshotForTesting(
+        route: ControlPlane.Route
+    ) -> ControlPlaneLoadTestSnapshot? {
+        windowLoads[route].map {
+            ControlPlaneLoadTestSnapshot(
+                loadID: $0.loadID,
+                waiterCount: $0.waiters.count,
+                foregroundWaiterCount: $0.waiters.count,
+                rpcHandle: $0.rpcHandle
+            )
+        }
+    }
+
+    @discardableResult
+    func timeoutForegroundToolsCatalogWaiterForTesting() -> Bool {
+        let loads = [toolsCatalogLoad, prewarmToolsCatalogLoad].compactMap { $0 }
+        guard loads.count == 1,
+              let load = loads.first,
+              let waiterID = load.waiters.first(where: {
+                  if case .foreground = $0.value.kind {
+                      return true
+                  }
+                  return false
+              })?.key else {
+            return false
+        }
+        timeoutToolsCatalogWaiter(loadID: load.loadID, waiterID: waiterID)
+        return true
+    }
+
+    @discardableResult
+    func timeoutWindowWaiterForTesting(route: ControlPlane.Route) -> Bool {
+        guard let load = windowLoads[route],
+              load.waiters.count == 1,
+              let waiterID = load.waiters.keys.first else {
+            return false
+        }
+        timeoutWindowWaiter(route: route, loadID: load.loadID, waiterID: waiterID)
+        return true
+    }
+
     func drainLoadsForTesting() async {
         while completionTasks.isEmpty == false {
             let tasks = Array(completionTasks.values)
@@ -609,6 +665,17 @@ extension ControlPlaneCoordinator {
             }
         }
     }
+
+    private func loadSnapshotForTesting(
+        _ load: ToolsCatalogLoadState
+    ) -> ControlPlaneLoadTestSnapshot {
+        ControlPlaneLoadTestSnapshot(
+            loadID: load.loadID,
+            waiterCount: load.waiters.count,
+            foregroundWaiterCount: load.foregroundWaiterCount,
+            rpcHandle: load.rpcHandle
+        )
+    }
 }
 
 func makeJSONRPCResponse(id: Int64, result: [String: Any]) throws -> Data {