From 9baeec9b5c4223aa08b68384dd1538751d12b003 Mon Sep 17 00:00:00 2001 From: HemSoft Date: Sat, 8 Aug 2026 03:26:43 -0400 Subject: [PATCH 1/3] Cover concurrent Gemini tier fetches --- CHANGELOG.md | 2 + .../Services/GeminiUsageProvider.swift | 37 ++++- CodexBarMacTests/GeminiProviderTests.swift | 126 ++++++++++++++++++ CodexBarMacTests/NetworkTestSupport.swift | 60 +++++++-- 4 files changed, 213 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 78864fa..358ee7e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -116,6 +116,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ### Developer Experience +- Added deterministic Gemini provider concurrency coverage for tier-fetch + coalescing, including the in-flight waiter and tier-derived project routing. - Added deterministic coordinator and GitHub Copilot concurrency coverage for account-scoped credential-refresh coalescing under parallel XCTest execution. - Gated provider test doubles now suspend in-flight refreshes deterministically, diff --git a/CodexBarMac/Services/GeminiUsageProvider.swift b/CodexBarMac/Services/GeminiUsageProvider.swift index 0631283..b085881 100644 --- a/CodexBarMac/Services/GeminiUsageProvider.swift +++ b/CodexBarMac/Services/GeminiUsageProvider.swift @@ -11,11 +11,12 @@ public final class GeminiUsageProvider: UsageProvider { private let projectsEndpoint: URL private let tokenEndpoint: URL private let now: @Sendable () -> Date + private let tierWaiterEnqueued: (@Sendable () -> Void)? private let tierCache = GeminiTierCache() public let providerID = ProviderID.gemini - public init( + public convenience init( session: URLSession = .shared, oauthFilePath: String = GeminiAuthFileStore.defaultPath(), settingsPath: String? = nil, @@ -24,6 +25,30 @@ public final class GeminiUsageProvider: UsageProvider { projectsEndpoint: URL = URL(string: "https://cloudresourcemanager.googleapis.com/v1/projects")!, tokenEndpoint: URL = GeminiTokenRefresh.tokenEndpoint, now: @escaping @Sendable () -> Date = { Date() } + ) { + self.init( + session: session, + oauthFilePath: oauthFilePath, + settingsPath: settingsPath, + quotaEndpoint: quotaEndpoint, + tierEndpoint: tierEndpoint, + projectsEndpoint: projectsEndpoint, + tokenEndpoint: tokenEndpoint, + now: now, + tierWaiterEnqueued: nil + ) + } + + init( + session: URLSession, + oauthFilePath: String, + settingsPath: String? = nil, + quotaEndpoint: URL, + tierEndpoint: URL, + projectsEndpoint: URL = URL(string: "https://cloudresourcemanager.googleapis.com/v1/projects")!, + tokenEndpoint: URL = GeminiTokenRefresh.tokenEndpoint, + now: @escaping @Sendable () -> Date = { Date() }, + tierWaiterEnqueued: (@Sendable () -> Void)? ) { self.session = session self.oauthFilePath = oauthFilePath @@ -37,6 +62,7 @@ public final class GeminiUsageProvider: UsageProvider { self.projectsEndpoint = projectsEndpoint self.tokenEndpoint = tokenEndpoint self.now = now + self.tierWaiterEnqueued = tierWaiterEnqueued } public func fetchUsage(for configuration: ProviderAccountConfiguration) async throws -> ProviderUsageResult { @@ -93,7 +119,7 @@ public final class GeminiUsageProvider: UsageProvider { ) async throws -> ProviderUsageResult { let fingerprint = credentialFingerprint(for: accessToken) await fetchTierIfNeeded(accessToken: accessToken, fingerprint: fingerprint) - await tierCache.waitForInFlightFetchIfNeeded() + await tierCache.waitForInFlightFetchIfNeeded(onWaiterEnqueued: tierWaiterEnqueued) let projectID = await resolveQuotaProjectID(accessToken: accessToken, fingerprint: fingerprint) let (data, response) = try await session.data(for: makeQuotaRequest(accessToken: accessToken, projectID: projectID)) @@ -593,20 +619,25 @@ private final class GeminiTierCache: @unchecked Sendable { resumeWaiters() } - func waitForInFlightFetchIfNeeded() async { + func waitForInFlightFetchIfNeeded(onWaiterEnqueued: (@Sendable () -> Void)?) async { let shouldWait = lock.withLock { fetchInProgress && !fetched } guard shouldWait else { return } await withCheckedContinuation { (continuation: CheckedContinuation) in + var didEnqueue = false lock.withLock { if !fetchInProgress || fetched { continuation.resume() } else { fetchWaiters.append(continuation) + didEnqueue = true } } + if didEnqueue { + onWaiterEnqueued?() + } } } diff --git a/CodexBarMacTests/GeminiProviderTests.swift b/CodexBarMacTests/GeminiProviderTests.swift index cad85be..b3256ae 100644 --- a/CodexBarMacTests/GeminiProviderTests.swift +++ b/CodexBarMacTests/GeminiProviderTests.swift @@ -899,6 +899,115 @@ final class GeminiProviderTests: XCTestCase { XCTAssertEqual(result.subtitle, "Live Gemini CLI usage") } + func testGeminiUsageProviderCoalescesConcurrentTierFetches() async throws { + let now = Date(timeIntervalSince1970: 2_000_000_000) + let directory = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString, isDirectory: true) + try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: directory) } + + let oauthFilePath = directory.appendingPathComponent("oauth_creds.json").path + try """ + { + "access_token": "redacted-access-token", + "refresh_token": "redacted-refresh-token", + "expiry_date": 4102444800000 + } + """.write(toFile: oauthFilePath, atomically: true, encoding: .utf8) + _ = chmod(oauthFilePath, 0o600) + + let tierResponseGate = TestAsyncGate() + let waiterEnqueued = TestSignal() + let requests = GeminiConcurrentRequestRecorder() + let isolatedSession = IsolatedTestURLSession { request in + guard let url = request.url else { + throw URLError(.badURL) + } + + switch url.path { + case "/gemini-tier": + await requests.recordTierRequest() + await tierResponseGate.wait() + return ( + HTTPURLResponse(url: url, statusCode: 200, httpVersion: nil, headerFields: nil)!, + Data( + #"{"currentTier":{"id":"standard-tier"},"cloudaicompanionProject":"gen-lang-client-coalesced"}"#.utf8 + ) + ) + case "/gemini-quota": + guard + let body = requestBodyData(from: request), + let json = try JSONSerialization.jsonObject(with: body) as? [String: String], + let projectID = json["project"] + else { + throw URLError(.cannotParseResponse) + } + await requests.recordQuotaRequest(projectID: projectID) + return ( + HTTPURLResponse(url: url, statusCode: 200, httpVersion: nil, headerFields: nil)!, + Data( + #"{"buckets":[{"tokenType":"REQUESTS","modelId":"gemini-2.5-pro","remainingFraction":0.8,"resetTime":"2026-07-17T12:00:00Z"}]}"#.utf8 + ) + ) + default: + throw URLError(.unsupportedURL) + } + } + defer { isolatedSession.invalidate() } + + let provider = GeminiUsageProvider( + session: isolatedSession.session, + oauthFilePath: oauthFilePath, + quotaEndpoint: URL(string: "https://example.test/gemini-quota")!, + tierEndpoint: URL(string: "https://example.test/gemini-tier")!, + now: { now }, + tierWaiterEnqueued: { + waiterEnqueued.signal() + } + ) + let configuration = ProviderAccountConfiguration.defaultConfiguration(for: .gemini) + + let results = try await withTestWatchdog( + timeout: .seconds(10), + failureMessage: "Concurrent Gemini tier fetches did not complete.", + onTimeout: { + Task { + await tierResponseGate.release() + } + } + ) { + let firstFetch = Task { + try await provider.fetchUsage(for: configuration) + } + await tierResponseGate.waitUntilBlocked() + + let secondFetch = Task { + try await provider.fetchUsage(for: configuration) + } + await waiterEnqueued.wait() + + let blockedSnapshot = await requests.snapshot() + XCTAssertEqual(blockedSnapshot.tierRequestCount, 1) + XCTAssertTrue(blockedSnapshot.quotaProjectIDs.isEmpty) + + await tierResponseGate.release() + return try await (firstFetch.value, secondFetch.value) + } + + let completedSnapshot = await requests.snapshot() + XCTAssertEqual(completedSnapshot.tierRequestCount, 1) + XCTAssertEqual( + completedSnapshot.quotaProjectIDs, + ["gen-lang-client-coalesced", "gen-lang-client-coalesced"] + ) + for result in [results.0, results.1] { + XCTAssertEqual(result.bars.count, 1) + XCTAssertEqual(result.bars[0].label, "Pro (Code Assist)") + XCTAssertEqual(result.bars[0].used, 0.2, accuracy: 0.0001) + XCTAssertEqual(result.subtitle, "Live Gemini CLI usage") + } + } + func testGeminiUsageProviderDiscoversProjectViaResourceManager() async throws { let now = Date(timeIntervalSince1970: 2_000_000_000) let directory = FileManager.default.temporaryDirectory @@ -1118,3 +1227,20 @@ final class GeminiProviderTests: XCTestCase { } } + +private actor GeminiConcurrentRequestRecorder { + private var tierRequestCount = 0 + private var quotaProjectIDs: [String] = [] + + func recordTierRequest() { + tierRequestCount += 1 + } + + func recordQuotaRequest(projectID: String) { + quotaProjectIDs.append(projectID) + } + + func snapshot() -> (tierRequestCount: Int, quotaProjectIDs: [String]) { + (tierRequestCount, quotaProjectIDs) + } +} diff --git a/CodexBarMacTests/NetworkTestSupport.swift b/CodexBarMacTests/NetworkTestSupport.swift index 612f0bb..849bc18 100644 --- a/CodexBarMacTests/NetworkTestSupport.swift +++ b/CodexBarMacTests/NetworkTestSupport.swift @@ -66,14 +66,28 @@ final class MockURLProtocol: URLProtocol, @unchecked Sendable { override func stopLoading() {} } -typealias IsolatedTestURLProtocolHandler = (URLRequest) throws -> (HTTPURLResponse, Data) +typealias IsolatedTestURLProtocolHandler = @Sendable (URLRequest) async throws -> (HTTPURLResponse, Data) private final class IsolatedTestURLProtocol: URLProtocol, @unchecked Sendable { + private final class LoadContext: @unchecked Sendable { + weak var protocolInstance: IsolatedTestURLProtocol? + weak var client: (any URLProtocolClient)? + + init(protocolInstance: IsolatedTestURLProtocol, client: (any URLProtocolClient)?) { + self.protocolInstance = protocolInstance + self.client = client + } + } + static let handlerIDHeader = "X-CodexBar-Test-Handler-ID" private static let lock = NSLock() nonisolated(unsafe) private static var handlers: [String: IsolatedTestURLProtocolHandler] = [:] + private let loadingTaskLock = NSLock() + private var loadingTask: Task? + private var didStopLoading = false + static func register(_ handler: @escaping IsolatedTestURLProtocolHandler, for handlerID: String) { lock.withLock { handlers[handlerID] = handler @@ -103,17 +117,45 @@ private final class IsolatedTestURLProtocol: URLProtocol, @unchecked Sendable { return } - do { - let (response, data) = try handler(request) - client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed) - client?.urlProtocol(self, didLoad: data) - client?.urlProtocolDidFinishLoading(self) - } catch { - client?.urlProtocol(self, didFailWithError: error) + let context = LoadContext(protocolInstance: self, client: client) + let capturedRequest = request + let task = Task { [context, capturedRequest, handler] in + do { + let (response, data) = try await handler(capturedRequest) + try Task.checkCancellation() + guard let protocolInstance = context.protocolInstance else { return } + context.client?.urlProtocol( + protocolInstance, + didReceive: response, + cacheStoragePolicy: .notAllowed + ) + context.client?.urlProtocol(protocolInstance, didLoad: data) + context.client?.urlProtocolDidFinishLoading(protocolInstance) + } catch is CancellationError { + return + } catch { + guard let protocolInstance = context.protocolInstance else { return } + context.client?.urlProtocol(protocolInstance, didFailWithError: error) + } + } + let shouldCancel = loadingTaskLock.withLock { + guard !didStopLoading else { return true } + loadingTask = task + return false + } + if shouldCancel { + task.cancel() } } - override func stopLoading() {} + override func stopLoading() { + let task = loadingTaskLock.withLock { + didStopLoading = true + defer { loadingTask = nil } + return loadingTask + } + task?.cancel() + } } final class IsolatedTestURLSession: @unchecked Sendable { From b3f25cf1ad61e00404c0429d60ab82ca8713e666 Mon Sep 17 00:00:00 2001 From: HemSoft Date: Sat, 8 Aug 2026 03:32:02 -0400 Subject: [PATCH 2/3] Harden isolated request cancellation --- CodexBarMacTests/NetworkTestSupport.swift | 77 +++++++++++++++++++---- 1 file changed, 64 insertions(+), 13 deletions(-) diff --git a/CodexBarMacTests/NetworkTestSupport.swift b/CodexBarMacTests/NetworkTestSupport.swift index 849bc18..194f327 100644 --- a/CodexBarMacTests/NetworkTestSupport.swift +++ b/CodexBarMacTests/NetworkTestSupport.swift @@ -69,14 +69,57 @@ final class MockURLProtocol: URLProtocol, @unchecked Sendable { typealias IsolatedTestURLProtocolHandler = @Sendable (URLRequest) async throws -> (HTTPURLResponse, Data) private final class IsolatedTestURLProtocol: URLProtocol, @unchecked Sendable { + deinit {} + private final class LoadContext: @unchecked Sendable { + deinit {} + + private let lock = NSRecursiveLock() weak var protocolInstance: IsolatedTestURLProtocol? weak var client: (any URLProtocolClient)? + private var isStopped = false init(protocolInstance: IsolatedTestURLProtocol, client: (any URLProtocolClient)?) { self.protocolInstance = protocolInstance self.client = client } + + func stop() { + lock.withLock { + isStopped = true + } + } + + func deliver(response: HTTPURLResponse, data: Data) { + lock.withLock { + guard + !isStopped, + let protocolInstance, + let client + else { return } + + client.urlProtocol( + protocolInstance, + didReceive: response, + cacheStoragePolicy: .notAllowed + ) + guard !isStopped else { return } + client.urlProtocol(protocolInstance, didLoad: data) + guard !isStopped else { return } + client.urlProtocolDidFinishLoading(protocolInstance) + } + } + + func deliver(error: any Error) { + lock.withLock { + guard + !isStopped, + let protocolInstance, + let client + else { return } + client.urlProtocol(protocolInstance, didFailWithError: error) + } + } } static let handlerIDHeader = "X-CodexBar-Test-Handler-ID" @@ -86,6 +129,7 @@ private final class IsolatedTestURLProtocol: URLProtocol, @unchecked Sendable { private let loadingTaskLock = NSLock() private var loadingTask: Task? + private var loadContext: LoadContext? private var didStopLoading = false static func register(_ handler: @escaping IsolatedTestURLProtocolHandler, for handlerID: String) { @@ -118,24 +162,26 @@ private final class IsolatedTestURLProtocol: URLProtocol, @unchecked Sendable { } let context = LoadContext(protocolInstance: self, client: client) + let shouldStart = loadingTaskLock.withLock { + guard !didStopLoading else { return false } + loadContext = context + return true + } + guard shouldStart else { + context.stop() + return + } + let capturedRequest = request let task = Task { [context, capturedRequest, handler] in do { let (response, data) = try await handler(capturedRequest) try Task.checkCancellation() - guard let protocolInstance = context.protocolInstance else { return } - context.client?.urlProtocol( - protocolInstance, - didReceive: response, - cacheStoragePolicy: .notAllowed - ) - context.client?.urlProtocol(protocolInstance, didLoad: data) - context.client?.urlProtocolDidFinishLoading(protocolInstance) + context.deliver(response: response, data: data) } catch is CancellationError { return } catch { - guard let protocolInstance = context.protocolInstance else { return } - context.client?.urlProtocol(protocolInstance, didFailWithError: error) + context.deliver(error: error) } } let shouldCancel = loadingTaskLock.withLock { @@ -144,16 +190,21 @@ private final class IsolatedTestURLProtocol: URLProtocol, @unchecked Sendable { return false } if shouldCancel { + context.stop() task.cancel() } } override func stopLoading() { - let task = loadingTaskLock.withLock { + let (task, context) = loadingTaskLock.withLock { didStopLoading = true - defer { loadingTask = nil } - return loadingTask + defer { + loadingTask = nil + loadContext = nil + } + return (loadingTask, loadContext) } + context?.stop() task?.cancel() } } From 712a5cffaf39c1bd13ae4a4c1581f33c7933b918 Mon Sep 17 00:00:00 2001 From: HemSoft Date: Sat, 8 Aug 2026 03:36:47 -0400 Subject: [PATCH 3/3] Let cancellation interrupt test callbacks --- CodexBarMacTests/NetworkTestSupport.swift | 42 +++++++++++------------ 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/CodexBarMacTests/NetworkTestSupport.swift b/CodexBarMacTests/NetworkTestSupport.swift index 194f327..a660e2f 100644 --- a/CodexBarMacTests/NetworkTestSupport.swift +++ b/CodexBarMacTests/NetworkTestSupport.swift @@ -74,7 +74,7 @@ private final class IsolatedTestURLProtocol: URLProtocol, @unchecked Sendable { private final class LoadContext: @unchecked Sendable { deinit {} - private let lock = NSRecursiveLock() + private let lock = NSLock() weak var protocolInstance: IsolatedTestURLProtocol? weak var client: (any URLProtocolClient)? private var isStopped = false @@ -91,33 +91,33 @@ private final class IsolatedTestURLProtocol: URLProtocol, @unchecked Sendable { } func deliver(response: HTTPURLResponse, data: Data) { - lock.withLock { - guard - !isStopped, - let protocolInstance, - let client - else { return } - - client.urlProtocol( - protocolInstance, - didReceive: response, - cacheStoragePolicy: .notAllowed - ) - guard !isStopped else { return } - client.urlProtocol(protocolInstance, didLoad: data) - guard !isStopped else { return } - client.urlProtocolDidFinishLoading(protocolInstance) - } + guard !Task.isCancelled, let (protocolInstance, client) = activeClient() else { return } + client.urlProtocol( + protocolInstance, + didReceive: response, + cacheStoragePolicy: .notAllowed + ) + + guard !Task.isCancelled, let (protocolInstance, client) = activeClient() else { return } + client.urlProtocol(protocolInstance, didLoad: data) + + guard !Task.isCancelled, let (protocolInstance, client) = activeClient() else { return } + client.urlProtocolDidFinishLoading(protocolInstance) } func deliver(error: any Error) { + guard !Task.isCancelled, let (protocolInstance, client) = activeClient() else { return } + client.urlProtocol(protocolInstance, didFailWithError: error) + } + + private func activeClient() -> (IsolatedTestURLProtocol, any URLProtocolClient)? { lock.withLock { guard !isStopped, let protocolInstance, let client - else { return } - client.urlProtocol(protocolInstance, didFailWithError: error) + else { return nil } + return (protocolInstance, client) } } } @@ -204,8 +204,8 @@ private final class IsolatedTestURLProtocol: URLProtocol, @unchecked Sendable { } return (loadingTask, loadContext) } - context?.stop() task?.cancel() + context?.stop() } }