Cover concurrent Gemini tier-fetch coalescing

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,

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<Void, Never>) in
+            var didEnqueue = false
             lock.withLock {
                 if !fetchInProgress || fetched {
                     continuation.resume()
                 } else {
                     fetchWaiters.append(continuation)
+                    didEnqueue = true
                 }
             }
+            if didEnqueue {
+                onWaiterEnqueued?()
+            }
         }
     }
 

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

CodexBarMacTests/NetworkTestSupport.swift

@@ -66,14 +66,72 @@ 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 {
+    deinit {}
+
+    private final class LoadContext: @unchecked Sendable {
+        deinit {}
+
+        private let lock = NSLock()
+        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) {
+            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 nil }
+                return (protocolInstance, 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<Void, Never>?
+    private var loadContext: LoadContext?
+    private var didStopLoading = false
+
     static func register(_ handler: @escaping IsolatedTestURLProtocolHandler, for handlerID: String) {
         lock.withLock {
             handlers[handlerID] = handler
@@ -103,17 +161,52 @@ 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 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()
+                context.deliver(response: response, data: data)
+            } catch is CancellationError {
+                return
+            } catch {
+                context.deliver(error: error)
+            }
+        }
+        let shouldCancel = loadingTaskLock.withLock {
+            guard !didStopLoading else { return true }
+            loadingTask = task
+            return false
+        }
+        if shouldCancel {
+            context.stop()
+            task.cancel()
         }
     }
 
-    override func stopLoading() {}
+    override func stopLoading() {
+        let (task, context) = loadingTaskLock.withLock {
+            didStopLoading = true
+            defer {
+                loadingTask = nil
+                loadContext = nil
+            }
+            return (loadingTask, loadContext)
+        }
+        task?.cancel()
+        context?.stop()
+    }
 }
 
 final class IsolatedTestURLSession: @unchecked Sendable {