From 216ead74aef53fac4086c3d6f191b3b6ef91a32d Mon Sep 17 00:00:00 2001
From: iamtoruk <hello@agentseal.org>
Date: Wed, 2 Sep 2026 07:52:18 -0700
Subject: [PATCH] fix(menubar): refresh(key:) reported failure instead of
 joining an in-flight fetch

The interactive refresh path bailed with false whenever the key was already
being fetched. recoverFromStuckLoading() reads that return value, so a healthy
fetch that was simply still running was announced as a failed recovery, and a
second interactive refresh threw away a result that was moments from landing.

Join the in-flight fetch and report on its outcome, which is what the quiet
path already did.

Adds a DataClient fetch hook under DEBUG: slot ownership and acceptance
ordering are only observable when two REAL fetches overlap, which needs a
fetch whose completion the test controls rather than a spawned CLI.
---
 mac/Sources/CodeBurnMenubar/AppStore.swift    |   9 +-
 .../CodeBurnMenubar/Data/DataClient.swift     |  13 ++
 .../AppStoreRefreshConcurrencyTests.swift     | 140 ++++++++++++++++++
 3 files changed, 161 insertions(+), 1 deletion(-)
 create mode 100644 mac/Tests/CodeBurnMenubarTests/AppStoreRefreshConcurrencyTests.swift

