From 1c64d43c2858569f09864edbeee7b20343d1089c Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 3 Sep 2026 08:41:04 +0900 Subject: [PATCH 1/3] fix(proxy): stabilize asynchronous control-plane waits --- .../Session/DocumentationProvider.swift | 18 +++-- .../RuntimeCoordinatorTests.swift | 76 +++++++++++++++---- 2 files changed, 71 insertions(+), 23 deletions(-) diff --git a/Sources/XcodeMCPProxyRuntime/Session/DocumentationProvider.swift b/Sources/XcodeMCPProxyRuntime/Session/DocumentationProvider.swift index c0a56495..1c24f7dc 100644 --- a/Sources/XcodeMCPProxyRuntime/Session/DocumentationProvider.swift +++ b/Sources/XcodeMCPProxyRuntime/Session/DocumentationProvider.swift @@ -913,15 +913,17 @@ private final class DocumentationSearchServiceRepairWaiter: @unchecked Sendable private var resolved = false func wait( - for task: Task, + 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 { @@ -929,25 +931,24 @@ private final class DocumentationSearchServiceRepairWaiter: @unchecked Sendable 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) { + private func setContinuation(_ continuation: CheckedContinuation) -> 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) { @@ -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, diff --git a/Tests/XcodeMCPProxyRuntimeTests/RuntimeCoordinatorTests.swift b/Tests/XcodeMCPProxyRuntimeTests/RuntimeCoordinatorTests.swift index 55969f2a..b0da3bf3 100644 --- a/Tests/XcodeMCPProxyRuntimeTests/RuntimeCoordinatorTests.swift +++ b/Tests/XcodeMCPProxyRuntimeTests/RuntimeCoordinatorTests.swift @@ -6996,6 +6996,11 @@ struct RuntimeCoordinatorRecoveryTests { } let firstRequest = try await sentValue(from: upstream, at: 2, timeout: .seconds(2)) #expect(methodName(from: firstRequest) == "tools/list") + _ = try await waitWithTimeout("waiting for first tools/list waiter to attach") { + try await manager.controlPlaneDebugMirror.waitForSnapshot { + $0.waiterCounts.toolsCatalog == 1 + } + } uptimeClock.advance(by: .nanoseconds(120_000_001)) @@ -7006,18 +7011,18 @@ struct RuntimeCoordinatorRecoveryTests { ) } - try await waitForSentCount(upstream, count: 5, timeoutSeconds: 2) - let firstCancellation = try await sentValue( - from: upstream, - at: 3, - timeout: .seconds(2) + let firstCancellation = try await upstream.nextSent( + startingAt: 3, + matching: { methodName(from: $0) == "notifications/cancelled" } ) #expect( try extractCancellationRequestID(from: firstCancellation) == extractUpstreamID(from: firstRequest) ) - let secondRequest = try await sentValue(from: upstream, at: 4, timeout: .seconds(2)) - #expect(methodName(from: secondRequest) == "tools/list") + _ = try await upstream.nextSent( + startingAt: 3, + matching: { methodName(from: $0) == "tools/list" } + ) firstTask.cancel() secondTask.cancel() @@ -7120,6 +7125,11 @@ struct RuntimeCoordinatorRecoveryTests { } let firstRequest = try await sentValue(from: upstream, at: 2, timeout: .seconds(2)) #expect(methodName(from: firstRequest) == "tools/list") + _ = try await waitWithTimeout("waiting for first tools/list waiter to attach") { + try await manager.controlPlaneDebugMirror.waitForSnapshot { + $0.waiterCounts.toolsCatalog == 1 + } + } uptimeClock.advance(by: .nanoseconds(120_000_001)) @@ -7130,15 +7140,22 @@ struct RuntimeCoordinatorRecoveryTests { ) } try await upstream.waitForBlockedCancellation() - let firstCancellation = try await sentValue( - from: upstream, - at: 3, - timeout: .seconds(2) + let firstCancellation = try await upstream.nextSent( + startingAt: 3, + matching: { methodName(from: $0) == "notifications/cancelled" } ) #expect( try extractCancellationRequestID(from: firstCancellation) == extractUpstreamID(from: firstRequest) ) + _ = try await waitWithTimeout("waiting for promoted tools/list waiters to attach") { + try await manager.controlPlaneDebugMirror.waitForSnapshot { + $0.waiterCounts.toolsCatalog == 2 + } + } + let sharedLoadID = try #require( + await manager.controlPlaneCoordinator.toolsCatalogLoad?.loadID + ) uptimeClock.advance(by: .nanoseconds(120_000_001)) @@ -7154,8 +7171,7 @@ struct RuntimeCoordinatorRecoveryTests { $0.waiterCounts.toolsCatalog == 3 } } - let sentCount = await upstream.sentCount() - #expect(sentCount == 4) + #expect(await manager.controlPlaneCoordinator.toolsCatalogLoad?.loadID == sharedLoadID) await upstream.releaseBlockedCancellation() firstTask.cancel() @@ -7201,6 +7217,11 @@ struct RuntimeCoordinatorRecoveryTests { ) } _ = try await sentValue(from: upstream, at: 2, timeout: .seconds(2)) + _ = try await waitWithTimeout("waiting for first promoted tools/list waiter to attach") { + try await manager.controlPlaneDebugMirror.waitForSnapshot { + $0.waiterCounts.toolsCatalog == 1 + } + } uptimeClock.advance(by: .nanoseconds(120_000_001)) @@ -7210,7 +7231,15 @@ struct RuntimeCoordinatorRecoveryTests { requestTimeoutOverride: .seconds(5) ) } - _ = try await sentValue(from: upstream, at: 3, timeout: .seconds(2)) + _ = try await upstream.nextSent( + startingAt: 3, + matching: { methodName(from: $0) == "notifications/cancelled" } + ) + _ = try await waitWithTimeout("waiting for promoted tools/list waiters to attach") { + try await manager.controlPlaneDebugMirror.waitForSnapshot { + $0.waiterCounts.toolsCatalog == 2 + } + } firstTask.cancel() do { @@ -13455,6 +13484,13 @@ struct RuntimeCoordinatorWindowRoutingTests { ) } _ = try await sentValue(from: upstream, at: 2, timeout: .seconds(2)) + _ = try await waitWithTimeout( + "waiting for first promoted XcodeListWindows waiter to attach" + ) { + try await manager.controlPlaneDebugMirror.waitForSnapshot { + $0.waiterCounts.windows == 1 + } + } uptimeClock.advance(by: .nanoseconds(120_000_001)) @@ -13464,7 +13500,17 @@ struct RuntimeCoordinatorWindowRoutingTests { requestTimeoutOverride: .seconds(5) ) } - _ = try await sentValue(from: upstream, at: 3, timeout: .seconds(2)) + _ = try await upstream.nextSent( + startingAt: 3, + matching: { methodName(from: $0) == "notifications/cancelled" } + ) + _ = try await waitWithTimeout( + "waiting for promoted XcodeListWindows waiters to attach" + ) { + try await manager.controlPlaneDebugMirror.waitForSnapshot { + $0.waiterCounts.windows == 2 + } + } firstTask.cancel() do { From f43ad60a69106fc38914b25f8f2a2324027684c2 Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 3 Sep 2026 08:45:47 +0900 Subject: [PATCH 2/3] test(proxy): bound asynchronous promotion assertions --- .../RuntimeCoordinatorTests.swift | 54 ++++++++++++------- 1 file changed, 34 insertions(+), 20 deletions(-) diff --git a/Tests/XcodeMCPProxyRuntimeTests/RuntimeCoordinatorTests.swift b/Tests/XcodeMCPProxyRuntimeTests/RuntimeCoordinatorTests.swift index b0da3bf3..9eca110e 100644 --- a/Tests/XcodeMCPProxyRuntimeTests/RuntimeCoordinatorTests.swift +++ b/Tests/XcodeMCPProxyRuntimeTests/RuntimeCoordinatorTests.swift @@ -7011,18 +7011,24 @@ struct RuntimeCoordinatorRecoveryTests { ) } - let firstCancellation = try await upstream.nextSent( - startingAt: 3, - matching: { methodName(from: $0) == "notifications/cancelled" } - ) + let firstCancellation = try await waitWithTimeout( + "waiting for promoted tools/list cancellation" + ) { + try await upstream.nextSent( + startingAt: 3, + matching: { methodName(from: $0) == "notifications/cancelled" } + ) + } #expect( try extractCancellationRequestID(from: firstCancellation) == extractUpstreamID(from: firstRequest) ) - _ = try await upstream.nextSent( - startingAt: 3, - matching: { methodName(from: $0) == "tools/list" } - ) + _ = try await waitWithTimeout("waiting for replacement tools/list request") { + try await upstream.nextSent( + startingAt: 3, + matching: { methodName(from: $0) == "tools/list" } + ) + } firstTask.cancel() secondTask.cancel() @@ -7140,10 +7146,14 @@ struct RuntimeCoordinatorRecoveryTests { ) } try await upstream.waitForBlockedCancellation() - let firstCancellation = try await upstream.nextSent( - startingAt: 3, - matching: { methodName(from: $0) == "notifications/cancelled" } - ) + let firstCancellation = try await waitWithTimeout( + "waiting for shared tools/list cancellation" + ) { + try await upstream.nextSent( + startingAt: 3, + matching: { methodName(from: $0) == "notifications/cancelled" } + ) + } #expect( try extractCancellationRequestID(from: firstCancellation) == extractUpstreamID(from: firstRequest) @@ -7231,10 +7241,12 @@ struct RuntimeCoordinatorRecoveryTests { requestTimeoutOverride: .seconds(5) ) } - _ = try await upstream.nextSent( - startingAt: 3, - matching: { methodName(from: $0) == "notifications/cancelled" } - ) + _ = try await waitWithTimeout("waiting for migrated tools/list cancellation") { + try await upstream.nextSent( + startingAt: 3, + matching: { methodName(from: $0) == "notifications/cancelled" } + ) + } _ = try await waitWithTimeout("waiting for promoted tools/list waiters to attach") { try await manager.controlPlaneDebugMirror.waitForSnapshot { $0.waiterCounts.toolsCatalog == 2 @@ -13500,10 +13512,12 @@ struct RuntimeCoordinatorWindowRoutingTests { requestTimeoutOverride: .seconds(5) ) } - _ = try await upstream.nextSent( - startingAt: 3, - matching: { methodName(from: $0) == "notifications/cancelled" } - ) + _ = try await waitWithTimeout("waiting for migrated XcodeListWindows cancellation") { + try await upstream.nextSent( + startingAt: 3, + matching: { methodName(from: $0) == "notifications/cancelled" } + ) + } _ = try await waitWithTimeout( "waiting for promoted XcodeListWindows waiters to attach" ) { From b37c7e56e8c248d7b69ac32df1ab6c2ea983b6eb Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Thu, 3 Sep 2026 09:22:00 +0900 Subject: [PATCH 3/3] fix(proxy): make control-plane timing deterministic --- .../ControlPlaneCoordinator.swift | 19 ++ .../RuntimeCoordinatorTestSupport.swift | 67 ++++ .../RuntimeCoordinatorTests.swift | 306 ++++++++++++------ 3 files changed, 302 insertions(+), 90 deletions(-) diff --git a/Sources/XcodeMCPProxyRuntime/Session/ControlPlane/ControlPlaneCoordinator.swift b/Sources/XcodeMCPProxyRuntime/Session/ControlPlane/ControlPlaneCoordinator.swift index 46f8410f..7e8d63e0 100644 --- a/Sources/XcodeMCPProxyRuntime/Session/ControlPlane/ControlPlaneCoordinator.swift +++ b/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 } diff --git a/Tests/XcodeMCPProxyRuntimeTests/RuntimeCoordinatorTestSupport.swift b/Tests/XcodeMCPProxyRuntimeTests/RuntimeCoordinatorTestSupport.swift index f6e0e3cd..c0764c31 100644 --- a/Tests/XcodeMCPProxyRuntimeTests/RuntimeCoordinatorTestSupport.swift +++ b/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 { diff --git a/Tests/XcodeMCPProxyRuntimeTests/RuntimeCoordinatorTests.swift b/Tests/XcodeMCPProxyRuntimeTests/RuntimeCoordinatorTests.swift index 9eca110e..c8ee9763 100644 --- a/Tests/XcodeMCPProxyRuntimeTests/RuntimeCoordinatorTests.swift +++ b/Tests/XcodeMCPProxyRuntimeTests/RuntimeCoordinatorTests.swift @@ -6848,10 +6848,13 @@ struct RuntimeCoordinatorRecoveryTests { ) #expect(methodName(from: firstRequest) == "tools/list") try await upstream.waitForBlockedSend() - try await advanceRuntimeCoordinatorTimeout( - timeoutClock: clocks.timeoutClock, - uptimeClock: clocks.uptimeClock, - by: .seconds(5) + _ = try await waitWithTimeout("waiting for first tools/list waiter to attach") { + try await manager.controlPlaneDebugMirror.waitForSnapshot { + $0.waiterCounts.toolsCatalog == 1 + } + } + #expect( + await manager.controlPlaneCoordinator.timeoutForegroundToolsCatalogWaiterForTesting() ) await #expect(throws: TimeoutError.self) { _ = try await firstTask.value @@ -6890,10 +6893,13 @@ struct RuntimeCoordinatorRecoveryTests { try await waitForSentCount(upstream, count: 5, timeoutSeconds: 2) let secondRequest = try await sentValue(from: upstream, at: 4, timeout: .seconds(2)) #expect(methodName(from: secondRequest) == "tools/list") - try await advanceRuntimeCoordinatorTimeout( - timeoutClock: clocks.timeoutClock, - uptimeClock: clocks.uptimeClock, - by: .seconds(5) + _ = try await waitWithTimeout("waiting for second tools/list waiter to attach") { + try await manager.controlPlaneDebugMirror.waitForSnapshot { + $0.waiterCounts.toolsCatalog == 1 + } + } + #expect( + await manager.controlPlaneCoordinator.timeoutForegroundToolsCatalogWaiterForTesting() ) await #expect(throws: TimeoutError.self) { _ = try await secondTask.value @@ -6921,9 +6927,13 @@ struct RuntimeCoordinatorRecoveryTests { _ = try await initFuture.get() try await waitForSentCount(upstream, count: 2, timeoutSeconds: 2) - manager.refreshToolsListIfNeeded() let prewarmRequest = try await sentValue(from: upstream, at: 2, timeout: .seconds(2)) #expect(methodName(from: prewarmRequest) == "tools/list") + let prewarmLoad = try #require( + await manager.controlPlaneCoordinator.prewarmToolsCatalogLoadSnapshotForTesting() + ) + #expect(prewarmLoad.waiterCount == 1) + #expect(prewarmLoad.foregroundWaiterCount == 0) let sessionID = "session-tools-prewarm" _ = manager.session(id: sessionID) @@ -6939,7 +6949,16 @@ struct RuntimeCoordinatorRecoveryTests { $0.waiterCounts.toolsCatalog == 1 } } - #expect(await upstream.sentCount() == 3) + let reusedLoad = try #require( + await manager.controlPlaneCoordinator.prewarmToolsCatalogLoadSnapshotForTesting() + ) + #expect(reusedLoad.loadID == prewarmLoad.loadID) + #expect(reusedLoad.waiterCount == 2) + #expect(reusedLoad.foregroundWaiterCount == 1) + #expect( + await manager.controlPlaneCoordinator.requestToolsCatalogLoadSnapshotForTesting()? + .loadID == nil + ) let prewarmUpstreamID = try extractUpstreamID(from: prewarmRequest) let response: [String: Any] = [ @@ -6957,6 +6976,64 @@ struct RuntimeCoordinatorRecoveryTests { #expect(object["tools"] != nil) } + @Test func controlPlaneRejectsExpiredLoadsBeforeStartingUnobservedWork() async throws { + let group = borrowSharedTestEventLoopGroup() + defer { shutdownAndWait(group) } + let eventLoop = group.next() + let upstream = TestUpstreamClient() + let clocks = makeRuntimeCoordinatorDeterministicClocks() + let manager = RuntimeCoordinator( + config: makeConfig(requestTimeout: 5), + eventLoop: eventLoop, + upstreams: [upstream], + clock: clocks.clock + ) + defer { manager.shutdownAndWait() } + + let initFuture = manager.registerInitialize( + originalID: JSONRPC.ID(any: NSNumber(value: 1))!, + requestObject: makeInitializeRequest(id: 1), + on: eventLoop + ) + let initRequest = try await sentValue(from: upstream, at: 0, timeout: .seconds(2)) + let initUpstreamID = try extractUpstreamID(from: initRequest) + await upstream.yield(.message(try makeInitializeResponse(id: initUpstreamID))) + _ = try await initFuture.get() + try await waitForSentCount(upstream, count: 2, timeoutSeconds: 2) + + let expiredDeadline = clocks.uptimeClock.now() + await #expect(throws: TimeoutError.self) { + _ = try await manager.controlPlaneCoordinator.toolsCatalog( + deadlineUptimeNs: expiredDeadline + ) + } + #expect( + await manager.controlPlaneCoordinator.prewarmToolsCatalogIfNeeded( + deadlineUptimeNs: expiredDeadline + ) == nil + ) + await #expect(throws: TimeoutError.self) { + _ = try await manager.controlPlaneCoordinator.listWindows( + route: .anyHealthy, + deadlineUptimeNs: expiredDeadline + ) + } + + #expect( + await manager.controlPlaneCoordinator.requestToolsCatalogLoadSnapshotForTesting()? + .loadID == nil + ) + #expect( + await manager.controlPlaneCoordinator.prewarmToolsCatalogLoadSnapshotForTesting()? + .loadID == nil + ) + #expect( + await manager.controlPlaneCoordinator.windowLoadSnapshotForTesting(route: .anyHealthy)? + .loadID == nil + ) + #expect(await upstream.sentCount() == 2) + } + @Test func sessionManagerSharedToolsListPromotesPartlyConsumedSharedTimeout() async throws { @@ -6965,12 +7042,12 @@ struct RuntimeCoordinatorRecoveryTests { let eventLoop = group.next() let upstream = TestUpstreamClient() let config = makeConfig(requestTimeout: 5) - let uptimeClock = TestUptimeClock() + let clocks = makeRuntimeCoordinatorDeterministicClocks() let manager = RuntimeCoordinator( config: config, eventLoop: eventLoop, upstreams: [upstream], - nowUptimeNanoseconds: uptimeClock.now + clock: clocks.clock ) defer { manager.shutdownAndWait() } @@ -7001,8 +7078,13 @@ struct RuntimeCoordinatorRecoveryTests { $0.waiterCounts.toolsCatalog == 1 } } + let firstLoad = try #require( + await manager.controlPlaneCoordinator.requestToolsCatalogLoadSnapshotForTesting() + ) + #expect(firstLoad.waiterCount == 1) + #expect(firstLoad.foregroundWaiterCount == 1) - uptimeClock.advance(by: .nanoseconds(120_000_001)) + clocks.uptimeClock.advance(by: .nanoseconds(120_000_001)) let secondTask = Task { try await manager.sharedToolsList( @@ -7010,25 +7092,18 @@ struct RuntimeCoordinatorRecoveryTests { requestTimeoutOverride: .seconds(5) ) } - - let firstCancellation = try await waitWithTimeout( - "waiting for promoted tools/list cancellation" - ) { - try await upstream.nextSent( - startingAt: 3, - matching: { methodName(from: $0) == "notifications/cancelled" } - ) + _ = try await waitWithTimeout("waiting for promoted tools/list waiters to attach") { + try await manager.controlPlaneDebugMirror.waitForSnapshot { + $0.waiterCounts.toolsCatalog == 2 + } } - #expect( - try extractCancellationRequestID(from: firstCancellation) - == extractUpstreamID(from: firstRequest) + let promotedLoad = try #require( + await manager.controlPlaneCoordinator.requestToolsCatalogLoadSnapshotForTesting() ) - _ = try await waitWithTimeout("waiting for replacement tools/list request") { - try await upstream.nextSent( - startingAt: 3, - matching: { methodName(from: $0) == "tools/list" } - ) - } + #expect(promotedLoad.loadID != firstLoad.loadID) + #expect(promotedLoad.waiterCount == 2) + #expect(promotedLoad.foregroundWaiterCount == 2) + #expect(firstLoad.rpcHandle.isCancelled()) firstTask.cancel() secondTask.cancel() @@ -7099,12 +7174,12 @@ struct RuntimeCoordinatorRecoveryTests { let eventLoop = group.next() let upstream = TestUpstreamClient() let config = makeConfig(requestTimeout: 5) - let uptimeClock = TestUptimeClock() + let clocks = makeRuntimeCoordinatorDeterministicClocks() let manager = RuntimeCoordinator( config: config, eventLoop: eventLoop, upstreams: [upstream], - nowUptimeNanoseconds: uptimeClock.now + clock: clocks.clock ) defer { manager.shutdownAndWait() } @@ -7121,7 +7196,6 @@ struct RuntimeCoordinatorRecoveryTests { let sessionID = "session-tools-shared-no-starvation" _ = manager.session(id: sessionID) - await upstream.blockNextCancellation() let firstTask = Task { try await manager.sharedToolsList( @@ -7136,8 +7210,11 @@ struct RuntimeCoordinatorRecoveryTests { $0.waiterCounts.toolsCatalog == 1 } } + let firstLoad = try #require( + await manager.controlPlaneCoordinator.requestToolsCatalogLoadSnapshotForTesting() + ) - uptimeClock.advance(by: .nanoseconds(120_000_001)) + clocks.uptimeClock.advance(by: .nanoseconds(120_000_001)) let secondTask = Task { try await manager.sharedToolsList( @@ -7145,29 +7222,20 @@ struct RuntimeCoordinatorRecoveryTests { requestTimeoutOverride: .seconds(5) ) } - try await upstream.waitForBlockedCancellation() - let firstCancellation = try await waitWithTimeout( - "waiting for shared tools/list cancellation" - ) { - try await upstream.nextSent( - startingAt: 3, - matching: { methodName(from: $0) == "notifications/cancelled" } - ) - } - #expect( - try extractCancellationRequestID(from: firstCancellation) - == extractUpstreamID(from: firstRequest) - ) _ = try await waitWithTimeout("waiting for promoted tools/list waiters to attach") { try await manager.controlPlaneDebugMirror.waitForSnapshot { $0.waiterCounts.toolsCatalog == 2 } } - let sharedLoadID = try #require( - await manager.controlPlaneCoordinator.toolsCatalogLoad?.loadID + let sharedLoad = try #require( + await manager.controlPlaneCoordinator.requestToolsCatalogLoadSnapshotForTesting() ) + #expect(sharedLoad.loadID != firstLoad.loadID) + #expect(sharedLoad.waiterCount == 2) + #expect(sharedLoad.foregroundWaiterCount == 2) + #expect(firstLoad.rpcHandle.isCancelled()) - uptimeClock.advance(by: .nanoseconds(120_000_001)) + clocks.uptimeClock.advance(by: .nanoseconds(120_000_001)) let thirdTask = Task { try await manager.sharedToolsList( @@ -7181,8 +7249,13 @@ struct RuntimeCoordinatorRecoveryTests { $0.waiterCounts.toolsCatalog == 3 } } - #expect(await manager.controlPlaneCoordinator.toolsCatalogLoad?.loadID == sharedLoadID) - await upstream.releaseBlockedCancellation() + let unchangedLoad = try #require( + await manager.controlPlaneCoordinator.requestToolsCatalogLoadSnapshotForTesting() + ) + #expect(unchangedLoad.loadID == sharedLoad.loadID) + #expect(unchangedLoad.waiterCount == 3) + #expect(unchangedLoad.foregroundWaiterCount == 3) + #expect(sharedLoad.rpcHandle.isCancelled() == false) firstTask.cancel() secondTask.cancel() @@ -7190,6 +7263,16 @@ struct RuntimeCoordinatorRecoveryTests { _ = try? await firstTask.value _ = try? await secondTask.value _ = try? await thirdTask.value + _ = try await waitWithTimeout("waiting for shared tools/list load cancellation") { + try await manager.controlPlaneDebugMirror.waitForSnapshot { + $0.waiterCounts.toolsCatalog == 0 && $0.inFlightControlPlaneRequests.isEmpty + } + } + #expect( + await manager.controlPlaneCoordinator.requestToolsCatalogLoadSnapshotForTesting()? + .loadID == nil + ) + #expect(sharedLoad.rpcHandle.isCancelled()) } @Test func sessionManagerPromotedToolsListCancellationRemovesMigratedWaiter() async throws { @@ -7198,12 +7281,12 @@ struct RuntimeCoordinatorRecoveryTests { let eventLoop = group.next() let upstream = TestUpstreamClient() let config = makeConfig(requestTimeout: 5) - let uptimeClock = TestUptimeClock() + let clocks = makeRuntimeCoordinatorDeterministicClocks() let manager = RuntimeCoordinator( config: config, eventLoop: eventLoop, upstreams: [upstream], - nowUptimeNanoseconds: uptimeClock.now + clock: clocks.clock ) defer { manager.shutdownAndWait() } @@ -7232,8 +7315,11 @@ struct RuntimeCoordinatorRecoveryTests { $0.waiterCounts.toolsCatalog == 1 } } + let firstLoad = try #require( + await manager.controlPlaneCoordinator.requestToolsCatalogLoadSnapshotForTesting() + ) - uptimeClock.advance(by: .nanoseconds(120_000_001)) + clocks.uptimeClock.advance(by: .nanoseconds(120_000_001)) let secondTask = Task { try await manager.sharedToolsList( @@ -7241,17 +7327,18 @@ struct RuntimeCoordinatorRecoveryTests { requestTimeoutOverride: .seconds(5) ) } - _ = try await waitWithTimeout("waiting for migrated tools/list cancellation") { - try await upstream.nextSent( - startingAt: 3, - matching: { methodName(from: $0) == "notifications/cancelled" } - ) - } _ = try await waitWithTimeout("waiting for promoted tools/list waiters to attach") { try await manager.controlPlaneDebugMirror.waitForSnapshot { $0.waiterCounts.toolsCatalog == 2 } } + let promotedLoad = try #require( + await manager.controlPlaneCoordinator.requestToolsCatalogLoadSnapshotForTesting() + ) + #expect(promotedLoad.loadID != firstLoad.loadID) + #expect(promotedLoad.waiterCount == 2) + #expect(promotedLoad.foregroundWaiterCount == 2) + #expect(firstLoad.rpcHandle.isCancelled()) firstTask.cancel() do { @@ -7267,6 +7354,13 @@ struct RuntimeCoordinatorRecoveryTests { $0.waiterCounts.toolsCatalog == 1 } } + let remainingLoad = try #require( + await manager.controlPlaneCoordinator.requestToolsCatalogLoadSnapshotForTesting() + ) + #expect(remainingLoad.loadID == promotedLoad.loadID) + #expect(remainingLoad.waiterCount == 1) + #expect(remainingLoad.foregroundWaiterCount == 1) + #expect(promotedLoad.rpcHandle.isCancelled() == false) secondTask.cancel() do { @@ -7282,6 +7376,11 @@ struct RuntimeCoordinatorRecoveryTests { $0.waiterCounts.toolsCatalog == 0 && $0.inFlightControlPlaneRequests.isEmpty } } + #expect( + await manager.controlPlaneCoordinator.requestToolsCatalogLoadSnapshotForTesting()? + .loadID == nil + ) + #expect(promotedLoad.rpcHandle.isCancelled()) #expect(manager.debugSnapshot().upstreams[0].activeCorrelatedRequestCount == 0) } @@ -7290,6 +7389,7 @@ struct RuntimeCoordinatorRecoveryTests { defer { shutdownAndWait(group) } let eventLoop = group.next() let upstream = TestUpstreamClient() + let prewarmCompletions = LockedRecordedValues() var config = makeConfig(requestTimeout: 5) config.prewarmToolsList = true let clocks = makeRuntimeCoordinatorDeterministicClocks() @@ -7297,7 +7397,10 @@ struct RuntimeCoordinatorRecoveryTests { config: config, eventLoop: eventLoop, upstreams: [upstream], - clock: clocks.clock + clock: clocks.clock, + testHooks: RuntimeCoordinatorTestHooks( + toolsListPrewarmCompleted: { prewarmCompletions.append(()) } + ) ) defer { manager.shutdownAndWait() } @@ -7312,7 +7415,6 @@ struct RuntimeCoordinatorRecoveryTests { _ = try await initFuture.get() try await waitForSentCount(upstream, count: 2, timeoutSeconds: 2) - manager.refreshToolsListIfNeeded() let prewarmRequest = try await sentValue(from: upstream, at: 2, timeout: .seconds(2)) #expect(methodName(from: prewarmRequest) == "tools/list") await upstream.blockNextCancellation() @@ -7331,11 +7433,8 @@ struct RuntimeCoordinatorRecoveryTests { $0.waiterCounts.toolsCatalog == 1 } } - try await advanceRuntimeCoordinatorTimeout( - timeoutClock: clocks.timeoutClock, - uptimeClock: clocks.uptimeClock, - by: .seconds(5), - suspendedSleepers: 2 + #expect( + await manager.controlPlaneCoordinator.timeoutForegroundToolsCatalogWaiterForTesting() ) await #expect(throws: TimeoutError.self) { _ = try await firstTask.value @@ -7351,7 +7450,11 @@ struct RuntimeCoordinatorRecoveryTests { == extractUpstreamID(from: prewarmRequest) ) await upstream.releaseBlockedCancellation() - await manager.drainRuntimeTasksForTesting() + _ = try await waitForRecordedValue( + prewarmCompletions, + at: 0, + description: "waiting for timed-out tools/list prewarm completion" + ) #expect(manager.debugSnapshot().upstreams[0].activeCorrelatedRequestCount == 0) let secondTask = Task { @@ -7363,10 +7466,13 @@ struct RuntimeCoordinatorRecoveryTests { try await waitForSentCount(upstream, count: 5, timeoutSeconds: 2) let secondRequest = try await sentValue(from: upstream, at: 4, timeout: .seconds(2)) #expect(methodName(from: secondRequest) == "tools/list") - try await advanceRuntimeCoordinatorTimeout( - timeoutClock: clocks.timeoutClock, - uptimeClock: clocks.uptimeClock, - by: .seconds(5) + _ = try await waitWithTimeout("waiting for fresh tools/list waiter to attach") { + try await manager.controlPlaneDebugMirror.waitForSnapshot { + $0.waiterCounts.toolsCatalog == 1 + } + } + #expect( + await manager.controlPlaneCoordinator.timeoutForegroundToolsCatalogWaiterForTesting() ) await #expect(throws: TimeoutError.self) { _ = try await secondTask.value @@ -7515,10 +7621,13 @@ struct RuntimeCoordinatorRecoveryTests { try await waitForSentCount(upstream, count: 3, timeoutSeconds: 2) let firstRequest = try await sentValue(from: upstream, at: 2, timeout: .seconds(2)) #expect(methodName(from: firstRequest) == "tools/call") - try await advanceRuntimeCoordinatorTimeout( - timeoutClock: clocks.timeoutClock, - uptimeClock: clocks.uptimeClock, - by: .seconds(5) + _ = try await waitWithTimeout("waiting for first XcodeListWindows waiter to attach") { + try await manager.controlPlaneDebugMirror.waitForSnapshot { + $0.waiterCounts.windows == 1 + } + } + #expect( + await manager.controlPlaneCoordinator.timeoutWindowWaiterForTesting(route: .anyHealthy) ) await #expect(throws: TimeoutError.self) { _ = try await firstTask.value @@ -7552,10 +7661,13 @@ struct RuntimeCoordinatorRecoveryTests { try await waitForSentCount(upstream, count: 5, timeoutSeconds: 2) let secondRequest = try await sentValue(from: upstream, at: 4, timeout: .seconds(2)) #expect(methodName(from: secondRequest) == "tools/call") - try await advanceRuntimeCoordinatorTimeout( - timeoutClock: clocks.timeoutClock, - uptimeClock: clocks.uptimeClock, - by: .seconds(5) + _ = try await waitWithTimeout("waiting for second XcodeListWindows waiter to attach") { + try await manager.controlPlaneDebugMirror.waitForSnapshot { + $0.waiterCounts.windows == 1 + } + } + #expect( + await manager.controlPlaneCoordinator.timeoutWindowWaiterForTesting(route: .anyHealthy) ) await #expect(throws: TimeoutError.self) { _ = try await secondTask.value @@ -13469,12 +13581,12 @@ struct RuntimeCoordinatorWindowRoutingTests { let eventLoop = group.next() let upstream = TestUpstreamClient() let config = makeConfig(requestTimeout: 5) - let uptimeClock = TestUptimeClock() + let clocks = makeRuntimeCoordinatorDeterministicClocks() let manager = RuntimeCoordinator( config: config, eventLoop: eventLoop, upstreams: [upstream], - nowUptimeNanoseconds: uptimeClock.now + clock: clocks.clock ) defer { manager.shutdownAndWait() } @@ -13503,8 +13615,11 @@ struct RuntimeCoordinatorWindowRoutingTests { $0.waiterCounts.windows == 1 } } + let firstLoad = try #require( + await manager.controlPlaneCoordinator.windowLoadSnapshotForTesting(route: .anyHealthy) + ) - uptimeClock.advance(by: .nanoseconds(120_000_001)) + clocks.uptimeClock.advance(by: .nanoseconds(120_000_001)) let secondTask = Task { try await manager.liveXcodeListWindowsResult( @@ -13512,12 +13627,6 @@ struct RuntimeCoordinatorWindowRoutingTests { requestTimeoutOverride: .seconds(5) ) } - _ = try await waitWithTimeout("waiting for migrated XcodeListWindows cancellation") { - try await upstream.nextSent( - startingAt: 3, - matching: { methodName(from: $0) == "notifications/cancelled" } - ) - } _ = try await waitWithTimeout( "waiting for promoted XcodeListWindows waiters to attach" ) { @@ -13525,6 +13634,12 @@ struct RuntimeCoordinatorWindowRoutingTests { $0.waiterCounts.windows == 2 } } + let promotedLoad = try #require( + await manager.controlPlaneCoordinator.windowLoadSnapshotForTesting(route: .anyHealthy) + ) + #expect(promotedLoad.loadID != firstLoad.loadID) + #expect(promotedLoad.waiterCount == 2) + #expect(firstLoad.rpcHandle.isCancelled()) firstTask.cancel() do { @@ -13540,6 +13655,12 @@ struct RuntimeCoordinatorWindowRoutingTests { $0.waiterCounts.windows == 1 } } + let remainingLoad = try #require( + await manager.controlPlaneCoordinator.windowLoadSnapshotForTesting(route: .anyHealthy) + ) + #expect(remainingLoad.loadID == promotedLoad.loadID) + #expect(remainingLoad.waiterCount == 1) + #expect(promotedLoad.rpcHandle.isCancelled() == false) secondTask.cancel() do { @@ -13555,6 +13676,11 @@ struct RuntimeCoordinatorWindowRoutingTests { $0.waiterCounts.windows == 0 && $0.inFlightControlPlaneRequests.isEmpty } } + #expect( + await manager.controlPlaneCoordinator.windowLoadSnapshotForTesting(route: .anyHealthy)? + .loadID == nil + ) + #expect(promotedLoad.rpcHandle.isCancelled()) #expect(manager.debugSnapshot().upstreams[0].activeCorrelatedRequestCount == 0) }