Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

### Fixed
- Command Code: parse and display 5-hour and weekly rolling limits alongside monthly credits and reset times (#2466). Thanks @derekszen!
- OpenCode Go: include Zen balance in CLI usage reads without waiting beyond five seconds (#2583). Thanks @Yuxin-Qiao!
- Usage & Spend: keep validated Codex totals visible while the local scanner catches up, with refresh indicators in the dashboard and menu cost rows (#2397). Thanks @hhh2210!
- ZoomMate: preserve browser cookie scope so parent-domain sessions reach both API hosts without leaking host-only cookies (fixes #2507). Thanks @weddle!
- Sync: propagate provider configuration edits made by the CLI or directly in `config.json` to the iCloud fleet without echoing remotely applied writes.
Expand Down
1 change: 1 addition & 0 deletions Sources/CodexBarCLI/CLIUsageCommand.swift
Original file line number Diff line number Diff line change
Expand Up @@ -477,6 +477,7 @@ extension CodexBarCLI {
runtime: command.providerRuntime,
sourceMode: effectiveSourceMode,
includeCredits: command.includeCredits,
requiresOptionalUsageCompleteness: true,
webTimeout: command.webTimeout,
webDebugDumpHTML: command.webDebugDumpHTML,
verbose: command.verbose,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -172,7 +172,11 @@ struct OpenCodeGoLocalUsageFetchStrategy: ProviderFetchStrategy {
return nil
}
}
let zenBalance = try await OpenCodeGoUsageFetcher.completedOptionalZenBalance(from: zenBalanceTask)
let zenBalance = try await OpenCodeGoUsageFetcher.completedOptionalZenBalance(
from: zenBalanceTask,
timeout: OpenCodeGoUsageFetchStrategy.shouldWaitForZenBalance(context: context)
? .seconds(OpenCodeGoUsageFetcher.optionalZenBalanceTimeout)
: OpenCodeGoUsageFetcher.optionalZenBalanceJoinGrace)
return (snapshot.withZenBalanceUSD(zenBalance), false)
}

Expand All @@ -187,7 +191,8 @@ struct OpenCodeGoLocalUsageFetchStrategy: ProviderFetchStrategy {
cookieHeader: cookieHeader,
timeout: context.webTimeout,
workspaceIDOverride: workspaceOverride,
includeZenBalance: context.includeOptionalUsage)
includeZenBalance: context.includeOptionalUsage,
waitForZenBalance: OpenCodeGoUsageFetchStrategy.shouldWaitForZenBalance(context: context))
} catch OpenCodeGoUsageError.invalidCredentials {
throw OpenCodeGoUsageError.invalidCredentials
} catch is CancellationError {
Expand Down Expand Up @@ -226,6 +231,14 @@ struct OpenCodeGoUsageFetchStrategy: ProviderFetchStrategy {
let id: String = "opencodego.web"
let kind: ProviderFetchKind = .web

/// Usage-snapshot reads (`codexbar usage`, `codexbar serve`) are foreground commands, so a
/// Zen balance that is merely slower than the subscription page is worth waiting for, bounded
/// by the optional-balance timeout. Guard and diagnostic commands keep the short optional join
/// grace so a slow balance cannot consume their deadline.
static func shouldWaitForZenBalance(context: ProviderFetchContext) -> Bool {
context.requiresOptionalUsageCompleteness
}

func isAvailable(_ context: ProviderFetchContext) async -> Bool {
guard context.settings?.opencodego?.cookieSource != .off else { return false }
return true
Expand All @@ -241,7 +254,8 @@ struct OpenCodeGoUsageFetchStrategy: ProviderFetchStrategy {
cookieHeader: cookieHeader,
timeout: context.webTimeout,
workspaceIDOverride: workspaceOverride,
includeZenBalance: context.includeOptionalUsage)
includeZenBalance: context.includeOptionalUsage,
waitForZenBalance: Self.shouldWaitForZenBalance(context: context))
return self.makeResult(
usage: snapshot.toUsageSnapshot(),
sourceLabel: "web")
Expand All @@ -253,7 +267,8 @@ struct OpenCodeGoUsageFetchStrategy: ProviderFetchStrategy {
cookieHeader: cookieHeader,
timeout: context.webTimeout,
workspaceIDOverride: workspaceOverride,
includeZenBalance: context.includeOptionalUsage)
includeZenBalance: context.includeOptionalUsage,
waitForZenBalance: Self.shouldWaitForZenBalance(context: context))
return self.makeResult(
usage: snapshot.toUsageSnapshot(),
sourceLabel: "web")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,7 @@ public struct OpenCodeGoUsageFetcher: Sendable {
now: Date = Date(),
workspaceIDOverride: String? = nil,
includeZenBalance: Bool = true,
waitForZenBalance: Bool = false,
session: URLSession? = nil) async throws -> OpenCodeGoUsageSnapshot
{
let session = session ?? self.redirectGuardSession
Expand Down Expand Up @@ -194,7 +195,11 @@ public struct OpenCodeGoUsageFetcher: Sendable {
guard let zenBalanceTask else {
return snapshot
}
let zenBalance = try await self.completedOptionalZenBalance(from: zenBalanceTask)
let zenBalance = try await self.completedOptionalZenBalance(
from: zenBalanceTask,
timeout: waitForZenBalance
? .seconds(self.optionalZenBalanceTimeout)
: Self.optionalZenBalanceJoinGrace)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Start the optional Zen timeout when its task starts

When the subscription page is slow and the Zen request hangs, this five-second timer is not armed until after the subscription task has completed and been parsed, even though the Zen task started concurrently and still receives the caller's potentially 60-second timeout at line 157. In the current revision, a four-second subscription fetch can therefore be followed by another five-second join, making the CLI wait roughly nine seconds for a supposedly five-second-bounded optional request. Start the optional deadline with the Zen task, while restarting it with the configured timeout if it becomes the required Zen-only fallback, or subtract the task's elapsed time here.

Useful? React with 👍 / 👎.

return snapshot.withZenBalanceUSD(zenBalance)
}

Expand Down Expand Up @@ -236,18 +241,19 @@ public struct OpenCodeGoUsageFetcher: Sendable {
guard let requestCookieHeader = OpenCodeWebCookieSupport.requestCookieHeader(from: cookieHeader) else {
throw OpenCodeGoUsageError.invalidCredentials
}
let requestTimeout = min(timeout, self.optionalZenBalanceTimeout)
let workspaceID: String = if let override = self.normalizeWorkspaceID(workspaceIDOverride) {
override
} else {
try await self.fetchWorkspaceID(
cookieHeader: requestCookieHeader,
timeout: timeout,
timeout: requestTimeout,
session: session)
}
return try await self.fetchOptionalZenBalance(
workspaceID: workspaceID,
cookieHeader: requestCookieHeader,
timeout: min(timeout, self.optionalZenBalanceTimeout),
timeout: requestTimeout,
session: session)
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -115,10 +115,13 @@ extension OpenCodeGoUsageFetcher {
}
}

static func completedOptionalZenBalance(from task: Task<Double?, Error>) async throws -> Double? {
static func completedOptionalZenBalance(
from task: Task<Double?, Error>,
timeout: Duration? = Self.optionalZenBalanceJoinGrace) async throws -> Double?
{
let race = OpenCodeGoZenBalanceTaskRace(sourceTask: task)
do {
return try await race.value(timeout: self.optionalZenBalanceJoinGrace)
return try await race.value(timeout: timeout)
} catch is CancellationError {
throw CancellationError()
} catch {
Expand Down
7 changes: 7 additions & 0 deletions Sources/CodexBarCore/Providers/ProviderFetchPlan.swift
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,11 @@ public struct ProviderFetchContext: Sendable {
public let sourceMode: ProviderSourceMode
public let includeCredits: Bool
public let includeOptionalUsage: Bool
/// Whether this fetch should wait for optional usage data (such as prepaid balances) to
/// complete instead of bounding it with the short optional join grace. Usage-snapshot
/// reads enable this; guard and diagnostic commands keep the bounded join so a slow
/// optional request cannot consume their deadline.
public let requiresOptionalUsageCompleteness: Bool
public let webTimeout: TimeInterval
public let webDebugDumpHTML: Bool
public let verbose: Bool
Expand Down Expand Up @@ -55,6 +60,7 @@ public struct ProviderFetchContext: Sendable {
sourceMode: ProviderSourceMode,
includeCredits: Bool,
includeOptionalUsage: Bool = true,
requiresOptionalUsageCompleteness: Bool = false,
webTimeout: TimeInterval,
webDebugDumpHTML: Bool,
verbose: Bool,
Expand All @@ -75,6 +81,7 @@ public struct ProviderFetchContext: Sendable {
self.sourceMode = sourceMode
self.includeCredits = includeCredits
self.includeOptionalUsage = includeOptionalUsage
self.requiresOptionalUsageCompleteness = requiresOptionalUsageCompleteness
self.webTimeout = webTimeout
self.webDebugDumpHTML = webDebugDumpHTML
self.verbose = verbose
Expand Down
102 changes: 102 additions & 0 deletions Tests/CodexBarTests/OpenCodeGoOptionalZenBalanceTimeoutTests.swift
Original file line number Diff line number Diff line change
@@ -0,0 +1,102 @@
import Foundation
import Testing
@testable import CodexBarCore

private final class OptionalZenBalanceTimeoutRecorder<Value: Sendable>: @unchecked Sendable {
private let lock = NSLock()
private var storage: [Value] = []

func append(_ value: Value) {
self.lock.lock()
defer { self.lock.unlock() }
self.storage.append(value)
}

var values: [Value] {
self.lock.lock()
defer { self.lock.unlock() }
return self.storage
}
}

struct OpenCodeGoOptionalZenBalanceTimeoutTests {
@Test
func `optional zen balance caps workspace and balance request timeouts`() async throws {
defer {
OptionalZenBalanceTimeoutURLProtocol.handler = nil
}

let timeouts = OptionalZenBalanceTimeoutRecorder<TimeInterval>()
OptionalZenBalanceTimeoutURLProtocol.handler = { request in
guard let url = request.url else { throw URLError(.badURL) }
timeouts.append(request.timeoutInterval)
if url.path == "/_server" {
return Self.makeResponse(
url: url,
body: #"{"workspaces":[{"id":"wrk_TEST123"}]}"#,
contentType: "application/json")
}
#expect(url.path == "/workspace/wrk_TEST123")
return Self.makeResponse(
url: url,
body: #"<html><body><h2>Current balance $98.76</h2></body></html>"#,
contentType: "text/html")
}

let configuration = URLSessionConfiguration.ephemeral
configuration.protocolClasses = [OptionalZenBalanceTimeoutURLProtocol.self]
let balance = try await OpenCodeGoUsageFetcher.fetchOptionalZenBalance(
cookieHeader: "auth=test",
timeout: 60,
session: URLSession(configuration: configuration))

#expect(balance == 98.76)
#expect(timeouts.values == [5, 5])
}

private static func makeResponse(
url: URL,
body: String,
contentType: String) -> (HTTPURLResponse, Data)
{
let response = HTTPURLResponse(
url: url,
statusCode: 200,
httpVersion: "HTTP/1.1",
headerFields: ["Content-Type": contentType])!
return (response, Data(body.utf8))
}
}

private final class OptionalZenBalanceTimeoutURLProtocol: URLProtocol {
private static let handlerBox = LockIsolated<((URLRequest) throws -> (HTTPURLResponse, Data))?>(nil)
static var handler: ((URLRequest) throws -> (HTTPURLResponse, Data))? {
get { Self.handlerBox.value }
set { Self.handlerBox.setValue(newValue) }
}

override static func canInit(with request: URLRequest) -> Bool {
request.url?.host == "opencode.ai"
}

override static func canonicalRequest(for request: URLRequest) -> URLRequest {
request
}

override func startLoading() {
guard let handler = Self.handler else {
self.client?.urlProtocol(self, didFailWithError: URLError(.badServerResponse))
return
}
do {
let (response, data) = try handler(self.request)
self.client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
self.client?.urlProtocol(self, didLoad: data)
self.client?.urlProtocolDidFinishLoading(self)
} catch {
self.client?.urlProtocol(self, didFailWithError: error)
}
}

override func stopLoading() {}
}
15 changes: 14 additions & 1 deletion Tests/CodexBarTests/OpenCodeGoProviderStrategyTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -18,15 +18,18 @@ struct OpenCodeGoProviderStrategyTests {
}

private func makeContext(
runtime: ProviderRuntime = .app,
sourceMode: ProviderSourceMode = .auto,
requiresOptionalUsageCompleteness: Bool = false,
env: [String: String] = [:],
settings: ProviderSettingsSnapshot? = nil,
selectedTokenAccountID: UUID? = nil) -> ProviderFetchContext
{
ProviderFetchContext(
runtime: .app,
runtime: runtime,
sourceMode: sourceMode,
includeCredits: false,
requiresOptionalUsageCompleteness: requiresOptionalUsageCompleteness,
webTimeout: 1,
webDebugDumpHTML: false,
verbose: false,
Expand Down Expand Up @@ -141,4 +144,14 @@ struct OpenCodeGoProviderStrategyTests {
#expect(!strategy.shouldFallback(on: OpenCodeGoUsageError.networkError("timeout"), context: autoContext))
#expect(!strategy.shouldFallback(on: OpenCodeGoSettingsError.missingCookie, context: webContext))
}

@Test
func `web strategy waits for zen balance only on usage completeness reads`() {
#expect(!OpenCodeGoUsageFetchStrategy.shouldWaitForZenBalance(
context: self.makeContext(runtime: .app)))
#expect(!OpenCodeGoUsageFetchStrategy.shouldWaitForZenBalance(
context: self.makeContext(runtime: .cli)))
#expect(OpenCodeGoUsageFetchStrategy.shouldWaitForZenBalance(
context: self.makeContext(runtime: .cli, requiresOptionalUsageCompleteness: true)))
}
}
Loading