diff --git a/mac/Sources/CodeBurnMenubar/AppStore.swift b/mac/Sources/CodeBurnMenubar/AppStore.swift
index b31eeebe..da7beec2 100644
--- a/mac/Sources/CodeBurnMenubar/AppStore.swift
+++ b/mac/Sources/CodeBurnMenubar/AppStore.swift
@@ -1087,7 +1087,14 @@ final class AppStore {
            let cached = cache[key],
            cached.isFresh,
            !providerPayloadContradictsAll(cached.payload, for: key) { return true }
-        if inFlightKeys[key] != nil { return false }
+        // Join an in-flight fetch instead of reporting failure. Returning false
+        // here made recoverFromStuckLoading() announce a failed recovery while a
+        // perfectly healthy fetch was still running, which is what the quiet
+        // path already avoided by waiting.
+        if inFlightKeys[key] != nil {
+            await waitForInFlight(key)
+            return consistentCachedPayload(for: key) != nil
+        }
         inFlightKeys[key] = Date()
         attemptedKeys.insert(key)
         lastErrorByKey[key] = nil
diff --git a/mac/Sources/CodeBurnMenubar/Data/DataClient.swift b/mac/Sources/CodeBurnMenubar/Data/DataClient.swift
index fe86e171..3a8dc246 100644
--- a/mac/Sources/CodeBurnMenubar/Data/DataClient.swift
+++ b/mac/Sources/CodeBurnMenubar/Data/DataClient.swift
@@ -47,6 +47,14 @@ struct CLIDecodeFailure: Error, CustomStringConvertible {
 /// Runs the CLI via argv (no shell interpretation). See `CodeburnCLI` for why we never route
 /// commands through `/bin/zsh -c` anymore.
 struct DataClient {
+#if DEBUG
+    /// Test seam. The in-flight slot and the all-provider acceptance ordering
+    /// are only observable when two REAL fetches overlap, which needs a fetch
+    /// whose completion the test controls rather than a spawned CLI.
+    @MainActor
+    static var fetchHookForTesting: (@Sendable (ProviderFilter, Bool) async throws -> MenubarPayload)?
+#endif
+
     static func fetch(period: Period,
                       day: String? = nil,
                       days: Set<String> = [],
@@ -56,6 +64,11 @@ struct DataClient {
                       claudeConfigSourceId: String? = nil,
                       bypassResident: Bool = false,
                       qualityOfService: QualityOfService = .userInitiated) async throws -> MenubarPayload {
+#if DEBUG
+        if let hook = await MainActor.run(body: { fetchHookForTesting }) {
+            return try await hook(provider, bypassResident)
+        }
+#endif
         let subcommand = statusSubcommand(
             period: period,
             day: day,
diff --git a/mac/Tests/CodeBurnMenubarTests/AppStoreRefreshConcurrencyTests.swift b/mac/Tests/CodeBurnMenubarTests/AppStoreRefreshConcurrencyTests.swift
new file mode 100644
index 00000000..b65708ab
--- /dev/null
+++ b/mac/Tests/CodeBurnMenubarTests/AppStoreRefreshConcurrencyTests.swift
@@ -0,0 +1,140 @@
+import Foundation
+import Testing
+@testable import CodeBurnMenubar
+
+/// A fetch whose completion the test controls, so two real refresh tasks can be
+/// held in flight at once. Every attempt records the provider it was issued for
+/// and parks until the test releases it by index.
+@MainActor
+private final class ScriptedFetch {
+    private var parked: [Int: CheckedContinuation<Void, Never>] = [:]
+    private(set) var attempts: [(index: Int, provider: ProviderFilter)] = []
+    private(set) var released: [Int] = []
+    var payloadForProvider: (ProviderFilter) -> MenubarPayload = { _ in scriptedPayload(cost: 1) }
+
+    func install() {
+        DataClient.fetchHookForTesting = { [weak self] provider, _ in
+            guard let self else { return scriptedPayload(cost: 0) }
+            let index = await self.begin(provider)
+            await self.park(index)
+            return await self.payload(for: provider)
+        }
+    }
+
+    func uninstall() {
+        DataClient.fetchHookForTesting = nil
+        for (_, continuation) in parked { continuation.resume() }
+        parked.removeAll()
+    }
+
+    private func begin(_ provider: ProviderFilter) -> Int {
+        let index = attempts.count
+        attempts.append((index: index, provider: provider))
+        return index
+    }
+
+    private func park(_ index: Int) async {
+        await withCheckedContinuation { parked[index] = $0 }
+    }
+
+    private func payload(for provider: ProviderFilter) -> MenubarPayload {
+        payloadForProvider(provider)
+    }
+
+    func isParked(_ index: Int) -> Bool { parked[index] != nil }
+
+    func release(_ index: Int) {
+        guard let continuation = parked.removeValue(forKey: index) else { return }
+        released.append(index)
+        continuation.resume()
+    }
+
+    func indices(for provider: ProviderFilter) -> [Int] {
+        attempts.filter { $0.provider == provider }.map(\.index)
+    }
+}
+
+private func scriptedPayload(cost: Double,
+                             calls: Int = 1,
+                             providers: [String: Double]? = nil,
+                             providerDetails: [ProviderDetail] = []) -> MenubarPayload {
+    MenubarPayload(
+        generated: "test",
+        current: CurrentBlock(
+            label: "Today",
+            cost: cost,
+            calls: calls,
+            sessions: 1,
+            oneShotRate: nil,
+            inputTokens: 1,
+            outputTokens: 1,
+            cacheHitPercent: 0,
+            codexCredits: nil,
+            topActivities: [],
+            topModels: [],
+            localModelSavings: LocalModelSavings(totalUSD: 0, calls: 0, byModel: [], byProvider: []),
+            providers: providers ?? ["claude": cost],
+            providerDetails: providerDetails,
+            topProjects: [],
+            modelEfficiency: [],
+            topSessions: [],
+            retryTax: RetryTax(totalUSD: 0, retries: 0, editTurns: 0, byModel: []),
+            routingWaste: RoutingWaste(totalSavingsUSD: 0, baselineModel: "", baselineCostPerEdit: 0, byModel: []),
+            tools: [],
+            skills: [],
+            subagents: [],
+            mcpServers: []
+        ),
+        optimize: OptimizeBlock(findingCount: 0, savingsUSD: 0, topFindings: []),
+        history: HistoryBlock(daily: []),
+        combined: nil,
+        claudeConfigs: nil
+    )
+}
+
+/// Yield repeatedly until `condition` holds. The tasks under test suspend on
+/// real continuations, so progress needs actual scheduler turns, not a sleep.
+@MainActor
+private func settle(until condition: () -> Bool, turns: Int = 500) async -> Bool {
+    for _ in 0..<turns {
+        if condition() { return true }
+        await Task.yield()
+    }
+    return condition()
+}
+
+// Serialized: the DataClient fetch hook is process-global, so two of these
+// running at once would each see the other's attempts.
+@Suite("AppStore refresh concurrency", .serialized)
+@MainActor
+struct AppStoreRefreshConcurrencyTests {
+
+    @Test("an interactive refresh joins an in-flight fetch instead of reporting failure")
+    func interactiveRefreshJoinsInFlightFetch() async {
+        let script = ScriptedFetch()
+        script.install()
+        defer { script.uninstall() }
+
+        let store = AppStore()
+        store.setCacheDateToTodayForTesting()
+
+        let first = Task { await store.refresh(includeOptimize: false, force: true) }
+        #expect(await settle(until: { script.attempts.count == 1 }))
+
+        // Second interactive refresh on the same key. It used to return false
+        // immediately, which is what made recoverFromStuckLoading() announce a
+        // failed recovery while a perfectly healthy fetch was still running.
+        var secondResult: Bool?
+        let second = Task {
+            let value = await store.refresh(includeOptimize: false, force: true)
+            secondResult = value
+            return value
+        }
+        #expect(await settle(until: { store.inFlightWaiterCountForTesting(period: .today, provider: .all) == 1 }))
+        #expect(secondResult == nil)
+
+        script.release(0)
+        #expect(await first.value)
+        #expect(await second.value)
+    }
+}