Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
85 changes: 61 additions & 24 deletions Sources/CodexBarCore/Providers/Gemini/GeminiStatusProbe.swift
Original file line number Diff line number Diff line change
Expand Up @@ -337,25 +337,10 @@ public struct GeminiStatusProbe: Sendable {

let snapshot = try Self.parseAPIResponse(data, email: claims.email)

// Plan display strings with tier mapping:
// - standard-tier: Paid subscription (AI Pro, AI Ultra, Code Assist
// Standard/Enterprise, Developer Program Premium)
// - free-tier + hd claim: Workspace account (Gemini included free since Jan 2025)
// - free-tier: Personal free account (1000 req/day limit)
// - legacy-tier: Unknown legacy/grandfathered tier
// - nil (API failed): Leave blank (no display)
let plan: String? = switch (caStatus.tier, claims.hostedDomain) {
case (.standard, _):
"Paid"
case let (.free, .some(domain)):
{ Self.log.info("Workspace account detected", metadata: ["domain": domain]); return "Workspace" }()
case (.free, .none):
{ Self.log.info("Personal free account"); return "Free" }()
case (.legacy, _):
"Legacy"
case (.none, _):
{ Self.log.info("Tier detection failed, leaving plan blank"); return nil }()
}
let plan = Self.resolveAccountPlan(
tier: caStatus.tier,
hostedDomain: claims.hostedDomain,
paidTierName: caStatus.paidTierName)

return GeminiStatusSnapshot(
modelQuotas: snapshot.modelQuotas,
Expand Down Expand Up @@ -411,8 +396,9 @@ public struct GeminiStatusProbe: Sendable {
private struct CodeAssistStatus {
let tier: GeminiUserTierId?
let projectId: String?
let paidTierName: String?

static let empty = CodeAssistStatus(tier: nil, projectId: nil)
static let empty = CodeAssistStatus(tier: nil, projectId: nil, paidTierName: nil)
}

private static func loadCodeAssistStatus(
Expand Down Expand Up @@ -484,20 +470,26 @@ public struct GeminiStatusProbe: Sendable {
}

let tierId = (json["currentTier"] as? [String: Any])?["id"] as? String
let paidTierName = Self.parsePaidTierName(from: json)

guard let tierId else {
Self.log.warning("loadCodeAssist: no currentTier.id in response", metadata: [
"json": "\(json)",
])
return CodeAssistStatus(tier: nil, projectId: projectId)
return CodeAssistStatus(tier: nil, projectId: projectId, paidTierName: paidTierName)
}

guard let tier = GeminiUserTierId(rawValue: tierId) else {
Self.log.warning("loadCodeAssist: unknown tier ID", metadata: ["tierId": tierId])
return CodeAssistStatus(tier: nil, projectId: projectId)
return CodeAssistStatus(tier: nil, projectId: projectId, paidTierName: paidTierName)
}

Self.log.info("loadCodeAssist: success", metadata: ["tier": tierId, "projectId": projectId ?? "nil"])
return CodeAssistStatus(tier: tier, projectId: projectId)
Self.log.info("loadCodeAssist: success", metadata: [
"tier": tierId,
"projectId": projectId ?? "nil",
"paidTierName": paidTierName ?? "nil",
])
return CodeAssistStatus(tier: tier, projectId: projectId, paidTierName: paidTierName)
}

private struct OAuthCredentials {
Expand Down Expand Up @@ -1142,6 +1134,51 @@ public struct GeminiStatusProbe: Sendable {
}
}

extension GeminiStatusProbe {
/// Plan display strings with tier mapping:
/// - standard-tier: Paid subscription (Code Assist Standard/Enterprise, Developer Program Premium)
/// - free-tier + hd claim: Workspace account (Gemini included free since Jan 2025)
/// - free-tier + paidTier.name: Consumer subscription (Google AI Pro/Ultra, shown as Plus/Ultra)
/// - free-tier: Personal free account
/// - legacy-tier: Unknown legacy/grandfathered tier
/// - nil (API failed): Leave blank (no display)
fileprivate static func resolveAccountPlan(
tier: GeminiUserTierId?,
hostedDomain: String?,
paidTierName: String?) -> String?
{
switch (tier, hostedDomain) {
case (.standard, _):
return "Paid"
Comment thread
steipete marked this conversation as resolved.
case let (.free, .some(domain)):
Self.log.info("Workspace account detected", metadata: ["domain": domain])
return "Workspace"
case (.free, .none):
if let paidTierName {
self.log.info("Consumer paid tier detected", metadata: ["plan": paidTierName])
return paidTierName
}
Self.log.info("Personal free account")
return "Free"
case (.legacy, _):
return "Legacy"
case (.none, _):
self.log.info("Tier detection failed, leaving plan blank")
return nil
}
}

fileprivate static func parsePaidTierName(from json: [String: Any]) -> String? {
guard let paidTier = json["paidTier"] as? [String: Any],
let rawName = paidTier["name"] as? String
Comment on lines +1178 to +1179

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 Preserve paid-tier signals when names are omitted

When loadCodeAssist returns a paidTier object that has an id but no name, this parser returns nil, so resolveAccountPlan falls through to currentTier and can still display a paid consumer account as Free when currentTier.id is free-tier. Preserve paidTier.id or at least the presence of paidTier as a paid fallback when the specific display name is absent.

Useful? React with 👍 / 👎.

else {
return nil
}
let trimmed = rawName.trimmingCharacters(in: .whitespacesAndNewlines)
return trimmed.isEmpty ? nil : trimmed
}
}

extension GeminiStatusProbe {
private static let processTimeoutQueue = DispatchQueue(
label: "com.steipete.codexbar.gemini-process-timeout",
Expand Down
19 changes: 18 additions & 1 deletion Tests/CodexBarTests/GeminiAPITestHelpers.swift
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,11 @@ enum GeminiAPITestHelpers {
return "header.\(encoded).sig"
}

static func loadCodeAssistResponse(tierId: String, projectId: String? = nil) -> Data {
static func loadCodeAssistResponse(
tierId: String,
projectId: String? = nil,
paidTierName: String? = nil) -> Data
{
var payload: [String: Any] = [
"currentTier": [
"id": tierId,
Expand All @@ -86,9 +90,22 @@ enum GeminiAPITestHelpers {
if let projectId {
payload["cloudaicompanionProject"] = projectId
}
if let paidTierName {
payload["paidTier"] = [
"id": tierId,
"name": paidTierName,
]
}
return self.jsonData(payload)
}

static func loadCodeAssistConsumerPlusResponse(projectId: String? = "cloudaicompanion-123") -> Data {
self.loadCodeAssistResponse(
tierId: "free-tier",
projectId: projectId,
paidTierName: "Plus")
}

static func loadCodeAssistFreeTierResponse() -> Data {
self.loadCodeAssistResponse(tierId: "free-tier")
}
Expand Down
45 changes: 45 additions & 0 deletions Tests/CodexBarTests/GeminiStatusProbePlanTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -252,6 +252,51 @@ struct GeminiStatusProbePlanTests {
#expect(snapshot.accountPlan == "Workspace")
}

@Test
func `detects consumer plus from free tier with paid tier name`() async throws {
let env = try GeminiTestEnvironment()
defer { env.cleanup() }
let idToken = GeminiAPITestHelpers.makeIDToken(email: "user@gmail.com")
try env.writeCredentials(
accessToken: "token",
refreshToken: nil,
expiry: Date().addingTimeInterval(3600),
idToken: idToken)

let dataLoader = GeminiAPITestHelpers.dataLoader { request in
guard let url = request.url, let host = url.host else {
throw URLError(.badURL)
}
switch host {
case "cloudresourcemanager.googleapis.com":
return GeminiAPITestHelpers.response(
url: url.absoluteString,
status: 200,
body: GeminiAPITestHelpers.jsonData(["projects": []]))
case "cloudcode-pa.googleapis.com":
if url.path == "/v1internal:loadCodeAssist" {
return GeminiAPITestHelpers.response(
url: url.absoluteString,
status: 200,
body: GeminiAPITestHelpers.loadCodeAssistConsumerPlusResponse())
}
if url.path != "/v1internal:retrieveUserQuota" {
return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())
}
return GeminiAPITestHelpers.response(
url: url.absoluteString,
status: 200,
body: GeminiAPITestHelpers.sampleQuotaResponse())
default:
return GeminiAPITestHelpers.response(url: url.absoluteString, status: 404, body: Data())
}
}

let probe = GeminiStatusProbe(timeout: 1, homeDirectory: env.homeURL.path, dataLoader: dataLoader)
let snapshot = try await probe.fetch()
#expect(snapshot.accountPlan == "Plus")
}

@Test
func `detects free from free tier without hosted domain`() async throws {
let env = try GeminiTestEnvironment()
Expand Down
1 change: 1 addition & 0 deletions docs/gemini.md
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ Gemini uses the Gemini CLI OAuth credentials and private quota APIs. No browser
- Tier from `loadCodeAssist`:
- `standard-tier` → "Paid"
- `free-tier` + `hd` claim → "Workspace"
- `free-tier` + `paidTier.name` → consumer subscription label (e.g. "Plus", "Ultra")
- `free-tier` → "Free"
- `legacy-tier` → "Legacy"
- Email from `id_token` JWT claims.
Expand Down