From 1d0703e0b90c7c739384383d27c67516869629a8 Mon Sep 17 00:00:00 2001
From: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com>
Date: Wed, 5 Aug 2026 12:24:44 +0800
Subject: [PATCH] Fix Antigravity agy cold-start quota readiness wait

---
 .../AntigravityProviderDescriptor.swift       | 33 ++++++++++++--
 ...ntigravityCLIHTTPSFetchStrategyTests.swift | 45 +++++++++++++++++++
 2 files changed, 74 insertions(+), 4 deletions(-)

diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift
index 790e74809e..0b6df19d09 100644
--- a/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift
+++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift
@@ -396,7 +396,8 @@ struct AntigravityCLIHTTPSFetchStrategy: ProviderFetchStrategy {
                 try await self.fetchBySpawning(
                     binary: binary,
                     idleWindow: idleWindow,
-                    resetAfterFetch: resetAfterFetch)
+                    resetAfterFetch: resetAfterFetch,
+                    expectedAccountEmail: expectedAccountEmail)
             })
     }
 
@@ -448,17 +449,23 @@ struct AntigravityCLIHTTPSFetchStrategy: ProviderFetchStrategy {
     private func fetchBySpawning(
         binary: String,
         idleWindow: TimeInterval?,
-        resetAfterFetch: Bool) async throws -> ProviderFetchResult
+        resetAfterFetch: Bool,
+        expectedAccountEmail: String?) async throws -> ProviderFetchResult
     {
         let session = AntigravityCLISession.shared
         let pid = try await session.beginProbe(binary: binary, idleWindow: idleWindow)
-        let deadline = Date().addingTimeInterval(5.0)
+        // Fresh `agy` processes take a few seconds to complete macOS keyring
+        // authentication, then more time before quota endpoints answer. A 5s
+        // window reliably missed that cold-start window in live tests, so keep
+        // the readiness deadline long enough for a cold spawn.
+        let deadline = Date().addingTimeInterval(15.0)
         let snap: AntigravityStatusSnapshot
         let usage: UsageSnapshot
         do {
             snap = try await Self.waitForSnapshot(
                 pid: pid,
                 deadline: deadline,
+                expectedAccountEmail: expectedAccountEmail,
                 dependencies: SnapshotWaitDependencies(
                     pollIntervalNanoseconds: 200_000_000,
                     listeningPorts: { pid, timeout in
@@ -500,6 +507,7 @@ struct AntigravityCLIHTTPSFetchStrategy: ProviderFetchStrategy {
     static func waitForSnapshot(
         pid: pid_t,
         deadline: Date,
+        expectedAccountEmail: String? = nil,
         dependencies: SnapshotWaitDependencies) async throws -> AntigravityStatusSnapshot
     {
         var lastFetchError: Error?
@@ -534,7 +542,24 @@ struct AntigravityCLIHTTPSFetchStrategy: ProviderFetchStrategy {
                 }
                 if let readySnapshot {
                     try await Self.checkAuthenticationPrompt(dependencies)
-                    return readySnapshot
+                    if AntigravitySelectedAccountGuard.matches(
+                        snapshotAccountEmail: readySnapshot.accountEmail,
+                        expectedAccountEmail: expectedAccountEmail)
+                    {
+                        return readySnapshot
+                    }
+                    // Fresh `agy` processes can answer quota endpoints before the
+                    // signed-in account email is available; keep polling so the
+                    // account guard does not reject the cold-start snapshot.
+                    lastFetchError = AntigravityStatusProbeError.accountMismatch(
+                        expected: expectedAccountEmail,
+                        found: readySnapshot.accountEmail)
+                    Self.log.debug(
+                        "Antigravity CLI HTTPS snapshot account not ready yet",
+                        metadata: [
+                            "pid": "\(pid)",
+                            "ports": ports.map(String.init).joined(separator: ","),
+                        ])
                 }
             }
 
diff --git a/Tests/CodexBarTests/AntigravityCLIHTTPSFetchStrategyTests.swift b/Tests/CodexBarTests/AntigravityCLIHTTPSFetchStrategyTests.swift
index 298229fa7b..9ee69a83e1 100644
--- a/Tests/CodexBarTests/AntigravityCLIHTTPSFetchStrategyTests.swift
+++ b/Tests/CodexBarTests/AntigravityCLIHTTPSFetchStrategyTests.swift
@@ -487,6 +487,51 @@ struct AntigravityCLIHTTPSFetchStrategyTests {
         #expect(snapshot.modelQuotas.first?.modelId == "claude-sonnet")
     }
 
+    @Test
+    func `cli HTTPS keeps waiting while snapshot account is not ready yet`() async throws {
+        let fetchAttempts = AntigravityCLICounter()
+
+        let snapshot = try await AntigravityCLIHTTPSFetchStrategy.waitForSnapshot(
+            pid: 123,
+            deadline: Date().addingTimeInterval(5),
+            expectedAccountEmail: "user@example.com",
+            dependencies: AntigravityCLIHTTPSFetchStrategy.SnapshotWaitDependencies(
+                pollIntervalNanoseconds: 0,
+                listeningPorts: { _, _ in [50080] },
+                drainOutput: { Data() },
+                fetchSnapshot: { _ in
+                    if fetchAttempts.increment() == 1 {
+                        return AntigravityStatusSnapshot(
+                            modelQuotas: [
+                                AntigravityModelQuota(
+                                    label: "Claude Sonnet",
+                                    modelId: "claude-sonnet",
+                                    remainingFraction: 0.5,
+                                    resetTime: nil,
+                                    resetDescription: nil),
+                            ],
+                            accountEmail: nil,
+                            accountPlan: "Pro",
+                            source: .local)
+                    }
+                    return AntigravityStatusSnapshot(
+                        modelQuotas: [
+                            AntigravityModelQuota(
+                                label: "Claude Sonnet",
+                                modelId: "claude-sonnet",
+                                remainingFraction: 0.5,
+                                resetTime: nil,
+                                resetDescription: nil),
+                        ],
+                        accountEmail: "user@example.com",
+                        accountPlan: "Pro",
+                        source: .local)
+                }))
+
+        #expect(fetchAttempts.value == 2)
+        #expect(snapshot.accountEmail == "user@example.com")
+    }
+
     @Test
     func `cli HTTPS drains output before ports appear`() async throws {
         let portPolls = AntigravityCLICounter()